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 bf24e37a6e [REFACTOR][TIRx] Make alignment assumptions explicit and
retire legacy bounds checks (#20557)
bf24e37a6e is described below
commit bf24e37a6ecd8e4d9e5e3c648d6f396cf87154e0
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 6 09:59:17 2026 -0400
[REFACTOR][TIRx] Make alignment assumptions explicit and retire legacy
bounds checks (#20557)
Represent tensor alignment facts as an ordinary leaf operation with
typed operands. Keep alignment information with physical aliases and
bound its scope to the paths where the assumption holds. Remove the
obsolete bounds instrumentation pass and its configuration protocol.
---
include/tvm/s_tir/stmt.h | 3 -
include/tvm/s_tir/transform.h | 6 -
include/tvm/tirx/builtin.h | 8 +
include/tvm/tirx/stmt.h | 2 -
python/tvm/s_tir/backend/adreno/pipeline.py | 2 -
python/tvm/s_tir/pipeline.py | 2 -
python/tvm/s_tir/transform/transform.py | 11 -
python/tvm/tirx/__init__.py | 2 +-
python/tvm/tirx/op.py | 10 +
python/tvm/tirx/script/ir_builder/op.py | 2 +
src/s_tir/transform/bound_checker.cc | 299 ---------------------
src/target/llvm/codegen_llvm.cc | 40 ++-
src/target/source/codegen_c.cc | 4 +
src/tirx/ir/transform.cc | 1 -
src/tirx/op/builtin.cc | 10 +
src/tirx/transform/make_packed_api.cc | 7 +-
src/tirx/transform/tvm_ffi_binder.cc | 20 +-
tests/python/codegen/test_target_codegen_llvm.py | 29 --
.../s_tir/script/test_s_tir_script_printer.py | 9 +-
.../test_tir_transform_make_packed_api.py | 4 +-
20 files changed, 78 insertions(+), 393 deletions(-)
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index 5685694ec2..c60da8af42 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -230,9 +230,6 @@ constexpr const char* async_wait_inflight_count =
tirx::attr::async_wait_infligh
*/
constexpr const char* async_scope = "async_scope";
-/*! \brief Mark stores/loads with their bounds. */
-constexpr const char* buffer_bound = "buffer_bound";
-
/*!
* \brief SBlock annotation selecting a write-buffer index for double
buffering in
* InjectSoftwarePipeline.
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index 851f1457d6..19bf6d23f5 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -272,12 +272,6 @@ TVM_DLL Pass RenormalizeSplitPattern();
*/
TVM_DLL Pass RewriteUnsafeSelect();
-/*!
- * \brief Instruments bound checkers.
- * \return The pass.
- */
-TVM_DLL Pass InstrumentBoundCheckers();
-
/*!
* \brief Lower VTCM allocations.
* \return The pass.
diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h
index 784753f1c7..bef5314dcf 100644
--- a/include/tvm/tirx/builtin.h
+++ b/include/tvm/tirx/builtin.h
@@ -621,6 +621,14 @@ TVM_DLL const Op& texture2d_load();
*/
TVM_DLL const Op& assume();
+/*!
+ * \brief Assume a tensor's base address has the given constant byte alignment.
+ *
+ * This leaf operation carries a compiler fact, without checking or changing
+ * the address. Alignment must be a power of two between 1 and 2^27 bytes.
+ */
+TVM_DLL const Op& assume_aligned();
+
/*!
* \brief Returns an initialized but arbitrary value
*
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 86ea304051..2d2bd24736 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -875,8 +875,6 @@ constexpr const char* auto_unroll_max_step =
"auto_unroll_max_step";
* Integer policy inherited by nested loops unless they override it. Consumed
by UnrollLoop.
*/
constexpr const char* unroll_explicit = "unroll_explicit";
-/*! \brief Mark storage alignment requirement of buffers */
-constexpr const char* storage_alignment = "storage_alignment";
/*! \brief Shared execution attributes consumed before and after block
lowering. */
constexpr const char* async_wait_queue_scope = "async_wait_queue_scope";
constexpr const char* async_wait_inflight_count = "async_wait_inflight_count";
diff --git a/python/tvm/s_tir/backend/adreno/pipeline.py
b/python/tvm/s_tir/backend/adreno/pipeline.py
index 913bcdb933..5cac658c12 100644
--- a/python/tvm/s_tir/backend/adreno/pipeline.py
+++ b/python/tvm/s_tir/backend/adreno/pipeline.py
@@ -71,8 +71,6 @@ def default_tir_pipeline(*, prepare_only=False):
]
)
# Additional passes based on configuration.
- if bool(config.get("tirx.instrument_bound_checkers", False)):
- passes.append(s_tir.transform.InstrumentBoundCheckers())
if not bool(config.get("tirx.disable_cse_tir", False)):
passes.append(tirx.transform.CommonSubexprElim())
passes.extend(
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index c989d0c274..df7814ba3b 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -71,8 +71,6 @@ def default_s_tir_pipeline(*, prepare_only=False):
]
)
# Additional passes based on configuration.
- if bool(config.get("tirx.instrument_bound_checkers", False)):
- passes.append(s_tir.transform.InstrumentBoundCheckers())
if not bool(config.get("tirx.disable_cse_tir", False)):
passes.append(tirx.transform.CommonSubexprElim())
passes.extend(
diff --git a/python/tvm/s_tir/transform/transform.py
b/python/tvm/s_tir/transform/transform.py
index 02413bfd31..e89513371a 100644
--- a/python/tvm/s_tir/transform/transform.py
+++ b/python/tvm/s_tir/transform/transform.py
@@ -312,17 +312,6 @@ def RewriteUnsafeSelect():
return _ffi_api.RewriteUnsafeSelect() # type: ignore
-def InstrumentBoundCheckers():
- """Instruments bound checkers.
-
- Returns
- -------
- fpass : tvm.transform.Pass
- The result pass
- """
- return _ffi_api.InstrumentBoundCheckers() # type: ignore
-
-
def VerifyVTCMLimit(default_target=None):
"""Verify if the size of the allocated vtcm memory satisfies the limit.
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 01989b0a35..60da49ed28 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -64,7 +64,7 @@ from .op import CallFFIKernelAttr, call_ffi_kernel,
TensorMapEncodeTiledAttr, te
from .op import call_llvm_intrin, call_llvm_pure_intrin, all, any, min_value,
max_value
from .op import tvm_stack_alloca, tvm_stack_make_shape, tvm_stack_make_array
from .op import handle_add_byte_offset, tvm_struct_get, tvm_struct_set
-from .op import address_of, assume, undef
+from .op import address_of, assume, assume_aligned, undef
from .op import tvm_thread_allreduce, type_annotation, tvm_access_ptr,
ptr_byte_offset
from .op import tvm_throw_last_error, cpu_parallel_barrier
from .op import (
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 6220d308da..3ab04608b7 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -673,6 +673,16 @@ def assume(cond=None):
return call_intrin("bool", "tirx.assume", cond)
+def assume_aligned(tensor, alignment_bytes):
+ """Assume the tensor's base address is aligned to ``alignment_bytes``.
+
+ This compiler fact does not check or modify the address. The tensor must
+ be a tensor variable and alignment a scalar integer constant, a power of
+ two between 1 and 2**27 bytes (inclusive).
+ """
+ return call_intrin("void", "tirx.assume_aligned", tensor, alignment_bytes)
+
+
def undef():
"""Returns an initialized but arbitrary value
diff --git a/python/tvm/tirx/script/ir_builder/op.py
b/python/tvm/tirx/script/ir_builder/op.py
index 1fab8ba767..c7201bc36a 100644
--- a/python/tvm/tirx/script/ir_builder/op.py
+++ b/python/tvm/tirx/script/ir_builder/op.py
@@ -658,6 +658,7 @@ cooperative_tensor_multiply_accumulate =
_tir_op.cooperative_tensor_multiply_acc
assume = _tir_op.assume
+assume_aligned = _tir_op.assume_aligned
undef = _tir_op.undef
@@ -910,6 +911,7 @@ __all__ = [
"asin",
"asinh",
"assume",
+ "assume_aligned",
"atan",
"atan2",
"atanh",
diff --git a/src/s_tir/transform/bound_checker.cc
b/src/s_tir/transform/bound_checker.cc
deleted file mode 100644
index 8aa149e985..0000000000
--- a/src/s_tir/transform/bound_checker.cc
+++ /dev/null
@@ -1,299 +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 bounds_checker.cc
- */
-// Instrument checkers for out of the bounds access.
-
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/builtin.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/s_tir/stmt.h>
-#include <tvm/s_tir/stmt_functor.h>
-#include <tvm/s_tir/transform.h>
-#include <tvm/sym/analyzer.h>
-#include <tvm/tirx/builtin.h>
-#include <tvm/tirx/op.h>
-
-#include <unordered_map>
-#include <utility>
-#include <vector>
-
-namespace tvm {
-namespace s_tir {
-using namespace tvm::prim;
-using namespace tvm::tirx;
-
-// TODO(Lunderberg): Move this pass to be before
-// FlattenBuffer. That will simplify this pass,
-// because it can check directly against the buffer limits.
-class BoundCollector : public StmtExprVisitor {
- public:
- using StmtExprVisitor::Visit_;
- ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
- if (value.as<ExprNode>()) return std::nullopt;
- return StmtExprVisitor::Visit(value);
- }
- BoundCollector() {}
-
- ffi::Optional<VisitInterrupt> Visit_(const AttrStmtNode* op) final {
- if (op->attr_key == s_tir::attr::buffer_bound) {
- const VarNode* key = op->node.as<VarNode>();
- const CallNode* container = op->value.as<CallNode>();
- if (key && container) {
- ffi::Array<PrimExpr> shape =
container->args.as_or_throw<ffi::Array<PrimExpr>>();
- mem_to_shape[key] = shape;
- }
- }
- return StmtExprVisitor::Visit_(op);
- }
- // Hashtable which maps buffer_var to shape.
- std::unordered_map<const VarNode*, ffi::Array<PrimExpr>> mem_to_shape;
-};
-
-class BoundChecker : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
-
- explicit BoundChecker(
- const std::unordered_map<const VarNode*, ffi::Array<PrimExpr>>&
mem_to_shape)
- : mem_to_shape_(mem_to_shape) {}
-
- UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final {
- if (const auto* call = op->value.as<CallNode>();
- call && call->op.same_as(tirx::builtin::alloc_tensor())) {
- return Mutate_AllocTensor(op, call, inplace_mode);
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
-
- UnchangedOr<Stmt> Mutate_AllocTensor(const BindNode* op, const CallNode*
call,
- InplaceMode inplace_mode) {
- if (UpdateIsNeeded(op->var.as_or_throw<TensorVar>().var())) {
- tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
- DLDataType dtype = call->args[1].as_or_throw<DataTypeImm>()->value;
- Update(op->var, shape->fields.as_or_throw<ffi::Array<PrimExpr>>(),
PrimType(dtype));
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
-
- UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode)
final {
- if (process_store_ && op->op.same_as(prim::builtin::if_then_else())) {
- unsafe_rewritten_ = true;
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
-
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode
inplace_mode) final {
- store_scope_bound_collector_.clear();
- process_store_ = true;
- unsafe_rewritten_ = false;
- StmtExprMutator::Mutate_(op, InplaceMode::kDisallow);
- process_store_ = false;
- if (CanInstrument(op->indices, op->buffer.var())) {
- Collect(op->indices, op->buffer.var());
- }
- // The collector should has at least one item.
- if (store_scope_bound_collector_.size()) {
- auto condition = MakeCondition();
- if (condition.has_value()) {
- Stmt then_case = ffi::GetRef<Stmt>(op);
- Stmt else_case =
- AssertStmt(condition.value(), StringImm("RuntimeError"),
{StringImm(error_message_)});
- Stmt body = IfThenElse(condition.value(), then_case, else_case);
- return body;
- }
- }
- return ffi::Unchanged();
- }
-
- UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* op, InplaceMode
inplace_mode) final {
- if (CanInstrument(op->indices,
op->source.as_or_throw<tvm::tirx::TensorVar>().var())) {
- Collect(op->indices,
op->source.as_or_throw<tvm::tirx::TensorVar>().var());
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
-
- private:
- bool UpdateIsNeeded(const Var& buffer_var) const {
- return (buffer_var.defined() && mem_to_shape_.count(buffer_var.get()));
- }
-
- void Update(const Var& buffer_var, ffi::Array<PrimExpr> new_shape, PrimType
dtype) {
- // Sanity check at first.
- if (!ShapeIsValid(new_shape)) {
- return;
- }
-
- int lanes = dtype.lanes();
- TVM_FFI_ICHECK_GE(lanes, 0);
- new_shape.MutateByApply([&](const PrimExpr& dim) {
- // Cast to uint64 to avoid potential overflow.
- return IntImm(PrimType::UInt(64), lanes) * dim;
- });
- mem_to_shape_[buffer_var.get()] = new_shape;
- }
-
- bool ShapeIsValid(const ffi::Array<PrimExpr>& shape) const {
- if (!shape.defined()) {
- return false;
- }
- for (const auto& dim : shape) {
- if (!IsValidScalar(dim) || is_negative_const(dim)) {
- return false;
- }
- }
-
- return true;
- }
-
- bool IndicesAreValid(const ffi::Array<PrimExpr>& indices) const {
- if (!indices.defined()) {
- return false;
- }
-
- for (const auto& index : indices) {
- if (!index.defined()) {
- return false;
- }
-
- if (const RampNode* ramp_index = index.as<RampNode>()) {
- if (!IsValidScalar(ramp_index->base)) {
- return false;
- }
- if (!IsValidScalar(ramp_index->stride)) {
- return false;
- }
- bool lanes_int = ramp_index->lanes->IsInstance<IntImmNode>();
- if (!lanes_int) {
- return false;
- }
- int lanes =
ramp_index->lanes.as_or_throw<IntImm>()->value.as<int>().value();
- if (lanes <= 0) {
- return false;
- }
- }
- }
- return true;
- }
-
- bool IsValidScalar(const PrimExpr& expr) const {
- if (!expr.defined()) return false;
- return expr.ty().IsScalar();
- }
-
- bool CanInstrument(const ffi::Array<PrimExpr>& indices, const Var&
buffer_var) const {
- return buffer_var.defined() && mem_to_shape_.count(buffer_var.get()) &&
- IndicesAreValid(indices) && !unsafe_rewritten_;
- }
-
- void Collect(ffi::Array<PrimExpr> indices, Var buffer_var) {
- store_scope_bound_collector_.push_back(
- std::make_pair(indices, mem_to_shape_[buffer_var.get()]));
- }
-
- ffi::Optional<PrimExpr> MakeCondition() {
- ffi::Optional<PrimExpr> condition;
- for (const auto& pair : store_scope_bound_collector_) {
- ffi::Array<PrimExpr> indices = pair.first;
- ffi::Array<PrimExpr> shape = pair.second;
-
- TVM_FFI_ICHECK_EQ(indices.size(), shape.size())
- << "Mismatch between dimension of physical shape and physical
indices";
-
- for (size_t i = 0; i < indices.size(); i++) {
- PrimExpr index = indices[i];
- PrimExpr upper_bound = shape[i];
-
- if (const RampNode* ramp_index = index.as<RampNode>()) {
- index = ramp_index->base + ramp_index->lanes * ramp_index->stride;
- }
-
- // Try to simplify index and bound.
- index = analyzer_->Simplify(index);
- upper_bound = analyzer_->Simplify(upper_bound);
-
- // Cast to the same type - signed, to be able to check lower bound.
- index = Cast(PrimType::Int(64), index);
- upper_bound = Cast(PrimType::Int(64), upper_bound);
-
- // Looks like a lower bound should always be zero after normalization.
- PrimExpr lower_bound = IntImm::Int64(0);
-
- PrimExpr current_condition = And(GE(index, lower_bound), LT(index,
upper_bound));
- condition =
- condition.has_value() ? And(condition.value(), current_condition)
: current_condition;
- }
- }
- return condition;
- }
-
- // Whether we process store value recursively.
- bool process_store_{false};
- // Whether we face tvm_if_then_else intrinsic.
- bool unsafe_rewritten_{false};
- // Pool which collects the pair of index and shape for specific store/load.
- std::vector<std::pair<ffi::Array<PrimExpr>, ffi::Array<PrimExpr>>>
store_scope_bound_collector_;
- // Error message.
- const char* const error_message_ = "OUT OF THE BOUNDS";
- // Hashtable which maps buffer_var to shape.
- std::unordered_map<const VarNode*, ffi::Array<PrimExpr>> mem_to_shape_;
- // internal analyzer
- sym::Analyzer analyzer_;
-};
-
-Stmt InstrumentBoundCheckers(Stmt stmt) {
- auto bound_collector = ffi::make_object<BoundCollector>();
- // At first walk recursively and collect bound attributes.
- bound_collector->Visit(stmt);
- return ffi::make_object<BoundChecker>(bound_collector->mem_to_shape)
- ->Mutate(stmt, InplaceMode::kAllow)
- .ValueOrUnchanged(std::move(stmt));
-}
-
-namespace transform {
-
-Pass InstrumentBoundCheckers() {
- auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
- auto* n = f.CopyOnWrite();
- auto bound_collector = ffi::make_object<BoundCollector>();
- // At first walk recursively and collect bound attributes.
- bound_collector->Visit(n->body);
- n->body = ffi::make_object<BoundChecker>(bound_collector->mem_to_shape)
- ->Mutate(n->body, InplaceMode::kAllow)
- .ValueOrUnchanged(std::move(n->body));
- return f;
- };
- return CreatePrimFuncPass(pass_func, 0, "s_tir.InstrumentBoundCheckers", {});
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.transform.InstrumentBoundCheckers",
- static_cast<Pass (*)()>(InstrumentBoundCheckers));
-}
-
-} // namespace transform
-
-} // namespace s_tir
-} // namespace tvm
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 5b1119f2ef..33a00a7d10 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -345,18 +345,9 @@ void CodeGenLLVM::AddFunctionInternal(const GlobalVar&
gvar, const PrimFunc& f)
builder_->SetInsertPoint(entry);
this->Dispatch(f->body.value());
- // Add alignment attribute if needed.
- for (size_t i = 0; i < f->params.size(); ++i) {
- const Var& var = f->params[i];
- auto f = alloc_storage_info_.find(var.get());
- if (f != alloc_storage_info_.end()) {
- unsigned align = f->second.alignment;
- if (align > 1) {
- auto attr = llvm::Attribute::get(*ctx, llvm::Attribute::Alignment,
align);
- function_->addParamAttr(i, attr);
- }
- }
- }
+ // Body-local alignment assumptions need not hold on paths that returned
+ // before reaching them. Keep them as llvm.assume instead of strengthening
+ // the function's parameter contract retroactively.
EmitDebugLocation(f->span);
@@ -879,9 +870,12 @@ void CodeGenLLVM::CreateSerialFor(llvm::Value* begin,
llvm::Value* end, llvm::Va
builder_->SetInsertPoint(for_body);
EmitDebugLocation(body->span);
+ // Facts established inside a loop do not dominate its exit.
+ auto outer_storage_info = alloc_storage_info_;
PushLoopFrame(for_next, for_end);
this->Dispatch(body);
PopLoopFrame();
+ alloc_storage_info_ = std::move(outer_storage_info);
var_map_.erase(loop_var.get());
builder_->CreateBr(for_next);
@@ -1463,6 +1457,13 @@ llvm::Value* CodeGenLLVM::CreateIntrinsic(const
CallNode* op) {
} else if (op->op.same_as(tirx::builtin::atomic_add())) {
// TODO(masahi): Support atomic for CPU backend
TVM_FFI_THROW(InternalError) << "CPU backend does not support atomic add
yet.";
+ } else if (op->op.same_as(tirx::builtin::assume_aligned())) {
+ const VarNode* tensor = args[0].as_or_throw<TensorVar>().get();
+ const VarNode* root = GetBufferPhysicalRoot(tensor);
+ int alignment = args[1].as_or_throw<IntImm>()->value.as<int>().value();
+ StorageInfo& info = alloc_storage_info_[root];
+ info.alignment = std::max(info.alignment, alignment);
+ return builder_->CreateAlignmentAssumption(*data_layout_,
GetVarValue(tensor), alignment);
} else if (op->op.same_as(tirx::builtin::assume())) {
llvm::Value* cond = MakeValue(args[0]);
return builder_->CreateAssumption(cond);
@@ -2119,9 +2120,11 @@ void CodeGenLLVM::Dispatch_(const WhileNode* op) {
builder_->SetInsertPoint(while_cond);
builder_->CreateCondBr(MakeValue(op->condition), while_body, while_merge);
builder_->SetInsertPoint(while_body);
+ auto outer_storage_info = alloc_storage_info_;
PushLoopFrame(while_cond, while_merge);
this->Dispatch(op->body);
PopLoopFrame();
+ alloc_storage_info_ = std::move(outer_storage_info);
builder_->CreateBr(while_cond);
builder_->SetInsertPoint(while_merge);
}
@@ -2159,6 +2162,8 @@ void CodeGenLLVM::Dispatch_(const ContinueNode* op) {
void CodeGenLLVM::Dispatch_(const IfThenElseNode* op) {
EmitDebugLocation(op);
+ // A branch-local assumption cannot strengthen its sibling or the join.
+ auto outer_storage_info = alloc_storage_info_;
llvm::Value* cond = MakeValue(op->condition);
llvm::LLVMContext* ctx = llvm_target_->GetContext();
auto* then_block = llvm::BasicBlock::Create(*ctx, "if_then", function_);
@@ -2169,6 +2174,7 @@ void CodeGenLLVM::Dispatch_(const IfThenElseNode* op) {
builder_->SetInsertPoint(then_block);
this->Dispatch(op->then_case);
builder_->CreateBr(end_block);
+ alloc_storage_info_ = outer_storage_info;
builder_->SetInsertPoint(else_block);
this->Dispatch(op->else_case.value());
builder_->CreateBr(end_block);
@@ -2178,6 +2184,7 @@ void CodeGenLLVM::Dispatch_(const IfThenElseNode* op) {
this->Dispatch(op->then_case);
builder_->CreateBr(end_block);
}
+ alloc_storage_info_ = std::move(outer_storage_info);
builder_->SetInsertPoint(end_block);
}
@@ -2258,15 +2265,6 @@ void CodeGenLLVM::Dispatch_(const AttrStmtNode* op) {
TVM_FFI_CHECK(op->attr_key != "thread_extent" && op->attr_key !=
"virtual_thread", ValueError)
<< "Launch attributes are retired; use tirx.launch_thread RegionStmt";
EmitDebugLocation(op);
- if (op->attr_key == tirx::attr::storage_alignment) {
- const VarNode* v = op->node.as<VarNode>();
- TVM_FFI_ICHECK(v);
- alloc_storage_info_[v].alignment =
op->value.as<IntImmNode>()->value.as<int>().value();
- if (var_map_.count(v) && alloc_storage_info_[v].alignment > 1) {
- builder_->CreateAlignmentAssumption(*data_layout_, GetVarValue(v),
- alloc_storage_info_[v].alignment);
- }
- }
this->Dispatch(op->body);
}
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index a0fdc0efba..aa6313b76b 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -1478,6 +1478,10 @@ void CodeGenC::Dispatch_(const EvaluateNode* op) {
if (auto value = op->value.as<PrimExpr>(); value &&
is_const_int(value.value())) return;
const CallNode* call = op->value.as<CallNode>();
if (call) {
+ if (call->op.same_as(tirx::builtin::assume_aligned())) {
+ // Alignment facts do not require a runtime statement on C-family
targets.
+ return;
+ }
if (call->op.same_as(tirx::builtin::tvm_storage_sync())) {
this->PrintStorageSync(call);
return;
diff --git a/src/tirx/ir/transform.cc b/src/tirx/ir/transform.cc
index 78bcf8f8eb..4662a52b8b 100644
--- a/src/tirx/ir/transform.cc
+++ b/src/tirx/ir/transform.cc
@@ -33,7 +33,6 @@ namespace transform {
// Register build pipeline related options
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.noalias", bool);
-TVM_REGISTER_PASS_CONFIG_OPTION("tirx.instrument_bound_checkers", bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.disable_assert", bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.disable_vectorize", bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.enable_buffer_level_predication", bool);
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index 1974e79cc0..bbc78eba31 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -246,6 +246,7 @@ TVM_DEFINE_CACHED_OP_GETTER(nd_mem_alloc_with_scope,
"tirx.nd_mem_alloc_with_sco
TVM_DEFINE_CACHED_OP_GETTER(texture2d_store, "tirx.texture2d_store")
TVM_DEFINE_CACHED_OP_GETTER(texture2d_load, "tirx.texture2d_load")
TVM_DEFINE_CACHED_OP_GETTER(assume, "tirx.assume")
+TVM_DEFINE_CACHED_OP_GETTER(assume_aligned, "tirx.assume_aligned")
TVM_DEFINE_CACHED_OP_GETTER(undef, "tirx.undef")
TVM_DEFINE_CACHED_OP_GETTER(get_active_lane_mask, "tirx.get_active_lane_mask")
TVM_DEFINE_CACHED_OP_GETTER(masked_load, "tirx.masked_load")
@@ -734,6 +735,15 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kEmbedInfo));
+ OpDef("tirx.assume_aligned")
+ .signature(sig::arg<TensorVar>("tensor", "The tensor whose base address
is aligned."),
+ sig::arg<IntImm>("alignment_bytes", "The constant byte
alignment."))
+ .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void())
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.assume_aligned"))
+ .set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
+ .set_attr<TCallEffectKind>("TCallEffectKind",
+
static_cast<int64_t>(CallEffectKind::kEmbedInfo));
+
OpDef("tirx.undef")
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.undef"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
diff --git a/src/tirx/transform/make_packed_api.cc
b/src/tirx/transform/make_packed_api.cc
index 1762a884a1..7ccca70c4c 100644
--- a/src/tirx/transform/make_packed_api.cc
+++ b/src/tirx/transform/make_packed_api.cc
@@ -284,9 +284,10 @@ PrimFunc MakePackedAPI(PrimFunc func) {
// Return error code of zero on success
body = SeqStmt({body, Return(IntImm::Int32(0))});
- body = MergeNest({std::move(result.init_nest), seq_check,
std::move(result.asserts),
- std::move(result.decl_buffers)},
- body);
+ // Tensor declarations and alignment assumptions are ordinary statements,
+ // not scopes with a body hole for MergeNest to fill.
+ body = SeqStmt::Flatten(result.decl_buffers, body);
+ body = MergeNest({std::move(result.init_nest), seq_check,
std::move(result.asserts)}, body);
func_ptr->body = body;
func_ptr->params = args;
diff --git a/src/tirx/transform/tvm_ffi_binder.cc
b/src/tirx/transform/tvm_ffi_binder.cc
index 25c78c0f03..08ca86da02 100644
--- a/src/tirx/transform/tvm_ffi_binder.cc
+++ b/src/tirx/transform/tvm_ffi_binder.cc
@@ -626,6 +626,18 @@ void TVMFFIABIBuilder::DecodeAllParams() {
{data, tvm::Tuple(buffer.value()->shape),
DataTypeImm(buffer.value()->dtype->dtype),
StringImm(buffer.value().scope())},
{})));
+ PrimExpr size = IntImm(PrimType(buffer.value()->DefaultIndexType()), 1);
+ for (const auto& extent : buffer.value()->shape) size *= extent;
+ Stmt alignment =
+ Evaluate(Call(PrimType::Void(), builtin::assume_aligned(),
+ {buffer.value(),
IntImm::Int32(buffer.value()->data_alignment)}));
+ // Empty tensors may have arbitrary data pointers. Their checks must not
+ // be strengthened into an unconditional alignment promise.
+ if (is_const_int(size)) {
+ if (!is_zero(size)) decl_buffers_.push_back(alignment);
+ } else {
+ decl_buffers_.push_back(IfThenElse(size != 0, alignment));
+ }
}
}
}
@@ -869,14 +881,6 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const
TensorVar& buffer, const PrimEx
StringImm(std::to_string(buffer->data_alignment)),
StringImm(" bytes")})));
}
- // mark alignment of external bufs — must be after the alignment
assertion
- // so the compiler does not emit aligned loads before the check fires.
- asserts_.emplace_back(AttrStmt(buffer.var(),
tirx::attr::storage_alignment,
- IntImm::Int32(buffer->data_alignment),
Evaluate(0)));
- } else {
- // Even without alignment check, mark alignment for the compiler.
- init_nest_.emplace_back(AttrStmt(buffer.var(),
tirx::attr::storage_alignment,
-
IntImm::Int32(buffer->data_alignment), Evaluate(0)));
}
}
return typed_data;
diff --git a/tests/python/codegen/test_target_codegen_llvm.py
b/tests/python/codegen/test_target_codegen_llvm.py
index cf3548c775..c3fa1331df 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -407,35 +407,6 @@ def test_rank_zero():
tvm.testing.assert_allclose(d.numpy(), d_np)
[email protected](not env.has_llvm(), reason="need llvm")
-def test_rank_zero_bound_checkers():
- @I.ir_module
- class Module:
- @T.prim_func
- def main(
- A: T.Tensor((64,), "float32"),
- scale: T.Tensor((), "float32"),
- compute: T.Tensor((), "float32"),
- ):
- T.func_attr({"tirx.noalias": True})
- C = T.alloc_tensor(())
- C[()] = T.float32(0.0)
- for k in range(64):
- C[()] = C[()] + A[k] * scale[()]
- compute[()] = C[()] + T.float32(1.0)
-
- n = 64
- with tvm.transform.PassContext(config={"tirx.instrument_bound_checkers":
True}):
- f = tvm.compile(Module, target="llvm")
- dev = tvm.cpu(0)
- a = tvm.runtime.tensor(np.random.randint(0, 2,
size=(n,)).astype("float32"), dev)
- sc = tvm.runtime.tensor(np.random.randint(0, 2,
size=()).astype("float32"), dev)
- d = tvm.runtime.empty((), "float32", dev)
- f(a, sc, d)
- d_np = np.sum(a.numpy()) * sc.numpy() + 1
- tvm.testing.assert_allclose(d.numpy(), d_np)
-
-
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_alignment():
@I.ir_module
diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py
b/tests/python/s_tir/script/test_s_tir_script_printer.py
index e3787402c5..7d9e1ec37b 100644
--- a/tests/python/s_tir/script/test_s_tir_script_printer.py
+++ b/tests/python/s_tir/script/test_s_tir_script_printer.py
@@ -1075,7 +1075,8 @@ def opt_conv_tensorcore_mod_host():
arg2: T.let[T.handle] = T.tvm_struct_get(args, 2, 12, dtype="handle")
A: T.let[T.handle] = T.tvm_struct_get(arg0, 0, 1, dtype="handle")
- T.attr(A, "storage_alignment", 128)
+ A_tensor = T.decl_tensor([1], "float16", data=A)
+ T.assume_aligned(A_tensor, 128)
arg0_shape_data: T.let[T.handle("int64")] = T.tvm_struct_get(
arg0, 0, 2, dtype=T.handle("int64").ty
)
@@ -1088,10 +1089,12 @@ def opt_conv_tensorcore_mod_host():
dev_id: T.let[T.int32] = T.tvm_struct_get(arg0, 0, 9, dtype="int32")
W: T.let[T.handle] = T.tvm_struct_get(arg1, 0, 1, dtype="handle")
- T.attr(W, "storage_alignment", 128)
+ W_tensor = T.decl_tensor([1], "float16", data=W)
+ T.assume_aligned(W_tensor, 128)
Conv: T.let[T.handle] = T.tvm_struct_get(arg2, 0, 1, dtype="handle")
- T.attr(Conv, "storage_alignment", 128)
+ Conv_tensor = T.decl_tensor([1], "float32", data=Conv)
+ T.assume_aligned(Conv_tensor, 128)
assert (((arg0_code == 3) or (arg0_code == 13)) or (arg0_code == 7))
or (arg0_code == 4), (
"default_function: Expect arg[0] to be pointer"
diff --git a/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
b/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
index 809e9a3f6b..bcb5ce14b9 100644
--- a/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
+++ b/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
@@ -487,8 +487,8 @@ def test_buffer_alignment_attached_to_buffer_var():
declared_buffers = []
def collect(node):
- if isinstance(node, tirx.AttrStmt) and node.attr_key ==
"storage_alignment":
- alignment_nodes.append(node.node)
+ if isinstance(node, tvm.ir.Call) and node.op ==
tvm.ir.Op.get("tirx.assume_aligned"):
+ alignment_nodes.append(node.args[0])
if _is_buffer_binding(node, "tirx.decl_tensor"):
declared_buffers.append(node.var)
