diff --git a/extra/fp8/fp8_linear.py b/extra/fp8/fp8_linear.py index 48e7097543..757a88e78d 100644 --- a/extra/fp8/fp8_linear.py +++ b/extra/fp8/fp8_linear.py @@ -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) diff --git a/extra/gemm/amd_copy_matmul.py b/extra/gemm/amd_copy_matmul.py index cd86c7cea0..5dbea993d4 100644 --- a/extra/gemm/amd_copy_matmul.py +++ b/extra/gemm/amd_copy_matmul.py @@ -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] diff --git a/extra/gemm/amd_flash_attention.py b/extra/gemm/amd_flash_attention.py index 7e134208d8..d0896d5d78 100644 --- a/extra/gemm/amd_flash_attention.py +++ b/extra/gemm/amd_flash_attention.py @@ -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] diff --git a/extra/gemm/amd_uop_matmul.py b/extra/gemm/amd_uop_matmul.py index 8db8a608ee..8091ba79b7 100644 --- a/extra/gemm/amd_uop_matmul.py +++ b/extra/gemm/amd_uop_matmul.py @@ -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: diff --git a/extra/gemm/cdna_asm_gemm.py b/extra/gemm/cdna_asm_gemm.py index 6823b1fc02..cc7add0c52 100644 --- a/extra/gemm/cdna_asm_gemm.py +++ b/extra/gemm/cdna_asm_gemm.py @@ -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) diff --git a/extra/gemm/mi350x_uop_matmul_2.py b/extra/gemm/mi350x_uop_matmul_2.py index 26bb43ff4d..f3efd581a6 100644 --- a/extra/gemm/mi350x_uop_matmul_2.py +++ b/extra/gemm/mi350x_uop_matmul_2.py @@ -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 diff --git a/extra/gemm_fragment.py b/extra/gemm_fragment.py index 6f1e2f5c8f..c1d5900692 100644 --- a/extra/gemm_fragment.py +++ b/extra/gemm_fragment.py @@ -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). diff --git a/extra/llama_kernels/quantize_fp8_delayed/__init__.py b/extra/llama_kernels/quantize_fp8_delayed/__init__.py index e7b00c8b61..7f2f189c79 100644 --- a/extra/llama_kernels/quantize_fp8_delayed/__init__.py +++ b/extra/llama_kernels/quantize_fp8_delayed/__init__.py @@ -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 diff --git a/extra/thunder/tiny/tk/kernel.py b/extra/thunder/tiny/tk/kernel.py index 1c3d9af9ab..a6ad42272c 100644 --- a/extra/thunder/tiny/tk/kernel.py +++ b/extra/thunder/tiny/tk/kernel.py @@ -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 diff --git a/test/external/external_benchmark_op_conv.py b/test/external/external_benchmark_op_conv.py index 2ab421cb81..bc7e774862 100644 --- a/test/external/external_benchmark_op_conv.py +++ b/test/external/external_benchmark_op_conv.py @@ -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) diff --git a/test/mockgpu/amd/emu.py b/test/mockgpu/amd/emu.py index 0dea6e175a..61dda4f20e 100644 --- a/test/mockgpu/amd/emu.py +++ b/test/mockgpu/amd/emu.py @@ -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).""" diff --git a/test/null/test_linearizer_failures.py b/test/null/test_linearizer_failures.py index 4a51c173ed..fe23315741 100644 --- a/test/null/test_linearizer_failures.py +++ b/test/null/test_linearizer_failures.py @@ -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) diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index e1608d16e8..4f60284230 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -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) 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) diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py index ece76a78ed..c9b133fa5b 100644 --- a/tinygrad/codegen/opt/heuristic.py +++ b/tinygrad/codegen/opt/heuristic.py @@ -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 diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index e0805e0e0e..dd4e5f03e7 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -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") diff --git a/tinygrad/nn/__init__.py b/tinygrad/nn/__init__.py index 27ef6877b3..e98186111e 100644 --- a/tinygrad/nn/__init__.py +++ b/tinygrad/nn/__init__.py @@ -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 diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 90d24b6552..b9ab1ce1f9 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -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) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 80833ba322..68a359e23a 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -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) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index a19a9eee7e..10fc425cf9 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -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