Benchmark with/without the BoundUsingReciprocal function

Author: Lunderberg

src/arith/const_int_bound.cc

@@ -28,6 +28,7 @@
 #include <algorithm>
 #include <optional>
 
+#include "../support/utils.h"
 #include "constraint_extract.h"
 #include "int_operator.h"
 #include "pattern_match.h"
@@ -240,8 +241,10 @@ class ConstIntBoundAnalyzer::Impl
     ret.min_value = InfAwareAdd(a.min_value, b.min_value);
     ret.max_value = InfAwareAdd(a.max_value, b.max_value);
 
-    if (auto bound = BoundUsingReciprocal(GetRef<PrimExpr>(op))) {
-      ret = Intersect(ret, bound.value());
+    if (support::BoolEnvironmentVar("TVM_ENABLE_RECIPROCAL_PATTERN_MATCH")) {
+      if (auto bound = BoundUsingReciprocal(GetRef<PrimExpr>(op))) {
+        ret = Intersect(ret, bound.value());
+      }
     }
 
     return ret;
@@ -254,11 +257,13 @@ class ConstIntBoundAnalyzer::Impl
     ret.min_value = InfAwareAdd(a.min_value, -b.max_value);
     ret.max_value = InfAwareAdd(a.max_value, -b.min_value);
 
-    if (auto bound = BoundUsingReciprocal(GetRef<Sub>(op))) {
-      ret = Intersect(ret, bound.value());
-    }
-    if (auto bound = BoundUsingReciprocal(Sub(op->b, op->a))) {
-      ret = Intersect(ret, Negative(bound.value()));
+    if (support::BoolEnvironmentVar("TVM_ENABLE_RECIPROCAL_PATTERN_MATCH")) {
+      if (auto bound = BoundUsingReciprocal(GetRef<Sub>(op))) {
+        ret = Intersect(ret, bound.value());
+      }
+      if (auto bound = BoundUsingReciprocal(Sub(op->b, op->a))) {
+        ret = Intersect(ret, Negative(bound.value()));
+      }
     }
     return ret;
   }

tests/python/arith/test_arith_const_int_bound.py

@@ -43,8 +43,33 @@ def __name__(self):
         return str(self.expr)
 
 
+with_reciprocal_pattern_match = tvm.testing.parameter(
+    by_dict={
+        "with_updated_const_int_analyzer": True,
+        "without_updated_const_int_analyzer": False,
+    }
+)
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def set_reciprocal_pattern_match(with_reciprocal_pattern_match):
+    import os
+
+    var_name = "TVM_ENABLE_RECIPROCAL_PATTERN_MATCH"
+    old_value = os.environ.get(var_name)
+    os.environ[var_name] = str(int(with_reciprocal_pattern_match))
+    yield
+
+    if old_value is None:
+        del os.environ[var_name]
+    else:
+        os.environ = old_value
+
+
 class BaseCompare:
-    def test_const_bounds(self, test_case):
+    def test_const_bounds(self, test_case, benchmark):
         analyzer = tvm.arith.Analyzer()
 
         for var, bounds in test_case.known_bounds.items():
@@ -54,7 +79,7 @@ def test_const_bounds(self, test_case):
             if test_case.constraint is not None:
                 stack.enter_context(analyzer.constraint_scope(test_case.constraint))
 
-            bounds = analyzer.const_int_bound(test_case.expr)
+            bounds = benchmark(analyzer.const_int_bound, test_case.expr)
 
         if test_case.expected_bounds[0] is None:
             assert bounds.max_value == test_case.expected_bounds[1]

tests/python/arith/test_arith_rewrite_simplify.py

@@ -50,18 +50,43 @@ def __name__(self):
         return str(self.before)
 
 
+with_reciprocal_pattern_match = tvm.testing.parameter(
+    by_dict={
+        "with_updated_const_int_analyzer": True,
+        "without_updated_const_int_analyzer": False,
+    }
+)
+
+
+@pytest.fixture(autouse=True)
+def set_reciprocal_pattern_match(with_reciprocal_pattern_match):
+    import os
+
+    var_name = "TVM_ENABLE_RECIPROCAL_PATTERN_MATCH"
+    old_value = os.environ.get(var_name)
+    os.environ[var_name] = str(int(with_reciprocal_pattern_match))
+    yield
+
+    if old_value is None:
+        del os.environ[var_name]
+    else:
+        os.environ = old_value
+
+
 class BaseCompare:
-    def test_simplify(self, test_case):
+    def test_simplify(self, test_case, benchmark):
         analyzer = tvm.arith.Analyzer()
 
         if inspect.isclass(test_case.expected) and issubclass(test_case.expected, Exception):
             with pytest.raises(test_case.expected):
                 with analyzer.constraint_scope(test_case.constraint):
-                    analyzer.rewrite_simplify(test_case.before)
+                    # analyzer.rewrite_simplify(test_case.before)
+                    benchmark(analyzer.rewrite_simplify, test_case.before)
         else:
 
             with analyzer.constraint_scope(test_case.constraint):
-                after = analyzer.rewrite_simplify(test_case.before)
+                # after = analyzer.rewrite_simplify(test_case.before)
+                after = benchmark(analyzer.rewrite_simplify, test_case.before)
 
             assert tvm.ir.structural_equal(after, test_case.expected), (
                 f"Rewrite didn't match expected.\n"