Merge pull request #1 from Flamefire/tf215

Add patch fixing compilation on AVX512 systems
VRehnberg · Jun 13, 2024 · dd74cf8 · dd74cf8
2 parents 0bcde97 + 8e6aa82
commit dd74cf8
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 0 deletions.
diff --git a/easybuild/easyconfigs/t/TensorFlow/TensorFlow-2.15.1-foss-2023a-CUDA-12.1.1.eb b/easybuild/easyconfigs/t/TensorFlow/TensorFlow-2.15.1-foss-2023a-CUDA-12.1.1.eb
@@ -150,6 +150,7 @@ exts_list = [
             'TensorFlow-2.15.1_remove-libclang-dep.patch',
             'TensorFlow-2.15.1_remove-io-gcs-filesystem-dep.patch',
             'TensorFlow-2.15.1_add-default-shell-env.patch',
+            'TensorFlow-2.15.1_fix-AVX512-eigen-compilation.patch',
             'TensorFlow-2.15.1_fix-flatbuffer-license.patch',
             'TensorFlow-2.15.1_fix-pybind11-build.patch',
             'TensorFlow-2.15.1_fix-cuda_build_defs.patch',
@@ -208,6 +209,8 @@ exts_list = [
              'eba7351a4b0696c589b9c507bacb0257ebce8c39fde39ab72d5d6a69deaaec02'},
             {'TensorFlow-2.15.1_add-default-shell-env.patch':
              '3d5196b4bf2e91048dc8a18f9e8f487a223fcd973d6302e80b0d4000ea3d652b'},
+            {'TensorFlow-2.15.1_fix-AVX512-eigen-compilation.patch':
+             '761059e5f5f5eeeef8aed5517a7685a0eb0a9193d4afe8d45237527681c9c0a3'},
             {'TensorFlow-2.15.1_fix-flatbuffer-license.patch':
              '2c04d5095977a628a238dbf93c5fada7159c86752a7183e64e0cf7c7ab00caf4'},
             {'TensorFlow-2.15.1_fix-pybind11-build.patch':

diff --git a/easybuild/easyconfigs/t/TensorFlow/TensorFlow-2.15.1-foss-2023a.eb b/easybuild/easyconfigs/t/TensorFlow/TensorFlow-2.15.1-foss-2023a.eb
@@ -144,6 +144,7 @@ exts_list = [
             'TensorFlow-2.15.1_remove-libclang-dep.patch',
             'TensorFlow-2.15.1_remove-io-gcs-filesystem-dep.patch',
             'TensorFlow-2.15.1_add-default-shell-env.patch',
+            'TensorFlow-2.15.1_fix-AVX512-eigen-compilation.patch',
             'TensorFlow-2.15.1_fix-flatbuffer-license.patch',
             'TensorFlow-2.15.1_fix-pybind11-build.patch',
             'TensorFlow-2.15.1_disable-avx512-extensions.patch',
@@ -195,6 +196,8 @@ exts_list = [
              'eba7351a4b0696c589b9c507bacb0257ebce8c39fde39ab72d5d6a69deaaec02'},
             {'TensorFlow-2.15.1_add-default-shell-env.patch':
              '3d5196b4bf2e91048dc8a18f9e8f487a223fcd973d6302e80b0d4000ea3d652b'},
+            {'TensorFlow-2.15.1_fix-AVX512-eigen-compilation.patch':
+             '761059e5f5f5eeeef8aed5517a7685a0eb0a9193d4afe8d45237527681c9c0a3'},
             {'TensorFlow-2.15.1_fix-flatbuffer-license.patch':
              '2c04d5095977a628a238dbf93c5fada7159c86752a7183e64e0cf7c7ab00caf4'},
             {'TensorFlow-2.15.1_fix-pybind11-build.patch':

diff --git a/easybuild/easyconfigs/t/TensorFlow/TensorFlow-2.15.1_fix-AVX512-eigen-compilation.patch b/easybuild/easyconfigs/t/TensorFlow/TensorFlow-2.15.1_fix-AVX512-eigen-compilation.patch
@@ -0,0 +1,47 @@
+Fix a compilation error for CPUs with AVX512 features where an issue in Eigen leads to
+> invalid 'static_cast' from type 'const Eigen::internal::eigen_packet_wrapper<__vector(4) long long int, 1>' to type '__vector(16) float'
+
+See https://gitlab.com/libeigen/eigen/-/issues/2829
+
+Author: Alexander Grund (TU Dresden)
+
+diff --git a/third_party/eigen3/fix-avx512.patch b/third_party/eigen3/fix-avx512.patch
+new file mode 100644
+index 00000000000..0650d52a0c9
+--- /dev/null
++++ b/third_party/eigen3/fix-avx512.patch
+@@ -0,0 +1,22 @@
++diff --git a/Eigen/src/Core/arch/AVX512/TypeCasting.h b/Eigen/src/Core/arch/AVX512/TypeCasting.h
++index 02c56282f..e253e6b49 100644
++--- a/Eigen/src/Core/arch/AVX512/TypeCasting.h
+++++ b/Eigen/src/Core/arch/AVX512/TypeCasting.h
++@@ -145,8 +145,6 @@ template<> EIGEN_STRONG_INLINE Packet8bf preinterpret<Packet8bf, Packet16bf>(con
++   return _mm256_castsi256_si128(a);
++ }
++ 
++-#ifndef EIGEN_VECTORIZE_AVX512FP16
++-
++ template<> EIGEN_STRONG_INLINE Packet16f pcast<Packet16h, Packet16f>(const Packet16h& a) {
++   return half2float(a);
++ }
++@@ -155,8 +153,6 @@ template<> EIGEN_STRONG_INLINE Packet16h pcast<Packet16f, Packet16h>(const Packe
++   return float2half(a);
++ }
++ 
++-#endif
++-
++ template<> EIGEN_STRONG_INLINE Packet16f pcast<Packet16bf, Packet16f>(const Packet16bf& a) {
++   return Bf16ToF32(a);
++ }
+diff --git a/third_party/eigen3/workspace.bzl b/third_party/eigen3/workspace.bzl
+index d1d8d4ac486..da549e37432 100644
+--- a/third_party/eigen3/workspace.bzl
++++ b/third_party/eigen3/workspace.bzl
+@@ -14,6 +14,7 @@ def repo():
+     tf_http_archive(
+         name = "eigen_archive",
+         build_file = "//third_party/eigen3:eigen_archive.BUILD",
++        patch_file = ["//third_party/eigen3:fix-avx512.patch"],
+         sha256 = EIGEN_SHA256,
+         strip_prefix = "eigen-{commit}".format(commit = EIGEN_COMMIT),
+         urls = tf_mirror_urls("https://gitlab.com/libeigen/eigen/-/archive/{commit}/eigen-{commit}.tar.gz".format(commit = EIGEN_COMMIT)),