working on kernel 15

extra/onnx.py

@@ -745,13 +745,16 @@ def get_onnx_ops():
       # 3x3 depthwise (C,1,3,3)
       # "width multiple of 4 depth multiple of 32 aligned to 128bytes"
       x = x.pad(((0,0), (0,0), (0,0), (0,1)))
-      if x.shape[0]%32 == 0 and False:
-        print("HERE")
+      if x.shape[0]%32 == 0:
+        # depth/32 is a loop -- lsr(depth, #5)
+        # width/4 is a loop -- lsr(out_width, #2)
+        # height is a loop
         x = x.reshape(-1, 32, 1, 3, 4)
         order = (0,3,1,2,4)
         x = x.permute(*order).contiguous().permute(*argsort(order))
         x = x.reshape(-1, 1, 3, 4)
       else:
+        print("HERE", x.shape)
         order = (2,0,1,3)
         x = x.permute(*order).contiguous().permute(*argsort(order))
       x = x[:, :, :, :3]
@@ -764,7 +767,7 @@ def get_onnx_ops():
       else:
         # DSP swizzle memory
         x = x.reshape(x.shape[0], x.shape[1]//WEIGHT_SHIFT, WEIGHT_SHIFT).permute(1,0,2).contiguous().permute(1,0,2).reshape(x.shape)
-    if getenv("NHWC") and x.shape == (1000, 1280):
+    if False and getenv("NHWC") and x.shape == (1000, 1280):
       x = x.reshape(-1, 320, 4)
       order = (1,0,2)
       x = x.permute(*order).contiguous().permute(*argsort(order))

extra/replay_pkl.py

@@ -114,6 +114,12 @@ if __name__ == "__main__":
             k.apply_opt(Opt(OptOps.UPCAST, 2, 0))
             k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
             #k.apply_opt(Opt(OptOps.UNROLL, 0, 0))
+          elif knum == 15:
+            # 28x28, 192 chan, 3x3 dwconv
+            k.apply_opt(Opt(OptOps.UNROLL, 1, 0))
+            k.apply_opt(Opt(OptOps.UNROLL, 0, 0))
+            k.apply_opt(Opt(OptOps.UPCAST, 2, 32))
+            #k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
           elif knum == 4:
             k.apply_opt(Opt(OptOps.UNROLL, 0, 8))
             k.apply_opt(Opt(OptOps.UPCAST, 1, 96))

tinygrad/runtime/ops_dsp.py

@@ -67,8 +67,11 @@ def multi_add_int32(**aa):
     swizzle.append(32+i)
     swizzle.append(64+i)
     swizzle.append(96+i)
+  for x in aa.values():
+    assert x.src[0].dtype.scalar() is dtypes.uchar
+    assert x.op is Ops.CAST
   swizzle = tuple(swizzle)
-  m0 = UOp(Ops.CAT, dtypes.int.vec(128), src=tuple(x.src[0] for x in aa.values())).gep(swizzle)
+  m0 = UOp(Ops.CAT, dtypes.uchar.vec(128), src=tuple(x.src[0] for x in aa.values())).gep(swizzle)
   return UOp(Ops.CUSTOMI, dtypes.int.vec(32), (m0, UOp.const(dtypes.uint, 0x01010101)), "__builtin_HEXAGON_V6_vrmpybus_128B({0}, {1})")
 
 def multi_add_int2(**aa):