fixup

EnzymeAD · Dec 17, 2024 · d62d01e · d62d01e
1 parent 6121847
commit d62d01e
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 4 deletions.
diff --git a/deps/ReactantExtra/API.cpp b/deps/ReactantExtra/API.cpp
@@ -50,7 +50,6 @@
 #include "xla/pjrt/pjrt_executable.h"
 #include "xla/pjrt/pjrt_api.h"
 #include "xla/pjrt/pjrt_c_api_client.h"
-#include "xla/service/cpu/simple_orc_jit.h"
 
 #include "xla/python/ifrt/hlo/hlo_program.h"
 #include "llvm/MC/TargetRegistry.h"

diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl
@@ -352,10 +352,29 @@ Reactant.@reactant_override @noinline function (func::LLVMFunc{F,tt})(args...; c
 
     fname = Reactant.TracedUtils.get_attribute_by_name(func.entry, "sym_name")
     # Force public for now while we don't have real users
-    MLIR.IR.rmattr!(func.entry, "sym_visibility")
+    # MLIR.IR.rmattr!(func.entry, "sym_visibility")
+
+    op_ty_results = IR.Type[result_0...,]
+    operands = Value[inputs...,]
+    owned_regions = MLIR.IR.Region[]
+    successors = MLIR.IR.Block[]
+    attributes = MLIR.IR.NamedAttribute[
+	MLIR.IR.namedattribute("fn", fname),
+	MLIR.IR.namedattribute("output_operand_aliases", output_operand_aliases)
+    ]
+
+    location = MLIR.IR.Location()
+    call = MLIR.IR.create_operation(
+        "stablehlo.custom_call",
+        location;
+        mlir_args,
+        owned_regions,
+        successors,
+        attributes,
+        results=restys
+        result_inference=false,
+    )
 
-    call = MLIR.Dialects.stablehlo.custom_call(mlir_args; result_0=restys, call_target_name="reactant_gpu_call", output_operand_aliases, backend_config=MLIR.IR.Attribute(fname))
-    # call = MLIR.Dialects.stablehlo.custom_call(mlir_args; result_0=restys, call_target_name="reactant_gpu_call", output_operand_aliases, backend_config=MLIR.IR.Attribute(func.mod))
     for (i, res) in enumerate(rarrays)
        res.mlir_data = transpose_val(MLIR.IR.result(call, i))
     end