This is an automated email from the ASF dual-hosted git repository.
tlopex 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 9d5253e2ea [Fix][Relax] Skip ReorderPermuteDimsAfterConcat for
unknown-rank inputs (#20216)
9d5253e2ea is described below
commit 9d5253e2ea59897664b92c0e8b5508fd83475fdd
Author: yanght27 <[email protected]>
AuthorDate: Sat Aug 29 12:23:27 2026 +0800
[Fix][Relax] Skip ReorderPermuteDimsAfterConcat for unknown-rank inputs
(#20216)
Fixes #20204.
`ReorderPermuteDimsAfterConcat` derives the permutation for
`permute_dims(..., axes=None)` from the input rank. When an input tensor
has unknown rank, its `ndim` is `-1`, but the pass previously asserted
that `ndim >= 0`. As a result, valid Relax IR could raise an
`InternalError`.
This PR changes the axis derivation helper to return
`ffi::Optional<ffi::Array<int64_t>>`. When the axes are implicit and the
input rank is unknown, a concrete permutation cannot be derived, so the
pass conservatively leaves the original expression unchanged. Existing
behavior is preserved for known-rank inputs and explicit axes.
A regression test covers unknown-rank inputs followed by `concat` and
verifies that the pass neither rewrites the expression nor raises an
error.
Testing:
- Built the modified C++ target successfully.
- `test_transform_reorder_permute_dims_after_concat.py`: 7 passed.
- Pre-commit hooks passed on both modified files.
- The original issue reproducer completes without an `InternalError`.
---
.../transform/reorder_permute_dims_after_concat.cc | 66 ++++++++++++++--------
..._transform_reorder_permute_dims_after_concat.py | 20 +++++++
2 files changed, 61 insertions(+), 25 deletions(-)
diff --git a/src/relax/transform/reorder_permute_dims_after_concat.cc
b/src/relax/transform/reorder_permute_dims_after_concat.cc
index e1dbbab826..91f7633d76 100644
--- a/src/relax/transform/reorder_permute_dims_after_concat.cc
+++ b/src/relax/transform/reorder_permute_dims_after_concat.cc
@@ -91,40 +91,55 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr,
ffi::Map<DFPattern, Expr>)>>
return attrs->axes;
};
- auto get_permute_dims_axes =
- [get_permute_dims_optional_axes](const Expr& expr) ->
ffi::Array<int64_t> {
+ auto try_get_permute_dims_axes =
+ [get_permute_dims_optional_axes](const Expr& expr) ->
ffi::Optional<ffi::Array<int64_t>> {
if (auto opt_axes = get_permute_dims_optional_axes(expr)) {
- return opt_axes.value();
- } else {
- auto call = expr.as_or_throw<Call>();
- ffi::Array<int64_t> permutation;
- auto arg_ty = call->args[0]->ty.as<TensorTypeNode>();
- TVM_FFI_ICHECK(arg_ty) << "Expected permute_dims to have a single tensor
argument, "
- << "but argument " << call->args[0] << " has type
"
- << call->args[0]->ty;
- TVM_FFI_ICHECK_GE(arg_ty->ndim, 0);
- size_t ndim = arg_ty->ndim;
- for (size_t i = 0; i < ndim; i++) {
- permutation.push_back(static_cast<int64_t>(ndim - i - 1));
- }
- return permutation;
+ return opt_axes;
+ }
+
+ auto call = expr.as_or_throw<Call>();
+ ffi::Array<int64_t> permutation;
+ auto arg_ty = call->args[0]->ty.as<TensorTypeNode>();
+ TVM_FFI_ICHECK(arg_ty) << "Expected permute_dims to have a single tensor
argument, "
+ << "but argument " << call->args[0] << " has type "
<< call->args[0]->ty;
+
+ // An implicit permutation depends on the input rank, so it cannot be
materialized here.
+ if (arg_ty->IsUnknownNdim()) {
+ return std::nullopt;
+ }
+
+ size_t ndim = arg_ty->ndim;
+ for (size_t i = 0; i < ndim; i++) {
+ permutation.push_back(static_cast<int64_t>(ndim - i - 1));
}
+ return permutation;
};
- auto permute_dims_axes_are_compatible = [&](const ffi::Array<Expr>&
permute_dims) -> bool {
- auto first_axes = get_permute_dims_axes(permute_dims[0]);
+ auto try_get_compatible_permute_dims_axes =
+ [try_get_permute_dims_axes](
+ const ffi::Array<Expr>& permute_dims) ->
ffi::Optional<ffi::Array<int64_t>> {
+ auto opt_first_axes = try_get_permute_dims_axes(permute_dims[0]);
+ if (!opt_first_axes) {
+ return std::nullopt;
+ }
+ const auto& first_axes = opt_first_axes.value();
+
for (size_t i_arg = 1; i_arg < permute_dims.size(); i_arg++) {
- auto i_axes = get_permute_dims_axes(permute_dims[i_arg]);
+ auto opt_i_axes = try_get_permute_dims_axes(permute_dims[i_arg]);
+ if (!opt_i_axes) {
+ return std::nullopt;
+ }
+ const auto& i_axes = opt_i_axes.value();
if (i_axes.size() != first_axes.size()) {
- return false;
+ return std::nullopt;
}
for (size_t i_axis = 0; i_axis < first_axes.size(); i_axis++) {
if (i_axes[i_axis] != first_axes[i_axis]) {
- return false;
+ return std::nullopt;
}
}
}
- return true;
+ return opt_first_axes;
};
auto rewriter = [=](Expr expr, ffi::Map<DFPattern, Expr> matches) -> Expr {
@@ -141,9 +156,12 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr,
ffi::Map<DFPattern, Expr>)>>
<< "Pattern match should return at least " << min_concat << " items,
but only found "
<< all_permute_dims.size() << ": " << all_permute_dims;
- if (!permute_dims_axes_are_compatible(all_permute_dims)) {
+ auto opt_permute_dims_axes =
try_get_compatible_permute_dims_axes(all_permute_dims);
+ if (!opt_permute_dims_axes) {
return expr;
}
+ const auto& permute_dims_axes = opt_permute_dims_axes.value();
+
ffi::Optional<ffi::Array<int64_t>> permute_axes =
get_permute_dims_optional_axes(all_permute_dims[0]);
@@ -151,8 +169,6 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr,
ffi::Map<DFPattern, Expr>)>>
auto concat_attrs = concat_call->attrs.as<ConcatAttrs>();
TVM_FFI_ICHECK(concat_attrs);
- auto permute_dims_axes = get_permute_dims_axes(all_permute_dims[0]);
-
int64_t old_concat_axis = concat_attrs->axis.value_or(0);
int64_t ndim = static_cast<int64_t>(permute_dims_axes.size());
if (old_concat_axis < 0) {
diff --git
a/tests/python/relax/test_transform_reorder_permute_dims_after_concat.py
b/tests/python/relax/test_transform_reorder_permute_dims_after_concat.py
index 2da6cfcda9..4f8cd848d4 100644
--- a/tests/python/relax/test_transform_reorder_permute_dims_after_concat.py
+++ b/tests/python/relax/test_transform_reorder_permute_dims_after_concat.py
@@ -70,6 +70,26 @@ class TestSimple(Base):
return out
+class TestDoNotRewriteUnknownRank(Base):
+ """Do not rewrite implicit permutations when their rank is unknown."""
+
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(
+ x: R.Tensor(dtype="float32"),
+ y: R.Tensor(dtype="float32"),
+ ):
+ with R.dataflow():
+ x_t = R.permute_dims(x)
+ y_t = R.permute_dims(y)
+ out = R.concat([x_t, y_t], axis=0)
+ R.output(out)
+ return out
+
+ Expected = Before
+
+
class TestCombineExplicitAndImplicitAxes(Base):
"""Check for explicit axes to be permuted
