remove custom kernel contiguous

extra/gemm/metal_uop_matmul.py

@@ -20,8 +20,8 @@ def hand_spec_tc_cores():
 
   gk = UOp.range(N // 8, 0, AxisType.REDUCE)
 
-  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)])
+  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)])
 
   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.stack(acc.after(gk)[0], acc.after(gk)[1])
+  acc_load = UOp.vectorize(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.gep(i)) for i in range(2)]).end(gk)

extra/thunder/tiny/tk/group.py

@@ -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.stack(*[a[height, inner, i] for i in range(4)])
-            b_in = UOp.stack(*[b[inner, width, i] for i in range(4)])
+            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)])
           elif a_base_shape.cols == 32:
-            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)])
+            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)])
           else: raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}")
-          d_in = UOp.stack(*[c[height, width, i] for i in range(4)])
+          d_in = UOp.vectorize(*[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.gep(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.stack(*[a[height, inner, i] for i in range(4)])
-            b_in = UOp.stack(*[b[width, inner, i] for i in range(4)])
+            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)])
           elif a_base_shape.cols == 32:
-            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)])
+            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)])
           else: raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}")
-          d_in = UOp.stack(*[c[height, width, i] for i in range(4)])
+          d_in = UOp.vectorize(*[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.gep(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.stack(*[a[inner, height, i] for i in range(4)])
-            b_in = UOp.stack(*[b[inner, width, i] for i in range(4)])
+            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)])
           elif a_base_shape.cols == 32:
-            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)])
+            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)])
           else: raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}")
-          d_in = UOp.stack(*[c[height, width, i] for i in range(4)])
+          d_in = UOp.vectorize(*[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.gep(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.stack(*[a[inner, height, i] for i in range(4)])
-            b_in = UOp.stack(*[b[width, inner, i] for i in range(4)])
+            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)])
           elif a_base_shape.cols == 32:
-            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)])
+            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)])
           else: raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}")
-          d_in = UOp.stack(*[c[height, width, i] for i in range(4)])
+          d_in = UOp.vectorize(*[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.gep(i)) for i in range(4)]

test/backend/test_custom_kernel.py

@@ -321,7 +321,6 @@ class TestCustomKernel(unittest.TestCase):
     self.assertEqual(GlobalCounters.kernel_count, 2)
     self.assertEqual(z.tolist(), x.add(2).tolist())
 
-  @unittest.expectedFailure
   def test_custom_kernel_sched_copy(self): self.test_custom_kernel_sched(use_custom=True)
 
 class TestUOpReduce(unittest.TestCase):

test/backend/test_randomness.py

@@ -361,7 +361,7 @@ class TestRandomness(unittest.TestCase):
     _check_with_torch(w=[0.231, 0., 1., 0.5], num_samples=300, replacement=True)
     _check_with_torch(w=[[0.2, 0.8]], num_samples=300, replacement=True)  # 2D but only 1 row
     _check_with_torch(w=[[0.453, 0., 1., 0.81], [0.1, 0.8, 0., 0.1]], num_samples=300, replacement=True)
-    # no-replacement
+    # no-replacement isn't supported, unless taking only one sample
     w = [0.1, 0.9]
     self.assertRaises(AssertionError, lambda: Tensor(w).multinomial(100, replacement=False))
 
@@ -372,23 +372,6 @@ class TestRandomness(unittest.TestCase):
     torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(1000)]
     self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_samples), lambda _: torch.tensor(torch_samples)))
 
-    w = list(range(32))
-    s1 = Tensor(w).multinomial(5, replacement=False).numpy()
-    self.assertEqual(len(set(s1.tolist())), 5)
-    s2 = Tensor(w).multinomial(5, replacement=False).numpy()
-    self.assertFalse(np.array_equal(s1, s2))
-    full = Tensor(w).multinomial(len(w), replacement=False).numpy()
-    self.assertEqual(sorted(full.tolist()), w)
-
-    w = [0.1, 0.2, 0.3, 0.4]
-    @TinyJit
-    def sample_three(): return Tensor(w).multinomial(3, replacement=False).realize()
-
-    tiny_draws = np.array([sample_three().numpy() for _ in range(1000)])
-    torch_draws = np.array([torch.tensor(w).multinomial(3, replacement=False).numpy() for _ in range(1000)])
-    for pos in range(3):
-      self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_draws[:, pos]), lambda _: torch.tensor(torch_draws[:, pos])))
-
   @unittest.skip("this test is flaky")
   def test_multinomial_counterexample(self):
     tiny_res = Tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True)

tinygrad/codegen/late/devectorizer.py

@@ -68,7 +68,7 @@ def expand_index(buf:UOp, vec:UOp):
     # search for dims that drop the most valid statements
     best_drop, cands = -1, []
     for ch, cw in ImageDType.valid_dims(dt):
-      if (dropped:=len(_drop_valid_stmts(valid, cidx:=uop_given_valid(valid, UOp.stack((x//4)%cw, x//(4*cw))), ch, cw))) > best_drop:
+      if (dropped:=len(_drop_valid_stmts(valid, cidx:=uop_given_valid(valid, UOp.vectorize((x//4)%cw, x//(4*cw))), ch, cw))) > best_drop:
         best_drop, cands = dropped, [(ch, cw, cidx)]
       elif dropped == best_drop: cands.append((ch, cw, cidx))
     # and tiebreak with indexing complexity (ie. number of nodes)
@@ -197,9 +197,8 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
   return UOp(Ops.VCAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)
 
 def get_image_idx(idx:UOp, width:int):
-  x, valid = idx.src[1].get_idx(), idx.src[1].get_valid()
-  idx_x, idx_y = (x // 4) % width, x // (4*width)
-  return idx.replace(src=(idx.src[0], UOp.stack(idx_x, idx_y).valid(valid)))
+  oidx = UOp(Ops.STACK, dtypes.weakint.vec(2), (((x:=idx.src[1].get_idx()) // 4) % width, (x // (4*width))))
+  return idx.replace(src=(idx.src[0], oidx.valid(idx.src[1].get_valid())))
 
 def image_fixup(ls:UOp):
   # normal image load or store, with the CAST from expand_index
@@ -386,7 +385,7 @@ def make_image(ls, buf, off):
   if (vcount:=buf.dtype.vcount) != 1: buf = buf.src[0]
   if buf.op == Ops.PARAM and not isinstance(dt:=buf.dtype, ImageDType) and (dims:=ImageDType.valid_dims(dt)):
     buf = buf.replace(dtype=(dtypes.imageh if dt.base == dtypes.half else dtypes.imagef)((*dims[0], 4)))
-    if vcount != 1: buf = UOp.stack(*([buf] * vcount))
+    if vcount != 1: buf = UOp.vectorize(*([buf] * vcount))
     if ls.op is Ops.LOAD: return ls.replace(src=(buf.index(off, ptr=True),), dtype=dtypes.float.vec(ls.dtype.vcount)).cast(dt.base)
     return buf.index(off, ptr=True).store(pm_imageh_store.rewrite(ls.src[1]) if dt.base == dtypes.half else ls.src[1])
 

tinygrad/tensor.py

@@ -822,27 +822,19 @@ class Tensor(OpMixin):
     """
     Returns a tensor with `num_samples` indices sampled from a multinomial distribution weighted by `self`.
 
+    NOTE: `replacement=False` for `num_samples > 1` is not supported yet.
     ```python exec="true" source="above" session="tensor" result="python"
     Tensor.manual_seed(42)
     t = Tensor([1, 2, 3, 4])
     print(t.multinomial(20, replacement=True).numpy())
     ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    Tensor.manual_seed(42)
-    t = Tensor([1, 2, 3, 4])
-    print(t.multinomial(3, replacement=False).numpy())
-    ```
     """
     assert 1 <= self.ndim <= 2 and num_samples > 0, f"{self.ndim=} must be 1 or 2 dim, {num_samples=} must be positive"
+    assert replacement or num_samples == 1, "no replacement only supports num_samples = 1"
     weight = self.unsqueeze(0) if self.ndim == 1 else self
-    assert replacement or num_samples <= weight.shape[1], "no replacement samples must not exceed population size"
-    if replacement or num_samples == 1:
-      cdf = (cw := weight.cumsum(1).float()) / cw[:, -1].unsqueeze(1)
-      unif_samples = Tensor.rand(num_samples, cdf.shape[0], 1).to(self.device)
-      indices = (unif_samples.expand((-1, -1, cdf.shape[1])) >= cdf).sum(2).permute((1, 0))
-    else:
-      # Efraimidis–Spirakis
-      indices = (weight.rand_like(dtype=dtypes.float32).log2() / weight).topk(num_samples, dim=1)[1]
+    cdf = (cw := weight.cumsum(1).float()) / cw[:, -1].unsqueeze(1)
+    unif_samples = Tensor.rand(num_samples, cdf.shape[0], 1).to(self.device)
+    indices = (unif_samples.expand((-1, -1, cdf.shape[1])) >= cdf).sum(2).permute((1, 0))
     return (indices.squeeze(0) if self.ndim == 1 else indices).cast(dtypes.int32)
 
   # ***** toposort and backward pass *****

tinygrad/uop/decompositions.py

@@ -418,7 +418,7 @@ def f2f_clamp(val:UOp, dt:DType) -> UOp:
 
 def f2f_load(x: UOp, fr:DType, to:DType) -> UOp:
   if (n:=x.dtype.count) == 1: return f2f(x.replace(dtype=f2f_dt[fr]), fr, to)
-  return UOp.stack(*(f2f(x.replace(dtype=f2f_dt[fr], src=(reindex(x.src[0].src[0], i, 1),)), fr, to) for i in range(n)))
+  return UOp.vectorize(*(f2f(x.replace(dtype=f2f_dt[fr], src=(reindex(x.src[0].src[0], i, 1),)), fr, to) for i in range(n)))
 
 def f2f_store(st, idx, val, fr:DType, to:DType):
   if (n:=val.dtype.count) == 1: return st.replace(src=(idx, f2f(val.bitcast(f2f_dt[to]), to, fr)))

tinygrad/uop/ops.py

@@ -419,7 +419,7 @@ class UOp(OpMixin, 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, dtypes.void, tuple([x for x in srcs if x is not None]))
-  def stack(self, *srcs:UOp, **kwargs):
+  def vectorize(self, *srcs, **kwargs):
     return UOp(Ops.STACK, self.dtype.vec(len(srcs)+1), (self,)+srcs, **kwargs)
   def index(self, *srcs:UOp|None, ptr=False, **kwargs):
     return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
@@ -953,10 +953,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     body = self if self.op is Ops.TUPLE else UOp.maketuple(self)
     return UOp(Ops.FUNCTION, dtypes.void, (body,)+srcs, CallInfo(grad_fxn, metadata, name, precompile, precompile_backward))
   def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
-    contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs)
-    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)]
-    kernel = fxn(*placeholders).call(*contig_srcs, grad_fxn=grad_fxn)
-    return [s.after(kernel) for s in contig_srcs]
+    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(srcs)]
+    kernel = fxn(*placeholders).call(*srcs, grad_fxn=grad_fxn)
+    return [s.after(kernel) for s in srcs]
 
 @dataclass(frozen=True)
 class KernelInfo:
@@ -1136,8 +1135,8 @@ class UPat(OpMixin):
 
   # copied from UOp
   def sink(self, *srcs:UPat|None, **kwargs): return UPat(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
-  def index(self, *srcs:UPat|None, **kwargs):
-    return UPat(Ops.INDEX, self.match_dtype, (self,)+tuple(x for x in srcs if x is not None), **kwargs)
+  def index(self, idx:UPat, valid:UPat|None=None, **kwargs):
+    return UPat(Ops.INDEX, self.match_dtype, (self,idx,valid) if valid is not None else (self,idx), **kwargs)
   def cast(self, dtype=None, **kwargs):
     if dtype is not None and self.match_dtype == (dtype,): return self
     return UPat(Ops.CAST, dtype, (self,), **kwargs)
@@ -1533,7 +1532,7 @@ pm_lower_index_dtype = PatternMatcher([
    lambda n: n.replace(src=tuple(s.src[0] if s.op is Ops.CAST and s.dtype == dtypes.weakint else s for s in n.src))),
   # vectorized indexes (ie. images) must be int
   (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.STACK, dtypes.long, name="vec")), allow_any_len=True, name="idx"),
-   lambda idx,vec: idx.replace(src=(idx.src[0], UOp.stack(*(u.cast(dtypes.int) for u in vec.src)), *idx.src[2:])))
+   lambda idx,vec: idx.replace(src=(idx.src[0], UOp.vectorize(*(u.cast(dtypes.int) for u in vec.src)), *idx.src[2:])))
 ])
 def _index_to_concrete_int(u:UOp) -> UOp: return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]