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 057449d33f [REFACTOR][TIRx] Use function targets in VectorizeLoop
(#20550)
057449d33f is described below
commit 057449d33faaca5d109b2c393ec25015baa438c2
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 5 13:48:58 2026 -0400
[REFACTOR][TIRx] Use function targets in VectorizeLoop (#20550)
VectorizeLoop now selects target-dependent behavior from each PrimFunc's
target attribute. Removing ambient and scoped target overrides gives
each function one vectorization target while preserving
target-independent fixed-width vectorization without a target attribute.
Update the pass documentation and target-dependent vectorization cases
to use the function target contract. Attribute statements remain intact
for other passes.
---
include/tvm/tirx/transform.h | 5 +
python/tvm/tirx/transform/transform.py | 5 +
src/tirx/transform/vectorize_loop.cc | 20 +-
.../tirx-transform/test_tir_transform_vectorize.py | 201 +++++++++------------
4 files changed, 100 insertions(+), 131 deletions(-)
diff --git a/include/tvm/tirx/transform.h b/include/tvm/tirx/transform.h
index 1dbbe63441..dbbfe60689 100644
--- a/include/tvm/tirx/transform.h
+++ b/include/tvm/tirx/transform.h
@@ -62,6 +62,11 @@ TVM_DLL Pass CreatePrimFuncPass(
/*!
* \brief Lower vectorization loops.
*
+ * Target-dependent vectorization uses the PrimFunc's target attribute, not an
+ * ambient target context or nested target attributes. Target-independent
+ * fixed-width loops do not require a target. Code needing different
vectorization
+ * targets must be separated into functions before this pass.
+ *
* \param enable_vectorize Whether vectorization is enabled.
*
* \return The pass.
diff --git a/python/tvm/tirx/transform/transform.py
b/python/tvm/tirx/transform/transform.py
index 365a894343..57d46d03dd 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -52,6 +52,11 @@ def Apply(ftransform):
def VectorizeLoop(enable_vectorize: bool = True):
"""Lower vectorization loops.
+ Target-dependent vectorization uses the PrimFunc's ``target`` attribute,
+ not an ambient target context or nested target attributes.
Target-independent
+ fixed-width loops do not require a target. Code needing different
vectorization
+ targets must be separated into functions before this pass.
+
Parameters
----------
enable_vectorize : bool
diff --git a/src/tirx/transform/vectorize_loop.cc
b/src/tirx/transform/vectorize_loop.cc
index 7cdf830474..4e526e19b1 100644
--- a/src/tirx/transform/vectorize_loop.cc
+++ b/src/tirx/transform/vectorize_loop.cc
@@ -1308,11 +1308,8 @@ class LoopVectorizer : public StmtExprMutator {
return StmtExprMutator::Mutate(value, inplace_mode);
}
- explicit LoopVectorizer(DictAttrs attrs) {
- if (auto opt_target = attrs.GetAttr<Target>(tvm::attr::kTarget)) {
- target_ = opt_target.value();
- }
- }
+ explicit LoopVectorizer(DictAttrs attrs)
+ :
target_(attrs.GetAttr<Target>(tvm::attr::kTarget).value_or(Target(nullptr))) {}
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
if (op->kind == ForKind::kVectorized) {
@@ -1347,17 +1344,6 @@ class LoopVectorizer : public StmtExprMutator {
}
}
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- if (op->attr_key == tvm::attr::kTarget) {
- Target previous_target = target_;
- target_ = op->node.as<Target>().value();
- auto result = StmtExprMutator::Mutate_(op, inplace_mode);
- target_ = previous_target;
- return result;
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
-
private:
Stmt VectorizeFixedLoopForRVV(const ForNode* op, int64_t extent) {
// Match the existing TIRx scalable-vector convention. LLVM/RVV still
@@ -1393,7 +1379,7 @@ class LoopVectorizer : public StmtExprMutator {
return this->Mutate(loop, InplaceMode::kDisallow).ValueOrUnchanged(loop);
}
- Target target_ = Target::Current();
+ Target target_;
};
class VectorizeSkipper : public StmtExprMutator {
diff --git a/tests/python/tirx-transform/test_tir_transform_vectorize.py
b/tests/python/tirx-transform/test_tir_transform_vectorize.py
index e675d4890f..17958fea2e 100644
--- a/tests/python/tirx-transform/test_tir_transform_vectorize.py
+++ b/tests/python/tirx-transform/test_tir_transform_vectorize.py
@@ -41,6 +41,7 @@ def test_vectorize_loop(extent, target):
class Before:
@T.prim_func
def main(A: T.Tensor((16,), "float32")):
+ T.func_attr({"target": target})
for j in T.vectorized(0, extent):
A[j] = 1
@@ -48,11 +49,11 @@ def test_vectorize_loop(extent, target):
class After:
@T.prim_func
def main(A: T.Tensor((16,), "float32")):
+ T.func_attr({"target": target})
A[T.Ramp(0, 1, extent)] = T.Broadcast(1, extent)
- with tvm.target.Target(target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
def test_vectorize_vector():
@@ -79,13 +80,13 @@ def test_vectorize_vector_scalable_error():
class Module:
@T.prim_func
def main(A: T.Tensor((25,), "float32")):
+ T.func_attr({"target": sve_target})
for j in T.vectorized(T.vscale() * 4):
A[T.ramp(j * 4, 1, 4)] = T.Broadcast(T.float32(1), 4)
error_msg = "Creating scalable vectors from existing vectors is not
supported."
- with tvm.target.Target(sve_target):
- with pytest.raises(tvm.error.InternalError, match=error_msg):
- tvm.tirx.transform.VectorizeLoop()(Module)
+ with pytest.raises(tvm.error.InternalError, match=error_msg):
+ tvm.tirx.transform.VectorizeLoop()(Module)
def test_vectorize_vector_scalable_error2():
@@ -106,6 +107,7 @@ def test_vectorize_vector_scalable_error3():
class Module:
@T.prim_func
def main(A: T.Tensor((25,), "float32")):
+ T.func_attr({"target": sve_target})
for j in T.vectorized(4):
A[T.ramp(j * T.vscale() * 4, 1, T.vscale() * 4)] = T.Broadcast(
T.float32(1), T.vscale() * 4
@@ -113,8 +115,7 @@ def test_vectorize_vector_scalable_error3():
error_msg = "Vectorizing over existing scalable vectors is not supported."
with pytest.raises(tvm.error.InternalError, match=error_msg):
- with tvm.target.Target(sve_target):
- tvm.tirx.transform.VectorizeLoop()(Module)
+ tvm.tirx.transform.VectorizeLoop()(Module)
def test_vectorize_vector_scalable_error4():
@@ -122,6 +123,7 @@ def test_vectorize_vector_scalable_error4():
class Module:
@T.prim_func(private=True)
def main(A: T.Tensor((25,), "float32")):
+ T.func_attr({"target": sve_target})
for j in T.vectorized(T.vscale() * 4):
A[T.ramp(j * T.vscale() * 4, 1, T.vscale() * 4)] = T.Broadcast(
T.float32(1), T.vscale() * 4
@@ -129,8 +131,7 @@ def test_vectorize_vector_scalable_error4():
error_msg = "Creating scalable vectors from existing vectors is not
supported."
with pytest.raises(tvm.error.InternalError, match=error_msg):
- with tvm.target.Target(sve_target):
- tvm.tirx.transform.VectorizeLoop()(Module)
+ tvm.tirx.transform.VectorizeLoop()(Module)
def test_vectorize_with_if():
@@ -141,6 +142,7 @@ def test_vectorize_with_if():
class Before:
@T.prim_func
def main(A: T.Tensor((25,), "float32"), n: T.int32, x: T.int32):
+ T.func_attr({"target": target})
for i in T.vectorized(extent):
if x < n:
A[i] = A[i] + T.float32(1)
@@ -152,6 +154,7 @@ def test_vectorize_with_if():
class After:
@T.prim_func
def main(A: T.Tensor((25,), "float32"), n: T.int32, x: T.int32):
+ T.func_attr({"target": target})
if x < n:
A[T.Ramp(0, 1, extent)] = A[T.Ramp(0, 1, extent)] +
T.Broadcast(
T.float32(1), extent
@@ -161,9 +164,8 @@ def test_vectorize_with_if():
if i_s < n:
A[i_s] = T.float32(2)
- with tvm.target.Target(target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
def test_vectorize_if_scalable_extent():
@@ -174,6 +176,7 @@ def test_vectorize_if_scalable_extent():
class Before:
@T.prim_func
def main(A: T.Tensor((25,), "float32"), n: T.int32, x: T.int32):
+ T.func_attr({"target": target})
for i in T.vectorized(extent):
if x < n:
A[i] = A[i] + T.float32(1)
@@ -185,6 +188,7 @@ def test_vectorize_if_scalable_extent():
class After:
@T.prim_func
def main(A: T.Tensor((25,), "float32"), n: T.int32, x: T.int32):
+ T.func_attr({"target": target})
if x < n:
A[T.Ramp(0, 1, extent)] = A[T.Ramp(0, 1, extent)] +
T.Broadcast(
T.float32(1), extent
@@ -201,9 +205,8 @@ def test_vectorize_if_scalable_extent():
)
)
- with tvm.target.Target(target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() *
4, sve_target)])
@@ -212,6 +215,7 @@ def test_vectorize_let(extent, target):
class Before:
@T.prim_func
def main(A: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
for i in T.vectorized(extent):
v: T.let = A[i] + T.float32(1)
A[i] = v + T.float32(2)
@@ -220,12 +224,12 @@ def test_vectorize_let(extent, target):
class After:
@T.prim_func
def main(A: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
v: T.let = A[T.Ramp(0, 1, extent)] + T.Broadcast(T.float32(1),
extent)
A[T.Ramp(0, 1, extent)] = v + T.Broadcast(T.float32(2), extent)
- with tvm.target.Target(target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() *
4, sve_target)])
@@ -234,15 +238,15 @@ def test_vectorize_with_le_cond(extent, target):
class Module:
@T.prim_func
def main(A: T.Tensor((16,), "float32"), n: T.int32):
+ T.func_attr({"target": target})
for i in T.vectorized(extent):
if i <= n:
A[i] = A[i] + T.float32(1)
- with tvm.target.Target(target):
- stmt = tvm.tirx.transform.VectorizeLoop()(Module)["main"].body
+ stmt = tvm.tirx.transform.VectorizeLoop()(Module)["main"].body
- # Check that the loop wasn't vectorised
- assert isinstance(stmt, tvm.tirx.For)
+ # Check that the loop wasn't vectorised
+ assert isinstance(stmt, tvm.tirx.For)
@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() *
4, sve_target)])
@@ -251,15 +255,15 @@ def test_vectorize_with_ge_cond(extent, target):
class Module:
@T.prim_func
def main(A: T.Tensor((16,), "float32"), n: T.int32):
+ T.func_attr({"target": target})
for i in T.vectorized(extent):
if i >= n:
A[i] = A[i] + T.float32(1)
- with tvm.target.Target(target):
- stmt = tvm.tirx.transform.VectorizeLoop()(Module)["main"].body
+ stmt = tvm.tirx.transform.VectorizeLoop()(Module)["main"].body
- # Check that the loop wasn't vectorised
- assert isinstance(stmt, tvm.tirx.For)
+ # Check that the loop wasn't vectorised
+ assert isinstance(stmt, tvm.tirx.For)
@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() *
4, sve_target)])
@@ -268,6 +272,7 @@ def test_vectorize_if_then_else_scalarize(extent, target):
class Before:
@T.prim_func
def main(A: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
for i in T.vectorized(extent):
A[i] = T.if_then_else(i > 0, A[i] + T.float32(1), A[i])
@@ -275,12 +280,12 @@ def test_vectorize_if_then_else_scalarize(extent, target):
class After:
@T.prim_func
def main(A: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
for i_s in range(extent):
A[i_s] = T.if_then_else(i_s > 0, A[i_s] + T.float32(1), A[i_s])
- with tvm.target.Target(target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() *
4, sve_target)])
@@ -289,6 +294,7 @@ def test_vectorize_if_then_else_vector(extent, target):
class Before:
@T.prim_func
def main(A: T.Tensor((25,), "float32"), n: T.int32):
+ T.func_attr({"target": target})
for i in range(n):
for j in T.vectorized(extent):
A[i * extent + j] = T.if_then_else(i > 0, A[i * extent +
j], 0)
@@ -297,14 +303,14 @@ def test_vectorize_if_then_else_vector(extent, target):
class After:
@T.prim_func
def main(A: T.Tensor((25,), "float32"), n: T.int32):
+ T.func_attr({"target": target})
for i in range(n):
A[T.Ramp(i * extent, 1, extent)] = T.if_then_else(
i > 0, A[T.Ramp(i * extent, 1, extent)], T.Broadcast(0,
extent)
)
- with tvm.target.Target(target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
def test_vectorize_let_if_then_else():
@@ -312,6 +318,7 @@ def test_vectorize_let_if_then_else():
class Before:
@T.prim_func
def main():
+ T.func_attr({"target": simple_target})
for i in T.vectorized(4):
if i < 2:
result: T.let[T.int32] = T.if_then_else(i < 1, 1, 2)
@@ -320,14 +327,14 @@ def test_vectorize_let_if_then_else():
class After:
@T.prim_func
def main():
+ T.func_attr({"target": simple_target})
for i_s in range(4):
if i_s < 2:
result: T.let[T.int32] = T.if_then_else(i_s < 1, 1, 2)
T.evaluate(0)
- with tvm.target.Target(simple_target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
def test_vectorize_while_fail():
@@ -371,6 +378,7 @@ def test_vectorize_with_reinterpret(extent, vec_str,
target):
class Before:
@T.prim_func
def main(A: T.Tensor((16,), "int32"), B: T.Tensor((16,), "float32")):
+ T.func_attr({"target": target})
for i in T.vectorized(0, extent):
B[i] = T.reinterpret("float32", A[i])
@@ -378,11 +386,11 @@ def test_vectorize_with_reinterpret(extent, vec_str,
target):
class After:
@T.prim_func
def main(A: T.Tensor((16,), "int32"), B: T.Tensor((16,), "float32")):
+ T.func_attr({"target": target})
B[T.Ramp(0, 1, extent)] = T.reinterpret(vec_str, A[T.Ramp(0, 1,
extent)])
- with tvm.target.Target(target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() *
4, sve_target)])
@@ -411,6 +419,7 @@ def test_vectorize_binary(op, extent, target):
class Before:
@T.prim_func
def main(A: T.Tensor((25,), "float32"), B: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
for j in T.vectorized(extent):
A[j] = op(T.float32(3), B[j])
@@ -418,11 +427,11 @@ def test_vectorize_binary(op, extent, target):
class After:
@T.prim_func
def main(A: T.Tensor((25,), "float32"), B: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
A[T.Ramp(0, 1, extent)] = op(T.Broadcast(T.float32(3), extent),
B[T.Ramp(0, 1, extent)])
- with tvm.target.Target(target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() *
4, sve_target)])
@@ -432,6 +441,7 @@ def test_vectorize_logical(op, extent, target):
class Before:
@T.prim_func
def main(A: T.Tensor((25,), "bool"), B: T.Tensor((25,), "bool")):
+ T.func_attr({"target": target})
for j in T.vectorized(extent):
A[j] = op(T.bool(1), B[j])
@@ -439,11 +449,11 @@ def test_vectorize_logical(op, extent, target):
class After:
@T.prim_func
def main(A: T.Tensor((25,), "bool"), B: T.Tensor((25,), "bool")):
+ T.func_attr({"target": target})
A[T.Ramp(0, 1, extent)] = op(T.Broadcast(T.bool(1), extent),
B[T.Ramp(0, 1, extent)])
- with tvm.target.Target(target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
@pytest.mark.parametrize("extent, target", [(4, simple_target), (T.vscale() *
4, sve_target)])
@@ -452,6 +462,7 @@ def test_vectorize_select(extent, target):
class Before:
@T.prim_func
def main(A: T.Tensor((25,), "float32"), B: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
for j in T.vectorized(extent):
A[j] = T.Select(T.bool(True), A[j], B[j])
@@ -459,15 +470,15 @@ def test_vectorize_select(extent, target):
class After:
@T.prim_func
def main(A: T.Tensor((25,), "float32"), B: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
A[T.Ramp(0, 1, extent)] = T.Select(
T.Broadcast(T.bool(True), extent),
A[T.Ramp(0, 1, extent)],
B[T.Ramp(0, 1, extent)],
)
- with tvm.target.Target(target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
@pytest.mark.parametrize(
@@ -479,6 +490,7 @@ def test_vectorize_cast(extent, vec_str, target):
class Before:
@T.prim_func
def main(A: T.Tensor((25,), "int32"), B: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
for j in T.vectorized(extent):
A[j] = T.Cast("int32", B[j])
@@ -486,11 +498,11 @@ def test_vectorize_cast(extent, vec_str, target):
class After:
@T.prim_func
def main(A: T.Tensor((25,), "int32"), B: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
A[T.Ramp(0, 1, extent)] = T.Cast(vec_str, B[T.Ramp(0, 1, extent)])
- with tvm.target.Target(target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
def test_illegal_extent():
@@ -513,13 +525,13 @@ def test_illegal_vscale_in_non_sve_compilation():
class Mod:
@T.prim_func
def main(A: T.Tensor((16,), "float32")):
+ T.func_attr({"target": simple_target})
for j in T.vectorized(0, 4 * T.vscale()):
A[j] = 13
msg = "Failed to vectorize loop with extent T.vscale\\(\\) \\* 4 for
target"
- with tvm.target.Target(simple_target):
- with pytest.raises(tvm.error.InternalError, match=msg):
- tvm.tirx.transform.VectorizeLoop()(Mod)
+ with pytest.raises(tvm.error.InternalError, match=msg):
+ tvm.tirx.transform.VectorizeLoop()(Mod)
def test_vectorize_and_predicate_all_buffer_loads_stores():
@@ -633,6 +645,7 @@ def test_vectorize_nested_predicates_preserve_both_masks():
@T.prim_func
def before(A: T.Tensor((16,), "float32"), B: T.Tensor((16,), "float32")):
+ T.func_attr({"target": rvv_target})
for i_0 in T.serial(4):
for i_1 in T.vectorized(4):
if i_0 * 4 + i_1 < 15:
@@ -640,8 +653,7 @@ def test_vectorize_nested_predicates_preserve_both_masks():
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"]
+ after =
tvm.tirx.transform.VectorizeLoop()(tvm.IRModule.from_expr(before))["before"]
predicates = []
@@ -693,7 +705,7 @@ def
test_vectorize_with_explicitly_disabled_buffer_level_predication():
# option, so no buffer-level predicates should be added.
@T.prim_func
def before(A: T.Tensor((16,), "float32"), B: T.Tensor((16,), "float32")):
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+ T.func_attr({"target": sve_target, "global_symbol": "main",
"tirx.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:
@@ -701,15 +713,14 @@ def
test_vectorize_with_explicitly_disabled_buffer_level_predication():
@T.prim_func
def expected(A: T.Tensor((16,), "float32"), B: T.Tensor((16,), "float32")):
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+ T.func_attr({"target": sve_target, "global_symbol": "main",
"tirx.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={"tirx.enable_buffer_level_predication":
False}):
- with tvm.target.Target(sve_target):
- after = tvm.tirx.transform.VectorizeLoop()(mod)["main"]
+ after = tvm.tirx.transform.VectorizeLoop()(mod)["main"]
tvm.ir.assert_structural_equal(after, expected)
@@ -751,46 +762,6 @@ def
test_vectorize_and_predicate_buffer_load_stores_with_sve_func_attr_target():
tvm.ir.assert_structural_equal(after, expected)
-def
test_vectorize_and_predicate_buffer_load_stores_with_sve_attr_scope_target():
- @T.prim_func
- def before(A: T.Tensor((16,), "float32"), B: T.Tensor((16,), "float32")):
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
- with T.attr(sve_target, "target", 0):
- 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.Tensor((16,), "float32"), B: T.Tensor((16,), "float32")):
- T.func_attr({"global_symbol": "main", "tirx.noalias": True})
- with T.attr(sve_target, "target", 0):
- for i_0 in range(4):
- T.evaluate(
- T.call_intrin(
- "void",
- "tirx.masked_store",
- B,
- (
- 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),
- )
- + T.Broadcast(T.float32(1), 4)
- ),
- T.Ramp(i_0 * 4, 1, 4),
- T.get_active_lane_mask("uint1x4", i_0 * 4, 14),
- )
- )
-
- mod = tvm.IRModule.from_expr(before)
- after = tvm.tirx.transform.VectorizeLoop()(mod)["main"]
- tvm.ir.assert_structural_equal(after, expected)
-
-
@pytest.mark.parametrize(
"extent, vec_str, target",
[(4, "float32x4", simple_target)],
@@ -800,6 +771,7 @@ def test_vectorize_llvm_pure_intrin(extent, vec_str,
target):
class Before:
@T.prim_func
def main(A: T.Tensor((25,), "float32"), B: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
for j in T.vectorized(extent):
A[j] = T.call_llvm_pure_intrin("float32", "llvm.sqrt", B[j])
@@ -807,14 +779,14 @@ def test_vectorize_llvm_pure_intrin(extent, vec_str,
target):
class After:
@T.prim_func
def main(A: T.Tensor((25,), "float32"), B: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
A[T.Ramp(0, 1, extent)] = T.call_llvm_pure_intrin(
vec_str, "llvm.sqrt", B[T.Ramp(0, 1, extent)]
)
- with tvm.target.Target(target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
- mod = tvm.compile(mod, target=target)
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
+ mod = tvm.compile(mod, target=target)
@pytest.mark.parametrize(
@@ -826,6 +798,7 @@ def test_vectorize_llvm_pure_intrin_fail(extent, vec_str,
target):
class Before:
@T.prim_func
def main(A: T.Tensor((25,), "int32"), B: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
for j in T.vectorized(extent):
A[j] = T.call_llvm_pure_intrin("int32", "llvm.lround", B[j])
@@ -833,21 +806,21 @@ def test_vectorize_llvm_pure_intrin_fail(extent, vec_str,
target):
class After:
@T.prim_func
def main(A: T.Tensor((25,), "int32"), B: T.Tensor((25,), "float32")):
+ T.func_attr({"target": target})
A[T.Ramp(0, 1, extent)] = T.call_llvm_pure_intrin(
vec_str, "llvm.lround", B[T.Ramp(0, 1, extent)]
)
- with tvm.target.Target(target):
- mod = tvm.tirx.transform.VectorizeLoop()(Before)
- tvm.ir.assert_structural_equal(mod, After)
- # LLVM 20 added vector support for llvm.lround/llvm.llround. The IR
Verifier's
- # "Intrinsic does not support vectors" check was removed in
release/20.x, so
- # compilation only fails on LLVM <= 19.
- if llvm_version_major() >= 20:
+ mod = tvm.tirx.transform.VectorizeLoop()(Before)
+ tvm.ir.assert_structural_equal(mod, After)
+ # LLVM 20 added vector support for llvm.lround/llvm.llround. The IR
Verifier's
+ # "Intrinsic does not support vectors" check was removed in release/20.x,
so
+ # compilation only fails on LLVM <= 19.
+ if llvm_version_major() >= 20:
+ tvm.compile(mod, target=target)
+ else:
+ with pytest.raises(Exception, match="Intrinsic does not support
vectors"):
tvm.compile(mod, target=target)
- else:
- with pytest.raises(Exception, match="Intrinsic does not support
vectors"):
- tvm.compile(mod, target=target)
if __name__ == "__main__":
