reshape doesn't copy anymore
MODIFY
tinygrad/opsgpu.py
tinygrad/opsgpu.py
@@ -350,21 +350,19 @@ class Pad2D(Function):
return ret
register('pad2d', Pad2D, gpu=True)
-# TODO: Reshape shouldn't make a copy, but this is a big change since data tensor can't have shape
class Reshape(Function):
@staticmethod
def forward(ctx, x, shape):
ctx.save_for_backward(x.shape)
- r = unary_op(ctx, 'a', x)
- r.shape = tuple(-np.prod(x.shape) // np.prod(shape) if s == -1 else s for s in shape)
+ shape = tuple(-np.prod(x.shape) // np.prod(shape) if s == -1 else s for s in shape)
+ r = GPUBuffer(shape, hostbuf=x)
assert np.prod(x.shape) == np.prod(r.shape)
return r
@staticmethod
def backward(ctx, grad_output):
in_shape, = ctx.saved_tensors
- grad_output = unary_op(ctx, 'a', grad_output)
- grad_output.shape = in_shape
+ grad_output = GPUBuffer(in_shape, hostbuf=grad_output)
return grad_output
register('reshape', Reshape, gpu=True)

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -50,6 +50,9 @@ class GPUBuffer:
self.shape = tuple(shape)
self.dtype = dtype
if hostbuf is not None:
+ if isinstance(hostbuf, GPUBuffer):
+ self.cl = hostbuf.cl
+ else:
self.cl = cl.Buffer(cl_ctx, cl.mem_flags.READ_ONLY | cl.mem_flags.COPY_HOST_PTR, hostbuf=hostbuf.ravel())
else:
self.cl = cl.Buffer(cl_ctx, cl.mem_flags.READ_WRITE, 4*np.prod(shape))

reshape doesn't copy anymore