clean

.github/workflows/benchmark.yml

@@ -429,15 +429,13 @@ jobs:
     #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
     - name: Test speed vs theoretical
       run: AMD=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
-    - name: Test tensor cores AMD_LLVM=0
-      run: AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
-    # TODO: this is flaky
-    # - name: Test tensor cores AMD_LLVM=1
-    #   run: AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
-    - name: Run Tensor Core GEMM (AMD)
+    - name: Test tensor cores
       run: |
+        AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
+        AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
         AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-        AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
+    - name: Run Tensor Core GEMM (AMD)
+      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
     - name: Test AMD=1
       run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
     #- name: Test HIP=1

.github/workflows/test.yml

@@ -241,8 +241,9 @@ jobs:
       run: |
         python -m mypy --strict-equality --lineprecision-report .
         cat lineprecision.txt
-    - name: Run TYPED=1
-      run: TYPED=1 python -c "import tinygrad"
+    # broken because of UPatAny
+    #- name: Run TYPED=1
+    #  run: TYPED=1 python -c "import tinygrad"
 
   unittest:
     name: Unit Tests
@@ -641,7 +642,7 @@ jobs:
         if: matrix.backend=='amdllvm'
         run: python test/device/test_amd_llvm.py
       - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py --durations=20
+        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py --durations=20
       - name: Run pytest (amd)
         run: python -m pytest test/external/external_test_am.py --durations=20
       - name: Run TRANSCENDENTAL math

docs/abstractions3.py

@@ -45,9 +45,10 @@ print(f"The schedule contains {len(schedule)} items.")
 for si in schedule: print(str(si)[:80])
 
 # *****
-# 4. Lower and run the schedule.
+# 4. Run the schedule.
 
-for si in tqdm(schedule): si.run()
+from tinygrad.engine.realize import run_schedule
+run_schedule(schedule)
 
 # *****
 # 5. Print the weight change

docs/developer/developer.md

@@ -13,7 +13,7 @@ There's also a [doc describing speed](../developer/speed.md)
 
 Everything in [Tensor](../tensor/index.md) is syntactic sugar around constructing a graph of [UOps](../developer/uop.md).
 
-The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not all UOps will actually become realized. There's two types of UOps, base and view. base contains compute into a contiguous buffer, and view is a view. Inputs to a base can be either base or view, inputs to a view can only be a single base.
+The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not all UOps will actually become realized. There's two types of UOps, base and view. base contains compute into a contiguous buffer, and view is a view (specified by a ShapeTracker). Inputs to a base can be either base or view, inputs to a view can only be a single base.
 
 ## Scheduling
 

docs/developer/layout.md

@@ -26,9 +26,9 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi
 
 ## tinygrad/codegen
 
-Transform the optimized ast into a linearized and rendered program.
+Transform the optimized ast into a linearized list of UOps.
 
-::: tinygrad.codegen.get_program
+::: tinygrad.codegen.full_rewrite
     options:
         members: false
         show_labels: false

eslint.config.mjs

@@ -0,0 +1,9 @@
+import globals from "globals";
+import pluginJs from "@eslint/js";
+import pluginHtml from "eslint-plugin-html";
+
+export default [
+  {files: ["**/*.html"], plugins: {html: pluginHtml}, rules:{"max-len": ["error", {"code": 150}]}},
+  {languageOptions: {globals: globals.browser}},
+  pluginJs.configs.recommended,
+];

examples/mlperf/dataloader.py

@@ -763,26 +763,48 @@ class BlendedGPTDataset:
 
     return dataset_idx, dataset_sample_idx
 
-def get_llama3_dataset(samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False) -> BlendedGPTDataset:
-  if small:
-    if val:
-      return BlendedGPTDataset(
-        [base_dir / "c4-validation-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
-    return BlendedGPTDataset(
-      [base_dir / "c4-train.en_6_text_document"], [1.0], samples, seqlen, seed, shuffle=True)
+def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
   if val:
-    return BlendedGPTDataset(
-      [base_dir / "validation" / "c4-validationn-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
-  return BlendedGPTDataset(
-    [base_dir / "c4-train.en_6_text_document", base_dir / "c4-train.en_7_text_document"], [1.0, 1.0], samples, seqlen, seed, shuffle=True)
+    dataset = BlendedGPTDataset([
+      base_dir / "validation" / "c4-validationn-91205-samples.en_text_document",
+    ], [
+      1.0
+    ], samples, seqlen, seed, False)
+  else:
+    dataset = BlendedGPTDataset([
+      base_dir / "c4-train.en_6_text_document",
+      base_dir / "c4-train.en_7_text_document",
+    ], [
+      1.0, 1.0
+    ], samples, seqlen, seed, True)
 
-def iterate_llama3_dataset(dataset:BlendedGPTDataset, bs:int):
-  for b in range(math.ceil(dataset.samples / bs)):
-    batch = [dataset.get(b * bs + i) for i in range(bs)]
+  for b in range(math.ceil(samples / bs)):
+    batch = []
+    for i in range(bs):
+      tokens = dataset.get(b * bs + i)
+      batch.append(tokens)
     yield Tensor.stack(batch, dim=0)
 
-def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False):
-  return iterate_llama3_dataset(get_llama3_dataset(samples, seqlen, base_dir, seed, val, small), bs)
+def batch_load_llama3_small(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
+  if val:
+    dataset = BlendedGPTDataset([
+      base_dir / "c4-validation-91205-samples.en_text_document",
+    ], [
+      1.0
+    ], samples, seqlen, seed, False)
+  else:
+    dataset = BlendedGPTDataset([
+      base_dir / "c4-train.en_6_text_document",
+    ], [
+      1.0
+    ], samples, seqlen, seed, True)
+
+  for b in range(math.ceil(samples / bs)):
+    batch = []
+    for i in range(bs):
+      tokens = dataset.get(b * bs + i)
+      batch.append(tokens)
+    yield Tensor.stack(batch, dim=0)
 
 if __name__ == "__main__":
   def load_unet3d(val):

examples/mlperf/model_eval.py

@@ -234,9 +234,12 @@ def eval_llama3():
     loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
     return loss.flatten().float()
 
-  from examples.mlperf.dataloader import get_llama3_dataset, iterate_llama3_dataset
-  eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
-  iter = iterate_llama3_dataset(eval_dataset, BS)
+  if SMALL:
+    from examples.mlperf.dataloader import batch_load_llama3_small
+    iter = batch_load_llama3_small(BS, 5760, SEQLEN, BASEDIR, val=True)
+  else:
+    from examples.mlperf.dataloader import batch_load_llama3
+    iter = batch_load_llama3(BS, 5760, SEQLEN, BASEDIR, val=True)
 
   losses = []
   for tokens in tqdm(iter, total=5760//BS):

examples/mlperf/model_train.py

@@ -1314,13 +1314,6 @@ def train_llama3():
   opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
   opt_end_learning_rate = getenv("END_LR", 8e-7)
 
-  # ** init wandb **
-  WANDB = getenv("WANDB")
-  if WANDB:
-    import wandb
-    wandb_args = {"id": wandb_id, "resume": "must"} if (wandb_id := getenv("WANDB_RESUME", "")) else {}
-    wandb.init(config=config, **wandb_args, project="MLPerf-LLaMA3")
-
   model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
   # vocab_size from the mixtral tokenizer
   if not SMALL: model_params |= {"vocab_size": 32000}
@@ -1424,56 +1417,49 @@ def train_llama3():
     if getenv("FAKEDATA", 0):
       return fake_data(BS, SAMPLES)
     else:
-      from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
-
-  if getenv("FAKEDATA", 0):
-    eval_dataset = None
-  else:
-    from examples.mlperf.dataloader import get_llama3_dataset
-    eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
+      if SMALL:
+        from examples.mlperf.dataloader import batch_load_llama3_small
+        return batch_load_llama3_small(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
+      else:
+        from examples.mlperf.dataloader import batch_load_llama3
+        return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
 
   def get_eval_iter():
-    if eval_dataset is None:
+    if getenv("FAKEDATA", 0):
       return fake_data(EVAL_BS, 5760)
-    from examples.mlperf.dataloader import iterate_llama3_dataset
-    return iterate_llama3_dataset(eval_dataset, EVAL_BS)
+    else:
+      if SMALL:
+        from examples.mlperf.dataloader import batch_load_llama3_small
+        return batch_load_llama3_small(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
+      else:
+        from examples.mlperf.dataloader import batch_load_llama3
+        return batch_load_llama3(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
 
   iter = get_train_iter()
   i, sequences_seen = resume_ckpt, 0
   for tokens in tqdm(iter, total=SAMPLES//GBS):
+    t = time.perf_counter()
     GlobalCounters.reset()
-    if getenv("TRAIN", 1):
-      t = time.perf_counter()
-      loss, lr = train_step(model, tokens)
-      loss = loss.float().item()
-      lr = lr.item()
+    loss, lr = train_step(model, tokens)
+    loss = loss.float().item()
 
-      i += 1
-      sequences_seen += tokens.shape[0]
+    i += 1
+    sequences_seen += tokens.shape[0]
 
-      sec = time.perf_counter()-t
-      mem_gb = GlobalCounters.mem_used / 1e9
-      gflops = GlobalCounters.global_ops / 1e9 / sec
-      tqdm.write(
-        f"{i:5} {sec:.2f} s run, {loss:.4f} loss, {lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS")
+    tqdm.write(f"{loss:.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
+    if (fname:=getenv("LOSS_FILE", "")):
+      with open(fname, "a") as f:
+        f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
 
-      if (fname:=getenv("LOSS_FILE", "")):
-        with open(fname, "a") as f:
-          f.write(f"{i} {loss:.4f} {lr:.12f} {mem_gb:.2f}\n")
+    if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
+      tqdm.write("saving checkpoint")
+      if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
+      fn = f"{ckpt_dir}/llama3_{i}.safe"
+      safe_save(get_state_dict(model), fn)
 
-      if WANDB:
-        wandb.log({"lr": lr, "train/loss": loss, "train/step_time": sec, "train/GFLOPS": gflops, "train/sequences_seen": sequences_seen})
-
-      if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
-        tqdm.write("saving checkpoint")
-        if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-        fn = f"{ckpt_dir}/llama3_{i}.safe"
-        safe_save(get_state_dict(model), fn)
-
-        tqdm.write("saving optim checkpoint")
-        fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
-        safe_save(get_state_dict(scheduler), fn)
+      tqdm.write("saving optim checkpoint")
+      fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
+      safe_save(get_state_dict(scheduler), fn)
 
     if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
       tqdm.write(f"evaluating after {sequences_seen} sequences")
@@ -1489,9 +1475,6 @@ def train_llama3():
 
       tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
 
-      if WANDB:
-        wandb.log({"eval/log_perplexity": log_perplexity, "eval/sequences_seen": sequences_seen})
-
       if log_perplexity < EVAL_TARGET:
         tqdm.write(f"target achieved after {sequences_seen} sequences")
         if getenv("CKPT"):

extra/amdpci/am_smi.py

@@ -184,7 +184,7 @@ class SMICtx:
         if compact: return {k: temps[k] for k in ("Hotspot", "HBM") if temps.get(k, 0) != 0}
         return {k: v for k, v in temps.items() if v != 0}
       case _:
-        temps_keys = [(k, name) for k, name in dev.smu.smu_mod.TEMP_e.items()
+        temps_keys = [(k, name) for k, name in dev.smu.smu_mod.c__EA_TEMP_e__enumvalues.items()
                       if k < dev.smu.smu_mod.TEMP_COUNT and metrics.SmuMetrics.AvgTemperature[k] != 0]
         if compact: temps_keys = [(k, name) for k, name in temps_keys if k in (dev.smu.smu_mod.TEMP_HOTSPOT, dev.smu.smu_mod.TEMP_MEM)]
         return {name: metrics.SmuMetrics.AvgTemperature[k] for k, name in temps_keys}
@@ -193,7 +193,7 @@ class SMICtx:
     match dev.ip_ver[am.MP1_HWIP]:
       case (13,0,6): return {}
       case _:
-        voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.SVI_PLANE_e.items()
+        voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.c__EA_SVI_PLANE_e__enumvalues.items()
                         if k < dev.smu.smu_mod.SVI_PLANE_COUNT and metrics.SmuMetrics.AvgVoltage[k] != 0]
         return {name: metrics.SmuMetrics.AvgVoltage[k] for k, name in voltage_keys}
 

extra/assembly/rdna3/alu.py

@@ -1,254 +0,0 @@
-# Pure combinational ALU functions for RDNA3 emulation
-from __future__ import annotations
-import struct, math
-from typing import Callable
-from extra.assembly.rdna3.autogen import SOP1Op, SOP2Op, SOPCOp, SOPKOp, VOP1Op, VOP2Op, VOP3Op
-
-# Format base offsets for unified opcode space
-SOP2_BASE, SOP1_BASE, SOPC_BASE, SOPK_BASE = 0x000, 0x100, 0x200, 0x300
-VOP2_BASE, VOP1_BASE = 0x100, 0x180
-
-# Float conversion helpers
-_I, _f, _H, _e = struct.Struct('<I'), struct.Struct('<f'), struct.Struct('<H'), struct.Struct('<e')
-def f32(i: int) -> float: return _f.unpack(_I.pack(i & 0xffffffff))[0]
-def i32(f: float) -> int:
-  if math.isinf(f): return 0x7f800000 if f > 0 else 0xff800000
-  try: return _I.unpack(_f.pack(f))[0]
-  except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
-def f16(i: int) -> float: return _e.unpack(_H.pack(i & 0xffff))[0]
-def i16(f: float) -> int:
-  if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
-  try: return _H.unpack(_e.pack(f))[0]
-  except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
-def sext(v: int, b: int) -> int: return v - (1 << b) if v & (1 << (b-1)) else v
-def clz(x: int) -> int: return 32 - x.bit_length() if x else 32
-def cls(x: int) -> int: x &= 0xffffffff; return 31 if x in (0, 0xffffffff) else clz(~x & 0xffffffff if x >> 31 else x) - 1
-def _cvt_i32_f32(v): return (0x7fffffff if v > 0 else 0x80000000) if math.isinf(v) else (0 if math.isnan(v) else max(-0x80000000, min(0x7fffffff, int(v))) & 0xffffffff)
-def _cvt_u32_f32(v): return (0xffffffff if v > 0 else 0) if math.isinf(v) else (0 if math.isnan(v) or v < 0 else min(0xffffffff, int(v)))
-
-# SALU: op -> fn(s0, s1, scc_in) -> (result, scc_out)
-SALU: dict[int, Callable] = {
-  # SOP2
-  SOP2_BASE + SOP2Op.S_ADD_U32: lambda a, b, scc: ((a + b) & 0xffffffff, int((a + b) >= 0x100000000)),
-  SOP2_BASE + SOP2Op.S_SUB_U32: lambda a, b, scc: ((a - b) & 0xffffffff, int(b > a)),
-  SOP2_BASE + SOP2Op.S_ADDC_U32: lambda a, b, scc: ((r := a + b + scc) & 0xffffffff, int(r >= 0x100000000)),
-  SOP2_BASE + SOP2Op.S_SUBB_U32: lambda a, b, scc: ((a - b - scc) & 0xffffffff, int((b + scc) > a)),
-  SOP2_BASE + SOP2Op.S_ADD_I32: lambda a, b, scc: ((r := sext(a, 32) + sext(b, 32)) & 0xffffffff, int(((a >> 31) == (b >> 31)) and ((a >> 31) != ((r >> 31) & 1)))),
-  SOP2_BASE + SOP2Op.S_SUB_I32: lambda a, b, scc: ((r := sext(a, 32) - sext(b, 32)) & 0xffffffff, int(((a >> 31) != (b >> 31)) and ((a >> 31) != ((r >> 31) & 1)))),
-  SOP2_BASE + SOP2Op.S_AND_B32: lambda a, b, scc: ((r := a & b), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_OR_B32: lambda a, b, scc: ((r := a | b), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_XOR_B32: lambda a, b, scc: ((r := a ^ b), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_AND_NOT1_B32: lambda a, b, scc: ((r := a & (~b & 0xffffffff)), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_OR_NOT1_B32: lambda a, b, scc: ((r := a | (~b & 0xffffffff)), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_LSHL_B32: lambda a, b, scc: ((r := (a << (b & 0x1f)) & 0xffffffff), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_LSHR_B32: lambda a, b, scc: ((r := a >> (b & 0x1f)), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_ASHR_I32: lambda a, b, scc: ((r := sext(a, 32) >> (b & 0x1f)) & 0xffffffff, int(r != 0)),
-  SOP2_BASE + SOP2Op.S_MUL_I32: lambda a, b, scc: ((sext(a, 32) * sext(b, 32)) & 0xffffffff, scc),
-  SOP2_BASE + SOP2Op.S_MUL_HI_U32: lambda a, b, scc: (((a * b) >> 32) & 0xffffffff, scc),
-  SOP2_BASE + SOP2Op.S_MUL_HI_I32: lambda a, b, scc: (((sext(a, 32) * sext(b, 32)) >> 32) & 0xffffffff, scc),
-  SOP2_BASE + SOP2Op.S_MIN_I32: lambda a, b, scc: (a, 1) if sext(a, 32) < sext(b, 32) else (b, 0),
-  SOP2_BASE + SOP2Op.S_MIN_U32: lambda a, b, scc: (a, 1) if a < b else (b, 0),
-  SOP2_BASE + SOP2Op.S_MAX_I32: lambda a, b, scc: (a, 1) if sext(a, 32) > sext(b, 32) else (b, 0),
-  SOP2_BASE + SOP2Op.S_MAX_U32: lambda a, b, scc: (a, 1) if a > b else (b, 0),
-  SOP2_BASE + SOP2Op.S_CSELECT_B32: lambda a, b, scc: (a if scc else b, scc),
-  SOP2_BASE + SOP2Op.S_BFE_U32: lambda a, b, scc: ((r := ((a >> (b & 0x1f)) & ((1 << ((b >> 16) & 0x7f)) - 1)) if (b >> 16) & 0x7f else 0), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_BFE_I32: lambda a, b, scc: ((r := sext((a >> (b & 0x1f)) & ((1 << w) - 1), w) & 0xffffffff if (w := (b >> 16) & 0x7f) else 0), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_PACK_LL_B32_B16: lambda a, b, scc: ((a & 0xffff) | ((b & 0xffff) << 16), scc),
-  SOP2_BASE + SOP2Op.S_PACK_LH_B32_B16: lambda a, b, scc: ((a & 0xffff) | (b & 0xffff0000), scc),
-  SOP2_BASE + SOP2Op.S_PACK_HH_B32_B16: lambda a, b, scc: (((a >> 16) & 0xffff) | (b & 0xffff0000), scc),
-  SOP2_BASE + SOP2Op.S_PACK_HL_B32_B16: lambda a, b, scc: (((a >> 16) & 0xffff) | ((b & 0xffff) << 16), scc),
-  SOP2_BASE + SOP2Op.S_ADD_F32: lambda a, b, scc: (i32(f32(a) + f32(b)), scc),
-  SOP2_BASE + SOP2Op.S_SUB_F32: lambda a, b, scc: (i32(f32(a) - f32(b)), scc),
-  SOP2_BASE + SOP2Op.S_MUL_F32: lambda a, b, scc: (i32(f32(a) * f32(b)), scc),
-  # SOP1
-  SOP1_BASE + SOP1Op.S_MOV_B32: lambda a, b, scc: (a, scc),
-  SOP1_BASE + SOP1Op.S_NOT_B32: lambda a, b, scc: ((r := (~a) & 0xffffffff), int(r != 0)),
-  SOP1_BASE + SOP1Op.S_BREV_B32: lambda a, b, scc: (int(f'{a & 0xffffffff:032b}'[::-1], 2), scc),
-  SOP1_BASE + SOP1Op.S_CLZ_I32_U32: lambda a, b, scc: (clz(a), scc),
-  SOP1_BASE + SOP1Op.S_CLS_I32: lambda a, b, scc: (cls(a), scc),
-  SOP1_BASE + SOP1Op.S_SEXT_I32_I8: lambda a, b, scc: (sext(a & 0xff, 8) & 0xffffffff, scc),
-  SOP1_BASE + SOP1Op.S_SEXT_I32_I16: lambda a, b, scc: (sext(a & 0xffff, 16) & 0xffffffff, scc),
-  SOP1_BASE + SOP1Op.S_ABS_I32: lambda a, b, scc: ((r := abs(sext(a, 32)) & 0xffffffff), int(r != 0)),
-  SOP1_BASE + SOP1Op.S_CVT_F32_I32: lambda a, b, scc: (i32(float(sext(a, 32))), scc),
-  SOP1_BASE + SOP1Op.S_CVT_F32_U32: lambda a, b, scc: (i32(float(a)), scc),
-  SOP1_BASE + SOP1Op.S_CVT_I32_F32: lambda a, b, scc: (_cvt_i32_f32(f32(a)), scc),
-  SOP1_BASE + SOP1Op.S_CVT_U32_F32: lambda a, b, scc: (_cvt_u32_f32(f32(a)), scc),
-  SOP1_BASE + SOP1Op.S_CEIL_F32: lambda a, b, scc: (i32(math.ceil(f32(a))), scc),
-  SOP1_BASE + SOP1Op.S_FLOOR_F32: lambda a, b, scc: (i32(math.floor(f32(a))), scc),
-  SOP1_BASE + SOP1Op.S_TRUNC_F32: lambda a, b, scc: (i32(math.trunc(f32(a))), scc),
-  SOP1_BASE + SOP1Op.S_RNDNE_F32: lambda a, b, scc: (i32(round(f32(a))), scc),
-  SOP1_BASE + SOP1Op.S_CVT_F16_F32: lambda a, b, scc: (i16(f32(a)), scc),
-  SOP1_BASE + SOP1Op.S_CVT_F32_F16: lambda a, b, scc: (i32(f16(a)), scc),
-  # SOPC
-  SOPC_BASE + SOPCOp.S_CMP_EQ_I32: lambda a, b, scc: (0, int(sext(a, 32) == sext(b, 32))),
-  SOPC_BASE + SOPCOp.S_CMP_LG_I32: lambda a, b, scc: (0, int(sext(a, 32) != sext(b, 32))),
-  SOPC_BASE + SOPCOp.S_CMP_GT_I32: lambda a, b, scc: (0, int(sext(a, 32) > sext(b, 32))),
-  SOPC_BASE + SOPCOp.S_CMP_GE_I32: lambda a, b, scc: (0, int(sext(a, 32) >= sext(b, 32))),
-  SOPC_BASE + SOPCOp.S_CMP_LT_I32: lambda a, b, scc: (0, int(sext(a, 32) < sext(b, 32))),
-  SOPC_BASE + SOPCOp.S_CMP_LE_I32: lambda a, b, scc: (0, int(sext(a, 32) <= sext(b, 32))),
-  SOPC_BASE + SOPCOp.S_CMP_EQ_U32: lambda a, b, scc: (0, int(a == b)),
-  SOPC_BASE + SOPCOp.S_CMP_LG_U32: lambda a, b, scc: (0, int(a != b)),
-  SOPC_BASE + SOPCOp.S_CMP_GT_U32: lambda a, b, scc: (0, int(a > b)),
-  SOPC_BASE + SOPCOp.S_CMP_GE_U32: lambda a, b, scc: (0, int(a >= b)),
-  SOPC_BASE + SOPCOp.S_CMP_LT_U32: lambda a, b, scc: (0, int(a < b)),
-  SOPC_BASE + SOPCOp.S_CMP_LE_U32: lambda a, b, scc: (0, int(a <= b)),
-  SOPC_BASE + SOPCOp.S_BITCMP0_B32: lambda a, b, scc: (0, int((a & (1 << (b & 0x1f))) == 0)),
-  SOPC_BASE + SOPCOp.S_BITCMP1_B32: lambda a, b, scc: (0, int((a & (1 << (b & 0x1f))) != 0)),
-  # SOPK
-  SOPK_BASE + SOPKOp.S_MOVK_I32: lambda a, b, scc: (sext(b, 16) & 0xffffffff, scc),
-  SOPK_BASE + SOPKOp.S_CMOVK_I32: lambda a, b, scc: ((sext(b, 16) & 0xffffffff) if scc else a, scc),
-  SOPK_BASE + SOPKOp.S_ADDK_I32: lambda a, b, scc: ((r := sext(a, 32) + sext(b, 16)) & 0xffffffff, int(((a >> 31) == ((b >> 15) & 1)) and ((a >> 31) != ((r >> 31) & 1)))),
-  SOPK_BASE + SOPKOp.S_MULK_I32: lambda a, b, scc: ((sext(a, 32) * sext(b, 16)) & 0xffffffff, scc),
-  SOPK_BASE + SOPKOp.S_CMPK_EQ_I32: lambda a, b, scc: (0, int(sext(a, 32) == sext(b, 16))),
-  SOPK_BASE + SOPKOp.S_CMPK_LG_I32: lambda a, b, scc: (0, int(sext(a, 32) != sext(b, 16))),
-  SOPK_BASE + SOPKOp.S_CMPK_GT_I32: lambda a, b, scc: (0, int(sext(a, 32) > sext(b, 16))),
-  SOPK_BASE + SOPKOp.S_CMPK_GE_I32: lambda a, b, scc: (0, int(sext(a, 32) >= sext(b, 16))),
-  SOPK_BASE + SOPKOp.S_CMPK_LT_I32: lambda a, b, scc: (0, int(sext(a, 32) < sext(b, 16))),
-  SOPK_BASE + SOPKOp.S_CMPK_LE_I32: lambda a, b, scc: (0, int(sext(a, 32) <= sext(b, 16))),
-  SOPK_BASE + SOPKOp.S_CMPK_EQ_U32: lambda a, b, scc: (0, int(a == (b & 0xffff))),
-  SOPK_BASE + SOPKOp.S_CMPK_LG_U32: lambda a, b, scc: (0, int(a != (b & 0xffff))),
-  SOPK_BASE + SOPKOp.S_CMPK_GT_U32: lambda a, b, scc: (0, int(a > (b & 0xffff))),
-  SOPK_BASE + SOPKOp.S_CMPK_GE_U32: lambda a, b, scc: (0, int(a >= (b & 0xffff))),
-  SOPK_BASE + SOPKOp.S_CMPK_LT_U32: lambda a, b, scc: (0, int(a < (b & 0xffff))),
-  SOPK_BASE + SOPKOp.S_CMPK_LE_U32: lambda a, b, scc: (0, int(a <= (b & 0xffff))),
-}
-
-# VALU: op -> fn(s0, s1, s2) -> result
-VALU: dict[int, Callable] = {
-  # VOP2
-  VOP2_BASE + VOP2Op.V_ADD_F32: lambda a, b, c: i32(f32(a) + f32(b)),
-  VOP2_BASE + VOP2Op.V_SUB_F32: lambda a, b, c: i32(f32(a) - f32(b)),
-  VOP2_BASE + VOP2Op.V_SUBREV_F32: lambda a, b, c: i32(f32(b) - f32(a)),
-  VOP2_BASE + VOP2Op.V_MUL_F32: lambda a, b, c: i32(f32(a) * f32(b)),
-  VOP2_BASE + VOP2Op.V_MIN_F32: lambda a, b, c: i32(min(f32(a), f32(b))),
-  VOP2_BASE + VOP2Op.V_MAX_F32: lambda a, b, c: i32(max(f32(a), f32(b))),
-  VOP2_BASE + VOP2Op.V_ADD_NC_U32: lambda a, b, c: (a + b) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_SUB_NC_U32: lambda a, b, c: (a - b) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_SUBREV_NC_U32: lambda a, b, c: (b - a) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_AND_B32: lambda a, b, c: a & b,
-  VOP2_BASE + VOP2Op.V_OR_B32: lambda a, b, c: a | b,
-  VOP2_BASE + VOP2Op.V_XOR_B32: lambda a, b, c: a ^ b,
-  VOP2_BASE + VOP2Op.V_XNOR_B32: lambda a, b, c: (~(a ^ b)) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_LSHLREV_B32: lambda a, b, c: (b << (a & 0x1f)) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_LSHRREV_B32: lambda a, b, c: b >> (a & 0x1f),
-  VOP2_BASE + VOP2Op.V_ASHRREV_I32: lambda a, b, c: (sext(b, 32) >> (a & 0x1f)) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_MIN_I32: lambda a, b, c: a if sext(a, 32) < sext(b, 32) else b,
-  VOP2_BASE + VOP2Op.V_MAX_I32: lambda a, b, c: a if sext(a, 32) > sext(b, 32) else b,
-  VOP2_BASE + VOP2Op.V_MIN_U32: lambda a, b, c: min(a, b),
-  VOP2_BASE + VOP2Op.V_MAX_U32: lambda a, b, c: max(a, b),
-  VOP2_BASE + VOP2Op.V_MUL_I32_I24: lambda a, b, c: (sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24)) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_MUL_HI_I32_I24: lambda a, b, c: ((sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24)) >> 32) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_MUL_U32_U24: lambda a, b, c: ((a & 0xffffff) * (b & 0xffffff)) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_MUL_HI_U32_U24: lambda a, b, c: (((a & 0xffffff) * (b & 0xffffff)) >> 32) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_CVT_PK_RTZ_F16_F32: lambda a, b, c: i16(f32(a)) | (i16(f32(b)) << 16),
-  VOP2_BASE + VOP2Op.V_LDEXP_F16: lambda a, b, c: i16(math.ldexp(f16(a), sext(b, 32))),
-  VOP2_BASE + VOP2Op.V_ADD_F16: lambda a, b, c: i16(f16(a) + f16(b)),
-  VOP2_BASE + VOP2Op.V_SUB_F16: lambda a, b, c: i16(f16(a) - f16(b)),
-  VOP2_BASE + VOP2Op.V_MUL_F16: lambda a, b, c: i16(f16(a) * f16(b)),
-  VOP2_BASE + VOP2Op.V_MIN_F16: lambda a, b, c: i16(min(f16(a), f16(b))),
-  VOP2_BASE + VOP2Op.V_MAX_F16: lambda a, b, c: i16(max(f16(a), f16(b))),
-  # VOP1
-  VOP1_BASE + VOP1Op.V_MOV_B32: lambda a, b, c: a,
-  VOP1_BASE + VOP1Op.V_NOT_B32: lambda a, b, c: (~a) & 0xffffffff,
-  VOP1_BASE + VOP1Op.V_BFREV_B32: lambda a, b, c: int(f'{a & 0xffffffff:032b}'[::-1], 2),
-  VOP1_BASE + VOP1Op.V_CLZ_I32_U32: lambda a, b, c: clz(a),
-  VOP1_BASE + VOP1Op.V_CLS_I32: lambda a, b, c: cls(a),
-  VOP1_BASE + VOP1Op.V_CVT_F32_I32: lambda a, b, c: i32(float(sext(a, 32))),
-  VOP1_BASE + VOP1Op.V_CVT_F32_U32: lambda a, b, c: i32(float(a)),
-  VOP1_BASE + VOP1Op.V_CVT_I32_F32: lambda a, b, c: _cvt_i32_f32(f32(a)),
-  VOP1_BASE + VOP1Op.V_CVT_U32_F32: lambda a, b, c: _cvt_u32_f32(f32(a)),
-  VOP1_BASE + VOP1Op.V_CVT_F16_F32: lambda a, b, c: i16(f32(a)),
-  VOP1_BASE + VOP1Op.V_CVT_F32_F16: lambda a, b, c: i32(f16(a)),
-  VOP1_BASE + VOP1Op.V_RCP_F32: lambda a, b, c: i32(1.0 / f32(a) if f32(a) != 0 else math.copysign(float('inf'), f32(a))),
-  VOP1_BASE + VOP1Op.V_RCP_IFLAG_F32: lambda a, b, c: i32(1.0 / f32(a) if f32(a) != 0 else math.copysign(float('inf'), f32(a))),
-  VOP1_BASE + VOP1Op.V_RSQ_F32: lambda a, b, c: i32(1.0 / math.sqrt(f32(a)) if f32(a) > 0 else (float('nan') if f32(a) < 0 else float('inf'))),
-  VOP1_BASE + VOP1Op.V_SQRT_F32: lambda a, b, c: i32(math.sqrt(f32(a)) if f32(a) >= 0 else float('nan')),
-  VOP1_BASE + VOP1Op.V_LOG_F32: lambda a, b, c: i32(math.log2(f32(a)) if f32(a) > 0 else (float('-inf') if f32(a) == 0 else float('nan'))),
-  VOP1_BASE + VOP1Op.V_EXP_F32: lambda a, b, c: i32(float('inf') if f32(a) > 128 else (0.0 if f32(a) < -150 else math.pow(2.0, f32(a)))),
-  VOP1_BASE + VOP1Op.V_SIN_F32: lambda a, b, c: i32(math.sin(f32(a) * 2 * math.pi)),
-  VOP1_BASE + VOP1Op.V_COS_F32: lambda a, b, c: i32(math.cos(f32(a) * 2 * math.pi)),
-  VOP1_BASE + VOP1Op.V_FLOOR_F32: lambda a, b, c: i32(math.floor(f32(a))),
-  VOP1_BASE + VOP1Op.V_CEIL_F32: lambda a, b, c: i32(math.ceil(f32(a))),
-  VOP1_BASE + VOP1Op.V_TRUNC_F32: lambda a, b, c: i32(math.trunc(f32(a))),
-  VOP1_BASE + VOP1Op.V_RNDNE_F32: lambda a, b, c: i32(round(f32(a))),
-  VOP1_BASE + VOP1Op.V_FRACT_F32: lambda a, b, c: i32((v := f32(a)) - math.floor(v)),
-  VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE0: lambda a, b, c: i32(float(a & 0xff)),
-  VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE1: lambda a, b, c: i32(float((a >> 8) & 0xff)),
-  VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE2: lambda a, b, c: i32(float((a >> 16) & 0xff)),
-  VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE3: lambda a, b, c: i32(float((a >> 24) & 0xff)),
-  VOP1_BASE + VOP1Op.V_FREXP_MANT_F32: lambda a, b, c: i32(math.frexp(v)[0] if (v := f32(a)) != 0 else 0.0),
-  VOP1_BASE + VOP1Op.V_FREXP_EXP_I32_F32: lambda a, b, c: (math.frexp(v)[1] if (v := f32(a)) != 0 else 0) & 0xffffffff,
-  # VOP3
-  VOP3Op.V_FMA_F32: lambda a, b, c: i32(f32(a) * f32(b) + f32(c)),
-  VOP3Op.V_DIV_FMAS_F32: lambda a, b, c: i32(f32(a) * f32(b) + f32(c)),
-  VOP3Op.V_ADD3_U32: lambda a, b, c: (a + b + c) & 0xffffffff,
-  VOP3Op.V_LSHL_ADD_U32: lambda a, b, c: ((a << (b & 0x1f)) + c) & 0xffffffff,
-  VOP3Op.V_ADD_LSHL_U32: lambda a, b, c: ((a + b) << (c & 0x1f)) & 0xffffffff,
-  VOP3Op.V_XOR3_B32: lambda a, b, c: a ^ b ^ c,
-  VOP3Op.V_OR3_B32: lambda a, b, c: a | b | c,
-  VOP3Op.V_AND_OR_B32: lambda a, b, c: (a & b) | c,
-  VOP3Op.V_LSHL_OR_B32: lambda a, b, c: ((a << (b & 0x1f)) | c) & 0xffffffff,
-  VOP3Op.V_XAD_U32: lambda a, b, c: ((a ^ b) + c) & 0xffffffff,
-  VOP3Op.V_MAD_U32_U24: lambda a, b, c: ((a & 0xffffff) * (b & 0xffffff) + c) & 0xffffffff,
-  VOP3Op.V_MAD_I32_I24: lambda a, b, c: (sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24) + sext(c, 32)) & 0xffffffff,
-  VOP3Op.V_BFE_U32: lambda a, b, c: (a >> (b & 0x1f)) & ((1 << (c & 0x1f)) - 1) if c & 0x1f else 0,
-  VOP3Op.V_BFE_I32: lambda a, b, c: sext((a >> (b & 0x1f)) & ((1 << w) - 1), w) & 0xffffffff if (w := c & 0x1f) else 0,
-  VOP3Op.V_ALIGNBIT_B32: lambda a, b, c: (((a << 32) | b) >> (c & 0x1f)) & 0xffffffff,
-  VOP3Op.V_MUL_LO_U32: lambda a, b, c: (a * b) & 0xffffffff,
-  VOP3Op.V_MUL_HI_U32: lambda a, b, c: ((a * b) >> 32) & 0xffffffff,
-  VOP3Op.V_MUL_HI_I32: lambda a, b, c: ((sext(a, 32) * sext(b, 32)) >> 32) & 0xffffffff,
-  VOP3Op.V_LDEXP_F32: lambda a, b, c: i32(math.ldexp(f32(a), sext(b, 32))),
-  VOP3Op.V_DIV_FIXUP_F32: lambda a, b, c: i32(math.copysign(float('inf'), f32(c)) if f32(b) == 0.0 else f32(c) / f32(b)),
-  VOP3Op.V_PACK_B32_F16: lambda a, b, c: (a & 0xffff) | ((b & 0xffff) << 16),
-  VOP3Op.V_CVT_PK_RTZ_F16_F32: lambda a, b, c: i16(f32(a)) | (i16(f32(b)) << 16),
-  VOP3Op.V_LSHLREV_B16: lambda a, b, c: ((b & 0xffff) << (a & 0xf)) & 0xffff,
-  VOP3Op.V_LSHRREV_B16: lambda a, b, c: (b & 0xffff) >> (a & 0xf),
-  VOP3Op.V_ASHRREV_I16: lambda a, b, c: (sext(b & 0xffff, 16) >> (a & 0xf)) & 0xffff,
-  VOP3Op.V_ADD_NC_U16: lambda a, b, c: ((a & 0xffff) + (b & 0xffff)) & 0xffff,
-  VOP3Op.V_SUB_NC_U16: lambda a, b, c: ((a & 0xffff) - (b & 0xffff)) & 0xffff,
-  VOP3Op.V_MUL_LO_U16: lambda a, b, c: ((a & 0xffff) * (b & 0xffff)) & 0xffff,
-  VOP3Op.V_MIN_U16: lambda a, b, c: min(a & 0xffff, b & 0xffff),
-  VOP3Op.V_MAX_U16: lambda a, b, c: max(a & 0xffff, b & 0xffff),
-  VOP3Op.V_MIN_I16: lambda a, b, c: (a & 0xffff) if sext(a & 0xffff, 16) < sext(b & 0xffff, 16) else (b & 0xffff),
-  VOP3Op.V_MAX_I16: lambda a, b, c: (a & 0xffff) if sext(a & 0xffff, 16) > sext(b & 0xffff, 16) else (b & 0xffff),
-  VOP3Op.V_MAD_U16: lambda a, b, c: ((a & 0xffff) * (b & 0xffff) + (c & 0xffff)) & 0xffff,
-  VOP3Op.V_MAD_I16: lambda a, b, c: (sext(a & 0xffff, 16) * sext(b & 0xffff, 16) + sext(c & 0xffff, 16)) & 0xffff,
-  VOP3Op.V_FMA_F16: lambda a, b, c: i16(f16(a) * f16(b) + f16(c)),
-  VOP3Op.V_MIN3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[0] & 0xffffffff,
-  VOP3Op.V_MAX3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[2] & 0xffffffff,
-  VOP3Op.V_MED3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[1] & 0xffffffff,
-  VOP3Op.V_MIN3_F16: lambda a, b, c: i16(min(f16(a), f16(b), f16(c))),
-  VOP3Op.V_MAX3_F16: lambda a, b, c: i16(max(f16(a), f16(b), f16(c))),
-  VOP3Op.V_MED3_F16: lambda a, b, c: i16(sorted([f16(a), f16(b), f16(c)])[1]),
-  VOP3Op.V_MIN3_U16: lambda a, b, c: min(a & 0xffff, b & 0xffff, c & 0xffff),
-  VOP3Op.V_MAX3_U16: lambda a, b, c: max(a & 0xffff, b & 0xffff, c & 0xffff),
-  VOP3Op.V_MED3_U16: lambda a, b, c: sorted([a & 0xffff, b & 0xffff, c & 0xffff])[1],
-  VOP3Op.V_MIN3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[0] & 0xffff,
-  VOP3Op.V_MAX3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[2] & 0xffff,
-  VOP3Op.V_MED3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[1] & 0xffff,
-}
-
-def _cmp8(a, b): return [False, a < b, a == b, a <= b, a > b, a != b, a >= b, True]
-def _cmp6(a, b): return [a < b, a == b, a <= b, a > b, a != b, a >= b]
-
-def vopc(op: int, s0: int, s1: int, s0_hi: int = 0, s1_hi: int = 0) -> int:
-  base = op & 0x7f
-  if 16 <= base <= 31:  # F32
-    f0, f1, cmp, nan = f32(s0), f32(s1), base - 16, math.isnan(f32(s0)) or math.isnan(f32(s1))
-    return int([False, f0<f1, f0==f1, f0<=f1, f0>f1, f0!=f1, f0>=f1, not nan, nan, f0<f1 or nan, f0==f1 or nan, f0<=f1 or nan, f0>f1 or nan, f0!=f1 or nan, f0>=f1 or nan, True][cmp])
-  if 49 <= base <= 54: return int(_cmp6(sext(s0 & 0xffff, 16), sext(s1 & 0xffff, 16))[base - 49])  # I16
-  if 57 <= base <= 62: return int(_cmp6(s0 & 0xffff, s1 & 0xffff)[base - 57])  # U16
-  if 64 <= base <= 79:  # I32/U32
-    cmp = (base - 64) % 8
-    return int(_cmp8(sext(s0, 32), sext(s1, 32))[cmp] if base < 72 else _cmp8(s0, s1)[cmp])
-  if 80 <= base <= 95:  # I64/U64
-    s0_64, s1_64 = s0 | (s0_hi << 32), s1 | (s1_hi << 32)
-    return int(_cmp8(sext(s0_64, 64), sext(s1_64, 64))[(base - 80) % 8] if base < 88 else _cmp8(s0_64, s1_64)[(base - 80) % 8])
-  if base == 126:  # CLASS_F32
-    f, mask = f32(s0), s1
-    if math.isnan(f): return int(bool(mask & 0x3))
-    if math.isinf(f): return int(bool(mask & (0x4 if f < 0 else 0x200)))
-    if f == 0.0: return int(bool(mask & (0x20 if (s0 >> 31) & 1 else 0x40)))
-    exp, sign = (s0 >> 23) & 0xff, (s0 >> 31) & 1
-    return int(bool(mask & ((0x10 if sign else 0x80) if exp == 0 else (0x8 if sign else 0x100))))
-  raise NotImplementedError(f"VOPC op {op} (base {base})")

extra/assembly/rdna3/asm.py

@@ -1,602 +0,0 @@
-# RDNA3 assembler and disassembler
-from __future__ import annotations
-import re
-from extra.assembly.rdna3.lib import Inst, RawImm, Reg, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory, FLOAT_ENC, SRC_FIELDS, unwrap
-
-# Decoding helpers
-SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
-SPECIAL_DEC = {**SPECIAL_GPRS, **{v: str(k) for k, v in FLOAT_ENC.items()}}
-SPECIAL_PAIRS = {106: "vcc", 126: "exec"}  # Special register pairs (for 64-bit ops)
-# GFX11 hwreg names (IDs 16-17 are TBA - not supported, IDs 18-19 are PERF_SNAPSHOT)
-HWREG_NAMES = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
-               6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
-               19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI',
-               22: 'HW_REG_XNACK_MASK', 23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
-HWREG_IDS = {v.lower(): k for k, v in HWREG_NAMES.items()}  # Reverse map for assembler
-MSG_NAMES = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
-             131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA'}
-_16BIT_TYPES = ('f16', 'i16', 'u16', 'b16')
-def _is_16bit(s: str) -> bool: return any(s.endswith(x) for x in _16BIT_TYPES)
-
-def decode_src(val: int) -> str:
-  if val <= 105: return f"s{val}"
-  if val in SPECIAL_DEC: return SPECIAL_DEC[val]
-  if 108 <= val <= 123: return f"ttmp{val - 108}"
-  if 128 <= val <= 192: return str(val - 128)
-  if 193 <= val <= 208: return str(-(val - 192))
-  if 256 <= val <= 511: return f"v{val - 256}"
-  return "lit" if val == 255 else f"?{val}"
-
-def _reg(prefix: str, base: int, cnt: int = 1) -> str: return f"{prefix}{base}" if cnt == 1 else f"{prefix}[{base}:{base+cnt-1}]"
-def _sreg(base: int, cnt: int = 1) -> str: return _reg("s", base, cnt)
-def _vreg(base: int, cnt: int = 1) -> str: return _reg("v", base, cnt)
-
-def _fmt_sdst(v: int, cnt: int = 1) -> str:
-  """Format SGPR destination with special register names."""
-  if v == 124: return "null"
-  if 108 <= v <= 123: return _reg("ttmp", v - 108, cnt)
-  if cnt > 1 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
-  if cnt > 1: return _sreg(v, cnt)
-  return {126: "exec_lo", 127: "exec_hi", 106: "vcc_lo", 107: "vcc_hi", 125: "m0"}.get(v, f"s{v}")
-
-def _fmt_ssrc(v: int, cnt: int = 1) -> str:
-  """Format SGPR source with special register names and pairs."""
-  if cnt == 2:
-    if v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
-    if v <= 105: return _sreg(v, 2)
-    if 108 <= v <= 123: return _reg("ttmp", v - 108, 2)
-  return decode_src(v)
-
-def _fmt_src_n(v: int, cnt: int) -> str:
-  """Format source with given register count (1, 2, or 4)."""
-  if cnt == 1: return decode_src(v)
-  if v >= 256: return _vreg(v - 256, cnt)
-  if v <= 105: return _sreg(v, cnt)
-  if cnt == 2 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
-  if 108 <= v <= 123: return _reg("ttmp", v - 108, cnt)
-  return decode_src(v)
-
-def _fmt_src64(v: int) -> str:
-  """Format 64-bit source (VGPR pair, SGPR pair, or special pair)."""
-  return _fmt_src_n(v, 2)
-
-def _parse_sop_sizes(op_name: str) -> tuple[int, ...]:
-  """Parse dst and src sizes from SOP instruction name. Returns (dst_cnt, src0_cnt) or (dst_cnt, src0_cnt, src1_cnt)."""
-  if op_name in ('s_bitset0_b64', 's_bitset1_b64'): return (2, 1)
-  if op_name in ('s_lshl_b64', 's_lshr_b64', 's_ashr_i64', 's_bfe_u64', 's_bfe_i64'): return (2, 2, 1)
-  if op_name in ('s_bfm_b64',): return (2, 1, 1)
-  # SOPC: s_bitcmp0_b64, s_bitcmp1_b64 - 64-bit src0, 32-bit src1 (bit index)
-  if op_name in ('s_bitcmp0_b64', 's_bitcmp1_b64'): return (1, 2, 1)
-  if m := re.search(r'_(b|i|u)(32|64)_(b|i|u)(32|64)$', op_name):
-    return (2 if m.group(2) == '64' else 1, 2 if m.group(4) == '64' else 1)
-  if m := re.search(r'_(b|i|u)(32|64)$', op_name):
-    sz = 2 if m.group(2) == '64' else 1
-    return (sz, sz)
-  return (1, 1)
-
-# Waitcnt helpers (RDNA3 format: bits 15:10=vmcnt, bits 9:4=lgkmcnt, bits 3:0=expcnt)
-def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
-  return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
-def decode_waitcnt(val: int) -> tuple[int, int, int]:
-  return (val >> 10) & 0x3f, val & 0xf, (val >> 4) & 0x3f  # vmcnt, expcnt, lgkmcnt
-
-# VOP3SD opcodes (shared encoding with VOP3 but different field layout)
-# Note: opcodes 0-255 are VOPC promoted to VOP3 - never treat as VOP3SD
-VOP3SD_OPCODES = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
-
-# Disassembler
-def disasm(inst: Inst) -> str:
-  op_val = unwrap(inst._values.get('op', 0))
-  cls_name = inst.__class__.__name__
-  # VOP3 and VOP3SD share encoding - check opcode to determine which
-  is_vop3sd = cls_name == 'VOP3' and op_val in VOP3SD_OPCODES
-  try:
-    from extra.assembly.rdna3 import autogen
-    if is_vop3sd:
-      op_name = autogen.VOP3SDOp(op_val).name.lower()
-    else:
-      op_name = getattr(autogen, f"{cls_name}Op")(op_val).name.lower() if hasattr(autogen, f"{cls_name}Op") else f"op_{op_val}"
-  except (ValueError, KeyError): op_name = f"op_{op_val}"
-  def fmt_src(v): return f"0x{inst._literal:x}" if v == 255 and getattr(inst, '_literal', None) else decode_src(v)
-
-  # VOP1
-  if cls_name == 'VOP1':
-    vdst, src0 = unwrap(inst._values['vdst']), unwrap(inst._values['src0'])
-    if op_name == 'v_nop': return 'v_nop'
-    if op_name == 'v_pipeflush': return 'v_pipeflush'
-    parts = op_name.split('_')
-    is_16bit_dst = any(p in _16BIT_TYPES for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in _16BIT_TYPES and 'cvt' not in op_name)
-    is_16bit_src = parts[-1] in _16BIT_TYPES and 'sat_pk' not in op_name
-    _F64_OPS = ('v_ceil_f64', 'v_floor_f64', 'v_fract_f64', 'v_frexp_mant_f64', 'v_rcp_f64', 'v_rndne_f64', 'v_rsq_f64', 'v_sqrt_f64', 'v_trunc_f64')
-    is_f64_dst = op_name in _F64_OPS or op_name in ('v_cvt_f64_f32', 'v_cvt_f64_i32', 'v_cvt_f64_u32')
-    is_f64_src = op_name in _F64_OPS or op_name in ('v_cvt_f32_f64', 'v_cvt_i32_f64', 'v_cvt_u32_f64', 'v_frexp_exp_i32_f64')
-    if op_name == 'v_readfirstlane_b32':
-      return f"v_readfirstlane_b32 {decode_src(vdst)}, v{src0 - 256 if src0 >= 256 else src0}"
-    dst_str = _vreg(vdst, 2) if is_f64_dst else f"v{vdst & 0x7f}.{'h' if vdst >= 128 else 'l'}" if is_16bit_dst else f"v{vdst}"
-    src_str = _fmt_src64(src0) if is_f64_src else f"v{(src0 - 256) & 0x7f}.{'h' if src0 >= 384 else 'l'}" if is_16bit_src and src0 >= 256 else fmt_src(src0)
-    return f"{op_name}_e32 {dst_str}, {src_str}"
-
-  # VOP2
-  if cls_name == 'VOP2':
-    vdst, src0_raw, vsrc1 = unwrap(inst._values['vdst']), unwrap(inst._values['src0']), unwrap(inst._values['vsrc1'])
-    suffix = "" if op_name == "v_dot2acc_f32_f16" else "_e32"
-    is_16bit_op = ('_f16' in op_name or '_i16' in op_name or '_u16' in op_name) and '_f32' not in op_name and '_i32' not in op_name and 'pk_' not in op_name
-    if is_16bit_op:
-      dst_str = f"v{vdst & 0x7f}.{'h' if vdst >= 128 else 'l'}"
-      src0_str = f"v{(src0_raw - 256) & 0x7f}.{'h' if src0_raw >= 384 else 'l'}" if src0_raw >= 256 else fmt_src(src0_raw)
-      vsrc1_str = f"v{vsrc1 & 0x7f}.{'h' if vsrc1 >= 128 else 'l'}"
-    else:
-      dst_str, src0_str, vsrc1_str = f"v{vdst}", fmt_src(src0_raw), f"v{vsrc1}"
-    return f"{op_name}{suffix} {dst_str}, {src0_str}, {vsrc1_str}" + (", vcc_lo" if op_name == "v_cndmask_b32" else "")
-
-  # VOPC
-  if cls_name == 'VOPC':
-    src0, vsrc1 = unwrap(inst._values['src0']), unwrap(inst._values['vsrc1'])
-    is_64bit = any(x in op_name for x in ('f64', 'i64', 'u64'))
-    is_64bit_vsrc1 = is_64bit and 'class' not in op_name
-    is_16bit = any(x in op_name for x in ('_f16', '_i16', '_u16')) and 'f32' not in op_name
-    is_cmpx = op_name.startswith('v_cmpx')  # VOPCX writes to exec, no vcc destination
-    src0_str = _fmt_src64(src0) if is_64bit else f"v{(src0 - 256) & 0x7f}.{'h' if src0 >= 384 else 'l'}" if is_16bit and src0 >= 256 else fmt_src(src0)
-    vsrc1_str = _vreg(vsrc1, 2) if is_64bit_vsrc1 else f"v{vsrc1 & 0x7f}.{'h' if vsrc1 >= 128 else 'l'}" if is_16bit else f"v{vsrc1}"
-    return f"{op_name}_e32 {src0_str}, {vsrc1_str}" if is_cmpx else f"{op_name}_e32 vcc_lo, {src0_str}, {vsrc1_str}"
-
-  # SOPP
-  if cls_name == 'SOPP':
-    simm16 = unwrap(inst._values.get('simm16', 0))
-    # No-operand instructions (simm16 is ignored)
-    no_imm_ops = ('s_endpgm', 's_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_ttracedata_imm',
-                  's_wait_idle', 's_endpgm_saved', 's_code_end', 's_endpgm_ordered_ps_done')
-    if op_name in no_imm_ops: return op_name
-    if op_name == 's_waitcnt':
-      vmcnt, expcnt, lgkmcnt = decode_waitcnt(simm16)
-      parts = []
-      if vmcnt != 0x3f: parts.append(f"vmcnt({vmcnt})")
-      if expcnt != 0x7: parts.append(f"expcnt({expcnt})")
-      if lgkmcnt != 0x3f: parts.append(f"lgkmcnt({lgkmcnt})")
-      return f"s_waitcnt {' '.join(parts)}" if parts else "s_waitcnt 0"
-    if op_name == 's_delay_alu':
-      dep_names = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2','TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3']
-      skip_names = ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
-      id0, skip, id1 = simm16 & 0xf, (simm16 >> 4) & 0x7, (simm16 >> 7) & 0xf
-      def dep_name(v): return dep_names[v-1] if 0 < v <= len(dep_names) else str(v)
-      parts = [f"instid0({dep_name(id0)})"] if id0 else []
-      if skip: parts.append(f"instskip({skip_names[skip]})")
-      if id1: parts.append(f"instid1({dep_name(id1)})")
-      return f"s_delay_alu {' | '.join(p for p in parts if p)}" if parts else "s_delay_alu 0"
-    if op_name.startswith('s_cbranch') or op_name.startswith('s_branch'):
-      return f"{op_name} {simm16}"
-    # Most SOPP ops require immediate (s_nop, s_setkill, s_sethalt, s_sleep, s_setprio, s_sendmsg*, etc.)
-    return f"{op_name} 0x{simm16:x}"
-
-  # SMEM
-  if cls_name == 'SMEM':
-    if op_name in ('s_gl1_inv', 's_dcache_inv'): return op_name
-    sdata, sbase, soffset, offset = unwrap(inst._values['sdata']), unwrap(inst._values['sbase']), unwrap(inst._values['soffset']), unwrap(inst._values.get('offset', 0))
-    glc, dlc = unwrap(inst._values.get('glc', 0)), unwrap(inst._values.get('dlc', 0))
-    # Format offset: "soffset offset:X" if both, "0x{offset:x}" if only imm, or decode_src(soffset)
-    off_str = f"{decode_src(soffset)} offset:0x{offset:x}" if offset and soffset != 124 else f"0x{offset:x}" if offset else decode_src(soffset)
-    sbase_idx, sbase_cnt = sbase * 2, 4 if (8 <= op_val <= 12 or op_name == 's_atc_probe_buffer') else 2
-    sbase_str = _fmt_ssrc(sbase_idx, sbase_cnt) if sbase_cnt == 2 else _sreg(sbase_idx, sbase_cnt) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_cnt)
-    if op_name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{op_name} {sdata}, {sbase_str}, {off_str}"
-    width = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op_val, 1)
-    mods = [m for m in ["glc" if glc else "", "dlc" if dlc else ""] if m]
-    return f"{op_name} {_fmt_sdst(sdata, width)}, {sbase_str}, {off_str}" + (" " + " ".join(mods) if mods else "")
-
-  # FLAT
-  if cls_name == 'FLAT':
-    vdst, addr, data, saddr, offset, seg = [unwrap(inst._values.get(f, 0)) for f in ['vdst', 'addr', 'data', 'saddr', 'offset', 'seg']]
-    instr = f"{['flat', 'scratch', 'global'][seg] if seg < 3 else 'flat'}_{op_name.split('_', 1)[1] if '_' in op_name else op_name}"
-    width = {'b32':1, 'b64':2, 'b96':3, 'b128':4, 'u8':1, 'i8':1, 'u16':1, 'i16':1}.get(op_name.split('_')[-1], 1)
-    addr_str = _vreg(addr, 2) if saddr == 0x7F else _vreg(addr)
-    saddr_str = "" if saddr == 0x7F else f", {_sreg(saddr, 2)}" if saddr < 106 else ", off" if saddr == 124 else f", {decode_src(saddr)}"
-    off_str = f" offset:{offset}" if offset else ""
-    vdata_str = _vreg(data if 'store' in op_name else vdst, width)
-    return f"{instr} {addr_str}, {vdata_str}{saddr_str}{off_str}" if 'store' in op_name else f"{instr} {vdata_str}, {addr_str}{saddr_str}{off_str}"
-
-  # VOP3: vector ops with modifiers (can be 1, 2, or 3 sources depending on opcode range)
-  if cls_name == 'VOP3':
-    # Handle VOP3SD opcodes (same encoding, different field layout)
-    if is_vop3sd:
-      vdst = unwrap(inst._values.get('vdst', 0))
-      # VOP3SD: sdst is at bits [14:8], but VOP3 decodes opsel at [14:11], abs at [10:8], clmp at [15]
-      # We need to reconstruct sdst from these fields
-      opsel_raw = unwrap(inst._values.get('opsel', 0))
-      abs_raw = unwrap(inst._values.get('abs', 0))
-      clmp_raw = unwrap(inst._values.get('clmp', 0))
-      sdst = (clmp_raw << 7) | (opsel_raw << 3) | abs_raw
-      src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
-      neg = unwrap(inst._values.get('neg', 0))
-      omod = unwrap(inst._values.get('omod', 0))
-      omod_str = {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
-      is_f64 = 'f64' in op_name
-      # v_mad_i64_i32/v_mad_u64_u32: 64-bit dst and src2, 32-bit src0/src1
-      is_mad64 = 'mad_i64_i32' in op_name or 'mad_u64_u32' in op_name
-      def fmt_sd_src(v, neg_bit, is_64bit=False):
-        s = _fmt_src64(v) if (is_64bit or is_f64) else fmt_src(v)
-        return f"-{s}" if neg_bit else s
-      src0_str, src1_str = fmt_sd_src(src0, neg & 1), fmt_sd_src(src1, neg & 2)
-      src2_str = fmt_sd_src(src2, neg & 4, is_mad64)
-      dst_str = _vreg(vdst, 2) if (is_f64 or is_mad64) else f"v{vdst}"
-      sdst_str = _fmt_sdst(sdst, 1)
-      # v_add_co_u32, v_sub_co_u32, v_subrev_co_u32, v_add_co_ci_u32, etc. only use 2 sources
-      if op_name in ('v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32', 'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'):
-        return f"{op_name} {dst_str}, {sdst_str}, {src0_str}, {src1_str}"
-      # v_div_scale uses 3 sources
-      return f"{op_name} {dst_str}, {sdst_str}, {src0_str}, {src1_str}, {src2_str}" + omod_str
-
-    vdst = unwrap(inst._values.get('vdst', 0))
-    src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
-    neg, abs_, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('abs', 0)), unwrap(inst._values.get('clmp', 0))
-    opsel = unwrap(inst._values.get('opsel', 0))
-    # Check if 64-bit op (needs register pairs)
-    is_f64 = 'f64' in op_name or 'i64' in op_name or 'u64' in op_name or 'b64' in op_name
-    # v_cmp_class_* has 64-bit src0 but 32-bit src1 (class mask)
-    is_class = 'class' in op_name
-    # Shift ops: v_*rev_*64 have 32-bit shift amount (src0), 64-bit value (src1)
-    is_shift64 = 'rev' in op_name and '64' in op_name and op_name.startswith('v_')
-    # v_ldexp_f64: 64-bit src0 (mantissa), 32-bit src1 (exponent)
-    is_ldexp64 = op_name == 'v_ldexp_f64'
-    # v_trig_preop_f64: 64-bit dst/src0, 32-bit src1 (exponent/scale)
-    is_trig_preop = op_name == 'v_trig_preop_f64'
-    # v_readlane_b32: destination is SGPR (despite vdst field)
-    is_readlane = op_name == 'v_readlane_b32'
-    # SAD/QSAD/MQSAD instructions have mixed sizes
-    # v_qsad_pk_u16_u8, v_mqsad_pk_u16_u8: 64-bit dst/src0/src2, 32-bit src1
-    # v_mqsad_u32_u8: 128-bit (4 reg) dst/src2, 64-bit src0, 32-bit src1
-    is_sad64 = any(x in op_name for x in ('qsad_pk', 'mqsad_pk'))
-    is_mqsad_u32 = 'mqsad_u32' in op_name
-    # Detect 16-bit and 64-bit operand sizes for various instruction patterns
-    if 'cvt_pk' in op_name:
-      is_f16_dst, is_f16_src, is_f16_src2 = False, op_name.endswith('16'), False
-    elif m := re.match(r'v_(?:cvt|frexp_exp)_([a-z0-9_]+)_([a-z0-9]+)', op_name):
-      dst_type, src_type = m.group(1), m.group(2)
-      is_f16_dst, is_f16_src, is_f16_src2 = _is_16bit(dst_type), _is_16bit(src_type), _is_16bit(src_type)
-      is_f64_dst, is_f64_src, is_f64 = '64' in dst_type, '64' in src_type, False
-    elif re.match(r'v_mad_[iu]32_[iu]16', op_name):
-      is_f16_dst, is_f16_src, is_f16_src2 = False, True, False  # 32-bit dst, 16-bit src0/src1, 32-bit src2
-    elif 'pack_b32' in op_name:
-      is_f16_dst, is_f16_src, is_f16_src2 = False, True, True  # 32-bit dst, 16-bit sources
-    else:
-      is_16bit_op = any(x in op_name for x in _16BIT_TYPES) and not any(x in op_name for x in ('dot2', 'pk_', 'sad', 'msad', 'qsad', 'mqsad'))
-      is_f16_dst = is_f16_src = is_f16_src2 = is_16bit_op
-    # Check if any opsel bit is set (any operand uses .h) - if so, we need explicit .l for low-half
-    any_hi = opsel != 0
-    def fmt_vop3_src(v, neg_bit, abs_bit, hi_bit=False, reg_cnt=1, is_16=False):
-      s = _fmt_src_n(v, reg_cnt) if reg_cnt > 1 else f"v{v - 256}.h" if is_16 and v >= 256 and hi_bit else f"v{v - 256}.l" if is_16 and v >= 256 and any_hi else fmt_src(v)
-      if abs_bit: s = f"|{s}|"
-      return f"-{s}" if neg_bit else s
-    # Determine register count for each source (check for cvt-specific 64-bit flags first)
-    is_src0_64 = locals().get('is_f64_src', is_f64 and not is_shift64) or is_sad64 or is_mqsad_u32
-    is_src1_64 = is_f64 and not is_class and not is_ldexp64 and not is_trig_preop
-    src0_cnt = 2 if is_src0_64 else 1
-    src1_cnt = 2 if is_src1_64 else 1
-    src2_cnt = 4 if is_mqsad_u32 else 2 if (is_f64 or is_sad64) else 1
-    src0_str = fmt_vop3_src(src0, neg & 1, abs_ & 1, opsel & 1, src0_cnt, is_f16_src)
-    src1_str = fmt_vop3_src(src1, neg & 2, abs_ & 2, opsel & 2, src1_cnt, is_f16_src)
-    src2_str = fmt_vop3_src(src2, neg & 4, abs_ & 4, opsel & 4, src2_cnt, is_f16_src2)
-    # Format destination - for 16-bit ops, use .h/.l suffix; readlane uses SGPR dest
-    is_dst_64 = locals().get('is_f64_dst', is_f64) or is_sad64
-    dst_cnt = 4 if is_mqsad_u32 else 2 if is_dst_64 else 1
-    if is_readlane:
-      dst_str = _fmt_sdst(vdst, 1)
-    elif dst_cnt > 1:
-      dst_str = _vreg(vdst, dst_cnt)
-    elif is_f16_dst:
-      dst_str = f"v{vdst}.h" if (opsel & 8) else f"v{vdst}.l" if any_hi else f"v{vdst}"
-    else:
-      dst_str = f"v{vdst}"
-    clamp_str = " clamp" if clmp else ""
-    omod = unwrap(inst._values.get('omod', 0))
-    omod_str = {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
-    # op_sel for non-VGPR sources (when opsel bits are set but source is not a VGPR)
-    # For 16-bit ops with VGPR sources, opsel is encoded in .h/.l suffix
-    # For non-VGPR sources or non-16-bit ops, we need explicit op_sel
-    has_nonvgpr_opsel = (src0 < 256 and (opsel & 1)) or (src1 < 256 and (opsel & 2)) or (src2 < 256 and (opsel & 4))
-    need_opsel = has_nonvgpr_opsel or (opsel and not is_f16_src)
-    # Helper to format opsel string based on source count
-    def fmt_opsel(num_src):
-      if not need_opsel: return ""
-      # When dst is .h (for 16-bit ops) and non-VGPR sources have opsel, use all 1s
-      if is_f16_dst and (opsel & 8):  # dst is .h
-        return f" op_sel:[1,1,1{',1' if num_src == 3 else ''}]"
-      # Otherwise output actual opsel values
-      if num_src == 3:
-        return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{(opsel >> 3) & 1}]"
-      return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1}]"
-    # Determine number of sources based on opcode range:
-    # 0-255: VOPC promoted (comparison, 2 src, sdst)
-    # 256-383: VOP2 promoted (2 src)
-    # 384-511: VOP1 promoted (1 src)
-    # 512+: Native VOP3 (2 or 3 src depending on instruction)
-    if op_val < 256:  # VOPC promoted
-      # VOPCX (v_cmpx_*) writes to exec, no explicit destination
-      if op_name.startswith('v_cmpx'):
-        return f"{op_name}_e64 {src0_str}, {src1_str}"
-      return f"{op_name}_e64 {_fmt_sdst(vdst, 1)}, {src0_str}, {src1_str}"
-    elif op_val < 384:  # VOP2 promoted
-      # v_cndmask_b32 in VOP3 format has 3 sources (src2 is mask selector)
-      if 'cndmask' in op_name:
-        return f"{op_name}_e64 {dst_str}, {src0_str}, {src1_str}, {src2_str}" + fmt_opsel(3) + clamp_str + omod_str
-      return f"{op_name}_e64 {dst_str}, {src0_str}, {src1_str}" + fmt_opsel(2) + clamp_str + omod_str
-    elif op_val < 512:  # VOP1 promoted
-      if op_name in ('v_nop', 'v_pipeflush'): return f"{op_name}_e64"
-      return f"{op_name}_e64 {dst_str}, {src0_str}" + fmt_opsel(1) + clamp_str + omod_str
-    else:  # Native VOP3 - determine 2 vs 3 sources based on instruction name
-      # 3-source ops: fma, mad, min3, max3, med3, div_fixup, div_fmas, sad, msad, qsad, mqsad, lerp, alignbit/byte, cubeid/sc/tc/ma, bfe, bfi, perm_b32, permlane, cndmask
-      # Note: v_writelane_b32 is 2-src (src0, src1 with vdst as 3rd operand - read-modify-write)
-      is_3src = any(x in op_name for x in ('fma', 'mad', 'min3', 'max3', 'med3', 'div_fix', 'div_fmas', 'sad', 'lerp', 'align', 'cube',
-                                            'bfe', 'bfi', 'perm_b32', 'permlane', 'cndmask', 'xor3', 'or3', 'add3', 'lshl_or', 'and_or', 'lshl_add',
-                                            'add_lshl', 'xad', 'maxmin', 'minmax', 'dot2', 'cvt_pk_u8', 'mullit'))
-      if is_3src:
-        return f"{op_name} {dst_str}, {src0_str}, {src1_str}, {src2_str}" + fmt_opsel(3) + clamp_str + omod_str
-      return f"{op_name} {dst_str}, {src0_str}, {src1_str}" + fmt_opsel(2) + clamp_str + omod_str
-
-  # VOP3SD: 3-source with scalar destination (v_div_scale_*, v_add_co_u32, v_mad_*64_*32, etc.)
-  if cls_name == 'VOP3SD':
-    vdst, sdst = unwrap(inst._values.get('vdst', 0)), unwrap(inst._values.get('sdst', 0))
-    src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
-    neg, omod, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('omod', 0)), unwrap(inst._values.get('clmp', 0))
-    is_f64, is_mad64 = 'f64' in op_name, 'mad_i64_i32' in op_name or 'mad_u64_u32' in op_name
-    def fmt_neg(v, neg_bit, is_64=False): return f"-{_fmt_src64(v) if (is_64 or is_f64) else fmt_src(v)}" if neg_bit else _fmt_src64(v) if (is_64 or is_f64) else fmt_src(v)
-    srcs = [fmt_neg(src0, neg & 1), fmt_neg(src1, neg & 2), fmt_neg(src2, neg & 4, is_mad64)]
-    dst_str, sdst_str = _vreg(vdst, 2) if (is_f64 or is_mad64) else f"v{vdst}", _fmt_sdst(sdst, 1)
-    clamp_str, omod_str = " clamp" if clmp else "", {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
-    is_2src = op_name in ('v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32')
-    suffix = "_e64" if op_name.startswith('v_') and 'co_' in op_name else ""
-    return f"{op_name}{suffix} {dst_str}, {sdst_str}, {', '.join(srcs[:2] if is_2src else srcs)}" + clamp_str + omod_str
-
-  # VOPD: dual-issue instructions
-  if cls_name == 'VOPD':
-    from extra.assembly.rdna3 import autogen
-    opx, opy, vdstx, vdsty_enc = [unwrap(inst._values.get(f, 0)) for f in ('opx', 'opy', 'vdstx', 'vdsty')]
-    srcx0, vsrcx1, srcy0, vsrcy1 = [unwrap(inst._values.get(f, 0)) for f in ('srcx0', 'vsrcx1', 'srcy0', 'vsrcy1')]
-    vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1)  # Decode vdsty
-    def fmt_vopd(op, vdst, src0, vsrc1):
-      try: name = autogen.VOPDOp(op).name.lower()
-      except (ValueError, KeyError): name = f"op_{op}"
-      return f"{name} v{vdst}, {fmt_src(src0)}" if 'mov' in name else f"{name} v{vdst}, {fmt_src(src0)}, v{vsrc1}"
-    return f"{fmt_vopd(opx, vdstx, srcx0, vsrcx1)} :: {fmt_vopd(opy, vdsty, srcy0, vsrcy1)}"
-
-  # VOP3P: packed vector ops
-  if cls_name == 'VOP3P':
-    vdst, clmp = unwrap(inst._values.get('vdst', 0)), unwrap(inst._values.get('clmp', 0))
-    src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
-    neg, neg_hi = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('neg_hi', 0))
-    opsel, opsel_hi, opsel_hi2 = unwrap(inst._values.get('opsel', 0)), unwrap(inst._values.get('opsel_hi', 0)), unwrap(inst._values.get('opsel_hi2', 0))
-    is_wmma, is_3src = 'wmma' in op_name, any(x in op_name for x in ('fma', 'mad', 'dot', 'wmma'))
-    def fmt_bits(name, val, n): return f"{name}:[{','.join(str((val >> i) & 1) for i in range(n))}]"
-    # WMMA: f16/bf16 use 8-reg sources, iu8 uses 4-reg, iu4 uses 2-reg; all have 8-reg dst
-    if is_wmma:
-      src_cnt = 2 if 'iu4' in op_name else 4 if 'iu8' in op_name else 8
-      src0_str, src1_str, src2_str = _fmt_src_n(src0, src_cnt), _fmt_src_n(src1, src_cnt), _fmt_src_n(src2, 8)
-      dst_str = _vreg(vdst, 8)
-    else:
-      src0_str, src1_str, src2_str = _fmt_src_n(src0, 1), _fmt_src_n(src1, 1), _fmt_src_n(src2, 1)
-      dst_str = f"v{vdst}"
-    n = 3 if is_3src else 2
-    full_opsel_hi = opsel_hi | (opsel_hi2 << 2)
-    mods = [fmt_bits("op_sel", opsel, n)] if opsel else []
-    if full_opsel_hi != (0b111 if is_3src else 0b11): mods.append(fmt_bits("op_sel_hi", full_opsel_hi, n))
-    if neg: mods.append(fmt_bits("neg_lo", neg, n))
-    if neg_hi: mods.append(fmt_bits("neg_hi", neg_hi, n))
-    if clmp: mods.append("clamp")
-    mod_str = " " + " ".join(mods) if mods else ""
-    return f"{op_name} {dst_str}, {src0_str}, {src1_str}, {src2_str}{mod_str}" if is_3src else f"{op_name} {dst_str}, {src0_str}, {src1_str}{mod_str}"
-
-  # VINTERP: interpolation instructions
-  if cls_name == 'VINTERP':
-    vdst = unwrap(inst._values.get('vdst', 0))
-    src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
-    neg, waitexp, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('waitexp', 0)), unwrap(inst._values.get('clmp', 0))
-    def fmt_neg_vi(v, neg_bit): return f"-{v}" if neg_bit else v
-    srcs = [fmt_neg_vi(f"v{s - 256}" if s >= 256 else fmt_src(s), neg & (1 << i)) for i, s in enumerate([src0, src1, src2])]
-    mods = [m for m in [f"wait_exp:{waitexp}" if waitexp else "", "clamp" if clmp else ""] if m]
-    return f"{op_name} v{vdst}, {', '.join(srcs)}" + (" " + " ".join(mods) if mods else "")
-
-  # MUBUF/MTBUF helpers
-  def _buf_vaddr(vaddr, offen, idxen): return _vreg(vaddr, 2) if offen and idxen else f"v{vaddr}" if offen or idxen else "off"
-  def _buf_srsrc(srsrc): srsrc_base = srsrc * 4; return _reg("ttmp", srsrc_base - 108, 4) if 108 <= srsrc_base <= 123 else _sreg(srsrc_base, 4)
-
-  # MUBUF: buffer load/store
-  if cls_name == 'MUBUF':
-    vdata, vaddr, srsrc, soffset = [unwrap(inst._values.get(f, 0)) for f in ('vdata', 'vaddr', 'srsrc', 'soffset')]
-    offset, offen, idxen = unwrap(inst._values.get('offset', 0)), unwrap(inst._values.get('offen', 0)), unwrap(inst._values.get('idxen', 0))
-    glc, dlc, slc, tfe = [unwrap(inst._values.get(f, 0)) for f in ('glc', 'dlc', 'slc', 'tfe')]
-    if op_name in ('buffer_gl0_inv', 'buffer_gl1_inv'): return op_name
-    # Determine data width from op name
-    if 'd16' in op_name: width = 2 if any(x in op_name for x in ('xyz', 'xyzw')) else 1
-    elif 'atomic' in op_name:
-      base_width = 2 if any(x in op_name for x in ('b64', 'u64', 'i64')) else 1
-      width = base_width * 2 if 'cmpswap' in op_name else base_width
-    else: width = {'b32':1, 'b64':2, 'b96':3, 'b128':4, 'b16':1, 'x':1, 'xy':2, 'xyz':3, 'xyzw':4}.get(op_name.split('_')[-1], 1)
-    if tfe: width += 1
-    mods = [m for m in ["offen" if offen else "", "idxen" if idxen else "", f"offset:{offset}" if offset else "",
-                        "glc" if glc else "", "dlc" if dlc else "", "slc" if slc else "", "tfe" if tfe else ""] if m]
-    return f"{op_name} {_vreg(vdata, width)}, {_buf_vaddr(vaddr, offen, idxen)}, {_buf_srsrc(srsrc)}, {decode_src(soffset)}" + (" " + " ".join(mods) if mods else "")
-
-  # MTBUF: typed buffer load/store
-  if cls_name == 'MTBUF':
-    vdata, vaddr, srsrc, soffset = [unwrap(inst._values.get(f, 0)) for f in ('vdata', 'vaddr', 'srsrc', 'soffset')]
-    offset, tbuf_fmt, offen, idxen = [unwrap(inst._values.get(f, 0)) for f in ('offset', 'format', 'offen', 'idxen')]
-    glc, dlc, slc = [unwrap(inst._values.get(f, 0)) for f in ('glc', 'dlc', 'slc')]
-    mods = [f"format:{tbuf_fmt}"] + [m for m in ["idxen" if idxen else "", "offen" if offen else "", f"offset:{offset}" if offset else "",
-                                                  "glc" if glc else "", "dlc" if dlc else "", "slc" if slc else ""] if m]
-    width = 2 if 'd16' in op_name and any(x in op_name for x in ('xyz', 'xyzw')) else 1 if 'd16' in op_name else {'x':1, 'xy':2, 'xyz':3, 'xyzw':4}.get(op_name.split('_')[-1], 1)
-    return f"{op_name} {_vreg(vdata, width)}, {_buf_vaddr(vaddr, offen, idxen)}, {_buf_srsrc(srsrc)}, {decode_src(soffset)} {' '.join(mods)}"
-
-  # SOP1/SOP2/SOPC/SOPK
-  if cls_name in ('SOP1', 'SOP2', 'SOPC', 'SOPK'):
-    sizes = _parse_sop_sizes(op_name)
-    dst_cnt, src0_cnt = sizes[0], sizes[1]
-    src1_cnt = sizes[2] if len(sizes) > 2 else src0_cnt
-    if cls_name == 'SOP1':
-      sdst, ssrc0 = unwrap(inst._values.get('sdst', 0)), unwrap(inst._values.get('ssrc0', 0))
-      if op_name == 's_getpc_b64': return f"{op_name} {_fmt_sdst(sdst, 2)}"
-      if op_name in ('s_setpc_b64', 's_rfe_b64'): return f"{op_name} {_fmt_ssrc(ssrc0, 2)}"
-      if op_name == 's_swappc_b64': return f"{op_name} {_fmt_sdst(sdst, 2)}, {_fmt_ssrc(ssrc0, 2)}"
-      if op_name in ('s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'):
-        return f"{op_name} {_fmt_sdst(sdst, 2 if 'b64' in op_name else 1)}, sendmsg({MSG_NAMES.get(ssrc0, str(ssrc0))})"
-      ssrc0_str = fmt_src(ssrc0) if src0_cnt == 1 else _fmt_ssrc(ssrc0, src0_cnt)
-      return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, {ssrc0_str}"
-    if cls_name == 'SOP2':
-      sdst, ssrc0, ssrc1 = [unwrap(inst._values.get(f, 0)) for f in ('sdst', 'ssrc0', 'ssrc1')]
-      return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, {_fmt_ssrc(ssrc0, src0_cnt)}, {_fmt_ssrc(ssrc1, src1_cnt)}"
-    if cls_name == 'SOPC':
-      return f"{op_name} {_fmt_ssrc(unwrap(inst._values.get('ssrc0', 0)), src0_cnt)}, {_fmt_ssrc(unwrap(inst._values.get('ssrc1', 0)), src1_cnt)}"
-    if cls_name == 'SOPK':
-      sdst, simm16 = unwrap(inst._values.get('sdst', 0)), unwrap(inst._values.get('simm16', 0))
-      if op_name == 's_version': return f"{op_name} 0x{simm16:x}"
-      if op_name in ('s_setreg_b32', 's_getreg_b32'):
-        hwreg_id, hwreg_offset, hwreg_size = simm16 & 0x3f, (simm16 >> 6) & 0x1f, ((simm16 >> 11) & 0x1f) + 1
-        hwreg_str = f"0x{simm16:x}" if hwreg_id in (16, 17) else f"hwreg({HWREG_NAMES.get(hwreg_id, str(hwreg_id))}, {hwreg_offset}, {hwreg_size})"
-        return f"{op_name} {hwreg_str}, {_fmt_sdst(sdst, 1)}" if op_name == 's_setreg_b32' else f"{op_name} {_fmt_sdst(sdst, 1)}, {hwreg_str}"
-      return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, 0x{simm16:x}"
-
-  # Generic fallback
-  def fmt_field(n, v):
-    v = unwrap(v)
-    if n in SRC_FIELDS: return fmt_src(v) if v != 255 else "0xff"
-    if n in ('sdst', 'vdst'): return f"{'s' if n == 'sdst' else 'v'}{v}"
-    return f"v{v}" if n == 'vsrc1' else f"0x{v:x}" if n == 'simm16' else str(v)
-  ops = [fmt_field(n, inst._values.get(n, 0)) for n in inst._fields if n not in ('encoding', 'op')]
-  return f"{op_name} {', '.join(ops)}" if ops else op_name
-
-# Assembler
-SPECIAL_REGS = {'vcc_lo': RawImm(106), 'vcc_hi': RawImm(107), 'null': RawImm(124), 'off': RawImm(124), 'm0': RawImm(125), 'exec_lo': RawImm(126), 'exec_hi': RawImm(127), 'scc': RawImm(253)}
-FLOAT_CONSTS = {'0.5': 0.5, '-0.5': -0.5, '1.0': 1.0, '-1.0': -1.0, '2.0': 2.0, '-2.0': -2.0, '4.0': 4.0, '-4.0': -4.0}
-REG_MAP: dict[str, _RegFactory] = {'s': s, 'v': v, 't': ttmp, 'ttmp': ttmp}
-
-def parse_operand(op: str) -> tuple:
-  op = op.strip().lower()
-  neg = op.startswith('-') and not op[1:2].isdigit(); op = op[1:] if neg else op
-  abs_ = op.startswith('|') and op.endswith('|') or op.startswith('abs(') and op.endswith(')')
-  op = op[1:-1] if op.startswith('|') else op[4:-1] if op.startswith('abs(') else op
-  hi_half = op.endswith('.h')
-  op = re.sub(r'\.[lh]$', '', op)
-  if op in FLOAT_CONSTS: return (FLOAT_CONSTS[op], neg, abs_, hi_half)
-  if re.match(r'^-?\d+$', op): return (int(op), neg, abs_, hi_half)
-  if m := re.match(r'^-?0x([0-9a-f]+)$', op):
-    v = -int(m.group(1), 16) if op.startswith('-') else int(m.group(1), 16)
-    return (v, neg, abs_, hi_half)
-  if op in SPECIAL_REGS: return (SPECIAL_REGS[op], neg, abs_, hi_half)
-  if op == 'lit': return (RawImm(255), neg, abs_, hi_half)  # literal marker (actual value comes from literal word)
-  if m := re.match(r'^([svt](?:tmp)?)\[(\d+):(\d+)\]$', op): return (REG_MAP[m.group(1)][int(m.group(2)):int(m.group(3))], neg, abs_, hi_half)
-  if m := re.match(r'^([svt](?:tmp)?)(\d+)$', op):
-    reg = REG_MAP[m.group(1)][int(m.group(2))]
-    reg.hi = hi_half
-    return (reg, neg, abs_, hi_half)
-  # hwreg(name, offset, size) or hwreg(name) -> simm16 encoding
-  if m := re.match(r'^hwreg\((\w+)(?:,\s*(\d+),\s*(\d+))?\)$', op):
-    name_str = m.group(1).lower()
-    hwreg_id = HWREG_IDS.get(name_str, int(name_str) if name_str.isdigit() else None)
-    if hwreg_id is None: raise ValueError(f"unknown hwreg name: {name_str}")
-    offset, size = int(m.group(2)) if m.group(2) else 0, int(m.group(3)) if m.group(3) else 32
-    return (((size - 1) << 11) | (offset << 6) | hwreg_id, neg, abs_, hi_half)
-  raise ValueError(f"cannot parse operand: {op}")
-
-SMEM_OPS = {'s_load_b32', 's_load_b64', 's_load_b128', 's_load_b256', 's_load_b512',
-            's_buffer_load_b32', 's_buffer_load_b64', 's_buffer_load_b128', 's_buffer_load_b256', 's_buffer_load_b512'}
-SOP1_SRC_ONLY = {'s_setpc_b64', 's_rfe_b64'}
-SOP1_MSG_IMM = {'s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'}
-SOPK_IMM_ONLY = {'s_version'}
-SOPK_IMM_FIRST = {'s_setreg_b32'}
-SOPK_UNSUPPORTED = {'s_setreg_imm32_b32'}
-
-def asm(text: str) -> Inst:
-  from extra.assembly.rdna3 import autogen
-  text = text.strip()
-  clamp = 'clamp' in text.lower()
-  if clamp: text = re.sub(r'\s+clamp\s*$', '', text, flags=re.I)
-  modifiers = {}
-  if m := re.search(r'\s+wait_exp:(\d+)', text, re.I): modifiers['waitexp'] = int(m.group(1)); text = text[:m.start()] + text[m.end():]
-  parts = text.replace(',', ' ').split()
-  if not parts: raise ValueError("empty instruction")
-  mnemonic, op_str = parts[0].lower(), text[len(parts[0]):].strip()
-  # Handle s_waitcnt specially before operand parsing
-  if mnemonic == 's_waitcnt':
-    vmcnt, expcnt, lgkmcnt = 0x3f, 0x7, 0x3f
-    for part in op_str.replace(',', ' ').split():
-      if m := re.match(r'vmcnt\((\d+)\)', part): vmcnt = int(m.group(1))
-      elif m := re.match(r'expcnt\((\d+)\)', part): expcnt = int(m.group(1))
-      elif m := re.match(r'lgkmcnt\((\d+)\)', part): lgkmcnt = int(m.group(1))
-      elif re.match(r'^0x[0-9a-f]+$|^\d+$', part): return autogen.s_waitcnt(simm16=int(part, 0))
-    return autogen.s_waitcnt(simm16=waitcnt(vmcnt, expcnt, lgkmcnt))
-  # Handle VOPD dual-issue instructions: opx dst, src :: opy dst, src
-  if '::' in text:
-    x_part, y_part = text.split('::')
-    x_parts, y_parts = x_part.strip().replace(',', ' ').split(), y_part.strip().replace(',', ' ').split()
-    opx_name, opy_name = x_parts[0].upper(), y_parts[0].upper()
-    opx, opy = autogen.VOPDOp[opx_name], autogen.VOPDOp[opy_name]
-    x_ops, y_ops = [parse_operand(p)[0] for p in x_parts[1:]], [parse_operand(p)[0] for p in y_parts[1:]]
-    vdstx, srcx0 = x_ops[0], x_ops[1] if len(x_ops) > 1 else 0
-    vsrcx1 = x_ops[2] if len(x_ops) > 2 else VGPR(0)
-    vdsty, srcy0 = y_ops[0], y_ops[1] if len(y_ops) > 1 else 0
-    vsrcy1 = y_ops[2] if len(y_ops) > 2 else VGPR(0)
-    # Handle fmaak/fmamk literals (4th operand on x or y side)
-    lit = None
-    if 'fmaak' in opx_name.lower() and len(x_ops) > 3: lit = unwrap(x_ops[3])
-    elif 'fmamk' in opx_name.lower() and len(x_ops) > 3: lit, vsrcx1 = unwrap(x_ops[2]), x_ops[3]
-    elif 'fmaak' in opy_name.lower() and len(y_ops) > 3: lit = unwrap(y_ops[3])
-    elif 'fmamk' in opy_name.lower() and len(y_ops) > 3: lit, vsrcy1 = unwrap(y_ops[2]), y_ops[3]
-    return autogen.VOPD(opx, opy, vdstx=vdstx, vdsty=vdsty, srcx0=srcx0, vsrcx1=vsrcx1, srcy0=srcy0, vsrcy1=vsrcy1, literal=lit)
-  operands, current, depth, in_pipe = [], "", 0, False
-  for ch in op_str:
-    if ch in '[(': depth += 1
-    elif ch in '])': depth -= 1
-    elif ch == '|': in_pipe = not in_pipe
-    if ch == ',' and depth == 0 and not in_pipe: operands.append(current.strip()); current = ""
-    else: current += ch
-  if current.strip(): operands.append(current.strip())
-  parsed = [parse_operand(op) for op in operands]
-  values = [p[0] for p in parsed]
-  neg_bits = sum((1 << (i-1)) for i, p in enumerate(parsed) if i > 0 and p[1])
-  abs_bits = sum((1 << (i-1)) for i, p in enumerate(parsed) if i > 0 and p[2])
-  opsel_bits = (8 if len(parsed) > 0 and parsed[0][3] else 0) | sum((1 << i) for i, p in enumerate(parsed[1:4]) if p[3])
-  lit = None
-  if mnemonic in ('v_fmaak_f32', 'v_fmaak_f16') and len(values) == 4: lit, values = unwrap(values[3]), values[:3]
-  elif mnemonic in ('v_fmamk_f32', 'v_fmamk_f16') and len(values) == 4: lit, values = unwrap(values[2]), [values[0], values[1], values[3]]
-  vcc_ops = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32', 'v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32'}
-  if mnemonic.replace('_e32', '') in vcc_ops and len(values) >= 5: values = [values[0], values[2], values[3]]
-  if mnemonic.startswith('v_cmp') and len(values) >= 3 and operands[0].strip().lower() in ('vcc_lo', 'vcc_hi', 'vcc'):
-    values = values[1:]
-  # CMPX instructions with _e64 suffix: prepend implicit EXEC_LO destination (vdst=126)
-  if 'cmpx' in mnemonic and mnemonic.endswith('_e64') and len(values) == 2:
-    values = [VGPR(126, 1)] + values
-    # Recalculate modifiers: parsed[0]=src0, parsed[1]=src1 (no vdst in user input)
-    neg_bits = sum((1 << i) for i, p in enumerate(parsed[:3]) if p[1])
-    abs_bits = sum((1 << i) for i, p in enumerate(parsed[:3]) if p[2])
-    opsel_bits = sum((1 << i) for i, p in enumerate(parsed[:2]) if p[3])
-  vop3sd_ops = {'v_div_scale_f32', 'v_div_scale_f64'}
-  if mnemonic in vop3sd_ops and len(parsed) >= 5:
-    neg_bits = sum((1 << i) for i, p in enumerate(parsed[2:5]) if p[1])
-    abs_bits = sum((1 << i) for i, p in enumerate(parsed[2:5]) if p[2])
-  if mnemonic in SOPK_UNSUPPORTED: raise ValueError(f"unsupported instruction: {mnemonic}")
-  elif mnemonic in SOP1_SRC_ONLY:
-    return getattr(autogen, mnemonic)(ssrc0=values[0])
-  elif mnemonic in SOP1_MSG_IMM:
-    return getattr(autogen, mnemonic)(sdst=values[0], ssrc0=RawImm(unwrap(values[1])))
-  elif mnemonic in SOPK_IMM_ONLY:
-    return getattr(autogen, mnemonic)(simm16=values[0])
-  elif mnemonic in SOPK_IMM_FIRST:
-    return getattr(autogen, mnemonic)(simm16=values[0], sdst=values[1])
-  elif mnemonic in SMEM_OPS and len(operands) >= 3 and re.match(r'^-?[0-9]|^-?0x', operands[2].strip().lower()):
-    return getattr(autogen, mnemonic)(sdata=values[0], sbase=values[1], offset=values[2], soffset=RawImm(124))
-  elif mnemonic.startswith('buffer_') and len(operands) >= 2 and operands[1].strip().lower() == 'off':
-    return getattr(autogen, mnemonic)(vdata=values[0], vaddr=0, srsrc=values[2], soffset=RawImm(unwrap(values[3])) if len(values) > 3 else RawImm(0))
-  elif (mnemonic.startswith('flat_load') or mnemonic.startswith('global_load') or mnemonic.startswith('scratch_load')) and len(values) >= 3:
-    offset = int(m.group(1)) if (m := re.search(r'offset:(-?\d+)', op_str)) else 0
-    return getattr(autogen, mnemonic)(vdst=values[0], addr=values[1], saddr=values[2], offset=offset)
-  elif (mnemonic.startswith('flat_store') or mnemonic.startswith('global_store') or mnemonic.startswith('scratch_store')) and len(values) >= 3:
-    offset = int(m.group(1)) if (m := re.search(r'offset:(-?\d+)', op_str)) else 0
-    return getattr(autogen, mnemonic)(addr=values[0], data=values[1], saddr=values[2], offset=offset)
-  for suffix in (['_e32', ''] if not (neg_bits or abs_bits or clamp) else ['', '_e32']):
-    if hasattr(autogen, name := mnemonic.replace('.', '_') + suffix):
-      use_opsel = 'opsel' in getattr(autogen, name).func._fields
-      vals = [type(v)(v.idx, v.count, False) if isinstance(v, Reg) and v.hi and use_opsel else v for v in values]
-      inst = getattr(autogen, name)(*vals, literal=lit, **modifiers)
-      if neg_bits and 'neg' in inst._fields: inst._values['neg'] = neg_bits
-      if opsel_bits and use_opsel: inst._values['opsel'] = opsel_bits
-      if abs_bits and 'abs' in inst._fields: inst._values['abs'] = abs_bits
-      if clamp and 'clmp' in inst._fields: inst._values['clmp'] = 1
-      return inst
-  raise ValueError(f"unknown instruction: {mnemonic}")

extra/assembly/rdna3/emu.py

@@ -1,505 +0,0 @@
-# RDNA3 emulator - pure Python implementation for testing
-from __future__ import annotations
-import ctypes, struct, math
-from typing import Callable
-from extra.assembly.rdna3.lib import Inst, Inst32, Inst64, RawImm
-
-Program = dict[int, Inst]  # pc (word offset) -> instruction
-from extra.assembly.rdna3.autogen import (
-  SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, VOPD, SrcEnum,
-  SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, DSOp, FLATOp, GLOBALOp, VOPDOp
-)
-from extra.assembly.rdna3.alu import (
-  f32, i32, f16, i16, sext, vopc, SALU, VALU,
-  SOP1_BASE, SOP2_BASE, SOPC_BASE, SOPK_BASE, VOP1_BASE, VOP2_BASE
-)
-
-WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
-VCC_LO, VCC_HI, NULL, M0, EXEC_LO, EXEC_HI, SCC = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.M0, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.SCC
-# Pre-computed inline constant table for src operands 128-254 (index = src - 128)
-_INLINE_CONSTS = [0] * 127
-for _i in range(65): _INLINE_CONSTS[_i] = _i  # 128-192 -> 0-64
-for _i in range(1, 17): _INLINE_CONSTS[64 + _i] = ((-_i) & 0xffffffff)  # 193-208 -> -1 to -16
-for _k, _v in {SrcEnum.POS_HALF: 0x3f000000, SrcEnum.NEG_HALF: 0xbf000000, SrcEnum.POS_ONE: 0x3f800000, SrcEnum.NEG_ONE: 0xbf800000,
-               SrcEnum.POS_TWO: 0x40000000, SrcEnum.NEG_TWO: 0xc0000000, SrcEnum.POS_FOUR: 0x40800000, SrcEnum.NEG_FOUR: 0xc0800000,
-               SrcEnum.INV_2PI: 0x3e22f983}.items(): _INLINE_CONSTS[_k - 128] = _v
-
-_valid_mem_ranges: list[tuple[int, int]] = []
-def set_valid_mem_ranges(ranges: set[tuple[int, int]]) -> None: global _valid_mem_ranges; _valid_mem_ranges = list(ranges)
-def _mem_valid(addr: int, size: int) -> bool:
-  for s, z in _valid_mem_ranges:
-    if s <= addr and addr + size <= s + z: return True
-  return not _valid_mem_ranges
-def _ctypes_at(addr: int, size: int): return (ctypes.c_uint8 if size == 1 else ctypes.c_uint16 if size == 2 else ctypes.c_uint32).from_address(addr)
-def mem_read(addr: int, size: int) -> int: return _ctypes_at(addr, size).value if _mem_valid(addr, size) else 0
-def mem_write(addr: int, size: int, val: int) -> None:
-  if _mem_valid(addr, size): _ctypes_at(addr, size).value = val
-
-# Memory op tables - (cnt, sz, sign) for loads, (cnt, sz) for stores
-def _mem_ops(ops, suffix_map):
-  return {getattr(e, f"{p}_{s}"): v for e in ops for s, v in suffix_map.items() for p in [e.__name__.replace("Op", "")]}
-_LOAD_MAP = {'LOAD_B32': (1,4,0), 'LOAD_B64': (2,4,0), 'LOAD_B96': (3,4,0), 'LOAD_B128': (4,4,0), 'LOAD_U8': (1,1,0), 'LOAD_I8': (1,1,1), 'LOAD_U16': (1,2,0), 'LOAD_I16': (1,2,1)}
-_STORE_MAP = {'STORE_B32': (1,4), 'STORE_B64': (2,4), 'STORE_B96': (3,4), 'STORE_B128': (4,4), 'STORE_B8': (1,1), 'STORE_B16': (1,2)}
-FLAT_LOAD = _mem_ops([GLOBALOp, FLATOp], _LOAD_MAP)
-FLAT_STORE = _mem_ops([GLOBALOp, FLATOp], _STORE_MAP)
-DS_LOAD: dict[int, tuple[int,int,int]] = {DSOp.DS_LOAD_B32: (1,4,0), DSOp.DS_LOAD_B64: (2,4,0), DSOp.DS_LOAD_B128: (4,4,0), DSOp.DS_LOAD_U8: (1,1,0), DSOp.DS_LOAD_I8: (1,1,1), DSOp.DS_LOAD_U16: (1,2,0), DSOp.DS_LOAD_I16: (1,2,1)}
-DS_STORE: dict[int, tuple[int,int]] = {DSOp.DS_STORE_B32: (1,4), DSOp.DS_STORE_B64: (2,4), DSOp.DS_STORE_B128: (4,4), DSOp.DS_STORE_B8: (1,1), DSOp.DS_STORE_B16: (1,2)}
-FLAT_D16_LO = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('LOAD_D16_U8', (1, 0)), ('LOAD_D16_I8', (1, 1)), ('LOAD_D16_B16', (2, 0))]}
-FLAT_D16_HI = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('LOAD_D16_HI_U8', (1, 0)), ('LOAD_D16_HI_I8', (1, 1)), ('LOAD_D16_HI_B16', (2, 0))]}
-FLAT_D16_STORE = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('STORE_D16_HI_B8', 1), ('STORE_D16_HI_B16', 2)]}
-SMEM_LOAD: dict[int, int] = {SMEMOp.S_LOAD_B32: 1, SMEMOp.S_LOAD_B64: 2, SMEMOp.S_LOAD_B128: 4, SMEMOp.S_LOAD_B256: 8, SMEMOp.S_LOAD_B512: 16}
-SOPK_WAIT = {SOPKOp.S_WAITCNT_VSCNT, SOPKOp.S_WAITCNT_VMCNT, SOPKOp.S_WAITCNT_EXPCNT, SOPKOp.S_WAITCNT_LGKMCNT}
-
-class WaveState:
-  __slots__ = ('sgpr', 'vgpr', 'scc', 'pc', 'literal', '_pend_sgpr')
-  def __init__(self):
-    self.sgpr, self.vgpr = [0] * SGPR_COUNT, [[0] * VGPR_COUNT for _ in range(WAVE_SIZE)]
-    self.sgpr[EXEC_LO] = 0xffffffff  # wave32: all lanes active
-    self.scc = self.pc = self.literal = 0
-    self._pend_sgpr = {}
-
-  @property
-  def vcc(self) -> int: return self.sgpr[VCC_LO] | (self.sgpr[VCC_HI] << 32)
-  @vcc.setter
-  def vcc(self, v: int) -> None: self.sgpr[VCC_LO] = v & 0xffffffff; self.sgpr[VCC_HI] = (v >> 32) & 0xffffffff
-  @property
-  def exec_mask(self) -> int: return self.sgpr[EXEC_LO] | (self.sgpr[EXEC_HI] << 32)
-  @exec_mask.setter
-  def exec_mask(self, v: int) -> None: self.sgpr[EXEC_LO] = v & 0xffffffff; self.sgpr[EXEC_HI] = (v >> 32) & 0xffffffff
-
-  def rsgpr(self, i: int) -> int:
-    if i == NULL: return 0
-    if i == SCC: return self.scc
-    return self.sgpr[i] if i < SGPR_COUNT else 0
-  def wsgpr(self, i: int, v: int) -> None:
-    if i < SGPR_COUNT and i != NULL: self.sgpr[i] = v & 0xffffffff
-  def rsgpr64(self, i: int) -> int: return self.rsgpr(i) | (self.rsgpr(i+1) << 32)
-  def wsgpr64(self, i: int, v: int) -> None: self.wsgpr(i, v & 0xffffffff); self.wsgpr(i+1, (v >> 32) & 0xffffffff)
-
-  def rsrc(self, v: int, lane: int) -> int:
-    if v < SGPR_COUNT: return self.sgpr[v]
-    if v == SCC: return self.scc
-    if v < 255: return _INLINE_CONSTS[v - 128]
-    if v == 255: return self.literal
-    return self.vgpr[lane][v - 256] if v <= 511 else 0
-
-  def rsrc64(self, v: int, lane: int) -> int:
-    return self.rsrc(v, lane) | ((self.rsrc(v+1, lane) if v < VCC_LO or 256 <= v <= 511 else 0) << 32)
-
-  def pend_sgpr_lane(self, reg: int, lane: int, val: int) -> None:
-    if reg not in self._pend_sgpr: self._pend_sgpr[reg] = 0
-    if val: self._pend_sgpr[reg] |= (1 << lane)
-
-  def commit_pends(self) -> None:
-    for reg, val in self._pend_sgpr.items(): self.sgpr[reg] = val
-    self._pend_sgpr.clear()
-
-def decode_format(word: int) -> tuple[type[Inst] | None, bool]:
-  hi2 = (word >> 30) & 0x3
-  if hi2 == 0b11:
-    enc = (word >> 26) & 0xf
-    if enc == 0b1101: return SMEM, True
-    if enc == 0b0101:
-      op = (word >> 16) & 0x3ff
-      return (VOP3SD, True) if op in (288, 289, 290, 764, 765, 766, 767, 768, 769, 770) else (VOP3, True)
-    return {0b0011: (VOP3P, True), 0b0110: (DS, True), 0b0111: (FLAT, True), 0b0010: (VOPD, True)}.get(enc, (None, True))
-  if hi2 == 0b10:
-    enc = (word >> 23) & 0x7f
-    return {0b1111101: (SOP1, False), 0b1111110: (SOPC, False), 0b1111111: (SOPP, False)}.get(enc, (SOPK, False) if ((word >> 28) & 0xf) == 0b1011 else (SOP2, False))
-  enc = (word >> 25) & 0x7f
-  return (VOPC, False) if enc == 0b0111110 else (VOP1, False) if enc == 0b0111111 else (VOP2, False)
-
-def _unwrap(v) -> int: return v.val if isinstance(v, RawImm) else v.value if hasattr(v, 'value') else v
-
-def decode_program(data: bytes) -> Program:
-  result: Program = {}
-  i = 0
-  while i < len(data):
-    word = int.from_bytes(data[i:i+4], 'little')
-    inst_class, is_64 = decode_format(word)
-    if inst_class is None: i += 4; continue
-    base_size = 8 if is_64 else 4
-    inst = inst_class.from_bytes(data[i:i+base_size])
-    for name, val in inst._values.items(): setattr(inst, name, _unwrap(val))
-    has_literal = any(getattr(inst, fld, None) == 255 for fld in ('src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'srcx0', 'srcy0'))
-    if inst_class == VOP2 and inst.op in (44, 45, 55, 56): has_literal = True
-    if inst_class == VOPD and (inst.opx in (1, 2) or inst.opy in (1, 2)): has_literal = True
-    if inst_class == SOP2 and inst.op in (69, 70): has_literal = True
-    if has_literal: inst._literal = int.from_bytes(data[i+base_size:i+base_size+4], 'little')
-    inst._words = inst.size() // 4  # cache size for step_wave
-    result[i // 4] = inst
-    i += inst._words * 4
-  return result
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# SCALAR EXECUTION
-# ═══════════════════════════════════════════════════════════════════════════════
-def exec_sop1(st: WaveState, inst: SOP1) -> int:
-  s0, op = st.rsrc(inst.ssrc0, 0), inst.op
-  # 64-bit and special ops handled inline
-  if op == SOP1Op.S_MOV_B64: st.wsgpr64(inst.sdst, st.rsrc64(inst.ssrc0, 0)); return 0
-  if op == SOP1Op.S_NOT_B64: r = (~st.rsrc64(inst.ssrc0, 0)) & 0xffffffffffffffff; st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
-  if op == SOP1Op.S_BITSET0_B32: st.wsgpr(inst.sdst, st.rsgpr(inst.sdst) & ~(1 << (s0 & 0x1f))); return 0
-  if op == SOP1Op.S_BITSET1_B32: st.wsgpr(inst.sdst, st.rsgpr(inst.sdst) | (1 << (s0 & 0x1f))); return 0
-  if op == SOP1Op.S_AND_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 & old; st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
-  if op == SOP1Op.S_OR_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 | old; st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
-  if op == SOP1Op.S_AND_NOT1_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 & (~old & 0xffffffff); st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
-  if op == SOP1Op.S_GETPC_B64: return -3
-  if op == SOP1Op.S_SETPC_B64: return -4
-  if op == SOP1Op.S_SWAPPC_B64: return -5
-  if (fn := SALU.get(SOP1_BASE + op)) is None: raise NotImplementedError(f"SOP1 op {op}")
-  r, scc = fn(s0, 0, st.scc); st.wsgpr(inst.sdst, r); st.scc = scc; return 0
-
-_SOP2_64: dict[int, Callable[[int, int], int]] = {SOP2Op.S_AND_B64: lambda a, b: a & b, SOP2Op.S_OR_B64: lambda a, b: a | b, SOP2Op.S_XOR_B64: lambda a, b: a ^ b}
-def exec_sop2(st: WaveState, inst: SOP2) -> int:
-  s0, s1, op = st.rsrc(inst.ssrc0, 0), st.rsrc(inst.ssrc1, 0), inst.op
-  # 64-bit ops handled inline
-  if op == SOP2Op.S_LSHL_B64: r = (st.rsrc64(inst.ssrc0, 0) << (s1 & 0x3f)) & 0xffffffffffffffff; st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
-  if op == SOP2Op.S_LSHR_B64: r = st.rsrc64(inst.ssrc0, 0) >> (s1 & 0x3f); st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
-  if op == SOP2Op.S_ASHR_I64: r = sext(st.rsrc64(inst.ssrc0, 0), 64) >> (s1 & 0x3f); st.wsgpr64(inst.sdst, r & 0xffffffffffffffff); st.scc = int(r != 0); return 0
-  if (fn := _SOP2_64.get(op)): r = fn(st.rsrc64(inst.ssrc0, 0), st.rsrc64(inst.ssrc1, 0)); st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
-  if op == SOP2Op.S_CSELECT_B64: st.wsgpr64(inst.sdst, st.rsrc64(inst.ssrc0, 0) if st.scc else st.rsrc64(inst.ssrc1, 0)); return 0
-  if op == SOP2Op.S_FMAC_F32: st.wsgpr(inst.sdst, i32(f32(st.rsgpr(inst.sdst)) + f32(s0) * f32(s1))); return 0
-  if op == SOP2Op.S_FMAAK_F32: st.wsgpr(inst.sdst, i32(f32(s0) * f32(s1) + f32(inst._literal or 0))); return 0
-  if op == SOP2Op.S_FMAMK_F32: st.wsgpr(inst.sdst, i32(f32(s0) * f32(inst._literal or 0) + f32(s1))); return 0
-  if (fn := SALU.get(SOP2_BASE + op)) is None: raise NotImplementedError(f"SOP2 op {op}")
-  r, scc = fn(s0, s1, st.scc); st.wsgpr(inst.sdst, r); st.scc = scc; return 0
-
-def exec_sopc(st: WaveState, inst: SOPC) -> int:
-  s0, s1, op = st.rsrc(inst.ssrc0, 0), st.rsrc(inst.ssrc1, 0), inst.op
-  if op == SOPCOp.S_CMP_EQ_U64: st.scc = int(st.rsrc64(inst.ssrc0, 0) == st.rsrc64(inst.ssrc1, 0)); return 0
-  if op == SOPCOp.S_CMP_LG_U64: st.scc = int(st.rsrc64(inst.ssrc0, 0) != st.rsrc64(inst.ssrc1, 0)); return 0
-  if (fn := SALU.get(SOPC_BASE + op)) is None: raise NotImplementedError(f"SOPC op {op}")
-  st.scc = fn(s0, s1, st.scc)[1]; return 0
-
-_SOPK_CMP = frozenset((SOPKOp.S_CMPK_EQ_I32, SOPKOp.S_CMPK_LG_I32, SOPKOp.S_CMPK_GT_I32, SOPKOp.S_CMPK_GE_I32,
-  SOPKOp.S_CMPK_LT_I32, SOPKOp.S_CMPK_LE_I32, SOPKOp.S_CMPK_EQ_U32, SOPKOp.S_CMPK_LG_U32,
-  SOPKOp.S_CMPK_GT_U32, SOPKOp.S_CMPK_GE_U32, SOPKOp.S_CMPK_LT_U32, SOPKOp.S_CMPK_LE_U32))
-def exec_sopk(st: WaveState, inst: SOPK) -> int:
-  simm, s0, op = inst.simm16, st.rsgpr(inst.sdst), inst.op
-  if op in SOPK_WAIT: return 0
-  if (fn := SALU.get(SOPK_BASE + op)) is None: raise NotImplementedError(f"SOPK op {op}")
-  r, scc = fn(s0, simm, st.scc)
-  if op not in _SOPK_CMP: st.wsgpr(inst.sdst, r)
-  st.scc = scc; return 0
-
-def exec_sopp(st: WaveState, inst: SOPP) -> int:
-  if inst.op == SOPPOp.S_ENDPGM: return -1
-  if inst.op == SOPPOp.S_BARRIER: return -2
-  if inst.op == SOPPOp.S_BRANCH: return sext(inst.simm16, 16)
-  if inst.op == SOPPOp.S_CBRANCH_SCC0: return sext(inst.simm16, 16) if st.scc == 0 else 0
-  if inst.op == SOPPOp.S_CBRANCH_SCC1: return sext(inst.simm16, 16) if st.scc == 1 else 0
-  # In wave32 mode, only VCC_LO is used for lane masks; VCC_HI is a free SGPR
-  if inst.op == SOPPOp.S_CBRANCH_VCCZ: return sext(inst.simm16, 16) if (st.vcc & 0xffffffff) == 0 else 0
-  if inst.op == SOPPOp.S_CBRANCH_VCCNZ: return sext(inst.simm16, 16) if (st.vcc & 0xffffffff) != 0 else 0
-  if inst.op == SOPPOp.S_CBRANCH_EXECZ: return sext(inst.simm16, 16) if st.exec_mask == 0 else 0
-  if inst.op == SOPPOp.S_CBRANCH_EXECNZ: return sext(inst.simm16, 16) if st.exec_mask != 0 else 0
-  # Scheduling hints and wait instructions are no-ops in emulation
-  if inst.op <= 31: return 0  # S_NOP, S_CLAUSE, S_DELAY_ALU, S_WAITCNT, etc.
-  # S_WAKEUP(52), S_SETPRIO(53), S_SENDMSG(54), S_SENDMSGHALT(55), perf counters, S_ICACHE_INV(60) are no-ops
-  if inst.op in (52, 53, 54, 55, 56, 57, 60): return 0
-  raise NotImplementedError(f"SOPP op {inst.op}")
-
-def exec_smem(st: WaveState, inst: SMEM) -> int:
-  addr = st.rsgpr64(inst.sbase * 2) + sext(inst.offset, 21)
-  if inst.soffset not in (NULL, 0x7f): addr += st.rsrc(inst.soffset, 0)
-  if (cnt := SMEM_LOAD.get(inst.op)) is None: raise NotImplementedError(f"SMEM op {inst.op}")
-  for i in range(cnt): st.wsgpr(inst.sdata + i, mem_read((addr + i * 4) & 0xffffffffffffffff, 4))
-  return 0
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# VECTOR EXECUTION
-# ═══════════════════════════════════════════════════════════════════════════════
-def f64(hi: int, lo: int) -> float: return struct.unpack('<d', struct.pack('<Q', (hi << 32) | lo))[0]
-def i64_parts(f: float) -> tuple[int, int]:
-  if math.isnan(f): val = 0x7ff8000000000000
-  elif math.isinf(f): val = 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
-  else: val = struct.unpack('<Q', struct.pack('<d', f))[0]
-  return val & 0xffffffff, (val >> 32) & 0xffffffff
-
-def exec_vop1(st: WaveState, inst: VOP1, lane: int) -> None:
-  if inst.op == VOP1Op.V_NOP: return
-  V, s0 = st.vgpr[lane], st.rsrc(inst.src0, lane)
-  if inst.op == VOP1Op.V_READFIRSTLANE_B32:
-    first = (st.exec_mask & -st.exec_mask).bit_length() - 1 if st.exec_mask else 0
-    st.wsgpr(inst.vdst, st.rsrc(inst.src0, first) if inst.src0 >= 256 else s0); return
-  # F64 ops handled inline
-  if inst.op == VOP1Op.V_CVT_F64_F32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(f32(s0))); return
-  if inst.op == VOP1Op.V_CVT_F64_I32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(sext(s0, 32))); return
-  if inst.op == VOP1Op.V_CVT_F64_U32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(s0)); return
-  if inst.op in (VOP1Op.V_CVT_F32_F64, VOP1Op.V_CVT_I32_F64, VOP1Op.V_CVT_U32_F64):
-    src = inst.src0 - 256 if inst.src0 >= 256 else inst.src0
-    lo, hi = (V[src], V[src+1]) if inst.src0 >= 256 else (st.sgpr[src], st.sgpr[src+1])
-    v = f64(hi, lo)
-    if inst.op == VOP1Op.V_CVT_F32_F64: V[inst.vdst] = i32(v)
-    elif inst.op == VOP1Op.V_CVT_I32_F64: V[inst.vdst] = (max(-0x80000000, min(0x7fffffff, int(v))) & 0xffffffff) if math.isfinite(v) else 0
-    else: V[inst.vdst] = max(0, min(0xffffffff, int(v))) if math.isfinite(v) and v == v else 0
-    return
-  if (fn := VALU.get(VOP1_BASE + inst.op)): V[inst.vdst] = fn(s0, 0, 0); return
-  raise NotImplementedError(f"VOP1 op {inst.op}")
-
-def exec_vop2(st: WaveState, inst: VOP2, lane: int) -> None:
-  V, s0, s1, op = st.vgpr[lane], st.rsrc(inst.src0, lane), st.vgpr[lane][inst.vsrc1], inst.op
-  if op == VOP2Op.V_CNDMASK_B32: V[inst.vdst] = s1 if (st.vcc >> lane) & 1 else s0; return
-  if op == VOP2Op.V_FMAC_F32: V[inst.vdst] = i32(f32(s0)*f32(s1)+f32(V[inst.vdst])); return
-  if op == VOP2Op.V_FMAMK_F32: V[inst.vdst] = i32(f32(s0)*f32(st.literal)+f32(s1)); return
-  if op == VOP2Op.V_FMAAK_F32: V[inst.vdst] = i32(f32(s0)*f32(s1)+f32(st.literal)); return
-  if op == VOP2Op.V_FMAC_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(s1)+f16(V[inst.vdst])); return
-  if op == VOP2Op.V_FMAMK_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(st.literal)+f16(s1)); return
-  if op == VOP2Op.V_FMAAK_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(s1)+f16(st.literal)); return
-  if op == VOP2Op.V_PK_FMAC_F16:
-    lo = i16(f16(s0 & 0xffff) * f16(s1 & 0xffff) + f16(V[inst.vdst] & 0xffff))
-    hi = i16(f16((s0 >> 16) & 0xffff) * f16((s1 >> 16) & 0xffff) + f16((V[inst.vdst] >> 16) & 0xffff))
-    V[inst.vdst] = lo | (hi << 16); return
-  if op == VOP2Op.V_ADD_CO_CI_U32: r = s0+s1+((st.vcc>>lane)&1); st.pend_sgpr_lane(VCC_LO, lane, r >= 0x100000000); V[inst.vdst] = r & 0xffffffff; return
-  if op == VOP2Op.V_SUB_CO_CI_U32: b = (st.vcc>>lane)&1; st.pend_sgpr_lane(VCC_LO, lane, s1+b > s0); V[inst.vdst] = (s0-s1-b) & 0xffffffff; return
-  if (fn := VALU.get(VOP2_BASE + op)): V[inst.vdst] = fn(s0, s1, 0); return
-  raise NotImplementedError(f"VOP2 op {op}")
-
-def vop3_mod(val: int, neg: int, abs_: int, idx: int) -> int:
-  if (abs_ >> idx) & 1: val = i32(abs(f32(val)))
-  if (neg >> idx) & 1: val = i32(-f32(val))
-  return val
-
-def exec_vop3(st: WaveState, inst: VOP3, lane: int) -> None:
-  op, src0, src1, src2, vdst, neg, abs_ = inst.op, inst.src0, inst.src1, inst.src2, inst.vdst, inst.neg, getattr(inst, 'abs', 0)
-  V = st.vgpr[lane]
-  # VOPC encoded in VOP3 (0-255)
-  if 0 <= op <= 255:
-    base = op & 0x7f
-    # For 64-bit comparisons (I64: 80-87, U64: 88-95), read raw 64-bit values (no float modifiers)
-    if 80 <= base <= 95:
-      s0_64, s1_64 = st.rsrc64(src0, lane), st.rsrc64(src1, lane)
-      result = vopc(op, s0_64 & 0xffffffff, s1_64 & 0xffffffff, (s0_64 >> 32) & 0xffffffff, (s1_64 >> 32) & 0xffffffff)
-    else:
-      s0, s1 = vop3_mod(st.rsrc(src0, lane), neg, abs_, 0), vop3_mod(st.rsrc(src1, lane), neg, abs_, 1)
-      result = vopc(op, s0, s1)
-    is_cmpx = op >= 128
-    st.pend_sgpr_lane(vdst, lane, result)
-    if is_cmpx: st.pend_sgpr_lane(EXEC_LO, lane, result)
-    return
-  s0, s1, s2 = vop3_mod(st.rsrc(src0, lane), neg, abs_, 0), vop3_mod(st.rsrc(src1, lane), neg, abs_, 1), vop3_mod(st.rsrc(src2, lane), neg, abs_, 2)
-  # Special ops
-  if op == VOP3Op.V_FMAC_F32: V[vdst] = i32(f32(s0)*f32(s1)+f32(V[vdst])); return
-  if op == VOP3Op.V_READLANE_B32: st.wsgpr(vdst, st.vgpr[s1 & 0x1f][src0 - 256] if src0 >= 256 else s0); return
-  if op == VOP3Op.V_WRITELANE_B32: st.vgpr[s1 & 0x1f][vdst] = s0; return
-  if op == VOP3Op.V_CNDMASK_B32:
-    mask = st.rsgpr(src2) if src2 < 256 else st.vcc
-    V[vdst] = s1 if (mask >> lane) & 1 else s0; return
-  if op in (VOP3Op.V_LSHLREV_B64, VOP3Op.V_LSHRREV_B64, VOP3Op.V_ASHRREV_I64):
-    v64 = st.rsrc64(src1, lane)
-    r = ((v64 << (s0 & 0x3f)) & 0xffffffffffffffff if op == VOP3Op.V_LSHLREV_B64 else
-         v64 >> (s0 & 0x3f) if op == VOP3Op.V_LSHRREV_B64 else sext(v64, 64) >> (s0 & 0x3f))
-    V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff; return
-  if op in (VOP3Op.V_ADD_F64, VOP3Op.V_MUL_F64, VOP3Op.V_FMA_F64, VOP3Op.V_MAX_F64, VOP3Op.V_MIN_F64):
-    a, b = f64(st.rsrc(src0+1, lane), s0), f64(st.rsrc(src1+1, lane), s1)
-    c = f64(st.rsrc(src2+1, lane), s2) if op == VOP3Op.V_FMA_F64 else 0.0
-    rf = a + b if op == VOP3Op.V_ADD_F64 else a * b if op == VOP3Op.V_MUL_F64 else a * b + c if op == VOP3Op.V_FMA_F64 else max(a, b) if op == VOP3Op.V_MAX_F64 else min(a, b)
-    V[vdst], V[vdst+1] = i64_parts(rf); return
-  if (fn := VALU.get(op)): V[vdst] = fn(s0, s1, s2); return
-  raise NotImplementedError(f"VOP3 op {op}")
-
-def exec_vopc(st: WaveState, inst: VOPC, lane: int) -> None:
-  result, is_cmpx = vopc(inst.op, st.rsrc(inst.src0, lane), st.vgpr[lane][inst.vsrc1]), inst.op >= 128
-  st.pend_sgpr_lane(EXEC_LO if is_cmpx else VCC_LO, lane, result)
-
-def exec_vop3sd(st: WaveState, inst: VOP3SD, lane: int) -> None:
-  op, src0, src1, src2, vdst, sdst, neg = inst.op, inst.src0, inst.src1, inst.src2, inst.vdst, inst.sdst, inst.neg
-  s0, s1, s2 = st.rsrc(src0, lane), st.rsrc(src1, lane), st.rsrc(src2, lane)
-  if (neg >> 0) & 1: s0 = i32(-f32(s0))
-  if (neg >> 1) & 1: s1 = i32(-f32(s1))
-  if (neg >> 2) & 1: s2 = i32(-f32(s2))
-  V = st.vgpr[lane]
-  if op == VOP3SDOp.V_ADD_CO_U32: r = s0 + s1; V[vdst] = r & 0xffffffff; st.pend_sgpr_lane(sdst, lane, r >= 0x100000000)
-  elif op == VOP3SDOp.V_SUB_CO_U32: V[vdst] = (s0 - s1) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s1 > s0)
-  elif op == VOP3SDOp.V_SUBREV_CO_U32: V[vdst] = (s1 - s0) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s0 > s1)
-  elif op == VOP3SDOp.V_ADD_CO_CI_U32:
-    cin = (st.rsgpr(src2) >> lane) & 1 if src2 < 256 else (st.vcc >> lane) & 1
-    r = s0 + s1 + cin; V[vdst] = r & 0xffffffff; st.pend_sgpr_lane(sdst, lane, r >= 0x100000000)
-  elif op == VOP3SDOp.V_SUB_CO_CI_U32:
-    cin = (st.rsgpr(src2) >> lane) & 1 if src2 < 256 else (st.vcc >> lane) & 1
-    V[vdst] = (s0 - s1 - cin) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s1 + cin > s0)
-  elif op == VOP3SDOp.V_MAD_U64_U32:
-    s2_64 = s2 | (st.rsrc(src2+1, lane) << 32); r = s0 * s1 + s2_64
-    V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff
-  elif op == VOP3SDOp.V_MAD_I64_I32:
-    s2_64 = sext(s2 | (st.rsrc(src2+1, lane) << 32), 64)
-    r = (sext(s0, 32) * sext(s1, 32) + s2_64) & 0xffffffffffffffff
-    V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff
-  elif op == VOP3SDOp.V_DIV_SCALE_F32: V[vdst] = 0; st.pend_sgpr_lane(sdst, lane, False)
-  elif op == VOP3SDOp.V_DIV_SCALE_F64: V[vdst], V[vdst+1] = s0, st.rsrc(src0+1, lane); st.pend_sgpr_lane(VCC_LO, lane, s0 == s2)
-  else: raise NotImplementedError(f"VOP3SD op {op}")
-
-def exec_flat(st: WaveState, inst: FLAT, lane: int) -> None:
-  op, addr_reg, data_reg, vdst, offset, saddr, V = inst.op, inst.addr, inst.data, inst.vdst, sext(inst.offset, 13), inst.saddr, st.vgpr[lane]
-  addr = V[addr_reg] | (V[addr_reg+1] << 32)
-  addr = (st.rsgpr64(saddr) + V[addr_reg] + offset) & 0xffffffffffffffff if saddr not in (NULL, 0x7f) else (addr + offset) & 0xffffffffffffffff
-  if op in FLAT_LOAD:
-    cnt, sz, sign = FLAT_LOAD[op]
-    for i in range(cnt): val = mem_read(addr + i * sz, sz); V[vdst + i] = sext(val, sz * 8) & 0xffffffff if sign else val
-  elif op in FLAT_STORE:
-    cnt, sz = FLAT_STORE[op]
-    for i in range(cnt): mem_write(addr + i * sz, sz, V[data_reg + i] & ((1 << (sz * 8)) - 1))
-  elif op in FLAT_D16_LO: sz, sign = FLAT_D16_LO[op]; val = mem_read(addr, sz); V[vdst] = (V[vdst] & 0xffff0000) | ((sext(val, sz * 8) & 0xffff) if sign else (val & 0xffff))
-  elif op in FLAT_D16_HI: sz, sign = FLAT_D16_HI[op]; val = mem_read(addr, sz); V[vdst] = (V[vdst] & 0x0000ffff) | (((sext(val, sz * 8) & 0xffff) if sign else (val & 0xffff)) << 16)
-  elif op in FLAT_D16_STORE: mem_write(addr, FLAT_D16_STORE[op], (V[data_reg] >> 16) & ((1 << (FLAT_D16_STORE[op] * 8)) - 1))
-  else: raise NotImplementedError(f"FLAT op {op}")
-
-def exec_ds(st: WaveState, inst: DS, lane: int, lds: bytearray) -> None:
-  op, addr, vdst, V = inst.op, (st.vgpr[lane][inst.addr] + inst.offset0) & 0xffff, inst.vdst, st.vgpr[lane]
-  if op in DS_LOAD:
-    cnt, sz, sign = DS_LOAD[op]
-    for i in range(cnt): val = int.from_bytes(lds[addr+i*sz:addr+i*sz+sz], 'little'); V[vdst + i] = sext(val, sz * 8) & 0xffffffff if sign else val
-  elif op in DS_STORE:
-    cnt, sz = DS_STORE[op]
-    for i in range(cnt): lds[addr+i*sz:addr+i*sz+sz] = (V[inst.data0 + i] & ((1 << (sz * 8)) - 1)).to_bytes(sz, 'little')
-  else: raise NotImplementedError(f"DS op {op}")
-
-VOPD_OPS: dict[int, Callable[[int, int, int, int, int], int]] = {
-  VOPDOp.V_DUAL_MUL_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)), VOPDOp.V_DUAL_ADD_F32: lambda a, b, d, l, lit: i32(f32(a)+f32(b)),
-  VOPDOp.V_DUAL_SUB_F32: lambda a, b, d, l, lit: i32(f32(a)-f32(b)), VOPDOp.V_DUAL_SUBREV_F32: lambda a, b, d, l, lit: i32(f32(b)-f32(a)),
-  VOPDOp.V_DUAL_MAX_F32: lambda a, b, d, l, lit: i32(max(f32(a), f32(b))), VOPDOp.V_DUAL_MIN_F32: lambda a, b, d, l, lit: i32(min(f32(a), f32(b))),
-  VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: lambda a, b, d, l, lit: i32(0.0 if f32(a) == 0.0 or f32(b) == 0.0 else f32(a)*f32(b)),
-  VOPDOp.V_DUAL_MOV_B32: lambda a, b, d, l, lit: a, VOPDOp.V_DUAL_ADD_NC_U32: lambda a, b, d, l, lit: (a + b) & 0xffffffff,
-  VOPDOp.V_DUAL_LSHLREV_B32: lambda a, b, d, l, lit: (b << (a & 0x1f)) & 0xffffffff, VOPDOp.V_DUAL_AND_B32: lambda a, b, d, l, lit: a & b,
-  VOPDOp.V_DUAL_FMAC_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)+f32(d)), VOPDOp.V_DUAL_FMAAK_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)+f32(lit)),
-  VOPDOp.V_DUAL_FMAMK_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(lit)+f32(b)), VOPDOp.V_DUAL_CNDMASK_B32: lambda a, b, d, l, lit: b if l else a,
-}
-def exec_vopd(st: WaveState, inst: VOPD, lane: int) -> None:
-  V, vdsty, vcc_lane = st.vgpr[lane], (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), (st.vcc >> lane) & 1
-  sx0, sx1, sy0, sy1, dstx = st.rsrc(inst.srcx0, lane), V[inst.vsrcx1], st.rsrc(inst.srcy0, lane), V[inst.vsrcy1], inst.vdstx
-  if (fn := VOPD_OPS.get(inst.opx)): V[dstx] = fn(sx0, sx1, V[dstx], vcc_lane, st.literal)
-  else: raise NotImplementedError(f"VOPD opx {inst.opx}")
-  if (fn := VOPD_OPS.get(inst.opy)): V[vdsty] = fn(sy0, sy1, V[vdsty], vcc_lane, st.literal)
-  else: raise NotImplementedError(f"VOPD opy {inst.opy}")
-
-def exec_vop3p(st: WaveState, inst: VOP3P, lane: int) -> None:
-  op, vdst, V = inst.op, inst.vdst, st.vgpr[lane]
-  s0, s1, s2 = st.rsrc(inst.src0, lane), st.rsrc(inst.src1, lane), st.rsrc(inst.src2, lane)
-  opsel, opsel_hi = [(inst.opsel >> i) & 1 for i in range(3)], [(inst.opsel_hi >> i) & 1 for i in range(2)] + [inst.opsel_hi2]
-  neg, neg_hi = inst.neg, inst.neg_hi
-  def get_src(src: int, idx: int, for_mix: bool = False) -> float:
-    if for_mix:
-      if not opsel_hi[idx]: return abs(f32(src)) if (neg_hi >> idx) & 1 else f32(src)
-      return float(f16((src >> 16) & 0xffff) if opsel[idx] else f16(src & 0xffff))
-    use_hi = opsel[idx]
-    val = ((src >> 16) & 0xffff) if use_hi else (src & 0xffff)
-    f = f16(val)
-    if use_hi and (neg >> idx) & 1: f = -f
-    elif not use_hi and (neg_hi >> idx) & 1: f = -f
-    return f
-  if op == VOP3POp.V_FMA_MIX_F32: V[vdst] = i32(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True))
-  elif op == VOP3POp.V_FMA_MIXLO_F16: V[vdst] = (V[vdst] & 0xffff0000) | i16(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True))
-  elif op == VOP3POp.V_FMA_MIXHI_F16: V[vdst] = (V[vdst] & 0x0000ffff) | (i16(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True)) << 16)
-  else: raise NotImplementedError(f"VOP3P op {op}")
-
-def exec_wmma_f32_16x16x16_f16(st: WaveState, inst: VOP3P, n_lanes: int) -> None:
-  src0_base, src1_base, src2_base = (inst.src0 - 256) if inst.src0 >= 256 else inst.src0, (inst.src1 - 256) if inst.src1 >= 256 else inst.src1, (inst.src2 - 256) if inst.src2 >= 256 else inst.src2
-  src0_is_vgpr, src1_is_vgpr, src2_is_vgpr, vdst = inst.src0 >= 256, inst.src1 >= 256, inst.src2 >= 256, inst.vdst
-  A, B, C = [[0.0] * 16 for _ in range(16)], [[0.0] * 16 for _ in range(16)], [[0.0] * 16 for _ in range(16)]
-  for lane in range(min(n_lanes, 16)):
-    V = st.vgpr[lane]
-    for reg in range(8):
-      val = V[src0_base + reg] if src0_is_vgpr else st.sgpr[src0_base + reg]
-      A[lane][reg * 2], A[lane][reg * 2 + 1] = f16(val & 0xffff), f16((val >> 16) & 0xffff)
-      val = V[src1_base + reg] if src1_is_vgpr else st.sgpr[src1_base + reg]
-      B[reg * 2][lane], B[reg * 2 + 1][lane] = f16(val & 0xffff), f16((val >> 16) & 0xffff)
-  for row in range(16):
-    for col in range(16):
-      idx, lane_idx, reg = row * 16 + col, (row * 16 + col) % 32, (row * 16 + col) // 32
-      if lane_idx < n_lanes:
-        val = st.vgpr[lane_idx][src2_base + reg] if src2_is_vgpr else st.sgpr[src2_base + reg]
-        C[row][col] = f32(val)
-  for row in range(16):
-    for col in range(16):
-      for k in range(16): C[row][col] += A[row][k] * B[k][col]
-  for row in range(16):
-    for col in range(16):
-      idx, lane_idx, reg = row * 16 + col, (row * 16 + col) % 32, (row * 16 + col) // 32
-      if lane_idx < n_lanes and (st.exec_mask & (1 << lane_idx)): st.vgpr[lane_idx][vdst + reg] = i32(C[row][col])
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# MAIN EXECUTION LOOP
-# ═══════════════════════════════════════════════════════════════════════════════
-SCALAR: dict[type, Callable[..., int]] = {SOP1: exec_sop1, SOP2: exec_sop2, SOPC: exec_sopc, SOPK: exec_sopk, SOPP: exec_sopp, SMEM: exec_smem}
-VECTOR: dict[type, Callable[..., None]] = {VOP1: exec_vop1, VOP2: exec_vop2, VOP3: exec_vop3, VOP3SD: exec_vop3sd, VOPC: exec_vopc, FLAT: exec_flat, DS: exec_ds, VOPD: exec_vopd, VOP3P: exec_vop3p}
-
-_WMMA_OPS = frozenset((VOP3POp.V_WMMA_F32_16X16X16_F16, VOP3POp.V_WMMA_F32_16X16X16_BF16, VOP3POp.V_WMMA_F16_16X16X16_F16,
-                       VOP3POp.V_WMMA_BF16_16X16X16_BF16, VOP3POp.V_WMMA_I32_16X16X16_IU8, VOP3POp.V_WMMA_I32_16X16X16_IU4))
-
-def step_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int) -> int:
-  inst = program.get(st.pc)
-  if inst is None: return 1
-  inst_words, st.literal, inst_type = inst._words, inst._literal or 0, type(inst)
-  if (handler := SCALAR.get(inst_type)) is not None:
-    delta = handler(st, inst)
-    if delta == -1: return -1
-    if delta == -2: st.pc += inst_words; return -2
-    if delta == -3:  # S_GETPC_B64
-      sop1 = inst if isinstance(inst, SOP1) else None
-      assert sop1 is not None
-      next_pc = (st.pc + inst_words) * 4; st.wsgpr(sop1.sdst, next_pc & 0xffffffff); st.wsgpr(sop1.sdst + 1, (next_pc >> 32) & 0xffffffff); st.pc += inst_words; return 0
-    if delta == -4:  # S_SETPC_B64
-      sop1 = inst if isinstance(inst, SOP1) else None
-      assert sop1 is not None
-      st.pc = st.rsrc64(sop1.ssrc0, 0) // 4; return 0
-    if delta == -5:  # S_SWAPPC_B64
-      sop1 = inst if isinstance(inst, SOP1) else None
-      assert sop1 is not None
-      next_pc = (st.pc + inst_words) * 4; st.wsgpr(sop1.sdst, next_pc & 0xffffffff); st.wsgpr(sop1.sdst + 1, (next_pc >> 32) & 0xffffffff); st.pc = st.rsrc64(sop1.ssrc0, 0) // 4; return 0
-    st.pc += inst_words + delta
-  else:
-    vec_handler, exec_mask = VECTOR[inst_type], st.exec_mask
-    if inst_type is DS:
-      for lane in range(n_lanes):
-        if exec_mask & (1 << lane): vec_handler(st, inst, lane, lds)
-    elif inst_type is VOP3P:
-      vop3p = inst if isinstance(inst, VOP3P) else None
-      assert vop3p is not None
-      if vop3p.op in _WMMA_OPS:
-        exec_wmma_f32_16x16x16_f16(st, vop3p, n_lanes)
-      else:
-        for lane in range(n_lanes):
-          if exec_mask & (1 << lane): vec_handler(st, vop3p, lane)
-    else:
-      for lane in range(n_lanes):
-        if exec_mask & (1 << lane): vec_handler(st, inst, lane)
-    st.commit_pends(); st.pc += inst_words
-  return 0
-
-def exec_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int, wg_id: tuple[int,int,int]=(0,0,0), local_size: tuple[int,int,int]=(1,1,1), wave_start: int=0) -> int:
-  while st.pc in program:
-    result = step_wave(program, st, lds, n_lanes)
-    if result == -1: return 0
-    if result == -2: return -2
-  return 0
-
-def exec_workgroup(program: Program, workgroup_id: tuple[int, int, int], local_size: tuple[int, int, int], args_ptr: int, dispatch_dim: int) -> None:
-  lx, ly, lz = local_size
-  total_threads, lds = lx * ly * lz, bytearray(65536)
-  waves: list[tuple[WaveState, int, int]] = []
-  for wave_start in range(0, total_threads, WAVE_SIZE):
-    n_lanes, st = min(WAVE_SIZE, total_threads - wave_start), WaveState()
-    st.exec_mask = (1 << n_lanes) - 1
-    st.wsgpr64(0, args_ptr)
-    gx, gy, gz = workgroup_id
-    if dispatch_dim >= 3: st.sgpr[13], st.sgpr[14], st.sgpr[15] = gx, gy, gz
-    elif dispatch_dim == 2: st.sgpr[14], st.sgpr[15] = gx, gy
-    else: st.sgpr[15] = gx
-    for i in range(n_lanes):
-      tid = wave_start + i
-      st.vgpr[i][0] = tid if local_size == (lx, 1, 1) else ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
-    waves.append((st, n_lanes, wave_start))
-  has_barrier = any(isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER for inst in program.values())
-  for _ in range(2 if has_barrier else 1):
-    for st, n_lanes, wave_start in waves: exec_wave(program, st, lds, n_lanes, workgroup_id, local_size, wave_start)
-
-def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int) -> int:
-  data = (ctypes.c_char * lib_sz).from_address(lib).raw
-  program = decode_program(data)
-  if not program: return -1
-  dispatch_dim = 3 if gz > 1 else (2 if gy > 1 else 1)
-  for gidz in range(gz):
-    for gidy in range(gy):
-      for gidx in range(gx): exec_workgroup(program, (gidx, gidy, gidz), (lx, ly, lz), args_ptr, dispatch_dim)
-  return 0

extra/assembly/rdna3/gen.py

@@ -1,191 +0,0 @@
-#!/usr/bin/env python3
-# generates autogen/__init__.py by parsing the AMD RDNA3.5 ISA PDF
-import re, pdfplumber, pathlib
-from tinygrad.helpers import fetch
-
-PDF_URL = "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content"
-FIELD_TYPES = {'SSRC0': 'SSrc', 'SSRC1': 'SSrc', 'SOFFSET': 'SSrc', 'SADDR': 'SSrc', 'SRC0': 'Src', 'SRC1': 'Src', 'SRC2': 'Src',
-  'SDST': 'SGPRField', 'SBASE': 'SGPRField', 'SDATA': 'SGPRField', 'SRSRC': 'SGPRField', 'VDST': 'VGPRField', 'VSRC1': 'VGPRField', 'VDATA': 'VGPRField',
-  'VADDR': 'VGPRField', 'ADDR': 'VGPRField', 'DATA': 'VGPRField', 'DATA0': 'VGPRField', 'DATA1': 'VGPRField', 'SIMM16': 'SImm', 'OFFSET': 'Imm',
-  'OPX': 'VOPDOp', 'OPY': 'VOPDOp', 'SRCX0': 'Src', 'SRCY0': 'Src', 'VSRCX1': 'VGPRField', 'VSRCY1': 'VGPRField', 'VDSTX': 'VGPRField', 'VDSTY': 'VDSTYEnc'}
-FIELD_ORDER = {
-  'SOP2': ['op', 'sdst', 'ssrc0', 'ssrc1'], 'SOP1': ['op', 'sdst', 'ssrc0'], 'SOPC': ['op', 'ssrc0', 'ssrc1'],
-  'SOPK': ['op', 'sdst', 'simm16'], 'SOPP': ['op', 'simm16'], 'VOP1': ['op', 'vdst', 'src0'], 'VOPC': ['op', 'src0', 'vsrc1'],
-  'VOP2': ['op', 'vdst', 'src0', 'vsrc1'], 'VOP3SD': ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2', 'clmp'],
-  'SMEM': ['op', 'sdata', 'sbase', 'soffset', 'offset', 'glc', 'dlc'], 'DS': ['op', 'vdst', 'addr', 'data0', 'data1'],
-  'VOP3': ['op', 'vdst', 'src0', 'src1', 'src2', 'omod', 'neg', 'abs', 'clmp', 'opsel'],
-  'VOP3P': ['op', 'vdst', 'src0', 'src1', 'src2', 'neg', 'neg_hi', 'opsel', 'opsel_hi', 'clmp'],
-  'FLAT': ['op', 'vdst', 'addr', 'data', 'saddr', 'offset', 'seg', 'dlc', 'glc', 'slc'],
-  'MUBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
-  'MTBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
-  'MIMG': ['op', 'vdata', 'vaddr', 'srsrc', 'ssamp', 'dmask', 'dim', 'unrm', 'dlc', 'glc', 'slc'],
-  'EXP': ['en', 'target', 'vsrc0', 'vsrc1', 'vsrc2', 'vsrc3', 'done', 'row'],
-  'VINTERP': ['op', 'vdst', 'src0', 'src1', 'src2', 'waitexp', 'clmp', 'opsel', 'neg'],
-  'VOPD': ['opx', 'opy', 'vdstx', 'vdsty', 'srcx0', 'vsrcx1', 'srcy0', 'vsrcy1'],
-  'LDSDIR': ['op', 'vdst', 'attr', 'attr_chan', 'wait_va']}
-SRC_EXTRAS = {233: 'DPP8', 234: 'DPP8FI', 250: 'DPP16', 251: 'VCCZ', 252: 'EXECZ', 254: 'LDS_DIRECT'}
-FLOAT_MAP = {'0.5': 'POS_HALF', '-0.5': 'NEG_HALF', '1.0': 'POS_ONE', '-1.0': 'NEG_ONE', '2.0': 'POS_TWO', '-2.0': 'NEG_TWO',
-  '4.0': 'POS_FOUR', '-4.0': 'NEG_FOUR', '1/(2*PI)': 'INV_2PI', '0': 'ZERO'}
-
-def parse_bits(s: str) -> tuple[int, int] | None:
-  return (int(m.group(1)), int(m.group(2) or m.group(1))) if (m := re.match(r'\[(\d+)(?::(\d+))?\]', s)) else None
-
-def parse_fields_table(table: list, fmt: str, enums: set[str]) -> list[tuple]:
-  fields = []
-  for row in table[1:]:
-    if not row or not row[0]: continue
-    name, bits_str = row[0].split('\n')[0].strip(), (row[1] or '').split('\n')[0].strip()
-    if not (bits := parse_bits(bits_str)): continue
-    enc_val, hi, lo = None, bits[0], bits[1]
-    if name == 'ENCODING' and row[2] and (m := re.search(r"'b([01_]+)", row[2])):
-      enc_bits = m.group(1).replace('_', '')
-      enc_val = int(enc_bits, 2)
-      declared_width, actual_width = hi - lo + 1, len(enc_bits)
-      if actual_width > declared_width: lo = hi - actual_width + 1
-    ftype = f"{fmt}Op" if name == 'OP' and f"{fmt}Op" in enums else FIELD_TYPES.get(name.upper())
-    fields.append((name, hi, lo, enc_val, ftype))
-  return fields
-
-def generate(output_path: pathlib.Path|str|None = None) -> dict:
-  """Generate RDNA3.5 instruction definitions from the AMD ISA PDF. Returns dict with formats for testing."""
-  pdf = pdfplumber.open(fetch(PDF_URL))
-  pages = pdf.pages[150:200]
-  page_texts = [p.extract_text() or '' for p in pages]
-  page_tables = [[t.extract() for t in p.find_tables()] for p in pages]
-  full_text = '\n'.join(page_texts)
-
-  # parse SSRC encoding from first page with VCC_LO
-  src_enum = dict(SRC_EXTRAS)
-  for text in page_texts[:10]:
-    if 'SSRC0' in text and 'VCC_LO' in text:
-      for m in re.finditer(r'^(\d+)\s+(\S+)', text, re.M):
-        val, name = int(m.group(1)), m.group(2).rstrip('.:')
-        if name in FLOAT_MAP: src_enum[val] = FLOAT_MAP[name]
-        elif re.match(r'^[A-Z][A-Z0-9_]*$', name): src_enum[val] = name
-      break
-
-  # parse opcode tables
-  enums: dict[str, dict[int, str]] = {}
-  for m in re.finditer(r'Table \d+\. (\w+) Opcodes(.*?)(?=Table \d+\.|\n\d+\.\d+\.\d+\.\s+\w+\s*\nDescription|$)', full_text, re.S):
-    if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+([A-Z][A-Z0-9_]+)', m.group(2))}:
-      enums[m.group(1) + "Op"] = ops
-  if vopd_m := re.search(r'Table \d+\. VOPD Y-Opcodes\n(.*?)(?=Table \d+\.|15\.\d)', full_text, re.S):
-    if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+(V_DUAL_\w+)', vopd_m.group(1))}:
-      enums["VOPDOp"] = ops
-  enum_names = set(enums.keys())
-
-  def is_fields_table(t) -> bool: return t and len(t) > 1 and t[0] and 'Field' in str(t[0][0] or '')
-  def has_encoding(fields) -> bool: return any(f[0] == 'ENCODING' for f in fields)
-  def has_header_before_fields(text) -> bool:
-    return (pos := text.find('Field Name')) != -1 and bool(re.search(r'\d+\.\d+\.\d+\.\s+\w+\s*\n', text[:pos]))
-
-  # find format headers with their page indices
-  format_headers = []  # (fmt_name, page_idx)
-  for i, text in enumerate(page_texts):
-    for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n?Description', text): format_headers.append((m.group(1), i, m.start()))
-    for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n', text):
-      if m.start() > len(text) - 200 and 'Description' not in text[m.end():] and i + 1 < len(page_texts):
-        next_text = page_texts[i + 1].lstrip()
-        if next_text.startswith('Description') or (next_text.startswith('"RDNA') and 'Description' in next_text[:200]):
-          format_headers.append((m.group(1), i, m.start()))
-
-  # parse instruction formats
-  formats: dict[str, list] = {}
-  for fmt_name, page_idx, header_pos in format_headers:
-    if fmt_name in formats: continue
-    text, tables = page_texts[page_idx], page_tables[page_idx]
-    field_pos = text.find('Field Name', header_pos)
-
-    # find fields table with ENCODING (same page or up to 2 pages ahead)
-    fields = None
-    for offset in range(3):
-      if page_idx + offset >= len(pages): break
-      if offset > 0 and has_header_before_fields(page_texts[page_idx + offset]): break
-      for t in page_tables[page_idx + offset] if offset > 0 or field_pos > header_pos else []:
-        if is_fields_table(t) and (f := parse_fields_table(t, fmt_name, enum_names)) and has_encoding(f):
-          fields = f
-          break
-      if fields: break
-
-    # for modifier formats (no ENCODING), accept first fields table on same page
-    if not fields and field_pos > header_pos:
-      for t in tables:
-        if is_fields_table(t) and (f := parse_fields_table(t, fmt_name, enum_names)):
-          fields = f
-          break
-
-    if not fields: continue
-    field_names = {f[0] for f in fields}
-
-    # check next pages for continuation fields (tables without ENCODING)
-    for pg_offset in range(1, 3):
-      if page_idx + pg_offset >= len(pages) or has_header_before_fields(page_texts[page_idx + pg_offset]): break
-      for t in page_tables[page_idx + pg_offset]:
-        if is_fields_table(t) and (extra := parse_fields_table(t, fmt_name, enum_names)) and not has_encoding(extra):
-          for ef in extra:
-            if ef[0] not in field_names:
-              fields.append(ef)
-              field_names.add(ef[0])
-          break
-    formats[fmt_name] = fields
-
-  # fix known PDF errors (verified against LLVM test vectors)
-  # SMEM: PDF says DLC=bit14, GLC=bit16 but actual encoding is DLC=bit13, GLC=bit14
-  if 'SMEM' in formats:
-    formats['SMEM'] = [(n, 13 if n == 'DLC' else 14 if n == 'GLC' else h, 13 if n == 'DLC' else 14 if n == 'GLC' else l, e, t)
-                       for n, h, l, e, t in formats['SMEM']]
-
-  # generate output
-  def enum_lines(name, items):
-    return [f"class {name}(IntEnum):"] + [f"  {n} = {v}" for v, n in sorted(items.items())] + [""]
-  def field_key(f): return order.index(f[0].lower()) if f[0].lower() in order else 1000
-  lines = ["# autogenerated from AMD RDNA3.5 ISA PDF by gen.py - do not edit", "from enum import IntEnum",
-           "from typing import Annotated",
-           "from extra.assembly.rdna3.lib import bits, BitField, Inst32, Inst64, SGPR, VGPR, TTMP as TTMP, s as s, v as v, ttmp as ttmp, SSrc, Src, SImm, Imm, VDSTYEnc, SGPRField, VGPRField",
-           "import functools", ""]
-  lines += enum_lines("SrcEnum", src_enum) + sum([enum_lines(n, ops) for n, ops in sorted(enums.items())], [])
-  # Format-specific field defaults (verified against LLVM test vectors)
-  format_defaults = {'VOP3P': {'opsel_hi': 3, 'opsel_hi2': 1}}
-  lines.append("# instruction formats")
-  for fmt_name, fields in sorted(formats.items()):
-    base = "Inst64" if max(f[1] for f in fields) > 31 or fmt_name == 'VOP3SD' else "Inst32"
-    order = FIELD_ORDER.get(fmt_name, [])
-    lines.append(f"class {fmt_name}({base}):")
-    if enc := next((f for f in fields if f[0] == 'ENCODING'), None):
-      enc_str = f"bits[{enc[1]}:{enc[2]}] == 0b{enc[3]:b}" if enc[1] != enc[2] else f"bits[{enc[1]}] == {enc[3]}"
-      lines.append(f"  encoding = {enc_str}")
-    if defaults := format_defaults.get(fmt_name):
-      lines.append(f"  _defaults = {defaults}")
-    for name, hi, lo, _, ftype in sorted([f for f in fields if f[0] != 'ENCODING'], key=field_key):
-      # Wrap IntEnum types (ending in Op) with Annotated[BitField, ...] for correct typing
-      if ftype and ftype.endswith('Op'):
-        ann = f":Annotated[BitField, {ftype}]"
-      else:
-        ann = f":{ftype}" if ftype else ""
-      lines.append(f"  {name.lower()}{ann} = bits[{hi}]" if hi == lo else f"  {name.lower()}{ann} = bits[{hi}:{lo}]")
-    lines.append("")
-  lines.append("# instruction helpers")
-  for cls_name, ops in sorted(enums.items()):
-    fmt = cls_name[:-2]
-    for op_val, name in sorted(ops.items()):
-      seg = {"GLOBAL": ", seg=2", "SCRATCH": ", seg=2"}.get(fmt, "")
-      tgt = {"GLOBAL": "FLAT, GLOBALOp", "SCRATCH": "FLAT, SCRATCHOp"}.get(fmt, f"{fmt}, {cls_name}")
-      if fmt in formats or fmt in ("GLOBAL", "SCRATCH"):
-        # VOP1/VOP2/VOPC get _e32 suffix, VOP3 promoted ops (< 512) get _e64 suffix
-        if fmt in ("VOP1", "VOP2", "VOPC"):
-          suffix = "_e32"
-        elif fmt == "VOP3" and op_val < 512:
-          suffix = "_e64"
-        else:
-          suffix = ""
-        lines.append(f"{name.lower()}{suffix} = functools.partial({tgt}.{name}{seg})")
-  # export SrcEnum values, but skip DPP8/DPP16 which conflict with class names
-  skip_exports = {'DPP8', 'DPP16'}
-  lines += [""] + [f"{name} = SrcEnum.{name}" for _, name in sorted(src_enum.items()) if name not in skip_exports] + ["OFF = NULL\n"]
-
-  if output_path is not None: pathlib.Path(output_path).write_text('\n'.join(lines))
-  return {"formats": formats, "enums": enums, "src_enum": src_enum}
-
-if __name__ == "__main__":
-  result = generate("extra/assembly/rdna3/autogen/__init__.py")
-  print(f"generated SrcEnum ({len(result['src_enum'])}) + {len(result['enums'])} opcode enums + {len(result['formats'])} format classes")

extra/assembly/rdna3/generate.py

@@ -0,0 +1,136 @@
+import os, sys, struct
+sys.path.append(os.getcwd())
+# PROFILE=1 to use
+#os.environ["PROFILE"] = "1"
+os.environ["SQTT"] = "1"
+os.environ["SQTT_ITRACE_SE_MASK"] = "1"
+os.environ["SQTT_LIMIT_SE"] = "1"
+import xml.etree.ElementTree as ET
+
+from tinygrad import nn, Tensor, Device
+from tinygrad.helpers import get_single_element
+from tinygrad.runtime.support.elf import elf_loader
+from tinygrad.runtime.ops_amd import ProfileSQTTEvent
+from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
+
+def disassemble(text, root:ET.Element):
+  i = 0
+  while i < len(text):
+    ins = struct.unpack("I", text[i:i+4])[0]
+
+    # 1. Get the encoding
+    did_match = False
+    for enc_el in root.findall("./ISA/Encodings/Encoding"):
+      mask = enc_el.findtext("EncodingIdentifierMask")
+      assert len(mask)%32 == 0
+      bit_mask = int(mask, 2)
+      iden = [int(x.text, 2) for x in enc_el.find("EncodingIdentifiers").findall("EncodingIdentifier")]
+      for ide in iden:
+        if ins&bit_mask == ide:
+          did_match = True
+          break
+      if did_match: break
+    if not did_match: raise RuntimeError(f"unknown instruction {ins:08X}")
+    if len(mask) >= 64: ins = (struct.unpack("I", text[i+4:i+8])[0]<<32) | ins
+    if len(mask) >= 96: ins = (struct.unpack("I", text[i+8:i+12])[0]<<64) | ins
+    encoding_name = enc_el.findtext("EncodingName")
+
+    #print(ET.tostring(enc_el).decode())
+
+    # 2. Parse the Fields for this Encoding
+    field_data = {}
+    for field in enc_el.findall("MicrocodeFormat/BitMap/Field"):
+      # Fields can be split into multiple ranges (RangeCount > 1)
+      ranges = sorted(field.findall("BitLayout/Range"), key=lambda x: int(x.attrib.get('Order')))
+      val = 0
+      current_shift = 0
+      for rng in ranges:
+        width = int(rng.find("BitCount").text)
+        chunk = (ins >> int(rng.find("BitOffset").text)) & ((1 << width) - 1)
+        val |= (chunk << current_shift)
+        current_shift += width
+      field_data[field.find("FieldName").text] = val
+    # this is already used
+    del field_data["ENCODING"]
+
+    # 3. Extract the instruction
+    did_match = False
+    for ins_el in root.findall("./ISA/Instructions/Instruction"):
+      ins_name = ins_el.findtext("InstructionName")
+      for ins_enc in ins_el.findall("InstructionEncodings/InstructionEncoding"):
+        if ins_enc.findtext("EncodingName") == encoding_name:
+          opcode = int(ins_enc.findtext("Opcode"))
+          if "OP" in field_data and opcode == field_data["OP"]:
+            did_match = True
+            del field_data["OP"]
+            break
+        if did_match: break
+      if did_match: break
+
+    #print(ET.tostring(ins_enc).decode())
+    #print()
+    #print(field_data)
+    if not did_match:
+      print(f"{i:4X} : {ins:16x} -- {encoding_name}")
+    elif did_match:
+      params = []
+      #print(ET.tostring(ins_el).decode())
+
+      # 4. Extract the opcodes
+      for op_ins in ins_enc.findall("Operands/Operand"):
+        op_type = op_ins.findtext("OperandType")
+        op_size = op_ins.findtext("OperandSize")
+        op_fmt = op_ins.findtext("DataFormatName")
+        op_field_name = op_ins.findtext("FieldName")
+        if op_field_name is None: continue
+        assert op_field_name in field_data
+        # loop through operands for compare
+        for op_el in root.findall("./ISA/OperandTypes/OperandType"):
+          test_op_type = op_el.findtext("OperandTypeName")
+          val_dict = {}
+          for op_val in op_el.findall("OperandPredefinedValues/PredefinedValue"):
+            val_dict[int(op_val.findtext("Value"))] = op_val.findtext("Name")
+          if op_type == test_op_type:
+            if field_data[op_field_name] in val_dict:
+              print(op_type, op_size, op_fmt)
+              params.append(val_dict[field_data[op_field_name]])
+            else:
+              params.append(f"{op_type}({field_data[op_field_name]})")
+            del field_data[op_field_name]
+            #print(op_type, op_size, op_fmt, op_el, op_field_name,
+            #      field_data[op_field_name],
+            #      val_dict.get(field_data[op_field_name], "<UNK>"))
+            #print(ET.tostring(op_el).decode())
+
+      print(f"{i:4X} : {ins:16x} -- {ins_name.lower()} {', '.join(params)}", field_data)
+
+    # advance
+    i += len(mask) // 8
+
+  #print(ET.tostring(root).decode())
+
+if __name__ == "__main__":
+  # human readable manual at https://docs.amd.com/v/u/en-US/rdna35_instruction_set_architecture
+  fns = nn.state.zip_extract(Tensor.from_url("https://gpuopen.com/download/machine-readable-isa/latest/"))
+  xml_str = fns['amdgpu_isa_rdna3_5.xml'].to("CPU").data()
+  with open("/tmp/rdna35.xml", "wb") as f: f.write(bytes(xml_str))
+  root = ET.fromstring(xml_str)
+
+  a = Tensor.empty(16)+1
+  for ei in a.schedule():
+    ei.lower()
+    # get text
+    _, hdr, _ = elf_loader(ei.prg.lib)
+    text = get_single_element([x for x in hdr if x.name==".text"]).content
+
+    # llvm disassembler
+    Device["AMD"].compiler.disassemble(ei.prg.lib)
+
+    # run program
+    ei.run()
+
+  sqtt_events = [e for e in Device["AMD"].profile_events if isinstance(e, ProfileSQTTEvent)]
+  for e in sqtt_events[0:1]: # only the first SE
+    parse_sqtt_print_packets(e.blob)
+
+  disassemble(text[:0x40], root)

extra/assembly/rdna3/lib.py

@@ -1,254 +0,0 @@
-# library for RDNA3 assembly DSL
-from __future__ import annotations
-from enum import IntEnum
-from typing import overload, Annotated, TypeVar, Generic
-
-# Bit field DSL
-class BitField:
-  def __init__(self, hi: int, lo: int, name: str | None = None): self.hi, self.lo, self.name = hi, lo, name
-  def __set_name__(self, owner, name): self.name, self._owner = name, owner
-  def __eq__(self, val: int) -> tuple[BitField, int]: return (self, val)  # type: ignore
-  def mask(self) -> int: return (1 << (self.hi - self.lo + 1)) - 1
-  @property
-  def marker(self) -> type | None:
-    # Get marker from Annotated type hint if present
-    import typing
-    if hasattr(self, '_owner') and self.name:
-      hints = typing.get_type_hints(self._owner, include_extras=True)
-      if self.name in hints:
-        hint = hints[self.name]
-        if typing.get_origin(hint) is Annotated:
-          args = typing.get_args(hint)
-          return args[1] if len(args) > 1 else None
-    return None
-  @overload
-  def __get__(self, obj: None, objtype: type) -> BitField: ...
-  @overload
-  def __get__(self, obj: object, objtype: type | None = None) -> int: ...
-  def __get__(self, obj, objtype=None):
-    if obj is None: return self
-    val = unwrap(obj._values.get(self.name, 0))
-    # Convert to IntEnum if marker is an IntEnum subclass
-    if self.marker and isinstance(self.marker, type) and issubclass(self.marker, IntEnum):
-      try: return self.marker(val)
-      except ValueError: pass
-    return val
-
-class _Bits:
-  def __getitem__(self, key) -> BitField: return BitField(key.start, key.stop) if isinstance(key, slice) else BitField(key, key)
-bits = _Bits()
-
-# Register types
-class Reg:
-  def __init__(self, idx: int, count: int = 1, hi: bool = False, neg: bool = False): self.idx, self.count, self.hi, self.neg = idx, count, hi, neg
-  def __repr__(self): return f"{self.__class__.__name__.lower()[0]}[{self.idx}]" if self.count == 1 else f"{self.__class__.__name__.lower()[0]}[{self.idx}:{self.idx + self.count}]"
-  def __neg__(self): return self.__class__(self.idx, self.count, self.hi, neg=not self.neg)
-
-T = TypeVar('T', bound=Reg)
-class _RegFactory(Generic[T]):
-  def __init__(self, cls: type[T], name: str): self._cls, self._name = cls, name
-  @overload
-  def __getitem__(self, key: int) -> Reg: ...
-  @overload
-  def __getitem__(self, key: slice) -> Reg: ...
-  def __getitem__(self, key: int | slice) -> Reg:
-    return self._cls(key.start, key.stop - key.start + 1) if isinstance(key, slice) else self._cls(key)
-  def __repr__(self): return f"<{self._name} factory>"
-
-class SGPR(Reg): pass
-class VGPR(Reg): pass
-class TTMP(Reg): pass
-s: _RegFactory[SGPR] = _RegFactory(SGPR, "SGPR")
-v: _RegFactory[VGPR] = _RegFactory(VGPR, "VGPR")
-ttmp: _RegFactory[TTMP] = _RegFactory(TTMP, "TTMP")
-
-# Field type markers (runtime classes for validation)
-class _SSrc: pass
-class _Src: pass
-class _Imm: pass
-class _SImm: pass
-class _VDSTYEnc: pass  # VOPD vdsty: encoded = actual >> 1, actual = (encoded << 1) | ((vdstx & 1) ^ 1)
-class _SGPRField: pass
-class _VGPRField: pass
-
-# Type aliases for annotations - tells mypy it's a BitField while preserving marker info
-SSrc = Annotated[BitField, _SSrc]
-Src = Annotated[BitField, _Src]
-Imm = Annotated[BitField, _Imm]
-SImm = Annotated[BitField, _SImm]
-VDSTYEnc = Annotated[BitField, _VDSTYEnc]
-SGPRField = Annotated[BitField, _SGPRField]
-VGPRField = Annotated[BitField, _VGPRField]
-class RawImm:
-  def __init__(self, val: int): self.val = val
-  def __repr__(self): return f"RawImm({self.val})"
-  def __eq__(self, other): return isinstance(other, RawImm) and self.val == other.val
-
-def unwrap(val) -> int:
-  return val.val if isinstance(val, RawImm) else val.value if hasattr(val, 'value') else val.idx if hasattr(val, 'idx') else val
-
-# Encoding helpers
-FLOAT_ENC = {0.5: 240, -0.5: 241, 1.0: 242, -1.0: 243, 2.0: 244, -2.0: 245, 4.0: 246, -4.0: 247}
-SRC_FIELDS = {'src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'soffset', 'srcx0', 'srcy0'}
-RAW_FIELDS = {'vdata', 'vdst', 'vaddr', 'addr', 'data', 'data0', 'data1', 'sdst', 'sdata'}
-
-def _encode_reg(val) -> int:
-  if isinstance(val, TTMP): return 108 + val.idx
-  return val.idx | (0x80 if val.hi else 0)
-
-def encode_src(val) -> int:
-  if isinstance(val, VGPR): return 256 + _encode_reg(val)
-  if isinstance(val, Reg): return _encode_reg(val)
-  if hasattr(val, 'value'): return val.value
-  if isinstance(val, float): return 128 if val == 0.0 else FLOAT_ENC.get(val, 255)
-  return 128 + val if isinstance(val, int) and 0 <= val <= 64 else 192 + (-val) if isinstance(val, int) and -16 <= val <= -1 else 255
-
-# Instruction base class
-class Inst:
-  _fields: dict[str, BitField]
-  _encoding: tuple[BitField, int] | None = None
-  _defaults: dict[str, int] = {}
-  _values: dict[str, int | RawImm]
-  _words: int  # size in 32-bit words, set by decode_program
-  _literal: int | None
-
-  def __init_subclass__(cls, **kwargs):
-    super().__init_subclass__(**kwargs)
-    cls._fields = {n: v[0] if isinstance(v, tuple) else v for n, v in cls.__dict__.items() if isinstance(v, BitField) or (isinstance(v, tuple) and len(v) == 2 and isinstance(v[0], BitField))}
-    if 'encoding' in cls._fields and isinstance(cls.__dict__.get('encoding'), tuple): cls._encoding = cls.__dict__['encoding']
-
-  def __init__(self, *args, literal: int | None = None, **kwargs):
-    self._values, self._literal = dict(self._defaults), literal
-    # Map positional args to field names
-    field_names = [n for n in self._fields if n != 'encoding']
-    orig_args = dict(zip(field_names, args))
-    orig_args.update(kwargs)
-    self._values.update(orig_args)
-    # Validate register counts for SMEM instructions (before encoding)
-    if self.__class__.__name__ == 'SMEM':
-      op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
-      if op_val is not None:
-        if hasattr(op_val, 'value'): op_val = op_val.value
-        expected_cnt = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op_val)
-        sdata_val = orig_args.get('sdata')
-        if expected_cnt is not None and isinstance(sdata_val, Reg) and sdata_val.count != expected_cnt:
-          raise ValueError(f"SMEM op {op_val} expects {expected_cnt} registers, got {sdata_val.count}")
-    # Validate register counts for SOP1 instructions (b32 = 1 reg, b64 = 2 regs)
-    if self.__class__.__name__ == 'SOP1':
-      op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
-      if op_val is not None and hasattr(op_val, 'name'):
-        expected = 2 if op_val.name.endswith('_B64') else 1
-        sdst_val, ssrc0_val = orig_args.get('sdst'), orig_args.get('ssrc0')
-        if isinstance(sdst_val, Reg) and sdst_val.count != expected:
-          raise ValueError(f"SOP1 {op_val.name} expects {expected} destination register(s), got {sdst_val.count}")
-        if isinstance(ssrc0_val, Reg) and ssrc0_val.count != expected:
-          raise ValueError(f"SOP1 {op_val.name} expects {expected} source register(s), got {ssrc0_val.count}")
-    # Type check and encode values
-    for name, val in list(self._values.items()):
-      if name == 'encoding': continue
-      # For RawImm, only process RAW_FIELDS to unwrap to int
-      if isinstance(val, RawImm):
-        if name in RAW_FIELDS: self._values[name] = val.val
-        continue
-      field = self._fields.get(name)
-      marker = field.marker if field else None
-      # Type validation
-      if marker is _SGPRField:
-        if isinstance(val, VGPR): raise TypeError(f"field '{name}' requires SGPR, got VGPR")
-        if not isinstance(val, (SGPR, TTMP, int, RawImm)): raise TypeError(f"field '{name}' requires SGPR, got {type(val).__name__}")
-      if marker is _VGPRField:
-        if not isinstance(val, VGPR): raise TypeError(f"field '{name}' requires VGPR, got {type(val).__name__}")
-      if marker is _SSrc and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires scalar source, got VGPR")
-      # Encode source fields as RawImm for consistent disassembly
-      if name in SRC_FIELDS:
-        encoded = encode_src(val)
-        self._values[name] = RawImm(encoded)
-        # Handle negation modifier for VOP3 instructions
-        if isinstance(val, Reg) and val.neg and 'neg' in self._fields:
-          neg_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
-          cur_neg = self._values.get('neg', 0)
-          self._values['neg'] = (cur_neg.val if isinstance(cur_neg, RawImm) else cur_neg) | neg_bit
-        # Track literal value if needed (encoded as 255)
-        if encoded == 255 and self._literal is None and isinstance(val, int) and not isinstance(val, IntEnum):
-          self._literal = val
-        elif encoded == 255 and self._literal is None and isinstance(val, float):
-          import struct
-          self._literal = struct.unpack('<I', struct.pack('<f', val))[0]
-      # Encode raw register fields for consistent repr
-      elif name in RAW_FIELDS:
-        if isinstance(val, Reg): self._values[name] = _encode_reg(val)
-        elif hasattr(val, 'value'): self._values[name] = val.value  # IntEnum like SrcEnum.NULL
-      # Encode sbase (divided by 2) and srsrc/ssamp (divided by 4)
-      elif name == 'sbase' and isinstance(val, Reg):
-        self._values[name] = val.idx // 2
-      elif name in {'srsrc', 'ssamp'} and isinstance(val, Reg):
-        self._values[name] = val.idx // 4
-      # VOPD vdsty: encode as actual >> 1 (constraint: vdsty parity must be opposite of vdstx)
-      elif marker is _VDSTYEnc and isinstance(val, VGPR):
-        self._values[name] = val.idx >> 1
-
-  def _encode_field(self, name: str, val) -> int:
-    if isinstance(val, RawImm): return val.val
-    if name in {'srsrc', 'ssamp'}: return val.idx // 4 if isinstance(val, Reg) else val
-    if name == 'sbase': return val.idx // 2 if isinstance(val, Reg) else val
-    if name in RAW_FIELDS: return _encode_reg(val) if isinstance(val, Reg) else val
-    if isinstance(val, Reg) or name in SRC_FIELDS: return encode_src(val)
-    return val.value if hasattr(val, 'value') else val
-
-  def to_int(self) -> int:
-    word = (self._encoding[1] & self._encoding[0].mask()) << self._encoding[0].lo if self._encoding else 0
-    for n, bf in self._fields.items():
-      if n != 'encoding' and n in self._values: word |= (self._encode_field(n, self._values[n]) & bf.mask()) << bf.lo
-    return word
-
-  def _get_literal(self) -> int | None:
-    for n in SRC_FIELDS:
-      if n in self._values and not isinstance(v := self._values[n], RawImm) and isinstance(v, int) and not isinstance(v, IntEnum) and not (0 <= v <= 64 or -16 <= v <= -1): return v
-    return None
-
-  def to_bytes(self) -> bytes:
-    result = self.to_int().to_bytes(self._size(), 'little')
-    return result + (lit & 0xffffffff).to_bytes(4, 'little') if (lit := self._get_literal() or getattr(self, '_literal', None)) else result
-
-  @classmethod
-  def _size(cls) -> int: return 4 if issubclass(cls, Inst32) else 8
-  def size(self) -> int: return self._size() + (4 if self._literal is not None else 0)
-
-  @classmethod
-  def from_int(cls, word: int):
-    inst = object.__new__(cls)
-    inst._values = {n: RawImm(v) if n in SRC_FIELDS else v for n, bf in cls._fields.items() if n != 'encoding' for v in [(word >> bf.lo) & bf.mask()]}
-    inst._literal = None
-    return inst
-
-  @classmethod
-  def from_bytes(cls, data: bytes):
-    inst = cls.from_int(int.from_bytes(data[:cls._size()], 'little'))
-    op_val = inst._values.get('op', 0)
-    has_literal = cls.__name__ == 'VOP2' and op_val in (44, 45, 55, 56)
-    has_literal = has_literal or (cls.__name__ == 'SOP2' and op_val in (69, 70))
-    for n in SRC_FIELDS:
-      if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255: has_literal = True
-    if has_literal and len(data) >= cls._size() + 4: inst._literal = int.from_bytes(data[cls._size():cls._size()+4], 'little')
-    return inst
-
-  def __repr__(self):
-    # Use _fields order and exclude fields that are 0/default (for consistent repr after roundtrip)
-    def is_zero(v): return (isinstance(v, int) and v == 0) or (isinstance(v, VGPR) and v.idx == 0 and v.count == 1)
-    items = [(k, self._values[k]) for k in self._fields if k in self._values and k != 'encoding'
-             and not (is_zero(self._values[k]) and k not in {'op'})]
-    lit = f", literal={hex(self._literal)}" if self._literal is not None else ""
-    return f"{self.__class__.__name__}({', '.join(f'{k}={v}' for k, v in items)}{lit})"
-
-  def __eq__(self, other):
-    if not isinstance(other, Inst): return NotImplemented
-    return self.__class__ == other.__class__ and self._values == other._values and self._literal == other._literal
-
-  def __hash__(self): return hash((self.__class__.__name__, tuple(sorted((k, repr(v)) for k, v in self._values.items())), self._literal))
-
-  def disasm(self) -> str:
-    from extra.assembly.rdna3.asm import disasm
-    return disasm(self)
-
-class Inst32(Inst): pass
-class Inst64(Inst): pass

extra/assembly/rdna3/parse.py

@@ -0,0 +1,15 @@
+from tinygrad import Tensor, nn
+import xml.etree.ElementTree as ET
+
+if __name__ == "__main__":
+  # human readable manual at https://docs.amd.com/v/u/en-US/rdna35_instruction_set_architecture
+  fns = nn.state.zip_extract(Tensor.from_url("https://gpuopen.com/download/machine-readable-isa/latest/"))
+  xml_str = fns['amdgpu_isa_rdna3_5.xml'].to("CPU").data()
+  root = ET.fromstring(xml_str)
+
+  for op_el in root.findall("./ISA/OperandTypes/OperandType"):
+    op_name = op_el.findtext("OperandTypeName")
+    val_dict = {}
+    for op_val in op_el.findall("OperandPredefinedValues/PredefinedValue"):
+      val_dict[int(op_val.findtext("Value"))] = op_val.findtext("Name")
+    print(op_name, val_dict)

extra/assembly/rdna3/test/bench_emu.py

@@ -1,294 +0,0 @@
-#!/usr/bin/env python3
-"""Benchmark comparing Python vs Rust RDNA3 emulators on synthetic and real tinygrad kernels."""
-import ctypes, time, os, struct, cProfile, pstats, io
-from pathlib import Path
-from typing import Callable
-
-# Set AMD=1 before importing tinygrad
-os.environ["AMD"] = "1"
-
-from extra.assembly.rdna3.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program, step_wave, WaveState, WAVE_SIZE
-
-REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
-if not REMU_PATH.exists():
-  REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.dylib"
-
-def get_rust_remu():
-  """Load the Rust libremu shared library."""
-  if not REMU_PATH.exists(): return None
-  remu = ctypes.CDLL(str(REMU_PATH))
-  remu.run_asm.restype = ctypes.c_int32
-  remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
-                           ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
-  return remu
-
-def count_instructions(kernel: bytes) -> int:
-  """Count instructions in a kernel."""
-  return len(decode_program(kernel))
-
-def setup_buffers(buf_sizes: list[int], init_data: dict[int, bytes] | None = None):
-  """Allocate buffers and return args pointer + valid ranges."""
-  if init_data is None: init_data = {}
-  buffers = []
-  for i, size in enumerate(buf_sizes):
-    padded = ((size + 15) // 16) * 16 + 16
-    data = init_data.get(i, b'\x00' * padded)
-    data_list = list(data) + [0] * (padded - len(data))
-    buf = (ctypes.c_uint8 * padded)(*data_list[:padded])
-    buffers.append(buf)
-  args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
-  args_ptr = ctypes.addressof(args)
-  ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
-  ranges.add((args_ptr, ctypes.sizeof(args)))
-  return buffers, args, args_ptr, ranges
-
-def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size, args_ptr, iterations: int = 5):
-  """Benchmark an emulator and return average time."""
-  gx, gy, gz = global_size
-  lx, ly, lz = local_size
-  kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-  lib_ptr = ctypes.addressof(kernel_buf)
-
-  # Warmup
-  run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
-
-  # Timed runs
-  times = []
-  for _ in range(iterations):
-    start = time.perf_counter()
-    result = run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
-    end = time.perf_counter()
-    if result != 0:
-      print(f"  {name} returned error: {result}")
-      return None
-    times.append(end - start)
-
-  return sum(times) / len(times)
-
-def create_synthetic_kernel(n_ops: int) -> bytes:
-  """Create a synthetic kernel with n_ops vector operations."""
-  instructions = []
-  # VOP2 instructions: v_add_f32, v_mul_f32, v_max_f32, v_min_f32
-  ops = [
-    (0b0000011 << 25) | (1 << 17) | (0 << 9) | 256,  # v_add_f32 v0, v0, v1
-    (0b0001000 << 25) | (1 << 17) | (0 << 9) | 256,  # v_mul_f32 v0, v0, v1
-    (0b0010000 << 25) | (1 << 17) | (0 << 9) | 256,  # v_max_f32 v0, v0, v1
-    (0b0001111 << 25) | (1 << 17) | (0 << 9) | 256,  # v_min_f32 v0, v0, v1
-  ]
-  for i in range(n_ops):
-    instructions.append(ops[i % len(ops)])
-  # S_ENDPGM
-  instructions.append((0b101111111 << 23) | (48 << 16) | 0)
-  return b''.join(struct.pack('<I', inst) for inst in instructions)
-
-def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes]] | None:
-  """Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data)."""
-  try:
-    from tinygrad import Tensor
-    from tinygrad.runtime.support.elf import elf_loader
-    import numpy as np
-    np.random.seed(42)
-
-    ops = {
-      "add": lambda: Tensor.empty(1024) + Tensor.empty(1024),
-      "mul": lambda: Tensor.empty(1024) * Tensor.empty(1024),
-      "matmul_small": lambda: Tensor.empty(16, 16) @ Tensor.empty(16, 16),
-      "matmul_medium": lambda: Tensor.empty(64, 64) @ Tensor.empty(64, 64),
-      "reduce_sum": lambda: Tensor.empty(4096).sum(),
-      "reduce_max": lambda: Tensor.empty(4096).max(),
-      "softmax": lambda: Tensor.empty(256).softmax(),
-      "layernorm": lambda: Tensor.empty(32, 64).layernorm(),
-      "conv2d": lambda: Tensor.empty(1, 4, 16, 16).conv2d(Tensor.empty(4, 4, 3, 3)),
-      "gelu": lambda: Tensor.empty(1024).gelu(),
-      "exp": lambda: Tensor.empty(1024).exp(),
-      "sin": lambda: Tensor.empty(1024).sin(),
-    }
-
-    if op_name not in ops: return None
-    out = ops[op_name]()
-    sched = out.schedule()
-
-    for ei in sched:
-      lowered = ei.lower()
-      if ei.ast.op.name == 'SINK' and 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_sizes = [b.nbytes for b in lowered.bufs]
-            # Get initial data from numpy arrays if available
-            buf_data = {}
-            for i, buf in enumerate(lowered.bufs):
-              if hasattr(buf, 'base') and buf.base is not None and hasattr(buf.base, '_buf'):
-                try: buf_data[i] = bytes(buf.base._buf)
-                except: pass
-            return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data)
-    return None
-  except Exception as e:
-    print(f"  Error getting kernel: {e}")
-    return None
-
-def profile_python_emu(kernel: bytes, global_size, local_size, args_ptr, n_runs: int = 1):
-  """Profile the Python emulator to find bottlenecks."""
-  gx, gy, gz = global_size
-  lx, ly, lz = local_size
-  kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-  lib_ptr = ctypes.addressof(kernel_buf)
-
-  pr = cProfile.Profile()
-  pr.enable()
-  for _ in range(n_runs):
-    python_run_asm(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
-  pr.disable()
-
-  s = io.StringIO()
-  ps = pstats.Stats(pr, stream=s).sort_stats('cumulative')
-  ps.print_stats(20)
-  return s.getvalue()
-
-def measure_step_rate(kernel: bytes, n_steps: int = 10000) -> float:
-  """Measure raw step_wave() performance (steps per second)."""
-  program = decode_program(kernel)
-  if not program: return 0.0
-
-  st = WaveState()
-  st.exec_mask = 0xffffffff
-  lds = bytearray(65536)
-  n_lanes = 32
-
-  # Reset PC for each measurement
-  start = time.perf_counter()
-  for _ in range(n_steps):
-    st.pc = 0
-    while st.pc in program:
-      result = step_wave(program, st, lds, n_lanes)
-      if result == -1: break
-  elapsed = time.perf_counter() - start
-  return n_steps / elapsed if elapsed > 0 else 0
-
-# Test configurations
-SYNTHETIC_TESTS = [
-  ("synthetic_10ops", 10, (1, 1, 1), (32, 1, 1)),
-  ("synthetic_100ops", 100, (1, 1, 1), (32, 1, 1)),
-  ("synthetic_500ops", 500, (1, 1, 1), (32, 1, 1)),
-  ("synthetic_100ops_4wg", 100, (4, 1, 1), (32, 1, 1)),
-  ("synthetic_100ops_16wg", 100, (16, 1, 1), (32, 1, 1)),
-]
-
-TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
-
-def main():
-  import argparse
-  parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
-  parser.add_argument("--profile", action="store_true", help="Profile Python emulator")
-  parser.add_argument("--synthetic-only", action="store_true", help="Only run synthetic tests")
-  parser.add_argument("--tinygrad-only", action="store_true", help="Only run tinygrad tests")
-  parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
-  args = parser.parse_args()
-
-  rust_remu = get_rust_remu()
-  if rust_remu is None:
-    print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
-    print("Running Python-only benchmarks...\n")
-
-  print("=" * 90)
-  print("RDNA3 Emulator Benchmark: Python vs Rust")
-  print("=" * 90)
-
-  results = []
-
-  # Synthetic workloads
-  if not args.tinygrad_only:
-    print("\n[SYNTHETIC WORKLOADS]")
-    print("-" * 90)
-
-    for name, n_ops, global_size, local_size in SYNTHETIC_TESTS:
-      kernel = create_synthetic_kernel(n_ops)
-      n_insts = count_instructions(kernel)
-      n_workgroups = global_size[0] * global_size[1] * global_size[2]
-      n_threads = local_size[0] * local_size[1] * local_size[2]
-      total_work = n_insts * n_workgroups * n_threads
-
-      print(f"\n{name}: {n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
-
-      buf_sizes = [4096]
-      buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes)
-      set_valid_mem_ranges(ranges)
-
-      # Benchmark
-      py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
-      rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
-
-      if py_time:
-        py_rate = total_work / py_time / 1e6
-        print(f"  Python: {py_time*1000:8.3f} ms  ({py_rate:7.2f} M ops/s)")
-      if rust_time:
-        rust_rate = total_work / rust_time / 1e6
-        speedup = py_time / rust_time if py_time else 0
-        print(f"  Rust:   {rust_time*1000:8.3f} ms  ({rust_rate:7.2f} M ops/s)  [{speedup:.1f}x faster]")
-
-      results.append(("synthetic", name, n_insts, n_workgroups, py_time, rust_time))
-
-  # Tinygrad kernels
-  if not args.synthetic_only:
-    print("\n[TINYGRAD KERNELS]")
-    print("-" * 90)
-
-    for op_name in TINYGRAD_TESTS:
-      print(f"\n{op_name}:", end=" ", flush=True)
-      kernel_info = get_tinygrad_kernel(op_name)
-      if kernel_info is None:
-        print("failed to compile")
-        continue
-
-      kernel, global_size, local_size, buf_sizes, buf_data = kernel_info
-      n_insts = count_instructions(kernel)
-      n_workgroups = global_size[0] * global_size[1] * global_size[2]
-      n_threads = local_size[0] * local_size[1] * local_size[2]
-      total_work = n_insts * n_workgroups * n_threads
-
-      print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
-
-      buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
-      set_valid_mem_ranges(ranges)
-
-      py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
-      rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
-
-      if py_time:
-        py_rate = total_work / py_time / 1e6
-        print(f"  Python: {py_time*1000:8.3f} ms  ({py_rate:7.2f} M ops/s)")
-      if rust_time:
-        rust_rate = total_work / rust_time / 1e6
-        speedup = py_time / rust_time if py_time else 0
-        print(f"  Rust:   {rust_time*1000:8.3f} ms  ({rust_rate:7.2f} M ops/s)  [{speedup:.1f}x faster]")
-
-      results.append(("tinygrad", op_name, n_insts, n_workgroups, py_time, rust_time))
-
-      # Optional profiling
-      if args.profile and py_time:
-        print("\n  [PROFILE - Top 10 functions]")
-        profile_output = profile_python_emu(kernel, global_size, local_size, args_ptr)
-        for line in profile_output.split('\n')[5:15]:
-          if line.strip(): print(f"    {line}")
-
-  # Summary table
-  print("\n" + "=" * 90)
-  print("SUMMARY")
-  print("=" * 90)
-  print(f"{'Type':<10} {'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
-  print("-" * 90)
-
-  for test_type, name, n_insts, n_wgs, py_time, rust_time in results:
-    py_ms = f"{py_time*1000:.3f}" if py_time else "error"
-    if rust_time:
-      rust_ms = f"{rust_time*1000:.3f}"
-      speedup = f"{py_time/rust_time:.1f}x" if py_time else "N/A"
-    else:
-      rust_ms, speedup = "N/A", "N/A"
-    print(f"{test_type:<10} {name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
-
-
-
-if __name__ == "__main__":
-  main()

extra/assembly/rdna3/test/test_compare_emulators.py

@@ -1,466 +0,0 @@
-# Test to compare Python and Rust RDNA3 emulators by running real tinygrad kernels
-import unittest, ctypes, os
-from dataclasses import dataclass
-from pathlib import Path
-
-# Set environment before any tinygrad imports to use MOCKGPU
-# This allows generating AMD GPU kernels without requiring real hardware
-os.environ["AMD"] = "1"
-os.environ["MOCKGPU"] = "1"
-
-from extra.assembly.rdna3.emu import WaveState, decode_program, step_wave, WAVE_SIZE
-
-REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
-
-@dataclass
-class KernelInfo:
-  code: bytes
-  global_size: tuple[int, int, int]
-  local_size: tuple[int, int, int]
-  buf_idxs: list[int]  # indices into shared buffer pool
-  buf_sizes: list[int]  # sizes for each buffer index
-
-@dataclass
-class StateSnapshot:
-  pc: int
-  scc: int
-  vcc: int
-  exec_mask: int
-  sgpr: list[int]
-  vgpr: list[list[int]]
-
-  def diff(self, other: 'StateSnapshot', n_lanes: int) -> list[str]:
-    """Return list of differences between two states."""
-    diffs = []
-    if self.pc != other.pc: diffs.append(f"pc: {self.pc} vs {other.pc}")
-    if self.scc != other.scc: diffs.append(f"scc: {self.scc} vs {other.scc}")
-    if self.vcc != other.vcc: diffs.append(f"vcc: 0x{self.vcc:08x} vs 0x{other.vcc:08x}")
-    if self.exec_mask != other.exec_mask: diffs.append(f"exec: 0x{self.exec_mask:08x} vs 0x{other.exec_mask:08x}")
-    for i, (a, b) in enumerate(zip(self.sgpr, other.sgpr)):
-      # Skip VCC_LO/HI (106/107) and EXEC_LO/HI (126/127) as they alias vcc/exec_mask which are compared separately
-      if i in (106, 107, 126, 127): continue
-      if a != b: diffs.append(f"sgpr[{i}]: 0x{a:08x} vs 0x{b:08x}")
-    for lane in range(n_lanes):
-      for i, (a, b) in enumerate(zip(self.vgpr[lane], other.vgpr[lane])):
-        if a != b: diffs.append(f"vgpr[{lane}][{i}]: 0x{a:08x} vs 0x{b:08x}")
-    return diffs
-
-class CStateSnapshot(ctypes.Structure):
-  _fields_ = [("pc", ctypes.c_uint32), ("scc", ctypes.c_uint32), ("vcc", ctypes.c_uint32), ("exec_mask", ctypes.c_uint32),
-              ("sgpr", ctypes.c_uint32 * 128), ("vgpr", (ctypes.c_uint32 * 256) * 32)]
-
-  def to_snapshot(self) -> StateSnapshot:
-    return StateSnapshot(pc=self.pc, scc=self.scc, vcc=self.vcc, exec_mask=self.exec_mask,
-                         sgpr=list(self.sgpr), vgpr=[list(self.vgpr[i]) for i in range(32)])
-
-class RustEmulator:
-  def __init__(self):
-    self.lib = ctypes.CDLL(str(REMU_PATH))
-    self.lib.wave_create.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
-    self.lib.wave_create.restype = ctypes.c_void_p
-    self.lib.wave_step.argtypes = [ctypes.c_void_p]
-    self.lib.wave_step.restype = ctypes.c_int32
-    self.lib.wave_get_snapshot.argtypes = [ctypes.c_void_p, ctypes.POINTER(CStateSnapshot)]
-    self.lib.wave_set_sgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
-    self.lib.wave_set_vgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32]
-    self.lib.wave_init_lds.argtypes = [ctypes.c_void_p, ctypes.c_uint32]
-    self.lib.wave_free.argtypes = [ctypes.c_void_p]
-    self.ctx = None
-
-  def create(self, kernel: bytes, n_lanes: int):
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    self.ctx = self.lib.wave_create(ctypes.addressof(kernel_buf), len(kernel), n_lanes)
-    self._kernel_buf = kernel_buf
-
-  def step(self) -> int: return self.lib.wave_step(self.ctx)
-  def set_sgpr(self, idx: int, val: int): self.lib.wave_set_sgpr(self.ctx, idx, val)
-  def set_vgpr(self, lane: int, idx: int, val: int): self.lib.wave_set_vgpr(self.ctx, lane, idx, val)
-  def init_lds(self, size: int): self.lib.wave_init_lds(self.ctx, size)
-
-  def get_snapshot(self) -> StateSnapshot:
-    snap = CStateSnapshot()
-    self.lib.wave_get_snapshot(self.ctx, ctypes.byref(snap))
-    return snap.to_snapshot()
-
-  def free(self):
-    if self.ctx: self.lib.wave_free(self.ctx); self.ctx = None
-
-class PythonEmulator:
-  def __init__(self):
-    self.state: WaveState | None = None
-    self.program: dict | None = None
-    self.lds: bytearray | None = None
-    self.n_lanes = 0
-
-  def create(self, kernel: bytes, n_lanes: int):
-    self.program = decode_program(kernel)
-    self.state = WaveState()
-    self.state.exec_mask = (1 << n_lanes) - 1
-    self.lds = bytearray(65536)
-    self.n_lanes = n_lanes
-
-  def step(self) -> int:
-    assert self.program is not None and self.state is not None and self.lds is not None
-    return step_wave(self.program, self.state, self.lds, self.n_lanes)
-  def set_sgpr(self, idx: int, val: int):
-    assert self.state is not None
-    self.state.sgpr[idx] = val & 0xffffffff
-  def set_vgpr(self, lane: int, idx: int, val: int):
-    assert self.state is not None
-    self.state.vgpr[lane][idx] = val & 0xffffffff
-
-  def get_snapshot(self) -> StateSnapshot:
-    assert self.state is not None
-    return StateSnapshot(pc=self.state.pc, scc=self.state.scc, vcc=self.state.vcc & 0xffffffff,
-                         exec_mask=self.state.exec_mask & 0xffffffff, sgpr=list(self.state.sgpr),
-                         vgpr=[list(self.state.vgpr[i]) for i in range(WAVE_SIZE)])
-
-def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: tuple[int, int, int],
-                      program, max_steps: int, debug: bool, trace_len: int, kernel_idx: int = 0,
-                      max_workgroups: int = 64) -> tuple[bool, str, int]:
-  """Run a single kernel through both emulators. Returns (success, message, total_steps)."""
-  gx, gy, gz = global_size
-  total_steps = 0
-  wg_count = 0
-
-  for gidz in range(gz):
-    for gidy in range(gy):
-      for gidx in range(gx):
-        if wg_count >= max_workgroups: return True, f"Completed {wg_count} workgroups (limit reached)", total_steps
-        wg_count += 1
-        rust = RustEmulator()
-        python = PythonEmulator()
-        rust.create(kernel, n_lanes)
-        python.create(kernel, n_lanes)
-
-        # Initialize LDS (64KB, standard size for AMD GPUs)
-        rust.init_lds(65536)
-
-        for emu in (rust, python):
-          emu.set_sgpr(0, args_ptr & 0xffffffff)
-          emu.set_sgpr(1, (args_ptr >> 32) & 0xffffffff)
-          emu.set_sgpr(13, gidx)
-          emu.set_sgpr(14, gidy)
-          emu.set_sgpr(15, gidz)
-
-        step = 0
-        trace: list[tuple[int, int, str, StateSnapshot, StateSnapshot]] = []
-        try:
-          while step < max_steps:
-            rust_before = rust.get_snapshot()
-            python_before = python.get_snapshot()
-
-            inst = program.get(python_before.pc)
-            inst_str = inst.disasm() if inst else f"unknown at PC={python_before.pc}"
-            trace.append((step, python_before.pc, inst_str, rust_before, python_before))
-            if len(trace) > trace_len: trace.pop(0)
-
-            if debug: print(f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: PC={python_before.pc}, inst={inst_str}")
-
-            diffs = rust_before.diff(python_before, n_lanes)
-            if diffs:
-              trace_lines = []
-              for s, pc, d, rb, pb in trace[:-1]:
-                trace_lines.append(f"    step {s}: PC={pc:3d} {d}")
-                if trace.index((s, pc, d, rb, pb)) < len(trace) - 2:
-                  next_rb, next_pb = trace[trace.index((s, pc, d, rb, pb)) + 1][3:5]
-                  inst_diffs = rb.diff(next_rb, n_lanes)
-                  if inst_diffs: trace_lines.append(f"             rust changes: {', '.join(inst_diffs[:3])}")
-              trace_str = "\n".join(trace_lines)
-              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ:\n  " + "\n  ".join(diffs[:10]) + f"\n  Recent instructions:\n{trace_str}", total_steps
-
-            rust_result = rust.step()
-            python_result = python.step()
-
-            if rust_result != python_result:
-              trace_str = "\n".join(f"    step {s}: PC={pc:3d} {d}" for s, pc, d, _, _ in trace)
-              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: different return codes: rust={rust_result}, python={python_result}, inst={inst_str}\n  Recent instructions:\n{trace_str}", total_steps
-
-            if rust_result == -1:
-              total_steps += step + 1
-              break
-            if rust_result == 1:
-              total_steps += step + 1
-              break
-            if rust_result < 0 and rust_result != -2:
-              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: error code {rust_result}", total_steps
-
-            step += 1
-          else:
-            return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Max steps ({max_steps}) reached", total_steps
-        finally:
-          rust.free()
-
-  return True, f"Completed {gx*gy*gz} workgroups", total_steps
-
-def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int, int], max_steps: int = 1000,
-                                    debug: bool = False, trace_len: int = 10, buf_data: dict[int, bytes] | None = None) -> tuple[bool, str]:
-  """Run all kernels through both emulators with shared buffer pool."""
-  from extra.assembly.rdna3.emu import set_valid_mem_ranges, decode_program
-  if buf_data is None: buf_data = {}
-
-  # Allocate shared buffer pool with padding for over-reads (GPU loads up to 16 bytes at once)
-  buf_id_to_ptr: dict[int, int] = {}
-  buffers = []
-  for buf_id, size in buf_pool.items():
-    padded_size = ((size + 15) // 16) * 16 + 16  # round up to 16 bytes + extra padding
-    # Initialize with data from COPY if available
-    init_data = buf_data.get(buf_id, b'\x00' * padded_size)
-    init_list = list(init_data) + [0] * (padded_size - len(init_data))
-    buf = (ctypes.c_uint8 * padded_size)(*init_list[:padded_size])
-    buffers.append((buf, padded_size))
-    buf_id_to_ptr[buf_id] = ctypes.addressof(buf)
-
-  # Set up valid memory ranges
-  ranges = {(ctypes.addressof(b), size) for b, size in buffers}
-
-  total_steps = 0
-  for ki, kernel in enumerate(kernels):
-    # Create args array for this kernel's buffers
-    args = (ctypes.c_uint64 * len(kernel.buf_idxs))(*[buf_id_to_ptr[bid] for bid in kernel.buf_idxs])
-    args_ptr = ctypes.addressof(args)
-
-    # Update valid ranges to include this args array
-    kernel_ranges = ranges | {(args_ptr, ctypes.sizeof(args))}
-    set_valid_mem_ranges(kernel_ranges)
-
-    program = decode_program(kernel.code)
-    n_lanes = kernel.local_size[0] * kernel.local_size[1] * kernel.local_size[2]
-
-    ok, msg, steps = run_single_kernel(
-      kernel.code, min(n_lanes, 32), args_ptr, kernel.global_size,
-      program, max_steps, debug, trace_len, ki
-    )
-    total_steps += steps
-    if not ok:
-      return False, msg
-
-  return True, f"Completed {len(kernels)} kernels, {total_steps} total steps"
-
-def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list, max_steps: int = 1000, debug: bool = False,
-                                   global_size: tuple[int, int, int] = (1, 1, 1), trace_len: int = 10) -> tuple[bool, str]:
-  """Run both emulators with memory set up for tinygrad kernels, executing all workgroups. Legacy wrapper."""
-  from extra.assembly.rdna3.emu import set_valid_mem_ranges, decode_program
-
-  # Allocate buffers
-  buffers = []
-  for size in buf_sizes:
-    buf = (ctypes.c_uint8 * size)(*[0] * size)
-    buffers.append(buf)
-
-  # Create args array with buffer pointers
-  args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
-  args_ptr = ctypes.addressof(args)
-
-  # Set up valid memory ranges for Python emulator
-  ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
-  ranges.add((args_ptr, ctypes.sizeof(args)))
-  set_valid_mem_ranges(ranges)
-
-  program = decode_program(kernel)
-  ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, program, max_steps, debug, trace_len)
-  return ok, msg
-
-def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelInfo], dict[int, int], dict[int, bytes]]:
-  """Compile a tinygrad operation and extract all kernels with their buffer mappings."""
-  from tinygrad import Tensor
-  from tinygrad.runtime.support.elf import elf_loader
-
-  out = op_fn(Tensor)
-  sched = out.schedule()
-  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(KernelInfo(
-              code=bytes(sec.content),
-              global_size=tuple(lowered.prg.p.global_size),
-              local_size=tuple(lowered.prg.p.local_size),
-              buf_idxs=buf_idxs,
-              buf_sizes=buf_sizes
-            ))
-  if not kernels: raise RuntimeError("No kernel found")
-  return kernels, buf_pool, buf_data
-
-def get_kernel_from_tinygrad(op_fn) -> tuple[bytes, tuple[int, int, int], tuple[int, int, int], list]:
-  """Compile a tinygrad operation and extract the last (main) kernel binary. Legacy wrapper."""
-  kernels, _, _ = get_kernels_from_tinygrad(op_fn)
-  k = kernels[-1]
-  return k.code, k.global_size, k.local_size, k.buf_sizes
-
-@unittest.skipUnless(REMU_PATH.exists(), "libremu.so not found")
-class TestTinygradKernels(unittest.TestCase):
-  """Compare emulators on real tinygrad-compiled kernels."""
-
-  def _test_kernel(self, op_fn, max_steps=10000):
-    kernels, buf_pool, buf_data = get_kernels_from_tinygrad(op_fn)
-    ok, msg = compare_emulators_multi_kernel(kernels, buf_pool, max_steps=max_steps, buf_data=buf_data)
-    self.assertTrue(ok, msg)
-
-  # Basic unary ops
-  def test_neg(self): self._test_kernel(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
-  def test_relu(self): self._test_kernel(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
-  def test_exp(self): self._test_kernel(lambda T: T([0.0, 1.0, 2.0]).exp())
-  def test_log(self): self._test_kernel(lambda T: T([1.0, 2.0, 3.0]).log())
-  def test_sin(self): self._test_kernel(lambda T: T([0.0, 1.0, 2.0]).sin())
-  def test_sqrt(self): self._test_kernel(lambda T: T([1.0, 4.0, 9.0]).sqrt())
-  def test_recip(self): self._test_kernel(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
-
-  # Binary ops
-  def test_add(self): self._test_kernel(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
-  def test_sub(self): self._test_kernel(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
-  def test_mul(self): self._test_kernel(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
-  def test_div(self): self._test_kernel(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
-  def test_max_binary(self): self._test_kernel(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
-
-  # Reductions
-  def test_sum_reduce(self): self._test_kernel(lambda T: T.empty(64).sum())
-  def test_max_reduce(self): self._test_kernel(lambda T: T.empty(64).max())
-  def test_mean_reduce(self): self._test_kernel(lambda T: T.empty(32).mean())
-
-  # Matmul - various sizes
-  def test_gemm_4x4(self): self._test_kernel(lambda T: T.empty(4, 4) @ T.empty(4, 4), max_steps=100000)
-  def test_gemm_8x8(self): self._test_kernel(lambda T: T.empty(8, 8) @ T.empty(8, 8), max_steps=200000)
-  @unittest.skip("too slow")
-  def test_gemm_16x16(self): self._test_kernel(lambda T: T.empty(16, 16) @ T.empty(16, 16), max_steps=500000)
-  def test_gemv(self): self._test_kernel(lambda T: T.empty(1, 16) @ T.empty(16, 16), max_steps=100000)
-
-  # Complex ops
-  def test_softmax(self): self._test_kernel(lambda T: T.empty(16).softmax())
-  def test_layernorm(self): self._test_kernel(lambda T: T.empty(8, 8).layernorm())
-
-  # Memory patterns
-  def test_contiguous(self): self._test_kernel(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
-  def test_reshape(self): self._test_kernel(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
-  def test_expand(self): self._test_kernel(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
-
-  # Cast ops
-  def test_cast_int(self): self._test_kernel(lambda T: T.empty(16).int().float())
-  def test_cast_half(self): self._test_kernel(lambda T: T.empty(16).half().float())
-
-  # Min/max (uses comparison internally)
-  def test_min_binary(self): self._test_kernel(lambda T: T([1.0, 5.0, 3.0]).minimum(T([3.0, 2.0, 4.0])))
-
-  # Comparison ops (test VOPC instructions) - use 32+ elements to force vector instructions
-  def test_cmp_lt(self): self._test_kernel(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
-  def test_cmp_eq(self): self._test_kernel(lambda T: (T.empty(64) == T.empty(64)).where(T.empty(64), T.empty(64)))
-  def test_where(self): self._test_kernel(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
-
-  # Bitwise ops
-  def test_bitwise_and(self): self._test_kernel(lambda T: T([0xF0, 0x0F, 0xFF]).int() & T([0x0F, 0x0F, 0x00]).int())
-  def test_bitwise_or(self): self._test_kernel(lambda T: T([0xF0, 0x0F, 0x00]).int() | T([0x0F, 0x0F, 0xFF]).int())
-  def test_bitwise_xor(self): self._test_kernel(lambda T: T([0xFF, 0x0F, 0xF0]).int() ^ T([0x0F, 0xF0, 0xF0]).int())
-
-  # Integer ops - use 32+ elements to force vector instructions
-  def test_int_add(self): self._test_kernel(lambda T: (T.empty(64).int() + T.empty(64).int()).float())
-  def test_int_mul(self): self._test_kernel(lambda T: (T.empty(64).int() * T.empty(64).int()).float())
-  def test_int_mod(self): self._test_kernel(lambda T: (T.empty(64).int().abs() % (T.empty(64).int().abs() + 1)).float())
-
-  # More math ops - use 32+ elements to force vector instructions
-  def test_abs(self): self._test_kernel(lambda T: T.empty(64).abs())
-  def test_floor(self): self._test_kernel(lambda T: T.empty(64).floor())
-  def test_ceil(self): self._test_kernel(lambda T: T.empty(64).ceil())
-  def test_trunc(self): self._test_kernel(lambda T: T.empty(64).trunc())
-
-  # Fused ops
-  def test_fma(self): self._test_kernel(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
-
-  # Argmax/argmin (tests different reduction pattern) - use 32+ elements to force vector instructions
-  def test_argmax(self): self._test_kernel(lambda T: T.empty(64).argmax())
-  def test_argmin(self): self._test_kernel(lambda T: T.empty(64).argmin())
-
-  # Exact value tests - use 32+ elements to force vector instructions (small tensors use scalar ops which Rust emu doesn't fully support)
-  def test_abs_exact(self): self._test_kernel(lambda T: T([-1., 0., 1.]*11).abs())  # 33 elements
-  def test_neg_exact(self): self._test_kernel(lambda T: -T([-1., 0., 1.]*11))
-  def test_log_special(self): self._test_kernel(lambda T: T([1., 2., 0.5]*11).log())
-  def test_exp_exact(self): self._test_kernel(lambda T: T([0., 1., -1.]*11).exp())
-  def test_reciprocal_exact(self): self._test_kernel(lambda T: T([1., 2., 0.5]*11).reciprocal())
-
-  # Integer division and mod - use 32+ elements
-  def test_int_div(self): self._test_kernel(lambda T: (T([10, 20, 30]*11).int() // T([3, 4, 5]*11).int()).float())
-  def test_int_neg(self): self._test_kernel(lambda T: (-T([1, -2, 3]*11).int()).float())
-
-  # Mixed precision - use 32+ elements
-  def test_half_add(self): self._test_kernel(lambda T: (T([1., 2.]*16).half() + T([3., 4.]*16).half()).float())
-  def test_half_mul(self): self._test_kernel(lambda T: (T([2., 3.]*16).half() * T([4., 5.]*16).half()).float())
-
-  # Matrix ops - patterns from test_ops.py failures
-  def test_cat(self): self._test_kernel(lambda T: T.empty(32, 64).cat(T.empty(32, 64), dim=1))
-  def test_gather(self): self._test_kernel(lambda T: T.empty(64).gather(0, T.arange(32).int()))
-
-  # Tests from test_ops.py that are failing
-  def test_permute(self): self._test_kernel(lambda T: T.empty(3, 4, 5, 6).permute((3, 2, 1, 0)).contiguous())
-  def test_cat_large(self): self._test_kernel(lambda T: T.empty(45, 65, 9).cat(T.empty(45, 65, 9), T.empty(45, 65, 9), dim=1))
-  def test_gather_small(self): self._test_kernel(lambda T: T.empty(10).gather(0, T.arange(5).int()))
-  @unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
-  def test_cross_entropy(self): self._test_kernel(lambda T: T.randn(32, 10).softmax().log().sum())
-  def test_cross_entropy_class(self):
-    import numpy as np
-    np.random.seed(0)
-    classes = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
-    x_np = np.random.randn(32, 10).astype(np.float32)
-    self._test_kernel(lambda T: (T(x_np.tolist()).reshape(32,10) + 0).cross_entropy((T(classes).int().reshape(32) + 0)))
-
-  # Regression tests for BFE operations with width=0 (walrus operator bug)
-  def test_topk(self): self._test_kernel(lambda T: T.empty(64).topk(3)[0])
-  def test_interpolate_uint8(self): self._test_kernel(lambda T: T.empty(2,3,64,64).relu().cast('uint8').interpolate((10,10), mode="linear"))
-
-  # Regression test for 64-bit comparison (V_CMP_GT_I64, V_CMP_LT_U64, etc.) with rsrc64
-  def test_index_int64(self):
-    from tinygrad import dtypes
-    self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
-
-  @unittest.skip("only works with mock GPU")
-  def test_index_int64_2d(self):
-    from tinygrad import dtypes
-    # Tests 64-bit compare with inline constants (comparing against 0)
-    self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), T.arange(4).cast(dtypes.int64)])
-
-  # Pooling operations - regression test for VCC wave32 mode (S_CBRANCH_VCCZ should only check VCC_LO)
-  def test_avg_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4), stride=2))
-  @unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
-  def test_avg_pool3d(self):
-    import numpy as np
-    np.random.seed(0)
-    self._test_kernel(lambda T: T(np.random.randn(1, 1, 16, 16, 16).astype(np.float32).tolist()).avg_pool2d(kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False))
-  def test_max_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4), stride=2))
-
-  # Convolution operations - multi-kernel tests
-  def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 4, 8, 8).conv2d(T.empty(4, 4, 3, 3)), max_steps=100000)
-  def test_conv_transpose2d(self): self._test_kernel(lambda T: T.empty(1, 4, 8, 8).conv_transpose2d(T.empty(4, 4, 3, 3)), max_steps=200000)
-  @unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
-  def test_conv_transpose3d(self):
-    import numpy as np
-    np.random.seed(0)
-    self._test_kernel(lambda T: T(np.random.randn(2, 4, 9, 9, 9).astype(np.float32).tolist()).conv_transpose2d(
-      T(np.random.randn(4, 4, 3, 3, 3).astype(np.float32).tolist())), max_steps=500000)
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_emu.py

@@ -1,845 +0,0 @@
-# Unit tests for RDNA3 Python emulator
-import unittest
-import ctypes
-import struct
-import math
-from extra.assembly.rdna3.emu import (
-  WaveState, decode_program, exec_wave, exec_workgroup, run_asm,
-  i32, f32, sext, WAVE_SIZE, set_valid_mem_ranges
-)
-from extra.assembly.rdna3.autogen import *
-from extra.assembly.rdna3.lib import RawImm
-
-def run_kernel(kernel: bytes, n_threads: int = 1, n_outputs: int = 1) -> list[int]:
-  """Helper to run a kernel and return output values."""
-  output = (ctypes.c_uint32 * (n_threads * n_outputs))(*[0xdead] * (n_threads * n_outputs))
-  output_ptr = ctypes.addressof(output)
-  args = (ctypes.c_uint64 * 1)(output_ptr)
-  args_ptr = ctypes.addressof(args)
-  kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-  kernel_ptr = ctypes.addressof(kernel_buf)
-  # Register valid memory ranges for bounds checking
-  set_valid_mem_ranges({
-    (output_ptr, ctypes.sizeof(output)),
-    (args_ptr, ctypes.sizeof(args)),
-    (kernel_ptr, len(kernel)),
-  })
-  result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, n_threads, 1, 1, args_ptr)
-  assert result == 0, f"run_asm failed with {result}"
-  return [output[i] for i in range(n_threads * n_outputs)]
-
-def make_store_kernel(setup_instrs: list, store_vreg: int = 1) -> bytes:
-  """Create a kernel that runs setup instructions then stores v[store_vreg] to output[tid]."""
-  kernel = b''
-  # Load output pointer
-  kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
-  kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-  # Run setup instructions
-  for instr in setup_instrs:
-    kernel += instr.to_bytes()
-  # Compute offset: v3 = tid * 4
-  kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
-  # Store result
-  kernel += global_store_b32(addr=v[3], data=v[store_vreg], saddr=s[2]).to_bytes()
-  kernel += s_endpgm().to_bytes()
-  return kernel
-
-class TestScalarOps(unittest.TestCase):
-  def test_s_mov_b32(self):
-    state = WaveState()
-    kernel = s_mov_b32(s[5], 42).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[5], 42)
-
-  def test_s_add_u32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 100, 50
-    kernel = s_add_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 150)
-    self.assertEqual(state.scc, 0)  # no carry
-
-  def test_s_add_u32_carry(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 0xffffffff, 1
-    kernel = s_add_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 0)
-    self.assertEqual(state.scc, 1)  # carry
-
-  def test_s_sub_u32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 100, 30
-    kernel = s_sub_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 70)
-    self.assertEqual(state.scc, 0)  # no borrow
-
-  def test_s_and_b32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 0xff00, 0x0ff0
-    kernel = s_and_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 0x0f00)
-
-  def test_s_or_b32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 0xff00, 0x00ff
-    kernel = s_or_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 0xffff)
-
-  def test_s_lshl_b32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 1, 4
-    kernel = s_lshl_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 16)
-
-  def test_s_lshr_b32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 256, 4
-    kernel = s_lshr_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 16)
-
-  def test_s_mul_i32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 7, 6
-    kernel = s_mul_i32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 42)
-
-  def test_s_cmp_eq_u32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 42, 42
-    kernel = s_cmp_eq_u32(s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.scc, 1)
-
-  def test_s_cmp_lg_u32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 42, 43
-    kernel = s_cmp_lg_u32(s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.scc, 1)
-
-class TestVectorOps(unittest.TestCase):
-  def test_v_mov_b32(self):
-    kernel = make_store_kernel([v_mov_b32_e32(v[1], 42)])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [42])
-
-  def test_v_add_nc_u32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 10),
-      v_mov_b32_e32(v[2], 32),
-      v_add_nc_u32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [42])
-
-  def test_v_sub_nc_u32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 50),
-      v_mov_b32_e32(v[2], 8),
-      v_sub_nc_u32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [42])
-
-  def test_v_mul_lo_u32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 6),
-      v_mov_b32_e32(v[2], 7),
-      v_mul_lo_u32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [42])
-
-  def test_v_and_b32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 0xff0f),
-      v_mov_b32_e32(v[2], 0x0fff),
-      v_and_b32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [0x0f0f])
-
-  def test_v_or_b32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 0xf000),
-      v_mov_b32_e32(v[2], 0x000f),
-      v_or_b32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [0xf00f])
-
-  def test_v_lshlrev_b32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 1),
-      v_lshlrev_b32_e32(v[1], 5, v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [32])
-
-  def test_v_lshrrev_b32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 128),
-      v_lshrrev_b32_e32(v[1], 3, v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [16])
-
-  def test_v_add_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(1.5)),
-      v_mov_b32_e32(v[2], i32(2.5)),
-      v_add_f32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 4.0)
-
-  def test_v_mul_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(3.0)),
-      v_mov_b32_e32(v[2], i32(4.0)),
-      v_mul_f32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 12.0)
-
-  def test_v_max_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(3.0)),
-      v_mov_b32_e32(v[2], i32(5.0)),
-      v_max_f32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 5.0)
-
-  def test_v_min_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(3.0)),
-      v_mov_b32_e32(v[2], i32(5.0)),
-      v_min_f32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 3.0)
-
-class TestThreading(unittest.TestCase):
-  def test_thread_id(self):
-    """Each thread should get its own thread ID in v0."""
-    kernel = make_store_kernel([v_mov_b32_e32(v[1], v[0])], store_vreg=1)
-    out = run_kernel(kernel, n_threads=4)
-    self.assertEqual(out, [0, 1, 2, 3])
-
-  def test_thread_local_ops(self):
-    """Each thread computes tid * 10."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[2], 10),
-      v_mul_lo_u32(v[1], v[0], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=4)
-    self.assertEqual(out, [0, 10, 20, 30])
-
-  def test_exec_mask(self):
-    """Test that exec mask controls which lanes execute."""
-    kernel = b''
-    kernel += s_load_b64(s[2:3], s[0:1], 0, soffset=NULL).to_bytes()
-    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-    kernel += v_mov_b32_e32(v[1], 100).to_bytes()  # default value
-    kernel += s_mov_b32(EXEC_LO, 0b0101).to_bytes()  # only lanes 0 and 2
-    kernel += v_mov_b32_e32(v[1], 42).to_bytes()   # only for active lanes
-    kernel += s_mov_b32(EXEC_LO, 0xf).to_bytes()   # restore all lanes
-    kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
-    kernel += global_store_b32(addr=v[3], data=v[1], saddr=s[2]).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    out = run_kernel(kernel, n_threads=4)
-    self.assertEqual(out, [42, 100, 42, 100])
-
-class TestBranching(unittest.TestCase):
-  def test_s_branch(self):
-    """Test unconditional branch."""
-    state = WaveState()
-    kernel = b''
-    kernel += s_mov_b32(s[0], 1).to_bytes()
-    kernel += s_branch(1).to_bytes()  # skip next instruction
-    kernel += s_mov_b32(s[0], 2).to_bytes()  # should be skipped
-    kernel += s_mov_b32(s[1], 3).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[0], 1)  # not overwritten
-    self.assertEqual(state.sgpr[1], 3)
-
-  def test_s_cbranch_scc0(self):
-    """Test conditional branch on SCC=0."""
-    state = WaveState()
-    state.scc = 0
-    kernel = b''
-    kernel += s_mov_b32(s[0], 1).to_bytes()
-    kernel += s_cbranch_scc0(1).to_bytes()  # branch if scc=0
-    kernel += s_mov_b32(s[0], 2).to_bytes()  # should be skipped
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[0], 1)
-
-  def test_s_cbranch_scc1(self):
-    """Test conditional branch on SCC=1."""
-    state = WaveState()
-    state.scc = 1
-    kernel = b''
-    kernel += s_mov_b32(s[0], 1).to_bytes()
-    kernel += s_cbranch_scc1(1).to_bytes()  # branch if scc=1
-    kernel += s_mov_b32(s[0], 2).to_bytes()  # should be skipped
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[0], 1)
-
-  def test_unknown_sopp_opcode(self):
-    """Regression test: unknown SOPP opcodes should be ignored, not crash."""
-    state = WaveState()
-    # Create a raw SOPP instruction with opcode 8 (undefined in our enum)
-    # SOPP format: bits[31:23] = 0b101111111, bits[22:16] = op, bits[15:0] = simm16
-    unknown_sopp = (0b101111111 << 23) | (8 << 16) | 0  # op=8, simm16=0
-    kernel = unknown_sopp.to_bytes(4, 'little') + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    # Should not raise an exception
-    exec_wave(prog, state, bytearray(65536), 1)
-
-class TestMemory(unittest.TestCase):
-  def test_global_load_store(self):
-    """Test global load followed by store."""
-    # Create input buffer
-    input_buf = (ctypes.c_uint32 * 4)(10, 20, 30, 40)
-    input_ptr = ctypes.addressof(input_buf)
-    output_buf = (ctypes.c_uint32 * 4)(*[0]*4)
-    output_ptr = ctypes.addressof(output_buf)
-    args = (ctypes.c_uint64 * 2)(output_ptr, input_ptr)
-    args_ptr = ctypes.addressof(args)
-
-    # Kernel: load from input[tid], add 1, store to output[tid]
-    kernel = b''
-    kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()  # output ptr
-    kernel += s_load_b64(s[4:5], s[0:1], soffset=NULL, offset=8).to_bytes()  # input ptr
-    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-    kernel += v_lshlrev_b32_e32(v[2], 2, v[0]).to_bytes()  # offset = tid * 4
-    kernel += global_load_b32(vdst=v[1], addr=v[2], saddr=s[4]).to_bytes()
-    kernel += s_waitcnt(vmcnt=0).to_bytes()
-    kernel += v_add_nc_u32_e32(v[1], 1, v[1]).to_bytes()  # add 1
-    kernel += global_store_b32(addr=v[2], data=v[1], saddr=s[2]).to_bytes()
-    kernel += s_endpgm().to_bytes()
-
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    kernel_ptr = ctypes.addressof(kernel_buf)
-    set_valid_mem_ranges({
-      (input_ptr, ctypes.sizeof(input_buf)),
-      (output_ptr, ctypes.sizeof(output_buf)),
-      (args_ptr, ctypes.sizeof(args)),
-      (kernel_ptr, len(kernel)),
-    })
-    result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, 4, 1, 1, args_ptr)
-    self.assertEqual(result, 0)
-    self.assertEqual([output_buf[i] for i in range(4)], [11, 21, 31, 41])
-
-class TestFloatOps(unittest.TestCase):
-  def test_v_rcp_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(4.0)),
-      v_rcp_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertAlmostEqual(f32(out[0]), 0.25, places=5)
-
-  def test_v_sqrt_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(16.0)),
-      v_sqrt_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertAlmostEqual(f32(out[0]), 4.0, places=5)
-
-  def test_v_floor_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(3.7)),
-      v_floor_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 3.0)
-
-  def test_v_ceil_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(3.2)),
-      v_ceil_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 4.0)
-
-  def test_v_cvt_f32_i32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 42),
-      v_cvt_f32_i32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 42.0)
-
-  def test_v_cvt_i32_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(42.9)),
-      v_cvt_i32_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0], 42)
-
-class TestVOP3(unittest.TestCase):
-  def test_v_fma_f32(self):
-    """Test fused multiply-add: a*b + c"""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(2.0)),
-      v_mov_b32_e32(v[2], i32(3.0)),
-      v_mov_b32_e32(v[4], i32(4.0)),
-      v_fma_f32(v[1], v[1], v[2], v[4]),  # 2*3+4 = 10
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 10.0)
-
-  def test_v_add3_u32(self):
-    """Test 3-operand add."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 10),
-      v_mov_b32_e32(v[2], 20),
-      v_mov_b32_e32(v[4], 12),
-      v_add3_u32(v[1], v[1], v[2], v[4]),  # 10+20+12 = 42
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0], 42)
-
-  def test_v_neg_modifier(self):
-    """Test VOP3 negation modifier."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(5.0)),
-      v_mov_b32_e32(v[2], i32(3.0)),
-      # v_add_f32 with neg on src1: 5 + (-3) = 2
-      v_add_f32_e64(v[1], v[1], v[2], neg=0b010),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 2.0)
-
-  def test_v_ldexp_f32(self):
-    """Regression test: V_LDEXP_F32 used by exp()."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(1.5)),
-      v_mov_b32_e32(v[2], 3),  # exponent
-      v_ldexp_f32(v[1], v[1], v[2]),  # 1.5 * 2^3 = 12.0
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 12.0)
-
-  def test_v_xad_u32(self):
-    """Regression test: V_XAD_U32 (xor-add) used by random number generation."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 3),
-      v_mov_b32_e32(v[2], 4),
-      v_mov_b32_e32(v[4], 5),
-      v_xad_u32(v[1], v[1], v[2], v[4]),  # (3^4)+5 = 7+5 = 12
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0], 12)
-
-  def test_v_lshl_or_b32(self):
-    """Regression test: V_LSHL_OR_B32 operand order is (s0 << s1) | s2, not (s0 << s2) | s1."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 5),   # s0 = value to shift
-      v_mov_b32_e32(v[2], 2),   # s1 = shift amount
-      v_mov_b32_e32(v[4], 3),   # s2 = value to OR
-      v_lshl_or_b32(v[1], v[1], v[2], v[4]),  # (5 << 2) | 3 = 20 | 3 = 23
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0], 23)
-
-  def test_v_sqrt_f32_negative(self):
-    """Regression test: V_SQRT_F32 should return NaN for negative inputs, not 0."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(-1.0)),
-      v_sqrt_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertTrue(math.isnan(f32(out[0])))
-
-  def test_v_rsq_f32_negative(self):
-    """Regression test: V_RSQ_F32 should return NaN for negative inputs, not inf."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(-1.0)),
-      v_rsq_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertTrue(math.isnan(f32(out[0])))
-
-class TestVOPD(unittest.TestCase):
-  def test_vopd_add_nc_u32(self):
-    """Test VOPD V_DUAL_ADD_NC_U32."""
-    state = WaveState()
-    state.vgpr[0][1] = 100
-    state.vgpr[0][2] = 50
-    # vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
-    kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
-                  opy=VOPDOp.V_DUAL_ADD_NC_U32, srcy0=v[1], vsrcy1=VGPR(2), vdsty=VGPR(4)).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.vgpr[0][3], 100)  # MOV result
-    self.assertEqual(state.vgpr[0][4], 150)  # 100 + 50
-
-  def test_vopd_lshlrev(self):
-    """Test VOPD V_DUAL_LSHLREV_B32."""
-    state = WaveState()
-    state.vgpr[0][1] = 0x10
-    state.vgpr[0][2] = 0
-    # vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
-    kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
-                  opy=VOPDOp.V_DUAL_LSHLREV_B32, srcy0=4, vsrcy1=VGPR(1), vdsty=VGPR(4)).to_bytes()  # V4 = V1 << 4
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.vgpr[0][3], 0x10)  # MOV result
-    self.assertEqual(state.vgpr[0][4], 0x100)  # 0x10 << 4 = 0x100
-
-  def test_vopd_and(self):
-    """Test VOPD V_DUAL_AND_B32."""
-    state = WaveState()
-    state.vgpr[0][1] = 0xff
-    state.vgpr[0][2] = 0x0f
-    # vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
-    kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
-                  opy=VOPDOp.V_DUAL_AND_B32, srcy0=v[1], vsrcy1=VGPR(2), vdsty=VGPR(4)).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.vgpr[0][3], 0xff)
-    self.assertEqual(state.vgpr[0][4], 0x0f)  # 0xff & 0x0f = 0x0f
-
-  def test_vopd_parallel_read(self):
-    """Regression: VOPD must read all inputs before writing - Y op reads register that X op writes."""
-    state = WaveState()
-    state.vgpr[0][4] = 0
-    state.vgpr[0][7] = 5  # Y op reads v7 as vsrcy1, X op writes to v7
-    # X: MOV v7, v0 (v0=0, so v7 becomes 0)
-    # Y: ADD v6, v4, v7 (should use original v7=5, not the overwritten 0)
-    # vdsty_enc=3 with vdstx=7 (odd) -> vdsty = (3 << 1) | (7&1)^1 = 6 | 0 = 6
-    kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[0], vsrcx1=VGPR(0), vdstx=VGPR(7),
-                  opy=VOPDOp.V_DUAL_ADD_NC_U32, srcy0=v[4], vsrcy1=VGPR(7), vdsty=VGPR(6)).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.vgpr[0][7], 0)  # X op: v7 = v0 = 0
-    self.assertEqual(state.vgpr[0][6], 5)  # Y op: v6 = v4 + v7 = 0 + 5 (original v7)
-
-class TestDecoder(unittest.TestCase):
-  def test_vopd_literal_handling(self):
-    """Regression test: VOPD srcx0/srcy0 with literal (255) wasn't consuming the literal dword."""
-    state = WaveState()
-    # Create VOPD with srcx0=255 (literal), followed by literal value 0x12345678
-    vopd_bytes = VOPD(opx=8, srcx0=RawImm(255), vsrcx1=VGPR(0), vdstx=VGPR(1),  # MOV: V1 = literal
-                      opy=8, srcy0=RawImm(128), vsrcy1=VGPR(0), vdsty=VGPR(2)).to_bytes()  # MOV: V2 = 0
-    literal_bytes = (0x12345678).to_bytes(4, 'little')
-    kernel = vopd_bytes + literal_bytes + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    # Should decode as 3 instructions: VOPD (with literal), then S_ENDPGM
-    # The literal should NOT be decoded as a separate instruction
-    self.assertEqual(len(prog), 2)  # VOPD + S_ENDPGM
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.vgpr[0][1], 0x12345678)
-
-  def test_s_endpgm_stops_decode(self):
-    """Regression test: decoder should stop at S_ENDPGM, not read past into metadata."""
-    # Create a kernel followed by garbage that looks like an invalid instruction
-    kernel = s_mov_b32(s[0], 42).to_bytes() + s_endpgm().to_bytes()
-    garbage = bytes([0xff] * 16)  # garbage after kernel
-    prog = decode_program(kernel + garbage)
-    # Should only have 2 instructions (s_mov_b32 and s_endpgm)
-    self.assertEqual(len(prog), 2)
-
-class TestFloatConversion(unittest.TestCase):
-  """Unit tests for i32/i16/f32/f16 float conversion functions."""
-
-  def test_i32_preserves_nan_sign(self):
-    """NaN sign bit should be preserved when converting float to int bits."""
-    from extra.assembly.rdna3.emu import i32, f32
-    # 0 * -inf produces a negative NaN
-    neg_nan = 0.0 * float('-inf')
-    bits = i32(neg_nan)
-    # Should have sign bit set (0xffc00000), not canonical positive NaN (0x7fc00000)
-    self.assertEqual(bits & 0x80000000, 0x80000000, f"Expected negative NaN, got 0x{bits:08x}")
-    self.assertTrue(math.isnan(f32(bits)))
-
-  def test_i32_preserves_positive_nan(self):
-    """Positive NaN should remain positive."""
-    from extra.assembly.rdna3.emu import i32, f32
-    pos_nan = float('nan')
-    bits = i32(pos_nan)
-    # Standard Python NaN is positive (0x7fc00000)
-    self.assertEqual(bits & 0x80000000, 0, f"Expected positive NaN, got 0x{bits:08x}")
-    self.assertTrue(math.isnan(f32(bits)))
-
-  def test_i32_overflow_to_inf(self):
-    """Values too large for f32 should become inf."""
-    from extra.assembly.rdna3.emu import i32, f32
-    big = 2.0 ** 200
-    self.assertEqual(i32(big), 0x7f800000)  # +inf
-    self.assertEqual(i32(-big), 0xff800000)  # -inf
-
-  def test_i32_inf(self):
-    """Infinity should be preserved."""
-    from extra.assembly.rdna3.emu import i32
-    self.assertEqual(i32(float('inf')), 0x7f800000)
-    self.assertEqual(i32(float('-inf')), 0xff800000)
-
-  def test_i32_normal_values(self):
-    """Normal float values should round-trip correctly (within f32 precision)."""
-    from extra.assembly.rdna3.emu import i32, f32
-    # Use values exactly representable in float32
-    for val in [0.0, 1.0, -1.0, 0.5, -0.5, 100.0, -100.0, 1e10]:
-      bits = i32(val)
-      self.assertAlmostEqual(f32(bits), val, places=5)
-
-  def test_i16_overflow_to_inf(self):
-    """Values too large for f16 should become inf."""
-    from extra.assembly.rdna3.emu import i16
-    big = 100000.0  # way larger than f16 max (65504)
-    self.assertEqual(i16(big), 0x7c00)  # +inf
-    self.assertEqual(i16(-big), 0xfc00)  # -inf
-
-  def test_i16_inf(self):
-    """Infinity should be preserved."""
-    from extra.assembly.rdna3.emu import i16
-    self.assertEqual(i16(float('inf')), 0x7c00)
-    self.assertEqual(i16(float('-inf')), 0xfc00)
-
-  def test_fma_nan_sign_preserved(self):
-    """FMA producing NaN should preserve the correct sign bit."""
-    from extra.assembly.rdna3.emu import i32, f32
-    # 0 * (-inf) + 1.0 = NaN (from 0 * -inf)
-    a, b, c = 0.0, float('-inf'), 1.0
-    result = i32(a * b + c)
-    # The NaN should be negative since 0 * -inf produces negative NaN
-    self.assertEqual(result & 0x80000000, 0x80000000, f"Expected negative NaN, got 0x{result:08x}")
-
-class TestMultiWave(unittest.TestCase):
-  def test_all_waves_execute(self):
-    """Regression test: all waves in a workgroup must execute, not just the first."""
-    n_threads = 64  # 2 waves of 32 threads each
-    output = (ctypes.c_uint32 * n_threads)(*[0xdead] * n_threads)
-    output_ptr = ctypes.addressof(output)
-    args = (ctypes.c_uint64 * 1)(output_ptr)
-    args_ptr = ctypes.addressof(args)
-
-    # Simple kernel: store tid to output[tid]
-    kernel = b''
-    kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
-    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-    kernel += v_lshlrev_b32_e32(v[1], 2, v[0]).to_bytes()  # offset = tid * 4
-    kernel += global_store_b32(addr=v[1], data=v[0], saddr=s[2]).to_bytes()
-    kernel += s_endpgm().to_bytes()
-
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    kernel_ptr = ctypes.addressof(kernel_buf)
-    set_valid_mem_ranges({
-      (output_ptr, ctypes.sizeof(output)),
-      (args_ptr, ctypes.sizeof(args)),
-      (kernel_ptr, len(kernel)),
-    })
-    result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, n_threads, 1, 1, args_ptr)
-    self.assertEqual(result, 0)
-    # All threads should have written their tid
-    for i in range(n_threads):
-      self.assertEqual(output[i], i, f"Thread {i} didn't execute")
-
-class TestRegressions(unittest.TestCase):
-  """Regression tests for bugs fixed in the emulator."""
-
-  def test_v_fmac_f16(self):
-    """V_FMAC_F16: fused multiply-add for FP16. Regression for VOP2 op 54."""
-    from extra.assembly.rdna3.emu import i16, f16
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i16(2.0)),  # v1.lo = 2.0 (fp16)
-      v_mov_b32_e32(v[2], i16(3.0)),  # v2.lo = 3.0 (fp16)
-      # v1 = v1 * v2 + v1 = 2.0 * 3.0 + 2.0 = 8.0
-      VOP2(VOP2Op.V_FMAC_F16, v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertAlmostEqual(f16(out[0] & 0xffff), 8.0, places=2)
-
-  def test_v_cvt_f64_f32(self):
-    """V_CVT_F64_F32: convert float32 to float64. Regression for VOP1 op 16."""
-    kernel = b''
-    kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
-    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-    kernel += v_mov_b32_e32(v[1], i32(3.14159)).to_bytes()
-    kernel += VOP1(VOP1Op.V_CVT_F64_F32, v[4], v[1]).to_bytes()  # v4:v5 = f64(v1)
-    kernel += v_lshlrev_b32_e32(v[3], 3, v[0]).to_bytes()  # offset = tid * 8
-    kernel += global_store_b64(addr=v[3], data=v[4], saddr=s[2]).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    output = (ctypes.c_double * 1)(0.0)
-    output_ptr = ctypes.addressof(output)
-    args = (ctypes.c_uint64 * 1)(output_ptr)
-    args_ptr = ctypes.addressof(args)
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    kernel_ptr = ctypes.addressof(kernel_buf)
-    set_valid_mem_ranges({(output_ptr, 8), (args_ptr, 8), (kernel_ptr, len(kernel))})
-    run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
-    self.assertAlmostEqual(output[0], 3.14159, places=4)
-
-  def test_v_add_f64(self):
-    """V_ADD_F64: add two float64 values. Regression for VOP3 op 807."""
-    from extra.assembly.rdna3.emu import i64_parts
-    kernel = b''
-    kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
-    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-    # Load 1.5 into v1:v2
-    lo, hi = i64_parts(1.5)
-    kernel += v_mov_b32_e32(v[1], lo).to_bytes()
-    kernel += v_mov_b32_e32(v[2], hi).to_bytes()
-    # Load 2.5 into v3:v4
-    lo, hi = i64_parts(2.5)
-    kernel += v_mov_b32_e32(v[3], lo).to_bytes()
-    kernel += v_mov_b32_e32(v[4], hi).to_bytes()
-    # v5:v6 = v1:v2 + v3:v4 = 1.5 + 2.5 = 4.0
-    kernel += VOP3(VOP3Op.V_ADD_F64, v[5], v[1], v[3]).to_bytes()
-    kernel += v_lshlrev_b32_e32(v[7], 3, v[0]).to_bytes()
-    kernel += global_store_b64(addr=v[7], data=v[5], saddr=s[2]).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    output = (ctypes.c_double * 1)(0.0)
-    output_ptr = ctypes.addressof(output)
-    args = (ctypes.c_uint64 * 1)(output_ptr)
-    args_ptr = ctypes.addressof(args)
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    kernel_ptr = ctypes.addressof(kernel_buf)
-    set_valid_mem_ranges({(output_ptr, 8), (args_ptr, 8), (kernel_ptr, len(kernel))})
-    run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
-    self.assertAlmostEqual(output[0], 4.0, places=10)
-
-  def test_flat_load_d16_hi_b16(self):
-    """FLAT_LOAD_D16_HI_B16: load 16-bit to high half. Regression for FLAT op 35."""
-    from extra.assembly.rdna3.emu import i16
-    # Create a buffer with test data
-    src_data = (ctypes.c_uint16 * 1)(0x1234)
-    src_ptr = ctypes.addressof(src_data)
-    output = (ctypes.c_uint32 * 1)(0xABCD0000)  # preset low bits
-    output_ptr = ctypes.addressof(output)
-    args = (ctypes.c_uint64 * 2)(output_ptr, src_ptr)
-    args_ptr = ctypes.addressof(args)
-
-    kernel = b''
-    kernel += s_load_b128(s[0:3], s[0:1], soffset=NULL, offset=0).to_bytes()
-    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-    kernel += v_mov_b32_e32(v[1], 0xDEAD).to_bytes()  # initial value with low bits set
-    kernel += v_mov_b32_e32(v[2], 0).to_bytes()  # offset = 0
-    kernel += FLAT(FLATOp.FLAT_LOAD_D16_HI_B16, v[1], v[2], saddr=s[2], offset=0).to_bytes()
-    kernel += s_waitcnt(vmcnt=0).to_bytes()
-    kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
-    kernel += global_store_b32(addr=v[3], data=v[1], saddr=s[0]).to_bytes()
-    kernel += s_endpgm().to_bytes()
-
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    kernel_ptr = ctypes.addressof(kernel_buf)
-    set_valid_mem_ranges({(output_ptr, 4), (src_ptr, 2), (args_ptr, 16), (kernel_ptr, len(kernel))})
-    run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
-    # High 16 bits should be 0x1234, low 16 bits preserved as 0xDEAD
-    self.assertEqual(output[0], 0x1234DEAD)
-
-  def test_v_mad_u16(self):
-    """V_MAD_U16: multiply-add unsigned 16-bit. Regression for VOP3 op 577."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 10),  # a = 10
-      v_mov_b32_e32(v[2], 20),  # b = 20
-      v_mov_b32_e32(v[4], 5),   # c = 5
-      VOP3(VOP3Op.V_MAD_U16, v[1], v[1], v[2], v[4]),  # v1 = 10*20+5 = 205
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0] & 0xffff, 205)
-
-  def test_v_lshrrev_b16(self):
-    """V_LSHRREV_B16: logical shift right 16-bit. Regression for VOP3 op 825."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 0x8000),  # value to shift
-      v_mov_b32_e32(v[2], 4),       # shift amount
-      VOP3(VOP3Op.V_LSHRREV_B16, v[1], v[2], v[1]),  # v1 = 0x8000 >> 4 = 0x0800
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0] & 0xffff, 0x0800)
-
-  def test_v_min_u16(self):
-    """V_MIN_U16: minimum of two unsigned 16-bit values. Regression for VOP3 op 779."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 100),
-      v_mov_b32_e32(v[2], 50),
-      VOP3(VOP3Op.V_MIN_U16, v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0] & 0xffff, 50)
-
-class TestWMMA(unittest.TestCase):
-  """Tests for WMMA (Wave Matrix Multiply Accumulate) instructions."""
-
-  def test_wmma_f32_16x16x16_f16_identity(self):
-    """V_WMMA_F32_16X16X16_F16 with identity matrix. Regression for VOP3P op 64."""
-    from extra.assembly.rdna3.emu import i16, f16, exec_wmma_f32_16x16x16_f16, WaveState
-    # Test using direct emulator call rather than full kernel to simplify
-    st = WaveState()
-    st.exec_mask = 0xffffffff  # all 32 lanes active
-
-    # Set up A as identity matrix: A[i][i] = 1.0, rest = 0.0
-    # Lane i holds row i of A in 8 regs (2 fp16 per reg)
-    for lane in range(16):
-      for reg in range(8):
-        col0, col1 = reg * 2, reg * 2 + 1
-        val0 = i16(1.0) if col0 == lane else 0
-        val1 = i16(1.0) if col1 == lane else 0
-        st.vgpr[lane][0 + reg] = val0 | (val1 << 16)  # src0 = v0:v7
-
-    # Set up B as identity matrix: lane i holds column i of B
-    for lane in range(16):
-      for reg in range(8):
-        row0, row1 = reg * 2, reg * 2 + 1
-        val0 = i16(1.0) if row0 == lane else 0
-        val1 = i16(1.0) if row1 == lane else 0
-        st.vgpr[lane][8 + reg] = val0 | (val1 << 16)  # src1 = v8:v15
-
-    # Set up C as zeros
-    for lane in range(32):
-      for reg in range(8):
-        st.vgpr[lane][16 + reg] = 0  # src2 = v16:v23
-
-    # Create a fake VOP3P instruction
-    inst = VOP3P(VOP3POp.V_WMMA_F32_16X16X16_F16, v[24], src0=VGPR(0), src1=VGPR(8), src2=VGPR(16))
-
-    # Execute WMMA
-    exec_wmma_f32_16x16x16_f16(st, inst, 32)
-
-    # Check result: C should be identity (since A @ B where both are identity)
-    # Output i = row*16+col goes to lane (i%32), reg (i//32)
-    for row in range(16):
-      for col in range(16):
-        idx = row * 16 + col
-        lane, reg = idx % 32, idx // 32
-        result = st.vgpr[lane][24 + reg]
-        expected = 1.0 if row == col else 0.0
-        self.assertAlmostEqual(f32(result), expected, places=3,
-          msg=f"C[{row},{col}] = {f32(result)}, expected {expected}")
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_formats.py

@@ -1,332 +0,0 @@
-#!/usr/bin/env python3
-"""Test MUBUF, MTBUF, MIMG, EXP, DS formats against LLVM."""
-import unittest
-from extra.assembly.rdna3.autogen import *
-from extra.assembly.rdna3.lib import encode_src
-
-class TestMUBUF(unittest.TestCase):
-  """Test MUBUF (buffer) instructions."""
-
-  def test_buffer_load_b32_basic(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_b32_idxen(self):
-    # buffer_load_b32 v5, v0, s[8:11], s3 idxen offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x82,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, idxen=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x82,0x03]))
-
-  def test_buffer_load_b32_offen(self):
-    # buffer_load_b32 v5, v0, s[8:11], s3 offen offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x42,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, offen=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x42,0x03]))
-
-  def test_buffer_load_b32_glc(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 glc
-    # GFX11: encoding: [0xff,0x4f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x4f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_b32_slc(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 slc
-    # GFX11: encoding: [0xff,0x1f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, slc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x1f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_b32_dlc(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 dlc
-    # GFX11: encoding: [0xff,0x2f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, dlc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x2f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_b32_all_flags(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 glc slc dlc
-    # GFX11: encoding: [0xff,0x7f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1, slc=1, dlc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x7f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_store_b32(self):
-    # buffer_store_b32 v1, off, s[12:15], s4 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x68,0xe0,0x00,0x01,0x03,0x04]
-    inst = buffer_store_b32(vdata=v[1], vaddr=v[0], srsrc=s[12:16], soffset=s[4], offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x68,0xe0,0x00,0x01,0x03,0x04]))
-
-  def test_buffer_load_b64(self):
-    # buffer_load_b64 v[5:6], off, s[8:11], s3 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x54,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b64(vdata=v[5:7], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x54,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_soffset_m0(self):
-    # buffer_load_b32 v5, off, s[8:11], m0 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x7d]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=M0, offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x7d]))
-
-  def test_buffer_load_soffset_inline_const(self):
-    # buffer_load_b32 v5, off, s[8:11], 0 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x80]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=0, offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x80]))
-
-  def test_buffer_disasm_roundtrip(self):
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1)
-    decoded = MUBUF.from_bytes(inst.to_bytes())
-    self.assertEqual(decoded.to_bytes(), inst.to_bytes())
-
-
-class TestMTBUF(unittest.TestCase):
-  """Test MTBUF (typed buffer) instructions."""
-
-  def test_tbuffer_load_format_x(self):
-    # tbuffer_load_format_x v5, off, s[8:11], s3 format:[BUF_FMT_32_FLOAT] offset:4095
-    # BUF_FMT_32_FLOAT = 22
-    # GFX11: encoding: [0xff,0x0f,0xb0,0xe8,0x00,0x05,0x02,0x03]
-    inst = tbuffer_load_format_x(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=22)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0xb0,0xe8,0x00,0x05,0x02,0x03]))
-
-  def test_tbuffer_store_format_x(self):
-    # tbuffer_store_format_x v5, off, s[8:11], s3 format:[BUF_FMT_32_FLOAT] offset:4095
-    # BUF_FMT_32_FLOAT = 22
-    # GFX11: encoding: [0xff,0x0f,0xb2,0xe8,0x00,0x05,0x02,0x03]
-    inst = tbuffer_store_format_x(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=22)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0xb2,0xe8,0x00,0x05,0x02,0x03]))
-
-  def test_tbuffer_load_format_xy(self):
-    # tbuffer_load_format_xy v[5:6], off, s[8:11], s3 format:[BUF_FMT_32_32_FLOAT] offset:4095
-    # BUF_FMT_32_32_FLOAT = 50
-    # GFX11: encoding: [0xff,0x8f,0x90,0xe9,0x00,0x05,0x02,0x03]
-    inst = tbuffer_load_format_xy(vdata=v[5:7], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=50)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x8f,0x90,0xe9,0x00,0x05,0x02,0x03]))
-
-
-class TestMIMG(unittest.TestCase):
-  """Test MIMG (image) instructions."""
-
-  def test_image_load_2d(self):
-    # image_load v[0:3], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D
-    # GFX11: encoding: [0x04,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]
-    inst = image_load(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1)  # dim=1 is SQ_RSRC_IMG_2D
-    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]))
-
-  def test_image_store_2d(self):
-    # image_store v[0:3], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D
-    # GFX11: encoding: [0x04,0x0f,0x18,0xf0,0x04,0x00,0x00,0x00]
-    inst = image_store(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x18,0xf0,0x04,0x00,0x00,0x00]))
-
-  def test_image_load_1d(self):
-    # image_load v[0:3], v4, s[0:7] dmask:0xf dim:SQ_RSRC_IMG_1D
-    # GFX11: encoding: [0x00,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]
-    inst = image_load(vdata=v[0:4], vaddr=v[4], srsrc=s[0:8], dmask=0xf, dim=0)  # dim=0 is SQ_RSRC_IMG_1D
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]))
-
-  def test_image_sample(self):
-    # image_sample v[0:3], v[4:5], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
-    # GFX11: encoding: [0x04,0x0f,0x6c,0xf0,0x04,0x00,0x00,0x08]
-    inst = image_sample(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], ssamp=s[8:12], dmask=0xf, dim=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x6c,0xf0,0x04,0x00,0x00,0x08]))
-
-  def test_image_load_d16(self):
-    # image_load v[0:1], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D d16
-    # GFX11: encoding: [0x04,0x0f,0x02,0xf0,0x04,0x00,0x00,0x00]
-    inst = image_load(vdata=v[0:2], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1, d16=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x02,0xf0,0x04,0x00,0x00,0x00]))
-
-
-class TestEXP(unittest.TestCase):
-  """Test EXP (export) instructions."""
-
-  def test_exp_mrt0(self):
-    # exp mrt0 v0, v1, v2, v3
-    # GFX11: encoding: [0x0f,0x00,0x00,0xf8,0x00,0x01,0x02,0x03]
-    inst = EXP(en=0xf, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[2], vsrc3=v[3])
-    self.assertEqual(inst.to_bytes(), bytes([0x0f,0x00,0x00,0xf8,0x00,0x01,0x02,0x03]))
-
-  def test_exp_mrtz(self):
-    # exp mrtz v4, v3, v2, v1
-    # GFX11: encoding: [0x8f,0x00,0x00,0xf8,0x04,0x03,0x02,0x01]
-    inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[1])
-    self.assertEqual(inst.to_bytes(), bytes([0x8f,0x00,0x00,0xf8,0x04,0x03,0x02,0x01]))
-
-  def test_exp_mrtz_done(self):
-    # exp mrtz v4, v3, v2, v1 done
-    # GFX11: encoding: [0x8f,0x08,0x00,0xf8,0x04,0x03,0x02,0x01]
-    inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[3], done=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x8f,0x08,0x00,0xf8,0x04,0x03,0x02,0x03]))
-
-  def test_exp_partial_mask(self):
-    # exp mrt0 v0, v1, off, off (en=0x3, only first two components)
-    # GFX11: encoding: [0x03,0x00,0x00,0xf8,0x00,0x01,0x00,0x00]
-    inst = EXP(en=0x3, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[0], vsrc3=v[0])
-    self.assertEqual(inst.to_bytes(), bytes([0x03,0x00,0x00,0xf8,0x00,0x01,0x00,0x00]))
-
-  def test_exp_row_en(self):
-    # exp mrtz v4, v3, v2, v1 row_en
-    # GFX11: encoding: [0x8f,0x20,0x00,0xf8,0x04,0x03,0x02,0x01]
-    inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[1], row=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x8f,0x20,0x00,0xf8,0x04,0x03,0x02,0x01]))
-
-
-class TestDS(unittest.TestCase):
-  """Test DS (data share / LDS) instructions."""
-
-  def test_ds_store_b32(self):
-    # ds_store_b32 v0, v1
-    # GFX11: encoding: [0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
-    inst = ds_store_b32(addr=v[0], data0=v[1])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
-
-  def test_ds_load_b32(self):
-    # ds_load_b32 v0, v1
-    # GFX11: encoding: [0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]
-    inst = ds_load_b32(vdst=v[0], addr=v[1])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]))
-
-  def test_ds_store_b32_offset(self):
-    # ds_store_b32 v0, v1 offset:64
-    # GFX11: encoding: [0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
-    inst = ds_store_b32(addr=v[0], data0=v[1], offset0=64)
-    self.assertEqual(inst.to_bytes(), bytes([0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
-
-  def test_ds_load_b64(self):
-    # ds_load_b64 v[0:1], v2
-    # GFX11: encoding: [0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]
-    inst = ds_load_b64(vdst=v[0:2], addr=v[2])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]))
-
-  def test_ds_add_u32(self):
-    # ds_add_u32 v0, v1
-    # GFX11: encoding: [0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]
-    inst = ds_add_u32(addr=v[0], data0=v[1])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]))
-
-  def test_ds_store_b32_gds(self):
-    # ds_store_b32 v0, v1 gds
-    # GFX11: encoding: [0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]
-    inst = ds_store_b32(addr=v[0], data0=v[1], gds=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]))
-
-
-class TestVOP3(unittest.TestCase):
-  """Test VOP3 (3-operand vector) instructions."""
-
-  def test_v_fma_f32(self):
-    # v_fma_f32 v0, v1, v2, v3
-    # GFX11: encoding: [0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]
-    inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]))
-
-  def test_v_mad_f32(self):
-    # v_fmac_f32_e64 v0, v1, v2 (fmac is fma with implicit dst as src2)
-    # Use v_fma_f32 with vdst == src2
-    inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[0])
-    self.assertEqual(inst.to_bytes()[:4], bytes([0x00,0x00,0x13,0xd6]))
-
-  def test_v_add3_u32(self):
-    # v_add3_u32 v0, v1, v2, v3
-    # GFX11: encoding: [0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]
-    inst = v_add3_u32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]))
-
-
-class TestFLAT(unittest.TestCase):
-  """Test FLAT/GLOBAL/SCRATCH memory instructions."""
-
-  def test_global_load_b32(self):
-    # global_load_b32 v0, v[1:2], off (seg=2 for global)
-    # GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]
-    inst = global_load_b32(vdst=v[0], addr=v[1:3], saddr=OFF)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]))
-
-  def test_global_store_b32(self):
-    # global_store_b32 v[0:1], v2, off (seg=2 for global)
-    # GFX11: encoding: [0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]
-    inst = global_store_b32(addr=v[0:2], data=v[2], saddr=OFF)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]))
-
-  def test_global_load_b32_saddr(self):
-    # global_load_b32 v0, v1, s[0:1] (seg=2 for global)
-    # GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]
-    inst = global_load_b32(vdst=v[0], addr=v[1], saddr=s[0:2])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]))
-
-  def test_global_load_b32_offset(self):
-    # global_load_b32 v0, v[1:2], off offset:256 (seg=2 for global)
-    # GFX11: encoding: [0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]
-    inst = global_load_b32(vdst=v[0], addr=v[1:3], saddr=OFF, offset=256)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]))
-
-  def test_global_load_b64(self):
-    # global_load_b64 v[0:1], v[2:3], off (seg=2 for global)
-    # GFX11: encoding: [0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]
-    inst = global_load_b64(vdst=v[0:2], addr=v[2:4], saddr=OFF)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]))
-
-
-class TestSMEM(unittest.TestCase):
-  """Test SMEM (scalar memory) instructions - regression tests for glc/dlc bit positions."""
-
-  def test_smem_dlc_bit_position(self):
-    # s_load_b32 s5, s[2:3], s0 dlc - tests that DLC is at bit 13 (not bit 14)
-    # GFX11: encoding: [0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]
-    inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], dlc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]))
-
-  def test_smem_glc_bit_position(self):
-    # s_load_b32 s5, s[2:3], s0 glc - tests that GLC is at bit 14 (not bit 16)
-    # GFX11: encoding: [0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]
-    inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], glc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]))
-
-  def test_smem_glc_dlc_combined(self):
-    # s_load_b32 s5, s[2:3], s0 glc dlc - tests both flags together
-    # GFX11: encoding: [0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]
-    inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], glc=1, dlc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]))
-
-  def test_smem_disasm_roundtrip_dlc(self):
-    # Test that disassembly/reassembly preserves DLC bit correctly
-    data = bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00])
-    decoded = SMEM.from_bytes(data)
-    self.assertEqual(decoded.to_bytes(), data)
-
-  def test_smem_disasm_roundtrip_glc_dlc(self):
-    # Test that disassembly/reassembly preserves GLC+DLC bits correctly
-    data = bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00])
-    decoded = SMEM.from_bytes(data)
-    self.assertEqual(decoded.to_bytes(), data)
-
-
-class TestVOP3Literal(unittest.TestCase):
-  """Test VOP3 literal handling - regression tests for Inst64 literal encoding."""
-
-  def test_vop3_with_literal(self):
-    # v_add3_u32 v5, vcc_hi, 0xaf123456, v255
-    # GFX11: encoding: [0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf]
-    from extra.assembly.rdna3.lib import RawImm
-    inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=RawImm(107), src1=0xaf123456, src2=v[255])
-    expected = bytes([0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf])
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_vop3_literal_null_operand(self):
-    # v_add3_u32 v5, null, exec_lo, 0xaf123456
-    # GFX11: encoding: [0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf]
-    from extra.assembly.rdna3.lib import RawImm
-    inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=NULL, src1=RawImm(126), src2=0xaf123456)
-    expected = bytes([0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf])
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_vop3p_with_literal(self):
-    # Test VOP3P literal encoding (also uses Inst64)
-    from extra.assembly.rdna3.lib import RawImm
-    inst = VOP3P(VOP3POp.V_PK_ADD_F16, vdst=v[5], src0=RawImm(240), src1=0x12345678, src2=v[0])
-    self.assertEqual(len(inst.to_bytes()), 12)  # 8 bytes + 4 byte literal
-
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_handwritten.py

@@ -1,178 +0,0 @@
-# do not change these tests. we need to fix bugs to make them pass
-# the Inst constructor should be looking at the types of the fields to correctly set the value
-
-import unittest, struct
-from extra.assembly.rdna3.autogen import *
-from extra.assembly.rdna3.lib import Inst
-from extra.assembly.rdna3.asm import asm
-from extra.assembly.rdna3.test.test_roundtrip import compile_asm
-
-class TestIntegration(unittest.TestCase):
-  inst: Inst
-  def tearDown(self):
-    if not hasattr(self, 'inst'): return
-    b = self.inst.to_bytes()
-    st = self.inst.disasm()
-    reasm = asm(st)
-    desc = f"{st:25s} {self.inst} {b!r} {reasm}"
-    self.assertEqual(b, compile_asm(st), desc)
-    # TODO: this compare should work for valid things
-    #self.assertEqual(self.inst, reasm)
-    self.assertEqual(repr(self.inst), repr(reasm))
-    print(desc)
-
-  def test_load_b128(self):
-    self.inst = s_load_b128(s[4:7], s[0:1], NULL, 0)
-
-  def test_load_b128_wrong_size(self):
-    # this should have to be 4 regs on the loaded to
-    with self.assertRaises(Exception):
-      self.inst = s_load_b128(s[4:6], s[0:1], NULL, 0)
-
-  def test_mov_b32(self):
-    self.inst = s_mov_b32(s[80], s[0])
-
-  def test_mov_b64(self):
-    self.inst = s_mov_b64(s[80:81], s[0:1])
-
-  def test_mov_b32_wrong(self):
-    with self.assertRaises(Exception):
-      self.inst = s_mov_b32(s[80:81], s[0:1])
-    with self.assertRaises(Exception):
-      self.inst = s_mov_b32(s[80:81], s[0])
-    with self.assertRaises(Exception):
-      self.inst = s_mov_b32(s[80], s[0:1])
-
-  def test_mov_b64_wrong(self):
-    with self.assertRaises(Exception):
-      self.inst = s_mov_b64(s[80], s[0])
-    with self.assertRaises(Exception):
-      self.inst = s_mov_b64(s[80], s[0:1])
-    with self.assertRaises(Exception):
-      self.inst = s_mov_b64(s[80:81], s[0])
-
-  def test_load_b128_no_0(self):
-    self.inst = s_load_b128(s[4:7], s[0:1], NULL)
-
-  def test_load_b128_s(self):
-    self.inst = s_load_b128(s[4:7], s[0:1], s[8], 0)
-
-  def test_load_b128_v(self):
-    with self.assertRaises(TypeError):
-      self.inst = s_load_b128(s[4:7], s[0:1], v[8], 0)
-
-  def test_load_b128_off(self):
-    self.inst = s_load_b128(s[4:7], s[0:1], NULL, 3)
-
-  def test_simple_stos(self):
-    self.inst = s_mov_b32(s[0], s[1])
-
-  def test_simple_wrong(self):
-    with self.assertRaises(TypeError):
-      self.inst = s_mov_b32(v[0], s[1])
-
-  def test_simple_vtov(self):
-    self.inst = v_mov_b32_e32(v[0], v[1])
-
-  def test_simple_stov(self):
-    self.inst = v_mov_b32_e32(v[0], s[2])
-
-  def test_simple_float_to_v(self):
-    self.inst = v_mov_b32_e32(v[0], 1.0)
-
-  def test_simple_v_to_float(self):
-    with self.assertRaises(TypeError):
-      self.inst = v_mov_b32_e32(1, v[0])
-
-  def test_simple_int_to_v(self):
-    self.inst = v_mov_b32_e32(v[0], 1)
-
-  def test_three_add(self):
-    self.inst = v_add_co_ci_u32_e32(v[3], s[7], v[3])
-
-  def test_three_add_v(self):
-    self.inst = v_add_co_ci_u32_e32(v[3], v[7], v[3])
-
-  def test_three_add_const(self):
-    self.inst = v_add_co_ci_u32_e32(v[3], 2.0, v[3])
-
-  def test_swaitcnt_lgkm(self): self.inst = s_waitcnt(0xfc07)
-  def test_swaitcnt_vm(self): self.inst = s_waitcnt(0x03f7)
-
-  def test_vmad(self):
-    self.inst = v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2])
-
-  def test_large_imm(self):
-    self.inst = v_mov_b32_e32(v[0], 0x1234)
-
-  def test_dual_mov(self):
-    self.inst = VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], srcy0=v[4])
-
-  def test_dual_mul(self):
-    self.inst = v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
-
-  def test_simple_int_to_s(self):
-    self.inst = s_mov_b32(s[0], 3)
-
-  def test_complex_int_to_s(self):
-    self.inst = s_mov_b32(s[0], 0x235646)
-
-  def test_simple_float_to_s(self):
-    self.inst = s_mov_b32(s[0], 1.0)
-
-  def test_complex_float_to_s(self):
-    self.inst = s_mov_b32(s[0], 1337.0)
-    int_inst = s_mov_b32(s[0], struct.unpack("I", struct.pack("f", 1337.0))[0])
-    self.assertEqual(self.inst, int_inst)
-
-class TestRegisterSliceSyntax(unittest.TestCase):
-  """
-  Issue: Register slice syntax should use AMD assembly convention (inclusive end).
-
-  In AMD assembly, s[4:7] means registers s4, s5, s6, s7 (4 registers, inclusive).
-  The DSL should match this convention so that:
-  - s[4:7] gives 4 registers
-  - Disassembler output can be copied directly back into DSL code
-
-  Fix: Change _RegFactory.__getitem__ to use inclusive end:
-    key.stop - key.start + 1  (instead of key.stop - key.start)
-  """
-  def test_register_slice_count(self):
-    # s[4:7] should give 4 registers: s4, s5, s6, s7 (AMD convention, inclusive)
-    reg = s[4:7]
-    self.assertEqual(reg.count, 4, "s[4:7] should give 4 registers (s4, s5, s6, s7)")
-
-  def test_register_slice_roundtrip(self):
-    # Round-trip: DSL -> disasm -> DSL should preserve register count
-    reg = s[4:7]  # 4 registers in AMD convention
-    inst = s_load_b128(reg, s[0:1], NULL, 0)
-    disasm = inst.disasm()
-    # Disasm shows s[4:7] - user should be able to copy this back
-    self.assertIn("s[4:7]", disasm)
-    # And s[4:7] in DSL should give the same 4 registers
-    reg_from_disasm = s[4:7]
-    self.assertEqual(reg_from_disasm.count, 4, "s[4:7] from disasm should give 4 registers")
-
-class TestInstructionEquality(unittest.TestCase):
-  """
-  Issue: No __eq__ method - instruction comparison requires repr() workaround.
-
-  Two identical instructions should compare equal with ==, but currently:
-    inst1 == inst2 returns False
-
-  The test_handwritten.py works around this with:
-    self.assertEqual(repr(self.inst), repr(reasm))
-  """
-  def test_identical_instructions_equal(self):
-    inst1 = v_mov_b32_e32(v[0], v[1])
-    inst2 = v_mov_b32_e32(v[0], v[1])
-    self.assertEqual(inst1, inst2, "identical instructions should be equal")
-
-  def test_different_instructions_not_equal(self):
-    inst1 = v_mov_b32_e32(v[0], v[1])
-    inst2 = v_mov_b32_e32(v[0], v[2])
-    self.assertNotEqual(inst1, inst2, "different instructions should not be equal")
-
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_integration.py

@@ -1,343 +0,0 @@
-#!/usr/bin/env python3
-"""Integration test: round-trip RDNA3 assembly through AMD toolchain."""
-import unittest, re, io, sys
-from extra.assembly.rdna3.autogen import *
-from extra.assembly.rdna3.asm import waitcnt, asm
-
-def get_amd_toolchain():
-  """Check if AMD toolchain is available."""
-  try:
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    HIPCompiler("gfx1100").compile(".text\ns_endpgm")
-    return True
-  except Exception:
-    return False
-
-def disassemble(lib: bytes, arch: str = "gfx1100") -> str:
-  """Disassemble ELF binary using tinygrad's compiler, return raw output."""
-  from tinygrad.runtime.support.compiler_amd import HIPCompiler
-  old_stdout = sys.stdout
-  sys.stdout = io.StringIO()
-  HIPCompiler(arch).disassemble(lib)
-  output = sys.stdout.getvalue()
-  sys.stdout = old_stdout
-  return output
-
-def parse_disassembly(raw: str) -> list[str]:
-  """Parse disassembly output to list of instruction mnemonics."""
-  lines = []
-  for line in raw.splitlines():
-    if line.startswith('\t'):
-      instr = line.split('//')[0].strip()
-      if instr: lines.append(instr)
-  return lines
-
-def assemble_and_disassemble(instructions: list, arch: str = "gfx1100") -> list[str]:
-  """Assemble instructions with our DSL, then disassemble with AMD toolchain."""
-  from tinygrad.runtime.support.compiler_amd import HIPCompiler
-
-  # Generate bytes from our DSL
-  code_bytes = b''.join(inst.to_bytes() for inst in instructions)
-
-  # Wrap in minimal ELF-compatible assembly with .byte directives
-  byte_str = ', '.join(f'0x{b:02x}' for b in code_bytes)
-  asm_src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n.byte {byte_str}\n"
-
-  # Assemble with AMD COMGR and disassemble
-  lib = HIPCompiler(arch).compile(asm_src)
-  return parse_disassembly(disassemble(lib, arch))
-
-@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
-class TestIntegration(unittest.TestCase):
-  """Test our assembler output matches LLVM disassembly."""
-
-  def test_simple_sop1(self):
-    """Test SOP1 instructions round-trip."""
-    instructions = [
-      s_mov_b32(s[0], s[1]),
-      s_mov_b32(s[2], 0),
-      s_not_b32(s[3], s[4]),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('s_mov_b32', disasm[0])
-    self.assertIn('s_mov_b32', disasm[1])
-    self.assertIn('s_not_b32', disasm[2])
-
-  def test_simple_sop2(self):
-    """Test SOP2 instructions round-trip."""
-    instructions = [
-      s_add_u32(s[0], s[1], s[2]),
-      s_sub_u32(s[3], s[4], 10),
-      s_and_b32(s[5], s[6], s[7]),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('s_add_u32', disasm[0])
-    self.assertIn('s_sub_u32', disasm[1])
-    self.assertIn('s_and_b32', disasm[2])
-
-  def test_simple_vop2(self):
-    """Test VOP2 instructions round-trip."""
-    instructions = [
-      v_add_f32_e32(v[0], v[1], v[2]),
-      v_mul_f32_e32(v[3], 1.0, v[4]),  # 1.0 is inline constant
-      v_and_b32_e32(v[5], 10, v[6]),  # small inline constant
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('v_add_f32', disasm[0])
-    self.assertIn('v_mul_f32', disasm[1])
-
-  def test_control_flow(self):
-    """Test control flow instructions."""
-    instructions = [
-      s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
-      s_endpgm(),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('s_waitcnt', disasm[0])
-    self.assertIn('s_endpgm', disasm[1])
-
-  def test_memory_ops(self):
-    """Test memory instructions."""
-    instructions = [
-      s_load_b32(s[0], s[0:2], NULL),
-      s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
-      global_store_b32(addr=v[0:2], data=v[2], saddr=OFF),
-      s_endpgm(),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('s_load_b32', disasm[0])
-    self.assertIn('s_waitcnt', disasm[1])
-    self.assertIn('global_store_b32', disasm[2])
-
-  def test_full_kernel(self):
-    """Test a complete kernel similar to tinygrad output."""
-    # Simple kernel: load value, add 1, store back
-    instructions = [
-      # Get thread ID
-      v_mov_b32_e32(v[0], s[0]),  # base addr low
-      v_mov_b32_e32(v[1], s[1]),  # base addr high
-      # Load value
-      global_load_b32(vdst=v[2], addr=v[0:2], saddr=OFF),
-      s_waitcnt(simm16=waitcnt(vmcnt=0)),
-      # Add 1.0
-      v_add_f32_e32(v[2], 1.0, v[2]),
-      # Store result
-      global_store_b32(addr=v[0:2], data=v[2], saddr=OFF),
-      s_endpgm(),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    # Verify key instructions are present
-    self.assertTrue(any('global_load' in d for d in disasm))
-    self.assertTrue(any('v_add_f32' in d for d in disasm))
-    self.assertTrue(any('global_store' in d for d in disasm))
-    self.assertTrue(any('s_endpgm' in d for d in disasm))
-
-  def test_bytes_roundtrip(self):
-    """Test that our bytes match what AMD assembler produces."""
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-
-    # Simple instruction
-    inst = s_mov_b32(s[0], s[1])
-    our_bytes = inst.to_bytes()
-
-    # Assemble same instruction with AMD toolchain
-    asm_src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\ns_mov_b32 s0, s1\n"
-    compiler = HIPCompiler("gfx1100")
-    lib = compiler.compile(asm_src)
-    raw = disassemble(lib)
-
-    for line in raw.splitlines():
-      if 's_mov_b32' in line and '//' in line:
-        # Extract hex bytes from comment: "// 000000001300: BE800001"
-        comment = line.split('//')[1].strip()
-        hex_str = comment.split(':')[1].strip()
-        # Convert big-endian hex string to little-endian bytes
-        amd_bytes = bytes.fromhex(hex_str)[::-1]  # reverse for little-endian
-        self.assertEqual(our_bytes, amd_bytes, f"Bytes mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
-        return
-    self.fail("Could not find s_mov_b32 in disassembly")
-
-@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
-class TestAsm(unittest.TestCase):
-  """Test asm() string parsing."""
-
-  def test_asm_basic(self):
-    """Test basic instruction parsing."""
-    inst = asm('s_mov_b32 s0, s1')
-    self.assertEqual(inst.to_bytes(), s_mov_b32(s[0], s[1]).to_bytes())
-
-  def test_asm_with_immediates(self):
-    """Test parsing with immediate values."""
-    inst = asm('s_add_u32 s0, s1, 10')
-    self.assertEqual(inst.to_bytes(), s_add_u32(s[0], s[1], 10).to_bytes())
-
-  def test_asm_float_const(self):
-    """Test parsing float constants."""
-    inst = asm('v_mul_f32_e32 v0, 1.0, v1')
-    self.assertEqual(inst.to_bytes(), v_mul_f32_e32(v[0], 1.0, v[1]).to_bytes())
-
-  def test_asm_hex_immediate(self):
-    """Test parsing hex immediates."""
-    inst = asm('s_waitcnt 0xfc07')
-    self.assertEqual(inst.to_bytes(), s_waitcnt(simm16=0xfc07).to_bytes())
-
-  def test_asm_special_regs(self):
-    """Test parsing special registers."""
-    inst = asm('s_mov_b32 s0, vcc_lo')
-    self.assertEqual(inst.to_bytes(), s_mov_b32(s[0], VCC_LO).to_bytes())
-
-  def test_asm_register_range(self):
-    """Test parsing register ranges."""
-    inst = asm('s_load_b128 s[4:7], s[0:1], null')
-    self.assertEqual(inst.to_bytes(), s_load_b128(s[4:7], s[0:1], NULL).to_bytes())
-
-  def test_asm_matches_llvm(self):
-    """Test asm() output matches LLVM assembler."""
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    compiler = HIPCompiler('gfx1100')
-
-    def get_llvm_bytes(instr: str) -> bytes:
-      src = f'.text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n{instr}\n'
-      lib = compiler.compile(src)
-      raw = disassemble(lib)
-      for line in raw.splitlines():
-        if instr.split()[0] in line and '//' in line:
-          hex_str = line.split('//')[1].strip().split(':')[1].strip()
-          return bytes.fromhex(hex_str)[::-1]
-      return b''
-
-    tests = ['s_mov_b32 s0, s1', 's_endpgm', 'v_add_f32_e32 v0, v1, v2']
-    for t in tests:
-      self.assertEqual(asm(t).to_bytes(), get_llvm_bytes(t), f"mismatch for: {t}")
-
-  def test_asm_vop3_modifiers(self):
-    """Test asm() with VOP3 modifiers (neg, abs, clamp)."""
-    import subprocess, re
-
-    def get_llvm_encoding(instr: str) -> str:
-      result = subprocess.run(['llvm-mc', '-triple=amdgcn', '-mcpu=gfx1100', '-show-encoding'],
-                              input=instr, capture_output=True, text=True)
-      if m := re.search(r'encoding:\s*\[(.*?)\]', result.stdout):
-        return m.group(1).replace('0x','').replace(',','').replace(' ','')
-      return ''
-
-    tests = [
-      'v_fma_f32 v0, -v1, v2, v3',       # neg on src0
-      'v_fma_f32 v0, v1, |v2|, v3',      # abs on src1
-      'v_fma_f32 v0, v1, v2, v3 clamp',  # clamp
-      'v_fma_f32 v0, -v1, |v2|, v3 clamp',  # all modifiers
-      'v_fma_f32 v0, -|v1|, v2, v3',     # neg+abs on same operand
-    ]
-    for t in tests:
-      our_hex = asm(t).to_bytes().hex()
-      llvm_hex = get_llvm_encoding(t)
-      self.assertEqual(our_hex, llvm_hex, f"mismatch for: {t}")
-
-@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
-class TestTinygradIntegration(unittest.TestCase):
-  """Test that we can parse disassembled tinygrad kernels."""
-
-  def test_simple_add_kernel(self):
-    """Generate a simple add kernel from tinygrad and verify disassembly."""
-    from tinygrad import Tensor
-    from tinygrad.codegen import get_program
-    from tinygrad.renderer.cstyle import AMDHIPRenderer
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    from tinygrad.uop.ops import Ops
-
-    # Create a computation that generates a real kernel
-    a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
-    b = Tensor([5.0, 6.0, 7.0, 8.0]).realize()
-    c = a + b
-
-    # Get schedule and find SINK
-    schedule = c.schedule()
-    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
-    self.assertTrue(len(sink_items) > 0, "No SINK in schedule")
-
-    # Generate program
-    renderer = AMDHIPRenderer('gfx1100')
-    prg = get_program(sink_items[0].ast, renderer)
-    self.assertIsNotNone(prg.src)
-
-    # Compile and disassemble
-    compiler = HIPCompiler('gfx1100')
-    lib = compiler.compile(prg.src)
-    raw_disasm = disassemble(lib)
-    instrs = parse_disassembly(raw_disasm)
-
-    # Verify we got some instructions
-    self.assertTrue(len(instrs) > 0, "No instructions in disassembly")
-    # Should have an endpgm
-    self.assertTrue(any('s_endpgm' in i for i in instrs), "Missing s_endpgm")
-
-  def test_matmul_kernel(self):
-    """Generate a matmul kernel and verify disassembly has expected patterns."""
-    from tinygrad import Tensor
-    from tinygrad.codegen import get_program
-    from tinygrad.renderer.cstyle import AMDHIPRenderer
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    from tinygrad.uop.ops import Ops
-
-    # Create a small matmul
-    a = Tensor.rand(4, 4).realize()
-    b = Tensor.rand(4, 4).realize()
-    c = a @ b
-
-    # Get schedule
-    schedule = c.schedule()
-    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
-    self.assertTrue(len(sink_items) > 0)
-
-    # Generate and compile
-    renderer = AMDHIPRenderer('gfx1100')
-    prg = get_program(sink_items[0].ast, renderer)
-    compiler = HIPCompiler('gfx1100')
-    lib = compiler.compile(prg.src)
-    raw_disasm = disassemble(lib)
-    instrs = parse_disassembly(raw_disasm)
-
-    # Matmul should have multiply and add instructions
-    has_mul = any('mul' in i.lower() for i in instrs)
-    has_add = any('add' in i.lower() for i in instrs)
-    self.assertTrue(has_mul or has_add, "Matmul should have mul/add ops")
-
-  def test_disasm_to_bytes_roundtrip(self):
-    """Parse disassembled instructions and verify we can re-encode some of them."""
-    from tinygrad import Tensor
-    from tinygrad.codegen import get_program
-    from tinygrad.renderer.cstyle import AMDHIPRenderer
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    from tinygrad.uop.ops import Ops
-
-    # Simple kernel
-    a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
-    b = (a * 2.0)
-
-    schedule = b.schedule()
-    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
-    if not sink_items: return  # skip if no kernel
-
-    renderer = AMDHIPRenderer('gfx1100')
-    prg = get_program(sink_items[0].ast, renderer)
-    compiler = HIPCompiler('gfx1100')
-    lib = compiler.compile(prg.src)
-    raw_disasm = disassemble(lib)
-
-    # Find s_endpgm and verify we can encode it
-    for line in raw_disasm.splitlines():
-      if 's_endpgm' in line and '//' in line:
-        # Extract bytes from comment
-        comment = line.split('//')[1].strip()
-        hex_str = comment.split(':')[1].strip()
-        amd_bytes = bytes.fromhex(hex_str)[::-1]
-
-        # Our encoding
-        our_inst = s_endpgm()
-        our_bytes = our_inst.to_bytes()
-
-        self.assertEqual(our_bytes, amd_bytes, f"s_endpgm mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
-        return
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_llvm.py

@@ -1,173 +0,0 @@
-#!/usr/bin/env python3
-"""Test RDNA3 assembler/disassembler against LLVM test vectors."""
-import unittest, re
-from tinygrad.helpers import fetch
-from extra.assembly.rdna3.autogen import *
-from extra.assembly.rdna3.asm import asm
-from extra.assembly.rdna3.test.test_roundtrip import compile_asm, disassemble_lib
-
-LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/main/llvm/test/MC/AMDGPU"
-
-# Format info: (filename, format_class, op_enum)
-LLVM_TEST_FILES = {
-  # Scalar ALU
-  'sop1': ('gfx11_asm_sop1.s', SOP1, SOP1Op),
-  'sop2': ('gfx11_asm_sop2.s', SOP2, SOP2Op),
-  'sopp': ('gfx11_asm_sopp.s', SOPP, SOPPOp),
-  'sopk': ('gfx11_asm_sopk.s', SOPK, SOPKOp),
-  'sopc': ('gfx11_asm_sopc.s', SOPC, SOPCOp),
-  # Vector ALU
-  'vop1': ('gfx11_asm_vop1.s', VOP1, VOP1Op),
-  'vop2': ('gfx11_asm_vop2.s', VOP2, VOP2Op),
-  'vopc': ('gfx11_asm_vopc.s', VOPC, VOPCOp),
-  'vop3': ('gfx11_asm_vop3.s', VOP3, VOP3Op),
-  'vop3p': ('gfx11_asm_vop3p.s', VOP3P, VOP3POp),
-  'vop3sd': ('gfx11_asm_vop3.s', VOP3SD, VOP3SDOp),  # VOP3SD shares file with VOP3
-  'vinterp': ('gfx11_asm_vinterp.s', VINTERP, VINTERPOp),
-  'vopd': ('gfx11_asm_vopd.s', VOPD, VOPDOp),
-  'vopcx': ('gfx11_asm_vopcx.s', VOPC, VOPCOp),  # VOPCX uses VOPC format
-  # VOP3 promotions (VOP1/VOP2/VOPC promoted to VOP3 encoding)
-  'vop3_from_vop1': ('gfx11_asm_vop3_from_vop1.s', VOP3, VOP3Op),
-  'vop3_from_vop2': ('gfx11_asm_vop3_from_vop2.s', VOP3, VOP3Op),
-  'vop3_from_vopc': ('gfx11_asm_vop3_from_vopc.s', VOP3, VOP3Op),
-  'vop3_from_vopcx': ('gfx11_asm_vop3_from_vopcx.s', VOP3, VOP3Op),
-  # Memory
-  'ds': ('gfx11_asm_ds.s', DS, DSOp),
-  'smem': ('gfx11_asm_smem.s', SMEM, SMEMOp),
-  'flat': ('gfx11_asm_flat.s', FLAT, FLATOp),
-  'mubuf': ('gfx11_asm_mubuf.s', MUBUF, MUBUFOp),
-  'mtbuf': ('gfx11_asm_mtbuf.s', MTBUF, MTBUFOp),
-  'mimg': ('gfx11_asm_mimg.s', MIMG, MIMGOp),
-  # WMMA (matrix multiply)
-  'wmma': ('gfx11_asm_wmma.s', VOP3P, VOP3POp),
-  # Additional features
-  'vop3_features': ('gfx11_asm_vop3_features.s', VOP3, VOP3Op),
-  'vop3p_features': ('gfx11_asm_vop3p_features.s', VOP3P, VOP3POp),
-  'vopd_features': ('gfx11_asm_vopd_features.s', VOPD, VOPDOp),
-  # Alias files (alternative mnemonics)
-  'vop3_alias': ('gfx11_asm_vop3_alias.s', VOP3, VOP3Op),
-  'vop3p_alias': ('gfx11_asm_vop3p_alias.s', VOP3P, VOP3POp),
-  'vopc_alias': ('gfx11_asm_vopc_alias.s', VOPC, VOPCOp),
-  'vopcx_alias': ('gfx11_asm_vopcx_alias.s', VOPC, VOPCOp),
-  'vinterp_alias': ('gfx11_asm_vinterp_alias.s', VINTERP, VINTERPOp),
-  'smem_alias': ('gfx11_asm_smem_alias.s', SMEM, SMEMOp),
-  'mubuf_alias': ('gfx11_asm_mubuf_alias.s', MUBUF, MUBUFOp),
-  'mtbuf_alias': ('gfx11_asm_mtbuf_alias.s', MTBUF, MTBUFOp),
-}
-
-def parse_llvm_tests(text: str) -> list[tuple[str, bytes]]:
-  """Parse LLVM test format into (asm, expected_bytes) pairs."""
-  tests, lines = [], text.split('\n')
-  for i, line in enumerate(lines):
-    line = line.strip()
-    if not line or line.startswith(('//', '.', ';')): continue
-    asm_text = line.split('//')[0].strip()
-    if not asm_text: continue
-    for j in range(i, min(i + 3, len(lines))):
-      # Match GFX11, W32, or W64 encodings (all valid for gfx11)
-      if m := re.search(r'(?:GFX11|W32|W64)[^:]*:.*?encoding:\s*\[(.*?)\]', lines[j]):
-        hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
-        if hex_bytes:
-          try: tests.append((asm_text, bytes.fromhex(hex_bytes)))
-          except ValueError: pass
-        break
-  return tests
-
-def try_assemble(text: str):
-  """Try to assemble instruction text, return bytes or None on failure."""
-  try: return asm(text).to_bytes()
-  except: return None
-
-class TestLLVM(unittest.TestCase):
-  """Test assembler and disassembler against all LLVM test vectors."""
-  tests: dict[str, list[tuple[str, bytes]]] = {}
-
-  @classmethod
-  def setUpClass(cls):
-    for name, (filename, _, _) in LLVM_TEST_FILES.items():
-      try:
-        data = fetch(f"{LLVM_BASE}/{filename}").read_bytes()
-        cls.tests[name] = parse_llvm_tests(data.decode('utf-8', errors='ignore'))
-      except Exception as e:
-        print(f"Warning: couldn't fetch {filename}: {e}")
-        cls.tests[name] = []
-
-# Generate test methods dynamically for each format
-def _make_asm_test(name):
-  def test(self):
-    passed, failed, skipped = 0, 0, 0
-    for asm_text, expected in self.tests.get(name, []):
-      result = try_assemble(asm_text)
-      if result is None: skipped += 1
-      elif result == expected: passed += 1
-      else: failed += 1
-    print(f"{name.upper()} asm: {passed} passed, {failed} failed, {skipped} skipped")
-    self.assertEqual(failed, 0)
-  return test
-
-def _make_disasm_test(name):
-  def test(self):
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    compiler = HIPCompiler('gfx1100')
-    _, fmt_cls, op_enum = LLVM_TEST_FILES[name]
-    passed, failed, skipped, failures = 0, 0, 0, []
-    # VOP3SD opcodes that share encoding with VOP3 (only for vop3sd test, not vopc promotions)
-    # Note: opcodes 0-255 are VOPC promoted to VOP3, never VOP3SD
-    vop3sd_opcodes = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
-    # vop3_from_vopc/vopcx tests have VOPC opcodes 0-255, not VOP3SD - don't detect as VOP3SD
-    is_vopc_promotion = name in ('vop3_from_vopc', 'vop3_from_vopcx')
-    # Undocumented opcodes not in AMD ISA PDF - skip these
-    undocumented = {'smem': {34, 35}, 'sopk': {22, 23}, 'sopp': {8, 58, 59}}  # s_atc_probe*, s_subvector_loop*, s_waitcnt_depctr, unknown
-    for asm_text, data in self.tests.get(name, []):
-      if len(data) > fmt_cls._size(): continue  # skip literals (need different handling)
-      # Skip undocumented opcodes
-      temp_inst = fmt_cls.from_bytes(data)
-      temp_op = temp_inst._values.get('op', 0)
-      temp_op = temp_op.val if hasattr(temp_op, 'val') else temp_op
-      if temp_op in undocumented.get(name, set()): skipped += 1; continue
-      # Skip SOPP no-imm instructions with non-zero simm16 (can't roundtrip through LLVM)
-      if name == 'sopp':
-        simm16 = temp_inst._values.get('simm16', 0)
-        simm16 = simm16.val if hasattr(simm16, 'val') else simm16
-        sopp_no_imm = {48, 54, 53, 55, 60, 61, 62}  # s_endpgm, s_barrier, s_wakeup, s_icache_inv, s_wait_idle, s_endpgm_saved, s_code_end
-        if temp_op in sopp_no_imm and simm16 != 0: skipped += 1; continue
-      try:
-        # VOP3 and VOP3SD share encoding - peek at opcode to determine which class to use
-        if fmt_cls.__name__ in ('VOP3', 'VOP3SD'):
-          temp = VOP3.from_bytes(data)
-          op_val = temp._values.get('op', 0)
-          op_val = op_val.val if hasattr(op_val, 'val') else op_val
-          is_vop3sd = (op_val in vop3sd_opcodes) and not is_vopc_promotion
-          decoded = VOP3SD.from_bytes(data) if is_vop3sd else VOP3.from_bytes(data)
-          # Validate opcode with appropriate enum
-          if is_vop3sd:
-            VOP3SDOp(op_val)
-          else:
-            VOP3Op(op_val)
-        else:
-          decoded = fmt_cls.from_bytes(data)
-          op_val = decoded._values.get('op', 0)
-          op_val = op_val.val if hasattr(op_val, 'val') else op_val
-          op_enum(op_val)  # validate opcode
-        if decoded.to_bytes()[:len(data)] != data:
-          failed += 1; failures.append(f"decode roundtrip failed for {data.hex()}"); continue
-        disasm_str = decoded.disasm()
-        # Test: LLVM should assemble our disasm output to the same bytes
-        llvm_bytes = compile_asm(disasm_str, compiler)
-        if llvm_bytes is None:
-          failed += 1; failures.append(f"LLVM failed to assemble: '{disasm_str}' (from '{asm_text}')")
-        elif llvm_bytes == data: passed += 1
-        else: failed += 1; failures.append(f"'{disasm_str}': expected={data.hex()} got={llvm_bytes.hex()}")
-      except Exception as e:
-        failed += 1; failures.append(f"exception for {data.hex()}: {e}")
-    print(f"{name.upper()} disasm: {passed} passed, {failed} failed" + (f", {skipped} skipped" if skipped else ""))
-    if failures[:10]: print("  " + "\n  ".join(failures[:10]))
-    self.assertEqual(failed, 0)
-  return test
-
-for name in LLVM_TEST_FILES:
-  setattr(TestLLVM, f'test_{name}_asm', _make_asm_test(name))
-  setattr(TestLLVM, f'test_{name}_disasm', _make_disasm_test(name))
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_mockgpu_invalid.py

@@ -1,54 +0,0 @@
-#!/usr/bin/env python3
-"""Test that invalid instructions raise exceptions through the mock GPU stack."""
-import unittest, subprocess, os, time
-
-class TestMockGPUInvalidInstruction(unittest.TestCase):
-  def test_unsupported_instruction_raises(self):
-    """Test that unsupported instructions raise immediately through the full MOCKGPU stack."""
-    test_code = '''
-import struct
-from tinygrad import Device, Tensor
-from tinygrad.engine.realize import get_runner
-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)
-
-prg = runner._prg
-lib = bytearray(prg.lib)
-
-# Find s_endpgm (0xBFB00000) and replace with invalid SOPP op=127 (0xBFFF0000)
-found = False
-for i in range(0, len(lib) - 4, 4):
-  if struct.unpack("<I", lib[i:i+4])[0] == 0xBFB00000:
-    lib[i:i+4] = struct.pack("<I", 0xBFFF0000)
-    found = True
-    break
-assert found, "s_endpgm not found"
-
-patched_prg = AMDProgram(dev, "patched", bytes(lib))
-b.uop.buffer.allocate()
-patched_prg(b.uop.buffer._buf, a.uop.buffer._buf, global_size=(1,1,1), local_size=(1,1,1))
-dev.synchronize()
-'''
-
-    env = os.environ.copy()
-    env["AMD"] = "1"
-    env["MOCKGPU"] = "1"
-    env["PYTHON_REMU"] = "1"
-    env["HCQDEV_WAIT_TIMEOUT_MS"] = "10000"
-
-    st = time.perf_counter()
-    result = subprocess.run(["python", "-c", test_code], env=env, capture_output=True, text=True, timeout=60)
-    elapsed = time.perf_counter() - st
-
-    self.assertNotEqual(result.returncode, 0, "should have raised")
-    self.assertIn("NotImplementedError", result.stderr)
-    # Should exit immediately, not wait for the full timeout
-    self.assertLess(elapsed, 5.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_pdf_parser.py

@@ -1,158 +0,0 @@
-#!/usr/bin/env python3
-"""Test that PDF parser correctly extracts format fields."""
-import unittest
-from extra.assembly.rdna3.autogen import (
-  SOP1, SOP2, SOPK, SOPP, VOP1, VOP2, VOP3SD, VOPC, FLAT, VOPD,
-  SOP1Op, SOP2Op, VOP1Op, VOP3Op
-)
-
-# expected formats with key fields and whether they have ENCODING
-EXPECTED_FORMATS = {
-  'DPP16': (['SRC0', 'DPP_CTRL', 'BANK_MASK', 'ROW_MASK'], False),
-  'DPP8': (['SRC0', 'LANE_SEL0', 'LANE_SEL7'], False),
-  'DS': (['OP', 'ADDR', 'DATA0', 'DATA1', 'VDST'], True),
-  'EXP': (['EN', 'TARGET', 'VSRC0', 'VSRC1', 'VSRC2', 'VSRC3'], True),
-  'FLAT': (['OP', 'ADDR', 'DATA', 'SADDR', 'VDST', 'OFFSET'], True),
-  'LDSDIR': (['VDST', 'OP'], True),
-  'MIMG': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'DMASK'], True),
-  'MTBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'FORMAT', 'SOFFSET'], True),
-  'MUBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'SOFFSET'], True),
-  'SMEM': (['OP', 'SBASE', 'SDATA', 'OFFSET', 'SOFFSET'], True),
-  'SOP1': (['OP', 'SDST', 'SSRC0'], True),
-  'SOP2': (['OP', 'SDST', 'SSRC0', 'SSRC1'], True),
-  'SOPC': (['OP', 'SSRC0', 'SSRC1'], True),
-  'SOPK': (['OP', 'SDST', 'SIMM16'], True),
-  'SOPP': (['OP', 'SIMM16'], True),
-  'VINTERP': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
-  'VOP1': (['OP', 'VDST', 'SRC0'], True),
-  'VOP2': (['OP', 'VDST', 'SRC0', 'VSRC1'], True),
-  'VOP3': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
-  'VOP3P': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
-  'VOP3SD': (['OP', 'VDST', 'SDST', 'SRC0', 'SRC1', 'SRC2'], True),
-  'VOPC': (['OP', 'SRC0', 'VSRC1'], True),
-  'VOPD': (['OPX', 'OPY', 'SRCX0', 'SRCY0', 'VDSTX', 'VDSTY'], True),
-}
-
-class TestPDFParserGenerate(unittest.TestCase):
-  """Test the PDF parser by running generate() and checking results."""
-  result: dict
-
-  @classmethod
-  def setUpClass(cls):
-    from extra.assembly.rdna3.gen import generate
-    cls.result = generate()
-
-  def test_all_formats_present(self):
-    """All expected formats should be parsed."""
-    for fmt_name in EXPECTED_FORMATS:
-      self.assertIn(fmt_name, self.result["formats"], f"missing format {fmt_name}")
-
-  def test_format_count(self):
-    """Should have exactly 23 formats."""
-    self.assertEqual(len(self.result["formats"]), 23)
-
-  def test_no_duplicate_fields(self):
-    """No format should have duplicate field names."""
-    for fmt_name, fields in self.result["formats"].items():
-      field_names = [f[0] for f in fields]
-      self.assertEqual(len(field_names), len(set(field_names)), f"{fmt_name} has duplicate fields: {field_names}")
-
-  def test_expected_fields(self):
-    """Each format should have its expected key fields."""
-    for fmt_name, (expected_fields, has_encoding) in EXPECTED_FORMATS.items():
-      fields = {f[0] for f in self.result["formats"].get(fmt_name, [])}
-      for field in expected_fields:
-        self.assertIn(field, fields, f"{fmt_name} missing {field}")
-      if has_encoding:
-        self.assertIn("ENCODING", fields, f"{fmt_name} should have ENCODING")
-      else:
-        self.assertNotIn("ENCODING", fields, f"{fmt_name} should not have ENCODING")
-
-  def test_vopd_no_dpp16_fields(self):
-    """VOPD should not have DPP16-specific fields (parser boundary bug)."""
-    vopd_fields = {f[0] for f in self.result["formats"].get("VOPD", [])}
-    for field in ['DPP_CTRL', 'BANK_MASK', 'ROW_MASK']:
-      self.assertNotIn(field, vopd_fields, f"VOPD should not have {field}")
-
-  def test_dpp16_no_vinterp_fields(self):
-    """DPP16 should not have VINTERP-specific fields."""
-    dpp16_fields = {f[0] for f in self.result["formats"].get("DPP16", [])}
-    for field in ['VDST', 'WAITEXP']:
-      self.assertNotIn(field, dpp16_fields, f"DPP16 should not have {field}")
-
-  def test_sopp_no_smem_fields(self):
-    """SOPP should not have SMEM fields (page break bug)."""
-    sopp_fields = {f[0] for f in self.result["formats"].get("SOPP", [])}
-    for field in ['SBASE', 'SDATA']:
-      self.assertNotIn(field, sopp_fields, f"SOPP should not have {field}")
-
-class TestPDFParser(unittest.TestCase):
-  """Verify format classes have correct fields from PDF parsing."""
-
-  def test_sop2_fields(self):
-    """SOP2 should have op, sdst, ssrc0, ssrc1."""
-    for field in ['op', 'sdst', 'ssrc0', 'ssrc1']:
-      self.assertIn(field, SOP2._fields)
-    self.assertEqual(SOP2._fields['op'].hi, 29)
-    self.assertEqual(SOP2._fields['op'].lo, 23)
-
-  def test_sop1_fields(self):
-    """SOP1 should have op, sdst, ssrc0 with correct bit positions."""
-    for field in ['op', 'sdst', 'ssrc0']:
-      self.assertIn(field, SOP1._fields)
-    self.assertNotIn('simm16', SOP1._fields)
-    self.assertEqual(SOP1._fields['ssrc0'].hi, 7)
-    self.assertEqual(SOP1._fields['ssrc0'].lo, 0)
-    assert SOP1._encoding is not None
-    self.assertEqual(SOP1._encoding[0].hi, 31)
-    self.assertEqual(SOP1._encoding[1], 0b101111101)
-
-  def test_vop3sd_fields(self):
-    """VOP3SD should have all fields including src0/src1/src2 from page continuation."""
-    for field in ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2']:
-      self.assertIn(field, VOP3SD._fields)
-    self.assertEqual(VOP3SD._fields['src0'].hi, 40)
-    self.assertEqual(VOP3SD._fields['src0'].lo, 32)
-    self.assertEqual(VOP3SD._size(), 8)
-
-  def test_flat_has_vdst(self):
-    """FLAT should have vdst field."""
-    self.assertIn('vdst', FLAT._fields)
-    self.assertEqual(FLAT._fields['vdst'].hi, 63)
-    self.assertEqual(FLAT._fields['vdst'].lo, 56)
-
-  def test_encoding_bits(self):
-    """Verify encoding bits are correct for major formats."""
-    tests = [
-      (SOP2, 31, 30, 0b10),
-      (SOPK, 31, 28, 0b1011),
-      (SOPP, 31, 23, 0b101111111),
-      (VOP1, 31, 25, 0b0111111),
-      (VOP2, 31, 31, 0b0),
-      (VOPC, 31, 25, 0b0111110),
-      (FLAT, 31, 26, 0b110111),
-    ]
-    for cls, hi, lo, val in tests:
-      assert cls._encoding is not None
-      self.assertEqual(cls._encoding[0].hi, hi, f"{cls.__name__} encoding hi")
-      self.assertEqual(cls._encoding[0].lo, lo, f"{cls.__name__} encoding lo")
-      self.assertEqual(cls._encoding[1], val, f"{cls.__name__} encoding val")
-
-  def test_opcode_enums_exist(self):
-    """Verify opcode enums are generated with expected counts."""
-    self.assertGreater(len(SOP1Op), 50)
-    self.assertGreater(len(SOP2Op), 50)
-    self.assertGreater(len(VOP1Op), 50)
-    self.assertGreater(len(VOP3Op), 200)
-
-  def test_vopd_no_duplicate_fields(self):
-    """VOPD should not have duplicate fields and should not include DPP16 fields."""
-    field_names = list(VOPD._fields.keys())
-    self.assertEqual(len(field_names), len(set(field_names)))
-    for field in ['srcx0', 'srcy0', 'opx', 'opy']:
-      self.assertIn(field, VOPD._fields)
-    for field in ['dpp_ctrl', 'bank_mask', 'row_mask']:
-      self.assertNotIn(field, VOPD._fields)
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_rdna3_asm.py

@@ -1,94 +0,0 @@
-#!/usr/bin/env python3
-import unittest, subprocess
-from extra.assembly.rdna3.autogen import *
-
-def llvm_assemble(asm: str) -> bytes:
-  """Assemble using llvm-mc and return bytes."""
-  result = subprocess.run(
-    ["llvm-mc", "-triple=amdgcn", "-mcpu=gfx1100", "-show-encoding"],
-    input=asm, capture_output=True, text=True
-  )
-  out = b''
-  for line in result.stdout.split('\n'):
-    if 'encoding:' in line:
-      enc = line.split('encoding:')[1].strip()
-      enc = enc.strip('[]').replace('0x', '').replace(',', '')
-      out += bytes.fromhex(enc)
-  if not out: raise ValueError(f"no encoding found: {result.stdout} {result.stderr}")
-  return out
-
-class TestRDNA3Asm(unittest.TestCase):
-  def test_full_program(self):
-    """Test the full program from rdna3fun.py matches llvm-mc output."""
-    program = [
-      v_bfe_u32(v[1], v[0], 10, 10),
-      s_load_b128(s[4:7], s[0:1], NULL),
-      v_and_b32_e32(v[0], 0x3FF, v[0]),
-      s_mulk_i32(s[3], 0x87),
-      v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2]),
-      v_mul_u32_u24_e32(v[0], 45, v[0]),
-      v_ashrrev_i32_e32(v[2], 31, v[1]),
-      v_add3_u32(v[0], v[0], s[3], v[1]),
-      v_lshlrev_b64(v[2:3], 2, v[1:2]),
-      v_ashrrev_i32_e32(v[1], 31, v[0]),
-      v_lshlrev_b64(v[0:1], 2, v[0:1]),
-      s_waitcnt(0xfc07),  # lgkmcnt(0)
-      v_add_co_u32(v[2], VCC_LO, s[6], v[2]),
-      v_add_co_ci_u32_e32(v[3], s[7], v[3]),
-      v_add_co_u32(v[0], VCC_LO, s[4], v[0]),
-      global_load_b32(vdst=v[2], addr=v[2], saddr=OFF),
-      v_add_co_ci_u32_e32(v[1], s[5], v[1]),
-      s_waitcnt(0x03f7),  # vmcnt(0)
-      global_store_b32(addr=v[0], data=v[2], saddr=OFF),
-      s_endpgm(),
-    ]
-
-    asm = """
-v_bfe_u32 v1, v0, 10, 10
-s_load_b128 s[4:7], s[0:1], null
-v_and_b32_e32 v0, 0x3FF, v0
-s_mulk_i32 s3, 0x87
-v_mad_u64_u32 v[1:2], null, s2, 3, v[1:2]
-v_mul_u32_u24_e32 v0, 45, v0
-v_ashrrev_i32_e32 v2, 31, v1
-v_add3_u32 v0, v0, s3, v1
-v_lshlrev_b64 v[2:3], 2, v[1:2]
-v_ashrrev_i32_e32 v1, 31, v0
-v_lshlrev_b64 v[0:1], 2, v[0:1]
-s_waitcnt lgkmcnt(0)
-v_add_co_u32 v2, vcc_lo, s6, v2
-v_add_co_ci_u32_e32 v3, vcc_lo, s7, v3, vcc_lo
-v_add_co_u32 v0, vcc_lo, s4, v0
-global_load_b32 v2, v[2:3], off
-v_add_co_ci_u32_e32 v1, vcc_lo, s5, v1, vcc_lo
-s_waitcnt vmcnt(0)
-global_store_b32 v[0:1], v2, off
-s_endpgm
-"""
-    expected = llvm_assemble(asm)
-    for inst,rt in zip(program, asm.strip().split("\n")): print(f"{inst.disasm():50s} {rt}")
-    actual = b''.join(inst.to_bytes() for inst in program)
-    self.assertEqual(actual, expected)
-
-  def test_sop2_s_add_u32(self):
-    inst = SOP2(SOP2Op.S_ADD_U32, s[3], s[0], s[1])
-    expected = llvm_assemble("s_add_u32 s3, s0, s1")
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_vop2_v_and_b32_inline_const(self):
-    inst = v_and_b32_e32(v[0], 10, v[0])
-    expected = llvm_assemble("v_and_b32_e32 v0, 10, v0")
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_sopp_s_endpgm(self):
-    inst = s_endpgm()
-    expected = llvm_assemble("s_endpgm")
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_sop1_s_mov_b32(self):
-    inst = s_mov_b32(s[0], s[1])
-    expected = llvm_assemble("s_mov_b32 s0, s1")
-    self.assertEqual(inst.to_bytes(), expected)
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_roundtrip.py

@@ -1,234 +0,0 @@
-#!/usr/bin/env python3
-"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
-import unittest, io, sys, re
-from extra.assembly.rdna3.autogen import *
-from extra.assembly.rdna3.lib import Inst
-from extra.assembly.rdna3.asm import asm
-
-# Instruction format detection based on encoding bits
-def detect_format(data: bytes) -> type[Inst] | None:
-  """Detect instruction format from machine code bytes."""
-  if len(data) < 4: return None
-  word = int.from_bytes(data[:4], 'little')
-  enc_9bit = (word >> 23) & 0x1FF  # 9-bit encoding for SOP1/SOPC/SOPP
-  enc_8bit = (word >> 24) & 0xFF
-
-  # Check 9-bit encodings first (most specific)
-  if enc_9bit == 0x17D: return SOP1  # bits 31:23 = 101111101
-  if enc_9bit == 0x17E: return SOPC  # bits 31:23 = 101111110
-  if enc_9bit == 0x17F: return SOPP  # bits 31:23 = 101111111
-  # SOPK: bits 31:28 = 1011, bits 27:23 = opcode (check after SOP1/SOPC/SOPP)
-  if enc_8bit in range(0xB0, 0xC0): return SOPK
-  # SOP2: bits 31:23 in range 0x100-0x17C (0x80-0xBE in bits 31:24, but not SOPK)
-  if 0x80 <= enc_8bit <= 0x9F: return SOP2
-  # VOP1: bits 31:25 = 0111111 (0x3F)
-  if (word >> 25) == 0x3F: return VOP1
-  # VOPC: bits 31:25 = 0111110 (0x3E)
-  if (word >> 25) == 0x3E: return VOPC
-  # VOP2: bits 31:30 = 00
-  if (word >> 30) == 0: return VOP2
-
-  # Check 64-bit formats
-  if len(data) >= 8:
-    if enc_8bit in (0xD4, 0xD5, 0xD7): return VOP3
-    if enc_8bit == 0xD6: return VOP3SD
-    if enc_8bit == 0xCC: return VOP3P
-    if enc_8bit == 0xCD: return VINTERP
-    if enc_8bit in (0xC8, 0xC9): return VOPD
-    if enc_8bit == 0xF4: return SMEM
-    if enc_8bit == 0xD8: return DS
-    if enc_8bit in (0xDC, 0xDD, 0xDE, 0xDF): return FLAT
-    if enc_8bit in (0xE0, 0xE1, 0xE2, 0xE3): return MUBUF
-    if enc_8bit in (0xE8, 0xE9, 0xEA, 0xEB): return MTBUF
-
-  return None
-
-def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
-  """Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
-  old_stdout = sys.stdout
-  sys.stdout = io.StringIO()
-  compiler.disassemble(lib)
-  output = sys.stdout.getvalue()
-  sys.stdout = old_stdout
-
-  results = []
-  for line in output.splitlines():
-    if '//' not in line: continue
-    instr = line.split('//')[0].strip()
-    if not instr: continue
-    comment = line.split('//')[1].strip()
-    if ':' not in comment: continue
-    hex_str = comment.split(':')[1].strip().split()[0]
-    try:
-      machine_bytes = bytes.fromhex(hex_str)[::-1]  # big-endian to little-endian
-      results.append((instr, machine_bytes))
-    except ValueError:
-      continue
-  return results
-
-def compile_asm(instr: str, compiler=None) -> bytes | None:
-  """Compile a single instruction with llvm-mc and return the machine code bytes."""
-  import subprocess
-  try:
-    result = subprocess.run(
-      ['llvm-mc', '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
-      input=f".text\n{instr}\n", capture_output=True, text=True)
-    if result.returncode != 0: return None
-    # Parse encoding: [0x01,0x39,0x0a,0x7e]
-    for line in result.stdout.split('\n'):
-      if 'encoding:' in line:
-        enc = line.split('encoding:')[1].strip()
-        if enc.startswith('[') and enc.endswith(']'):
-          hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
-          return bytes.fromhex(hex_vals)
-  except Exception:
-    pass
-  return None
-
-class TestTinygradKernelRoundtrip(unittest.TestCase):
-  """Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
-
-  def _test_kernel_roundtrip(self, op_fn):
-    """Generate kernel from op_fn, test:
-    1. decode -> reencode matches original bytes
-    2. asm(disasm()) matches LLVM output
-    3. our disasm() matches LLVM's disassembly string exactly
-    """
-    from extra.assembly.rdna3.test.test_compare_emulators import get_kernels_from_tinygrad
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-
-    kernels, _, _ = get_kernels_from_tinygrad(op_fn)
-    compiler = HIPCompiler('gfx1100')
-
-    decode_passed, decode_failed, decode_skipped = 0, 0, 0
-    asm_passed, asm_failed, asm_skipped = 0, 0, 0
-    disasm_passed, disasm_failed, disasm_skipped = 0, 0, 0
-    decode_failures, asm_failures, disasm_failures = [], [], []
-
-    for ki, kernel in enumerate(kernels):
-      offset = 0
-      while offset < len(kernel.code):
-        remaining = kernel.code[offset:]
-        fmt = detect_format(remaining)
-        if fmt is None:
-          decode_skipped += 1
-          asm_skipped += 1
-          disasm_skipped += 1
-          offset += 4
-          continue
-
-        size = fmt._size()
-        if len(remaining) < size:
-          break
-
-        orig_bytes = remaining[:size]
-
-        # Test 1: decode -> reencode roundtrip
-        try:
-          decoded = fmt.from_bytes(orig_bytes)
-          reencoded = decoded.to_bytes()
-          if reencoded[:size] == orig_bytes:
-            decode_passed += 1
-          else:
-            decode_failed += 1
-            decode_failures.append(f"K{ki}@{offset}: {decoded.disasm()}: orig={orig_bytes.hex()} reenc={reencoded[:size].hex()}")
-
-          our_disasm = decoded.disasm()
-
-          # Test 2: asm(disasm()) matches LLVM output
-          try:
-            our_bytes = asm(our_disasm).to_bytes()
-            llvm_bytes = compile_asm(our_disasm, compiler)
-            if llvm_bytes is None:
-              asm_skipped += 1
-            elif our_bytes[:len(llvm_bytes)] == llvm_bytes:
-              asm_passed += 1
-            else:
-              asm_failed += 1
-              asm_failures.append(f"K{ki}@{offset}: '{our_disasm}': ours={our_bytes[:len(llvm_bytes)].hex()} llvm={llvm_bytes.hex()}")
-          except Exception:
-            asm_skipped += 1
-
-          # Test 3: our disasm() matches LLVM's disassembly string exactly
-          # Skip if instruction uses op_XX (unknown opcode) or looks malformed (many raw field values)
-          if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm):
-            disasm_skipped += 1
-          else:
-            try:
-              # Get LLVM's disassembly of our instruction
-              src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n  {our_disasm}\n"
-              lib = compiler.compile(src)
-              llvm_instrs = disassemble_lib(lib, compiler)
-              if llvm_instrs:
-                llvm_disasm = llvm_instrs[0][0]
-                if our_disasm == llvm_disasm:
-                  disasm_passed += 1
-                else:
-                  disasm_failed += 1
-                  disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm}'")
-              else:
-                disasm_skipped += 1
-            except Exception:
-              disasm_skipped += 1
-
-        except Exception:
-          decode_skipped += 1
-          asm_skipped += 1
-          disasm_skipped += 1
-
-        offset += size
-
-    print(f"decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
-    print(f"asm vs llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
-    print(f"disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
-    self.assertEqual(decode_failed, 0, f"Decode failures:\n" + "\n".join(decode_failures[:20]))
-    self.assertEqual(asm_failed, 0, f"Asm failures:\n" + "\n".join(asm_failures[:20]))
-    self.assertEqual(disasm_failed, 0, f"Disasm failures:\n" + "\n".join(disasm_failures[:20]))
-
-  # Basic unary ops
-  def test_neg(self): self._test_kernel_roundtrip(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
-  def test_relu(self): self._test_kernel_roundtrip(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
-  def test_exp(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).exp())
-  def test_log(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 3.0]).log())
-  def test_sin(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).sin())
-  def test_sqrt(self): self._test_kernel_roundtrip(lambda T: T([1.0, 4.0, 9.0]).sqrt())
-  def test_recip(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
-
-  # Binary ops
-  def test_add(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
-  def test_sub(self): self._test_kernel_roundtrip(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
-  def test_mul(self): self._test_kernel_roundtrip(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
-  def test_div(self): self._test_kernel_roundtrip(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
-  def test_max_binary(self): self._test_kernel_roundtrip(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
-
-  # Reductions
-  def test_sum_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).sum())
-  def test_max_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).max())
-  def test_mean_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(32).mean())
-
-  # Matmul
-  def test_gemm_4x4(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4) @ T.empty(4, 4))
-  def test_gemv(self): self._test_kernel_roundtrip(lambda T: T.empty(1, 16) @ T.empty(16, 16))
-
-  # Complex ops
-  def test_softmax(self): self._test_kernel_roundtrip(lambda T: T.empty(16).softmax())
-  def test_layernorm(self): self._test_kernel_roundtrip(lambda T: T.empty(8, 8).layernorm())
-
-  # Memory patterns
-  def test_contiguous(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
-  def test_reshape(self): self._test_kernel_roundtrip(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
-  def test_expand(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
-
-  # Cast ops
-  def test_cast_int(self): self._test_kernel_roundtrip(lambda T: T.empty(16).int().float())
-  def test_cast_half(self): self._test_kernel_roundtrip(lambda T: T.empty(16).half().float())
-
-  # Comparison ops
-  def test_cmp_lt(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
-  def test_where(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
-
-  # Fused ops
-  def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
-
-if __name__ == "__main__":
-  unittest.main()

extra/gemm/.gitignore

@@ -1,2 +1,3 @@
+*.s
 *.ll
 fp32_sgemm_amd

extra/gemm/asm/test.py

@@ -1,65 +0,0 @@
-# Run assembly on the AMD runtime and check correctness
-# VIZ=2 to profile
-import pathlib
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.engine.realize import ExecItem, CompiledRunner
-from tinygrad.renderer import ProgramSpec
-from tinygrad.uop.ops import track_rewrites, UOp
-from tinygrad.helpers import TracingKey
-
-fp = pathlib.Path(__file__).parent/"gemm.s"
-
-# ** generate inputs on CPU
-
-N = 8192
-scale = 10.0
-
-import torch
-torch.manual_seed(0)
-A = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
-B = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
-Bt = B.t().contiguous() # transpose B for the baseline gemm
-C_torch = A@Bt
-
-# ** copy buffers to AMD
-
-# input creation and validation run on the copy engine for simpler tracing
-
-def from_torch(t:torch.Tensor) -> Tensor:
-  return Tensor.from_blob(t.data_ptr(), t.shape, dtype=dtypes.bfloat16, device="cpu").to(Device.DEFAULT).realize()
-
-C_tiny = Tensor.matmul(from_torch(A), from_torch(Bt), dtype=dtypes.float32).cast(dtypes.bfloat16)
-C_asm = Tensor.empty_like(C_tiny)
-C_asm.uop.buffer.allocate()
-
-# ** run gemms
-
-@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name,), ret=ret))
-def get_asm_gemm(ast:UOp, fp:pathlib.Path) -> ProgramSpec:
-  src = fp.read_text()
-  lib = Device[Device.DEFAULT].compiler.compile(src)
-  return ProgramSpec("gemm", src, Device.DEFAULT, ast, lib=lib, global_size=[1024, 1, 1], local_size=[256, 1, 1], globals=[0, 1, 2])
-
-sched = C_tiny.schedule()
-assert len(sched) == 1
-eis:list[ExecItem] = [sched[-1].lower()]
-ast = eis[0].ast
-prg = get_asm_gemm(ast, fp)
-eis.append(ExecItem(ast, [C_asm.uop.buffer, from_torch(B).uop.buffer, from_torch(A).uop.buffer], prg=CompiledRunner(prg)))
-
-
-for ei in eis:
-  et = ei.run(wait=True)
-  print(f"{(N*N*N*2 / et)*1e-12:.2f} REAL TFLOPS")
-
-# ** correctness
-
-import ctypes
-
-def torch_bf16(t:Tensor) -> torch.tensor:
-  asm_out = t.to("cpu").realize().uop.buffer._buf
-  buf = (ctypes.c_uint16*C_asm.uop.size).from_address(asm_out.va_addr)
-  return torch.frombuffer(buf, dtype=torch.bfloat16, count=C_asm.uop.size).reshape(C_asm.shape)
-
-assert torch.allclose(torch_bf16(C_asm), C_torch, rtol=1e-2, atol=1e-3)
-assert torch.allclose(torch_bf16(C_tiny), C_torch, rtol=1e-2, atol=1e-3)

extra/gemm/asm/unpack_kd.py

@@ -1,179 +0,0 @@
-# unpack the complete kernel descriptor of an amdgpu ELF of for gfx950
-# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#code-object-v3-kernel-descriptor
-import struct, pathlib
-from tinygrad.runtime.support.elf import elf_loader
-
-def bits(x, lo, hi): return (x >> lo) & ((1 << (hi - lo + 1)) - 1)
-def assert_zero(x, lo, hi): assert bits(x, lo, hi) == 0
-
-with open(fp:=pathlib.Path(__file__).parent/"lib", "rb") as f:
-  lib = f.read()
-
-image, sections, relocs = elf_loader(lib)
-rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"))
-
-# rodata is exactly 64 bytes
-kd = image[rodata_entry:rodata_entry+64]
-desc = int.from_bytes(kd, byteorder="little")
-
-group_segment_fixed_size = bits(desc, 0, 31)
-private_segment_fixed_size = bits(desc, 32, 63)
-kernarg_size = bits(desc, 64, 95)
-reserved_127_96 = bits(desc, 96, 127)
-assert reserved_127_96 == 0
-
-print("GROUP_SEGMENT_FIXED_SIZE:", group_segment_fixed_size)
-print("PRIVATE_SEGMENT_FIXED_SIZE:", private_segment_fixed_size)
-print("KERNARG_SIZE:", kernarg_size)
-print("RESERVED 127:96:", reserved_127_96)
-
-entry_off = bits(desc, 128, 191)
-
-# sign-extend manually if needed
-if entry_off & (1 << 63):
-  entry_off -= 1 << 64
-
-print("KERNEL_CODE_ENTRY_BYTE_OFFSET:", entry_off)
-
-kd_addr = 0x1840
-entry_addr = kd_addr + entry_off
-
-print("Computed entry address: 0x%016x" % entry_addr)
-print("256B aligned:", entry_addr % 256 == 0)
-
-pgm_rsrc3 = bits(desc, 352, 383)
-pgm_rsrc1 = bits(desc, 384, 415)
-pgm_rsrc2 = bits(desc, 416, 447)
-
-print("COMPUTE_PGM_RSRC3: 0x%08x" % pgm_rsrc3)
-print("COMPUTE_PGM_RSRC1: 0x%08x" % pgm_rsrc1)
-print("COMPUTE_PGM_RSRC2: 0x%08x" % pgm_rsrc2)
-
-# rsrc 3
-
-accum_offset_raw = bits(pgm_rsrc3, 0, 5)
-assert_zero(pgm_rsrc3, 6, 15)
-tg_split = bits(pgm_rsrc3, 16, 16)
-accum_offset_vgprs = (accum_offset_raw + 1) * 4
-print("RSRC3.ACCUM_OFFSET (AccVGPR index):", accum_offset_vgprs)
-print("RSRC3.TG_SPLIT:", tg_split)
-
-# rsrc 1
-
-vgpr_gran = bits(pgm_rsrc1, 0, 5)
-sgpr_gran = bits(pgm_rsrc1, 6, 9)
-assert_zero(pgm_rsrc1, 27, 28)
-
-# NOTE: this is vgprs + agprs
-vgprs_used = (vgpr_gran + 1) * 8
-assert 0 <= vgprs_used <= 512
-
-k = sgpr_gran // 2
-sgprs_used = (k + 1) * 16
-
-print("RSRC1.VGPRS:", vgprs_used)
-print("RSRC1.SGPRS:", sgprs_used)
-
-assert_zero(pgm_rsrc1, 10, 11)
-
-float_round_mode_32 = bits(pgm_rsrc1, 12, 13)
-float_round_mode_16_64 = bits(pgm_rsrc1, 15, 14)
-float_denorm_mode_32 = bits(pgm_rsrc1, 16, 17)
-float_denorm_mode_16_64 = bits(pgm_rsrc1, 18, 19)
-
-priv = bits(pgm_rsrc1, 20, 20)
-assert priv == 0
-enable_dx10_clamp_wg_rr_en = bits(pgm_rsrc1, 21, 21)
-debug_mode = bits(pgm_rsrc1, 22, 22)
-enable_ieee_mode = bits(pgm_rsrc1, 23, 23)
-bulky = bits(pgm_rsrc1, 24, 24)
-assert bulky == 0
-cdbg_user = bits(pgm_rsrc1, 25, 25)
-assert cdbg_user == 0
-fp16_ovfl = bits(pgm_rsrc1, 26, 26)
-assert_zero(pgm_rsrc1, 27, 28)  # reserved
-assert_zero(pgm_rsrc1, 29, 29)  # WGP_MODE (reserved on gfx9)
-assert_zero(pgm_rsrc1, 30, 30)  # MEM_ORDERED (reserved on gfx9)
-assert_zero(pgm_rsrc1, 31, 31)  # FWD_PROGRESS (reserved on gfx9)
-
-# rsrc 2
-
-enable_private_segment = bits(pgm_rsrc2, 0, 0)  # SCRATCH_EN
-user_sgpr_count = bits(pgm_rsrc2, 1, 5)         # USER_SGPR
-enable_trap_handler = bits(pgm_rsrc2, 6, 6)     # TRAP_PRESENT (must be 0 here)
-assert enable_trap_handler == 0
-
-enable_sgpr_workgroup_id_x = bits(pgm_rsrc2, 7, 7)
-enable_sgpr_workgroup_id_y = bits(pgm_rsrc2, 8, 8)
-enable_sgpr_workgroup_id_z = bits(pgm_rsrc2, 9, 9)
-enable_sgpr_workgroup_info = bits(pgm_rsrc2, 10, 10)
-
-enable_vgpr_workitem_id = bits(pgm_rsrc2, 11, 12)  # TIDIG_CMP_CNT enum (0..3)
-
-enable_exception_address_watch = bits(pgm_rsrc2, 13, 13)
-assert enable_exception_address_watch == 0
-enable_exception_memory = bits(pgm_rsrc2, 14, 14)
-assert enable_exception_memory == 0
-
-granulated_lds_size = bits(pgm_rsrc2, 15, 23)
-assert granulated_lds_size == 0  # spec: must be 0; CP uses dispatch packet rounding
-
-enable_exception_fp_invalid = bits(pgm_rsrc2, 24, 24)
-enable_exception_fp_denorm_src = bits(pgm_rsrc2, 25, 25)
-enable_exception_fp_div0 = bits(pgm_rsrc2, 26, 26)
-enable_exception_fp_overflow = bits(pgm_rsrc2, 27, 27)
-enable_exception_fp_underflow = bits(pgm_rsrc2, 28, 28)
-enable_exception_fp_inexact = bits(pgm_rsrc2, 29, 29)
-enable_exception_int_div0 = bits(pgm_rsrc2, 30, 30)
-
-assert_zero(pgm_rsrc2, 31, 31)
-
-print("RSRC2.ENABLE_PRIVATE_SEGMENT:", enable_private_segment)
-print("RSRC2.USER_SGPR_COUNT:", user_sgpr_count)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_X:", enable_sgpr_workgroup_id_x)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Y:", enable_sgpr_workgroup_id_y)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Z:", enable_sgpr_workgroup_id_z)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_INFO:", enable_sgpr_workgroup_info)
-print("RSRC2.ENABLE_VGPR_WORKITEM_ID (enum):", enable_vgpr_workitem_id)
-
-print("RSRC2.EXC_FP_INVALID:", enable_exception_fp_invalid)
-print("RSRC2.EXC_FP_DENORM_SRC:", enable_exception_fp_denorm_src)
-print("RSRC2.EXC_FP_DIV0:", enable_exception_fp_div0)
-print("RSRC2.EXC_FP_OVERFLOW:", enable_exception_fp_overflow)
-print("RSRC2.EXC_FP_UNDERFLOW:", enable_exception_fp_underflow)
-print("RSRC2.EXC_FP_INEXACT:", enable_exception_fp_inexact)
-print("RSRC2.EXC_INT_DIV0:", enable_exception_int_div0)
-
-# user sgprs
-
-enable_sgpr_private_segment_buffer = bits(desc, 448, 448)
-enable_sgpr_dispatch_ptr = bits(desc, 449, 449)
-enable_sgpr_queue_ptr = bits(desc, 450, 450)
-enable_sgpr_kernarg_segment_ptr = bits(desc, 451, 451)
-enable_sgpr_dispatch_id = bits(desc, 452, 452)
-enable_sgpr_flat_scratch_init = bits(desc, 453, 453)
-enable_sgpr_private_segment_size = bits(desc, 454, 454)
-
-assert_zero(desc, 455, 457)
-
-print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER:", enable_sgpr_private_segment_buffer)
-print("DESC.ENABLE_SGPR_DISPATCH_PTR:", enable_sgpr_dispatch_ptr)
-print("DESC.ENABLE_SGPR_QUEUE_PTR:", enable_sgpr_queue_ptr)
-print("DESC.ENABLE_SGPR_KERNARG_SEGMENT_PTR:", enable_sgpr_kernarg_segment_ptr)
-print("DESC.ENABLE_SGPR_DISPATCH_ID:", enable_sgpr_dispatch_id)
-print("DESC.ENABLE_SGPR_FLAT_SCRATCH_INIT:", enable_sgpr_flat_scratch_init)
-print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_SIZE:", enable_sgpr_private_segment_size)
-
-assert_zero(desc, 458, 459)
-
-uses_dynamic_stack = bits(desc, 459, 460)
-print("DESC.USES_DYNAMIC_STACK:", uses_dynamic_stack)
-
-assert_zero(desc, 460, 463)
-kernarg_preload_spec_length = bits(desc, 464, 470)
-print("DESC.KERNARG_PRELOAD_SPEC_LENGTH:", kernarg_preload_spec_length)
-
-kernarg_preload_spec_offset = bits(desc, 471, 479)
-print("DESC.KERNARG_PRELOAD_SPEC_OFFSET:", kernarg_preload_spec_offset)
-
-assert_zero(desc, 480, 511)

extra/models/llama.py

@@ -245,11 +245,6 @@ def convert_from_huggingface(weights:dict[str, Tensor], n_layers: int, n_heads:
       continue
     sd[keymap[k]] = v
   for k,v in experts.items(): sd[k] = Tensor.stack(*[v[i] for i in range(len(v))])
-
-  # Handle tied embeddings (e.g., Llama 3.2 1B Instruct where lm_head shares weights with embed_tokens)
-  if "output.weight" not in sd and "tok_embeddings.weight" in sd:
-    sd["output.weight"] = sd["tok_embeddings.weight"]
-
   return sd
 
 def convert_from_gguf(weights:dict[str, Tensor], n_layers:int):

extra/optimization/extract_policynet.py

@@ -0,0 +1,114 @@
+import os, sys, sqlite3, pickle, random
+from tqdm import tqdm, trange
+from copy import deepcopy
+from tinygrad.nn import Linear
+from tinygrad.tensor import Tensor
+from tinygrad.nn.optim import Adam
+from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
+from tinygrad.codegen.opt.search import actions
+from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.helpers import getenv
+
+# stuff needed to unpack a kernel
+from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
+from tinygrad.dtype import dtypes
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View
+from tinygrad.uop.ops import Variable
+inf, nan = float('inf'), float('nan')
+from tinygrad.codegen.opt.kernel import Opt, OptOps
+
+INNER = 256
+class PolicyNet:
+  def __init__(self):
+    self.l1 = Linear(1021,INNER)
+    self.l2 = Linear(INNER,INNER)
+    self.l3 = Linear(INNER,1+len(actions))
+  def __call__(self, x):
+    x = self.l1(x).relu()
+    x = self.l2(x).relu().dropout(0.9)
+    return self.l3(x).log_softmax()
+
+def dataset_from_cache(fn):
+  conn = sqlite3.connect(fn)
+  cur = conn.cursor()
+  cur.execute("SELECT * FROM beam_search")
+  X,A = [], []
+  for f in tqdm(cur.fetchall()):
+    Xs,As = [], []
+    try:
+      lin = Kernel(eval(f[0]))
+      opts = pickle.loads(f[-1])
+      for o in opts:
+        Xs.append(lin_to_feats(lin, use_sts=True))
+        As.append(actions.index(o))
+        lin.apply_opt(o)
+      Xs.append(lin_to_feats(lin, use_sts=True))
+      As.append(0)
+    except Exception:
+      pass
+    X += Xs
+    A += As
+  return X,A
+
+if __name__ == "__main__":
+  if getenv("REGEN"):
+    X,V = dataset_from_cache(sys.argv[1] if len(sys.argv) > 1 else "/tmp/tinygrad_cache")
+    safe_save({"X": Tensor(X), "V": Tensor(V)}, "/tmp/dataset_policy")
+  else:
+    ld = safe_load("/tmp/dataset_policy")
+    X,V = ld['X'].numpy(), ld['V'].numpy()
+
+  print(X.shape, V.shape)
+  order = list(range(X.shape[0]))
+  random.shuffle(order)
+  X, V = X[order], V[order]
+
+  ratio = -256
+  X_test, V_test = Tensor(X[ratio:]), Tensor(V[ratio:])
+  X,V = X[:ratio], V[:ratio]
+  print(X.shape, V.shape)
+
+  net = PolicyNet()
+  #if os.path.isfile("/tmp/policynet.safetensors"): load_state_dict(net, safe_load("/tmp/policynet.safetensors"))
+  optim = Adam(get_parameters(net))
+
+  def get_minibatch(X,Y,bs):
+    xs, ys = [], []
+    for _ in range(bs):
+      sel = random.randint(0, len(X)-1)
+      xs.append(X[sel])
+      ys.append(Y[sel])
+    return Tensor(xs), Tensor(ys)
+
+  Tensor.training = True
+  losses = []
+  test_losses = []
+  test_accuracy = 0
+  test_loss = float('inf')
+  for i in (t:=trange(500)):
+    x,y = get_minibatch(X,V,bs=256)
+    out = net(x)
+    loss = out.sparse_categorical_crossentropy(y)
+    optim.zero_grad()
+    loss.backward()
+    optim.step()
+    cat = out.argmax(axis=-1)
+    accuracy = (cat == y).mean()
+    t.set_description(f"loss {loss.numpy():7.2f} accuracy {accuracy.numpy()*100:7.2f}%, test loss {test_loss:7.2f} test accuracy {test_accuracy*100:7.2f}%")
+
+    losses.append(loss.numpy().item())
+    test_losses.append(test_loss)
+    if i % 10:
+      out = net(X_test)
+      test_loss = out.sparse_categorical_crossentropy(V_test).square().mean().numpy().item()
+      cat = out.argmax(axis=-1)
+      test_accuracy = (cat == y).mean().numpy()
+
+  safe_save(get_state_dict(net), "/tmp/policynet.safetensors")
+
+  import matplotlib.pyplot as plt
+  plt.plot(losses[10:])
+  plt.plot(test_losses[10:])
+  plt.show()

extra/optimization/extract_sa_pairs.py

@@ -0,0 +1,129 @@
+import sys, sqlite3, pickle, math
+from collections import defaultdict
+from tqdm import tqdm, trange
+import numpy as np
+
+# stuff needed to unpack a kernel
+from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
+from tinygrad.dtype import dtypes
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View
+from tinygrad.uop.ops import Variable
+inf, nan = float('inf'), float('nan')
+from tinygrad.codegen.opt.kernel import Opt, OptOps
+
+# more stuff
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.search import actions
+from extra.optimization.helpers import lin_to_feats
+from extra.optimization.pretrain_valuenet import ValueNet
+from tinygrad.nn.optim import Adam
+from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
+import random
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import getenv
+
+def dataset_from_cache(fn):
+  conn = sqlite3.connect(fn)
+  cur = conn.cursor()
+  cur.execute("SELECT * FROM time_linearizer")
+  grouped = defaultdict(dict)
+  for f in tqdm(cur.fetchall()): grouped[f[0]][f[1:-1]] = pickle.loads(f[-1])
+
+  opts_to_outcome = {}
+
+  for ast,sk in grouped.items():
+    cnts = defaultdict(int)
+    for sks,tm in sk.items():
+      if sks[1] != 1: continue
+      opts = eval(sks[0])
+      cnts[(len(opts), sks[1])] += 1
+      opts_to_outcome[(ast, tuple(opts))] = tm
+    #print(cnts)
+
+  S,A,V = [], [], []
+  for ast,k in tqdm(opts_to_outcome):
+    if len(k) == 0: continue
+    old_tm = min(opts_to_outcome[(ast,k[:-1])])
+    new_tm = min(opts_to_outcome[(ast,k)])
+    if math.isinf(old_tm) or math.isinf(new_tm) or old_tm < 1e-9 or new_tm < 1e-9: continue
+    try:
+      lin = Kernel(eval(ast))
+    except Exception:
+      continue
+    lin.apply_opts(k[:-1])
+    act = k[-1]
+    log_ratio = math.log(old_tm/new_tm)
+    #print(f"ratio: {old_tm/new_tm:6.2f}x (log {log_ratio:5.2f}) from {str(act):50s} on {lin.colored_shape()}")
+    S.append(lin_to_feats(lin, use_sts=True))
+    A.append(actions.index(act))
+    V.append([log_ratio])  # NOTE: i have written the bug many times with this having the wrong dim
+
+  S, A, V = np.array(S), np.array(A), np.array(V, dtype=np.float32)
+  X = np.zeros((S.shape[0], S.shape[1]+len(actions)), dtype=np.float32)
+  X[:, :S.shape[1]] = S
+  X[range(S.shape[0]), S.shape[1]+A] = 1.0
+  return X, V
+
+def log_likelihood(x:Tensor, mu:Tensor, log_sigma:Tensor):
+  #print(x.shape, mu.shape, log_sigma.shape)
+  #return (x-mu).abs() * (-log_sigma).exp() + log_sigma
+  return (x-mu).square() * (-2*log_sigma).exp() / 2 + log_sigma
+
+if __name__ == "__main__":
+  if getenv("REGEN"):
+    X,V = dataset_from_cache(sys.argv[1] if len(sys.argv) > 1 else "/tmp/tinygrad_cache")
+    safe_save({"X": Tensor(X), "V": Tensor(V)}, "/tmp/dataset")
+  else:
+    ld = safe_load("/tmp/dataset")
+    X,V = ld['X'].numpy(), ld['V'].numpy()
+
+  print(X.shape, V.shape)
+  order = list(range(X.shape[0]))
+  random.shuffle(order)
+  X, V = X[order], V[order]
+
+  ratio = -512
+  X_test, V_test = Tensor(X[ratio:]), Tensor(V[ratio:])
+  X,V = X[:ratio], V[:ratio]
+  print(X.shape, V.shape)
+
+  #print(X[0], V[0])
+  #print(X[-1], V[-1])
+  print(X.shape)
+
+  net = ValueNet(X.shape[1], 2)
+  optim = Adam(get_parameters(net))
+
+  def get_minibatch(X,Y,bs):
+    xs, ys = [], []
+    #random.seed(1337)
+    for _ in range(bs):
+      sel = random.randint(0, len(X)-1)
+      xs.append(X[sel])
+      ys.append(Y[sel])
+    return Tensor(xs), Tensor(ys)
+
+  Tensor.training = True
+  losses = []
+  test_losses = []
+  test_loss = float('inf')
+  for i in (t:=trange(2000)):
+    x,y = get_minibatch(X,V,bs=256)
+    out = net(x)
+    #loss = (out-y).square().mean()
+    loss = log_likelihood(y, out[:, 0:1], out[:, 1:2]).mean()
+    optim.zero_grad()
+    loss.backward()
+    optim.step()
+    t.set_description(f"loss {loss.numpy():7.2f}, test loss {test_loss:7.2f}")
+    losses.append(loss.numpy().item())
+    test_losses.append(test_loss)
+    if i % 10: test_loss = (net(X_test)[:, 0:1]-V_test).square().mean().numpy().item()
+
+  safe_save(get_state_dict(net), "/tmp/qnet.safetensors")
+
+  import matplotlib.pyplot as plt
+  plt.plot(losses[20:])
+  plt.plot(test_losses[20:])
+  plt.show()

extra/optimization/helpers.py

@@ -0,0 +1,124 @@
+# stuff needed to unpack a kernel
+from tinygrad import Variable
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.dtype import dtypes, PtrDType
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View
+from tinygrad.helpers import getenv
+from tinygrad.engine.realize import get_program
+inf, nan = float('inf'), float('nan')
+UOps = Ops
+
+# kernel unpacker
+from tinygrad.codegen.opt.kernel import Kernel
+def ast_str_to_ast(ast_str:str) -> UOp: return eval(ast_str)
+def ast_str_to_lin(ast_str:str, opts=None): return Kernel(ast_str_to_ast(ast_str), opts=opts)
+def kern_str_to_lin(kern_str:str, opts=None):
+  (ast, applied_opts,) = eval(kern_str)
+  k = Kernel(ast, opts=opts)
+  k.apply_opts(applied_opts)
+  return k
+
+# load worlds, a dataset of about 12k kernels
+import gzip
+from pathlib import Path
+import random
+from tinygrad.helpers import dedup, DEBUG
+def load_worlds(filter_reduce=True, filter_noimage=True, filter_novariable=True):
+  fn = Path(__file__).parent.parent / "datasets/sops.gz"
+  ast_strs = dedup(gzip.open(fn).read().decode('utf-8').strip().split("\n"))
+  assert len(ast_strs) >= getenv("MIN_ASTS", 1000), f"dataset size = {len(ast_strs)} is too small"
+  if DEBUG >= 1: print(f"loaded {len(ast_strs)=} before filters")
+  if filter_reduce: ast_strs = [x for x in ast_strs if "REDUCE_AXIS" in x]
+  if filter_noimage: ast_strs = [x for x in ast_strs if "dtypes.image" not in x]
+  if filter_novariable: ast_strs = [x for x in ast_strs if "DEFINE_VAR" not in x]
+  if DEBUG >= 1: print(f"loaded {len(ast_strs)=} after filters {filter_reduce=}, {filter_noimage=}, {filter_novariable=}")
+  random.seed(1337)
+  random.shuffle(ast_strs)
+  return ast_strs
+
+def assert_same_lin(l1, l2):
+  assert l1.colored_shape() == l2.colored_shape()
+  assert all(x==y for x,y in zip(l1.sts, l2.sts))
+
+# get features
+import math
+
+MAX_DIMS = 16
+MAX_BUFS = 9
+def lin_to_feats(lin:Kernel, use_sts=True):
+  assert lin.shape_len < MAX_DIMS, "too many dims"
+
+  all_colors = ["blue", "cyan", "white", "green", "red", "magenta", "yellow"]
+  lc = [all_colors.index(x) for x in lin.colors()]
+
+  ret = []
+  # before, some generic linearizer stuff
+  ret.append(lin.upcasted)
+  ret.append(lin.local_dims)
+
+  # first, the full shape, including the colors
+  for s,os,c in zip(lin.full_shape,lin.output_shape,lc):
+    if isinstance(s, UOp):
+      ret.append(False)
+      ret += [0]*9
+    else:
+      ret.append(True)
+      ret.append(math.log2(s))
+      ret.append(min(33, s))
+      ret.append(math.log2(os))
+      ret.append(min(33, os))
+      ret.append(s%2 == 0)
+      ret.append(s%3 == 0)
+      ret.append(s%4 == 0)
+      ret.append(s%8 == 0)
+      ret.append(s%16 == 0)
+    cc = [0]*7
+    cc[c] = 1
+    ret += cc
+  ret += [0] * (17*(MAX_DIMS-len(lin.full_shape)))
+  ret = [float(x) for x in ret]
+
+  if use_sts:
+    my_sts = dedup([(x.shape == lin.full_shape, x.is_expanded(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
+    assert len(my_sts) < MAX_BUFS
+    sts_len = 3 + 5*MAX_DIMS
+    for s in my_sts:
+      ret.append(s[0])  # reduce
+      ret.append(s[2])  # has mask
+      ret.append(s[3])  # len views
+      for d in s[1]:
+        ret.append(d is None)
+        ret.append(d == 0)
+        ret.append(d == 1)
+        ret.append(min(33, d) if d is not None else -1)
+        if d is not None and d >= 1: ret.append(math.log2(d))
+        else: ret.append(-1)
+      ret += [0] * (5*(MAX_DIMS - len(s[1])))
+    ret += [0] * (sts_len*(MAX_BUFS - len(my_sts)))
+    assert len(ret) == 1021, f"wrong len {len(ret)}"
+  else:
+    assert len(ret) == 274, f"wrong len {len(ret)}"
+  return ret
+
+from tinygrad.device import Device, Buffer
+from tinygrad.codegen.opt.search import _ensure_buffer_alloc, _time_program
+from tinygrad.helpers import to_function_name, CACHELEVEL, diskcache_get, diskcache_put
+
+def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_global_size=65536, cnt=3, disable_cache=False, clear_l2=False) -> float:  # noqa: E501
+  key = {"ast": lin.ast.key, "opts": str(lin.applied_opts), "allow_test_size": allow_test_size,
+         "max_global_size": max_global_size, "clear_l2": clear_l2, "device": lin.opts.device, "suffix": lin.opts.suffix}
+  if not disable_cache and CACHELEVEL >= 2 and (val:=diskcache_get("time_linearizer", key)) is not None: return min(val)
+
+  dev = Device[lin.opts.device]
+  assert dev.compiler is not None
+
+  rawbufs = _ensure_buffer_alloc(rawbufs)
+  var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
+  p = get_program(lin.get_optimized_ast(), lin.opts)
+  tms = _time_program(p, dev.compiler.compile(p.src), var_vals, rawbufs,
+                      max_global_size=max_global_size if allow_test_size else None, clear_l2=clear_l2, cnt=cnt, name=to_function_name(lin.name))
+
+  if CACHELEVEL >= 2: diskcache_put("time_linearizer", key, tms)
+  return min(tms)

extra/optimization/pretrain_valuenet.py

@@ -0,0 +1,88 @@
+from tinygrad.codegen.opt.kernel import Kernel
+from tqdm import tqdm, trange
+import math
+import random
+from tinygrad.tensor import Tensor
+from tinygrad.nn import Linear
+from tinygrad.nn.optim import Adam
+from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
+
+# stuff needed to unpack a kernel
+from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
+from tinygrad.dtype import dtypes
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View
+from tinygrad.uop.ops import Variable
+inf, nan = float('inf'), float('nan')
+from tinygrad.codegen.opt.kernel import Opt, OptOps
+
+from extra.optimization.helpers import lin_to_feats, MAX_DIMS
+
+# NOTE: this is not real value of the state, it's just a prediction of the runtime
+INNER = 512
+class ValueNet:
+  def __init__(self, feats=240, out=1):
+    self.l1 = Linear(feats,INNER)
+    self.l2 = Linear(INNER,INNER)
+    self.l3 = Linear(INNER,INNER)
+    self.l4 = Linear(INNER,out)
+  def __call__(self, x):
+    x = self.l1(x).relu()
+    x = self.l2(x).relu()
+    x = self.l3(x).relu().dropout(0.8)
+    return self.l4(x)
+
+if __name__ == "__main__":
+  net = ValueNet()
+  optim = Adam(get_parameters(net))
+
+  TEST_SIZE = 256
+
+  dset = open("/tmp/logtm").read().strip().split("\n")
+  random.seed(1337)
+  random.shuffle(dset)
+
+  X,Y = [], []
+  for i,x in enumerate(tqdm(dset)):
+    ast, opts, tms = eval(x)
+    lin = Kernel(ast)
+    for o in opts: lin.apply_opt(o)
+    if lin.shape_len >= MAX_DIMS: continue
+    if min(tms) == float('inf'): continue
+    X.append(lin_to_feats(lin))
+    Y.append([math.log(min(tms))])
+  print(f"got {len(X)} samples")
+
+  X_test,Y_test = Tensor(X[-TEST_SIZE:]), Tensor(Y[-TEST_SIZE:])
+  X,Y = X[:-TEST_SIZE], Y[:-TEST_SIZE]
+
+  def get_minibatch(X,Y,bs):
+    xs, ys = [], []
+    for _ in range(bs):
+      sel = random.randint(0, len(X)-1)
+      xs.append(X[sel])
+      ys.append(Y[sel])
+    return Tensor(xs), Tensor(ys)
+
+  Tensor.training = True
+  losses = []
+  test_losses = []
+  test_loss = float('inf')
+  for i in (t:=trange(2000)):
+    x,y = get_minibatch(X,Y,bs=256)
+    out = net(x)
+    loss = (out-y).square().mean()
+    optim.zero_grad()
+    loss.backward()
+    optim.step()
+    t.set_description(f"loss {loss.numpy():7.2f}, test loss {test_loss:7.2f}")
+    losses.append(loss.numpy().item())
+    test_losses.append(test_loss)
+    if i % 10: test_loss = (net(X_test)-Y_test).square().mean().numpy().item()
+
+  safe_save(get_state_dict(net), "/tmp/valuenet.safetensors")
+
+  import matplotlib.pyplot as plt
+  plt.plot(losses[200:])
+  plt.plot(test_losses[200:])
+  plt.show()

extra/remu/src/lib.rs

@@ -1,5 +1,4 @@
-use crate::state::StateSnapshot;
-use crate::work_group::{WaveContext, WorkGroup};
+use crate::work_group::WorkGroup;
 use std::os::raw::c_char;
 use std::slice;
 mod helpers;
@@ -31,47 +30,3 @@ pub extern "C" fn run_asm(lib: *const c_char, lib_sz: u32, gx: u32, gy: u32, gz:
     }
     0
 }
-
-// FFI functions for single-stepping comparison tests
-
-#[no_mangle]
-pub extern "C" fn wave_create(lib: *const c_char, lib_sz: u32, n_lanes: u32) -> *mut WaveContext {
-    if lib.is_null() || (lib_sz % 4) != 0 { return std::ptr::null_mut(); }
-    let kernel = unsafe { slice::from_raw_parts(lib as *const u32, (lib_sz / 4) as usize).to_vec() };
-    Box::into_raw(Box::new(WaveContext::new(kernel, n_lanes as usize)))
-}
-
-#[no_mangle]
-pub extern "C" fn wave_step(ctx: *mut WaveContext) -> i32 {
-    if ctx.is_null() { return -99; }
-    unsafe { (*ctx).step() }
-}
-
-#[no_mangle]
-pub extern "C" fn wave_get_snapshot(ctx: *const WaveContext, out: *mut StateSnapshot) {
-    if ctx.is_null() || out.is_null() { return; }
-    unsafe { *out = (*ctx).get_snapshot(); }
-}
-
-#[no_mangle]
-pub extern "C" fn wave_set_sgpr(ctx: *mut WaveContext, idx: u32, val: u32) {
-    if ctx.is_null() || idx >= 128 { return; }
-    unsafe { (*ctx).scalar_reg[idx as usize] = val; }
-}
-
-#[no_mangle]
-pub extern "C" fn wave_set_vgpr(ctx: *mut WaveContext, lane: u32, idx: u32, val: u32) {
-    if ctx.is_null() || lane >= 32 || idx >= 256 { return; }
-    unsafe { (*ctx).vec_reg.get_lane_mut(lane as usize)[idx as usize] = val; }
-}
-
-#[no_mangle]
-pub extern "C" fn wave_init_lds(ctx: *mut WaveContext, size: u32) {
-    if ctx.is_null() { return; }
-    unsafe { (*ctx).lds.data.resize(size as usize, 0); }
-}
-
-#[no_mangle]
-pub extern "C" fn wave_free(ctx: *mut WaveContext) {
-    if !ctx.is_null() { unsafe { drop(Box::from_raw(ctx)); } }
-}

extra/remu/src/state.rs

@@ -96,24 +96,6 @@ impl WaveValue {
     }
 }
 
-/// C-compatible state snapshot for FFI - used for comparing emulator states
-#[repr(C)]
-#[derive(Clone, Debug)]
-pub struct StateSnapshot {
-    pub pc: u32,
-    pub scc: u32,
-    pub vcc: u32,
-    pub exec_mask: u32,
-    pub sgpr: [u32; 128],
-    pub vgpr: [[u32; 256]; 32],
-}
-
-impl StateSnapshot {
-    pub fn new() -> Self {
-        Self { pc: 0, scc: 0, vcc: 0, exec_mask: 0, sgpr: [0; 128], vgpr: [[0; 256]; 32] }
-    }
-}
-
 #[derive(Clone, Debug)]
 pub struct VecDataStore {
     pub data: Vec<u8>,

extra/remu/src/work_group.rs

@@ -1,5 +1,5 @@
 use crate::helpers::{colored, DEBUG};
-use crate::state::{Register, StateSnapshot, VecDataStore, WaveValue, VGPR};
+use crate::state::{Register, VecDataStore, WaveValue, VGPR};
 use crate::thread::{Thread, END_PRG, SGPR_COUNT};
 use std::collections::HashMap;
 
@@ -28,96 +28,6 @@ struct WaveState {
 
 const SYNCS: [u32; 4] = [0xBF89FC07, 0xBC7C0000, 0xBF890007, 0xbFB60003];
 const S_BARRIER: u32 = 0xBFBD0000;
-
-/// Context for single-stepping through a wave - holds all mutable state
-pub struct WaveContext {
-    pub kernel: Vec<u32>,
-    pub scalar_reg: [u32; SGPR_COUNT],
-    pub scc: u32,
-    pub pc: usize,
-    pub vec_reg: VGPR,
-    pub vcc: WaveValue,
-    pub exec: WaveValue,
-    pub lds: VecDataStore,
-    pub sds: HashMap<usize, VecDataStore>,
-    pub n_lanes: usize,
-}
-
-impl WaveContext {
-    pub fn new(kernel: Vec<u32>, n_lanes: usize) -> Self {
-        let active = (!0u32).wrapping_shr(32 - (n_lanes as u32));
-        Self {
-            kernel,
-            scalar_reg: [0; SGPR_COUNT],
-            scc: 0,
-            pc: 0,
-            vec_reg: VGPR::new(),
-            vcc: WaveValue::new(0, n_lanes),
-            exec: WaveValue::new(active, n_lanes),
-            lds: VecDataStore::new(),
-            sds: (0..=31).map(|i| (i, VecDataStore::new())).collect(),
-            n_lanes,
-        }
-    }
-
-    /// Execute a single instruction. Returns: 0=continue, -1=endpgm, -2=barrier, 1=done (pc past program), negative=error
-    pub fn step(&mut self) -> i32 {
-        if self.pc >= self.kernel.len() { return 1; }
-        if self.kernel[self.pc] == END_PRG { return -1; }
-        if self.kernel[self.pc] == S_BARRIER { self.pc += 1; return -2; }
-        // Skip sync/nop instructions
-        if SYNCS.contains(&self.kernel[self.pc]) || self.kernel[self.pc] >> 20 == 0xbf8 || self.kernel[self.pc] == 0x7E000000 {
-            self.pc += 1;
-            return 0;
-        }
-
-        let mut sgpr_co = None;
-        for lane_id in 0..self.n_lanes {
-            self.vec_reg.default_lane = Some(lane_id);
-            self.vcc.default_lane = Some(lane_id);
-            self.exec.default_lane = Some(lane_id);
-            let mut thread = Thread {
-                scalar_reg: &mut self.scalar_reg,
-                scc: &mut self.scc,
-                vec_reg: &mut self.vec_reg,
-                vcc: &mut self.vcc,
-                exec: &mut self.exec,
-                lds: &mut self.lds,
-                sds: &mut self.sds.get_mut(&lane_id).unwrap(),
-                pc_offset: 0,
-                stream: self.kernel[self.pc..].to_vec(),
-                scalar: false,
-                simm: None,
-                warp_size: self.n_lanes,
-                sgpr_co: &mut sgpr_co,
-            };
-            if let Err(e) = thread.interpret() { return e; }
-            if thread.scalar {
-                self.pc = ((self.pc as isize) + 1 + (thread.pc_offset as isize)) as usize;
-                break;
-            }
-            if lane_id == self.n_lanes - 1 {
-                self.pc = ((self.pc as isize) + 1 + (thread.pc_offset as isize)) as usize;
-            }
-        }
-        if self.vcc.mutations.is_some() { self.vcc.apply_muts(); self.vcc.mutations = None; }
-        if self.exec.mutations.is_some() { self.exec.apply_muts(); self.exec.mutations = None; }
-        if let Some((idx, mut wv)) = sgpr_co.take() { wv.apply_muts(); self.scalar_reg[idx] = wv.value; }
-        0
-    }
-
-    pub fn get_snapshot(&self) -> StateSnapshot {
-        let mut snap = StateSnapshot::new();
-        snap.pc = self.pc as u32;
-        snap.scc = self.scc;
-        snap.vcc = self.vcc.value;
-        snap.exec_mask = self.exec.value;
-        snap.sgpr = self.scalar_reg;
-        for lane in 0..32 { snap.vgpr[lane] = self.vec_reg.get_lane(lane); }
-        snap
-    }
-}
-
 impl<'a> WorkGroup<'a> {
     pub fn new(dispatch_dim: u32, id: [u32; 3], launch_bounds: [u32; 3], kernel: &'a Vec<u32>, kernel_args: *const u64) -> Self {
         Self { dispatch_dim, id, kernel, launch_bounds, kernel_args, lds: VecDataStore::new(), wave_state: HashMap::new() }

extra/remu/test/hwtest.py

@@ -1,37 +1,32 @@
-# ruff: noqa: F405, F403
-# allow define from star imports
-
 import numpy as np
 import unittest
-import subprocess, struct, math, textwrap
+import subprocess, struct, math
 from tinygrad import Tensor, dtypes, Device, UOp
-from tinygrad.uop.ops import Ops
 from tinygrad.helpers import getenv
 from tinygrad.runtime.support.compiler_amd import amdgpu_disassemble
 from tinygrad.renderer import ProgramSpec
 from tinygrad.engine.realize import CompiledRunner
 
-from extra.assembly.rdna3.autogen import *
-from extra.assembly.rdna3.asm import waitcnt
-from test.testextra.test_cfg_viz import template
-
-def get_output(asm:list, n_threads:int=1, vdst:VGPR=v[1]):
-  out = Tensor([0]*n_threads, dtype=dtypes.uint32).realize()
-  src = "\n".join(inst.disasm() for inst in [
-    s_load_b64(s[0:1], s[0:1], NULL),
-    *asm,
-    v_lshlrev_b32_e32(v[0], 2, v[0]),
-    s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
-    #global_store_b32(v[0], v[1], s[0:1]),
-    global_store_b32(addr=v[0], data=vdst, saddr=s[0:1]),
-    s_endpgm()
-  ])
-  prg = ProgramSpec("test", template.replace("fn_name", "test").replace("INSTRUCTION", textwrap.dedent(src)), Device.DEFAULT, UOp(Ops.SINK),
-                    global_size=[1, 1, 1], local_size=[n_threads, 1, 1], globals=[0])
+def get_output(asm:str, n_threads:int=1):
+  input_asm = "\n".join([ln if ln.strip().startswith('asm volatile') else f'asm volatile("{ln.strip().lstrip()}" : "+v"(a), "+v"(b));'
+                         for ln in asm.strip().splitlines() if ln.strip()])
+  src = f"""
+  typedef long unsigned int size_t;
+  extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
+  extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, {n_threads}))) test(unsigned int* data0_1) {{
+    int l = __ockl_get_local_id(0);
+    unsigned a = 0, b = 0, c = 0;
+    {input_asm}
+    unsigned res;
+    asm volatile("v_mov_b32 %0, %1" : "=v"(res) : "v"(a));
+    *(data0_1+l) = res;
+  }}"""
+  t = Tensor.zeros(n_threads, dtype=dtypes.uint32).contiguous().realize()
+  prg = ProgramSpec("test", src, Device.DEFAULT, UOp.sink(t), global_size=[1, 1, 1], local_size=[n_threads, 1, 1])
   car = CompiledRunner(prg)
   if getenv("PRINT_ASM"): amdgpu_disassemble(car.lib)
-  car([out.uop.buffer], {}, wait=True)
-  return out.tolist()
+  car([t.uop.buffer], {}, wait=True)
+  return t.numpy()
 
 def f16_to_bits(x:float) -> int: return struct.unpack('<H', struct.pack('<e', x))[0]
 def f32_from_bits(x:int) -> float: return struct.unpack('<f', struct.pack('<I', x))[0]
@@ -42,89 +37,83 @@ class TestHW(unittest.TestCase):
   def setUp(self):
     if getenv("MOCKGPU"): subprocess.run(["cargo", "build", "--release", "--manifest-path", "./extra/remu/Cargo.toml"], check=True)
 
-  def test_simple_v_mov(self):
-    out = get_output([
-      v_mov_b32_e32(v[1], 2),
-    ])
-    self.assertEqual(out, [2])
-
-  def test_simple_s_mov(self):
-    out = get_output([
-      s_mov_b32(s[7], 0x7fffffff),
-      v_mov_b32_e32(v[1], s[7]),
-    ])
-    self.assertEqual(out, [0x7fffffff])
+  def test_simple(self):
+    out = get_output("""
+    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 %2 %1
+    """)[0]
+    np.testing.assert_equal(out, 42)
 
   def test_exec_mov(self):
-    out = get_output([
-      v_mov_b32_e32(v[1], 42),
-      s_mov_b32(EXEC_LO, 0b10),
-      v_mov_b32_e32(v[1], 10),
-      s_mov_b32(EXEC_LO, 0b11),
-    ], n_threads=2)
+    out = get_output("""
+    v_mov_b32_e32 %1 42
+    s_mov_b32_e32 exec_lo 0b10
+    v_mov_b32_e32 %1 10
+    s_mov_b32_e32 exec_lo 0b11
+    v_mov_b32_e32 %2 %1
+    """, n_threads=2)
     np.testing.assert_equal(out, [42, 10])
 
   def test_exec_cmp_vopc(self):
-    out = get_output([
-      s_mov_b32(VCC_LO, 0), # reset vcc
-      v_mov_b32_e32(v[1], 42),
-      v_mov_b32_e32(v[2], 10),
-      s_mov_b32(EXEC_LO, 0b01),
-      v_cmp_ne_u32_e32(v[1], v[2]),
-      s_mov_b32(EXEC_LO, 0b11),
-      v_mov_b32_e32(v[1], VCC_LO),
-    ], n_threads=2)[0]
-    np.testing.assert_equal(out, 1)
+    out = get_output("""
+    s_mov_b32 vcc_lo 0 // reset vcc
+    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 %2 10
+    s_mov_b32_e32 exec_lo 0b01
+    v_cmp_ne_u32 %1 %2
+    s_mov_b32_e32 exec_lo 0b11
+    v_mov_b32_e32 %2 vcc_lo
+    """, n_threads=2)
+    np.testing.assert_equal(out, 0b01)
 
   def test_exec_cmpx_vop3(self):
-    out = get_output([
-      s_mov_b32(EXEC_LO, 0b11),
-      v_mov_b32_e32(v[1], 42),
-      v_mov_b32_e32(v[2], 10),
-      s_mov_b32(EXEC_LO, 0b01),
-      v_cmpx_ne_u32_e32(v[1], v[2]),
-      s_mov_b32(s[10], EXEC_LO),
-      s_mov_b32(EXEC_LO, 0b11),
-      v_mov_b32_e32(v[1], s[10]),
-    ], n_threads=2)[0]
+    out = get_output("""
+    s_mov_b32_e32 exec_lo 0b11
+    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 %2 10
+    s_mov_b32_e32 exec_lo 0b01
+    v_cmpx_ne_u32 %1 %2
+    s_mov_b32_e32 s10 exec_lo
+    s_mov_b32_e32 exec_lo 0b11
+    v_mov_b32_e32 %2 s10
+    """, n_threads=2)[0]
     np.testing.assert_equal(out & 0b11, 0b01)
 
   def test_fmac_vop3_modifier(self):
-    init_state = [
-      v_mov_b32_e32(a:=v[1], f16_to_bits(4.0)),
-      v_mov_b32_e32(b:=v[2], f16_to_bits(3.0)),
-      v_mov_b32_e32(c:=v[3], f16_to_bits(2.0)),
-    ]
-    def run_fmac(a, b): return get_output(init_state+[v_fmac_f16_e64(c, a, b)], vdst=c)[0]
-    self.assertEqual(run_fmac(a, b), f16_to_bits(14.0))
-    self.assertEqual(run_fmac(a, -b), f16_to_bits(-10.0))
-    self.assertEqual(run_fmac(-a, -b), f16_to_bits(14.0))
+    init_state = f"""
+    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(4.0)}" : "+v"(a));
+    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(3.0)}" : "+v"(b));
+    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(2.0)}" : "+v"(c));
+    """
+    mov = """asm volatile("v_mov_b32_e32 %1, %2" : "+v"(c), "+v"(a));"""
+    def fmac(a, b, c): return f"""asm volatile("v_fmac_f16_e64 {c}, {a}, {b}" : "+v"(c) : "v"(a), "v"(b));"""+"\n"+mov
+    self.assertEqual(get_output(init_state+"\n"+fmac("%1", "%2", "%3")), f16_to_bits(14.))
+    self.assertEqual(get_output(init_state+"\n"+fmac("%1", "-%2", "%3")), f16_to_bits(-10.))
+    self.assertEqual(get_output(init_state+"\n"+fmac("-%1", "-%2", "%3")), f16_to_bits(14.))
 
   def test_s_abs_i32(self):
-    def check(x, y, dst=s[10], scc=0):
-      for reg,val in [(dst, y), (SCC, scc)]:
-        self.assertEqual(get_output([
-          s_mov_b32(dst, x),
-          s_abs_i32(dst, dst),
-          v_mov_b32_e32(v[1], reg)
-        ])[0], val)
-
-    check(0x00000001, 0x00000001, scc=1)
-    check(0x7fffffff, 0x7fffffff, scc=1)
-    check(0x80000000, 0x80000000, scc=1)
-    check(0x80000001, 0x7fffffff, scc=1)
-    check(0x80000002, 0x7ffffffe, scc=1)
-    check(0xffffffff, 0x00000001, scc=1)
-    check(0, 0, scc=0)
+    def s_abs_i32(x, y, dst="s10", scc=0):
+      for reg,val in [(dst, y), ("scc", scc)]:
+        self.assertEqual(get_output(f"""
+        s_mov_b32_e32 {dst} {x}
+        s_abs_i32 {dst} {dst}
+        v_mov_b32_e32 %2 {reg}
+        """)[0], val)
+    s_abs_i32(0x00000001, 0x00000001, scc=1)
+    s_abs_i32(0x7fffffff, 0x7fffffff, scc=1)
+    s_abs_i32(0x80000000, 0x80000000, scc=1)
+    s_abs_i32(0x80000001, 0x7fffffff, scc=1)
+    s_abs_i32(0x80000002, 0x7ffffffe, scc=1)
+    s_abs_i32(0xffffffff, 0x00000001, scc=1)
+    s_abs_i32(0, 0, scc=0)
 
   def test_v_rcp_f32_neg_vop3(self):
     def v_neg_rcp_f32(x:float, y:float):
-      out = get_output([
-        v_mov_b32_e32(v[2], f32_to_bits(x)),
-        v_rcp_f32_e64(v[2], -v[2]),
-      ], vdst=v[2])[0]
+      out = get_output(f"""
+      v_mov_b32_e32 %2 {f32_to_bits(x)}
+      v_rcp_f32_e64 %2, -%2
+      """)[0]
       assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
-
     v_neg_rcp_f32(math.inf, -0.0)
     v_neg_rcp_f32(-math.inf, 0.0)
     v_neg_rcp_f32(0.0, -math.inf)
@@ -133,28 +122,25 @@ class TestHW(unittest.TestCase):
     v_neg_rcp_f32(2.0, -0.5)
 
   def test_v_cndmask_b32_neg(self):
-    def v_neg(x:float, y:float):
-      out = get_output([
-        v_mov_b32_e32(v[1], f32_to_bits(x)),
-        s_mov_b32(s[10], 1),
-        v_cndmask_b32_e64(v[1], v[1], -v[1], s[10]),
-      ])[0]
+    def v_neg(x:int|float, y:float):
+      # always pick -v1
+      out = get_output(f"""
+      v_mov_b32_e32 %2 {f32_to_bits(x)}
+      s_mov_b32_e32 s10 1
+      v_cndmask_b32 %2, %2, -%2 s10
+      """)[0]
       assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
-
     v_neg(-0.0, 0.0)
     v_neg(0.0, -0.0)
     v_neg(2.0, -2.0)
     v_neg(math.inf, -math.inf)
     v_neg(-math.inf, math.inf)
 
-  @unittest.skip("how does VOPD work in the dsl")
   def test_v_subrev_wrap(self):
-    out = get_output([
-      #v_dual_mov_b32(v[1], 0xffffffff, v[2], 0x0),
-      #v_dual_mov_b32(vdstx=v[1], srcx=0xffffffff, vdsty=v[2], srcy=0x0),
-      #VOPD(opx=VOPDOp.V_DUAL_MOV_B32, opy=VOPDOp.V_DUAL_MOV_B32, vdstx=v[1], srcx=0xffffffff, vdsty=v[2], srcy=0x0),
-      v_subrev_co_u32(v[2], VCC_LO, v[2], v[1]),
-    ], vdst=v[2])[0]
+    out = get_output("""
+    v_dual_mov_b32 %1, 0xffffffff :: v_dual_mov_b32 %2, 0x0
+    v_subrev_co_u32 %2, vcc_lo, %2, %1
+    """)[0]
     self.assertEqual(out, 0xffff_ffff)
 
 if __name__ == "__main__":

extra/sqtt/attempt_sqtt_parse.py

@@ -534,8 +534,8 @@ def parse_sqtt_print_packets(data: bytes, filter=DEFAULT_FILTER, verbose=True) -
 
 def parse(fn:str):
   with Timing(f"unpickle {fn}: "): dat = pickle.load(open(fn, "rb"))
-  #if getenv("ROCM", 0):
-  #  with Timing(f"decode {fn}: "): ctx = decode(dat)
+  if getenv("ROCM", 0):
+    with Timing(f"decode {fn}: "): ctx = decode(dat)
   dat_sqtt = [x for x in dat if isinstance(x, ProfileSQTTEvent)]
   print(f"got {len(dat_sqtt)} SQTT events in {fn}")
   return dat_sqtt

extra/sqtt/roc.py

@@ -5,7 +5,29 @@ from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEven
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent
 from tinygrad.runtime.autogen import llvm, rocprof
 from tinygrad.runtime.support.elf import elf_loader
-from tinygrad.viz.serve import llvm_disasm
+
+# to pass NULL to callbacks
+llvm.LLVMCreateDisasmCPUFeatures.argtypes = tuple(llvm.LLVMCreateDisasmCPUFeatures.argtypes[:5]) + (ctypes.c_void_p, ctypes.c_void_p)
+def llvm_disasm(arch:str, lib:bytes) -> dict[int, tuple[str, int]]:
+  llvm.LLVMInitializeAMDGPUTargetInfo()
+  llvm.LLVMInitializeAMDGPUTargetMC()
+  llvm.LLVMInitializeAMDGPUAsmParser()
+  llvm.LLVMInitializeAMDGPUDisassembler()
+  ctx = llvm.LLVMCreateDisasmCPUFeatures("amdgcn-amd-amdhsa".encode(), arch.encode(), "".encode(), None, 0, None, None)
+
+  image, sections, relocs = elf_loader(lib)
+  text = next((sh.header for sh in sections if sh.name == ".text"), None)
+  off, sz = unwrap(text).sh_addr, unwrap(text).sh_size
+
+  addr_table:dict[int, tuple[str, int]] = {}
+  out = ctypes.create_string_buffer(128)
+  cur_off = off
+  while cur_off < sz + off:
+    view = (ctypes.c_ubyte * ((sz + off) - cur_off)).from_buffer_copy(memoryview(image)[cur_off:])
+    instr_sz = llvm.LLVMDisasmInstruction(ctx, view, ctypes.c_uint64(len(view)), ctypes.c_uint64(0), out, ctypes.c_size_t(128))
+    addr_table[cur_off] = (out.value.decode("utf-8", "replace").strip(), instr_sz)
+    cur_off += instr_sz
+  return addr_table
 
 @dataclasses.dataclass(frozen=True)
 class InstExec:
@@ -48,11 +70,17 @@ class OccEvent(WaveSlot):
 RunKey = tuple[str, int]
 
 class _ROCParseCtx:
-  def __init__(self, sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, tuple[str, int]]]):
-    self.sqtt_evs, self.disasms = iter(sqtt_evs), disasms
+  def __init__(self, dev_evs:dict[str, ProfileDeviceEvent], sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
+    self.dev_evs, self.sqtt_evs, self.prog_evs = dev_evs, iter(sqtt_evs), prog_evs
+    self.disasms:dict[str, dict[int, tuple[str, int]]] = {}
     self.inst_execs:dict[RunKey, list[WaveExec]] = {}
     self.occ_events:dict[RunKey, list[OccEvent]] = {}
 
+    for prog in prog_evs:
+      arch = "gfx%d%x%x" % ((trgt:=unwrap(dev_evs[prog.device].props)['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
+      base = unwrap(prog.base)
+      self.disasms[prog.name] = asm = {base+addr:info for addr,info in llvm_disasm(arch, unwrap(prog.lib)).items()}
+
   def next_sqtt(self):
     x = next(self.sqtt_evs, None)
     self.active_run = (x.kern, x.exec_tag) if x is not None else None
@@ -75,8 +103,16 @@ class _ROCParseCtx:
     self.inst_execs.setdefault(unwrap(self.active_run), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.begin_time,
                                                                              ev.end_time, insts_blob))
 
-def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, tuple[str, int]]]) -> _ROCParseCtx:
-  ROCParseCtx = _ROCParseCtx(sqtt_evs, disasms)
+def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
+  dev_events:dict[str, ProfileDeviceEvent] = {}
+  sqtt_events:list[ProfileSQTTEvent] = []
+  prog_events:list[ProfileProgramEvent] = []
+  for e in profile:
+    if isinstance(e, ProfileDeviceEvent): dev_events[e.device] = e
+    if isinstance(e, ProfileSQTTEvent): sqtt_events.append(e)
+    if isinstance(e, ProfileProgramEvent) and e.device.startswith("AMD"): prog_events.append(e)
+
+  ROCParseCtx = _ROCParseCtx(dev_events, sqtt_events, prog_events)
 
   @rocprof.rocprof_trace_decoder_se_data_callback_t
   def copy_cb(buf, buf_size, _):
@@ -136,7 +172,7 @@ if __name__ == "__main__":
   args = parser.parse_args()
 
   with args.profile.open("rb") as f: profile = pickle.load(f)
-  #rctx = decode(profile, disasm)
-  #print('SQTT:', rctx.inst_execs.keys())
+  rctx = decode(profile)
+  print('SQTT:', rctx.inst_execs.keys())
 
   print_pmc([ev for ev in profile if isinstance(ev, ProfilePMCEvent)])

extra/test_pyrender.py

@@ -0,0 +1,20 @@
+from extra.optimization.helpers import load_worlds, ast_str_to_ast
+from tinygrad.helpers import tqdm
+from tinygrad.uop.ops import pyrender, UOp, Ops
+from tinygrad import dtypes
+from tinygrad.shape.shapetracker import ShapeTracker, View
+inf, nan = float('inf'), float('nan')
+
+if __name__ == "__main__":
+  ast_strs = load_worlds()
+  for i, ast_str in enumerate(tqdm(ast_strs)):
+    good_ast = ast_str_to_ast(ast_str)
+    code = '\n'.join(pyrender(good_ast))
+    print("\n***************\n\n"+code)
+    exec(code)
+    if str(good_ast) != str(ast):
+      print(code)
+      print("MISMATCH")
+      print(good_ast)
+      print(ast)
+      break

opencode.json

@@ -1 +1 @@
-{"$schema": "https://opencode.ai/config.json", "formatter": false, "lsp": false}
+{"$schema": "https://opencode.ai/config.json", "formatter": false}

test/external/external_benchmark_op_conv.py

@@ -1,9 +1,10 @@
 # ruff: noqa: E501 E712 F401
-from dataclasses import replace
 from tinygrad import dtypes, Device
 from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
+from tinygrad.codegen import full_rewrite
 from tinygrad.codegen.opt import Opt, OptOps # pylint: disable=unused-import
-from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.renderer import ProgramSpec
+from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import dedup, getenv
 from tinygrad.device import Buffer
 from tinygrad.dtype import ImageDType, Invalid
@@ -85,11 +86,16 @@ def dm_conv_172():
 
 ast = {143: vision_conv_143, 153: vision_conv_153, 172: dm_conv_172}[getenv("NUM", 143)]()
 
+compiler = Device.default.compiler
 renderer = Device.default.renderer
 allocator = Device.default.allocator
 
-ps = get_program(ast, renderer)
-cr = CompiledRunner(replace(ps, device=Device.DEFAULT))
+uops = full_rewrite(ast, renderer)
+src = renderer.render(uops)
+
+lib = compiler.compile(src)
+ps = ProgramSpec("conv", src, Device.DEFAULT, ast, uops)
+cr = CompiledRunner(ps, precompiled=lib)
 
 gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.DEFINE_GLOBAL]), key=lambda u: u.arg)
 # print(len(gs))

test/external/external_benchmark_resnet.py

@@ -71,8 +71,8 @@ class BenchmarkResnetTrain(unittest.TestCase):
 
       y = x.sequential(layer).contiguous().contiguous_backward()
       y.sum().backward()
-      if getenv("ASSIGN", 1): sched, _ = Tensor.schedule_with_vars(y, x.grad, *optim.schedule_step())
-      else: sched, _ = Tensor.schedule_with_vars(y, x.grad, *[t.grad for t in optim.params])
+      if getenv("ASSIGN", 1): sched, _, _ = Tensor.schedule_with_vars(y, x.grad, *optim.schedule_step())
+      else: sched, _, _ = Tensor.schedule_with_vars(y, x.grad, *[t.grad for t in optim.params])
 
       for _ in range(JITCNT):
         run_schedule(list(sched))

test/external/external_hip_compiler_bug.py

@@ -0,0 +1,225 @@
+# [<buf device:HIP size:1605632 dtype:dtypes.float>, <buf device:HIP size:301506 dtype:dtypes.float>, <buf device:HIP size:9408 dtype:dtypes.float>]
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer, CompiledRunner
+
+import ctypes
+import gpuctypes.hip as hip
+from tinygrad.helpers import to_char_p_p, init_c_var
+def get_bytes(arg, get_sz, get_str, check) -> bytes: return (sz := init_c_var(ctypes.c_size_t(), lambda x: check(get_sz(arg, ctypes.byref(x)))), ctypes.string_at(init_c_var(ctypes.create_string_buffer(sz.value), lambda x: check(get_str(arg, x))), size=sz.value))[1]  # noqa: E501
+def check(status):
+  if status != 0: raise RuntimeError(f"HIP Error {status}, {ctypes.string_at(hip.hipGetErrorString(status)).decode()}")
+def compile_hip(prg:str, arch="gfx1100") -> bytes:
+  check(hip.hiprtcCreateProgram(ctypes.byref(prog := hip.hiprtcProgram()), prg.encode(), "<null>".encode(), 0, None, None))
+  compile_options = [f'--offload-arch={arch}', '-I/opt/rocm/include']
+  status = hip.hiprtcCompileProgram(prog, len(compile_options), to_char_p_p([o.encode() for o in compile_options]))
+  if status != 0: raise RuntimeError(f"compile failed: {get_bytes(prog, hip.hiprtcGetProgramLogSize, hip.hiprtcGetProgramLog, check).decode()}")
+  return get_bytes(prog, hip.hiprtcGetCodeSize, hip.hiprtcGetCode, check)
+
+prefix = """
+typedef long unsigned int size_t;
+extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_local_id(unsigned int);
+extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_group_id(unsigned int);
+extern "C" __attribute__((device)) __attribute__((const)) size_t __ockl_get_local_size(unsigned int);
+typedef float float2 __attribute__((ext_vector_type(2)));
+static inline __attribute__((device)) float2 make_float2(float x, float y) { return {x, y}; }
+"""
+
+code = """
+extern "C" __attribute__((global))void r_2_8_7_7_4_8_3_7_7_4_4_2_2(float* data0, const float* data1, const float* data2) {
+  int gidx0 = __ockl_get_group_id(2); /* 2 */
+  int gidx1 = __ockl_get_group_id(1); /* 8 */
+  int gidx2 = __ockl_get_group_id(0); /* 49 */
+  int lidx4 = __ockl_get_local_id(1); /* 4 */
+  int lidx5 = __ockl_get_local_id(0); /* 8 */
+  float2 acc0 = make_float2(0.0f,0.0f);
+  float2 acc1 = make_float2(0.0f,0.0f);
+  float2 acc2 = make_float2(0.0f,0.0f);
+  float2 acc3 = make_float2(0.0f,0.0f);
+  float2 acc4 = make_float2(0.0f,0.0f);
+  float2 acc5 = make_float2(0.0f,0.0f);
+  float2 acc6 = make_float2(0.0f,0.0f);
+  float2 acc7 = make_float2(0.0f,0.0f);
+  float2 acc8 = make_float2(0.0f,0.0f);
+  float2 acc9 = make_float2(0.0f,0.0f);
+  float2 acc10 = make_float2(0.0f,0.0f);
+  float2 acc11 = make_float2(0.0f,0.0f);
+  float2 acc12 = make_float2(0.0f,0.0f);
+  float2 acc13 = make_float2(0.0f,0.0f);
+  float2 acc14 = make_float2(0.0f,0.0f);
+  float2 acc15 = make_float2(0.0f,0.0f);
+  float2 acc16 = make_float2(0.0f,0.0f);
+  float2 acc17 = make_float2(0.0f,0.0f);
+  float2 acc18 = make_float2(0.0f,0.0f);
+  float2 acc19 = make_float2(0.0f,0.0f);
+  float2 acc20 = make_float2(0.0f,0.0f);
+  float2 acc21 = make_float2(0.0f,0.0f);
+  float2 acc22 = make_float2(0.0f,0.0f);
+  float2 acc23 = make_float2(0.0f,0.0f);
+  float2 acc24 = make_float2(0.0f,0.0f);
+  float2 acc25 = make_float2(0.0f,0.0f);
+  float2 acc26 = make_float2(0.0f,0.0f);
+  float2 acc27 = make_float2(0.0f,0.0f);
+  float2 acc28 = make_float2(0.0f,0.0f);
+  float2 acc29 = make_float2(0.0f,0.0f);
+  float2 acc30 = make_float2(0.0f,0.0f);
+  float2 acc31 = make_float2(0.0f,0.0f);
+  int alu0 = (gidx2/7);
+  int alu1 = (gidx2%7);
+  int alu2 = (alu1*32);
+  int alu3 = (lidx5*4);
+  int alu4 = ((gidx0*802816)+(gidx1*100352)+(alu0*1792)+(alu1*16)+(lidx4*448)+(lidx5*2));
+  for (int ridx0 = 0; ridx0 < 3; ridx0++) {
+    for (int ridx1 = 0; ridx1 < 7; ridx1++) {
+      int alu5 = ((alu0*(-32))+(lidx4*(-8))+(ridx1*(-1)));
+      bool alu6 = (alu5<(-2));
+      bool alu7 = (alu5<0);
+      bool alu8 = (((alu0*32)+(lidx4*8)+ridx1)<221);
+      for (int ridx2 = 0; ridx2 < 7; ridx2++) {
+        int alu9 = ((gidx0*150528)+(ridx0*50176)+(alu0*7168)+(lidx4*1792)+(ridx1*224)+alu2+alu3+ridx2);
+        int alu10 = ((alu1*(-32))+(lidx5*(-4))+(ridx2*(-1)));
+        bool alu11 = (alu10<(-2));
+        float val0 = 0.0f;
+        if ((alu6*alu11)) { val0 = data1[alu9+(-675)]; }
+        float val1 = 0.0f;
+        if ((alu7*alu11)) { val1 = data1[alu9+(-227)]; }
+        float val2 = 0.0f;
+        if (alu11) { val2 = data1[alu9+221]; }
+        float val3 = 0.0f;
+        if ((alu8*alu11)) { val3 = data1[alu9+669]; }
+        bool alu12 = (alu10<0);
+        bool alu13 = ((alu2+alu3+ridx2)<225);
+        float val4 = 0.0f;
+        if ((alu6*alu12*alu13)) { val4 = data1[alu9+(-673)]; }
+        float val5 = 0.0f;
+        if ((alu7*alu12*alu13)) { val5 = data1[alu9+(-225)]; }
+        float val6 = 0.0f;
+        if ((alu12*alu13)) { val6 = data1[alu9+223]; }
+        float val7 = 0.0f;
+        if ((alu8*alu12*alu13)) { val7 = data1[alu9+671]; }
+        int alu14 = ((gidx1*1176)+(ridx0*49)+(ridx1*7)+ridx2);
+        float val8 = data2[alu14];
+        float val9 = data2[alu14+147];
+        float val10 = data2[alu14+294];
+        float val11 = data2[alu14+441];
+        float val12 = data2[alu14+588];
+        float val13 = data2[alu14+735];
+        float val14 = data2[alu14+882];
+        float val15 = data2[alu14+1029];
+        (acc0).x = ((val0*val8)+(acc0).x);
+        (acc1).x = ((val0*val9)+(acc1).x);
+        (acc2).x = ((val0*val10)+(acc2).x);
+        (acc3).x = ((val0*val11)+(acc3).x);
+        (acc4).x = ((val1*val8)+(acc4).x);
+        (acc5).x = ((val1*val9)+(acc5).x);
+        (acc6).x = ((val1*val10)+(acc6).x);
+        (acc7).x = ((val1*val11)+(acc7).x);
+        (acc8).x = ((val2*val8)+(acc8).x);
+        (acc9).x = ((val2*val9)+(acc9).x);
+        (acc10).x = ((val2*val10)+(acc10).x);
+        (acc11).x = ((val2*val11)+(acc11).x);
+        (acc12).x = ((val3*val8)+(acc12).x);
+        (acc13).x = ((val3*val9)+(acc13).x);
+        (acc14).x = ((val3*val10)+(acc14).x);
+        (acc15).x = ((val3*val11)+(acc15).x);
+        (acc16).x = ((val0*val12)+(acc16).x);
+        (acc17).x = ((val0*val13)+(acc17).x);
+        (acc18).x = ((val0*val14)+(acc18).x);
+        (acc19).x = ((val0*val15)+(acc19).x);
+        (acc20).x = ((val1*val12)+(acc20).x);
+        (acc21).x = ((val1*val13)+(acc21).x);
+        (acc22).x = ((val1*val14)+(acc22).x);
+        (acc23).x = ((val1*val15)+(acc23).x);
+        (acc24).x = ((val2*val12)+(acc24).x);
+        (acc25).x = ((val2*val13)+(acc25).x);
+        (acc26).x = ((val2*val14)+(acc26).x);
+        (acc27).x = ((val2*val15)+(acc27).x);
+        (acc28).x = ((val3*val12)+(acc28).x);
+        (acc29).x = ((val3*val13)+(acc29).x);
+        (acc30).x = ((val3*val14)+(acc30).x);
+        (acc31).x = ((val3*val15)+(acc31).x);
+        (acc0).y = ((val4*val8)+(acc0).y);
+        (acc1).y = ((val4*val9)+(acc1).y);
+        (acc2).y = ((val4*val10)+(acc2).y);
+        (acc3).y = ((val4*val11)+(acc3).y);
+        (acc4).y = ((val5*val8)+(acc4).y);
+        (acc5).y = ((val5*val9)+(acc5).y);
+        (acc6).y = ((val5*val10)+(acc6).y);
+        (acc7).y = ((val5*val11)+(acc7).y);
+        (acc8).y = ((val6*val8)+(acc8).y);
+        (acc9).y = ((val6*val9)+(acc9).y);
+        (acc10).y = ((val6*val10)+(acc10).y);
+        (acc11).y = ((val6*val11)+(acc11).y);
+        (acc12).y = ((val7*val8)+(acc12).y);
+        (acc13).y = ((val7*val9)+(acc13).y);
+        (acc14).y = ((val7*val10)+(acc14).y);
+        (acc15).y = ((val7*val11)+(acc15).y);
+        (acc16).y = ((val4*val12)+(acc16).y);
+        (acc17).y = ((val4*val13)+(acc17).y);
+        (acc18).y = ((val4*val14)+(acc18).y);
+        (acc19).y = ((val4*val15)+(acc19).y);
+        (acc20).y = ((val5*val12)+(acc20).y);
+        (acc21).y = ((val5*val13)+(acc21).y);
+        (acc22).y = ((val5*val14)+(acc22).y);
+        (acc23).y = ((val5*val15)+(acc23).y);
+        (acc24).y = ((val6*val12)+(acc24).y);
+        (acc25).y = ((val6*val13)+(acc25).y);
+        (acc26).y = ((val6*val14)+(acc26).y);
+        (acc27).y = ((val6*val15)+(acc27).y);
+        (acc28).y = ((val7*val12)+(acc28).y);
+        (acc29).y = ((val7*val13)+(acc29).y);
+        (acc30).y = ((val7*val14)+(acc30).y);
+        (acc31).y = ((val7*val15)+(acc31).y);
+      }
+    }
+  }
+  *((float2*)(data0+alu4)) = acc0;
+  *((float2*)(data0+alu4+12544)) = acc1;
+  *((float2*)(data0+alu4+25088)) = acc2;
+  *((float2*)(data0+alu4+37632)) = acc3;
+  *((float2*)(data0+alu4+112)) = acc4;
+  *((float2*)(data0+alu4+12656)) = acc5;
+  *((float2*)(data0+alu4+25200)) = acc6;
+  *((float2*)(data0+alu4+37744)) = acc7;
+  *((float2*)(data0+alu4+224)) = acc8;
+  *((float2*)(data0+alu4+12768)) = acc9;
+  *((float2*)(data0+alu4+25312)) = acc10;
+  *((float2*)(data0+alu4+37856)) = acc11;
+  *((float2*)(data0+alu4+336)) = acc12;
+  *((float2*)(data0+alu4+12880)) = acc13;
+  *((float2*)(data0+alu4+25424)) = acc14;
+  *((float2*)(data0+alu4+37968)) = acc15;
+  *((float2*)(data0+alu4+50176)) = acc16;
+  *((float2*)(data0+alu4+62720)) = acc17;
+  *((float2*)(data0+alu4+75264)) = acc18;
+  *((float2*)(data0+alu4+87808)) = acc19;
+  *((float2*)(data0+alu4+50288)) = acc20;
+  *((float2*)(data0+alu4+62832)) = acc21;
+  *((float2*)(data0+alu4+75376)) = acc22;
+  *((float2*)(data0+alu4+87920)) = acc23;
+  *((float2*)(data0+alu4+50400)) = acc24;
+  *((float2*)(data0+alu4+62944)) = acc25;
+  *((float2*)(data0+alu4+75488)) = acc26;
+  *((float2*)(data0+alu4+88032)) = acc27;
+  *((float2*)(data0+alu4+50512)) = acc28;
+  *((float2*)(data0+alu4+63056)) = acc29;
+  *((float2*)(data0+alu4+75600)) = acc30;
+  *((float2*)(data0+alu4+88144)) = acc31;
+}
+"""
+
+dev = "HIP"
+lib = Device[dev].compiler.compile(prefix+code)
+#lib = compile_hip(code)
+b0 = Buffer(dev, 1605632, dtypes.float)
+b1 = Buffer(dev, 301506, dtypes.float)
+b2 = Buffer(dev, 9408, dtypes.float)
+print(hex(b0._buf.value), hex(b0._buf.value+1605632*4))
+print(hex(b1._buf.value))
+print(hex(b2._buf.value))
+#prg = CompiledRunner("r_2_8_7_7_4_8_3_7_7_4_4_2_2", "", dev, [7, 1, 1], [8, 4, 1], precompiled=lib)
+prg = CompiledRunner("r_2_8_7_7_4_8_3_7_7_4_4_2_2", "", dev, [49, 8, 2], [8, 4, 1], precompiled=lib)
+print("compiled")
+prg([b0, b1, b2], {})
+print("ran")
+Device[dev].synchronize()
+print("sync")

test/external/external_test_am.py

@@ -2,7 +2,7 @@ import unittest
 from tinygrad.runtime.support.am.amdev import AMMemoryManager, AMPageTableEntry
 from tinygrad.runtime.support.am.ip import AM_GMC
 from tinygrad.runtime.support.hcq import MMIOInterface
-from tinygrad.runtime.support.memory import PageTableTraverseContext, AddrSpace
+from tinygrad.runtime.support.memory import PageTableTraverseContext
 from tinygrad.runtime.autogen.am import am
 from tinygrad.helpers import mv_address
 
@@ -70,7 +70,7 @@ class TestAMPageTable(unittest.TestCase):
 
     for va,sz in [(0x10000, 0x3000), (0x11000, 0x300000), (0x10000, 0x2000), (0x11000, 0x5000),
                   (0x2000000, 0x2000), (0x4000000, 0x4000000), (0x38000, 0x303000), (0x8000, 0x1000)]:
-      mm.map_range(vaddr=helper_va(va), size=sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+      mm.map_range(vaddr=helper_va(va), size=sz, paddrs=[(va, sz)])
 
       ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, helper_va(va))
       results = list(ctx.next(sz))
@@ -102,8 +102,8 @@ class TestAMPageTable(unittest.TestCase):
     mm0 = self.d[0].mm
 
     for (va1,sz1),(va2,sz2) in [((0x10000, (0x1000)), (0x11000, (2 << 20)))]:
-      mm0.map_range(vaddr=helper_va(va1), size=sz1, paddrs=[(va1, sz1)], aspace=AddrSpace.PHYS)
-      mm0.map_range(vaddr=helper_va(va2), size=sz2, paddrs=[(va2, sz2)], aspace=AddrSpace.PHYS)
+      mm0.map_range(vaddr=helper_va(va1), size=sz1, paddrs=[(va1, sz1)])
+      mm0.map_range(vaddr=helper_va(va2), size=sz2, paddrs=[(va2, sz2)])
       mm0.unmap_range(helper_va(va2), sz2)
       mm0.unmap_range(helper_va(va1), sz1)
 
@@ -112,24 +112,24 @@ class TestAMPageTable(unittest.TestCase):
 
     for va,sz in [(0x10000, 0x3000), (0x1000000, 0x1000000), (0x12000, 0x4000)]:
       exteranl_va = helper_va(va)
-      mm0.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+      mm0.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)])
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va, size=0x1000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+        mm0.map_range(vaddr=exteranl_va, size=0x1000, paddrs=[(va, sz)])
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va, size=0x100000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+        mm0.map_range(vaddr=exteranl_va, size=0x100000, paddrs=[(va, sz)])
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va + 0x1000, size=0x1000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+        mm0.map_range(vaddr=exteranl_va + 0x1000, size=0x1000, paddrs=[(va, sz)])
 
       with self.assertRaises(AssertionError):
-        mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+        mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(va, sz)])
 
       mm0.unmap_range(vaddr=exteranl_va, size=sz)
 
       # Finally can map and check paddrs
-      mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(0xdead0000, 0x1000), (0xdead1000, 0xff000)], aspace=AddrSpace.PHYS)
+      mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(0xdead0000, 0x1000), (0xdead1000, 0xff000)])
 
       ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, exteranl_va + 0x2000)
       for tup in ctx.next(0x100000):
@@ -147,13 +147,13 @@ class TestAMPageTable(unittest.TestCase):
     with self.assertRaises(AssertionError):
       mm0.unmap_range(helper_va(0x10000), 0x3000)
 
-    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)], aspace=AddrSpace.PHYS)
+    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)])
     mm0.unmap_range(helper_va(0x10000), 0x3000)
 
     with self.assertRaises(AssertionError):
       mm0.unmap_range(helper_va(0x10000), 0x3000)
 
-    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)], aspace=AddrSpace.PHYS)
+    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)])
     mm0.unmap_range(helper_va(0x10000), 0x3000)
 
     with self.assertRaises(AssertionError):
@@ -164,16 +164,16 @@ class TestAMPageTable(unittest.TestCase):
 
     # offset from start
     for off in [0, 0x3000, 0x10000]:
-      mm0.map_range(helper_va(0x1000000) + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000), aspace=AddrSpace.PHYS)
+      mm0.map_range(helper_va(0x1000000) + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
       mm0.unmap_range(helper_va(0x1000000) + off, (2 << 20) - off)
-      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)], aspace=AddrSpace.PHYS)
+      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
       mm0.unmap_range(helper_va(0x1000000), 2 << 20)
 
     # offset from end
     for off in [0x1000, 0x20000]:
-      mm0.map_range(helper_va(0x1000000), (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000), aspace=AddrSpace.PHYS)
+      mm0.map_range(helper_va(0x1000000), (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
       mm0.unmap_range(helper_va(0x1000000), (2 << 20) - off)
-      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)], aspace=AddrSpace.PHYS)
+      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
       mm0.unmap_range(helper_va(0x1000000), 2 << 20)
 
   def test_frag_size(self):

test/external/external_test_nv.py

@@ -0,0 +1,60 @@
+import unittest, struct, array, ctypes
+from tinygrad import Device, dtypes, Tensor
+from tinygrad.helpers import to_mv
+from tinygrad.runtime.ops_nv import NVDevice, HWQueue
+from tinygrad.codegen.opt.search import Opt, OptOps
+from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
+from test.external.fuzz_linearizer import get_fuzz_rawbufs
+
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.uop.ops import LazyOp, Ops, ReduceOps, BufferOps, MemBuffer
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View
+
+@unittest.skipUnless(Device.DEFAULT == "NV", "NV specific tests/fixes")
+class TestNV(unittest.TestCase):
+  @classmethod
+  def setUpClass(self):
+    TestNV.d0: NVDevice = Device["NV"]
+    TestNV.a = Tensor([0.,1.], device="NV").realize()
+    TestNV.b = self.a + 1
+    si = self.b.schedule()[-1]
+    TestNV.d0_runner = get_runner(TestNV.d0.device, si.ast)
+    TestNV.b.uop.buffer.allocate()
+    TestNV.addr = struct.pack("QQ", TestNV.b.uop.buffer._buf.va_addr, TestNV.a.uop.buffer._buf.va_addr)
+
+  def test_error_on_huge_dims(self):
+    ast = LazyOp(op=BufferOps.STORE, src=(LazyOp(op=ReduceOps.SUM, src=(LazyOp(op=Ops.CAST, src=(LazyOp(op=Ops.MUL, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=1, dtype=dtypes.half, st=ShapeTracker(views=(View(shape=(1, 1, 1024, 683), strides=(0, 0, 0, 1), offset=0, mask=None, contiguous=False),)))), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=2, dtype=dtypes.half, st=ShapeTracker(views=(View(shape=(1, 1, 1024, 683), strides=(0, 0, 683, 1), offset=0, mask=None, contiguous=True),))))), arg=None),), arg=dtypes.float),), arg=(3,)),), arg=MemBuffer(idx=0, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(1, 1, 1024, 1), strides=(0, 0, 1, 0), offset=0, mask=None, contiguous=True),)))) # noqa: E501
+    opts = [Opt(op=OptOps.GROUP, axis=0, arg=0), Opt(op=OptOps.PADTO, axis=1, arg=32), Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=2), Opt(op=OptOps.LOCAL, axis=0, arg=2)] # noqa: E501
+    with self.assertRaises(RuntimeError) as cm:
+      lin = Kernel(ast)
+      lin.apply_opts(opts)
+      rawbufs = get_fuzz_rawbufs(lin)
+      prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
+      prg(rawbufs, {}, wait=True)
+    self.assertEqual(str(cm.exception), "This is a runtime error message")
+
+  def test_buf4_usage(self):
+    TestNV.along = Tensor([105615], device="NV").realize()
+    ast = LazyOp(op=BufferOps.STORE, src=(LazyOp(op=Ops.SIN, src=(LazyOp(op=Ops.CAST, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=1, dtype=dtypes.ulong, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))),), arg=dtypes.float),), arg=None),), arg=MemBuffer(idx=0, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))) # noqa: E501
+    temp_runner = get_runner(TestNV.d0.device, (ast,))
+    temp_runner([TestNV.b.uop.buffer, TestNV.along.uop.buffer], var_vals={})
+    val = TestNV.b.uop.buffer.as_buffer().cast("f")[0]
+    assert abs(val - 0.80647) < 0.001, f"got val {val}"
+
+  def test_kernargs_no_oob_access(self):
+    kernargs_start = TestNV.d0._gpu_alloc((2 << 20), map_to_cpu=True).va_addr
+    kernargs = kernargs_start + ((2 << 20) - TestNV.d0_runner._prg.kernargs_alloc_size)
+    to_mv(kernargs, 0x160).cast('I')[:] = array.array('I', TestNV.d0_runner._prg.constbuffer_0)
+    ctypes.memmove(kernargs + TestNV.d0_runner._prg.kernargs_offset, TestNV.addr, len(TestNV.addr))
+
+    q = HWQueue()
+    q.exec(TestNV.d0_runner._prg, kernargs, TestNV.d0_runner.global_size, TestNV.d0_runner.local_size)
+    q.signal(TestNV.d0.timeline_signal, TestNV.d0.timeline_value).submit(TestNV.d0)
+    TestNV.d0._wait_signal(TestNV.d0.timeline_signal, TestNV.d0.timeline_value)
+    TestNV.d0.timeline_value += 1
+    val = TestNV.b.uop.buffer.as_buffer().cast("f")[0]
+    assert val == 1.0, f"got val {val}"
+
+if __name__ == "__main__":
+  unittest.main()

test/external/fuzz_shapetracker.py

@@ -0,0 +1,61 @@
+import random
+from tinygrad.helpers import DEBUG, getenv
+from test.unit.test_shapetracker import CheckingShapeTracker
+
+def do_permute(st):
+  perm = list(range(0, len(st.shape)))
+  random.shuffle(perm)
+  perm = tuple(perm)
+  if DEBUG >= 1: print("st.permute(", perm, ")")
+  st.permute(perm)
+
+def do_pad(st):
+  c = random.randint(0, len(st.shape)-1)
+  pad = tuple((random.randint(0,2), random.randint(0,2)) if i==c else (0,0) for i in range(len(st.shape)))
+  if DEBUG >= 1: print("st.pad(", pad, ")")
+  st.pad(pad)
+
+def do_reshape_split_one(st):
+  c = random.randint(0, len(st.shape)-1)
+  poss = [n for n in [1,2,3,4,5] if st.shape[c]%n == 0]
+  spl = random.choice(poss)
+  shp = st.shape[0:c] + (st.shape[c]//spl, spl) + st.shape[c+1:]
+  if DEBUG >= 1: print("st.reshape(", shp, ")")
+  st.reshape(shp)
+
+def do_reshape_combine_two(st):
+  if len(st.shape) < 2: return
+  c = random.randint(0, len(st.shape)-2)
+  shp = st.shape[:c] + (st.shape[c] * st.shape[c+1], ) + st.shape[c+2:]
+  if DEBUG >= 1: print("st.reshape(", shp, ")")
+  st.reshape(shp)
+
+def do_shrink(st):
+  c = random.randint(0, len(st.shape)-1)
+  while 1:
+    shrink = tuple((random.randint(0,s), random.randint(0,s)) if i == c else (0,s) for i,s in enumerate(st.shape))
+    if all(x<y for (x,y) in shrink): break
+  if DEBUG >= 1: print("st.shrink(", shrink, ")")
+  st.shrink(shrink)
+
+def do_flip(st):
+  flip = tuple(random.random() < 0.5 for _ in st.shape)
+  if DEBUG >= 1: print("st.flip(", flip, ")")
+  st.flip(flip)
+
+def do_expand(st):
+  c = [i for i,s in enumerate(st.shape) if s==1]
+  if len(c) == 0: return
+  c = random.choice(c)
+  expand = tuple(random.choice([2,3,4]) if i==c else s for i,s in enumerate(st.shape))
+  if DEBUG >= 1: print("st.expand(", expand, ")")
+  st.expand(expand)
+
+shapetracker_ops = [do_permute, do_pad, do_shrink, do_reshape_split_one, do_reshape_combine_two, do_flip, do_expand]
+
+if __name__ == "__main__":
+  random.seed(42)
+  for _ in range(getenv("CNT", 200)):
+    st = CheckingShapeTracker((random.randint(2, 10), random.randint(2, 10), random.randint(2, 10)))
+    for i in range(8): random.choice(shapetracker_ops)(st)
+    st.assert_same()

test/external/fuzz_shapetracker_math.py

@@ -0,0 +1,34 @@
+import random
+from tinygrad.helpers import getenv, DEBUG, colored, trange
+from tinygrad.shape.shapetracker import ShapeTracker
+from test.external.fuzz_shapetracker import shapetracker_ops
+from test.unit.test_shapetracker_math import st_equal, MultiShapeTracker
+
+def fuzz_plus() -> tuple[ShapeTracker, ShapeTracker]:
+  m = MultiShapeTracker([ShapeTracker.from_shape((random.randint(1, 10), random.randint(1, 10), random.randint(1, 10)))])
+  for _ in range(4): random.choice(shapetracker_ops)(m)
+  backup = m.sts[0]
+  m.sts.append(ShapeTracker.from_shape(m.sts[0].shape))
+  for _ in range(4): random.choice(shapetracker_ops)(m)
+  st_sum = backup + m.sts[1]
+  return m.sts[0], st_sum
+
+if __name__ == "__main__":
+  if seed:=getenv("SEED"): random.seed(seed)
+  total = getenv("CNT", 1000)
+  for fuzz in [globals()[f'fuzz_{x}'] for x in getenv("FUZZ", "plus").split(",")]:
+    same_but_neq = 0
+    for _ in trange(total, desc=f"{fuzz}"):
+      st1, st2 = fuzz()
+      eq = st_equal(st1, st2)
+      if getenv("CHECK_NEQ") and eq and st1.simplify() != st2.simplify():
+        print(colored("same but unequal", "yellow"))
+        print(st1.simplify())
+        print(st2.simplify())
+        same_but_neq += 1
+      if DEBUG >= 1:
+        print(f"EXP: {st1}")
+        print(f"GOT: {st2}")
+        print(colored("****", "green" if eq else "red"))
+      if not eq: exit(0)
+    if getenv("CHECK_NEQ"): print(f"same but unequal {same_but_neq}/{total} = {(same_but_neq/total)*100:.2f}%")

test/external/mlperf_bert/external_benchmark_bert.py

@@ -49,8 +49,8 @@ class BenchmarkBertTrain(unittest.TestCase):
 
       y = layer(*inputs).contiguous().contiguous_backward()
       y.sum().backward()
-      if getenv("ASSIGN", 1): sched, _ = Tensor.schedule_with_vars(y, *list(inputs), *optim.schedule_step())
-      else: sched, _ = Tensor.schedule_with_vars(y, *list(inputs), *[t.grad for t in optim.params])
+      if getenv("ASSIGN", 1): sched, _, _ = Tensor.schedule_with_vars(y, *list(inputs), *optim.schedule_step())
+      else: sched, _, _ = Tensor.schedule_with_vars(y, *list(inputs), *[t.grad for t in optim.params])
 
       for _ in range(JITCNT):
         run_schedule(sched)

test/helpers.py

@@ -2,27 +2,18 @@ import os, time, struct, functools, unittest
 from typing import Any, Callable
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.uop.ops import UOp, Ops
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.engine.realize import Runner, get_program
+from tinygrad.engine.realize import Runner
 from tinygrad.dtype import DType
 from tinygrad.nn.state import get_parameters
 from tinygrad.helpers import T, CI
-from tinygrad.renderer import Renderer
-from tinygrad.codegen import full_rewrite_to_sink, line_rewrite, pm_linearize_cleanups
-from tinygrad.codegen.late.linearizer import linearize
+from tinygrad.codegen import full_rewrite
 
 # decorator to skip slow tests by default, run with RUN_SLOW=1 to include them
 slow = unittest.skipUnless(os.getenv("RUN_SLOW"), "slow test, set RUN_SLOW=1 to run")
 from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler
 
-def get_uops(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
-  """Extract linearized UOps from a sink. Test helper that only does linearization (no render)."""
-  if ren is None: ren = Renderer()
-  if sink.arg is None: sink = sink.replace(arg=KernelInfo())
-  full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
-  return line_rewrite(linearize(full_sink), pm_linearize_cleanups)
-
 def derandomize_model(model):
   for p in get_parameters(model):
     p.replace(Tensor.empty(p.shape, device=p.device, dtype=p.dtype))
@@ -60,12 +51,13 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
   bufs = []
   for buf_dt, data in inputs or []:
     bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
-    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
+    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + buf_dt.fmt, *data)))
   g = UOp(Ops.DEFINE_GLOBAL, uop.dtype.ptr(), arg=0, src=())
-  prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), PythonRenderer())
-  prog = PythonProgram("run", PythonCompiler().compile(prg.src))
+  opts = PythonRenderer()
+  lst = full_rewrite(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), opts)
+  prog = PythonProgram("run", PythonCompiler().compile(opts.render(lst)))
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs)
-  return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
+  return out_buf.cast(uop.dtype.fmt).tolist()[0]
 
 def not_support_multi_device():
   # CL and CUDA don't support multi device if in CI

test/mockgpu/amd/amdgpu.py

@@ -185,8 +185,6 @@ class PM4Executor(AMDQueue):
       if st <= prg_addr < st+sz: prg_sz = sz - (prg_addr - st)
 
     assert prg_sz > 0, "Invalid prg ptr (not found in mapped ranges)"
-    # Pass valid memory ranges to Python emulator for bounds checking
-    if hasattr(remu, 'valid_mem_ranges'): remu.valid_mem_ranges = self.gpu.mapped_ranges
     err = remu.run_asm(prg_addr, prg_sz, *gl, *lc, args_addr)
     if err != 0: raise RuntimeError("remu does not support the new instruction introduced in this kernel")
 

test/mockgpu/helpers.py

@@ -1,5 +1,4 @@
 import ctypes, ctypes.util
-from tinygrad.helpers import getenv
 
 def _try_dlopen_gpuocelot():
   GPUOCELOT_PATHS = [ctypes.util.find_library("gpuocelot")] if ctypes.util.find_library("gpuocelot") is not None else []
@@ -15,20 +14,7 @@ def _try_dlopen_gpuocelot():
   print("Could not find libgpuocelot.so")
   return None
 
-class PythonRemu:
-  """Python RDNA3 emulator wrapper that matches the libremu.so interface."""
-  valid_mem_ranges: set[tuple[int, int]] = set()
-
-  def run_asm(self, lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int) -> int:
-    from extra.assembly.rdna3.emu import run_asm, set_valid_mem_ranges
-    # Pad ranges to handle GPU loads that may read past small buffers (e.g. s_load_b128 on 12-byte buffer)
-    set_valid_mem_ranges({(start, size + 4096) for start, size in self.valid_mem_ranges})
-    return run_asm(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr)
-
 def _try_dlopen_remu():
-  # Use Python emulator only if PYTHON_REMU=1
-  if getenv("PYTHON_REMU"):
-    return PythonRemu()
   REMU_PATHS = ["extra/remu/target/release/libremu.so", "libremu.so", "/usr/local/lib/libremu.so",
                "extra/remu/target/release/libremu.dylib", "libremu.dylib", "/usr/local/lib/libremu.dylib", "/opt/homebrew/lib/libremu.dylib"]
   for path in REMU_PATHS:

test/test_multitensor.py

@@ -57,18 +57,12 @@ class TestMultiTensor(unittest.TestCase):
       assert lb.shape == (128,)
     (X + X).realize()
 
-  @unittest.expectedFailure # TODO: fix
   def test_shard_empty(self):
     GlobalCounters.reset()
     X = Tensor.empty(256).shard(devices_2, 0).realize()
     assert GlobalCounters.kernel_count == 0
     (X + X).realize()
 
-  def test_arange_shrink(self):
-    x = Tensor.arange(4)
-    self.assertEqual(x.shard(devices_2, 0).realize().shrink(((2, 4),)).tolist(), [2, 3])
-    self.assertEqual(x.shard(devices_2, 0).realize().shrink(((0, 2),)).tolist(), [0, 1])
-
   def test_shard_like(self):
     X = Tensor.ones(256).shard(devices_2, 0)
     Y = Tensor.zeros(256).shard_like(X)

test/test_renderer_failures.py

@@ -1,26 +1,30 @@
 import unittest
 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.helpers import prod
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.helpers import dedup, flatten, prod
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.runtime.ops_python import PythonRenderer
 from tinygrad.uop.ops import UOp, Ops, python_alu
+from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Tensor, _to_np_dtype
+from tinygrad.codegen import full_rewrite
 
-def _test_uop_result(inputs:list[Tensor], prg, local_size=None):
+def _test_uop_result(inputs:list[Tensor], stores:list[UOp], local_size=None):
   for x in inputs: x.realize()
-  uops = prg.uops
+  # NOTE: we only toposort the stores
+  uops: list[UOp] = []
+  def _recursive_add(uop:UOp) -> list[UOp]: return flatten([_recursive_add(x) for x in uop.src])+[uop]
+  uops = dedup(flatten(_recursive_add(st) for st in stores))
   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)
+  src = Device[Device.DEFAULT].renderer.render(uops)
+  ei = CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test",
+                                  src, Device.DEFAULT, uops[-1], uops=uops, local_size=local_size))
   ei.exec(outbufs+inbufs)
   return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]
 
@@ -33,8 +37,8 @@ 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,))
-  prg = get_program(sink, Device[Device.DEFAULT].renderer)
-  return _test_uop_result([Tensor([input_val])], prg)[0]
+  uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
+  return _test_uop_result([Tensor([input_val])], uops)[0]
 
 class TestRendererFailures(unittest.TestCase):
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
@@ -43,8 +47,8 @@ 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,))
-    prg = get_program(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0]
+    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
+    ret = _test_uop_result([], uops, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])
 
   @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
@@ -54,8 +58,8 @@ 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,))
-    prg = get_program(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], prg, local_size=[4, 2, 1])[0]
+    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
+    ret = _test_uop_result([], uops, local_size=[4, 2, 1])[0]
     np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1])
 
 @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage), "uops are for cstyle")
@@ -100,8 +104,8 @@ 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,))
-    prg = get_program(sink, Device[Device.DEFAULT].renderer)
-    ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0]
+    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
+    ret = _test_uop_result([], uops, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])
 
   @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")

test/test_tensor.py

@@ -10,7 +10,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 full_rewrite
 from tinygrad.dtype import DType
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
@@ -865,12 +865,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()
+    schedule, _buffer_map, _var_vals = 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
+        uops = full_rewrite(s.ast, renderer)
+        renderer.render(uops)
+        return uops
 
   def _assert(self, dtype: DType, a: Tensor):
     uops = self._schedule_render(a)

test/test_uops.py

@@ -7,27 +7,30 @@ from tinygrad.dtype import dtypes, DType, AddrSpace
 from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu, AxisType
 from tinygrad.uop.spec import shared_spec
+from tinygrad.renderer import ProgramSpec
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.engine.realize import CompiledRunner, get_program, get_runner
 from tinygrad.engine.schedule import ExecItem
+from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
 from tinygrad.device import is_dtype_supported
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import get_uops
-from dataclasses import replace
 
 def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
   sink = UOp.group(*u)
   for r in sink.ranges: sink = sink.end(r)
   # we strip the SINK here for legacy reasons
-  ret = get_uops(sink.sink(arg=KernelInfo(opts_to_apply=())), ren)
+  ret = full_rewrite(sink.sink(arg=KernelInfo(opts_to_apply=())), ren)
   assert ret[-1].op is Ops.SINK
   return ret[:-1]
 
 def _uops_to_prg(uops_list):
-  prg = get_program(UOp.sink(*uops_list), Device[Device.DEFAULT].renderer)
-  return CompiledRunner(replace(prg, device=Device.DEFAULT))
+  uops = full_rewrite(ast:=UOp.sink(*uops_list), ren=Device[Device.DEFAULT].renderer)
+  src = Device[Device.DEFAULT].renderer.render(uops)
+  has_local = Device[Device.DEFAULT].renderer.has_local
+  return CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test", src, Device.DEFAULT, ast, uops=uops,
+                                global_size=[1,1,1] if has_local else None, local_size=[1,1,1] if has_local else None))
 
 def uop(uops:list[UOp], uop:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
   uops.append(UOp(uop, dtype, tuple(src), arg))
@@ -478,7 +481,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).schedule_with_vars()
     self.assertEqual(len(var_vals), 1)
     self.assertEqual(list(var_vals)[0], a.expr)
 
@@ -544,12 +547,6 @@ class TestUopsObject(unittest.TestCase):
     self.assertEqual(a.device, Device.DEFAULT)
 
 class TestUOpRender(unittest.TestCase):
-  def test_render_vectorize_empty(self):
-    u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(0), src=())
-    self.assertEqual(u.render(simplify=False), "{}")
-  def test_render_vectorize_empty_simplified(self):
-    u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(0), src=())
-    self.assertEqual(u.render(), "{}")
   def test_render_vectorize_same(self):
     u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
     self.assertEqual(u.render(simplify=False), "{0, ...}")

test/test_uops_stats.py

@@ -4,6 +4,7 @@ from tinygrad.helpers import getenv, GlobalCounters, EMULATE
 from tinygrad.engine.realize import get_program
 from tinygrad.renderer import ProgramSpec
 from tinygrad.renderer import Estimates
+from tinygrad.codegen import full_rewrite
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
@@ -145,7 +146,7 @@ class TestUOpsStats(unittest.TestCase):
     u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
     u4 = UOp(Ops.MUL, dtypes.int, (u1,u2))
     u5 = UOp(Ops.ADD, dtypes.int, (u4,u3))
-    uops = list(u5.toposort())
+    uops = full_rewrite(u5.sink())
 
     globl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), tuple())
     o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
@@ -154,7 +155,7 @@ class TestUOpsStats(unittest.TestCase):
     u2 = globl.index(o2)
     u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
     u4 = UOp(Ops.MULACC, dtypes.int, (u1,u2,u3))
-    uops_fma = list(u4.toposort())
+    uops_fma = full_rewrite(u4.sink())
 
     self.assertEqual(flops_mem(uops), flops_mem(uops_fma))
 

test/testextra/test_cfg_viz.py

@@ -1,180 +0,0 @@
-# ruff: noqa: F405, F403
-# allow define from star imports
-
-import unittest
-import textwrap
-
-from tinygrad import Device, Tensor
-from tinygrad.uop.ops import UOp, Ops, track_rewrites
-from tinygrad.renderer import ProgramSpec
-from tinygrad.helpers import TracingKey, getenv
-from tinygrad.engine.realize import ExecItem, CompiledRunner
-
-from extra.assembly.rdna3.autogen import *
-
-# TODO: use the RDNA3 renderer when it's in master
-template = """.text
-.globl fn_name
-.p2align 8
-.type fn_name,@function
-fn_name:
-  INSTRUCTION
-
-.rodata
-.p2align 6
-.amdhsa_kernel fn_name
-  .amdhsa_kernarg_size 8
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
-  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
-  .amdhsa_wavefront_size32 1
-.end_amdhsa_kernel
-
-.amdgpu_metadata
----
-amdhsa.version:
-  - 1
-  - 0
-amdhsa.kernels:
-  - .name: fn_name
-    .symbol: fn_name.kd
-    .group_segment_fixed_size: 0
-    .private_segment_fixed_size: 0
-    .wavefront_size: 32
-    .sgpr_count: 8
-    .vgpr_count: 8
-    .max_flat_workgroup_size: 1024
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 8
-    .args:
-      - .address_space:  global
-        .name:           a
-        .offset:         0
-        .size:           8
-        .type_name:      'float*'
-        .value_kind:     global_buffer
-...
-.end_amdgpu_metadata
-"""
-
-@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, ret=ret))
-def run_asm(name:str, insts:list) -> ProgramSpec:
-  src = "\n".join([inst if isinstance(inst, str) else inst.disasm() for inst in insts])
-  prg = ProgramSpec(name, src:=template.replace("fn_name", name).replace("INSTRUCTION", textwrap.dedent(src)), Device.DEFAULT, UOp(Ops.SINK),
-                    lib=Device[Device.DEFAULT].compiler.compile(src), global_size=[1, 1, 1], local_size=[1, 1, 1], globals=[0])
-  ei = ExecItem(UOp(Ops.SINK), [Tensor.empty(1).uop.buffer.ensure_allocated()], prg=CompiledRunner(prg))
-  ei.run()
-  return prg
-
-@unittest.skipUnless(Device.DEFAULT == "AMD" and not getenv("AMD_LLVM"), "only on AMD with comgr")
-class TestCfg(unittest.TestCase):
-  def setUp(self):
-    arch = Device["AMD"].arch
-    if not any(arch.startswith(a) for a in {"gfx11", "gfx12"}):
-      self.skipTest(f"tests written for RDNA, got arch {arch}")
-
-  def test_simple(self):
-    run_asm("simple", [
-      "entry:",
-        s_branch("bb1"),
-      "bb1:",
-        s_endpgm(),
-    ])
-
-  def test_diamond(self):
-    run_asm("diamond", [
-      "entry:",
-        s_cmp_eq_i32(s[0], 0),
-        s_cbranch_scc1("if"),
-        s_branch("else"),
-      "if:",
-        s_nop(1),
-        s_branch("end"),
-      "else:",
-        s_nop(0),
-      "end:",
-        s_endpgm(),
-    ])
-
-  def test_loop(self):
-    run_asm("simple_loop", [
-      "entry:",
-        s_mov_b32(s[1], 4),
-      "loop:",
-        s_add_u32(s[1], s[1], -1),
-        s_cmp_eq_i32(s[1], 0),
-        s_cbranch_scc0("loop"),
-        s_endpgm(),
-    ])
-
-  def test_loop_branch(self):
-    run_asm("loop_if", [
-      "entry:",
-        s_mov_b32(s[1], 4),
-      "loop:",
-        s_add_u32(s[1], s[1], -1),
-        s_cmp_eq_i32(s[1], 2),
-        s_cbranch_scc1("cond"),
-        s_branch("cont"),
-      "cond:",
-        s_add_u32(s[1], s[1], -2),
-      "cont:",
-        s_cmp_eq_i32(s[1], 0),
-        s_cbranch_scc0("loop"),
-        s_endpgm(),
-    ])
-
-  def test_loop_break(self):
-    run_asm("loop_break", [
-      "entry:",
-        s_mov_b32(s[1], 8),
-      "loop:",
-        s_add_u32(s[1], s[1], -1),
-        s_cmp_eq_i32(s[1], 5),
-        s_cbranch_scc1("break"),
-        s_cmp_eq_i32(s[1], 0),
-        s_cbranch_scc0("loop"),
-      "break:",
-        s_endpgm(),
-    ])
-
-  def test_switch(self):
-    run_asm("switch_case", [
-      "entry:",
-        s_cmp_eq_i32(s[0], 0),
-        s_cbranch_scc1("case0"),
-        s_cmp_eq_i32(s[0], 1),
-        s_cbranch_scc1("case1"),
-        s_branch("case2"),
-      "case0:",
-        s_nop(0),
-        s_branch("join"),
-      "case1:",
-        s_nop(1),
-        s_branch("join"),
-      "case2:",
-        s_nop(2),
-        s_branch("join"),
-      "join:",
-        s_endpgm(),
-    ])
-
-  def test_ping_pong(self):
-    run_asm("ping_pong", [
-      "entry:",
-        s_cmp_eq_i32(s[0], 0),
-        s_cbranch_scc1("ping"),
-        s_branch("pong"),
-      "ping:",
-        s_cmp_eq_i32(s[1], 0),
-        s_cbranch_scc1("pong"),
-        s_branch("end"),
-      "pong:",
-        s_cmp_eq_i32(s[2], 0),
-        s_cbranch_scc1("ping"),
-      "end:",
-        s_endpgm(),
-    ])
-
-if __name__ == "__main__":
-  unittest.main()

test/unit/test_allreduce.py

@@ -9,7 +9,7 @@ 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]
+      schedules, _, _ = t.sum(0).schedule_with_vars()
       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]
       # N*(N-1) scatter reduce, and N*(N-1) allgather

test/unit/test_dtype.py

@@ -66,21 +66,5 @@ class TestDtypeTolist(unittest.TestCase):
     # 57344
     self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e5m2).tolist(), [-28672.0, 1.5, 3.0, 28672.0])
 
-class TestCanLosslessCast(unittest.TestCase):
-  def test_can_lossless_cast(self):
-    from tinygrad.dtype import can_lossless_cast
-    # signed -> unsigned is NOT lossless (negative values wrap)
-    self.assertFalse(can_lossless_cast(dtypes.int8, dtypes.uint64))
-    self.assertFalse(can_lossless_cast(dtypes.int32, dtypes.uint32))
-    # unsigned -> larger signed is lossless
-    self.assertTrue(can_lossless_cast(dtypes.uint8, dtypes.int16))
-    self.assertTrue(can_lossless_cast(dtypes.uint32, dtypes.int64))
-    # large ints don't fit in floats
-    self.assertFalse(can_lossless_cast(dtypes.int32, dtypes.float))
-    self.assertFalse(can_lossless_cast(dtypes.int64, dtypes.double))
-    # half has more mantissa bits
-    self.assertTrue(can_lossless_cast(dtypes.int8, dtypes.half))
-    self.assertFalse(can_lossless_cast(dtypes.int8, dtypes.bfloat16))
-
 if __name__ == "__main__":
   unittest.main()

test/unit/test_dtype_spec.py

@@ -8,7 +8,7 @@ from hypothesis import given, settings, strategies as strat
 import numpy as np
 import torch
 
-settings.register_profile("my_profile", max_examples=50, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
+settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
 
 core_dtypes = list(DTYPES_DICT.values())

test/unit/test_elf.py

@@ -1,5 +1,5 @@
 import unittest, subprocess, platform
-from tinygrad.runtime.support.compiler_cpu import ClangJITCompiler
+from tinygrad.runtime.ops_cpu import ClangJITCompiler
 from tinygrad.runtime.support.elf import elf_loader
 
 class TestElfLoader(unittest.TestCase):

test/unit/test_helpers.py

@@ -1,6 +1,6 @@
-import ctypes, gzip, unittest, timeit, pickle
+import ctypes, gzip, unittest, timeit
 from tinygrad import Variable
-from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, get_contraction, count
+from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, get_contraction
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
 from tinygrad.tensor import Tensor, get_shape
 import numpy as np
@@ -120,18 +120,6 @@ class TestRoundUp(unittest.TestCase):
     self.assertEqual(round_up(232, 24984), 24984)
     self.assertEqual(round_up(24984, 232), 25056)
 
-class TestCount(unittest.TestCase):
-  def test_count_basic(self):
-    c = count(3)
-    self.assertEqual(next(c), 3)
-    self.assertEqual(next(c), 4)
-
-  def test_count_step_pickle(self):
-    c = count(1, 2)
-    self.assertEqual(next(c), 1)
-    c2 = pickle.loads(pickle.dumps(c))
-    self.assertEqual(next(c2), 3)
-
 @unittest.skip("no fetch tests because they need internet")
 class TestFetch(unittest.TestCase):
   def test_fetch_bad_http(self):

test/unit/test_llm_moe.py

@@ -1,53 +0,0 @@
-import unittest
-import numpy as np
-from tinygrad import Tensor
-
-class TestMoEFeedForward(unittest.TestCase):
-  def test_moe_feed_forward(self):
-    from tinygrad.apps.llm import TransformerBlock
-    dim, hidden, n_heads = 8, 16, 2
-    num_experts, k = 4, 2
-
-    block = TransformerBlock(dim, hidden, n_heads, n_heads, norm_eps=1e-5, head_dim=dim//n_heads,
-                             rope_theta=10000, max_context=16, num_experts=num_experts, num_experts_per_tok=k)
-
-    # set up weights: gate scales by (expert_id+1), up/down are identity-ish, router picks experts 0,2
-    block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
-    block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
-    block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
-    block.ffn_gate_inp.weight = Tensor([[1, 0, 1, 0]] * dim).T  # router strongly prefers experts 0 and 2
-    block.ffn_norm.weight = Tensor.ones(dim)  # identity norm
-
-    # input of ones -> after norm still ~ones -> experts 0,2 selected -> weighted sum of silu outputs
-    h = Tensor.ones(1, 1, dim)
-    out = block._feed_forward(h)
-
-    # expected: residual + moe_output ≈ 1 + avg(silu(1), silu(3))
-    expected = 1 + (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
-    np.testing.assert_allclose(out.numpy()[0, 0, 0], expected, rtol=1e-2)
-
-  def test_moe_feed_forward_batched(self):
-    from tinygrad.apps.llm import TransformerBlock
-    dim, hidden, n_heads = 8, 16, 2
-    num_experts, k = 4, 2
-
-    block = TransformerBlock(dim, hidden, n_heads, n_heads, norm_eps=1e-5, head_dim=dim//n_heads,
-                             rope_theta=10000, max_context=16, num_experts=num_experts, num_experts_per_tok=k)
-
-    # same setup as BS=1 test
-    block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
-    block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
-    block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
-    block.ffn_gate_inp.weight = Tensor([[1, 0, 1, 0]] * dim).T
-    block.ffn_norm.weight = Tensor.ones(dim)
-
-    # test with BS=2, T=3
-    h = Tensor.ones(2, 3, dim)
-    out = block._feed_forward(h)
-
-    # all outputs should match the BS=1 expected value
-    expected = 1 + (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
-    np.testing.assert_allclose(out.numpy(), expected, rtol=1e-2)
-
-if __name__ == '__main__':
-  unittest.main()

test/unit/test_llm_server.py

@@ -27,8 +27,8 @@ class TestLLMServer(unittest.TestCase):
     from tinygrad.apps.llm import Handler
     from tinygrad.helpers import TCPServerWithReuse
 
-    cls.server = TCPServerWithReuse(('127.0.0.1', 0), Handler)
-    cls.port = cls.server.server_address[1]
+    cls.port = 11435
+    cls.server = TCPServerWithReuse(('127.0.0.1', cls.port), Handler)
     cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
     cls.server_thread.start()
     time.sleep(0.1)

test/unit/test_schedule_cache.py

@@ -23,7 +23,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.schedule_with_vars()
     self.assertEqual(var_vals, {'pos': 42})
 
   def test_simple(self):

test/unit/test_tqdm.py

@@ -3,7 +3,7 @@ from unittest.mock import patch
 from io import StringIO
 from collections import namedtuple
 from tqdm import tqdm
-from tinygrad.helpers import tqdm as tinytqdm, trange
+from tinygrad.helpers import tqdm as tinytqdm, trange as tinytrange
 import numpy as np
 
 def _get_iter_per_second(raw:str) -> float:
@@ -165,7 +165,7 @@ class TestProgressBar(unittest.TestCase):
       mock_stderr.truncate(0)
 
       # compare bars at each iteration (only when tinytqdm bar has been updated)
-      for n in (bar := trange(total, desc="Test")):
+      for n in (bar := tinytrange(total, desc="Test")):
         if bar.i % bar.skip != 0: continue
         tiny_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
         iters_per_sec = float(tiny_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0

test/unit/test_uop_symbolic.py

@@ -3,8 +3,8 @@ import unittest, pickle, functools, math
 import z3
 
 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
+from tinygrad.codegen import full_rewrite
 from tinygrad.helpers import Context
-from test.helpers import get_uops
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid
 from tinygrad.uop.validate import uops_to_z3
@@ -747,7 +747,7 @@ class TestSymbolic(unittest.TestCase):
 
     # TODO: copied from render, render does not support cast
     glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
-    uops = get_uops(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
+    uops = full_rewrite(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
     rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
 
     self.assertEqual(rewritten_uop, cond.where(a.cast(dtypes.half), b.cast(dtypes.half)))

test/web/test_viz.js

@@ -0,0 +1,35 @@
+const { spawn } = require("child_process");
+const puppeteer = require("puppeteer");
+
+async function main() {
+  // ** start viz server
+  const proc = spawn("python", ["-u", "-c", "from tinygrad import Tensor; Tensor.arange(4).realize()"], { env: { ...process.env, VIZ:"1" },
+                      stdio: ["inherit", "pipe", "inherit"]});
+  await new Promise(resolve => proc.stdout.on("data", r => {
+    if (r.includes("ready")) resolve();
+  }));
+
+  // ** run browser tests
+  let browser, page;
+  try {
+    browser = await puppeteer.launch({ headless: true });
+    page = await browser.newPage();
+    const res = await page.goto("http://localhost:8000", { waitUntil:"domcontentloaded" });
+    if (res.status() !== 200) throw new Error("Failed to load page");
+    const scheduleSelector = await page.waitForSelector("ul:nth-of-type(2)");
+    scheduleSelector.click();
+    await page.waitForSelector("rect");
+    await page.waitForFunction(() => {
+      const nodes = document.querySelectorAll("#nodes > g").length;
+      const edges = document.querySelectorAll("#edges > path").length;
+      return nodes > 0 && edges > 0;
+    });
+  } finally {
+    // ** cleanups
+    if (page != null) await page.close();
+    if (browser != null) await browser.close();
+    proc.kill();
+  }
+}
+
+main();

tinygrad/apps/llm.py

@@ -5,14 +5,14 @@ from tinygrad.helpers import partition, TCPServerWithReuse, HTTPRequestHandler,
 
 class SimpleTokenizer:
   def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int], preset:str="llama3"):
-    if preset not in ("llama3","llama-v3","llama-bpe","qwen2","olmo"): raise ValueError(f"Invalid tokenizer preset '{preset}'")
+    if preset not in ("llama3","llama-v3","llama-bpe","qwen2"): raise ValueError(f"Invalid tokenizer preset '{preset}'")
     # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
     bs = [*range(33, 127), *range(161, 173), *range(174, 256)]  # bytes that map to themselves
     self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
 
     # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
-    # 0x323b0 is one past the max codepoint in unicode categories L/N/Z (0x323af is max L)
-    def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
+    # TODO: ucat_range is slow
+    def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(sys.maxunicode + 1) if unicodedata.category(chr(cp)).startswith(pre))
     r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
     self._split_to_word = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
       f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")
@@ -52,13 +52,9 @@ class SimpleTokenizer:
 
   def decode(self, ids:list[int]) -> str: return b''.join(self._tok2bytes[tid] for tid in ids).decode(errors='replace')
   def role(self, role:str):
-    if self.preset == 'olmo': return self.encode("<|" + role + "|>\n")  # OLMoE Instruct format
     if self.preset == 'qwen2': return self.encode("<|im_start|>" + role + "\n")
     return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")
-  def end_turn(self, eos_id:int):
-    if self.preset == 'olmo': return self.encode("\n")
-    if self.preset == 'qwen2': return [eos_id] + self.encode("\n")
-    return [eos_id]
+  def end_turn(self, eos_id:int): return [eos_id] + self.encode("\n") if self.preset == 'qwen2' else [eos_id]
 
 @functools.cache
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
@@ -66,14 +62,6 @@ def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
   freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
   return freqs.cos().cat(freqs.sin(), dim=-1).contiguous()
 
-class ExpertWeights:
-  """Like nn.Linear but with num_experts dimension. Weight shape: (num_experts, out_features, in_features)."""
-  def __init__(self, num_experts:int, in_features:int, out_features:int):
-    self.weight = Tensor.zeros(num_experts, out_features, in_features)
-  def __call__(self, sel:Tensor, x:Tensor) -> Tensor:
-    # sel: (B, T, k), x: (B, T, 1, in) or (B, T, k, in) -> output: (B, T, k, out)
-    return (x.unsqueeze(-2) @ self.weight[sel].transpose(-1, -2)).squeeze(-2)
-
 def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
   assert x.shape[-1] % 2 == 0
   cos, sin = freqs_cis.reshape(1, 1, x.shape[2], -1).chunk(2, dim=-1)
@@ -82,13 +70,12 @@ def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
 
 class TransformerBlock:
   def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, head_dim:int, rope_theta:float,
-               max_context:int=0, qk_norm:int=0, num_experts:int=0, num_experts_per_tok:int=0):
+               max_context:int=0, qk_norm:bool=False):
     self.n_heads      = n_heads
     self.n_kv_heads   = n_kv_heads
     self.head_dim     = head_dim
-    self.rope_theta   = rope_theta
     self.max_context  = max_context
-    self.qk_norm      = qk_norm
+    self.rope_theta   = rope_theta
 
     # --- attention projections (all linear, bias-free) ------------------
     q_proj_out       = self.head_dim * n_heads
@@ -101,30 +88,23 @@ class TransformerBlock:
     # --- RMSNorms --------------------------------------------------------
     self.attn_norm   = nn.RMSNorm(dim, norm_eps)
     self.ffn_norm    = nn.RMSNorm(dim, norm_eps)
-    if qk_norm: self.attn_q_norm, self.attn_k_norm = nn.RMSNorm(qk_norm, norm_eps), nn.RMSNorm(qk_norm, norm_eps)
+    if qk_norm: self.attn_q_norm, self.attn_k_norm = nn.RMSNorm(self.head_dim, norm_eps), nn.RMSNorm(self.head_dim, norm_eps)
 
-    # --- feed-forward (MoE or dense) -------------------------------------
-    if num_experts > 0:
-      self.num_experts_per_tok = num_experts_per_tok
-      self.ffn_gate_inp = nn.Linear(dim, num_experts, bias=False)  # router
-      self.ffn_gate_exps = ExpertWeights(num_experts, dim, hidden_dim)
-      self.ffn_up_exps = ExpertWeights(num_experts, dim, hidden_dim)
-      self.ffn_down_exps = ExpertWeights(num_experts, hidden_dim, dim)
-    else:
-      self.ffn_gate    = nn.Linear(dim, hidden_dim, bias=False)
-      self.ffn_up      = nn.Linear(dim, hidden_dim, bias=False)
-      self.ffn_down    = nn.Linear(hidden_dim, dim, bias=False)
+    # --- feed-forward ----------------------------------------------------
+    self.ffn_gate    = nn.Linear(dim, hidden_dim, bias=False)
+    self.ffn_up      = nn.Linear(dim, hidden_dim, bias=False)
+    self.ffn_down    = nn.Linear(hidden_dim, dim, bias=False)
 
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
     x_norm = self.attn_norm(x)                       # (B,T,D)
     q, k, v = self.attn_q(x_norm), self.attn_k(x_norm), self.attn_v(x_norm)
-    if self.qk_norm and self.qk_norm != self.head_dim: q, k = self.attn_q_norm(q), self.attn_k_norm(k)
 
     B, T, _ = x.shape
     q = q.reshape(B, T, self.n_heads,    self.head_dim).transpose(1, 2)  # (B,H,T,Hd)
     k = k.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
     v = v.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
-    if self.qk_norm == self.head_dim: q, k = self.attn_q_norm(q), self.attn_k_norm(k)
+
+    if hasattr(self, 'attn_q_norm'): q, k = self.attn_q_norm(q), self.attn_k_norm(k)
 
     # TODO: make UOp have SupportsIndex
     freqs_cis = precompute_freqs_cis(self.head_dim, self.max_context, self.rope_theta)[start_pos:start_pos+T]  # type: ignore
@@ -147,11 +127,6 @@ class TransformerBlock:
 
   def _feed_forward(self, h: Tensor) -> Tensor:
     h_norm = self.ffn_norm(h)
-    if hasattr(self, 'ffn_gate_exps'):
-      x = h_norm.unsqueeze(2)  # (B, T, 1, D) - add expert dim for broadcasting
-      probs, sel = self.ffn_gate_inp(h_norm).softmax(-1).topk(self.num_experts_per_tok)  # (B, T, k) each
-      x_down = self.ffn_down_exps(sel, self.ffn_gate_exps(sel, x).silu() * self.ffn_up_exps(sel, x))  # (B, T, k, D)
-      return h + (x_down * probs.unsqueeze(-1)).sum(axis=2)  # (B, T, D)
     # TODO: remove the need for this contiguous
     gated  = self.ffn_gate(h_norm).silu().contiguous() * self.ffn_up(h_norm)
     return h + self.ffn_down(gated)
@@ -161,9 +136,9 @@ class TransformerBlock:
 
 class Transformer:
   def __init__(self, *, num_blocks, dim, hidden_dim, n_heads, n_kv_heads, norm_eps, vocab_size, head_dim:int, rope_theta:float,
-               max_context:int=0, qk_norm:int=0, num_experts:int=0, num_experts_per_tok:int=0):
-    self.blk = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, head_dim, rope_theta, max_context, qk_norm,
-                                 num_experts, num_experts_per_tok) for _ in range(num_blocks)]
+               max_context:int=0, qk_norm:bool=False):
+    self.blk = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, head_dim, rope_theta, max_context, qk_norm)
+                for _ in range(num_blocks)]
     self.token_embd  = nn.Embedding(vocab_size, dim)
     self.output_norm = nn.RMSNorm(dim, norm_eps)
     self.output = nn.Linear(dim, vocab_size, bias=False)
@@ -195,20 +170,16 @@ class Transformer:
     max_context = min(max_context, kv[f'{arch}.context_length']) if max_context is not None else kv[f'{arch}.context_length']
     n_heads, n_kv_heads = kv[f'{arch}.attention.head_count'], kv[f'{arch}.attention.head_count_kv']
 
-    # Permute Q/K weights from interleaved to half-split RoPE layout (llama-style models only)
-    if arch == 'llama':
+    # permute Q/K weights from interleaved to half-split RoPE layout: [0,1,2,3,4,5...] -> [0,2,4,...,1,3,5,...]
+    if arch != 'qwen3':
       for name in state_dict:
         if 'attn_q.weight' in name: state_dict[name] = state_dict[name].rearrange("(n h two) d -> (n two h) d", n=n_heads, two=2)
         if 'attn_k.weight' in name: state_dict[name] = state_dict[name].rearrange("(n h two) d -> (n two h) d", n=n_kv_heads, two=2)
 
-    model = Transformer(num_blocks=kv[f'{arch}.block_count'], dim=kv[f'{arch}.embedding_length'],
-                        hidden_dim=kv.get(f'{arch}.expert_feed_forward_length', kv[f'{arch}.feed_forward_length']),
+    model = Transformer(num_blocks=kv[f'{arch}.block_count'], dim=kv[f'{arch}.embedding_length'], hidden_dim=kv[f'{arch}.feed_forward_length'],
                         n_heads=n_heads, n_kv_heads=n_kv_heads, norm_eps=kv[f'{arch}.attention.layer_norm_rms_epsilon'],
-                        vocab_size=len(kv['tokenizer.ggml.tokens']),
-                        head_dim=kv.get(f'{arch}.attention.key_length', kv[f'{arch}.embedding_length'] // n_heads),
-                        rope_theta=kv[f'{arch}.rope.freq_base'], max_context=max_context,
-                        qk_norm=int(state_dict['blk.0.attn_q_norm.weight'].shape[0]) if 'blk.0.attn_q_norm.weight' in state_dict else 0,
-                        num_experts=kv.get(f'{arch}.expert_count', 0), num_experts_per_tok=kv.get(f'{arch}.expert_used_count', 0))
+                        vocab_size=len(kv['tokenizer.ggml.tokens']), head_dim=kv[f'{arch}.attention.key_length'],
+                        rope_theta=kv[f'{arch}.rope.freq_base'], max_context=max_context, qk_norm='blk.0.attn_q_norm.weight' in state_dict)
     nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)  # NOTE: rope_freqs.weight (32,) is unused
     # NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster
     for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
@@ -236,8 +207,6 @@ models = {
   "qwen3:0.6b": "https://huggingface.co/Qwen/Qwen3-0.6B-GGUF/resolve/main/Qwen3-0.6B-Q8_0.gguf",
   "qwen3:1.7b": "https://huggingface.co/unsloth/Qwen3-1.7B-GGUF/resolve/main/Qwen3-1.7B-Q4_K_M.gguf",
   "qwen3:8b": "https://huggingface.co/Qwen/Qwen3-8B-GGUF/resolve/main/Qwen3-8B-Q4_K_M.gguf",
-  "qwen3:30b-a3b": "https://huggingface.co/Qwen/Qwen3-30B-A3B-GGUF/resolve/main/Qwen3-30B-A3B-Q4_K_M.gguf",
-  "olmoe": "https://huggingface.co/allenai/OLMoE-1B-7B-0924-Instruct-GGUF/resolve/main/olmoe-1b-7b-0924-instruct-q4_k_m.gguf",
 }
 
 # *** simple OpenAI compatible server on 11434 to match ollama ***

tinygrad/codegen/__init__.py

@@ -1,12 +1,11 @@
 from typing import cast
 import itertools
-from tinygrad.helpers import DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, getenv, TracingKey
-from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, pyrender
+from tinygrad.helpers import DEVECTORIZE, TRANSCENDENTAL, SPEC
+from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
-from tinygrad.renderer import Renderer, ProgramSpec
+from tinygrad.renderer import Renderer
 from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.helpers import panic
-from tinygrad.codegen.opt import Opt
 
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
@@ -29,8 +28,6 @@ pm_syntactic_sugar = PatternMatcher([
 def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -> UOp:
   if ren is None: ren = Renderer()
 
-  if getenv("VIZ"): graph_rewrite(sink, PatternMatcher([]), name="View Base AST")
-  if DEBUG >= 5: print(pyrender(sink))
   if SPEC: type_verify(sink, kernel_spec)
 
   # preprocess
@@ -126,49 +123,20 @@ def line_rewrite(lst:list[UOp], pm:PatternMatcher) -> list[UOp]:
     newlst.extend(ret[1])
   return newlst
 
-def do_linearize(prg:UOp, sink:UOp) -> UOp:
-  lst = line_rewrite(linearize(sink), pm_linearize_cleanups)
-  if SPEC: type_verify(lst, program_spec)
-  return prg.replace(src=prg.src + (UOp(Ops.LINEAR, src=tuple(lst)),))
-
-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),))
-
-def do_compile(ctx:Renderer, prg:UOp, source:UOp) -> UOp|None:
-  if ctx.compiler is None: return None
-  lib = ctx.compiler.compile_cached(source.arg)
-  return prg.replace(src=prg.src + (UOp(Ops.BINARY, arg=lib),))
-
-pm_to_program = PatternMatcher([
-  (UPat(Ops.PROGRAM, src=(UPat(Ops.SINK, name="sink"), UPat(Ops.DEVICE)), name="prg"), do_linearize),
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR, name="lin")), name="prg"), do_render),
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile),
-])
-
-@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name, ret.ast), ret=ret), replay=True)
-def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> ProgramSpec:
+def full_rewrite(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
   """
-  Transform an AST into a ProgramSpec. May trigger BEAM search.
+  Function to transform the Kernel UOp graph into a linearized program.
 
   Args:
-    ast: The Ops.SINK rooted AST
-    renderer: The renderer used to generate the code
+    sink: The Ops.SINK rooting the Kernel graph.
+    ren: The Renderer (can change how things are processed, fix this).
 
   Returns:
-    The ProgramSpec of the program.
+    Linear program in UOps.
   """
 
-  # fix up KernelInfo
-  if opts is not None:
-    assert ast.arg is None, "can't apply opts if sink has an arg"
-    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
-  if ast.arg is None: ast = ast.replace(arg=KernelInfo())
-
-  # rewrite to prg
-  full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
-  prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.device)))
-  prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
-
-  # create the ProgramSpec
-  return ProgramSpec.from_uop(prg)
+  full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
+  assert len(full_sink.ranges) == 0, f"all ranges must end by the sink, {full_sink.ranges}"
+  lst = line_rewrite(linearize(full_sink), pm_linearize_cleanups)
+  if SPEC: type_verify(lst, program_spec)
+  return lst

tinygrad/codegen/opt/postrange.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import math, itertools
 from collections import defaultdict
 from typing import cast, Final
@@ -6,7 +7,7 @@ from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes, ImageDType
 from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
-from tinygrad.helpers import ALLOW_TF32, count
+from tinygrad.helpers import ALLOW_TF32
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check
 from tinygrad.codegen.simplify import pm_flatten_range
 from tinygrad.renderer import Renderer
@@ -18,7 +19,7 @@ class Scheduler:
     self.ast, self.ren = ast, ren
     self.dont_use_locals = self.ast.arg.dont_use_locals if self.ast.arg is not None else False
     self.applied_opts = list(self.ast.arg.applied_opts) if self.ast.arg is not None else []
-    self.opt_range = count(start=max([x.arg[0] for x in self.rngs], default=0)+1)
+    self.opt_range = itertools.count(start=max([x.arg[0] for x in self.rngs], default=0)+1)
 
   @property
   def rngs(self):

tinygrad/codegen/opt/search.py

@@ -6,8 +6,7 @@ from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, di
 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.engine.realize import CompiledRunner, get_program
 from tinygrad.renderer import ProgramSpec
 from tinygrad.codegen.opt.postrange import Scheduler
 
@@ -38,10 +37,10 @@ def get_test_global_size(global_size, max_global_size, var_vals):
 def _time_program(p:ProgramSpec, 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") -> list[float]:
   factor = 1
-  if allow_test_size and max_global_size is not None:
+  if allow_test_size and p.global_size is not None 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))
+  try: car = CompiledRunner(p, precompiled=lib)
   except AssertionError: return [math.inf] * cnt
   tms = []
   input_bufs = [rawbufs[i] for i in car.p.globals]
@@ -72,7 +71,7 @@ def _try_compile(x:tuple[int,Scheduler], compiler:Compiler) -> tuple[int, tuple[
       if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many uops. {len(p.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 = compiler.compile(p.src)
     et = time.perf_counter() - st
     ret = (p, prog, et)
   except RuntimeError:

tinygrad/device.py

@@ -199,7 +199,8 @@ class Buffer:
     return mv
   def view(self, size:int, dtype:DType, offset:int) -> Buffer:
     assert offset < self.nbytes, "offset must be less than nbytes"
-    return Buffer(self.device, size, dtype, base=self.base, offset=self.offset+offset)
+    if self._base is not None: return Buffer(self.device, size, dtype, base=self._base, offset=self.offset+offset)
+    return Buffer(self.device, size, dtype, base=self, offset=offset)
 
 @dataclass(frozen=True)
 class DMACPURef:
@@ -277,7 +278,7 @@ class Compiler:
   def disassemble(self, lib:bytes): pass
 
 @dataclass(frozen=True)
-class CompilerPair: renderer:type[Renderer]|functools.partial; compiler:type[Compiler]|functools.partial|None; ctrl_var:ContextVar|None = None # noqa: E702
+class CompilerPair: renderer:type[Renderer]|functools.partial; compiler:type[Compiler]|functools.partial; ctrl_var:ContextVar|None = None # noqa: E702
 
 @dataclass(frozen=True)
 class CompilerSet: cset:list[CompilerPair]; ctrl_var:ContextVar|None = None # noqa: E702
@@ -289,25 +290,21 @@ class Compiled:
     self.device, self.allocator, self.runtime, self.graph, self.group_id = device, allocator, runtime, graph, group_id
 
     self.comps_ctrl_var = compilers.ctrl_var if compilers is not None else None
-    self.comp_sets:dict[Any, tuple[ContextVar|None, tuple[type[Renderer]|functools.partial, type[Compiler]|functools.partial|None]]] = {}
-    self.cached_pair:dict[Any, tuple[Renderer, Compiler|None]] = {}
+    self.comp_sets:dict[Any, tuple[ContextVar|None, tuple[type[Renderer]|functools.partial, type[Compiler]|functools.partial]]] = {}
+    self.cached_pair:dict[Any, tuple[Renderer, Compiler]] = {}
     for cpair in (compilers.cset if compilers is not None else [CompilerPair(Renderer, Compiler)]):
-      self.comp_sets[self._compiler_name(cpair.renderer, cpair.compiler)] = (cpair.ctrl_var, (cpair.renderer, cpair.compiler))
+      self.comp_sets[self._compiler_name(cpair.compiler)] = (cpair.ctrl_var, (cpair.renderer, cpair.compiler))
 
   @property
   def renderer(self) -> Renderer: return self._select_compiler_pair()[0]
 
   @property
-  def compiler(self) -> Compiler:
-    if (ret:=self.renderer.compiler or self._select_compiler_pair()[1]) is None: raise RuntimeError(f"no compiler for {self.device}")
-    return ret
+  def compiler(self) -> Compiler: return self._select_compiler_pair()[1]
 
-  def _compiler_name(self, r:type[Renderer]|functools.partial, c:type[Compiler]|functools.partial|None) -> str:
-    devname = self.device.split(':')[0].upper()
-    if c is None: return unwrap_class_type(r).__name__.upper().removesuffix("RENDERER").removeprefix(devname) or devname
-    return unwrap_class_type(c).__name__.upper().removesuffix("COMPILER").removeprefix(devname) or devname
+  def _compiler_name(self, c:type[Compiler]|functools.partial) -> str:
+    return unwrap_class_type(c).__name__.upper().removesuffix("COMPILER").removeprefix(devname:=self.device.split(':')[0].upper()) or devname
 
-  def _select_compiler_pair(self) -> tuple[Renderer, Compiler|None]:
+  def _select_compiler_pair(self) -> tuple[Renderer, Compiler]:
     # select forced compiler from global env var.
     forced_comps = set([self.comp_sets[val][1]] if self.comps_ctrl_var is not None and (val:=self.comps_ctrl_var.value) else [])
 
@@ -401,7 +398,7 @@ def enumerate_devices_str() -> Generator[str, None, None]:
             # d.renderer, d.compiler = r(), c()
             with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist()
             if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]")
-            set_text = f'({cc_ctrl_var.key}={d._compiler_name(r, c)} to make default)' if cc_ctrl_var is not None else ''
+            set_text = f'({cc_ctrl_var.key}={d._compiler_name(c)} to make default)' if cc_ctrl_var is not None else ''
             default_text = '(default)' if type(default_compiler) is type(d.compiler) else set_text
             compilers_results.append(f"{colored('+', 'green')} {unwrap_class_type(c).__name__} {default_text}")
             any_works = True

tinygrad/dtype.py

@@ -218,19 +218,17 @@ DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType)
 INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void", "index":"index"}
 
 @functools.cache
-def can_lossless_cast(dt0:DType, dt1:DType) -> bool:
+def can_safe_cast(dt0:DType, dt1:DType) -> bool:
   # return if dt1 preserves value of dt0
-  # similar to https://numpy.org/doc/stable/reference/generated/numpy.can_cast.html
+  # https://numpy.org/doc/stable/reference/generated/numpy.can_cast.html
   if dt0 == dt1 or dt0 == dtypes.bool: return True
   match dt1:
     case dtypes.index: return dt0 in dtypes.ints
     case dtypes.double: return dt0 in (dtypes.float, dtypes.half, dtypes.bfloat16, *dtypes.fp8s,
       dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8)
     case dtypes.float: return dt0 in (dtypes.half, dtypes.bfloat16, *dtypes.fp8s, dtypes.uint16, dtypes.uint8, dtypes.int16, dtypes.int8)
-    case dtypes.half: return dt0 in (*dtypes.fp8s, dtypes.uint8, dtypes.int8)
     case dtypes.uint64: return dt0 in (dtypes.uint32, dtypes.uint16, dtypes.uint8)
     case dtypes.uint32: return dt0 in (dtypes.uint16, dtypes.uint8)
-    case dtypes.uint16: return dt0 in (dtypes.uint8,)
     case dtypes.int64: return dt0 in (dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8)
     case dtypes.int32: return dt0 in (dtypes.uint16, dtypes.uint8, dtypes.int16, dtypes.int8)
     case dtypes.int16: return dt0 in (dtypes.uint8, dtypes.int8)
@@ -321,8 +319,11 @@ def fp8_to_float(x: int, dtype: DType) -> float:
 truncate: dict[DType, Callable] = {dtypes.bool: bool,
   dtypes.float16: float_to_fp16, dtypes.bfloat16: lambda x: float_to_bf16(float(x)),
   **{fp8: (lambda x, dtype=fp8: fp8_to_float(float_to_fp8(x, dtype), dtype)) for fp8 in dtypes.fp8s},
-  **{getattr(dtypes, n): (lambda x, c=getattr(ctypes, f'c_{n}'): c(x).value)
-     for n in ('float', 'double', 'int8', 'int16', 'int32', 'int64', 'uint8', 'uint16', 'uint32', 'uint64')}}
+  dtypes.float32: lambda x: ctypes.c_float(x).value, dtypes.float64: lambda x: ctypes.c_double(x).value,
+  dtypes.uint8: lambda x: ctypes.c_uint8(x).value, dtypes.uint16: lambda x: ctypes.c_uint16(x).value,
+  dtypes.uint32: lambda x: ctypes.c_uint32(x).value, dtypes.uint64: lambda x: ctypes.c_uint64(x).value,
+  dtypes.int8: lambda x: ctypes.c_int8(x).value, dtypes.int16: lambda x: ctypes.c_int16(x).value, dtypes.int32: lambda x: ctypes.c_int32(x).value,
+  dtypes.int64: lambda x: ctypes.c_int64(x).value}
 
 # numpy and torch dtype interop
 

tinygrad/engine/execution.py

@@ -0,0 +1,170 @@
+from __future__ import annotations
+from typing import Any, cast
+from tinygrad.helpers import DEBUG, GlobalCounters, all_same, colored, ansilen, PROFILE, ProfilePointEvent, cpu_events, time_to_str, TRACEMETA
+from tinygrad.uop.ops import UOp, Ops, sym_infer
+from tinygrad.device import Device, Buffer
+
+# **************** ExecutionUnit ****************
+
+class ExecutionUnit:
+  """
+  A bound, ready-to-execute unit. Replaces CapturedJit.
+
+  Takes ExecItems and binds them to real device resources on execution.
+  """
+
+  def __init__(self, items: list):
+    """
+    Create an ExecutionUnit from ExecItems.
+
+    Args:
+      items: ExecItems (with bufs as Buffers or UOps, lib optionally set)
+    """
+    from tinygrad.engine.realize import ExecItem
+
+    self.items: list[ExecItem] = items
+    self.buffer_map: dict[UOp, Buffer] = {}
+    self.var_vals: dict[str, int] = {}
+
+    # Create bound items with runners - lazy, done on first call
+    self._bound_items: list[tuple[Any, list[Buffer], tuple, dict[str, int]]]|None = None
+    self._graphs: list|None = None
+    self._first_run = True
+
+  def _bind(self):
+    """Create runners from lib and bind buffers."""
+    from tinygrad.engine.realize import CompiledRunner, BufferCopy, BufferXfer, ViewOp, EncDec, get_runner, get_program
+
+    self._bound_items = []
+    for item in self.items:
+      # Get buffers - prefer buf_uops with buffer_map, fall back to bufs for backwards compatibility
+      bufs: list[Buffer] = []
+      if item.buf_uops:
+        for uop in item.buf_uops:
+          bufs.append(cast(Buffer, self.buffer_map.get(uop) or uop.buffer))
+      else:
+        for buf in item.bufs:
+          if buf is not None: bufs.append(buf)
+
+      # Create runner from lib or use existing prg
+      if item.prg is not None:
+        runner = item.prg
+      elif item.ast.op is Ops.SINK:
+        device = bufs[0].device
+        if item.lib is not None:
+          # Create runner from cached lib
+          prg = get_program(item.ast, Device[device].renderer)
+          runner = CompiledRunner(prg, item.lib)
+        else:
+          # Compile and create runner
+          runner = get_runner(device, item.ast)
+      elif item.ast.op is Ops.BUFFER_VIEW:
+        runner = ViewOp(bufs[0])
+      elif item.ast.op is Ops.COPY:
+        if hasattr(Device[bufs[0].device].allocator, '_transfer') and all_same([x.device.split(":")[0] for x in bufs]):
+          runner = BufferXfer(bufs[0].nbytes, bufs[0].device, bufs[1].device)
+        else:
+          runner = BufferCopy(bufs[0].nbytes, bufs[0].device, bufs[1].device)
+      elif item.ast.op is Ops.ENCDEC:
+        runner = EncDec(item.ast, bufs[0].nbytes, bufs[1].device)
+      else:
+        raise RuntimeError(f"unknown op {item.ast.op}")
+
+      self._bound_items.append((runner, bufs, item.metadata, item.fixedvars))
+
+  def update(self, buffers: dict[UOp, Buffer]|None = None, var_vals: dict[str, int]|None = None) -> ExecutionUnit:
+    """Update buffer mapping and/or var_vals for next run. Returns self for chaining."""
+    if buffers is not None:
+      self.buffer_map.update(buffers)
+      # Need to rebind if we update buffers
+      self._bound_items = None
+    if var_vals is not None:
+      self.var_vals.update(var_vals)
+    return self
+
+  def __call__(self, var_vals: dict[str, int]|None = None, wait=False, do_update_stats=True, jit=False) -> float|None:
+    """Execute all items."""
+    from tinygrad.engine.realize import CompiledRunner
+
+    if var_vals is not None:
+      self.var_vals.update(var_vals)
+
+    # Lazy bind on first call
+    if self._bound_items is None:
+      self._bind()
+
+    assert self._bound_items is not None
+
+    # TODO: create graphs on first run
+    # if self._first_run:
+    #   self._create_graphs()
+    #   self._first_run = False
+
+    # Execute all items
+    total_et = 0.0
+    for runner, bufs, metadata, fixedvars in self._bound_items:
+      merged_var_vals = self.var_vals | fixedvars
+
+      # Ensure buffers are allocated (skip if jit - already allocated)
+      if not jit:
+        for b in bufs:
+          b.ensure_allocated()
+
+      # Reorder bufs to match program globals if needed
+      if isinstance(runner, CompiledRunner):
+        ordered_bufs = [bufs[i] for i in runner.p.globals]
+      else:
+        ordered_bufs = bufs
+
+      # PROFILE events
+      if PROFILE:
+        payload = {"metadata":metadata, "var_vals":merged_var_vals, "bufs":[b.trace_num for b in ordered_bufs], "name":runner.display_name}
+        payload["outputs"], payload["inputs"] = (runner.p.outs, runner.p.ins) if isinstance(runner, CompiledRunner) else ([0], [1])
+        cpu_events.append(ProfilePointEvent(runner.device, "exec", len(cpu_events), payload))
+
+      et = runner(ordered_bufs, merged_var_vals, wait=wait or DEBUG >= 2)
+      if et is not None:
+        total_et += et
+
+      # Update stats
+      if do_update_stats:
+        GlobalCounters.kernel_count += 1
+        op_est = sym_infer(runner.estimates.ops, merged_var_vals)
+        mem_est = sym_infer(runner.estimates.mem, merged_var_vals)
+        GlobalCounters.global_ops += op_est
+        GlobalCounters.global_mem += mem_est
+        if et is not None:
+          GlobalCounters.time_sum_s += et
+        if DEBUG >= 2:
+          lds_est = sym_infer(runner.estimates.lds, merged_var_vals)
+          mem_est = min(mem_est, lds_est)   # there can't be more memory accessed than loads/stores
+          header_color = 'magenta' if jit else ('green' if runner.first_run else None)
+          ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
+          flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
+          flops_str = f"{flops*1e-9:7.0f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:7.0f} TFLOPS", 'green')
+          mem_str = f"{membw*1e-9:4.0f}|{ldsbw*1e-9:<6.0f} GB/s" if membw < 1e13 and ldsbw < 1e15 else \
+            colored(f"{membw*1e-12:4.0f}|{ldsbw*1e-12:<6.0f} TB/s", 'green')
+          print(f"{colored(f'*** {runner.device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
+            f" {runner.display_name+' '*(46-ansilen(runner.display_name))} arg {len(ordered_bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
+            ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})")+
+            f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in metadata] if metadata else ''}")
+        runner.first_run = False
+
+    return total_et if wait else None
+
+  def __add__(self, other: ExecutionUnit) -> ExecutionUnit:
+    """Combine two ExecutionUnits, rebuild graph lazily."""
+    combined = ExecutionUnit(self.items + other.items)
+    combined.buffer_map = {**self.buffer_map, **other.buffer_map}
+    combined.var_vals = {**self.var_vals, **other.var_vals}
+    return combined
+
+  def free_intermediates(self):
+    """Deallocate internal buffers."""
+    for buf in self.buffer_map.values():
+      if buf.is_allocated():
+        buf.deallocate()
+    # Reset bound state
+    self._bound_items = None
+    self._graphs = None
+    self._first_run = True

tinygrad/engine/jit.py

@@ -97,7 +97,7 @@ class GraphRunner(Runner):
         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
+          assert ji.prg.p.global_size is not None and 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))
 
     # used in MultiGraphRunner. the ints are id() of _bufs
@@ -195,6 +195,8 @@ class CapturedJit(Generic[ReturnType]):
 
   # jit exec
   def __call__(self, input_buffers:list[Buffer], var_vals:dict[str, int]) -> ReturnType:
+    from tinygrad.engine.execution import ExecutionUnit
+
     # assign inputs
     for idx, offset, device, size, dtype in self.extra_view_inputs:
       input_buffers.append(Buffer(device, size, dtype, base=input_buffers[idx], offset=offset).ensure_allocated())
@@ -213,7 +215,8 @@ class CapturedJit(Generic[ReturnType]):
       self._first_run = False
 
     if DEBUG >= 1 and len(self._jit_cache) >= 10: print(f"jit execs {len(self._jit_cache)} kernels")
-    for ei in self._jit_cache: ei.run(var_vals, jit=True)
+    # Use ExecutionUnit for execution
+    ExecutionUnit(self._jit_cache).update(var_vals=var_vals)(jit=True)
     self._clear_inputs()
     return self.ret
 
@@ -251,7 +254,7 @@ class TinyJit(Generic[ReturnType]):
     return ret
 
   def add(self, ei:ExecItem):
-    self._jit_cache.append(ExecItem(ei.ast, [self.add_buffer(buf) for buf in ei.bufs if buf is not None], ei.metadata, ei.fixedvars, ei.prg))
+    self._jit_cache.append(ExecItem(ei.ast, [self.add_buffer(buf) for buf in ei.bufs if buf is not None], ei.metadata, ei.fixedvars, ei.lib, ei.prg))
 
   def reset(self):
     assert self.fxn is not None, "can't reset without function"
@@ -308,14 +311,16 @@ class TinyJit(Generic[ReturnType]):
 
       # prune independent kernels (optional)
       if self.prune:
+        from tinygrad.engine.execution import ExecutionUnit
         depends = set(input_buffers)
         update_depends(depends, jit_cache)
         pruned, onetime = partition(jit_cache, lambda ei: any(b in depends for b in get_out_buffers_for_ei(ei)))
         if DEBUG >= 1: print(f"pruned from {len(jit_cache)} -> {len(pruned)} kernels")
         # run the onetime kernels here
-        for ei in onetime:
-          for b in ei.bufs: cast(Buffer, b).ensure_allocated()
-          ei.run(var_vals, jit=True)
+        if onetime:
+          for ei in onetime:
+            for b in ei.bufs: cast(Buffer, b).ensure_allocated()
+          ExecutionUnit(onetime).update(var_vals=var_vals)(jit=True)
         jit_cache = pruned
 
       # memory planning (optional)

tinygrad/engine/memory.py

@@ -3,7 +3,7 @@ from collections import defaultdict
 from tinygrad.engine.realize import ExecItem
 from tinygrad.device import Device, Buffer
 from tinygrad.helpers import NO_MEMORY_PLANNER, dedup, DEBUG, round_up
-from tinygrad.uop.ops import Ops
+from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes, ImageDType
 from tinygrad.runtime.support.memory import TLSFAllocator
 
@@ -63,8 +63,16 @@ def _internal_memory_planner(buffers:list[list[Buffer]], noopt_buffers=None, ign
 
   return assigned
 
-def memory_planner(schedule:list[ExecItem]) -> list[ExecItem]:
+def memory_planner(schedule:list[ExecItem]) -> tuple[list[ExecItem], dict[UOp, Buffer]]:
   # Exclude buffers involved in load ops (e.g transfers) to preserve parallelism in graphs.
   assigned = _internal_memory_planner([[b for b in si.bufs if b is not None] for si in schedule],
                                       noopt_buffers={b for si in schedule if si.ast.op is not Ops.SINK for b in si.bufs if b is not None})
-  return [ExecItem(si.ast, [assigned.get(x, x) if x is not None else None for x in si.bufs], si.metadata, si.fixedvars) for si in schedule]
+  new_schedule = [ExecItem(si.ast, [assigned.get(x, x) if x is not None else None for x in si.bufs], si.metadata, si.fixedvars, buf_uops=si.buf_uops)
+                  for si in schedule]
+  # Build buffer_map from buf_uops -> assigned buffers
+  buffer_map: dict[UOp, Buffer] = {}
+  for si in new_schedule:
+    for i, uop in enumerate(si.buf_uops):
+      if uop not in buffer_map and i < len(si.bufs) and si.bufs[i] is not None:
+        buffer_map[uop] = si.bufs[i]  # type: ignore
+  return new_schedule, buffer_map

tinygrad/engine/realize.py

@@ -2,12 +2,52 @@ from typing import cast, Callable
 import time, pprint, random, itertools, math
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
-from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context
+from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context
 from tinygrad.helpers import unwrap
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, graph_rewrite, print_uops, track_rewrites, KernelInfo, pyrender
 from tinygrad.device import Device, Buffer
-from tinygrad.renderer import ProgramSpec, Estimates
-from tinygrad.codegen import get_program
+from tinygrad.renderer import Renderer, ProgramSpec, Estimates
+from tinygrad.codegen import full_rewrite
+from tinygrad.codegen.opt import Opt
+
+# **************** Program Creation ****************
+
+@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name, ret.ast), ret=ret), replay=True)
+def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> ProgramSpec:
+  """
+  Transform an AST into a ProgramSpec. May trigger BEAM search.
+
+  Args:
+    ast: The Ops.SINK rooted AST
+    renderer: The renderer used to generate the code
+
+  Returns:
+    The ProgramSpec of the program.
+  """
+
+  if getenv("VIZ"): graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
+  if DEBUG >= 5: print(pyrender(ast))
+
+  # linearize
+  if opts is not None:
+    assert ast.arg is None, "can't apply opts if sink has an arg"
+    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
+  try:
+    uops = full_rewrite(ast, renderer)
+  except RuntimeError as e:
+    print("***** LINEARIZE FAILURE *****")
+    print(e)
+    print(pyrender(ast))
+    raise
+  assert uops[-1].op is Ops.SINK, "last uop must be sink"
+
+  # print and render
+  if DEBUG >= 6: print_uops(uops)
+  src = renderer.render(uops)
+
+  return ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test", src, renderer.device, ast, uops,
+                     global_size=[1,1,1] if renderer.has_local or renderer.has_threads else None,
+                     local_size=[1,1,1] if renderer.has_local else None)
 
 # **************** Runners ****************
 
@@ -36,29 +76,36 @@ def optimize_local_size(_prg:Callable, global_size:list[int], rawbufs:list[Buffe
   return ret[1]
 
 class CompiledRunner(Runner):
-  def __init__(self, p:ProgramSpec, prg=None):
+  def __init__(self, p:ProgramSpec, precompiled:bytes|None=None, prg=None):
     if DEBUG >= 3: 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) if prg is None else prg
+    if precompiled is not None: self.lib = precompiled
+    else:
+      with cpu_profile(TracingKey(f"compile {p.name}", (p.function_name,)), "TINY"):
+        self.lib = Device[p.device].compiler.compile_cached(p.src)
+    if DEBUG >= 7: Device[p.device].compiler.disassemble(self.lib)
+    self._prg = Device[p.device].runtime(p.function_name, self.lib) if prg is None else prg
     super().__init__(p.name, p.device, p.estimates)
 
-  def __reduce__(self): return self.__class__, (self.p,)
+  def __reduce__(self): return self.__class__, (self.p, self.lib)
 
   def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None, wait=False) -> float|None:
     if var_vals is None: var_vals = {}
+    has_local = Device[self.p.device].renderer.has_local
     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): # type: ignore[arg-type]
+    if has_local and global_size is not None and local_size is None and all_int(self.p.global_size): # type: ignore[arg-type]
       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] for k in self.p.vars), wait=wait)
+    lra = {}
+    if global_size:
+      lra['global_size'] = tuple(global_size)
+      assert len(global_size) == 3, "global size must have len 3"
+    if local_size:
+      lra['local_size'] = tuple(local_size)
+      assert len(local_size) == 3, "local size must have len 3"
+    return self._prg(*[x._buf for x in rawbufs], **lra, vals=tuple(var_vals[k.expr] for k in self.p.vars), wait=wait)
 
 class ViewOp(Runner):
   def __init__(self, buf:Buffer): super().__init__(colored(f"view {buf.nbytes:8d} @ {buf.offset:<10d}", "yellow"), buf.device)
@@ -67,13 +114,13 @@ class ViewOp(Runner):
 
 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}"
+    if total_sz >= 1e6: name = f"{type(self).__name__[6:].lower()} {total_sz/1e6:7.2f}M, {dest_device[:7]:>7s} <- {src_device[:7]:7s}"
+    else: name = f"{type(self).__name__[6:].lower()} {total_sz:8d}, {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 hasattr(src.allocator.dev, 'io_uring') and \
       getattr(src.allocator.dev, 'fd', None) is not None and dest.allocator.supports_copy_from_disk
-    if disk_supports_fast_copyout and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096:
+    if src.device.startswith("DISK") and hasattr(dest.allocator, 'copy_from_disk') and disk_supports_fast_copyout and src.nbytes >= 4096:
       dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
     elif (src.device.startswith("DISK") or src.device.startswith("TINYFS")) and hasattr(dest.allocator, '_as_buffer'):
       # fast(ish) path, uses readinto in diskbuffers
@@ -115,7 +162,7 @@ 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(replace(bret.p, device=device), bret.lib)
   else:
     prg: ProgramSpec = get_program(ast, Device[device].renderer)
     method_cache[ckey] = method_cache[bkey] = ret = CompiledRunner(replace(prg, device=device))
@@ -136,15 +183,21 @@ si_lowerer = PatternMatcher([
 @dataclass
 class ExecItem:
   ast: UOp
-  bufs: list[Buffer|None] = field(default_factory=list)
+  bufs: list[Buffer|None] = field(default_factory=list)  # TODO: deprecate, use buf_uops + buffer_map
   metadata: tuple[Metadata, ...] = ()
   fixedvars: dict[str, int] = field(default_factory=dict)
+  lib: bytes|None = None  # compiled binary, None for COPY/VIEW/ENCDEC
   prg: Runner|None = None
+  buf_uops: tuple[UOp, ...] = ()  # buffer UOps, binding happens in ExecutionUnit
 
   def lower(self):
-    """Populate self.prg by lowering the AST."""
+    """Populate self.prg and self.lib by lowering the AST."""
     if self.prg is not None: return self
-    try: self.prg = cast(Runner, si_lowerer.rewrite(self.ast, self.bufs))
+    try:
+      self.prg = cast(Runner, si_lowerer.rewrite(self.ast, self.bufs))
+      # Store lib for SINK ops (compiled kernels)
+      if isinstance(self.prg, CompiledRunner):
+        self.lib = self.prg.lib
     except Exception as e:
       if DEBUG >= 2:
         print(f"error lowering {self.ast.op}")
@@ -190,24 +243,39 @@ class ExecItem:
 
 capturing: list = []  # put classes with an add method in here
 
-def run_schedule(schedule:list[ExecItem], var_vals:dict[str, int]|None=None, do_update_stats=True):
-  while len(schedule):
-    ei = schedule.pop(0).lower()
+def run_schedule(schedule:list[ExecItem], buffer_map:dict[UOp, Buffer]|None=None, var_vals:dict[str, int]|None=None, do_update_stats=True):
+  from tinygrad.engine.execution import ExecutionUnit
+  if buffer_map is None: buffer_map = {}
+
+  # Lower all items first
+  lowered: list[ExecItem] = []
+  for ei in schedule:
+    ei = ei.lower()
     if len(capturing) and CAPTURING: capturing[0].add(ei)
-    if VALIDATE_WITH_CPU and ei.ast.op is Ops.SINK:
-      # copy in allocated buffers from the GPU
-      bufs = [b for b in ei.bufs if b is not None]
-      nb: list[Buffer|None] = [Buffer("CPU", b.size, b.dtype) for b in bufs]
-      for cpu_b, gpu_b in zip(nb, bufs):
-        if cpu_b is not None and gpu_b.is_allocated(): cpu_b.ensure_allocated().copyin(gpu_b.as_buffer())
+    lowered.append(ei)
 
-      # run on GPU
-      ei.run(var_vals, do_update_stats=do_update_stats)
+  if VALIDATE_WITH_CPU:
+    # Run item by item with CPU validation
+    for ei in lowered:
+      if ei.ast.op is Ops.SINK:
+        # copy in allocated buffers from the GPU
+        bufs = [b for b in ei.bufs if b is not None]
+        nb: list[Buffer|None] = [Buffer("CPU", b.size, b.dtype) for b in bufs]
+        for cpu_b, gpu_b in zip(nb, bufs):
+          if cpu_b is not None and gpu_b.is_allocated(): cpu_b.ensure_allocated().copyin(gpu_b.as_buffer())
 
-      # validate the output buffers match (NOTE: this is assuming the output is buffer 0)
-      with Context(BEAM=0): ExecItem(ei.ast, nb, ei.metadata, ei.fixedvars).run(var_vals, do_update_stats=do_update_stats)
-      import numpy as np
-      assert nb[0] is not None
-      np.testing.assert_allclose(bufs[0].numpy(), nb[0].numpy(), rtol=1e-3, atol=1e-3)
-    else:
-      ei.run(var_vals, do_update_stats=do_update_stats)
+        # run on GPU
+        ExecutionUnit([ei]).update(buffers=buffer_map, var_vals=var_vals)(do_update_stats=do_update_stats)
+
+        # validate the output buffers match (NOTE: this is assuming the output is buffer 0)
+        with Context(BEAM=0):
+          ExecutionUnit([ExecItem(ei.ast, nb, ei.metadata, ei.fixedvars)]).update(var_vals=var_vals)(do_update_stats=do_update_stats)
+        import numpy as np
+        assert nb[0] is not None
+        np.testing.assert_allclose(bufs[0].numpy(), nb[0].numpy(), rtol=1e-3, atol=1e-3)
+      else:
+        ExecutionUnit([ei]).update(buffers=buffer_map, var_vals=var_vals)(do_update_stats=do_update_stats)
+  else:
+    # Use ExecutionUnit for batched execution
+    if lowered:
+      ExecutionUnit(lowered).update(buffers=buffer_map, var_vals=var_vals)(do_update_stats=do_update_stats)

tinygrad/engine/schedule.py

@@ -1,7 +1,9 @@
+from __future__ import annotations
 import time
 from typing import cast
 from collections import deque
-from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map
+from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites
+from tinygrad.uop.ops import PatternMatcher, UPat, graph_rewrite, graph_rewrite_map
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Buffer, MultiBuffer
 from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize
@@ -88,29 +90,23 @@ from tinygrad.engine.memory import memory_planner
 from tinygrad.schedule.rangeify import get_rangeify_map
 from tinygrad.schedule.multi import get_multi_map
 
-def replace_input_buffer(ctx:tuple[dict[UOp, UOp], dict[str, int]], b:UOp):
-  if (ret:=ctx[0].get(b, None)) is None:
+def replace_input_buffer(ctx:dict[UOp, UOp], b:UOp):
+  if (ret:=ctx.get(b, None)) is None:
     if b.op is Ops.BUFFER:
-      ctx[0][b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=len(ctx[0])), b.src[1]))
+      ctx[b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=len(ctx)), b.src[1]))
     else:
       # TODO: flip args in CONST
       assert b.op is Ops.CONST
-      ctx[0][b] = ret = b.replace(src=(b.src[0], UOp(Ops.LUNIQUE, arg=len(ctx[0]))))
+      ctx[b] = ret = b.replace(src=(b.src[0], UOp(Ops.LUNIQUE, arg=len(ctx))))
   return ret
 
-def strip_bind(ctx:tuple[dict[UOp, UOp], dict[str, int]], b:UOp):
-  var, val = b.src[0], b.src[1].arg
-  assert var.expr not in ctx[1] or ctx[1][var.expr] == val, f"bind mismatch on {var}, {ctx[1][var.expr]} != {val}"
-  ctx[1][var.expr] = val
-  return ctx[0].setdefault(b, b.replace(src=(b.src[0],)))
-
 pm_pre_sched_cache = PatternMatcher([
   # replace input buffers
   (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer),
   # remove unique consts
   (UPat(Ops.CONST, src=(UPat(Ops.DEVICE), UPat(Ops.UNIQUE)), name="b"), replace_input_buffer),
   # strip value from BIND for cache key normalization, so different values hit same cache
-  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), strip_bind),
+  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), lambda ctx,b: ctx.setdefault(b, b.replace(src=(b.src[0],)))),
 ])
 
 def replace_input_buffer_back(ctx:dict[UOp, UOp], b:UOp):
@@ -129,14 +125,13 @@ pm_post_sched_cache = PatternMatcher([
 
 schedule_cache: dict[bytes, tuple[list[ExecItem], UOp]] = {}
 @track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[1]))}")
-def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], list[ExecItem], dict[str, int]]:
+def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], list[ExecItem], dict[UOp, Buffer], dict[str, int]]:
   # big_sink srcs are all the Tensors
   st = time.perf_counter()
 
-  # replace all UNIQUE buffers with LUNIQUE, strip BIND values for cache key, extract var_vals
+  # replace all UNIQUE buffers with LUNIQUE, strip BIND values for cache key
   input_buffers: dict[UOp, UOp] = {}
-  var_vals: dict[str, int] = {}
-  big_sink_cache = graph_rewrite(big_sink, pm_pre_sched_cache, ctx=(input_buffers, var_vals), name="rewrite for sched cache")
+  big_sink_cache = graph_rewrite(big_sink, pm_pre_sched_cache, ctx=input_buffers, name="rewrite for sched cache")
   sched_cache_key = big_sink_cache.key
 
   if (sc_ret:=schedule_cache.get(sched_cache_key, None)) is None:
@@ -144,7 +139,7 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
     if SPEC: type_verify(big_sink, tensor_spec)
 
     # hack to preserve metadata
-    graph_rewrite_map(big_sink, pm_pre_sched_cache, ctx=({}, {}), name="preserve metadata")
+    graph_rewrite_map(big_sink, pm_pre_sched_cache, ctx={}, name="preserve metadata")
 
     # tensor map is what we return
     tensor_map: dict[UOp, UOp] = {}
@@ -190,14 +185,23 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
       assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
       dnums = [x for x in si.ast.variables() if x.arg[0] == '_device_num']
       for j, bufs in enumerate(zip(*[x.bufs for x in cast(tuple[MultiBuffer, ...], ubufs)])):
-        schedule.append(ExecItem(si.ast, list(bufs), si.metadata, si.fixedvars | ({dnums[0].expr:j} if len(dnums) else {})))
+        schedule.append(ExecItem(si.ast, list(bufs), si.metadata, si.fixedvars | ({dnums[0].expr:j} if len(dnums) else {}), buf_uops=buf_uops))
     else:
       # ONE -> ONE
-      schedule.append(ExecItem(si.ast, list(ubufs), si.metadata, si.fixedvars))
-  with cpu_profile(TracingKey("memory planner")): schedule = memory_planner(schedule)
+      schedule.append(ExecItem(si.ast, list(ubufs), si.metadata, si.fixedvars, buf_uops=buf_uops))
+  with cpu_profile(TracingKey("memory planner")): schedule, buffer_map = memory_planner(schedule)
+
+  # extract var_vals from BINDs that were stripped (only if there are kernels)
+  var_vals: dict[str, int] = {}
+  if schedule:
+    for u in input_buffers:
+      if u.op is Ops.BIND:
+        var, val = u.unbind()
+        assert var.expr not in var_vals or var_vals[var.expr] == val, f"bind mismatch on {var}, {var_vals[var.expr]} != {val}"
+        var_vals[var.expr] = val
 
   if (DEBUG >= 1 and len(schedule) > 1) or DEBUG >= 3:
     print(f"scheduled {len(schedule):4d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
           f" | {' cache hit' if sc_ret is not None else 'CACHE MISS'} {sched_cache_key.hex()[:8]}"+\
           f" | {len(UOpMetaClass.ucache)} uops in cache")
-  return tensor_map, schedule, var_vals if schedule else {}
+  return tensor_map, schedule, buffer_map, var_vals

tinygrad/gradient.py

@@ -59,7 +59,7 @@ def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp
   grads = {root: root_grad}
   for t0 in reversed(_deepwalk(root, targets)):
     if t0 not in grads: continue
-    lgrads: tuple[UOp|None, ...]|None = cast(tuple[UOp|None, ...]|None, pm_gradient.rewrite(t0, ctx=grads[t0]))
+    lgrads: tuple[UOp|None, ...]|None = cast(tuple[UOp, ...]|None, pm_gradient.rewrite(t0, ctx=grads[t0]))
     if lgrads is None: raise RuntimeError(f"failed to compute gradient for {t0.op}\n\nin {str(t0)[0:1000]}...")
     assert len(lgrads) == len(t0.src), f"got {len(lgrads)} gradient, expected {len(t0.src)}"
     for k,v in zip(t0.src, lgrads):

tinygrad/helpers.py

@@ -115,12 +115,12 @@ def suppress_finalizing(func):
       if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
   return wrapper
 
-def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]|None]], err_msg:str, cache:dict|None=None):
+def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]]], err_msg:str, cache:dict|None=None) -> tuple[T,...]|T:
   excs = []
   for typ in candidates:
     if cache is not None and typ in cache: return cache[typ]
     try:
-      x = tuple([cast(Callable, t)() if t is not None else None for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)()
+      x = tuple([cast(Callable, t)() for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)()
       if cache is not None: cache[typ] = x
       return x
     except Exception as e: excs.append(e)
@@ -523,7 +523,8 @@ class tqdm(Generic[T]):
   @classmethod
   def write(cls, s:str): print(f"\r\033[K{s}", flush=True, file=sys.stderr)
 
-def trange(n:int, **kwargs) -> tqdm[int]: return tqdm(range(n), total=n, **kwargs)
+class trange(tqdm):
+  def __init__(self, n:int, **kwargs): super().__init__(iterable=range(n), total=n, **kwargs)
 
 class disable_gc(contextlib.ContextDecorator):
   def __enter__(self):
@@ -542,11 +543,3 @@ copyreg.pickle(types.CodeType, _serialize_code)
 
 def _serialize_module(module:types.ModuleType): return importlib.import_module, (module.__name__,)
 copyreg.pickle(types.ModuleType, _serialize_module)
-
-class count:
-  def __init__(self, start:int=0, step:int=1):
-    self.n, self.step = start, step
-  def __next__(self) -> int:
-    cur = self.n
-    self.n += self.step
-    return cur

tinygrad/mixin/movement.py

@@ -1,4 +1,5 @@
 # mixins add syntactic sugar to Tensor and UOp
+import functools
 from typing import TypeAlias, TYPE_CHECKING, Self
 from tinygrad.uop import Ops
 from tinygrad.helpers import prod, argfix, flatten, dedup, make_tuple, ceildiv
@@ -281,38 +282,38 @@ class MovementMixin:
     ```
     """
 
-    def parse_side(s: str) -> tuple[list[str], list[tuple[int, int]]]:
-      """Parse one side of formula into (axis_names, dims) where dims are (start, end) index pairs for parens."""
-      tokens = f" {s} ".replace("…", "...").replace("(", " ( ").replace(")", " ) ").replace(" ", "  ").replace(" 1 ", " ( ) ").split()
-      lparens, rparens = [i for i, tok in enumerate(tokens) if tok == "("], [i for i, tok in enumerate(tokens) if tok == ")"]
+    def parse_formula(formula: str):
+      tokens = f" {formula} ".replace("…", "...").replace("(", " ( ").replace(")", " ) ").replace(" ", "  ").replace(" 1 ", " ( ) ").split()
+      lparens, rparens = map(lambda x: [i for i, ch in enumerate(tokens) if ch == x], ("(", ")"))
       pairs = list(zip(lparens, rparens))
       assert len(lparens) == len(rparens) and sorted(flatten(pairs)) == flatten(pairs), "bracket mismatch"
-      return [tok for tok in tokens if tok not in ("(", ")")], [(lp - 2*i, rp - 1 - 2*i) for i, (lp, rp) in enumerate(pairs)]
+      return [name for name in tokens if name not in ("(", ")")], [(s - 2 * i, e - 1 - 2 * i) for i, (s, e) in enumerate(pairs)]
 
     assert formula.count("->") == 1, 'need exactly one "->" in formula'
-    (lhs, unflatten_dims), (rhs, flatten_dims) = map(parse_side, formula.split("->"))
 
-    for name in sizes: assert name in lhs, f"axis {name} is not used in transform"
+    (lhs, unflatten_dims), (rhs, flatten_dims) = map(parse_formula, formula.split("->"))
+
+    for name in sizes:
+      assert name in lhs, f"axis {name} is not used in transform"
     assert sorted(lhs) == sorted(rhs) and len(lhs) == len(set(lhs)), f"name mismatch in {formula}"
-    for name in lhs+rhs: assert name == "..." or (name.isidentifier() and "_" not in (name[0], name[-1])), f"invalid axis name {name}"
+    for name in flatten((lhs, rhs)):
+      assert name == "..." or (name.isidentifier() and "_" not in (name[0], name[-1])), f"invalid axis name {name}"
     assert "..." not in flatten([lhs[s:e] for s, e in unflatten_dims]), f"cannot have collapsed ellipsis (...) in lhs of {formula}"
     assert lhs.count("...") <= 1, f"too many ellipses in {formula}"
 
     # resolve ellipsis
     if "..." in lhs:
       ell_len = len(self.shape) - len(lhs) + 1 + sum(e - s - 1 for s, e in unflatten_dims)
-      lhs, rhs = map(lambda l: l[:(i := l.index("..."))] + [f"...{j}" for j in range(ell_len)] + l[i + 1:] if "..." in l else l, (lhs, rhs))
-      def newdims(side, s, e): return (s + (ell_len - 1 if "...0" in side[:s] else 0), e + (ell_len - 1 if "...0" in side[:e] else 0))
-      unflatten_dims, flatten_dims = [newdims(lhs, s, e) for s, e in unflatten_dims], [newdims(rhs, s, e) for s, e in flatten_dims]
+    lhs, rhs = map(lambda l: l[: (i := l.index("..."))] + [f"...{j}" for j in range(ell_len)] + l[i + 1 :] if "..." in l else l, (lhs, rhs))
+    unflatten_dims = [(s + (ell_len - 1 if "...0" in lhs[:s] else 0), e + (ell_len - 1 if "...0" in lhs[:e] else 0)) for s, e in unflatten_dims]
+    flatten_dims = [(s + (ell_len - 1 if "...0" in rhs[:s] else 0), e + (ell_len - 1 if "...0" in rhs[:e] else 0)) for s, e in flatten_dims]
 
-    # unflatten -> permute -> flatten
-    t = self
-    for start, end in unflatten_dims: t = t.unflatten(start, tuple(sizes.get(lhs[i], -1) for i in range(start, end)))
+    # apply movement ops in order unflatten -> permute -> flatten/unsqueeze
+    t = functools.reduce(lambda x, dims: x.unflatten(dims[0], tuple(sizes.get(lhs[d], -1) for d in range(*dims))), unflatten_dims, self)
     for i, name in enumerate(lhs):
-      if name in sizes: assert sizes[name] == t.shape[i], f"size provided for dimension {name} incorrect"
+      assert (name not in sizes) or sizes[name] == t.shape[i], f"size provided for dimension {name} incorrect"
     t = t.permute([lhs.index(name) for name in rhs])
-    for start, end in reversed(flatten_dims): t = t.flatten(start, end - 1) if start < end else t.unsqueeze(start)
-    return t
+    return functools.reduce(lambda x, dims: x.flatten(dims[0], dims[1] - 1) if dims[0] < dims[1] else x.unsqueeze(dims[0]), reversed(flatten_dims), t)
 
   # *** movement ops with expand ***
 

tinygrad/renderer/__init__.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
-from typing import Callable, cast, TYPE_CHECKING
+from typing import Callable, cast
 import functools
 from dataclasses import dataclass, field
-from tinygrad.helpers import to_function_name, dedup, prod, DEBUG
-from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp, PatternMatcher, print_uops
+from tinygrad.helpers import to_function_name, dedup, prod
+from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp, PatternMatcher
 from tinygrad.dtype import AddrSpace, PtrDType
 from tinygrad.codegen.opt.tc import TensorCore
 from tinygrad.codegen.opt import Opt
-if TYPE_CHECKING: from tinygrad.device import Compiler
 
 @dataclass(frozen=True)
 class Estimates:
@@ -65,15 +64,36 @@ class ProgramSpec:
   device:str
   ast:UOp  # save the base ast (this is method cache key)
   uops:list[UOp]|None=None
-  lib:bytes|None=None
 
-  # filled in from uops (via from_uop)
-  global_size:list[int]=field(default_factory=lambda: [1,1,1])
+  # filled in from uops (if we have uops)
+  global_size:list[int]|None=None
   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)
+  _ran_post_init:bool=False  # NOTE: this is needed if you call replace on the Program
+
+  def __post_init__(self):
+    if not self._ran_post_init and self.uops is not None:
+      # single pass through the uops
+      for u in self.uops:
+        if u.op is Ops.DEFINE_VAR: self.vars.append(u)
+        if u.op is Ops.DEFINE_GLOBAL: self.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.DEFINE_GLOBAL: (self.outs if u.op is Ops.STORE else self.ins).append(buf.arg)
+          # TODO: can else happen?
+        if u.op is Ops.SPECIAL:
+          # NOTE: you have to set local_size and global_size to the base [1,1,1] outside this
+          if u.arg[0] == 'i': self.local_size = None
+          special_size = self.local_size if u.arg[0] == 'l' else self.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())
+      self.vars = sorted(self.vars, key=lambda v: v.arg)
+      self.outs = sorted(dedup(self.outs))
+      self.ins = sorted(dedup(self.ins))
+      self._ran_post_init = True
 
   @functools.cached_property
   def estimates(self) -> Estimates:
@@ -89,43 +109,10 @@ class ProgramSpec:
     return self.uops[-1].arg.applied_opts
 
   def launch_dims(self, var_vals:dict[str, int]):
-    global_size = [sym_infer(sz, var_vals) for sz in self.global_size]
+    global_size = [sym_infer(sz, var_vals) for sz in self.global_size] if self.global_size is not None else None
     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 uops:
-      if u.op is Ops.DEFINE_VAR: _vars.append(u)
-      if u.op is Ops.DEFINE_GLOBAL: _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.DEFINE_GLOBAL: (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())
-
-    return ProgramSpec(sink.arg.name, source.arg, device.arg, sink, uops, lib, global_size, local_size,
-                       sorted(_vars, key=lambda v: v.arg), sorted(dedup(_globals)), sorted(dedup(outs)), sorted(dedup(ins)))
-
 class Renderer:
   device: str = ""
   suffix: str = ""
@@ -142,7 +129,6 @@ class Renderer:
   pre_matcher: PatternMatcher|None = None
   extra_matcher: PatternMatcher|None = None
   code_for_op: dict[Ops, Callable] = {}
-  compiler: Compiler|None = None
 
   def __reduce__(self): return self.__class__, ()
   def render(self, uops:list[UOp]) -> str: raise NotImplementedError("needs a renderer")

tinygrad/renderer/cstyle.py

@@ -8,7 +8,6 @@ from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, trunc
 from tinygrad.renderer import Renderer
 from tinygrad.codegen.late.devectorizer import no_vectorized_alu
 
-
 base_rewrite = PatternMatcher([
   (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),
   (UPat(Ops.IF, name="x"), lambda ctx,x: f"if ({ctx[x.src[0]]}) {{"),
@@ -279,11 +278,6 @@ class ClangRenderer(CStyleLanguage):
     defines = '\n'.join(self._render_defines(uops))
     return defines + "\n" + self._render_body(function_name, kernel, bufs, uops, prefix) + "\n" + self._render_entry(function_name, bufs)
 
-class ClangJITRenderer(ClangRenderer):
-  def __init__(self):
-    from tinygrad.runtime.support.compiler_cpu import ClangJITCompiler
-    self.compiler = ClangJITCompiler()
-
 class OpenCLRenderer(CStyleLanguage):
   device = "CL"
 
@@ -382,8 +376,8 @@ class CUDARenderer(CStyleLanguage):
   shared_max = 49152
 
   def __init__(self, arch:str):
-    self.arch, arch_ver = arch, int(arch[3:])
-    self.tensor_cores = tc.cuda_sm89 if arch_ver >= 89 else tc.cuda_sm80 if arch_ver >= 80 else tc.cuda_sm75 if arch_ver >= 75 else []
+    self.arch = arch
+    self.tensor_cores = tc.cuda_sm89 if int(arch[3:]) >= 89 else tc.cuda_sm80 if int(arch[3:]) >= 80 else tc.cuda_sm75 if int(arch[3:]) >= 75 else []
   def __reduce__(self): return self.__class__, (self.arch,)
 
   # language options
@@ -446,7 +440,7 @@ class CUDARenderer(CStyleLanguage):
 
     return super().render_kernel(function_name, kernel, bufs, uops, prefix=prefix)
 
-class AMDHIPRenderer(CStyleLanguage):
+class AMDRenderer(CStyleLanguage):
   device = "AMD"
   shared_max = 65536
   # NOTE: this is only really needed on gfx12, even though gfx11 reports the same limitation
@@ -458,8 +452,7 @@ class AMDHIPRenderer(CStyleLanguage):
   @staticmethod
   def is_cdna(arch): return arch.split(":")[0] in {"gfx942", "gfx950"}
   def __init__(self, arch:str): # gfx942 => MI300, gfx1100 => RX 7900, gfx1201 => RX 9700
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    self.arch, self.compiler = arch, HIPCompiler(arch)
+    self.arch = arch
     self.tensor_cores = self.get_tensor_cores(arch)
     if self.is_cdna(self.arch):
       self.string_rewrite = PatternMatcher([
@@ -540,11 +533,5 @@ class AMDHIPRenderer(CStyleLanguage):
     return super().render_kernel(function_name, kernel, bufs, uops, prefix)
 
 class NVRenderer(CUDARenderer): device = "NV"
-class HIPRenderer(AMDHIPRenderer): device = "HIP"
-class AMDHIPCCRenderer(AMDHIPRenderer):
-  def __init__(self, arch:str):
-    from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
-    super().__init__(arch)
-    self.compiler = HIPCCCompiler(arch)
-
+class HIPRenderer(AMDRenderer): device = "HIP"
 class QCOMRenderer(OpenCLRenderer): device = "QCOM"

tinygrad/renderer/llvmir.py

@@ -2,7 +2,7 @@ from typing import cast
 import math, struct, sys
 from tinygrad.codegen.opt import tc
 from tinygrad.renderer import Renderer
-from tinygrad.renderer.cstyle import AMDHIPRenderer, create_non_native_float_pats, pm_manual_bf16_cast
+from tinygrad.renderer.cstyle import AMDRenderer, create_non_native_float_pats, pm_manual_bf16_cast
 from tinygrad.uop.decompositions import xexp2, xlog2
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, range_str
 from tinygrad.dtype import dtypes, float_to_fp8, DType, PtrDType, truncate
@@ -209,8 +209,8 @@ llvm_intrinsics = {Ops.SQRT: "sqrt", Ops.LOG2: "log2", Ops.EXP2: "exp2"}
 class AMDLLVMRenderer(LLVMRenderer):
   device = "AMD"
   has_local = True
-  shared_max = AMDHIPRenderer.shared_max
-  global_max = AMDHIPRenderer.global_max
+  shared_max = AMDRenderer.shared_max
+  global_max = AMDRenderer.global_max
   abi = "amdgpu_kernel"
   code_for_op = {**LLVMRenderer.code_for_op, **{op: lambda: None for op in llvm_intrinsics}}
   string_rewrite = PatternMatcher([
@@ -254,10 +254,9 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
                   f'"amdgpu-flat-work-group-size"="1,{requiredMaxThreadsPerBlock}"', '"no-trapping-math"="true"']
     return 'attributes #0 = { ' + ' '.join(attributes) + ' }'
   def __init__(self, arch:str):
-    from tinygrad.runtime.support.compiler_amd import AMDLLVMCompiler
-    self.arch, self.compiler = arch, AMDLLVMCompiler(arch)
-    self.tensor_cores = AMDHIPRenderer.get_tensor_cores(arch)
-    self.is_cdna = AMDHIPRenderer.is_cdna(arch)
+    self.arch = arch
+    self.tensor_cores = AMDRenderer.get_tensor_cores(arch)
+    self.is_cdna = AMDRenderer.is_cdna(arch)
     self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, cdna=self.is_cdna: render_wmma_amd(ctx, wmma, cdna))])
     if self.is_cdna:
       self.extra_matcher += PatternMatcher([

tinygrad/renderer/ptx.py

@@ -145,8 +145,8 @@ class PTXRenderer(Renderer):
   code_for_op = asm_for_op
   extra_matcher = ptx_matcher
   def __init__(self, arch:str, device="CUDA"):
-    self.device, self.arch, arch_ver = device, arch, int(arch[3:])
-    self.tensor_cores = PTXRenderer.tc_sm80 if arch_ver >= 80 else tc.cuda_sm75 if arch_ver >= 75 else []
+    self.device, self.arch = device, arch
+    self.tensor_cores = PTXRenderer.tc_sm80 if int(arch[3:]) >= 80 else tc.cuda_sm75 if int(arch[3:]) >= 75 else []
   def __reduce__(self): return self.__class__, (self.arch, self.device)
 
   # language options

tinygrad/runtime/ops_amd.py

@@ -9,15 +9,15 @@ from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerSet, CompilerPair
 from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar
 from tinygrad.helpers import VIZ, AMD_CC, AMD_LLVM, ceildiv
-from tinygrad.renderer.cstyle import AMDHIPRenderer, AMDHIPCCRenderer
+from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
 from tinygrad.runtime.autogen.am import am
+from tinygrad.runtime.support.compiler_amd import HIPCompiler, HIPCCCompiler, AMDLLVMCompiler
 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, PCIDevice, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
-from tinygrad.runtime.support.memory import AddrSpace
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
 
 SQTT = ContextVar("SQTT", abs(VIZ.value)>=2)
@@ -648,11 +648,7 @@ class AMDQueueDesc:
 
     # Flush hdp if queue is in dev mem.
     if dev.is_am() and not dev.is_usb(): dev.iface.dev_impl.gmc.flush_hdp()
-    try:
-      for doorbell in self.doorbells: doorbell[0] = self.put_value if doorbell_value is None else doorbell_value
-    except Exception as e:
-      dev.error_state = e
-      raise
+    for doorbell in self.doorbells: doorbell[0] = self.put_value if doorbell_value is None else doorbell_value
 
 class KFDIface:
   kfd:FileIOInterface|None = None
@@ -834,8 +830,8 @@ class PCIIface(PCIIfaceBase):
                                     doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0), pipe=0, queue=0)
     else:
       pv = self.dev_impl.gfx.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr, wptr_addr=gart.va_addr+wptr,
-        eop_addr=eop_buffer.va_addr, eop_size=eop_buffer.size, doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_MEC_RING0), pipe=0,
-        queue=int(is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL)), aql=is_aql)
+        eop_addr=eop_buffer.va_addr, eop_size=eop_buffer.size, doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_MEC_RING0), pipe=0, queue=0,
+        aql=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL))
 
     return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbells=[self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q')],
       read_ptrs=[gart.cpu_view().view(offset=rptr, size=8, fmt='Q')], write_ptrs=[gart.cpu_view().view(offset=wptr, size=8, fmt='Q')], put_value=pv)
@@ -863,7 +859,7 @@ class USBIface(PCIIface):
     self.sys_buf, self.sys_next_off = self._dma_region(ctrl_addr=0xa000, sys_addr=0x820000, size=0x1000), 0x800
 
   def _dma_region(self, ctrl_addr, sys_addr, size):
-    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], aspace=AddrSpace.SYS, uncached=True)
+    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], system=True, uncached=True)
     return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)
 
   def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, **kwargs) -> HCQBuffer:
@@ -934,9 +930,9 @@ class AMDDevice(HCQCompiled):
     max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
     self.sdma_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20))
 
-    compilers = CompilerSet([CompilerPair(functools.partial(AMDHIPRenderer, self.arch), None),
-                             CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), None, AMD_LLVM),
-                             CompilerPair(functools.partial(AMDHIPCCRenderer, self.arch), None)], ctrl_var=AMD_CC)
+    compilers = CompilerSet([CompilerPair(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCompiler, self.arch)),
+                             CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch), AMD_LLVM),
+                             CompilerPair(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCCCompiler, self.arch))], ctrl_var=AMD_CC)
 
     super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
                      functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),

tinygrad/runtime/ops_cpu.py

@@ -5,10 +5,10 @@ from tinygrad.helpers import CPU_CC, CPU_LVP, CPU_LLVM
 from tinygrad.device import BufferSpec, DMACPURef, CompilerSet, CompilerPair
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocatorBase, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
 from tinygrad.runtime.support.hcq import CLikeArgsState
-from tinygrad.renderer.cstyle import ClangJITRenderer
+from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.renderer.llvmir import LLVMRenderer
 from tinygrad.renderer.nir import LVPRenderer
-from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler
+from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler
 from tinygrad.runtime.support.compiler_mesa import LVPCompiler
 from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.uop.ops import sint
@@ -136,6 +136,6 @@ class CPUDevice(HCQCompiled):
   def __init__(self, device:str=""):
     self.tasks:queue.Queue = queue.Queue()
     CPUWorker(self, self.tasks, thread_id=0).start()
-    compilers = CompilerSet([CompilerPair(ClangJITRenderer, None), CompilerPair(LLVMRenderer, CPULLVMCompiler, ctrl_var=CPU_LLVM),
+    compilers = CompilerSet([CompilerPair(ClangRenderer, ClangJITCompiler), CompilerPair(LLVMRenderer, CPULLVMCompiler, ctrl_var=CPU_LLVM),
                              CompilerPair(LVPRenderer, LVPCompiler, ctrl_var=CPU_LVP)], ctrl_var=CPU_CC)
     super().__init__(device, CPUAllocator(self), compilers, functools.partial(CPUProgram, self), CPUSignal, CPUComputeQueue)

tinygrad/runtime/ops_dsp.py

@@ -83,7 +83,7 @@ class DSPProgram:
   def __init__(self, dev:DSPDevice, name:str, lib:bytes):
     self.dev, self.lib = dev, lib
 
-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
+  def __call__(self, *bufs, vals:tuple[int, ...]=(), wait=False):
     if len(bufs) >= 16: raise RuntimeError(f"Too many buffers to execute: {len(bufs)}")
 
     pra, fds, attrs, _ = rpc_prep_args(ins=[var_vals_mv:=memoryview(bytearray((len(bufs)+len(vals))*4)), off_mv:=memoryview(bytearray(len(bufs)*4))],
@@ -289,7 +289,7 @@ class MockDSPRenderer(DSPRenderer):
 
 class MockDSPProgram:
   def __init__(self, name:str, lib:bytes): self.lib = lib
-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
+  def __call__(self, *bufs, vals:tuple[int, ...]=(), wait=False):
     with tempfile.NamedTemporaryFile(suffix=".out") as dsp_lib:
       dsp_lib.write(self.lib)
       dsp_lib.flush()

tinygrad/runtime/ops_hip.py

@@ -2,6 +2,7 @@ import ctypes, functools
 from tinygrad.helpers import init_c_var, mv_address, init_c_struct_t, getenv
 from tinygrad.device import Compiled, LRUAllocator, BufferSpec, CompilerSet, CompilerPair
 from tinygrad.runtime.autogen import hip
+from tinygrad.runtime.support.compiler_amd import HIPCompiler
 from tinygrad.renderer.cstyle import HIPRenderer
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
 
@@ -14,7 +15,7 @@ class HIPDevice(Compiled):
     self.arch = init_c_var(hip.hipDeviceProp_t(), lambda x: check(hip.hipGetDeviceProperties(x, self.device_id))).gcnArchName.decode()
     self.time_event_st, self.time_event_en = [init_c_var(hip.hipEvent_t(), lambda x: hip.hipEventCreate(ctypes.byref(x), 0)) for _ in range(2)]
 
-    compilers = CompilerSet([CompilerPair(functools.partial(HIPRenderer, self.arch), None)])
+    compilers = CompilerSet([CompilerPair(functools.partial(HIPRenderer, self.arch), functools.partial(HIPCompiler, self.arch))])
     super().__init__(device, HIPAllocator(self), compilers, functools.partial(HIPProgram, self))
   def synchronize(self):
     check(hip.hipSetDevice(self.device_id))

tinygrad/runtime/ops_nv.py

@@ -508,7 +508,7 @@ class PCIIface(PCIIfaceBase):
     # PCIIface's MAP_FIXED mmap will overwrite UVM allocations made by NVKIface, so don't try PCIIface if kernel driver was already used.
     if NVKIface.root is not None: raise RuntimeError("Cannot use PCIIface after NVKIface has been initialized (would corrupt UVM memory)")
     super().__init__(dev, dev_id, vendor=0x10de, devices=[(0xff00, [0x2200, 0x2400, 0x2500, 0x2600, 0x2700, 0x2800, 0x2b00, 0x2c00, 0x2d00, 0x2f00])],
-      base_class=0x03, bars=[0, 1], vram_bar=1, va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size)
+      bars=[0, 1], vram_bar=1, va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size)
     if not OSX: System.reserve_hugepages(64)
 
     self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)