diff --git a/extra/gemm_fragment.py b/extra/gemm_fragment.py index c1d5900692..efa229c3e6 100644 --- a/extra/gemm_fragment.py +++ b/extra/gemm_fragment.py @@ -2,10 +2,12 @@ tilelang-style matmul_relu written with tinygrad UOp APIs. Demonstrates that tilelang's T.alloc_fragment is expressible with existing -tinygrad primitives: a per-thread REG buffer, wrapped in Ops.UNSHARD over a -LOCAL thread range to form the full logical tile. The kernel is written -against the full-tile UNSHARD view, and multi_pm (the same pass that lowers -multi-device UNSHARDs) resolves it into per-thread shard code. +tinygrad primitives: a per-thread REG buffer, wrapped in one Ops.UNSHARD per +sharded axis over the LOCAL thread-grid ranges to form the full logical tile. +Here the 64 threads are an 8x8 grid and each thread owns an 8x8 sub-tile -- +the 2-D fragment layout tilelang infers. The kernel is written against the +full-tile UNSHARD view, and multi_pm (the same pass that lowers multi-device +UNSHARDs) resolves it into per-thread shard code. Reference tilelang kernel: @@ -30,14 +32,15 @@ Reference tilelang kernel: API mapping (tilelang -> tinygrad UOps, idioms from test/backend/test_custom_kernel.py): - T.Kernel(gx, gy, threads=T) -> AxisType.GLOBAL ranges (blocks) + AxisType.LOCAL range (threads) + T.Kernel(gx, gy, threads=T) -> AxisType.GLOBAL ranges (blocks) + AxisType.LOCAL ranges (thread grid) T.alloc_shared(shape, dtype) -> UOp.placeholder(shape, dtype, slot, AddrSpace.LOCAL) - T.alloc_fragment(shape, dt) -> per-thread REG placeholder, wrapped in Ops.UNSHARD over the AxisType.LOCAL - range: fragment.unshard(axis, tid). The full logical tile is shard x threads, - exactly like device sharding, but the sharding axis is the thread axis - carried by the RANGE metadata instead of a device tuple. C_local[i, j] with - i in this thread's shard is INDEX on the UNSHARD, which multi_pm resolves - into INDEX on the per-thread REG shard. + T.alloc_fragment(shape, dt) -> per-thread REG placeholder, wrapped in one Ops.UNSHARD per sharded axis over + the AxisType.LOCAL ranges: fragment.unshard((axis_y, axis_x), (ty, tx)). + The full logical tile is the shard with each sharded axis multiplied by its + range size, exactly like device sharding, but the sharding axes are thread + axes carried by the RANGE metadata instead of a device tuple. C_local[i, j] + with [i, j] in this thread's shard is INDEX on the UNSHARD, which multi_pm + resolves into INDEX on the per-thread REG shard, axis by axis. T.copy(gmem_slice, smem) -> smem[thread_idx].set(gmem_slice[thread_idx], end=copy_rng). set returns the smem tile AFTER the copy; the implicit-barrier pass turns the store->load dependency of the loop that consumes it into a workgroup barrier @@ -51,7 +54,7 @@ API mapping (tilelang -> tinygrad UOps, idioms from test/backend/test_custom_ker from tinygrad.dtype import dtypes, AddrSpace, DType from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo -from tinygrad.helpers import cdiv +from tinygrad.helpers import cdiv, getenv from tinygrad.tensor import Tensor # --------------------------------------------------------------------------- @@ -62,35 +65,37 @@ def alloc_shared(shape:tuple[int, ...], dtype:DType) -> UOp: """T.alloc_shared: one LOCAL buffer shared by all threads in the block.""" return UOp.placeholder(tuple(shape), dtype, next(UOp.unique_num), AddrSpace.LOCAL) -def alloc_fragment(shape:tuple[int, ...], dtype:DType, axis:int, tnum:UOp) -> UOp: - """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL axis. +def alloc_fragment(shape:tuple[int, ...], dtype:DType, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp: + """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid. - Each thread privately owns shape[axis]//threads elements along `axis` in a - REG buffer. The UNSHARD over the LOCAL thread range presents the full - logical tile: full_shape = shard_shape with `axis` multiplied by the number - of threads. This is exactly how UNSHARD carries the DEVICE axis today, - except the sharding axis is the thread axis carried by the RANGE metadata. + Each thread privately owns shape[axis]//threads elements along every sharded + axis in a REG buffer. The UNSHARDs over the LOCAL thread ranges present the + full logical tile: full_shape = shard_shape with each sharded axis multiplied + by its range size. This is exactly how UNSHARD carries a DEVICE axis today, + except the sharding axes are thread axes carried by the RANGE metadata. """ - nthreads = int(tnum.vmax) + 1 - assert tnum.op is Ops.RANGE and tnum.arg[-1] is AxisType.LOCAL, "fragment must shard over a LOCAL range" - assert shape[axis] % nthreads == 0 - shard_shape = tuple(s // nthreads if i == axis else s for i, s in enumerate(shape)) + assert len(axes) == len(rngs) + assert all(tnum.op is Ops.RANGE and tnum.arg[-1] is AxisType.LOCAL for tnum in rngs), "fragments shard over LOCAL ranges" + assert all(shape[a] % (int(rng.vmax)+1) == 0 for a, rng in zip(axes, rngs)) + by_axis = dict(zip(axes, rngs)) + shard_shape = tuple(s // (int(by_axis[i].vmax)+1) if i in by_axis else s for i, s in enumerate(shape)) fragment = UOp.placeholder(shard_shape, dtype, next(UOp.unique_num), AddrSpace.REG) - return fragment.unshard(axis, tnum) + return fragment.unshard(axes, rngs) # --------------------------------------------------------------------------- # GEMM kernel: C = relu(A @ B), float inputs (fp16 or fp32), fp32 fragment accumulator, no WMMA # --------------------------------------------------------------------------- -# 64x64 output tile per block, 64 threads (tilelang uses 128 with a 2D per-thread fragment layout; -# an UNSHARD splits one axis, so here each thread owns BLOCK_M//THREADS full tile rows of BLOCK_N fp32) +# 64x64 output tile per block, 64 threads as an 8x8 grid; each thread owns an 8x8 fragment sub-tile +# (the 2-D per-thread layout tilelang infers for this GEMM) BLOCK_M = BLOCK_N = BLOCK_K = 64 -THREADS = 64 -ROWS = BLOCK_M // THREADS # fragment rows per thread -ROWS_B = BLOCK_N // THREADS # B tile columns copied per thread +TY = TX = 8 +THREADS = TY * TX +TM = BLOCK_M // TY # fragment rows per thread +TN = BLOCK_N // TX # fragment columns per thread def matmul_relu_kernel(c:UOp, a:UOp, b:UOp) -> UOp: - """C[M, N] = relu(A[M, K] @ B[K, N]) -- one 64x64 tile per block, locals + fragments.""" + """C[M, N] = relu(A[M, K] @ B[K, N]) -- one 64x64 tile per block, locals + a 2-D fragment.""" M, K = a.shape K2, N = b.shape assert K == K2 and a.dtype == b.dtype == c.dtype and not dtypes.is_int(a.dtype) @@ -99,53 +104,56 @@ def matmul_relu_kernel(c:UOp, a:UOp, b:UOp) -> UOp: # with T.Kernel(T.ceildiv(N, BLOCK_N), T.ceildiv(M, BLOCK_M), threads=64) as (bx, by): bx = UOp.range(cdiv(N, BLOCK_N), 0, AxisType.GLOBAL) by = UOp.range(cdiv(M, BLOCK_M), 1, AxisType.GLOBAL) - tid = UOp.range(THREADS, 2, AxisType.LOCAL) + ty = UOp.range(TY, 2, AxisType.LOCAL) + tx = UOp.range(TX, 3, AxisType.LOCAL) # A_shared = T.alloc_shared((BLOCK_M, BLOCK_K), dtype) # B_shared = T.alloc_shared((BLOCK_K, BLOCK_N), dtype) A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype) B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype) - # C_local = T.alloc_fragment((BLOCK_M, BLOCK_N), accum_dtype) - C_local = alloc_fragment((BLOCK_M, BLOCK_N), dtypes.float32, axis=0, tnum=tid) + # C_local = T.alloc_fragment((BLOCK_M, BLOCK_N), accum_dtype) -- an 8x8 REG tile per thread of the 8x8 grid + C_local = alloc_fragment((BLOCK_M, BLOCK_N), dtypes.float32, (0, 1), (ty, tx)) - # T.clear(C_local) -- each thread zeroes its own fragment rows - ir0, j0 = UOp.range(ROWS, 4, AxisType.LOOP), UOp.range(BLOCK_N, 5, AxisType.LOOP) - C_loc = C_local[tid*ROWS + ir0, j0].set(0.0, end=(ir0, j0)) + # T.clear(C_local) -- each thread zeroes its own fragment sub-tile + ic, jc = UOp.range(TM, 4, AxisType.LOOP), UOp.range(TN, 5, AxisType.LOOP) + C_loc = C_local[ty*TM + ic, tx*TN + jc].set(0.0, end=(ic, jc)) # for ko in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3): # (num_stages pipelining is async copy + multi-buffering; this is the synchronous single-buffer version) - ko = UOp.range(cdiv(K, BLOCK_K), 3, AxisType.LOOP) + ko = UOp.range(cdiv(K, BLOCK_K), 6, AxisType.LOOP) - # T.copy(A[by * BLOCK_M, ko * BLOCK_K], A_shared) -- each thread copies ROWS row(s) + # T.copy(A[by * BLOCK_M, ko * BLOCK_K], A_shared) -- each thread copies its own 8x8 sub-tile # set returns the tile AFTER the copy; codegen turns that store->load dependency into a workgroup barrier - iar, ka = UOp.range(ROWS, 6, AxisType.LOOP), UOp.range(BLOCK_K, 7, AxisType.LOOP) - A_shared = A_shared[tid*ROWS + iar, ka].set(a[by*BLOCK_M + tid*ROWS + iar, ko*BLOCK_K + ka], end=(iar, ka)) + iar, ka = UOp.range(TM, 7, AxisType.LOOP), UOp.range(TN, 8, AxisType.LOOP) + A_store = A_shared[ty*TM + iar, tx*TN + ka].store(a[by*BLOCK_M + ty*TM + iar, ko*BLOCK_K + tx*TN + ka]).end(iar, ka) - # T.copy(B[ko * BLOCK_K, bx * BLOCK_N], B_shared) -- each thread copies ROWS_B column(s) - kb, ibr = UOp.range(BLOCK_K, 8, AxisType.LOOP), UOp.range(ROWS_B, 9, AxisType.LOOP) - B_shared = B_shared[kb, tid*ROWS_B + ibr].set(b[ko*BLOCK_K + kb, bx*BLOCK_N + tid*ROWS_B + ibr], end=(kb, ibr)) + # T.copy(B[ko * BLOCK_K, bx * BLOCK_N], B_shared) + kb, ibr = UOp.range(TM, 9, AxisType.LOOP), UOp.range(TN, 10, AxisType.LOOP) + B_store = B_shared[ty*TM + kb, tx*TN + ibr].store(b[ko*BLOCK_K + ty*TM + kb, bx*BLOCK_N + tx*TN + ibr]).end(kb, ibr) - # T.gemm(A_shared, B_shared, C_local), no WMMA -- per-thread accumulate over its fragment rows. + # get the shared after the stores (single barrier) + A_shared = A_shared.after(A_store, B_store) + B_shared = B_shared.after(A_store, B_store) + + # T.gemm(A_shared, B_shared, C_local), no WMMA -- per-thread accumulate over its fragment sub-tile. # identical to custom_gemm: a self-referential store over the loop-carried kk range, # which codegen turns into a register accumulator - # kk nests outside the fragment row/col loops (lower ids nest outer), so B loads are contiguous and - # the A row value is loaded once per kk - ir, kk = UOp.range(ROWS, 10, AxisType.LOOP), UOp.range(BLOCK_K, 11, AxisType.LOOP) - jj = UOp.range(BLOCK_N, 12, AxisType.LOOP) - acc = C_loc.after(kk)[tid*ROWS + ir, jj] + A_shared[tid*ROWS + ir, kk].cast(dtypes.float32) * B_shared[kk, jj].cast(dtypes.float32) + ir, kk = UOp.range(TM, 11, AxisType.LOOP), UOp.range(BLOCK_K, 12, AxisType.LOOP) + jj = UOp.range(TN, 13, AxisType.LOOP) + acc = C_loc.after(kk)[ty*TM + ir, tx*TN + jj] + A_shared[ty*TM + ir, kk].cast(dtypes.float32) * B_shared[kk, tx*TN + jj].cast(dtypes.float32) # closing the ko loop here too; codegen adds the barrier so no thread overwrites the tiles while others still read them - C_loc = C_loc[tid*ROWS + ir, jj].set(acc, end=(kk, ir, jj, ko)) + C_loc = C_loc[ty*TM + ir, tx*TN + jj].set(acc, end=(kk, ir, jj, ko)) # for i, j in T.Parallel(BLOCK_M, BLOCK_N): C_local[i, j] = T.max(C_local[i, j], 0) # T.copy(C_local, C[by * BLOCK_M, bx * BLOCK_N]) -- per-thread store of the fragment shard (relu fused into it) # LOOP: these loops are the per-thread output layout; convert_loop_to_global must not globalize them - ie, je = UOp.range(ROWS, 13, AxisType.LOOP), UOp.range(BLOCK_N, 14, AxisType.LOOP) - c_st = c[by*BLOCK_M + tid*ROWS + ie, bx*BLOCK_N + je].store(C_loc[tid*ROWS + ie, je].relu().cast(c.dtype)) + ie, je = UOp.range(TM, 14, AxisType.LOOP), UOp.range(TN, 15, AxisType.LOOP) + c_st = c[by*BLOCK_M + ty*TM + ie, bx*BLOCK_N + tx*TN + je].store(C_loc[ty*TM + ie, tx*TN + je].relu().cast(c.dtype)) # all open ranges are closed at the final store (ko was closed above). # the fragment UNSHARDs go to codegen as is: multi_pm there resolves the full-tile view into per-thread shard code - return c_st.end(je, ie, tid, by, bx).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=())) + return c_st.end(je, ie, tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=())) # --------------------------------------------------------------------------- # python wrapper: same signature as the tilelang function @@ -163,7 +171,7 @@ def matmul_relu(a:Tensor, b:Tensor) -> Tensor: if __name__ == "__main__": from tinygrad import Device assert Device[Device.DEFAULT].renderer.has_local, "this GPU-style kernel needs a backend with local memory (LOCAL ranges + barriers)" - M = K = N = 256 # 4x4 grid of 64x64 tiles, 4 K chunks + M = K = N = getenv("N", 256) # 4x4 grid of 64x64 tiles, 4 K chunks a = Tensor.randn(M, K, dtype=dtypes.float16).contiguous() b = Tensor.randn(K, N, dtype=dtypes.float16).contiguous() diff --git a/spec/tinyspec.pdf b/spec/tinyspec.pdf index 8f358d4e82..4e380f5d9d 100644 Binary files a/spec/tinyspec.pdf and b/spec/tinyspec.pdf differ diff --git a/spec/tinyspec.tex b/spec/tinyspec.tex index e76455a3c2..4444d84563 100644 --- a/spec/tinyspec.tex +++ b/spec/tinyspec.tex @@ -80,7 +80,7 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op \op{Index} & $(T, i_0, i_1, \ldots)$ & --- & Index from left. $()$-shaped $i$ removes dim; $(k,)$-shaped makes it $k$. \\ \op{Stack} & $(T_0, T_1, \ldots)$ & --- & Join along a newly created leading axis. All shapes must match. \\ \op{Bitcast} & $(T,)$ & dtype & Reinterpret storage as target dtype; preserve total bytes. \\ -\op{Unshard} & $(T, R)$ & axis $a$ & Concatenate the shards indexed by \op{Range} $R$ along $a$; $R$ is outer. \\ +\op{Unshard} & $(T, R_0, R_1, \ldots)$ & axes $(a_0, a_1, \ldots)$ & Concatenate shards of \op{Range} $R_k$ along axis $a_k$; $R_k$ is outer. \\ \bottomrule \end{tabular} @@ -260,7 +260,7 @@ Every UOp has a \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{addrspa \op{Const} & from arg & $()$ & \textsc{null} & $[v, v]$ \\ \op{Param} & from arg & from $\mathrm{src}[0]$ & from arg & from src or dtype range \\[3pt] Movement ops & $\mathrm{src}[0].\mathrm{dtype}$ & (see op) & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\ -\op{Unshard} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0]$, axis $\times n$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\ +\op{Unshard} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0]$, each $a_k \times n_k$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\ \op{Reduce} & $\mathrm{src}[0].\mathrm{dtype}$ & remove first $n$ axes & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\[3pt] \op{Cast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & clamped to dtype \\ \op{Bitcast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\ @@ -286,9 +286,9 @@ $[a,A]$, $[b,B]$, $[c,C]$ denote min\_max of $\mathrm{src}[0]$, $\mathrm{src}[1] Default \emph{dtype range}: $[\mathrm{dtype\_min},\, \mathrm{dtype\_max}]$. \medskip -\textbf{axis} tracks the multi-device sharding dimension. \op{Unshard} defines it (axis $=$ arg). \op{Buffer} with $n$-tuple device: axis $= 0$ (device dim). -\op{Reshape} remaps axis to preserve the shard boundary. \op{Permute} follows the permutation. \op{Expand} shifts axis right by $|\mathbf{n}|$. -\op{Reduce} on the shard axis $\to$ \textsc{null} (shard axis is among the first $n$ axes). \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}. +\textbf{sharding} tracks multi-device sharding as a set of (axis, \op{Range}) pairs. \op{Unshard} defines it: arg is the tuple of sharded axes, one \op{Range} in src per axis (positional: the $k$-th \op{Range} shards the $k$-th axis). \op{Buffer} with $n$-tuple device: sharded on axis $0$ (device dim). The single-axis convenience \textbf{axis} is \textsc{null} unless exactly one axis is sharded. +\op{Reshape} remaps each sharded axis to preserve its shard boundary. \op{Permute} follows the permutation. \op{Expand} shifts all sharded axes right by $|\mathbf{n}|$. +\op{Reduce} on a sharded axis drops it. \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}. %% ============================================================ \subsection*{Kernel Optimizations (OptOps) \normalfont\small--- schedule-level transforms on kernel ranges} @@ -382,7 +382,7 @@ def scatter_add(T, idx, val): Let $D = (d_0, \ldots, d_{n-1})$ be an $n$-tuple device. \op{Copy} to an $n$-tuple device reshards with axis $= 0$. \op{Copy} never changes shape. -\textbf{Sharding} splits a tensor along an axis across $n$ devices. It opens a \op{Range} of type \texttt{DEVICE} (a symbolic per-device index $d$), shrinks each device's view to its piece, then closes the range with \op{Unshard}$(T, R, a)$. The result is a logical tensor whose shape along axis $a$ is the full size; each device holds $1/n$ of it. \op{Unshard} is the inverse of sharding --- it marks the boundary between per-device computation and the logical multi-device tensor. The range need not be \texttt{DEVICE}; e.g.\ a \texttt{WARP} range closes the same way, concatenating per-lane shards along $a$ with the range as the outer factor. +\textbf{Sharding} splits a tensor along an axis across $n$ devices. It opens a \op{Range} of type \texttt{DEVICE} (a symbolic per-device index $d$), shrinks each device's view to its piece, then closes the range with \op{Unshard}$(T, R, a)$. The result is a logical tensor whose shape along axis $a$ is the full size; each device holds $1/n$ of it. \op{Unshard} is the inverse of sharding --- it marks the boundary between per-device computation and the logical multi-device tensor. The range need not be \texttt{DEVICE}; e.g.\ a \texttt{WARP} range closes the same way, concatenating per-lane shards along $a$ with the range as the outer factor. A tensor may be sharded along several axes at once: \op{Unshard}$(T, R_0, R_1, \ldots;\; a_0, a_1, \ldots)$ carries one \op{Range} per sharded axis, and every movement op maps each sharded axis independently. \begin{lstlisting} # T has shape (s,) on a single device. diff --git a/test/backend/test_multitensor.py b/test/backend/test_multitensor.py index 8b58f4c554..ce218ac7a9 100644 --- a/test/backend/test_multitensor.py +++ b/test/backend/test_multitensor.py @@ -425,6 +425,49 @@ class TestMultiBufferView(unittest.TestCase): run_linear(linear, var_vals) np.testing.assert_equal(out.numpy(), ref[5].numpy()) +@unittest.skipIf(not_support_multi_device(), "need multi") +class Test2DShard(unittest.TestCase): + def setUp(self): + self.devices_4 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4)) + self.rng = UOp.range(4, -1, AxisType.DEVICE) + self.rng0, self.rng1 = self.rng // 2, self.rng % 2 + + def _shard_2d(self, t:Tensor) -> Tensor: + u = t.uop.copy_to_device(self.devices_4)._shard(0, self.rng0)._shard(1, self.rng1).unshard((0, 1), (self.rng0, self.rng1)) + return Tensor(u) + + def test_2d_shard_basic(self): + ref = Tensor.arange(16).reshape(4, 4).contiguous().realize() + t = self._shard_2d(ref) + out = t.contiguous().realize() + np.testing.assert_equal(out.numpy(), ref.numpy()) + + def test_2d_shard_elementwise(self): + ref = Tensor.arange(16).reshape(4, 4).contiguous().realize() + t = self._shard_2d(ref) + out = (t + 1).contiguous().realize() + np.testing.assert_equal(out.numpy(), ref.numpy() + 1) + + def test_2d_shard_sum_all(self): + ref = Tensor.arange(16).reshape(4, 4).contiguous().realize() + t = self._shard_2d(ref) + out = t.sum().contiguous().realize() + np.testing.assert_equal(out.numpy(), np.array(ref.numpy().sum())) + + def test_2d_shard_sum_non_sharded_axis(self): + ref = Tensor.arange(4*4*2).reshape(4, 4, 2).contiguous().realize() + t = self._shard_2d(ref) + out = t.sum(axis=2).contiguous().realize() + np.testing.assert_equal(out.numpy(), ref.numpy().sum(axis=2)) + + def test_2d_shard_matmul(self): + a = Tensor.arange(16).reshape(4, 4).contiguous().realize() + b = Tensor.arange(16).reshape(4, 4).contiguous().realize() + a_s = self._shard_2d(a) + b_s = self._shard_2d(b) + out = (a_s @ b_s).contiguous().realize() + np.testing.assert_equal(out.numpy(), a.numpy() @ b.numpy()) + @unittest.skipIf(not_support_multi_device(), "need multi") class TestMultiTransformer(unittest.TestCase): @needs_second_gpu diff --git a/test/unit/test_multitensor.py b/test/unit/test_multitensor.py index d87214934d..60f4e6701c 100644 --- a/test/unit/test_multitensor.py +++ b/test/unit/test_multitensor.py @@ -597,7 +597,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase): t = Tensor.arange(64).reshape(8, 8).clone().realize() t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0) - with self.assertRaises(AssertionError): + with self.assertRaises(RuntimeError): # sharded axis shrink on non-device boundry is not allowed a = t.shrink(((0, 3), (0, 8))).contiguous() a.schedule_linear() diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 2877eb119f..2c9345af9b 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -83,7 +83,7 @@ def contiguous_mops_to_view(c:UOp, src:UOp): resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view") if resolved.op is not Ops.UNSHARD: return None if (view := _make_buffer_view(resolved.src[0])) is None: return None - return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1]).contiguous(tag=c.tag) + return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:]).contiguous(tag=c.tag) return None diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index b85ebe4639..32cc82485c 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -1,12 +1,9 @@ from tinygrad.helpers import all_same, prod, getenv, ALLREDUCE_CAST from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, AxisType, graph_rewrite, broadcast_axes, _broadcast_shape, sint_to_uop +from tinygrad.uop.ops import sint, ssimplify from tinygrad.dtype import dtypes from tinygrad.schedule.allreduce import handle_allreduce -def shard_count(multi:UOp) -> int: - # the shard count: the device count for multi-device UNSHARDs, the range size otherwise (e.g. threads for LOCAL fragments) - return len(multi.device) if isinstance(multi.device, tuple) else int(multi.src[1].vmax)+1 - # ***** multi rewrite MSELECT/MSTACK ***** def _apply_shrink(marg, s:UOp, i:int) -> UOp: @@ -75,6 +72,13 @@ def shard_srcs(msrcs:tuple[UOp, ...], axis:int) -> list[UOp]: return srcs def alu_multi(root:UOp): + multis = [m for m in root.src if m.op is Ops.UNSHARD] + if not multis: return None + sharding = multis[0].sharding + if len(multis) == len(root.src) and all(m.sharding == sharding for m in multis): + srcs = [m.src[0] for m in root.src] + return srcs[0].alu(root.op, *srcs[1:]).unshard(multis[0].arg, multis[0].src[1:]) + # resharding: single-axis fallback via shard_srcs axis = root.axis assert axis is not None srcs = shard_srcs(root.src, axis) @@ -82,86 +86,145 @@ def alu_multi(root:UOp): def reduce_multi(root:UOp, multi:UOp): op, num_axes = root.arg - if multi.axis is not None and multi.axis < num_axes: - local = multi.src[0]._rop(op, tuple(range(num_axes))) - # allreduce in pre-cast dtype when sum_acc_dtype promoted from bf16/half + sharding = multi.sharding + reduced = [(ax, rng) for ax, rng in sharding if ax < num_axes] + remaining = [(ax, rng) for ax, rng in sharding if ax >= num_axes] + local = multi.src[0]._rop(op, tuple(range(num_axes))) + if reduced: + assert not remaining, f"partial allreduce not supported for multi-axis sharding {sharding}" + # all sharded axes are reduced: full allreduce if ALLREDUCE_CAST and multi.src[0].op is Ops.CAST and multi.src[0].src[0].dtype in (dtypes.bfloat16, dtypes.half): orig_dtype = multi.src[0].src[0].dtype return local.cast(orig_dtype).allreduce(op, multi.device).cast(local.dtype) return local.allreduce(op, multi.device) - # reduce on non sharded axes, piecewise is fine. if axis is None this is also correct - new_axis = multi.axis - num_axes if multi.axis is not None else None - return multi.src[0]._rop(op, tuple(range(num_axes))).unshard(new_axis, multi.src[1]) + # no sharded axes reduced: piecewise, keep all remaining sharding + new_axes = tuple(ax - num_axes for ax, _ in remaining) + new_rngs = tuple(rng for _, rng in remaining) + return local.unshard(new_axes, new_rngs) def reshape_multi(root:UOp, multi:UOp): if prod(multi.shape) != prod(new_shape:=root.marg): raise RuntimeError("reshape must maintain prod(shape)") - if (new_axis:=root.axis) is not None: new_shape = tuple(s//shard_count(multi) if a==new_axis else s for a,s in enumerate(new_shape)) - return multi.src[0].reshape(new_shape).unshard(new_axis, multi.src[1]) + # map every sharded axis through the reshape: the axis boundary must survive intact and stay divisible by its shard count + arg_acc:list[sint] = [1] + for s in new_shape: arg_acc.append(ssimplify(arg_acc[-1]*s)) + new_shardings = [] + for ax, rng in multi.sharding: + count = int(rng.vmax)+1 + target = prod(multi.shape[:ax]) + if target not in arg_acc: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards") + new_ax = len(arg_acc) - arg_acc[::-1].index(target) - 1 + if new_shape[new_ax] % count != 0: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards") + new_shardings.append((new_ax, rng)) + new_axs = {a for a, _ in new_shardings} + new_shape = tuple(s//(int(rng.vmax)+1) if a in new_axs else s for a,s in enumerate(new_shape)) + return multi.src[0].reshape(new_shape).unshard(tuple(a for a,_ in new_shardings), tuple(r for _,r in new_shardings)) def expand_multi(root:UOp, multi:UOp): - new_axis = None if multi.axis is None else multi.axis + len(root.marg) - return multi.src[0]._mop(Ops.EXPAND, arg=root.marg).unshard(new_axis, multi.src[1]) + shift = len(root.marg) + return multi.src[0]._mop(Ops.EXPAND, arg=root.marg) \ + .unshard(tuple(ax+shift for ax,_ in multi.sharding), tuple(r for _,r in multi.sharding)) def pad_multi(root:UOp, multi:UOp): - assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]), f"padding not supported for {root.marg=}" - local_pad = tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.marg)) - return multi.src[0]._mop(Ops.PAD, local_pad).unshard(multi.axis, multi.src[1]) + for ax, _ in multi.sharding: + assert root.marg[ax] == (0, multi.shape[ax]), f"padding not supported for {root.marg=}" + counts = {a for a,_ in multi.sharding} + local_pad = tuple((0, multi.src[0].shape[a]) if a in counts else s for a,s in enumerate(root.marg)) + return multi.src[0]._mop(Ops.PAD, local_pad).unshard(multi.arg, multi.src[1:]) def permute_multi(root:UOp, multi:UOp): # all permutes supported! - return multi.src[0].permute(root.marg).unshard(root.axis, multi.src[1]) + return multi.src[0].permute(root.marg) \ + .unshard(tuple(root.marg.index(ax) for ax,_ in multi.sharding), tuple(r for _,r in multi.sharding)) def shrink_multi(root:UOp, multi:UOp): - if multi.axis is not None: - shard_sz = multi.src[0].shape[multi.axis] - s, l = root.marg[multi.axis] # SHRINK marg is (start, length) - # shrink to exactly this range's own shard: this resolves the UNSHARD into its per-shard view - # (e.g. a fragment indexed by its LOCAL thread range becomes that thread's REG shard, no copy needed) - if sint_to_uop(l).ssimplify() == shard_sz and (sint_to_uop(s)-multi.src[1]*shard_sz).ssimplify() == 0: - non_shard_shrink = tuple((0, multi.src[0].shape[i]) if i == multi.axis else t for i, t in enumerate(root.marg)) - return multi.src[0]._mop(Ops.SHRINK, non_shard_shrink) - shard_bounds = tuple((s,e-s) for s,e in multi.bounds) if multi.axis is not None else () - assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]) or root.marg[multi.axis] in shard_bounds, \ - f"shrinking not supported for {root.marg=}" - if multi.axis is not None and root.marg[multi.axis] in shard_bounds and root.marg[multi.axis] != (0, multi.shape[multi.axis]): - # NOTE: shrink on the shard axis is only allowed when result is a single partition, denoted by the new real - # we just copy it to all the devices, no real. this will be optimized out later - non_shard_shrink = tuple((0, multi.src[0].shape[i]) if i == multi.axis else s for i, s in enumerate(root.marg)) - return multi.src[0].copy_to_device(multi.device, arg=shard_bounds.index(root.marg[multi.axis]))._mop(Ops.SHRINK, non_shard_shrink) - local_shrink = tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.marg)) - return multi.src[0]._mop(Ops.SHRINK, local_shrink).unshard(multi.axis, multi.src[1]) + # resolve each sharded axis independently: a shrink to exactly this range's own shard resolves the UNSHARD along + # that axis (e.g. a fragment indexed by its LOCAL thread range becomes that thread's REG shard, no copy needed) + local_marg = list(root.marg) + remaining = list(multi.sharding) + for ax, rng in multi.sharding: + shard_sz = multi.src[0].shape[ax] + s, l = root.marg[ax] # SHRINK marg is (start, length) + if sint_to_uop(l).ssimplify() == shard_sz and (sint_to_uop(s)-rng*shard_sz).ssimplify() == 0: + local_marg[ax] = (0, shard_sz) + remaining.remove((ax, rng)) + continue + part_bounds = tuple((i*shard_sz, shard_sz) for i in range(int(rng.vmax)+1)) + if (s, l) == (0, multi.shape[ax]): local_marg[ax] = (0, shard_sz) # full axis stays sharded, shrink the other axes locally + else: + # NOTE: otherwise a shrink on the shard axis is only allowed on the legacy device path, selecting a single + # partition (which is copied to all the devices and optimized out later) + if len(multi.sharding) != 1 or not isinstance(multi.device, tuple) or (s, l) not in part_bounds: + raise RuntimeError(f"shrinking not supported for {root.marg=}") + non_shard_shrink = tuple((0, shard_sz) if i == ax else t for i, t in enumerate(root.marg)) + return multi.src[0].copy_to_device(multi.device, arg=part_bounds.index((s, l)))._mop(Ops.SHRINK, non_shard_shrink) + val = multi.src[0]._mop(Ops.SHRINK, tuple(local_marg)) + return val if not remaining else val.unshard(tuple(a for a,_ in remaining), tuple(r for _,r in remaining)) def flip_multi(root:UOp, multi:UOp): - assert multi.axis is None or not root.marg[multi.axis], "flipping not supported on sharded axis" - return multi.src[0].flip([i for i,x in enumerate(root.marg) if x]).unshard(multi.axis, multi.src[1]) + for ax, _ in multi.sharding: + if root.marg[ax]: raise RuntimeError(f"flipping not supported on sharded axis {ax}") + return multi.src[0].flip([i for i,x in enumerate(root.marg) if x]).unshard(multi.arg, multi.src[1:]) def stack_multi(root:UOp): # STACK adds a leading axis: srcs are sharded one axis below the output + multis = [m for m in root.src if m.op is Ops.UNSHARD] + if not multis: return None + sharding = multis[0].sharding + if all(m.sharding == sharding for m in multis): + srcs = [m.src[0] if m.op is Ops.UNSHARD else m for m in root.src] + new_sharding = tuple((ax+1, rng) for ax, rng in sharding) + return UOp(Ops.STACK, src=tuple(srcs)).unshard(tuple(a for a,_ in new_sharding), tuple(r for _,r in new_sharding)) + # resharding: single-axis fallback axis = root.axis assert axis is not None return UOp(Ops.STACK, src=tuple(shard_srcs(root.src, axis-1))).unshard(axis, next(m.src[1] for m in root.src if m.op is Ops.UNSHARD)) def index_multi(root:UOp, multi:UOp): - # INDEX on UNSHARD: resolve the sharded axis into this range's own shard (idx - rng*shard_sz) - if multi.axis is None: return None - shard_sz = multi.src[0].shape[multi.axis] - local = (root.src[1+multi.axis] - multi.src[1]*shard_sz).simplify() - # the index along the sharded axis must be provably inside this shard - if local.vmin < 0 or local.vmax >= shard_sz: return None - return multi.src[0].index(*root.src[1:1+multi.axis], local, *root.src[2+multi.axis:]) + # INDEX on UNSHARD: resolve each sharded axis into this range's own shard (idx - rng*shard_sz) + idxs = list(root.src[1:]) + for ax, rng in multi.sharding: + shard_sz = multi.src[0].shape[ax] + local = (idxs[ax] - rng*shard_sz).simplify() + # the index along each sharded axis must be provably inside this shard + if local.vmin < 0 or local.vmax >= shard_sz: return None + idxs[ax] = local + return multi.src[0].index(*idxs) + +def _shard_idx(rng:UOp, dev_idx:int) -> int: + drngs = [r for r in rng.ranges if r.arg[-1] is AxisType.DEVICE] + return 0 if not drngs else int(rng.substitute({drngs[0]: drngs[0].const_like(dev_idx)}).ssimplify()) def copy_multi(multi:UOp, device:str | tuple[str, ...]): - assert multi.axis is not None, "all multi ops have axis" + sharding = multi.sharding if isinstance(device, str): - pieces = [multi.src[0].mselect(i).copy_to_device(device) for i in range(len(multi.device))] - return pieces[0].cat(*pieces[1:], dim=multi.axis) - return multi.src[0]._unshard(multi.axis).allreduce(Ops.ADD, device) + # reconstruct by concatenating along each axis from last to first + piece_info: list[tuple[tuple, UOp]] = [] + for i in range(len(multi.device)): + idxs = tuple(_shard_idx(r, i) for _, r in sharding) + piece_info.append((idxs, multi.src[0].mselect(i).copy_to_device(device))) + for j in range(len(sharding) - 1, -1, -1): + ax, rng = sharding[j] + groups: dict[tuple, list[tuple[int, UOp]]] = {} + for idxs, p in piece_info: + key = idxs[:j] + idxs[j+1:] + groups.setdefault(key, []).append((idxs[j], p)) + piece_info = [] + for key in sorted(groups): + grp = sorted(groups[key], key=lambda x: x[0]) + piece_info.append((key, grp[0][1].cat(*[x[1] for x in grp[1:]], dim=ax))) + return piece_info[0][1] + # multi-device target: unshard all axes and allreduce + val = multi.src[0] + for ax, rng in sharding: + bsz = val.shape[ax] + val = val.pad(tuple((0,0) if a != ax else (bsz*rng, bsz*int(rng.vmax) - bsz*rng) for a in range(len(val.shape)))) + return val.allreduce(Ops.ADD, device) -def store_after_multi(dest:UOp, src:UOp): return dest.after(dest.store(src.src[0])).unshard(src.axis, src.src[1]) +def store_after_multi(dest:UOp, src:UOp): return dest.after(dest.store(src.src[0])).unshard(src.arg, src.src[1:]) def passthrough_multi(root:UOp, multi:UOp): new_src = (multi.src[0],)+tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src[1:]) - return UOp(root.op, root.dtype, src=new_src, arg=root.arg).unshard(multi.axis, multi.src[1]) + return UOp(root.op, root.dtype, src=new_src, arg=root.arg).unshard(multi.arg, multi.src[1:]) def rewrite_into_function(call:UOp): if call.arg.precompile: return None @@ -171,7 +234,7 @@ def rewrite_into_function(call:UOp): assert new_body.op is Ops.TUPLE if any(s.op is Ops.UNSHARD for s in new_body.src): shard_call = call.replace(src=(UOp.maketuple(*[s.src[0] if s.op is Ops.UNSHARD else s for s in new_body.src]),)+new_args) - return UOp.maketuple(*[shard_call.gettuple(i).unshard(s.axis, s.src[1]) if s.op is Ops.UNSHARD else shard_call.gettuple(i) + return UOp.maketuple(*[shard_call.gettuple(i).unshard(s.arg, s.src[1:]) if s.op is Ops.UNSHARD else shard_call.gettuple(i) for i, s in enumerate(new_body.src)]) return call.replace(src=(new_body,)+new_args) @@ -195,14 +258,13 @@ multi_pm = PatternMatcher([ (UPat(Ops.AFTER, src=(UPat(Ops.UNSHARD), UPat(Ops.STORE, src=(UPat(Ops.UNSHARD, name="dest"), UPat(Ops.UNSHARD, name="src"))))), store_after_multi), (UPat(Ops.COPY, src=(UPat(Ops.UNSHARD, name="multi"),), name="copy"), lambda multi,copy: copy_multi(multi, copy.arg)), (UPat(Ops.ALLREDUCE, src=(UPat(Ops.UNSHARD, name="multi"),), name="red"), - lambda multi,red: multi.src[0].allreduce(*red.arg).unshard(multi.axis, multi.src[1])), + lambda multi,red: multi.src[0].allreduce(*red.arg).unshard(multi.arg, multi.src[1:])), # resolve TUPLE+GETTUPLE (needed in multi) (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]), # GETTUPLE on UNSHARD: passthrough UNSHARD (e.g. when FUNCTION was replaced by UNSHARD(GETTUPLE(...))) (UPat(Ops.GETTUPLE, src=(UPat(Ops.UNSHARD, name="multi"),), name="g"), - lambda g, multi: multi.src[0].gettuple(g.arg).unshard(multi.axis, multi.src[1]) if multi.src[0].op in {Ops.FUNCTION, Ops.TUPLE} - else multi), + lambda g, multi: multi.src[0].gettuple(g.arg).unshard(multi.arg, multi.src[1:]) if multi.src[0].op in {Ops.FUNCTION, Ops.TUPLE} else multi), # rewrite into FUNCTION calls explicitly for UNSHARD (value-producing) (UPat(Ops.FUNCTION, name="call"), rewrite_into_function), (UPat((Ops.CALL, Ops.FUNCTION, Ops.AFTER), src=(UPat(Ops.UNSHARD, name="multi"), ), name="root", allow_any_len=True), passthrough_multi), diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 5c446cf989..a917297e06 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -439,7 +439,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): case Ops.FLIP: if len(ps) != len(self.marg) or not all(isinstance(x, bool) for x in self.marg): raise ValueError(f"bad flip on {ps}, {self.marg}") return ps - case Ops.UNSHARD: return tuple(s*(int(self.src[1].vmax)+1) if a == self.axis else s for a,s in enumerate(ps)) + case Ops.UNSHARD: return tuple(s*(int(self.src[1:][self.arg.index(a)].vmax)+1) if a in self.arg else s for a,s in enumerate(ps)) case Ops.REDUCE: num_axes = self.arg[1] if not isinstance(num_axes, int) or num_axes < 0 or num_axes > len(ps): @@ -665,15 +665,24 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # *** multi-device helpers *** - def unshard(self, axis:int|None, device_range:UOp|None=None): + def unshard(self, axis:int|tuple[int, ...]|None, device_range:UOp|tuple[UOp, ...]|None=None): assert axis is not None, "multi None is no longer supported" - # an UNSHARD always has two srcs: the value and the range it ends (defaults to a DEVICE range over the devices) - # the range need not be DEVICE, e.g. a LOCAL range shards a kernel tile into per-thread fragments + # an UNSHARD carries the value and one sharding range per sharded axis (arg is the tuple of sharded axes, + # sorted). the single-axis axis form defaults the range to a DEVICE range over the devices; a range need not + # be DEVICE, e.g. a LOCAL range shards a kernel tile into per-thread fragments + if isinstance(axis, int): axis = (axis,) if device_range is None: assert isinstance(self.device, tuple), f"multi device must be tuple, {self.device} isn't" - device_range = UOp.range(len(self.device), -1, AxisType.DEVICE) - assert device_range.op is Ops.RANGE - return UOp(Ops.UNSHARD, src=(self, device_range), arg=axis) + device_range = (UOp.range(len(self.device), -1, AxisType.DEVICE),) + if isinstance(device_range, UOp): device_range = (device_range,) + assert isinstance(device_range, tuple) and len(axis) == len(device_range) and len(set(axis)) == len(axis) + axis, device_range = map(tuple, zip(*sorted(zip(axis, device_range)))) + return UOp(Ops.UNSHARD, src=(self, *device_range), arg=axis) + + @property + def sharding(self) -> tuple[tuple[int, UOp], ...]: + """(axis, RANGE) pairs this value is sharded over (the source of truth for shard bounds/counts).""" + return tuple(zip(self.arg, self.src[1:])) if self.op is Ops.UNSHARD else () @property def bounds(self): @@ -685,7 +694,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def axis(self) -> int|None: # COPY removes axis. TODO: add more tests for this, and consider MSELECT/MSTACK if self.op is Ops.COPY: return None - if self.op is Ops.UNSHARD: return self.arg + if self.op is Ops.UNSHARD: + if len(self.arg) != 1: raise RuntimeError(f"UOp is sharded on multiple axes {self.arg}, use .sharding") + return self.arg[0] # GETTUPLE: axis comes from the specific TUPLE element, not src[0] if self.op is Ops.GETTUPLE: in_tuple = self.src[0].src[0] if self.src[0].op is Ops.FUNCTION else self.src[0] @@ -723,7 +734,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return self.pad(tuple((0,0) if a != axis else (bsz*dnum, bsz*(dcount-1) - bsz*dnum) for a in range(len(self.shape)))) def _shard(self, axis:int, rng:UOp) -> UOp: - assert rng.op is Ops.RANGE, f"_shard requires a RANGE, got {rng.op}" if len(self.shape) == 0: return self # scalars broadcast, no sharding needed dcount = int(rng.vmax)+1 if self.shape[axis] % dcount != 0: raise RuntimeError(f"multi axis uneven: {self.shape[axis]=} {axis=} {dcount=}") @@ -879,7 +889,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # CL 1.1 provides the clCreateSubBuffer API, but at the time of writing, relevant CL runtimes (rusticl, adreno, nvidia, amd) do not provide # reasonable values for CL_DEVICE_MEM_BASE_ADDR_ALIGN. cl_ext_buffer_device_address could potentially help, but this extension is not provided # by relevant CL runtimes at time of writing. - if any(d.startswith(("WEBGPU", "CL")) for d in ((self.device,) if isinstance(self.device, str) else self.device)): return None + if (dev:=self.device) is not None and any(d.startswith(("WEBGPU", "CL")) for d in ((dev,) if isinstance(dev, str) else dev)): return None idx = self.flatten().index(UOp.range(self.numel(), 0)) out = graph_rewrite(idx, pm_mops+symbolic+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset") diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 1596526d9b..5298129d0a 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -176,9 +176,9 @@ spec_tensor = PatternMatcher([ len(red.arg) == 2 and red.arg[0] in GroupOp.Reduce and is_device(red.arg[1])), # UNSHARD/MSELECT/MSTACK - # an UNSHARD always has two srcs: the value and the range it ends (usually DEVICE, but any typed range can be sharded over) - (UPat(Ops.UNSHARD, name="multi"), lambda multi: len(multi.src) == 2 and matches_dtype(multi.src[0], multi.dtype) - and isinstance(multi.arg, int) and multi.src[1].op is Ops.RANGE and isinstance(multi.src[1].arg[-1], AxisType)), + # an UNSHARD carries the value and one sharding range per sharded axis (usually a DEVICE RANGE, but can be a derived expression) + (UPat(Ops.UNSHARD, name="multi"), lambda multi: len(multi.src) == 1+len(multi.arg) and matches_dtype(multi.src[0], multi.dtype) + and all(isinstance(a, int) for a in multi.arg) and all(r.dtype in dtypes.weaks for r in multi.src[1:])), (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)), (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)),