Added support for gfx950 architecture platform

zahiqbal · zahiqbal · commit ed6bf74c1ba3 · 2025-02-11T08:29:59.000-08:00
diff --git a/tensorflow/core/grappler/optimizers/auto_mixed_precision.cc b/tensorflow/core/grappler/optimizers/auto_mixed_precision.cc
@@ -116,7 +116,8 @@ bool HasFastFP16Support(const DeviceProperties& props) {
 #elif TENSORFLOW_USE_ROCM
   absl::flat_hash_set<std::string> FP16SupportedDevices = {
       {"gfx906"}, {"gfx908"}, {"gfx90a"}, {"gfx910"}, {"gfx940"}, {"gfx941"}, 
-      {"gfx942"}, {"gfx1010"}, {"gfx1012"}, {"gfx1030"},
+      {"gfx942"}, {"gfx950"},
+      {"gfx1010"}, {"gfx1012"}, {"gfx1030"},
       {"gfx1100"}, {"gfx1101"}, {"gfx1102"},
       {"gfx1200"}, {"gfx1201"}
   };
diff --git a/tensorflow/core/grappler/optimizers/generic_layout_optimizer.cc b/tensorflow/core/grappler/optimizers/generic_layout_optimizer.cc
@@ -73,6 +73,7 @@ inline GpuStats GetNumGPUs(const Cluster& cluster) {
          compute_capability_it->second == "gfx940" ||
          compute_capability_it->second == "gfx941" ||
          compute_capability_it->second == "gfx942" ||
+         compute_capability_it->second == "gfx950" ||
          compute_capability_it->second == "gfx1101" ||
          compute_capability_it->second == "gfx1102" ||
          compute_capability_it->second == "gfx1200" ||
diff --git a/tensorflow/core/util/gpu_device_functions.h b/tensorflow/core/util/gpu_device_functions.h
@@ -743,7 +743,7 @@ __device__ inline double GpuAtomicAdd(double* ptr, double value) {
 }
 #endif
 
-#if __gfx908__ || __gfx90a__ || __gfx940__ || __gfx941__ || __gfx942__ || __gfx1101__ || __gfx1102__ || __gfx1200__ || __gfx1201__
+#if __gfx908__ || __gfx90a__ || __gfx940__ || __gfx941__ || __gfx942__ || __gfx950__ || __gfx1101__ || __gfx1102__ || __gfx1200__ || __gfx1201__
 
 #define ADDRSP1 __attribute__((address_space(1)))
 __device__ float
@@ -963,7 +963,7 @@ __device__ inline int64_t GpuAtomicMin(int64_t* ptr, int64_t value) {
 }
 #endif
 
-#if __gfx908__ || __gfx90a__ || __gfx940__ || __gfx941__ || __gfx942__ || __gfx1101__ || __gfx1102__ || __gfx1200__ || __gfx1201__
+#if __gfx908__ || __gfx90a__ || __gfx940__ || __gfx941__ || __gfx942__ || __gfx950__ || __gfx1101__ || __gfx1102__ || __gfx1200__ || __gfx1201__
 // Low level instructions don't return. For now, assume that return value
 // is always unused.
 __device__ float GpuAtomicAdd(float* dst, float val) {
@@ -978,7 +978,7 @@ __device__ inline T GpuAtomicAddShared(T* ptr, T value) {
   return GpuAtomicAdd(ptr, value);
 }
 
-#if __gfx908__ || __gfx90a__ || __gfx940__ || __gfx941__ || __gfx942__ || __gfx1101__ || __gfx1102__ || __gfx1200__ || __gfx1201__
+#if __gfx908__ || __gfx90a__ || __gfx940__ || __gfx941__ || __gfx942__ || __gfx950__ || __gfx1101__ || __gfx1102__ || __gfx1200__ || __gfx1201__
 __device__ float GpuAtomicAddShared(float* dst, float val) {
   atomicAdd(dst, val);
   return val;
diff --git a/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile.rocm.manylinux2014 b/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile.rocm.manylinux2014
@@ -8,7 +8,7 @@ COPY setup.packages.rocm.cs7.sh setup.packages.rocm.cs7.sh
 COPY builder.packages.rocm.cs7.txt builder.packages.rocm.cs7.txt
 RUN /setup.packages.rocm.cs7.sh /builder.packages.rocm.cs7.txt
 
-ARG GPU_DEVICE_TARGETS="gfx908 gfx90a gfx940 gfx941 gfx942 gfx1030 gfx1100 gfx1101 gfx1102 gfx1200 gfx1201"
+ARG GPU_DEVICE_TARGETS="gfx908 gfx90a gfx940 gfx941 gfx942 gfx950 gfx1030 gfx1100 gfx1101 gfx1102 gfx1200 gfx1201"
 ENV GPU_DEVICE_TARGETS=${GPU_DEVICE_TARGETS}
 
 # Install ROCM
diff --git a/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile.rocm.manylinux_2_28 b/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile.rocm.manylinux_2_28
@@ -17,7 +17,7 @@ COPY setup.packages.rocm.el8.sh setup.packages.rocm.el8.sh
 COPY builder.packages.rocm.el8.txt builder.packages.rocm.el8.txt
 RUN /setup.packages.rocm.el8.sh /builder.packages.rocm.el8.txt
 
-ARG GPU_DEVICE_TARGETS="gfx908 gfx90a gfx940 gfx941 gfx942 gfx1030 gfx1100 gfx1101 gfx1102 gfx1200 gfx1201"
+ARG GPU_DEVICE_TARGETS="gfx908 gfx90a gfx940 gfx941 gfx942 gfx950 gfx1030 gfx1100 gfx1101 gfx1102 gfx1200 gfx1201"
 ENV GPU_DEVICE_TARGETS=${GPU_DEVICE_TARGETS}
 
 # Install ROCM
diff --git a/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile.rocm.ub20 b/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile.rocm.ub20
@@ -2,7 +2,7 @@
 FROM ubuntu:20.04
 ################################################################################
 
-ARG GPU_DEVICE_TARGETS="gfx908 gfx90a gfx940 gfx941 gfx942 gfx1030 gfx1100 gfx1101 gfx1102 gfx1200 gfx1201"
+ARG GPU_DEVICE_TARGETS="gfx908 gfx90a gfx940 gfx941 gfx942 gfx950 gfx1030 gfx1100 gfx1101 gfx1102 gfx1200 gfx1201"
 ENV GPU_DEVICE_TARGETS=${GPU_DEVICE_TARGETS}
 
 # Install build dependencies
diff --git a/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile.rocm.ub22 b/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile.rocm.ub22
@@ -2,7 +2,7 @@
 FROM ubuntu:22.04
 ################################################################################
 
-ARG GPU_DEVICE_TARGETS="gfx908 gfx90a gfx940 gfx941 gfx942 gfx1030 gfx1100 gfx1101 gfx1102 gfx1200 gfx1201"
+ARG GPU_DEVICE_TARGETS="gfx908 gfx90a gfx940 gfx941 gfx942 gfx950 gfx1030 gfx1100 gfx1101 gfx1102 gfx1200 gfx1201"
 ENV GPU_DEVICE_TARGETS=${GPU_DEVICE_TARGETS}
 
 # Install build dependencies
diff --git a/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile.rocm.ub24 b/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile.rocm.ub24
@@ -2,7 +2,7 @@
 FROM ubuntu:24.04
 ################################################################################
 
-ARG GPU_DEVICE_TARGETS="gfx908 gfx90a gfx940 gfx941 gfx942 gfx1030 gfx1100 gfx1101 gfx1102 gfx1200 gfx1201"
+ARG GPU_DEVICE_TARGETS="gfx908 gfx90a gfx940 gfx941 gfx942 gfx950 gfx1030 gfx1100 gfx1101 gfx1102 gfx1200 gfx1201"
 ENV GPU_DEVICE_TARGETS=${GPU_DEVICE_TARGETS}
 
 # Install build dependencies
diff --git a/tensorflow/tools/tf_sig_build_dockerfiles/setup.rocm.sh b/tensorflow/tools/tf_sig_build_dockerfiles/setup.rocm.sh
@@ -145,7 +145,7 @@ then
   echo "build:rocm_base --copt=-fclang-abi-compat=17" >> /etc/bazel.bazelrc
 fi
 
-GPU_DEVICE_TARGETS=${GPU_DEVICE_TARGETS:-"gfx908 gfx90a gfx940 gfx941 gfx942 gfx1030 gfx1100"}
+GPU_DEVICE_TARGETS=${GPU_DEVICE_TARGETS:-"gfx908 gfx90a gfx940 gfx941 gfx942 gfx950 gfx1030 gfx1100"}
 
 echo $ROCM_VERSION
 echo $ROCM_REPO
diff --git a/third_party/xla/xla/service/gpu/llvm_gpu_backend/amdgpu_backend.cc b/third_party/xla/xla/service/gpu/llvm_gpu_backend/amdgpu_backend.cc
@@ -340,7 +340,8 @@ std::string MapGCNArchNameTokenToFeatureStr(const std::string& token,
     return "+sramecc";
   } else if (token == "sramecc-") {
     if (gfx == "gfx90a" || gfx == "gfx940" || gfx == "gfx941" ||
-        gfx == "gfx942" || gfx == "gfx1101" || gfx == "gfx1102" ||
+        gfx == "gfx942" || gfx == "gfx950" ||
+        gfx == "gfx1101" || gfx == "gfx1102" ||
         gfx == "gfx1200" || gfx == "gfx1201")
       return "";
     return "-sramecc";
diff --git a/third_party/xla/xla/service/gpu/transforms/conv_rewriter_test.cc b/third_party/xla/xla/service/gpu/transforms/conv_rewriter_test.cc
@@ -791,6 +791,13 @@ TEST_F(ConvRewriterTest, TestInvalidTypes) {
                      ::testing::HasSubstr(
                          "FP8 convolutions are only supported on CUDA GPUs")));
 
+  s = ConvRewriter(se::RocmComputeCapability{"gfx950"}).Run(m.get()).status();
+  EXPECT_THAT(s, tsl::testing::StatusIs(
+                     absl::StatusCode::kUnimplemented,
+                     ::testing::HasSubstr(
+                         "FP8 convolutions are only supported on CUDA GPUs")));
+
+
   // Test unsupported FP8 type
   module_with_type = absl::StrReplaceAll(module_str, {{"TYPE", "f8e4m3fnuz"}});
   TF_ASSERT_OK_AND_ASSIGN(m, ParseAndReturnVerifiedModule(module_with_type));
diff --git a/third_party/xla/xla/stream_executor/device_description.h b/third_party/xla/xla/stream_executor/device_description.h
@@ -78,13 +78,13 @@ class RocmComputeCapability {
 
   bool gfx9_mi100_or_later() const {
     static constexpr absl::string_view kList[] = {"gfx908", "gfx90a", "gfx940",
-                                                  "gfx941", "gfx942"};
+                                                  "gfx941", "gfx942", "gfx950"};
     return absl::c_count(kList, gfx_version()) != 0;
   }
 
   bool gfx9_mi200_or_later() const {
     static constexpr absl::string_view kList[] = {"gfx90a", "gfx940", "gfx941",
-                                                  "gfx942"};
+                                                  "gfx942", "gfx950"};
     return absl::c_count(kList, gfx_version()) != 0;
   }
 
@@ -157,6 +157,7 @@ class RocmComputeCapability {
       "gfx908",                         // MI100
       "gfx90a",                         // MI200
       "gfx940",  "gfx941", "gfx942",    // MI300
+      "gfx950",                        
       "gfx1030",                        // RX68xx / RX69xx
       "gfx1100", "gfx1101", "gfx1102",  // RX7900
       "gfx1200", "gfx1201",             // RX8900      
