C++ Interface

The C++ and Python interfaces of Python are designed to be as similar as possible. The following subsections explain how to translate from one language to the other along with a few unavoidable differences.

Use as a header library

By default, Dr.Jit operates as a header-only library without further dependencies.

#include <drjit/array.h>

namespace dr = drjit;

using Array3f  = dr::Array<float, 3>;
using Matrix3f = dr::Matrix<float, 3>;

int main(int, char**) {
    Array3f x(1, 2, 3),
            y(1, 0, 1);

    Array3f z = dr::normalize(dr::cross(x, y)) * .5f;
    printf("Result = %s\n", dr::string(z).c_str());
}

Vectorized loops

Analogous to the vectorized loop interface in Python (drjit.while_loop()), it is possible to run symbolic and evaluated loops in C++ via the drjit::while_loop() function. Use of this API requires the header file

#include <drjit/while_loop.h>

An example usage is shown below, which is essentially a translation of the Python interface to C++.

using Int = dr::CUDAArray<int>;

Int i = ..., j = ....;

std::tie(i, j) = dr::while_loop(
    // Initial loop state
    std::make_tuple(i, j),

    // while(..) condition of the loop
    [](const UInt &i, const UInt &) {
        return i < 5;
    },

    /// Loop body
    [](UInt &i, UInt &j) {
        i += 1;
        j = j * j;
    }
);

The short-hand notation provided through the @drjit.syntax decorator is not available in C++.

This syntax can be cumbersome when the loop state consists of many variables, since the associated type and argument names must be repeated several times. In this case, it can be preferable to locally define a custom struct representing the entirety of the loop state. See the explanation of custom traversable types for details. Dr.Jit optimizes away superfluous loop state variables, so there isn’t any harm in specifying loop state that isn’t actually modified by the loop.

The detailed interface of this function is as follows:

template<typename State, typename Cond, typename Body>
std::decay_t<State> while_loop(State &&state, Cond &&cond, Body &&body, const char *label = nullptr)

This function takes an instance state of the tuple type State (which could be a std::pair, std::tuple, or the lighter-weight alternative drjit::tuple created via drjit::make_tuple()).

It invokes the loop body body with an unpacked version of the tuple elements (i.e., body(std::get<0>(state), ...)) until the loop condition cond(std::get<0>(state), ...) equals false.

When the loop condition returns a scalar C++ bool, the operation compiles into an ordinary C++ loop. When it is a Dr.Jit array, the loop either runs in symbolic or evaluated mode. Please see the Python equivalent of this function (drjit.while_loop()) for details on what this means.

The label argument can be used to optionally specify a human-readable name that will be included in both low-level IR and GraphViz output.

Both cond and body may specify arbitrary callables (lambda functions, types with a custom operator() implementation). When such callables capture state from the surrounding call frame, it is important to note that Dr.Jit’s AD system may need to re-evaluate the loop at a later time, at which point the function which originally called drjit::while_loop() has itself returned. The &alpha variable captured by reference below would lead to undefined behavior in this case (i.e., it would likely crash your program).

int step = 123;

dr::while_loop(
    ...
    /// Loop body
    [&step](UInt &i) {
        i += step;
        ...
    }
    ...
);

Instead, capture relevant variable state by value or include it as part of state. Dr.Jit will move the two functions (cond and body including captured state) into a persistent object that will eventually be released by the AD backend when it is no longer needed.

Vectorized conditionals

Analogous to the vectorized conditional statement interface in Python (drjit.if_stmt()), it is possible to evaluate symbolic and evaluated conditionals in C++ via the drjit::if_stmt() function. Use of this API requires the header file

#include <drjit/if_stmt.h>

An example usage is shown below, which is essentially a translation of the Python interface to C++.

using Int = dr::CUDAArray<int>;

Int i = ..., j = ....;

Int abs_diff = dr::if_stmt(
    // 'args': arguments to forward to 'true_fn' and 'false_fn'
    std::make_tuple(i, j),

    // 'cond': conditional expression
    i < j,

    // 'true_fn': to be called for elements with 'cond == true'
    [](UInt i, UInt j) {
        return j - i;
    }

    // 'false_fn': to be called for elements with 'cond == false'
    [](UInt i, UInt j) {
        return i - j;
    }
);

The argument args must always be a tuple that will be unpacked and passed as arguments of true_fn and false_fn. The return value of these function can be any tree of arbitrarily nested arrays, tuples, and other custom data structures.

The short-hand notation provided through the @drjit.syntax decorator is not available in C++.

As with the loop, the C++ syntax can be cumbersome when there are many input/output variables, in which case it is preferable to locally define custom struct types representing all inputs/outputs, or both. Dr.Jit later optimizes away superfluous variables of drjit.if_stmt(), so there isn’t any harm in being redundant.

The detailed interface of this function is as follows:

template<typename Args, typename Mask, typename Body>
auto if_stmt(Args &&state, const Mask &cond, TrueFn &&true_fn, FalseFn &&false_fn, const char *label = nullptr)

This function takes an instance args of the tuple type Args (which could be a std::pair, std::tuple, or the lighter-weight alternative drjit::tuple created via drjit::make_tuple()).

It invokes true_fn and false_fn with an unpacked version of the tuple elements (i.e., true_fn(std::get<0>(state), ...)) and combines them based on the values of cond.

When the loop condition returns a scalar C++ bool, the operation compiles into an ordinary C++ conditional statement. When it is a Dr.Jit array, the loop either runs in symbolic or evaluated mode. Please see the Python equivalent of this function (drjit.if_stmt()) for details on what this means.

The label argument can be used to optionally specify a human-readable name that will be included in both low-level IR and GraphViz output.

The arguments true_fn and false_fn can be used to pass arbitrary callables (lambda functions, types with a custom operator() implementation). When such callables capture state from the surrounding call frame, it is important to note that Dr.Jit’s AD system may need to re-evaluate the conditional statement at a later time, at which point the function which originally called drjit::if_stmt() has itself returned. The &step variable captured by reference below would lead to undefined behavior in this case (i.e., it would likely crash your program).

int step = 123;

dr::if_stmt(
    ...
    /// true_fn
    [&step](UInt i) {
        return i + step;
    }
    ...
);

Instead, capture relevant variable state by value or include it as part of args. Dr.Jit will move the two functions (true_fn and false_fn including captured state) into a persistent object that will eventually be released by the AD backend when it is no longer needed.

Vectorized method calls

Given a pointer Foo* to a user-defined type Foo, a common operation in C++ entails dispatching a method or virtual method call:

Foo *ptr = ...;
float result = ptr->method(arg_1, arg_2, ...);

Dr.Jit also supports this operation in a vectorized form to dispatch method or virtual method calls to a large set instances in parallel:

using FooPtr = dr::CUDAArray<Foo *>;
using Float = dr::CUDAArray<float>;

FooPtr ptr = ...;
Float result = ptrs->f(arg_1, arg_2, ...);

It does so efficiently using at most a single invocation of each callable.

A limitation of array-based method calls is that input/output or output-only parameters passed using mutable references or pointers are not supported. Parameters are all inputs, and the function return value is the sole output—use pairs, tuples, or custom structures to return multiple values.

To enable this functionality for a new class, the following changes to its implementation are necessary:

First, include the header file

#include <drjit/call.h>

Next, modify the constructors and destructor of the class so that they register/unregister themselves with the Dr.Jit instance registry.

struct Foo {
    using Float = CUDAArray<float>;

    Foo() {
        jit_registry_put(dr::backend_v<Float>, "Foo", this);
    }

    virtual ~Foo() { jit_registry_remove(this); }

    /// Suppose this is a function implemented by subclasses of the ``Foo`` interface.
    virtual Float f(Float x) = 0;
};

The call to jit_registry_put must pass the backend (which can be manually specified or determined from a Dr.Jit array type via drjit::backend_v), a class name, and the this pointer.

Next, you use the following macros to describe the interface of the type. They must appear at the top level (i.e., outside of classes and namespaces) and simply list all function names that Dr.Jit should intercept.

DRJIT_CALL_BEGIN(Foo)
    DRJIT_CALL_METHOD(f)
    // Specify other methods here
DRJIT_CALL_END()

There is no need to specify return values, argument types, or multiple overloads. Just be sure to list each function that you want to be able to call on a Dr.Jit instance arrays. Below is an overview of the available macros:

DRJIT_CALL_BEGIN(Name)

Demarcates the start of an interface block. The Name parameter must refer to the type in question. The jit_registry_put call in the earlier snippet should provide the string-quoted equivalent of Name including namespace prefixes.

DRJIT_CALL_TEMPLATE_BEGIN(Name)

A variant of the above macro that should be used when Name refers to a template class.

DRJIT_CALL_END()

Demarcates the end of an interface block.

DRJIT_CALL_TEMPLATE_INHERITED_BEGIN(Name, Parent)

A variant of the above macro that should be used when Name refers to a template class that inherits from another template class Parent. The class Name should only be registered under the registry domain of Parent. The interface of Parent must have already been exposed with a call to DRJIT_CALL_TEMPLATE_BEGIN(Parent) ... DRJIT_CALL_END(). The interface will then be extended with any methods listed in this block. Note this feature is only safe to use with polymorphic types, where C++ can use RTTI to determine if a particular instance is of type Name.

DRJIT_CALL_INHERITED_END()

Demarcates the end of an interface block that started with DRJIT_CALL_TEMPLATE_INHERITED_BEGIN.

DRJIT_CALL_METHOD(Name)

Indicates to Dr.Jit that Name is the name of a method provided by the orginal type.

DRJIT_CALL_GETTER(Name)

This is an optimized form of the above macro that should be used when the function in question is a getter. This refers to a function that does not take in put arguments, and which is pure (i.e., causes no side effects). The implementation can then avoid the cost of an actual indirect jump.

Following these declarations, the following code performs a vectorized method or virtual method call.

dr::CUDAArray<Foo*> instances = ...;
Float x = ....;
Float y = instances->f(x);

All of the commentary about function calls in Python (see drjit.switch()) applies here as well. The call can be done in symbolic or evaluated mode, and it supports derivative propagation in forward and reverse modes.

Masks passed as the last function argument are treated specially and apply to the entire operation. Masked elements of the call effectively don’t perform the function call at all, and their return value is zero. Side effects performed by the called functions are also disabled for these elements.

It is legal to perform a function call on an array containing nullptr pointers. These elements are considered to be masked as well.

Besides vectorizing method calls, there is an alternative call interface named drjit::dispatch:

template<typename Self, typename Func, typename ...Args>
auto dispatch(const Self &self, const Func &func, const Args&... args)

This C++ interface is analogous the Python API function drjit.dispatch(). Please review the documentation of the Python variant first.

The C++ interface takes an instance array self and invokes a provided callable func once for each unique instance. The callable should take a scalar instance pointer as first input argument, followed by args.... Dr.Jit assembles the traced computation into an indirect jump to one of several subroutines.

The dispatch interface is convenient whenever adding a method or virtual method to a class is undesirable. Also, dynamic dispatch is a relatively costly operation. When multiple calls are performed on the same set of instances, it may be preferable to merge them into a single and potentially signficantly faster use of drjit::dispatch().

The following snippet shows how merge two method calls into a combined call:

dr::CUDAArray<Foo*> instances = ...;
Float x = ....;

Float y = dr::dispatch(
    instances,
    [](Foo *self, Float x) { return self->f(x) + self->g(x); },
    x
);

Python bindings

Regular arrays

It is possible to register additional bindings for custom array types from within user code. The test case tests/custom_type_ext.cpp shows how to define a Color3f type that exposes .r, .g, and .b members.

Instance arrays

Suppose you have created a C++ type with the following signature:

using Float = dr::DiffArray<JitBackend::CUDA, float>;

struct Foo {
    virtual Float f(Float input) const = 0;
    virtual ~Foo() = default;
};

The nanobind description to expose this type in Python is as follows:

nb::class_<Foo>(m, "Foo")
    .def("f", &Foo::f);

It can also be useful to create similar bindings for Dr.Jit Foo instance arrays that automatically dispatch function calls to the f method. To do so, include

#include <drjit/python.h>

and append the following binding declarations:

using FooPtr = dr::CUDADiffArray<Foo *>;

dr::ArrayBinding b;
auto base_ptr = dr::bind_array_t<FooPtr>(b, m, "FooPtr")
    .def("f", [](FooPtr &self, Float a) { return self->f(a); })

Custom data structures

The ability to traverse through members of custom data structures was previously discussed here in the context of Python.

This feature also exists on the C++ side. For this, you must include the header file

#include <drjit/struct.h>

Following this, you can use the variable-argument DRJIT_STRUCT(...) macro to list the available fields.

using Float = dr::CUDADiffArray<float>;

struct MyPoint2f {
    Float x;
    Float y;

    DRJIT_STRUCT(x, y);
};

Custom data structures can be defined globally, locally, and they can be arbitrarily nested. Dr.Jit static arrays, tuples (drjit::tuple), STL std::tuple<...> and std::pair<T1, T2> are all automatically traversable.

It is also easy to define custom functions that recursively process arbitrary trees:

template <typename T> void visit_jit_pairs(T &v0, T &v1) {
    if constexpr (dr::is_jit_v<T> && dr::depth_v<T> == 1) {
        /// Do something with 'v0' and 'v1'
    } else if constexpr (dr::is_traversable_v<T>) {
        /// Recurse and try again if the object is traversable
        dr::traverse_2(
            /// Extract the fields of 'v0' and 'v1'
            dr::fields(v0), dr::fields(v1),
            // .. and call the following lambda function on them
            [&](auto &x, auto &y) { visit_jit_pairs(x, y); }
        );
    }
}

The type trait dr::is_traversable<T> checks if an instance of a particular type can be traversed.

The helper functions drjit::traverse_1() drjit::traverse_2(), drjit::traverse_3(), respectively traverse individual objects, pairs, or triples, at the same time. They take an arbitrary stateless or stateful lambda function and simply apply it to the input tuples that can be extracted using the function drjit::fields() that takes a traversable as input.