Skip to content

[TIRx][Schedule] Support rfactor for arg reducers selecting last index #19909

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ZephyrLi-pro wants to merge 1 commit into
base: main
Choose a base branch
from
+190 −0
Open

[TIRx][Schedule] Support rfactor for arg reducers selecting last index #19909

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions src/s_tir/schedule/primitive/reduction.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -396,6 +396,36 @@ struct ReducerRegistry {
PrimExpr val = Select(less(x[1], y[1]), x[1], y[1]);
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], y[1]), And(equal(x[1], y[1]), greater(x[0], y[0]))),
x[0], y[0]);
PrimExpr val = Select(greater(x[1], y[1]), x[1], y[1]);
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], y[1]), And(equal(x[1], y[1]), greater(x[0], y[0]))),
x[0], y[0]);
PrimExpr val = Select(less(x[1], y[1]), x[1], y[1]);
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())};
Expand Down
160 changes: 160 additions & 0 deletions tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1329,6 +1329,136 @@ def argmin_topi_rfactor(
placeholder_red[ax0] = placeholder_red_temp_v0[ax0]


@T.prim_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]


@T.prim_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


Expand Down Expand Up @@ -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)
Expand Down
Loading