Refactor VisitExpr_

Author: cchung100m

src/relax/transform/canonicalize_shape_expr.cc

@@ -62,53 +62,153 @@ bool IsCanonicalPrimExpr(const PrimExpr& expr) {
  * \brief Mutator to canonicalize ShapeExpr in struct info
  *
  * This pass handles ShapeExpr canonicalization by:
- * 1. Detecting compound PrimExpr in ShapeExpr dimensions
- * 2. Lifting them into separate ShapeExpr bindings
+ * 1. Detecting compound PrimExpr in variable struct_info
+ * 2. Emitting ShapeExpr bindings to compute expressions
  * 3. Using MatchCast to extract values into fresh symbolic tir::Var
- * 4. Replacing compound expressions with these canonical vars
+ * 4. Replacing compound expressions with these canonical vars in struct_info
  */
 class ShapeExprCanonicalizer : public ExprMutator {
  public:
   using ExprMutator::VisitExpr_;
 
   Expr VisitExpr_(const FunctionNode* func) override {
     // Reset state for each function
-    auto cached_compound_to_var = compound_expr_to_var_;
-    auto cached_counter = symbolic_var_counter_;
+    symbolic_var_counter_ = 0;
+    compound_expr_to_var_.clear();
+    emitted_bindings_.clear();
 
-    auto result = ExprMutator::VisitExpr_(func);
+    // Process the function body
+    Expr new_body = VisitExpr(func->body);
 
-    compound_expr_to_var_ = cached_compound_to_var;
-    symbolic_var_counter_ = cached_counter;
+    if (new_body.same_as(func->body)) {
+      return ffi::GetRef<Function>(func);
+    }
 
-    return result;
+    return Function(func->params, new_body, func->ret_struct_info, func->is_pure, func->attrs,
+                    func->span);
+  }
+
+  void VisitBinding_(const VarBindingNode* binding) override {
+    // Emit canonicalization bindings before processing the binding
+    auto original_sinfo = GetStructInfo(binding->var);
+    if (NeedsCanonicalization(original_sinfo)) {
+      CanonicalizeStructInfoAndEmit(original_sinfo);
+    }
+
+    // Compute the updated var and value
+    Expr new_value = this->VisitExpr(binding->value);
+    Var new_var = this->VisitVarDef(binding->var);
+
+    // Update var_remap_ to register the new variable
+    this->var_remap_[binding->var->vid] = new_var;
+    this->var_remap_[new_var->vid] = new_var;
+
+    // If the var's struct_info changed and no longer matches the value, emit MatchCast
+    StructInfo var_sinfo = GetStructInfo(new_var);
+    StructInfo value_sinfo = GetStructInfo(new_value);
+    StructuralEqual struct_equal;
+    if (!struct_equal(var_sinfo, value_sinfo)) {
+      MatchCast match_cast = MatchCast(new_var, new_value, var_sinfo);
+      builder_->EmitNormalized(match_cast);
+      builder_->AddDefinitionToScope(match_cast->var);
+    } else {
+      builder_->EmitNormalized(VarBinding(new_var, new_value));
+    }
+  }
+
+  void VisitBinding_(const MatchCastNode* binding) override {
+    // Emit canonicalization bindings before processing the binding
+    auto original_sinfo = GetStructInfo(binding->var);
+    if (NeedsCanonicalization(original_sinfo)) {
+      CanonicalizeStructInfoAndEmit(original_sinfo);
+    }
+
+    ExprMutator::VisitBinding_(binding);
   }
 
-  /*!
-   * \brief Override VisitVarDef to canonicalize struct_info
-   *
-   * This is where we intercept variable definitions and canonicalize any
-   * compound PrimExpr in their TensorStructInfo shapes.
-   */
   Var VisitVarDef(const Var& var) override {
     auto sinfo = GetStructInfo(var);
-
-    // Check if we need to canonicalize the struct_info
     auto canonical_sinfo = CanonicalizeStructInfo(sinfo);
 
     if (canonical_sinfo.same_as(sinfo)) {
-      // No changes needed
       return ExprMutator::VisitVarDef(var);
     }
 
-    // Create a new var with canonicalized strcut_info
+    // Create a new var with canonicalized struct_info
+    Var canonical_var;
     if (var->IsInstance<DataflowVarNode>()) {
-      return DataflowVar(var->vid, canonical_sinfo, var->span);
+      canonical_var = DataflowVar(var->vid, canonical_sinfo, var->span);
+    } else {
+      canonical_var = Var(var->vid, canonical_sinfo, var->span);
     }
-    return Var(var->vid, canonical_sinfo, var->span);
+
+    return ExprMutator::VisitVarDef(canonical_var);
   }
 
  private:
+  /*!
+   * \brief Check if struct_info needs canonicalization
+   */
+  bool NeedsCanonicalization(const StructInfo& sinfo) {
+    if (auto tensor_sinfo = sinfo.as<TensorStructInfoNode>()) {
+      if (!tensor_sinfo->shape.defined()) {
+        return false;
+      }
+      auto shape_expr = tensor_sinfo->shape.as<ShapeExprNode>();
+      if (!shape_expr) {
+        return false;
+      }
+      for (const PrimExpr& dim : shape_expr->values) {
+        if (!IsCanonicalPrimExpr(dim)) {
+          return true;
+        }
+      }
+      return false;
+    } else if (auto tuple_sinfo = sinfo.as<TupleStructInfoNode>()) {
+      for (const StructInfo& field : tuple_sinfo->fields) {
+        if (NeedsCanonicalization(field)) {
+          return true;
+        }
+      }
+      return false;
+    }
+    return false;
+  }
+
+  /*!
+   * Canonicalize struct info and emit necessary bindings
+   */
+  void CanonicalizeStructInfoAndEmit(const StructInfo& sinfo) {
+    if (auto tensor_sinfo = sinfo.as<TensorStructInfoNode>()) {
+      CanonicalizeTensorStructInfoAndEmit(ffi::GetRef<TensorStructInfo>(tensor_sinfo));
+    } else if (auto tuple_sinfo = sinfo.as<TupleStructInfoNode>()) {
+      for (const StructInfo& field : tuple_sinfo->fields) {
+        CanonicalizeStructInfoAndEmit(field);
+      }
+    }
+  }
+
+  /*!
+   * Canonicalize tensor struct info and emit necessary bindings
+   */
+  void CanonicalizeTensorStructInfoAndEmit(const TensorStructInfo& sinfo) {
+    if (!sinfo->shape.defined()) {
+      return;
+    }
+
+    auto shape_expr = sinfo->shape.as<ShapeExprNode>();
+    if (!shape_expr) {
+      return;
+    }
+
+    // Emit bindings for each compound dimension
+    for (const PrimExpr& dim : shape_expr->values) {
+      if (!IsCanonicalPrimExpr(dim)) {
+        CanonicalizeDimension(dim);
+      }
+    }
+  }
+
   /*!
    * \brief Canonicalize struct info by lifting compound shape expressions
    */
@@ -140,7 +240,7 @@ class ShapeExprCanonicalizer : public ExprMutator {
     bool changed = false;
 
     for (const PrimExpr& dim : shape_expr->values) {
-      PrimExpr canonical_dim = CanonicalizeDimension(dim);
+      PrimExpr canonical_dim = GetCanonicalDimension(dim);
       canonical_dims.push_back(canonical_dim);
       changed |= !canonical_dim.same_as(dim);
     }
@@ -174,15 +274,9 @@ class ShapeExprCanonicalizer : public ExprMutator {
   }
 
   /*!
-   * \brief Canonicalize a single shape dimension
-   *
-   * If the dimension is a compound PrimExpr:
-   * 1. Emit a ShapeExpr binding containing the compound expression
-   * 2. Create a fresh symbolic tir::Var
-   * 3. Emit a MatchCast to bind the computed value to the symbolic var
-   * 4. Return the symbolic var
+   * \brief Get the canonical form of a dimension (returns the symbolic var if already emitted)
    */
-  PrimExpr CanonicalizeDimension(const PrimExpr& dim) {
+  PrimExpr GetCanonicalDimension(const PrimExpr& dim) {
     // If already canonical, return as is
     if (IsCanonicalPrimExpr(dim)) {
       return dim;
@@ -193,9 +287,42 @@ class ShapeExprCanonicalizer : public ExprMutator {
       return it->second;
     }
 
-    // Create a fresh symbolic variable
+    // Create a fresh symbolic variable, but don't emit yet
     tir::Var symbolic_var = CreateFreshSymbolicVar(dim->dtype);
 
+    compound_expr_to_var_[dim] = symbolic_var;
+
+    return symbolic_var;
+  }
+
+  /*!
+   * \brief Emit bindings for a single compound dimension
+   *
+   * If the dimension is a compound PrimExpr:
+   * 1. Emit a ShapeExpr binding containing the compound expression
+   * 2. Create a fresh symbolic tir::Var
+   * 3. Emit a MatchCast to bind the computed value to the symbolic var
+   */
+  void CanonicalizeDimension(const PrimExpr& dim) {
+    // If already canonical, nothing to emit
+    if (IsCanonicalPrimExpr(dim)) {
+      return;
+    }
+
+    // Check If we've already emitted bindings for this expression
+    auto it = compound_expr_to_var_.find(dim);
+    if (it == compound_expr_to_var_.end()) {
+      // This should not happen if GetCanonicalDimension was called first
+      return;
+    }
+
+    // Check if we've already emitted the bindings
+    if (emitted_bindings_.count(dim)) {
+      return;
+    }
+
+    tir::Var symbolic_var = it->second;
+
     // Emit shape binding: shape_var = R.shape([compound_expr])
     ShapeExpr shape_value({dim});
     Var shape_var = builder_->Emit(shape_value);
@@ -206,10 +333,8 @@ class ShapeExprCanonicalizer : public ExprMutator {
     Var match_cast_var("_", match_sinfo);
     builder_->EmitNormalized(MatchCast(match_cast_var, shape_var, match_sinfo));
 
-    // Cache the mapping to avoid duplicate bindings
-    compound_expr_to_var_[dim] = symbolic_var;
-
-    return symbolic_var;
+    // Mark as emitted
+    emitted_bindings_.insert(dim);
   }
 
   /*!
@@ -223,6 +348,9 @@ class ShapeExprCanonicalizer : public ExprMutator {
   // Cache to avoid creating duplicate bindings for the same compound expression
   std::unordered_map<PrimExpr, tir::Var, StructuralHash, StructuralEqual> compound_expr_to_var_;
 
+  // Track which compound expressions have had their bindings emitted
+  std::unordered_set<PrimExpr, StructuralHash, StructuralEqual> emitted_bindings_;
+
   // Counter for generating unique symbolic variable names
   int symbolic_var_counter_ = 0;
 };