| 331 | @tvm.testing.parametrize_targets("cuda", "rocm") |
| 332 | def test_crossthread_reduction1(target, dev): |
| 333 | def sched(nthd): |
| 334 | @I.ir_module(s_tir=True) |
| 335 | class Module: |
| 336 | @T.prim_func(s_tir=True) |
| 337 | def main(var_A: T.handle, var_B: T.handle): |
| 338 | T.func_attr({"tirx.noalias": True}) |
| 339 | n, m = T.int32(), T.int32() |
| 340 | A = T.match_buffer(var_A, (n, m)) |
| 341 | B = T.match_buffer(var_B, (n,)) |
| 342 | for i in T.thread_binding(n, thread="blockIdx.x"): |
| 343 | for m_0 in T.thread_binding(nthd, thread="threadIdx.x"): |
| 344 | for m_1 in range((m + nthd - 1) // nthd): |
| 345 | with T.sblock("B"): |
| 346 | v_i = T.axis.spatial(n, i) |
| 347 | v_m = T.axis.reduce(m, m_0 * ((m + nthd - 1) // nthd) + m_1) |
| 348 | T.where(m_0 * ((m + nthd - 1) // nthd) + m_1 < m) |
| 349 | T.reads(A[v_i, v_m]) |
| 350 | T.writes(B[v_i]) |
| 351 | with T.init(): |
| 352 | B[v_i] = T.float32(0.0) |
| 353 | B[v_i] = B[v_i] + A[v_i, v_m] |
| 354 | |
| 355 | fun = tvm.compile(Module, target="cuda") |
| 356 | return fun |
| 357 | |
| 358 | def verify(nthd): |
| 359 | func = sched(nthd) |