more work

Author: William Moses

Project.toml

@@ -6,7 +6,6 @@ version = "0.2.9"
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"

ext/ReactantCUDAExt.jl

@@ -437,7 +437,8 @@ function compiler_cache(ctx::MLIR.IR.Context)
     return cache
 end
 
-function recufunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT}
+Reactant.@overlay function CUDA.cufunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT}
+    @show "recufunction", f, tt
     res = Base.@lock CUDA.cufunction_lock begin
         # compile the function
 	cache = compiler_cache(MLIR.IR.context())
@@ -450,43 +451,4 @@ function recufunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT}
     res
 end
 
-const CC = Core.Compiler
-
-import Core.Compiler:
-    AbstractInterpreter,
-    abstract_call,
-    abstract_call_known,
-    ArgInfo,
-    StmtInfo,
-    AbsIntState,
-    get_max_methods,
-    CallMeta,
-    Effects,
-    NoCallInfo,
-    widenconst,
-    mapany,
-    MethodResultPure
-
-
-function Reactant.set_reactant_abi(
-    interp,
-    f::typeof(CUDA.cufunction),
-    arginfo::ArgInfo,
-    si::StmtInfo,
-    sv::AbsIntState,
-    max_methods::Int=get_max_methods(interp, f, sv),
-)
-    (; fargs, argtypes) = arginfo
-	
-	arginfo2 = ArgInfo(
-	    if fargs isa Nothing
-		nothing
-	    else
-		[:($(recufunction)), fargs[2:end]...]
-	    end,
-	    [Core.Const(recufunction), argtypes[2:end]...],
-	)
-	return abstract_call_known(interp, recufunction, arginfo2, si, sv, max_methods)
-end
-
 end # module ReactantCUDAExt

src/Interpreter.jl

@@ -21,6 +21,15 @@ import Core.Compiler:
     mapany,
     MethodResultPure
 
+
+Base.Experimental.@MethodTable REACTANT_METHOD_TABLE
+
+macro overlay(method_expr)
+    def = splitdef(method_expr)
+    def[:name] = Expr(:overlay, :(Reactant.REACTANT_METHOD_TABLE), def[:name])
+    return esc(combinedef(def))
+end
+
 function set_reactant_abi(
     interp,
     @nospecialize(f),
@@ -54,50 +63,6 @@ function set_reactant_abi(
         end
     end
 
-    if length(argtypes) >= 5 &&
-        f === Core.kwcall &&
-        (
-            widenconst(argtypes[3]) == typeof(Enzyme.gradient) ||
-            widenconst(argtypes[3]) == typeof(Enzyme.jacobian)
-        ) &&
-        widenconst(argtypes[4]) <: Enzyme.Mode
-        newmode = Enzyme.set_abi(widenconst(argtypes[4]), ReactantABI)
-        if newmode != widenconst(argtypes[4])
-            newmodev = newmode()
-            arginfo2 = ArgInfo(
-                if fargs isa Nothing
-                    nothing
-                else
-                    [fargs[1:3]..., :($(newmodev)), fargs[5:end]...]
-                end,
-                [argtypes[1:3]..., Core.Const(newmodev), argtypes[5:end]...],
-            )
-            return abstract_call_known(interp, f, arginfo2, si, sv, max_methods)
-        end
-    end
-
-    if length(argtypes) >= 5 &&
-        methods(f)[1].module == Enzyme &&
-        widenconst(argtypes[5]) <: Enzyme.Mode &&
-        (
-            widenconst(argtypes[4]) == typeof(Enzyme.gradient) ||
-            widenconst(argtypes[4]) == typeof(Enzyme.jacobian)
-        )
-        newmode = Enzyme.set_abi(widenconst(argtypes[5]), ReactantABI)
-        if newmode != widenconst(argtypes[5])
-            newmodev = newmode()
-            arginfo2 = ArgInfo(
-                if fargs isa Nothing
-                    nothing
-                else
-                    [fargs[1:4]..., :($(newmodev)), fargs[6:end]...]
-                end,
-                [argtypes[1:4]..., Core.Const(newmodev), argtypes[6:end]...],
-            )
-            return abstract_call_known(interp, f, arginfo2, si, sv, max_methods)
-        end
-    end
-
     return Base.@invoke abstract_call_known(
         interp::AbstractInterpreter,
         f::Any,
@@ -116,7 +81,7 @@ function set_reactant_abi end
     function ReactantInterpreter(; world::UInt=Base.get_world_counter())
         return Enzyme.Compiler.Interpreter.EnzymeInterpreter(
             ReactantCacheToken(),
-            nothing,            #=mt=#
+            REACTANT_METHOD_TABLE,
             world,
             true,            #=forward_rules=#
             true,            #=reverse_rules=#
@@ -132,7 +97,7 @@ else
     )
         return Enzyme.Compiler.Interpreter.EnzymeInterpreter(
             REACTANT_CACHE,
-            nothing,            #=mt=#
+            REACTANT_METHOD_TABLE,
             world,
             true,            #=forward_rules=#
             true,            #=forward_rules=#

src/Reactant.jl

@@ -128,4 +128,7 @@ function set_default_backend(backend::String)
     return set_default_backend(XLA.backends[backend])
 end
 
+# include("../ext/ReactantCUDAExt.jl")
+
 end # module
+

src/utils.jl

@@ -37,6 +37,88 @@ function apply(f, args...; kwargs...)
     return f(args...; kwargs...)
 end
 
+function call_with_reactant end
+
+function rewrite_inst(inst)
+  @show inst
+  if Meta.isexpr(inst, :call)
+    rep = Expr(:call, call_with_reactant, inst.args...)
+    @show rep
+    return rep
+  end
+  return inst
+end
+
+function call_with_reactant_generator(world::UInt, source::LineNumberNode, @nospecialize(F::Type), @nospecialize(N::Int), self, @nospecialize(f::Type), @nospecialize(args))
+    @nospecialize
+    @show f, args
+
+    stub = Core.GeneratedFunctionStub(identity, Core.svec(:methodinstance, :f, :args), Core.svec())
+
+    # look up the method match
+    method_error = :(throw(MethodError(f, args, $world)))
+
+    interp = ReactantInterpreter(; world)
+    
+    mt = interp.method_table
+
+    sig = Tuple{F, args...}
+    min_world = Ref{UInt}(typemin(UInt))
+    max_world = Ref{UInt}(typemax(UInt))
+    match = ccall(:jl_gf_invoke_lookup_worlds, Any,
+                  (Any, Any, Csize_t, Ref{Csize_t}, Ref{Csize_t}),
+                  sig, mt, world, min_world, max_world)
+    match === nothing && return stub(world, source, method_error)
+
+    # look up the method and code instance
+    mi = ccall(:jl_specializations_get_linfo, Ref{Core.MethodInstance},
+               (Any, Any, Any), match.method, match.spec_types, match.sparams)
+  
+    result = Core.Compiler.InferenceResult(mi, Core.Compiler.typeinf_lattice(interp))
+    frame = Core.Compiler.InferenceState(result, #=cache_mode=#:global, interp)
+    @assert frame !== nothing
+    Core.Compiler.typeinf(interp, frame)
+    @assert Core.Compiler.is_inferred(frame)
+
+    #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)
+	caller = frame.result
+	@static if VERSION < v"1.11-"
+	  ir = Core.Compiler.run_passes(opt.src, opt, caller)
+	else
+	  ir = Core.Compiler.run_passes_ipo_safe(opt.src, opt, caller)
+	  Core.Compiler.ipo_dataflow_analysis!(interp, opt, ir, caller)
+	end
+	@show ir
+	  for (i, inst) in enumerate(ir.stmts)
+	     @static if VERSION < v"1.11"
+	        Core.Compiler.setindex!(ir.stmts[i], rewrite_inst(inst[:inst]), :inst)
+	     else
+	        Core.Compiler.setindex!(ir.stmts[i], rewrite_inst(inst[:stmt]), :stmt)
+	     end
+	  end
+	@show ir
+    	Core.Compiler.finish(interp, opt, ir, caller)
+        src = Core.Compiler.ir_to_codeinf!(opt)
+    #end
+
+    new_ci = copy(src)
+    new_ci.slotnames = Symbol[Symbol("#self#"), :f, :args]
+    new_ci.edges = Core.MethodInstance[mi]
+    new_ci.min_world = min_world[]
+    new_ci.max_world = max_world[]
+
+    return new_ci
+end
+
+@eval function call_with_reactant(f::F, args::Vararg{Any, N}) where {F, N}
+    $(Expr(:meta, :generated_only))
+    $(Expr(:meta, :generated, call_with_reactant_generator))
+end
+
 function make_mlir_fn(
     f,
     args,
@@ -131,36 +213,13 @@ function make_mlir_fn(
         interp = ReactantInterpreter()
 
         # TODO replace with `Base.invoke_within` if julia#52964 lands        
-        # TODO fix it for kwargs
-        ircoderes = Base.code_ircode(f, map(typeof, traced_args); interp)
-
-        if length(ircoderes) != 1
-            throw(
-                AssertionError(
-                    "Could not find unique ircode for $f $traced_args, found $ircoderes"
-                ),
-            )
-        end
-        ir, ty = ircoderes[1]
-        oc = Core.OpaqueClosure(ir)
+        # TODO fix it for kwargs	
+	oc = call_with_reactant # Core.OpaqueClosure(ir)
 
         if f === Reactant.apply
-            oc(traced_args[1], (traced_args[2:end]...,))
+            oc(f, traced_args[1], (traced_args[2:end]...,))
         else
-            if (length(traced_args) + 1 != length(ir.argtypes)) || (
-                length(traced_args) > 0 &&
-                length(ir.argtypes) > 0 &&
-                !(last(ir.argtypes) isa Core.Const) &&
-                last(ir.argtypes) != typeof(traced_args[end])
-            )
-                @assert ir.argtypes[end] <: Tuple
-                oc(
-                    traced_args[1:(length(ir.argtypes) - 2)]...,
-                    (traced_args[(length(ir.argtypes) - 1):end]...,),
-                )
-            else
-                oc(traced_args...)
-            end
+            oc(f, traced_args...)
         end
     end
 

test/cuda.jl

@@ -11,7 +11,9 @@ end
 
 # basic squaring on GPU
 function square!(x)
-    @cuda blocks = 1 threads = length(x) square_kernel!(x)
+    # @cuda blocks = 1 threads = length(x) square_kernel!(x)
+    cr = @cuda launch=false square_kernel!(x)
+    @show cr
     return nothing
 end