I would like to apply/use NNs within Reactant/GPU kernels. Currently this doesn't seem to work. Searching a bit through the repo, I found an old documentation on this https://github.com/LuxDL/Lux.jl/blob/10ea255face7fdd6b198cf9e2e641a8414afe037/docs/src/manual/nn_inside_gpu_kernels.md
Is this not supported anymore? Or did I do something wrong.
Below the example I tried (with LLM help)
using Reactant, CUDA, Lux, Random, StaticArrays
using KernelAbstractions: @kernel, @index, get_backend
import Lux.MLDataDevices as MLDD
const dev = MLDD.reactant_device()
# A tiny MLP (1 → 4 → 1); move parameters and data to the device the standard Lux way.
mlp = Lux.Chain(Lux.Dense(1 => 4, tanh), Lux.Dense(4 => 1))
ps, st = Lux.setup(Random.Xoshiro(0), mlp) |> dev
x = collect(Float32, -3:0.5:3) |> dev
out = zero(x)
# ---------------------------------------------------------------------------------------------
# Control: the MLP's forward pass written out by hand (from the Lux parameters) compiles and runs.
# ---------------------------------------------------------------------------------------------
@kernel function manual_kernel!(out, @Const(x), @Const(W1), @Const(b1), @Const(W2), @Const(b2))
i = @index(Global)
xi = @inbounds x[i]
acc = @inbounds b2[1]
for h in 1:length(b1)
acc += @inbounds W2[1, h] * tanh(W1[h, 1] * xi + b1[h])
end
@inbounds out[i] = acc
end
run_manual!(out, x, W1, b1, W2, b2) =
(manual_kernel!(get_backend(out))(out, x, W1, b1, W2, b2; ndrange = length(out)); nothing)
W1, b1 = ps.layer_1.weight, ps.layer_1.bias
W2, b2 = ps.layer_2.weight, ps.layer_2.bias
compiled_manual! = @compile raise = true run_manual!(out, x, W1, b1, W2, b2)
compiled_manual!(out, x, W1, b1, W2, b2)
@info "control OK: hand-written forward compiled and ran" result = Array(out)
# ---------------------------------------------------------------------------------------------
# The failure: the same MLP applied with `Lux.apply` inside the kernel fails to compile.
# ---------------------------------------------------------------------------------------------
@kernel function apply_kernel!(out, @Const(x), model, ps, st)
i = @index(Global)
y, _ = Lux.apply(model, SA[x[i]], ps, st)
@inbounds out[i] = y[1]
end
run_apply!(out, x, model, ps, st) =
(apply_kernel!(get_backend(out))(out, x, model, ps, st; ndrange = length(out)); nothing)
@info "compiling the Lux.apply kernel (expected to fail with an InvalidIRError) ..."
@compile raise = true run_apply!(out, x, mlp, ps, st)This gives a lot of InvalidIRError with Reason: unsupported call to a lazy-initialized function (call to utf8proc_isupper) from different parts of the code.
About the scope/context of this: I am interested in hybrid model combining physics with ML. One kind would just replace a point wise / column-wise parameterisation with an MLPs. The easiest way to integrate this within our models would be to just call them within the existing kernels. This approach is really limited to simple NN architectures like MLPs or maybe simple RNNs as well.
I would like to apply/use NNs within Reactant/GPU kernels. Currently this doesn't seem to work. Searching a bit through the repo, I found an old documentation on this https://github.com/LuxDL/Lux.jl/blob/10ea255face7fdd6b198cf9e2e641a8414afe037/docs/src/manual/nn_inside_gpu_kernels.md
Is this not supported anymore? Or did I do something wrong.
Below the example I tried (with LLM help)
This gives a lot of
InvalidIRErrorwithReason: unsupported call to a lazy-initialized function (call to utf8proc_isupper)from different parts of the code.About the scope/context of this: I am interested in hybrid model combining physics with ML. One kind would just replace a point wise / column-wise parameterisation with an MLPs. The easiest way to integrate this within our models would be to just call them within the existing kernels. This approach is really limited to simple NN architectures like MLPs or maybe simple RNNs as well.