Increase the absolute tolerance between gradients in CTCLossDeterministicTest.testForwardAndBackward due to the floating-point arithmetics for gfx11xx and gfx12xx. Implement test_util.gpu_gcn_arch() which returns the information on the current gcn architecture. Implement unit test for test_util.gpu_gcn_arch(). Modify GetShortDeviceDescription() in gpu_device.cc to return gcn arch as well.

Author: AleksaArsic

tensorflow/core/common_runtime/gpu/gpu_device.cc

@@ -1970,6 +1970,7 @@ static string GetShortDeviceDescription(
 #elif TENSORFLOW_USE_ROCM
   return strings::StrCat("device: ", platform_device_id.value(),
                          ", name: ", desc.name(),
+                         ", gcn arch: ", desc.rocm_compute_capability().gcn_arch_name(),
                          ", pci bus id: ", desc.pci_bus_id());
 #endif
 }

tensorflow/python/framework/test_util.py

@@ -189,6 +189,29 @@ def gpu_device_name() -> str:
       return compat.as_str(x.name)
   return ""
 
+@tf_export("test.gpu_gcn_arch")
+def gpu_gcn_arch() -> str:
+  """ Returns the GCN Arch if GPU available or an empty string.
+
+  This method should only be used in tests written with tf.test.TestCase
+
+  >>> class MyTest(tf.test.TestCase):
+  ...
+  ...   if not tf_test.is_built_with_rocm():
+  ...    self.skipTest("Test is only applicable for GPUs built with ROCm")
+  ...
+  ...   self.assertNotEqual("", test_util.gpu_gcn_arch())
+
+  """
+  for x in device_lib.list_local_devices():
+    if x.device_type == "GPU":
+      desc = getattr(x, "physical_device_desc", "")
+      gcn_arch = re.search(r"gfx[0-9]+", desc)
+      
+      if gcn_arch:
+        return compat.as_str(gcn_arch.group(0))
+
+  return ""
 
 def assert_ops_in_graph(
     expected_ops: dict[str, str], graph: ops.Graph

tensorflow/python/framework/test_util_test.py

@@ -45,6 +45,7 @@
 from tensorflow.python.framework import tensor
 from tensorflow.python.framework import test_ops
 from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test as tf_test
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_assert
 from tensorflow.python.ops import lookup_ops
@@ -1094,6 +1095,11 @@ def some_test(self):
     some_test(None)
     self.assertEqual(tested_codepaths, set(["present", "future"]))
 
+  def test_assert_gcn_arch(self):
+    if not tf_test.is_built_with_rocm():
+      self.skipTest("Test is only applicable for GPUs built with ROCm")
+
+    self.assertNotEqual("", test_util.gpu_gcn_arch())
 
 class SkipTestTest(test_util.TensorFlowTestCase):
 

tensorflow/python/kernel_tests/nn_ops/ctc_loss_op_test.py

@@ -1233,8 +1233,16 @@ def testForwardAndBackward(self, sparse_labels, logits_time_major):
         loss_a, loss_b, gradient_a, gradient_b = self.evaluate(
             (loss_a, loss_b, gradient_a, gradient_b))
         self.assertAllEqual(loss_a, loss_b, "Loss mismatch")
-        # self.assertAllEqual(gradient_a, gradient_b, "Gradient mismatch")
-        self.assertAllClose(gradient_a, gradient_b, atol=5e-05)
+        # Determine which gcn architecture is the GPU and set the absolute 
+        # tolerance based on that information.
+        # Needed on gfx11 and gfx12 due to the floating point arithmetic.
+        gcn_arch =  test_util.gpu_gcn_arch()
+        if "gfx11" or "gfx12" in gcn_arch:
+            abs_tolerance = 1e-4
+        else:
+            abs_tolerance = 5e-5
+
+        self.assertAllClose(gradient_a, gradient_b, atol=abs_tolerance)
 
 
 if __name__ == "__main__":