JIT backend

Overview

JIT backends execute a recorded JITGraph. XAD exposes:

• JITBackend<Scalar>: abstract execution interface
• JITGraphInterpreter<Scalar>: reference backend that interprets the graph

JITBackend

template <class Scalar> class JITBackend

Abstract base class for JIT execution backends. The Scalar template parameter specifies the floating-point type used for computations and must match the type used by JITCompiler<Real, N>:

• JITBackend<double> for use with JITCompiler<double, N>
• JITBackend<float> for use with JITCompiler<float, N>

Most applications use double precision. The float option is available for performance-sensitive applications where single precision is sufficient.

Backends are responsible for:

• compiling a JITGraph into an executable form (optional)
• executing forward evaluation
• optionally executing forward+backward to produce input adjoints

Main virtual functions

compile

virtual void compile(const JITGraph& graph) = 0;

Prepare the backend for executing the given graph.

vectorWidth

virtual std::size_t vectorWidth() const = 0;

Returns the number of parallel evaluations per execution (1 for scalar backends, 4 for AVX2).

numInputs / numOutputs

virtual std::size_t numInputs() const = 0; virtual std::size_t numOutputs() const = 0;

Return the number of inputs/outputs in the compiled graph.

setInput

virtual void setInput(std::size_t inputIndex, const Scalar* values) = 0;

Set input values for an input variable. The values array must contain vectorWidth() elements.

forward

virtual void forward(Scalar* outputs) = 0;

Run a forward pass. The outputs array must have space for numOutputs() * vectorWidth() elements.

forwardAndBackward

virtual void forwardAndBackward(Scalar* outputs, Scalar* inputGradients) = 0;

Run forward and backward passes combined, with all output adjoints seeded to 1. The outputs array must have space for numOutputs() * vectorWidth() elements, and inputGradients must have space for numInputs() * vectorWidth() elements.

forwardAndBackwardSeeded

virtual void forwardAndBackwardSeeded(Scalar* outputs, const Scalar* outputAdjoints, Scalar* inputGradients);

As forwardAndBackward, but the backward pass starts from the given output adjoints (numOutputs() * vectorWidth() values). JITCompiler::computeAdjoints uses it.

It is not pure virtual: the default implementation calls forwardAndBackward and scales the gradients, which works only when all output adjoints are equal; otherwise it throws std::runtime_error. Backends that support different output adjoints override it, as JITGraphInterpreter does.

reset

virtual void reset() = 0;

Reset any cached backend state.

JITGraphInterpreter

template <class Scalar> class JITGraphInterpreter : public JITBackend<Scalar>

Reference backend that interprets the graph directly, one evaluation at a time (vectorWidth() is 1). It is mainly intended as:

• a correctness reference implementation
• a simple fallback backend
• a baseline for testing and debugging

Example Usage

For double backend:

auto backend = std::make_unique<xad::JITGraphInterpreter<double>>();
xad::JITCompiler<double, 1> jit(std::move(backend));
// ... record graph ...
jit.compile();

For float backend:

auto backend = std::make_unique<xad::JITGraphInterpreter<float>>();
xad::JITCompiler<float, 1> jit(std::move(backend));
// ... record graph ...
jit.compile();