This is an automated email from the ASF dual-hosted git repository.
lunderberg pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 20d769617f [Relax] Express dynamic arguments of strided_slice as
arguments (#16826)
20d769617f is described below
commit 20d769617fa6ab561d7ed2b7cd61ed2b6b4710ba
Author: Eric Lunderberg <[email protected]>
AuthorDate: Wed May 1 09:12:43 2024 -0500
[Relax] Express dynamic arguments of strided_slice as arguments (#16826)
* [Relax] Express dynamic arguments of strided_slice as arguments
Prior to this commit, `relax.op.strided_slice` stored the `axes`,
`begin`, `end`, and `strides` in the `CallNode::attrs`. However, the
attributes are only intended to store static values. The indices used
used for `relax.op.strided_slice` must frequently be in terms of
symbolic shape variables, which should not be stored in the
attributes. While some utilities have special handling for
`relax.op.strided_slice` (e.g. `tvm::relax::Bind`), many do
not (e.g. `tvm::relax::WellFormed` and
`tvm::relax::FreeSymbolicVars`). As a result, the symbolic
expressions in `relax.op.strided_slice` will fail to be updated in
generic utilities, and will fail to trigger safeguards when this
occurs.
This commit changes the representation of `relax.op.strided_slice` to
store all arguments in the `relax::CallNode::args`, rather than the
`relax::CallNode::attrs`. As mentioned in a comment from
https://github.com/apache/tvm/pull/13987, which initially implemented
`relax.op.strided_slice`, this was an intended refactor once
`relax::PrimValue` was fully supported.
* Undo unnecessary changes in const_int_bound
* Remove unnecessary changes to rewrite_simplify
* lint fixes
* Fix unit tests
* Improve error message
* Fix additional unit tests
* Mark MSC tests with xfail
* remove commented-out code
* Resolve failing unit test
* Remove unused imports
---
include/tvm/relax/attrs/index.h | 11 -
python/tvm/relax/__init__.py | 6 +
python/tvm/relax/op/index.py | 12 +-
python/tvm/relax/transform/legalize_ops/index.py | 39 +-
python/tvm/relax/type_converter.py | 179 +++++++++
python/tvm/relax/utils.py | 150 +-------
.../msc/framework/tensorrt/transform_tensorrt.cc | 10 +-
src/relax/analysis/struct_info_analysis.cc | 34 +-
src/relax/op/tensor/index.cc | 403 ++++++++++++++++-----
src/relax/op/tensor/index.h | 6 +-
src/relax/transform/convert_layout.cc | 19 +-
src/relax/transform/infer_layout_utils.h | 4 +-
src/relax/utils.cc | 43 ---
src/script/ir_builder/relax/ir.cc | 10 +
tests/python/contrib/test_msc/test_graph_build.py | 3 +
.../contrib/test_msc/test_translate_relax.py | 3 +
.../contrib/test_msc/test_translate_tensorflow.py | 4 +
.../contrib/test_msc/test_translate_torch.py | 3 +
tests/python/relax/test_dataflow_pattern.py | 28 +-
tests/python/relax/test_op_index.py | 43 ++-
20 files changed, 653 insertions(+), 357 deletions(-)
diff --git a/include/tvm/relax/attrs/index.h b/include/tvm/relax/attrs/index.h
index 1043fe30ce..aa6c2e1461 100644
--- a/include/tvm/relax/attrs/index.h
+++ b/include/tvm/relax/attrs/index.h
@@ -40,20 +40,9 @@ struct TakeAttrs : public tvm::AttrsNode<TakeAttrs> {
/*! \brief Attributes used in strided_slice operator */
struct StridedSliceAttrs : public tvm::AttrsNode<StridedSliceAttrs> {
- Array<Integer> axes;
- Array<PrimExpr> begin;
- Array<PrimExpr> end;
- Optional<Array<PrimExpr>> strides;
bool assume_inbound;
TVM_DECLARE_ATTRS(StridedSliceAttrs, "relax.attrs.StridedSliceAttrs") {
- TVM_ATTR_FIELD(axes).describe("Axes along which slicing is applied.");
- TVM_ATTR_FIELD(begin).describe("The indices to begin with in the slicing,
inclusive.");
- TVM_ATTR_FIELD(end).describe("The indices indicating end of the slice,
exclusive.");
- TVM_ATTR_FIELD(strides).describe(
- "Specifies the stride values, it can be negative in that case, the
input tensor will be "
- "reversed in that particular axis. If not specified, it by default is
an list of ones of "
- "the same length as `axes`.");
TVM_ATTR_FIELD(assume_inbound)
.set_default(true)
.describe(
diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py
index 23cfaf2935..dd3245441b 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -19,6 +19,8 @@
from tvm.runtime import relax_vm as vm
from tvm.runtime.relax_vm import VirtualMachine, VMInstrumentReturnKind
+from .type_converter import args_converter
+
# Expr
from .expr import (
Expr,
@@ -92,6 +94,9 @@ from .struct_info import (
from .pipeline import get_pipeline
from .pipeline import register_pipeline
+# utils
+from .utils import convert_to_expr
+
# Import submodules in the last to avoid dependency
from . import exec_builder
from . import expr
@@ -105,6 +110,7 @@ from . import backend
from . import training
from . import distributed
from . import frontend
+from . import utils
# VM
from .vm_build import build, Executable
diff --git a/python/tvm/relax/op/index.py b/python/tvm/relax/op/index.py
index 8504b4d683..ec68bd585c 100644
--- a/python/tvm/relax/op/index.py
+++ b/python/tvm/relax/op/index.py
@@ -15,12 +15,13 @@
# specific language governing permissions and limitations
# under the License.
"""Indexing operators."""
-from typing import List, Optional, Union
+from typing import Optional, Union
from tvm.ir.expr import PrimExpr
from . import _ffi_api
from ..expr import Expr
+from .. import args_converter
PrimExprLike = Union[int, PrimExpr]
@@ -52,12 +53,13 @@ def take(x: Expr, indices: Expr, axis: Optional[int] =
None) -> Expr:
return _ffi_api.take(x, indices, axis) # type: ignore
+@args_converter.auto
def strided_slice(
x: Expr,
- axes: List[int],
- begin: List[PrimExprLike],
- end: List[PrimExprLike],
- strides: Optional[List[PrimExprLike]] = None,
+ axes: Expr,
+ begin: Expr,
+ end: Expr,
+ strides: Optional[Expr] = None,
assume_inbound: bool = False,
) -> Expr:
"""Strided slice of a tensor.
diff --git a/python/tvm/relax/transform/legalize_ops/index.py
b/python/tvm/relax/transform/legalize_ops/index.py
index 5889da9487..a4fac46a13 100644
--- a/python/tvm/relax/transform/legalize_ops/index.py
+++ b/python/tvm/relax/transform/legalize_ops/index.py
@@ -20,7 +20,7 @@ from tvm import topi, tir, te
from ...op import call_pure_packed
from ...block_builder import BlockBuilder
from ...expr import Call, Expr
-from ...struct_info import ShapeStructInfo
+from ...struct_info import ShapeStructInfo, PrimStructInfo
from .common import register_legalize
@@ -35,18 +35,37 @@ def _take(bb: BlockBuilder, call: Call) -> Expr:
@register_legalize("relax.strided_slice")
def _strided_slice(bb: BlockBuilder, call: Call) -> Expr:
- strides = (
- [tir.IntImm("int64", 1)] * len(call.attrs.axes)
- if call.attrs.strides is None
- else call.attrs.strides
- )
+ def _relax_tuple_to_tir(relax_tuple):
+ output = []
+ for field in relax_tuple.struct_info.fields:
+ assert isinstance(field, PrimStructInfo)
+ assert field.value is not None
+ output.append(field.value)
+ return output
+
+ if len(call.args) == 4:
+ data, axes, begin, end = call.args
+ strides = [tir.IntImm("int64", 1)] * len(axes.struct_info.fields)
+ elif len(call.args) == 5:
+ data, axes, begin, end, strides = call.args
+ strides = _relax_tuple_to_tir(strides)
+ else:
+ raise ValueError(
+ f"Expression {call} provides {len(call.args)} arguments, "
+ f"but {call.op} requires either 4 or 5 arguments."
+ )
+
+ axes = _relax_tuple_to_tir(axes)
+ begin = _relax_tuple_to_tir(begin)
+ end = _relax_tuple_to_tir(end)
+
return bb.call_te(
topi.strided_slice,
- call.args[0],
- call.attrs.begin,
- call.attrs.end,
+ data,
+ begin,
+ end,
strides,
- call.attrs.axes,
+ axes,
slice_mode="end",
)
diff --git a/python/tvm/relax/type_converter.py
b/python/tvm/relax/type_converter.py
new file mode 100644
index 0000000000..b29555f687
--- /dev/null
+++ b/python/tvm/relax/type_converter.py
@@ -0,0 +1,179 @@
+# 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,too-many-locals
+
+"""Argument converter utility for Relax
+
+This utility is used to decorate constructors of `tvm.relax.Expr`, and
+must be able to be imported before `tvm.relax.Expr` or its subtypes
+have been defined. Neither the class definitions nor any type
+signature in this file may reference relax types. All references must
+be exclusively in function bodies to avoid having a circular reference
+during module imports.
+"""
+
+import functools
+import inspect
+
+from typing import List, Optional, Callable, TypeVar, Any
+
+import tvm
+
+FType = TypeVar("FType", bound=Callable[..., "tvm.relax.Expr"])
+
+
+class _ArgsConverter:
+ """A helper class to convert the arguments to Expr."""
+
+ @staticmethod
+ def convert(args_to_expr: List[str], args_to_list_expr: List[str]):
+ """Convert the arguments to Expr.
+
+ Parameters
+ ----------
+ args_to_expr : List[str]
+ The argument names to be converted to Expr.
+
+ args_to_list_expr : List[str]
+ The argument names to be converted to List[Expr].
+
+ Returns
+ -------
+ output : Callable[[FType], FType]
+ The decorator.
+ """
+
+ if any([x in args_to_list_expr for x in args_to_expr]):
+ raise ValueError("`args_to_expr` and `args_to_list_expr` should be
disjoint.")
+
+ def _convert(name: str, value: Any) -> Any:
+ if value is None:
+ return value
+ if name in args_to_expr:
+ try:
+ return tvm.relax.utils.convert_to_expr(value)
+ except Exception as err:
+ raise TypeError(
+ f"Argument `{name}` is expected to be converted to
`Expr`, "
+ f"but failed with input value: {value}"
+ ) from err
+ elif name in args_to_list_expr:
+ try:
+ return [tvm.relax.utils.convert_to_expr(x) for x in value]
+ except Exception as err:
+ raise TypeError(
+ f"Argument `{name}` is expected to be converted to
`List[Expr]`, "
+ f"but failed with input value: {value}"
+ ) from err
+ else:
+ return value
+
+ def inner(func: FType) -> FType:
+ sig = inspect.signature(func)
+ param_names = list(sig.parameters.keys())
+ for name in args_to_expr + args_to_list_expr:
+ if name not in param_names:
+ raise ValueError(f"Argument `{name}` is not found in
function signature.")
+
+ @functools.wraps(func)
+ def wrapper(*args, **kwargs):
+ bound = sig.bind(*args, **kwargs)
+ bound.apply_defaults()
+ for param in sig.parameters.values():
+ if param.kind == param.VAR_POSITIONAL:
+ # *args case
+ values = [_convert(param.name, x) for x in
bound.arguments[param.name]]
+ bound.arguments[param.name] = tuple(values)
+ elif param.kind == param.VAR_KEYWORD:
+ # **kwargs case
+ key_value = {
+ key: _convert(param.name, value)
+ for key, value in
bound.arguments[param.name].items()
+ }
+ bound.arguments[param.name] = key_value
+ else:
+ bound.arguments[param.name] = _convert(
+ param.name, bound.arguments[param.name]
+ )
+ return func(*bound.args, **bound.kwargs)
+
+ return wrapper # type: ignore
+
+ return inner
+
+ @staticmethod
+ def to_expr(*arg_names: str) -> Callable:
+ """Convert the arguments to Expr.
+
+ Parameters
+ ----------
+ *arg_names: str
+ The list of argument names that need to be converted to Expr.
+
+ Returns
+ -------
+ output: Callable
+ The decorator.
+ """
+
+ return _ArgsConverter.convert(args_to_expr=list(arg_names),
args_to_list_expr=[])
+
+ @staticmethod
+ def to_list_expr(*arg_names: str) -> Callable:
+ """Convert the arguments to List of Expr.
+
+ Parameters
+ ----------
+ *arg_names: str
+ The list of argument names that need to be converted to List of
Expr.
+
+ Returns
+ -------
+ output: Callable
+ The decorator.
+ """
+
+ return _ArgsConverter.convert(args_to_expr=[],
args_to_list_expr=list(arg_names))
+
+ @staticmethod
+ def auto(func: FType) -> FType:
+ """Decorator for automatically convert the arguments to Expr according
to type annotation.
+ Only two patterns are supported:
+
+ 1. The argument is Expr or Optional[Expr].
+
+ 2. The argument is List[Expr] or Optional[List[Expr]].
+
+ """
+ sig = inspect.signature(func)
+ args_to_expr = []
+ args_to_list_expr = []
+
+ from . import Expr # pylint: disable=import-outside-toplevel
+
+ for param in sig.parameters.values():
+ anno = param.annotation
+ if anno in (Expr, Optional[Expr]):
+ args_to_expr.append(param.name)
+ if anno in (List[Expr], Optional[List[Expr]]):
+ args_to_list_expr.append(param.name)
+
+ return _ArgsConverter.convert(args_to_expr, args_to_list_expr)(func)
+
+
+args_converter = _ArgsConverter() # pylint: disable=invalid-name
diff --git a/python/tvm/relax/utils.py b/python/tvm/relax/utils.py
index 48beeed8da..9323bc40da 100644
--- a/python/tvm/relax/utils.py
+++ b/python/tvm/relax/utils.py
@@ -19,13 +19,11 @@
"""Utility functions for Relax"""
-import functools
-import inspect
import itertools
import string
from typing import Tuple as typing_Tuple
-from typing import Any, Callable, List, Dict, Optional, TypeVar
+from typing import Any, Callable, List, Dict, Optional
import tvm
from .. import tir
@@ -38,6 +36,9 @@ from ..te import Tensor as te_Tensor, create_prim_func
from ..ir import Array, Attrs, Type, Map, VDevice
from .struct_info import PrimStructInfo, ShapeStructInfo, TensorStructInfo
+# Re-export `args_converter` here for backwards compatibility
+from .type_converter import args_converter # pylint: disable=unused-import
+
def metadata_partitioner(rx_txt: str) -> List[str]:
"""Extract Relax program and metadata section.
@@ -119,149 +120,6 @@ def convert_to_expr(value: Any) -> Expr:
raise TypeError(f"Cannot convert {value} with type {type(value)} to
`relax.Expr`")
-FType = TypeVar("FType", bound=Callable[..., Expr])
-
-
-class _ArgsConverter:
- """A helper class to convert the arguments to Expr."""
-
- @staticmethod
- def convert(args_to_expr: List[str], args_to_list_expr: List[str]):
- """Convert the arguments to Expr.
-
- Parameters
- ----------
- args_to_expr : List[str]
- The argument names to be converted to Expr.
-
- args_to_list_expr : List[str]
- The argument names to be converted to List[Expr].
-
- Returns
- -------
- output : Callable[[FType], FType]
- The decorator.
- """
-
- if any([x in args_to_list_expr for x in args_to_expr]):
- raise ValueError("`args_to_expr` and `args_to_list_expr` should be
disjoint.")
-
- def _convert(name: str, value: Any) -> Any:
- if value is None:
- return value
- if name in args_to_expr:
- try:
- return convert_to_expr(value)
- except:
- raise TypeError(
- f"Argument `{name}` is expected to be converted to
`Expr`, "
- f"but failed with input value: {value}"
- )
- elif name in args_to_list_expr:
- try:
- return [convert_to_expr(x) for x in value]
- except:
- raise TypeError(
- f"Argument `{name}` is expected to be converted to
`List[Expr]`, "
- f"but failed with input value: {value}"
- )
- else:
- return value
-
- def inner(func: FType) -> FType:
- sig = inspect.signature(func)
- param_names = list(sig.parameters.keys())
- for name in args_to_expr + args_to_list_expr:
- if name not in param_names:
- raise ValueError(f"Argument `{name}` is not found in
function signature.")
-
- @functools.wraps(func)
- def wrapper(*args, **kwargs):
- bound = sig.bind(*args, **kwargs)
- bound.apply_defaults()
- for param in sig.parameters.values():
- if param.kind == param.VAR_POSITIONAL:
- # *args case
- values = [_convert(param.name, x) for x in
bound.arguments[param.name]]
- bound.arguments[param.name] = tuple(values)
- elif param.kind == param.VAR_KEYWORD:
- # **kwargs case
- key_value = {
- key: _convert(param.name, value)
- for key, value in
bound.arguments[param.name].items()
- }
- bound.arguments[param.name] = key_value
- else:
- bound.arguments[param.name] = _convert(
- param.name, bound.arguments[param.name]
- )
- return func(*bound.args, **bound.kwargs)
-
- return wrapper # type: ignore
-
- return inner
-
- @staticmethod
- def to_expr(*arg_names: str) -> Callable:
- """Convert the arguments to Expr.
-
- Parameters
- ----------
- *arg_names: str
- The list of argument names that need to be converted to Expr.
-
- Returns
- -------
- output: Callable
- The decorator.
- """
-
- return _ArgsConverter.convert(args_to_expr=list(arg_names),
args_to_list_expr=[])
-
- @staticmethod
- def to_list_expr(*arg_names: str) -> Callable:
- """Convert the arguments to List of Expr.
-
- Parameters
- ----------
- *arg_names: str
- The list of argument names that need to be converted to List of
Expr.
-
- Returns
- -------
- output: Callable
- The decorator.
- """
-
- return _ArgsConverter.convert(args_to_expr=[],
args_to_list_expr=list(arg_names))
-
- @staticmethod
- def auto(func: FType) -> FType:
- """Decorator for automatically convert the arguments to Expr according
to type annotation.
- Only two patterns are supported:
-
- 1. The argument is Expr or Optional[Expr].
-
- 2. The argument is List[Expr] or Optional[List[Expr]].
-
- """
- sig = inspect.signature(func)
- args_to_expr = []
- args_to_list_expr = []
-
- for param in sig.parameters.values():
- anno = param.annotation
- if anno in (Expr, Optional[Expr]):
- args_to_expr.append(param.name)
- if anno in (List[Expr], Optional[List[Expr]]):
- args_to_list_expr.append(param.name)
-
- return _ArgsConverter.convert(args_to_expr, args_to_list_expr)(func)
-
-
-args_converter = _ArgsConverter() # pylint: disable=invalid-name
-
-
def copy_with_new_vars(func: Function) -> Function:
"""Copy the given function. All variables that are bound inside the
original function
would be copied to satisfy the restriction in the well-formed check:
Variables in
diff --git a/src/contrib/msc/framework/tensorrt/transform_tensorrt.cc
b/src/contrib/msc/framework/tensorrt/transform_tensorrt.cc
index c71cb60501..3f85309cd8 100644
--- a/src/contrib/msc/framework/tensorrt/transform_tensorrt.cc
+++ b/src/contrib/msc/framework/tensorrt/transform_tensorrt.cc
@@ -644,15 +644,11 @@ Expr RewriteSplit(BlockBuilder builder, const Var& var,
const Call& src_call,
<< src_attrs->indices_or_sections->GetTypeKey() << ")";
}
// create strided_slices
- static const Op& slice_op = Op::Get("relax.strided_slice");
Array<Expr> outputs;
for (size_t i = 0; i < split_begins.size(); i++) {
- auto slice_attrs = make_object<StridedSliceAttrs>();
- slice_attrs->axes.push_back(Integer(axis));
- slice_attrs->begin.push_back(Integer(split_begins[i]));
- slice_attrs->end.push_back(Integer(split_ends[i]));
- const auto& slice = MakeCall(builder, call->span, "slice_" +
std::to_string(i), slice_op,
- {call->args[0]}, Attrs(slice_attrs));
+ auto slice = strided_slice(call->args[0],
Tuple(Array<Expr>{PrimValue(Integer(axis))}),
+
Tuple(Array<Expr>{PrimValue(Integer(split_begins[i]))}),
+
Tuple(Array<Expr>{PrimValue(Integer(split_ends[i]))}));
outputs.push_back(slice);
}
return Tuple(outputs, call->span);
diff --git a/src/relax/analysis/struct_info_analysis.cc
b/src/relax/analysis/struct_info_analysis.cc
index 08e2acfbd0..0432c96e2e 100644
--- a/src/relax/analysis/struct_info_analysis.cc
+++ b/src/relax/analysis/struct_info_analysis.cc
@@ -1163,19 +1163,29 @@ class TIRVarsDetector : public StructInfoVisitor {
Array<tir::Var> GetTIRVars() const { return tir_vars_; }
private:
- void VisitShape(Array<PrimExpr> shape) {
- for (const PrimExpr& value : shape) {
- if (collection_type == VarType::Definition) {
- if (auto opt = value.as<tir::Var>()) {
- RecordTIRVar(opt.value());
- }
- } else if (collection_type == VarType::Usage) {
- for (const tir::Var& tir_var : tir::UndefinedVars(value)) {
- RecordTIRVar(tir_var);
- }
- } else {
- LOG(FATAL) << "Invalid value for VarType enum, " <<
static_cast<int>(collection_type);
+ void VisitPrimExpr(PrimExpr expr) {
+ if (collection_type == VarType::Definition) {
+ if (auto opt = expr.as<tir::Var>()) {
+ RecordTIRVar(opt.value());
}
+ } else if (collection_type == VarType::Usage) {
+ for (const tir::Var& tir_var : tir::UndefinedVars(expr)) {
+ RecordTIRVar(tir_var);
+ }
+ } else {
+ LOG(FATAL) << "Invalid value for VarType enum, " <<
static_cast<int>(collection_type);
+ }
+ }
+
+ void VisitShape(Array<PrimExpr> shape) {
+ for (const PrimExpr& expr : shape) {
+ VisitPrimExpr(expr);
+ }
+ }
+
+ void VisitStructInfo_(const PrimStructInfoNode* prim_sinfo) final {
+ if (prim_sinfo->value.defined()) {
+ VisitPrimExpr(prim_sinfo->value.value());
}
}
diff --git a/src/relax/op/tensor/index.cc b/src/relax/op/tensor/index.cc
index d052c2a64f..022ef31c66 100644
--- a/src/relax/op/tensor/index.cc
+++ b/src/relax/op/tensor/index.cc
@@ -24,6 +24,11 @@
#include "index.h"
+#include <tvm/relax/analysis.h>
+
+#include <algorithm>
+#include <optional>
+#include <tuple>
#include <utility>
#include <vector>
@@ -122,117 +127,323 @@ TVM_REGISTER_OP("relax.take")
/* relax.strided_slice */
TVM_REGISTER_NODE_TYPE(StridedSliceAttrs);
-Expr strided_slice(Expr x, //
- Array<Integer> axes, //
- Array<PrimExpr> begin, //
- Array<PrimExpr> end, //
- Optional<Array<PrimExpr>> strides, //
+Expr strided_slice(Expr x, Expr axes, Expr begin, Expr end, Optional<Expr>
strides,
bool assume_inbound) {
- int n_axis = axes.size();
- CHECK_EQ(static_cast<int>(begin.size()), n_axis)
- << "StridedSlice requires the number of begin indices to equal the
number of axes.";
- CHECK_EQ(static_cast<int>(end.size()), n_axis)
- << "StridedSlice requires the number of end indices to equal the number
of axes.";
- if (strides.defined()) {
- CHECK_EQ(static_cast<int>(strides.value().size()), n_axis)
- << "StridedSlice requires the number of strides to equal the number of
axes.";
- }
-
- // Todo(relax-team): We are going to support dynamic strided slice, where
- // begin/end/stride can be not static at compile time. Therefore,
begin/end/stride
- // should not be part of StridedSliceAttrs, as we only allow static values to
- // reside in attributes. However, using ShapeExpr to represent these
- // arrays is not conceptually right, because they are not describing a
- // concrete shape. The proper way to support dynamic strided slice is to use
- // Tuple of PrimValue to represent begin/end/stride. Since at this moment
- // we have no support for PrimValue, we store begin/end/stride as attribute
- // fields as a workaround.
- // Will switch to Tuple of PrimValue after introducing PrimValue.
- auto f_convert_to_int64 = [](const PrimExpr& value) {
- if (value->IsInstance<IntImmNode>()) {
- return cast(DataType::Int(64), value);
+ // Initial validation of the arguments. A more complete validation
+ // will be done when inferring the StructInfo, but that requires the
+ // StructInfo of all arguments to be populated.
+
+ std::optional<std::tuple<const char*, size_t, Expr>> known_length;
+ auto check_tuple = [&known_length](const char* name, Expr expr) {
+ if (const auto* tuple = expr.as<TupleNode>()) {
+ size_t length = tuple->fields.size();
+ if (known_length.has_value()) {
+ const auto& prev = known_length.value();
+ CHECK_EQ(length, std::get<size_t>(prev))
+ << "The strided_slice operator requires that "
+ << "the axes, begin, end, and strides tuples are all the same
length. "
+ << "However, the " << std::get<const char*>(prev) << " argument ("
+ << std::get<Expr>(prev) << ") has " << std::get<size_t>(prev) << "
elements, while the "
+ << name << " argument (" << expr << ") has " << length << "
elements.";
+ } else {
+ known_length = std::tuple{name, length, expr};
+ }
}
- CHECK(value.dtype() == DataType::Int(64)) << "strided_slice expects the
input begin/end/stride "
- "values to be all int64.
However, the given "
- << value << " has dtype " <<
value->dtype;
- return value;
};
+ check_tuple("axes", axes);
+ check_tuple("begin", begin);
+ check_tuple("end", end);
+ if (strides.defined()) check_tuple("strides", strides.value());
ObjectPtr<StridedSliceAttrs> attrs = make_object<StridedSliceAttrs>();
- attrs->axes = std::move(axes);
- attrs->begin = begin.Map(f_convert_to_int64);
- attrs->end = end.Map(f_convert_to_int64);
- attrs->strides = strides.defined() ? strides.value().Map(f_convert_to_int64)
: strides;
attrs->assume_inbound = assume_inbound;
+ Array<Expr> args = {x, axes, begin, end};
+ if (strides.defined()) {
+ args.push_back(strides.value());
+ }
+
static const Op& op = Op::Get("relax.strided_slice");
- return Call(op, {std::move(x)}, Attrs(attrs), {});
+ auto call = Call(op, args, Attrs(attrs));
+
+ return call;
}
TVM_REGISTER_GLOBAL("relax.op.strided_slice").set_body_typed(strided_slice);
-inline PrimExpr CanonicalizeIndex(PrimExpr index, PrimExpr extent, int64_t
stride,
- bool assume_inbound) {
- // Same as topi strided slice CanonicalizeIndex function in
- // include/tvm/topi/detail/strided_slice.h
- PrimExpr begin_range = stride < 0 ? -1 : 0;
- PrimExpr end_range = stride < 0 ? extent - 1 : extent;
+inline PrimExpr CanonicalizeIndex(PrimExpr index, PrimExpr extent, PrimExpr
stride) {
+ // Handle Python-style negative indices
index = if_then_else(index < 0, index + extent, index);
- return assume_inbound ? index : min(max(index, begin_range), end_range); //
NOLINT
+ // Clamp the result to valid indices
+ PrimExpr lower_bound = tvm::if_then_else(stride < 0, -1, 0);
+ PrimExpr upper_bound = tvm::if_then_else(stride < 0, extent - 1, extent);
+ index = tvm::min(tvm::max(index, lower_bound), upper_bound);
+
+ return index;
}
-PrimExpr GetLength(PrimExpr begin, PrimExpr end, const int64_t stride, const
PrimExpr& length,
+PrimExpr GetLength(PrimExpr begin, PrimExpr end, PrimExpr stride, PrimExpr
extent,
bool assume_inbound) {
- begin = CanonicalizeIndex(begin, length, stride, assume_inbound);
- end = CanonicalizeIndex(end, length, stride, assume_inbound);
- arith::Analyzer ana;
- if (stride < 0) {
- return ana.Simplify(ceildiv(begin - end, IntImm(DataType::Int(64),
-stride)));
+ if (assume_inbound) {
+ return ceildiv(end - begin, stride);
} else {
- return ana.Simplify(ceildiv(end - begin, IntImm(DataType::Int(64),
stride)));
+ begin = CanonicalizeIndex(begin, extent, stride);
+ end = CanonicalizeIndex(end, extent, stride);
+ return tvm::if_then_else(stride < 0, ceildiv(begin - end, -stride),
+ ceildiv(end - begin, stride));
}
}
-StructInfo InferStructInfoStridedSlice(const Call& call, const BlockBuilder&
ctx) {
- TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
- const auto* attrs = call->attrs.as<StridedSliceAttrs>();
- if (attrs->axes.empty()) {
- return data_sinfo;
- }
+/* \brief Helper function to unpack a relax::Tuple
+ *
+ * A `relax::Tuple` may be provided to an operator as an in-line
+ * expression, as a variable bound to known tuple within the current
+ * function, as a function argument, etc. The StructInfo of the tuple
+ * tracks the known values of any `PrimValue` elements, but it can be
+ * tedious to extract. This utility extracts the `PrimExpr` contents
+ * of a `relax::Tuple`.
+ *
+ * If the StructInfo cannot contain a tuple of the type specified,
+ * this function will throw an exception. (e.g. Attempting to extract
+ * a tuple from a `TensorStructInfo`.)
+ *
+ * \tparam PrimType The subtype of PrimExpr to extract. For example,
+ * extracting an `Array<Integer>`
+ *
+ * \param sinfo The StructInfo to inspect
+ *
+ * \returns An array of the `PrimType`, if it can be extracted.
+ * Otherwise, `NullOpt`.
+ */
+template <typename PrimType = PrimExpr,
+ typename = std::enable_if_t<std::is_base_of_v<PrimExpr, PrimType>>>
+Optional<Array<PrimType>> UnpackTupleOfPrimValue(Optional<StructInfo> sinfo) {
+ if (!sinfo) return NullOpt;
- if (data_sinfo->IsUnknownNdim()) {
- return TensorStructInfo(data_sinfo->dtype, kUnknownNDim,
data_sinfo->vdevice);
+ // An ObjectStructInfo may contain a tuple of the desired type, but
+ // it isn't yet known whether it does. Return early, as we cannot
+ // provide a known `Array<PrimType>` to the caller.
+ if (sinfo.as<ObjectStructInfoNode>()) return NullOpt;
+
+ auto tuple = sinfo.as<TupleStructInfoNode>();
+ CHECK(tuple) << "TypeError: "
+ << "The struct info " << sinfo << " cannot contain a tuple
whose elements are "
+ << PrimType::ContainerType::_type_key;
+
+ Array<PrimType> output;
+ for (size_t i = 0; i < tuple->fields.size(); i++) {
+ auto field = tuple->fields[i];
+
+ if (field.as<ObjectStructInfoNode>()) return NullOpt;
+
+ auto prim_sinfo = field.as<PrimStructInfoNode>();
+ CHECK(prim_sinfo) << "TypeError: "
+ << "The struct info " << sinfo
+ << " cannot contain a tuple whose elements are "
+ << PrimType::ContainerType::_type_key << ", because
element " << i
+ << " has struct info " << field;
+
+ if (!prim_sinfo->value.defined()) return NullOpt;
+
+ Optional<PrimType> element = prim_sinfo->value.as<PrimType>();
+ if (!element) return NullOpt;
+
+ output.push_back(element.value());
}
+ return output;
+}
- std::vector<int> axes = NormalizeAxes(call, ctx, data_sinfo->ndim,
attrs->axes);
- const auto* data_shape = data_sinfo->shape.as<ShapeExprNode>();
- if (data_shape == nullptr) {
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice);
+/* \brief Helper function to unpack a relax::Tuple
+ *
+ * A `relax::Tuple` may be provided to an operator as an in-line
+ * expression, as a variable bound to known tuple within the current
+ * function, as a function argument, etc. The StructInfo of the tuple
+ * tracks the known values of any `PrimValue` elements, but it can be
+ * tedious to extract. This utility extracts the `PrimExpr` contents
+ * of a `relax::Tuple`.
+ *
+ * If the StructInfo cannot contain a tuple of the type specified,
+ * this function will throw an exception. (e.g. Attempting to extract
+ * a tuple from a `TensorStructInfo`.)
+ *
+ * \tparam PrimType The subtype of PrimExpr to extract. For example,
+ * extracting an `Array<Integer>`
+ *
+ * \param expr The `relax::Expr` to inspect
+ *
+ * \returns An array of the `PrimType`, if it can be extracted.
+ * Otherwise, `NullOpt`.
+ */
+template <typename PrimType = PrimExpr,
+ typename = std::enable_if_t<std::is_base_of_v<PrimExpr, PrimType>>>
+Optional<Array<PrimType>> UnpackTupleOfPrimValue(Optional<Expr> expr) {
+ if (expr) {
+ return UnpackTupleOfPrimValue<PrimType>(GetStructInfo(expr.value()));
+ } else {
+ return NullOpt;
}
+}
- int n_axis = axes.size();
- Array<PrimExpr> strides = attrs->strides.defined()
- ? attrs->strides.value()
- : Array<PrimExpr>(n_axis,
IntImm(DataType::Int(64), 1));
- std::vector<int64_t> int_strides;
- int_strides.reserve(n_axis);
- // Only do output shape inference when all the begin/end/strides values are
integers.
- for (int i = 0; i < n_axis; ++i) {
- const auto* int_stride = strides[i].as<IntImmNode>();
- if (!int_stride) {
- return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim,
data_sinfo->vdevice);
+StructInfo InferStructInfoStridedSlice(const Call& call, const BlockBuilder&
ctx) {
+ size_t n_args = call->args.size();
+ CHECK(4 <= n_args && n_args <= 5)
+ << "Operator " << call->op << " accepts either three arguments (data,
axes, begin, end) "
+ << " or four arguments (data, axes, begin, end, strides), "
+ << "but received " << n_args << " in expression " << call;
+
+ Expr data = call->args[0];
+ Expr axes = call->args[1];
+ Expr begin = call->args[2];
+ Expr end = call->args[3];
+ Optional<Expr> strides = [&]() -> Optional<Expr> {
+ if (n_args > 4) {
+ return call->args[4];
+ } else {
+ return NullOpt;
}
- int_strides.push_back(int_stride->value);
+ }();
+
+ auto axes_sinfo = GetStructInfo(call->args[1]);
+ auto begin_sinfo = GetStructInfo(call->args[2]);
+ auto end_sinfo = GetStructInfo(call->args[3]);
+ auto strides_sinfo = [&]() -> Optional<StructInfo> {
+ if (n_args > 4) {
+ return GetStructInfo(call->args[4]);
+ } else {
+ return NullOpt;
+ }
+ }();
+
+ CHECK(IsBaseOf(relax::TensorStructInfo(DataType::Void(), kUnknownNDim),
GetStructInfo(data)))
+ << "Operator " << call->op << " requires the first argument to be a
tensor. "
+ << "However, in expression " << call << ", the first argument " << data
<< " has struct info "
+ << GetStructInfo(data);
+
+ // TODO(Lunderberg): Implement this check using `IsBaseOf`. Doing
+ // so will require a way to represent a `relax::TupleStructInfo` of
+ // unknown length, where each element has the same `StructInfo`.
+ auto is_base_of_tuple_of_int64 = [&](const StructInfo& sinfo) -> bool {
+ if (sinfo.as<ObjectStructInfoNode>()) {
+ return true;
+ }
+
+ const auto* tuple = sinfo.as<TupleStructInfoNode>();
+ if (!tuple) return false;
+
+ return std::all_of(tuple->fields.begin(), tuple->fields.end(), [](const
StructInfo& field) {
+ return IsBaseOf(relax::PrimStructInfo(DataType::Int(64)), field);
+ });
+ };
+ auto check_tuple = [&](const char* name, Expr expr) {
+ auto sinfo = GetStructInfo(expr);
+
+ CHECK(is_base_of_tuple_of_int64(sinfo)) << "Operator " << call->op << "
requires the " << name
+ << " argument to be a tuple of
int64 PrimValues. "
+ << "However, in expression " <<
call << ", the " << name
+ << " argument " << expr << " has
struct info " << sinfo;
+ };
+ check_tuple("axes", call->args[1]);
+ check_tuple("begin", call->args[2]);
+ check_tuple("end", call->args[3]);
+ if (call->args.size() > 4) {
+ check_tuple("strides", call->args[4]);
}
- Array<PrimExpr> output_shape = data_shape->values;
- for (int i = 0; i < n_axis; ++i) {
- ICHECK_NE(int_strides[i], 0)
- << "Strided slice requires strides to be non-zero but got 0 for axis "
<< axes[i] << ".";
- output_shape.Set(axes[i], GetLength(attrs->begin[i], attrs->end[i],
int_strides[i],
- data_shape->values[axes[i]],
attrs->assume_inbound));
+ const auto* data_sinfo = data->struct_info_.as<TensorStructInfoNode>();
+
+ DataType dtype = DataType::Void();
+ Optional<VDevice> vdevice = NullOpt;
+ int ndim = kUnknownNDim;
+ if (data_sinfo) {
+ dtype = data_sinfo->dtype;
+ vdevice = data_sinfo->vdevice;
+ ndim = data_sinfo->ndim;
+ }
+
+ Optional<Expr> shape = [&]() -> Optional<Expr> {
+ if (!data_sinfo) return NullOpt;
+ if (!data_sinfo->shape) return NullOpt;
+
+ auto opt_axes_tuple = UnpackTupleOfPrimValue<Integer>(axes);
+ if (!opt_axes_tuple) return NullOpt;
+ auto axes_tuple = opt_axes_tuple.value();
+
+ auto opt_begin_tuple = UnpackTupleOfPrimValue(begin);
+ if (!opt_begin_tuple) return NullOpt;
+ auto begin_tuple = opt_begin_tuple.value();
+
+ CHECK_EQ(axes_tuple.size(), begin_tuple.size())
+ << "For operator " << call->op << ", "
+ << "the number of axes provided must match the number of 'begin'
indices. "
+ << "However, there are " << axes_tuple.size() << " axes specified ("
<< axes_tuple
+ << ") and " << begin_tuple.size() << " 'begin' indices specified (" <<
begin_tuple << ")";
+
+ auto opt_end_tuple = UnpackTupleOfPrimValue(end);
+ if (!opt_end_tuple) return NullOpt;
+ auto end_tuple = opt_end_tuple.value();
+
+ CHECK_EQ(axes_tuple.size(), end_tuple.size())
+ << "For operator " << call->op << ", "
+ << "the number of axes provided must match the number of 'end'
indices. "
+ << "However, there are " << axes_tuple.size() << " axes specified ("
<< axes_tuple
+ << ") and " << end_tuple.size() << " 'end' indices specified (" <<
end_tuple << ")";
+
+ Array<PrimExpr> strides_tuple;
+ if (strides.defined()) {
+ auto opt_strides_tuple = UnpackTupleOfPrimValue(strides);
+ if (!opt_strides_tuple) return NullOpt;
+
+ strides_tuple = opt_strides_tuple.value();
+ } else {
+ strides_tuple = Array<PrimExpr>(axes_tuple.size(),
IntImm(DataType::Int(64), 1));
+ }
+
+ CHECK_EQ(axes_tuple.size(), strides_tuple.size())
+ << "For operator " << call->op << ", "
+ << "when the optional 'strides' argument is provided, "
+ << "the number of axes provided must match the number of strides
provided. "
+ << "However, there are " << axes_tuple.size() << " axes specified ("
<< axes_tuple
+ << ") and " << strides_tuple.size() << " strides specified (" <<
strides_tuple << ")";
+
+ auto opt_data_shape = data_sinfo->GetShape();
+
+ if (axes_tuple.empty() && !opt_data_shape.defined()) {
+ return data_sinfo->shape.value();
+ } else if (!opt_data_shape.defined()) {
+ return NullOpt;
+ }
+
+ std::vector<int> axes = NormalizeAxes(call, ctx, data_sinfo->ndim,
axes_tuple);
+ auto attrs = call->attrs.as<StridedSliceAttrs>();
+
+ Array<PrimExpr> output_shape = data_sinfo->GetShape().value();
+ for (size_t i = 0; i < axes.size(); i++) {
+ size_t axis = axes[i];
+ PrimExpr input_dim = output_shape[axis];
+ PrimExpr begin = begin_tuple[i];
+ PrimExpr end = end_tuple[i];
+
+ PrimExpr output_dim =
+ GetLength(begin, end, strides_tuple[i], input_dim,
attrs->assume_inbound);
+
+ arith::Analyzer* analyzer = ctx->GetAnalyzer();
+ std::optional<With<arith::ConstraintContext>> context;
+ if (attrs->assume_inbound) {
+ context.emplace(analyzer, 0 <= begin && begin <= input_dim && 0 <= end
&& end <= input_dim);
+ }
+
+ output_dim = analyzer->Simplify(output_dim);
+
+ output_shape.Set(axis, output_dim);
+ }
+ return ShapeExpr(output_shape);
+ }();
+
+ if (shape.defined()) {
+ return TensorStructInfo(shape.value(), dtype, vdevice);
+ } else {
+ return TensorStructInfo(dtype, ndim, vdevice);
}
- return TensorStructInfo(ShapeExpr(output_shape), data_sinfo->dtype,
data_sinfo->vdevice);
}
InferLayoutOutput InferLayoutStridedSlice(const Call& call,
@@ -242,17 +453,29 @@ InferLayoutOutput InferLayoutStridedSlice(const Call&
call,
const auto* attrs = call->attrs.as<StridedSliceAttrs>();
ICHECK(attrs != nullptr) << "Invalid Call";
+
const auto* tensor_sinfo =
GetStructInfoAs<TensorStructInfoNode>(call->args[0]);
- ICHECK(tensor_sinfo != nullptr) << "Invalid Call";
- ICHECK(!tensor_sinfo->IsUnknownNdim()) << "Only support known ndim";
+ CHECK(tensor_sinfo) << "Invalid Call";
+ CHECK(!tensor_sinfo->IsUnknownNdim()) << "Layout inference only supports
known dimensionality, "
+ << "but expression " << call << " has
argument "
+ << call->args[0] << " of unknown
dimensionality.";
LayoutDecision existing_layout = GetLayoutDecision(var_layout_map,
call->args[0]);
- std::vector<Integer> new_axes;
- for (const auto& axis : attrs->axes) {
- new_axes.push_back(FindAxis(existing_layout->layout, axis->value));
+
+ auto opt_axes_tuple =
UnpackTupleOfPrimValue<Integer>(GetStructInfo(call->args[1]));
+ CHECK(opt_axes_tuple) << "Layout inference of " << call->op
+ << " requires slices to be along static axes. "
+ << "However, expression " << call << " slices along
non-static axes "
+ << call->args[1];
+ Array<Integer> axes_tuple = opt_axes_tuple.value();
+
+ Array<Expr> new_axes;
+ for (const auto& axis : axes_tuple) {
+ int new_axis = FindAxis(existing_layout->layout, axis->value);
+ new_axes.push_back(relax::PrimValue::Int64(new_axis));
}
- ObjectPtr<StridedSliceAttrs> new_attrs =
make_object<StridedSliceAttrs>(*attrs);
- new_attrs->axes = std::move(new_axes);
- return InferLayoutOutput({existing_layout}, {existing_layout},
Attrs(new_attrs));
+
+ return InferLayoutOutput({existing_layout}, {existing_layout}, call->attrs,
+ {{1, relax::Tuple(new_axes)}});
}
TVM_REGISTER_OP("relax.strided_slice")
diff --git a/src/relax/op/tensor/index.h b/src/relax/op/tensor/index.h
index c8c7428f48..3f0e5d227b 100644
--- a/src/relax/op/tensor/index.h
+++ b/src/relax/op/tensor/index.h
@@ -54,11 +54,7 @@ Expr take(Expr x, Expr indices, Optional<Integer> axis);
* \param assume_inbound Whether to assume the indices are in bound.
* \return The sliced result
*/
-Expr strided_slice(Expr x, //
- Array<Integer> axes, //
- Array<PrimExpr> begin, //
- Array<PrimExpr> end, //
- Optional<Array<PrimExpr>> strides, //
+Expr strided_slice(Expr x, Expr axes, Expr begin, Expr end, Optional<Expr>
strides = NullOpt,
bool assume_inbound = false);
} // namespace relax
diff --git a/src/relax/transform/convert_layout.cc
b/src/relax/transform/convert_layout.cc
index 6530d0d2cf..2f437545b6 100644
--- a/src/relax/transform/convert_layout.cc
+++ b/src/relax/transform/convert_layout.cc
@@ -107,11 +107,22 @@ class LayoutConvertMutator : public ExprMutator {
}
Array<Expr> RewriteArgs(const Array<Expr>& args, const Array<NLayout>& to) {
- ICHECK(args.size() == to.size());
+ // The `Array<Expr> args` array contains both tensor and
+ // non-tensor arguments, where the `Array<NLayout> to` array only
+ // contains tensor arguments. The number of tensor arguments in
+ // `args` should match the full extent of `to`.
+
+ ICHECK_LE(to.size(), args.size());
+
std::vector<Expr> new_args;
for (size_t i = 0; i < args.size(); ++i) {
- new_args.push_back(RewriteExpr(args[i], to[i]));
+ Expr arg = args[i];
+ if (i < to.size()) {
+ arg = RewriteExpr(arg, to[i]);
+ }
+ new_args.push_back(arg);
}
+
return std::move(new_args);
}
@@ -189,7 +200,11 @@ class LayoutConvertMutator : public ExprMutator {
} else {
// Convert the layout according to the inferred layout output.
Array<Expr> new_args = RewriteArgs(call_node->args,
res.value()->input_layouts);
+ for (const auto& [i, arg] : res.value()->new_args) {
+ new_args.Set(i->value, arg);
+ }
new_call->args = std::move(new_args);
+
new_call->attrs = std::move(res.value()->new_attrs);
Expr cur_call = builder_->Normalize(Call(new_call));
if (binding->var->IsInstance<DataflowVarNode>()) {
diff --git a/src/relax/transform/infer_layout_utils.h
b/src/relax/transform/infer_layout_utils.h
index 2cbbe23ede..4e54d92544 100644
--- a/src/relax/transform/infer_layout_utils.h
+++ b/src/relax/transform/infer_layout_utils.h
@@ -102,6 +102,7 @@ class InferLayoutOutputNode : public Object {
Array<NLayout> input_layouts;
Array<NLayout> output_layouts;
Attrs new_attrs;
+ Map<Integer, Expr> new_args;
void VisitAttrs(tvm::AttrVisitor* v) {
v->Visit("input_layouts", &input_layouts);
@@ -117,11 +118,12 @@ class InferLayoutOutputNode : public Object {
class InferLayoutOutput : public ObjectRef {
public:
explicit InferLayoutOutput(Array<NLayout> input_layouts, Array<NLayout>
output_layouts,
- Attrs new_attrs) {
+ Attrs new_attrs, Map<Integer, Expr> new_args =
{}) {
auto n = make_object<InferLayoutOutputNode>();
n->input_layouts = std::move(input_layouts);
n->output_layouts = std::move(output_layouts);
n->new_attrs = std::move(new_attrs);
+ n->new_args = std::move(new_args);
data_ = n;
}
TVM_DEFINE_OBJECT_REF_METHODS(InferLayoutOutput, ObjectRef,
InferLayoutOutputNode);
diff --git a/src/relax/utils.cc b/src/relax/utils.cc
index 77e6b33f0c..f0239e424f 100644
--- a/src/relax/utils.cc
+++ b/src/relax/utils.cc
@@ -65,49 +65,6 @@ class ExprBinder : public ExprMutator {
}
}
- Expr VisitExpr_(const CallNode* op) final {
- auto call_node = Downcast<Call>(ExprMutator::VisitExpr_(op));
-
- // Special case for strided_slice
- //
- // The strided_slice operator currently stores the begins/ends in
- // the CallNode::attrs. Because the CallNode::attrs is only
- // intended to store static information, any PrimExpr members in
- // the attributes are not visited by `ExprMutator::VisitPrimExpr`.
- // Therefore, these must be explicitly visited.
- //
- // When the strided_slice operator is updated to store begins/ends
- // as a tuple of `relax::PrimValue` in the arguments, this special
- // case can be removed.
- static auto strided_slice_op = Op::Get("relax.strided_slice");
- if (call_node->op.same_as(strided_slice_op)) {
- auto attrs = call_node->attrs.as<StridedSliceAttrs>();
-
- auto visit_prim_expr = [this](const auto& expr) { return
VisitPrimExpr(expr); };
-
- Array<PrimExpr> begin = attrs->begin.Map(visit_prim_expr);
- Array<PrimExpr> end = attrs->end.Map(visit_prim_expr);
- auto strides = attrs->strides;
- if (strides.defined()) {
- strides = strides.value().Map(visit_prim_expr);
- }
-
- bool all_same = begin.same_as(attrs->begin) && end.same_as(attrs->end) &&
- (!strides.defined() || strides.same_as(attrs->strides));
- if (!all_same) {
- ObjectPtr<StridedSliceAttrs> new_attrs =
make_object<StridedSliceAttrs>();
- new_attrs->axes = attrs->axes;
- new_attrs->begin = std::move(begin);
- new_attrs->end = std::move(end);
- new_attrs->strides = std::move(strides);
- new_attrs->assume_inbound = attrs->assume_inbound;
- call_node.CopyOnWrite()->attrs = Attrs(new_attrs);
- }
- }
-
- return std::move(call_node);
- }
-
Expr VisitExpr_(const VarNode* op) final {
auto id = GetRef<Var>(op);
auto it = args_map_.find(id);
diff --git a/src/script/ir_builder/relax/ir.cc
b/src/script/ir_builder/relax/ir.cc
index 60f78c0f58..2e94ae420a 100644
--- a/src/script/ir_builder/relax/ir.cc
+++ b/src/script/ir_builder/relax/ir.cc
@@ -70,6 +70,16 @@ tvm::relax::Var Arg(const String& name, const
tvm::relax::StructInfo& struct_inf
FunctionFrame frame = FindFunctionFrame("R.Arg");
tvm::relax::Var var(name, struct_info);
frame->params.push_back(var);
+
+ // This constraint would normally be provided as part of
+ // `BlockBuilder::BeginScope`. However, because the frame and its
+ // scope are initialized before the arguments are known, the scope
+ // doesn't have access to these constraints.
+ auto* analyzer = frame->block_builder->GetAnalyzer();
+ for (const auto& tir_var : DefinableTIRVarsInStructInfo(struct_info)) {
+ analyzer->MarkGlobalNonNegValue(tir_var);
+ }
+
return var;
}
diff --git a/tests/python/contrib/test_msc/test_graph_build.py
b/tests/python/contrib/test_msc/test_graph_build.py
index 3b1cfc4057..315d6813ea 100644
--- a/tests/python/contrib/test_msc/test_graph_build.py
+++ b/tests/python/contrib/test_msc/test_graph_build.py
@@ -17,6 +17,8 @@
""" Test graph builder && graph. """
+import pytest
+
import torch
from torch import fx
from torch.nn import Module
@@ -1099,6 +1101,7 @@ def test_getattr():
verify_model(GetAttr1(), input_info, expected)
[email protected](reason="MSC does not support Tuple of PrimValue")
def test_getitem():
"""test graph builder for getitem"""
diff --git a/tests/python/contrib/test_msc/test_translate_relax.py
b/tests/python/contrib/test_msc/test_translate_relax.py
index fdc1577715..00975be85e 100644
--- a/tests/python/contrib/test_msc/test_translate_relax.py
+++ b/tests/python/contrib/test_msc/test_translate_relax.py
@@ -17,6 +17,8 @@
""" Test translate from relax. """
+import pytest
+
import torch
from torch import fx
from torch.nn import Module
@@ -622,6 +624,7 @@ def test_getattr():
_verify_model(GetAttr1(), input_info)
[email protected](reason="MSC does not support Tuple of PrimValue")
def test_getitem():
"""test relax translator for getitem"""
diff --git a/tests/python/contrib/test_msc/test_translate_tensorflow.py
b/tests/python/contrib/test_msc/test_translate_tensorflow.py
index cb4ea3c02e..61f8ce1a97 100644
--- a/tests/python/contrib/test_msc/test_translate_tensorflow.py
+++ b/tests/python/contrib/test_msc/test_translate_tensorflow.py
@@ -18,6 +18,8 @@
""" Test translate from tensorflow. """
+import pytest
+
from packaging import version as package_version
import numpy as np
@@ -502,6 +504,7 @@ def _test_stridedslice(
verify_model(graph_def, golden, **io_info)
[email protected](reason="MSC does not support Tuple of PrimValue")
def test_stridedslice():
"""test tensorflow translator for stridedslice"""
@@ -1062,6 +1065,7 @@ def _test_slice_operation_input(input_value, begin_value,
size_value):
verify_model(graph_def, golden, **io_info)
[email protected](reason="MSC does not support Tuple of PrimValue")
def test_slice():
"""test tensorflow translator for slice"""
diff --git a/tests/python/contrib/test_msc/test_translate_torch.py
b/tests/python/contrib/test_msc/test_translate_torch.py
index 949c5669f9..81c6031ce1 100644
--- a/tests/python/contrib/test_msc/test_translate_torch.py
+++ b/tests/python/contrib/test_msc/test_translate_torch.py
@@ -17,6 +17,8 @@
""" Test translate from torch. """
+import pytest
+
import numpy as np
import torch
@@ -587,6 +589,7 @@ def test_getattr():
verify_model(GetAttr1(), input_info)
[email protected](reason="MSC does not support Tuple of PrimValue")
def test_getitem():
"""test torch translator for getitem"""
diff --git a/tests/python/relax/test_dataflow_pattern.py
b/tests/python/relax/test_dataflow_pattern.py
index 24c36d20dc..f67b0530ca 100644
--- a/tests/python/relax/test_dataflow_pattern.py
+++ b/tests/python/relax/test_dataflow_pattern.py
@@ -1563,23 +1563,37 @@ def test_iterative_rewrite_without_trivial_binding():
return c
pattern_arg = wildcard()
- pattern = is_op("relax.strided_slice")(pattern_arg).has_attr(
- {
- "axes": [0],
- "strides": [T.int64(1)],
- }
+ pattern_axes = wildcard()
+ pattern_begin = wildcard()
+ pattern_end = wildcard()
+ pattern_strides = wildcard()
+ pattern = is_op("relax.strided_slice")(
+ pattern_arg, pattern_axes, pattern_begin, pattern_end, pattern_strides
)
def rewriter(expr, matches):
arg = matches[pattern_arg]
+ axes = matches[pattern_axes]
+ begin = matches[pattern_begin]
+ end = matches[pattern_end]
+ strides = matches[pattern_strides]
strided_slice = matches[pattern]
if arg.struct_info.shape is None:
return expr
+ if len(axes) != 1:
+ return expr
+
+ axis = axes[0].value
+ begin = begin[0].value
+ end = end[0].value
+ stride = strides[0].value
+
+ if stride != 1:
+ return expr
+
size = arg.struct_info.shape[0]
- begin = strided_slice.attrs.begin[0]
- end = strided_slice.attrs.end[0]
if (
isinstance(size, tir.IntImm)
and isinstance(begin, tir.IntImm)
diff --git a/tests/python/relax/test_op_index.py
b/tests/python/relax/test_op_index.py
index 1455b4182a..57e7a14b70 100644
--- a/tests/python/relax/test_op_index.py
+++ b/tests/python/relax/test_op_index.py
@@ -528,7 +528,7 @@ def test_strided_slice_infer_struct_info_shape_var():
_check_inference(
bb,
relax.op.strided_slice(x0, axes=[0], begin=[0], end=[8]),
- relax.TensorStructInfo(dtype="float32", ndim=2),
+ relax.TensorStructInfo(shape=[8, 10], dtype="float32"),
)
_check_inference(
bb,
@@ -543,7 +543,7 @@ def test_strided_slice_infer_struct_info_shape_var():
_check_inference(
bb,
relax.op.strided_slice(x3, axes=[0], begin=[0], end=[8]),
- relax.TensorStructInfo(dtype="", ndim=2),
+ relax.TensorStructInfo(shape=[8, 10], dtype=""),
)
_check_inference(
bb,
@@ -614,12 +614,15 @@ def
test_strided_slice_infer_struct_info_symbolic_begin_end_strides():
_check_inference(
bb,
relax.op.strided_slice(x, axes=[0], begin=[0], end=[8], strides=[var]),
- relax.TensorStructInfo(dtype="float32", ndim=2),
+ relax.TensorStructInfo(
+ [tir.if_then_else(var < 0, -8 // (0 - var) + 1, (var + 7) // var),
9],
+ dtype="float32",
+ ),
)
_check_inference(
bb,
relax.op.strided_slice(x, axes=[0], begin=[0], end=[8],
strides=[size_var]),
- relax.TensorStructInfo(dtype="float32", ndim=2),
+ relax.TensorStructInfo([7 // size_var + 1, 9], dtype="float32"),
)
@@ -633,7 +636,7 @@ def
test_strided_slice_infer_struct_info_symbolic_begin_end_strides_inbound():
bb,
relax.op.strided_slice(x, axes=[0], begin=[var], end=[8],
assume_inbound=True),
relax.TensorStructInfo(
- (8 - tir.if_then_else(var < 0, var + 8, var), 9),
+ (8 - var, 9),
dtype="float32",
),
)
@@ -645,7 +648,7 @@ def
test_strided_slice_infer_struct_info_symbolic_begin_end_strides_inbound():
_check_inference(
bb,
relax.op.strided_slice(x, axes=[0], begin=[0], end=[var],
assume_inbound=True),
- relax.TensorStructInfo((tir.if_then_else(var < 0, var + 8, var), 9),
dtype="float32"),
+ relax.TensorStructInfo((var, 9), dtype="float32"),
)
_check_inference(
bb,
@@ -655,12 +658,12 @@ def
test_strided_slice_infer_struct_info_symbolic_begin_end_strides_inbound():
_check_inference(
bb,
relax.op.strided_slice(x, axes=[0], begin=[0], end=[8], strides=[var],
assume_inbound=True),
- relax.TensorStructInfo(dtype="float32", ndim=2),
+ relax.TensorStructInfo([(var + 7) // var, 9], dtype="float32"),
)
_check_inference(
bb,
relax.op.strided_slice(x, axes=[0], begin=[0], end=[8], strides=[var],
assume_inbound=True),
- relax.TensorStructInfo(dtype="float32", ndim=2),
+ relax.TensorStructInfo([(var + 7) // var, 9], dtype="float32"),
)
@@ -696,7 +699,7 @@ def test_strided_slice_infer_struct_info_no_axis():
_check_inference(
bb,
relax.op.strided_slice(x3, axes=[], begin=[], end=[]),
- relax.TensorStructInfo(s0, "float32"),
+ relax.TensorStructInfo([m, n], "float32"),
)
_check_inference(
bb,
@@ -716,15 +719,19 @@ def test_strided_slice_begin_end_strides_int64():
x, axes=[0, 1, 3], begin=[1, 0, 8], end=[8, 9, 0], strides=[2, 1, -3]
)
- assert strided_slice.attrs.begin[0].dtype == "int64"
- assert strided_slice.attrs.begin[1].dtype == "int64"
- assert strided_slice.attrs.begin[2].dtype == "int64"
- assert strided_slice.attrs.end[0].dtype == "int64"
- assert strided_slice.attrs.end[1].dtype == "int64"
- assert strided_slice.attrs.end[2].dtype == "int64"
- assert strided_slice.attrs.strides[0].dtype == "int64"
- assert strided_slice.attrs.strides[1].dtype == "int64"
- assert strided_slice.attrs.strides[2].dtype == "int64"
+ begins = strided_slice.args[1]
+ ends = strided_slice.args[2]
+ strides = strided_slice.args[3]
+
+ assert begins[0].struct_info.dtype == "int64"
+ assert begins[1].struct_info.dtype == "int64"
+ assert begins[2].struct_info.dtype == "int64"
+ assert ends[0].struct_info.dtype == "int64"
+ assert ends[1].struct_info.dtype == "int64"
+ assert ends[2].struct_info.dtype == "int64"
+ assert strides[0].struct_info.dtype == "int64"
+ assert strides[1].struct_info.dtype == "int64"
+ assert strides[2].struct_info.dtype == "int64"
def test_strided_slice_inconsistent_axes_begin_end_strides_length():
