gemini-code-assist[bot] commented on code in PR #18491: URL: https://github.com/apache/tvm/pull/18491#discussion_r3409919635
########## src/relax/transform/canonicalize_shape_expr.cc: ########## @@ -0,0 +1,122 @@ +/* + * 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<tirx::VarNode>(); +} + +class ShapeExprCanonicalizer : public ExprMutator { + public: + using ExprMutator::VisitExpr_; + + BindingBlock VisitBindingBlock(const BindingBlock& block) final { + bool prev = inside_binding_block_; + inside_binding_block_ = true; + BindingBlock ret = ExprMutator::VisitBindingBlock(block); + inside_binding_block_ = prev; + return ret; + } Review Comment:  If a local function (lambda) is defined inside a binding block, `inside_binding_block_` will remain `true` when visiting the local function's parameters and return type. This can lead to incorrect canonicalization of the nested function's parameters in the outer function's scope, emitting `MatchCast` bindings into the wrong block. We should override `VisitExpr_(const FunctionNode* op)` to temporarily reset `inside_binding_block_` to `false` when entering a local function. ```c BindingBlock VisitBindingBlock(const BindingBlock& block) final { bool prev = inside_binding_block_; inside_binding_block_ = true; BindingBlock ret = ExprMutator::VisitBindingBlock(block); inside_binding_block_ = prev; return ret; } Expr VisitExpr_(const FunctionNode* op) final { bool prev = inside_binding_block_; inside_binding_block_ = false; Expr ret = ExprMutator::VisitExpr_(op); inside_binding_block_ = prev; return ret; } ``` ########## tests/python/relax/test_transform_canonicalize_shape_expr.py: ########## @@ -0,0 +1,76 @@ +# 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, tirx +from tvm.script import ir as I +from tvm.script import relax as R +from tvm.script import tirx as T + + [email protected]_module +class Before: + @R.function + def main(x: R.Tensor(("x_0", "x_1", "x_2", "x_3"), "float32")): + R.func_attr({"relax.force_pure": True}) + 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, (tirx.IntImm, tirx.Var)): + composite_dims.append(dim) + + relax.analysis.post_order_visit(mod["main"], _visit) + assert not composite_dims + + +def test_canonicalize_shape_expr_unblocks_vm_shape_lower(): + mod = Before + mod = relax.transform.CanonicalizeShapeExpr()(mod) + mod = relax.transform.ComputePrimValue()(mod) + mod = relax.transform.VMShapeLower()(mod) + + assert any("compute_symbolic_expr" in gv.name_hint for gv in mod.get_global_vars()) + + [email protected]_module +class ParamCompoundShape: + @R.function + def main(x: R.Tensor(("A", "B", "A + B"), "float32")) -> R.Tensor((1,), "float32"): + out: R.Tensor((1,), "float32") = R.zeros(R.shape([1]), dtype="float32") + return out + + +def test_canonicalize_shape_expr_skips_parameter_struct_info(): + mod = relax.transform.CanonicalizeShapeExpr()(ParamCompoundShape) + param_shape = mod["main"].params[0].struct_info.shape + + assert any(not isinstance(dim, (tirx.IntImm, tirx.Var)) for dim in param_shape.values) + + +if __name__ == "__main__": + tvm.testing.main() Review Comment:  Add a test case to verify that nested functions do not have their parameters incorrectly canonicalized in the outer function's scope. ```python @I.ir_module class NestedFunc: @R.function def main(x: R.Tensor(("n",), "float32")): with R.dataflow(): @R.function def local_func(y: R.Tensor(("n + 1",), "float32")): out = R.zeros(R.shape([n + 1]), dtype="float32") return out res = local_func(x) R.output(res) return res def test_canonicalize_shape_expr_nested_function(): mod = relax.transform.CanonicalizeShapeExpr()(NestedFunc) local_func = mod["main"].body.blocks[0].bindings[0].value param_shape = local_func.params[0].struct_info.shape assert any(not isinstance(dim, (tirx.IntImm, tirx.Var)) for dim in param_shape.values) if __name__ == "__main__": tvm.testing.main() ``` ########## src/relax/transform/canonicalize_shape_expr.cc: ########## @@ -0,0 +1,122 @@ +/* + * 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<tirx::VarNode>(); +} + +class ShapeExprCanonicalizer : public ExprMutator { + public: + using ExprMutator::VisitExpr_; + + BindingBlock VisitBindingBlock(const BindingBlock& block) final { + bool prev = inside_binding_block_; + inside_binding_block_ = true; + BindingBlock ret = ExprMutator::VisitBindingBlock(block); + inside_binding_block_ = prev; + return ret; + } + + Expr VisitExpr_(const ShapeExprNode* op) final { + if (!inside_binding_block_) { + return ffi::GetRef<ShapeExpr>(op); + } + + 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: + tirx::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_++); + tirx::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)); + MatchCast binding(match_var, PrimValue(expr), target_sinfo); Review Comment:  The local variable `binding` of type `MatchCast` is constructed but never used. Since `builder_->EmitNormalized` already creates and emits the `MatchCast` binding, this line is completely redundant and can be safely removed. ```suggestion builder_->EmitNormalized(MatchCast(match_var, PrimValue(expr), target_sinfo)); ``` -- 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]
