Merge origin/master

1.00kit/s and not 1000it/s for value 999.5

.github/workflows/benchmark.yml

@@ -49,19 +49,19 @@ jobs:
     - name: Print macOS version
       run: sw_vers
     - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
     - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing
     - name: Run Stable Diffusion v2
       # TODO: very slow step time
-      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing
     # process replay can't capture this, the graph is too large
     - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing
     - name: Run model inference benchmark
       run: METAL=1 NOCLANG=1 python3.11 test/external/external_model_benchmark.py
     - name: Test speed vs torch
-      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
+      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py
     - name: Test tensor cores
       run: METAL=1 python3.11 test/opt/test_tensor_cores.py
     - name: Test AMX tensor cores
@@ -71,84 +71,59 @@ jobs:
         DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
         DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
     - name: Run Tensor Core GEMM (float)
-      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py
     - name: Run Tensor Core GEMM (half)
-      run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py
     - name: Run Tensor Core GEMM (bfloat16)
-      run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py
     - name: Fuzz Padded Tensor Core GEMM
       run: METAL=1 M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3.11 ./extra/gemm/fuzz_matmul.py
     - name: Run LLaMA
       run: |
-        BENCHMARK_LOG=llama_nojit JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
-        BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+        BENCHMARK_LOG=llama_nojit JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
     - name: Run LLaMA with BEAM
-      run: BENCHMARK_LOG=llama_beam JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
+      run: BENCHMARK_LOG=llama_beam JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
     - name: Run quantized LLaMA
       run: |
-        BENCHMARK_LOG=llama_int8 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8 | tee llama_int8.txt
-        BENCHMARK_LOG=llama_nf4 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4 | tee llama_nf4.txt
+        BENCHMARK_LOG=llama_int8 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8
+        BENCHMARK_LOG=llama_nf4 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4
     - name: Run quantized LLaMA3
       run: |
-        BENCHMARK_LOG=llama3_int8 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize int8 | tee llama3_int8.txt
-        BENCHMARK_LOG=llama3_nf4 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize nf4 | tee llama3_nf4.txt
+        BENCHMARK_LOG=llama3_int8 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize int8
+        BENCHMARK_LOG=llama3_nf4 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize nf4
     #- name: Run LLaMA 7B on 4 (virtual) GPUs
-    #  run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
+    #  run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing
     - name: Run GPT2
       run: |
-        BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
     - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing
     - name: Run GPT2 w HALF/BEAM
-      run: BENCHMARK_LOG=gpt2_half_beam HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+      run: BENCHMARK_LOG=gpt2_half_beam HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing
     - name: Run OLMoE
       run: BENCHMARK_LOG=olmoe python3.11 examples/olmoe.py
     - name: Train MNIST
-      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
+      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py
 
     # NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
     # the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
     #- name: Run 10 CIFAR training steps
-    #  run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
+    #  run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py
     #- name: Run 10 CIFAR training steps w HALF
-    #  run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
+    #  run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py
 
     #- name: Run 10 CIFAR training steps w BF16
-    #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+    #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py
     # TODO: too slow
     # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (Mac)
         path: |
           onnx_inference_speed.csv
-          torch_speed.txt
-          llama_unjitted.txt
-          llama_jitted.txt
-          llama_beam.txt
-          llama_int8.txt
-          llama_nf4.txt
-          llama3_int8.txt
-          llama3_nf4.txt
-          llama_four_gpu.txt
-          gpt2_unjitted.txt
-          gpt2_jitted.txt
-          gpt2_half.txt
-          gpt2_half_beam.txt
-          matmul.txt
-          matmul_half.txt
-          matmul_bfloat16.txt
-          sd.txt
-          sd_no_fp16.txt
-          sdv2.txt
-          sdxl.txt
-          beautiful_mnist.txt
-          train_cifar.txt
-          train_cifar_half.txt
-          train_cifar_bf16.txt
-          train_cifar_wino.txt
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3.11 process_replay.py
 
@@ -215,7 +190,7 @@ jobs:
     - name: Run model inference benchmark
       run: NV=1 CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
     - name: Test speed vs torch
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
+      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
     - name: Test speed vs theoretical
       run: NV=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 benchmark allreduce
@@ -226,79 +201,58 @@ jobs:
         NV=1 NV_PTX=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
     - name: Run Tensor Core GEMM (CUDA)
       run: |
-        CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
-        CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
-        CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
-        CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_fp8.txt
+        CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
+        CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
     - name: Run Tensor Core GEMM (PTX)
-      run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
+      run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
     - name: Run Tensor Core GEMM (NV)
-      run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
+      run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
     - name: Test NV=1
       run: DEBUG=2 NV=1 python -m pytest -rA test/test_tiny.py
     - name: Test CUDA=1
       run: DEBUG=2 CUDA=1 python -m pytest -rA test/test_tiny.py
     - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
     # TODO: too slow
     # - name: Run SDXL
-    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing
     - name: Run LLaMA
       run: |
-        BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
-        BENCHMARK_LOG=llama NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+        BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=llama NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
     - name: Run LLaMA with BEAM
-      run: BENCHMARK_LOG=llama_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
+      run: BENCHMARK_LOG=llama_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
     # - name: Run LLaMA 7B on 4 GPUs
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing
     # - name: Run LLaMA 7B on 6 GPUs
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
     - name: Run LLaMA-3 8B BEAM
-      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
+      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
     - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
+      run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
     - name: Run quantized LLaMA3
-      run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8 | tee llama3_fp8.txt
+      run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8
     # - name: Run LLaMA-3 8B on 6 GPUs
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
     # - name: Run LLaMA-2 70B
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
     - name: Run Mixtral 8x7B
-      run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
+      run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing
     - name: Run GPT2
       run: |
-        BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
     - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing
     - name: Run GPT2 w HALF/BEAM
-      run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+      run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (NVIDIA)
         path: |
           onnx_inference_speed.csv
-          torch_speed.txt
-          matmul.txt
-          matmul_bfloat16.txt
-          matmul_tf32.txt
-          matmul_ptx.txt
-          matmul_nv.txt
-          sd.txt
-          sdxl.txt
-          llama_unjitted.txt
-          llama_jitted.txt
-          llama_beam.txt
-          llama3_beam.txt
-          llama3_four_gpu.txt
-          llama3_six_gpu.txt
-          llama3_fp8.txt
-          llama_2_70B.txt
-          mixtral.txt
-          gpt2_unjitted.txt
-          gpt2_jitted.txt
-          gpt2_half.txt
-          gpt2_half_beam.txt
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
@@ -339,42 +293,28 @@ jobs:
     - name: HEVC Decode Benchmark
       run: VALIDATE=1 MAX_FRAMES=100 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
     - name: Train MNIST
-      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
+      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
     - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py
     - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
     - name: Run 10 CIFAR training steps w BF16
-      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
     # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
     - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
     - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
     - name: Run MLPerf resnet eval on training data
       run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
     - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+      run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
     - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+      run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
     - name: Run 10 MLPerf Bert training steps (6 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (NVIDIA Training)
-        path: |
-          beautiful_mnist.txt
-          train_cifar.txt
-          train_cifar_half.txt
-          train_cifar_bf16.txt
-          train_cifar_wino.txt
-          train_cifar_one_gpu.txt
-          train_cifar_six_gpu.txt
-          train_resnet.txt
-          train_resnet_one_gpu.txt
-          train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
@@ -426,16 +366,18 @@ jobs:
     #- name: Test speed vs torch
     #  run: |
     #    python3 -c "import torch; print(torch.__version__)"
-    #    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
+    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
     - 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
-      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
+    - 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)
-      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
+      run: |
+        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
     - name: Test AMD=1
       run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
     #- name: Test HIP=1
@@ -450,61 +392,39 @@ jobs:
     - name: Test AM warm start time
       run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
     - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
     - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing
     - name: Run LLaMA 7B
       run: |
-        BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
-        BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+        BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
     - name: Run LLaMA 7B with BEAM
-      run: BENCHMARK_LOG=llama_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
+      run: BENCHMARK_LOG=llama_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
     # - name: Run LLaMA 7B on 4 GPUs
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing
     # - name: Run LLaMA 7B on 6 GPUs
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
     - name: Run LLaMA-3 8B BEAM
-      run: BENCHMARK_LOG=llama3_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
+      run: BENCHMARK_LOG=llama3_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
     - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: BENCHMARK_LOG=llama3_beam_4gpu AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
+      run: BENCHMARK_LOG=llama3_beam_4gpu AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
     # - name: Run LLaMA-3 8B on 6 GPUs
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
     #- name: Restore amdgpu
     #  run: sudo modprobe amdgpu
     # - name: Run LLaMA-2 70B
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
     - name: Run Mixtral 8x7B
-      run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
+      run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing
     - name: Run GPT2
       run: |
-        BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 AMD=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=gpt2 AMD=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
     - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing
     - name: Run GPT2 w HALF/BEAM
-      run: BENCHMARK_LOG=gpt2_half_beam AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (AMD)
-        path: |
-          onnx_inference_speed.csv
-          torch_speed.txt
-          llama_unjitted.txt
-          llama_jitted.txt
-          llama_beam.txt
-          llama3_beam.txt
-          llama3_four_gpu.txt
-          llama3_six_gpu.txt
-          llama_2_70B.txt
-          gpt2_unjitted.txt
-          gpt2_jitted.txt
-          gpt2_half.txt
-          gpt2_half_beam.txt
-          matmul.txt
-          matmul_amd.txt
-          sd.txt
-          sdxl.txt
-          mixtral.txt
+      run: BENCHMARK_LOG=gpt2_half_beam AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
@@ -541,31 +461,20 @@ jobs:
     - name: reset process replay
       run: test/external/process_replay/reset.py
     - name: Train MNIST
-      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
+      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
     - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py
     - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
     # - name: Run 10 CIFAR training steps w BF16
-    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
     # TODO: too slow
     # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
     - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
     - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (AMD Training)
-        path: |
-          beautiful_mnist.txt
-          train_cifar.txt
-          train_cifar_half.txt
-          train_cifar_bf16.txt
-          train_cifar_wino.txt
-          train_cifar_one_gpu.txt
-          train_cifar_six_gpu.txt
+      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
@@ -604,19 +513,12 @@ jobs:
     - name: Run MLPerf resnet eval
       run: time BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
     - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+      run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
     - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+      run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
     - name: Run 10 MLPerf Bert training steps (6 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (AMD MLPerf)
-        path: |
-          train_resnet.txt
-          train_resnet_one_gpu.txt
-          train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
@@ -706,7 +608,7 @@ jobs:
     #     AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
     #     AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
     - 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 am_matmul_amd.txt
+      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
     - name: Test AMD=1
       run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
     - name: Test DISK copy time
@@ -716,20 +618,12 @@ jobs:
         AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
         AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
     - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
     # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-    #   run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
+    #   run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
     - name: Run 10 MLPerf Bert training steps (1 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (AM Driver)
-        path: |
-          am_matmul_amd.txt
-          am_train_cifar_one_gpu.txt
-          am_train_resnet_one_gpu.txt
-          am_train_bert_one_gpu.txt
+      run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
@@ -776,21 +670,13 @@ jobs:
         NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
         NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
     - name: Test LLAMA-3
-      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
+      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0
     - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
     - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
+      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
     - name: Run 10 MLPerf Bert training steps (1 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (NV Driver)
-        path: |
-          nv_llama3_beam.txt
-          nv_train_cifar_one_gpu.txt
-          nv_train_resnet_one_gpu.txt
-          nv_train_bert_one_gpu.txt
+      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

.github/workflows/test.yml

@@ -5,6 +5,7 @@ env:
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
+  IGNORE_OOB: 0
 
 on:
   push:
@@ -36,6 +37,8 @@ jobs:
     name: Docs
     runs-on: ubuntu-22.04
     timeout-minutes: 10
+    env:
+      IGNORE_OOB: 1
     steps:
     - name: Checkout Code
       uses: actions/checkout@v4
@@ -102,15 +105,11 @@ jobs:
       run: |
         sudo apt update || true
         sudo apt install -y --no-install-recommends ninja-build
-    - name: Lint with ruff
-      run: |
-        pip3 install --upgrade --force-reinstall ruff==0.11.0
-        python3 -m ruff check extra/torch_backend/backend.py
     - name: Test one op
       run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
     - name: Test ResNet-18
       run: DEBUG=2 python3 extra/torch_backend/example.py
-    - name: My (custom) tests
+    - name: custom tests
       run: python3 extra/torch_backend/test.py
     - name: Test one op in torch tests
       run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
@@ -219,7 +218,6 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 10
 
-    # TODO: run the pre-commit hook to replace a lot of this
     steps:
     - name: Checkout Code
       uses: actions/checkout@v4
@@ -233,17 +231,17 @@ jobs:
       run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string='  ' --recursive=y .
     - name: Lint with ruff
       run: |
-        pip3 install --upgrade --force-reinstall ruff==0.11.0
-        python3 -m ruff check .
+        pip3 install --upgrade --force-reinstall ruff==0.14.10
+        pre-commit run ruff --all-files
         python3 -m ruff check examples/mlperf/ --ignore E501
         python3 -m ruff check extra/thunder/tiny/ --ignore E501 --ignore F841 --ignore E722
+        python3 -m ruff check extra/torch_backend/backend.py
     - name: Run mypy
       run: |
-        python -m mypy --strict-equality --lineprecision-report .
+        python -m mypy --lineprecision-report .
         cat lineprecision.txt
-    # broken because of UPatAny
-    #- name: Run TYPED=1
-    #  run: TYPED=1 python -c "import tinygrad"
+    - name: Run TYPED=1
+      run: TYPED=1 python -c "import tinygrad"
 
   unittest:
     name: Unit Tests
@@ -311,7 +309,7 @@ jobs:
         deps: testing_unit
         python-version: '3.14'
     - name: Test SPEC=2
-      run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
+      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
 
   fuzzing:
     name: Fuzzing
@@ -474,6 +472,8 @@ jobs:
     name: Test LLM
     runs-on: ubuntu-24.04
     timeout-minutes: 15
+    env:
+      IGNORE_OOB: 1
     steps:
       - name: Checkout Code
         uses: actions/checkout@v4
@@ -642,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 --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 test/testextra/test_cfg_viz.py --durations=20
       - name: Run pytest (amd)
         run: python -m pytest test/external/external_test_am.py --durations=20
       - name: Run TRANSCENDENTAL math
@@ -655,6 +655,53 @@ jobs:
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
+  testamdasm:
+    name: AMD ASM IDE
+    runs-on: ubuntu-24.04
+    timeout-minutes: 10
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: rdna3-emu
+          deps: testing_minimal
+          amd: 'true'
+      - name: Install LLVM 21
+        run: |
+          wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
+          echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-21 main" | sudo tee /etc/apt/sources.list.d/llvm.list
+          sudo apt-get update
+          sudo apt-get install llvm-21 llvm-21-tools cloc
+      - name: RDNA3 Line Count
+        run: cloc --by-file extra/assembly/amd/*.py
+      - name: Run RDNA3 emulator tests
+        run: python -m pytest -n=auto extra/assembly/amd/ --durations 20
+      - name: Run RDNA3 emulator tests (AMD_LLVM=1)
+        run: AMD_LLVM=1 python -m pytest -n=auto extra/assembly/amd/ --durations 20
+      - name: Run RDNA3 dtype tests
+        run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
+      - name: Run RDNA3 dtype tests (AMD_LLVM=1)
+        run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
+
+  testamdautogen:
+    name: AMD autogen
+    runs-on: ubuntu-24.04
+    timeout-minutes: 10
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: rdna3-autogen
+          pydeps: "pdfplumber"
+      - name: Verify AMD autogen is up to date
+        run: |
+          python -m extra.assembly.amd.pdf --arch all
+          git diff --exit-code extra/assembly/amd/autogen/
+
   testnvidia:
     strategy:
       fail-fast: false

.pre-commit-config.yaml

@@ -16,7 +16,7 @@ repos:
         pass_filenames: false
       - id: mypy
         name: mypy
-        entry: python3 -m mypy tinygrad/ --strict-equality
+        entry: python3 -m mypy
         language: system
         always_run: true
         pass_filenames: false

CLAUDE.md

@@ -34,33 +34,6 @@ result = graph_rewrite(uop, pm)
 ### Schedule Cache
 Schedules are cached by graph structure. BIND nodes (variables with bound values) are unbound before cache key computation so different values hit the same cache.
 
-## Directory Structure
-
-```
-tinygrad/
-├── tensor.py          # Tensor class, user API
-├── device.py          # Buffer, device management
-├── dtype.py           # Data types
-├── helpers.py         # Utilities, environment vars
-├── uop/
-│   ├── ops.py         # UOp class, Ops enum, PatternMatcher
-│   ├── spec.py        # UOp type verification
-│   └── symbolic.py    # Symbolic math simplification
-├── engine/
-│   ├── schedule.py    # Schedule creation, caching
-│   ├── realize.py     # Tensor realization
-│   ├── jit.py         # JIT compilation
-│   └── memory.py      # Memory planning
-├── schedule/
-│   ├── rangeify.py    # Convert movements to ranges
-│   └── indexing.py    # Index calculations
-├── codegen/
-│   ├── kernel.py      # Kernel optimization
-│   └── uopgraph.py    # UOp graph transformations
-├── renderer/          # Code generation (CUDA, Metal, etc.)
-└── runtime/           # Device backends
-```
-
 ## Testing
 
 ```bash
@@ -79,7 +52,7 @@ VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
 
 ## Common Environment Variables
 
-- `DEBUG=1-4` - Increasing verbosity
+- `DEBUG=1-7` - Increasing verbosity (7 shows assembly output)
 - `VIZ=1` - Enable graph visualization
 - `SPEC=1` - Enable UOp spec verification
 - `NOOPT=1` - Disable optimizations
@@ -100,6 +73,16 @@ VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
 - Run tests before proposing commits
 - Test with `SPEC=2` when modifying UOp-related code
 
+## Auto-generated Files (DO NOT EDIT)
+
+The following files are auto-generated and should never be edited manually:
+- `extra/assembly/amd/autogen/{arch}/__init__.py` - Generated by `python -m extra.assembly.amd.dsl --arch {arch}`
+- `extra/assembly/amd/autogen/{arch}/gen_pcode.py` - Generated by `python -m extra.assembly.amd.pcode --arch {arch}`
+
+Where `{arch}` is one of: `rdna3`, `rdna4`, `cdna`
+
+To add missing instruction implementations, add them to `extra/assembly/amd/emu.py` instead.
+
 ## Style Notes
 
 - 2-space indentation, 150 char line limit
@@ -225,3 +208,9 @@ Key patterns to watch (from ResNet50 benchmark):
 - `vmin==vmax folding`: ~55ms, 0.33% match rate - checks 52K ops but rarely matches
 
 Patterns with 0% match rate are workload-specific overhead. They may be useful in other workloads, so don't remove them without understanding their purpose.
+
+## AMD Performance Counter Profiling
+
+Set VIZ to `-2` to save performance counters traces for the AMD backend.
+
+Use the CLI in `./extra/sqtt/roc.py` to explore the trace.

docs/developer/layout.md

@@ -28,7 +28,7 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi
 
 Transform the optimized ast into a linearized and rendered program.
 
-::: tinygrad.codegen.full_rewrite_to_program
+::: tinygrad.codegen.get_program
     options:
         members: false
         show_labels: false

examples/mlperf/dataloader.py

@@ -763,48 +763,26 @@ class BlendedGPTDataset:
 
     return dataset_idx, dataset_sample_idx
 
-def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
+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)
   if val:
-    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)
+    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)
 
-  for b in range(math.ceil(samples / bs)):
-    batch = []
-    for i in range(bs):
-      tokens = dataset.get(b * bs + i)
-      batch.append(tokens)
+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)]
     yield Tensor.stack(batch, dim=0)
 
-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)
+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)
 
 if __name__ == "__main__":
   def load_unet3d(val):

examples/mlperf/model_eval.py

@@ -234,12 +234,9 @@ def eval_llama3():
     loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
     return loss.flatten().float()
 
-  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)
+  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)
 
   losses = []
   for tokens in tqdm(iter, total=5760//BS):

examples/mlperf/model_train.py

@@ -1314,6 +1314,13 @@ 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}
@@ -1417,23 +1424,20 @@ def train_llama3():
     if getenv("FAKEDATA", 0):
       return fake_data(BS, SAMPLES)
     else:
-      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))
+      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))
 
   def get_eval_iter():
-    if getenv("FAKEDATA", 0):
+    if eval_dataset is None:
       return fake_data(EVAL_BS, 5760)
-    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)
+    from examples.mlperf.dataloader import iterate_llama3_dataset
+    return iterate_llama3_dataset(eval_dataset, EVAL_BS)
 
   iter = get_train_iter()
   i, sequences_seen = resume_ckpt, 0
@@ -1449,13 +1453,17 @@ def train_llama3():
       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, {GlobalCounters.mem_used / 1e9:.2f} GB used, "
-        f"{GlobalCounters.global_ops * 1e-9 / sec:9.2f} GFLOPS")
+        f"{i:5} {sec:.2f} s run, {loss:.4f} loss, {lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS")
 
       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")
+          f.write(f"{i} {loss:.4f} {lr:.12f} {mem_gb:.2f}\n")
+
+      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")
@@ -1481,6 +1489,9 @@ 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/assembly/amd/README

@@ -0,0 +1,31 @@
+An integrated environment for AMD GPU assembly and emulation
+
+Test with `PYTHONPATH="." pytest -n12 extra/assembly/amd/`
+`AMD_LLVM=1 PYTHONPATH="." pytest -n12 extra/assembly/amd/`
+
+* pdf.py -- extract assembly format + instruction psuedocode from AMD PDF
+* dsl.py -- helpers for the autogen instruction classes in `__init__.py`. should be standalone with init
+* pcode.py -- psuedocode execution environment. psuedocode should be transformed as little as possible.
+* asm.py -- an asm/disasm function to transform to and from AMD assembly syntax
+* emu.py -- an emulator for RDNA that runs in tinygrad with `AMD=1 MOCKGPU=1 PYTHON_REMU=1`
+
+The code should be as readable and deduplicated as possible. asm and emu shouldn't be required for dsl.
+
+test_emu.py has a good set of instruction tests for the emulation, with USE_HW=1 it will compare to real hardware.
+Whenever an instruction is fixed, regression tests should be added here and confirmed with real hardware.
+
+test_llvm.py tests asm/disasm on the LLVM tests, confirming it behaves the same as LLVM.
+
+tinygrad's dtype tests should pass with and without LLVM. they run in about 12 seconds.
+
+`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`
+`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`
+
+The ops tests also pass, but they are very slow, so you should run them one at a time.
+
+`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_ops.py`
+`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_ops.py`
+
+When something is caught by main tinygrad tests, a local regression test should be added to `extra/assembly/amd/test`. While working with tinygrad, you can dump the assembly with `DEBUG=7`. These tests all pass on real hardware, so if a test is failing with `AMD=1 PYTHON_REMU=1 MOCKGPU=1` it's likely because an instruction is emulated incorrectly. You can test without `MOCKGPU=1` to test on real hardware, if it works on real hardware there's a bug in the emulator.
+
+Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~2000 lines.

extra/assembly/amd/dsl.py

@@ -0,0 +1,574 @@
+# library for RDNA3 assembly DSL
+# mypy: ignore-errors
+from __future__ import annotations
+import struct, math, re
+from enum import IntEnum
+from functools import cache, cached_property
+from typing import overload, Annotated, TypeVar, Generic
+from extra.assembly.amd.autogen.rdna3.enum import (VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, VOPDOp, SOP1Op, SOP2Op,
+  SOPCOp, SOPKOp, SOPPOp, SMEMOp, DSOp, FLATOp, MUBUFOp, MTBUFOp, MIMGOp, VINTERPOp)
+
+# Common masks and bit conversion functions
+MASK32, MASK64 = 0xffffffff, 0xffffffffffffffff
+_struct_f, _struct_I = struct.Struct("<f"), struct.Struct("<I")
+_struct_e, _struct_H = struct.Struct("<e"), struct.Struct("<H")
+_struct_d, _struct_Q = struct.Struct("<d"), struct.Struct("<Q")
+def _f32(i): return _struct_f.unpack(_struct_I.pack(i & MASK32))[0]
+def _i32(f):
+  if isinstance(f, int): f = float(f)
+  if math.isnan(f): return 0xffc00000 if math.copysign(1.0, f) < 0 else 0x7fc00000
+  if math.isinf(f): return 0x7f800000 if f > 0 else 0xff800000
+  try: return _struct_I.unpack(_struct_f.pack(f))[0]
+  except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
+def _sext(v, b): return v - (1 << b) if v & (1 << (b - 1)) else v
+def _f16(i): return _struct_e.unpack(_struct_H.pack(i & 0xffff))[0]
+def _i16(f):
+  if math.isnan(f): return 0x7e00
+  if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
+  try: return _struct_H.unpack(_struct_e.pack(f))[0]
+  except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
+def _f64(i): return _struct_d.unpack(_struct_Q.pack(i & MASK64))[0]
+def _i64(f):
+  if math.isnan(f): return 0x7ff8000000000000
+  if math.isinf(f): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
+  try: return _struct_Q.unpack(_struct_d.pack(f))[0]
+  except (OverflowError, struct.error): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
+
+# Instruction spec - register counts and dtypes derived from instruction names
+_REGS = {'B32': 1, 'B64': 2, 'B96': 3, 'B128': 4, 'B256': 8, 'B512': 16,
+         'F32': 1, 'I32': 1, 'U32': 1, 'F64': 2, 'I64': 2, 'U64': 2,
+         'F16': 1, 'I16': 1, 'U16': 1, 'B16': 1, 'I8': 1, 'U8': 1, 'B8': 1}
+_CVT_RE = re.compile(r'CVT_([FIUB]\d+)_([FIUB]\d+)$')
+_MAD_MUL_RE = re.compile(r'(?:MAD|MUL)_([IU]\d+)_([IU]\d+)$')
+_PACK_RE = re.compile(r'PACK_([FIUB]\d+)_([FIUB]\d+)$')
+_DST_SRC_RE = re.compile(r'_([FIUB]\d+)_([FIUB]\d+)$')
+_SINGLE_RE = re.compile(r'_([FIUB](?:32|64|16|8|96|128|256|512))$')
+@cache
+def _suffix(name: str) -> tuple[str | None, str | None]:
+  name = name.upper()
+  if m := _CVT_RE.search(name): return m.group(1), m.group(2)
+  if m := _MAD_MUL_RE.search(name): return m.group(1), m.group(2)
+  if m := _PACK_RE.search(name): return m.group(1), m.group(2)
+  if m := _DST_SRC_RE.search(name): return m.group(1), m.group(2)
+  if m := _SINGLE_RE.search(name): return m.group(1), m.group(1)
+  return None, None
+_SPECIAL_REGS = {
+  'V_LSHLREV_B64': (2, 1, 2, 1), 'V_LSHRREV_B64': (2, 1, 2, 1), 'V_ASHRREV_I64': (2, 1, 2, 1),
+  'S_LSHL_B64': (2, 2, 1, 1), 'S_LSHR_B64': (2, 2, 1, 1), 'S_ASHR_I64': (2, 2, 1, 1),
+  'S_BFE_U64': (2, 2, 1, 1), 'S_BFE_I64': (2, 2, 1, 1), 'S_BFM_B64': (2, 1, 1, 1),
+  'S_BITSET0_B64': (2, 1, 1, 1), 'S_BITSET1_B64': (2, 1, 1, 1),
+  'S_BITCMP0_B64': (1, 2, 1, 1), 'S_BITCMP1_B64': (1, 2, 1, 1),
+  'V_LDEXP_F64': (2, 2, 1, 1), 'V_TRIG_PREOP_F64': (2, 2, 1, 1),
+  'V_CMP_CLASS_F64': (1, 2, 1, 1), 'V_CMPX_CLASS_F64': (1, 2, 1, 1),
+  'V_CMP_CLASS_F32': (1, 1, 1, 1), 'V_CMPX_CLASS_F32': (1, 1, 1, 1),
+  'V_CMP_CLASS_F16': (1, 1, 1, 1), 'V_CMPX_CLASS_F16': (1, 1, 1, 1),
+  'V_MAD_U64_U32': (2, 1, 1, 2), 'V_MAD_I64_I32': (2, 1, 1, 2), 'V_MAD_CO_U64_U32': (2, 1, 1, 2), 'V_MAD_CO_I64_I32': (2, 1, 1, 2),
+  'V_QSAD_PK_U16_U8': (2, 2, 1, 2), 'V_MQSAD_PK_U16_U8': (2, 2, 1, 2), 'V_MQSAD_U32_U8': (4, 2, 1, 4),
+  # RDNA4 CVT_PK_F32 instructions output 2 F32 values (64-bit)
+  'V_CVT_PK_F32_BF8': (2, 1, 1, 1), 'V_CVT_PK_F32_FP8': (2, 1, 1, 1),
+}
+_SPECIAL_DTYPE = {
+  'V_LSHLREV_B64': ('B64', 'U32', 'B64', None), 'V_LSHRREV_B64': ('B64', 'U32', 'B64', None), 'V_ASHRREV_I64': ('I64', 'U32', 'I64', None),
+  'S_LSHL_B64': ('B64', 'B64', 'U32', None), 'S_LSHR_B64': ('B64', 'B64', 'U32', None), 'S_ASHR_I64': ('I64', 'I64', 'U32', None),
+  'S_BFE_U64': ('U64', 'U64', 'U32', None), 'S_BFE_I64': ('I64', 'I64', 'U32', None),
+  'S_BFM_B64': ('B64', 'U32', 'U32', None), 'S_BITSET0_B64': ('B64', 'U32', None, None), 'S_BITSET1_B64': ('B64', 'U32', None, None),
+  'S_BITCMP0_B64': ('SCC', 'B64', 'U32', None), 'S_BITCMP1_B64': ('SCC', 'B64', 'U32', None),
+  'V_LDEXP_F64': ('F64', 'F64', 'I32', None), 'V_TRIG_PREOP_F64': ('F64', 'F64', 'U32', None),
+  'V_CMP_CLASS_F64': ('VCC', 'F64', 'U32', None), 'V_CMPX_CLASS_F64': ('EXEC', 'F64', 'U32', None),
+  'V_CMP_CLASS_F32': ('VCC', 'F32', 'U32', None), 'V_CMPX_CLASS_F32': ('EXEC', 'F32', 'U32', None),
+  'V_CMP_CLASS_F16': ('VCC', 'F16', 'U32', None), 'V_CMPX_CLASS_F16': ('EXEC', 'F16', 'U32', None),
+  'V_MAD_U64_U32': ('U64', 'U32', 'U32', 'U64'), 'V_MAD_I64_I32': ('I64', 'I32', 'I32', 'I64'),
+  'V_QSAD_PK_U16_U8': ('B64', 'B64', 'B64', 'B64'), 'V_MQSAD_PK_U16_U8': ('B64', 'B64', 'B64', 'B64'),
+  'V_MQSAD_U32_U8': ('B128', 'B64', 'B64', 'B128'),
+  # RDNA4 CVT_PK_F32 instructions: source is 8-bit packed as 16-bit operand
+  'V_CVT_PK_F32_BF8': ('F32', 'B16', None, None), 'V_CVT_PK_F32_FP8': ('F32', 'B16', None, None),
+}
+@cache
+def spec_regs(name: str) -> tuple[int, int, int, int]:
+  uname = name.upper()
+  if uname in _SPECIAL_REGS: return _SPECIAL_REGS[uname]
+  if 'SAD' in uname and 'U8' in uname and 'QSAD' not in uname and 'MQSAD' not in uname: return 1, 1, 1, 1
+  dst_suf, src_suf = _suffix(name)
+  return _REGS.get(dst_suf, 1), _REGS.get(src_suf, 1), _REGS.get(src_suf, 1), _REGS.get(src_suf, 1)
+@cache
+def spec_dtype(name: str) -> tuple[str | None, str | None, str | None, str | None]:
+  uname = name.upper()
+  if uname in _SPECIAL_DTYPE: return _SPECIAL_DTYPE[uname]
+  if 'SAD' in uname and ('U8' in uname or 'U16' in uname) and 'QSAD' not in uname and 'MQSAD' not in uname: return 'U32', 'U32', 'U32', 'U32'
+  if '_CMP_' in uname or '_CMPX_' in uname:
+    dst_suf, src_suf = _suffix(name)
+    return 'EXEC' if '_CMPX_' in uname else 'VCC', src_suf, src_suf, None
+  dst_suf, src_suf = _suffix(name)
+  return dst_suf, src_suf, src_suf, src_suf
+_F16_RE = re.compile(r'_[FIUB]16(?:_|$)')
+_F64_RE = re.compile(r'_[FIUB]64(?:_|$)')
+@cache
+def spec_is_16bit(name: str) -> bool:
+  uname = name.upper()
+  if 'SAD' in uname or 'PACK' in uname or '_PK_' in uname or 'SAT_PK' in uname or 'DOT2' in uname: return False
+  if '_F32' in uname or '_I32' in uname or '_U32' in uname or '_B32' in uname: return False
+  return bool(_F16_RE.search(uname))
+@cache
+def spec_is_64bit(name: str) -> bool: return bool(_F64_RE.search(name.upper()))
+_3SRC = {'FMA', 'MAD', 'MIN3', 'MAX3', 'MED3', 'DIV_FIX', 'DIV_FMAS', 'DIV_SCALE', 'SAD', 'LERP', 'ALIGN', 'CUBE', 'BFE', 'BFI',
+         'PERM_B32', 'PERMLANE', 'CNDMASK', 'XOR3', 'OR3', 'ADD3', 'LSHL_OR', 'AND_OR', 'LSHL_ADD', 'ADD_LSHL', 'XAD', 'MAXMIN',
+         'MINMAX', 'MAXIMUMMINIMUM', 'MINIMUMMAXIMUM', 'MAXIMUM3', 'MINIMUM3', 'DOT2', 'DOT4', 'DOT8', 'WMMA', 'CVT_PK_U8', 'MULLIT', 'CO_CI'}
+_2SRC = {'FMAC'}  # FMAC uses dst as implicit accumulator, so only 2 explicit sources
+def spec_num_srcs(name: str) -> int:
+  name = name.upper()
+  if any(k in name for k in _2SRC): return 2
+  return 3 if any(k in name for k in _3SRC) else 2
+def is_dtype_16(dt: str | None) -> bool: return dt is not None and '16' in dt
+def is_dtype_64(dt: str | None) -> bool: return dt is not None and '64' in dt
+
+# Bit field DSL
+class BitField:
+  def __init__(self, hi: int, lo: int, name: str | None = None): self.hi, self.lo, self.name, self._marker = hi, lo, name, None
+  def __set_name__(self, owner, name):
+    import typing
+    self.name, self._owner = name, owner
+    # Cache marker at class definition time
+    hints = typing.get_type_hints(owner, include_extras=True)
+    if name in hints:
+      hint = hints[name]
+      if typing.get_origin(hint) is Annotated:
+        args = typing.get_args(hint)
+        self._marker = args[1] if len(args) > 1 else None
+  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: return self._marker
+  @overload
+  def __get__(self, obj: None, objtype: type) -> BitField: ...
+  @overload
+  def __get__(self, obj: object, objtype: type | None = None) -> int: ...
+  # Map RDNA4 class names to their corresponding enum names for op field dynamic lookup
+  _RDNA4_OP_ENUMS = {'VDS': 'DSOp', 'VBUFFER': 'VBUFFEROp', 'VEXPORT': 'EXPOp', 'VFLAT': 'VFLATOp', 'VGLOBAL': 'VGLOBALOp',
+                     'VSCRATCH': 'VSCRATCHOp', 'VIMAGE': 'VIMAGEOp', 'VSAMPLE': 'VSAMPLEOp', 'VDSDIR': 'VDSDIROp'}
+
+  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):
+      # VOP3 with VOPC opcodes (0-255) -> VOPCOp, VOP3SD opcodes -> VOP3SDOp
+      # Check by name to handle both RDNA3 and RDNA4 enums
+      if self.marker.__name__ == 'VOP3Op':
+        # Get the appropriate enums from the same module as the marker
+        marker_mod = self.marker.__module__
+        import importlib
+        enum_mod = importlib.import_module(marker_mod)
+        if val < 256: return enum_mod.VOPCOp(val)
+        if val in Inst._VOP3SD_OPS: return enum_mod.VOP3SDOp(val)
+      try: return self.marker(val)
+      except ValueError: pass
+    # For RDNA4 op fields without type annotations, dynamically look up enum
+    elif self.name == 'op' and 'rdna4' in obj.__class__.__module__:
+      import importlib
+      enum_mod = importlib.import_module('extra.assembly.amd.autogen.rdna4.enum')
+      cls_name = obj.__class__.__name__
+      enum_name = self._RDNA4_OP_ENUMS.get(cls_name, cls_name + 'Op')
+      if hasattr(enum_mod, enum_name):
+        try: return getattr(enum_mod, enum_name)(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()
+
+# Source operand with modifiers - base class for anything that can be a src with neg/abs
+class SrcMod:
+  __slots__ = ('val', 'neg', 'abs_')
+  def __init__(self, val: int, neg: bool = False, abs_: bool = False): self.val, self.neg, self.abs_ = val, neg, abs_
+  def __repr__(self): return f"{'-' if self.neg else ''}{'|' if self.abs_ else ''}{self.val}{'|' if self.abs_ else ''}"
+  def __neg__(self): return SrcMod(self.val, not self.neg, self.abs_)
+  def __abs__(self): return SrcMod(self.val, self.neg, True)
+
+# Register types
+class Reg(SrcMod):
+  __slots__ = ('idx', 'count', 'hi')
+  def __init__(self, idx: int, count: int = 1, hi: bool = False, neg: bool = False, abs_: bool = False):
+    self.idx, self.count, self.hi = idx, count, hi
+    super().__init__(idx, neg, abs_)
+  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, not self.neg, self.abs_)
+  def __abs__(self): return self.__class__(self.idx, self.count, self.hi, self.neg, True)
+  @property
+  def l(self): return self.__class__(self.idx, self.count, False, self.neg, self.abs_)
+  @property
+  def h(self): return self.__class__(self.idx, self.count, True, self.neg, self.abs_)
+
+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")
+
+# Special registers as SrcMod objects (support -VCC_LO, abs(EXEC_LO), etc.)
+VCC_LO, VCC_HI, VCC = SrcMod(106), SrcMod(107), SrcMod(106)
+EXEC_LO, EXEC_HI, EXEC = SrcMod(126), SrcMod(127), SrcMod(126)
+SCC, M0, NULL, OFF = SrcMod(253), SrcMod(125), SrcMod(124), SrcMod(124)
+
+# 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:
+  if isinstance(val, RawImm): return val.val
+  if isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val  # Special registers like VCC_LO, NULL
+  if hasattr(val, 'value'): return val.value  # IntEnum
+  if hasattr(val, 'idx'): return val.idx  # Reg
+  return val
+
+# Encoding/decoding constants
+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}
+FLOAT_DEC = {v: str(k) for k, v in FLOAT_ENC.items()}
+SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
+SPECIAL_PAIRS = {106: "vcc", 126: "exec"}
+SRC_FIELDS = {'src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'soffset', 'srcx0', 'srcy0'}
+RAW_FIELDS = {'vdata', 'vdst', 'vaddr', 'addr', 'data', 'data0', 'data1', 'sdst', 'sdata', 'vsrc1'}
+
+def _encode_reg(val: Reg) -> int: return (108 if isinstance(val, TTMP) else 0) + val.idx
+
+def _is_inline_const(v: int) -> bool: return 0 <= v <= 127 or 128 <= v <= 208 or 240 <= v <= 255
+
+def encode_src(val) -> int:
+  if isinstance(val, VGPR): return 256 + _encode_reg(val)
+  if isinstance(val, Reg): return _encode_reg(val)
+  if isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val if _is_inline_const(val.val) else 255
+  if hasattr(val, 'value'): return val.value  # IntEnum
+  if isinstance(val, float): return 128 if val == 0.0 else FLOAT_ENC.get(val, 255)
+  if isinstance(val, int): return 128 + val if 0 <= val <= 64 else 192 - val if -16 <= val <= -1 else 255
+  return 255
+
+def decode_src(val: int) -> str:
+  if val <= 105: return f"s{val}"
+  if val in SPECIAL_GPRS: return SPECIAL_GPRS[val]
+  if val in FLOAT_DEC: return FLOAT_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}"
+
+# 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 _or_field(self, name: str, bit: int):
+    cur = self._values.get(name, 0)
+    self._values[name] = (cur.val if isinstance(cur, RawImm) else cur) | bit
+
+  def _encode_src(self, name: str, val):
+    """Encode a source field, handling modifiers and literals."""
+    encoded = encode_src(val)
+    has_opsel = 'opsel' in self._fields
+    if isinstance(val, Reg) and val.hi and not has_opsel: encoded |= 0x80  # hi bit in src for VOP1/2/C
+    self._values[name] = RawImm(encoded)
+    # Handle neg/abs/opsel modifiers
+    if isinstance(val, SrcMod):
+      mod_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
+      if val.neg and 'neg' in self._fields: self._or_field('neg', mod_bit)
+      if val.abs_ and 'abs' in self._fields: self._or_field('abs', mod_bit)
+    if isinstance(val, Reg) and val.hi and has_opsel:
+      self._or_field('opsel', {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0))
+    # Track literal value if needed
+    if encoded == 255 and self._literal is None:
+      import struct
+      # Check if THIS source uses 64-bit encoding (not just src0)
+      src_idx = {'src0': 0, 'src1': 1, 'src2': 2, 'ssrc0': 0, 'ssrc1': 1}.get(name, 0)
+      src_regs = self.src_regs(src_idx)
+      is_64 = src_regs == 2
+      if isinstance(val, SrcMod) and not isinstance(val, Reg): lit32 = val.val & MASK32
+      elif isinstance(val, int) and not isinstance(val, IntEnum): lit32 = val & MASK32
+      elif isinstance(val, float): lit32 = (_i64(val) >> 32) if is_64 else _i32(val)  # f64: high 32 bits of f64 repr
+      else: return
+      self._literal = (lit32 << 32) if is_64 else lit32
+
+  def _encode_raw(self, name: str, val):
+    """Encode a raw register field (vdst, vdata, etc.)."""
+    if isinstance(val, Reg):
+      encoded = _encode_reg(val)
+      if val.hi and 'opsel' not in self._fields: encoded |= 0x80
+      self._values[name] = encoded
+      if name == 'vdst' and val.hi and 'opsel' in self._fields: self._or_field('opsel', 8)
+    elif hasattr(val, 'value'): self._values[name] = val.value
+
+  def _validate(self, orig_args: dict):
+    """Format-specific validation. Override in subclass or check by class name."""
+    cls_name, op = self.__class__.__name__, orig_args.get('op')
+    if hasattr(op, 'value'): op = op.value
+    # SMEM: register count must match opcode
+    if cls_name == 'SMEM' and op is not None:
+      expected = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op)
+      sdata = orig_args.get('sdata')
+      if expected and isinstance(sdata, Reg) and sdata.count != expected:
+        raise ValueError(f"SMEM op {op} expects {expected} registers, got {sdata.count}")
+    # SOP1: b32=1 reg, b64=2 regs
+    if cls_name == 'SOP1' and hasattr(orig_args.get('op'), 'name'):
+      expected = 2 if orig_args['op'].name.endswith('_B64') else 1
+      for fld in ('sdst', 'ssrc0'):
+        if isinstance(orig_args.get(fld), Reg) and orig_args[fld].count != expected:
+          raise ValueError(f"SOP1 {orig_args['op'].name} expects {expected} register(s) for {fld}, got {orig_args[fld].count}")
+
+  def __init__(self, *args, literal: int | None = None, **kwargs):
+    self._values, self._literal = dict(self._defaults), None
+    field_names = [n for n in self._fields if n != 'encoding']
+    orig_args = dict(zip(field_names, args)) | kwargs
+    self._values.update(orig_args)
+    self._validate(orig_args)
+    # Pre-shift literal for 64-bit sources (literal param is always raw 32-bit value from user)
+    if literal is not None:
+      # Find which source uses the literal (255) and check its register count
+      for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
+        v = orig_args.get(n)
+        if (isinstance(v, RawImm) and v.val == 255) or (isinstance(v, int) and v == 255):
+          self._literal = (literal << 32) if self.src_regs(idx) == 2 else literal
+          break
+      else:
+        self._literal = literal  # fallback if no literal source found
+    cls_name = self.__class__.__name__
+
+    # Format-specific setup
+    if cls_name == 'FLAT' and 'sve' in self._fields:
+      seg = self._values.get('seg', 0)
+      if (seg.val if isinstance(seg, RawImm) else seg) == 1 and isinstance(orig_args.get('addr'), VGPR): self._values['sve'] = 1
+    if cls_name == 'VOP3P':
+      op = orig_args.get('op')
+      if hasattr(op, 'value'): op = op.value
+      if op in (32, 33, 34) and 'opsel_hi' not in orig_args: self._values['opsel_hi'] = self._values['opsel_hi2'] = 0
+
+    # Encode all fields
+    for name, val in list(self._values.items()):
+      if name == 'encoding': continue
+      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 and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires SGPR, got VGPR")
+      if marker is _VGPRField and 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 by field type
+      if name in SRC_FIELDS: self._encode_src(name, val)
+      elif name in RAW_FIELDS: self._encode_raw(name, val)
+      elif name == 'sbase': self._values[name] = (val.idx if isinstance(val, Reg) else val.val if isinstance(val, SrcMod) else val * 2) // 2
+      elif name in {'srsrc', 'ssamp'} and isinstance(val, Reg): self._values[name] = val.idx // 4
+      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 isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val  # Special regs like VCC_LO
+    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.val // 2 if isinstance(val, SrcMod) 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 _is_64bit_op(self) -> bool:
+    """Check if this instruction uses 64-bit operands (and thus 64-bit literals)."""
+    op = self._values.get('op')
+    if op is None: return False
+    op_name = op.name if hasattr(op, 'name') else None
+    # Look up op name from int if needed (happens in from_bytes path)
+    if op_name is None and self.__class__.__name__ == 'VOP3':
+      try: op_name = VOP3Op(op).name
+      except ValueError: pass
+    if op_name is None and self.__class__.__name__ == 'VOPC':
+      try: op_name = VOPCOp(op).name
+      except ValueError: pass
+    if op_name is None: return False
+    # V_LDEXP_F64 has 32-bit integer src1, so literal is 32-bit
+    return op_name != 'V_LDEXP_F64' and op_name.endswith(('_F64', '_B64', '_I64', '_U64'))
+
+  def to_bytes(self) -> bytes:
+    result = self.to_int().to_bytes(self._size(), 'little')
+    lit = self._get_literal() or getattr(self, '_literal', None)
+    if lit is None: return result
+    # For 64-bit sources, literal is stored in high 32 bits internally, but encoded as 4 bytes
+    # Find which source uses the literal (255) and check its register count
+    lit_src_is_64 = False
+    for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
+      if n not in self._values: continue
+      v = self._values[n]
+      if (isinstance(v, RawImm) and v.val == 255) or (isinstance(v, int) and v == 255):
+        lit_src_is_64 = self.is_src_64(idx)
+        break
+    lit32 = (lit >> 32) if lit_src_is_64 else lit
+    return result + (lit32 & MASK32).to_bytes(4, 'little')
+
+  @classmethod
+  def _size(cls) -> int: return 4 if issubclass(cls, Inst32) else 12 if issubclass(cls, Inst96) else 8
+  def size(self) -> int:
+    # Literal is always 4 bytes in the binary (for 64-bit ops, it's in high 32 bits)
+    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))
+    # VOPD fmaak/fmamk always have a literal (opx/opy value 1 or 2)
+    opx, opy = inst._values.get('opx', 0), inst._values.get('opy', 0)
+    has_literal = has_literal or (cls.__name__ == 'VOPD' and (opx in (1, 2) or opy in (1, 2)))
+    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:
+      # For 64-bit ops, the literal is 32 bits placed in the HIGH 32 bits of the 64-bit value
+      # (low 32 bits are zero). This is how AMD hardware interprets 32-bit literals for 64-bit ops.
+      # Check which source uses the literal and whether THAT source is 64-bit
+      if len(data) >= cls._size() + 4:
+        lit32 = int.from_bytes(data[cls._size():cls._size()+4], 'little')
+        # Find which source has literal (255) and check its register count
+        lit_src_is_64 = False
+        for n, idx in [('src0', 0), ('src1', 1), ('src2', 2)]:
+          if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255:
+            lit_src_is_64 = inst.src_regs(idx) == 2
+            break
+        inst._literal = (lit32 << 32) if lit_src_is_64 else lit32
+    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 __getattr__(self, name: str):
+    if name.startswith('_'): raise AttributeError(name)
+    return unwrap(self._values.get(name, 0))
+
+  def lit(self, v: int, neg: bool = False) -> str:
+    s = f"0x{self._literal:x}" if v == 255 and self._literal else decode_src(v)
+    return f"-{s}" if neg else s
+
+  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, wave_size: int = 32) -> str:
+    from extra.assembly.amd.asm import disasm
+    return disasm(self, wave_size)
+
+  _enum_map = {'VOP1': VOP1Op, 'VOP2': VOP2Op, 'VOP3': VOP3Op, 'VOP3SD': VOP3SDOp, 'VOP3P': VOP3POp, 'VOPC': VOPCOp,
+               'SOP1': SOP1Op, 'SOP2': SOP2Op, 'SOPC': SOPCOp, 'SOPK': SOPKOp, 'SOPP': SOPPOp,
+               'SMEM': SMEMOp, 'DS': DSOp, 'FLAT': FLATOp, 'MUBUF': MUBUFOp, 'MTBUF': MTBUFOp, 'MIMG': MIMGOp,
+               'VOPD': VOPDOp, 'VINTERP': VINTERPOp}
+  _VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
+
+  # Map RDNA4 class names to their corresponding enum names
+  _rdna4_enum_names = {'VDS': 'DSOp', 'VBUFFER': 'VBUFFEROp', 'VEXPORT': 'EXPOp', 'VFLAT': 'VFLATOp', 'VGLOBAL': 'VGLOBALOp',
+                       'VSCRATCH': 'VSCRATCHOp', 'VIMAGE': 'VIMAGEOp', 'VSAMPLE': 'VSAMPLEOp', 'VDSDIR': 'VDSDIROp'}
+
+  @property
+  def op(self):
+    """Return the op as an enum (e.g., VOP1Op.V_MOV_B32). VOP3 returns VOPCOp/VOP3SDOp for those op ranges."""
+    val = self._values.get('op')
+    if val is None: return None
+    if hasattr(val, 'name'): return val  # already an enum
+    cls_name = self.__class__.__name__
+    # First check if op field has an annotated enum type
+    import typing
+    if 'op' in self.__class__.__annotations__:
+      ann = self.__class__.__annotations__['op']
+      if hasattr(ann, '__metadata__'):
+        for m in typing.get_args(ann)[1:]:
+          if isinstance(m, type) and issubclass(m, IntEnum): return m(val)
+    # Check if this is an RDNA4 class (module path contains rdna4) and get enum from its module
+    if 'rdna4' in self.__class__.__module__:
+      import importlib
+      enum_mod = importlib.import_module('extra.assembly.amd.autogen.rdna4.enum')
+      enum_name = self._rdna4_enum_names.get(cls_name, cls_name + 'Op')
+      if hasattr(enum_mod, enum_name): return getattr(enum_mod, enum_name)(val)
+    # Fall back to static enum map
+    assert cls_name in self._enum_map, f"no enum map for {cls_name}"
+    return self._enum_map[cls_name](val)
+
+  @cached_property
+  def op_name(self) -> str:
+    op = self.op
+    return op.name if hasattr(op, 'name') else ''
+
+  @cached_property
+  def _spec_regs(self) -> tuple[int, int, int, int]: return spec_regs(self.op_name)
+  @cached_property
+  def _spec_dtype(self) -> tuple[str | None, str | None, str | None, str | None]: return spec_dtype(self.op_name)
+  def dst_regs(self) -> int: return self._spec_regs[0]
+  def src_regs(self, n: int) -> int: return self._spec_regs[n + 1]
+  def num_srcs(self) -> int: return spec_num_srcs(self.op_name)
+  def dst_dtype(self) -> str | None: return self._spec_dtype[0]
+  def src_dtype(self, n: int) -> str | None: return self._spec_dtype[n + 1]
+  def is_src_16(self, n: int) -> bool: return self._spec_regs[n + 1] == 1 and is_dtype_16(self._spec_dtype[n + 1])
+  def is_src_64(self, n: int) -> bool: return self._spec_regs[n + 1] == 2
+  def is_16bit(self) -> bool: return spec_is_16bit(self.op_name)
+  def is_64bit(self) -> bool: return spec_is_64bit(self.op_name)
+  def is_dst_16(self) -> bool: return self._spec_regs[0] == 1 and is_dtype_16(self._spec_dtype[0])
+
+class Inst32(Inst): pass
+class Inst64(Inst): pass
+class Inst96(Inst): pass

extra/assembly/amd/emu.py

@@ -0,0 +1,482 @@
+# RDNA3 emulator - executes compiled pseudocode from AMD ISA PDF
+# mypy: ignore-errors
+from __future__ import annotations
+import ctypes
+from extra.assembly.amd.dsl import Inst, unwrap, FLOAT_ENC, MASK32, MASK64, _f32, _i32, _sext, _f16, _i16, _f64, _i64
+from extra.assembly.amd.pcode import Reg
+from extra.assembly.amd.asm import detect_format
+from extra.assembly.amd.autogen.rdna3.gen_pcode import get_compiled_functions
+from extra.assembly.amd.autogen.rdna3.ins 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)
+
+Program = dict[int, Inst]
+WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
+VCC_LO, VCC_HI, NULL, EXEC_LO, EXEC_HI, SCC = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.SCC
+
+# Inline constants for src operands 128-254. Build tables for f32, f16, and f64 formats.
+_FLOAT_CONSTS = {v: k for k, v in FLOAT_ENC.items()} | {248: 0.15915494309189535}  # INV_2PI
+def _build_inline_consts(mask, to_bits):
+  tbl = list(range(65)) + [((-i) & mask) for i in range(1, 17)] + [0] * (127 - 81)
+  for k, v in _FLOAT_CONSTS.items(): tbl[k - 128] = to_bits(v)
+  return tbl
+_INLINE_CONSTS = _build_inline_consts(MASK32, _i32)
+_INLINE_CONSTS_F16 = _build_inline_consts(0xffff, _i16)
+_INLINE_CONSTS_F64 = _build_inline_consts(MASK64, _i64)
+
+# Helper: extract/write 16-bit half from/to 32-bit value
+def _src16(raw: int, is_hi: bool) -> int: return ((raw >> 16) & 0xffff) if is_hi else (raw & 0xffff)
+def _dst16(cur: int, val: int, is_hi: bool) -> int: return (cur & 0x0000ffff) | ((val & 0xffff) << 16) if is_hi else (cur & 0xffff0000) | (val & 0xffff)
+def _vgpr_hi(src: int) -> bool: return src >= 256 and ((src - 256) & 0x80) != 0
+def _vgpr_masked(src: int) -> int: return ((src - 256) & 0x7f) + 256 if src >= 256 else src
+
+# Memory access
+_valid_mem_ranges: list[tuple[int, int]] = []
+def set_valid_mem_ranges(ranges: set[tuple[int, int]]) -> None: _valid_mem_ranges.clear(); _valid_mem_ranges.extend(ranges)
+def _mem_valid(addr: int, size: int) -> bool:
+  return not _valid_mem_ranges or any(s <= addr and addr + size <= s + z for s, z in _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 (not pseudocode - these are format descriptions)
+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, FLAT_STORE = _mem_ops([GLOBALOp, FLATOp], _LOAD_MAP), _mem_ops([GLOBALOp, FLATOp], _STORE_MAP)
+# D16 ops: load/store 16-bit to lower or upper half of VGPR. Format: (size, sign, hi) where hi=1 means upper 16 bits
+_D16_LOAD_MAP = {'LOAD_D16_U8': (1,0,0), 'LOAD_D16_I8': (1,1,0), 'LOAD_D16_B16': (2,0,0),
+                 'LOAD_D16_HI_U8': (1,0,1), 'LOAD_D16_HI_I8': (1,1,1), 'LOAD_D16_HI_B16': (2,0,1)}
+_D16_STORE_MAP = {'STORE_D16_HI_B8': (1,1), 'STORE_D16_HI_B16': (2,1)}  # (size, hi)
+FLAT_D16_LOAD = _mem_ops([GLOBALOp, FLATOp], _D16_LOAD_MAP)
+FLAT_D16_STORE = _mem_ops([GLOBALOp, FLATOp], _D16_STORE_MAP)
+DS_LOAD = {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 = {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)}
+# 2ADDR ops: load/store two values using offset0 and offset1
+DS_LOAD_2ADDR = {DSOp.DS_LOAD_2ADDR_B32: 4, DSOp.DS_LOAD_2ADDR_B64: 8}
+DS_STORE_2ADDR = {DSOp.DS_STORE_2ADDR_B32: 4, DSOp.DS_STORE_2ADDR_B64: 8}
+SMEM_LOAD = {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}
+
+# VOPD op -> VOP3 op mapping (VOPD is dual-issue of VOP1/VOP2 ops, use VOP3 enums for pseudocode lookup)
+_VOPD_TO_VOP = {
+  VOPDOp.V_DUAL_FMAC_F32: VOP3Op.V_FMAC_F32, VOPDOp.V_DUAL_FMAAK_F32: VOP2Op.V_FMAAK_F32, VOPDOp.V_DUAL_FMAMK_F32: VOP2Op.V_FMAMK_F32,
+  VOPDOp.V_DUAL_MUL_F32: VOP3Op.V_MUL_F32, VOPDOp.V_DUAL_ADD_F32: VOP3Op.V_ADD_F32, VOPDOp.V_DUAL_SUB_F32: VOP3Op.V_SUB_F32,
+  VOPDOp.V_DUAL_SUBREV_F32: VOP3Op.V_SUBREV_F32, VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: VOP3Op.V_MUL_DX9_ZERO_F32,
+  VOPDOp.V_DUAL_MOV_B32: VOP3Op.V_MOV_B32, VOPDOp.V_DUAL_CNDMASK_B32: VOP3Op.V_CNDMASK_B32,
+  VOPDOp.V_DUAL_MAX_F32: VOP3Op.V_MAX_F32, VOPDOp.V_DUAL_MIN_F32: VOP3Op.V_MIN_F32,
+  VOPDOp.V_DUAL_ADD_NC_U32: VOP3Op.V_ADD_NC_U32, VOPDOp.V_DUAL_LSHLREV_B32: VOP3Op.V_LSHLREV_B32, VOPDOp.V_DUAL_AND_B32: VOP3Op.V_AND_B32,
+}
+
+# Compiled pseudocode functions (lazy loaded)
+_COMPILED: dict | None = None
+
+def _get_compiled() -> dict:
+  global _COMPILED
+  if _COMPILED is None: _COMPILED = get_compiled_functions()
+  return _COMPILED
+
+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], self.scc, self.pc, self.literal, self._pend_sgpr = 0xffffffff, 0, 0, 0, {}
+
+  @property
+  def vcc(self) -> int: return self.sgpr[VCC_LO] | (self.sgpr[VCC_HI] << 32)
+  @vcc.setter
+  def vcc(self, v: int): self.sgpr[VCC_LO], self.sgpr[VCC_HI] = v & MASK32, (v >> 32) & MASK32
+  @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): self.sgpr[EXEC_LO], self.sgpr[EXEC_HI] = v & MASK32, (v >> 32) & MASK32
+
+  def rsgpr(self, i: int) -> int: return 0 if i == NULL else self.scc if i == SCC else self.sgpr[i] if i < SGPR_COUNT else 0
+  def wsgpr(self, i: int, v: int):
+    if i < SGPR_COUNT and i != NULL: self.sgpr[i] = v & MASK32
+  def rsgpr64(self, i: int) -> int: return self.rsgpr(i) | (self.rsgpr(i+1) << 32)
+  def wsgpr64(self, i: int, v: int): self.wsgpr(i, v & MASK32); self.wsgpr(i+1, (v >> 32) & MASK32)
+
+  def _rsrc_base(self, v: int, lane: int, consts):
+    if v < SGPR_COUNT: return self.sgpr[v]
+    if v == SCC: return self.scc
+    if v < 255: return consts[v - 128]
+    if v == 255: return self.literal
+    return self.vgpr[lane][v - 256] if v <= 511 else 0
+  def rsrc(self, v: int, lane: int) -> int: return self._rsrc_base(v, lane, _INLINE_CONSTS)
+  def rsrc_f16(self, v: int, lane: int) -> int: return self._rsrc_base(v, lane, _INLINE_CONSTS_F16)
+  def rsrc64(self, v: int, lane: int) -> int:
+    if 128 <= v < 255: return _INLINE_CONSTS_F64[v - 128]
+    if v == 255: return self.literal  # literal is already shifted in from_bytes for 64-bit ops
+    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):
+    if reg not in self._pend_sgpr: self._pend_sgpr[reg] = 0
+    if val: self._pend_sgpr[reg] |= (1 << lane)
+  def commit_pends(self):
+    for reg, val in self._pend_sgpr.items(): self.sgpr[reg] = val
+    self._pend_sgpr.clear()
+
+
+def decode_program(data: bytes) -> Program:
+  result: Program = {}
+  i = 0
+  while i < len(data):
+    try: inst_class = detect_format(data[i:])
+    except ValueError: break  # stop at invalid instruction (padding/metadata after code)
+    if inst_class is None: i += 4; continue
+    base_size = inst_class._size()
+    # Pass enough data for potential 64-bit literal (base + 8 bytes max)
+    inst = inst_class.from_bytes(data[i:i+base_size+8])
+    for name, val in inst._values.items():
+      if name != 'op': setattr(inst, name, unwrap(val))  # skip op to preserve property access
+    inst._words = inst.size() // 4
+    result[i // 4] = inst
+    i += inst._words * 4
+  return result
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# EXECUTION - All ALU ops use pseudocode from PDF
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def exec_scalar(st: WaveState, inst: Inst) -> int:
+  """Execute scalar instruction. Returns PC delta or negative for special cases."""
+  compiled = _get_compiled()
+
+  # SOPP: special cases for control flow that has no pseudocode
+  if isinstance(inst, SOPP):
+    if inst.op == SOPPOp.S_ENDPGM: return -1
+    if inst.op == SOPPOp.S_BARRIER: return -2
+
+  # SMEM: memory loads (not ALU)
+  if isinstance(inst, SMEM):
+    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) & MASK64, 4))
+    return 0
+
+  # Get op enum and lookup compiled function
+  if isinstance(inst, SOP1): ssrc0, sdst = inst.ssrc0, inst.sdst
+  elif isinstance(inst, SOP2): ssrc0, sdst = inst.ssrc0, inst.sdst
+  elif isinstance(inst, SOPC): ssrc0, sdst = inst.ssrc0, None
+  elif isinstance(inst, SOPK): ssrc0, sdst = inst.sdst, inst.sdst  # sdst is both src and dst
+  elif isinstance(inst, SOPP): ssrc0, sdst = None, None
+  else: raise NotImplementedError(f"Unknown scalar type {type(inst)}")
+
+  # SOPP has gaps in the opcode enum - treat unknown opcodes as no-ops
+  try: op = inst.op
+  except ValueError:
+    if isinstance(inst, SOPP): return 0
+    raise
+  fn = compiled.get(type(op), {}).get(op)
+  if fn is None:
+    # SOPP instructions without pseudocode (waits, hints, nops) are no-ops
+    if isinstance(inst, SOPP): return 0
+    raise NotImplementedError(f"{op.name} not in pseudocode")
+
+  # Build context - use inst methods to determine operand sizes
+  s0 = st.rsrc64(ssrc0, 0) if inst.is_src_64(0) else (st.rsrc(ssrc0, 0) if not isinstance(inst, (SOPK, SOPP)) else (st.rsgpr(inst.sdst) if isinstance(inst, SOPK) else 0))
+  s1 = st.rsrc64(inst.ssrc1, 0) if inst.is_src_64(1) else (st.rsrc(inst.ssrc1, 0) if isinstance(inst, (SOP2, SOPC)) else inst.simm16 if isinstance(inst, SOPK) else 0)
+  d0 = st.rsgpr64(sdst) if inst.dst_regs() == 2 and sdst is not None else (st.rsgpr(sdst) if sdst is not None else 0)
+  literal = inst.simm16 if isinstance(inst, (SOPK, SOPP)) else st.literal
+
+  # Create Reg objects for compiled function - mask VCC/EXEC to 32 bits for wave32
+  result = fn(Reg(s0), Reg(s1), None, Reg(d0), Reg(st.scc), Reg(st.vcc & MASK32), 0, Reg(st.exec_mask & MASK32), literal, None, PC=Reg(st.pc * 4))
+
+  # Apply results - extract values from returned Reg objects
+  if sdst is not None and 'D0' in result:
+    (st.wsgpr64 if inst.dst_regs() == 2 else st.wsgpr)(sdst, result['D0']._val)
+  if 'SCC' in result: st.scc = result['SCC']._val & 1
+  if 'EXEC' in result: st.exec_mask = result['EXEC']._val
+  if 'PC' in result:
+    # Convert absolute byte address to word delta
+    pc_val = result['PC']._val
+    new_pc = pc_val if pc_val < 0x8000000000000000 else pc_val - 0x10000000000000000
+    new_pc_words = new_pc // 4
+    return new_pc_words - st.pc - 1  # -1 because emulator adds inst_words (1 for scalar)
+  return 0
+
+def exec_vector(st: WaveState, inst: Inst, lane: int, lds: bytearray | None = None) -> None:
+  """Execute vector instruction for one lane."""
+  compiled = _get_compiled()
+  V = st.vgpr[lane]
+
+  # Memory ops (not ALU pseudocode)
+  if isinstance(inst, FLAT):
+    op, addr_reg, data_reg, vdst, offset, saddr = inst.op, inst.addr, inst.data, inst.vdst, _sext(inst.offset, 13), inst.saddr
+    addr = V[addr_reg] | (V[addr_reg+1] << 32)
+    addr = (st.rsgpr64(saddr) + V[addr_reg] + offset) & MASK64 if saddr not in (NULL, 0x7f) else (addr + offset) & MASK64
+    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) & MASK32 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_LOAD:
+      sz, sign, hi = FLAT_D16_LOAD[op]
+      val = mem_read(addr, sz)
+      if sign: val = _sext(val, sz * 8) & 0xffff
+      V[vdst] = _dst16(V[vdst], val, hi)
+    elif op in FLAT_D16_STORE:
+      sz, hi = FLAT_D16_STORE[op]
+      mem_write(addr, sz, _src16(V[data_reg], hi) & ((1 << (sz * 8)) - 1))
+    else: raise NotImplementedError(f"FLAT op {op}")
+    return
+
+  if isinstance(inst, DS):
+    op, addr0, vdst = inst.op, (V[inst.addr] + inst.offset0) & 0xffff, inst.vdst
+    if op in DS_LOAD:
+      cnt, sz, sign = DS_LOAD[op]
+      for i in range(cnt): val = int.from_bytes(lds[addr0+i*sz:addr0+i*sz+sz], 'little'); V[vdst + i] = _sext(val, sz * 8) & MASK32 if sign else val
+    elif op in DS_STORE:
+      cnt, sz = DS_STORE[op]
+      for i in range(cnt): lds[addr0+i*sz:addr0+i*sz+sz] = (V[inst.data0 + i] & ((1 << (sz * 8)) - 1)).to_bytes(sz, 'little')
+    elif op in DS_LOAD_2ADDR:
+      # Load two values from addr+offset0*sz and addr+offset1*sz into vdst (B32: 1 dword each, B64: 2 dwords each)
+      # Note: offsets are scaled by data size (4 for B32, 8 for B64) per AMD ISA
+      sz = DS_LOAD_2ADDR[op]
+      addr0 = (V[inst.addr] + inst.offset0 * sz) & 0xffff
+      addr1 = (V[inst.addr] + inst.offset1 * sz) & 0xffff
+      cnt = sz // 4  # 1 for B32, 2 for B64
+      for i in range(cnt): V[vdst + i] = int.from_bytes(lds[addr0+i*4:addr0+i*4+4], 'little')
+      for i in range(cnt): V[vdst + cnt + i] = int.from_bytes(lds[addr1+i*4:addr1+i*4+4], 'little')
+    elif op in DS_STORE_2ADDR:
+      # Store two values from data0 and data1 to addr+offset0*sz and addr+offset1*sz
+      # Note: offsets are scaled by data size (4 for B32, 8 for B64) per AMD ISA
+      sz = DS_STORE_2ADDR[op]
+      addr0 = (V[inst.addr] + inst.offset0 * sz) & 0xffff
+      addr1 = (V[inst.addr] + inst.offset1 * sz) & 0xffff
+      cnt = sz // 4
+      for i in range(cnt): lds[addr0+i*4:addr0+i*4+4] = (V[inst.data0 + i] & MASK32).to_bytes(4, 'little')
+      for i in range(cnt): lds[addr1+i*4:addr1+i*4+4] = (V[inst.data1 + i] & MASK32).to_bytes(4, 'little')
+    else: raise NotImplementedError(f"DS op {op}")
+    return
+
+  # VOPD: dual-issue, execute two ops simultaneously (read all inputs before writes)
+  if isinstance(inst, VOPD):
+    vdsty = (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1)
+    inputs = [(inst.opx, st.rsrc(inst.srcx0, lane), V[inst.vsrcx1], V[inst.vdstx], inst.vdstx),
+              (inst.opy, st.rsrc(inst.srcy0, lane), V[inst.vsrcy1], V[vdsty], vdsty)]
+    def exec_vopd(vopd_op, s0, s1, d0):
+      op = _VOPD_TO_VOP[vopd_op]
+      return compiled[type(op)][op](Reg(s0), Reg(s1), None, Reg(d0), Reg(st.scc), Reg(st.vcc), lane, Reg(st.exec_mask), st.literal, None)['D0']._val
+    for vopd_op, s0, s1, d0, dst in inputs: V[dst] = exec_vopd(vopd_op, s0, s1, d0)
+    return
+
+  # VOP3SD: has extra scalar dest for carry output
+  if isinstance(inst, VOP3SD):
+    fn = compiled[VOP3SDOp][inst.op]
+    # Read sources based on register counts from inst properties
+    def rsrc_n(src, regs): return st.rsrc64(src, lane) if regs == 2 else st.rsrc(src, lane)
+    s0, s1, s2 = rsrc_n(inst.src0, inst.src_regs(0)), rsrc_n(inst.src1, inst.src_regs(1)), rsrc_n(inst.src2, inst.src_regs(2))
+    # Carry-in ops use src2 as carry bitmask instead of VCC
+    vcc = st.rsgpr64(inst.src2) if 'CO_CI' in inst.op_name else st.vcc
+    result = fn(Reg(s0), Reg(s1), Reg(s2), Reg(V[inst.vdst]), Reg(st.scc), Reg(vcc), lane, Reg(st.exec_mask), st.literal, None)
+    d0_val = result['D0']._val
+    V[inst.vdst] = d0_val & MASK32
+    if inst.dst_regs() == 2: V[inst.vdst + 1] = (d0_val >> 32) & MASK32
+    if 'VCC' in result: st.pend_sgpr_lane(inst.sdst, lane, (result['VCC']._val >> lane) & 1)
+    return
+
+  # Get op enum and sources (None means "no source" for that operand)
+  # dst_hi: for VOP1/VOP2 16-bit dst ops, bit 7 of vdst indicates .h (high 16-bit) destination
+  dst_hi = False
+  if isinstance(inst, VOP1):
+    if inst.op == VOP1Op.V_NOP: return
+    src0, src1, src2 = inst.src0, None, None
+    dst_hi = (inst.vdst & 0x80) != 0 and inst.is_dst_16()
+    vdst = inst.vdst & 0x7f if inst.is_dst_16() else inst.vdst
+  elif isinstance(inst, VOP2):
+    src0, src1, src2 = inst.src0, inst.vsrc1 + 256, None
+    dst_hi = (inst.vdst & 0x80) != 0 and inst.is_dst_16()
+    vdst = inst.vdst & 0x7f if inst.is_dst_16() else inst.vdst
+  elif isinstance(inst, VOP3):
+    # VOP3 ops 0-255 are VOPC comparisons encoded as VOP3 - inst.op returns VOPCOp for these
+    src0, src1, src2, vdst = inst.src0, inst.src1, (None if inst.op.value < 256 else inst.src2), inst.vdst
+  elif isinstance(inst, VOPC):
+    # For 16-bit VOPC, vsrc1 uses same encoding as VOP2 16-bit: bit 7 selects hi(1) or lo(0) half
+    # vsrc1 field is 8 bits: [6:0] = VGPR index, [7] = hi flag
+    src0, src1, src2, vdst = inst.src0, inst.vsrc1 + 256, None, VCC_LO
+  elif isinstance(inst, VOP3P):
+    # VOP3P: Packed 16-bit operations using compiled functions
+    # WMMA: wave-level matrix multiply-accumulate (special handling - needs cross-lane access)
+    if 'WMMA' in inst.op_name:
+      if lane == 0:  # Only execute once per wave, write results for all lanes
+        exec_wmma(st, inst, inst.op)
+      return
+    # V_FMA_MIX: Mixed precision FMA - opsel_hi controls f32(0) vs f16(1), opsel selects which f16 half
+    if 'FMA_MIX' in inst.op_name:
+      opsel, opsel_hi, opsel_hi2 = getattr(inst, 'opsel', 0), getattr(inst, 'opsel_hi', 0), getattr(inst, 'opsel_hi2', 0)
+      neg, abs_ = getattr(inst, 'neg', 0), getattr(inst, 'neg_hi', 0)  # neg_hi reused as abs
+      raws = [st.rsrc(inst.src0, lane), st.rsrc(inst.src1, lane), st.rsrc(inst.src2, lane) if inst.src2 is not None else 0]
+      is_f16 = [opsel_hi & 1, opsel_hi & 2, opsel_hi2]
+      srcs = [_f16(_src16(raws[i], bool(opsel & (1<<i)))) if is_f16[i] else _f32(raws[i]) for i in range(3)]
+      for i in range(3):
+        if abs_ & (1<<i): srcs[i] = abs(srcs[i])
+        if neg & (1<<i): srcs[i] = -srcs[i]
+      result = srcs[0] * srcs[1] + srcs[2]
+      st.vgpr[lane][inst.vdst] = _i32(result) if inst.op == VOP3POp.V_FMA_MIX_F32 else _dst16(V[inst.vdst], _i16(result), inst.op == VOP3POp.V_FMA_MIXHI_F16)
+      return
+    # VOP3P packed ops: opsel selects halves for lo, opsel_hi for hi; neg toggles f16 sign
+    raws = [st.rsrc_f16(inst.src0, lane), st.rsrc_f16(inst.src1, lane), st.rsrc_f16(inst.src2, lane) if inst.src2 is not None else 0]
+    opsel, opsel_hi, opsel_hi2 = getattr(inst, 'opsel', 0), getattr(inst, 'opsel_hi', 3), getattr(inst, 'opsel_hi2', 1)
+    neg, neg_hi = getattr(inst, 'neg', 0), getattr(inst, 'neg_hi', 0)
+    hi_sels = [opsel_hi & 1, opsel_hi & 2, opsel_hi2]
+    srcs = [((_src16(raws[i], hi_sels[i]) ^ (0x8000 if neg_hi & (1<<i) else 0)) << 16) |
+            (_src16(raws[i], opsel & (1<<i)) ^ (0x8000 if neg & (1<<i) else 0)) for i in range(3)]
+    result = compiled[VOP3POp][inst.op](Reg(srcs[0]), Reg(srcs[1]), Reg(srcs[2]), Reg(0), Reg(st.scc), Reg(st.vcc), lane, Reg(st.exec_mask), st.literal, None)
+    st.vgpr[lane][inst.vdst] = result['D0']._val & MASK32
+    return
+  else: raise NotImplementedError(f"Unknown vector type {type(inst)}")
+
+  op_cls = type(inst.op)
+  if (fn := compiled.get(op_cls, {}).get(inst.op)) is None: raise NotImplementedError(f"{inst.op_name} not in pseudocode")
+
+  # Read sources (with VOP3 modifiers if applicable)
+  neg, abs_ = (getattr(inst, 'neg', 0), getattr(inst, 'abs', 0)) if isinstance(inst, VOP3) else (0, 0)
+  opsel = getattr(inst, 'opsel', 0) if isinstance(inst, VOP3) else 0
+  def mod_src(val: int, idx: int, is64=False) -> int:
+    to_f, to_i = (_f64, _i64) if is64 else (_f32, _i32)
+    if (abs_ >> idx) & 1: val = to_i(abs(to_f(val)))
+    if (neg >> idx) & 1: val = to_i(-to_f(val))
+    return val
+
+  # Use inst methods to determine operand sizes (inst.is_src_16, inst.is_src_64, etc.)
+  is_vop2_16bit = isinstance(inst, VOP2) and inst.is_16bit()
+
+  # Read sources based on register counts and dtypes from inst properties
+  def read_src(src, idx, regs, is_src_16):
+    if src is None: return 0
+    if regs == 2: return mod_src(st.rsrc64(src, lane), idx, is64=True)
+    if is_src_16 and isinstance(inst, VOP3):
+      raw = st.rsrc_f16(src, lane) if 128 <= src < 255 else st.rsrc(src, lane)
+      val = _src16(raw, bool(opsel & (1 << idx)))
+      if abs_ & (1 << idx): val &= 0x7fff
+      if neg & (1 << idx): val ^= 0x8000
+      return val
+    if is_src_16 and isinstance(inst, (VOP1, VOP2, VOPC)):
+      if src >= 256: return _src16(mod_src(st.rsrc(_vgpr_masked(src), lane), idx), _vgpr_hi(src))
+      return mod_src(st.rsrc_f16(src, lane), idx) & 0xffff
+    return mod_src(st.rsrc(src, lane), idx)
+
+  s0 = read_src(src0, 0, inst.src_regs(0), inst.is_src_16(0))
+  s1 = read_src(src1, 1, inst.src_regs(1), inst.is_src_16(1)) if src1 is not None else 0
+  s2 = read_src(src2, 2, inst.src_regs(2), inst.is_src_16(2)) if src2 is not None else 0
+  # Read destination (accumulator for VOP2 f16, 64-bit for 64-bit ops)
+  d0 = _src16(V[vdst], dst_hi) if is_vop2_16bit else (V[vdst] | (V[vdst + 1] << 32)) if inst.dst_regs() == 2 else V[vdst]
+
+  # V_CNDMASK_B32/B16: VOP3 encoding uses src2 as mask (not VCC); VOP2 uses VCC implicitly
+  # Pass the correct mask as vcc to the function so pseudocode VCC.u64[laneId] works correctly
+  vcc_for_fn = st.rsgpr64(src2) if inst.op in (VOP3Op.V_CNDMASK_B32, VOP3Op.V_CNDMASK_B16) and isinstance(inst, VOP3) and src2 is not None and src2 < 256 else st.vcc
+
+  # Execute compiled function - pass src0_idx and vdst_idx for lane instructions
+  # For VGPR access: src0 index is the VGPR number (src0 - 256 if VGPR, else src0 for SGPR)
+  src0_idx = (src0 - 256) if src0 is not None and src0 >= 256 else (src0 if src0 is not None else 0)
+  result = fn(Reg(s0), Reg(s1), Reg(s2), Reg(d0), Reg(st.scc), Reg(vcc_for_fn), lane, Reg(st.exec_mask), st.literal, st.vgpr, src0_idx, vdst)
+
+  # Apply results - extract values from returned Reg objects
+  if 'vgpr_write' in result:
+    # Lane instruction wrote to VGPR: (lane, vgpr_idx, value)
+    wr_lane, wr_idx, wr_val = result['vgpr_write']
+    st.vgpr[wr_lane][wr_idx] = wr_val
+  if 'VCC' in result:
+    # VOP2 carry ops write to VCC implicitly; VOPC/VOP3 write to vdst
+    st.pend_sgpr_lane(VCC_LO if isinstance(inst, VOP2) and 'CO_CI' in inst.op_name else vdst, lane, (result['VCC']._val >> lane) & 1)
+  if 'EXEC' in result:
+    # V_CMPX instructions write to EXEC per-lane (not to vdst)
+    st.pend_sgpr_lane(EXEC_LO, lane, (result['EXEC']._val >> lane) & 1)
+  elif op_cls is VOPCOp:
+    # VOPC comparison result stored in D0 bitmask, extract lane bit (non-CMPX only)
+    st.pend_sgpr_lane(vdst, lane, (result['D0']._val >> lane) & 1)
+  if op_cls is not VOPCOp and 'vgpr_write' not in result:
+    writes_to_sgpr = 'READFIRSTLANE' in inst.op_name or 'READLANE' in inst.op_name
+    d0_val = result['D0']._val
+    if writes_to_sgpr: st.wsgpr(vdst, d0_val & MASK32)
+    elif inst.dst_regs() == 2: V[vdst], V[vdst + 1] = d0_val & MASK32, (d0_val >> 32) & MASK32
+    elif inst.is_dst_16(): V[vdst] = _dst16(V[vdst], d0_val, bool(opsel & 8) if isinstance(inst, VOP3) else dst_hi)
+    else: V[vdst] = d0_val & MASK32
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# WMMA (Wave Matrix Multiply-Accumulate)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def exec_wmma(st: WaveState, inst, op: VOP3POp) -> None:
+  """Execute WMMA instruction - 16x16x16 matrix multiply across the wave."""
+  src0, src1, src2, vdst = inst.src0, inst.src1, inst.src2, inst.vdst
+  # Read 16x16 f16 matrix from 16 lanes × 8 VGPRs (2 f16 per VGPR)
+  def read_f16_mat(src):
+    return [f for l in range(16) for r in range(8) for v in [st.vgpr[l][src-256+r] if src >= 256 else st.rsgpr(src+r)] for f in [_f16(v&0xffff), _f16((v>>16)&0xffff)]]
+  mat_a, mat_b = read_f16_mat(src0), read_f16_mat(src1)
+  # Read matrix C (16x16 f32) from lanes 0-31, VGPRs src2 to src2+7
+  mat_c = [_f32(st.vgpr[i % 32][src2 - 256 + i // 32] if src2 >= 256 else st.rsgpr(src2 + i // 32)) for i in range(256)]
+  # Compute D = A × B + C (16x16 matrix multiply)
+  mat_d = [sum(mat_a[row*16+k] * mat_b[col*16+k] for k in range(16)) + mat_c[row*16+col] for row in range(16) for col in range(16)]
+  # Write result - f16 packed or f32
+  if op == VOP3POp.V_WMMA_F16_16X16X16_F16:
+    for i in range(0, 256, 2):
+      st.vgpr[(i//2) % 32][vdst + (i//2)//32] = ((_i16(mat_d[i+1]) & 0xffff) << 16) | (_i16(mat_d[i]) & 0xffff)
+  else:
+    for i in range(256): st.vgpr[i % 32][vdst + i//32] = _i32(mat_d[i])
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# MAIN EXECUTION LOOP
+# ═══════════════════════════════════════════════════════════════════════════════
+
+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._words, getattr(inst, '_literal', None) or 0
+
+  if isinstance(inst, (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM)):
+    delta = exec_scalar(st, inst)
+    if delta == -1: return -1  # endpgm
+    if delta == -2: st.pc += inst_words; return -2  # barrier
+    st.pc += inst_words + delta
+  else:
+    # V_READFIRSTLANE/V_READLANE write to SGPR, execute once; others execute per-lane with exec_mask
+    is_readlane = isinstance(inst, (VOP1, VOP3)) and ('READFIRSTLANE' in inst.op_name or 'READLANE' in inst.op_name)
+    exec_mask = 1 if is_readlane else st.exec_mask
+    for lane in range(1 if is_readlane else n_lanes):
+      if exec_mask & (1 << lane): exec_vector(st, inst, lane, lds)
+    st.commit_pends()
+    st.pc += inst_words
+  return 0
+
+def exec_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int) -> 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,
+                   wg_id_sgpr_base: int, wg_id_enables: tuple[bool, bool, bool]) -> 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)
+    # Set workgroup IDs in SGPRs based on USER_SGPR_COUNT and enable flags from COMPUTE_PGM_RSRC2
+    sgpr_idx = wg_id_sgpr_base
+    for wg_id, enabled in zip(workgroup_id, wg_id_enables):
+      if enabled: st.sgpr[sgpr_idx] = wg_id; sgpr_idx += 1
+    # Set workitem IDs in VGPR0 using packed method: v0 = (Z << 20) | (Y << 10) | X
+    for i in range(n_lanes):
+      tid = wave_start + i
+      st.vgpr[i][0] = ((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, _ in waves: exec_wave(program, st, lds, n_lanes)
+
+def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int, rsrc2: int = 0x19c) -> int:
+  program = decode_program((ctypes.c_char * lib_sz).from_address(lib).raw)
+  if not program: return -1
+  wg_id_enables = tuple(bool((rsrc2 >> (7+i)) & 1) for i in range(3))
+  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, (rsrc2 >> 1) & 0x1f, wg_id_enables)
+  return 0

extra/assembly/amd/pcode.py

@@ -0,0 +1,542 @@
+# DSL for RDNA3 pseudocode - makes pseudocode expressions work directly as Python
+import struct, math
+from extra.assembly.amd.dsl import MASK32, MASK64, _f32, _i32, _sext, _f16, _i16, _f64, _i64
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# HELPER FUNCTIONS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def _div(a, b):
+  try: return a / b
+  except ZeroDivisionError:
+    if a == 0.0 or math.isnan(a): return float("nan")
+    return math.copysign(float("inf"), a * b) if b == 0.0 else float("inf") if a > 0 else float("-inf")
+def _to_f16_bits(v): return v if isinstance(v, int) else _i16(v)
+def _isnan(x):
+  try: return math.isnan(float(x))
+  except (TypeError, ValueError): return False
+def _check_nan_type(x, quiet_bit_expected, default):
+  """Check NaN type by examining quiet bit. Returns default if can't determine."""
+  try:
+    if not math.isnan(float(x)): return False
+    if hasattr(x, '_reg') and hasattr(x, '_bits'):
+      bits = x._reg._val & ((1 << x._bits) - 1)
+      # NaN format: exponent all 1s, quiet bit, mantissa != 0
+      # f16: exp[14:10]=31, quiet=bit9, mant[8:0]  |  f32: exp[30:23]=255, quiet=bit22, mant[22:0]  |  f64: exp[62:52]=2047, quiet=bit51, mant[51:0]
+      exp_bits, quiet_pos, mant_mask = {16: (0x1f, 9, 0x3ff), 32: (0xff, 22, 0x7fffff), 64: (0x7ff, 51, 0xfffffffffffff)}.get(x._bits, (0,0,0))
+      exp_shift = {16: 10, 32: 23, 64: 52}.get(x._bits, 0)
+      if exp_bits and ((bits >> exp_shift) & exp_bits) == exp_bits and (bits & mant_mask) != 0:
+        return ((bits >> quiet_pos) & 1) == quiet_bit_expected
+    return default
+  except (TypeError, ValueError): return False
+def _isquietnan(x): return _check_nan_type(x, 1, True)  # quiet NaN has quiet bit = 1
+def _issignalnan(x): return _check_nan_type(x, 0, False)  # signaling NaN has quiet bit = 0
+def _gt_neg_zero(a, b): return (a > b) or (a == 0 and b == 0 and not math.copysign(1, a) < 0 and math.copysign(1, b) < 0)
+def _lt_neg_zero(a, b): return (a < b) or (a == 0 and b == 0 and math.copysign(1, a) < 0 and not math.copysign(1, b) < 0)
+def _fma(a, b, c): return a * b + c
+def _signext(v): return v
+def _fpop(fn): return lambda x: (x := float(x), x if math.isnan(x) or math.isinf(x) else float(fn(x)))[1]
+trunc, floor, ceil = _fpop(math.trunc), _fpop(math.floor), _fpop(math.ceil)
+class _SafeFloat(float):
+  """Float subclass that uses _div for division to handle 0/inf correctly."""
+  def __truediv__(self, o): return _div(float(self), float(o))
+  def __rtruediv__(self, o): return _div(float(o), float(self))
+def sqrt(x): return _SafeFloat(math.sqrt(x)) if x >= 0 else _SafeFloat(float("nan"))
+def log2(x): return math.log2(x) if x > 0 else (float("-inf") if x == 0 else float("nan"))
+i32_to_f32 = u32_to_f32 = i32_to_f64 = u32_to_f64 = f32_to_f64 = f64_to_f32 = float
+def _f_to_int(f, lo, hi): f = float(f); return 0 if math.isnan(f) else (hi if f >= hi else lo if f <= lo else int(f))
+def f32_to_i32(f): return _f_to_int(f, -2147483648, 2147483647)
+def f32_to_u32(f): return _f_to_int(f, 0, 4294967295)
+f64_to_i32, f64_to_u32 = f32_to_i32, f32_to_u32
+def f32_to_f16(f):
+  f = float(f)
+  if math.isnan(f): return 0x7e00  # f16 NaN
+  if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00  # f16 ±infinity
+  try: return struct.unpack("<H", struct.pack("<e", f))[0]
+  except OverflowError: return 0x7c00 if f > 0 else 0xfc00  # overflow -> ±infinity
+def _f16_to_f32_bits(bits): return struct.unpack("<e", struct.pack("<H", int(bits) & 0xffff))[0]
+def f16_to_f32(v): return v if isinstance(v, float) else _f16_to_f32_bits(v)
+def i16_to_f16(v): return f32_to_f16(float(_sext(int(v) & 0xffff, 16)))
+def u16_to_f16(v): return f32_to_f16(float(int(v) & 0xffff))
+def f16_to_i16(bits): f = _f16_to_f32_bits(bits); return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
+def f16_to_u16(bits): f = _f16_to_f32_bits(bits); return max(0, min(65535, int(f))) if not math.isnan(f) else 0
+def u8_to_u32(v): return int(v) & 0xff
+def u4_to_u32(v): return int(v) & 0xf
+def _sign(f): return 1 if math.copysign(1.0, f) < 0 else 0
+def _mantissa_f32(f): return struct.unpack("<I", struct.pack("<f", f))[0] & 0x7fffff if not (math.isinf(f) or math.isnan(f)) else 0
+def _ldexp(m, e): return math.ldexp(m, e)
+def isEven(x):
+  x = float(x)
+  if math.isinf(x) or math.isnan(x): return False
+  return int(x) % 2 == 0
+def fract(x): return x - math.floor(x)
+PI = math.pi
+def _trig(fn, x):
+  # V_SIN/COS_F32: hardware does frac on input cycles before computing
+  if math.isinf(x) or math.isnan(x): return float("nan")
+  frac_cycles = fract(x / (2 * math.pi))
+  return fn(frac_cycles * 2 * math.pi)
+def sin(x): return _trig(math.sin, x)
+def cos(x): return _trig(math.cos, x)
+def pow(a, b):
+  try: return a ** b
+  except OverflowError: return float("inf") if b > 0 else 0.0
+def _brev(v, bits): return int(bin(v & ((1 << bits) - 1))[2:].zfill(bits)[::-1], 2)
+def _brev32(v): return _brev(v, 32)
+def _brev64(v): return _brev(v, 64)
+def _ctz(v, bits):
+  v, n = int(v) & ((1 << bits) - 1), 0
+  if v == 0: return bits
+  while (v & 1) == 0: v >>= 1; n += 1
+  return n
+def _ctz32(v): return _ctz(v, 32)
+def _ctz64(v): return _ctz(v, 64)
+def _exponent(f):
+  # Handle TypedView (f16/f32/f64) to get correct exponent for that type
+  if hasattr(f, '_bits') and hasattr(f, '_float') and f._float:
+    raw = f._val
+    if f._bits == 16: return (raw >> 10) & 0x1f  # f16: 5-bit exponent
+    if f._bits == 32: return (raw >> 23) & 0xff  # f32: 8-bit exponent
+    if f._bits == 64: return (raw >> 52) & 0x7ff  # f64: 11-bit exponent
+  # Fallback: convert to f32 and get exponent
+  f = float(f)
+  if math.isinf(f) or math.isnan(f): return 255
+  if f == 0.0: return 0
+  try: bits = struct.unpack("<I", struct.pack("<f", f))[0]; return (bits >> 23) & 0xff
+  except: return 0
+def _is_denorm_f32(f):
+  if not isinstance(f, float): f = _f32(int(f) & 0xffffffff)
+  if math.isinf(f) or math.isnan(f) or f == 0.0: return False
+  bits = struct.unpack("<I", struct.pack("<f", float(f)))[0]
+  return (bits >> 23) & 0xff == 0
+def _is_denorm_f64(f):
+  if not isinstance(f, float): f = _f64(int(f) & 0xffffffffffffffff)
+  if math.isinf(f) or math.isnan(f) or f == 0.0: return False
+  bits = struct.unpack("<Q", struct.pack("<d", float(f)))[0]
+  return (bits >> 52) & 0x7ff == 0
+def v_min_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _lt_neg_zero(a, b) else b)
+def v_max_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _gt_neg_zero(a, b) else b)
+v_min_f16, v_max_f16 = v_min_f32, v_max_f32
+v_min_i32, v_max_i32 = min, max
+v_min_i16, v_max_i16 = min, max
+def v_min_u32(a, b): return min(a & MASK32, b & MASK32)
+def v_max_u32(a, b): return max(a & MASK32, b & MASK32)
+def v_min_u16(a, b): return min(a & 0xffff, b & 0xffff)
+def v_max_u16(a, b): return max(a & 0xffff, b & 0xffff)
+def v_min3_f32(a, b, c): return v_min_f32(v_min_f32(a, b), c)
+def v_max3_f32(a, b, c): return v_max_f32(v_max_f32(a, b), c)
+v_min3_f16, v_max3_f16 = v_min3_f32, v_max3_f32
+v_min3_i32, v_max3_i32, v_min3_i16, v_max3_i16 = min, max, min, max
+def v_min3_u32(a, b, c): return min(a & MASK32, b & MASK32, c & MASK32)
+def v_max3_u32(a, b, c): return max(a & MASK32, b & MASK32, c & MASK32)
+def v_min3_u16(a, b, c): return min(a & 0xffff, b & 0xffff, c & 0xffff)
+def v_max3_u16(a, b, c): return max(a & 0xffff, b & 0xffff, c & 0xffff)
+def ABSDIFF(a, b): return abs(int(a) - int(b))
+
+# BF16 (bfloat16) conversion functions
+def _bf16(i):
+  """Convert bf16 bits to float. BF16 is just the top 16 bits of f32."""
+  return struct.unpack("<f", struct.pack("<I", (i & 0xffff) << 16))[0]
+def _ibf16(f):
+  """Convert float to bf16 bits (truncate to top 16 bits of f32)."""
+  if math.isnan(f): return 0x7fc0  # bf16 quiet NaN
+  if math.isinf(f): return 0x7f80 if f > 0 else 0xff80  # bf16 ±infinity
+  try: return (struct.unpack("<I", struct.pack("<f", float(f)))[0] >> 16) & 0xffff
+  except (OverflowError, struct.error): return 0x7f80 if f > 0 else 0xff80
+def bf16_to_f32(v): return _bf16(v) if isinstance(v, int) else float(v)
+def f32_to_bf16(f): return _ibf16(f)
+
+# BYTE_PERMUTE for V_PERM_B32 - select bytes from 64-bit data based on selector
+def BYTE_PERMUTE(data, sel):
+  """Select a byte from 64-bit data based on selector value.
+  sel 0-7: select byte from data (S1 is bytes 0-3, S0 is bytes 4-7 in {S0,S1})
+  sel 8-11: sign-extend from specific bytes (8->byte1, 9->byte3, 10->byte5, 11->byte7)
+  sel 12: constant 0x00
+  sel >= 13: constant 0xFF"""
+  sel = int(sel) & 0xff
+  if sel <= 7: return (int(data) >> (sel * 8)) & 0xff
+  if sel == 8: return 0xff if ((int(data) >> 15) & 1) else 0x00  # sign of byte 1
+  if sel == 9: return 0xff if ((int(data) >> 31) & 1) else 0x00  # sign of byte 3
+  if sel == 10: return 0xff if ((int(data) >> 47) & 1) else 0x00  # sign of byte 5
+  if sel == 11: return 0xff if ((int(data) >> 63) & 1) else 0x00  # sign of byte 7
+  if sel == 12: return 0x00
+  return 0xff  # sel >= 13
+
+# v_sad_u8 helper for V_SAD instructions (sum of absolute differences of 4 bytes)
+def v_sad_u8(s0, s1, s2):
+  """V_SAD_U8: Sum of absolute differences of 4 byte pairs plus accumulator."""
+  s0, s1, s2 = int(s0), int(s1), int(s2)
+  result = s2
+  for i in range(4):
+    a = (s0 >> (i * 8)) & 0xff
+    b = (s1 >> (i * 8)) & 0xff
+    result += abs(a - b)
+  return result & 0xffffffff
+
+# v_msad_u8 helper (masked SAD - skip when reference byte is 0)
+def v_msad_u8(s0, s1, s2):
+  """V_MSAD_U8: Masked sum of absolute differences (skip if reference byte is 0)."""
+  s0, s1, s2 = int(s0), int(s1), int(s2)
+  result = s2
+  for i in range(4):
+    a = (s0 >> (i * 8)) & 0xff
+    b = (s1 >> (i * 8)) & 0xff
+    if b != 0:  # Only add diff if reference (s1) byte is non-zero
+      result += abs(a - b)
+  return result & 0xffffffff
+def f16_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
+def f16_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
+def f32_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
+def f32_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
+def v_cvt_i16_f32(f): return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
+def v_cvt_u16_f32(f): return max(0, min(65535, int(f))) if not math.isnan(f) else 0
+def u32_to_u16(u): return int(u) & 0xffff
+def i32_to_i16(i): return ((int(i) + 32768) & 0xffff) - 32768
+def SAT8(v): return max(0, min(255, int(v)))
+def f32_to_u8(f): return max(0, min(255, int(f))) if not math.isnan(f) else 0
+def mantissa(f):
+  if f == 0.0 or math.isinf(f) or math.isnan(f): return f
+  m, _ = math.frexp(f)
+  return m  # AMD V_FREXP_MANT returns mantissa in [0.5, 1.0) range
+def signext_from_bit(val, bit):
+  bit = int(bit)
+  if bit == 0: return 0
+  mask = (1 << bit) - 1
+  val = int(val) & mask
+  if val & (1 << (bit - 1)): return val - (1 << bit)
+  return val
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# DSL EXPORTS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+__all__ = [
+  # Classes
+  'Reg', 'SliceProxy', 'TypedView',
+  # Pack functions
+  '_pack', '_pack32', 'pack', 'pack32',
+  # Constants
+  'WAVE32', 'WAVE64', 'MASK32', 'MASK64', 'WAVE_MODE', 'DENORM', 'OVERFLOW_F32', 'UNDERFLOW_F32',
+  'OVERFLOW_F64', 'UNDERFLOW_F64', 'MAX_FLOAT_F32', 'ROUND_MODE', 'cvtToQuietNAN', 'DST', 'INF', 'PI',
+  'TWO_OVER_PI_1201',
+  # Aliases for pseudocode
+  's_ff1_i32_b32', 's_ff1_i32_b64', 'GT_NEG_ZERO', 'LT_NEG_ZERO',
+  'isNAN', 'isQuietNAN', 'isSignalNAN', 'fma', 'ldexp', 'sign', 'exponent', 'F', 'signext',
+  # Conversion functions
+  '_f32', '_i32', '_f16', '_i16', '_f64', '_i64', '_sext', '_to_f16_bits', '_f16_to_f32_bits',
+  'i32_to_f32', 'u32_to_f32', 'i32_to_f64', 'u32_to_f64', 'f32_to_f64', 'f64_to_f32',
+  'f32_to_i32', 'f32_to_u32', 'f64_to_i32', 'f64_to_u32', 'f32_to_f16', 'f16_to_f32',
+  'i16_to_f16', 'u16_to_f16', 'f16_to_i16', 'f16_to_u16', 'u32_to_u16', 'i32_to_i16',
+  'f16_to_snorm', 'f16_to_unorm', 'f32_to_snorm', 'f32_to_unorm', 'v_cvt_i16_f32', 'v_cvt_u16_f32',
+  'SAT8', 'f32_to_u8', 'u8_to_u32', 'u4_to_u32',
+  # BF16 conversion functions
+  '_bf16', '_ibf16', 'bf16_to_f32', 'f32_to_bf16',
+  # Math functions
+  'trunc', 'floor', 'ceil', 'sqrt', 'log2', 'sin', 'cos', 'pow', 'fract', 'isEven', 'mantissa',
+  # Min/max functions
+  'v_min_f32', 'v_max_f32', 'v_min_i32', 'v_max_i32', 'v_min_u32', 'v_max_u32',
+  'v_min_f16', 'v_max_f16', 'v_min_i16', 'v_max_i16', 'v_min_u16', 'v_max_u16',
+  'v_min3_f32', 'v_max3_f32', 'v_min3_i32', 'v_max3_i32', 'v_min3_u32', 'v_max3_u32',
+  'v_min3_f16', 'v_max3_f16', 'v_min3_i16', 'v_max3_i16', 'v_min3_u16', 'v_max3_u16',
+  'ABSDIFF',
+  # Byte/SAD helper functions
+  'BYTE_PERMUTE', 'v_sad_u8', 'v_msad_u8',
+  # Bit manipulation
+  '_brev32', '_brev64', '_ctz32', '_ctz64', '_exponent', '_is_denorm_f32', '_is_denorm_f64',
+  '_sign', '_mantissa_f32', '_div', '_isnan', '_isquietnan', '_issignalnan', '_gt_neg_zero', '_lt_neg_zero', '_fma', '_ldexp', '_signext',
+  'signext_from_bit',
+]
+
+# Aliases used in pseudocode
+s_ff1_i32_b32, s_ff1_i32_b64 = _ctz32, _ctz64
+GT_NEG_ZERO, LT_NEG_ZERO = _gt_neg_zero, _lt_neg_zero
+isNAN = _isnan
+isQuietNAN = _isquietnan
+isSignalNAN = _issignalnan
+fma, ldexp, sign, exponent = _fma, _ldexp, _sign, _exponent
+def F(x):
+  """32'F(x) or 64'F(x) - interpret x as float. If x is int, treat as bit pattern."""
+  if isinstance(x, int): return _f32(x)  # int -> interpret as f32 bits
+  if isinstance(x, TypedView): return x  # preserve TypedView for bit-pattern checks
+  return float(x)  # already a float or float-like
+signext = lambda x: int(x)  # sign-extend to full width - already handled by Python's arbitrary precision ints
+pack = lambda hi, lo: ((int(hi) & 0xffff) << 16) | (int(lo) & 0xffff)
+pack32 = lambda hi, lo: ((int(hi) & 0xffffffff) << 32) | (int(lo) & 0xffffffff)
+_pack, _pack32 = pack, pack32  # Aliases for internal use
+WAVE32, WAVE64 = True, False
+
+# Float overflow/underflow constants
+OVERFLOW_F32 = float('inf')
+UNDERFLOW_F32 = 0.0
+OVERFLOW_F64 = float('inf')
+UNDERFLOW_F64 = 0.0
+MAX_FLOAT_F32 = 3.4028235e+38  # Largest finite float32
+
+# INF object that supports .f16/.f32/.f64 access and comparison with floats
+class _Inf:
+  f16 = f32 = f64 = float('inf')
+  def __neg__(self): return _NegInf()
+  def __pos__(self): return self
+  def __float__(self): return float('inf')
+  def __eq__(self, other): return float(other) == float('inf') if not isinstance(other, _NegInf) else False
+  def __req__(self, other): return self.__eq__(other)
+class _NegInf:
+  f16 = f32 = f64 = float('-inf')
+  def __neg__(self): return _Inf()
+  def __pos__(self): return self
+  def __float__(self): return float('-inf')
+  def __eq__(self, other): return float(other) == float('-inf') if not isinstance(other, _Inf) else False
+  def __req__(self, other): return self.__eq__(other)
+INF = _Inf()
+
+# Rounding mode placeholder
+class _RoundMode:
+  NEAREST_EVEN = 0
+ROUND_MODE = _RoundMode()
+
+# Helper functions for pseudocode
+def cvtToQuietNAN(x): return float('nan')
+DST = None  # Placeholder, will be set in context
+
+# 2/PI with 1201 bits of precision for V_TRIG_PREOP_F64
+# Computed as: int((2/pi) * 2^1201) - this is the fractional part of 2/pi scaled to integer
+# The MSB (bit 1200) corresponds to 2^0 position in the fraction 0.b1200 b1199 ... b1 b0
+_TWO_OVER_PI_1201_RAW = 0x0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6
+
+class _BigInt:
+  """Wrapper for large integers that supports bit slicing [high:low]."""
+  __slots__ = ('_val',)
+  def __init__(self, val): self._val = val
+  def __getitem__(self, key):
+    if isinstance(key, slice):
+      high, low = key.start, key.stop
+      if high < low: high, low = low, high  # Handle reversed slice
+      mask = (1 << (high - low + 1)) - 1
+      return (self._val >> low) & mask
+    return (self._val >> key) & 1
+  def __int__(self): return self._val
+  def __index__(self): return self._val
+  def __lshift__(self, n): return self._val << int(n)
+  def __rshift__(self, n): return self._val >> int(n)
+  def __and__(self, n): return self._val & int(n)
+  def __or__(self, n): return self._val | int(n)
+
+TWO_OVER_PI_1201 = _BigInt(_TWO_OVER_PI_1201_RAW)
+
+class _WaveMode:
+  IEEE = False
+WAVE_MODE = _WaveMode()
+
+class _DenormChecker:
+  """Comparator for denormalized floats. x == DENORM.f32 checks if x is denormalized."""
+  def __init__(self, bits): self._bits = bits
+  def _check(self, other):
+    return _is_denorm_f64(float(other)) if self._bits == 64 else _is_denorm_f32(float(other))
+  def __eq__(self, other): return self._check(other)
+  def __req__(self, other): return self._check(other)
+  def __ne__(self, other): return not self._check(other)
+
+class _Denorm:
+  f32 = _DenormChecker(32)
+  f64 = _DenormChecker(64)
+DENORM = _Denorm()
+
+def _brev(v, bits):
+  """Bit-reverse a value."""
+  result = 0
+  for i in range(bits): result |= ((v >> i) & 1) << (bits - 1 - i)
+  return result
+
+class SliceProxy:
+  """Proxy for D0[31:16] that supports .f16/.u16 etc getters and setters."""
+  __slots__ = ('_reg', '_high', '_low', '_reversed')
+  def __init__(self, reg, high, low):
+    self._reg = reg
+    # Handle reversed slices like [0:31] which means bit-reverse
+    if high < low: self._high, self._low, self._reversed = low, high, True
+    else: self._high, self._low, self._reversed = high, low, False
+  def _nbits(self): return self._high - self._low + 1
+  def _mask(self): return (1 << self._nbits()) - 1
+  def _get(self):
+    v = (self._reg._val >> self._low) & self._mask()
+    return _brev(v, self._nbits()) if self._reversed else v
+  def _set(self, v):
+    v = int(v)
+    if self._reversed: v = _brev(v, self._nbits())
+    self._reg._val = (self._reg._val & ~(self._mask() << self._low)) | ((v & self._mask()) << self._low)
+
+  u8 = property(lambda s: s._get() & 0xff)
+  u16 = property(lambda s: s._get() & 0xffff, lambda s, v: s._set(v))
+  u32 = property(lambda s: s._get() & MASK32, lambda s, v: s._set(v))
+  i16 = property(lambda s: _sext(s._get() & 0xffff, 16), lambda s, v: s._set(v))
+  i32 = property(lambda s: _sext(s._get() & MASK32, 32), lambda s, v: s._set(v))
+  f16 = property(lambda s: _f16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _i16(float(v))))
+  f32 = property(lambda s: _f32(s._get()), lambda s, v: s._set(_i32(float(v))))
+  bf16 = property(lambda s: _bf16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _ibf16(float(v))))
+  b16, b32 = u16, u32
+
+  def __int__(self): return self._get()
+  def __index__(self): return self._get()
+
+  # Comparison operators (compare as integers)
+  def __eq__(s, o): return s._get() == int(o)
+  def __ne__(s, o): return s._get() != int(o)
+  def __lt__(s, o): return s._get() < int(o)
+  def __le__(s, o): return s._get() <= int(o)
+  def __gt__(s, o): return s._get() > int(o)
+  def __ge__(s, o): return s._get() >= int(o)
+
+class TypedView:
+  """View for S0.u32 that supports [4:0] slicing and [bit] access."""
+  __slots__ = ('_reg', '_bits', '_signed', '_float', '_bf16')
+  def __init__(self, reg, bits, signed=False, is_float=False, is_bf16=False):
+    self._reg, self._bits, self._signed, self._float, self._bf16 = reg, bits, signed, is_float, is_bf16
+
+  @property
+  def _val(self):
+    mask = MASK64 if self._bits == 64 else MASK32 if self._bits == 32 else (1 << self._bits) - 1
+    return self._reg._val & mask
+
+  def __getitem__(self, key):
+    if isinstance(key, slice):
+      high, low = int(key.start), int(key.stop)
+      return SliceProxy(self._reg, high, low)
+    return (self._val >> int(key)) & 1
+
+  def __setitem__(self, key, value):
+    if isinstance(key, slice):
+      high, low = int(key.start), int(key.stop)
+      if high < low: high, low, value = low, high, _brev(int(value), low - high + 1)
+      mask = (1 << (high - low + 1)) - 1
+      self._reg._val = (self._reg._val & ~(mask << low)) | ((int(value) & mask) << low)
+    elif value: self._reg._val |= (1 << int(key))
+    else: self._reg._val &= ~(1 << int(key))
+
+  def __int__(self): return _sext(self._val, self._bits) if self._signed else self._val
+  def __index__(self): return int(self)
+  def __trunc__(self): return int(float(self)) if self._float else int(self)
+  def __float__(self):
+    if self._float:
+      if self._bf16: return _bf16(self._val)  # bf16 uses different conversion
+      return _f16(self._val) if self._bits == 16 else _f32(self._val) if self._bits == 32 else _f64(self._val)
+    return float(int(self))
+
+  # Arithmetic - floats use float(), ints use int()
+  def __add__(s, o): return float(s) + float(o) if s._float else int(s) + int(o)
+  def __radd__(s, o): return float(o) + float(s) if s._float else int(o) + int(s)
+  def __sub__(s, o): return float(s) - float(o) if s._float else int(s) - int(o)
+  def __rsub__(s, o): return float(o) - float(s) if s._float else int(o) - int(s)
+  def __mul__(s, o): return float(s) * float(o) if s._float else int(s) * int(o)
+  def __rmul__(s, o): return float(o) * float(s) if s._float else int(o) * int(s)
+  def __truediv__(s, o): return _div(float(s), float(o)) if s._float else _div(int(s), int(o))
+  def __rtruediv__(s, o): return _div(float(o), float(s)) if s._float else _div(int(o), int(s))
+  def __pow__(s, o): return float(s) ** float(o) if s._float else int(s) ** int(o)
+  def __rpow__(s, o): return float(o) ** float(s) if s._float else int(o) ** int(s)
+  def __neg__(s): return -float(s) if s._float else -int(s)
+  def __abs__(s): return abs(float(s)) if s._float else abs(int(s))
+
+  # Bitwise - GPU shifts mask the shift amount to valid range
+  def __and__(s, o): return int(s) & int(o)
+  def __or__(s, o): return int(s) | int(o)
+  def __xor__(s, o): return int(s) ^ int(o)
+  def __invert__(s): return ~int(s)
+  def __lshift__(s, o): n = int(o); return int(s) << n if 0 <= n < 64 else 0
+  def __rshift__(s, o): n = int(o); return int(s) >> n if 0 <= n < 64 else 0
+  def __rand__(s, o): return int(o) & int(s)
+  def __ror__(s, o): return int(o) | int(s)
+  def __rxor__(s, o): return int(o) ^ int(s)
+  def __rlshift__(s, o): n = int(s); return int(o) << n if 0 <= n < 64 else 0
+  def __rrshift__(s, o): n = int(s); return int(o) >> n if 0 <= n < 64 else 0
+
+  # Comparison - handle _DenormChecker specially
+  def __eq__(s, o):
+    if isinstance(o, _DenormChecker): return o._check(s)
+    return float(s) == float(o) if s._float else int(s) == int(o)
+  def __ne__(s, o):
+    if isinstance(o, _DenormChecker): return not o._check(s)
+    return float(s) != float(o) if s._float else int(s) != int(o)
+  def __lt__(s, o): return float(s) < float(o) if s._float else int(s) < int(o)
+  def __le__(s, o): return float(s) <= float(o) if s._float else int(s) <= int(o)
+  def __gt__(s, o): return float(s) > float(o) if s._float else int(s) > int(o)
+  def __ge__(s, o): return float(s) >= float(o) if s._float else int(s) >= int(o)
+
+  def __bool__(s): return bool(int(s))
+
+  # Allow chained type access like jump_addr.i64 when jump_addr is already a TypedView
+  # These just return self or convert appropriately
+  @property
+  def i64(s): return s if s._bits == 64 and s._signed else int(s)
+  @property
+  def u64(s): return s if s._bits == 64 and not s._signed else int(s) & MASK64
+  @property
+  def i32(s): return s if s._bits == 32 and s._signed else _sext(int(s) & MASK32, 32)
+  @property
+  def u32(s): return s if s._bits == 32 and not s._signed else int(s) & MASK32
+
+class Reg:
+  """GPU register: D0.f32 = S0.f32 + S1.f32 just works."""
+  __slots__ = ('_val',)
+  def __init__(self, val=0): self._val = int(val) & MASK64
+
+  # Typed views
+  u64 = property(lambda s: TypedView(s, 64), lambda s, v: setattr(s, '_val', int(v) & MASK64))
+  i64 = property(lambda s: TypedView(s, 64, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK64))
+  b64 = property(lambda s: TypedView(s, 64), lambda s, v: setattr(s, '_val', int(v) & MASK64))
+  f64 = property(lambda s: TypedView(s, 64, is_float=True), lambda s, v: setattr(s, '_val', v if isinstance(v, int) else _i64(float(v))))
+  u32 = property(lambda s: TypedView(s, 32), lambda s, v: setattr(s, '_val', int(v) & MASK32))
+  i32 = property(lambda s: TypedView(s, 32, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK32))
+  b32 = property(lambda s: TypedView(s, 32), lambda s, v: setattr(s, '_val', int(v) & MASK32))
+  f32 = property(lambda s: TypedView(s, 32, is_float=True), lambda s, v: setattr(s, '_val', _i32(float(v))))
+  u24 = property(lambda s: TypedView(s, 24))
+  i24 = property(lambda s: TypedView(s, 24, signed=True))
+  u16 = property(lambda s: TypedView(s, 16), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
+  i16 = property(lambda s: TypedView(s, 16, signed=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
+  b16 = property(lambda s: TypedView(s, 16), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
+  f16 = property(lambda s: TypedView(s, 16, is_float=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _i16(float(v))) & 0xffff)))
+  bf16 = property(lambda s: TypedView(s, 16, is_float=True, is_bf16=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _ibf16(float(v))) & 0xffff)))
+  u8 = property(lambda s: TypedView(s, 8))
+  i8 = property(lambda s: TypedView(s, 8, signed=True))
+  u1 = property(lambda s: TypedView(s, 1))  # single bit
+
+  def __getitem__(s, key):
+    if isinstance(key, slice): return SliceProxy(s, int(key.start), int(key.stop))
+    return (s._val >> int(key)) & 1
+
+  def __setitem__(s, key, value):
+    if isinstance(key, slice):
+      high, low = int(key.start), int(key.stop)
+      mask = (1 << (high - low + 1)) - 1
+      s._val = (s._val & ~(mask << low)) | ((int(value) & mask) << low)
+    elif value: s._val |= (1 << int(key))
+    else: s._val &= ~(1 << int(key))
+
+  def __int__(s): return s._val
+  def __index__(s): return s._val
+  def __bool__(s): return bool(s._val)
+
+  # Arithmetic (for tmp = tmp + 1 patterns). Float operands trigger f32 interpretation.
+  def __add__(s, o): return (_f32(s._val) + float(o)) if isinstance(o, float) else s._val + int(o)
+  def __radd__(s, o): return (float(o) + _f32(s._val)) if isinstance(o, float) else int(o) + s._val
+  def __sub__(s, o): return (_f32(s._val) - float(o)) if isinstance(o, float) else s._val - int(o)
+  def __rsub__(s, o): return (float(o) - _f32(s._val)) if isinstance(o, float) else int(o) - s._val
+  def __mul__(s, o): return (_f32(s._val) * float(o)) if isinstance(o, float) else s._val * int(o)
+  def __rmul__(s, o): return (float(o) * _f32(s._val)) if isinstance(o, float) else int(o) * s._val
+  def __and__(s, o): return s._val & int(o)
+  def __rand__(s, o): return int(o) & s._val
+  def __or__(s, o): return s._val | int(o)
+  def __ror__(s, o): return int(o) | s._val
+  def __xor__(s, o): return s._val ^ int(o)
+  def __rxor__(s, o): return int(o) ^ s._val
+  def __lshift__(s, o): n = int(o); return s._val << n if 0 <= n < 64 else 0
+  def __rshift__(s, o): n = int(o); return s._val >> n if 0 <= n < 64 else 0
+  def __invert__(s): return ~s._val
+
+  # Comparison (for tmp >= 0x100000000 patterns)
+  def __lt__(s, o): return s._val < int(o)
+  def __le__(s, o): return s._val <= int(o)
+  def __gt__(s, o): return s._val > int(o)
+  def __ge__(s, o): return s._val >= int(o)
+  def __eq__(s, o): return s._val == int(o)
+  def __ne__(s, o): return s._val != int(o)
+
+

extra/assembly/amd/pdf.py

@@ -0,0 +1,671 @@
+# Generate AMD ISA autogen files from PDF documentation
+# Combines format/enum generation (previously in dsl.py) and pseudocode compilation (previously in pcode.py)
+# Usage: python -m extra.assembly.amd.pdf [--arch rdna3|rdna4|cdna|all]
+import re, functools
+from pathlib import Path
+from concurrent.futures import ProcessPoolExecutor
+
+PDF_URLS = {
+  "rdna3": "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content",
+  "rdna4": "https://docs.amd.com/api/khub/documents/uQpkEvk3pv~kfAb2x~j4uw/content",
+  "cdna": ["https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-mi300-cdna3-instruction-set-architecture.pdf",
+           "https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-cdna4-instruction-set-architecture.pdf"],
+}
+
+# Field type mappings and ordering
+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'}
+INST_PATTERN = re.compile(r'^([SV]_[A-Z0-9_]+)\s+(\d+)\s*$', re.M)
+
+# Patterns that can't be handled by the DSL (require special handling in emu.py)
+UNSUPPORTED = ['SGPR[', 'V_SWAP', 'eval ', 'FATAL_HALT', 'HW_REGISTERS',
+               'vscnt', 'vmcnt', 'expcnt', 'lgkmcnt',
+               'CVT_OFF_TABLE', 'ThreadMask',
+               'S1[i', 'C.i32', 'S[i]', 'in[',
+               'if n.', 'DST.u32', 'addrd = DST', 'addr = DST',
+               'BARRIER_STATE', 'ReallocVgprs',
+               'GPR_IDX', 'VSKIP', 'specified in', 'TTBL',
+               'fp6', 'bf6']  # Malformed pseudocode from PDF
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# COMPILER: pseudocode -> Python (minimal transforms)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def compile_pseudocode(pseudocode: str) -> str:
+  """Compile pseudocode to Python. Transforms are minimal - most syntax just works."""
+  pseudocode = re.sub(r'\bpass\b', 'pass_', pseudocode)  # 'pass' is Python keyword
+  raw_lines = pseudocode.strip().split('\n')
+  joined_lines: list[str] = []
+  for line in raw_lines:
+    line = line.strip()
+    if joined_lines and (joined_lines[-1].rstrip().endswith(('||', '&&', '(', ',')) or
+                         (joined_lines[-1].count('(') > joined_lines[-1].count(')'))):
+      joined_lines[-1] = joined_lines[-1].rstrip() + ' ' + line
+    else:
+      joined_lines.append(line)
+
+  lines = []
+  indent, need_pass, in_first_match_loop = 0, False, False
+  for line in joined_lines:
+    line = line.strip()
+    if not line or line.startswith('//'): continue
+    if line.startswith('if '):
+      lines.append('  ' * indent + f"if {_expr(line[3:].rstrip(' then'))}:")
+      indent += 1
+      need_pass = True
+    elif line.startswith('elsif '):
+      if need_pass: lines.append('  ' * indent + "pass")
+      indent -= 1
+      lines.append('  ' * indent + f"elif {_expr(line[6:].rstrip(' then'))}:")
+      indent += 1
+      need_pass = True
+    elif line == 'else':
+      if need_pass: lines.append('  ' * indent + "pass")
+      indent -= 1
+      lines.append('  ' * indent + "else:")
+      indent += 1
+      need_pass = True
+    elif line.startswith('endif'):
+      if need_pass: lines.append('  ' * indent + "pass")
+      indent -= 1
+      need_pass = False
+    elif line.startswith('endfor'):
+      if need_pass: lines.append('  ' * indent + "pass")
+      indent -= 1
+      need_pass, in_first_match_loop = False, False
+    elif line.startswith('declare '):
+      pass
+    elif m := re.match(r'for (\w+) in (.+?)\s*:\s*(.+?) do', line):
+      start, end = _expr(m[2].strip()), _expr(m[3].strip())
+      lines.append('  ' * indent + f"for {m[1]} in range({start}, int({end})+1):")
+      indent += 1
+      need_pass, in_first_match_loop = True, True
+    elif '=' in line and not line.startswith('=='):
+      need_pass = False
+      line = line.rstrip(';')
+      if m := re.match(r'\{\s*D1\.[ui]1\s*,\s*D0\.[ui]64\s*\}\s*=\s*(.+)', line):
+        rhs = _expr(m[1])
+        lines.append('  ' * indent + f"_full = {rhs}")
+        lines.append('  ' * indent + f"D0.u64 = int(_full) & 0xffffffffffffffff")
+        lines.append('  ' * indent + f"D1 = Reg((int(_full) >> 64) & 1)")
+      elif any(op in line for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^=')):
+        for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^='):
+          if op in line:
+            lhs, rhs = line.split(op, 1)
+            lines.append('  ' * indent + f"{lhs.strip()} {op} {_expr(rhs.strip())}")
+            break
+      else:
+        lhs, rhs = line.split('=', 1)
+        lhs_s, rhs_s = _expr(lhs.strip()), rhs.strip()
+        stmt = _assign(lhs_s, _expr(rhs_s))
+        if in_first_match_loop and rhs_s == 'i' and (lhs_s == 'tmp' or lhs_s == 'D0.i32'):
+          stmt += "; break"
+        lines.append('  ' * indent + stmt)
+  if need_pass: lines.append('  ' * indent + "pass")
+  return '\n'.join(lines)
+
+def _assign(lhs: str, rhs: str) -> str:
+  if lhs in ('tmp', 'SCC', 'VCC', 'EXEC', 'D0', 'D1', 'saveexec', 'PC'):
+    return f"{lhs} = Reg({rhs})"
+  return f"{lhs} = {rhs}"
+
+def _expr(e: str) -> str:
+  e = e.strip()
+  e = e.replace('&&', ' and ').replace('||', ' or ').replace('<>', ' != ')
+  e = re.sub(r'!([^=])', r' not \1', e)
+  e = re.sub(r'\{\s*(\w+\.u32)\s*,\s*(\w+\.u32)\s*\}', r'_pack32(\1, \2)', e)
+  def pack(m):
+    hi, lo = _expr(m[1].strip()), _expr(m[2].strip())
+    return f'_pack({hi}, {lo})'
+  e = re.sub(r'\{\s*([^,{}]+)\s*,\s*([^,{}]+)\s*\}', pack, e)
+  e = re.sub(r"1201'B\(2\.0\s*/\s*PI\)", "TWO_OVER_PI_1201", e)
+  e = re.sub(r"\d+'([0-9a-fA-Fx]+)[UuFf]*", r'\1', e)
+  e = re.sub(r"\d+'[FIBU]\(", "(", e)
+  e = re.sub(r'\bB\(', '(', e)
+  e = re.sub(r'([0-9a-fA-Fx])ULL\b', r'\1', e)
+  e = re.sub(r'([0-9a-fA-Fx])LL\b', r'\1', e)
+  e = re.sub(r'([0-9a-fA-Fx])U\b', r'\1', e)
+  e = re.sub(r'(\d\.?\d*)F\b', r'\1', e)
+  e = re.sub(r'(\[laneId\])\.[uib]\d+', r'\1', e)
+  e = e.replace('+INF', 'INF').replace('-INF', '(-INF)')
+  e = re.sub(r'NAN\.f\d+', 'float("nan")', e)
+  def convert_verilog_slice(m):
+    start, width = m.group(1).strip(), m.group(2).strip()
+    return f'[({start}) + ({width}) - 1 : ({start})]'
+  e = re.sub(r'\[([^:\[\]]+)\s*\+:\s*([^:\[\]]+)\]', convert_verilog_slice, e)
+  def process_brackets(s):
+    result, i = [], 0
+    while i < len(s):
+      if s[i] == '[':
+        depth, start = 1, i + 1
+        j = start
+        while j < len(s) and depth > 0:
+          if s[j] == '[': depth += 1
+          elif s[j] == ']': depth -= 1
+          j += 1
+        inner = _expr(s[start:j-1])
+        result.append('[' + inner + ']')
+        i = j
+      else:
+        result.append(s[i])
+        i += 1
+    return ''.join(result)
+  e = process_brackets(e)
+  while '?' in e:
+    depth, bracket, q = 0, 0, -1
+    for i, c in enumerate(e):
+      if c == '(': depth += 1
+      elif c == ')': depth -= 1
+      elif c == '[': bracket += 1
+      elif c == ']': bracket -= 1
+      elif c == '?' and depth == 0 and bracket == 0: q = i; break
+    if q < 0: break
+    depth, bracket, col = 0, 0, -1
+    for i in range(q + 1, len(e)):
+      if e[i] == '(': depth += 1
+      elif e[i] == ')': depth -= 1
+      elif e[i] == '[': bracket += 1
+      elif e[i] == ']': bracket -= 1
+      elif e[i] == ':' and depth == 0 and bracket == 0: col = i; break
+    if col < 0: break
+    cond, t, f = e[:q].strip(), e[q+1:col].strip(), e[col+1:].strip()
+    e = f'(({t}) if ({cond}) else ({f}))'
+  return e
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PDF PARSING WITH PAGE CACHING
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class CachedPDF:
+  """PDF wrapper with page text/table caching for faster repeated access."""
+  def __init__(self, pdf):
+    self._pdf, self._text_cache, self._table_cache = pdf, {}, {}
+  def __len__(self): return len(self._pdf.pages)
+  def text(self, i):
+    if i not in self._text_cache: self._text_cache[i] = self._pdf.pages[i].extract_text() or ''
+    return self._text_cache[i]
+  def tables(self, i):
+    if i not in self._table_cache: self._table_cache[i] = [t.extract() for t in self._pdf.pages[i].find_tables()]
+    return self._table_cache[i]
+
+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]:
+      desc = row[2]
+      # Handle shared FLAT/GLOBAL/SCRATCH table: look for format-specific encoding
+      fmt_key = fmt.lstrip('V').lower().capitalize()  # VFLAT -> Flat, VGLOBAL -> Global
+      if m := re.search(rf"{fmt_key}='b([01_]+)", desc):
+        enc_bits = m.group(1).replace('_', '')
+      elif m := re.search(r"(?:'b|Must be:\s*)([01_]+)", desc):
+        enc_bits = m.group(1).replace('_', '')
+      else:
+        enc_bits = None
+      if enc_bits:
+        enc_val, declared_width, actual_width = int(enc_bits, 2), 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 _parse_single_pdf(url: str):
+  """Parse a single PDF and return (formats, enums, src_enum, doc_name, instructions)."""
+  import pdfplumber
+  from tinygrad.helpers import fetch
+
+  pdf = CachedPDF(pdfplumber.open(fetch(url)))
+  total_pages = len(pdf)
+
+  # Auto-detect document type
+  first_page = pdf.text(0)
+  is_cdna4, is_cdna3 = 'CDNA4' in first_page or 'CDNA 4' in first_page, 'CDNA3' in first_page or 'MI300' in first_page
+  is_cdna, is_rdna4 = is_cdna3 or is_cdna4, 'RDNA4' in first_page or 'RDNA 4' in first_page
+  is_rdna35, is_rdna3 = 'RDNA3.5' in first_page or 'RDNA 3.5' in first_page, 'RDNA3' in first_page and 'RDNA3.5' not in first_page
+  doc_name = "CDNA4" if is_cdna4 else "CDNA3" if is_cdna3 else "RDNA4" if is_rdna4 else "RDNA3.5" if is_rdna35 else "RDNA3" if is_rdna3 else "Unknown"
+
+  # Find Microcode Formats section (for formats/enums)
+  microcode_start = next((i for i in range(int(total_pages * 0.2), total_pages)
+                          if re.search(r'\d+\.\d+\.\d+\.\s+SOP2\b|Chapter \d+\.\s+Microcode Formats', pdf.text(i))), int(total_pages * 0.9))
+  # Find Instructions section (for pseudocode)
+  instr_start = next((i for i in range(int(total_pages * 0.1), int(total_pages * 0.5))
+                      if re.search(r'Chapter \d+\.\s+Instructions\b', pdf.text(i))), total_pages // 3)
+  instr_end = next((i for start in [int(total_pages * 0.6), int(total_pages * 0.5), instr_start]
+                    for i in range(start, min(start + 100, total_pages))
+                    if re.search(r'Chapter \d+\.\s+Microcode Formats', pdf.text(i))), total_pages)
+
+  # Parse src enum from SSRC encoding table
+  src_enum = dict(SRC_EXTRAS)
+  for i in range(microcode_start, min(microcode_start + 10, total_pages)):
+    text = pdf.text(i)
+    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
+  full_text = '\n'.join(pdf.text(i) for i in range(microcode_start, min(microcode_start + 50, total_pages)))
+  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())
+
+  # Parse instruction formats
+  def is_fields_table(t): return t and len(t) > 1 and t[0] and 'Field' in str(t[0][0] or '')
+  def has_encoding(fields): return any(f[0] == 'ENCODING' for f in fields)
+  def has_header_before_fields(text): return (pos := text.find('Field Name')) != -1 and bool(re.search(r'\d+\.\d+\.\d+\.\s+\w+\s*\n', text[:pos]))
+
+  format_headers = []
+  for i in range(50):
+    if microcode_start + i >= total_pages: break
+    text = pdf.text(microcode_start + i)
+    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):
+      fmt_name = m.group(1)
+      if is_cdna and fmt_name.isupper() and len(fmt_name) >= 2: format_headers.append((fmt_name, i, m.start()))
+      elif m.start() > len(text) - 200 and 'Description' not in text[m.end():] and i + 1 < 50:
+        next_text = pdf.text(microcode_start + i + 1).lstrip()
+        if next_text.startswith('Description') or (next_text.startswith('"RDNA') and 'Description' in next_text[:200]):
+          format_headers.append((fmt_name, i, m.start()))
+    # RDNA4: Look for "Table X. Y Fields" patterns (e.g., VIMAGE, VSAMPLE, or shared FLAT/GLOBAL/SCRATCH)
+    for m in re.finditer(r'Table \d+\.\s+([\w,\s]+?)\s+Fields', text):
+      table_name = m.group(1).strip()
+      # Handle shared table like "FLAT, GLOBAL and SCRATCH"
+      if ',' in table_name or ' and ' in table_name:
+        for part in re.split(r',\s*|\s+and\s+', table_name):
+          fmt_name = 'V' + part.strip()
+          if fmt_name not in [h[0] for h in format_headers]: format_headers.append((fmt_name, i, m.start()))
+      elif table_name.startswith('V'):
+        if table_name not in [h[0] for h in format_headers]: format_headers.append((table_name, i, m.start()))
+
+  formats: dict[str, list] = {}
+  for fmt_name, rel_idx, header_pos in format_headers:
+    if fmt_name in formats: continue
+    page_idx = microcode_start + rel_idx
+    text = pdf.text(page_idx)
+    field_pos = text.find('Field Name', header_pos)
+    fields = None
+    for offset in range(3):
+      if page_idx + offset >= total_pages: break
+      if offset > 0 and has_header_before_fields(pdf.text(page_idx + offset)): break
+      for t in pdf.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
+    if not fields and field_pos > header_pos:
+      for t in pdf.tables(page_idx):
+        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}
+    for pg_offset in range(1, 3):
+      if page_idx + pg_offset >= total_pages or has_header_before_fields(pdf.text(page_idx + pg_offset)): break
+      for t in pdf.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
+  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']]
+  # RDNA4: VFLAT/VGLOBAL/VSCRATCH OP field is [20:14] not [20:13] (PDF documentation error)
+  for fmt_name in ['VFLAT', 'VGLOBAL', 'VSCRATCH']:
+    if fmt_name in formats:
+      formats[fmt_name] = [(n, h, 14 if n == 'OP' else l, e, t) for n, h, l, e, t in formats[fmt_name]]
+  if doc_name in ('RDNA3', 'RDNA3.5'):
+    if 'SOPPOp' in enums: assert 8 not in enums['SOPPOp']; enums['SOPPOp'][8] = 'S_WAITCNT_DEPCTR'
+    if 'DSOp' in enums:
+      for k, v in {24: 'DS_GWS_SEMA_RELEASE_ALL', 25: 'DS_GWS_INIT', 26: 'DS_GWS_SEMA_V', 27: 'DS_GWS_SEMA_BR', 28: 'DS_GWS_SEMA_P', 29: 'DS_GWS_BARRIER'}.items():
+        assert k not in enums['DSOp']; enums['DSOp'][k] = v
+    if 'FLATOp' in enums:
+      for k, v in {40: 'GLOBAL_LOAD_ADDTID_B32', 41: 'GLOBAL_STORE_ADDTID_B32', 55: 'FLAT_ATOMIC_CSUB_U32'}.items():
+        assert k not in enums['FLATOp']; enums['FLATOp'][k] = v
+
+  # Extract pseudocode for instructions
+  all_text = '\n'.join(pdf.text(i) for i in range(instr_start, instr_end))
+  matches = list(INST_PATTERN.finditer(all_text))
+  raw_pseudocode: dict[tuple[str, int], str] = {}
+  for i, match in enumerate(matches):
+    name, opcode = match.group(1), int(match.group(2))
+    start, end = match.end(), matches[i + 1].start() if i + 1 < len(matches) else match.end() + 2000
+    snippet = all_text[start:end].strip()
+    if pseudocode := _extract_pseudocode(snippet): raw_pseudocode[(name, opcode)] = pseudocode
+
+  return {"formats": formats, "enums": enums, "src_enum": src_enum, "doc_name": doc_name, "pseudocode": raw_pseudocode, "is_cdna": is_cdna}
+
+def _extract_pseudocode(text: str) -> str | None:
+  """Extract pseudocode from an instruction description snippet."""
+  lines, result, depth, in_lambda = text.split('\n'), [], 0, 0
+  for line in lines:
+    s = line.strip()
+    if not s or re.match(r'^\d+ of \d+$', s) or re.match(r'^\d+\.\d+\..*Instructions', s): continue
+    if s.startswith(('Notes', 'Functional examples')): break
+    if s.startswith(('"RDNA', 'AMD ', 'CDNA')): continue
+    if '= lambda(' in s: in_lambda += 1; continue
+    if in_lambda > 0:
+      if s.endswith(');'): in_lambda -= 1
+      continue
+    if s.startswith('if '): depth += 1
+    elif s.startswith('endif'): depth = max(0, depth - 1)
+    if s.endswith('.') and not any(p in s for p in ['D0', 'D1', 'S0', 'S1', 'S2', 'SCC', 'VCC', 'tmp', '=']): continue
+    if re.match(r'^[a-z].*\.$', s) and '=' not in s: continue
+    is_code = (any(p in s for p in ['D0.', 'D1.', 'S0.', 'S1.', 'S2.', 'SCC =', 'SCC ?', 'VCC', 'EXEC', 'tmp =', 'tmp[', 'lane =', 'PC =',
+                                    'D0[', 'D1[', 'S0[', 'S1[', 'S2[']) or
+               s.startswith(('if ', 'else', 'elsif', 'endif', 'declare ', 'for ', 'endfor', '//')) or
+               re.match(r'^[a-z_]+\s*=', s) or re.match(r'^[a-z_]+\[', s) or (depth > 0 and '=' in s))
+    if is_code: result.append(s)
+  return '\n'.join(result) if result else None
+
+def _merge_results(results: list[dict]) -> dict:
+  """Merge multiple PDF parse results into a superset."""
+  merged = {"formats": {}, "enums": {}, "src_enum": dict(SRC_EXTRAS), "doc_names": [], "pseudocode": {}, "is_cdna": False}
+  for r in results:
+    merged["doc_names"].append(r["doc_name"])
+    merged["is_cdna"] = merged["is_cdna"] or r["is_cdna"]
+    for val, name in r["src_enum"].items():
+      if val in merged["src_enum"]: assert merged["src_enum"][val] == name
+      else: merged["src_enum"][val] = name
+    for enum_name, ops in r["enums"].items():
+      if enum_name not in merged["enums"]: merged["enums"][enum_name] = {}
+      for val, name in ops.items():
+        if val in merged["enums"][enum_name]: assert merged["enums"][enum_name][val] == name
+        else: merged["enums"][enum_name][val] = name
+    for fmt_name, fields in r["formats"].items():
+      if fmt_name not in merged["formats"]: merged["formats"][fmt_name] = list(fields)
+      else:
+        existing = {f[0]: (f[1], f[2]) for f in merged["formats"][fmt_name]}
+        for f in fields:
+          if f[0] in existing: assert existing[f[0]] == (f[1], f[2])
+          else: merged["formats"][fmt_name].append(f)
+    for key, pc in r["pseudocode"].items():
+      if key not in merged["pseudocode"]: merged["pseudocode"][key] = pc
+  return merged
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# CODE GENERATION
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def _generate_enum_py(enums, src_enum, doc_name) -> str:
+  """Generate enum.py content (just enums, no dsl.py dependency)."""
+  def enum_lines(name, items): return [f"class {name}(IntEnum):"] + [f"  {n} = {v}" for v, n in sorted(items.items())] + [""]
+  lines = [f"# autogenerated from AMD {doc_name} ISA PDF by pdf.py - do not edit", "from enum import IntEnum", ""]
+  lines += enum_lines("SrcEnum", src_enum) + sum([enum_lines(n, ops) for n, ops in sorted(enums.items())], [])
+  return '\n'.join(lines)
+
+def _generate_ins_py(formats, enums, src_enum, doc_name) -> str:
+  """Generate ins.py content (instruction formats and helpers, imports dsl.py and enum.py)."""
+  def field_key(f, order): return order.index(f[0].lower()) if f[0].lower() in order else 1000
+  lines = [f"# autogenerated from AMD {doc_name} ISA PDF by pdf.py - do not edit",
+           "# ruff: noqa: F401,F403", "from typing import Annotated",
+           "from extra.assembly.amd.dsl import bits, BitField, Inst32, Inst64, Inst96, SGPR, VGPR, TTMP as TTMP, s as s, v as v, ttmp as ttmp, SSrc, Src, SImm, Imm, VDSTYEnc, SGPRField, VGPRField",
+           "from extra.assembly.amd.autogen.{arch}.enum import *",
+           "import functools", ""]
+  format_defaults = {'VOP3P': {'opsel_hi': 3, 'opsel_hi2': 1}}
+  lines.append("# instruction formats")
+  for fmt_name, fields in sorted(formats.items()):
+    max_bit = max(f[1] for f in fields)
+    base = "Inst96" if max_bit > 63 else "Inst64" if max_bit > 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):
+      lines.append(f"  encoding = bits[{enc[1]}:{enc[2]}] == 0b{enc[3]:b}" if enc[1] != enc[2] else f"  encoding = bits[{enc[1]}] == {enc[3]}")
+    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=lambda f: field_key(f, order)):
+      ann = f":Annotated[BitField, {ftype}]" if ftype and ftype.endswith('Op') else 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=1"}.get(fmt, "")
+      tgt = {"GLOBAL": "FLAT, GLOBALOp", "SCRATCH": "FLAT, SCRATCHOp"}.get(fmt, f"{fmt}, {cls_name}")
+      if fmt in formats or fmt in ("GLOBAL", "SCRATCH"):
+        suffix = "_e32" if fmt in ("VOP1", "VOP2", "VOPC") else "_e64" if fmt == "VOP3" and op_val < 512 else ""
+        if name in ('V_FMAMK_F32', 'V_FMAMK_F16'):
+          lines.append(f"def {name.lower()}{suffix}(vdst, src0, K, vsrc1): return {fmt}({cls_name}.{name}, vdst, src0, vsrc1, literal=K)")
+        elif name in ('V_FMAAK_F32', 'V_FMAAK_F16'):
+          lines.append(f"def {name.lower()}{suffix}(vdst, src0, vsrc1, K): return {fmt}({cls_name}.{name}, vdst, src0, vsrc1, literal=K)")
+        else: lines.append(f"{name.lower()}{suffix} = functools.partial({tgt}.{name}{seg})")
+  src_names = {name for _, name in src_enum.items()}
+  lines += [""] + [f"{name} = SrcEnum.{name}" for _, name in sorted(src_enum.items()) if name not in {'DPP8', 'DPP16'}]
+  if "NULL" in src_names: lines.append("OFF = NULL\n")
+  return '\n'.join(lines)
+
+def _generate_gen_pcode_py(enums, pseudocode, arch) -> str:
+  """Generate gen_pcode.py content (compiled pseudocode functions)."""
+  # Get op enums for this arch (import from .ins which re-exports from .enum)
+  import importlib
+  autogen = importlib.import_module(f"extra.assembly.amd.autogen.{arch}.ins")
+  OP_ENUMS = [getattr(autogen, name) for name in ['SOP1Op', 'SOP2Op', 'SOPCOp', 'SOPKOp', 'SOPPOp', 'VOP1Op', 'VOP2Op', 'VOP3Op', 'VOP3SDOp', 'VOP3POp', 'VOPCOp', 'VOP3AOp', 'VOP3BOp'] if hasattr(autogen, name)]
+
+  # Build defined ops mapping
+  defined_ops: dict[tuple, list] = {}
+  for enum_cls in OP_ENUMS:
+    for op in enum_cls:
+      if op.name.startswith(('S_', 'V_')): defined_ops.setdefault((op.name, op.value), []).append((enum_cls, op))
+
+  enum_names = [e.__name__ for e in OP_ENUMS]
+  lines = [f'''# autogenerated by pdf.py - do not edit
+# to regenerate: python -m extra.assembly.amd.pdf --arch {arch}
+# ruff: noqa: E501,F405,F403
+# mypy: ignore-errors
+from extra.assembly.amd.autogen.{arch}.enum import {", ".join(enum_names)}
+from extra.assembly.amd.pcode import *
+''']
+
+  instructions: dict = {cls: {} for cls in OP_ENUMS}
+  for key, pc in pseudocode.items():
+    if key in defined_ops:
+      for enum_cls, enum_val in defined_ops[key]: instructions[enum_cls][enum_val] = pc
+
+  for enum_cls in OP_ENUMS:
+    cls_name = enum_cls.__name__
+    if not instructions.get(enum_cls): continue
+    fn_entries = []
+    for op, pc in instructions[enum_cls].items():
+      if any(p in pc for p in UNSUPPORTED): continue
+      try:
+        code = compile_pseudocode(pc)
+        code = _apply_pseudocode_fixes(op, code)
+        fn_name, fn_code = _generate_function(cls_name, op, pc, code)
+        lines.append(fn_code)
+        fn_entries.append((op, fn_name))
+      except Exception as e: print(f"  Warning: Failed to compile {op.name}: {e}")
+    if fn_entries:
+      lines.append(f'{cls_name}_FUNCTIONS = {{')
+      for op, fn_name in fn_entries: lines.append(f"  {cls_name}.{op.name}: {fn_name},")
+      lines.append('}\n')
+
+  # Add V_WRITELANE_B32 if VOP3Op exists
+  if 'VOP3Op' in enum_names:
+    lines.append('''
+# V_WRITELANE_B32: Write scalar to specific lane's VGPR (not in PDF pseudocode)
+def _VOP3Op_V_WRITELANE_B32(s0, s1, s2, d0, scc, vcc, lane, exec_mask, literal, VGPR, _vars, src0_idx=0, vdst_idx=0):
+  wr_lane = s1 & 0x1f
+  return {'d0': d0, 'scc': scc, 'vgpr_write': (wr_lane, vdst_idx, s0 & 0xffffffff)}
+VOP3Op_FUNCTIONS[VOP3Op.V_WRITELANE_B32] = _VOP3Op_V_WRITELANE_B32
+''')
+
+  lines.append('COMPILED_FUNCTIONS = {')
+  for enum_cls in OP_ENUMS:
+    if instructions.get(enum_cls): lines.append(f'  {enum_cls.__name__}: {enum_cls.__name__}_FUNCTIONS,')
+  lines.append('}\n\ndef get_compiled_functions(): return COMPILED_FUNCTIONS')
+  return '\n'.join(lines)
+
+def _apply_pseudocode_fixes(op, code: str) -> str:
+  """Apply known fixes for PDF pseudocode bugs."""
+  if op.name == 'V_DIV_FMAS_F32':
+    code = code.replace('D0.f32 = 2.0 ** 32 * fma(S0.f32, S1.f32, S2.f32)',
+                        'D0.f32 = (2.0 ** 64 if exponent(S2.f32) > 127 else 2.0 ** -64) * fma(S0.f32, S1.f32, S2.f32)')
+  if op.name == 'V_DIV_FMAS_F64':
+    code = code.replace('D0.f64 = 2.0 ** 64 * fma(S0.f64, S1.f64, S2.f64)',
+                        'D0.f64 = (2.0 ** 128 if exponent(S2.f64) > 1023 else 2.0 ** -128) * fma(S0.f64, S1.f64, S2.f64)')
+  if op.name == 'V_DIV_SCALE_F32':
+    code = code.replace('D0.f32 = float("nan")', 'VCC = Reg(0x1); D0.f32 = float("nan")')
+    code = code.replace('elif S1.f32 == DENORM.f32:\n  D0.f32 = ldexp(S0.f32, 64)', 'elif False:\n  pass')
+    code += '\nif S1.f32 == DENORM.f32:\n  D0.f32 = float("nan")'
+    code = code.replace('elif exponent(S2.f32) <= 23:\n  D0.f32 = ldexp(S0.f32, 64)', 'elif exponent(S2.f32) <= 23:\n  VCC = Reg(0x1); D0.f32 = ldexp(S0.f32, 64)')
+    code = code.replace('elif S2.f32 / S1.f32 == DENORM.f32:\n  VCC = Reg(0x1)\n  if S0.f32 == S2.f32:\n    D0.f32 = ldexp(S0.f32, 64)', 'elif S2.f32 / S1.f32 == DENORM.f32:\n  VCC = Reg(0x1)')
+  if op.name == 'V_DIV_SCALE_F64':
+    code = code.replace('D0.f64 = float("nan")', 'VCC = Reg(0x1); D0.f64 = float("nan")')
+    code = code.replace('elif S1.f64 == DENORM.f64:\n  D0.f64 = ldexp(S0.f64, 128)', 'elif False:\n  pass')
+    code += '\nif S1.f64 == DENORM.f64:\n  D0.f64 = float("nan")'
+    code = code.replace('elif exponent(S2.f64) <= 52:\n  D0.f64 = ldexp(S0.f64, 128)', 'elif exponent(S2.f64) <= 52:\n  VCC = Reg(0x1); D0.f64 = ldexp(S0.f64, 128)')
+    code = code.replace('elif S2.f64 / S1.f64 == DENORM.f64:\n  VCC = Reg(0x1)\n  if S0.f64 == S2.f64:\n    D0.f64 = ldexp(S0.f64, 128)', 'elif S2.f64 / S1.f64 == DENORM.f64:\n  VCC = Reg(0x1)')
+  if op.name == 'V_DIV_FIXUP_F32':
+    code = code.replace('D0.f32 = ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))',
+                        'D0.f32 = ((-OVERFLOW_F32) if (sign_out) else (OVERFLOW_F32)) if isNAN(S0.f32) else ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))')
+  if op.name == 'V_DIV_FIXUP_F64':
+    code = code.replace('D0.f64 = ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))',
+                        'D0.f64 = ((-OVERFLOW_F64) if (sign_out) else (OVERFLOW_F64)) if isNAN(S0.f64) else ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))')
+  if op.name == 'V_TRIG_PREOP_F64':
+    code = code.replace('result = F((TWO_OVER_PI_1201[1200 : 0] << shift.u32) & 0x1fffffffffffff)',
+                        'result = float(((TWO_OVER_PI_1201[1200 : 0] << int(shift)) >> (1201 - 53)) & 0x1fffffffffffff)')
+  return code
+
+def _generate_function(cls_name: str, op, pc: str, code: str) -> tuple[str, str]:
+  """Generate a single compiled pseudocode function."""
+  has_d1 = '{ D1' in pc
+  is_cmpx = (cls_name in ('VOPCOp', 'VOP3Op')) and 'EXEC.u64[laneId]' in pc
+  is_div_scale = 'DIV_SCALE' in op.name
+  has_sdst = cls_name == 'VOP3SDOp' and ('VCC.u64[laneId]' in pc or is_div_scale)
+  combined = code + pc
+
+  fn_name = f"_{cls_name}_{op.name}"
+  # Function accepts Reg objects directly (uppercase names), laneId is passed directly as int
+  lines = [f"def {fn_name}(S0, S1, S2, D0, SCC, VCC, laneId, EXEC, literal, VGPR, src0_idx=0, vdst_idx=0, PC=None):"]
+
+  # Registers that need special handling (not passed directly)
+  # Only init if used but not first assigned as `name = Reg(...)` in the compiled code
+  def needs_init(name): return name in combined and not re.search(rf'^\s*{name}\s*=\s*Reg\(', code, re.MULTILINE)
+  special_regs = [('D1', 'Reg(0)'), ('SIMM16', 'Reg(literal)'), ('SIMM32', 'Reg(literal)'),
+                  ('SRC0', 'Reg(src0_idx)'), ('VDST', 'Reg(vdst_idx)')]
+  if needs_init('tmp'): special_regs.insert(0, ('tmp', 'Reg(0)'))
+  if needs_init('saveexec'): special_regs.insert(0, ('saveexec', 'Reg(EXEC._val)'))
+  used = {name for name, _ in special_regs if name in combined}
+
+  # Detect which registers are modified (not just read) - look for assignments
+  modifies_d0 = is_div_scale or bool(re.search(r'\bD0\b[.\[]', combined))
+  modifies_exec = is_cmpx or bool(re.search(r'EXEC\.(u32|u64|b32|b64)\s*=', combined))
+  modifies_vcc = has_sdst or bool(re.search(r'VCC\.(u32|u64|b32|b64)\s*=|VCC\.u64\[laneId\]\s*=', combined))
+  modifies_scc = bool(re.search(r'\bSCC\s*=', combined))
+  modifies_pc = bool(re.search(r'\bPC\s*=', combined))
+
+  # Build init code for special registers
+  init_lines = []
+  if is_div_scale: init_lines.append("  D0 = Reg(S0._val)")
+  for name, init in special_regs:
+    if name in used: init_lines.append(f"  {name} = {init}")
+  if 'EXEC_LO' in code: init_lines.append("  EXEC_LO = SliceProxy(EXEC, 31, 0)")
+  if 'EXEC_HI' in code: init_lines.append("  EXEC_HI = SliceProxy(EXEC, 63, 32)")
+  if 'VCCZ' in code and not re.search(r'^\s*VCCZ\s*=', code, re.MULTILINE): init_lines.append("  VCCZ = Reg(1 if VCC._val == 0 else 0)")
+  if 'EXECZ' in code and not re.search(r'^\s*EXECZ\s*=', code, re.MULTILINE): init_lines.append("  EXECZ = Reg(1 if EXEC._val == 0 else 0)")
+  code_lines = [line for line in code.split('\n') if line.strip()]
+  if init_lines:
+    lines.extend(init_lines)
+    if code_lines: lines.append("  # --- compiled pseudocode ---")
+  for line in code_lines:
+    lines.append(f"  {line}")
+
+  # Build result dict - only include registers that are modified
+  result_items = []
+  if modifies_d0: result_items.append("'D0': D0")
+  if modifies_scc: result_items.append("'SCC': SCC")
+  if modifies_vcc: result_items.append("'VCC': VCC")
+  if modifies_exec: result_items.append("'EXEC': EXEC")
+  if has_d1: result_items.append("'D1': D1")
+  if modifies_pc: result_items.append("'PC': PC")
+  lines.append(f"  return {{{', '.join(result_items)}}}\n")
+  return fn_name, '\n'.join(lines)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# MAIN GENERATION
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def generate_arch(arch: str) -> dict:
+  """Generate enum.py, ins.py and gen_pcode.py for a single architecture."""
+  urls = PDF_URLS[arch]
+  if isinstance(urls, str): urls = [urls]
+
+  print(f"\n{'='*60}\nGenerating {arch}...")
+  print(f"Parsing {len(urls)} PDF(s)...")
+  results = [_parse_single_pdf(url) for url in urls]
+  merged = _merge_results(results) if len(results) > 1 else results[0]
+  doc_name = "+".join(merged["doc_names"]) if len(results) > 1 else merged["doc_name"]
+
+  base_path = Path(f"extra/assembly/amd/autogen/{arch}")
+  base_path.mkdir(parents=True, exist_ok=True)
+  (base_path / "__init__.py").touch()
+
+  # Write enum.py (enums only, no dsl.py dependency)
+  enum_path = base_path / "enum.py"
+  enum_content = _generate_enum_py(merged["enums"], merged["src_enum"], doc_name)
+  enum_path.write_text(enum_content)
+  print(f"Generated {enum_path}: SrcEnum ({len(merged['src_enum'])}) + {len(merged['enums'])} enums")
+
+  # Write ins.py (instruction formats and helpers, imports dsl.py and enum.py)
+  ins_path = base_path / "ins.py"
+  ins_content = _generate_ins_py(merged["formats"], merged["enums"], merged["src_enum"], doc_name).replace("{arch}", arch)
+  ins_path.write_text(ins_content)
+  print(f"Generated {ins_path}: {len(merged['formats'])} formats")
+
+  # Write gen_pcode.py (needs enum.py to exist first for imports)
+  pcode_path = base_path / "gen_pcode.py"
+  pcode_content = _generate_gen_pcode_py(merged["enums"], merged["pseudocode"], arch)
+  pcode_path.write_text(pcode_content)
+  print(f"Generated {pcode_path}: {len(merged['pseudocode'])} instructions")
+
+  return merged
+
+def _generate_arch_wrapper(arch: str):
+  """Wrapper for multiprocessing - returns arch name for ordering."""
+  generate_arch(arch)
+  return arch
+
+def generate_all():
+  """Generate all architectures in parallel."""
+  with ProcessPoolExecutor() as executor:
+    list(executor.map(_generate_arch_wrapper, PDF_URLS.keys()))
+
+if __name__ == "__main__":
+  import argparse
+  parser = argparse.ArgumentParser(description="Generate AMD ISA autogen files from PDF documentation")
+  parser.add_argument("--arch", choices=list(PDF_URLS.keys()) + ["all"], default="rdna3")
+  args = parser.parse_args()
+  if args.arch == "all": generate_all()
+  else: generate_arch(args.arch)

extra/assembly/amd/test/bench_emu.py

@@ -0,0 +1,294 @@
+#!/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.amd.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/amd/test/external_test_usability.py

@@ -0,0 +1,196 @@
+# Usability tests for the RDNA3 ASM DSL
+# These tests demonstrate how the DSL *should* work for a good user experience
+# Currently many of these tests fail - they document desired behavior
+
+import unittest
+from extra.assembly.amd.autogen.rdna3.ins import *
+from extra.assembly.amd.dsl import Inst, RawImm, SGPR, VGPR
+
+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 TestReprReadability(unittest.TestCase):
+  """
+  Issue: repr() leaks internal RawImm type and omits zero-valued fields.
+
+  When you create v_mov_b32_e32(v[0], v[1]), the repr shows:
+    VOP1(op=1, src0=RawImm(257))
+
+  Problems:
+  1. vdst=v[0] is omitted because 0 is treated as "default"
+  2. src0 shows RawImm(257) instead of v[1]
+  3. User sees encoded values (257 = 256 + 1) instead of register names
+
+  Expected repr: VOP1(op=1, vdst=v[0], src0=v[1])
+  """
+  def test_repr_shows_registers_not_raw_imm(self):
+    inst = v_mov_b32_e32(v[0], v[1])
+    # Should show v[1], not RawImm(257)
+    self.assertNotIn("RawImm", repr(inst), "repr should not expose RawImm internal type")
+    self.assertIn("v[1]", repr(inst), "repr should show register name")
+
+  def test_repr_includes_zero_dst(self):
+    inst = v_mov_b32_e32(v[0], v[1])
+    # v[0] is a valid destination register, should be shown
+    self.assertIn("vdst", repr(inst), "repr should include vdst even when 0")
+
+  def test_repr_roundtrip(self):
+    # repr should produce something that can be eval'd back
+    inst = v_mov_b32_e32(v[0], v[1])
+    # This would require repr to output valid Python, e.g.:
+    # "VOP1(op=VOP1Op.V_MOV_B32, vdst=v[0], src0=v[1])"
+    r = repr(inst)
+    # At minimum, it should be human-readable
+    self.assertIn("v[", r, "repr should show register syntax")
+
+
+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")
+
+
+class TestVOPDHelperSignature(unittest.TestCase):
+  """
+  Issue: VOPD helper functions have confusing semantics.
+
+  v_dual_mul_f32 is defined as:
+    v_dual_mul_f32 = functools.partial(VOPD, VOPDOp.V_DUAL_MUL_F32)
+
+  This binds VOPDOp.V_DUAL_MUL_F32 to the FIRST positional arg of VOPD.__init__,
+  which is 'opx'. So v_dual_mul_f32 sets the X operation.
+
+  But then test_dual_mul in test_handwritten.py does:
+    v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], ...)
+
+  This passes V_DUAL_MUL_F32 as the SECOND positional arg (opy), making both
+  X and Y operations the same. This is confusing because:
+  1. The function name suggests it handles the X operation
+  2. But you still pass an opcode as the first arg (which becomes opy)
+
+  Expected: Either make the helper fully specify both ops, or make the
+  signature clearer about what the positional arg means.
+  """
+  def test_vopd_helper_opy_should_be_required(self):
+    # Using only keyword args "works" but opy silently defaults to 0
+    inst = v_dual_mul_f32(vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
+    self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32)
+    # Bug: opy defaults to 0 (V_DUAL_FMAC_F32) silently - should require explicit opy
+    # This test documents the bug - it should fail once fixed
+    self.assertNotEqual(inst.opy, VOPDOp.V_DUAL_FMAC_F32, "opy should not silently default to FMAC")
+
+  def test_vopd_helper_positional_arg_is_opy(self):
+    # The first positional arg after the partial becomes opy, not a second opx
+    inst = v_dual_mul_f32(VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
+    self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32)  # From partial
+    self.assertEqual(inst.opy, VOPDOp.V_DUAL_MOV_B32)  # From first positional arg
+
+
+class TestFieldAccessPreservesType(unittest.TestCase):
+  """
+  Issue: Field access loses type information.
+
+  After creating an instruction, accessing fields returns encoded int values:
+    inst = v_mov_b32_e32(v[0], v[1])
+    inst.vdst  # returns 0, not VGPR(0)
+
+  This makes it impossible to round-trip register types through field access.
+  """
+  def test_vdst_returns_register(self):
+    inst = v_mov_b32_e32(v[5], v[1])
+    vdst = inst.vdst
+    # Should return a VGPR, not an int
+    self.assertIsInstance(vdst, (VGPR, int), "vdst should return VGPR or at least be usable")
+    # Ideally: self.assertIsInstance(vdst, VGPR)
+
+  def test_src_returns_register_for_vgpr_source(self):
+    inst = v_mov_b32_e32(v[0], v[1])
+    # src0 is encoded as 257 (256 + 1 for v1)
+    # Ideally it should decode back to v[1]
+    src0_raw = inst._values.get('src0')
+    # Currently returns RawImm(257), should return VGPR(1) or similar
+    self.assertNotIsInstance(src0_raw, RawImm, "source should not be RawImm internally")
+
+
+class TestArgumentDiscoverability(unittest.TestCase):
+  """
+  Issue: No clear signature for positional arguments.
+
+  inspect.signature(s_load_b128) shows: (*args, literal=None, **kwargs)
+
+  Users have no way to know the argument order without reading source code.
+  The order is implicitly defined by the class field definition order.
+
+  Possible fixes:
+  1. Add explicit parameter names to functools.partial
+  2. Generate type stubs with proper signatures
+  3. Add docstrings listing the expected arguments
+  """
+  def test_signature_has_named_params(self):
+    import inspect
+    sig = inspect.signature(s_load_b128)
+    params = list(sig.parameters.keys())
+    # Currently: ['args', 'literal', 'kwargs'] (from *args, literal=None, **kwargs)
+    # Expected: something like ['sdata', 'sbase', 'soffset', 'offset', 'literal']
+    self.assertIn('sdata', params, "signature should show field names")
+
+
+class TestSpecialConstants(unittest.TestCase):
+  """
+  Issue: NULL and other constants are IntEnum values that might be confusing.
+
+  NULL = SrcEnum.NULL = 124, but users might expect NULL to be a special object
+  that clearly represents "no register" rather than a magic number.
+  """
+  def test_null_has_clear_repr(self):
+    # NULL should have a clear string representation
+    self.assertIn("NULL", str(NULL) or repr(NULL), "NULL should be clearly identifiable")
+
+  def test_null_is_distinguishable_from_int(self):
+    # NULL should be distinguishable from the raw integer 124
+    self.assertNotEqual(type(NULL), int, "NULL should not be plain int")
+
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/amd/test/helpers.py

@@ -0,0 +1,66 @@
+"""Shared test helpers for RDNA3 tests."""
+import shutil
+from dataclasses import dataclass
+
+@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
+
+# LLVM tool detection (shared across test files)
+def get_llvm_mc():
+  """Find llvm-mc executable, preferring newer versions."""
+  for p in ['llvm-mc', 'llvm-mc-21', 'llvm-mc-20']:
+    if shutil.which(p): return p
+  raise FileNotFoundError("llvm-mc not found")
+
+def get_llvm_objdump():
+  """Find llvm-objdump executable, preferring newer versions."""
+  for p in ['llvm-objdump', 'llvm-objdump-21', 'llvm-objdump-20']:
+    if shutil.which(p): return p
+  raise FileNotFoundError("llvm-objdump not found")
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# EXECUTION CONTEXT (for testing compiled pseudocode)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class ExecContext:
+  """Context for running compiled pseudocode in tests."""
+  def __init__(self, s0=0, s1=0, s2=0, d0=0, scc=0, vcc=0, lane=0, exec_mask=0xffffffff, literal=0, vgprs=None, src0_idx=0, vdst_idx=0):
+    from extra.assembly.amd.pcode import Reg, MASK32, MASK64, SliceProxy
+    self._Reg, self._MASK64, self._SliceProxy = Reg, MASK64, SliceProxy
+    self.S0, self.S1, self.S2 = Reg(s0), Reg(s1), Reg(s2)
+    self.D0, self.D1 = Reg(d0), Reg(0)
+    self.SCC, self.VCC, self.EXEC = Reg(scc), Reg(vcc), Reg(exec_mask)
+    self.tmp, self.saveexec = Reg(0), Reg(exec_mask)
+    self.lane, self.laneId, self.literal = lane, lane, literal
+    self.SIMM16, self.SIMM32 = Reg(literal), Reg(literal)
+    self.VGPR = vgprs if vgprs is not None else {}
+    self.SRC0, self.VDST = Reg(src0_idx), Reg(vdst_idx)
+
+  def run(self, code: str):
+    """Execute compiled code."""
+    import extra.assembly.amd.pcode as pcode
+    ns = {k: getattr(pcode, k) for k in dir(pcode) if not k.startswith('_')}
+    # Also include underscore-prefixed helpers that compiled pseudocode uses
+    for k in ['_pack', '_pack32']:
+      if hasattr(pcode, k): ns[k] = getattr(pcode, k)
+    ns.update({
+      'S0': self.S0, 'S1': self.S1, 'S2': self.S2, 'D0': self.D0, 'D1': self.D1,
+      'SCC': self.SCC, 'VCC': self.VCC, 'EXEC': self.EXEC,
+      'EXEC_LO': self._SliceProxy(self.EXEC, 31, 0), 'EXEC_HI': self._SliceProxy(self.EXEC, 63, 32),
+      'tmp': self.tmp, 'saveexec': self.saveexec,
+      'lane': self.lane, 'laneId': self.laneId, 'literal': self.literal,
+      'SIMM16': self.SIMM16, 'SIMM32': self.SIMM32, 'VGPR': self.VGPR, 'SRC0': self.SRC0, 'VDST': self.VDST,
+    })
+    exec(code, ns)
+    def _sync(ctx_reg, ns_val):
+      if isinstance(ns_val, self._Reg): ctx_reg._val = ns_val._val
+      else: ctx_reg._val = int(ns_val) & self._MASK64
+    for name in ('SCC', 'VCC', 'EXEC', 'D0', 'D1', 'tmp', 'saveexec'):
+      if ns.get(name) is not getattr(self, name): _sync(getattr(self, name), ns[name])
+
+  def result(self) -> dict: return {"d0": self.D0._val, "scc": self.SCC._val & 1}

extra/assembly/amd/test/test_compare_emulators.py

@@ -0,0 +1,402 @@
+# 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"
+os.environ["PYTHON_REMU"] = "1"
+
+from extra.assembly.amd.emu import WaveState, decode_program, step_wave, WAVE_SIZE, set_valid_mem_ranges
+from extra.assembly.amd.test.helpers import KernelInfo
+
+REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
+
+def _is_f32_nan(bits: int) -> bool:
+  """Check if 32-bit value is a NaN (exponent all 1s, mantissa non-zero)."""
+  return (bits & 0x7f800000) == 0x7f800000 and (bits & 0x007fffff) != 0
+
+def _vals_equal(a: int, b: int) -> bool:
+  """Compare two 32-bit values, treating all NaN bit patterns as equal."""
+  if a == b: return True
+  return _is_f32_nan(a) and _is_f32_nan(b)
+
+@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, arrow: str = " vs ") -> list[str]:
+    """Return list of differences between two states."""
+    diffs = []
+    if self.pc != other.pc: diffs.append(f"pc: {self.pc}{arrow}{other.pc}")
+    if self.scc != other.scc: diffs.append(f"scc: {self.scc}{arrow}{other.scc}")
+    if self.vcc != other.vcc: diffs.append(f"vcc: 0x{self.vcc:08x}{arrow}0x{other.vcc:08x}")
+    if self.exec_mask != other.exec_mask: diffs.append(f"exec: 0x{self.exec_mask:08x}{arrow}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 not _vals_equal(a, b): diffs.append(f"sgpr[{i}]: 0x{a:08x}{arrow}0x{b:08x}")
+    for lane in range(n_lanes):
+      for i, (a, b) in enumerate(zip(self.vgpr[lane], other.vgpr[lane])):
+        if not _vals_equal(a, b): diffs.append(f"vgpr[{lane}][{i}]: 0x{a:08x}{arrow}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 = 8) -> 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}")
+
+            # Instructions with known Rust emulator bugs - sync Python to Rust after execution
+            # v_div_scale/v_div_fixup: Rust has different VCC handling
+            # v_cvt_f16_f32: Rust clears high 16 bits, but hardware (and Python) preserves them
+            sync_after = any(x in inst_str for x in ('v_div_scale_f32', 'v_div_scale_f64', 'v_div_fixup_f32', 'v_div_fixup_f64',
+                                                      'v_cvt_f16_f32'))
+            diffs = rust_before.diff(python_before, n_lanes)
+            if diffs:
+              trace_lines = []
+              for idx, (s, pc, d, rb, pb) in enumerate(trace):
+                trace_lines.append(f"    step {s}: PC={pc:3d} {d}")
+                if idx < len(trace) - 1:
+                  next_rb, next_pb = trace[idx + 1][3:5]
+                  rust_diffs = rb.diff(next_rb, n_lanes, "->")
+                  python_diffs = pb.diff(next_pb, n_lanes, "->")
+                  if rust_diffs: trace_lines.append(f"             rust:   {', '.join(rust_diffs[:5])}")
+                  if python_diffs: trace_lines.append(f"             python: {', '.join(python_diffs[:5])}")
+                  elif rust_diffs: trace_lines.append(f"             python: (no changes)")
+                else:
+                  # Last traced instruction - compare with current state
+                  rust_diffs = rb.diff(rust_before, n_lanes, "->")
+                  python_diffs = pb.diff(python_before, n_lanes, "->")
+                  if rust_diffs: trace_lines.append(f"             rust:   {', '.join(rust_diffs[:5])}")
+                  if python_diffs: trace_lines.append(f"             python: {', '.join(python_diffs[:5])}")
+                  elif rust_diffs: trace_lines.append(f"             python: (no changes)")
+              trace_str = "\n".join(trace_lines)
+              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ (rust vs python):\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:
+              # Rust returns 1 for unsupported instructions - skip test
+              if rust_result == 1 and python_result == 0:
+                raise unittest.SkipTest(f"Rust emulator doesn't support instruction: {inst_str}")
+              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
+
+            # Sync Python state to Rust after instructions with known Rust emulator differences
+            if sync_after:
+              rust_after = rust.get_snapshot()
+              for i in range(128): python.set_sgpr(i, rust_after.sgpr[i])
+              for lane in range(n_lanes):
+                for i in range(256): python.set_vgpr(lane, i, rust_after.vgpr[lane][i])
+              assert python.state is not None
+              python.state.pc, python.state.scc, python.state.vcc, python.state.exec_mask = rust_after.pc, rust_after.scc, rust_after.vcc, rust_after.exec_mask
+
+            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."""
+  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."""
+  # 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
+
+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 ops - consolidated tests covering key instruction patterns
+  def test_unary_ops(self): self._test_kernel(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu().exp().log().sqrt().reciprocal())
+  def test_binary_ops(self): self._test_kernel(lambda T: (T([1.0, 2.0]) + T([3.0, 4.0])) * T([0.5, 0.5]) - T([1.0, 1.0]))
+  def test_trig(self): self._test_kernel(lambda T: T([0.1, 1.0, 3.14, -1.0]*8).sin() + T([0.1, 1.0, 3.14, -1.0]*8).cos())
+  def test_compare(self): self._test_kernel(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
+  def test_bitwise(self): self._test_kernel(lambda T: (T([0xF0, 0x0F, 0xFF]*11).int() & T([0x0F, 0x0F, 0x00]*11).int()) | T([1]*33).int())
+  def test_int_ops(self): self._test_kernel(lambda T: ((T.empty(64).int() + T.empty(64).int()) * T.empty(64).int()).float())
+
+  # Reductions
+  def test_reduce(self): self._test_kernel(lambda T: T.empty(64).sum() + T.empty(64).max())
+  def test_argmax(self): self._test_kernel(lambda T: T.empty(64).argmax())
+
+  # Matmul
+  def test_gemm(self): self._test_kernel(lambda T: T.empty(8, 8) @ T.empty(8, 8), max_steps=100000)
+  @unittest.skip("Rust emulator crashes on this kernel (assertion failure in thread.rs)")
+  def test_gemm_fp16(self): self._test_kernel(lambda T: T.empty(16, 16).half() @ T.empty(16, 16).half(), 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_memory(self): self._test_kernel(lambda T: T.empty(4, 4).permute(1, 0).contiguous() + T.empty(4, 1).expand(4, 4))
+
+  # Cast ops
+  def test_cast(self): self._test_kernel(lambda T: T.empty(32).half().float() + T.empty(32).int().float())
+
+  # Pooling - regression for VCC wave32 mode
+  def test_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4)) + T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4)))
+
+  # Convolution
+  def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 2, 8, 8).conv2d(T.empty(2, 2, 3, 3)), max_steps=50000)
+
+  # Regression tests
+  def test_topk(self): self._test_kernel(lambda T: T.empty(64).topk(3)[0])
+  def test_interpolate(self): self._test_kernel(lambda T: T.empty(1,2,16,16).relu().cast('uint8').interpolate((8,8), mode="linear"))
+  def test_index_int64(self):
+    from tinygrad import dtypes
+    self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
+  def test_gelu(self): self._test_kernel(lambda T: T.empty(32, 32).gelu())
+  def test_cross_entropy(self):
+    import numpy as np
+    np.random.seed(0)
+    classes = np.random.randint(0, 10, (16,), dtype=np.int32).tolist()
+    x_np = np.random.randn(16, 10).astype(np.float32)
+    self._test_kernel(lambda T: (T(x_np.tolist()).reshape(16,10) + 0).cross_entropy((T(classes).int().reshape(16) + 0)))
+  def test_isinf(self): self._test_kernel(lambda T: T([float('-inf'), 0., float('inf'), 1.1]*8).isinf())
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/amd/test/test_formats.py

@@ -0,0 +1,407 @@
+#!/usr/bin/env python3
+"""Test MUBUF, MTBUF, MIMG, EXP, DS formats against LLVM."""
+import unittest
+from extra.assembly.amd.autogen.rdna3.ins import *
+from extra.assembly.amd.dsl import encode_src, RawImm
+from extra.assembly.amd.asm import detect_format
+
+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.amd.dsl 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.amd.dsl 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.amd.dsl 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
+
+
+class TestDetectFormat(unittest.TestCase):
+  """Test detect_format uses encoding from autogen classes."""
+
+  def test_detect_sopp(self):
+    self.assertEqual(detect_format(s_endpgm().to_bytes()), SOPP)
+    self.assertEqual(detect_format(s_nop(0).to_bytes()), SOPP)
+    self.assertEqual(detect_format(s_barrier().to_bytes()), SOPP)
+
+  def test_detect_sop1(self):
+    self.assertEqual(detect_format(s_mov_b32(s[0], 0).to_bytes()), SOP1)
+    self.assertEqual(detect_format(s_mov_b64(s[0:1], 0).to_bytes()), SOP1)
+
+  def test_detect_sop2(self):
+    self.assertEqual(detect_format(s_add_u32(s[0], s[1], s[2]).to_bytes()), SOP2)
+    self.assertEqual(detect_format(s_mul_i32(s[0], s[1], s[2]).to_bytes()), SOP2)
+
+  def test_detect_sopc(self):
+    self.assertEqual(detect_format(s_cmp_eq_i32(s[0], s[1]).to_bytes()), SOPC)
+
+  def test_detect_sopk(self):
+    self.assertEqual(detect_format(s_movk_i32(s[0], 0x1234).to_bytes()), SOPK)
+
+  def test_detect_vop1(self):
+    self.assertEqual(detect_format(v_mov_b32_e32(v[0], 0).to_bytes()), VOP1)
+    self.assertEqual(detect_format(v_rcp_f32_e32(v[0], v[1]).to_bytes()), VOP1)
+
+  def test_detect_vop2(self):
+    self.assertEqual(detect_format(v_add_f32_e32(v[0], v[1], v[2]).to_bytes()), VOP2)
+    self.assertEqual(detect_format(v_mul_f32_e32(v[0], v[1], v[2]).to_bytes()), VOP2)
+
+  def test_detect_vopc(self):
+    self.assertEqual(detect_format(v_cmp_eq_f32_e32(v[0], v[1]).to_bytes()), VOPC)
+    self.assertEqual(detect_format(v_cmp_lt_i32_e32(v[0], v[1]).to_bytes()), VOPC)
+
+  def test_detect_vop3(self):
+    self.assertEqual(detect_format(v_add_f32_e64(v[0], v[1], v[2]).to_bytes()), VOP3)
+    self.assertEqual(detect_format(v_fma_f32(v[0], v[1], v[2], v[3]).to_bytes()), VOP3)
+
+  def test_detect_vop3p(self):
+    self.assertEqual(detect_format(VOP3P(VOP3POp.V_PK_ADD_F16, v[0], v[1], v[2], v[3]).to_bytes()), VOP3P)
+
+  def test_detect_smem(self):
+    self.assertEqual(detect_format(s_load_b32(s[0], s[2:3], 0).to_bytes()), SMEM)
+    self.assertEqual(detect_format(s_load_b64(s[0:1], s[2:3], s[5]).to_bytes()), SMEM)
+
+  def test_detect_ds(self):
+    self.assertEqual(detect_format(ds_load_b32(v[0], v[1]).to_bytes()), DS)
+    self.assertEqual(detect_format(ds_store_b32(v[0], v[1]).to_bytes()), DS)
+
+  def test_detect_flat(self):
+    self.assertEqual(detect_format(global_load_b32(v[0], v[1:3], RawImm(124)).to_bytes()), FLAT)
+    self.assertEqual(detect_format(global_store_b32(v[0:2], v[2], RawImm(124)).to_bytes()), FLAT)
+
+  def test_detect_mubuf(self):
+    self.assertEqual(detect_format(buffer_load_b32(v[0], v[1], s[0:4], s[5]).to_bytes()), MUBUF)
+
+  def test_detect_mtbuf(self):
+    self.assertEqual(detect_format(tbuffer_load_format_x(v[0], v[1], s[0:4], s[5], format=22).to_bytes()), MTBUF)
+
+  def test_detect_mimg(self):
+    self.assertEqual(detect_format(image_load(v[0:4], v[4:6], s[0:8], dmask=0xf, dim=1).to_bytes()), MIMG)
+
+  def test_detect_exp(self):
+    self.assertEqual(detect_format(EXP(en=0xf, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[2], vsrc3=v[3]).to_bytes()), EXP)
+
+  def test_detect_vopd(self):
+    inst = VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=0, srcy0=0)
+    self.assertEqual(detect_format(inst.to_bytes()), VOPD)
+
+  def test_detect_vinterp(self):
+    inst = VINTERP(VINTERPOp.V_INTERP_P10_F32, vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
+    self.assertEqual(detect_format(inst.to_bytes()), VINTERP)
+
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/amd/test/test_handwritten.py

@@ -0,0 +1,178 @@
+# 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.amd.autogen.rdna3.ins import *
+from extra.assembly.amd.dsl import Inst
+from extra.assembly.amd.asm import asm
+from extra.assembly.amd.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/amd/test/test_integration.py

@@ -0,0 +1,330 @@
+#!/usr/bin/env python3
+"""Integration test: round-trip RDNA3 assembly through AMD toolchain."""
+import unittest, re, io, sys, subprocess
+from extra.assembly.amd.autogen.rdna3.ins import *
+from extra.assembly.amd.asm import waitcnt, asm
+from extra.assembly.amd.test.helpers import get_llvm_mc
+
+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))
+
+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")
+
+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)."""
+    def get_llvm_encoding(instr: str) -> str:
+      result = subprocess.run([get_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}")
+
+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/amd/test/test_llvm.py

@@ -0,0 +1,206 @@
+#!/usr/bin/env python3
+"""Test RDNA3/RDNA4 assembler/disassembler against LLVM test vectors."""
+import unittest, re, subprocess
+from tinygrad.helpers import fetch
+from extra.assembly.amd.test.helpers import get_llvm_mc
+
+LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/main/llvm/test/MC/AMDGPU"
+
+RDNA3_TEST_FILES = {
+  'sop1': 'gfx11_asm_sop1.s', 'sop2': 'gfx11_asm_sop2.s', 'sopp': 'gfx11_asm_sopp.s', 'sopk': 'gfx11_asm_sopk.s', 'sopc': 'gfx11_asm_sopc.s',
+  'vop1': 'gfx11_asm_vop1.s', 'vop2': 'gfx11_asm_vop2.s', 'vopc': 'gfx11_asm_vopc.s', 'vop3': 'gfx11_asm_vop3.s', 'vop3p': 'gfx11_asm_vop3p.s',
+  'vinterp': 'gfx11_asm_vinterp.s', 'vopd': 'gfx11_asm_vopd.s', 'vopcx': 'gfx11_asm_vopcx.s',
+  'vop3_from_vop1': 'gfx11_asm_vop3_from_vop1.s', 'vop3_from_vop2': 'gfx11_asm_vop3_from_vop2.s',
+  'vop3_from_vopc': 'gfx11_asm_vop3_from_vopc.s', 'vop3_from_vopcx': 'gfx11_asm_vop3_from_vopcx.s',
+  'ds': 'gfx11_asm_ds.s', 'smem': 'gfx11_asm_smem.s', 'flat': 'gfx11_asm_flat.s',
+  'mubuf': 'gfx11_asm_mubuf.s', 'mtbuf': 'gfx11_asm_mtbuf.s', 'mimg': 'gfx11_asm_mimg.s', 'mimg_features': 'gfx11_asm_mimg_features.s', 'ldsdir': 'gfx11_asm_ldsdir.s',
+  'exp': 'gfx11_asm_exp.s', 'wmma': 'gfx11_asm_wmma.s',
+  'vop3_features': 'gfx11_asm_vop3_features.s', 'vop3p_features': 'gfx11_asm_vop3p_features.s', 'vopd_features': 'gfx11_asm_vopd_features.s',
+  'vop3_alias': 'gfx11_asm_vop3_alias.s', 'vop3p_alias': 'gfx11_asm_vop3p_alias.s', 'vopc_alias': 'gfx11_asm_vopc_alias.s',
+  'vopcx_alias': 'gfx11_asm_vopcx_alias.s', 'vinterp_alias': 'gfx11_asm_vinterp_alias.s',
+  'smem_alias': 'gfx11_asm_smem_alias.s', 'mubuf_alias': 'gfx11_asm_mubuf_alias.s', 'mtbuf_alias': 'gfx11_asm_mtbuf_alias.s',
+}
+
+RDNA4_TEST_FILES = {
+  'sop1': 'gfx12_asm_sop1.s', 'sop2': 'gfx12_asm_sop2.s', 'sop2_alias': 'gfx12_asm_sop2_alias.s',
+  'sopp': 'gfx12_asm_sopp.s', 'sopk': 'gfx12_asm_sopk.s', 'sopk_alias': 'gfx12_asm_sopk_alias.s', 'sopc': 'gfx12_asm_sopc.s',
+  'vop1': 'gfx12_asm_vop1.s', 'vop2': 'gfx12_asm_vop2.s', 'vop2_aliases': 'gfx12_asm_vop2_aliases.s',
+  'vopc': 'gfx12_asm_vopc.s', 'vopcx': 'gfx12_asm_vopcx.s',
+  'vop3': 'gfx12_asm_vop3.s', 'vop3_aliases': 'gfx12_asm_vop3_aliases.s', 'vop3c': 'gfx12_asm_vop3c.s', 'vop3cx': 'gfx12_asm_vop3cx.s',
+  'vop3p': 'gfx12_asm_vop3p.s', 'vop3p_aliases': 'gfx12_asm_vop3p_aliases.s', 'vop3p_features': 'gfx12_asm_vop3p_features.s',
+  'vopd': 'gfx12_asm_vopd.s', 'vopd_features': 'gfx12_asm_vopd_features.s',
+  'vop3_from_vop1': 'gfx12_asm_vop3_from_vop1.s', 'vop3_from_vop2': 'gfx12_asm_vop3_from_vop2.s',
+  'ds': 'gfx12_asm_ds.s', 'ds_alias': 'gfx12_asm_ds_alias.s', 'smem': 'gfx12_asm_smem.s',
+  'vflat': 'gfx12_asm_vflat.s', 'vflat_alias': 'gfx12_asm_vflat_alias.s',
+  'vglobal': 'gfx12_asm_vflat.s', 'vglobal_alias': 'gfx12_asm_vflat_alias.s',  # global instructions in vflat files
+  'vscratch': 'gfx12_asm_vflat.s',  # scratch instructions in vflat file
+  'vbuffer_mubuf': 'gfx12_asm_vbuffer_mubuf.s', 'vbuffer_mubuf_alias': 'gfx12_asm_vbuffer_mubuf_alias.s',
+  'vbuffer_mtbuf': 'gfx12_asm_vbuffer_mtbuf.s', 'vbuffer_mtbuf_alias': 'gfx12_asm_vbuffer_mtbuf_alias.s',
+  'vimage': 'gfx12_asm_vimage.s', 'vimage_alias': 'gfx12_asm_vimage_alias.s', 'vsample': 'gfx12_asm_vsample.s',
+  'vdsdir': 'gfx12_asm_vdsdir.s', 'vdsdir_alias': 'gfx12_asm_vdsdir_alias.s',
+  'exp': 'gfx12_asm_exp.s', 'wmma_w32': 'gfx12_asm_wmma_w32.s', 'wmma_w64': 'gfx12_asm_wmma_w64.s',
+  'global_load_tr': 'gfx12_asm_global_load_tr.s',
+  # NOTE: 'features' (gfx12_asm_features.s) tests DPP instruction variants which require separate format decoders
+}
+
+def parse_llvm_tests(text: str, gfx_prefix: str) -> list[tuple[str, bytes]]:
+  """Parse LLVM test format into (asm, expected_bytes) pairs."""
+  tests, lines = [], text.split('\n')
+  pattern = rf'(?:{gfx_prefix}|W32|W64)[^:]*:.*?encoding:\s*\[(.*?)\]'
+  pattern2 = rf'(?:{gfx_prefix}|W32|W64)[^:]*:\s*\[(0x[0-9a-fA-F,x\s]+)\]'
+  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))):
+      if m := re.search(pattern, lines[j]):
+        hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
+      elif m := re.search(pattern2, lines[j]):
+        hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
+      else: continue
+      if hex_bytes:
+        try: tests.append((asm_text, bytes.fromhex(hex_bytes)))
+        except ValueError: pass
+      break
+  return tests
+
+def compile_asm_batch(instrs: list[str], mcpu: str, mattr: str = '+real-true16,+wavefrontsize32') -> list[bytes]:
+  """Compile multiple instructions with a single llvm-mc call."""
+  if not instrs: return []
+  result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', f'-mattr={mattr}', '-show-encoding'],
+                          input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True, timeout=30)
+  if result.returncode != 0: raise RuntimeError(f"llvm-mc failed: {result.stderr.strip()}")
+  results = []
+  for line in result.stdout.split('\n'):
+    if 'encoding:' not in line: continue
+    enc = line.split('encoding:')[1].strip()
+    if enc.startswith('[') and enc.endswith(']'):
+      results.append(bytes.fromhex(enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')))
+  if len(results) != len(instrs): raise RuntimeError(f"expected {len(instrs)} encodings, got {len(results)}")
+  return results
+
+def matches_encoding(data: bytes, fmt) -> bool:
+  """Check if instruction bytes match format's expected encoding bits."""
+  if not hasattr(fmt, '_encoding') or fmt._encoding is None: return True
+  bf, expected = fmt._encoding
+  val = int.from_bytes(data[:fmt._size()], 'little')
+  return ((val >> bf.lo) & bf.mask()) == expected
+
+class TestLLVMBase(unittest.TestCase):
+  """Base class for LLVM assembler tests."""
+  tests: dict[str, list[tuple[str, bytes]]] = {}
+  formats: dict[str, type] = {}
+  gfx_prefix: str = ""
+  mcpu: str = ""
+  arch_name: str = ""
+
+  @classmethod
+  def _load_tests(cls, test_files: dict[str, str]):
+    for name, filename in test_files.items():
+      try:
+        data = fetch(f"{LLVM_BASE}/{filename}").read_bytes()
+        cls.tests[name] = parse_llvm_tests(data.decode('utf-8', errors='ignore'), cls.gfx_prefix)
+      except Exception as e:
+        print(f"Warning: couldn't fetch {filename}: {e}")
+        cls.tests[name] = []
+
+  def _test_disasm(self, name: str):
+    """Test decoding instructions and verify disassembly produces correct bytes."""
+    if name not in self.tests or not self.tests[name]: self.skipTest(f"No test data for {name}")
+    fmt_cls = self.formats.get(name)
+    if fmt_cls is None: self.skipTest(f"No format class for {name}")
+
+    # Determine wave size from test name (w64 = wave64, otherwise wave32)
+    wave_size = 64 if 'w64' in name else 32
+    mattr = f'+real-true16,+wavefrontsize{wave_size}'
+
+    to_test: list[tuple[str, bytes, str | None, str | None]] = []
+    for asm_text, data in self.tests.get(name, []):
+      if len(data) > fmt_cls._size(): continue
+      if not matches_encoding(data, fmt_cls): continue
+      try:
+        decoded = fmt_cls.from_bytes(data)
+        if decoded.to_bytes()[:len(data)] != data:
+          to_test.append((asm_text, data, None, "decode roundtrip failed"))
+          continue
+        to_test.append((asm_text, data, decoded.disasm(wave_size), None))
+      except Exception as e:
+        to_test.append((asm_text, data, None, f"exception: {e}"))
+
+    disasm_strs = [(i, t[2]) for i, t in enumerate(to_test) if t[2] is not None]
+    llvm_map = {}
+    if disasm_strs:
+      llvm_results = compile_asm_batch([s for _, s in disasm_strs], self.mcpu, mattr)
+      llvm_map = {i: llvm_results[j] for j, (i, _) in enumerate(disasm_strs)}
+
+    passed, failed, failures = 0, 0, []
+    for idx, (asm_text, data, disasm_str, error) in enumerate(to_test):
+      if error:
+        failed += 1; failures.append(f"{error} for {data.hex()}")
+      elif disasm_str is not None and idx in llvm_map:
+        llvm_bytes = llvm_map[idx]
+        if llvm_bytes == data: passed += 1
+        else: failed += 1; failures.append(f"'{disasm_str}': expected={data.hex()} got={llvm_bytes.hex()}")
+
+    print(f"{self.arch_name} {name.upper()} disasm: {passed} passed, {failed} failed")
+    if failures[:5]: print("  " + "\n  ".join(failures[:5]))
+    self.assertGreater(passed, 0, f"No tests passed for {name}")
+
+class TestLLVMRDNA3(TestLLVMBase):
+  """Test RDNA3 assembler against LLVM test vectors."""
+  gfx_prefix, mcpu, arch_name = "GFX11", "gfx1100", "RDNA3"
+
+  @classmethod
+  def setUpClass(cls):
+    from extra.assembly.amd.autogen.rdna3.ins import SOP1, SOP2, SOPC, SOPK, SOPP, VOP1, VOP2, VOP3, VOP3P, VOPC, VOPD, VINTERP, DS, SMEM, FLAT, MUBUF, MTBUF, MIMG, LDSDIR, EXP
+    cls.formats = {
+      'sop1': SOP1, 'sop2': SOP2, 'sopc': SOPC, 'sopk': SOPK, 'sopp': SOPP,
+      'vop1': VOP1, 'vop2': VOP2, 'vopc': VOPC, 'vopcx': VOPC, 'vop3': VOP3, 'vop3p': VOP3P,
+      'vinterp': VINTERP, 'vopd': VOPD, 'ds': DS, 'smem': SMEM, 'flat': FLAT,
+      'mubuf': MUBUF, 'mtbuf': MTBUF, 'mimg': MIMG, 'mimg_features': MIMG, 'wmma': VOP3P, 'ldsdir': LDSDIR, 'exp': EXP,
+      'vop3_from_vop1': VOP3, 'vop3_from_vop2': VOP3, 'vop3_from_vopc': VOP3, 'vop3_from_vopcx': VOP3,
+      'vop3_features': VOP3, 'vop3p_features': VOP3P, 'vopd_features': VOPD,
+      'vop3_alias': VOP3, 'vop3p_alias': VOP3P, 'vopc_alias': VOPC, 'vopcx_alias': VOPC,
+      'vinterp_alias': VINTERP, 'smem_alias': SMEM, 'mubuf_alias': MUBUF, 'mtbuf_alias': MTBUF,
+    }
+    cls._load_tests(RDNA3_TEST_FILES)
+
+class TestLLVMRDNA4(TestLLVMBase):
+  """Test RDNA4 assembler against LLVM test vectors."""
+  gfx_prefix, mcpu, arch_name = "GFX12", "gfx1200", "RDNA4"
+
+  @classmethod
+  def setUpClass(cls):
+    import extra.assembly.amd.autogen.rdna4.ins as rdna4
+    get = lambda n: getattr(rdna4, n, None)
+    cls.formats = {
+      'sop1': get('SOP1'), 'sop2': get('SOP2'), 'sop2_alias': get('SOP2'), 'sopc': get('SOPC'),
+      'sopk': get('SOPK'), 'sopk_alias': get('SOPK'), 'sopp': get('SOPP'),
+      'vop1': get('VOP1'), 'vop2': get('VOP2'), 'vop2_aliases': get('VOP2'), 'vopc': get('VOPC'), 'vopcx': get('VOPC'),
+      'vop3': get('VOP3'), 'vop3_aliases': get('VOP3'), 'vop3c': get('VOP3'), 'vop3cx': get('VOP3'),
+      'vop3p': get('VOP3P'), 'vop3p_aliases': get('VOP3P'), 'vop3p_features': get('VOP3P'),
+      'vopd': get('VOPD'), 'vopd_features': get('VOPD'),
+      'vop3_from_vop1': get('VOP3'), 'vop3_from_vop2': get('VOP3'),
+      'ds': get('VDS'), 'ds_alias': get('VDS'), 'smem': get('SMEM'), 'vinterp': get('VINTERP'), 'exp': get('VEXPORT'),
+      'vbuffer_mubuf': get('VBUFFER'), 'vbuffer_mubuf_alias': get('VBUFFER'),
+      'vbuffer_mtbuf': get('VBUFFER'), 'vbuffer_mtbuf_alias': get('VBUFFER'),
+      'vdsdir': get('VDSDIR'), 'vdsdir_alias': get('VDSDIR'),
+      'vflat': get('VFLAT'), 'vflat_alias': get('VFLAT'),
+      'vglobal': get('VGLOBAL'), 'vglobal_alias': get('VGLOBAL'),
+      'vscratch': get('VSCRATCH'),
+      'vimage': get('VIMAGE'), 'vimage_alias': get('VIMAGE'), 'vsample': get('VSAMPLE'),
+      'wmma_w32': get('VOP3P'), 'wmma_w64': get('VOP3P'),
+      'global_load_tr': get('VGLOBAL'),
+    }
+    cls._load_tests(RDNA4_TEST_FILES)
+
+# Generate test methods dynamically
+def _make_test(name):
+  def test(self): self._test_disasm(name)
+  return test
+
+for name in RDNA3_TEST_FILES: setattr(TestLLVMRDNA3, f'test_{name}_disasm', _make_test(name))
+for name in RDNA4_TEST_FILES: setattr(TestLLVMRDNA4, f'test_{name}_disasm', _make_test(name))
+
+if __name__ == "__main__": unittest.main()

extra/assembly/amd/test/test_mockgpu_invalid.py

@@ -0,0 +1,55 @@
+#!/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.assertTrue("NotImplementedError" in result.stderr or "ValueError" in result.stderr,
+                    f"expected NotImplementedError or ValueError in stderr")
+    # Should exit immediately, not wait for the full timeout
+    self.assertLess(elapsed, 9.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/amd/test/test_pcode.py

@@ -0,0 +1,404 @@
+#!/usr/bin/env python3
+"""Tests for the RDNA3 pseudocode DSL."""
+import unittest
+from extra.assembly.amd.pcode import (Reg, TypedView, SliceProxy, MASK32, MASK64,
+                                       _f32, _i32, _f16, _i16, f32_to_f16, _isnan, _bf16, _ibf16, bf16_to_f32, f32_to_bf16,
+                                       BYTE_PERMUTE, v_sad_u8, v_msad_u8)
+from extra.assembly.amd.pdf import compile_pseudocode, _expr
+from extra.assembly.amd.test.helpers import ExecContext
+from extra.assembly.amd.autogen.rdna3.gen_pcode import _VOP3SDOp_V_DIV_SCALE_F32, _VOPCOp_V_CMP_CLASS_F32
+
+class TestReg(unittest.TestCase):
+  def test_u32_read(self):
+    r = Reg(0xDEADBEEF)
+    self.assertEqual(int(r.u32), 0xDEADBEEF)
+
+  def test_u32_write(self):
+    r = Reg(0)
+    r.u32 = 0x12345678
+    self.assertEqual(r._val, 0x12345678)
+
+  def test_f32_read(self):
+    r = Reg(0x40400000)  # 3.0f
+    self.assertAlmostEqual(float(r.f32), 3.0)
+
+  def test_f32_write(self):
+    r = Reg(0)
+    r.f32 = 3.0
+    self.assertEqual(r._val, 0x40400000)
+
+  def test_i32_signed(self):
+    r = Reg(0xFFFFFFFF)  # -1 as signed
+    self.assertEqual(int(r.i32), -1)
+
+  def test_u64(self):
+    r = Reg(0xDEADBEEFCAFEBABE)
+    self.assertEqual(int(r.u64), 0xDEADBEEFCAFEBABE)
+
+  def test_f64(self):
+    r = Reg(0x4008000000000000)  # 3.0 as f64
+    self.assertAlmostEqual(float(r.f64), 3.0)
+
+class TestTypedView(unittest.TestCase):
+  def test_bit_slice(self):
+    r = Reg(0xDEADBEEF)
+    # Slices return SliceProxy which supports .u32, .u16 etc (matching pseudocode like S1.u32[1:0].u32)
+    self.assertEqual(r.u32[7:0].u32, 0xEF)
+    self.assertEqual(r.u32[15:8].u32, 0xBE)
+    self.assertEqual(r.u32[23:16].u32, 0xAD)
+    self.assertEqual(r.u32[31:24].u32, 0xDE)
+    # Also works with int() for arithmetic
+    self.assertEqual(int(r.u32[7:0]), 0xEF)
+
+  def test_single_bit_read(self):
+    r = Reg(0b11010101)
+    self.assertEqual(r.u32[0], 1)
+    self.assertEqual(r.u32[1], 0)
+    self.assertEqual(r.u32[2], 1)
+    self.assertEqual(r.u32[3], 0)
+
+  def test_single_bit_write(self):
+    r = Reg(0)
+    r.u32[5] = 1
+    r.u32[3] = 1
+    self.assertEqual(r._val, 0b00101000)
+
+  def test_nested_bit_access(self):
+    # S0.u32[S1.u32[4:0]] - access bit at position from another register
+    s0 = Reg(0b11010101)
+    s1 = Reg(3)
+    bit_pos = s1.u32[4:0]  # SliceProxy, int value = 3
+    bit_val = s0.u32[int(bit_pos)]  # bit 3 of s0 = 0
+    self.assertEqual(int(bit_pos), 3)
+    self.assertEqual(bit_val, 0)
+
+  def test_arithmetic(self):
+    r1 = Reg(0x40400000)  # 3.0f
+    r2 = Reg(0x40800000)  # 4.0f
+    result = r1.f32 + r2.f32
+    self.assertAlmostEqual(result, 7.0)
+
+  def test_comparison(self):
+    r1 = Reg(5)
+    r2 = Reg(3)
+    self.assertTrue(r1.u32 > r2.u32)
+    self.assertFalse(r1.u32 < r2.u32)
+    self.assertTrue(r1.u32 != r2.u32)
+
+class TestSliceProxy(unittest.TestCase):
+  def test_slice_read(self):
+    r = Reg(0x56781234)
+    self.assertEqual(r[15:0].u16, 0x1234)
+    self.assertEqual(r[31:16].u16, 0x5678)
+
+  def test_slice_write(self):
+    r = Reg(0)
+    r[15:0].u16 = 0x1234
+    r[31:16].u16 = 0x5678
+    self.assertEqual(r._val, 0x56781234)
+
+  def test_slice_f16(self):
+    r = Reg(0)
+    r[15:0].f16 = 3.0
+    self.assertAlmostEqual(_f16(r._val & 0xffff), 3.0, places=2)
+
+class TestCompiler(unittest.TestCase):
+  def test_ternary(self):
+    result = _expr("a > b ? 1 : 0")
+    self.assertIn("if", result)
+    self.assertIn("else", result)
+
+  def test_type_prefix_strip(self):
+    self.assertEqual(_expr("1'0U"), "0")
+    self.assertEqual(_expr("32'1"), "1")
+    self.assertEqual(_expr("16'0xFFFF"), "0xFFFF")
+
+  def test_suffix_strip(self):
+    self.assertEqual(_expr("0ULL"), "0")
+    self.assertEqual(_expr("1LL"), "1")
+    self.assertEqual(_expr("5U"), "5")
+    self.assertEqual(_expr("3.14F"), "3.14")
+
+  def test_boolean_ops(self):
+    self.assertIn("and", _expr("a && b"))
+    self.assertIn("or", _expr("a || b"))
+    self.assertIn("!=", _expr("a <> b"))
+
+  def test_pack16(self):
+    result = _expr("{ a, b }")
+    self.assertIn("_pack", result)
+
+  def test_type_cast_strip(self):
+    self.assertEqual(_expr("64'U(x)"), "(x)")
+    self.assertEqual(_expr("32'I(y)"), "(y)")
+
+class TestExecContext(unittest.TestCase):
+  def test_float_add(self):
+    ctx = ExecContext(s0=0x40400000, s1=0x40800000)  # 3.0f, 4.0f
+    ctx.D0.f32 = ctx.S0.f32 + ctx.S1.f32
+    self.assertAlmostEqual(_f32(ctx.D0._val), 7.0)
+
+  def test_float_mul(self):
+    ctx = ExecContext(s0=0x40400000, s1=0x40800000)  # 3.0f, 4.0f
+    ctx.run("D0.f32 = S0.f32 * S1.f32")
+    self.assertAlmostEqual(_f32(ctx.D0._val), 12.0)
+
+  def test_scc_comparison(self):
+    ctx = ExecContext(s0=42, s1=42)
+    ctx.run("SCC = S0.u32 == S1.u32")
+    self.assertEqual(ctx.SCC._val, 1)
+
+  def test_scc_comparison_false(self):
+    ctx = ExecContext(s0=42, s1=43)
+    ctx.run("SCC = S0.u32 == S1.u32")
+    self.assertEqual(ctx.SCC._val, 0)
+
+  def test_ternary(self):
+    code = compile_pseudocode("D0.u32 = S0.u32 > S1.u32 ? 1'1U : 1'0U")
+    ctx = ExecContext(s0=5, s1=3)
+    ctx.run(code)
+    self.assertEqual(ctx.D0._val, 1)
+
+  def test_pack(self):
+    code = compile_pseudocode("D0 = { S1[15:0].u16, S0[15:0].u16 }")
+    ctx = ExecContext(s0=0x1234, s1=0x5678)
+    ctx.run(code)
+    self.assertEqual(ctx.D0._val, 0x56781234)
+
+  def test_tmp_with_typed_access(self):
+    code = compile_pseudocode("""tmp = S0.u32 + S1.u32
+D0.u32 = tmp.u32""")
+    ctx = ExecContext(s0=100, s1=200)
+    ctx.run(code)
+    self.assertEqual(ctx.D0._val, 300)
+
+  def test_s_add_u32_pattern(self):
+    # Real pseudocode pattern from S_ADD_U32
+    code = compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
+SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
+D0.u32 = tmp.u32""")
+    # Test overflow case
+    ctx = ExecContext(s0=0xFFFFFFFF, s1=0x00000001)
+    ctx.run(code)
+    self.assertEqual(ctx.D0._val, 0)  # Wraps to 0
+    self.assertEqual(ctx.SCC._val, 1)  # Carry set
+
+  def test_s_add_u32_no_overflow(self):
+    code = compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
+SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
+D0.u32 = tmp.u32""")
+    ctx = ExecContext(s0=100, s1=200)
+    ctx.run(code)
+    self.assertEqual(ctx.D0._val, 300)
+    self.assertEqual(ctx.SCC._val, 0)  # No carry
+
+  def test_vcc_lane_read(self):
+    ctx = ExecContext(vcc=0b1010, lane=1)
+    # Lane 1 is set
+    self.assertEqual(ctx.VCC.u64[1], 1)
+    self.assertEqual(ctx.VCC.u64[2], 0)
+
+  def test_vcc_lane_write(self):
+    ctx = ExecContext(vcc=0, lane=0)
+    ctx.VCC.u64[3] = 1
+    ctx.VCC.u64[1] = 1
+    self.assertEqual(ctx.VCC._val, 0b1010)
+
+  def test_for_loop(self):
+    # CTZ pattern - find first set bit
+    code = compile_pseudocode("""tmp = -1
+for i in 0 : 31 do
+  if S0.u32[i] == 1 then
+    tmp = i
+  endif
+endfor
+D0.i32 = tmp""")
+    ctx = ExecContext(s0=0b1000)  # Bit 3 is set
+    ctx.run(code)
+    self.assertEqual(ctx.D0._val & MASK32, 3)
+
+  def test_result_dict(self):
+    ctx = ExecContext(s0=5, s1=3)
+    ctx.D0.u32 = 42
+    ctx.SCC._val = 1
+    result = ctx.result()
+    self.assertEqual(result['d0'], 42)
+    self.assertEqual(result['scc'], 1)
+
+class TestPseudocodeRegressions(unittest.TestCase):
+  """Regression tests for pseudocode instruction emulation bugs."""
+
+  def test_v_div_scale_f32_vcc_always_returned(self):
+    """V_DIV_SCALE_F32 must always return VCC, even when VCC=0 (no scaling needed).
+    Bug: when VCC._val == vcc (both 0), VCC wasn't returned, so VCC bits weren't written.
+    This caused division to produce wrong results for multiple lanes."""
+    # Normal case: 1.0 / 3.0, no scaling needed, VCC should be 0
+    S0 = Reg(0x3f800000)  # 1.0
+    S1 = Reg(0x40400000)  # 3.0
+    S2 = Reg(0x3f800000)  # 1.0 (numerator)
+    D0, SCC, VCC, EXEC = Reg(0), Reg(0), Reg(0), Reg(0xffffffff)
+    result = _VOP3SDOp_V_DIV_SCALE_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
+    # Must always have VCC in result
+    self.assertIn('VCC', result, "V_DIV_SCALE_F32 must always return VCC")
+    self.assertEqual(result['VCC']._val & 1, 0, "VCC lane 0 should be 0 when no scaling needed")
+
+  def test_v_cmp_class_f32_detects_quiet_nan(self):
+    """V_CMP_CLASS_F32 must correctly identify quiet NaN vs signaling NaN.
+    Bug: isQuietNAN and isSignalNAN both used math.isnan which can't distinguish them."""
+    quiet_nan = 0x7fc00000   # quiet NaN: exponent=255, bit22=1
+    signal_nan = 0x7f800001  # signaling NaN: exponent=255, bit22=0
+    # Test quiet NaN detection (bit 1 in mask)
+    s1_quiet = 0b0000000010  # bit 1 = quiet NaN
+    S0, S1, S2, D0, SCC, VCC, EXEC = Reg(quiet_nan), Reg(s1_quiet), Reg(0), Reg(0), Reg(0), Reg(0), Reg(0xffffffff)
+    result = _VOPCOp_V_CMP_CLASS_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
+    self.assertEqual(result['D0']._val & 1, 1, "Should detect quiet NaN with quiet NaN mask")
+    # Test signaling NaN detection (bit 0 in mask)
+    s1_signal = 0b0000000001  # bit 0 = signaling NaN
+    S0, S1 = Reg(signal_nan), Reg(s1_signal)
+    result = _VOPCOp_V_CMP_CLASS_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
+    self.assertEqual(result['D0']._val & 1, 1, "Should detect signaling NaN with signaling NaN mask")
+    # Test that quiet NaN doesn't match signaling NaN mask
+    S0, S1 = Reg(quiet_nan), Reg(s1_signal)
+    result = _VOPCOp_V_CMP_CLASS_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
+    self.assertEqual(result['D0']._val & 1, 0, "Quiet NaN should not match signaling NaN mask")
+    # Test that signaling NaN doesn't match quiet NaN mask
+    S0, S1 = Reg(signal_nan), Reg(s1_quiet)
+    result = _VOPCOp_V_CMP_CLASS_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
+    self.assertEqual(result['D0']._val & 1, 0, "Signaling NaN should not match quiet NaN mask")
+
+  def test_isnan_with_typed_view(self):
+    """_isnan must work with TypedView objects, not just Python floats.
+    Bug: _isnan checked isinstance(x, float) which returned False for TypedView."""
+    nan_reg = Reg(0x7fc00000)  # quiet NaN
+    normal_reg = Reg(0x3f800000)  # 1.0
+    inf_reg = Reg(0x7f800000)  # +inf
+    self.assertTrue(_isnan(nan_reg.f32), "_isnan should return True for NaN TypedView")
+    self.assertFalse(_isnan(normal_reg.f32), "_isnan should return False for normal TypedView")
+    self.assertFalse(_isnan(inf_reg.f32), "_isnan should return False for inf TypedView")
+
+class TestBF16(unittest.TestCase):
+  """Tests for BF16 (bfloat16) support."""
+
+  def test_bf16_conversion(self):
+    """Test bf16 <-> f32 conversion."""
+    # bf16 is just the top 16 bits of f32
+    # 1.0f = 0x3f800000, bf16 = 0x3f80
+    self.assertAlmostEqual(_bf16(0x3f80), 1.0, places=2)
+    self.assertEqual(_ibf16(1.0), 0x3f80)
+    # 2.0f = 0x40000000, bf16 = 0x4000
+    self.assertAlmostEqual(_bf16(0x4000), 2.0, places=2)
+    self.assertEqual(_ibf16(2.0), 0x4000)
+    # -1.0f = 0xbf800000, bf16 = 0xbf80
+    self.assertAlmostEqual(_bf16(0xbf80), -1.0, places=2)
+    self.assertEqual(_ibf16(-1.0), 0xbf80)
+
+  def test_bf16_special_values(self):
+    """Test bf16 special values (inf, nan)."""
+    import math
+    # +inf: f32 = 0x7f800000, bf16 = 0x7f80
+    self.assertTrue(math.isinf(_bf16(0x7f80)))
+    self.assertEqual(_ibf16(float('inf')), 0x7f80)
+    # -inf: f32 = 0xff800000, bf16 = 0xff80
+    self.assertTrue(math.isinf(_bf16(0xff80)))
+    self.assertEqual(_ibf16(float('-inf')), 0xff80)
+    # NaN: quiet NaN bf16 = 0x7fc0
+    self.assertTrue(math.isnan(_bf16(0x7fc0)))
+    self.assertEqual(_ibf16(float('nan')), 0x7fc0)
+
+  def test_bf16_register_property(self):
+    """Test Reg.bf16 property."""
+    r = Reg(0)
+    r.bf16 = 3.0  # 3.0f = 0x40400000, bf16 = 0x4040
+    self.assertEqual(r._val & 0xffff, 0x4040)
+    self.assertAlmostEqual(float(r.bf16), 3.0, places=1)
+
+  def test_bf16_slice_property(self):
+    """Test SliceProxy.bf16 property."""
+    r = Reg(0x40404040)  # Two bf16 3.0 values
+    self.assertAlmostEqual(r[15:0].bf16, 3.0, places=1)
+    self.assertAlmostEqual(r[31:16].bf16, 3.0, places=1)
+
+class TestBytePermute(unittest.TestCase):
+  """Tests for BYTE_PERMUTE helper function (V_PERM_B32)."""
+
+  def test_byte_select_0_to_7(self):
+    """Test selecting bytes 0-7 from 64-bit data."""
+    # data = {s0, s1} where s0 is bytes 0-3, s1 is bytes 4-7
+    # Combined: 0x0706050403020100 (byte 0 = 0x00, byte 7 = 0x07)
+    data = 0x0706050403020100
+    for i in range(8):
+      self.assertEqual(BYTE_PERMUTE(data, i), i, f"byte {i} should be {i}")
+
+  def test_sign_extend_bytes(self):
+    """Test sign extension selectors 8-11."""
+    # sel 8: sign of byte 1 (bits 15:8)
+    # sel 9: sign of byte 3 (bits 31:24)
+    # sel 10: sign of byte 5 (bits 47:40)
+    # sel 11: sign of byte 7 (bits 63:56)
+    data = 0x8000800080008000  # All relevant bytes have sign bit set
+    self.assertEqual(BYTE_PERMUTE(data, 8), 0xff)
+    self.assertEqual(BYTE_PERMUTE(data, 9), 0xff)
+    self.assertEqual(BYTE_PERMUTE(data, 10), 0xff)
+    self.assertEqual(BYTE_PERMUTE(data, 11), 0xff)
+    data = 0x7f007f007f007f00  # No sign bits set
+    self.assertEqual(BYTE_PERMUTE(data, 8), 0x00)
+    self.assertEqual(BYTE_PERMUTE(data, 9), 0x00)
+    self.assertEqual(BYTE_PERMUTE(data, 10), 0x00)
+    self.assertEqual(BYTE_PERMUTE(data, 11), 0x00)
+
+  def test_constant_zero(self):
+    """Test selector 12 returns 0x00."""
+    self.assertEqual(BYTE_PERMUTE(0xffffffffffffffff, 12), 0x00)
+
+  def test_constant_ff(self):
+    """Test selectors >= 13 return 0xFF."""
+    for sel in [13, 14, 15, 255]:
+      self.assertEqual(BYTE_PERMUTE(0, sel), 0xff, f"sel {sel} should be 0xff")
+
+class TestSADHelpers(unittest.TestCase):
+  """Tests for V_SAD_U8 and V_MSAD_U8 helper functions."""
+
+  def test_v_sad_u8_basic(self):
+    """Test v_sad_u8 with simple values."""
+    # s0 = 0x04030201, s1 = 0x04030201 -> diff = 0 for all bytes
+    result = v_sad_u8(0x04030201, 0x04030201, 0)
+    self.assertEqual(result, 0)
+    # s0 = 0x05040302, s1 = 0x04030201 -> diff = 1+1+1+1 = 4
+    result = v_sad_u8(0x05040302, 0x04030201, 0)
+    self.assertEqual(result, 4)
+
+  def test_v_sad_u8_with_accumulator(self):
+    """Test v_sad_u8 with non-zero accumulator."""
+    # s0 = 0x05040302, s1 = 0x04030201, s2 = 100 -> 4 + 100 = 104
+    result = v_sad_u8(0x05040302, 0x04030201, 100)
+    self.assertEqual(result, 104)
+
+  def test_v_sad_u8_large_diff(self):
+    """Test v_sad_u8 with maximum byte differences."""
+    # s0 = 0xffffffff, s1 = 0x00000000 -> diff = 255*4 = 1020
+    result = v_sad_u8(0xffffffff, 0x00000000, 0)
+    self.assertEqual(result, 1020)
+
+  def test_v_msad_u8_basic(self):
+    """Test v_msad_u8 masks when reference byte is 0."""
+    # s0 = 0x10101010, s1 = 0x00000000 -> all masked, result = 0
+    result = v_msad_u8(0x10101010, 0x00000000, 0)
+    self.assertEqual(result, 0)
+    # s0 = 0x10101010, s1 = 0x01010101 -> diff = |0x10-0x01|*4 = 15*4 = 60
+    result = v_msad_u8(0x10101010, 0x01010101, 0)
+    self.assertEqual(result, 60)
+
+  def test_v_msad_u8_partial_mask(self):
+    """Test v_msad_u8 with partial masking."""
+    # s0 = 0x10101010, s1 = 0x00010001 -> bytes 1 and 3 masked
+    # diff = |0x10-0x01| + |0x10-0x01| = 15 + 15 = 30
+    result = v_msad_u8(0x10101010, 0x00010001, 0)
+    self.assertEqual(result, 30)
+
+  def test_v_msad_u8_with_accumulator(self):
+    """Test v_msad_u8 with non-zero accumulator."""
+    result = v_msad_u8(0x10101010, 0x01010101, 50)
+    self.assertEqual(result, 110)  # 60 + 50
+
+if __name__ == '__main__':
+  unittest.main()

extra/assembly/amd/test/test_pdf_parser.py

@@ -0,0 +1,150 @@
+#!/usr/bin/env python3
+"""Test that PDF parser correctly extracts format fields."""
+import unittest, os
+from extra.assembly.amd.autogen.rdna3.ins 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),
+}
+
+# Skip PDF parsing tests by default - only run with TEST_PDF_PARSER=1
+# These are slow (~5s) and only needed when regenerating autogen/
+@unittest.skipUnless(os.environ.get("TEST_PDF_PARSER"), "set TEST_PDF_PARSER=1 to run PDF parser tests")
+class TestPDFParserGenerate(unittest.TestCase):
+  """Test the PDF parser by running generate() and checking results."""
+
+  def test_pdf_parser(self):
+    """Single test that validates all PDF parser outputs."""
+    from extra.assembly.amd.dsl import generate
+    result = generate()
+
+    # test_all_formats_present
+    for fmt_name in EXPECTED_FORMATS:
+      self.assertIn(fmt_name, result["formats"], f"missing format {fmt_name}")
+
+    # test_format_count
+    self.assertEqual(len(result["formats"]), 23)
+
+    # test_no_duplicate_fields
+    for fmt_name, fields in 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}")
+
+    # test_expected_fields
+    for fmt_name, (expected_fields, has_encoding) in EXPECTED_FORMATS.items():
+      fields = {f[0] for f in 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")
+
+    # test_vopd_no_dpp16_fields
+    vopd_fields = {f[0] for f in result["formats"].get("VOPD", [])}
+    for field in ['DPP_CTRL', 'BANK_MASK', 'ROW_MASK']:
+      self.assertNotIn(field, vopd_fields, f"VOPD should not have {field}")
+
+    # test_dpp16_no_vinterp_fields
+    dpp16_fields = {f[0] for f in result["formats"].get("DPP16", [])}
+    for field in ['VDST', 'WAITEXP']:
+      self.assertNotIn(field, dpp16_fields, f"DPP16 should not have {field}")
+
+    # test_sopp_no_smem_fields
+    sopp_fields = {f[0] for f in 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/amd/test/test_rdna3_asm.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+import unittest, subprocess
+from extra.assembly.amd.autogen.rdna3.ins import *
+from extra.assembly.amd.test.helpers import get_llvm_mc
+
+def llvm_assemble(asm: str) -> bytes:
+  """Assemble using llvm-mc and return bytes."""
+  result = subprocess.run(
+    [get_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/amd/test/test_roundtrip.py

@@ -0,0 +1,246 @@
+#!/usr/bin/env python3
+"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
+import unittest, io, sys, re, subprocess, os
+from extra.assembly.amd.dsl import Inst
+from extra.assembly.amd.test.helpers import get_llvm_mc, get_llvm_objdump
+
+def compile_asm_batch(instrs: list[str], mcpu: str = 'gfx1100') -> list[bytes]:
+  """Compile multiple instructions with a single llvm-mc call."""
+  if not instrs: return []
+  result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
+                          input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True)
+  if result.returncode != 0: raise RuntimeError(f"llvm-mc batch failed: {result.stderr.strip()}")
+  encodings = []
+  for line in result.stdout.split('\n'):
+    if 'encoding:' in line:
+      enc = line.split('encoding:')[1].strip()
+      if enc.startswith('[') and enc.endswith(']'):
+        encodings.append(bytes.fromhex(enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')))
+  if len(encodings) != len(instrs): raise RuntimeError(f"expected {len(instrs)} encodings, got {len(encodings)}")
+  return encodings
+
+def compile_and_disasm_batch(instrs: list[str], mcpu: str = 'gfx1100') -> list[str]:
+  """Compile instructions with LLVM and get LLVM's disassembly."""
+  import tempfile
+  if not instrs: return []
+  src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n" + "\n".join(f"  {instr}" for instr in instrs) + "\n"
+  with tempfile.NamedTemporaryFile(suffix='.o', delete=False) as f:
+    obj_path = f.name
+  try:
+    result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', '-mattr=+real-true16,+wavefrontsize32', '-filetype=obj', '-o', obj_path],
+                            input=src, capture_output=True, text=True)
+    if result.returncode != 0: raise RuntimeError(f"llvm-mc failed: {result.stderr.strip()}")
+    result = subprocess.run([get_llvm_objdump(), '-d', f'--mcpu={mcpu}', obj_path], capture_output=True, text=True)
+    if result.returncode != 0: raise RuntimeError(f"llvm-objdump failed: {result.stderr.strip()}")
+    results: list[str] = []
+    for line in result.stdout.splitlines():
+      if '//' not in line: continue
+      instr = line.split('//')[0].strip()
+      if instr: results.append(instr)
+    return results[:len(instrs)]
+  finally:
+    os.unlink(obj_path)
+
+class TestRoundtripBase(unittest.TestCase):
+  """Base class for roundtrip tests."""
+  mcpu: str = 'gfx1100'
+  arch: str = 'rdna3'
+
+  @classmethod
+  def _get_modules(cls):
+    if cls.arch == 'rdna3':
+      from extra.assembly.amd.autogen.rdna3 import ins
+      from extra.assembly.amd.asm import detect_format, asm
+    else:
+      import extra.assembly.amd.autogen.rdna4.ins as ins
+      from extra.assembly.amd.asm import asm
+      detect_format = None  # RDNA4 uses different detection
+    return ins, detect_format, asm
+
+  def _test_kernel_roundtrip(self, op_fn):
+    """Generate kernel from op_fn, test decode -> reencode and asm(disasm()) matches LLVM."""
+    from extra.assembly.amd.test.test_compare_emulators import get_kernels_from_tinygrad
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+
+    ins, detect_format, asm = self._get_modules()
+    kernels, _, _ = get_kernels_from_tinygrad(op_fn)
+    compiler = HIPCompiler(self.mcpu)
+
+    # First pass: decode all instructions
+    decoded_instrs: list[tuple] = []
+    for ki, kernel in enumerate(kernels):
+      offset = 0
+      while offset < len(kernel.code):
+        remaining = kernel.code[offset:]
+        if len(remaining) < 4: break
+
+        # Try to detect format
+        if detect_format is not None:
+          try:
+            fmt = detect_format(remaining)
+          except ValueError:
+            decoded_instrs.append((ki, offset, None, None, None, False, "no format"))
+            offset += 4
+            continue
+        else:
+          # For RDNA4, try formats in order
+          fmt = None
+          from extra.assembly.amd.autogen.rdna4.ins import SOP1, SOP2, SOPC, SOPK, SOPP, VOP1, VOP2, VOP3, VOP3P, VOPC, VOPD, VDS, SMEM, VFLAT, VBUFFER, VIMAGE, VSAMPLE, VEXPORT, VDSDIR
+          word = int.from_bytes(remaining[:4], 'little')
+          for cls in [VOPD, VOP3P, VOP3, VDS, VFLAT, VBUFFER, VIMAGE, VSAMPLE, SMEM, VEXPORT, SOP1, SOPC, SOPP, SOPK, VOPC, VOP1, SOP2, VOP2, VDSDIR]:
+            if cls._encoding is not None:
+              bf, val = cls._encoding
+              if ((word >> bf.lo) & bf.mask()) == val:
+                fmt = cls
+                break
+          if fmt is None:
+            decoded_instrs.append((ki, offset, None, None, None, False, "no format"))
+            offset += 4
+            continue
+
+        base_size = fmt._size()
+        if len(remaining) < base_size: break
+
+        try:
+          decoded = fmt.from_bytes(remaining)
+          size = decoded.size()
+          orig_bytes = remaining[:size]
+          reencoded = decoded.to_bytes()
+          our_disasm = decoded.disasm()
+          decode_ok = reencoded == orig_bytes
+          decode_err = None if decode_ok else f"orig={orig_bytes.hex()} reenc={reencoded.hex()}"
+          decoded_instrs.append((ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err))
+        except Exception as e:
+          decoded_instrs.append((ki, offset, remaining[:base_size], None, None, False, str(e)))
+          size = base_size
+        offset += size
+
+    # Collect disasm strings for batched LLVM calls
+    asm_test_instrs: list[tuple[int, str]] = []
+    for idx, (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err) in enumerate(decoded_instrs):
+      if our_disasm is None: continue
+      if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm): continue
+      asm_test_instrs.append((idx, our_disasm))
+
+    # Batch compile for asm test
+    asm_llvm_results = compile_asm_batch([d for _, d in asm_test_instrs], self.mcpu)
+    asm_llvm_map = {idx: result for (idx, _), result in zip(asm_test_instrs, asm_llvm_results)}
+
+    # Batch compile+disasm for disasm comparison test
+    disasm_llvm_results = compile_and_disasm_batch([d for _, d in asm_test_instrs], self.mcpu)
+    disasm_llvm_map = {idx: result for (idx, _), result in zip(asm_test_instrs, disasm_llvm_results)}
+
+    # Evaluate results
+    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 idx, (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err) in enumerate(decoded_instrs):
+      if decode_ok: decode_passed += 1
+      elif decode_err == "no format": decode_skipped += 1
+      else:
+        decode_failed += 1
+        decode_failures.append(f"K{ki}@{offset}: {our_disasm}: {decode_err}")
+
+      if our_disasm is None:
+        asm_skipped += 1
+        disasm_skipped += 1
+      elif idx in asm_llvm_map:
+        llvm_bytes = asm_llvm_map[idx]
+        try:
+          our_bytes = asm(our_disasm).to_bytes()
+          if 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
+
+        if idx in disasm_llvm_map:
+          if our_disasm == disasm_llvm_map[idx]: disasm_passed += 1
+          else:
+            disasm_failed += 1
+            disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{disasm_llvm_map[idx]}'")
+        else:
+          disasm_skipped += 1
+      else:
+        asm_skipped += 1
+        disasm_skipped += 1
+
+    print(f"{self.arch.upper()} decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
+    print(f"{self.arch.upper()} asm vs llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
+    print(f"{self.arch.upper()} 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]))
+
+class TestRoundtripRDNA3(TestRoundtripBase):
+  """Roundtrip tests for RDNA3 (gfx1100)."""
+  mcpu, arch = 'gfx1100', 'rdna3'
+
+  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())
+  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])))
+  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())
+  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))
+  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())
+  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())
+  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())
+  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)))
+  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])))
+
+@unittest.skipUnless(os.environ.get("TEST_RDNA4"), "RDNA4 roundtrip tests require TEST_RDNA4=1 and gfx1200 hardware")
+class TestRoundtripRDNA4(TestRoundtripBase):
+  """Roundtrip tests for RDNA4 (gfx1200)."""
+  mcpu, arch = 'gfx1200', 'rdna4'
+
+  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())
+  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])))
+  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())
+  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))
+  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())
+  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())
+  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())
+  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)))
+  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])))
+
+# Keep old class name for backwards compatibility
+TestTinygradKernelRoundtrip = TestRoundtripRDNA3
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/rdna3/generate.py

@@ -1,136 +0,0 @@
-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/parse.py

@@ -1,15 +0,0 @@
-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/rocm/.gitignore

@@ -1,4 +0,0 @@
-*.deb
-build
-src
-sniffer/sniff.so

extra/assembly/rocm/NOTES

@@ -1,20 +0,0 @@
-Built ROCT-Thunk-Interface (hsakmt)
-  hsakmt-roct-dev_5.4.4.99999-local_amd64.deb
-  note: installs to /opt/rocm
-Built ROCm-Device-Libs
-  Works with ROCM_PATH=/home/tiny/build/ROCm-Device-Libs/build/dist
-  rocm-device-libs_1.0.0.99999-local_amd64.deb
-Built ROCm-CompilerSupport (amd_comgr)
-  no deb, sudo make install to /usr/local
-Built ROCR-Runtime
-  hsa-rocr_1.8.0-local_amd64.deb
-  hsa-rocr-dev_1.8.0-local_amd64.deb
-Built ROCm-OpenCL-Runtime
-  rocm-ocl-icd_2.0.0-local_amd64.deb
-  ISSUE: these depend on "comgr"
-  rocm-opencl_2.0.0-local_amd64.deb
-  rocm-opencl-dev_2.0.0-local_amd64.deb
-  Did sudo make install
-
-
-

extra/assembly/rocm/kernel_crashes/dump

@@ -1,41 +0,0 @@
-# run two "rocm-bandwidth-test" in a loop
-# amdgpu-6.0.5-1581431.20.04
-# fixed in kernel 6.2.14
-
-[   72.153646] RIP: 0010:pm_send_runlist+0x4a/0x630 [amdgpu]
-[   72.153815] Code: 30 65 48 8b 04 25 28 00 00 00 48 89 45 d0 31 c0 80 fb 01 0f 87 aa 9d 49 00 83 e3 01 0f 85 1c 05 00 00 49 8b 3f b8 01 00 00 00 <48> 8b 97 30 01 00 00 44 8b b7 6c 01 00 00 8b 9f 70 01 00 00 8b 8a
-[   72.153900] RSP: 0018:ffffb48445c03c30 EFLAGS: 00010246
-[   72.153928] RAX: 0000000000000001 RBX: 0000000000000000 RCX: 0000000000000000
-[   72.153962] RDX: 000000000000007b RSI: ffff9395e1562558 RDI: 0000000000000000
-[   72.153996] RBP: ffffb48445c03cb8 R08: 0000000000000000 R09: 0000000000000001
-[   72.154030] R10: ffff9395c900d840 R11: 0000000000000000 R12: 0000000000000000
-[   72.154065] R13: ffff9395c9e00400 R14: 0000000000000001 R15: ffff9395e15624e0
-[   72.154099] FS:  00007f345c6463c0(0000) GS:ffff93a4aee80000(0000) knlGS:0000000000000000
-[   72.154137] CS:  0010 DS: 0000 ES: 0000 CR0: 0000000080050033
-[   72.154165] CR2: 0000000000000130 CR3: 0000000112840000 CR4: 0000000000750ee0
-[   72.154201] PKRU: 55555554
-[   72.154215] Call Trace:
-[   72.154230]  <TASK>
-[   72.154244]  map_queues_cpsch+0x75/0xc0 [amdgpu]
-[   72.154365]  debug_map_and_unlock+0x51/0x90 [amdgpu]
-[   72.154480]  debug_refresh_runlist+0x1f/0x30 [amdgpu]
-[   72.154591]  kfd_dbg_runtime_disable+0x13c/0x240 [amdgpu]
-[   72.154705]  kfd_ioctl_dbg_set_debug_trap+0x69d/0x8b0 [amdgpu]
-[   72.154820]  kfd_ioctl+0x24a/0x5b0 [amdgpu]
-[   72.154925]  ? kfd_ioctl_create_queue+0x770/0x770 [amdgpu]
-[   72.155035]  ? syscall_exit_to_user_mode+0x27/0x50
-[   72.155061]  ? exit_to_user_mode_prepare+0x3d/0x1c0
-[   72.155088]  __x64_sys_ioctl+0x95/0xd0
-[   72.155109]  do_syscall_64+0x5c/0xc0
-[   72.155128]  ? syscall_exit_to_user_mode+0x27/0x50
-[   72.155151]  ? do_syscall_64+0x69/0xc0
-[   72.155172]  entry_SYSCALL_64_after_hwframe+0x61/0xcb
-[   72.155198] RIP: 0033:0x7f345c7f63ab
-[   72.155218] Code: 0f 1e fa 48 8b 05 e5 7a 0d 00 64 c7 00 26 00 00 00 48 c7 c0 ff ff ff ff c3 66 0f 1f 44 00 00 f3 0f 1e fa b8 10 00 00 00 0f 05 <48> 3d 01 f0 ff ff 73 01 c3 48 8b 0d b5 7a 0d 00 f7 d8 64 89 01 48
-[   72.155301] RSP: 002b:00007ffc97cc89f8 EFLAGS: 00000246 ORIG_RAX: 0000000000000010
-[   72.155339] RAX: ffffffffffffffda RBX: 00007ffc97cc8a30 RCX: 00007f345c7f63ab
-[   72.155375] RDX: 00007ffc97cc8a30 RSI: 00000000c0284b82 RDI: 0000000000000003
-[   72.155411] RBP: 00000000c0284b82 R08: 0000000000000000 R09: 0000000000000000
-[   72.155447] R10: 00007f345cd4ddb0 R11: 0000000000000246 R12: 00007ffc97cc8a30
-[   72.155481] R13: 0000000000000003 R14: 00007ffc97cc8d20 R15: 0000000000000000
-[   72.155517]  </TASK>

extra/assembly/rocm/kernel_crashes/dump2

@@ -1,41 +0,0 @@
-# run two tinygrad matrix example in a loop
-# amdgpu-6.0.5-1581431.20.04
-# NOT fixed in kernel 6.2.14
-
-[  553.016624] gmc_v11_0_process_interrupt: 30 callbacks suppressed
-[  553.016631] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:9 pasid:32770, for process python3 pid 10001 thread python3 pid 10001)
-[  553.016790] amdgpu 0000:0b:00.0: amdgpu:   in page starting at address 0x00007f0000000000 from client 10
-[  553.016892] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00901A30
-[  553.016974] amdgpu 0000:0b:00.0: amdgpu:      Faulty UTCL2 client ID: SDMA0 (0xd)
-[  553.017051] amdgpu 0000:0b:00.0: amdgpu:      MORE_FAULTS: 0x0
-[  553.017111] amdgpu 0000:0b:00.0: amdgpu:      WALKER_ERROR: 0x0
-[  553.017173] amdgpu 0000:0b:00.0: amdgpu:      PERMISSION_FAULTS: 0x3
-[  553.017238] amdgpu 0000:0b:00.0: amdgpu:      MAPPING_ERROR: 0x0
-[  553.017300] amdgpu 0000:0b:00.0: amdgpu:      RW: 0x0
-[  553.123921] [drm:mes_v11_0_submit_pkt_and_poll_completion.constprop.0 [amdgpu]] *ERROR* MES failed to response msg=2
-[  553.124153] amdgpu: failed to add hardware queue to MES, doorbell=0x1a16
-[  553.124195] amdgpu: MES might be in unrecoverable state, issue a GPU reset
-[  553.124237] amdgpu: Failed to restore queue 2
-[  553.124266] amdgpu: Failed to restore process queues
-[  553.124270] amdgpu: Failed to evict queue 3
-[  553.124297] amdgpu: amdgpu_amdkfd_restore_userptr_worker: Failed to resume KFD
-
-# alternative crash in kernel 6.2.14
-
-[  151.097948] gmc_v11_0_process_interrupt: 30 callbacks suppressed
-[  151.097953] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:8 pasid:32771, for process python3 pid 7525 thread python3 pid 7525)
-[  151.097993] amdgpu 0000:0b:00.0: amdgpu:   in page starting at address 0x00007f0000000000 from client 10
-[  151.098008] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00801A30
-[  151.098020] amdgpu 0000:0b:00.0: amdgpu:      Faulty UTCL2 client ID: SDMA0 (0xd)
-[  151.098032] amdgpu 0000:0b:00.0: amdgpu:      MORE_FAULTS: 0x0
-[  151.098042] amdgpu 0000:0b:00.0: amdgpu:      WALKER_ERROR: 0x0
-[  151.098052] amdgpu 0000:0b:00.0: amdgpu:      PERMISSION_FAULTS: 0x3
-[  151.098062] amdgpu 0000:0b:00.0: amdgpu:      MAPPING_ERROR: 0x0
-[  151.098071] amdgpu 0000:0b:00.0: amdgpu:      RW: 0x0
-[  151.209517] [drm:mes_v11_0_submit_pkt_and_poll_completion.constprop.0 [amdgpu]] *ERROR* MES failed to response msg=2
-[  151.209724] amdgpu: failed to add hardware queue to MES, doorbell=0x1002
-[  151.209734] amdgpu: MES might be in unrecoverable state, issue a GPU reset
-[  151.209743] amdgpu: Failed to restore queue 1
-[  151.209751] amdgpu: Failed to restore process queues
-[  151.209759] amdgpu: amdgpu_amdkfd_restore_userptr_worker: Failed to resume KFD
-[  151.209858] amdgpu 0000:0b:00.0: amdgpu: GPU reset begin!

extra/assembly/rocm/kernel_crashes/dump3

@@ -1,20 +0,0 @@
-# two tinygrad + two bandwidth test
-# RDNA2, driver 6.0.5
-# recovered from this!
-
-[  136.971209] gmc_v10_0_process_interrupt: 39 callbacks suppressed
-[  136.971218] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:11 pasid:32773, for process rocm-bandwidth- pid 20281 thread rocm-bandwidth- pid 20281)
-[  136.971228] amdgpu 0000:0b:00.0: amdgpu:   in page starting at address 0x00007f5c2b800000 from client 0x1b (UTCL2)
-[  136.971232] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00B01A31
-[  136.971233] amdgpu 0000:0b:00.0: amdgpu:      Faulty UTCL2 client ID: SDMA0 (0xd)
-[  136.971235] amdgpu 0000:0b:00.0: amdgpu:      MORE_FAULTS: 0x1
-[  136.971236] amdgpu 0000:0b:00.0: amdgpu:      WALKER_ERROR: 0x0
-[  136.971236] amdgpu 0000:0b:00.0: amdgpu:      PERMISSION_FAULTS: 0x3
-[  136.971237] amdgpu 0000:0b:00.0: amdgpu:      MAPPING_ERROR: 0x0
-[  136.971238] amdgpu 0000:0b:00.0: amdgpu:      RW: 0x0
-...
-[  136.993979] amdgpu 0000:0b:00.0: amdgpu: IH ring buffer overflow (0x000BE5A0, 0x0003C480, 0x0003E5C0)
-[  138.209072] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00004000 flags=0x0000]
-[  138.209078] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00004d80 flags=0x0000]
-[  138.209081] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00005000 flags=0x0000]
-[  138.209084] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00005d80 flags=0x0000]

extra/assembly/rocm/kernel_crashes/dump4

@@ -1,33 +0,0 @@
-# ROCK-Kernel-Driver 0b579de9622f5c93021dcb7927d13926313740a2
-# non fatal "crash"
-
-[  127.418045] ------------[ cut here ]------------
-[  127.418046] User pages unexpectedly invalid
-[  127.418056] WARNING: CPU: 16 PID: 260 at drivers/gpu/drm/amd/amdgpu/amdgpu_amdkfd_gpuvm.c:3000 amdgpu_amdkfd_restore_userptr_worker+0x4d9/0x500 [amdgpu]
-[  127.418235] Modules linked in: rfcomm cmac algif_hash algif_skcipher af_alg bnep nls_iso8859_1 iwlmvm mac80211 intel_rapl_msr intel_rapl_common edac_mce_amd snd_hda_codec_realtek snd_hda_codec_generic snd_hda_codec_hdmi kvm_amd binfmt_misc snd_hda_intel snd_intel_dspcfg kvm libarc4 snd_intel_sdw_acpi snd_hda_codec btusb iwlwifi btrtl snd_hda_core btbcm btintel irqbypass btmtk snd_hwdep crct10dif_pclmul snd_pcm polyval_clmulni bluetooth snd_seq_midi snd_seq_midi_event snd_rawmidi snd_seq polyval_generic cfg80211 ghash_clmulni_intel eeepc_wmi snd_seq_device snd_timer aesni_intel asus_wmi ecdh_generic snd platform_profile crypto_simd ledtrig_audio cryptd ecc ccp soundcore sparse_keymap rapl k10temp wmi_bmof mac_hid sch_fq_codel msr parport_pc ppdev lp parport ramoops pstore_blk efi_pstore reed_solomon pstore_zone ip_tables x_tables autofs4 amdgpu hid_generic usbhid hid i2c_algo_bit drm_ttm_helper ttm video iommu_v2 drm_buddy gpu_sched drm_display_helper drm_kms_helper syscopyarea
-[  127.418276]  sysfillrect sysimgblt fb_sys_fops drm nvme nvme_core cec r8169 ahci crc32_pclmul rc_core i2c_piix4 xhci_pci libahci nvme_common xhci_pci_renesas realtek wmi
-[  127.418284] CPU: 16 PID: 260 Comm: kworker/16:1 Tainted: G        W          6.0.0 #4
-[  127.418286] Hardware name: System manufacturer System Product Name/TUF GAMING X570-PLUS (WI-FI), BIOS 3603 03/20/2021
-[  127.418287] Workqueue: events amdgpu_amdkfd_restore_userptr_worker [amdgpu]
-[  127.418455] RIP: 0010:amdgpu_amdkfd_restore_userptr_worker+0x4d9/0x500 [amdgpu]
-[  127.418601] Code: ff e8 2b 8a 96 d1 e9 66 fe ff ff 48 c7 c7 40 4f f5 c0 e8 56 7b 8a d1 0f 0b e9 2e ff ff ff 48 c7 c7 d8 d0 ed c0 e8 43 7b 8a d1 <0f> 0b e9 0a fe ff ff 4c 89 ef e8 f8 89 96 d1 e9 cb fd ff ff e8 ce
-[  127.418603] RSP: 0018:ffffb36740a83dc8 EFLAGS: 00010282
-[  127.418604] RAX: 0000000000000000 RBX: ffff9d159ee9df30 RCX: 0000000000000027
-[  127.418605] RDX: 0000000000000027 RSI: ffffb36740a83c88 RDI: ffff9d242a220568
-[  127.418606] RBP: ffffb36740a83e58 R08: ffff9d242a220560 R09: 0000000000000001
-[  127.418607] R10: 0000000000000001 R11: 0000000000000020 R12: ffff9d159ee9df98
-[  127.418607] R13: ffff9d159ee9df70 R14: ffff9d159ee9dee0 R15: ffff9d159ee9dee0
-[  127.418608] FS:  0000000000000000(0000) GS:ffff9d242a200000(0000) knlGS:0000000000000000
-[  127.418609] CS:  0010 DS: 0000 ES: 0000 CR0: 0000000080050033
-[  127.418610] CR2: 00007fd5d4715000 CR3: 0000000120ffe000 CR4: 0000000000750ee0
-[  127.418611] PKRU: 55555554
-[  127.418611] Call Trace:
-[  127.418612]  <TASK>
-[  127.418613]  process_one_work+0x21f/0x3f0
-[  127.418615]  worker_thread+0x4a/0x3c0
-[  127.418617]  ? process_one_work+0x3f0/0x3f0
-[  127.418618]  kthread+0xf0/0x120
-[  127.418619]  ? kthread_complete_and_exit+0x20/0x20
-[  127.418620]  ret_from_fork+0x22/0x30
-[  127.418622]  </TASK>
-[  127.418623] ---[ end trace 0000000000000000 ]---

extra/assembly/rocm/rdna3/asm.py

@@ -1,80 +0,0 @@
-import numpy as np
-import pathlib
-from hexdump import hexdump
-from tinygrad.helpers import colored
-from extra.helpers import enable_early_exec
-early_exec = enable_early_exec()
-
-from tinygrad.runtime.ops_cl import CLProgram, CLBuffer, ROCM_LLVM_PATH
-
-ENABLE_NON_ASM = False
-
-WMMA = True
-DUAL_ALU = True
-F32 = True
-
-if ENABLE_NON_ASM:
-  buf = CLBuffer.fromCPU(np.zeros(10, np.float32))
-  prg_empty = CLProgram("code", "__kernel void code(__global float *a) { a[0] = 1; }")
-  asm_real = prg_empty.binary()
-  with open("/tmp/cc.elf", "wb") as f:
-    f.write(asm_real)
-  prg_empty([1], [1], buf, wait=True)
-  print(buf.toCPU())
-
-print(colored("creating CLBuffer", "green"))
-buf = CLBuffer.fromCPU(np.zeros(10, np.float32))
-code = open(pathlib.Path(__file__).parent / "prog.s", "r").read()
-
-gen = []
-FLOPS = 0
-MAX_REG = 251
-for j in range(1):
-  if WMMA:
-    KY, KX = 4, 4
-    for y in range(KY):
-      for x in range(KX):
-        c = (y*KX+x)*8
-        a = (KY*KX*8) + y*8
-        b = (KY*KX*8) + (KY*8) + x*8
-        gen.append(f"v_wmma_f32_16x16x16_f16 v[{c}:{c+7}], v[{a}:{a+7}], v[{b}:{b+7}], v[{c}:{c+7}]")
-        FLOPS += 16*8*2
-  else:
-    for i in range(0, MAX_REG, 6):
-      if DUAL_ALU:
-        if F32:
-          gen.append(f"v_dual_fmac_f32 v{i+0}, v{i+1}, v{i+2} :: v_dual_fmac_f32 v{i+3}, v{i+4}, v{i+5}")
-          FLOPS += 4
-        else:
-          gen.append(f"v_dual_dot2acc_f32_f16 v{i+0}, v{i+1}, v{i+2} :: v_dual_dot2acc_f32_f16 v{i+3}, v{i+4}, v{i+5}")
-          FLOPS += 8
-      else:
-        assert F32
-        gen.append(f"v_fmac_f32 v{i+0}, v{i+1}, v{i+2}")
-        gen.append(f"v_fmac_f32 v{i+3}, v{i+4}, v{i+5}")
-code = code.replace("// FLOPS", '\n'.join(gen))
-print(code)
-
-
-# fix: COMGR failed to get code object ISA name. set triple to 'amdgcn-amd-amdhsa'
-
-object = early_exec(([ROCM_LLVM_PATH / "llvm-mc", '--arch=amdgcn', '--mcpu=gfx1100', '--triple=amdgcn-amd-amdhsa', '--filetype=obj', '-'], code.encode("utf-8")))
-asm = early_exec(([ROCM_LLVM_PATH / "ld.lld", "/dev/stdin", "-o", "/dev/stdout", "--pie"], object))
-
-with open("/tmp/cc2.o", "wb") as f:
-  f.write(object)
-with open("/tmp/cc2.elf", "wb") as f:
-  f.write(asm)
-
-print(colored("creating CLProgram", "green"))
-prg = CLProgram("code", asm)
-
-print(colored("running program", "green"))
-G = 512
-FLOPS *= 100000*G*G  # loop * global_size
-for i in range(3):
-  tm = prg(buf, global_size=[G//256, G, 1], local_size=[256, 1, 1], wait=True)
-  print(f"ran in {tm*1e3:.2f} ms, {FLOPS/(tm*1e9):.2f} GFLOPS")
-
-print(colored("transferring buffer", "green"))
-print(buf.toCPU())

extra/assembly/rocm/rdna3/prog.s

@@ -1,80 +0,0 @@
-.global _start
-_start:
-.rodata
-.align 0x10
-.global code.kd
-.type code.kd,STT_OBJECT
-# amd_kernel_code_t (must be at 0x440 for kernel_code_entry_byte_offset to be right)
-code.kd:
-# amd_kernel_..., amd_machine_...
-.long 0,0,0,0
-# kernel_code_entry_byte_offset, kernel_code_prefetch_byte_offset
-.long 0x00000bc0,0x00000000,0x00000000,0x00000000
-# kernel_code_prefetch_byte_size, max_scratch_backing_memory_byte_size
-.long 0,0,0,0
-# compute_pgm_rsrc1, compute_pgm_rsrc2, kernel_code_properties, workitem_private_segment_byte_size
-.long 0x60af0000,0x0000009e,0x00000408,0x00000000
-# compute_pgm_rsrc1 |= AMD_COMPUTE_PGM_RSRC_ONE_FLOAT_DENORM_MODE_32 | AMD_COMPUTE_PGM_RSRC_ONE_FLOAT_DENORM_MODE_16_64
-# compute_pgm_rsrc1 |= AMD_COMPUTE_PGM_RSRC_ONE_ENABLE_DX10_CLAMP | AMD_COMPUTE_PGM_RSRC_ONE_ENABLE_IEEE_MODE
-# compute_pgm_rsrc2 |= AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT = 0xF
-# compute_pgm_rsrc2 |= AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X
-# kernel_code_properties |= AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_KERNARG_SEGMENT_PTR = 1
-# kernel_code_properties |= AMD_KERNEL_CODE_PROPERTIES_RESERVED1 = 1
-.text
-.global code
-.type code,STT_FUNC
-code:
-# https://llvm.org/docs/AMDGPUUsage.html#initial-kernel-execution-state
-# s[0:1] contains the kernarg_address
-# TODO: can we use s[2:3] if this was really a wave since we only alloced 2 SGPRs?
-s_load_b64 s[2:3], s[0:1], null
-
-s_mov_b32 s8, 0
-loop:
-s_addk_i32 s8, 1
-s_cmp_eq_u32 s8, 100000
-// FLOPS
-s_cbranch_scc0 loop
-
-# wait for the s_load_b64
-s_waitcnt lgkmcnt(0)
-
-v_dual_mov_b32 v0, 4 :: v_dual_mov_b32 v1, 2.0
-global_store_b32 v0, v1, s[2:3]
-
-# Deallocate all VGPRs for this wave. Use only when next instruction is S_ENDPGM.
-s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
-s_endpgm
-s_code_end
-
-.amdgpu_metadata
-amdhsa.kernels:
-  - .args:
-      - .address_space:  global
-        .name:           a
-        .offset:         0
-        .size:           8
-        .type_name:      'float*'
-        .value_kind:     global_buffer
-    .group_segment_fixed_size: 0
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 8
-    .language:       OpenCL C
-    .language_version:
-      - 1
-      - 2
-    .max_flat_workgroup_size: 256
-    .name:           code
-    .private_segment_fixed_size: 0
-    .sgpr_count:     2
-    .sgpr_spill_count: 0
-    .symbol:         code.kd
-    .uses_dynamic_stack: false
-    .vgpr_count:     256
-    .vgpr_spill_count: 0
-    .wavefront_size: 32
-amdhsa.target:   amdgcn-amd-amdhsa--gfx1100
-amdhsa.version:
-  - 1
-  - 2
-.end_amdgpu_metadata

extra/assembly/rocm/rocm_clone.sh

@@ -1,11 +0,0 @@
-#!/bin/bash
-mkdir -p src
-cd src
-git clone https://github.com/RadeonOpenCompute/ROCT-Thunk-Interface.git -b rocm-5.5.0
-git clone https://github.com/RadeonOpenCompute/ROCm-Device-Libs.git -b rocm-5.5.0
-git clone https://github.com/RadeonOpenCompute/llvm-project.git -b rocm-5.5.0 --depth 1
-git clone https://github.com/RadeonOpenCompute/ROCR-Runtime.git -b rocm-5.5.0
-git clone https://github.com/ROCm-Developer-Tools/ROCclr.git -b rocm-5.5.0
-git clone https://github.com/RadeonOpenCompute/ROCm-CompilerSupport.git -b rocm-5.5.0
-git clone https://github.com/RadeonOpenCompute/ROCm-OpenCL-Runtime.git -b rocm-5.5.0
-cd ../

extra/assembly/rocm/rocm_from_scratch.sh

@@ -1,69 +0,0 @@
-#!/bin/bash
-mkdir -p build/debs
-cd build
-
-# ROCT-Thunk-Interface (hsakmt)
-if [ ! -f debs/hsakmt-roct-dev_5.5.0.99999-local_amd64.deb ]
-then
-  mkdir -p ROCT-Thunk-Interface
-  cd ROCT-Thunk-Interface
-  cmake ../../src/ROCT-Thunk-Interface
-  make -j32 package
-  cp hsakmt-roct-dev_5.5.0.99999-local_amd64.deb ../debs
-  cd ../
-fi
-
-
-# build custom LLVM
-if [ ! -f llvm-project/bin/clang ]
-then
-  mkdir -p llvm-project
-  cd llvm-project
-  cmake -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_PROJECTS="llvm;clang;lld" -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" ../../src/llvm-project/llvm
-  make -j32
-  cd ..
-fi
-
-# use custom LLVM
-export PATH="$PWD/llvm-project/bin:$PATH"
-
-# ROCm-Device-Libs
-if [ ! -f debs/rocm-device-libs_1.0.0.99999-local_amd64.deb ]
-then
-  mkdir -p ROCm-Device-Libs
-  cd ROCm-Device-Libs
-  cmake ../../src/ROCm-Device-Libs
-  make -j32 package
-  cp rocm-device-libs_1.0.0.99999-local_amd64.deb ../debs
-  cd ../
-fi
-
-# ROCR-Runtime
-if [ ! -f debs/hsa-rocr_1.8.0-local_amd64.deb ]
-then
-  mkdir -p ROCR-Runtime
-  cd ROCR-Runtime
-  cmake ../../src/ROCR-Runtime/src
-  make -j32 package
-  cp hsa-rocr_1.8.0-local_amd64.deb ../debs
-  cp hsa-rocr-dev_1.8.0-local_amd64.deb ../debs
-  cd ../
-fi
-
-# ROCm-OpenCL-Runtime (needs ROCclr)
-if [ ! -f debs/rocm-opencl_2.0.0-local_amd64.deb ]
-then
-  mkdir -p ROCm-OpenCL-Runtime
-  cd ROCm-OpenCL-Runtime
-  cmake ../../src/ROCm-OpenCL-Runtime
-  make -j32 package
-  cp rocm-opencl_2.0.0-local_amd64.deb ../debs
-  cp rocm-opencl-dev_2.0.0-local_amd64.deb ../debs
-  cp rocm-ocl-icd_2.0.0-local_amd64.deb ../debs
-fi
-
-# ROCm-CompilerSupport (broken)
-#mkdir -p ROCm-CompilerSupport
-#cd ROCm-CompilerSupport
-#cmake ../../src/ROCm-CompilerSupport/lib/comgr
-#make -j32

extra/assembly/rocm/rocm_setup.sh

@@ -1,14 +0,0 @@
-#!/bin/bash
-rm amdgpu-install_5.5.50500-1_all.deb
-wget https://repo.radeon.com/amdgpu-install/5.5/ubuntu/$(lsb_release -cs)/amdgpu-install_5.5.50500-1_all.deb
-sudo dpkg -i amdgpu-install_5.5.50500-1_all.deb
-sudo apt-get update
-
-# kernel driver
-sudo apt-get install amdgpu-dkms
-
-# for opencl
-sudo apt-get install rocm-opencl-runtime
-
-# for HIP
-sudo apt-get install hip-runtime-amd rocm-device-libs hip-dev

extra/assembly/rocm/sniffer/build.sh

@@ -1,11 +0,0 @@
-#!/bin/bash -e
-clang sniff.cc -Werror -shared -fPIC -I../src/ -I../src/ROCT-Thunk-Interface/include -I../src/ROCm-Device-Libs/ockl/inc -o sniff.so -lstdc++
-#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so /home/tiny/build/HIP-Examples/HIP-Examples-Applications/HelloWorld/HelloWorld
-#AMD_LOG_LEVEL=4 LD_PRELOAD=$PWD/sniff.so $HOME/build/HIP-Examples/HIP-Examples-Applications/HelloWorld/HelloWorld
-#AMD_LOG_LEVEL=5 LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
-DEBUG=5 LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
-#AMD_LOG_LEVEL=5 HSAKMT_DEBUG_LEVEL=7 DEBUG=5 LD_PRELOAD=$PWD/sniff.so strace -F python3 ../rdna3/asm.py
-#LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
-#AMD_LOG_LEVEL=4 LD_PRELOAD=$PWD/sniff.so FORWARD_ONLY=1 DEBUG=2 python3 ../../../test/test_ops.py TestOps.test_add
-#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so rocm-bandwidth-test -s 0 -d 1 -m 1
-#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so rocm-bandwidth-test -s 1 -d 2 -m 1

extra/assembly/rocm/sniffer/sniff.cc

@@ -1,282 +0,0 @@
-// template copied from https://github.com/geohot/cuda_ioctl_sniffer/blob/master/sniff.cc
-
-#include <stdio.h>
-#include <stdlib.h>
-#include <string.h>
-#include <dlfcn.h>
-#include <signal.h>
-#include <ucontext.h>
-
-#include <sys/mman.h>
-
-// includes from the ROCm sources
-#include <linux/kfd_ioctl.h>
-#include <hsa.h>
-#include <amd_hsa_kernel_code.h>
-#include <ROCR-Runtime/src/core/inc/sdma_registers.h>
-using namespace rocr::AMD;
-
-#include <string>
-#include <map>
-std::map<int, std::string> files;
-std::map<uint64_t, uint64_t> ring_base_addresses;
-
-#define D(args...) fprintf(stderr, args)
-
-uint64_t doorbell_offset = -1;
-std::map<uint64_t, int> queue_types;
-
-void hexdump(void *d, int l) {
-  for (int i = 0; i < l; i++) {
-    if (i%0x10 == 0 && i != 0) printf("\n");
-    if (i%0x10 == 8) printf(" ");
-    if (i%0x10 == 0) printf("%8X: ", i);
-    printf("%2.2X ", ((uint8_t*)d)[i]);
-  }
-  printf("\n");
-}
-
-extern "C" {
-
-// https://defuse.ca/online-x86-assembler.htm#disassembly2
-static void handler(int sig, siginfo_t *si, void *unused) {
-  ucontext_t *u = (ucontext_t *)unused;
-  uint8_t *rip = (uint8_t*)u->uc_mcontext.gregs[REG_RIP];
-
-  int store_size = 0;
-  uint64_t value;
-  if (rip[0] == 0x48 && rip[1] == 0x89 && rip[2] == 0x30) {
-    // 0:  48 89 30                mov    QWORD PTR [rax],rsi
-    store_size = 8;
-    value = u->uc_mcontext.gregs[REG_RSI];
-    u->uc_mcontext.gregs[REG_RIP] += 3;
-  } else if (rip[0] == 0x4c && rip[1] == 0x89 && rip[2] == 0x28) {
-    // 0:  4c 89 28                mov    QWORD PTR [rax],r13
-    store_size = 8;
-    value = u->uc_mcontext.gregs[REG_R13];
-    u->uc_mcontext.gregs[REG_RIP] += 3;
-  } else {
-    D("segfault %02X %02X %02X %02X %02X %02X %02X %02X rip: %p addr: %p\n", rip[0], rip[1], rip[2], rip[3], rip[4], rip[5], rip[6], rip[7], rip, si->si_addr);
-    D("rax: %llx rcx: %llx rdx: %llx rsi: %llx rbx: %llx\n", u->uc_mcontext.gregs[REG_RAX], u->uc_mcontext.gregs[REG_RCX], u->uc_mcontext.gregs[REG_RDX], u->uc_mcontext.gregs[REG_RSI], u->uc_mcontext.gregs[REG_RBX]);
-    exit(-1);
-  }
-
-  uint64_t ring_base_address = ring_base_addresses[((uint64_t)si->si_addr)&0xFFF];
-  int queue_type = queue_types[((uint64_t)si->si_addr)&0xFFF];
-  D("%16p: \u001b[31mDING DONG\u001b[0m (queue_type %d) store(%d): 0x%8lx -> %p ring_base_address:0x%lx\n", rip, queue_type, store_size, value, si->si_addr, ring_base_address);
-
-  if (queue_type == KFD_IOC_QUEUE_TYPE_SDMA) {
-    uint8_t *sdma_ptr = (uint8_t*)(ring_base_address);
-    while (sdma_ptr < ((uint8_t*)(ring_base_address)+value)) {
-      D("0x%3lx: ", sdma_ptr-(uint8_t*)(ring_base_address));
-      if (sdma_ptr[0] == SDMA_OP_TIMESTAMP) {
-        D("SDMA_PKT_TIMESTAMP\n");
-        sdma_ptr += sizeof(SDMA_PKT_TIMESTAMP);
-      } else if (sdma_ptr[0] == SDMA_OP_GCR) {
-        D("SDMA_PKT_GCR\n");
-        sdma_ptr += sizeof(SDMA_PKT_GCR);
-      } else if (sdma_ptr[0] == SDMA_OP_ATOMIC) {
-        D("SDMA_PKT_ATOMIC\n");
-        sdma_ptr += sizeof(SDMA_PKT_ATOMIC);
-      } else if (sdma_ptr[0] == SDMA_OP_FENCE) {
-        D("SDMA_PKT_FENCE\n");
-        sdma_ptr += sizeof(SDMA_PKT_FENCE);
-      } else if (sdma_ptr[0] == SDMA_OP_TRAP) {
-        D("SDMA_PKT_TRAP\n");
-        sdma_ptr += sizeof(SDMA_PKT_TRAP);
-      } else if (sdma_ptr[0] == SDMA_OP_COPY && sdma_ptr[1] == SDMA_SUBOP_COPY_LINEAR) {
-        SDMA_PKT_COPY_LINEAR *pkt = (SDMA_PKT_COPY_LINEAR *)sdma_ptr;
-        D("SDMA_PKT_COPY_LINEAR: count:0x%x src:0x%lx dst:0x%lx\n", pkt->COUNT_UNION.count+1,
-          (uint64_t)pkt->SRC_ADDR_LO_UNION.src_addr_31_0 | ((uint64_t)pkt->SRC_ADDR_HI_UNION.src_addr_63_32 << 32),
-          (uint64_t)pkt->DST_ADDR_LO_UNION.dst_addr_31_0 | ((uint64_t)pkt->DST_ADDR_HI_UNION.dst_addr_63_32 << 32)
-        );
-        sdma_ptr += sizeof(SDMA_PKT_COPY_LINEAR);
-      } else {
-        D("unhandled packet type %d %d, exiting\n", sdma_ptr[0], sdma_ptr[1]);
-        break;
-      }
-    }
-
-    //hexdump((void*)(ring_base_address), 0x100);
-  } else if (queue_type == KFD_IOC_QUEUE_TYPE_COMPUTE_AQL) {
-    hsa_kernel_dispatch_packet_t *pkt = (hsa_kernel_dispatch_packet_t *)(ring_base_address+value*0x40);
-    if ((pkt->header&0xFF) == HSA_PACKET_TYPE_KERNEL_DISPATCH) {
-      D("HSA_PACKET_TYPE_KERNEL_DISPATCH -- setup:%d workgroup[%d, %d, %d] grid[%d, %d, %d] kernel_object:0x%lx kernarg_address:%p\n", pkt->setup, pkt->workgroup_size_x, pkt->workgroup_size_y, pkt->workgroup_size_z, pkt->grid_size_x, pkt->grid_size_y, pkt->grid_size_z, pkt->kernel_object, pkt->kernarg_address);
-      amd_kernel_code_t *code = (amd_kernel_code_t *)pkt->kernel_object;
-      D("kernel_code_entry_byte_offset:%lx\n", code->kernel_code_entry_byte_offset);
-      uint32_t *kernel_code = (uint32_t*)(pkt->kernel_object + code->kernel_code_entry_byte_offset);
-      int code_len = 0;
-      while (kernel_code[code_len] != 0xbf9f0000 && kernel_code[code_len] != 0) code_len++;
-      hexdump(kernel_code, code_len*4);
-      /*FILE *f = fopen("/tmp/kernel_code", "wb");
-      fwrite(kernel_code, 4, code_len, f);
-      fclose(f);
-      system("python -c 'print(\" \".join([(\"0x%02X\"%x) for x in open(\"/tmp/kernel_code\", \"rb\").read()]))' | ../build/llvm-project/bin/llvm-mc --disassemble --arch=amdgcn --mcpu=gfx1100 --show-encoding");*/
-      D("kernargs (kernarg_segment_byte_size:0x%lx)\n", code->kernarg_segment_byte_size);
-      // get length
-      int i;
-      for (i = 0; i < 0x400; i+=0x10) {
-        if (memcmp((void*)((uint64_t)pkt->kernarg_address+i), "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", 0x10) == 0) break;
-      }
-      hexdump((void*)pkt->kernarg_address, i+0x10);
-    } else if ((pkt->header&0xFF) == HSA_PACKET_TYPE_BARRIER_AND) {
-      hsa_barrier_and_packet_t *pkt_and = (hsa_barrier_and_packet_t *)(ring_base_address+value*0x40);
-      D("HSA_PACKET_TYPE_BARRIER_AND completion_signal:0x%lx\n", pkt_and->completion_signal.handle);
-      //hexdump((void*)(ring_base_address+value*0x40), 0x40);
-    } else if ((pkt->header&0xFF) == HSA_PACKET_TYPE_VENDOR_SPECIFIC) {
-      D("HSA_PACKET_TYPE_VENDOR_SPECIFIC\n");
-      hexdump((void*)(ring_base_address+value*0x40), 0x40);
-    } else {
-      hexdump((void*)(ring_base_address+value*0x40), 0x40);
-    }
-  }
-
-  mprotect((void *)((uint64_t)si->si_addr & ~0xFFF), 0x2000, PROT_READ | PROT_WRITE);
-  if (store_size == 8) {
-    *(volatile uint64_t*)(si->si_addr) = value;
-  } else if (store_size == 4) {
-    *(volatile uint32_t*)(si->si_addr) = value;
-  } else if (store_size == 2) {
-    *(volatile uint16_t*)(si->si_addr) = value;
-  } else {
-    D("store size not supported\n");
-    exit(-1);
-  }
-  mprotect((void *)((uint64_t)si->si_addr & ~0xFFF), 0x2000, PROT_NONE);
-}
-
-void register_sigsegv_handler() {
-  struct sigaction sa = {0};
-  sa.sa_flags = SA_SIGINFO;
-  sigemptyset(&sa.sa_mask);
-  sa.sa_sigaction = handler;
-  if (sigaction(SIGSEGV, &sa, NULL) == -1) {
-    D("ERROR: failed to register sigsegv handler");
-    exit(-1);
-  }
-  // NOTE: python (or ocl runtime?) blocks the SIGSEGV signal
-  sigset_t x;
-  sigemptyset(&x);
-  sigaddset(&x, SIGSEGV);
-  sigprocmask(SIG_UNBLOCK, &x, NULL);
-}
-
-int (*my_open)(const char *pathname, int flags, mode_t mode);
-#undef open
-int open(const char *pathname, int flags, mode_t mode) {
-  if (my_open == NULL) my_open = reinterpret_cast<decltype(my_open)>(dlsym(RTLD_NEXT, "open"));
-  int ret = my_open(pathname, flags, mode);
-  //D("open %s (0o%o) = %d\n", pathname, flags, ret);
-  files[ret] = pathname;
-  return ret;
-}
-
-
-int (*my_open64)(const char *pathname, int flags, mode_t mode);
-#undef open
-int open64(const char *pathname, int flags, mode_t mode) {
-  if (my_open64 == NULL) my_open64 = reinterpret_cast<decltype(my_open64)>(dlsym(RTLD_NEXT, "open64"));
-  int ret = my_open64(pathname, flags, mode);
-  //D("open %s (0o%o) = %d\n", pathname, flags, ret);
-  files[ret] = pathname;
-  return ret;
-}
-
-void *(*my_mmap)(void *addr, size_t length, int prot, int flags, int fd, off_t offset);
-#undef mmap
-void *mmap(void *addr, size_t length, int prot, int flags, int fd, off_t offset) {
-  if (my_mmap == NULL) my_mmap = reinterpret_cast<decltype(my_mmap)>(dlsym(RTLD_NEXT, "mmap"));
-  void *ret = my_mmap(addr, length, prot, flags, fd, offset);
-
-  if (doorbell_offset != -1 && offset == doorbell_offset) {
-    D("HIDDEN DOORBELL %p, handled by %p\n", addr, handler);
-    register_sigsegv_handler();
-    mprotect(addr, length, PROT_NONE);
-  }
-
-  if (fd != -1) D("mmapped %p (target %p) with flags 0x%x length 0x%zx fd %d %s offset 0x%lx\n", ret, addr, flags, length, fd, files[fd].c_str(), offset);
-  return ret;
-}
-
-void *(*my_mmap64)(void *addr, size_t length, int prot, int flags, int fd, off_t offset);
-#undef mmap64
-void *mmap64(void *addr, size_t length, int prot, int flags, int fd, off_t offset) { return mmap(addr, length, prot, flags, fd, offset); }
-
-int ioctl_num = 1;
-int (*my_ioctl)(int filedes, unsigned long request, void *argp) = NULL;
-#undef ioctl
-int ioctl(int filedes, unsigned long request, void *argp) {
-  if (my_ioctl == NULL) my_ioctl = reinterpret_cast<decltype(my_ioctl)>(dlsym(RTLD_NEXT, "ioctl"));
-  int ret = 0;
-  ret = my_ioctl(filedes, request, argp);
-  if (!files.count(filedes)) return ret;
-
-  uint8_t type = (request >> 8) & 0xFF;
-  uint8_t nr = (request >> 0) & 0xFF;
-  uint16_t size = (request >> 16) & 0xFFF;
-
-  D("%3d: %d = %3d(%20s) 0x%3x ", ioctl_num, ret, filedes, files[filedes].c_str(), size);
-
-  if (request == AMDKFD_IOC_SET_EVENT) {
-    kfd_ioctl_set_event_args *args = (kfd_ioctl_set_event_args *)argp;
-    D("AMDKFD_IOC_SET_EVENT event_id:%d", args->event_id);
-  } else if (request == AMDKFD_IOC_ALLOC_MEMORY_OF_GPU) {
-    kfd_ioctl_alloc_memory_of_gpu_args *args = (kfd_ioctl_alloc_memory_of_gpu_args *)argp;
-    D("AMDKFD_IOC_ALLOC_MEMORY_OF_GPU va_addr:0x%llx size:0x%llx handle:%llX gpu_id:0x%x", args->va_addr, args->size, args->handle, args->gpu_id);
-  } else if (request == AMDKFD_IOC_MAP_MEMORY_TO_GPU) {
-    kfd_ioctl_map_memory_to_gpu_args *args = (kfd_ioctl_map_memory_to_gpu_args *)argp;
-    D("AMDKFD_IOC_MAP_MEMORY_TO_GPU handle:%llX", args->handle);
-  } else if (request == AMDKFD_IOC_CREATE_EVENT) {
-    kfd_ioctl_create_event_args *args = (kfd_ioctl_create_event_args *)argp;
-    D("AMDKFD_IOC_CREATE_EVENT event_page_offset:0x%llx event_type:%d event_id:%d", args->event_page_offset, args->event_type, args->event_id);
-  } else if (request == AMDKFD_IOC_WAIT_EVENTS) {
-    D("AMDKFD_IOC_WAIT_EVENTS");
-  } else if (request == AMDKFD_IOC_SET_XNACK_MODE) {
-    D("AMDKFD_IOC_SET_XNACK_MODE");
-  } else if (request == AMDKFD_IOC_SVM || (type == 0x4b && nr == 0x20)) {
-    // NOTE: this one is variable length
-    kfd_ioctl_svm_args *args = (kfd_ioctl_svm_args *)argp;
-    D("AMDKFD_IOC_SVM start_addr:0x%llx size:0x%llx op:%d", args->start_addr, args->size, args->op);
-  } else if (request == AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU) {
-    kfd_ioctl_unmap_memory_from_gpu_args *args = (kfd_ioctl_unmap_memory_from_gpu_args *)argp;
-    D("AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU handle:%llX", args->handle);
-  } else if (request == AMDKFD_IOC_FREE_MEMORY_OF_GPU) {
-    D("AMDKFD_IOC_FREE_MEMORY_OF_GPU");
-  } else if (request == AMDKFD_IOC_SET_SCRATCH_BACKING_VA) {
-    D("AMDKFD_IOC_SET_SCRATCH_BACKING_VA");
-  } else if (request == AMDKFD_IOC_GET_TILE_CONFIG) {
-    D("AMDKFD_IOC_GET_TILE_CONFIG");
-  } else if (request == AMDKFD_IOC_SET_TRAP_HANDLER) {
-    D("AMDKFD_IOC_SET_TRAP_HANDLER");
-  } else if (request == AMDKFD_IOC_GET_VERSION) {
-    kfd_ioctl_get_version_args *args = (kfd_ioctl_get_version_args *)argp;
-    D("AMDKFD_IOC_GET_VERSION major_version:%d minor_version:%d", args->major_version, args->minor_version);
-  } else if (request == AMDKFD_IOC_GET_PROCESS_APERTURES_NEW) {
-    D("AMDKFD_IOC_GET_PROCESS_APERTURES_NEW");
-  } else if (request == AMDKFD_IOC_ACQUIRE_VM) {
-    D("AMDKFD_IOC_ACQUIRE_VM");
-  } else if (request == AMDKFD_IOC_SET_MEMORY_POLICY) {
-    D("AMDKFD_IOC_SET_MEMORY_POLICY");
-  } else if (request == AMDKFD_IOC_GET_CLOCK_COUNTERS) {
-    D("AMDKFD_IOC_GET_CLOCK_COUNTERS");
-  } else if (request == AMDKFD_IOC_CREATE_QUEUE) {
-    kfd_ioctl_create_queue_args *args = (kfd_ioctl_create_queue_args *)argp;
-    D("AMDKFD_IOC_CREATE_QUEUE\n");
-    D("queue_type:%d ring_base_address:0x%llx\n", args->queue_type, args->ring_base_address);
-    D("eop_buffer_address:0x%llx ctx_save_restore_address:0x%llx\n", args->eop_buffer_address, args->ctx_save_restore_address);
-    D("ring_size:0x%x queue_priority:%d\n", args->ring_size, args->queue_priority);
-    D("RETURNS write_pointer_address:0x%llx read_pointer_address:0x%llx doorbell_offset:0x%llx queue_id:%d\n", args->write_pointer_address, args->read_pointer_address, args->doorbell_offset, args->queue_id);
-    //D("RETURNS *write_pointer_address:0x%llx *read_pointer_address:0x%llx\n", *(uint64_t*)args->write_pointer_address, *(uint64_t*)args->read_pointer_address);
-    ring_base_addresses[args->doorbell_offset&0xFFF] = args->ring_base_address;
-    queue_types[args->doorbell_offset&0xFFF] = args->queue_type;
-    doorbell_offset = args->doorbell_offset&~0xFFF;
-  } else {
-    D("type:0x%x nr:0x%x size:0x%x", type, nr, size);
-  }
-
-  D("\n");
-  ioctl_num++;
-  return ret;
-}
-
-}

extra/export_model.py

@@ -1,11 +1,12 @@
 from typing import Tuple, Dict, List, Optional
 from tinygrad.dtype import DType
+from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Device, Tensor
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
-from tinygrad.helpers import Context, to_mv, to_function_name
+from tinygrad.helpers import Context, to_mv
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import Ops, UOp
+from tinygrad.uop.ops import Ops
 import json
 from collections import OrderedDict
 
@@ -14,13 +15,8 @@ EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]
 def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
   functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
   for ji in run.jit_cache:
-    prg: UOp = ji.prg.p
-    name = prg.src[0].arg.name
-    function_name = to_function_name(name)
-    src = prg.src[3].arg
-    global_size, local_size = prg.sizes
-    prg_vars = prg.variables()
-    functions[function_name] = src   # NOTE: this assumes all with the same name are the same
+    fxn: ProgramSpec = ji.prg.p
+    functions[fxn.function_name] = fxn.src   # NOTE: this assumes all with the same name are the same
     cargs = []
     for i,arg in enumerate(ji.bufs):
       key = id(arg)
@@ -32,8 +28,8 @@ def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str]
           bufnum += 1
           if i > 0: bufs_to_save[bufs[key][0]] = arg   # if first usage of a buffer is not an output, and it's not a special name
       cargs.append(bufs[key][0])
-    cargs += [var for var in prg_vars if getattr(var, "op", None) is Ops.DEFINE_VAR] # symbolic vars; is it necessary or sufficient to check for DEFINE_VAR?
-    statements.append((function_name, cargs, list(global_size) if global_size else None, list(local_size) if local_size else None))
+    cargs += [var for var in fxn.vars if getattr(var, "op", None) is Ops.DEFINE_VAR] # symbolic vars; is it necessary or sufficient to check for DEFINE_VAR?
+    statements.append((fxn.function_name, cargs, fxn.global_size, fxn.local_size))
 
   return functions, statements, {name:(size, dtype, key) for (name,size,dtype,key) in bufs.values()}, bufs_to_save
 

extra/gemm/.gitignore

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

extra/gemm/asm/cdna/template.s

@@ -0,0 +1,83 @@
+.text
+.section	.text.
+.global	gemm
+.p2align	8
+.type	gemm,@function
+
+gemm:
+INSTRUCTIONS
+
+.section .rodata,"a",@progbits
+.p2align 6, 0x0
+.amdhsa_kernel gemm
+  # basic memory requirements
+  .amdhsa_group_segment_fixed_size 133120
+  .amdhsa_private_segment_fixed_size 0
+  .amdhsa_kernarg_size 32
+  # register usage (RSRC1)
+  .amdhsa_next_free_vgpr 504
+  .amdhsa_next_free_sgpr 96
+  # workgroup / workitem IDs (RSRC2)
+  .amdhsa_system_sgpr_workgroup_id_x 1
+  .amdhsa_system_sgpr_workgroup_id_y 1
+  .amdhsa_system_sgpr_workgroup_id_z 1
+  # user SGPRs, we only specify the kernel args ptr in s[0:1]
+  .amdhsa_user_sgpr_kernarg_segment_ptr 1
+  .amdhsa_user_sgpr_count 2
+  .amdhsa_user_sgpr_kernarg_preload_length 0
+  .amdhsa_user_sgpr_kernarg_preload_offset 0
+  # gfx90a / gfx940 specifics (RSRC3)
+  .amdhsa_accum_offset 248
+  .amdhsa_uses_dynamic_stack 0
+  .amdhsa_tg_split 0
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.kernels:
+  - .name: gemm
+    .symbol: gemm.kd
+    .args:
+      - .name: C
+        .address_space: global
+        .offset: 0
+        .size: 8
+        .value_kind: global_buffer
+        .value_type: bf16
+      - .name: B
+        .address_space: global
+        .offset: 8
+        .size: 8
+        .value_kind: global_buffer
+        .value_type: bf16
+      - .name: A
+        .address_space: global
+        .offset: 16
+        .size: 8
+        .value_kind: global_buffer
+        .value_type: bf16
+      - .name: sz
+        .offset: 24
+        .size: 4
+        .value_kind: by_value
+        .value_type: u32
+      - .name: num_wg
+        .offset: 28
+        .size: 4
+        .value_kind: by_value
+        .value_type: u32
+    .group_segment_fixed_size: 133120
+    .private_segment_fixed_size: 0
+    .kernarg_segment_align: 8
+    .kernarg_segment_size: 32
+    .max_flat_workgroup_size: 256
+    .sgpr_count: 88
+    .sgpr_spill_count: 0
+    .vgpr_count: 248
+    .vgpr_spill_count: 0
+    .wavefront_size: 64
+amdhsa.version:
+  - 1
+  - 0
+...
+.end_amdgpu_metadata

extra/gemm/asm/cdna/test.py

@@ -0,0 +1,73 @@
+# Run assembly on the AMD runtime and check correctness
+# VIZ=2 to profile
+import pathlib
+from tinygrad import Tensor, Device, dtypes, Context
+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, getenv
+
+fp = pathlib.Path(__file__).parent/"gemm.s"
+
+N = getenv("N", 8192)
+THREADS_PER_WG = 256
+NUM_WG = N//THREADS_PER_WG * N//THREADS_PER_WG
+
+assert N % THREADS_PER_WG == 0, "N must be divisible by THREADS_PER_WG"
+
+# ** generate inputs on CPU
+
+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
+
+# baseline tinygrad
+sched = C_tiny.schedule()
+assert len(sched) == 1
+eis:list[ExecItem] = [sched[-1].lower()]
+ast = sched[-1].ast
+
+# assembly gemm
+@track_rewrites(name=lambda ret: TracingKey(ret.name, (ret.function_name,), ret))
+def get_asm_prg() -> ProgramSpec:
+  src = (pathlib.Path(__file__).parent/"template.s").read_text().replace("INSTRUCTIONS", fp.read_text())
+  lib = Device[Device.DEFAULT].compiler.compile(src)
+  return ProgramSpec("gemm", src, Device.DEFAULT, ast, lib=lib, global_size=[NUM_WG, 1, 1], local_size=[THREADS_PER_WG, 1, 1],
+                     globals=[0, 1, 2], vars=[UOp.variable("SZ", 256, 8192), UOp.variable("NUM_WG", 1, 1024)])
+eis.append(ExecItem(ast, [C_asm.uop.buffer, from_torch(A).uop.buffer, from_torch(B).uop.buffer], fixedvars={"SZ":N, "NUM_WG":NUM_WG},
+                    prg=CompiledRunner(get_asm_prg())))
+
+with Context(DEBUG=2):
+  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

@@ -0,0 +1,179 @@
+# unpack the complete kernel descriptor of an amdgpu ELF
+# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#code-object-v3-kernel-descriptor
+import struct, pathlib, sys
+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(sys.argv[1], "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 (gfx950)
+
+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)
+
+# gfx950 only
+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/gemm/triton_nv_matmul.py

@@ -4,8 +4,7 @@ import triton.language as tl
 from triton.compiler import AttrsDescriptor, ASTSource, compile as triton_compile
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.realize import CompiledRunner, ExecItem
+from tinygrad.engine.realize import CompiledRunner, ExecItem, ProgramSpec
 from tinygrad.helpers import getenv
 np.set_printoptions(suppress=True)
 
@@ -86,12 +85,9 @@ if __name__ == "__main__":
   # remove debug sections
   src = src.split("\t.file")[0]
   assert '.extern .shared' not in src
-  # Create linearized uops with SPECIAL for global/local sizes
-  global_size = [M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1]
-  local_size = [32*compiled.metadata.num_warps, 1, 1]
-  uops = [UOp(Ops.SPECIAL, arg=('g', i), src=(UOp.const(dtypes.int, global_size[i]),)) for i in range(3)]
-  uops += [UOp(Ops.SPECIAL, arg=('l', i), src=(UOp.const(dtypes.int, local_size[i]),)) for i in range(3)]
-  prg = UOp.new_program("matmul_kernel", src, Device.DEFAULT, si.ast, uops)
+  prg = ProgramSpec("matmul_kernel", src, device=Device.DEFAULT,
+                global_size=[M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1], local_size=[32*compiled.metadata.num_warps, 1, 1],
+                mem_estimate=A.nbytes() + B.nbytes() + C.nbytes())
   ei = ExecItem(si.ast, [x.ensure_allocated() for x in si.bufs], si.metadata, prg=CompiledRunner(prg))
   tflops = []
   for i in range(5):

extra/mmapeak/mmapeak.py

@@ -4,7 +4,8 @@ import os, pathlib
 os.environ["AMD_AQL"] = "1"
 
 from tinygrad.device import Device
-from tinygrad.runtime.ops_amd import AMDProgram, HIPCompiler
+from tinygrad.runtime.support.compiler_amd import HIPCompiler
+from tinygrad.runtime.ops_amd import AMDProgram
 
 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32

extra/optimization/test_beam_search.py

@@ -1,9 +1,13 @@
 import unittest
 import numpy as np
 
-from tinygrad.helpers import BEAM, Timing, CI, Context
-from tinygrad import Variable, Tensor
+from tinygrad.helpers import BEAM, Timing, CI, prod
+from tinygrad import Variable, Device, Tensor
 from tinygrad.nn import Conv2d
+from tinygrad.uop.ops import AxisType
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.codegen.opt.postrange import Scheduler
+from tinygrad.codegen.opt.search import get_kernel_actions
 
 def rand(*shape):
   return Tensor(np.random.rand(*shape).astype(np.float32))
@@ -75,5 +79,27 @@ class TestBeamSearch(unittest.TestCase):
         a = (a + a) * a
     a.realize()
 
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  def test_tc_up(self):
+    tc = Device[Device.DEFAULT].renderer.tensor_cores[0]
+    size = max(tc.dims[0], tc.dims[1]) * 8
+    a, b = Tensor.rand(size, size, dtype=tc.dtype_in), Tensor.rand(size, size, dtype=tc.dtype_in)
+    ast = a.matmul(b, dtype=tc.dtype_out).schedule()[-1].ast
+    s = Scheduler(ast, Device[Device.DEFAULT].renderer)
+    s.apply_opt(Opt(OptOps.TC, 0, (-1, 0, 1)))
+    up = prod([x for x, t in zip(s.full_shape, s.axis_types) if t in (AxisType.UPCAST, AxisType.UNROLL)])
+    actions = get_kernel_actions(s, include_0=False, max_up=int(up))
+    upcasted = [s for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]
+    assert len(upcasted) > 0, f"expected upcast/unroll actions after TC with max_up={up}, but got none"
+
+  def test_max_up(self):
+    a = Tensor.rand(16, 16)
+    ast = a.schedule()[-1].ast
+    s = Scheduler(ast, Device[Device.DEFAULT].renderer)
+    for max_up in (2, 4):
+      actions = get_kernel_actions(s, include_0=False, max_up=max_up)
+      for up_opts in [s.applied_opts for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]:
+        assert len([opt for opt in up_opts if opt.arg > max_up]) == 0 and len([op for op in up_opts if op.arg <= max_up]) > 0
+
 if __name__ == '__main__':
   unittest.main()

extra/remu/src/lib.rs

@@ -1,4 +1,5 @@
-use crate::work_group::WorkGroup;
+use crate::state::StateSnapshot;
+use crate::work_group::{WaveContext, WorkGroup};
 use std::os::raw::c_char;
 use std::slice;
 mod helpers;
@@ -30,3 +31,47 @@ 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,6 +96,24 @@ 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, VecDataStore, WaveValue, VGPR};
+use crate::state::{Register, StateSnapshot, VecDataStore, WaveValue, VGPR};
 use crate::thread::{Thread, END_PRG, SGPR_COUNT};
 use std::collections::HashMap;
 
@@ -28,6 +28,96 @@ 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,34 +1,38 @@
+# ruff: noqa: F405, F403
+# allow define from star imports
+
 import numpy as np
 import unittest
-import subprocess, struct, math
-from tinygrad import Tensor, dtypes, Device
-from tinygrad.uop.ops import UOp, Ops
+import subprocess, struct, math, textwrap, functools
+from tinygrad import Tensor, dtypes, Device, UOp
+from tinygrad.uop.ops import Ops, KernelInfo
 from tinygrad.helpers import getenv
-from tinygrad.runtime.support.compiler_amd import amdgpu_disassemble
-from tinygrad.engine.realize import CompiledRunner
 
-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()
-  # Create linearized uops with SPECIAL for local size
-  uops = [UOp(Ops.SPECIAL, arg=('l', 0), src=(UOp.const(dtypes.int, n_threads),))]
-  prg = UOp.new_program("test", src, Device.DEFAULT, UOp.sink(t.uop), uops)
-  car = CompiledRunner(prg)
-  if getenv("PRINT_ASM"): amdgpu_disassemble(car.lib)
-  car([t.uop.buffer], {}, wait=True)
-  return t.numpy()
+from extra.assembly.amd.autogen.rdna3.ins import *
+from extra.assembly.amd.asm import waitcnt
+from test.testextra.test_cfg_viz import template
+
+def custom_src(out:UOp, src:str, device:str, n_threads:int=1, n_workgroups:int=1) -> UOp:
+  lidx = UOp.special(n_threads, "lidx0")
+  gidx = UOp.special(n_workgroups, "gidx0")
+  sink = UOp.sink(out, lidx, gidx, arg=KernelInfo(name="test"))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=device), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
+
+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()
+  ])
+  src = template.replace("fn_name", "test").replace("INSTRUCTION", textwrap.dedent(src))
+  out = Tensor.custom_kernel(out, fxn=functools.partial(custom_src, src=src, device=out.device, n_threads=n_threads))[0]
+  out.realize()
+  return out.tolist()
 
 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]
@@ -39,83 +43,89 @@ 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(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_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_exec_mov(self):
-    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)
+    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)
     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 %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)
+    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)
 
   def test_exec_cmpx_vop3(self):
-    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]
+    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]
     np.testing.assert_equal(out & 0b11, 0b01)
 
   def test_fmac_vop3_modifier(self):
-    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.))
+    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))
 
   def test_s_abs_i32(self):
-    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 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 test_v_rcp_f32_neg_vop3(self):
     def v_neg_rcp_f32(x:float, y:float):
-      out = get_output(f"""
-      v_mov_b32_e32 %2 {f32_to_bits(x)}
-      v_rcp_f32_e64 %2, -%2
-      """)[0]
+      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]
       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)
@@ -124,25 +134,28 @@ class TestHW(unittest.TestCase):
     v_neg_rcp_f32(2.0, -0.5)
 
   def test_v_cndmask_b32_neg(self):
-    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]
+    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]
       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 %1, 0xffffffff :: v_dual_mov_b32 %2, 0x0
-    v_subrev_co_u32 %2, vcc_lo, %2, %1
-    """)[0]
+    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]
     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

@@ -1,11 +1,9 @@
-import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools, threading
+#!/usr/bin/env python3
+import ctypes, pathlib, argparse, pickle, dataclasses, threading
 from typing import Generator
 from tinygrad.helpers import temp, unwrap, DEBUG
-from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
-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
+from tinygrad.runtime.ops_amd import ProfileSQTTEvent
+from tinygrad.runtime.autogen import rocprof
 
 @dataclasses.dataclass(frozen=True)
 class InstExec:
@@ -48,17 +46,11 @@ class OccEvent(WaveSlot):
 RunKey = tuple[str, int]
 
 class _ROCParseCtx:
-  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]]] = {}
+  def __init__(self, sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, tuple[str, int]]]):
+    self.sqtt_evs, self.disasms = iter(sqtt_evs), disasms
     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
@@ -81,16 +73,8 @@ 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(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)
+def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, tuple[str, int]]]) -> _ROCParseCtx:
+  ROCParseCtx = _ROCParseCtx(sqtt_evs, disasms)
 
   @rocprof.rocprof_trace_decoder_se_data_callback_t
   def copy_cb(buf, buf_size, _):
@@ -131,26 +115,52 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
 
   def worker():
     try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
-    except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
+    except AttributeError as e:
+      raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
   (t:=threading.Thread(target=worker, daemon=True)).start()
   t.join()
   return ROCParseCtx
 
-def print_pmc(events:list[ProfilePMCEvent]) -> None:
-  from tinygrad.viz.serve import unpack_pmc
+def print_data(data:dict) -> None:
   from tabulate import tabulate
-  for e in events:
-    print("**", e.kern)
-    data = unpack_pmc(e)
-    print(tabulate([r[:-1] for r in data["rows"]], headers=data["cols"], tablefmt="github"))
+  # plaintext
+  if "src" in data: print(data["src"])
+  # table format
+  elif "cols" in data:
+    print(tabulate([r[:len(data["cols"])] for r in data["rows"]], headers=data["cols"], tablefmt="github"))
+
+def main() -> None:
+  import tinygrad.viz.serve as viz
+  viz.ctxs = []
 
-if __name__ == "__main__":
   parser = argparse.ArgumentParser()
-  parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  parser.add_argument('--profile', type=pathlib.Path, metavar="PATH", help='Path to profile (optional file, default: latest profile)',
+                      default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  parser.add_argument('--kernel', type=str, default=None, metavar="NAME", help='Kernel to focus on (optional name, default: all kernels)')
+  parser.add_argument('-n', type=int, default=3, metavar="NUM", help='Max traces to print (optional number, default: 3 traces)')
   args = parser.parse_args()
 
   with args.profile.open("rb") as f: profile = pickle.load(f)
-  rctx = decode(profile)
-  print('SQTT:', rctx.inst_execs.keys())
 
-  print_pmc([ev for ev in profile if isinstance(ev, ProfilePMCEvent)])
+  viz.get_profile(profile)
+
+  # List all kernels
+  if args.kernel is None:
+    for c in viz.ctxs:
+      print(c["name"])
+      for s in c["steps"]: print("  "+s["name"])
+    return None
+
+  # Find kernel trace
+  trace = next((c for c in viz.ctxs if c["name"] == f"Exec {args.kernel}"), None)
+  if not trace: raise RuntimeError(f"no matching trace for {args.kernel}")
+  n = 0
+  for s in trace["steps"]:
+    print(s["name"])
+    data = viz.get_render(s["query"])
+    print_data(data)
+    n += 1
+    if n > args.n: break
+
+if __name__ == "__main__":
+  main()

pyproject.toml

@@ -49,7 +49,7 @@ arm = ["unicorn"]
 triton = ["triton-nightly>=2.1.0.dev20231014192330"]
 linting = [
   "pylint",
-  "mypy==1.18.1",
+  "mypy==1.19.1",
   "typing-extensions",
   "pre-commit",
   "ruff",
@@ -61,7 +61,7 @@ linting = [
 # ]
 testing_minimal = [
   "numpy",
-  "torch==2.9.0",
+  "torch==2.9.1",
   "pytest",
   "pytest-xdist",
   "pytest-timeout",
@@ -135,9 +135,14 @@ check_untyped_defs = true
 explicit_package_bases = true
 warn_unreachable = true
 warn_redundant_casts = true
+strict_equality = true
 # NOTE: had to comment this out to make mypy pass on both CI and OSX
 #warn_unused_ignores = true
 
+[[tool.mypy.overrides]]
+module = "extra.*"
+follow_imports = "skip"
+
 [tool.pytest.ini_options]
 norecursedirs = [
   "extra",
@@ -188,7 +193,6 @@ select = [
   "E72",
   "E112",   # no-indented-block
   "E113",   # unexpected-indentation
-  # "E124",
   "E203",   # whitespace-before-punctuation
   "E272",   # multiple-spaces-before-keyword
   "E275",   # missing-whitespace-after-keyword

test/device/test_hcq.py

@@ -76,7 +76,7 @@ class TestHCQ(unittest.TestCase):
         TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
         TestHCQ.d0.timeline_value += 1
 
-  @unittest.skipIf(MOCKGPU or Device.DEFAULT in {"CPU"}, "Can't handle async update on MOCKGPU for now")
+  @unittest.skipIf(Device.DEFAULT in {"CPU"}, "Can't handle async update on CPU device")
   def test_wait_late_set(self):
     for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
       if queue_type is None: continue

test/device/test_qcom.py

@@ -1,74 +0,0 @@
-#!/usr/bin/env python
-import unittest
-from tinygrad.device import Device, BufferSpec
-from tinygrad.dtype import dtypes
-
-@unittest.skipUnless(Device.DEFAULT == "QCOM", "QCOM device required to run")
-class TestQcom(unittest.TestCase):
-  def test_image_pitch(self):
-    dev = Device["QCOM"]
-
-    def __validate(imgdt, expected_pitch):
-      img = dev.allocator.alloc(imgdt.shape[0] * imgdt.shape[1] * 16, options:=BufferSpec(image=imgdt))
-      pitch = img.texture_info.pitch
-      assert pitch == expected_pitch, f"Failed pitch for image: {imgdt}. Got 0x{pitch:X}, expected 0x{expected_pitch:X}"
-      dev.allocator.free(img, imgdt.shape[0] * imgdt.shape[1] * 16, options)
-
-    # Match opencl pitches for perf
-    __validate(dtypes.imageh((1, 201)), 0x680)
-    __validate(dtypes.imageh((16, 216)), 0x700)
-    __validate(dtypes.imageh((16, 9)), 0x80)
-    __validate(dtypes.imageh((48, 64)), 0x200)
-    __validate(dtypes.imageh((32, 128)), 0x400)
-    __validate(dtypes.imageh((96, 128)), 0x400)
-    __validate(dtypes.imageh((64, 256)), 0x840)
-    __validate(dtypes.imageh((64, 9)), 0x80)
-    __validate(dtypes.imageh((192, 256)), 0x840)
-    __validate(dtypes.imageh((64, 768)), 0x1840)
-    __validate(dtypes.imageh((256, 49)), 0x1C0)
-    __validate(dtypes.imageh((128, 9)), 0x80)
-    __validate(dtypes.imageh((16, 1024)), 0x2080)
-    __validate(dtypes.imageh((64, 512)), 0x1040)
-    __validate(dtypes.imageh((16, 512)), 0x1080)
-    __validate(dtypes.imageh((132, 64)), 0x200)
-    __validate(dtypes.imageh((4, 512)), 0x1200)
-    __validate(dtypes.imageh((8, 512)), 0x1100)
-    __validate(dtypes.imageh((128, 128)), 0x400)
-    __validate(dtypes.imageh((32, 512)), 0x1040)
-    __validate(dtypes.imageh((26, 64)), 0x200)
-    __validate(dtypes.imageh((32, 516)), 0x1040)
-    __validate(dtypes.imageh((32, 1024)), 0x2040)
-    __validate(dtypes.imageh((16, 2048)), 0x4080)
-    __validate(dtypes.imageh((8, 2048)), 0x4100)
-    __validate(dtypes.imageh((4, 4096)), 0x8200)
-
-    __validate(dtypes.imagef((16, 49)), 0x380)
-    __validate(dtypes.imagef((16, 1024)), 0x4080)
-    __validate(dtypes.imagef((256, 64)), 0x400)
-    __validate(dtypes.imagef((64, 512)), 0x2040)
-    __validate(dtypes.imagef((16, 512)), 0x2080)
-    __validate(dtypes.imagef((132, 64)), 0x400)
-    __validate(dtypes.imagef((4, 512)), 0x2200)
-    __validate(dtypes.imagef((4, 16)), 0x200)
-    __validate(dtypes.imagef((2, 16)), 0x400)
-    __validate(dtypes.imagef((8, 512)), 0x2100)
-    __validate(dtypes.imagef((12, 64)), 0x400)
-    __validate(dtypes.imagef((3, 32)), 0x400)
-    __validate(dtypes.imagef((128, 128)), 0x840)
-    __validate(dtypes.imagef((32, 512)), 0x2040)
-    __validate(dtypes.imagef((8, 3072)), 0xC100)
-    __validate(dtypes.imagef((4, 2048)), 0x8200)
-    __validate(dtypes.imagef((4, 1024)), 0x4200)
-    __validate(dtypes.imagef((4, 4096)), 0x10200)
-    __validate(dtypes.imagef((10, 384)), 0x1900)
-    __validate(dtypes.imagef((24, 64)), 0x400)
-    __validate(dtypes.imagef((128, 12)), 0xC0)
-    __validate(dtypes.imagef((10, 24)), 0x200)
-    __validate(dtypes.imagef((1, 129)), 0x840)
-    __validate(dtypes.imagef((1, 32)), 0x200)
-    __validate(dtypes.imagef((1, 64)), 0x400)
-    __validate(dtypes.imagef((1, 1239)), 0x4D80)
-    __validate(dtypes.imagef((1, 1)), 0x40)
-
-if __name__ == "__main__":
-  unittest.main()

test/external/external_benchmark_multitensor_allreduce.py

@@ -1,7 +1,7 @@
 from tinygrad import Tensor, Device, GlobalCounters, TinyJit, dtypes
-from tinygrad.helpers import getenv, Context, RING, DEBUG
+from tinygrad.helpers import getenv, Context, DEBUG
 
-def test(devs: list[str], N: int, iters:int = 10):
+def test(devs: list[str], N: int, iters:int = 10, name:str = "allreduce"):
   @TinyJit
   def f(t: Tensor) -> Tensor: t.sum(0).realize()
 
@@ -17,39 +17,33 @@ def test(devs: list[str], N: int, iters:int = 10):
     i_secs = GlobalCounters.time_sum_s
     i_gflops = GlobalCounters.global_ops/i_secs/10**9
     i_gbs = (N*4)/i_secs/10**9
-    print(f"{'ring_allreduce' if RING >= 2 else 'naive_allreduce'} iter {i+1}/{iters}: {i_secs:.6f} sec {i_gflops:.2f} GFLOP/s {i_gbs:.2f} GB/s")
+    print(f"{name} iter {i+1}/{iters}: {i_secs:.6f} sec {i_gflops:.2f} GFLOP/s {i_gbs:.2f} GB/s")
     secs += i_secs
     gflops += i_gflops
     gbs += i_gbs
 
   return (gflops/iters, gbs/iters, secs/iters)
 
-def run(sz, n_gpus=6, iters=10, use_ring=False):
+def run(sz, n_gpus=6, iters=10, ring=0, all2all=0):
   devs = tuple([f"{Device.DEFAULT}:{x}" for x in range(n_gpus)])
   N = sz // dtypes.float32.itemsize
-  with Context(RING=(2 if use_ring else 0), DEBUG=max(DEBUG.value, 2)): return test(devs, N, iters=iters)
+  name = "all2all" if all2all else ("ring" if ring else "naive")
+  with Context(RING=(2 if ring else 0), ALL2ALL=(2 if all2all else 0), JIT_BATCH_SIZE=0, DEBUG=max(DEBUG.value, 2)):
+    return test(devs, N, iters=iters, name=name)
 
 def main():
-  ONLY_RING = getenv("ONLY_RING", 0)
   n_gpus = getenv("GPUS", 6)
   iters = getenv("ITERS", 10)
+  sz = getenv("SZ", 1000) * 10**6 # size of data on each gpu
+  print(f"Using {sz/10**9:.2f} GB of numbers on each of {n_gpus} GPUs, {n_gpus*sz/10**9:.2f} GB total.")
 
-  if getenv("BENCHMARK_SPLIT"):
-    l, r = 0, 512
-    while r - l > 1:
-      m = (l + r) // 2
-      (ring_gflops, ring_gbs, ring_secs) = run(m * 1024 * 4, n_gpus=n_gpus, iters=100, use_ring=True)
-      (naive_gflops, naive_gbs, naive_secs) = run(m * 1024 * 4, n_gpus=n_gpus, iters=100, use_ring=False)
-      if ring_secs > naive_secs: l = m
-      else: r = m
-    print("Better split", r * 1024, "elements")
-  else:
-    sz = getenv("SZ", 1000) * 10**6 # size of data on each gpu
-    print(f"Using {sz/10**9:.2f} GB of numbers on each of {n_gpus} GPUs, {n_gpus*sz/10**9:.2f} GB total.")
-    (ring_gflops, ring_gbs, ring_secs) = run(sz, use_ring=True, n_gpus=n_gpus, iters=iters)
-    if not ONLY_RING: (naive_gflops, naive_gbs, naive_secs) = run(sz, use_ring=False, n_gpus=n_gpus, iters=iters)
-    print(f"Ring:\n  {ring_secs:.6f} seconds/iter\n  {ring_gflops:.2f} GFLOP/s\n  {ring_gbs:.2f} GB/s")
-    if not ONLY_RING: print(f"Naive:\n  {naive_secs:.6f} seconds/iter\n  {naive_gflops:.2f} GFLOP/s\n  {naive_gbs:.2f} GB/s")
+  results = {}
+  for name, kwargs in [("naive", {}), ("ring", {"ring": 2}), ("all2all", {"all2all": 2})]:
+    results[name] = run(sz, n_gpus=n_gpus, iters=iters, **kwargs)
+
+  print("\n=== RESULTS ===")
+  for name, (gflops, gbs, secs) in results.items():
+    print(f"{name.upper()}:\n  {secs:.6f} seconds/iter\n  {gflops:.2f} GFLOP/s\n  {gbs:.2f} GB/s")
 
 if __name__ == "__main__":
   main()

test/external/external_benchmark_op_conv.py

@@ -1,4 +1,5 @@
 # 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.opt import Opt, OptOps # pylint: disable=unused-import
@@ -88,9 +89,7 @@ renderer = Device.default.renderer
 allocator = Device.default.allocator
 
 ps = get_program(ast, renderer)
-# update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
-ps = ps.replace(src=(ps.src[0], UOp(Ops.DEVICE, arg=Device.DEFAULT), *ps.src[2:]))
-cr = CompiledRunner(ps)
+cr = CompiledRunner(replace(ps, device=Device.DEFAULT))
 
 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_hip_compiler_bug.py

@@ -1,225 +0,0 @@
-# [<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/process_replay/process_replay.py

@@ -9,7 +9,7 @@ if not int(os.getenv("ASSERT_PROCESS_REPLAY", "1")): ASSERT_DIFF = 0
 
 try:
   from tinygrad.schedule.rangeify import get_rangeify_map
-  from tinygrad.renderer import Renderer
+  from tinygrad.renderer import Renderer, ProgramSpec
   from tinygrad.engine.realize import get_program
   from tinygrad.uop.ops import UOp, Ops, KernelInfo
   from tinygrad.codegen.opt import Opt
@@ -51,20 +51,15 @@ def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str,
     return "\n".join([f"{len(asts)} kernels", *asts])
   return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,)
 
-def replay_get_program(p:UOp, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
+def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
   # the ast.arg is non None if we are inside of search.py
-  # p is a PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
-  p_name = p.src[0].arg.name
-  p_applied_opts = p.src[0].arg.applied_opts
-  sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p_applied_opts if BEAM>=1 else None)
-  input_ast = ast.replace(arg=replace(sink_arg, name=p_name))
+  sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p.applied_opts if BEAM>=1 else None)
+  input_ast = ast.replace(arg=replace(sink_arg, name=p.name))
   p2 = get_program(input_ast, renderer=renderer)
-  def to_str(ret:UOp) -> str:
+  def to_str(ret:ProgramSpec) -> str:
     # PYTHON renderer pickles UOps, first unpickle and decode here
-    ret_src = ret.src[3].arg
-    ret_device = ret.device
-    if ret_device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(ret_src))])
-    return ret_src
+    if p.device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(ret.src))])
+    return ret.src
   # properly color the name arg
   ast_repr = codecs.decode(str(input_ast), "unicode_escape")
   return to_str(p2), to_str(p), (ast_repr, renderer)

test/helpers.py

@@ -63,7 +63,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
     allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *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[3].arg))  # source code is in src[3].arg
+  prog = PythonProgram("run", PythonCompiler().compile(prg.src))
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs)
   return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
 

test/mockgpu/amd/amddriver.py

@@ -52,6 +52,7 @@ class AMDDriver(VirtDriver):
     self.doorbells = {}
     self.next_doorbell = collections.defaultdict(int)
     self.mmu_event_ids = []
+    self._executing = False  # re-entrancy guard for _emulate_execute
 
     for i in range(gpus): self._prepare_gpu(i+1)
 
@@ -125,6 +126,9 @@ class AMDDriver(VirtDriver):
       if struct.gpu_id not in self.gpus: return -1
       struct.handle = self._alloc_handle()
       self.object_by_handle[struct.handle] = copy.deepcopy(struct) # save memory struct to know what mem it is
+      # Track signal memory (uncached + coherent) - progress queues when written to
+      if struct.flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED:
+        self.track_address(struct.va_addr, struct.va_addr + struct.size, lambda mv,off: None, lambda mv, off: self._emulate_execute())
     elif nr == kfd_ioctls.AMDKFD_IOC_FREE_MEMORY_OF_GPU:
       self.object_by_handle.pop(struct.handle)
     elif nr == kfd_ioctls.AMDKFD_IOC_MAP_MEMORY_TO_GPU:
@@ -173,9 +177,14 @@ class AMDDriver(VirtDriver):
     return 0
 
   def _emulate_execute(self):
-    any_progress = True
-    while any_progress:
-      any_progress = False
-      for gpu in self.gpus.values():
-        for q in gpu.queues:
-          if q.executing: any_progress |= q.execute() > 0
+    if self._executing: return  # prevent re-entrancy
+    self._executing = True
+    try:
+      any_progress = True
+      while any_progress:
+        any_progress = False
+        for gpu in self.gpus.values():
+          for q in gpu.queues:
+            if q.executing: any_progress |= q.execute() > 0
+    finally:
+      self._executing = False

test/mockgpu/amd/amdgpu.py

@@ -7,6 +7,7 @@ import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.
 SDMA_MAX_COPY_SIZE = 0x400000
 
 regCOMPUTE_PGM_LO = 0x1bac + amd_gpu.GC_BASE__INST0_SEG0
+regCOMPUTE_PGM_RSRC2 = 0x1bb3 + amd_gpu.GC_BASE__INST0_SEG0
 regCOMPUTE_USER_DATA_0 = 0x1be0 + amd_gpu.GC_BASE__INST0_SEG0
 regCOMPUTE_NUM_THREAD_X = 0x1ba7 + amd_gpu.GC_BASE__INST0_SEG0
 regGRBM_GFX_INDEX = 0x2200 + amd_gpu.GC_BASE__INST0_SEG1
@@ -179,12 +180,16 @@ class PM4Executor(AMDQueue):
     prg_addr = (self.gpu.regs[regCOMPUTE_PGM_LO] + (self.gpu.regs[regCOMPUTE_PGM_LO + 1] << 32)) << 8
     args_addr = self.gpu.regs[regCOMPUTE_USER_DATA_0] + (self.gpu.regs[regCOMPUTE_USER_DATA_0 + 1] << 32)
     lc = [self.gpu.regs[i] for i in range(regCOMPUTE_NUM_THREAD_X, regCOMPUTE_NUM_THREAD_X+3)]
+    rsrc2 = self.gpu.regs[regCOMPUTE_PGM_RSRC2]
 
     prg_sz = 0
     for st,sz in self.gpu.mapped_ranges:
       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 and rsrc2 to Python emulator for bounds checking and SGPR/VGPR layout
+    if hasattr(remu, 'valid_mem_ranges'): remu.valid_mem_ranges = self.gpu.mapped_ranges
+    if hasattr(remu, 'rsrc2'): remu.rsrc2 = rsrc2
     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,4 +1,5 @@
 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 []
@@ -14,7 +15,21 @@ 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()
+  rsrc2: int = 0x19c  # Default: USER_SGPR_COUNT=14, enable X and Y workgroup IDs
+
+  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.amd.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, self.rsrc2)
+
 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/mockgpu/nv/nvdriver.py

@@ -15,7 +15,7 @@ libc.munmap.restype = ctypes.c_int
 NVSubDevice = collections.namedtuple('NVSubDevice', ['device'])
 NVUserMode = collections.namedtuple('NVUserMode', ['subdevice'])
 NVVASpace = collections.namedtuple('NVVASpace', ['device'])
-NVAllocation = collections.namedtuple('NVAllocation', ['device', 'size'])
+NVAllocation = collections.namedtuple('NVAllocation', ['device', 'size', 'is_signal'])
 NVChannelGroup = collections.namedtuple('NVChannelGroup', ['device'])
 NVContextShare = collections.namedtuple('NVContextShare', ['channel_group'])
 NVGPFIFO = collections.namedtuple('NVGPFIFO', ['device', 'token'])
@@ -41,12 +41,14 @@ class NVDevFileDesc(VirtFileDesc):
     super().__init__(fd)
     self.driver, self.gpu = driver, gpu
     self._mapping_userland = False
+    self._mapping_signal = False
 
   def ioctl(self, fd, request, argp): return self.driver.dev_ioctl(self.gpu, request, argp)
   def mmap(self, start, sz, prot, flags, fd, offset):
     start = libc.mmap(start, sz, prot, flags|mmap.MAP_ANONYMOUS, -1, 0)
-    if self._mapping_userland:
+    if self._mapping_userland or self._mapping_signal:
       self.driver.track_address(start, start+sz, lambda mv,off: None, lambda mv, off: self.driver._gpu_mmio_write(mv, off, self.gpu))
+      self._mapping_signal = False
     return start
 
 class NVDriver(VirtDriver):
@@ -65,6 +67,7 @@ class NVDriver(VirtDriver):
     self.object_by_handle = {}
     self.opened_fds = {}
     self.next_doorbell = collections.defaultdict(int)
+    self._executing = False  # re-entrancy guard for _gpu_mmio_write
 
     for i in range(gpus): self._prepare_gpu(i)
 
@@ -115,7 +118,8 @@ class NVDriver(VirtDriver):
       assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPU)
       params = nv_gpu.NV_MEMORY_ALLOCATION_PARAMS.from_address(params_ptr)
       struct.hObjectNew = self._alloc_handle()
-      self.object_by_handle[struct.hObjectNew] = NVAllocation(self.object_by_handle[struct.hObjectParent], params.size)
+      is_signal = struct.hClass == nv_gpu.NV1_MEMORY_SYSTEM  # signal memory uses NV1_MEMORY_SYSTEM (uncached)
+      self.object_by_handle[struct.hObjectNew] = NVAllocation(self.object_by_handle[struct.hObjectParent], params.size, is_signal)
     elif struct.hClass == nv_gpu.KEPLER_CHANNEL_GROUP_A:
       assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPU)
       struct.hObjectNew = self._alloc_handle()
@@ -206,7 +210,6 @@ class NVDriver(VirtDriver):
   def ctl_ioctl(self, req, argp):
     nr = req & 0xff
     if nr == nv_gpu.NV_ESC_RM_ALLOC: return self.rm_alloc(argp)
-    elif nr == nv_gpu.NV_ESC_RM_ALLOC_MEMORY: pass
     elif nr == nv_gpu.NV_ESC_RM_CONTROL: return self.rm_control(argp)
     elif nr == nv_gpu.NV_ESC_RM_MAP_MEMORY:
       st:Any = nv_gpu.nv_ioctl_nvos33_parameters_with_fd.from_address(argp)
@@ -215,6 +218,10 @@ class NVDriver(VirtDriver):
         file = self.opened_fds[st.fd]
         assert isinstance(file, NVDevFileDesc)
         file._mapping_userland = True
+      elif isinstance(obj, NVAllocation) and obj.is_signal:
+        file = self.opened_fds[st.fd]
+        assert isinstance(file, NVDevFileDesc)
+        file._mapping_signal = True
     elif nr == nv_gpu.NV_ESC_RM_FREE:
       st = nv_gpu.NVOS00_PARAMETERS.from_address(argp)
       self.object_by_handle.pop(st.hObjectOld)
@@ -256,12 +263,26 @@ class NVDriver(VirtDriver):
     else: raise RuntimeError(f"Unknown {nr} to nvidia-uvm")
     return 0
 
-  def dev_ioctl(self, dev, req, argp): return 0
+  def dev_ioctl(self, dev, req, argp):
+    nr = req & 0xff
+    # Handle NV_ESC_RM_ALLOC_MEMORY for host/signal memory
+    if nr == nv_gpu.NV_ESC_RM_ALLOC_MEMORY:
+      st:Any = nv_gpu.nv_ioctl_nvos02_parameters_with_fd.from_address(argp)
+      # Track host memory (signal memory) - progress queues when written to
+      if st.params.hClass == nv_gpu.NV01_MEMORY_SYSTEM_OS_DESCRIPTOR:
+        self.track_address(st.params.pMemory, st.params.pMemory + st.params.limit + 1,
+                           lambda mv,off: None, lambda mv, off: self._gpu_mmio_write(mv, off, None))
+    return 0
   def _gpu_mmio_write(self, mv, off, gpu):
-    any_progress = True
-    while any_progress:
-      any_progress = False
-      for gpu in self.gpus.values():
-        for q in gpu.queues:
-          if q.ctrl.GPGet != q.ctrl.GPPut:
-            any_progress |= q.execute()
+    if self._executing: return  # prevent re-entrancy
+    self._executing = True
+    try:
+      any_progress = True
+      while any_progress:
+        any_progress = False
+        for gpu in self.gpus.values():
+          for q in gpu.queues:
+            if q.ctrl.GPGet != q.ctrl.GPPut:
+              any_progress |= q.execute()
+    finally:
+      self._executing = False

test/models/test_bert.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env python
+import os
+os.environ['USE_TF'] = '0'  # prevent transformers from importing tensorflow
 import unittest
 from tinygrad import Tensor
 import numpy as np

test/models/test_efficientnet.py

@@ -68,7 +68,7 @@ class TestEfficientNet(unittest.TestCase):
     self.assertEqual(_LABELS[labels[0]], "sports car, sport car")
 
   def test_chicken_car(self):
-    labels = _infer(self.model, np.concat([chicken_img, car_img], axis=0))
+    labels = _infer(self.model, np.concatenate([chicken_img, car_img], axis=0))
     self.assertEqual(_LABELS[labels[0]], "hen")
     self.assertEqual(_LABELS[labels[1]], "sports car, sport car")
 

test/opt/test_tensor_cores.py

@@ -1,9 +1,10 @@
 import numpy as np
 import unittest
+from dataclasses import replace
 
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.uop.ops import Ops, UOp
+from tinygrad.uop.ops import Ops
 from tinygrad.dtype import DType
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import AMX, AMD_LLVM, CPU_LLVM, Context
@@ -43,10 +44,7 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi
   if dtype_in == dtypes.bfloat16: r = r.float()
   realized_ast, bufs = helper_realized_ast(r)
   opts = [Opt(op=OptOps.TC, axis=axis, arg=(tc_select, tc_opt, use_tensor_cores))]
-  p = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts)
-  # update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
-  p = p.replace(src=(p.src[0], UOp(Ops.DEVICE, arg=Device.DEFAULT), *p.src[2:]))
-  prg = CompiledRunner(p)
+  prg = CompiledRunner(replace(get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts), device=Device.DEFAULT))
   if use_tensor_cores == 1: assert len([uop for uop in prg.p.uops if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
   assert len([x for x in prg.p.uops[-1].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
   prg.exec(bufs)

test/test_custom_kernel.py

@@ -193,6 +193,16 @@ class TestCustomKernel(unittest.TestCase):
     err = (tst - (a@b)).square().max()
     self.assertLess(err.item(), 1e-6)
 
+  def test_gemm_multi(self):
+    devs = ("CPU:0", "CPU:1")
+    N = 16
+    a = Tensor.randn(N, N).shard_(devs, axis=0)
+    b = Tensor.randn(N, N).to(devs)
+    c = Tensor(Tensor.empty(N//2, N, device=devs).uop.multi(0), device=devs)
+    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
+    err = (tst - (a@b)).square().max()
+    self.assertLess(err.item(), 1e-6)
+
   def test_gemm_backward_custom(self): self.test_gemm_backward(True)
   # NOTE: grad_fxn doesn't work with pyrender
   def test_gemm_backward(self, custom_backward_gemm=False):

test/test_dtype_alu.py

@@ -1,7 +1,7 @@
 import unittest, operator, math
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.dtype import DType, truncate
-from tinygrad.helpers import CI, getenv, CPU_LLVM
+from tinygrad.helpers import CI, getenv
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
 from tinygrad.runtime.ops_python import from_storage_scalar
@@ -9,7 +9,7 @@ from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
 import numpy as np
 import pytest
-from hypothesis import assume, given, strategies as strat, settings, HealthCheck
+from hypothesis import assume, given, strategies as strat, settings
 
 pytestmark = pytest.mark.filterwarnings("ignore")
 
@@ -48,7 +48,7 @@ class ht:
   int32 = strat.integers(-2147483648, 2147483647)
   int64 = strat.integers(-9223372036854775808, 9223372036854775807)
   bool = strat.booleans()
-ht.bfloat16 = ht.uint16
+ht.bfloat16 = ht.uint16.filter(lambda x: ((x >> 7) & 0xFF) != 0)  # filter subnormal bfloat16
 ht.fp8e4m3 = ht.uint8
 ht.fp8e5m2 = ht.uint8
 
@@ -138,7 +138,6 @@ class TestDTypeALU(unittest.TestCase):
   def test_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)
 
   @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
-  @unittest.skipIf(CPU_LLVM, "bfloat16 precision issues with CPU_LLVM")
   @given(ht.bfloat16, strat.sampled_from(unary_operations))
   def test_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
 
@@ -206,29 +205,23 @@ class TestDTypeALU(unittest.TestCase):
   @given(ht.int32, strat.sampled_from(dtypes_float+dtypes_int+dtypes_bool))
   def test_int32_cast(self, a, dtype): universal_test_cast(a, dtypes.int32, dtype)
 
-  @settings(suppress_health_check=[HealthCheck.filter_too_much])
-  @given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
+  @given(strat.floats(width=32, min_value=1.0, max_value=254.0, allow_subnormal=False),
+         strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned(self, a, float_dtype, unsigned_dtype):
     if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
-    float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
-    float_strat = float_strat.filter(lambda x: 0 < x < dtypes.max(unsigned_dtype))
-    universal_test_cast(a.draw(float_strat), float_dtype, unsigned_dtype)
+    universal_test_cast(a, float_dtype, unsigned_dtype)
 
-  @settings(suppress_health_check=[HealthCheck.filter_too_much])
-  @given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
+  @given(strat.floats(width=32, min_value=256.0, max_value=65000.0, allow_subnormal=False),
+         strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned_overflow(self, a, float_dtype, unsigned_dtype):
     if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
-    float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
-    overflow_strat = float_strat.filter(lambda x: x > dtypes.max(unsigned_dtype) and x <= dtypes.max(dtypes.int32))
-    universal_test_cast(a.draw(overflow_strat), float_dtype, unsigned_dtype)
+    universal_test_cast(a, float_dtype, unsigned_dtype)
 
-  @settings(suppress_health_check=[HealthCheck.filter_too_much])
-  @given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
+  @given(strat.floats(width=32, min_value=-65000.0, max_value=-1.0, allow_subnormal=False),
+         strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned_underflow(self, a, float_dtype, unsigned_dtype):
     if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
-    float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
-    underflow_strat = float_strat.filter(lambda x: x < 0 and x >= dtypes.min(dtypes.int32))
-    universal_test_cast(a.draw(underflow_strat), float_dtype, unsigned_dtype)
+    universal_test_cast(a, float_dtype, unsigned_dtype)
 
   @unittest.expectedFailure
   def test_unsafe_cast_float_to_int_failure(self):

test/test_fusion_op.py

@@ -28,7 +28,7 @@ class TestFusionOp(unittest.TestCase):
     sched = a.schedule()
     sched[-1].lower()
     self.assertLess(time.perf_counter()-st, 2.0)
-    assert len(sched[-1].prg.p.src[3].arg.splitlines()) < 250  # source code is in src[3].arg
+    assert len(sched[-1].prg.p.src.splitlines()) < 250
 
   def test_recursive_add_cmp(self):
     st = time.perf_counter()

test/test_image_dtype.py

@@ -44,6 +44,66 @@ class TestImageCopy(unittest.TestCase):
 
 @unittest.skipUnless(REAL_DEV in IMAGE_SUPPORTED_DEVICES, "Images not supported")
 class TestImageDType(unittest.TestCase):
+  def test_image_pitch(self):
+    def __validate(imgdt, expected_pitch):
+      assert imgdt.pitch == expected_pitch, f"Failed pitch for image: {imgdt}. Got 0x{imgdt.pitch:X}, expected 0x{expected_pitch:X}"
+
+    # Match opencl pitches for perf
+    __validate(dtypes.imageh((1, 201)), 0x680)
+    __validate(dtypes.imageh((16, 216)), 0x700)
+    __validate(dtypes.imageh((16, 9)), 0x80)
+    __validate(dtypes.imageh((48, 64)), 0x200)
+    __validate(dtypes.imageh((32, 128)), 0x400)
+    __validate(dtypes.imageh((96, 128)), 0x400)
+    __validate(dtypes.imageh((64, 256)), 0x840)
+    __validate(dtypes.imageh((64, 9)), 0x80)
+    __validate(dtypes.imageh((192, 256)), 0x840)
+    __validate(dtypes.imageh((64, 768)), 0x1840)
+    __validate(dtypes.imageh((256, 49)), 0x1C0)
+    __validate(dtypes.imageh((128, 9)), 0x80)
+    __validate(dtypes.imageh((16, 1024)), 0x2080)
+    __validate(dtypes.imageh((64, 512)), 0x1040)
+    __validate(dtypes.imageh((16, 512)), 0x1080)
+    __validate(dtypes.imageh((132, 64)), 0x200)
+    __validate(dtypes.imageh((4, 512)), 0x1200)
+    __validate(dtypes.imageh((8, 512)), 0x1100)
+    __validate(dtypes.imageh((128, 128)), 0x400)
+    __validate(dtypes.imageh((32, 512)), 0x1040)
+    __validate(dtypes.imageh((26, 64)), 0x200)
+    __validate(dtypes.imageh((32, 516)), 0x1040)
+    __validate(dtypes.imageh((32, 1024)), 0x2040)
+    __validate(dtypes.imageh((16, 2048)), 0x4080)
+    __validate(dtypes.imageh((8, 2048)), 0x4100)
+    __validate(dtypes.imageh((4, 4096)), 0x8200)
+
+    __validate(dtypes.imagef((16, 49)), 0x380)
+    __validate(dtypes.imagef((16, 1024)), 0x4080)
+    __validate(dtypes.imagef((256, 64)), 0x400)
+    __validate(dtypes.imagef((64, 512)), 0x2040)
+    __validate(dtypes.imagef((16, 512)), 0x2080)
+    __validate(dtypes.imagef((132, 64)), 0x400)
+    __validate(dtypes.imagef((4, 512)), 0x2200)
+    __validate(dtypes.imagef((4, 16)), 0x200)
+    __validate(dtypes.imagef((2, 16)), 0x400)
+    __validate(dtypes.imagef((8, 512)), 0x2100)
+    __validate(dtypes.imagef((12, 64)), 0x400)
+    __validate(dtypes.imagef((3, 32)), 0x400)
+    __validate(dtypes.imagef((128, 128)), 0x840)
+    __validate(dtypes.imagef((32, 512)), 0x2040)
+    __validate(dtypes.imagef((8, 3072)), 0xC100)
+    __validate(dtypes.imagef((4, 2048)), 0x8200)
+    __validate(dtypes.imagef((4, 1024)), 0x4200)
+    __validate(dtypes.imagef((4, 4096)), 0x10200)
+    __validate(dtypes.imagef((10, 384)), 0x1900)
+    __validate(dtypes.imagef((24, 64)), 0x400)
+    __validate(dtypes.imagef((128, 12)), 0xC0)
+    __validate(dtypes.imagef((10, 24)), 0x200)
+    __validate(dtypes.imagef((1, 129)), 0x840)
+    __validate(dtypes.imagef((1, 32)), 0x200)
+    __validate(dtypes.imagef((1, 64)), 0x400)
+    __validate(dtypes.imagef((1, 1239)), 0x4D80)
+    __validate(dtypes.imagef((1, 1)), 0x40)
+
   def test_image_and_back(self):
     data = Tensor.randn(9*27*4).realize()
     tst = data.numpy()

test/test_linearizer.py

@@ -1,5 +1,6 @@
 import numpy as np
 import unittest
+from dataclasses import replace
 
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.codegen.gpudims import get_grouped_dims
@@ -167,8 +168,8 @@ class TestLinearizer(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
   def test_upcast_with_locals_cpu(self):
     out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
-    prg = get_program(out.schedule()[-1].ast, opts=[Opt(OptOps.LOCAL, axis=0, arg=4)])
-    self.assertEqual(len(prg.src[3].arg.split("for")), 5)  # source code is in src[3].arg
+    prg = get_program(out.schedule()[-1].ast, opts=[Opt(OptOps.LOCAL, axis=0, arg=4)]).uops
+    self.assertEqual(len(prg.src.split("for")), 5)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -516,11 +517,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
   device = real_bufs[0].device
   wanna_output = [np.array(x).flatten() for x in wanna_output]
 
-  def get_prg(opts):
-    prg = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts)
-    # update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
-    prg = prg.replace(src=(prg.src[0], UOp(Ops.DEVICE, arg=device), *prg.src[2:]))
-    return CompiledRunner(prg)
+  def get_prg(opts): return CompiledRunner(replace(get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts), device=device))
 
   def check_opt(opts):
     prg = get_prg(opts=opts)

test/test_multitensor.py

@@ -57,12 +57,18 @@ 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)
@@ -250,6 +256,11 @@ class TestMultiTensor(unittest.TestCase):
       a,b = _test_allreduce(Tensor.rand(256, 256))
       np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
 
+  def test_allreduce_all2all(self):
+    with Context(ALL2ALL=2):
+      a,b = _test_allreduce(Tensor.rand(256, 256))
+      np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
+
   def test_copy_jit(self):
     @TinyJit
     def copy_tensor(x:Tensor): return (x.to(f"{x.device.split(':')[0]}:1") + 1)

test/test_ops.py

@@ -848,7 +848,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], lambda x: x.cos())
     helper_test_op([()], lambda x: x.cos())
     if not ((getenv("MOCKGPU") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
-      helper_test_op(None, lambda x: x.sin(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
+      helper_test_op(None, lambda x: x.cos(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
       helper_test_op(None, lambda x: x.cos(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                     atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
   @unittest.skipIf(Device.DEFAULT == "WEBGPU" and platform.system() == "Windows", "Not accurate enough with DirectX backend")
@@ -859,8 +859,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], lambda x: x.tan(), low=-5, high=5)
     helper_test_op([()], lambda x: x.tan())
     if not ((getenv("MOCKGPU") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
-      helper_test_op(None, lambda x: x.sin(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
-      helper_test_op(None, lambda x: x.cos(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
+      helper_test_op(None, lambda x: x.tan(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
+      helper_test_op(None, lambda x: x.tan(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                     atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
 
   def test_asin(self):
@@ -1655,7 +1655,7 @@ class TestOps(unittest.TestCase):
   def test_broadcast_full(self):
     for torch_op, tinygrad_op in [(torch.add, Tensor.add), (torch.sub, Tensor.sub), (torch.mul, Tensor.mul),
                                   (torch.div, Tensor.div), (torch.pow, Tensor.pow)]:
-      for shapes in [((5,13,24,16), (5,1,24,1)), ((1,3,1,7,1), (2,1,5,1,8))]:
+      for shapes in [((5,3,14,16), (5,1,14,1)), ((1,3,1,7,1), (2,1,5,1,8))]:
         with self.subTest(op=torch_op.__name__, shapes=shapes):
           if tinygrad_op != Tensor.pow:
             helper_test_op(shapes, torch_op, tinygrad_op)
@@ -2078,7 +2078,7 @@ class TestOps(unittest.TestCase):
       lambda x,w: Tensor.conv2d(x,w,padding=[1,1,1,1,1,1]), grad_rtol=1e-5)
 
   def test_simple_conv2d_m4(self):
-    helper_test_op([(1,16,18,18), (16,16,3,3)],
+    helper_test_op([(1,16,9,9), (16,16,3,3)],
       lambda x,w: torch.nn.functional.conv2d(x,w),
       lambda x,w: Tensor.conv2d(x,w), atol=1e-05, grad_rtol=1e-5)
 
@@ -2535,7 +2535,7 @@ class TestOps(unittest.TestCase):
 
   @slow_test
   def test_avg_pool2d(self):
-    shape = (32,2,111,28)
+    shape = (32,2,11,28)
     for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]:
       with self.subTest(kernel_size=ksz):
         helper_test_op([shape],
@@ -2549,7 +2549,7 @@ class TestOps(unittest.TestCase):
 
   @slow_test
   def test_avg_pool2d_padding(self):
-    shape = (32,2,111,28)
+    shape = (32,2,11,28)
     for ksz in [(2,2), (3,3), 2, 3, (3,2)]:
       for p in [1, (1,0), (0,1)]:
         with self.subTest(kernel_size=ksz, padding=p):
@@ -2557,10 +2557,10 @@ class TestOps(unittest.TestCase):
             lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=ksz, padding=p),
             lambda x: Tensor.avg_pool2d(x, kernel_size=ksz, padding=p), rtol=1e-5)
     with self.assertRaises(ValueError):
-      Tensor.avg_pool2d(Tensor.randn((32,2,111,28)), kernel_size=(2,2), padding=(1,1,1))
+      Tensor.avg_pool2d(Tensor.randn((32,2,11,28)), kernel_size=(2,2), padding=(1,1,1))
 
   def test_avg_pool2d_asymmetric_padding(self):
-    shape = (32,2,111,28)
+    shape = (32,2,11,28)
     for p in [(0,1,0,1), (2,1,2,1), (2,0,2,1)]:
       with self.subTest(padding=p):
         helper_test_op([shape],
@@ -2571,7 +2571,7 @@ class TestOps(unittest.TestCase):
 
   @slow_test
   def test_avg_pool2d_padding_not_counted(self):
-    shape = (32,2,111,28)
+    shape = (32,2,11,28)
     for ksz in [(2,2), (3,3), 2, 3, (3,2)]:
       with self.subTest(kernel_size=ksz):
         helper_test_op([shape],
@@ -2607,9 +2607,9 @@ class TestOps(unittest.TestCase):
       lambda x: Tensor.avg_pool2d(x, kernel_size=(3,3), stride=3, padding=1, ceil_mode=True, count_include_pad=True))
 
   def test_global_avg_pool2d(self):
-    helper_test_op([(32,2,111,28)],
-      lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=(111,28)),
-      lambda x: Tensor.avg_pool2d(x, kernel_size=(111,28)), rtol=1e-5)
+    helper_test_op([(32,2,11,28)],
+      lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=(11,28)),
+      lambda x: Tensor.avg_pool2d(x, kernel_size=(11,28)), rtol=1e-5)
 
   def test_avg_pool3d(self):
     # TODO: AMD_LLVM has larger atol
@@ -3142,10 +3142,10 @@ class TestOps(unittest.TestCase):
         lambda x: x.log_softmax(axis=1).nll_loss(Tensor(target), Tensor(weight), reduction=r))
 
   def test_nll_loss_3d_weight(self):
-    target = np.random.randint(0, 10, (32,3,3,3), dtype=np.int32).tolist()
+    target = np.random.randint(0, 10, (16,3,3,3), dtype=np.int32).tolist()
     weight = np.random.normal(0, 1, (10,)).astype(np.float32).tolist()
     for r in ("mean", "sum", "none"):
-      helper_test_op([(32,10,3,3,3)],
+      helper_test_op([(16,10,3,3,3)],
           lambda x: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1), torch.tensor(target), torch.tensor(weight), reduction=r),
           lambda x: x.log_softmax(axis=1).nll_loss(Tensor(target), Tensor(weight), reduction=r))
 

test/test_opts.py

@@ -14,7 +14,7 @@ class TestOpts(unittest.TestCase):
     self.assertEqual(s[-1].ast.arg.opts_to_apply, opts)
     if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM and not CPU_LVP:
       prg = get_program(s[-1].ast, renderer=Device[Device.DEFAULT].renderer)
-      self.assertIn('float4', prg.src[3].arg)  # source code is in src[3].arg
+      self.assertIn('float4', prg.src)
 
 if __name__ == '__main__':
   unittest.main()

test/test_renderer_failures.py

@@ -1,5 +1,6 @@
 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
@@ -17,8 +18,8 @@ def _test_uop_result(inputs:list[Tensor], prg, local_size=None):
   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]
-  # update device in PROGRAM UOp: (SINK, DEVICE, LINEAR, SOURCE)
-  prg = prg.replace(src=(prg.src[0], UOp(Ops.DEVICE, arg=Device.DEFAULT), *prg.src[2:]))
+  prg = replace(prg, device=Device.DEFAULT)
+  if local_size is not None: prg = replace(prg, local_size=local_size)
   ei = CompiledRunner(prg)
   ei.exec(outbufs+inbufs)
   return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]
@@ -71,7 +72,7 @@ class TestCStyleFailures(unittest.TestCase):
     schedule = ret.schedule()
     assert len(schedule) == 1
     schedule[0].lower()
-    src = schedule[0].prg.p.src[3].arg  # source code is in src[3].arg
+    src = schedule[0].prg.p.src
     self.assertEqual("("*5 not in src, should_strip_paren)
 
   def test_repeat_add(self): self._test_src_strip_paren(Ops.ADD)

test/test_symbolic_ops.py

@@ -288,17 +288,22 @@ class TestSymbolicOps(unittest.TestCase):
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)
 
   def test_conv2d_ceildiv_edge_case(self):
-    v = Variable('v', 11, 50_000)
-    val = 39601
-    x = Tensor.randn(1, 22, 50_000)[:, :, :v.bind(val)]
-    weight = Tensor.randn(256, 22, 12)
+    # tests symbolic ceildiv in conv2d output shape calculation
+    # val=79 triggers the edge case where old ceildiv simplifies incorrectly: old gives floor=12, correct ceildiv=13
+    v = Variable('v', 11, 100)
+    val = 79
+    x_full = Tensor.randn(1, 8, 100)
+    weight = Tensor.randn(16, 8, 12)
 
-    result = x.conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))
+    # symbolic version
+    result = x_full[:, :, :v.bind(val)].conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))
     var_val = {v.expr: val}
     shape = tuple(sym_infer(s, var_val) for s in result.shape)
-    with self.assertRaises(AssertionError):
-      self.assertEqual(shape, (1, 256, 6600))  # TODO: fails if ceildiv is incorrect
-    # TODO: test output is correct
+    self.assertEqual(shape, (1, 16, 13))
+
+    # concrete version for comparison
+    expected = x_full[:, :, :val].conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))
+    np.testing.assert_allclose(result[:, :, :13].numpy(), expected.numpy(), atol=1e-5, rtol=1e-5)
 
 if __name__ == '__main__':
   unittest.main()

test/test_tensor_uop.py

@@ -1,117 +0,0 @@
-#!/usr/bin/env python
-import numpy as np
-import unittest
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.engine.realize import run_schedule
-from tinygrad.uop.ops import UOp
-from tinygrad.helpers import SPLIT_REDUCEOP
-
-class TestTensorUOp(unittest.TestCase):
-  def test_fromcpu_shape_tracker(self):
-    def helper(a: np.ndarray):
-      print(a.shape, a.strides, a.flags.c_contiguous)
-      b = Tensor(a).uop
-      assert b.shape == a.shape
-      np.testing.assert_equal(a, Tensor(b).numpy())
-
-    for ndims in range(1, 4):
-      a = np.random.randn(*(4,)*ndims).astype(np.float32)
-      for stride in [-2, 1, 2]:
-        for start in [0, 1]:
-          helper(a[(slice(start, None, stride),)*ndims])
-
-  def test_shuffle_pad_ops_cmpeq(self):
-    y = Tensor([1]).cat(Tensor([1]) == 0).numpy()
-    z = Tensor([1, 0]).numpy()
-    np.testing.assert_allclose(y, z)
-
-  def test_shuffle_pad_ops_div(self):
-    y = Tensor([1]).cat(Tensor([1]).div(Tensor([2.0]))).numpy()
-    z = Tensor([1, 0.5]).numpy()
-    np.testing.assert_allclose(y, z)
-
-  def test_shuffle_pad_ops_log(self):
-    y = Tensor([1]).cat(Tensor([1]).log()).numpy()
-    z = Tensor([1, 0]).numpy()
-    np.testing.assert_allclose(y, z)
-
-  def test_shuffle_pad_ops_exp(self):
-    y = Tensor([1]).cat(Tensor([1]).exp()).numpy()
-    z = Tensor([1, np.e]).numpy()
-    np.testing.assert_allclose(y, z)
-
-  def test_device_0_is_the_same_device(self):
-    a = Tensor([1, 2, 3], f"{Device.DEFAULT}")
-    b = Tensor([1, 2, 3], f"{Device.DEFAULT}:0")
-    assert a.device == b.device
-
-  def test_shrink_const_into_zero(self):
-    # regression test to make sure the shapetracker is preserved
-    a = Tensor.zeros(4,4,4).shrink((None, (0,0), None))
-    b = Tensor.zeros(4,1,4)
-    c = a.cat(b, dim=1)
-    np.testing.assert_allclose(c.numpy(), np.concatenate((a.numpy(), b.numpy()), axis=1))
-
-  def test_shrink_const_then_cast(self):
-    # regression test to make sure the shapetracker is preserved
-    a = Tensor.zeros(4,4,4).shrink((None, (0,0), None)).cast(dtypes.int32)
-    b = Tensor.zeros(4,1,4)
-    c = a.cat(b, dim=1)
-    np.testing.assert_allclose(c.numpy(), np.concatenate((a.numpy(), b.numpy()), axis=1))
-
-  def test_const_dtype(self):
-    lb: UOp = Tensor([1], dtype=dtypes.int).uop
-    assert lb.const_like(1).base.arg == 1
-    assert type(lb.const_like(1).base.arg) is int
-
-    lb: UOp = Tensor([1], dtype=dtypes.float).uop
-    assert lb.const_like(1).base.arg == 1.0
-    assert type(lb.const_like(1).base.arg) is float
-
-  def test_contiguous_alu(self):
-    a = Tensor.randn(2, 2).realize()
-    b = Tensor.randn(2, 2).realize()
-    add = (a+b).contiguous()
-    out = add+2
-    sched = out.schedule()
-    self.assertEqual(len(sched), 2)
-    run_schedule(sched)
-    np.testing.assert_allclose(out.numpy(), a.numpy()+b.numpy()+2)
-
-  # NOTE: contiguous on a buffer collapses
-  @unittest.skip("contiguous on a buffer no longer collapses")
-  def test_contiguous_empty(self):
-    empty = Tensor.empty(1).contiguous()
-    sched = empty.schedule()
-    self.assertEqual(len(sched), 0)
-
-  def test_contiguous_folded_alu(self):
-    a = Tensor.empty(8, 8)
-    # NOTE: the buffer for mul_0 late folds to just a CONST
-    mul_0 = a*0
-    out = mul_0.shrink(((4, 8), (0, 8))).contiguous()
-    out.realize()
-    self.assertEqual(out.tolist(), Tensor.zeros(4, 8).tolist())
-
-@unittest.skipUnless(SPLIT_REDUCEOP, "only for SPLIT_REDUCEOP")
-class TestReduceOp(unittest.TestCase):
-  def test_no_split_reduce_kernel(self):
-    a = Tensor.rand(4, 4).realize()
-    a = a.sum()
-    sched = a.schedule()
-    assert len(sched) == 1
-
-  def test_split_reduce_kernel_dim0(self):
-    a = Tensor.rand(256, 255).realize()
-    a = a.sum()
-    sched = a.schedule()
-    assert len(sched) == 2
-
-  def test_split_reduce_kernel_dim1(self):
-    a = Tensor.rand(255, 256).realize()
-    a = a.sum()
-    sched = a.schedule()
-    assert len(sched) == 2
-
-if __name__ == "__main__":
-  unittest.main()

test/test_tensor_variable.py

@@ -73,8 +73,6 @@ class TestTensorVariable(unittest.TestCase):
     ret = Tensor.arange(vv.bind(4), 7)
     self.assertListEqual(ret[:3].tolist(), [4,5,6])
 
-  # TODO: add vmin/vmax pattern for symbolic denominator
-  @unittest.expectedFailure
   def test_symbolic_arange_sym_step(self):
     vv = Variable("step", 1, 3)
     ret = Tensor.arange(0, 10, vv.bind(2))
@@ -86,6 +84,18 @@ class TestTensorVariable(unittest.TestCase):
     ret = Tensor.arange(begin.bind(4), end.bind(7))
     self.assertListEqual(ret[:3].tolist(), [4,5,6])
 
+  def test_symbolic_arange_three_vars(self):
+    begin = Variable("b", 0, 5)
+    end = Variable("e", 10, 20)
+    step = Variable("s", 1, 3)
+    ret = Tensor.arange(begin.bind(2), end.bind(14), step.bind(3))
+    self.assertListEqual(ret[:4].tolist(), [2,5,8,11])
+
+  def test_symbolic_full(self):
+    vv = Variable("x", 1, 10).bind(5)
+    t = Tensor.full((3,), vv)
+    self.assertListEqual(t.tolist(), [5,5,5])
+
   def test_variable_empty(self):
     v = Variable("i", 1, 10)
     # TODO: Tensor creation from unbound variable should assert

test/test_transcendental.py

@@ -101,6 +101,44 @@ class TestFromFuzzer(unittest.TestCase):
     _test_value(0)
     _test_value(0.0000009)
 
+class TestFloat16Log2(unittest.TestCase):
+  """Tests for native float16 log2 implementation (no float32 cast)"""
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16, Device.DEFAULT), f"no float16 on {Device.DEFAULT}")
+  def test_float16_log2_basic(self):
+    # basic values
+    test_values = [1.0, 2.0, 4.0, 0.5, 0.25, 10.0, 100.0, 1000.0]
+    with Context(TRANSCENDENTAL=2):
+      for val in test_values:
+        result = Tensor([val], dtype=dtypes.float16).log2().numpy()[0]
+        expected = np.log2(np.float16(val))
+        np.testing.assert_allclose(result, expected, rtol=1e-3, err_msg=f"log2({val})")
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16, Device.DEFAULT), f"no float16 on {Device.DEFAULT}")
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and CI, "Nan handling differs on Vulkan")
+  def test_float16_log2_special(self):
+    # special values: inf, -inf, nan, 0, negative
+    with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'):
+      # log2(inf) = inf
+      assert np.isinf(Tensor([np.inf], dtype=dtypes.float16).log2().numpy()[0])
+      # log2(0) = -inf
+      assert Tensor([0.0], dtype=dtypes.float16).log2().numpy()[0] == -np.inf
+      # log2(negative) = nan
+      assert np.isnan(Tensor([-1.0], dtype=dtypes.float16).log2().numpy()[0])
+      # log2(nan) = nan
+      assert np.isnan(Tensor([np.nan], dtype=dtypes.float16).log2().numpy()[0])
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16, Device.DEFAULT), f"no float16 on {Device.DEFAULT}")
+  def test_float16_log2_denormal(self):
+    # test values near and below float16 min normal (6.1e-5)
+    # these exercise the denormal handling path with 2^10 scaling
+    test_values = [1e-4, 6e-5, 1e-5]
+    with Context(TRANSCENDENTAL=2):
+      for val in test_values:
+        result = Tensor([val], dtype=dtypes.float16).log2().numpy()[0]
+        expected = np.log2(np.float16(val))
+        # denormals have lower precision due to float16 limitations
+        np.testing.assert_allclose(result, expected, rtol=5e-2, err_msg=f"log2({val})")
+
 class TestTranscendentalSchedule(unittest.TestCase):
   @unittest.skipUnless(is_dtype_supported(dtypes.ulong), "Needs ulong")
   def test_transcendental_sin_fusion(self):

test/test_uop_graph.py

@@ -478,143 +478,6 @@ class TestUOpGraph(unittest.TestCase):
     for u in uops:
       self.assertNotEqual(u.dtype, dtypes.long)
 
-  def test_in_out_of_bounds_access(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 0), ptr=True),))
-      to_uops_list([ld0])
-      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 15), ptr=True),))
-      to_uops_list([ld1])
-      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 7), ptr=True),))
-      to_uops_list([ld1])
-
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 42), ptr=True),))
-      with self.assertRaises(RuntimeError): to_uops_list([ld0])
-
-  def test_in_out_of_bounds_access_symbolic(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 1, 10), ptr=True),))
-      to_uops_list([ld0])
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 15), ptr=True),))
-      to_uops_list([ld0])
-
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 20), ptr=True),))
-      with self.assertRaises(RuntimeError): to_uops_list([ld0])
-
-  def test_in_out_of_bounds_access_gated_store(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), src=(), arg=0)
-      v = Variable("v", 0, 20)
-      st0 = UOp(Ops.STORE, dtypes.void, src=(glbl0.index(v.valid(v<16)), UOp.const(dtypes.int, 0)))
-      to_uops_list([st0])
-
-      st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v.valid(v<20)), v))
-      with self.assertRaises(RuntimeError): to_uops_list([st1])
-
-  @unittest.skip("if not allowed in graph")
-  def test_in_bounds_access_gated_local(self):
-    with Context(IGNORE_OOB=0):
-      # Define buffers
-      gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.uint.ptr(400), (), 0)
-      sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.uint.ptr(8, addrspace=AddrSpace.LOCAL), (), "temp0")
-
-      # Define indices, valids and barrier
-      gidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 416),), "gidx0")
-      lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 10),), "lidx0")
-
-      gate = (gidx<400) & (lidx<8)
-
-      local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx, lidx<8), UOp.const(dtypes.uint, 1)))
-
-      barrier = UOp(Ops.BARRIER, dtypes.void, (local_store,))
-      if_barrier = UOp(Ops.IF, dtypes.void, (gate, barrier))
-
-      # Load from local memory (after the IF/barrier)
-      local_load = UOp(Ops.LOAD, dtypes.uint, (sbuf.index(lidx, ptr=True), if_barrier))
-
-      # Store to global memory
-      global_store = UOp(Ops.STORE, dtypes.void, (gbuf.index(gidx), local_load))
-      to_uops_list([global_store])
-
-  def test_load_with_float_in_index(self):
-    with Context(IGNORE_OOB=0):
-      ridx = UOp.range(20, 0)
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      i = (ridx.cast(dtypes.float)*0.68).trunc().cast(dtypes.int)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i.valid((0<=i)&(i<16)), ptr=True),))
-      to_uops_list([ld0])
-      glblfloat = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(20), (), 0)
-      ldfloat = UOp(Ops.LOAD, dtypes.float, (glblfloat.index(ridx),))
-      i = (ldfloat+3.14).cast(dtypes.int)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, ((0<=i)&(i<16)), ptr=True),))
-
-  def test_load_cast_to_bool(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), (), 0)
-      ridx = UOp.range(20, 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(ridx.valid(ridx.cast(dtypes.bool).logical_not()), ptr=True),))
-      to_uops_list([ld0])
-
-  @unittest.skip("Bool load is not supported yet")
-  def test_load_mask(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      mask = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
-      ridx = UOp.range(20, 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask), ptr=True)))
-      to_uops_list([ld0])
-
-  def test_out_of_bounds_off_by_one_access(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 16), ptr=True),))
-      with self.assertRaises(RuntimeError): to_uops_list([ld0])
-
-  def test_in_out_bounds_access_with_mask(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      gidx0 = UOp.range(42, 0, AxisType.GLOBAL)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0.valid((5<gidx0)&(gidx0<16)), ptr=True),))
-      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0.valid(gidx0<16), ptr=True),))
-      to_uops_list([ld0, ld1])
-
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0.valid(gidx0<17), ptr=True),))
-      with self.assertRaises(RuntimeError): to_uops_list([ld0])
-
-  def test_in_out_of_bounds_access_symbolic_mask(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      i = Variable("i", 1, 80)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i.valid(i<10), ptr=True),))
-      to_uops_list([ld0])
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i.valid(i<15), ptr=True),))
-      to_uops_list([ld0])
-
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i.valid(i<20), ptr=True),))
-      with self.assertRaises(RuntimeError): to_uops_list([ld0])
-
-  def test_in_out_of_bounds_access_index_load(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 0)
-      gidx0 = UOp.range(42, 0, AxisType.GLOBAL)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0.valid(gidx0<8), ptr=True),)).cast(dtypes.index)
-      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index((ld0*2).valid((ld0>=0)&(ld0<32)), ptr=True),))
-      to_uops_list([ld1])
-
-      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index((ld0*2).valid((ld0>=0)&(ld0<64)), ptr=True),))
-      with self.assertRaises(RuntimeError): to_uops_list([ld1])
-
-  def test_bounds_with_loaded_bool(self):
-    with Context(IGNORE_OOB=0):
-      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
-      glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(8), (), 0)
-      gidx0 = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 16),), "gidx0")
-      ld0 = glbl0.index(gidx0, ptr=True).load()
-      ld1 = glbl1.index(gidx0.valid(ld0), ptr=True).load()
-      with self.assertRaises(RuntimeError): to_uops_list([ld1])
-
   def test_fold_gated_load(self):
     glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
     glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 1)