From d65ea465edfd561d69a5d5f465fd1c820f733acd Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Thu, 30 Jul 2026 20:53:28 -0700 Subject: [PATCH] cleanup gemm fragment + add store unshard (#17313) * cleanup gemm fragment + add store unshard * multi * fix --- extra/gemm_fragment.py | 156 +++++++++-------------------- test/backend/test_custom_kernel.py | 72 +++++++++++++ tinygrad/schedule/multi.py | 41 ++++++-- 3 files changed, 157 insertions(+), 112 deletions(-) diff --git a/extra/gemm_fragment.py b/extra/gemm_fragment.py index d9ad4c5d22..457439e3d1 100644 --- a/extra/gemm_fragment.py +++ b/extra/gemm_fragment.py @@ -1,55 +1,25 @@ """ 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 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: - - @tilelang.jit - def matmul_relu(A, B, block_M=64, block_N=64, block_K=64, - dtype=T.float16, accum_dtype=T.float32): - M, N, K = T.const('M, N, K') - C = T.empty([M, N], dtype) - with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): - A_shared = T.alloc_shared((block_M, block_K), dtype) - B_shared = T.alloc_shared((block_K, block_N), dtype) - C_local = T.alloc_fragment((block_M, block_N), accum_dtype) - T.clear(C_local) - for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): - T.copy(A[by * block_M, ko * block_K], A_shared) - T.copy(B[ko * block_K, bx * block_N], B_shared) - T.gemm(A_shared, B_shared, C_local) - 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]) - return C - -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 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 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 - T.gemm (no WMMA) -> C_local[..].set(C_local.after(k)[..] + a_shared[..] * b_shared[..], end=k) - with k a loop-carried LOOP range (codegen builds the register accumulator - from this self-referential store automatically) - T.copy(fragment, gmem) -> gmem.index(gidx).store(C_local[..]).end(all_ranges) - UNSHARD lowering -> multi_pm in codegen (full_rewrite_to_sink): INDEX/AFTER/STORE ops on the - full-tile view become per-thread shard ops, no UNSHARD survives into the program. + @tilelang.jit + def matmul_relu(A, B, block_M=64, block_N=64, block_K=64, + dtype=T.float16, accum_dtype=T.float32): + M, N, K = T.const('M, N, K') + C = T.empty([M, N], dtype) + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): + A_shared = T.alloc_shared((block_M, block_K), dtype) + B_shared = T.alloc_shared((block_K, block_N), dtype) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + T.clear(C_local) + for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): + T.copy(A[by * block_M, ko * block_K], A_shared) + T.copy(B[ko * block_K, bx * block_N], B_shared) + T.gemm(A_shared, B_shared, C_local) + 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]) + return C """ from tinygrad.dtype import dtypes, AddrSpace, DType @@ -61,25 +31,17 @@ from tinygrad.tensor import Tensor # tilelang builtins, expressed with tinygrad UOp APIs # --------------------------------------------------------------------------- -def alloc_shared(shape:tuple[int, ...], dtype:DType) -> UOp: +def alloc_shared(shape:tuple[int, ...], dtype:DType, slot:int) -> 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) + return UOp.placeholder(tuple(shape), dtype, slot, AddrSpace.LOCAL) -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 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. - """ +def alloc_fragment(shape:tuple[int, ...], dtype:DType, slot:int, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp: + """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid.""" 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) + fragment = UOp.placeholder(shard_shape, dtype, slot, AddrSpace.REG) return fragment.unshard(axes, rngs) # --------------------------------------------------------------------------- @@ -106,64 +68,45 @@ 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=128) as (bx, by): bx = UOp.range(cdiv(N, BLOCK_N), 0, AxisType.GLOBAL) by = UOp.range(cdiv(M, BLOCK_M), 1, AxisType.GLOBAL) - # tx (N, 16) is the fast/inner LOCAL axis so a warp covers 16 cols x 2 rows -- - # matching tilelang's (tidx>>4, tidx&15) warp composition. This keeps the 8 A_shared - # reads in a warp on only 2 row-groups (broadcast across 16 cols) instead of 8 rows - # (8-way bank conflict), since A_shared[row*512 + ...] all map to the same bank when 8 - # distinct rows land in one warp. + + # 16*8 threads = 128 threads tx = UOp.range(TX, 2, AxisType.LOCAL) ty = UOp.range(TY, 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) + # shared + fragment (regs) + A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype, 0) + B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype, 1) + C_local = alloc_fragment((TM, TY, TX, TN), dtypes.float32, 0, (1, 2), (ty, tx)) - # C_local = T.alloc_fragment((BLOCK_M, BLOCK_N), accum_dtype) -- an 8x4 REG tile per thread of the 8x16 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 sub-tile - ic, jc = UOp.range(TM, 4, AxisType.LOOP), UOp.range(TN, 5, AxisType.UPCAST) - C_loc = C_local[ic*TM + ty, tx*TN + jc].set(0.0, end=(ic, jc)) + # zero out the regs to start. this is expanded by the devectorizer + C_local = C_local.after(C_local.store(0.0)) # 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), 6, AxisType.LOOP) - # T.copy(A[by * BLOCK_M, ko * BLOCK_K], A_shared) -- each thread copies its own 8x4 sub-tile. - # Row index is iar*TM + ty (strided by TM across ty), matching tilelang's layout: thread ty owns - # rows {ty, ty+8, ..., ty+56} not {ty*8, ..., ty*8+7}. - iar, ka = UOp.range(TM, 7, AxisType.LOOP), UOp.range(TN, 8, AxisType.UPCAST) - A_store = A_shared[iar*TM + ty, tx*TN + ka].store(a[by*BLOCK_M + iar*TM + ty, ko*BLOCK_K + tx*TN + ka]).end(iar, ka) + # index the outer matrices + a = a.rearrange("(m bm) (k bk) -> m k bm bk", bm=BLOCK_M, bk=BLOCK_K)[by, ko] + b = b.rearrange("(k bk) (n bn) -> k n bk bn", bk=BLOCK_K, bn=BLOCK_N)[ko, bx] + c = c.rearrange("(m bm) (n bn) -> m n bm bn", bm=BLOCK_M, bn=BLOCK_N)[by, bx] - # T.copy(B[ko * BLOCK_K, bx * BLOCK_N], B_shared) - kb, ibr = UOp.range(TM, 9, AxisType.LOOP), UOp.range(TN, 10, AxisType.UPCAST) - B_store = B_shared[kb*TM + ty, tx*TN + ibr].store(b[ko*BLOCK_K + kb*TM + ty, bx*BLOCK_N + tx*TN + ibr]).end(kb, ibr) + # T.copy: A_shared <- a, B_shared <- b + def with_threads(x:UOp): return x.rearrange("(tm ty) (tx tn) -> ty tx tm tn", tm=TM, tn=TN)[ty, tx] + A_shared = A_shared.after(with_threads(A_shared).store(with_threads(a))) + B_shared = B_shared.after(with_threads(B_shared).store(with_threads(b))) - # 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 is the outer compute loop (axis 11) so that for each kk we read all 8 A rows and reuse - # the B[kk] read across them -- matching tilelang's ko > kk > row > col access order exactly. - kk, ir = UOp.range(BLOCK_K, 11, AxisType.LOOP), UOp.range(TM, 12, AxisType.LOOP) + # T.gemm(A_shared, B_shared, C_local), no WMMA + kk = UOp.range(BLOCK_K, 11, AxisType.LOOP) + ir = UOp.range(TM, 12, AxisType.LOOP) jj = UOp.range(TN, 13, AxisType.UPCAST) - acc = C_loc.after(kk)[ir*TM + ty, tx*TN + jj] + A_shared[ir*TM + ty, kk].cast(dtypes.float32) * B_shared[kk, tx*TN + jj].cast(dtypes.float32) + acc = C_local.after(kk)[ir, ty, tx, jj] + A_shared[ir*TM + ty, 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[ir*TM + ty, tx*TN + jj].set(acc, end=(kk, ir, jj, ko)) + C_local = C_local[ir, ty, tx, 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(TM, 14, AxisType.LOOP), UOp.range(TN, 15, AxisType.UPCAST) - c_st = c[by*BLOCK_M + ie*TM + ty, bx*BLOCK_N + tx*TN + je].store(C_loc[ie*TM + ty, tx*TN + je].relu().cast(c.dtype)) + # c <- C_local (with relu and cast): every thread stores its shard's sub-view of the output tile + c_st = c.reshape(C_local.shape).store(C_local.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, tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=())) + # close the locals and globals + return c_st.end(tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=())) # --------------------------------------------------------------------------- # python wrapper: same signature as the tilelang function @@ -188,7 +131,8 @@ if __name__ == "__main__": b = Tensor.randn(K, N, dtype=dtype_in).contiguous() ref = (a @ b).relu().realize() - out = matmul_relu(a, b).realize() + for _ in range(10): + out = matmul_relu(a, b).realize() import numpy as np np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=1e-1, rtol=1e-2) diff --git a/test/backend/test_custom_kernel.py b/test/backend/test_custom_kernel.py index 182e619899..7e087da373 100644 --- a/test/backend/test_custom_kernel.py +++ b/test/backend/test_custom_kernel.py @@ -1,5 +1,6 @@ import unittest from tinygrad import Tensor, UOp, GlobalCounters, Context, Device +import numpy as np from tinygrad.dtype import AddrSpace, dtypes, Invalid from tinygrad.uop.ops import KernelInfo, AxisType, Ops from tinygrad.renderer.ptx import PTXRenderer @@ -527,6 +528,77 @@ class TestUnshardIndex(unittest.TestCase): with self.assertRaisesRegex(RuntimeError, "cannot shard index"): self._run(kernel, (64, 8)) +def _run_fragment_kernel(testcase, kernel, out_shape, inputs=()): + c = Tensor.empty(*out_shape) + out = Tensor.custom_kernel(c, *inputs, fxn=kernel)[0] + try: return out.numpy() + except RuntimeError as e: + if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e): + testcase.skipTest("PTX does not support dynamic register indexing") + raise + +class TestUnshardAlu(unittest.TestCase): + """Tests for ALU on (fragment) UNSHARD values in schedule/multi.py's alu_multi. + + An ALU with UNSHARD srcs lowers to per-shard ops when every src is one of: + same sharding: peel the UNSHARD, keep the layout + scalar: broadcast to every shard + whole unsharded same-shape value: takes its per-shard sub-view (shard_subview) + """ + @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges") + def test_alu_scalar_broadcast(self): + # scalar srcs broadcast to every shard: frag*2.0 where frag is 1.5 per thread -> 3.0 everywhere + def kernel(C:UOp) -> UOp: + ty = UOp.range(8, 0, AxisType.LOCAL) + # 8 values per thread, 8 threads -> 64-value full view + frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,)) + v = frag.after(frag.store(1.5)) * 2.0 + return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_scalar", opts_to_apply=())) + out = _run_fragment_kernel(self, kernel, (64,)) + np.testing.assert_allclose(out, 3.0) + + @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges") + def test_alu_whole_value_subview(self): + # UNSHARD + whole unsharded same-shape value: each shard adds its own sub-view of A. + def kernel(C:UOp, A:UOp) -> UOp: + ty = UOp.range(8, 0, AxisType.LOCAL) + frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,)) + v = frag.after(frag.store(0.0)) + A + return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_subview", opts_to_apply=())) + a = Tensor(np.arange(64, dtype=np.float32)) + out = _run_fragment_kernel(self, kernel, (64,), inputs=(a,)) + np.testing.assert_allclose(out, a.numpy(), atol=1e-4) + +class TestUnshardStore(unittest.TestCase): + """Tests for STORE of a sharded value into an unsharded dest (store_value_multi in schedule/multi.py). + + Every shard stores its value into its own contiguous sub-view of the dest, one SHRINK per sharded axis. + """ + @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges") + def test_store_unshard_value(self): + # single-axis: 8 threads each own 8 values of the 64-value output tile + def kernel(C:UOp) -> UOp: + ty = UOp.range(8, 0, AxisType.LOCAL) + frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,)) + v = frag.after(frag.store(0.0)) + 2.5 + return C.store(v).end(ty).sink(arg=KernelInfo(name="store_unshard", opts_to_apply=())) + out = _run_fragment_kernel(self, kernel, (64,)) + np.testing.assert_allclose(out, 2.5) + + @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges") + def test_store_unshard_value_2axis(self): + # two sharded axes (the gemm fragment layout): thread (ty, tx) owns the (2, 1, 1, 2) sub-view of the + # (2, 4, 2, 2) output tile; the store must SHRINK dest on both sharded axes + def kernel(C:UOp, A:UOp) -> UOp: + ty = UOp.range(4, 0, AxisType.LOCAL) + tx = UOp.range(2, 1, AxisType.LOCAL) + frag = UOp.placeholder((2, 1, 1, 2), dtypes.float32, 0, AddrSpace.REG).unshard((1, 2), (ty, tx)) + v = frag.after(frag.store(0.0)) + A + return C.store(v).end(tx, ty).sink(arg=KernelInfo(name="store_unshard_2axis", opts_to_apply=())) + a = Tensor(np.arange(32, dtype=np.float32).reshape(2, 4, 2, 2)) + out = _run_fragment_kernel(self, kernel, (2, 4, 2, 2), inputs=(a,)) + np.testing.assert_allclose(out, a.numpy(), atol=1e-4) + class TestUOpReduce(unittest.TestCase): def test_uop_sum(self): a = Tensor([1.0, 2, 3, 4, 5]) diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index 275cffc95d..440b1565f9 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -71,13 +71,29 @@ def shard_srcs(msrcs:tuple[UOp, ...], axis:int) -> list[UOp]: srcs.append(full if axis in broadcast_axes(mlb.shape, out_shape) else full._shard(src_axis, sharding_rng)) return srcs +def shard_subview(full:UOp, multi:UOp) -> UOp: + """the sub-view of an unsharded full-shape value (shape == multi.shape) that belongs to this shard: + _shard along every sharded axis (contiguous blocks, like the device path).""" + assert tuple(full.shape) == tuple(multi.shape), f"shard sub-view shape mismatch {full.shape} != {multi.shape}" + # an EXPAND of a scalar over the full shape is the same broadcast on every shard: re-expand over the shard shape + if full.op is Ops.EXPAND and full.src[0].shape == (): return full.src[0].expand(multi.src[0].shape) + for ax, rng in multi.sharding: full = full._shard(ax, rng) + return full + 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:]) + target = multis[0] + def can_handle(m:UOp) -> bool: + # same sharding (peel the UNSHARD), or a whole unsharded value of the full tile shape (takes its per-shard + # sub-view), or a broadcast scalar + if m.sharding: return m.sharding == sharding + return m.shape == () or tuple(m.shape) == tuple(target.shape) + if all(can_handle(m) for m in root.src): + # every src either has the target sharding or is whole on every shard: run the alu per-shard + srcs = [m.src[0] if m.op is Ops.UNSHARD else m if m.shape == () else shard_subview(m, target) for m in root.src] + return srcs[0].alu(root.op, *srcs[1:]).unshard(target.arg, target.src[1:]) # resharding: single-axis fallback via shard_srcs axis = root.axis assert axis is not None @@ -233,6 +249,18 @@ def copy_multi(multi:UOp, device:str | tuple[str, ...]): def store_after_multi(dest:UOp, src:UOp): return dest.after(dest.store(src.src[0])).unshard(src.arg, src.src[1:]) +def store_value_multi(dest:UOp, multi:UOp): + # storing a sharded value into an unsharded dest: every shard stores into its own sub-view of the dest + return shard_subview(dest, multi).store(multi.src[0]) + +def store_dest_multi(root:UOp, multi:UOp): + # STORE with a sharded dest: every shard stores into its own shard of the dest. + # the value is handled like in alu_multi: UNSHARD srcs peel, full-shape values take their per-shard sub-view + # (scalars arrive EXPANDed to the full shape by UOp.store's const_like, so they sub-view like everything else) + srcs = [multi.src[0]] + [x.src[0] if x.op is Ops.UNSHARD else shard_subview(x, multi) if tuple(x.shape) == tuple(multi.shape) else x + for x in root.src[1:]] + return UOp(root.op, root.dtype, tuple(srcs), root.arg) + 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.arg, multi.src[1:]) @@ -284,7 +312,8 @@ multi_pm = PatternMatcher([ UOp(root.op, root.dtype, tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src), root.arg)), (UPat((Ops.CAST, Ops.BITCAST, Ops.CONTIGUOUS, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), passthrough_multi), - # remove UNSHARD from STORE - (UPat(Ops.STORE, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root", allow_any_len=True), - lambda root,multi: UOp(root.op, root.dtype, (multi.src[0],)+tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src[1:]), root.arg)), + # STORE of a sharded value into an unsharded dest (e.g. a fragment into a full output tile) + (UPat(Ops.STORE, src=(UPat.var("dest"), UPat(Ops.UNSHARD, name="multi"))), store_value_multi), + # STORE into a sharded dest (e.g. the fragment init): every shard stores into its own shard + (UPat(Ops.STORE, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root", allow_any_len=True), store_dest_multi), ])+replace_allreduce