Repository navigation
[Bug][S-TIR] compute_at after rfactor drops the tail predicate: wrong sums and out-of-bounds reads #20812
Description
split of a reduction loop by a factor that does not divide its extent gives the block a predicate
(k_0 * 64 + k_1 < 77). rfactor keeps it on Y_rf, but compute_at(Y_rf, ...) then rebuilds the loops of Y_rf
without it. No error is raised, and the kernel reads past each row (and past the end of X).
import numpy as np, tvm
from tvm import te, s_tir
create = getattr(te, "create_prim_func", None) or te.create_function # renamed on main
R, K = 4, 77
X = te.placeholder((R, K), "float32", name="X")
k = te.reduce_axis((0, K), name="k")
Y = te.compute((R,), lambda r: te.sum(X[r, k], axis=k), name="Y")
sch = s_tir.Schedule(create([X, Y]))
blk = sch.get_sblock("Y")
_, k0 = sch.get_loops(blk)
sch.split(k0, factors=[None, 64]) # 2 * 64 > 77: predicate k_0 * 64 + k_1 < 77
rf = sch.rfactor(sch.get_loops(blk)[-1], 0) # Y_rf keeps the predicate
sch.compute_at(rf, sch.get_loops(blk)[-1]) # Y_rf loses it
f = tvm.compile(sch.mod, target="llvm")
x = (np.arange(R * K, dtype="float32") % 7).reshape(R, K)
y = tvm.runtime.tensor(np.zeros(R, "float32"))
f(tvm.runtime.tensor(x), y)
print("expected", x.sum(1))
print("got ", y.numpy())Output (apache-tvm 0.27.0.post1, LLVM target; the last value is memory past the buffer and varies):
expected [231. 231. 231. 231.]
got [3.79e+02 3.79e+02 3.79e+02 1.72e+32]
Same TIR (no T.where on Y_rf after compute_at) on main e345029; the CUDA kernel has the same unguarded read.
Probable cause: rfactor marks the new blocks as affine_binding / region_cover / stage_pipeline (reduction.cc,
"Update scope information"), and compute_at trusts these flags (ScopeReconstructor::MakeNewLoop, compute_at.cc).
With debug_mask="all" the schedule fails verification right after rfactor. Rejecting the schedule would be a fine fix.
Activity