Fix aarch build (#64)

* More prints

* fix

* for yggy

* more fixes

* cleanup

Author: wsmoses

deps/ReactantExtra/API.cpp

@@ -45,9 +45,16 @@
 #include "xla/pjrt/cpu/cpu_client.h"
 #include "xla/pjrt/status_casters.h"
 #include "xla/pjrt/pjrt_executable.h"
+#include "xla/pjrt/pjrt_api.h"
+#include "xla/pjrt/pjrt_c_api_client.h"
 #include "xla/python/ifrt/executable.h"
+#include "xla/service/cpu/simple_orc_jit.h"
 
 #include "xla/python/ifrt/hlo/hlo_program.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/TargetParser/Host.h"
+#include "llvm/ExecutionEngine/ExecutionEngine.h"
+#include "llvm/Support/Process.h"
 
 using namespace mlir;
 using namespace llvm;
@@ -58,7 +65,17 @@ using namespace xla;
 
 extern "C" void InitializeLogs() {
     absl::InitializeLog();
+    LLVMInitializeX86Target();
+    LLVMInitializeX86TargetInfo();
+    LLVMInitializeX86TargetMC();
+    LLVMInitializeX86AsmPrinter();
+    LLVMInitializeX86AsmParser();
+
     LLVMInitializeAArch64Target();
+    LLVMInitializeAArch64TargetInfo();
+    LLVMInitializeAArch64TargetMC();
+    LLVMInitializeAArch64AsmPrinter();
+    LLVMInitializeAArch64AsmParser();
 }
 
 extern "C"
@@ -102,6 +119,39 @@ extern "C" PjRtClient* MakeGPUClient(int node_id, int num_nodes, int* allowed_de
     }
 }
 
+const char* const kEnvTpuLibraryPath = "TPU_LIBRARY_PATH";
+
+extern "C" PjRtClient* MakeTPUClient(const char* tpu_path , const char** error) {
+    // Prefer $TPU_LIBRARY_PATH if set
+    std::string tpu_library_path;
+    if (tpu_path) {
+        tpu_library_path = std::string(tpu_path);
+    } else if (auto path = llvm::sys::Process::GetEnv(kEnvTpuLibraryPath)) {
+        tpu_library_path = *path;
+    } else {
+        *error = "Could not find TPU path";
+        return nullptr;
+    }
+
+    absl::StatusOr<const PJRT_Api *> pluginLoad = pjrt::LoadPjrtPlugin("tpu", tpu_library_path);
+    if (!pluginLoad.ok()) {
+      auto str = pluginLoad.status().message();
+      char* err = (char*)malloc(str.size()+1);
+      memcpy(err, str.data(), str.size()+1);
+      *error = err;
+      return nullptr;
+    }
+    absl::Status tpu_status = pjrt::InitializePjrtPlugin("tpu");
+    if (!tpu_status.ok()) {
+      auto str = tpu_status.message();
+      char* err = (char*)malloc(str.size()+1);
+      memcpy(err, str.data(), str.size()+1);
+      *error = err;
+      return nullptr;
+    }
+    return xla::GetCApiClient("TPU").value().release();
+}
+
 extern "C" int ClientNumDevices(PjRtClient* client) {
     return client->device_count();
 }

deps/ReactantExtra/BUILD

@@ -319,9 +319,13 @@ cc_library(
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:AArch64AsmParser",
         "@llvm-project//llvm:AArch64CodeGen",
+        "@llvm-project//llvm:X86AsmParser",
+        "@llvm-project//llvm:X86CodeGen",
         "@enzyme_ad//src/enzyme_ad/jax:TransformOps",
         "@enzyme_ad//src/enzyme_ad/jax:XLADerivatives",
         "@stablehlo//:chlo_ops",
+        "@xla//xla/pjrt:pjrt_api",
+        "@xla//xla/pjrt:pjrt_c_api_client",
         "@xla//xla/pjrt/cpu:cpu_client",
         "@xla//xla/pjrt/gpu:se_gpu_pjrt_client",