diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 9277b16886..ffb67b7d5e 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -623,6 +623,8 @@ jobs: run: test/external/process_replay/reset.py - name: openpilot compile3 0.11.0 driving_vision run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx + - name: openpilot compile3 0.11.0 driving_vision (from pickle) + run: BENCHMARK_LOG=openpilot_0_11_0_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py - name: IR3 openpilot compile3 0.11.0 driving_vision run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx - name: openpilot compile3 0.11.0 driving_policy @@ -668,6 +670,8 @@ jobs: run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx - name: openpilot load_pickle 0.10.1 driving_vision run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py + - name: openpilot run_pickle 0.10.1 driving_vision + run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py testreddriverbenchmark: name: AM Benchmark diff --git a/.gitignore b/.gitignore index 1332ef84d8..12333d702b 100644 --- a/.gitignore +++ b/.gitignore @@ -68,3 +68,4 @@ mutants .mutmut-cache dagre/ graphlib/ +uv.lock diff --git a/docs/abstractions3.py b/docs/abstractions3.py index 2a316daeaa..e0aef07fbb 100644 --- a/docs/abstractions3.py +++ b/docs/abstractions3.py @@ -1,6 +1,4 @@ # abstractions2 goes from back to front, here we will go from front to back -from typing import List -from tinygrad.helpers import tqdm # ***** # 0. Load mnist on the device @@ -33,21 +31,21 @@ model(X).sparse_categorical_crossentropy(Y).backward() optim.schedule_step() # this will step the optimizer without running realize # ***** -# 3. Create a schedule. +# 3. Create a schedule (linear uop). # The weight Tensors have been assigned to, but not yet realized. Everything is still lazy at this point # l1.uop and l2.uop define a computation graph -from tinygrad.schedule import ExecItem -schedule: List[ExecItem] = Tensor.schedule(l1, l2) +from tinygrad.engine.realize import run_linear +linear = Tensor.schedule_linear(l1, l2) -print(f"The schedule contains {len(schedule)} items.") -for si in schedule: print(str(si)[:80]) +print(f"The schedule contains {len(linear.src)} items.") +for call in linear.src: print(str(call)[:80]) # ***** -# 4. Lower and run the schedule. +# 4. Lower and run the schedule (linear uop). -for si in tqdm(schedule): si.run() +run_linear(linear) # ***** # 5. Print the weight change diff --git a/docs/developer/developer.md b/docs/developer/developer.md index 2d4eb78241..59d3c8c701 100644 --- a/docs/developer/developer.md +++ b/docs/developer/developer.md @@ -17,13 +17,11 @@ The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not al ## Scheduling -The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/schedule/__init__.py) converts the graph of UOps into a list of `ExecItem`. One `ExecItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on. - -::: tinygrad.schedule.ExecItem +The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/schedule/__init__.py) converts the graph of UOps into a `LINEAR` UOp whose `src` is a list of `CALL` UOps. One `CALL` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. The `CALL`'s `src[0]` (a `SINK` ast) specifies what compute to run, and the remaining `src` are the buffers to run it on. ## Lowering -The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers `ExecItem` by populating its `prg` field with +The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers each `CALL` by compiling its ast into a `PROGRAM` and running it. ::: tinygrad.engine.realize.run_linear @@ -35,13 +33,7 @@ Then we render the UOps into code with a `Renderer`, then we compile the code to ## Execution -Creating `ExecItem`, which has a run method - -::: tinygrad.engine.realize.ExecItem - options: - members: true - -Lists of `ExecItem` can be condensed into a single ExecItem with the Graph API (rename to Queue?) +`run_linear` walks the `LINEAR` UOp, dispatching each `CALL` to a runner (kernel, copy, view, encdec, or graph). ## Runtime diff --git a/docs/developer/layout.md b/docs/developer/layout.md index a74d239fe7..9b95ed54c9 100644 --- a/docs/developer/layout.md +++ b/docs/developer/layout.md @@ -28,7 +28,7 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi Transform the optimized ast into a linearized and rendered program. -::: tinygrad.codegen.get_program +::: tinygrad.codegen.to_program options: members: false show_labels: false @@ -53,7 +53,7 @@ Transform the linearized list of UOps into a program, represented as a string. Abstracted high level interface to the runtimes. -::: tinygrad.engine.realize.get_program +::: tinygrad.engine.realize.to_program options: members: false show_labels: false diff --git a/docs/tensor/properties.md b/docs/tensor/properties.md index a2c0a1ba23..7fb2034d00 100644 --- a/docs/tensor/properties.md +++ b/docs/tensor/properties.md @@ -19,8 +19,8 @@ ## tinygrad ops -::: tinygrad.Tensor.schedule_with_vars -::: tinygrad.Tensor.schedule +::: tinygrad.Tensor.linear_with_vars +::: tinygrad.Tensor.schedule_linear ::: tinygrad.Tensor.realize ::: tinygrad.Tensor.replace ::: tinygrad.Tensor.assign diff --git a/examples/anthropic_challenge.py b/examples/anthropic_challenge.py index 0539452364..0dafd21992 100644 --- a/examples/anthropic_challenge.py +++ b/examples/anthropic_challenge.py @@ -173,16 +173,16 @@ if __name__ == "__main__": # *** render to device *** - from tinygrad.codegen import get_program + from tinygrad.codegen import to_program with Context(PCONTIG=2, DEVECTORIZE=2, SPEC=0): out = tree_traversal(forest_t, val_t, height, rounds) - sink = out.schedule()[-1].ast - prg = get_program(sink, VLIWRenderer()) + sink = out.schedule_linear().src[-1].src[0] + prg = to_program(sink, VLIWRenderer()) # *** run on Machine and compare *** # NOTE: the scratch size needs to be reduced to 1536 when you have a register allocator - src = eval(prg.src) + src = eval(prg.src[3].arg) max_regs = max(t[1] for instr in src for v in instr.values() for t in v if len(t) > 1) + 8 print(f"{max_regs:5d} regs used" + ("" if max_regs <= 1536 else " <-- WARNING: TOO MANY REGISTERS, MUST BE <= 1536")) machine = problem.Machine(mem, src, problem.DebugInfo(scratch_map={}), n_cores=1, trace=False, scratch_size=max_regs) diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index 62084a5533..43fc11aae3 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -1282,7 +1282,7 @@ def train_bert(): previous_step = i def train_llama3(): - from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8, FP8_DTYPE + from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE from examples.llama3 import MODEL_PARAMS from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup from examples.mlperf.optim import GradAccClipAdamW @@ -1395,7 +1395,7 @@ def train_llama3(): params = get_parameters(model) - if getenv("FAKEDATA"): + if getenv("EMPTYWEIGHT"): for v in get_parameters(model): v = v.assign(Tensor.empty(v.shape, dtype=v.dtype)) @@ -1432,18 +1432,17 @@ def train_llama3(): print(f"loading optim checkpoint from {fn}") load_state_dict(scheduler, safe_load(fn), realize=False) - fp8_amax = [t for ts in model._fp8_amax.values() for t in ts] if FP8 else [] - fp8_inv_scales = list(model._fp8_inv_scale.values()) if FP8 else [] + fp8_amax = [t for ts in model._fp8_amax.values() for t in ts] + fp8_inv_scales = list(model._fp8_inv_scale.values()) - if FP8: - from tinygrad.nn.state import get_state_dict - model_state = get_state_dict(model) - for wname in ["wqkv", "wo", "w13", "w2"]: - w = model_state[wname] - w._inv_scale = model._fp8_inv_scale[wname] - if optim.master_params: - idx = next(j for j, p in enumerate(optim.params) if p is w) - optim.master_params[idx].assign((optim.master_params[idx] * w._inv_scale.reshape(-1, *([1]*(w.ndim-1)))).contiguous()) + from tinygrad.nn.state import get_state_dict + model_state = get_state_dict(model) + for wname in ["wqkv", "wo", "w13", "w2"]: + w = model_state[wname] + w._inv_scale = model._fp8_inv_scale[wname] + if optim.master_params: + idx = next(j for j, p in enumerate(optim.params) if p is w) + optim.master_params[idx].assign((optim.master_params[idx] * w._inv_scale.reshape(-1, *([1]*(w.ndim-1)))).contiguous()) @TinyJit def minibatch(tokens:Tensor): @@ -1451,7 +1450,11 @@ def train_llama3(): if is_mp: tokens = tokens.shard(device) if not is_sharding: tokens = tokens.to(None) logits:Tensor = model(tokens[:, :-1]) - loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:]) + if getenv("FAST_CE", 0): + from extra.llama_kernels.fused_ce import fused_ce_loss + loss = fused_ce_loss(logits.cast(dtypes.bfloat16), tokens[:, 1:], label_smoothing=0.0) + else: + loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:]) for g, new_g in zip(grads, loss.gradient(*optim.params)): apply_grad(g, new_g.uop) @@ -1555,7 +1558,7 @@ def train_llama3(): mem_gb = GlobalCounters.mem_used / 1e9 gflops = GlobalCounters.global_ops / 1e9 / dev_time - mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * (4.6e15 if FP8 else 2.3e15))) * 100 + mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 4.6e15)) * 100 tqdm.write( f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \ f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU") diff --git a/examples/mlperf/models/flat_llama.py b/examples/mlperf/models/flat_llama.py index 08ac53e48e..fe7f55671d 100644 --- a/examples/mlperf/models/flat_llama.py +++ b/examples/mlperf/models/flat_llama.py @@ -1,4 +1,4 @@ -import math, os, functools +import math, os if __name__ == "__main__": os.environ["DEFAULT_FLOAT"] = "bfloat16" os.environ["OPTIM_DTYPE"] = "bfloat16" @@ -16,65 +16,60 @@ from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker from tinygrad.uop.ops import Ops, UOp from extra.models.llama import apply_rotary_emb, precompute_freqs_cis +from extra.llama_kernels.rmsnorm import rmsnorm +from extra.llama_kernels import FP8_MAX, local_abs_max -FP8 = getenv("FP8", 0) +ASM_GEMM = getenv("ASM_GEMM", 0) FP8_DTYPE = dtypes.fp8e4m3 FP8_GRAD_DTYPE = dtypes.fp8e5m2 -FP8_MAX = 448.0 - -# per-device abs max without allreduce (matches TE delayed scaling behavior) -@functools.cache -def _local_abs_max_fxn(x_p, device): - x = Tensor(x_p, device=device) - inner = Tensor(x.uop.src[0]) if x.uop.op is Ops.MULTI else x - return (inner.abs().max(),) - -def _local_abs_max(x:Tensor) -> Tensor: - param = x.as_param(0) - fxn = _local_abs_max_fxn(param.uop, x.device) - return Tensor(fxn[0].uop.call(x.uop).gettuple(0)) def quantize_fp8(x:Tensor, amax_state:Tensor|None=None): - new_amax = (_local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach() + new_amax = (local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach() scale = FP8_MAX / ((amax_state if amax_state is not None else new_amax) + 1e-8) x_scaled = x * scale x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach()) # STE return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax -def matmul(x:Tensor, w:Tensor, fp8=FP8, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None, +def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None, x_fp8:Tensor|None=None, x_scale:Tensor|None=None, x_new_amax:Tensor|None=None) -> tuple[Tensor,...]: if not fp8: - if getenv("ASM_GEMM"): + if ASM_GEMM: from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),) return (x @ w.T,) assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)" if x_fp8 is None: x_fp8, x_scale, x_new_amax = quantize_fp8(x, amax_state=amax_x) - if getenv("ASM_GEMM"): + if ASM_GEMM: from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm if can_use_asm_gemm(x_fp8, w.T): return asm_gemm(x_fp8, w.T, x_scale=x_scale, w_scale=w_inv_scale), x_new_amax, x_fp8, w - return x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale, x_new_amax, x_fp8, w + return (x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale).cast(dtypes.bfloat16), x_new_amax, x_fp8, w -def _rmsnorm_fwd(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]: - x = x_in.float() - rrms = (x.square().mean(-1, keepdim=True) + eps).rsqrt() - return (x * rrms).cast(x_in.dtype), rrms +def norm_mul_quantize_matmul(x:Tensor, norm:Tensor, amax_x, w_inv_scale, w:Tensor, eps:float): + FUSED_NORM_MUL_QUANTIZE = getenv("FUSED_NORM_MUL_QUANTIZE", 0) + normed, rrms = rmsnorm(x, eps) + if FUSED_NORM_MUL_QUANTIZE: + from extra.llama_kernels.fused_mul_quantize_fp8 import fused_mul_quantize_fp8 + amax_s = amax_x if amax_x is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=normed.device) + x_fp8, x_inv_scale, new_amax = fused_mul_quantize_fp8(normed, norm, amax_s, FP8_DTYPE) + out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, x_scale=x_inv_scale, x_new_amax=new_amax) + else: + x = normed * norm + out, *ret = matmul(x, w, amax_x=amax_x, w_inv_scale=w_inv_scale) + return out, normed, rrms, ret -@functools.cache -def _rmsnorm_fwd_fxn(x_in_p, eps, device): - return _rmsnorm_fwd(Tensor(x_in_p, device=device), eps) - -def _rmsnorm_bwd(grad:UOp, call:UOp) -> tuple: - x_normed = Tensor(call.gettuple(0)).float() - do_float = Tensor(grad).float() - d_x = Tensor(call.gettuple(1)) * (do_float - x_normed * (do_float * x_normed).mean(-1, keepdim=True)) - return (d_x.cast(call.src[1].dtype).uop,) - -def rmsnorm(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]: - fxn = _rmsnorm_fwd_fxn(x_in.as_param(0).uop, eps, x_in.device) - call = UOp.maketuple(fxn[0].uop, fxn[1].uop).call(x_in.uop, grad_fxn=_rmsnorm_bwd) - return Tensor(call.gettuple(0)), Tensor(call.gettuple(1)) +def silu_w13_matmul(x_w13:Tensor, w2:Tensor, amax_x2, s_2): + FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0) + if FUSED_SILU_W13: + from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13 + amax_s = amax_x2 if amax_x2 is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=x_w13.device) + x2_fp8, x2_inv_scale, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_s, FP8_DTYPE) + out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, x_scale=x2_inv_scale, x_new_amax=new_amax_x2) + else: + hidden_dim = x_w13.shape[-1] // 2 + x_w1, x_w3 = x_w13[..., :hidden_dim], x_w13[..., hidden_dim:] + out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2) + return out, ret class FlatTransformer: def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None, @@ -90,7 +85,7 @@ class FlatTransformer: scaled_std = 0.02 / math.sqrt(2 * n_layers) # Attention - self._init_inv_scales = [] # populated by lin_per_layer when FP8 + self._init_inv_scales = [] # populated by lin_per_layer self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2) self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std) @@ -109,21 +104,19 @@ class FlatTransformer: self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16) self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False) - if FP8: - def _amax(): return Tensor.full((), FP8_MAX).contiguous().requires_grad_(False) - names = ["xqkv", "xo", "x13", "x2"] - self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names} - # per-weight inv_scale: single (n_layers,) float32 tensor per weight (kernel reads float* pointers) - w_names = ["wqkv", "wo", "w13", "w2"] - self._fp8_inv_scale = {} - for wname, inv_scales in zip(w_names, self._init_inv_scales): - self._fp8_inv_scale[wname] = inv_scales.float().contiguous().requires_grad_(False) - del self._init_inv_scales + def _amax(): return Tensor.full((), FP8_MAX).contiguous().requires_grad_(False) + names = ["xqkv", "xo", "x13", "x2"] + self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names} + # per-weight inv_scale: single (n_layers,) float32 tensor per weight (kernel reads float* pointers) + w_names = ["wqkv", "wo", "w13", "w2"] + self._fp8_inv_scale = {} + for wname, inv_scales in zip(w_names, self._init_inv_scales): + self._fp8_inv_scale[wname] = inv_scales.float().contiguous().requires_grad_(False) + del self._init_inv_scales def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02): - if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features) + if getenv("ZEROS", 0): w = Tensor.zeros(self.n_layers, out_features, in_features) else: w = Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std) - if not FP8: return w # per-layer scaled fp8 cast: fill the fp8 range for best precision amax = w.abs().flatten(1).max(1).detach() scale = FP8_MAX / (amax + 1e-8) @@ -135,18 +128,8 @@ class FlatTransformer: bsz, seqlen, _ = x.shape new_amaxs, saves = [], [] - x, rrms = rmsnorm(x, self.norm_eps) - saves.extend([x, rrms]) - - if FP8 and getenv("FUSED_NORM_MUL_QUANTIZE", 1): - from extra.amax.cast_amax import fused_mul_quantize_fp8 - amax_s = amax_xqkv if amax_xqkv is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=x.device) - x_fp8, x_inv_scale, new_amax_xqkv = fused_mul_quantize_fp8(x, attention_norm, amax_s, FP8_DTYPE) - xqkv, *ret = matmul(None, wqkv, w_inv_scale=s_qkv, x_fp8=x_fp8, x_scale=x_inv_scale, x_new_amax=new_amax_xqkv) - else: - x = x * attention_norm - xqkv, *ret = matmul(x, wqkv, amax_x=amax_xqkv, w_inv_scale=s_qkv) - + xqkv, normed, rrms, ret = norm_mul_quantize_matmul(x, attention_norm, amax_xqkv, s_qkv, wqkv, self.norm_eps) + saves.extend([normed, rrms]) new_amaxs.extend(ret[:1]) saves.extend(ret[1:] + [xqkv]) xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim) @@ -155,7 +138,7 @@ class FlatTransformer: xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim) xq, xk = apply_rotary_emb(xq, xk, freqs_cis) - if FP8: xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16) + xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16) xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2) if getenv("HK_FLASH_ATTENTION"): from extra.thunder.amd.fa import flash_attention @@ -174,28 +157,12 @@ class FlatTransformer: amax_x13=None, amax_x2=None, s_13=None, s_2=None): new_amaxs, saves = [], [] - x, rrms = rmsnorm(x, self.norm_eps) - saves.extend([x, rrms]) - - if FP8 and getenv("FUSED_NORM_MUL_QUANTIZE", 1): - from extra.amax.cast_amax import fused_mul_quantize_fp8 - amax_s13 = amax_x13 if amax_x13 is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=x.device) - x_fp8_13, x_inv_scale_13, new_amax_x13 = fused_mul_quantize_fp8(x, ffn_norm, amax_s13, FP8_DTYPE) - x_w13, *ret = matmul(None, w13, w_inv_scale=s_13, x_fp8=x_fp8_13, x_scale=x_inv_scale_13, x_new_amax=new_amax_x13) - else: - x = x * ffn_norm - x_w13, *ret = matmul(x, w13, amax_x=amax_x13, w_inv_scale=s_13) + x_w13, normed, rrms, ret = norm_mul_quantize_matmul(x, ffn_norm, amax_x13, s_13, w13, self.norm_eps) + saves.extend([normed, rrms]) new_amaxs.extend(ret[:1]) saves.extend(ret[1:] + [x_w13]) - if FP8 and getenv("FUSED_SILU_W13", 1): - from extra.amax.cast_amax import fused_quantize_fp8_w13 - amax_s = amax_x2 if amax_x2 is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=x_w13.device) - x2_fp8, x2_inv_scale, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_s, FP8_DTYPE) - out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, x_scale=x2_inv_scale, x_new_amax=new_amax_x2) - else: - x_w1, x_w3 = x_w13[..., :self.hidden_dim], x_w13[..., self.hidden_dim:] - out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2) + out, ret = silu_w13_matmul(x_w13, w2, amax_x2, s_2) new_amaxs.extend(ret[:1]) saves.extend(ret[1:] + [out]) return (out, *new_amaxs, *saves) @@ -226,41 +193,35 @@ class FlatTransformer: else: # flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer self.wqkv.shard_(device, axis=1).realize() # (n_layers, out, dim) shard out - self.wo.shard_(device, axis=2).realize() # (n_layers, dim, in) shard in - self.w13.shard_(device, axis=1).realize() # (n_layers, hidden*2, dim) shard out - self.w2.shard_(device, axis=2).realize() # (n_layers, dim, hidden) shard in + self.wo.shard_(device, axis=2).realize() # (n_layers, dim, in) shard in + self.w13.shard_(device, axis=1).realize() # (n_layers, hidden*2, dim) shard out + self.w2.shard_(device, axis=2).realize() # (n_layers, dim, hidden) shard in self.attention_norm.shard_(device, axis=None).realize() self.ffn_norm.shard_(device, axis=None).realize() self.norm.weight.shard_(device, axis=None).realize() self.tok_embeddings.weight.shard_(device, axis=0).realize() self.output.shard_(device, axis=1).realize() self.freqs_cis.shard_(device, axis=None).realize() - if FP8: - for name in self._fp8_amax: - for i in range(len(self._fp8_amax[name])): - self._fp8_amax[name][i] = self._fp8_amax[name][i].to(device).contiguous().requires_grad_(False) - for name in self._fp8_inv_scale: - self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().requires_grad_(False) + for name in self._fp8_amax: + for i in range(len(self._fp8_amax[name])): + self._fp8_amax[name][i] = self._fp8_amax[name][i].to(device).contiguous().requires_grad_(False) + for name in self._fp8_inv_scale: + self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().requires_grad_(False) def __call__(self, tokens:Tensor): h = self.tok_embeddings(tokens) freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :] - a = self._fp8_amax if FP8 else None - s = self._fp8_inv_scale if FP8 else None + amaxs, inv_scales = self._fp8_amax, self._fp8_inv_scale for i in range(self.n_layers): - amax_layer = {"amax_xqkv": a["xqkv"][i], "amax_xo": a["xo"][i], - "amax_x13": a["x13"][i], "amax_x2": a["x2"][i]} if a else {} - scale_layer = {"s_qkv": s["wqkv"][i], "s_o": s["wo"][i], - "s_13": s["w13"][i], "s_2": s["w2"][i]} if s else {} h, *ret = self.run_layer(h, freqs_cis, self.attention_norm[i], self.wqkv[i], self.wo[i], self.ffn_norm[i], self.w13[i], self.w2[i], - **amax_layer, **scale_layer) - if a: - amaxs = ret[:5] - amax_names = ["xqkv", "xo", "x13", "x2"] - for name, new_val in zip(amax_names, amaxs): - a[name][i].assign(new_val) + amax_xqkv=amaxs["xqkv"][i], amax_xo=amaxs["xo"][i], + amax_x13=amaxs["x13"][i], amax_x2=amaxs["x2"][i], + s_qkv=inv_scales["wqkv"][i], s_o=inv_scales["wo"][i], + s_13=inv_scales["w13"][i], s_2=inv_scales["w2"][i]) + for name, new_val in zip(["xqkv", "xo", "x13", "x2"], ret[:5]): + amaxs[name][i].assign(new_val) logits = matmul(self.norm(h).contiguous().contiguous_backward(), self.output[0], fp8=False)[0].contiguous_backward() return logits diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh index 6e69feb640..22f15f5f87 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh @@ -15,9 +15,10 @@ export WQKV=${WQKV:-1} export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1} export FP8=${FP8:-1} export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1} +export FAST_CE=${FASE_CE:-1} export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16" -export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-16} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2} +export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2} export GBS=$((BS * GRADIENT_ACC_STEPS)) export MODEL="llama3" @@ -36,7 +37,7 @@ export DATA_SEED=${DATA_SEED:-5760} export JITBEAM=${JITBEAM:-3} export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1 -export FAKEDATA=1 BENCHMARK=${BENCHMARK:-10} +export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10} if [ -z "$FULL_LAYERS" ]; then export LLAMA_LAYERS=2 fi diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh index b958a99aef..cff289dde5 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh @@ -15,9 +15,10 @@ export WQKV=${WQKV:-1} export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1} export FP8=${FP8:-1} export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1} +export FAST_CE=${FASE_CE:-1} export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16" -export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-16} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2} +export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2} export GBS=$((BS * GRADIENT_ACC_STEPS)) export MODEL="llama3" diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh index e55dc6f90a..de9f641120 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh @@ -3,4 +3,4 @@ export BENCHMARK=5 export EVAL_BS=0 VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL" -python -m tinygrad.viz.cli --profile -s "$SRC" --top 20 +python -m tinygrad.viz.cli -s "$SRC" --top 20 diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/run_and_time.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/run_and_time.sh index 90bac94511..0851d92928 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/run_and_time.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/run_and_time.sh @@ -16,9 +16,10 @@ export WQKV=1 export MASTER_WEIGHTS=1 export FP8=1 export ALLREDUCE_CAST=1 +export FAST_CE=1 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16" -export DP=8 MP=1 BS=16 EVAL_BS=16 GRADIENT_ACC_STEPS=2 +export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2 export GBS=$((BS * GRADIENT_ACC_STEPS)) export MODEL="llama3" diff --git a/examples/openpilot/compile3.py b/examples/openpilot/compile3.py index ffed2ae32a..c8a4502a8f 100644 --- a/examples/openpilot/compile3.py +++ b/examples/openpilot/compile3.py @@ -5,7 +5,6 @@ if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0" from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes from tinygrad.helpers import DEBUG, getenv from tinygrad.uop.ops import Ops -from tinygrad.engine.realize import get_runner from tinygrad.nn.onnx import OnnxRunner OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx" @@ -37,7 +36,7 @@ def compile(onnx_file): # copy i == 1 so use of JITBEAM is okay if i == 1: test_val = np.copy(ret) # iterate kernel CALLs in the captured LINEAR UOp; toposort descends into batched graph CUSTOM_FUNCTIONs - kernel_asts = {Ops.SINK, Ops.PROGRAM} + kernel_asts = {Ops.PROGRAM} kernel_calls = [u for u in run_onnx_jit.captured.linear.toposort(gate=lambda x: x.op not in kernel_asts) if u.op is Ops.CALL and u.src[0].op in kernel_asts] print(f"captured {len(kernel_calls)} kernels") @@ -49,8 +48,8 @@ def compile(onnx_file): read_image_count = 0 gated_read_image_count = 0 for call in kernel_calls: - device = next(b.device for b in call.src[1:] if b.op is not Ops.BIND) - src = get_runner(device, call.src[0]).p.src + _, _, _, source, _ = call.src[0].src + src = source.arg kernel_count += 1 read_image_count += src.count("read_image") gated_read_image_count += src.count("?read_image") @@ -134,14 +133,20 @@ def bench(run, inputs): run(**inputs).numpy() if __name__ == "__main__": - onnx_file = fetch(OPENPILOT_MODEL) - inputs, outputs = compile(onnx_file) + if getenv("RUN_PICKLE"): + with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f) + inputs = {name: Tensor(Tensor.randn(*[int(s) for s in view.src[1].arg], dtype=dtype).numpy(), device=device) + for name, (view, _vars, dtype, device) in zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_input_info)} + test_vs_compile(pickle_loaded, inputs) + else: + onnx_file = fetch(OPENPILOT_MODEL) + inputs, outputs = compile(onnx_file) - with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f) + with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f) - test_vs_compile(pickle_loaded, inputs, outputs) - if getenv("SELFTEST"): - test_vs_onnx(inputs, outputs, onnx_file, 1e-4) + test_vs_compile(pickle_loaded, inputs, outputs) + if getenv("SELFTEST"): + test_vs_onnx(inputs, outputs, onnx_file, 1e-4) if getenv("BENCHMARK_LOG", ""): bench(pickle_loaded, inputs) diff --git a/examples/webgpu/stable_diffusion/compile.py b/examples/webgpu/stable_diffusion/compile.py index fd926a988f..cfa2689705 100644 --- a/examples/webgpu/stable_diffusion/compile.py +++ b/examples/webgpu/stable_diffusion/compile.py @@ -114,7 +114,7 @@ if __name__ == "__main__": linear, output_bufs = jit_model(step, *step.input) functions, statements, bufs, _ = compile_net(linear, output_bufs) state = get_state_dict(model) - weights = {id(x.uop.base.realized): name for name, x in state.items()} + weights = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None} kernel_code = '\n\n'.join([f"const {key} = `{fixup_code(code, key)}`;" for key, code in functions.items()]) kernel_names = ', '.join([name for (name, _, _, _) in statements]) input_names = [f"input{i}" for i in range(len(step.input))] diff --git a/extra/amax/cast_amax.py b/extra/amax/cast_amax.py deleted file mode 100644 index b0c987cfec..0000000000 --- a/extra/amax/cast_amax.py +++ /dev/null @@ -1,133 +0,0 @@ -import functools, pathlib -from tinygrad import Tensor, dtypes -from tinygrad.uop.ops import UOp, Ops, KernelInfo -from tinygrad.renderer import Estimates -from tinygrad.runtime.support.compiler_amd import HIPCCCompiler - -FP8_MAX = 448.0 -NUM_WG, THREADS_PER_WG = 1024, 256 - -def _compile(cpp_name:str, n_elems:int, hidden:int): - src = (pathlib.Path(__file__).parent/cpp_name).read_text() - defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={hidden}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"] - return src, HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src) - -def _shard_shape(shape:tuple, axis:int, ndev:int) -> list: - s = list(shape); s[axis] //= ndev; return s - -def _scalar_amax(amax_buf:Tensor) -> Tensor: - if isinstance(amax_buf.device, tuple): - from examples.mlperf.models.flat_llama import _local_abs_max - return _local_abs_max(amax_buf).detach() - return amax_buf.max().detach() - -# ** fused silu*mul -> fp8 cast + amax (w13 layout) - -@functools.cache -def _custom_fused_bwd_w13(grad_xw13:UOp, xw13:UOp, grad_x2:UOp, amax_state:UOp, dname:str) -> UOp: - hidden = xw13.shape[2] // 2 - n_elems = xw13.shape[0] * xw13.shape[1] * hidden - threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0") - mem = n_elems * 2 * 5 - sink = UOp.sink(grad_xw13.base, xw13.base, grad_x2.base, amax_state.base, threads, workgroups, - arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=8*n_elems, mem=mem))) - src, lib = _compile("cast_amax_bwd_w13.cpp", n_elems, hidden) - return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), - UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib))) - -@functools.cache -def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, dname:str) -> UOp: - hidden = xw13.shape[2] // 2 - n_elems = xw13.shape[0] * xw13.shape[1] * hidden - threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0") - mem = n_elems * 2 * 2 + n_elems + NUM_WG * 2 - sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups, - arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem))) - src, lib = _compile("cast_amax_fwd_w13.cpp", n_elems, hidden) - return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), - UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib))) - -def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp): - _, _, xw13, amax_state = kernel.src[1:] - device = xw13.device - if isinstance(device, tuple): - axis, ndev = xw13.axis, len(device) - assert axis in (0, 1), f"unsupported sharding axis={axis}" - grad_xw13 = Tensor(Tensor.invalids(*_shard_shape(xw13.shape, axis, ndev), dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device) - dname = device[0].split(":")[0] - else: - grad_xw13 = Tensor.invalids(*xw13.shape, dtype=dtypes.bfloat16, device=device) - dname = device.split(":")[0] if isinstance(device, str) else device - grad_x2_t = Tensor(gradient, device=device).cast(dtypes.bfloat16) - fxn = functools.partial(_custom_fused_bwd_w13, dname=dname) - grad_xw13, *_ = Tensor.custom_kernel(grad_xw13, Tensor(xw13, device=device), grad_x2_t, Tensor(amax_state, device=device), fxn=fxn) - return (None, None, grad_xw13.uop, None) - -def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[Tensor, Tensor, Tensor]: - # silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, inv_scale, new_amax). - assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}" - MBS, SEQ, H2 = xw13.shape - assert H2 % 2 == 0, f"w13 last-axis must be even, got {H2}" - HIDDEN = H2 // 2 - if isinstance(xw13.device, tuple): - axis, ndev = xw13.uop.axis, len(xw13.device) - assert axis in (0, 1), f"unsupported sharding axis={axis}" - fp8_out = Tensor(Tensor.invalids(*_shard_shape((MBS, SEQ, HIDDEN), axis, ndev), dtype=fp8_dtype, device=xw13.device).uop.multi(axis), device=xw13.device) - amax_buf = Tensor(Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device).uop.multi(0), device=xw13.device) - dname = xw13.device[0].split(":")[0] - else: - fp8_out = Tensor.invalids(MBS, SEQ, HIDDEN, dtype=fp8_dtype, device=xw13.device) - amax_buf = Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device) - dname = xw13.device.split(":")[0] if isinstance(xw13.device, str) else xw13.device - fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname) - fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, fxn=fxn, grad_fxn=_fused_quantize_bwd_w13) - inv_scale = (amax_state.float() + 1e-8) / FP8_MAX - return fp8_out, inv_scale, _scalar_amax(amax_buf) - -# ** fused (x * weight) -> fp8 cast + amax (norm-mul-quantize) - -@functools.cache -def _custom_mul_quantize_fp8(fp8_out:UOp, amax_buf:UOp, x:UOp, weight:UOp, amax_state:UOp, dname:str) -> UOp: - MBS, SEQ, HIDDEN = x.shape - n_elems = MBS * SEQ * HIDDEN - threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0") - mem = n_elems * 2 + HIDDEN * 2 + n_elems + NUM_WG * 2 - sink = UOp.sink(fp8_out.base, amax_buf.base, x.base, weight.base, amax_state.base, threads, workgroups, - arg=KernelInfo(f"fused_mul_quantize_fp8_{n_elems}_h{HIDDEN}", estimates=Estimates(ops=3*n_elems, mem=mem))) - src, lib = _compile("fused_mul_quantize_fp8.cpp", n_elems, HIDDEN) - return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), - UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib))) - -def _fused_mul_quantize_fp8_bwd(gradient:UOp, kernel:UOp): - # inputs: (fp8_out, amax_buf, x, weight, amax_state); grads for x and weight - _, _, x_u, weight_u, amax_state_u = kernel.src[1:] - device = x_u.device - grad_t = Tensor(gradient, device=device).cast(dtypes.bfloat16) - x_t, weight_t = Tensor(x_u, device=device), Tensor(weight_u, device=device) - scale = FP8_MAX / (Tensor(amax_state_u, device=device).float() + 1e-8) - grad_scaled = grad_t.float() * scale - # grad_x stays bf16 to avoid CSE materializing a (MBS, SEQ, HIDDEN) fp32 intermediate - grad_x = (grad_scaled * weight_t.float()).cast(dtypes.bfloat16) - grad_weight = (grad_scaled * x_t.float()).sum(axis=(0, 1)).cast(dtypes.bfloat16) - return (None, None, grad_x.uop, grad_weight.uop, None) - -def fused_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[Tensor, Tensor, Tensor]: - # (x * weight) -> fp8 + amax, delayed scaling. Returns (fp8, inv_scale, new_amax). - assert x.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16 - assert x.shape[-1] == weight.shape[-1], f"HIDDEN mismatch: x={x.shape}, weight={weight.shape}" - MBS, SEQ, HIDDEN = x.shape - if isinstance(x.device, tuple): - axis, ndev = x.uop.axis, len(x.device) - assert axis in (0, 1), f"unsupported sharding axis={axis}" - fp8_out = Tensor(Tensor.invalids(*_shard_shape((MBS, SEQ, HIDDEN), axis, ndev), dtype=fp8_dtype, device=x.device).uop.multi(axis), device=x.device) - amax_buf = Tensor(Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=x.device).uop.multi(0), device=x.device) - dname = x.device[0].split(":")[0] - else: - fp8_out = Tensor.invalids(MBS, SEQ, HIDDEN, dtype=fp8_dtype, device=x.device) - amax_buf = Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=x.device) - dname = x.device.split(":")[0] if isinstance(x.device, str) else x.device - fxn = functools.partial(_custom_mul_quantize_fp8, dname=dname) - fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_mul_quantize_fp8_bwd) - new_amax = _scalar_amax(amax_buf) - inv_scale = (amax_state.float() + 1e-8) / FP8_MAX - return fp8_out, inv_scale, new_amax diff --git a/extra/export_model.py b/extra/export_model.py index a9a56a48a9..4ff53218a6 100644 --- a/extra/export_model.py +++ b/extra/export_model.py @@ -6,7 +6,7 @@ from tinygrad.engine.jit import TinyJit from tinygrad.nn.state import get_state_dict from tinygrad.helpers import Context, to_mv, prod from tinygrad.uop.ops import Ops, UOp -from tinygrad.codegen import get_program +from tinygrad.codegen import to_program import json from collections import OrderedDict @@ -36,10 +36,11 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li for call in iter_kernel_calls(linear): arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND] - prg = get_program(call.src[0], Device[arg_uops[0].device].renderer) - functions[prg.function_name] = prg.src - cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + [v for v in prg.vars if v.op is Ops.DEFINE_VAR] - statements.append((prg.function_name, cargs, prg.global_size, prg.local_size)) + prg = to_program(call.src[0], Device[arg_uops[0].device].renderer) + info = prg.arg + functions[info.function_name] = prg.src[3].arg + cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + [v for v in info.vars if v.op is Ops.DEFINE_VAR] + statements.append((info.function_name, cargs, info.global_size, info.local_size)) return functions, statements, {name:(size, dtype, key) for name, size, dtype, key in bufs.values()}, bufs_to_save @@ -244,7 +245,7 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model" with Context(JIT=2, CPU_COUNT=1): linear, output_bufs = jit_model(model, *inputs) functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs) state = get_state_dict(model) - weight_names = {id(x.uop.base.realized): name for name, x in state.items()} + weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None} input_names = [f"input{i}" for i in range(len(inputs))] output_names = [f"output{i}" for i in range(len(output_bufs))] diff --git a/extra/gemm/amd_asm_matmul.py b/extra/gemm/amd_asm_matmul.py index 52b74b582d..337908909b 100644 --- a/extra/gemm/amd_asm_matmul.py +++ b/extra/gemm/amd_asm_matmul.py @@ -13,7 +13,7 @@ from tinygrad import Tensor, Device, Context, GlobalCounters from tinygrad.uop.ops import UOp, Ops, KernelInfo from tinygrad.helpers import getenv, colored from tinygrad.dtype import dtypes, AddrSpace -from tinygrad.engine.realize import Estimates +from tinygrad.engine.realize import Estimates, run_linear from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL from tinygrad.runtime.autogen.amd.rdna3.ins import * @@ -463,11 +463,14 @@ def test_matmul(): estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3))) return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts])))) c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2] - ei = c.schedule()[0].lower() + linear = c.schedule_linear() ets = [] with Context(DEBUG=2): - for _ in range(getenv("CNT", 5)): ets.append(ei.run(wait=True)) + for _ in range(getenv("CNT", 5)): + start = GlobalCounters.time_sum_s + run_linear(linear) + ets.append(GlobalCounters.time_sum_s - start) print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}") if getenv("VERIFY", 1): diff --git a/extra/gemm/amd_matmul.py b/extra/gemm/amd_matmul.py index d72034fef5..4def9c2493 100644 --- a/extra/gemm/amd_matmul.py +++ b/extra/gemm/amd_matmul.py @@ -1,31 +1,39 @@ # kernel8_batched_gmem.s from https://seb-v.github.io/optimization/update/2025/01/20/Fast-GPU-Matrix-multiplication.html # sudo PATH=/opt/homebrew/Cellar/llvm/20.1.6/bin:$PATH AMD_LLVM=0 AMD=1 DEBUG=2 python3 extra/gemm/amd_matmul.py import pathlib -from dataclasses import replace from tinygrad import Tensor, Device, Context, GlobalCounters from tinygrad.helpers import getenv -from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program +from tinygrad.uop.ops import UOp, Ops, KernelInfo +from tinygrad.renderer import Estimates +from tinygrad.engine.realize import run_linear N = 4096 run_count = 5 -if __name__ == "__main__": - ast = (Tensor.empty(N, N)@Tensor.empty(N, N)).schedule()[-1].ast - prg = get_program(ast, Device.default.renderer) +def make_matmul_kernel(name:str, src:str, local_size:int): + def fxn(a:UOp, b:UOp, c:UOp) -> UOp: + threads = UOp.special(local_size, "lidx0") + wg_x = UOp.special(N//128, "gidx0") + wg_y = UOp.special(N//128, "gidx1") + sink = UOp.sink(a.base, b.base, c.base, threads, wg_x, wg_y, arg=KernelInfo(name, estimates=Estimates(ops=2*N**3, mem=3*N*N*4))) + lib = Device[Device.DEFAULT].compiler.compile_cached(src) + return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)), + UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib))) + return fxn +if __name__ == "__main__": if getenv("ASM") == 1: src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel8_batched_gmem.s").read_text() - prgfast = replace(prg, name="kernel", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1]) + name, local_size = "kernel", 128 elif getenv("ASM") == -1: src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel3_registers.cpp").read_text() - prgfast = replace(prg, name="kernel3_registers", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1]) + name, local_size = "kernel3_registers", 256 elif getenv("ASM") == -2: src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel4_gmem_df.cpp").read_text() - prgfast = replace(prg, name="kernel4_gmem_db", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1]) + name, local_size = "kernel4_gmem_db", 256 else: src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel5_lds_optim.cpp").read_text() - prgfast = replace(prg, name="kernel5_lds_optim", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1]) - runner = CompiledRunner(prgfast) + name, local_size = "kernel5_lds_optim", 128 a = Tensor.randn(N, N).realize() b = Tensor.randn(N, N).realize() @@ -35,8 +43,8 @@ if __name__ == "__main__": with Context(DEBUG=2): for _ in range(run_count): tc = (a@b).realize() + linear = Tensor.custom_kernel(a, b, c, fxn=make_matmul_kernel(name, src, local_size))[2].schedule_linear() GlobalCounters.reset() - ei = ExecItem(ast, [a.uop.buffer, b.uop.buffer, c.uop.buffer], prg=runner) with Context(DEBUG=2): - for _ in range(run_count): ei.run(wait=True) + for _ in range(run_count): run_linear(linear) print(f"custom {(c-tc).square().mean().item()}") diff --git a/extra/gemm/max_matmul.py b/extra/gemm/max_matmul.py index 0d1bb9e7c5..5a41fe17f8 100644 --- a/extra/gemm/max_matmul.py +++ b/extra/gemm/max_matmul.py @@ -1,7 +1,6 @@ import numpy as np, os from tinygrad.helpers import getenv, flat_mv from tinygrad import dtypes -from tinygrad.engine.realize import get_program # for copied uops from tinygrad import dtypes diff --git a/extra/gemm/rdna4_asm_matmul.py b/extra/gemm/rdna4_asm_matmul.py index 04ed8ce908..a317020e06 100644 --- a/extra/gemm/rdna4_asm_matmul.py +++ b/extra/gemm/rdna4_asm_matmul.py @@ -4,7 +4,7 @@ from tinygrad import Tensor, Device, Context, GlobalCounters from tinygrad.uop.ops import UOp, Ops, KernelInfo from tinygrad.helpers import getenv, colored from tinygrad.dtype import dtypes, AddrSpace -from tinygrad.engine.realize import Estimates +from tinygrad.engine.realize import Estimates, run_linear from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL, src, ttmp from tinygrad.runtime.autogen.amd.rdna4.ins import * @@ -225,11 +225,14 @@ def test_matmul(): return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts])))) c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2] - ei = c.schedule()[0].lower() + linear = c.schedule_linear() ets = [] with Context(DEBUG=2): - for _ in range(getenv("CNT", 5)): ets.append(ei.run(wait=True)) + for _ in range(getenv("CNT", 5)): + start = GlobalCounters.time_sum_s + run_linear(linear) + ets.append(GlobalCounters.time_sum_s - start) print(f"REAL TFLOPS {N*N*N*2 / min(ets) * 1e-12:.2f}") if getenv("VERIFY", 1): diff --git a/extra/gemm/simple_matmul.py b/extra/gemm/simple_matmul.py index 379b50474e..605ffe161c 100644 --- a/extra/gemm/simple_matmul.py +++ b/extra/gemm/simple_matmul.py @@ -2,6 +2,7 @@ import numpy as np from tinygrad import dtypes, Tensor from tinygrad.helpers import getenv, get_single_element from tinygrad.dtype import _to_np_dtype +from tinygrad.engine.realize import compile_linear from tinygrad.codegen.opt import OptOps dtype_in = (dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else @@ -38,10 +39,10 @@ if __name__ == "__main__": c = a.matmul(b, dtype=acc_dtype).realize() if getenv("SHOULD_USE_TC"): - sched = a.matmul(b, dtype=acc_dtype).schedule() - ei = get_single_element(sched) - ei.lower() - assert any(opt.op is OptOps.TC for opt in ei.prg.p.applied_opts), f"TC not triggered, {ei.prg.p.applied_opts}" + linear = compile_linear(a.matmul(b, dtype=acc_dtype).schedule_linear()) + call = get_single_element(list(linear.src)) + applied_opts = call.src[0].src[0].arg.applied_opts + assert any(opt.op is OptOps.TC for opt in applied_opts), f"TC not triggered, {applied_opts}" ref = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32) res = c.numpy() diff --git a/extra/gemm/tinygrad_nv_matmul.py b/extra/gemm/tinygrad_nv_matmul.py index 1b2e34efa7..adc9a747e8 100644 --- a/extra/gemm/tinygrad_nv_matmul.py +++ b/extra/gemm/tinygrad_nv_matmul.py @@ -1,7 +1,7 @@ -from tinygrad import Tensor, dtypes, Device -from tinygrad.helpers import getenv, DEBUG -from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps -from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program +from tinygrad import Tensor, dtypes, Context +from tinygrad.helpers import getenv +from tinygrad.codegen.opt import Opt, OptOps +from tinygrad.engine.realize import run_linear from dataclasses import replace N = 4096 @@ -11,9 +11,6 @@ if __name__ == "__main__": else: A, B = Tensor.empty(N, N, dtype=dtypes.float16), Tensor.empty(N, N, dtype=dtypes.float16) C = A.matmul(B) - si = C.schedule()[-1] - ast = si.ast - k = Kernel(ast, opts=Device[Device.DEFAULT].renderer) if getenv("GEMV"): opts = [ Opt(op=OptOps.UNROLL, axis=0, amt=8), @@ -28,10 +25,10 @@ if __name__ == "__main__": Opt(op=OptOps.LOCAL, axis=1, amt=2), Opt(op=OptOps.LOCAL, axis=0, amt=2), ] - k.apply_opts(opts) - prg = get_program(k.ast.replace(arg=replace(k.ast.arg, opts_to_apply=tuple(k.applied_opts))), k.opts) - new_src = prg.src - # can mod source here - prg = replace(prg, src=new_src) - ei = ExecItem(si.ast, [x.ensure_allocated() for x in si.bufs], si.metadata, prg=CompiledRunner(prg)) - for i in range(5): ei.run(wait=True) + linear = C.schedule_linear() + call = linear.src[-1] + new_ast = call.src[0].replace(arg=replace(call.src[0].arg, opts_to_apply=tuple(opts))) + new_call = call.replace(src=(new_ast, *call.src[1:])) + linear = linear.replace(src=tuple(new_call if c is call else c for c in linear.src)) + with Context(DEBUG=2): + for i in range(5): run_linear(linear) diff --git a/extra/gemm/triton_nv_matmul.py b/extra/gemm/triton_nv_matmul.py index f6ee932641..60fe9d5c82 100644 --- a/extra/gemm/triton_nv_matmul.py +++ b/extra/gemm/triton_nv_matmul.py @@ -4,7 +4,8 @@ import triton.language as tl from triton.compiler import AttrsDescriptor, ASTSource, compile as triton_compile import numpy as np from tinygrad import Tensor, dtypes, Device -from tinygrad.engine.realize import CompiledRunner, ExecItem, ProgramSpec +from tinygrad.engine.realize import CompiledRunner +from tinygrad.uop.ops import Ops, UOp, KernelInfo, ProgramInfo from tinygrad.helpers import getenv np.set_printoptions(suppress=True) @@ -73,9 +74,11 @@ if __name__ == "__main__": A, B = Tensor.normal(M, K, std=1e-1, dtype=dtypes.float16).realize(), Tensor.normal(K, N, std=1e-1, dtype=dtypes.float16).realize() C = A.matmul(B) - from tinygrad.schedule import linear_to_schedule + from tinygrad.uop.ops import Ops linear, var_vals = C.linear_with_vars() - si = linear_to_schedule(linear)[-1] + last_call = linear.src[-1] + ast = last_call.src[0] + bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND] src = compiled.asm["ptx"] # specify the shared memory here so we don't need to do it dynamically @@ -86,13 +89,16 @@ if __name__ == "__main__": # remove debug sections src = src.split("\t.file")[0] assert '.extern .shared' not in src - prg = ProgramSpec("matmul_kernel", src, device=Device.DEFAULT, - global_size=[M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1], local_size=[32*compiled.metadata.num_warps, 1, 1], - mem_estimate=A.nbytes() + B.nbytes() + C.nbytes()) - ei = ExecItem(si.ast, [x.ensure_allocated() for x in si.bufs], si.metadata, prg=CompiledRunner(prg)) + info = ProgramInfo(name="matmul_kernel", + global_size=(M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1), local_size=(32*compiled.metadata.num_warps, 1, 1)) + sink = UOp.sink(arg=KernelInfo(name="matmul_kernel")) + prg_uop = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=src)), arg=info) + runner = CompiledRunner(prg_uop, Device.DEFAULT) + all_bufs = [x.ensure_allocated() for x in bufs] + prg_bufs = [all_bufs[i] for i in runner.p.globals] tflops = [] for i in range(5): - tm = ei.run(wait=True) + tm = runner(prg_bufs, {}, wait=True) tflops.append((2*M*K*N/tm)*1e-12) print(f"TFLOPS: {max(tflops):.2f}") diff --git a/extra/gemm/tvm_gemm.py b/extra/gemm/tvm_gemm.py index da58df2438..f13f35682e 100644 --- a/extra/gemm/tvm_gemm.py +++ b/extra/gemm/tvm_gemm.py @@ -36,10 +36,10 @@ A = Tensor.rand(M, K, device="CPU") B = Tensor.rand(K, N, device="CPU") C = (A.reshape(M, 1, K) * B.permute(1,0).reshape(1, N, K)).sum(axis=2) -sched = C.schedule() +linear = C.schedule_linear() from tinygrad.codegen.opt.kernel import Kernel from tinygrad.device import CompilerOptions -lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False)) +lin = Kernel(linear.src[-1].src[0], CompilerOptions(has_local=False, supports_float4=False)) lin.to_program() from tinygrad.runtime.ops_cpu import renderer src = renderer("mmult", lin.uops) diff --git a/extra/llama_kernels/__init__.py b/extra/llama_kernels/__init__.py new file mode 100644 index 0000000000..00a6b1c98b --- /dev/null +++ b/extra/llama_kernels/__init__.py @@ -0,0 +1,35 @@ +from __future__ import annotations +import functools, pathlib +from tinygrad import Tensor, dtypes +from tinygrad.uop.ops import Ops +from tinygrad.runtime.support.compiler_amd import HIPCCCompiler + +FP8_MAX = 448.0 +NUM_WG, THREADS_PER_WG = 1024, 256 + +# per-device abs max without allreduce +@functools.cache +def _local_abs_max_fxn(x_p, device): + x = Tensor(x_p, device=device) + inner = Tensor(x.uop.src[0]) if x.uop.op is Ops.MULTI else x + return (inner.abs().max(),) + +def local_abs_max(x:Tensor) -> Tensor: + param = x.as_param(0) + fxn = _local_abs_max_fxn(param.uop, x.device) + return Tensor(fxn[0].uop.call(x.uop).gettuple(0)) + +def scalar_amax(amax_buf:Tensor) -> Tensor: + if isinstance(amax_buf.device, tuple): + return local_abs_max(amax_buf).detach() + return amax_buf.max().detach() + +def shard_shape(shape:tuple, axis:int, ndev:int) -> list: + s = list(shape) + s[axis] //= ndev + return s + +def compile_cpp(cpp_dir:pathlib.Path, cpp_name:str, n_elems:int, hidden:int): + src = (cpp_dir/cpp_name).read_text() + defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={hidden}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"] + return src, HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src) diff --git a/extra/llama_kernels/cast_amax/__init__.py b/extra/llama_kernels/cast_amax/__init__.py new file mode 100644 index 0000000000..64db3d905a --- /dev/null +++ b/extra/llama_kernels/cast_amax/__init__.py @@ -0,0 +1,73 @@ +from __future__ import annotations +import functools, pathlib +from tinygrad import Tensor, dtypes +from tinygrad.uop.ops import UOp, Ops, KernelInfo +from tinygrad.renderer import Estimates +from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, compile_cpp, shard_shape, scalar_amax + +@functools.cache +def _custom_fused_bwd_w13(grad_xw13:UOp, xw13:UOp, grad_x2:UOp, amax_state:UOp, dname:str) -> UOp: + hidden = xw13.shape[2] // 2 + n_elems = xw13.shape[0] * xw13.shape[1] * hidden + threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0") + mem = n_elems * 2 * 5 + sink = UOp.sink(grad_xw13.base, xw13.base, grad_x2.base, amax_state.base, threads, workgroups, + arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=8*n_elems, mem=mem))) + src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden) + return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), + UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib))) + +@functools.cache +def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, dname:str) -> UOp: + hidden = xw13.shape[2] // 2 + n_elems = xw13.shape[0] * xw13.shape[1] * hidden + threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0") + mem = n_elems * 2 * 2 + n_elems + NUM_WG * 2 + sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups, + arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem))) + src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_fwd_w13.cpp", n_elems, hidden) + return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), + UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib))) + +def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp): + # NOTE: inputs are (fp8_out, amax_buf, xw13, amax_state); grad for xw13 only + _, _, xw13, amax_state = kernel.src[1:] + device = xw13.device + if isinstance(device, tuple): + axis, ndev = xw13.axis, len(device) + assert axis in (0, 1), f"unsupported sharding axis={axis}" + grad_xw13 = Tensor(Tensor.invalids(*shard_shape(xw13.shape, axis, ndev), dtype=dtypes.bfloat16, + device=device).uop.multi(axis), device=device) + dname = device[0].split(":")[0] + else: + grad_xw13 = Tensor.invalids(*xw13.shape, dtype=dtypes.bfloat16, device=device) + dname = device.split(":")[0] if isinstance(device, str) else device + grad_x2_t = Tensor(gradient, device=device).cast(dtypes.bfloat16) + fxn = functools.partial(_custom_fused_bwd_w13, dname=dname) + grad_xw13, *_ = Tensor.custom_kernel(grad_xw13, Tensor(xw13, device=device), grad_x2_t, + Tensor(amax_state, device=device), fxn=fxn) + return (None, None, grad_xw13.uop, None) + +def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[Tensor, Tensor, Tensor]: + # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, inv_scale, new_amax) + assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}" + MBS, SEQ, H2 = xw13.shape + assert H2 % 2 == 0, f"w13 last-axis must be even, got {H2}" + HIDDEN = H2 // 2 + if isinstance(xw13.device, tuple): + axis, ndev = xw13.uop.axis, len(xw13.device) + assert axis in (0, 1), f"unsupported sharding axis={axis}" + fp8_out = Tensor(Tensor.invalids(*shard_shape((MBS, SEQ, HIDDEN), axis, ndev), dtype=fp8_dtype, + device=xw13.device).uop.multi(axis), device=xw13.device) + amax_buf = Tensor(Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device).uop.multi(0), + device=xw13.device) + dname = xw13.device[0].split(":")[0] + else: + fp8_out = Tensor.invalids(MBS, SEQ, HIDDEN, dtype=fp8_dtype, device=xw13.device) + amax_buf = Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device) + dname = xw13.device.split(":")[0] if isinstance(xw13.device, str) else xw13.device + fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname) + fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, fxn=fxn, + grad_fxn=_fused_quantize_bwd_w13) + inv_scale = (amax_state.float() + 1e-8) / FP8_MAX + return fp8_out, inv_scale, scalar_amax(amax_buf) diff --git a/extra/amax/cast_amax_bwd_w13.cpp b/extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp similarity index 100% rename from extra/amax/cast_amax_bwd_w13.cpp rename to extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp diff --git a/extra/amax/cast_amax_fwd_w13.cpp b/extra/llama_kernels/cast_amax/cast_amax_fwd_w13.cpp similarity index 100% rename from extra/amax/cast_amax_fwd_w13.cpp rename to extra/llama_kernels/cast_amax/cast_amax_fwd_w13.cpp diff --git a/extra/llama_kernels/fused_ce/__init__.py b/extra/llama_kernels/fused_ce/__init__.py new file mode 100644 index 0000000000..dc2ad39401 --- /dev/null +++ b/extra/llama_kernels/fused_ce/__init__.py @@ -0,0 +1,98 @@ +from __future__ import annotations +import functools, pathlib +from tinygrad import Tensor, dtypes +from tinygrad.uop.ops import UOp, Ops, KernelInfo +from tinygrad.renderer import Estimates +from tinygrad.runtime.support.compiler_amd import HIPCCCompiler + +THREADS_PER_WG = 256 + +@functools.cache +def _custom_fused_ce_loss_fwd(loss_out:UOp, max_out:UOp, lse_out:UOp, logits:UOp, targets:UOp, + dname:str, vocab:int, rows:int, label_smoothing:float) -> UOp: + threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(rows, "gidx0") + mem = rows * vocab * 2 + rows * 12 + rows * 4 + sink = UOp.sink(loss_out.base, max_out.base, lse_out.base, logits.base, targets.base, + threads, workgroups, + arg=KernelInfo(f"fused_ce_loss_fwd", estimates=Estimates(ops=6*rows*vocab, mem=mem))) + src = (pathlib.Path(__file__).parent/"fused_ce_loss.cpp").read_text() + defines = [f"-DVOCAB={vocab}", f"-DTHREADS_PER_WG={THREADS_PER_WG}", + f"-DLABEL_SMOOTHING={label_smoothing}f"] + lib = HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src) + return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), + UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib))) + +@functools.cache +def _custom_fused_ce_loss_bwd(d_logits:UOp, logits:UOp, lse:UOp, targets:UOp, scale:UOp, + dname:str, vocab:int, rows:int, label_smoothing:float) -> UOp: + threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(rows, "gidx0") + mem = rows * vocab * 4 + rows * 8 + 4 + sink = UOp.sink(d_logits.base, logits.base, lse.base, targets.base, scale.base, + threads, workgroups, + arg=KernelInfo(f"fused_ce_loss_bwd", estimates=Estimates(ops=4*rows*vocab, mem=mem))) + src = (pathlib.Path(__file__).parent/"fused_ce_loss_bwd.cpp").read_text() + defines = [f"-DVOCAB={vocab}", f"-DTHREADS_PER_WG={THREADS_PER_WG}", + f"-DLABEL_SMOOTHING={label_smoothing}f"] + lib = HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src) + return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), + UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib))) + +def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float): + # NOTE: forward inputs are (loss_out, max_out, lse_out, logits, targets) + # gradient is the upstream grad w.r.t. per-row loss (shape: (rows,) fp32) + _, _, lse_u, logits_u, targets_u = kernel.src[1:] + device = logits_u.device + rows_vocab = logits_u.shape # (rows, VOCAB) after reshape + rows, VOCAB = rows_vocab + if isinstance(device, tuple): + axis = logits_u.axis + ndev = len(device) + d_logits = Tensor(Tensor.invalids(rows // ndev, VOCAB, dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device) + dname = device[0].split(":")[0] + rows_per_dev = rows // ndev + else: + d_logits = Tensor.invalids(rows, VOCAB, dtype=dtypes.bfloat16, device=device) + dname = device.split(":")[0] if isinstance(device, str) else device + rows_per_dev = rows + grad_t = Tensor(gradient, device=device).float().reshape(-1) # (rows,) fp32 + # NOTE: .mean() backward gives same grad per row (1/N), so broadcast is safe; take scalar + scale = grad_t[0:1].contiguous() + logits_t = Tensor(logits_u.after(kernel), device=device) + lse_t = Tensor(lse_u.after(kernel), device=device) + targets_t = Tensor(targets_u, device=device) + fxn = functools.partial(_custom_fused_ce_loss_bwd, dname=dname, vocab=VOCAB, rows=rows_per_dev, label_smoothing=label_smoothing) + d_logits, *_ = Tensor.custom_kernel(d_logits, logits_t, lse_t, targets_t, scale, fxn=fxn) + return (None, None, None, d_logits.uop, None) + +def fused_ce_loss(logits:Tensor, targets:Tensor, label_smoothing:float=0.1) -> Tensor: + # NOTE: fused sparse_categorical_crossentropy with label smoothing, returns mean loss scalar + assert logits.dtype == dtypes.bfloat16, f"expected bf16, got {logits.dtype}" + assert logits.ndim == 3, f"expected (MBS, SEQ, VOCAB), got {logits.shape}" + MBS, SEQ, VOCAB = logits.shape + rows = MBS * SEQ + if isinstance(logits.device, tuple): + axis = logits.uop.axis + assert axis in (0, 1), f"unsupported sharding axis={axis} for CE loss" + ndev = len(logits.device) + loss_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0), + device=logits.device) + max_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0), + device=logits.device) + lse_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0), + device=logits.device) + dname = logits.device[0].split(":")[0] + rows_per_dev = rows // ndev + else: + loss_out = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device) + max_out = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device) + lse_out = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device) + dname = logits.device.split(":")[0] if isinstance(logits.device, str) else logits.device + rows_per_dev = rows + logits_flat = logits.reshape(rows, VOCAB) + targets_flat = targets.reshape(-1).cast(dtypes.int32) + fxn = functools.partial(_custom_fused_ce_loss_fwd, dname=dname, vocab=VOCAB, rows=rows_per_dev, + label_smoothing=label_smoothing) + loss_out, max_out, lse_out, *_ = Tensor.custom_kernel( + loss_out, max_out, lse_out, logits_flat, targets_flat, + fxn=fxn, grad_fxn=functools.partial(_fused_ce_loss_bwd, label_smoothing=label_smoothing)) + return loss_out.mean() diff --git a/extra/llama_kernels/fused_ce/fused_ce_loss.cpp b/extra/llama_kernels/fused_ce/fused_ce_loss.cpp new file mode 100644 index 0000000000..cdbd2d630a --- /dev/null +++ b/extra/llama_kernels/fused_ce/fused_ce_loss.cpp @@ -0,0 +1,104 @@ +#include +#include + +// Fused forward sparse-CE with label smoothing. +// SINGLE-PASS online softmax + vectorized 8-wide bf16 loads for HBM coalescing. + +#ifndef VOCAB +#define VOCAB 128256 +#endif +#ifndef THREADS_PER_WG +#define THREADS_PER_WG 256 +#endif +#ifndef LABEL_SMOOTHING +#define LABEL_SMOOTHING 0.1f +#endif + +constexpr int VEC = 8; + +extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void +fused_ce_loss_fwd( + float* __restrict__ loss_out, // out: fp32, ROWS + float* __restrict__ max_out, // out: fp32, ROWS + float* __restrict__ lse_out, // out: fp32, ROWS + const __hip_bfloat16* __restrict__ logits, // in: bf16, ROWS*VOCAB + const int* __restrict__ targets) // in: int32, ROWS +{ + __shared__ float sdata_m[THREADS_PER_WG]; + __shared__ float sdata_s[THREADS_PER_WG]; + __shared__ float sdata_sumx[THREADS_PER_WG]; + __shared__ float sdata_tgt[THREADS_PER_WG]; + + const int tid = threadIdx.x; + const int row = blockIdx.x; + const int target = targets[row]; + const __hip_bfloat16* row_logits = logits + (size_t)row * VOCAB; + + float m = -INFINITY; + float s = 0.0f; + float sum_x = 0.0f; + float target_logit = 0.0f; + constexpr bool needs_sum_x = (LABEL_SMOOTHING != 0.0f); + + // Vectorized stride: each iter loads 8 bf16 = 16 bytes. Warp loads 32*16 = 512 bytes (4 cache lines). + const int VOCAB_VEC = VOCAB & ~(VEC - 1); // round down to multiple of VEC + for (int i = tid * VEC; i < VOCAB_VEC; i += THREADS_PER_WG * VEC) { + float4 raw = *reinterpret_cast(&row_logits[i]); + const __hip_bfloat16* xi = reinterpret_cast(&raw); + #pragma unroll + for (int k = 0; k < VEC; k++) { + const float x = static_cast(xi[k]); + if constexpr (needs_sum_x) sum_x += x; + if (i + k == target) target_logit = x; + if (x > m) { + s = s * __expf(m - x) + 1.0f; + m = x; + } else { + s += __expf(x - m); + } + } + } + // tail (VOCAB not divisible by VEC): + for (int i = VOCAB_VEC + tid; i < VOCAB; i += THREADS_PER_WG) { + const float x = static_cast(row_logits[i]); + if constexpr (needs_sum_x) sum_x += x; + if (i == target) target_logit = x; + if (x > m) { s = s * __expf(m - x) + 1.0f; m = x; } + else { s += __expf(x - m); } + } + + sdata_m[tid] = m; + sdata_s[tid] = s; + sdata_sumx[tid] = sum_x; + sdata_tgt[tid] = target_logit; + __syncthreads(); + + for (int step = THREADS_PER_WG / 2; step > 0; step >>= 1) { + if (tid < step) { + const float m1 = sdata_m[tid]; + const float m2 = sdata_m[tid + step]; + const float s1 = sdata_s[tid]; + const float s2 = sdata_s[tid + step]; + const float m_new = fmaxf(m1, m2); + const float s_new = s1 * __expf(m1 - m_new) + s2 * __expf(m2 - m_new); + sdata_m[tid] = m_new; + sdata_s[tid] = s_new; + sdata_sumx[tid] += sdata_sumx[tid + step]; + sdata_tgt[tid] += sdata_tgt[tid + step]; + } + __syncthreads(); + } + + if (tid == 0) { + const float row_max = sdata_m[0]; + const float row_sum_exp = sdata_s[0]; + const float row_sum_x = sdata_sumx[0]; + const float tgt = sdata_tgt[0]; + const float row_lse = logf(row_sum_exp) + row_max; + const float mean_logits = row_sum_x / static_cast(VOCAB); + const float loss = row_lse - (1.0f - LABEL_SMOOTHING) * tgt - LABEL_SMOOTHING * mean_logits; + loss_out[row] = loss; + max_out[row] = row_max; + lse_out[row] = row_lse; + } +} diff --git a/extra/llama_kernels/fused_ce/fused_ce_loss_bwd.cpp b/extra/llama_kernels/fused_ce/fused_ce_loss_bwd.cpp new file mode 100644 index 0000000000..70e70cb726 --- /dev/null +++ b/extra/llama_kernels/fused_ce/fused_ce_loss_bwd.cpp @@ -0,0 +1,58 @@ +#include +#include + +// Vectorized CE bwd: 8-wide bf16 loads + stores. + +#ifndef VOCAB +#define VOCAB 128256 +#endif +#ifndef THREADS_PER_WG +#define THREADS_PER_WG 256 +#endif +#ifndef LABEL_SMOOTHING +#define LABEL_SMOOTHING 0.1f +#endif + +constexpr int VEC = 8; + +extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void +fused_ce_loss_bwd( + __hip_bfloat16* __restrict__ d_logits, + const __hip_bfloat16* __restrict__ logits, + const float* __restrict__ lse, + const int* __restrict__ targets, + const float* __restrict__ scale_in) +{ + const int tid = threadIdx.x; + const int row = blockIdx.x; + const int target = targets[row]; + const float lse_r = lse[row]; + const __hip_bfloat16* row_logits = logits + (size_t)row * VOCAB; + __hip_bfloat16* row_dlogits = d_logits + (size_t)row * VOCAB; + const float inv_vocab = 1.0f / static_cast(VOCAB); + const float scale = *scale_in; + const float ls_term = LABEL_SMOOTHING * inv_vocab; + + const int VOCAB_VEC = VOCAB & ~(VEC - 1); + for (int i = tid * VEC; i < VOCAB_VEC; i += THREADS_PER_WG * VEC) { + float4 raw = *reinterpret_cast(&row_logits[i]); + const __hip_bfloat16* xi = reinterpret_cast(&raw); + __hip_bfloat16 out[VEC]; + #pragma unroll + for (int k = 0; k < VEC; k++) { + const float x = static_cast(xi[k]); + float g = __expf(x - lse_r); + if (i + k == target) g -= (1.0f - LABEL_SMOOTHING); + g -= ls_term; + out[k] = static_cast<__hip_bfloat16>(g * scale); + } + *reinterpret_cast(&row_dlogits[i]) = *reinterpret_cast(out); + } + for (int i = VOCAB_VEC + tid; i < VOCAB; i += THREADS_PER_WG) { + const float x = static_cast(row_logits[i]); + float g = __expf(x - lse_r); + if (i == target) g -= (1.0f - LABEL_SMOOTHING); + g -= ls_term; + row_dlogits[i] = static_cast<__hip_bfloat16>(g * scale); + } +} diff --git a/extra/llama_kernels/fused_mul_quantize_fp8/__init__.py b/extra/llama_kernels/fused_mul_quantize_fp8/__init__.py new file mode 100644 index 0000000000..dfba644efa --- /dev/null +++ b/extra/llama_kernels/fused_mul_quantize_fp8/__init__.py @@ -0,0 +1,54 @@ +from __future__ import annotations +import functools, pathlib +from tinygrad import Tensor, dtypes +from tinygrad.uop.ops import UOp, Ops, KernelInfo +from tinygrad.renderer import Estimates +from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, compile_cpp, shard_shape, scalar_amax + +@functools.cache +def _custom_mul_quantize_fp8(fp8_out:UOp, amax_buf:UOp, x:UOp, weight:UOp, amax_state:UOp, dname:str) -> UOp: + MBS, SEQ, HIDDEN = x.shape + n_elems = MBS * SEQ * HIDDEN + threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0") + mem = n_elems * 2 + HIDDEN * 2 + n_elems + NUM_WG * 2 + sink = UOp.sink(fp8_out.base, amax_buf.base, x.base, weight.base, amax_state.base, threads, workgroups, + arg=KernelInfo(f"fused_mul_quantize_fp8_{n_elems}_h{HIDDEN}", estimates=Estimates(ops=3*n_elems, mem=mem))) + src, lib = compile_cpp(pathlib.Path(__file__).parent, "fused_mul_quantize_fp8.cpp", n_elems, HIDDEN) + return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), + UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib))) + +def _fused_mul_quantize_fp8_bwd(gradient:UOp, kernel:UOp): + # NOTE: inputs are (fp8_out, amax_buf, x, weight, amax_state); grads for x and weight + _, _, x_u, weight_u, amax_state_u = kernel.src[1:] + device = x_u.device + grad_t = Tensor(gradient, device=device).cast(dtypes.bfloat16) + x_t, weight_t = Tensor(x_u, device=device), Tensor(weight_u, device=device) + scale = FP8_MAX / (Tensor(amax_state_u, device=device).float() + 1e-8) + grad_scaled = grad_t.float() * scale + # NOTE: grad_x stays bf16 to avoid CSE materializing a (MBS, SEQ, HIDDEN) fp32 intermediate + grad_x = (grad_scaled * weight_t.float()).cast(dtypes.bfloat16) + grad_weight = (grad_scaled * x_t.float()).sum(axis=(0, 1)).cast(dtypes.bfloat16) + return (None, None, grad_x.uop, grad_weight.uop, None) + +def fused_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[Tensor, Tensor, Tensor]: + # NOTE: (x * weight) -> fp8 + amax, delayed scaling. Returns (fp8, inv_scale, new_amax) + assert x.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16 + assert x.shape[-1] == weight.shape[-1], f"HIDDEN mismatch: x={x.shape}, weight={weight.shape}" + MBS, SEQ, HIDDEN = x.shape + if isinstance(x.device, tuple): + axis, ndev = x.uop.axis, len(x.device) + assert axis in (0, 1), f"unsupported sharding axis={axis}" + fp8_out = Tensor(Tensor.invalids(*shard_shape((MBS, SEQ, HIDDEN), axis, ndev), dtype=fp8_dtype, + device=x.device).uop.multi(axis), device=x.device) + amax_buf = Tensor(Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=x.device).uop.multi(0), device=x.device) + dname = x.device[0].split(":")[0] + else: + fp8_out = Tensor.invalids(MBS, SEQ, HIDDEN, dtype=fp8_dtype, device=x.device) + amax_buf = Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=x.device) + dname = x.device.split(":")[0] if isinstance(x.device, str) else x.device + fxn = functools.partial(_custom_mul_quantize_fp8, dname=dname) + fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, x, weight, amax_state, fxn=fxn, + grad_fxn=_fused_mul_quantize_fp8_bwd) + new_amax = scalar_amax(amax_buf) + inv_scale = (amax_state.float() + 1e-8) / FP8_MAX + return fp8_out, inv_scale, new_amax diff --git a/extra/amax/fused_mul_quantize_fp8.cpp b/extra/llama_kernels/fused_mul_quantize_fp8/fused_mul_quantize_fp8.cpp similarity index 100% rename from extra/amax/fused_mul_quantize_fp8.cpp rename to extra/llama_kernels/fused_mul_quantize_fp8/fused_mul_quantize_fp8.cpp diff --git a/extra/llama_kernels/rmsnorm/__init__.py b/extra/llama_kernels/rmsnorm/__init__.py new file mode 100644 index 0000000000..d7c4350794 --- /dev/null +++ b/extra/llama_kernels/rmsnorm/__init__.py @@ -0,0 +1,24 @@ +from __future__ import annotations +import functools +from tinygrad import Tensor +from tinygrad.uop.ops import UOp + +def rmsnorm_fwd(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]: + x = x_in.float() + rrms = (x.square().mean(-1, keepdim=True) + eps).rsqrt() + return (x * rrms).cast(x_in.dtype), rrms + +@functools.cache +def _rmsnorm_fwd_fxn(x_in_p, eps, device): + return rmsnorm_fwd(Tensor(x_in_p, device=device), eps) + +def _rmsnorm_bwd(grad:UOp, call:UOp) -> tuple: + x_normed = Tensor(call.gettuple(0)).float() + do_float = Tensor(grad).float() + d_x = Tensor(call.gettuple(1)) * (do_float - x_normed * (do_float * x_normed).mean(-1, keepdim=True)) + return (d_x.cast(call.src[1].dtype).uop,) + +def rmsnorm(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]: + fxn = _rmsnorm_fwd_fxn(x_in.as_param(0).uop, eps, x_in.device) + call = UOp.maketuple(fxn[0].uop, fxn[1].uop).call(x_in.uop, grad_fxn=_rmsnorm_bwd) + return Tensor(call.gettuple(0)), Tensor(call.gettuple(1)) diff --git a/extra/mmapeak/mmapeak.py b/extra/mmapeak/mmapeak.py index b5bb581db6..0fd6ade5bd 100644 --- a/extra/mmapeak/mmapeak.py +++ b/extra/mmapeak/mmapeak.py @@ -3,11 +3,12 @@ import os # TODO: there is a timing bug without this os.environ["AMD_AQL"] = "1" -from tinygrad import Tensor, Device +from tinygrad import Tensor, Device, GlobalCounters, Context from tinygrad.helpers import getenv, DEV from tinygrad.uop.ops import UOp, Ops, KernelInfo from tinygrad.renderer import Estimates from tinygrad.renderer.amd.dsl import Reg, Inst, s, v +from tinygrad.engine.realize import run_linear NUM_WORKGROUPS = 96 WAVE_SIZE = 32 @@ -36,11 +37,17 @@ def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs) gidx = UOp.special(NUM_WORKGROUPS, "gidx0") FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP sink = UOp.sink(A.base, threads, gidx, arg=KernelInfo(inst.op.name.lower(), estimates=Estimates(ops=FLOPs, mem=0))) - return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts])))) + return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts])))) dummy = Tensor.zeros(1).contiguous().realize() out = Tensor.custom_kernel(dummy, fxn=fxn)[0] - ei = out.schedule()[-1].lower() - elapsed = min([ei.run(wait=True) for _ in range(2)]) + linear = out.schedule_linear() + ets = [] + with Context(DEBUG=2): + for _ in range(2): + start = GlobalCounters.time_sum_s + run_linear(linear) + ets.append(GlobalCounters.time_sum_s - start) + elapsed = min(ets) FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP print(f"{inst.op_name.lower():<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS") diff --git a/extra/optimization/test_beam_search.py b/extra/optimization/test_beam_search.py index 6ff2bbd36c..a7cc18f83f 100644 --- a/extra/optimization/test_beam_search.py +++ b/extra/optimization/test_beam_search.py @@ -84,7 +84,7 @@ class TestBeamSearch(unittest.TestCase): tc = Device[Device.DEFAULT].renderer.tensor_cores[0] size = max(tc.dims[0], tc.dims[1]) * 8 a, b = Tensor.rand(size, size, dtype=tc.dtype_in), Tensor.rand(size, size, dtype=tc.dtype_in) - ast = a.matmul(b, dtype=tc.dtype_out).schedule()[-1].ast + ast = a.matmul(b, dtype=tc.dtype_out).schedule_linear().src[-1].src[0] s = Scheduler(ast, Device[Device.DEFAULT].renderer) s.apply_opt(Opt(OptOps.TC, 0, (-1, 0, 1))) up = prod([x for x, t in zip(s.full_shape, s.axis_types) if t in (AxisType.UPCAST, AxisType.UNROLL)]) @@ -94,7 +94,7 @@ class TestBeamSearch(unittest.TestCase): def test_max_up(self): a = Tensor.rand(16, 16) - ast = a.schedule()[-1].ast + ast = a.schedule_linear().src[-1].src[0] s = Scheduler(ast, Device[Device.DEFAULT].renderer) for max_up in (2, 4): actions = get_kernel_actions(s, include_0=False, max_up=max_up) diff --git a/extra/sqtt/examples/gfx1200/profile_handwritten_run_0.pkl b/extra/sqtt/examples/gfx1200/profile_handwritten_run_0.pkl index f0c862b6b9..4add9d0c58 100644 Binary files a/extra/sqtt/examples/gfx1200/profile_handwritten_run_0.pkl and b/extra/sqtt/examples/gfx1200/profile_handwritten_run_0.pkl differ diff --git a/extra/sqtt/examples/gfx1200/profile_handwritten_run_1.pkl b/extra/sqtt/examples/gfx1200/profile_handwritten_run_1.pkl index 8f8dc9f114..9ffaa18f0a 100644 Binary files a/extra/sqtt/examples/gfx1200/profile_handwritten_run_1.pkl and b/extra/sqtt/examples/gfx1200/profile_handwritten_run_1.pkl differ diff --git a/test/amd/test_custom_kernel.py b/test/amd/test_custom_kernel.py index a003f455c8..be01156659 100644 --- a/test/amd/test_custom_kernel.py +++ b/test/amd/test_custom_kernel.py @@ -3,12 +3,14 @@ import functools import numpy as np from tinygrad import Tensor, Device, dtypes from tinygrad.uop.ops import UOp, Ops, KernelInfo +from tinygrad.engine.realize import run_linear, estimate_uop from tinygrad.renderer import Estimates from tinygrad.dtype import AddrSpace +from tinygrad.helpers import getenv from tinygrad.runtime.autogen.amd.rdna3.ins import * import tinygrad.runtime.autogen.amd.rdna3.ins as r3 import tinygrad.runtime.autogen.amd.rdna4.ins as r4 -from tinygrad.renderer.amd.dsl import s, v +from tinygrad.renderer.amd.dsl import s, v, NULL from test.amd.helpers import TARGET_TO_ARCH from extra.gemm.amd_asm_matmul import Kernel @@ -100,31 +102,44 @@ def custom_lds_sync(A:UOp, arch:str) -> UOp: def custom_handwritten(A:UOp, arch:str) -> UOp: A = A.flatten() threads = UOp.special(128, "lidx0") - wg = UOp.special(256, "gidx0") + wg = UOp.special(1, "gidx0") lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=512, addrspace=AddrSpace.LOCAL), (), 'lds') # 128 * 4 bytes + pipes = {getenv("PIPE", "")} if getenv("PIPE", "") else {"SALU", "VALU", "TRANSCENDENTAL", "WMMA"} k = Kernel(arch) - k.emit(r4.s_nop(0)) - k.emit(r4.v_mov_b32_e32(v[1], 4)) - def emit_alt(): - for i in range(2): - k.emit(r4.v_mov_b32_e32(v[20+i], 4.0)) - k.emit(r4.v_rcp_f32_e32(v[22+i], v[20+i])) - k.emit(r4.s_mov_b32(s[20+i], i)) - k.emit(r4.s_mul_i32(s[14+i], s[12+i], 32)) - def emit_wmma(): - for _ in range(2): - k.emit(r4.v_wmma_f32_16x16x16_f16(v[0:7], v[8:11], v[8:11], 1)) - k.label("start") - k.emit(s_mov_b32(s[1], 10)) + # wrap in loop to filter out icache misses + LOOP_N, UNROLL_N = 8, 5 + k.emit(r4.s_mov_b32(s[1], LOOP_N)) k.label("loop") - # wmma should've overlapped here if it was a different unit? - for _ in range(2): - emit_wmma() - emit_alt() - for _ in range(8): k.emit(s_nop(1)) - k.emit(s_add_u32(s[1], s[1], -1)) - k.emit(s_cmp_eq_i32(s[1], 0)) - k.emit(s_cbranch_scc0(), target="loop") + if "SALU" in pipes: + for i in range(UNROLL_N): + k.emit(r4.s_mov_b32(s[20+i], i)) + k.emit(r4.s_min_i32(s[30+i], i)) + k.emit(r4.s_mov_b32(s[40+i], i)) + k.emit(r4.s_mul_i32(s[14+i], s[12+i], 32)) + if "VALU" in pipes: + for i in range(UNROLL_N): + k.emit(r4.v_mov_b32_e32(v[20+i], i)) + k.emit(r4.v_lshlrev_b64_e32(v[30+2*i:31+2*i], 2, v[12+i:13+i])) + k.emit(r4.v_mad_co_u64_u32(v[40+2*i:41+2*i], NULL, v[12+i], v[13+i], v[14+i:15+i])) + if "TRANSCENDENTAL" in pipes: + # transcendental VALU runs on the TFU, it can run regular VALU at the same time + for i in range(UNROLL_N): + k.emit(r4.v_mov_b32_e32(v[20+i], i)) + k.emit(r4.v_s_rcp_f32(s[20+i], s[12+i])) + k.emit(r4.v_rcp_f32_e32(v[30+i], v[12+i])) + k.emit(r4.v_s_exp_f32(s[30+i], s[12+i])) + if "WMMA" in pipes: + base = 30 + for i in range(UNROLL_N): + a = base + i*40 + b, cd = a + 4, a + 8 + k.emit(r4.v_wmma_f32_16x16x16_f16(v[cd:cd+7], v[a:a+3], v[b:b+3], v[cd:cd+7])) + a = base + i*40 + 16 + b, cd = a + 2, a + 4 + k.emit(r4.v_wmma_i32_16x16x16_iu8(v[cd:cd+7], v[a:a+1], v[b:b+1], v[cd:cd+7])) + k.emit(r4.s_add_co_i32(s[1], s[1], -1)) + k.emit(r4.s_cmp_eq_i32(s[1], 0)) + k.emit(r4.s_cbranch_scc0(), target="loop") k.emit(r4.s_endpgm()) insts = k.finalize() sink = UOp.sink(A.base, threads, wg, lds, arg=KernelInfo("custom_handwritten")) @@ -154,10 +169,11 @@ class TestCustomKernel(unittest.TestCase): if self.arch != "rdna3": self.skipTest("only rdna3") a = Tensor.full((16, 16), 1.).contiguous().realize() a = Tensor.custom_kernel(a, fxn=custom_add_one)[0] - ei = a.schedule()[-1].lower() - self.assertEqual(ei.prg.estimates.ops, a.numel()) - self.assertEqual(ei.prg.estimates.mem, a.nbytes()*2) - ei.run() + linear = a.schedule_linear() + est = estimate_uop(linear.src[-1]) + self.assertEqual(est.ops, a.numel()) + self.assertEqual(est.mem, a.nbytes()*2) + run_linear(linear) self.assertTrue((a.numpy() == 2.).all()) def test_variable(self): @@ -165,9 +181,9 @@ class TestCustomKernel(unittest.TestCase): b = Tensor.full((16, 16), 1, dtype=dtypes.uint32).contiguous().realize() a = Tensor.zeros_like(b).contiguous().realize() a = Tensor.custom_kernel(a, b, fxn=custom_add_var)[0] - ei = a.schedule()[-1].lower() + linear = a.schedule_linear() for i in range(4): - ei.run({"var":i}) + run_linear(linear, var_vals={"var":i}) self.assertTrue((a.numpy() == 1+i).all()) def test_lds_sync(self): diff --git a/test/amd/test_integration.py b/test/amd/test_integration.py index 5ca10c16b0..a6ac267b1c 100644 --- a/test/amd/test_integration.py +++ b/test/amd/test_integration.py @@ -78,18 +78,18 @@ class TestTinygradIntegration(unittest.TestCase): def _get_kernel_code(self, op_fn) -> bytes: from tinygrad import Tensor from tinygrad.helpers import Target - from tinygrad.codegen import get_program + from tinygrad.codegen import to_program from tinygrad.renderer.llvmir import AMDLLVMRenderer from tinygrad.runtime.support.elf import elf_loader from tinygrad.uop.ops import Ops result = op_fn(Tensor) - schedule = result.schedule() - sink_items = [si for si in schedule if si.ast.op == Ops.SINK] + linear = result.schedule_linear() + sink_items = [call for call in linear.src if call.src[0].op == Ops.SINK] assert len(sink_items) > 0, "No SINK in schedule" renderer = AMDLLVMRenderer(Target("AMD", arch='gfx1100')) - prg = get_program(sink_items[0].ast, renderer) - lib = renderer.compiler.compile(prg.src) + prg = to_program(sink_items[0].src[0], renderer) + lib = renderer.compiler.compile(prg.src[3].arg) return next(s.content for s in elf_loader(lib)[1] if s.name == ".text") def test_simple_add_kernel(self): diff --git a/test/amd/test_mockgpu_invalid.py b/test/amd/test_mockgpu_invalid.py index 91c8f412fb..f5692cb68f 100644 --- a/test/amd/test_mockgpu_invalid.py +++ b/test/amd/test_mockgpu_invalid.py @@ -14,8 +14,8 @@ from tinygrad.runtime.ops_amd import AMDProgram dev = Device["AMD"] a = Tensor([1.0]).realize() b = a + 1 -si = b.schedule()[-1] -runner = get_runner(dev.device, si.ast) +si = b.schedule_linear().src[-1] +runner = get_runner(dev.device, si.src[0]) prg = runner._prg lib = bytearray(prg.lib) diff --git a/test/amd/test_roundtrip.py b/test/amd/test_roundtrip.py index f3dde9ad20..3f1d85b52b 100644 --- a/test/amd/test_roundtrip.py +++ b/test/amd/test_roundtrip.py @@ -57,49 +57,53 @@ class KernelSnapshot: def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, int], dict[int, bytes]]: """Compile a tinygrad operation and extract all kernels with their buffer mappings.""" from tinygrad import Tensor + from tinygrad.uop.ops import Ops + from tinygrad.engine.realize import compile_linear, resolve_params, unwrap_multi from tinygrad.runtime.support.elf import elf_loader out = op_fn(Tensor) - sched = out.schedule() + linear = compile_linear(out.schedule_linear()) kernels = [] buf_pool: dict[int, int] = {} # buffer id -> size buf_data: dict[int, bytes] = {} # buffer id -> initial data from COPY - for ei in sched: - lowered = ei.lower() - if ei.ast.op.name == 'COPY': - # Handle COPY: extract source data to initialize destination buffer - if len(lowered.bufs) >= 2: - dst_buf, src_buf = lowered.bufs[0], lowered.bufs[1] - dst_id = id(dst_buf) - if dst_id not in buf_pool: - buf_pool[dst_id] = dst_buf.nbytes - # Get source data if it's from numpy/CPU - if hasattr(src_buf, 'base') and src_buf.base is not None and hasattr(src_buf.base, '_buf'): - src_data = bytes(src_buf.base._buf) - buf_data[dst_id] = src_data - elif ei.ast.op.name == 'SINK': - if lowered.prg and lowered.prg.p.lib: - lib = bytes(lowered.prg.p.lib) - _, sections, _ = elf_loader(lib) - for sec in sections: - if sec.name == '.text': - buf_idxs = [] - buf_sizes = [] - for b in lowered.bufs: - buf_id = id(b) - if buf_id not in buf_pool: - buf_pool[buf_id] = b.nbytes - buf_idxs.append(buf_id) - buf_sizes.append(b.nbytes) - kernels.append(KernelSnapshot( - code=bytes(sec.content), - src=lowered.prg.p.src, - global_size=tuple(lowered.prg.p.global_size), - local_size=tuple(lowered.prg.p.local_size), - buf_idxs=buf_idxs, - buf_sizes=buf_sizes - )) + for call in linear.src: + ast = call.src[0] + for bufs, _ in unwrap_multi(call, resolve_params(call, ())): + if ast.op is Ops.COPY: + # Handle COPY: extract source data to initialize destination buffer + if len(bufs) >= 2: + dst_buf, src_buf = bufs[0], bufs[1] + dst_id = id(dst_buf) + if dst_id not in buf_pool: + buf_pool[dst_id] = dst_buf.nbytes + # Get source data if it's from numpy/CPU + if hasattr(src_buf, 'base') and src_buf.base is not None and hasattr(src_buf.base, '_buf'): + src_data = bytes(src_buf.base._buf) + buf_data[dst_id] = src_data + elif ast.op is Ops.PROGRAM: + info = ast.arg + if len(ast.src) > 4 and ast.src[4].op is Ops.BINARY: + lib = bytes(ast.src[4].arg) + _, sections, _ = elf_loader(lib) + for sec in sections: + if sec.name == '.text': + buf_idxs = [] + buf_sizes = [] + for b in bufs: + buf_id = id(b) + if buf_id not in buf_pool: + buf_pool[buf_id] = b.nbytes + buf_idxs.append(buf_id) + buf_sizes.append(b.nbytes) + kernels.append(KernelSnapshot( + code=bytes(sec.content), + src=ast.src[3].arg, + global_size=tuple(info.global_size), + local_size=tuple(info.local_size), + buf_idxs=buf_idxs, + buf_sizes=buf_sizes + )) if not kernels: raise RuntimeError("No kernel found") return kernels, buf_pool, buf_data diff --git a/test/amd/test_sqtt_examples.py b/test/amd/test_sqtt_examples.py index 5d5e47c4f8..f8326c2f93 100644 --- a/test/amd/test_sqtt_examples.py +++ b/test/amd/test_sqtt_examples.py @@ -8,10 +8,11 @@ from tinygrad.runtime.support.elf import elf_loader from tinygrad.renderer.amd import decode_inst from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp -from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, INST, INST_RDNA4, VALUINST, +from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, WAVEEND_RDNA4, INST, INST_RDNA4, VALUINST, IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4, PACKET_TYPES_CDNA, CDNA_WAVESTART, print_packets, CDNA_WAVEEND, CDNA_INST) from test.amd.helpers import TARGET_TO_ARCH +from test.amd.test_sqttmap import needs_rocprof import tinygrad EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples" @@ -132,7 +133,7 @@ class SQTTExamplesTestBase(unittest.TestCase): with self.subTest(example=name): all_packets = [p for e in events for p in decode(e.blob)] self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_RDNA4, CDNA_WAVESTART))]), 0, f"no WAVESTART in {name}") - self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVEEND, CDNA_WAVEEND))]), 0, f"no WAVEEND in {name}") + self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND))]), 0, f"no WAVEEND in {name}") def test_time_monotonic(self): for name, (events, *_) in self.examples.items(): @@ -160,6 +161,7 @@ class SQTTExamplesTestBase(unittest.TestCase): counts = [len(list(decode(e.blob))) for e in events] self.assertEqual(counts, self.expected[name], f"packet count mismatch in {name}") + @needs_rocprof def test_rocprof_wave_times_match(self): """Wave start/end times must match rocprof exactly.""" for name, (events, lib, base) in self.examples.items(): @@ -180,7 +182,7 @@ class SQTTExamplesTestBase(unittest.TestCase): for p in decode(event.blob): if first_timestamp is None: first_timestamp = p._time if isinstance(p, (WAVESTART, CDNA_WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time - elif isinstance(p, (WAVEEND, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts: + elif isinstance(p, (WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts: our_waves.append((wave_starts[key], p._time)) for st in wave_starts.values(): self.assertGreater(st, first_timestamp, "wave start must be after the first packet") @@ -189,6 +191,7 @@ class SQTTExamplesTestBase(unittest.TestCase): for st, et in our_waves: self.assertGreater(et, st, "wave end must be after start") + @needs_rocprof def test_rocprof_inst_times_match(self): """Instruction times must match rocprof exactly (excluding s_endpgm).""" for name, (events, lib, base) in self.examples.items(): diff --git a/test/amd/test_sqtt_profiler.py b/test/amd/test_sqtt_profiler.py index 5f8334b89a..3738bcfcc5 100644 --- a/test/amd/test_sqtt_profiler.py +++ b/test/amd/test_sqtt_profiler.py @@ -1,6 +1,8 @@ import unittest, contextlib from tinygrad import Device, Tensor, Context, TinyJit from tinygrad.device import Compiled, ProfileProgramEvent, ProfileDeviceEvent +from tinygrad.engine.realize import run_linear +from tinygrad.codegen import to_program from tinygrad.viz.serve import load_amd_counters, VizData @contextlib.contextmanager @@ -26,39 +28,41 @@ class TestSQTTProfiler(unittest.TestCase): def test_simple(self): t = Tensor.empty(1) + 1 with save_sqtt() as sqtt: - ei = t.schedule()[0].lower() - ei.run() + linear = t.schedule_linear() + run_linear(linear) + fn_name = to_program(linear.src[0].src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name self.assertEqual(len(sqtt), 1) - self.assertEqual(sqtt[0]["name"], f"SQTT {ei.prg.p.function_name}") + self.assertEqual(sqtt[0]["name"], f"SQTT {fn_name}") def test_multiple_runs(self): t = Tensor.empty(1) + 1 with save_sqtt() as sqtt: - ei = t.schedule()[0].lower() - for _ in range(N:=3): - ei.run() + linear = t.schedule_linear() + for _ in range(N:=3): run_linear(linear) + fn_name = to_program(linear.src[0].src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name self.assertEqual(len(sqtt), N) for i in range(1, N): - self.assertEqual(sqtt[i]["name"], f"SQTT {ei.prg.p.function_name} n{i+1}") + self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name} n{i+1}") def test_multiple_kernels(self): t = ((Tensor.empty(1) + 1).contiguous() + 2) - sched = t.schedule() + linear = t.schedule_linear() with save_sqtt() as sqtt: - for si in sched: si.lower().run() - self.assertEqual(len(sqtt), len(sched)) - for i,k in enumerate(sched): - self.assertEqual(sqtt[i]["name"], f"SQTT {k.lower().prg.p.function_name}") + run_linear(linear) + self.assertEqual(len(sqtt), len(linear.src)) + for i,call in enumerate(linear.src): + fn_name = to_program(call.src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name + self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name}") def test_multiple_kernels_lower(self): t = ((Tensor.empty(1) + 1).contiguous() + 2) - sched = t.schedule() + linear = t.schedule_linear() with save_sqtt() as sqtt: - prgs = [si.lower() for si in sched] - for p in prgs: p.run() - self.assertEqual(len(sqtt), len(sched)) - for i,ei in enumerate(prgs): - self.assertEqual(sqtt[i]["name"], f"SQTT {ei.prg.p.function_name}") + run_linear(linear) + self.assertEqual(len(sqtt), len(linear.src)) + for i,call in enumerate(linear.src): + fn_name = to_program(call.src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name + self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name}") def test_jit(self): @TinyJit diff --git a/test/amd/test_sqttmap.py b/test/amd/test_sqttmap.py index 869befa2b7..490db7b57e 100644 --- a/test/amd/test_sqttmap.py +++ b/test/amd/test_sqttmap.py @@ -1,19 +1,36 @@ # test to compare every packet with the rocprof decoder -import unittest, pickle +import unittest, pickle, functools from typing import Iterator from pathlib import Path -from tinygrad.helpers import DEBUG, getenv, temp, ansistrip +from tinygrad.helpers import DEBUG, getenv, temp, ansistrip, Context from tinygrad.renderer.amd.sqtt import print_packets, map_insts from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm -from tinygrad.viz.serve import sqtt_timeline +from tinygrad.viz.serve import sqtt_timeline, amd_decode from test.amd.disasm import disasm from test.null.test_viz import run_cli import tinygrad EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples" +def needs_rocprof(fn): + @functools.wraps(fn) + def wrapper(self, *args, **kwargs): + # check if latest rocprof is available, if not, skip rocprof comparison tests + # rocprof doesn't have a version string, decode a known pickle to validate it's the latest + try: + from extra.sqtt.roc import decode as roc_decode + with open(EXAMPLES_DIR/"gfx1200"/"profile_plus_run_0.pkl", "rb") as f: + data = pickle.load(f) + sqtt = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"][1] + kern = {e.tag:e for e in data if type(e).__name__ == "ProfileProgramEvent"}[sqtt.kern] + rctx = roc_decode([sqtt], {kern.tag:{addr+kern.base:inst for addr,inst in amd_decode(kern.lib, "gfx1200").items()}}) + insts = [e.time for e in list(rctx.inst_execs.values())[0][0].unpack_insts()] + self.assertListEqual(insts, [28178, 28179, 28180, 28181, 28182, 29882, 29883, 29884, 29885, 30966, 30983, 30985, 30992, 30993]) + except Exception as e: self.skipTest(f"latest rocprof not available, install with extra/sqtt/install_rocprof_decoder.py: {e}") + return fn(self, *args, **kwargs) + return wrapper + def rocprof_inst_traces_match(sqtt, prg, target): - from tinygrad.viz.serve import amd_decode from extra.sqtt.roc import decode as roc_decode, InstExec addr_table = amd_decode(prg.lib, target) disasm_map = {addr+prg.base:inst for addr,inst in addr_table.items()} @@ -63,6 +80,7 @@ class TestSQTTMapBase(unittest.TestCase): if sqtt_events and kern_events: cls.examples[pkl_path.stem] = (sqtt_events, kern_events, cls.target) + @needs_rocprof def test_rocprof_inst_traces_match(self): for name, (events, kern_events, target) in self.examples.items(): if "sync" in name and self.target.startswith("gfx12"): @@ -94,7 +112,7 @@ class TestSQTTMapBase(unittest.TestCase): elif "WAVE" in e.device: # sopk/immediates don't get ALU/MEM EXEC if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE", "BARRIER", "BARRIER_SIGNAL", - "WAVEEND", "WAVERDY"} and not e.name.display_name.startswith("OTHER_"): insts += 1 + "WAVEEND", "WAVEEND_RDNA4", "WAVERDY"} and not e.name.display_name.startswith("OTHER_"): insts += 1 else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}") self.assertEqual(execs, insts) @@ -111,15 +129,18 @@ class TestSQTTMapBase(unittest.TestCase): def test_sqtt_cli(self): for pkl_path in sorted((EXAMPLES_DIR/self.target).glob("*.pkl")): - out = run_cli("--profile", "--profile-path", str(pkl_path)) + out = run_cli("--profile-path", str(pkl_path), "--ls") sqtt_traces = [l.strip() for l in out.split("\n") if "SQTT" in l] for name in sqtt_traces: - out = run_cli("--profile", "--profile-path", str(pkl_path), "-s", ansistrip(name)) + out = run_cli("--profile-path", str(pkl_path), "-s", ansistrip(name)) lines = out.split("\n") self.assertIn("Clk", lines[0]) for r in lines[2:]: parts = r.split() self.assertTrue(parts[0].isdigit(), f"expected clock timestamp, got {parts[0]}") + with Context(DEBUG=2): + kernels = run_cli("--profile-path", str(pkl_path), "-s", "AMD").split("\n") + self.assertEqual(len(kernels), len(self.examples[pkl_path.stem][1])) class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100" @@ -127,14 +148,25 @@ class TestSQTTMapRDNA4(TestSQTTMapBase): target = "gfx1200" @unittest.expectedFailure - def test_rdna4_wmma(self): + def test_pipes(self): events, kernels, target = self.examples["profile_handwritten_run_0"] - row_ends = {} - for e in sqtt_timeline(events[0].blob, list(kernels.values())[0].lib, target): - if type(e).__name__ != "ProfileRangeEvent" or e.device != "ALUEXEC:0 WMMA": continue - if (et:=row_ends.get(e.device)) is not None and e.st < et: - raise RuntimeError(f"WMMA exec overlaps in {e.device}: {e.st} {et}.") - row_ends[e.device] = e.en + lib = list(kernels.values())[0].lib + dispatch_st:dict[str, int] = {} + row_ends:dict[str, int] = {} + row_counts:dict[str, int] = {} + for e in sqtt_timeline(events[1].blob, lib, target): + if type(e).__name__ != "ProfileRangeEvent": continue + info = e.name.ret or "" + if e.device.startswith("WAVE"): + idx = row_counts.get(e.device, 0) + dispatch_st[f"{e.device}-{idx}"] = int(e.st) + row_counts[e.device] = idx + 1 + elif info.startswith("LINK:"): + delay = int(e.st) - dispatch_st[info[len("LINK:"):]] + self.assertGreaterEqual(delay, 1, f"EXEC {e.device} starts before DISPATCH: delay={delay}") + if (prev_en:=row_ends.get(e.device)) is not None: + self.assertGreaterEqual(e.st, prev_en, f"EXEC overlap in {e.device}: {e.st} < prev end {prev_en}") + row_ends[e.device] = int(e.en) class TestSQTTMapCDNA(TestSQTTMapBase): target = "gfx950" diff --git a/test/backend/test_arange.py b/test/backend/test_arange.py index d789e243d8..b6c1af0bbd 100644 --- a/test/backend/test_arange.py +++ b/test/backend/test_arange.py @@ -2,23 +2,18 @@ import unittest import numpy as np from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable from tinygrad.helpers import Context, getenv, DEV -from tinygrad.engine.realize import run_linear -from tinygrad.schedule import linear_to_schedule -from tinygrad.engine.realize import CompiledRunner, get_program -from tinygrad.schedule import ExecItem -from tinygrad.renderer import Estimates +from tinygrad.engine.realize import run_linear, estimate_uop from tinygrad.renderer.ptx import PTXRenderer from test.helpers import needs_second_gpu class TestArange(unittest.TestCase): def _get_flops(self, tensor, desired): GlobalCounters.reset() - sched = tensor.schedule() - self.assertEqual(len(sched), 1) - p = get_program(sched[-1].ast, renderer=Device[Device.DEFAULT].renderer) - ExecItem(sched[-1].ast, [tensor.uop.buffer], prg=CompiledRunner(p)).run() + linear = tensor.schedule_linear() + self.assertEqual(len(linear.src), 1) + run_linear(linear) np.testing.assert_equal(tensor.numpy(), desired) - return p.estimates.ops + return estimate_uop(linear.src[-1]).ops def test_arange_complexity(self): self.assertEqual(self._get_flops(Tensor.arange(256), np.arange(256)), 0) @@ -41,9 +36,8 @@ class TestArange(unittest.TestCase): def test_tri_complexity(self): with Context(NOOPT=1): t = Tensor.ones(256, 256).contiguous().realize() - sched = t.triu().schedule() - p = get_program(sched[-1].ast, renderer=Device[Device.DEFAULT].renderer) - self.assertLessEqual(Estimates.from_uops(p.uops).ops, 4 * 256 * 256) + linear = t.triu().schedule_linear() + self.assertLessEqual(estimate_uop(linear.src[-1]).ops, 4 * 256 * 256) DSET, DDIM = 2048, 32 @@ -56,7 +50,7 @@ class TestIndexing(unittest.TestCase): GlobalCounters.reset() out = ((Tensor.arange(1,16385)-1)*needle).sum() linear, var_vals = out.linear_with_vars() - self.assertEqual(len(linear_to_schedule(linear)), 1) + self.assertEqual(len(linear.src), 1) run_linear(linear, var_vals) self.assertEqual(out.item(), 1337) @@ -73,7 +67,7 @@ class TestIndexing(unittest.TestCase): full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1)) X = full.sum(axis=(2,3)) linear, var_vals = X.linear_with_vars() - self.assertEqual(len(linear_to_schedule(linear)), 1) + self.assertEqual(len(linear.src), 1) run_linear(linear, var_vals) assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}" np.testing.assert_allclose(real_index, X.numpy()) @@ -99,7 +93,7 @@ class TestIndexing(unittest.TestCase): X = dataset[idxs] assert X.shape == (4,DDIM) linear, var_vals = X.linear_with_vars() - self.assertEqual(len(linear_to_schedule(linear)), 1) + self.assertEqual(len(linear.src), 1) run_linear(linear, var_vals) assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}" np.testing.assert_allclose(real_index, X.numpy()) @@ -114,7 +108,7 @@ class TestIndexing(unittest.TestCase): X = dataset[idxs] assert X.shape == (4,DDIM) linear, var_vals = X.linear_with_vars() - self.assertEqual(len(linear_to_schedule(linear)), 1) + self.assertEqual(len(linear.src), 1) run_linear(linear, var_vals) assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops} != {4*DSET}" np.testing.assert_allclose(real_index, X.numpy()) @@ -235,10 +229,9 @@ class TestIndexing(unittest.TestCase): xq = xq.reshape(bs, seqlen, n_heads, head_dim) xq_rope, _ = apply_rotary_emb(xq, xq, freqs_cis) xq_rope.sum().backward() - sched = wq.grad.schedule() - assert len(sched) == 1, f"expected one kernel for backward, got: {len(sched)}" - prg = sched[0].lower().prg.p - bwd_ops = prg.estimates.ops + linear = wq.grad.schedule_linear() + assert len(linear.src) == 1, f"expected one kernel for backward, got: {len(linear.src)}" + bwd_ops = estimate_uop(linear.src[0]).ops # bfloat16 on non CDNA4 has ~10x ops overhead because of the software emulation if dtype == dtypes.bfloat16 and not Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"): ops_scale = 10 else: ops_scale = 1 diff --git a/test/backend/test_asm_gemm.py b/test/backend/test_asm_gemm.py index 7d006a4122..518aaeeb3e 100644 --- a/test/backend/test_asm_gemm.py +++ b/test/backend/test_asm_gemm.py @@ -46,9 +46,9 @@ def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=N np.testing.assert_allclose(tst.numpy(), ref.numpy(), atol=atol, rtol=rtol) np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), atol=grad_atol, rtol=grad_rtol) np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), atol=grad_atol, rtol=grad_rtol) - assert tst.allclose(ref, atol=atol, rtol=rtol), "forward mismatch" - assert a.grad.allclose(a_ref.grad, atol=grad_atol, rtol=grad_rtol), "grad_a mismatch" - assert b.grad.allclose(b_ref.grad, atol=grad_atol, rtol=grad_rtol), "grad_b mismatch" + assert tst.allclose(ref, atol=atol, rtol=rtol).item(), "forward mismatch" + assert a.grad.allclose(a_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_a mismatch" + assert b.grad.allclose(b_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_b mismatch" def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None: run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus) diff --git a/test/backend/test_const_folding.py b/test/backend/test_const_folding.py index 2e2095f75a..1d62e6813e 100644 --- a/test/backend/test_const_folding.py +++ b/test/backend/test_const_folding.py @@ -8,8 +8,8 @@ from test.helpers import not_support_multi_device def _check_ast_count(desired_count:int, t:Tensor): # NOTE: this has side effect because everything can be scheduled only once - schedule = t.schedule() - asts = [s for s in schedule if s.ast.op is Ops.SINK] + schedule = t.schedule_linear() + asts = [s for s in schedule.src if s.src[0].op is Ops.SINK] len(asts) # NOT SUPPORTED ANYMORE #assert len(asts) == desired_count, f"{len(asts)} != {desired_count}" @@ -28,8 +28,8 @@ class TestMovedConstFolding(unittest.TestCase): _check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),))) def test_copy_padded_const(self): - schedule = Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").schedule() - assert not any(si.ast.op is Ops.COPY for si in schedule), "const copy should be folded" + schedule = Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").schedule_linear() + assert not any(si.src[0].op is Ops.COPY for si in schedule.src), "const copy should be folded" np.testing.assert_equal(Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").numpy(), [0, 1, 1, 1, 1, 0]) def test_cast_padded(self): diff --git a/test/backend/test_custom_kernel.py b/test/backend/test_custom_kernel.py index ba4d834200..262522f3b7 100644 --- a/test/backend/test_custom_kernel.py +++ b/test/backend/test_custom_kernel.py @@ -189,7 +189,7 @@ class TestCustomKernel(unittest.TestCase): A = Tensor.randn(16, 16).contiguous() B = Tensor.empty(16) B = Tensor.custom_kernel(B, A, fxn=slice_sum_kernel)[0] - self.assertTrue(B.allclose(A.sum(1))) + self.assertTrue(B.allclose(A.sum(1)).item()) def test_gemm(self): N = 16 @@ -273,12 +273,12 @@ class TestCustomKernel(unittest.TestCase): C, D, _, _ = Tensor.custom_kernel(C, D, A2, B2, fxn=custom_elementwise_addmul_kernel) # depends on A2 AND B2 E = (A2 * 3).contiguous() # kernel 2: depends only on A2 result = (C + D + E).sum() # kernel 3: custom_addmul, then kernel 4: sum - schedule = result.schedule() + schedule = result.schedule_linear().src # Find the custom_addmul kernel position custom_idx = next((i for i, item in enumerate(schedule) - if hasattr(item.ast, "arg") and hasattr(item.ast.arg, "name") - and "custom_addmul" in item.ast.arg.name), None) + if hasattr(item.src[0], "arg") and hasattr(item.src[0].arg, "name") + and "custom_addmul" in item.src[0].arg.name), None) self.assertIsNotNone(custom_idx, "custom_addmul kernel not found in schedule") self.assertEqual(custom_idx, 3, f"custom_addmul should be at index 3, got {custom_idx}") diff --git a/test/backend/test_graph.py b/test/backend/test_graph.py index 031e327ed3..554c3db646 100644 --- a/test/backend/test_graph.py +++ b/test/backend/test_graph.py @@ -6,7 +6,7 @@ from tinygrad.tensor import Tensor from tinygrad.helpers import Context, from_mv from tinygrad.dtype import dtypes from tinygrad.engine.jit import MultiGraphRunner -from tinygrad.engine.realize import run_linear +from tinygrad.engine.realize import run_linear, compile_linear from tinygrad.uop.ops import UOp, Ops, buffers from test.helpers import needs_second_gpu @@ -24,7 +24,7 @@ def get_ast(device:str, num_inputs:int) -> UOp: fst = [Tensor.randn(BUF_SIZE, dtype=dtypes.int).realize() for _ in range(num_inputs)] s = fst[0] for i in range(1, num_inputs): s = s.bitwise_xor(fst[i]) - cached_asts[(device, num_inputs)] = s.schedule()[-1].ast + cached_asts[(device, num_inputs)] = s.schedule_linear().src[-1].src[0] return cached_asts[(device, num_inputs)] def make_buffer(device, size=BUF_SIZE, fill=False): @@ -44,7 +44,7 @@ def get_buf_uop(buf:Buffer, cache:dict[Buffer,UOp]) -> UOp: return cache[buf] def make_graph(graph_cls, calls:list[UOp]): - linear = UOp(Ops.LINEAR, src=tuple(calls)) + linear = compile_linear(UOp(Ops.LINEAR, src=tuple(calls))) cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(linear,), arg="graph") return graph_cls(cf, []) diff --git a/test/backend/test_jit.py b/test/backend/test_jit.py index 63f6e1b6c7..44bb2193c7 100644 --- a/test/backend/test_jit.py +++ b/test/backend/test_jit.py @@ -92,6 +92,20 @@ class TestJit(unittest.TestCase): np.testing.assert_allclose(e.numpy(), a.numpy()*b.numpy(), atol=1e-4, rtol=1e-5) assert_jit_cache_len(f, 3) + def test_global_counters_jit(self): + @TinyJit + def f(a, b): + c = (a + b).realize() + d = (c * 2).realize() + return (d - a).realize() + a, b = Tensor.randn(64, 64).realize(), Tensor.randn(64, 64).realize() + for _ in range(4): + GlobalCounters.reset() + f(a, b) + Device[a.device].synchronize() + self.assertGreater(GlobalCounters.global_mem, 0) + self.assertGreater(GlobalCounters.global_ops, 0) + def test_nothing_jitted(self): @TinyJit def add(a, b): return None diff --git a/test/backend/test_linearizer.py b/test/backend/test_linearizer.py index 85629d800d..9529460e75 100644 --- a/test/backend/test_linearizer.py +++ b/test/backend/test_linearizer.py @@ -1,13 +1,12 @@ import numpy as np import unittest -from dataclasses import replace from tinygrad.codegen.opt import Opt, OptOps from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType from tinygrad.device import Device, Buffer, is_dtype_supported from tinygrad.tensor import Tensor, _to_np_dtype -from tinygrad.engine.realize import run_linear, CompiledRunner, get_program -from tinygrad.schedule import linear_to_schedule +from tinygrad.engine.realize import run_linear, CompiledRunner +from tinygrad.codegen import to_program from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, DEV from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace from tinygrad.renderer.ptx import PTXRenderer @@ -26,9 +25,9 @@ class TestLinearizer(unittest.TestCase): a, b = Tensor.randn(4).realize(), Tensor.randn(4).realize() np_a, np_b = a.numpy(), b.numpy() c = ((a.shrink(((0, 2),)) - a.shrink(((2, 4),))) - (b.shrink(((0, 2),)) - b.shrink(((2, 4),)))) - sched = c.schedule() - for si in sched: si.run() - rawbufs = sched[-1].bufs + linear = c.schedule_linear() + run_linear(linear) + rawbufs = [s.buffer for s in linear.src[-1].src[1:] if s.op is not Ops.BIND] assert len(rawbufs) == 3 and set(rawbufs[1:]) == {a.uop.base.realized, b.uop.base.realized} np_c = (np_a[:2] - np_a[2:]) - (np_b[:2] - np_b[2:]) np.testing.assert_allclose(np_c, c.numpy(), atol=1e-4, rtol=1e-4) @@ -46,7 +45,7 @@ class TestLinearizer(unittest.TestCase): tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize() out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2 ast = helper_linearizer_opt(out) - uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src) self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1) @unittest.expectedFailure @@ -54,7 +53,7 @@ class TestLinearizer(unittest.TestCase): tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize() out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2 ast = helper_linearizer_opt(out) - uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src) self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0) @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx") @@ -64,7 +63,7 @@ class TestLinearizer(unittest.TestCase): b = Tensor.empty(16) out = img.conv2d(w, b) ast = helper_linearizer_opt(out) - uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src) # slice at the last loop end uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1] # only valid test if outermost range is the reduce @@ -85,7 +84,7 @@ class TestLinearizer(unittest.TestCase): a = Tensor.randn(2, ).realize() out = a.reshape(2, 1).expand(2, 3).sum() ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)).sum()]) - uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src) ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE] assert len(ranges) == 1 # NOTE: it collapses now @@ -93,7 +92,7 @@ class TestLinearizer(unittest.TestCase): a = Tensor.randn(2, ).realize() out = a.reshape(2, 1).expand(2, 3).expand(2, 2, 3).sum() ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)), (2, 2, 3)).sum()]) - uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src) ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE] assert len(ranges) == 1 # NOTE: it collapses now @@ -101,7 +100,7 @@ class TestLinearizer(unittest.TestCase): a = Tensor([2, 2]).realize() out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum() ast = helper_linearizer_opt(out, wanna_output=[24]) - uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src) ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE] # RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]]) @@ -114,7 +113,7 @@ class TestLinearizer(unittest.TestCase): b = Tensor.randn(1, 1).realize() out = (a + b[0]).sum() + b[0] ast = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()]) - uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src) ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE] # LOAD -> RANGE -> LOAD -> STORE assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1 @@ -124,7 +123,7 @@ class TestLinearizer(unittest.TestCase): b = Tensor.randn(1, 1).realize() out = (a.reshape(2, 1).expand(2, 3) + b[0]).sum() + b[0] ast = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)) + b.numpy()[0]).sum() + b.numpy()]) - uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src) ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE] assert len(ranges) == 1 # NOTE: it collapses now @@ -135,7 +134,8 @@ class TestLinearizer(unittest.TestCase): # these are of size 3 to avoid float4 coalesce r = a[:-1] + a[1:] - uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), + renderer=Device[Device.DEFAULT].renderer).src[2].src) num_loads = len([uop for uop in uops if uop.op is Ops.LOAD]) assert num_loads <= 4, "more load uops than needed" assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?" @@ -147,7 +147,8 @@ class TestLinearizer(unittest.TestCase): a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize() r = a.expand([2]) + b.expand([2]) - uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), + renderer=Device[Device.DEFAULT].renderer).src[2].src) num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU]) assert num_ops <= 1, "more alu uops than needed" @@ -156,8 +157,8 @@ class TestLinearizer(unittest.TestCase): x, w = Tensor.randn((1,1,3)).realize(), Tensor.randn((1,1,2)).realize() r = Tensor.conv2d(x,w,padding=1).relu() - uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), - renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], + [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).src[2].src) accs = [u for u in uops if u.op is Ops.DEFINE_REG] stores = [u for u in uops if u.op is Ops.STORE] assert len(accs) == 0 # it's removed now @@ -169,8 +170,9 @@ class TestLinearizer(unittest.TestCase): @unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU") def test_upcast_with_locals_cpu(self): out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous() - prg = get_program(replace_opts(out.schedule()[-1].ast, [Opt(OptOps.LOCAL, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops - self.assertEqual(len(prg.src.split("for")), 5) + prg = to_program(replace_opts(out.schedule_linear().src[-1].src[0], [Opt(OptOps.LOCAL, axis=0, arg=4)]), + renderer=Device[Device.DEFAULT].renderer) + self.assertEqual(len(prg.src[3].arg.split("for")), 5) @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals") @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared") @@ -180,9 +182,9 @@ class TestLinearizer(unittest.TestCase): x, y = Tensor.rand(1,128), Tensor.rand(128, 128) r = (x@y).relu() opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)] - program = get_program(replace_opts(r.schedule()[-1].ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer) + program = to_program(replace_opts(r.schedule_linear().src[-1].src[0], opts_to_apply), renderer=Device[Device.DEFAULT].renderer) - stores = [u for u in program.uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG] + stores = [u for u in tuple(program.src[2].src) if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG] # the first store is to lds and can be upcasted assert stores[0].src[1].dtype == dtypes.float.vec(4) @@ -194,7 +196,8 @@ class TestLinearizer(unittest.TestCase): def test_zero_fold(self): a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize() r = Tensor.stack(a, b) - uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), + renderer=Device[Device.DEFAULT].renderer).src[2].src) num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU]) assert num_ops == 0, "more alu uops than needed" @@ -203,16 +206,16 @@ class TestLinearizer(unittest.TestCase): (dtypes.bool, dtypes.int), (dtypes.int16, dtypes.int), (dtypes.float16, dtypes.float), (dtypes.bfloat16, dtypes.float)): if is_dtype_supported(tensor_dtype) and is_dtype_supported(acc_dtype): a = Tensor([1, 2, 3], dtype=tensor_dtype).sum() - realized_ast = a.schedule()[-1].ast - program = get_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer) - local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG] + realized_ast = a.schedule_linear().src[-1].src[0] + program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer) + local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.DEFINE_REG] assert local[0].dtype.base == acc_dtype def test_arg_acc_dtype(self): def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType): - realized_ast = c.schedule()[-1].ast - program = get_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer) - local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG] + realized_ast = c.schedule_linear().src[-1].src[0] + program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer) + local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.DEFINE_REG] self.assertEqual(local[0].dtype.base, expected_dtype) tests = ( @@ -239,7 +242,7 @@ class TestLinearizer(unittest.TestCase): opt = [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)] ast = helper_linearizer_opt(r, [opt]) # the uops graph is DEFINE_REG -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE - uops = get_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src) begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1] end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0] for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype) @@ -259,7 +262,7 @@ class TestLinearizer(unittest.TestCase): # shrink so that the dims do not collapse t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6))) ast = helper_linearizer_opt(t+1) - uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src) idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL]) idxs = sorted(idxs, key=lambda uop: uop.arg) assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0] @@ -268,10 +271,10 @@ class TestLinearizer(unittest.TestCase): def test_sum_collapse(self): t = Tensor([2]).reshape(1, 1).expand(256, 256).sum() - sched = [si for si in t.schedule() if si.ast.op is Ops.SINK] + sched = [si for si in t.schedule_linear().src if si.src[0].op is Ops.SINK] # sum_collapse is a full collapse now assert len(sched) == 1 - assert not any(u.op is Ops.REDUCE_AXIS for u in sched[0].ast.toposort()), "found reduce in sum collapse" + assert not any(u.op is Ops.REDUCE_AXIS for u in sched[0].src[0].toposort()), "found reduce in sum collapse" #lin = Kernel(sched[0].ast) #assert not any(u.op is Ops.RANGE for u in lin.linearize().uops), "found loop in sum collapse" @@ -288,17 +291,16 @@ class TestLinearizer(unittest.TestCase): b = a.shrink(((1, 2), None)).pad(((1, 2), None)) a.assign(b.where(2, a)) linear, var_vals = a.linear_with_vars() - sched_copy = linear_to_schedule(linear) - assert len(sched_copy) == 1 + assert len(linear.src) == 1 run_linear(linear, var_vals) np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.]) - program = get_program(replace_opts(sched_copy[-1].ast, []), renderer=Device[Device.DEFAULT].renderer) - assert not any(u.op == Ops.WHERE for u in program.uops), "found where where where should be folded" + program = to_program(replace_opts(linear.src[-1].src[0], []), renderer=Device[Device.DEFAULT].renderer) + assert not any(u.op == Ops.WHERE for u in tuple(program.src[2].src)), "found where where where should be folded" def test_phi_simplification(self): def helper(t, max_ops=0): ast = helper_linearizer_opt(t) - uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src) # ignore kernel optimized IF statements for now if if_op:=next((u for u in uops if u.op is Ops.IF), None): uops = uops[:uops.index(if_op)] @@ -330,7 +332,7 @@ class TestLinearizer(unittest.TestCase): out = x.matmul(y) with Context(TC=0): ast = helper_linearizer_opt(out) - uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src) # check that the float4 cast collapses store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG] for val in store_vals: @@ -341,7 +343,7 @@ class TestLinearizer(unittest.TestCase): x = Tensor.randn((4,3,6,6)).realize() out = x.flip((0,1)).contiguous() ast = helper_linearizer_opt(out) - store_val = [u.src[1] for u in get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops if u.op is Ops.STORE][0] + store_val = [u.src[1] for u in tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src) if u.op is Ops.STORE][0] assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.STACK @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals") @@ -354,7 +356,7 @@ class TestLinearizer(unittest.TestCase): Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces ast = helper_linearizer_opt(out, opts=[opt]) def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src]) - uops = get_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src) local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))] global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.PARAM for x in get_recursive(u.src[0]))] barrier = [u for u in uops if u.op is Ops.BARRIER] @@ -375,7 +377,7 @@ class TestLinearizer(unittest.TestCase): x, y = Tensor.rand(1,128), Tensor.rand(128, 128) r = (x@y).relu() ast = helper_linearizer_opt(r) - uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops + uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src) stores = [u for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG] # the float4 value stores directly in lds and we skip upcast @@ -390,15 +392,17 @@ class TestLinearizer(unittest.TestCase): def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]: if isinstance(r, Tensor): r = [r] linear, var_vals = Tensor.linear_with_vars(*r) - s = linear_to_schedule(linear) run_linear(UOp(Ops.LINEAR, src=linear.src[:-1]), var_vals) # run all kernels except the last one - assert s[-1].ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {s[-1]}" - # now all input buffers in s[-1] should be realized + last_call = linear.src[-1] + ast = last_call.src[0] + assert ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {last_call}" + last_bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND] + # now all input buffers in last_call should be realized # create fresh buffers for the outputs - bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)] + bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(ast.src) else x for i,x in enumerate(last_bufs)] # ensure buffers are allocated for b in bufs: b.ensure_allocated() - return s[-1].ast, bufs + return ast, bufs def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs): assert isinstance(ast, UOp), "ast must be UOp" @@ -425,7 +429,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[] def get_prg(opts): ast = realized_ast if opts is None else replace_opts(realized_ast, list(opts)) - return CompiledRunner(replace(get_program(ast, renderer=Device[Device.DEFAULT].renderer), device=device)) + return CompiledRunner(to_program(ast, renderer=Device[Device.DEFAULT].renderer), device) def check_opt(opts): prg = get_prg(opts=opts) diff --git a/test/backend/test_linearizer_dumb.py b/test/backend/test_linearizer_dumb.py index b56fa7bbb6..da5160d4e6 100644 --- a/test/backend/test_linearizer_dumb.py +++ b/test/backend/test_linearizer_dumb.py @@ -6,7 +6,7 @@ import unittest from tinygrad import Device, dtypes from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo from tinygrad.codegen.opt.search import Opt, OptOps -from tinygrad.engine.realize import get_program +from tinygrad.codegen import to_program class TestLinearizerFailure(unittest.TestCase): @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL") @@ -25,7 +25,7 @@ class TestLinearizerFailure(unittest.TestCase): c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.weakint, 6000))+c6)+((c7*UOp.const(dtypes.weakint, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.weakint, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD) c12 = c0.index((((c1*UOp.const(dtypes.weakint, 7840))+(c2*UOp.const(dtypes.weakint, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3) ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None)) - _ = get_program(ast, Device["METAL"].renderer) + _ = to_program(ast, Device["METAL"].renderer) if __name__ == '__main__': unittest.main() diff --git a/test/backend/test_multitensor.py b/test/backend/test_multitensor.py index 38471b3842..9e0aa718d6 100644 --- a/test/backend/test_multitensor.py +++ b/test/backend/test_multitensor.py @@ -4,8 +4,7 @@ from tinygrad.device import is_dtype_supported from tinygrad.uop.ops import Ops, UOp from tinygrad.helpers import getenv, prod, Context from tinygrad.nn.state import get_parameters, get_state_dict -from tinygrad.engine.realize import CompiledRunner, run_linear -from tinygrad.schedule import linear_to_schedule +from tinygrad.engine.realize import run_linear, compile_linear import numpy as np from hypothesis import given, strategies as strat, settings from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph @@ -128,12 +127,9 @@ class TestMultiTensor(unittest.TestCase): X = Tensor.ones(256).contiguous().realize() X.shard_(devices_2, 0) out = (X + X) - sched = out.schedule() - names = [] - for si in sched: - si.lower() - if isinstance(si.prg, CompiledRunner): names.append(si.prg.p.name) - si.run() + linear = compile_linear(out.schedule_linear()) + names = [call.src[0].src[0].arg.name for call in linear.src if call.src[0].op is Ops.PROGRAM] + run_linear(linear) self.assertEqual(len(set(names)), 1, "function was relinearized") def test_shard_same_device(self): @@ -194,9 +190,9 @@ class TestMultiTensor(unittest.TestCase): for i in range(2): xt = X[i*2:i*2+2].contiguous() linear, var_vals = xt.linear_with_vars() - #kernels = [s for s in linear_to_schedule(linear) if s.ast.op is Ops.SINK] + #kernels = [call for call in linear.src if call.src[0].op is Ops.SINK] #self.assertEqual(len(kernels), 1) - #self.assertEqual(kernels[0].bufs[0].device, devices_2[i]) + #self.assertEqual(kernels[0].src[1].buffer.device, devices_2[i]) run_linear(linear, var_vals) np.testing.assert_equal(xt.numpy(), X_np[i*2:i*2+2]) @@ -555,6 +551,21 @@ class TestMultiTensor(unittest.TestCase): np.testing.assert_allclose(c.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5) assert jf.captured is not None + def test_multi_tensor_jit_graph_assign_updates_each_shard(self): + @TinyJit + def jf(out: Tensor) -> Tensor: + tmp = (Tensor.arange(4, dtype=dtypes.float).shard(devices_2, 0) + 1).contiguous().realize() + out.assign((tmp + 1).contiguous()).realize() + return out + + out = Tensor.full((4,), -1.0).shard(devices_2, 0).contiguous().realize() + expected = np.arange(4, dtype=np.float32) + 2 + for _ in range(5): + out.assign(Tensor.full((4,), -1.0).shard(devices_2, 0).contiguous()).realize() + jf(out) + np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-5) + assert jf.captured is not None + def test_multi_tensor_jit_body(self): @TinyJit def jf() -> Tensor: @@ -640,18 +651,15 @@ class TestMultiTensor(unittest.TestCase): for p in get_parameters(bn): p.shard_(devices_4).realize() out = bn(t) - scheds = [sched for sched in out.schedule() if sched.bufs[0].device in devices_4 and sched.ast.op is not Ops.COPY] - assert set(sched.bufs[0].device for sched in scheds) == set(devices_4), "should have ast on each shard device" - asts = [sched.ast for sched in scheds] - self.assertEqual(len(asts), 4) - # ast are the same on devices - self.assertEqual(len(set(asts)), 1) + scheds = [call for call in out.schedule_linear().src if call.src[0].op is not Ops.COPY and set(call.device) <= set(devices_4)] + self.assertEqual(set(scheds[0].device), set(devices_4), "should have ast on each shard device") + self.assertEqual(len(set(s.src[0] for s in scheds)), 1) def test_flip(self): rng = Tensor.rand((10, 10, 10)) t0 = rng.shard(devices_2, axis=1) out = t0.flip(0) + 1 - self.assertTrue((rng.flip(0)+1).allclose(out.to(rng.device))) + self.assertTrue((rng.flip(0)+1).allclose(out.to(rng.device)).item()) @unittest.skip("flaky") def test_reshape_on_axis(self): @@ -685,7 +693,7 @@ class TestMultiTensor(unittest.TestCase): # test no left join with self.assertRaises((AssertionError, ValueError)): - t0.reshape((26*15,7)).contiguous().schedule() + t0.reshape((26*15,7)).contiguous().schedule_linear() # it doesn't work like this anymore # NOTE: this never failed in assign_multi, it failed tensor spec because MULTI was never pushed in the graph @@ -696,7 +704,7 @@ class TestMultiTensor(unittest.TestCase): with self.assertRaises(RuntimeError): # don't allow assigns that change axes t_none.assign(t_zero) - t_none.schedule() + t_none.schedule_linear() def test_init_rand_with_multiple_devices_fail(self): # init rand with multi device is not allowed @@ -794,7 +802,7 @@ class TestMultiTensor(unittest.TestCase): t = Tensor.ones(16, 16, dtype=dtypes.int).shard(devices_2, axis=0) out = Tensor.full_like(t, 2)[:, :8] linear, var_vals = out.linear_with_vars() - self.assertEqual(len(linear_to_schedule(linear)), 0) + self.assertEqual(len(linear.src), 0) run_linear(linear, var_vals) self.assertEqual(out.tolist(), [[2]*8]*16) @@ -837,7 +845,7 @@ class TestMultiTensor(unittest.TestCase): a = Tensor.arange(3).realize() zeros = Tensor.zeros(3).realize() b = a.to(devices_2)*zeros.to(devices_2) - sched = b.schedule() + sched = b.schedule_linear().src self.assertEqual(len(sched), 0) self.assertListEqual(b.tolist(), [0, 0, 0]) @@ -848,7 +856,7 @@ class TestHandleData(unittest.TestCase): device = (d0, d1, d2, d3) t = Tensor([1, 2, 3, 4]).shard(device).realize() not_covered = t.to(d5) - sched = not_covered.schedule() + sched = not_covered.schedule_linear().src assert len(sched) == 1 # setup again because create_schedule has side effect t = Tensor([1, 2, 3, 4]).shard(device).realize() @@ -858,7 +866,7 @@ class TestHandleData(unittest.TestCase): for d in device: t = Tensor([1, 2, 3, 4]).shard(device).realize() covered = t.to(d) - sched = covered.schedule() + sched = covered.schedule_linear().src # TODO: this isn't optimized out anymore #assert len(sched) == 0 # setup again because create_schedule has side effect @@ -879,18 +887,18 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase): with self.assertRaises(AssertionError): # sharded axis shrink on non-device boundry is not allowed a = t.shrink(((0, 3), (0, 8))).contiguous() - a.schedule() + a.schedule_linear() a = t.shrink(((0, 2), (2, 4))) assert a.shape == (2, 2) ref = Tensor.arange(64).reshape(8, 8).shrink(((0, 2), (2, 4))) np.testing.assert_equal(a.numpy(), ref.numpy()) a = t.shrink(((0, 2), (0, 8))).contiguous() - a.schedule() + a.schedule_linear() assert a.shape == (2, 8) p = a.pad(((0, 6), (0, 0))).contiguous() - p.schedule() + p.schedule_linear() assert p.shape == (8, 8) @given(strat.sampled_from([dtypes.float, dtypes.int, dtypes.int64, dtypes.int16])) @@ -1105,9 +1113,9 @@ class TestBatchNorm(unittest.TestCase): p.to_(devices) synced_out = synced_bn(x) - synced_si = list(synced_out.schedule()) + synced_si = list(synced_out.schedule_linear().src) unsynced_out = unsynced_bn(x) - unsynced_si = list(unsynced_out.schedule()) + unsynced_si = list(unsynced_out.schedule_linear().src) # TODO: test synced / unsynced batchnorm cross device kernel and copies assert synced_si @@ -1144,12 +1152,12 @@ class TestMultiBufferView(unittest.TestCase): def setUp(self): pass def _check(self, a_ref:Tensor, a_multi:Tensor, view_fn): - """Apply view_fn to both, verify zero compiled kernels and matching values.""" b_ref = view_fn(a_ref) b_multi = view_fn(a_multi).contiguous() linear, var_vals = b_multi.linear_with_vars() - compiled = [si for si in linear_to_schedule(linear) if isinstance(si.prg, CompiledRunner)] - self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}") + if all(hasattr(Device[d].allocator, "_offset") for d in b_multi.device): + compiled = [call for call in linear.src if call.src[0].op is Ops.SINK] + self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}") run_linear(linear, var_vals) np.testing.assert_equal(b_multi.numpy(), b_ref.numpy()) @@ -1177,11 +1185,13 @@ class TestMultiBufferView(unittest.TestCase): def test_4_devices(self): ref = Tensor.arange(8*12).reshape(8, 12).contiguous().realize() a = Tensor.arange(8*12).reshape(8, 12).contiguous().shard(devices_4, axis=1).realize() - linear, var_vals = a[5].contiguous().linear_with_vars() - compiled = [si for si in linear_to_schedule(linear) if isinstance(si.prg, CompiledRunner)] - self.assertEqual(len(compiled), 0) + out = a[5].contiguous() + linear, var_vals = out.linear_with_vars() + if all(hasattr(Device[d].allocator, "_offset") for d in out.device): + compiled = [call for call in linear.src if call.src[0].op is Ops.SINK] + self.assertEqual(len(compiled), 0) run_linear(linear, var_vals) - np.testing.assert_equal(a[5].contiguous().numpy(), ref[5].numpy()) + np.testing.assert_equal(out.numpy(), ref[5].numpy()) @unittest.skipIf(not_support_multi_device(), "need multi") class TestMultiFromUnrenderable(unittest.TestCase): diff --git a/test/backend/test_nn.py b/test/backend/test_nn.py index 4d2b8cfbe5..263ee5760e 100644 --- a/test/backend/test_nn.py +++ b/test/backend/test_nn.py @@ -9,7 +9,6 @@ from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell from tinygrad.nn.state import load_state_dict from tinygrad.engine.realize import run_linear -from tinygrad.schedule import linear_to_schedule from test.helpers import not_support_multi_device, needs_second_gpu, slow @slow @@ -433,7 +432,7 @@ class TestNN(unittest.TestCase): [12, 19, 8, 1]]) result = layer(a) linear, var_vals = result.linear_with_vars() - self.assertEqual(len([item for item in linear_to_schedule(linear) if item.ast.op is Ops.SINK]), kcount, + self.assertEqual(len([call for call in linear.src if call.src[0].op is Ops.SINK]), kcount, "first run realizes weight and embedding") run_linear(linear, var_vals) @@ -442,7 +441,7 @@ class TestNN(unittest.TestCase): [7, 8, 9]]) result = layer(b) linear, var_vals = result.linear_with_vars() - self.assertEqual(1, len([item for item in linear_to_schedule(linear) if item.ast.op is Ops.SINK]), + self.assertEqual(1, len([call for call in linear.src if call.src[0].op is Ops.SINK]), "second run realizes embedding only") run_linear(linear, var_vals) print(f"Embedding used {GlobalCounters.global_ops} ops") diff --git a/test/backend/test_opt_gemm.py b/test/backend/test_opt_gemm.py index c810d16e29..244e3df889 100644 --- a/test/backend/test_opt_gemm.py +++ b/test/backend/test_opt_gemm.py @@ -1,10 +1,10 @@ import numpy as np import unittest -from tinygrad import Tensor, Device +from tinygrad import Tensor from tinygrad.helpers import get_single_element from tinygrad.codegen.opt import Opt, OptOps -from tinygrad.engine.realize import CompiledRunner, get_program -from tinygrad.schedule import ExecItem +from tinygrad.engine.realize import run_linear +from tinygrad.uop.ops import Ops, UOp from test.helpers import replace_opts class TestOptGemm(unittest.TestCase): @@ -19,10 +19,10 @@ class TestOptGemm(unittest.TestCase): def _test_gemm_unrolled_permute_l(self, opts=[]): t = self.a.T @ self.b.T # TODO: this should be a generic test helper - si = get_single_element(t.schedule()) - run = CompiledRunner(get_program(replace_opts(si.ast, opts), renderer=Device[Device.DEFAULT].renderer)) - ExecItem(si.ast, list(si.bufs), prg=run).run() - test = si.bufs[0].numpy().reshape(self.res.shape) + call = get_single_element(t.schedule_linear().src) + new_call = call.replace(src=(replace_opts(call.src[0], opts), *call.src[1:])) + run_linear(UOp(Ops.LINEAR, src=(new_call,))) + test = call.src[1].buffer.numpy().reshape(self.res.shape) np.testing.assert_allclose(self.res, test, atol=1e-4) def test_gemm_unrolled_permute_l_44(self): diff --git a/test/backend/test_pickle.py b/test/backend/test_pickle.py index 3cc617d272..d2db7ec37d 100644 --- a/test/backend/test_pickle.py +++ b/test/backend/test_pickle.py @@ -142,10 +142,10 @@ class TestPickle(unittest.TestCase): def test_pickle_schedule(self): a = Tensor([1,2]) out = a + 2 - sched = out.schedule() + sched = out.schedule_linear() pk = pickle.dumps(sched) sched_pk = pickle.loads(pk) - self.assertEqual(sched_pk[-1].ast, sched[-1].ast) + self.assertEqual(sched_pk.src[-1].src[0], sched.src[-1].src[0]) def test_pickle_renderer(self): from tinygrad.device import Device diff --git a/test/backend/test_profiler.py b/test/backend/test_profiler.py index bcb75574a9..3132432e40 100644 --- a/test/backend/test_profiler.py +++ b/test/backend/test_profiler.py @@ -44,9 +44,9 @@ class TestProfiler(unittest.TestCase): TestProfiler.a = Tensor([0.,1.], device=Device.DEFAULT).realize() TestProfiler.b = self.a + 1 - si = self.b.schedule()[-1] + si = self.b.schedule_linear().src[-1] - TestProfiler.runner = get_runner(TestProfiler.d0.device, si.ast) + TestProfiler.runner = get_runner(TestProfiler.d0.device, si.src[0]) TestProfiler.b.uop.buffer.allocate() def test_profile_kernel_run(self): diff --git a/test/backend/test_quantize_onnx.py b/test/backend/test_quantize_onnx.py index 5ef356b8b4..e4440fb019 100644 --- a/test/backend/test_quantize_onnx.py +++ b/test/backend/test_quantize_onnx.py @@ -1,12 +1,11 @@ # ruff: noqa: E501 import numpy as np import unittest -from dataclasses import replace -from tinygrad import Tensor, Context, Device, dtypes +from tinygrad import Tensor, Context, Device, dtypes, UOp from tinygrad.uop.ops import Ops from tinygrad.codegen.opt import Opt, OptOps -from tinygrad.engine.realize import CompiledRunner, get_program -from tinygrad.schedule import ExecItem +from tinygrad.engine.realize import run_linear +from tinygrad.codegen import to_program from test.helpers import replace_opts N = 512 @@ -39,13 +38,16 @@ def create_gemm_model(model_path:str, batch_size=N, in_size=N, out_size=N, bias= return model_path def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3): - si = out.schedule()[-1] - prg = get_program(replace_opts(si.ast, opts), renderer=Device[Device.DEFAULT].renderer) + linear = out.schedule_linear() + call = linear.src[-1] + prg = to_program(replace_opts(call.src[0], opts), renderer=Device[Device.DEFAULT].renderer) if replace_src is not None: - old_name = prg.src.split("__attribute__((noinline)) void ")[1].split("(")[0] - prg = replace(prg, src=replace_src + "/* DSP boilerplate */" + prg.src.split("/* DSP boilerplate */")[1].replace(old_name, "fxn")) - new_si = ExecItem(si.ast, [x.ensure_allocated() for x in si.bufs], si.metadata, prg=CompiledRunner(prg)) - for _ in range(run_count): new_si.run(wait=True) + old_name = prg.src[3].arg.split("__attribute__((noinline)) void ")[1].split("(")[0] + new_src = replace_src + "/* DSP boilerplate */" + prg.src[3].arg.split("/* DSP boilerplate */")[1].replace(old_name, "fxn") + # drop BINARY and replace SOURCE so run_linear recompiles + prg = prg.replace(src=prg.src[:3] + (UOp(Ops.SOURCE, arg=new_src),)) + linear = linear.replace(src=linear.src[:-1] + (call.replace(src=(prg, *call.src[1:])),)) + for _ in range(run_count): run_linear(linear) def get_quantized_model(sz): from onnxruntime.quantization import quantize_static, QuantFormat, QuantType, CalibrationDataReader @@ -75,9 +77,9 @@ class TestQuantizeOnnxCPU(unittest.TestCase): run_onnx = OnnxRunner(out_file) inp = Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32)) with Context(QUANTIZE=1): - sched = run_onnx({"input":inp})["output"].schedule() - sched[-2].lower() - daccs = [u for u in sched[-2].prg.p.uops if u.op is Ops.DEFINE_REG] + linear = run_onnx({"input":inp})["output"].schedule_linear() + prg = to_program(linear.src[-2].src[0], renderer=Device[Device.DEFAULT].renderer) + daccs = [u for u in tuple(prg.src[2].src) if u.op is Ops.DEFINE_REG] assert all(u.dtype.scalar() is dtypes.int for u in daccs) @unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP") diff --git a/test/backend/test_randomness.py b/test/backend/test_randomness.py index fc874b5aaf..119c3789ae 100644 --- a/test/backend/test_randomness.py +++ b/test/backend/test_randomness.py @@ -4,7 +4,9 @@ from functools import partial from tinygrad import nn, dtypes, Tensor, Device, TinyJit, Variable from tinygrad.helpers import getenv, CI, OSX from tinygrad.device import is_dtype_supported -from tinygrad.engine.realize import CompiledRunner +from tinygrad.codegen import to_program + +from tinygrad.uop.ops import Ops from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.nir import NIRRenderer from test.helpers import not_support_multi_device, needs_second_gpu @@ -117,12 +119,13 @@ class TestRandomness(unittest.TestCase): @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "PTX and NIR use pointer arithmetic") def test_threefry_doesnt_use_long(self): - sched = Tensor.rand(20).schedule() - for si in sched: - si.lower() - if isinstance(si.prg, CompiledRunner): - for u in si.prg.p.uops: - self.assertNotIn(u.dtype, {dtypes.long, dtypes.ulong}, msg=f"long found in {si.prg.p.name}") + linear = Tensor.rand(20).schedule_linear() + for call in linear.src: + ast = call.src[0] + if ast.op is Ops.SINK: + prg = to_program(ast, renderer=Device[Device.DEFAULT].renderer) + for u in tuple(prg.src[2].src): + self.assertNotIn(u.dtype, {dtypes.long, dtypes.ulong}, msg=f"long found in {prg.arg.name}") def test_threefry_against_reference_full(self): Tensor.manual_seed(1337) @@ -187,24 +190,24 @@ class TestRandomness(unittest.TestCase): Tensor.rand(1).realize() - s = Tensor.rand(20).schedule() - s2 = Tensor.rand(20).schedule() + s = Tensor.rand(20).schedule_linear().src + s2 = Tensor.rand(20).schedule_linear().src assert len(s) == len(s2), f"{len(s)} != {len(s2)}" for x,y in zip(s, s2): - if not (x.ast == y.ast): - print(f"{x.ast} != {y.ast}") + if not (x.src[0] == y.src[0]): + print(f"{x.src[0]} != {y.src[0]}") Tensor.rand(1, device=f"{Device.DEFAULT}:1").realize() - s3 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule() - s4 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule() + s3 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule_linear().src + s4 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule_linear().src assert len(s3) == len(s4), f"{len(s3)} != {len(s4)}" assert len(s2) == len(s4), f"{len(s)} != {len(s3)}" for x,y in zip(s3, s4): - if not (x.ast == y.ast): - print(f"{x.ast} != {y.ast}") + if not (x.src[0] == y.src[0]): + print(f"{x.src[0]} != {y.src[0]}") @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "need bfloat16 support") def test_rand_bfloat16(self): diff --git a/test/backend/test_renderer_failures.py b/test/backend/test_renderer_failures.py index 3f9fa2afaa..11abf6c603 100644 --- a/test/backend/test_renderer_failures.py +++ b/test/backend/test_renderer_failures.py @@ -3,7 +3,8 @@ import numpy as np from dataclasses import replace from tinygrad.device import Buffer, Device, is_dtype_supported from tinygrad.dtype import dtypes, ConstType -from tinygrad.engine.realize import CompiledRunner, get_program +from tinygrad.engine.realize import CompiledRunner +from tinygrad.codegen import to_program from tinygrad.helpers import prod from tinygrad.renderer.cstyle import CStyleLanguage from tinygrad.renderer.ptx import PTXRenderer @@ -12,15 +13,15 @@ from tinygrad.runtime.ops_python import PythonRenderer from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu from tinygrad.tensor import Tensor, _to_np_dtype -def _test_uop_result(inputs:list[Tensor], prg, local_size=None): +def _test_uop_result(inputs:list[Tensor], prg:UOp, local_size=None): for x in inputs: x.realize() - uops = prg.uops + uops = prg.src[2].src outbufs = [Buffer(Device.DEFAULT, sz:=(1 if local_size is None else prod(local_size)), (dtype:=u.src[1].dtype), \ initial_value=np.zeros(sz, dtype=_to_np_dtype(dtype)).data) for u in uops if u.op is Ops.STORE] inbufs = [x.uop.base.buffer for x in inputs] - prg = replace(prg, device=Device.DEFAULT) - if local_size is not None: prg = replace(prg, local_size=local_size) - ei = CompiledRunner(prg) + info = prg.arg + if local_size is not None: info = replace(info, local_size=tuple(local_size)) + ei = CompiledRunner(prg.replace(arg=info), Device.DEFAULT) ei.exec(outbufs+inbufs) return [np.frombuffer(x.as_memoryview(), _to_np_dtype(x.dtype)) for x in outbufs] @@ -33,7 +34,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp): alu = ld.alu(alu_op, *alu_src_uops) store = UOp.store(a.index(idx), alu) sink = UOp(Ops.SINK, dtypes.void, (store,), arg=KernelInfo()) - prg = get_program(sink, Device[Device.DEFAULT].renderer) + prg = to_program(sink, Device[Device.DEFAULT].renderer) return _test_uop_result([Tensor([input_val])], prg)[0] class TestRendererFailures(unittest.TestCase): @@ -43,7 +44,7 @@ class TestRendererFailures(unittest.TestCase): gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0) gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1))) sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo()) - prg = get_program(sink, Device[Device.DEFAULT].renderer) + prg = to_program(sink, Device[Device.DEFAULT].renderer) ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0] np.testing.assert_equal(ret, [0, 1, 1, 1]) @@ -54,7 +55,7 @@ class TestRendererFailures(unittest.TestCase): gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0) gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1))) sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo()) - prg = get_program(sink, Device[Device.DEFAULT].renderer) + prg = to_program(sink, Device[Device.DEFAULT].renderer) ret = _test_uop_result([], prg, local_size=[4, 2, 1])[0] np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1]) @@ -69,10 +70,9 @@ class TestCStyleFailures(unittest.TestCase): dtype = "bool" if op in (Ops.OR, Ops.XOR, Ops.AND) else None ret = Tensor.empty(1, dtype=dtype) for _ in range(5): ret = python_alu[op](ret, Tensor.empty(1, dtype=dtype)) - schedule = ret.schedule() - assert len(schedule) == 1 - schedule[0].lower() - src = schedule[0].prg.p.src + linear = ret.schedule_linear() + assert len(linear.src) == 1 + src = to_program(linear.src[0].src[0], Device[Device.DEFAULT].renderer).src[3].arg self.assertEqual("("*5 not in src, should_strip_paren) def test_repeat_add(self): self._test_src_strip_paren(Ops.ADD) @@ -102,7 +102,7 @@ class TestPTXFailures(unittest.TestCase): if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,)) gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0, if_uop), val)) sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo()) - prg = get_program(sink, Device[Device.DEFAULT].renderer) + prg = to_program(sink, Device[Device.DEFAULT].renderer) ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0] np.testing.assert_equal(ret, [0, 1, 1, 1]) diff --git a/test/backend/test_schedule.py b/test/backend/test_schedule.py index 9fc8ca1a06..22e5f19e29 100644 --- a/test/backend/test_schedule.py +++ b/test/backend/test_schedule.py @@ -12,8 +12,7 @@ from tinygrad.device import is_dtype_supported from tinygrad.dtype import DType from tinygrad.uop.ops import UOp, Ops, UPat from tinygrad.helpers import CI, DEBUG, OSX, GlobalCounters, Context, getenv, all_same, temp -from tinygrad.engine.realize import CompiledRunner, run_linear -from tinygrad.schedule import linear_to_schedule +from tinygrad.engine.realize import compile_linear, run_linear class KernelCountException(Exception): pass def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True): @@ -24,17 +23,17 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te else: assert isinstance(t, UOp), f"can't schedule {t}" linear, var_vals = Tensor(t).linear_with_vars() - # test lowering all the ExecItems - sched = linear_to_schedule(linear) - for si in sched: si.lower() - kernel_cnt = len([si for si in sched if isinstance(si.prg, CompiledRunner) or not filter_sink]) + kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1) + for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink) if kernel_cnt != allowed: print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}") if DEBUG >= 3: - for i,s in enumerate(sched): + for i,call in enumerate(linear.src): print("kernel", i+1) - print(s.ast) + print(call.src[0]) raise KernelCountException(f"{kernel_cnt} != {allowed}") + # test compiling the linear + compile_linear(linear) return linear, var_vals def _realize_weights(m): @@ -50,9 +49,8 @@ def _test_conv2d(allowed:int, dtype:DType=dtypes.float): ret = Tensor.conv2d(img, w).relu().mean().backward() dtypes.default_float = old_default_float linear, var_vals = Tensor.linear_with_vars(ret, img.grad, w.grad) - s = linear_to_schedule(linear) run_linear(linear, var_vals) - cnt = len([si for si in s if si.ast.op is Ops.SINK]) + cnt = len([call for call in linear.src if call.src[0].op is Ops.SINK]) assert cnt == allowed, f"expected {allowed} kernels, got {cnt}" if getenv("CHECK", 1): import torch @@ -74,7 +72,7 @@ class TestSchedule(unittest.TestCase): x = Tensor.arange(25).reshape(1,1,5,5).cast(dtypes.float32) t = x.avg_pool2d(padding=1) linear, var_vals = t.linear_with_vars() - self.assertEqual(len(linear_to_schedule(linear)), kcount) + self.assertEqual(len(linear.src), kcount) run_linear(linear, var_vals) import torch torch_out = torch.nn.functional.avg_pool2d(torch.arange(25).reshape(1,1,5,5).float(), kernel_size=(2,2), padding=1).numpy() @@ -789,7 +787,7 @@ class TestSchedule(unittest.TestCase): gc.collect() base = GlobalCounters.mem_used x = Tensor.ones(256).contiguous().realize() - (x+Tensor.ones(256).contiguous()).schedule() + (x+Tensor.ones(256).contiguous()).schedule_linear() gc.collect() self.assertEqual(GlobalCounters.mem_used-base, 1024) @@ -799,9 +797,8 @@ class TestSchedule(unittest.TestCase): def cnt(): x, y, z = Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float') a = (x @ y).relu() - sched = ((a @ z).relu() + a).schedule() - for si in sched: si.lower() - return len([si for si in sched if isinstance(si.prg, CompiledRunner)]) + linear = compile_linear(((a @ z).relu() + a).schedule_linear()) + return len([call for call in linear.src if call.src[0].op is Ops.PROGRAM]) with Context(IMAGE=1): self.assertEqual(cnt(), 5) @@ -816,9 +813,8 @@ class TestSchedule(unittest.TestCase): rb = (((((inp @ b1) + c1).relu() @ b2) + c2).relu() + inp).relu() b16, c16 = Tensor.empty((512, 16), dtype='float'), Tensor.empty((16,), dtype='float') b32, c32 = Tensor.empty((512, 32), dtype='float'), Tensor.empty((32,), dtype='float') - sched = Tensor.schedule((rb @ b16 + c16).relu(), (rb @ b32 + c32).relu()) - for si in sched: si.lower() - return len([si for si in sched if isinstance(si.prg, CompiledRunner)]) + linear = compile_linear(Tensor.schedule_linear((rb @ b16 + c16).relu(), (rb @ b32 + c32).relu())) + return len([call for call in linear.src if call.src[0].op is Ops.PROGRAM]) with Context(IMAGE=1): self.assertEqual(cnt(), 9) @@ -830,9 +826,8 @@ class TestSchedule(unittest.TestCase): x, y, z = Tensor.empty((1, 4, 3, 3)), Tensor.empty((4, 1, 3, 3)), Tensor.empty((4, 1, 7, 7)) a = x.conv2d(y, Tensor.empty(4), groups=4, padding=1) b = a.conv2d(z, groups=4, padding=3) - sched = (a + b).schedule() - for si in sched: si.lower() - return len([si for si in sched if isinstance(si.prg, CompiledRunner)]) + linear = compile_linear((a + b).schedule_linear()) + return len([call for call in linear.src if call.src[0].op is Ops.PROGRAM]) with Context(IMAGE=1): self.assertEqual(cnt(), 5) @@ -1055,7 +1050,7 @@ class TestSchedule(unittest.TestCase): expected = (a+a2).tolist() a.assign(a+a2) linear, var_vals = a.linear_with_vars() - kcount = len(linear_to_schedule(linear)) + kcount = len(linear.src) run_linear(linear, var_vals) self.assertListEqual(a.tolist(), expected) self.assertEqual(kcount, expected_kcount) @@ -1334,7 +1329,7 @@ class TestCopyFolding(unittest.TestCase): b = Tensor.empty(4, device="CPU") add = a+b assert all_same([x.device for x in add.uop.src]), f"ALU has different devices! {[x.device for x in add.src]}" - add.schedule() + add.schedule_linear() def test_alu_before_copy(self): buf = Tensor.ones(1).contiguous().realize() @@ -1356,7 +1351,7 @@ class TestCopyFolding(unittest.TestCase): a = Tensor.ones(4).contiguous().realize().uop.buf_uop t = Tensor(a.copy_to_device(a.device)) linear, var_vals = t.linear_with_vars() - assert len([s for s in linear_to_schedule(linear) if s.ast.op is Ops.COPY]) == 0 + assert len([call for call in linear.src if call.src[0].op is Ops.COPY]) == 0 run_linear(linear, var_vals) assert t.uop.is_realized, f"didn't realize Tensor {t}" self.assertListEqual(t.tolist(), [1.,1.,1.,1.]) diff --git a/test/backend/test_softmax_fusion.py b/test/backend/test_softmax_fusion.py index 7ee57853f1..dcfc283f8a 100644 --- a/test/backend/test_softmax_fusion.py +++ b/test/backend/test_softmax_fusion.py @@ -2,6 +2,7 @@ import unittest import numpy as np from tinygrad import Tensor, GlobalCounters, Context, Device from tinygrad.dtype import DTypeLike, dtypes +from tinygrad.engine.realize import run_linear from tinygrad.helpers import DEBUG, get_single_element from tinygrad.device import is_dtype_supported @@ -26,7 +27,10 @@ def single_kernel_softmax(x_in:Tensor, axis=-1, dtype:DTypeLike|None=None) -> Te out = e.div(ss).reshape(x_in.shape) return out -def run_one_schedule_item(out): get_single_element(out.schedule()).run() +def run_one_schedule_item(out): + linear = out.schedule_linear() + get_single_element(linear.src) + run_linear(linear) class TestFuse(unittest.TestCase): def _test_fuse(self, fxn, *args, atol=1e-6, allow_multiple=False, **kwargs): @@ -100,8 +104,8 @@ class TestFuse(unittest.TestCase): k = (x @ wk).contiguous() v = (x @ wv).contiguous() attn = q.scaled_dot_product_attention(k, v) - s = attn.schedule() - self.assertEqual(len(s), 4) # 3 matmul and 1 attention + s = attn.schedule_linear() + self.assertEqual(len(s.src), 4) # 3 matmul and 1 attention @unittest.skip("needs RANGEIFY>1") def test_flash_attention(self): diff --git a/test/backend/test_uops.py b/test/backend/test_uops.py index a66edcc20a..40fd9cf6a1 100644 --- a/test/backend/test_uops.py +++ b/test/backend/test_uops.py @@ -7,17 +7,16 @@ from tinygrad.dtype import dtypes, DType, AddrSpace, ConstFloat # noqa: F401 from tinygrad.device import Buffer, Device from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType from tinygrad.renderer.cstyle import CStyleLanguage -from tinygrad.engine.realize import CompiledRunner, get_program, get_runner -from tinygrad.schedule import ExecItem +from tinygrad.engine.realize import CompiledRunner, run_linear +from tinygrad.codegen import to_program from tinygrad.device import is_dtype_supported from tinygrad.codegen.opt import Opt, OptOps from tinygrad.renderer.ptx import PTXRenderer from test.helpers import to_uops_list -from dataclasses import replace def _uops_to_prg(uops_list): - prg = get_program(UOp.sink(*uops_list, arg=KernelInfo()), Device[Device.DEFAULT].renderer) - return CompiledRunner(replace(prg, device=Device.DEFAULT)) + prg = to_program(UOp.sink(*uops_list, arg=KernelInfo()), Device[Device.DEFAULT].renderer) + return CompiledRunner(prg, Device.DEFAULT) def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp: if op is Ops.CONST: uops.append(UOp.const(dtype, arg)) @@ -246,11 +245,11 @@ class TestAssembly(unittest.TestCase): a = Tensor.empty(1024) b = Tensor.empty(1024) c = (a*b).sum() - ast = c.schedule()[-1].ast + ast = c.schedule_linear().src[-1].src[0] opts_to_apply = [Opt(OptOps.UNROLL, 0, 4)] ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply))) - program = get_program(ast, Device[Device.DEFAULT].renderer) - uops = program.uops + program = to_program(ast, Device[Device.DEFAULT].renderer) + uops = tuple(program.src[2].src) self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4) def test_mulacc_shl(self): @@ -281,7 +280,7 @@ class TestZeroRange(unittest.TestCase): class TestUOpPrograms(unittest.TestCase): def _run(self, prog:UOp, *tensors:Tensor): - ExecItem(prog, [t.uop.buffer for t in tensors], prg=get_runner(Device.DEFAULT, prog)).run(wait=True) + run_linear(UOp(Ops.LINEAR, src=(prog.call(*[t.uop.buf_uop for t in tensors]),)), do_update_stats=False) def test_simple(self): out = Tensor.empty(10,10,dtype=dtypes.int) diff --git a/test/device/test_hcq.py b/test/device/test_hcq.py index 65cd294997..edba1d6ea2 100644 --- a/test/device/test_hcq.py +++ b/test/device/test_hcq.py @@ -6,7 +6,8 @@ from tinygrad.device import Buffer, BufferSpec from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer from tinygrad.runtime.autogen import libc from tinygrad.runtime.support.system import PCIIfaceBase -from tinygrad.engine.realize import get_runner, CompiledRunner, get_program +from tinygrad.engine.realize import get_runner, CompiledRunner +from tinygrad.codegen import to_program from tinygrad.codegen.opt import Opt, OptOps from tinygrad import Variable @@ -19,9 +20,9 @@ class TestHCQ(unittest.TestCase): TestHCQ.d0 = Device[Device.DEFAULT] TestHCQ.a = Tensor([0.,1.], device=Device.DEFAULT).realize() TestHCQ.b = self.a + 1 - si = self.b.schedule()[-1] + si = self.b.schedule_linear().src[-1] - TestHCQ.runner = get_runner(TestHCQ.d0.device, si.ast) + TestHCQ.runner = get_runner(TestHCQ.d0.device, si.src[0]) TestHCQ.b.uop.buffer.allocate() TestHCQ.kernargs_ba_ptr = TestHCQ.runner._prg.fill_kernargs([TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf]) @@ -163,9 +164,10 @@ class TestHCQ(unittest.TestCase): a = Tensor.randint((3, 3, 3), dtype=dtypes.int, device=Device.DEFAULT).realize() b = a + 1 - si = b.schedule()[-1] + si = b.schedule_linear().src[-1] - runner = CompiledRunner(get_program(replace_opts(si.ast, [Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]), TestHCQ.d0.renderer)) + runner = CompiledRunner(to_program(replace_opts(si.src[0], [Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]), TestHCQ.d0.renderer), + Device.DEFAULT) zb = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated() zt = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated() @@ -468,7 +470,7 @@ class TestHCQ(unittest.TestCase): def test_memory_barrier(self): a = Tensor([0, 1], device=Device.DEFAULT, dtype=dtypes.int8).realize() b = a + 1 - runner = get_runner(TestHCQ.d0.device, b.schedule()[-1].ast) + runner = get_runner(TestHCQ.d0.device, b.schedule_linear().src[-1].src[0]) buf1 = Buffer(Device.DEFAULT, 2, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated() buf2 = Buffer(Device.DEFAULT, 2, dtypes.int8, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated() diff --git a/test/device/test_metal.py b/test/device/test_metal.py index 80844cbc01..e86f643c36 100644 --- a/test/device/test_metal.py +++ b/test/device/test_metal.py @@ -50,6 +50,19 @@ kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgr compiled = compiled[:40] # corrupt the compiled program MetalProgram(device, "r_5", compiled) + def test_wait_skips_in_flight(self): + device = MetalDevice("metal") + compiled = MetalCompiler().compile(""" +#include +kernel void noop(uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {} +""") + prg = MetalProgram(device, "noop", compiled) + self.assertIsInstance(prg(wait=True), float) + self.assertEqual(device.mtl_buffers_in_flight, []) + self.assertIsNone(prg(wait=False)) + self.assertEqual(len(device.mtl_buffers_in_flight), 1) + device.synchronize() + def test_free(self): size = 2**16 device = Device['METAL'] diff --git a/test/device/test_validate_with_cpu.py b/test/device/test_validate_with_cpu.py new file mode 100644 index 0000000000..8b4710d4c8 --- /dev/null +++ b/test/device/test_validate_with_cpu.py @@ -0,0 +1,40 @@ +import unittest +from tinygrad import Tensor, Context, Variable, Device +from test.helpers import needs_second_gpu + +class TestValidateWithCPU(unittest.TestCase): + def setUp(self): + self.ctx = Context(VALIDATE_WITH_CPU=1) + self.ctx.__enter__() + def tearDown(self): self.ctx.__exit__(None, None, None) + + def test_add(self): self.assertListEqual((Tensor([1.,2,3])+Tensor([4.,5,6])).tolist(), [5.0, 7.0, 9.0]) + def test_mul(self): self.assertListEqual((Tensor([1.,2,3])*Tensor([4.,5,6])).tolist(), [4.0, 10.0, 18.0]) + def test_sum(self): self.assertEqual(Tensor([1.,2,3,4]).sum().item(), 10.0) + def test_reduce_then_op(self): self.assertEqual((Tensor([1.,2,3,4]).sum() * 2).item(), 20.0) + + def test_assign(self): + a = Tensor([1.,2,3]).realize() + a.assign(a + 1).realize() + self.assertListEqual(a.tolist(), [2.0, 3.0, 4.0]) + + def test_buffer_view(self): + self.assertListEqual((Tensor([1.,2,3,4,5,6,7,8])[2:6] + 1).tolist(), [4.0, 5.0, 6.0, 7.0]) + + def test_symbolic(self): + i = Variable('i', 1, 10) + ones = Tensor.ones(10).contiguous() + self.assertListEqual((ones[:i.bind(5)] + 1).contiguous()[:5].tolist(), [2.0]*5) + + def test_multi_kernel(self): + a = (Tensor([1.,2,3]) + 1).contiguous() + b = (a * 2).contiguous() + self.assertListEqual((b - 1).tolist(), [3.0, 5.0, 7.0]) + + @needs_second_gpu + def test_sharded(self): + t = Tensor([1.,2,3,4]).shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"), axis=0) + self.assertListEqual((t + 1).tolist(), [2.0, 3.0, 4.0, 5.0]) + +if __name__ == "__main__": + unittest.main() diff --git a/test/external/external_benchmark_op_conv.py b/test/external/external_benchmark_op_conv.py index 144cda7661..edaa47448b 100644 --- a/test/external/external_benchmark_op_conv.py +++ b/test/external/external_benchmark_op_conv.py @@ -3,7 +3,8 @@ from dataclasses import replace from tinygrad import dtypes, Device from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo from tinygrad.codegen.opt import Opt, OptOps # pylint: disable=unused-import -from tinygrad.engine.realize import CompiledRunner, get_program +from tinygrad.engine.realize import CompiledRunner +from tinygrad.codegen import to_program from tinygrad.helpers import dedup, getenv from tinygrad.device import Buffer from tinygrad.dtype import ImageDType, Invalid @@ -88,13 +89,13 @@ ast = {143: vision_conv_143, 153: vision_conv_153, 172: dm_conv_172}[getenv("NUM renderer = Device.default.renderer allocator = Device.default.allocator -ps = get_program(ast, renderer) -cr = CompiledRunner(replace(ps, device=Device.DEFAULT)) +ps = to_program(ast, renderer) +cr = CompiledRunner(ps, Device.DEFAULT) gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.PARAM]), key=lambda u: u.arg) # print(len(gs)) # print([g.dtype for g in gs]) -bufs = [Buffer(ps.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs] +bufs = [Buffer(ps.arg.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs] t = cr(bufs, wait=True) print(f"{t*1e6:.2f} us") \ No newline at end of file diff --git a/test/external/external_benchmark_schedule.py b/test/external/external_benchmark_schedule.py index 86879e7489..d7efc7be87 100644 --- a/test/external/external_benchmark_schedule.py +++ b/test/external/external_benchmark_schedule.py @@ -23,10 +23,10 @@ if __name__ == "__main__": if not FORWARD_ONLY: with Timing("***** model schedule in "): with Profiling(PROFILE >= 3): - sched = out.schedule() + linear = out.schedule_linear() if not SCHEDULE_ONLY: - asts = list({x.ast.key:x.ast for x in sched if x.ast.op is Ops.SINK}.values()) + asts = list({call.src[0].key:call.src[0] for call in linear.src if call.src[0].op is Ops.SINK}.values()) if (restrict_kernel := getenv("RESTRICT_KERNEL", -1)) != -1: asts = asts[restrict_kernel:restrict_kernel+1] with Profiling(PROFILE, fn="/tmp/rewrite.prof"): diff --git a/test/external/external_test_hcq.py b/test/external/external_test_hcq.py index 6e1da82703..ad2a5f0cbe 100644 --- a/test/external/external_test_hcq.py +++ b/test/external/external_test_hcq.py @@ -20,8 +20,8 @@ class TestHCQ(unittest.TestCase): #TestHCQ.d1: AMDDevice = Device["AMD:1"] TestHCQ.a = Tensor([0.,1.], device=Device.DEFAULT).realize() TestHCQ.b = self.a + 1 - si = self.b.schedule()[-1] - TestHCQ.runner = get_runner(TestHCQ.d0.device, si.ast) + linear = self.b.schedule_linear() + TestHCQ.runner = get_runner(TestHCQ.d0.device, linear.src[-1].src[0]) TestHCQ.b.uop.buffer.allocate() # wow that's a lot of abstraction layers TestHCQ.addr = struct.pack("QQ", TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf) diff --git a/test/external/external_test_onnx_runner.py b/test/external/external_test_onnx_runner.py index 3d58f9c323..d91a2aec59 100644 --- a/test/external/external_test_onnx_runner.py +++ b/test/external/external_test_onnx_runner.py @@ -10,8 +10,8 @@ from hypothesis import given, strategies as st # copied from test_const_folding.py def _check_ast_count(desired_count:int, t:Tensor): # NOTE: this has side effect because everything can be scheduled only once - schedule = t.schedule() - asts = [s for s in schedule if s.ast.op is Ops.SINK] + linear = t.schedule_linear() + asts = [call for call in linear.src if call.src[0].op is Ops.SINK] assert len(asts) == desired_count, f"{len(asts)} != {desired_count}" def build_onnx(nodes, from_disk:bool=True, **kwargs): diff --git a/test/external/external_test_opt.py b/test/external/external_test_opt.py index 4f57694c9f..0ae7ae32bc 100644 --- a/test/external/external_test_opt.py +++ b/test/external/external_test_opt.py @@ -7,7 +7,6 @@ from tinygrad import GlobalCounters, Tensor, Device from tinygrad.helpers import getenv from tinygrad.nn.state import get_parameters from tinygrad.engine.realize import capturing, run_linear -from tinygrad.schedule import linear_to_schedule from tinygrad.tensor import _to_np_dtype class CLCache: @@ -15,7 +14,7 @@ class CLCache: self.allowed, self.strict, self.preclear, self.var_vals = allowed, strict, preclear, var_vals if var_vals is not None else {} self.count = 0 def add_linear(self, linear, var_vals): - self.count += len(linear_to_schedule(linear)) + self.count += len(linear.src) run_linear(linear, var_vals) def __enter__(self): if self.preclear: diff --git a/test/external/external_test_schedule_scaling.py b/test/external/external_test_schedule_scaling.py index 2e07ea4e25..e7026508ee 100644 --- a/test/external/external_test_schedule_scaling.py +++ b/test/external/external_test_schedule_scaling.py @@ -6,7 +6,7 @@ class TestScheduleScaling(unittest.TestCase): def _assert_linear(self, fn, n_small=200, n_large=1000): """Assert schedule time scales at most ~linearly: time(n_large)/time(n_small) should be close to n_large/n_small.""" - fn(n_small).schedule() # warmup + fn(n_small).schedule_linear() # warmup t_small = min(self._time_schedule(fn, n) for n in [n_small]*3) t_large = min(self._time_schedule(fn, n) for n in [n_large]*3) size_ratio = n_large / n_small # 5.0 @@ -19,7 +19,7 @@ class TestScheduleScaling(unittest.TestCase): @staticmethod def _time_schedule(fn, n) -> float: st = time.perf_counter() - fn(n).schedule() + fn(n).schedule_linear() return time.perf_counter() - st # *** rangeify: ending_ranges accumulation and consumer merge *** diff --git a/test/external/external_uop_gc.py b/test/external/external_uop_gc.py index 74712a09e3..3077009ae8 100644 --- a/test/external/external_uop_gc.py +++ b/test/external/external_uop_gc.py @@ -1,7 +1,8 @@ import gc from tinygrad import Tensor, UOp, Device, nn from tinygrad.schedule import schedule_cache -from tinygrad.engine.realize import method_cache, get_program +from tinygrad.engine.realize import method_cache +from tinygrad.codegen import to_program, to_program_cache from tinygrad.schedule.indexing import apply_movement_op, _apply_reshape from tinygrad.uop.divandmod import fold_divmod_general from test.test_tiny import TestTiny @@ -14,13 +15,13 @@ def print_uops(): def start(): pass def single_tensor(): Tensor([2]) def two_plus_two(): Tensor([2])+Tensor([2]) -def two_plus_two_schedule(): (Tensor([2])+Tensor([2])).schedule() +def two_plus_two_schedule(): (Tensor([2])+Tensor([2])).schedule_linear() def two_plus_two_kernel(): - si = (Tensor([2])+Tensor([2])).schedule()[-1] - get_program(si.ast, Device.default.renderer) + linear = (Tensor([2])+Tensor([2])).schedule_linear() + to_program(linear.src[-1].src[0], Device.default.renderer) def two_plus_two_linearize(): - si = (Tensor([2])+Tensor([2])).schedule()[-1] - get_program(si.ast, Device.default.renderer) + linear = (Tensor([2])+Tensor([2])).schedule_linear() + to_program(linear.src[-1].src[0], Device.default.renderer) def two_plus_two_realize(): (Tensor([2])+Tensor([2])).realize() def two_plus_two_item(): (Tensor([2])+Tensor([2])).item() def gradient_test(): @@ -36,8 +37,8 @@ def kernel_matmul(): x = Tensor.eye(3, requires_grad=True) y = Tensor([[2.0,0,-2.0]], requires_grad=True) z = y.matmul(x) - si = z.schedule()[-1] - get_program(si.ast, Device.default.renderer) + linear = z.schedule_linear() + to_program(linear.src[-1].src[0], Device.default.renderer) def realized_matmul(): x = Tensor.eye(3, requires_grad=True) y = Tensor([[2.0,0,-2.0]], requires_grad=True) @@ -71,6 +72,7 @@ if __name__ == "__main__": # these caches will keep uops alive schedule_cache.clear() method_cache.clear() + to_program_cache.clear() apply_movement_op.cache_clear() _apply_reshape.cache_clear() fold_divmod_general.cache_clear() diff --git a/test/external/fuzz_graph.py b/test/external/fuzz_graph.py index c168e63b51..edc2c6faee 100644 --- a/test/external/fuzz_graph.py +++ b/test/external/fuzz_graph.py @@ -4,8 +4,7 @@ from tinygrad.device import Buffer, Device from tinygrad.helpers import Context, getenv, from_mv from tinygrad.dtype import dtypes from tinygrad.tensor import Tensor, _to_np_dtype -from tinygrad.engine.realize import BufferXfer, get_runner -from tinygrad.schedule import ExecItem +from tinygrad.engine.realize import BufferXfer, get_runner, ExecItem from tinygrad.uop.ops import UOp, Ops from tinygrad.engine.jit import apply_graph_to_jit @@ -20,8 +19,8 @@ def gen_prg(device, inputs_cnt): s = fst[0] for i in range(1, inputs_cnt): s = s.bitwise_xor(fst[i]) - si = s.schedule()[-1] - prg = get_runner(device, si.ast) + linear = s.schedule_linear() + prg = get_runner(device, linear.src[-1].src[0]) cached_prgs[(device, inputs_cnt)] = prg return prg diff --git a/test/helpers.py b/test/helpers.py index ed4f55b19e..80227b2ae1 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -5,7 +5,7 @@ import numpy as np from tinygrad import Tensor, dtypes, Device from tinygrad.uop.ops import UOp, Ops, KernelInfo from tinygrad.tensor import _to_np_dtype -from tinygrad.engine.realize import get_program +from tinygrad.codegen import to_program from tinygrad.dtype import DType from tinygrad.nn.state import get_parameters from tinygrad.helpers import T, CI, Target @@ -80,8 +80,8 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None): bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize)) allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data))) g = UOp(Ops.PARAM, uop.dtype.ptr(), arg=0, src=()) - prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON"))) - prog = PythonProgram("run", PythonCompiler().compile(prg.src)) + prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON"))) + prog = PythonProgram("run", PythonCompiler().compile(prg.src[3].arg)) prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs) return out_buf.cast(uop.dtype.fmt or "").tolist()[0] diff --git a/test/mockgpu/usb.py b/test/mockgpu/usb.py index cd59a9ce43..4946178aa8 100644 --- a/test/mockgpu/usb.py +++ b/test/mockgpu/usb.py @@ -202,6 +202,8 @@ class MockASM24State: return None class MockUSB3: + @classmethod + def list_devices(cls, vendor, dev): return [(0, "usb:mock")] def __init__(self, *args, **kwargs): self.product, self.is_custom = "", False def send_batch(self, cdbs:list[bytes], idata:list[int]|None=None, odata:list[bytes|None]|None=None) -> list[bytes|None]: diff --git a/test/null/test_attention.py b/test/null/test_attention.py index a42558333e..aa23608de9 100644 --- a/test/null/test_attention.py +++ b/test/null/test_attention.py @@ -16,9 +16,9 @@ class TestAttention(unittest.TestCase): k = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize() v = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize() attn = q.scaled_dot_product_attention(k, v) - sched = attn.schedule() + sched = attn.schedule_linear() # attention has 4 kernels now - self.assertEqual(len(sched), 4) + self.assertEqual(len(sched.src), 4) def test_apply_rope_jit_prune(self): def rope_fn(x_in, pos): return apply_rope(x_in, pos) diff --git a/test/null/test_compile_failures.py b/test/null/test_compile_failures.py index ba8245862a..3d4a5c3b34 100644 --- a/test/null/test_compile_failures.py +++ b/test/null/test_compile_failures.py @@ -3,11 +3,12 @@ from contextlib import redirect_stdout from tinygrad import Tensor, dtypes, Device from tinygrad.helpers import OSX, DEV from tinygrad.device import is_dtype_supported -from tinygrad.engine.realize import get_program +from tinygrad.engine.realize import compile_linear +from tinygrad.codegen import to_program class TestCompileFailures(unittest.TestCase): def compile(self, out:Tensor): - for si in out.schedule(): si.lower() + compile_linear(out.schedule_linear()) @unittest.skipUnless(is_dtype_supported(dtypes.uchar), f"no uint8 on {Device.DEFAULT}") def test_interpolate_atari(self): @@ -21,9 +22,9 @@ class TestDisassembly(unittest.TestCase): @unittest.skipUnless(Device.DEFAULT in ("CPU",) and DEV.renderer not in ("LLVM", "LVP") and OSX, "m series cpus support fp16 arithmetic") def test_float16_alu(self): c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16) - s = c.schedule()[-1] - p = get_program(s.ast, Device[Device.DEFAULT].renderer) - lib = Device[Device.DEFAULT].compiler.compile(p.src) + s = c.schedule_linear().src[-1] + p = to_program(s.src[0], Device[Device.DEFAULT].renderer) + lib = Device[Device.DEFAULT].compiler.compile(p.src[3].arg) out = io.StringIO() with redirect_stdout(out): Device[Device.DEFAULT].compiler.disassemble(lib) assert "fcvt" not in out.getvalue() diff --git a/test/null/test_const_folding.py b/test/null/test_const_folding.py index b9e71a1cce..1d244609af 100644 --- a/test/null/test_const_folding.py +++ b/test/null/test_const_folding.py @@ -7,8 +7,8 @@ import numpy as np def _check_ast_count(desired_count:int, t:Tensor): # NOTE: this has side effect because everything can be scheduled only once - schedule = t.schedule() - asts = [s for s in schedule if s.ast.op is Ops.SINK] + linear = t.schedule_linear() + asts = [s for s in linear.src if s.src[0].op is Ops.SINK] len(asts) # NOT SUPPORTED ANYMORE #assert len(asts) == desired_count, f"{len(asts)} != {desired_count}" diff --git a/test/null/test_device.py b/test/null/test_device.py index 90f8588fef..e798d653f1 100644 --- a/test/null/test_device.py +++ b/test/null/test_device.py @@ -153,7 +153,7 @@ class TestDevVar(unittest.TestCase): self.assertEqual(DEV.target("CPU"), Target("CPU")) def test_dev_arch_override(self): - with Context(DEV="NULL:HIP:gfx1100"): + with Context(DEV="NULL::gfx1100"): self.assertEqual(Device["NULL"].renderer.target.arch, "gfx1100") class MockCompiler(Compiler): diff --git a/test/null/test_gc.py b/test/null/test_gc.py index 46d76962fb..21682e395f 100644 --- a/test/null/test_gc.py +++ b/test/null/test_gc.py @@ -60,7 +60,7 @@ class TestGC(unittest.TestCase): init = bufs_allocated() x = Tensor.ones(256).contiguous().realize() y = Tensor.ones(5, 5).contiguous() - y.schedule() + y.schedule_linear() del x del y self.assertEqual(bufs_allocated()-init, 0) diff --git a/test/null/test_linearizer_failures.py b/test/null/test_linearizer_failures.py index b8636c8ac2..c65f27ab69 100644 --- a/test/null/test_linearizer_failures.py +++ b/test/null/test_linearizer_failures.py @@ -2,8 +2,8 @@ import unittest from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo from tinygrad.dtype import dtypes -from tinygrad.engine.realize import get_program from tinygrad.device import Device +from tinygrad.codegen import to_program class TestLinearizerFailures(unittest.TestCase): def test_fail_1(self): @@ -19,7 +19,7 @@ class TestLinearizerFailures(unittest.TestCase): c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal() c10 = c0.index(c3).store(c9).end(c1, c2) ast = c10.sink(arg=KernelInfo()) - get_program(ast, renderer=Device[Device.DEFAULT].renderer) + to_program(ast, renderer=Device[Device.DEFAULT].renderer) if __name__ == '__main__': unittest.main() diff --git a/test/null/test_linearizer_rewrite.py b/test/null/test_linearizer_rewrite.py index 35c61b6ab1..93c506d25a 100644 --- a/test/null/test_linearizer_rewrite.py +++ b/test/null/test_linearizer_rewrite.py @@ -1,6 +1,6 @@ import unittest from tinygrad import Tensor, Context, Device -from tinygrad.engine.realize import get_program +from tinygrad.codegen import to_program from tinygrad.codegen.opt import Opt, OptOps from tinygrad.uop.ops import KernelInfo @@ -9,37 +9,37 @@ class TestLinearizerRewrite(unittest.TestCase): t = Tensor.ones((64,64), device="NULL").contiguous().realize() out = (t*2).sum(axis=1) with Context(SPLIT_REDUCEOP=0, DEVECTORIZE=0): - si = out.schedule()[-1] + si = out.schedule_linear().src[-1] opts_to_apply = [] opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4)) opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4)) - ast = si.ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply))) - prg = get_program(ast, Device["CPU"].renderer) - print(prg.src) + ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply))) + prg = to_program(ast, Device["CPU"].renderer) + print(prg.src[3].arg) def test_arange(self): out = Tensor.arange(32, device="NULL") with Context(SPLIT_REDUCEOP=0, DEVECTORIZE=0): - si = out.schedule()[-1] + si = out.schedule_linear().src[-1] opts_to_apply = [] opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4)) - ast = si.ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply))) - prg = get_program(ast, Device["CPU"].renderer) - print(prg.src) + ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply))) + prg = to_program(ast, Device["CPU"].renderer) + print(prg.src[3].arg) def test_kernel_info(self): out = Tensor.arange(4, device="NULL") - si = out.schedule()[-1] + si = out.schedule_linear().src[-1] - ast = si.ast.replace(arg=KernelInfo(opts_to_apply=())) - prg = get_program(ast, Device["CPU"].renderer) - assert prg.applied_opts == (), f"expected no opts, got {prg}" + ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=())) + prg = to_program(ast, Device["CPU"].renderer) + assert prg.src[0].arg.applied_opts == (), f"expected no opts, got {prg}" - prg = get_program(ast.replace(arg=KernelInfo()), Device["CPU"].renderer) - assert prg.applied_opts != (), f"expected opts to apply, got {prg.applied_opts}" + prg = to_program(ast.replace(arg=KernelInfo()), Device["CPU"].renderer) + assert prg.src[0].arg.applied_opts != (), f"expected opts to apply, got {prg.src[0].arg.applied_opts}" - prg = get_program(ast.replace(arg=KernelInfo(name="custom")), Device["CPU"].renderer) - self.assertEqual(prg.name, "custom") + prg = to_program(ast.replace(arg=KernelInfo(name="custom")), Device["CPU"].renderer) + self.assertEqual(prg.arg.name, "custom") if __name__ == '__main__': unittest.main() diff --git a/test/null/test_process_replay.py b/test/null/test_process_replay.py index 87c032ed88..b75e4b09cf 100644 --- a/test/null/test_process_replay.py +++ b/test/null/test_process_replay.py @@ -9,7 +9,7 @@ N = 16 class TestProcessReplay(unittest.TestCase): @classmethod def setUpClass(cls): - cls.ast = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1].ast + cls.ast = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0] cls.renderer = Device[Device.DEFAULT].renderer def test_replay_no_opts(self): @@ -35,9 +35,9 @@ class TestProcessReplay(unittest.TestCase): def test_beam(self): with Context(BEAM=1): - si = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1] - p = do_to_program(si.ast, self.renderer) - good, compare, _ = replay_to_program(p, si.ast, self.renderer) + ast = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0] + p = do_to_program(ast, self.renderer) + good, compare, _ = replay_to_program(p, ast, self.renderer) self.assertEqual(good, compare) if __name__ == '__main__': diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py index a5095fd360..cef6054e88 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -3,8 +3,8 @@ import gc, unittest, time from tinygrad import nn, dtypes, Device, Tensor from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, KernelInfo from tinygrad.helpers import DEBUG, GlobalCounters, Context -from tinygrad.engine.realize import CompiledRunner, run_linear -from tinygrad.schedule import linear_to_schedule +from tinygrad.engine.realize import compile_linear, run_linear +from tinygrad.codegen import to_program class KernelCountException(Exception): pass def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True): @@ -15,17 +15,17 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te else: assert isinstance(t, UOp), f"can't schedule {t}" linear, var_vals = Tensor(t).linear_with_vars() - # test lowering all the ExecItems - sched = linear_to_schedule(linear) - for si in sched: si.lower() - kernel_cnt = len([si for si in sched if isinstance(si.prg, CompiledRunner) or not filter_sink]) + kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1) + for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink) if kernel_cnt != allowed: print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}") if DEBUG >= 3: - for i,s in enumerate(sched): + for i,call in enumerate(linear.src): print("kernel", i+1) - print(s.ast) + print(call.src[0]) raise KernelCountException(f"{kernel_cnt} != {allowed}") + # test compiling the linear + compile_linear(linear) return linear, var_vals def _realize_weights(m): @@ -88,7 +88,7 @@ class TestBufferUOp(unittest.TestCase): # unused variable should not appear in var_vals even when there's other work a = Tensor(UOp.variable("unused", 0, 10).bind(1)) b = Tensor.empty(3) + 1 - _, var_vals = Tensor.schedule_with_vars(a, b) + _, var_vals = Tensor.linear_with_vars(a, b) self.assertEqual(var_vals, {}) self.assertIsNone(a.uop.base.realized) @@ -142,7 +142,7 @@ class TestSimpleSchedule(unittest.TestCase): a = Tensor.empty(16,16).sum(axis=1) a1 = a.reshape(4,4) a2 = a.reshape(16,1,1) - self.assertEqual(len(Tensor.schedule(a1, a2)), 1) + self.assertEqual(len(Tensor.schedule_linear(a1, a2).src), 1) class TestSchedule(unittest.TestCase): def test_create_schedule_handles_multi_kernel_after_and_after_deps(self): @@ -167,8 +167,8 @@ class TestSchedule(unittest.TestCase): kc = Tensor.custom_kernel(out, src_after, fxn=named_copy("kc"))[0] out_after = Tensor(kc.uop.src[0].after(*kc.uop.src[1:], kd.uop)) - schedule = out_after.schedule() - names = [si.ast.arg.name for si in schedule] + linear = out_after.schedule_linear() + names = [call.src[0].arg.name for call in linear.src] self.assertEqual(set(names), {"ka", "kb", "kc", "kd"}) self.assertEqual(names[-1], "kc") self.assertLess(names.index("ka"), names.index("kc")) @@ -209,10 +209,10 @@ class TestSchedule(unittest.TestCase): t = Tensor.zeros((3, 3)).contiguous().realize() v = t[1] # view - is_realized but not has_buffer_identity assert v.uop.is_realized - sched, _ = Tensor.schedule_with_vars(v) - self.assertEqual(len(sched), 0) + linear, _ = Tensor.linear_with_vars(v) + self.assertEqual(len(linear.src), 0) - # NOTE: because empty does not have a lowered ExecItem if realize is called on a childless empty, it never gets allocated. + # NOTE: because empty does not have a lowered kernel if realize is called on a childless empty, it never gets allocated. def test_childless_empty_never_allocates(self): a = Tensor.empty(10) a.realize() @@ -668,9 +668,9 @@ class TestSchedule(unittest.TestCase): check_schedule(c, 2) def _alu_from_tensor(self, t:Tensor): - s = [s for s in t.schedule() if s.ast.op is Ops.SINK] + s = [s for s in t.schedule_linear().src if s.src[0].op is Ops.SINK] self.assertEqual(len(s), 1) - return [u.op for u in s[0].ast.toposort() if u.op in GroupOp.ALU] + return [u.op for u in s[0].src[0].toposort() if u.op in GroupOp.ALU] def test_2_pow_is_exp2(self): t = 2.0 ** Tensor([1.0, 2.0, 3.0]) @@ -799,12 +799,12 @@ class TestSchedule(unittest.TestCase): Tensor.manual_seed(0) x = Tensor.randn(4, 12, 64, 64, dtype=dtypes.half).realize() out = x.softmax(dtype=dtypes.float) - sched = out.schedule() - self.assertEqual(len(sched), 3) + linear = out.schedule_linear() + self.assertEqual(len(linear.src), 3) # max reduction stays in input dtype (no numerical loss), upcast happens after subtracting max - self.assertEqual(sched[0].bufs[0].dtype, dtypes.half) - self.assertEqual(sched[1].bufs[0].dtype, dtypes.float) - self.assertEqual(sched[2].bufs[0].dtype, dtypes.float) + self.assertEqual(linear.src[0].src[1].dtype, dtypes.half) + self.assertEqual(linear.src[1].src[1].dtype, dtypes.float) + self.assertEqual(linear.src[2].src[1].dtype, dtypes.float) def test_softmax_backward(self): Tensor.manual_seed(0) @@ -961,7 +961,7 @@ class TestSchedule(unittest.TestCase): gc.collect() base = GlobalCounters.mem_used Tensor.ones(256).contiguous().realize() - Tensor.ones(5, 5).contiguous().schedule() + Tensor.ones(5, 5).contiguous().schedule_linear() gc.collect() self.assertEqual(GlobalCounters.mem_used-base, 0) @@ -1174,24 +1174,24 @@ class TestFusionOp(unittest.TestCase): st = time.perf_counter() a = Tensor([1,2,3,4]) for _ in range(24): a = a + a - sched = a.schedule() - sched[-1].lower() + linear = a.schedule_linear() + prg = to_program(linear.src[-1].src[0], renderer=Device[Device.DEFAULT].renderer) self.assertLess(time.perf_counter()-st, 2.0) - assert len(sched[-1].prg.p.src.splitlines()) < 250 + assert len(prg.src[3].arg.splitlines()) < 250 def test_recursive_add_cmp(self): st = time.perf_counter() a = Tensor([1,2,3,4]) for _ in range(24): a = a + a - sched1 = a.schedule() + linear1 = a.schedule_linear() b = Tensor([1,2,3,4]) for _ in range(24): b = b + b - sched2 = b.schedule() + linear2 = b.schedule_linear() c = Tensor([1,2,3,4]) for _ in range(23): c = c + c - sched3 = c.schedule() - self.assertEqual(sched1[-1].ast, sched2[-1].ast) - with self.assertRaises(AssertionError): self.assertEqual(sched1[-1].ast, sched3[-1].ast) + linear3 = c.schedule_linear() + self.assertEqual(linear1.src[-1].src[0], linear2.src[-1].src[0]) + with self.assertRaises(AssertionError): self.assertEqual(linear1.src[-1].src[0], linear3.src[-1].src[0]) self.assertLess(time.perf_counter()-st, 2.0) def test_recursive_pad(self): @@ -1199,8 +1199,8 @@ class TestFusionOp(unittest.TestCase): val = 1.0 a = Tensor(val) for _ in range(24): a = Tensor.stack(a, a)[0] - sched = a.schedule() - self.assertLessEqual(len(sched), 1) + linear = a.schedule_linear() + self.assertLessEqual(len(linear.src), 1) self.assertLess(time.perf_counter()-st, 2.0) def test_recursive_reshape(self): @@ -1209,8 +1209,8 @@ class TestFusionOp(unittest.TestCase): b = Tensor.empty(16, 2).realize() r = a.sum(1) for _ in range(24): r = r.reshape(16, 2) + b - sched = r.schedule() - self.assertEqual(len(sched), 1) + linear = r.schedule_linear() + self.assertEqual(len(linear.src), 1) self.assertLess(time.perf_counter()-st, 2.0) # NOTE: the NULL backend supports BUFFER_VIEW diff --git a/test/null/test_schedule_cache.py b/test/null/test_schedule_cache.py index e8897b1a07..99717404a9 100644 --- a/test/null/test_schedule_cache.py +++ b/test/null/test_schedule_cache.py @@ -4,7 +4,7 @@ from tinygrad.helpers import cpu_events from tinygrad.schedule import schedule_cache def schedule_one(): - Tensor([1]).schedule() + Tensor([1]).schedule_linear() class TestScheduleCache(unittest.TestCase): def test_bound_variable_var_vals(self): @@ -12,7 +12,7 @@ class TestScheduleCache(unittest.TestCase): x = Tensor.ones(10).contiguous().realize() t = x + Tensor(v.bind(42)) - _, var_vals = t.schedule_with_vars() + _, var_vals = t.linear_with_vars() self.assertEqual(var_vals, {'pos': 42}) def test_disable_schedule_cache(self): diff --git a/test/null/test_tensor.py b/test/null/test_tensor.py index 6dee065480..209d122e42 100644 --- a/test/null/test_tensor.py +++ b/test/null/test_tensor.py @@ -6,7 +6,7 @@ from tinygrad.device import is_dtype_supported from tinygrad.uop.ops import Ops, UOp from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.nir import NIRRenderer -from tinygrad.engine.realize import get_program +from tinygrad.codegen import to_program from tinygrad.dtype import DType x_init = np.random.randn(1,3).astype(np.float32) @@ -62,12 +62,13 @@ class TestIdxUpcast(unittest.TestCase): for src in ast.src: if (ret:=self._find_op(src, op)) is not None: return ret def _schedule_render(self, a: Tensor): - schedule, _ = a.schedule_with_vars() - for s in schedule: - if s.ast.op is Ops.SINK: - renderer = Device[s.bufs[0].device].renderer - prg = get_program(s.ast, renderer) - return prg.uops + linear, _ = a.linear_with_vars() + for si in linear.src: + ast = si.src[0] + if ast.op is Ops.SINK: + renderer = Device[si.src[1].buffer.device].renderer + prg = to_program(ast, renderer) + return tuple(prg.src[2].src) def _assert(self, dtype: DType, a: Tensor): uops = self._schedule_render(a) @@ -162,9 +163,9 @@ class TestRand(unittest.TestCase): def test_rand_large_tensor(self): # large tensor rand (num > uint32.max) should not crash in frontend Tensor.manual_seed(0) - Tensor.rand(2**17, 2**17).schedule() - Tensor.rand(2**17, 2**17).schedule() - Tensor.rand(2**17, 2**17).schedule() + Tensor.rand(2**17, 2**17).schedule_linear() + Tensor.rand(2**17, 2**17).schedule_linear() + Tensor.rand(2**17, 2**17).schedule_linear() class TestTensorConstLike(unittest.TestCase): def test_const_like_shape(self): diff --git a/test/null/test_tensor_metadata.py b/test/null/test_tensor_metadata.py index 665358e660..0ffa9ee8b4 100644 --- a/test/null/test_tensor_metadata.py +++ b/test/null/test_tensor_metadata.py @@ -2,7 +2,6 @@ import unittest from tinygrad import Tensor, dtypes from tinygrad.tensor import _METADATA from tinygrad.engine.realize import capturing -from tinygrad.schedule import linear_to_schedule from tinygrad.helpers import Context @unittest.skip("tensor metadata is no longer supported") @@ -18,41 +17,41 @@ class TestTensorMetadata(unittest.TestCase): def test_exclude_noop_metadata(self): a = Tensor.rand(4, 4)*1 self.assertEqual(a.uop.metadata[0].name, "__mul__") - k = a.schedule()[-1] - self.assertEqual([m.name for m in k.metadata], ["rand"]) + k = a.schedule_linear().src[-1] + self.assertEqual([m.name for m in k.arg.metadata], ["rand"]) @unittest.skip("metadata not reaching kernel schedule") def test_exclude_const_metadata(self): a = Tensor.arange(4) b = Tensor.full((4,), -1, dtype=dtypes.int).contiguous() - sched = Tensor.schedule(a, b) - self.assertEqual([m.name for m in sched[0].metadata], ["arange"]) - self.assertEqual([m.name for m in sched[1].metadata], ["contiguous"]) + sched = a.schedule_linear(b) + self.assertEqual([m.name for m in sched.src[0].arg.metadata], ["arange"]) + self.assertEqual([m.name for m in sched.src[1].arg.metadata], ["contiguous"]) def test_matmul(self): x = Tensor.rand(3, requires_grad=True) W = Tensor.rand(3, 3, requires_grad=True) out = x.matmul(W) self.assertEqual(out.uop.metadata[0].name, "matmul") - si = out.schedule()[-1] - self.assertEqual(len(si.metadata), 1) - self.assertEqual(si.metadata[0].name, "matmul") + si = out.schedule_linear().src[-1] + self.assertEqual(len(si.arg.metadata), 1) + self.assertEqual(si.arg.metadata[0].name, "matmul") def test_relu(self): x = Tensor.rand(3, requires_grad=True) out = x.relu() self.assertEqual(out.uop.metadata[0].name, "relu") - si = out.schedule()[-1] - self.assertEqual(len(si.metadata), 1) - self.assertEqual(si.metadata[0].name, "relu") + si = out.schedule_linear().src[-1] + self.assertEqual(len(si.arg.metadata), 1) + self.assertEqual(si.arg.metadata[0].name, "relu") @unittest.skip("assign metadata no longer captured") def test_assign(self): x = Tensor.empty(10, 10).realize() x.assign(Tensor.ones(10, 10).contiguous()) - si = x.schedule()[-1] - self.assertEqual(len(si.metadata), 1) - self.assertEqual(si.metadata[0].name, "assign") + si = x.schedule_linear().src[-1] + self.assertEqual(len(si.arg.metadata), 1) + self.assertEqual(si.arg.metadata[0].name, "assign") def test_complex(self): x = Tensor.rand(3, requires_grad=True) @@ -61,9 +60,9 @@ class TestTensorMetadata(unittest.TestCase): self.assertEqual(out.uop.metadata[0].name, "__mul__") self.assertEqual(out.uop.src[0].metadata[0].name, "relu") self.assertEqual(out.uop.src[1].metadata[0].name, "sigmoid") - si = out.schedule()[-1] - self.assertEqual(len(si.metadata), 3) - self.assertEqual(set(m.name for m in si.metadata), {"relu", "sigmoid", "__mul__"}) + si = out.schedule_linear().src[-1] + self.assertEqual(len(si.arg.metadata), 3) + self.assertEqual(set(m.name for m in si.arg.metadata), {"relu", "sigmoid", "__mul__"}) @unittest.skip("flaky") def test_complex_backward(self): @@ -76,10 +75,10 @@ class TestTensorMetadata(unittest.TestCase): #self.assertTrue(x.grad.uop.metadata[0].backward) # TODO: backward flag is False self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid") #self.assertTrue(y.grad.uop.metadata[0].backward) # TODO: backward flag is False - si = Tensor.schedule(out, x.grad, y.grad)[-1] - #self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}") + si = out.schedule_linear(x.grad, y.grad).src[-1] + #self.assertEqual(len(si.arg.metadata), 3, f"failed with {si.arg.metadata}") # skip numpy, this is schedule cache - self.assertSetEqual(set(m.name for m in si.metadata if m.name != "numpy"), {"sigmoid", "relu"}) + self.assertSetEqual(set(m.name for m in si.arg.metadata if m.name != "numpy"), {"sigmoid", "relu"}) #bw = [m for m in si.metadata if m.backward] #self.assertEqual(len(bw), 1) #self.assertEqual(bw[0].name, "sigmoid") @@ -91,16 +90,16 @@ class TestTensorMetadata(unittest.TestCase): out = (x.relu() * y.sigmoid()).sum() self.assertIsNone(out.uop.metadata) self.assertIsNone(out.uop.src[0].metadata) - si = out.schedule()[-1] - self.assertEqual(si.metadata, ()) + si = out.schedule_linear().src[-1] + self.assertEqual(si.arg.metadata, ()) def _has_metadata(self, h, name): linears = [] capturing.append(type("", (), {"add_linear": lambda _, linear, var_vals: linears.append(linear)})()) try: h.realize() finally: capturing.clear() - items = [ei for linear in linears for ei in linear_to_schedule(linear)] - return any(m.name == name for ei in items for m in ei.metadata) + calls = [call for linear in linears for call in linear.src] + return any(m.name == name for call in calls for m in call.arg.metadata) def test_metadata_survives_realize_pending_assign(self): shared = Tensor.rand(4) diff --git a/test/null/test_tensor_uop_mixin.py b/test/null/test_tensor_uop_mixin.py index 14b69028bc..02fbafb9ae 100644 --- a/test/null/test_tensor_uop_mixin.py +++ b/test/null/test_tensor_uop_mixin.py @@ -112,6 +112,31 @@ class TestTensorUOpOneHot(unittest.TestCase): t = _t(5) self.assertIs(_strip_unique(t.one_hot(5).uop), _strip_unique(t.uop.one_hot(5))) +class TestTensorUOpSort(unittest.TestCase): + def _check(self, t, **kw): + tv, ti = t.sort(**kw) + uv, ui = t.uop.sort(**kw) + self.assertIs(_strip_unique(tv.uop), _strip_unique(uv)) + self.assertIs(_strip_unique(ti.uop), _strip_unique(ui)) + def test_sort_1d(self): self._check(Tensor([0.5, 0.1, 0.3]).float()) + def test_sort_descending(self): self._check(Tensor([0.5, 0.1, 0.3]).float(), descending=True) + def test_sort_2d(self): self._check(_t(2, 4).float()) + def test_sort_single(self): self._check(Tensor([1.0]).float()) + def test_argsort(self): + t = Tensor([0.5, 0.1, 0.3]).float() + self.assertIs(_strip_unique(t.argsort().uop), _strip_unique(t.uop.argsort())) + def test_topk(self): + t = _t(2, 4).float() + tv, ti = t.topk(2) + uv, ui = t.uop.topk(2) + self.assertIs(_strip_unique(tv.uop), _strip_unique(uv)) + self.assertIs(_strip_unique(ti.uop), _strip_unique(ui)) + +class TestTensorUOpAllclose(unittest.TestCase): + def test_allclose(self): + a, b = _t(4).float(), _t(4).float() + self.assertIs(_strip_unique(a.allclose(b).uop), _strip_unique(a.uop.allclose(b.uop))) + class TestTensorUOpGather(unittest.TestCase): def _check(self, t, dim, idx): self.assertIs(_strip_unique(t.gather(dim, idx).uop), _strip_unique(t.uop.gather(dim, idx.uop))) @@ -138,6 +163,48 @@ class TestTensorUOpLoss(unittest.TestCase): t, Y = _t(2, 3).float(), Tensor([1, 2], dtype=dtypes.int32) self.assertIs(_strip_unique(t.sparse_categorical_crossentropy(Y, ignore_index=0).uop), _strip_unique(t.uop.sparse_categorical_crossentropy(Y.uop, ignore_index=0))) + def test_nll_loss(self): + t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32) + self.assertIs(_strip_unique(t.nll_loss(Y).uop), _strip_unique(t.uop.nll_loss(Y.uop))) + def test_nll_loss_weight(self): + t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float() + self.assertIs(_strip_unique(t.nll_loss(Y, weight=w).uop), _strip_unique(t.uop.nll_loss(Y.uop, weight=w.uop))) + def test_nll_loss_ignore_index(self): + t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32) + self.assertIs(_strip_unique(t.nll_loss(Y, ignore_index=1).uop), _strip_unique(t.uop.nll_loss(Y.uop, ignore_index=1))) + def test_nll_loss_none_reduction(self): + t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32) + self.assertIs(_strip_unique(t.nll_loss(Y, reduction="none").uop), _strip_unique(t.uop.nll_loss(Y.uop, reduction="none"))) + def test_nll_loss_weight_ignore_index(self): + t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float() + self.assertIs(_strip_unique(t.nll_loss(Y, weight=w, ignore_index=1).uop), + _strip_unique(t.uop.nll_loss(Y.uop, weight=w.uop, ignore_index=1))) + +class TestTensorUOpScatter(unittest.TestCase): + def test_scatter(self): + x, idx, src = _t(3, 4).float(), Tensor([[0, 1, 2, 0]], dtype=dtypes.int32), _t(1, 4).float() + self.assertIs(_strip_unique(x.scatter(0, idx, src).uop), _strip_unique(x.uop.scatter(0, idx.uop, src.uop))) + def test_scatter_scalar_src(self): + x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32) + self.assertIs(_strip_unique(x.scatter(1, idx, 3.14).uop), _strip_unique(x.uop.scatter(1, idx.uop, 3.14))) + # inf cannot be cast to int — this regresses if scalar src is routed through index.dtype first + def test_scatter_inf_src(self): + x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32) + self.assertIs(_strip_unique(x.scatter(1, idx, float("inf")).uop), + _strip_unique(x.uop.scatter(1, idx.uop, float("inf")))) + def test_scatter_add(self): + x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32) + self.assertIs(_strip_unique(x.scatter(1, idx, 3.14, reduce="add").uop), + _strip_unique(x.uop.scatter(1, idx.uop, 3.14, reduce="add"))) + def test_scatter_multiply(self): + x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32) + self.assertIs(_strip_unique(x.scatter(1, idx, 3.14, reduce="multiply").uop), + _strip_unique(x.uop.scatter(1, idx.uop, 3.14, reduce="multiply"))) + # tensor src with reduce hits the "elif reduce: raise" branch in both Tensor and UOp paths + def test_scatter_tensor_src_with_reduce_raises(self): + x, idx, src = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32), _t(1, 2).float() + with self.assertRaises(TypeError): x.scatter(1, idx, src, reduce="add") + with self.assertRaises(TypeError): x.uop.scatter(1, idx.uop, src.uop, reduce="add") class TestTensorUOpScatterReduce(unittest.TestCase): def _check(self, x, idx, src, **kw): @@ -176,6 +243,17 @@ class TestTensorUOpCat(unittest.TestCase): def test_cat_3tensors(self): _check(self, _t(2, 3), lambda x: x.cat(x, x, dim=0)) def test_cat_neg_dim(self): _check(self, _t(2, 3, 4), lambda x: x.cat(x, dim=-1)) +class TestTensorUOpPad(unittest.TestCase): + def test_pad_flat(self): _check(self, _t(4, 5), lambda x: x.pad((1, 2, 0, 3))) + def test_pad_flat_negative(self): _check(self, _t(4, 5), lambda x: x.pad((1, -1, 0, 2), value=-1.0)) + def test_pad_grouped_none(self): _check(self, _t(4, 5), lambda x: x.pad((None, (0, 3)))) + def test_pad_circular(self): _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="circular")) + def test_pad_circular_zero_after(self):_check(self, _t(4, 5), lambda x: x.pad(((1, 0), (2, 0)), mode="circular")) + def test_pad_reflect(self): _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="reflect")) + def test_pad_reflect_negative(self): _check(self, _t(4, 5), lambda x: x.pad(((1, -1), (0, 2)), mode="reflect")) + def test_pad_replicate(self): _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="replicate")) + def test_pad_replicate_negative(self): _check(self, _t(4, 5), lambda x: x.pad(((1, -1), (0, 2)), mode="replicate")) + class TestTensorUOpStack(unittest.TestCase): def test_stack_dim0(self): _check(self, _t(2, 3), lambda x: x.stack(x, dim=0)) def test_stack_dim1(self): _check(self, _t(2, 3), lambda x: x.stack(x, dim=1)) diff --git a/test/null/test_tensor_uop_representation.py b/test/null/test_tensor_uop_representation.py index f393ecd2ff..24d2161d86 100644 --- a/test/null/test_tensor_uop_representation.py +++ b/test/null/test_tensor_uop_representation.py @@ -16,7 +16,7 @@ class TestTensorMutates(unittest.TestCase): pa = a.uop pb = b.uop pr = ret.uop - ret.schedule() + ret.schedule_linear() self.assertIsNot(pa, a.uop) self.assertIsNot(pb, b.uop) self.assertIsNot(pr, ret.uop) diff --git a/test/null/test_tinyfs.py b/test/null/test_tinyfs.py index 1f2457e2fc..aeb215a9fb 100644 --- a/test/null/test_tinyfs.py +++ b/test/null/test_tinyfs.py @@ -5,22 +5,22 @@ class TestLoadStore(unittest.TestCase): def test_load_shape(self): t = Tensor(bytes(16)).fs_load(1024) assert t.shape == (1024,), t.shape - t.schedule() + t.schedule_linear() def test_store_shape(self): t = Tensor.zeros(1024).fs_store() assert t.shape == (16,), t.shape - t.schedule() + t.schedule_linear() def test_load_large_shape(self): t = Tensor(bytes(16)).fs_load(10_000_000) assert t.shape == (10_000_000,), t.shape - t.schedule() + t.schedule_linear() def test_store_large_shape(self): t = Tensor.zeros(10_000_000).fs_store() assert t.shape == (16,), t.shape - t.schedule() + t.schedule_linear() if __name__ == "__main__": unittest.main() diff --git a/test/null/test_uops.py b/test/null/test_uops.py index fe1bd0e6ca..9b20ead8c5 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -228,7 +228,7 @@ class TestUOpMethod(unittest.TestCase): a = UOp.variable("a", 1, 10) uop_var = Tensor(a.bind(1)) st_var = Tensor.empty((2, 10))[:, :a.bind(1)] - _, var_vals = (uop_var+st_var).schedule_with_vars() + _, var_vals = (uop_var+st_var).linear_with_vars() self.assertEqual(len(var_vals), 1) self.assertEqual(list(var_vals)[0], a.expr) diff --git a/test/null/test_uops_stats.py b/test/null/test_uops_stats.py index 2d46066f1f..0cf729cc58 100644 --- a/test/null/test_uops_stats.py +++ b/test/null/test_uops_stats.py @@ -1,8 +1,8 @@ import unittest from tinygrad import Tensor from tinygrad.helpers import GlobalCounters, DEV -from tinygrad.engine.realize import get_program -from tinygrad.renderer import ProgramSpec +from tinygrad.engine.realize import compile_linear, estimate_uop +from tinygrad.codegen import to_program from tinygrad.renderer import Estimates from tinygrad.uop.ops import Ops, UOp from tinygrad.dtype import dtypes @@ -18,8 +18,8 @@ def flops_mem(uops, ignore_indexing=False): # **************** new FlopCounter **************** def get_stats(x:Tensor): - si = x.schedule()[-1].lower() - return si.prg.estimates.ops, si.prg.estimates.mem + est = estimate_uop(compile_linear(x.schedule_linear()).src[-1]) + return est.ops, est.mem @unittest.skipIf(Device.DEFAULT == "WEBGPU", "webgpu does extra load/store for packed types") class TestMemoryCount(unittest.TestCase): @@ -165,75 +165,75 @@ N = 64 class TestStatsOptimized(unittest.TestCase): @classmethod def setUpClass(cls): - cls.ast_gemm = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1].ast - cls.ast_reduce = (Tensor.empty(N*N).sum()).schedule()[-1].ast + cls.ast_gemm = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0] + cls.ast_reduce = (Tensor.empty(N*N).sum()).schedule_linear().src[-1].src[0] - def check_gemm(self, p:ProgramSpec, extra_flops=0): - #p.uops.print() - #print(p.src) - print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds) - self.assertEqual(p.estimates.ops, 2*N*N*N + extra_flops) # N**3 mulaccs - self.assertEqual(p.estimates.mem, 3*N*N*4) # 3 NxN mats with floats + def check_gemm(self, p:UOp, extra_flops=0): + est = p.src[0].arg.estimates + print(p.arg.name, est.ops, est.mem, est.lds) + self.assertEqual(est.ops, 2*N*N*N + extra_flops) # N**3 mulaccs + self.assertEqual(est.mem, 3*N*N*4) # 3 NxN mats with floats def test_gemm(self): - p = get_program(replace_opts(self.ast_gemm, []), renderer=Device[Device.DEFAULT].renderer) + p = to_program(replace_opts(self.ast_gemm, []), renderer=Device[Device.DEFAULT].renderer) self.check_gemm(p) - self.assertEqual(p.estimates.lds, 2*N*N*N*4 + 4*N*N) + self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4 + 4*N*N) def test_gemm_tc_unroll(self): try: - p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]), + p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]), renderer=Device[Device.DEFAULT].renderer) except KernelOptError: raise unittest.SkipTest("no tensor cores") - print(p.src) + print(p.src[3].arg) self.check_gemm(p) # this is a good lesson about why UPCASTing is a good idea def test_gemm_one_upcasted(self): - p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4)]), renderer=Device[Device.DEFAULT].renderer) + p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4)]), renderer=Device[Device.DEFAULT].renderer) self.check_gemm(p) - self.assertEqual(p.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N) + self.assertEqual(p.src[0].arg.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N) def test_gemm_upcasted(self): - p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)]), + p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)]), renderer=Device[Device.DEFAULT].renderer) self.check_gemm(p) - self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N) + self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4//4 + 4*N*N) def test_gemm_upcasted_locals(self): try: - p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 0, 4), + p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4)]), renderer=Device[Device.DEFAULT].renderer) except KernelOptError: raise unittest.SkipTest("no locals") self.check_gemm(p) - self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N) + self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4//4 + 4*N*N) def test_gemm_group(self): try: - p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.GROUP, 0, 4)]), renderer=Device[Device.DEFAULT].renderer) + p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.GROUP, 0, 4)]), renderer=Device[Device.DEFAULT].renderer) except KernelOptError: raise unittest.SkipTest("no locals") SZ = N*N*4 # NOTE: these are sort of wrong. they aren't honoring the IF statement self.check_gemm(p, extra_flops=SZ*4) - self.assertEqual(p.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4) + self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4) def test_reduce(self): - p = get_program(replace_opts(self.ast_reduce, []), renderer=Device[Device.DEFAULT].renderer) - print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds) - self.assertEqual(p.estimates.ops, N*N) - self.assertEqual(p.estimates.mem, N*N*4 + 4) + p = to_program(replace_opts(self.ast_reduce, []), renderer=Device[Device.DEFAULT].renderer) + est = p.src[0].arg.estimates + print(p.arg.name, est.ops, est.mem, est.lds) + self.assertEqual(est.ops, N*N) + self.assertEqual(est.mem, N*N*4 + 4) def test_reduce_group(self): try: - p = get_program(replace_opts(self.ast_reduce, [Opt(OptOps.GROUP, 0, 50)]), renderer=Device[Device.DEFAULT].renderer) + p = to_program(replace_opts(self.ast_reduce, [Opt(OptOps.GROUP, 0, 50)]), renderer=Device[Device.DEFAULT].renderer) except KernelOptError: raise unittest.SkipTest("no locals") - # NOTE: these are wrong, they don't respect the if statement - print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds) + est = p.src[0].arg.estimates + print(p.arg.name, est.ops, est.mem, est.lds) if __name__ == '__main__': unittest.main(verbosity=2) diff --git a/test/null/test_viz.py b/test/null/test_viz.py index aa7ddfc36a..90efc84e35 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -13,6 +13,7 @@ from tinygrad.device import Buffer from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, active_group, _name_cnt, RewriteTrace from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData from tinygrad.codegen import to_program_cache +from tinygrad.codegen import to_program @track_rewrites(name=True) def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=None) -> UOp: @@ -320,27 +321,27 @@ class TestVizGC(unittest.TestCase): # VIZ integrates with other parts of tinygrad from tinygrad import Tensor, Device -from tinygrad.engine.realize import get_program +from tinygrad.engine.realize import get_runner class TestVizIntegration(unittest.TestCase): # codegen supports rendering of code blocks def test_codegen_tracing(self): with save_viz() as viz: - ast = Tensor.schedule(Tensor.empty(4)+Tensor.empty(4))[0].ast - prg = get_program(ast, Device[Device.DEFAULT].renderer) + ast = (Tensor.empty(4)+Tensor.empty(4)).schedule_linear().src[0].src[0] + prg = to_program(ast, Device[Device.DEFAULT].renderer) lst = viz.list_items() self.assertEqual(len(lst), 3) self.assertEqual(lst[0]["name"], "Callify 1 Buffer n1") self.assertEqual(lst[1]["name"], "Schedule 1 Kernel n1") - self.assertEqual(lst[2]["name"], prg.name) + self.assertEqual(lst[2]["name"], prg.arg.name) # schedule graph CALL nodes have a link to jump to codegen def test_link_sched_codegen(self): with save_viz() as viz: c1 = Tensor.empty(4).add(1) c2 = Tensor.empty(8).add(1) - sched = Tensor.schedule(c1, c2) - prgs = [get_program(si.ast, Device[Device.DEFAULT].renderer).name for si in sched] + sched = c1.schedule_linear(c2) + prgs = [to_program(si.src[0], Device[Device.DEFAULT].renderer).arg.name for si in sched.src] lst = viz.list_items() sched_idx = next(i for i,l in enumerate(lst) if l["name"].startswith("Schedule")) viz_kernel = next(i for i,s in enumerate(lst[sched_idx]["steps"]) if s["name"] == "View Kernel Graph") @@ -356,7 +357,7 @@ class TestVizIntegration(unittest.TestCase): a = Tensor.empty(1) b = Tensor.empty(1) metadata = (alu:=a+b).uop.metadata - alu.schedule() + alu.schedule_linear() graph = next(viz.get_details(0, 0))["graph"] self.assertEqual(len([n for n in graph.values() if repr(metadata) in n["label"]]), 1) @@ -722,10 +723,10 @@ class TestCfg(unittest.TestCase): gidx = UOp.special(1, "gidx0") sink = UOp.sink(out.base, lidx, gidx, arg=KernelInfo(name=name)) return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="NULL"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts])))) - with Context(DEV=f"NULL:HIP:{self.arch}"): + with Context(DEV=f"NULL::{self.arch}"): out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0] - prg = out.schedule()[-1].lower().prg.p - return amdgpu_cfg(prg.lib, self.arch) + runner = get_runner(out.device, out.schedule_linear().src[-1].src[0]) + return amdgpu_cfg(runner.prg.src[4].arg, self.arch) def test_simple(self): k = Kernel(arch=self.arch) @@ -924,7 +925,7 @@ class TestCLI(unittest.TestCase): (p:=Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events)) # reconstruct DEBUG=4 output and see all markers. with Context(DEBUG=4): - kernels = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-p", "-s", "NULL") + kernels = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-s", "NULL") self.assertIn("void custom_empty_n0", kernels) self.assertIn("marker @ 1", kernels) self.assertIn("void custom_empty_n1", kernels) @@ -933,11 +934,11 @@ class TestCLI(unittest.TestCase): self.assertIn("UOp.const", kernels) # get the top slowest functions across all devices with Context(DEBUG=2): - times = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-p", "-s", "ALL", "--top", "-1") + times = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-s", "ALL", "--top", "-1") self.assertIn("TINY", times) self.assertIn("NULL", times) with Context(DEBUG=3): - json_lines = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-p", "-s", "ALL", "--json") + json_lines = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-s", "ALL", "--json") for line in json_lines.split("\n"): _ = json.loads(line) if __name__ == "__main__": diff --git a/test/null/test_winograd.py b/test/null/test_winograd.py index 7e1a24d4c7..8e3402e1f2 100644 --- a/test/null/test_winograd.py +++ b/test/null/test_winograd.py @@ -18,14 +18,14 @@ class TestWinograd(unittest.TestCase): def test_forward_kernels(self): x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize() out = Tensor.conv2d(x,w) - self.assertEqual(len(out.schedule()), 2) + self.assertEqual(len(out.schedule_linear().src), 2) def test_backward_kernels(self): x,w = Tensor.empty(1,4,9,9,requires_grad=True).realize(), Tensor.empty(4,4,3,3,requires_grad=True).realize() out = Tensor.conv2d(x,w, padding=1) out.mean().backward() - backward_schedule = Tensor.schedule(x.grad, w.grad) - self.assertEqual(len(backward_schedule), 4) + backward_schedule = x.grad.schedule_linear(w.grad) + self.assertEqual(len(backward_schedule.src), 4) def test_counters(self): IC, OC, X, Y = 4,4,9,9 diff --git a/test/opt/test_gen_float4.py b/test/opt/test_gen_float4.py index 61f66f2ecf..66c3c848ad 100644 --- a/test/opt/test_gen_float4.py +++ b/test/opt/test_gen_float4.py @@ -1,8 +1,9 @@ import unittest from tinygrad import Device, Tensor, dtypes from tinygrad.uop.ops import UOp, Ops +from tinygrad.codegen import to_program from tinygrad.codegen.opt import Opt, OptOps -from tinygrad.engine.realize import get_program + from tinygrad.helpers import DEV from test.helpers import replace_opts @@ -24,12 +25,12 @@ class TestFloat4(unittest.TestCase): b = Tensor.empty(2, 8).realize() c = a + b - s = c.schedule()[0] - realized_ast = s.ast + s = c.schedule_linear().src[0] + realized_ast = s.src[0] opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)] - program = get_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer) + program = to_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer) - assert TestFloat4.count_float4(program.uops) == (2, 1) + assert TestFloat4.count_float4(tuple(program.src[2].src)) == (2, 1) @unittest.skipIf(Device.DEFAULT in {"CPU"} and AMX, "CPU with AMX upcasts float up to size 16") def test_float4_multidim(self): @@ -37,9 +38,9 @@ class TestFloat4(unittest.TestCase): b = Tensor.empty(2, 8).realize() c = a + b - s = c.schedule()[0] - uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]), - renderer=Device[Device.DEFAULT].renderer).uops + s = c.schedule_linear().src[0] + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]), + renderer=Device[Device.DEFAULT].renderer).src[2].src) assert TestFloat4.count_float4(uops) == (4, 2) @unittest.skipUnless(Device.DEFAULT in {"CPU"} and AMX, "Only CPU with AMX upcasts float up to size 16") @@ -49,9 +50,9 @@ class TestFloat4(unittest.TestCase): b = Tensor.empty(2, size).realize() c = a + b - s = c.schedule()[0] - return get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]), - renderer=Device[Device.DEFAULT].renderer).uops + s = c.schedule_linear().src[0] + return tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]), + renderer=Device[Device.DEFAULT].renderer).src[2].src) sizes = [12, 8, 16] shifts = [3, 2, 4] @@ -66,12 +67,12 @@ class TestFloat4(unittest.TestCase): b = Tensor.empty(9).realize().shrink(((1, 9),)) c = a + b - s = c.schedule()[0] - realized_ast = s.ast + s = c.schedule_linear().src[0] + realized_ast = s.src[0] opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)] - program = get_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer) + program = to_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer) - assert TestFloat4.count_float4(program.uops) == (0, 1) + assert TestFloat4.count_float4(tuple(program.src[2].src)) == (0, 1) @unittest.skipIf(Device.DEFAULT in {"CPU"} and AMX, "CPU with AMX upcasts float up to size 16") def test_float4_multidim_unaligned_load(self): @@ -79,9 +80,9 @@ class TestFloat4(unittest.TestCase): b = Tensor.empty(2, 9).realize().shrink(((0, 2), (1, 9),)) c = a + b - s = c.schedule()[0] - uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]), - renderer=Device[Device.DEFAULT].renderer).uops + s = c.schedule_linear().src[0] + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]), + renderer=Device[Device.DEFAULT].renderer).src[2].src) assert TestFloat4.count_float4(uops) == (0, 2) @@ -92,9 +93,9 @@ class TestFloat4(unittest.TestCase): b = Tensor.empty(2, size).realize().shrink(((0, 2), (1, size),)) c = a + b - s = c.schedule()[0] - return get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]), - renderer=Device[Device.DEFAULT].renderer).uops + s = c.schedule_linear().src[0] + return tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]), + renderer=Device[Device.DEFAULT].renderer).src[2].src) sizes = [13, 9, 17] shifts = [3, 2, 4] @@ -111,8 +112,8 @@ class TestFloat4(unittest.TestCase): # only the first and last conv dot products are aligned in a, and b is never aligned, so no # float4 should be emitted (the reduce axis of size 4 is the float4 axis here) - s = c.schedule()[0] - uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UNROLL, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops + s = c.schedule_linear().src[0] + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UNROLL, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[2].src) assert TestFloat4.count_float4(uops) == (0, 0) @@ -125,9 +126,9 @@ class TestFloat4(unittest.TestCase): # don't. # UPDATE: now we do this fusion - s = c.schedule()[0] - uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), - renderer=Device[Device.DEFAULT].renderer).uops + s = c.schedule_linear().src[0] + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), + renderer=Device[Device.DEFAULT].renderer).src[2].src) assert TestFloat4.count_float4(uops) in {(0,1), (1,1)} @@ -139,8 +140,8 @@ class TestFloat4(unittest.TestCase): # we will upcast the top axis of sz 4. they should not be coalesced into float4, # since the top axis is not contiguous. - s = c.schedule()[0] - uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops + s = c.schedule_linear().src[0] + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[2].src) assert TestFloat4.count_float4(uops) == (0, 1) @@ -151,8 +152,8 @@ class TestFloat4(unittest.TestCase): # should float4 b but not a - s = c.schedule()[0] - uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops + s = c.schedule_linear().src[0] + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[2].src) assert TestFloat4.count_float4(uops) == (1, 1) diff --git a/test/opt/test_tensor_cores.py b/test/opt/test_tensor_cores.py index 2d04574441..669afbbfd9 100644 --- a/test/opt/test_tensor_cores.py +++ b/test/opt/test_tensor_cores.py @@ -1,6 +1,5 @@ import numpy as np import unittest -from dataclasses import replace from tinygrad import Device, Tensor, dtypes from tinygrad.tensor import _to_np_dtype @@ -9,14 +8,15 @@ from tinygrad.dtype import DType from tinygrad.device import is_dtype_supported from tinygrad.helpers import DEV, Context from test.helpers import slow, replace_opts -from tinygrad.engine.realize import CompiledRunner, get_program +from tinygrad.engine.realize import CompiledRunner +from tinygrad.codegen import to_program from tinygrad.codegen.opt import Opt, OptOps, KernelOptError from tinygrad.codegen.opt.tc import amd_cdna_1616128 # TODO: write a clean version of this from test.backend.test_linearizer import helper_realized_ast, helper_linearizer_opt -# NOTE: get_program always passes in Device[Device.DEFAULT].renderer explicitly for process_replay!!! +# NOTE: to_program always passes in Device[Device.DEFAULT].renderer explicitly for process_replay!!! AMX = "AMX" in DEV.arch @@ -24,19 +24,19 @@ def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, d ensure_triggered:bool=True): a, b = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in) r = a.matmul(b, dtype=dtype_out) - sched = r.schedule() - realized_ast = sched[-1].ast + sched = r.schedule_linear() + realized_ast = sched.src[-1].src[0] opts_to_apply = [Opt(OptOps.TC, axis, (tc_select, tc_opt, 1))] if ensure_triggered: - program = get_program(replace_opts(realized_ast, opts_to_apply), Device[Device.DEFAULT].renderer) - wmmas = len([uop for uop in program.uops if uop.op is Ops.WMMA]) - tcs = len([x for x in program.applied_opts if x.op is OptOps.TC]) + program = to_program(replace_opts(realized_ast, opts_to_apply), Device[Device.DEFAULT].renderer) + wmmas = len([uop for uop in tuple(program.src[2].src) if uop.op is Ops.WMMA]) + tcs = len([x for x in program.src[0].arg.applied_opts if x.op is OptOps.TC]) assert wmmas > 0, "tensor core not triggered" assert tcs == 1, "tensor core opt not included" else: try: - program = get_program(replace_opts(realized_ast, opts_to_apply), Device[Device.DEFAULT].renderer) + program = to_program(replace_opts(realized_ast, opts_to_apply), Device[Device.DEFAULT].renderer) assert False, "OptOps.TC triggered, expected KernelOptError" except KernelOptError: pass @@ -47,9 +47,10 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi if dtype_in == dtypes.bfloat16: r = r.float() realized_ast, bufs = helper_realized_ast(r) opts = [Opt(op=OptOps.TC, axis=axis, arg=(tc_select, tc_opt, use_tensor_cores))] - prg = CompiledRunner(replace(get_program(replace_opts(realized_ast, opts), Device[Device.DEFAULT].renderer), device=Device.DEFAULT)) - if use_tensor_cores == 1: assert len([uop for uop in prg.p.uops if uop.op is Ops.WMMA]) > 0, "wmma not triggered" - assert len([x for x in prg.p.uops[-1].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included" + pu = to_program(replace_opts(realized_ast, opts), Device[Device.DEFAULT].renderer) + if use_tensor_cores == 1: assert len([uop for uop in pu.src[2].src if uop.op is Ops.WMMA]) > 0, "wmma not triggered" + assert len([x for x in pu.src[0].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included" + prg = CompiledRunner(pu, Device.DEFAULT) prg.exec(bufs) if dtype_in == dtypes.half: tc_atol, tc_rtol = 1e-2, 1e-3 elif dtype_in == dtypes.bfloat16: tc_atol, tc_rtol = (1e-1, 2e-2) if dtype_out == dtypes.bfloat16 else (1e-2, 1e-2) @@ -76,15 +77,16 @@ class TestTensorCores(unittest.TestCase): n, m, k = tc.dims[0], tc.dims[1], 2 if AMX else tc.dims[2] a, b = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in) r = a.matmul(b, dtype=tc.dtype_out) - prg = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))]), Device[Device.DEFAULT].renderer) + prg = to_program(replace_opts(r.schedule_linear().src[-1].src[0], + [Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))]), Device[Device.DEFAULT].renderer) if Device.DEFAULT == "CPU" and DEV.renderer == "LLVM": - assert "0x201000" in prg.src + assert "0x201000" in prg.src[3].arg elif Device.DEFAULT == "AMD" and DEV.renderer == "LLVM": - assert "@llvm.amdgcn.wmma" in prg.src + assert "@llvm.amdgcn.wmma" in prg.src[3].arg elif Device[Device.DEFAULT].renderer.suffix == "PTX": - assert "mma.sync.aligned" in prg.src + assert "mma.sync.aligned" in prg.src[3].arg else: - assert "__WMMA_" in prg.src + assert "__WMMA_" in prg.src[3].arg @Context(ALLOW_TF32=1) @unittest.skipIf((Device.DEFAULT == "AMD") or (Device.DEFAULT == "PYTHON" and Device.default.renderer.target.device == "AMD"), "broken for AMD") @@ -143,11 +145,11 @@ class TestTensorCores(unittest.TestCase): c = a.conv2d(b, padding=1, dtype=tc.dtype_out) realized_ast, real_bufs = helper_realized_ast(c) - program = get_program(replace_opts(realized_ast, [Opt(OptOps.TC, axis, (-1, 2, 1))]), Device[Device.DEFAULT].renderer) - assert len([uop for uop in program.uops if uop.op is Ops.WMMA]) > 0, "tensor core not triggered" - assert len([x for x in program.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included" + program = to_program(replace_opts(realized_ast, [Opt(OptOps.TC, axis, (-1, 2, 1))]), Device[Device.DEFAULT].renderer) + assert len([uop for uop in tuple(program.src[2].src) if uop.op is Ops.WMMA]) > 0, "tensor core not triggered" + assert len([x for x in program.src[0].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included" - prg = CompiledRunner(program) + prg = CompiledRunner(program, Device.DEFAULT) # TODO: support this even if numpy doesn't if _to_np_dtype(real_bufs[0].dtype) is None: continue real_bufs[0].copyin(np.zeros((real_bufs[0].size, ), dtype=_to_np_dtype(real_bufs[0].dtype)).data) # Zero to check that all values are filled @@ -167,7 +169,7 @@ class TestTensorCores(unittest.TestCase): r = x.matmul(y, dtype=tc.dtype_out) opts = [Opt(OptOps.UNROLL, 0, 4)] ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3) - for u in get_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).uops: + for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[2].src): if u.op is Ops.WMMA: assert u.src[-1].src[0].op != Ops.STORE @@ -181,7 +183,7 @@ class TestTensorCores(unittest.TestCase): r = x.matmul(y, dtype=tc.dtype_out) opts = [Opt(OptOps.UNROLL, 0, 4)] ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3) - for u in get_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).uops: + for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[2].src): if u.op is Ops.WMMA: #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2])) assert u.src[-1].src[0].op != Ops.STORE @@ -197,7 +199,7 @@ class TestTensorCores(unittest.TestCase): r = x.matmul(y, dtype=tc.dtype_out).relu() opts = [Opt(OptOps.UNROLL, 0, 4)] ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3) - for u in get_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).uops: + for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[2].src): if u.op is Ops.WMMA: #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2])) assert u.src[-1].src[0].op != Ops.STORE diff --git a/test/testextra/test_export_model.py b/test/testextra/test_export_model.py index dc246c526b..8b87ba9020 100644 --- a/test/testextra/test_export_model.py +++ b/test/testextra/test_export_model.py @@ -2,6 +2,8 @@ import unittest from extra.export_model import export_model, EXPORT_SUPPORTED_DEVICE from tinygrad.tensor import Tensor from tinygrad.device import Device +from tinygrad.nn import Linear +from tinygrad.nn.state import get_state_dict from tinygrad import dtypes import json @@ -66,5 +68,15 @@ class TextModelExportWebGPU(unittest.TestCase): self.assertIn(f"const resultBuffer{i} = new {expected_arr_prefix}Array(gpuReadBuffer{i}.size/{dt.itemsize});", prg) self.assertIn(f"resultBuffer{i}.set(new {expected_arr_prefix}Array(gpuReadBuffer{i}.getMappedRange()));", prg) + def test_weights_bound_to_safetensor(self): + # regression test: every weight ended up as createEmptyBuf (zero-init) instead of createWeightBuf + class MyModel: + def __init__(self): self.fc1, self.fc2 = Linear(4, 8), Linear(8, 2) + def forward(self, x): return self.fc2(self.fc1(x).relu()) + model = MyModel() + for t in get_state_dict(model).values(): t.realize() + prg, _, _, _ = export_model(model, "webgpu", Tensor.randn(1, 4)) + self.assertEqual(prg.count("createWeightBuf("), len(get_state_dict(model))) + if __name__ == '__main__': unittest.main() diff --git a/test/testextra/test_tk.py b/test/testextra/test_tk.py index 3cccdb9ad3..a24af1e89d 100644 --- a/test/testextra/test_tk.py +++ b/test/testextra/test_tk.py @@ -1,9 +1,8 @@ import unittest, math, time -from tinygrad import Tensor, Device, dtypes, Context +from tinygrad import Tensor, Device, dtypes, Context, GlobalCounters from tinygrad.uop.ops import UOp, Ops -from tinygrad.engine.realize import get_runner -from tinygrad.schedule import ExecItem +from tinygrad.engine.realize import run_linear from tinygrad.engine.jit import TinyJit import numpy as np @@ -67,8 +66,9 @@ class TestTK(unittest.TestCase): c = Tensor.empty(1, 1, N, N, dtype="float32") Tensor.realize(a, b, c) - ei = ExecItem(sink, [t.uop.buffer for t in (c, a, b)], prg=get_runner(Device.DEFAULT, sink)) - for _ in range(5): ei.run(wait=True) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (c, a, b)]),)) + + for _ in range(5): run_linear(linear, do_update_stats=False) c = c.float() ref = a.matmul(b, dtype=dtypes.float32).float() @@ -115,8 +115,9 @@ class TestTK(unittest.TestCase): c = Tensor.empty(1, 1, N, N, dtype="float32") Tensor.realize(a, b, c) - ei = ExecItem(sink, [t.uop.buffer for t in (c, a, b)], prg=get_runner(Device.DEFAULT, sink)) - for _ in range(5): ei.run(wait=True) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (c, a, b)]),)) + + for _ in range(5): run_linear(linear, do_update_stats=False) c = c.float() ref = a.matmul(b.transpose(2, 3), dtype=dtypes.float32).float() @@ -151,8 +152,9 @@ class TestTK(unittest.TestCase): b = Tensor.empty(1, 1, N, N, dtype="float32") Tensor.realize(a, b) - ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink)) - for _ in range(5): ei.run(wait=True) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),)) + + for _ in range(5): run_linear(linear, do_update_stats=False) b = b.float() ref = a.float() @@ -190,8 +192,9 @@ class TestTK(unittest.TestCase): b = Tensor.empty(1, 1, N, N, dtype="float32") Tensor.realize(a, b) - ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink)) - for _ in range(5): ei.run(wait=True) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),)) + + for _ in range(5): run_linear(linear, do_update_stats=False) b = b.float() ref = a.float() @@ -232,8 +235,9 @@ class TestTK(unittest.TestCase): c = Tensor.empty(1, 1, N, N, dtype="float32") Tensor.realize(a, b, c) - ei = ExecItem(sink, [t.uop.buffer for t in (b, c, a)], prg=get_runner(Device.DEFAULT, sink)) - for _ in range(5): ei.run(wait=True) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, c, a)]),)) + + for _ in range(5): run_linear(linear, do_update_stats=False) b = b.float() c = c.float() @@ -272,8 +276,9 @@ class TestTK(unittest.TestCase): b = Tensor.empty(1, 1, N, N, dtype="float32") Tensor.realize(a, b) - ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink)) - for _ in range(5): ei.run(wait=True) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),)) + + for _ in range(5): run_linear(linear, do_update_stats=False) b = b.float() ref = a.float() @@ -309,8 +314,9 @@ class TestTK(unittest.TestCase): b = Tensor.empty(1, 1, N, N, dtype="float32") Tensor.realize(a, b) - ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink)) - for _ in range(5): ei.run(wait=True) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),)) + + for _ in range(5): run_linear(linear, do_update_stats=False) b = b.float() ref = a.float() + 1 @@ -354,8 +360,9 @@ class TestTK(unittest.TestCase): b = Tensor.empty(1, 1, N, N, dtype="float32") Tensor.realize(a, b) - ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink)) - for _ in range(5): ei.run(wait=True) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),)) + + for _ in range(5): run_linear(linear, do_update_stats=False) b = b.float() ref = a.float().max(axis=2, keepdim=True).expand(a.shape) @@ -399,8 +406,9 @@ class TestTK(unittest.TestCase): b = Tensor.empty(1, 1, N, M, dtype="float32") Tensor.realize(a, b) - ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink)) - for _ in range(5): ei.run(wait=True) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),)) + + for _ in range(5): run_linear(linear, do_update_stats=False) b = b.float() ref = a.float().max(axis=2, keepdim=True).expand(a.shape) @@ -444,8 +452,9 @@ class TestTK(unittest.TestCase): b = Tensor.empty(1, 1, N, N, dtype="float32") Tensor.realize(a, b) - ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink)) - for _ in range(5): ei.run(wait=True) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),)) + + for _ in range(5): run_linear(linear, do_update_stats=False) b = b.float() ref = a.float().sum(axis=2, keepdim=True).expand(a.shape) @@ -489,8 +498,9 @@ class TestTK(unittest.TestCase): b = Tensor.empty(1, 1, N, M, dtype="float32") Tensor.realize(a, b) - ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink)) - for _ in range(5): ei.run(wait=True) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),)) + + for _ in range(5): run_linear(linear, do_update_stats=False) b = b.float() ref = a.float().sum(axis=2, keepdim=True).expand(a.shape) @@ -549,8 +559,9 @@ class TestTK(unittest.TestCase): b = Tensor.empty(1, 1, BLOCK_SIZE, N, dtype="float32") Tensor.realize(a, b) - ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink)) - for _ in range(5): ei.run(wait=True) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),)) + + for _ in range(5): run_linear(linear, do_update_stats=False) b = b.float() ref = a.float().softmax(axis=3) @@ -609,8 +620,9 @@ class TestTK(unittest.TestCase): b = Tensor.empty(1, 1, N, BLOCK_SIZE, dtype="float32") Tensor.realize(a, b) - ei = ExecItem(sink, [t.uop.buffer for t in (b, a)], prg=get_runner(Device.DEFAULT, sink)) - for _ in range(5): ei.run(wait=True) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),)) + + for _ in range(5): run_linear(linear, do_update_stats=False) b = b.float() ref = a.float().softmax(axis=2) @@ -719,9 +731,11 @@ class TestTK(unittest.TestCase): out = Tensor.empty(B, N, H, D, dtype=dtypes.bfloat16) Tensor.realize(q, k, v, out) - ei = ExecItem(sink, [t.uop.buffer for t in (out, q, k, v)], prg=get_runner(Device.DEFAULT, sink)) + linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (out, q, k, v)]),)) for _ in range(5): - et = ei.run(wait=True) + GlobalCounters.reset() + with Context(DEBUG=2): run_linear(linear) + et = GlobalCounters.time_sum_s attn_flops = 2 * B * H * N * N * D + \ 4 * B * H * N * N + \ 2 * B * H * N * N * D diff --git a/test/unit/test_allreduce.py b/test/unit/test_allreduce.py index e894268643..b24fe07c0b 100644 --- a/test/unit/test_allreduce.py +++ b/test/unit/test_allreduce.py @@ -9,9 +9,9 @@ class TestRingAllReduce(unittest.TestCase): N = 4 ds = tuple(f"CPU:{i}" for i in range(N)) t = Tensor.empty(N, N*100).shard(ds, axis=0).realize() - schedules = t.sum(0).schedule_with_vars()[0] - copies = [si for si in schedules if si.ast.op is Ops.COPY] - pairs = [(c.bufs[0].device, c.bufs[1].device) for c in copies] + linear = t.sum(0).linear_with_vars()[0] + copies = [si for si in linear.src if si.src[0].op is Ops.COPY] + pairs = [(c.src[1].buffer.device, c.src[2].buffer.device) for c in copies] # N*(N-1) scatter reduce, and N*(N-1) allgather self.assertEqual(len(pairs), N*(N-1)*2) # copy topology forms a ring @@ -30,8 +30,8 @@ class TestAllreduceCast(unittest.TestCase): ds = tuple(f"CPU:{i}" for i in range(2)) with Context(ALLREDUCE_CAST=allreduce_cast, RING=0, SCACHE=0): t = Tensor.empty(4, 4, dtype=dtype).shard(ds, axis=0) - schedules = t.sum(0).schedule_with_vars()[0] - return {si.bufs[0].dtype.scalar() for si in schedules if si.ast.op is Ops.COPY} + linear = t.sum(0).linear_with_vars()[0] + return {si.src[1].buffer.dtype.scalar() for si in linear.src if si.src[0].op is Ops.COPY} def test_allreduce_cast_bf16(self): # with ALLREDUCE_CAST, allreduce copies stay in bfloat16 instead of promoting to float32 diff --git a/test/unit/test_cpu.py b/test/unit/test_cpu.py index ea7668d4f2..2921bb9c93 100644 --- a/test/unit/test_cpu.py +++ b/test/unit/test_cpu.py @@ -3,19 +3,19 @@ from contextlib import redirect_stdout from tinygrad import Tensor, Device from tinygrad.helpers import Target from tinygrad.renderer.nir import LVPRenderer -from tinygrad.engine.realize import get_program +from tinygrad.codegen import to_program @unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU") class TestCPU(unittest.TestCase): def test_arch_feats(self): - ast = (Tensor.empty(16) + Tensor.empty(16)).schedule()[-1].ast + ast = (Tensor.empty(16) + Tensor.empty(16)).schedule_linear().src[-1].src[0] for ren in Device[Device.DEFAULT].renderers: for arch, expect_vmov in [("x86_64,x86-64,avx", True), ("x86_64,x86-64,-avx", False)]: with self.subTest(arch=arch): if ren is LVPRenderer: continue # LVP does not play nice with cross compilation r = ren(Target(device="CPU", arch=arch)) - p = get_program(ast, r) - lib = r.compiler.compile(p.src) + p = to_program(ast, r) + lib = r.compiler.compile(p.src[3].arg) out = io.StringIO() with redirect_stdout(out): r.compiler.disassemble(lib) self.assertEqual("vmov" in out.getvalue(), expect_vmov, out.getvalue()) diff --git a/test/unit/test_hcq_graph.py b/test/unit/test_hcq_graph.py index 9f7fb8d472..d8f0322a4b 100644 --- a/test/unit/test_hcq_graph.py +++ b/test/unit/test_hcq_graph.py @@ -20,9 +20,9 @@ class TestHCQUnit(unittest.TestCase): inp, inp_cpu = Tensor.randn(10, 10, device=Device.DEFAULT).realize(), Tensor.randn(10, 10, device="CPU").realize() for _ in range(5): f(inp, inp_cpu) - # construct minimal CALL UOps for supports_exec_item - gpu_call = UOp(Ops.SINK).call(UOp.new_buffer(Device.DEFAULT, 1, dtypes.float)) - cpu_call = UOp(Ops.SINK).call(UOp.new_buffer("CPU", 1, dtypes.float)) + # construct minimal CALL UOps for supports_exec_item (graphs only see PROGRAMs after compile_linear) + gpu_call = UOp(Ops.PROGRAM).call(UOp.new_buffer(Device.DEFAULT, 1, dtypes.float)) + cpu_call = UOp(Ops.PROGRAM).call(UOp.new_buffer("CPU", 1, dtypes.float)) gpu_devs = [d0] # local MMIO: GPU works alone and with CPU in batch (cpu_support=True) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 95042992fb..eb6cfce057 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -1,11 +1,11 @@ from typing import cast from dataclasses import replace -import itertools, weakref -from tinygrad.helpers import DISABLE_FAST_IDIV, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES -from tinygrad.helpers import TracingKey, Context, Target, panic -from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, pyrender +import itertools +from tinygrad.helpers import DISABLE_FAST_IDIV, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC +from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, Target, panic +from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, pyrender from tinygrad.uop.spec import type_verify, program_spec, kernel_spec -from tinygrad.renderer import Renderer, ProgramSpec, Estimates +from tinygrad.renderer import Renderer, Estimates from tinygrad.dtype import dtypes # import all pattern matchers here @@ -139,7 +139,8 @@ def do_assemble(ctx:Renderer, prg:UOp, lin:UOp) -> UOp: def do_render(ctx:Renderer, prg:UOp, lin:UOp) -> UOp: src = ctx.render(list(lin.src)) - return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=src),), arg=ctx.aux(list(lin.src)) if ctx.has_aux else prg.arg) + new_arg = replace(prg.arg, aux=tuple(ctx.aux(list(lin.src)))) if ctx.has_aux else prg.arg + return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=src),), arg=new_arg) def do_compile(ctx:Renderer, prg:UOp, source:UOp) -> UOp|None: lib = ctx.compiler.compile_cached(source.arg) @@ -170,17 +171,18 @@ 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" full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None, beam=ast.arg.beam) - prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.target.device))) + prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.target.device)), arg=ProgramInfo.from_sink(full_sink)) else: raise RuntimeError(f"can't call to_program on {ast.op}") + if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0])) prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render") if VIZ: graph_rewrite(prg, PatternMatcher([]), name="View Program") return prg -to_program_cache: weakref.WeakValueDictionary[tuple, UOp] = weakref.WeakValueDictionary() +to_program_cache: dict[tuple, UOp] = {} def to_program(ast:UOp, renderer:Renderer) -> UOp: - if ast.op is Ops.PROGRAM and len(ast.src) >= 5 and ast.src[4].op is Ops.BINARY: return ast - key = (ast.key, type(renderer), renderer.target, NOOPT.value, DEVECTORIZE.value, EMULATED_DTYPES.value) + if ast.op is Ops.PROGRAM and len(ast.src) >= 5 and ast.src[4].op is Ops.BINARY: + return ast if isinstance(ast.arg, ProgramInfo) else ast.replace(arg=ProgramInfo.from_sink(ast.src[0])) + config = (NOOPT, DEVECTORIZE, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32) + key = (ast.key, type(renderer), renderer.target, *[x.value for x in config]) if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer) return prg - -def get_program(ast:UOp, renderer:Renderer) -> ProgramSpec: return ProgramSpec.from_uop(to_program(ast, renderer)) diff --git a/tinygrad/codegen/opt/search.py b/tinygrad/codegen/opt/search.py index b92d431ea0..e431aebe12 100644 --- a/tinygrad/codegen/opt/search.py +++ b/tinygrad/codegen/opt/search.py @@ -1,14 +1,13 @@ import functools, math, time, multiprocessing, traceback, signal, atexit from dataclasses import replace -from tinygrad.uop.ops import sym_infer, AxisType, pyrender +from tinygrad.uop.ops import sym_infer, AxisType, pyrender, UOp, Ops from tinygrad.device import Device, Buffer, Compiler from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str, unwrap from tinygrad.helpers import IGNORE_BEAM_CACHE from tinygrad.codegen.opt import Opt, OptOps, KernelOptError from tinygrad.tensor import Tensor from tinygrad.engine.realize import CompiledRunner -from tinygrad.codegen import get_program -from tinygrad.renderer import ProgramSpec +from tinygrad.codegen import to_program from tinygrad.codegen.opt.postrange import Scheduler actions = [Opt(op=OptOps.UPCAST, axis=axis, arg=amt) for amt in [0,2,3,4,5,7] for axis in range(8)] @@ -35,20 +34,22 @@ def get_test_global_size(global_size, max_global_size, var_vals): break return test_global_size, input_size / prod(test_global_size) -def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[str, int], rawbufs:list[Buffer], early_stop:float|None=None, +def _time_program(prg:UOp, lib:bytes, var_vals:dict[str, int], rawbufs:list[Buffer], early_stop:float|None=None, allow_test_size:int=True, max_global_size:int|None=65536, clear_l2=False, cnt=3, name="test", dev_timeout=False) -> list[float]: timeout = int(early_stop * 1e3) if dev_timeout and early_stop is not None and early_stop < math.inf else None factor = 1 + info = prg.arg if allow_test_size and max_global_size is not None: - global_size, factor = get_test_global_size(p.global_size, max_global_size, var_vals) - p = replace(p, global_size=global_size) - try: car = CompiledRunner(replace(p, lib=lib)) + global_size, factor = get_test_global_size(info.global_size, max_global_size, var_vals) + prg = prg.replace(arg=replace(info, global_size=tuple(global_size))) + if len(prg.src) <= 4 or prg.src[4].op is not Ops.BINARY: prg = prg.replace(src=prg.src + (UOp(Ops.BINARY, arg=lib),)) + try: car = CompiledRunner(prg, prg.src[1].arg) except AssertionError: return [math.inf] * cnt tms = [] input_bufs = [rawbufs[i] for i in car.p.globals] for _ in range(cnt): if clear_l2: - if hasattr(dev:=Device[p.device], 'invalidate_caches'): dev.invalidate_caches() + if hasattr(dev:=Device[prg.src[1].arg], 'invalidate_caches'): dev.invalidate_caches() else: with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024,1024).contiguous().realize(do_update_stats=False) tms.append(unwrap(car(input_bufs, var_vals, wait=True, timeout=timeout))*factor) @@ -60,22 +61,22 @@ def timeout_handler(signum, frame): if DEBUG >= 2: print("*** BEAM COMPILE TIMEOUT") raise TimeoutException() -def _try_compile(x:tuple[int,Scheduler], compiler:Compiler) -> tuple[int, tuple[ProgramSpec, bytes, float]|None]: +def _try_compile(x:tuple[int,Scheduler], compiler:Compiler) -> tuple[int, tuple[UOp, bytes, float]|None]: if hasattr(signal, "alarm"): signal.signal(getattr(signal, 'SIGALRM'), timeout_handler) # set timeout signal.alarm(getenv("BEAM_TIMEOUT_SEC", 10)) ret = None try: - p = get_program(x[1].copy().get_optimized_ast(name_override="test"), x[1].ren) - assert p.uops is not None, "uop list wasn't generated?" - if len(p.uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 3000)) > 0: - if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many uops. {len(p.uops)=}, {uops_max=}") + prg = to_program(x[1].copy().get_optimized_ast(name_override="test"), x[1].ren) + uops = prg.src[2].src + if len(uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 3000)) > 0: + if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many uops. {len(uops)=}, {uops_max=}") raise RuntimeError("too many uops") st = time.perf_counter() - prog = p.lib if p.lib is not None else compiler.compile(p.src) + prog = prg.src[4].arg if len(prg.src) > 4 and prg.src[4].op is Ops.BINARY else compiler.compile(prg.src[3].arg) et = time.perf_counter() - st - ret = (p, prog, et) + ret = (prg, prog, et) except RuntimeError: if DEBUG >= 4: traceback.print_exc() except Exception as e: @@ -153,15 +154,16 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True least_compute_ops = math.inf for i,proc in (map(_compile_fn, enumerate(candidates)) if beam_pool is None else beam_pool.imap_unordered(_compile_fn, enumerate(candidates))): if proc is None: continue - p, lib, compile_et = proc + prg, lib, compile_et = proc if lib in seen_libs: continue # filter out kernels that use 1000x more compute than the smallest - least_compute_ops = min(this_compute_ops:=sym_infer(p.estimates.ops, var_vals), least_compute_ops) + estimates = prg.src[0].arg.estimates + least_compute_ops = min(this_compute_ops:=sym_infer(estimates.ops if estimates is not None else 0, var_vals), least_compute_ops) if least_compute_ops*1000 < this_compute_ops: if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too much compute. {this_compute_ops} when least is {least_compute_ops}") continue seen_libs.add(lib) - try: tms = _time_program(p, lib, var_vals, rawbufs, early_stop=beam[0][1]*3 if len(beam) else 1.0, + try: tms = _time_program(prg, lib, var_vals, rawbufs, early_stop=beam[0][1]*3 if len(beam) else 1.0, allow_test_size=allow_test_size, clear_l2=hasattr(dev, 'invalidate_caches'), dev_timeout=getenv("BEAM_DEV_TIMEOUT", 1)) except Exception as e: @@ -170,7 +172,7 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True raise timed.append((candidates[i], min(tms))) if BEAM_DEBUG > 1: - print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(unwrap(p.uops)):5d} uops", + print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(prg.src[2].src):5d} uops", f"{time_to_str(compile_et, w=12)} compile/{time_to_str(timed[-1][1], w=12)} run", f" {len(timed):4d}/{len(candidates):4d} {timed[-1][0].colored_shape()}") elif DEBUG >= 2: diff --git a/tinygrad/device.py b/tinygrad/device.py index 83f1dec41a..dc3da08789 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -6,6 +6,7 @@ import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re from tinygrad.helpers import BENCHMARKS, CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, 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, EMULATED_DTYPES, IMAGE, FLOAT16, TracingKey, size_to_str, Target, VIZ +from tinygrad.helpers import pluralize from tinygrad.dtype import DType, PtrDType, dtypes, _to_np_dtype if TYPE_CHECKING: from tinygrad.renderer import Renderer @@ -293,7 +294,7 @@ class Compiled: f"{self.device}_{rn}=1 is deprecated, use DEV={self.device}:{rn} or {self.device}_CC={rn} instead" t = DEV.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, target=t) + f"No renderer for {self.device} is available", self.cached_renderer, t) def count(self) -> int: """ @@ -378,23 +379,26 @@ def enumerate_devices_str() -> Generator[str, None, None]: from tinygrad import Tensor, Device for device in ALL_DEVICES: - compilers_results, any_works = [], False + ren_results, iface_results = [], [] try: d = Device[device] - default_renderer = d.renderer + for iface in [i for i in getattr(d, 'ifaces', []) if not i.__name__.startswith("MOCK")]: + try: + name = iface.__name__[:-5] + default_text, count = ("(default)", d.count()) if type(d.iface) is iface else (f"(DEV={name}+{device} to make default)", iface(d, 0).count) # type: ignore + iface_results.append(f"{colored('+', 'green')} {name}: {pluralize('device', count)} {default_text}") + except Exception as e: iface_results.append(f"{colored('-', 'red')} {iface.__name__[:-5]}: {e}") for r in d.renderers: try: - # d.renderer, d.compiler = r(), c() with Context(CACHELEVEL=0, DEV=f"{device}:{d._renderer_name(r)}"): test = (Tensor([1,2,3], device=device) * 2).tolist() if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]") - default_text = '(default)' if type(default_renderer) is type(d.renderer) else f'(DEV={device}:{d._renderer_name(r)} to make default)' - compilers_results.append(f"{colored('+', 'green')} {d._renderer_name(r)} {default_text}") - any_works = True - except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {d._renderer_name(r)}: {e}") - result = (colored('PASS', 'green') if any_works else f"{colored('FAIL', 'yellow')}") + ''.join([f'\n{" "*16} {x}' for x in compilers_results]) - except Exception as e: - result = f"{colored('FAIL', 'red')} {e}" - yield f"{'*' if device == Device.DEFAULT else ' '} {device:10s}: {result}" + default_text = '(default)' if type(d.renderer) is r else f'(DEV={device}:{d._renderer_name(r)} to make default)' + ren_results.append(f"{colored('+', 'green')} {d._renderer_name(r)} {default_text}") + except Exception as e: ren_results.append(f"{colored('-', 'red')} {d._renderer_name(r)}: {e}") + result = (colored('PASS', 'green') + ("\n"+" "*12+"interfaces:\n" if iface_results else "") + '\n'.join([" "*13+x for x in iface_results]) + + (("\n"+" "*12+"renderers:\n") + '\n'.join([" "*13+x for x in ren_results]) if len(ren_results) > 1 else "")) + except Exception as e: result = f"{colored('FAIL', 'red')} {e}" + yield f"{'*' if device == Device.DEFAULT else ' '} {device:8s}: {result}" if __name__ == "__main__": for s in enumerate_devices_str(): print(s) diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index f0006caf1e..7187c2ac42 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -5,10 +5,9 @@ from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, from tinygrad.device import Buffer, Compiled, Device, MultiBuffer from tinygrad.dtype import DType, dtypes from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, track_rewrites, graph_rewrite -from tinygrad.engine.realize import ExecItem, capturing, CompiledRunner, Runner, Estimates, compile_linear, run_linear, get_runner, graph_cache +from tinygrad.engine.realize import capturing, CompiledRunner, Runner, Estimates, compile_linear, run_linear, get_runner, graph_cache, estimate_uop from tinygrad.engine.realize import unwrap_multi, resolve_params from tinygrad.schedule.memory import memory_plan_rewrite, _collect_bufs -from tinygrad.schedule import linear_to_schedule from tinygrad.nn.state import get_parameters from tinygrad.schedule.rangeify import mop_cleanup from dataclasses import dataclass @@ -63,9 +62,7 @@ def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp: def _call_outs_ins(call:UOp) -> tuple[set[int], set[int]]: non_bind = [s for s in call.src[1:] if s.op is not Ops.BIND] ast = call.src[0] - if ast.op in (Ops.SINK, Ops.PROGRAM): - prg = get_runner(non_bind[0].device if isinstance(non_bind[0].device, str) else non_bind[0].device[0], call.src[0]) - return set(prg.p.outs), set(prg.p.ins) + if ast.op is Ops.PROGRAM: return set(ast.arg.outs), set(ast.arg.ins) if ast.op in (Ops.COPY, Ops.BUFFER_VIEW): return {0}, {1} if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return {0}, set(range(1, len(non_bind))) return set(), set() @@ -98,22 +95,9 @@ def _check_no_non_tensor_return(ret): def graph_class(dev): return dev.graph.func if isinstance(dev.graph, functools.partial) else dev.graph -def get_input_replace(jit_cache: list[ExecItem], input_buffers:list[Buffer]) -> dict[tuple[int, int], int]: - input_replace: dict[tuple[int, int], int] = {} - for j,ji in enumerate(jit_cache): - for i,a in enumerate(ji.bufs): - if a in input_buffers: input_replace[(j,i)] = input_buffers.index(a) - return input_replace - class GraphRunner(Runner): - def __init__(self, linear:UOp, input_buffers:list[Buffer], input_uops:tuple[UOp, ...]=()): + def __init__(self, linear:UOp, input_uops:tuple[UOp, ...]=()): self.linear = linear.src[0] - self.jit_cache = [ei.lower() for ei in linear_to_schedule(self.linear.substitute({p: input_uops[p.arg] for p in linear.src[1:]}))] - for ei in self.jit_cache: - for b in ei.bufs: - if b is not None: b.ensure_allocated() - self.input_replace = get_input_replace(self.jit_cache, input_buffers) if input_buffers else {} - self.calls: list[tuple[int, UOp, list[Buffer], dict[str, int]]] = [] self.progs: list[CompiledRunner|None] = [] self.uop_replace: list[list[tuple[int, int]]] = [] @@ -121,42 +105,37 @@ class GraphRunner(Runner): replace = [(p, b.arg) for p, b in enumerate(b for b in call.src[1:] if b.op is not Ops.BIND) if b.op is Ops.PARAM] for dev_idx, (bufs, device_vars) in enumerate(unwrap_multi(call, resolve_params(call, input_uops))): self.calls.append((dev_idx, call.src[0], [b.ensure_allocated() for b in bufs], device_vars)) - self.progs.append(get_runner(bufs[0].device, call.src[0]) if call.src[0].op in (Ops.SINK, Ops.PROGRAM) else None) + self.progs.append(get_runner(bufs[0].device, call.src[0]) if call.src[0].op is Ops.PROGRAM else None) self.uop_replace.append(replace) self.var_vals_replace:dict[int, list[tuple[int, int]]] = {} self.launch_dims_replace:dict[int, tuple[int|None, int|None]] = {} - self.launch_dims_base:dict[int, tuple[tuple[int, ...], tuple[int, ...]]] = {} + self.launch_dims_base:dict[int, tuple[tuple[int|float, ...], tuple[int, ...]]] = {} def is_sym_dim(dim) -> bool: return not all(isinstance(d, (int, float)) for d in dim) - crs = [(ji, ji.prg) for ji in self.jit_cache if isinstance(ji.prg, CompiledRunner)] - self.vars = sorted({v.expr for ji,p in crs for v in p.p.vars if v.expr not in ji.fixedvars | p.p.runtimevars}) - self.symbolic_dims = dedup([tuple(d) for _,p in crs if (d:=p.p.local_size) and is_sym_dim(d)] + - [tuple(d) for _,p in crs if (d:=p.p.global_size) and is_sym_dim(d)]) + crs = [(j, p, self.calls[j][3]) for j,p in enumerate(self.progs) if isinstance(p, CompiledRunner)] + self.vars = sorted({v.expr for _,p,dv in crs for v in p.p.vars if v.expr not in dv | p.p.runtimevars}) + self.symbolic_dims = dedup(tuple(d) for _,p,_ in crs for d in (p.p.local_size, p.p.global_size) if d and is_sym_dim(d)) def find_symbolic_dim(dim): return self.symbolic_dims.index(tuple(dim)) if dim is not None and tuple(dim) in self.symbolic_dims else None - estimates = Estimates() - for j,ji in enumerate(self.jit_cache): - assert ji.prg is not None - estimates += ji.prg.estimates - if isinstance(ji.prg, CompiledRunner): - if (replace:=[(i, self.vars.index(v.expr)) for i, v in enumerate(ji.prg.p.vars) if v.expr not in ji.fixedvars | ji.prg.p.runtimevars]): - self.var_vals_replace[j] = replace + for j,p,dv in crs: + if (replace:=[(i, self.vars.index(v.expr)) for i, v in enumerate(p.p.vars) if v.expr not in dv | p.p.runtimevars]): + self.var_vals_replace[j] = replace + global_dim_idx, local_dim_idx = find_symbolic_dim(p.p.global_size), find_symbolic_dim(p.p.local_size) + if global_dim_idx is not None or local_dim_idx is not None: + self.launch_dims_replace[j] = (global_dim_idx, local_dim_idx) + assert p.p.local_size is not None + self.launch_dims_base[j] = (tuple(p.p.global_size), tuple(p.p.local_size)) - global_dim_idx, local_dim_idx = find_symbolic_dim(ji.prg.p.global_size), find_symbolic_dim(ji.prg.p.local_size) - if global_dim_idx is not None or local_dim_idx is not None: - self.launch_dims_replace[j] = (global_dim_idx, local_dim_idx) - assert ji.prg.p.local_size is not None - self.launch_dims_base[j] = (tuple(ji.prg.p.global_size), tuple(ji.prg.p.local_size)) + estimates = sum((estimate_uop(call) for call in self.linear.src), Estimates()) # used in MultiGraphRunner. tracks (offset, end, dep) ranges per base buffer id to handle suballocated buffers correctly. self.w_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list) self.r_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list) - assert self.jit_cache[0].prg is not None - super().__init__(colored(f"", "cyan"), self.jit_cache[0].prg.device.split(":")[0], estimates.simplify()) + super().__init__(colored(f"", "cyan"), self.calls[0][2][0].device.split(":")[0], estimates.simplify()) def updated_vars(self, var_vals: dict[str, int]): vals = [var_vals[v] for v in self.vars] @@ -194,14 +173,14 @@ class GraphRunner(Runner): @staticmethod def supports_exec_item(batch_devs:list[Compiled], new_call:UOp) -> bool: - return new_call.src[0].op in (Ops.SINK, Ops.PROGRAM) and len(GraphRunner._all_devs(batch_devs, new_call)) == 1 + return new_call.src[0].op is Ops.PROGRAM and len(GraphRunner._all_devs(batch_devs, new_call)) == 1 # a marker for your graph supporting multiple devices of the same type class MultiGraphRunner(GraphRunner): @staticmethod def supports_exec_item(batch_devs:list[Compiled], new_call:UOp) -> bool: # Devices must be the same type - return new_call.src[0].op in (Ops.SINK, Ops.PROGRAM, Ops.COPY) and len(dedup([type(d) for d in GraphRunner._all_devs(batch_devs, new_call)])) == 1 + return new_call.src[0].op in (Ops.PROGRAM, Ops.COPY) and len(dedup([type(d) for d in GraphRunner._all_devs(batch_devs, new_call)])) == 1 ReturnType = TypeVar('ReturnType') @dataclass diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index eec31fa130..0b45cd8d87 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -1,16 +1,30 @@ -from typing import cast, Callable, Iterator -import time, pprint, random, itertools, math, contextlib, weakref +from typing import cast, Iterator +import time, random, itertools, math, contextlib, weakref from dataclasses import dataclass, replace, field -from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, NOOPT, all_int, Metadata, TRACEMETA, TracingKey +from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, NOOPT, all_int, Metadata, TRACEMETA, TracingKey from tinygrad.helpers import BEAM, DEVECTORIZE, size_to_str, time_to_str, VALIDATE_WITH_CPU, cpu_profile, PROFILE, ProfilePointEvent, cpu_events -from tinygrad.helpers import prod, unwrap, EMULATED_DTYPES, flatten -from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite +from tinygrad.helpers import prod, EMULATED_DTYPES, flatten +from tinygrad.dtype import dtypes +from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite, ProgramInfo from tinygrad.device import Device, Buffer, MultiBuffer -from tinygrad.renderer import ProgramSpec, Estimates -from tinygrad.codegen import get_program, to_program +from tinygrad.renderer import Estimates +from tinygrad.codegen import to_program +from tinygrad.codegen.opt.postrange import bufs_from_ast # **************** Stat **************** +def estimate_uop(call:UOp) -> Estimates: + if call.src[0].op is Ops.SINK: call = pm_compile.rewrite(call) + + ast = call.src[0] + if ast.op is Ops.PROGRAM: return ast.src[0].arg.estimates or Estimates() + if ast.op is Ops.COPY or (ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec"): + nbytes = prod(call.src[1].shape) * call.src[1].dtype.itemsize + return Estimates(lds=nbytes, mem=nbytes) + if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": + return runner.estimates if (runner:=graph_cache.get(ast)) is not None else Estimates() + return Estimates() + def update_stats(display_name:str, device:str, estimates:Estimates, var_vals:dict[str, int], et:float|None, buf_count:int, jit=False, metadata:tuple[Metadata, ...]=(), first_run=False): GlobalCounters.kernel_count += 1 @@ -42,89 +56,50 @@ class Runner: def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False) -> float|None: raise NotImplementedError("override this") -def optimize_local_size(_prg:Callable, global_size:list[int], rawbufs:list[Buffer]) -> list[int]: - test_rawbuffers = [Buffer(rawbufs[0].device, rawbufs[0].size, rawbufs[0].dtype).allocate(), *rawbufs[1:]] if rawbufs[0] in rawbufs[1:] else rawbufs - MAX_WORKGROUP = 1024 - local_dims = [[x for x in set([sz, 1, 2, 4, 8, 16, 32, 64, 128, 256, MAX_WORKGROUP]) if x<=sz] for sz in global_size] - local_sizes = [list(x) for x in itertools.product(*local_dims) if prod(x) <= MAX_WORKGROUP] * 2 # try each valid size twice - def try_exec(local_size): - try: - return _prg(*[x._buf for x in test_rawbuffers],global_size=[g//l if g%l == 0 else g/l for g,l in zip(global_size, local_size)], - local_size=local_size, wait=True) - except Exception: return float('inf') - ret = min([(try_exec(local_size), local_size) for local_size in random.sample(local_sizes, len(local_sizes))]) - assert not math.isinf(ret[0]), "all optimize_local_size exec failed" - return ret[1] +local_size_cache: dict[bytes, tuple[int, ...]] = {} +def optimize_local_size(call:UOp, prg:UOp) -> UOp|None: + device = prg.src[1].arg + if prg.arg.local_size is not None or not Device[device].renderer.has_local or not all_int(prg.arg.global_size): return None + + if (local_size:=local_size_cache.get(prg.key)) is None: + bufs = [b._buf for b in (b.allocate() for b in bufs_from_ast(prg.src[0], device))] + rt = Device[device].runtime(prg.arg.function_name, prg.src[4].arg, *prg.arg.aux, runtimevars=prg.arg.runtimevars) + def try_exec(local_size): + try: return rt(*bufs, global_size=[g//l if g%l == 0 else g/l for g,l in zip(prg.arg.global_size, local_size)], local_size=local_size, wait=True) + except Exception: return float('inf') + + MAX_WORKGROUP = 1024 + local_dims = [[x for x in set([sz, 1, 2, 4, 8, 16, 32, 64, 128, 256, MAX_WORKGROUP]) if x<=sz] for sz in prg.arg.global_size] + local_sizes = [list(x) for x in itertools.product(*local_dims) if prod(x) <= MAX_WORKGROUP] * 2 # try each valid size twice + best_time, best = min([(try_exec(ls), ls) for ls in random.sample(local_sizes, len(local_sizes))]) + assert not math.isinf(best_time), "all optimize_local_size exec failed" + local_size = local_size_cache[prg.key] = tuple(best) + + new_global = tuple(g//l if g%l == 0 else g/l for g,l in zip(prg.arg.global_size, local_size)) + return call.replace(src=(prg.replace(arg=replace(prg.arg, global_size=new_global, local_size=local_size)), *call.src[1:])) class CompiledRunner(Runner): - def __init__(self, p:ProgramSpec, prg=None): - if DEBUG >= 3 and p.applied_opts: print(p.applied_opts) - if DEBUG >= 4: print(p.src) - if p.lib is None: - with cpu_profile(TracingKey(f"compile {p.name}", (p.function_name,)), "TINY"): - p = replace(p, lib=Device[p.device].compiler.compile_cached(p.src)) - self.p:ProgramSpec = p - assert self.p.lib is not None - if DEBUG >= 7: Device[p.device].compiler.disassemble(self.p.lib) - self._prg = Device[p.device].runtime(p.function_name, self.p.lib, *p.aux, runtimevars=p.runtimevars) if prg is None else prg - super().__init__(p.name, p.device, p.estimates) - - def __reduce__(self): return self.__class__, (self.p,) + def __init__(self, prg:UOp, device:str): + info: ProgramInfo = prg.arg + sink = prg.src[0] + if DEBUG >= 3 and sink.arg.applied_opts: print(sink.arg.applied_opts) + if DEBUG >= 4: print(prg.src[3].arg) + if len(prg.src) <= 4 or prg.src[4].op is not Ops.BINARY: + with cpu_profile(TracingKey(f"compile {info.name}", (info.function_name,)), "TINY"): + lib = Device[device].compiler.compile_cached(prg.src[3].arg) + prg = prg.replace(src=prg.src + (UOp(Ops.BINARY, arg=lib),)) + self.prg:UOp = prg + self.p:ProgramInfo = info + if DEBUG >= 7: Device[device].compiler.disassemble(prg.src[4].arg) + self._prg = Device[device].runtime(info.function_name, prg.src[4].arg, *info.aux, runtimevars=info.runtimevars) + super().__init__(info.name, device, sink.arg.estimates or Estimates()) def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None, wait=False, timeout:int|None=None) -> float|None: if var_vals is None: var_vals = {} global_size, local_size = self.p.launch_dims(var_vals) - if Device[self.p.device].renderer.has_local and local_size is None and all_int(self.p.global_size): - local_size = optimize_local_size(self._prg, global_size, rawbufs) - global_size = [g//l if g%l == 0 else g/l for g,l in zip(global_size, local_size)] - self.p = replace(self.p, global_size=global_size, local_size=local_size) return self._prg(*[x._buf for x in rawbufs], global_size=tuple(global_size), local_size=tuple(local_size) if local_size else None, vals=tuple(var_vals[k.expr] if k.expr not in self.p.runtimevars else None for k in self.p.vars), wait=wait, timeout=timeout) -class ViewOp(Runner): - def __init__(self, buf:Buffer): super().__init__(colored(f"view {buf.nbytes:8d} @ {buf.offset:<10d}", "yellow"), buf.device) - def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False): - assert rawbufs[0]._base is not None and rawbufs[0]._base == rawbufs[1].base, f"must be base {rawbufs}" - -class BufferCopy(Runner): - def __init__(self, total_sz, dest_device, src_device): - sz = f"{total_sz/1e6:7.2f}M" if total_sz >= 1e6 else f"{total_sz:8d}" - name = f"{type(self).__name__[6:].lower()} {sz}, {dest_device[:7]:>7s} <- {src_device[:7]:7s}" - super().__init__(colored(name, "yellow"), dest_device, Estimates(lds=total_sz, mem=total_sz)) - def copy(self, dest, src): - disk_supports_fast_copyout = src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None - if disk_supports_fast_copyout 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 isinstance(src.device, str) and src.device.startswith(("DISK", "TINYFS")) and hasattr(dest.allocator, '_as_buffer'): - # fast(ish) path, uses readinto in diskbuffers - src.allocator._copyout(dest.allocator._as_buffer(dest._buf), src._buf) - else: - dest.copyin(src.as_memoryview(allow_zero_copy=True)) # may allocate a CPU buffer depending on allow_zero_copy - def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False): - dest, src = rawbufs[0:2] - assert dest.size == src.size and dest.dtype == src.dtype, f"buffer copy mismatch, {dest.size} != {src.size}, {dest.dtype} != {src.dtype}" - st = time.perf_counter() - self.copy(dest, src) - if wait: - Device[dest.device].synchronize() - return time.perf_counter() - st - -class BufferXfer(BufferCopy): - def copy(self, dest, src): dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev) - -class EncDec(Runner): - def __init__(self, cf:UOp, total_sz:int, device:str): - self.shape, self.pos_var = tuple(s.arg for s in cf.src if s.op is Ops.CONST), cf.variables()[0].expr - name = f"enc/dec {total_sz/1e6:7.2f}M, HEVC" if total_sz >= 1e6 else f"enc/dec {total_sz:8d}, HEVC" - super().__init__(colored(name, "yellow"), device, Estimates(lds=total_sz, mem=total_sz)) - def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False): - st = time.perf_counter() - rawbufs[0].allocator._encode_decode(rawbufs[0]._buf, rawbufs[1]._buf, rawbufs[2]._buf, - [x._buf for x in rawbufs[3:]], self.shape, var_vals[self.pos_var]) - if wait: - Device[rawbufs[0].device].synchronize() - return time.perf_counter() - st - # **************** method cache **************** method_cache: dict[tuple[str, type, bytes, tuple, bool], CompiledRunner] = {} @@ -135,62 +110,12 @@ def get_runner(device:str, ast:UOp) -> CompiledRunner: if cret:=method_cache.get(ckey): return cret bkey = (device.split(":")[0], type(Device[device].compiler), ast.key, context, True) if bret:=method_cache.get(bkey): - method_cache[ckey] = ret = CompiledRunner(replace(bret.p, device=device)) + method_cache[ckey] = ret = CompiledRunner(bret.prg, device) else: - prg: ProgramSpec = get_program(ast, Device[device].renderer) - method_cache[ckey] = method_cache[bkey] = ret = CompiledRunner(replace(prg, device=device)) + prg = to_program(ast, Device[device].renderer) + method_cache[ckey] = method_cache[bkey] = ret = CompiledRunner(prg, device) return ret -# **************** lowering functions **************** - -# NOTE: ctx is the buffers -si_lowerer = PatternMatcher([ - (UPat((Ops.SINK, Ops.PROGRAM), name="sink"), lambda ctx,sink: get_runner(ctx[0].device, sink)), - (UPat(Ops.BUFFER_VIEW), lambda ctx: ViewOp(ctx[0])), - (UPat(Ops.COPY), lambda ctx: (BufferXfer(ctx[0].nbytes, ctx[0].device, ctx[1].device) \ - if hasattr(alc:=Device[ctx[0].device].allocator, '_transfer') and alc.supports_transfer and all_same([x.device.split(":")[0] for x in ctx]) \ - else BufferCopy(ctx[0].nbytes, ctx[0].device, ctx[1].device))), - (UPat(Ops.CUSTOM_FUNCTION, arg="encdec", name="cf"), lambda ctx,cf: EncDec(cf, ctx[0].nbytes, ctx[0].device)), - (UPat(Ops.CUSTOM_FUNCTION, arg="graph", name="cf"), lambda ctx,cf: Device[cf.device if isinstance(cf.device,str) else cf.device[0]].graph(cf, ctx)) -]) - -@dataclass -class ExecItem: - ast: UOp - bufs: list[Buffer|None] = field(default_factory=list) - metadata: tuple[Metadata, ...] = () - fixedvars: dict[str, int] = field(default_factory=dict) - prg: Runner|None = None - - def lower(self): - """Populate self.prg by lowering the AST.""" - if self.prg is not None: return self - try: self.prg = cast(Runner, si_lowerer.rewrite(self.ast, self.bufs)) - except Exception as e: - if DEBUG >= 2: - print(f"error lowering {self.ast.op}") - print("tensor operations:") - pprint.pprint(self.metadata, indent=2) - raise e - return self - - def run(self, _var_vals:dict[str, int]|None=None, wait=False, jit=False, do_update_stats=True) -> float|None: - if self.prg is None: self.lower() - assert self.prg is not None - var_vals = self.fixedvars if _var_vals is None else (_var_vals|self.fixedvars) - # reorder bufs to match program globals if needed - _bufs = [self.bufs[i] for i in self.prg.p.globals] if isinstance(self.prg, CompiledRunner) else self.bufs - bufs = [unwrap(x) for x in _bufs] if jit else [unwrap(x).ensure_allocated() for x in _bufs] - if PROFILE: - payload = {"metadata":self.metadata, "var_vals":var_vals, "bufs":[b.trace_num for b in bufs], "name":self.prg.display_name} - payload["outputs"], payload["inputs"] = (self.prg.p.outs, self.prg.p.ins) if isinstance(self.prg, CompiledRunner) else ([0], [1]) - cpu_events.append(ProfilePointEvent(self.prg.device, "exec", len(cpu_events), payload)) - et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2) - if do_update_stats: - update_stats(self.prg.display_name, self.prg.device, self.prg.estimates, var_vals, et, len(bufs), jit, self.metadata, self.prg.first_run) - self.prg.first_run = False - return et - # **************** run linear **************** capturing: list = [] # put classes with an add_linear method in here @@ -208,7 +133,7 @@ def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp: def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in call.src[1:] if b.op is not Ops.BIND] @contextlib.contextmanager -def track_stats(ctx:ExecContext, call:UOp, device:str, display_name:str, estimates:Estimates, bufs:list[Buffer], var_vals:dict[str, int], +def track_stats(ctx:ExecContext, call:UOp, device:str, display_name:str, bufs:list[Buffer], var_vals:dict[str, int], outputs=(0,), inputs=(1,), first_run=False): if PROFILE: cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"metadata": call.arg.metadata, "var_vals": var_vals, "bufs": [b.trace_num for b in bufs], "name": display_name, "outputs": outputs, "inputs": inputs})) @@ -219,7 +144,7 @@ def track_stats(ctx:ExecContext, call:UOp, device:str, display_name:str, estimat if DEBUG >= 2 and timing[0] is None: Device[device].synchronize() timing[0] = time.perf_counter() - st - update_stats(display_name, device, estimates, var_vals, timing[0], len(bufs), jit=ctx.jit, metadata=call.arg.metadata, first_run=first_run) + update_stats(display_name, device, estimate_uop(call), var_vals, timing[0], len(bufs), jit=ctx.jit, metadata=call.arg.metadata, first_run=first_run) def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], dict[str, int]]]: bufs = [b.buffer for b in resolved] @@ -232,16 +157,23 @@ def exec_view(ctx:ExecContext, call, ast): resolved = resolve_params(call, ctx.input_uops) bufs = [cast(Buffer, b.buffer) for b in resolved] bv = bufs[1].view(resolved[0].arg, ast.dtype, ast.arg[1]*bufs[1].dtype.itemsize) - with track_stats(ctx, call, bv.device, colored(f"view {bv.nbytes:8d} @ {bv.offset:<10d}", "yellow"), Estimates(), [bv, bufs[1]], ctx.var_vals): + with track_stats(ctx, call, bv.device, colored(f"view {bv.nbytes:8d} @ {bv.offset:<10d}", "yellow"), [bv, bufs[1]], ctx.var_vals): buffers[resolved[0]] = bv def exec_copy(ctx:ExecContext, call, ast): for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)): dest, src = bufs[0].ensure_allocated(), bufs[1].ensure_allocated() - xfer = hasattr(alc:=Device[dest.device].allocator,'_transfer') and alc.supports_transfer and dest.device.split(":")[0]==src.device.split(":")[0] - prg = (BufferXfer if xfer else BufferCopy)(dest.nbytes, dest.device, src.device) - with track_stats(ctx, call, dest.device, prg.display_name, Estimates(lds=dest.nbytes, mem=dest.nbytes), [dest, src], ctx.var_vals): - prg.copy(dest, src) + xfer = hasattr(dest.allocator,'_transfer') and dest.allocator.supports_transfer and dest.device.split(":")[0] == src.device.split(":")[0] + name = colored(f"{'xfer' if xfer else 'copy'} {size_to_str(bufs[0].nbytes):>10}, {dest.device[:7]:>7s} <- {src.device[:7]:7s}", "yellow") + with track_stats(ctx, call, dest.device, name, [dest, src], ctx.var_vals): + if xfer: + dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev) # type:ignore[attr-defined] + 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 src.device.startswith(("DISK", "TINYFS")) and hasattr(dest.allocator, '_as_buffer'): + src.allocator._copyout(dest.allocator._as_buffer(dest._buf), src._buf) + else: dest.copyin(src.as_memoryview(allow_zero_copy=True)) def exec_kernel(ctx:ExecContext, call, ast): for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)): @@ -249,35 +181,46 @@ def exec_kernel(ctx:ExecContext, call, ast): prg = get_runner(bufs[0].device, ast) prg_bufs = [bufs[i].ensure_allocated() for i in prg.p.globals] - if VALIDATE_WITH_CPU and ast.op is Ops.SINK: - cpu_bufs = [Buffer("CPU", b.size, b.dtype).ensure_allocated().copyin(b.ensure_allocated().as_memoryview()) for b in bufs] - - with track_stats(ctx, call, prg.device, prg.display_name, prg.estimates, prg_bufs, var_vals, + with track_stats(ctx, call, prg.device, prg.display_name, prg_bufs, var_vals, outputs=tuple(prg.p.outs), inputs=tuple(prg.p.ins), first_run=prg.first_run) as timing: timing[0] = prg(prg_bufs, var_vals, wait=DEBUG >= 2) prg.first_run = False - if VALIDATE_WITH_CPU and ast.op is Ops.SINK: - import numpy as np - cpu_prg = get_runner("CPU", ast) - cpu_prg([cpu_bufs[i] for i in cpu_prg.p.globals], var_vals, wait=False) - for i in prg.p.outs: np.testing.assert_allclose(prg_bufs[i].numpy(), cpu_bufs[i].numpy(), rtol=1e-3, atol=1e-3) +def exec_validate(ctx:ExecContext, call, ast): + import numpy as np + for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)): + cpu_bufs, dev_bufs = bufs[:len(bufs)//2], bufs[len(bufs)//2:] + cpu_prg = get_runner("CPU", ast.src[0]) + cpu_prg([cpu_bufs[i].ensure_allocated() for i in cpu_prg.p.globals], {**ctx.var_vals, **device_vars}, wait=False) + for i in cpu_prg.p.outs: np.testing.assert_allclose(dev_bufs[i].ensure_allocated().numpy(), cpu_bufs[i].numpy(), rtol=1e-3, atol=1e-3) def exec_encdec(ctx:ExecContext, call, ast): bufs = [cast(Buffer, b.buffer).ensure_allocated() for b in resolve_params(call, ctx.input_uops)] shape, pos_var = tuple(s.arg for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr - with track_stats(ctx, call, bufs[0].device, colored(f"enc/dec {size_to_str(bufs[0].nbytes)}", "yellow"), - Estimates(lds=bufs[0].nbytes, mem=bufs[0].nbytes), bufs, ctx.var_vals): + with track_stats(ctx, call, bufs[0].device, colored(f"enc/dec {size_to_str(bufs[0].nbytes)}", "yellow"), bufs, ctx.var_vals): bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var]) graph_cache:weakref.WeakKeyDictionary[UOp, Runner] = weakref.WeakKeyDictionary() def exec_graph(ctx:ExecContext, call, cf): bufs = flatten([b.bufs if isinstance(b, MultiBuffer) else [b] for b in (u.buffer for u in resolve_params(call, ctx.input_uops))]) if (runner:=graph_cache.get(cf)) is None: - graph_cache[cf] = runner = Device[cf.device if isinstance(cf.device, str) else cf.device[0]].graph(cf, bufs, input_uops=ctx.input_uops) - with track_stats(ctx, call, runner.device, runner.display_name, runner.estimates, bufs, ctx.var_vals) as t: + graph_cache[cf] = runner = Device[cf.device if isinstance(cf.device, str) else cf.device[0]].graph(cf, input_uops=ctx.input_uops) + with track_stats(ctx, call, runner.device, runner.display_name, bufs, ctx.var_vals) as t: t[0] = runner(bufs, ctx.var_vals, wait=DEBUG >= 2, input_uops=ctx.input_uops) # type: ignore[call-arg] +# flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src +pm_flatten_linear = PatternMatcher([ + (UPat(Ops.LINEAR, custom_early_reject={Ops.LINEAR}, name="lin"), + lambda lin: lin.replace(src=tuple(flatten(c.src if c.op is Ops.LINEAR else (c,) for c in lin.src)))), +]) + +def _validate(call:UOp, sink:UOp) -> UOp: + params = tuple(p for p in call.src[1:] if p.op is not Ops.BIND) + shadows = tuple(UOp.new_buffer(("CPU",)*len(p.device) if isinstance(p.device, tuple) else "CPU", prod(p.max_shape), p.dtype.base) for p in params) + copies = tuple(p.copy_to_device(s.device).call(s, p) for s, p in zip(shadows, params)) + return UOp(Ops.LINEAR, src=copies + (call, UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(sink,), arg="validate").call(*shadows, *params))) +pm_validate = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.SINK, name="sink"),), name="call", allow_any_len=True), _validate)]) + pm_flatten_linear + # ctx is beam value pm_beam = PatternMatcher([ (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="sink"),), name="call", allow_any_len=True), @@ -289,19 +232,26 @@ pm_compile = PatternMatcher([ call.replace(src=(to_program(ast, Device[call.device if isinstance(call.device, str) else call.device[0]].renderer), *call.src[1:]))), ]) +pm_optimize_local_size = PatternMatcher([ + (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), optimize_local_size), +]) + pm_exec = PatternMatcher([ (UPat(Ops.CALL, src=(UPat(Ops.BUFFER_VIEW, name="ast"),), name="call", allow_any_len=True), exec_view), (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="ast"),), name="call", allow_any_len=True), exec_copy), - (UPat(Ops.CALL, src=(UPat((Ops.PROGRAM, Ops.SINK), name="ast"),), name="call", allow_any_len=True), exec_kernel), + (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="ast"),), name="call", allow_any_len=True), exec_kernel), (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="encdec", name="ast"),), name="call", allow_any_len=True), exec_encdec), (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="graph", name="cf"),), name="call", allow_any_len=True), exec_graph), + (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate), ]) -def compile_linear(linear:UOp, beam=0) -> UOp: +def compile_linear(linear:UOp, beam=0, validate=False) -> UOp: + if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True) if (beam_val:=(beam or BEAM.value)) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True) - return graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True) if not VALIDATE_WITH_CPU else linear + linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True) + return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True) def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:tuple[UOp, ...]=(), do_update_stats=True, jit=False): - if not jit: linear = compile_linear(linear) + if not jit: linear = compile_linear(linear, validate=VALIDATE_WITH_CPU) ctx = ExecContext(var_vals or {}, input_uops, do_update_stats, jit) for call in linear.src: pm_exec.rewrite(call, ctx) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 2cc58cc68d..5cf6827c1c 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -133,13 +133,13 @@ def select_by_name(candidates:Sequence[T], get_name:Callable[...,str], query:str raise RuntimeError(err_msg + (f", did you mean: {m[0]!r}?" if (m:=difflib.get_close_matches(query, map(get_name, candidates))) else "")) return ret -def select_first_inited(candidates:Sequence[Callable[...,T]], err_msg:str, cache:dict|None=None, **kwargs): +def select_first_inited(candidates:Sequence[Callable[...,T]], err_msg:str, cache:dict|None=None, *args): excs = [] for typ in candidates: - if cache is not None and typ in cache: return cache[typ] + if cache is not None and (typ,) + args in cache: return cache[(typ,) + args] try: - x = typ(**kwargs) - if cache is not None: cache[typ] = x + x = typ(*args) + if cache is not None: cache[(typ,) + args] = x return x except Exception as e: excs.append(e) raise excs[0] if len(excs) == 1 else ExceptionGroup(err_msg + " is available", excs) diff --git a/tinygrad/mixin/__init__.py b/tinygrad/mixin/__init__.py index 5962746cf8..437ad81768 100644 --- a/tinygrad/mixin/__init__.py +++ b/tinygrad/mixin/__init__.py @@ -202,6 +202,73 @@ class OpMixin(ElementwiseMixin, ReduceMixin): base = base.cast(least_upper_dtype(base.dtype, dtypes.from_py(value))) return base + MovementMixin.pad(X.ones_like(), pads).cast(dtypes.bool).where(base.zeros_like(), base.full_like(value)) + def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self: + if any(pB>sh or pA>sh for (pB,pA),sh in zip(pX, self.shape)): raise ValueError('Padding value causes wrapping around more than once.') + if any(pB<0 or pA<0 for pB,pA in pX): raise NotImplementedError("Negative pads with circular pads is not supported") + orig_shape, X = self.shape, self.repeat(tuple(1 + bool(pB) + bool(pA) for pB,pA in pX)) + return X.shrink(tuple((0 if pB == 0 else osh-pB, xsh if pA == 0 else xsh-osh+pA) for (pB,pA),osh,xsh in zip(pX, orig_shape, X.shape))) + + def _pad_reflect_replicate(self, pX:tuple[tuple[sint, sint], ...], mode:str) -> Self: + X, pads = self, tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX) + for d,(pB,pA) in enumerate(pads): + if mode == "reflect": + if pB >= (s:=X.shape[d]) or pA>=s: raise ValueError(f"Padding ({pB}, {pA}) should be less than the input size={s} for dim={d}.") + slcB, slcA = slice(pB,0,-1), slice(s-2 if s-2>=0 else None, s-2-pA if s-2-pA>=0 else None, -1) + xB, xA = (X[[slc if i == d else slice(None) for i in range(X.ndim)]] if p > 0 else None for slc, p in ((slcB, pB), (slcA, pA))) + else: + shrB, shrA = tuple((0,1) if i==d else None for i in range(X.ndim)), tuple((X.shape[i]-1,X.shape[i]) if i==d else None for i in range(X.ndim)) + xB, xA = (X.shrink(shr).expand(tuple(p if i==d else None for i in range(X.ndim))) if p > 0 else None for shr, p in ((shrB, pB), (shrA, pA))) + pieces = [X_ for X_ in (xB, X, xA) if X_ is not None] + X = pieces[0].cat(*pieces[1:], dim=d) + # shrink after for negative pads (reflection/replication must see full data first) + return X.shrink(tuple((-min(pB,0), min(pA+s,s)) for (pB,pA),s in zip(pX, X.shape))) + + def pad(self, padding:Sequence[sint]|Sequence[tuple[sint, sint]|None], mode:str="constant", value:float=0.0) -> Self: + """ + Returns a tensor with padding applied based on the input `padding`. + + `padding` supports two padding structures: + + 1. Flat padding: `(padding_left, padding_right, padding_top, padding_bottom, ...)` + - This structure matches PyTorch's pad. + - `padding` length must be even. + + 2. Group padding: `(..., (padding_top, padding_bottom), (padding_left, padding_right))` + - This structure matches pad for JAX, NumPy, TensorFlow, and others. + - For each axis, padding can be `None`, meaning no padding, or a tuple `(start, end)`. + - `padding` must have the same length as `self.ndim`. + + Padding values can be negative, resulting in dimension shrinks that work similarly to Python negative slices. + Padding modes is selected with `mode` which supports `constant`, `reflect` and `replicate`. + + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor.arange(9).reshape(1, 1, 3, 3) + print(t.numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.pad((1, 2, 0, -1)).numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.pad(((None, None, (0, -1), (1, 2)))).numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.pad((1, 2, 0, -1), value=-float('inf')).numpy()) + ``` + """ + # normalize to grouped format + pX: tuple[tuple[sint, sint], ...] + if not any(isinstance(p, (tuple, type(None))) for p in padding): + if len(padding)%2 != 0: raise ValueError("Flat padding must have even number of pads") + pX = ((0,0),)*(self.ndim - len(padding)//2) + flat_to_grouped(padding) # type: ignore[arg-type] + else: pX = tuple((0,0) if p is None else p for p in padding) # type: ignore[misc] + if len(pX) != self.ndim: raise ValueError(f"padding length is improper, {padding=} {self.ndim=}") + # dispatch + if mode == "constant": return self._pad_constant(pX, value) + assert all_int(self.shape), f"does not support symbolic shape {self.shape}" + if mode == "circular": return self._pad_circular(pX) + if mode in {"reflect", "replicate"}: return self._pad_reflect_replicate(pX, mode) + raise NotImplementedError(f"{mode=} is not supported") + def _ufix_keep_dtype(self, x) -> bool: # matches Tensor scalar-wrapping behavior: keep self.dtype for float self, or for int self with int/Invalid scalar return dtypes.is_float(self.dtype) or (dtypes.is_int(self.dtype) and isinstance(x, (int, InvalidType))) @@ -702,6 +769,96 @@ class OpMixin(ElementwiseMixin, ReduceMixin): """ return self._inverse().argmax(axis=axis, keepdim=keepdim) + def sort(self, dim:int=-1, descending:bool=False) -> tuple[Self, Self]: + """ + Performs a bitonic sort on the tensor along the specified dimension. + + Order of indices for equivalent elements is always preserved. + + See: https://en.wikipedia.org/wiki/Bitonic_sorter + + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor([[0.1, 0.5, 1.2, 3.4, 2.1], [2.2, 1.9, 0.3, 4.5, 0.8]]) + print(t.numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + sorted_values, indices = t.sort(dim=1, descending=True) + print(sorted_values.numpy()) + print(indices.numpy()) + ``` + """ + x, dim = self, self._resolve_dim(dim) + if (orig_len := int(x.shape[dim])) <= 1: return x, x.zeros_like(dtype=dtypes.default_int) + # pad to power of 2 + n_stages = (orig_len-1).bit_length() + pads = tuple((0, 2**n_stages - orig_len) if i == dim else None for i in range(x.ndim)) + x = x._pad_constant(pads, x.dtype.min if descending else x.dtype.max).unflatten(dim, (2,)*n_stages) + # https://en.wikipedia.org/wiki/Bitonic_sorter#/media/File:BitonicSort1.svg + for stage in range(1, n_stages+1): + if stage != n_stages: + # flip so arrows of green boxes point the same way as blue boxes + crossover_dim = dim + n_stages - stage - 1 + blue_box, green_box = x.split(1, crossover_dim) + flip_dims = tuple(-i for i in range(1, stage+1+(self.ndim-dim))) + x = (blue_box.cat(green_box.flip(flip_dims), dim=crossover_dim)).contiguous() + for substage in range(stage-1, -1, -1): + partner_dim = dim + n_stages - substage - 1 + x_top, x_bottom = x.split(1, partner_dim) + x_larger, x_smaller = x_top.maximum(x_bottom), x_top.minimum(x_bottom) + x = (x_larger.cat(x_smaller, dim=partner_dim) if descending else x_smaller.cat(x_larger, dim=partner_dim)).contiguous() + if stage != n_stages: + # flip wires back to undo the crossover + blue_box, flipped_green_box = x.split(1, crossover_dim) + x = blue_box.cat(flipped_green_box.flip(flip_dims), dim=crossover_dim) + x = x.flatten(dim, dim+n_stages-1).shrink_to(self.shape) + # compute indices for sorted values + mask = type(self).ones(orig_len, orig_len, dtype=dtypes.bool, device=self.device).tril().reshape((None, None) + (1,)*(self.ndim-dim-1)) + def compute_counts(t:Self): return (mask & t.unsqueeze(dim).eq(t.unsqueeze(dim+1))).sum(dim+1) + count_orig, count_sorted = compute_counts(self), compute_counts(x) + cond = self.unsqueeze(dim+1).eq(x.unsqueeze(dim)) & count_orig.unsqueeze(dim+1).eq(count_sorted.unsqueeze(dim)) + idx = type(self).arange(orig_len, device=self.device).reshape(tuple(orig_len if i == dim else 1 for i in range(x.ndim))) + idx = (cond * idx.unsqueeze(dim+1)).sum(dim) + return x, idx + + def argsort(self, dim:int=-1, descending:bool=False) -> Self: + """ + Returns the indices that sort input tensor along given `dimension` in given `descending` order by value. + + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor([[2, 3, 4, 1], [1, 4, 3, 2]]) + print(t.argsort().numpy()) + ``` + """ + return self.sort(dim, descending)[1] + + def topk(self, k:int, dim:int=-1, largest:bool=True, sorted_:bool=True) -> tuple[Self, Self]: + """ + Computes the top-k elements of the tensor along the specified `dim`. + + Order of indices for equivalent elements is always preserved. + + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor([[0.1, 0.5, 1.2, 3.4, 2.1], [2.2, 1.9, 0.3, 4.5, 0.8]]) + print(t.numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + topk_values, topk_indices = t.topk(2, dim=1) + print(topk_values.numpy()) + print(topk_indices.numpy()) + ``` + """ + if not sorted_: raise NotImplementedError("topk with sorted_=False is not supported") + if k > self.shape[dim:=self._resolve_dim(dim)]: raise ValueError(f"selected index {k=} is out of range") + x, idx = self.sort(dim, descending=largest) + topk_shape = tuple(k if i == dim else None for i in range(self.ndim)) + return x.shrink_to(topk_shape), idx.shrink_to(topk_shape) + + def allclose(self, other:Self, rtol:float=1e-05, atol:float=1e-08, equal_nan=False) -> Self: + """ + Check if all self and other are close. + """ + return self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan).all() + # helper function commonly used for indexing def _one_hot_along_dim(self, num_classes:sint, dim:int=-1) -> Self: from tinygrad.uop.ops import sint_to_uop @@ -831,6 +988,49 @@ class OpMixin(ElementwiseMixin, ReduceMixin): return mask.where(src, 0).sum(-1).add(self if include_self else _inv_mask(self, 0)).div(count) raise RuntimeError(f"{reduce=} must be one of 'sum', 'prod', 'mean', 'amax', 'amin'") + def scatter(self, dim:int, index:Self, src:Self|PyConst, reduce:Literal['multiply', 'add']|None=None) -> Self: + """ + Scatters `src` values along an axis specified by `dim`. + Apply `add` or `multiply` reduction operation with `reduce`. + + NOTE: To use the `reduce` argument with a Tensor `src`, see `Tensor.scatter_reduce`. + + ```python exec="true" source="above" session="tensor" result="python" + src = Tensor.arange(1, 11).reshape(2, 5) + print(src.numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + index = Tensor([[0, 1, 2, 0]]) + print(Tensor.zeros(3, 5, dtype=src.dtype).scatter(0, index, src).numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + index = Tensor([[0, 1, 2], [0, 1, 4]]) + print(Tensor.zeros(3, 5, dtype=src.dtype).scatter(1, index, src).numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(Tensor.full((2, 4), 2.0).scatter(1, Tensor([[2], [3]]), 1.23, reduce='multiply').numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(Tensor.full((2, 4), 2.0).scatter(1, Tensor([[2], [3]]), 1.23, reduce='add').numpy()) + ``` + """ + if reduce not in {None, "add", "multiply"}: raise TypeError(f"{reduce=} must be one of None, 'multiply', or 'add'") + if isinstance(src, (int, float, bool)): src = type(self).full(index.shape, src, dtype=self.dtype, device=self.device) + elif reduce: raise TypeError("non-scalar src is not supported with reduce arg. use scatter_reduce") + if reduce == "add": return self.scatter_reduce(dim, index, src, "sum", include_self=True) + if reduce == "multiply": return self.scatter_reduce(dim, index, src, "prod", include_self=True) + src, mask = self._pre_scatter(dim, index, src) + return self._masked_merge(src, mask, (-1,)) + + def _masked_merge(self, values:Self, mask:Self, axes:tuple[int, ...]) -> Self: + # reduce such that if mask contains repeated indices the last one remains + 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) + # select from values for each True element in mask else select from self + return mask.where(values, self) + # ***** functional nn ops ***** def sequential(self, ll:list[Callable[[Self], Self]]) -> Self: @@ -1185,6 +1385,30 @@ class OpMixin(ElementwiseMixin, ReduceMixin): Y = (1 - label_smoothing)*Y + label_smoothing / int(Y.shape[classes_dim]) return -self.log_softmax(classes_dim).mul(Y).sum(classes_dim)._do_reduction(reduction) + def nll_loss(self, Y:Self, weight:Self|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean") -> Self: + """ + Computes the negative log likelihood loss between log-probabilities and target labels. + + NOTE: `self` is log-probabilities and `Y` is the Y labels or class probabilities. + + See: https://pytorch.org/docs/stable/generated/torch.nn.functional.nll_loss.html + + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor([[-1, 2, -3], [1, -2, 3]]) + Y = Tensor([1, 2]) + print(t.log_softmax().nll_loss(Y).item()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor([[-1, 2, -3], [1, -2, 3]]) + Y = Tensor([1, 2]) + print(t.log_softmax().nll_loss(Y, reduction='none').numpy()) + ``` + """ + weight = Y.ones_like() if weight is None else weight.gather(0, Y.flatten()).reshape(Y.shape) + masked_weight = weight if ignore_index is None else weight * Y.ne(ignore_index) + nll = -self.gather(1, Y.unsqueeze(1)).squeeze(1) * masked_weight + return nll.sum() / masked_weight.sum() if reduction == "mean" else nll._do_reduction(reduction) + # ***** matrix ops ***** def newton_schulz(self, steps:int, params:tuple[int, ...], eps:float=1.0e-7) -> Self: diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index 6a2571f76e..1f55b0afd5 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -39,6 +39,8 @@ class ElementwiseMixin(DTypeMixin, CreationMixin): """ return self.cast(dtypes.bool).ne(True) + def contiguous(self, *args, **kwargs) -> Self: raise NotImplementedError + def contiguous_backward(self) -> Self: """ Inserts a contiguous operation in the backward pass. diff --git a/tinygrad/nn/onnx.py b/tinygrad/nn/onnx.py index 3a1d149bb2..ecc8de6ee3 100644 --- a/tinygrad/nn/onnx.py +++ b/tinygrad/nn/onnx.py @@ -2,7 +2,8 @@ from typing import Any, Sequence, cast, Literal, NamedTuple, Generator import dataclasses, functools, io, math, types, warnings, pathlib, sys, os, struct, enum from tinygrad.nn.state import TensorIO -from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr +from tinygrad.tensor import Tensor, _broadcast_shape +from tinygrad.mixin import ReductionStr from tinygrad.helpers import getenv, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element, polyN from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype, truncate, least_upper_dtype, DTYPES_DICT from tinygrad.device import is_dtype_supported, Device diff --git a/tinygrad/nn/optim.py b/tinygrad/nn/optim.py index 2ee634286b..a912bd6501 100644 --- a/tinygrad/nn/optim.py +++ b/tinygrad/nn/optim.py @@ -94,7 +94,7 @@ def Muon(params: list[Tensor], lr=0.001, momentum=0.95, weight_decay=0.1, ns_ste """ assert not fused, "FUSE_OPTIM not allowed for Muon optimizer" return LARS(params, lr, momentum, weight_decay, ns_steps, ns_coefficients, nesterov, - classic=False, pre_wd=False, tcoef=0.0, device=None, fused=fused) + classic=False, pre_wd=False, tcoef=0.0, device=device, fused=fused) class LARS(Optimizer): """ diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index cc3f66e520..044d6b28f2 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -1,12 +1,10 @@ from __future__ import annotations from typing import Callable, cast -import functools -from dataclasses import dataclass, field -from tinygrad.helpers import to_function_name, dedup, prod, Target, DEBUG -from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, smin, GroupOp, PatternMatcher, print_uops +from dataclasses import dataclass +from tinygrad.helpers import prod, Target +from tinygrad.uop.ops import Ops, UOp, sint, ssimplify, smin, GroupOp, PatternMatcher from tinygrad.dtype import AddrSpace, PtrDType from tinygrad.codegen.opt.tc import TensorCore -from tinygrad.codegen.opt import Opt from tinygrad.device import Compiler @dataclass(frozen=True) @@ -61,76 +59,6 @@ class Estimates: elif u.op is Ops.WMMA and u not in dont_count: flops += 2 * prod(u.arg[1]) // u.arg[5] * mults return Estimates(flops, lds, sum(mem.values())) -@dataclass -class ProgramSpec: - name:str - src:str - device:str - ast:UOp # save the base ast (this is method cache key) - prg:UOp|None=None - uops:list[UOp]|None=None - lib:bytes|None=None - aux:list=field(default_factory=list) - - # filled in from uops (via from_uop) - global_size:list[int]=field(default_factory=lambda: [1,1,1]) - local_size:list[int]|None=None - vars:list[Variable]=field(default_factory=list) - globals:list[int]=field(default_factory=list) - outs:list[int]=field(default_factory=list) - ins:list[int]=field(default_factory=list) - - @property - def estimates(self) -> Estimates: return self.ast.arg.estimates if self.ast.arg is not None and self.ast.arg.estimates is not None else Estimates() - - @functools.cached_property - def function_name(self) -> str: return to_function_name(self.name) - - @functools.cached_property - def runtimevars(self) -> dict[str, int]: return {v.expr: i for i, v in enumerate(self.vars) if v.expr == 'core_id'} - - @property - def applied_opts(self) -> tuple[Opt, ...]|None: return self.ast.arg.applied_opts if self.ast.arg is not None else None - - def launch_dims(self, var_vals:dict[str, int]): - global_size = [sym_infer(sz, var_vals) for sz in self.global_size] - local_size = [sym_infer(sz, var_vals) for sz in self.local_size] if self.local_size is not None else None - return global_size, local_size - - @staticmethod - def from_uop(prg:UOp) -> ProgramSpec: - """Construct ProgramSpec from a PROGRAM UOp.""" - assert prg.op is Ops.PROGRAM, f"expected PROGRAM, got {prg.op}" - # SINK/DEVICE/LINEAR/SOURCE/BINARY? - sink, device, linear, source = prg.src[:4] - lib = prg.src[4].arg if len(prg.src) > 4 else None - uops = list(linear.src) - if DEBUG >= 6: print_uops(uops) # LINEAR is src[2] - - # single pass through the uops to extract metadata - _vars: list[Variable] = [] - _globals: list[int] = [] - outs: list[int] = [] - ins: list[int] = [] - global_size: list[int] = [1, 1, 1] - local_size: list[int]|None = [1, 1, 1] - for u in sink.toposort(): - if u.op is Ops.DEFINE_VAR: _vars.append(u) - if u.op is Ops.PARAM: _globals.append(u.arg) - if u.op in (Ops.STORE, Ops.LOAD): - if (idx:=u.src[0]).op is Ops.INDEX or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX): - if (buf:=idx.src[0]).op is Ops.PARAM: (outs if u.op is Ops.STORE else ins).append(buf.arg) - # TODO: can else happen? - if u.op is Ops.SPECIAL: - if u.arg[0] == 'i': local_size = None - special_size = local_size if u.arg[0] == 'l' else global_size - # TODO: this cast is wrong, u.src[0].ssimplify() can be sint - if special_size is not None: special_size[int(u.arg[-1])] = cast(int, u.src[0].ssimplify()) - if u.op is Ops.DEFINE_VAR and u.arg[0] == 'core_id': global_size[0] = u.arg[2] + 1 - - return ProgramSpec(sink.arg.name, source.arg, device.arg, sink, prg, uops, lib, list(prg.arg) if prg.arg else [], global_size, local_size, - sorted(_vars, key=lambda v: v.arg), sorted(dedup(_globals)), sorted(dedup(outs)), sorted(dedup(ins))) - class Renderer: target: Target suffix: str = "" diff --git a/tinygrad/renderer/amd/elf.py b/tinygrad/renderer/amd/elf.py index 4383d2a811..a2a5aa081d 100644 --- a/tinygrad/renderer/amd/elf.py +++ b/tinygrad/renderer/amd/elf.py @@ -11,7 +11,7 @@ from tinygrad.runtime.autogen.amd.rdna3.ins import s_code_end # same encoding as from tinygrad.runtime.autogen.amd.cdna.ins import s_nop as s_nop_cdna _arch_map = {"gfx9": "cdna", "gfx10": "rdna3", "gfx11": "rdna3", "gfx12": "rdna4"} -def assemble_linear(ctx, prg:UOp, lin:UOp) -> bytes: +def assemble_linear(prg:UOp, lin:UOp, arch:str) -> bytes: insts = [u.arg for u in lin.src] # ** scan for max vgpr/sgpr/accvgpr @@ -41,7 +41,7 @@ def assemble_linear(ctx, prg:UOp, lin:UOp) -> bytes: elif u.op is Ops.DEFINE_LOCAL: lds_size += u.ptrdtype.size * u.ptrdtype.base.itemsize elif u.op is Ops.SPECIAL and u.arg.startswith("gidx"): gids.add(int(u.arg[-1])) code_bytes = b"".join(inst.to_bytes() for inst in insts) - arch = next(v for k, v in _arch_map.items() if ctx.target.arch.startswith(k)) + arch = next(v for k, v in _arch_map.items() if arch.startswith(k)) is_cdna, is_rdna4 = arch == "cdna", arch == "rdna4" # ** pad text to ISA alignment diff --git a/tinygrad/renderer/amd/sqtt.py b/tinygrad/renderer/amd/sqtt.py index b2dad15434..caf58a4c42 100644 --- a/tinygrad/renderer/amd/sqtt.py +++ b/tinygrad/renderer/amd/sqtt.py @@ -105,11 +105,11 @@ class InstOpRDNA4(Enum): SALU_NO_EXEC = 0x7 MESSAGE = 0x9 VALU_1 = 0xa - VALU_TRANS = 0xb - VALU_B1 = 0xc - VALU_B2 = 0xd - VALU_B4 = 0xe - VALU_B16 = 0xf + VALUT_4 = 0xb + VALUB_1 = 0xc + VALUB_2 = 0xd + VALUB_4 = 0xe + VALUB_16 = 0xf VINTERP = 0x12 BARRIER_WAIT = 0x13 FLAT_RD_2 = 0x1c @@ -143,7 +143,7 @@ class InstOpRDNA4(Enum): LDS_PARAM_LOAD = 0x6f SALU_WR_EXEC = 0x72 VALU1_WR_EXEC = 0x73 - VALU_B2_WR_EXEC = 0x74 + VALU_WR_EXEC_2 = 0x74 OTHER_LDS_6 = 0x77 OTHER_LDS_10 = 0x78 BARRIER_SIGNAL = 0x7a @@ -154,7 +154,7 @@ class InstOpRDNA4(Enum): WMMA_32 = 0x8e WMMA_64 = 0x8f VALU_DPFP = 0x92 - SALU_FLOAT3 = 0x98 + SALU_FLOAT_3 = 0x98 VALU_SCL_TRANS = 0x99 SALU_2 = 0x9b SALU_5 = 0x9c @@ -210,7 +210,7 @@ class PacketType: class TS_DELTA_S8_W3(PacketType): encoding = bits[6:0] == 0b0100001 delta = bits[10:8] - _padding = bits[63:11] + _padding = bits[71:11] class TS_DELTA_S5_W3(PacketType): encoding = bits[4:0] == 0b00110 @@ -295,6 +295,16 @@ class WAVEEND(PacketType): # exclude: 1 << 4 @property def cu(self) -> int: return self.wgp | (self.sa << 3) +class WAVEEND_RDNA4(PacketType): + encoding = bits[4:0] == 0b10101 + delta = bits[7:5] + sa = bits[8:8] + simd = bits[10:9] + wgp = bits[14:11] + wave = bits[19:15] + @property + def cu(self) -> int: return self.wgp | (self.sa << 4) + class WAVESTART(PacketType): # exclude: 1 << 4 encoding = bits[4:0] == 0b01100 delta = bits[6:5] @@ -306,16 +316,16 @@ class WAVESTART(PacketType): # exclude: 1 << 4 @property def cu(self) -> int: return self.wgp | (self.sa << 3) -class WAVESTART_RDNA4(PacketType): # Layout 4 has wave field at different position +class WAVESTART_RDNA4(PacketType): # Layout 4: wgp is 4 bits, wave shifted to bits 15-19 encoding = bits[4:0] == 0b01100 delta = bits[6:5] sa = bits[7:7] simd = bits[9:8] - wgp = bits[12:10] + wgp = bits[13:10] wave = bits[19:15] id7 = bits[31:20] @property - def cu(self) -> int: return self.wgp | (self.sa << 3) + def cu(self) -> int: return self.wgp | (self.sa << 4) class WAVEALLOC(PacketType): # exclude: 1 << 10 encoding = bits[4:0] == 0b00101 @@ -415,7 +425,7 @@ PACKET_TYPES_RDNA3: dict[int, type[PacketType]] = { } PACKET_TYPES_RDNA4: dict[int, type[PacketType]] = { **PACKET_TYPES_RDNA3, - 9: WAVESTART_RDNA4, 10: TS_DELTA_S5_W2_RDNA4, 11: WAVEALLOC_RDNA4, + 8: WAVEEND_RDNA4, 9: WAVESTART_RDNA4, 10: TS_DELTA_S5_W2_RDNA4, 11: WAVEALLOC_RDNA4, 12: TS_DELTA_S5_W3_RDNA4, 13: PERF_RDNA4, 22: TS_DELTA_OR_MARK_RDNA4, 24: INST_RDNA4, } @@ -654,7 +664,7 @@ def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, I if isinstance(p, (WAVESTART, WAVESTART_RDNA4, CDNA_WAVESTART)): assert p.wave not in wave_pc, "only one inflight wave per unit" wave_pc[p.wave] = next(iter(pc_map)) - elif isinstance(p, WAVEEND): + elif isinstance(p, (WAVEEND, WAVEEND_RDNA4)): pc = wave_pc.pop(p.wave) yield (p, InstructionInfo(pc, p.wave, s_endpgm())) elif isinstance(p, IMMEDIATE_MASK): @@ -703,7 +713,7 @@ def format_packet(p) -> str: elif isinstance(p, VALUINST): fields = f"wave={p.wave}" + (" flag" if p.flag else "") elif isinstance(p, ALUEXEC): fields = f"src={p.src.name if isinstance(p.src, AluSrc) else p.src}" elif isinstance(p, VMEMEXEC): fields = f"src={p.src.name if isinstance(p.src, MemSrc) else p.src}" - elif isinstance(p, (WAVESTART, WAVESTART_RDNA4, WAVEEND)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}" + elif isinstance(p, (WAVESTART, WAVESTART_RDNA4, WAVEEND, WAVEEND_RDNA4)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}" elif hasattr(p, '_fields'): filt = {'delta', 'encoding'} if not isinstance(p, (TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4)) else {'encoding'} fields = " ".join(f"{k}=0x{getattr(p, k):x}" if k in {'snap', 'val32'} else f"{k}={getattr(p, k)}" diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 03bf27d465..50ba73d340 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -473,10 +473,10 @@ class HIPRenderer(CStyleLanguage): def is_cdna(arch): return arch.split(":")[0] in {"gfx942", "gfx950"} @staticmethod def is_cdna4(arch): return arch.split(":")[0] == "gfx950" - def __init__(self, target:Target): # gfx942 => MI300, gfx1100 => RX 7900, gfx1201 => RX 9700 + 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 - self.compiler, self.tensor_cores = HIPCompiler(target.arch), tc.get_amd(target.arch) + from tinygrad.runtime.support.compiler_amd import HIPCompiler, HIPCCCompiler + self.compiler, self.tensor_cores = (HIPCCCompiler if use_hipcc else HIPCompiler)(target.arch), tc.get_amd(target.arch) if not self.is_cdna4(target.arch): self.extra_matcher += pm_manual_bf16_cast + extra_pm if self.is_cdna(target.arch): self.string_rewrite = PatternMatcher([ @@ -512,7 +512,7 @@ class HIPRenderer(CStyleLanguage): def asm(self, prg:UOp, lin:UOp) -> bytes: from tinygrad.renderer.amd.elf import assemble_linear - return assemble_linear(self, prg, lin) + return assemble_linear(prg, lin, self.target.arch) def render_vector_prefix(self, dtype:DType) -> str: vec, scal = self.render_dtype(dtype), self.render_dtype(dtype.scalar()) @@ -560,10 +560,7 @@ class HIPRenderer(CStyleLanguage): return super().render_kernel(function_name, kernel, bufs, uops, prefix) class HIPCCRenderer(HIPRenderer): - def __init__(self, target:Target): - super().__init__(target) - from tinygrad.runtime.support.compiler_amd import HIPCCCompiler - self.compiler = HIPCCCompiler(target.arch) + def __init__(self, target:Target): super().__init__(target, use_hipcc=True) class QCOMCLRenderer(OpenCLRenderer): def __init__(self, target:Target): diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 20078efbf5..e763701b9e 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -238,7 +238,9 @@ class AMDLLVMRenderer(LLVMRenderer): (UPat(Ops.LOG2, dtype=dtypes.double, src=(UPat.var("d"),)), xlog2), (UPat(Ops.EXP2, dtype=dtypes.double, src=(UPat.var("d"),)), xexp2), ]) - def asm(self, prg: UOp, lin: UOp) -> bytes: return HIPRenderer(self.target).asm(prg, lin) + def asm(self, prg: UOp, lin: UOp) -> bytes: + from tinygrad.renderer.amd.elf import assemble_linear + return assemble_linear(prg, lin, self.target.arch) def render(self, uops: list[UOp]) -> str: prefix = ["""define i8 @f32_to_fp8(float %val, i1 %is_bf8) { entry: %ival = bitcast float %val to i32\n %exp = and i32 %ival, 2139095040\n %is_special = icmp eq i32 %exp, 2139095040 diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index a82a6a8ad3..5bc9894a23 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -6,7 +6,7 @@ from tinygrad.renderer.cstyle import CUDARenderer, OpenCLRenderer from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str from tinygrad.runtime.autogen import mesa from tinygrad.runtime.support.c import POINTER -import base64, ctypes, ctypes.util, struct, functools, inspect, contextlib, itertools +import base64, ctypes, ctypes.util, struct, functools, inspect, itertools def g(s:str): return getattr(mesa, s) def nsrc(d:mesa.nir_def) -> mesa.nir_src: return mesa.nir_src(ssa=ctypes.pointer(d)) @@ -169,9 +169,10 @@ class NIRRenderer(Renderer): self.compiler = fromimport("tinygrad.runtime.support.compiler_mesa", self.__class__.__name__.replace("Renderer", "Compiler"))(target.arch) if hasattr(self.compiler, "nir_options"): self.nir_options = self.compiler.nir_options mesa.glsl_type_singleton_init_or_ref() + self._deinit_types = True def __del__(self): - with contextlib.suppress(AttributeError): mesa.glsl_type_singleton_decref() + if getattr(self, "_deinit_types", False): mesa.glsl_type_singleton_decref() def param(self, b:mesa.nir_builder, x, sz:int) -> mesa.nir_def: raise NotImplementedError("needs param") def prerender(self, uops:list[UOp]): diff --git a/tinygrad/runtime/graph/cuda.py b/tinygrad/runtime/graph/cuda.py index 1519f4c93a..e738abc203 100644 --- a/tinygrad/runtime/graph/cuda.py +++ b/tinygrad/runtime/graph/cuda.py @@ -8,14 +8,14 @@ from tinygrad.runtime.ops_cuda import CUDADevice, check, encode_args, cu_time_ex from tinygrad.engine.jit import MultiGraphRunner class CUDAGraph(MultiGraphRunner): - def __init__(self, linear, input_buffers, input_uops=()): - super().__init__(linear, input_buffers, input_uops) + def __init__(self, linear, input_uops=()): + super().__init__(linear, input_uops) self.nodes: list[tuple[Any, ...]] = [] # list of tuple(graph node, node params, c_args/context, is memcpy) self.graph = init_c_var(cuda.CUgraph, lambda x: check(cuda.cuGraphCreate(ctypes.byref(x), 0))) for (dev_idx, ast, bufs, device_vars), prg in zip(self.calls, self.progs): - if ast.op in (Ops.SINK, Ops.PROGRAM): + if ast.op is Ops.PROGRAM: assert prg is not None global_size, local_size = prg.p.launch_dims({v: 0 for v in self.vars}) diff --git a/tinygrad/runtime/graph/hcq.py b/tinygrad/runtime/graph/hcq.py index d2ad964fac..4bee02eac1 100644 --- a/tinygrad/runtime/graph/hcq.py +++ b/tinygrad/runtime/graph/hcq.py @@ -2,45 +2,44 @@ import collections, time from typing import Any, cast from tinygrad.helpers import round_up, PROFILE, ALL2ALL, merge_dicts, getenv, suppress_finalizing, TracingKey, unwrap from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWQueue, HCQArgsState, BumpAllocator, MMIOInterface -from tinygrad.device import Buffer, BufferSpec, Compiled, Device, ProfileGraphEntry, ProfileGraphEvent +from tinygrad.device import Buffer, BufferSpec, Compiled, Device, MultiBuffer, ProfileGraphEntry, ProfileGraphEvent from tinygrad.dtype import dtypes from tinygrad.uop.ops import UOp, Ops, Variable -from tinygrad.engine.realize import BufferXfer, CompiledRunner, BufferCopy from tinygrad.engine.jit import GraphRunner, MultiGraphRunner from tinygrad.runtime.ops_rdma import RDMACopyQueue class HCQGraph(MultiGraphRunner): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.devices = list(set(cast(HCQCompiled, d) for ji in self.jit_cache for d in [Device[cast(Buffer, x).device] for x in ji.bufs])) + self.devices = list({cast(HCQCompiled, Device[b.device]) for (_,_,bufs,_) in self.calls for b in bufs}) # CPU Device is always last self.devices = sorted(self.devices, key=lambda x: 1 if x._is_cpu() else 0) # Replace input buffers with variables. - self.hcq_bufs = [[cast(Buffer, x)._buf for x in ji.bufs] for ji in self.jit_cache] + self.hcq_bufs = [[b._buf for b in bufs] for (_,_,bufs,_) in self.calls] self.input_replace_to_var: dict[tuple[int, int], Variable] = {} - for (j,i), input_idx in self.input_replace.items(): - x = self.input_replace_to_var.setdefault((j,i), UOp.variable(f"input_{input_idx}", 0, 0xffffffffffffffff, dtype=dtypes.uint64)) - self.hcq_bufs[j][i] = HCQBuffer(x, self.hcq_bufs[j][i].size) # Create fake buffer with variable + for j, replace in enumerate(self.uop_replace): + for pos, iidx in replace: + x = self.input_replace_to_var.setdefault((j,pos), UOp.variable(f"inp_{iidx}_{self.calls[j][0]}", 0, 0xffffffffffffffff, dtype=dtypes.uint64)) + self.hcq_bufs[j][pos] = HCQBuffer(x, self.hcq_bufs[j][pos].size) # Create fake buffer with variable # Allocate kernel args. kernargs_size: dict[Compiled, int] = collections.defaultdict(int) - for ji in self.jit_cache: - if not isinstance(ji.prg, CompiledRunner): continue - kernargs_size[ji.prg.dev] += round_up(ji.prg._prg.kernargs_alloc_size, 16) + for prg in self.progs: + if prg is None: continue + kernargs_size[prg.dev] += round_up(prg._prg.kernargs_alloc_size, 16) self.kernargs_bufs: dict[Compiled, HCQBuffer] = {d:d.allocator._alloc(max(sz, 1), BufferSpec(cpu_access=True)) for d,sz in kernargs_size.items()} # Fill initial arguments. self.ji_args: dict[int, HCQArgsState] = {} kargs_alloc: dict[Compiled, BumpAllocator] = {dev:BumpAllocator(buf.size) for dev,buf in self.kernargs_bufs.items()} - for j,ji in enumerate(self.jit_cache): - if not isinstance(ji.prg, CompiledRunner): continue - - argsbuf = self.kernargs_bufs[ji.prg.dev].offset(kargs_alloc[ji.prg.dev].alloc(ji.prg._prg.kernargs_alloc_size, 16)) - self.ji_args[j] = ji.prg._prg.fill_kernargs(self.hcq_bufs[j], ji.prg.p.vars, argsbuf) + for j, prg in enumerate(self.progs): + if prg is None: continue + argsbuf = self.kernargs_bufs[prg.dev].offset(kargs_alloc[prg.dev].alloc(prg._prg.kernargs_alloc_size, 16)) + self.ji_args[j] = prg._prg.fill_kernargs(self.hcq_bufs[j], prg.p.vars, argsbuf) # Schedule Dependencies. # There are two types of queues on each device: copy and compute. Both must synchronize with all external operations before launching any @@ -81,32 +80,34 @@ class HCQGraph(MultiGraphRunner): for dev, queue in self.comp_queues.items(): self.dev_access[queue].add(dev) - self.input_replace_map: dict[HCQCompiled, set[int]] = collections.defaultdict(set) + self.input_replace_map: dict[HCQCompiled, set[tuple[int, int]]] = collections.defaultdict(set) self.device_vars: dict[HCQCompiled, dict[str, int]] = {} - for j,ji in enumerate(self.jit_cache): - ji_devs = [cast(HCQCompiled, Device[cast(Buffer, b).device]) for b in ji.bufs] if isinstance(ji.prg, BufferXfer) else [] + for j, ((_, ast, bufs, device_vars), prg) in enumerate(zip(self.calls, self.progs)): + is_xfer = ast.op is Ops.COPY and hasattr(alc:=Device[bufs[0].device].allocator, '_transfer') and alc.supports_transfer \ + and bufs[0].device.split(":")[0] == bufs[1].device.split(":")[0] + ji_devs = [cast(HCQCompiled, Device[b.device]) for b in bufs] if is_xfer else [] is_rdma = len(ji_devs) > 0 and not any(d._is_cpu() for d in ji_devs) and len(set(d.peer_group for d in ji_devs)) > 1 - if is_exec_prg:=isinstance(ji.prg, CompiledRunner): enqueue_dev: HCQCompiled = ji.prg.dev + if prg is not None: enqueue_dev: HCQCompiled = prg.dev else: # For copy ops prioritize enqeueuing on the src device, so reverse the buffers. - for b in cast(list[Buffer], ji.bufs[::-1]): + for b in bufs[::-1]: if (enqueue_dev:=cast(HCQCompiled, Device[b.device])).hw_copy_queue_t is not None: break # set any fixedvars on the device - self.device_vars[enqueue_dev] = merge_dicts([self.device_vars.get(enqueue_dev, {}), ji.fixedvars]) - if is_exec_prg: self.device_vars[enqueue_dev] = merge_dicts([self.device_vars[enqueue_dev], cast(CompiledRunner, ji.prg).p.runtimevars]) + self.device_vars[enqueue_dev] = merge_dicts([self.device_vars.get(enqueue_dev, {}), device_vars]) + if prg is not None: self.device_vars[enqueue_dev] = merge_dicts([self.device_vars[enqueue_dev], prg.p.runtimevars]) - if is_exec_prg: + if prg is not None: enqueue_queue = self.comp_queues[enqueue_dev] elif is_rdma: enqueue_queue = self.comp_queues[enqueue_dev] - rdma_key = (cast(HCQCompiled, Device[cast(Buffer, ji.bufs[0]).device]).rdma_dev(), enqueue_dev.rdma_dev()) + rdma_key = (cast(HCQCompiled, Device[bufs[0].device]).rdma_dev(), enqueue_dev.rdma_dev()) self.rdma_queues.setdefault(rdma_key, RDMACopyQueue(enqueue_dev.rdma_dev())) else: assert (enqueue_dev.hw_copy_queue_t is not None), "device must implement a copy queue" - queue_idx = self.devices.index(cast(HCQCompiled, Device[cast(Buffer, ji.bufs[0]).device])) % self.num_copy_queues + queue_idx = self.devices.index(cast(HCQCompiled, Device[bufs[0].device])) % self.num_copy_queues enqueue_queue = self.copy_queues.setdefault((enqueue_dev, queue_idx), enqueue_dev.hw_copy_queue_t(queue_idx=queue_idx).wait(self.kick_signals[enqueue_dev.peer_group], self.kickoff_var)) @@ -115,19 +116,19 @@ class HCQGraph(MultiGraphRunner): # Get dependencies based on input and output buffers. if is_rdma: src_qp, dest_qp = rdma_key[1].iface.connect(rdma_key[0])[:2] - sync_signals, opt_deps, rdeps = self._resolve_deps(ji.bufs[1:], [], enqueue_queue, enqueue_dev, out_signal, j, - is_copy=isinstance(ji.prg, BufferXfer), rdma_qp=src_qp) - peer_queue = self.comp_queues[peer_dev:=cast(HCQCompiled, Device[cast(Buffer, ji.bufs[0]).device])] + sync_signals, opt_deps, rdeps = self._resolve_deps(bufs[1:], [], enqueue_queue, enqueue_dev, out_signal, j, + is_copy=is_xfer, rdma_qp=src_qp) + peer_queue = self.comp_queues[peer_dev:=cast(HCQCompiled, Device[bufs[0].device])] peer_out_signal = self.signals.setdefault(peer_queue, self.pg_dev[peer_dev.peer_group].new_signal(value=0)) - peer_sync_signals, peer_opt_deps, peer_rdeps = self._resolve_deps(ji.bufs[:1], [0], peer_queue, peer_dev, peer_out_signal, j, - is_copy=isinstance(ji.prg, BufferXfer), rdma_qp=dest_qp) + peer_sync_signals, peer_opt_deps, peer_rdeps = self._resolve_deps(bufs[:1], [0], peer_queue, peer_dev, peer_out_signal, j, + is_copy=is_xfer, rdma_qp=dest_qp) self.rdma_deps[j] = (peer_queue, peer_sync_signals + peer_opt_deps, peer_out_signal, j + 1) self.last_j[peer_queue] = j else: - sync_signals, opt_deps, rdeps = self._resolve_deps(ji.bufs, cast(CompiledRunner, ji.prg).p.outs if is_exec_prg else [0], enqueue_queue, - enqueue_dev, out_signal, j, is_copy=isinstance(ji.prg, BufferXfer)) + sync_signals, opt_deps, rdeps = self._resolve_deps(bufs, prg.p.outs if prg is not None else [0], enqueue_queue, + enqueue_dev, out_signal, j, is_copy=is_xfer) - self.ji_schedule[j] = (enqueue_dev, enqueue_queue, sync_signals, opt_deps[::-1], out_signal, None if is_exec_prg else (j + 1)) + self.ji_schedule[j] = (enqueue_dev, enqueue_queue, sync_signals, opt_deps[::-1], out_signal, None if prg is not None else (j + 1)) # Collect profile information if profiling is enabled. if PROFILE: @@ -135,9 +136,9 @@ class HCQGraph(MultiGraphRunner): sig_st = prev_ji * 2 + 1 if len(opt_deps) == 0 and (prev_ji:=self.last_j[enqueue_queue]) is not None else j * 2 # Description based on the command. - prof_ji_desc = ji.prg._prg.name if is_exec_prg else TracingKey(f"{ji.bufs[1].device} -> {ji.bufs[0].device}", ret=ji.bufs[0].nbytes) # type: ignore + prof_ji_desc = prg._prg.name if prg is not None else TracingKey(f"{bufs[1].device} -> {bufs[0].device}", ret=bufs[0].nbytes) # type: ignore - prof_name = f"{enqueue_dev.device}:SDMA:{queue_idx}" if not is_exec_prg else enqueue_dev.device + prof_name = enqueue_dev.device if prg is not None else f"{enqueue_dev.device}:SDMA:{queue_idx}" self.prof_graph_entries.append(ProfileGraphEntry(prof_name, prof_ji_desc, sig_st, j * 2 + 1)) self.prof_graph_deps.append([d - 1 for _, d in rdeps]) @@ -158,7 +159,7 @@ class HCQGraph(MultiGraphRunner): self.comp_queues[dev].memory_barrier().wait(self.virt_timeline_signals[dev], self.virt_timeline_vals[dev]) \ .wait(self.kick_signals[dev.peer_group], self.kickoff_var).signal(self.signals[dev], self.kickoff_var) - for j,ji in enumerate(self.jit_cache): + for j, ((dev_idx, ast, bufs, _), prg) in enumerate(zip(self.calls, self.progs)): enqueue_dev, enqueue_queue, sync_signals, deps, signal, signal_val = self.ji_schedule[j] # Lazy allocate signals @@ -170,13 +171,13 @@ class HCQGraph(MultiGraphRunner): if PROFILE and j * 2 in self.prof_signal_is_used: enqueue_queue.timestamp(self.prof_signals[j * 2]) # Encode main commands based on ji type. - if isinstance(ji.prg, CompiledRunner): - enqueue_queue.exec(ji.prg._prg, self.ji_args[j], tuple(ji.prg.p.global_size or (1,1,1)), tuple(ji.prg.p.local_size or (1,1,1))) - elif isinstance(ji.prg, BufferXfer) and len(set(cast(HCQCompiled, Device[cast(Buffer, b).device]).peer_group for b in ji.bufs)) > 1: + if prg is not None: + enqueue_queue.exec(prg._prg, self.ji_args[j], tuple(prg.p.global_size or (1,1,1)), tuple(prg.p.local_size or (1,1,1))) # type: ignore[arg-type] + elif j in self.rdma_deps: dest_queue, dest_deps, dest_out_signal, dest_out_val = self.rdma_deps[j] for sig, val in dest_deps: dest_queue.wait(sig, val) - dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]] + dest, src = bufs[0], bufs[1] dest_dev, src_dev = cast(HCQCompiled, Device[dest.device]), cast(HCQCompiled, Device[src.device]) dest_rdma, src_rdma = dest_dev.rdma_dev(), src_dev.rdma_dev() @@ -194,10 +195,11 @@ class HCQGraph(MultiGraphRunner): dest_queue.signal(dest_out_signal, dest_out_val) self.num_rdma_ops[(dest_rdma, src_rdma)] += 1 - elif isinstance(ji.prg, (BufferXfer, BufferCopy)): - dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]] - for bufid, src in enumerate(cast(list[Buffer], ji.bufs)): - if (inprep_idx:=self.input_replace.get((j, bufid))) is not None: self.input_replace_map[enqueue_dev].add(inprep_idx) + elif ast.op is Ops.COPY: + dest, src = bufs[0], bufs[1] + uop_replace_j = dict(self.uop_replace[j]) + for bufid in range(len(bufs)): + if (replace_iidx:=uop_replace_j.get(bufid)) is not None: self.input_replace_map[enqueue_dev].add((replace_iidx, dev_idx)) else: cast(HCQAllocator, enqueue_dev.allocator).map(self.hcq_bufs[j][bufid]) enqueue_queue.copy(self.hcq_bufs[j][0], self.hcq_bufs[j][1], dest.nbytes) self.copy_to_devs[cast(HCQCompiled, Device[dest.device])].add(cast(HCQCompiled, Device[src.device])) @@ -261,7 +263,9 @@ class HCQGraph(MultiGraphRunner): def __call__(self, input_buffers: list[Buffer], var_vals: dict[str, int], wait=False, input_uops=None) -> float|None: # Map input buffers for dev in self.devices: - for idx_to_map in self.input_replace_map[dev]: cast(HCQAllocator, dev.allocator).map(input_buffers[idx_to_map]._buf) + for iidx, dev_idx in self.input_replace_map[dev]: + buf = b.bufs[dev_idx] if isinstance(b:=input_uops[iidx].buffer, MultiBuffer) else b + cast(HCQAllocator, dev.allocator).map(buf._buf) # Wait and restore signals self.kickoff_value += 1 @@ -273,8 +277,11 @@ class HCQGraph(MultiGraphRunner): **{sig.base_buf.va_addr.expr: dev.timeline_signal.base_buf.va_addr for dev, sig in self.virt_timeline_signals.items()}} # Update buffers - for (j,i),input_idx in self.input_replace.items(): - hcq_var_vals[self.input_replace_to_var[(j,i)].expr] = input_buffers[input_idx]._buf.va_addr + for j, replace in enumerate(self.uop_replace): + dev_idx = self.calls[j][0] + for pos, iidx in replace: + buf = b.bufs[dev_idx] if isinstance(b:=input_uops[iidx].buffer, MultiBuffer) else b + hcq_var_vals[self.input_replace_to_var[(j,pos)].expr] = buf._buf.va_addr for (var, qp) in self.rdma_vars.values(): hcq_var_vals[var.expr] = qp.head for q in self.rdma_queues.values(): q.submit(q.dev, hcq_var_vals) @@ -326,4 +333,4 @@ class HCQGraph(MultiGraphRunner): # MOCKGPU is not supported, since it can't execute commands in parallel is_xfer = len(set(type(d) for d in all_devs)) == 1 and hasattr(alc:=all_devs[0].allocator, '_transfer') and alc.supports_transfer return is_xfer or (all_devs[0].hw_copy_queue_t is not None and not getattr(all_devs[0], 'iface', None).__class__.__name__.startswith("MOCK")) - return new_call.src[0].op in (Ops.SINK, Ops.PROGRAM) + return new_call.src[0].op is Ops.PROGRAM diff --git a/tinygrad/runtime/graph/metal.py b/tinygrad/runtime/graph/metal.py index 6ed8dd9d1c..e09211b843 100644 --- a/tinygrad/runtime/graph/metal.py +++ b/tinygrad/runtime/graph/metal.py @@ -9,8 +9,8 @@ from tinygrad.runtime.ops_metal import wait_check, to_ns_str from tinygrad.runtime.autogen import metal class MetalGraph(GraphRunner): - def __init__(self, linear, input_buffers, input_uops=()): - super().__init__(linear, input_buffers, input_uops) + def __init__(self, linear, input_uops=()): + super().__init__(linear, input_uops) # create metal batch exec icb_descriptor = metal.MTLIndirectCommandBufferDescriptor.new() diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index 250782b080..36906b4379 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -8,7 +8,7 @@ from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filte from tinygrad.uop.ops import sint from tinygrad.device import Compiled, BufferSpec from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar, TracingKey -from tinygrad.helpers import VIZ, ceildiv, unwrap +from tinygrad.helpers import VIZ, ceildiv, unwrap, pluralize from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer from tinygrad.renderer.llvmir import AMDLLVMRenderer from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm @@ -17,6 +17,7 @@ from tinygrad.runtime.support.elf import elf_loader from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_ip_offsets, import_pmc from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE +from tinygrad.runtime.support.usb import USB3 from tinygrad.runtime.support.memory import AddrSpace if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl # noqa: F401 # pylint: disable=unused-import @@ -644,7 +645,7 @@ class AMDProgram(HCQProgram): class AMDAllocator(HCQAllocator['AMDDevice']): def __init__(self, dev:AMDDevice): super().__init__(dev, copy_bufs=getattr(dev.iface, 'copy_bufs', None), max_copyout_size=0x1000 if dev.is_usb() else None, - supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue) + supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb()) def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer: return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue) @@ -912,10 +913,10 @@ class PCIIface(PCIIfaceBase): def device_fini(self): self.dev_impl.fini() class USBIface(PCIIface): - count = 1 # TODO: support multiple usbgpus, see usb.py - def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called - self.dev, self.pci_dev, self.vram_bar = dev, USBPCIDevice(dev.__class__.__name__[:2], f"usb:{dev_id}"), 0 + if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001), "AMD")): + raise RuntimeError(f"AMD:{dev_id} does not exist ({pluralize('device', len(visible))} available)") + self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible) self.dev_impl = AMDev(self.pci_dev) self._compute_props() self.pci_dev.usb._pci_cacheable += [self.pci_dev.bar_info(2)] # doorbell region is cacheable @@ -945,15 +946,18 @@ class USBIface(PCIIface): def sleep(self, timeout): pass +def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {}) + class AMDDevice(HCQCompiled): + ifaces = [KFDIface, PCIIface, USBIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface), _mock(USBIface)] + def is_am(self) -> bool: return isinstance(self.iface, (PCIIface, USBIface)) def is_usb(self) -> bool: return isinstance(self.iface, USBIface) def __init__(self, device:str=""): self.device_id = int(device.split(":")[1]) if ":" in device else 0 - def mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {}) - self.iface = self._select_iface(KFDIface, PCIIface, USBIface, mock(KFDIface, "MOCKIface"), mock(KFDIface), mock(PCIIface), mock(USBIface)) + self.iface = self._select_iface() self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100) self.arch = "gfx%d%x%x" % self.target diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index 2acb91ad44..a5c8a1c77e 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -44,7 +44,7 @@ class CPUComputeQueue(HWQueue): def _exec(self, tid, prg, bufs, *args): vals = list(args[bufs:]) if 'core_id' in prg.runtimevars: vals[prg.runtimevars['core_id']] = tid - prg.fxn(*map(ctypes.c_uint64, args[:bufs]), *map(ctypes.c_int64 if platform.machine() == "arm64" else ctypes.c_int32, vals)) + prg.fxn(*map(ctypes.c_uint64, args[:bufs]), *map(ctypes.c_int64 if platform.machine().lower() == "arm64" else ctypes.c_int32, vals)) def _signal(self, tid, signal_addr, value): to_mv(signal_addr, 4).cast('I')[0] = value def _wait(self, tid, tmpl_sig, signal_addr, value): tmpl_sig.base_buf = HCQBuffer(signal_addr, 16, view=MMIOInterface(signal_addr, 16)) @@ -136,4 +136,4 @@ class CPUDevice(HCQCompiled): self.tasks:queue.Queue = queue.Queue() CPUWorker(self, self.tasks, thread_id=0).start() super().__init__(device, CPUAllocator(self), [ClangJITRenderer, CPULLVMRenderer, LVPRenderer], functools.partial(CPUProgram, self), CPUSignal, - CPUComputeQueue, arch={'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m)+",native") + CPUComputeQueue, arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native") diff --git a/tinygrad/runtime/ops_metal.py b/tinygrad/runtime/ops_metal.py index bc9747ba76..9eda5d580d 100644 --- a/tinygrad/runtime/ops_metal.py +++ b/tinygrad/runtime/ops_metal.py @@ -147,10 +147,10 @@ class MetalProgram: encoder.endEncoding() command_buffer.setLabel(to_ns_str(self.name)) # TODO: is this always needed? command_buffer.commit() - self.dev.mtl_buffers_in_flight.append(command_buffer) if wait: wait_check(command_buffer) return command_buffer.GPUEndTime() - command_buffer.GPUStartTime() + self.dev.mtl_buffers_in_flight.append(command_buffer) class MetalBuffer: def __init__(self, buf:metal.MTLBuffer, size:int, offset=0): self.buf, self.size, self.offset = buf, size, offset diff --git a/tinygrad/runtime/ops_null.py b/tinygrad/runtime/ops_null.py index b068e3a405..baac2fafe4 100644 --- a/tinygrad/runtime/ops_null.py +++ b/tinygrad/runtime/ops_null.py @@ -3,14 +3,18 @@ from tinygrad.device import Compiled, Allocator from tinygrad.engine.jit import MultiGraphRunner from tinygrad.renderer import Renderer, cstyle, nir, ptx, llvmir, wgsl from tinygrad.renderer.cstyle import CStyleLanguage -from tinygrad.uop.ops import Ops -from tinygrad.helpers import cpu_profile, getenv, NULL_ALLOW_COPYOUT +from tinygrad.uop.ops import UOp, Ops +from tinygrad.helpers import cpu_profile, getenv, dedup, NULL_ALLOW_COPYOUT class NullRenderer(CStyleLanguage): has_local = False float4 = "float4" barrier = "// BARRIER" code_for_op = {**CStyleLanguage.code_for_op, Ops.THREEFRY: lambda a,b,dtype: f"threefry({a},{b})", Ops.MAX: lambda a,b,dtype: f"max({a},{b})"} + def asm(self, prg: UOp, lin: UOp) -> bytes: + assert self.target.arch.startswith("gfx"), "only amd supports assembly" + from tinygrad.renderer.amd.elf import assemble_linear + return assemble_linear(prg, lin, self.target.arch) class NullProgram: def __init__(self, device:str, name:str, lib:bytes, *args, **kwargs): self.device, self.name = device, name @@ -35,4 +39,4 @@ class NullDevice(Compiled): "EMULATE is deprecated, use DEV=NULL:HIP:"+{"AMD":"gfx1100", "AMD_RDNA4":"gfx1201", "AMD_CDNA4":"gfx950"}.get(emu, "") renderers = [NullRenderer] + [r for m in [cstyle, nir, ptx, llvmir, wgsl] for r in m.__dict__.values() if inspect.isclass(r) and issubclass(r, Renderer)] - super().__init__(device, NullAllocator(self), renderers, functools.partial(NullProgram, device), NullGraph) + super().__init__(device, NullAllocator(self), dedup(renderers), functools.partial(NullProgram, device), NullGraph) diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index 0087979dac..de70781e9d 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -291,12 +291,12 @@ class NVProgram(HCQProgram): if not NAK: self.cbuf_0[188:192], self.cbuf_0[223] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window)], 0xfffdc0 qmd = {'qmd_major_version':5, 'qmd_type':nv_gpu.NVCEC0_QMDV05_00_QMD_TYPE_GRID_CTA, 'program_address_upper_shifted4':hi32(prog_addr>>4), 'program_address_lower_shifted4':lo32(prog_addr>>4), 'register_count':self.regs_usage, 'shared_memory_size_shifted7':self.shmem_usage>>7, - 'shader_local_memory_high_size_shifted4':self.lcmem_usage>>4 if NAK else self.dev.slm_per_thread>>4} + f'shader_local_memory_{"low" if NAK else "high"}_size_shifted4': self.dev.slm_per_thread>>4} else: if not NAK: self.cbuf_0[6:12] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window), *data64_le(0xfffdc0)] qmd = {'qmd_major_version':3, 'sm_global_caching_enable':1, 'program_address_upper':hi32(prog_addr), 'program_address_lower':lo32(prog_addr), 'shared_memory_size':self.shmem_usage, 'register_count_v':self.regs_usage, - **({'shader_local_memory_low_size':self.lcmem_usage} if NAK else {'shader_local_memory_high_size':self.dev.slm_per_thread})} + f'shader_local_memory_{"low" if NAK else "high"}_size':self.dev.slm_per_thread} smem_cfg = min(shmem_conf * 1024 for shmem_conf in [32, 64, 100] if shmem_conf * 1024 >= self.shmem_usage) // 4096 + 1 @@ -581,11 +581,13 @@ class PCIIface(PCIIfaceBase): class MOCKIface(NVKIface): count = 1 class NVDevice(HCQCompiled[NVSignal]): + ifaces = [NVKIface, PCIIface, MOCKIface] + def is_nvd(self) -> bool: return isinstance(self.iface, PCIIface) def __init__(self, device:str=""): self.device_id = int(device.split(":")[1]) if ":" in device else 0 - self.iface = self._select_iface(NVKIface, PCIIface, MOCKIface) + self.iface = self._select_iface() device_params = nv_gpu.NV0080_ALLOC_PARAMETERS(deviceId=self.iface.gpu_instance, hClientShare=self.iface.root, vaMode=nv_gpu.NV_DEVICE_ALLOCATION_VAMODE_OPTIONAL_MULTIPLE_VASPACES) diff --git a/tinygrad/runtime/ops_qcom.py b/tinygrad/runtime/ops_qcom.py index cd0686e806..119e6ab25e 100644 --- a/tinygrad/runtime/ops_qcom.py +++ b/tinygrad/runtime/ops_qcom.py @@ -19,14 +19,14 @@ BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2 @functools.cache def dcache_flush(): from tinygrad.uop.ops import UOp, Ops, KernelInfo - from tinygrad.codegen import get_program + from tinygrad.codegen import to_program buf, n = UOp(Ops.PARAM, dtypes.uint8.ptr(), arg=0), UOp(Ops.PARAM, dtypes.uint8.ptr(), arg=1) i = UOp.range(n.cast(dtypes.int), 0, dtype=dtypes.int) flush = UOp(Ops.CUSTOM, dtypes.void, (buf.cast(dtypes.ulong) + i.cast(dtypes.ulong) * UOp.const(dtypes.ulong, 64),), arg='__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");') sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, dtypes.void, (), arg='__asm__ volatile("dsb sy" ::: "memory");'), arg=KernelInfo(name="dcache_flush")) - ps = get_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="CPU"), UOp(Ops.LINEAR, src=tuple(sink.toposort())))), Device["CPU"].renderer) - return Device["CPU"].runtime(ps.function_name, ps.lib) + prg = to_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="CPU"), UOp(Ops.LINEAR, src=tuple(sink.toposort())))), Device["CPU"].renderer) + return Device["CPU"].runtime(prg.arg.function_name, prg.src[4].arg) #Parse C-style defines: ___SHIFT and ___MASK from the adreno module into the following format: # qreg.(=..., =..., ..., =...) @@ -320,10 +320,6 @@ class QCOMProgram(HCQProgram): reg_desc_off = _read_lib(lib, 0x34) self.fregs, self.hregs = _read_lib(lib, reg_desc_off + 0x14), _read_lib(lib, reg_desc_off + 0x18) -class QCOMTextureInfo: - def __init__(self, pitch:int, real_stride:int, desc:list[int], ibo:list[int]): - self.pitch, self.real_stride, self.desc, self.ibo = pitch, real_stride, desc, ibo - class QCOMAllocator(HCQAllocatorBase): def _alloc(self, size:int, opts:BufferSpec) -> HCQBuffer: return self.dev._gpu_map(opts.external_ptr, size) if opts.external_ptr else self.dev._gpu_alloc(size) diff --git a/tinygrad/runtime/support/c.py b/tinygrad/runtime/support/c.py index c195d945f7..bc5c9d2657 100644 --- a/tinygrad/runtime/support/c.py +++ b/tinygrad/runtime/support/c.py @@ -112,7 +112,7 @@ class DLL(ctypes.CDLL): if f.read(4) == b'\x7FELF': return str(l) def __init__(self, nm:str, paths:str|list[str], extra_paths=[], emsg="", **kwargs): - self.nm, self.emsg = nm, emsg + self.nm, self.emsg = nm, emsg or f"try setting {nm.upper()+'_PATH'}?" if (path:= DLL.findlib(nm, paths if isinstance(paths, list) else [paths], extra_paths if isinstance(extra_paths, list) else [extra_paths])): if DEBUG >= 3: print(f"loading {nm} from {path}") try: @@ -126,6 +126,7 @@ class DLL(ctypes.CDLL): def bind(self, restype, *argtypes): def wrap(fn): cfunc = None + @functools.wraps(fn) def wrapper(*args): nonlocal cfunc if cfunc is None: (cfunc:=getattr(self, fn.__name__)).argtypes, cfunc.restype = argtypes, restype @@ -134,6 +135,5 @@ class DLL(ctypes.CDLL): return wrap def __getattr__(self, nm): - if self.nm not in self._loaded_: - raise AttributeError(f"failed to load library {self.nm}: " + (self.emsg or f"try setting {self.nm.upper()+'_PATH'}?")) + if self.nm not in self._loaded_: raise AttributeError(f"failed to load library {self.nm}: {self.emsg}") return super().__getattr__(nm) diff --git a/tinygrad/runtime/support/compiler_amd.py b/tinygrad/runtime/support/compiler_amd.py index 0b9fb212b8..88cc2c6771 100644 --- a/tinygrad/runtime/support/compiler_amd.py +++ b/tinygrad/runtime/support/compiler_amd.py @@ -92,6 +92,7 @@ def compile_hip(prg:str, arch="gfx1100", asm=False) -> bytes: class HIPCompiler(Compiler): def __init__(self, arch:str): + assert comgr.dll.nm in c.DLL._loaded_, f"comgr not available: {comgr.dll.emsg}" self.arch = arch super().__init__(f"compile_hip_{self.arch}") def compile(self, src:str) -> bytes: diff --git a/tinygrad/runtime/support/hcq.py b/tinygrad/runtime/support/hcq.py index 1f60f95acf..01cb6c69aa 100644 --- a/tinygrad/runtime/support/hcq.py +++ b/tinygrad/runtime/support/hcq.py @@ -70,7 +70,6 @@ SignalType = TypeVar('SignalType', bound='HCQSignal') HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQCompiled') ProgramType = TypeVar('ProgramType', bound='HCQProgram') ArgsStateType = TypeVar('ArgsStateType', bound='HCQArgsState') -QueueType = TypeVar('QueueType', bound='HWQueue') class HWQueue(Generic[SignalType, HCQDeviceType, ProgramType, ArgsStateType]): """ @@ -490,11 +489,12 @@ class HCQCompiled(Compiled, Generic[SignalType]): buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False return buf, realloced - def _select_iface(self, *ifaces:Type): + def _select_iface(self): assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "", \ f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead" + assert hasattr(self, "ifaces"), "must have ifaces to select an iface" t = DEV.target(dev:=type(self).__name__[:-6]) - filtered = select_by_name(ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}") + filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}") filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fallback to mock ifaces return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in filtered], f"No interface for {dev}:{self.device_id} is available") diff --git a/tinygrad/runtime/support/system.py b/tinygrad/runtime/support/system.py index 6ea7000b27..52b85bc762 100644 --- a/tinygrad/runtime/support/system.py +++ b/tinygrad/runtime/support/system.py @@ -214,12 +214,13 @@ class PCIDevice: except OSError as e: raise RuntimeError(f"Cannot resize BAR {bar_idx}: {e}. Ensure the resizable BAR option is enabled.") from e class USBPCIDevice(PCIDevice): - def __init__(self, devpref:str, pcibus:str): + def __init__(self, devpref:str, dev, pcibus): + self.pcibus, self.peer_group = pcibus, f"USBPCIDevice_{pcibus}" self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock") - usb = USB3(0xADD1, 0x0001, 0x81, 0x83, 0x02, 0x04) - if DEBUG >= 1: print(f"am usb: product string: {usb.product!r}") + usb = USB3(dev, 0x81, 0x83, 0x02, 0x04) + if DEBUG >= 1: print(f"am {self.pcibus}: product string: {usb.product!r}") self.usb: CustomASM24Controller | ASM24Controller = CustomASM24Controller(usb) if usb.is_custom else ASM24Controller(usb) - self.pcibus, self._bar_info = pcibus, System.pci_setup_usb_bars(self.usb, gpu_bus=4, mem_base=0x10000000, pref_mem_base=(32 << 30)) + self._bar_info = System.pci_setup_usb_bars(self.usb, gpu_bus=4, mem_base=0x10000000, pref_mem_base=(32 << 30)) self.sram = BumpAllocator(size=0x80000, wrap=False) # asm24 controller sram def dma_view(self, ctrl_addr, size): return USBMMIOInterface(self.usb, ctrl_addr, size, fmt='B', pcimem=False) diff --git a/tinygrad/runtime/support/usb.py b/tinygrad/runtime/support/usb.py index 5e4cf586e3..c90612ca9c 100644 --- a/tinygrad/runtime/support/usb.py +++ b/tinygrad/runtime/support/usb.py @@ -1,59 +1,78 @@ -import ctypes, struct, dataclasses, array, itertools, time +import ctypes, struct, dataclasses, array, itertools, time, functools from typing import Sequence from tinygrad.runtime.autogen import libusb from tinygrad.helpers import DEBUG, DEV, to_mv, round_up, OSX, getenv, ceildiv from tinygrad.runtime.support.hcq import MMIOInterface +from tinygrad.runtime.support import c def alloc_cbuffer(sz:int) -> tuple[ctypes.Array, memoryview]: return (buf:=(ctypes.c_ubyte * sz)()), to_mv(ctypes.addressof(buf), sz) +def checked(fn, msg=None): + @functools.wraps(fn) + def wrapper(*args): + if (rc:=fn(*args)) < 0: raise RuntimeError(f"{msg or fn.__name__}: {ctypes.string_at(libusb.libusb_strerror(rc)).decode()}") + return rc + return wrapper class USB3: - def __init__(self, vendor:int, dev:int, ep_data_in:int, ep_stat_in:int, ep_data_out:int, ep_cmd_out:int, max_streams:int=31, use_bot=False): - self.vendor, self.dev = vendor, dev + @staticmethod + @functools.cache + def ctx(): + ctx = c.init_c_var(ctypes.POINTER(libusb.struct_libusb_context), checked(libusb.libusb_init)) + if DEBUG >= 6: checked(libusb.libusb_set_option)(ctx, libusb.LIBUSB_OPTION_LOG_LEVEL, 4) + return ctx + + @classmethod + @functools.cache + def list_devices(cls, vendor:int, dev:int) -> list[tuple[c.POINTER[libusb.struct_libusb_device], str]]: + ret = [] + for i in range(checked(libusb.libusb_get_device_list)(cls.ctx(), devs:=ctypes.POINTER(ctypes.POINTER(libusb.struct_libusb_device))())): + desc = c.init_c_var(libusb.struct_libusb_device_descriptor, lambda x: checked(libusb.libusb_get_device_descriptor)(devs[i], x)) + if (desc.idVendor, desc.idProduct) == (vendor, dev): + ret.append((libusb.libusb_ref_device(devs[i]), f"usb:{libusb.libusb_get_bus_number(devs[i])}-{libusb.libusb_get_device_address(devs[i])}")) + libusb.libusb_free_device_list(devs, 1) + return ret + + def __init__(self, dev:c.POINTER[libusb.struct_libusb_device], ep_data_in:int, ep_stat_in:int, ep_data_out:int, ep_cmd_out:int, + max_streams:int=31, use_bot=False): self.ep_data_in, self.ep_stat_in, self.ep_data_out, self.ep_cmd_out = ep_data_in, ep_stat_in, ep_data_out, ep_cmd_out self.max_streams, self.use_bot = max_streams, use_bot self._transferred = ctypes.c_int(0) self._bulk_in_buf, self._bulk_in_mv = alloc_cbuffer(4 << 20) self._bulk_out_buf, self._bulk_out_mv = alloc_cbuffer(4 << 20) - self.ctx = ctypes.POINTER(libusb.struct_libusb_context)() - if libusb.libusb_init(ctypes.byref(self.ctx)): raise RuntimeError("libusb_init failed") - if DEBUG >= 6: libusb.libusb_set_option(self.ctx, libusb.LIBUSB_OPTION_LOG_LEVEL, 4) - - self.handle = libusb.libusb_open_device_with_vid_pid(self.ctx, self.vendor, self.dev) - if not self.handle: raise RuntimeError(f"device {self.vendor:04x}:{self.dev:04x} not found. sudo required?") + self.handle = c.init_c_var(c.POINTER[libusb.struct_libusb_device_handle], lambda x: checked(libusb.libusb_open)(dev, x)) # Read product string descriptor _buf = (ctypes.c_ubyte * 256)() _desc = libusb.struct_libusb_device_descriptor() - libusb.libusb_get_device_descriptor(libusb.libusb_get_device(self.handle), ctypes.byref(_desc)) - _ret = libusb.libusb_get_string_descriptor_ascii(self.handle, _desc.iProduct, _buf, 256) - self.product = bytes(_buf[:max(_ret, 0)]).decode("ascii", errors="replace") if _ret > 0 else "" + checked(libusb.libusb_get_device_descriptor)(libusb.libusb_get_device(self.handle), ctypes.byref(_desc)) + _ret = checked(libusb.libusb_get_string_descriptor_ascii)(self.handle, _desc.iProduct, _buf, 256) + self.product = bytes(_buf[:_ret]).decode("ascii", errors="replace") self.is_custom = self.product.startswith("custom") if self.is_custom: self.use_bot = use_bot = True # Detach kernel driver if needed - if libusb.libusb_kernel_driver_active(self.handle, 0): - libusb.libusb_detach_kernel_driver(self.handle, 0) - libusb.libusb_reset_device(self.handle) + if checked(libusb.libusb_kernel_driver_active)(self.handle, 0): + checked(libusb.libusb_detach_kernel_driver)(self.handle, 0) + checked(libusb.libusb_reset_device)(self.handle) # Set configuration and claim interface - if libusb.libusb_set_configuration(self.handle, 1): raise RuntimeError("set_configuration failed") - if libusb.libusb_claim_interface(self.handle, 0): raise RuntimeError("claim_interface failed. sudo required?") + checked(libusb.libusb_set_configuration)(self.handle, 1) + checked(libusb.libusb_claim_interface)(self.handle, 0) if use_bot: - libusb.libusb_set_interface_alt_setting(self.handle, 0, 0) + checked(libusb.libusb_set_interface_alt_setting)(self.handle, 0, 0) self._tag = 0 else: - if libusb.libusb_set_interface_alt_setting(self.handle, 0, 1): raise RuntimeError("alt_setting failed") + checked(libusb.libusb_set_interface_alt_setting)(self.handle, 0, 1) # Clear any stalled endpoints all_eps = (self.ep_data_out, self.ep_data_in, self.ep_stat_in, self.ep_cmd_out) - for ep in all_eps: libusb.libusb_clear_halt(self.handle, ep) + for ep in all_eps: checked(libusb.libusb_clear_halt)(self.handle, ep) # Allocate streams stream_eps = (ctypes.c_uint8 * 3)(self.ep_data_out, self.ep_data_in, self.ep_stat_in) - if (rc:=libusb.libusb_alloc_streams(self.handle, self.max_streams * len(stream_eps), stream_eps, len(stream_eps))) < 0: - raise RuntimeError(f"alloc_streams failed: {rc}") + checked(libusb.libusb_alloc_streams)(self.handle, self.max_streams * len(stream_eps), stream_eps, len(stream_eps)) # Base cmd cmd_template = bytes([0x01, 0x00, 0x00, 0x01, *([0] * 12), 0xE4, 0x24, 0x00, 0xB2, 0x1A, 0x00, 0x00, 0x00, *([0] * 8)]) @@ -77,11 +96,11 @@ class USB3: return tr def _submit_and_wait(self, cmds): - for tr in cmds: libusb.libusb_submit_transfer(tr) + for tr in cmds: checked(libusb.libusb_submit_transfer)(tr) running = len(cmds) while running: - libusb.libusb_handle_events(self.ctx) + checked(libusb.libusb_handle_events)(USB3.ctx()) running = len(cmds) for tr in cmds: if tr.contents.status == libusb.LIBUSB_TRANSFER_COMPLETED: running -= 1 @@ -90,14 +109,12 @@ class USB3: def _bulk_out(self, ep: int, payload: bytes, timeout: int = 1000): if len(payload) > len(self._bulk_out_mv): self._bulk_out_buf, self._bulk_out_mv = alloc_cbuffer(len(payload)) self._bulk_out_mv[:len(payload)] = payload - rc = libusb.libusb_bulk_transfer(self.handle, ep, self._bulk_out_buf, len(payload), ctypes.byref(self._transferred), timeout) - assert rc == 0, f"bulk OUT 0x{ep:02X} failed: {rc}" + checked(libusb.libusb_bulk_transfer, f"bulk OUT 0x{ep:02X} failed")(self.handle, ep, self._bulk_out_buf, len(payload), self._transferred, timeout) assert self._transferred.value == len(payload), f"bulk OUT short write on 0x{ep:02X}: {self._transferred.value}/{len(payload)} bytes" def _bulk_in(self, ep: int, length: int, timeout: int = 1000) -> memoryview: if length > len(self._bulk_in_mv): self._bulk_in_buf, self._bulk_in_mv = alloc_cbuffer(length) - rc = libusb.libusb_bulk_transfer(self.handle, ep, self._bulk_in_buf, length, ctypes.byref(self._transferred), timeout) - assert rc == 0, f"bulk IN 0x{ep:02X} failed: {rc}" + checked(libusb.libusb_bulk_transfer, f"bulk IN 0x{ep:02X} failed")(self.handle, ep, self._bulk_in_buf, length, self._transferred, timeout) return self._bulk_in_mv[:self._transferred.value] def send_batch(self, cdbs:list[bytes], idata:list[int]|None=None, odata:list[bytes|None]|None=None) -> list[bytes|None]: @@ -172,7 +189,11 @@ class ScsiWriteOp: data:bytes; lba:int=0 # noqa: E702 class CustomASM24Controller: def __init__(self, usb:USB3|None=None): - self.usb = usb or USB3(0xADD1, 0x0001, 0x81, 0x83, 0x02, 0x04, use_bot=True) + if not usb: + devs = USB3.list_devices(0xADD1, 0x0001) + assert len(devs), "no ASM24 controller found" + self.usb = USB3(devs[0][0], 0x81, 0x83, 0x02, 0x04, use_bot=True) + else: self.usb = usb self._pci_cacheable: list[tuple[int, int]] = [] self._pci_cache: dict[int, int|None] = {} @@ -186,8 +207,8 @@ class CustomASM24Controller: if ltssm != 0x78: raise RuntimeError(f"PCIe link not up (LTSSM=0x{ltssm:02X}), custom firmware not ready") def set_pcie_power(self, enabled:bool, timeout:int=10000): - ret = libusb.libusb_control_transfer(self.usb.handle, 0x40, 0xF3, int(enabled), 0, None, 0, timeout) - assert ret >= 0, f"F3 PCIe power {'on' if enabled else 'off'} failed: {ret}" + checked(libusb.libusb_control_transfer, + f"F3 PCIe power {'on' if enabled else 'off'} failed")(self.usb.handle, 0x40, 0xF3, int(enabled), 0, None, 0, timeout) # === PCIe TLP via 0xF0 vendor command === @@ -267,8 +288,8 @@ class CustomASM24Controller: def write(self, base_addr:int, data:bytes, **kwargs): """Write to chip XDATA via vendor control OUT (bRequest=0xE5). wValue=addr, wIndex=val.""" for off, val in enumerate(data): - ret = libusb.libusb_control_transfer(self.usb.handle, 0x40, 0xE5, base_addr + off, val, None, 0, 1000) - assert ret >= 0, f"write(0x{base_addr + off:04X}, 0x{val:02X}) failed: {ret}" + checked(libusb.libusb_control_transfer, + f"write(0x{base_addr + off:04X}, 0x{val:02X}) failed")(self.usb.handle, 0x40, 0xE5, base_addr + off, val, None, 0, 1000) def scsi_write(self, buf:bytes, lba:int=0): """Write to SRAM via 0xF2 vendor command + bulk OUT.""" @@ -277,20 +298,23 @@ class CustomASM24Controller: num_slots = round_up(len(buf_padded), 0x4000) // 0x4000 # 16KB per slot # 0xF2 OUT: wValue=sectors, wIndex=start_slot|(num_slots<<8) windex = (num_slots & 0xFF) << 8 - ret = libusb.libusb_control_transfer(self.usb.handle, 0x40, 0xF2, sectors, windex, None, 0, 1000) - assert ret >= 0, f"F2 setup failed: {ret}" + checked(libusb.libusb_control_transfer, "F2 setup failed")(self.usb.handle, 0x40, 0xF2, sectors, windex, None, 0, 1000) self.usb._bulk_out(0x02, buf_padded) def scsi_read_arm(self, size:int): windex = (ceildiv(size, 0x4000) & 0xFF) << 8 - ret = libusb.libusb_control_transfer(self.usb.handle, 0x40, 0xF2, (ceildiv(size, 512) & 0x7FFF) | 0x8000, windex, None, 0, 1000) - assert ret >= 0, f"F2 read arm failed: {ret}" + checked(libusb.libusb_control_transfer, + "F2 read arm failed")(self.usb.handle, 0x40, 0xF2, (ceildiv(size, 512) & 0x7FFF) | 0x8000, windex, None, 0, 1000) def scsi_read(self, size:int) -> memoryview: return self.usb._bulk_in(0x81, round_up(size, 512), timeout=10000)[:size] class ASM24Controller: def __init__(self, usb:USB3|None=None): - self.usb = usb or USB3(0xADD1, 0x0001, 0x81, 0x83, 0x02, 0x04, use_bot=bool(getenv("USE_BOT", 0))) + if not usb: + devs = USB3.list_devices(0xADD1, 0x0001) + assert len(devs), "no ASM24 controller found" + self.usb = USB3(devs[0][0], 0x81, 0x83, 0x02, 0x04, use_bot=bool(getenv("USE_BOT", 0))) + else: self.usb = usb self._cache: dict[int, int|None] = {} self._pci_cacheable: list[tuple[int, int]] = [] self._pci_cache: dict[int, int|None] = {} diff --git a/tinygrad/schedule/__init__.py b/tinygrad/schedule/__init__.py index 8cb578f6e8..d4d4589f0a 100644 --- a/tinygrad/schedule/__init__.py +++ b/tinygrad/schedule/__init__.py @@ -1,12 +1,8 @@ import time, inspect -from typing import cast from collections import deque -from dataclasses import replace -from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink, KernelInfo +from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink, KernelInfo from tinygrad.uop.spec import type_verify, tensor_spec -from tinygrad.device import Buffer, MultiBuffer -from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, flatten, BEAM, partition -from tinygrad.engine.realize import ExecItem +from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, partition # **** schedule linearizer @@ -70,33 +66,8 @@ def create_schedule(sched_sink:UOp) -> UOp: if in_degree[x] == 0: queue.append(x) return UOp(Ops.LINEAR, src=tuple(linearized)) -def linear_to_schedule(linear:UOp) -> list[ExecItem]: - """Convert a LINEAR UOp to a list of ExecItems.""" - schedule: list[ExecItem] = [] - for si in linear.src: - ast, buf_uops = si.src[0], si.src[1:] - # create subbuffers if needed - if ast.op is Ops.BUFFER_VIEW: - base = buf_uops[1].buffer - assert isinstance(base, Buffer), "base can't be MultiBuffer" - buffers[buf_uops[0]] = base.view(buf_uops[0].arg, ast.dtype, ast.arg[1]*base.dtype.itemsize) - # set beam on KernelInfo when beam search is enabled - if ast.op is Ops.SINK and BEAM >= 1 and ast.arg.beam == 0: ast = ast.replace(arg=replace(ast.arg, beam=BEAM.value)) - ubufs = [b.buffer for b in buf_uops if b.op is not Ops.BIND] - metadata = si.arg.metadata - if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": - schedule.append(ExecItem(ast, flatten([b.bufs if isinstance(b, MultiBuffer) else [b] for b in ubufs]), metadata)) - elif any(isinstance(x, MultiBuffer) for x in ubufs): - assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer" - dnums = [x for x in ast.variables() if x.expr == '_device_num'] - for j, bufs in enumerate(zip(*[x.bufs for x in cast(tuple[MultiBuffer, ...], ubufs)])): - schedule.append(ExecItem(ast, list(bufs), metadata, {dnums[0].expr:j} if len(dnums) else {})) - else: - schedule.append(ExecItem(ast, cast(list[Buffer|None], ubufs), metadata)) - return schedule - from tinygrad.schedule.memory import memory_plan_rewrite -from tinygrad.engine.realize import capturing +from tinygrad.engine.realize import capturing, pm_flatten_linear from tinygrad.schedule.rangeify import get_kernel_graph from tinygrad.helpers import CAPTURING from tinygrad.uop.ops import PatternMatcher, UPat @@ -115,10 +86,7 @@ pm_resolve_linear_call = PatternMatcher([ # call LINEAR is resolved here (UPat(Ops.CALL, src=(UPat(Ops.LINEAR),), name="linear_call", allow_any_len=True), lambda linear_call: graph_rewrite(linear_call.src[0], pm_post_sched_cache, ctx=({}, linear_call.src[1:]), walk=True, name="params to buffers")), - # LINEAR on LINEAR - (UPat(Ops.LINEAR, custom_early_reject={Ops.LINEAR}, name="x"), - lambda x: x.replace(src=tuple(flatten(x.src if x.op is Ops.LINEAR else (x,) for x in x.src)))), -]) +])+pm_flatten_linear schedule_cache: dict[bytes, UOp] = {} # ctx is just for DEBUG on inner diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 699aa1423e..e7833e03fc 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -10,9 +10,9 @@ from tinygrad.helpers import argfix, flatten, prod, all_int, round_up, getenv, a from tinygrad.helpers import resolve_pool_pads, IMAGE, FLOAT16, WINO, Metadata, TRACEMETA, is_numpy_ndarray, TracingKey, cpu_profile from tinygrad.helpers import suppress_finalizing, disable_gc from tinygrad.gradient import compute_gradient -from tinygrad.mixin import OpMixin, ReductionStr -from tinygrad.uop.ops import smax, UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable, _broadcast_shape -from tinygrad.schedule import ExecItem, create_linear_with_vars, linear_to_schedule +from tinygrad.mixin import OpMixin +from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable, _broadcast_shape +from tinygrad.schedule import create_linear_with_vars from tinygrad.device import Buffer, canonicalize_device from tinygrad.engine.realize import run_linear from tinygrad.callify import transform_to_call @@ -77,15 +77,6 @@ def _apply_winograd_matrix(mat, t:Tensor, dims:int) -> Tensor: assert isinstance(ret, Tensor), "sum didn't return a Tensor" return ret -def _masked_setitem(target:Tensor, values:Tensor, mask:Tensor, axes:tuple[int, ...]) -> Tensor: - # reduce such that if mask contains repeated indices the last one remains - 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) - # select from values for each True element in mask else select from target - return mask.where(values, target) - class Tensor(OpMixin): """ A `Tensor` is a multi-dimensional matrix containing elements of a single data type. @@ -241,20 +232,11 @@ class Tensor(OpMixin): _apply_map_to_tensors(becomes_map, name="buffers") return create_linear_with_vars(big_sink) - def schedule_with_vars(self, *lst:Tensor) -> tuple[list[ExecItem], dict[str, int]]: - """ - Creates the schedule needed to realize these Tensor(s), with Variables. - - NOTE: A Tensor can only be scheduled once. - """ - linear, var_vals = self.linear_with_vars(*lst) - return linear_to_schedule(linear), var_vals - - def schedule(self, *lst:Tensor) -> list[ExecItem]: + def schedule_linear(self, *lst:Tensor) -> UOp: """Creates the schedule needed to realize these Tensor(s).""" - schedule, var_vals = self.schedule_with_vars(*lst) + linear, var_vals = self.linear_with_vars(*lst) assert len(var_vals) == 0 - return schedule + return linear @disable_gc() def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor: @@ -928,71 +910,6 @@ class Tensor(OpMixin): def _mop(self, op:Ops, arg) -> Tensor: return self._apply_uop(UOp._mop, extra_args=(op,), arg=arg) def _rop(self, op:Ops, axis:tuple[int, ...]) -> Tensor: return self._apply_uop(UOp._rop, op=op, axis=axis) - def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Tensor: - if any(pB>sh or pA>sh for (pB,pA),sh in zip(pX, self.shape)): raise ValueError('Padding value causes wrapping around more than once.') - if any(pB<0 or pA<0 for pB,pA in pX): raise NotImplementedError("Negative pads with circular pads is not supported") - orig_shape, X = self.shape, self.repeat(tuple(1 + bool(pB) + bool(pA) for pB,pA in pX)) - return X.shrink(tuple((0 if pB == 0 else osh-pB, xsh if pA == 0 else xsh-osh+pA) for (pB,pA),osh,xsh in zip(pX, orig_shape, X.shape))) - - def _pad_reflect_replicate(self, pX:tuple[tuple[sint, sint], ...], mode:str) -> Tensor: - X, pads = self, tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX) - for d,(pB,pA) in enumerate(pads): - if mode == "reflect": - if pB >= (s:=X.shape[d]) or pA>=s: raise ValueError(f"Padding ({pB}, {pA}) should be less than the input size={s} for dim={d}.") - slcB, slcA = slice(pB,0,-1), slice(s-2 if s-2>=0 else None, s-2-pA if s-2-pA>=0 else None, -1) - xB, xA = (X[[slc if i == d else slice(None) for i in range(X.ndim)]] if p > 0 else None for slc, p in ((slcB, pB), (slcA, pA))) - else: - shrB, shrA = tuple((0,1) if i==d else None for i in range(X.ndim)), tuple((X.shape[i]-1,X.shape[i]) if i==d else None for i in range(X.ndim)) - xB, xA = (X.shrink(shr).expand(tuple(p if i==d else None for i in range(X.ndim))) if p > 0 else None for shr, p in ((shrB, pB), (shrA, pA))) - X = Tensor.cat(*(X_ for X_ in (xB, X, xA) if X_ is not None), dim=d) - # shrink after for negative pads (reflection/replication must see full data first) - return X.shrink(tuple((-min(pB,0), min(pA+s,s)) for (pB,pA),s in zip(pX, X.shape))) - - def pad(self, padding:Sequence[sint]|Sequence[tuple[sint, sint]|None], mode:str="constant", value:float=0.0) -> Tensor: - """ - Returns a tensor with padding applied based on the input `padding`. - - `padding` supports two padding structures: - - 1. Flat padding: `(padding_left, padding_right, padding_top, padding_bottom, ...)` - - This structure matches PyTorch's pad. - - `padding` length must be even. - - 2. Group padding: `(..., (padding_top, padding_bottom), (padding_left, padding_right))` - - This structure matches pad for JAX, NumPy, TensorFlow, and others. - - For each axis, padding can be `None`, meaning no padding, or a tuple `(start, end)`. - - `padding` must have the same length as `self.ndim`. - - Padding values can be negative, resulting in dimension shrinks that work similarly to Python negative slices. - Padding modes is selected with `mode` which supports `constant`, `reflect` and `replicate`. - - ```python exec="true" source="above" session="tensor" result="python" - t = Tensor.arange(9).reshape(1, 1, 3, 3) - print(t.numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - print(t.pad((1, 2, 0, -1)).numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - print(t.pad(((None, None, (0, -1), (1, 2)))).numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - print(t.pad((1, 2, 0, -1), value=-float('inf')).numpy()) - ``` - """ - # normalize to grouped format - if all(isinstance(p, (int,UOp)) for p in padding): - if len(padding)%2 != 0: raise ValueError("Flat padding must have even number of pads") - pX = ((0,0),)*(self.ndim - len(padding)//2) + flat_to_grouped(cast(Sequence[sint], padding)) - else: pX = tuple((0,0) if p is None else p for p in cast(Sequence[tuple[sint, sint]|None], padding)) - if len(pX) != self.ndim: raise ValueError(f"padding length is improper, {padding=} {self.ndim=}") - # dispatch - if mode == "constant": return self._pad_constant(pX, value) - assert all_int(self.shape), f"does not support symbolic shape {self.shape}" - if mode == "circular": return self._pad_circular(pX) - if mode in {"reflect", "replicate"}: return self._pad_reflect_replicate(pX, mode) - raise NotImplementedError(f"{mode=} is not supported") - def _getitem(self, indices, v: Tensor|None = None) -> Tensor: # view-only indexing (no Tensor/list indices, no setitem) is handled by MovementMixin.__getitem__ if v is None and not any(isinstance(i, (Tensor, list, tuple)) for i in (indices if isinstance(indices, tuple) else (indices,))): @@ -1063,7 +980,7 @@ class Tensor(OpMixin): vb = v.cast(self.dtype)._broadcast_to(_broadcast_shape(x.shape, v.shape)) for dim in sum_axis: vb = vb.unsqueeze(dim) # add back reduced dims from sum start = dims[0] if not permuted else 0 - vb = _masked_setitem(x_pre, vb, mask, tuple(range(start, start + len(big_shape)))) + vb = x_pre._masked_merge(vb, mask, tuple(range(start, start + len(big_shape)))) elif v is None: return x # basic getitem # basic setitem: broadcast v, reshape to self.ndim (unsqueeze int dims, squeeze None dims) else: vb = v.cast(self.dtype)._broadcast_to(x.shape) @@ -1203,12 +1120,6 @@ class Tensor(OpMixin): # ***** reduce ops ***** - def allclose(self, other:Tensor, rtol:float=1e-05, atol:float=1e-08, equal_nan=False) -> bool: - """ - Check if all self and other are close. Return True or False. - """ - return bool(self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan).all().item()) - def keccak(self, cfg:str|tuple[int, int]="sha3_256"): """ Calculates a Keccak hash over the last dimension. Uses "sha3_256" by default. @@ -1371,124 +1282,6 @@ class Tensor(OpMixin): if IMAGE: return self.image_dot(w, dtype) return super().dot(w, dtype) - def scatter(self, dim:int, index:Tensor, src:Tensor|PyConst, reduce:Literal['multiply', 'add']|None=None) -> Tensor: - """ - Scatters `src` values along an axis specified by `dim`. - Apply `add` or `multiply` reduction operation with `reduce`. - - NOTE: To use the `reduce` argument with a Tensor `src`, see `Tensor.scatter_reduce`. - - ```python exec="true" source="above" session="tensor" result="python" - src = Tensor.arange(1, 11).reshape(2, 5) - print(src.numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - index = Tensor([[0, 1, 2, 0]]) - print(Tensor.zeros(3, 5, dtype=src.dtype).scatter(0, index, src).numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - index = Tensor([[0, 1, 2], [0, 1, 4]]) - print(Tensor.zeros(3, 5, dtype=src.dtype).scatter(1, index, src).numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - print(Tensor.full((2, 4), 2.0).scatter(1, Tensor([[2], [3]]), 1.23, reduce='multiply').numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - print(Tensor.full((2, 4), 2.0).scatter(1, Tensor([[2], [3]]), 1.23, reduce='add').numpy()) - ``` - """ - if reduce not in {None, "add", "multiply"}: raise TypeError(f"{reduce=} must be one of None, 'multiply', or 'add'") - if reduce and isinstance(src, Tensor): raise TypeError("Tensor src is not supported with reduce arg. see scatter_reduce") - if not isinstance(src, Tensor): src = index.full_like(src, device=self.device, dtype=self.dtype) - if reduce == "add": return self.scatter_reduce(dim, index, src, "sum", include_self=True) - if reduce == "multiply": return self.scatter_reduce(dim, index, src, "prod", include_self=True) - src, mask = self._pre_scatter(dim, index, src) - return _masked_setitem(self, src, mask, (-1,)) - - def sort(self, dim:int=-1, descending:bool=False) -> tuple[Tensor, Tensor]: - """ - Performs a bitonic sort on the tensor along the specified dimension. - - Order of indices for equivalent elements is always preserved. - - See: https://en.wikipedia.org/wiki/Bitonic_sorter - - ```python exec="true" source="above" session="tensor" result="python" - t = Tensor([[0.1, 0.5, 1.2, 3.4, 2.1], [2.2, 1.9, 0.3, 4.5, 0.8]]) - print(t.numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - sorted_values, indices = t.sort(dim=1, descending=True) - print(sorted_values.numpy()) - print(indices.numpy()) - ``` - """ - x, dim = self, self._resolve_dim(dim) - if (orig_len := int(x.shape[dim])) <= 1: return x, x.zeros_like(dtype=dtypes.default_int) - # pad to power of 2 - n_stages = (orig_len-1).bit_length() - pads = tuple((0, 2**n_stages - orig_len) if i == dim else None for i in range(x.ndim)) - x = x.pad(pads, value=x.dtype.min if descending else x.dtype.max).unflatten(dim, (2,)*n_stages) - # https://en.wikipedia.org/wiki/Bitonic_sorter#/media/File:BitonicSort1.svg - for stage in range(1, n_stages+1): - if stage != n_stages: - # flip so arrows of green boxes point the same way as blue boxes - crossover_dim = dim + n_stages - stage - 1 - blue_box, green_box = x.split(1, crossover_dim) - flip_dims = tuple(-i for i in range(1, stage+1+(self.ndim-dim))) - x = (blue_box.cat(green_box.flip(flip_dims), dim=crossover_dim)).contiguous() - for substage in range(stage-1, -1, -1): - partner_dim = dim + n_stages - substage - 1 - x_top, x_bottom = x.split(1, partner_dim) - x_larger, x_smaller = x_top.maximum(x_bottom), x_top.minimum(x_bottom) - x = (x_larger.cat(x_smaller, dim=partner_dim) if descending else x_smaller.cat(x_larger, dim=partner_dim)).contiguous() - if stage != n_stages: - # flip wires back to undo the crossover - blue_box, flipped_green_box = x.split(1, crossover_dim) - x = blue_box.cat(flipped_green_box.flip(flip_dims), dim=crossover_dim) - x = x.flatten(dim, dim+n_stages-1).shrink_to(self.shape) - # compute indices for sorted values - mask = Tensor.ones(orig_len, orig_len, dtype=dtypes.bool, device=self.device).tril().reshape((None, None) + (1,)*(self.ndim-dim-1)) - def compute_counts(t:Tensor): return (mask & (t.unsqueeze(dim) == t.unsqueeze(dim+1))).sum(dim+1) - count_orig, count_sorted = compute_counts(self), compute_counts(x) - cond = (self.unsqueeze(dim+1) == x.unsqueeze(dim)) & (count_orig.unsqueeze(dim+1) == count_sorted.unsqueeze(dim)) - idx = Tensor.arange(orig_len, device=self.device).reshape(tuple(orig_len if i == dim else 1 for i in range(x.ndim))) - idx = (cond * idx.unsqueeze(dim+1)).sum(dim) - return x, idx - - def argsort(self, dim:int=-1, descending:bool=False) -> Tensor: - """ - Returns the indices that sort input tensor along given `dimension` in given `descending` order by value. - - ```python exec="true" source="above" session="tensor" result="python" - t = Tensor([[2, 3, 4, 1], [1, 4, 3, 2]]) - print(t.argsort().numpy()) - ``` - """ - return self.sort(dim, descending)[1] - - def topk(self, k:int, dim:int=-1, largest:bool=True, sorted_:bool=True) -> tuple[Tensor, Tensor]: - """ - Computes the top-k elements of the tensor along the specified `dim`. - - Order of indices for equivalent elements is always preserved. - - ```python exec="true" source="above" session="tensor" result="python" - t = Tensor([[0.1, 0.5, 1.2, 3.4, 2.1], [2.2, 1.9, 0.3, 4.5, 0.8]]) - print(t.numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - topk_values, topk_indices = t.topk(2, dim=1) - print(topk_values.numpy()) - print(topk_indices.numpy()) - ``` - """ - if not sorted_: raise NotImplementedError("topk with sorted_=False is not supported") - if k > self.shape[dim:=self._resolve_dim(dim)]: raise ValueError(f"selected index {k=} is out of range") - x, idx = self.sort(dim, descending=largest) - topk_shape = tuple(k if i == dim else None for i in range(self.ndim)) - return x.shrink_to(topk_shape), idx.shrink_to(topk_shape) - # ***** unary ops ***** def contiguous(self, *args, **kwargs) -> Tensor: @@ -1666,30 +1459,6 @@ class Tensor(OpMixin): qk = qk + attn_mask return qk.cast(self.dtype).softmax(-1).dropout(dropout_p) @ value - def nll_loss(self, Y:Tensor, weight:Tensor|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean") -> Tensor: - """ - Computes the negative log likelihood loss between log-probabilities and target labels. - - NOTE: `self` is log-probabilities and `Y` is the Y labels or class probabilities. - - See: https://pytorch.org/docs/stable/generated/torch.nn.functional.nll_loss.html - - ```python exec="true" source="above" session="tensor" result="python" - t = Tensor([[-1, 2, -3], [1, -2, 3]]) - Y = Tensor([1, 2]) - print(t.log_softmax().nll_loss(Y).item()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - t = Tensor([[-1, 2, -3], [1, -2, 3]]) - Y = Tensor([1, 2]) - print(t.log_softmax().nll_loss(Y, reduction='none').numpy()) - ``` - """ - weight = Y.ones_like(requires_grad=False) if weight is None else weight[Y] - masked_weight = weight if ignore_index is None else weight * (Y != ignore_index) - nll = -self.gather(1, Y.unsqueeze(1)).squeeze(1) * masked_weight - return nll.sum() / masked_weight.sum() if reduction == "mean" else nll._do_reduction(reduction) - def qr(self) -> tuple[Tensor, Tensor]: assert self.ndim > 1, f"expected two or more dimensions, got {self.ndim}" b_shape, m, n = self.shape[:-2], int(self.shape[-2]), int(self.shape[-1]) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 775e06d8d6..1b10724b52 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -348,9 +348,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass): if not isinstance(self.device, tuple) or self.axis is None: return self.max_shape return tuple(x//len(self.device) if i == self.axis else x for i,x in enumerate(self.max_shape)) - @property - def shard_size(self) -> int: return prod(self.max_shard_shape) - @functools.cached_property def ended_ranges(self) -> tuple[UOp, ...]: if self.op in range_start: return self.src[range_start[self.op]:] @@ -978,6 +975,51 @@ class KernelInfo: @property def function_name(self): return to_function_name(self.name) +@dataclass(frozen=True) +class ProgramInfo: + name: str = "test" + global_size: tuple[int|float, ...] = (1, 1, 1) + local_size: tuple[int, ...]|None = None + vars: tuple[UOp, ...] = () + globals: tuple[int, ...] = () + outs: tuple[int, ...] = () + ins: tuple[int, ...] = () + aux: tuple = () + + @property + def function_name(self): return to_function_name(self.name) + + @property + def runtimevars(self) -> dict[str, int]: return {v.expr: i for i, v in enumerate(self.vars) if v.expr == 'core_id'} + + def launch_dims(self, var_vals:dict[str, int]): + global_size = [sym_infer(sz, var_vals) for sz in self.global_size] # type: ignore[arg-type] + local_size = [sym_infer(sz, var_vals) for sz in self.local_size] if self.local_size is not None else None + return global_size, local_size + + @staticmethod + def from_sink(sink:UOp, aux:tuple=()) -> ProgramInfo: + _vars: list[UOp] = [] + _globals: list[int] = [] + outs: list[int] = [] + ins: list[int] = [] + global_size: list[int] = [1, 1, 1] + local_size: list[int]|None = [1, 1, 1] + for u in sink.toposort(): + if u.op is Ops.DEFINE_VAR: _vars.append(u) + if u.op is Ops.PARAM: _globals.append(u.arg) + if u.op in (Ops.STORE, Ops.LOAD): + if (idx:=u.src[0]).op is Ops.INDEX or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX): + if (buf:=idx.src[0]).op is Ops.PARAM: (outs if u.op is Ops.STORE else ins).append(buf.arg) + if u.op is Ops.SPECIAL: + if u.arg[0] == 'i': local_size = None + special_size = local_size if u.arg[0] == 'l' else global_size + if special_size is not None: special_size[int(u.arg[-1])] = cast(int, u.src[0].ssimplify()) + if u.op is Ops.DEFINE_VAR and u.arg[0] == 'core_id': global_size[0] = u.arg[2] + 1 + return ProgramInfo(sink.arg.name if isinstance(sink.arg, KernelInfo) else "test", tuple(global_size), + tuple(local_size) if local_size is not None else None, tuple(sorted(_vars, key=lambda v: v.arg)), + tuple(sorted(dedup(_globals))), tuple(sorted(dedup(outs))), tuple(sorted(dedup(ins))), aux) + @dataclass(frozen=True) class CallInfo: grad_fxn: Callable|None = None diff --git a/tinygrad/viz/README b/tinygrad/viz/README index 02a73b8841..3227159b52 100644 --- a/tinygrad/viz/README +++ b/tinygrad/viz/README @@ -21,7 +21,7 @@ By default, VIZ UIs automatically load the latest files. user story: viewing profiling data * tinygrad ran 32 LLM decode steps: web: click "profiler", view the timeline of all python codegen and GPU kernels. -cli: Run `DEBUG=3 python -m tinygrad.viz.cli --profile -s ALL --json` to extract kernel timing info and ASTs in JSON format. +cli: Run `DEBUG=3 python -m tinygrad.viz.cli --json` to extract kernel timing info and ASTs in JSON format. - note: Make sure to add NO_COLOR=1 to disable colored output. user story: viewing code @@ -30,11 +30,11 @@ user story: viewing code * schedule 2 (97) = main.py:97 * schedule 3 (10) = main.py:145 * web: click "schedule 1", get list of kernels (like DEBUG=2) -* cli: `python -m tinygrad.viz.cli --rewrites -s "schedule 1"` +* cli: `python -m tinygrad.viz.cli -s TINY -i "Schedule 3 Kernels n1"` * kernel 1 "E_34_34" -- 'sin' * kernel 2 "R_4545" * web: click "E_34_34" -* cli: `python -m tinygrad.viz.cli --rewrites -s "E_34_34" -i "initial symbolic"` +* cli: `python -m tinygrad.viz.cli -s TINY -i "do_to_program for E_34_34" "initial symbolic"` * pre-rewritten UOp graph (step through rewrite here) * post-rewritten UOp graph * UOp list @@ -54,8 +54,8 @@ note: SQTT has additional overhead, to enable it, set VIZ=2. * tinygrad ran custom assembly GEMM kernel. * web: click "SQTT gemm SE:1 PKTS", see wave instruction scheduling and CU execution unit occupancy at every clock cycle. -* cli: python -m tinygrad.viz.cli --profile -s "kernel SQTT SE:0 PKTS" +* cli: python -m tinygrad.viz.cli -s "kernel SQTT SE:0 PKTS" * get bank conflicts: * web: click "gemm PMC" -* cli: python -m tinygrad.viz.cli -p -s "gemm PMC" | rg -A 16 SQC_LDS_BANK_CONFLICT +* cli: python -m tinygrad.viz.cli -s "gemm PMC" | rg -A 16 SQC_LDS_BANK_CONFLICT diff --git a/tinygrad/viz/cli.py b/tinygrad/viz/cli.py index e678349e98..9da63fbdb9 100755 --- a/tinygrad/viz/cli.py +++ b/tinygrad/viz/cli.py @@ -1,12 +1,12 @@ #!/usr/bin/env python3 -import argparse, pathlib, signal, sys, struct, json, os, itertools, heapq +import argparse, pathlib, signal, struct, json, os, itertools, heapq os.environ["VIZ"] = "0" if hasattr(signal, "SIGPIPE"): signal.signal(signal.SIGPIPE, signal.SIG_DFL) from typing import Iterator from tinygrad.viz import serve as viz from tinygrad.uop.ops import RewriteTrace from tinygrad.helpers import temp, ansistrip, colored, time_to_str, ansilen, ProfilePointEvent, ProfileRangeEvent, TracingKey, unwrap, NO_COLOR -from tinygrad.helpers import DEBUG +from tinygrad.helpers import DEBUG, Context # profile decoder used in CLI and tests def decode_profile(data:bytes) -> dict: @@ -72,21 +72,12 @@ def main(args) -> None: for line in m["diff"]: print(fmt(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None))) if data.get("src") is not None: print(fmt(data["src"])) - # ** Graph rewrites printer - if args.rewrites: - if args.src is None: return print("Select a source with -s"+"\n"+"\n".join([f" {fmt_colored(k)}" for k in rewrites])) - steps = get(rewrites, args.src) - if args.item is None: - for k,v in steps.items(): print(" "*v["depth"]+k+(f" - {v['match_count']}" if v.get('match_count', 0) else '')) - else: print_step(get(steps, args.item)) - return None - events:list = viz.load_pickle(args.profile_path, default=[]) if (profile_bytes:=viz.get_profile(viz_data, events)) is None: raise RuntimeError(f"empty profile in {args.profile_path}") profile = decode_profile(profile_bytes) profile["layout"].update([(f'{c["name"][5:]}{" SQTT" if s["name"].endswith("PKTS") else ""} {s["name"]}', s["data"]) for c in viz_data.ctxs if c["name"].startswith("SQTT") for s in c["steps"] if s["name"].endswith(("PMC", "PKTS"))]) - if args.src is None: return print("Select a source with -s"+"\n ALL\n"+"\n".join([f" {fmt_colored(k)}" for k in profile["layout"]])) + if args.list and args.src == "ALL": return print("ALL\n"+"\n".join(fmt_colored(k) for k in profile["layout"])) # ** SQTT printer data = None if args.src == "ALL" else get(profile["layout"], args.src) @@ -118,7 +109,7 @@ def main(args) -> None: phase, delay = "EXEC", int(e.st) - dispatch_st if inst and phase: info = f"{phase:<8} {inst}" unit = e.device.replace(" ", "-") - row = {"clk":int(e.st)-inst_st, "unit":unit, "op":op_name, "dur":int(unwrap(e.en)-e.st), "delay":delay or "", "info":info} + row = {"clk":int(e.st)-inst_st, "cycle":int(e.st), "unit":unit, "op":op_name, "dur":int(unwrap(e.en)-e.st), "delay":delay or "", "info":info} print(fmt(row, lambda _: f"{row['clk']:<12} {unit:<20} {op_str}{' '*(22-ansilen(op_str))} {row['dur']:<4} {str(row['delay']):<4} {info}")) # ** PMC printer @@ -139,7 +130,7 @@ def main(args) -> None: # ** Profiler printer else: - timelines = [(n,l) for n,l in profile["layout"].items() if l.get("event_type") == 0] + timelines = [(n,l) for n,l in profile["layout"].items() if isinstance(l, dict) and l.get("event_type") == 0] def produce_top_kernels() -> Iterator[dict]: tagged = ((n,e) for n,l in timelines for e in l["events"]) if args.src == "ALL" else ((args.src,e) for e in unwrap(data)["events"]) agg:dict[tuple[str,str], tuple[float, int, int|None]] = {} # map (device, kernel name) to (total time, count and ref) @@ -161,6 +152,9 @@ def main(args) -> None: def produce_all_kernels() -> Iterator[dict]: event_streams = [[(e["st"], n, e) for e in l["events"]] for n,l in timelines] if args.src == "ALL" \ else [[(e["st"], args.src, e) for e in unwrap(data)["events"]]] + if args.src == "ALL": + for n,l in profile["layout"].items(): + if not isinstance(l, dict) or l.get("event_type") != 0: yield {"device":"SOURCE", "name":n, "st_ms":0, "ref":None, "ext":None} marker_stream = sorted([(m["ts"], "MARKER", m) for m in profile.get("markers", [])], key=lambda t:t[0]) for ts,dev,e in heapq.merge(*event_streams, marker_stream, key=lambda t:t[0]): if dev == "MARKER": @@ -179,30 +173,39 @@ def main(args) -> None: def fmt_top(k:dict) -> str: return f"{fmt_colored(k['name'])}{' ' * max(0, 36-ansilen(k['name']))} {time_to_str(k['dur_ms']*1e-3, w=9)} {k['count']:7d} {k['pct']:6.2f}%" def fmt_all(k:dict) -> str: - if k["device"] == "MARKER": return f"--- MARKER {k['name']} /{k['st_ms']:9.2f}ms" + if k["device"] in {"MARKER", "SOURCE"}: return f"--- {k['device']} {k['name']}"+(f"/{k['st_ms']:9.2f}ms" if k['st_ms'] else "") ptm = colored(time_to_str(k["dur_ms"]*1e-3, w=9), "yellow" if k["dur_ms"] > 10 else None) fmt_str = " ".join(p+" "*max(0, 14-ansilen(p)) for p in k["fmt"].split("\n")) name = f"*** {k['device'][:7]:7s} "+k["name"]+" "*(46-ansilen(k["name"])) return f"{name} tm {ptm}/{k['st_ms']:9.2f}ms"+(f" ({fmt_str})" if k["fmt"] else "") fmt_row = fmt_top if args.top else fmt_all seen_refs:set[int] = set() - for k in (produce_top_kernels if args.top else produce_all_kernels)(): + def render_event(k:dict, ls=args.list) -> None: print(fmt(k, to_str=fmt_row)) if k["ref"] is not None and k["ref"] not in seen_refs: seen_refs.add(k["ref"]) - steps = rewrites[viz_data.ctxs[k["ref"]]["name"]] - if DEBUG >= 3 and (ast_step:=steps.get("View Base AST")) is not None: print_step(ast_step) - if DEBUG >= 4 and (src_step:=steps.get("View Source")) is not None: print_step(src_step) + for s in viz_data.ctxs[k["ref"]]["steps"]: + if DEBUG >= 3 and s["name"] == "View Base AST": print_step(s) + if DEBUG >= 4 and s["name"] == "View Source": print_step(s) + if DEBUG >= 5 or ls: print(fmt(" "*s["depth"]+s["name"]+(f" - {s['match_count']}" if s.get('match_count', 0) else ''))) + if DEBUG >= 6: print_step(s) elif DEBUG >= 3 and k.get("ext"): print(fmt(k["ext"])) + produce = produce_top_kernels if args.top else produce_all_kernels + if args.item: + if len(args.item) > 2: raise RuntimeError(f"-i takes at most 2 names (got {args.item})") + k = get({r["name"]:r for r in produce()}, args.item[0]) + if len(args.item) == 1: + with Context(DEBUG=max(DEBUG.value, 3)): render_event(k, ls=True) + else: print_step(get(rewrites[viz_data.ctxs[k["ref"]]["name"]], args.item[1])) + else: + for k in produce(): render_event(k) def get_arg_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(add_help=False, prog="python -m tinygrad.viz.cli") - g_mode = parser.add_argument_group("mode") - g_mode.add_argument("-p", "--profile", action="store_true", help="View profile") - g_mode.add_argument("-r", "--rewrites", action="store_true", help="View graph rewrites") g_opts = parser.add_argument_group("optional args") - g_opts.add_argument("-s", "--src", type=str, default=None, metavar="NAME", help="Select a data source (default: list all sources)") - g_opts.add_argument("-i", "--item", type=str, default=None, metavar="NAME", help="Select an item within the source (default: list all items)") + g_opts.add_argument("-s", "--src", type=str, default="ALL", metavar="NAME", help="Select a data source (default: ALL)") + g_opts.add_argument("-i", "--item", nargs="+", default=None, metavar="NAME", help="Select an item within the source (default: list all items)") + g_opts.add_argument("--list", "--ls", dest="list", action="store_true", help="List sources") g_opts.add_argument("-t", "--top", nargs="?", type=int, const=20, metavar="COUNT", help="Aggregate top kernels (optional count, default 20)") g_opts.add_argument("--profile-path", type=pathlib.Path, metavar="PATH", help="Optional path to profile.pkl (default: latest profile)", default=pathlib.Path(temp("profile.pkl", append_user=True))) @@ -213,10 +216,5 @@ def get_arg_parser() -> argparse.ArgumentParser: return parser if __name__ == "__main__": - args = get_arg_parser().parse_args() - if not args.profile and not args.rewrites: - get_arg_parser().print_help() - sys.exit(0) - - try: main(args) + try: main(get_arg_parser().parse_args()) except KeyboardInterrupt: pass diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index 8bef2c992c..d650e42b73 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -204,7 +204,7 @@ const waveColor = (op) => { if (op.includes("LDS_")) { ret = darkenHex(ret, 25) } return ret }; -const colorScheme = {TINY:new Map([["Schedule","#1b5745"],["get_program","#1d2e62"],["compile","#63b0cd"],["DEFAULT","#354f52"]]), +const colorScheme = {TINY:new Map([["Schedule","#1b5745"],["precompile","#1d2e62"],["compile","#63b0cd"],["DEFAULT","#354f52"]]), DEFAULT:["#2b2e39", "#2c2f3a", "#31343f", "#323544", "#2d303a", "#2e313c", "#343746", "#353847", "#3c4050", "#404459", "#444862", "#4a4e65"], BUFFER:["#342483", "#3E2E94", "#4938A4", "#5442B4", "#5E4CC2", "#674FCA"], SIMD:new Map([["OCC", "#101725"], ["INST", "#0A2042"]]), GPC:new Map([["NONE","#1a7a2e"],["MEMORY_DEPENDENCY","#8b1a00"],["EXEC_DEPENDENCY","#006b6b"],["INST_FETCH","#7a7a00"],["SYNC","#6b006b"], @@ -1180,7 +1180,7 @@ document.addEventListener("keydown", (event) => { // r key toggles indexing if (event.key === "r") showIndexing.toggle.click(); // c key toggles CALL src - if (event.key === "c") showCallSrc.toggle.click(); + if (event.key === "c" && !event.ctrlKey && !event.metaKey && !event.altKey) showCallSrc.toggle.click(); // s key toggles SINK if (event.key === "s") showSink.toggle.click(); // g key toggles graph diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 964d7afd89..45628ee876 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -358,7 +358,7 @@ wave_colors = {"WMMA": "#1F7857", **{x:"#ffffc0" for x in ["VALU", "VINTERP"]}, def sqtt_timeline(data:bytes, lib:bytes, target:str) -> Generator[ProfileEvent, None, None]: from tinygrad.renderer.amd.sqtt import (map_insts, InstructionInfo, PacketType, INST, InstOp, VALUINST, IMMEDIATE, IMMEDIATE_MASK, VMEMEXEC, ALUEXEC, INST_RDNA4, InstOpRDNA4, TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4, CDNA_INST, InstOpCDNA, - WAVEEND, CDNA_WAVEEND, WAVERDY) + WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND, WAVERDY) pc_map = {addr:str(inst) for addr,inst in amd_decode(lib, target).items()} row_ends:dict[str, Decimal] = {} row_counts:dict[str, itertools.count] = {} @@ -374,20 +374,24 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> Generator[ProfileEvent, link = f"PC:{info.pc}" if info else None if isinstance(p, (ALUEXEC, VMEMEXEC)): dispatch_id, op_type = exec_pending[name].pop(0) - # get the number of cycles from the op type + # wmma exec gets its own color and its own row on rdna4 + if op_type.startswith("WMMA"): + name = name+"_WMMA" + if not op_type.startswith("WMMA_VALU"): row = "ALUEXEC:0 WMMA" + # transcendental valu gets its own row + if op_type.startswith("VALUT"): row = "ALUEXEC:0 TFU" + # extend execs by the op type's known duration, p._time marks the first or last cycle based on the op type duration = int(dur_match.group(1)) if (dur_match:=re.match(r".*_(\d+)$", op_type)) else 1 - # for execs, extend end time by the duration - start_time, end_time = p._time, p._time+duration + if any(ss in row for ss in ("SALU", "TFU", "VMEM", "LDS")): start_time, end_time = p._time, p._time+duration + else: start_time, end_time = p._time-duration, p._time link = f"LINK:{dispatch_id}" - # wmma exec gets its own row and color - if op_type.startswith("WMMA"): name, row = name+"_WMMA", "ALUEXEC:0 WMMA" # queue inst dispatches idx = next(row_counts.setdefault(row, itertools.count(0))) if isinstance(p, (VALUINST, INST, INST_RDNA4)) and (exec_type:=dispatch_to_exec.get(name.replace("OTHER_", "").split("_")[0])) is not None: if name.startswith("OTHER_"): exec_type = f"{exec_type}_ALT" # detect rdna3 wmma from the asm, only rdna4 has an op type for it if isinstance(p, VALUINST) and (asm:=getattr(unwrap(info).inst, "op_name", "")).startswith("V_WMMA"): - name = f"WMMA_{16 if 'IU4' in asm else 32}" + name = f"WMMA_VALU_{16 if 'IU4' in asm else 32}" exec_pending.setdefault(exec_type, []).append((f"{row}-{idx}", name)) # construct and yield the event for this packet if row not in row_ends: yield ProfilePointEvent(row, "JSON", "pcMap", pc_map, ts=Decimal(0)) @@ -412,7 +416,7 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> Generator[ProfileEvent, if isinstance(p, (INST, INST_RDNA4, CDNA_INST)): name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4, InstOpCDNA)) else f"0x{p.op:02x}" yield from add(name, p, info=info) - if isinstance(p, (VALUINST, IMMEDIATE, WAVEEND, CDNA_WAVEEND)): yield from add(p.__class__.__name__, p, info=info) + if isinstance(p, (VALUINST, IMMEDIATE, WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND)): yield from add(p.__class__.__name__, p, info=info) if isinstance(p, IMMEDIATE_MASK): yield from add("IMMEDIATE", p, wave=unwrap(info).wave, info=info) if isinstance(p, WAVERDY): for wave in range(16):