gemini-code-assist[bot] commented on code in PR #19650:
URL: https://github.com/apache/tvm/pull/19650#discussion_r3330713773
##########
src/relax/transform/adjust_matmul_order.cc:
##########
@@ -141,20 +141,44 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr,
ffi::Map<DFPattern, Expr>)>>
auto shape_b = opt_shape_b.value();
auto shape_c = opt_shape_c.value();
+ auto permute_last_two_dims = [&](Expr expr) -> Expr {
+ auto opt_shape = get_shape(expr);
+ if (!opt_shape) return expr;
+
+ size_t ndim = opt_shape.value().size();
+ TVM_FFI_ICHECK_GE(ndim, 2);
+
+ ffi::Optional<ffi::Array<int64_t>> axes;
+
+ if (ndim == 2) {
+ // Pass none axes to permute_dims for simple transpose of 2D tensors.
+ axes = std::nullopt;
+ } else {
+ ffi::Array<int64_t> axes_array;
+ for (size_t i = 0; i < ndim; ++i) axes_array.push_back(i);
+ axes_array.Set(ndim - 1, ndim - 2);
+ axes_array.Set(ndim - 2, ndim - 1);
+ axes = ffi::Optional<ffi::Array<int64_t>>(axes_array);
+ }
+ return permute_dims(std::move(expr), axes);
+ };
+
+ auto transpose_shape_last_two_dims = [&](ffi::Array<PrimExpr>& shape) {
+ PrimExpr last_dim_shape = shape[shape.size() - 1];
+ shape.Set(shape.size() - 1, shape[shape.size() - 2]);
+ shape.Set(shape.size() - 2, last_dim_shape);
+ };
+
if (matches.count(pat_permuted_matmul_on_lhs)) {
- expr_a = permute_dims(expr_a, std::nullopt);
- expr_b = permute_dims(expr_b, std::nullopt);
- TVM_FFI_ICHECK_EQ(shape_a.size(), 2);
- TVM_FFI_ICHECK_EQ(shape_b.size(), 2);
- shape_a = {shape_a[1], shape_a[0]};
- shape_b = {shape_b[1], shape_b[0]};
+ expr_a = permute_last_two_dims(expr_a);
+ expr_b = permute_last_two_dims(expr_b);
+ transpose_shape_last_two_dims(shape_a);
+ transpose_shape_last_two_dims(shape_b);
} else if (matches.count(pat_permuted_matmul_on_rhs)) {
- expr_b = permute_dims(expr_b, std::nullopt);
- expr_c = permute_dims(expr_c, std::nullopt);
- TVM_FFI_ICHECK_EQ(shape_b.size(), 2);
- TVM_FFI_ICHECK_EQ(shape_c.size(), 2);
- shape_b = {shape_b[1], shape_b[0]};
- shape_c = {shape_c[1], shape_c[0]};
+ expr_b = permute_last_two_dims(expr_b);
+ expr_c = permute_last_two_dims(expr_c);
+ transpose_shape_last_two_dims(shape_b);
+ transpose_shape_last_two_dims(shape_c);
}
Review Comment:
If any of the operands have a dimension less than 2 (e.g., 1D tensors),
calling `transpose_shape_last_two_dims` will cause an underflow on
`shape.size() - 2` (since `shape.size()` is 1, `1 - 2` wraps around to a very
large number), leading to out-of-bounds memory access or a crash. Additionally,
`permute_last_two_dims` will trigger an FFI check failure.
Since 1D tensors are explicitly supported and padded later in this pass, we
should gracefully skip the reordering optimization if any of the permuted
operands are less than 2D.
```suggestion
if (matches.count(pat_permuted_matmul_on_lhs)) {
if (shape_a.size() < 2 || shape_b.size() < 2) return expr;
expr_a = permute_last_two_dims(expr_a);
expr_b = permute_last_two_dims(expr_b);
transpose_shape_last_two_dims(shape_a);
transpose_shape_last_two_dims(shape_b);
} else if (matches.count(pat_permuted_matmul_on_rhs)) {
if (shape_b.size() < 2 || shape_c.size() < 2) return expr;
expr_b = permute_last_two_dims(expr_b);
expr_c = permute_last_two_dims(expr_c);
transpose_shape_last_two_dims(shape_b);
transpose_shape_last_two_dims(shape_c);
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
