ane refactor, fix examples
MODIFY
ane/2_compile/simple/neuron.plist
ane/2_compile/simple/neuron.plist
<key>InputHeight</key>
<integer>1</integer>
<key>InputWidth</key>
- <integer>678</integer>
+ <integer>200</integer>
<key>InputType</key>
<string>Float16</string>

MODIFY
ane/lib/ane.py
ane/lib/ane.py
@@ -3,6 +3,7 @@ import os
from ctypes import *
import numpy as np
import faulthandler
+import struct
faulthandler.enable()
libane = cdll.LoadLibrary(os.path.join(
@@ -67,6 +68,10 @@ ANE_Struct = [
("u8", 0x273, "OutputInterleave")]
+ANE_Struct_Dict = {}
+for typ, num, nam in ANE_Struct:
+ styp = {"u32": "I", "u16": "H", "u8": "B"}[typ]
+ ANE_Struct_Dict[nam] = (styp, num)
class ANETensor:
def __init__(self, *shape):
@@ -101,6 +106,18 @@ class ANE:
def tensor(self, shape):
return ANETensor(shape)
+ def filln(self, dat, nvdict, base=0x4000):
+ for n,v in nvdict.items():
+ styp, num = ANE_Struct_Dict[n]
+ dat = self.fill(dat, [num], styp, v)
+ return dat
+
+ def fill(self, dat, addrs, type, val, base=0x4000):
+ x = struct.pack(type, val)
+ for a in addrs:
+ dat[base+a:base+a+len(x)] = x
+ return dat
+
if __name__ == "__main__":
ane = ANE()
@@ -120,3 +137,4 @@ if __name__ == "__main__":
print("** after **")
print(tind)
print(toutd)
+

MODIFY
examples/benchmark.py
examples/benchmark.py
@@ -7,7 +7,7 @@ import time
ba = Tensor(np.random.normal(size=(0x4000,)))
for dev in ["CPU", "GPU", "ANE"]:
if dev == "GPU":
- baa = ba.cuda()
+ baa = ba.gpu()
elif dev == "ANE":
baa = ba.ane()
else:

MODIFY
tinygrad/ops_ane.py
tinygrad/ops_ane.py
from .tensor import Device, Function, register
from functools import lru_cache
-import struct
@lru_cache
def compile_wrapper(ane, dat):
@@ -9,22 +8,17 @@ def compile_wrapper(ane, dat):
def roundup(x, v):
return x + (v-x)%v
-def fill(dat, addrs, type, val, base=0x4000):
- x = struct.pack(type, val)
- for a in addrs:
- dat[base+a:base+a+len(x)] = x
- return dat
-
@lru_cache
def compile_relu(ane, sz):
dat = list(open("ane/ops/relu.hwx", "rb").read())
# TODO: make this all nice and once
- # number of relus
- dat = fill(dat, [0x128, 0x13C], "H", sz)
# number of engines? (max 0x100)
- dat = fill(dat, [0x1ec, 0x1f0, 0x1f4, 0x1f8], "I", max(0x100, roundup(sz*2, 0x10)))
- # strides?
- dat = fill(dat, [0x260, 0x264, 0x268], "I", roundup(sz*2, 0x40))
+ dat = ane.fill(dat, [0x1ec, 0x1f0, 0x1f4, 0x1f8], "I", max(0x100, roundup(sz*2, 0x10)))
+ stride = roundup(sz*2, 0x40)
+ dat = ane.filln(dat, {
+ "InputWidth": sz, "OutputWidth": sz,
+ "InputRowStride": stride, "InputPlaneStride": stride, "InputDepthStride": stride,
+ "OutputRowStride": stride, "OutputPlaneStride": stride, "OutputDepthStride": stride})
return compile_wrapper(ane, bytes(dat))
class ReLU(Function):

ane refactor, fix examples