Tracing check within KernelAbstractions

Author: albertomercurio

ext/ReactantKernelAbstractionsExt.jl

@@ -1,6 +1,7 @@
 module ReactantKernelAbstractionsExt
 
 using Reactant: Reactant
+using ReactantCore: ReactantCore
 
 using Adapt: Adapt
 using KernelAbstractions: KernelAbstractions
@@ -101,6 +102,14 @@ function tokw(ndrange, workgroupsize, obj, args...)
 end
 
 function (obj::KA.Kernel{ReactantBackend})(args...; ndrange=nothing, workgroupsize=nothing)
+    # If we're already inside a compilation/tracing context, or if any arguments are traced,
+    # we should trace through this kernel call instead of trying to compile it again.
+    if Reactant.within_compile() || any(ReactantCore.is_traced, args)
+        return Reactant.call_with_reactant(
+            Reactant.ka_with_reactant, ndrange, workgroupsize, obj, args...
+        )
+    end
+    
     if Reactant.precompiling()
         Reactant.@code_hlo optimize = false tokw(ndrange, workgroupsize, obj, args...)
     else

src/Overlay.jl

@@ -133,7 +133,7 @@ for (cT, aT, bT) in (
                 # Inference barrier is required when calling function recursively within
                 # overload. This is required since otherwise type inference will think this
                 # is a recursive edge rather than a call to the base method
-                Base.inferencebarrier(LinearAlgebra.mul!)(C, A, B, α, β)
+                Base.inferencebarrier(LinearAlgebra.mul!)(C2, A2, B2, α, β)
             end
             return C
         end