gemini-code-assist[bot] commented on code in PR #18491: URL: https://github.com/apache/tvm/pull/18491#discussion_r3343143701
########## src/relax/transform/canonicalize_shape_expr.cc: ########## @@ -0,0 +1,109 @@ +/* + * 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 src/relax/transform/canonicalize_shape_expr.cc + * \brief Canonicalize ShapeExpr by replacing composite PrimExpr dimensions with symbolic vars. + */ + +#include <tvm/ffi/reflection/registry.h> +#include <tvm/relax/analysis.h> +#include <tvm/relax/expr.h> +#include <tvm/relax/expr_functor.h> +#include <tvm/relax/transform.h> + +#include <string> +#include <unordered_map> + +namespace tvm { +namespace relax { + +namespace { + +bool IsSimpleShapeDim(const PrimExpr& expr) { + return expr->IsInstance<IntImmNode>() || expr->IsInstance<tir::VarNode>(); +} + +class ShapeExprCanonicalizer : public ExprMutator { + public: + using ExprMutator::VisitExpr_; + + Expr VisitExpr_(const ShapeExprNode* op) final { + ffi::Array<PrimExpr> new_values; + bool changed = false; + for (const PrimExpr& dim : op->values) { + if (IsSimpleShapeDim(dim)) { + new_values.push_back(dim); + continue; + } + + changed = true; + new_values.push_back(GetOrCreateSymbol(dim)); + } + + if (!changed) { + return ffi::GetRef<ShapeExpr>(op); + } + return ShapeExpr(new_values, op->span); + } + + private: + tir::Var GetOrCreateSymbol(const PrimExpr& expr) { + auto it = expr_to_var_.find(expr); + if (it != expr_to_var_.end()) { + return it->second; + } + + std::string base_name = "shape_expr_symbol_" + std::to_string(symbol_counter_++); + tir::Var sym_var(base_name, expr->dtype); + expr_to_var_.emplace(expr, sym_var); + + PrimStructInfo target_sinfo(sym_var); + Var match_var(base_name + "_pv", target_sinfo); + builder_->EmitNormalized(MatchCast(match_var, PrimValue(expr), target_sinfo)); + + return sym_var; + } + + int symbol_counter_ = 0; + std::unordered_map<PrimExpr, tir::Var, StructuralHash, StructuralEqual> expr_to_var_; +}; Review Comment:  If a Relax function has a parameter with a compound shape expression (e.g., `x: R.Tensor((n + 1,), "float32")`), `CanonicalizeShapeExpr` will attempt to canonicalize it. However, since parameter mutation happens outside of any active binding block, calling `builder_->EmitNormalized` inside `GetOrCreateSymbol` will crash because there is no active block in the block builder.\n\nTo prevent this, we should track whether we are mutating parameters, accumulate any created `MatchCast` bindings, and then prepend them as a new `BindingBlock` at the beginning of the function body before mutating it. This ensures that the bindings are properly normalized and emitted inside the function body. ```c class ShapeExprCanonicalizer : public ExprMutator {\n public:\n using ExprMutator::VisitExpr_;\n\n Expr VisitExpr_(const FunctionNode* op) final {\n bool old_is_mutating_params = is_mutating_params_;\n tvm::Array<Binding> old_accumulated_bindings = accumulated_bindings_;\n\n is_mutating_params_ = true;\n accumulated_bindings_ = {};\n\n tvm::Array<Var> params;\n bool params_changed = false;\n for (Var param : op->params) {\n Var new_param = this->VisitVarDef(param);\n params.push_back(new_param);\n if (!new_param.same_as(param)) {\n params_changed = true;\n }\n }\n is_mutating_params_ = false;\n\n Expr body = op->body;\n if (!accumulated_bindings_.empty()) {\n const SeqExprNode* seq = op->body.as<SeqExprNode>();\n ICHECK(seq != nullptr);\n tvm::Array<BindingBlock> new_blocks;\n new_blocks.push_back(BindingBlock(accumulated_bindings_));\n for (auto block : seq->blocks) {\n new_blocks.push_b ack(block);\n }\n body = SeqExpr(new_blocks, seq->body, seq->span);\n }\n\n Expr new_body = this->VisitWithNewScope(body, params);\n\n is_mutating_params_ = old_is_mutating_params;\n accumulated_bindings_ = old_accumulated_bindings;\n\n if (params_changed || !new_body.same_as(op->body)) {\n return Function(params, new_body, op->ret_struct_info, op->is_pure, op->attrs, op->span);\n }\n return GetRef<Expr>(op);\n }\n\n Expr VisitExpr_(const ShapeExprNode* op) final {\n ffi::Array<PrimExpr> new_values;\n bool changed = false;\n for (const PrimExpr& dim : op->values) {\n if (IsSimpleShapeDim(dim)) {\n new_values.push_back(dim);\n continue;\n }\n\n changed = true;\n new_values.push_back(GetOrCreateSymbol(dim));\n }\n\n if (!changed) {\n return ffi::GetRef<ShapeExpr>(op);\n }\n return ShapeExpr(new_values, op->span);\n }\n\n private:\n tir::Var GetOrCreateSymbol(const PrimExpr& expr) {\n auto it = expr_to_var_.find(expr);\n if (it != expr_to_var_.end()) {\n return it->second;\n }\n\n std::string base_name = "shape_expr_symbol_" + std::to_string(symbol_counter_++);\n tir::Var sym_var(base_name, expr->dtype);\n expr_to_var_.emplace(expr, sym_var);\n\n PrimStructInfo target_sinfo(sym_var);\n Var match_var(base_name + "_pv", target_sinfo);\n MatchCast match_cast(match_var, PrimValue(expr), target_sinfo);\n if (is_mutating_params_) {\n accumulated_bindings_.push_back(match_cast);\n } else {\n builder_->EmitNormalized(match_cast);\n }\n\n return sym_var;\n }\n\n int symbol_counter_ = 0;\n bool is_mutating_params_ = false;\n tvm::Array<Binding> accumulated_bindings_;\n std::unordered_map<PrimExpr, tir::Var, StructuralHash, StructuralEqual> expr_to_var_;\n}; ``` ########## tests/python/relax/test_transform_canonicalize_shape_expr.py: ########## @@ -0,0 +1,60 @@ +# 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. + +import tvm.testing +from tvm import relax, tir +from tvm.script import ir as I +from tvm.script import relax as R +from tvm.script import tir as T + + [email protected]_module +class Before: + @R.function + def main(x: R.Tensor(("x_0", "x_1", "x_2", "x_3"), "float32")): + x_0, x_1, x_2, x_3 = T.int64(), T.int64(), T.int64(), T.int64() + out: R.Tensor((T.int64(4) * (x_0 * x_1 * x_2 * x_3),), "float32") = R.zeros( + R.shape([T.int64(4) * (x_0 * x_1 * x_2 * x_3)]), dtype="float32" + ) + return out + + +def test_canonicalize_shape_expr_removes_composite_dims(): + mod = relax.transform.CanonicalizeShapeExpr()(Before) + composite_dims = [] + + def _visit(expr): + if isinstance(expr, relax.ShapeExpr): + for dim in expr.values: + if not isinstance(dim, tir.IntImm | tir.Var): Review Comment:  Using the `|` operator for union types in `isinstance` (PEP 604) requires Python 3.10 or newer. Since TVM supports older Python versions (such as Python 3.8 and 3.9), this will raise a `TypeError` in those environments. Please use a tuple of types instead: `isinstance(dim, (tir.IntImm, tir.Var))`. ```suggestion if not isinstance(dim, (tir.IntImm, tir.Var)): ``` ########## python/tvm/relax/transform/transform.py: ########## @@ -739,6 +739,32 @@ def FoldConstant() -> tvm.ir.transform.Pass: return _ffi_api.FoldConstant() # type: ignore +def CanonicalizeShapeExpr() -> tvm.ir.transform.Pass: + """Canonicalize ShapeExpr by replacing compound PrimExpr with fresh symbolic variables. + + VMShapeLower can only handle ShapeExpr where each dimension is either: + - IntImm (concrete integer constant) + - tir::Var (symbolic variable from function parameters or match_cast) + + This pass transforms compound PrimExpr (e.g., n+1, 4*n*m) by: + 1. Creating a fresh tir::Var for each compound expression + 2. Emitting a MatchCast that binds the fresh var to a PrimValue computing the expression + 3. Replacing the compound expression in ShapeExpr with teh fresh var Review Comment:  Fix typo: 'teh' -> 'the'. ```suggestion 3. Replacing the compound expression in ShapeExpr with the fresh var ``` -- 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]
