From e8368993362bd151f810ae3c3916dfb33abe33b0 Mon Sep 17 00:00:00 2001 From: chenyu Date: Mon, 29 Jun 2026 09:58:21 -0400 Subject: [PATCH] remove Tensor conv2d and dot override [PR] (#16792) --- test/null/test_tensor_uop_mixin.py | 20 ++- tinygrad/mixin/__init__.py | 190 ++++++++++++++++++++++++++- tinygrad/tensor.py | 201 +---------------------------- 3 files changed, 210 insertions(+), 201 deletions(-) diff --git a/test/null/test_tensor_uop_mixin.py b/test/null/test_tensor_uop_mixin.py index ed7efa2fc5..f5d3ca6db5 100644 --- a/test/null/test_tensor_uop_mixin.py +++ b/test/null/test_tensor_uop_mixin.py @@ -1,6 +1,6 @@ import math, unittest from dataclasses import replace -from tinygrad import Tensor, dtypes +from tinygrad import Tensor, dtypes, Context from tinygrad.uop.ops import ParamArg, UOp, UPat, Ops, PatternMatcher, graph_rewrite _strip_unique_pm = PatternMatcher([ @@ -395,6 +395,24 @@ class TestTensorUOpConv2d(unittest.TestCase): def test_conv2d_3d(self): w = _t(1, 1, 2, 2, 2).float() _check(self, _t(1, 1, 3, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop)) + def test_conv2d_winograd(self): + w, a = _t(2, 2, 3, 3).float(), _t(1, 2, 6, 6).float() + with Context(WINO=0): direct = a.conv2d(w).uop + with Context(WINO=1): + self.assertIsNot(a.conv2d(w).uop, direct) + _check(self, a, lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop)) + def test_conv2d_image(self): + w, a = _t(4, 4, 3, 3).float(), _t(1, 4, 8, 8).float() + with Context(IMAGE=0): direct = a.conv2d(w).uop + with Context(IMAGE=1): + self.assertIsNot(a.conv2d(w).uop, direct) + _check(self, a, lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop)) + def test_dot_image(self): + y, a = _t(4, 3).float(), _t(2, 4).float() + with Context(IMAGE=0): direct = a.dot(y).uop + with Context(IMAGE=1): + self.assertIsNot(a.dot(y).uop, direct) + _check(self, a, lambda x: x.dot(y if isinstance(x, Tensor) else y.uop)) def test_conv_transpose2d_basic(self): w = _t(1, 1, 2, 2).float() _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv_transpose2d(w if isinstance(x, Tensor) else w.uop)) diff --git a/tinygrad/mixin/__init__.py b/tinygrad/mixin/__init__.py index 2a57ba6cb0..ff94401589 100644 --- a/tinygrad/mixin/__init__.py +++ b/tinygrad/mixin/__init__.py @@ -1,14 +1,14 @@ from __future__ import annotations -import functools, itertools, string +import functools, itertools, math, string from typing import TYPE_CHECKING, Callable, Self, Sequence, Literal, get_args from tinygrad.mixin.elementwise import ElementwiseMixin from tinygrad.mixin.movement import MovementMixin from tinygrad.mixin.reduce import ReduceMixin from tinygrad.uop import Ops from tinygrad.uop.ops import _broadcast_shape, resolve, smax, smin, identity_element -from tinygrad.dtype import ConstType, DTypeLike, Invalid, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype +from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype from tinygrad.helpers import all_int, argfix, argsort, ceildiv, flatten, flat_to_grouped, fully_flatten, get_shape, make_tuple, merge_dicts, prod -from tinygrad.helpers import resolve_pool_pads, round_up +from tinygrad.helpers import resolve_pool_pads, round_up, IMAGE, FLOAT16, WINO if TYPE_CHECKING: from tinygrad.uop.ops import sint, UOp @@ -388,6 +388,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): print(a.dot(b).numpy()) ``` """ + if IMAGE: return self.image_dot(w, dtype) x, dx, dw = self, self.ndim, w.ndim if not (dx > 0 and dw > 0): raise RuntimeError(f"both tensors need to be at least 1D, got {dx}D and {dw}D") if x.shape[-1] != w.shape[axis_w:=-min(w.ndim,2)]: raise RuntimeError(f"cannot dot {x.shape} and {w.shape}") @@ -1411,8 +1412,191 @@ 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" + return ret + + # TODO: winograd can be a rewrite rule like split_reduceop + def _conv2d_winograd(self, weight:Self, bias:Self|None, groups:int, padding:int|Sequence[int], dtype:DTypeLike|None) -> Self: + (bs,cin_), (cout,cin), HW = self.shape[:2], weight.shape[:2], weight.shape[2:] + padding_ = resolve_pool_pads(padding, len(HW)) + assert groups*cin == cin_ and len(self.shape) == len(weight.shape),\ + f"Input Tensor shape {self.shape} does not match the shape of the weights {weight.shape}. ({groups*cin} vs. {cin_})" + rcout, oyx = cout//groups, self.pad(padding_)._pool(HW, 1, 1).shape[2:-len(HW)] + HWI, HWO = (6,) * len(HW), (4,) * len(HW) # F(4x4,3x3) winograd tiles + winograd_G = [[1/4, 0, 0], [-1/6, -1/6, -1/6], [-1/6, 1/6, -1/6], [1/24, 1/12, 1/6], [1/24, -1/12, 1/6], [0, 0, 1]] + winograd_Bt = [[4, 0, -5, 0, 1, 0], [0, -4, -4, 1, 1, 0], [0, 4, -4, -1, 1, 0], [0, -2, -1, 2, 1, 0], [0, 2, -1, -2, 1, 0], [0, 4, 0, -5, 0, 1]] + winograd_At = [[1, 1, 1, 1, 1, 0], [0, 1, -1, 2, -2, 0], [0, 1, 1, 4, 4, 0], [0, 1, -1, 8, -8, 1]] # applying At in pre-order doubles compile time + + # TODO: stride == dilation + # use padding to round up to 4x4 output tiles + # (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))) + 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 + + # compute 6x6 winograd tiles: GgGt, BtdB. contiguous so the transforms are materialized once + # (HWI, groups * rcout, cin) -> (HWI, bs=1, groups, rcout, cin, tyx=(1,1)) + gfactors = g._apply_winograd_matrix(winograd_G, len(HW)).contiguous().reshape(*HWI, 1, groups, rcout, cin, *([1]*len(tyx))) + # (HWI, bs, cin_, tyx) -> (HWI, bs, groups, 1 ,cin, *tyx) + dfactors = d._apply_winograd_matrix(winograd_Bt, len(HW)).contiguous().reshape(*HWI, bs, groups, 1, cin, *tyx) + + # matmul; sum across cin: (HWI, bs, groups, rcout, *tyx); then HWI -> HWO: (HWO, bs, groups, rcout, *tyx) + ret = (gfactors * dfactors).sum(axis=-1-len(HW), dtype=dtype)._apply_winograd_matrix(winograd_At, len(HW)) + + # interleave tyx and HWO: (bs, groups, rcout, oy, HO, ox, WO) + ret = ret.permute([*range(len(HW), len(ret.shape)-len(HW)), *[i+o for i in range(len(HW)) for o in [len(ret.shape)-len(HW),0]]]) + # merge groups and rcout, tyx and HWO: (bs, groups, cout, *yx), shrink to final + ret = ret.reshape(bs, cout, *[c * HWO[i] for i, c in enumerate(tyx)]).shrink_to(bs, cout, *oyx) + + return (ret if bias is None else ret.add(bias.reshape(1, -1, *[1 for _ in range(len(HW))]))).contiguous().contiguous_backward() + + # *** image function replacements (used when IMAGE is set) *** + + def image_dot(self, w:Self, dtype:DTypeLike|None=None) -> Self: + # NOTE: we use a 1x1 conv2d to do the matmul. mxk @ kxn = (1,k,m,1).conv2d(n,k,1,1) + if not (self.ndim > 0 and w.ndim > 0): raise RuntimeError(f"both tensors need to be at least 1D, got {self.ndim=}, {w.ndim=}") + if self.shape[-1] != w.shape[-min(w.ndim, 2)]: raise RuntimeError(f"cannot image_dot {self.shape} and {w.shape}") + + bs, groups, cin, cout = prod(self.shape[0:-2]), prod(w.shape[0:-2]), w.shape[-2], w.shape[-1] + out_shape_t = self.shape[0:-2] + (cout,-1) if len(self.shape) > 1 else (cout,) + + # NOTE: with NHWC we can remove the transposes + # bs x groups*cin x H x W + cx = self.transpose(self.ndim-1, self.ndim-2).reshape(bs//groups, groups*cin, -1, 1) + # groups*cout x cin x H, W + cw = w.transpose(w.ndim-1, w.ndim-2).reshape(groups*cout, cin, 1, 1) + return cx.image_conv2d(cw, groups=groups, dtype=dtype).reshape(out_shape_t).transpose(self.ndim-1, self.ndim-2) + + def image_conv2d(self, weight:Self, bias:Self|None=None, groups=1, stride=1, dilation=1, padding=0, dtype=None) -> Self: + dtsz = 2 if FLOAT16 else 4 + + (bs,_,_,_), (cout,cin,H,W) = self.shape, weight.shape + assert isinstance(cin, int) and isinstance(cout, int) + x, w = self, weight.reshape(groups, (rcout := cout//groups), cin, H, W) + + padding_neg, padding_pos = [min(0, p) for p in resolve_pool_pads(padding, 2)], [max(0, p) for p in resolve_pool_pads(padding, 2)] + x = x.pad(padding_neg) + iy, ix = x.shape[2:] + + # hack for non multiples of 4 on cin + if cin % 4 != 0 and not (cin == 1 and groups%4 == 0): + new_cin = round_up(cin, 4) + w = w.pad_to(None, None, new_cin, None, None) + x = x.reshape(bs, groups, cin, iy, ix) + x = x.pad_to(None, None, new_cin, None, None).reshape(bs, groups*new_cin, iy, ix) + cin = new_cin + + # hack for non multiples of 4 on rcout + added_output_channels = 0 + if rcout % 4 != 0 and not (rcout == 1 and groups%4 == 0): + added_output_channels = 4 - (rcout % 4) + rcout += added_output_channels + cout = groups * rcout + w = w.pad_to(None, rcout, None, None, None) + + # packed (note: flipping bs and iy would make the auto-padding work) + x = x.permute(0,2,3,1) + cin_last = iy == 1 and ix == 1 + if cin == 1: w = w.reshape(cout//4,4,H,W).permute(0,2,3,1) + elif cin_last: w = w.reshape(cout//4,4,cin//4,4,H,W).permute(0,4,2,5,1,3) + else: w = w.reshape(cout//4,4,cin//4,4,H,W).permute(0,4,2,5,3,1) + + def is_pow2(v): return v > 0 and v & (v - 1) == 0 + # pad dimension i to amt with invalids + def ipad(t, i, amt): + return t.pad(tuple(None if d != i else (0, amt-s) for d,s in enumerate(t.shape)), value=Invalid) if amt != t.shape[i] else t + # align a dimension, use at to specify the dimension to pad in, defaults to first + def pad_align(t, dim, at=None, force=False): + # align to 64 pixels when height is real, otherwise 64 bytes is sufficient + align = (64 // dtsz) if prod(t.shape[:dim]) == 1 or prod(t.shape) < 16384 * 4 else 256 + return ipad(t, at:=at or dim, round_up(t.shape[at] + int(force), align // math.gcd(prod(t.shape[dim:]) // t.shape[at], align))) + + # bank conflicts + bank_conflict = cin >= 8 and is_pow2(cin // 4) + if bank_conflict: + x, w = pad_align(x.reshape(bs, iy, ix, groups, cin // 4, 4), 2, at=4, force=True), pad_align(w, 1, at=2, force=True) + else: x, w = pad_align(x, 2), pad_align(w, 1) + + # contiguous creates the image, and early realize static weights (TODO: test for the static weight) + if FLOAT16: x, w = x.cast(dtypes.half).contiguous().cast(dtypes.float), w.cast(dtypes.half).contiguous().cast(dtypes.float) + else: x, w = x.contiguous(), w.contiguous() + + # undo alignment hacks + if bank_conflict: x, w = x[:, :, :, :, :cin // 4, :], w[:, :, :cin // 4, ...] + else: x, w = x[:, :, :ix, :], w[:, :H, ...] + + # expand out + rcin_hi, rcin_lo = (cin//4, 4) if cin >= 4 else (1, 1) + group_shape, rcout_expand = (groups//4, 4) if cin == 1 else (groups, 1), (rcout//4, 4) if rcout >= 4 else (1, 1) + x = x.reshape(bs, iy, -1, groups, rcin_hi, rcin_lo) + if cin_last: w = w.reshape(cout//4, H, rcin_hi, W, 4, rcin_lo) + else: w = w.reshape(cout//4, H, rcin_hi, W, rcin_lo, 4).permute(0,1,2,3,5,4) + + # prepare input + x = x.permute(0,3,4,5,1,2).pad(padding_pos)._pool((H,W), stride, dilation)# -> (bs, groups, rcin_hi, rcin_lo, oy, ox, H, W) + x = x.permute(0,4,5,1,2,3,6,7).reshape(bs, (oy := x.shape[4]), (ox := x.shape[5]), *group_shape, 1, 1, rcin_hi, rcin_lo, H, W) + + # prepare weights + w = w.permute(0,4,2,5,1,3).reshape((1, 1, 1, *group_shape, *rcout_expand, rcin_hi, rcin_lo, H, W)) + + # the conv! + ret = (x*w).cast(dtypes.float32).sum((-4, -3, -2, -1), dtype=dtype) + + ret = ret.reshape(bs, oy, ox, groups, rcout) + # undo hack for non multiples of 4 on C.rcout + if added_output_channels: ret = ret[:, :, :, :, :-added_output_channels] + # NCHW output + ret = ret.reshape(bs, oy, ox, groups * (rcout - added_output_channels)).permute(0,3,1,2) + return ret if bias is None else ret.add(bias.reshape(1, -1, 1, 1)) + def conv2d(self, weight:Self, bias:Self|None=None, groups=1, stride=1, dilation=1, padding:int|Sequence[int]=0, dtype:DTypeLike|None=None) -> Self: + """ + Applies a convolution over a tensor with a given `weight` and optional `bias`. + + This function supports three different types of `padding` + + 1. `int` (single value): + Applies the same padding value uniformly to all spatial dimensions. + + 2. `tuple[int, ...]` (length = number of spatial dimensions): + Specifies a distinct padding value for each spatial dimension in the form `(padding_height, padding_width, ...)`. + + 3. `tuple[int, ...]` (length = 2 * number of spatial dimensions): + Specifies explicit padding for each side of each spatial dimension in the form + `(padding_left, padding_right, padding_top, padding_bottom, ...)`. + + NOTE: unlike PyTorch, this implementation is not limited to only 2d convolutions and instead works for any number of dimensions. + + See: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html + + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor.arange(9).reshape(1, 1, 3, 3) + w = Tensor.ones(1, 1, 2, 2) + print(t.conv2d(w).numpy()) + ``` + """ + if IMAGE: return self.image_conv2d(weight, bias, groups, stride, dilation, padding, dtype) + if WINO and all(x == 3 for x in weight.shape[2:]) and stride == dilation == 1: return self._conv2d_winograd(weight, bias, groups, padding, dtype) (bs,cin_), (cout,cin), HW = self.shape[:2], weight.shape[:2], weight.shape[2:] padding_ = resolve_pool_pads(padding, len(HW)) assert groups*cin == cin_ and len(self.shape) == len(weight.shape),\ diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 6c68fe9420..2f72873ec3 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1,13 +1,11 @@ # inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py from __future__ import annotations -import time, math, itertools, functools, sys, inspect, pathlib, hashlib, weakref +import time, math, functools, sys, inspect, pathlib, hashlib, weakref from typing import Any, Callable, Sequence, cast, get_args, ParamSpec, TypeVar, Generic, TYPE_CHECKING if TYPE_CHECKING: import numpy -from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, to_dtype -from tinygrad.dtype import _from_np_dtype, _to_np_dtype, PyConst, Invalid -from tinygrad.helpers import argfix, flatten, prod, all_int, round_up, getenv, fully_flatten, ceildiv, fetch, flat_to_grouped -from tinygrad.helpers import resolve_pool_pads, IMAGE, FLOAT16, WINO, Metadata, TRACEMETA, is_numpy_ndarray, TracingKey, cpu_profile -from tinygrad.helpers import suppress_finalizing, disable_gc +from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, to_dtype, _from_np_dtype, _to_np_dtype, PyConst +from tinygrad.helpers import argfix, prod, all_int, getenv, fully_flatten, ceildiv, fetch, Metadata, TRACEMETA, is_numpy_ndarray, TracingKey +from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable, _broadcast_shape from tinygrad.mixin.rand import RandMixin from tinygrad.schedule import create_linear_with_vars @@ -42,22 +40,6 @@ def _fromnp(x: 'numpy.ndarray') -> UOp: ret.buffer.allocate(x) return ret.reshape(x.shape) -def _get_winograd_matcols(mat, dims:int, shp:tuple[sint, ...], dtype:DType) -> list[list[Tensor]]: - return [[Tensor.cat(*[Tensor.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(mat, t:Tensor, dims:int) -> Tensor: - # 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_ = t.reshape(t.shape[:dims] + (1,) * dims + t.shape[dims:]).expand(t.shape[:dims] + (len(mat),) * dims + t.shape[dims:]) # add output dims - # precalculate mat columns for each dim; prod(itertools.product(matcols)) gives the columns of kron(mat, mat, ...) - matcols = _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 isinstance(ret, Tensor), "sum didn't return a Tensor" - return ret - class Tensor(RandMixin): """ A `Tensor` is a multi-dimensional matrix containing elements of a single data type. @@ -636,82 +618,6 @@ class Tensor(RandMixin): chunks = ceildiv(chunks, 65536) return data.pad_to(2**20).unsqueeze(0)._hash_1mb().flatten()[:16] - # ***** processing ops ***** - - # TODO: winograd can be a rewrite rule like split_reduceop - def _conv2d_winograd(self, weight:Tensor, bias:Tensor|None, groups:int, padding:int|Sequence[int], dtype:DTypeLike|None) -> Tensor: - (bs,cin_), (cout,cin), HW = self.shape[:2], weight.shape[:2], weight.shape[2:] - padding_ = resolve_pool_pads(padding, len(HW)) - assert groups*cin == cin_ and len(self.shape) == len(weight.shape),\ - f"Input Tensor shape {self.shape} does not match the shape of the weights {weight.shape}. ({groups*cin} vs. {cin_})" - rcout, oyx = cout//groups, self.pad(padding_)._pool(HW, 1, 1).shape[2:-len(HW)] - HWI, HWO = (6,) * len(HW), (4,) * len(HW) # F(4x4,3x3) winograd tiles - winograd_G = [[1/4, 0, 0], [-1/6, -1/6, -1/6], [-1/6, 1/6, -1/6], [1/24, 1/12, 1/6], [1/24, -1/12, 1/6], [0, 0, 1]] - winograd_Bt = [[4, 0, -5, 0, 1, 0], [0, -4, -4, 1, 1, 0], [0, 4, -4, -1, 1, 0], [0, -2, -1, 2, 1, 0], [0, 2, -1, -2, 1, 0], [0, 4, 0, -5, 0, 1]] - winograd_At = [[1, 1, 1, 1, 1, 0], [0, 1, -1, 2, -2, 0], [0, 1, 1, 4, 4, 0], [0, 1, -1, 8, -8, 1]] # applying At in pre-order doubles compile time - - # TODO: stride == dilation - # use padding to round up to 4x4 output tiles - # (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))) - 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 - - # compute 6x6 winograd tiles: GgGt, BtdB. contiguous so the transforms are materialized once - # (HWI, groups * rcout, cin) -> (HWI, bs=1, groups, rcout, cin, tyx=(1,1)) - gfactors = _apply_winograd_matrix(winograd_G, g, len(HW)).contiguous().reshape(*HWI, 1, groups, rcout, cin, *([1]*len(tyx))) - # (HWI, bs, cin_, tyx) -> (HWI, bs, groups, 1 ,cin, *tyx) - dfactors = _apply_winograd_matrix(winograd_Bt, d, len(HW)).contiguous().reshape(*HWI, bs, groups, 1, cin, *tyx) - - # matmul; sum across cin: (HWI, bs, groups, rcout, *tyx); then HWI -> HWO: (HWO, bs, groups, rcout, *tyx) - ret = _apply_winograd_matrix(winograd_At, (gfactors * dfactors).sum(axis=-1-len(HW), dtype=dtype), len(HW)) - - # interleave tyx and HWO: (bs, groups, rcout, oy, HO, ox, WO) - ret = ret.permute([*range(len(HW), len(ret.shape)-len(HW)), *[i+o for i in range(len(HW)) for o in [len(ret.shape)-len(HW),0]]]) - # merge groups and rcout, tyx and HWO: (bs, groups, cout, *yx), shrink to final - ret = ret.reshape(bs, cout, *[c * HWO[i] for i, c in enumerate(tyx)]).shrink_to(bs, cout, *oyx) - - return (ret if bias is None else ret.add(bias.reshape(1, -1, *[1 for _ in range(len(HW))]))).contiguous().contiguous_backward() - - def conv2d(self, weight:Tensor, bias:Tensor|None=None, groups=1, stride=1, dilation=1, padding:int|Sequence[int]=0, - dtype:DTypeLike|None=None) -> Tensor: - """ - Applies a convolution over a tensor with a given `weight` and optional `bias`. - - This function supports three different types of `padding` - - 1. `int` (single value): - Applies the same padding value uniformly to all spatial dimensions. - - 2. `tuple[int, ...]` (length = number of spatial dimensions): - Specifies a distinct padding value for each spatial dimension in the form `(padding_height, padding_width, ...)`. - - 3. `tuple[int, ...]` (length = 2 * number of spatial dimensions): - Specifies explicit padding for each side of each spatial dimension in the form - `(padding_left, padding_right, padding_top, padding_bottom, ...)`. - - NOTE: unlike PyTorch, this implementation is not limited to only 2d convolutions and instead works for any number of dimensions. - - See: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html - - ```python exec="true" source="above" session="tensor" result="python" - t = Tensor.arange(9).reshape(1, 1, 3, 3) - w = Tensor.ones(1, 1, 2, 2) - print(t.conv2d(w).numpy()) - ``` - """ - if IMAGE: return self.image_conv2d(weight, bias, groups, stride, dilation, padding, dtype) - if WINO and all(x == 3 for x in weight.shape[2:]) and stride == dilation == 1: return self._conv2d_winograd(weight, bias, groups, padding, dtype) - return super().conv2d(weight, bias, groups, stride, dilation, padding, dtype) - - def dot(self, w:Tensor, dtype:DTypeLike|None=None) -> Tensor: - if IMAGE: return self.image_dot(w, dtype) - return super().dot(w, dtype) - # ***** broadcasted elementwise ops ***** def where(self:Tensor, x:Tensor|ConstType|sint, y:Tensor|ConstType|sint) -> Tensor: @@ -798,105 +704,6 @@ class Tensor(RandMixin): return Tensor.stack(*(tmp>>8*i*ns for i in range(os//ns)), dim=-1).flatten(-2).cast(new_uint).bitcast(dtype) return self._apply_uop(UOp.bitcast, dtype=dt) if self.dtype != dt else self - # *** image Tensor function replacements *** - - def image_dot(self, w:Tensor, dtype:DTypeLike|None=None) -> Tensor: - # NOTE: we use a 1x1 conv2d to do the matmul. mxk @ kxn = (1,k,m,1).conv2d(n,k,1,1) - if not (self.ndim > 0 and w.ndim > 0): raise RuntimeError(f"both tensors need to be at least 1D, got {self.ndim=}, {w.ndim=}") - if self.shape[-1] != w.shape[-min(w.ndim, 2)]: raise RuntimeError(f"cannot image_dot {self.shape} and {w.shape}") - - bs, groups, cin, cout = prod(self.shape[0:-2]), prod(w.shape[0:-2]), w.shape[-2], w.shape[-1] - out_shape_t = self.shape[0:-2] + (cout,-1) if len(self.shape) > 1 else (cout,) - - # NOTE: with NHWC we can remove the transposes - # bs x groups*cin x H x W - cx = self.transpose(self.ndim-1, self.ndim-2).reshape(bs//groups, groups*cin, -1, 1) - # groups*cout x cin x H, W - cw = w.transpose(w.ndim-1, w.ndim-2).reshape(groups*cout, cin, 1, 1) - return cx.image_conv2d(cw, groups=groups, dtype=dtype).reshape(out_shape_t).transpose(self.ndim-1, self.ndim-2) - - def image_conv2d(self, weight:Tensor, bias:Tensor|None=None, groups=1, stride=1, dilation=1, padding=0, dtype=None) -> Tensor: - dtsz = 2 if FLOAT16 else 4 - - (bs,_,_,_), (cout,cin,H,W) = self.shape, weight.shape - assert isinstance(cin, int) and isinstance(cout, int) - x, w = self, weight.reshape(groups, (rcout := cout//groups), cin, H, W) - - padding_neg, padding_pos = [min(0, p) for p in resolve_pool_pads(padding, 2)], [max(0, p) for p in resolve_pool_pads(padding, 2)] - x = x.pad(padding_neg) - iy, ix = x.shape[2:] - - # hack for non multiples of 4 on cin - if cin % 4 != 0 and not (cin == 1 and groups%4 == 0): - new_cin = round_up(cin, 4) - w = w.pad_to(None, None, new_cin, None, None) - x = x.reshape(bs, groups, cin, iy, ix) - x = x.pad_to(None, None, new_cin, None, None).reshape(bs, groups*new_cin, iy, ix) - cin = new_cin - - # hack for non multiples of 4 on rcout - added_output_channels = 0 - if rcout % 4 != 0 and not (rcout == 1 and groups%4 == 0): - added_output_channels = 4 - (rcout % 4) - rcout += added_output_channels - cout = groups * rcout - w = w.pad_to(None, rcout, None, None, None) - - # packed (note: flipping bs and iy would make the auto-padding work) - x = x.permute(0,2,3,1) - cin_last = iy == 1 and ix == 1 - if cin == 1: w = w.reshape(cout//4,4,H,W).permute(0,2,3,1) - elif cin_last: w = w.reshape(cout//4,4,cin//4,4,H,W).permute(0,4,2,5,1,3) - else: w = w.reshape(cout//4,4,cin//4,4,H,W).permute(0,4,2,5,3,1) - - def is_pow2(v): return v > 0 and v & (v - 1) == 0 - # pad dimension i to amt with invalids - def ipad(t, i, amt): - return t.pad(tuple(None if d != i else (0, amt-s) for d,s in enumerate(t.shape)), value=Invalid) if amt != t.shape[i] else t - # align a dimension, use at to specify the dimension to pad in, defaults to first - def pad_align(t, dim, at=None, force=False): - # align to 64 pixels when height is real, otherwise 64 bytes is sufficient - align = (64 // dtsz) if prod(t.shape[:dim]) == 1 or prod(t.shape) < 16384 * 4 else 256 - return ipad(t, at:=at or dim, round_up(t.shape[at] + int(force), align // math.gcd(prod(t.shape[dim:]) // t.shape[at], align))) - - # bank conflicts - bank_conflict = cin >= 8 and is_pow2(cin // 4) - if bank_conflict: - x, w = pad_align(x.reshape(bs, iy, ix, groups, cin // 4, 4), 2, at=4, force=True), pad_align(w, 1, at=2, force=True) - else: x, w = pad_align(x, 2), pad_align(w, 1) - - # contiguous creates the image, and early realize static weights (TODO: test for the static weight) - if FLOAT16: x, w = x.cast(dtypes.half).contiguous().cast(dtypes.float), w.cast(dtypes.half).contiguous().cast(dtypes.float) - else: x, w = x.contiguous(), w.contiguous() - - # undo alignment hacks - if bank_conflict: x, w = x[:, :, :, :, :cin // 4, :], w[:, :, :cin // 4, ...] - else: x, w = x[:, :, :ix, :], w[:, :H, ...] - - # expand out - rcin_hi, rcin_lo = (cin//4, 4) if cin >= 4 else (1, 1) - group_shape, rcout_expand = (groups//4, 4) if cin == 1 else (groups, 1), (rcout//4, 4) if rcout >= 4 else (1, 1) - x = x.reshape(bs, iy, -1, groups, rcin_hi, rcin_lo) - if cin_last: w = w.reshape(cout//4, H, rcin_hi, W, 4, rcin_lo) - else: w = w.reshape(cout//4, H, rcin_hi, W, rcin_lo, 4).permute(0,1,2,3,5,4) - - # prepare input - x = x.permute(0,3,4,5,1,2).pad(padding_pos)._pool((H,W), stride, dilation)# -> (bs, groups, rcin_hi, rcin_lo, oy, ox, H, W) - x = x.permute(0,4,5,1,2,3,6,7).reshape(bs, (oy := x.shape[4]), (ox := x.shape[5]), *group_shape, 1, 1, rcin_hi, rcin_lo, H, W) - - # prepare weights - w = w.permute(0,4,2,5,1,3).reshape((1, 1, 1, *group_shape, *rcout_expand, rcin_hi, rcin_lo, H, W)) - - # the conv! - ret = (x*w).cast(dtypes.float32).sum((-4, -3, -2, -1), dtype=dtype) - - ret = ret.reshape(bs, oy, ox, groups, rcout) - # undo hack for non multiples of 4 on C.rcout - if added_output_channels: ret = ret[:, :, :, :, :-added_output_channels] - # NCHW output - ret = ret.reshape(bs, oy, ox, groups * (rcout - added_output_channels)).permute(0,3,1,2) - return ret if bias is None else ret.add(bias.reshape(1, -1, 1, 1)) - P = ParamSpec("P") T = TypeVar("T")