Recover non-standard bool test for msort (#139870)

xw285cornell · pytorchmergebot · commit 565a7942eee1 · 2024-11-11T02:00:34.000Z
Summary: I was looking into why the non-standard bool value will fail for msort - it makes sense for argsort and sort to fail, because we're randomly generating uint8 so the order will be different (and thus the indices will be different). But msort should work. After some digging, it's interesting that even though scalar_t is bool, when the actual value is a uint8_t, the comparison will treat them as signed. I tried lhs=255 and rhs=0: lhs < rhs is equivalent to -1 < 0 which is true (but it's supposed to be False) Therefore we add an explicit type cast. Test Plan: Remove the test skip Differential Revision: D65472170 Pull Request resolved: #139870 Approved by: https://github.com/Skylion007, https://github.com/davidberard98
diff --git a/aten/src/ATen/native/cuda/SortingCommon.cuh b/aten/src/ATen/native/cuda/SortingCommon.cuh
@@ -46,14 +46,16 @@ inline bool getGridFromTiles(int64_t gridTiles, dim3& grid) {
 template <typename scalar_t, bool handleNaN = false>
 struct GTOp {
   __device__ bool operator()(const scalar_t& lhs, const scalar_t& rhs) const {
-    return (handleNaN && at::_isnan(lhs) && !at::_isnan(rhs)) || (lhs > rhs);
+    return (handleNaN && at::_isnan(lhs) && !at::_isnan(rhs)) ||
+        (static_cast<scalar_t>(lhs) > static_cast<scalar_t>(rhs));
   }
 };
 
 template <typename scalar_t, bool handleNaN = false>
 struct LTOp {
   __device__ bool operator()(const scalar_t& lhs, const scalar_t& rhs) const {
-    return (handleNaN && at::_isnan(rhs) && !at::_isnan(lhs)) || (lhs < rhs);
+    return (handleNaN && at::_isnan(rhs) && !at::_isnan(lhs)) ||
+        (static_cast<scalar_t>(lhs) < static_cast<scalar_t>(rhs));
   }
 };
 
diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py
@@ -19517,8 +19517,9 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
            supports_fwgrad_bwgrad=True,
            sample_inputs_func=sample_inputs_msort,
            skips=(
+               # https://github.com/pytorch/pytorch/issues/139972
                DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_non_standard_bool_values',
-                            dtypes=[torch.bool], device_type='cuda'),
+                            dtypes=[torch.bool], device_type='cuda', active_if=TEST_WITH_ROCM),
            )),
     OpInfo('movedim',
            aliases=('moveaxis',),
