This is an automated email from the ASF dual-hosted git repository.
hongyij pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 0cf5f47a1e [Unity] Dispatch cumsum and sort (#16254)
0cf5f47a1e is described below
commit 0cf5f47a1e27b42a6734b2c63331b71fab96ac0a
Author: Yong Wu <[email protected]>
AuthorDate: Wed Jan 3 13:01:29 2024 -0800
[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'
---
include/tvm/relax/attrs/sort.h | 52 +++++
python/tvm/relax/backend/__init__.py | 1 +
python/tvm/relax/backend/dispatch_sort_scan.py | 102 +++++++++
python/tvm/relax/op/__init__.py | 1 +
python/tvm/relax/op/op_attrs.py | 5 +
.../tvm/relax/{backend/__init__.py => op/sort.py} | 31 ++-
python/tvm/relax/pipeline.py | 3 +-
python/tvm/relax/vm_build.py | 7 +-
python/tvm/script/ir_builder/relax/ir.py | 2 +
python/tvm/topi/cuda/sort.py | 1 -
src/relax/op/tensor/sort.cc | 56 +++++
src/relax/op/tensor/sort.h | 49 ++++
src/relax/transform/meta_schedule.cc | 10 +-
.../relax/test_backend_dispatch_sort_scan.py | 253 +++++++++++++++++++++
tests/python/relax/test_dataflow_pattern.py | 5 +-
tests/python/relax/test_op_sort.py | 102 +++++++++
.../python/relax/test_tvmscript_parser_op_sort.py | 54 +++++
17 files changed, 723 insertions(+), 11 deletions(-)
diff --git a/include/tvm/relax/attrs/sort.h b/include/tvm/relax/attrs/sort.h
new file mode 100644
index 0000000000..fc0c4e7189
--- /dev/null
+++ b/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_
diff --git a/python/tvm/relax/backend/__init__.py
b/python/tvm/relax/backend/__init__.py
index c3786591e3..e4a89bdb95 100644
--- a/python/tvm/relax/backend/__init__.py
+++ b/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
diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py
b/python/tvm/relax/backend/dispatch_sort_scan.py
new file mode 100644
index 0000000000..f0f1aa9063
--- /dev/null
+++ b/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()
diff --git a/python/tvm/relax/op/__init__.py b/python/tvm/relax/op/__init__.py
index 60a4332d83..085761f15d 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/__init__.py
@@ -99,6 +99,7 @@ from .mask import masked_fill
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 (
diff --git a/python/tvm/relax/op/op_attrs.py b/python/tvm/relax/op/op_attrs.py
index 848c2094a0..4dbbc17cf2 100644
--- a/python/tvm/relax/op/op_attrs.py
+++ b/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"""
diff --git a/python/tvm/relax/backend/__init__.py b/python/tvm/relax/op/sort.py
similarity index 54%
copy from python/tvm/relax/backend/__init__.py
copy to python/tvm/relax/op/sort.py
index c3786591e3..b139eefcdf 100644
--- a/python/tvm/relax/backend/__init__.py
+++ b/python/tvm/relax/op/sort.py
@@ -14,7 +14,32 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Relax backends"""
+"""Sortings operators."""
-from . import contrib
-from .pattern_registry import get_pattern, get_patterns_with_prefix
+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
diff --git a/python/tvm/relax/pipeline.py b/python/tvm/relax/pipeline.py
index ebcbd2d609..474833bdfd 100644
--- a/python/tvm/relax/pipeline.py
+++ b/python/tvm/relax/pipeline.py
@@ -24,7 +24,7 @@ as it is or serves as a basis to do further composition.
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(),
diff --git a/python/tvm/relax/vm_build.py b/python/tvm/relax/vm_build.py
index 9120f74e13..ca756d4dc6 100644
--- a/python/tvm/relax/vm_build.py
+++ b/python/tvm/relax/vm_build.py
@@ -328,7 +328,12 @@ def build(
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()
diff --git a/python/tvm/script/ir_builder/relax/ir.py
b/python/tvm/script/ir_builder/relax/ir.py
index 142d0e6d96..7c0be2a722 100644
--- a/python/tvm/script/ir_builder/relax/ir.py
+++ b/python/tvm/script/ir_builder/relax/ir.py
@@ -135,6 +135,7 @@ from tvm.relax.op import (
sign,
sin,
sinh,
+ sort,
split,
square,
squeeze,
@@ -758,6 +759,7 @@ __all__ = [
"sign",
"sin",
"sinh",
+ "sort",
"split",
"square",
"squeeze",
diff --git a/python/tvm/topi/cuda/sort.py b/python/tvm/topi/cuda/sort.py
index 24b46f9a88..058584a302 100644
--- a/python/tvm/topi/cuda/sort.py
+++ b/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
diff --git a/src/relax/op/tensor/sort.cc b/src/relax/op/tensor/sort.cc
new file mode 100644
index 0000000000..31de102ccf
--- /dev/null
+++ b/src/relax/op/tensor/sort.cc
@@ -0,0 +1,56 @@
+/*
+ * 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 sort.cc
+ * \brief sorting operators.
+ */
+
+#include "sort.h"
+
+namespace tvm {
+namespace relax {
+
+/* relax.sort */
+TVM_REGISTER_NODE_TYPE(SortAttrs);
+
+Expr sort(Expr data, int axis, bool descending) {
+ auto attrs = make_object<SortAttrs>();
+ attrs->axis = std::move(axis);
+ attrs->descending = std::move(descending);
+
+ static const Op& op = Op::Get("relax.sort");
+ return Call(op, {std::move(data)}, Attrs{attrs}, {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.sort").set_body_typed(sort);
+
+StructInfo InferStructInfoSort(const Call& call, const BlockBuilder& ctx) {
+ return GetUnaryInputTensorStructInfo(call, ctx);
+}
+
+TVM_REGISTER_OP("relax.sort")
+ .set_attrs_type<SortAttrs>()
+ .set_num_inputs(1)
+ .add_argument("data", "Tensor", "The input tensor.")
+ .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoSort)
+ .set_attr<Bool>("FPurity", Bool(true));
+
+} // namespace relax
+} // namespace tvm
diff --git a/src/relax/op/tensor/sort.h b/src/relax/op/tensor/sort.h
new file mode 100644
index 0000000000..92203034aa
--- /dev/null
+++ b/src/relax/op/tensor/sort.h
@@ -0,0 +1,49 @@
+/*
+ * 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 sort.h
+ * \brief The functions to make Relax tensor sorting operator calls.
+ */
+#ifndef TVM_RELAX_OP_TENSOR_SORT_H_
+#define TVM_RELAX_OP_TENSOR_SORT_H_
+
+#include <tvm/relax/attrs/sort.h>
+
+#include <algorithm>
+#include <utility>
+
+#include "../op_common.h"
+
+namespace tvm {
+namespace relax {
+
+/*!
+ * \brief Reverses the order of elements along given axis.
+ * \param data The input tensor.
+ * \param axis The axis to sort on.
+ * \param descending Whether to sort in descending order.
+ * \return The computed result.
+ */
+Expr sort(Expr data, int axis, bool descending);
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_OP_TENSOR_SORT_H_
diff --git a/src/relax/transform/meta_schedule.cc
b/src/relax/transform/meta_schedule.cc
index 2fef2dc933..c54e75e9cb 100644
--- a/src/relax/transform/meta_schedule.cc
+++ b/src/relax/transform/meta_schedule.cc
@@ -152,9 +152,13 @@ Pass MetaScheduleApplyDatabase(Optional<String> work_dir,
bool enable_warning =
ICHECK_EQ(new_mod->functions.size(), 1);
BaseFunc new_base_func = (*new_mod->functions.begin()).second;
ICHECK(new_base_func->IsInstance<tir::PrimFuncNode>());
- tir::PrimFunc new_prim_func = Downcast<tir::PrimFunc>(new_base_func);
- // copy the original attrs
- new_prim_func = WithAttrs(std::move(new_prim_func),
{prim_func->attrs->dict});
+ tir::PrimFunc tuned_prim_func =
Downcast<tir::PrimFunc>(new_base_func);
+ // maintain the original attributes
+ tir::PrimFunc new_prim_func =
tir::PrimFunc(/*params=*/tuned_prim_func->params,
+
/*body=*/tuned_prim_func->body,
+
/*ret_type=*/tuned_prim_func->ret_type,
+
/*buffer_map=*/tuned_prim_func->buffer_map,
+
/*attrs=*/prim_func->attrs);
new_prim_func = WithAttr(std::move(new_prim_func),
tir::attr::kIsScheduled, Bool(true));
result.Set(gv, new_prim_func);
continue;
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py
b/tests/python/relax/test_backend_dispatch_sort_scan.py
new file mode 100644
index 0000000000..79234c2af8
--- /dev/null
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -0,0 +1,253 @@
+# 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.
+
+import pytest
+
+import tvm
+from tvm import topi, relax, tir
+import tvm.script
+import tvm.testing
+from tvm.script import relax as R, tir as T, ir as I
+from tvm.contrib.thrust import can_use_thrust
+
+
+from tvm.relax.backend import DispatchSortScan
+from tvm.ir.base import assert_structural_equal
+
+
+def test_dispatch_cumsum():
+ @I.ir_module
+ class Before:
+ I.module_global_infos({"vdevice": [I.vdevice("cuda", 0),
I.vdevice("llvm", 0)]})
+
+ @R.function
+ def foo(x: R.Tensor((2, 3), "float32", "llvm")):
+ with R.dataflow():
+ gv = R.cumsum(x, axis=1, dtype="float64")
+ R.output(gv)
+ return gv
+
+ @I.ir_module
+ class Expected:
+ I.module_global_infos({"vdevice": [I.vdevice("cuda", 0),
I.vdevice("llvm", 0)]})
+
+ @T.prim_func(private=True)
+ def cumsum(var_A: T.handle, out_buf: T.Buffer((T.int64(2),
T.int64(3)), "float64")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ A = T.match_buffer(var_A, (T.int64(2), T.int64(3)),
offset_factor=1)
+ with T.block("cumsum_generic"):
+ T.reads(A[T.int64(0) : T.int64(2), T.int64(0) : T.int64(3)])
+ T.writes(out_buf[T.int64(0) : T.int64(2), T.int64(0) :
T.int64(3)])
+ for fused in T.parallel(T.int64(2)):
+ out_buf[
+ fused * T.int64(3) // T.int64(3), fused * T.int64(3) %
T.int64(3)
+ ] = T.Cast(
+ "float64",
+ A[fused * T.int64(3) // T.int64(3), fused * T.int64(3)
% T.int64(3)],
+ )
+ for _k in range(T.int64(2)):
+ out_buf[
+ (fused * T.int64(3) + (_k + T.int64(1))) //
T.int64(3),
+ (fused * T.int64(3) + (_k + T.int64(1))) %
T.int64(3),
+ ] = out_buf[
+ (fused * T.int64(3) + (_k + T.int64(1) -
T.int64(1))) // T.int64(3),
+ (fused * T.int64(3) + (_k + T.int64(1) -
T.int64(1))) % T.int64(3),
+ ] + T.Cast(
+ "float64",
+ A[
+ (fused * T.int64(3) + (_k + T.int64(1))) //
T.int64(3),
+ (fused * T.int64(3) + (_k + T.int64(1))) %
T.int64(3),
+ ],
+ )
+
+ @R.function
+ def foo(
+ x: R.Tensor((2, 3), dtype="float32", vdevice="llvm")
+ ) -> R.Tensor((2, 3), dtype="float64", vdevice="llvm"):
+ cls = Expected
+ with R.dataflow():
+ gv = R.call_tir(cls.cumsum, (x,), out_sinfo=R.Tensor((2, 3),
"float64", "llvm"))
+ R.output(gv)
+ return gv
+
+ mod = DispatchSortScan()(Before)
+ assert_structural_equal(mod, Expected)
+
+
+def test_dispatch_cumsum_cuda():
+ @I.ir_module
+ class Before:
+ I.module_global_infos({"vdevice": [I.vdevice("cuda", 0),
I.vdevice("llvm", 0)]})
+
+ @R.function
+ def main(x: R.Tensor(("m", 3), "float32", "cuda")):
+ with R.dataflow():
+ lv = R.cumsum(x, axis=1)
+ gv = lv
+ R.output(gv)
+ return gv
+
+ target = tvm.target.Target("cuda", host="llvm")
+
+ vdevices = [I.vdevice("cuda", 0), I.vdevice("llvm", 0)]
+ m = tir.Var("m", "int64")
+ x = relax.Var("x", R.Tensor((m, 3), "float32", vdevices[0]))
+ bb = relax.BlockBuilder()
+ with target:
+ with bb.function("main", (x,), {"global_symbol": "main"}):
+ with bb.dataflow():
+ out = bb.emit_te(
+ topi.cuda.cumsum,
+ x,
+ axis=1,
+ )
+ out = bb.emit_output(out)
+ bb.emit_func_output(out)
+ expected_mod = bb.finalize()
+ expected_mod.update_global_info("vdevice", vdevices)
+
+ with target:
+ mod = DispatchSortScan()(Before)
+
+ assert_structural_equal(mod, expected_mod)
+
+
+def test_dispatch_sort():
+ @I.ir_module
+ class Before:
+ I.module_global_infos({"vdevice": [I.vdevice("cuda", 0),
I.vdevice("llvm", 0)]})
+
+ @R.function
+ def foo(x: R.Tensor(("m", 3), "float32", "llvm")):
+ m = T.int64()
+ with R.dataflow():
+ gv = R.sort(x, axis=1, descending=False)
+ R.output(gv)
+ return gv
+
+ @I.ir_module
+ class Expected:
+ I.module_global_infos({"vdevice": [I.vdevice("cuda", 0),
I.vdevice("llvm", 0)]})
+
+ @T.prim_func(private=True)
+ def sort(var_A: T.handle, var_sort_cpu: T.handle):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ m = T.int64()
+ data_buf = T.match_buffer(var_A, (m, T.int64(3)), align=8)
+ out_buf = T.match_buffer(var_sort_cpu, (m, T.int64(3)), align=8)
+ with T.block("sort_cpu"):
+ T.reads(data_buf[T.int64(0) : m, T.int64(0) : T.int64(3)])
+ T.writes(out_buf[T.int64(0) : m, T.int64(0) : T.int64(3)])
+ T.call_packed(
+ "tvm.contrib.sort.sort",
+ T.tvm_stack_make_array(
+ data_buf.data,
+ T.tvm_stack_make_shape(m, T.int64(3)),
+ 0,
+ 2,
+ T.float32(0),
+ T.int64(0),
+ ),
+ T.tvm_stack_make_array(
+ out_buf.data,
+ T.tvm_stack_make_shape(m, T.int64(3)),
+ 0,
+ 2,
+ T.float32(0),
+ T.int64(0),
+ ),
+ 1,
+ T.bool(True),
+ )
+
+ @R.function
+ def foo(
+ x: R.Tensor(("m", 3), dtype="float32", vdevice="llvm")
+ ) -> R.Tensor(("m", 3), dtype="float32", vdevice="llvm"):
+ m = T.int64()
+ cls = Expected
+ with R.dataflow():
+ gv = R.call_tir(
+ cls.sort, (x,), out_sinfo=R.Tensor((m, 3),
dtype="float32", vdevice="llvm")
+ )
+ R.output(gv)
+ return gv
+
+ mod = DispatchSortScan()(Before)
+ assert_structural_equal(mod, Expected)
+
+
+def test_dispatch_sort_cuda():
+ @I.ir_module
+ class Before:
+ I.module_global_infos({"vdevice": [I.vdevice("cuda"),
I.vdevice("llvm")]})
+
+ @R.function
+ def foo(x: R.Tensor((2, 3), "float32", "cuda")):
+ with R.dataflow():
+ lv = R.sort(x, axis=1, descending=False)
+ gv = lv
+ R.output(gv)
+ return gv
+
+ @R.function
+ def foo2(y: R.Tensor((2, 3), "float32")):
+ with R.dataflow():
+ lv = R.sort(y, axis=0, descending=True)
+ gv = lv
+ R.output(gv)
+ return gv
+
+ target = tvm.target.Target("cuda -libs=thrust", host="llvm")
+
+ vdevices = [I.vdevice("cuda", 0), I.vdevice("llvm", 0)]
+ x = relax.Var("x", R.Tensor((2, 3), "float32", vdevices[0]))
+ y = relax.Var("y", R.Tensor((2, 3), "float32"))
+ bb = relax.BlockBuilder()
+ with target:
+ with bb.function("foo", (x,), {"global_symbol": "foo"}):
+ with bb.dataflow():
+ out = bb.emit_te(
+ topi.cuda.sort,
+ x,
+ axis=1,
+ )
+ out = bb.emit_output(out)
+ bb.emit_func_output(out)
+ with bb.function("foo2", (y,), {"global_symbol": "foo2"}):
+ with bb.dataflow():
+ out = bb.emit_te(
+ topi.cuda.sort_thrust
+ if can_use_thrust(target, "tvm.contrib.thrust.sort")
+ else topi.cuda.sort,
+ y,
+ 0,
+ False,
+ )
+ out = bb.emit_output(out)
+ bb.emit_func_output(out)
+ expected_mod = bb.finalize()
+ expected_mod.update_global_info("vdevice", vdevices)
+
+ with target:
+ mod = DispatchSortScan()(Before)
+
+ assert_structural_equal(mod, expected_mod, map_free_vars=True)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_dataflow_pattern.py
b/tests/python/relax/test_dataflow_pattern.py
index edd3bd1610..7f2cb241bb 100644
--- a/tests/python/relax/test_dataflow_pattern.py
+++ b/tests/python/relax/test_dataflow_pattern.py
@@ -39,7 +39,7 @@ class Module:
B = T.match_buffer(y, (32, 32))
C = T.match_buffer(z, (32, 32))
- for (i0, j0, k0) in T.grid(32, 32, 32):
+ for i0, j0, k0 in T.grid(32, 32, 32):
with T.block():
i, j, k = T.axis.remap("SSR", [i0, j0, k0])
with T.init():
@@ -51,7 +51,7 @@ class Module:
T.func_attr({"global_symbol": "tir_relu"})
A = T.match_buffer(x, (32, 32))
B = T.match_buffer(y, (32, 32))
- for (i, j) in T.grid(32, 32):
+ for i, j in T.grid(32, 32):
with T.block():
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = T.max(A[vi, vj], 0.0)
@@ -69,6 +69,7 @@ class Module:
main_fn = Module["main"]
bindings = main_fn.body.blocks[0].bindings
+
## Node-wise Matching
def test_expr_pattern():
ep = is_expr(rx.Var("x"))
diff --git a/tests/python/relax/test_op_sort.py
b/tests/python/relax/test_op_sort.py
new file mode 100644
index 0000000000..b6a064a641
--- /dev/null
+++ b/tests/python/relax/test_op_sort.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.
+import pytest
+import tvm
+import tvm.testing
+from tvm import relax, tir
+from tvm import TVMError
+from tvm.ir import Op, VDevice
+from tvm.script import relax as R
+
+
+def test_op_correctness():
+ x = relax.Var("x", R.Tensor((3, 4, 5), "float32"))
+ assert relax.op.sort(x, axis=1).op == Op.get("relax.sort")
+
+
+def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_sinfo:
relax.StructInfo):
+ ret = bb.normalize(call)
+ tvm.ir.assert_structural_equal(ret.struct_info, expected_sinfo)
+
+
+def test_sort_infer_struct_info():
+ bb = relax.BlockBuilder()
+ vdev0 = VDevice("llvm")
+ x0 = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
+ x1 = relax.Var("x", R.Tensor("float32", ndim=3))
+ x2 = relax.Var("x", R.Tensor("float32"))
+ x3 = relax.Var("x", R.Tensor((2, 10, 4)))
+ x4 = relax.Var("x", R.Tensor(ndim=3))
+ x5 = relax.Var("x", R.Tensor())
+ x6 = relax.Var("x", R.Tensor((2, 10, 4), "float32", vdev0))
+
+ _check_inference(bb, relax.op.sort(x0, axis=1), relax.TensorStructInfo((2,
10, 4), "float32"))
+ _check_inference(
+ bb, relax.op.sort(x6, axis=1), relax.TensorStructInfo((2, 10, 4),
"float32", vdev0)
+ )
+ _check_inference(bb, relax.op.sort(x1, axis=1),
relax.TensorStructInfo(dtype="float32", ndim=3))
+ _check_inference(bb, relax.op.sort(x2, axis=1),
relax.TensorStructInfo(dtype="float32"))
+ _check_inference(bb, relax.op.sort(x3, axis=1), relax.TensorStructInfo((2,
10, 4), dtype=""))
+ _check_inference(bb, relax.op.sort(x4, axis=1),
relax.TensorStructInfo(dtype="", ndim=3))
+ _check_inference(bb, relax.op.sort(x5, axis=1),
relax.TensorStructInfo(dtype=""))
+ _check_inference(bb, relax.op.sort(x0), relax.TensorStructInfo((2, 10, 4),
"float32"))
+ _check_inference(
+ bb,
+ relax.op.sort(x0, axis=1, descending=False),
+ relax.TensorStructInfo((2, 10, 4), "float32"),
+ )
+
+
+def test_sort_infer_struct_info_shape_symbolic():
+ bb = relax.BlockBuilder()
+ a = tir.Var("a", "int64")
+ b = tir.Var("b", "int64")
+ c = tir.Var("c", "int64")
+ x = relax.Var("x", R.Tensor((a, b, c), "float32"))
+
+ _check_inference(bb, relax.op.sort(x, axis=1), relax.TensorStructInfo((a,
b, c), "float32"))
+ _check_inference(bb, relax.op.sort(x), relax.TensorStructInfo((a, b, c),
"float32"))
+
+
+def test_sort_infer_struct_info_more_input_dtype():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", R.Tensor((2, 3, 4), "float16"))
+ x1 = relax.Var("x", R.Tensor((2, 3, 4), "int8"))
+
+ _check_inference(bb, relax.op.sort(x0, axis=1), relax.TensorStructInfo((2,
3, 4), "float16"))
+ _check_inference(bb, relax.op.sort(x1, axis=1), relax.TensorStructInfo((2,
3, 4), "int8"))
+
+
+def test_sort_wrong_input():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", relax.ShapeStructInfo((2, 3, 4, 5)))
+ x1 = relax.Var("x", relax.FuncStructInfo([], R.Tensor((2, 3, 4, 5),
"float32")))
+ x = relax.Var("x", R.Tensor((3, 4, 5), "float32"))
+ y = relax.Var("y", R.Tensor((2, 3, 4), "float32"))
+
+ with pytest.raises(TVMError):
+ relax.op.sort(x, y)
+
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.sort(x0, axis=1))
+
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.sort(x1, axis=1))
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_tvmscript_parser_op_sort.py
b/tests/python/relax/test_tvmscript_parser_op_sort.py
new file mode 100644
index 0000000000..8b94fa0ab9
--- /dev/null
+++ b/tests/python/relax/test_tvmscript_parser_op_sort.py
@@ -0,0 +1,54 @@
+# 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.
+
+from typing import Optional, Union
+
+import tvm
+import tvm.script
+import tvm.testing
+from tvm import IRModule, relax
+from tvm.script import relax as R
+
+
+def _check(
+ parsed: Union[relax.Function, IRModule],
+ expect: Optional[Union[relax.Function, IRModule]],
+):
+ test = parsed.script(show_meta=True)
+ roundtrip_mod = tvm.script.from_source(test)
+ tvm.ir.assert_structural_equal(parsed, roundtrip_mod)
+ if expect:
+ tvm.ir.assert_structural_equal(parsed, expect)
+
+
+def test_sort():
+ @R.function
+ def foo(x: R.Tensor((2, 3), "int32")) -> R.Tensor((2, 3), "int32"):
+ r = R.sort(x, axis=1)
+ return r
+
+ x = relax.Var("x", R.Tensor((2, 3), "int32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x,)):
+ tensor = bb.emit(relax.op.sort(x, axis=1))
+ bb.emit_func_output(tensor)
+
+ _check(foo, bb.get()["foo"])
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
