Fix build with newer CUDA

4 files changed, 161 insertions, 9 deletions

diff --git a/.SRCINFO b/.SRCINFO
index 4c15c62861c9..e00c1b8b120b 100644
--- a/.SRCINFO
+++ b/.SRCINFO
@@ -1,12 +1,12 @@
 pkgbase = python-tensorflow-addons-cuda-git
 	pkgdesc = Useful extra functionality for TensorFlow 2.x, built with CUDA support
-	pkgver = r1399.334cd7ca
-	pkgrel = 2
+	pkgver = r1490.14af819a
+	pkgrel = 1
 	url = https://github.com/tensorflow/addons
 	arch = any
 	license = APACHE
 	makedepends = python-setuptools
-	makedepends = bazel3
+	makedepends = bazel4
 	makedepends = rsync
 	makedepends = python-pip
 	depends = tensorflow-cuda
@@ -16,7 +16,11 @@ pkgbase = python-tensorflow-addons-cuda-git
 	depends = cudnn
 	source = git+https://github.com/tensorflow/addons.git
 	source = bazel-cuda-build-fix.patch
+	source = fix_cuda.patch
+	source = use_cub_from_cuda.patch
 	sha256sums = SKIP
 	sha256sums = 987e4933a298616a2ede9f2a938e047df20287ab7f514f104703ac7af5e8ef9d
+	sha256sums = ade7b591f96339dffacccf838006a206ccb0eb3758ae571ac77aef94520f001f
+	sha256sums = de727ac26b8353de6ff7972ae084d3c5ca692cebc05aac1d73bca221cb0cd45a
 
 pkgname = python-tensorflow-addons-cuda-git
diff --git a/PKGBUILD b/PKGBUILD
index f659db801b04..59fd9e640274 100644
--- a/PKGBUILD
+++ b/PKGBUILD
@@ -1,18 +1,22 @@
 # Maintainer: piernov <piernov@piernov.org>
 
 pkgname=python-tensorflow-addons-cuda-git
-pkgver=r1399.334cd7ca
-pkgrel=2
+pkgver=r1490.14af819a
+pkgrel=1
 pkgdesc="Useful extra functionality for TensorFlow 2.x, built with CUDA support"
 url="https://github.com/tensorflow/addons"
 arch=('any')
 license=('APACHE')
 depends=('tensorflow-cuda' 'python-tensorflow' 'python-typeguard' 'cuda' 'cudnn')
-makedepends=('python-setuptools' 'bazel3' 'rsync' 'python-pip')
+makedepends=('python-setuptools' 'bazel4' 'rsync' 'python-pip')
 source=("git+https://github.com/tensorflow/addons.git"
-	"bazel-cuda-build-fix.patch")
+	"bazel-cuda-build-fix.patch"
+	"fix_cuda.patch"
+	"use_cub_from_cuda.patch")
 sha256sums=('SKIP'
-            '987e4933a298616a2ede9f2a938e047df20287ab7f514f104703ac7af5e8ef9d')
+            '987e4933a298616a2ede9f2a938e047df20287ab7f514f104703ac7af5e8ef9d'
+            'ade7b591f96339dffacccf838006a206ccb0eb3758ae571ac77aef94520f001f'
+            'de727ac26b8353de6ff7972ae084d3c5ca692cebc05aac1d73bca221cb0cd45a')
 
 pkgver() {
   cd addons
@@ -25,6 +29,13 @@ prepare() {
   # Force OpenJDK-11
   export JAVA_HOME=/usr/lib/jvm/java-11-openjdk
 
+  # CUDA 11.8 uses GCC 11
+  export CC=gcc-11
+  export CXX=g++-11
+  export GCC_HOST_COMPILER_PATH=/usr/bin/${CC}
+  export HOST_C_COMPILER=/usr/bin/${CC}
+  export HOST_CXX_COMPILER=/usr/bin/${CXX}
+
   export TF_NEED_CUDA=1
   export TF_CUDA_CLANG=0
   export CUDA_TOOLKIT_PATH=/opt/cuda
@@ -33,9 +44,11 @@ prepare() {
 
   export TF_CUDA_VERSION=$(/opt/cuda/bin/nvcc --version | sed -n 's/^.*release \(.*\),.*/\1/p')
   export TF_CUDNN_VERSION=$(sed -n 's/^#define CUDNN_MAJOR\s*\(.*\).*/\1/p' /usr/include/cudnn.h)
-  export TF_CUDA_COMPUTE_CAPABILITIES=3.5,3.7,5.0,5.2,5.3,6.0,6.1,6.2,7.0,7.2,7.5
+  export TF_CUDA_COMPUTE_CAPABILITIES=3.5,3.7,5.0,5.2,5.3,6.0,6.1,6.2,7.0,7.2,7.5,8.0,8.6
 
   patch -p1 < "${srcdir}/bazel-cuda-build-fix.patch"
+  patch -p1 < "${srcdir}/fix_cuda.patch"
+  patch -p1 < "${srcdir}/use_cub_from_cuda.patch"
 }
 
 build() {
diff --git a/fix_cuda.patch b/fix_cuda.patch
new file mode 100644
index 000000000000..601beacc22e2
--- /dev/null
+++ b/fix_cuda.patch
@@ -0,0 +1,26 @@
+From 2f32601be926472f142bffbe820a28d05682219a Mon Sep 17 00:00:00 2001
+From: Bernhard Bermeitinger <bernhard.bermeitinger@unisg.ch>
+Date: Fri, 27 May 2022 10:52:13 +0200
+Subject: [PATCH] fix compilation on cuda
+
+Signed-off-by: Bernhard Bermeitinger <bernhard.bermeitinger@unisg.ch>
+---
+ .../crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc.tpl | 2 +-
+ configure.py                                                    | 2 +-
+ 2 files changed, 2 insertions(+), 2 deletions(-)
+
+diff --git a/configure.py b/configure.py
+index 0d65e88..24fd2d5 100644
+--- a/configure.py
++++ b/configure.py
+@@ -185,7 +185,7 @@ def configure_cuda():
+     write("build --config=cuda")
+     write("build:cuda --define=using_cuda=true --define=using_cuda_nvcc=true")
+     write(
+-        "build:cuda --crosstool_top=@ubuntu20.04-gcc9_manylinux2014-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain"
++        "build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain"
+     )
+ 
+ 
+-- 
+2.36.1
diff --git a/use_cub_from_cuda.patch b/use_cub_from_cuda.patch
new file mode 100644
index 000000000000..f2ac948c7183
--- /dev/null
+++ b/use_cub_from_cuda.patch
@@ -0,0 +1,109 @@
+diff --git a/build_deps/toolchains/gpu/cub.BUILD b/build_deps/toolchains/gpu/cub.BUILD
+index cdc9e4f3..159e5f33 100644
+--- a/build_deps/toolchains/gpu/cub.BUILD
++++ b/build_deps/toolchains/gpu/cub.BUILD
+@@ -18,7 +18,6 @@ filegroup(
+ cc_library(
+     name = "cub",
+     hdrs = if_cuda([":cub_header_files"]),
+-    include_prefix = "gpu",
+     deps = [
+         "@local_config_cuda//cuda:cuda_headers",
+     ],
+diff --git a/build_deps/toolchains/gpu/cuda/BUILD.tpl b/build_deps/toolchains/gpu/cuda/BUILD.tpl
+index 1ac5643f..2761abac 100644
+--- a/build_deps/toolchains/gpu/cuda/BUILD.tpl
++++ b/build_deps/toolchains/gpu/cuda/BUILD.tpl
+@@ -175,6 +175,11 @@ cc_library(
+     ],
+ )
+ 
++alias(
++    name = "cub_headers",
++    actual = "%{cub_actual}",
++)
++
+ cc_library(
+     name = "cupti_headers",
+     hdrs = [
+diff --git a/build_deps/toolchains/gpu/cuda/BUILD.windows.tpl b/build_deps/toolchains/gpu/cuda/BUILD.windows.tpl
+index 3ed4fd41..14554a38 100644
+--- a/build_deps/toolchains/gpu/cuda/BUILD.windows.tpl
++++ b/build_deps/toolchains/gpu/cuda/BUILD.windows.tpl
+@@ -134,6 +134,11 @@ cc_library(
+     ],
+ )
+ 
++alias(
++    name = "cub_headers",
++    actual = "%{cub_actual}",
++)
++
+ cc_library(
+     name = "cupti_headers",
+     hdrs = [
+diff --git a/build_deps/toolchains/gpu/cuda_configure.bzl b/build_deps/toolchains/gpu/cuda_configure.bzl
+index ba38c6b5..b30c67ac 100644
+--- a/build_deps/toolchains/gpu/cuda_configure.bzl
++++ b/build_deps/toolchains/gpu/cuda_configure.bzl
+@@ -668,6 +668,7 @@ def _get_cuda_config(repository_ctx):
+     return struct(
+         cuda_toolkit_path = toolkit_path,
+         cuda_version = cuda_version,
++        cuda_version_major = cuda_major,
+         cudart_version = cudart_version,
+         cublas_version = cublas_version,
+         cusolver_version = cusolver_version,
+@@ -725,6 +726,7 @@ def _create_dummy_repository(repository_ctx):
+             "%{cufft_lib}": lib_name("cufft", cpu_value),
+             "%{curand_lib}": lib_name("curand", cpu_value),
+             "%{cupti_lib}": lib_name("cupti", cpu_value),
++            "%{cub_actual}": ":cuda_headers",
+             "%{copy_rules}": "",
+             "%{cuda_headers}": "",
+         },
+@@ -950,6 +952,10 @@ def _create_local_cuda_repository(repository_ctx):
+         },
+     )
+ 
++    cub_actual = "@cub_archive//:cub"
++    if int(cuda_config.cuda_version_major) >= 11:
++        cub_actual = ":cuda_headers"
++
+     _tpl(
+         repository_ctx,
+         "cuda:BUILD",
+@@ -964,6 +970,7 @@ def _create_local_cuda_repository(repository_ctx):
+             "%{cufft_lib}": cuda_libs["cufft"].basename,
+             "%{curand_lib}": cuda_libs["curand"].basename,
+             "%{cupti_lib}": cuda_libs["cupti"].basename,
++            "%{cub_actual}": cub_actual,
+             "%{copy_rules}": "\n".join(copy_rules),
+             "%{cuda_headers}": (
+                 '":cuda-include",\n' + '        ":cudnn-include",'
+diff --git a/tensorflow_addons/custom_ops/layers/BUILD b/tensorflow_addons/custom_ops/layers/BUILD
+index 44d87710..a039f25d 100644
+--- a/tensorflow_addons/custom_ops/layers/BUILD
++++ b/tensorflow_addons/custom_ops/layers/BUILD
+@@ -12,7 +12,7 @@ custom_op_library(
+         "cc/ops/correlation_cost_op.cc",
+     ],
+     cuda_deps = [
+-        "@cub_archive//:cub",
++        "@local_config_cuda//cuda:cub_headers",
+     ],
+     cuda_srcs = [
+         "cc/kernels/correlation_cost_op.h",
+diff --git a/tensorflow_addons/custom_ops/layers/cc/kernels/correlation_cost_op_gpu.cu.cc b/tensorflow_addons/custom_ops/layers/cc/kernels/correlation_cost_op_gpu.cu.cc
+index 16b012a7..722f317d 100644
+--- a/tensorflow_addons/custom_ops/layers/cc/kernels/correlation_cost_op_gpu.cu.cc
++++ b/tensorflow_addons/custom_ops/layers/cc/kernels/correlation_cost_op_gpu.cu.cc
+@@ -17,7 +17,7 @@ limitations under the License.
+ 
+ #define EIGEN_USE_GPU
+ 
+-#include "gpu/cub/device/device_reduce.cuh"
++#include "cub/device/device_reduce.cuh"
+ #include "tensorflow/core/framework/tensor.h"
+ #include "tensorflow/core/framework/tensor_shape.h"
+ #include "tensorflow/core/util/gpu_kernel_helper.h"