[JLL] Add conversion from llvm to mlir (#324)

* [JLL] Add conversion from llvm to mlir

* And call registrator

* export

Author: wsmoses

deps/ReactantExtra/API.cpp

@@ -361,6 +361,22 @@ extern "C" void FreeClient(PjRtClient * client) {
     delete client;
 }
 
+#include "xla/service/custom_call_target_registry.h"
+extern "C" void RegisterCustomCallTarget(const char* name, void* address,
+                                         const char* platform) {
+    CustomCallTargetRegistry::Global()->Register(std::string(name), address, std::string(platform));
+}
+
+#include "mlir/Target/LLVMIR/Import.h"
+extern "C" MlirModule ConvertLLVMToMLIR(LLVMModuleRef lmod, MlirContext cctx) {
+    auto llvmModule = std::unique_ptr<llvm::Module>(unwrap(lmod));
+    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) {
     auto program = std::make_unique<xla::ifrt::HloProgram>(cast<ModuleOp>(*unwrap(cmod)));

deps/ReactantExtra/BUILD

@@ -423,6 +423,8 @@ cc_library(
 "-Wl,-exported_symbol,_RegisterDialects",
 "-Wl,-exported_symbol,_InitializeRegistryAndPasses",
 "-Wl,-exported_symbol,_ifrt_*",
+"-Wl,-exported_symbol,RegisterCustomCallTarget",
+"-Wl,-exported_symbol,ConvertLLVMToMLIR",
     ]}),
     deps = [
                 "@enzyme//:EnzymeMLIR",
@@ -435,6 +437,7 @@ cc_library(
         "@llvm-project//mlir:ConversionPasses",
         "@llvm-project//mlir:DLTIDialect",
         "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FromLLVMIRTranslation",
         "@llvm-project//mlir:GPUDialect",
         "@llvm-project//mlir:LinalgDialect",
         "@llvm-project//mlir:LLVMDialect",