sum in metal
MODIFY
accel/metal/ops_metal.py
accel/metal/ops_metal.py
@@ -6,7 +6,12 @@ import MetalPerformanceShaders
device = Metal.MTLCreateSystemDefaultDevice()
mtl_queue = device.newCommandQueue()
-mtl_buffer = mtl_queue.commandBuffer()
+mtl_buffers = []
+
+def cmd_buffer():
+ ret = mtl_queue.commandBuffer()
+ mtl_buffers.append(ret)
+ return ret
class MetalBuffer:
def __init__(self, shape, hostbuf=None):
@@ -36,19 +41,48 @@ class MetalBuffer:
return MetalBuffer(data.shape, data)
def toCPU(self):
- status = mtl_buffer.status()
- #print("waiting", status)
- if status != 0:
- mtl_buffer.waitUntilCompleted()
+ global mtl_buffers
+ for b in mtl_buffers:
+ b.waitUntilCompleted()
+ mtl_buffers = []
return np.frombuffer(b''.join(self.mtl.contents()[0:self.sz]), dtype=self.dtype).reshape(self.shape)
relu_shader = MetalPerformanceShaders.MPSImageThresholdToZero.alloc().initWithDevice_thresholdValue_linearGrayColorTransform_(device, 0, None)
+inv_relu_shader = MetalPerformanceShaders.MPSImageThresholdBinary.alloc().initWithDevice_thresholdValue_maximumValue_linearGrayColorTransform_(device, 0, 1, None)
add_shader = MetalPerformanceShaders.MPSImageAdd.alloc().initWithDevice_(device)
+mul_shader = MetalPerformanceShaders.MPSImageMultiply.alloc().initWithDevice_(device)
+sum_shader = MetalPerformanceShaders.MPSImageReduceRowSum.alloc().initWithDevice_(device)
+
+class Sum(Function):
+ def forward(ctx, input, axis=None):
+ assert axis is None
+ ctx.save_for_backward(input.shape, axis)
+ out = MetalBuffer((1,), None)
+ mtl_buffer = cmd_buffer()
+ sum_shader.encodeToCommandBuffer_sourceTexture_destinationTexture_(
+ mtl_buffer, input.texture, out.texture
+ )
+ mtl_buffer.commit()
+ return out
+
+ def backward(ctx, grad_output):
+ shape, axis = ctx.saved_tensors
+ out = MetalBuffer(shape, None)
+ ret = MetalBuffer(shape, None)
+ mtl_buffer = cmd_buffer()
+ add_shader.setPrimaryEdgeMode_(MetalPerformanceShaders.MPSImageEdgeModeClamp)
+ add_shader.setSecondaryEdgeMode_(MetalPerformanceShaders.MPSImageEdgeModeClamp)
+ add_shader.encodeToCommandBuffer_primaryTexture_secondaryTexture_destinationTexture_(
+ mtl_buffer, out.texture, grad_output.texture, ret.texture
+ )
+ mtl_buffer.commit()
+ return ret
class ReLU(Function):
def forward(ctx, input):
ctx.save_for_backward(input)
out = MetalBuffer(input.shape, None)
+ mtl_buffer = cmd_buffer()
relu_shader.encodeToCommandBuffer_sourceTexture_destinationTexture_(
mtl_buffer, input.texture, out.texture
)
@@ -57,6 +91,21 @@ class ReLU(Function):
def backward(ctx, grad_output):
input, = ctx.saved_tensors
+ out = MetalBuffer(input.shape, None)
+ mtl_buffer = mtl_queue.commandBuffer()
+ inv_relu_shader.encodeToCommandBuffer_sourceTexture_destinationTexture_(
+ mtl_buffer, input.texture, out.texture
+ )
+ # TODO: make in place work
+ #mul_shader.encodeToCommandBuffer_inPlacePrimaryTexture_secondaryTexture_fallbackCopyAllocator_(
+ # mtl_buffers, out.texture, grad_output.texture, None)
+ ret = MetalBuffer(input.shape, None)
+ mul_shader.encodeToCommandBuffer_primaryTexture_secondaryTexture_destinationTexture_(
+ mtl_buffer, grad_output.texture, out.texture, ret.texture
+ )
+ mtl_buffer.commit()
+ return ret
+
"""
class Add(Function):
@@ -73,6 +122,7 @@ if __name__ == "__main__":
b2 = MetalBuffer(10, np.ones(10))
out = MetalBuffer(10, None)
+ mtl_buffer = cmd_buffer()
add_shader.encodeToCommandBuffer_primaryTexture_secondaryTexture_destinationTexture_(
mtl_buffer, b1.texture, b2.texture, out.texture
)
@@ -84,8 +134,11 @@ if __name__ == "__main__":
from tinygrad.tensor import Tensor, Device
- r1 = Tensor([-2,-1,0,1,2], device=Device.METAL)
+ r1 = Tensor([-2,-1,0,2,4], device=Device.METAL)
r2 = r1.relu()
+ r3 = r2.sum()
+ r3.backward()
print(r1.cpu())
print(r2.cpu())
+ print(r3.cpu())

sum in metal