diff --git a/test/null/test_winograd.py b/test/null/test_winograd.py index 8f08a2ff8a..1530675828 100644 --- a/test/null/test_winograd.py +++ b/test/null/test_winograd.py @@ -15,12 +15,22 @@ class TestWinograd(unittest.TestCase): out = Tensor.conv2d(x,w) self.assertEqual(len(out.schedule_linear().src), 4) - def test_backward_kernels(self): - x,w = Tensor.empty(1,4,9,9).realize(), Tensor.empty(4,4,3,3).realize() - out = Tensor.conv2d(x,w, padding=1) - out.mean().backward() - backward_schedule = x.grad.schedule_linear(w.grad) - self.assertEqual(len(backward_schedule.src), 4) + def test_backward_counters(self): + # contiguous_backward on the pooled input keeps the input-transform adjoint out of the overlap accumulation, so + # winograd backward runs in a fraction of the direct-conv flops; NOOPT=1 keeps the raw flop ratio from drifting with the optimizer + IC, OC, H = 64, 64, 28 + x,w = Tensor.empty(1,IC,H,H,device="NULL").realize(), Tensor.empty(OC,IC,3,3,device="NULL").realize() + x.requires_grad = w.requires_grad = True + def backward_ops(wino): + x.grad = w.grad = None + GlobalCounters.reset() + with Context(NOOPT=1, WINO=wino): + Tensor.conv2d(x,w,padding=1).mean().backward() + Tensor.realize(x.grad, w.grad) + return GlobalCounters.global_ops + ops_wino, ops_normal = backward_ops(1), backward_ops(0) + print(f"backward ops: normal {ops_normal} wino {ops_wino} ratio {ops_wino/ops_normal:.2f}") + self.assertLess(ops_wino/ops_normal, 0.35) def test_counters(self): IC, OC, H = 64, 64, 28 diff --git a/test/unit/test_multitensor.py b/test/unit/test_multitensor.py index 2dbfbaedd9..5d48f5dfe5 100644 --- a/test/unit/test_multitensor.py +++ b/test/unit/test_multitensor.py @@ -227,7 +227,6 @@ class TestMultiTensor(unittest.TestCase): optim.step() out.numpy() - @slow def test_backprop_conv_wino(self): with Context(WINO=1): self.test_backprop_conv() diff --git a/tinygrad/mixin/op.py b/tinygrad/mixin/op.py index 4b0ba3ee4f..2a64b17995 100644 --- a/tinygrad/mixin/op.py +++ b/tinygrad/mixin/op.py @@ -1382,22 +1382,17 @@ class OpMixin(ElementwiseMixin, ReduceMixin): ret = (indices.reshape(bs,c,1,-1)._one_hot_along_dim(prod(output_size), 2).where(self.reshape(bs,c,1,-1), 0)).sum(3) return ret.reshape(bs,c,*output_size) - @classmethod - def _get_winograd_matcols(cls, mat, dims:int, shp:tuple[sint, ...], dtype:DType) -> list[list[Self]]: - return [[cls.cat(*[cls.full(shp[:dim] + (1,) + shp[dim+1:], float(m[k]), dtype=dtype, buffer=False) for m in mat], dim=dim) - for k in range(len(mat[0]))] for dim in range(dims)] - # winograd conv 3 kernel f(4x4,3x3) see: http://arxiv.org/abs/1509.09308 def _apply_winograd_matrix(self, mat, dims:int) -> Self: - # multiply mat_1 @ mat_2 @ t with foldable constants, where mat_i acts on vector t along dimension i; roughly kron(mat, mat) @ t - # due to realize-before-expand rule in lazy.py, we must operate in this order: reshape -> expand -> arithmetic - t_ = self.reshape(self.shape[:dims] + (1,) * dims + self.shape[dims:]).expand( - self.shape[:dims] + (len(mat),) * dims + self.shape[dims:]) # add output dims - # precalculate mat columns for each dim; prod(itertools.product(matcols)) gives the columns of kron(mat, mat, ...) - matcols = type(self)._get_winograd_matcols(mat, dims, t_.shape[dims:], t_.dtype) - # multiply each element of t_ by the corresponding stacked column of kron(mat, mat), producing only one view for each element of t - ret = sum(prod(col[idx] for col, idx in zip(matcols, mat_is)) * t_[mat_is] for mat_is in itertools.product(range(len(mat[0])), repeat=dims)) - assert not isinstance(ret, int), "sum over empty winograd matrix" + # apply mat along each of the first `dims` axes: the separable transform kron(mat, ..., mat) @ self + # column k of mat is a stacked-CONST vector that folds into the arithmetic, so no constant is materialized + ret = self + for dim in range(dims): + ret = ret.transpose(0, dim) + ret = sum(type(self).const(ret.dtype, tuple(float(m[k]) for m in mat)).reshape((len(mat),)+(1,)*(ret.ndim-1)) * ret[k] + for k in range(len(mat[0]))) + assert not isinstance(ret, int), "sum over empty winograd matrix" + ret = ret.transpose(0, dim) return ret # TODO: winograd can be a rewrite rule like split_reduceop @@ -1417,8 +1412,8 @@ class OpMixin(ElementwiseMixin, ReduceMixin): # (bs, cin_, tyx, HWI) pads = [(pB, pA + (-(s + pB + pA - 2) % 4)) for (pB, pA), s in zip(flat_to_grouped(padding_), self.shape[-len(HW):])] d = self.pad(flatten(reversed(pads)))._pool(HWI, HWO) - # move HW to the front: # (HWI, bs, cin_, tyx) - d = d.permute(*range(len(d.shape)-len(HW),len(d.shape)), *range(len(d.shape)-len(HW))) + # move HW to the front: # (HWI, bs, cin_, tyx); contiguous_backward keeps the input transform's adjoint out of the overlap accumulation + d = d.permute(*range(len(d.shape)-len(HW),len(d.shape)), *range(len(d.shape)-len(HW))).contiguous_backward() tyx = d.shape[-len(HWI):] # dim of tiling g = weight.permute(*range(len(weight.shape)-len(HW),len(weight.shape)), *range(len(weight.shape)-len(HW))) # move HW to the front @@ -1881,8 +1876,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): # https://keccak.team/keccak_specs_summary.html def ctensor(l: Sequence[PyConst], dtype: DType = dtypes.uint64): - # TODO: contiguous is here for compile speed - return type(self).stack(*(type(self).const(dtype, v) for v in l)).contiguous() + return type(self).const(dtype, tuple(l)) rot_offsets = [44, 43, 21, 14, 28, 20, 3, 45, 61, 1, 6, 25, 8, 18, 27, 36, 10, 15, 56, 62, 55, 39, 41, 2] rot_offsets_v0, rot_offsets_v1 = ctensor([0] + [1 << v for v in rot_offsets]), ctensor([1] + [1 << (64 - v) for v in rot_offsets])