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 81bb3dcfe6 [REFACTOR][TIR] Split masked buffer access into special
calls (#20244)
81bb3dcfe6 is described below
commit 81bb3dcfe6bcacee62a9639558fea5bd6eb4cc90
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Aug 31 10:06:24 2026 -0400
[REFACTOR][TIR] Split masked buffer access into special calls (#20244)
Move lane-mask semantics out of ordinary BufferLoad and BufferStore
nodes into dedicated tirx.masked_load and tirx.masked_store calls.
Construct the intrinsic Calls directly, update the semantic consumers
that require masked-memory knowledge, use BufferType-based reasoning
elsewhere, and keep LLVM lowering in dedicated masked-load/store helpers
while rejecting unsupported source backends.
Validation:
- LLVM-enabled build
- 756 focused and broad masked/script/LLVM/legalization tests passed, 10
skipped
- Clang-format, Ruff, and git diff --check passed
---
include/tvm/tirx/buffer.h | 10 +-
include/tvm/tirx/builtin.h | 16 +++
include/tvm/tirx/expr.h | 8 +-
include/tvm/tirx/script/builder/ir.h | 5 +-
include/tvm/tirx/stmt.h | 6 +-
python/tvm/tirx/buffer.py | 17 +--
python/tvm/tirx/expr.py | 5 -
python/tvm/tirx/expr_functor.py | 2 +-
python/tvm/tirx/op.py | 46 +++++++
python/tvm/tirx/script/builder/ir.py | 11 +-
python/tvm/tirx/script/parser/parser.py | 2 +-
python/tvm/tirx/stmt.py | 8 --
python/tvm/tirx/stmt_functor.py | 8 +-
src/backend/vulkan/codegen/codegen_spirv.cc | 6 +-
src/backend/webgpu/codegen/codegen_webgpu.cc | 6 +-
src/relax/ir/expr_functor.cc | 13 +-
.../analysis/sblock_access_region_detector.cc | 28 +++++
src/s_tir/transform/inject_virtual_thread.cc | 37 +++++-
src/s_tir/transform/lower_match_buffer.cc | 14 ++-
.../manifest_shared_memory_local_stage.cc | 3 -
src/s_tir/transform/storage_access.cc | 24 +++-
src/target/llvm/codegen_llvm.cc | 69 ++++++++++-
src/target/llvm/codegen_llvm.h | 2 +
src/target/source/codegen_c.cc | 6 +-
src/te/operation/create_primfunc.cc | 6 +-
src/tirx/analysis/deep_equal.cc | 3 +-
src/tirx/analysis/verify_memory.cc | 8 ++
src/tirx/ir/buffer.cc | 15 +--
src/tirx/ir/expr.cc | 33 +----
src/tirx/ir/expr_functor.cc | 2 +-
src/tirx/ir/stmt.cc | 28 +----
src/tirx/op/builtin.cc | 9 ++
src/tirx/script/builder/ir.cc | 5 +-
src/tirx/script/printer/buffer.cc | 21 +---
src/tirx/script/printer/utils.h | 9 +-
src/tirx/transform/flatten_buffer.cc | 15 +++
src/tirx/transform/remove_no_op.cc | 3 +-
src/tirx/transform/storage_rewrite.cc | 122 +++++++++++++-----
src/tirx/transform/unsupported_dtype_legalize.cc | 97 +++++++++++----
src/tirx/transform/vectorize_loop.cc | 68 ++++++++---
tests/python/codegen/test_target_codegen.py | 38 +++++-
tests/python/codegen/test_target_codegen_llvm.py | 29 ++++-
.../s_tir/analysis/test_sblock_access_region.py | 22 ++++
.../test_s_tir_transform_inject_virtual_thread.py | 35 ++++++
.../test_s_tir_transform_lower_match_buffer.py | 16 +++
tests/python/tirx-base/test_tir_nodes.py | 71 +----------
.../test_tir_transform_bf16_legalize.py | 80 ++++++++++++
.../tirx-transform/test_tir_transform_vectorize.py | 136 ++++++++++++++++-----
tests/python/tirx/transform/test_stmt_functor.py | 2 +-
.../tvmscript/test_tvmscript_ir_builder_tir.py | 14 ---
.../python/tvmscript/test_tvmscript_printer_tir.py | 81 +++++++++---
tests/python/tvmscript/test_tvmscript_roundtrip.py | 19 ++-
52 files changed, 930 insertions(+), 409 deletions(-)
diff --git a/include/tvm/tirx/buffer.h b/include/tvm/tirx/buffer.h
index 24cc3b9f1a..c54102276c 100644
--- a/include/tvm/tirx/buffer.h
+++ b/include/tvm/tirx/buffer.h
@@ -232,20 +232,14 @@ class BufferVar : public Var {
* \brief Create an Expr that does a vector load at begin index.
* \param begin The beginning index
* \param dtype The data type to be loaded.
- * \param predicate A vector mask of boolean values indicating which lanes
of a vector are to be
- * loaded. The number lanes of the mask must be equal to the number of lanes
in being loaded.
*/
- TVM_DLL PrimExpr vload(ffi::Array<PrimExpr> begin, PrimType dtype,
- ffi::Optional<PrimExpr> predicate = std::nullopt)
const;
+ TVM_DLL PrimExpr vload(ffi::Array<PrimExpr> begin, PrimType dtype) const;
/*!
* \brief Create a Stmt that does a vector store at begin index.
* \param begin The beginning index
* \param value The value to be stored.
- * \param predicate A vector mask of boolean values indicating which lanes
of a vector are to be
- * stored. The number lanes of the mask must be equal to the number of lanes
in value.
*/
- TVM_DLL Stmt vstore(ffi::Array<PrimExpr> begin, PrimExpr value,
- ffi::Optional<PrimExpr> predicate = std::nullopt) const;
+ TVM_DLL Stmt vstore(ffi::Array<PrimExpr> begin, PrimExpr value) const;
/*!
* \brief Get a flattened version of the buffer.
diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h
index 3026894016..f209b381f7 100644
--- a/include/tvm/tirx/builtin.h
+++ b/include/tvm/tirx/builtin.h
@@ -771,6 +771,22 @@ TVM_DLL const Op& vscale();
*/
TVM_DLL const Op& get_active_lane_mask();
+/*!
+ * \brief Masked buffer load.
+ *
+ * Arguments are the buffer variable, one or more indices, and a trailing
boolean lane mask.
+ * The result type is the vector type loaded from the selected lanes.
+ */
+TVM_DLL const Op& masked_load();
+
+/*!
+ * \brief Masked buffer store.
+ *
+ * Arguments are the buffer variable, value, one or more indices, and a
trailing boolean lane
+ * mask. The result type is void.
+ */
+TVM_DLL const Op& masked_store();
+
/*! \brief Annotate a predicate not be considered as target condition of loop
partition. */
TVM_DLL const Op& ignore_loop_partition();
/*!
diff --git a/include/tvm/tirx/expr.h b/include/tvm/tirx/expr.h
index 4a4d4b7ce8..6696533ccf 100644
--- a/include/tvm/tirx/expr.h
+++ b/include/tvm/tirx/expr.h
@@ -548,14 +548,11 @@ class BufferLoadNode : public ExprNode {
BufferVar buffer;
/*! \brief The indices location to be loaded. */
ffi::Array<PrimExpr> indices;
- /*! \brief The predicate mask for loading values. */
- ffi::Optional<PrimExpr> predicate;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BufferLoadNode>()
.def_ro("buffer", &BufferLoadNode::buffer,
refl::AttachFieldFlag::SEqHashDefRecursive())
- .def_ro("indices", &BufferLoadNode::indices)
- .def_ro("predicate", &BufferLoadNode::predicate);
+ .def_ro("indices", &BufferLoadNode::indices);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferLoad", BufferLoadNode,
ExprNode);
@@ -582,8 +579,7 @@ class BufferLoadNode : public ExprNode {
*/
class BufferLoad : public PrimExpr {
public:
- TVM_DLL explicit BufferLoad(BufferVar buffer, ffi::Array<PrimExpr> indices,
- ffi::Optional<PrimExpr> predicate =
std::nullopt, Span span = Span());
+ TVM_DLL explicit BufferLoad(BufferVar buffer, ffi::Array<PrimExpr> indices,
Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BufferLoad, PrimExpr,
BufferLoadNode);
static constexpr bool _type_container_is_exact = true;
TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferLoadNode);
diff --git a/include/tvm/tirx/script/builder/ir.h
b/include/tvm/tirx/script/builder/ir.h
index fcb6fb4b10..affdbfa690 100644
--- a/include/tvm/tirx/script/builder/ir.h
+++ b/include/tvm/tirx/script/builder/ir.h
@@ -498,11 +498,8 @@ Var EnvThread(ffi::String thread_tag, PrimType dtype =
PrimType::Int(32));
* \param buffer The buffer.
* \param value The value to be stored.
* \param indices The indices location to be stored.
- * \param predicate A vector mask of boolean values indicating which lanes of
a vector are to be
- * stored. The number lanes of the mask must be equal to the number of lanes
in value.
*/
-void BufferStore(BufferVar buffer, PrimExpr value, ffi::Array<PrimExpr>
indices,
- ffi::Optional<PrimExpr> predicate);
+void BufferStore(BufferVar buffer, PrimExpr value, ffi::Array<PrimExpr>
indices);
/*!
* \brief Evaluate the input expression.
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index d95c1af501..85ae71d9fe 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -207,16 +207,13 @@ class BufferStoreNode : public StmtNode {
PrimExpr value;
/*! \brief The indices location to be stored. */
ffi::Array<PrimExpr> indices;
- /*! \brief The predicate mask for storing values. */
- ffi::Optional<PrimExpr> predicate;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BufferStoreNode>()
.def_ro("buffer", &BufferStoreNode::buffer,
refl::AttachFieldFlag::SEqHashDefRecursive())
.def_ro("value", &BufferStoreNode::value)
- .def_ro("indices", &BufferStoreNode::indices)
- .def_ro("predicate", &BufferStoreNode::predicate);
+ .def_ro("indices", &BufferStoreNode::indices);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferStore", BufferStoreNode,
StmtNode);
};
@@ -228,7 +225,6 @@ class BufferStoreNode : public StmtNode {
class BufferStore : public Stmt {
public:
TVM_DLL explicit BufferStore(BufferVar buffer, PrimExpr value,
ffi::Array<PrimExpr> indices,
- ffi::Optional<PrimExpr> predicate =
std::nullopt,
Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BufferStore, Stmt,
BufferStoreNode);
diff --git a/python/tvm/tirx/buffer.py b/python/tvm/tirx/buffer.py
index 376707a97f..d832aa7602 100644
--- a/python/tvm/tirx/buffer.py
+++ b/python/tvm/tirx/buffer.py
@@ -148,7 +148,7 @@ class _BufferMethods:
extent, # type: ignore
)
- def vload(self, begin, dtype=None, predicate=None):
+ def vload(self, begin, dtype=None):
"""Generate an Expr that loads dtype from begin index.
Parameters
@@ -160,10 +160,6 @@ class _BufferMethods:
The data type to be loaded,
can be vector type which have lanes that is multiple of
Buffer.dtype
- predicate : Optional[Expr]
- A vector mask of boolean values indicating which lanes of a vector
are to be
- loaded. The number lanes of the mask must be equal to the number
of lanes being loaded.
-
Returns
-------
load : Expr
@@ -171,9 +167,9 @@ class _BufferMethods:
"""
begin = (begin,) if isinstance(begin, int) or
tvm.ir.is_prim_expr(begin) else begin
dtype = dtype if dtype else self.ty.dtype
- return _ffi_api.BufferVLoad(self, begin, dtype, predicate) # type:
ignore
+ return _ffi_api.BufferVLoad(self, begin, dtype) # type: ignore
- def vstore(self, begin, value, predicate=None):
+ def vstore(self, begin, value):
"""Generate a Stmt that store value into begin index.
Parameters
@@ -184,18 +180,13 @@ class _BufferMethods:
value : Expr
The value to be stored.
- predicate : Optional[Expr]
- A vector mask of boolean values indicating which lanes of a vector
are to be
- stored. The number lanes of the mask must be equal to the number
of lanes in
- value.
-
Returns
-------
store : Stmt
The corresponding store stmt.
"""
begin = (begin,) if isinstance(begin, int) or
tvm.ir.is_prim_expr(begin) else begin
- return _ffi_api.BufferVStore(self, begin, value, predicate) # type:
ignore
+ return _ffi_api.BufferVStore(self, begin, value) # type: ignore
def scope(self):
"""Return the storage scope associated with this buffer.
diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py
index 6e29d94444..978f1acb09 100644
--- a/python/tvm/tirx/expr.py
+++ b/python/tvm/tirx/expr.py
@@ -1186,9 +1186,6 @@ class BufferLoad(ExprWithOp):
span : Optional[Span]
The location of this expression in the source code.
- predicate : Optional[Expr]
- A vector mask of boolean values indicating which lanes of a vector are
to be
- loaded. The number lanes of the mask must be equal to the number of
lanes being loaded.
"""
buffer: Buffer
@@ -1198,14 +1195,12 @@ class BufferLoad(ExprWithOp):
self,
buffer: Buffer,
indices: list[Expr],
- predicate: Expr | None = None,
span: Span | None = None,
) -> None:
self.__init_handle_by_constructor__(
_ffi_api.BufferLoad,
buffer,
indices,
- predicate,
span, # type: ignore
)
diff --git a/python/tvm/tirx/expr_functor.py b/python/tvm/tirx/expr_functor.py
index def3b18bda..4fb502cb8b 100644
--- a/python/tvm/tirx/expr_functor.py
+++ b/python/tvm/tirx/expr_functor.py
@@ -472,7 +472,7 @@ class ExprMutator(ExprFunctor):
if all(old_index is new_index for old_index, new_index in
zip(op.indices, indices)):
return op
else:
- return tvm.tirx.BufferLoad(op.buffer, indices, op.predicate)
+ return tvm.tirx.BufferLoad(op.buffer, indices)
def visit_opaque_expr_(self, op):
"""Mutator implementation for an opaque construction-time
expression."""
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 39939803a6..a4963f54e1 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -3003,6 +3003,52 @@ def get_active_lane_mask(dtype, base, limit):
return call_intrin(dtype, "tirx.get_active_lane_mask", base, limit)
+def masked_load(dtype, buffer, *indices_and_mask):
+ """Load vector lanes selected by a predicate mask.
+
+ Parameters
+ ----------
+ dtype : str
+ The vector data type to load.
+
+ buffer : Buffer
+ The buffer to load.
+
+ indices_and_mask : Expr
+ The buffer indices followed by a boolean lane mask. The mask must
match the
+ lane count and scalability of the loaded vector.
+
+ Returns
+ -------
+ call : Expr
+ A ``tirx.masked_load`` call with result type ``dtype``.
+ """
+ return call_intrin(dtype, "tirx.masked_load", buffer, *indices_and_mask)
+
+
+def masked_store(buffer, value, *indices_and_mask):
+ """Store vector lanes selected by a predicate mask.
+
+ Parameters
+ ----------
+ buffer : Buffer
+ The buffer to update.
+
+ value : Expr
+ The vector value to store.
+
+ indices_and_mask : Expr
+ The buffer indices followed by a boolean lane mask. The mask must
match the
+ lane count and scalability of ``value``.
+
+ Returns
+ -------
+ call : Expr
+ A void-typed ``tirx.masked_store`` call.
+ """
+ return call_intrin("void", "tirx.masked_store", buffer, value,
*indices_and_mask)
+
+
def get_vscale_expr(dtype: str | tvm_ffi.dtype, min_size: int = 128) -> Expr:
"""
Create a datatype dependent scalable expression.
diff --git a/python/tvm/tirx/script/builder/ir.py
b/python/tvm/tirx/script/builder/ir.py
index d7897647fd..e2ed32a925 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -2336,7 +2336,6 @@ def buffer_store(
buffer: Buffer, # pylint: disable=redefined-outer-name
value: Expr,
indices: list[Expr | slice],
- predicate: Expr | None = None,
) -> None:
"""Buffer store node.
@@ -2351,10 +2350,6 @@ def buffer_store(
indices : List[Union[Expr, slice]]
The indices location to be stored.
- predicate : Optional[Expr]
- A vector mask of boolean values indicating which lanes of a vector are
to be
- stored. The number lanes of the mask must be equal to the number of
lanes in
- value.
"""
from tvm.arith import Analyzer # pylint: disable=import-outside-toplevel
@@ -2377,7 +2372,7 @@ def buffer_store(
if isinstance(value, bool) and buffer.ty.dtype == "bool":
value = IntImm("bool", value)
return _ffi_api.BufferStore( # type: ignore[attr-defined] # pylint:
disable=no-member
- buffer, value, expr_indices, predicate
+ buffer, value, expr_indices
)
@@ -3223,6 +3218,8 @@ vectorlow = _dtype_forward(_tir_op.vectorlow)
vectorhigh = _dtype_forward(_tir_op.vectorhigh)
vectorcombine = _dtype_forward(_tir_op.vectorcombine)
get_active_lane_mask = _dtype_forward(_tir_op.get_active_lane_mask)
+masked_load = _dtype_forward(_tir_op.masked_load)
+masked_store = _op_wrapper(_tir_op.masked_store)
dp4a = _dtype_forward(_tir_op.dp4a)
@@ -3589,6 +3586,8 @@ __all__ = [
"Range",
"vscale",
"get_active_lane_mask",
+ "masked_load",
+ "masked_store",
"call_kernel",
"ignore_loop_partition",
]
diff --git a/python/tvm/tirx/script/parser/parser.py
b/python/tvm/tirx/script/parser/parser.py
index 2efef79d07..4e17808970 100644
--- a/python/tvm/tirx/script/parser/parser.py
+++ b/python/tvm/tirx/script/parser/parser.py
@@ -1044,7 +1044,7 @@ def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
# Ignore docstrings
pass
elif isinstance(res, tvm.tirx.stmt.BufferStore):
- T.buffer_store(res.buffer, res.value, res.indices, res.predicate)
+ T.buffer_store(res.buffer, res.value, res.indices)
elif is_buffer_var(res):
# ``T.match_buffer(...)`` used as a bare statement (no LHS) — the
# buffer object is discarded; the underlying side effect (the
diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py
index fe4c98a402..fdccf5ee14 100644
--- a/python/tvm/tirx/stmt.py
+++ b/python/tvm/tirx/stmt.py
@@ -288,11 +288,6 @@ class BufferStore(Stmt):
indices : List[Expr]
The indices location to be stored.
- predicate : Optional[Expr]
- A vector mask of boolean values indicating which lanes of a vector are
to be
- stored. The number lanes of the mask must be equal to the number of
lanes in
- value.
-
span : Optional[Span]
The location of the stmt in the source code.
"""
@@ -300,7 +295,6 @@ class BufferStore(Stmt):
buffer: Buffer
value: Expr
indices: list[Expr]
- predicate: Expr | None
span: Span | None
def __init__(
@@ -308,7 +302,6 @@ class BufferStore(Stmt):
buffer: Buffer,
value: Expr,
indices: list[Expr],
- predicate: Expr | None = None,
span: Span | None = None,
) -> None:
self.__init_handle_by_constructor__(
@@ -316,7 +309,6 @@ class BufferStore(Stmt):
buffer,
value,
indices,
- predicate,
span, # type: ignore
)
diff --git a/python/tvm/tirx/stmt_functor.py b/python/tvm/tirx/stmt_functor.py
index 1f9755b22a..a0a7c4a553 100644
--- a/python/tvm/tirx/stmt_functor.py
+++ b/python/tvm/tirx/stmt_functor.py
@@ -292,8 +292,6 @@ class StmtVisitor(StmtFunctor):
"""Visitor implementation for BufferStore."""
self.visit_expr(op.value)
_visit_array(op.indices, lambda x: self.visit_expr(x))
- if op.predicate is not None:
- self.visit_expr(op.predicate)
def visit_assert_(self, op):
"""Visitor implementation for AssertStmt."""
@@ -546,14 +544,12 @@ class StmtMutator(StmtFunctor):
"""Mutator implementation for BufferStore."""
value = self.visit_expr(op.value)
indices = [self.visit_expr(idx) for idx in op.indices]
- predicate = self.visit_expr(op.predicate) if op.predicate is not None
else None
-
indices_changed = any(old is not new for old, new in zip(op.indices,
indices))
- if value is op.value and not indices_changed and predicate is
op.predicate:
+ if value is op.value and not indices_changed:
return op
- return tvm.tirx.BufferStore(op.buffer, value, indices, predicate,
op.span)
+ return tvm.tirx.BufferStore(op.buffer, value, indices, op.span)
def visit_buffer_realize_(self, op):
"""Mutator implementation for BufferRealize."""
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc
b/src/backend/vulkan/codegen/codegen_spirv.cc
index 69958d7d1b..af8ade1743 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -331,6 +331,10 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const LetNode* op) {
}
spirv::Value CodeGenSPIRV::VisitExpr_(const CallNode* op) {
+ TVM_FFI_ICHECK(!op->op.same_as(builtin::masked_load()))
+ << "Predicated buffer load is not supported.";
+ TVM_FFI_ICHECK(!op->op.same_as(builtin::masked_store()))
+ << "Predicated buffer store is not supported.";
if (op->op.same_as(builtin::buffer_data())) {
TVM_FFI_ICHECK_EQ(op->args.size(), 1U);
return MakeValue(op->args[0]);
@@ -594,7 +598,6 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const BroadcastNode*
op) {
spirv::Value CodeGenSPIRV::VisitExpr_(const BufferLoadNode* op) {
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "SPIR-V codegen expects flat
memory buffers";
- TVM_FFI_ICHECK(!op->predicate.has_value()) << "Predicated buffer load is not
supported.";
Var buffer_var = op->buffer.var();
PrimExpr prim_index = op->indices[0];
@@ -681,7 +684,6 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const ShuffleNode*
op) {
void CodeGenSPIRV::VisitStmt_(const BufferStoreNode* op) {
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "SPIR-V codegen expects flat
memory buffers";
- TVM_FFI_ICHECK(!op->predicate.has_value()) << "Predicated buffer store is
not supported.";
Var buffer_var = op->buffer.var();
PrimExpr prim_index = op->indices[0];
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index a82f9f96f1..5d20f26b2d 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -453,6 +453,10 @@ PrimExpr CodeGenWebGPU::EnforceU32(PrimExpr value) {
}
void CodeGenWebGPU::VisitExpr_(const CallNode* op, std::ostream& os) { //
NOLINT(*)
+ TVM_FFI_ICHECK(!op->op.same_as(builtin::masked_load()))
+ << "Predicated buffer load is not supported.";
+ TVM_FFI_ICHECK(!op->op.same_as(builtin::masked_store()))
+ << "Predicated buffer store is not supported.";
if (op->op.same_as(builtin::reinterpret())) {
// generate bitcast<TYPE>(ARG)
os << "bitcast<";
@@ -583,7 +587,6 @@ void CodeGenWebGPU::VisitExpr_(const BufferLoadNode* op,
std::ostream& os) { //
// to ensure correctness in the case of nested-expression
// do not try to lift common printings from each case
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "Load from non-flat memory not
supported.";
- TVM_FFI_ICHECK(!op->predicate.has_value()) << "Predicated buffer load is not
supported.";
PrimType value_ty = op->ty.as_or_throw<PrimType>();
PrimExpr index = op->indices[0];
@@ -655,7 +658,6 @@ void CodeGenWebGPU::VisitStmt_(const BindNode* op) {
void CodeGenWebGPU::VisitStmt_(const BufferStoreNode* op) {
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "Store to non-flat memory not
supported.";
- TVM_FFI_ICHECK(!op->predicate.has_value()) << "Predicated buffer store is
not supported.";
PrimType value_ty = op->value.ty();
const PrimType& element_ty = op->buffer->dtype;
diff --git a/src/relax/ir/expr_functor.cc b/src/relax/ir/expr_functor.cc
index 76c84cdb4d..bdb16ea1ba 100644
--- a/src/relax/ir/expr_functor.cc
+++ b/src/relax/ir/expr_functor.cc
@@ -180,9 +180,6 @@ void ExprVisitor::VisitExpr_(const tirx::BufferLoadNode*
op) {
for (const PrimExpr& index : op->indices) {
this->VisitExpr(index);
}
- if (op->predicate.has_value()) {
- this->VisitExpr(op->predicate.value());
- }
VisitExprDepTypeFieldIfNeeded(this, op->ty);
}
@@ -536,16 +533,10 @@ Expr ExprMutatorBase::VisitExpr_(const CallNode*
call_node) {
Expr ExprMutatorBase::VisitExpr_(const tirx::BufferLoadNode* op) {
ffi::Array<PrimExpr> indices = op->indices.Map(
[this](const PrimExpr& e) { return
this->VisitExpr(e).as_or_throw<PrimExpr>(); });
- ffi::Optional<PrimExpr> predicate = op->predicate;
- if (predicate.has_value()) {
- predicate = this->VisitExpr(predicate.value()).as_or_throw<PrimExpr>();
- }
- bool predicate_unchanged =
- !op->predicate.has_value() ||
predicate.value().same_as(op->predicate.value());
- if (indices.same_as(op->indices) && predicate_unchanged) {
+ if (indices.same_as(op->indices)) {
return ffi::GetRef<Expr>(op);
}
- return tirx::BufferLoad(op->buffer, indices, predicate, op->span);
+ return tirx::BufferLoad(op->buffer, indices, op->span);
}
#define RELAX_MUTATE_TIRX_BINOP(OP) \
diff --git a/src/s_tir/analysis/sblock_access_region_detector.cc
b/src/s_tir/analysis/sblock_access_region_detector.cc
index 3feb2c42ea..8256a19247 100644
--- a/src/s_tir/analysis/sblock_access_region_detector.cc
+++ b/src/s_tir/analysis/sblock_access_region_detector.cc
@@ -210,6 +210,34 @@ void BlockReadWriteDetector::VisitStmt_(const BindNode*
op) {
}
void BlockReadWriteDetector::VisitExpr_(const CallNode* op) {
+ auto update_masked_access = [this](const BufferVar& buffer, const
ffi::Array<PrimExpr>& indices,
+ std::vector<BufferVar>* buffers,
+ std::vector<std::vector<arith::IntSet>>*
regions) {
+ std::vector<arith::IntSet> relaxed_region;
+ for (PrimExpr index : indices) {
+ PrimExpr remapped_index = Substitute(index, let_bindings_);
+ while (!remapped_index.same_as(index)) {
+ index = remapped_index;
+ remapped_index = Substitute(index, let_bindings_);
+ }
+
relaxed_region.push_back(arith::EvalSet(arith::IntSet::Vector(remapped_index),
dom_map_));
+ }
+ Update(buffers, regions, buffer, relaxed_region);
+ };
+ if (op->op.same_as(builtin::masked_load()) ||
op->op.same_as(builtin::masked_store())) {
+ bool is_load = op->op.same_as(builtin::masked_load());
+ BufferVar buffer(op->args[0].as_or_throw<Var>());
+ ffi::Array<PrimExpr> indices;
+ for (size_t i = is_load ? 1 : 2; i + 1 < op->args.size(); ++i) {
+ indices.push_back(op->args[i].as_or_throw<PrimExpr>());
+ }
+ update_masked_access(buffer, indices, is_load ? &read_buffers_ :
&writes_buffers_,
+ is_load ? &read_regions_ : &write_regions_);
+ for (size_t i = 1; i < op->args.size(); ++i) {
+ VisitExpr(op->args[i]);
+ }
+ return;
+ }
if (op->op.same_as(builtin::tvm_access_ptr())) {
const VarNode* buffer_var = op->args[1].as<VarNode>();
if (const auto* data = op->args[1].as<CallNode>();
diff --git a/src/s_tir/transform/inject_virtual_thread.cc
b/src/s_tir/transform/inject_virtual_thread.cc
index e6b0396190..b6c8fe7162 100644
--- a/src/s_tir/transform/inject_virtual_thread.cc
+++ b/src/s_tir/transform/inject_virtual_thread.cc
@@ -75,7 +75,18 @@ class ExprTouched final : public StmtExprVisitor {
}
void VisitExpr_(const VarNode* op) final { HandleUseVar(op); }
void VisitExpr_(const CallNode* op) final {
- if (op->op.same_as(builtin::tvm_access_ptr())) {
+ if (op->op.same_as(builtin::masked_load()) ||
op->op.same_as(builtin::masked_store())) {
+ bool is_load = op->op.same_as(builtin::masked_load());
+ const VarNode* buffer = op->args[0].as_or_throw<Var>().get();
+ if (is_load) {
+ HandleUseVar(buffer);
+ } else {
+ HandleWriteVar(buffer);
+ }
+ for (size_t i = 1; i < op->args.size(); ++i) {
+ this->VisitExpr(op->args[i]);
+ }
+ } else if (op->op.same_as(builtin::tvm_access_ptr())) {
const auto* rw_mask = op->args[4].as<IntImmNode>();
auto buffer = GetBufferDataVar(op->args[1]);
if (!buffer.has_value()) {
@@ -240,7 +251,29 @@ class VTInjector : public arith::IRMutatorWithAnalyzer {
}
// Expression.
Expr VisitExpr_(const CallNode* op) final {
- if (op->op.same_as(builtin::buffer_data())) {
+ if (op->op.same_as(builtin::masked_load()) ||
op->op.same_as(builtin::masked_store())) {
+ bool is_load = op->op.same_as(builtin::masked_load());
+ BufferVar buffer(op->args[0].as_or_throw<Var>());
+ PrimExpr value;
+ if (!is_load) value =
this->VisitPrimExpr(op->args[1].as_or_throw<PrimExpr>());
+ ffi::Array<PrimExpr> indices;
+ for (size_t i = is_load ? 1 : 2; i + 1 < op->args.size(); ++i) {
+
indices.push_back(this->VisitPrimExpr(op->args[i].as_or_throw<PrimExpr>()));
+ }
+ PrimExpr predicate =
this->VisitPrimExpr(op->args.back().as_or_throw<PrimExpr>());
+ if (is_load) {
+ BufferLoad access = VisitBufferAccess(BufferLoad(buffer, indices,
op->span));
+ ffi::Array<Expr> args{access->buffer.var()};
+ for (const PrimExpr& index : access->indices) args.push_back(index);
+ args.push_back(predicate);
+ return Call(op->ty, op->op, args, op->attrs, op->ty_args, op->span);
+ }
+ BufferStore access = VisitBufferAccess(BufferStore(buffer, value,
indices, op->span));
+ ffi::Array<Expr> args{access->buffer.var(), access->value};
+ for (const PrimExpr& index : access->indices) args.push_back(index);
+ args.push_back(predicate);
+ return Call(op->ty, op->op, args, op->attrs, op->ty_args, op->span);
+ } else if (op->op.same_as(builtin::buffer_data())) {
auto buffer = GetBufferDataVar(ffi::GetRef<Call>(op)).value();
auto it = alloc_remap_.find(buffer.get());
if (it == alloc_remap_.end()) {
diff --git a/src/s_tir/transform/lower_match_buffer.cc
b/src/s_tir/transform/lower_match_buffer.cc
index 98343099ea..9fbcfc25c3 100644
--- a/src/s_tir/transform/lower_match_buffer.cc
+++ b/src/s_tir/transform/lower_match_buffer.cc
@@ -63,7 +63,7 @@ class MatchBufferLower : public StmtExprMutator {
for (const auto& kv : match_buffers_) {
orig_buffers.push_back(kv.first);
}
- Stmt stmt = StmtExprMutator ::VisitStmt_(op);
+ Stmt stmt = StmtExprMutator::VisitStmt_(op);
// Add remapped buffer keys to match_buffers_
for (const BufferVar& orig_buf : orig_buffers) {
if (auto remap_it = buffer_remap_.find(orig_buf); remap_it !=
buffer_remap_.end()) {
@@ -107,6 +107,14 @@ class MatchBufferLower : public StmtExprMutator {
}
Expr VisitExpr_(const CallNode* op) final {
+ if ((op->op.same_as(builtin::masked_load()) ||
op->op.same_as(builtin::masked_store())) &&
+ !op->args.empty()) {
+ if (auto var = op->args[0].as<Var>(); var &&
var.value()->ty.as<BufferTypeNode>()) {
+ BufferVar buffer(var.value());
+ TVM_FFI_ICHECK(!match_buffers_.count(buffer))
+ << "Predicated buffer access is not currently supported in lower
match buffer pass.";
+ }
+ }
if (op->op.same_as(builtin::buffer_data()) && op->args.size() == 1) {
if (auto var = op->args[0].as<Var>();
var.has_value() && var.value()->ty.as<BufferTypeNode>()) {
@@ -137,8 +145,6 @@ class MatchBufferLower : public StmtExprMutator {
auto n = CopyOnWrite(op);
n->indices = ConvertIndices(MatchBufferRegion(buffer, source),
op->indices);
n->buffer = source->buffer;
- TVM_FFI_ICHECK(!op->predicate.has_value())
- << "Predicated buffer store is not currently supported in lower
match buffer pass.";
return Stmt(n);
}
}
@@ -157,8 +163,6 @@ class MatchBufferLower : public StmtExprMutator {
const BufferVar& buffer = (*it).first;
const BufferRegion& source = (*it).second;
ffi::Array<PrimExpr> indices = ConvertIndices(MatchBufferRegion(buffer,
source), op->indices);
- TVM_FFI_ICHECK(!op->predicate.has_value())
- << "Predicated buffer load is not currently supported in lower match
buffer pass.";
return BufferLoad(source->buffer, indices);
}
}
diff --git a/src/s_tir/transform/manifest_shared_memory_local_stage.cc
b/src/s_tir/transform/manifest_shared_memory_local_stage.cc
index 5dd428dece..4a3be8ee71 100644
--- a/src/s_tir/transform/manifest_shared_memory_local_stage.cc
+++ b/src/s_tir/transform/manifest_shared_memory_local_stage.cc
@@ -72,9 +72,6 @@ class IntermediateStageRewriter {
Stmt local_stage = MakeLocalStage(block, new_buffer, buffer_indices,
relaxed_loops, store);
// Step 3: Create BufferLoad from the intermediate buffer
- TVM_FFI_ICHECK(!store->predicate.has_value())
- << "Predicated buffer store is not currently supported in "
- "manifest shared memory local stage pass.";
BufferLoad new_buffer_load = BufferLoad(new_buffer, buffer_indices);
BufferStore new_buffer_store = block->body.as_or_throw<BufferStore>();
new_buffer_store.CopyOnWrite()->value = new_buffer_load;
diff --git a/src/s_tir/transform/storage_access.cc
b/src/s_tir/transform/storage_access.cc
index 897ccfa57d..2ad37af585 100644
--- a/src/s_tir/transform/storage_access.cc
+++ b/src/s_tir/transform/storage_access.cc
@@ -260,7 +260,29 @@ void StorageAccessVisitor::VisitStmt_(const WhileNode* op)
{
}
void StorageAccessVisitor::VisitExpr_(const CallNode* op) {
- if (op->op.same_as(builtin::address_of())) {
+ Call call = ffi::GetRef<Call>(op);
+ if (op->op.same_as(builtin::masked_load()) ||
op->op.same_as(builtin::masked_store())) {
+ bool is_load = op->op.same_as(builtin::masked_load());
+ BufferVar buffer(op->args[0].as_or_throw<Var>());
+ PrimType value_dtype =
+ is_load ? op->ty.as_or_throw<PrimType>() :
op->args[1].as_or_throw<PrimExpr>().ty();
+ Var buf = ResolveBuffer(buffer.var());
+ StorageScope scope = StorageScope::Create(buffer.scope());
+ if (Enabled(buf.get(), scope)) {
+ TVM_FFI_ICHECK(allow_append_) << call << " " << scope.to_string();
+ AccessEntry e;
+ e.threads = env_threads();
+ e.buffer = buf;
+ e.dtype = value_dtype.WithLanes(1);
+ for (size_t i = is_load ? 1 : 2; i + 1 < op->args.size(); ++i) {
+
e.touched.push_back(arith::IntSet::Vector(op->args[i].as_or_throw<PrimExpr>()));
+ }
+ e.type = is_load ? kRead : kWrite;
+ e.scope = scope;
+ curr_stmt_.access.emplace_back(std::move(e));
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ } else if (op->op.same_as(builtin::address_of())) {
if (const auto* load = op->args[0].as<BufferLoadNode>()) {
// Taking an address does not read the buffer value. Visit only the
// load's children so index expressions still contribute accesses.
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 24da88e24e..74f5bac291 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -1871,7 +1871,7 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const
BufferLoadNode* op) {
// Pass all indices into BufferAccessHelper. In CodeGenLLVM,
// non-flat indices will result in an error in CreateBufferPtr, but
// a subclass may override CreateBufferPtr.
- BufferAccessHelper(op->buffer, op->indices, op->predicate, access_dtype,
make_load);
+ BufferAccessHelper(op->buffer, op->indices, std::nullopt, access_dtype,
make_load);
llvm::Value* ret;
if (loads.size() == 1) {
@@ -1888,8 +1888,73 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const
BufferLoadNode* op) {
return ret;
}
+llvm::Value* CodeGenLLVM::CreateMaskedLoad(const CallNode* op) {
+ TVM_FFI_ICHECK_GE(op->args.size(), 3U);
+ BufferVar buffer(op->args[0].as_or_throw<Var>());
+ ffi::Array<PrimExpr> indices;
+ for (size_t i = 1; i + 1 < op->args.size(); ++i) {
+ indices.push_back(op->args[i].as_or_throw<PrimExpr>());
+ }
+ PrimExpr predicate = op->args.back().as_or_throw<PrimExpr>();
+ PrimType value_dtype = op->ty.as_or_throw<PrimType>();
+ PrimType access_dtype = BufferAccessType(value_dtype);
+ std::vector<llvm::Value*> loads;
+ auto make_load =
+ [this, &loads](TypedPointer buffer_ptr, int /*subelement_i*/,
llvm::Value* predicate,
+ int alignment, bool is_volatile) {
+ TVM_FFI_ICHECK(!is_volatile)
+ << "The masked load intrinsic does not support declaring load as
volatile.";
+ llvm::Instruction* value = builder_->CreateMaskedLoad(buffer_ptr.type,
buffer_ptr.addr,
+
llvm::Align(alignment), predicate);
+ loads.push_back(value);
+ return value;
+ };
+ BufferAccessHelper(buffer, indices, predicate, access_dtype, make_load);
+ llvm::Value* ret =
+ loads.size() == 1 ? loads[0] :
llvm::UndefValue::get(DTypeToLLVMType(access_dtype));
+ for (size_t i = 0; loads.size() > 1 && i < loads.size(); ++i) {
+ ret = builder_->CreateInsertElement(ret, loads[i], ConstInt32(i));
+ }
+ return access_dtype.same_as(value_dtype) ? ret : CreateCast(access_dtype,
value_dtype, ret);
+}
+
+llvm::Value* CodeGenLLVM::CreateMaskedStore(const CallNode* op) {
+ TVM_FFI_ICHECK_GE(op->args.size(), 4U);
+ BufferVar buffer(op->args[0].as_or_throw<Var>());
+ PrimExpr value_expr = op->args[1].as_or_throw<PrimExpr>();
+ ffi::Array<PrimExpr> indices;
+ for (size_t i = 2; i + 1 < op->args.size(); ++i) {
+ indices.push_back(op->args[i].as_or_throw<PrimExpr>());
+ }
+ PrimExpr predicate = op->args.back().as_or_throw<PrimExpr>();
+ PrimType value_dtype = value_expr.ty();
+ llvm::Value* value = MakeValue(value_expr);
+ PrimType access_dtype = BufferAccessType(value_dtype);
+ if (!access_dtype.same_as(value_dtype)) {
+ value = CreateCast(value_dtype, access_dtype, value);
+ value_dtype = access_dtype;
+ }
+ llvm::Instruction* last_store = nullptr;
+ auto make_store =
+ [this, value, &last_store](TypedPointer buffer_ptr, int subelement_i,
llvm::Value* predicate,
+ int alignment, bool is_volatile) {
+ TVM_FFI_ICHECK(!is_volatile)
+ << "The masked store intrinsic does not support declaring store as
volatile.";
+ llvm::Value* to_store =
+ subelement_i == -1 ? value : builder_->CreateExtractElement(value,
subelement_i);
+ last_store = builder_->CreateMaskedStore(to_store, buffer_ptr.addr,
llvm::Align(alignment),
+ predicate);
+ return last_store;
+ };
+ BufferAccessHelper(buffer, indices, predicate, value_dtype, make_store);
+ TVM_FFI_ICHECK(last_store != nullptr);
+ return last_store;
+}
+
llvm::Value* CodeGenLLVM::VisitExpr_(const CallNode* op) {
const ffi::Array<Expr>& args = op->args;
+ if (op->op.same_as(builtin::masked_load())) return CreateMaskedLoad(op);
+ if (op->op.same_as(builtin::masked_store())) return CreateMaskedStore(op);
if (op->op.same_as(builtin::buffer_data())) {
TVM_FFI_ICHECK_EQ(args.size(), 1U);
return MakeValue(args[0]);
@@ -2037,7 +2102,7 @@ void CodeGenLLVM::VisitStmt_(const BufferStoreNode* op) {
// Pass all indices into BufferAccessHelper. In CodeGenLLVM,
// non-flat indices will result in an error in CreateBufferPtr, but
// a subclass may override CreateBufferPtr.
- BufferAccessHelper(op->buffer, op->indices, op->predicate, value_dtype,
make_store);
+ BufferAccessHelper(op->buffer, op->indices, std::nullopt, value_dtype,
make_store);
}
void CodeGenLLVM::VisitStmt_(const ForNode* op) {
diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h
index 8a7e0bd1b1..ea9763448f 100644
--- a/src/target/llvm/codegen_llvm.h
+++ b/src/target/llvm/codegen_llvm.h
@@ -363,6 +363,8 @@ class CodeGenLLVM : public ExprFunctor<llvm::Value*(const
Expr&)>,
std::function<llvm::Instruction*(TypedPointer buffer_ptr, int
subelement_i,
llvm::Value* predicate, int alignment,
bool is_volatile)>
make_instruction);
+ llvm::Value* CreateMaskedLoad(const CallNode* op);
+ llvm::Value* CreateMaskedStore(const CallNode* op);
const VarNode* GetBufferPhysicalRoot(const VarNode* buffer) const;
// Initialize target
virtual void InitTarget();
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index da8a6cb849..2b62d55dc2 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -671,6 +671,10 @@ void CodeGenC::PrintCallExtern(Type ret_type, ffi::String
global_symbol,
}
void CodeGenC::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT(*)
+ TVM_FFI_ICHECK(!op->op.same_as(builtin::masked_load()))
+ << "Predicated buffer load is not supported.";
+ TVM_FFI_ICHECK(!op->op.same_as(builtin::masked_store()))
+ << "Predicated buffer store is not supported.";
if (auto opt_call_op = op->op.as<Op>()) {
auto call_op = opt_call_op.value();
@@ -951,7 +955,6 @@ void CodeGenC::VisitStmt_(const DeclBufferNode* op) {
void CodeGenC::VisitExpr_(const BufferLoadNode* op, std::ostream& os) { //
NOLINT(*)
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "Load from non-flat memory not
supported.";
- TVM_FFI_ICHECK(!op->predicate.has_value()) << "Predicated buffer load is not
supported.";
PrimType value_ty = op->ty.as_or_throw<PrimType>();
PrimExpr index = op->indices[0];
@@ -1026,7 +1029,6 @@ void CodeGenC::VisitExpr_(const BufferLoadNode* op,
std::ostream& os) { // NOLI
void CodeGenC::VisitStmt_(const BufferStoreNode* op) {
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "Store to non-flat memory not
supported.";
- TVM_FFI_ICHECK(!op->predicate.has_value()) << "Predicated buffer store is
not supported.";
PrimType value_ty = op->value.ty();
const PrimType& element_ty = op->buffer->dtype;
diff --git a/src/te/operation/create_primfunc.cc
b/src/te/operation/create_primfunc.cc
index 2332276f10..70ab9e0d9f 100644
--- a/src/te/operation/create_primfunc.cc
+++ b/src/te/operation/create_primfunc.cc
@@ -84,7 +84,7 @@ class TensorLoadToBufferTransformer : public StmtExprMutator {
auto it = tensor2buffers_.find(tensor);
TVM_FFI_ICHECK(it != tensor2buffers_.end()) << "IndexError: Cannot find
the tensor " << tensor;
const BufferVar& buffer = it->second;
- return BufferLoad(buffer, te::GetTensorLoadIndices(call), std::nullopt,
call->span);
+ return BufferLoad(buffer, te::GetTensorLoadIndices(call), call->span);
}
private:
@@ -111,7 +111,7 @@ class BufferSubstituter : public StmtExprMutator {
auto load = StmtExprMutator::VisitExpr_(op).as_or_throw<BufferLoad>();
auto it = buffer_map_.find(load->buffer.get());
if (it != buffer_map_.end()) {
- return BufferLoad(it->second, load->indices, load->predicate,
load->span);
+ return BufferLoad(it->second, load->indices, load->span);
}
return load;
}
@@ -120,7 +120,7 @@ class BufferSubstituter : public StmtExprMutator {
auto store = StmtExprMutator::VisitStmt_(op).as_or_throw<BufferStore>();
auto it = buffer_map_.find(store->buffer.get());
if (it != buffer_map_.end()) {
- return BufferStore(it->second, store->value, store->indices,
store->predicate, store->span);
+ return BufferStore(it->second, store->value, store->indices,
store->span);
}
return store;
}
diff --git a/src/tirx/analysis/deep_equal.cc b/src/tirx/analysis/deep_equal.cc
index e915aad1b8..aecf22e0ab 100644
--- a/src/tirx/analysis/deep_equal.cc
+++ b/src/tirx/analysis/deep_equal.cc
@@ -142,8 +142,7 @@ class ExprDeepEqualChecker : private ExprFunctor<bool(const
Expr&, const PrimExp
const auto* prhs = rhs.as<BufferLoadNode>();
// we run pointer comparison of the buffer
return plhs->ty.as_or_throw<PrimType>() ==
prhs->ty.as_or_throw<PrimType>() &&
- plhs->buffer.same_as(prhs->buffer) && ArrayDeepEqual(plhs->indices,
prhs->indices) &&
- OptionalDeepEqual(plhs->predicate, prhs->predicate);
+ plhs->buffer.same_as(prhs->buffer) && ArrayDeepEqual(plhs->indices,
prhs->indices);
}
bool VisitExpr_(const LetNode* plhs, const PrimExpr& rhs) final {
diff --git a/src/tirx/analysis/verify_memory.cc
b/src/tirx/analysis/verify_memory.cc
index 5694284614..3f79148985 100644
--- a/src/tirx/analysis/verify_memory.cc
+++ b/src/tirx/analysis/verify_memory.cc
@@ -98,6 +98,14 @@ class MemoryAccessVerifier final : protected StmtExprVisitor
{
HandleLoadStoreToVariable(op->buffer.var());
return StmtExprVisitor::VisitStmt_(op);
}
+
+ void VisitExpr_(const CallNode* op) final {
+ if ((op->op.same_as(builtin::masked_load()) ||
op->op.same_as(builtin::masked_store())) &&
+ !op->args.empty()) {
+ HandleLoadStoreToVariable(op->args[0].as_or_throw<Var>());
+ }
+ StmtExprVisitor::VisitExpr_(op);
+ }
//@}
/// Check if the value of a Variable comes from function argument.
diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc
index c0a29ff8ba..14e0358001 100644
--- a/src/tirx/ir/buffer.cc
+++ b/src/tirx/ir/buffer.cc
@@ -375,8 +375,7 @@ BufferVar BufferVar::GetFlattenedBuffer() const {
}
}
-PrimExpr BufferVar::vload(ffi::Array<PrimExpr> begin, PrimType value_dtype,
- ffi::Optional<PrimExpr> predicate) const {
+PrimExpr BufferVar::vload(ffi::Array<PrimExpr> begin, PrimType value_dtype)
const {
const BufferTypeNode* n = operator->();
TVM_FFI_ICHECK(n != nullptr);
PrimType buffer_dtype(n->dtype);
@@ -397,11 +396,10 @@ PrimExpr BufferVar::vload(ffi::Array<PrimExpr> begin,
PrimType value_dtype,
indices.Set(indices.size() - 1, Ramp(base, 1, factor));
}
}
- return BufferLoad(*this, indices, predicate);
+ return BufferLoad(*this, indices);
}
-Stmt BufferVar::vstore(ffi::Array<PrimExpr> begin, PrimExpr value,
- ffi::Optional<PrimExpr> predicate) const {
+Stmt BufferVar::vstore(ffi::Array<PrimExpr> begin, PrimExpr value) const {
const BufferTypeNode* n = operator->();
TVM_FFI_ICHECK(n != nullptr);
PrimType value_dtype = value.ty();
@@ -423,7 +421,7 @@ Stmt BufferVar::vstore(ffi::Array<PrimExpr> begin, PrimExpr
value,
indices.Set(indices.size() - 1, Ramp(base, 1, factor));
}
}
- return BufferStore(*this, value, indices, predicate);
+ return BufferStore(*this, value, indices);
}
ffi::String BufferVar::scope() const { return (*this)->storage_scope; }
@@ -581,10 +579,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def_method("tirx.BufferGetFlattenedBuffer",
&BufferVar::GetFlattenedBuffer)
.def_method("tirx.BufferOffsetOf", &BufferVar::OffsetOf)
.def_method("tirx.BufferOffsetOfp", &BufferVar::OffsetOf_p)
- .def_method(
- "tirx.BufferVLoad",
- static_cast<PrimExpr (BufferVar::*)(ffi::Array<PrimExpr>, PrimType,
- ffi::Optional<PrimExpr>)
const>(&BufferVar::vload))
+ .def_method("tirx.BufferVLoad", &BufferVar::vload)
.def_method("tirx.BufferVStore", &BufferVar::vstore)
.def_method("tirx.BufferStorageScope", &BufferVar::scope)
.def_method("tirx.BufferWithAllocatedAddr",
&BufferVar::with_allocated_addr)
diff --git a/src/tirx/ir/expr.cc b/src/tirx/ir/expr.cc
index 860a13f7eb..1c7f946771 100644
--- a/src/tirx/ir/expr.cc
+++ b/src/tirx/ir/expr.cc
@@ -748,38 +748,15 @@ void BufferLoadNode::LegalizeDType() {
}
}
-BufferLoad::BufferLoad(BufferVar buffer, ffi::Array<PrimExpr> indices,
- ffi::Optional<PrimExpr> predicate, Span span) {
+BufferLoad::BufferLoad(BufferVar buffer, ffi::Array<PrimExpr> indices, Span
span) {
TVM_FFI_ICHECK_EQ(buffer->shape.size(), indices.size())
<< "BufferVar " << buffer.name() << " is " << buffer->shape.size()
<< "-dimensional, cannot be indexed with the " << indices.size()
<< "-dimensional indices provided.";
- if (predicate.has_value()) {
- PrimType predicate_ty = predicate.value().ty();
- bool is_index_scalable = indices.empty() ? false :
indices.back().ty().IsScalableVector();
- bool is_predicate_scalable = predicate_ty.IsScalableVector();
- TVM_FFI_ICHECK_EQ(is_index_scalable, is_predicate_scalable)
- << "Predicate mask dtype and load indices must both be scalable.";
-
- int16_t buffer_encoded_lanes =
static_cast<int16_t>(buffer->dtype->dtype.lanes);
- int buffer_lanes = buffer_encoded_lanes < -1 ? -buffer_encoded_lanes :
buffer_encoded_lanes;
- int index_lanes = indices.empty() ? 1 :
GetLanesOrVScaleFactor(indices.back().ty());
- int predicate_lanes = GetLanesOrVScaleFactor(predicate_ty);
- TVM_FFI_ICHECK_EQ(index_lanes * buffer_lanes, predicate_lanes)
- << "Got a predicate mask with " << predicate_lanes
- << " lanes, but trying to load a vector with " << index_lanes
- << " lanes. The number of lanes must match.";
-
- TVM_FFI_ICHECK(predicate_ty.MatchesCode(DLDataTypeCode::kDLBool) ||
- predicate_ty.MatchesElementType(DLDataTypeCode::kDLUInt, 1))
- << "Predicate mask elements must be boolean values, but got " <<
predicate_ty->dtype << ".";
- }
-
ffi::ObjectPtr<BufferLoadNode> node = ffi::make_object<BufferLoadNode>();
node->buffer = std::move(buffer);
node->indices = std::move(indices);
- node->predicate = std::move(predicate);
node->span = std::move(span);
node->LegalizeDType();
data_ = std::move(node);
@@ -787,10 +764,10 @@ BufferLoad::BufferLoad(BufferVar buffer,
ffi::Array<PrimExpr> indices,
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.BufferLoad", [](BufferVar buffer,
ffi::Array<PrimExpr> indices,
- ffi::Optional<PrimExpr>
predicate, Span span) {
- return BufferLoad(buffer, indices, predicate, span);
- });
+ refl::GlobalDef().def("tirx.BufferLoad",
+ [](BufferVar buffer, ffi::Array<PrimExpr> indices,
Span span) {
+ return BufferLoad(buffer, indices, span);
+ });
}
} // namespace tirx
diff --git a/src/tirx/ir/expr_functor.cc b/src/tirx/ir/expr_functor.cc
index 9a73caf2c8..f739891473 100644
--- a/src/tirx/ir/expr_functor.cc
+++ b/src/tirx/ir/expr_functor.cc
@@ -124,7 +124,7 @@ Expr ExprMutator::VisitExpr_(const BufferLoadNode* op) {
if (indices.same_as(op->indices)) {
return ffi::GetRef<PrimExpr>(op);
} else {
- return BufferLoad(op->buffer, indices, op->predicate);
+ return BufferLoad(op->buffer, indices);
}
}
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index a35791c60a..84d70b8fc1 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -422,7 +422,7 @@ TVM_FFI_INLINE int GetLanesOrVScaleFactor(const PrimType&
ty) {
}
BufferStore::BufferStore(BufferVar buffer, PrimExpr value,
ffi::Array<PrimExpr> indices,
- ffi::Optional<PrimExpr> predicate, Span span) {
+ Span span) {
TVM_FFI_ICHECK_EQ(buffer->shape.size(), indices.size())
<< "BufferVar " << buffer.name() << " is " << buffer->shape.size()
<< "-dimensional, cannot be indexed with the " << indices.size()
@@ -442,12 +442,6 @@ BufferStore::BufferStore(BufferVar buffer, PrimExpr value,
ffi::Array<PrimExpr>
TVM_FFI_ICHECK(!(is_index_scalable && is_buffer_dtype_scalable))
<< "Index dtype and buffer dtype can't both be scalable.";
- if (predicate.has_value()) {
- bool is_predicate_dtype_scalable =
predicate.value().ty().IsScalableVector();
- TVM_FFI_ICHECK_EQ(is_value_dtype_scalable, is_predicate_dtype_scalable)
- << "Predicate mask dtype and value dtype must both be scalable.";
- }
-
if (is_index_scalable || is_buffer_dtype_scalable) {
TVM_FFI_ICHECK(is_value_dtype_scalable) << "Can't store non-scalable data
into scalable buffer";
}
@@ -461,21 +455,6 @@ BufferStore::BufferStore(BufferVar buffer, PrimExpr value,
ffi::Array<PrimExpr>
<< index_lanes * buffer_lanes << " (" << index_lanes << " index lanes *
" << buffer_lanes
<< " buffer element lanes)";
- if (predicate.has_value()) {
- PrimType predicate_ty = predicate.value().ty();
- int predicate_dtype_lanes = GetLanesOrVScaleFactor(predicate_ty);
- TVM_FFI_ICHECK_EQ(value_dtype_lanes, predicate_dtype_lanes)
- << "Got a predicate mask with " << predicate_dtype_lanes
- << " lanes, but trying to store a value with " << value_dtype_lanes
- << " lanes. The number of lanes must match.";
-
- PrimType predicate_element_ty = predicate_ty.WithLanes(1);
- TVM_FFI_ICHECK(predicate_element_ty.MatchesCode(DLDataTypeCode::kDLBool) ||
-
predicate_element_ty.MatchesElementType(DLDataTypeCode::kDLUInt, 1))
- << "Predicate mask elements must be boolean values, but got "
- << ffi::DLDataTypeToString(predicate_element_ty->dtype) << ".";
- }
-
PrimType buffer_dtype = PrimType::Void();
if (is_index_scalable || is_buffer_dtype_scalable) {
buffer_dtype = PrimType::ScalableVector(buffer->dtype.code(),
buffer->dtype.bits(),
@@ -495,7 +474,6 @@ BufferStore::BufferStore(BufferVar buffer, PrimExpr value,
ffi::Array<PrimExpr>
node->buffer = std::move(buffer);
node->value = std::move(value);
node->indices = std::move(indices);
- node->predicate = std::move(predicate);
node->span = std::move(span);
data_ = std::move(node);
}
@@ -504,9 +482,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def("tirx.BufferStore",
[](BufferVar buffer, PrimExpr value,
ffi::Array<PrimExpr> indices,
- ffi::Optional<PrimExpr> predicate, Span span) {
- return BufferStore(buffer, value, indices,
predicate, span);
- });
+ Span span) { return BufferStore(buffer, value,
indices, span); });
}
// BufferRegion
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index 737d857e09..df9fe03382 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -379,6 +379,15 @@ TIR_DEFINE_BUILTIN_FUNC(get_active_lane_mask)
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
+TIR_DEFINE_BUILTIN_FUNC(masked_load)
+ .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kReadState))
+ .set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
+
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
+
+TIR_DEFINE_BUILTIN_FUNC(masked_store)
+ .set_attr<TCallEffectKind>("TCallEffectKind",
+
static_cast<int64_t>(CallEffectKind::kUpdateState));
+
TIR_DEFINE_BUILTIN_FUNC(ignore_loop_partition)
.set_num_inputs(1)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure))
diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc
index f4a6a81f84..a94b97fb79 100644
--- a/src/tirx/script/builder/ir.cc
+++ b/src/tirx/script/builder/ir.cc
@@ -784,8 +784,7 @@ Var EnvThread(ffi::String thread_tag, PrimType dtype) {
return var;
}
-void BufferStore(BufferVar buffer, PrimExpr value, ffi::Array<PrimExpr>
indices,
- ffi::Optional<PrimExpr> predicate = std::nullopt) {
+void BufferStore(BufferVar buffer, PrimExpr value, ffi::Array<PrimExpr>
indices) {
PrimType buffer_dtype = buffer->dtype;
PrimType index_ty = indices.empty() ? PrimType::Int(32) :
indices.back().ty();
bool is_index_scalable = !indices.empty() && index_ty.IsScalableVector();
@@ -833,7 +832,7 @@ void BufferStore(BufferVar buffer, PrimExpr value,
ffi::Array<PrimExpr> indices,
}
value = tvm::cast(lhs_dtype, value);
}
- tvm::tirx::BufferStore store(buffer, value, indices, predicate);
+ tvm::tirx::Stmt store = tvm::tirx::BufferStore(buffer, value, indices);
if (lhs_dtype != rhs_dtype) {
if (lhs_dtype.code() != rhs_dtype.code()) {
if ((lhs_dtype.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt)) &&
diff --git a/src/tirx/script/printer/buffer.cc
b/src/tirx/script/printer/buffer.cc
index c13a06aa0d..34750447d3 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -413,8 +413,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
ExprDoc value = d->AsDoc<ExprDoc>(store->value, p->Attr("value"));
// special case for scalar buffers
- if ((store->buffer.IsScalar(true) || store->buffer.IsScalar(false))
&&
- !store->predicate.has_value()) {
+ if (store->buffer.IsScalar(true) || store->buffer.IsScalar(false)) {
// TVM_FFI_ICHECK(store->indices.size() == 1 &&
tirx::is_zero(store->indices[0]))
// << "1-dim buffer with shape (1,) store with indices other
than [0] is not "
// "supported";
@@ -424,14 +423,6 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
return AssignDoc(doc.value(), value, std::nullopt);
}
- // Use .vstore(...) syntax when there is a predicate
- if (store->predicate.has_value()) {
- ExprDoc indices = d->AsDoc<ExprDoc>(store->indices,
p->Attr("indices"));
- ExprDoc predicate = d->AsDoc<ExprDoc>(store->predicate,
p->Attr("predicate"));
- return ExprStmtDoc(
- buffer->Attr("vstore")->Call({indices, value}, {"predicate"},
{predicate}));
- }
-
return AssignDoc(
/*lhs=*/buffer[BufferIndices(store->indices, p->Attr("indices"),
d)],
/*rhs=*/value, std::nullopt);
@@ -443,8 +434,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
ExprDoc buffer = d->AsDoc<ExprDoc>(load->buffer, p->Attr("buffer"));
// special case for scalar
- if ((load->buffer.IsScalar(true) || load->buffer.IsScalar(false)) &&
- !load->predicate.has_value()) {
+ if (load->buffer.IsScalar(true) || load->buffer.IsScalar(false)) {
// TVM_FFI_ICHECK(load->indices.size() == 1 &&
tirx::is_zero(load->indices[0]))
// << "Scalar buffer load with indices other than [0] is not
supported";
ffi::Optional<ExprDoc> doc = d->GetVarDoc(load->buffer);
@@ -453,13 +443,6 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
return doc.value();
}
- // Use .vload(...) syntax when there is a predicate
- if (load->predicate.has_value()) {
- ExprDoc indices = d->AsDoc<ExprDoc>(load->indices,
p->Attr("indices"));
- ExprDoc predicate = d->AsDoc<ExprDoc>(load->predicate,
p->Attr("predicate"));
- return buffer->Attr("vload")->Call({indices}, {"predicate"},
{predicate});
- }
-
return buffer[BufferIndices(load->indices, p->Attr("indices"), d)];
});
diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h
index b2246a667f..e98e09c1ff 100644
--- a/src/tirx/script/printer/utils.h
+++ b/src/tirx/script/printer/utils.h
@@ -123,6 +123,12 @@ inline void AsDocBody(const tirx::Stmt& stmt, AccessPath
p, TIRFrameNode* f, con
if (load->buffer.same_as(buffer)) {
found = true;
}
+ } else if (const auto* call = node.as<CallNode>()) {
+ if (call->op.same_as(tirx::builtin::masked_load()) &&
!call->args.empty()) {
+ if (auto var = call->args[0].as<Var>(); var &&
var.value().same_as(buffer.var())) {
+ found = true;
+ }
+ }
}
});
return found;
@@ -137,8 +143,7 @@ inline void AsDocBody(const tirx::Stmt& stmt, AccessPath p,
TIRFrameNode* f, con
if (d->cfg->syntax_sugar && alloc != nullptr &&
alloc->buffer.IsScalar(true) && i + 1 < n) {
const auto* store = body[i + 1].as<tirx::BufferStoreNode>();
bool can_merge_init = store != nullptr &&
store->buffer.same_as(alloc->buffer) &&
- !store->predicate.has_value() &&
store->indices.size() == 1 &&
- tirx::is_zero(store->indices[0]) &&
+ store->indices.size() == 1 &&
tirx::is_zero(store->indices[0]) &&
!value_refs_buffer(store->value, alloc->buffer);
if (can_merge_init) {
Doc alloc_doc = d->AsDoc(body[i], item_p);
diff --git a/src/tirx/transform/flatten_buffer.cc
b/src/tirx/transform/flatten_buffer.cc
index d901d91434..978ecec1f9 100644
--- a/src/tirx/transform/flatten_buffer.cc
+++ b/src/tirx/transform/flatten_buffer.cc
@@ -248,6 +248,21 @@ class BufferFlattener : public
arith::IRMutatorWithAnalyzer {
}
Expr VisitExpr_(const CallNode* op) final {
+ if (op->op.same_as(builtin::masked_load()) ||
op->op.same_as(builtin::masked_store())) {
+ bool is_load = op->op.same_as(builtin::masked_load());
+ BufferVar original(op->args[0].as_or_throw<Var>());
+ ffi::Array<PrimExpr> indices;
+ for (size_t i = is_load ? 1 : 2; i + 1 < op->args.size(); ++i) {
+
indices.push_back(this->VisitPrimExpr(op->args[i].as_or_throw<PrimExpr>()));
+ }
+ buffers_used_.insert(original);
+ const FlatInfo& info = Lookup(original);
+ ffi::Array<Expr> args{info.flattened.var()};
+ if (!is_load) args.push_back(this->VisitExpr(op->args[1]));
+ for (const PrimExpr& index : FoldIndices(info, indices))
args.push_back(index);
+ args.push_back(this->VisitExpr(op->args.back()));
+ return Call(op->ty, op->op, args, op->attrs, op->ty_args, op->span);
+ }
if (op->op.same_as(builtin::buffer_data()) && op->args.size() == 1) {
if (auto var = op->args[0].as<Var>()) {
if (var.value()->ty.as<BufferTypeNode>()) {
diff --git a/src/tirx/transform/remove_no_op.cc
b/src/tirx/transform/remove_no_op.cc
index 9311a4e1ea..57dec79eae 100644
--- a/src/tirx/transform/remove_no_op.cc
+++ b/src/tirx/transform/remove_no_op.cc
@@ -193,8 +193,7 @@ class NoOpRemover : public arith::IRMutatorWithAnalyzer {
// A write whose destination is known to already contain the
// values to be written is a no-op.
- PrimExpr stores_existing_value =
- store->value - BufferLoad(store->buffer, store->indices,
store->predicate) == 0;
+ PrimExpr stores_existing_value = store->value - BufferLoad(store->buffer,
store->indices) == 0;
stores_existing_value = analyzer_->Simplify(stores_existing_value);
if (is_one(stores_existing_value)) {
return only_side_effects();
diff --git a/src/tirx/transform/storage_rewrite.cc
b/src/tirx/transform/storage_rewrite.cc
index 00116b968f..5d5c0bfe9d 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -145,20 +145,7 @@ class LinearAccessPatternFinder final : public
StmtExprVisitor {
scope_.push_back(StmtEntry());
// visit subexpr
StmtExprVisitor::VisitStmt_(op);
- // Add write access.
- const VarNode* buffer_var =
- buffer_aliases_.Get(op->buffer.var()).value_or(op->buffer.var()).get();
- auto it = alloc_info_.find(buffer_var);
- if (it != alloc_info_.end() && it->second.alloc) {
- TVM_FFI_ICHECK_LT(it->second.level, scope_.size());
- scope_[it->second.level].touched.push_back(buffer_var);
-
- TVM_FFI_ICHECK_EQ(1, it->second.num_physical_dimensions)
- << "BufferVar " << op->buffer.name() << " is allocated with "
- << it->second.num_physical_dimensions
- << " physical dimensions, but is accessed as having "
- << "1 physical dimension" << std::endl;
- }
+ RecordAccess(op->buffer);
StmtEntry e = scope_.back();
scope_.pop_back();
if (e.touched.size() != 0) {
@@ -168,23 +155,8 @@ class LinearAccessPatternFinder final : public
StmtExprVisitor {
}
void VisitExpr_(const BufferLoadNode* op) final {
- // Add write access.
StmtExprVisitor::VisitExpr_(op);
-
- const VarNode* buffer_var =
- buffer_aliases_.Get(op->buffer.var()).value_or(op->buffer.var()).get();
- auto it = alloc_info_.find(buffer_var);
- if (it != alloc_info_.end() && it->second.alloc) {
- TVM_FFI_ICHECK_LT(it->second.level, scope_.size())
- << "Load memory in places other than store.";
- scope_[it->second.level].touched.push_back(buffer_var);
-
- TVM_FFI_ICHECK_EQ(1, it->second.num_physical_dimensions)
- << "BufferVar " << op->buffer.name() << " is allocated with "
- << it->second.num_physical_dimensions
- << " physical dimensions, but is accessed as having "
- << "1 physical dimension" << std::endl;
- }
+ RecordAccess(op->buffer);
}
void VisitStmt_(const EvaluateNode* op) final {
@@ -279,6 +251,19 @@ class LinearAccessPatternFinder final : public
StmtExprVisitor {
}
}
+ void RecordAccess(const BufferVar& buffer) {
+ const VarNode* buffer_var =
buffer_aliases_.Get(buffer.var()).value_or(buffer.var()).get();
+ auto it = alloc_info_.find(buffer_var);
+ if (it == alloc_info_.end() || !it->second.alloc) return;
+ TVM_FFI_ICHECK_LT(it->second.level, scope_.size())
+ << "Buffer access occurs outside a statement scope.";
+ scope_[it->second.level].touched.push_back(buffer_var);
+ TVM_FFI_ICHECK_EQ(1, it->second.num_physical_dimensions)
+ << "BufferVar " << buffer.name() << " is allocated with "
+ << it->second.num_physical_dimensions << " physical dimensions, but is
accessed as having "
+ << "1 physical dimension" << std::endl;
+ }
+
// linearized access sequence.
std::vector<StmtEntry> linear_seq_;
// The storage scope of each buffer
@@ -549,7 +534,31 @@ class StoragePlanRewriter : public StmtExprMutator {
}
}
Expr VisitExpr_(const CallNode* op) final {
- if (op->op.same_as(builtin::tvm_access_ptr())) {
+ if (op->op.same_as(builtin::masked_load()) ||
op->op.same_as(builtin::masked_store())) {
+ bool is_load = op->op.same_as(builtin::masked_load());
+ BufferVar buffer(op->args[0].as_or_throw<Var>());
+ PrimExpr value;
+ if (!is_load) value =
this->VisitPrimExpr(op->args[1].as_or_throw<PrimExpr>());
+ ffi::Array<PrimExpr> indices;
+ for (size_t i = is_load ? 1 : 2; i + 1 < op->args.size(); ++i) {
+
indices.push_back(this->VisitPrimExpr(op->args[i].as_or_throw<PrimExpr>()));
+ }
+ if (is_load) {
+ BufferLoad access(buffer, indices, op->span);
+ access = VisitBufferAccess(std::move(access));
+ ffi::Array<Expr> args{access->buffer.var()};
+ for (const PrimExpr& index : access->indices) args.push_back(index);
+ args.push_back(this->VisitExpr(op->args.back()));
+ return Call(access->ty, op->op, args, op->attrs, op->ty_args,
op->span);
+ } else {
+ BufferStore access(buffer, value, indices, op->span);
+ access = VisitBufferAccess(std::move(access));
+ ffi::Array<Expr> args{access->buffer.var(), access->value};
+ for (const PrimExpr& index : access->indices) args.push_back(index);
+ args.push_back(this->VisitExpr(op->args.back()));
+ return Call(PrimType::Void(), op->op, args, op->attrs, op->ty_args,
op->span);
+ }
+ } else if (op->op.same_as(builtin::tvm_access_ptr())) {
TVM_FFI_ICHECK_EQ(op->args.size(), 5U);
PrimExpr dtype_marker = op->args[0].as_or_throw<PrimExpr>();
PrimType dtype = dtype_marker.ty();
@@ -1311,7 +1320,17 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
}
void VisitExpr_(const CallNode* op) final {
- if (op->op.same_as(builtin::tvm_access_ptr())) {
+ if (op->op.same_as(builtin::masked_load()) ||
op->op.same_as(builtin::masked_store())) {
+ bool is_load = op->op.same_as(builtin::masked_load());
+ BufferVar buffer(op->args[0].as_or_throw<Var>());
+ PrimType dtype =
+ is_load ? op->ty.as_or_throw<PrimType>() :
op->args[1].as_or_throw<PrimExpr>().ty();
+ ffi::Array<PrimExpr> indices;
+ for (size_t i = is_load ? 1 : 2; i + 1 < op->args.size(); ++i) {
+ indices.push_back(op->args[i].as_or_throw<PrimExpr>());
+ }
+ OnArrayAccess(dtype, buffer.get(), indices, is_load);
+ } else if (op->op.same_as(builtin::tvm_access_ptr())) {
PrimType dtype = op->args[0].as_or_throw<PrimExpr>().ty();
auto buffer_var = GetBufferDataVar(op->args[1]);
PrimExpr index = op->args[2].as_or_throw<PrimExpr>();
@@ -1699,6 +1718,42 @@ class VectorTypeRewriter : public StmtExprMutator {
return modified;
}
+ ffi::Optional<Expr> RewriteMaskedCall(const CallNode* op) {
+ if (op->op.same_as(builtin::masked_load())) {
+ BufferVar buffer(op->args[0].as_or_throw<Var>());
+ ffi::Array<PrimExpr> indices;
+ for (size_t i = 1; i + 1 < op->args.size(); ++i) {
+
indices.push_back(this->VisitPrimExpr(op->args[i].as_or_throw<PrimExpr>()));
+ }
+ BufferLoad access(buffer, indices, op->span);
+ auto [modified, shuffle_index] = VisitBufferAccess(access);
+ TVM_FFI_ICHECK_LT(shuffle_index, 0)
+ << "A masked vector load cannot be rewritten into a scalar shuffle.";
+ if (!modified.same_as(access)) modified.CopyOnWrite()->LegalizeDType();
+ ffi::Array<Expr> args{modified->buffer.var()};
+ for (const PrimExpr& index : modified->indices) args.push_back(index);
+ args.push_back(this->VisitExpr(op->args.back()));
+ return Call(modified->ty, op->op, args, op->attrs, op->ty_args,
op->span);
+ }
+ if (op->op.same_as(builtin::masked_store())) {
+ BufferVar buffer(op->args[0].as_or_throw<Var>());
+ PrimExpr value =
this->VisitPrimExpr(op->args[1].as_or_throw<PrimExpr>());
+ ffi::Array<PrimExpr> indices;
+ for (size_t i = 2; i + 1 < op->args.size(); ++i) {
+
indices.push_back(this->VisitPrimExpr(op->args[i].as_or_throw<PrimExpr>()));
+ }
+ BufferStore access(buffer, value, indices, op->span);
+ auto [modified, shuffle_index] = VisitBufferAccess(std::move(access));
+ TVM_FFI_ICHECK_LT(shuffle_index, 0)
+ << "A masked vector store cannot be rewritten into a scalar
shuffle.";
+ ffi::Array<Expr> args{modified->buffer.var(), modified->value};
+ for (const PrimExpr& index : modified->indices) args.push_back(index);
+ args.push_back(this->VisitExpr(op->args.back()));
+ return Call(PrimType::Void(), op->op, args, op->attrs, op->ty_args,
op->span);
+ }
+ return std::nullopt;
+ }
+
Stmt VisitStmt_(const BindNode* op) final {
auto it = rewrite_map_.find(op->var.get());
Expr value = this->VisitExpr(op->value);
@@ -1764,6 +1819,9 @@ class VectorTypeRewriter : public StmtExprMutator {
}
Expr VisitExpr_(const CallNode* op) final {
+ if (auto rewritten = RewriteMaskedCall(op)) {
+ return rewritten.value();
+ }
if (op->op.same_as(builtin::buffer_data()) && op->args.size() == 1) {
if (auto var = op->args[0].as<Var>();
var.has_value() && var.value()->ty.as<BufferTypeNode>()) {
diff --git a/src/tirx/transform/unsupported_dtype_legalize.cc
b/src/tirx/transform/unsupported_dtype_legalize.cc
index c9ba1736b0..815db4a864 100644
--- a/src/tirx/transform/unsupported_dtype_legalize.cc
+++ b/src/tirx/transform/unsupported_dtype_legalize.cc
@@ -140,7 +140,9 @@ class ComputeLegalizePlanner : public StmtExprVisitor {
void VisitExpr_(const VarNode* op) final {
StmtExprVisitor::VisitExpr_(op);
Var buffer_var = ffi::GetRef<Var>(op);
- if (buffer_var->ty.as<PointerTypeNode>()) {
+ if (buffer_var->ty.as<BufferTypeNode>()) {
+ this->PopulateBufferRemap(BufferVar(buffer_var));
+ } else if (buffer_var->ty.as<PointerTypeNode>()) {
opaque_var_access_.insert(buffer_var);
}
}
@@ -265,6 +267,39 @@ class ComputeLegalizer : public StmtExprMutator {
}
Expr VisitExpr_(const CallNode* op) final {
+ if (op->op.same_as(builtin::masked_load()) ||
op->op.same_as(builtin::masked_store())) {
+ bool is_load = op->op.same_as(builtin::masked_load());
+ BufferVar original(op->args[0].as_or_throw<Var>());
+ BufferVar buffer = GetRemappedBuffer(original);
+ ffi::Array<Expr> args{buffer.var()};
+ PrimExpr value;
+ if (!is_load) {
+ value = this->VisitPrimExpr(op->args[1].as_or_throw<PrimExpr>());
+ }
+ ffi::Array<PrimExpr> indices;
+ for (size_t i = is_load ? 1 : 2; i + 1 < op->args.size(); ++i) {
+
indices.push_back(this->VisitPrimExpr(op->args[i].as_or_throw<PrimExpr>()));
+ }
+ PrimExpr predicate =
this->VisitPrimExpr(op->args.back().as_or_throw<PrimExpr>());
+ if (is_load) {
+ for (const PrimExpr& index : indices) args.push_back(index);
+ args.push_back(predicate);
+ Type type = BufferLoad(buffer, indices).ty();
+ return Call(type, op->op, args, op->attrs, op->ty_args, op->span);
+ }
+ if (MatchType(buffer->dtype)) {
+ value = CastTargetToDType(value, BufferLoad(buffer, indices).ty());
+ }
+ PrimType storage_dtype = BufferLoad(buffer, indices).ty();
+ if (value.ty() != storage_dtype) {
+ TVM_FFI_ICHECK(MatchType(value.ty()));
+ value = DTypeConversion(value, storage_dtype);
+ }
+ args.push_back(value);
+ for (const PrimExpr& index : indices) args.push_back(index);
+ args.push_back(predicate);
+ return Call(PrimType::Void(), op->op, args, op->attrs, op->ty_args,
op->span);
+ }
if (!op->ty.as<PrimTypeNode>()) {
return StmtExprMutator::VisitExpr_(op);
}
@@ -365,30 +400,22 @@ class ComputeLegalizer : public StmtExprMutator {
auto fmutate = [this](const PrimExpr& e) { return this->VisitPrimExpr(e);
};
ffi::Array<PrimExpr> indices = op->indices.Map(fmutate);
- ffi::Optional<PrimExpr> predicate = std::nullopt;
- if (op->predicate.has_value()) {
- predicate = this->VisitPrimExpr(op->predicate.value());
- }
-
BufferVar new_buf = GetRemappedBuffer(op->buffer);
- if (value.same_as(op->value) && indices.same_as(op->indices) &&
- predicate.same_as(op->predicate) && new_buf.same_as(op->buffer)) {
+ if (value.same_as(op->value) && indices.same_as(op->indices) &&
new_buf.same_as(op->buffer)) {
return ffi::GetRef<Stmt>(op);
} else {
if (MatchType(new_buf->dtype)) {
- int index_lanes = indices.size() ? indices.back().ty().lanes() : 1;
- int buffer_lanes = new_buf->dtype.lanes();
- PrimType legalized_dtype = new_buf->dtype.WithLanes(index_lanes *
buffer_lanes);
- value = CastTargetToDType(value, legalized_dtype);
+ value = CastTargetToDType(value, BufferLoad(new_buf, indices).ty());
}
- if (value.ty() != new_buf->dtype) {
+ PrimType storage_dtype = BufferLoad(new_buf, indices).ty();
+ if (value.ty() != storage_dtype) {
// this happens when buffer get rewritten to f32
// but values remain as fp8/bf16
TVM_FFI_ICHECK(MatchType(value.ty()));
- value = DTypeConversion(value,
new_buf->dtype.WithLanes(value.ty().lanes()));
+ value = DTypeConversion(value, storage_dtype);
}
- return BufferStore(new_buf, value, indices, predicate);
+ return BufferStore(new_buf, value, indices);
}
}
@@ -477,7 +504,7 @@ class ComputeLegalizer : public StmtExprMutator {
if (new_buf.same_as(op->buffer)) {
return ret;
} else {
- return BufferLoad(new_buf, op->indices, op->predicate);
+ return BufferLoad(new_buf, op->indices);
}
}
@@ -648,18 +675,13 @@ class StorageLegalizer : public StmtExprMutator {
PrimExpr value = this->ChangeToUInt(VisitPrimExpr(op->value));
BufferVar new_buf = GetRemappedBuffer(op->buffer);
auto indices = op->indices.Map([this](PrimExpr expr) { return
this->VisitPrimExpr(expr); });
- ffi::Optional<PrimExpr> predicate = std::nullopt;
- if (op->predicate.has_value()) {
- predicate = this->VisitPrimExpr(op->predicate.value());
- }
- if (new_buf.same_as(op->buffer) && indices.same_as(op->indices) &&
- predicate.same_as(op->predicate) && value.same_as(op->value)) {
+ if (new_buf.same_as(op->buffer) && indices.same_as(op->indices) &&
value.same_as(op->value)) {
return ffi::GetRef<Stmt>(op);
} else {
if (MatchType(op->value.ty())) {
TVM_FFI_ICHECK(new_buf->dtype.MatchesCode(DLDataTypeCode::kDLUInt));
}
- return BufferStore(new_buf, value, indices, predicate);
+ return BufferStore(new_buf, value, indices);
}
}
@@ -688,11 +710,38 @@ class StorageLegalizer : public StmtExprMutator {
if (new_buf.same_as(op->buffer)) {
return ret;
} else {
- return BufferLoad(new_buf, op->indices, op->predicate);
+ return BufferLoad(new_buf, op->indices);
}
}
Expr VisitExpr_(const CallNode* op) final {
+ if (op->op.same_as(builtin::masked_load()) ||
op->op.same_as(builtin::masked_store())) {
+ bool is_load = op->op.same_as(builtin::masked_load());
+ BufferVar buffer =
GetRemappedBuffer(BufferVar(op->args[0].as_or_throw<Var>()));
+ ffi::Array<Expr> args{buffer.var()};
+ PrimExpr value;
+ if (!is_load) {
+ PrimExpr original_value = op->args[1].as_or_throw<PrimExpr>();
+ value = this->ChangeToUInt(this->VisitPrimExpr(original_value));
+ if (MatchType(original_value.ty())) {
+ TVM_FFI_ICHECK(buffer->dtype.MatchesCode(DLDataTypeCode::kDLUInt));
+ }
+ args.push_back(value);
+ }
+ ffi::Array<PrimExpr> indices;
+ for (size_t i = is_load ? 1 : 2; i + 1 < op->args.size(); ++i) {
+ PrimExpr index =
this->VisitPrimExpr(op->args[i].as_or_throw<PrimExpr>());
+ indices.push_back(index);
+ args.push_back(index);
+ }
+ args.push_back(this->VisitExpr(op->args.back()));
+ if (is_load) {
+ Type type = BufferLoad(buffer, indices).ty();
+ return Call(type, op->op, args, op->attrs, op->ty_args, op->span);
+ } else {
+ return Call(PrimType::Void(), op->op, args, op->attrs, op->ty_args,
op->span);
+ }
+ }
if (const auto* pointer_type = op->ty.as<PointerTypeNode>()) {
Expr ret = StmtExprMutator::VisitExpr_(op);
const auto* element_type = pointer_type->element_type.as<PrimTypeNode>();
diff --git a/src/tirx/transform/vectorize_loop.cc
b/src/tirx/transform/vectorize_loop.cc
index 0cc74d5405..715ac307e2 100644
--- a/src/tirx/transform/vectorize_loop.cc
+++ b/src/tirx/transform/vectorize_loop.cc
@@ -158,8 +158,10 @@ bool EnableBufferLevelPredication(Target target) {
* After:
* for i_0 in T.serial(4):
* predicate = T.get_active_lane_mask("uint1x4", i_0 * 4, 14)
- * A_load = T.meta_var(A.vload([T.Ramp(i_0 * 4, 1, 4)], predicate=predicate))
- * B.vstore([T.Ramp(i_0 * 4, 1, 4)], A_load, predicate=predicate)
+ * A_load = T.meta_var(T.call_intrin("float32x4", "tirx.masked_load", A,
+ * T.Ramp(i_0 * 4, 1, 4), predicate))
+ * T.evaluate(T.call_intrin("void", "tirx.masked_store", B, A_load,
+ * T.Ramp(i_0 * 4, 1, 4), predicate))
*/
class TryPredicateBufferAccesses : public StmtExprMutator {
public:
@@ -214,20 +216,39 @@ class TryPredicateBufferAccesses : public StmtExprMutator
{
return TryPredicateBufferAccess(store);
}
- template <typename AccessNode>
- AccessNode TryPredicateBufferAccess(AccessNode node) {
+ Expr VisitExpr_(const CallNode* op) final {
+ Call call = StmtExprMutator::VisitExpr_(op).as_or_throw<Call>();
+ if (!call->op.same_as(builtin::masked_load()) &&
!call->op.same_as(builtin::masked_store())) {
+ return call;
+ }
+
+ bool is_load = call->op.same_as(builtin::masked_load());
+ ffi::Array<PrimExpr> indices;
+ for (size_t i = is_load ? 1 : 2; i + 1 < call->args.size(); ++i) {
+ indices.push_back(call->args[i].as_or_throw<PrimExpr>());
+ }
+ if (auto lane_mask = GetLaneMask(indices)) {
+ PrimExpr predicate = call->args.back().as_or_throw<PrimExpr>();
+ predicate = allow_offset_predication_ ? predicate & lane_mask.value() :
lane_mask.value();
+ ffi::Array<Expr> args = call->args;
+ args.Set(args.size() - 1, predicate);
+ return Call(call->ty, call->op, args, call->attrs, call->ty_args,
call->span);
+ }
+ return call;
+ }
+
+ ffi::Optional<PrimExpr> GetLaneMask(const ffi::Array<PrimExpr>& indices) {
num_accesses_analyzed_ += 1;
// Do not try to predicate non-vectorized accesses
- ffi::Array<PrimExpr> indices = node->indices;
if (!indices.size() || !indices[0]->IsInstance<RampNode>()) {
- return node;
+ return std::nullopt;
}
- Ramp ramp = node->indices[0].template as_or_throw<Ramp>();
+ Ramp ramp = indices[0].as_or_throw<Ramp>();
if (!ffi::StructuralEqual()(ramp->stride, stride_) ||
!ffi::StructuralEqual()(ramp->lanes, lanes_)) {
- return node;
+ return std::nullopt;
}
bool same_base = ffi::StructuralEqual()(ramp->base, base_);
@@ -236,7 +257,7 @@ class TryPredicateBufferAccesses : public StmtExprMutator {
// memory base. This covers accesses such as A[offset + i] guarded by
// a predicate over i.
if (!allow_offset_predication_) {
- return node;
+ return std::nullopt;
}
}
@@ -249,15 +270,28 @@ class TryPredicateBufferAccesses : public StmtExprMutator
{
.as_or_throw<PrimExpr>();
num_accesses_rewritten_ += 1;
- auto writer = node.CopyOnWrite();
- if (node->predicate.has_value() && allow_offset_predication_) {
- // BufferVar predicates are uint1 lane masks, so mask merging uses
bitwise
- // and rather than logical &&.
- writer->predicate = node->predicate.value() & lane_mask;
- } else {
- writer->predicate = lane_mask;
+ return lane_mask;
+ }
+
+ Expr TryPredicateBufferAccess(BufferLoad load) {
+ if (auto mask = GetLaneMask(load->indices)) {
+ ffi::Array<Expr> args{load->buffer.var()};
+ for (const PrimExpr& index : load->indices) args.push_back(index);
+ args.push_back(mask.value());
+ return Call(load->ty, builtin::masked_load(), args, {}, {}, load->span);
}
- return node;
+ return load;
+ }
+
+ Stmt TryPredicateBufferAccess(BufferStore store) {
+ if (auto mask = GetLaneMask(store->indices)) {
+ ffi::Array<Expr> args{store->buffer.var(), store->value};
+ for (const PrimExpr& index : store->indices) args.push_back(index);
+ args.push_back(mask.value());
+ return Evaluate(Call(PrimType::Void(), builtin::masked_store(), args,
{}, {}, store->span),
+ store->span);
+ }
+ return store;
}
/*! \brief The variable base expr of the predicate. */
diff --git a/tests/python/codegen/test_target_codegen.py
b/tests/python/codegen/test_target_codegen.py
index b7fa32575c..16e9cf1fcc 100644
--- a/tests/python/codegen/test_target_codegen.py
+++ b/tests/python/codegen/test_target_codegen.py
@@ -30,7 +30,16 @@ def test_buffer_store_predicate_not_supported():
@T.prim_func(s_tir=True)
def func(b: T.handle):
B = T.match_buffer(b, (8,), "float32")
- B.vstore([T.Ramp(0, 2, 4)], T.Broadcast(1.0, 4),
predicate=T.Broadcast(T.bool(True), 4))
+ T.evaluate(
+ T.call_intrin(
+ "void",
+ "tirx.masked_store",
+ B,
+ T.Broadcast(1.0, 4),
+ T.Ramp(0, 2, 4),
+ T.Broadcast(T.bool(True), 4),
+ )
+ )
err_msg = "Predicated buffer store is not supported."
with pytest.raises(RuntimeError, match=err_msg):
@@ -58,8 +67,15 @@ def test_buffer_store_predicate_not_supported_gpu(target):
B = T.match_buffer(b, (6,), "float32")
T.func_attr({"global_symbol": "main"})
for i_0 in T.thread_binding(3, thread="threadIdx.x"):
- B.vstore(
- [T.Ramp(i_0, 1, 4)], T.Broadcast(1.0, 4),
predicate=T.Broadcast(T.bool(True), 4)
+ T.evaluate(
+ T.call_intrin(
+ "void",
+ "tirx.masked_store",
+ B,
+ T.Broadcast(1.0, 4),
+ T.Ramp(i_0, 1, 4),
+ T.Broadcast(T.bool(True), 4),
+ )
)
err_msg = "Predicated buffer store is not supported."
@@ -78,7 +94,13 @@ def test_buffer_load_predicate_not_supported():
for i_0 in range(4):
B.vstore(
[T.Ramp(0, 2, 4)],
- A.vload([T.Ramp(i_0, 1, 4)],
predicate=T.Broadcast(T.bool(True), 4)),
+ T.call_intrin(
+ "float32x4",
+ "tirx.masked_load",
+ A,
+ T.Ramp(i_0, 1, 4),
+ T.Broadcast(T.bool(True), 4),
+ ),
)
err_msg = "Predicated buffer load is not supported."
@@ -108,7 +130,13 @@ def test_buffer_load_predicate_not_supported_gpu(target):
for i_0 in T.thread_binding(3, thread="threadIdx.x"):
B.vstore(
[T.Ramp(0, 2, 4)],
- A.vload([T.Ramp(i_0, 1, 4)],
predicate=T.Broadcast(T.bool(True), 4)),
+ T.call_intrin(
+ "float32x4",
+ "tirx.masked_load",
+ A,
+ T.Ramp(i_0, 1, 4),
+ T.Broadcast(T.bool(True), 4),
+ ),
)
err_msg = "Predicated buffer load is not supported."
diff --git a/tests/python/codegen/test_target_codegen_llvm.py
b/tests/python/codegen/test_target_codegen_llvm.py
index 6f9f128a98..bdf617ff54 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -1255,7 +1255,13 @@ def test_invalid_volatile_masked_buffer_load():
def main(b: T.handle):
B = T.match_buffer(b, [4])
A = T.alloc_buffer((4,), annotations={"tirx.volatile": True})
- B[0:4] = A.vload([T.Ramp(0, 1, 4)],
predicate=T.Broadcast(T.bool(True), 4))
+ B[0:4] = T.call_intrin(
+ "float32x4",
+ "tirx.masked_load",
+ A,
+ T.Ramp(0, 1, 4),
+ T.Broadcast(T.bool(True), 4),
+ )
err_msg = "The masked load intrinsic does not support declaring load as
volatile."
with pytest.raises(RuntimeError, match=err_msg):
@@ -1271,7 +1277,13 @@ def test_invalid_volatile_masked_decl_buffer_load():
B = T.match_buffer(b, [4])
A = T.alloc_buffer((4,), annotations={"tirx.volatile": True})
A_alias = T.decl_buffer((4,), data=A.data)
- B[0:4] = A_alias.vload([T.Ramp(0, 1, 4)],
predicate=T.Broadcast(T.bool(True), 4))
+ B[0:4] = T.call_intrin(
+ "float32x4",
+ "tirx.masked_load",
+ A_alias,
+ T.Ramp(0, 1, 4),
+ T.Broadcast(T.bool(True), 4),
+ )
err_msg = "The masked load intrinsic does not support declaring load as
volatile."
with pytest.raises(RuntimeError, match=err_msg):
@@ -1285,10 +1297,15 @@ def test_invalid_volatile_masked_buffer_store():
@T.prim_func(s_tir=True)
def main():
A = T.alloc_buffer((4,), annotations={"tirx.volatile": True})
- A.vstore(
- [T.Ramp(0, 1, 4)],
- T.Broadcast(0.0, 4),
- predicate=T.Broadcast(T.bool(True), 4),
+ T.evaluate(
+ T.call_intrin(
+ "void",
+ "tirx.masked_store",
+ A,
+ T.Broadcast(0.0, 4),
+ T.Ramp(0, 1, 4),
+ T.Broadcast(T.bool(True), 4),
+ )
)
err_msg = "The masked store intrinsic does not support declaring store as
volatile."
diff --git a/tests/python/s_tir/analysis/test_sblock_access_region.py
b/tests/python/s_tir/analysis/test_sblock_access_region.py
index 15a84fae30..b9266be09a 100644
--- a/tests/python/s_tir/analysis/test_sblock_access_region.py
+++ b/tests/python/s_tir/analysis/test_sblock_access_region.py
@@ -45,6 +45,16 @@ def func() -> None:
T.evaluate(D.data)
[email protected]_func(s_tir=True)
+def masked_access_func() -> None:
+ A = T.sblock_alloc_buffer((16,), "float32")
+ B = T.sblock_alloc_buffer((16,), "float32")
+ with T.sblock():
+ mask = T.meta_var(T.Broadcast(T.bool(True), 4))
+ value = T.meta_var(T.masked_load("float32x4", A, T.Ramp(4, 1, 4),
mask))
+ T.masked_store(B, value, T.Ramp(8, 1, 4), mask)
+
+
@T.prim_func(s_tir=True)
def match_buffer_func() -> None:
with T.sblock("root"):
@@ -232,6 +242,18 @@ def test_block_access_region_detector():
)
+def test_masked_access_is_not_opaque():
+ root = masked_access_func.body.block
+ block = root.body.block
+ A, B = root.alloc_buffers
+ reads, writes, opaque = s_tir.analysis.get_sblock_access_region(block, {A:
A, B: B})
+ tvm.ir.assert_structural_equal(reads, [tvm.tirx.BufferRegion(A,
[Range.from_min_extent(4, 4)])])
+ tvm.ir.assert_structural_equal(
+ writes, [tvm.tirx.BufferRegion(B, [Range.from_min_extent(8, 4)])]
+ )
+ tvm.ir.assert_structural_equal(opaque, [])
+
+
def test_opaque_block():
alloc_buffers = opaque_block_func.body.block.alloc_buffers
buffer_var_map = {buf: buf for buf in alloc_buffers}
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
index 5eaf1f5d37..6ba06251c9 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
@@ -213,5 +213,40 @@ def test_vthread_vectorized():
assert allocate_node.buffer.ty.dtype == "int32x4"
+def test_vthread_rewrites_masked_accesses():
+ @T.prim_func(s_tir=True)
+ def before_func():
+ vthread = T.env_thread("vthread")
+ T.launch_thread(vthread, 2)
+ B = T.alloc_buffer((4,), "float32", scope="shared")
+ mask = T.meta_var(T.Broadcast(T.bool(True), 4))
+ loaded = T.meta_var(T.masked_load("float32x4", B, T.Ramp(0, 1, 4),
mask))
+ value = T.meta_var(loaded + T.Broadcast(T.Cast("float32", vthread), 4))
+ T.masked_store(B, value, T.Ramp(0, 1, 4), mask)
+
+ after = tvm.s_tir.transform.InjectVirtualThread()(
+ tvm.IRModule.from_expr(before_func.with_attr("global_symbol", "main"))
+ )["main"]
+ masked_calls = []
+
+ def visitor(node):
+ if isinstance(node, tvm.ir.Call) and node.op.name in {
+ "tirx.masked_load",
+ "tirx.masked_store",
+ }:
+ masked_calls.append(node)
+
+ tvm.tirx.stmt_functor.post_order_visit(after.body, visitor)
+ assert len(masked_calls) == 4
+ assert all(list(call.args[0].ty.shape) == [8] for call in masked_calls)
+ analyzer = tvm.arith.Analyzer()
+ assert sorted(int(analyzer.simplify(call.args[-2].base)) for call in
masked_calls) == [
+ 0,
+ 0,
+ 4,
+ 4,
+ ]
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py
b/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py
index 601931ab52..2f5184c69f 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py
@@ -587,5 +587,21 @@ def test_scalar_match_buffer_type_coercion():
_check(scalar_match_buffer_type_coercion,
transformed_scalar_match_buffer_type_coercion)
[email protected]_func(s_tir=True)
+def masked_match_buffer(a: T.handle) -> None:
+ A = T.match_buffer(a, (8,), "float32")
+ with T.sblock():
+ T.reads(A[2:6])
+ sub_A = T.match_buffer(A[2:6], (4,), offset_factor=1)
+ mask = T.meta_var(T.Broadcast(T.bool(True), 4))
+ T.evaluate(T.masked_load("float32x4", sub_A, T.Ramp(0, 1, 4), mask))
+
+
+def test_masked_match_buffer_fails_explicitly():
+ mod = tvm.IRModule.from_expr(masked_match_buffer)
+ with pytest.raises(RuntimeError, match="Predicated buffer access is not
currently supported"):
+ tvm.s_tir.transform.LowerMatchBuffer()(mod)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/tirx-base/test_tir_nodes.py
b/tests/python/tirx-base/test_tir_nodes.py
index 4af595a783..3f35ae5bf6 100644
--- a/tests/python/tirx-base/test_tir_nodes.py
+++ b/tests/python/tirx-base/test_tir_nodes.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: F811, F841
+# ruff: noqa: F841
import numpy as np
import pytest
@@ -427,75 +427,6 @@ def test_buffer_store_scalable_vec():
assert store.value.ty.dtype == "int32xvscalex4"
-def test_buffer_store_predicate_invalid_scalability():
- b = tvm.tirx.decl_buffer((24,), "int32")
- value = tvm.tirx.expr.Broadcast(1, 4 * tvm.tirx.vscale())
- index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
- predicate = tvm.tirx.expr.Broadcast(tvm.tirx.IntImm("int1", 1), 4)
-
- err_msg = "Predicate mask dtype and value dtype must both be scalable."
- with pytest.raises(RuntimeError, match=err_msg):
- tvm.tirx.BufferStore(b, value, [index], predicate)
-
-
-def test_buffer_store_predicate_invalid_lanes():
- b = tvm.tirx.decl_buffer((24,), "int32")
- value = tvm.tirx.expr.Broadcast(1, 4 * tvm.tirx.vscale())
- index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
- predicate = tvm.tirx.expr.Broadcast(tvm.tirx.IntImm("int1", 1), 8 *
tvm.tirx.vscale())
-
- err_msg = (
- "Got a predicate mask with 8 lanes, but trying to store a "
- "value with 4 lanes. The number of lanes must match."
- )
- with pytest.raises(RuntimeError, match=err_msg):
- tvm.tirx.BufferStore(b, value, [index], predicate)
-
-
-def test_buffer_store_predicate_elements_invalid_type():
- b = tvm.tirx.decl_buffer((24,), "int32")
- value = tvm.tirx.expr.Broadcast(1, 4 * tvm.tirx.vscale())
- index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
- predicate = tvm.tirx.expr.Broadcast(1, 4 * tvm.tirx.vscale())
-
- err_msg = "Predicate mask elements must be boolean values, but got int32."
- with pytest.raises(RuntimeError, match=err_msg):
- tvm.tirx.BufferStore(b, value, [index], predicate)
-
-
-def test_buffer_load_predicate_elements_invalid_type():
- b = tvm.tirx.decl_buffer((24,), "int32")
- index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
- predicate = tvm.tirx.expr.Broadcast(1, 4 * tvm.tirx.vscale())
-
- err_msg = "Predicate mask elements must be boolean values, but got int32."
- with pytest.raises(RuntimeError, match=err_msg):
- tvm.tirx.BufferLoad(b, [index], predicate)
-
-
-def test_buffer_store_predicate_invalid_scalability():
- b = tvm.tirx.decl_buffer((24,), "int32")
- index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
- predicate = tvm.tirx.expr.Broadcast(tvm.tirx.IntImm("int1", 1), 4)
-
- err_msg = "Predicate mask dtype and load indices must both be scalable."
- with pytest.raises(RuntimeError, match=err_msg):
- tvm.tirx.BufferLoad(b, [index], predicate)
-
-
-def test_buffer_store_predicate_invalid_lanes():
- b = tvm.tirx.decl_buffer((24,), "int32")
- index = tvm.tirx.expr.Ramp(0, 1, 4 * tvm.tirx.vscale())
- predicate = tvm.tirx.expr.Broadcast(tvm.tirx.IntImm("int1", 1), 8 *
tvm.tirx.vscale())
-
- err_msg = (
- "Got a predicate mask with 8 lanes, but trying to load a "
- "vector with 4 lanes. The number of lanes must match."
- )
- with pytest.raises(RuntimeError, match=err_msg):
- tvm.tirx.BufferLoad(b, [index], predicate)
-
-
def test_scalable_vec_cast():
b = tvm.tirx.decl_buffer((24,), "float32")
value = tvm.tirx.expr.Broadcast(1, 12 *
tvm.tirx.vscale()).astype("float32xvscalex12")
diff --git a/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py
b/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py
index a75b87ec56..ce14378598 100644
--- a/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py
+++ b/tests/python/tirx-transform/test_tir_transform_bf16_legalize.py
@@ -106,6 +106,86 @@ def test_bf16_simple_store_will_legalize():
tvm.ir.assert_structural_equal(after_storage,
BindTarget(target)(after_storage_legalize()))
+def test_bf16_masked_load_store_will_legalize():
+ def get_before():
+ @tvm.script.ir_module
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(Aptr: T.handle("bfloat16"), Cptr: T.handle("bfloat16")):
+ T.func_attr({"global_symbol": "main"})
+ A = T.decl_buffer((16,), "bfloat16", data=Aptr)
+ B = T.decl_buffer((16,), "bfloat16")
+ C = T.decl_buffer((16,), "bfloat16", data=Cptr)
+ mask = T.Broadcast(T.bool(True), 4)
+ T.evaluate(
+ T.call_intrin(
+ "void",
+ "tirx.masked_store",
+ B,
+ T.call_intrin("bfloat16x4", "tirx.masked_load", A,
T.Ramp(0, 1, 4), mask),
+ T.Ramp(0, 1, 4),
+ mask,
+ )
+ )
+ T.evaluate(
+ T.call_intrin(
+ "void",
+ "tirx.masked_store",
+ C,
+ T.call_intrin("bfloat16x4", "tirx.masked_load", B,
T.Ramp(0, 1, 4), mask),
+ T.Ramp(0, 1, 4),
+ mask,
+ )
+ )
+
+ return Before
+
+ target = Target("nvidia/geforce-rtx-2080-ti")
+ before = BindTarget(target)(get_before())
+ after_compute = tvm.tirx.transform.BF16ComputeLegalize()(before)
+ after_storage = tvm.tirx.transform.BF16StorageLegalize()(after_compute)
+
+ def collect(mod):
+ nodes = []
+ tvm.tirx.stmt_functor.post_order_visit(mod["main"].body, nodes.append)
+ buffers = {
+ node.buffer.name: str(node.buffer.dtype)
+ for node in nodes
+ if isinstance(node, tvm.tirx.DeclBuffer | tvm.tirx.AllocBuffer)
+ }
+ masked_loads = [
+ node
+ for node in nodes
+ if isinstance(node, tvm.ir.Call) and node.op.name ==
"tirx.masked_load"
+ ]
+ masked_stores = [
+ node
+ for node in nodes
+ if isinstance(node, tvm.ir.Call) and node.op.name ==
"tirx.masked_store"
+ ]
+ return buffers, masked_loads, masked_stores
+
+ compute_buffers, compute_loads, compute_stores = collect(after_compute)
+ assert compute_buffers == {"A": "bfloat16", "B": "float32", "C":
"bfloat16", "mask": "boolx4"}
+ assert sorted(str(load.ty) for load in compute_loads) == ["bfloat16x4",
"float32x4"]
+ assert sorted(str(store.args[1].ty) for store in compute_stores) == [
+ "bfloat16x4",
+ "float32x4",
+ ]
+
+ storage_buffers, storage_loads, storage_stores = collect(after_storage)
+ assert storage_buffers == {"A": "uint16", "B": "float32", "C": "uint16",
"mask": "boolx4"}
+ assert sorted(str(load.ty) for load in storage_loads) == [
+ "float32x4",
+ "float32x4",
+ "uint16x4",
+ ]
+ assert sorted(str(store.args[1].ty) for store in storage_stores) == [
+ "float32x4",
+ "uint16x4",
+ ]
+
+
def test_bf16_storage_compute_scope_will_legalize():
def get_before():
@tvm.script.ir_module
diff --git a/tests/python/tirx-transform/test_tir_transform_vectorize.py
b/tests/python/tirx-transform/test_tir_transform_vectorize.py
index 140655dbd5..94620ce7f5 100644
--- a/tests/python/tirx-transform/test_tir_transform_vectorize.py
+++ b/tests/python/tirx-transform/test_tir_transform_vectorize.py
@@ -193,10 +193,15 @@ def test_vectorize_if_scalable_extent():
T.float32(1), extent
)
else:
- A.vstore(
- [T.Ramp(0, 1, T.vscale() * 4)],
- T.Broadcast(T.float32(2), T.vscale() * 4),
- predicate=T.get_active_lane_mask("uint1xvscalex4", 0, n),
+ T.evaluate(
+ T.call_intrin(
+ "void",
+ "tirx.masked_store",
+ A,
+ T.Broadcast(T.float32(2), T.vscale() * 4),
+ T.Ramp(0, 1, T.vscale() * 4),
+ T.get_active_lane_mask("uint1xvscalex4", 0, n),
+ )
)
with tvm.target.Target(target):
@@ -537,16 +542,24 @@ def
test_vectorize_and_predicate_all_buffer_loads_stores():
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
for i_0 in range(4):
load_a = T.meta_var(
- A.vload(
- [T.Ramp(i_0 * 4, 1, 4)],
- predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+ T.call_intrin(
+ "float32x4",
+ "tirx.masked_load",
+ A,
+ T.Ramp(i_0 * 4, 1, 4),
+ T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
)
)
add_1 = T.meta_var(load_a + T.Broadcast(T.float32(1), 4))
- B.vstore(
- [T.Ramp(i_0 * 4, 1, 4)],
- add_1,
- predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+ T.evaluate(
+ T.call_intrin(
+ "void",
+ "tirx.masked_store",
+ B,
+ add_1,
+ T.Ramp(i_0 * 4, 1, 4),
+ T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+ )
)
mod = tvm.IRModule.from_expr(before)
@@ -601,15 +614,25 @@ def
test_vectorize_and_predicate_multiple_access_statements():
B = T.match_buffer(b, (16,), "float32")
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
for i_0 in range(4):
- A.vstore(
- [T.Ramp(i_0 * 4, 1, 4)],
- T.Broadcast(T.float32(2), 4),
- predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+ T.evaluate(
+ T.call_intrin(
+ "void",
+ "tirx.masked_store",
+ A,
+ T.Broadcast(T.float32(2), 4),
+ T.Ramp(i_0 * 4, 1, 4),
+ T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+ )
)
- B.vstore(
- [T.Ramp(i_0 * 4, 1, 4)],
- T.Broadcast(T.float32(1), 4),
- predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+ T.evaluate(
+ T.call_intrin(
+ "void",
+ "tirx.masked_store",
+ B,
+ T.Broadcast(T.float32(1), 4),
+ T.Ramp(i_0 * 4, 1, 4),
+ T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+ )
)
before_mod = tvm.IRModule.from_expr(before)
@@ -618,6 +641,37 @@ def
test_vectorize_and_predicate_multiple_access_statements():
tvm.ir.assert_structural_equal(after, expected)
+def test_vectorize_nested_predicates_preserve_both_masks():
+ rvv_target = tvm.target.Target(
+ {"kind": "llvm", "mtriple": "riscv64-unknown-linux-gnu", "mattr":
["+v"]}
+ )
+
+ @T.prim_func(s_tir=True)
+ def before(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")):
+ for i_0 in T.serial(4):
+ for i_1 in T.vectorized(4):
+ if i_0 * 4 + i_1 < 15:
+ if i_0 * 4 + i_1 < 14:
+ A[i_0 * 4 + i_1] = T.float32(1)
+ B[i_0 * 4 + i_1] = T.float32(2)
+
+ with tvm.target.Target(rvv_target):
+ after =
tvm.tirx.transform.VectorizeLoop()(tvm.IRModule.from_expr(before))["before"]
+
+ predicates = []
+
+ def collect_predicates(node):
+ if isinstance(node, tvm.ir.Call) and node.op.name ==
"tirx.masked_store":
+ predicates.append(node.args[-1])
+
+ tvm.tirx.stmt_functor.post_order_visit(after.body, collect_predicates)
+ assert len(predicates) == 2
+ assert any(
+ isinstance(predicate, tvm.ir.Call) and predicate.op.name ==
"tirx.bitwise_and"
+ for predicate in predicates
+ )
+
+
def test_vectorize_and_predicate_invalid_conditions():
@T.prim_func(s_tir=True)
def before(a: T.handle, b: T.handle):
@@ -703,16 +757,24 @@ def
test_vectorize_and_predicate_buffer_load_stores_with_sve_func_attr_target():
T.func_attr({"global_symbol": "main", "tirx.noalias": True, "target":
sve_target})
for i_0 in range(4):
load_a = T.meta_var(
- A.vload(
- [T.Ramp(i_0 * 4, 1, 4)],
- predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+ T.call_intrin(
+ "float32x4",
+ "tirx.masked_load",
+ A,
+ T.Ramp(i_0 * 4, 1, 4),
+ T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
)
)
add_1 = T.meta_var(load_a + T.Broadcast(T.float32(1), 4))
- B.vstore(
- [T.Ramp(i_0 * 4, 1, 4)],
- add_1,
- predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+ T.evaluate(
+ T.call_intrin(
+ "void",
+ "tirx.masked_store",
+ B,
+ add_1,
+ T.Ramp(i_0 * 4, 1, 4),
+ T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+ )
)
mod = tvm.IRModule.from_expr(before)
@@ -740,16 +802,24 @@ def
test_vectorize_and_predicate_buffer_load_stores_with_sve_attr_scope_target()
with T.attr(sve_target, "target", 0):
for i_0 in range(4):
load_a = T.meta_var(
- A.vload(
- [T.Ramp(i_0 * 4, 1, 4)],
- predicate=T.get_active_lane_mask("uint1x4", i_0 * 4,
14),
+ T.call_intrin(
+ "float32x4",
+ "tirx.masked_load",
+ A,
+ T.Ramp(i_0 * 4, 1, 4),
+ T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
)
)
add_1 = T.meta_var(load_a + T.Broadcast(T.float32(1), 4))
- B.vstore(
- [T.Ramp(i_0 * 4, 1, 4)],
- add_1,
- predicate=T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+ T.evaluate(
+ T.call_intrin(
+ "void",
+ "tirx.masked_store",
+ B,
+ add_1,
+ T.Ramp(i_0 * 4, 1, 4),
+ T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
+ )
)
mod = tvm.IRModule.from_expr(before)
diff --git a/tests/python/tirx/transform/test_stmt_functor.py
b/tests/python/tirx/transform/test_stmt_functor.py
index bf845c6516..889a536a6b 100644
--- a/tests/python/tirx/transform/test_stmt_functor.py
+++ b/tests/python/tirx/transform/test_stmt_functor.py
@@ -1217,7 +1217,7 @@ def test_op_call_pointer_config_visited_and_mutated():
def visit_buffer_load_(self, op):
new_op = super().visit_buffer_load_(op)
if op.buffer.same_as(mbar_buffer):
- return tir.BufferLoad(replacement, new_op.indices,
new_op.predicate)
+ return tir.BufferLoad(replacement, new_op.indices)
return new_op
updated = ReplaceMbarLoad().visit_stmt(op_call)
diff --git a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py
b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py
index dd02640f30..1927e07e34 100644
--- a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py
+++ b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py
@@ -446,20 +446,6 @@ def test_ir_builder_tir_buffer_store_scalable_vec():
assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
-def test_ir_builder_tir_buffer_store_predicate():
- buffer_a = T.Buffer((30,), "float32")
- value = T.broadcast(0.11, T.vscale() * 4)
- index = T.ramp(0, 1, T.vscale() * 4)
- predicate = T.broadcast(T.bool(True), T.vscale() * 4)
-
- with IRBuilder() as ib:
- T.buffer_store(buffer_a, value, [index], predicate)
-
- ir_actual = ib.get()
- ir_expected = tirx.BufferStore(buffer_a, value, [index], predicate)
- assert_structural_equal(ir_actual, ir_expected, map_free_vars=True)
-
-
def test_ir_builder_tir_evaluate():
with IRBuilder() as ib:
T.evaluate(0)
diff --git a/tests/python/tvmscript/test_tvmscript_printer_tir.py
b/tests/python/tvmscript/test_tvmscript_printer_tir.py
index 0cd538fb65..0ecadaba67 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_tir.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_tir.py
@@ -993,8 +993,27 @@ def test_predicated_load_store():
A = T.match_buffer(a, (128, 128), "float32")
B = T.match_buffer(b, (256, 256), "float32")
T.func_attr({"global_symbol": "func"})
- a_load = T.meta_var(A.vload([0, T.Ramp(0, 4, 4)],
predicate=T.Broadcast(T.bool(False), 4)))
- A.vstore([0, T.Ramp(0, 2, 4)], a_load,
predicate=T.Broadcast(T.bool(False), 4))
+ a_load = T.meta_var(
+ T.call_intrin(
+ "float32x4",
+ "tirx.masked_load",
+ A,
+ 0,
+ T.Ramp(0, 4, 4),
+ T.Broadcast(T.bool(False), 4),
+ )
+ )
+ T.evaluate(
+ T.call_intrin(
+ "void",
+ "tirx.masked_store",
+ A,
+ a_load,
+ 0,
+ T.Ramp(0, 2, 4),
+ T.Broadcast(T.bool(False), 4),
+ )
+ )
expected_output = """
# from tvm.script import tirx as T
@@ -1002,7 +1021,7 @@ def test_predicated_load_store():
@T.prim_func(s_tir=True)
def func(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256),
"float32")):
- A.vstore([0, T.Ramp(0, 2, 4)], A.vload([0, T.Ramp(0, 4, 4)],
predicate=T.Broadcast(T.bool(False), 4)), predicate=T.Broadcast(T.bool(False),
4))
+ T.masked_store(A, T.masked_load("float32x4", A, 0, T.Ramp(0, 4, 4),
T.Broadcast(T.bool(False), 4)), 0, T.Ramp(0, 2, 4), T.Broadcast(T.bool(False),
4))
"""
_assert_print(main, expected_output)
@@ -1014,16 +1033,24 @@ def test_predicated_buffer_load_store():
a: tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A"),
b: tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B"),
}
- buffer_load = tirx.BufferLoad(
- buffer=buffers[b],
- indices=[0, tirx.Ramp(0, 4, 4)],
- predicate=tirx.Broadcast(tirx.IntImm("bool", 0), 4),
+ buffer_load = tirx.call_intrin(
+ "float32x4",
+ "tirx.masked_load",
+ buffers[b],
+ 0,
+ tirx.Ramp(0, 4, 4),
+ tirx.Broadcast(tirx.IntImm("bool", 0), 4),
)
- body = tirx.BufferStore(
- buffer=buffers[a],
- value=buffer_load,
- indices=[0, tirx.Ramp(0, 2, 4)],
- predicate=tirx.Broadcast(tirx.IntImm("bool", 0), 4),
+ body = tirx.Evaluate(
+ tirx.call_intrin(
+ "void",
+ "tirx.masked_store",
+ buffers[a],
+ buffer_load,
+ 0,
+ tirx.Ramp(0, 2, 4),
+ tirx.Broadcast(tirx.IntImm("bool", 0), 4),
+ )
)
func = tirx.PrimFunc(
params=[buffers[a], buffers[b]],
@@ -1037,7 +1064,7 @@ def test_predicated_buffer_load_store():
@T.prim_func(private=True, s_tir=True)
def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256),
"float32")):
- A.vstore([0, T.Ramp(0, 2, 4)], B.vload([0, T.Ramp(0, 4, 4)],
predicate=T.Broadcast(T.bool(False), 4)), predicate=T.Broadcast(T.bool(False),
4))
+ T.masked_store(A, T.masked_load("float32x4", B, 0, T.Ramp(0, 4, 4),
T.Broadcast(T.bool(False), 4)), 0, T.Ramp(0, 2, 4), T.Broadcast(T.bool(False),
4))
"""
_assert_print(func, expected_output)
@@ -1051,8 +1078,16 @@ def test_predicated_scalable_load_store():
B = T.match_buffer(b, (256, 256), "float32")
T.func_attr({"global_symbol": "func"})
mask = T.meta_var(T.get_active_lane_mask("uint1xvscalex4", 0, 13))
- a_load = T.meta_var(A.vload([0, T.Ramp(0, 4, T.vscale() * 4)],
predicate=mask))
- A.vstore([0, T.Ramp(0, 2, T.vscale() * 4)], a_load, predicate=mask)
+ a_load = T.meta_var(
+ T.call_intrin(
+ "float32xvscalex4", "tirx.masked_load", A, 0, T.Ramp(0, 4,
T.vscale() * 4), mask
+ )
+ )
+ T.evaluate(
+ T.call_intrin(
+ "void", "tirx.masked_store", A, a_load, 0, T.Ramp(0, 2,
T.vscale() * 4), mask
+ )
+ )
expected_output = """
# from tvm.script import tirx as T
@@ -1060,11 +1095,25 @@ def test_predicated_scalable_load_store():
@T.prim_func(s_tir=True)
def func(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256),
"float32")):
- A.vstore([0, T.Ramp(0, 2, T.vscale() * 4)], A.vload([0, T.Ramp(0, 4,
T.vscale() * 4)], predicate=T.get_active_lane_mask("uint1xvscalex4", 0, 13)),
predicate=T.get_active_lane_mask("uint1xvscalex4", 0, 13))
+ T.masked_store(A, T.masked_load("float32xvscalex4", A, 0, T.Ramp(0, 4,
T.vscale() * 4), T.get_active_lane_mask("uint1xvscalex4", 0, 13)), 0, T.Ramp(0,
2, T.vscale() * 4), T.get_active_lane_mask("uint1xvscalex4", 0, 13))
"""
_assert_print(main, expected_output)
+def test_masked_load_prevents_scalar_allocation_init_fusion():
+ from tvm.script import tirx as T
+
+ @T.prim_func(s_tir=True)
+ def main():
+ A = T.alloc_buffer((1,), "float32x4")
+ A[0] = T.masked_load("float32x4", A, 0, T.Broadcast(T.bool(True), 4))
+
+ source = main.script()
+ assert "A = T.alloc_buffer" in source
+ assert "A[0] = T.masked_load" in source
+ tvm.ir.assert_structural_equal(tvm.script.from_source(source), main)
+
+
def test_vload_with_explicit_scalable_data_type():
from tvm.script import tirx as T
diff --git a/tests/python/tvmscript/test_tvmscript_roundtrip.py
b/tests/python/tvmscript/test_tvmscript_roundtrip.py
index 94d4169e38..ddc2c09603 100644
--- a/tests/python/tvmscript/test_tvmscript_roundtrip.py
+++ b/tests/python/tvmscript/test_tvmscript_roundtrip.py
@@ -2456,9 +2456,24 @@ def predicated_buffer_load_store():
B = T.match_buffer(b, (8,), "float32")
for i_0 in range(4):
load_a = T.meta_var(
- A.vload([T.Ramp(i_0, 1, 4)],
predicate=T.Broadcast(T.bool(True), 4))
+ T.call_intrin(
+ "float32x4",
+ "tirx.masked_load",
+ A,
+ T.Ramp(i_0, 1, 4),
+ T.Broadcast(T.bool(True), 4),
+ )
+ )
+ T.evaluate(
+ T.call_intrin(
+ "void",
+ "tirx.masked_store",
+ B,
+ load_a,
+ T.Ramp(0, 2, 4),
+ T.Broadcast(T.bool(True), 4),
+ )
)
- B.vstore([T.Ramp(0, 2, 4)], load_a,
predicate=T.Broadcast(T.bool(True), 4))
return func
