Interpreter fixes (prerequisite to CUDA support)

Project.toml

@@ -41,7 +41,7 @@ Adapt = "4"
 ArrayInterface = "7.10"
 CEnum = "0.4, 0.5"
 Downloads = "1.6"
-Enzyme = "0.13.21"
+Enzyme = "0.13.22"
 EnzymeCore = "0.8.8"
 GPUArraysCore = "0.1.6, 0.2"
 LinearAlgebra = "1.10"

deps/ReactantExtra/API.cpp

@@ -376,6 +376,16 @@ extern "C" MlirModule ConvertLLVMToMLIR(LLVMModuleRef lmod, MlirContext cctx) {
     return wrap(res);
 }
 
+#include "llvm/IRReader/IRReader.h"
+extern "C" MlirModule ConvertLLVMStrToMLIR(const char* lmod, MlirContext cctx) {
+    LLVMContext Context;
+    SMDiagnostic Err;
+    auto llvmModule = llvm::parseIR(llvm::MemoryBufferRef(lmod, "conversion"), Err, Context);
+    mlir::MLIRContext &context = *unwrap(cctx);
+    auto res = mlir::translateLLVMIRToModule(std::move(llvmModule), &context, /*emitExpensiveWarnings*/false, /*dropDICompositeElements*/false).release();
+    return wrap(res);
+}
+
 
 /* Note that this */
 extern "C" xla::PjRtLoadedExecutable* ClientCompile(PjRtClient * client, MlirModule cmod) {

deps/ReactantExtra/BUILD

@@ -450,6 +450,8 @@ cc_library(
         "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:TransformDialect",
         "@llvm-project//mlir:Transforms",
+
+        "@llvm-project//llvm:IRReader",
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:AArch64AsmParser",
         "@llvm-project//llvm:AArch64CodeGen",

src/Interpreter.jl

@@ -21,6 +21,14 @@ import Core.Compiler:
     mapany,
     MethodResultPure
 
+Base.Experimental.@MethodTable(REACTANT_METHOD_TABLE)
+
+function var"@reactant_override"(__source__::LineNumberNode, __module__::Module, def)
+    return Base.Experimental.var"@overlay"(
+        __source__, __module__, :(Reactant.REACTANT_METHOD_TABLE), def
+    )
+end
+
 function set_reactant_abi(
     interp,
     @nospecialize(f),
@@ -54,50 +62,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 +80,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 +96,7 @@ else
     )
         return Enzyme.Compiler.Interpreter.EnzymeInterpreter(
             REACTANT_CACHE,
-            nothing,            #=mt=#
+            REACTANT_METHOD_TABLE,
             world,
             true,            #=forward_rules=#
             true,            #=forward_rules=#