Fix ARM CUDA build (#85)

Makefile

@@ -33,11 +33,13 @@ $(BUILD_ARCHIVE): $(OPENXLA_DIR) extension/BUILD
 # Clones OPENXLA
 $(OPENXLA_DIR):
 	mkdir -p $(OPENXLA_DIR) && \
+	  cp -r extension/patches $(OPENXLA_DIR) && \
 		cd $(OPENXLA_DIR) && \
 		git init && \
 		git remote add origin $(OPENXLA_GIT_REPO) && \
 		git fetch --depth 1 origin $(OPENXLA_GIT_REV) && \
 		git checkout FETCH_HEAD && \
+		bash patches/apply.sh && \
 		rm $(OPENXLA_DIR)/.bazelversion
 
 # Print OPENXLA Dir

extension/patches/absl_neon.patch

@@ -0,0 +1,13 @@
+diff --git a/absl/base/config.h b/absl/base/config.h
+index 5fa9f0efe5a4..741e320fe40c 100644
+--- a/absl/base/config.h
++++ b/absl/base/config.h
+@@ -962,7 +962,7 @@ static_assert(ABSL_INTERNAL_INLINE_NAMESPACE_STR[0] != 'h' ||
+ // https://llvm.org/docs/CompileCudaWithLLVM.html#detecting-clang-vs-nvcc-from-code
+ #ifdef ABSL_INTERNAL_HAVE_ARM_NEON
+ #error ABSL_INTERNAL_HAVE_ARM_NEON cannot be directly set
+-#elif defined(__ARM_NEON) && !defined(__CUDA_ARCH__)
++#elif defined(__ARM_NEON) && !defined(__CUDACC__)
+ #define ABSL_INTERNAL_HAVE_ARM_NEON 1
+ #endif
+

extension/patches/absl_workspace.patch

@@ -0,0 +1,12 @@
+diff --git a/third_party/tsl/third_party/absl/workspace.bzl b/third_party/tsl/third_party/absl/workspace.bzl
+index 06f75166ce4b..56d146d65abe 100644
+--- a/third_party/tsl/third_party/absl/workspace.bzl
++++ b/third_party/tsl/third_party/absl/workspace.bzl
+@@ -42,6 +42,7 @@ def repo():
+         build_file = "//third_party/absl:com_google_absl.BUILD",
+         system_build_file = "//third_party/absl:system.BUILD",
+         system_link_files = SYS_LINKS,
++       patch_file = ["//third_party/absl:absl_neon.patch"],
+         strip_prefix = "abseil-cpp-{commit}".format(commit = ABSL_COMMIT),
+         urls = tf_mirror_urls("https://github.com/abseil/abseil-cpp/archive/{commit}.tar.gz".format(commit = ABSL_COMMIT)),
+     )

extension/patches/apply.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+set -ex
+
+dir="$(cd "$(dirname "$0")"; pwd)"
+arch="$(uname -m)"
+
+if [[ $arch == 'aarch64' ]]; then
+  # See https://github.com/tensorflow/tensorflow/issues/62490#issuecomment-2077646521
+  git apply $dir/absl_workspace.patch
+  cp $dir/absl_neon.patch third_party/tsl/third_party/absl
+fi