This is an automated email from the ASF dual-hosted git repository.
tqchen 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 e1d71b3720 [Unity] Add dlight.gpu.Fallback in DispatchSortScan, add
argsort, topk, and cumprod (#16351)
e1d71b3720 is described below
commit e1d71b3720347ba566e5100a0dc7c4fc7fc054a5
Author: Yong Wu <[email protected]>
AuthorDate: Tue Jan 9 17:47:43 2024 -0800
[Unity] Add dlight.gpu.Fallback in DispatchSortScan, add argsort, topk, and
cumprod (#16351)
---
include/tvm/relax/attrs/sort.h | 52 ----
include/tvm/relax/attrs/sorting.h | 99 +++++++
include/tvm/relax/attrs/statistical.h | 20 +-
python/tvm/relax/backend/dispatch_sort_scan.py | 124 ++++++---
python/tvm/relax/op/__init__.py | 4 +-
python/tvm/relax/op/op_attrs.py | 16 +-
python/tvm/relax/op/sort.py | 45 ---
python/tvm/relax/op/sorting.py | 116 ++++++++
python/tvm/relax/op/statistical.py | 72 ++++-
.../relax/transform/legalize_ops/statistical.py | 11 +-
python/tvm/script/ir_builder/relax/ir.py | 6 +
src/relax/op/tensor/sort.cc | 56 ----
src/relax/op/tensor/sorting.cc | 155 +++++++++++
src/relax/op/tensor/{sort.h => sorting.h} | 32 ++-
src/relax/op/tensor/statistical.cc | 55 ++--
src/relax/op/tensor/statistical.h | 21 +-
.../relax/test_backend_dispatch_sort_scan.py | 301 +++++++++++++++++++--
tests/python/relax/test_op_sort.py | 192 +++++++++++++
tests/python/relax/test_op_statistical.py | 53 ++--
.../python/relax/test_tvmscript_parser_op_sort.py | 14 +-
.../relax/test_tvmscript_parser_op_statistical.py | 8 +-
21 files changed, 1178 insertions(+), 274 deletions(-)
diff --git a/include/tvm/relax/attrs/sort.h b/include/tvm/relax/attrs/sort.h
deleted file mode 100644
index fc0c4e7189..0000000000
--- a/include/tvm/relax/attrs/sort.h
+++ /dev/null
@@ -1,52 +0,0 @@
-/*
- * 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/include/tvm/relax/attrs/sorting.h
b/include/tvm/relax/attrs/sorting.h
new file mode 100644
index 0000000000..4daf7a45b2
--- /dev/null
+++ b/include/tvm/relax/attrs/sorting.h
@@ -0,0 +1,99 @@
+/*
+ * 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/sorting.h
+ * \brief Attributes for sorting operators.
+ */
+#ifndef TVM_RELAX_ATTRS_SORTING_H_
+#define TVM_RELAX_ATTRS_SORTING_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
+
+/*! \brief Attributes used in argsort operator */
+struct ArgsortAttrs : public tvm::AttrsNode<ArgsortAttrs> {
+ int axis;
+ bool descending;
+ DataType dtype;
+
+ TVM_DECLARE_ATTRS(ArgsortAttrs, "relax.attrs.ArgsortAttrs") {
+ TVM_ATTR_FIELD(axis).set_default(-1).describe(
+ "Axis along which the argsort is computed."
+ "The default the last axis is used.");
+ TVM_ATTR_FIELD(descending)
+ .set_default(false)
+ .describe(
+ "Whether to argsort in descending order."
+ "If it is not specified, it defaults to the ascending order.");
+ TVM_ATTR_FIELD(dtype)
+ .set_default(NullValue<DataType>())
+ .describe("DType of the output indices.");
+ }
+}; // struct ArgsortAttrs
+
+/*! \brief Attributes used in topk operator */
+struct TopKAttrs : public tvm::AttrsNode<TopKAttrs> {
+ int k;
+ int axis;
+ bool largest;
+ String ret_type;
+ DataType dtype;
+
+ TVM_DECLARE_ATTRS(TopKAttrs, "relax.attrs.TopKAttrs") {
+ TVM_ATTR_FIELD(k).describe("Number of top elements to select");
+ TVM_ATTR_FIELD(axis).set_default(-1).describe("Axis along which to sort
the input tensor.");
+ TVM_ATTR_FIELD(ret_type).set_default("both").describe(
+ "The return type [both, values, indices]."
+ "both - return both top k data and indices."
+ "values - return top k data only."
+ "indices - return top k indices only.");
+ TVM_ATTR_FIELD(largest).set_default(true).describe(
+ "Whether to return largest or smallest elements."
+ "By default, return the largest k elements.");
+ TVM_ATTR_FIELD(dtype)
+ .set_default(NullValue<DataType>())
+ .describe("Data type of the output indices.");
+ }
+}; // struct TopKAttrs
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_ATTRS_SORTING_H_
diff --git a/include/tvm/relax/attrs/statistical.h
b/include/tvm/relax/attrs/statistical.h
index d5d8d982b5..9f9a2fa870 100644
--- a/include/tvm/relax/attrs/statistical.h
+++ b/include/tvm/relax/attrs/statistical.h
@@ -42,20 +42,24 @@ struct StatisticalAttrs : public
tvm::AttrsNode<StatisticalAttrs> {
}
}; // struct StatisticalAttrs
-/*! \brief Attributes used in cumsum operators */
-struct CumsumAttrs : public tvm::AttrsNode<CumsumAttrs> {
+/*! \brief Attributes used in scan operators like cumsum, cumprod */
+struct ScanopAttrs : public tvm::AttrsNode<ScanopAttrs> {
Optional<Integer> axis;
DataType dtype;
+ Bool exclusive = Bool(false);
- TVM_DECLARE_ATTRS(CumsumAttrs, "relax.attrs.CumsumAttrs") {
+ TVM_DECLARE_ATTRS(ScanopAttrs, "relax.attrs.ScanopAttrs") {
TVM_ATTR_FIELD(axis).describe(
- "Axis along which the cumulative sum is computed."
- "The default (None) is to compute the cumsum over the flattened
array.");
+ "The axis along which to perform the scan computation."
+ "The default (None) is to compute over the flattened array.");
TVM_ATTR_FIELD(dtype).describe(
- "Type of the returned array and of the accumulator in which the
elements are summed."
- "If dtype is not specified, it defaults to the dtype of data.");
+ "The output data type."
+ "If dtype is not specified, it defaults to the dtype of input data.");
+ TVM_ATTR_FIELD(exclusive)
+ .describe("The first element is not included")
+ .set_default(Bool(false));
}
-}; // struct CumsumAttrs
+}; // struct ScanopAttrs
} // namespace relax
} // namespace tvm
diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py
b/python/tvm/relax/backend/dispatch_sort_scan.py
index f0f1aa9063..bb3f57ce96 100644
--- a/python/tvm/relax/backend/dispatch_sort_scan.py
+++ b/python/tvm/relax/backend/dispatch_sort_scan.py
@@ -17,13 +17,13 @@
# 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 import topi, dlight, relax
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
+from tvm.relax import PyExprMutator, expr_functor
@expr_functor.mutator
@@ -36,13 +36,17 @@ class SortScanDispatcher(PyExprMutator):
def __init__(self, mod):
super().__init__(mod)
- def _get_target(self, expr: Expr) -> Target:
- sinfo = expr.struct_info
+ def _get_target(self, sinfo: relax.StructInfo) -> Target:
# Get target information from TensorStructInfo
- if isinstance(sinfo, TensorStructInfo):
+ if isinstance(sinfo, relax.TensorStructInfo):
vdevice = sinfo.vdevice
if vdevice is not None:
return vdevice.target
+ elif isinstance(sinfo, relax.TupleStructInfo):
+ for f in sinfo.fields:
+ tgt = self._get_target(f)
+ if tgt != Target.current():
+ return tgt
# Return the target in current context
target = Target.current()
if target is None:
@@ -52,38 +56,94 @@ class SortScanDispatcher(PyExprMutator):
)
return target
- def visit_call_(self, call: Call) -> Expr:
+ def _apply_dlight_gpu_fallback(self, target: Target, tir_call: relax.Call)
-> None:
+ # Apply dlight.gpu.Fallback() on GPU
+ gvar = tir_call.args[0]
+ assert isinstance(gvar, relax.GlobalVar)
+ scan_prim_func = self.builder_.get()[gvar]
+ sch = dlight.base.transform._apply_rules(
+ scan_prim_func,
+ target,
+ [
+ dlight.gpu.Fallback(),
+ ],
+ False,
+ )
+ if sch is not None:
+ assert len(sch) == 1
+ self.builder_.update_func(gvar,
sch[0].mod["main"].with_attr("tir.is_scheduled", 1))
+
+ def visit_call_(self, call: relax.Call) -> relax.Expr:
if not isinstance(call.op, Op):
return super().visit_call_(call)
if call.op.name == "relax.sort":
- tgt = self._get_target(call)
+ tgt = self._get_target(call.struct_info)
+ te_func = topi.sort
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
+ te_func = topi.cuda.sort_thrust
+ elif tgt.kind.name == "cuda":
+ te_func = topi.cuda.sort
+ return self.builder_.call_te(
+ te_func,
+ call.args[0],
+ call.attrs.axis,
+ not call.attrs.descending,
+ )
+ if call.op.name == "relax.argsort":
+ tgt = self._get_target(call.struct_info)
+ te_func = topi.argsort
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,
- )
-
+ if can_use_thrust(tgt, "tvm.contrib.thrust.sort"):
+ te_func = topi.cuda.argsort_thrust
+ elif tgt.kind.name == "cuda":
+ te_func = topi.cuda.argsort
+ return self.builder_.call_te(
+ te_func,
+ call.args[0],
+ axis=call.attrs.axis,
+ is_ascend=not call.attrs.descending,
+ dtype=call.attrs.dtype,
+ )
+ if call.op.name == "relax.topk":
+ tgt = self._get_target(call.struct_info)
+ te_func = topi.topk
+ if can_use_thrust(tgt, "tvm.contrib.thrust.sort"):
+ te_func = topi.cuda.topk_thrust
+ elif tgt.kind.name == "cuda":
+ te_func = topi.cuda.topk
+ tir_call = self.builder_.call_te(
+ te_func,
+ call.args[0],
+ axis=call.attrs.axis,
+ ret_type=call.attrs.ret_type,
+ is_ascend=not call.attrs.largest,
+ dtype=call.attrs.dtype,
+ )
+ if tgt.kind.name != "cuda":
+ return tir_call
+ # apply dlight gpu fallback
+ self._apply_dlight_gpu_fallback(tgt, tir_call)
+ return tir_call
+ if call.op.name in ("relax.cumprod", "relax.cumsum"):
+ tgt = self._get_target(call.struct_info)
+ axis = int(call.attrs.axis) if call.attrs.axis is not None else
call.attrs.axis
+ te_func = topi.cuda.cumsum if tgt.kind.name == "cuda" else
topi.cumsum
+ if call.op.name == "relax.cumprod":
+ te_func = topi.cuda.cumprod if tgt.kind.name == "cuda" else
topi.cumprod
+ tir_call = self.builder_.call_te(
+ te_func,
+ call.args[0],
+ axis,
+ call.attrs.dtype,
+ call.attrs.exclusive,
+ )
+ if tgt.kind.name != "cuda":
+ return tir_call
+ # apply dlight gpu fallback
+ self._apply_dlight_gpu_fallback(tgt, tir_call)
+ return tir_call
return super().visit_call_(call)
@@ -96,7 +156,7 @@ class DispatchSortScan:
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):
+ if isinstance(func, relax.Function):
func = sort_scan_dispater.visit_expr(func)
sort_scan_dispater.builder_.update_func(gv, func)
- return sort_scan_dispater.builder_.get()
+ return sort_scan_dispater.builder_.finalize()
diff --git a/python/tvm/relax/op/__init__.py b/python/tvm/relax/op/__init__.py
index 085761f15d..5b585e18b4 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/__init__.py
@@ -99,8 +99,8 @@ 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 .sorting import sort, argsort, topk
+from .statistical import cumsum, cumprod, max, mean, min, prod, std, sum,
variance
from .ternary import ewise_fma
from .unary import (
abs,
diff --git a/python/tvm/relax/op/op_attrs.py b/python/tvm/relax/op/op_attrs.py
index 4dbbc17cf2..a3d46428c5 100644
--- a/python/tvm/relax/op/op_attrs.py
+++ b/python/tvm/relax/op/op_attrs.py
@@ -119,6 +119,11 @@ class SortAttrs(Attrs):
"""Attributes for sort operator"""
+@tvm._ffi.register_object("relax.attrs.ArgsortAttrs")
+class ArgsortAttrs(Attrs):
+ """Attributes for argsort operator"""
+
+
@tvm._ffi.register_object("relax.attrs.SplitAttrs")
class SplitAttrs(Attrs):
"""Attributes used in split operator"""
@@ -154,9 +159,14 @@ class TileAttrs(Attrs):
"""Attributes for tile operator"""
-@tvm._ffi.register_object("relax.attrs.CumsumAttrs")
-class CumsumAttrs(Attrs):
- """Attributes for cumsum operator"""
+@tvm._ffi.register_object("relax.attrs.ScanopAttrs")
+class ScanopAttrs(Attrs):
+ """Attributes for scan operators"""
+
+
+@tvm._ffi.register_object("relax.attrs.TopKAttrs")
+class TopKAttrs(Attrs):
+ """Attributes for topk operators"""
@tvm._ffi.register_object("relax.attrs.EinsumAttrs")
diff --git a/python/tvm/relax/op/sort.py b/python/tvm/relax/op/sort.py
deleted file mode 100644
index b139eefcdf..0000000000
--- a/python/tvm/relax/op/sort.py
+++ /dev/null
@@ -1,45 +0,0 @@
-# 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
diff --git a/python/tvm/relax/op/sorting.py b/python/tvm/relax/op/sorting.py
new file mode 100644
index 0000000000..13937933c4
--- /dev/null
+++ b/python/tvm/relax/op/sorting.py
@@ -0,0 +1,116 @@
+# 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, Constant
+
+
+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
+
+
+def argsort(data: Expr, axis: int = -1, descending: bool = False, dtype: str =
"int32"):
+ """Performs sorting along the given axis and returns an array of indices
+ having same shape as an input array that index data in sorted order.
+
+ Parameters
+ ----------
+ data : relax.Expr
+ The input data tensor.
+
+ axis : int
+ Axis long which to sort the input tensor.
+
+ descending : bool
+ Whether to sort in descending order, the default is False
+
+ dtype : str
+ The data type of the output indices.
+
+ Returns
+ -------
+ out : relax.Expr
+ Tensor with same shape as data.
+ """
+ return _ffi_api.argsort(data, axis, descending, dtype) # type: ignore
+
+
+def topk(
+ data: Expr,
+ k: int = 1,
+ axis: int = -1,
+ ret_type: str = "both",
+ largest: bool = True,
+ dtype: str = "int32",
+):
+ """Get the top k elements in an input tensor along the given axis.
+
+ ret_type specifies the return type, can be one of ("both", "values",
"indices").
+
+ Parameters
+ ----------
+ data : relax.Expr
+ The input data tensor.
+
+ k : int
+ Number of top elements to select. Return all elements if k < 1.
+
+ axis : int
+ Axis long which to sort the input tensor.
+
+ ret_type: str
+ The return type [both, values, indices].
+ "both": return both top k data and indices.
+ "values": return top k data only.
+ "indices": return top k indices only.
+
+ largest : bool
+ Whether to return largest or smallest elements.
+ The k smallest elements are returned if largest is False.
+
+ dtype : str
+ The data type of the indices output.
+
+ Returns
+ -------
+ out : relax.Expr or List[relax.Expr]
+ The computed result.
+ """
+ if isinstance(k, Constant):
+ k = k.data.numpy().item()
+ return _ffi_api.topk(data, k, axis, ret_type, largest, dtype) # type:
ignore
diff --git a/python/tvm/relax/op/statistical.py
b/python/tvm/relax/op/statistical.py
index f187f9d456..eb44696871 100644
--- a/python/tvm/relax/op/statistical.py
+++ b/python/tvm/relax/op/statistical.py
@@ -191,7 +191,71 @@ def sum(x: Expr, axis: Optional[Union[int, List[int]]] =
None, keepdims: bool =
return _ffi_api.sum(x, axis, keepdims) # type: ignore
-def cumsum(data: Expr, axis: Optional[int] = None, dtype: Optional[Union[str,
DataType]] = None):
+def cumprod(
+ data: Expr,
+ axis: Optional[int] = None,
+ dtype: Optional[Union[str, DataType]] = None,
+ exclusive: Optional[bool] = None,
+):
+ """Numpy style cumprod op. Return the cumulative product of the elements
along
+ a given axis.
+
+ Parameters
+ ----------
+ data : relax.Expr
+ The input data to the operator.
+
+ axis : Optional[int]
+ Axis along which the cumulative product is computed. The default
(None) is to compute
+ the cumprod over the flattened array.
+
+ dtype : Optional[Union[str, DataType]]
+ Type of the returned array and of the accumulator in which the
elements are computed.
+ If dtype is not specified, it defaults to the dtype of data.
+
+ exclusive : Optional[bool]
+ If true will return exclusive sum in which the first element is not
+ included.
+
+ Returns
+ -------
+ result : relax.Expr
+ The result has the same size as data, and the same shape as data if
axis is not None.
+ If axis is None, the result is a 1-d array.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ a = [[1, 2, 3], [4, 5, 6]]
+
+ cumprod(a) # if axis is not provided, cumprod is done over the
flattened input.
+ -> [ 1, 2, 6, 24, 120, 720]
+
+ cumprod(a, dtype="float32")
+ -> [ 1., 2., 6., 24., 120., 720.]
+
+ cumprod(a, axis=0) # multiply over rows for each of the 3 columns
+ -> [[1, 2, 3],
+ [4, 10, 18]]
+
+ cumprod(a, axis=1)
+ -> [[ 1, 2, 6],
+ [ 4, 20, 120]]
+
+ a = [1, 1, 1, 0, 1, 1, 0] # a is a boolean array
+ cumprod(a, dtype=int32) # dtype should be provided to get the
expected results
+ -> [1, 1, 1, 0, 0, 0, 0]
+ """
+ return _ffi_api.cumprod(data, axis, dtype, exclusive) # type: ignore
+
+
+def cumsum(
+ data: Expr,
+ axis: Optional[int] = None,
+ dtype: Optional[Union[str, DataType]] = None,
+ exclusive: Optional[bool] = None,
+):
"""Numpy style cumsum op. Return the cumulative inclusive sum of the
elements along
a given axis.
@@ -208,6 +272,10 @@ def cumsum(data: Expr, axis: Optional[int] = None, dtype:
Optional[Union[str, Da
Type of the returned array and of the accumulator in which the
elements are summed.
If dtype is not specified, it defaults to the dtype of data.
+ exclusive : Optional[bool]
+ If true will return exclusive sum in which the first element is not
+ included.
+
Returns
-------
result : relax.Expr
@@ -238,7 +306,7 @@ def cumsum(data: Expr, axis: Optional[int] = None, dtype:
Optional[Union[str, Da
cumsum(a, dtype=int32) # dtype should be provided to get the expected
results
-> [1, 1, 2, 2, 3, 4, 4]
"""
- return _ffi_api.cumsum(data, axis, dtype) # type: ignore
+ return _ffi_api.cumsum(data, axis, dtype, exclusive) # type: ignore
def variance(x: Expr, axis: Optional[Union[int, List[int]]] = None, keepdims:
bool = False) -> Expr:
diff --git a/python/tvm/relax/transform/legalize_ops/statistical.py
b/python/tvm/relax/transform/legalize_ops/statistical.py
index e1f273bda0..1181b3b2a7 100644
--- a/python/tvm/relax/transform/legalize_ops/statistical.py
+++ b/python/tvm/relax/transform/legalize_ops/statistical.py
@@ -89,4 +89,13 @@ register_legalize("relax.sum", _statistical(topi.sum))
@register_legalize("relax.cumsum")
def _cumsum(bb: BlockBuilder, call: Call) -> Expr:
- return bb.call_te(topi.cumsum, call.args[0], call.attrs.axis,
call.attrs.dtype)
+ return bb.call_te(
+ topi.cumsum, call.args[0], call.attrs.axis, call.attrs.dtype,
call.attrs.exclusive
+ )
+
+
+@register_legalize("relax.cumprod")
+def _cumprod(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(
+ topi.cumprod, call.args[0], call.attrs.axis, call.attrs.dtype,
call.attrs.exclusive
+ )
diff --git a/python/tvm/script/ir_builder/relax/ir.py
b/python/tvm/script/ir_builder/relax/ir.py
index 7c0be2a722..9105fce00f 100644
--- a/python/tvm/script/ir_builder/relax/ir.py
+++ b/python/tvm/script/ir_builder/relax/ir.py
@@ -43,6 +43,7 @@ from tvm.relax.op import (
arange,
argmax,
argmin,
+ argsort,
assert_op,
astype,
bitwise_and,
@@ -65,6 +66,7 @@ from tvm.relax.op import (
concat,
cos,
cosh,
+ cumprod,
cumsum,
einsum,
scatter_elements,
@@ -145,6 +147,7 @@ from tvm.relax.op import (
tanh,
erf,
tile,
+ topk,
tril,
triu,
unique,
@@ -644,6 +647,7 @@ __all__ = [
"arg",
"argmax",
"argmin",
+ "argsort",
"assert_op",
"astype",
"bitwise_and",
@@ -670,6 +674,7 @@ __all__ = [
"const",
"cpu",
"cuda",
+ "cumprod",
"cumsum",
"einsum",
"scatter_elements",
@@ -773,6 +778,7 @@ __all__ = [
"tan",
"tanh",
"tile",
+ "topk",
"to_vdevice",
"tril",
"triu",
diff --git a/src/relax/op/tensor/sort.cc b/src/relax/op/tensor/sort.cc
deleted file mode 100644
index 31de102ccf..0000000000
--- a/src/relax/op/tensor/sort.cc
+++ /dev/null
@@ -1,56 +0,0 @@
-/*
- * 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/sorting.cc b/src/relax/op/tensor/sorting.cc
new file mode 100644
index 0000000000..c4c4c5a614
--- /dev/null
+++ b/src/relax/op/tensor/sorting.cc
@@ -0,0 +1,155 @@
+/*
+ * 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 sorting.cc
+ * \brief sorting operators.
+ */
+
+#include "sorting.h"
+
+#include <vector>
+
+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));
+
+/* relax.argsort */
+TVM_REGISTER_NODE_TYPE(ArgsortAttrs);
+
+Expr argsort(Expr data, int axis, bool descending, DataType dtype) {
+ auto attrs = make_object<ArgsortAttrs>();
+ attrs->axis = std::move(axis);
+ attrs->descending = std::move(descending);
+ attrs->dtype = std::move(dtype);
+
+ static const Op& op = Op::Get("relax.argsort");
+ return Call(op, {std::move(data)}, Attrs{attrs}, {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.argsort").set_body_typed(argsort);
+
+StructInfo InferStructInfoArgsort(const Call& call, const BlockBuilder& ctx) {
+ TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
+ const auto* attrs = call->attrs.as<ArgsortAttrs>();
+ DataType out_type = attrs->dtype.is_void() ? data_sinfo->dtype :
attrs->dtype;
+ if (data_sinfo->shape.defined()) {
+ return TensorStructInfo(data_sinfo->shape.value(), out_type,
data_sinfo->vdevice);
+ }
+ return TensorStructInfo(out_type, data_sinfo->ndim, data_sinfo->vdevice);
+}
+
+TVM_REGISTER_OP("relax.argsort")
+ .set_attrs_type<ArgsortAttrs>()
+ .set_num_inputs(1)
+ .add_argument("data", "Tensor", "The input tensor.")
+ .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoArgsort)
+ .set_attr<Bool>("FPurity", Bool(true));
+
+/* relax.topk */
+TVM_REGISTER_NODE_TYPE(TopKAttrs);
+
+Expr topk(Expr data, int k, int axis, String ret_type, bool largest, DataType
dtype) {
+ auto attrs = make_object<TopKAttrs>();
+ attrs->k = std::move(k);
+ attrs->axis = std::move(axis);
+ attrs->ret_type = std::move(ret_type);
+ attrs->largest = std::move(largest);
+ attrs->dtype = std::move(dtype);
+
+ static const Op& op = Op::Get("relax.topk");
+ return Call(op, {std::move(data)}, Attrs{attrs}, {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.topk").set_body_typed(topk);
+
+StructInfo InferStructInfoTopK(const Call& call, const BlockBuilder& ctx) {
+ TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
+ const auto* data_shape = data_sinfo->shape.as<ShapeExprNode>();
+ const auto* attrs = call->attrs.as<TopKAttrs>();
+ DataType indices_type = attrs->dtype.is_void() ? data_sinfo->dtype :
attrs->dtype;
+ int ndim = data_sinfo->ndim;
+ int k = attrs->k;
+ String ret_type = attrs->ret_type;
+ int axis = attrs->axis;
+ if (axis < 0 && ndim > 0) {
+ axis += ndim;
+ }
+
+ std::vector<StructInfo> output_sinfos;
+ output_sinfos.reserve(2);
+ if (data_shape == nullptr) {
+ output_sinfos.push_back(
+ TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice));
+ output_sinfos.push_back(TensorStructInfo(indices_type, data_sinfo->ndim,
data_sinfo->vdevice));
+ } else {
+ Array<PrimExpr> out_shape = data_shape->values;
+ const auto* int_dim = out_shape[axis].as<IntImmNode>();
+ if (k > 0 && (int_dim == nullptr || k < int_dim->value)) {
+ out_shape.Set(axis, k);
+ }
+ output_sinfos.push_back(
+ TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype,
data_sinfo->vdevice));
+ output_sinfos.push_back(
+ TensorStructInfo(ShapeExpr(out_shape), indices_type,
data_sinfo->vdevice));
+ }
+
+ if (ret_type == "both") {
+ return TupleStructInfo(output_sinfos);
+ } else if (ret_type == "values") {
+ return output_sinfos[0];
+ } else if (ret_type == "indices") {
+ return output_sinfos[1];
+ }
+ LOG(FATAL) << "Unsupported ret type: " << ret_type;
+}
+
+TVM_REGISTER_OP("relax.topk")
+ .set_attrs_type<TopKAttrs>()
+ .set_num_inputs(1)
+ .add_argument("data", "Tensor", "The input tensor.")
+ .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoTopK)
+ .set_attr<Bool>("FPurity", Bool(true));
+
+} // namespace relax
+} // namespace tvm
diff --git a/src/relax/op/tensor/sort.h b/src/relax/op/tensor/sorting.h
similarity index 55%
rename from src/relax/op/tensor/sort.h
rename to src/relax/op/tensor/sorting.h
index 92203034aa..8a785bc4e2 100644
--- a/src/relax/op/tensor/sort.h
+++ b/src/relax/op/tensor/sorting.h
@@ -18,13 +18,13 @@
*/
/*!
- * \file sort.h
+ * \file sorting.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_
+#ifndef TVM_RELAX_OP_TENSOR_SORTING_H_
+#define TVM_RELAX_OP_TENSOR_SORTING_H_
-#include <tvm/relax/attrs/sort.h>
+#include <tvm/relax/attrs/sorting.h>
#include <algorithm>
#include <utility>
@@ -43,7 +43,29 @@ namespace relax {
*/
Expr sort(Expr data, int axis, bool descending);
+/*!
+ * \brief Performs sorting along the given axis and returns an array of
indices.
+ * \param data The input tensor.
+ * \param axis The axis to sort on.
+ * \param descending Whether to sort in descending order.
+ * \param dtype The data type of the output indices.
+ * \return The computed result.
+ */
+Expr argsort(Expr data, int axis, bool descending, DataType dtype);
+
+/*!
+ * \brief Get the top k elements in an input tensor along the given axis.
+ * \param data The input tensor.
+ * \param k Number of top elements.
+ * \param axis The axis to sort on.
+ * \param ret_type The return type, can be set to one of [both, values,
indices].
+ * \param largest Whether to return largest or smallest elements.
+ * \param dtype The data type of the indices output.
+ * \return The computed result.
+ */
+Expr topk(Expr data, int k, int axis, String ret_type, bool largest, DataType
dtype);
+
} // namespace relax
} // namespace tvm
-#endif // TVM_RELAX_OP_TENSOR_SORT_H_
+#endif // TVM_RELAX_OP_TENSOR_SORTING_H_
diff --git a/src/relax/op/tensor/statistical.cc
b/src/relax/op/tensor/statistical.cc
index b861aafe21..24ccde4559 100644
--- a/src/relax/op/tensor/statistical.cc
+++ b/src/relax/op/tensor/statistical.cc
@@ -135,23 +135,11 @@ InferLayoutOutput InferLayoutStatistical(const Call& call,
Attrs(new_attrs));
}
-/* relax.cumsum */
-TVM_REGISTER_NODE_TYPE(CumsumAttrs);
-
-Expr cumsum(Expr data, Optional<Integer> axis, DataType dtype) {
- auto attrs = make_object<CumsumAttrs>();
- attrs->axis = std::move(axis);
- attrs->dtype = std::move(dtype);
-
- static const Op& op = Op::Get("relax.cumsum");
- return Call(op, {std::move(data)}, Attrs{attrs}, {});
-}
+TVM_REGISTER_NODE_TYPE(ScanopAttrs);
-TVM_REGISTER_GLOBAL("relax.op.cumsum").set_body_typed(cumsum);
-
-StructInfo InferStructInfoCumsum(const Call& call, const BlockBuilder& ctx) {
+StructInfo InferStructInfoScan(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
- const auto* attrs = call->attrs.as<CumsumAttrs>();
+ const auto* attrs = call->attrs.as<ScanopAttrs>();
DataType out_type = attrs->dtype.is_void() ? data_sinfo->dtype :
attrs->dtype;
@@ -177,11 +165,44 @@ StructInfo InferStructInfoCumsum(const Call& call, const
BlockBuilder& ctx) {
}
}
+/* relax.cumprod */
+Expr cumprod(Expr data, Optional<Integer> axis, DataType dtype, Bool
exclusive) {
+ auto attrs = make_object<ScanopAttrs>();
+ attrs->axis = std::move(axis);
+ attrs->dtype = std::move(dtype);
+ attrs->exclusive = std::move(exclusive);
+
+ static const Op& op = Op::Get("relax.cumprod");
+ return Call(op, {std::move(data)}, Attrs{attrs}, {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.cumprod").set_body_typed(cumprod);
+
+TVM_REGISTER_OP("relax.cumprod")
+ .set_attrs_type<ScanopAttrs>()
+ .set_num_inputs(1)
+ .add_argument("data", "Tensor", "The input tensor.")
+ .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoScan)
+ .set_attr<Bool>("FPurity", Bool(true));
+
+/* relax.cumsum */
+Expr cumsum(Expr data, Optional<Integer> axis, DataType dtype, Bool exclusive)
{
+ auto attrs = make_object<ScanopAttrs>();
+ attrs->axis = std::move(axis);
+ attrs->dtype = std::move(dtype);
+ attrs->exclusive = std::move(exclusive);
+
+ static const Op& op = Op::Get("relax.cumsum");
+ return Call(op, {std::move(data)}, Attrs{attrs}, {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.cumsum").set_body_typed(cumsum);
+
TVM_REGISTER_OP("relax.cumsum")
- .set_attrs_type<CumsumAttrs>()
+ .set_attrs_type<ScanopAttrs>()
.set_num_inputs(1)
.add_argument("data", "Tensor", "The input tensor.")
- .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoCumsum)
+ .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoScan)
.set_attr<Bool>("FPurity", Bool(true));
TVM_REGISTER_NODE_TYPE(StatisticalAttrs);
diff --git a/src/relax/op/tensor/statistical.h
b/src/relax/op/tensor/statistical.h
index 23a6da99f1..310c87f7d6 100644
--- a/src/relax/op/tensor/statistical.h
+++ b/src/relax/op/tensor/statistical.h
@@ -85,6 +85,22 @@ Expr std(Expr x, Optional<Array<Integer>> axis, bool
keepdims);
/*! \brief Computes the sum of tensor elements over given axes. */
Expr sum(Expr x, Optional<Array<Integer>> axis, bool keepdims);
+/*!
+ * \brief Numpy style cumprod op. Return the cumulative inclusive product of
the elements along
+ * a given axis.
+ * \param data The input tensor.
+ * \param axis Axis along which the cumulative product is computed. The
default (None) is to compute
+ * the cumprod over the flattened array.
+ * \param dtype Type of the returned array and of the accumulator in which the
elements are
+ * computed. If dtype is not specified, it defaults to the dtype of data.
+ * \param exclusive Whehter the first element is exclusive. If true will
return exclusive sum in
+ * which the first element is not included.
+ * \return The computed
+ * result.
+ */
+Expr cumprod(Expr data, Optional<Integer> axis = NullOpt, DataType dtype =
DataType::Void(),
+ Bool exclusive = Bool(false));
+
/*!
* \brief Numpy style cumsum op. Return the cumulative inclusive sum of the
elements along
* a given axis.
@@ -93,9 +109,12 @@ Expr sum(Expr x, Optional<Array<Integer>> axis, bool
keepdims);
* the cumsum over the flattened array.
* \param dtype Type of the returned array and of the accumulator in which the
elements are summed.
* If dtype is not specified, it defaults to the dtype of data.
+ * \param exclusive Whehter the first element is exclusive. If true will
return exclusive sum in
+ * which the first element is not included.
* \return The computed result.
*/
-Expr cumsum(Expr data, Optional<Integer> axis = NullOpt, DataType dtype =
DataType::Void());
+Expr cumsum(Expr data, Optional<Integer> axis = NullOpt, DataType dtype =
DataType::Void(),
+ Bool exclusive = Bool(false));
/*! \brief Computes the variance of tensor elements over given axes. */
Expr variance(Expr x, Optional<Array<Integer>> axis, bool keepdims);
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py
b/tests/python/relax/test_backend_dispatch_sort_scan.py
index c21dd4882f..8921372f2f 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -18,7 +18,7 @@
import pytest
import tvm
-from tvm import topi, relax, tir
+from tvm import topi, relax, tir, dlight
import tvm.script
import tvm.testing
from tvm.script import relax as R, tir as T, ir as I
@@ -29,21 +29,22 @@ from tvm.relax.backend import DispatchSortScan
from tvm.ir.base import assert_structural_equal
-def test_dispatch_cumsum():
+def test_dispatch_scanop():
@I.ir_module
class Before:
- I.module_global_infos({"vdevice": [I.vdevice("cuda", 0),
I.vdevice("llvm", 0)]})
+ I.module_global_infos({"vdevice": [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")
+ lv0 = R.cumsum(x, axis=1, dtype="float64", exclusive=False)
+ gv = R.cumprod(lv0, axis=1, dtype="float64", exclusive=False)
R.output(gv)
return gv
@I.ir_module
class Expected:
- I.module_global_infos({"vdevice": [I.vdevice("cuda", 0),
I.vdevice("llvm", 0)]})
+ I.module_global_infos({"vdevice": [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")):
@@ -72,13 +73,44 @@ def test_dispatch_cumsum():
],
)
+ @T.prim_func(private=True)
+ def cumprod(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)), "float64",
offset_factor=1)
+ with T.block("cumprod_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)] = 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),
+ ]
+ * 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"))
+ lv0 = R.call_tir(cls.cumsum, (x,), out_sinfo=R.Tensor((2, 3),
"float64", "llvm"))
+ gv = R.call_tir(
+ cls.cumprod,
+ (lv0,),
+ out_sinfo=R.Tensor((2, 3), dtype="float64",
vdevice="llvm"),
+ )
R.output(gv)
return gv
@@ -86,33 +118,39 @@ def test_dispatch_cumsum():
assert_structural_equal(mod, Expected)
-def test_dispatch_cumsum_cuda():
+def test_dispatch_scanop_cuda():
@I.ir_module
class Before:
- I.module_global_infos({"vdevice": [I.vdevice("cuda", 0),
I.vdevice("llvm", 0)]})
+ I.module_global_infos({"vdevice": [I.vdevice("cuda", 0)]})
@R.function
def main(x: R.Tensor(("m", 3), "float32", "cuda")):
with R.dataflow():
- lv = R.cumsum(x, axis=1)
- gv = lv
+ lv0 = R.cumsum(x, axis=1)
+ lv1 = R.cumprod(lv0, axis=1)
+ gv = lv1
R.output(gv)
return gv
target = tvm.target.Target("cuda", host="llvm")
- vdevices = [I.vdevice("cuda", 0), I.vdevice("llvm", 0)]
+ vdevices = [I.vdevice("cuda", 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(
+ lv = bb.emit_te(
topi.cuda.cumsum,
x,
axis=1,
)
+ out = bb.emit_te(
+ topi.cuda.cumprod,
+ lv,
+ axis=1,
+ )
out = bb.emit_output(out)
bb.emit_func_output(out)
expected_mod = bb.finalize()
@@ -120,14 +158,15 @@ def test_dispatch_cumsum_cuda():
with target:
mod = DispatchSortScan()(Before)
+ expected_mod =
dlight.ApplyDefaultSchedule(dlight.gpu.Fallback())(expected_mod)
- assert_structural_equal(mod, expected_mod)
+ assert_structural_equal(mod, expected_mod, map_free_vars=True)
def test_dispatch_sort():
@I.ir_module
class Before:
- I.module_global_infos({"vdevice": [I.vdevice("cuda", 0),
I.vdevice("llvm", 0)]})
+ I.module_global_infos({"vdevice": [I.vdevice("llvm", 0)]})
@R.function
def foo(x: R.Tensor(("m", 3), "float32", "llvm")):
@@ -139,7 +178,7 @@ def test_dispatch_sort():
@I.ir_module
class Expected:
- I.module_global_infos({"vdevice": [I.vdevice("cuda", 0),
I.vdevice("llvm", 0)]})
+ I.module_global_infos({"vdevice": [I.vdevice("llvm", 0)]})
@T.prim_func(private=True)
def sort(var_A: T.handle, var_sort_cpu: T.handle):
@@ -192,7 +231,7 @@ def test_dispatch_sort():
def test_dispatch_sort_cuda():
@I.ir_module
class Before:
- I.module_global_infos({"vdevice": [I.vdevice("cuda"),
I.vdevice("llvm")]})
+ I.module_global_infos({"vdevice": [I.vdevice("cuda")]})
@R.function
def foo(x: R.Tensor((2, 3), "float32", "cuda")):
@@ -212,7 +251,7 @@ def test_dispatch_sort_cuda():
target = tvm.target.Target("cuda -libs=thrust", host="llvm")
- vdevices = [I.vdevice("cuda", 0), I.vdevice("llvm", 0)]
+ vdevices = [I.vdevice("cuda", 0)]
x = relax.Var("x", R.Tensor((2, 3), "float32", vdevices[0]))
y = relax.Var("y", R.Tensor((2, 3), "float32"))
bb = relax.BlockBuilder()
@@ -244,7 +283,233 @@ def test_dispatch_sort_cuda():
with target:
mod = DispatchSortScan()(Before)
- assert_structural_equal(mod, expected_mod, map_free_vars=True)
+ assert_structural_equal(mod, expected_mod)
+
+
+def test_dispatch_argsort():
+ @I.ir_module
+ class Before:
+ I.module_global_infos({"vdevice": [I.vdevice("llvm", 0)]})
+
+ @R.function
+ def foo(x: R.Tensor(("m", 3), "float32", "llvm")):
+ m = T.int64()
+ with R.dataflow():
+ gv = R.argsort(x, axis=1, descending=False)
+ R.output(gv)
+ return gv
+
+ @I.ir_module
+ class Expected:
+ I.module_global_infos({"vdevice": [I.vdevice("llvm", 0)]})
+
+ @T.prim_func(private=True)
+ def argsort(var_A: T.handle, var_argsort_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_argsort_cpu, (m, T.int64(3)),
"int32", align=8)
+ with T.block("argsort_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.argsort",
+ 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, 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="int32", vdevice="llvm"):
+ m = T.int64()
+ cls = Expected
+ with R.dataflow():
+ gv = R.call_tir(
+ cls.argsort, (x,), out_sinfo=R.Tensor((m, 3),
dtype="int32", vdevice="llvm")
+ )
+ R.output(gv)
+ return gv
+
+ mod = DispatchSortScan()(Before)
+ assert_structural_equal(mod, Expected)
+
+
+def test_dispatch_argsort_cuda():
+ @I.ir_module
+ class Before:
+ I.module_global_infos({"vdevice": [I.vdevice("cuda")]})
+
+ @R.function
+ def foo(x: R.Tensor((2, 3), "float32", "cuda")):
+ with R.dataflow():
+ lv = R.argsort(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.argsort(y, axis=0, descending=True, dtype="int64")
+ gv = lv
+ R.output(gv)
+ return gv
+
+ target = tvm.target.Target("cuda -libs=thrust", host="llvm")
+
+ vdevices = [I.vdevice("cuda", 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.argsort, x, axis=1, is_ascend=True,
dtype="int32")
+ 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.argsort_thrust
+ if can_use_thrust(target, "tvm.contrib.thrust.sort")
+ else topi.cuda.argsort,
+ y,
+ 0,
+ False,
+ "int64",
+ )
+ 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_topk():
+ @I.ir_module
+ class Before:
+ I.module_global_infos({"vdevice": [I.vdevice("llvm", 0)]})
+
+ @R.function
+ def foo(x: R.Tensor(("m", 3), "float32", "llvm")):
+ m = T.int64()
+ with R.dataflow():
+ gv = R.topk(x, k=2, axis=1, largest=True)
+ R.output(gv)
+ return gv
+
+ @I.ir_module
+ class Expected:
+ I.module_global_infos({"vdevice": [I.vdevice("llvm", 0)]})
+
+ @T.prim_func(private=True)
+ def topk(var_A: T.handle, var_topk_cpu_v0: T.handle, var_topk_cpu_v1:
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)
+ value_buf = T.match_buffer(var_topk_cpu_v0, (m, T.int64(1)),
align=8)
+ indices_buf = T.match_buffer(var_topk_cpu_v1, (m, T.int64(1)),
"int32", align=8)
+ with T.block("topk_cpu"):
+ T.reads(data_buf[T.int64(0) : m, T.int64(0) : T.int64(3)])
+ T.writes(
+ value_buf[T.int64(0) : m, T.int64(0)],
indices_buf[T.int64(0) : m, T.int64(0)]
+ )
+ T.call_packed(
+ "tvm.contrib.sort.topk",
+ 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(
+ value_buf.data, T.tvm_stack_make_shape(m, 1), 0, 2,
T.float32(0), T.int64(0)
+ ),
+ T.tvm_stack_make_array(
+ indices_buf.data, T.tvm_stack_make_shape(m, 1), 0, 2,
0, T.int64(0)
+ ),
+ 1,
+ 1,
+ "both",
+ T.bool(False),
+ )
+
+ @R.function
+ def foo(
+ x: R.Tensor(("m", 3), dtype="float32", vdevice="llvm")
+ ) -> R.Tuple(
+ R.Tensor(("m", 1), dtype="float32", vdevice="llvm"),
+ R.Tensor(("m", 1), dtype="int32", vdevice="llvm"),
+ ):
+ m = T.int64()
+ cls = Expected
+ with R.dataflow():
+ gv = R.call_tir(
+ cls.topk,
+ (x,),
+ out_sinfo=[
+ R.Tensor((m, 1), dtype="float32", vdevice="llvm"),
+ R.Tensor((m, 1), dtype="int32", vdevice="llvm"),
+ ],
+ )
+ R.output(gv)
+ return gv
+
+ mod = DispatchSortScan()(Before)
+ assert_structural_equal(mod, Expected)
+
+
+def test_dispatch_topk_cuda():
+ @I.ir_module
+ class Before:
+ I.module_global_infos({"vdevice": [I.vdevice("cuda")]})
+
+ @R.function
+ def foo(x: R.Tensor((2, 3), "float32", "cuda")):
+ with R.dataflow():
+ lv = R.topk(x, k=2, axis=1, largest=True)
+ gv = lv
+ R.output(gv)
+ return gv
+
+ target = tvm.target.Target("cuda -libs=thrust", host="llvm")
+
+ vdevices = [I.vdevice("cuda", 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.topk, x, axis=1, is_ascend=False,
dtype="int32")
+ 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)
+ expected_mod =
dlight.ApplyDefaultSchedule(dlight.gpu.Fallback())(expected_mod)
+
+ assert_structural_equal(mod, expected_mod)
if __name__ == "__main__":
diff --git a/tests/python/relax/test_op_sort.py
b/tests/python/relax/test_op_sort.py
index b6a064a641..ed47570b82 100644
--- a/tests/python/relax/test_op_sort.py
+++ b/tests/python/relax/test_op_sort.py
@@ -26,6 +26,8 @@ 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")
+ assert relax.op.argsort(x, axis=1).op == Op.get("relax.argsort")
+ assert relax.op.topk(x, k=1, axis=1).op == Op.get("relax.topk")
def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_sinfo:
relax.StructInfo):
@@ -98,5 +100,195 @@ def test_sort_wrong_input():
bb.normalize(relax.op.sort(x1, axis=1))
+def test_argsort_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.argsort(x0, axis=1, descending=False, dtype="int64"),
+ relax.TensorStructInfo((2, 10, 4), "int64"),
+ )
+ _check_inference(
+ bb, relax.op.argsort(x6, axis=1), relax.TensorStructInfo((2, 10, 4),
"int32", vdev0)
+ )
+ _check_inference(
+ bb, relax.op.argsort(x1, axis=1),
relax.TensorStructInfo(dtype="int32", ndim=3)
+ )
+ _check_inference(
+ bb, relax.op.argsort(x2, axis=1, dtype="float16"),
relax.TensorStructInfo(dtype="float16")
+ )
+ _check_inference(
+ bb, relax.op.argsort(x3, axis=1), relax.TensorStructInfo((2, 10, 4),
dtype="int32")
+ )
+ _check_inference(
+ bb, relax.op.argsort(x4, axis=1),
relax.TensorStructInfo(dtype="int32", ndim=3)
+ )
+ _check_inference(bb, relax.op.argsort(x5, axis=1),
relax.TensorStructInfo(dtype="int32"))
+ _check_inference(bb, relax.op.argsort(x0), relax.TensorStructInfo((2, 10,
4), "int32"))
+ _check_inference(
+ bb,
+ relax.op.argsort(x0, axis=1, descending=False),
+ relax.TensorStructInfo((2, 10, 4), "int32"),
+ )
+
+
+def test_argsort_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.argsort(x, axis=1),
relax.TensorStructInfo((a, b, c), "int32"))
+ _check_inference(bb, relax.op.argsort(x), relax.TensorStructInfo((a, b,
c), "int32"))
+
+
+def test_topk_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.topk(x0, k=5, axis=1, ret_type="both", largest=False,
dtype="int64"),
+ relax.TupleStructInfo(
+ [
+ relax.TensorStructInfo((2, 5, 4), "float32"),
+ relax.TensorStructInfo((2, 5, 4), "int64"),
+ ]
+ ),
+ )
+ _check_inference(
+ bb,
+ relax.op.topk(x6),
+ relax.TupleStructInfo(
+ [
+ relax.TensorStructInfo((2, 10, 1), "float32", vdev0),
+ relax.TensorStructInfo((2, 10, 1), "int32", vdev0),
+ ]
+ ),
+ )
+ _check_inference(
+ bb,
+ relax.op.topk(x1, k=3, axis=1),
+ relax.TupleStructInfo(
+ [
+ relax.TensorStructInfo(dtype="float32", ndim=3),
+ relax.TensorStructInfo(dtype="int32", ndim=3),
+ ]
+ ),
+ )
+ _check_inference(
+ bb,
+ relax.op.topk(x2),
+ relax.TupleStructInfo(
+ [relax.TensorStructInfo(dtype="float32"),
relax.TensorStructInfo(dtype="int32")]
+ ),
+ )
+ _check_inference(
+ bb,
+ relax.op.topk(x3, axis=0),
+ relax.TupleStructInfo(
+ [
+ relax.TensorStructInfo((1, 10, 4), None),
+ relax.TensorStructInfo((1, 10, 4), dtype="int32"),
+ ]
+ ),
+ )
+ _check_inference(
+ bb,
+ relax.op.topk(x4, axis=1),
+ relax.TupleStructInfo(
+ [
+ relax.TensorStructInfo(ndim=3, dtype=None),
+ relax.TensorStructInfo(dtype="int32", ndim=3),
+ ]
+ ),
+ )
+ _check_inference(
+ bb,
+ relax.op.topk(x5, axis=1),
+ relax.TupleStructInfo(
+ [
+ relax.TensorStructInfo(dtype=None),
+ relax.TensorStructInfo(dtype="int32"),
+ ]
+ ),
+ )
+ _check_inference(
+ bb,
+ relax.op.topk(x0),
+ relax.TupleStructInfo(
+ [
+ relax.TensorStructInfo((2, 10, 1), "float32"),
+ relax.TensorStructInfo((2, 10, 1), "int32"),
+ ]
+ ),
+ )
+ _check_inference(
+ bb,
+ relax.op.topk(x0, k=-1),
+ relax.TupleStructInfo(
+ [
+ relax.TensorStructInfo((2, 10, 4), "float32"),
+ relax.TensorStructInfo((2, 10, 4), "int32"),
+ ]
+ ),
+ )
+ _check_inference(
+ bb,
+ relax.op.topk(x0, k=6),
+ relax.TupleStructInfo(
+ [
+ relax.TensorStructInfo((2, 10, 4), "float32"),
+ relax.TensorStructInfo((2, 10, 4), "int32"),
+ ]
+ ),
+ )
+
+
+def test_topk_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.topk(x, axis=1),
+ relax.TupleStructInfo(
+ [
+ relax.TensorStructInfo((a, 1, c), "float32"),
+ relax.TensorStructInfo((a, 1, c), "int32"),
+ ]
+ ),
+ )
+ _check_inference(
+ bb,
+ relax.op.topk(x, k=3),
+ relax.TupleStructInfo(
+ [
+ relax.TensorStructInfo((a, b, 3), "float32"),
+ relax.TensorStructInfo((a, b, 3), "int32"),
+ ]
+ ),
+ )
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/test_op_statistical.py
b/tests/python/relax/test_op_statistical.py
index 5c7d56556c..0f32c964f4 100644
--- a/tests/python/relax/test_op_statistical.py
+++ b/tests/python/relax/test_op_statistical.py
@@ -14,6 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+from typing import Callable
import pytest
import tvm
import tvm.testing
@@ -205,7 +206,13 @@ def test_statistical_infer_struct_info_wrong_input_type():
bb.normalize(relax.op.variance(x1))
-def test_cumsum_infer_struct_info():
+(scan_op,) = tvm.testing.parameters(
+ (relax.op.cumprod,),
+ (relax.op.cumsum,),
+)
+
+
+def test_scan_op_infer_struct_info(scan_op: Callable):
bb = relax.BlockBuilder()
vdev0 = VDevice("llvm")
x0 = relax.Var("x", R.Tensor((2, 10, 4), "float32"))
@@ -216,60 +223,56 @@ def test_cumsum_infer_struct_info():
x5 = relax.Var("x", R.Tensor())
x6 = relax.Var("x", R.Tensor((2, 10, 4), "float32", vdev0))
- _check_inference(bb, relax.op.cumsum(x0, axis=1),
relax.TensorStructInfo((2, 10, 4), "float32"))
- _check_inference(
- bb, relax.op.cumsum(x6, axis=1), relax.TensorStructInfo((2, 10, 4),
"float32", vdev0)
- )
- _check_inference(
- bb, relax.op.cumsum(x1, axis=1),
relax.TensorStructInfo(dtype="float32", ndim=3)
- )
- _check_inference(bb, relax.op.cumsum(x2, axis=1),
relax.TensorStructInfo(dtype="float32"))
- _check_inference(bb, relax.op.cumsum(x3, axis=1),
relax.TensorStructInfo((2, 10, 4), dtype=""))
- _check_inference(bb, relax.op.cumsum(x4, axis=1),
relax.TensorStructInfo(dtype="", ndim=3))
- _check_inference(bb, relax.op.cumsum(x5, axis=1),
relax.TensorStructInfo(dtype=""))
- _check_inference(bb, relax.op.cumsum(x0), relax.TensorStructInfo((80,),
"float32"))
+ _check_inference(bb, scan_op(x0, axis=1), relax.TensorStructInfo((2, 10,
4), "float32"))
+ _check_inference(bb, scan_op(x6, axis=1), relax.TensorStructInfo((2, 10,
4), "float32", vdev0))
+ _check_inference(bb, scan_op(x1, axis=1),
relax.TensorStructInfo(dtype="float32", ndim=3))
+ _check_inference(bb, scan_op(x2, axis=1),
relax.TensorStructInfo(dtype="float32"))
+ _check_inference(bb, scan_op(x3, axis=1), relax.TensorStructInfo((2, 10,
4), dtype=""))
+ _check_inference(bb, scan_op(x4, axis=1), relax.TensorStructInfo(dtype="",
ndim=3))
+ _check_inference(bb, scan_op(x5, axis=1), relax.TensorStructInfo(dtype=""))
+ _check_inference(bb, scan_op(x0), relax.TensorStructInfo((80,), "float32"))
_check_inference(
- bb, relax.op.cumsum(x0, axis=1, dtype="int32"),
relax.TensorStructInfo((2, 10, 4), "int32")
+ bb, scan_op(x0, axis=1, dtype="int32"), relax.TensorStructInfo((2, 10,
4), "int32")
)
-def test_cumsum_infer_struct_info_shape_symbolic():
+def test_scan_op_infer_struct_info_shape_symbolic(scan_op: Callable):
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.cumsum(x, axis=1),
relax.TensorStructInfo((a, b, c), "float32"))
- _check_inference(bb, relax.op.cumsum(x), relax.TensorStructInfo((a * b *
c,), "float32"))
+ _check_inference(bb, scan_op(x, axis=1), relax.TensorStructInfo((a, b, c),
"float32"))
+ _check_inference(bb, scan_op(x), relax.TensorStructInfo((a * b * c,),
"float32"))
-def test_cumsum_infer_struct_info_more_input_dtype():
+def test_scan_op_infer_struct_info_more_input_dtype(scan_op: Callable):
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.cumsum(x0, axis=1),
relax.TensorStructInfo((2, 3, 4), "float16"))
- _check_inference(bb, relax.op.cumsum(x1, axis=1),
relax.TensorStructInfo((2, 3, 4), "int8"))
+ _check_inference(bb, scan_op(x0, axis=1), relax.TensorStructInfo((2, 3,
4), "float16"))
+ _check_inference(bb, scan_op(x1, axis=1), relax.TensorStructInfo((2, 3,
4), "int8"))
-def test_cumsum_wrong_input_number():
+def test_scan_op_wrong_input_number(scan_op: Callable):
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.cumsum(x, y)
+ scan_op(x, y)
-def test_cumsum_infer_struct_info_wrong_input_type():
+def test_scan_opinfer_struct_info_wrong_input_type(scan_op: Callable):
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")))
with pytest.raises(TVMError):
- bb.normalize(relax.op.cumsum(x0, axis=1))
+ bb.normalize(scan_op(x0, axis=1))
with pytest.raises(TVMError):
- bb.normalize(relax.op.cumsum(x1, axis=1))
+ bb.normalize(scan_op(x1, axis=1))
if __name__ == "__main__":
diff --git a/tests/python/relax/test_tvmscript_parser_op_sort.py
b/tests/python/relax/test_tvmscript_parser_op_sort.py
index 8b94fa0ab9..044fba3d8d 100644
--- a/tests/python/relax/test_tvmscript_parser_op_sort.py
+++ b/tests/python/relax/test_tvmscript_parser_op_sort.py
@@ -37,15 +37,21 @@ def _check(
def test_sort():
@R.function
- def foo(x: R.Tensor((2, 3), "int32")) -> R.Tensor((2, 3), "int32"):
- r = R.sort(x, axis=1)
+ def foo(
+ x: R.Tensor((2, 3), "int32")
+ ) -> R.Tuple(R.Tensor((2, 2), dtype="int32"), R.Tensor((2, 2),
dtype="int32")):
+ lv0 = R.sort(x, axis=1)
+ lv1 = R.argsort(lv0)
+ r = R.topk(lv1, axis=1, k=2)
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)
+ lv0 = bb.emit(relax.op.sort(x, axis=1))
+ lv1 = bb.emit(relax.op.argsort(lv0))
+ r = bb.emit(relax.op.topk(lv1, axis=1, k=2))
+ bb.emit_func_output(r)
_check(foo, bb.get()["foo"])
diff --git a/tests/python/relax/test_tvmscript_parser_op_statistical.py
b/tests/python/relax/test_tvmscript_parser_op_statistical.py
index 87446cedf3..910c08bf1e 100644
--- a/tests/python/relax/test_tvmscript_parser_op_statistical.py
+++ b/tests/python/relax/test_tvmscript_parser_op_statistical.py
@@ -170,16 +170,18 @@ def test_std():
_check(foo, bb.get()["foo"])
-def test_cumsum():
+def test_scan():
@R.function
def foo(x: R.Tensor((2, 3, 4), "float32")):
- gv = R.cumsum(x, axis=1, dtype="int32")
+ lv = R.cumsum(x, axis=1, dtype="int32")
+ gv = R.cumprod(lv, axis=1, dtype="int32")
return gv
x = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
bb = relax.BlockBuilder()
with bb.function("foo", [x]):
- gv = bb.emit(relax.op.cumsum(x, axis=1, dtype="int32"))
+ lv = bb.emit(relax.op.cumsum(x, axis=1, dtype="int32"))
+ gv = bb.emit(relax.op.cumprod(lv, axis=1, dtype="int32"))
bb.emit_func_output(gv)
_check(foo, bb.get()["foo"])
