Fix export of cpu/gpu handler (#1311)

* Fix export of cpu/gpu handler

* fix header import

Author: wsmoses

src/enzyme_ad/jax/BUILD

@@ -26,6 +26,7 @@ cc_library(
     srcs = ["gpu.cc"],
     copts = ["-Wno-vla-cxx-extension"],
     deps = [
+        "@llvm-project//mlir:CAPIIR",
         "@xla//xla/ffi:ffi_api",
         "@xla//xla/ffi/api:ffi",
     ],
@@ -35,6 +36,7 @@ cc_library(
     name = "cpu",
     srcs = ["cpu.cc"],
     deps = [
+        "@llvm-project//mlir:CAPIIR",
         "@xla//xla/service:custom_call_status",
         "@xla//xla/service:custom_call_target_registry",
     ],

src/enzyme_ad/jax/cpu.cc

@@ -2,6 +2,8 @@
 #include "xla/service/custom_call_target_registry.h"
 #include <cstring>
 
+#include "mlir-c/Support.h"
+
 template <bool withError> struct CallInfo;
 
 template <> struct CallInfo<false> {
@@ -28,7 +30,7 @@ void forwarding_custom_call(void *out, const void **in, const void *opaque_ptr,
   }
 }
 
-extern "C" void RegisterEnzymeXLACPUHandler() {
+extern "C" MLIR_CAPI_EXPORTED void RegisterEnzymeXLACPUHandler() {
   xla::CustomCallTargetRegistry::Global()->Register(
       "enzymexla_compile_cpu", (void *)&forwarding_custom_call<false>, "Host");
   xla::CustomCallTargetRegistry::Global()->Register(

src/enzyme_ad/jax/gpu.cc

@@ -1,6 +1,8 @@
 #include "xla/ffi/api/ffi.h"
 #include "xla/ffi/ffi_api.h"
 
+#include "mlir-c/Support.h"
+
 template <bool withError> struct CallInfo;
 
 template <> struct CallInfo<false> {
@@ -138,7 +140,7 @@ XLA_FFI_Error *execute(XLA_FFI_CallFrame *call_frame) {
   return nullptr;
 }
 
-extern "C" void RegisterEnzymeXLAGPUHandler() {
+extern "C" MLIR_CAPI_EXPORTED void RegisterEnzymeXLAGPUHandler() {
   XLA_FFI_Handler_Bundle bundle = {instantiate, prepare, initialize<false>,
                                    execute<false>};