Fix atomic store in cuda (#1263)

* Fix atomic store in cuda

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update cuda.jl

* Update cuda.jl

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: wsmoses

ext/ReactantCUDAExt.jl

@@ -46,7 +46,12 @@ struct CuTracedRNumber{T,A} <: Number
     end
 end
 
-CuTracedRNumber{T,A}(val::Number) where {T,A} = convert(CuTracedRNumber{T,A}, val)
+Base.@nospecializeinfer Reactant.is_traced_number(
+    @nospecialize(T::Type{<:CuTracedRNumber})
+) = true
+Reactant.unwrapped_eltype(::Type{<:CuTracedRNumber{T}}) where {T} = T
+
+@inline CuTracedRNumber{T,A}(val::Number) where {T,A} = convert(CuTracedRNumber{T,A}, val)
 
 function Base.getindex(RN::CuTracedRNumber{T,A}) where {T,A}
     align = alignment(RN)
@@ -99,13 +104,13 @@ Base.OneTo(x::CuTracedRNumber{<:Integer}) = Base.OneTo(x[])
     end
 end
 
-function Base.convert(CT::Type{CuTracedRNumber{Float64,1}}, x::Number)
+@inline function Base.convert(CT::Type{CuTracedRNumber{Float64,1}}, x::Number)
     return CT(
         Base.llvmcall(
             (
                 """define double addrspace(1)* @entry(double %d) alwaysinline {
           %a = alloca double
-          store double %d, double* %a
+          store atomic double %d, double* %a release, align 8
           %ac = addrspacecast double* %a to double addrspace(1)*
           ret double addrspace(1)* %ac
                     }
@@ -119,13 +124,13 @@ function Base.convert(CT::Type{CuTracedRNumber{Float64,1}}, x::Number)
     )
 end
 
-function Base.convert(CT::Type{CuTracedRNumber{Float32,1}}, x::Number)
+@inline function Base.convert(CT::Type{CuTracedRNumber{Float32,1}}, x::Number)
     return CT(
         Base.llvmcall(
             (
                 """define float addrspace(1)* @entry(float %d) alwaysinline {
           %a = alloca float
-          store float %d, float* %a
+          store atomic float %d, float* %a release, align 4
           %ac = addrspacecast float* %a to float addrspace(1)*
           ret float addrspace(1)* %ac
                     }
@@ -1070,6 +1075,7 @@ Reactant.@reactant_overlay @noinline function (func::LLVMFunc{F,tt})(
     # linearize kernel arguments
     seen = Reactant.OrderedIdDict()
     kernelargsym = gensym("kernelarg")
+
     for (i, prev) in enumerate(Any[func.f, args...])
         Reactant.make_tracer(seen, prev, (kernelargsym, i), Reactant.NoStopTracedTrack)
     end

src/Tracing.jl

@@ -12,6 +12,9 @@ struct VisitedObject
     id::Int
 end
 
+is_traced_number(x::Type) = false
+Base.@nospecializeinfer is_traced_number(@nospecialize(T::Type{<:TracedRNumber})) = true
+
 function traced_type_inner end
 
 Base.@nospecializeinfer function traced_type_inner(
@@ -1018,7 +1021,7 @@ Base.@nospecializeinfer function make_tracer_via_immutable_constructor(
             end
             FT = fieldtype(TT, i)
             if mode != TracedToTypes && !(Core.Typeof(xi2) <: FT)
-                if FT <: TracedRNumber && xi2 isa unwrapped_eltype(FT)
+                if is_traced_number(FT) && xi2 isa unwrapped_eltype(FT)
                     xi2 = FT(xi2)
                     xi2 = Core.Typeof(xi2)((newpath,), xi2.mlir_data)
                     seen[xi2] = xi2
@@ -1139,7 +1142,7 @@ Base.@nospecializeinfer function make_tracer_unknown(
             end
             FT = fieldtype(TT, i)
             if mode != TracedToTypes && !(Core.Typeof(xi2) <: FT)
-                if FT <: TracedRNumber && xi2 isa unwrapped_eltype(FT)
+                if is_traced_number(FT) && xi2 isa unwrapped_eltype(FT)
                     xi2 = FT(xi2)
                     xi2 = Core.Typeof(xi2)((newpath,), xi2.mlir_data)
                     seen[xi2] = xi2

test/integration/cuda.jl

@@ -203,3 +203,22 @@ end
     @jit searchsorted!(A, B)
     @test all(Array(A) .≈ 311)
 end
+
+function convert_mul_kernel!(Gu, w::FT) where {FT}
+    r = FT(0.5) * w
+    @inbounds Gu[1, 1, 1] = r
+    return nothing
+end
+
+function convert_mul!(Gu, w)
+    @cuda blocks = 1 threads = 1 convert_mul_kernel!(Gu, w)
+    return nothing
+end
+
+@testset "Convert mul" begin
+    w = Reactant.ConcreteRNumber(0.6)
+    Gu = Reactant.to_rarray(ones(24, 24, 24))
+    @jit convert_mul!(Gu, w)
+    Gui = Array((Gu))
+    @test Gui[1] ≈ 0.3
+end