forked from tinygrad/tinygrad
support 2d on UNSHARD (kimi) (#17285)
* support 2d on UNSHARD * fixes * Fix test and spec * single barrier * 2d sharding works for devices too * cleanups * no _rewrap
This commit is contained in:
George Hotz
+62
-54
@@ -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()
|
||||
|
||||
Binary file not shown.
+6
-6
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
|
||||
+117
-55
@@ -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),
|
||||
|
||||
+20
-10
@@ -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")
|
||||
|
||||
@@ -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)),
|
||||
|
||||
|
||||
Reference in New Issue
Block a user