override sort operators for UT functionality (#4956) (#4967)

* fix example case test_mode in test_torch_mode.py

* fix comments

* fix ut

Author: huaiyuzh

csrc/gpu/aten/operators/Sort.cpp

@@ -1,6 +1,9 @@
 #include <core/Memory.h>
 #include <core/detail/IndexUtils.h>
 #include <core/detail/TensorInfo.h>
+#ifdef USE_OVERRIDE_OP
+#include "utils/CustomOperatorRegistration.h"
+#endif
 #include <utils/DPCPP.h>
 #include "comm/ATDispatch.h"
 #include "comm/Numerics.h"
@@ -55,6 +58,56 @@ Tensor argsort(const Tensor& self, bool stable, int64_t dim, bool descending) {
   return std::get<1>(
       sort_out_stable(self, stable, dim, descending, sorted, indices));
 }
+#ifdef USE_OVERRIDE_OP
+std::tuple<at::Tensor, at::Tensor> sort_ipex(
+    const at::Tensor& self,
+    int64_t dim,
+    bool descending) {
+  return at::AtenIpexTypeXPU::sort(self, dim, descending);
+}
+
+std::tuple<Tensor&, Tensor&> sort_values(
+    const Tensor& input,
+    int64_t dim,
+    bool order,
+    Tensor& sorted,
+    Tensor& indices) {
+  return at::AtenIpexTypeXPU::sort_out(input, dim, order, sorted, indices);
+}
+
+std::tuple<Tensor, Tensor> sort_stable(
+    const Tensor& self,
+    c10::optional<bool> stable,
+    int64_t dim,
+    bool descending) {
+  return at::AtenIpexTypeXPU::sort(self, stable, dim, descending);
+}
 
+std::tuple<Tensor&, Tensor&> sort_values_stable(
+    const Tensor& self,
+    c10::optional<bool> stable,
+    int64_t dim,
+    bool descending,
+    Tensor& values,
+    Tensor& indices) {
+  return at::AtenIpexTypeXPU::sort_out(
+      self, stable, dim, descending, values, indices);
+}
+#endif
 } // namespace AtenIpexTypeXPU
 } // namespace at
+
+#ifdef USE_OVERRIDE_OP
+namespace {
+
+IPEX_TORCH_LIBRARY_IMPL(aten, XPU, m) {
+  m.impl("sort", TORCH_FN((&at::AtenIpexTypeXPU::sort_ipex)));
+  m.impl("sort.stable", TORCH_FN((&at::AtenIpexTypeXPU::sort_stable)));
+  m.impl("sort.values", TORCH_FN((&at::AtenIpexTypeXPU::sort_values)));
+  m.impl(
+      "sort.values_stable",
+      TORCH_FN((&at::AtenIpexTypeXPU::sort_values_stable)));
+}
+
+} // namespace
+#endif

scripts/tools/torchgen/yaml/xpu_functions.yaml

@@ -18,6 +18,10 @@ supported:
 # - col2im.out
 # - im2col
 # - im2col.out
+#  - sort
+#  - sort.stable
+#  - sort.values
+#  - sort.values_stable
 ################## override below ops due to performance issues
 # - convolution_overrideable
 # - convolution_backward_overrideable
@@ -585,10 +589,6 @@ supported:
 #  - softplus_backward.grad_input
 #  - softshrink.out
 #  - softshrink_backward.grad_input
-  - sort
-#  - sort.stable
-  - sort.values
-#  - sort.values_stable
   - special_entr.out
   - special_erfcx.out
   - special_i0e.out

tests/gpu/examples/test_torch_mode.py

@@ -2,15 +2,12 @@
 from torch.testing._internal.common_utils import TestCase
 import intel_extension_for_pytorch  # noqa
 
-import pytest
-
 cpu_device = torch.device("cpu")
 xpu_device = torch.device("xpu")
 value_range = 30
 
 
 class TestTorchMethod(TestCase):
-    @pytest.mark.skip(reason="PT2.5: Scalars are not equal!")
     def test_mode(self):
         def mode_test_helper(input_list):
             for input_cpu in input_list: