lhutton1 commented on code in PR #16966:
URL: https://github.com/apache/tvm/pull/16966#discussion_r1599972721
##########
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:
Done
##########
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:
Done. There was a bit of confusion here since LLVM intrinsics expect
predicate masks of `int1` values. However, boolean types are already converted
to `int1` in the codegen backend. So I have required `uint1`/boolean types in
TVM.
##########
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:
Done
##########
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:
Thanks for the detailed description. To save this PR exploding further, I've
added the minimum changes necessary for the vectorizer pass. Perhaps we can
clean this up and fix the target checks in the analyzer in a separate PR?
##########
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:
Got it, thanks. I think part of the confusion here was that I wasn't sure
whether it was intended to be used this way - I couldn't find vload/vstore used
this way in the codebase
##########
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:
Done
##########
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:
Done
##########
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:
Agreed, thanks. I've added support for previous versions of LLVM. I've
checked the build with the following versions of LLVM: 7*, 8*, 9*, 10, 11, 12,
13, 17
* fails to build due to other seemingly unrelated errors
##########
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:
Done
--
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]
