[ARITH] Allow Analyzer to MarkGlobalPositiveValue

This PR introduces an utility function MarkGlobalPositiveValue.
This function allows analyzer to mark buffer shapes in function arguments
as positive globally and opens doors for more symbolic simplification.

Author: tqchen

include/tvm/arith/analyzer.h

@@ -618,6 +618,22 @@ class TVM_DLL Analyzer {
   TransitiveComparisonAnalyzer transitive_comparisons;
   /*! \brief constructor */
   Analyzer();
+  /*!
+   * \brief Mark the value as positive value globally in analyzer.
+   *
+   * Only call this function if the positive condition is global and
+   * not context-dependent.
+   *
+   * This function does best-effort propagations to the sub-analyzers
+   *
+   * \note We expose this function because positive global values,
+   * such as symbolic buffer shapes in function arguments are really
+   * important to ensure the best simplification, and usually they
+   * can be handled in a simpler way than the generic constraints.
+   *
+   * This function may call into the Update function of the sub-analyzers.
+   */
+  void MarkGlobalPositiveValue(const PrimExpr& value);
   /*!
    * \brief Notify all the sub-analyzers that var
    *        is created and binded to expr.

src/arith/analyzer.cc

@@ -25,6 +25,7 @@
 #include <tvm/tir/expr.h>
 #include <tvm/tir/op.h>
 
+#include "const_fold.h"
 #include "product_normal_form.h"
 
 namespace tvm {
@@ -63,6 +64,35 @@ void Analyzer::Bind(const Var& var, const Range& range, bool allow_override) {
   // skip rewrite simplify
 }
 
+void Analyzer::MarkGlobalPositiveValue(const PrimExpr& value) {
+  // split out the symbolic and non-symbolic part
+  int64_t cscale = 1;
+  PrimExpr symbolic = tir::make_const(value.dtype(), 1);
+  auto fcollect = [&](PrimExpr val) {
+    if (const auto* intimm = val.as<IntImmNode>()) {
+      cscale *= intimm->value;
+    } else {
+      symbolic = symbolic * val;
+    }
+  };
+  UnpackReduction<tir::MulNode>(value, fcollect);
+  if (cscale <= 0) return;
+  // override the constant int bound by marking it as non-negative
+  // NOTE: there might be future opportunities of more bound hint
+  // this is a simple step and covers all the current needs
+  //
+  // We may consider enhance the sub analyzer to directly take
+  // MarkPositiveVar so their bounds do not overlap
+  if (const auto* var_ptr = symbolic.as<VarNode>()) {
+    Var var = GetRef<Var>(var_ptr);
+    bool allow_override = true;
+    // mark the constant bound is sufficient
+    this->const_int_bound.Update(var, ConstIntBound(1, ConstIntBound::kPosInf), allow_override);
+    this->int_set.Update(var, IntSet::Interval(tir::make_const(var.dtype(), 1), pos_inf()),
+                         allow_override);
+  }
+}
+
 void Analyzer::Bind(const Map<Var, Range>& variables, bool allow_override) {
   for (const auto& iter : variables) {
     this->Bind(iter.first, iter.second, allow_override);

src/arith/ir_mutator_with_analyzer.cc

@@ -30,6 +30,15 @@ namespace arith {
 
 using namespace tir;
 
+void IRMutatorWithAnalyzer::MarkBufferMapShapes(const tir::PrimFunc& func) {
+  // Mark the all the symbolic buffer shape values in the buffer map as positive value.
+  for (auto kv : func->buffer_map) {
+    for (PrimExpr shape : kv.second->shape) {
+      analyzer_->MarkGlobalPositiveValue(shape);
+    }
+  }
+}
+
 Stmt IRMutatorWithAnalyzer::VisitStmt_(const ForNode* op) {
   // record the loop variable as iterators
   Range dom = Range::FromMinExtent(op->min, op->extent);

src/arith/ir_mutator_with_analyzer.h

@@ -62,6 +62,14 @@ class IRMutatorWithAnalyzer : public tir::StmtExprMutator {
   PrimExpr VisitExpr_(const tir::ReduceNode* op) override;
 
  protected:
+  /*!
+   * \brief Mark the all the buffer shape values in the buffer map as positive value.
+   *
+   * \note call this function before Visit function's body to maximize
+   *       simplification efficiency
+   */
+  void MarkBufferMapShapes(const tir::PrimFunc& func);
+
   /*! \brief internal analyzer field. */
   Analyzer* analyzer_;
   // the following two fields are useful in case we want

src/tir/transforms/flatten_buffer.cc

@@ -42,6 +42,7 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer {
     arith::Analyzer ana;
     auto pass = BufferFlattener(&ana);
     auto writer = func.CopyOnWrite();
+    pass.MarkBufferMapShapes(func);
     writer->body = pass.VisitStmt(func->body);
     // The buffers in func->buffer_map are deliberately left
     // unflattened, as they are used for validation of user-provided

src/tir/transforms/simplify.cc

@@ -142,20 +142,24 @@ TVM_REGISTER_PASS_CONFIG_OPTION("tir.Simplify", SimplifyConfig);
 
 class StmtSimplifier : public IRMutatorWithAnalyzer {
  public:
-  static Stmt Apply(Stmt stmt, Analyzer* analyzer, Optional<SimplifyConfig> config_opt = NullOpt) {
+  static PrimFunc Apply(PrimFunc func, Analyzer* analyzer,
+                        Optional<SimplifyConfig> config_opt = NullOpt) {
     auto config = config_opt.value_or(AttrsWithDefaultValues<arith::SimplifyConfig>());
     analyzer->rewrite_simplify.SetEnabledExtensions(config->GetEnabledExtensions());
 
     std::optional<ControlFlowGraph> touch_pattern = std::nullopt;
     if (config->propagate_knowns_to_prove_conditional ||
         config->propagate_knowns_to_simplify_expressions) {
-      touch_pattern = ControlFlowGraph(stmt);
+      touch_pattern = ControlFlowGraph(func->body);
     }
 
-    std::unordered_set<const VarNode*> used_in_buffer_def = CollectVarsUsedInBufferDefinition(stmt);
+    std::unordered_set<const VarNode*> used_in_buffer_def =
+        CollectVarsUsedInBufferDefinition(func->body);
     StmtSimplifier simplifier(analyzer, config, std::move(touch_pattern),
                               std::move(used_in_buffer_def));
-    return simplifier(std::move(stmt));
+    simplifier.MarkBufferMapShapes(func);
+    func.CopyOnWrite()->body = simplifier(func->body);
+    return func;
   }
 
  private:
@@ -335,21 +339,14 @@ class StmtSimplifier : public IRMutatorWithAnalyzer {
 }  // namespace arith
 
 namespace tir {
-
-Stmt Simplify(Stmt stmt, arith::Analyzer* analyzer) {
-  return arith::StmtSimplifier::Apply(stmt, analyzer);
-}
-
 namespace transform {
 
 Pass Simplify() {
   auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
     arith::Analyzer analyzer;
     auto cfg = ctx->GetConfig<arith::SimplifyConfig>("tir.Simplify");
 
-    auto* n = f.CopyOnWrite();
-    n->body = arith::StmtSimplifier::Apply(std::move(n->body), &analyzer, cfg);
-    return f;
+    return arith::StmtSimplifier::Apply(f, &analyzer, cfg);
   };
   return CreatePrimFuncPass(pass_func, 0, "tir.Simplify", {});
 }

tests/python/unittest/test_tir_transform_simplify.py

@@ -1733,5 +1733,21 @@ def before(A_ptr: T.handle("float32"), A_stride: T.int32):
     expected = before
 
 
+class TestBufferShapeConstraint(BaseBeforeAfter):
+    """If enabled, rewrite boolean expressions into AND of OR"""
+
+    convert_boolean_to_and_of_ors = True
+
+    def before(a: T.handle):
+        n = T.int64()
+        A = T.match_buffer(a, (n * 32,), "float32")
+        A[T.min(T.int64(0), n)] = T.float32(0)
+
+    def expected(a: T.handle):
+        n = T.int64()
+        A = T.match_buffer(a, (n * 32,), "float32")
+        A[T.int64(0)] = T.float32(0)
+
+
 if __name__ == "__main__":
     tvm.testing.main()