tinyfs: use assign (#16947)

test/null/test_tinyfs.py

@@ -1,27 +0,0 @@
-import unittest
-from tinygrad import Tensor
-from tinygrad.nn.state import fs_store, fs_load
-
-class TestLoadStore(unittest.TestCase):
-  def test_load_shape(self):
-    t = fs_load(Tensor(bytes(16)), 1024)
-    assert t.shape == (1024,), t.shape
-    t.schedule_linear()
-
-  def test_store_shape(self):
-    t = fs_store(Tensor.zeros(1024))
-    assert t.shape == (16,), t.shape
-    t.schedule_linear()
-
-  def test_load_large_shape(self):
-    t = fs_load(Tensor(bytes(16)), 10_000_000)
-    assert t.shape == (10_000_000,), t.shape
-    t.schedule_linear()
-
-  def test_store_large_shape(self):
-    t = fs_store(Tensor.zeros(10_000_000))
-    assert t.shape == (16,), t.shape
-    t.schedule_linear()
-
-if __name__ == "__main__":
-  unittest.main()

tinygrad/nn/state.py

@@ -102,7 +102,8 @@ def fs_store(t:Tensor) -> Tensor:
 
   level_chunks = base_chunks
   for _ in range(tree_depth + 1):
-    data = data.to("tinyfs:store")[:level_chunks * 16].contiguous().to(to_device)
+    # assign data into tinyfs:store and read back hashes
+    data = Tensor.empty(data.shape[0], dtype=dtypes.uint8, device="tinyfs:store").assign(data)[:level_chunks * 16].to(to_device)
     if (tsize := data.shape[0]) % CHUNK_SIZE != 0: data = data.pad((0, CHUNK_SIZE - tsize % CHUNK_SIZE))
     level_chunks = math.ceil(data.shape[0] / CHUNK_SIZE)
 
@@ -123,18 +124,16 @@ def fs_load(t:Tensor, size:int) -> Tensor:
   tree_depth = math.ceil(math.log(base_chunks, CHUNK_SIZE // 16))
   data, level_chunks = h, 0
   for i in reversed(range(tree_depth + 1)):
-    data = data.to("tinyfs:load")
-
     # if not last level, its still hashes
     if i > 0 or tree_depth == 0:
       level_chunks = max(1, math.ceil(base_chunks / (CHUNK_SIZE // 16)**(i-1)))
-      pad_amt = 16 * level_chunks
-    else: pad_amt = CHUNK_SIZE * level_chunks
-    if (tsize := data.shape[0]) < pad_amt: data = data.pad((0, pad_amt - tsize))
-    data = data[:pad_amt].contiguous()
-    if i != 0: data = data.to(t.device)
+      out_sz = 16 * level_chunks
+    else: out_sz = CHUNK_SIZE * level_chunks
+    # assign hash into tinyfs:load and read back data
+    (load:=Tensor.empty(out_sz, dtype=dtypes.uint8, device="tinyfs:load"))[:data.shape[0]].assign(data)
+    data = load
 
-  return data[:size]
+  return data.to(t.device)[:size]
 
 # state dict
 

tinygrad/runtime/ops_tinyfs.py

@@ -74,8 +74,8 @@ class TinyFSDevice(Compiled):
 class TinyFSBuffer:
   def __init__(self, device:TinyFSDevice, size:int, offset=0, copyout_queue=None, hash_buf=None):
     self.device, self.size, self.offset = device, size, offset
-    self.copyout_queue = copyout_queue or []
-    self.hash_buf = hash_buf or bytearray()
+    self.copyout_queue = [] if copyout_queue is None else copyout_queue
+    self.hash_buf = bytearray() if hash_buf is None else hash_buf
   def __repr__(self): return f"<TinyFSBuffer size={self.size} offset={self.offset}>"
 
 class TinyFSAllocator(Allocator[TinyFSDevice]):
@@ -91,11 +91,11 @@ class TinyFSAllocator(Allocator[TinyFSDevice]):
 
     if dest.device.op == "LOAD":
       locs = self.dev.sfile.readline()
-      dest.copyout_queue = json.loads(locs)
-      dest.hash_buf = src.tobytes()
+      dest.copyout_queue[:] = json.loads(locs)
+      dest.hash_buf[:] = src.tobytes()
     elif dest.device.op == "STORE":
       expected_hashes = math.ceil(dest.size / CHUNK_SIZE)
-      dest.hash_buf = bytearray(expected_hashes * 16)
+      dest.hash_buf[:] = bytearray(expected_hashes * 16)
       self.dev.sfile.readinto(dest.hash_buf)
 
   def _copyout(self, dest:memoryview, src:TinyFSBuffer):

tinygrad/tensor.py

@@ -204,7 +204,7 @@ class Tensor(RandMixin):
     return self
 
   def assign(self, x:Tensor|PyConst|list|tuple) -> Tensor:
-    is_disk = isinstance(self.device, str) and self.device.startswith("DISK")
+    is_disk = isinstance(self.device, str) and self.device.startswith(("DISK", "TINYFS"))
     if not isinstance(x, Tensor): x = Tensor(x, device="CPU" if is_disk else self.device, dtype=self.dtype)
     if self.uop is x.uop: return self  # a self assign is a NOOP
     # broadcast x (shape only, dtype must match)