[Unity] Dispatch cumsum and sort (#16254)

* [Unity] Add dispatch for scan and sort

* add test cases

* Use pass instead of pattern rewriter

* add test case

* fix lint

* fix comments

* fix lint

* Add target context for default pipeline

* fix tests

* remove 'is_scheduled'

Author: yongwww

include/tvm/relax/attrs/sort.h

@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/relax/attrs/sort.h
+ * \brief Attributes for sorting operators.
+ */
+#ifndef TVM_RELAX_ATTRS_SORT_H_
+#define TVM_RELAX_ATTRS_SORT_H_
+
+#include <tvm/relax/expr.h>
+#include <tvm/tir/index_map.h>
+
+namespace tvm {
+namespace relax {
+
+/*! \brief Attributes used in sort operator */
+struct SortAttrs : public tvm::AttrsNode<SortAttrs> {
+  int axis;
+  bool descending;
+
+  TVM_DECLARE_ATTRS(SortAttrs, "relax.attrs.SortAttrs") {
+    TVM_ATTR_FIELD(axis).set_default(-1).describe(
+        "Axis along which the sort is computed."
+        "The default the last axis is used.");
+    TVM_ATTR_FIELD(descending)
+        .set_default(false)
+        .describe(
+            "Whether to sort in descending order."
+            "If it is not specified, it defaults to the ascending order.");
+  }
+};  // struct SortAttrs
+}  // namespace relax
+}  // namespace tvm
+
+#endif  // TVM_RELAX_ATTRS_SORT_H_

python/tvm/relax/backend/__init__.py

@@ -18,3 +18,4 @@
 
 from . import contrib
 from .pattern_registry import get_pattern, get_patterns_with_prefix
+from .dispatch_sort_scan import DispatchSortScan

python/tvm/relax/backend/dispatch_sort_scan.py

@@ -0,0 +1,102 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, unused-argument, redefined-argument-from-local
+"""Dispatch sort and scan operators to platform dependent implementation."""
+
+from tvm import topi
+from tvm.ir import Op
+from tvm.ir.module import IRModule
+from tvm.ir.transform import PassContext, module_pass
+from tvm.target import Target
+from tvm.contrib.thrust import can_use_thrust
+from tvm.relax import Expr, Function, Call, PyExprMutator, expr_functor, TensorStructInfo
+
+
+@expr_functor.mutator
+class SortScanDispatcher(PyExprMutator):
+    """
+    Dispatcher to dispatch sort and scan.
+
+    """
+
+    def __init__(self, mod):
+        super().__init__(mod)
+
+    def _get_target(self, expr: Expr) -> Target:
+        sinfo = expr.struct_info
+        # Get target information from TensorStructInfo
+        if isinstance(sinfo, TensorStructInfo):
+            vdevice = sinfo.vdevice
+            if vdevice is not None:
+                return vdevice.target
+        # Return the target in current context
+        target = Target.current()
+        if target is None:
+            raise ValueError(
+                "Target not found. Please ensure that the target is annotated within the module, "
+                "or alternatively, execute this within a specified target context."
+            )
+        return target
+
+    def visit_call_(self, call: Call) -> Expr:
+        if not isinstance(call.op, Op):
+            return super().visit_call_(call)
+
+        if call.op.name == "relax.sort":
+            tgt = self._get_target(call)
+            with tgt:
+                if can_use_thrust(tgt, "tvm.contrib.thrust.sort"):
+                    return self.builder_.call_te(
+                        topi.cuda.sort_thrust,
+                        call.args[0],
+                        call.attrs.axis,
+                        not call.attrs.descending,
+                    )
+                return self.builder_.call_te(
+                    topi.cuda.sort if tgt.kind.name == "cuda" else topi.sort,
+                    call.args[0],
+                    call.attrs.axis,
+                    not call.attrs.descending,
+                )
+
+        if call.op.name == "relax.cumsum":
+            tgt = self._get_target(call)
+            axis = int(call.attrs.axis) if call.attrs.axis is not None else call.attrs.axis
+            with tgt:
+                return self.builder_.call_te(
+                    topi.cuda.cumsum if tgt.kind.name == "cuda" else topi.cumsum,
+                    call.args[0],
+                    axis,
+                    call.attrs.dtype,
+                )
+
+        return super().visit_call_(call)
+
+
+@module_pass(opt_level=0, name="DispatchSortScan")
+class DispatchSortScan:
+    """
+    Pass to dispatch scan and sort operators to platform dependent implementation.
+    """
+
+    def transform_module(self, mod: IRModule, ctx: PassContext) -> IRModule:
+        sort_scan_dispater = SortScanDispatcher(mod)
+        for gv, func in mod.functions_items():
+            if isinstance(func, Function):
+                func = sort_scan_dispater.visit_expr(func)
+                sort_scan_dispater.builder_.update_func(gv, func)
+        return sort_scan_dispater.builder_.get()

python/tvm/relax/op/__init__.py

@@ -99,6 +99,7 @@
 from .qdq import quantize, dequantize
 from .search import argmax, argmin, where
 from .set import unique
+from .sort import sort
 from .statistical import cumsum, max, mean, min, prod, std, sum, variance
 from .ternary import ewise_fma
 from .unary import (

python/tvm/relax/op/op_attrs.py

@@ -114,6 +114,11 @@ class PermuteDimsAttrs(Attrs):
     """Attributes for permute_dims operator"""
 
 
+@tvm._ffi.register_object("relax.attrs.SortAttrs")
+class SortAttrs(Attrs):
+    """Attributes for sort operator"""
+
+
 @tvm._ffi.register_object("relax.attrs.SplitAttrs")
 class SplitAttrs(Attrs):
     """Attributes used in split operator"""

python/tvm/relax/op/sort.py

@@ -0,0 +1,45 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Sortings operators."""
+
+from . import _ffi_api
+from ..expr import Expr
+
+
+def sort(x: Expr, axis: int = -1, descending: bool = False):
+    """Performs sorting along the given axis and returns an array
+    in sorted order.
+
+    Parameters
+    ----------
+    x : relax.Expr
+        The input tensor.
+
+    axis : int
+        Axis along which to sort the input tensor.
+        By default the last axis of the input is used.
+
+    descending : bool
+        Whether to sort in descending order, the default is False
+
+    Returns
+    -------
+    out : relax.Expr
+        Sorted tensor.
+
+    """
+    return _ffi_api.sort(x, axis, descending)  # type: ignore

python/tvm/relax/pipeline.py

@@ -24,7 +24,7 @@
 import tvm
 from tvm import meta_schedule as ms
 
-from . import transform
+from . import transform, backend
 
 
 def zero_pipeline(*, enable_warning: bool = False):
@@ -81,6 +81,7 @@ def default_build_pipeline():
     def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
         seq = tvm.transform.Sequential(
             [
+                backend.DispatchSortScan(),
                 transform.LegalizeOps(),
                 transform.RewriteDataflowReshape(),
                 transform.ToNonDataflow(),

python/tvm/relax/vm_build.py

@@ -328,7 +328,12 @@ def _extract_attrs(mod: tvm.IRModule):
     if pipeline is not None:
         if isinstance(pipeline, str):
             pipeline = relax.get_pipeline(pipeline)
-        mod = pipeline(mod)
+        if target is None:
+            mod = pipeline(mod)
+        else:
+            with target:
+                mod = pipeline(mod)
+
     ext_libs, constants = _extract_attrs(mod)
     params.update(dict(constants))
     builder = relax.ExecBuilder()

python/tvm/script/ir_builder/relax/ir.py

@@ -135,6 +135,7 @@
     sign,
     sin,
     sinh,
+    sort,
     split,
     square,
     squeeze,
@@ -758,6 +759,7 @@ def dtype(value: Union[py_str, DataType]) -> Expr:
     "sign",
     "sin",
     "sinh",
+    "sort",
     "split",
     "square",
     "squeeze",

python/tvm/topi/cuda/sort.py

@@ -120,7 +120,6 @@ def _odd_even_sort(
     values=None,
     values_swap=None,
 ):
-
     nthread_tx = block_size // 2
     nthread_bx = ceil_div(size, block_size)
     nthread_by = axis_mul_before