mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 12:56:07 +00:00
Merge branch 'master' into x86_numel
This commit is contained in:
George Hotz
@@ -20,8 +20,8 @@ def hand_spec_tc_cores():
|
||||
|
||||
gk = UOp.range(N // 8, 0, AxisType.REDUCE)
|
||||
|
||||
a_tc = UOp.vectorize(*[mat_idx(a, gx, gk, warp, i) for i in range(2)])
|
||||
b_tc = UOp.vectorize(*[mat_idx(b, gk, gy, warp, i) for i in range(2)])
|
||||
a_tc = UOp.stack(*[mat_idx(a, gx, gk, warp, i) for i in range(2)])
|
||||
b_tc = UOp.stack(*[mat_idx(b, gk, gy, warp, i) for i in range(2)])
|
||||
|
||||
acc = UOp.placeholder((2,), dtypes.float, slot=0, addrspace=AddrSpace.REG)
|
||||
acc = acc[0].set(0.0)
|
||||
@@ -30,7 +30,7 @@ def hand_spec_tc_cores():
|
||||
# TODO: make this simple
|
||||
wmma_arg = ('WMMA_8_8_8_float_float', (8, 8, 8), dtypes.float, dtypes.float, 'METAL', 32, (((3, 2),), ((3, 2),), ((3, 2),)), ())
|
||||
|
||||
acc_load = UOp.vectorize(acc.after(gk)[0], acc.after(gk)[1])
|
||||
acc_load = UOp.stack(acc.after(gk)[0], acc.after(gk)[1])
|
||||
out = UOp(Ops.WMMA, dtypes.float.vec(2), (a_tc, b_tc, acc_load), arg=wmma_arg)
|
||||
|
||||
end_loop = UOp.group(*[acc[i].store(out.index(i)) for i in range(2)]).end(gk)
|
||||
|
||||
@@ -192,7 +192,7 @@ acc = UOp.placeholder((4,), dtypes.float, 0, AddrSpace.REG)
|
||||
acc = acc[init_l:=UOp.range(4, 1)].set(0.0, end=init_l)
|
||||
|
||||
# do the wmma
|
||||
acc_load = UOp.vectorize(*[acc.after(K_loop)[i] for i in range(4)])
|
||||
acc_load = UOp.stack(*[acc.after(K_loop)[i] for i in range(4)])
|
||||
wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
|
||||
out = UOp(Ops.WMMA, dtypes.float.vec(4), (A_in, B_in, acc_load), arg=wmma_arg)
|
||||
|
||||
|
||||
@@ -84,13 +84,13 @@ class Group:
|
||||
for width in self.ker.range(c.shape[-2], track=False):
|
||||
for inner in self.ker.range(a.shape[-2], axis_type=AxisType.REDUCE, track=False):
|
||||
if a_base_shape.cols == 16:
|
||||
a_in = UOp.vectorize(*[a[height, inner, i] for i in range(4)])
|
||||
b_in = UOp.vectorize(*[b[inner, width, i] for i in range(4)])
|
||||
a_in = UOp.stack(*[a[height, inner, i] for i in range(4)])
|
||||
b_in = UOp.stack(*[b[inner, width, i] for i in range(4)])
|
||||
elif a_base_shape.cols == 32:
|
||||
a_in = UOp.vectorize(*[a[height, inner, i] for i in range(8)])
|
||||
b_in = UOp.vectorize(*[b[inner, width, i] for i in range(8)])
|
||||
a_in = UOp.stack(*[a[height, inner, i] for i in range(8)])
|
||||
b_in = UOp.stack(*[b[inner, width, i] for i in range(8)])
|
||||
else: raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}")
|
||||
d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
|
||||
d_in = UOp.stack(*[c[height, width, i] for i in range(4)])
|
||||
|
||||
out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
|
||||
c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
|
||||
@@ -114,13 +114,13 @@ class Group:
|
||||
for width in self.ker.range(c.shape[-2], track=False):
|
||||
for inner in self.ker.range(a.shape[-2], axis_type=AxisType.REDUCE, track=False):
|
||||
if a_base_shape.cols == 16:
|
||||
a_in = UOp.vectorize(*[a[height, inner, i] for i in range(4)])
|
||||
b_in = UOp.vectorize(*[b[width, inner, i] for i in range(4)])
|
||||
a_in = UOp.stack(*[a[height, inner, i] for i in range(4)])
|
||||
b_in = UOp.stack(*[b[width, inner, i] for i in range(4)])
|
||||
elif a_base_shape.cols == 32:
|
||||
a_in = UOp.vectorize(*[a[height, inner, i] for i in range(8)])
|
||||
b_in = UOp.vectorize(*[b[width, inner, i] for i in range(8)])
|
||||
a_in = UOp.stack(*[a[height, inner, i] for i in range(8)])
|
||||
b_in = UOp.stack(*[b[width, inner, i] for i in range(8)])
|
||||
else: raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}")
|
||||
d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
|
||||
d_in = UOp.stack(*[c[height, width, i] for i in range(4)])
|
||||
|
||||
out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
|
||||
c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
|
||||
@@ -144,13 +144,13 @@ class Group:
|
||||
for width in self.ker.range(c.shape[-2], track=False):
|
||||
for inner in self.ker.range(a.shape[-3], axis_type=AxisType.REDUCE, track=False):
|
||||
if a_base_shape.cols == 16:
|
||||
a_in = UOp.vectorize(*[a[inner, height, i] for i in range(4)])
|
||||
b_in = UOp.vectorize(*[b[inner, width, i] for i in range(4)])
|
||||
a_in = UOp.stack(*[a[inner, height, i] for i in range(4)])
|
||||
b_in = UOp.stack(*[b[inner, width, i] for i in range(4)])
|
||||
elif a_base_shape.cols == 32:
|
||||
a_in = UOp.vectorize(*[a[inner, height, i] for i in range(8)])
|
||||
b_in = UOp.vectorize(*[b[inner, width, i] for i in range(8)])
|
||||
a_in = UOp.stack(*[a[inner, height, i] for i in range(8)])
|
||||
b_in = UOp.stack(*[b[inner, width, i] for i in range(8)])
|
||||
else: raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}")
|
||||
d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
|
||||
d_in = UOp.stack(*[c[height, width, i] for i in range(4)])
|
||||
|
||||
out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
|
||||
c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
|
||||
@@ -174,13 +174,13 @@ class Group:
|
||||
for width in self.ker.range(c.shape[-2], track=False):
|
||||
for inner in self.ker.range(a.shape[-3], axis_type=AxisType.REDUCE, track=False):
|
||||
if a_base_shape.cols == 16:
|
||||
a_in = UOp.vectorize(*[a[inner, height, i] for i in range(4)])
|
||||
b_in = UOp.vectorize(*[b[width, inner, i] for i in range(4)])
|
||||
a_in = UOp.stack(*[a[inner, height, i] for i in range(4)])
|
||||
b_in = UOp.stack(*[b[width, inner, i] for i in range(4)])
|
||||
elif a_base_shape.cols == 32:
|
||||
a_in = UOp.vectorize(*[a[inner, height, i] for i in range(8)])
|
||||
b_in = UOp.vectorize(*[b[width, inner, i] for i in range(8)])
|
||||
a_in = UOp.stack(*[a[inner, height, i] for i in range(8)])
|
||||
b_in = UOp.stack(*[b[width, inner, i] for i in range(8)])
|
||||
else: raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}")
|
||||
d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
|
||||
d_in = UOp.stack(*[c[height, width, i] for i in range(4)])
|
||||
|
||||
out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
|
||||
c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
|
||||
|
||||
@@ -39,8 +39,8 @@ class TestIselX86(unittest.TestCase):
|
||||
c = UOp.variable("c", 0, 0, dtypes.float32.vec(4))
|
||||
d = UOp.variable("e", 0, 0, dtypes.float32)
|
||||
|
||||
valid = [UOp.vectorize(lane(a, 0), lane(b, 1), lane(a, 2), lane(b, 3)),
|
||||
UOp.vectorize(lane(a, 3), lane(b, 2), lane(c, 1), d)]
|
||||
valid = [UOp.stack(lane(a, 0), lane(b, 1), lane(a, 2), lane(b, 3)),
|
||||
UOp.stack(lane(a, 3), lane(b, 2), lane(c, 1), d)]
|
||||
for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VINSERTPS)
|
||||
|
||||
# complex address is [base + index*scale + displacement]
|
||||
|
||||
@@ -127,6 +127,15 @@ class TestMultiTensor(unittest.TestCase):
|
||||
fn = f(n)
|
||||
np.testing.assert_allclose(fX.numpy(), fn, rtol=1e-6, atol=1e-6)
|
||||
|
||||
def test_stack(self):
|
||||
X = Tensor.rand(4, 4).shard_(devices_2, 0)
|
||||
Y = Tensor.rand(4, 4).shard_(devices_2, 0)
|
||||
Z = Tensor.rand(4, 4).shard_(devices_2, 1) # mismatched shard axis gets resharded
|
||||
for dim in (0, 1):
|
||||
np.testing.assert_allclose(Tensor.stack(X, Y, Z, dim=dim).numpy(), np.stack([X.numpy(), Y.numpy(), Z.numpy()], axis=dim))
|
||||
grad = Tensor.stack(X, Y).sum().gradient(X)[0]
|
||||
np.testing.assert_allclose(grad.numpy(), 1)
|
||||
|
||||
def test_allreduce_naive(self):
|
||||
with Context(RING=0):
|
||||
a,b = _test_allreduce(Tensor.rand(256, 256))
|
||||
|
||||
@@ -188,7 +188,7 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase):
|
||||
def test_gep_tuple_extraction(self):
|
||||
# GEP on a vector dtype to extract multiple elements as a vector
|
||||
base_vector = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
|
||||
self.assertEqual(list(apply_rewrite_values(UOp.vectorize(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0])
|
||||
self.assertEqual(list(apply_rewrite_values(UOp.stack(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0])
|
||||
|
||||
def test_gep_on_const_stack(self):
|
||||
# GEP on a const STACK to extract a single element
|
||||
@@ -198,7 +198,7 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase):
|
||||
def test_gep_tuple_on_const_stack(self):
|
||||
# GEP on a const STACK using a tuple to extract multiple elements
|
||||
const_stack = UOp.const(dtypes.float32, (7.0, 8.0, 9.0, 10.0))
|
||||
self.assertEqual(list(apply_rewrite_values(UOp.vectorize(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])
|
||||
self.assertEqual(list(apply_rewrite_values(UOp.stack(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])
|
||||
|
||||
def test_vectorize_multiple_elements(self):
|
||||
# Vectorizing multiple elements using GEP
|
||||
|
||||
@@ -376,6 +376,13 @@ class TestTensorUOpStack(unittest.TestCase):
|
||||
def test_stack_dim1(self): _check(self, _t(2, 3), lambda x: x.stack(x, dim=1))
|
||||
def test_stack_3tensors(self): _check(self, _t(2, 3), lambda x: x.stack(x, x, dim=0))
|
||||
def test_stack_new_last(self): _check(self, _t(2, 3), lambda x: x.stack(x, dim=-1))
|
||||
def test_stack_mixed_dtype(self):
|
||||
w = _t(2, 3).float()
|
||||
_check(self, _t(2, 3), lambda x: x.stack(w if isinstance(x, Tensor) else w.uop))
|
||||
self.assertIs(_t(2, 3).uop.stack(w.uop).dtype, dtypes.float32)
|
||||
def test_stack_index_dtype(self):
|
||||
# index is outside the promotion lattice, equal dtypes bypass promotion
|
||||
self.assertEqual(UOp.const(dtypes.index, 1).stack(UOp.const(dtypes.index, 2)).shape, (2,))
|
||||
|
||||
class TestTensorUOpConv2d(unittest.TestCase):
|
||||
def test_conv2d_basic(self):
|
||||
|
||||
@@ -268,12 +268,12 @@ class TestViz(unittest.TestCase):
|
||||
def test_stack_movement_not_folded_unless_all_const(self):
|
||||
a = UOp.variable("a", 0, 10, dtype=dtypes.int)
|
||||
c = UOp.const(dtypes.int, 1)
|
||||
stack = a.vectorize(c)
|
||||
stack = a.stack(c)
|
||||
reshaped = stack.reshape((1, 2))
|
||||
graph = uop_to_json(VizData(), reshaped)
|
||||
self.assertFalse(graph[id(stack)]["exclude"])
|
||||
|
||||
const_stack = c.vectorize(UOp.const(dtypes.int, 2))
|
||||
const_stack = c.stack(UOp.const(dtypes.int, 2))
|
||||
const_reshaped = const_stack.reshape((1, 2))
|
||||
const_graph = uop_to_json(VizData(), const_reshaped)
|
||||
self.assertTrue(const_graph[id(const_stack)]["exclude"])
|
||||
|
||||
@@ -111,7 +111,7 @@ def broadcast_and_devec_wmma(b:UOp):
|
||||
for idx in itertools.product(*[range(i) for i in b.shape[:-1]]):
|
||||
idx_c = [UOp.const(dtypes.index, i) for i in idx]
|
||||
src.append(b.replace(src=tuple([x.index(*idx_c) for x in src_reshaped])))
|
||||
return UOp.vectorize(*src).reshape(b.shape)
|
||||
return UOp.stack(*src).reshape(b.shape)
|
||||
|
||||
pm_wmma_add = PatternMatcher([
|
||||
(UPat(Ops.WMMA, name="wmma") + UPat.var("add"),
|
||||
@@ -136,7 +136,7 @@ def do_devectorize(b:UOp):
|
||||
for idx in itertools.product(*[range(x) for x in b.shape]):
|
||||
idx_c = [UOp.const(dtypes.index, i) for i in idx]
|
||||
src.append(b.replace(src=tuple([x.index(*idx_c) for x in b.src])))
|
||||
return UOp.vectorize(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
|
||||
return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
|
||||
|
||||
def do_stack_wmma(u:UOp):
|
||||
if all(x.op in (Ops.STACK, Ops.WMMA) for x in u.src): return None
|
||||
@@ -144,7 +144,7 @@ def do_stack_wmma(u:UOp):
|
||||
src = []
|
||||
for b in u.src:
|
||||
if b.op != Ops.STACK:
|
||||
src.append(UOp._stack(*[b.index(UOp.const(dtypes.index, i)) for i in range(b.max_numel())]))
|
||||
src.append(UOp.stack(*[b.index(UOp.const(dtypes.index, i)) for i in range(b.max_numel())]))
|
||||
else:
|
||||
src.append(b)
|
||||
return u.replace(src=tuple(src))
|
||||
@@ -163,7 +163,7 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
|
||||
(UPat(Ops.WMMA, name="u"), do_stack_wmma),
|
||||
# stacked INDEX is many INDEX
|
||||
(UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.STACK, name="s"))),
|
||||
lambda b,s: UOp.vectorize(*[b.index(u) for u in s.src])),
|
||||
lambda b,s: UOp.stack(*[b.index(u) for u in s.src])),
|
||||
# INDEX into RESHAPE moves the RESHAPE
|
||||
(UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.RESHAPE, name="s"))),
|
||||
lambda b,s: b.index(s.src[0]).reshape(s.shape)),
|
||||
@@ -173,7 +173,7 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
|
||||
(UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(UOp.const(dtypes.index, 0)) if x.marg == () and x.src[0].shape == (1,) else None),
|
||||
# EXPAND on scalar -> STACK
|
||||
(UPat(Ops.EXPAND, src=(UPat.var("x"), UPat()), name="out"),
|
||||
lambda x,out: UOp.vectorize(*([x]*out.max_numel())) if x.shape == () and out.shape == (out.max_numel(),) else None),
|
||||
lambda x,out: UOp.stack(*([x]*out.max_numel())) if x.shape == () and out.shape == (out.max_numel(),) else None),
|
||||
# INDEX on INDEX is INDEX
|
||||
(UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="idx1", allow_any_len=True),), allow_any_len=True, name="idx2"),
|
||||
lambda idx1, idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:])),
|
||||
|
||||
@@ -41,7 +41,7 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
|
||||
def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|None:
|
||||
if not is_image_shape(buf._shape): return None
|
||||
if idx_x.dtype != idx_y.dtype: idx_x, idx_y = idx_x.cast(dtypes.int), idx_y.cast(dtypes.int)
|
||||
start_idx = idx_x._stack(idx_y)
|
||||
start_idx = idx_x.stack(idx_y)
|
||||
idx = uop_given_valid(valid, start_idx)
|
||||
drop_stmt = _drop_valid_stmts(valid, idx, buf._shape[0], buf._shape[1])
|
||||
|
||||
@@ -74,7 +74,7 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
|
||||
# search for dims that drop the most valid statements
|
||||
best_drop, cands = -1, []
|
||||
for ch, cw in [shapes[buf.arg.slot]] if buf.arg.slot in shapes else image_valid_dims(buf.dtype, buf.max_numel(), ren.target.arch):
|
||||
cidx = uop_given_valid(valid, ((x//4)%cw)._stack(x//(4*cw)))
|
||||
cidx = uop_given_valid(valid, ((x//4)%cw).stack(x//(4*cw)))
|
||||
dropped = len(_drop_valid_stmts(valid, cidx, ch, cw))
|
||||
if dropped > best_drop: best_drop, cands = dropped, [(ch, cw, cidx)]
|
||||
elif dropped == best_drop: cands.append((ch, cw, cidx))
|
||||
@@ -152,7 +152,7 @@ def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp:
|
||||
for i,g in enumerate(grp):
|
||||
assert len(offsets[g]) == 1, f"attempting multiple stores: {len(offsets[g])}"
|
||||
datas.append(offsets[g][0].src[1])
|
||||
store = idx.store(UOp._stack(*datas) if len(datas) > 1 else datas[0])
|
||||
store = idx.store(UOp.stack(*datas) if len(datas) > 1 else datas[0])
|
||||
for i,g in enumerate(grp): replacements[offsets[g][0]] = store
|
||||
else:
|
||||
ld = idx.load()
|
||||
|
||||
@@ -74,6 +74,7 @@ pm_gradient = PatternMatcher([
|
||||
(UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[0]-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
|
||||
(UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.marg)),)),
|
||||
(UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip([i for i,x in enumerate(ret.marg) if x]),)),
|
||||
(UPat(Ops.STACK, name="ret"), lambda ctx, ret: tuple(ctx[i] for i in range(len(ret.src)))),
|
||||
(UPat(Ops.COPY, name="ret"), lambda ctx, ret: (ctx.copy_to_device(ret.src[0].device),)),
|
||||
(UPat(Ops.MULTI, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
|
||||
(UPat(Ops.TUPLE), lambda ctx: ctx.src),
|
||||
|
||||
@@ -241,6 +241,24 @@ class MovementMixin:
|
||||
flip_arg = tuple([i in axis_arg for i in range(len(self.shape))])
|
||||
return self._mop(Ops.FLIP, arg=flip_arg) if any(flip_arg) else self
|
||||
|
||||
def stack(self, *args: Self, dim: int = 0) -> Self:
|
||||
"""
|
||||
Concatenates self with other tensors in `args` along a new dimension specified by `dim`.
|
||||
|
||||
```python exec="true" source="above" session="tensor" result="python"
|
||||
t0, t1, t2 = Tensor([1, 2]), Tensor([3, 4]), Tensor([5, 6])
|
||||
print(t0.stack(t1, t2, dim=0).numpy())
|
||||
```
|
||||
```python exec="true" source="above" session="tensor" result="python"
|
||||
print(t0.stack(t1, t2, dim=1).numpy())
|
||||
```
|
||||
"""
|
||||
tensors = argfix(self, *args)
|
||||
dim = tensors[0]._resolve_dim(dim, extra=True)
|
||||
assert all(t.shape == tensors[0].shape for t in tensors), f"all shapes must match for stack, got {[t.shape for t in tensors]}"
|
||||
ret = tensors[0]._mop(Ops.STACK, arg=tuple(t._uop for t in tensors[1:]))
|
||||
return ret if dim == 0 else ret.permute(tuple(range(1, dim+1)) + (0,) + tuple(range(dim+1, ret.ndim)))
|
||||
|
||||
# **** high level ****
|
||||
|
||||
def shrink_to(self, shape, *args) -> Self:
|
||||
|
||||
@@ -733,22 +733,6 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
|
||||
padded = [t.pad(tuple((dim_cumsum[i], dim_cumsum[-1]-dim_cumsum[i+1]) if j==dim else None for j in range(t.ndim))) for i,t in enumerate(tensors)]
|
||||
return padded[0].usum(*padded[1:])
|
||||
|
||||
def stack(self, *args:Self, dim:int=0) -> Self:
|
||||
"""
|
||||
Concatenates self with other tensors in `args` along a new dimension specified by `dim`.
|
||||
|
||||
```python exec="true" source="above" session="tensor" result="python"
|
||||
t0, t1, t2 = Tensor([1, 2]), Tensor([3, 4]), Tensor([5, 6])
|
||||
print(t0.stack(t1, t2, dim=0).numpy())
|
||||
```
|
||||
```python exec="true" source="above" session="tensor" result="python"
|
||||
print(t0.stack(t1, t2, dim=1).numpy())
|
||||
```
|
||||
"""
|
||||
# checks for shapes and number of dimensions delegated to cat
|
||||
unsqueezed = [t.unsqueeze(dim) for t in argfix(self, *args)]
|
||||
return unsqueezed[0].cat(*unsqueezed[1:], dim=dim)
|
||||
|
||||
def _cumalu(self, axis:int, op:Ops) -> Self:
|
||||
assert self.shape[axis] != 0 and op in (Ops.ADD, Ops.MAX, Ops.MUL)
|
||||
pads = (None,)*(self.ndim-1) + ((self.shape[axis]-1, 0),)
|
||||
|
||||
@@ -53,8 +53,7 @@ class IndexingContext:
|
||||
# if a range has a 1 src, it's the same as UOp.const(dtypes.index, 0)
|
||||
return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0)
|
||||
|
||||
def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
|
||||
if x.op in {Ops.STAGE, Ops.INDEX}: return None
|
||||
def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
|
||||
new_srcs = []
|
||||
for i, s in enumerate(x.src):
|
||||
new_src = s
|
||||
@@ -78,7 +77,11 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
|
||||
new_src = UOp(Ops.STAGE, src=(new_src,)+closed_ranges, arg=opts)
|
||||
if x in ctx.range_map: new_src = new_src.index(*[r for i,r in enumerate(ctx.range_map[x][0]) if i in realized_ranges])
|
||||
new_srcs.append(new_src)
|
||||
return x.replace(src=tuple(new_srcs))
|
||||
return new_srcs
|
||||
|
||||
def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
|
||||
if x.op in {Ops.STAGE, Ops.INDEX}: return None
|
||||
return x.replace(src=tuple(create_bufferize_and_index_srcs(ctx, x)))
|
||||
|
||||
def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp):
|
||||
if x not in ctx.range_map: return None
|
||||
@@ -93,6 +96,16 @@ def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
|
||||
new_ranges = list(ctx.range_map[x][0][:x.arg[1]])
|
||||
return UOp(Ops.REDUCE, src=(bx.src[0],)+tuple(new_ranges), arg=(x.arg[0], 0))
|
||||
|
||||
def convert_stack_to_where(ctx:IndexingContext, x:UOp):
|
||||
# only data STACKs: shape tuple STACKs aren't in range_map, the empty shape tuple is void
|
||||
if x not in ctx.range_map or x.dtype == dtypes.void: return None
|
||||
# use the src list directly, a transient STACK of mid-rangeify srcs violates the spec shape rule
|
||||
srcs = create_bufferize_and_index_srcs(ctx, x)
|
||||
r0 = ctx.range_map[x][1][0]
|
||||
ret = srcs[-1]
|
||||
for k in range(len(srcs)-2, -1, -1): ret = r0.eq(k).where(srcs[k], ret)
|
||||
return ret
|
||||
|
||||
def remove_movement_op_after_rangeify(ctx:IndexingContext, x:UOp):
|
||||
if x in ctx.range_map or x.src[0].op is Ops.INDEX: return x.src[0]
|
||||
|
||||
@@ -101,6 +114,8 @@ pm_apply_rangeify = PatternMatcher([
|
||||
(UPat(Ops.REDUCE, name="x"), convert_reduce_to_reduce_with_ranges),
|
||||
# PAD -> WHERE
|
||||
(UPat(Ops.PAD, name="x"), convert_pad_to_where_to_keep_behavior_local),
|
||||
# STACK -> WHERE select on the leading range
|
||||
(UPat(Ops.STACK, name="x"), convert_stack_to_where),
|
||||
# finally, apply_rangeify
|
||||
(UPat(GroupOp.All, name="x"), create_bufferize_and_index_based_on_ranges),
|
||||
# remove movement op
|
||||
@@ -245,6 +260,8 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
|
||||
|
||||
# apply movement ops
|
||||
if x.op in GroupOp.Movement: rngs = apply_movement_op(x.op, x.src[0].shape, x.marg, rngs)
|
||||
# STACK: the leading range selects the src, srcs get the trailing ranges
|
||||
if x.op is Ops.STACK: rngs = out_rngs[1:]
|
||||
# if the EXPAND is used to inject a range, we don't mark it as ending_ranges. otherwise we do.
|
||||
# NOTE: this doesn't actually always end a range, but this is why convs are realized, so for now we need it
|
||||
if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape):
|
||||
|
||||
@@ -41,13 +41,11 @@ if not getenv("LATE_ALLREDUCE", 1): replace_allreduce = _early_allreduce + repla
|
||||
|
||||
# ***** multi functions *****
|
||||
|
||||
def alu_multi(root:UOp):
|
||||
msrcs = root.src
|
||||
def shard_srcs(msrcs:tuple[UOp, ...], axis:int) -> list[UOp]:
|
||||
# normalize srcs to local shards on axis
|
||||
devices = [x.device for x in msrcs if x.device is not None]
|
||||
assert all_same(devices), f"all buffers must have the same device {devices}"
|
||||
dcount = len(devices[0])
|
||||
axis = root.axis
|
||||
assert axis is not None
|
||||
|
||||
srcs:list[UOp] = []
|
||||
for mlb in msrcs:
|
||||
@@ -63,6 +61,12 @@ def alu_multi(root:UOp):
|
||||
else:
|
||||
# axis mismatch, copy to all devices, and shard it correctly
|
||||
srcs.append(copy_multi(mlb, mlb.device)._shard(axis, dcount))
|
||||
return srcs
|
||||
|
||||
def alu_multi(root:UOp):
|
||||
axis = root.axis
|
||||
assert axis is not None
|
||||
srcs = shard_srcs(root.src, axis)
|
||||
return srcs[0].alu(root.op, *srcs[1:]).multi(axis)
|
||||
|
||||
def reduce_multi(root:UOp, multi:UOp):
|
||||
@@ -112,6 +116,12 @@ 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]).multi(multi.axis)
|
||||
|
||||
def stack_multi(root:UOp):
|
||||
# STACK adds a leading axis: srcs are sharded one axis below the output
|
||||
axis = root.axis
|
||||
assert axis is not None
|
||||
return UOp(Ops.STACK, src=tuple(shard_srcs(root.src, axis-1))).multi(axis)
|
||||
|
||||
def copy_multi(multi:UOp, device:str | tuple[str, ...]):
|
||||
assert multi.axis is not None, "all multi ops have axis"
|
||||
if isinstance(device, str):
|
||||
@@ -151,6 +161,7 @@ multi_pm = PatternMatcher([
|
||||
(UPat(Ops.SHRINK, src=(UPat(Ops.MULTI, name="multi"), UPat(), UPat()), name="root"), shrink_multi),
|
||||
(UPat(Ops.PERMUTE, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), permute_multi),
|
||||
(UPat(Ops.FLIP, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), flip_multi),
|
||||
(UPat(Ops.STACK, name="root", custom_early_reject=set([Ops.MULTI])), stack_multi),
|
||||
(UPat(Ops.AFTER, src=(UPat(Ops.MULTI), UPat(Ops.STORE, src=(UPat(Ops.MULTI, name="dest"), UPat(Ops.MULTI, name="src"))))), store_after_multi),
|
||||
(UPat(Ops.COPY, src=(UPat(Ops.MULTI, name="multi"),), name="copy"), lambda multi,copy: copy_multi(multi, copy.arg)),
|
||||
(UPat(Ops.ALLREDUCE, src=(UPat(Ops.MULTI, name="multi"),), name="red"),
|
||||
|
||||
+12
-10
@@ -125,8 +125,8 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
|
||||
return src[1].dtype
|
||||
case Ops.STACK:
|
||||
if len(src) == 0: return dtypes.void
|
||||
if not all_same([x.dtype for x in src]): raise RuntimeError("stack must have matching dtype")
|
||||
return src[0].dtype
|
||||
if all_same(dts:=[x.dtype for x in src]): return dts[0]
|
||||
return least_upper_dtype(*dts)
|
||||
case Ops.BIND:
|
||||
assert src[0].dtype == src[1].dtype, f"bind dtype mismatch {src[0].dtype} != {src[1].dtype}"
|
||||
return src[0].dtype
|
||||
@@ -523,10 +523,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
def group(*srcs:UOp|None): # pylint: disable=no-self-argument
|
||||
if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
|
||||
return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None]))
|
||||
def _stack(self, *srcs):
|
||||
# TODO: this should become the real stack
|
||||
return UOp(Ops.STACK, src=(self,)+srcs)
|
||||
def vectorize(self, *srcs): return self._stack(*srcs)
|
||||
def index(self, *srcs:UOp|int|None, **kwargs):
|
||||
new_srcs: list[UOp] = [UOp.const(dtypes.index, x) if isinstance(x, int) else x for x in srcs if x is not None]
|
||||
if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].arg]
|
||||
@@ -584,7 +580,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
def ins(self, arg, **kwargs): return UOp(Ops.INS, kwargs.pop("dtype", self.dtype), kwargs.pop("src", self.src), arg, kwargs.pop("tag", self.tag))
|
||||
def contract(self, *rngs:UOp):
|
||||
assert all(x.arg[-1] == AxisType.UPCAST for x in rngs), "all contract ranges must be upcast"
|
||||
return UOp.vectorize(*[self.substitute(dict(zip(rngs, [r.const_like(i) for r,i in zip(rngs, idx)])))
|
||||
return UOp.stack(*[self.substitute(dict(zip(rngs, [r.const_like(i) for r,i in zip(rngs, idx)])))
|
||||
for idx in itertools.product(*[range(int(r.vmax)+1) for r in rngs])])
|
||||
@staticmethod
|
||||
def wmma(a:UOp, b:UOp, acc:UOp, arg:tuple[tuple[int, int, int], str, int]):
|
||||
@@ -606,7 +602,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
# NOTE: it always has to be STACK now, even if they are all the same
|
||||
if isinstance(b, tuple):
|
||||
stk = [UOp(Ops.CONST, dtype, arg=dtype.const(c), src=()) for c in b]
|
||||
ret = UOp.vectorize(*stk)
|
||||
ret = UOp.stack(*stk)
|
||||
else:
|
||||
ret = UOp(Ops.CONST, dtype, arg=dtype.const(b), src=())
|
||||
return ret._mop(Ops.EXPAND, arg=shape) if shape is not None and shape != () and ret.shape != shape else ret
|
||||
@@ -631,11 +627,11 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
return cond.where(self, self.const_like(Invalid))
|
||||
def get_idx(self) -> UOp:
|
||||
assert dtypes.is_int(self.dtype), "Can only call get_idx on index dtype"
|
||||
if self.op is Ops.STACK: return UOp.vectorize(*(x.get_idx() for x in self.src))
|
||||
if self.op is Ops.STACK: return UOp.stack(*(x.get_idx() for x in self.src))
|
||||
return self.src[1] if self.op is Ops.WHERE and self.src[2].arg is Invalid else self
|
||||
def get_valid(self) -> UOp:
|
||||
assert dtypes.is_int(self.dtype), "Can only call get_valid on index dtype"
|
||||
if self.op is Ops.STACK: return UOp.vectorize(*(x.get_valid() for x in self.src))
|
||||
if self.op is Ops.STACK: return UOp.stack(*(x.get_valid() for x in self.src))
|
||||
return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(dtypes.bool, self.arg is not Invalid)
|
||||
def reduce(self, *src:UOp, **kwargs):
|
||||
arg = kwargs.pop('arg', None)
|
||||
@@ -682,6 +678,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
if self.op is Ops.PARAM: return self.arg.axis
|
||||
# NOTE: they all have to share an axis, we always choose [-1]
|
||||
if self.op in GroupOp.ALU: return axes[-1] if (axes := dedup([x.axis for x in self.src if x.axis is not None])) else None
|
||||
# STACK adds a leading axis
|
||||
if self.op is Ops.STACK: return axes[-1]+1 if (axes := dedup([x.axis for x in self.src if x.axis is not None])) else None
|
||||
if len(self.src) == 0: return None
|
||||
src_axis = self.src[0].axis
|
||||
if self.op is Ops.SHRINK and src_axis is not None and self.marg[src_axis] != (0, self.src[0].shape[src_axis]):
|
||||
@@ -758,6 +756,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
case Ops.RESHAPE | Ops.EXPAND: src_args = [arg]
|
||||
case Ops.PAD | Ops.SHRINK: src_args = list(zip(*arg))
|
||||
case Ops.PERMUTE | Ops.FLIP: src_args = []
|
||||
case Ops.STACK:
|
||||
# arg is the other srcs; all are cast to the promoted dtype, spec requires STACK srcs to match its dtype
|
||||
srcs = (self,)+tuple(arg)
|
||||
return UOp(Ops.STACK, src=tuple(u.cast(dtype_from_uop(Ops.STACK, srcs, None)) for u in srcs))
|
||||
case _: raise RuntimeError(f"{op} is not a MovementOp")
|
||||
usrcs = [shape_to_shape_arg(arg) for arg in src_args]
|
||||
if len(usrcs) == 0: return UOp(op, src=(self,), arg=arg)
|
||||
|
||||
@@ -58,7 +58,8 @@ spec_shared = PatternMatcher([
|
||||
|
||||
# STACK is everywhere too
|
||||
(UPat(Ops.STACK, dtype=dtypes.void, src=()), lambda: True),
|
||||
(UPat(Ops.STACK, src=(UPat(),), allow_any_len=True, name="s"), lambda s: all_same([x.shape for x in s.src])),
|
||||
(UPat(Ops.STACK, src=(UPat(),), allow_any_len=True, name="s"),
|
||||
lambda s: all_same([x.shape for x in s.src]) and all(x.dtype == s.dtype for x in s.src)),
|
||||
|
||||
# ALUs: most ALUs have all matching dtypes, except CMPLT, CMPNE, and WHERE
|
||||
(UPat(Ops.WHERE, name="w", src=(UPat(dtype=dtypes.bool), UPat.var("x"), UPat.var("y"))), lambda w,x,y: w.dtype == x.dtype == y.dtype),
|
||||
|
||||
Reference in New Issue
Block a user