Computational kernels

Suppose your code base contains custom computational kernels, such as GPU kernels defined with KernelAbstractions.jl or directly with a backend like CUDA.jl.

Example

using KernelAbstractions

Here we define a very simple squaring kernel:

@kernel function square_kernel!(y, @Const(x))
    i = @index(Global)
    @inbounds y[i] = x[i] * x[i]
end

function square(x)
    y = similar(x)
    backend = KernelAbstractions.get_backend(x)
    kernel! = square_kernel!(backend)
    kernel!(y, x; ndrange=length(x))
    return y
end

square (generic function with 1 method)

Let's test it to make sure it works:

x = float.(1:5)
y = square(x)

Kernel compilation

To compile this kernel with Reactant, the CUDA.jl package needs to be loaded (even on non-NVIDIA hardware).

import CUDA
using Reactant

The rest of the compilation works as usual:

xr = ConcreteRArray(x)
square_compiled = @compile square(xr)

Reactant compiled function square (with tag ##square_reactant#171485)

yr = square_compiled(xr)

The Reactant-compiled function square_compiled now runs on whatever device you request from Reactant, including CPU, GPU, TPU or distributed settings. It will not run on the device it was written for, nor will it require a CUDA-enabled device.

Kernel raising

Reactant has the ability to detect and optimize high-level tensor representations of existing kernels, through a process called raising. For more information, see the corresponding documentation.

Kernel differentiation

You can combine with Enzyme.jl to take derivatives of kernels. For more information of differntiating generic code with Enzyme+Reactant, see the corresponding documentation.

import Enzyme

When differentiating computational kernels, Enzyme can either differentiate the kernel directly, or differentiated the raised linear algebra. Using Enzyme outside of Reactant, Enzyme.jl differentiates the kernel directly. Currently within Reactant, only differentiating the raised code is supported (though supporting both is in progress). As a result, you must use the raise = true and raise_first = true compilation options to make sure the kernel is raised before Enzyme performs automatic differentiation on the program.

sumsquare(x) = sum(square(x))
gradient_compiled = @compile raise=true raise_first=true Enzyme.gradient(Enzyme.Reverse, sumsquare, xr)

Reactant compiled function gradient (with tag ##gradient_reactant#171498)

Note that the mode and function argument are partially evaluated at compilation time, but we still need to provide them again at execution time:

gr = gradient_compiled(Enzyme.Reverse, sumsquare, xr)[1]