feat: tracing fully functional

Author: avik-pal

ext/ReactantPythonCallExt/pycall.jl

@@ -106,20 +106,19 @@ function overlayed_pycall_with_triton(
         ),
     )
 
+    # Currently we are doing a double compilation here. can we do better?
+    # we are compiling here + lowering again inside enzymejax
     ccinfo = triton.compile(src; target=target, options=options.__dict__)
 
-    @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
+    @opcall triton_call(
+        pyconvert(String, ccinfo.asm["ttir"]),
+        filter(x -> x isa Reactant.TracedType, args)...;
+        func_name=pyconvert(String, ccinfo.metadata.name),
+        grid_x=@opcall(constant(grid[1])),
+        grid_y=@opcall(constant(grid[2])),
+        grid_z=@opcall(constant(grid[3])),
+        shmem=@opcall(constant(pyconvert(Int, ccinfo.metadata.shared))),
     )
 
-    return error("TODO: implement triton")
+    return nothing
 end

src/Ops.jl

@@ -1808,14 +1808,27 @@ function triton_call(
     mlir_code::String,
     args::Union{TracedRArray,TracedRNumber,Number}...;
     func_name::String="main",
+    grid_x::TracedRNumber{<:Integer},
+    grid_y::TracedRNumber{<:Integer},
+    grid_z::TracedRNumber{<:Integer},
+    shmem::TracedRNumber{<:Integer},
     location=mlir_stacktrace("triton_call", @__FILE__, @__LINE__),
+    # TODO: other kwargs
 )
     name_to_call = _extract_function(mlir_code; func_name, func_op_kind="tt.func")
 
-    @show name_to_call
-    display(MLIR.IR.mmodule())
+    enzymexla.triton_call(
+        grid_x.mlir_data,
+        grid_y.mlir_data,
+        grid_z.mlir_data,
+        shmem.mlir_data,
+        [Reactant.TracedUtils.get_mlir_data(a) for a in args];
+        fn=MLIR.IR.FlatSymbolRefAttribute(name_to_call),
+        result_0=MLIR.IR.Type[],
+        location,
+    )
 
-    error("TODO: implement triton_call")
+    return nothing
 end
 
 """