diff --git a/python/flydsl/expr/gpu.py b/python/flydsl/expr/gpu.py index 417253096..f0398202f 100644 --- a/python/flydsl/expr/gpu.py +++ b/python/flydsl/expr/gpu.py @@ -20,8 +20,9 @@ from .._mlir.dialects import gpu from .._mlir.dialects._fly_enum_gen import AddressSpace from ..compiler.protocol import dsl_align_of, dsl_size_of +from .math import dsl_math_wrap_result from .meta import dsl_loc_tracing -from .numeric import Numeric, Uint8 +from .numeric import Int32, Numeric, Uint8 from .primitive import get_dyn_shared, make_ptr from .struct import ( Arena, @@ -32,7 +33,7 @@ is_composite_type, is_struct_type, ) -from .typing import Array, PointerType, Tuple3D +from .typing import Array, PointerType, Tuple3D, as_ir_value @dsl_loc_tracing @@ -50,6 +51,45 @@ def barrier(*args, **kwargs): return gpu.barrier(*args, **kwargs) +@dsl_loc_tracing +@dsl_math_wrap_result +def shuffle(value, offset, width, mode="xor"): + """Move ``value`` across lanes of a subgroup (warp) via ``gpu.shuffle``. + + ``width`` is the number of participating lanes and must be uniform across + the subgroup. + """ + if mode not in ("xor", "up", "down", "idx"): + raise ValueError(f"invalid shuffle mode {mode!r}; expected one of (xor, up, down, idx)") + + return gpu.ShuffleOp( + as_ir_value(value), + Int32(offset).ir_value(), + Int32(width).ir_value(), + mode=mode, + ).shuffleResult + + +def shuffle_xor(value, offset, width): + """``shuffle`` in ``"xor"`` mode: lane ``k`` reads lane ``k ^ offset``.""" + return shuffle(value, offset, width, mode="xor") + + +def shuffle_up(value, offset, width): + """``shuffle`` in ``"up"`` mode: lane ``k`` reads lane ``k - offset``.""" + return shuffle(value, offset, width, mode="up") + + +def shuffle_down(value, offset, width): + """``shuffle`` in ``"down"`` mode: lane ``k`` reads lane ``k + offset``.""" + return shuffle(value, offset, width, mode="down") + + +def shuffle_idx(value, lane, width): + """``shuffle`` in ``"idx"`` mode: every lane reads lane ``lane``.""" + return shuffle(value, lane, width, mode="idx") + + thread_idx = Tuple3D(gpu.thread_id) block_idx = Tuple3D(gpu.block_id) block_dim = Tuple3D(gpu.block_dim) @@ -199,6 +239,10 @@ def _allocate_static_shared(self, nbytes: int, align: int): "block_dim", "grid_dim", "barrier", + "shuffle_xor", + "shuffle_up", + "shuffle_down", + "shuffle_idx", "smem_space", "lds_space", "SharedAllocator",