tvm_tirx.tirx.transform.LowerThreadAllreduce

tvm_tirx.tirx.transform.LowerThreadAllreduce()

Lower gpu_thread_allreduce to target-specific communication.

Requires a function target attribute and explicit thread extents. Selects warp shuffles or shared-memory reductions as supported.

Returns:

fpass (tvm_tirx.ir.transform.Pass) – The result pass.

Examples

A single participating thread reduces to a direct store.

import tvm_tirx as tvm
from tvm_tirx.script import ir as I, tirx as T
from tvm_tirx.tirx import transform as tr

@I.ir_module
class Before:
    @T.function(private=True)
    def main(A: T.Tensor((1,), T.float32)):
        T.func_attr({"target": T.target({"kind": "cuda", "arch": "sm_75"})})
        R = T.alloc_tensor(ty_args=[T.Tensor((1,), T.float32, scope="local")])
        with T.launch_thread("threadIdx.x", 1) as tx:
            T.gpu_thread_allreduce(
                T.Lambda([T.float32, T.float32], lambda arg0, arg1: (arg0 + arg1,)),
                (T.float32(0.0),),
                (A[tx],),
                T.bool(True),
                (R[0],),
                (tx,),
            )

@I.ir_module
class Expected:
    @T.function(private=True)
    def main(A: T.Tensor((1,), T.float32)):
        T.func_attr({"target": T.target({"kind": "cuda", "arch": "sm_75"})})
        R = T.alloc_tensor(ty_args=[T.Tensor((1,), T.float32, scope="local")])
        with T.launch_thread("threadIdx.x", 1) as tx:
            R[0] = A[tx]

def run_pass():
    return tr.LowerThreadAllreduce()(Before)

tvm.ir.assert_structural_equal(run_pass(), Expected)