Kernel: support constant input arg (#522)

* Kernel: support constant input arg

* Update utils.jl

* Apply suggestions from code review

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

---------

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

Author: wsmoses

ext/ReactantCUDAExt.jl

@@ -9,6 +9,12 @@ using Adapt
 
 struct CuTracedArray{T,N,A,Size} <: DenseArray{T,N}
     ptr::Core.LLVMPtr{T,A}
+
+    function CuTracedArray{T,N,A,Size}(xs::TracedRArray) where {T,N,A,Size}
+        push!(Reactant.Compiler.context_gc_vector[MLIR.IR.context()], xs)
+        ptr = Base.reinterpret(Core.LLVMPtr{T,CUDA.AS.Global}, Base.pointer_from_objref(xs))
+        return new(ptr)
+    end
 end
 
 function Base.show(io::IO, a::AT) where {AT<:CuTracedArray}
@@ -211,10 +217,34 @@ function Base.reshape(a::CuTracedArray{T,M,A}, dims::NTuple{N,Int}) where {T,N,M
     return _derived_array(a, T, dims)
 end
 
-function Adapt.adapt_storage(::CUDA.KernelAdaptor, xs::TracedRArray{T,N}) where {T,N}
-    res = CuTracedArray{T,N,CUDA.AS.Global,size(xs)}(
-        Base.reinterpret(Core.LLVMPtr{T,CUDA.AS.Global}, Base.pointer_from_objref(xs))
+struct ReactantKernelAdaptor end
+
+function Adapt.adapt_storage(to::ReactantKernelAdaptor, p::CUDA.CuPtr)
+    return error("Cannot convert CuPtr argument of Reactant Kernel")
+end
+function Adapt.adapt_storage(ka::ReactantKernelAdaptor, xs::DenseCuArray)
+    return Adapt.adapt_storage(ka, Array(xs))
+end
+function Adapt.adapt_storage(ka::ReactantKernelAdaptor, xs::Array)
+    return Adapt.adapt_storage(ka, Reactant.Ops.constant(xs))
+end
+function Adapt.adapt_structure(to::ReactantKernelAdaptor, ref::Base.RefValue)
+    return error("Cannot convert RefValue argument of Reactant Kernel")
+end
+function Adapt.adapt_structure(
+    to::ReactantKernelAdaptor, bc::Broadcast.Broadcasted{Style,<:Any,Type{T}}
+) where {Style,T}
+    return Broadcast.Broadcasted{Style}(
+        (x...) -> T(x...), Adapt.adapt(to, bc.args), bc.axes
     )
+end
+
+Reactant.@reactant_overlay @noinline function CUDA.cudaconvert(arg)
+    return adapt(ReactantKernelAdaptor(), arg)
+end
+
+function Adapt.adapt_storage(::ReactantKernelAdaptor, xs::TracedRArray{T,N}) where {T,N}
+    res = CuTracedArray{T,N,CUDA.AS.Global,size(xs)}(xs)
     return res
 end
 
@@ -383,7 +413,8 @@ end
 function Reactant.make_tracer(
     seen, @nospecialize(prev::CuTracedArray), @nospecialize(path), mode; kwargs...
 )
-    x = Base.unsafe_pointer_to_objref(Base.reinterpret(Ptr{Cvoid}, prev.ptr))::TracedRArray
+    x = Base.unsafe_pointer_to_objref(Base.reinterpret(Ptr{Cvoid}, prev.ptr))
+    x = x::TracedRArray
     Reactant.make_tracer(seen, x, path, mode; kwargs...)
     return prev
 end
@@ -441,9 +472,10 @@ Reactant.@reactant_overlay @noinline function (func::LLVMFunc{F,tt})(
 
     # linearize kernel arguments
     seen = Reactant.OrderedIdDict()
-    prev = Any[func.f, args...]
     kernelargsym = gensym("kernelarg")
-    Reactant.make_tracer(seen, prev, (kernelargsym,), Reactant.NoStopTracedTrack)
+    for (i, prev) in enumerate(Any[func.f, args...])
+        Reactant.make_tracer(seen, prev, (kernelargsym, i), Reactant.NoStopTracedTrack)
+    end
     wrapper_tys = MLIR.IR.Type[]
     for arg in values(seen)
         if !(arg isa TracedRArray || arg isa TracedRNumber)

src/Compiler.jl

@@ -318,26 +318,34 @@ function run_pass_pipeline!(mod, pass_pipeline; enable_verifier=true)
     return mod
 end
 
+const context_gc_vector = Dict{MLIR.IR.Context,Vector{TracedRArray}}()
+
 # helper for debug purposes: String -> Text
 function run_pass_pipeline_on_source(source, pass_pipeline; enable_verifier=true)
     ctx = MLIR.IR.Context(Reactant.registry[], false)
+    context_gc_vector[ctx] = Vector{TracedRArray}(undef, 0)
     @ccall MLIR.API.mlir_c.RegisterDialects(ctx::MLIR.API.MlirContext)::Cvoid
-    MLIR.IR.context!(ctx) do
+    result = MLIR.IR.context!(ctx) do
         mod = parse(MLIR.IR.Module, source)
         run_pass_pipeline!(mod, pass_pipeline; enable_verifier)
         MLIR.IR.verifyall(MLIR.IR.Operation(mod); debug=true)
         Text(repr(mod))
     end
+    Base.delete!(context_gc_vector, ctx)
+    return result
 end
 
 function compile_mlir(f, args; kwargs...)
     ctx = MLIR.IR.Context(Reactant.registry[], false)
+    context_gc_vector[ctx] = Vector{TracedRArray}(undef, 0)
     @ccall MLIR.API.mlir_c.RegisterDialects(ctx::MLIR.API.MlirContext)::Cvoid
-    MLIR.IR.context!(ctx) do
+    results = MLIR.IR.context!(ctx) do
         mod = MLIR.IR.Module(MLIR.IR.Location())
         evalinfo = compile_mlir!(mod, f, args; kwargs...)
-        return mod, evalinfo...
+        return (mod, evalinfo...)
     end
+    Base.delete!(context_gc_vector, ctx)
+    return results
 end
 
 const cuLaunch = Ref{UInt}(0)
@@ -859,10 +867,11 @@ end
 function compile_xla(f, args; client=nothing, optimize=true, no_nan=false)
     # register MLIR dialects
     ctx = MLIR.IR.Context(Reactant.registry[], false)
+    context_gc_vector[ctx] = Vector{TracedRArray}(undef, 0)
     @ccall MLIR.API.mlir_c.RegisterDialects(ctx::MLIR.API.MlirContext)::Cvoid
 
     MLIR.IR.activate!(ctx)
-    return try
+    results = try
         # compile function to MLIR module
         mod = MLIR.IR.Module(MLIR.IR.Location())
         linear_args, linear_results, preserved_args, seen_args, concrete_result, isclosure = compile_mlir!(
@@ -888,6 +897,8 @@ function compile_xla(f, args; client=nothing, optimize=true, no_nan=false)
     finally
         MLIR.IR.deactivate!(ctx)
     end
+    Base.delete!(context_gc_vector, ctx)
+    return results
 end
 
 function compile(f, args; client=nothing, optimize=true, sync=false, no_nan=false)

src/utils.jl

@@ -195,7 +195,11 @@ struct MustThrowError end
 @generated function applyiterate_with_reactant(
     iteratefn, applyfn, args::Vararg{Any,N}
 ) where {N}
-    @assert iteratefn == typeof(Base.iterate)
+    if iteratefn != typeof(Base.iterate)
+        return quote
+            error("Unhandled apply_iterate with iteratefn=$iteratefn")
+        end
+    end
     newargs = Vector{Expr}(undef, N)
     for i in 1:N
         @inbounds newargs[i] = :(args[$i]...)

test/integration/cuda.jl

@@ -152,3 +152,28 @@ end
         end
     end
 end
+
+using Reactant, CUDA
+
+function cmul!(a, b)
+    b[1] *= a[1]
+    return nothing
+end
+
+function mixed(a, b)
+    @cuda threads = 1 cmul!(a, b)
+    return nothing
+end
+
+@static if !Sys.isapple()
+    @testset "Non-traced argument" begin
+        if CUDA.functional()
+            a = CuArray([4])
+            b = ConcreteRArray([3])
+
+            @jit mixed(a, b)
+            @test all(Array(a) == 4)
+            @test all(Array(b) == 12)
+        end
+    end
+end