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 3e50db91c8 [REFACTOR][ARITH] Use StructuralWalk and Map in arith
(#20300)
3e50db91c8 is described below
commit 3e50db91c8b4fc75595304cf06e83d541139833a
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Sep 9 21:46:42 2026 -0400
[REFACTOR][ARITH] Use StructuralWalk and Map in arith (#20300)
Arith only needs expression-rooted substitution and idempotent variable
walks, so the generic structural utilities can replace its remaining
statement-functor calls. This removes the statement-layer dependency
from arith source files and its public bound header.
- Replace Substitute with preorder StructuralMap callbacks.
- Replace the two idempotent PostOrderVisit sites with typed postorder
StructuralWalk callbacks.
- Move DomainTouched, the only real statement consumer, to
s_tir/analysis and follow its FFI namespace mechanically.
The arith conversion sites name no buffer type; TrySolveCondition
excludes TensorLoad through its is_simple gate; and every
substitution-map key is a PrimVar with a PrimType, making type-field
traversal inert.
---
include/tvm/arith/bound.h | 17 ----------
include/tvm/s_tir/analysis.h | 11 +++++++
python/tvm/te/operation.py | 2 +-
src/arith/detect_linear_equation.cc | 15 ++++-----
src/arith/int_constraints.cc | 36 +++++++++++++++-------
src/arith/iter_affine_map.cc | 1 -
src/arith/presburger_set.cc | 22 +++++++------
src/arith/solve_linear_equation.cc | 9 ++++--
src/arith/solve_linear_inequality.cc | 21 +++++++++++--
src/{arith => s_tir/analysis}/domain_touched.cc | 23 ++++++++------
src/s_tir/analysis/identify_memcpy.cc | 7 +++--
.../{arith => s_tir}/test_arith_domain_touched.py | 20 ++++++------
12 files changed, 109 insertions(+), 75 deletions(-)
diff --git a/include/tvm/arith/bound.h b/include/tvm/arith/bound.h
index 69f72eafea..6004dff535 100644
--- a/include/tvm/arith/bound.h
+++ b/include/tvm/arith/bound.h
@@ -26,18 +26,12 @@
#include <tvm/arith/int_set.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/stmt.h>
#include <unordered_map>
namespace tvm {
namespace arith {
-using tirx::Region;
-using tirx::Stmt;
-using tirx::Var;
-using tirx::VarNode;
-
/*!
* \brief Deduce the bound of the target variable in a expression,
* give the domain of each variables. Return undefined IntSet to
@@ -69,17 +63,6 @@ IntSet DeduceBound(PrimExpr v, PrimExpr cond,
const std::unordered_map<const VarNode*, IntSet>& hint_map,
const std::unordered_map<const VarNode*, IntSet>&
relax_map);
-/*!
- * \brief Infer a regular domain that covers all the calls or provides within
the given statement.
- * \param body The given statement.
- * \param buffer The buffer to check the access info.
- * \param consider_loads If loads are considered.
- * \param consider_stores If stores are considered.
- * \return The domain that covers all the calls or provides within the given
statement.
- */
-Region DomainTouched(const Stmt& body, const tirx::BufferVar& buffer, bool
consider_loads,
- bool consider_stores);
-
} // namespace arith
} // namespace tvm
#endif // TVM_ARITH_BOUND_H_
diff --git a/include/tvm/s_tir/analysis.h b/include/tvm/s_tir/analysis.h
index d10818e191..a44c9fed82 100644
--- a/include/tvm/s_tir/analysis.h
+++ b/include/tvm/s_tir/analysis.h
@@ -144,6 +144,17 @@ struct MemCpyDetails {
TVM_DLL std::optional<MemCpyDetails> IdentifyMemCpy(const For& loop,
const arith::Analyzer&
analyzer);
+/*!
+ * \brief Infer the domain touched by buffer accesses within a statement.
+ * \param body The statement to analyze.
+ * \param buffer The buffer whose accesses are analyzed.
+ * \param consider_loads Whether to include loads.
+ * \param consider_stores Whether to include stores.
+ * \return The domain covering the selected accesses.
+ */
+TVM_DLL Region DomainTouched(const Stmt& body, const BufferVar& buffer, bool
consider_loads,
+ bool consider_stores);
+
/*!
* \brief Calculate the allocated memory per scope in bytes needed inside the
TIR PrimFunc
* \param func The TIR PrimFunc for which the allocated memory size to be
calculated
diff --git a/python/tvm/te/operation.py b/python/tvm/te/operation.py
index f7d88f8131..fe69b3e42b 100644
--- a/python/tvm/te/operation.py
+++ b/python/tvm/te/operation.py
@@ -390,7 +390,7 @@ def extern_primfunc(input_tensors: list[_tensor.Tensor],
primfunc: tvm.tirx.Prim
"""
# Preserve the function parameter order while selecting BufferType
annotations.
- dt_access_map = tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc)
+ dt_access_map = tvm.s_tir._ffi_api.DomainTouchedAccessMap(primfunc)
ordered_buffers = [param for param in primfunc.params if
tvm.tirx.is_buffer_var(param)]
in_buffers = [buf for buf in ordered_buffers if len(dt_access_map[buf][0])]
out_buffers = [buf for buf in ordered_buffers if
len(dt_access_map[buf][1])]
diff --git a/src/arith/detect_linear_equation.cc
b/src/arith/detect_linear_equation.cc
index 3855dcfbe8..f00bb889e6 100644
--- a/src/arith/detect_linear_equation.cc
+++ b/src/arith/detect_linear_equation.cc
@@ -22,13 +22,13 @@
* \brief Utility to detect patterns in the expression.
*/
#include <tvm/arith/analyzer.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/tirx/analysis.h>
#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
namespace tvm {
namespace arith {
@@ -174,22 +174,23 @@ bool DetectClipBound(const PrimExpr& cond,
std::unordered_map<const VarNode*, IntervalEntry>* bmap) {
int flag = 0;
PrimVar var;
- auto fvisit = [&bmap, &flag, &var](const ffi::ObjectRef& n) {
- if (auto prim_var = n.as<PrimVar>()) {
- const VarNode* v = prim_var->get();
- if (bmap->count(v)) {
+ auto fvisit = [&bmap, &flag, &var](const Var& v) ->
ffi::Expected<ffi::WalkResult> {
+ if (auto prim_var = v.as<PrimVar>()) {
+ const VarNode* var_node = prim_var->get();
+ if (bmap->count(var_node)) {
if (flag == 0) {
var = *prim_var;
flag = 1;
} else if (flag == 1) {
- if (!var.same_as(n)) {
+ if (!var.same_as(*prim_var)) {
flag = -1;
}
}
}
}
+ return ffi::WalkResult::Advance();
};
- PostOrderVisit(cond, fvisit);
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(cond, fvisit);
if (flag != 1) return false;
// canonical form: exp >= 0
bool is_eq = false;
diff --git a/src/arith/int_constraints.cc b/src/arith/int_constraints.cc
index f615654cd7..0d42fb0d50 100644
--- a/src/arith/int_constraints.cc
+++ b/src/arith/int_constraints.cc
@@ -23,19 +23,17 @@
*/
#include <tvm/arith/analyzer.h>
#include <tvm/arith/int_solver.h>
+#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <algorithm>
#include <unordered_map>
#include <utility>
-#include "../tirx/transform/ir_utils.h"
-
namespace tvm {
namespace arith {
@@ -119,11 +117,16 @@ IntGroupBounds IntGroupBounds::operator+(const Range& r) {
}
IntGroupBounds IntGroupBounds::Substitute(const ffi::Map<Var, PrimExpr>&
subst) const {
- auto apply_fun = [&subst](const PrimExpr& e) { return tirx::Substitute(e,
subst); };
- return IntGroupBounds(tirx::Substitute(operator->()->coef, subst),
- tirx::UpdateArray(operator->()->lower, apply_fun),
- tirx::UpdateArray(operator->()->equal, apply_fun),
- tirx::UpdateArray(operator->()->upper, apply_fun));
+ auto f_subst = [&subst](const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
+ if (auto repl = subst.Get(var)) return ffi::Any(repl.value());
+ return ffi::Unchanged();
+ };
+ auto apply_fun = [&f_subst](const PrimExpr& e) {
+ return ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(e,
f_subst).cast<PrimExpr>();
+ };
+ return IntGroupBounds(apply_fun(operator->()->coef),
operator->()->lower.Map(apply_fun),
+
operator->()->equal.Map(apply_fun),
+
operator->()->upper.Map(apply_fun));
}
Range IntGroupBounds::FindBestRange(const ffi::Map<Var, Range>& vranges_addl)
const {
@@ -272,15 +275,26 @@ IntConstraintsTransform
IntConstraintsTransform::operator+(
Analyzer ana_first;
ana_first->Bind(operator->()->src->ranges);
+ auto f_dst_to_src = [this](const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
+ if (auto repl = operator->()->dst_to_src.Get(var)) return
ffi::Any(repl.value());
+ return ffi::Unchanged();
+ };
for (auto p : other->dst_to_src) {
- dst_to_src.Set(p.first,
- ana_first->Simplify(tirx::Substitute(p.second,
operator->()->dst_to_src)));
+ dst_to_src.Set(p.first,
ana_first->Simplify(ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(
+ p.second, f_dst_to_src)
+ .cast<PrimExpr>()));
}
Analyzer ana_second;
ana_second->Bind(other->dst->ranges);
+ auto f_src_to_dst = [&other](const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
+ if (auto repl = other->src_to_dst.Get(var)) return ffi::Any(repl.value());
+ return ffi::Unchanged();
+ };
for (auto p : operator->()->src_to_dst) {
- src_to_dst.Set(p.first, ana_second->Simplify(tirx::Substitute(p.second,
other->src_to_dst)));
+ src_to_dst.Set(p.first,
ana_second->Simplify(ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(
+ p.second, f_src_to_dst)
+ .cast<PrimExpr>()));
}
return IntConstraintsTransform(operator->()->src, other->dst, src_to_dst,
dst_to_src);
}
diff --git a/src/arith/iter_affine_map.cc b/src/arith/iter_affine_map.cc
index 2a766dccf8..3b8daaa7f4 100644
--- a/src/arith/iter_affine_map.cc
+++ b/src/arith/iter_affine_map.cc
@@ -28,7 +28,6 @@
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <utility>
diff --git a/src/arith/presburger_set.cc b/src/arith/presburger_set.cc
index 2076aca214..ab4cc2048d 100644
--- a/src/arith/presburger_set.cc
+++ b/src/arith/presburger_set.cc
@@ -27,11 +27,11 @@
#include <tvm/arith/int_solver.h>
#include <tvm/arith/pattern.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/tirx/expr_functor.h>
-#include <tvm/tirx/stmt_functor.h>
#include <algorithm>
#include <unordered_map>
@@ -94,15 +94,17 @@ static void Update(const PrimExpr& constraint,
PresburgerSetNode* intset) {
PresburgerSet::PresburgerSet(const PrimExpr& constraint) {
ffi::Array<PrimVar> vars;
- PostOrderVisit(constraint, [&vars](const ffi::ObjectRef& obj) {
- if (auto prim_var = obj.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);
- }
- }
- });
+ 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);
diff --git a/src/arith/solve_linear_equation.cc
b/src/arith/solve_linear_equation.cc
index a0bb501e6f..90c3404e78 100644
--- a/src/arith/solve_linear_equation.cc
+++ b/src/arith/solve_linear_equation.cc
@@ -25,12 +25,12 @@
#include <tvm/arith/int_solver.h>
#include <tvm/arith/pattern.h>
#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <unordered_set>
@@ -449,8 +449,13 @@ IntConstraintsTransform SolveLinearEquations(const
IntConstraints& system_to_sol
}
// Add the rest conditions
+ auto f_subst = [&old_to_new_map](const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
+ if (auto repl = old_to_new_map.Get(var)) return ffi::Any(repl.value());
+ return ffi::Unchanged();
+ };
for (const PrimExpr& cond : rest) {
- new_relations.push_back(tirx::Substitute(cond, old_to_new_map));
+ new_relations.push_back(
+ ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(cond,
f_subst).cast<PrimExpr>());
}
IntConstraints solution(new_vars, new_ranges, new_relations);
diff --git a/src/arith/solve_linear_inequality.cc
b/src/arith/solve_linear_inequality.cc
index 61885aaf11..44b97a30ab 100644
--- a/src/arith/solve_linear_inequality.cc
+++ b/src/arith/solve_linear_inequality.cc
@@ -25,12 +25,13 @@
#include <tvm/arith/int_solver.h>
#include <tvm/arith/pattern.h>
#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/tirx/analysis.h>
+#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include "int_operator.h"
@@ -467,6 +468,12 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const
IntConstraints& inequ
}
analyzer->Bind(vranges);
+ auto f_dst_to_src =
+ [&res_dst_to_src](const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
+ if (auto repl = res_dst_to_src.Get(var)) return ffi::Any(repl.value());
+ return ffi::Unchanged();
+ };
+
// We process variables in the reverse direction to start with the most
independent one.
// This order is needed to compute new ranges.
for (auto it = inequalities->variables.rbegin(); it !=
inequalities->variables.rend(); ++it) {
@@ -511,7 +518,9 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const
IntConstraints& inequ
// that is we have to substitute new with old in best_range here
res_dst_to_src.Set(new_var,
analyzer->Simplify(var.as_or_throw<PrimExpr>() -
- Substitute(best_range->min,
res_dst_to_src)));
+
ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(
+ best_range->min,
f_dst_to_src)
+ .cast<PrimExpr>()));
// Add the new var to the resulting axis
auto range = Range(IntImm(new_var->ty.as_or_throw<PrimType>(), 0),
best_range->extent);
@@ -525,9 +534,15 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const
IntConstraints& inequ
}
// Add the original conditions (with variables substituted) to the resulting
conditions
+ auto f_src_to_dst =
+ [&res_src_to_dst](const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
+ if (auto repl = res_src_to_dst.Get(var)) return ffi::Any(repl.value());
+ return ffi::Unchanged();
+ };
for (const PrimExpr& old_cond :
AsConditions(inequalities->variables, solved_bounds,
solved_other_relations)) {
- PrimExpr new_cond = analyzer->Simplify(Substitute(old_cond,
res_src_to_dst));
+ PrimExpr new_cond = analyzer->Simplify(
+ ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(old_cond,
f_src_to_dst).cast<PrimExpr>());
if (!is_const_int(new_cond, 1)) {
// those not represented in vranges (res_ranges)
res_relations.push_back(new_cond);
diff --git a/src/arith/domain_touched.cc b/src/s_tir/analysis/domain_touched.cc
similarity index 90%
rename from src/arith/domain_touched.cc
rename to src/s_tir/analysis/domain_touched.cc
index be69aaac9c..f3d22c1ecc 100644
--- a/src/arith/domain_touched.cc
+++ b/src/s_tir/analysis/domain_touched.cc
@@ -18,13 +18,15 @@
*/
/*!
- * \file bound_deducer.cc
- * \brief Utility to deduce bound of expression
+ * \file domain_touched.cc
+ * \brief Analyze buffer domains touched by a statement
*/
+#include <tvm/arith/int_set.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/logging.h>
+#include <tvm/s_tir/analysis.h>
#include <tvm/te/tensor.h>
#include <tvm/tirx/stmt_functor.h>
@@ -32,12 +34,13 @@
#include <unordered_map>
#include <unordered_set>
-#include "ir_visitor_with_analyzer.h"
+#include "../../arith/ir_visitor_with_analyzer.h"
namespace tvm {
-namespace arith {
+namespace s_tir {
using namespace tirx;
+using arith::IntSet;
namespace {
@@ -60,7 +63,7 @@ using BufferDomainAccess = std::tuple<LoadAccess,
StoreAccess, CombinedAccess>;
} // namespace
// Find Read region of the tensor in the stmt.
-class BufferTouchedDomain final : public IRVisitorWithAnalyzer {
+class BufferTouchedDomain final : public arith::IRVisitorWithAnalyzer {
public:
BufferTouchedDomain(const Stmt& stmt) { operator()(stmt); }
@@ -72,7 +75,7 @@ class BufferTouchedDomain final : public
IRVisitorWithAnalyzer {
Region ret;
auto kv = buffer_access_map_.find(buffer.get());
if (kv == buffer_access_map_.end()) {
- LOG(WARNING) << "[arith::BufferDomainTouched] "
+ LOG(WARNING) << "[s_tir::BufferDomainTouched] "
<< "The requested buffer is not contained in the provided
stmt body: " << buffer;
return ret;
}
@@ -96,7 +99,7 @@ class BufferTouchedDomain final : public
IRVisitorWithAnalyzer {
}
private:
- using Parent = IRVisitorWithAnalyzer;
+ using Parent = arith::IRVisitorWithAnalyzer;
using Parent::VisitExpr_;
using Parent::VisitStmt_;
@@ -171,9 +174,9 @@ ffi::Map<BufferVar, ffi::Array<ffi::ObjectRef>>
DomainTouchedAccessMap(const Pri
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
- .def("arith.DomainTouched", DomainTouched)
- .def("arith.DomainTouchedAccessMap", DomainTouchedAccessMap);
+ .def("s_tir.DomainTouched", DomainTouched)
+ .def("s_tir.DomainTouchedAccessMap", DomainTouchedAccessMap);
}
-} // namespace arith
+} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/analysis/identify_memcpy.cc
b/src/s_tir/analysis/identify_memcpy.cc
index 8beeb42a49..7b94596c65 100644
--- a/src/s_tir/analysis/identify_memcpy.cc
+++ b/src/s_tir/analysis/identify_memcpy.cc
@@ -22,11 +22,12 @@
* \brief Check if a loop nest is equivalent to memcpy
*/
-#include <tvm/arith/bound.h>
+#include <tvm/arith/int_set.h>
#include <tvm/arith/iter_affine_map.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/analysis.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/buffer.h>
#include <tvm/tirx/op.h>
@@ -275,8 +276,8 @@ std::variant<MemCpyDetails, std::string>
IdentifyMemCpyImpl(const For& loop,
BufferRegion src_region(
load->source.as_or_throw<tvm::tirx::BufferVar>(),
- arith::DomainTouched(loop,
load->source.as_or_throw<tvm::tirx::BufferVar>(), true, true));
- BufferRegion dst_region(store->buffer, arith::DomainTouched(loop,
store->buffer, true, true));
+ DomainTouched(loop, load->source.as_or_throw<tvm::tirx::BufferVar>(),
true, true));
+ BufferRegion dst_region(store->buffer, DomainTouched(loop, store->buffer,
true, true));
return MemCpyDetails{src_region, dst_region};
}
diff --git a/tests/python/arith/test_arith_domain_touched.py
b/tests/python/s_tir/test_arith_domain_touched.py
similarity index 81%
rename from tests/python/arith/test_arith_domain_touched.py
rename to tests/python/s_tir/test_arith_domain_touched.py
index a5deb70587..4f3604fe52 100644
--- a/tests/python/arith/test_arith_domain_touched.py
+++ b/tests/python/s_tir/test_arith_domain_touched.py
@@ -37,20 +37,20 @@ def test_domain_touched():
a, b = [var for var in func.params if tvm.tirx.is_buffer_var(var)]
ir = func.body
- a_domain_r = tvm.arith._ffi_api.DomainTouched(ir, a, True, False)
+ a_domain_r = tvm.s_tir._ffi_api.DomainTouched(ir, a, True, False)
assert a_domain_r[0].min.value == -1
assert a_domain_r[0].extent.value == 100
assert a_domain_r[1].min.value == -1
assert a_domain_r[1].extent.name == "m"
- a_domain_w = tvm.arith._ffi_api.DomainTouched(ir, a, False, True)
+ a_domain_w = tvm.s_tir._ffi_api.DomainTouched(ir, a, False, True)
assert a_domain_w[0].min.value == 0
assert a_domain_w[0].extent.value == 100
assert a_domain_w[1].min.value == 0
assert a_domain_w[1].extent.name == "m"
- a_domain_rw = tvm.arith._ffi_api.DomainTouched(ir, a, True, True)
+ a_domain_rw = tvm.s_tir._ffi_api.DomainTouched(ir, a, True, True)
assert a_domain_rw[0].min.value == -1
assert a_domain_rw[0].extent.value == 101
assert a_domain_rw[1].min.value == -1
@@ -58,14 +58,14 @@ def test_domain_touched():
assert a_domain_rw[1].extent.a.name == "m"
assert a_domain_rw[1].extent.b.value == 1
- b_domain_r = tvm.arith._ffi_api.DomainTouched(ir, b, True, False)
+ b_domain_r = tvm.s_tir._ffi_api.DomainTouched(ir, b, True, False)
assert b_domain_r
assert b_domain_r[0].min.value == -1
assert b_domain_r[0].extent.value == 100
assert b_domain_r[1].min.value == 1
assert b_domain_r[1].extent.name == "m"
- b_domain_w = tvm.arith._ffi_api.DomainTouched(ir, b, False, True)
+ b_domain_w = tvm.s_tir._ffi_api.DomainTouched(ir, b, False, True)
assert isinstance(b_domain_w, tvm_ffi.Array)
assert len(b_domain_w) == 0
@@ -84,11 +84,11 @@ def test_domain_touched_vector():
a, b = [var for var in func.params[:2] if tvm.tirx.is_buffer_var(var)]
- assert tvm.arith._ffi_api.DomainTouched(func.body, a, True,
False)[0].extent.value == 128
- assert tvm.arith._ffi_api.DomainTouched(func.body, a, True,
False)[0].extent.value == 128
- assert tvm.arith._ffi_api.DomainTouched(func.body, a, True,
True)[0].extent.value == 128
- assert tvm.arith._ffi_api.DomainTouched(func.body, b, True,
False)[0].extent.value == 128
- assert tvm.arith._ffi_api.DomainTouched(func.body, b, True,
False)[0].extent.value == 128
+ assert tvm.s_tir._ffi_api.DomainTouched(func.body, a, True,
False)[0].extent.value == 128
+ assert tvm.s_tir._ffi_api.DomainTouched(func.body, a, True,
False)[0].extent.value == 128
+ assert tvm.s_tir._ffi_api.DomainTouched(func.body, a, True,
True)[0].extent.value == 128
+ assert tvm.s_tir._ffi_api.DomainTouched(func.body, b, True,
False)[0].extent.value == 128
+ assert tvm.s_tir._ffi_api.DomainTouched(func.body, b, True,
False)[0].extent.value == 128
if __name__ == "__main__":
