Merge branch 'master' into check_dims

tinygrad/opt/kernel.py

@@ -449,22 +449,7 @@ class Kernel:
       if op.op in GroupOp.Buffer and op in self.bufs:
         st = self.sts[self.bufs.index(op)]
         # replace the VIEW source
-        if op.op is Ops.LOAD:
-          global_buf = ret.src[0].src[0]
-          # add locals cache
-          local_shape = [s if self.axis_types[i] not in (AxisType.GLOBAL, AxisType.REDUCE)
-                         and ss != 0 else 1 for i,(s,ss) in enumerate(zip(st.shape, st.real_strides()))]
-          # NOTE: this can have any permutation here
-          lst = lst_store = ShapeTracker.from_shape(tuple(local_shape)).expand(st.shape)
-          lbuf = UOp(Ops.DEFINE_LOCAL, dtype=global_buf.dtype.base.ptr(prod(local_shape), AddrSpace.LOCAL), arg=1000+global_buf.arg)
-          # TODO: permute to place any UPCASTs in 0 stride LOCALs
-          # any permutes of st + lst_store together are fine
-          print(list(zip(self.shape_str(), st.shape, st.real_strides())))
-          ret = ret.replace(src=(ret.src[0].replace(arg=st),)+ret.src[1:])
-          ret = lbuf.view(lst).load(lbuf.view(lst_store).store(ret))
-        else:
-          ret = ret.replace(src=(ret.src[0].replace(arg=st),)+ret.src[1:])
-        return ret
+        return ret.replace(src=(ret.src[0].replace(arg=st),)+ret.src[1:])
       if op.op is Ops.SINK:
         # NOTE: should group_for_reduces be added to the local_dims?
         kernel_name = ret.arg.name if ret.arg is not None else self.name if name_override is None else name_override

tinygrad/opt/tc.py

@@ -37,6 +37,10 @@ class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x
     assert 2**local_axes == self.threads, f"{self.threads} threads construct the warp but found {2**local_axes} in {self.opts}"
     assert 2**upcast_axes == self.elements_per_thread[2], \
       f"{self.elements_per_thread[2]} elements from C are processed per thread but found {2**upcast_axes} in {self.opts}"
+    # check dims match opts
+    assert self.dims[0] == 2**len(gd:=[x for x in self.opts if x[1] == '0']), f"opts wrong on dims[0], {self.dims[0]} vs {gd}"
+    assert self.dims[1] == 2**len(gd:=[x for x in self.opts if x[1] == '1']), f"opts wrong on dims[1], {self.dims[1]} vs {gd}"
+    # NOTE: the K opts is implictly set by the dim
     # check swizzle
     assert len(self.swizzle[0]) == 3 and len(self.swizzle[1]) == 3, "swizzle has wrong part count"
     assert len(self.swizzle[0][0]) == len(self.swizzle[1][0]) == local_axes, "local swizzle size is wrong"
@@ -52,6 +56,7 @@ class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x
       if o[1] == '1': zero_stride_1.append(o[0] + str(un if o[0] == 'u' else ln))
       if o[0] == 'u': un += 1
       if o[0] == 'l': ln += 1
+    # NOTE: all the zero_stride dims can be placed in any order in the swizzle
     upcasted_0 = [x for x in (self.swizzle[0][1] + self.swizzle[0][2]) if x not in zero_stride_0 and x[0] != 'l']
     upcasted_1 = [x for x in (self.swizzle[1][1] + self.swizzle[1][2]) if x not in zero_stride_1 and x[0] != 'l']
     assert 2**len(upcasted_0) == self.elements_per_thread[0], f"mismatch in elements_per_thread[0], {upcasted_0} vs {self.elements_per_thread[0]}"