forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Prebuilt Release @ ab07000
This commit is contained in:
IQ.Lvbs history cleanup
6
tinygrad_repo/extra/thunder/tiny/tk/__init__.py
Normal file
6
tinygrad_repo/extra/thunder/tiny/tk/__init__.py
Normal file
@@ -0,0 +1,6 @@
|
||||
from tinygrad.device import Device
|
||||
|
||||
if Device.DEFAULT == "AMD":
|
||||
WARP_THREADS = 64
|
||||
else:
|
||||
WARP_THREADS = 32
|
||||
491
tinygrad_repo/extra/thunder/tiny/tk/group.py
Normal file
491
tinygrad_repo/extra/thunder/tiny/tk/group.py
Normal file
@@ -0,0 +1,491 @@
|
||||
import math
|
||||
from typing import cast, Callable
|
||||
from tinygrad import dtypes
|
||||
from tinygrad.uop.ops import AxisType, UOp, Ops
|
||||
from tinygrad.dtype import AddrSpace, PtrDType
|
||||
from tinygrad.helpers import prod
|
||||
|
||||
from extra.thunder.tiny.tk import WARP_THREADS
|
||||
from extra.thunder.tiny.tk.tiles import ALL_TILES, ST, RT, RV, TileLayout, VecLayout
|
||||
|
||||
class Group:
|
||||
def __init__(self, warps:int, ker):
|
||||
self.warps = warps
|
||||
self.group_threads = warps * WARP_THREADS
|
||||
self.ker = ker
|
||||
|
||||
# helpers
|
||||
@property
|
||||
def laneid(self): return self.ker.threadIdx_x % self.group_threads
|
||||
@property
|
||||
def warpid(self): return self.laneid // WARP_THREADS
|
||||
@property
|
||||
def groupid(self): return self.ker.threadIdx_x // self.group_threads
|
||||
|
||||
# ops that only work on a single warp
|
||||
|
||||
def clear(self, reg:ALL_TILES, value:float=0):
|
||||
reg = cast(UOp, reg)
|
||||
assert self.warps == 1
|
||||
|
||||
rngs_for_shape = tuple(self.ker.raw_range(dim) for dim in reg.shape)
|
||||
|
||||
reg_store = reg[*rngs_for_shape].store(value).end(*rngs_for_shape)
|
||||
|
||||
self.ker.push_store(reg_store, reg)
|
||||
return reg.after(reg_store).reshape(reg.shape)
|
||||
|
||||
def zero(self, reg:ALL_TILES): return self.clear(reg, 0)
|
||||
def ones(self, reg:ALL_TILES): return self.clear(reg, 1)
|
||||
def neg_inf(self, reg:ALL_TILES): return self.clear(reg, -math.inf)
|
||||
|
||||
def copy(self, dst:ALL_TILES, src:ALL_TILES):
|
||||
dst, src = cast(UOp, dst), cast(UOp, src)
|
||||
assert self.warps == 1
|
||||
assert dst.shape == src.shape
|
||||
|
||||
rngs_for_shape = tuple(self.ker.raw_range(dim) for dim in dst.shape)
|
||||
|
||||
src_load = src[*rngs_for_shape]
|
||||
if src.dtype.base != dst.dtype.base:
|
||||
src_load = src_load.cast(dst.dtype.base)
|
||||
dst_store = dst[*rngs_for_shape].store(src_load).end(*rngs_for_shape)
|
||||
|
||||
self.ker.push_store(dst_store, dst)
|
||||
return dst.after(dst_store).reshape(dst.shape)
|
||||
|
||||
def transpose(self, dst:UOp|RT, src:UOp|RT):
|
||||
dst, src = cast(UOp, dst), cast(UOp, src)
|
||||
assert self.warps == 1
|
||||
|
||||
for height in self.ker.range(src.shape[-3], track=False):
|
||||
for width in self.ker.range(src.shape[-2], track=False):
|
||||
for inner in self.ker.range(src.shape[-1], track=False):
|
||||
src_load = src[height, width, inner]
|
||||
if src.dtype.base != dst.dtype.base:
|
||||
src_load = src_load.cast(dst.dtype.base)
|
||||
dst_store = dst[width, height, inner].store(src_load).end(height, width, inner)
|
||||
|
||||
self.ker.push_store(dst_store, dst)
|
||||
return dst.after(dst_store).reshape(dst.shape)
|
||||
|
||||
def mma_AB(self, c:UOp|RT, a:UOp|RT, b:UOp|RT):
|
||||
c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
|
||||
assert self.warps == 1
|
||||
|
||||
a_base_shape = cast(RT, a).base_shape
|
||||
if a_base_shape.cols == 16:
|
||||
wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
|
||||
elif a_base_shape.cols == 32:
|
||||
wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
|
||||
else: raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}")
|
||||
|
||||
for height in self.ker.range(c.shape[-3], track=False):
|
||||
for width in self.ker.range(c.shape[-2], track=False):
|
||||
for inner in self.ker.range(a.shape[-2], axis_type=AxisType.REDUCE, track=False):
|
||||
if a_base_shape.cols == 16:
|
||||
a_in = UOp.vectorize(*[a[height, inner, i] for i in range(4)])
|
||||
b_in = UOp.vectorize(*[b[inner, width, i] for i in range(4)])
|
||||
elif a_base_shape.cols == 32:
|
||||
a_in = UOp.vectorize(*[a[height, inner, i] for i in range(8)])
|
||||
b_in = UOp.vectorize(*[b[inner, width, i] for i in range(8)])
|
||||
else: raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}")
|
||||
d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
|
||||
|
||||
out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
|
||||
c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
|
||||
c_store = UOp.group(*c_i).end(height, width, inner)
|
||||
|
||||
self.ker.push_store(c_store, c)
|
||||
return c.after(c_store).reshape(c.shape)
|
||||
|
||||
def mma_ABt(self, c:UOp|RT, a:UOp|RT, b:UOp|RT):
|
||||
c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
|
||||
assert self.warps == 1
|
||||
|
||||
a_base_shape = cast(RT, a).base_shape
|
||||
if a_base_shape.cols == 16:
|
||||
wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
|
||||
elif a_base_shape.cols == 32:
|
||||
wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
|
||||
else: raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}")
|
||||
|
||||
for height in self.ker.range(c.shape[-3], track=False):
|
||||
for width in self.ker.range(c.shape[-2], track=False):
|
||||
for inner in self.ker.range(a.shape[-2], axis_type=AxisType.REDUCE, track=False):
|
||||
if a_base_shape.cols == 16:
|
||||
a_in = UOp.vectorize(*[a[height, inner, i] for i in range(4)])
|
||||
b_in = UOp.vectorize(*[b[width, inner, i] for i in range(4)])
|
||||
elif a_base_shape.cols == 32:
|
||||
a_in = UOp.vectorize(*[a[height, inner, i] for i in range(8)])
|
||||
b_in = UOp.vectorize(*[b[width, inner, i] for i in range(8)])
|
||||
else: raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}")
|
||||
d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
|
||||
|
||||
out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
|
||||
c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
|
||||
c_store = UOp.group(*c_i).end(height, width, inner)
|
||||
|
||||
self.ker.push_store(c_store, c)
|
||||
return c.after(c_store).reshape(c.shape)
|
||||
|
||||
def mma_AtB(self, c:UOp|RT, a:UOp|RT, b:UOp|RT):
|
||||
c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
|
||||
assert self.warps == 1
|
||||
|
||||
a_base_shape = cast(RT, a).base_shape
|
||||
if a_base_shape.cols == 16:
|
||||
wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
|
||||
elif a_base_shape.cols == 32:
|
||||
wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
|
||||
else: raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}")
|
||||
|
||||
for height in self.ker.range(c.shape[-3], track=False):
|
||||
for width in self.ker.range(c.shape[-2], track=False):
|
||||
for inner in self.ker.range(a.shape[-3], axis_type=AxisType.REDUCE, track=False):
|
||||
if a_base_shape.cols == 16:
|
||||
a_in = UOp.vectorize(*[a[inner, height, i] for i in range(4)])
|
||||
b_in = UOp.vectorize(*[b[inner, width, i] for i in range(4)])
|
||||
elif a_base_shape.cols == 32:
|
||||
a_in = UOp.vectorize(*[a[inner, height, i] for i in range(8)])
|
||||
b_in = UOp.vectorize(*[b[inner, width, i] for i in range(8)])
|
||||
else: raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}")
|
||||
d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
|
||||
|
||||
out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
|
||||
c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
|
||||
c_store = UOp.group(*c_i).end(height, width, inner)
|
||||
|
||||
self.ker.push_store(c_store, c)
|
||||
return c.after(c_store).reshape(c.shape)
|
||||
|
||||
def mma_AtBt(self, c:UOp|RT, a:UOp|RT, b:UOp|RT):
|
||||
c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
|
||||
assert self.warps == 1
|
||||
|
||||
a_base_shape = cast(RT, a).base_shape
|
||||
if a_base_shape.cols == 16:
|
||||
wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
|
||||
elif a_base_shape.cols == 32:
|
||||
wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
|
||||
else: raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}")
|
||||
|
||||
for height in self.ker.range(c.shape[-3], track=False):
|
||||
for width in self.ker.range(c.shape[-2], track=False):
|
||||
for inner in self.ker.range(a.shape[-3], axis_type=AxisType.REDUCE, track=False):
|
||||
if a_base_shape.cols == 16:
|
||||
a_in = UOp.vectorize(*[a[inner, height, i] for i in range(4)])
|
||||
b_in = UOp.vectorize(*[b[width, inner, i] for i in range(4)])
|
||||
elif a_base_shape.cols == 32:
|
||||
a_in = UOp.vectorize(*[a[inner, height, i] for i in range(8)])
|
||||
b_in = UOp.vectorize(*[b[width, inner, i] for i in range(8)])
|
||||
else: raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}")
|
||||
d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
|
||||
|
||||
out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
|
||||
c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
|
||||
c_store = UOp.group(*c_i).end(height, width, inner)
|
||||
|
||||
self.ker.push_store(c_store, c)
|
||||
return c.after(c_store).reshape(c.shape)
|
||||
|
||||
def map(self, a:ALL_TILES, op:Callable[[UOp], UOp]|Callable[[UOp, tuple], UOp]):
|
||||
a = cast(UOp, a)
|
||||
assert self.warps == 1
|
||||
|
||||
rngs_for_shape = tuple(self.ker.raw_range(dim) for dim in a.shape)
|
||||
|
||||
if op.__code__.co_argcount == 1:
|
||||
to_store = op(a[*rngs_for_shape]) # type: ignore
|
||||
else:
|
||||
to_store = op(a[*rngs_for_shape], rngs_for_shape) # type: ignore
|
||||
|
||||
a_store = a[*rngs_for_shape].store(to_store).end(*rngs_for_shape)
|
||||
|
||||
self.ker.push_store(a_store, a)
|
||||
return a.after(a_store).reshape(a.shape)
|
||||
|
||||
def row_reduce(self, vec:UOp|RV, src:UOp|RT, op:Callable[[UOp, UOp], UOp], init_value:float=0.0):
|
||||
vec, src = cast(UOp, vec), cast(UOp, src)
|
||||
assert self.warps == 1
|
||||
|
||||
red_local = self.ker.alloc((self.group_threads,), src.dtype.base, AddrSpace.LOCAL)
|
||||
red_reg = self.ker.alloc((1,), src.dtype.base, AddrSpace.REG)
|
||||
|
||||
for height in self.ker.range(src.shape[-3], track=False):
|
||||
i = self.ker.raw_range(red_reg.size)
|
||||
red_reg = red_reg.after(height, *[tkr._rng for tkr in self.ker.range_stack])
|
||||
reg_store = red_reg.flatten()[i].store(init_value).end(i)
|
||||
red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
|
||||
|
||||
for width in self.ker.range(src.shape[-2], axis_type=AxisType.REDUCE, track=False):
|
||||
for inner in self.ker.range(4, axis_type=AxisType.REDUCE, track=False):
|
||||
reg_store = red_reg[0].store(op(red_reg[0], src[height, width, inner])).end(width, inner)
|
||||
red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
|
||||
|
||||
# store to shared memory
|
||||
red_local_store = red_local[self.laneid].store(red_reg[0])
|
||||
red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
|
||||
|
||||
# reduce from shared memory
|
||||
for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
|
||||
offset = (self.laneid + (1 + inner) * 16) % self.group_threads
|
||||
reg_store = red_reg[0].store(op(red_reg[0], red_local[offset])).end(inner)
|
||||
red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
|
||||
|
||||
# reduce with vec
|
||||
vec_store = vec[height, 0].store(op(vec[height, 0], red_reg[0])).end(height)
|
||||
|
||||
self.ker.push_store(vec_store, vec)
|
||||
return vec.after(vec_store).reshape(vec.shape)
|
||||
|
||||
def col_reduce(self, vec:UOp|RV, src:UOp|RT, op:Callable[[UOp, UOp], UOp], init_value:float=0.0):
|
||||
vec, src = cast(UOp, vec), cast(UOp, src)
|
||||
assert self.warps == 1
|
||||
|
||||
red_local = self.ker.alloc((self.group_threads,), src.dtype.base, AddrSpace.LOCAL)
|
||||
red_reg = self.ker.alloc((1,), src.dtype.base, AddrSpace.REG)
|
||||
|
||||
for width in self.ker.range(src.shape[-2], track=False):
|
||||
i = self.ker.raw_range(red_reg.size)
|
||||
red_reg = red_reg.after(width, *[tkr._rng for tkr in self.ker.range_stack])
|
||||
reg_store = red_reg.flatten()[i].store(init_value).end(i)
|
||||
red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
|
||||
|
||||
for height in self.ker.range(src.shape[-3], axis_type=AxisType.REDUCE, track=False):
|
||||
for inner in self.ker.range(4, axis_type=AxisType.REDUCE, track=False):
|
||||
reg_store = red_reg[0].store(op(red_reg[0], src[height, width, inner])).end(height, inner)
|
||||
red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
|
||||
|
||||
# store to shared memory
|
||||
red_local_store = red_local[self.laneid].store(red_reg[0])
|
||||
red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
|
||||
|
||||
# reduce from shared memory
|
||||
for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
|
||||
offset = (self.laneid + (1 + inner) * 16) % self.group_threads
|
||||
reg_store = red_reg[0].store(op(red_reg[0], red_local[offset])).end(inner)
|
||||
red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
|
||||
|
||||
# reduce with vec
|
||||
vec_store = vec[width, 0].store(op(vec[width, 0], red_reg[0])).end(width)
|
||||
|
||||
self.ker.push_store(vec_store, vec)
|
||||
return vec.after(vec_store).reshape(vec.shape)
|
||||
|
||||
# ops that can work across multiple warps
|
||||
|
||||
def load(self, dst:ALL_TILES, src:ALL_TILES, dst_idxs:tuple[UOp|int,...]=(), idxs:tuple[UOp|int,...]=(), axis:int=0):
|
||||
dst, src = cast(UOp, dst), cast(UOp, src)
|
||||
assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
|
||||
dst_dtype, src_dtype = dst.dtype, src.dtype
|
||||
if dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.LOCAL:
|
||||
laneid = self.ker.laneid
|
||||
rt, st = cast(RT, dst), cast(ST, src)
|
||||
elements_per_thread = rt.base_shape.elements_per_thread
|
||||
|
||||
for height in self.ker.range(dst.shape[-3], track=False):
|
||||
for width in self.ker.range(dst.shape[-2], track=False):
|
||||
for inner in self.ker.range(elements_per_thread, track=False):
|
||||
if rt.layout != st.layout:
|
||||
row = rt.base_shape.stride * (laneid // rt.base_shape.cols) + inner
|
||||
col = laneid % rt.base_shape.cols
|
||||
else:
|
||||
row = laneid % rt.base_shape.rows
|
||||
col = rt.base_shape.stride * (laneid // rt.base_shape.rows) + inner
|
||||
|
||||
sheight = height
|
||||
swidth = width
|
||||
if len(idxs) == 2:
|
||||
row_idx = idxs[0] * dst.shape[-3] * rt.base_shape.rows
|
||||
col_idx = idxs[1] * dst.shape[-2] * rt.base_shape.cols
|
||||
|
||||
row += row_idx % st.base_shape.rows
|
||||
col += col_idx % st.base_shape.cols
|
||||
sheight += row_idx // st.base_shape.rows
|
||||
swidth += col_idx // st.base_shape.cols
|
||||
|
||||
srow, scol = cast(ST, src).swizzle(row, col)
|
||||
|
||||
src_load = src[*idxs[:-2], sheight, swidth, srow, scol]
|
||||
if src.dtype.base != dst.dtype.base:
|
||||
src_load = src_load.cast(dst.dtype.base)
|
||||
dst_store = dst[*dst_idxs, height, width, inner].store(src_load)
|
||||
dst_store = dst_store.end(height, width, inner)
|
||||
elif dst_dtype.addrspace == AddrSpace.LOCAL and src_dtype.addrspace == AddrSpace.GLOBAL:
|
||||
srcf = src.flatten()
|
||||
row_stride = prod(src.shape[axis+1:])
|
||||
|
||||
st = cast(ST, dst)
|
||||
idxs = tuple(idx * st.rows if i == axis else idx for i, idx in enumerate(idxs))
|
||||
idxs = tuple(idx * st.cols if i == 3 else idx for i, idx in enumerate(idxs))
|
||||
src_i = ((idxs[0] * src.shape[-3] + idxs[1]) * src.shape[-2] + idxs[2]) * src.shape[-1] + idxs[3]
|
||||
|
||||
elements_per_thread = st.base_shape.elements_per_thread
|
||||
memcpy_per_row = st.cols // elements_per_thread
|
||||
total_calls = (dst.shape[-4] * dst.shape[-3] * st.base_shape.num_elements) // (self.group_threads * elements_per_thread)
|
||||
|
||||
for outer in self.ker.range(total_calls, track=False):
|
||||
for inner in self.ker.range(elements_per_thread, axis_type=AxisType.UPCAST, track=False):
|
||||
load_idx = outer * self.group_threads + self.laneid
|
||||
row = load_idx // memcpy_per_row
|
||||
col = (load_idx * elements_per_thread) % st.cols + inner
|
||||
height = row // st.base_shape.rows
|
||||
width = col // st.base_shape.cols
|
||||
|
||||
row = row % st.base_shape.rows
|
||||
col = col % st.base_shape.cols
|
||||
|
||||
srow, scol = cast(ST, dst).swizzle(row, col)
|
||||
|
||||
src_i += height * st.base_shape.rows * row_stride + width * st.base_shape.cols
|
||||
src_i += row * row_stride + col
|
||||
|
||||
src_load = srcf[src_i]
|
||||
if src.dtype.base != dst.dtype.base:
|
||||
src_load = src_load.cast(dst.dtype.base)
|
||||
dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load)
|
||||
dst_store = dst_store.end(height, width, outer, inner).barrier()
|
||||
elif dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT):
|
||||
srcf = src.flatten()
|
||||
row_stride = prod(src.shape[axis+1:])
|
||||
|
||||
laneid = self.ker.laneid
|
||||
rt = cast(RT, dst)
|
||||
elements_per_thread = rt.base_shape.elements_per_thread
|
||||
|
||||
idxs = tuple(idx * dst.shape[-3] * rt.base_shape.rows if i == axis else idx for i, idx in enumerate(idxs))
|
||||
idxs = tuple(idx * dst.shape[-2] * rt.base_shape.cols if i == 3 else idx for i, idx in enumerate(idxs))
|
||||
src_i = ((idxs[0] * src.shape[-3] + idxs[1]) * src.shape[-2] + idxs[2]) * src.shape[-1] + idxs[3]
|
||||
|
||||
for height in self.ker.range(dst.shape[-3], track=False):
|
||||
for width in self.ker.range(dst.shape[-2], track=False):
|
||||
for inner in self.ker.range(elements_per_thread, track=False):
|
||||
base_row = height * rt.base_shape.rows
|
||||
base_col = width * rt.base_shape.cols
|
||||
|
||||
if rt.layout == TileLayout.COL:
|
||||
row = rt.base_shape.stride * (laneid // rt.base_shape.cols) + inner
|
||||
col = laneid % rt.base_shape.cols
|
||||
else:
|
||||
row = laneid % rt.base_shape.rows
|
||||
col = rt.base_shape.stride * (laneid // rt.base_shape.rows) + inner
|
||||
|
||||
srow, scol = base_row + row, base_col + col
|
||||
|
||||
src_i += srow * row_stride + scol
|
||||
|
||||
src_load = srcf[src_i]
|
||||
if src.dtype.base != dst.dtype.base:
|
||||
src_load = src_load.cast(dst.dtype.base)
|
||||
dst_store = dst[*dst_idxs, height, width, inner].store(src_load).end(height, width, inner)
|
||||
elif dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.GLOBAL and isinstance(dst, RV):
|
||||
srcf = src.flatten()
|
||||
row_stride = prod(src.shape[axis+1:])
|
||||
|
||||
laneid = self.ker.laneid
|
||||
rv = cast(RV, dst)
|
||||
reductions = rv.base_shape.rows
|
||||
|
||||
assert rv.layout == VecLayout.ORTHO, "only ortho layout supported"
|
||||
|
||||
idxs = tuple(idx * rv.length if i == 3 else idx for i, idx in enumerate(idxs))
|
||||
src_i = ((idxs[0] * src.shape[-3] + idxs[1]) * src.shape[-2] + idxs[2]) * src.shape[-1] + idxs[3]
|
||||
|
||||
for outer in self.ker.range(dst.shape[-2], track=False):
|
||||
src_i += outer * reductions + (laneid % reductions)
|
||||
|
||||
src_load = srcf[src_i]
|
||||
if src.dtype.base != dst.dtype.base:
|
||||
src_load = src_load.cast(dst.dtype.base)
|
||||
dst_store = dst[outer, 0].store(src_load).end(outer)
|
||||
else:
|
||||
raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented for {type(dst)=}")
|
||||
|
||||
self.ker.push_store(dst_store, dst)
|
||||
return dst.after(dst_store).reshape(dst.shape)
|
||||
|
||||
def store(self, dst:ALL_TILES, src:ALL_TILES, idxs:tuple[UOp|int,...]=(), src_idxs:tuple[UOp|int,...]=(), axis:int=0):
|
||||
dst, src = cast(UOp, dst), cast(UOp, src)
|
||||
assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
|
||||
dst_dtype, src_dtype = dst.dtype, src.dtype
|
||||
if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.LOCAL:
|
||||
laneid = self.ker.laneid
|
||||
st, rt = cast(ST, dst), cast(RT, src)
|
||||
elements_per_thread = rt.base_shape.elements_per_thread
|
||||
|
||||
for height in self.ker.range(src.shape[-3], track=False):
|
||||
for width in self.ker.range(src.shape[-2], track=False):
|
||||
for inner in self.ker.range(elements_per_thread, track=False):
|
||||
if rt.layout != st.layout:
|
||||
row = rt.base_shape.stride * (laneid // rt.base_shape.cols) + inner
|
||||
col = laneid % rt.base_shape.cols
|
||||
else:
|
||||
row = laneid % rt.base_shape.rows
|
||||
col = rt.base_shape.stride * (laneid // rt.base_shape.rows) + inner
|
||||
|
||||
srow, scol = cast(ST, dst).swizzle(row, col)
|
||||
|
||||
src_load = src[*src_idxs, height, width, inner]
|
||||
if src.dtype.base != dst.dtype.base:
|
||||
src_load = src_load.cast(dst.dtype.base)
|
||||
dst_store = dst[*idxs[:-2], height, width, srow, scol].store(src_load)
|
||||
dst_store = dst_store.end(height, width, inner)
|
||||
elif src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL and isinstance(src, RT):
|
||||
dstf = dst.flatten()
|
||||
row_stride = prod(dst.shape[axis+1:])
|
||||
|
||||
laneid = self.ker.laneid
|
||||
rt = cast(RT, src)
|
||||
elements_per_thread = rt.base_shape.elements_per_thread
|
||||
|
||||
idxs = tuple(idx * src.shape[-3] * rt.base_shape.rows if i == axis else idx for i, idx in enumerate(idxs))
|
||||
idxs = tuple(idx * src.shape[-2] * rt.base_shape.cols if i == 3 else idx for i, idx in enumerate(idxs))
|
||||
dst_i = ((idxs[0] * dst.shape[-3] + idxs[1]) * dst.shape[-2] + idxs[2]) * dst.shape[-1] + idxs[3]
|
||||
|
||||
for height in self.ker.range(src.shape[-3], track=False):
|
||||
for width in self.ker.range(src.shape[-2], track=False):
|
||||
for inner in self.ker.range(elements_per_thread, track=False):
|
||||
base_row = height * rt.base_shape.rows
|
||||
base_col = width * rt.base_shape.cols
|
||||
|
||||
if rt.layout == TileLayout.COL:
|
||||
row = rt.base_shape.stride * (laneid // rt.base_shape.cols) + inner
|
||||
col = laneid % rt.base_shape.cols
|
||||
else:
|
||||
row = laneid % rt.base_shape.rows
|
||||
col = rt.base_shape.stride * (laneid // rt.base_shape.rows) + inner
|
||||
|
||||
srow, scol = base_row + row, base_col + col
|
||||
|
||||
dst_i += srow * row_stride + scol
|
||||
|
||||
src_load = src[*src_idxs, height, width, inner]
|
||||
if src.dtype.base != dst.dtype.base:
|
||||
src_load = src_load.cast(dst.dtype.base)
|
||||
dst_store = dstf[dst_i].store(src_load).end(height, width, inner)
|
||||
elif src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL and isinstance(src, RV):
|
||||
dstf = dst.flatten()
|
||||
row_stride = prod(dst.shape[axis+1:])
|
||||
|
||||
laneid = self.ker.laneid
|
||||
rv = cast(RV, src)
|
||||
reductions = rv.base_shape.rows
|
||||
|
||||
assert rv.layout == VecLayout.ORTHO, "only ortho layout supported"
|
||||
|
||||
idxs = tuple(idx * rv.length if i == 3 else idx for i, idx in enumerate(idxs))
|
||||
dst_i = ((idxs[0] * dst.shape[-3] + idxs[1]) * dst.shape[-2] + idxs[2]) * dst.shape[-1] + idxs[3]
|
||||
|
||||
for outer in self.ker.range(src.shape[-2], track=False):
|
||||
dst_i += outer * reductions + (laneid % reductions)
|
||||
|
||||
src_load = src[outer, 0]
|
||||
if src.dtype.base != dst.dtype.base:
|
||||
src_load = src_load.cast(dst.dtype.base)
|
||||
dst_store = dstf[dst_i].store(src_load).end(outer)
|
||||
else:
|
||||
raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented for {type(src)=}")
|
||||
|
||||
self.ker.push_store(dst_store, dst)
|
||||
return dst.after(dst_store).reshape(dst.shape)
|
||||
113
tinygrad_repo/extra/thunder/tiny/tk/kernel.py
Normal file
113
tinygrad_repo/extra/thunder/tiny/tk/kernel.py
Normal file
@@ -0,0 +1,113 @@
|
||||
from contextlib import AbstractContextManager
|
||||
from tinygrad.uop.ops import UOp, KernelInfo, AxisType, AddrSpace
|
||||
from extra.thunder.tiny.tk import WARP_THREADS
|
||||
from extra.thunder.tiny.tk.group import Group
|
||||
from extra.thunder.tiny.tk.tiles import GL, ST_16X16, ST, RT_16X16, RT, RV, TileLayout, VecLayout
|
||||
|
||||
class _tk_range:
|
||||
def __init__(self, start:int, end:int, step:int, axis_type:AxisType, rid:int):
|
||||
self.start, self.end, self.step = start, end, step
|
||||
self.axis_type, self.rid, self.done = axis_type, rid, False
|
||||
def __iter__(self): return self
|
||||
def __next__(self):
|
||||
if not self.done:
|
||||
self.done = True
|
||||
self._rng = UOp.range((self.end - self.start) // self.step, self.rid, axis_type=self.axis_type) * self.step + self.start
|
||||
return self._rng
|
||||
raise StopIteration
|
||||
|
||||
class Kernel(AbstractContextManager):
|
||||
def __init__(self, name:str, grid_size:tuple[int, int, int], block_size:int):
|
||||
self.name = name
|
||||
|
||||
self.blockIdx_x = UOp.special(grid_size[0], "gidx0")
|
||||
self.blockIdx_y = UOp.special(grid_size[1], "gidx1")
|
||||
self.blockIdx_z = UOp.special(grid_size[2], "gidx2")
|
||||
self.threadIdx_x = UOp.special(block_size, "lidx0")
|
||||
|
||||
self.range_stack: list[_tk_range] = []
|
||||
self.store_stack: list[tuple[UOp, UOp]] = []
|
||||
|
||||
self.global_slot = 0
|
||||
self.shared_slot = 0
|
||||
self.register_slot = 0
|
||||
self.range_id = 0
|
||||
self.allocs: dict[tuple[str, tuple], UOp] = {}
|
||||
|
||||
@property
|
||||
def warpid(self): return self.threadIdx_x // WARP_THREADS
|
||||
@property
|
||||
def laneid(self): return self.threadIdx_x % WARP_THREADS
|
||||
|
||||
def __enter__(self): return self
|
||||
def __exit__(self, exc_type, exc_value, traceback): pass
|
||||
|
||||
def group(self, size:int): return Group(size, self)
|
||||
@property
|
||||
def warp(self): return self.group(1)
|
||||
@property
|
||||
def warpgroup(self): return self.group(4)
|
||||
|
||||
def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.LOOP, track:bool=True):
|
||||
if end == 0: start, end = 0, start
|
||||
rng = _tk_range(start, end, step, axis_type, self.range_id)
|
||||
self.range_id += 1
|
||||
if track: self.range_stack.append(rng)
|
||||
return rng
|
||||
|
||||
def raw_range(self, end:int=0, axis_type:AxisType=AxisType.LOOP):
|
||||
rng = UOp.range(end, self.range_id, axis_type=axis_type)
|
||||
self.range_id += 1
|
||||
return rng
|
||||
|
||||
def alloc(self, shape, dtype, addrspace:AddrSpace, name:str|None=None):
|
||||
match addrspace:
|
||||
case AddrSpace.GLOBAL:
|
||||
slot = self.global_slot
|
||||
self.global_slot += 1
|
||||
case AddrSpace.LOCAL:
|
||||
slot = self.shared_slot
|
||||
self.shared_slot += 1
|
||||
case AddrSpace.REG:
|
||||
slot = self.register_slot
|
||||
self.register_slot += 1
|
||||
|
||||
uop = UOp.placeholder(shape, dtype, slot=slot, addrspace=addrspace)
|
||||
|
||||
if name:
|
||||
if (name, shape) in self.allocs: return self.allocs[(name, shape)]
|
||||
self.allocs[(name, shape)] = uop
|
||||
|
||||
return uop
|
||||
|
||||
def gl(self, shape, dtype): return GL.create(shape, dtype, self)
|
||||
def st(self, shape, dtype, layout=TileLayout.ROW, base_shape=ST_16X16): return ST.create(shape, dtype, layout, base_shape, self)
|
||||
def rt(self, shape, dtype, layout=TileLayout.ROW, base_shape=RT_16X16): return RT.create(shape, dtype, layout, base_shape, self)
|
||||
def rv(self, length, dtype, layout=VecLayout.ORTHO, rt_base_shape=RT_16X16): return RV.create(length, dtype, layout, rt_base_shape, self)
|
||||
|
||||
def push_store(self, store:UOp, uop:UOp): self.store_stack.append((store, uop))
|
||||
|
||||
def finish(self, stores:int=1):
|
||||
# end all ranges
|
||||
rngs = []
|
||||
while self.range_stack: rngs.append(self.range_stack.pop(0)._rng)
|
||||
|
||||
# end stores stores
|
||||
store_uops = []
|
||||
for _ in range(stores):
|
||||
store = self.store_stack.pop()[0]
|
||||
if hasattr(store, '_uop'): store_uops.append(store._uop)
|
||||
else: store_uops.append(store)
|
||||
uop = UOp.group(*store_uops)
|
||||
|
||||
return uop.end(*rngs).sink(arg=KernelInfo(name=self.name, opts_to_apply=())).simplify()
|
||||
|
||||
def endrange(self, ranges:int=1):
|
||||
last_store = self.store_stack.pop()
|
||||
|
||||
rngs = []
|
||||
for _ in range(ranges):
|
||||
last_range = self.range_stack.pop()
|
||||
rngs.append(last_range._rng)
|
||||
|
||||
return last_store[1].after(last_store[0].end(*rngs)).reshape(last_store[1].shape)
|
||||
276
tinygrad_repo/extra/thunder/tiny/tk/tiles.py
Normal file
276
tinygrad_repo/extra/thunder/tiny/tk/tiles.py
Normal file
@@ -0,0 +1,276 @@
|
||||
from enum import Enum, auto
|
||||
import functools
|
||||
from typing import Callable
|
||||
from dataclasses import dataclass
|
||||
from tinygrad.dtype import AddrSpace, DType
|
||||
from tinygrad.mixin import ElementwiseMixin
|
||||
from tinygrad.uop.ops import UOp, Ops
|
||||
|
||||
from extra.thunder.tiny.tk import WARP_THREADS
|
||||
|
||||
def unwrap(x):
|
||||
if hasattr(x, "_uop"): return x._uop
|
||||
if isinstance(x, (list, tuple)): return type(x)(unwrap(y) for y in x)
|
||||
if isinstance(x, dict): return {k: unwrap(v) for k,v in x.items()}
|
||||
return x
|
||||
|
||||
def wrap(x, s):
|
||||
if isinstance(x, UOp): return s.ruop(x)
|
||||
if isinstance(x, (list, tuple)): return type(x)(wrap(y, s) for y in x)
|
||||
return x
|
||||
|
||||
def autowrap(source_cls, blacklist=None):
|
||||
if blacklist is None:
|
||||
blacklist = {
|
||||
"__init__", "__new__", "__str__", "__del__", "__repr__", "__dict__", "__getattribute__",
|
||||
"__setattr__", "__delattr__", "__weakref__", "__slots__", "__class__",
|
||||
"__reduce__", "__reduce_ex__", "__getstate__", "__setstate__", "__hash__"
|
||||
}
|
||||
|
||||
def decorator(cls):
|
||||
def __getattr__(self, name):
|
||||
uop = object.__getattribute__(self, "_uop")
|
||||
val = getattr(uop, name)
|
||||
if callable(val):
|
||||
@functools.wraps(val)
|
||||
def proxy(*args, **kwargs):
|
||||
return wrap(val(*unwrap(args), **unwrap(kwargs)), self)
|
||||
return proxy
|
||||
if name in UOp.__slots__: return val # type: ignore
|
||||
return wrap(val, self)
|
||||
cls.__getattr__ = __getattr__
|
||||
|
||||
for name in dir(source_cls):
|
||||
if name in blacklist or not name.startswith("__"): continue
|
||||
|
||||
for base in cls.mro():
|
||||
if base is source_cls: break
|
||||
if name in base.__dict__: break
|
||||
else:
|
||||
original = getattr(source_cls, name)
|
||||
if callable(original):
|
||||
def make_proxy(_, func):
|
||||
def proxy(self, *args, **kwargs):
|
||||
return wrap(func(self._uop, *unwrap(args), **unwrap(kwargs)), self)
|
||||
return proxy
|
||||
setattr(cls, name, make_proxy(name, original))
|
||||
|
||||
return cls
|
||||
return decorator
|
||||
|
||||
class TileMathMixin(ElementwiseMixin):
|
||||
def alu(self, op, *src, inner_op=lambda x:x):
|
||||
assert isinstance(self, (RT, RV))
|
||||
if len(src) == 0:
|
||||
if self._uop._shape is None: uop = UOp.alu(self._uop, op)
|
||||
else: uop = self.ker.warp.map(self._uop, lambda x: UOp.alu(x, op))
|
||||
elif len(src) == 1:
|
||||
if self._uop._shape is None: uop = UOp.alu(self._uop, op, inner_op(self._uop.ufix(src[0])))
|
||||
elif isinstance(src[0], (int,float,bool)): uop = self.ker.warp.map(self._uop, lambda x: UOp.alu(x, op, inner_op(x.ufix(src[0]))))
|
||||
elif src[0]._shape is None: uop = UOp.alu(self._uop, op, inner_op(self._uop.ufix(src[0])))
|
||||
else:
|
||||
if isinstance(self, RT) and isinstance(src[0], RV):
|
||||
match self.layout:
|
||||
case TileLayout.ROW: uop = self.ker.warp.map(self._uop, lambda x, idx: UOp.alu(x, op, inner_op(src[0]._uop[idx[0], 0])))
|
||||
case TileLayout.COL: uop = self.ker.warp.map(self._uop, lambda x, idx: UOp.alu(x, op, inner_op(src[0]._uop[idx[1], 0])))
|
||||
else: uop = self.ker.warp.map(self._uop, lambda x, idx: UOp.alu(x, op, inner_op(src[0]._uop[*idx])))
|
||||
else: raise NotImplementedError
|
||||
return self.ruop(uop)
|
||||
def const_like(self, b): return b
|
||||
|
||||
@property
|
||||
def dtype(self): return self._uop.dtype
|
||||
def cast(self, dtype:DType): return self.ruop(self._uop.cast(dtype))
|
||||
|
||||
# override ops that do compute on the src uop
|
||||
def sub(self, x, reverse=False):
|
||||
return self.ufix(x).alu(Ops.ADD, self, inner_op=lambda y: -y) if reverse else self.alu(Ops.ADD, self.ufix(x), inner_op=lambda y: -y)
|
||||
def div(self, x, reverse=False):
|
||||
return self.ufix(x).alu(Ops.MUL, self, inner_op=lambda y: 1/y) if reverse else self.alu(Ops.MUL, self.ufix(x), inner_op=lambda y: 1/y)
|
||||
|
||||
@autowrap(UOp)
|
||||
class GL:
|
||||
def __init__(self, uop:UOp, ker):
|
||||
self._uop, self.ker = uop, ker
|
||||
|
||||
def ruop(self, uop:UOp):
|
||||
return GL(uop, self.ker)
|
||||
|
||||
@classmethod
|
||||
def create(cls, shape, dtype:DType, ker):
|
||||
uop = ker.alloc(shape, dtype, AddrSpace.GLOBAL)
|
||||
return cls(uop, ker)
|
||||
|
||||
class TileLayout(Enum):
|
||||
ROW = auto()
|
||||
COL = auto()
|
||||
|
||||
class VecLayout(Enum):
|
||||
ORTHO = auto()
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BaseShape:
|
||||
rows: int
|
||||
cols: int
|
||||
|
||||
@property
|
||||
def num_elements(self): return self.rows * self.cols
|
||||
@property
|
||||
def elements_per_thread(self): return self.num_elements // WARP_THREADS
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class STBaseShape(BaseShape):
|
||||
_swizzle: Callable[[UOp, DType], UOp]
|
||||
bytes_per_thread: Callable[[DType], int]
|
||||
|
||||
def swizzle(self, row, col, dtype:DType):
|
||||
offset = row * self.cols + col
|
||||
offset *= dtype.itemsize
|
||||
offset = self._swizzle(offset, dtype)
|
||||
offset //= dtype.itemsize
|
||||
return offset
|
||||
|
||||
def st_16x16_swizzle(offset:UOp, _): return offset
|
||||
def st_16x16_bpt(dtype:DType):
|
||||
if dtype.itemsize == 2 or dtype.itemsize == 4: return 16
|
||||
else: raise NotImplementedError
|
||||
ST_16X16 = STBaseShape(16, 16, st_16x16_swizzle, st_16x16_bpt)
|
||||
|
||||
def st_16x16_swizzled_swizzle(offset:UOp, dtype:DType):
|
||||
if dtype.itemsize == 2:
|
||||
swizzle = ((offset % 512) >> 7) << 3
|
||||
return offset ^ swizzle
|
||||
elif dtype.itemsize == 4:
|
||||
return offset
|
||||
else: raise NotImplementedError
|
||||
def st_16x16_swizzled_bpt(dtype:DType):
|
||||
if dtype.itemsize == 2: return 4
|
||||
elif dtype.itemsize == 4: return 16
|
||||
else: raise NotImplementedError
|
||||
ST_16X16_SWIZZLED = STBaseShape(16, 16, st_16x16_swizzled_swizzle, st_16x16_swizzled_bpt)
|
||||
|
||||
def st_32x32_swizzle(offset:UOp, dtype:DType):
|
||||
if dtype.itemsize == 2:
|
||||
first_swizzle = ((offset % 1024) >> 9) << 5
|
||||
second_swizzle = ((offset % 2048) >> 10) << 4
|
||||
return offset ^ first_swizzle ^ second_swizzle
|
||||
elif dtype.itemsize == 4:
|
||||
return offset
|
||||
else: raise NotImplementedError
|
||||
def st_32x32_bpt(dtype:DType):
|
||||
if dtype.itemsize == 2 or dtype.itemsize == 4: return 16
|
||||
else: raise NotImplementedError
|
||||
ST_32X32 = STBaseShape(32, 32, st_32x32_swizzle, st_32x32_bpt)
|
||||
|
||||
def st_16x32_swizzle(offset:UOp, dtype:DType):
|
||||
if dtype.itemsize == 2:
|
||||
swizzle = ((offset % 1024) >> 9) << 5
|
||||
return offset ^ swizzle
|
||||
elif dtype.itemsize == 4:
|
||||
return offset
|
||||
else: raise NotImplementedError
|
||||
def st_16x32_bpt(dtype:DType):
|
||||
if dtype.itemsize == 2 or dtype.itemsize == 4: return 16
|
||||
else: raise NotImplementedError
|
||||
ST_16X32 = STBaseShape(16, 32, st_16x32_swizzle, st_16x32_bpt)
|
||||
|
||||
def st_32x16_swizzle(offset:UOp, dtype:DType):
|
||||
if dtype.itemsize == 2:
|
||||
swizzle = ((offset % 1024) >> 9) << 4
|
||||
return offset ^ swizzle
|
||||
elif dtype.itemsize == 4:
|
||||
return offset
|
||||
else: raise NotImplementedError
|
||||
def st_32x16_bpt(dtype:DType):
|
||||
if dtype.itemsize == 2 or dtype.itemsize == 4: return 16
|
||||
else: raise NotImplementedError
|
||||
ST_32X16 = STBaseShape(32, 16, st_32x16_swizzle, st_32x16_bpt)
|
||||
|
||||
@autowrap(UOp)
|
||||
class ST:
|
||||
def __init__(self, uop:UOp, rows:int, cols:int, layout:TileLayout, base_shape:STBaseShape, ker):
|
||||
self._uop, self.rows, self.cols, self.layout, self.base_shape, self.ker = uop, rows, cols, layout, base_shape, ker
|
||||
|
||||
def ruop(self, uop:UOp):
|
||||
return ST(uop, self.rows, self.cols, self.layout, self.base_shape, self.ker)
|
||||
|
||||
@classmethod
|
||||
def create(cls, shape, dtype:DType, layout:TileLayout, base_shape:STBaseShape, ker):
|
||||
rows = shape[-2]
|
||||
cols = shape[-1]
|
||||
assert rows % base_shape.rows == 0
|
||||
assert cols % base_shape.cols == 0
|
||||
assert cols % base_shape.elements_per_thread == 0
|
||||
|
||||
height = rows // base_shape.rows
|
||||
width = cols // base_shape.cols
|
||||
|
||||
uop = ker.alloc(shape[:-2] + (height, width, base_shape.rows, base_shape.cols), dtype, AddrSpace.LOCAL)
|
||||
return cls(uop, rows, cols, layout, base_shape, ker)
|
||||
|
||||
def swizzle(self, row, col):
|
||||
swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype.base.scalar())
|
||||
|
||||
row = swizzled_offset // self.base_shape.cols
|
||||
col = swizzled_offset % self.base_shape.cols
|
||||
|
||||
return row, col
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RTBaseShape(BaseShape):
|
||||
stride: int
|
||||
|
||||
@property
|
||||
def num_strides(self):
|
||||
return self.elements_per_thread // self.stride
|
||||
|
||||
RT_16X16 = RTBaseShape(rows=16, cols=16, stride=4)
|
||||
RT_32X32 = RTBaseShape(rows=32, cols=32, stride=4)
|
||||
RT_32X32_8 = RTBaseShape(rows=32, cols=32, stride=8)
|
||||
RT_16X32 = RTBaseShape(rows=16, cols=32, stride=8)
|
||||
RT_32X16 = RTBaseShape(rows=32, cols=16, stride=8)
|
||||
RT_32X16_4 = RTBaseShape(rows=32, cols=16, stride=4)
|
||||
RT_16X32_4 = RTBaseShape(rows=16, cols=32, stride=4)
|
||||
|
||||
@autowrap(UOp)
|
||||
class RT(TileMathMixin):
|
||||
def __init__(self, uop:UOp, layout:TileLayout, base_shape:RTBaseShape, ker):
|
||||
self._uop, self.layout, self.base_shape, self.ker = uop, layout, base_shape, ker
|
||||
|
||||
def ruop(self, uop:UOp):
|
||||
return RT(uop, self.layout, self.base_shape, self.ker)
|
||||
|
||||
@classmethod
|
||||
def create(cls, shape, dtype:DType, layout:TileLayout, base_shape:RTBaseShape, ker):
|
||||
assert len(shape) == 2
|
||||
assert shape[0] % base_shape.rows == 0
|
||||
assert shape[1] % base_shape.cols == 0
|
||||
|
||||
height = shape[0] // base_shape.rows
|
||||
width = shape[1] // base_shape.cols
|
||||
|
||||
uop = ker.alloc((height, width, base_shape.elements_per_thread), dtype, AddrSpace.REG)
|
||||
return cls(uop, layout, base_shape, ker)
|
||||
|
||||
@autowrap(UOp)
|
||||
class RV(TileMathMixin):
|
||||
def __init__(self, uop:UOp, length:int, layout:VecLayout, base_shape:RTBaseShape, ker):
|
||||
self._uop, self.ker = uop, ker
|
||||
self.length, self.layout, self.base_shape = length, layout, base_shape
|
||||
|
||||
def ruop(self, uop:UOp):
|
||||
return RV(uop, self.length, self.layout, self.base_shape, self.ker)
|
||||
|
||||
@classmethod
|
||||
def create(cls, length, dtype:DType, layout:VecLayout, base_shape:RTBaseShape, ker):
|
||||
tiles = length // base_shape.rows
|
||||
|
||||
match layout:
|
||||
case VecLayout.ORTHO:
|
||||
inner_dim = 1
|
||||
outer_dim = tiles
|
||||
|
||||
uop = ker.alloc((outer_dim, inner_dim), dtype, AddrSpace.REG)
|
||||
return RV(uop, length, layout, base_shape, ker)
|
||||
|
||||
ALL_TILES = UOp | GL | ST | RT | RV
|
||||
Reference in New Issue
Block a user