J3ff/IQ.Pilot
1
0
Fork 0
forked from IQ.Lvbs/IQ.Pilot
Code Pull Requests Activity

IQ.Pilot Release Commit @ 5bc9cd3

Browse Source
This commit is contained in:
IQ.Lvbs history cleanup
2026-08-22 23:42:42 -05:00
commit d037016281
4504 changed files with 1129827 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

416
tinygrad_repo/extra/thunder/tiny/fa.py Normal file
View File

@@ -0,0 +1,416 @@
import math
from tinygrad import Tensor, dtypes
from tinygrad.helpers import DEBUG
from tinygrad.uop.ops import UOp, Ops
from extra.thunder.tiny.tk import WARP_THREADS
from extra.thunder.tiny.tk.kernel import Kernel
from extra.thunder.tiny.tk.tiles import GL, TileLayout
NUM_WORKERS = 1
Q_BLOCK_SIZE = 32
KV_BLOCK_SIZE = 32
def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None) -> Tensor:
if not isinstance(ref.device, tuple): return Tensor.empty(*shape, dtype=ref.dtype, device=ref.device)
shape = tuple(s // len(ref.device) if i == ref.uop.axis else s for i, s in enumerate(shape))
axis = ref.uop.axis if axis is None else axis
return Tensor(Tensor.empty(*shape, dtype=ref.dtype, device=ref.device).uop.multi(axis), dtype=ref.dtype, device=ref.device)
def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
return _sharded_empty(ref.shape, ref, axis)
def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False):
if len(xq.shape) == 3: xq, xk, xv = xq.unsqueeze(0), xk.unsqueeze(0), xv.unsqueeze(0)
odtype = xq.dtype
xq, xk, xv = xq.transpose(1, 2).cast(dtypes.bfloat16), xk.transpose(1, 2).cast(dtypes.bfloat16), xv.transpose(1, 2).cast(dtypes.bfloat16)
_, N_, _, D_ = xq.shape
block_size = max(Q_BLOCK_SIZE, KV_BLOCK_SIZE)
assert D_ % block_size == 0, f"embedding dimension must be multiple of block size, got {D_=} {block_size=}"
# pad to multiple of block size
xq = xq.pad(((0, 0), (0, (block_size - (xq.shape[1] % block_size)) % block_size), (0, 0), (0, 0)))
xk = xk.pad(((0, 0), (0, (block_size - (xk.shape[1] % block_size)) % block_size), (0, 0), (0, 0)))
xv = xv.pad(((0, 0), (0, (block_size - (xv.shape[1] % block_size)) % block_size), (0, 0), (0, 0)))
B, N, H, D = xq.shape
H_KV = xk.shape[2]
GROUP_SIZE = H // H_KV
num_devices = len(xq.device) if isinstance(xq.device, tuple) else 1
B_local = B // num_devices
if DEBUG >= 2: print(f"Flash Attention {B=} {B_local=} {N=} {H=} {D=} {H_KV=} {GROUP_SIZE=}")
def _custom_forward_impl(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp|None) -> UOp:
with Kernel("fa_custom_forward", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
warp = ker.warp
o, q, k, v, l_vec = GL(ou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(l_vecu, ker)
mask = GL(masku, ker) if masku is not None else None
head = ker.blockIdx_x
head_kv = head // GROUP_SIZE
batch = ker.blockIdx_z
q_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
q_reg_fl = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
q_reg_transposed = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
k_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
k_reg_transposed = ker.rt((D, KV_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
v_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
o_reg = ker.rt((D, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
o_reg_transposed = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
att_block = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
att_block_mma = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
mask_reg = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.float32)
mask_reg_transposed = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
max_vec_last = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
max_vec = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
norm_vec = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
scale_vec = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
max_vec = warp.neg_inf(max_vec)
norm_vec = warp.zero(norm_vec)
o_reg = warp.zero(o_reg)
scale_vec = warp.ones(scale_vec)
# load q tile
q_reg_fl = warp.load(q_reg_fl, q, (), (batch, q_seq, head, 0), axis=1)
q_reg_fl *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
q_reg = warp.copy(q_reg, q_reg_fl)
q_reg_transposed = warp.transpose(q_reg_transposed, q_reg)
num_kv_blocks = (q_seq + 1) if is_causal else (N // KV_BLOCK_SIZE)
for kv_idx in ker.range(num_kv_blocks):
k_reg = warp.load(k_reg, k, (), (batch, kv_idx, head_kv, 0), axis=1)
v_reg = warp.load(v_reg, v, (), (batch, kv_idx, head_kv, 0), axis=1)
# mma qk^t
att_block = warp.zero(att_block.after(kv_idx))
k_reg_transposed = warp.transpose(k_reg_transposed, k_reg)
att_block = warp.mma_AtB(att_block, k_reg_transposed, q_reg_transposed)
# apply attention mask
if is_causal:
bs_rows, bs_cols, bs_stride = att_block.base_shape.rows, att_block.base_shape.cols, att_block.base_shape.stride
q_base = q_seq * Q_BLOCK_SIZE + (warp.laneid % bs_cols)
kv_base = kv_idx * KV_BLOCK_SIZE + (warp.laneid // bs_cols) * bs_stride
att_block = warp.map(att_block,
lambda x, idx: ((kv_base + idx[0]*bs_rows + idx[2]) > (q_base + idx[1]*bs_cols)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
elif mask is not None:
mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
att_block += mask_reg_transposed
# softmax
max_vec_last = warp.copy(max_vec_last.after(kv_idx), max_vec)
max_vec = warp.col_reduce(max_vec.after(max_vec_last), att_block, lambda a, b: a.maximum(b), init_value=-math.inf)
scale_vec = warp.map(scale_vec.after(max_vec_last, max_vec), lambda _, idx: max_vec_last[*idx] - max_vec[*idx])
scale_vec = scale_vec.exp2()
o_reg *= scale_vec
norm_vec *= scale_vec
att_block -= max_vec
att_block = att_block.exp2()
norm_vec = warp.col_reduce(norm_vec.after(scale_vec), att_block, lambda a, b: a + b)
# mma av
att_block_mma = warp.copy(att_block_mma.after(kv_idx, norm_vec), att_block)
o_reg = warp.mma_AtB(o_reg, v_reg, att_block_mma)
o_reg = ker.endrange()
norm_vec = norm_vec.after(o_reg)
max_vec = max_vec.after(o_reg)
o_reg /= norm_vec
o_reg_transposed = warp.transpose(o_reg_transposed, o_reg)
o = warp.store(o, o_reg_transposed, (batch, q_seq, head, 0), (), axis=1)
norm_vec = norm_vec.after(o)
max_vec = max_vec.after(o)
max_vec *= math.log(2)
norm_vec = norm_vec.log2() * math.log(2)
norm_vec += max_vec
l_vec = warp.store(l_vec, norm_vec, (batch, head, 0, q_seq), (), axis=2)
o = o.after(l_vec)
return ker.finish()
def custom_forward_causal(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp) -> UOp:
return _custom_forward_impl(ou, l_vecu, qu, ku, vu, None)
def custom_forward_masked(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp) -> UOp:
return _custom_forward_impl(ou, l_vecu, qu, ku, vu, masku)
def _custom_backward_q_impl(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp|None, l_vecu:UOp, delta_vecu:UOp) -> UOp:
with Kernel("fa_custom_backward_q", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
warp = ker.warp
dq, do, q, k, v = GL(dqu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker)
mask = GL(masku, ker) if masku is not None else None
l_vec, delta_vec = GL(l_vecu, ker), GL(delta_vecu, ker)
head = ker.blockIdx_x
head_kv = head // GROUP_SIZE
batch = ker.blockIdx_z
q_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
q_reg_fl = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
q_reg_t = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
k_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
k_reg_t = ker.rt((D, KV_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
k_reg_col = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
k_reg_col_t = ker.rt((D, KV_BLOCK_SIZE), dtypes.bfloat16)
v_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
mask_reg = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.float32)
mask_reg_transposed = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
dq_reg = ker.rt((D, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
dq_reg_transposed = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
do_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
dp_block = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
att_block = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
att_block_mma = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
l_vec_reg = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
delta_vec_reg = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
dq_reg = warp.zero(dq_reg)
# load q tile
q_reg_fl = warp.load(q_reg_fl, q, (), (batch, q_seq, head, 0), axis=1)
q_reg_fl *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
q_reg = warp.copy(q_reg, q_reg_fl)
q_reg_t = warp.transpose(q_reg_t, q_reg)
# load do tile
do_reg = warp.load(do_reg, do, (), (batch, q_seq, head, 0), axis=1)
# load l_vec
l_vec_reg = warp.load(l_vec_reg, l_vec, (), (batch, head, 0, q_seq), axis=2)
l_vec_reg *= 1.0 / math.log(2)
delta_vec_reg = warp.load(delta_vec_reg, delta_vec, (), (batch, head, 0, q_seq), axis=2)
num_kv_blocks = (q_seq + 1) if is_causal else (N // KV_BLOCK_SIZE)
for kv_idx in ker.range(num_kv_blocks):
k_reg = warp.load(k_reg, k, (), (batch, kv_idx, head_kv, 0), axis=1)
k_reg_col = warp.load(k_reg_col, k, (), (batch, kv_idx, head_kv, 0), axis=1)
v_reg = warp.load(v_reg, v, (), (batch, kv_idx, head_kv, 0), axis=1)
k_reg_t = warp.transpose(k_reg_t, k_reg)
k_reg_col_t = warp.transpose(k_reg_col_t, k_reg_col)
# mma qk^t
att_block = warp.zero(att_block.after(kv_idx))
att_block = warp.mma_AtB(att_block, k_reg_t, q_reg_t)
# apply attention mask
if is_causal:
bs_rows, bs_cols, bs_stride = att_block.base_shape.rows, att_block.base_shape.cols, att_block.base_shape.stride
q_base = q_seq * Q_BLOCK_SIZE + (warp.laneid % bs_cols)
kv_base = kv_idx * KV_BLOCK_SIZE + (warp.laneid // bs_cols) * bs_stride
att_block = warp.map(att_block,
lambda x, idx: ((kv_base + idx[0]*bs_rows + idx[2]) > (q_base + idx[1]*bs_cols)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
elif mask is not None:
mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
att_block += mask_reg_transposed
att_block -= l_vec_reg
att_block = att_block.exp2()
dp_block = warp.zero(dp_block.after(kv_idx, att_block))
dp_block = warp.mma_ABt(dp_block, v_reg, do_reg)
dp_block -= delta_vec_reg
att_block *= dp_block
att_block *= 1.0 / math.sqrt(D)
att_block_mma = warp.copy(att_block_mma, att_block)
dq_reg = warp.mma_AB(dq_reg, k_reg_col_t, att_block_mma)
dq_reg = ker.endrange()
dq_reg_transposed = warp.transpose(dq_reg_transposed, dq_reg)
dq = warp.store(dq, dq_reg_transposed, (batch, q_seq, head, 0), axis=1)
return ker.finish()
def custom_backward_q_causal(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
return _custom_backward_q_impl(dqu, dou, qu, ku, vu, None, l_vecu, delta_vecu)
def custom_backward_q_masked(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
return _custom_backward_q_impl(dqu, dou, qu, ku, vu, masku, l_vecu, delta_vecu)
def _custom_backward_kv_impl(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp|None, l_vecu:UOp, delta_vecu:UOp):
with Kernel("fa_custom_backward_kv", (H_KV, N // (KV_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
warp = ker.warp
dk, dv, do, q, k, v = GL(dku, ker), GL(dvu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker)
mask = GL(masku, ker) if masku is not None else None
l_vec, delta_vec = GL(l_vecu, ker), GL(delta_vecu, ker)
head_kv = ker.blockIdx_x
batch = ker.blockIdx_z
kv_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
att_smem = ker.st((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16)
q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
q_reg_t = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
q_reg_col = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
k_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
k_reg_t = ker.rt((D, KV_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
v_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
mask_reg = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.float32)
mask_reg_transposed = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
dk_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.float32, TileLayout.COL)
dv_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.float32, TileLayout.COL)
do_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
do_reg_col = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
dp_block = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
att_block = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
att_block_mma = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
att_block_transposed = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
att_block_row = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16)
l_vec_reg = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
delta_vec_reg = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
dk_reg = warp.zero(dk_reg)
dv_reg = warp.zero(dv_reg)
# load kv tile
k_reg = warp.load(k_reg, k, (), (batch, kv_seq, head_kv, 0), axis=1)
k_reg_t = warp.transpose(k_reg_t, k_reg)
v_reg = warp.load(v_reg, v, (), (batch, kv_seq, head_kv, 0), axis=1)
q_start = kv_seq if is_causal else 0
for q_idx in ker.range(q_start, N // Q_BLOCK_SIZE):
for g in ker.range(GROUP_SIZE):
head_q = head_kv * GROUP_SIZE + g
q_reg = warp.load(q_reg, q, (), (batch, q_idx, head_q, 0), axis=1)
q_reg_col = warp.load(q_reg_col, q, (), (batch, q_idx, head_q, 0), axis=1)
do_reg = warp.load(do_reg, do, (), (batch, q_idx, head_q, 0), axis=1)
do_reg_col = warp.load(do_reg_col, do, (), (batch, q_idx, head_q, 0), axis=1)
q_reg_t = warp.transpose(q_reg_t, q_reg)
# load l_vec and delta_vec
l_vec_reg = warp.load(l_vec_reg, l_vec, (), (batch, head_q, 0, q_idx), axis=2)
l_vec_reg *= 1.0 / math.log(2)
delta_vec_reg = warp.load(delta_vec_reg, delta_vec, (), (batch, head_q, 0, q_idx), axis=2)
# mma qk^t
att_block = warp.zero(att_block.after(g))
att_block = warp.mma_AtB(att_block, k_reg_t, q_reg_t)
att_block *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
# apply attention mask
if is_causal:
bs_rows, bs_cols, bs_stride = att_block.base_shape.rows, att_block.base_shape.cols, att_block.base_shape.stride
q_base = q_idx * Q_BLOCK_SIZE + (warp.laneid % bs_cols)
kv_base = kv_seq * KV_BLOCK_SIZE + (warp.laneid // bs_cols) * bs_stride
att_block = warp.map(att_block,
lambda x, idx: ((kv_base + idx[0]*bs_rows + idx[2]) > (q_base + idx[1]*bs_cols)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
elif mask is not None:
mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_idx, kv_seq), axis=2)
mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
att_block += mask_reg_transposed
att_block -= l_vec_reg
att_block = att_block.exp2()
att_block_mma = warp.copy(att_block_mma, att_block)
att_block_transposed = warp.transpose(att_block_transposed, att_block_mma)
att_smem = warp.store(att_smem, att_block_transposed)
att_block_row = warp.load(att_block_row, att_smem)
dv_reg_ = warp.mma_AtB(dv_reg, att_block_row, do_reg_col)
dp_block = warp.zero(dp_block.after(g, q_idx, dv_reg_))
dp_block = warp.mma_ABt(dp_block, v_reg, do_reg)
dp_block -= delta_vec_reg
att_block *= dp_block
att_block *= 1.0 / math.sqrt(D)
att_block_mma = warp.copy(att_block_mma, att_block)
att_block_transposed = warp.transpose(att_block_transposed, att_block_mma)
att_smem = warp.store(att_smem, att_block_transposed)
att_block_row = warp.load(att_block_row, att_smem)
dk_reg = warp.mma_AtB(dk_reg, att_block_row, q_reg_col)
dk_reg = ker.endrange(2)
dv_reg = dv_reg.after(dk_reg)
dv_reg = warp.map(dv_reg, lambda x, idx: x + v_reg[*idx].cast(dtypes.float32) * 1e-30)
dk = warp.store(dk, dk_reg, (batch, kv_seq, head_kv, 0), axis=1)
dv = warp.store(dv, dv_reg, (batch, kv_seq, head_kv, 0), axis=1)
return ker.finish(2)
def custom_backward_kv_causal(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, l_vecu:UOp, delta_vecu:UOp):
return _custom_backward_kv_impl(dku, dvu, dou, qu, ku, vu, None, l_vecu, delta_vecu)
def custom_backward_kv_masked(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp):
return _custom_backward_kv_impl(dku, dvu, dou, qu, ku, vu, masku, l_vecu, delta_vecu)
single_device = xq.device[0] if isinstance(xq.device, tuple) else xq.device
if is_causal:
if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True")
elif attn_mask is not None:
if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf"))
if attn_mask.shape != (B, 1, N, N):
attn_mask = attn_mask.expand(B, 1, N, N)
if isinstance(xq.device, tuple) and not isinstance(attn_mask.device, tuple):
attn_mask = attn_mask.shard(xq.device, axis=0)
else:
attn_mask = Tensor.zeros((B, 1, N, N), device=single_device, dtype=dtypes.float32)
if isinstance(xq.device, tuple):
attn_mask = attn_mask.shard(xq.device, axis=0)
attn = _sharded_empty_like(xq, axis=0)
l_vec = _sharded_empty((B, H, 1, N), xq, axis=0)
def grad_causal(gradu:UOp, _) -> tuple[None, None, UOp, UOp, UOp]:
grad = Tensor(gradu, device=gradu.device)
grad_q = _sharded_empty_like(xq, axis=0)
grad_k = _sharded_empty_like(xk, axis=0)
grad_v = _sharded_empty_like(xv, axis=0)
delta_vec = (grad * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
grad_q = Tensor.custom_kernel(grad_q, grad, xq, xk, xv, l_vec, delta_vec, fxn=custom_backward_q_causal)[0]
grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, xq, xk, xv, l_vec, delta_vec, fxn=custom_backward_kv_causal)[:2]
return (None, None, grad_q.uop, grad_k.uop, grad_v.uop)
def grad_masked(gradu:UOp, _) -> tuple[None, None, UOp, UOp, UOp, None]:
grad = Tensor(gradu, device=gradu.device)
grad_q = _sharded_empty_like(xq, axis=0)
grad_k = _sharded_empty_like(xk, axis=0)
grad_v = _sharded_empty_like(xv, axis=0)
delta_vec = (grad * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
grad_q = Tensor.custom_kernel(grad_q, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_q_masked)[0]
grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_kv_masked)[:2]
return (None, None, grad_q.uop, grad_k.uop, grad_v.uop, None)
if is_causal:
attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, fxn=custom_forward_causal, grad_fxn=grad_causal)[:2]
else:
attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, attn_mask, fxn=custom_forward_masked, grad_fxn=grad_masked)[:2]
attn_ = attn[:, :N_, :, :D_]
return attn_.transpose(1, 2).cast(odtype)

6
tinygrad_repo/extra/thunder/tiny/tk/__init__.py Normal file
View 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
View 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
View 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
View 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

156
tinygrad_repo/extra/thunder/tiny/visualize_tile.py Normal file
View File

@@ -0,0 +1,156 @@
from tinygrad.helpers import colored
WARP_THREADS = 64
BASE_TILE_ROWS = 16
BASE_TILE_COLS = 16
BASE_TILE_NEPT = (BASE_TILE_ROWS * BASE_TILE_COLS) // WARP_THREADS
DTYPE_SIZE = 2
INST = "ds_read_b64"
def row_col(threadIdx_x):
local_warpid = threadIdx_x // WARP_THREADS
warp_laneid = threadIdx_x % WARP_THREADS
ret = []
for inner in range(BASE_TILE_NEPT):
if BASE_TILE_ROWS == 16 and BASE_TILE_COLS == 16:
row = warp_laneid % 16
col = 4 * (warp_laneid // 16)
elif BASE_TILE_ROWS == 16 and BASE_TILE_COLS == 32:
row = warp_laneid % 16
col = 8 * (warp_laneid // 16)
row_offset = 0
col_offset = inner
# swizzle then find row and col
offset = (row + row_offset) * BASE_TILE_COLS + (col + col_offset)
offset *= DTYPE_SIZE
if BASE_TILE_ROWS == 16 and BASE_TILE_COLS == 16:
swizzle = ((offset % 512) >> 7) << 3
offset = offset ^ swizzle
elif BASE_TILE_ROWS == 16 and BASE_TILE_COLS == 32:
swizzle = ((offset % 1024) >> 9) << 5
offset = offset ^ swizzle
offset //= DTYPE_SIZE
row = offset // BASE_TILE_COLS
col = offset % BASE_TILE_COLS
ret.append((row, col))
return ret
# ===
def shm_phase(inst, threadIdx_x):
match inst:
case "ds_read_b128":
match threadIdx_x:
case 0 | 1 | 2 | 3 | 12 | 13 | 14 | 15 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27: return 0
case 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 16 | 17 | 18 | 19 | 28 | 29 | 30 | 31: return 1
case 32 | 33 | 34 | 35 | 44 | 45 | 46 | 47 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59: return 2
case 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 48 | 49 | 50 | 51 | 60 | 61 | 62 | 63: return 3
case "ds_read_b64":
if threadIdx_x < 32: return 0
else: return 1
case "ds_write_b64":
if threadIdx_x < 16: return 0
elif threadIdx_x < 32: return 1
elif threadIdx_x < 48: return 2
else: return 3
def shm_bank(inst, row, col):
bank = row * (BASE_TILE_COLS // 2) + (col // 2)
match inst:
case "ds_read_b128": bank = bank % 64
case "ds_read_b64": bank = bank % 64
case "ds_write_b64": bank = bank % 32
return bank
def map_range(value, from_min, from_max, to_min, to_max):
ratio = (value - from_min) / (from_max - from_min)
return to_min + ratio * (to_max - to_min)
def shm_bank_gradient(inst, bank):
# rgb color for each bank
# for 16 bit elements, two elements per bank row wise
# gradient from blue to red
amount = map_range(bank, 0, (64 if inst != "ds_write_b64" else 32) - 1, 0, 120)
amount = int(amount)
return (amount, amount // 2, 120 - amount)
def color_code(phase):
match phase:
case 0: return "red"
case 1: return "green"
case 2: return "blue"
case 3: return "yellow"
def rgb_bg(text, color):
return f"\033[48;2;{color[0]};{color[1]};{color[2]}m{text}\033[0m"
def visualize_threads(inst=INST):
for threadIdx_x in range(WARP_THREADS):
row, col = zip(*row_col(threadIdx_x))
print(f"Thread {threadIdx_x:2}: ", end="")
for r, c in zip(row, col):
phase = shm_phase(inst, threadIdx_x)
color = color_code(phase)
print(f"{color}({r:3},{c:3})\033[0m ", end="")
print()
unique_pairs = set()
for threadIdx_x in range(WARP_THREADS):
rc_list = row_col(threadIdx_x)
for rc in rc_list:
unique_pairs.add(rc)
assert len(unique_pairs) == 64 * BASE_TILE_NEPT, f"Expected {64 * BASE_TILE_NEPT} unique pairs, got {len(unique_pairs)}"
def visualize_tile(inst=INST):
tile = [[-1 for _ in range(BASE_TILE_COLS)] for _ in range(BASE_TILE_ROWS)]
for threadIdx_x in range(WARP_THREADS):
rc_list = row_col(threadIdx_x)
for r, c in rc_list:
try:
tile[r][c] = threadIdx_x
except:
pass
bank_conflicts = {}
print("\nTile layout (each number indicates the thread holding that position):")
for r in range(BASE_TILE_ROWS):
for c in range(BASE_TILE_COLS):
phase = shm_phase(inst, tile[r][c])
bank = shm_bank(inst, r, c)
color = color_code(phase)
bank_color = shm_bank_gradient(inst, bank)
if (bank, phase) not in bank_conflicts:
bank_conflicts[(bank, phase)] = []
bank_conflicts[(bank, phase)].append((r, c, tile[r][c]))
if phase == -1:
bank_color = (0, 0, 0)
text = colored(f"{tile[r][c]:2}", color)
text = rgb_bg(text, bank_color)
print(f"{text:2}", end=" ")
print()
for (bank, phase), positions in bank_conflicts.items():
if len(positions) > 1:
unique_threads = set(pos[2] for pos in positions)
if len(unique_threads) > 1:
print(f"{len(unique_threads)} way bank conflict: bank {bank}")
if __name__ == "__main__":
visualize_tile()
# visualize_threads()