From a7932cb93053e881d544060c89332c411afd7606 Mon Sep 17 00:00:00 2001 From: Olli Lupton Date: Thu, 13 Feb 2025 18:57:08 +0100 Subject: [PATCH] Do not emit arch-conditional PTX Arch-conditional PTX cannot be JIT compiled to newer architectures. These lists are only applies to a few kernels built into JAX/XLA, so emitting arch-conditional code should not be important. Also, remove deprecated support for pre-Turing architectures. --- .github/container/build-jax.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/container/build-jax.sh b/.github/container/build-jax.sh index 278221c56..8260b7dbf 100755 --- a/.github/container/build-jax.sh +++ b/.github/container/build-jax.sh @@ -11,9 +11,9 @@ print_var() { supported_compute_capabilities() { ARCH=$1 if [[ "${ARCH}" == "amd64" ]]; then - echo "5.2,6.0,6.1,7.0,7.5,8.0,8.6,8.9,9.0,9.0a,10.0,10.0a" + echo "sm_75,sm_80,sm_86,sm_90,sm_100,compute_120" elif [[ "${ARCH}" == "arm64" ]]; then - echo "5.3,6.2,7.0,7.2,7.5,8.0,8.6,8.7,8.9,9.0,9.0a,10.0,10.0a" + echo "sm_80,sm_86,sm_90,sm_100,compute_120" else echo "Invalid arch '$ARCH' (expected 'amd64' or 'arm64')" 1>&2 return 1