github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-02-15 00:55:11 -05:00
Code Issues Packages Projects Releases Wiki Activity

ops_ext to replace cpu import (#3406)

Browse Source 
* ops_ext to replace cpu import

* don't allow zero copy with as buffer

* memoryview(bytearray

* reenable test
This commit is contained in:
George Hotz
2024-02-15 12:14:58 +01:00
committed by GitHub
parent 49cb1fee54
commit 91eb93f85a
9 changed files with 29 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tinygrad/runtime/ops_ext.py Normal file
View File

@@ -0,0 +1,12 @@
from typing import Tuple, Any
from tinygrad.device import Compiled, Allocator
# the Any is an arbitrary object that's kept in scope with the memoryview
class ExtAllocator(Allocator):
# NOTE: this doesn't work with allow_zero_copy, it's read only somehow
#def as_buffer(self, src:Tuple[memoryview, Any]) -> memoryview: return src[0]
def copyin(self, dest:Tuple[memoryview, Any], src:memoryview): dest[0][:] = src
def copyout(self, dest:memoryview, src:Tuple[memoryview, Any]): dest[:] = src[0]
class ExtDevice(Compiled):
def __init__(self, device:str): super().__init__(device, ExtAllocator(), None, None)

2
tinygrad/runtime/ops_python.py
View File

@@ -93,7 +93,7 @@ class PythonProgram:
ul[i] = [pbufs.pop(0).cast(dtype.fmt)] * warp_size
elif uop is UOps.DEFINE_LOCAL:
assert dtype.fmt is not None
lbuf = memoryview(bytearray(arg[1]*dtype.sz))
lbuf = memoryview(bytearray(arg[1]*dtype.itemsize))
ul[i] = [lbuf.cast(dtype.fmt)] * warp_size
elif uop is UOps.SPECIAL:
if arg[1][0] == 'g':