EnzymeAD / EnzymeAD/Reactant.jl

Sincos NVIDIA intrinsic raising issue `__nv_sincos`

Open
#2,202 0 comments 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.