forked from tinygrad/tinygrad
remove Tensor conv2d and dot override [PR] (#16792)
This commit is contained in:
chenyu
@@ -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))
|
||||
|
||||
+187
-3
@@ -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),\
|
||||
|
||||
+4
-197
@@ -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")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user