CUDA: fix gc issues (#685)

* CUDA: fix gc issues

* Update ext/ReactantCUDAExt.jl

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

@@ -18,7 +18,9 @@ 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)
+        gc_vec = Reactant.Compiler.context_gc_vector[MLIR.IR.context()]
+        push!(gc_vec, xs)
+        @assert gc_vec[end] === xs
         ptr = Base.reinterpret(Core.LLVMPtr{T,CUDA.AS.Global}, Base.pointer_from_objref(xs))
         return new(ptr)
     end

src/TracedRArray.jl

@@ -27,6 +27,12 @@ Base.strides(x::TracedRArray) = Base.size_to_strides(1, size(x)...)
 
 Base.IndexStyle(::Type{<:TracedRArray}) = Base.IndexLinear()
 
+# This is required otherwise we will copy a tracedrarray each time
+# we use it
+function Base.convert(::Type{TracedRArray}, x::TracedRArray)
+    return x
+end
+
 function Base.convert(::Type{TracedRArray}, x::AnyTracedRArray)
     return Base.convert(TracedRArray{unwrapped_eltype(x),ndims(x)}, x)
 end