[release/2.6] fix failed tests on MI350

Author: dnikolaev-amd

aten/src/ATen/native/cuda/SortStable.cu

@@ -226,8 +226,9 @@ void launch_stable_sort_kernel(
     return;
   }
 
-  int64_t numel_or_intmax =
-      std::min(numel, static_cast<int64_t>(std::numeric_limits<int>::max()));
+  // On ROCm, std::min -> ::min did not work as expected on when input values >= 2147483648
+  const int64_t intmax = static_cast<int64_t>(std::numeric_limits<int>::max());
+  int64_t numel_or_intmax = numel < intmax ? numel : intmax;
   int64_t nsort = self.size(dim);
   int64_t nbatch = (numel_or_intmax / nsort) * nsort;
   TORCH_CHECK(nbatch > 0, "Cannot sort dimension of length ", nsort);
@@ -239,7 +240,8 @@ void launch_stable_sort_kernel(
         scalar_t* values_ptr = values.mutable_data_ptr<scalar_t>();
         int64_t remaining = numel;
         while (remaining > 0) {
-          int64_t n = std::min(remaining, nbatch);
+          // On ROCm, std::min -> ::min did not work as expected on when input values >= 2147483648
+          int64_t n = remaining < nbatch ? remaining : nbatch;
           int64_t nsegments = n / nsort;
 
           if (nsegments == 1 ||

test/test_sort_and_select.py

@@ -209,21 +209,21 @@ def test_stable_sort(self, device, dtype):
             )
 
     @onlyCUDA
-    @dtypes(torch.uint8)
+    @dtypes(torch.float16)
     @largeTensorTest("200GB")  # Unfortunately 80GB A100 is not large enough
     def test_sort_large(self, device, dtype):
         t0 = torch.randperm(8192, device=device).to(dtype)
         t = t0.view(1, 8192).expand(2**18 + 1, -1).contiguous()
         v, i = t.sort()
         del t
-        iv, im = i.var_mean(dim=0)
+        iv, im = torch.var_mean(i.to(dtype), dim=0)
         del i
-        vv, vm = v.var_mean(dim=0)
+        vv, vm = torch.var_mean(v.to(dtype), dim=0)
         del v
         self.assertEqual(vv, torch.zeros_like(vv))
         self.assertEqual(iv, torch.zeros_like(iv))
-        self.assertEqual(vm, torch.arange(255, dtype=dtype, device=device))
-        self.assertEqual(im, t0.sort().indices)
+        self.assertEqual(vm, torch.arange(8192, dtype=dtype, device=device))
+        self.assertEqual(im, t0.sort().indices, exact_dtype=False)
 
     @dtypes(torch.float32)
     def test_sort_restride(self, device, dtype):