tlopex commented on code in PR #20243:
URL: https://github.com/apache/tvm/pull/20243#discussion_r3927755040
##########
src/relax/transform/combine_parallel_matmul.cc:
##########
@@ -117,19 +117,26 @@ Patterns CreatePatterns(const BranchInfo& branch_info) {
/*! \brief Create a rewriter for the given parallel matmul branches. */
ffi::TypedFunction<ffi::Map<Var, Expr>(ffi::Map<DFPattern, Var>, ffi::Map<Var,
Expr>)> GetRewriter(
const Patterns& patterns, const BranchInfo& branch_info, FCheck check) {
- auto batch_dims_compatible = [](size_t rhs_dim, const std::vector<size_t>&
indices,
- const std::vector<ffi::Array<PrimExpr>>&
rhs_shapes) {
- arith::Analyzer ana;
- for (auto ind : indices) {
- TVM_FFI_ICHECK_EQ(static_cast<int>(rhs_shapes[ind].size()), rhs_dim);
- // -2 for reduction and concat axes
- for (size_t i = 0; i < rhs_dim - 2; ++i) {
- if (!ana->CanProve(rhs_shapes[indices[0]][i] == rhs_shapes[ind][i])) {
- return false;
+ auto shapes_compatible_excluding_trailing_axes =
+ [](const std::vector<ffi::Array<PrimExpr>>& shapes, size_t
num_trailing_axes_excluded) {
+ arith::Analyzer ana;
+ size_t ndim = shapes[0].size();
+ for (const auto& shape : shapes) {
+ TVM_FFI_ICHECK_EQ(shape.size(), ndim);
+ for (size_t i = 0; i < ndim - num_trailing_axes_excluded; ++i) {
+ if (!ana->CanProve(shapes[0][i] == shape[i])) {
+ return false;
+ }
+ }
}
- }
- }
- return true;
+ return true;
+ };
+ auto batch_dims_compatible = [&](const std::vector<size_t>& indices,
+ const std::vector<ffi::Array<PrimExpr>>&
rhs_shapes) {
+ std::vector<ffi::Array<PrimExpr>> selected;
+ selected.reserve(indices.size());
+ for (size_t ind : indices) selected.push_back(rhs_shapes[ind]);
+ return shapes_compatible_excluding_trailing_axes(selected, 2);
Review Comment:
Please guard against ranks smaller than `num_trailing_axes_excluded` here.
Scalar biases are valid broadcast operands, and `std::optional<int>{0}` still
enters the bias-fusion branch. For two matmul branches with scalar biases,
`ndim` is 0 and this call excludes one trailing axis, so `ndim -
num_trailing_axes_excluded` underflows as a `size_t` and the loop accesses a
nonexistent dimension. The later `shape[shape.size() - 1]` check has the same
issue.
Could this helper return `false` when `ndim < num_trailing_axes_excluded`
(and preferably treat rank mismatches as incompatible rather than `ICHECK`),
with a regression test using two scalar biases? Such branches should skip bias
fusion because rank-0 tensors have no concat axis.
--
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]
