that turned out to be dumb
MODIFY
extra/thneed.py
extra/thneed.py
@@ -16,8 +16,7 @@ FLOAT16 = getenv("FLOAT16", 0)
class Thneed:
def __init__(self, cl_cache=[], inputs={}):
- self.cl_cache = [(prg, (list(args[0:2]) + [(x._cl if isinstance(x, (CLImage, CLBuffer)) else x) for x in args[2:]])) for prg,args in cl_cache]
- self.inputs = {k:v._cl for k,v in inputs.items()}
+ self.cl_cache, self.inputs = cl_cache[:], inputs
self.gobj = 0
# build graph

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -340,7 +340,7 @@ class GPUBuffer(ExplicitExecAST):
if self._backing is not None:
self._buf.copyin(self._backing)
self._backing = None
- return self._buf
+ return self._buf._cl
def __repr__(self): return f"GPUBuffer(shape={self.st}, hostbuf=GPUBuffer(shape={self._base_shape}" + (f", backing=np.array({self._backing}, dtype=np.float32)))" if self._backing else ", force_create=True))")

MODIFY
tinygrad/runtime/cuda.py
tinygrad/runtime/cuda.py
@@ -9,9 +9,9 @@ class CLImage:
def __init__(self, shape): raise NotImplementedError("CUDA runtime doesn't support images")
class CLBuffer:
- def __init__(self, size): self._cuda = cuda.mem_alloc(size)
- def copyin(self, b:np.ndarray, stream:Optional[cuda.Stream]=None): cuda.memcpy_htod_async(self._cuda, b, stream)
- def copyout(self, a:np.ndarray): cuda.memcpy_dtoh(a, self._cuda)
+ def __init__(self, size): self._cl = cuda.mem_alloc(size)
+ def copyin(self, b:np.ndarray, stream:Optional[cuda.Stream]=None): cuda.memcpy_htod_async(self._cl, b, stream)
+ def copyout(self, a:np.ndarray): cuda.memcpy_dtoh(a, self._cl)
class CLProgram:
def __init__(self, name:str, prg:str, op_estimate:int=0, mem_estimate:int=0):
@@ -25,7 +25,7 @@ class CLProgram:
assert all(x%y == 0 for x,y in zip(global_size, local_size)), f"local:{local_size} must divide global:{global_size}"
global_size = [x//y for x,y in zip(global_size, local_size)]
if DEBUG >= 2: print("CUDA launch", global_size, local_size)
- self.prg([x._cuda for x in args], block=tuple(local_size), grid=tuple(global_size))
+ self.prg(*args, block=tuple(local_size), grid=tuple(global_size))
GlobalCounters.kernel_count += 1
GlobalCounters.global_ops += self.op_estimate
GlobalCounters.global_mem += self.mem_estimate
\ No newline at end of file

MODIFY
tinygrad/runtime/opencl.py
tinygrad/runtime/opencl.py
@@ -81,7 +81,7 @@ class CLProgram:
# print the PTX for NVIDIA. TODO: probably broken for everything else
if DEBUG >= 5: print(self.clprogram.get_info(cl.program_info.BINARIES)[0].decode('utf-8'))
if CL.CACHE is not None: CL.CACHE.append((self, args))
- else: e = self.clprg(CL().cl_queue, args[0], args[1], *[(x._cl if isinstance(x, (CLImage, CLBuffer)) else x) for x in args[2:]])
+ else: e = self.clprg(CL().cl_queue, *args)
if DEBUG >= 2 and CL.CACHE is None:
CL.cl_queue.finish()
# NOTE: Profiling is not in ns in OS X, we multiply by a computed ratio

that turned out to be dumb