diff --git a/tinygrad_repo/.github/actions/setup-tinygrad/action.yml b/tinygrad_repo/.github/actions/setup-tinygrad/action.yml index 372f227e3..4fbfb7488 100644 --- a/tinygrad_repo/.github/actions/setup-tinygrad/action.yml +++ b/tinygrad_repo/.github/actions/setup-tinygrad/action.yml @@ -258,10 +258,11 @@ runs: if: inputs.cuda == 'true' shell: bash run: | - sudo mkdir -p /usr/local/cuda/targets/x86_64-linux + sudo mkdir -p /usr/local/cuda-11.5/targets/x86_64-linux curl -fL https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvrtc/linux-x86_64/cuda_nvrtc-linux-x86_64-11.5.119-archive.tar.xz \ - | sudo tar -xJ -C /usr/local/cuda/targets/x86_64-linux --strip-components=1 - echo /usr/local/cuda/targets/x86_64-linux/lib | sudo tee /etc/ld.so.conf.d/cuda-nvrtc.conf + | sudo tar -xJ -C /usr/local/cuda-11.5/targets/x86_64-linux --strip-components=1 + sudo ln -sfnT /usr/local/cuda-11.5 /usr/local/cuda + echo /usr/local/cuda-11.5/targets/x86_64-linux/lib | sudo tee /etc/ld.so.conf.d/cuda-11.5.conf sudo ldconfig # **** gpuocelot **** diff --git a/tinygrad_repo/docs/tensor/ops.md b/tinygrad_repo/docs/tensor/ops.md index d102c697a..4a4e2e2b0 100644 --- a/tinygrad_repo/docs/tensor/ops.md +++ b/tinygrad_repo/docs/tensor/ops.md @@ -52,6 +52,7 @@ ::: tinygrad.Tensor.linear ::: tinygrad.Tensor.sequential ::: tinygrad.Tensor.layernorm +::: tinygrad.Tensor.rmsnorm ::: tinygrad.Tensor.batchnorm ::: tinygrad.Tensor.dropout ::: tinygrad.Tensor.one_hot diff --git a/tinygrad_repo/examples/openpilot/check_image_reads.py b/tinygrad_repo/examples/openpilot/check_image_reads.py index 53a194adf..3aa31e2ee 100644 --- a/tinygrad_repo/examples/openpilot/check_image_reads.py +++ b/tinygrad_repo/examples/openpilot/check_image_reads.py @@ -1,16 +1,16 @@ import pickle, sys from tinygrad.helpers import getenv, is_image_shape, temp from tinygrad.uop.ops import Ops -from tinygrad.viz.serve import VizData, _reconstruct +from tinygrad.viz.serve import VizData, _reconstruct, canonicalize_ast with open(sys.argv[1] if len(sys.argv) > 1 else temp("rewrites.pkl", append_user=True), "rb") as f: data = VizData(pickle.load(f)) steps = [s for key, ctx in zip(data.trace.keys, data.trace.rewrites) if key.display_name.startswith("JIT ") for s in ctx if s.name == "View captured linear"] assert len(steps) == 1, f"expected one JIT compilation, found {len(steps)}" asts = [call.without_after.src[0] for call in _reconstruct(data, steps[0].sink).src] -by_ast = {key.keys[1]: _reconstruct(data, s.sink) for key, ctx in zip(data.trace.keys, data.trace.rewrites) - for s in ctx if s.name == "View Program" and key.keys[1] in asts} -programs = [by_ast[ast] if ast.op is Ops.SINK else ast for ast in asts if ast.op in {Ops.SINK, Ops.PROGRAM}] +by_ast = {canonicalize_ast(_reconstruct(data, b.sink)): _reconstruct(data, p.sink) for ctx in data.trace.rewrites + for b in ctx if b.name == "View Base AST" for p in ctx if p.name == "View Program"} +programs = [by_ast[canonicalize_ast(ast)] if ast.op is Ops.SINK else ast for ast in asts if ast.op in {Ops.SINK, Ops.PROGRAM}] read_image, gated_read_image = 0, 0 for prg in programs: uops = prg.src[1].src diff --git a/tinygrad_repo/extra/gemm/mi350x_uop_matmul_2.py b/tinygrad_repo/extra/gemm/mi350x_uop_matmul_2.py index f3efd581a..b44b6ca70 100644 --- a/tinygrad_repo/extra/gemm/mi350x_uop_matmul_2.py +++ b/tinygrad_repo/extra/gemm/mi350x_uop_matmul_2.py @@ -114,7 +114,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp: C_view = C[gx, :, gy, :].reshape(BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M, BLOCK_N//TC_N, TC_N)[:, warpgroup, warp%16, :, (warp//16)*4] sink = copy(C_view, acc.after(sink), rng=300) else: - sink = C.after(barrier.end(K_outer_loop))[0,0,0,0].store(As[0,0]+Bs[0,0]) + sink = C.after(barrier.end(K_outer_loop))[0,0,0,0].store((As[0,0]+Bs[0,0]).float()) return sink.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify() diff --git a/tinygrad_repo/extra/models/t5.py b/tinygrad_repo/extra/models/t5.py index 7a0fff23f..60dd1fe34 100644 --- a/tinygrad_repo/extra/models/t5.py +++ b/tinygrad_repo/extra/models/t5.py @@ -62,8 +62,7 @@ class T5LayerNorm: # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for # half-precision inputs is done in fp32 - variance = hidden_states.cast(dtypes.float32).pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * Tensor.rsqrt(variance + self.variance_epsilon) + hidden_states = hidden_states.cast(dtypes.float32).rmsnorm(eps=self.variance_epsilon) # convert into half-precision if necessary if self.weight.dtype in [dtypes.float16, dtypes.bfloat16]: diff --git a/tinygrad_repo/extra/thunder/amd/fa.py b/tinygrad_repo/extra/thunder/amd/fa.py index 4498b30e7..30e36967a 100644 --- a/tinygrad_repo/extra/thunder/amd/fa.py +++ b/tinygrad_repo/extra/thunder/amd/fa.py @@ -201,12 +201,15 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False @functools.cache def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True, window:int=0): - code = (pathlib.Path(__file__).parent / "fa_fwd_causal.cpp").read_text() + w4 = bool(getenv("GPTOSS_FA_FWD_W4", 0) and arch == "gfx950" and has_sink and (B, N, H, H_KV, D, window) == (2, 8192, 64, 8, 64, 128)) + full_w4 = bool(getenv("GPTOSS_FA_FULL_FWD_W4", 0) and arch == "gfx950" and has_sink and (B, N, H, H_KV, D, window) == (2, 8192, 64, 8, 64, 0)) + source_name = "fa_fwd_causal_d64.cpp" if full_w4 else "fa_fwd_causal_d64_window.cpp" if w4 else "fa_fwd_causal.cpp" + code = (pathlib.Path(__file__).parent / source_name).read_text() compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math", f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}", f"-DWINDOW={window}"] Q_BLOCK_SIZE = 32 - NUM_WARPS = 8 + NUM_WARPS = 4 if w4 or full_w4 else 8 NUM_THREADS = 64 * NUM_WARPS gsz = (H, (math.ceil((N // Q_BLOCK_SIZE) / NUM_WARPS)), B) lsz = (NUM_THREADS, 1, 1) @@ -217,15 +220,16 @@ def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None mem = (2*B*N*H*D + 2*B*N*H_KV*D) * el + B*H*N * l_vec.dtype.itemsize estimates = Estimates(ops=2*B*H*N*N*D, lds=mem, mem=mem) buf_inputs = (o.base, l_vec.base, q.base, k.base, v.base) + ((sinks.base,) if has_sink else ()) + name = "custom_fa_forward_d64_full_w4" if full_w4 else "custom_fa_forward_d64_window_w4" if w4 else "custom_fa_forward" sink = UOp.sink(*buf_inputs, threadIdx_x, blockIdx_x, blockIdx_y, blockIdx_z, - arg=KernelInfo(name="custom_fa_forward", estimates=estimates)) + arg=KernelInfo(name=name, estimates=estimates)) lib = HIPCCCompiler(arch, compile_args).compile_cached(code) if not getenv("NO_HIPCC"): lib = bytearray(lib) rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata") - struct.pack_into(' q_reg_fl; load<1, qo_tile, _gl_QKVO>(q_reg_fl, g.Qg, {batch_idx, tile_idx, head_idx, 0}); - #if !WINDOW + #if !WINDOW && ATTN_D != 64 // D64 builds scale AFTER the MMA in fp32 to keep forward LSE consistent with backward P mul(q_reg_fl, q_reg_fl, TEMPERATURE_SCALE); // Use sqrtf for clarity #endif copy(q_reg, q_reg_fl); @@ -290,8 +290,8 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p zero(att_block[0]); transpose(k_reg_transposed, k_reg); mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]); - #if WINDOW - mul(att_block[0], att_block[0], TEMPERATURE_SCALE); + #if WINDOW || ATTN_D == 64 + mul(att_block[0], att_block[0], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE #endif __builtin_amdgcn_sched_barrier(0); if constexpr (causal) { @@ -342,8 +342,8 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p zero(att_block[1]); transpose(k_reg_transposed, k_reg); mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]); - #if WINDOW - mul(att_block[1], att_block[1], TEMPERATURE_SCALE); + #if WINDOW || ATTN_D == 64 + mul(att_block[1], att_block[1], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE #endif #if WINDOW // window masks interior tiles that causal skips @@ -409,8 +409,8 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p zero(att_block[0]); transpose(k_reg_transposed, k_reg); mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]); - #if WINDOW - mul(att_block[0], att_block[0], TEMPERATURE_SCALE); + #if WINDOW || ATTN_D == 64 + mul(att_block[0], att_block[0], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE #endif // Finish softmax for QK1 exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]); @@ -480,8 +480,8 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p zero(att_block[1]); transpose(k_reg_transposed, k_reg); mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]); - #if WINDOW - mul(att_block[1], att_block[1], TEMPERATURE_SCALE); + #if WINDOW || ATTN_D == 64 + mul(att_block[1], att_block[1], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE #endif // Finish softmax for QK2 exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]); @@ -549,8 +549,8 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p zero(att_block[0]); transpose(k_reg_transposed, k_reg); mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]); - #if WINDOW - mul(att_block[0], att_block[0], TEMPERATURE_SCALE); + #if WINDOW || ATTN_D == 64 + mul(att_block[0], att_block[0], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE #endif // Finish softmax for QK3 exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]); @@ -614,8 +614,8 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p zero(att_block[1]); transpose(k_reg_transposed, k_reg); mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]); - #if WINDOW - mul(att_block[1], att_block[1], TEMPERATURE_SCALE); + #if WINDOW || ATTN_D == 64 + mul(att_block[1], att_block[1], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE #endif // Finish softmax for QK4 exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]); diff --git a/tinygrad_repo/extra/thunder/amd/fa_fwd_causal_d64.cpp b/tinygrad_repo/extra/thunder/amd/fa_fwd_causal_d64.cpp new file mode 100644 index 000000000..400667ddb --- /dev/null +++ b/tinygrad_repo/extra/thunder/amd/fa_fwd_causal_d64.cpp @@ -0,0 +1,728 @@ +#include "kittens.cuh" + +#ifndef ATTN_B +constexpr int ATTN_B = 16; // batch size +#endif + +#ifndef ATTN_H +constexpr int ATTN_H = 32; // number of heads +#endif + +#ifndef ATTN_H_KV +constexpr int ATTN_H_KV = 8; // number of heads for key and value +#endif + +constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV; // queries per KV head group + +#ifndef ATTN_N +constexpr int ATTN_N = 8192; // sequence length +#endif + +#ifndef ATTN_D +constexpr int ATTN_D = 128; // dimension +#endif +#ifndef ATTN_SINK +#define ATTN_SINK 0 +#endif +#if ATTN_D == 64 +#define FA_VM2 "2" +#define FA_VM4 "4" +#else +#define FA_VM2 "2" +#define FA_VM4 "4" +#endif +constexpr int Q_BLOCK_SIZE = 32; // q block size +constexpr int KV_BLOCK_SIZE = 64; // kv block size +constexpr bool causal = true; +// WINDOW>0: sliding-window attention, query i attends keys in [i-WINDOW+1, i] +#ifndef WINDOW +#define WINDOW 0 +#endif + +#define NUM_WARPS 4 +#define NUM_THREADS (kittens::WARP_THREADS * NUM_WARPS) + +#define MFMA_MASK 0x08 +#define VALU_MASK 0x02 +#define EXP_MASK 0x400 + +using namespace kittens; +using _gl_QKVO = gl; + +using G = kittens::group; + +#define SCHED_BARRIER(mask, cnt, group) __builtin_amdgcn_sched_group_barrier(mask, cnt, group) + +template +__device__ __forceinline__ void sched_barrier_pairs() { + SCHED_BARRIER(MFMA_MASK, 1, Group); + SCHED_BARRIER(VALU_MASK, VALU_CNT, Group); + if constexpr (Pairs > 1) sched_barrier_pairs(); +} + +template +__device__ __forceinline__ void sched_barrier_exp_pairs() { + SCHED_BARRIER(MFMA_MASK, 1, Group); + SCHED_BARRIER(EXP_MASK, EXP_CNT, Group); + if constexpr (Pairs > 1) sched_barrier_exp_pairs(); +} + +template +__device__ inline void exp2(rt_base &dst, const rt_base &src) { + static_assert(std::is_same_v, "Only 32x32 tiles are supported"); + + #pragma unroll + for(int k = 0; k < dst.packed_per_thread; k++) { + dst.data[k] = base_ops::exp2::op(src.data[k]); + } + +} + +template using qo_tile = rt; +template using qo_tile_transposed = rt; +template using kv_tile = rt; +template using kv_tile_transposed = rt; +template using attn_tile = rt; + +/**********************************************************/ +template +__device__ inline void mask_vec2_imm(uint32_t rel_vgpr, uint32_t rel_hi_vgpr, uint32_t neg_inf_vgpr, + uint32_t& x_ref, uint32_t& y_ref) { + + uint64_t x_mask, y_mask; +#if WINDOW + // causal+window in one asm block to not disturb register allocation + asm volatile( + "v_cmp_lt_i32_e64 %0, %4, %5\n\t" + "v_cmp_lt_i32_e64 %1, %4, %7\n\t" + "v_cndmask_b32_e64 %2, %2, %6, %0\n\t" + "v_cndmask_b32_e64 %3, %3, %6, %1\n\t" + "v_cmp_ge_i32_e64 %0, %8, %5\n\t" + "v_cmp_ge_i32_e64 %1, %8, %7\n\t" + "v_cndmask_b32_e64 %2, %2, %6, %0\n\t" + "v_cndmask_b32_e64 %3, %3, %6, %1\n\t" + : "=s"(x_mask), "=s"(y_mask), "+v"(x_ref), "+v"(y_ref) + : "v"(rel_vgpr), "n"(THR_X), "v"(neg_inf_vgpr), "n"(THR_Y), "v"(rel_hi_vgpr) + : "vcc" + ); +#else + asm volatile( + // x: rel < THR_X ? + "v_cmp_lt_i32_e64 %0, %6, %7\n\t" + // y: rel < THR_Y ? + "v_cmp_lt_i32_e64 %1, %6, %9\n\t" + "v_cndmask_b32_e64 %2, %4, %8, %0\n\t" + "v_cndmask_b32_e64 %3, %5, %8, %1\n\t" + : "=s"(x_mask), "=s"(y_mask), "=v"(x_ref), "=v"(y_ref) + : "v"(x_ref), "v"(y_ref), "v"(rel_vgpr), + "n"(THR_X), "v"(neg_inf_vgpr), "n"(THR_Y) + : "vcc" + ); +#endif +} + +template +__device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_inf_v, int lane) { + const int col = lane & 31; // 0..31 column within the 32-wide col tile + + // Absolute positions + const int q_base = q_abs * Q_BLOCK_SIZE; // start index for this Q tile + const int k_base = k_abs * KV_BLOCK_SIZE; // start index for this K/V tile + + // q position for this lane's column + const int q_pos = q_base + col; + + #pragma unroll + for (int i = 0; i < dst.height; ++i) { + // Row base of the 32x* chunk produced by MFMA + const int row_base = (i * 32) + ((lane >> 5) << 2); // multiplesof 4 + + // Relative index of the FIRST element in this row-chunk w.r.t. q_pos + // (smaller rel ⇒ more "future" keys that must be -inf) + const int rel0 = q_pos - (k_base + row_base); + const uint32_t rel = static_cast(rel0); + // rel-WINDOW keeps THR within the inline-constant range + const uint32_t rel_hi = static_cast(rel0 - WINDOW); + + #pragma unroll + for (int j = 0; j < dst.width; ++j) { + auto& d0x = *reinterpret_cast(&dst.tiles[i][j].data[0].x); + auto& d0y = *reinterpret_cast(&dst.tiles[i][j].data[0].y); + auto& d1x = *reinterpret_cast(&dst.tiles[i][j].data[1].x); + auto& d1y = *reinterpret_cast(&dst.tiles[i][j].data[1].y); + auto& d2x = *reinterpret_cast(&dst.tiles[i][j].data[2].x); + auto& d2y = *reinterpret_cast(&dst.tiles[i][j].data[2].y); + auto& d3x = *reinterpret_cast(&dst.tiles[i][j].data[3].x); + auto& d3y = *reinterpret_cast(&dst.tiles[i][j].data[3].y); + auto& d4x = *reinterpret_cast(&dst.tiles[i][j].data[4].x); + auto& d4y = *reinterpret_cast(&dst.tiles[i][j].data[4].y); + auto& d5x = *reinterpret_cast(&dst.tiles[i][j].data[5].x); + auto& d5y = *reinterpret_cast(&dst.tiles[i][j].data[5].y); + auto& d6x = *reinterpret_cast(&dst.tiles[i][j].data[6].x); + auto& d6y = *reinterpret_cast(&dst.tiles[i][j].data[6].y); + auto& d7x = *reinterpret_cast(&dst.tiles[i][j].data[7].x); + auto& d7y = *reinterpret_cast(&dst.tiles[i][j].data[7].y); + + // - reuse a single neg_inf register + // - keep VCC live across the pair + // - avoid reloading -inf or recomputing rel + mask_vec2_imm< 0, 1 >(rel, rel_hi, neg_inf_v, d0x, d0y); + mask_vec2_imm< 2, 3 >(rel, rel_hi, neg_inf_v, d1x, d1y); + mask_vec2_imm< 8, 9 >(rel, rel_hi, neg_inf_v, d2x, d2y); + mask_vec2_imm<10,11 >(rel, rel_hi, neg_inf_v, d3x, d3y); + mask_vec2_imm<16,17 >(rel, rel_hi, neg_inf_v, d4x, d4y); + mask_vec2_imm<18,19 >(rel, rel_hi, neg_inf_v, d5x, d5y); + mask_vec2_imm<24,25 >(rel, rel_hi, neg_inf_v, d6x, d6y); + mask_vec2_imm<26,27 >(rel, rel_hi, neg_inf_v, d7x, d7y); + } + } +} + +/**********************************************************/ + +template struct attn_globals { + _gl_QKVO Qg, Kg, Vg, Og; + gl L_vec; + dim3 grid() { return dim3(ATTN_H, ((ATTN_N / Q_BLOCK_SIZE + NUM_WARPS - 1) / NUM_WARPS), ATTN_B); } + dim3 block() { return dim3(NUM_THREADS); } + size_t dynamic_shared_memory() { return MAX_SHARED_MEMORY; } +}; + +template __launch_bounds__(NUM_THREADS, 2) +__global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_ptr, bf16 *V_ptr +#if ATTN_SINK + , float *Sinks_ptr +#endif + ) { + _gl_QKVO Og{O_ptr, ATTN_B, ATTN_N, ATTN_H, ATTN_D}; + _gl_QKVO Qg{Q_ptr, ATTN_B, ATTN_N, ATTN_H, ATTN_D}; + _gl_QKVO Kg{K_ptr, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D}; + _gl_QKVO Vg{V_ptr, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D}; + gl L_vec{L_vec_ptr, ATTN_B, ATTN_H, 1, ATTN_N}; + attn_globals g{Qg, Kg, Vg, Og, L_vec}; + + extern __shared__ alignment_dummy __shm[]; + shared_allocator al((int*)&__shm[0]); + st_bf (&k_smem)[2] = al.allocate, 2>(); + st_bf (&v_smem)[2] = al.allocate, 2>(); + + const int head_idx = (blockIdx.x % ATTN_H_KV) * GROUP_SIZE + (blockIdx.x / ATTN_H_KV); + const int batch_idx = blockIdx.z; + const int head_idx_kv = head_idx / GROUP_SIZE; + const int block_tile_idx = blockIdx.y; + const int tile_idx = block_tile_idx * NUM_WARPS + warpid(); + const int stagger = warpid() / 4; + const int lane = laneid(); + + const int num_tiles = ATTN_N / KV_BLOCK_SIZE; + const int max_tile_idx = block_tile_idx * NUM_WARPS + NUM_WARPS - 1; + const int max_q_end_pos = (max_tile_idx + 1) * Q_BLOCK_SIZE; + int max_num_tiles = (max_q_end_pos + KV_BLOCK_SIZE - 1) / KV_BLOCK_SIZE; + if constexpr (causal) max_num_tiles = min(max_num_tiles, num_tiles); + else max_num_tiles = num_tiles; + // The existing four-tile prologue/epilogue needs at least four tiles even for the first 128 queries. + max_num_tiles = ((max_num_tiles + 3) / 4) * 4; + const int q_start_pos = tile_idx * Q_BLOCK_SIZE; + +#if WINDOW + // start at the first in-window tile; clamp keeps >=4 tiles for the pipeline unroll + const int block_min_q = block_tile_idx * NUM_WARPS * Q_BLOCK_SIZE; + int min_tile = (block_min_q - WINDOW + 1) / KV_BLOCK_SIZE; + if (min_tile < 0) min_tile = 0; + if (min_tile > max_num_tiles - 4) min_tile = max(0, max_num_tiles - 4); +#else + constexpr int min_tile = 0; +#endif + + constexpr float TEMPERATURE_SCALE = (D == 128) ? 0.08838834764f*1.44269504089f : 0.125f*1.44269504089f; + uint32_t neg_inf_v = 0xff800000; + + // Initialize all of the register tiles. + qo_tile q_reg; // Q and K are both row layout, as we use mma_ABt. + qo_tile_transposed q_reg_transposed; + kv_tile k_reg; + kv_tile_transposed k_reg_transposed; + + kv_tile v_reg; + qo_tile_transposed o_reg; // Output tile. + attn_tile att_block[2]; // attention tile, in float. + attn_tile att_block_bf16; + attn_tile att_block_bf16_in; + typename attn_tile::row_vec max_vec, norm_vec, max_vec_prev, scale_vec; + + zero(o_reg); + zero(norm_vec); + zero(scale_vec); + + using T = typename st_bf::dtype; + constexpr int bytes_per_thread = st_32x32_s::template bytes_per_thread(); + constexpr int bytes_per_memcpy = bytes_per_thread * NUM_THREADS; + constexpr int memcpy_per_tile = KV_BLOCK_SIZE * ATTN_D * sizeof(T) / bytes_per_memcpy; + + uint32_t swizzled_offsets_V[memcpy_per_tile]; + uint32_t swizzled_offsets_K[memcpy_per_tile]; + G::prefill_swizzled_offsets<1, false>(k_smem[0], g.Kg, swizzled_offsets_K); + G::prefill_swizzled_offsets<1, false>(v_smem[0], g.Vg, swizzled_offsets_V); + + G::load<1, false>(k_smem[0], g.Kg, {batch_idx, min_tile, head_idx_kv, 0}, swizzled_offsets_K); + __builtin_amdgcn_s_waitcnt(0); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + + qo_tile q_reg_fl; + load<1, qo_tile, _gl_QKVO>(q_reg_fl, g.Qg, {batch_idx, tile_idx, head_idx, 0}); + #if !WINDOW && ATTN_D != 64 // D64 builds scale AFTER the MMA in fp32 to keep forward LSE consistent with backward P + mul(q_reg_fl, q_reg_fl, TEMPERATURE_SCALE); // Use sqrtf for clarity + #endif + copy(q_reg, q_reg_fl); + transpose(q_reg_transposed, q_reg); + + // All warps then collaboratively load in the first slice of V (V0) and the second slice of K (K1) into shared memory + G::load<1, false>(k_smem[1], g.Kg, {batch_idx, min_tile + 1, head_idx_kv, 0}, swizzled_offsets_K); + // All warps then load in the first slice of K (K0) + G::load<1, false>(v_smem[0], g.Vg, {batch_idx, min_tile, head_idx_kv, 0}, swizzled_offsets_V); + load(k_reg, k_smem[0]); + __builtin_amdgcn_sched_barrier(0); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM2 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + + // Each warp performs QK0 + zero(att_block[0]); + transpose(k_reg_transposed, k_reg); + mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]); + #if WINDOW || ATTN_D == 64 + mul(att_block[0], att_block[0], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE + #endif + __builtin_amdgcn_sched_barrier(0); + if constexpr (causal) { + const int kv_end_pos = (min_tile + 1) * KV_BLOCK_SIZE; + if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, WINDOW ? 1 : 0)) { + mask_kv_tile(att_block[0], tile_idx, min_tile, neg_inf_v, lane); + } + } + // Each warp performs a partial softmax of QK0 (i.e. some of the online softmax up until but not including the second exponential scaling of the attention block likely) +#if WINDOW + // floor the max: min_tile can be fully masked, which would NaN via exp2(-inf - -inf) + zero(max_vec_prev); + add(max_vec_prev, max_vec_prev, -1e4f); + col_max(max_vec, att_block[0], max_vec_prev); +#else + col_max(max_vec, att_block[0]); +#endif + + copy(max_vec_prev, max_vec); + exp2(scale_vec, scale_vec); + + sub_col(att_block[0], att_block[0], max_vec); + exp2(att_block[0].tiles[0][0], att_block[0].tiles[0][0]); + __builtin_amdgcn_sched_barrier(0); + mul_col(o_reg, o_reg, scale_vec); + + if (stagger) { + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + } + + // __builtin_amdgcn_sched_barrier(0); + // All warps then load in the second slice of K (K1) + load(k_reg, k_smem[1]); + // All warps then collaboratively load in the third slice of K (K2) into shared memory + G::load<1, false>(k_smem[0], g.Kg, {batch_idx, min_tile + 2, head_idx_kv, 0}, swizzled_offsets_K); + // All warps then collaboratively load in the second slice of V (V1) into shared memory + G::load<1, false>(v_smem[1], g.Vg, {batch_idx, min_tile + 1, head_idx_kv, 0}, swizzled_offsets_V); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + + // hot loop + for (int j = min_tile + 3; j < max_num_tiles - 1; j += 2) { + // Cluster 0: + // QK1 + zero(att_block[1]); + transpose(k_reg_transposed, k_reg); + mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]); + #if WINDOW || ATTN_D == 64 + mul(att_block[1], att_block[1], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE + #endif +#if WINDOW + // window masks interior tiles that causal skips + mask_kv_tile(att_block[1], tile_idx, j - 2, neg_inf_v, lane); +#endif + // Finish softmax for QK0 + exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]); + mul(norm_vec, norm_vec, scale_vec); + col_sum(norm_vec, att_block[0], norm_vec); + copy(att_block_bf16, att_block[0]); + att_block_bf16_in = *reinterpret_cast*>(&att_block_bf16); + sched_barrier_exp_pairs<6, 3, 1>(); + sched_barrier_pairs<10, 5, 1>(); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 1: + // Load K3 into shared + G::load<1, false>(k_smem[1], g.Kg, {batch_idx, j, head_idx_kv, 0}, swizzled_offsets_K); + // Load V0 into registers + load(v_reg, v_smem[0]); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 2: + // A0V0 + __builtin_amdgcn_s_setprio(1); + mma_AtB(o_reg, v_reg, att_block_bf16_in, o_reg); + // Partial softmax for QK1 + col_max(max_vec, att_block[1], max_vec_prev); + sub(scale_vec, max_vec_prev, max_vec); + copy(max_vec_prev, max_vec); + exp2(scale_vec, scale_vec); + sub_col(att_block[1], att_block[1], max_vec); + exp2(att_block[1].tiles[0][0], att_block[1].tiles[0][0]); + sched_barrier_pairs<10, 5, 2>(); + sched_barrier_exp_pairs<6, 3, 2>(); + __builtin_amdgcn_sched_barrier(0); + mul_col(o_reg, o_reg, scale_vec); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 3: + // Load V2 into shared + G::load<1, false>(v_smem[0], g.Vg, {batch_idx, j - 1, head_idx_kv, 0}, swizzled_offsets_V); + // Load K2 into registers + load(k_reg, k_smem[0]); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 4: + // QK2 + __builtin_amdgcn_s_setprio(1); + zero(att_block[0]); + transpose(k_reg_transposed, k_reg); + mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]); + #if WINDOW || ATTN_D == 64 + mul(att_block[0], att_block[0], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE + #endif + // Finish softmax for QK1 + exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]); + mul(norm_vec, norm_vec, scale_vec); + col_sum(norm_vec, att_block[1], norm_vec); + copy(att_block_bf16, att_block[1]); + att_block_bf16_in = *reinterpret_cast*>(&att_block_bf16); + sched_barrier_exp_pairs<6, 3, 3>(); + sched_barrier_pairs<10, 5, 3>(); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 5: + // Load K4 into shared + G::load<1, false>(k_smem[0], g.Kg, {batch_idx, j + 1, head_idx_kv, 0}, swizzled_offsets_K); + // Load V1 into registers + load(v_reg, v_smem[1]); + if constexpr (causal) { + const int kv_end_pos = (j) * KV_BLOCK_SIZE; + if (WINDOW || q_start_pos < kv_end_pos) { + mask_kv_tile(att_block[0], tile_idx, j - 1, neg_inf_v, lane); + } + } + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 6: + // A1V1 + __builtin_amdgcn_s_setprio(1); + mma_AtB(o_reg, v_reg, att_block_bf16_in, o_reg); + // Partial softmax for QK2 + col_max(max_vec, att_block[0], max_vec_prev); + sub(scale_vec, max_vec_prev, max_vec); + copy(max_vec_prev, max_vec); + exp2(scale_vec, scale_vec); + sub_col(att_block[0], att_block[0], max_vec); + exp2(att_block[0].tiles[0][0], att_block[0].tiles[0][0]); + sched_barrier_pairs<10, 5, 4>(); + sched_barrier_exp_pairs<6, 3, 4>(); + __builtin_amdgcn_sched_barrier(0); + mul_col(o_reg, o_reg, scale_vec); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 7: + // Load V3 into shared + G::load<1, false>(v_smem[1], g.Vg, {batch_idx, j, head_idx_kv, 0}, swizzled_offsets_V); + // Load K3 into registers + load(k_reg, k_smem[1]); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + } + + // Epilogue + // Cluster 0: + // QK3 + zero(att_block[1]); + transpose(k_reg_transposed, k_reg); + mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]); + #if WINDOW || ATTN_D == 64 + mul(att_block[1], att_block[1], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE + #endif + // Finish softmax for QK2 + exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]); + mul(norm_vec, norm_vec, scale_vec); + + col_sum(norm_vec, att_block[0], norm_vec); + copy(att_block_bf16, att_block[0]); + att_block_bf16_in = *reinterpret_cast*>(&att_block_bf16); + sched_barrier_exp_pairs<6, 3, 5>(); + sched_barrier_pairs<10, 5, 5>(); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 1: + // Load K5 into shared + G::load<1, false>(k_smem[1], g.Kg, {batch_idx, max_num_tiles - 1, head_idx_kv, 0}, swizzled_offsets_K); + // Load V2 into registers + load(v_reg, v_smem[0]); + if constexpr (causal) { + const int kv_end_pos = (max_num_tiles - 2) * KV_BLOCK_SIZE; + if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, WINDOW ? 1 : 0)) { + mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 3, neg_inf_v, lane); + } + } + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 2: + // A2V2 + __builtin_amdgcn_s_setprio(1); + mma_AtB(o_reg, v_reg, att_block_bf16_in, o_reg); + // Partial softmax for QK3 + col_max(max_vec, att_block[1], max_vec_prev); + sub(scale_vec, max_vec_prev, max_vec); + copy(max_vec_prev, max_vec); + exp2(scale_vec, scale_vec); + sub_col(att_block[1], att_block[1], max_vec); + exp2(att_block[1].tiles[0][0], att_block[1].tiles[0][0]); + sched_barrier_pairs<10, 5, 6>(); + sched_barrier_exp_pairs<6, 3, 6>(); + __builtin_amdgcn_sched_barrier(0); + mul_col(o_reg, o_reg, scale_vec); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 3: + // Load V4 into shared + G::load<1, false>(v_smem[0], g.Vg, {batch_idx, max_num_tiles - 2, head_idx_kv, 0}, swizzled_offsets_V); + // Load K4 into registers + load(k_reg, k_smem[0]); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 4: + // QK4 + zero(att_block[0]); + transpose(k_reg_transposed, k_reg); + mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]); + #if WINDOW || ATTN_D == 64 + mul(att_block[0], att_block[0], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE + #endif + // Finish softmax for QK3 + exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]); + mul(norm_vec, norm_vec, scale_vec); + col_sum(norm_vec, att_block[1], norm_vec); + copy(att_block_bf16, att_block[1]); + att_block_bf16_in = *reinterpret_cast*>(&att_block_bf16); + sched_barrier_exp_pairs<6, 3, 7>(); + sched_barrier_pairs<10, 5, 7>(); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 5: + // Load V3 into registers + load(v_reg, v_smem[1]); + if constexpr (causal) { + const int kv_end_pos = (max_num_tiles - 1) * KV_BLOCK_SIZE; + if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, 1)) { + mask_kv_tile(att_block[0], tile_idx, max_num_tiles - 2, neg_inf_v, lane); + } + } + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM2 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 6: + // A3V3 + __builtin_amdgcn_s_setprio(1); + mma_AtB(o_reg, v_reg, att_block_bf16_in, o_reg); + // Partial softmax for QK4 + col_max(max_vec, att_block[0], max_vec_prev); + sub(scale_vec, max_vec_prev, max_vec); + copy(max_vec_prev, max_vec); + exp2(scale_vec, scale_vec); + sub_col(att_block[0], att_block[0], max_vec); + exp2(att_block[0].tiles[0][0], att_block[0].tiles[0][0]); + sched_barrier_pairs<10, 5, 8>(); + sched_barrier_exp_pairs<6, 3, 8>(); + __builtin_amdgcn_sched_barrier(0); + mul_col(o_reg, o_reg, scale_vec); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 7: + // Load V5 into shared + G::load<1, false>(v_smem[1], g.Vg, {batch_idx, max_num_tiles - 1, head_idx_kv, 0}, swizzled_offsets_V); + // Load K5 into registers + load(k_reg, k_smem[1]); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM2 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 8: + // QK5 + zero(att_block[1]); + transpose(k_reg_transposed, k_reg); + mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]); + #if WINDOW || ATTN_D == 64 + mul(att_block[1], att_block[1], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE + #endif + // Finish softmax for QK4 + exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]); + mul(norm_vec, norm_vec, scale_vec); + col_sum(norm_vec, att_block[0], norm_vec); + copy(att_block_bf16, att_block[0]); + att_block_bf16_in = *reinterpret_cast*>(&att_block_bf16); + sched_barrier_exp_pairs<6, 3, 9>(); + sched_barrier_pairs<10, 5, 9>(); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 9: + // Load V4 into registers + load(v_reg, v_smem[0]); + if constexpr (causal) { + const int kv_end_pos = (max_num_tiles) * KV_BLOCK_SIZE; + if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, 1)) { + mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 1, neg_inf_v, lane); + } + } + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(0)"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 10: + // A4V4 + mma_AtB(o_reg, v_reg, att_block_bf16_in, o_reg); + // Full softmax for QK5 + col_max(max_vec, att_block[1], max_vec_prev); + sub(scale_vec, max_vec_prev, max_vec); + copy(max_vec_prev, max_vec); + exp2(scale_vec, scale_vec); + + sub_col(att_block[1], att_block[1], max_vec); + exp2(att_block[1].tiles[0][0], att_block[1].tiles[0][0]); + sched_barrier_pairs<10, 5, 10>(); + sched_barrier_exp_pairs<6, 3, 10>(); + __builtin_amdgcn_sched_barrier(0); + + exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]); + mul(norm_vec, norm_vec, scale_vec); + + col_sum(norm_vec, att_block[1], norm_vec); + copy(att_block_bf16, att_block[1]); + att_block_bf16_in = *reinterpret_cast*>(&att_block_bf16); + + __builtin_amdgcn_sched_barrier(0); + mul_col(o_reg, o_reg, scale_vec); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 11: + // Load V5 into registers + load(v_reg, v_smem[1]); + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 12: + // A5V5 + mma_AtB(o_reg, v_reg, att_block_bf16_in, o_reg); +#if ATTN_SINK + { + const float sink_l2 = Sinks_ptr[head_idx] * 1.44269504089f; + typename attn_tile::row_vec sink_term; + mul(sink_term, max_vec, -1.0f); + add(sink_term, sink_term, sink_l2); + exp2(sink_term, sink_term); + add(norm_vec, norm_vec, sink_term); + } +#endif + div_col(o_reg, o_reg, norm_vec); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Conclusion + if (!stagger) { + __builtin_amdgcn_s_barrier(); + } + + qo_tile o_reg_transposed; + transpose(o_reg_transposed, o_reg); + store<1>(g.Og, o_reg_transposed, {batch_idx, tile_idx, head_idx, 0}); + + // multiply by ln(2) + mul(max_vec, max_vec, 0.69314718056f); + log(norm_vec, norm_vec); + add(norm_vec, norm_vec, max_vec); + store(g.L_vec, norm_vec, {batch_idx, head_idx, 0, tile_idx}); +} + +template __global__ void attend_ker(bf16*, float*, bf16*, bf16*, bf16* +#if ATTN_SINK + , float* +#endif +); + +static_assert(ATTN_D == 64 && WINDOW == 0 && NUM_WARPS == 4 && ATTN_N == 8192); +static_assert(sizeof(st_bf<64,64,st_32x32_s>) == 8192); +static_assert(sizeof(st_bf<64,64,st_8x32_s>) == 8192); +static_assert(alignof(st_bf<64,64,st_32x32_s>) <= 16); +static_assert(alignof(st_bf<64,64,st_8x32_s>) <= 16); diff --git a/tinygrad_repo/extra/thunder/amd/fa_fwd_causal_d64_window.cpp b/tinygrad_repo/extra/thunder/amd/fa_fwd_causal_d64_window.cpp new file mode 100644 index 000000000..b1ee3e843 --- /dev/null +++ b/tinygrad_repo/extra/thunder/amd/fa_fwd_causal_d64_window.cpp @@ -0,0 +1,728 @@ +#include "kittens.cuh" + +#ifndef ATTN_B +constexpr int ATTN_B = 16; // batch size +#endif + +#ifndef ATTN_H +constexpr int ATTN_H = 32; // number of heads +#endif + +#ifndef ATTN_H_KV +constexpr int ATTN_H_KV = 8; // number of heads for key and value +#endif + +constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV; // queries per KV head group + +#ifndef ATTN_N +constexpr int ATTN_N = 8192; // sequence length +#endif + +#ifndef ATTN_D +constexpr int ATTN_D = 128; // dimension +#endif +#ifndef ATTN_SINK +#define ATTN_SINK 0 +#endif +#if ATTN_D == 64 +#define FA_VM2 "2" +#define FA_VM4 "4" +#else +#define FA_VM2 "2" +#define FA_VM4 "4" +#endif +constexpr int Q_BLOCK_SIZE = 32; // q block size +constexpr int KV_BLOCK_SIZE = 64; // kv block size +constexpr bool causal = true; +// WINDOW>0: sliding-window attention, query i attends keys in [i-WINDOW+1, i] +#ifndef WINDOW +#define WINDOW 0 +#endif + +#define NUM_WARPS 4 +#define NUM_THREADS (kittens::WARP_THREADS * NUM_WARPS) + +#define MFMA_MASK 0x08 +#define VALU_MASK 0x02 +#define EXP_MASK 0x400 + +using namespace kittens; +using _gl_QKVO = gl; + +using G = kittens::group; + +#define SCHED_BARRIER(mask, cnt, group) __builtin_amdgcn_sched_group_barrier(mask, cnt, group) + +template +__device__ __forceinline__ void sched_barrier_pairs() { + SCHED_BARRIER(MFMA_MASK, 1, Group); + SCHED_BARRIER(VALU_MASK, VALU_CNT, Group); + if constexpr (Pairs > 1) sched_barrier_pairs(); +} + +template +__device__ __forceinline__ void sched_barrier_exp_pairs() { + SCHED_BARRIER(MFMA_MASK, 1, Group); + SCHED_BARRIER(EXP_MASK, EXP_CNT, Group); + if constexpr (Pairs > 1) sched_barrier_exp_pairs(); +} + +template +__device__ inline void exp2(rt_base &dst, const rt_base &src) { + static_assert(std::is_same_v, "Only 32x32 tiles are supported"); + + #pragma unroll + for(int k = 0; k < dst.packed_per_thread; k++) { + dst.data[k] = base_ops::exp2::op(src.data[k]); + } + +} + +template using qo_tile = rt; +template using qo_tile_transposed = rt; +template using kv_tile = rt; +template using kv_tile_transposed = rt; +template using attn_tile = rt; + +/**********************************************************/ +template +__device__ inline void mask_vec2_imm(uint32_t rel_vgpr, uint32_t rel_hi_vgpr, uint32_t neg_inf_vgpr, + uint32_t& x_ref, uint32_t& y_ref) { + + uint64_t x_mask, y_mask; +#if WINDOW + // causal+window in one asm block to not disturb register allocation + asm volatile( + "v_cmp_lt_i32_e64 %0, %4, %5\n\t" + "v_cmp_lt_i32_e64 %1, %4, %7\n\t" + "v_cndmask_b32_e64 %2, %2, %6, %0\n\t" + "v_cndmask_b32_e64 %3, %3, %6, %1\n\t" + "v_cmp_ge_i32_e64 %0, %8, %5\n\t" + "v_cmp_ge_i32_e64 %1, %8, %7\n\t" + "v_cndmask_b32_e64 %2, %2, %6, %0\n\t" + "v_cndmask_b32_e64 %3, %3, %6, %1\n\t" + : "=s"(x_mask), "=s"(y_mask), "+v"(x_ref), "+v"(y_ref) + : "v"(rel_vgpr), "n"(THR_X), "v"(neg_inf_vgpr), "n"(THR_Y), "v"(rel_hi_vgpr) + : "vcc" + ); +#else + asm volatile( + // x: rel < THR_X ? + "v_cmp_lt_i32_e64 %0, %6, %7\n\t" + // y: rel < THR_Y ? + "v_cmp_lt_i32_e64 %1, %6, %9\n\t" + "v_cndmask_b32_e64 %2, %4, %8, %0\n\t" + "v_cndmask_b32_e64 %3, %5, %8, %1\n\t" + : "=s"(x_mask), "=s"(y_mask), "=v"(x_ref), "=v"(y_ref) + : "v"(x_ref), "v"(y_ref), "v"(rel_vgpr), + "n"(THR_X), "v"(neg_inf_vgpr), "n"(THR_Y) + : "vcc" + ); +#endif +} + +template +__device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_inf_v, int lane) { + const int col = lane & 31; // 0..31 column within the 32-wide col tile + + // Absolute positions + const int q_base = q_abs * Q_BLOCK_SIZE; // start index for this Q tile + const int k_base = k_abs * KV_BLOCK_SIZE; // start index for this K/V tile + + // q position for this lane's column + const int q_pos = q_base + col; + + #pragma unroll + for (int i = 0; i < dst.height; ++i) { + // Row base of the 32x* chunk produced by MFMA + const int row_base = (i * 32) + ((lane >> 5) << 2); // multiplesof 4 + + // Relative index of the FIRST element in this row-chunk w.r.t. q_pos + // (smaller rel ⇒ more "future" keys that must be -inf) + const int rel0 = q_pos - (k_base + row_base); + const uint32_t rel = static_cast(rel0); + // rel-WINDOW keeps THR within the inline-constant range + const uint32_t rel_hi = static_cast(rel0 - WINDOW); + + #pragma unroll + for (int j = 0; j < dst.width; ++j) { + auto& d0x = *reinterpret_cast(&dst.tiles[i][j].data[0].x); + auto& d0y = *reinterpret_cast(&dst.tiles[i][j].data[0].y); + auto& d1x = *reinterpret_cast(&dst.tiles[i][j].data[1].x); + auto& d1y = *reinterpret_cast(&dst.tiles[i][j].data[1].y); + auto& d2x = *reinterpret_cast(&dst.tiles[i][j].data[2].x); + auto& d2y = *reinterpret_cast(&dst.tiles[i][j].data[2].y); + auto& d3x = *reinterpret_cast(&dst.tiles[i][j].data[3].x); + auto& d3y = *reinterpret_cast(&dst.tiles[i][j].data[3].y); + auto& d4x = *reinterpret_cast(&dst.tiles[i][j].data[4].x); + auto& d4y = *reinterpret_cast(&dst.tiles[i][j].data[4].y); + auto& d5x = *reinterpret_cast(&dst.tiles[i][j].data[5].x); + auto& d5y = *reinterpret_cast(&dst.tiles[i][j].data[5].y); + auto& d6x = *reinterpret_cast(&dst.tiles[i][j].data[6].x); + auto& d6y = *reinterpret_cast(&dst.tiles[i][j].data[6].y); + auto& d7x = *reinterpret_cast(&dst.tiles[i][j].data[7].x); + auto& d7y = *reinterpret_cast(&dst.tiles[i][j].data[7].y); + + // - reuse a single neg_inf register + // - keep VCC live across the pair + // - avoid reloading -inf or recomputing rel + mask_vec2_imm< 0, 1 >(rel, rel_hi, neg_inf_v, d0x, d0y); + mask_vec2_imm< 2, 3 >(rel, rel_hi, neg_inf_v, d1x, d1y); + mask_vec2_imm< 8, 9 >(rel, rel_hi, neg_inf_v, d2x, d2y); + mask_vec2_imm<10,11 >(rel, rel_hi, neg_inf_v, d3x, d3y); + mask_vec2_imm<16,17 >(rel, rel_hi, neg_inf_v, d4x, d4y); + mask_vec2_imm<18,19 >(rel, rel_hi, neg_inf_v, d5x, d5y); + mask_vec2_imm<24,25 >(rel, rel_hi, neg_inf_v, d6x, d6y); + mask_vec2_imm<26,27 >(rel, rel_hi, neg_inf_v, d7x, d7y); + } + } +} + +/**********************************************************/ + +template struct attn_globals { + _gl_QKVO Qg, Kg, Vg, Og; + gl L_vec; + dim3 grid() { return dim3(ATTN_H, ((ATTN_N / Q_BLOCK_SIZE + NUM_WARPS - 1) / NUM_WARPS), ATTN_B); } + dim3 block() { return dim3(NUM_THREADS); } + size_t dynamic_shared_memory() { return MAX_SHARED_MEMORY; } +}; + +template __launch_bounds__(NUM_THREADS, 2) +__global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_ptr, bf16 *V_ptr +#if ATTN_SINK + , float *Sinks_ptr +#endif + ) { + _gl_QKVO Og{O_ptr, ATTN_B, ATTN_N, ATTN_H, ATTN_D}; + _gl_QKVO Qg{Q_ptr, ATTN_B, ATTN_N, ATTN_H, ATTN_D}; + _gl_QKVO Kg{K_ptr, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D}; + _gl_QKVO Vg{V_ptr, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D}; + gl L_vec{L_vec_ptr, ATTN_B, ATTN_H, 1, ATTN_N}; + attn_globals g{Qg, Kg, Vg, Og, L_vec}; + + extern __shared__ alignment_dummy __shm[]; + shared_allocator al((int*)&__shm[0]); + st_bf (&k_smem)[2] = al.allocate, 2>(); + st_bf (&v_smem)[2] = al.allocate, 2>(); + + const int head_idx = (blockIdx.x % ATTN_H_KV) * GROUP_SIZE + (blockIdx.x / ATTN_H_KV); + const int batch_idx = blockIdx.z; + const int head_idx_kv = head_idx / GROUP_SIZE; + const int block_tile_idx = blockIdx.y; + const int tile_idx = block_tile_idx * NUM_WARPS + warpid(); + const int stagger = warpid() / 4; + const int lane = laneid(); + + const int num_tiles = ATTN_N / KV_BLOCK_SIZE; + const int max_tile_idx = block_tile_idx * NUM_WARPS + NUM_WARPS - 1; + const int max_q_end_pos = (max_tile_idx + 1) * Q_BLOCK_SIZE; + int max_num_tiles = (max_q_end_pos + KV_BLOCK_SIZE - 1) / KV_BLOCK_SIZE; + if constexpr (causal) max_num_tiles = min(max_num_tiles, num_tiles); + else max_num_tiles = num_tiles; + // The existing four-tile prologue/epilogue needs at least four tiles even for the first 128 queries. + max_num_tiles = max(max_num_tiles, 4); + const int q_start_pos = tile_idx * Q_BLOCK_SIZE; + +#if WINDOW + // start at the first in-window tile; clamp keeps >=4 tiles for the pipeline unroll + const int block_min_q = block_tile_idx * NUM_WARPS * Q_BLOCK_SIZE; + int min_tile = (block_min_q - WINDOW + 1) / KV_BLOCK_SIZE; + if (min_tile < 0) min_tile = 0; + if (min_tile > max_num_tiles - 4) min_tile = max(0, max_num_tiles - 4); +#else + constexpr int min_tile = 0; +#endif + + constexpr float TEMPERATURE_SCALE = (D == 128) ? 0.08838834764f*1.44269504089f : 0.125f*1.44269504089f; + uint32_t neg_inf_v = 0xff800000; + + // Initialize all of the register tiles. + qo_tile q_reg; // Q and K are both row layout, as we use mma_ABt. + qo_tile_transposed q_reg_transposed; + kv_tile k_reg; + kv_tile_transposed k_reg_transposed; + + kv_tile v_reg; + qo_tile_transposed o_reg; // Output tile. + attn_tile att_block[2]; // attention tile, in float. + attn_tile att_block_bf16; + attn_tile att_block_bf16_in; + typename attn_tile::row_vec max_vec, norm_vec, max_vec_prev, scale_vec; + + zero(o_reg); + zero(norm_vec); + zero(scale_vec); + + using T = typename st_bf::dtype; + constexpr int bytes_per_thread = st_32x32_s::template bytes_per_thread(); + constexpr int bytes_per_memcpy = bytes_per_thread * NUM_THREADS; + constexpr int memcpy_per_tile = KV_BLOCK_SIZE * ATTN_D * sizeof(T) / bytes_per_memcpy; + + uint32_t swizzled_offsets_V[memcpy_per_tile]; + uint32_t swizzled_offsets_K[memcpy_per_tile]; + G::prefill_swizzled_offsets<1, false>(k_smem[0], g.Kg, swizzled_offsets_K); + G::prefill_swizzled_offsets<1, false>(v_smem[0], g.Vg, swizzled_offsets_V); + + G::load<1, false>(k_smem[0], g.Kg, {batch_idx, min_tile, head_idx_kv, 0}, swizzled_offsets_K); + __builtin_amdgcn_s_waitcnt(0); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + + qo_tile q_reg_fl; + load<1, qo_tile, _gl_QKVO>(q_reg_fl, g.Qg, {batch_idx, tile_idx, head_idx, 0}); + #if !WINDOW && ATTN_D != 64 // D64 builds scale AFTER the MMA in fp32 to keep forward LSE consistent with backward P + mul(q_reg_fl, q_reg_fl, TEMPERATURE_SCALE); // Use sqrtf for clarity + #endif + copy(q_reg, q_reg_fl); + transpose(q_reg_transposed, q_reg); + + // All warps then collaboratively load in the first slice of V (V0) and the second slice of K (K1) into shared memory + G::load<1, false>(k_smem[1], g.Kg, {batch_idx, min_tile + 1, head_idx_kv, 0}, swizzled_offsets_K); + // All warps then load in the first slice of K (K0) + G::load<1, false>(v_smem[0], g.Vg, {batch_idx, min_tile, head_idx_kv, 0}, swizzled_offsets_V); + load(k_reg, k_smem[0]); + __builtin_amdgcn_sched_barrier(0); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM2 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + + // Each warp performs QK0 + zero(att_block[0]); + transpose(k_reg_transposed, k_reg); + mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]); + #if WINDOW || ATTN_D == 64 + mul(att_block[0], att_block[0], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE + #endif + __builtin_amdgcn_sched_barrier(0); + if constexpr (causal) { + const int kv_end_pos = (min_tile + 1) * KV_BLOCK_SIZE; + if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, WINDOW ? 1 : 0)) { + mask_kv_tile(att_block[0], tile_idx, min_tile, neg_inf_v, lane); + } + } + // Each warp performs a partial softmax of QK0 (i.e. some of the online softmax up until but not including the second exponential scaling of the attention block likely) +#if WINDOW + // floor the max: min_tile can be fully masked, which would NaN via exp2(-inf - -inf) + zero(max_vec_prev); + add(max_vec_prev, max_vec_prev, -1e4f); + col_max(max_vec, att_block[0], max_vec_prev); +#else + col_max(max_vec, att_block[0]); +#endif + + copy(max_vec_prev, max_vec); + exp2(scale_vec, scale_vec); + + sub_col(att_block[0], att_block[0], max_vec); + exp2(att_block[0].tiles[0][0], att_block[0].tiles[0][0]); + __builtin_amdgcn_sched_barrier(0); + mul_col(o_reg, o_reg, scale_vec); + + if (stagger) { + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + } + + // __builtin_amdgcn_sched_barrier(0); + // All warps then load in the second slice of K (K1) + load(k_reg, k_smem[1]); + // All warps then collaboratively load in the third slice of K (K2) into shared memory + G::load<1, false>(k_smem[0], g.Kg, {batch_idx, min_tile + 2, head_idx_kv, 0}, swizzled_offsets_K); + // All warps then collaboratively load in the second slice of V (V1) into shared memory + G::load<1, false>(v_smem[1], g.Vg, {batch_idx, min_tile + 1, head_idx_kv, 0}, swizzled_offsets_V); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + + // hot loop + for (int j = min_tile + 3; j < max_num_tiles - 1; j += 2) { + // Cluster 0: + // QK1 + zero(att_block[1]); + transpose(k_reg_transposed, k_reg); + mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]); + #if WINDOW || ATTN_D == 64 + mul(att_block[1], att_block[1], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE + #endif +#if WINDOW + // window masks interior tiles that causal skips + mask_kv_tile(att_block[1], tile_idx, j - 2, neg_inf_v, lane); +#endif + // Finish softmax for QK0 + exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]); + mul(norm_vec, norm_vec, scale_vec); + col_sum(norm_vec, att_block[0], norm_vec); + copy(att_block_bf16, att_block[0]); + att_block_bf16_in = *reinterpret_cast*>(&att_block_bf16); + sched_barrier_exp_pairs<6, 3, 1>(); + sched_barrier_pairs<10, 5, 1>(); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 1: + // Load K3 into shared + G::load<1, false>(k_smem[1], g.Kg, {batch_idx, j, head_idx_kv, 0}, swizzled_offsets_K); + // Load V0 into registers + load(v_reg, v_smem[0]); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 2: + // A0V0 + __builtin_amdgcn_s_setprio(1); + mma_AtB(o_reg, v_reg, att_block_bf16_in, o_reg); + // Partial softmax for QK1 + col_max(max_vec, att_block[1], max_vec_prev); + sub(scale_vec, max_vec_prev, max_vec); + copy(max_vec_prev, max_vec); + exp2(scale_vec, scale_vec); + sub_col(att_block[1], att_block[1], max_vec); + exp2(att_block[1].tiles[0][0], att_block[1].tiles[0][0]); + sched_barrier_pairs<10, 5, 2>(); + sched_barrier_exp_pairs<6, 3, 2>(); + __builtin_amdgcn_sched_barrier(0); + mul_col(o_reg, o_reg, scale_vec); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 3: + // Load V2 into shared + G::load<1, false>(v_smem[0], g.Vg, {batch_idx, j - 1, head_idx_kv, 0}, swizzled_offsets_V); + // Load K2 into registers + load(k_reg, k_smem[0]); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 4: + // QK2 + __builtin_amdgcn_s_setprio(1); + zero(att_block[0]); + transpose(k_reg_transposed, k_reg); + mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]); + #if WINDOW || ATTN_D == 64 + mul(att_block[0], att_block[0], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE + #endif + // Finish softmax for QK1 + exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]); + mul(norm_vec, norm_vec, scale_vec); + col_sum(norm_vec, att_block[1], norm_vec); + copy(att_block_bf16, att_block[1]); + att_block_bf16_in = *reinterpret_cast*>(&att_block_bf16); + sched_barrier_exp_pairs<6, 3, 3>(); + sched_barrier_pairs<10, 5, 3>(); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 5: + // Load K4 into shared + G::load<1, false>(k_smem[0], g.Kg, {batch_idx, j + 1, head_idx_kv, 0}, swizzled_offsets_K); + // Load V1 into registers + load(v_reg, v_smem[1]); + if constexpr (causal) { + const int kv_end_pos = (j) * KV_BLOCK_SIZE; + if (WINDOW || q_start_pos < kv_end_pos) { + mask_kv_tile(att_block[0], tile_idx, j - 1, neg_inf_v, lane); + } + } + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 6: + // A1V1 + __builtin_amdgcn_s_setprio(1); + mma_AtB(o_reg, v_reg, att_block_bf16_in, o_reg); + // Partial softmax for QK2 + col_max(max_vec, att_block[0], max_vec_prev); + sub(scale_vec, max_vec_prev, max_vec); + copy(max_vec_prev, max_vec); + exp2(scale_vec, scale_vec); + sub_col(att_block[0], att_block[0], max_vec); + exp2(att_block[0].tiles[0][0], att_block[0].tiles[0][0]); + sched_barrier_pairs<10, 5, 4>(); + sched_barrier_exp_pairs<6, 3, 4>(); + __builtin_amdgcn_sched_barrier(0); + mul_col(o_reg, o_reg, scale_vec); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 7: + // Load V3 into shared + G::load<1, false>(v_smem[1], g.Vg, {batch_idx, j, head_idx_kv, 0}, swizzled_offsets_V); + // Load K3 into registers + load(k_reg, k_smem[1]); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + } + + // Epilogue + // Cluster 0: + // QK3 + zero(att_block[1]); + transpose(k_reg_transposed, k_reg); + mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]); + #if WINDOW || ATTN_D == 64 + mul(att_block[1], att_block[1], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE + #endif + // Finish softmax for QK2 + exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]); + mul(norm_vec, norm_vec, scale_vec); + + col_sum(norm_vec, att_block[0], norm_vec); + copy(att_block_bf16, att_block[0]); + att_block_bf16_in = *reinterpret_cast*>(&att_block_bf16); + sched_barrier_exp_pairs<6, 3, 5>(); + sched_barrier_pairs<10, 5, 5>(); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 1: + // Load K5 into shared + G::load<1, false>(k_smem[1], g.Kg, {batch_idx, max_num_tiles - 1, head_idx_kv, 0}, swizzled_offsets_K); + // Load V2 into registers + load(v_reg, v_smem[0]); + if constexpr (causal) { + const int kv_end_pos = (max_num_tiles - 2) * KV_BLOCK_SIZE; + if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, WINDOW ? 1 : 0)) { + mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 3, neg_inf_v, lane); + } + } + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 2: + // A2V2 + __builtin_amdgcn_s_setprio(1); + mma_AtB(o_reg, v_reg, att_block_bf16_in, o_reg); + // Partial softmax for QK3 + col_max(max_vec, att_block[1], max_vec_prev); + sub(scale_vec, max_vec_prev, max_vec); + copy(max_vec_prev, max_vec); + exp2(scale_vec, scale_vec); + sub_col(att_block[1], att_block[1], max_vec); + exp2(att_block[1].tiles[0][0], att_block[1].tiles[0][0]); + sched_barrier_pairs<10, 5, 6>(); + sched_barrier_exp_pairs<6, 3, 6>(); + __builtin_amdgcn_sched_barrier(0); + mul_col(o_reg, o_reg, scale_vec); + __builtin_amdgcn_s_setprio(0); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 3: + // Load V4 into shared + G::load<1, false>(v_smem[0], g.Vg, {batch_idx, max_num_tiles - 2, head_idx_kv, 0}, swizzled_offsets_V); + // Load K4 into registers + load(k_reg, k_smem[0]); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM4 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 4: + // QK4 + zero(att_block[0]); + transpose(k_reg_transposed, k_reg); + mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]); + #if WINDOW || ATTN_D == 64 + mul(att_block[0], att_block[0], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE + #endif + // Finish softmax for QK3 + exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]); + mul(norm_vec, norm_vec, scale_vec); + col_sum(norm_vec, att_block[1], norm_vec); + copy(att_block_bf16, att_block[1]); + att_block_bf16_in = *reinterpret_cast*>(&att_block_bf16); + sched_barrier_exp_pairs<6, 3, 7>(); + sched_barrier_pairs<10, 5, 7>(); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 5: + // Load V3 into registers + load(v_reg, v_smem[1]); + if constexpr (causal) { + const int kv_end_pos = (max_num_tiles - 1) * KV_BLOCK_SIZE; + if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, 1)) { + mask_kv_tile(att_block[0], tile_idx, max_num_tiles - 2, neg_inf_v, lane); + } + } + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM2 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 6: + // A3V3 + __builtin_amdgcn_s_setprio(1); + mma_AtB(o_reg, v_reg, att_block_bf16_in, o_reg); + // Partial softmax for QK4 + col_max(max_vec, att_block[0], max_vec_prev); + sub(scale_vec, max_vec_prev, max_vec); + copy(max_vec_prev, max_vec); + exp2(scale_vec, scale_vec); + sub_col(att_block[0], att_block[0], max_vec); + exp2(att_block[0].tiles[0][0], att_block[0].tiles[0][0]); + sched_barrier_pairs<10, 5, 8>(); + sched_barrier_exp_pairs<6, 3, 8>(); + __builtin_amdgcn_sched_barrier(0); + mul_col(o_reg, o_reg, scale_vec); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 7: + // Load V5 into shared + G::load<1, false>(v_smem[1], g.Vg, {batch_idx, max_num_tiles - 1, head_idx_kv, 0}, swizzled_offsets_V); + // Load K5 into registers + load(k_reg, k_smem[1]); + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(" FA_VM2 ")"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 8: + // QK5 + zero(att_block[1]); + transpose(k_reg_transposed, k_reg); + mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]); + #if WINDOW || ATTN_D == 64 + mul(att_block[1], att_block[1], TEMPERATURE_SCALE); // fp32 scale-after -> unbiased online LSE + #endif + // Finish softmax for QK4 + exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]); + mul(norm_vec, norm_vec, scale_vec); + col_sum(norm_vec, att_block[0], norm_vec); + copy(att_block_bf16, att_block[0]); + att_block_bf16_in = *reinterpret_cast*>(&att_block_bf16); + sched_barrier_exp_pairs<6, 3, 9>(); + sched_barrier_pairs<10, 5, 9>(); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 9: + // Load V4 into registers + load(v_reg, v_smem[0]); + if constexpr (causal) { + const int kv_end_pos = (max_num_tiles) * KV_BLOCK_SIZE; + if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, 1)) { + mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 1, neg_inf_v, lane); + } + } + asm volatile("s_waitcnt lgkmcnt(0)"); + asm volatile("s_waitcnt vmcnt(0)"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 10: + // A4V4 + mma_AtB(o_reg, v_reg, att_block_bf16_in, o_reg); + // Full softmax for QK5 + col_max(max_vec, att_block[1], max_vec_prev); + sub(scale_vec, max_vec_prev, max_vec); + copy(max_vec_prev, max_vec); + exp2(scale_vec, scale_vec); + + sub_col(att_block[1], att_block[1], max_vec); + exp2(att_block[1].tiles[0][0], att_block[1].tiles[0][0]); + sched_barrier_pairs<10, 5, 10>(); + sched_barrier_exp_pairs<6, 3, 10>(); + __builtin_amdgcn_sched_barrier(0); + + exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]); + mul(norm_vec, norm_vec, scale_vec); + + col_sum(norm_vec, att_block[1], norm_vec); + copy(att_block_bf16, att_block[1]); + att_block_bf16_in = *reinterpret_cast*>(&att_block_bf16); + + __builtin_amdgcn_sched_barrier(0); + mul_col(o_reg, o_reg, scale_vec); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 11: + // Load V5 into registers + load(v_reg, v_smem[1]); + asm volatile("s_waitcnt lgkmcnt(0)"); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Cluster 12: + // A5V5 + mma_AtB(o_reg, v_reg, att_block_bf16_in, o_reg); +#if ATTN_SINK + { + const float sink_l2 = Sinks_ptr[head_idx] * 1.44269504089f; + typename attn_tile::row_vec sink_term; + mul(sink_term, max_vec, -1.0f); + add(sink_term, sink_term, sink_l2); + exp2(sink_term, sink_term); + add(norm_vec, norm_vec, sink_term); + } +#endif + div_col(o_reg, o_reg, norm_vec); + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + + // Conclusion + if (!stagger) { + __builtin_amdgcn_s_barrier(); + } + + qo_tile o_reg_transposed; + transpose(o_reg_transposed, o_reg); + store<1>(g.Og, o_reg_transposed, {batch_idx, tile_idx, head_idx, 0}); + + // multiply by ln(2) + mul(max_vec, max_vec, 0.69314718056f); + log(norm_vec, norm_vec); + add(norm_vec, norm_vec, max_vec); + store(g.L_vec, norm_vec, {batch_idx, head_idx, 0, tile_idx}); +} + +template __global__ void attend_ker(bf16*, float*, bf16*, bf16*, bf16* +#if ATTN_SINK + , float* +#endif +); + +static_assert(ATTN_D == 64 && WINDOW == 128 && NUM_WARPS == 4 && ATTN_N == 8192); +static_assert(sizeof(st_bf<64,64,st_32x32_s>) == 8192); +static_assert(sizeof(st_bf<64,64,st_8x32_s>) == 8192); +static_assert(alignof(st_bf<64,64,st_32x32_s>) <= 16); +static_assert(alignof(st_bf<64,64,st_8x32_s>) <= 16); diff --git a/tinygrad_repo/extra/torch_backend/backend.py b/tinygrad_repo/extra/torch_backend/backend.py index 85466aa82..49ca13d1e 100644 --- a/tinygrad_repo/extra/torch_backend/backend.py +++ b/tinygrad_repo/extra/torch_backend/backend.py @@ -604,7 +604,7 @@ tiny_backend = {**tiny_backend_out, **{ "aten.topk": Tensor.topk, "aten.constant_pad_nd": lambda self, padding, value=0.0: self.pad(padding, mode="constant", value=value).contiguous(), "aten.cumsum": lambda self, dim: self.cumsum(dim), - "aten.logsumexp": lambda self, axis, keepdim=False: self.logsumexp(axis[0], keepdim=keepdim), + "aten.logsumexp": lambda self, axis, keepdim=False: self.logsumexp(axis, keepdim=keepdim), "aten.roll": Tensor.roll, "aten.logcumsumexp": Tensor.logcumsumexp, "aten.lerp.Tensor": Tensor.lerp, diff --git a/tinygrad_repo/spec/tinyspec.pdf b/tinygrad_repo/spec/tinyspec.pdf index 7dec94f0b..3e062de52 100644 Binary files a/tinygrad_repo/spec/tinyspec.pdf and b/tinygrad_repo/spec/tinyspec.pdf differ diff --git a/tinygrad_repo/spec/tinyspec.tex b/tinygrad_repo/spec/tinyspec.tex index 1a31a03eb..9241c4d5e 100644 --- a/tinygrad_repo/spec/tinyspec.tex +++ b/tinygrad_repo/spec/tinyspec.tex @@ -87,7 +87,6 @@ A value \op{Call} is void: its \op{Sink} body stores to \op{Param}s bound positi \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\ \midrule \op{Permute} & $(T,)$ & axis order $\pi$ & Reorder axes. $\pi = (1,0)$ is transpose. \\ -\op{Flip} & $(T,)$ & bools $\mathbf{f}$ & Reverse along flagged axes. \\ \op{Reshape} & $(T, \mathbf{s'})$ & --- & Reinterpret in row-major order. $\prod s_k = \prod s'_k$. \\ \op{Expand} & $(T, \mathbf{n})$ & --- & Prepend axes $\mathbf{n}$ on the left. Output shape is $\mathbf{n} + T.\mathrm{shape}$. \\ \op{Pad} & $(T, \mathbf{o}, \mathbf{s'})$ & --- & Place $T$ at offset $o_k$ in an invalid-filled output of shape $s'_k$. \\ diff --git a/tinygrad_repo/test/amd/test_emu2_pcode.py b/tinygrad_repo/test/amd/test_emu2_pcode.py index ff6366372..071d6a940 100644 --- a/tinygrad_repo/test/amd/test_emu2_pcode.py +++ b/tinygrad_repo/test/amd/test_emu2_pcode.py @@ -83,6 +83,13 @@ class TestParseExpr(unittest.TestCase): """Test parsing hex literals.""" self.assertIs(parse_expr('0xFF', {}), UOp.const(255, dtypes.uint32)) + def test_24bit_operands(self): + for value, signed, unsigned in ((0x3FFFFFFF, -1, 0xFFFFFF), (0xFF800000, -8388608, 0x800000), (0xAB000001, 1, 1)): + with self.subTest(value=value): + src = {'S0': UOp.const(value, dtypes.uint32)} + self.assertEqual(parse_expr('S0.i24', src).ssimplify(), signed) + self.assertEqual(parse_expr('S0.u24', src).ssimplify(), unsigned) + def test_variable_lookup(self): """Test variable lookup in parse_expr.""" vrs = {'x': UOp.const(42, dtypes.uint32)} diff --git a/tinygrad_repo/test/backend/test_ops.py b/tinygrad_repo/test/backend/test_ops.py index 5483cb35d..6967ea6e7 100644 --- a/tinygrad_repo/test/backend/test_ops.py +++ b/tinygrad_repo/test/backend/test_ops.py @@ -920,6 +920,17 @@ class TestOps(unittest.TestCase): self.assertEqual(a, b) self.assertEqual(Tensor(-1).contiguous().div(4, rounding_mode="trunc").item(), 0) # NOTE this is trunc-div behaviour + @Context(TRANSCENDENTAL=2) + def test_sin_transcendental(self): self.test_sin() + @Context(TRANSCENDENTAL=2) + def test_cos_transcendental(self): self.test_cos() + @Context(TRANSCENDENTAL=2) + def test_tan_transcendental(self): self.test_tan() + @Context(TRANSCENDENTAL=2) + def test_exp_transcendental(self): self.test_exp() + @Context(TRANSCENDENTAL=2) + def test_log_transcendental(self): self.test_log() + @unittest.skipIf(DEV.renderer == "NAK", "MUFU.SIN is not accurate enough") def test_sin(self): helper_test_op([(45,65)], lambda x: x.sin()) @@ -1778,6 +1789,10 @@ class TestOps(unittest.TestCase): helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=3, dim=0), lambda x: x.normalize(p=3, dim=0), atol=1e-7, grad_atol=1e-7) helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=0), lambda x: x.normalize(p=0), atol=1e-7, grad_atol=1e-7) helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=-1), lambda x: x.normalize(p=-1), atol=1e-7, grad_atol=1e-7) + def test_rmsnorm(self): + helper_test_op([(45,65)], lambda x: torch.nn.functional.rms_norm(x, (65,), eps=1e-6), lambda x: x.rmsnorm(), atol=1e-6, grad_atol=1e-6) + helper_test_op([(10,10,10)], lambda x: torch.nn.functional.rms_norm(x, (10,10), eps=1e-3), lambda x: x.rmsnorm((1,2), eps=1e-3), + atol=1e-6, grad_atol=1e-6) def test_normalize_int(self): helper_test_op(None, lambda x: torch.nn.functional.normalize(x.float(), p=2), lambda x: x.normalize(p=2), forward_only=True, vals=[[[3, 4], [6, 8]]]) @@ -1790,6 +1805,7 @@ class TestOps(unittest.TestCase): helper_test_op([(6,6,6)], lambda x: torch.logsumexp(x, dim=2), lambda x: x.logsumexp(2), atol=1e-7, grad_atol=1e-7) helper_test_op([(6,6,6,6)], lambda x: torch.logsumexp(x, dim=2), lambda x: x.logsumexp(2), atol=1e-7, grad_atol=1e-7) helper_test_op([(6,6,6,6)], lambda x: torch.logsumexp(x, dim=3), lambda x: x.logsumexp(3), atol=1e-7, grad_atol=1e-7) + helper_test_op([(6,6,6,6)], lambda x: torch.logsumexp(x, dim=(1,3)), lambda x: x.logsumexp((1,3)), atol=1e-7, grad_atol=1e-7) helper_test_op([(45)], lambda x: torch.logsumexp(x, dim=0), lambda x: x.logsumexp(0), atol=1e-7, grad_atol=1e-7) helper_test_op([()], lambda x: torch.logsumexp(x, dim=0), lambda x: x.logsumexp(0), atol=1e-7, grad_atol=1e-7) helper_test_op([()], lambda x: torch.logsumexp(x, dim=-1), lambda x: x.logsumexp(-1), atol=1e-7, grad_atol=1e-7) @@ -2145,6 +2161,9 @@ class TestOps(unittest.TestCase): helper_test_op([(4,3,6,6)], lambda x: x.squeeze(3)) self.helper_test_exception([(4,3,6,6)], lambda x: x.squeeze(50), expected=IndexError) helper_test_op([(4,3,6,1)], lambda x: x.squeeze(-1)) + helper_test_op([(1,3,1,6)], lambda x: x.squeeze((0,2))) + helper_test_op([(1,3,1,6)], lambda x: x.squeeze((-2,1))) + helper_test_op([(1,3,1,6)], lambda x: x.squeeze(())) helper_test_op([(4,3,6,6)], lambda x: x.squeeze()) helper_test_op([(1,3,6,6)], lambda x: x.squeeze()) helper_test_op([(2,3,1)], lambda x: x.squeeze()) diff --git a/tinygrad_repo/test/backend/test_transcendental.py b/tinygrad_repo/test/backend/test_transcendental.py index 20f008154..68788ae1f 100644 --- a/tinygrad_repo/test/backend/test_transcendental.py +++ b/tinygrad_repo/test/backend/test_transcendental.py @@ -27,27 +27,29 @@ class TestTranscendentalMath(unittest.TestCase): atol=3e-2, rtol=1e-5) # sin can have bigger atol for very big x @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "crashed") - @given(ht.float32, strat.sampled_from([(Tensor.exp, np.exp),(Tensor.log, np.log)] + - ([(Tensor.sin, np.sin)] if dtypes.ulong in supported_dtypes else []))) + @given(ht.float32, strat.sampled_from([(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.sin)])) def test_float32(self, x, op): - # wrong nan behavior on Vulkan - if (math.isnan(x) or (x < 0 and op[0] == Tensor.log)) and Device.DEFAULT == "WEBGPU" and not OSX: return with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'): np.testing.assert_allclose(op[0](Tensor([x], dtype=dtypes.float32)).numpy(), op[1](np.array([x], dtype=_to_np_dtype(dtypes.float32))), atol=2e-5, rtol=1e-5) @unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}") - @given(ht.float16, strat.sampled_from([(Tensor.exp, np.exp),(Tensor.log, np.log)] + - ([(Tensor.sin, np.sin)] if dtypes.ulong in supported_dtypes else []))) + @given(ht.float16, strat.sampled_from([(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.sin)])) def test_float16(self, x, op): - # wrong nan behavior on Vulkan - if (math.isnan(x) or (x < 0 and op[0] == Tensor.log)) and Device.DEFAULT == "WEBGPU" and not OSX: return with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'): np.testing.assert_allclose(op[0](Tensor([x], dtype=dtypes.float16)).numpy(), op[1](np.array([x], dtype=_to_np_dtype(dtypes.float16))), atol=1e-2, rtol=5e-3) # exp can have bigger rtol + @Context(TRANSCENDENTAL=2) + def test_log2_special_values(self): + for dtype, np_dtype in ((dtypes.float16, np.float16), (dtypes.float32, np.float32)): + if dtype not in supported_dtypes: continue + with self.subTest(dtype=dtype), np.errstate(all='ignore'): + values = np.array([0.0, -0.0, math.inf, -math.inf, math.nan, -math.nan, 1.0, -1.0], dtype=np_dtype) + np.testing.assert_allclose(Tensor(values).log2().numpy(), np.log2(values)) + # TODO: WEBGPU produces incorrect values near infinity @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU incorrect values near inf") @given(strat.sampled_from([(dtypes.float64, 709.5), (dtypes.float32, 88.7), (dtypes.float16, 11)])) @@ -62,7 +64,6 @@ class TestTranscendentalMath(unittest.TestCase): class TestFromFuzzer(unittest.TestCase): @given(strat.sampled_from(dtypes_float)) - @unittest.skipUnless(dtypes.ulong in supported_dtypes, "Needs ulong") def test_sin(self, dtype): if dtype not in supported_dtypes: return if dtype == dtypes.float64: @@ -116,7 +117,6 @@ class TestFloat16Log2(unittest.TestCase): np.testing.assert_allclose(result, expected, rtol=1e-3, err_msg=f"log2({val})") @unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}") - @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "Nan handling differs on Vulkan") def test_float16_log2_special(self): # special values: inf, -inf, nan, 0, negative with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'): diff --git a/tinygrad_repo/test/external/external_test_example.py b/tinygrad_repo/test/external/external_test_example.py index 37554fc22..4ce98740d 100644 --- a/tinygrad_repo/test/external/external_test_example.py +++ b/tinygrad_repo/test/external/external_test_example.py @@ -8,7 +8,7 @@ def multidevice_test(fxn): def ret(self): for device in Device._devices: # broken on OSX USB AMD, why? - if device in ["DISK", "NPY", "FAKE", "DSP", "NULL"] or (OSX and device in ["AMD"]): continue + if device in ["DISK", "NPY", "FAKE", "DSP", "NULL", "RDMA"] or (OSX and device in ["AMD"]): continue if sys.stdout.isatty(): print(device) if device in exclude_devices: if sys.stdout.isatty(): print(f"WARNING: {device} test is excluded") diff --git a/tinygrad_repo/test/external/process_replay/process_replay.py b/tinygrad_repo/test/external/process_replay/process_replay.py index 2c056aa58..c8bed9c39 100755 --- a/tinygrad_repo/test/external/process_replay/process_replay.py +++ b/tinygrad_repo/test/external/process_replay/process_replay.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 # compare kernels created by HEAD against master import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, functools, base64, codecs +from concurrent.futures import ProcessPoolExecutor from dataclasses import replace from typing import Callable, Any @@ -107,12 +108,9 @@ def _pmap(fxns:dict[str, Callable]) -> None: finally: cur.close() - with multiprocessing.get_context("spawn").Pool(multiprocessing.cpu_count()) as pool: - bar = tqdm(total=row_count) - for _ in pool.imap_unordered(functools.partial(diff, fxns=fxns), range(0, row_count, s:=min(PAGE_SIZE, row_count))): bar.update(s) - pool.close() - pool.join() - pool.terminate() + with ProcessPoolExecutor(multiprocessing.cpu_count(), mp_context=multiprocessing.get_context("spawn")) as pool: + with tqdm(total=row_count) as bar: + for _ in pool.map(functools.partial(diff, fxns=fxns), range(0, row_count, s:=min(PAGE_SIZE, row_count))): bar.update(min(s, row_count-bar.n)) # *** main loop diff --git a/tinygrad_repo/test/mockgpu/amd/amddriver.py b/tinygrad_repo/test/mockgpu/amd/amddriver.py index 0ec6d9814..edd1ef35f 100644 --- a/tinygrad_repo/test/mockgpu/amd/amddriver.py +++ b/tinygrad_repo/test/mockgpu/amd/amddriver.py @@ -128,7 +128,8 @@ class AMDDriver(VirtDriver): if struct.flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED: self.track_address(struct.va_addr, struct.va_addr + struct.size, lambda mv,off: None, lambda mv, off: self._emulate_execute()) elif nr == _ioctl_nr(kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU): - self.object_by_handle.pop(struct.handle) + va = self.object_by_handle.pop(struct.handle).va_addr + self.tracked_addresses = [x for x in self.tracked_addresses if x[0] != va] elif nr == _ioctl_nr(kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU): dev_ids = (ctypes.c_int32 * struct.n_devices).from_address(struct.device_ids_array_ptr) for i in range(struct.n_devices): diff --git a/tinygrad_repo/test/mockgpu/amd/amdgpu.py b/tinygrad_repo/test/mockgpu/amd/amdgpu.py index 820c595a2..c12cde835 100644 --- a/tinygrad_repo/test/mockgpu/amd/amdgpu.py +++ b/tinygrad_repo/test/mockgpu/amd/amdgpu.py @@ -338,7 +338,7 @@ class AMDGPURegisters: self.regs: dict[tuple[int, int], int] = {} def __getitem__(self, addr:int) -> int: if addr == regGRBM_GFX_INDEX: return self.grbm_index - if regCPG_PERFCOUNTER1_LO < addr < regGUS_PERFCOUNTER_HI: + if regCPG_PERFCOUNTER1_LO <= addr <= regGUS_PERFCOUNTER_HI: assert self.regs[(regCP_PERFMON_CNTL, 0)] == 0x401, "read mode should be enabled" return addr << 16 | self.grbm_index return self.regs[(addr, getbits(self.grbm_index, 16, 23))] diff --git a/tinygrad_repo/test/mockgpu/amd/pcode.py b/tinygrad_repo/test/mockgpu/amd/pcode.py index 8d49b6b42..f3f454cf1 100644 --- a/tinygrad_repo/test/mockgpu/amd/pcode.py +++ b/tinygrad_repo/test/mockgpu/amd/pcode.py @@ -704,6 +704,9 @@ class Parser: dt_name = self.eat('IDENT').val return result.cast(DTYPES.get(dt_name, dtypes.uint32)) return result + if field in ('i24', 'u24'): + val = base.cast(dtypes.uint32) & _u32(0xFFFFFF) + return _signext_from_bit(val, _u32(24)).bitcast(dtypes.int32) if field == 'i24' else val dt = DTYPES.get(field) if dt is None: return base if dt == base.dtype: return base diff --git a/tinygrad_repo/test/mockgpu/nv/nvgpu.py b/tinygrad_repo/test/mockgpu/nv/nvgpu.py index 84893f581..c9c413218 100644 --- a/tinygrad_repo/test/mockgpu/nv/nvgpu.py +++ b/tinygrad_repo/test/mockgpu/nv/nvgpu.py @@ -57,14 +57,14 @@ class GPFIFO: self.buf_sz = sz // 4 def execute(self) -> bool: - initial_off = self.buf_ptr + initial_get, initial_off = self.ctrl.GPGet, self.buf_ptr while self.ctrl.GPGet != self.ctrl.GPPut: self._set_buf_state(self.gpfifo[self.ctrl.GPGet]) if not self.execute_buf(): # Buffer isn't executed fully, check if any progress and report. # Do not move GPGet in this case, will continue from the same state next time. - return self.buf_ptr != initial_off + return self.ctrl.GPGet != initial_get or self.buf_ptr != initial_off self.ctrl.GPGet = (self.ctrl.GPGet + 1) % self.entries_cnt self._reset_buf_state() diff --git a/tinygrad_repo/test/null/test_dtype.py b/tinygrad_repo/test/null/test_dtype.py index 35c157b5f..b6c17ba8e 100644 --- a/tinygrad_repo/test/null/test_dtype.py +++ b/tinygrad_repo/test/null/test_dtype.py @@ -56,7 +56,11 @@ class TestCanLosslessCast(unittest.TestCase): self.assertFalse(can_lossless_cast(dtypes.int64, dtypes.double)) # half has more mantissa bits self.assertTrue(can_lossless_cast(dtypes.int8, dtypes.half)) - self.assertFalse(can_lossless_cast(dtypes.int8, dtypes.bfloat16)) + self.assertTrue(can_lossless_cast(dtypes.int8, dtypes.bfloat16)) + self.assertTrue(can_lossless_cast(dtypes.uint8, dtypes.bfloat16)) + self.assertFalse(can_lossless_cast(dtypes.int16, dtypes.bfloat16)) + self.assertTrue(can_lossless_cast(dtypes.fp8e4m3fnuz, dtypes.bfloat16)) + self.assertFalse(can_lossless_cast(dtypes.fp8e4m3, dtypes.fp8e4m3fnuz)) class TestInvalidSingleton(unittest.TestCase): def test_singleton(self): diff --git a/tinygrad_repo/test/null/test_simplify_valid_idx.py b/tinygrad_repo/test/null/test_simplify_valid_idx.py index 05f831fb8..5adabc645 100644 --- a/tinygrad_repo/test/null/test_simplify_valid_idx.py +++ b/tinygrad_repo/test/null/test_simplify_valid_idx.py @@ -236,6 +236,8 @@ class TestImageSimplification(unittest.TestCase): # same thing, valid has a div load = get_load_image_uop((10, 10, 4), gidx1//2<5, (gidx0, gidx1)) self.check(load, None, "gidx0", "gidx1") + load = get_load_image_uop((10, 10, 4), gidx1<=9, (gidx0, gidx1)) + self.check(load, None, "gidx0", "gidx1") # 10x20 image, not out of bound load = get_load_image_uop((20, 10, 4), gidx1<10, (gidx0, gidx1)) diff --git a/tinygrad_repo/test/null/test_uop_graph.py b/tinygrad_repo/test/null/test_uop_graph.py index 577e39afe..7d6846d22 100644 --- a/tinygrad_repo/test/null/test_uop_graph.py +++ b/tinygrad_repo/test/null/test_uop_graph.py @@ -3,6 +3,7 @@ from tinygrad import dtypes, Variable, Device from tinygrad.dtype import AddrSpace from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes, KernelInfo from tinygrad.uop.symbolic import sym +from tinygrad.codegen.simplify import pm_reduce_unparented from test.helpers import full_rewrite, to_uops_list from tinygrad.codegen import full_rewrite_to_sink @@ -423,6 +424,15 @@ class TestReduceCollapse(unittest.TestCase): # Should become add of two separate reduces self.assertEqual(result.op, Ops.ADD) + def test_bool_reduce_unparented_keeps_dtype(self): + for op in (Ops.ADD, Ops.MUL): + self.assertIs(graph_rewrite(UOp.const(True).reduce(UOp.range(4, 0), arg=op), pm_reduce_unparented).dtype, dtypes.bool) + + def test_max_reduce_unparented_empty_range(self): + n = Variable("n", 0, 4) + ret = graph_rewrite(UOp.const(1.0).reduce(UOp.range(n, 0), arg=Ops.MAX), pm_reduce_unparented) + self.assertEqual(ret.substitute({n: n.const_like(0)}).simplify().arg, float("-inf")) + def test_reduce_shapeless_const_unroll(self): """a REDUCE over a shapeless CONST (e.g. x*0 folded late in codegen) must collapse before the expander""" out = UOp.param(0, dtypes.float, 1) diff --git a/tinygrad_repo/test/null/test_uop_vmin_vmax.py b/tinygrad_repo/test/null/test_uop_vmin_vmax.py index fcc783408..9b8a4df91 100644 --- a/tinygrad_repo/test/null/test_uop_vmin_vmax.py +++ b/tinygrad_repo/test/null/test_uop_vmin_vmax.py @@ -150,6 +150,14 @@ class TestVminVmaxProperties(unittest.TestCase): self.assertEqual(uop3.vmin, ~6) self.assertEqual(uop3.vmax, ~(-5)) + def test_vmin_vmax_xor_nonneg(self): + # non-negative operands stay below the wider operand's next power of two + self.assertEqual((UOp.variable('x', 0, 31) ^ 16)._min_max, (0, 31)) + self.assertEqual((UOp.variable('x', 3, 10) ^ UOp.variable('y', 0, 40))._min_max, (0, 63)) + self.assertEqual((UOp.variable('x', 0, 200, dtypes.uint8) ^ UOp.variable('y', 0, 3, dtypes.uint8))._min_max, (0, 255)) + # a possibly-negative operand is not bounded + self.assertEqual((UOp.variable('x', -1, 10) ^ 16)._min_max, (dtypes.weakint.min, dtypes.weakint.max)) + def test_vmin_vmax_cast(self): x = UOp.variable('x', -10, 10, dtypes.int) x_float = x.cast(dtypes.float) diff --git a/tinygrad_repo/test/null/test_uops_stats.py b/tinygrad_repo/test/null/test_uops_stats.py index 315f2a657..02ff07a55 100644 --- a/tinygrad_repo/test/null/test_uops_stats.py +++ b/tinygrad_repo/test/null/test_uops_stats.py @@ -131,6 +131,9 @@ class TestUOpsStats(unittest.TestCase): unequal = equal[:3] + [Tensor.empty(256, 129)] self.assertGreater(get_stats(Tensor.cat(*unequal, dim=1))[0], 0) + def test_xor_gather(self): + self.assertEqual(get_stats(Tensor.empty(32)[Tensor.arange(32) ^ 16])[0], 0) + def test_simple_matmul(self, M=1024, N=1024, K=1024): a = Tensor.empty(M,N) b = Tensor.empty(N,K) diff --git a/tinygrad_repo/test/unit/test_attention.py b/tinygrad_repo/test/unit/test_attention.py index 7a3947378..aa49c4adc 100644 --- a/tinygrad_repo/test/unit/test_attention.py +++ b/tinygrad_repo/test/unit/test_attention.py @@ -118,8 +118,8 @@ class TestGatedDeltaNetBlock(unittest.TestCase): x_float = x.astype(np.float32) return (x_float / np.sqrt((x_float * x_float).mean(axis=-1, keepdims=True) + eps)) * weight.astype(np.float32) - def _normalize_np(self, x:np.ndarray, eps:float=1e-6) -> np.ndarray: - return x / np.maximum(np.sqrt((x * x).sum(axis=-1, keepdims=True)), eps) + def _l2norm_np(self, x:np.ndarray, eps:float=1e-6) -> np.ndarray: # as in the FLA kernels + return x / np.sqrt((x * x).sum(axis=-1, keepdims=True) + eps) def _softplus_np(self, x:np.ndarray) -> np.ndarray: return np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0) @@ -154,8 +154,8 @@ class TestGatedDeltaNetBlock(unittest.TestCase): conv_window = np.concatenate([conv_state, self._linear_np(x_half, qkv_weight)], axis=1) conv_out = self._silu_np((conv_window * conv_weight).sum(axis=1)) q, k, v = np.split(conv_out, [block.q_dim, 2 * block.q_dim], axis=-1) - q = self._normalize_np(q.reshape(B, block.num_k_heads, block.head_k_dim)) - k = self._normalize_np(k.reshape(B, block.num_k_heads, block.head_k_dim)) + q = self._l2norm_np(q.reshape(B, block.num_k_heads, block.head_k_dim)) + k = self._l2norm_np(k.reshape(B, block.num_k_heads, block.head_k_dim)) v = v.reshape(B, block.num_v_heads, block.head_v_dim) if block.num_v_heads != block.num_k_heads: k_repeat = block.num_v_heads // block.num_k_heads @@ -214,6 +214,19 @@ class TestGatedDeltaNetBlock(unittest.TestCase): np.testing.assert_allclose(recurrent_state, expected_recurrent[step], rtol=1e-3, atol=1e-3, err_msg=f"GatedDeltaNet reset recurrent cache mismatch at step {step}") + def test_kda_qk_l2norm(self): + # q/k are normalized as x/sqrt(sum(x²)+1e-6) like FLA: tiny vectors aren't scaled up to unit length + ssm = SSMConfig(conv_kernel=4, state_size=8, group_count=2, time_step_rank=2, inner_size=16, kda=True) + block = GatedDeltaNetBlock(self._make_config(dim=16, n_heads=2, ssm=ssm), ssm) + block.attn_qkv.weight = Tensor([1e-4]*32 + [1.0]*16).reshape(48, 1).expand(48, 16) / 16 # x of ones: q = k = 1e-4, v = 1 + block.ssm_conv1d["weight"], block.ssm_beta.weight = Tensor.ones(48, 4), Tensor.zeros(2, 16) + x = Tensor.ones(1, 1, 16) + block._init_state(x) + block._attention(x, 0).realize() + # the first token's state is beta * v (x) l2norm(k), with beta = sigmoid(0) and v, k through the conv and silu + def silu(z): return z / (1 + np.exp(-z)) + np.testing.assert_allclose(block.recurrent_state.numpy(), 0.5 * silu(1.0) * silu(1e-4) / np.sqrt(8 * silu(1e-4)**2 + 1e-6), rtol=1e-3) + def test_kda_channel_decay(self): config = self._make_config(dim=4, hidden_dim=8, n_heads=2, head_dim=4, rope_dim=4, v_head_dim=4, ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True)) diff --git a/tinygrad_repo/test/unit/test_dtype_weak.py b/tinygrad_repo/test/unit/test_dtype_weak.py index b8b022829..dcffd69db 100644 --- a/tinygrad_repo/test/unit/test_dtype_weak.py +++ b/tinygrad_repo/test/unit/test_dtype_weak.py @@ -130,7 +130,7 @@ class TestWeakPromotion(unittest.TestCase): self.assertEqual((out.src[1], out.src[2]), (UOp.const(5.0, dtypes.bfloat16), gate)) def test_weak_srcs_commit_only_at_a_concrete_lub(self): - weak_lub = UOp(Ops.ADD, src=(UOp.const(1), UOp.const(1.0))) + weak_lub = UOp(Ops.ADD, src=(UOp.const(2.0), UOp.const(1.0))) self.assertIs(graph_rewrite(weak_lub, pm_commit_weak), weak_lub) concrete = UOp.const(2.0).cast(dtypes.float16) # the weak arm stays bare: its sibling states the width, so the WHERE already derives float16 for it diff --git a/tinygrad_repo/tinygrad/codegen/__init__.py b/tinygrad_repo/tinygrad/codegen/__init__.py index 1e2bba733..4c7008778 100644 --- a/tinygrad_repo/tinygrad/codegen/__init__.py +++ b/tinygrad_repo/tinygrad/codegen/__init__.py @@ -269,7 +269,7 @@ pm_implicit_barriers = PatternMatcher([ ]) def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: - if DEBUG >= 5: print(render_uir(list(ast.toposort()))) + if DEBUG >= 5: print(render_uir(ast)) if SPEC: type_verify(ast, spec_tensor) # resolve UNSHARDs (multi-device UNSHARDs are already resolved by the scheduler; this handles in-kernel shards, e.g. fragments) @@ -350,11 +350,12 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: sink = graph_rewrite(sink, pm_decomp, name="early decompositions") # late decomps + move gates from unrenderable INVALID where - sink = graph_rewrite(sink, pm_dtype_decomps+pm_commit_weak, ctx=(set(), ren), name="decomp dtypes") pm_decomp = pm_decomp+\ get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))+\ get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2) sink = graph_rewrite(sink, pm_decomp, ctx=ren, name="late decompositions") + # Transcendental decompositions can introduce unsupported dtypes, e.g. uint64 in Payne-Hanek reduction. + sink = graph_rewrite(sink, pm_dtype_decomps+pm_commit_weak, ctx=(set(), ren), name="decomp dtypes") sink = graph_rewrite(sink, pm_move_gates_from_index, name="move gates from index") # final rules for the renderer (without sym) @@ -384,7 +385,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: if os.environ.get("DBGTV"): try: type_verify(sink, spec_program) except RuntimeError: - print(render_uir(list(sink.toposort()))) + print(render_uir(sink)) raise else: type_verify(sink, spec_program) @@ -413,20 +414,18 @@ def line_rewrite(lst:list[UOp], pm:PatternMatcher, ctx=None) -> list[UOp]: newlst.extend(ret[1]) return newlst -pm_lower_calls = PatternMatcher([ - (UPat(Ops.CALL, src=(UPat(Ops.SINK),), allow_any_len=True, name="call"), - lambda ctx,call: call.replace(src=(full_rewrite_to_sink(call.body, ctx, optimize=False),)+call.src[1:])), -]) +def lower_call(ctx:Renderer, call:UOp) -> UOp: + body = graph_rewrite(call.body, pm_lower_calls, ctx=ctx, name="lower calls", walk=True) + body = full_rewrite_to_sink(body, ctx, optimize=False) + body = UOp(Ops.LINEAR, src=tuple(line_rewrite(linearize(body), pm_linearize_cleanups+pm_alloc_to_buf)), arg=to_function_name(call.arg.name)) + return call.replace(src=(body,)+call.src[1:]) -pm_call_fixup = PatternMatcher([ - (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="sink"),), allow_any_len=True, name="call"), - lambda call,sink: call.replace(src=(UOp(Ops.LINEAR, src=tuple(line_rewrite(linearize(sink), pm_linearize_cleanups+pm_alloc_to_buf)), - arg=to_function_name(call.arg.name)),)+call.src[1:])), +pm_lower_calls = PatternMatcher([ + (UPat(Ops.CALL, src=(UPat(Ops.SINK),), allow_any_len=True, name="call"), lower_call), ]) def do_linearize(ctx:Renderer, prg:UOp, sink:UOp) -> UOp: if DEBUG >= 3 and sink.arg.applied_opts: print(f"{sink.arg.function_name:<25} opts: {sink.arg.applied_opts}") - sink = graph_rewrite(sink, pm_call_fixup, name="call fixup", enter_calls=True) lst = line_rewrite(linearize(sink), pm_linearize_cleanups+pm_alloc_to_buf) prg = prg.replace(src=(lst[-1],)) # isa renderers need to allocate registers @@ -469,7 +468,7 @@ pm_to_program = PatternMatcher([ (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile), ]) -@rewrite_group(name=lambda ast,renderer,ret,**_: TracingKey((k:=ret.src[0].arg).name,(k.function_name, ast, ret.key),ret=renderer), replay=True) +@rewrite_group(name=lambda ast,renderer,ret,**_: TracingKey((k:=ret.src[0].arg).name,(k.function_name, ret.key),ret=renderer), replay=True) @Context(ALLOW_DEVICE_USAGE=0) def do_to_program(ast:UOp, renderer:Renderer) -> UOp: """ @@ -486,7 +485,7 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp: elif ast.op is Ops.SINK: assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to to_program" if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST") - ast = graph_rewrite(ast, pm_lower_calls, ctx=renderer, name="lower calls", walk=True, enter_calls=True) + ast = graph_rewrite(ast, pm_lower_calls, ctx=renderer, name="lower calls", walk=True) full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None) prog_info = ProgramInfo.from_sink(full_sink, renderer.target) # instruction selection diff --git a/tinygrad_repo/tinygrad/codegen/decomp/dtype.py b/tinygrad_repo/tinygrad/codegen/decomp/dtype.py index 20644eaf7..548a9697e 100644 --- a/tinygrad_repo/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad_repo/tinygrad/codegen/decomp/dtype.py @@ -2,7 +2,7 @@ from dataclasses import replace from tinygrad.dtype import dtypes, DType, AddrSpace, truncate from tinygrad.helpers import flatten, DEBUG, EMULATED_DTYPES from tinygrad.uop import GroupOp -from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite +from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite, shape_to_shape_arg from tinygrad.uop.weak import commit_weak_consts from tinygrad.renderer import Renderer from tinygrad.codegen.decomp.transcendental import exponent_bias, shl, shr @@ -83,7 +83,8 @@ def l2i(op: Ops, dt: DType, *uops:UOp): def l2i_define(x:UOp) -> UOp: # cannot decomp a Variable if x.addrspace == AddrSpace.ALU: raise RuntimeError(f"long decomposition of variable {x.arg.name} unsupported") - return x.replace(arg=replace(x.arg, dtype=l2i_dt[x.dtype]), src=(UOp.const(x.src[0].val*2) if x.shape else x.src[0], *x.src[1:])) + size = shape_to_shape_arg((x.max_numel()*2,)) if x.shape else x.src[0] + return x.replace(arg=replace(x.arg, dtype=l2i_dt[x.dtype]), src=(size, *x.src[1:])) def split_l2i(ctx:dict, op: Ops, dt: DType, *uops:UOp): # l2i does arithmetic on its inputs; rules enter here to split them to 32-bit words first, l2i recurses on itself. diff --git a/tinygrad_repo/tinygrad/codegen/decomp/transcendental.py b/tinygrad_repo/tinygrad/codegen/decomp/transcendental.py index 0cf4c6102..70841c562 100644 --- a/tinygrad_repo/tinygrad/codegen/decomp/transcendental.py +++ b/tinygrad_repo/tinygrad/codegen/decomp/transcendental.py @@ -249,10 +249,10 @@ def xlog2(d:UOp) -> UOp: r = d.ne(0.0).where(r, r.const_like(-math.inf)) # log2(x) = NaN for x < 0 r = (d<-0.0).where(r.const_like(math.nan), r) - # log2(NaN) = NaN - r = d.ne(d).where(r.const_like(math.nan), r) # log2(-0.0) = -Inf. In certain devices like PTX, x == -0.0 won't be true. so making reciprocal. - return d.reciprocal().ne(-math.inf).where(r, r.const_like(-math.inf)) + r = d.reciprocal().ne(-math.inf).where(r, r.const_like(-math.inf)) + # Check NaN last: comparisons of its reciprocal are not reliable on every backend. + return d.ne(d).where(r.const_like(math.nan), r) def xpow(base:UOp, exponent:UOp) -> UOp: # start with b ** e = exp2(e * log2(b)) diff --git a/tinygrad_repo/tinygrad/codegen/late/coalesce.py b/tinygrad_repo/tinygrad/codegen/late/coalesce.py index ec57edd4e..7ec5fa235 100644 --- a/tinygrad_repo/tinygrad/codegen/late/coalesce.py +++ b/tinygrad_repo/tinygrad/codegen/late/coalesce.py @@ -111,7 +111,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: # collect memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp, object], dict[int, list[UOp]]] = defaultdict(dict) - for u in sink.toposort(enter_calls=False): + for u in sink.toposort(): # TODO: this should handle images too, it's just memory coalescing if u.op in {Ops.LOAD, Ops.STORE}: assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores" diff --git a/tinygrad_repo/tinygrad/codegen/late/linearizer.py b/tinygrad_repo/tinygrad/codegen/late/linearizer.py index cdff4a64d..0838d70c9 100644 --- a/tinygrad_repo/tinygrad/codegen/late/linearizer.py +++ b/tinygrad_repo/tinygrad/codegen/late/linearizer.py @@ -7,9 +7,9 @@ from tinygrad.helpers import prod, getenv, dedup, TUPLE_ORDER def linearize(sink:UOp) -> list[UOp]: # this is a toposort with priority - lst = list(sink.toposort(enter_calls=False)) + lst = list(sink.toposort()) out_degree:defaultdict[UOp, int] = defaultdict(int) - priorities:dict[UOp, tuple[int, int, Any]] = {} + priorities:dict[UOp, tuple[int, int, int, Any]] = {} # get consumers and assign priorities # NOTE: this requires the lst be locally toposorted @@ -30,7 +30,7 @@ def linearize(sink:UOp) -> list[UOp]: case Ops.RANGE: priority = 5 # placing RANGE is good case Ops.END | Ops.BACKEDGE: priority = -5 # placing loop exits is bad case _: priority = 0 # everything else has priority 0 - priorities[u] = (run_count, priority, extra) + priorities[u] = (run_count, len(u.ranges), priority, extra) # number the uops in "ideal" order nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if TUPLE_ORDER else ())))} @@ -62,7 +62,7 @@ class CFGContext: for u in sink.toposort(): # get the deps from the src deps[u] = {} - for s in u.src: deps[u] |= deps[s] + for s in u.src_without_body: deps[u] |= deps[s] if u.op in (Ops.END, Ops.BACKEDGE, Ops.SINK): nesting |= {x:u for x in deps[u] if x.op in (Ops.END, Ops.BACKEDGE) and (u.op is Ops.SINK or u.src[1] in deps[x]) and x not in nesting} diff --git a/tinygrad_repo/tinygrad/codegen/simplify.py b/tinygrad_repo/tinygrad/codegen/simplify.py index c9878b8fc..46caa869e 100644 --- a/tinygrad_repo/tinygrad/codegen/simplify.py +++ b/tinygrad_repo/tinygrad/codegen/simplify.py @@ -1,6 +1,6 @@ import itertools from typing import Callable -from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, AxisType +from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, AxisType, identity_element from tinygrad.uop.symbolic import symbolic from tinygrad.helpers import partition, dedup from tinygrad.dtype import dtypes @@ -84,10 +84,10 @@ def reduce_unparented(red:UOp) -> UOp|None: reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].ranges) if len(reduce_unparented) == 0: return None ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) else red.src[0] - if red.arg[0] is Ops.ADD: - for r in reduce_unparented: ret = ret * r.src[0] - if red.arg[0] is Ops.MUL: - for r in reduce_unparented: ret = ret ** r.src[0] + for r in reduce_unparented: + # an idempotent op (MAX, and bool ADD/MUL which are OR/AND) applied N times is the op once, or the identity when N is 0 + if red.arg[0] is Ops.MAX or red.dtype is dtypes.bool: ret = (r.src[0] > 0).where(ret, ret.const_like(identity_element(red.arg[0], red.dtype))) + else: ret = ret * r.src[0] if red.arg[0] is Ops.ADD else ret ** r.src[0] return ret pm_reduce_unparented = PatternMatcher([ diff --git a/tinygrad_repo/tinygrad/device.py b/tinygrad_repo/tinygrad/device.py index 36053a04c..a4a763251 100644 --- a/tinygrad_repo/tinygrad/device.py +++ b/tinygrad_repo/tinygrad/device.py @@ -2,11 +2,11 @@ from __future__ import annotations from dataclasses import dataclass, replace, field from collections import defaultdict from typing import Any, Callable, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING -import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal, subprocess, struct, mmap, time, statistics +import importlib, inspect, functools, itertools, pathlib, os, contextlib, re, atexit, pickle, decimal, subprocess, struct, mmap, time, statistics from tinygrad.helpers import mv_address, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up, is_numpy_ndarray -from tinygrad.helpers import cpu_profile, perf_counter_us, to_name, HCQ_RUNTIME_DEV +from tinygrad.helpers import cpu_profile, perf_counter_us, to_name, HCQ_RUNTIME_DEV, to_tuple from tinygrad.dtype import dtypes, DType from tinygrad.runtime.support.memory import BumpAllocator, MMIOInterface if TYPE_CHECKING: @@ -64,7 +64,7 @@ def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]: return canonical[0] if len(canonical:=tuple(Device.canonicalize(d) for d in device)) == 1 else canonical def is_disk_device(device:str|tuple[str, ...]) -> bool: - return any(d.split(":", 1)[0].upper() == "DISK" for d in ((device,) if isinstance(device, str) else device)) + return any(d.split(":", 1)[0].upper() == "DISK" for d in to_tuple(device)) # **************** Profile **************** @@ -108,7 +108,7 @@ class Buffer: def __init__(self, device:str, nbytes:int, opaque:Any=None, options:BufferSpec|None=None, initial_value:bytes|pickle.PickleBuffer|None=None, base:Buffer|None=None, offset:int=0, preallocate=False, allocator:Allocator|None=None): - self.device, self.nbytes, self.offset, self.allocated_views, self._base = Device.canonicalize(device), nbytes, offset, 0, base + self.device, self.nbytes, self.offset, self._base = Device.canonicalize(device), nbytes, offset, base if allocator is not None: self.allocator = allocator self.options = options if options is not None else BufferSpec() self._storage:BufferStorage|None = None @@ -185,7 +185,6 @@ class Buffer: GlobalCounters.mem_used += self.nbytes GlobalCounters.mem_used_per_device[self.device] += self.nbytes if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"nbytes":self.nbytes})) - elif self._storage is None: self.base.allocated_views += 1 self._storage, self._base_storage = storage, self.base._storage if self._base else None return self @@ -198,7 +197,6 @@ class Buffer: GlobalCounters.mem_used_per_device[self.device] -= self.nbytes if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self.trace_num)) self.allocator.free(self._storage, self.nbytes, self.options) - else: self.base.allocated_views -= 1 self._storage, self._base_storage = None, None def __reduce_ex__(self, protocol): @@ -250,11 +248,10 @@ DeviceType = TypeVar('DeviceType', bound='Compiled') class Allocator(Generic[DeviceType]): lru = True - def __init__(self, dev:DeviceType, supports_copy_from_disk:bool=True, supports_transfer:bool=True): + def __init__(self, dev:DeviceType): self.dev: DeviceType = dev self.default_buffer_spec: BufferSpec = BufferSpec() self.cache:dict[tuple[int, BufferSpec|None], list[BufferStorage]] = defaultdict(list) - self.supports_copy_from_disk, self.supports_transfer = supports_copy_from_disk, supports_transfer def alloc(self, size:int, options:BufferSpec|None=None) -> BufferStorage: assert size > 0, f"alloc size must be positive, getting {size}" @@ -291,10 +288,8 @@ class Allocator(Generic[DeviceType]): def _map(self, buf) -> BufferStorage: raise NotImplementedError("need map") def _unmap(self, mb): pass # default no-op; override if _map allocates iface-side state def _offset(self, buf, size:int, offset:int): raise NotImplementedError("need offset") - # def _transfer(self, dest, src, sz:int, src_dev, dest_dev): class HostAllocator(Allocator): - def __init__(self, dev): super().__init__(dev, supports_copy_from_disk=False, supports_transfer=False) def _alloc(self, size:int, options:BufferSpec) -> BufferStorage: if options.external_ptr is not None: view, meta = self._view(options.external_ptr, size), None elif (remote:=getattr(self.dev, "remote", None)) is not None: view, meta = remote.alloc_sysmem(round_up(size, mmap.PAGESIZE)) @@ -395,7 +390,7 @@ class Compiled: def has_copy_queue(self) -> bool: return True @property - def host(self) -> str: return f"CPU:{self.peer_group[7:]}" if self.peer_group.startswith("remote:") else HCQ_RUNTIME_DEV.value + def host(self) -> str: return f"CPU:{self.peer_group[7:]}" if self.peer_group.startswith("remote:") else HCQ_RUNTIME_DEV.device @property def renderer(self) -> Renderer: return self._select_renderer() @@ -463,10 +458,10 @@ class Compiled: def _renderer_name(self, r:type[Renderer]) -> str: return r.__name__.upper().removesuffix("RENDERER").removeprefix(devname:=self.device.split(':')[0].upper()) or devname - def _select_renderer(self) -> Renderer: + def _select_renderer(self, cv=DEV) -> Renderer: assert (rn:=next((self._renderer_name(r) for r in self.renderers if getenv(f"{self.device}_{self._renderer_name(r)}")), None)) is None, \ f"{self.device}_{rn}=1 is deprecated, use DEV={self.device}:{rn} instead" - t = DEV.target(self.device.split(':')[0], **({"arch":self.arch} if self.arch else {})) + t = cv.target(self.device.split(':')[0], **({"arch":self.arch} if self.arch else {})) return select_first_inited(select_by_name(self.renderers, self._renderer_name, t.renderer, f"{self.device} has no renderer {t.renderer!r}"), f"No renderer for {self.device} is available", self.cached_renderer, t) @@ -483,10 +478,9 @@ class Compiled: # profiling def collect_prof(self): - if PROFILE: - es = list(self.prof_ents.items()) - sigs = [buf.host.view(fmt='Q')[i]/decimal.Decimal(self.timestamp_divider) for (buf, _), e in es for i in (e.st_id, e.en_id)] - Compiled.profile_events.append(ProfileGraphEvent([replace(e, st_id=2*i, en_id=2*i+1) for i,(_, e) in enumerate(es)], [], sigs)) + div = decimal.Decimal(self.timestamp_divider) + for buf, g in itertools.groupby(self.prof_ents.items(), key=lambda x: x[0][0]) if PROFILE else []: + Compiled.profile_events.append(ProfileGraphEvent([e for _, e in g], [], [t/div for t in buf.host.view(fmt='Q')[:]])) self.prof_ents.clear() def _at_profile_finalize(self): diff --git a/tinygrad_repo/tinygrad/dtype.py b/tinygrad_repo/tinygrad/dtype.py index 27a67a9ff..1e5dd2cf7 100644 --- a/tinygrad_repo/tinygrad/dtype.py +++ b/tinygrad_repo/tinygrad/dtype.py @@ -202,19 +202,13 @@ def can_lossless_cast(dt0:DType, dt1:DType) -> bool: # return if dt1 preserves value of dt0 # similar to https://numpy.org/doc/stable/reference/generated/numpy.can_cast.html if dt0 == dt1 or dt0 == dtypes.bool: return True - match dt1: - case dtypes.weakint: return dt0 in dtypes.ints - case dtypes.double: return dt0 in (dtypes.float, dtypes.half, dtypes.bfloat16, *dtypes.fp8s, - dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8) - case dtypes.float: return dt0 in (dtypes.half, dtypes.bfloat16, *dtypes.fp8s, dtypes.uint16, dtypes.uint8, dtypes.int16, dtypes.int8) - case dtypes.half: return dt0 in (*dtypes.fp8s, dtypes.uint8, dtypes.int8) - case dtypes.uint64: return dt0 in (dtypes.uint32, dtypes.uint16, dtypes.uint8) - case dtypes.uint32: return dt0 in (dtypes.uint16, dtypes.uint8) - case dtypes.uint16: return dt0 in (dtypes.uint8,) - case dtypes.int64: return dt0 in (dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8) - case dtypes.int32: return dt0 in (dtypes.uint16, dtypes.uint8, dtypes.int16, dtypes.int8) - case dtypes.int16: return dt0 in (dtypes.uint8, dtypes.int8) - case _: return False + if dtypes.is_int(dt1): return dtypes.is_int(dt0) and dt1.min <= dt0.min and dt0.max <= dt1.max + if dt1 not in dtypes.floats: return False + e1, m1 = dtypes.finfo(dt1) + if dtypes.is_int(dt0): return dt0.bitsize - (dt0 in dtypes.sints) <= m1 + 1 + if dt0 not in dtypes.floats: return False + e0, m0 = dtypes.finfo(dt0) + return e0 <= e1 and m0 <= m1 and (e0, m0) != (e1, m1) def sum_acc_dtype(dt:DType): # default acc dtype for sum diff --git a/tinygrad_repo/tinygrad/engine/jit.py b/tinygrad_repo/tinygrad/engine/jit.py index 9d6b68c32..271d8b61c 100644 --- a/tinygrad_repo/tinygrad/engine/jit.py +++ b/tinygrad_repo/tinygrad/engine/jit.py @@ -61,7 +61,7 @@ class CapturedJit(Generic[ReturnType]): def __reduce__(self): return self.__class__, (self.ret, self._linear, self.expected_names, self.expected_input_info) @functools.cached_property - def _jit_bufs(self) -> set[UOp]: return {u for u in self._linear.toposort(enter_calls=False) if u.op is Ops.BUFFER} + def _jit_bufs(self) -> set[UOp]: return {u for u in self._linear.toposort() if u.op is Ops.BUFFER} @functools.cached_property def _symbolic_ret(self) -> list[tuple[Tensor, UOp, dict[Variable, int]]]: diff --git a/tinygrad_repo/tinygrad/engine/realize.py b/tinygrad_repo/tinygrad/engine/realize.py index 7bfc9df2a..d03c7be3a 100644 --- a/tinygrad_repo/tinygrad/engine/realize.py +++ b/tinygrad_repo/tinygrad/engine/realize.py @@ -2,10 +2,10 @@ from __future__ import annotations from typing import cast, Iterator, Any, Sequence import decimal from dataclasses import dataclass, replace, field -from tinygrad.helpers import CAPTURE_PROCESS_REPLAY, colored, DEBUG, GlobalCounters, ansipad, prod, flatten, Context, to_tuple, tqdm +from tinygrad.helpers import CAPTURE_PROCESS_REPLAY, colored, DEBUG, GlobalCounters, ansipad, prod, flatten, Context, to_tuple, tqdm, unwrap from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us, cpu_profile from tinygrad.uop.ops import get_process_replay_loc, Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite, ProgramInfo, KernelInfo -from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry +from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry, DEV, HCQ_RUNTIME_DEV from tinygrad.renderer import Estimates, Renderer from tinygrad.codegen import to_program, to_program_cache, to_program_key, to_program_context from tinygrad.engine.worker import get_worker_pool, terminate_worker_pool @@ -32,7 +32,7 @@ def get_call_kernels(call:UOp) -> list[tuple[str, UOp, tuple|None]]: return kernels + [(d, call, (name, estimates, key, bufs, io)) for devices,name,estimates,_,key,bufs,io in call.arg.aux.kernels for d in devices] ast = call.body if ast.op is Ops.CUSTOM_FUNCTION and ast.arg.name == "validate": return [] - return [(d, call, None) for d in to_tuple(call.src[1].device)] + return [(d, call, None) for d in to_tuple(call.device)] def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|None=None) -> str: def _uop_sz_to_str(uop:UOp) -> str: return size_to_str(sym_infer(prod(uop.shape) * uop.dtype.itemsize, var_vals or {})) @@ -130,8 +130,10 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d if not any(isinstance(b, MultiBuffer) for b in bufs): yield cast(list[Buffer], bufs), {} else: # the DEVICE axis is bound per device at launch: it's a RANGE in the AST and the _device_num variable after codegen - has_dnum = any((x.op is Ops.RANGE and x.axis_type is AxisType.DEVICE) or (x.op is Ops.PARAM and x.arg.name == '_device_num') - for x in call.body.toposort()) + # on archs like x86, stack args (after the 6th) aren't PARAMs in the body, so also check the PROGRAM's vars + has_dnum = ((isinstance(call.body.arg, ProgramInfo) and any(v.arg.name == '_device_num' for v in call.body.arg.vars)) or + any((x.op is Ops.RANGE and x.axis_type is AxisType.DEVICE) or (x.op is Ops.PARAM and x.arg.name == '_device_num') + for x in call.body.toposort())) lanes = max(len(b.bufs) for b in bufs if isinstance(b, MultiBuffer)) # a single buffer is shared by every lane per_lane = [b.bufs if isinstance(b, MultiBuffer) else (b,)*lanes for b in bufs] for j, per_dev in enumerate(zip(*per_lane)): yield list(per_dev), {"_device_num": j} if has_dnum else {} @@ -139,12 +141,7 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]: for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)): dest, src = bufs[0].ensure_allocated(), bufs[1].ensure_allocated() - if hasattr(dest.allocator,'_transfer') and dest.allocator.supports_transfer and dest.device.split(":")[0] == src.device.split(":")[0]: - dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev) - elif src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None \ - and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk: - dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes) - elif dest.get_storage().host is not None and src.get_storage().host is not None: + if dest.get_storage().host is not None and src.get_storage().host is not None: for b in (dest, src): b.allocator.dev.synchronize() with cpu_profile(f"{src.device} -> {dest.device}", f"{src.device}:COPY"): dest.host[:] = src.host[:] elif dest._host_mv() is not None: src.allocator._copyout(dest.as_memoryview(allow_zero_copy=True), src._buf) @@ -187,12 +184,13 @@ def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]: info = call.arg.aux ctx = replace(ctx, wait=ctx.wait and not info.skip_wait, var_vals={**ctx.var_vals, **{k: v for d in info.device for k, v in cast(Any, Device[d]).var_vals.items()}}) + slots = {d: cast(Buffer, call.src[1 + i].buffer) for d, i in info.slots} + for d in {d for devs, _, _, prof, *_ in info.kernels for d in devs if prof and (slots[d], prof[0]) in Device[d].prof_ents}: Device[d].synchronize() ets = exec_kernel(ctx, call, ast, devices=(Device[info.device[0]].host,)) for d in info.device: if (err:=Device[d].error_state.host.view(fmt='q')[0]): raise RuntimeError(f"{d} failed with {err}") if not (ctx.wait or PROFILE): return ets - slots = {d: cast(Buffer, call.src[1 + i].buffer) for d, i in info.slots} for devs, name, _, prof, pkey, *_ in info.kernels: for d in (devs if prof else ()): cast(Any, Device[d]).prof_ents[(slots[d], prof[0])] = ProfileGraphEntry(d, name, prof[0], prof[1], pkey) if ctx.wait: @@ -231,12 +229,12 @@ def _get_call_to_compile(c:UOp) -> tuple[UOp, Renderer]|None: # a PROGRAM with a ProgramInfo and a BINARY is already compiled if (ast.op is Ops.SINK and isinstance(ast.arg, KernelInfo)) or \ (ast.op is Ops.PROGRAM and not (isinstance(ast.arg, ProgramInfo) and ast.src[-1].op is Ops.BINARY)): - return ast, Device[c.device if isinstance(c.device, str) else c.device[0]].renderer + return ast, Device[to_tuple(c.device)[0]]._select_renderer(HCQ_RUNTIME_DEV if isinstance(c.arg.aux, HCQInfo) else DEV) return None def lower_and_compile(linear:UOp, verbose=True) -> UOp: # collect the kernels to lower and compile, deduped by their compile cache key - if not len(ar:={c: a for c in linear.toposort(enter_calls=False) if c.op is Ops.CALL and (a:=_get_call_to_compile(c)) is not None}): return linear + if not len(ar:={c: a for c in linear.toposort() if c.op is Ops.CALL and (a:=_get_call_to_compile(c)) is not None}): return linear # lower and compile what's not cached, in parallel if there's a worker pool keys = {c: to_program_key(*a) for c, a in ar.items()} @@ -294,8 +292,8 @@ def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True, cache=False), allow_cache=ctx.cache) while True: if clear_l2: - if hasattr(dev:=Device[call.src[1].device], 'invalidate_caches'): dev.invalidate_caches() + if hasattr(dev:=Device[call.device], 'invalidate_caches'): dev.invalidate_caches() else: from tinygrad.tensor import Tensor with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False) - yield max(pm_exec.rewrite(linear.src[0].without_after, ctx) or [0.0]) + yield max(map(unwrap, pm_exec.rewrite(c:=linear.src[0].without_after, ctx)[isinstance(c.arg.aux, HCQInfo):])) diff --git a/tinygrad_repo/tinygrad/function.py b/tinygrad_repo/tinygrad/function.py index ec702cb9a..f94b69c74 100644 --- a/tinygrad_repo/tinygrad/function.py +++ b/tinygrad_repo/tinygrad/function.py @@ -21,7 +21,7 @@ pm_ctx = PatternMatcher([ (UPat(GroupOp.Defines, name="x"), lambda ctx,x: add_to_ctx(ctx,x) if is_implicit_storage(ctx, x) else None), (UPat((Ops.AFTER, Ops.STAGE), name="x"), lambda ctx,x: add_to_ctx(ctx,x) if not any(p.op is Ops.PARAM and p.arg.slot >= 0 for p in x.backward_slice) and - any(is_implicit_storage(ctx, b) for b in x.toposort(enter_calls=False)) else None), + any(is_implicit_storage(ctx, b) for b in x.toposort()) else None), ]) def invalid_outputs(uret:UOp) -> set[UOp]: @@ -33,7 +33,7 @@ def invalid_outputs(uret:UOp) -> set[UOp]: def renumber_invalid_outputs(uret:UOp) -> UOp: invalid = invalid_outputs(uret) return uret.substitute({b:b.replace(op=Ops.ALLOC, arg=replace(b.arg, slot=i, buffer=None, bind_on_realize=False)) - for i,b in enumerate(x for x in uret.toposort(enter_calls=False) if x in invalid)}) + for i,b in enumerate(x for x in uret.toposort() if x in invalid)}) ReturnType = TypeVar('ReturnType') class _function(Generic[ReturnType]): diff --git a/tinygrad_repo/tinygrad/helpers.py b/tinygrad_repo/tinygrad/helpers.py index 24788b8cd..3085bbc2e 100644 --- a/tinygrad_repo/tinygrad/helpers.py +++ b/tinygrad_repo/tinygrad/helpers.py @@ -249,7 +249,7 @@ RING, ALL2ALL, ALLREDUCE_CAST = ContextVar("RING", 1), ContextVar("ALL2ALL", 0), ALLREDUCE_NODE_NDEVS = ContextVar("ALLREDUCE_NODE_NDEVS", 0) # gpus per node, the nodes cabled gpu k to gpu k CACHELEVEL, IGNORE_BEAM_CACHE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0) VALIDATE_WITH_CPU = ContextVar("VALIDATE_WITH_CPU", 0) -HCQ_RUNTIME_DEV = ContextVar("HCQ_RUNTIME_DEV", "CPU") +HCQ_RUNTIME_DEV = _DEV("HCQ_RUNTIME_DEV", "CPU:LLVM") # TODO: this is broken for some indexing DISABLE_FAST_IDIV = ContextVar("DISABLE_FAST_IDIV", 1) FUSE_OPTIM = ContextVar("FUSE_OPTIM", 0) diff --git a/tinygrad_repo/tinygrad/llm/cli.py b/tinygrad_repo/tinygrad/llm/cli.py index b92ef47bc..2a008564e 100644 --- a/tinygrad_repo/tinygrad/llm/cli.py +++ b/tinygrad_repo/tinygrad/llm/cli.py @@ -170,7 +170,7 @@ def main(): if not args.no_chat_template and (ct := kv.get('tokenizer.chat_template')) is not None: try: import jinja2 - env = jinja2.Environment() + env = jinja2.Environment(extensions=['jinja2.ext.loopcontrols']) # (for {% break %}) env.filters['tojson'] = lambda obj, **kwargs: json.dumps(obj, **kwargs) # jinja2's tojson escapes <>& for HTML safety env.globals['raise_exception'] = lambda msg: (_ for _ in ()).throw(RuntimeError(msg)) env.globals['strftime_now'] = lambda fmt: time.strftime(fmt) diff --git a/tinygrad_repo/tinygrad/llm/kernels/amd.py b/tinygrad_repo/tinygrad/llm/kernels/amd.py index 3d7228995..d60cd086f 100644 --- a/tinygrad_repo/tinygrad/llm/kernels/amd.py +++ b/tinygrad_repo/tinygrad/llm/kernels/amd.py @@ -667,7 +667,7 @@ def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start: Q_lds = QP_lds[:, :D] Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid].store(q.reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid]) load_k = UOp.range(KV_ELEMS_PER_THREAD, 90) - kval = k.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_k].float() + kval = k.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_k] K_store = KV_lds.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_k].store(kval).end(load_k) Q_lds, KV_lds_k = Q_lds.after(Q_store, K_store), KV_lds.after(Q_store, K_store) S_reg = _reg((TM, TN), 0, n_tile) @@ -709,7 +709,7 @@ def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start: V_lds = UOp.alloc((D, BLOCK_N + LDS_PAD), dtypes.half, addrspace=AddrSpace.LOCAL)[:, :BLOCK_N] V_copy, load_v = V_lds.after(qk_done).permute(1, 0), UOp.range(KV_ELEMS_PER_THREAD, 390) v_pos = n_tile*BLOCK_N + (tid*KV_ELEMS_PER_THREAD + load_v)//D - vval = (v_pos < valid_kv_len).where(v.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_v].float(), 0) + vval = (v_pos < valid_kv_len).where(v.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_v], 0) V_store = V_copy.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_v].store(vval).end(load_v) P_lds, V_lds = P_lds.after(P_store, V_store), V_lds.after(P_store, V_store) pv_acc = _reg((TM, TD), 0, n_tile) diff --git a/tinygrad_repo/tinygrad/llm/model.py b/tinygrad_repo/tinygrad/llm/model.py index 8d19736cc..05447993c 100644 --- a/tinygrad_repo/tinygrad/llm/model.py +++ b/tinygrad_repo/tinygrad/llm/model.py @@ -106,7 +106,7 @@ class TransformerConfig: expert_bias: bool = False expert_proj_bias: bool = False swiglu_alpha: float = 1.0 - swiglu_clamp_exp: float|None = None + swiglu_limit: float|None = None swiglu_up_bias: float = 0.0 sliding_window: int = 0 sliding_window_pattern: int = 0 @@ -157,16 +157,16 @@ class FFNBlock: if normalize_topk: probs = probs / probs.sum(axis=-1, keepdim=True) probs = probs * self.config.routed_scaling_factor gate, up = self.ffn_gate_exps(sel, h), self.ffn_up_exps(sel, h) - act = gated_activation(gate, up, alpha=self.config.swiglu_alpha, limit=self.config.swiglu_clamp_exp, up_bias=self.config.swiglu_up_bias) + act = gated_activation(gate, up, alpha=self.config.swiglu_alpha, limit=self.config.swiglu_limit, up_bias=self.config.swiglu_up_bias) x_down = self.ffn_down_exps(sel, act.contiguous()) # (B, T, k, D) out = (x_down * probs.unsqueeze(-1)).sum(axis=2) # (B, T, D) if hasattr(self, 'ffn_gate_shexp'): - shexp = self.ffn_down_shexp(self.ffn_gate_shexp(x).silu().contiguous() * self.ffn_up_shexp(x)) + shexp = self.ffn_down_shexp(gated_activation(self.ffn_gate_shexp(x), self.ffn_up_shexp(x), limit=self.config.swiglu_limit).contiguous()) if hasattr(self, 'ffn_gate_inp_shexp'): shexp = shexp * (x * self.ffn_gate_inp_shexp["weight"]).sum(axis=-1, keepdim=True).sigmoid() out = out + shexp return out # TODO: remove the need for this contiguous - return self.ffn_down(self.ffn_gate(x).silu().contiguous() * self.ffn_up(x)) + return self.ffn_down(gated_activation(self.ffn_gate(x), self.ffn_up(x), limit=self.config.swiglu_limit).contiguous()) # given the token-prefix match, return how much cached state this block can still reuse def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return prefix_len @@ -286,14 +286,16 @@ class MLATransformerBlock(FFNBlock): k = Tensor(self.cache_k.uop.after(self.cache_k[:, :, start_pos:start_pos+T, :].uop.store(k_store.uop)))[:, :, 0:start_pos+T, :] v = k[..., :self.config.kv_lora_rank] - mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \ - if resolve(T != 1) else None attn = q @ k.transpose(-1, -2) * (1.0 / self.config.head_dim ** 0.5) - if mask is not None: attn = attn + mask + if (mask:=self._attn_mask(x, start_pos)) is not None: attn = attn + mask attn = attn.softmax(-1) attn = ((attn @ v) @ self.attn_v_b["weight"].transpose(-1, -2)).transpose(1, 2).reshape(B, T, -1) return self.attn_output(attn) + def _attn_mask(self, x:Tensor, start_pos:int|UOp) -> Tensor|None: # (1, 1, T, start_pos+T), None: everything is visible + T = x.shape[1] + return Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) if resolve(T != 1) else None + def _init_state(self, x:Tensor): if not hasattr(self, "cache_k"): self.cache_k = Tensor.empty(x.shape[0], 1, self.config.max_context, self.config.kv_lora_rank + self.config.rope_dim, device=x.device) @@ -335,8 +337,8 @@ class GatedDeltaNetBlock(FFNBlock): out_gate = out_gate.reshape(B, T, self.num_v_heads, self.head_v_dim) beta = self.ssm_beta(x).sigmoid().reshape(B, T, self.num_v_heads) alpha = self.ssm_f_b(self.ssm_f_a(x)) if is_kda else self.ssm_alpha(x) - log_alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, T, self.num_v_heads, -1) * - self.ssm_a.reshape(self.num_v_heads, -1)) + g, a = (alpha.float() + self.ssm_dt["bias"]).reshape(B, T, self.num_v_heads, -1), self.ssm_a.reshape(self.num_v_heads, -1) + log_alpha = g.softplus() * a # qkv conv, conv_state is reset when starting from position 0 conv_state = initial.where(0, self.conv_state) @@ -353,9 +355,9 @@ class GatedDeltaNetBlock(FFNBlock): out_gate = out_gate.pad_to((B, T_pad, self.num_v_heads, self.head_v_dim)) beta, log_alpha = beta.pad_to((B, T_pad, self.num_v_heads)), log_alpha.pad_to((B, T_pad, *log_alpha.shape[2:])) q, k, v = conv_out.split([self.q_dim, self.q_dim, self.conv_channels - 2*self.q_dim], dim=-1) - qk_eps = 1e-12 if is_kda else 1e-6 - q, k = (z.reshape(B, T_pad, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=qk_eps) - .repeat(1, 1, self.num_v_heads//self.num_k_heads, 1) for z in (q, k)) + q, k = (z.reshape(B, T_pad, self.num_k_heads, self.head_k_dim) for z in (q, k)) + # l2 norm x/sqrt(sum(x²)+1e-6) as in the FLA kernels, then the v heads share the k heads + q, k = ((z.rmsnorm(eps=1e-6/self.head_k_dim) / self.head_k_dim**0.5).repeat(1, 1, self.num_v_heads//self.num_k_heads, 1) for z in (q, k)) v = v.reshape(B, T_pad, self.num_v_heads, self.head_v_dim) # layout the per-step operands to broadcast against the (B, H, V, K) state q, k, v, beta = (z.transpose(1, 2).float() for z in (q, k, v, beta)) @@ -451,17 +453,19 @@ class Transformer: max_context = min(max_context, kv[f'{arch}.context_length']) if max_context is not None else kv[f'{arch}.context_length'] + num_blocks = kv[f'{arch}.block_count'] - kv.get(f'{arch}.nextn_predict_layers', 0) ssm = None ssm_layers: tuple[bool, ...] = () if arch in ('qwen35', 'qwen35moe'): ssm = SSMConfig(**{k: kv[f'{arch}.ssm.{k}'] for k in ('conv_kernel','state_size','group_count','time_step_rank','inner_size')}) - ssm_layers = tuple((i+1) % kv[f'{arch}.full_attention_interval'] != 0 for i in range(kv[f'{arch}.block_count'])) + ssm_layers = tuple((i+1) % kv[f'{arch}.full_attention_interval'] != 0 for i in range(num_blocks)) elif arch == 'kimi-linear': - ssm_layers = tuple(x == 0 for x in n_kv_heads) + ssm_layers = tuple(x == 0 for x in n_kv_heads)[:num_blocks] n_kv_heads = max(n_kv_heads) ssm = SSMConfig(kv[f'{arch}.ssm.conv_kernel'], kv[f'{arch}.kda.head_dim'], n_heads, n_heads, n_heads*kv[f'{arch}.kda.head_dim'], kda=True) for i, is_ssm in enumerate(ssm_layers): if not is_ssm: continue + state_dict[f"blk.{i}.ssm_a"] = state_dict[f"blk.{i}.ssm_a"].reshape(n_heads, 1) state_dict[f"blk.{i}.attn_qkv.weight"] = state_dict.pop(f"blk.{i}.attn_q.weight").cat( state_dict.pop(f"blk.{i}.attn_k.weight"), state_dict.pop(f"blk.{i}.attn_v.weight"), dim=0).contiguous() state_dict[f"blk.{i}.ssm_conv1d.weight"] = state_dict.pop(f"blk.{i}.ssm_conv1d_q.weight").cat( @@ -491,12 +495,12 @@ class Transformer: elif kv_lora_rank and 'attn_kv_a_mqa.weight' in name: state_dict[name] = state_dict[name][:kv_lora_rank].cat(state_dict[name][kv_lora_rank:].rearrange("(h two) d -> (two h) d", two=2), dim=0) config = TransformerConfig( - num_blocks=kv[f'{arch}.block_count'] - kv.get(f'{arch}.nextn_predict_layers', 0), dim=kv[f'{arch}.embedding_length'], + num_blocks=num_blocks, dim=kv[f'{arch}.embedding_length'], hidden_dim=kv.get(f'{arch}.expert_feed_forward_length', kv.get(f'{arch}.feed_forward_length', 0)), n_heads=n_heads, n_kv_heads=n_kv_heads, norm_eps=kv[f'{arch}.attention.layer_norm_rms_epsilon'], vocab_size=len(kv['tokenizer.ggml.tokens']), head_dim=head_dim, - rope_theta=kv[f'{arch}.rope.freq_base'], rope_dim=rope_dim, yarn=yarn, + rope_theta=kv.get(f'{arch}.rope.freq_base', 10000.0), rope_dim=rope_dim, yarn=yarn, v_head_dim=kv.get(f'{arch}.attention.value_length_mla', kv.get(f'{arch}.attention.value_length', head_dim)), max_context=max_context, qk_norm=int(state_dict['blk.0.attn_q_norm.weight'].shape[0]) if 'blk.0.attn_q_norm.weight' in state_dict else 0, @@ -516,7 +520,7 @@ class Transformer: qkv_bias='blk.0.attn_q.bias' in state_dict, expert_bias=f"blk.{kv.get(f'{arch}.leading_dense_block_count', 0)}.exp_probs_b.bias" in state_dict, expert_proj_bias='blk.0.ffn_gate_exps.bias' in state_dict, attn_output_bias='blk.0.attn_output.bias' in state_dict, - swiglu_alpha=1.702 if arch == 'gpt-oss' else 1.0, swiglu_clamp_exp=7.0 if arch == 'gpt-oss' else None, + swiglu_alpha=1.702 if arch == 'gpt-oss' else 1.0, swiglu_limit=7.0 if arch == 'gpt-oss' else kv.get(f'{arch}.swiglu_clamp_exp', [None])[0], swiglu_up_bias=1.0 if arch == 'gpt-oss' else 0.0, attn_sinks='blk.0.attn_sinks.weight' in state_dict, sliding_window=kv.get(f'{arch}.attention.sliding_window', 0), sliding_window_pattern=kv.get(f'{arch}.attention.sliding_window_pattern', 2 if arch == 'gpt-oss' else 0)) diff --git a/tinygrad_repo/tinygrad/mixin/creation.py b/tinygrad_repo/tinygrad/mixin/creation.py index 77df13c1b..ef5f2587c 100644 --- a/tinygrad_repo/tinygrad/mixin/creation.py +++ b/tinygrad_repo/tinygrad/mixin/creation.py @@ -1,6 +1,6 @@ from typing import TYPE_CHECKING, Callable, Self from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, dtypes, to_dtype -from tinygrad.helpers import argfix, prod +from tinygrad.helpers import argfix from tinygrad.mixin.dtype import DTypeMixin from tinygrad.mixin.movement import MovementMixin @@ -32,12 +32,11 @@ class CreationMixin(DTypeMixin, MovementMixin): print(t.shape) ``` """ - from tinygrad.uop.ops import UOp, Ops, ParamArg, to_max_shape + from tinygrad.uop.ops import UOp from tinygrad.device import canonicalize_device dt = to_dtype(dtype) if dtype is not None else dtypes.default_float if dt in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {dt}") new_shape = argfix(*shape) - max_shape = to_max_shape(new_shape) dev = canonicalize_device(device) if axis is not None and isinstance(dev, tuple): axis = range(max(1, len(new_shape)))[axis] @@ -45,10 +44,7 @@ class CreationMixin(DTypeMixin, MovementMixin): if new_shape[axis] % len(dev) != 0: raise RuntimeError(f"can't split {new_shape} on axis {axis} over {dev}") local = tuple(s//len(dev) if i == axis else s for i,s in enumerate(new_shape)) return cls._wrap_uop(cls.empty(*local, device=dev, dtype=dt)._uop.unshard(axis)) - u = UOp(Ops.ALLOC, src=(UOp.const(prod(max_shape)),)+UOp.device_range_src(dev), - arg=ParamArg(next(UOp.unique_num), dt, device=dev, bind_on_realize=True)) - u = u.reshape(max_shape).shrink_to(new_shape) - return cls._wrap_uop(u) + return cls._wrap_uop(UOp.alloc(new_shape, dt, device=dev, bind_on_realize=True)) def empty_like(self, dtype: DTypeLike|None=None, device: str|tuple[str, ...]|None=None) -> Self: """ diff --git a/tinygrad_repo/tinygrad/mixin/gradient.py b/tinygrad_repo/tinygrad/mixin/gradient.py index b89159e95..338815b9a 100644 --- a/tinygrad_repo/tinygrad/mixin/gradient.py +++ b/tinygrad_repo/tinygrad/mixin/gradient.py @@ -27,7 +27,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]: git = iter(k.arg.grad_fxn(*real, call=k) if len(real) > 1 else k.arg.grad_fxn(real[0], k)) return (None,) + tuple(None if i in outputs else next(git) for i in range(len(args))) assert outputs, f"expected a CALL with output STOREs or a grad_fxn, got {fxn.op}" - params = {p.arg.slot:p for p in fxn.toposort(enter_calls=False) if p.op is Ops.PARAM and p.arg.slot >= 0} + params = {p.arg.slot:p for p in fxn.toposort() if p.op is Ops.PARAM and p.arg.slot >= 0} grad_args = tuple(ctx.src[i] for i in outputs) root_grad = UOp.sink(*[g if g.device is None else g.param_like(len(args)+i) for i,g in enumerate(grad_args)]) grads = compute_gradient(UOp.sink(*[st.src[1] for st in outputs.values()]), root_grad, set(params.values())) @@ -41,7 +41,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]: args += grad_args bwd_body = renumber_invalid_outputs(bwd_body) # Compact this scope's PARAMs, capturing free Variables as inputs. - used = sorted((p for p in bwd_body.toposort(enter_calls=False) if p.op is Ops.PARAM), key=lambda p:p.arg.slot) + used = sorted((p for p in bwd_body.toposort() if p.op is Ops.PARAM), key=lambda p:p.arg.slot) bwd_args = [p if p.arg.slot < 0 else args[p.arg.slot] for p in used] bwd_body = bwd_body.substitute({p:p.replace(arg=dataclasses.replace(p.arg, slot=i, name=None, val=None)) for i,p in enumerate(used)}, walk=True) @@ -52,14 +52,13 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]: def partial_store_gradient(ctx:UOp, dest:UOp, view:UOp): # A write through a non-overlapping view replaces only that region of the returned state. path, base = [], view - while base is not dest and base.op in {Ops.RESHAPE, Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD}: + while base is not dest and base.op in {Ops.RESHAPE, Ops.SHRINK, Ops.PERMUTE, Ops.PAD}: path.append(base) base = base.src[0] if base is not dest: return None grad = ctx for mop in reversed(path): grad = mop.replace(src=(grad,)+mop.src[1:]) - mask = grad.const_like(1) - for mop in path: mask = pm_gradient.rewrite(mop, ctx=mask)[0] + mask = compute_gradient(view, grad.const_like(1), {dest})[dest] return mask.cast(dtypes.bool).where(0, ctx), grad # ctx is grad_output @@ -85,9 +84,8 @@ pm_gradient = PatternMatcher([ (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape), None)), (UPat(Ops.EXPAND), lambda ctx: (ctx, None)), (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[0], s+p[0]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)), - (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[0]-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)), + (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx._mop(Ops.PAD, tuple((p[0], s) for s,p in zip(ret.src[0].shape, ret.marg))), None, None)), (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.marg)),)), - (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip([i for i,x in enumerate(ret.marg) if x]),)), (UPat(Ops.STACK, name="ret"), lambda ctx, ret: tuple(ctx[i] for i in range(len(ret.src)))), (UPat(Ops.COPY, name="ret"), lambda ctx, ret: (ctx.copy_to_device(ret.src[0].device),) + (None,) * (len(ret.src)-1)), (UPat(Ops.UNSHARD, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src), diff --git a/tinygrad_repo/tinygrad/mixin/movement.py b/tinygrad_repo/tinygrad/mixin/movement.py index 35ebe0eb7..092b031aa 100644 --- a/tinygrad_repo/tinygrad/mixin/movement.py +++ b/tinygrad_repo/tinygrad/mixin/movement.py @@ -250,8 +250,22 @@ class MovementMixin: assert all(not isinstance(x, bool) and x >= 0 and x < self.ndim for x in axis_arg), f"flip args must be axis ints {axis_arg}" if len(axis_arg) != len(dedup(axis_arg)): raise RuntimeError(f"dim can appear at most once, getting {axis_arg}") - flip_arg = tuple([i in axis_arg for i in range(len(self.shape))]) - return self._mop(Ops.FLIP, arg=flip_arg) if any(flip_arg) else self + ret = self + for axis in axis_arg: + if resolve((n:=ret.shape[axis]) <= 1, False): continue + x = ret.transpose(0, axis) + bounds, crop = (None,) * (x.ndim-1), tuple((0, s) for s in x.shape[1:]) + tail = tuple(s if resolve(s > 0, False) else s+1 for s in x.shape[1:]) + m, tail_pad = n if resolve(n >= 2, False) else n+2, tuple((0, s) for s in tail) + # Place values on an anti-diagonal, transpose it, then extract the anti-diagonal. + # Pad possibly empty dimensions before reshaping, then crop them back afterwards. + x = x._mop(Ops.PAD, ((0, m),)+tail_pad).reshape(m, 1, *tail) + x = x._mop(Ops.PAD, ((0, m), (0, m-1))+tail_pad).reshape(m*(m-1), *tail) + x = x._mop(Ops.PAD, ((m-1, m*m),)+tail_pad).reshape(m, m, *tail).transpose(0, 1).reshape(m*m, *tail) + x = x._mop(Ops.SHRINK, ((m-1, m*(m-1)),)+tail_pad).reshape(m, m-1, *tail) + x = x.shrink((None, (0, 1))+bounds).reshape(m, *tail) + ret = x.shrink(((m-n, m),)+crop).transpose(0, axis) + return ret def stack(self, *args: Self, dim: int = 0) -> Self: """ @@ -284,10 +298,10 @@ class MovementMixin: """`.view` is an alias for `.reshape`.""" return self.reshape(shape, *args) - def squeeze(self, dim: int | None = None) -> Self: + def squeeze(self, dim: int | Sequence[int] | None = None) -> Self: """ Returns a tensor with specified dimensions of input of size 1 removed. - If `dim` is not specified, all dimensions with size 1 are removed. + `dim` can be an int or a sequence of ints. If `dim` is not specified, all dimensions with size 1 are removed. ```python exec="true" source="above" session="tensor" result="python" t = Tensor.zeros(2, 1, 2, 1, 2) @@ -299,11 +313,12 @@ class MovementMixin: ```python exec="true" source="above" session="tensor" result="python" print(t.squeeze(1).shape) ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.squeeze((1, 3)).shape) + ``` """ - if dim is None: - return self.reshape(tuple(dim for dim in self.shape if dim != 1)) - dim = self._resolve_dim(dim) - return self if not self.ndim or self.shape[dim] != 1 else self.reshape(self.shape[:dim] + self.shape[dim + 1 :]) + dims = range(self.ndim) if dim is None else [self._resolve_dim(d) for d in make_tuple(dim, 1)] + return self.reshape(tuple(s for i,s in enumerate(self.shape) if i not in dims or s != 1)) def unsqueeze(self, dim: int) -> Self: """ @@ -587,6 +602,7 @@ class MovementMixin: ``` """ repeats = argfix(repeats, *args) + if any(resolve(r < 0, False) for r in repeats): raise RuntimeError(f"repeats should be non negative, got {repeats} instead") base_shape = _align_left(self.shape, repeats)[0] unsqueezed_shape = flatten([[s] if r == 1 else [1, s] for r, s in zip(repeats, base_shape)]) expanded_shape = flatten([[s] if r == 1 else [r, s] for r, s in zip(repeats, base_shape)]) diff --git a/tinygrad_repo/tinygrad/mixin/op.py b/tinygrad_repo/tinygrad/mixin/op.py index 04412b5e2..0c45d6302 100644 --- a/tinygrad_repo/tinygrad/mixin/op.py +++ b/tinygrad_repo/tinygrad/mixin/op.py @@ -1172,7 +1172,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): for dim in reversed(axes): mask, values = functools.reduce(lambda x,y: (x[0]|y[0], y[0].where(y[1], x[1])), zip(mask.split(1, dim), values.split(1, dim))) # remove extra dims from reduce - for dim in reversed(axes): mask, values = mask.squeeze(dim), values.squeeze(dim) + mask, values = mask.squeeze(axes), values.squeeze(axes) # select from values for each True element in mask else select from self return mask.where(values, self) @@ -1641,6 +1641,23 @@ class OpMixin(ElementwiseMixin, ReduceMixin): padding = flatten((((k-1)*d-pB,(k-1)*d-pA+op) for k,d,(pB,pA),op in reversed(list(zip(HW, dilation, padding, output_padding))))) return x.conv2d(w.flatten(end_dim=1), groups=groups, bias=bias, dilation=dilation, padding=padding) + def rmsnorm(self, axis:int|tuple[int,...]=-1, eps:float=1e-6) -> Self: + """ + Applies Root Mean Square Normalization over a mini-batch of inputs. + + - Paper: https://arxiv.org/abs/1910.07467 + + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor.randn(8, 10, 16) * 2 + 8 + print(t.square().mean().sqrt().item()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + t = t.rmsnorm() + print(t.square().mean().sqrt().item()) + ``` + """ + return self.mul(self.square().mean(axis, keepdim=True).add(eps).rsqrt()) + def layernorm(self, axis:int|tuple[int,...]=-1, eps:float=1e-5) -> Self: """ Applies Layer Normalization over a mini-batch of inputs. @@ -1656,8 +1673,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): print(t.mean().item(), t.std().item()) ``` """ - y = (self - self.mean(axis, keepdim=True)) - return y.mul((y*y).mean(axis, keepdim=True).add(eps).rsqrt()) + return (self - self.mean(axis, keepdim=True)).rmsnorm(axis, eps) def batchnorm(self, weight:Self|None, bias:Self|None, mean:Self, invstd:Self, axis:int|tuple[int, ...]=1) -> Self: """ diff --git a/tinygrad_repo/tinygrad/nn/__init__.py b/tinygrad_repo/tinygrad/nn/__init__.py index 15e4b1358..646e9064d 100644 --- a/tinygrad_repo/tinygrad/nn/__init__.py +++ b/tinygrad_repo/tinygrad/nn/__init__.py @@ -2,7 +2,7 @@ from __future__ import annotations import math, functools from tinygrad.tensor import Tensor from tinygrad.dtype import dtypes -from tinygrad.helpers import prod, make_tuple, flatten, USE_ATOMICS, TRAINING +from tinygrad.helpers import prod, make_tuple, flatten, USE_ATOMICS, TRAINING, to_tuple from tinygrad.nn import optim, state, datasets # noqa: F401 class BatchNorm: @@ -297,10 +297,8 @@ class RMSNorm: self.eps = eps self.weight = Tensor.ones(dim) if elementwise_affine else None - def _norm(self, x:Tensor) -> Tensor: return x * (x.square().mean(-1, keepdim=True) + self.eps).rsqrt() - def __call__(self, x:Tensor) -> Tensor: - x = self._norm(x.float()).cast(x.dtype) + x = x.float().rmsnorm(eps=self.eps).cast(x.dtype) return x if self.weight is None else x * self.weight from tinygrad.uop.ops import UOp, KernelInfo, Ops, AxisType @@ -326,8 +324,7 @@ def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple: return out.flatten()[i].store(0).end(i).sink(arg=KernelInfo(name="zero")) grad_weight_uop = grad_weight_uop.custom_kernel(fxn=_zero_kernel)[0] - # TODO: do we have a universal helper for this? - device = call.device.split(":")[0] if not isinstance(call.device, tuple) else call.device[0].split(":")[0] + device = to_tuple(call.device)[0].split(":")[0] # this is the real atomic kernel def _embedding_bwd_kernel(grad_weight:UOp, grad_emb:UOp, idx:UOp) -> UOp: diff --git a/tinygrad_repo/tinygrad/nn/onnx.py b/tinygrad_repo/tinygrad/nn/onnx.py index d61c97882..24a3ff4f6 100644 --- a/tinygrad_repo/tinygrad/nn/onnx.py +++ b/tinygrad_repo/tinygrad/nn/onnx.py @@ -683,9 +683,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT def Shrink(x:Tensor, bias:float=0.0, lambd:float=0.5): return (x < -lambd).where(x+bias, (x > lambd).where(x-bias, 0)).cast(x.dtype) def Transpose(x:Tensor, perm:tuple[int, ...]|None=None): return x.permute(order=perm or list(range(x.ndim)[::-1])) - def Squeeze(data:Tensor, axes:Sequence[int]|None=None): - if axes is None: return data.squeeze() - return functools.reduce(lambda d, dim: d.squeeze(dim), sorted(map(data._resolve_dim, axes), reverse=True), data) + def Squeeze(data:Tensor, axes:Sequence[int]|None=None): return data.squeeze(axes) def Unsqueeze(data:Tensor, axes:Sequence[int]): return functools.reduce(lambda d, dim: d.unsqueeze(dim), sorted(data._resolve_dim(a, extra=len(axes)) for a in axes), data) @@ -1094,8 +1092,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT def RMSNormalization(X:Tensor, scale:Tensor, axis:int=-1, epsilon:float=1e-5, stash_type:int=1): assert stash_type == 1, "only float32 is supported" - norm = X.cast(dtypes.float).square().mean(axis=tuple(range(axis + X.ndim if axis < 0 else axis, X.ndim)), keepdim=True).add(epsilon).rsqrt() - return X * norm * scale + return X.cast(dtypes.float).rmsnorm(axis=tuple(range(axis + X.ndim if axis < 0 else axis, X.ndim)), eps=epsilon) * scale def RotaryEmbedding(X:Tensor, cos_cache:Tensor, sin_cache:Tensor, position_ids:Tensor|None=None, interleaved:int=0, num_heads:int|None=None, rotary_embedding_dim:int=0): diff --git a/tinygrad_repo/tinygrad/renderer/amd/__init__.py b/tinygrad_repo/tinygrad/renderer/amd/__init__.py index 875005f57..96a763c23 100644 --- a/tinygrad_repo/tinygrad/renderer/amd/__init__.py +++ b/tinygrad_repo/tinygrad/renderer/amd/__init__.py @@ -6,7 +6,7 @@ from tinygrad.renderer.amd.dsl import Inst, FixedBitField, EnumBitField # 0xf9 (249) = SDWA, 0xfa (250) = DPP16 for CDNA (GFX9) _VARIANT_SRC0 = {"_SDWA_SDST": 0xf9, "_SDWA": 0xf9, "_DPP16": 0xfa} -def _matches(data: bytes, cls: type[Inst]) -> bool: +def _matches(data: bytes|memoryview, cls: type[Inst]) -> bool: """Check if data matches all FixedBitFields and op is in allowed.""" for _, field in cls._fields: dword_idx = field.lo // 32 @@ -63,7 +63,7 @@ def _load_formats() -> dict[str, list[type[Inst]]]: } return _FORMATS -def detect_format(data: bytes, arch: str = "rdna3") -> type[Inst]: +def detect_format(data: bytes|memoryview, arch: str = "rdna3") -> type[Inst]: """Detect instruction format from machine code bytes.""" assert len(data) >= 4, f"need at least 4 bytes, got {len(data)}" for cls in _load_formats()[arch]: diff --git a/tinygrad_repo/tinygrad/renderer/amd/dsl.py b/tinygrad_repo/tinygrad/renderer/amd/dsl.py index 6526706be..211696f04 100644 --- a/tinygrad_repo/tinygrad/renderer/amd/dsl.py +++ b/tinygrad_repo/tinygrad/renderer/amd/dsl.py @@ -437,7 +437,7 @@ class Inst: return None @classmethod - def from_bytes(cls, data: bytes): + def from_bytes(cls, data: bytes|memoryview): inst = object.__new__(cls) inst._raw = int.from_bytes(data[:cls._base_size], 'little') # Upgrade to variant class if needed (_LIT, _DPP8, _DPP16, _SDWA) diff --git a/tinygrad_repo/tinygrad/renderer/cstyle.py b/tinygrad_repo/tinygrad/renderer/cstyle.py index 012b25277..67ebc3da8 100644 --- a/tinygrad_repo/tinygrad/renderer/cstyle.py +++ b/tinygrad_repo/tinygrad/renderer/cstyle.py @@ -267,8 +267,9 @@ class CStyleLanguage(Renderer): def param_type(self, p:UOp): return "volatile "*p.arg.volatile + self._render_dtype(p.dtype, 1, p.addrspace, True, p.addrspace != AddrSpace.ALU) def render(self, uops:list[UOp]) -> str: prefix, call_bodies, self.fn_names = [], [], dict[UOp, str]() - prefix += [f"extern void {f}();" for f in dedup(u.arg.name for u in UOp.sink(*uops).toposort() if u.op is Ops.CUSTOM_FUNCTION)] # symbols to link - for body in (u for u in UOp.sink(*uops).toposort() if u.op is Ops.LINEAR): + topo = UOp.sink(*uops).toposort(enter_calls=True) + prefix += [f"extern void {f}();" for f in dedup(u.arg.name for u in topo if u.op is Ops.CUSTOM_FUNCTION)] # symbols to link + for body in (u for u in topo if u.op is Ops.LINEAR): self.fn_names[body] = body.arg + (f"_{n}" if (n:=sum(b.arg == body.arg for b in self.fn_names)) else "") # a name traced with other args _, call, bufs = self._render(body.src) params = ', '.join(f"{self.param_type(p)} {n}" for n,(p,_) in bufs) @@ -502,9 +503,9 @@ class HIPRenderer(CStyleLanguage): global_prod_max = (0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF) @staticmethod - def is_cdna(arch): return arch.split(":")[0] in {"gfx942", "gfx950"} + def is_cdna(arch): return arch in {"gfx942", "gfx950"} @staticmethod - def is_cdna4(arch): return arch.split(":")[0] == "gfx950" + def is_cdna4(arch): return arch == "gfx950" def __init__(self, target:Target, use_hipcc=False): # gfx942 => MI300, gfx1100 => RX 7900, gfx1201 => RX 9700 super().__init__(target) from tinygrad.runtime.support.compiler_amd import HIPCompiler, HIPCCCompiler @@ -588,14 +589,16 @@ class HIPRenderer(CStyleLanguage): elif (N, M, K) == (16, 16, 32): type_map = {**type_map, dtypes.bf16: "_bf16", dtypes.f16: "_f16"} elif (N, M, K) == (16, 16, 128): type_map = {**type_map, dtypes.fp8e4m3: "_f8f6f4", dtypes.fp8e5m2: "_f8f6f4"} prefix.append(f"#define __{name} __builtin_amdgcn_mfma_{'scale_' if K == 128 else ''}f32_{N}x{M}x{K}{type_map[dtype_in]}") + elif dtype_out == dtypes.int32: + # RDNA4 uses 8 int8 values per lane in 2 VGPRs, RDNA3 uses 16 in 4 VGPRs + num_regs = 2 if (is_rdna4:=self.tensor_cores == tc.amd_rdna4) else 4 + prefix.append(f"typedef int wmma_int{num_regs} __attribute__((ext_vector_type({num_regs})));\n"+ + f"static inline __attribute__((device)) int8 __{name}"+f"""(signed_char{num_regs*4} a, signed_char{num_regs*4} b, int8 c) {{ + return __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32{'_gfx12' if is_rdna4 else ''}(true, __builtin_bit_cast(wmma_int{num_regs}, a), + true, __builtin_bit_cast(wmma_int{num_regs}, b), c, false);\n}}""") # #define __WMMA_16_16_16_f16_f16 __builtin_amdgcn_wmma_f16_16x16x16_f16_w32_gfx12 elif self.tensor_cores == tc.amd_rdna4: prefix.append(f"#define __{name} __builtin_amdgcn_wmma_{type_map[dtype_out]}_16x16x16_{type_map[dtype_in]}_w32_gfx12") - elif dtype_out == dtypes.int32: - prefix.append("typedef int wmma_int4 __attribute__((ext_vector_type(4)));\n"+ - f"static inline __attribute__((device)) int8 __{name}"+"""(signed_char16 a, signed_char16 b, int8 c) { - return __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, __builtin_bit_cast(wmma_int4, a), - true, __builtin_bit_cast(wmma_int4, b), c, false);\n}""") elif dtype_out == dtypes.float: prefix.append(f"#define __{name} __builtin_amdgcn_wmma_f32_16x16x16_{'f16' if dtype_in == dtypes.half else 'bf16'}_w32") else: prefix.append(f"static inline __attribute__((device)) half8 __{name}"+"""(half16 a, half16 b, half8 c) { diff --git a/tinygrad_repo/tinygrad/renderer/isa/x86.py b/tinygrad_repo/tinygrad/renderer/isa/x86.py index 56d4f0ebe..37e36c641 100644 --- a/tinygrad_repo/tinygrad/renderer/isa/x86.py +++ b/tinygrad_repo/tinygrad/renderer/isa/x86.py @@ -319,7 +319,7 @@ def copy_op(dt:DType) -> X86Ops: isel_matcher = PatternMatcher([ # **** Op -> Op **** # range is lowered to acc, cmp, jmp after regalloc - (UPat(Ops.RANGE, src=(UPat.cvar("c").cast(),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(x.dtype, c.val),) + x.src[1:])), + (UPat(Ops.RANGE, src=(UPat.cvar("c").cast(),), name="x"), lambda c,x: x.replace(src=(imm(x.dtype, c.val),))), # BACKEDGE becomes a conditional jump referencing the RANGE start label (UPat(Ops.BACKEDGE, src=(UPat(), UPat(), UPat(GroupOp.Comparison, name="cond")), name="x"), lambda x,cond: cond.ins(X86Ops.LOOP_CMP, tag=cond.op, src=cond.src + x.src[:2])), @@ -497,7 +497,7 @@ def lower_range(ctx, x:UOp) -> tuple[UOp, list[UOp]]: if x.dtype is dtypes.void: return (label, [label]) else: # an AFTER-wrapped bound carries the loop's ordering deps, they keep the acc unique per loop (hash-consing) - bound, deps = ((b:=x.src[0]).src[0], b.src[1:]) if (b:=x.src[0]).op is Ops.AFTER else (x.src[0], x.src[1:]) + bound, deps = ((b:=x.src[0]).src[0], b.src[1:]) if (b:=x.src[0]).op is Ops.AFTER else (x.src[0], ()) acc = x.ins(X86Ops.MOVi, src=(imm(x.dtype, 0),) + tuple(deps)) cmp = UOp(Ops.INS, arg=(X86Ops.CMPi if bound.op is Ops.CAST else X86Ops.CMP, dtypes.void), src=(acc, bound)) jump_out = UOp(Ops.INS, arg=(X86Ops.JGE, dtypes.void), src=(cmp,), tag=f".LOOP_OUT_{loop_label}") diff --git a/tinygrad_repo/tinygrad/renderer/llvmir.py b/tinygrad_repo/tinygrad/renderer/llvmir.py index 940bfc1b5..530dce9b5 100644 --- a/tinygrad_repo/tinygrad/renderer/llvmir.py +++ b/tinygrad_repo/tinygrad/renderer/llvmir.py @@ -16,7 +16,7 @@ def ldt(dt:DType, count=1, ptr=False): dtypes.uint8: "i8", dtypes.uint16: "i16", dtypes.uint32: "i32", dtypes.uint64: "i64", **{d: "i8" for d in dtypes.fp8s}, dtypes.float16: "half", dtypes.bfloat16: "bfloat", dtypes.float32: "float", dtypes.float64: "double"}[dt] -def lparam(u:UOp) -> str: return ldt(u.dtype, ptr=u.op is Ops.INDEX) # an external function's arg, an index is its address +def lparam(u:UOp) -> str: return ldt(u.dtype, ptr=u.addrspace not in (None, AddrSpace.ALU)) # an external function's arg, an index is its address def lconst(x, dtype:DType): if dtype in dtypes.floats: @@ -79,7 +79,7 @@ lop = {**{x:unsigned_lop for x in (dtypes.bool,)+dtypes.uints}, **{x:signed_lop base_rewrite = PatternMatcher([ # memory load/store - (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat((Ops.BUFFER, Ops.PARAM, Ops.AFTER)),), allow_any_len=True, name="x"), lambda ctx,x: + (UPat((Ops.INDEX, Ops.SHRINK), name="x"), lambda ctx,x: None if x.src[0].addrspace in (None, AddrSpace.ALU) else f" {ctx[x]} = getelementptr inbounds {ldt(x.dtype)}, {ldt(x.dtype, ptr=True)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}"), # register index (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.cvar("c").cast()), name="x"), lambda ctx,buf,c,x: @@ -196,8 +196,9 @@ class LLVMRenderer(Renderer): kernel.append(f" {r[u]} = addrspacecast [{size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{size} x {ldt(u.dtype)}]*") else: kernel.append(f" {r[u]} = alloca [{size} x {ldt(u.dtype)}], align 16") + elif u.op is Ops.BINARY: r[u] = f"@bin_{u.key.hex()}" elif u.op is Ops.CAST and u.src[0].op is Ops.CONST: r[u] = lconst(u.src[0].val, u.dtype) - elif u.op is Ops.CAST and ldt(u.dtype) == ldt(u.src[0].dtype): + elif (u.op is Ops.CAST and ldt(u.dtype) == ldt(u.src[0].dtype)) or (u.op is Ops.BITCAST and u.addrspace not in (None, AddrSpace.ALU)): r[u] = r[u.src[0]] # cast from signed to unsigned of the same size is a noop, or pointer cast else: vc += 1 @@ -216,13 +217,15 @@ class CPULLVMRenderer(LLVMRenderer): abi = 'win64cc' if sys.platform == 'win32' else None string_rewrite = base_rewrite def render(self, uops: list[UOp]) -> str: # the kernel is first, it is the entry. its functions follow, a name traced with other args gets a suffix - fns = {b: f"{b.arg}_{i}" for i, b in enumerate(b for b in UOp.sink(*uops).toposort() if b.op is Ops.LINEAR)} + fns = {b: f"{b.arg}_{i}" for i, b in enumerate(b for b in UOp.sink(*uops).toposort(enter_calls=True) if b.op is Ops.LINEAR)} defs = [self._render_kernel(b.src, name=n, fns=fns)[1] for b, n in fns.items()] return "\n".join((k:=self._render_kernel(uops, fns=fns))[0] + (k[1], *defs, self._render_footer(uops))) def _render_footer(self, uops: list[UOp]) -> str: decls = {x.src[0].arg.name: f"declare {ldt(x.dtype)} @{x.src[0].arg.name}({', '.join(map(lparam, x.src[1:]))})" - for x in UOp.sink(*uops).toposort() if x.op is Ops.CALL and x.src[0].op is Ops.CUSTOM_FUNCTION} - return "\n".join([*decls.values(), 'attributes #0 = { alwaysinline nounwind "no-builtins" "no-trapping-math"="true" }']) + for x in UOp.sink(*uops).toposort(enter_calls=True) if x.op is Ops.CALL and x.src[0].op is Ops.CUSTOM_FUNCTION} + blobs = {x.key: f"@bin_{x.key.hex()} = private constant [{len(x.arg)} x i8] c\"" + "".join("\\%02X" % b for b in x.arg) + '", align 16' + for x in UOp.sink(*uops).toposort(enter_calls=True) if x.op is Ops.BINARY} + return "\n".join([*decls.values(), *blobs.values(), 'attributes #0 = { alwaysinline nounwind "no-builtins" "no-trapping-math"="true" }']) def __init__(self, target:Target): super().__init__(target) from tinygrad.runtime.support.compiler_llvm import CPULLVMCompiler @@ -285,7 +288,7 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n %trunc f'"amdgpu-flat-work-group-size"="1,{requiredMaxThreadsPerBlock}"', '"no-trapping-math"="true"'] return 'attributes #0 = { ' + ' '.join(attributes) + ' }' @staticmethod - def is_rdna4(arch): return arch.split(':')[0] in {'gfx1200', 'gfx1201'} + def is_rdna4(arch): return arch in {'gfx1200', 'gfx1201'} def __init__(self, target:Target): super().__init__(target) from tinygrad.runtime.support.compiler_llvm import AMDLLVMCompiler diff --git a/tinygrad_repo/tinygrad/renderer/nir.py b/tinygrad_repo/tinygrad/renderer/nir.py index 4b23f6629..f6a07b32b 100644 --- a/tinygrad_repo/tinygrad/renderer/nir.py +++ b/tinygrad_repo/tinygrad/renderer/nir.py @@ -158,7 +158,7 @@ class NIRRenderer(Renderer): (UPat(Ops.STACK, name="x"), lambda ctx,x: nalu(ctx.b, f"vec{x.max_numel()}", *[ctx.r[src] for src in x.src])), (UPat(GroupOp.ALU, name="x"), lambda ctx,x: nalu(ctx.b, aop[x.src[0].dtype][x.op], *[ctx.r[src] for src in x.src])), (UPat(Ops.CAST, name="x"), lambda ctx,x: ncast(ctx.b, ctx.r[x.src[0]], x.src[0].dtype, x.dtype)), - (UPat(Ops.BITCAST, src=(UPat.var("a"),), allow_any_len=True), lambda ctx,a: ctx.r[a]), + (UPat(Ops.BITCAST, src=(UPat.var("a"),)), lambda ctx,a: ctx.r[a]), (UPat(Ops.BUFFER, name="x"), lambda ctx,x: mesa.nir_local_variable_create(ctx.b.impl, mesa.glsl_array_type(glsl_type(x.dtype), x.max_numel(), 0).contents, f"acc{x.arg.slot}".encode()).contents), (UPat(Ops.BARRIER), lambda ctx: nbarrier(ctx.b)), diff --git a/tinygrad_repo/tinygrad/renderer/ptx.py b/tinygrad_repo/tinygrad/renderer/ptx.py index 0b6db72c8..249cdc421 100644 --- a/tinygrad_repo/tinygrad/renderer/ptx.py +++ b/tinygrad_repo/tinygrad/renderer/ptx.py @@ -90,7 +90,7 @@ string_rewrite = PatternMatcher([ (UPat((Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ), name="x", allow_any_len=True, src=(UPat.var("src0"),)), lambda ctx, x, src0: ctx.code_for_op[x.op](ctx.r[x], *[ctx.r[v] for v in x.src], src0.dtype, ctx.types[src0.dtype])), (UPat(GroupOp.ALU, name="x"), lambda ctx, x: ctx.code_for_op[x.op](ctx.r[x], *[ctx.r[v] for v in x.src], x.dtype, ctx.types[x.dtype])), - (UPat(Ops.BITCAST, name="x", src=(UPat.var("a"),), allow_any_len=True), lambda ctx, x, a: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {ctx.r[a]};"), + (UPat(Ops.BITCAST, name="x", src=(UPat.var("a"),)), lambda ctx, x, a: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {ctx.r[a]};"), (UPat(Ops.CAST, name="x", src=(UPat(dtype=dtypes.bool, name="a"),)), lambda ctx, x, a: f"selp.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(1, x.dtype)}, {render_val(0, x.dtype)}, {ctx.r[a]};"), (UPat(Ops.CAST, name="x", src=(UPat.var("a"),)), diff --git a/tinygrad_repo/tinygrad/renderer/tc.py b/tinygrad_repo/tinygrad/renderer/tc.py index fa59f50c0..52b590460 100644 --- a/tinygrad_repo/tinygrad/renderer/tc.py +++ b/tinygrad_repo/tinygrad/renderer/tc.py @@ -78,7 +78,7 @@ amd_rdna3 = [TensorCore(dtype_in=di, dtype_out=do, frag_a=(("m0", "m1", "m2", "m amd_rdna4 = [TensorCore(dtype_in=di, dtype_out=do, frag_a=(("m0", "m1", "m2", "m3", "k2"), ("k0", "k1", "k3")), frag_b=(("n0", "n1", "n2", "n3", "k2"), ("k0", "k1", "k3")), frag_c=(("n0", "n1", "n2", "n3", "m3"), ("m0", "m1", "m2"))) for di,do in [(dtypes.half,dtypes.float),(dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float),(dtypes.bfloat16,dtypes.bfloat16), - (dtypes.fp8e4m3,dtypes.float),(dtypes.fp8e5m2,dtypes.float)]] + (dtypes.fp8e4m3,dtypes.float),(dtypes.fp8e5m2,dtypes.float),(dtypes.int8,dtypes.int32)]] # https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-cdna4-instruction-set-architecture.pdf def mfma(K:int, di:DType, do:DType) -> TensorCore: @@ -120,6 +120,9 @@ def pm_wmma_fp8(dtype:DType) -> PatternMatcher: ]) pm_validate_wmma_rdna4 = PatternMatcher([ + (UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace( + src=(x.src[0].bitcast(dtypes.int32), x.src[1].bitcast(dtypes.int32), x.src[2])) + if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 8 else None), (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace( src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16))) .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None), diff --git a/tinygrad_repo/tinygrad/renderer/wgsl.py b/tinygrad_repo/tinygrad/renderer/wgsl.py index bc34e7f63..ebf762e71 100644 --- a/tinygrad_repo/tinygrad/renderer/wgsl.py +++ b/tinygrad_repo/tinygrad/renderer/wgsl.py @@ -16,8 +16,9 @@ def packed_store(s:UOp): bidx, var, *gate = s.src idx, shift_am, mask = packed_field(bidx, var.dtype) new_v, wmask = (var.cast(dtypes.uint32) & mask) << shift_am, ((mask << shift_am) ^ 0xFFFFFFFF).cast(dtypes.uint32) - buf = idx.cast(dtypes.uint32).load(*((UOp.const(0, dtypes.uint32), *gate) if gate else ())) - return idx.store((buf & wmask) | new_v, *gate) + word = idx.cast(dtypes.uint32) + buf = word.load(*((UOp.const(0, dtypes.uint32), *gate) if gate else ())) + return word.store((buf & wmask) | new_v, *gate) # load for char: sign_extend(buf[idx/4] >> ((idx%4)*8)) def packed_load(root:UOp): @@ -85,7 +86,7 @@ class WGSLRenderer(CStyleLanguage): lambda ctx,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b], b.src[0])}, {ctx[gate]})"), (UPat.load(UPat.var("b")), lambda ctx, b: ctx.render_load(ctx[b], b)), # packed_store writes (load & wmask) | new_v: atomicAnd clears the field, atomicAdd sets it. new_v is gone when it is 0 - (UPat.store(UPat.var("b"), UPat.any(packed_rmw, packed_rmw | UPat.var("nv"))), lambda ctx,b,wmask,nv=None: + (UPat.store(UPat(Ops.CAST, src=(UPat.var("b"),)), UPat.any(packed_rmw, packed_rmw | UPat.var("nv"))), lambda ctx,b,wmask,nv=None: f"atomicAnd(&{ctx[b]},{ctx[wmask]});"+(f"\n atomicAdd(&{ctx[b]},{ctx[nv]});" if nv is not None else "") if is_packed(b) else None), (UPat.store(UPat.var("b"), UPat.var("v")), lambda ctx,b,v: f"{ctx[b]} = {ctx[v]};"), (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"))), diff --git a/tinygrad_repo/tinygrad/runtime/ops_amd.py b/tinygrad_repo/tinygrad/runtime/ops_amd.py index a5df720ac..853c5e317 100644 --- a/tinygrad_repo/tinygrad/runtime/ops_amd.py +++ b/tinygrad_repo/tinygrad/runtime/ops_amd.py @@ -164,7 +164,7 @@ class AMDComputeQueue(HWQueue): memsel_dw = self.pm4.DATA_SEL(data_sel) | self.pm4.INT_SEL(int_sel) ctxid = 0 addr_w = address if isinstance(address, UOp) else UOp.const(address, dtypes.uint64) - val_w = value.cast(dtypes.uint64) if isinstance(value, UOp) else UOp.const(value, dtypes.uint64) + val_w = UOp.const(value, dtypes.uint64) self.pkt3(self.pm4.PACKET3_RELEASE_MEM, event_dw | cache_flags_dw, memsel_dw, addr_w, val_w, ctxid) def memory_barrier(self): diff --git a/tinygrad_repo/tinygrad/runtime/ops_cuda.py b/tinygrad_repo/tinygrad/runtime/ops_cuda.py index 8f8a81d13..8d74851e2 100644 --- a/tinygrad_repo/tinygrad/runtime/ops_cuda.py +++ b/tinygrad_repo/tinygrad/runtime/ops_cuda.py @@ -1,5 +1,6 @@ from __future__ import annotations import ctypes, functools, mmap, struct, time +from typing import Any from tinygrad.helpers import DEBUG, DEV, getenv, unwrap from tinygrad.device import Buffer, BufferStorage, BufferSpec, Allocator, Compiled, MMIOInterface, HCQ_RUNTIME_DEV from tinygrad.dtype import dtypes @@ -20,7 +21,7 @@ def check(status:int): def as_int(handle) -> int: return unwrap(ctypes.cast(handle, ctypes.c_void_p).value) def host_stamp(slot:int): ctypes.c_uint64.from_address(slot).value = time.perf_counter_ns() -def extern(ptr:int, meta=None) -> Buffer: return Buffer(HCQ_RUNTIME_DEV.value, 8, opaque=BufferStorage(ptr, meta)) +def extern(ptr:int, meta=None) -> Buffer: return Buffer(HCQ_RUNTIME_DEV.device, 8, opaque=BufferStorage(ptr, meta)) # ***************** # queue @@ -47,7 +48,7 @@ class CUDAQueue(HWQueue): addr = UOp(Ops.LINEAR, src=tuple(pack_args([(0, UOp.const(size, dtypes.uint64))] + rows, 8 + size)), arg="kernargs").getaddr(self.devs) # extra: [buffer pointer, &args, buffer size, &size, end] self.call(cuda.cuLaunchKernel, self.extern((self.dev.tag("function"), obj.lib, obj.name)), *prg.arg.global_size, *prg.arg.local_size, 0, None, 0, - (self.q(*[UOp.const(v, dtypes.uint64) if isinstance(v, int) else v for v in (1, addr + 8, 2, addr, 0)]) - 40,)) + (self.q(*[UOp.const(v, dtypes.uint64) if isinstance(v, int) else v for v in (1, addr + 8, 2, addr, self.dev.launch_end)]) - 40,)) def copy(self, dst:UOp, src:UOp, sz:int): self.call(cuda.cuMemcpyAsync, dst.getaddr(self.devs), src.getaddr(self.devs), sz, None) def wait(self, sig:UOp, val:UOp): self.call(cuda.cuStreamWaitValue64_v2, None, sig.getaddr(self.devs), val, cuda.CU_STREAM_WAIT_VALUE_GEQ) @@ -64,29 +65,32 @@ class CUDAQueue(HWQueue): class CUDAAllocator(Allocator['CUDADevice']): def _alloc(self, size:int, options:BufferSpec) -> BufferStorage: - check(cuda.cuCtxSetCurrent(self.dev.context)) + check((api:=self.dev.api).cuCtxSetCurrent(self.dev.context)) if options.external_ptr: return BufferStorage(options.external_ptr) if options.host or options.cpu_access: - ptr = init_c_var(ctypes.c_void_p, lambda x: check(cuda.cuMemHostAlloc(ctypes.byref(x), size, 0))).value + ptr = init_c_var(ctypes.c_void_p, lambda x: check(api.cuMemHostAlloc(ctypes.byref(x), size, 0))).value return BufferStorage(ptr, None, MMIOInterface(ptr, size)) - return BufferStorage(init_c_var(cuda.CUdeviceptr, lambda x: check(cuda.cuMemAlloc_v2(ctypes.byref(x), size))).value) + return BufferStorage(init_c_var(api.CUdeviceptr, lambda x: check(api.cuMemAlloc_v2(ctypes.byref(x), size))).value) def _free(self, storage:BufferStorage, options:BufferSpec): self.dev.synchronize() - check((cuda.cuMemFreeHost if options.host or options.cpu_access else cuda.cuMemFree_v2)(storage.buf)) + check((self.dev.api.cuMemFreeHost if options.host or options.cpu_access else self.dev.api.cuMemFree_v2)(storage.buf)) def _map(self, buf:Buffer) -> BufferStorage: - if buf.device.startswith("CUDA"): return BufferStorage(buf._buf) + if buf.device.split(":")[0] == self.dev.device.split(":")[0]: return BufferStorage(buf._buf) if (host:=buf.get_storage().host) is None or host.addr % mmap.PAGESIZE: raise RuntimeError(f"{buf.device} memory is not page aligned host memory") - check(cuda.cuCtxSetCurrent(self.dev.context)) - if (status:=cuda.cuMemHostRegister_v2(host.addr, buf.nbytes, 0)) != cuda.CUDA_ERROR_HOST_MEMORY_ALREADY_REGISTERED: + check((api:=self.dev.api).cuCtxSetCurrent(self.dev.context)) + if (status:=api.cuMemHostRegister_v2(host.addr, buf.nbytes, 0)) != cuda.CUDA_ERROR_HOST_MEMORY_ALREADY_REGISTERED: check(status) # another device pinned it already return BufferStorage(host.addr, status == cuda.CUDA_SUCCESS) def _unmap(self, mapping:BufferStorage): - if mapping.meta: check(cuda.cuMemHostUnregister(mapping.buf)) + if mapping.meta: check(self.dev.api.cuMemHostUnregister(mapping.buf)) def _offset(self, buf:int, size:int, offset:int) -> int: return buf + offset class CUDADevice(Compiled): + api:Any = cuda + launch_end = 0 # CU_LAUNCH_PARAM_END + renderers = [CUDARenderer, PTXRenderer, NVCCRenderer] pm_encode = PatternMatcher([ (UPat(Ops.CALL, src=(UPat.custom_function("submit_cuda_compute"), UPat()), name="s"), lambda s: CUDAQueue(s).encode()), (UPat(Ops.CALL, src=(UPat.custom_function("submit_cuda_copy"), UPat()), name="s"), lambda s: CUDAQueue(s).encode()), @@ -94,34 +98,37 @@ class CUDADevice(Compiled): def __init__(self, device:str=""): device_id = int(device.split(":")[1]) if ":" in device else 0 - check(cuda.cuInit(0)) - self.cu_device = init_c_var(cuda.CUdevice, lambda x: check(cuda.cuDeviceGet(ctypes.byref(x), device_id))) - self.context = init_c_var(cuda.CUcontext, lambda x: check(cuda.cuCtxCreate_v2(ctypes.byref(x), 0, self.cu_device))) - check(cuda.cuDeviceComputeCapability(ctypes.byref(major:=ctypes.c_int()), ctypes.byref(minor:=ctypes.c_int()), device_id)) - self.streams = [init_c_var(cuda.CUstream, lambda x: check(cuda.cuStreamCreate(ctypes.byref(x), cuda.CU_STREAM_NON_BLOCKING))) for _ in range(2)] - super().__init__(device, CUDAAllocator(self), [CUDARenderer, PTXRenderer, NVCCRenderer], None, arch=f"sm_{major.value}{minor.value}") + check((api:=self.api).cuInit(0)) + self.cu_device = init_c_var(api.CUdevice, lambda x: check(api.cuDeviceGet(ctypes.byref(x), device_id))) + self.context = init_c_var(api.CUcontext, lambda x: check(api.cuCtxCreate_v2(ctypes.byref(x), 0, self.cu_device))) + self.streams = [init_c_var(api.CUstream, lambda x: check(api.cuStreamCreate(ctypes.byref(x), cuda.CU_STREAM_NON_BLOCKING))) for _ in range(2)] + super().__init__(device, CUDAAllocator(self), self.renderers, None, arch=self.get_arch(device_id)) Compiled.pm_bufferize += PatternMatcher([ (UPat(Ops.ALLOC, tag=self.tag("cuda")), lambda d=self: d.handles), (UPat(Ops.ALLOC, tag=self.tag("stamp")), lambda d=self: d.stamp), (UPat(Ops.ALLOC, name="b"), lambda b, d=self: d.function(*b.tag[1:]) if isinstance(b.tag, tuple) and b.tag[0] == d.tag("function") else None)]) + def get_arch(self, device_id:int) -> str: + check(cuda.cuDeviceComputeCapability(ctypes.byref(major:=ctypes.c_int()), ctypes.byref(minor:=ctypes.c_int()), device_id)) + return f"sm_{major.value}{minor.value}" + @functools.cached_property def handles(self) -> Buffer: - return Buffer(HCQ_RUNTIME_DEV.value, 32, initial_value=struct.pack("4Q", *[as_int(h) for h in (self.context, *self.streams)], 0)) + return Buffer(HCQ_RUNTIME_DEV.device, 32, initial_value=struct.pack("4Q", *[as_int(h) for h in (self.context, *self.streams)], 0)) @functools.cached_property def stamp(self) -> Buffer: return extern(as_int(fn:=cuda.CUhostFn(host_stamp)), fn) @functools.cache def function(self, lib:bytes, name:str) -> Buffer: - if DEBUG >= 5: print("\n".join([f"{i+1:>3} {line}" for i, line in enumerate(pretty_ptx(lib.decode()).split("\n"))])) - check(cuda.cuCtxSetCurrent(self.context)) - check(cuda.cuModuleLoadData(ctypes.byref(module:=cuda.CUmodule()), lib)) - check(cuda.cuModuleGetFunction(ctypes.byref(func:=cuda.CUfunction()), module, name.encode())) + if DEBUG >= 5 and self.api is cuda: print("\n".join([f"{i+1:>3} {line}" for i, line in enumerate(pretty_ptx(lib.decode()).split("\n"))])) + check(self.api.cuCtxSetCurrent(self.context)) + check(self.api.cuModuleLoadData(ctypes.byref(module:=self.api.CUmodule()), lib)) + check(self.api.cuModuleGetFunction(ctypes.byref(func:=self.api.CUfunction()), module, name.encode())) return extern(as_int(func)) - def count(self) -> int: return init_c_var(ctypes.c_int, lambda x: check(cuda.cuDeviceGetCount(ctypes.byref(x)))).value + def count(self) -> int: return init_c_var(ctypes.c_int, lambda x: check(self.api.cuDeviceGetCount(ctypes.byref(x)))).value def _wait_signal(self, sig:MMIOInterface|memoryview, value:int, timeout:int|None=None): # a fault raises here - check(cuda.cuCtxSetCurrent(self.context)) - check(cuda.cuCtxSynchronize()) + check(self.api.cuCtxSetCurrent(self.context)) + check(self.api.cuCtxSynchronize()) diff --git a/tinygrad_repo/tinygrad/runtime/ops_disk.py b/tinygrad_repo/tinygrad/runtime/ops_disk.py index 424feaad9..fa2fe1c36 100644 --- a/tinygrad_repo/tinygrad/runtime/ops_disk.py +++ b/tinygrad_repo/tinygrad/runtime/ops_disk.py @@ -1,22 +1,22 @@ -import os, sys, mmap, io, ctypes, contextlib, pathlib -from typing import Generator, Callable -from tinygrad.helpers import OSX, round_up, mv_address -from tinygrad.device import BufferStorage, MMIOInterface, Compiled, Allocator +import os, sys, mmap, io, ctypes, contextlib, pathlib, functools, collections, struct, itertools +from dataclasses import replace +from typing import cast +from tinygrad.helpers import OSX, mv_address, flatten, to_tuple, unwrap, ceildiv +from tinygrad.device import BufferStorage, MMIOInterface, Compiled, Allocator, Buffer, BufferSpec, Device, HCQ_RUNTIME_DEV +from tinygrad.dtype import dtypes +from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, uopfunc +from tinygrad.runtime.support.hcq2 import ccall, ins, patch with contextlib.suppress(ImportError): import _posixshmem from tinygrad.runtime.autogen import io_uring, libc class DiskDevice(Compiled): - _tried_io_uring_init = False - def synchronize(self, timeout:int|None=None): pass def __init__(self, device:str): - if not DiskDevice._tried_io_uring_init: self._iouring_setup() - self.size: int|None = None self.fd: int|None = None - self.refcount = 0 + self.refcount, self.info = 0, Buffer(HCQ_RUNTIME_DEV.device, 16, options=BufferSpec(nolru=True)) # info: [fd, mmap address] for batch reads super().__init__(device, DiskAllocator(self), [], None) def _might_open(self, size:int): assert self.size is None or size <= self.size, f"can't reopen Disk tensor with larger size, opened with {self.size}, tried to open with {size}" @@ -34,6 +34,7 @@ class DiskDevice(Compiled): except OSError: self.fd = os.open(filename, os.O_RDWR|os.O_CREAT) if not pathlib.Path(filename).is_block_device() and os.fstat(self.fd).st_size < size: os.ftruncate(self.fd, size) self.mem = mmap.mmap(self.fd, size) + self.info.host.view()[:16] = struct.pack("2Q", self.fd, mv_address(memoryview(self.mem))) self.size = size if hasattr(self.mem, 'madvise') and (hp := getattr(mmap, "MADV_HUGEPAGE", None)) is not None: with contextlib.suppress(OSError): self.mem.madvise(hp) # some systems have transparent_hugepage disabled @@ -41,34 +42,11 @@ class DiskDevice(Compiled): def _might_close(self): self.refcount -= 1 if self.refcount == 0: - if self.fd is not None: - os.close(self.fd) + if self.fd is not None: os.close(self.fd) if hasattr(self, "mem"): try: self.mem.close() except BufferError: pass self.size = None - def _iouring_setup(self): - DiskDevice._tried_io_uring_init = True - - if sys.platform == 'linux' and not hasattr(sys, "getandroidapilevel"): - fd = libc.syscall(io_uring.NR_io_uring_setup, 4096, ctypes.byref(p:=io_uring.struct_io_uring_params())) - if fd < 0: return - - sq_ptr = libc.mmap(0, p.sq_off.array + p.sq_entries * 4, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | MAP_POPULATE, fd, 0) - cq_ptr = libc.mmap(0, p.cq_off.cqes + p.cq_entries * ctypes.sizeof(io_uring.struct_io_uring_cqe), - mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | MAP_POPULATE, fd, io_uring.IORING_OFF_CQ_RING) - sqes = libc.mmap(0, p.sq_entries * ctypes.sizeof(io_uring.struct_io_uring_sqe), - mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | MAP_POPULATE, fd, io_uring.IORING_OFF_SQES) - - def u32ptr(val): return ctypes.cast(val, ctypes.POINTER(ctypes.c_uint32)) - sqdesc = io_uring.struct_io_uring_sq(khead=u32ptr(sq_ptr+p.sq_off.head), ktail=u32ptr(sq_ptr+p.sq_off.tail), - array=u32ptr(sq_ptr+p.sq_off.array), - kring_mask=u32ptr(sq_ptr+p.sq_off.ring_mask), sqes=ctypes.cast(sqes, ctypes.POINTER(io_uring.struct_io_uring_sqe))) - - cqdesc = io_uring.struct_io_uring_cq(khead=u32ptr(cq_ptr+p.cq_off.head), ktail=u32ptr(cq_ptr+p.cq_off.tail), - kring_mask=u32ptr(sq_ptr+p.cq_off.ring_mask), cqes=ctypes.cast(cq_ptr+p.cq_off.cqes, ctypes.POINTER(io_uring.struct_io_uring_cqe))) - - DiskDevice.io_uring = io_uring.struct_io_uring(ring_fd=fd, sq=sqdesc, cq=cqdesc) # type: ignore class DiskBuffer: def __init__(self, device:DiskDevice, size:int, offset=0): @@ -98,47 +76,95 @@ class DiskAllocator(Allocator): else: dest[:] = src._buf() - def _copyout_sharded(self, src:DiskBuffer, size:int, _get_free_buf:Callable, seg_len:int, - use_ioring:bool=True) -> Generator[tuple[int, int, int, int], None, None]: - fd_offset = src.offset - (minor_offset := src.offset % mmap.PAGESIZE) - processed_reqs_cnt, copied_in, next_read_offset, total_copy_size = 0, 0, 0, round_up(size + minor_offset, mmap.PAGESIZE) - - if not hasattr(DiskDevice, 'io_uring') or not use_ioring: - local_buf = memoryview(bytearray(seg_len)) - for off in range(0, total_copy_size, seg_len): - while (copy_batch := _get_free_buf()) is None: pass - read_size = min(seg_len, total_copy_size - off, src.device.size - fd_offset - off) - self._copyout(local_buf[:read_size], DiskBuffer(src.device, read_size, fd_offset + off)) - copy_batch[0].view(size=read_size)[:] = local_buf[:read_size] - real_copy_size = min(read_size - minor_offset, size - copied_in) - yield (copy_batch, copied_in, minor_offset, real_copy_size) - copied_in, minor_offset = copied_in + real_copy_size, 0 - return - - reqs: list[tuple[int, int, int, int]] = [] - while next_read_offset < total_copy_size or len(reqs) != processed_reqs_cnt: - if next_read_offset < total_copy_size and (copy_batch := _get_free_buf()) is not None: - # Prepare sqe - sqe_index = (tail:=DiskDevice.io_uring.sq.ktail[0]) & DiskDevice.io_uring.sq.kring_mask[0] - sqe = DiskDevice.io_uring.sq.sqes[sqe_index] - sqe.opcode, sqe.fd, sqe.off = io_uring.IORING_OP_READ, self.dev.fd, fd_offset + next_read_offset - sqe.addr, sqe.len, sqe.user_data = copy_batch[0].addr, min(seg_len, total_copy_size - next_read_offset), len(reqs) - - # Send sqe - DiskDevice.io_uring.sq.array[sqe_index] = sqe_index - DiskDevice.io_uring.sq.ktail[0] = tail + 1 - libc.syscall(io_uring.NR_io_uring_enter, DiskDevice.io_uring.ring_fd, 1, 1, io_uring.IORING_ENTER_GETEVENTS) - - reqs.append((copy_batch, copied_in, minor_offset, real_copy_size:=min(sqe.len - minor_offset, size - copied_in))) - next_read_offset += sqe.len - copied_in += real_copy_size - minor_offset = 0 - - if (head:=DiskDevice.io_uring.cq.khead[0]) != DiskDevice.io_uring.cq.ktail[0]: - cqe = DiskDevice.io_uring.cq.cqes[head & DiskDevice.io_uring.cq.kring_mask[0]] - assert cqe.res >= 0, f"read from disk failed, err: {cqe.res}" - yield reqs[cqe.user_data] - DiskDevice.io_uring.cq.khead[0] = head + 1 # advance - processed_reqs_cnt += 1 - def _offset(self, buf:DiskBuffer, size:int, offset:int): return DiskBuffer(buf.device, size, offset) + +# ***************** +# UOps implementation + +CHUNK_SZ, READ_SZ, SLOTS, RING_ENTRIES = 64 << 20, 1 << 20, 2, 256 +SLOT_SZ, STAGE_SZ = CHUNK_SZ + 4096, 4096 + SLOTS * (CHUNK_SZ + 4096) # +4096: reads are page aligned. stage: a page of counts, then the slots + +@uopfunc +def disk_read(file:UOp, srcs:UOp, counts:UOp, rings:UOp, sqes:UOp) -> UOp: # srcs: [address in the file's mmap, bytes] per copy + fd, p = unwrap(uring())[:2] + file, srcs, counts, rings, sqes = (b.replace(arg=replace(b.arg, volatile=True, device=None)) for b in (file, srcs, counts, rings, sqes)) + + # for each copy + src_va, nbytes = srcs[(copy:=UOp.range(srcs.max_numel() // 2, 0, dtype=dtypes.uint64)) * 2], srcs[copy * 2 + 1] + + # for each chunk: the next chunk n waits for the gpu to free its slot + n = counts.after(chunk:=UOp.range(ceildiv(nbytes, CHUNK_SZ), 1, dtype=dtypes.uint64))[0] + free = (copied:=counts.after(n, loop:=UOp.loop(3))[1]).backedge(loop, copied + SLOTS <= n) + pos = src_va + chunk * CHUNK_SZ - file[1] + span = (pos % 4096 + (nbytes - chunk * CHUNK_SZ).minimum(CHUNK_SZ) + 4095) // 4096 * 4096 + + # for each read: queues READ_SZ of the chunk's pages. an sqe is 8 words: fd << 32 | flags << 8 | opcode, file offset, address, bytes + tail, read = rings.after(free)[p.sq_off.tail // 4].cast(dtypes.uint64), UOp.range(nreads:=ceildiv(span, READ_SZ), 2, dtype=dtypes.uint64) + sqe = [file[0] << 32 | io_uring.IOSQE_ASYNC << 8 | io_uring.IORING_OP_READ, pos - pos % 4096 + read * READ_SZ, + counts.getaddr(HCQ_RUNTIME_DEV.device) + 4096 + n % SLOTS * SLOT_SZ + read * READ_SZ, (span - read * READ_SZ).minimum(READ_SZ)] + queued = UOp.group(*[sqes[(tail + read) % RING_ENTRIES * 8 + j].store(v) for j, v in enumerate(sqe)]).end(read) + + # submits the reads, waits for all, marks the chunk read + to_submit = nreads.cast(dtypes.int).after(rings.after(queued)[p.sq_off.tail // 4].store((tail + nreads).cast(dtypes.uint32))) + done = ccall(libc.syscall, io_uring.NR_io_uring_enter, fd, to_submit, to_submit, io_uring.IORING_ENTER_GETEVENTS, 0, 0) + reaped = rings.after(done)[p.cq_off.head // 4].store(rings.after(done)[p.cq_off.tail // 4]) + return counts.after(reaped)[0].store(n + 1).end(chunk).end(copy).sink() + +# ***************** +# 2. rewriter + +def is_disk_read(c:UOp) -> bool: + dev, src = [to_tuple(b.device)[0] for b in c.src[1:3]] if c.op is Ops.CALL and c.body.op is Ops.STORE else ("", "") + return src.startswith("DISK:") and not src.startswith("DISK:shm:") and Device[dev].has_copy_queue and not dev.startswith("NULL") and bool(uring()) + +def disk_copy_rewriter(s:UOp) -> UOp|None: + lins = [submit.without_after.src[1].without_after for submit in s.src] + if not (copies:=[c for lin in lins for c in lin.src if is_disk_read(c)]): return None + rings = UOp.alloc((unwrap(uring())[2].nbytes // 4,), dtypes.uint32, 0, device=HCQ_RUNTIME_DEV.device).rtag("uring_rings") + sqes = UOp.alloc((8 * RING_ENTRIES,), dtypes.uint64, 0, device=HCQ_RUNTIME_DEV.device).rtag("uring_sqes") + + # [mmap address, bytes] per copy in a table: an arg each would hit ctypes' 1024 limit + words = flatten((c.src[2].getaddr(HCQ_RUNTIME_DEV.device), UOp.const(c.src[2].nbytes(), dtypes.uint64)) for c in copies) + table = UOp.alloc((8 * len(words),), dtypes.uint8, device=HCQ_RUNTIME_DEV.device) + srcs = patch(table, [(8 * i, w) for i, w in enumerate(words)]).bitcast(dtypes.uint64) + + # per gpu: [chunks read, chunks copied] at the start of its stage + devs = [to_tuple(c.src[1].device)[0] for c in copies] + counts = {d: UOp.alloc((STAGE_SZ,), dtypes.uint8, 0, device=d).rtag("disk_stage")[:16].bitcast(dtypes.uint64) for d in devs} + + # cpu part + runs = [(key, list(ks)) for key, ks in itertools.groupby(range(len(copies)), lambda k: (devs[k], to_tuple(copies[k].src[2].device)[0]))] + done, offs, gpu_ops = s.src[-1], collections.Counter[str](), collections.defaultdict[UOp, list[UOp]](list) + for (dev, disk), ks in runs: + file = UOp.alloc((2,), dtypes.uint64, 0, device=disk).rtag("disk_info") + done = disk_read(file, srcs[2 * ks[0]:2 * ks[-1] + 2], counts[dev].after(done), rings, sqes) + + # gpu part + for k, (c, dev) in enumerate(zip(copies, devs)): + first, slots = counts[dev].index(0).load() + offs[dev], counts[dev].getaddr(dev) + 4096 + srcs.index(2 * k).load() % 4096 + offs[dev] += ceildiv(c.src[2].nbytes(), CHUNK_SZ) + for chunk in range(ceildiv(c.src[2].nbytes(), CHUNK_SZ)): + n, dst, nb = first + chunk, c.src[1].getaddr() + chunk * CHUNK_SZ, min(CHUNK_SZ, c.src[2].nbytes() - chunk * CHUNK_SZ) + gpu_ops[c] += [ins("wait", counts[dev], n + 1), ins("copy", dst, slots + n % SLOTS * SLOT_SZ, nb), ins("store", counts[dev][1:], n + 1)] + return s.replace(src=(*s.src, done)).substitute({lin: lin.replace(src=tuple(flatten(gpu_ops.get(c, [c]) for c in lin.src))) for lin in lins}) +Compiled.pm_batch = Compiled.pm_batch + PatternMatcher([(UPat(Ops.SINK, name="s"), disk_copy_rewriter)]) + +# ***************** +# 3. bufferize + +def shared(fd:int, at:int, n:int) -> Buffer: # n bytes of memory the kernel shares at offset at + return Buffer(HCQ_RUNTIME_DEV.device, n, options=BufferSpec(external_ptr=libc.mmap(0, n, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED, fd, at))) + +@functools.cache +def uring() -> tuple[int, io_uring.struct_io_uring_params, Buffer, Buffer]|None: # fd, offsets, the rings (heads, tails, cqes), the sqes + if sys.platform != "linux" or hasattr(sys, "getandroidapilevel"): return None + p = io_uring.struct_io_uring_params(flags=io_uring.IORING_SETUP_NO_SQARRAY) # type: ignore[unreachable] + if (fd:=libc.syscall(io_uring.NR_io_uring_setup, RING_ENTRIES, ctypes.byref(p))) < 0: return None + return fd, p, shared(fd, io_uring.IORING_OFF_SQ_RING, p.cq_off.cqes + 16 * p.cq_entries), shared(fd, io_uring.IORING_OFF_SQES, 64 * RING_ENTRIES) + +@functools.cache +def _stage(dev) -> Buffer: return Buffer(Device[dev].host, STAGE_SZ, options=BufferSpec(cpu_access=True), initial_value=bytes(STAGE_SZ)) + +Compiled.pm_bufferize += PatternMatcher([(UPat(Ops.ALLOC, tag="disk_info", name="b"), lambda b: cast(DiskDevice, Device[b.device]).info), + (UPat(Ops.ALLOC, tag="disk_stage", name="b"), lambda b: _stage(b.device)), + (UPat(Ops.ALLOC, tag="uring_rings"), lambda: unwrap(uring())[2]), (UPat(Ops.ALLOC, tag="uring_sqes"), lambda: unwrap(uring())[3])]) diff --git a/tinygrad_repo/tinygrad/runtime/ops_hip.py b/tinygrad_repo/tinygrad/runtime/ops_hip.py index b5fe55502..495c68367 100644 --- a/tinygrad_repo/tinygrad/runtime/ops_hip.py +++ b/tinygrad_repo/tinygrad/runtime/ops_hip.py @@ -1,71 +1,35 @@ -import ctypes -from tinygrad.helpers import mv_address, getenv, suppress_finalizing -from tinygrad.device import BufferStorage, Compiled, Allocator, BufferSpec, Program, TinyELF -from tinygrad.runtime.autogen import hip +from types import SimpleNamespace +from tinygrad.helpers import getenv +from tinygrad.dtype import dtypes +from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite from tinygrad.renderer.cstyle import HIPRenderer -from tinygrad.runtime.support.c import init_c_var, init_c_struct_t +from tinygrad.runtime.autogen import hip +from tinygrad.runtime.support.hcq2 import ccall +from tinygrad.runtime.support.c import init_c_var +from tinygrad.runtime.ops_cuda import CUDAQueue, CUDADevice, check if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl # noqa: F401 # pylint: disable=unused-import -def check(status): - if status != 0: raise RuntimeError(f"HIP Error {status}, {ctypes.string_at(hip.hipGetErrorString(status)).decode()}") +HIPIFY = dict(CUdevice=hip.hipDevice_t, CUcontext=hip.hipCtx_t, CUstream=hip.hipStream_t, CUdeviceptr=hip.hipDeviceptr_t, CUmodule=hip.hipModule_t, + CUfunction=hip.hipFunction_t, cuInit=hip.hipInit, cuDeviceGet=hip.hipDeviceGet, cuDeviceGetCount=hip.hipGetDeviceCount, + cuCtxCreate_v2=hip.hipCtxCreate, cuCtxSetCurrent=hip.hipCtxSetCurrent, cuCtxSynchronize=hip.hipDeviceSynchronize, + cuStreamCreate=hip.hipStreamCreateWithFlags, cuMemAlloc_v2=hip.hipMalloc, cuMemFree_v2=hip.hipFree, + cuMemHostAlloc=lambda p, sz, _: hip.hipHostMalloc(p, sz, hip.hipHostMallocCoherent), cuMemFreeHost=hip.hipHostFree, + cuMemHostRegister_v2=hip.hipHostRegister, cuMemHostUnregister=hip.hipHostUnregister, cuModuleLoadData=hip.hipModuleLoadData, + cuModuleGetFunction=hip.hipModuleGetFunction, cuLaunchKernel=hip.hipModuleLaunchKernel, cuMemcpyAsync=hip.hipMemcpyAsync, + cuStreamWaitValue64_v2=hip.hipStreamWaitValue64, cuStreamWriteValue64_v2=hip.hipStreamWriteValue64, cuLaunchHostFunc=hip.hipLaunchHostFunc) -class HIPDevice(Compiled): - def __init__(self, device:str=""): - self.device_id = int(device.split(":")[1]) if ":" in device else 0 - self.arch = init_c_var(hip.hipDeviceProp_t, lambda x: check(hip.hipGetDeviceProperties(x, self.device_id))).gcnArchName.decode() - self.time_event_st, self.time_event_en = [init_c_var(hip.hipEvent_t, lambda x: hip.hipEventCreate(ctypes.byref(x), 0)) for _ in range(2)] +def hipify(c:UOp, f:UOp) -> UOp|None: + if (fn:=HIPIFY.get(f.arg.name)) is None: return None + args = list(c.src[1:]) + if fn is hip.hipMemcpyAsync: args.insert(3, UOp.const(hip.hipMemcpyDefault, dtypes.int)) + if fn is hip.hipStreamWaitValue64: args.append(UOp.const((1<<64)-1, dtypes.uint64)) + return ccall(fn, *args) +pm_hipify = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, name="f"),), allow_any_len=True, name="c"), hipify)]) - super().__init__(device, HIPAllocator(self), [HIPRenderer], HIPProgram, arch=self.arch) +class HIPDevice(CUDADevice): + api, launch_end, renderers = SimpleNamespace(**HIPIFY), 3, [HIPRenderer] + pm_encode = PatternMatcher([(UPat(Ops.CALL, src=(UPat.custom_function(f"submit_hip_{q}"), UPat()), name="s"), + lambda s: graph_rewrite(CUDAQueue(s).encode(), pm_hipify, enter_calls=True, name="hipify")) for q in ("compute", "copy")]) - def count(self) -> int: return init_c_var(ctypes.c_int, lambda x: check(hip.hipGetDeviceCount(x))).value - - def synchronize(self, timeout:int|None=None): - check(hip.hipSetDevice(self.device_id)) - check(hip.hipDeviceSynchronize()) - -class HIPProgram(Program[HIPDevice]): - def __init__(self, dev:HIPDevice, obj:TinyELF): - self.dev, self.name, self.lib, self.signature = dev, obj.name, obj.lib, obj.signature - check(hip.hipSetDevice(self.dev.device_id)) - self.module = init_c_var(hip.hipModule_t, lambda x: check(hip.hipModuleLoadData(ctypes.byref(x), obj.lib))) - self.prg = init_c_var(hip.hipFunction_t, lambda x: check(hip.hipModuleGetFunction(ctypes.byref(x), self.module, obj.name.encode("utf-8")))) - - @suppress_finalizing - def __del__(self): - if hasattr(self, 'module'): check(hip.hipModuleUnload(self.module)) - - def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw): - check(hip.hipSetDevice(self.dev.device_id)) - if not hasattr(self, "vargs"): - fields = ([(f'f{i}', hip.hipDeviceptr_t, i*8) for i in range(len(args))] + - [(f'v{i}', getattr(ctypes, f"c_int{dt.bitsize}"), o) for i,(o,dt) in enumerate(TinyELF.iter_sig(self.signature[len(args):], len(args)*8))]) - self.c_args = init_c_struct_t(fields[-1][2] + ctypes.sizeof(fields[-1][1]) if len(fields) else 0, tuple(fields))(*args, *vals) - self.vargs = (ctypes.c_void_p * 5)(1, ctypes.cast(ctypes.byref(self.c_args), ctypes.c_void_p), 2, - ctypes.cast(ctypes.pointer(ctypes.c_size_t(ctypes.sizeof(self.c_args))), ctypes.c_void_p), 3) - - for i in range(len(args)): self.c_args.__setattr__(f'f{i}', args[i]) - for i in range(len(vals)): self.c_args.__setattr__(f'v{i}', vals[i]) - - if wait: check(hip.hipEventRecord(self.dev.time_event_st, None)) - - check(hip.hipModuleLaunchKernel(self.prg, *global_size, *local_size, 0, None, None, self.vargs)) - - if wait: - check(hip.hipEventRecord(self.dev.time_event_en, None)) - check(hip.hipEventSynchronize(self.dev.time_event_en)) - check(hip.hipEventElapsedTime(ctypes.byref(ret := ctypes.c_float()), self.dev.time_event_st, self.dev.time_event_en)) - return ret.value * 1e-3 - -class HIPAllocator(Allocator[HIPDevice]): - def _alloc(self, size:int, options:BufferSpec) -> BufferStorage: - check(hip.hipSetDevice(self.dev.device_id)) - return BufferStorage(init_c_var(hip.hipDeviceptr_t, lambda x: check(hip.hipMalloc(ctypes.byref(x), size)))) - - def _free(self, storage:BufferStorage, options:BufferSpec): check(hip.hipFree(storage.buf)) - def _copyin(self, dest, src: memoryview): - check(hip.hipSetDevice(self.dev.device_id)) - check(hip.hipMemcpy(dest, mv_address(src), len(src), hip.hipMemcpyHostToDevice)) - def _copyout(self, dest:memoryview, src): - self.dev.synchronize() - check(hip.hipMemcpy(mv_address(dest), src, len(dest), hip.hipMemcpyDeviceToHost)) - def _offset(self, buf, size:int, offset:int): return hip.hipDeviceptr_t(buf.value + offset) + def get_arch(self, device_id:int) -> str: + return init_c_var(hip.hipDeviceProp_t, lambda x: check(hip.hipGetDeviceProperties(x, device_id))).gcnArchName.decode().split(":")[0] diff --git a/tinygrad_repo/tinygrad/runtime/ops_metal.py b/tinygrad_repo/tinygrad/runtime/ops_metal.py index e9b1be707..b76d0c684 100644 --- a/tinygrad_repo/tinygrad/runtime/ops_metal.py +++ b/tinygrad_repo/tinygrad/runtime/ops_metal.py @@ -7,7 +7,7 @@ from tinygrad.dtype import dtypes, AddrSpace from tinygrad.renderer.cstyle import MetalRenderer from tinygrad.runtime.autogen import metal from tinygrad.runtime.support.c import DLL -from tinygrad.runtime.support.hcq2 import HWQueue, ccall, patch, layout_args, to_name +from tinygrad.runtime.support.hcq2 import HWQueue, ccall, patch, layout_args, to_name, bufferize_buf from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, uopfunc from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops @@ -141,7 +141,7 @@ class MetalQueue(HWQueue): dims = (*prg.arg.global_size, *prg.arg.local_size) if any(isinstance(d, UOp) for d in dims): self.sizes.append((len(self.cmds), at:=round_up(self.nbytes, 8))) - self.rows += layout_args([d.cast(dtypes.uint64) if isinstance(d, UOp) else UOp.const(d, dtypes.uint64) for d in dims], at) + self.rows += layout_args([UOp.const(d, dtypes.uint64) for d in dims], at) self.nbytes = at + 48 self.cmds.append((obj.lib, obj.name, tuple(1 if isinstance(d, UOp) else int(d) for d in dims), off)) @@ -217,7 +217,8 @@ class MetalDevice(Compiled): (UPat(Ops.ALLOC, tag=self.tag("handles")), lambda d=self: d.handles), (UPat(Ops.ALLOC, tag="slots", name="b"), # with stamps lambda b, d=self: d.new_slots(b.max_numel()) if b.device == d.device and b.max_numel() > 4 else None), - (UPat(Ops.ALLOC, name="b"), lambda b, d=self: d.new_icb(*b.tag[1:]) if isinstance(b.tag, tuple) and b.tag[0] == d.tag("mtl_icb") else None)]) + (UPat(Ops.ALLOC, name="b"), lambda ctx, b, d=self: + d.new_icb(bufferize_buf(ctx, b.rtag(None)).buffer, *b.tag[1:]) if isinstance(b.tag, tuple) and b.tag[0] == d.tag("mtl_icb") else None)]) @functools.cached_property def handles(self) -> Buffer: @@ -247,9 +248,8 @@ class MetalDevice(Compiled): descriptor.setSupportIndirectCommandBuffers(True) return checked(self.sysdevice.newComputePipelineStateWithDescriptor_options_reflection_error, descriptor, metal.MTLPipelineOptionNone, None) - def new_icb(self, cmds:tuple[tuple[bytes, str, tuple[int, ...], int], ...], header:int) -> Buffer: + def new_icb(self, buf:Buffer, cmds:tuple[tuple[bytes, str, tuple[int, ...], int], ...], header:int) -> Buffer: pipes = dedup(c[:2] for c in cmds) - buf = Buffer(self.device, header + 8 * (1 + len(cmds) + len(pipes)), options=BufferSpec(nolru=True), preallocate=True) desc = metal.MTLIndirectCommandBufferDescriptor.new() desc.setCommandTypes(metal.MTLIndirectCommandTypeConcurrentDispatch) desc.setMaxKernelBufferBindCount(1) @@ -259,7 +259,7 @@ class MetalDevice(Compiled): for cmd, (lib, name, dims, off) in zip(commands, cmds): cmd.setComputePipelineState(state:=self.pipeline(lib, name)) if prod(dims[3:]) > (mx:=state.maxTotalThreadsPerThreadgroup()): raise RuntimeError(f"local size {dims[3:]} bigger than {mx}") - cmd.setKernelBuffer_offset_atIndex(buf.get_storage().meta, off, 0) + cmd.setKernelBuffer_offset_atIndex(buf.meta, buf.offset + off, 0) cmd.concurrentDispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*dims[:3]), metal.MTLSize(*dims[3:])) cmd.setBarrier() buf.host.view(fmt='Q')[header // 8:] = array.array('Q', [icb.value, *[c.value for c in commands], *[self.pipeline(*p).value for p in pipes]]) diff --git a/tinygrad_repo/tinygrad/runtime/ops_qcom.py b/tinygrad_repo/tinygrad/runtime/ops_qcom.py index b699dfb77..e778c1f8b 100644 --- a/tinygrad_repo/tinygrad/runtime/ops_qcom.py +++ b/tinygrad_repo/tinygrad/runtime/ops_qcom.py @@ -201,7 +201,7 @@ class QCOMComputeQueue(HWQueue): ib, ib_off = unwrap_view(cmdbuf) fd, ctxid = [UOp.variable(n, 0, 2**31 - 1, dtypes.int32) for n in ("kgsl_fd", "kgsl_ctx")] obj = cstruct(kgsl.struct_kgsl_command_object, gpuaddr=ib.getaddr(self.devs) + ib_off, size=cmdbuf.max_numel(), flags=kgsl.KGSL_CMDLIST_IB) - req = cstruct(kgsl.struct_kgsl_gpu_command, cmdlist=obj.getaddr(HCQ_RUNTIME_DEV.value), cmdsize=ctypes.sizeof(kgsl.struct_kgsl_command_object), + req = cstruct(kgsl.struct_kgsl_gpu_command, cmdlist=obj.getaddr(HCQ_RUNTIME_DEV.device), cmdsize=ctypes.sizeof(kgsl.struct_kgsl_command_object), numcmds=1, context_id=ctxid) return qcom_submit(req.after(cmdbuf), UOp.alloc((1,), dtypes.int32, device=self.devs[0]).rtag("submit_ret"), fd) diff --git a/tinygrad_repo/tinygrad/runtime/support/hcq2.py b/tinygrad_repo/tinygrad/runtime/support/hcq2.py index 322eeabcd..652a19f01 100644 --- a/tinygrad_repo/tinygrad/runtime/support/hcq2.py +++ b/tinygrad_repo/tinygrad/runtime/support/hcq2.py @@ -4,11 +4,11 @@ import functools, itertools, weakref, ctypes, struct, time, hashlib, pathlib from dataclasses import replace, dataclass, field from collections import defaultdict from tinygrad.helpers import dedup, pluralize, unwrap, to_tuple, ContextVar, Context, panic, partition, getenv, to_name -from tinygrad.helpers import DEBUG, VIZ, DEV, ALL2ALL, PROFILE, diskcache_get, diskcache_put +from tinygrad.helpers import DEBUG, VIZ, ALL2ALL, PROFILE, diskcache_get, diskcache_put from tinygrad.device import Device, Buffer, BufferSpec, Compiled, TinyELF, HCQ_RUNTIME_DEV, ProfileProgramEvent from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo, GroupOp, graph_rewrite, rewrite_group, exec_alu, uopfunc, sym_infer from tinygrad.uop.ops import pm_renumber_slots -from tinygrad.dtype import dtypes, DType, DTYPES_DICT, AddrSpace +from tinygrad.dtype import dtypes, DType, AddrSpace, truncate, to_dtype from tinygrad.renderer import Estimates from tinygrad.engine.realize import get_call_arg_uops, get_call_name, get_call_outs_ins from tinygrad.engine.realize import estimate_uop, pm_flatten_linear, lower_and_compile @@ -17,8 +17,9 @@ from tinygrad.engine.realize import estimate_uop, pm_flatten_linear, lower_and_c # 0. helpers HCQ_CACHE_THRESH = ContextVar("HCQ_CACHE_THRESH", 64) -HCQ_DEVS = frozenset(("AMD", "NV", "QCOM", "CUDA", "NULL", "METAL")) +HCQ_DEVS = frozenset(("AMD", "NV", "QCOM", "CUDA", "HIP", "NULL", "METAL")) HOST_DEVS = frozenset(("CPU", "PYTHON")) +UVM_DEVS = frozenset(("QCOM", "METAL")) # the host maps their memory, so copies are host memcpys @dataclass(frozen=True) class HCQInfo: @@ -33,14 +34,12 @@ class HCQInfo: def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c -def get_enqueue_devs(call:UOp) -> Any|None: - if call.op is not Ops.CALL: return None # entries can be AFTER-wrapped calls - if call.body.op not in (Ops.PROGRAM, Ops.STORE) and not (call.body.op is Ops.CUSTOM_FUNCTION and call.body.arg.name == "encdec"): return None - if not (bufs:=get_call_arg_uops(call)): return None - if call.body.op is Ops.STORE: bufs = bufs[::-1] # copies push from the src device: p2p writes are faster than reads - devs = min(bufs, key=lambda b: not all_devices_in(b.device, HCQ_DEVS)).device - if not all_devices_in(devs, HCQ_DEVS): return None - if call.body.op is Ops.STORE and to_tuple(devs)[0].startswith(("QCOM", "METAL")): return None # unified memory uses host copies +def get_enqueue_devs(call:UOp) -> tuple[str, ...]: + if call.op is not Ops.CALL: return () # entries can be AFTER-wrapped calls + if call.body.op not in (Ops.PROGRAM, Ops.STORE) and not (call.body.op is Ops.CUSTOM_FUNCTION and call.body.arg.name == "encdec"): return () + runs_on = call.device if call.body.op is Ops.STORE and all_devices_in(call.device, HCQ_DEVS) else call.src[1].device + devs = tuple(Device.canonicalize(d) for d in to_tuple(runs_on)) + if not all_devices_in(devs, HCQ_DEVS) or (call.body.op is Ops.STORE and all_devices_in(devs, UVM_DEVS)): return () return devs def unwrap_view(v:UOp) -> tuple[UOp, int]: # look through views to (base, byte offset) @@ -57,13 +56,13 @@ def unwrap_lane(v:UOp) -> tuple[UOp, int|None, int]: # look through views and a def select_lane(u:UOp, lane:int) -> UOp: return u.src[lane] if u.op is Ops.MSTACK else u.mselect(lane) if len(to_tuple(u.device)) > 1 else u -def timeline(devs:tuple[str, ...]) -> UOp: return UOp.alloc((2,), dtypes.uint64, 0, device=devs[0]).rtag("timeline") +def timeline(devs:tuple[str, ...]) -> UOp: return UOp.alloc((2,), dtypes.u64, 0, device=devs[0]).rtag("timeline") def timeline_value(devs:tuple[str, ...]) -> UOp: return timeline(devs).index(1).load() def make_program(prg:UOp, size:int, device:str) -> UOp: slot = int.from_bytes(prg.key[:8], "little") if PROFILE: Compiled.profile_events.append(ProfileProgramEvent(device, prg.src[0].arg.function_name, prg.src[3].arg, None, slot, prg.key)) - return UOp.alloc((size,), dtypes.uint8, slot, device=device).rtag("program") + return UOp.alloc((size,), dtypes.u8, slot, device=device).rtag("program") def make_submit(*cmds, devs:str|tuple[str, ...], queue:str, fn:str|None=None, deps:tuple[UOp, ...]=()) -> UOp: # the order is on the arg lin = UOp(Ops.LINEAR, src=tuple(cmds), arg=(to_tuple(devs), queue)).after(*deps) @@ -71,14 +70,14 @@ def make_submit(*cmds, devs:str|tuple[str, ...], queue:str, fn:str|None=None, de def ins(name, *src) -> UOp: return UOp(Ops.INS, tuple(UOp.const(s, dtypes.u32) if isinstance(s, int) else s for s in src), (name, dtypes.void)) -def chunks(nbytes:int, chunk_sz:int, dtype:DType=dtypes.int) -> list[tuple[UOp|int, int]]: # (index, bytes): full chunks as one range, then the tail +def chunks(nbytes:int, chunk_sz:int, dtype:DType=dtypes.i32) -> list[tuple[UOp|int, int]]: # (index, bytes): full chunks as one range, then the tail full, tail = divmod(nbytes, chunk_sz) # no one-trip loops return ([(UOp.range(full, next(UOp.unique_num), dtype=dtype) if full > 1 else 0, chunk_sz)] if full else []) + ([(full, tail)] if tail else []) # C FFI def layout_args(args:Sequence[UOp|int], offset:int=0) -> list[tuple[int, UOp]]: - words = [a if isinstance(a, UOp) else UOp.const(a, dtypes.uint32) for a in args] + words = [a if isinstance(a, UOp) else UOp.const(a, dtypes.u32) for a in args] return [(offset + o, w) for (o, _), w in zip(TinyELF.iter_sig(tuple((None, i, w.dtype, ()) for i, w in enumerate(words))), words)] def pack_args(args:list[tuple[int, UOp]], size:int) -> list[UOp]: @@ -89,27 +88,22 @@ def pack_args(args:list[tuple[int, UOp]], size:int) -> list[UOp]: return words + [UOp(Ops.BINARY, arg=bytes(size - end))] def ccall(fn:Any, *args:UOp|int) -> UOp: - ret = dtypes.void if fn.restype is None else dtypes.uint64 if fn.restype is ctypes.c_void_p else \ - next(d for d in DTYPES_DICT.values() if d.fmt == fn.restype._type_) - return UOp.custom_function(fn.__name__, dtype=ret).call(*[UOp.const(a, dtypes.int) if isinstance(a, int) else a for a in args]) + ret = {None: dtypes.void, ctypes.c_void_p: dtypes.u64}.get(fn.restype) or to_dtype(fn.restype.__name__[2:]) # void, void*, c_int -> int + return UOp.custom_function(fn.__name__, dtype=ret).call(*[UOp.const(a, dtypes.i32) if isinstance(a, int) else a for a in args]) @uopfunc def do_get_time_ms(ms:UOp) -> UOp: - ts = UOp.placeholder((2,), dtypes.uint64, addrspace=AddrSpace.REG) - ts = ts.after(UOp.custom_function("clock_gettime", dtype=dtypes.int).call(UOp.const(time.CLOCK_MONOTONIC, dtypes.int), ts.index(0))) + ts = UOp.placeholder((2,), dtypes.u64, addrspace=AddrSpace.REG) + ts = ts.after(UOp.custom_function("clock_gettime", dtype=dtypes.i32).call(UOp.const(time.CLOCK_MONOTONIC, dtypes.i32), ts.index(0))) return ms.index(0).store(ts[0] * 1000 + ts[1] // 1000000).sink() -def get_time_ms(dep:UOp) -> UOp: return (ms:=UOp.placeholder((1,), dtypes.uint64, addrspace=AddrSpace.REG)).after(do_get_time_ms(ms.after(dep)))[0] - -CDTYPE = {1: dtypes.uchar, 2: dtypes.ushort, 4: dtypes.uint, 8: dtypes.ulong} # a C field as the unsigned int of its size +def get_time_ms(dep:UOp) -> UOp: return (ms:=UOp.placeholder((1,), dtypes.u64, addrspace=AddrSpace.REG)).after(do_get_time_ms(ms.after(dep)))[0] def cstruct(struct_t, **fields:UOp|int) -> UOp: - flds = {n: (o, CDTYPE[ctypes.sizeof(t)]) for n, t, o, *_ in struct_t._real_fields_ if ctypes.sizeof(t)} # skips zero length arrays - rows = [(flds[n][0], v.cast(flds[n][1]) if isinstance(v, UOp) else UOp.const(v, flds[n][1])) for n, v in fields.items()] - buf = UOp.alloc((ctypes.sizeof(struct_t),), dtypes.uint8, device=HCQ_RUNTIME_DEV.value).rtag(struct_t.__name__) + flds = {n: (o, {d.itemsize: d for d in dtypes.uints}[ctypes.sizeof(t)]) for n, t, o, *_ in struct_t._real_fields_ if ctypes.sizeof(t)} + rows = [(flds[n][0], UOp.const(v, flds[n][1])) for n, v in fields.items()] + buf = UOp.alloc((ctypes.sizeof(struct_t),), dtypes.u8, device=HCQ_RUNTIME_DEV.device).rtag(struct_t.__name__) return patch(buf, rows, bytes(ctypes.sizeof(struct_t))) -def cfield(buf:UOp, struct_t, name:str) -> UOp: return buf[(f:=getattr(struct_t, name)).offset:f.offset + f.size].bitcast(CDTYPE[f.size]).index(0) - # ***************** # 0.1. prep: eager buffers become tagged params @@ -124,8 +118,8 @@ pm_replace_buffers = PatternMatcher([(UPat(Ops.BUFFER, name="b"), lambda ctx, b: # 1.1. prep: unwrap multi def unwrap_call(call:UOp) -> UOp|None: - if get_enqueue_devs(call) is None or (n:=max(len(to_tuple(a.device)) for a in get_call_arg_uops(call))) == 1: return None - dnum = UOp.variable("_device_num", 0, n - 1, dtypes.int) + if not get_enqueue_devs(call) or (n:=max(len(to_tuple(a.device)) for a in get_call_arg_uops(call))) == 1: return None + dnum = UOp.variable("_device_num", 0, n - 1, dtypes.i32) return UOp(Ops.LINEAR, src=tuple(call.replace(src=(call.body, *[a if a.is_bound_var else select_lane(a, i) for a in call.src[1:]], dnum.bind(i))) for i in range(n))) pm_unwrap_multi = PatternMatcher([(UPat(Ops.CALL, name="call"), unwrap_call)]) @@ -133,11 +127,6 @@ pm_unwrap_multi = PatternMatcher([(UPat(Ops.CALL, name="call"), unwrap_call)]) # ***************** # 1.2. prep: staging copies -STAGING_SIZE, STAGING_SLOTS = (4 if DEV.interface.startswith("MOCK") else 128) << 20, 2 - -@functools.cache -def _staging(device:str) -> Buffer: return Buffer(device, STAGING_SIZE, preallocate=True) - def split_rdma(call:UOp, dst:UOp, src:UOp) -> UOp|None: devs = [to_tuple(b.device)[0] for b in (dst, src)] if not all(hasattr(Device[d], "iface") for d in devs) or Device[devs[0]].peer_group == Device[devs[1]].peer_group: return None # not 2 nodes @@ -153,21 +142,18 @@ def split_rdma(call:UOp, dst:UOp, src:UOp) -> UOp|None: def stage_copy(call:UOp, dst:UOp, src:UOp) -> UOp|None: devs = [to_tuple(b.device)[0] for b in (dst, src)] if any(d.startswith("RDMA") for d in devs): return None # over the nic + if not (enqueued:=get_enqueue_devs(call)): return None + device = enqueued[0] - if (device:=get_enqueue_devs(call)) is None: return None + from tinygrad.runtime.ops_disk import is_disk_read dev, host, usb_memcpys = Device[device], Device[device].host, getattr(Device[device], "is_usb", False) - mappable = {"CPU", "PYTHON"} | ({"NPY", "DISK"} if usb_memcpys else set()) + mappable = {"CPU", "PYTHON"} | ({"NPY", "DISK"} if usb_memcpys else {"DISK"} if is_disk_read(call) else set()) # ops_disk stages it in the batch if device != host and not all(Device[d].peer_group == dev.peer_group or (d.split(":")[0] in mappable and Device[d].host == host) for d in devs): - (staging:=_staging(host)).get_buf(device) - base, it, copies = UOp.from_buffer(staging, dtypes.uint8), src.dtype.itemsize, [] - chunk = (STAGING_SIZE // STAGING_SLOTS) // it - for i, off in enumerate(range(0, src.max_numel(), chunk)): - stage, part = base[(so:=(i % STAGING_SLOTS) * chunk * it):so + (n:=min(chunk, src.max_numel() - off)) * it], src[off:off+n] - copies += [stage.store_call(part), dst[off:off+n].store_call(stage)] - return UOp(Ops.LINEAR, src=tuple(copies)) + stage = UOp.alloc(src.max_shape, src.dtype, int.from_bytes(call.key[:8], "little"), device=host) + return UOp(Ops.LINEAR, src=(stage.store_call(src), dst.store_call(stage))) if Device[device].has_copy_queue: return None - out, inp = (UOp.param(i, dtypes.uint8, b.nbytes(), device=device) for i, b in enumerate((dst, src))) + out, inp = (UOp.param(i, dtypes.u8, b.nbytes(), device=device) for i, b in enumerate((dst, src))) ast = out.index(r:=UOp.range(src.nbytes(), 0)).store(inp.index(r).load()).end(r).sink(arg=KernelInfo()) return lower_and_compile(call.replace(src=(ast, *call.src[1:]))) @@ -181,36 +167,23 @@ pm_insert_copy_staging = PatternMatcher([ class DepsTracker: def __init__(self): - # tracks (offset, end, dep) ranges per base buffer/lane to handle suballocated buffers correctly. - self.w_dependency_map: dict[Any, list[tuple[int, int, Any]]] = defaultdict(list) - self.r_dependency_map: dict[Any, list[tuple[int, int, Any]]] = defaultdict(list) + self.writes:dict[Any, list[tuple[int, int, Any]]] = defaultdict(list) + self.reads:dict[Any, list[tuple[int, int, Any]]] = defaultdict(list) - def access_resources(self, bufs:Sequence[UOp|Buffer], write:list[int], new_dependency:Any): - ranges:list[tuple[Any, int, int]] = [] - for buf in bufs: - if isinstance(buf, Buffer): ranges.append((id(buf.base), buf.offset, buf.offset + buf.nbytes)) - else: - base, lane, off = unwrap_lane(buf) - ranges.append(((base, lane), off, off + buf.max_numel() * buf.dtype.itemsize)) - wait_nodes = [] - for i, (key, s, e) in enumerate(ranges): - wait_nodes += [dep for st,en,dep in self.w_dependency_map[key] if st < e and s < en] - if i in write: wait_nodes += [dep for st,en,dep in self.r_dependency_map[key] if st < e and s < en] - for i, (key, s, e) in enumerate(ranges): - if i in write: - for dmap in [self.w_dependency_map, self.r_dependency_map]: - kept = [] - for entry in dmap[key]: - st, en, dep = entry - if st == en: continue - if en <= s or e <= st: kept.append(entry) - else: - if st < s: kept.append((st, s, dep)) - if e < en: kept.append((e, en, dep)) - dmap[key] = kept - self.w_dependency_map[key].append((s, e, new_dependency)) - else: self.r_dependency_map[key].append((s, e, new_dependency)) - return list({id(x):x for x in wait_nodes}.values()) + def access_resources(self, bufs:Sequence[UOp], write:list[int], new_dependency:Any) -> list: + accesses = [(i in write, (base, lane), off, off + b.max_numel() * b.dtype.itemsize) + for i, b in enumerate(bufs) for base, lane, off in [unwrap_lane(b)]] + + waits = [dep for is_write, key, start, end in accesses for lo, hi, dep in (self.writes[key] + self.reads[key] if is_write else self.writes[key]) + if lo < end and start < hi] + + for is_write, key, start, end in accesses: + if is_write: + for ranges in (self.writes, self.reads): + before_and_after = [((lo, min(hi, start), dep), (max(lo, end), hi, dep)) for lo, hi, dep in ranges[key]] + ranges[key] = [(lo, hi, dep) for parts in before_and_after for lo, hi, dep in parts if lo < hi] + (self.writes if is_write else self.reads)[key].append((start, end, new_dependency)) + return dedup(waits) @dataclass class BatchCtx: @@ -236,7 +209,7 @@ class BatchCtx: self.queues.setdefault(d, []) self.signal_tags = {tag for (dev, q), tag in self.last.items() if q != self.epilogue_queue(dev) or (dev, q) in self.peers} # a slot is [signal][timestamp], 16 bytes: the queue signals, the timeline, then two per call if profiling - self.slots = {dev: UOp.alloc((2 * (len(qs) + 1 + (2 * len(self.batch) if self.profile else 0)),), dtypes.uint64, device=dev, + self.slots = {dev: UOp.alloc((2 * (len(qs) + 1 + (2 * len(self.batch) if self.profile else 0)),), dtypes.u64, device=dev, spec=BufferSpec(host=True, uncached=True, cpu_access=True)).rtag("slots") for dev, qs in self.queues.items()} def epilogue_queue(self, dev:str) -> str: return "COMPUTE:0" if len(self.queues[dev]) != 1 else self.queues[dev][0] # closes the device @@ -256,11 +229,10 @@ def _wait_ins(ctx:BatchCtx, call:UOp, device:str, queue:str, tag:int) -> list[UO if latest and device.split(":")[0] == "NV" and queue.startswith("COMPUTE") and (p:=ctx.prev[tag]) is not None: latest[(device, queue)] = p ctx.signal_tags |= set(latest.values()) - return [UOp(Ops.INS, arg=("wait", dtypes.void), src=(ctx.queue_signal((d,), q), UOp.const(t + 1, dtypes.uint64))) for (d, q), t in latest.items()] + return [ins("wait", ctx.queue_signal((d,), q), UOp.const(t + 1, dtypes.u64)) for (d, q), t in latest.items()] def _start_ins(ctx:BatchCtx, dev:str, queue:str) -> list[UOp]: # a queue first waits for prior work of its device and of the peers it touches - return [UOp(Ops.INS, arg=("barrier", dtypes.void), src=())] + \ - [UOp(Ops.INS, arg=("wait", dtypes.void), src=(timeline((d,)), timeline_value((d,)))) for d in [dev, *sorted(ctx.peers.get((dev, queue), ()))]] + return [ins("barrier")] + [ins("wait", timeline((d,)), timeline_value((d,))) for d in [dev, *sorted(ctx.peers.get((dev, queue), ()))]] def _build_queues(ctx:BatchCtx) -> dict[tuple[tuple[str, ...], str], list[UOp]]: # find all waits first to mark calls that must signal @@ -270,20 +242,19 @@ def _build_queues(ctx:BatchCtx) -> dict[tuple[tuple[str, ...], str], list[UOp]]: if not (q:=queues.setdefault((devices, queue), [])): q += _start_ins(ctx, devices[0], queue) # first use of a queue # dependency waits, then the call between its timestamps - ts_ins = [UOp(Ops.INS, arg=("timestamp", dtypes.void), src=(ctx.slot(devices, i),)) for i in ctx.stamps(devices, tag)] + ts_ins = [ins("timestamp", ctx.slot(devices, i)) for i in ctx.stamps(devices, tag)] q += waits + ts_ins[:1] + [call] + ts_ins[1:] # signal the queue if someone waits for us - if tag in ctx.signal_tags: - q += [UOp(Ops.INS, arg=("store", dtypes.void), src=(ctx.queue_signal(devices, queue), UOp.const(tag + 1, dtypes.uint64)))] + if tag in ctx.signal_tags: q += [ins("store", ctx.queue_signal(devices, queue), UOp.const(tag + 1, dtypes.u64))] # one queue advances the device timeline after all other queues finish, and after the queues of the peers that touched the device for dev in ctx.queues: touched = sorted(k for k, ds in ctx.peers.items() if dev in ds) qdev, queue = touched[0] if dev.split(":")[0] in HOST_DEVS else (dev, ctx.epilogue_queue(dev)) # the host has no queues, a toucher closes it - waits = [UOp(Ops.INS, arg=("wait", dtypes.void), src=(ctx.queue_signal((d,), q), UOp.const(ctx.last[(d, q)] + 1, dtypes.uint64))) + waits = [ins("wait", ctx.queue_signal((d,), q), UOp.const(ctx.last[(d, q)] + 1, dtypes.u64)) for d, q in [(dev, q) for q in ctx.queues[dev]] + touched if (d, q) != (qdev, queue)] - bump = UOp(Ops.INS, arg=("store", dtypes.void), src=(timeline((dev,)), timeline_value((dev,)) + UOp.const(1, dtypes.uint64))) + bump = ins("store", timeline((dev,)), timeline_value((dev,)) + 1) # multiple copy queues may need a new compute stream. a peer without calls of its own starts like any queue if not (q:=queues.setdefault(((qdev,), queue), [])) and qdev not in {d for d, _ in ctx.last}: q += _start_ins(ctx, qdev, queue) @@ -318,7 +289,7 @@ def _finalize_batch(ctx:BatchCtx, skip_wait:bool=False) -> UOp: @rewrite_group(new_ctx=False) def sched_batches(l:UOp, profile:bool) -> UOp: - devs = [() if (d:=get_enqueue_devs(c)) is None else tuple(Device.canonicalize(x) for x in to_tuple(d)) for c in l.src] + devs = [get_enqueue_devs(c) for c in l.src] # assign to queues peers = sorted({Device.canonicalize(d) for c in l.src if c.op is Ops.CALL and c.body.op is Ops.STORE @@ -326,7 +297,7 @@ def sched_batches(l:UOp, profile:bool) -> UOp: num_queues = max(1, getenv("HCQ_NUM_SDMA", min(len(peers), 8) if ALL2ALL >= 1 else 1)) queues = ["COMPUTE:0" if c.op is Ops.CALL and c.body.op is Ops.PROGRAM else "COPY:0" for c in l.src] for i, c in enumerate(l.src): - if c.op is Ops.CALL and c.body.op is Ops.CUSTOM_FUNCTION and c.body.arg.name == "encdec": queues[i] = "ENCDEC:0" + if c.op is Ops.CALL and c.body.op is Ops.CUSTOM_FUNCTION: queues[i] = "ENCDEC:0" if c.op is Ops.CALL and c.body.op is Ops.STORE and all(b.device in peers for b in get_call_arg_uops(c)): queues[i] = f"COPY:{(peers.index(c.src[1].device) - peers.index(c.src[2].device) - 1) % len(peers) % num_queues}" @@ -378,8 +349,8 @@ class HWQueue: self.patches.append((len(self.blob), w)) self.blob += bytes(w.dtype.itemsize) else: - v, n = (c.val, w.dtype.itemsize) if isinstance(w, UOp) else (c, 4) - self.blob += (v & (1 << 8 * n) - 1).to_bytes(n, 'little') + v, dt = (c.val, w.dtype) if isinstance(w, UOp) else (c, dtypes.u32) + self.blob += struct.pack(unwrap(dt.fmt), truncate[dt](v)) return len(self.blob) def loop(self, body:UOp, r:UOp): @@ -387,7 +358,7 @@ class HWQueue: for u in body.src: self.q_rewrite.rewrite(u, ctx=self) trip, words = len(self.blob) - start, self.patches[first:] # each trip gets its words as dwords: a trip need not be 8 bytes aligned - self.patches[first:] = [(o + 4 * k + r * trip, (w >> (32 * k)).cast(dtypes.uint32)) for o, w in words for k in range(w.dtype.itemsize // 4)] + self.patches[first:] = [(o + 4 * k + r * trip, (w >> (32 * k)).cast(dtypes.u32)) for o, w in words for k in range(w.dtype.itemsize // 4)] self.blob += self.blob[start:] * int(r.vmax) def memory_barrier(self): pass # a copy queue has nothing to flush @@ -424,7 +395,7 @@ def hcq_fence(slots:UOp, tl:UOp, tv:UOp, last:int) -> UOp: # wait for the previo loop = UOp.range(UOp(Ops.NOOP).after(start:=get_time_ms(target:=slots.index(last).load())), next(UOp.unique_num), dtype=dtypes.void) done = tl.after(target, loop).index(0).load() - bumped = tl.after(done.backedge(loop, (done < target) & (get_time_ms(done) - start < 30000))).index(1).store(nxt:=tv + UOp.const(1, dtypes.uint64)) + bumped = tl.after(done.backedge(loop, (done < target) & (get_time_ms(done) - start < 30000))).index(1).store(nxt:=tv + 1) return slots.after(bumped).index(last).store(nxt).sink() def encode_fence(f:UOp) -> UOp: @@ -453,7 +424,7 @@ def encode_cmdbuf(hq:HWQueue, lin:UOp|None=None, name:str="cmdbuf", device:str|t rt = [(o, w) for o, w in patches if isinstance(o, int) and o % 4 == 0 and w.dtype.itemsize in (4, 8) and not _is_link_patch(w)] uses = {w: [o for o, _ in grp] for w, grp in itertools.groupby(sorted(rt, key=lambda p: p[1].key), key=lambda p: p[1])} looped = {w: (at, [*w.ranges][0] if w.ranges else UOp.range(len(at), next(UOp.unique_num))) for w, at in uses.items() if len(at) > 1 or w.ranges} - dwords = [(UOp(Ops.BINARY, arg=struct.pack(f"<{len(at)}I", *at)).bitcast(dtypes.uint32).index(r).load() + 4 * k, (w >> 32 * k).cast(dtypes.uint32)) + dwords = [(UOp(Ops.BINARY, arg=struct.pack(f"<{len(at)}I", *at)).bitcast(dtypes.u32).index(r).load() + 4 * k, (w >> 32 * k).cast(dtypes.u32)) for w, (at, r) in looped.items() for k in range(w.dtype.itemsize // 4)] patches = [(o, w) for o, w in patches if w not in looped] + dwords nested = dedup([g.src[0] for _, w in patches for g in w.toposort() if g.op is Ops.GETADDR and g.src[0].op is Ops.LINEAR]) @@ -466,7 +437,7 @@ def encode_cmdbuf(hq:HWQueue, lin:UOp|None=None, name:str="cmdbuf", device:str|t bufs.append((offs, encode_cmdbuf(hq, name=lname))) views = {l: buf.without_after[o:e] for offs, buf in bufs for l, (o, e) in offs.items()} - buf = UOp.alloc((len(stream),), dtypes.uint8, device=device or hq.devs[0]).rtag(to_name(name, hq.queue)).after(*hq.deps) + buf = UOp.alloc((len(stream),), dtypes.u8, device=device or hq.devs[0]).rtag(to_name(name, hq.queue)).after(*hq.deps) words = UOp.sink(*[w for _, w in patches]).substitute(views).src return patch(buf, list(zip([o for o, _ in patches], words)), stream).after(*[b for _, b in bufs]) @@ -494,12 +465,12 @@ def _needs_arg(u:UOp, root:bool) -> bool: def _param_for(u:UOp, slot:int) -> UOp: if u.op is Ops.GETADDR or u.is_variable: - return UOp.param(slot, u.commit_dtype(dtypes.int), name=u.arg.name if u.is_variable else None, addrspace=AddrSpace.ALU).cast(u.dtype) - return UOp.param(slot, u.dtype, u.max_numel(), HCQ_RUNTIME_DEV.value, name=f"{u.tag}_{slot}" if isinstance(u.tag, str) else None) + return UOp.param(slot, u.commit_dtype(dtypes.i32), name=u.arg.name if u.is_variable else None, addrspace=AddrSpace.ALU).cast(u.dtype) + return UOp.param(slot, u.dtype, u.max_numel(), HCQ_RUNTIME_DEV.device, name=f"{u.tag}_{slot}" if isinstance(u.tag, str) else None) def lift(call:UOp, root:bool=False) -> UOp: # callees are lifted already body, args = graph_rewrite(call.body, pm_lift_deps + pm_renumber_slots, ctx=itertools.count(), walk=True, name="lift deps"), list(call.src[1:]) - nodes = body.toposort(gate=lambda u: u.op is not Ops.GETADDR, enter_calls=False) + nodes = body.toposort(gate=lambda u: u.op is not Ops.GETADDR) leaves = dedup([u for u in nodes if _needs_arg(u, root)] + [g for u in nodes for g in u.src if g.op is Ops.GETADDR]) slots = args + (new:=[u for u in leaves if u not in args]) body = body.substitute({u: _param_for(u, slots.index(u)) for u in leaves}, walk=True) @@ -543,7 +514,7 @@ def _source_stamp() -> str: return hashlib.sha256(str(files).encode()).hexdigest() def _disk_key(linear:UOp, profile:bool) -> str: - devs = sorted({d for c in linear.src if (ds:=get_enqueue_devs(c)) is not None for d in to_tuple(ds)} | {HCQ_RUNTIME_DEV.value}) + devs = sorted({d for c in linear.src if (ds:=get_enqueue_devs(c)) is not None for d in to_tuple(ds)} | {HCQ_RUNTIME_DEV.device}) return "|".join([_source_stamp(), linear.key.hex(), str(profile), str(ALL2ALL >= 1), *(f"{d}={Device[d].renderer.target!r}" for d in devs)]) @rewrite_group(lambda linear,input_uops,profile,cache=False,ret=None: f"HCQ Compile {pluralize('Kernel', len(ret.src))}") @@ -577,37 +548,32 @@ def bufferize_buf(ctx:LinkCtx, b:UOp) -> UOp: # ctx: a kept link (the jit's) own dev, spec = Device[b.device], b.arg.spec or BufferSpec(cpu_access=True) # data the device reads, unless the alloc says otherwise # a device owns the placeholders it names, the rest are allocated where they live - if (r:=cast(Buffer|None, Compiled.pm_bufferize.rewrite(b))) is not None: pass - elif not ctx.use_rt: r = Buffer(dev.device, max(b.max_numel(), 1) * b.dtype.itemsize, options=spec, preallocate=True) + if (r:=cast(Buffer|None, Compiled.pm_bufferize.rewrite(b, ctx))) is not None: pass + elif not ctx.use_rt or b.tag is None: r = Buffer(dev.device, max(b.max_numel(), 1) * b.dtype.itemsize, options=spec, preallocate=True) else: r = dev.rt_buffer(spec).view(b.nbytes(), dev.rt_allocator(spec).alloc(max(b.nbytes(), 1), alignment=256)).ensure_allocated() - return UOp.from_buffer(r, b.dtype, HCQ_RUNTIME_DEV.value) + return UOp.from_buffer(r, b.dtype, HCQ_RUNTIME_DEV.device) def resolve_getaddr(ctx:LinkCtx, g:UOp) -> UOp|None: if unwrap_lane(buf:=unwrap_view(g.src[0])[0])[0].op is not Ops.BUFFER: return None # input address, resolved per run ctx.refs.append(buf) # add to refs - return UOp.const(g.val, dtypes.uint64) + return UOp.const(g.val, dtypes.u64) -def fold_binary(buf:UOp, blob:UOp) -> UOp: - base, off = unwrap_view(buf) - cast(Buffer, base.buffer).ensure_allocated().host.view(fmt='B')[off:off + len(blob.arg)] = blob.arg - return UOp(Ops.NOOP) - -def write_words(buf:UOp, writes:list[tuple[int, int, int]]) -> UOp: # (word index, size, value) +def _write_host(buf:UOp, writes:list[tuple[int, bytes]]) -> UOp: base, off = unwrap_view(buf) mv = cast(Buffer, base.buffer).ensure_allocated().host.view(fmt='B') - for o, n, v in writes: mv[off + o * n:off + (o + 1) * n] = (v & (1 << 8 * n) - 1).to_bytes(n, 'little') + for o, data in writes: mv[off + o:off + o + len(data)] = data return UOp(Ops.NOOP) -def words(offs:UOp, ws:UOp): return zip(*[s.src if s.op is Ops.STACK else (s,) for s in (offs, ws)]) # a stack or one word - -def fold_words(buf:UOp, offs:UOp, ws:UOp) -> UOp: return write_words(buf, [(o.val, w.dtype.itemsize, w.val) for o, w in words(offs, ws)]) +def fold_words(buf:UOp, offs:UOp, ws:UOp) -> UOp: + return _write_host(buf, [(o.val * w.dtype.itemsize, struct.pack(w.dtype.fmt, truncate[w.dtype](w.val))) for o, w in zip(offs.src, ws.src)]) def fold_ranged(buf:UOp, offs:UOp, ws:UOp, e:UOp) -> UOp: # the words of every trip vs = {r: UOp.variable(f"r{i}", 0, r.vmax) for i, r in enumerate(e.src[1:])} - ows = list(words(offs.substitute(vs), ws.substitute(vs))) + ows = list(zip(offs.substitute(vs).src, ws.substitute(vs).src)) trips = [dict(zip([v.expr for v in vs.values()], t)) for t in itertools.product(*[range(int(r.vmax) + 1) for r in vs])] - return write_words(buf, [(sym_infer(o, t), w.dtype.itemsize, sym_infer(w, t)) for t in trips for o, w in ows]) + data = [(sym_infer(o, t) * w.dtype.itemsize, struct.pack(w.dtype.fmt, truncate[w.dtype](sym_infer(w, t)))) for t in trips for o, w in ows] + return _write_host(buf, data) pm_link = PatternMatcher([ # collapse committed const conversions @@ -620,10 +586,10 @@ pm_link = PatternMatcher([ # math on consts is a const (UPat(GroupOp.ALU, src=UPat.cvar().or_casted(), name="a"), lambda a: UOp.const(exec_alu(a.op, a.dtype, [s.val for s in a.src], False), a.dtype)), # fold rules - (UPat(name="buf").store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast())), fold_binary), + (UPat(name="buf").store(UPat(Ops.BINARY, name="blob").or_bitcasted()), lambda buf, blob: _write_host(buf, [(0, blob.arg)])), (UPat(name="buf").index(UPat(Ops.STACK, src=UPat.cvar(), name="offs")).store(UPat(Ops.STACK, src=UPat.cvar().or_casted(), name="ws")), fold_words), - (UPat(name="buf").index(UPat.cvar("offs")).store(UPat.cvar().or_casted("ws")), fold_words), - (UPat(name="buf").index(UPat(name="offs")).store(UPat(name="ws")).end(allow_any_len=True, name="e"), fold_ranged), + (UPat(name="buf").index(UPat.cvar("off")).store(UPat.cvar().or_casted("w")), lambda buf, off, w: fold_words(buf, UOp.stack(off), UOp.stack(w))), + (UPat(name="buf").index(UPat(Ops.STACK, name="offs")).store(UPat(Ops.STACK, name="ws")).end(allow_any_len=True, name="e"), fold_ranged), # a call keeps the deps that are not written yet (UPat(Ops.AFTER, src=(UPat(Ops.CALL),), allow_any_len=True, name="a"), lambda a: a.src[0].after(*(s for s in a.src[1:] if s.op is not Ops.NOOP))), (UPat(Ops.AFTER, name="a"), lambda a: None if a.without_after.op is Ops.CALL else @@ -637,7 +603,7 @@ def hcq_link(linear:UOp, input_uops:list[UOp]|None=None, allow_cache=True) -> UO if allow_cache and (linked:=link_linear_cache.get(linear)) is not None: return linked # if we have any link time buffers, do not cache this linear - cache = allow_cache and not any(u.tag == "lt_input" for u in linear.toposort(enter_calls=False) if u.op is Ops.PARAM) + cache = allow_cache and not any(u.tag == "lt_input" for u in linear.toposort() if u.op is Ops.PARAM) inputs = {UOp.param(i, b.dtype, b.max_numel(), b.device).replace(tag="lt_input"): b for i, b in enumerate(input_uops or ())} linked = graph_rewrite(linear, pm_link, ctx=(ctx:=LinkCtx(inputs, use_rt=allow_cache and not cache)), walk=True, name="link") diff --git a/tinygrad_repo/tinygrad/runtime/support/nv/nvdev.py b/tinygrad_repo/tinygrad/runtime/support/nv/nvdev.py index eafb18f76..e27af2fa1 100644 --- a/tinygrad_repo/tinygrad/runtime/support/nv/nvdev.py +++ b/tinygrad_repo/tinygrad/runtime/support/nv/nvdev.py @@ -95,9 +95,6 @@ class NVDev: def rreg(self, addr:int) -> int: return self.mmio[addr // 4] def _early_ip_init(self): - self.reg_names:set[str] = set() - self.reg_offsets:dict[str, tuple[int, int]] = {} - self.include("nv_ref", "") self.include("dev_fb", "tu102") self.include("dev_gc6_island", "ga102") diff --git a/tinygrad_repo/tinygrad/runtime/support/usb.py b/tinygrad_repo/tinygrad/runtime/support/usb.py index fe864c240..c85d93e8e 100644 --- a/tinygrad_repo/tinygrad/runtime/support/usb.py +++ b/tinygrad_repo/tinygrad/runtime/support/usb.py @@ -4,7 +4,7 @@ from tinygrad.helpers import DEBUG, DEV, HCQ_RUNTIME_DEV, to_mv, round_up, ceild from tinygrad.dtype import dtypes, DType, AddrSpace from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, uopfunc from tinygrad.device import Buffer, BufferSpec, Compiled -from tinygrad.runtime.support.hcq2 import HCQ_DEVS, CDTYPE, ccall, patch, unwrap_view, all_devices_in, to_name, get_time_ms, ins, chunks +from tinygrad.runtime.support.hcq2 import HCQ_DEVS, ccall, patch, unwrap_view, all_devices_in, to_name, get_time_ms, ins, chunks from tinygrad.runtime.support.memory import MMIOInterface from tinygrad.runtime.support import c @@ -207,7 +207,7 @@ class USBMMIOInterface(MMIOInterface): HALF, CHUNK, SLOT, STREAM = 0x40000, 0x40000 - 512, 0x4000, 1 << 20 # sram half, payload, slot, stream HOST_SIZE = 64 + 2 * HALF + STREAM # link, staging, zeros -def usb_host(dev:str) -> UOp: return UOp.alloc((HOST_SIZE,), dtypes.uint8, 0, device=HCQ_RUNTIME_DEV.value).rtag(to_name(dev, "usb_host")) +def usb_host(dev:str) -> UOp: return UOp.alloc((HOST_SIZE,), dtypes.uint8, 0, device=HCQ_RUNTIME_DEV.device).rtag(to_name(dev, "usb_host")) def usb_link(dev:str) -> UOp: return usb_host(dev)[:48].bitcast(dtypes.uint64) # handle, context, prev chunks, two transfers, error state def usb_stage(dev:str) -> UOp: return usb_host(dev)[64:64 + 2 * HALF] def usb_zeros(dev:str) -> UOp: return usb_host(dev)[64 + 2 * HALF:] @@ -222,9 +222,8 @@ def usb_sram(dev:str) -> UOp: return usb_asm24(dev)[0x5000:0x5000 + 2 * HALF] # ***************** # 2. transfers -def usb_word(v:UOp|int, dt:DType) -> UOp: return v.cast(dt) if isinstance(v, UOp) else UOp.const(v, dt) def usb_stack(dt:DType, *vals:UOp|int) -> UOp: - return (r:=UOp.placeholder((len(vals) or 1,), dt, addrspace=AddrSpace.REG)).after(*[r.index(i).store(usb_word(v, dt)) for i, v in enumerate(vals)]) + return (r:=UOp.placeholder((len(vals) or 1,), dt, addrspace=AddrSpace.REG)).after(*[r.index(i).store(UOp.const(v, dt)) for i, v in enumerate(vals)]) def usb_fail(link:UOp, code:UOp) -> UOp: return link.index(UOp.const(5).valid(link.after(code).index(5).load().eq(0))).store(code.cast(dtypes.uint64)) def usb_ctrl(link:UOp, rtype:int, req:int, val:UOp|int, idx:UOp|int, data:UOp, n:UOp|int, timeout:int=1000) -> UOp: @@ -283,9 +282,10 @@ def usb_store(b:UOp, idx:UOp, v:UOp) -> UOp: # kernargs: poke on change stored = usb_store(b, idx.src[0], v.src[0]) for i, w in zip(idx.src[1:], v.src[1:]): stored = usb_store(b.after(stored), i, w) return stored - link, addr, v = usb_link(b.device).after(*usb_deps(b)), usb_addr(b, idx, v.dtype), v.bitcast(CDTYPE[v.dtype.itemsize]) + link, addr = usb_link(b.device).after(*usb_deps(b)), usb_addr(b, idx, v.dtype) + v = v.bitcast({d.itemsize: d for d in dtypes.uints}[v.dtype.itemsize]) # bitcast to uint raw bits if str(unwrap_view(b)[0].tag).startswith("kernargs"): - cache = UOp.alloc((int(idx.vmax - idx.vmin) + 1,), v.dtype, device=HCQ_RUNTIME_DEV.value).rtag("usb_arg_cache") + cache = UOp.alloc((int(idx.vmax - idx.vmin) + 1,), v.dtype, device=HCQ_RUNTIME_DEV.device).rtag("usb_arg_cache") return usb_patch(link, addr, v, patch(cache, [], bytes(v.dtype.itemsize * cache.max_numel())).index(idx - idx.vmin)) return usb_poke_word(link, addr, v) @@ -306,7 +306,7 @@ pm_usb_lower = PatternMatcher([ def usb_wire(size:UOp|int) -> UOp|int: return (size + 512 + SLOT - 1) // SLOT * SLOT # payload + sentinel block, slot aligned def usb_sentinel(n:UOp) -> UOp: return ((n & 0xFFFFFF) | 0x51000000).cast(dtypes.uint32) def usb_put(link:UOp, ptr:UOp, dt:DType, *vals:UOp|int) -> UOp: # store through a pointer - return ccall(libc.memcpy, ptr, usb_stack(dt, *vals).after(link).index(0), dt.itemsize * len(vals)).cast(dtypes.void) + return ccall(libc.memcpy, ptr, usb_stack(dt, *vals).after(link).index(0), dt.itemsize * len(vals)) @uopfunc def usb_reap(link:UOp, xfer:UOp) -> UOp: # poll while pending (0xff), any other status but completed fails the link @@ -358,7 +358,7 @@ def usb_recv(link:UOp, table:UOp, i:UOp, run:UOp, go:UOp, stage:UOp) -> UOp: # c released = usb_poke(link.after(armed), go, (i + run + 1).cast(dtypes.uint32)) received = usb_bulk(link.after(released), 0x81, stage.index(0), wire) lower = ccall(libc.memcpy, addr, stage.after(received).index(0), first.cast(dtypes.uint64)) - return ccall(libc.memcpy, addr + CHUNK, stage.after(lower).index(HALF), second.cast(dtypes.uint64)).cast(dtypes.void).sink() + return ccall(libc.memcpy, addr + CHUNK, stage.after(lower).index(HALF), second.cast(dtypes.uint64)).sink() # ***************** # 5. batch and encode @@ -371,8 +371,8 @@ def usb_hostaddr(host:UOp, dev:str) -> UOp: def usb_table(hosts:list[tuple[UOp, int]], n:int, win:int, dev:str) -> UOp: # [address, bytes] per chunk rows = [(16 * (k + r) + o, w) for host, k in hosts for r, nb in chunks(host.nbytes(), win) - for o, w in ((0, usb_hostaddr(host, dev) + usb_word(r, dtypes.uint64) * win), (8, UOp.const(nb, dtypes.uint64)))] - return patch(UOp.alloc((2 * n,), dtypes.uint64, device=HCQ_RUNTIME_DEV.value).rtag("usb_table"), rows) + for o, w in ((0, usb_hostaddr(host, dev) + UOp.const(r, dtypes.uint64) * win), (8, UOp.const(nb, dtypes.uint64)))] + return patch(UOp.alloc((2 * n,), dtypes.uint64, device=HCQ_RUNTIME_DEV.device).rtag("usb_table"), rows) def usb_chunks(call:UOp, first:int, run:int) -> list[UOp]: # gpu side. first: chunk id of the copy and of its run dst, src = call.src[1:] @@ -380,7 +380,7 @@ def usb_chunks(call:UOp, first:int, run:int) -> list[UOp]: # gpu side. first: ch sram = usb_sram(dev:=vram.device).getaddr(dev) for r, nb in chunks(host.nbytes(), win): - n, va = (i:=usb_word(r, dtypes.uint64)) + first, vram.getaddr(dev) + i * win + n, va = (i:=UOp.const(r, dtypes.uint64)) + first, vram.getaddr(dev) + i * win if is_host(src): # copyin: wait sentinel, copy, release end = sram + (((n - run) & 1) + 1) * HALF cmds = [ins("wait_eq", end - 4, usb_sentinel(n)), ins("copy", va, end - usb_wire(nb), nb), ins("store", end - 4, 0)] diff --git a/tinygrad_repo/tinygrad/schedule/__init__.py b/tinygrad_repo/tinygrad/schedule/__init__.py index 4730e14ef..d8227b822 100644 --- a/tinygrad_repo/tinygrad/schedule/__init__.py +++ b/tinygrad_repo/tinygrad/schedule/__init__.py @@ -86,21 +86,19 @@ from tinygrad.schedule.rangeify import get_kernel_graph from tinygrad.helpers import CAPTURING from tinygrad.uop.ops import PatternMatcher, UPat -def create_new_buffer(ctx:tuple[dict[UOp, UOp], tuple[UOp, ...]], b:UOp): - if (ret:=ctx[0].get(b, None)) is None: - device = b.device if b.device is not None else next(a.device for a in ctx[1] if a.device is not None) - ctx[0][b] = ret = UOp.new_buffer(device, b.max_numel(), b.dtype) - return ret +def create_new_buffer(ctx:tuple[UOp, ...], b:UOp): + device = b.device if b.device is not None else next(a.device for a in ctx if a.device is not None) + return UOp.new_buffer(device, b.max_numel(), b.dtype) pm_post_sched_cache = PatternMatcher([ # Resolve positional arguments outside kernel bodies; free Variables have slot -1. - (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg.slot] if x.arg.slot >= 0 else None), + (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[x.arg.slot] if x.arg.slot >= 0 else None), # bind ALLOCs to fresh BUFFERs for this invocation (UPat(Ops.ALLOC, name="b"), create_new_buffer), ]) def resolve_linear_call(linear_call:UOp, outer_binds:dict[int, UOp]|None=None): - linear = graph_rewrite(linear_call.body, pm_post_sched_cache, ctx=({}, linear_call.src[1:]), walk=True, name="params to buffers") + linear = graph_rewrite(linear_call.body, pm_post_sched_cache, ctx=linear_call.src[1:], walk=True, name="params to buffers") # nested LINEAR calls are lexical scopes: their positional params shadow the enclosing scope, while calls without # scalar args (e.g. precompiled allreduce) inherit it binds = {**(outer_binds or {}), diff --git a/tinygrad_repo/tinygrad/schedule/allreduce.py b/tinygrad_repo/tinygrad/schedule/allreduce.py index aba5603e2..b92aa5d92 100644 --- a/tinygrad_repo/tinygrad/schedule/allreduce.py +++ b/tinygrad_repo/tinygrad/schedule/allreduce.py @@ -1,6 +1,6 @@ import functools, itertools from tinygrad.helpers import all_int, prod, DEBUG, RING, ALL2ALL, ALLREDUCE_NODE_NDEVS, getenv -from tinygrad.uop.ops import UOp, Ops, ParamArg +from tinygrad.uop.ops import UOp # *** allreduce implementation *** def handle_allreduce(buf:UOp, red:UOp) -> UOp|None: @@ -67,9 +67,7 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None: def create_allreduce_function(buf:UOp, red:UOp, output:UOp|None=None) -> UOp|None: if output is None: - output = UOp(Ops.ALLOC, src=(UOp.const(red.max_numel()),)+UOp.device_range_src(red.device), - arg=ParamArg(next(UOp.unique_num), red.dtype, device=red.device)) - output = output.reshape(red.max_shape).shrink_to(red.shape) + output = UOp.alloc(red.shape, red.dtype, device=red.device) to = red.param_like(0) src = buf.param_like(1) red = src.allreduce(*red.arg) diff --git a/tinygrad_repo/tinygrad/schedule/indexing.py b/tinygrad_repo/tinygrad/schedule/indexing.py index 3a705e6a5..2d50c5f05 100644 --- a/tinygrad_repo/tinygrad/schedule/indexing.py +++ b/tinygrad_repo/tinygrad/schedule/indexing.py @@ -33,7 +33,7 @@ def realize_srcs(ctx:IndexingContext, rb:UOp) -> None: def realize_store_after_src(ctx:IndexingContext, dest:UOp, src:UOp): # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce - if dest.base in src.toposort(enter_calls=False): ctx.realize_map[src] = None + if dest.base in src.toposort(): ctx.realize_map[src] = None def realize_custom_kernel_srcs(ctx:IndexingContext, c:UOp) -> None: for s in c.src[1:]: @@ -171,7 +171,6 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO match op: case Ops.SHRINK: rngs = tuple(a if off == 0 else a+off for a,(off,_) in zip(rngs, arg)) case Ops.PERMUTE: rngs = tuple(rngs[p] for p in argsort(arg)) - case Ops.FLIP: rngs = tuple(((s-1)-a) if f else a for a,s,f in zip(rngs, in_shape, arg)) case Ops.EXPAND: rngs = rngs[len(arg):] case Ops.PAD: # NOTE: the .where(r-s, i) is not inside the graph_rewrite so that `convert_pad_to_where_to_keep_behavior_local` @@ -180,7 +179,7 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO symbolic+pm_simplify_valid, name="pad")) for r,sh,(off,sz) in zip(rngs, in_shape, arg)) case Ops.RESHAPE: sink = UOp.sink(*rngs).simplify() # NOTE: this applies any commutative flips to the rngs early - sub_array = {r:r.replace(src=r.src[:1], arg=(AxisType.PLACEHOLDER, i)) for i,r in enumerate(sink.ranges)} + sub_array = {r:r.replace(arg=(AxisType.PLACEHOLDER, i)) for i,r in enumerate(sink.ranges)} rngs = _apply_reshape(in_shape, arg, sink.substitute(sub_array)).substitute({v:k for k,v in sub_array.items()}).src case _: raise RuntimeError(f"{op} is not a MovementOp") return rngs diff --git a/tinygrad_repo/tinygrad/schedule/memory.py b/tinygrad_repo/tinygrad/schedule/memory.py index 3a734727a..a13749ee3 100644 --- a/tinygrad_repo/tinygrad/schedule/memory.py +++ b/tinygrad_repo/tinygrad/schedule/memory.py @@ -1,5 +1,5 @@ from collections import defaultdict -from tinygrad.helpers import NO_MEMORY_PLANNER, DEBUG, round_up +from tinygrad.helpers import NO_MEMORY_PLANNER, DEBUG, round_up, to_tuple from tinygrad.uop.ops import UOp, Ops from tinygrad.dtype import dtypes from tinygrad.runtime.support.memory import TLSFAllocator @@ -11,9 +11,8 @@ def _collect_bufs(u:UOp) -> list[UOp]: def _can_plan(b:UOp, held_bufs:set[UOp]) -> bool: if b in held_bufs: return False - devs = (b.device,) if isinstance(b.device, str) else b.device # CL and WEBGPU do not support views, see explanation in contiguous_view_offset - return all(not d.startswith(("DISK", "CL", "WEBGPU")) for d in devs) + return all(not d.startswith(("DISK", "CL", "WEBGPU")) for d in to_tuple(b.device)) LaneKey = tuple[str, int] diff --git a/tinygrad_repo/tinygrad/schedule/multi.py b/tinygrad_repo/tinygrad/schedule/multi.py index 618ccb631..758753c7d 100644 --- a/tinygrad_repo/tinygrad/schedule/multi.py +++ b/tinygrad_repo/tinygrad/schedule/multi.py @@ -126,7 +126,7 @@ def reshape_multi(root:UOp, multi:UOp): # map every sharded axis through the reshape: the axis boundary must survive intact and stay divisible by its shard count arg_acc:list[sint] = [1] for s in new_shape: arg_acc.append(ssimplify(arg_acc[-1]*s)) - new_shardings = [] + new_shardings, new_counts = [], {} for ax, rng in multi.sharding: count = int(rng.vmax)+1 target = ssimplify(prod(multi.shape[:ax])) @@ -134,8 +134,8 @@ def reshape_multi(root:UOp, multi:UOp): new_ax = len(arg_acc) - arg_acc[::-1].index(target) - 1 if new_shape[new_ax] % count != 0: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards") new_shardings.append((new_ax, rng)) - new_axs = {a for a, _ in new_shardings} - new_shape = tuple(s//(int(rng.vmax)+1) if a in new_axs else s for a,s in enumerate(new_shape)) + new_counts[new_ax] = count + new_shape = tuple(s//new_counts.get(a, 1) for a,s in enumerate(new_shape)) return multi.src[0].reshape(new_shape).unshard(tuple(a for a,_ in new_shardings), tuple(r for _,r in new_shardings)) def expand_multi(root:UOp, multi:UOp): @@ -179,11 +179,6 @@ def shrink_multi(root:UOp, multi:UOp): val = multi.src[0]._mop(Ops.SHRINK, tuple(local_marg)) return val if not remaining else val.unshard(tuple(a for a,_ in remaining), tuple(r for _,r in remaining)) -def flip_multi(root:UOp, multi:UOp): - for ax, _ in multi.sharding: - if root.marg[ax]: raise RuntimeError(f"flipping not supported on sharded axis {ax}") - return multi.src[0].flip([i for i,x in enumerate(root.marg) if x]).unshard(multi.arg, multi.src[1:]) - def stack_multi(root:UOp): # STACK adds a leading axis: srcs are sharded one axis below the output multis = [m for m in root.src if m.op is Ops.UNSHARD] @@ -284,7 +279,6 @@ multi_pm = PatternMatcher([ (UPat(Ops.PAD, src=(UPat(Ops.UNSHARD, name="multi"), UPat(), UPat()), name="root"), pad_multi), (UPat(Ops.SHRINK, src=(UPat(Ops.UNSHARD, name="multi"), UPat(), UPat()), name="root"), shrink_multi), (UPat(Ops.PERMUTE, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), permute_multi), - (UPat(Ops.FLIP, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), flip_multi), (UPat(Ops.STACK, name="root", custom_early_reject=set([Ops.UNSHARD])), stack_multi), (UPat(Ops.INDEX, src=(UPat(Ops.UNSHARD, name="multi"),), name="root", allow_any_len=True), index_multi), (UPat(Ops.AFTER, src=(UPat(Ops.UNSHARD), UPat(Ops.STORE, src=(UPat(Ops.UNSHARD, name="dest"), UPat(Ops.UNSHARD, name="src"))))), store_after_multi), diff --git a/tinygrad_repo/tinygrad/schedule/prepare.py b/tinygrad_repo/tinygrad/schedule/prepare.py index 4db6c5d8a..82c514681 100644 --- a/tinygrad_repo/tinygrad/schedule/prepare.py +++ b/tinygrad_repo/tinygrad/schedule/prepare.py @@ -1,6 +1,6 @@ from dataclasses import replace from tinygrad.dtype import dtypes, to_dtype -from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, GroupOp, ParamArg +from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, GroupOp from tinygrad.uop.ops import graph_rewrite, rewrite_group, identity_element, resolve_returned_after from tinygrad.uop.movement import mop_cleanup from tinygrad.helpers import prod, getenv, all_int, DEBUG, SPLIT_REDUCEOP, OPENPILOT_HACKS, FLOAT16, argsort @@ -25,7 +25,7 @@ def forward_call_outputs(sink:UOp) -> UOp: # Forward the allocation, not just one view of it, so saved values and other aliases follow the same placement. key = base if base.op is Ops.ALLOC else src if key not in placed and (src.op is Ops.STAGE or target.has_buffer_identity()) and \ - target.storage_base not in st.src[1].toposort(enter_calls=False): + target.storage_base not in st.src[1].toposort(): if base.op is Ops.ALLOC and src.has_buffer_identity() and base.max_numel() == target.storage_base.max_numel(): placed[key] = target.storage_base elif src.op is Ops.STAGE: placed[key] = target.after(target.store(src.src[0])) @@ -76,6 +76,11 @@ def _mop_index(r:UOp, idx:UOp): pm_mops = PatternMatcher([ # handle movement ops on INDEX (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), _mop_index), + # A store through a padded view gates the destination index, not its loaded value. + (UPat.var("a").where(UPat.var("b").where(UPat.var("dst"), 0), UPat.var("pad")).store(UPat.var("val")), + lambda a,b,dst,pad,val: (a&b).where(dst, 0).store(val) if pad.is_invalid else None), + (UPat.var("cond").where(UPat.var("buf").index(UPat.var("idx")), 0).store(UPat.var("val")), + lambda cond,buf,idx,val: buf.index(idx.valid(cond)).store(val)), # move movement ops and INDEX after AFTER (UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True), lambda r,a: UOp(r.op, src=(a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], arg=r.arg)), @@ -92,9 +97,9 @@ def store_hazard_boundary(s:UOp): return True def fix_store_hazard(target:UOp, src:UOp): - if (base:=target.base) not in src.toposort(enter_calls=False): return None - # PERMUTE and FLIP reorder indices, SHRINK can have overlapping regions when dest is also shrunk - unsafe = {Ops.PERMUTE, Ops.FLIP} | ({Ops.SHRINK} if target.op_in_backward_slice_with_self(Ops.SHRINK) else set()) + if (base:=target.base) not in src.toposort(): return None + # PERMUTE reorders indices, SHRINK can have overlapping regions when dest is also shrunk + unsafe = {Ops.PERMUTE} | ({Ops.SHRINK} if target.op_in_backward_slice_with_self(Ops.SHRINK) else set()) reaches_base: dict[UOp, bool] = {} for s in src.toposort(gate=store_hazard_boundary): reaches_base[s] = s is base or any(reaches_base.get(c) for c in s.src) @@ -126,7 +131,7 @@ def split_reduceop(reduce:UOp, x:UOp): def resolve_function(c:UOp) -> UOp|None: if not c.is_inline_call: return None - nodes = c.body.toposort(enter_calls=False) + nodes = c.body.toposort() # Input and output PARAMs both bind to explicit arguments by slot; unused arguments are allowed. args = c.src[1:] @@ -140,7 +145,7 @@ def resolve_function(c:UOp) -> UOp|None: for p, a in dict_map.items(): if p.shape: n, flat = flat_storage(a) - if p.src[0].val != n: raise TypeError(f"arg {p.arg.slot} shape mismatch: expected size {p.src[0].val}, got {a.shape}") + if p.max_numel() != n: raise TypeError(f"arg {p.arg.slot} shape mismatch: expected size {p.max_numel()}, got {a.shape}") dict_map[p] = flat elif a.shape != (): raise TypeError(f"arg {p.arg.slot} shape mismatch: expected scalar, got {a.shape}") @@ -165,14 +170,13 @@ def copy_to_anon_store(x:UOp, copy:UOp): # copies are always cross device: pad to the max shape so the copy reads a whole buffer (SDMA can't do offset copies) x = x.pad_to(x.max_shape) # the buffer takes the DEVICE range from the copy (no-op for single device copies) - buf = UOp(Ops.ALLOC, src=(UOp.const(prod(x.max_shape)),)+copy.src[1:], - arg=ParamArg(next(UOp.unique_num), copy.dtype, device=copy.device)).reshape(x.max_shape) + buf = UOp.alloc((x.max_numel(),), copy.dtype, device=copy.device) + buf = buf.replace(src=buf.src[:1]+copy.src[1:]).reshape(x.max_shape) return buf.after(buf.store(x)).shrink_to(copy.shape) def stage_to_anon_store(x:UOp, stg:UOp): # the buffer created here is inside the call and is not persisted, like the buffers created for copies - buf = UOp(Ops.ALLOC, src=(UOp.const(prod(x.max_shape)),)+UOp.device_range_src(x.device), - arg=ParamArg(next(UOp.unique_num), stg.dtype, device=x.device)).reshape(x.max_shape) + buf = UOp.alloc(x.max_shape, stg.dtype, device=x.device) view = buf.shrink_to(stg.shape) return view.after(view.store(x)) diff --git a/tinygrad_repo/tinygrad/schedule/rangeify.py b/tinygrad_repo/tinygrad/schedule/rangeify.py index 26e28fa1d..bb4d8308a 100644 --- a/tinygrad_repo/tinygrad/schedule/rangeify.py +++ b/tinygrad_repo/tinygrad/schedule/rangeify.py @@ -2,7 +2,7 @@ from dataclasses import dataclass, field, replace from typing import cast import itertools from tinygrad.dtype import dtypes, AddrSpace, Invalid -from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg +from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group from tinygrad.uop.symbolic import symbolic from tinygrad.helpers import prod, dedup, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC @@ -222,7 +222,7 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp): return buf.after(*ended_stores) if x.arg.addrspace == AddrSpace.GLOBAL: - buf = UOp(Ops.ALLOC, src=(UOp.const(size),)+UOp.device_range_src(x.arg.device), arg=ParamArg(next(ctx), dtype, device=x.arg.device)) + buf = UOp.alloc((size,), dtype, slot=next(ctx), device=x.arg.device) do_store = buf.index(idx).store(x.src[0].cast(dtype)).end(*rngs) return buf.after(do_store).cast(x.dtype) @@ -276,7 +276,7 @@ class LocalAddBufferContext: range:int = 0 def debuf(ctx:LocalAddBufferContext, buf:UOp): - param = UOp(Ops.PARAM, src=(UOp.const(prod(buf.max_shape)),), arg=ParamArg(ctx.dg, buf.dtype, addrspace=buf.addrspace, device=buf.device)) + param = UOp.param(ctx.dg, buf.dtype, buf.max_numel(), addrspace=buf.addrspace, device=buf.device) ret = param.reshape(buf.max_shape) # if the buffer has symbolic shape, shrink the max-sized view to the actual shape if buf.max_shape != buf.shape: ret = ret.shrink(tuple((0, s) for s in buf.shape)) @@ -367,5 +367,5 @@ def get_kernel_graph(tsink:UOp) -> UOp: if SPEC: # validate the kernel graph from tinygrad.uop.spec import type_verify, spec_kernel_graph - type_verify(tsink, spec_kernel_graph, enter_calls=False) + type_verify(tsink, spec_kernel_graph) return tsink diff --git a/tinygrad_repo/tinygrad/tensor.py b/tinygrad_repo/tinygrad/tensor.py index f8ac0f292..cdad2de29 100644 --- a/tinygrad_repo/tinygrad/tensor.py +++ b/tinygrad_repo/tinygrad/tensor.py @@ -180,7 +180,7 @@ class Tensor(RandMixin): # Rebuild in dependency order: replacement values already reference the other outputs' storage. tensor_map:dict[UOp, UOp] = {} - for x in sink.toposort(enter_calls=False): + for x in sink.toposort(): u = x.replace(src=tuple(tensor_map.get(s, s) for s in x.src)) if x.op is Ops.ALLOC and (x.arg.bind_on_realize or x in bases): u = UOp.new_buffer(x.device, x.max_numel(), x.dtype) elif x in bases and u.needs_storage(): @@ -202,7 +202,7 @@ class Tensor(RandMixin): # Realized outputs become the storage their AFTER sequenced a store into. Compose with tensor_map before updating # Tensors so map values reference final storage. becomes_map = {u:graph_rewrite(u.src[0], pm_drop_after, bottom_up=True, name="drop after").shrink_to(u.shape) - for u in sink.toposort(enter_calls=False) if is_store_after(u)} + for u in sink.toposort() if is_store_after(u)} assert not any(x in becomes_map for x in becomes_map.values()) tensor_map = dict(zip(tensor_map, UOp.sink(*tensor_map.values()).substitute(becomes_map, walk=True).src)) _apply_map_to_tensors(becomes_map | tensor_map, name="bufferize") diff --git a/tinygrad_repo/tinygrad/uop/__init__.py b/tinygrad_repo/tinygrad/uop/__init__.py index dd785ec64..7991eae74 100644 --- a/tinygrad_repo/tinygrad/uop/__init__.py +++ b/tinygrad_repo/tinygrad/uop/__init__.py @@ -89,8 +89,8 @@ class Ops(FastEnum): # buffer ops STAGE = auto(); COPY = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto() - # the core 6 movement ops! these only exist in the tensor graph - RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); FLIP = auto() + # the core 5 movement ops! these only exist in the tensor graph + RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto() UNSHARD = auto() # UNSHARD is really a movement op # reduce @@ -114,7 +114,7 @@ class GroupOp: Defines = {Ops.PARAM, Ops.BUFFER, Ops.ALLOC} Irreducible = {Ops.CONST, Ops.SPECIAL, Ops.RANGE, Ops.PARAM, Ops.GETADDR} - Movement = {Ops.RESHAPE, Ops.EXPAND, Ops.PERMUTE, Ops.PAD, Ops.SHRINK, Ops.FLIP} + Movement = {Ops.RESHAPE, Ops.EXPAND, Ops.PERMUTE, Ops.PAD, Ops.SHRINK} # BinaryOps that can be flipped Commutative = {Ops.ADD, Ops.MUL, Ops.MAX, Ops.CMPNE, Ops.CMPEQ, Ops.XOR, Ops.AND, Ops.OR} diff --git a/tinygrad_repo/tinygrad/uop/ops.py b/tinygrad_repo/tinygrad/uop/ops.py index 6c59f3f2b..f07de64a5 100644 --- a/tinygrad_repo/tinygrad/uop/ops.py +++ b/tinygrad_repo/tinygrad/uop/ops.py @@ -212,7 +212,7 @@ class recursive_property: self.__doc__ = fxn.__doc__ def __get__(self, x:UOp|None, owner=None): if x is None: return self - for node in x.toposort(gate=lambda node: self.nm not in node.__dict__): node.__dict__[self.nm] = self.fxn(node) + for node in x.toposort(gate=lambda node: self.nm not in node.__dict__, enter_calls=True): node.__dict__[self.nm] = self.fxn(node) return x.__dict__[self.nm] class _UOpTuple(tuple): @@ -267,7 +267,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): @functools.cached_property def backward_slice(self:UOp) -> dict[UOp, None]: - res: dict[UOp, None] = self.toposort(enter_calls=False) + res: dict[UOp, None] = self.toposort() res.pop(self) return res @@ -283,7 +283,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): @property def bool_slice(self) -> frozenset[UOp]: return self._bool_slice | {self} if self.dtype is dtypes.bool else self._bool_slice - def toposort(self, gate:Callable|None=None, enter_calls=True) -> dict[UOp, None]: + def toposort(self, gate:Callable|None=None, enter_calls=False) -> dict[UOp, None]: cache: dict[UOp, None] = {} stack: list[tuple[UOp, bool]] = [(self, False)] # each stack entry is (node, visited_flag) while stack: @@ -409,9 +409,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if len(ps) != len(self.marg) or not all(resolve(0<=o) and resolve(sz>=0) and resolve(o+sz<=s) for s,(o,sz) in zip(ps, self.marg)): raise ValueError(f"invalid shrink {self.marg} for {ps}") return tuple(sz for _,sz in self.marg) - case Ops.FLIP: - if len(ps) != len(self.marg) or not all(isinstance(x, bool) for x in self.marg): raise ValueError(f"bad flip on {ps}, {self.marg}") - return ps case Ops.UNSHARD: return tuple(s*(int(self.src[1:][self.arg.index(a)].vmax)+1) if a in self.arg else s for a,s in enumerate(ps)) case Ops.REDUCE: num_axes = self.arg[1] @@ -793,7 +790,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): match self.op: case Ops.RESHAPE | Ops.EXPAND: return self.src[1].as_shape case Ops.PAD | Ops.SHRINK: return tuple(zip(self.src[1].as_shape, self.src[2].as_shape)) - case Ops.PERMUTE | Ops.FLIP: return self.arg + case Ops.PERMUTE: return self.arg case _: raise RuntimeError(f"{self.op} is not a MovementOp") def _mop(self, op:Ops, arg) -> UOp: @@ -805,7 +802,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): match op: case Ops.RESHAPE | Ops.EXPAND: src_args = [arg] case Ops.PAD | Ops.SHRINK: src_args = list(zip(*arg)) - case Ops.PERMUTE | Ops.FLIP: src_args = [] + case Ops.PERMUTE: src_args = [] case Ops.STACK: srcs = (self,)+tuple(arg) dtype = dtype_from_uop(Ops.STACK, srcs, None) @@ -828,21 +825,24 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # BUFFER/ALLOC/COPY carry a DEVICE range when targeting multiple devices return (UOp.range(len(device), 0, AxisType.DEVICE),) if isinstance(device, tuple) else () @staticmethod + def _buffer(size:int, dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL, device=None, buffer:Buffer|MultiBuffer|None=None) -> UOp: + return UOp(Ops.BUFFER, src=(shape_to_shape_arg((size,)),)+UOp.device_range_src(device), + arg=ParamArg(slot, dtype, addrspace=addrspace, device=device, buffer=buffer)) + @staticmethod def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None): if dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {dtype}") assert dtype.itemsize > 0, f"not a storage dtype: {dtype}" assert isinstance(size, int), f"new_buffer size must be a concrete int, got {size}" slot = next(UOp.unique_num) if num is None else num buf = MultiBuffer(device, size * dtype.itemsize) if isinstance(device, tuple) else Buffer(device, size * dtype.itemsize) - return UOp(Ops.BUFFER, src=(UOp.const(size),)+UOp.device_range_src(device), arg=ParamArg(slot, dtype, device=device, buffer=buf)) + return UOp._buffer(size, dtype, slot, device=device, buffer=buf) @staticmethod def from_buffer(opaque:Buffer|MultiBuffer, dtype:DType, device:str|tuple[str, ...]|None=None): if dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {dtype}") assert dtype.itemsize > 0, f"not a storage dtype: {dtype}" assert opaque.nbytes % dtype.itemsize == 0, f"{opaque.nbytes} bytes is not divisible by {dtype.itemsize} for {dtype}" # Both storage identity and interpretation are part of the UOp key. - return UOp(Ops.BUFFER, src=(UOp.const(opaque.nbytes // dtype.itemsize),)+UOp.device_range_src(device or opaque.device), - arg=ParamArg(-id(opaque), dtype, device=device or opaque.device, buffer=opaque)) + return UOp._buffer(opaque.nbytes // dtype.itemsize, dtype, -id(opaque), device=device or opaque.device, buffer=opaque) def empty_like(self, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None=None) -> UOp: device = canonicalize_device(self.device if device is None else device) axis = self.axis if isinstance(device, tuple) else None @@ -880,6 +880,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op in {Ops.BUFFER, Ops.ALLOC}: return self.arg.device if self.op is Ops.COPY: return self.arg if self.op is Ops.ALLREDUCE: return self.arg[1] + if self.op is Ops.CALL and self.body.op is Ops.STORE: return self.src[2].device # a copy pushes from its src for x in self.src: if x.device is not None: return x.device return None @@ -922,7 +923,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # CL 1.1 provides the clCreateSubBuffer API, but at the time of writing, relevant CL runtimes (rusticl, adreno, nvidia, amd) do not provide # reasonable values for CL_DEVICE_MEM_BASE_ADDR_ALIGN. cl_ext_buffer_device_address could potentially help, but this extension is not provided # by relevant CL runtimes at time of writing. - if (dev:=self.device) is not None and any(d.startswith(("WEBGPU", "CL")) for d in ((dev,) if isinstance(dev, str) else dev)): return None + if (dev:=self.device) is not None and any(d.startswith(("WEBGPU", "CL")) for d in to_tuple(dev)): return None rng = UOp.range(self.numel(), 0) out = graph_rewrite(self.flatten().index(rng), pm_mops+symbolic, name="contiguous_view_offset") @@ -959,10 +960,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def buffer(self) -> Buffer|MultiBuffer: # a bare STAGE (same-device materialization) keeps the source's buffer if self.op is Ops.STAGE and self.arg is None: return self.src[0].buffer - if self.op in {Ops.CONTIGUOUS_BACKWARD, Ops.RESHAPE, Ops.UNSHARD, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer + if self.op in {Ops.CONTIGUOUS_BACKWARD, Ops.RESHAPE, Ops.UNSHARD, Ops.DETACH, Ops.AFTER, Ops.BITCAST}: return self.src[0].buffer # this buffer can process disk tensors and simple movement ops. # NOTE: the view Buffer returned here is transient (short-lived), it only wraps an offset into the base BUFFER's storage - if self is not self.base or self.op is Ops.BITCAST: + if self is not self.base: src, offset = self._buffer_view if isinstance(buf:=src.buffer, MultiBuffer): mbuf = MultiBuffer.__new__(MultiBuffer) @@ -1120,6 +1121,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if s1_vmax < 0: return (s1_vmin+1, 0) if self.op is Ops.XOR and s1_vmin == s1_vmax == -1 and isinstance(s0_vmin, int) and isinstance(s0_vmax, int): return ~int(s0_vmax), ~int(s0_vmin) + if self.op is Ops.XOR and dtypes.is_int(self.dtype) and s0_vmin >= 0 and s1_vmin >= 0: + return 0, (1 << int(max(s0_vmax, s1_vmax)).bit_length()) - 1 if self.op is Ops.MAX: return max(s0_vmin, s1_vmin), max(s0_vmax, s1_vmax) if self.op is Ops.CMPLT: return (s0_vmax 1: ret = ret.reshape(shape) return ret def placeholder_like(self, slot:int, addrspace=AddrSpace.GLOBAL): @@ -1225,13 +1228,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return self.replace(arg=replace(self.arg, slot=slot, name=name, val=None)) # multi-device values become a per-shard sized param wrapped in UNSHARD: the sharding lives in the graph, not the arg if self.axis is not None and isinstance(self.device, tuple): - return UOp(Ops.PARAM, src=(UOp.const(prod(to_max_shape(self.shard_shape))),), - arg=ParamArg(slot, self.dtype, name=name, device=self.device)).view_as(self.shard_shape, self.axis) + return UOp.param(slot, self.dtype, prod(to_max_shape(self.shard_shape)), self.device, name=name).view_as(self.shard_shape, self.axis) return UOp.param(slot, self.dtype, self._shape, self.device, name=name) def view_as(self:UOp, shape:tuple[sint, ...], axis:int|None=None) -> UOp: """view flat storage as the given (possibly symbolic) shape, optionally sharded on axis, the UNSHARD gives back the multiplied shape""" max_shape = to_max_shape(shape) - ret = self.reshape(max_shape) if len(shape) > 1 else self + ret = self.reshape(max_shape) if shape else self if tuple(max_shape) != tuple(shape): ret = ret.shrink_to(shape) return ret if axis is None else ret.unshard(axis) @@ -1357,7 +1359,7 @@ class ProgramInfo: ins: list[int] = [] global_size: list[int] = [1, 1, 1] local_size: list[int] = [1, 1, 1] - for u in sink.toposort(enter_calls=False): + for u in sink.toposort(): if u.op is Ops.PARAM and u.addrspace == AddrSpace.ALU: _vars.append(u) if u.op is Ops.PARAM and u.addrspace != AddrSpace.ALU: _globals.append(u.arg.slot) if u.op in (Ops.STORE, Ops.LOAD): diff --git a/tinygrad_repo/tinygrad/uop/render.py b/tinygrad_repo/tinygrad/uop/render.py index 8817e91dc..03066faac 100644 --- a/tinygrad_repo/tinygrad/uop/render.py +++ b/tinygrad_repo/tinygrad/uop/render.py @@ -63,7 +63,7 @@ def _render_arg(x:UOp) -> str: def _inline(u:UOp) -> bool: return u.op is Ops.CONST or (u.op is Ops.STACK and all(s.op is Ops.CONST for s in u.src)) def render_uir(root:UOp|list[UOp]) -> str: - nodes = [u for u in (list(root.toposort()) if isinstance(root, UOp) else list(root)) if not _inline(u)] + nodes = [u for u in (list(root.toposort(enter_calls=True)) if isinstance(root, UOp) else list(root)) if not _inline(u)] table = {u:i for i,u in enumerate(nodes)} def src_str(u:UOp) -> str: if not _inline(u): return f"%{table[u]}" @@ -92,7 +92,6 @@ def marg_str(ctx, a:sint) -> str: return str(a) if not isinstance(a, UOp) else c def render_marg(ctx,x:UOp): if x.op is Ops.PERMUTE: return str(x.marg) - if x.op is Ops.FLIP: return str(tuple([i for i,x in enumerate(x.marg) if x])) pieces = [] if x.op in {Ops.RESHAPE, Ops.EXPAND}: pieces = [marg_str(ctx, a) for a in x.marg] if x.op in {Ops.PAD, Ops.SHRINK}: pieces = [f"({marg_str(ctx, a[0])}, {marg_str(ctx, a[1])})" for a in x.marg] @@ -148,7 +147,6 @@ def _render_with_splits(lst:list[UOp], pm:PatternMatcher, to_render:set[UOp], sp depth[u] = 0 if u in to_render else op_depth ren = pm.rewrite(u, ctx=r) assert isinstance(ren, str) - if u.tag is not None: ren += f".rtag({repr(u.tag)})" if u not in to_render: r[u] = ren else: r[u] = f"c{i}" if u is not lst[-1] else "ast" diff --git a/tinygrad_repo/tinygrad/uop/spec.py b/tinygrad_repo/tinygrad/uop/spec.py index cb2d3479d..66d63d6fc 100644 --- a/tinygrad_repo/tinygrad/uop/spec.py +++ b/tinygrad_repo/tinygrad/uop/spec.py @@ -1,7 +1,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, AxisType, ParamArg, CallInfo, OPAQUE_CALL_BODIES, \ CustomFunction from tinygrad.uop.render import render_uir -from tinygrad.dtype import DType, dtypes, AddrSpace, Invalid +from tinygrad.dtype import DType, dtypes, AddrSpace, Invalid, weak_dtype from tinygrad.helpers import DEBUG, Context, CHECK_OOB, all_same, is_image_shape from tinygrad.device import is_disk_device @@ -31,7 +31,7 @@ def valid_device_range(device:str|tuple[str, ...]|None, src:tuple[UOp, ...]) -> rng = src[0] return rng.op is Ops.RANGE and rng.axis_type is AxisType.DEVICE and int(rng.vmax)+1 == len(device) -def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher, enter_calls=True): +def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher, enter_calls=False): lst = list(ast.toposort(enter_calls=enter_calls)) if isinstance(ast, UOp) else ast with Context(TRACK_MATCH_STATS=0): @@ -58,20 +58,20 @@ spec_shared = PatternMatcher([ # STACK is everywhere too (UPat(Ops.STACK, dtype=dtypes.void, src=()), lambda: True), (UPat(Ops.STACK, src=(UPat(),), allow_any_len=True, name="s"), - lambda s: all_same([x.shape for x in s.src]) and all(matches_dtype(x, s.dtype) or x.dtype in dtypes.weaks for x in s.src)), + lambda s: all_same([x.shape for x in s.src]) and all(matches_dtype(x, s.dtype) or x.dtype == weak_dtype(s.dtype) for x in s.src)), # ALUs: operands match the result dtype, except comparisons/WHERE; renderer-lowered shifts may use a uint32 count - # a weak dtype matches any dtype until lowering commits its operand + # a weak dtype matches any dtype of its kind until lowering commits its operand (UPat(Ops.WHERE, name="w", src=(UPat(dtype=dtypes.bool), UPat(), UPat())), - lambda w: all(matches_dtype(s, w.dtype) or s.dtype in dtypes.weaks for s in w.src[1:])), + lambda w: all(matches_dtype(s, w.dtype) or s.dtype == weak_dtype(w.dtype) for s in w.src[1:])), (UPat(GroupOp.Comparison, dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))), - lambda x,y: matches_dtype(x, y.dtype) or matches_dtype(y, x.dtype) or x.dtype in dtypes.weaks or y.dtype in dtypes.weaks), + lambda x,y: matches_dtype(x, y.dtype) or matches_dtype(y, x.dtype) or x.dtype == weak_dtype(y.dtype) or y.dtype == weak_dtype(x.dtype)), (UPat((Ops.AND, Ops.OR, Ops.XOR, Ops.SHL, Ops.SHR), name="x"), lambda x: False if any(dtypes.is_float(s.dtype) for s in x.src) else None), (UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat.var("c")), name="a"), lambda a,x,c: matches_dtype(c, a.dtype) or c.dtype in (dtypes.uint, dtypes.weakint) or x.base.is_invalid), (UPat((Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD), name="x"), lambda x: None if dtypes.is_int(x.dtype) or any(s.base.is_invalid for s in x.src) else False), - (UPat(GroupOp.ALU, name="x"), lambda x: all(matches_dtype(y, x.dtype) or y.dtype in dtypes.weaks for y in x.src)), + (UPat(GroupOp.ALU, name="x"), lambda x: all(matches_dtype(y, x.dtype) or y.dtype == weak_dtype(x.dtype) for y in x.src)), # CAST (UPat((Ops.BITCAST, Ops.CAST), src=(UPat(),), name="x"), lambda x: isinstance(x.arg, DType)), @@ -121,7 +121,9 @@ spec_shared = PatternMatcher([ # assembly instruction (UPat(Ops.INS, name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 2 and isinstance(x.arg[1], DType)), - # LOAD(idx) / STORE(idx, val) with gates on the LOAD/STORE + # LOAD(idx) / STORE(idx, val) with gates on the LOAD/STORE; a STORE's value matches its target's dtype + (UPat(Ops.STORE, src=(UPat(name="t"), UPat(name="v")), allow_any_len=True), + lambda t,v: None if matches_dtype(v, t.dtype) or v.dtype == weak_dtype(t.dtype) else False), (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().load(), validate_index), (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().load(UPat.var("alt"), UPat.var("gate", dtype=dtypes.bool), name="load"), lambda uidx,gate,alt,load: validate_index(uidx, gate) if matches_dtype(alt, load.dtype) else False), @@ -161,9 +163,9 @@ spec_tensor = PatternMatcher([ (UPat(Ops.SPECIAL, src=(UPat(dtype=dtypes.weakint),), name="s"), lambda s: isinstance(s.arg, str)), # movement ops - (UPat((Ops.RESHAPE, Ops.EXPAND), src=(UPat(), UPat())), lambda: True), + (UPat((Ops.RESHAPE, Ops.EXPAND), src=(UPat(), UPat(dtype=dtypes.ints+(dtypes.weakint, dtypes.void)))), lambda: True), (UPat((Ops.PAD, Ops.SHRINK), src=(UPat(), UPat(), UPat()), name="x"), lambda x: x.src[1].shape == x.src[2].shape), - (UPat((Ops.PERMUTE, Ops.FLIP), name="mv", src=(UPat(),)), lambda mv: isinstance(mv.arg, tuple)), + (UPat(Ops.PERMUTE, name="mv", src=(UPat(),)), lambda mv: isinstance(mv.arg, tuple)), # REDUCE has arg=(op, num_axes), src[1:] are ranges after lowering (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), @@ -179,7 +181,7 @@ spec_tensor = PatternMatcher([ # UNSHARD/MSELECT/MSTACK # an UNSHARD carries the value and one sharding range per sharded axis (usually a DEVICE RANGE, but can be a derived expression) (UPat(Ops.UNSHARD, name="multi"), lambda multi: len(multi.src) == 1+len(multi.arg) - and all(isinstance(a, int) for a in multi.arg) and all(r.dtype in dtypes.weaks for r in multi.src[1:])), + and all(isinstance(a, int) for a in multi.arg) and all(r.dtype == dtypes.weakint for r in multi.src[1:])), (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)), (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)), diff --git a/tinygrad_repo/tinygrad/uop/symbolic.py b/tinygrad_repo/tinygrad/uop/symbolic.py index 9c6d130b3..bd12e42fc 100644 --- a/tinygrad_repo/tinygrad/uop/symbolic.py +++ b/tinygrad_repo/tinygrad/uop/symbolic.py @@ -3,7 +3,7 @@ import math from collections import defaultdict from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu, promo_dtype from tinygrad.dtype import PyConst, ConstFloat, dtypes, can_lossless_cast, Invalid, bitcast, truncate -from tinygrad.helpers import partition, all_same, prod, flatten, unwrap, IMAGE, dedup +from tinygrad.helpers import partition, all_same, prod, flatten, IMAGE, dedup, ceildiv from tinygrad.uop.divandmod import div_and_mod_symbolic from tinygrad.uop.movement import mop_cleanup from tinygrad.uop.weak import pm_uncast_const @@ -204,9 +204,10 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ # ******** phase 2 builds on phase 1, it includes the old "symbolic", rules that match deeper ******** def lt_folding(x:UOp, c:int) -> UOp|None: - p, np = partition(x.split_uop(Ops.ADD), lambda u: u.const_factor() == 1) - if np and (d:=math.gcd(*[u.const_factor() for u in np], c)) > 1 and 0 <= sum(u.vmin for u in p) and sum(u.vmax for u in p) < d: - return unwrap(UOp.usum(*np).divides(d))<(c//d) + p, np = partition([u.pop_const(Ops.MUL) for u in x.split_uop(Ops.ADD)], lambda uf: uf[1] == 1) + # d*y+p < c <=> y < ceildiv(c-p, d), so it folds when that bound is the same for p.vmin and p.vmax + if np and (d:=math.gcd(*[int(f) for _,f in np])) > 1 and (t:=ceildiv(c-sum(u.vmin for u,_ in p), d)) == ceildiv(c-sum(u.vmax for u,_ in p), d): + return UOp.usum(*[u*(f//d) for u,f in np]) UOp|None: @@ -286,9 +287,6 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ # ** lt ** # c0+x x < c1-c0 ((UPat.cvar("c0") + UPat.var("x", dtype=dtypes.ints+(dtypes.weakint,))) < UPat.cvar("c1"), lambda x,c0,c1: x<(c1-c0)), - # c0*x sign(c0)*x < ceil(c1/abs(c0)) - ((UPat.cvar("c0")*UPat.var("x", dtype=dtypes.weakint)) 0 else -x)<-(-c1.val//abs(c0.val)) if abs(c0.val) > 1 else None), # x//d x0, and -> c*d 0 else (x>c.val*d.val) if d.val < 0 else None), @@ -297,7 +295,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ ((UPat.var("x") * UPat.cvar("c1")) * UPat.var("y"), lambda x,c1,y: (x*y)*c1 if y.op is not Ops.CONST else None), # *** rules from symbolic *** # generic lt folding - (UPat.var("x", dtypes.weakint) 0. NOTE: not x < 1 means x > 0 ((UPat.var("x", dtypes.weakint)<1).ne(True), lambda x: (newx<1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None), @@ -332,15 +330,11 @@ def parse_valid(v:UOp) -> tuple[UOp, bool, int]|None: # if it's X <= c, returns X, True, c # if it's X >= c, returns X, False, c - if v.op is Ops.CMPNE and v.src[1].op is Ops.CONST and v.src[1].val == 1 and (s0:=v.src[0]).op is Ops.CMPLT and dtypes.is_int(s0.src[0].dtype): - # (X < c).ne(True) -> X >= c - return s0.src[0], False, int(s0.src[1].vmin) - if v.op is Ops.CMPLT and dtypes.is_int(v.src[0].dtype): - # c < X -> X >= c+1 (a const on the left is a lower bound on the right) - if v.src[0].op is Ops.CONST: return v.src[1], False, int(v.src[0].val)+1 - # X < c -> X <= c-1 - return v.src[0], True, int((v.src[1]).vmax)-1 - return None + # a=a+1 for a const a), its negation a>=b.vmin (b<=a for a const a) + neg = v.op is Ops.CMPNE and v.src[1].op is Ops.CONST and v.src[1].val == 1 + if (lt:=v.src[0] if neg else v).op is not Ops.CMPLT or not dtypes.is_int(lt.src[0].dtype): return None + if (a:=lt.src[0]).op is Ops.CONST: return lt.src[1], neg, int(a.val) + (0 if neg else 1) + return a, not neg, int(lt.src[1].vmin) if neg else int(lt.src[1].vmax)-1 def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp: # return simplified uop (might be the same as input) diff --git a/tinygrad_repo/tinygrad/viz/serve.py b/tinygrad_repo/tinygrad/viz/serve.py index 5223a62c8..e82a2b1d1 100755 --- a/tinygrad_repo/tinygrad/viz/serve.py +++ b/tinygrad_repo/tinygrad/viz/serve.py @@ -74,6 +74,7 @@ def load_rewrites(data:VizData) -> None: for j,s in enumerate(rewrites:=data.trace.rewrites[i]): steps.append(create_step(s.name, ("/graph-rewrites", i, j), loc=s.loc, match_count=len(s.matches), code_line=printable(s.loc), trace=k.tb if j==0 else None, depth=s.depth)) + if s.name == "View Base AST": data.ref_map[canonicalize_ast(_reconstruct(data, s.sink))] = i # get source and binary from Ops.PROGRAM if s.name == "linearize/render": lin_idx = j if lin_idx is not None and (j+1 == len(rewrites) or rewrites[j+1].depth <= rewrites[lin_idx].depth): @@ -82,7 +83,7 @@ def load_rewrites(data:VizData) -> None: steps.append(create_step("View Disassembly", ("/asm", i, len(steps)), (k.ret, lin_idx), depth=0)) lin_idx = None if s.name == "View Program": ki = _reconstruct(data, s.sink, depth=1).src[0].arg - for key in k.keys: data.ref_map[canonicalize_ast(key) if isinstance(key, UOp) else key] = i + for key in k.keys: data.ref_map[key] = i data.ctxs.append({"name":k.display_name, "steps":steps, "ki":ki}) # ** get the complete UOp graphs for one rewrite @@ -102,7 +103,7 @@ def fmt_colored(s:str) -> str: return ansistrip(s) if NO_COLOR else s def canonicalize_ast(u:UOp) -> UOp: return u.replace(arg=KernelInfo()) if u.op is Ops.SINK and isinstance(u.arg, KernelInfo) else u def tokenize_uir(data:VizData, root:UOp) -> list[dict]: - nodes = [u for u in root.toposort() if not _inline(u)] + nodes = [u for u in root.toposort(enter_calls=True) if not _inline(u)] refs = {f"%{i}":{"id":str(id(u))} for i,u in enumerate(nodes)} lines = [[{"st":s, **refs.get(s, {})} for s in re.split(r"( : [^\n]*|%\d+\b)", line) if s] for line in render_uir(root).split("\n")] for u,line in zip(nodes, lines): @@ -114,7 +115,7 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]: assert isinstance(x, UOp) graph: dict[int, dict] = {} excluded: set[UOp] = set() - for u in (toposort:=x.toposort()): + for u in (toposort:=x.toposort(enter_calls=True)): # always exclude CONST if u.op is Ops.CONST and u is not x: excluded.add(u) if u.op is Ops.STACK and len(u.src) == 0: excluded.add(u) @@ -187,7 +188,8 @@ def get_full_rewrite(data:VizData, ctx:TrackedGraphRewrite, depth:int|None=None, match_repr = f"# {dur*1e6:.2f} us\n"+printable(upat_loc) diff = difflib.unified_diff(u0.render_uir().splitlines(), u1.render_uir().splitlines()) yield {"graph":(sink_json:=uop_to_json(data, new_sink)), "uop":tokenize_uir(data, new_sink), "upat":(upat_loc, match_repr), "_sink":new_sink, - "change":[id(x) for x in u1.toposort() if id(x) in sink_json], "diff":[ansistrip(x) for x in diff if not x.startswith(("---","+++","@@"))]} + "change":[id(x) for x in u1.toposort(enter_calls=True) if id(x) in sink_json], + "diff":[ansistrip(x) for x in diff if not x.startswith(("---","+++","@@"))]} if not ctx.bottom_up: next_sink = new_sink def get_sink_at(upats:tuple[str, ...], viz_data:VizData, kernel_idx:int, lin_idx:int, depth:int|None=None, alt:str|None=None) -> UOp|None: @@ -280,7 +282,7 @@ def graph_layout(k:str, dev_events:list[tuple[int, int, float, DevEvent]], start if mem > peak: peak = mem if e.name == "exec" and e.arg["bufs"]: for b in e.arg["bufs"]: buf_ei.setdefault(b, []).append(e) - if e.name == "free": + if e.name == "free" and e.key in temp: events.append(encode_mem_free(e.key, rel_ts(e.ts, start_ts, f"free on {e.device}"), buf_ei.pop(e.key, []), scache)) mem -= temp.pop(e.key) for t in temp: events.append(encode_mem_free(t, rel_ts(end_ts, start_ts, f"end_ts for {k}"), buf_ei.pop(t, []), scache)) @@ -481,7 +483,7 @@ def get_profile(data:VizData, profile:list[ProfileEvent], sort_fn:Callable[[str] for ts,en,e in flatten_events(profile, device_ts_diffs): if isinstance(e, ProfilePointEvent) and e.name == "JSON": ext_data[e.key] = e.arg else: - dev_events.setdefault(e.device,[]).append((st:=int(ts), et:=int(en), float(en-ts), e)) + dev_events.setdefault(e.device,[]).append((st:=int(ts), et:=int(en) if isinstance(e,ProfilePointEvent) else max(st+1,int(en)), float(et-st),e)) if start_ts is None or st < start_ts: start_ts = st if end_ts is None or et > end_ts: end_ts = et if isinstance(e, ProfilePointEvent) and e.name == "marker": markers.append(e) @@ -550,18 +552,18 @@ def amd_decode(buf:bytes, arch:str, off:int=0) -> dict[int, Inst]: from tinygrad.runtime.autogen.amd.rdna3.ins import s_code_end code_end = s_code_end().to_bytes()*5 if arch.startswith("rdna") else None addr_table:dict[int, Inst] = {} - offset = 0 + offset, view = 0, memoryview(buf) while offset < len(buf): - remaining = buf[offset:] - if code_end is not None and remaining.startswith(code_end): break + if code_end is not None and buf.startswith(code_end, offset): break + remaining = view[offset:] fmt = detect_format(remaining, arch) decoded = fmt.from_bytes(remaining) addr_table[off+offset] = decoded offset += decoded.size() return addr_table -def parse_branch(inst) -> int|None: - if "branch" in getattr(inst, "op_name", "").lower(): +def parse_branch(inst, op:str) -> int|None: + if "branch" in op: x = inst.simm16 & 0xffff return (x - 0x10000 if x & 0x8000 else x)*4 return None @@ -580,7 +582,9 @@ def get_cfg(pc_table:dict[int, Inst]) -> dict: # get leaders leaders:set[int] = {next(iter(pc_table))} for pc, inst in pc_table.items(): - if (offset:=parse_branch(inst)) is not None: leaders.update((pc+inst.size()+offset, pc+inst.size())) + op = getattr(inst, "op_name", "").lower() + if (offset:=parse_branch(inst, op)) is not None: leaders.update((pc+inst.size()+offset, pc+inst.size())) + elif op in {"s_endpgm", "s_barrier", "s_barrier_wait"}: leaders.add(pc+inst.size()) # build the cfg curr:int|None = None blocks:dict[int, list[int]] = {} @@ -593,10 +597,10 @@ def get_cfg(pc_table:dict[int, Inst]) -> dict: blocks[curr].append(pc) # otherwise a basic block can have exactly one or two paths nx = pc+inst.size() - if (offset:=parse_branch(inst)) is not None: - if inst.op_name == "S_BRANCH": paths[curr][nx+offset] = UNCOND + if (offset:=parse_branch(inst, op:=getattr(inst, "op_name", "").lower())) is not None: + if op == "s_branch": paths[curr][nx+offset] = UNCOND else: paths[curr].update([(nx+offset, COND_TAKEN), (nx, COND_NOT_TAKEN)]) - elif nx in leaders: paths[curr][nx] = UNCOND + elif nx in leaders and op != "s_endpgm": paths[curr][nx] = UNCOND pc_tokens:dict[int, list[dict]] = {} for pc, inst in pc_table.items(): pc_tokens[pc] = tokens = [] @@ -626,7 +630,7 @@ def get_render(viz_data:VizData, query:str, **kwargs) -> dict: if (sink:=get_sink_at(("do_linearize",), viz_data, i, data, alt="View Program")) is None: return {"src":"No linear found"} if sink.op is Ops.REWRITE_ERROR: return {"src":sink.arg} ret:dict = {} - with soft_err(lambda err: ret.update(err)): ret["src"] = render_uir(list(sink.src[1].toposort())[:-1]) + with soft_err(lambda err: ret.update(err)): ret["src"] = render_uir(list(sink.src[1].toposort(enter_calls=True))[:-1]) return ret if fmt == "code": if (sink:=get_sink_at(("do_render",), viz_data, i, data, depth=1, alt="View Program")) is None: return {"src":"No source found"}