This is an automated email from the ASF dual-hosted git repository.
tlopex 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 ed932838d9 [TIRx][Schedule] Support rfactor for arg reducers selecting
last index (#19909)
ed932838d9 is described below
commit ed932838d9aa16a8d49b91414a9f4f7d68177837
Author: Zephyr <[email protected]>
AuthorDate: Fri Sep 11 04:34:27 2026 +0800
[TIRx][Schedule] Support rfactor for arg reducers selecting last index
(#19909)
This PR extends `rfactor` reducer pattern matching to support TOPI
`argmax`/`argmin` with `select_last_index=True`.
For `select_last_index=True`, ties should be resolved by choosing the
larger index. The added reducer patterns preserve the existing value
comparison semantics, then use `lhs_idx > rhs_idx` as the tie-break
condition.
Tests are added for applying `rfactor` to `topi.argmax` and
`topi.argmin` with `select_last_index=True`.
---------
Co-authored-by: lizhuoheng <[email protected]>
---
src/s_tir/schedule/primitive/reduction.cc | 38 +++++
.../s_tir/schedule/test_tir_schedule_rfactor.py | 160 +++++++++++++++++++++
2 files changed, 198 insertions(+)
diff --git a/src/s_tir/schedule/primitive/reduction.cc
b/src/s_tir/schedule/primitive/reduction.cc
index a96face461..75534ce1a6 100644
--- a/src/s_tir/schedule/primitive/reduction.cc
+++ b/src/s_tir/schedule/primitive/reduction.cc
@@ -432,6 +432,44 @@ struct ReducerRegistry {
x[1].as_or_throw<PrimExpr>(),
y[1].as_or_throw<PrimExpr>());
return ffi::Array<PrimExpr>{idx, val};
},
+ [](const ffi::Array<PrimExpr>& values) {
+ return ffi::Array<PrimExpr>{MakeConst(values[0].ty(), -1),
+ max_value(values[1].ty())};
+ }),
+ // argmax with `lhs_val > rhs_val` and tie-break `lhs_idx >
rhs_idx`, which corresponds
+ // to topi.argmax with `select_last_index=True` (preferring the
last occurrence).
+ CreateReducerGetter(
+ /*n_buffers=*/2,
+ [](const ffi::Array<Var>& x, const ffi::Array<Var>& y) {
+ PrimExpr idx = Select(
+ Or(greater(x[1].as_or_throw<PrimExpr>(),
y[1].as_or_throw<PrimExpr>()),
+ And(equal(x[1].as_or_throw<PrimExpr>(),
y[1].as_or_throw<PrimExpr>()),
+ greater(x[0].as_or_throw<PrimExpr>(),
y[0].as_or_throw<PrimExpr>()))),
+ x[0].as_or_throw<PrimExpr>(),
y[0].as_or_throw<PrimExpr>());
+ PrimExpr val =
+ Select(greater(x[1].as_or_throw<PrimExpr>(),
y[1].as_or_throw<PrimExpr>()),
+ x[1].as_or_throw<PrimExpr>(),
y[1].as_or_throw<PrimExpr>());
+ return ffi::Array<PrimExpr>{idx, val};
+ },
+ [](const ffi::Array<PrimExpr>& values) {
+ return ffi::Array<PrimExpr>{MakeConst(values[0].ty(), -1),
+ min_value(values[1].ty())};
+ }),
+ // argmin with `lhs_val < rhs_val` and tie-break `lhs_idx >
rhs_idx`, which corresponds
+ // to topi.argmin with `select_last_index=True` (preferring the
last occurrence).
+ CreateReducerGetter(
+ /*n_buffers=*/2,
+ [](const ffi::Array<Var>& x, const ffi::Array<Var>& y) {
+ PrimExpr idx = Select(
+ Or(less(x[1].as_or_throw<PrimExpr>(),
y[1].as_or_throw<PrimExpr>()),
+ And(equal(x[1].as_or_throw<PrimExpr>(),
y[1].as_or_throw<PrimExpr>()),
+ greater(x[0].as_or_throw<PrimExpr>(),
y[0].as_or_throw<PrimExpr>()))),
+ x[0].as_or_throw<PrimExpr>(),
y[0].as_or_throw<PrimExpr>());
+ PrimExpr val =
+ Select(less(x[1].as_or_throw<PrimExpr>(),
y[1].as_or_throw<PrimExpr>()),
+ x[1].as_or_throw<PrimExpr>(),
y[1].as_or_throw<PrimExpr>());
+ return ffi::Array<PrimExpr>{idx, val};
+ },
[](const ffi::Array<PrimExpr>& values) {
return ffi::Array<PrimExpr>{MakeConst(values[0].ty(), -1),
max_value(values[1].ty())};
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
b/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
index 306f4ed12d..c0d9c189c0 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
@@ -1329,6 +1329,136 @@ def argmin_topi_rfactor(
placeholder_red[ax0] = placeholder_red_temp_v0[ax0]
[email protected]_func(s_tir=True)
+def argmax_topi_select_last_rfactor(
+ placeholder: T.Buffer((1, 32), "int32"), placeholder_red: T.Buffer(1,
"int32")
+) -> None:
+ T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+ placeholder_red_temp_v0 = T.sblock_alloc_buffer([1], dtype="int32")
+ placeholder_red_temp_v1 = T.sblock_alloc_buffer([1], dtype="int32")
+ placeholder_red_temp_v0_rf = T.sblock_alloc_buffer([1, 8], dtype="int32")
+ placeholder_red_temp_v1_rf = T.sblock_alloc_buffer([1, 8], dtype="int32")
+ for i0, i1_0, i1_1 in T.grid(1, 4, 8):
+ with T.sblock("placeholder_red_temp_rf"):
+ vi1_1, ax0, vi1_0 = T.axis.remap("SSR", [i1_1, i0, i1_0])
+ T.reads(placeholder[ax0, vi1_0 * 8 + vi1_1])
+ T.writes(placeholder_red_temp_v0_rf[ax0, vi1_1],
placeholder_red_temp_v1_rf[ax0, vi1_1])
+ with T.init():
+ placeholder_red_temp_v0_rf[ax0, vi1_1] = -1
+ placeholder_red_temp_v1_rf[ax0, vi1_1] = -2147483648
+ v_placeholder_red_temp_v0_rf: T.let[T.int32] = T.Select(
+ placeholder_red_temp_v1_rf[ax0, vi1_1] > placeholder[ax0,
vi1_0 * 8 + vi1_1]
+ or (
+ placeholder_red_temp_v1_rf[ax0, vi1_1] == placeholder[ax0,
vi1_0 * 8 + vi1_1]
+ and placeholder_red_temp_v0_rf[ax0, vi1_1] > vi1_0 * 8 +
vi1_1
+ ),
+ placeholder_red_temp_v0_rf[ax0, vi1_1],
+ vi1_0 * 8 + vi1_1,
+ )
+ v_placeholder_red_temp_v1_rf: T.let[T.int32] = T.Select(
+ placeholder_red_temp_v1_rf[ax0, vi1_1] > placeholder[ax0,
vi1_0 * 8 + vi1_1],
+ placeholder_red_temp_v1_rf[ax0, vi1_1],
+ placeholder[ax0, vi1_0 * 8 + vi1_1],
+ )
+ placeholder_red_temp_v0_rf[ax0, vi1_1] =
v_placeholder_red_temp_v0_rf
+ placeholder_red_temp_v1_rf[ax0, vi1_1] =
v_placeholder_red_temp_v1_rf
+ for i0, i1_1 in T.grid(1, 8):
+ with T.sblock("placeholder_red_temp"):
+ vi1_1, ax0 = T.axis.remap("RS", [i1_1, i0])
+ T.reads(placeholder_red_temp_v0_rf[ax0, vi1_1],
placeholder_red_temp_v1_rf[ax0, vi1_1])
+ T.writes(placeholder_red_temp_v0[ax0],
placeholder_red_temp_v1[ax0])
+ with T.init():
+ placeholder_red_temp_v0[ax0] = -1
+ placeholder_red_temp_v1[ax0] = -2147483648
+ v_placeholder_red_temp_v0: T.let[T.int32] = T.Select(
+ placeholder_red_temp_v1[ax0] > placeholder_red_temp_v1_rf[ax0,
vi1_1]
+ or (
+ placeholder_red_temp_v1[ax0] ==
placeholder_red_temp_v1_rf[ax0, vi1_1]
+ and placeholder_red_temp_v0[ax0] >
placeholder_red_temp_v0_rf[ax0, vi1_1]
+ ),
+ placeholder_red_temp_v0[ax0],
+ placeholder_red_temp_v0_rf[ax0, vi1_1],
+ )
+ v_placeholder_red_temp_v1: T.let[T.int32] = T.Select(
+ placeholder_red_temp_v1[ax0] > placeholder_red_temp_v1_rf[ax0,
vi1_1],
+ placeholder_red_temp_v1[ax0],
+ placeholder_red_temp_v1_rf[ax0, vi1_1],
+ )
+ placeholder_red_temp_v0[ax0] = v_placeholder_red_temp_v0
+ placeholder_red_temp_v1[ax0] = v_placeholder_red_temp_v1
+ for i0 in T.serial(1):
+ with T.sblock("placeholder_red"):
+ ax0 = T.axis.spatial(1, i0)
+ T.reads(placeholder_red_temp_v0[ax0])
+ T.writes(placeholder_red[ax0])
+ placeholder_red[ax0] = placeholder_red_temp_v0[ax0]
+
+
[email protected]_func(s_tir=True)
+def argmin_topi_select_last_rfactor(
+ placeholder: T.Buffer((1, 32), "int32"), placeholder_red: T.Buffer(1,
"int32")
+) -> None:
+ T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+ placeholder_red_temp_v0 = T.sblock_alloc_buffer([1], dtype="int32")
+ placeholder_red_temp_v1 = T.sblock_alloc_buffer([1], dtype="int32")
+ placeholder_red_temp_v0_rf = T.sblock_alloc_buffer([1, 8], dtype="int32")
+ placeholder_red_temp_v1_rf = T.sblock_alloc_buffer([1, 8], dtype="int32")
+ for i0, i1_0, i1_1 in T.grid(1, 4, 8):
+ with T.sblock("placeholder_red_temp_rf"):
+ vi1_1, ax0, vi1_0 = T.axis.remap("SSR", [i1_1, i0, i1_0])
+ T.reads(placeholder[ax0, vi1_0 * 8 + vi1_1])
+ T.writes(placeholder_red_temp_v0_rf[ax0, vi1_1],
placeholder_red_temp_v1_rf[ax0, vi1_1])
+ with T.init():
+ placeholder_red_temp_v0_rf[ax0, vi1_1] = -1
+ placeholder_red_temp_v1_rf[ax0, vi1_1] = 2147483647
+ v_placeholder_red_temp_v0_rf: T.let[T.int32] = T.Select(
+ placeholder_red_temp_v1_rf[ax0, vi1_1] < placeholder[ax0,
vi1_0 * 8 + vi1_1]
+ or (
+ placeholder_red_temp_v1_rf[ax0, vi1_1] == placeholder[ax0,
vi1_0 * 8 + vi1_1]
+ and placeholder_red_temp_v0_rf[ax0, vi1_1] > vi1_0 * 8 +
vi1_1
+ ),
+ placeholder_red_temp_v0_rf[ax0, vi1_1],
+ vi1_0 * 8 + vi1_1,
+ )
+ v_placeholder_red_temp_v1_rf: T.let[T.int32] = T.Select(
+ placeholder_red_temp_v1_rf[ax0, vi1_1] < placeholder[ax0,
vi1_0 * 8 + vi1_1],
+ placeholder_red_temp_v1_rf[ax0, vi1_1],
+ placeholder[ax0, vi1_0 * 8 + vi1_1],
+ )
+ placeholder_red_temp_v0_rf[ax0, vi1_1] =
v_placeholder_red_temp_v0_rf
+ placeholder_red_temp_v1_rf[ax0, vi1_1] =
v_placeholder_red_temp_v1_rf
+ for i0, i1_1 in T.grid(1, 8):
+ with T.sblock("placeholder_red_temp"):
+ vi1_1, ax0 = T.axis.remap("RS", [i1_1, i0])
+ T.reads(placeholder_red_temp_v0_rf[ax0, vi1_1],
placeholder_red_temp_v1_rf[ax0, vi1_1])
+ T.writes(placeholder_red_temp_v0[ax0],
placeholder_red_temp_v1[ax0])
+ with T.init():
+ placeholder_red_temp_v0[ax0] = -1
+ placeholder_red_temp_v1[ax0] = 2147483647
+ v_placeholder_red_temp_v0: T.let[T.int32] = T.Select(
+ placeholder_red_temp_v1[ax0] < placeholder_red_temp_v1_rf[ax0,
vi1_1]
+ or (
+ placeholder_red_temp_v1[ax0] ==
placeholder_red_temp_v1_rf[ax0, vi1_1]
+ and placeholder_red_temp_v0[ax0] >
placeholder_red_temp_v0_rf[ax0, vi1_1]
+ ),
+ placeholder_red_temp_v0[ax0],
+ placeholder_red_temp_v0_rf[ax0, vi1_1],
+ )
+ v_placeholder_red_temp_v1: T.let[T.int32] = T.Select(
+ placeholder_red_temp_v1[ax0] < placeholder_red_temp_v1_rf[ax0,
vi1_1],
+ placeholder_red_temp_v1[ax0],
+ placeholder_red_temp_v1_rf[ax0, vi1_1],
+ )
+ placeholder_red_temp_v0[ax0] = v_placeholder_red_temp_v0
+ placeholder_red_temp_v1[ax0] = v_placeholder_red_temp_v1
+ for i0 in T.serial(1):
+ with T.sblock("placeholder_red"):
+ ax0 = T.axis.spatial(1, i0)
+ T.reads(placeholder_red_temp_v0[ax0])
+ T.writes(placeholder_red[ax0])
+ placeholder_red[ax0] = placeholder_red_temp_v0[ax0]
+
+
# pylint: enable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
@@ -1719,6 +1849,36 @@ def test_reduction_rfactor_topi_argmin():
verify_trace_roundtrip(s, mod=argmin_topi)
+def test_reduction_rfactor_topi_argmax_select_last_index():
+ A = te.placeholder((1, 32), dtype="int32")
+ B = topi.argmax(A, axis=1, select_last_index=True)
+ argmax_topi = te.create_prim_func([A, B])
+ s = tvm.s_tir.Schedule(argmax_topi, debug_mask="all")
+ argmax = s.get_sblock("placeholder_red_temp")
+ _, k = s.get_loops(argmax)
+ _, ki = s.split(k, [None, 8])
+ rf_block = s.rfactor(ki, 1)
+ assert_structural_equal_ignore_global_symbol(s.mod["main"],
argmax_topi_select_last_rfactor)
+ assert
s.get(rf_block).same_as(s.get(s.get_sblock("placeholder_red_temp_rf")))
+ assert s.get(argmax).same_as(s.get(s.get_sblock("placeholder_red_temp")))
+ verify_trace_roundtrip(s, mod=argmax_topi)
+
+
+def test_reduction_rfactor_topi_argmin_select_last_index():
+ A = te.placeholder((1, 32), dtype="int32")
+ B = topi.argmin(A, axis=1, select_last_index=True)
+ argmin_topi = te.create_prim_func([A, B])
+ s = tvm.s_tir.Schedule(argmin_topi, debug_mask="all")
+ argmin = s.get_sblock("placeholder_red_temp")
+ _, k = s.get_loops(argmin)
+ _, ki = s.split(k, [None, 8])
+ rf_block = s.rfactor(ki, 1)
+ assert_structural_equal_ignore_global_symbol(s.mod["main"],
argmin_topi_select_last_rfactor)
+ assert
s.get(rf_block).same_as(s.get(s.get_sblock("placeholder_red_temp_rf")))
+ assert s.get(argmin).same_as(s.get(s.get_sblock("placeholder_red_temp")))
+ verify_trace_roundtrip(s, mod=argmin_topi)
+
+
def test_reduction_rfactor_int64():
# fmt: off
@T.prim_func(s_tir=True)
