Lunderberg commented on code in PR #16966:
URL: https://github.com/apache/tvm/pull/16966#discussion_r1594285519
##########
include/tvm/script/ir_builder/tir/ir.h:
##########
@@ -411,8 +411,10 @@ Var EnvThread(String thread_tag, DataType dtype =
DataType::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 int1 values indicating which lanes of a
vector are to be
Review Comment:
Can we specify that the number of lanes for the `predicate` must be equal to
the number of lanes in the `value`?
##########
tests/python/tir-transform/test_tir_transform_vectorize.py:
##########
@@ -488,5 +529,170 @@ def main(A: T.Buffer((16,), "float32")):
tvm.tir.transform.VectorizeLoop()(Mod)
+def test_vectorize_and_predicate_all_buffer_loads_stores():
+ @T.prim_func
+ def before(a: T.handle, b: T.handle):
+ A = T.match_buffer(a, (16,), "float32")
+ B = T.match_buffer(b, (16,), "float32")
+ T.func_attr({"global_symbol": "main", "tir.noalias": True})
+ for i_0 in T.serial(T.ceildiv(14, 4)):
+ for i_1 in T.vectorized(4):
+ if i_0 * 4 + i_1 < 14:
+ B[i_0 * 4 + i_1] = A[i_0 * 4 + i_1] + 1.0
+
+ @T.prim_func
+ def expected(a: T.handle, b: T.handle):
+ A = T.match_buffer(a, (16,), "float32")
+ B = T.match_buffer(b, (16,), "float32")
+ T.func_attr({"global_symbol": "main", "tir.noalias": T.bool(True)})
+ for i_0 in range(4):
+ load_a = T.meta_var(
+ A.load(
+ [T.Ramp(i_0 * 4, 1, 4)],
predicate=T.get_active_lane_mask("int1x4", i_0 * 4, 14)
+ )
+ )
+ add_1 = T.meta_var(load_a + T.Broadcast(T.float32(1), 4))
+ B.store(
+ add_1,
+ [T.Ramp(i_0 * 4, 1, 4)],
+ predicate=T.get_active_lane_mask("int1x4", i_0 * 4, 14),
+ )
+
+ mod = tvm.IRModule.from_expr(before)
+ with
tvm.transform.PassContext(config={"tir.enable_buffer_level_predication": True}):
+ after = tvm.tir.transform.VectorizeLoop()(mod)["main"]
+ tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_vectorize_and_predicate_some_buffer_loads_stores():
+ # Currently revert to scalarizing the block if not all accesses
+ # have been predicated, otherwise incorrect code is generated.
+ @T.prim_func
+ def before(a: T.handle, b: T.handle):
+ A = T.match_buffer(a, (16,), "float32")
+ B = T.match_buffer(b, (16,), "float32")
+ T.func_attr({"global_symbol": "main", "tir.noalias": True})
+ for i_0 in T.serial(T.ceildiv(14, 4)):
+ for i_1 in T.vectorized(4):
+ if i_0 * 4 + i_1 < 14:
+ B[i_0 * 4 + i_1] = A[i_0] + 1.0
+
+ @T.prim_func
+ def expected(a: T.handle, b: T.handle):
+ A = T.match_buffer(a, (16,), "float32")
+ B = T.match_buffer(b, (16,), "float32")
+ T.func_attr({"global_symbol": "main", "tir.noalias": T.bool(True)})
+ for i_0, i_1_s in T.grid(4, 4):
+ if i_0 * 4 + i_1_s < 14:
+ B[i_0 * 4 + i_1_s] = A[i_0] + T.float32(1)
+
+ mod = tvm.IRModule.from_expr(before)
+ with
tvm.transform.PassContext(config={"tir.enable_buffer_level_predication": True}):
+ after = tvm.tir.transform.VectorizeLoop()(mod)["main"]
+ tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_vectorize_and_predicate_multiple_access_statements():
+ @T.prim_func
+ def before(a: T.handle, b: T.handle):
+ A = T.match_buffer(a, (16,), "float32")
+ B = T.match_buffer(b, (16,), "float32")
+ T.func_attr({"global_symbol": "main", "tir.noalias": True})
+ for i_0 in T.serial(T.ceildiv(14, 4)):
+ for i_1 in T.vectorized(4):
+ if i_0 * 4 + i_1 < 14:
+ A[i_0 * 4 + i_1] = 2.0
+ B[i_0 * 4 + i_1] = 1.0
+
+ @T.prim_func
+ def expected(a: T.handle, b: T.handle):
+ A = T.match_buffer(a, (16,), "float32")
+ B = T.match_buffer(b, (16,), "float32")
+ T.func_attr({"global_symbol": "main", "tir.noalias": T.bool(True)})
+ for i_0 in range(4):
+ A.store(
+ T.Broadcast(T.float32(2), 4),
+ [T.Ramp(i_0 * 4, 1, 4)],
+ predicate=T.get_active_lane_mask("int1x4", i_0 * 4, 14),
+ )
+ B.store(
+ T.Broadcast(T.float32(1), 4),
+ [T.Ramp(i_0 * 4, 1, 4)],
+ predicate=T.get_active_lane_mask("int1x4", i_0 * 4, 14),
+ )
+
+ before_mod = tvm.IRModule.from_expr(before)
+ with
tvm.transform.PassContext(config={"tir.enable_buffer_level_predication": True}):
+ after = tvm.tir.transform.VectorizeLoop()(before_mod)["main"]
+ tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_vectorize_and_predicate_invalid_conditions():
+ @T.prim_func
+ def before(a: T.handle, b: T.handle):
+ A = T.match_buffer(a, (16,), "float32")
+ B = T.match_buffer(b, (16,), "float32")
+ T.func_attr({"global_symbol": "main", "tir.noalias": True})
+ for i_0 in T.serial(T.ceildiv(14, 4)):
+ for i_1 in T.vectorized(4):
+ if i_0 * 4 + i_1 > 14:
+ A[i_0 * 4 + i_1] = 2.0
+ if 14 < i_0 * 4 + i_1:
+ A[i_0 * 4 + i_1] = 2.0
+ if i_0 * 4 + i_1 < i_0 * 4 + i_1:
+ A[i_0 * 4 + i_1] = 2.0
+
+ @T.prim_func
+ def expected(a: T.handle, b: T.handle):
+ A = T.match_buffer(a, (16,), "float32")
+ B = T.match_buffer(b, (16,), "float32")
+ T.func_attr({"global_symbol": "main", "tir.noalias": T.bool(True)})
+ for i_0 in range(4):
+ for i_1_s in range(4):
+ if i_0 * 4 + i_1_s > 14:
+ A[i_0 * 4 + i_1_s] = T.float32(2)
+ for i_1_s in range(4):
+ if 14 < i_0 * 4 + i_1_s:
+ A[i_0 * 4 + i_1_s] = T.float32(2)
+ for i_1_s in range(4):
+ if i_0 * 4 + i_1_s < i_0 * 4 + i_1_s:
+ A[i_0 * 4 + i_1_s] = T.float32(2)
+
+ before_mod = tvm.IRModule.from_expr(before)
+ with
tvm.transform.PassContext(config={"tir.enable_buffer_level_predication": True}):
+ after = tvm.tir.transform.VectorizeLoop()(before_mod)["main"]
+ tvm.ir.assert_structural_equal(after, expected)
+
+
+def test_vectorize_with_explicitly_disabled_buffer_level_predication():
+ # Since the target is has the SVe feature, buffer level predication is
enabled
+ # by default. However, it has been explicitely disabled by the pass context
+ # option, so no buffer-level predicates should be added.
+ @T.prim_func
+ def before(a: T.handle, b: T.handle):
+ A = T.match_buffer(a, (16,), "float32")
+ B = T.match_buffer(b, (16,), "float32")
+ T.func_attr({"global_symbol": "main", "tir.noalias": True})
+ for i_0 in T.serial(T.ceildiv(14, 4)):
+ for i_1 in T.vectorized(4):
+ if i_0 * 4 + i_1 < 14:
+ B[i_0 * 4 + i_1] = A[i_0 * 4 + i_1] + 1.0
+
+ @T.prim_func
+ def expected(a: T.handle, b: T.handle):
+ A = T.match_buffer(a, (16,), "float32")
+ B = T.match_buffer(b, (16,), "float32")
+ T.func_attr({"global_symbol": "main", "tir.noalias": True})
+ for i_0, i_1_s in T.grid(4, 4):
+ if i_0 * 4 + i_1_s < 14:
+ B[i_0 * 4 + i_1_s] = A[i_0 * 4 + i_1_s] + T.float32(1)
+
+ mod = tvm.IRModule.from_expr(before)
+ with
tvm.transform.PassContext(config={"tir.enable_buffer_level_predication":
False}):
+ with tvm.target.Target("llvm -mtriple=aarch64-linux-gnu -mattr=+sve"):
+ after = tvm.tir.transform.VectorizeLoop()(mod)["main"]
+ tvm.ir.assert_structural_equal(after, expected)
+
Review Comment:
Can we also add test cases for the target override both in the function
attributes (`T.func_attr({'target': T.target(...)})`) and in an attr stmt
(`T.attr(T.target(...), "target", 0)`)?
([Link](https://github.com/apache/tvm/blob/main/tests/python/tir-transform/test_tir_transform_annotate_device_regions.py#L27)
to example for the TVMScript syntax.)
##########
src/tir/transforms/vectorize_loop.cc:
##########
@@ -72,6 +72,126 @@ inline PrimExpr BroadcastTo(PrimExpr e, int lanes, bool
is_scalable) {
return Broadcast(e, CreateNewLanes(is_scalable, lanes));
}
+bool EnableBufferLevelPredication() {
Review Comment:
This will require updating `arith::TargetHasSVE()` as well. Apologies for
not catching it earlier in #16893.
##########
include/tvm/tir/expr.h:
##########
@@ -675,7 +679,8 @@ class BufferLoadNode : public PrimExprNode {
*/
class BufferLoad : public PrimExpr {
public:
- TVM_DLL explicit BufferLoad(Buffer buffer, Array<PrimExpr> indices, Span
span = Span());
+ TVM_DLL explicit BufferLoad(Buffer buffer, Array<PrimExpr> indices,
+ PrimExpr predicate = PrimExpr(), Span span =
Span());
Review Comment:
Using `PrimExpr()` as the default would have an undefined object as the
predicate. While this is legal, it is rather unexpected. While changing
`PrimExpr` to be non-nullable would involve more breakage than I'd want to
handle at the moment, we should avoid introducing additional cases where
`PrimExpr` contains null.
Changing this to `Optional<PrimExpr> predicate = NullOpt` would avoid this
issue.
##########
include/tvm/tir/stmt.h:
##########
@@ -261,7 +265,7 @@ class BufferStoreNode : public StmtNode {
class BufferStore : public Stmt {
public:
TVM_DLL explicit BufferStore(Buffer buffer, PrimExpr value, Array<PrimExpr>
indices,
- Span span = Span());
+ PrimExpr predicate = PrimExpr(), Span span =
Span());
Review Comment:
Same request here, `Optional<PrimExpr> predicate = NullOpt`.
##########
python/tvm/tir/expr.py:
##########
@@ -1093,20 +1093,27 @@ class BufferLoad(PrimExprWithOp):
The buffer to be loaded.
indices : List[PrimExpr]
- The buffer indices.
+ The buffer indices to load values from.
span : Optional[Span]
The location of this expression in the source code.
+
+ predicate : Optional[PrimExpr]
+ A vector mask of int1 values indicating which lanes of a vector are to
be loaded.
Review Comment:
At first glance, the `int1` reads like a typo. Can we explicitly call it
`boolean values` instead?
##########
tests/python/codegen/test_target_codegen_aarch64.py:
##########
@@ -700,5 +700,31 @@ def before(a: T.handle):
assert "get.active.lane.mask" in ll
[email protected](
Review Comment:
This validates that we have the correct output for architectures that
support SVE, but it doesn't test the behavior of other targets that do not
(yet) support predicated loads/stores. While the `VectorizeLoop` pass would
only insert a predicated load/store for targets that support it, the predicated
load/store could still be generated in hand-written kernels, or through other
transforms in the future.
Can we add a test, parametrized over each target tested in CI, which
attempts to compile a PrimFunc containing predicated loads/stores? For each
target that supports sve, `tvm.build` should compile without error, and for
each target that does not, `tvm.build` should raise an exception.
##########
src/tir/transforms/vectorize_loop.cc:
##########
@@ -72,6 +72,126 @@ inline PrimExpr BroadcastTo(PrimExpr e, int lanes, bool
is_scalable) {
return Broadcast(e, CreateNewLanes(is_scalable, lanes));
}
+bool EnableBufferLevelPredication() {
Review Comment:
I like the fallback, where it checks whether a target supports SVE in order
to apply the vectorization if not specified in the pass context. However, this
function should accept a `Target` argument from the calling scope rather than
using `Target::Current()`, because the `Target::Current()` may be overridden.
The override can be either at the function level
(`prim_func->GetAttr<Target>(tvm::attr::kTarget)`), or for individual blocks of
TIR (an `AttrStmt` with `attr_stmt->attr_key==tvm::attr::kTarget`).
##########
src/target/llvm/codegen_llvm.cc:
##########
@@ -1768,11 +1774,17 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const
BufferLoadNode* op) {
std::vector<llvm::Value*> loads;
- auto make_load = [this, &loads](TypedPointer buffer_ptr, int /* subelement_i
*/, int alignment,
- bool is_volatile) {
+ auto make_load = [this, &loads](TypedPointer buffer_ptr, int /* subelement_i
*/,
+ llvm::Value* predicate, int alignment, bool
is_volatile) {
#if TVM_LLVM_VERSION >= 110
- auto load = builder_->CreateAlignedLoad(buffer_ptr.type, buffer_ptr.addr,
- llvm::Align(alignment),
is_volatile);
+ llvm::Instruction* load = nullptr;
+ if (predicate != NULL) {
+ load = builder_->CreateMaskedLoad(buffer_ptr.type, buffer_ptr.addr,
llvm::Align(alignment),
+ predicate);
+ } else {
+ load = builder_->CreateAlignedLoad(buffer_ptr.type, buffer_ptr.addr,
llvm::Align(alignment),
+ is_volatile);
+ }
#elif TVM_LLVM_VERSION >= 80
auto load =
Review Comment:
The PR only adds `CreateMaskedLoad` when `TVM_LLVM_VERSION >= 110`. If
somebody is using an older version of LLVM, it would silently ignore the
predicate for the load/store. We should either support it, or throw an
exception.
It looks like `CreateMaskedLoad` has been supported in LLVM since [this
commit](https://github.com/llvm/llvm-project/commit/f1de34b84dea91b5060ee0fafbadaad5deaf199c),
so I'd lean toward adding it in the other `#elif` branches.
##########
python/tvm/tir/buffer.py:
##########
@@ -141,6 +141,57 @@ def vstore(self, begin, value):
begin = (begin,) if isinstance(begin, (int, PrimExpr)) else begin
return _ffi_api.BufferVStore(self, begin, value) # type: ignore
+ def load(self, indices, predicate=None):
Review Comment:
These look very similar to the existing `vload` and `vstore` methods.
Should they be extended instead with a `predicate` argument instead, rather
than introducing new methods?
##########
python/tvm/tir/stmt.py:
##########
@@ -224,24 +224,29 @@ class BufferStore(Stmt):
indices : List[PrimExpr]
The indices location to be stored.
+ predicate : Optional[PrimExpr]
+ A vector mask of int1 values indicating which lanes of a vector are to
be stored.
Review Comment:
Same nitpick here, using `boolean` instead of `int1`.
##########
src/tir/ir/stmt.cc:
##########
@@ -458,7 +458,8 @@
TVM_REGISTER_GLOBAL("tir.Evaluate").set_body_typed([](PrimExpr value, Span span)
TVM_REGISTER_NODE_TYPE(EvaluateNode);
// BufferStore
-BufferStore::BufferStore(Buffer buffer, PrimExpr value, Array<PrimExpr>
indices, Span span) {
+BufferStore::BufferStore(Buffer buffer, PrimExpr value, Array<PrimExpr>
indices, PrimExpr predicate,
+ Span span) {
ICHECK_EQ(buffer->shape.size(), indices.size())
Review Comment:
We should add an assert here, that if `predicate` is defined, the number of
lanes in `value` and in `predicate` are identical.
##########
src/target/llvm/codegen_llvm.cc:
##########
@@ -1768,11 +1774,17 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const
BufferLoadNode* op) {
std::vector<llvm::Value*> loads;
- auto make_load = [this, &loads](TypedPointer buffer_ptr, int /* subelement_i
*/, int alignment,
- bool is_volatile) {
+ auto make_load = [this, &loads](TypedPointer buffer_ptr, int /* subelement_i
*/,
+ llvm::Value* predicate, int alignment, bool
is_volatile) {
#if TVM_LLVM_VERSION >= 110
- auto load = builder_->CreateAlignedLoad(buffer_ptr.type, buffer_ptr.addr,
- llvm::Align(alignment),
is_volatile);
+ llvm::Instruction* load = nullptr;
+ if (predicate != NULL) {
+ load = builder_->CreateMaskedLoad(buffer_ptr.type, buffer_ptr.addr,
llvm::Align(alignment),
+ predicate);
Review Comment:
It looks like this ignores the `is_volatile` argument. Can we provide
`is_volatile` to `CreateMaskedLoad`? If not, we should specify that these two
arguments are incompatible and throw an exception if they are both provided.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
