remove MultiLazyBuffer.from_sharded [pr] (#8620)

it's eqivalent to taking the lazydata from Tensor.split, then copy to devices

test/test_multitensor.py

@@ -75,7 +75,7 @@ class TestMultiTensor(unittest.TestCase):
       ei.run()
     assert names[-2] == names[-1], "function was relinearized"
 
-  @unittest.skip("this doesn't fold because from_sharded calls contiguous on all lbs")
+  @unittest.skip("this doesn't fold because shard_ calls contiguous on all lbs")
   def test_sharded_memory(self):
     # Buffer may be stuck in track_cross_buffer
     for x in (d0, d1, d2, d3, d4): Device[x].synchronize()

tinygrad/multi.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import functools, itertools, operator
-from tinygrad.helpers import all_same, all_int, dedup, prod, DEBUG, RING, getenv, ceildiv
+from tinygrad.helpers import all_same, all_int, dedup, prod, DEBUG, RING, getenv
 from tinygrad.dtype import DType
 from tinygrad.ops import Ops, MathTrait, UOp, sint
 
@@ -63,18 +63,6 @@ class MultiLazyBuffer(MathTrait):
 
   def __repr__(self): return f"<MLB {self.axis=} {self.real=} {chr(10)}{chr(10).join([f'{x.device} {x.st}' for x in self.lbs])}>"
 
-  @staticmethod
-  def from_sharded(lb:UOp, devices:tuple[str, ...], axis:int|None):
-    if axis is not None:
-      if not isinstance(total:=lb.shape[axis], int): raise RuntimeError(f"cannot shard symbolic shape {lb.shape=}, {axis=}")
-      sz = ceildiv(total, len(devices))
-      splits = tuple([max(0, min(sz, total - sz*i)) for i in range(len(devices))])
-      bounds = tuple(itertools.pairwise(itertools.accumulate(splits, initial=0)))
-    lbs = [lb] * len(devices)
-    sharded_lbs = [lb.copy_to_device(d) for lb,d in zip(to_sharded(lbs, axis, bounds) if axis is not None and bounds is not None else lbs, devices)]
-    # NOTE: this contiguous is making it impossible for the scheduler to do late const folding
-    return MultiLazyBuffer([lb.contiguous(allow_buffer_view=False) for lb in sharded_lbs], axis)
-
   def copy_to_device(self, device:str) -> UOp:
     if self.axis is None:
       # if we already have a copy on the device, return that

tinygrad/tensor.py

@@ -179,7 +179,7 @@ class Tensor(SimpleMathTrait):
     # data might be on a different device
     if isinstance(device, str): self.lazydata:Union[UOp, MultiLazyBuffer] = data if data.device == device else data.copy_to_device(device)
     # if device is a tuple, we should have/construct a MultiLazyBuffer
-    elif isinstance(data, UOp): self.lazydata = MultiLazyBuffer.from_sharded(data, device, None)
+    elif isinstance(data, UOp): self.lazydata = Tensor(data).shard(device).lazydata
     else:
       assert data.device == device, f"MultiLazyBuffer device mismatch, {data.device} != {device}"
       self.lazydata = data
@@ -405,8 +405,16 @@ class Tensor(SimpleMathTrait):
     """
     assert isinstance(self.lazydata, UOp), "can't shard a MultiLazyBuffer"
     devices = tuple(Device.canonicalize(x) for x in devices)
-    if axis is not None: axis = self._resolve_dim(axis)
-    return Tensor(MultiLazyBuffer.from_sharded(self.lazydata, devices, axis), device=devices, requires_grad=self.requires_grad)
+    if axis is None: lbs = [self.lazydata] * len(devices)
+    else:
+      axis = self._resolve_dim(axis)
+      sz = ceildiv(self.shape[axis], len(devices))
+      sizes = [max(0, min(sz, self.shape[axis] - sz*i)) for i in range(len(devices))]
+      lbs = [cast(UOp, t.lazydata) for t in self.split(sizes, axis)]
+    sharded_lbs = [lb.copy_to_device(d) for lb,d in zip(lbs, devices)]
+    # NOTE: this contiguous is making it impossible for the scheduler to do late const folding
+    mlb = MultiLazyBuffer([lb.contiguous(allow_buffer_view=False) for lb in sharded_lbs], axis)
+    return Tensor(mlb, device=devices, requires_grad=self.requires_grad)
 
   def shard_(self, devices:tuple[str, ...], axis:Optional[int]=None):
     """