This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 9a4ae22c0b [REFACTOR][Sym] Simplify symbolic analysis dependencies
(#20694)
9a4ae22c0b is described below
commit 9a4ae22c0b9b6c7169149c17f6dcf6481f0daf1c
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 12 22:42:12 2026 -0400
[REFACTOR][Sym] Simplify symbolic analysis dependencies (#20694)
Expose an analyzer-aware `sym::Intersect` overload and use public
`IntSet` operations in loop partitioning. Replace native GCD/LCM helpers
with the standard library while retaining BigInt handling and localizing
extended Euclidean coefficients to modular-set analysis.
Use ordinary node checks for the S-TIR pattern consumers and direct
analyzer simplification in the Relax DFPattern matcher. Retire S-TIR
`UseAssumeToReduceBranches` and its public entry points, removing the
remaining private constraint-extraction dependency.
Remove the unused PresburgerSet C++/Python API and its MLIR build
plumbing while preserving optional Analyzer/Z3 support.
---
CMakeLists.txt | 1 -
cmake/config.cmake | 4 -
cmake/modules/LLVM.cmake | 3 -
cmake/utils/FindLLVM.cmake | 12 -
include/tvm/s_tir/transform.h | 6 -
include/tvm/sym/int_set.h | 9 +
python/tvm/ir/json_compact.py | 1 -
python/tvm/s_tir/transform/transform.py | 11 -
python/tvm/sym/__init__.py | 1 -
python/tvm/sym/int_set.py | 8 -
src/relax/ir/dataflow_matcher.cc | 33 +-
src/relax/ir/dataflow_matcher.h | 21 +-
src/s_tir/analysis/conditional_bounds.cc | 13 +-
src/s_tir/schedule/primitive/compute_at.cc | 42 +-
src/s_tir/schedule/primitive/decompose_padding.cc | 7 +-
src/s_tir/schedule/primitive/rolling_buffer.cc | 31 +-
src/s_tir/schedule/utils.h | 30 +-
src/s_tir/transform/compact_tensor_region.cc | 12 +-
src/s_tir/transform/loop_partition.cc | 40 +-
.../transform/using_assume_to_reduce_branches.cc | 422 --------------
src/sym/const_int_bound.cc | 5 +-
src/sym/int_operator.h | 74 +--
src/sym/int_set.cc | 4 +
src/sym/modular_set.cc | 51 +-
src/sym/presburger_set.cc | 293 ----------
src/sym/presburger_set.h | 201 -------
src/sym/rewrite_simplify.cc | 7 +-
...eliminate_pad_branch_using_buffer_assumption.py | 646 ---------------------
28 files changed, 173 insertions(+), 1815 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 1b2c4d764e..7a8663ffb2 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -69,7 +69,6 @@ tvm_option(USE_HEXAGON_EXTERNAL_LIBS "Path to git repo
containing external Hexag
tvm_option(USE_RPC "Build with RPC" ON)
tvm_option(USE_THREADS "Build with thread support" ON)
tvm_option(USE_LLVM "Build with LLVM, can be set to specific llvm-config path"
OFF)
-tvm_option(USE_MLIR "Build with MLIR support" OFF)
tvm_option(USE_OPENMP "Build with OpenMP thread pool implementation" OFF)
tvm_option(TVM_DEBUG_WITH_ABI_CHANGE "Enable debug code that may cause ABI
changes" OFF)
tvm_option(TVM_LOG_BEFORE_THROW "Whether log before throw, for debugging
purposes" OFF)
diff --git a/cmake/config.cmake b/cmake/config.cmake
index 068b17fbf0..83e81aa137 100644
--- a/cmake/config.cmake
+++ b/cmake/config.cmake
@@ -131,10 +131,6 @@ set(USE_IOS_RPC OFF)
# - /path/to/llvm-config: enable specific LLVM when multiple llvm-dev is
available.
set(USE_LLVM OFF)
-# Whether use MLIR to help analyze, requires USE_LLVM is enabled
-# Possible values: ON/OFF
-set(USE_MLIR OFF)
-
#---------------------------------------------
# Contrib libraries
#---------------------------------------------
diff --git a/cmake/modules/LLVM.cmake b/cmake/modules/LLVM.cmake
index a7fa4233ac..828526d8dc 100644
--- a/cmake/modules/LLVM.cmake
+++ b/cmake/modules/LLVM.cmake
@@ -51,9 +51,6 @@ if(NOT ${USE_LLVM} MATCHES ${IS_FALSE_PATTERN})
message(STATUS "Set TVM_LLVM_VERSION=" ${TVM_LLVM_VERSION})
# Set flags that are only needed for LLVM target
add_definitions(-DTVM_LLVM_VERSION=${TVM_LLVM_VERSION})
- if (${TVM_MLIR_VERSION})
- add_definitions(-DTVM_MLIR_VERSION=${TVM_MLIR_VERSION})
- endif()
add_definitions(-DTVM_LLVM_HAS_AARCH64_TARGET=${TVM_LLVM_HAS_AARCH64_TARGET})
tvm_file_glob(GLOB COMPILER_LLVM_SRCS
src/target/llvm/*.cc
diff --git a/cmake/utils/FindLLVM.cmake b/cmake/utils/FindLLVM.cmake
index 1f54ded1d5..fc6796a361 100644
--- a/cmake/utils/FindLLVM.cmake
+++ b/cmake/utils/FindLLVM.cmake
@@ -175,18 +175,6 @@ macro(find_llvm use_llvm)
if("AArch64" IN_LIST BUILT_TARGET_LIST)
set(TVM_LLVM_HAS_AARCH64_TARGET 1)
endif()
- if (${USE_MLIR})
- if (EXISTS "${__llvm_libdir}/libMLIRPresburger.a")
- if (EXISTS "${__llvm_libdir}/libMLIRSupport.a")
- message(STATUS "Found MLIR")
- list(APPEND LLVM_LIBS "${__llvm_libdir}/libMLIRPresburger.a")
- list(APPEND LLVM_LIBS "${__llvm_libdir}/libMLIRSupport.a")
- set(TVM_MLIR_VERSION ${TVM_LLVM_VERSION})
- message(STATUS "Build with MLIR")
- message(STATUS "Set TVM_MLIR_VERSION=" ${TVM_MLIR_VERSION})
- endif()
- endif()
- endif()
separate_arguments(__llvm_system_libs)
foreach(__flag IN ITEMS ${__llvm_system_libs})
if("${__flag}" STREQUAL "-lm")
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index a1ce06e6e5..4bc593d170 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -298,12 +298,6 @@ TVM_DLL Pass RemoveWeightLayoutRewriteBlock(bool
skip_tensor_rewrite = false);
*/
TVM_DLL Pass RemoveStoreUndef();
-/*!
- * \brief Eliminate branches by leveraging tensor assumptions (T.assume).
- * \return The pass.
- */
-TVM_DLL Pass UseAssumeToReduceBranches();
-
} // namespace transform
} // namespace s_tir
} // namespace tvm
diff --git a/include/tvm/sym/int_set.h b/include/tvm/sym/int_set.h
index d5a9c3f6be..10132b4f9c 100644
--- a/include/tvm/sym/int_set.h
+++ b/include/tvm/sym/int_set.h
@@ -262,6 +262,15 @@ ffi::Array<IntSet> UnionRegionLowerBound(const
ffi::Array<ffi::Array<IntSet>>& n
*/
IntSet Intersect(const ffi::Array<IntSet>& sets);
+/*!
+ * \brief Intersect two integer sets using the analyzer's bindings and
constraints.
+ * \param lhs The first set.
+ * \param rhs The second set.
+ * \param analyzer The analyzer used to prove whether the intersection is
empty.
+ * \return Nothing if emptiness is proved, or the intersected bounds otherwise.
+ */
+TVM_TIRX_DLL IntSet Intersect(const IntSet& lhs, const IntSet& rhs, const
Analyzer& analyzer);
+
/*!
* \brief Converts the Ranges to IntSets
* \param var_dom The ranges of variables
diff --git a/python/tvm/ir/json_compact.py b/python/tvm/ir/json_compact.py
index 8edc672588..66a76c5cfe 100644
--- a/python/tvm/ir/json_compact.py
+++ b/python/tvm/ir/json_compact.py
@@ -29,7 +29,6 @@ _PRIM_TYPE_KEY_RENAMES = {
"arith.IterSplitExpr": "sym.IterSplitExpr",
"arith.IterSumExpr": "sym.IterSumExpr",
"arith.ModularSet": "sym.ModularSet",
- "arith.PresburgerSet": "sym.PresburgerSet",
"arith.RewriteSimplifierStats": "sym.RewriteSimplifierStats",
"arith.SplitExpr": "sym.SplitExpr",
"arith.SumExpr": "sym.SumExpr",
diff --git a/python/tvm/s_tir/transform/transform.py
b/python/tvm/s_tir/transform/transform.py
index 9130ebcee2..991ccecde9 100644
--- a/python/tvm/s_tir/transform/transform.py
+++ b/python/tvm/s_tir/transform/transform.py
@@ -399,14 +399,3 @@ def RemoveStoreUndef():
The result pass
"""
return _ffi_api.RemoveStoreUndef() # type: ignore
-
-
-def UseAssumeToReduceBranches():
- """Eliminate layout specific pad branch by overcomputing values for padded
region.
-
- Returns
- -------
- fpass : tvm.transform.Pass
- The result pass
- """
- return _ffi_api.UseAssumeToReduceBranches() # type: ignore
diff --git a/python/tvm/sym/__init__.py b/python/tvm/sym/__init__.py
index 0f76bc14db..3a3e091eff 100644
--- a/python/tvm/sym/__init__.py
+++ b/python/tvm/sym/__init__.py
@@ -20,7 +20,6 @@
from .int_set import (
IntSet,
IntervalSet,
- PresburgerSet,
estimate_region_lower_bound,
estimate_region_strict_bound,
estimate_region_upper_bound,
diff --git a/python/tvm/sym/int_set.py b/python/tvm/sym/int_set.py
index 2b2307cd10..5e93f6362e 100644
--- a/python/tvm/sym/int_set.py
+++ b/python/tvm/sym/int_set.py
@@ -85,14 +85,6 @@ class IntervalSet(IntSet):
self.__init_handle_by_constructor__(_ffi_api.IntervalSet, min_value,
max_value)
-@tvm_ffi.register_object("sym.PresburgerSet")
-class PresburgerSet(IntSet):
- """Represent of Presburger Set"""
-
- def __init__(self):
- self.__init_handle_by_constructor__(_ffi_api.PresburgerSet)
-
-
def estimate_region_lower_bound(region, var_dom, predicate, analyzer=None):
"""Analyze the region with affine map, given the domain of variables and
their predicate
Some subregion may be discarded during the lower-bound analysis.
diff --git a/src/relax/ir/dataflow_matcher.cc b/src/relax/ir/dataflow_matcher.cc
index 8ec76a33d2..93d5d09450 100644
--- a/src/relax/ir/dataflow_matcher.cc
+++ b/src/relax/ir/dataflow_matcher.cc
@@ -48,7 +48,6 @@
#include <utility>
#include <vector>
-#include "../../sym/constraint_extract.h"
#include "../transform/utils.h"
namespace tvm {
@@ -442,7 +441,7 @@ bool DFPatternMatcher::VisitDFPattern_(const
TypePatternNode* op, const Expr& ex
return static_cast<bool>(as_int->value);
}
- symbolic_expr_condition_ = SimplifyCondition(symbolic_expr_condition_ &&
new_constraint);
+ symbolic_expr_condition_ = analyzer_->Simplify(symbolic_expr_condition_ &&
new_constraint);
if (auto* as_int = symbolic_expr_condition_.as<IntImmNode>()) {
return static_cast<bool>(as_int->value);
@@ -451,36 +450,6 @@ bool DFPatternMatcher::VisitDFPattern_(const
TypePatternNode* op, const Expr& ex
}
}
-PrimExpr DFPatternMatcher::SimplifyCondition(PrimExpr condition) {
- if (condition->IsInstance<IntImmNode>()) {
- return condition;
- }
-
- std::vector<PrimExpr> constraints = sym::ExtractConstraints(condition,
false);
- if (constraints.size() == 1) {
- return condition;
- }
-
- auto sort_key = [](PrimExpr expr) -> ffi::String {
- if (const auto* equal = expr.as<prim::EQNode>()) {
- if (auto var = equal->a.as<PrimVar>()) {
- return var.value()->name;
- }
- }
- return "";
- };
- std::stable_sort(
- constraints.begin(), constraints.end(),
- [&sort_key](const PrimExpr& a, const PrimExpr& b) { return sort_key(a) <
sort_key(b); });
-
- PrimExpr sorted_condition = IntImm::Bool(true);
- for (const PrimExpr& constraint : constraints) {
- sorted_condition = sorted_condition && constraint;
- }
-
- return analyzer_->Simplify(sorted_condition);
-}
-
static bool ShapeEqual(AnalyzerObj* analyzer, const ffi::Array<PrimExpr>& lhs,
const ffi::Array<PrimExpr>& rhs) {
if (lhs.size() != rhs.size()) return false;
diff --git a/src/relax/ir/dataflow_matcher.h b/src/relax/ir/dataflow_matcher.h
index 1869b2a4d3..677b5526b8 100644
--- a/src/relax/ir/dataflow_matcher.h
+++ b/src/relax/ir/dataflow_matcher.h
@@ -41,8 +41,12 @@ class DFPatternMatcher : public DFPatternFunctor<bool(const
DFPattern&, const Ex
public:
using var2val_t = ffi::Map<Var, Expr>;
- explicit DFPatternMatcher() {}
- explicit DFPatternMatcher(var2val_t var2val) : var2val_(std::move(var2val))
{}
+ explicit DFPatternMatcher() : DFPatternMatcher(var2val_t{}) {}
+ explicit DFPatternMatcher(var2val_t var2val) : var2val_(std::move(var2val)) {
+ // Propagate shape constraints across conjunctions regardless of operand
order.
+ analyzer_->rewrite_simplify.SetEnabledExtensions(
+ sym::RewriteSimplifier::kApplyConstraintsToBooleanBranches);
+ }
bool Match(const DFPattern& pattern, const Expr& expr);
ffi::Map<DFPattern, Expr> GetMemo() { return memo_; }
@@ -78,19 +82,6 @@ class DFPatternMatcher : public DFPatternFunctor<bool(const
DFPattern&, const Ex
const tvm::ffi::Array<Expr> fields,
std::vector<int8_t>& match_cache,
std::vector<bool>& matched);
- /* \brief Simplify a boolean condition using the analyzer
- *
- * Matching type can often produce conditions that do not
- * simplify cleanly. For example, while the rewrite simplifier can
- * recognize that `m==0 && m==1` can be simplifies to `false`, it
- * cannot recognize that `m==0 && n==0 && m==1` can be simplified to
- * false.
- *
- * This function applies additional simplification steps to handle
- * these cases, before delgating to `analyzer_.Simplify`.
- */
- PrimExpr SimplifyCondition(PrimExpr condition);
-
std::unordered_map<DFPattern, Expr, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
memo_;
var2val_t var2val_;
std::vector<DFPattern> matched_nodes_;
diff --git a/src/s_tir/analysis/conditional_bounds.cc
b/src/s_tir/analysis/conditional_bounds.cc
index c5d2340610..a5de8ec221 100644
--- a/src/s_tir/analysis/conditional_bounds.cc
+++ b/src/s_tir/analysis/conditional_bounds.cc
@@ -33,10 +33,9 @@
#include <algorithm>
#include <functional>
+#include <numeric>
#include <utility>
-#include "../../sym/int_operator.h"
-
namespace tvm {
namespace s_tir {
using namespace tvm::prim;
@@ -48,9 +47,6 @@ using sym::EvalSet;
using sym::IntSet;
namespace {
-using sym::ExtendedEuclidean;
-using sym::LeastCommonMultiple;
-
// The solver's intermediate representations remain local to this analysis.
struct IntGroupBounds {
PrimExpr coef;
@@ -418,10 +414,9 @@ PartialSolvedInequalities SolveLinearInequalities(const
IntConstraints& system_t
&coef_neg, analyzer.get());
// Combine each positive inequality with each negative one (by adding them
together)
- int64_t gcd_x, gcd_y;
for (const auto& pos : coef_pos) {
for (const auto& neg : coef_neg) {
- auto first_gcd = ExtendedEuclidean(pos.first, -neg.first, &gcd_x,
&gcd_y);
+ auto first_gcd = std::gcd(pos.first, -neg.first);
PrimType v_ty = v.ty();
PrimExpr c_pos = prim::MakeConst(v_ty, neg.first / first_gcd);
PrimExpr c_neg = IntImm(v_ty, pos.first / first_gcd);
@@ -443,10 +438,10 @@ PartialSolvedInequalities SolveLinearInequalities(const
IntConstraints& system_t
// We will generate formulas of the form coef_lcm*v <= bound
int64_t coef_lcm = 1;
for (const auto& pos : coef_pos) {
- coef_lcm = LeastCommonMultiple(coef_lcm, pos.first);
+ coef_lcm = std::lcm(coef_lcm, pos.first);
}
for (const auto& neg : coef_neg) {
- coef_lcm = LeastCommonMultiple(coef_lcm, -neg.first);
+ coef_lcm = std::lcm(coef_lcm, -neg.first);
}
// The resulting lower and upper bounds
diff --git a/src/s_tir/schedule/primitive/compute_at.cc
b/src/s_tir/schedule/primitive/compute_at.cc
index 6f7439ede9..89abb7330c 100644
--- a/src/s_tir/schedule/primitive/compute_at.cc
+++ b/src/s_tir/schedule/primitive/compute_at.cc
@@ -438,23 +438,29 @@ std::pair<Var, BlockVarDomainInfo>
SolveBlockVarDomain(const sym::IntSet& provid
PrimExpr required_max = analyzer->Simplify(required.max());
sym::IntSet var_dom{ffi::UnsafeInit{}}, var_bound{ffi::UnsafeInit{}};
ffi::Optional<Var> var;
- sym::PVar<Var> p_v;
- sym::PVar<PrimExpr> p_e;
- if ((p_v * p_e).Match(provided_min) || (p_e * p_v).Match(provided_min)) {
- PrimExpr e = p_e.Eval();
- var = p_v.Eval();
+ ffi::Optional<PrimExpr> factor;
+ if (const auto* mul = provided_min.as<prim::MulNode>()) {
+ if (auto lhs = mul->a.as<Var>()) {
+ var = lhs;
+ factor = mul->b;
+ } else if (auto rhs = mul->b.as<Var>()) {
+ var = rhs;
+ factor = mul->a;
+ }
+ }
+ if (factor.has_value()) {
+ PrimExpr e = factor.value();
var_dom = sym::IntSet::Interval(floordiv(required_min, e),
floordiv(required_max, e));
var_bound = sym::IntSet::Interval(0, floordiv(dim_max, e));
} else if (analyzer->CanProveEqual(provided_min, provided_max)) {
- if (p_v.Match(provided_min)) {
- var = p_v.Eval();
+ if (auto single_var = provided_min.as<Var>()) {
+ var = single_var;
var_dom = sym::IntSet::Interval(required_min, required_max);
var_bound = sym::IntSet::Interval(0, dim_max);
} else {
- sym::PVar<PrimExpr> p_f1, p_f2;
- if ((floordiv(p_f1, p_f2).Match(provided_min))) {
- PrimExpr var_expr = p_f1.Eval();
- PrimExpr fac = p_f2.Eval();
+ if (const auto* div = provided_min.as<prim::FloorDivNode>()) {
+ PrimExpr var_expr = div->a;
+ PrimExpr fac = div->b;
if (analyzer->CanProveGreaterEqual(fac, 1)) {
if (var_expr.as<PrimVar>()) {
// a <= (x // factor) <= b, fac > 0 ==> (a * fac) <= x <= (b * fac
+ fac - 1)
@@ -463,27 +469,27 @@ std::pair<Var, BlockVarDomainInfo>
SolveBlockVarDomain(const sym::IntSet& provid
analyzer->Simplify(required_max *
fac + fac - 1));
var_bound = sym::IntSet::Interval(0, analyzer->Simplify(dim_max *
fac + fac - 1));
} else {
- const sym::IntSet new_provided =
sym::IntSet::SinglePoint(p_f1.Eval());
+ const sym::IntSet new_provided = sym::IntSet::SinglePoint(div->a);
const sym::IntSet new_required = sym::IntSet::Interval(
required_min * fac, analyzer->Simplify(required_max * fac +
fac - 1));
return SolveBlockVarDomain(new_provided, new_required, dim_max,
analyzer);
}
}
- } else if ((floormod(p_f1, p_f2).Match(provided_min))) {
- PrimExpr var_expr = p_f1.Eval();
+ } else if (const auto* mod = provided_min.as<prim::FloorModNode>()) {
+ PrimExpr var_expr = mod->a;
if (var_expr.as<PrimVar>()) {
// generally domain of (x % fac) enforce no constraints to domain of
x
Var var_mod = var_expr.as_or_throw<Var>();
return {var_mod, BlockVarDomainInfo()};
} else {
- PrimExpr mod_1 = p_f1.Eval();
- PrimExpr mod_2 = p_f2.Eval();
+ PrimExpr mod_1 = mod->a;
+ PrimExpr mod_2 = mod->b;
if (analyzer->CanProveGreaterEqual(mod_1, 1) &&
analyzer->CanProveGreaterEqual(mod_2, 1)) {
- const sym::IntSet new_provided =
sym::IntSet::SinglePoint(p_f1.Eval());
+ const sym::IntSet new_provided = sym::IntSet::SinglePoint(mod->a);
if (analyzer->CanProveGreaterEqual(required_min, 0)) {
const sym::IntSet new_required =
- sym::IntSet::Interval(required_min,
sym::SymbolicLimits::pos_inf_);
+ sym::IntSet::Interval(required_min,
sym::IntSet::Everything().max());
return SolveBlockVarDomain(new_provided, new_required, dim_max,
analyzer);
}
}
diff --git a/src/s_tir/schedule/primitive/decompose_padding.cc
b/src/s_tir/schedule/primitive/decompose_padding.cc
index bf6a37ca62..544fd875c1 100644
--- a/src/s_tir/schedule/primitive/decompose_padding.cc
+++ b/src/s_tir/schedule/primitive/decompose_padding.cc
@@ -153,10 +153,9 @@ class PaddingInfoAnalyzer {
PrimExpr RewritePredicate(const PrimExpr& predicate) {
PrimExpr res = IntImm::Bool(true);
std::function<void(PrimExpr)> update = [&res, &update](PrimExpr e) {
- sym::PVar<PrimExpr> a, b;
- if ((a && b).Match(e)) {
- update(a.Eval());
- update(b.Eval());
+ if (const auto* and_node = e.as<prim::AndNode>()) {
+ update(and_node->a);
+ update(and_node->b);
} else {
if (const CallNode* call = e.as<CallNode>()) {
if (call->op.same_as(prim::likely_op())) {
diff --git a/src/s_tir/schedule/primitive/rolling_buffer.cc
b/src/s_tir/schedule/primitive/rolling_buffer.cc
index aaf8acc57f..9e4b5cf9b2 100644
--- a/src/s_tir/schedule/primitive/rolling_buffer.cc
+++ b/src/s_tir/schedule/primitive/rolling_buffer.cc
@@ -191,32 +191,39 @@ class RollingBufferInfoCollector {
std::vector<ffi::Optional<Var>> bound_iter_vars;
std::vector<int> bound_overlaps;
- sym::PVar<Var> p_var;
- sym::PVar<IntImm> p_stride, p_divisor;
for (auto bound : region) {
auto stride = 0;
auto divisor = 1;
ffi::Optional<Var> iter_var;
- if (floordiv((p_var * p_stride), p_divisor).Match(bound->min)) {
+ if (const auto* div = bound->min.as<prim::FloorDivNode>()) {
+ const auto* mul = div->a.as<prim::MulNode>();
+ if (!mul) return false;
+ iter_var = mul->a.as<Var>();
+ auto stride_imm = mul->b.as<IntImm>();
+ auto divisor_imm = div->b.as<IntImm>();
+ if (!iter_var.has_value() || !stride_imm.has_value() ||
!divisor_imm.has_value()) {
+ return false;
+ }
// Handle the case of fractional strides
- // They take this form: floordiv(hh.outer, 2)
+ // They take this form: floordiv(hh.outer * stride, divisor)
// Strip the floordiv and keep track of the divisor
- iter_var = p_var.Eval();
- auto divisor_value = p_divisor.Eval()->value.as<int>();
- auto stride_value = p_stride.Eval()->value.as<int>();
+ auto divisor_value = divisor_imm.value()->value.as<int>();
+ auto stride_value = stride_imm.value()->value.as<int>();
if (!divisor_value.has_value() || !stride_value.has_value()) return
false;
divisor = *divisor_value;
stride = std::ceil(static_cast<float>(*stride_value) / divisor);
- } else if ((p_var * p_stride).Match(bound->min)) {
+ } else if (const auto* mul = bound->min.as<prim::MulNode>()) {
// The bound is the iter var multiplied by the stride
- iter_var = p_var.Eval();
- auto stride_value = p_stride.Eval()->value.as<int>();
+ iter_var = mul->a.as<Var>();
+ auto stride_imm = mul->b.as<IntImm>();
+ if (!iter_var.has_value() || !stride_imm.has_value()) return false;
+ auto stride_value = stride_imm.value()->value.as<int>();
if (!stride_value.has_value()) return false;
stride = *stride_value;
- } else if (p_var.Match(bound->min)) {
+ } else if (auto var = bound->min.as<Var>()) {
// If the bound is just a Var, that implies the stride is 1
- iter_var = p_var.Eval();
+ iter_var = var;
stride = 1;
} else if (IsConstInt(bound->min)) {
// If the bound is an int, we can't roll over it
diff --git a/src/s_tir/schedule/utils.h b/src/s_tir/schedule/utils.h
index b6fe4b6342..745d5709f1 100644
--- a/src/s_tir/schedule/utils.h
+++ b/src/s_tir/schedule/utils.h
@@ -45,7 +45,6 @@
#include <utility>
#include "../../runtime/thread_storage_scope.h"
-#include "../../sym/pattern_match.h"
#include "../support/array_utils.h"
#include "../support/nd_int_set.h"
#include "./analysis.h"
@@ -212,18 +211,29 @@ inline ffi::Optional<Var> AnalyzeVarWithShift(const
PrimExpr& expr,
*constant = std::nullopt;
return static_cast<Var>(var.value());
}
- sym::PVar<Var> var;
- sym::PVar<IntImm> shift;
// match: "var + shift"
- if ((var + shift).Match(expr) || (shift + var).Match(expr)) {
- *constant = shift.Eval();
- return var.Eval();
+ if (const auto* add = expr.as<prim::AddNode>()) {
+ if (auto var = add->a.as<Var>()) {
+ if (auto shift = add->b.as<IntImm>()) {
+ *constant = shift;
+ return var;
+ }
+ }
+ if (auto shift = add->a.as<IntImm>()) {
+ if (auto var = add->b.as<Var>()) {
+ *constant = shift;
+ return var;
+ }
+ }
}
// match: "var - shift"
- if ((var - shift).Match(expr)) {
- IntImm result = shift.Eval();
- *constant = IntImm(result.ty(), -result->value);
- return var.Eval();
+ if (const auto* sub = expr.as<prim::SubNode>()) {
+ if (auto var = sub->a.as<Var>()) {
+ if (auto shift = sub->b.as<IntImm>()) {
+ *constant = IntImm(shift.value().ty(), -shift.value()->value);
+ return var;
+ }
+ }
}
return std::nullopt;
}
diff --git a/src/s_tir/transform/compact_tensor_region.cc
b/src/s_tir/transform/compact_tensor_region.cc
index ad30315c40..c45dbc94ff 100644
--- a/src/s_tir/transform/compact_tensor_region.cc
+++ b/src/s_tir/transform/compact_tensor_region.cc
@@ -134,6 +134,12 @@ class TensorAccessRegionCollector : public StmtExprVisitor
{
}
private:
+ static bool IsIndexTypedExpr(const PrimExpr& expr) {
+ PrimType type = expr.ty();
+ return type.IsScalar() && type.MatchesCode(DLDataTypeCode::kDLInt) &&
+ (type.bits() == 32 || type.bits() == 64);
+ }
+
struct TensorAccessInfo {
/*! \brief The tensor. */
TensorVar tensor;
@@ -212,7 +218,7 @@ class TensorAccessRegionCollector : public StmtExprVisitor {
call && call->op.same_as(tirx::decl_tensor_op()))
return StmtExprVisitor::Visit_(op);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->value));
- if (auto value = op->value.as<PrimExpr>(); value &&
sym::IsIndexTypedExpr(value.value())) {
+ if (auto value = op->value.as<PrimExpr>(); value &&
IsIndexTypedExpr(value.value())) {
dom_analyzer_->Bind(op->var, value.value());
dom_map_.emplace(op->var.get(), sym::IntSet::SinglePoint(value.value()));
}
@@ -221,12 +227,12 @@ class TensorAccessRegionCollector : public
StmtExprVisitor {
ffi::Optional<VisitInterrupt> Visit_(const LetNode* op) final {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->value));
- if (sym::IsIndexTypedExpr(op->value)) {
+ if (IsIndexTypedExpr(op->value)) {
dom_analyzer_->Bind(op->var, op->value);
dom_map_.emplace(op->var.get(), sym::IntSet::SinglePoint(op->value));
}
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->body));
- if (sym::IsIndexTypedExpr(op->value)) {
+ if (IsIndexTypedExpr(op->value)) {
dom_map_.erase(op->var.get());
}
return std::nullopt;
diff --git a/src/s_tir/transform/loop_partition.cc
b/src/s_tir/transform/loop_partition.cc
index 23273ab176..18e260b805 100644
--- a/src/s_tir/transform/loop_partition.cc
+++ b/src/s_tir/transform/loop_partition.cc
@@ -44,7 +44,6 @@
#include <unordered_map>
#include <unordered_set>
-#include "../../sym/interval_set.h"
#include "ir_utils.h"
namespace tvm {
@@ -586,7 +585,7 @@ class LoopPartitioner : public StmtExprMutator {
bool partition_thread_scope);
std::pair<IntSet, ExpressionSet> GetIntervalAndCondset(const Partition&
partitions,
- const
sym::IntervalSet& for_interval,
+ const IntSet&
for_interval,
bool cond_value, bool
has_partition_hint);
inline Stmt MakeFor(const ffi::Object* op, PrimExpr extent, Stmt body);
@@ -602,18 +601,19 @@ class LoopPartitioner : public StmtExprMutator {
// Returns an interval (in the first component) in which all the conditions
// given in the second component provably have value given by cond_value
-std::pair<IntSet, ExpressionSet> LoopPartitioner::GetIntervalAndCondset(
- const Partition& partitions, const sym::IntervalSet& for_interval, bool
cond_value,
- bool has_partition_hint) {
+std::pair<IntSet, ExpressionSet> LoopPartitioner::GetIntervalAndCondset(const
Partition& partitions,
+ const
IntSet& for_interval,
+ bool
cond_value,
+ bool
has_partition_hint) {
ffi::Array<IntSet> sets;
ExpressionSet cond_set;
for (const auto& kv : partitions) {
if (kv.first.second == cond_value) {
- sym::IntervalSet interval = kv.second.as_or_throw<sym::IntervalSet>();
- sym::IntervalSet intersection = sym::Intersect(analyzer_.get(),
interval, for_interval);
+ const IntSet& interval = kv.second;
+ IntSet intersection = sym::Intersect(interval, for_interval, analyzer_);
- if (!intersection->IsEmpty()) {
+ if (!intersection.IsNothing()) {
sets.push_back(kv.second);
cond_set.insert(kv.first.first);
}
@@ -624,21 +624,20 @@ std::pair<IntSet, ExpressionSet>
LoopPartitioner::GetIntervalAndCondset(
// Try to find the intersection of the cond_intervals until the intersection
// is nothing when has_partition_hint is true.
if (interval.IsNothing() && has_partition_hint) {
- sym::IntervalSet cond_intersection = sym::IntervalSet::Everything();
+ IntSet cond_intersection = IntSet::Everything();
cond_set.clear();
for (const auto& kv : partitions) {
if (kv.first.second == cond_value) {
- sym::IntervalSet cond_interval =
kv.second.as_or_throw<sym::IntervalSet>();
- sym::IntervalSet intersection =
- sym::Intersect(analyzer_.get(), cond_interval, for_interval);
- if (!intersection->IsEmpty()) {
- cond_intersection = sym::Intersect(analyzer_.get(),
cond_intersection, cond_interval);
+ const IntSet& cond_interval = kv.second;
+ IntSet intersection = sym::Intersect(cond_interval, for_interval,
analyzer_);
+ if (!intersection.IsNothing()) {
+ cond_intersection = sym::Intersect(cond_intersection, cond_interval,
analyzer_);
// Return the latest interval and cond_set if the cond_intersection
is nothing.
- if (!cond_intersection->IsEmpty()) {
+ if (!cond_intersection.IsNothing()) {
cond_set.insert(kv.first.first);
- interval =
sym::IntervalSet(analyzer_->Simplify(cond_intersection->min_value),
-
analyzer_->Simplify(cond_intersection->max_value));
+ interval =
IntSet::Interval(analyzer_->Simplify(cond_intersection.min()),
+
analyzer_->Simplify(cond_intersection.max()));
} else {
break;
}
@@ -711,7 +710,7 @@ ffi::Optional<Stmt> LoopPartitioner::TryPartition(const
Stmt& stmt, Var var, Pri
hint_map_.erase(var.get());
if (finder->partitions.empty()) return std::nullopt;
- sym::IntervalSet for_interval(min, max);
+ IntSet for_interval = IntSet::Interval(min, max);
auto [middle_interval, cond_set, opt_cond_value] =
[&]() -> std::tuple<ffi::Optional<IntSet>, ExpressionSet,
std::optional<bool>> {
@@ -788,7 +787,6 @@ ffi::Optional<Stmt> LoopPartitioner::TryPartition(const
Stmt& stmt, Var var, Pri
}
bool cond_value = opt_cond_value.value();
- IntervalSet middle_interval_i =
middle_interval.value().as_or_throw<IntervalSet>();
// middle_interval is the subrange of the loop variable range for which a
// set of conditions are true (or false resp.)
// The part of the loop variable range that is before (after resp.) that
@@ -799,7 +797,7 @@ ffi::Optional<Stmt> LoopPartitioner::TryPartition(const
Stmt& stmt, Var var, Pri
PrimExpr body_begin = min;
ffi::Optional<Stmt> pre_stmt;
bool pre_stmt_recurse = true;
- if (middle_interval_i->HasLowerBound()) {
+ if (middle_interval.value().HasLowerBound()) {
body_begin = analyzer_->Simplify(middle_interval.value().min());
if (!analyzer_->CanProve(body_begin == min)) {
PrimExpr extent = analyzer_->Simplify(body_begin - min);
@@ -832,7 +830,7 @@ ffi::Optional<Stmt> LoopPartitioner::TryPartition(const
Stmt& stmt, Var var, Pri
PrimExpr post_doubt_begin = max + 1;
ffi::Optional<Stmt> post_stmt;
bool post_stmt_recurse = true;
- if (middle_interval_i->HasUpperBound()) {
+ if (middle_interval.value().HasUpperBound()) {
post_doubt_begin = analyzer_->Simplify(middle_interval.value().max() + 1);
if (!analyzer_->CanProve(middle_interval.value().max() == max)) {
// require the extent to be non-negative
diff --git a/src/s_tir/transform/using_assume_to_reduce_branches.cc
b/src/s_tir/transform/using_assume_to_reduce_branches.cc
deleted file mode 100644
index 7f083e21f9..0000000000
--- a/src/s_tir/transform/using_assume_to_reduce_branches.cc
+++ /dev/null
@@ -1,422 +0,0 @@
-/*
- * 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 using_assume_to_reduce_branches.cc
- *
- * \brief Attempt to remove conditional branch statements by introducing
- * extra computations that do not impact the final results. Mainly
- * oriented for layout specific padding related branches.
- *
- * \note
- * 1. This pass works if the tensor assumption variable is in the branch
statement.
- * In case, the tensor assumption is not present in the branch statement
and
- * there are intermediate tensors then, inline the code.
- * 2. The assumptions leveraged here should be of the form
T.assume(condition_on_indices or
- * tensor_equals_to_some_value)
- * 3. Some part of the code are reused from the control_flow_graph.cc file
which also
- * handles eliminating branches in particular scenarios.
- * 4. This pass currently works for op_pattern kElemWise and kBroadcast.
- */
-
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/relax/expr.h>
-#include <tvm/relax/op_attr_types.h>
-#include <tvm/s_tir/stmt_functor.h>
-#include <tvm/s_tir/stmt_functor_with_analyzer.h>
-#include <tvm/s_tir/transform.h>
-#include <tvm/tirx/function.h>
-
-#include <unordered_map>
-
-#include "../../sym/constraint_extract.h"
-#include "tvm/ir/expr.h"
-namespace tvm {
-namespace s_tir {
-using namespace tvm::tirx;
-
-using namespace sym;
-
-class AssumeChecker : public StmtExprVisitor {
- public:
- using StmtExprVisitor::Visit_;
-
- /* This class checks if the function has assume statement.
- If yes, then only the FuncAnanlyzerMutator class runs. This is to ensure
speedup in the pass.*/
- bool has_assume = false;
-
- ffi::Optional<VisitInterrupt> Visit(ffi::AnyView stmt) final {
- if (has_assume) {
- return std::nullopt;
- }
- return StmtExprVisitor::Visit(stmt);
- }
- ffi::Optional<VisitInterrupt> Visit_(const CallNode* op) override {
- if (op->op.same_as(prim::assume_op())) {
- has_assume = true;
- }
- return std::nullopt;
- }
-};
-
-class ParseAssumeAndOvercompute : public StmtExprMutatorWithAnalyzer {
- public:
- using StmtExprMutatorWithAnalyzer::Mutate;
- using StmtExprMutatorWithAnalyzer::Mutate_;
-
- /* This class analyzes the complete function.
- It parses the tensor assumptions and eliminates the redundant branch
- introduced due to layout specific padding by leveraging from tensor
assumptions.
- On eliminating the branch there are more opportunities to vectorize the code
- and improve performance.
-
- Example:
- -------------
- Function Before :
- for (...)
- T.assume( assume_condition or A[i] == 0 )
- for (...)
- out = T.if_then_else(if_then_else_condition, 0, function(A))
- # here function(A) is some function on Var A
-
- Function After :
- for (...)
- T.assume( assume_condition or A[i] == 0 )
- for (...)
- out = function(A) # here function(A) is some function on the Var A
- --------------
- # High-level implementation details :
- 1. The pass parses the assume statement and stores the relevant
information.
- 2. The pass tries to evaluate the then_clause and else_clause in
then_condition_context
- and else_condition_context.
- It checks if the context of the assume statement (for condition indices and
- assume_condition) is same as the context of the if_then_else statement
(for condition indices
- and if_then_else condition). If context is same and the expression inside
if_then_else statement
- is a function of the tensor assumption (eg A in above example),
- then the pass substitutes the value from the tensor assumption and
simplifies the expression.
- 3. The pass then checks if then_clause and else_clause evaluate to same
value.
- If yes, then return the else_clause if we are in the
then_condition_context (since then_clause
- will be true in this context and if else_clause is also evaluating to true
then we can directly
- replace it with else_clause), similarly, we return the then_clause if we
are in the
- else_condition_context.
- This class handles all these scenarios.*/
-
- using Parent = StmtExprMutatorWithAnalyzer;
- explicit ParseAssumeAndOvercompute(const Analyzer& analyzer) :
Parent(analyzer) {}
-
- private:
- // This struct stores all the relevant data related to asssume statement
- struct assume_struct { // Consider the example : T.assume(i < 14 or
A[i] == 0)
- PrimExpr tensor_context; // The context of the assume statement (the
bound on the axis)
- PrimExpr tensor_predicate; // The condition inside assume statement (i <
14) excluding
- // tensorload expression (A[i] == 0)
- TensorLoad tensor_load; // Storing the tensor load Eg: A[i] in A[i] == 0
- PrimExpr tensor_value; // Storing the value for the tensor Eg : 0 in
A[i] == 0
- ffi::Array<PrimExpr> tensor_indices; // Storing the indices of the tensor
Eg : i
- };
- // List of conditions in a scope
- std::vector<PrimExpr> conditions_;
-
- // Storing all the tensor assumptions data in map
- std::unordered_map<tirx::TensorVar, assume_struct, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>
- map_tensor_assumption;
-
- struct InternalConstraintContext {
- /* This stuct appends the constraint passed to it in the conditions list.
- It keeps track of the bounds of the variables along with any conditions on
the variables */
- InternalConstraintContext(ParseAssumeAndOvercompute* self, PrimExpr
constraint)
- : self(self), analyzer_context(self->analyzer_, constraint) {
- old_num_constraints = self->conditions_.size();
-
- auto side_effect = SideEffect(constraint);
- if (side_effect <= CallEffectKind::kPure) {
- self->conditions_.push_back(constraint);
- } else if (side_effect <= CallEffectKind::kReadState) {
- assume = constraint;
- }
-
- new_num_constraints = self->conditions_.size();
- }
-
- ~InternalConstraintContext() {
- TVM_FFI_ICHECK_EQ(self->conditions_.size(), new_num_constraints)
- << "Internal error: Each condition should only be popped once.";
- self->conditions_.erase(self->conditions_.begin() + old_num_constraints,
- self->conditions_.end());
- }
-
- ParseAssumeAndOvercompute* self{nullptr};
- With<sym::ConstraintContext> analyzer_context;
- size_t old_num_constraints{0};
- size_t new_num_constraints{0};
- ffi::Optional<PrimExpr> assume{std::nullopt};
-
- // Disable default-generated copy/move assignment and constructors
- InternalConstraintContext(const InternalConstraintContext&) = delete;
- InternalConstraintContext& operator=(const InternalConstraintContext&) =
delete;
- InternalConstraintContext(InternalConstraintContext&&) = delete;
- InternalConstraintContext& operator=(InternalConstraintContext&&) = delete;
- };
-
- PrimExpr CurrentScopePredicate() const {
- /* This combines all the constraints in a scope */
- PrimExpr predicate = IntImm::Bool(true);
- for (const auto& condition : conditions_) {
- predicate = predicate && condition;
- }
- return predicate;
- }
-
- UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
- /* Create and delete the scope with bind.
- Add the minimum and maximum bound for the variables to the conditions_
list using
- InternalConstraintContext */
- analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent));
- InternalConstraintContext ctx1(this, op->loop_var >= op->min);
- InternalConstraintContext ctx2(this,
- static_cast<PrimExpr>(op->loop_var) <
op->min + op->extent);
- return Parent::Mutate_(op, inplace_mode);
- }
-
- UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* op, InplaceMode
inplace_mode) override {
- if
(map_tensor_assumption.find(op->source.as_or_throw<tvm::tirx::TensorVar>()) !=
- map_tensor_assumption.end()) {
- /* If the cuurent context where the tensor load is present is same as
- the context of the tensor assumption then, return the tensor value
present in the assumption.
- This will eventually replace the tensorload value in the complete
expresison */
-
- auto tensor_assumption =
-
map_tensor_assumption.at(op->source.as_or_throw<tvm::tirx::TensorVar>());
- PrimExpr current_predicate_and_context = CurrentScopePredicate();
- PrimExpr tensor_predicate_and_context =
- tensor_assumption.tensor_context &&
tensor_assumption.tensor_predicate;
- bool current_context_and_tensor_constraint_is_same =
ffi::StructuralEqual::Equal(
- current_predicate_and_context, tensor_predicate_and_context,
/*map_free_vars=*/true);
-
- if (current_context_and_tensor_constraint_is_same) {
- return tensor_assumption.tensor_value;
- }
- }
- return ffi::Unchanged();
- }
-
- UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode
inplace_mode) final {
- TensorStore store = Parent::Mutate_(op, inplace_mode)
- .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<TensorStore>();
- if (!op->unique()) inplace_mode = InplaceMode::kDisallow;
-
- // Eliminate the builtin if_then_else statement
- if (auto* call = op->value.as<CallNode>()) {
- if (call->op.same_as(prim::if_then_else_op())) {
- PrimExpr cond = call->args[0].as_or_throw<PrimExpr>();
- PrimExpr then_clause = call->args[1].as_or_throw<PrimExpr>();
- PrimExpr else_clause = call->args[2].as_or_throw<PrimExpr>();
-
- PrimExpr then_clause_in_then_context{ffi::UnsafeInit{}};
- PrimExpr else_clause_in_then_context{ffi::UnsafeInit{}};
- PrimExpr then_clause_in_else_context{ffi::UnsafeInit{}};
- PrimExpr else_clause_in_else_context{ffi::UnsafeInit{}};
- {
- // Simplifying expressions in " then context "
- InternalConstraintContext then_ctx(this, cond);
- // This will call the current class's appropriate VisitStmt function
- then_clause_in_then_context =
- Mutate(then_clause, inplace_mode).ValueOrUnchanged(then_clause);
- then_clause_in_then_context =
analyzer_->Simplify(then_clause_in_then_context);
-
- else_clause_in_then_context =
- Mutate(else_clause, inplace_mode).ValueOrUnchanged(else_clause);
- else_clause_in_then_context =
analyzer_->Simplify(else_clause_in_then_context);
- }
- {
- // Simplifying expressions in " else context "
- InternalConstraintContext else_ctx(this, !cond);
- // This will call the current class's appropriate VisitStmt function
- then_clause_in_else_context =
- Mutate(then_clause, inplace_mode).ValueOrUnchanged(then_clause);
- then_clause_in_else_context =
analyzer_->Simplify(then_clause_in_else_context);
-
- else_clause_in_else_context =
- Mutate(else_clause, inplace_mode).ValueOrUnchanged(else_clause);
- else_clause_in_else_context =
analyzer_->Simplify(else_clause_in_else_context);
- }
-
- if (ffi::StructuralEqual()(then_clause_in_then_context,
else_clause_in_then_context)) {
- PrimExpr value = analyzer_->Simplify(else_clause);
- if (inplace_mode == InplaceMode::kAllow) {
- const_cast<TensorStoreNode*>(op)->value = std::move(value);
- return ffi::Unchanged();
- } else {
- auto copy = ffi::make_object<TensorStoreNode>(*op);
- copy->value = std::move(value);
- return TensorStore(std::move(copy));
- }
- } else if (ffi::StructuralEqual()(then_clause_in_else_context,
- else_clause_in_else_context)) {
- PrimExpr value = analyzer_->Simplify(then_clause);
- if (inplace_mode == InplaceMode::kAllow) {
- const_cast<TensorStoreNode*>(op)->value = std::move(value);
- return ffi::Unchanged();
- } else {
- auto copy = ffi::make_object<TensorStoreNode>(*op);
- copy->value = std::move(value);
- return TensorStore(std::move(copy));
- }
- } else {
- return Parent::Mutate_(op, inplace_mode);
- }
- }
- }
- return Parent::Mutate_(op, inplace_mode);
- }
-
- UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode)
override {
- if (op->op.same_as(prim::assume_op())) {
- Assume(op->args[0].as_or_throw<PrimExpr>());
- }
- return Parent::Mutate_(op, inplace_mode);
- }
-
- void Assume(PrimExpr assumption) {
- for (const auto& expr : sym::ExtractConstraints(assumption, false)) {
- AssumeConstraintComponent(expr);
- }
- }
-
- void AssumeConstraintComponent(PrimExpr assumption) {
- PrimExpr additional_predicate = IntImm::Bool(true);
-
- std::vector<PrimExpr> tensor_exprs;
- for (const auto& expr : sym::ExtractComponents(assumption)) {
- auto side_effect = SideEffect(expr);
- if (side_effect <= CallEffectKind::kPure) {
- // Pulling out portions of the assumption that do not depend
- // on a tensor value allows the following two forms to be
- // treated identically.
- //
- // Option 1: if i < 3: T.assume(tensor[i] == value)
- // Option 2: T.assume(i>=3 or tensor[i] == value)
- additional_predicate = additional_predicate && logical_not(expr);
- } else if (side_effect == CallEffectKind::kReadState) {
- tensor_exprs.push_back(expr);
- } else {
- TVM_FFI_THROW(InternalError)
- << "Assumption must be pure or read-only, but contained expression
" << expr
- << " with side-effect \'" << side_effect << "\'";
- }
- }
-
- additional_predicate =
analyzer_->Simplify(std::move(additional_predicate));
- TVM_FFI_ICHECK_EQ(tensor_exprs.size(), 1)
- << "T.assume must contain only a single tensor expression";
-
- auto* as_equal_node = tensor_exprs[0].as<prim::EQNode>();
- TVM_FFI_ICHECK(as_equal_node)
- << "T.assume tensor constraint must be of the form 'tensor[indices] ==
"
- "value', but received "
- << assumption;
- if (!as_equal_node) {
- // This assumption is an inequality on a data-dependent
- // conditional. Not an error for this to occur, but also not
- // something that is currently supported.
- return;
- }
-
- // Parse the statement and store the desired values
- // Ex: A[i]==0, load = A[i], value = 0
- TensorLoad load{ffi::UnsafeInit{}};
- PrimExpr value{ffi::UnsafeInit{}};
- if (auto opt = as_equal_node->a.as<TensorLoad>()) {
- load = opt.value();
- value = as_equal_node->b;
- } else if (auto opt = as_equal_node->b.as<TensorLoad>()) {
- load = opt.value();
- value = as_equal_node->a;
- } else {
- TVM_FFI_THROW(InternalError)
- << "T.assume tensor constraint must be of the form 'tensor[indices]
== value'";
- }
-
- // Populating the assume statement predicate, tensor, value
- // and the context of the assume statement
- assume_struct tensor_data{CurrentScopePredicate(), additional_predicate,
load, value,
- load->indices};
- for (size_t i = 0; i < load->indices.size(); i++) {
-
tensor_data.tensor_indices.push_back(analyzer_->Simplify(load->indices[i]));
- }
- map_tensor_assumption.insert_or_assign(
- tensor_data.tensor_load->source.as_or_throw<tvm::tirx::TensorVar>(),
tensor_data);
-
- auto has_side_effect = SideEffect(value) > CallEffectKind::kPure;
- TVM_FFI_ICHECK(!has_side_effect)
- << "TensorVar value in constraint must be pure expression, but was "
<< value;
- if (has_side_effect) {
- return;
- }
- }
-};
-
-namespace transform {
-
-Pass UseAssumeToReduceBranches() {
- auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- auto* n = f.CopyOnWrite();
- sym::Analyzer analyzer;
-
- // The pass runs & eliminates pad branch with overcompute only if,
- // the function has op_pattern defined and is an elementwise op.
- // AnnotateTIROpPattern pass will set op_pattern in op attributes of the
function.
- if (n->attrs.GetAttr<int64_t>(tvm::relax::attr::kOpPattern).has_value()) {
- ffi::Optional<int64_t> opt_pattern =
f->GetAttr<int64_t>(tvm::relax::attr::kOpPattern);
- if (opt_pattern.has_value()) {
- relax::OpPatternKind pattern;
- pattern = static_cast<relax::OpPatternKind>(opt_pattern.value());
-
- if (pattern == relax::OpPatternKind::kElemWise ||
- pattern == relax::OpPatternKind::kBroadcast) {
- // If the function contains assume statement then, run the mutator
pass.
- auto assume_checker = ffi::make_object<AssumeChecker>();
- assume_checker->Visit(std::move(n->body));
-
- if (assume_checker->has_assume) {
- // Leverage from assume and eliminate the branch
- auto func_analyzer_mutator =
ffi::make_object<ParseAssumeAndOvercompute>(analyzer);
- n->body = func_analyzer_mutator->Mutate(n->body,
InplaceMode::kAllow)
- .ValueOrUnchanged(std::move(n->body));
- }
- }
- }
- }
- return f;
- };
- return CreateFunctionPass(pass_func, 0, "s_tir.UseAssumeToReduceBranches");
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.transform.UseAssumeToReduceBranches",
UseAssumeToReduceBranches);
-}
-
-} // namespace transform
-
-} // namespace s_tir
-} // namespace tvm
diff --git a/src/sym/const_int_bound.cc b/src/sym/const_int_bound.cc
index 7a9f5fae01..85f5078612 100644
--- a/src/sym/const_int_bound.cc
+++ b/src/sym/const_int_bound.cc
@@ -30,6 +30,7 @@
#include <tvm/sym/analyzer.h>
#include <algorithm>
+#include <numeric>
#include <optional>
#include "constraint_extract.h"
@@ -308,7 +309,7 @@ class ConstIntBoundAnalyzer::Impl
if (parent_ && b.min_value == b.max_value) {
ModularSet mod_a = parent_->modular_set(op->a);
int64_t modulus = b.min_value;
- int64_t gcd_coeff_mod = ZeroAwareGCD(mod_a->coeff, modulus);
+ int64_t gcd_coeff_mod = std::gcd(mod_a->coeff, modulus);
// If gcd_coeff_mod > 1, we can get tighter bounds.
// Since gcd_coeff_mod divides both mod_a->coeff and modulus, we know
@@ -409,7 +410,7 @@ class ConstIntBoundAnalyzer::Impl
if (parent_ && b.min_value == b.max_value) {
ModularSet mod_a = parent_->modular_set(op->a);
int64_t modulus = b.min_value;
- int64_t gcd_coeff_mod = ZeroAwareGCD(mod_a->coeff, modulus);
+ int64_t gcd_coeff_mod = std::gcd(mod_a->coeff, modulus);
// If gcd_coeff_mod > 1, we can get tighter bounds.
// Since gcd_coeff_mod divides both mod_a->coeff and modulus, we know
diff --git a/src/sym/int_operator.h b/src/sym/int_operator.h
index e9e6c33073..92d4597a99 100644
--- a/src/sym/int_operator.h
+++ b/src/sym/int_operator.h
@@ -28,6 +28,7 @@
#include <tvm/ir/prim/expr.h>
#include <limits>
+#include <numeric>
#include <utility>
#include "../ir/prim/int_operator.h"
@@ -91,78 +92,20 @@ using prim::detail::floormod;
using prim::detail::truncdiv;
using prim::detail::truncmod;
-/*!
- * \brief Use Extended Euclidean algorithm to solve ax + by = gcd(a, b)
- * \param a The first coefficient.
- * \param b The second coefficient.
- * \param x The solution of x.
- * \param y The solution of y.
- * \return The GCD of a and b.
- */
-inline int64_t ExtendedEuclidean(int64_t a, int64_t b, int64_t* x, int64_t* y)
{
- // Extended Euclidean algorithm
- // if a < 0, the problem can be convert into
- // |a|* (-x) + b * y = gcd(|a|, b)
- //
- // initial condition:
- // a * 0 + b * 1 = b
- // a * 1 + b * 0 = a
- int64_t s = 0, old_s = 1;
- int64_t r = b, old_r = a >= 0 ? a : -a;
- // Iteration (r2 < r1):
- // a * x1 + b * y1 = r1
- // a * x2 + b * y2 = r2
- // The above two eqs can derive the following eq (q = r1 / r2)
- // a * (x1 - x2 * q) + b * (y1 - y2 * q) = r1 - r2 * q = r3
- // Because r3 < r2, the iteration can eventually terminate
- while (r != 0) {
- int64_t q = old_r / r;
- int64_t tmp = old_r;
- old_r = r;
- r = tmp - q * r;
- tmp = old_s;
- old_s = s;
- s = tmp - q * s;
- }
-
- *x = a >= 0 ? old_s : -old_s;
- if (b != 0) {
- *y = (old_r - (*x) * a) / b;
- } else {
- *y = 1;
- }
-
- return old_r;
-}
-
/*!
* \brief Take GCD of a and b.
* \param a The first operand.
* \param b The second operand.
* \return The result.
*/
-inline int64_t ZeroAwareGCD(int64_t a, int64_t b) {
- if (a < 0) a = -a;
- if (b < 0) b = -b;
- if (a < b) std::swap(a, b);
- if (b == 0) return a;
- // perform GCD (greatest common divisor)
- // ax + by = gcd(a, b) z if a != 0, b != 0
- while (a % b != 0) {
- a = a % b;
- std::swap(a, b);
- }
- return b;
-}
-
inline ffi::BigInt ZeroAwareGCD(ffi::BigInt a, ffi::BigInt b) {
auto a_int64 = a.as<int64_t>();
auto b_int64 = b.as<int64_t>();
- // The native overload cannot negate INT64_MIN.
+ // std::gcd requires both absolute values to be representable as int64_t.
if (a_int64.has_value() && b_int64.has_value() &&
*a_int64 != std::numeric_limits<int64_t>::min() &&
*b_int64 != std::numeric_limits<int64_t>::min()) {
- return ZeroAwareGCD(*a_int64, *b_int64);
+ return std::gcd(*a_int64, *b_int64);
}
if (a < 0) a = -a;
if (b < 0) b = -b;
@@ -177,17 +120,6 @@ inline ffi::BigInt ZeroAwareGCD(ffi::BigInt a, ffi::BigInt
b) {
return b;
}
-/*!
- * \brief Calculate the least common multiple for two values.
- * \param a an integer number
- * \param b an integer number
- * \return the least common multiple.
- */
-inline int64_t LeastCommonMultiple(int64_t a, int64_t b) {
- int64_t x, y;
- return (a * b) / ExtendedEuclidean(a, b, &x, &y);
-}
-
} // namespace sym
} // namespace tvm
#endif // TVM_SYM_INT_OPERATOR_H_
diff --git a/src/sym/int_set.cc b/src/sym/int_set.cc
index beb21d9f58..2c9e3cb8ad 100644
--- a/src/sym/int_set.cc
+++ b/src/sym/int_set.cc
@@ -1012,6 +1012,10 @@ ffi::Array<IntSet> UnionRegionLowerBound(const
ffi::Array<ffi::Array<IntSet>>& n
return result;
}
+IntSet Intersect(const IntSet& lhs, const IntSet& rhs, const Analyzer&
analyzer) {
+ return Intersect(analyzer.get(), ToIntervalSet(lhs), ToIntervalSet(rhs));
+}
+
IntSet Intersect(const ffi::Array<IntSet>& sets) {
if (sets.size() == 0) return IntSet::Nothing();
if (sets.size() == 1) return sets[0];
diff --git a/src/sym/modular_set.cc b/src/sym/modular_set.cc
index dbbca70018..fd63276a39 100644
--- a/src/sym/modular_set.cc
+++ b/src/sym/modular_set.cc
@@ -29,6 +29,7 @@
#include <tvm/sym/analyzer.h>
#include <limits>
+#include <numeric>
#include <unordered_map>
#include <utility>
@@ -38,6 +39,44 @@ namespace tvm {
namespace sym {
using namespace tvm::prim;
+namespace {
+int64_t ExtendedEuclidean(int64_t a, int64_t b, int64_t* x, int64_t* y) {
+ // Extended Euclidean algorithm
+ // if a < 0, the problem can be convert into
+ // |a|* (-x) + b * y = gcd(|a|, b)
+ //
+ // initial condition:
+ // a * 0 + b * 1 = b
+ // a * 1 + b * 0 = a
+ int64_t s = 0, old_s = 1;
+ int64_t r = b, old_r = a >= 0 ? a : -a;
+ // Iteration (r2 < r1):
+ // a * x1 + b * y1 = r1
+ // a * x2 + b * y2 = r2
+ // The above two eqs can derive the following eq (q = r1 / r2)
+ // a * (x1 - x2 * q) + b * (y1 - y2 * q) = r1 - r2 * q = r3
+ // Because r3 < r2, the iteration can eventually terminate
+ while (r != 0) {
+ int64_t q = old_r / r;
+ int64_t tmp = old_r;
+ old_r = r;
+ r = tmp - q * r;
+ tmp = old_s;
+ old_s = s;
+ s = tmp - q * s;
+ }
+
+ *x = a >= 0 ? old_s : -old_s;
+ if (b != 0) {
+ *y = (old_r - (*x) * a) / b;
+ } else {
+ *y = 1;
+ }
+
+ return old_r;
+}
+} // namespace
+
TVM_FFI_STATIC_INIT_BLOCK() { ModularSetObj::RegisterReflection(); }
ModularSet::ModularSet(int64_t coeff, int64_t base) {
@@ -158,14 +197,14 @@ class ModularSetAnalyzer::Impl : public
tvm::ExprFunctor<ModularSetAnalyzer::Ent
Entry Dispatch_(const prim::AddNode* op) final {
Entry a = Dispatch(op->a);
Entry b = Dispatch(op->b);
- int64_t coeff = ZeroAwareGCD(a.coeff, b.coeff);
+ int64_t coeff = std::gcd(a.coeff, b.coeff);
return Entry(coeff, a.base + b.base);
}
Entry Dispatch_(const prim::SubNode* op) final {
Entry a = Dispatch(op->a);
Entry b = Dispatch(op->b);
- int64_t coeff = ZeroAwareGCD(a.coeff, b.coeff);
+ int64_t coeff = std::gcd(a.coeff, b.coeff);
return Entry(coeff, a.base - b.base);
}
@@ -179,7 +218,7 @@ class ModularSetAnalyzer::Impl : public
tvm::ExprFunctor<ModularSetAnalyzer::Ent
int64_t pq = a.coeff * b.coeff;
int64_t pm = a.coeff * b.base;
int64_t qn = a.base * b.coeff;
- int64_t coeff = ZeroAwareGCD(pq, ZeroAwareGCD(pm, qn));
+ int64_t coeff = std::gcd(pq, std::gcd(pm, qn));
return Entry(coeff, a.base * b.base);
}
@@ -237,7 +276,7 @@ class ModularSetAnalyzer::Impl : public
tvm::ExprFunctor<ModularSetAnalyzer::Ent
Entry ModByConst(const PrimExpr& lhs, int64_t val, bool round_down) {
Entry a = Dispatch(lhs);
TVM_FFI_ICHECK_NE(val, 0);
- int64_t coeff = ZeroAwareGCD(a.coeff, val);
+ int64_t coeff = std::gcd(a.coeff, val);
if (a.base % coeff == 0 ||
(a.base > 0 && (round_down || parent_->CanProveGreaterEqual(lhs, 0))))
{
return Entry(coeff, a.base % coeff);
@@ -330,7 +369,7 @@ class ModularSetAnalyzer::Impl : public
tvm::ExprFunctor<ModularSetAnalyzer::Ent
*/
static Entry Union(Entry a, Entry b) {
// {ax + y} \cup {bz + h} => {gcd(a, b) x + {y or h}}
- int64_t coeff = ZeroAwareGCD(a.coeff, b.coeff);
+ int64_t coeff = std::gcd(a.coeff, b.coeff);
if (coeff == 0) {
if (a.base == b.base) return a;
return Everything();
@@ -340,7 +379,7 @@ class ModularSetAnalyzer::Impl : public
tvm::ExprFunctor<ModularSetAnalyzer::Ent
if (base0 == base1) {
return Entry(coeff, base0);
} else {
- return Entry(ZeroAwareGCD(ZeroAwareGCD(base0, base1), coeff), base0);
+ return Entry(std::gcd(std::gcd(base0, base1), coeff), base0);
}
}
diff --git a/src/sym/presburger_set.cc b/src/sym/presburger_set.cc
deleted file mode 100644
index 2d19acf37b..0000000000
--- a/src/sym/presburger_set.cc
+++ /dev/null
@@ -1,293 +0,0 @@
-/*
- * 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 presburger_set.cc
- * \brief The presburger set functions
- */
-#include "presburger_set.h"
-
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/extra/structural_visit.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/sym/int_set.h>
-#include <tvm/sym/pattern.h>
-
-#include <algorithm>
-#include <unordered_map>
-#include <utility>
-#include <vector>
-
-#include "constraint_extract.h"
-#include "interval_set.h"
-
-namespace tvm {
-namespace sym {
-using namespace tvm::prim;
-
-#if defined(TVM_MLIR_VERSION) && TVM_MLIR_VERSION >= 150
-
-TVM_FFI_STATIC_INIT_BLOCK() { PresburgerSetObj::RegisterReflection(); }
-
-static int64_t GetPresburgerCoefficient(const PrimExpr& coefficient) {
- const auto* imm = coefficient.as<IntImmNode>();
- auto value = imm ? imm->value.as<int64_t>() : std::nullopt;
- TVM_FFI_ICHECK(value.has_value())
- << "Presburger coefficients must be constant integers that fit int64: "
<< coefficient;
- return *value;
-}
-
-static void Update(const PrimExpr& constraint, PresburgerSetObj* intset) {
- auto& space = intset->space;
- auto constraints_union = ExtractComponents(constraint);
- for (const PrimExpr& subconstraint : constraints_union) {
- auto entries = ExtractConstraints(subconstraint, false);
- auto vars = intset->GetVars();
- IntegerRelation disjunct(entries.size(), 0, vars.size() + 1, space);
- for (const PrimExpr& entry : entries) {
- // The expression is expect to be simplified to only contain ==, <= or <
- if (entry.as<prim::LENode>()) {
- auto coeffs_a = DetectLinearEquation(entry.as<prim::LENode>()->a,
vars);
- auto coeffs_b = DetectLinearEquation(entry.as<prim::LENode>()->b,
vars);
- std::vector<int64_t> int_coeffs;
- for (size_t i = 0; i < coeffs_a.size(); i++) {
- int_coeffs.push_back(GetPresburgerCoefficient(coeffs_b[i]) -
- GetPresburgerCoefficient(coeffs_a[i]));
- }
- disjunct.addInequality(int_coeffs);
- } else if (entry.as<prim::LTNode>()) {
- auto coeffs_a = DetectLinearEquation(entry.as<prim::LTNode>()->a,
vars);
- auto coeffs_b = DetectLinearEquation(entry.as<prim::LTNode>()->b,
vars);
- std::vector<int64_t> int_coeffs;
- for (size_t i = 0; i < coeffs_a.size(); i++) {
- int_coeffs.push_back(GetPresburgerCoefficient(coeffs_b[i]) -
- GetPresburgerCoefficient(coeffs_a[i]));
- }
- int_coeffs[int_coeffs.size() - 1] -= 1;
- disjunct.addInequality(int_coeffs);
- } else if (entry.as<prim::EQNode>()) {
- auto coeffs_a = DetectLinearEquation(entry.as<prim::EQNode>()->a,
vars);
- auto coeffs_b = DetectLinearEquation(entry.as<prim::EQNode>()->b,
vars);
- std::vector<int64_t> int_coeffs;
- for (size_t i = 0; i < coeffs_a.size(); i++) {
- int_coeffs.push_back(GetPresburgerCoefficient(coeffs_a[i]) -
- GetPresburgerCoefficient(coeffs_b[i]));
- }
- disjunct.addEquality(int_coeffs);
- } else {
- TVM_FFI_THROW(InternalError)
- << "Unsupported constraint expression: " << entry->GetTypeKey();
- }
- }
- intset->unionInPlace(disjunct);
- }
-}
-
-PresburgerSet::PresburgerSet(const PrimExpr& constraint) :
IntSet(ffi::UnsafeInit{}) {
- ffi::Array<PrimVar> vars;
- ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(
- constraint, [&vars](const Var& var) -> ffi::Expected<ffi::WalkResult> {
- if (auto prim_var = var.as<PrimVar>()) {
- PrimVar var = *prim_var;
- if (!std::any_of(vars.begin(), vars.end(),
- [&var](const PrimVar& v) { return v.same_as(var);
})) {
- vars.push_back(var);
- }
- }
- return ffi::WalkResult::Advance();
- });
- auto constraints_union = ExtractComponents(constraint);
- Analyzer analyzer;
- PrimExpr simplified_constraint = analyzer->Simplify(constraint,
kSimplifyRewriteCanonicalRewrite);
- auto space = PresburgerSpace::getRelationSpace(vars.size(), 0, 0, 0);
- auto node = ffi::make_object<PresburgerSetObj>(std::move(space), vars);
- node->SetVars(vars);
- Update(simplified_constraint, node.get());
- data_ = std::move(node);
-}
-
-PresburgerSet::PresburgerSet(const std::vector<IntegerRelation>& disjuncts,
- const ffi::Array<PrimVar>& vars)
- : IntSet(ffi::UnsafeInit{}) {
- auto node = ffi::make_object<PresburgerSetObj>(disjuncts,
disjuncts[0].getSpace(), vars);
- data_ = std::move(node);
-}
-
-void PresburgerSetObj::UpdateConstraint(const PrimExpr& constraint,
- const ffi::Array<PrimVar>& vars) {
- Analyzer analyzer;
- PrimExpr simplified_constraint = analyzer->Simplify(constraint,
kSimplifyRewriteCanonicalRewrite);
- Update(simplified_constraint, this);
- SetVars(vars);
-}
-
-PrimExpr PresburgerSetObj::GenerateConstraint() const {
- PrimExpr constraint = IntImm::Bool(false);
- for (const IntegerRelation& disjunct : disjuncts) {
- PrimExpr union_entry = IntImm::Bool(true);
- for (unsigned i = 0, e = disjunct.getNumEqualities(); i < e; ++i) {
- PrimExpr linear_eq = IntImm::Int64(0);
- if (disjunct.getNumCols() > 1) {
- for (unsigned j = 0, f = disjunct.getNumCols() - 1; j < f; ++j) {
-#if TVM_MLIR_VERSION >= 160
- auto coeff = int64_t(disjunct.atEq(i, j));
-#else
- auto coeff = disjunct.atEq(i, j);
-#endif
- if (coeff >= 0 || IsZero(linear_eq)) {
- linear_eq = linear_eq + IntImm::Int64(coeff) * vars[j];
- } else {
- linear_eq = linear_eq - IntImm::Int64(-coeff) * vars[j];
- }
- }
- }
-#if TVM_MLIR_VERSION >= 160
- auto c0 = int64_t(disjunct.atEq(i, disjunct.getNumCols() - 1));
-#else
- auto c0 = disjunct.atEq(i, disjunct.getNumCols() - 1);
-#endif
- linear_eq = linear_eq + IntImm::Int64(c0);
- union_entry = (union_entry && (linear_eq == 0));
- }
- for (unsigned i = 0, e = disjunct.getNumInequalities(); i < e; ++i) {
- PrimExpr linear_eq = IntImm::Int64(0);
- if (disjunct.getNumCols() > 1) {
- for (unsigned j = 0, f = disjunct.getNumCols() - 1; j < f; ++j) {
-#if TVM_MLIR_VERSION >= 160
- auto coeff = int64_t(disjunct.atIneq(i, j));
-#else
- auto coeff = disjunct.atIneq(i, j);
-#endif
- if (coeff >= 0 || IsZero(linear_eq)) {
- linear_eq = linear_eq + IntImm::Int64(coeff) * vars[j];
- } else {
- linear_eq = linear_eq - IntImm::Int64(-coeff) * vars[j];
- }
- }
- }
-#if TVM_MLIR_VERSION >= 160
- auto c0 = int64_t(disjunct.atIneq(i, disjunct.getNumCols() - 1));
-#else
- auto c0 = disjunct.atIneq(i, disjunct.getNumCols() - 1);
-#endif
- if (c0 >= 0) {
- linear_eq = linear_eq + IntImm::Int64(c0);
- } else {
- linear_eq = linear_eq - IntImm::Int64(-c0);
- }
- union_entry = (union_entry && (linear_eq >= 0));
- }
- constraint = constraint || union_entry;
- }
-
- return constraint;
-}
-
-PresburgerSet Union(const ffi::Array<PresburgerSet>& sets) {
- TVM_FFI_ICHECK_GT(sets.size(), 0);
- if (sets.size() == 1) return sets[0];
- auto relations = sets[0]->disjuncts;
- for (size_t i = 1; i < sets.size(); ++i) {
- for (const IntegerRelation& rel : sets[i]->disjuncts) {
- relations.push_back(rel);
- }
- }
- return PresburgerSet(std::move(relations), sets[0]->GetVars());
-}
-
-PresburgerSet Intersect(const ffi::Array<PresburgerSet>& sets) {
- TVM_FFI_ICHECK_GT(sets.size(), 0);
- if (sets.size() == 1) return sets[0];
- auto relations = sets[0]->disjuncts;
- const auto& space = sets[0]->space;
-
- for (size_t i = 1; i < sets.size(); ++i) {
- TVM_FFI_ICHECK(space.isCompatible(sets[i]->space)) << "Spaces should
match";
- for (const IntegerRelation& relA : sets[i]->disjuncts) {
- for (const IntegerRelation& relB : relations) {
- IntegerRelation intersection = relA.intersect(relB);
- if (!intersection.isEmpty()) relations.push_back(intersection);
- }
- }
- }
- return PresburgerSet(std::move(relations), sets[0]->GetVars());
-}
-
-IntSet EvalSet(const PrimExpr& e, const PresburgerSet& set) {
- ffi::Array<PrimExpr> tvm_coeffs = DetectLinearEquation(e, set->GetVars());
-#if TVM_MLIR_VERSION >= 190
- SmallVector<llvm::DynamicAPInt> coeffs;
-#elif TVM_MLIR_VERSION >= 160
- SmallVector<mlir::presburger::MPInt> coeffs;
-#else
- SmallVector<int64_t> coeffs;
-#endif
-
- coeffs.reserve(tvm_coeffs.size());
- for (const PrimExpr& it : tvm_coeffs) {
-#if TVM_MLIR_VERSION >= 190
- coeffs.push_back(llvm::DynamicAPInt(GetPresburgerCoefficient(it)));
-#elif TVM_MLIR_VERSION >= 160
- coeffs.push_back(mlir::presburger::MPInt(GetPresburgerCoefficient(it)));
-#else
- coeffs.push_back(GetPresburgerCoefficient(it));
-#endif
- }
-
- IntSet result = IntSet::Nothing();
- for (const IntegerRelation& it : set->disjuncts) {
- Simplex simplex(it);
- auto range = simplex.computeIntegerBounds(coeffs);
- auto maxRoundedDown(simplex.computeOptimum(Simplex::Direction::Up,
coeffs));
- auto opt = range.first.getOptimumIfBounded();
-#if TVM_MLIR_VERSION >= 160
- auto min = opt.has_value() ? IntImm::Int64(int64_t(opt.value())) :
neg_inf();
-#else
- auto min = opt.hasValue() ? IntImm::Int64(opt.getValue()) : neg_inf();
-#endif
- opt = range.second.getOptimumIfBounded();
-#if TVM_MLIR_VERSION >= 160
- auto max = opt.has_value() ? IntImm::Int64(int64_t(opt.value())) :
pos_inf();
-#else
- auto max = opt.hasValue() ? IntImm::Int64(opt.getValue()) : pos_inf();
-#endif
- auto interval = IntervalSet(min, max);
- result = Union({result, interval});
- }
- return result;
-}
-
-// Pattern A (RM): auto-default repr from reflection.
-
-#else // defined(TVM_MLIR_VERSION) && TVM_MLIR_VERSION >= 150
-
-PresburgerSet MakePresburgerSet(const PrimExpr& constraint) { return
PresburgerSet(constraint); }
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- PresburgerSetObj::RegisterReflection();
- refl::GlobalDef().def("sym.PresburgerSet", MakePresburgerSet);
-}
-
-#endif // defined(TVM_MLIR_VERSION) && TVM_MLIR_VERSION >= 150
-
-} // namespace sym
-} // namespace tvm
diff --git a/src/sym/presburger_set.h b/src/sym/presburger_set.h
deleted file mode 100644
index 85074f211a..0000000000
--- a/src/sym/presburger_set.h
+++ /dev/null
@@ -1,201 +0,0 @@
-/*
- * 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 presburger_set.h
- * \brief Integer set based on MLIR Presburger set
- */
-#ifndef TVM_SYM_PRESBURGER_SET_H_
-#define TVM_SYM_PRESBURGER_SET_H_
-
-#ifdef TVM_MLIR_VERSION
-#if TVM_MLIR_VERSION >= 150
-#include <mlir/Analysis/Presburger/IntegerRelation.h>
-#include <mlir/Analysis/Presburger/PresburgerRelation.h>
-#include <mlir/Analysis/Presburger/Simplex.h>
-#endif
-#endif
-
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/cow.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/sym/analyzer.h>
-
-#include <limits>
-#include <vector>
-
-#include "const_fold.h"
-
-namespace tvm {
-namespace sym {
-
-#ifdef TVM_MLIR_VERSION
-#if TVM_MLIR_VERSION >= 150
-using namespace mlir;
-using namespace presburger;
-
-// Acknowledgement: PresburgerSet is based on Presburger set of MLIR.
-/*!
- * \brief Symbolic integer set.
- *
- * \note PresburgerSet aims to provide compatible APIs with IntSet,
- * and some additional APIs that analyze and solve
- * multi-dimension interger set problems
- */
-class PresburgerSetObj : public IntSetObj {
- public:
- PresburgerSetObj() : space(PresburgerSpace::getRelationSpace()) {}
- explicit PresburgerSetObj(const PresburgerSpace& space, const
ffi::Array<PrimVar>& vars)
- : disjuncts({}), space(space), vars(vars) {}
- explicit PresburgerSetObj(const std::vector<IntegerRelation>& disjuncts,
- const PresburgerSpace& space, const
ffi::Array<PrimVar>& vars)
- : disjuncts(disjuncts), space(space), vars(vars) {}
-
- /*! \brief Represent the union of multiple IntegerRelation */
- std::vector<IntegerRelation> disjuncts;
- /*! \brief The space of all the disjuncts */
- PresburgerSpace space;
-
- // visitor overload.
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<PresburgerSetObj>();
- }
-
- /*!
- * \brief Do inplace union with given disjunct
- * \param disjunct The given disjunct to be union with
- */
- void unionInPlace(const IntegerRelation& disjunct) {
- assert(space.isCompatible(disjunct.getSpace()) && "Spaces should match");
- disjuncts.push_back(disjunct);
- }
-
- /*!
- * \brief Update int set with given constraint
- * \param constraint The added constraint to the PresburgerSet.
- * \param vars The specified domain vars in constraint expression.
- */
- void UpdateConstraint(const PrimExpr& constraint, const ffi::Array<PrimVar>&
vars);
-
- /*!
- * \brief Generate expression that represents the constraint
- * \return The generated expression
- */
- PrimExpr GenerateConstraint() const;
-
- /*!
- * \brief Set domain vars
- * \param new_vars Vars that will be taken as the domain vars
- */
- void SetVars(const ffi::Array<PrimVar>& new_vars) { vars = new_vars; }
-
- /*!
- * \brief Get the current domain vars
- * \return The current doamin vars
- */
- ffi::Array<PrimVar> GetVars() const { return vars; }
-
- /*! \return whether integer set is empty */
- bool IsEmpty() const {
- return std::all_of(disjuncts.begin(), disjuncts.end(),
- std::mem_fn(&IntegerRelation::isIntegerEmpty));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.PresburgerSet", PresburgerSetObj,
IntSetObj);
-
- private:
- ffi::Array<PrimVar> vars;
-};
-
-/*!
- * \brief Integer set used for multi-dimension integer analysis.
- * \sa PresburgerSetObj
- */
-class PresburgerSet : public IntSet {
- public:
- /*!
- * \brief Make a new instance of PresburgerSet.
- * \param disjuncts The disjunts to construct the set.
- * \param vars The variables that the constraint describes about.
- * \return The created PresburgerSet.
- */
- TVM_TIRX_DLL PresburgerSet(const std::vector<IntegerRelation>& disjuncts,
- const ffi::Array<PrimVar>& vars);
-
- /*!
- * \brief Make a new instance of PresburgerSet, collect all vars as space
vars.
- * \param constraint The constraint to construct the set.
- * \return The created PresburgerSet.
- */
- TVM_TIRX_DLL PresburgerSet(const PrimExpr& constraint);
-
- TVM_DEFINE_OBJECT_REF_COW_METHOD(PresburgerSetObj);
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PresburgerSet, IntSet,
PresburgerSetObj);
-};
-#endif // TVM_MLIR_VERSION >= 150
-#else // TVM_MLIR_VERSION
-// Class definition when MLIR is not enabled, to prevent compile-time error.
-class PresburgerSetObj : public IntSetObj {
- public:
- // dummy visitor overload.
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<PresburgerSetObj>();
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.PresburgerSet", PresburgerSetObj,
IntSetObj);
-};
-
-class PresburgerSet : public IntSet {
- public:
- /*!
- * \brief Constructor interface to prompt when MLIR is not enabled.
- * \param constraint The constraint to construct the set.
- * \return The created set.
- */
- TVM_TIRX_DLL PresburgerSet(const PrimExpr& constraint) :
IntSet(ffi::UnsafeInit{}) {
- TVM_FFI_THROW(InternalError) << "MLIR is not enabled!";
- }
-};
-#endif // TVM_MLIR_VERSION
-/*!
- * \brief Create a union set of all sets
- * \param sets The sets to be combined
- * \return the set after union
- */
-PresburgerSet Union(const ffi::Array<PresburgerSet>& sets);
-
-/*!
- * \brief Create an intersected set of all sets
- * \param sets The sets to be intersected
- * \return The intersect set
- */
-PresburgerSet Intersect(const ffi::Array<PresburgerSet>& sets);
-
-/*!
- * \brief Evaluate the range of given expression based on the constraint
- * in given PresburgerSet
- * \param e The target expresision to be evaluated.
- * \param set The PresburgerSet defining the constraint.
- */
-IntSet EvalSet(const PrimExpr& e, const PresburgerSet& set);
-
-} // namespace sym
-} // namespace tvm
-
-#endif // TVM_SYM_PRESBURGER_SET_H_
diff --git a/src/sym/rewrite_simplify.cc b/src/sym/rewrite_simplify.cc
index fc75e5f95f..be3b607fa2 100644
--- a/src/sym/rewrite_simplify.cc
+++ b/src/sym/rewrite_simplify.cc
@@ -32,6 +32,7 @@
#include <tvm/sym/analyzer.h>
#include <algorithm>
+#include <numeric>
#include <tuple>
#include <utility>
@@ -2104,9 +2105,9 @@ PrimExpr
RewriteSimplifier::Impl::ApplyRewriteRules(prim::LT ret, InplaceMode in
auto common_factor = [&]() -> int64_t {
auto modular_a = analyzer_->modular_set(ret->a);
auto modular_b = analyzer_->modular_set(ret->b);
- auto gcd_lhs = ZeroAwareGCD(modular_a->base, modular_a->coeff);
- auto gcd_rhs = ZeroAwareGCD(modular_b->base, modular_b->coeff);
- return ZeroAwareGCD(gcd_lhs, gcd_rhs);
+ auto gcd_lhs = std::gcd(modular_a->base, modular_a->coeff);
+ auto gcd_rhs = std::gcd(modular_b->base, modular_b->coeff);
+ return std::gcd(gcd_lhs, gcd_rhs);
}();
if (common_factor > 1) {
return RecursiveRewrite(floordiv(ret->a, common_factor) <
floordiv(ret->b, common_factor),
diff --git
a/tests/python/relax/test_eliminate_pad_branch_using_buffer_assumption.py
b/tests/python/relax/test_eliminate_pad_branch_using_buffer_assumption.py
deleted file mode 100644
index 69ce2d8ef6..0000000000
--- a/tests/python/relax/test_eliminate_pad_branch_using_buffer_assumption.py
+++ /dev/null
@@ -1,646 +0,0 @@
-# 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.
-# pylint: disable=missing-docstring, unused-variable
-# ruff: noqa: F841
-
-# The test attempts to eliminate redundant pad branch and overcompute the
value for elementwise ops.
-# This helps to expose more opportunities to vectorize the code.
-
-import tvm_ffi
-
-import tvm
-import tvm.script
-import tvm.testing
-from tvm.script import relax as R
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-
-
[email protected]_module
-class AddBefore:
- @Ts.function(private=True)
- def add(
- a: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- b: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- compute: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- ):
- T.func_attr(
- {
- "op_attrs": {"lhs_axis": 0, "op_name": "qnn.add", "rhs_axis":
0},
- "op_pattern": 0,
- "operator_name": "add",
- "tirx.noalias": True,
- }
- )
- # with Ts.sblock("root"):
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("buffer_A_assumptions"):
- v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6
= Ts.axis.remap(
- "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5,
axis6]
- )
- Ts.reads(a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- Ts.writes()
- T.assume(
- not (
- (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
- or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5)
- )
- or a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6]
- == T.uint8(0)
- )
-
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("buffer_B_assumptions"):
- v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6
= Ts.axis.remap(
- "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5,
axis6]
- )
- Ts.reads(b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- Ts.writes()
- T.assume(
- not (
- (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
- or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5)
- )
- or b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6]
- == T.uint8(0)
- )
-
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("compute"):
- v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6
= Ts.axis.remap(
- "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5,
axis6]
- )
- Ts.reads(
- a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6],
- b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6],
- )
- Ts.writes(compute[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- compute[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6] = (
- T.if_then_else(
- (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
- or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5),
- T.uint8(0),
- a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- + b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6],
- )
- )
-
- @R.function
- def main(
- a: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
- b: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
- ) -> R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"):
- out = R.call_tir(
- AddBefore.add,
- (a, b),
- ty_args=[R.Tensor((1, 4, 4, 16, 8, 8, 32), dtype="uint8")],
- )
- return out
-
-
[email protected]_module
-class AddExpected:
- @Ts.function(private=True)
- def add(
- a: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- b: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- compute: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- ):
- T.func_attr(
- {
- "op_attrs": {"lhs_axis": 0, "op_name": "qnn.add", "rhs_axis":
0},
- "op_pattern": 0,
- "operator_name": "add",
- "tirx.noalias": True,
- }
- )
- # with Ts.sblock("root"):
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("buffer_A_assumptions"):
- v_axis0 = Ts.axis.spatial(T.int64(1), T.int64(0))
- v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 =
Ts.axis.remap(
- "SSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6]
- )
- Ts.reads(a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- Ts.writes()
- T.assume(
- (
- (v_axis1 < T.int64(3) or v_axis4 < T.int64(4))
- and (v_axis2 < T.int64(3) or v_axis5 < T.int64(4))
- )
- or a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- == T.uint8(0)
- )
-
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("buffer_B_assumptions"):
- v_axis0 = Ts.axis.spatial(T.int64(1), T.int64(0))
- v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 =
Ts.axis.remap(
- "SSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6]
- )
- Ts.reads(b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- Ts.writes()
- T.assume(
- (
- (v_axis1 < T.int64(3) or v_axis4 < T.int64(4))
- and (v_axis2 < T.int64(3) or v_axis5 < T.int64(4))
- )
- or b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- == T.uint8(0)
- )
-
- for axis0, axis1, axis2, axis3, axis4, axis5_0 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(2)
- ):
- for axis5_1_axis6_fused in T.vectorized(T.int64(128)):
- with Ts.sblock("compute"):
- v_axis0 = Ts.axis.spatial(T.int64(1), T.int64(0))
- v_axis1, v_axis2, v_axis3, v_axis4 = Ts.axis.remap(
- "SSSS", [axis1, axis2, axis3, axis4]
- )
- v_axis5 = Ts.axis.spatial(
- T.int64(8), axis5_0 * T.int64(4) + axis5_1_axis6_fused
// T.int64(32)
- )
- v_axis6 = Ts.axis.spatial(T.int64(32), axis5_1_axis6_fused
% T.int64(32))
- Ts.reads(
- a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6],
- b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6],
- )
- Ts.writes(
- compute[T.int64(0), v_axis1, v_axis2, v_axis3,
v_axis4, v_axis5, v_axis6]
- )
- compute[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6] = (
- a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- + b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- )
-
- @R.function
- def main(
- a: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
- b: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
- ) -> R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"):
- out = R.call_tir(
- AddExpected.add,
- (a, b),
- ty_args=[R.Tensor((1, 4, 4, 16, 8, 8, 32), dtype="uint8")],
- )
- return out
-
-
[email protected]_module
-class SubBefore:
- @Ts.function(private=True)
- def sub(
- a: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- b: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- compute: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- ):
- T.func_attr(
- {
- "op_attrs": {"lhs_axis": 0, "op_name": "qnn.subtract",
"rhs_axis": 0},
- "op_pattern": 0,
- "operator_name": "sub",
- "tirx.noalias": True,
- }
- )
- # with Ts.sblock("root"):
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("buffer_A_assumptions"):
- v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6
= Ts.axis.remap(
- "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5,
axis6]
- )
- Ts.reads(a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- Ts.writes()
- T.assume(
- not (
- (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
- or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5)
- )
- or a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6]
- == T.uint8(0)
- )
-
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("buffer_B_assumptions"):
- v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6
= Ts.axis.remap(
- "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5,
axis6]
- )
- Ts.reads(b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- Ts.writes()
- T.assume(
- not (
- (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
- or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5)
- )
- or b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6]
- == T.uint8(0)
- )
-
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("compute"):
- v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6
= Ts.axis.remap(
- "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5,
axis6]
- )
- Ts.reads(
- a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6],
- b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6],
- )
- Ts.writes(compute[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- compute[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6] = (
- T.if_then_else(
- (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
- or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5),
- T.uint8(0),
- a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- - b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6],
- )
- )
-
- @R.function
- def main(
- a: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
- b: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
- ) -> R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"):
- out = R.call_tir(
- SubBefore.sub,
- (a, b),
- ty_args=[R.Tensor((1, 4, 4, 16, 8, 8, 32), dtype="uint8")],
- )
- return out
-
-
[email protected]_module
-class SubExpected:
- @Ts.function(private=True)
- def sub(
- a: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- b: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- compute: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- ):
- T.func_attr(
- {
- "op_attrs": {"lhs_axis": 0, "op_name": "qnn.subtract",
"rhs_axis": 0},
- "op_pattern": 0,
- "operator_name": "sub",
- "tirx.noalias": True,
- }
- )
- # with Ts.sblock("root"):
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("buffer_A_assumptions"):
- v_axis0 = Ts.axis.spatial(T.int64(1), T.int64(0))
- v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 =
Ts.axis.remap(
- "SSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6]
- )
- Ts.reads(a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- Ts.writes()
- T.assume(
- (
- (v_axis1 < T.int64(3) or v_axis4 < T.int64(4))
- and (v_axis2 < T.int64(3) or v_axis5 < T.int64(4))
- )
- or a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- == T.uint8(0)
- )
-
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("buffer_B_assumptions"):
- v_axis0 = Ts.axis.spatial(T.int64(1), T.int64(0))
- v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 =
Ts.axis.remap(
- "SSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6]
- )
- Ts.reads(b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- Ts.writes()
- T.assume(
- (
- (v_axis1 < T.int64(3) or v_axis4 < T.int64(4))
- and (v_axis2 < T.int64(3) or v_axis5 < T.int64(4))
- )
- or b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- == T.uint8(0)
- )
-
- for axis0, axis1, axis2, axis3, axis4, axis5_0 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(2)
- ):
- for axis5_1_axis6_fused in T.vectorized(T.int64(128)):
- with Ts.sblock("compute"):
- v_axis0 = Ts.axis.spatial(T.int64(1), T.int64(0))
- v_axis1, v_axis2, v_axis3, v_axis4 = Ts.axis.remap(
- "SSSS", [axis1, axis2, axis3, axis4]
- )
- v_axis5 = Ts.axis.spatial(
- T.int64(8), axis5_0 * T.int64(4) + axis5_1_axis6_fused
// T.int64(32)
- )
- v_axis6 = Ts.axis.spatial(T.int64(32), axis5_1_axis6_fused
% T.int64(32))
- Ts.reads(
- a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6],
- b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6],
- )
- Ts.writes(
- compute[T.int64(0), v_axis1, v_axis2, v_axis3,
v_axis4, v_axis5, v_axis6]
- )
- compute[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6] = (
- a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- - b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- )
-
- @R.function
- def main(
- a: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
- b: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
- ) -> R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"):
- out = R.call_tir(
- SubExpected.sub,
- (a, b),
- ty_args=[R.Tensor((1, 4, 4, 16, 8, 8, 32), dtype="uint8")],
- )
- return out
-
-
[email protected]_module
-class MulBefore:
- @Ts.function(private=True)
- def mul(
- a: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- b: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- compute: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- ):
- T.func_attr(
- {
- "op_attrs": {"lhs_axis": 0, "op_name": "qnn.mul", "rhs_axis":
0},
- "op_pattern": 0,
- "operator_name": "mul",
- "tirx.noalias": True,
- }
- )
- # with Ts.sblock("root"):
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("buffer_A_assumptions"):
- v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6
= Ts.axis.remap(
- "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5,
axis6]
- )
- Ts.reads(a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- Ts.writes()
- T.assume(
- not (
- (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
- or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5)
- )
- or a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6]
- == T.uint8(0)
- )
-
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("buffer_B_assumptions"):
- v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6
= Ts.axis.remap(
- "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5,
axis6]
- )
- Ts.reads(b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- Ts.writes()
- T.assume(
- not (
- (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
- or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5)
- )
- or b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6]
- == T.uint8(0)
- )
-
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("compute"):
- v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6
= Ts.axis.remap(
- "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5,
axis6]
- )
- Ts.reads(
- a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6],
- b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6],
- )
- Ts.writes(compute[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- compute[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5,
v_axis6] = (
- T.if_then_else(
- (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
- or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5),
- T.uint8(0),
- a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- * b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6],
- )
- )
-
- @R.function
- def main(
- a: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
- b: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
- ) -> R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"):
- out = R.call_tir(
- MulBefore.mul,
- (a, b),
- ty_args=[R.Tensor((1, 4, 4, 16, 8, 8, 32), dtype="uint8")],
- )
- return out
-
-
[email protected]_module
-class MulExpected:
- @Ts.function(private=True)
- def mul(
- a: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- b: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- compute: T.Tensor(
- (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)),
- "uint8",
- ),
- ):
- T.func_attr(
- {
- "op_attrs": {"lhs_axis": 0, "op_name": "qnn.mul", "rhs_axis":
0},
- "op_pattern": 0,
- "operator_name": "mul",
- "tirx.noalias": True,
- }
- )
- # with Ts.sblock("root"):
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("buffer_A_assumptions"):
- v_axis0 = Ts.axis.spatial(T.int64(1), T.int64(0))
- v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 =
Ts.axis.remap(
- "SSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6]
- )
- Ts.reads(a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- Ts.writes()
- T.assume(
- (
- (v_axis1 < T.int64(3) or v_axis4 < T.int64(4))
- and (v_axis2 < T.int64(3) or v_axis5 < T.int64(4))
- )
- or a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- == T.uint8(0)
- )
-
- for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(8), T.int64(32)
- ):
- with Ts.sblock("buffer_B_assumptions"):
- v_axis0 = Ts.axis.spatial(T.int64(1), T.int64(0))
- v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 =
Ts.axis.remap(
- "SSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6]
- )
- Ts.reads(b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6])
- Ts.writes()
- T.assume(
- (
- (v_axis1 < T.int64(3) or v_axis4 < T.int64(4))
- and (v_axis2 < T.int64(3) or v_axis5 < T.int64(4))
- )
- or b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- == T.uint8(0)
- )
-
- for axis0, axis1, axis2, axis3, axis4, axis5_0 in T.grid(
- T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8),
T.int64(2)
- ):
- for axis5_1_axis6_fused in T.vectorized(T.int64(128)):
- with Ts.sblock("compute"):
- v_axis0 = Ts.axis.spatial(T.int64(1), T.int64(0))
- v_axis1, v_axis2, v_axis3, v_axis4 = Ts.axis.remap(
- "SSSS", [axis1, axis2, axis3, axis4]
- )
- v_axis5 = Ts.axis.spatial(
- T.int64(8), axis5_0 * T.int64(4) + axis5_1_axis6_fused
// T.int64(32)
- )
- v_axis6 = Ts.axis.spatial(T.int64(32), axis5_1_axis6_fused
% T.int64(32))
- Ts.reads(
- a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6],
- b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6],
- )
- Ts.writes(
- compute[T.int64(0), v_axis1, v_axis2, v_axis3,
v_axis4, v_axis5, v_axis6]
- )
- compute[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6] = (
- a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- * b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4,
v_axis5, v_axis6]
- )
-
- @R.function
- def main(
- a: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
- b: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
- ) -> R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"):
- out = R.call_tir(
- MulExpected.mul,
- (a, b),
- ty_args=[R.Tensor((1, 4, 4, 16, 8, 8, 32), dtype="uint8")],
- )
- return out
-
-
-def test_add_function_overcompute():
- add_after = tvm.s_tir.transform.UseAssumeToReduceBranches()(AddBefore)
- tvm_ffi.structural_equal(add_after["add"], AddExpected["add"],
map_free_vars=True)
-
-
-def test_sub_function_overcompute():
- sub_after = tvm.s_tir.transform.UseAssumeToReduceBranches()(SubBefore)
- tvm_ffi.structural_equal(sub_after["sub"], SubExpected["sub"],
map_free_vars=True)
-
-
-def test_mul_function_overcompute():
- mul_after = tvm.s_tir.transform.UseAssumeToReduceBranches()(MulBefore)
- tvm_ffi.structural_equal(mul_after["mul"], MulExpected["mul"],
map_free_vars=True)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
