continuing

Author: William Moses

ext/ReactantCUDAExt.jl

@@ -222,8 +222,6 @@ function compile(job)
             asm, meta = CUDA.GPUCompiler.compile(:asm, job)
 	    mod = meta.ir
 	    modstr = string(mod)
-	    @show mod
-	    @show modstr
 	    # check if we'll need the device runtime
 	    undefined_fs = filter(collect(CUDA.LLVM.functions(meta.ir))) do f
 		CUDA.LLVM.isdeclaration(f) && !CUDA.LLVM.isintrinsic(f)
@@ -375,8 +373,7 @@ function compile(job)
 
 	    modstr, image, meta.entry
     end
-    
-    LLVMFunc{job.source.specTypes[1],job.source.specTypes}(nothing, modstr, image, LLVM.name(entry))
+    LLVMFunc{job.source.specTypes.parameters[1],job.source.specTypes}(nothing, modstr, image, CUDA.LLVM.name(entry))
 end
 
 # link into an executable kernel
@@ -385,20 +382,23 @@ function link(job, compiled)
     return compiled
 end
 
-function (func::LLVMFunc{F,tt})(args...; blocks::CUDA.CuDim=1, threads::CUDA.CuDim=1,
-					 shmem::Integer=0) where{F, tt}
+function (func::LLVMFunc{F,tt})(args...; convert=Val(false), blocks::CuDim=1, threads::CuDim=1,
+                cooperative::Bool=false, shmem::Integer=0, call_kwargs...) where{F, tt}
+    @show args
+    @show call_kwargs
+
     blockdim = CUDA.CuDim3(blocks)
     threaddim = CUDA.CuDim3(threads)
 
-    @show args
-
     mlir_args = MLIR.IR.Value[]
     restys = MLIR.IR.Type[]
     aliases = MLIR.API.MlirAttribute[]
+    rarrays = TracedRArray[]
     for (i, a) in enumerate(args)
 	@show a
-	@assert a isa CuDeviceArray
-	ta = Base.pointer_to_objref(a.ptr)::TracedRArray
+	@assert a isa CuTracedArray
+	ta = Base.unsafe_pointer_to_objref(Base.reinterpret(Ptr{Cvoid}, a.ptr))::TracedRArray
+	push!(rarrays, ta)
 	arg = ta.mlir_data
 	arg = Reactant.Compiler.transpose_val(arg)
 	push!(restys, MLIR.IR.Type(arg))
@@ -415,7 +415,10 @@ function (func::LLVMFunc{F,tt})(args...; blocks::CUDA.CuDim=1, threads::CUDA.CuD
     end
 
     output_operand_aliases=MLIR.ArrayAttr.get(MLIR.IR.context(), aliases)
-    MLIR.IR.Dialects.stablehlo.custom_call(mlir_args; result_0=restys, call_target_name="reactant_gpu_call", output_operand_aliases)
+    call = MLIR.IR.Dialects.stablehlo.custom_call(mlir_args; result_0=restys, call_target_name="reactant_gpu_call", output_operand_aliases)
+    for (i, res) in enumerate(rarrays)
+       ta.mlir_data = Reactant.Compiler.transpose_val(MLIR.IR.result(call, i-1))
+    end
     #CUDA.cuLaunchKernel(f,
     #	       blockdim.x, blockdim.y, blockdim.z,
     #	       threaddim.x, threaddim.y, threaddim.z,

src/utils.jl

@@ -63,13 +63,13 @@ function rewrite_inst(inst, ir)
     ft = Core.Compiler.widenconst(maybe_argextype(inst.args[1], ir))
     if !(ft <: Core.IntrinsicFunction) && !(ft <: Core.Builtin)
       rep = Expr(:call, call_with_reactant, inst.args...)
-      return rep
+      return true, rep
     end
   end
   if Meta.isexpr(inst, :invoke)
-    return Expr(:call, inst.args[2:end]...)
+    return false, Expr(:call, inst.args[2:end]...)
   end
-  return inst
+  return false, inst
 end
 
 const REDUB_ARGUMENTS_NAME = gensym("redub_arguments")
@@ -120,10 +120,14 @@ function _arg_partially_inline!(@nospecialize(x), slot_replacements::Vector{Any}
         return x
     end
     if isa(x, Core.ReturnNode)
-        return Core.ReturnNode(
+	if !isdefined(x, :val)
+	   return Core.ReturnNode(:nothing)
+	else
+	   return Core.ReturnNode(
             _arg_partially_inline!(x.val, slot_replacements, type_signature, static_param_values,
                                slot_offset, arg_offset, statement_offset, boundscheck),
-        )
+           )
+	end
     end
     if isa(x, Core.GotoIfNot)
         return Core.GotoIfNot(
@@ -257,12 +261,19 @@ function call_with_reactant_generator(world::UInt, source::LineNumberNode, self,
     Core.Compiler.typeinf(interp, frame)
     @assert Core.Compiler.is_inferred(frame)
 
+    method = match.method
+    @show mi
+    @show method
+
     #if Core.Compiler.result_is_constabi(interp, frame.result)
     #    rt = frame.result.result::Core.Compiler.Const
     #    src = Core.Compiler.codeinfo_for_const(interp, frame.linfo, rt.val)
     #else
         opt = Core.Compiler.OptimizationState(frame, interp)
 
+	@show Core.Compiler.retrieve_code_info(mi, world)
+	@show opt.src
+
 	caller = frame.result
 	@static if VERSION < v"1.11-"
 	  ir = Core.Compiler.run_passes(opt.src, opt, caller)
@@ -271,21 +282,35 @@ function call_with_reactant_generator(world::UInt, source::LineNumberNode, self,
 	  Core.Compiler.ipo_dataflow_analysis!(interp, opt, ir, caller)
 	end
 
-	  for (i, inst) in enumerate(ir.stmts)
+	@show ir
+	any_changed = false
+	for (i, inst) in enumerate(ir.stmts)
 
+	   
             @static if VERSION < v"1.11"
-               Core.Compiler.setindex!(ir.stmts[i], rewrite_inst(inst[:inst], ir), :inst)
+	       changed, next = rewrite_inst(inst[:inst], ir)
+               Core.Compiler.setindex!(ir.stmts[i], next, :inst)
             else
-               Core.Compiler.setindex!(ir.stmts[i], rewrite_inst(inst[:stmt], ir), :stmt)
+	       changed, next = rewrite_inst(inst[:stmt], ir)
+               Core.Compiler.setindex!(ir.stmts[i], next, :stmt)
             end
-            Core.Compiler.setindex!(ir.stmts[i], Any, :type)
+	    if changed
+		any_changed = true
+		Core.Compiler.setindex!(ir.stmts[i], Any, :type)
+	    end
          end
     	Core.Compiler.finish(interp, opt, ir, caller)
+	@show "post", ir
         src = Core.Compiler.ir_to_codeinf!(opt)
+	
+	@show any_changed, src
+	if !any_changed
+	   src = Core.Compiler.retrieve_code_info(mi, world)
+	   @show "post non change", src
+	end
 
     # prepare a new code info
     code_info = copy(src)
-    method = match.method
     static_params = match.sparams
     signature = sig
     is_invoke = args[1] === typeof(Core.invoke)
@@ -352,6 +377,7 @@ function call_with_reactant_generator(world::UInt, source::LineNumberNode, self,
 	push!(fn_args, Core.SSAValue(length(overdubbed_code)))
     end
 
+    @show code_info.code
     #=== finish initialization of `overdubbed_code`/`overdubbed_codelocs` ===#
 
     # substitute static parameters, offset slot numbers by number of added slots, and
@@ -383,6 +409,8 @@ function call_with_reactant_generator(world::UInt, source::LineNumberNode, self,
     code_info.ssavaluetypes = length(overdubbed_code)
     code_info.ssaflags = [0x00 for _ in 1:length(overdubbed_code)] # XXX we need to copy flags that are set for the original code
 
+    @show code_info
+
     return code_info
 end
 

test/cuda.jl

@@ -11,15 +11,14 @@ end
 
 # basic squaring on GPU
 function square!(x)
-    # @cuda blocks = 1 threads = length(x) square_kernel!(x)
-    cr = @cuda launch=false square_kernel!(x)
-    @show cr
+    @cuda blocks = 1 threads = length(x) square_kernel!(x)
     return nothing
 end
 
 @testset "Square Kernel" begin
     oA = collect(1:1:64)
     A = Reactant.to_rarray(oA)
+    @show @code_hlo square!(A)
     func = @compile square!(A)
-    @test all(A .≈ (oA .* oA))
+    @test all(Array(A) .≈ (oA .* oA))
 end