feat: copy tt.func into main module [skip ci]

Author: avik-pal

ext/ReactantPythonCallExt/pycall.jl

@@ -108,8 +108,18 @@ function overlayed_pycall_with_triton(
 
     ccinfo = triton.compile(src; target=target, options=options.__dict__)
 
-    println(pyconvert(String, ccinfo.asm["source"]))
-    println(pyconvert(String, ccinfo.asm["ttir"]))
+    @show ccinfo.metadata
+    @show ccinfo.asm.keys()
+    # shared = ccinfo.metadata["shared"]
+    kernel_name = pyconvert(String, ccinfo.metadata.name)
+    # cluster_dims = ccinfo.metadata["cluster_dims"]
+
+    # println(pyconvert(String, ccinfo.asm["source"]))
+    # println(pyconvert(String, ccinfo.asm["ttir"]))
+
+    res = @opcall triton_call(
+        pyconvert(String, ccinfo.asm["ttir"]), args...; func_name=kernel_name
+    )
 
     return error("TODO: implement triton")
 end

src/Ops.jl

@@ -1698,68 +1698,38 @@ end
 end
 
 # Generate a unique name given a module hash and a function name.
-function _hlo_call_name(orig_name, module_suffix)
-    return orig_name * "_hlo_call_" * module_suffix
-end
+_new_function_name(orig_name, module_suffix) = orig_name * "_call_" * module_suffix
 
-"""
-    hlo_call(mlir_code::String, args::Vararg{AnyTracedRArray}...; func_name::String="main") -> NTuple{N, AnyTracedRArray}
-
-Given a MLIR module given as a string, calls the function identified by the `func_name` keyword parameter (default "main")
-with the provided arguments and return a tuple for each result of the call.
-
-```julia-repl
-julia> Reactant.@jit(
-          hlo_call(
-              \"\"\"
-              module {
-                func.func @main(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor<3xf32> {
-                  %0 = stablehlo.add %arg0, %arg1 : tensor<3xf32>
-                  return %0 : tensor<3xf32>
-                }
-              }
-              \"\"\",
-              Reactant.to_rarray(Float32[1, 2, 3]),
-              Reactant.to_rarray(Float32[1, 2, 3]),
-          )
-       )
-(ConcretePJRTArray{Float32, 1}(Float32[2.0, 4.0, 6.0]),)
-```
-"""
-@noinline function hlo_call(
-    code,
-    args...;
-    func_name="main",
-    location=mlir_stacktrace("hlo_call", @__FILE__, @__LINE__),
+function _extract_function(
+    code::String; func_name::String="main", func_op_kind::String="func.func"
 )
     module_suffix = string(hash(code); base=16)
-    name_to_call = _hlo_call_name(func_name, module_suffix)
+    name_to_call = _new_function_name(func_name, module_suffix)
 
     current_module = MLIR.IR.mmodule()
     top_level_block = MLIR.IR.body(current_module)
 
     symbol_attr_name = String(MLIR.API.mlirSymbolTableGetSymbolAttributeName())
-
     fn = MLIR.IR.lookup(
         MLIR.IR.SymbolTable(MLIR.IR.Operation(current_module)), name_to_call
     )
+
     if isnothing(fn)
         new_mod = parse(MLIR.IR.Module, code)
         new_mod_op = MLIR.IR.Operation(new_mod)
         body = MLIR.IR.body(new_mod)
 
         operations = collect(MLIR.IR.OperationIterator(body))
         for op in operations
-            if MLIR.IR.name(op) == "func.func"
+            if MLIR.IR.name(op) == func_op_kind
                 fn_name = String(MLIR.IR.attr(op, symbol_attr_name))
                 if fn_name == func_name
                     fn = op
                 end
 
-                new_name = _hlo_call_name(fn_name, module_suffix)
                 res = MLIR.IR.LogicalResult(
                     MLIR.API.mlirSymbolTableReplaceAllSymbolUses(
-                        fn_name, new_name, new_mod_op
+                        fn_name, name_to_call, new_mod_op
                     ),
                 )
                 @assert res == MLIR.IR.success() "hlo_call: failed to rename $fn_name"
@@ -1772,7 +1742,7 @@ julia> Reactant.@jit(
                 )
 
                 # Change function name
-                MLIR.IR.attr!(op, symbol_attr_name, MLIR.IR.Attribute(new_name))
+                MLIR.IR.attr!(op, symbol_attr_name, MLIR.IR.Attribute(name_to_call))
             end
         end
 
@@ -1786,11 +1756,59 @@ julia> Reactant.@jit(
         error("hlo_call: could not find function $func_name in the provided module")
     end
 
+    return name_to_call
+end
+
+function triton_call(
+    mlir_code::String,
+    args::Union{TracedRArray,TracedRNumber,Number}...;
+    func_name::String="main",
+    location=mlir_stacktrace("triton_call", @__FILE__, @__LINE__),
+)
+    name_to_call = _extract_function(mlir_code; func_name, func_op_kind="tt.func")
+
+    @show name_to_call
+    display(MLIR.IR.mmodule())
+
+    error("TODO: implement triton_call")
+end
+
+"""
+    hlo_call(mlir_code::String, args::Vararg{AnyTracedRArray}...; func_name::String="main") -> NTuple{N, AnyTracedRArray}
+
+Given a MLIR module given as a string, calls the function identified by the `func_name` keyword parameter (default "main")
+with the provided arguments and return a tuple for each result of the call.
+
+```julia-repl
+julia> Reactant.@jit(
+          hlo_call(
+              \"\"\"
+              module {
+                func.func @main(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor<3xf32> {
+                  %0 = stablehlo.add %arg0, %arg1 : tensor<3xf32>
+                  return %0 : tensor<3xf32>
+                }
+              }
+              \"\"\",
+              Reactant.to_rarray(Float32[1, 2, 3]),
+              Reactant.to_rarray(Float32[1, 2, 3]),
+          )
+       )
+(ConcretePJRTArray{Float32, 1}(Float32[2.0, 4.0, 6.0]),)
+```
+"""
+@noinline function hlo_call(
+    code,
+    args::Union{TracedRArray,TracedRNumber}...;
+    func_name="main",
+    location=mlir_stacktrace("hlo_call", @__FILE__, @__LINE__),
+)
+    name_to_call = _extract_function(code; func_name, func_op_kind="func.func")
+
     ftype_attr = MLIR.IR.attr(fn, "function_type")
     ftype = MLIR.IR.Type(ftype_attr)
 
-    @assert all(Base.Fix2(isa, Union{TracedRArray,TracedRNumber}), args) "hlo_call: all inputs to hlo_call should be reactant arrays or numbers"
-    @assert MLIR.IR.ninputs(ftype) == length(args) "hlo_call: invalid number of arguments for function $func_name"
+    @assert MLIR.IR.ninputs(ftype) == length(args) "hlo_call: invalid number of arguments for function $func_name. Expected $(MLIR.IR.ninputs(ftype)), got $(length(args))"
 
     for (i, arg) in enumerate(args)
         expected_type = MLIR.IR.input(ftype, i)