indexer changes for master
MODIFY
extra/onnx.py
extra/onnx.py
@@ -6,6 +6,8 @@ from tinygrad.tensor import Tensor
from tinygrad.helpers import prod
from tinygrad.ops import DEBUG
+ONNXLIMIT = int(os.getenv("ONNXLIMIT", "-1"))
+
def get_run_onnx(onnx_model):
def shape_to_tuple(s): return tuple(x.dim_value for x in s.dim)
def buffer_parse(inp):
@@ -59,6 +61,7 @@ def get_run_onnx(onnx_model):
def run_onnx(inputs={}, debug=False):
input_tensors = {}
intermediate_tensors = {}
+ output_tensor_names = [x.name for x in onnx_model.graph.output]
# get inputs
for inp in onnx_model.graph.input:
@@ -172,6 +175,9 @@ def get_run_onnx(onnx_model):
if debug: print(ret.shape)
intermediate_tensors[n.output[0]] = ret
#print(ret.numpy().mean())
+ if num == ONNXLIMIT:
+ output_tensor_names = n.output
+ break
- return {outp.name:intermediate_tensors[outp.name] for outp in onnx_model.graph.output}
+ return {outp:intermediate_tensors[outp] for outp in output_tensor_names}
return run_onnx

MODIFY
extra/thneed.py
extra/thneed.py
@@ -276,11 +276,13 @@ class Thneed:
et = time.monotonic() - st
print(f"submit in {(mt-st)*1000.0:.2f} ms, total runtime is {et*1000.0:.2f} ms")
- if DEBUGCL:
+ if DEBUGCL >= 1:
total_runtime = 0
for i, ((prg, args), e) in enumerate(zip(self.cl_cache, events)):
runtime = (e.profile.end - e.profile.start)
- print(f"{i:3d} time {total_runtime/1e6:5.2f} ms running {prg.name:20s} with {str(args[0]):15s} {str(args[1]):15s} count {len(args)-2:2d} runtime {runtime/1e3:7.2f} us {prg.options} -> {args[2].shape if hasattr(args[2], 'shape') else args[2].size}")
+ print(f"{i:3d} time {total_runtime/1e6:5.2f} ms running {prg.name:20s} with {str(args[0]):15s} {str(args[1]):15s} count {len(args)-2:2d} runtime {runtime/1e3:7.2f} us {(prg.op_estimate)/runtime:9.2f} GFLOPS {prg.options} -> {args[2].shape if hasattr(args[2], 'shape') else args[2].size}")
+ if DEBUGCL >= 2 and int(os.getenv("PRINT_KERNEL")) == i:
+ print(prg.prg)
total_runtime += runtime
print(f"total runtime: {total_runtime/1e6:.2f} ms wall time: {et*1000.0:.2f} ms")
return et

MODIFY
test/test_ops.py
test/test_ops.py
@@ -6,7 +6,8 @@ import unittest
from tinygrad.tensor import Tensor, Device
FORWARD_ONLY = bool(int(os.getenv("FORWARD_ONLY", "0")))
-def helper_test_op(shps, torch_fxn, tinygrad_fxn, atol=1e-6, rtol=1e-3, grad_atol=1e-4, grad_rtol=1e-3, forward_only=False, vals=None, a=-0.5, b=3):
+def helper_test_op(shps, torch_fxn, tinygrad_fxn=None, atol=1e-6, rtol=1e-3, grad_atol=1e-4, grad_rtol=1e-3, forward_only=False, vals=None, a=-0.5, b=3):
+ if tinygrad_fxn is None: tinygrad_fxn = torch_fxn
torch.manual_seed(0)
np.random.seed(0)
if shps is None:
@@ -443,6 +444,9 @@ class TestOps(unittest.TestCase):
def test_clip(self):
helper_test_op([(45,65)], lambda x: x.clip(-2.3, 1.2), lambda x: x.clip(-2.3, 1.2))
+ def test_matvec(self):
+ helper_test_op([(1,128), (128,128), (128,128)], lambda x,y,z: (x@y).relu()@z, atol=1e-4)
+
if __name__ == '__main__':
np.random.seed(1337)
unittest.main(verbosity=2)

MODIFY
test/test_shapetracker.py
test/test_shapetracker.py
@@ -36,6 +36,14 @@ class DumbShapeTracker:
# Tensor.zeros(2, 4).permute(1,0).reshape(2, 4)
# (d1*4 + d0%4), d1=x//4, d0=x%4 = ((x//4)*4) + (x%4)%4
+class TestZeroViewShapeTracker(unittest.TestCase):
+ def test_pad(self):
+ self.st = ShapeTracker((4, 4))
+ self.st.pad((1, 1), (1, 1))
+ assert self.st.shape == (6,6)
+ print(self.st)
+ print(self.st.expr())
+
class TestComplexShapeTracker(unittest.TestCase):
def test_add_1s(self):
self.st = ShapeTracker((4, 4))

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -64,9 +64,9 @@ class CL:
@functools.lru_cache(maxsize=None)
class CLProgram:
kernel_cnt : Dict[str, int] = defaultdict(int)
- def __init__(self, name:str, prg:str, options:Tuple[str, ...]=tuple(), argdtypes=None, rename=True, binary=False):
+ def __init__(self, name:str, prg:str, options:Tuple[str, ...]=tuple(), argdtypes=None, rename=True, binary=False, op_estimate=0):
self.name = f"{name}{('_N'+str(CLProgram.kernel_cnt[name])) if CLProgram.kernel_cnt[name] else ''}" if rename else name
- self.prg, self.options, self.argdtypes = prg.replace(f"{name}(", f"{self.name}(") if rename else prg, options, argdtypes
+ self.prg, self.options, self.argdtypes, self.op_estimate = prg.replace(f"{name}(", f"{self.name}(") if rename else prg, options, argdtypes, op_estimate
self.clprogram = cl.Program(CL().cl_ctx, CL().cl_ctx.devices, [self.prg]) if binary else cl.Program(CL().cl_ctx, self.prg) # type: ignore
try:
self.clprg = self.clprogram.build(options=list(self.options)).__getattr__(self.name)
@@ -76,7 +76,7 @@ class CLProgram:
if self.argdtypes is not None:
self.clprg.set_scalar_arg_dtypes(self.argdtypes)
CLProgram.kernel_cnt[name] += 1
- def __call__(self, *args, op_estimate=0):
+ def __call__(self, *args):
CL.kernel_count += 1
if CL.CACHE is not None:
CL.CACHE.append((self, args))
@@ -88,10 +88,10 @@ class CLProgram:
CL.cl_queue.finish()
if DEBUG >= 1:
CL.time_sum += 0 if DEBUG <= 1 or CL.CACHE is not None else (e.profile.end - e.profile.start)
- CL.ops_sum += op_estimate
- print(f"**CL** {CL.kernel_count:6d} {self.name:28s} args {len(args[2:]):5d} kernels {str(args[0]):18s} {str(args[1]):12s} OPs {op_estimate/1e6:7.1f}M/{CL.ops_sum/1e9:7.2f}G mem {CL.mem_used/1e9:5.2f} GB " +
- ("" if DEBUG <= 1 or CL.CACHE is not None else f"tm {(e.profile.end - e.profile.start)/1e3:9.2f}us/{CL.time_sum/1e6:9.2f}ms ({op_estimate/(e.profile.end - e.profile.start):8.2f} GFLOPS)"))
- GlobalCounters.global_ops += op_estimate
+ CL.ops_sum += self.op_estimate
+ print(f"**CL** {CL.kernel_count:6d} {self.name:28s} args {len(args[2:]):5d} kernels {str(args[0]):18s} {str(args[1]):12s} OPs {self.op_estimate/1e6:7.1f}M/{CL.ops_sum/1e9:7.2f}G mem {CL.mem_used/1e9:5.2f} GB " +
+ ("" if DEBUG <= 1 or CL.CACHE is not None else f"tm {(e.profile.end - e.profile.start)/1e3:9.2f}us/{CL.time_sum/1e6:9.2f}ms ({self.op_estimate/(e.profile.end - e.profile.start):8.2f} GFLOPS)"))
+ GlobalCounters.global_ops += self.op_estimate
GlobalCounters.global_mem += sum([x.size//4 for x in args[2:] if isinstance(x, cl.Buffer)])
# **** end CL wrappers ****
@@ -304,11 +304,11 @@ class CLASTKernel(ASTKernel):
self.kernel = list(self.prekernel) + [f"__kernel void {function_name}(",] + [', '.join(f'{t} data{i}' for i,t in enumerate(buftypes) if i not in self.bufs_to_delete)] + [") {\n"] + self.kernel
# compile kernel
- fxn = CLProgram(function_name, ' '.join(self.kernel))
+ fxn = CLProgram(function_name, ' '.join(self.kernel), op_estimate=self.info.flops)
def runner(*bufs):
clbufs = [x.cl for i,x in enumerate(bufs) if i not in self.bufs_to_delete]
- return fxn(self.output_shape[::-1] if len(self.output_shape) > 0 else [1], None, *clbufs, op_estimate=self.info.flops)
+ return fxn(self.output_shape[::-1] if len(self.output_shape) > 0 else [1], None, *clbufs)
return runner
class GPUBuffer(ExplicitExecAST):

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -187,7 +187,7 @@ class ASTKernel:
st.strided(*zip(shape, stride))
# TODO: handle reduced shape here
st.reshape(*new_shape_fxn(shape))
- st.permute(*axis)
+ if axis is not None: st.permute(*axis)
assert len(st.views) == 1
new_shapes.append(st.shape)
new_strides.append(st.strides)

indexer changes for master