sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

rename LOOP -> WEAK and STRONGLOOP -> LOOP (#17283)

Browse Source
This commit is contained in:
George Hotz
2026-07-29 10:38:36 -07:00
committed by GitHub
parent 6c2b9fac08
commit 52c9e5a99e
20 changed files with 66 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/fp8/fp8_linear.py
+3 -3
View File
@@ -18,9 +18,9 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
SEQ = inp.shape[1]
OUT = weight.shape[0]
IN = weight.shape[-1]
seq_idx = UOp.range(SEQ, 2, AxisType.LOOP)
out_idx = UOp.range(OUT, 3, AxisType.LOOP)
batch_idx = UOp.range(output.size//SEQ//OUT, 1, AxisType.LOOP)
seq_idx = UOp.range(SEQ, 2)
out_idx = UOp.range(OUT, 3)
batch_idx = UOp.range(output.size//SEQ//OUT, 1)
reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
reduced = product.reduce(reduce_idx, arg=Ops.ADD)
extra/gemm/amd_copy_matmul.py
+2 -2
View File
@@ -70,8 +70,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
if use_wmma:
k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
tile_m = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
tile_n = UOp.range(TN, 201, AxisType.LOOP)
tile_m = UOp.range(TM // WMMA_ACC, 200)
tile_n = UOp.range(TN, 201)
acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0,2,1)[tile_m, tile_n]
a_frag = A_local.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_K // WMMA_K, WMMA_K)[wave_m, tile_m, lane_n, k]
extra/gemm/amd_flash_attention.py
+8 -8
View File
@@ -92,8 +92,8 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
S_reg = UOp.placeholder((TM, TN), dtypes.float, slot=6, addrspace=AddrSpace.REG)
S_reg = S_reg.after(S_reg.after(n_tile).store(S_reg.const_like(0)))
k_qk = UOp.range(D // WMMA_K, 101, AxisType.REDUCE)
tm1 = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
tn1 = UOp.range(TN, 201, AxisType.LOOP)
tm1 = UOp.range(TM // WMMA_ACC, 200)
tn1 = UOp.range(TN, 201)
S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
k_frag = KV_lds_k.reshape(WAVES_N, TN, WMMA_N, D // WMMA_K, WMMA_K)[wave_n, tn1, lane_n, k_qk]
@@ -110,7 +110,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
rm2 = UOp.range(TN, 261, AxisType.REDUCE)
m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
# warp reduce max (in-place)
ri_w = UOp.range(TM, 270, AxisType.LOOP)
ri_w = UOp.range(TM, 270)
m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce_max(m_ij[ri_w], lane)).end(ri_w))
# compute P = exp(S - m_ij) in S_reg
@@ -120,7 +120,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
p_local = p_local.after(p_local.after(n_tile).store(p_local.const_like(0)))
rp2 = UOp.range(TN, 291, AxisType.REDUCE)
p_local = p_local.after(p_local.store(p_local.after(rp2) + S_reg[:, rp2]).end(rp2))
ri_ws = UOp.range(TM, 295, AxisType.LOOP)
ri_ws = UOp.range(TM, 295)
p_sum = p_local.after(p_local[ri_ws].store(warp_reduce_sum(p_local[ri_ws], lane)).end(ri_ws))
# write P = exp(S - m_ij) to P_lds (reuses slot 0, Q no longer needed)
@@ -130,11 +130,11 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
P_store = P_write[tid].store(S_reg.cast(dtypes.half))
# -- online softmax correction --
ri4 = UOp.range(TM, 330, AxisType.LOOP)
ri4 = UOp.range(TM, 330)
m_new_val = m_i[ri4].maximum(m_ij[ri4])
alpha_val = ((m_i[ri4] - m_new_val) * LOG2E).exp2()
beta_val = ((m_ij[ri4] - m_new_val) * LOG2E).exp2()
rj4 = UOp.range(TD, 331, AxisType.LOOP)
rj4 = UOp.range(TD, 331)
correction = UOp.group(
acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4),
l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]),
@@ -153,8 +153,8 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
# -- acc += P @ V via WMMA --
k_pv = UOp.range(BLOCK_N // WMMA_K, 400, AxisType.REDUCE)
tm2 = UOp.range(TM // WMMA_ACC, 401, AxisType.LOOP)
tn2 = UOp.range(TD, 402, AxisType.LOOP)
tm2 = UOp.range(TM // WMMA_ACC, 401)
tn2 = UOp.range(TD, 402)
acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
p_frag = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
v_frag = KV_lds_v.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
extra/gemm/amd_uop_matmul.py
+2 -2
View File
@@ -28,10 +28,10 @@ REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.WEAK): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
def copy(dest:UOp, src:UOp, rng:int, upcast=False):
assert dest.shape == src.shape
rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.WEAK)
return dest[*rngs].store(src[*rngs]).end(*rngs)
def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
extra/gemm/cdna_asm_gemm.py
+2 -2
View File
@@ -171,8 +171,8 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
M, K = A.shape[0]*A.shape[1], A.shape[2]
K2, N = B.shape[(1 if B.ndim == 3 else 0):]
assert K == K2
m = UOp.range(M, 1, AxisType.LOOP)
n = UOp.range(N, 2, AxisType.LOOP)
m = UOp.range(M, 1)
n = UOp.range(N, 2)
k = UOp.range(K, 0, AxisType.REDUCE)
mul = (A.flatten().index((m*UOp.const(dtypes.weakint, K)+k))*
B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)
extra/gemm/mi350x_uop_matmul_2.py
+1 -1
View File
@@ -29,7 +29,7 @@ TID_SIZE = WARPGROUP_SIZE*WARP_SIZE
def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=()):
assert dest.shape == src.shape
rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.LOOP) for i,s in enumerate(src.shape)]
rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.WEAK) for i,s in enumerate(src.shape)]
copy = dest[*rngs].store(src[*rngs]).end(*rngs)
return dest.after(copy) if set else copy
extra/gemm_fragment.py
+9 -9
View File
@@ -42,7 +42,7 @@ API mapping (tilelang -> tinygrad UOps, idioms from test/backend/test_custom_ker
smem tile AFTER the copy; the implicit-barrier pass turns the store->load
dependency of the loop that consumes it into a workgroup barrier
T.gemm (no WMMA) -> C_local[..].set(C_local.after(k)[..] + a_shared[..] * b_shared[..], end=k)
with k a loop-carried STRONGLOOP range (codegen builds the register accumulator
with k a loop-carried LOOP range (codegen builds the register accumulator
from this self-referential store automatically)
T.copy(fragment, gmem) -> gmem.index(gidx).store(C_local[..]).end(all_ranges)
UNSHARD lowering -> multi_pm in codegen (full_rewrite_to_sink): INDEX/AFTER/STORE ops on the
@@ -110,20 +110,20 @@ def matmul_relu_kernel(c:UOp, a:UOp, b:UOp) -> UOp:
C_local = alloc_fragment((BLOCK_M, BLOCK_N), dtypes.float32, axis=0, tnum=tid)
# T.clear(C_local) -- each thread zeroes its own fragment rows
ir0, j0 = UOp.range(ROWS, 4, AxisType.STRONGLOOP), UOp.range(BLOCK_N, 5, AxisType.STRONGLOOP)
ir0, j0 = UOp.range(ROWS, 4, AxisType.LOOP), UOp.range(BLOCK_N, 5, AxisType.LOOP)
C_loc = C_local[tid*ROWS + ir0, j0].set(0.0, end=(ir0, j0))
# for ko in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3):
# (num_stages pipelining is async copy + multi-buffering; this is the synchronous single-buffer version)
ko = UOp.range(cdiv(K, BLOCK_K), 3, AxisType.STRONGLOOP)
ko = UOp.range(cdiv(K, BLOCK_K), 3, AxisType.LOOP)
# T.copy(A[by * BLOCK_M, ko * BLOCK_K], A_shared) -- each thread copies ROWS row(s)
# set returns the tile AFTER the copy; codegen turns that store->load dependency into a workgroup barrier
iar, ka = UOp.range(ROWS, 6, AxisType.STRONGLOOP), UOp.range(BLOCK_K, 7, AxisType.STRONGLOOP)
iar, ka = UOp.range(ROWS, 6, AxisType.LOOP), UOp.range(BLOCK_K, 7, AxisType.LOOP)
A_shared = A_shared[tid*ROWS + iar, ka].set(a[by*BLOCK_M + tid*ROWS + iar, ko*BLOCK_K + ka], end=(iar, ka))
# T.copy(B[ko * BLOCK_K, bx * BLOCK_N], B_shared) -- each thread copies ROWS_B column(s)
kb, ibr = UOp.range(BLOCK_K, 8, AxisType.STRONGLOOP), UOp.range(ROWS_B, 9, AxisType.STRONGLOOP)
kb, ibr = UOp.range(BLOCK_K, 8, AxisType.LOOP), UOp.range(ROWS_B, 9, AxisType.LOOP)
B_shared = B_shared[kb, tid*ROWS_B + ibr].set(b[ko*BLOCK_K + kb, bx*BLOCK_N + tid*ROWS_B + ibr], end=(kb, ibr))
# T.gemm(A_shared, B_shared, C_local), no WMMA -- per-thread accumulate over its fragment rows.
@@ -131,16 +131,16 @@ def matmul_relu_kernel(c:UOp, a:UOp, b:UOp) -> UOp:
# which codegen turns into a register accumulator
# kk nests outside the fragment row/col loops (lower ids nest outer), so B loads are contiguous and
# the A row value is loaded once per kk
ir, kk = UOp.range(ROWS, 10, AxisType.STRONGLOOP), UOp.range(BLOCK_K, 11, AxisType.STRONGLOOP)
jj = UOp.range(BLOCK_N, 12, AxisType.STRONGLOOP)
ir, kk = UOp.range(ROWS, 10, AxisType.LOOP), UOp.range(BLOCK_K, 11, AxisType.LOOP)
jj = UOp.range(BLOCK_N, 12, AxisType.LOOP)
acc = C_loc.after(kk)[tid*ROWS + ir, jj] + A_shared[tid*ROWS + ir, kk].cast(dtypes.float32) * B_shared[kk, jj].cast(dtypes.float32)
# closing the ko loop here too; codegen adds the barrier so no thread overwrites the tiles while others still read them
C_loc = C_loc[tid*ROWS + ir, jj].set(acc, end=(kk, ir, jj, ko))
# for i, j in T.Parallel(BLOCK_M, BLOCK_N): C_local[i, j] = T.max(C_local[i, j], 0)
# T.copy(C_local, C[by * BLOCK_M, bx * BLOCK_N]) -- per-thread store of the fragment shard (relu fused into it)
# STRONGLOOP: these loops are the per-thread output layout; convert_loop_to_global must not globalize them
ie, je = UOp.range(ROWS, 13, AxisType.STRONGLOOP), UOp.range(BLOCK_N, 14, AxisType.STRONGLOOP)
# LOOP: these loops are the per-thread output layout; convert_loop_to_global must not globalize them
ie, je = UOp.range(ROWS, 13, AxisType.LOOP), UOp.range(BLOCK_N, 14, AxisType.LOOP)
c_st = c[by*BLOCK_M + tid*ROWS + ie, bx*BLOCK_N + je].store(C_loc[tid*ROWS + ie, je].relu().cast(c.dtype))
# all open ranges are closed at the final store (ko was closed above).
extra/llama_kernels/quantize_fp8_delayed/__init__.py
+1 -1
View File
@@ -16,7 +16,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
wg = UOp.range(NUM_WG, 0, AxisType.GLOBAL)
tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.LOOP)
it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.WEAK)
lane = UOp.range(VEC, 3, AxisType.UNROLL)
idx = (((it * NUM_WG + wg) * THREADS_PER_WG + tid) * VEC) + lane
extra/thunder/tiny/tk/kernel.py
+2 -2
View File
@@ -48,14 +48,14 @@ class Kernel(AbstractContextManager):
@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):
def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.WEAK, 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):
def raw_range(self, end:int=0, axis_type:AxisType=AxisType.WEAK):
rng = UOp.range(end, self.range_id, axis_type=axis_type)
self.range_id += 1
return rng
test/external/external_benchmark_op_conv.py
Vendored
+7 -7
View File
@@ -13,9 +13,9 @@ from tinygrad.dtype import Invalid
def vision_conv_143():
c0 = UOp.param(0, dtypes.half, shape=(16, 1024, 4))
c2 = UOp.range(32, 3, AxisType.LOOP)
c5 = UOp.range(128, 4, AxisType.LOOP)
c8 = UOp.range(16, 2, AxisType.LOOP)
c2 = UOp.range(32, 3)
c5 = UOp.range(128, 4)
c8 = UOp.range(16, 2)
c16 = UOp.range(7, 0, AxisType.REDUCE)
c17 = c8*2+c16
c24 = ((c17<3)!=True)&(c17<35)
@@ -39,9 +39,9 @@ def vision_conv_143():
def vision_conv_153():
c0 = UOp.param(0, dtypes.half, shape=(8, 1024, 4))
c2 = UOp.range(16, 3, AxisType.LOOP)
c5 = UOp.range(256, 4, AxisType.LOOP)
c8 = UOp.range(8, 2, AxisType.LOOP)
c2 = UOp.range(16, 3)
c5 = UOp.range(256, 4)
c8 = UOp.range(8, 2)
c16 = UOp.range(7, 0, AxisType.REDUCE)
c17 = c8*2+c16
c24 = ((c17<3)!=True)&(c17<19)
@@ -65,7 +65,7 @@ def vision_conv_153():
def dm_conv_172():
c0 = UOp.param(0, dtypes.half, shape=(1, 240, 4))
c2 = UOp.range(960, 4, AxisType.LOOP)
c2 = UOp.range(960, 4)
c5 = UOp.param(1, dtypes.half, shape=(8, 384, 4))
c7 = UOp.range(32, 0, AxisType.REDUCE)
c10 = UOp.range(4, 1, AxisType.REDUCE)
test/mockgpu/amd/emu.py
+2 -2
View File
@@ -51,7 +51,7 @@ class _MXCSRContext:
if lib is None or not hasattr(self, '_saved'): return
lib.set_fpcr(self._saved)
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.dtype import dtypes, AddrSpace
from tinygrad.device import Buffer, BufferSpec, Device
from tinygrad.runtime.autogen import hsa
@@ -446,7 +446,7 @@ class _Ctx:
"""Create a lane range UOp with unique axis ID."""
if n is None: n = self.wave_size
self._axis_id += 1
return UOp.range(n, self._axis_id, AxisType.LOOP, dtype=dtypes.int)
return UOp.range(n, self._axis_id, dtype=dtypes.int)
def unroll_lanes(self, get_lane_bit, exec_mask: UOp, apply_exec: bool = True) -> UOp:
"""Combine lane bits into a mask using RANGE+REDUCE (32-bit for RDNA, 64-bit for CDNA)."""
test/null/test_linearizer_failures.py
+2 -2
View File
@@ -8,8 +8,8 @@ from tinygrad.codegen import to_program
class TestLinearizerFailures(unittest.TestCase):
def test_fail_1(self):
c0 = UOp.param(0, dtypes.float, (64,))
c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.LOOP)
c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.LOOP)
c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.WEAK)
c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.WEAK)
c3 = ((c1*UOp.const(dtypes.weakint, 32))+c2)
c4 = UOp.param(1, dtypes.float, (163840,))
c5 = UOp.range(UOp.const(dtypes.weakint, 2560), 0, AxisType.REDUCE)
test/null/test_uop_graph.py
+4 -4
View File
@@ -424,8 +424,8 @@ class TestUOpGraph(unittest.TestCase):
# mnist indexing with split reduceop
# Make sure we are not doign math on the loaded index, which would promote it to long
c0 = UOp.param(0, dtypes.uchar, (128000,))
c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.WEAK)
c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
c3 = UOp.param(1, dtypes.int, (512,))
c4 = c3.index(c1)
c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
@@ -441,8 +441,8 @@ class TestUOpGraph(unittest.TestCase):
def test_load_idx_no_math_on_loaded(self):
# test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
c0 = UOp.param(0, dtypes.uchar, (128000,))
c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.WEAK)
c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
c3 = UOp.param(1, dtypes.int, (512,))
c4 = c3.index(c1) # c4 is a load
c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
tinygrad/codegen/__init__.py
+1 -1
View File
@@ -268,7 +268,7 @@ def add_raw_barrier(after:UOp):
def add_war_barrier(end:UOp):
# a LOCAL buffer stored and loaded in the same loop needs a barrier at the end of the loop body
rngs = [r for r in end.src[1:] if r.op is Ops.RANGE and r.arg[1] in (AxisType.REDUCE, AxisType.LOOP, AxisType.STRONGLOOP) and r.vmax > 0]
rngs = [r for r in end.src[1:] if r.op is Ops.RANGE and r.arg[1] in (AxisType.REDUCE, AxisType.WEAK, AxisType.LOOP) and r.vmax > 0]
if not rngs or end.src[0].op is Ops.BARRIER: return None
sl = end.src[0].backward_slice_with_self
# only stores that are inside this loop body (not in the backward slice through AFTER chains from other loops)
tinygrad/codegen/opt/heuristic.py
+2 -2
View File
@@ -169,7 +169,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
else:
# prioritize making expand axes local
local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].get_idx().backward_slice for b in k.bufs), axis) \
for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
for axis in k.axes_of(AxisType.GLOBAL, AxisType.WEAK) if k.rngs[axis].src[0].op is Ops.CONST]
to_local: list[tuple[int, int]] = []
for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])):
local_size = prod(sz for _, sz in to_local)
@@ -188,7 +188,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
for threads in [32,16,12,8,6,5,4,3,2]:
# Skip if too many threads. Heuristic: use about 128K ops per thread
if threads > k.ren.global_max[0] or resolve(prod(k.full_shape) // (128 << 10) < threads): continue
for axis in k.axes_of(AxisType.LOOP):
for axis in k.axes_of(AxisType.WEAK):
if k.full_shape[axis] % threads == 0:
try: k.apply_opt(Opt(OptOps.THREAD, axis, threads))
except KernelOptError: pass
tinygrad/codegen/opt/postrange.py
+6 -6
View File
@@ -65,7 +65,7 @@ class Scheduler:
def _output_rngs(self) -> list[UOp]:
return flatten([[r for r in UOp.sink(*s.src[1:]).ranges if r.arg[-1] != AxisType.REDUCE] for s in self.ast.src if s.op is Ops.END])
def _globalizable_rngs(self) -> list[UOp]:
ret = [r for r in self._output_rngs() if r.arg[-1] == AxisType.LOOP]
ret = [r for r in self._output_rngs() if r.arg[-1] == AxisType.WEAK]
# exclude any output ranges from global that don't appear in all BUFFERIZE
for x in self.ast.toposort():
if x.op is Ops.STAGE:
@@ -86,8 +86,8 @@ class Scheduler:
ret = []
for x,r in zip(self.axis_types, self.rngs):
if self.dont_use_locals and x == AxisType.GLOBAL: ret.append("BLUE")
elif r not in output_rngs and x == AxisType.LOOP: ret.append("BLACK")
elif r not in globalizible_rngs and x == AxisType.LOOP: ret.append("white")
elif r not in output_rngs and x == AxisType.WEAK: ret.append("BLACK")
elif r not in globalizible_rngs and x == AxisType.WEAK: ret.append("white")
else: ret.append(axis_colors[x])
return ret
def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():>4s}', color) for x,color in zip(self.rngs, self.colors())])
@@ -108,7 +108,7 @@ class Scheduler:
# copied from kernel.py
@property
def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP) \
def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.WEAK) \
if isinstance(s:=self.full_shape[i], int) and s > 1]
@property
def unrollable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE) \
@@ -161,10 +161,10 @@ class Scheduler:
check(rng.arg[-1] in {AxisType.GROUP_REDUCE, AxisType.REDUCE}, "unroll is for GROUP_REDUCE/REDUCE")
if opt.op is OptOps.UPCAST:
check((self.ren is not None and self.ren.target.device == "DSP") or amt <= 16, "don't upcast more than 16")
check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP}, f"upcast is for GLOBAL/LOCAL/LOOP, not {rng.arg[-1]}")
check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOCAL, AxisType.WEAK}, f"upcast is for GLOBAL/LOCAL/LOOP, not {rng.arg[-1]}")
if opt.op is OptOps.LOCAL:
check(not self.dont_use_locals, "can't use locals")
check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOOP}, "local is for globals")
check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.WEAK}, "local is for globals")
if opt.op is OptOps.THREAD:
check(self.ren is not None and self.ren.has_threads, "target does not support threads")
check(self.ren is not None and self.ren.global_max is not None and amt <= self.ren.global_max[0], "too many threads")
tinygrad/nn/__init__.py
+1 -1
View File
@@ -337,7 +337,7 @@ def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple:
BLOCK_J = min(256, embed_size)
n_j_blocks = (embed_size + BLOCK_J - 1) // BLOCK_J
i = UOp.range(grad_emb_flat.shape[0], 0) # batch_size * sequence_length -> GLOBAL
j_inner = UOp.range(BLOCK_J, 2, AxisType.LOOP if device in ("CPU", "NULL") else AxisType.LOCAL) # BLOCK_J threads per workgroup
j_inner = UOp.range(BLOCK_J, 2, AxisType.WEAK if device in ("CPU", "NULL") else AxisType.LOCAL) # BLOCK_J threads per workgroup
j_outer = UOp.range(n_j_blocks, 1)
j = j_outer * BLOCK_J + j_inner
# mask padded embed
tinygrad/schedule/indexing.py
+1 -1
View File
@@ -50,7 +50,7 @@ class IndexingContext:
# create ranges
range_idx: Iterator[int] = field(default_factory=itertools.count)
def new_range(self, s:sint, axistype:AxisType=AxisType.LOOP) -> UOp:
def new_range(self, s:sint, axistype:AxisType=AxisType.WEAK) -> UOp:
if isinstance(s, UOp) and s.op is Ops.RANGE: return s
# if a range has a 1 src, it's the same as UOp.const(dtypes.weakint, 0)
return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(None, 0)
tinygrad/schedule/rangeify.py
+2 -2
View File
@@ -342,9 +342,9 @@ def limit_bufs(ctx:IndexingContext, root:UOp):
srcs = []
for s in root.src:
if s.op in GroupOp.Elementwise and s.device is not None:
# Insert bufferize: all AxisType.REDUCE before bufferize are AxisType.LOOP, the DEVICE range stays a launched axis
# Insert bufferize: all AxisType.REDUCE before bufferize are AxisType.WEAK, the DEVICE range stays a launched axis
orig_ranges = s.ranges
end_ranges = [x.replace(arg=(next(ctx.range_idx), AxisType.LOOP)) if x.op is Ops.RANGE and x.arg[-1] is not AxisType.DEVICE else x
end_ranges = [x.replace(arg=(next(ctx.range_idx), AxisType.WEAK)) if x.op is Ops.RANGE and x.arg[-1] is not AxisType.DEVICE else x
for x in s.ranges]
s = s.substitute(dict(zip(orig_ranges, end_ranges))).bufferize(*end_ranges, arg=BufferizeOpts(device=s.device)).index(*orig_ranges)
srcs.append(s)
tinygrad/uop/ops.py
+8 -8
View File
@@ -16,8 +16,8 @@ if TYPE_CHECKING:
class AxisType(Enum):
def __repr__(self): return str(self)
DEVICE = auto(); GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto() # noqa: E702
UNROLL = auto(); THREAD = auto(); PLACEHOLDER = auto(); STRONGLOOP = auto() # noqa: E702
DEVICE = auto(); GLOBAL = auto(); WARP = auto(); LOCAL = auto(); WEAK = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto() # noqa: E702
UNROLL = auto(); THREAD = auto(); PLACEHOLDER = auto(); LOOP = auto() # noqa: E702
@dataclass(frozen=True, order=True)
class ParamArg:
@@ -35,14 +35,14 @@ class ParamArg:
("volatile", False))
args = [repr(self.slot), repr(self.dtype)] + [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default]
return f"ParamArg({', '.join(args)})"
axis_letters = {AxisType.DEVICE: "d", AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L",
AxisType.STRONGLOOP: "S", AxisType.UPCAST: "u", AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
axis_letters = {AxisType.DEVICE: "d", AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.WEAK: "L",
AxisType.LOOP: "L", AxisType.UPCAST: "u", AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
axis_colors = {AxisType.DEVICE: "green", AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN",
AxisType.LOOP: "WHITE", AxisType.STRONGLOOP: "WHITE", AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red",
AxisType.WEAK: "WHITE", AxisType.LOOP: "WHITE", AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red",
AxisType.UNROLL: "magenta"}
# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
axis_to_pos = {AxisType.DEVICE: -2, AxisType.LOOP: -1, AxisType.STRONGLOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1,
axis_to_pos = {AxisType.DEVICE: -2, AxisType.WEAK: -1, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1,
AxisType.LOCAL: 2, AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
range_start = {Ops.STAGE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1,
@@ -617,10 +617,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
ret = UOp(Ops.CONST, dtype, arg=dtype.const(b), src=())
return ret._mop(Ops.EXPAND, arg=shape) if shape is not None and shape != () and ret.shape != shape else ret
@staticmethod
def range(end:sint, axis_id, axis_type=AxisType.LOOP, *arg, dtype=dtypes.weakint, src=(), **kwargs):
def range(end:sint, axis_id, axis_type=AxisType.WEAK, *arg, dtype=dtypes.weakint, src=(), **kwargs):
return UOp(Ops.RANGE, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs)
@staticmethod
def loop(axis_id:int, *arg): return UOp(Ops.RANGE, src=(UOp(Ops.NOOP),), arg=(axis_id, AxisType.LOOP)+arg)
def loop(axis_id:int, *arg): return UOp(Ops.RANGE, src=(UOp(Ops.NOOP),), arg=(axis_id, AxisType.WEAK)+arg)
@staticmethod
def special(end:sint, name:str, dtype=dtypes.weakint): return UOp(Ops.SPECIAL, src=(sint_to_uop(end, dtype),), arg=name)
@staticmethod