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 69e4624b5c [REFACTOR] Unify unsigned-index context and remove unused
helpers (#20711)
69e4624b5c is described below
commit 69e4624b5c9256c962267c6245cc8d97ae355e98
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 13 21:15:56 2026 -0400
[REFACTOR] Unify unsigned-index context and remove unused helpers (#20711)
Consolidate unsigned-index assumption scopes in AllowUIntAsIndex with
standard With hooks and encapsulated thread-local nesting state. Native
and Python contexts share the same scope implementation while preserving
the existing unsigned no-overflow semantics.
Remove unused helper macros and standard-library includes, and put
required includes on their direct consumers.
---
include/tvm/ir/transform.h | 1 +
include/tvm/ir/with_context.h | 1 -
include/tvm/relax/binding_rewrite.h | 4 ---
include/tvm/support/cuda/nvtx.h | 7 -----
include/tvm/sym/analyzer.h | 22 +++++++++++++++
python/tvm/sym/analyzer.py | 4 +--
.../hexagon/codegen/llvm/codegen_hexagon.cc | 6 ++--
src/backend/hexagon/runtime/hexagon_device_api.h | 1 -
src/relax/ir/binding_rewrite.cc | 2 ++
src/relax/transform/kill_after_last_use.cc | 5 ++--
src/runtime/extra/contrib/vllm/cache_kernels.cu | 1 -
src/s_tir/analysis/calculate_allocated_memory.cc | 2 +-
src/sym/analyzer.cc | 23 ++++++++++++----
src/sym/canonical_simplify.cc | 10 +++----
src/sym/const_fold.h | 32 +---------------------
src/sym/rewrite_simplify.cc | 8 +++---
16 files changed, 62 insertions(+), 67 deletions(-)
diff --git a/include/tvm/ir/transform.h b/include/tvm/ir/transform.h
index 7c7dfd1d76..a96c504537 100644
--- a/include/tvm/ir/transform.h
+++ b/include/tvm/ir/transform.h
@@ -36,6 +36,7 @@
#include <tvm/ir/module.h>
#include <tvm/ir/with_context.h>
+#include <functional>
#include <string>
#include <type_traits>
#include <utility>
diff --git a/include/tvm/ir/with_context.h b/include/tvm/ir/with_context.h
index 5c7fe6d0f2..f2e3fcbbbc 100644
--- a/include/tvm/ir/with_context.h
+++ b/include/tvm/ir/with_context.h
@@ -26,7 +26,6 @@
#define TVM_IR_WITH_CONTEXT_H_
#include <exception>
-#include <functional>
#include <memory>
#include <utility>
#include <vector>
diff --git a/include/tvm/relax/binding_rewrite.h
b/include/tvm/relax/binding_rewrite.h
index b0cc5d6385..b5884d19aa 100644
--- a/include/tvm/relax/binding_rewrite.h
+++ b/include/tvm/relax/binding_rewrite.h
@@ -29,11 +29,7 @@
#include <tvm/relax/analysis.h>
#include <tvm/relax/expr.h>
-#include <map>
-#include <set>
-#include <type_traits>
#include <utility>
-#include <vector>
namespace tvm {
namespace relax {
diff --git a/include/tvm/support/cuda/nvtx.h b/include/tvm/support/cuda/nvtx.h
index ef9083cfcd..a2e58523f7 100644
--- a/include/tvm/support/cuda/nvtx.h
+++ b/include/tvm/support/cuda/nvtx.h
@@ -65,13 +65,6 @@ class NVTXScopedRange {
NVTXScopedRange& operator=(NVTXScopedRange&& other) = delete;
};
-#ifdef _MSC_VER
-#define TVM_NVTX_FUNC_SCOPE() ::tvm::support::NVTXScopedRange
_nvtx_func_scope_(__FUNCSIG__);
-#else
-#define TVM_NVTX_FUNC_SCOPE() \
- ::tvm::support::NVTXScopedRange _nvtx_func_scope_(__PRETTY_FUNCTION__);
-#endif
-
} // namespace support
} // namespace tvm
diff --git a/include/tvm/sym/analyzer.h b/include/tvm/sym/analyzer.h
index 146277301b..c0d4988f9a 100644
--- a/include/tvm/sym/analyzer.h
+++ b/include/tvm/sym/analyzer.h
@@ -32,6 +32,7 @@
#include <tvm/ir/with_context.h>
#include <tvm/sym/int_set.h>
+#include <functional>
#include <limits>
#include <memory>
#include <unordered_map>
@@ -923,6 +924,27 @@ class Analyzer : public ffi::ObjectRef {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Analyzer, ffi::ObjectRef,
AnalyzerObj);
};
+/*!
+ * \brief Scoped no-overflow assumption for unsigned index arithmetic.
+ *
+ * Unsigned arithmetic normally wraps modulo 2^bits. Use
+ * With<AllowUIntAsIndex> only when unsigned values are guaranteed not to
+ * overflow, enabling the analyzer's index rewrites within that scope.
+ * The assumption is off by default, thread-local, and supports nested scopes.
+ * Constructing this context alone does not enable the assumption.
+ */
+class AllowUIntAsIndex {
+ public:
+ /*! \brief Whether unsigned index analysis is enabled in the current thread.
*/
+ TVM_TIRX_DLL static bool Enabled();
+
+ private:
+ friend class With<AllowUIntAsIndex>;
+ TVM_TIRX_DLL void EnterWithScope();
+ TVM_TIRX_DLL void ExitWithScope();
+ static thread_local int depth_;
+};
+
/*!
* \brief Constraint context.
*
diff --git a/python/tvm/sym/analyzer.py b/python/tvm/sym/analyzer.py
index baa302b66d..560b45a364 100644
--- a/python/tvm/sym/analyzer.py
+++ b/python/tvm/sym/analyzer.py
@@ -545,8 +545,8 @@ def allow_uint_as_index():
provably stay small), NOT for runtime codegen semantics where unsigned
wraparound is real behavior.
"""
- _ffi_api.EnterAllowUintAsIndex()
+ exit_scope = _ffi_api.EnterAllowUintAsIndex()
try:
yield
finally:
- _ffi_api.ExitAllowUintAsIndex()
+ exit_scope()
diff --git a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
index f4c3917aef..449db34d14 100644
--- a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
+++ b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
@@ -50,12 +50,14 @@
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/function.h>
+#include <algorithm>
#include <cstdio>
#include <cstdlib>
-#include <map>
+#include <cstring>
+#include <functional>
+#include <iterator>
#include <sstream>
#include <string>
-#include <unordered_map>
#include <utility>
#include <vector>
diff --git a/src/backend/hexagon/runtime/hexagon_device_api.h
b/src/backend/hexagon/runtime/hexagon_device_api.h
index dc870a1523..4a1b3e6c83 100644
--- a/src/backend/hexagon/runtime/hexagon_device_api.h
+++ b/src/backend/hexagon/runtime/hexagon_device_api.h
@@ -24,7 +24,6 @@
#include <tvm/runtime/memory_scope.h>
#include <condition_variable>
-#include <map>
#include <memory>
#include <mutex>
#include <string>
diff --git a/src/relax/ir/binding_rewrite.cc b/src/relax/ir/binding_rewrite.cc
index 1e84c88dfa..5ee1dbeaae 100644
--- a/src/relax/ir/binding_rewrite.cc
+++ b/src/relax/ir/binding_rewrite.cc
@@ -32,6 +32,8 @@
#include <functional>
#include <iterator>
+#include <set>
+#include <vector>
#include "../../support/ordered_set.h"
diff --git a/src/relax/transform/kill_after_last_use.cc
b/src/relax/transform/kill_after_last_use.cc
index 9f3701637c..8f0e6096ba 100644
--- a/src/relax/transform/kill_after_last_use.cc
+++ b/src/relax/transform/kill_after_last_use.cc
@@ -29,8 +29,9 @@
#include <tvm/sym/analyzer.h>
#include <tvm/tirx/stmt_functor.h>
-#include <map>
-#include <set>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
#include <vector>
#include "utils.h"
diff --git a/src/runtime/extra/contrib/vllm/cache_kernels.cu
b/src/runtime/extra/contrib/vllm/cache_kernels.cu
index 4e88222054..8af234c1f9 100644
--- a/src/runtime/extra/contrib/vllm/cache_kernels.cu
+++ b/src/runtime/extra/contrib/vllm/cache_kernels.cu
@@ -22,7 +22,6 @@
#include <algorithm>
#include <cassert>
-#include <map>
#include <vector>
namespace vllm {
diff --git a/src/s_tir/analysis/calculate_allocated_memory.cc
b/src/s_tir/analysis/calculate_allocated_memory.cc
index 0efc9924cb..1757f8eb40 100644
--- a/src/s_tir/analysis/calculate_allocated_memory.cc
+++ b/src/s_tir/analysis/calculate_allocated_memory.cc
@@ -37,7 +37,7 @@
#include <tvm/tirx/transform.h>
#include <algorithm>
-#include <map>
+#include <optional>
#include <unordered_map>
namespace tvm {
diff --git a/src/sym/analyzer.cc b/src/sym/analyzer.cc
index 8fad2e0b32..129d0f755e 100644
--- a/src/sym/analyzer.cc
+++ b/src/sym/analyzer.cc
@@ -34,6 +34,17 @@
namespace tvm {
namespace sym {
+thread_local int AllowUIntAsIndex::depth_ = 0;
+
+bool AllowUIntAsIndex::Enabled() { return depth_ > 0; }
+
+void AllowUIntAsIndex::EnterWithScope() { ++depth_; }
+
+void AllowUIntAsIndex::ExitWithScope() {
+ TVM_FFI_ICHECK(depth_ > 0) << "AllowUIntAsIndex exit without a matching
enter";
+ --depth_;
+}
+
AnalyzerObj::AnalyzerObj()
: const_int_bound(this),
modular_set(this),
@@ -343,14 +354,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
[](Analyzer analyzer, const PrimExpr& expr, int strength) {
return analyzer->CanProve(expr,
static_cast<ProofStrength>(strength));
})
- .def("sym.EnterAllowUintAsIndex", []() { ++sym::uint_as_index::g_depth;
})
- .def("sym.ExitAllowUintAsIndex",
+ .def("sym.EnterAllowUintAsIndex",
[]() {
- TVM_FFI_ICHECK(sym::uint_as_index::g_depth > 0)
- << "ExitAllowUintAsIndex without a matching Enter";
- --sym::uint_as_index::g_depth;
+ // With has a noexcept(false) destructor, so it cannot use
make_shared.
+ auto ctx = std::shared_ptr<With<AllowUIntAsIndex>>(new
With<AllowUIntAsIndex>());
+ return ffi::Function::FromPacked(
+ [ctx](ffi::PackedArgs, ffi::Any*) mutable { ctx.reset(); });
})
- .def("sym.GetAllowUintAsIndex", []() { return
sym::uint_as_index::Enabled(); })
+ .def("sym.GetAllowUintAsIndex", &AllowUIntAsIndex::Enabled)
.def("sym.AnalyzerSetMaximumRewriteSteps",
[](Analyzer analyzer, int64_t maximum) {
analyzer->rewrite_simplify.SetMaximumRewriteSteps(maximum);
diff --git a/src/sym/canonical_simplify.cc b/src/sym/canonical_simplify.cc
index f1c24bddce..4477977295 100644
--- a/src/sym/canonical_simplify.cc
+++ b/src/sym/canonical_simplify.cc
@@ -96,10 +96,10 @@ inline PrimExpr DivImpl(PrimExpr a, PrimExpr b, DivMode
mode) {
*/
bool CastIsSafe(PrimType dtype, PrimExpr value, AnalyzerObj* analyzer) {
// uint32/64 is not an index dtype by default; the no-overflow assertion
- // (uint_as_index) admits it for bounded non-negative upcast elimination.
+ // (AllowUIntAsIndex) admits it for bounded non-negative upcast elimination.
const bool is_index_dtype =
IsIndexType(dtype->dtype) ||
- (uint_as_index::Enabled() &&
+ (AllowUIntAsIndex::Enabled() &&
dtype->dtype.code == static_cast<uint8_t>(DLDataTypeCode::kDLUInt) &&
(dtype->dtype.bits == 32 || dtype->dtype.bits == 64) &&
dtype->dtype.lanes == 1);
if (!is_index_dtype) {
@@ -1362,10 +1362,10 @@ UnchangedOr<PrimExpr>
CanonicalSimplifier::Impl::Mutate_(const prim::FloorModNod
UnchangedOr<PrimExpr> CanonicalSimplifier::Impl::Mutate_(const prim::CastNode*
op,
InplaceMode
inplace_mode) {
// The cast reasoning below is index-centric; for unsigned operands it runs
- // only under the caller's no-overflow assertion (uint_as_index).
+ // only under the caller's no-overflow assertion (AllowUIntAsIndex).
if (!IsIndexTypedExpr(op) &&
- !(uint_as_index::Enabled() &&
op->ExprNode::ty.as<PrimTypeNode>()->dtype.code ==
-
static_cast<uint8_t>(DLDataTypeCode::kDLUInt))) {
+ !(AllowUIntAsIndex::Enabled() &&
op->ExprNode::ty.as<PrimTypeNode>()->dtype.code ==
+
static_cast<uint8_t>(DLDataTypeCode::kDLUInt))) {
return Rewriter::Mutate_(op, inplace_mode);
}
// normalize
diff --git a/src/sym/const_fold.h b/src/sym/const_fold.h
index 5da41dea44..e1b86f9d6c 100644
--- a/src/sym/const_fold.h
+++ b/src/sym/const_fold.h
@@ -19,7 +19,7 @@
/*!
* \file const_fold.h
- * \brief Arithmetic aliases for shared constant folding and scoped index
analysis.
+ * \brief Arithmetic aliases for shared constant folding.
*/
#ifndef TVM_SYM_CONST_FOLD_H_
#define TVM_SYM_CONST_FOLD_H_
@@ -27,9 +27,6 @@
#include "../ir/prim/const_fold.h"
#include "int_operator.h"
-#define TVM_SYM_CONST_PROPAGATION(BODY) TVM_PRIM_CONST_PROPAGATION(BODY)
-#define TVM_INDEX_CONST_PROPAGATION(BODY)
TVM_PRIM_INDEX_CONST_PROPAGATION(BODY)
-
namespace tvm {
namespace sym {
@@ -43,33 +40,6 @@ using prim::detail::pos_inf;
using prim::detail::SymbolicLimits;
using prim::detail::TryConstFold;
-/*!
- * \brief Scoped opt-in flag for unsigned index analysis (allow_u32).
- *
- * Unsigned arithmetic wraps modulo 2^bits, so most index rewrite rules are
- * unsound for it in general. When the caller can guarantee values never
- * approach the type limit (e.g. compile-time layout proofs over SMEM
- * offsets), enabling this flag activates a small set of otherwise-unsafe
- * rewrite rules for unsigned operands (see the unsigned blocks in
- * rewrite_simplify.cc). It is OFF by default; thread-local and nested-safe.
- */
-namespace uint_as_index {
-
-inline thread_local int g_depth = 0;
-
-inline bool Enabled() { return g_depth > 0; }
-
-} // namespace uint_as_index
-
-/*! \brief RAII guard enabling uint_as_index mode in the current scope. */
-class AllowUintAsIndexGuard {
- public:
- AllowUintAsIndexGuard() { ++uint_as_index::g_depth; }
- ~AllowUintAsIndexGuard() { --uint_as_index::g_depth; }
- AllowUintAsIndexGuard(const AllowUintAsIndexGuard&) = delete;
- AllowUintAsIndexGuard& operator=(const AllowUintAsIndexGuard&) = delete;
-};
-
} // namespace sym
} // namespace tvm
#endif // TVM_SYM_CONST_FOLD_H_
diff --git a/src/sym/rewrite_simplify.cc b/src/sym/rewrite_simplify.cc
index 94b146deaf..fa59082a4f 100644
--- a/src/sym/rewrite_simplify.cc
+++ b/src/sym/rewrite_simplify.cc
@@ -1192,12 +1192,12 @@ UnchangedOr<PrimExpr>
RewriteSimplifier::Impl::Mutate_(const prim::FloorDivNode*
}
}
- // Enabled only under the caller's no-overflow assertion (uint_as_index):
+ // Enabled only under the caller's no-overflow assertion
(AllowUIntAsIndex):
// these floordiv identities are unsound in general for unsigned
// wraparound (the quotient's high bits are lost when x*c1 overflows), but
// are exact in the asserted no-overflow domain. Each application is
// logged for audit.
- if (uint_as_index::Enabled()) {
+ if (AllowUIntAsIndex::Enabled()) {
if (floordiv(x * c1, c2).Match(ret) && c2.Eval()->value > 0 &&
c1.Eval()->value % c2.Eval()->value == 0) {
LOG(WARNING) << "sym: no-overflow floordiv rule on unsigned expr: " <<
ret;
@@ -1374,11 +1374,11 @@ UnchangedOr<PrimExpr>
RewriteSimplifier::Impl::Mutate_(const prim::FloorModNode*
}
}
- // Under the caller's no-overflow assertion (uint_as_index) the pow2
+ // Under the caller's no-overflow assertion (AllowUIntAsIndex) the pow2
// restriction above disappears: without wraparound the residues are the
// plain integer ones, so these hold for any c2 != 0. Each application is
// logged for audit.
- if (uint_as_index::Enabled()) {
+ if (AllowUIntAsIndex::Enabled()) {
if (floormod(x * c1, c2).Match(ret) && c2.Eval()->value != 0 &&
c1.Eval()->value % c2.Eval()->value == 0) {
LOG(WARNING) << "sym: no-overflow floormod rule on unsigned expr: " <<
ret;
