merge pr from apache#16560

Author: LeiWang199

src/tir/schedule/primitive/blockize_tensorize.cc

@@ -19,7 +19,7 @@
 #include <tvm/tir/data_type_rewriter.h>
 
 #include <functional>
-
+#include "../../transforms/simplify.h"
 #include "../ir_comparator.h"
 #include "../utils.h"
 
@@ -530,7 +530,8 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int
                << GetRef<Stmt>(sref->stmt);
     throw;
   }
-  PrimFunc intrin_desc = intrin->desc;
+  arith::Analyzer analyzer;
+  PrimFunc intrin_desc = Simplify(intrin->desc, &analyzer);
   PrimFunc intrin_impl = DeepCopy(intrin->impl);
 
   int index_dtype_bits = -1;

src/tir/transforms/simplify.cc

@@ -30,7 +30,7 @@
 #include <tvm/tir/transform.h>
 
 #include <optional>
-
+#include "../../tir/transforms/simplify.h"
 #include "../../arith/ir_mutator_with_analyzer.h"
 #include "../../tir/analysis/control_flow_graph.h"
 
@@ -270,6 +270,13 @@ class StmtSimplifier : public IRMutatorWithAnalyzer {
 
 namespace tir {
 
+
+PrimFunc Simplify(PrimFunc func, arith::Analyzer* analyzer) {
+  auto* n = func.CopyOnWrite();
+  n->body = arith::StmtSimplifier::Apply(std::move(n->body), analyzer);
+  return func;
+}
+
 Stmt Simplify(Stmt stmt, arith::Analyzer* analyzer) {
   return arith::StmtSimplifier::Apply(stmt, analyzer);
 }

src/tir/transforms/simplify.h

@@ -30,6 +30,12 @@
 namespace tvm {
 namespace tir {
 
+/* \brief Simplifies the prim func
+ *
+ * Applies the same behavior as the tir.transform.Simplify pass.
+ */
+PrimFunc Simplify(PrimFunc stmt, arith::Analyzer* analyzer);
+
 /* \brief Simplifies the statement
  *
  * Applies the same behavior as the tir.transform.Simplify pass, but