EnzymeAD / EnzymeAD/Reactant.jl
Sincos NVIDIA intrinsic raising issue `__nv_sincos`
- Dominant language
- Julia
- Stars
- 370
- Forks
- 74
- Avg merge
- 18h 47m
- Merged PRs (30d)
- 30
Description
Ties into existing Enzyme issue EnzymeAD/Enzyme#2552, for which I opened a PR with a potential fix EnzymeAD/Enzyme#2669.
In Enzyme, we can work around this issue by using `@noinline sincos`, but for Reactant, this fails as well.
MWE:
```julia
using Reactant: @compile, CompileOptions, to_rarray, set_default_backend
using KernelAbstractions: @kernel, @index, @Const, get_backend
import Enzyme, CUDA
set_default_backend("gpu")
@noinline sincos_noinline(x) = sincos(x)
@kernel function separate_sincos_kernel!(y, @Const(x))
i = @index(Global)
if i <= length(x)
s = sin(x[i])
c = cos(x[i])
y[i] = s + c
end
end
@kernel function sincos_kernel!(y, @Const(x))
i = @index(Global)
if i <= length(x)
s, c = sincos(x[i])
y[i] = s + c
end
end
@kernel function noinline_sincos_kernel!(y, @Const(x))
i = @index(Global)
if i <= length(x)
s, c = sincos_noinline(x[i])
y[i] = s + c
end
end
function apply_separate(x)
y = similar(x)
kernel! = separate_sincos_kernel!(get_backend(x), 256)
kernel!(y, x; ndrange=length(x))
return y
end
function apply_sincos(x)
y = similar(x)
kernel! = sincos_kernel!(get_backend(x), 256)
kernel!(y, x; ndrange=length(x))
return y
end
function apply_noinline(x)
y = similar(x)
kernel! = noinline_sincos_kernel!(get_backend(x), 256)
kernel!(y, x; ndrange=length(x))
return y
end
loss_separate(x) = sum(apply_separate(x))
loss_sincos(x) = sum(apply_sincos(x))
loss_noinline(x) = sum(apply_noinline(x))
x = to_rarray(randn(Float32, 100))
compile_opts = CompileOptions(; raise=true, raise_first=true, sync=true)
println("Test 1: Separate sin() + cos()")
try
f = @compile compile_options=compile_opts Enzyme.gradient(Enzyme.Reverse, loss_separate, x)
f(Enzyme.Reverse, loss_separate, x)
println("SUCCESS")
catch e
println("FAILED: $(typeof(e))")
end
println("\nTest 2: sincos()")
try
f = @compile compile_options=compile_opts Enzyme.gradient(Enzyme.Reverse, loss_sincos, x)
f(Enzyme.Reverse, loss_sincos, x)
println("SUCCESS")
catch e
println("FAILED: $(typeof(e))")
end
println("\nTest 3: @noinline sincos()")
try
f = @compile compile_options=compile_opts Enzyme.gradient(Enzyme.Reverse, loss_noinline, x)
f(Enzyme.Reverse, loss_noinline, x)
println("SUCCESS")
catch e
println("FAILED: $(typeof(e))")
end
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.