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,7 +366,7 @@ 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 AMD_LLVM=0
@@ -437,7 +377,7 @@ jobs:
     - name: Run Tensor Core GEMM (AMD)
       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 | tee matmul_amd.txt
+        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
@@ -452,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
 
@@ -543,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
 
@@ -606,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
 
@@ -708,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
@@ -718,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
 
@@ -778,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,13 +231,14 @@ 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
     - name: Run TYPED=1
       run: TYPED=1 python -c "import tinygrad"
@@ -310,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
@@ -473,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
@@ -680,9 +681,9 @@ jobs:
       - 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: PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
+        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: PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
+        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
@@ -698,8 +699,7 @@ jobs:
           pydeps: "pdfplumber"
       - name: Verify AMD autogen is up to date
         run: |
-          python -m extra.assembly.amd.dsl --arch all
-          python -m extra.assembly.amd.pcode --arch all
+          python -m extra.assembly.amd.pdf --arch all
           git diff --exit-code extra/assembly/amd/autogen/
 
   testnvidia:

.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

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/asm.py

@@ -1,98 +1,84 @@
 # RDNA3 assembler and disassembler
 from __future__ import annotations
 import re
-from extra.assembly.amd.dsl import Inst, RawImm, Reg, SrcMod, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory, FLOAT_ENC, SRC_FIELDS, unwrap
+from extra.assembly.amd.dsl import Inst, RawImm, Reg, SrcMod, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory, SRC_FIELDS, unwrap
 from extra.assembly.amd.dsl import VCC_LO, VCC_HI, VCC, EXEC_LO, EXEC_HI, EXEC, SCC, M0, NULL, OFF
-from extra.assembly.amd.autogen.rdna3 import VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, VOPD, VINTERP, SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, DS, FLAT, MUBUF, MTBUF, MIMG, EXP
-from extra.assembly.amd.autogen.rdna3 import VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, VOPDOp, VINTERPOp
-from extra.assembly.amd.autogen.rdna3 import SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, DSOp, FLATOp, MUBUFOp, MTBUFOp, MIMGOp
+from extra.assembly.amd.dsl import SPECIAL_GPRS, SPECIAL_PAIRS, FLOAT_DEC, FLOAT_ENC, decode_src
+from extra.assembly.amd.autogen.rdna3 import ins
+from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, VOPD, VINTERP, SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, DS, FLAT, MUBUF, MTBUF, MIMG, EXP, LDSDIR,
+  VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, VOPDOp, VINTERPOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, DSOp, FLATOp, MUBUFOp, MTBUFOp, MIMGOp)
 
 # VOP3SD opcodes that share VOP3 encoding
 VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
 
+def _matches_encoding(word: int, cls: type[Inst]) -> bool:
+  """Check if word matches the encoding pattern of an instruction class."""
+  if cls._encoding is None: return False
+  bf, val = cls._encoding
+  return ((word >> bf.lo) & bf.mask()) == val
+
+# Order matters: more specific encodings first, VOP2 last (it's a catch-all for bit31=0)
+_FORMATS_64 = [VOPD, VOP3P, VINTERP, VOP3, DS, FLAT, MUBUF, MTBUF, MIMG, SMEM, EXP]
+_FORMATS_32 = [SOP1, SOPC, SOPP, SOPK, VOPC, VOP1, SOP2, VOP2]  # SOP2/VOP2 are catch-alls
+
 def detect_format(data: bytes) -> type[Inst]:
   """Detect instruction format from machine code bytes."""
   assert len(data) >= 4, f"need at least 4 bytes, got {len(data)}"
   word = int.from_bytes(data[:4], 'little')
-  hi2 = (word >> 30) & 0x3
-  if hi2 == 0b11:
-    enc = (word >> 26) & 0xf
-    if enc == 0b0010: return VOPD
-    if enc == 0b0011: return VOP3P
-    if enc == 0b0100: return VINTERP
-    if enc == 0b0101: return VOP3SD if ((word >> 16) & 0x3ff) in VOP3SD_OPS else VOP3
-    if enc == 0b0110: return DS
-    if enc == 0b0111: return FLAT
-    if enc == 0b1000: return MUBUF
-    if enc == 0b1010: return MTBUF
-    if enc == 0b1100 or enc == 0b1111: return MIMG
-    if enc == 0b1101: return SMEM
-    if enc == 0b1110: return EXP
-    raise ValueError(f"unknown 64-bit format enc={enc:#06b} word={word:#010x}")
-  if hi2 == 0b10:
-    enc = (word >> 23) & 0x7f
-    if enc == 0b1111101: return SOP1
-    if enc == 0b1111110: return SOPC
-    if enc == 0b1111111: return SOPP
-    return SOPK if ((word >> 28) & 0xf) == 0b1011 else SOP2
-  # hi2 == 0b00 or 0b01: VOP1/VOP2/VOPC (bit 31 = 0)
-  assert (word >> 31) == 0, f"expected bit 31 = 0 for VOP, got word={word:#010x}"
-  enc = (word >> 25) & 0x7f
-  if enc == 0b0111110: return VOPC
-  if enc == 0b0111111: return VOP1
-  if enc <= 0b0111101: return VOP2
-  raise ValueError(f"unknown VOP format enc={enc:#09b} word={word:#010x}")
+  # Check 64-bit formats first (bits[31:30] == 0b11)
+  if (word >> 30) == 0b11:
+    for cls in _FORMATS_64:
+      if _matches_encoding(word, cls):
+        return VOP3SD if cls is VOP3 and ((word >> 16) & 0x3ff) in VOP3SD_OPS else cls
+    raise ValueError(f"unknown 64-bit format word={word:#010x}")
+  # 32-bit formats
+  for cls in _FORMATS_32:
+    if _matches_encoding(word, cls): return cls
+  raise ValueError(f"unknown 32-bit format word={word:#010x}")
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # CONSTANTS
 # ═══════════════════════════════════════════════════════════════════════════════
 
-SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
-SPECIAL_DEC = {**SPECIAL_GPRS, **{v: str(k) for k, v in FLOAT_ENC.items()}}
-SPECIAL_PAIRS = {106: "vcc", 126: "exec"}
+# GFX11 HWREG IDs
 HWREG = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
          6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
          19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI', 22: 'HW_REG_XNACK_MASK',
          23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
+# GFX12 HWREG IDs - use names that LLVM recognizes
+HWREG_GFX12 = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 5: 'HW_REG_GPR_ALLOC', 6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS',
+               18: 'HW_REG_EXCP_FLAG_USER', 19: 'HW_REG_TRAP_CTRL', 20: 'HW_REG_SCRATCH_BASE_LO', 21: 'HW_REG_SCRATCH_BASE_HI',
+               23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 29: 'HW_REG_SHADER_CYCLES_LO', 30: 'HW_REG_SHADER_CYCLES_HI'}
 HWREG_IDS = {v.lower(): k for k, v in HWREG.items()}
 MSG = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
        131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA'}
-VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # HELPERS
 # ═══════════════════════════════════════════════════════════════════════════════
 
-def decode_src(val: int) -> str:
-  if val <= 105: return f"s{val}"
-  if val in SPECIAL_DEC: return SPECIAL_DEC[val]
-  if 108 <= val <= 123: return f"ttmp{val - 108}"
-  if 128 <= val <= 192: return str(val - 128)
-  if 193 <= val <= 208: return str(-(val - 192))
-  if 256 <= val <= 511: return f"v{val - 256}"
-  return "lit" if val == 255 else f"?{val}"
-
 def _reg(p: str, b: int, n: int = 1) -> str: return f"{p}{b}" if n == 1 else f"{p}[{b}:{b+n-1}]"
 def _sreg(b: int, n: int = 1) -> str: return _reg("s", b, n)
 def _vreg(b: int, n: int = 1) -> str: return _reg("v", b, n)
-def _hl(v: int, hi_thresh: int = 128) -> str: return 'h' if v >= hi_thresh else 'l'
+def _ttmp(b: int, n: int = 1) -> str: return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
+def _sreg_or_ttmp(b: int, n: int = 1) -> str: return _ttmp(b, n) or _sreg(b, n)
 
 def _fmt_sdst(v: int, n: int = 1) -> str:
   if v == 124: return "null"
-  if 108 <= v <= 123: return _reg("ttmp", v - 108, n)
+  if t := _ttmp(v, n): return t
   if n > 1: return SPECIAL_PAIRS.get(v) or _sreg(v, n)
-  return {126: "exec_lo", 127: "exec_hi", 106: "vcc_lo", 107: "vcc_hi", 125: "m0"}.get(v, f"s{v}")
+  return SPECIAL_GPRS.get(v, f"s{v}")
 
 def _fmt_src(v: int, n: int = 1) -> str:
   if n == 1: return decode_src(v)
   if v >= 256: return _vreg(v - 256, n)
   if v <= 105: return _sreg(v, n)
   if n == 2 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
-  if 108 <= v <= 123: return _reg("ttmp", v - 108, n)
+  if t := _ttmp(v, n): return t
   return decode_src(v)
 
 def _fmt_v16(v: int, base: int = 256, hi_thresh: int = 384) -> str:
-  return f"v{(v - base) & 0x7f}.{_hl(v, hi_thresh)}"
+  return f"v{(v - base) & 0x7f}.{'h' if v >= hi_thresh else 'l'}"
 
 def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
   return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
@@ -101,6 +87,7 @@ def _has(op: str, *subs) -> bool: return any(s in op for s in subs)
 def _is16(op: str) -> bool: return _has(op, 'f16', 'i16', 'u16', 'b16') and not _has(op, '_f32', '_i32')
 def _is64(op: str) -> bool: return _has(op, 'f64', 'i64', 'u64', 'b64')
 def _omod(v: int) -> str: return {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(v, "")
+def _src16(inst, v: int) -> str: return _fmt_v16(v) if v >= 256 else inst.lit(v)  # format 16-bit src: vgpr.h/l or literal
 def _mods(*pairs) -> str: return " ".join(m for c, m in pairs if c)
 def _fmt_bits(label: str, val: int, count: int) -> str: return f"{label}:[{','.join(str((val >> i) & 1) for i in range(count))}]"
 
@@ -123,41 +110,59 @@ def _opsel_str(opsel: int, n: int, need: bool, is16_d: bool) -> str:
 # DISASSEMBLER
 # ═══════════════════════════════════════════════════════════════════════════════
 
+_VOP1_F64 = {VOP1Op.V_CEIL_F64, VOP1Op.V_FLOOR_F64, VOP1Op.V_FRACT_F64, VOP1Op.V_FREXP_MANT_F64, VOP1Op.V_RCP_F64, VOP1Op.V_RNDNE_F64, VOP1Op.V_RSQ_F64, VOP1Op.V_SQRT_F64, VOP1Op.V_TRUNC_F64}
+
 def _disasm_vop1(inst: VOP1) -> str:
-  op = VOP1Op(inst.op)
-  if op in (VOP1Op.V_NOP, VOP1Op.V_PIPEFLUSH): return op.name.lower()
-  F64_OPS = {VOP1Op.V_CEIL_F64, VOP1Op.V_FLOOR_F64, VOP1Op.V_FRACT_F64, VOP1Op.V_FREXP_MANT_F64, VOP1Op.V_RCP_F64, VOP1Op.V_RNDNE_F64, VOP1Op.V_RSQ_F64, VOP1Op.V_SQRT_F64, VOP1Op.V_TRUNC_F64}
-  is_f64_d = op in F64_OPS or op in (VOP1Op.V_CVT_F64_F32, VOP1Op.V_CVT_F64_I32, VOP1Op.V_CVT_F64_U32)
-  is_f64_s = op in F64_OPS or op in (VOP1Op.V_CVT_F32_F64, VOP1Op.V_CVT_I32_F64, VOP1Op.V_CVT_U32_F64, VOP1Op.V_FREXP_EXP_I32_F64)
-  name = op.name.lower()
+  # Use architecture-specific op enum
+  if 'rdna4' in inst.__class__.__module__:
+    from extra.assembly.amd.autogen.rdna4.enum import VOP1Op as OpEnum
+  else:
+    OpEnum = VOP1Op
+  op, name = OpEnum(inst.op), OpEnum(inst.op).name.lower()
+  if name in ('v_nop', 'v_pipeflush'): return name
+  if name == 'v_readfirstlane_b32': return f"v_readfirstlane_b32 {decode_src(inst.vdst)}, v{inst.src0 - 256 if inst.src0 >= 256 else inst.src0}"
   parts = name.split('_')
+  is_f64_d = 'f64' in name and any(x in name for x in ['ceil', 'floor', 'fract', 'frexp_mant', 'rcp', 'rndne', 'rsq', 'sqrt', 'trunc', 'cvt_f64_f32', 'cvt_f64_i32', 'cvt_f64_u32'])
+  is_f64_s = 'f64' in name and any(x in name for x in ['ceil', 'floor', 'fract', 'frexp_mant', 'rcp', 'rndne', 'rsq', 'sqrt', 'trunc', 'cvt_f32_f64', 'cvt_i32_f64', 'cvt_u32_f64', 'frexp_exp_i32_f64'])
+  # v_cvt_pk_f32_bf8/fp8 output 2 VGPRs (packed f32x2) and take packed 8-bit (16-bit VGPR with .l/.h) source
+  is_pk_f32 = 'cvt_pk_f32_bf8' in name or 'cvt_pk_f32_fp8' in name
   is_16d = any(p in ('f16','i16','u16','b16') for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in ('f16','i16','u16','b16') and 'cvt' not in name)
-  is_16s = parts[-1] in ('f16','i16','u16','b16') and 'sat_pk' not in name
-  if op == VOP1Op.V_READFIRSTLANE_B32: return f"v_readfirstlane_b32 {decode_src(inst.vdst)}, v{inst.src0 - 256 if inst.src0 >= 256 else inst.src0}"
-  dst = _vreg(inst.vdst, 2) if is_f64_d else _fmt_v16(inst.vdst, 0, 128) if is_16d else f"v{inst.vdst}"
-  src = _fmt_src(inst.src0, 2) if is_f64_s else _fmt_v16(inst.src0) if is_16s and inst.src0 >= 256 else inst.lit(inst.src0)
+  # Only packed bf8/fp8 (cvt_pk_*) use 16-bit VGPR encoding; non-packed versions use regular VGPRs
+  is_16s = (parts[-1] in ('f16','i16','u16','b16') and 'sat_pk' not in name) or (parts[-1] in ('bf8', 'fp8') and 'pk' in name)
+  dst = _vreg(inst.vdst, 2) if is_f64_d or is_pk_f32 else _fmt_v16(inst.vdst, 0, 128) if is_16d else f"v{inst.vdst}"
+  src = _fmt_src(inst.src0, 2) if is_f64_s else _src16(inst, inst.src0) if is_16s else inst.lit(inst.src0)
   return f"{name}_e32 {dst}, {src}"
 
 def _disasm_vop2(inst: VOP2) -> str:
-  op = VOP2Op(inst.op)
-  name = op.name.lower()
-  suf = "" if op == VOP2Op.V_DOT2ACC_F32_F16 else "_e32"
-  is16 = _is16(name) and 'pk_' not in name
+  # Use architecture-specific op enum
+  if 'rdna4' in inst.__class__.__module__:
+    from extra.assembly.amd.autogen.rdna4.enum import VOP2Op as OpEnum
+  else:
+    OpEnum = VOP2Op
+  op, name = OpEnum(inst.op), OpEnum(inst.op).name.lower()
+  suf, is16 = "" if name == 'v_dot2acc_f32_f16' else "_e32", _is16(name) and 'pk_' not in name
+  is64 = _is64(name)
+  # For shift ops with b64, src0 is 32-bit (shift amount), dst/vsrc1 are 64-bit
+  is_shift64 = 'lshlrev_b64' in name
   # fmaak: dst = src0 * vsrc1 + K, fmamk: dst = src0 * K + vsrc1
-  if op in (VOP2Op.V_FMAAK_F32, VOP2Op.V_FMAAK_F16): return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}, 0x{inst._literal:x}"
-  if op in (VOP2Op.V_FMAMK_F32, VOP2Op.V_FMAMK_F16): return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, 0x{inst._literal:x}, v{inst.vsrc1}"
-  if is16: return f"{name}{suf} {_fmt_v16(inst.vdst, 0, 128)}, {_fmt_v16(inst.src0) if inst.src0 >= 256 else inst.lit(inst.src0)}, {_fmt_v16(inst.vsrc1, 0, 128)}"
-  return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}" + (", vcc_lo" if op == VOP2Op.V_CNDMASK_B32 else "")
-
-VOPC_CLASS = {VOPCOp.V_CMP_CLASS_F16, VOPCOp.V_CMP_CLASS_F32, VOPCOp.V_CMP_CLASS_F64,
-              VOPCOp.V_CMPX_CLASS_F16, VOPCOp.V_CMPX_CLASS_F32, VOPCOp.V_CMPX_CLASS_F64}
+  if 'fmaak' in name: return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}, 0x{inst._literal:x}"
+  if 'fmamk' in name: return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, 0x{inst._literal:x}, v{inst.vsrc1}"
+  if is16: return f"{name}{suf} {_fmt_v16(inst.vdst, 0, 128)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1, 0, 128)}"
+  if is_shift64: return f"{name}{suf} {_vreg(inst.vdst, 2)}, {inst.lit(inst.src0)}, {_vreg(inst.vsrc1, 2)}"
+  if is64: return f"{name}{suf} {_vreg(inst.vdst, 2)}, {_fmt_src(inst.src0, 2)}, {_vreg(inst.vsrc1, 2)}"
+  return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}" + (", vcc_lo" if name == 'v_cndmask_b32' else "")
 
 def _disasm_vopc(inst: VOPC) -> str:
-  op = VOPCOp(inst.op)
-  name = op.name.lower()
+  # Use architecture-specific op enum
+  if 'rdna4' in inst.__class__.__module__:
+    from extra.assembly.amd.autogen.rdna4.enum import VOPCOp as OpEnum
+  else:
+    OpEnum = VOPCOp
+  op, name = OpEnum(inst.op), OpEnum(inst.op).name.lower()
   is64, is16 = _is64(name), _is16(name)
-  s0 = _fmt_src(inst.src0, 2) if is64 else _fmt_v16(inst.src0) if is16 and inst.src0 >= 256 else inst.lit(inst.src0)
-  s1 = _vreg(inst.vsrc1, 2) if is64 and op not in VOPC_CLASS else _fmt_v16(inst.vsrc1, 0, 128) if is16 else f"v{inst.vsrc1}"
+  is_class = 'class' in name
+  s0 = _fmt_src(inst.src0, 2) if is64 else _src16(inst, inst.src0) if is16 else inst.lit(inst.src0)
+  s1 = _vreg(inst.vsrc1, 2) if is64 and not is_class else _fmt_v16(inst.vsrc1, 0, 128) if is16 else f"v{inst.vsrc1}"
   return f"{name}_e32 {s0}, {s1}" if op.value >= 128 else f"{name}_e32 vcc_lo, {s0}, {s1}"
 
 NO_ARG_SOPP = {SOPPOp.S_ENDPGM, SOPPOp.S_BARRIER, SOPPOp.S_WAKEUP, SOPPOp.S_ICACHE_INV,
@@ -179,15 +184,41 @@ def _disasm_sopp(inst: SOPP) -> str:
   return f"{name} {inst.simm16}" if name.startswith(('s_cbranch', 's_branch')) else f"{name} 0x{inst.simm16:x}"
 
 def _disasm_smem(inst: SMEM) -> str:
-  op = SMEMOp(inst.op)
-  name = op.name.lower()
-  if op in (SMEMOp.S_GL1_INV, SMEMOp.S_DCACHE_INV): return name
-  off_s = f"{decode_src(inst.soffset)} offset:0x{inst.offset:x}" if inst.offset and inst.soffset != 124 else f"0x{inst.offset:x}" if inst.offset else decode_src(inst.soffset)
-  sbase_idx, sbase_count = inst.sbase * 2, 4 if (8 <= inst.op <= 12 or name == 's_atc_probe_buffer') else 2
-  sbase_str = _fmt_src(sbase_idx, sbase_count) if sbase_count == 2 else _sreg(sbase_idx, sbase_count) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_count)
-  if name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{name} {inst.sdata}, {sbase_str}, {off_s}"
-  width = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(inst.op, 1)
-  return f"{name} {_fmt_sdst(inst.sdata, width)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (inst.dlc, " dlc"))
+  is_rdna4 = 'rdna4' in inst.__class__.__module__
+  if is_rdna4:
+    from extra.assembly.amd.autogen.rdna4.enum import SMEMOp as SMEMOp4
+    op = SMEMOp4(inst.op)
+    name = op.name.lower()
+    if op == SMEMOp4.S_DCACHE_INV: return name
+    # RDNA4: s_buffer_* uses 4-SGPR descriptor, s_load/s_prefetch uses 2-SGPR
+    is_buffer = 'buffer' in name
+    sbase_idx = inst.sbase * 2
+    sbase_count = 4 if is_buffer else 2
+    sbase_str = _fmt_src(sbase_idx, sbase_count) if sbase_count == 2 else _sreg(sbase_idx, sbase_count) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_count)
+    # Format offset - ioffset is signed 24-bit, show as hex
+    ioff = inst.ioffset if inst.ioffset < 0x800000 else inst.ioffset - 0x1000000  # sign extend
+    off_hex = f"0x{ioff & 0xffffff:x}" if ioff >= 0 else f"-0x{(-ioff) & 0xffffff:x}"
+    off_s = f"{decode_src(inst.soffset)} offset:{off_hex}" if inst.soffset != 124 else off_hex
+    # Data width from opcode
+    width_map = {0:1, 1:2, 2:4, 3:8, 4:16, 5:3, 8:1, 9:1, 10:1, 11:1, 16:1, 17:2, 18:4, 19:8, 20:16, 21:3, 24:1, 25:1, 26:1, 27:1}
+    width = width_map.get(inst.op, 1)
+    if 'prefetch' in name:
+      # Prefetch has different format: s_prefetch_* sbase, offset, soffset, length
+      # But we need to handle various prefetch types differently
+      if name == 's_prefetch_inst_pc_rel' or name == 's_prefetch_data_pc_rel':
+        return f"{name} {off_hex}, {decode_src(inst.soffset)}, {inst.sdata}"
+      return f"{name} {sbase_str}, {off_hex}, {decode_src(inst.soffset)}, {inst.sdata}"
+    return f"{name} {_fmt_sdst(inst.sdata, width)}, {sbase_str}, {off_s}"
+  else:
+    op = SMEMOp(inst.op)
+    name = op.name.lower()
+    if op in (SMEMOp.S_GL1_INV, SMEMOp.S_DCACHE_INV): return name
+    off_s = f"{decode_src(inst.soffset)} offset:0x{inst.offset:x}" if inst.offset and inst.soffset != 124 else f"0x{inst.offset:x}" if inst.offset else decode_src(inst.soffset)
+    sbase_idx, sbase_count = inst.sbase * 2, 4 if (8 <= inst.op <= 12 or name == 's_atc_probe_buffer') else 2
+    sbase_str = _fmt_src(sbase_idx, sbase_count) if sbase_count == 2 else _sreg(sbase_idx, sbase_count) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_count)
+    if name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{name} {inst.sdata}, {sbase_str}, {off_s}"
+    width = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(inst.op, 1)
+    return f"{name} {_fmt_sdst(inst.sdata, width)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (inst.dlc, " dlc"))
 
 def _disasm_flat(inst: FLAT) -> str:
   name = FLATOp(inst.op).name.lower()
@@ -204,7 +235,7 @@ def _disasm_flat(inst: FLAT) -> str:
   elif inst.saddr == 124: saddr_s = ", off"
   elif seg == 'scratch': saddr_s = f", {decode_src(inst.saddr)}"
   elif inst.saddr in SPECIAL_PAIRS: saddr_s = f", {SPECIAL_PAIRS[inst.saddr]}"
-  elif 108 <= inst.saddr <= 123: saddr_s = f", {_reg('ttmp', inst.saddr - 108, 2)}"
+  elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
   else: saddr_s = f", {_sreg(inst.saddr, 2) if inst.saddr < 106 else decode_src(inst.saddr)}"
   # addtid: no addr
   if 'addtid' in name: return f"{instr} v{inst.data if 'store' in name else inst.vdst}{saddr_s}{mods}"
@@ -245,13 +276,18 @@ def _disasm_ds(inst: DS) -> str:
   return f"{name} {dst}, {addr}, {d0}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}{off}{gds}"
 
 def _disasm_vop3(inst: VOP3) -> str:
-  op = VOP3SDOp(inst.op) if inst.op in VOP3SD_OPS else VOP3Op(inst.op)
+  is_rdna4 = 'rdna4' in inst.__class__.__module__
+  if is_rdna4:
+    from extra.assembly.amd.autogen.rdna4.enum import VOP3Op as VOP3Op4, VOP3SDOp as VOP3SDOp4
+    op = VOP3SDOp4(inst.op) if inst.op in VOP3SD_OPS else VOP3Op4(inst.op)
+  else:
+    op = VOP3SDOp(inst.op) if inst.op in VOP3SD_OPS else VOP3Op(inst.op)
   name = op.name.lower()
 
   # VOP3SD (shared encoding)
   if inst.op in VOP3SD_OPS:
     sdst = (inst.clmp << 7) | (inst.opsel << 3) | inst.abs
-    is64, mad64 = 'f64' in name, _has(name, 'mad_i64_i32', 'mad_u64_u32')
+    is64, mad64 = 'f64' in name, _has(name, 'mad_i64_i32', 'mad_u64_u32', 'mad_co_i64_i32', 'mad_co_u64_u32')
     def src(v, neg, ext=False): s = _fmt_src(v, 2) if ext or is64 else inst.lit(v); return f"-{s}" if neg else s
     s0, s1, s2 = src(inst.src0, inst.neg & 1), src(inst.src1, inst.neg & 2), src(inst.src2, inst.neg & 4, mad64)
     dst = _vreg(inst.vdst, 2) if is64 or mad64 else f"v{inst.vdst}"
@@ -263,7 +299,9 @@ def _disasm_vop3(inst: VOP3) -> str:
   is64 = _is64(name)
   is64_src, is64_dst = False, False
   is16_d = is16_s = is16_s2 = False
-  if 'cvt_pk' in name: is16_s = name.endswith('16')
+  # v_cvt_pk_f32_bf8/fp8 outputs a VGPR pair (f32x2) from 16-bit packed input
+  if 'cvt_pk_f32_bf8' in name or 'cvt_pk_f32_fp8' in name: is64_dst = True
+  elif 'cvt_pk' in name: is16_s = name.endswith('16')
   elif m := re.match(r'v_(?:cvt|frexp_exp)_([a-z0-9_]+)_([a-z0-9]+)', name):
     is16_d, is16_s = _has(m.group(1), 'f16','i16','u16','b16'), _has(m.group(2), 'f16','i16','u16','b16')
     is64_src, is64_dst = '64' in m.group(2), '64' in m.group(1)
@@ -297,17 +335,41 @@ def _disasm_vop3(inst: VOP3) -> str:
   nonvgpr_opsel = (inst.src0 < 256 and (inst.opsel & 1)) or (inst.src1 < 256 and (inst.opsel & 2)) or (inst.src2 < 256 and (inst.opsel & 4))
   need_opsel = nonvgpr_opsel or (inst.opsel and not is16_s)
 
+  # RDNA4 v_s_* instructions (pseudo-scalar VOP1-like) have SGPR destination
+  if name.startswith('v_s_') and is_rdna4:
+    return f"{name} {_fmt_sdst(inst.vdst, 1)}, {s0}{cl}{om}"
+
   if inst.op < 256:  # VOPC
     return f"{name}_e64 {s0}, {s1}" if name.startswith('v_cmpx') else f"{name}_e64 {_fmt_sdst(inst.vdst, 1)}, {s0}, {s1}"
   if inst.op < 384:  # VOP2
     os = _opsel_str(inst.opsel, 3, need_opsel, is16_d) if 'cndmask' in name else _opsel_str(inst.opsel, 2, need_opsel, is16_d)
     return f"{name}_e64 {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if 'cndmask' in name else f"{name}_e64 {dst}, {s0}, {s1}{os}{cl}{om}"
   if inst.op < 512:  # VOP1
-    return f"{name}_e64" if op in (VOP3Op.V_NOP, VOP3Op.V_PIPEFLUSH) else f"{name}_e64 {dst}, {s0}{_opsel_str(inst.opsel, 1, need_opsel, is16_d)}{cl}{om}"
+    if name in ('v_nop', 'v_pipeflush'): return f"{name}_e64"
+    # Handle byte_sel for non-pk fp8/bf8 conversions
+    if ('cvt_f32_fp8' in name or 'cvt_f32_bf8' in name) and 'pk' not in name:
+      byte_sel = inst.opsel & 3
+      os = f" byte_sel:{byte_sel}" if byte_sel else ""
+    elif 'cvt_pk_f32_bf8' in name or 'cvt_pk_f32_fp8' in name:
+      os = _opsel_str(inst.opsel, 2, need_opsel, is16_d)  # 2-element for pk variants
+    else:
+      os = _opsel_str(inst.opsel, 1, need_opsel, is16_d)  # 1-element for other VOP1
+    return f"{name}_e64 {dst}, {s0}{os}{cl}{om}"
   # Native VOP3
   is3 = _has(name, 'fma', 'mad', 'min3', 'max3', 'med3', 'div_fix', 'div_fmas', 'sad', 'lerp', 'align', 'cube', 'bfe', 'bfi',
-             'perm_b32', 'permlane', 'cndmask', 'xor3', 'or3', 'add3', 'lshl_or', 'and_or', 'lshl_add', 'add_lshl', 'xad', 'maxmin', 'minmax', 'dot2', 'cvt_pk_u8', 'mullit')
-  os = _opsel_str(inst.opsel, 3 if is3 else 2, need_opsel, is16_d)
+             'perm_b32', 'cndmask', 'xor3', 'or3', 'add3', 'lshl_or', 'and_or', 'lshl_add', 'add_lshl', 'xad', 'maxmin', 'minmax', 'dot2', 'cvt_pk_u8', 'mullit',
+             'minimummaximum', 'maximumminimum', 'minimum3', 'maximum3')
+  # permlane16/permlanex16 have 3 sources, but _var variants have 2
+  if 'permlane' in name and 'var' not in name: is3 = True
+  # Handle byte_sel for fp8/bf8 instructions (opsel encodes byte_sel, not op_sel)
+  # For VOP1-encoded VOP3 (op < 512): cvt_f32_fp8, cvt_f32_bf8
+  # For native VOP3: cvt_sr_fp8, cvt_sr_bf8
+  if ('cvt_f32_fp8' in name or 'cvt_f32_bf8' in name or 'cvt_sr_fp8' in name or 'cvt_sr_bf8' in name) and 'pk' not in name:
+    # For VOP1 encoding (op < 512), byte_sel is in bits[1:0] of opsel; for native VOP3, it's bits[3:2]
+    byte_sel = (inst.opsel & 3) if inst.op < 512 else ((inst.opsel >> 2) & 3)
+    os = f" byte_sel:{byte_sel}" if byte_sel else ""
+  else:
+    os = _opsel_str(inst.opsel, 3 if is3 else 2, need_opsel, is16_d)
   return f"{name} {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if is3 else f"{name} {dst}, {s0}, {s1}{os}{cl}{om}"
 
 def _disasm_vop3sd(inst: VOP3SD) -> str:
@@ -320,20 +382,67 @@ def _disasm_vop3sd(inst: VOP3SD) -> str:
   return f"{name}{suffix} {dst}, {_fmt_sdst(inst.sdst, 1)}, {s0}, {s1}{'' if is2src else f', {s2}'}{' clamp' if inst.clmp else ''}{_omod(inst.omod)}"
 
 def _disasm_vopd(inst: VOPD) -> str:
+  is_rdna4 = 'rdna4' in inst.__class__.__module__
+  if is_rdna4:
+    from extra.assembly.amd.autogen.rdna4.enum import VOPDOp as VOPDOp4
+    OpEnum = VOPDOp4
+  else:
+    OpEnum = VOPDOp
   lit = inst._literal or inst.literal
-  vdst_y, nx, ny = (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), VOPDOp(inst.opx).name.lower(), VOPDOp(inst.opy).name.lower()
+  vdst_y, nx, ny = (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), OpEnum(inst.opx).name.lower(), OpEnum(inst.opy).name.lower()
   def half(n, vd, s0, vs1): return f"{n} v{vd}, {inst.lit(s0)}{f', 0x{lit:x}' if lit and _has(n, 'fmaak', 'fmamk') else ''}" if 'mov' in n else f"{n} v{vd}, {inst.lit(s0)}, v{vs1}{f', 0x{lit:x}' if lit and _has(n, 'fmaak', 'fmamk') else ''}"
   return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, vdst_y, inst.srcy0, inst.vsrcy1)}"
 
-def _disasm_vop3p(inst: VOP3P) -> str:
-  name = VOP3POp(inst.op).name.lower()
-  is_wmma, is_3src, is_fma_mix = 'wmma' in name, _has(name, 'fma', 'mad', 'dot', 'wmma'), 'fma_mix' in name
-  if is_wmma:
-    sc = 2 if 'iu4' in name else 4 if 'iu8' in name else 8
-    src0, src1, src2, dst = _fmt_src(inst.src0, sc), _fmt_src(inst.src1, sc), _fmt_src(inst.src2, 8), _vreg(inst.vdst, 8)
+def _disasm_vop3p(inst: VOP3P, wave_size: int = 32) -> str:
+  is_rdna4 = 'rdna4' in inst.__class__.__module__
+  if is_rdna4:
+    from extra.assembly.amd.autogen.rdna4.enum import VOP3POp as OpEnum
+  else:
+    OpEnum = VOP3POp
+  name = OpEnum(inst.op).name.lower()
+  is_wmma, is_swmmac = 'wmma' in name and 'swmmac' not in name, 'swmmac' in name
+  is_3src, is_fma_mix = _has(name, 'fma', 'mad', 'dot', 'wmma'), 'fma_mix' in name
+  # Wave64 uses half the register widths of wave32 for WMMA
+  wave_div = 2 if wave_size == 64 else 1
+  if is_swmmac and is_rdna4:
+    # SWMMAC (sparse WMMA): src2 is a single VGPR index, not an accumulator
+    # Determine src0/src1/dst sizes based on instruction type
+    if 'f16' in name or 'bf16' in name:
+      if 'f16_16x16x32_f16' in name or 'bf16_16x16x32_bf16' in name:
+        s0c, s1c, dc = 4, 8, 4  # f16/bf16 output
+      else:
+        s0c, s1c, dc = 4, 8, 8  # f32 output
+    elif 'iu8' in name: s0c, s1c, dc = 2, 4, 8
+    elif 'iu4' in name:
+      if '16x16x64' in name: s0c, s1c, dc = 2, 4, 8
+      else: s0c, s1c, dc = 1, 2, 8
+    elif 'fp8' in name or 'bf8' in name: s0c, s1c, dc = 2, 4, 8
+    else: s0c, s1c, dc = 4, 8, 8
+    s0c, s1c, dc = max(1, s0c // wave_div), max(1, s1c // wave_div), max(1, dc // wave_div)
+    src0, src1, src2, dst = _fmt_src(inst.src0, s0c), _fmt_src(inst.src1, s1c), _fmt_src(inst.src2, 1), _vreg(inst.vdst, dc)
+  elif is_wmma:
+    # RDNA4 WMMA uses smaller source register widths than RDNA3
+    if is_rdna4:
+      # RDNA4 wave32 source widths: iu4->1/2, iu8->2, fp8/bf8->2, f16/bf16->4
+      if 'iu4' in name:
+        sc = 2 if '16x16x32' in name else 1
+      elif 'iu8' in name or 'fp8' in name or 'bf8' in name: sc = 2
+      else: sc = 4  # f16/bf16
+      # Destination width: f16/bf16 output->4, f32/i32 output->8
+      dc = 4 if name.startswith('v_wmma_f16') or name.startswith('v_wmma_bf16') else 8
+    else:
+      # RDNA3: iu4->2, iu8->4, f16/bf16->8
+      sc = 2 if 'iu4' in name else 4 if 'iu8' in name else 8
+      dc = 8
+    sc, dc = max(1, sc // wave_div), max(1, dc // wave_div)
+    src0, src1, src2, dst = _fmt_src(inst.src0, sc), _fmt_src(inst.src1, sc), _fmt_src(inst.src2, dc), _vreg(inst.vdst, dc)
   else: src0, src1, src2, dst = _fmt_src(inst.src0, 1), _fmt_src(inst.src1, 1), _fmt_src(inst.src2, 1), f"v{inst.vdst}"
   n, opsel_hi = 3 if is_3src else 2, inst.opsel_hi | (inst.opsel_hi2 << 2)
-  if is_fma_mix:
+  if is_swmmac and is_rdna4:
+    # SWMMAC uses index_key instead of op_sel; opsel bits encode the key value
+    mods = ([f"index_key:{inst.opsel & 7}"] if inst.opsel else []) + \
+           ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else [])
+  elif is_fma_mix:
     def m(s, neg, abs_): return f"-{f'|{s}|' if abs_ else s}" if neg else (f"|{s}|" if abs_ else s)
     src0, src1, src2 = m(src0, inst.neg & 1, inst.neg_hi & 1), m(src1, inst.neg & 2, inst.neg_hi & 2), m(src2, inst.neg & 4, inst.neg_hi & 4)
     mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi else []) + (["clamp"] if inst.clmp else [])
@@ -351,7 +460,7 @@ def _disasm_buf(inst: MUBUF | MTBUF) -> str:
       {'b32':1,'b64':2,'b96':3,'b128':4,'b16':1,'x':1,'xy':2,'xyz':3,'xyzw':4}.get(name.split('_')[-1], 1)
   if inst.tfe: w += 1
   vaddr = _vreg(inst.vaddr, 2) if inst.offen and inst.idxen else f"v{inst.vaddr}" if inst.offen or inst.idxen else "off"
-  srsrc = _reg("ttmp", inst.srsrc*4 - 108, 4) if 108 <= inst.srsrc*4 <= 123 else _sreg(inst.srsrc*4, 4)
+  srsrc = _sreg_or_ttmp(inst.srsrc*4, 4)
   mods = ([f"format:{inst.format}"] if isinstance(inst, MTBUF) else []) + [m for c, m in [(inst.idxen,"idxen"),(inst.offen,"offen"),(inst.offset,f"offset:{inst.offset}"),(inst.glc,"glc"),(inst.dlc,"dlc"),(inst.slc,"slc"),(inst.tfe,"tfe")] if c]
   return f"{name} {_vreg(inst.vdata, w)}, {vaddr}, {srsrc}, {decode_src(inst.soffset)}{' ' + ' '.join(mods) if mods else ''}"
 
@@ -375,12 +484,11 @@ def _mimg_vaddr_width(name: str, dim: int, a16: bool) -> int:
 def _disasm_mimg(inst: MIMG) -> str:
   name = MIMGOp(inst.op).name.lower()
   srsrc_base = inst.srsrc * 4
-  srsrc_str = _reg("ttmp", srsrc_base - 108, 8) if 108 <= srsrc_base <= 123 else _sreg(srsrc_base, 8)
+  srsrc_str = _sreg_or_ttmp(srsrc_base, 8)
   # BVH intersect ray: special case with 4 SGPR srsrc
   if 'bvh' in name:
     vaddr = (9 if '64' in name else 8) if inst.a16 else (12 if '64' in name else 11)
-    srsrc = _reg("ttmp", srsrc_base - 108, 4) if 108 <= srsrc_base <= 123 else _sreg(srsrc_base, 4)
-    return f"{name} {_vreg(inst.vdata, 4)}, {_vreg(inst.vaddr, vaddr)}, {srsrc}{' a16' if inst.a16 else ''}"
+    return f"{name} {_vreg(inst.vdata, 4)}, {_vreg(inst.vaddr, vaddr)}, {_sreg_or_ttmp(srsrc_base, 4)}{' a16' if inst.a16 else ''}"
   # vdata width from dmask (gather4/msaa_load always 4), d16 packs, tfe adds 1
   vdata = 4 if 'gather4' in name or 'msaa_load' in name else (bin(inst.dmask).count('1') or 1)
   if inst.d16: vdata = (vdata + 1) // 2
@@ -390,8 +498,8 @@ def _disasm_mimg(inst: MIMG) -> str:
   dim = dim_names[inst.dim] if inst.dim < len(dim_names) else f"dim_{inst.dim}"
   vaddr = _mimg_vaddr_width(name, inst.dim, inst.a16)
   vaddr_str = f"v{inst.vaddr}" if vaddr == 1 else _vreg(inst.vaddr, vaddr)
-  # modifiers
-  mods = [f"dmask:0x{inst.dmask:x}"] if inst.dmask and (inst.dmask != 15 or 'atomic' in name) else []
+  # modifiers - always include dmask for image load/store/atomic (LLVM uses it for vdata size validation)
+  mods = [f"dmask:0x{inst.dmask:x}"] if inst.dmask else []
   mods.append(f"dim:SQ_RSRC_IMG_{dim.upper()}")
   for flag, mod in [(inst.unrm,"unorm"),(inst.glc,"glc"),(inst.slc,"slc"),(inst.dlc,"dlc"),(inst.r128,"r128"),
                     (inst.a16,"a16"),(inst.tfe,"tfe"),(inst.lwe,"lwe"),(inst.d16,"d16")]:
@@ -399,8 +507,7 @@ def _disasm_mimg(inst: MIMG) -> str:
   # ssamp for sample/gather/get_lod
   ssamp_str = ""
   if 'sample' in name or 'gather' in name or 'get_lod' in name:
-    ssamp_base = inst.ssamp * 4
-    ssamp_str = ", " + (_reg("ttmp", ssamp_base - 108, 4) if 108 <= ssamp_base <= 123 else _sreg(ssamp_base, 4))
+    ssamp_str = ", " + _sreg_or_ttmp(inst.ssamp * 4, 4)
   return f"{name} {_vreg(inst.vdata, vdata)}, {vaddr_str}, {srsrc_str}{ssamp_str} {' '.join(mods)}"
 
 def _sop_widths(name: str) -> tuple[int, int, int]:
@@ -432,12 +539,23 @@ def _disasm_sopc(inst: SOPC) -> str:
   return f"{name} {_fmt_src(inst.ssrc0, s0n)}, {_fmt_src(inst.ssrc1, s1n)}"
 
 def _disasm_sopk(inst: SOPK) -> str:
-  op, name = SOPKOp(inst.op), SOPKOp(inst.op).name.lower()
-  if op == SOPKOp.S_VERSION: return f"{name} 0x{inst.simm16:x}"
-  if op in (SOPKOp.S_SETREG_B32, SOPKOp.S_GETREG_B32):
+  # Use architecture-specific SOPK op enum
+  if 'rdna4' in inst.__class__.__module__:
+    from extra.assembly.amd.autogen.rdna4.enum import SOPKOp as OpEnum
+    hwreg_map = HWREG_GFX12
+  else:
+    OpEnum = SOPKOp
+    hwreg_map = HWREG
+  op, name = OpEnum(inst.op), OpEnum(inst.op).name.lower()
+  if name == 's_version': return f"{name} 0x{inst.simm16:x}"
+  if name in ('s_setreg_b32', 's_getreg_b32'):
     hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
-    hs = f"0x{inst.simm16:x}" if hid in (16, 17) else f"hwreg({HWREG.get(hid, str(hid))}, {hoff}, {hsz})"
-    return f"{name} {hs}, {_fmt_sdst(inst.sdst, 1)}" if op == SOPKOp.S_SETREG_B32 else f"{name} {_fmt_sdst(inst.sdst, 1)}, {hs}"
+    hreg_name = hwreg_map.get(hid, str(hid))
+    # If offset=0 and size=32, use short form hwreg(NAME), otherwise hwreg(NAME, off, sz)
+    if hid in (16, 17): hs = f"0x{inst.simm16:x}"
+    elif hoff == 0 and hsz == 32: hs = f"hwreg({hreg_name})"
+    else: hs = f"hwreg({hreg_name}, {hoff}, {hsz})"
+    return f"{name} {hs}, {_fmt_sdst(inst.sdst, 1)}" if name == 's_setreg_b32' else f"{name} {_fmt_sdst(inst.sdst, 1)}, {hs}"
   dn, _, _ = _sop_widths(name)
   return f"{name} {_fmt_sdst(inst.sdst, dn)}, 0x{inst.simm16:x}"
 
@@ -449,20 +567,272 @@ def _disasm_vinterp(inst: VINTERP) -> str:
   mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
   return f"{name} v{inst.vdst}, {src0}, {src1}, {src2}" + (" " + mods if mods else "")
 
-def _disasm_generic(inst: Inst) -> str:
-  name = f"op_{inst.op}"
-  def format_field(field_name, val):
-    val = unwrap(val)
-    if field_name in SRC_FIELDS: return inst.lit(val) if val != 255 else "0xff"
-    return f"{'s' if field_name == 'sdst' else 'v'}{val}" if field_name in ('sdst', 'vdst') else f"v{val}" if field_name == 'vsrc1' else f"0x{val:x}" if field_name == 'simm16' else str(val)
-  operands = [format_field(field_name, inst._values.get(field_name, 0)) for field_name in inst._fields if field_name not in ('encoding', 'op')]
-  return f"{name} {', '.join(operands)}" if operands else name
+# Export targets: mrt0-7, mrtz, pos0-4, prim, dual_src_blend0/1
+_EXP_TARGETS = {**{i: f'mrt{i}' for i in range(8)}, 8: 'mrtz', **{i+12: f'pos{i}' for i in range(5)}, 20: 'prim', 21: 'dual_src_blend0', 22: 'dual_src_blend1'}
+
+def _disasm_exp(inst) -> str:
+  target = _EXP_TARGETS.get(inst.target, f"invalid_target_{inst.target}")
+  en = inst.en
+  vsrc = lambda i, v: f"v{v}" if (en >> i) & 1 else "off"
+  srcs = f"{vsrc(0, inst.vsrc0)}, {vsrc(1, inst.vsrc1)}, {vsrc(2, inst.vsrc2)}, {vsrc(3, inst.vsrc3)}"
+  mods = _mods((inst.done, "done"), (inst.row, "row_en"))
+  prefix = "export" if 'rdna4' in inst.__class__.__module__ else "exp"
+  return f"{prefix} {target} {srcs}" + (" " + mods if mods else "")
+
+def _disasm_ldsdir(inst) -> str:
+  is_rdna4 = 'rdna4' in inst.__class__.__module__
+  if is_rdna4:
+    # RDNA4 uses ds_* prefix and wait_va_vdst/wait_vm_vsrc modifiers
+    wait = f" wait_va_vdst:{inst.wait_va} wait_vm_vsrc:{inst.wait_vm}"
+    if inst.op == 1: return f"ds_direct_load v{inst.vdst}{wait}"
+    if inst.op == 0: return f"ds_param_load v{inst.vdst}, attr{inst.attr}.{['x','y','z','w'][inst.attr_chan]}{wait}"
+  else:
+    # RDNA3 uses lds_* prefix and wait_vdst modifier
+    wait = f" wait_vdst:{inst.wait_va}" if inst.wait_va != 0 else ""
+    if inst.op == 1: return f"lds_direct_load v{inst.vdst}{wait}"
+    if inst.op == 0: return f"lds_param_load v{inst.vdst}, attr{inst.attr}.{['x','y','z','w'][inst.attr_chan]}{wait}"
+  raise ValueError(f"unknown LDSDIR op: {inst.op}")
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# RDNA4-specific disassemblers (GFX12)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+# th values for RDNA4 memory instructions (based on AMDGPU ISA docs and LLVM SIDefines.h)
+# Load: 0=RT(default), 1=NT, 2=HT, 3=LU, 4=NT_RT, 5=RT_NT, 6=NT_HT, 7=BYPASS(only with scope=SYS)
+_TH_LOAD = {0: '', 1: 'th:TH_LOAD_NT', 2: 'th:TH_LOAD_HT', 3: 'th:TH_LOAD_LU', 4: 'th:TH_LOAD_NT_RT', 5: 'th:TH_LOAD_RT_NT', 6: 'th:TH_LOAD_NT_HT', 7: 'th:TH_LOAD_BYPASS'}
+# Store: 0=RT(default), 1=NT, 2=HT, 3=WB, 4=NT_RT, 5=RT_NT, 6=NT_HT, 7=NT_WB (BYPASS is th=7 + scope=SYS)
+_TH_STORE = {0: '', 1: 'th:TH_STORE_NT', 2: 'th:TH_STORE_HT', 3: 'th:TH_STORE_WB', 4: 'th:TH_STORE_NT_RT', 5: 'th:TH_STORE_RT_NT', 6: 'th:TH_STORE_NT_HT', 7: 'th:TH_STORE_NT_WB'}
+# Atomic: bit0=RETURN, bit1=NT, bit2=CASCADE -> 0=none, 1=RETURN, 2=NT, 3=NT_RETURN, 4=CASCADE_RT, 5=RT_RETURN(N/A), 6=CASCADE_NT, 7=N/A
+_TH_ATOMIC = {0: '', 1: 'th:TH_ATOMIC_RETURN', 2: 'th:TH_ATOMIC_NT', 3: 'th:TH_ATOMIC_NT_RETURN', 4: 'th:TH_ATOMIC_CASCADE_RT', 5: 'th:TH_ATOMIC_RT_RETURN', 6: 'th:TH_ATOMIC_CASCADE_NT', 7: 'th:TH_ATOMIC_CASCADE_NT'}
+_SCOPE = {0: '', 1: 'scope:SCOPE_SE', 2: 'scope:SCOPE_DEV', 3: 'scope:SCOPE_SYS'}
+
+def _rdna4_mem_mods(th: int, scope: int, is_store: bool, is_atomic: bool) -> str:
+  th_map = _TH_ATOMIC if is_atomic else _TH_STORE if is_store else _TH_LOAD
+  # Special case: th=3 with scope=SYS means BYPASS for load/store (otherwise th=3 means LU/WB)
+  if th == 3 and scope == 3 and not is_atomic:
+    th_s = 'th:TH_STORE_BYPASS' if is_store else 'th:TH_LOAD_BYPASS'
+  else:
+    th_s = th_map.get(th, f'th:{th}' if th else '')
+  scope_s = _SCOPE.get(scope, f'scope:{scope}' if scope else '')
+  return ' '.join(x for x in [th_s, scope_s] if x)
+
+def _disasm_vflat(inst) -> str:
+  """Disassemble RDNA4 VFLAT/VGLOBAL/VSCRATCH instructions."""
+  from extra.assembly.amd.autogen.rdna4.enum import VFLATOp, VGLOBALOp, VSCRATCHOp
+  cls_name = type(inst).__name__
+  if cls_name == 'VGLOBAL': op_enum, prefix = VGLOBALOp, 'global'
+  elif cls_name == 'VSCRATCH': op_enum, prefix = VSCRATCHOp, 'scratch'
+  else: op_enum, prefix = VFLATOp, 'flat'
+  name = op_enum(inst.op).name.lower()
+
+  # global_wb, global_wbinv, global_inv are cache control instructions with no operands
+  if name in ('global_wb', 'global_wbinv', 'global_inv'):
+    mods = _rdna4_mem_mods(inst.th, inst.scope, False, False)
+    return f"{name}" + (f" {mods}" if mods else "")
+
+  # addtid instructions use thread ID as address offset, no vaddr operand
+  is_addtid = 'addtid' in name
+
+  # Data width based on instruction name suffix
+  suffix = name.split('_')[-1]
+  # block loads/stores use 32 VGPRs
+  if 'block' in name:
+    base_w = 32
+  else:
+    base_w = {'b32':1,'b64':2,'b96':3,'b128':4,'u8':1,'i8':1,'u16':1,'i16':1,'u32':1,'i32':1,'u64':2,'i64':2,'f32':1,'f64':2}.get(suffix, 1)
+  # For cmpswap: vsrc holds cmp+data pairs (2x base), vdst is base width
+  vsrc_w = base_w * 2 if 'cmpswap' in name else base_w
+  vdst_w = base_w
+
+  # Offset: signed 24-bit (stored as unsigned, needs sign extension)
+  off = inst.ioffset if inst.ioffset < 0x800000 else inst.ioffset - 0x1000000
+  off_s = f" offset:{off}" if off else ""
+
+  # Memory modifiers
+  is_store, is_atomic = 'store' in name, 'atomic' in name
+  mods = _rdna4_mem_mods(inst.th, inst.scope, is_store, is_atomic)
+
+  # saddr handling - VGLOBAL and VSCRATCH need explicit "off" when saddr=124
+  if inst.saddr == 124: saddr_s = "off" if prefix in ('global', 'scratch') else ""
+  elif inst.saddr in SPECIAL_PAIRS: saddr_s = SPECIAL_PAIRS[inst.saddr]
+  else: saddr_s = _sreg(inst.saddr, 2) if prefix == 'global' else decode_src(inst.saddr)
+
+  # Address width: 1 for scratch with saddr, 2 otherwise
+  addr_w = 1 if (prefix == 'scratch' or (inst.saddr != 124 and prefix != 'flat')) else 2
+  vaddr_s = _vreg(inst.vaddr, addr_w)
+  vsrc_s = _vreg(inst.vsrc, vsrc_w)
+  vdst_s = _vreg(inst.vdst, vdst_w)
+
+  # addtid instructions don't have vaddr, just vdata and saddr
+  if is_addtid:
+    if is_store: return f"{name} {vsrc_s}, {saddr_s}{off_s}" + (f" {mods}" if mods else "")
+    return f"{name} {vdst_s}, {saddr_s}{off_s}" + (f" {mods}" if mods else "")
+
+  # Regular instructions need comma before saddr
+  saddr_s = f", {saddr_s}" if saddr_s else ""
+
+  if is_atomic:
+    if inst.th == 1:  # TH_ATOMIC_RETURN
+      return f"{name} {vdst_s}, {vaddr_s}, {vsrc_s}{saddr_s}{off_s}" + (f" {mods}" if mods else "")
+    return f"{name} {vaddr_s}, {vsrc_s}{saddr_s}{off_s}" + (f" {mods}" if mods else "")
+  if is_store: return f"{name} {vaddr_s}, {vsrc_s}{saddr_s}{off_s}" + (f" {mods}" if mods else "")
+  return f"{name} {vdst_s}, {vaddr_s}{saddr_s}{off_s}" + (f" {mods}" if mods else "")
+
+def _disasm_vbuffer(inst) -> str:
+  """Disassemble RDNA4 VBUFFER instructions."""
+  from extra.assembly.amd.autogen.rdna4.enum import VBUFFEROp
+  name = VBUFFEROp(inst.op).name.lower()
+
+  # Determine if this is a typed buffer instruction (MTBUF format)
+  is_format = 'format' in name
+
+  # Data width based on instruction name
+  if is_format:
+    w = {'x': 1, 'xy': 2, 'xyz': 3, 'xyzw': 4}.get(name.split('_')[-1], 1)
+    if 'd16' in name: w = (w + 1) // 2
+  else:
+    suffix = name.split('_')[-1]
+    w = {'b32':1,'b64':2,'b96':3,'b128':4,'u8':1,'i8':1,'u16':1,'i16':1,'u32':1,'i32':1,'u64':2,'i64':2,'b8':1,'b16':1,'f16':1,'f32':1,'bf16':1}.get(suffix, 1)
+  if 'cmpswap' in name: w *= 2
+  if inst.tfe: w += 1
+
+  vdata_s = _vreg(inst.vdata, w)
+  vaddr_s = _vreg(inst.vaddr, 2) if inst.offen and inst.idxen else (f"v{inst.vaddr}" if inst.offen or inst.idxen else "off")
+  # RDNA4 VBUFFER rsrc field stores the SGPR index directly (not /4 like RDNA3)
+  srsrc_s = _sreg_or_ttmp(inst.rsrc, 4)
+  soffset_s = decode_src(inst.soffset)
+
+  off = inst.ioffset if inst.ioffset < 0x800000 else inst.ioffset - 0x1000000
+  is_store, is_atomic = 'store' in name, 'atomic' in name
+  mods = _rdna4_mem_mods(inst.th, inst.scope, is_store, is_atomic)
+
+  # Format field is only for MTBUF (tbuffer_*) instructions, not buffer_*_format_* instructions
+  # We don't output format for buffer instructions since they use implicit format
+  parts = []
+  if inst.idxen: parts.append("idxen")
+  if inst.offen: parts.append("offen")
+  if off: parts.append(f"offset:{off}")
+  if mods: parts.append(mods)
+  if inst.tfe: parts.append("tfe")
+
+  return f"{name} {vdata_s}, {vaddr_s}, {srsrc_s}, {soffset_s}" + (f" {' '.join(parts)}" if parts else "")
+
+def _disasm_vimage(inst) -> str:
+  """Disassemble RDNA4 VIMAGE instructions."""
+  from extra.assembly.amd.autogen.rdna4.enum import VIMAGEOp
+  name = VIMAGEOp(inst.op).name.lower()
+
+  # RDNA4 VIMAGE rsrc field stores the SGPR index directly (not /4 like RDNA3)
+  if 'bvh' in name:
+    # BVH intersect ray: special format with individual/range vaddr components
+    # Format: [node_ptr, ray_extent, ray_origin(3), ray_dir(3), ray_inv_dir(3)]
+    # bvh64 has 2-VGPR node_ptr, a16 removes ray_inv_dir
+    if 'dual' in name or 'bvh8' in name:
+      # dual/bvh8: [node_ptr(2), ray_extent(2), ray_origin(3), ray_dir(3), ...]
+      parts = [_vreg(inst.vaddr0, 2), _vreg(inst.vaddr1, 2), _vreg(inst.vaddr2, 3), _vreg(inst.vaddr3, 3)]
+      if not inst.a16: parts.append(_vreg(inst.vaddr4, 1 if 'bvh8' in name else 2))
+      dst_w = 10
+    elif '64' in name:
+      parts = [_vreg(inst.vaddr0, 2), f"v{inst.vaddr1}", _vreg(inst.vaddr2, 3), _vreg(inst.vaddr3, 3)]
+      if not inst.a16: parts.append(_vreg(inst.vaddr4, 3))
+      dst_w = 4
+    else:
+      parts = [f"v{inst.vaddr0}", f"v{inst.vaddr1}", _vreg(inst.vaddr2, 3), _vreg(inst.vaddr3, 3)]
+      if not inst.a16: parts.append(_vreg(inst.vaddr4, 3))
+      dst_w = 4
+    return f"{name} {_vreg(inst.vdata, dst_w)}, [{', '.join(parts)}], {_sreg_or_ttmp(inst.rsrc, 4)}{' a16' if inst.a16 else ''}"
+
+  # vdata width - msaa_load always uses 4 VGPRs per channel
+  vdata = 4 if 'gather4' in name or 'msaa_load' in name else (bin(inst.dmask).count('1') or 1)
+  if inst.d16: vdata = (vdata + 1) // 2
+  if inst.tfe: vdata += 1
+
+  # dim names
+  dim_names = ['1d', '2d', '3d', 'cube', '1d_array', '2d_array', '2d_msaa', '2d_msaa_array']
+  dim = dim_names[inst.dim] if inst.dim < len(dim_names) else f"{inst.dim}"
+
+  # vaddr width calculation (RDNA4 uses vaddr0-4 for address components)
+  vaddr_w = _mimg_vaddr_width(name, inst.dim, inst.a16)
+  if vaddr_w == 1: vaddr_s = f"v{inst.vaddr0}"
+  elif vaddr_w == 2: vaddr_s = f"[v{inst.vaddr0}, v{inst.vaddr1}]"
+  elif vaddr_w == 3: vaddr_s = f"[v{inst.vaddr0}, v{inst.vaddr1}, v{inst.vaddr2}]"
+  elif vaddr_w == 4: vaddr_s = f"[v{inst.vaddr0}, v{inst.vaddr1}, v{inst.vaddr2}, v{inst.vaddr3}]"
+  else: vaddr_s = f"[v{inst.vaddr0}, v{inst.vaddr1}, v{inst.vaddr2}, v{inst.vaddr3}, v{inst.vaddr4}]"
+
+  srsrc_s = _sreg_or_ttmp(inst.rsrc, 8)
+  # RDNA4 always requires dmask for size calculation (even if 0xf)
+  mods = [f"dmask:0x{inst.dmask:x}"]
+  mods.append(f"dim:SQ_RSRC_IMG_{dim.upper()}")
+  # Add th/scope before other modifiers, then r128, then a16/tfe/d16 (LLVM expects this order)
+  if inst.th or inst.scope:
+    is_store, is_atomic = 'store' in name, 'atomic' in name
+    mem_mods = _rdna4_mem_mods(inst.th, inst.scope, is_store, is_atomic)
+    if mem_mods: mods.append(mem_mods)
+  if inst.r128: mods.append("r128")
+  mods.extend([m for c, m in [(inst.a16, "a16"), (inst.tfe, "tfe"), (inst.d16, "d16")] if c])
+
+  return f"{name} {_vreg(inst.vdata, vdata)}, {vaddr_s}, {srsrc_s} {' '.join(mods)}"
+
+def _disasm_vsample(inst) -> str:
+  """Disassemble RDNA4 VSAMPLE instructions."""
+  from extra.assembly.amd.autogen.rdna4.enum import VSAMPLEOp
+  name = VSAMPLEOp(inst.op).name.lower()
+
+  # vdata width
+  vdata = 4 if 'gather4' in name or 'msaa_load' in name else (bin(inst.dmask).count('1') or 1)
+  if inst.d16: vdata = (vdata + 1) // 2
+  if inst.tfe: vdata += 1
+
+  dim_names = ['1d', '2d', '3d', 'cube', '1d_array', '2d_array', '2d_msaa', '2d_msaa_array']
+  dim = dim_names[inst.dim] if inst.dim < len(dim_names) else f"{inst.dim}"
+
+  vaddr = _mimg_vaddr_width(name, inst.dim, inst.a16)
+  if vaddr == 1: vaddr_s = f"v{inst.vaddr0}"
+  elif vaddr == 2: vaddr_s = f"[v{inst.vaddr0}, v{inst.vaddr1}]"
+  elif vaddr == 3: vaddr_s = f"[v{inst.vaddr0}, v{inst.vaddr1}, v{inst.vaddr2}]"
+  elif vaddr == 4: vaddr_s = f"[v{inst.vaddr0}, v{inst.vaddr1}, v{inst.vaddr2}, v{inst.vaddr3}]"
+  else:
+    # More than 4 vaddrs: vaddr3 becomes start of a contiguous range for the remaining coords
+    extra = vaddr - 3  # vaddr0-2 are individual, vaddr3 starts range of remaining
+    vaddr_s = f"[v{inst.vaddr0}, v{inst.vaddr1}, v{inst.vaddr2}, {_vreg(inst.vaddr3, extra)}]"
+
+  # RDNA4 VSAMPLE rsrc/samp fields store the SGPR index directly (not /4 like RDNA3)
+  srsrc_s = _sreg_or_ttmp(inst.rsrc, 8)
+
+  # msaa_load doesn't use a sampler (it's a load, not a sample), but gather4h does
+  uses_sampler = 'msaa_load' not in name
+  ssamp_s = f", {_sreg_or_ttmp(inst.samp, 4)}" if uses_sampler else ""
+
+  mods = [f"dmask:0x{inst.dmask:x}"] if inst.dmask else []
+  mods.append(f"dim:SQ_RSRC_IMG_{dim.upper()}")
+  if inst.unrm: mods.append("unorm")
+  # th/scope must come before r128, a16, tfe, lwe, d16
+  if inst.th or inst.scope:
+    mem_mods = _rdna4_mem_mods(inst.th, inst.scope, False, False)
+    if mem_mods: mods.append(mem_mods)
+  mods.extend([m for c, m in [(inst.r128, "r128"), (inst.a16, "a16"), (inst.tfe, "tfe"), (inst.lwe, "lwe"), (inst.d16, "d16")] if c])
+
+  return f"{name} {_vreg(inst.vdata, vdata)}, {vaddr_s}, {srsrc_s}{ssamp_s} {' '.join(mods)}"
 
 DISASM_HANDLERS = {VOP1: _disasm_vop1, VOP2: _disasm_vop2, VOPC: _disasm_vopc, VOP3: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOPD: _disasm_vopd, VOP3P: _disasm_vop3p,
                    VINTERP: _disasm_vinterp, SOPP: _disasm_sopp, SMEM: _disasm_smem, DS: _disasm_ds, FLAT: _disasm_flat, MUBUF: _disasm_buf, MTBUF: _disasm_buf,
-                   MIMG: _disasm_mimg, SOP1: _disasm_sop1, SOP2: _disasm_sop2, SOPC: _disasm_sopc, SOPK: _disasm_sopk}
+                   MIMG: _disasm_mimg, SOP1: _disasm_sop1, SOP2: _disasm_sop2, SOPC: _disasm_sopc, SOPK: _disasm_sopk, EXP: _disasm_exp, LDSDIR: _disasm_ldsdir}
 
-def disasm(inst: Inst) -> str: return DISASM_HANDLERS.get(type(inst), _disasm_generic)(inst)
+# RDNA4 uses different class names, dispatch by name for cross-arch support
+_DISASM_BY_NAME = {'VOP1': _disasm_vop1, 'VOP2': _disasm_vop2, 'VOPC': _disasm_vopc, 'VOP3': _disasm_vop3, 'VOP3SD': _disasm_vop3sd,
+                   'VOPD': _disasm_vopd, 'VOP3P': _disasm_vop3p, 'VINTERP': _disasm_vinterp, 'SOPP': _disasm_sopp, 'SMEM': _disasm_smem,
+                   'DS': _disasm_ds, 'VDS': _disasm_ds, 'FLAT': _disasm_flat, 'MUBUF': _disasm_buf, 'MTBUF': _disasm_buf, 'MIMG': _disasm_mimg,
+                   'SOP1': _disasm_sop1, 'SOP2': _disasm_sop2, 'SOPC': _disasm_sopc, 'SOPK': _disasm_sopk,
+                   'EXP': _disasm_exp, 'LDSDIR': _disasm_ldsdir, 'VEXPORT': _disasm_exp, 'VDSDIR': _disasm_ldsdir,
+                   'VFLAT': _disasm_vflat, 'VGLOBAL': _disasm_vflat, 'VSCRATCH': _disasm_vflat,
+                   'VBUFFER': _disasm_vbuffer, 'VIMAGE': _disasm_vimage, 'VSAMPLE': _disasm_vsample}
+
+def disasm(inst: Inst, wave_size: int = 32) -> str:
+  handler = DISASM_HANDLERS.get(type(inst)) or _DISASM_BY_NAME.get(type(inst).__name__)
+  if handler is None: raise KeyError(f"no disasm handler for {type(inst).__name__}")
+  # For VOP3P (includes WMMA), pass wave_size if handler supports it
+  if handler == _disasm_vop3p: return _disasm_vop3p(inst, wave_size)
+  return handler(inst)
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # ASSEMBLER
@@ -470,7 +840,7 @@ def disasm(inst: Inst) -> str: return DISASM_HANDLERS.get(type(inst), _disasm_ge
 
 SPEC_REGS = {'vcc_lo': RawImm(106), 'vcc_hi': RawImm(107), 'vcc': RawImm(106), 'null': RawImm(124), 'off': RawImm(124), 'm0': RawImm(125),
              'exec_lo': RawImm(126), 'exec_hi': RawImm(127), 'exec': RawImm(126), 'scc': RawImm(253), 'src_scc': RawImm(253)}
-FLOATS = {'0.5': 0.5, '-0.5': -0.5, '1.0': 1.0, '-1.0': -1.0, '2.0': 2.0, '-2.0': -2.0, '4.0': 4.0, '-4.0': -4.0}
+FLOATS = {str(k): k for k in FLOAT_ENC}  # Valid float literal strings: '0.5', '-0.5', '1.0', etc.
 REG_MAP: dict[str, _RegFactory] = {'s': s, 'v': v, 't': ttmp, 'ttmp': ttmp}
 SMEM_OPS = {'s_load_b32', 's_load_b64', 's_load_b128', 's_load_b256', 's_load_b512',
             's_buffer_load_b32', 's_buffer_load_b64', 's_buffer_load_b128', 's_buffer_load_b256', 's_buffer_load_b512'}
@@ -656,9 +1026,8 @@ def get_dsl(text: str) -> str:
   return f"{fn}({a_str}, {kw_str})" if kw_str and a_str else f"{fn}({kw_str})" if kw_str else f"{fn}({a_str})"
 
 def asm(text: str) -> Inst:
-  from extra.assembly.amd.autogen import rdna3 as ag
   dsl = get_dsl(text)
-  ns = {n: getattr(ag, n) for n in dir(ag) if not n.startswith('_')}
+  ns = {n: getattr(ins, n) for n in dir(ins) if not n.startswith('_')}
   ns.update({'s': s, 'v': v, 'ttmp': ttmp, 'abs': abs, 'RawImm': RawImm, 'SrcMod': SrcMod, 'VGPR': VGPR, 'SGPR': SGPR, 'TTMP': TTMP,
              'VCC_LO': VCC_LO, 'VCC_HI': VCC_HI, 'VCC': VCC, 'EXEC_LO': EXEC_LO, 'EXEC_HI': EXEC_HI, 'EXEC': EXEC, 'SCC': SCC, 'M0': M0, 'NULL': NULL, 'OFF': OFF})
   try: return eval(dsl, ns)

extra/assembly/amd/dsl.py

@@ -1,8 +1,125 @@
 # 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:
@@ -25,13 +142,35 @@ class BitField:
   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:
@@ -112,31 +251,36 @@ def unwrap(val) -> int:
   if hasattr(val, 'idx'): return val.idx  # Reg
   return val
 
-# Encoding helpers
+# 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:
-  if isinstance(val, TTMP): return 108 + val.idx
-  return val.idx  # hi bit is handled via opsel, not in register encoding
+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):
-    # SrcMod wraps either special registers (VCC_LO=106, EXEC_LO=126, etc.) or literals
-    # Special register values are in valid encoding ranges - return as-is
-    # Literals (large integers) need 255 marker
-    v = val.val
-    # Valid source encoding ranges: 0-127 (SGPRs/special), 128-192 (inline const), 193-208 (neg inline), 240-247 (float), 251-253 (special)
-    if 0 <= v <= 127 or 240 <= v <= 255: return v  # SGPRs, special regs, float constants
-    if 128 <= v <= 192: return v  # Inline positive constants (0-64)
-    if 193 <= v <= 208: return v  # Inline negative constants (-1 to -16)
-    return 255  # Literal marker - value stored separately
+  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)
-  return 128 + val if isinstance(val, int) and 0 <= val <= 64 else 192 + (-val) if isinstance(val, int) and -16 <= val <= -1 else 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:
@@ -152,119 +296,107 @@ class Inst:
     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), literal
-    # Map positional args to field names
+    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))
-    orig_args.update(kwargs)
+    orig_args = dict(zip(field_names, args)) | kwargs
     self._values.update(orig_args)
-    # Validate register counts for SMEM instructions (before encoding)
-    if self.__class__.__name__ == 'SMEM':
-      op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
-      if op_val is not None:
-        if hasattr(op_val, 'value'): op_val = op_val.value
-        expected_cnt = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op_val)
-        sdata_val = orig_args.get('sdata')
-        if expected_cnt is not None and isinstance(sdata_val, Reg) and sdata_val.count != expected_cnt:
-          raise ValueError(f"SMEM op {op_val} expects {expected_cnt} registers, got {sdata_val.count}")
-    # Validate register counts for SOP1 instructions (b32 = 1 reg, b64 = 2 regs)
-    if self.__class__.__name__ == 'SOP1':
-      op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
-      if op_val is not None and hasattr(op_val, 'name'):
-        expected = 2 if op_val.name.endswith('_B64') else 1
-        sdst_val, ssrc0_val = orig_args.get('sdst'), orig_args.get('ssrc0')
-        if isinstance(sdst_val, Reg) and sdst_val.count != expected:
-          raise ValueError(f"SOP1 {op_val.name} expects {expected} destination register(s), got {sdst_val.count}")
-        if isinstance(ssrc0_val, Reg) and ssrc0_val.count != expected:
-          raise ValueError(f"SOP1 {op_val.name} expects {expected} source register(s), got {ssrc0_val.count}")
-    # FLAT: set sve=1 when addr is a VGPR for scratch only
-    # For scratch (seg=1), sve=1 means addr VGPR is used; sve=0 means addr is "off"
-    # For global (seg=2) and flat (seg=0), sve is always 0
-    if self.__class__.__name__ == 'FLAT' and 'sve' in self._fields:
-      seg_val = self._values.get('seg', 0)
-      if isinstance(seg_val, RawImm): seg_val = seg_val.val
-      addr_val = orig_args.get('addr')
-      if seg_val == 1 and isinstance(addr_val, VGPR): self._values['sve'] = 1
-    # VOP3P: v_fma_mix* instructions (opcodes 32-34) have opsel_hi default of 0, not 7
-    if self.__class__.__name__ == 'VOP3P':
-      op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
-      if hasattr(op_val, 'value'): op_val = op_val.value
-      if op_val in (32, 33, 34) and 'opsel_hi' not in orig_args and 'opsel_hi2' not in orig_args:
-        self._values['opsel_hi'] = 0
-        self._values['opsel_hi2'] = 0
-    # Type check and encode values
+    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
-      # For RawImm, only process RAW_FIELDS to unwrap to int
       if isinstance(val, RawImm):
         if name in RAW_FIELDS: self._values[name] = val.val
         continue
       field = self._fields.get(name)
       marker = field.marker if field else None
       # Type validation
-      if marker is _SGPRField:
-        if isinstance(val, VGPR): raise TypeError(f"field '{name}' requires SGPR, got VGPR")
-        if not isinstance(val, (SGPR, TTMP, SrcMod, int, RawImm)): raise TypeError(f"field '{name}' requires SGPR, got {type(val).__name__}")
-      if marker is _VGPRField:
-        if not isinstance(val, VGPR): raise TypeError(f"field '{name}' requires VGPR, got {type(val).__name__}")
+      if marker is _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 source fields as RawImm for consistent disassembly
-      if name in SRC_FIELDS:
-        encoded = encode_src(val)
-        # For VOP1/VOP2/VOPC (no opsel field), encode hi bit in src value
-        if isinstance(val, Reg) and val.hi and 'opsel' not in self._fields:
-          encoded |= 0x80
-        self._values[name] = RawImm(encoded)
-        # Handle neg/abs/opsel modifiers for VOP3 instructions
-        if isinstance(val, SrcMod):
-          if val.neg and 'neg' in self._fields:
-            neg_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
-            cur_neg = self._values.get('neg', 0)
-            self._values['neg'] = (cur_neg.val if isinstance(cur_neg, RawImm) else cur_neg) | neg_bit
-          if val.abs_ and 'abs' in self._fields:
-            abs_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
-            cur_abs = self._values.get('abs', 0)
-            self._values['abs'] = (cur_abs.val if isinstance(cur_abs, RawImm) else cur_abs) | abs_bit
-        # Handle hi (opsel) for 16-bit ops - only for formats with opsel field
-        if isinstance(val, Reg) and val.hi and 'opsel' in self._fields:
-          opsel_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
-          cur_opsel = self._values.get('opsel', 0)
-          self._values['opsel'] = (cur_opsel.val if isinstance(cur_opsel, RawImm) else cur_opsel) | opsel_bit
-        # Track literal value if needed (encoded as 255)
-        # For 64-bit ops, store literal in high 32 bits (to match from_bytes decoding and to_bytes encoding)
-        if encoded == 255 and self._literal is None:
-          if isinstance(val, SrcMod) and not isinstance(val, Reg):
-            # SrcMod wrapping a literal value
-            self._literal = (val.val << 32) if self._is_64bit_op() else val.val
-          elif isinstance(val, int) and not isinstance(val, IntEnum):
-            self._literal = (val << 32) if self._is_64bit_op() else val
-          elif isinstance(val, float):
-            import struct
-            lit32 = struct.unpack('<I', struct.pack('<f', val))[0]
-            self._literal = (lit32 << 32) if self._is_64bit_op() else lit32
-      # Encode raw register fields for consistent repr
-      elif name in RAW_FIELDS:
-        if isinstance(val, Reg):
-          encoded = _encode_reg(val)
-          # For VOP1/VOP2/VOPC (no opsel field), encode hi bit in register value
-          if val.hi and 'opsel' not in self._fields:
-            encoded |= 0x80
-          self._values[name] = encoded
-          # Handle vdst hi (opsel bit 3) for 16-bit ops - only for formats with opsel field
-          if name == 'vdst' and val.hi and 'opsel' in self._fields:
-            cur_opsel = self._values.get('opsel', 0)
-            self._values['opsel'] = (cur_opsel.val if isinstance(cur_opsel, RawImm) else cur_opsel) | 8
-        elif hasattr(val, 'value'): self._values[name] = val.value  # IntEnum like SrcEnum.NULL
-      # Encode sbase (divided by 2) and srsrc/ssamp (divided by 4)
-      elif name == 'sbase':
-        if isinstance(val, Reg): self._values[name] = val.idx // 2
-        elif isinstance(val, SrcMod): self._values[name] = val.val // 2  # Special regs like VCC_LO
-      elif name in {'srsrc', 'ssamp'} and isinstance(val, Reg):
-        self._values[name] = val.idx // 4
-      # VOPD vdsty: encode as actual >> 1 (constraint: vdsty parity must be opposite of vdstx)
-      elif marker is _VDSTYEnc and isinstance(val, VGPR):
-        self._values[name] = val.idx >> 1
+      # 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
@@ -287,35 +419,39 @@ class Inst:
     return None
 
   def _is_64bit_op(self) -> bool:
-    """Check if this instruction uses 64-bit operands (and thus 64-bit literals).
-    Exception: V_LDEXP_F64 has 32-bit integer src1, so its literal is 32-bit."""
+    """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 may be an enum (from __init__) or an int (from from_int)
     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':
-      from extra.assembly.amd.autogen.rdna3 import VOP3Op
       try: op_name = VOP3Op(op).name
       except ValueError: pass
     if op_name is None and self.__class__.__name__ == 'VOPC':
-      from extra.assembly.amd.autogen.rdna3 import VOPCOp
       try: op_name = VOPCOp(op).name
       except ValueError: pass
     if op_name is None: return False
-    # V_LDEXP_F64 has 32-bit integer exponent in src1, so literal is 32-bit
-    if op_name == 'V_LDEXP_F64': return False
-    return op_name.endswith(('_F64', '_B64', '_I64', '_U64'))
+    # 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 ops, literal is stored in high 32 bits internally, but encoded as 4 bytes
-    lit32 = (lit >> 32) if self._is_64bit_op() else lit
-    return result + (lit32 & 0xffffffff).to_bytes(4, 'little')
+    # 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 8
+  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)
@@ -341,9 +477,16 @@ class Inst:
     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')
-        inst._literal = (lit32 << 32) if inst._is_64bit_op() else lit32
+        # 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):
@@ -358,9 +501,9 @@ class Inst:
     if name.startswith('_'): raise AttributeError(name)
     return unwrap(self._values.get(name, 0))
 
-  def lit(self, v: int) -> str:
-    from extra.assembly.amd.asm import decode_src
-    return f"0x{self._literal:x}" if v == 255 and self._literal else decode_src(v)
+  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
@@ -368,316 +511,64 @@ class Inst:
 
   def __hash__(self): return hash((self.__class__.__name__, tuple(sorted((k, repr(v)) for k, v in self._values.items())), self._literal))
 
-  def disasm(self) -> str:
+  def disasm(self, wave_size: int = 32) -> str:
     from extra.assembly.amd.asm import disasm
-    return disasm(self)
+    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
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# CODE GENERATION: generates autogen/__init__.py by parsing AMD ISA PDFs
-# Supports both RDNA3.5 and CDNA4 instruction set PDFs - auto-detects format
-# ═══════════════════════════════════════════════════════════════════════════════
-
-PDF_URLS = {
-  "rdna3": "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content",  # RDNA3.5
-  "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_TYPES = {'SSRC0': 'SSrc', 'SSRC1': 'SSrc', 'SOFFSET': 'SSrc', 'SADDR': 'SSrc', 'SRC0': 'Src', 'SRC1': 'Src', 'SRC2': 'Src',
-  'SDST': 'SGPRField', 'SBASE': 'SGPRField', 'SDATA': 'SGPRField', 'SRSRC': 'SGPRField', 'VDST': 'VGPRField', 'VSRC1': 'VGPRField', 'VDATA': 'VGPRField',
-  'VADDR': 'VGPRField', 'ADDR': 'VGPRField', 'DATA': 'VGPRField', 'DATA0': 'VGPRField', 'DATA1': 'VGPRField', 'SIMM16': 'SImm', 'OFFSET': 'Imm',
-  'OPX': 'VOPDOp', 'OPY': 'VOPDOp', 'SRCX0': 'Src', 'SRCY0': 'Src', 'VSRCX1': 'VGPRField', 'VSRCY1': 'VGPRField', 'VDSTX': 'VGPRField', 'VDSTY': 'VDSTYEnc'}
-FIELD_ORDER = {
-  'SOP2': ['op', 'sdst', 'ssrc0', 'ssrc1'], 'SOP1': ['op', 'sdst', 'ssrc0'], 'SOPC': ['op', 'ssrc0', 'ssrc1'],
-  'SOPK': ['op', 'sdst', 'simm16'], 'SOPP': ['op', 'simm16'], 'VOP1': ['op', 'vdst', 'src0'], 'VOPC': ['op', 'src0', 'vsrc1'],
-  'VOP2': ['op', 'vdst', 'src0', 'vsrc1'], 'VOP3SD': ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2', 'clmp'],
-  'SMEM': ['op', 'sdata', 'sbase', 'soffset', 'offset', 'glc', 'dlc'], 'DS': ['op', 'vdst', 'addr', 'data0', 'data1'],
-  'VOP3': ['op', 'vdst', 'src0', 'src1', 'src2', 'omod', 'neg', 'abs', 'clmp', 'opsel'],
-  'VOP3P': ['op', 'vdst', 'src0', 'src1', 'src2', 'neg', 'neg_hi', 'opsel', 'opsel_hi', 'clmp'],
-  'FLAT': ['op', 'vdst', 'addr', 'data', 'saddr', 'offset', 'seg', 'dlc', 'glc', 'slc'],
-  'MUBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
-  'MTBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
-  'MIMG': ['op', 'vdata', 'vaddr', 'srsrc', 'ssamp', 'dmask', 'dim', 'unrm', 'dlc', 'glc', 'slc'],
-  'EXP': ['en', 'target', 'vsrc0', 'vsrc1', 'vsrc2', 'vsrc3', 'done', 'row'],
-  'VINTERP': ['op', 'vdst', 'src0', 'src1', 'src2', 'waitexp', 'clmp', 'opsel', 'neg'],
-  'VOPD': ['opx', 'opy', 'vdstx', 'vdsty', 'srcx0', 'vsrcx1', 'srcy0', 'vsrcy1'],
-  'LDSDIR': ['op', 'vdst', 'attr', 'attr_chan', 'wait_va']}
-SRC_EXTRAS = {233: 'DPP8', 234: 'DPP8FI', 250: 'DPP16', 251: 'VCCZ', 252: 'EXECZ', 254: 'LDS_DIRECT'}
-FLOAT_MAP = {'0.5': 'POS_HALF', '-0.5': 'NEG_HALF', '1.0': 'POS_ONE', '-1.0': 'NEG_ONE', '2.0': 'POS_TWO', '-2.0': 'NEG_TWO',
-  '4.0': 'POS_FOUR', '-4.0': 'NEG_FOUR', '1/(2*PI)': 'INV_2PI', '0': 'ZERO'}
-
-def _parse_bits(s: str) -> tuple[int, int] | None:
-  import re
-  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]:
-  import re
-  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]:
-      # Handle both RDNA3 ('bXX) and CDNA4 (Must be: XX) encoding formats
-      if m := re.search(r"(?:'b|Must be:\s*)([01_]+)", row[2]):
-        enc_bits = m.group(1).replace('_', '')
-        enc_val = int(enc_bits, 2)
-        declared_width, actual_width = hi - lo + 1, len(enc_bits)
-        if actual_width > declared_width: lo = hi - actual_width + 1
-    ftype = f"{fmt}Op" if name == 'OP' and f"{fmt}Op" in enums else FIELD_TYPES.get(name.upper())
-    fields.append((name, hi, lo, enc_val, ftype))
-  return fields
-
-def _parse_single_pdf(url: str) -> dict:
-  """Parse a single PDF and return raw data (formats, enums, src_enum, doc_name, is_cdna)."""
-  import re, pdfplumber
-  from tinygrad.helpers import fetch
-
-  pdf = pdfplumber.open(fetch(url))
-
-  # Auto-detect document type from first page
-  first_page_text = pdf.pages[0].extract_text() or ''
-  is_cdna4 = 'CDNA4' in first_page_text or 'CDNA 4' in first_page_text
-  is_cdna3 = 'CDNA3' in first_page_text or 'CDNA 3' in first_page_text or 'MI300' in first_page_text
-  is_cdna = is_cdna3 or is_cdna4
-  is_rdna4 = 'RDNA4' in first_page_text or 'RDNA 4' in first_page_text
-  is_rdna35 = 'RDNA3.5' in first_page_text or 'RDNA 3.5' in first_page_text  # Check 3.5 before 3
-  is_rdna3 = not is_rdna35 and ('RDNA3' in first_page_text or 'RDNA 3' in first_page_text)
-  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 the "Microcode Formats" section - search for SOP2 format definition
-  microcode_start = None
-  total_pages = len(pdf.pages)
-  # Search from likely locations (formats are typically 20-95% through the document - RDNA3 has them at ~25%)
-  for i in range(int(total_pages * 0.2), total_pages):
-    text = pdf.pages[i].extract_text() or ''
-    # Look for "X.Y.Z. SOP2" section header or "Chapter X. Microcode Formats"
-    if re.search(r'\d+\.\d+\.\d+\.\s+SOP2\b', text) or re.search(r'Chapter \d+\.\s+Microcode Formats', text):
-      microcode_start = i
-      break
-  if microcode_start is None: microcode_start = int(total_pages * 0.9)
-
-  pages = pdf.pages[microcode_start:microcode_start + 50]
-  page_texts = [p.extract_text() or '' for p in pages]
-  page_tables = [[t.extract() for t in p.find_tables()] for p in pages]
-  full_text = '\n'.join(page_texts)
-
-  # parse SSRC encoding from first page with VCC_LO
-  src_enum = dict(SRC_EXTRAS)
-  for text in page_texts[:10]:
-    if 'SSRC0' in text and 'VCC_LO' in text:
-      for m in re.finditer(r'^(\d+)\s+(\S+)', text, re.M):
-        val, name = int(m.group(1)), m.group(2).rstrip('.:')
-        if name in FLOAT_MAP: src_enum[val] = FLOAT_MAP[name]
-        elif re.match(r'^[A-Z][A-Z0-9_]*$', name): src_enum[val] = name
-      break
-
-  # parse opcode tables
-  enums: dict[str, dict[int, str]] = {}
-  for m in re.finditer(r'Table \d+\. (\w+) Opcodes(.*?)(?=Table \d+\.|\n\d+\.\d+\.\d+\.\s+\w+\s*\nDescription|$)', full_text, re.S):
-    if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+([A-Z][A-Z0-9_]+)', m.group(2))}:
-      enums[m.group(1) + "Op"] = ops
-  if vopd_m := re.search(r'Table \d+\. VOPD Y-Opcodes\n(.*?)(?=Table \d+\.|15\.\d)', full_text, re.S):
-    if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+(V_DUAL_\w+)', vopd_m.group(1))}:
-      enums["VOPDOp"] = ops
-  enum_names = set(enums.keys())
-
-  def is_fields_table(t) -> bool: return t and len(t) > 1 and t[0] and 'Field' in str(t[0][0] or '')
-  def has_encoding(fields) -> bool: return any(f[0] == 'ENCODING' for f in fields)
-  def has_header_before_fields(text) -> bool:
-    return (pos := text.find('Field Name')) != -1 and bool(re.search(r'\d+\.\d+\.\d+\.\s+\w+\s*\n', text[:pos]))
-
-  # find format headers with their page indices
-  format_headers = []
-  for i, text in enumerate(page_texts):
-    for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n?Description', text): format_headers.append((m.group(1), i, m.start()))
-    for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n', text):
-      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 < len(page_texts):
-        next_text = page_texts[i + 1].lstrip()
-        if next_text.startswith('Description') or (next_text.startswith('"RDNA') and 'Description' in next_text[:200]):
-          format_headers.append((fmt_name, i, m.start()))
-
-  # parse instruction formats
-  formats: dict[str, list] = {}
-  for fmt_name, page_idx, header_pos in format_headers:
-    if fmt_name in formats: continue
-    text, tables = page_texts[page_idx], page_tables[page_idx]
-    field_pos = text.find('Field Name', header_pos)
-
-    fields = None
-    for offset in range(3):
-      if page_idx + offset >= len(pages): break
-      if offset > 0 and has_header_before_fields(page_texts[page_idx + offset]): break
-      for t in page_tables[page_idx + offset] if offset > 0 or field_pos > header_pos else []:
-        if is_fields_table(t) and (f := _parse_fields_table(t, fmt_name, enum_names)) and has_encoding(f):
-          fields = f
-          break
-      if fields: break
-
-    if not fields and field_pos > header_pos:
-      for t in tables:
-        if is_fields_table(t) and (f := _parse_fields_table(t, fmt_name, enum_names)):
-          fields = f
-          break
-
-    if not fields: continue
-    field_names = {f[0] for f in fields}
-
-    for pg_offset in range(1, 3):
-      if page_idx + pg_offset >= len(pages) or has_header_before_fields(page_texts[page_idx + pg_offset]): break
-      for t in page_tables[page_idx + pg_offset]:
-        if is_fields_table(t) and (extra := _parse_fields_table(t, fmt_name, enum_names)) and not has_encoding(extra):
-          for ef in extra:
-            if ef[0] not in field_names:
-              fields.append(ef)
-              field_names.add(ef[0])
-          break
-    formats[fmt_name] = fields
-
-  # fix known PDF errors - assert if already present (so we know when the bug is fixed)
-  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']]
-  # add missing opcodes not in PDF tables (RDNA3/RDNA3.5 specific)
-  if doc_name in ('RDNA3', 'RDNA3.5'):
-    if 'SOPPOp' in enums:
-      assert 8 not in enums['SOPPOp'], "S_WAITCNT_DEPCTR now in PDF, remove workaround"
-      enums['SOPPOp'][8] = 'S_WAITCNT_DEPCTR'
-    if 'DSOp' in enums:
-      gws_ops = {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'}
-      for k in gws_ops: assert k not in enums['DSOp'], f"{gws_ops[k]} now in PDF, remove workaround"
-      enums['DSOp'].update(gws_ops)
-    if 'FLATOp' in enums:
-      flat_ops = {40: 'GLOBAL_LOAD_ADDTID_B32', 41: 'GLOBAL_STORE_ADDTID_B32', 55: 'FLAT_ATOMIC_CSUB_U32'}
-      for k in flat_ops: assert k not in enums['FLATOp'], f"{flat_ops[k]} now in PDF, remove workaround"
-      enums['FLATOp'].update(flat_ops)
-
-  return {"formats": formats, "enums": enums, "src_enum": src_enum, "doc_name": doc_name, "is_cdna": is_cdna}
-
-def _merge_results(results: list[dict]) -> dict:
-  """Merge multiple PDF parse results into a superset. Asserts if any conflicts."""
-  merged = {"formats": {}, "enums": {}, "src_enum": dict(SRC_EXTRAS), "doc_names": [], "is_cdna": False}
-  for r in results:
-    merged["doc_names"].append(r["doc_name"])
-    merged["is_cdna"] = merged["is_cdna"] or r["is_cdna"]
-    # Merge src_enum (union, assert no conflicts)
-    for val, name in r["src_enum"].items():
-      if val in merged["src_enum"]:
-        assert merged["src_enum"][val] == name, f"SrcEnum conflict: {val} = {merged['src_enum'][val]} vs {name}"
-      else:
-        merged["src_enum"][val] = name
-    # Merge enums (union of ops per enum, assert no conflicts)
-    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, f"{enum_name} conflict: {val} = {merged['enums'][enum_name][val]} vs {name}"
-        else:
-          merged["enums"][enum_name][val] = name
-    # Merge formats (union of fields, assert no bit position conflicts for same field 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]}  # name -> (hi, lo)
-        for f in fields:
-          name, hi, lo = f[0], f[1], f[2]
-          if name in existing:
-            assert existing[name] == (hi, lo), f"Format {fmt_name} field {name} conflict: bits {existing[name]} vs ({hi}, {lo})"
-          else:
-            merged["formats"][fmt_name].append(f)
-  return merged
-
-def generate(output_path: str | None = None, arch: str = "rdna3") -> dict:
-  """Generate instruction definitions from AMD ISA PDF(s). Returns dict with formats for testing."""
-  urls = PDF_URLS[arch]
-  if isinstance(urls, str): urls = [urls]
-
-  # Parse all PDFs and merge
-  results = [_parse_single_pdf(url) for url in urls]
-  if len(results) == 1:
-    merged = results[0]
-    doc_name = merged["doc_name"]
-  else:
-    merged = _merge_results(results)
-    doc_name = "+".join(merged["doc_names"])
-
-  formats, enums, src_enum = merged["formats"], merged["enums"], merged["src_enum"]
-
-  # generate output
-  def enum_lines(name, items):
-    return [f"class {name}(IntEnum):"] + [f"  {n} = {v}" for v, n in sorted(items.items())] + [""]
-  def field_key(f): return order.index(f[0].lower()) if f[0].lower() in order else 1000
-  lines = [f"# autogenerated from AMD {doc_name} ISA PDF by dsl.py - do not edit", "from enum import IntEnum",
-           "from typing import Annotated",
-           "from extra.assembly.amd.dsl import bits, BitField, Inst32, Inst64, SGPR, VGPR, TTMP as TTMP, s as s, v as v, ttmp as ttmp, SSrc, Src, SImm, Imm, VDSTYEnc, SGPRField, VGPRField",
-           "import functools", ""]
-  lines += enum_lines("SrcEnum", src_enum) + sum([enum_lines(n, ops) for n, ops in sorted(enums.items())], [])
-  # Format-specific field defaults (verified against LLVM test vectors)
-  format_defaults = {'VOP3P': {'opsel_hi': 3, 'opsel_hi2': 1}}
-  lines.append("# instruction formats")
-  for fmt_name, fields in sorted(formats.items()):
-    base = "Inst64" if max(f[1] for f in fields) > 31 or fmt_name == 'VOP3SD' else "Inst32"
-    order = FIELD_ORDER.get(fmt_name, [])
-    lines.append(f"class {fmt_name}({base}):")
-    if enc := next((f for f in fields if f[0] == 'ENCODING'), None):
-      enc_str = f"bits[{enc[1]}:{enc[2]}] == 0b{enc[3]:b}" if enc[1] != enc[2] else f"bits[{enc[1]}] == {enc[3]}"
-      lines.append(f"  encoding = {enc_str}")
-    if defaults := format_defaults.get(fmt_name):
-      lines.append(f"  _defaults = {defaults}")
-    for name, hi, lo, _, ftype in sorted([f for f in fields if f[0] != 'ENCODING'], key=field_key):
-      if ftype and ftype.endswith('Op'):
-        ann = f":Annotated[BitField, {ftype}]"
-      else:
-        ann = f":{ftype}" if ftype else ""
-      lines.append(f"  {name.lower()}{ann} = bits[{hi}]" if hi == lo else f"  {name.lower()}{ann} = bits[{hi}:{lo}]")
-    lines.append("")
-  lines.append("# instruction helpers")
-  for cls_name, ops in sorted(enums.items()):
-    fmt = cls_name[:-2]
-    for op_val, name in sorted(ops.items()):
-      seg = {"GLOBAL": ", seg=2", "SCRATCH": ", seg=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"):
-        if fmt in ("VOP1", "VOP2", "VOPC"):
-          suffix = "_e32"
-        elif fmt == "VOP3" and op_val < 512:
-          suffix = "_e64"
-        else:
-          suffix = ""
-        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})")
-  skip_exports = {'DPP8', 'DPP16'}
-  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 skip_exports]
-  if "NULL" in src_names: lines.append("OFF = NULL\n")
-
-  if output_path is not None:
-    import pathlib
-    pathlib.Path(output_path).write_text('\n'.join(lines))
-  return {"formats": formats, "enums": enums, "src_enum": src_enum}
-
-if __name__ == "__main__":
-  import argparse
-  parser = argparse.ArgumentParser(description="Generate instruction definitions from AMD ISA PDF")
-  parser.add_argument("--arch", choices=list(PDF_URLS.keys()) + ["all"], default="rdna3", help="Target architecture (default: rdna3)")
-  args = parser.parse_args()
-  if args.arch == "all":
-    for arch in PDF_URLS.keys():
-      result = generate(f"extra/assembly/amd/autogen/{arch}/__init__.py", arch=arch)
-      print(f"{arch}: generated SrcEnum ({len(result['src_enum'])}) + {len(result['enums'])} opcode enums + {len(result['formats'])} format classes")
-  else:
-    result = generate(f"extra/assembly/amd/autogen/{args.arch}/__init__.py", arch=args.arch)
-    print(f"generated SrcEnum ({len(result['src_enum'])}) + {len(result['enums'])} opcode enums + {len(result['formats'])} format classes")
+class Inst96(Inst): pass

extra/assembly/amd/emu.py

@@ -1,63 +1,39 @@
 # RDNA3 emulator - executes compiled pseudocode from AMD ISA PDF
 # mypy: ignore-errors
 from __future__ import annotations
-import ctypes, os
-from extra.assembly.amd.dsl import Inst, RawImm
+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.pcode import _f32, _i32, _sext, _f16, _i16, _f64, _i64
 from extra.assembly.amd.autogen.rdna3.gen_pcode import get_compiled_functions
-from extra.assembly.amd.autogen.rdna3 import (
-  SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, VOPD, SrcEnum,
-  SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, DSOp, FLATOp, GLOBALOp, VOPDOp
-)
+from extra.assembly.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
 
-# VOP3 ops that use 64-bit operands (and thus 64-bit literals when src is 255)
-# Exception: V_LDEXP_F64 has 32-bit integer src1, so literal should NOT be 64-bit when src1=255
-_VOP3_64BIT_OPS = {op.value for op in VOP3Op if op.name.endswith(('_F64', '_B64', '_I64', '_U64'))}
-_VOPC_64BIT_OPS = {op.value for op in VOPCOp if op.name.endswith(('_F64', '_B64', '_I64', '_U64'))}
-# Ops where src1 is 32-bit (exponent/shift amount) even though the op name suggests 64-bit
-_VOP3_64BIT_OPS_32BIT_SRC1 = {VOP3Op.V_LDEXP_F64.value}
-# Ops with 16-bit types in name (for source/dest handling)
-# Exception: SAD/MSAD ops take 32-bit packed sources and extract 16-bit/8-bit chunks internally
-_VOP3_16BIT_OPS = {op for op in VOP3Op if any(s in op.name for s in ('_F16', '_B16', '_I16', '_U16')) and 'SAD' not in op.name}
-_VOP1_16BIT_OPS = {op for op in VOP1Op if any(s in op.name for s in ('_F16', '_B16', '_I16', '_U16'))}
-_VOP2_16BIT_OPS = {op for op in VOP2Op if any(s in op.name for s in ('_F16', '_B16', '_I16', '_U16'))}
-_VOPC_16BIT_OPS = {op for op in VOPCOp if any(s in op.name for s in ('_F16', '_B16', '_I16', '_U16'))}
-# CVT ops with 32/64-bit source (despite 16-bit in name)
-_CVT_32_64_SRC_OPS = {op for op in VOP3Op if op.name.startswith('V_CVT_') and op.name.endswith(('_F32', '_I32', '_U32', '_F64', '_I64', '_U64'))} | \
-                     {op for op in VOP1Op if op.name.startswith('V_CVT_') and op.name.endswith(('_F32', '_I32', '_U32', '_F64', '_I64', '_U64'))}
-# CVT ops with 32-bit destination (convert FROM 16-bit TO 32-bit): V_CVT_F32_F16, V_CVT_I32_I16, V_CVT_U32_U16
-_CVT_32_DST_OPS = {op for op in VOP3Op if op.name.startswith('V_CVT_') and any(s in op.name for s in ('F32_F16', 'I32_I16', 'U32_U16', 'I32_F16', 'U32_F16'))} | \
-                  {op for op in VOP1Op if op.name.startswith('V_CVT_') and any(s in op.name for s in ('F32_F16', 'I32_I16', 'U32_U16', 'I32_F16', 'U32_F16'))}
-# 16-bit dst ops (PACK has 32-bit dst despite F16 in name, CVT to 32-bit has 32-bit dst)
-_VOP3_16BIT_DST_OPS = {op for op in _VOP3_16BIT_OPS if 'PACK' not in op.name} - _CVT_32_DST_OPS
-_VOP1_16BIT_DST_OPS = {op for op in _VOP1_16BIT_OPS if 'PACK' not in op.name} - _CVT_32_DST_OPS
-# VOP1 16-bit source ops (excluding CVT ops with 32/64-bit source) - for VOP1 e32, .h encoded in register index
-_VOP1_16BIT_SRC_OPS = _VOP1_16BIT_OPS - _CVT_32_64_SRC_OPS
-
 # Inline constants for src operands 128-254. Build tables for f32, f16, and f64 formats.
-import struct as _struct
-_FLOAT_CONSTS = {SrcEnum.POS_HALF: 0.5, SrcEnum.NEG_HALF: -0.5, SrcEnum.POS_ONE: 1.0, SrcEnum.NEG_ONE: -1.0,
-                 SrcEnum.POS_TWO: 2.0, SrcEnum.NEG_TWO: -2.0, SrcEnum.POS_FOUR: 4.0, SrcEnum.NEG_FOUR: -4.0, SrcEnum.INV_2PI: 0.15915494309189535}
-def _build_inline_consts(neg_mask, float_to_bits):
-  tbl = list(range(65)) + [((-i) & neg_mask) for i in range(1, 17)] + [0] * (127 - 81)
-  for k, v in _FLOAT_CONSTS.items(): tbl[k - 128] = float_to_bits(v)
+_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(0xffffffff, lambda f: _struct.unpack('<I', _struct.pack('<f', f))[0])
-_INLINE_CONSTS_F16 = _build_inline_consts(0xffff, lambda f: _struct.unpack('<H', _struct.pack('<e', f))[0])
-_INLINE_CONSTS_F64 = _build_inline_consts(0xffffffffffffffff, lambda f: _struct.unpack('<Q', _struct.pack('<d', f))[0])
+_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:
-  for s, z in _valid_mem_ranges:
-    if s <= addr and addr + size <= s + z: return True
-  return not _valid_mem_ranges
+  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:
@@ -77,6 +53,9 @@ 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)
@@ -106,38 +85,29 @@ class WaveState:
   @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 & 0xffffffff, (v >> 32) & 0xffffffff
+  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 & 0xffffffff, (v >> 32) & 0xffffffff
+  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 & 0xffffffff
+    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 & 0xffffffff); self.wsgpr(i+1, (v >> 32) & 0xffffffff)
+  def wsgpr64(self, i: int, v: int): self.wsgpr(i, v & MASK32); self.wsgpr(i+1, (v >> 32) & MASK32)
 
-  def rsrc(self, v: int, lane: int) -> int:
+  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 _INLINE_CONSTS[v - 128]
+    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_f16(self, v: int, lane: int) -> int:
-    """Read source operand for VOP3P packed f16 operations. Uses f16 inline constants."""
-    if v < SGPR_COUNT: return self.sgpr[v]
-    if v == SCC: return self.scc
-    if v < 255: return _INLINE_CONSTS_F16[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:
-    """Read 64-bit source operand. For inline constants, returns 64-bit representation."""
-    # Inline constants 128-254 need special handling for 64-bit ops
     if 128 <= v < 255: return _INLINE_CONSTS_F64[v - 128]
-    if v == 255: return self.literal  # 32-bit literal, caller handles extension
+    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):
@@ -148,9 +118,6 @@ class WaveState:
     self._pend_sgpr.clear()
 
 
-
-def _unwrap(v) -> int: return v.val if isinstance(v, RawImm) else v.value if hasattr(v, 'value') else v
-
 def decode_program(data: bytes) -> Program:
   result: Program = {}
   i = 0
@@ -161,24 +128,8 @@ def decode_program(data: bytes) -> Program:
     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(): setattr(inst, name, _unwrap(val))
-    # from_bytes already handles literal reading - only need fallback for cases it doesn't handle
-    if inst._literal is None:
-      has_literal = any(getattr(inst, fld, None) == 255 for fld in ('src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'srcx0', 'srcy0'))
-      if inst_class == VOP2 and inst.op in (44, 45, 55, 56): has_literal = True
-      if inst_class == VOPD and (inst.opx in (1, 2) or inst.opy in (1, 2)): has_literal = True
-      if inst_class == SOP2 and inst.op in (69, 70): has_literal = True
-      if has_literal:
-        # For 64-bit ops, the 32-bit literal is placed in HIGH 32 bits (low 32 bits = 0)
-        # Exception: some ops have mixed src sizes (e.g., V_LDEXP_F64 has 32-bit src1)
-        op_val = inst._values.get('op')
-        if hasattr(op_val, 'value'): op_val = op_val.value
-        is_64bit = (inst_class is VOP3 and op_val in _VOP3_64BIT_OPS) or (inst_class is VOPC and op_val in _VOPC_64BIT_OPS)
-        # Don't treat literal as 64-bit if the op has 32-bit src1 and src1 is the literal
-        if is_64bit and op_val in _VOP3_64BIT_OPS_32BIT_SRC1 and getattr(inst, 'src1', None) == 255:
-          is_64bit = False
-        lit32 = int.from_bytes(data[i+base_size:i+base_size+4], 'little')
-        inst._literal = (lit32 << 32) if is_64bit else lit32
+    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
@@ -191,84 +142,74 @@ def decode_program(data: bytes) -> Program:
 def exec_scalar(st: WaveState, inst: Inst) -> int:
   """Execute scalar instruction. Returns PC delta or negative for special cases."""
   compiled = _get_compiled()
-  inst_type = type(inst)
 
   # SOPP: special cases for control flow that has no pseudocode
-  if inst_type is SOPP:
-    op = inst.op
-    if op == SOPPOp.S_ENDPGM: return -1
-    if op == SOPPOp.S_BARRIER: return -2
+  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 inst_type is SMEM:
+  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) & 0xffffffffffffffff, 4))
+    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 inst_type is SOP1: op_cls, ssrc0, sdst = SOP1Op, inst.ssrc0, inst.sdst
-  elif inst_type is SOP2: op_cls, ssrc0, sdst = SOP2Op, inst.ssrc0, inst.sdst
-  elif inst_type is SOPC: op_cls, ssrc0, sdst = SOPCOp, inst.ssrc0, None
-  elif inst_type is SOPK: op_cls, ssrc0, sdst = SOPKOp, inst.sdst, inst.sdst  # sdst is both src and dst
-  elif inst_type is SOPP: op_cls, ssrc0, sdst = SOPPOp, None, None
-  else: raise NotImplementedError(f"Unknown scalar type {inst_type}")
+  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 = op_cls(inst.op)
+  try: op = inst.op
   except ValueError:
-    if inst_type is SOPP: return 0
+    if isinstance(inst, SOPP): return 0
     raise
-  fn = compiled.get(op_cls, {}).get(op)
+  fn = compiled.get(type(op), {}).get(op)
   if fn is None:
     # SOPP instructions without pseudocode (waits, hints, nops) are no-ops
-    if inst_type is SOPP: return 0
+    if isinstance(inst, SOPP): return 0
     raise NotImplementedError(f"{op.name} not in pseudocode")
 
-  # Build context - handle 64-bit ops that need 64-bit source reads
-  # 64-bit source ops: name ends with _B64, _I64, _U64 or contains _U64, _I64 before last underscore
-  is_64bit_s0 = op.name.endswith(('_B64', '_I64', '_U64')) or '_U64_' in op.name or '_I64_' in op.name
-  is_64bit_s0s1 = op_cls is SOPCOp and op in (SOPCOp.S_CMP_EQ_U64, SOPCOp.S_CMP_LG_U64)
-  s0 = st.rsrc64(ssrc0, 0) if is_64bit_s0 or is_64bit_s0s1 else (st.rsrc(ssrc0, 0) if inst_type not in (SOPK, SOPP) else (st.rsgpr(inst.sdst) if inst_type is SOPK else 0))
-  is_64bit_sop2 = is_64bit_s0 and inst_type is SOP2
-  s1 = st.rsrc64(inst.ssrc1, 0) if (is_64bit_sop2 or is_64bit_s0s1) else (st.rsrc(inst.ssrc1, 0) if inst_type in (SOP2, SOPC) else inst.simm16 if inst_type is SOPK else 0)
-  d0 = st.rsgpr64(sdst) if (is_64bit_s0 or is_64bit_s0s1) and sdst is not None else (st.rsgpr(sdst) if sdst is not None else 0)
-  exec_mask = st.exec_mask
-  literal = inst.simm16 if inst_type in (SOPK, SOPP) else st.literal
+  # 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
 
-  # Execute compiled function - pass PC in bytes for instructions that need it
-  pc_bytes = st.pc * 4
-  result = fn(s0, s1, 0, d0, st.scc, st.vcc, 0, exec_mask, literal, None, {}, pc=pc_bytes)
+  # 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
-  if sdst is not None:
-    if result.get('d0_64'):
-      st.wsgpr64(sdst, result['d0'])
-    else:
-      st.wsgpr(sdst, result['d0'])
-  if 'scc' in result: st.scc = result['scc']
-  if 'exec' in result: st.exec_mask = result['exec']
-  if 'new_pc' in result:
+  # 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
-    # new_pc is where we want to go, st.pc is current position, inst._words will be added after
-    new_pc_words = result['new_pc'] // 4
+    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()
-  inst_type, V = type(inst), st.vgpr[lane]
+  V = st.vgpr[lane]
 
   # Memory ops (not ALU pseudocode)
-  if inst_type is FLAT:
+  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) & 0xffffffffffffffff if saddr not in (NULL, 0x7f) else (addr + offset) & 0xffffffffffffffff
+    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) & 0xffffffff if sign else val
+      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))
@@ -276,406 +217,185 @@ def exec_vector(st: WaveState, inst: Inst, lane: int, lds: bytearray | None = No
       sz, sign, hi = FLAT_D16_LOAD[op]
       val = mem_read(addr, sz)
       if sign: val = _sext(val, sz * 8) & 0xffff
-      if hi: V[vdst] = (V[vdst] & 0xffff) | (val << 16)  # upper 16 bits
-      else: V[vdst] = (V[vdst] & 0xffff0000) | (val & 0xffff)  # lower 16 bits
+      V[vdst] = _dst16(V[vdst], val, hi)
     elif op in FLAT_D16_STORE:
       sz, hi = FLAT_D16_STORE[op]
-      val = (V[data_reg] >> 16) & 0xffff if hi else V[data_reg] & 0xffff
-      mem_write(addr, sz, val & ((1 << (sz * 8)) - 1))
+      mem_write(addr, sz, _src16(V[data_reg], hi) & ((1 << (sz * 8)) - 1))
     else: raise NotImplementedError(f"FLAT op {op}")
     return
 
-  if inst_type is DS:
-    op, addr, vdst = inst.op, (V[inst.addr] + inst.offset0) & 0xffff, inst.vdst
+  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[addr+i*sz:addr+i*sz+sz], 'little'); V[vdst + i] = _sext(val, sz * 8) & 0xffffffff if sign else val
+      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[addr+i*sz:addr+i*sz+sz] = (V[inst.data0 + i] & ((1 << (sz * 8)) - 1)).to_bytes(sz, 'little')
+      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 using VOP2/VOP3 compiled functions
-  # Both ops execute simultaneously using pre-instruction values, so read all inputs first
-  if inst_type is VOPD:
+  # VOPD: dual-issue, execute two ops simultaneously (read all inputs before writes)
+  if isinstance(inst, VOPD):
     vdsty = (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1)
-    # Read all source operands BEFORE any writes (dual-issue semantics)
-    sx0, sx1 = st.rsrc(inst.srcx0, lane), V[inst.vsrcx1]
-    sy0, sy1 = st.rsrc(inst.srcy0, lane), V[inst.vsrcy1]
-    dx0, dy0 = V[inst.vdstx], V[vdsty]
-    # Execute X op
-    res_x = None
-    if (op_x := _VOPD_TO_VOP.get(inst.opx)):
-      if (fn_x := compiled.get(type(op_x), {}).get(op_x)):
-        res_x = fn_x(sx0, sx1, 0, dx0, st.scc, st.vcc, lane, st.exec_mask, st.literal, None, {})
-    # Execute Y op
-    res_y = None
-    if (op_y := _VOPD_TO_VOP.get(inst.opy)):
-      if (fn_y := compiled.get(type(op_y), {}).get(op_y)):
-        res_y = fn_y(sy0, sy1, 0, dy0, st.scc, st.vcc, lane, st.exec_mask, st.literal, None, {})
-    # Write results after both ops complete
-    if res_x is not None: V[inst.vdstx] = res_x['d0']
-    if res_y is not None: V[vdsty] = res_y['d0']
+    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 inst_type is VOP3SD:
-    op = VOP3SDOp(inst.op)
-    fn = compiled.get(VOP3SDOp, {}).get(op)
-    if fn is None: raise NotImplementedError(f"{op.name} not in pseudocode")
-    # VOP3SD has both 32-bit ops (V_ADD_CO_CI_U32, etc.) and 64-bit ops (V_DIV_SCALE_F64, V_MAD_U64_U32, etc.)
-    div_scale_64_ops = (VOP3SDOp.V_DIV_SCALE_F64,)
-    mad64_ops = (VOP3SDOp.V_MAD_U64_U32, VOP3SDOp.V_MAD_I64_I32)
-    if op in div_scale_64_ops:
-      # V_DIV_SCALE_F64: all sources are 64-bit
-      s0, s1, s2 = st.rsrc64(inst.src0, lane), st.rsrc64(inst.src1, lane), st.rsrc64(inst.src2, lane)
-    elif op in mad64_ops:
-      # V_MAD_U64_U32, V_MAD_I64_I32: src0/src1 are 32-bit, src2 is 64-bit
-      s0, s1 = st.rsrc(inst.src0, lane), st.rsrc(inst.src1, lane)
-      if inst.src2 >= 256:  # VGPR
-        s2 = V[inst.src2 - 256] | (V[inst.src2 - 256 + 1] << 32)
-      else:  # SGPR - read 64-bit from consecutive SGPRs
-        s2 = st.rsgpr64(inst.src2)
-    else:
-      # Default: 32-bit sources
-      s0, s1, s2 = st.rsrc(inst.src0, lane), st.rsrc(inst.src1, lane), st.rsrc(inst.src2, lane)
-    d0 = V[inst.vdst]
-    # For carry-in operations (V_*_CO_CI_*), src2 register contains the carry bitmask (not VCC).
-    # The pseudocode uses VCC but in VOP3SD encoding, the actual carry source is inst.src2.
-    # We pass the src2 register value as 'vcc' to the interpreter so it reads the correct carry.
-    carry_ops = (VOP3SDOp.V_ADD_CO_CI_U32, VOP3SDOp.V_SUB_CO_CI_U32, VOP3SDOp.V_SUBREV_CO_CI_U32)
-    vcc_for_exec = st.rsgpr64(inst.src2) if op in carry_ops else st.vcc
-    result = fn(s0, s1, s2, d0, st.scc, vcc_for_exec, lane, st.exec_mask, st.literal, None, {})
-    # Write result - handle 64-bit destinations
-    if result.get('d0_64'):
-      V[inst.vdst] = result['d0'] & 0xffffffff
-      V[inst.vdst + 1] = (result['d0'] >> 32) & 0xffffffff
-    else:
-      V[inst.vdst] = result['d0'] & 0xffffffff
-    if result.get('vcc_lane') is not None:
-      st.pend_sgpr_lane(inst.sdst, lane, result['vcc_lane'])
+  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)
-  # vop1_dst_hi/vop2_dst_hi: for VOP1/VOP2 16-bit dst ops, bit 7 of vdst indicates .h (high 16-bit) destination
-  vop1_dst_hi, vop2_dst_hi = False, False
-  if inst_type is VOP1:
+  # 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
-    op_cls, op, src0, src1, src2 = VOP1Op, VOP1Op(inst.op), inst.src0, None, None
-    # For 16-bit dst ops, vdst encodes .h in bit 7
-    if op in _VOP1_16BIT_DST_OPS:
-      vop1_dst_hi = (inst.vdst & 0x80) != 0
-      vdst = inst.vdst & 0x7f
-    else:
-      vdst = inst.vdst
-  elif inst_type is VOP2:
-    op_cls, op, src0, src1, src2 = VOP2Op, VOP2Op(inst.op), inst.src0, inst.vsrc1 + 256, None
-    # For 16-bit dst ops, vdst encodes .h in bit 7
-    if op in _VOP2_16BIT_OPS:
-      vop2_dst_hi = (inst.vdst & 0x80) != 0
-      vdst = inst.vdst & 0x7f
-    else:
-      vdst = inst.vdst
-  elif inst_type is VOP3:
-    # VOP3 ops 0-255 are VOPC comparisons encoded as VOP3 (use VOPCOp pseudocode)
-    if inst.op < 256:
-      op_cls, op, src0, src1, src2, vdst = VOPCOp, VOPCOp(inst.op), inst.src0, inst.src1, None, inst.vdst
-    else:
-      op_cls, op, src0, src1, src2, vdst = VOP3Op, VOP3Op(inst.op), inst.src0, inst.src1, inst.src2, inst.vdst
-  elif inst_type is VOPC:
-    op = VOPCOp(inst.op)
+    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
-    src1 = inst.vsrc1 + 256  # convert to standard VGPR encoding (256 + vgpr_idx)
-    op_cls, src0, src2, vdst = VOPCOp, inst.src0, None, VCC_LO
-  elif inst_type is VOP3P:
+    src0, src1, src2, vdst = inst.src0, inst.vsrc1 + 256, None, VCC_LO
+  elif isinstance(inst, VOP3P):
     # VOP3P: Packed 16-bit operations using compiled functions
-    op = VOP3POp(inst.op)
     # WMMA: wave-level matrix multiply-accumulate (special handling - needs cross-lane access)
-    if op in (VOP3POp.V_WMMA_F32_16X16X16_F16, VOP3POp.V_WMMA_F32_16X16X16_BF16, VOP3POp.V_WMMA_F16_16X16X16_F16):
+    if 'WMMA' in inst.op_name:
       if lane == 0:  # Only execute once per wave, write results for all lanes
-        exec_wmma(st, inst, op)
+        exec_wmma(st, inst, inst.op)
       return
-    # V_FMA_MIX: Mixed precision FMA - inputs can be f16 or f32 controlled by opsel_hi/opsel_hi2
-    # opsel_hi[0]: src0 is f32 (0) or f16 from hi bits (1)
-    # opsel_hi[1]: src1 is f32 (0) or f16 from hi bits (1)
-    # opsel_hi2: src2 is f32 (0) or f16 from hi bits (1)
-    # opsel[i]: when source is f16, use lo (0) or hi (1) 16 bits - BUT for V_FMA_MIX, opsel selects lo/hi when opsel_hi=1
-    # neg_hi[i]: abs modifier for source i (reuses neg_hi field for abs in V_FMA_MIX)
-    if op in (VOP3POp.V_FMA_MIX_F32, VOP3POp.V_FMA_MIXLO_F16, VOP3POp.V_FMA_MIXHI_F16):
-      opsel = getattr(inst, 'opsel', 0)
-      opsel_hi = getattr(inst, 'opsel_hi', 0)
-      opsel_hi2 = getattr(inst, 'opsel_hi2', 0)
-      neg = getattr(inst, 'neg', 0)
-      abs_ = getattr(inst, 'neg_hi', 0)  # neg_hi field is reused as abs for V_FMA_MIX
-      vdst = inst.vdst
-      # Read raw 32-bit values
-      s0_raw = st.rsrc(inst.src0, lane)
-      s1_raw = st.rsrc(inst.src1, lane)
-      s2_raw = st.rsrc(inst.src2, lane) if inst.src2 is not None else 0
-      # Decode sources based on opsel_hi (controls f32 vs f16) and opsel (controls which half for f16)
-      # src0: opsel_hi[0]=1 means f16, opsel[0] selects hi(1) or lo(0) half
-      if opsel_hi & 1:
-        s0 = _f16((s0_raw >> 16) & 0xffff) if (opsel & 1) else _f16(s0_raw & 0xffff)
-      else:
-        s0 = _f32(s0_raw)
-      # src1: opsel_hi[1]=1 means f16, opsel[1] selects hi(1) or lo(0) half
-      if opsel_hi & 2:
-        s1 = _f16((s1_raw >> 16) & 0xffff) if (opsel & 2) else _f16(s1_raw & 0xffff)
-      else:
-        s1 = _f32(s1_raw)
-      # src2: opsel_hi2=1 means f16, opsel[2] selects hi(1) or lo(0) half
-      if opsel_hi2:
-        s2 = _f16((s2_raw >> 16) & 0xffff) if (opsel & 4) else _f16(s2_raw & 0xffff)
-      else:
-        s2 = _f32(s2_raw)
-      # Apply abs modifiers (abs_ field reuses neg_hi position)
-      if abs_ & 1: s0 = abs(s0)
-      if abs_ & 2: s1 = abs(s1)
-      if abs_ & 4: s2 = abs(s2)
-      # Apply neg modifiers
-      if neg & 1: s0 = -s0
-      if neg & 2: s1 = -s1
-      if neg & 4: s2 = -s2
-      # Compute FMA: d = s0 * s1 + s2
-      result = s0 * s1 + s2
-      V = st.vgpr[lane]
-      if op == VOP3POp.V_FMA_MIX_F32:
-        V[vdst] = _i32(result)
-      elif op == VOP3POp.V_FMA_MIXLO_F16:
-        lo = _i16(result) & 0xffff
-        V[vdst] = (V[vdst] & 0xffff0000) | lo
-      else:  # V_FMA_MIXHI_F16
-        hi = _i16(result) & 0xffff
-        V[vdst] = (V[vdst] & 0x0000ffff) | (hi << 16)
+    # 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
-    # Use rsrc_f16 for VOP3P to get correct f16 inline constants
-    s0_raw = st.rsrc_f16(inst.src0, lane)
-    s1_raw = st.rsrc_f16(inst.src1, lane)
-    s2_raw = st.rsrc_f16(inst.src2, lane) if inst.src2 is not None else 0
-    # Handle opsel (which 16-bit halves to use for each source)
-    opsel = getattr(inst, 'opsel', 0)
-    opsel_hi = getattr(inst, 'opsel_hi', 3)  # Default: use hi for hi result
-    opsel_hi2 = getattr(inst, 'opsel_hi2', 1)  # Default for src2
-    # Handle neg modifiers for VOP3P
-    # neg applies to lo result inputs, neg_hi applies to hi result inputs
-    neg = getattr(inst, 'neg', 0)
-    neg_hi = getattr(inst, 'neg_hi', 0)
-    # Build "virtual" sources with halves arranged for pseudocode: lo half goes to [15:0], hi half goes to [31:16]
-    # opsel bit 0/1/2 selects which half of src0/1/2 goes to the LO result
-    # opsel_hi bit 0/1 selects which half of src0/1 goes to the HI result
-    s0_lo = (s0_raw >> 16) & 0xffff if (opsel & 1) else s0_raw & 0xffff
-    s1_lo = (s1_raw >> 16) & 0xffff if (opsel & 2) else s1_raw & 0xffff
-    s2_lo = (s2_raw >> 16) & 0xffff if (opsel & 4) else s2_raw & 0xffff
-    s0_hi = (s0_raw >> 16) & 0xffff if (opsel_hi & 1) else s0_raw & 0xffff
-    s1_hi = (s1_raw >> 16) & 0xffff if (opsel_hi & 2) else s1_raw & 0xffff
-    s2_hi = (s2_raw >> 16) & 0xffff if opsel_hi2 else s2_raw & 0xffff
-    # Apply neg to lo result inputs (toggle f16 sign bit)
-    if neg & 1: s0_lo ^= 0x8000
-    if neg & 2: s1_lo ^= 0x8000
-    if neg & 4: s2_lo ^= 0x8000
-    # Apply neg_hi to hi result inputs
-    if neg_hi & 1: s0_hi ^= 0x8000
-    if neg_hi & 2: s1_hi ^= 0x8000
-    if neg_hi & 4: s2_hi ^= 0x8000
-    # Pack into format expected by pseudocode: [31:16] = hi input, [15:0] = lo input
-    s0 = (s0_hi << 16) | s0_lo
-    s1 = (s1_hi << 16) | s1_lo
-    s2 = (s2_hi << 16) | s2_lo
-    op_cls, vdst = VOP3POp, inst.vdst
-    fn = compiled.get(op_cls, {}).get(op)
-    if fn is None: raise NotImplementedError(f"{op.name} not in pseudocode")
-    result = fn(s0, s1, s2, 0, st.scc, st.vcc, lane, st.exec_mask, st.literal, None, {})
-    st.vgpr[lane][vdst] = result['d0'] & 0xffffffff
+    # 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 {inst_type}")
+  else: raise NotImplementedError(f"Unknown vector type {type(inst)}")
 
-  fn = compiled.get(op_cls, {}).get(op)
-  if fn is None: raise NotImplementedError(f"{op.name} not in pseudocode")
+  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 inst_type is VOP3 else (0, 0)
-  opsel = getattr(inst, 'opsel', 0) if inst_type is VOP3 else 0
-  def mod_src(val: int, idx: int) -> int:
-    if (abs_ >> idx) & 1: val = _i32(abs(_f32(val)))
-    if (neg >> idx) & 1: val = _i32(-_f32(val))
-    return val
-  def mod_src64(val: int, idx: int) -> int:
-    if (abs_ >> idx) & 1: val = _i64(abs(_f64(val)))
-    if (neg >> idx) & 1: val = _i64(-_f64(val))
+  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
 
-  # Determine if sources are 64-bit based on instruction type
-  # For 64-bit shift ops: src0 is 32-bit (shift amount), src1 is 64-bit (value to shift)
-  # For most other _B64/_I64/_U64/_F64 ops: all sources are 64-bit
-  is_64bit_op = op.name.endswith(('_B64', '_I64', '_U64', '_F64'))
-  # V_LDEXP_F64, V_TRIG_PREOP_F64, V_CMP_CLASS_F64, V_CMPX_CLASS_F64: src0 is 64-bit, src1 is 32-bit
-  is_ldexp_64 = op in (VOP3Op.V_LDEXP_F64, VOP3Op.V_TRIG_PREOP_F64, VOP3Op.V_CMP_CLASS_F64, VOP3Op.V_CMPX_CLASS_F64,
-                       VOPCOp.V_CMP_CLASS_F64, VOPCOp.V_CMPX_CLASS_F64)
-  is_shift_64 = op in (VOP3Op.V_LSHLREV_B64, VOP3Op.V_LSHRREV_B64, VOP3Op.V_ASHRREV_I64)
-  # 16-bit source ops: use precomputed sets instead of string checks
-  # Note: must check op_cls to avoid cross-enum value collisions
-  is_16bit_src = op_cls is VOP3Op and op in _VOP3_16BIT_OPS and op not in _CVT_32_64_SRC_OPS
-  # VOP2 16-bit ops use f16 inline constants for src0 (vsrc1 is always a VGPR, no inline constants)
-  is_vop2_16bit = op_cls is VOP2Op and op in _VOP2_16BIT_OPS
+  # 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()
 
-  if is_shift_64:
-    s0 = mod_src(st.rsrc(src0, lane), 0)  # shift amount is 32-bit
-    s1 = st.rsrc64(src1, lane) if src1 is not None else 0  # value to shift is 64-bit
-    s2 = mod_src(st.rsrc(src2, lane), 2) if src2 is not None else 0
-  elif is_ldexp_64:
-    s0 = mod_src64(st.rsrc64(src0, lane), 0)  # mantissa is 64-bit float
-    # src1 is 32-bit int. For 64-bit ops (like V_CMP_CLASS_F64), the literal is stored shifted left by 32.
-    # For V_LDEXP_F64/V_TRIG_PREOP_F64, _is_64bit_op() returns False so literal is stored as-is.
-    s1_raw = st.rsrc(src1, lane) if src1 is not None else 0
-    # Only shift if src1 is literal AND this is a true 64-bit op (V_CMP_CLASS ops, not LDEXP/TRIG_PREOP)
-    is_class_op = op in (VOP3Op.V_CMP_CLASS_F64, VOP3Op.V_CMPX_CLASS_F64, VOPCOp.V_CMP_CLASS_F64, VOPCOp.V_CMPX_CLASS_F64)
-    s1 = mod_src((s1_raw >> 32) if src1 == 255 and is_class_op else s1_raw, 1)
-    s2 = mod_src(st.rsrc(src2, lane), 2) if src2 is not None else 0
-  elif is_64bit_op:
-    # 64-bit ops: apply neg/abs modifiers using f64 interpretation for float ops
-    s0 = mod_src64(st.rsrc64(src0, lane), 0)
-    s1 = mod_src64(st.rsrc64(src1, lane), 1) if src1 is not None else 0
-    s2 = mod_src64(st.rsrc64(src2, lane), 2) if src2 is not None else 0
-  elif is_16bit_src:
-    # For 16-bit source ops, opsel bits select which half to use
-    # Inline constants (128-254) must use f16 encoding, not f32
-    def rsrc_16bit(src, lane): return st.rsrc_f16(src, lane) if 128 <= src < 255 else st.rsrc(src, lane)
-    s0_raw = rsrc_16bit(src0, lane)
-    s1_raw = rsrc_16bit(src1, lane) if src1 is not None else 0
-    s2_raw = rsrc_16bit(src2, lane) if src2 is not None else 0
-    # opsel[0] selects hi(1) or lo(0) for src0, opsel[1] for src1, opsel[2] for src2
-    s0 = ((s0_raw >> 16) & 0xffff) if (opsel & 1) else (s0_raw & 0xffff)
-    s1 = ((s1_raw >> 16) & 0xffff) if (opsel & 2) else (s1_raw & 0xffff)
-    s2 = ((s2_raw >> 16) & 0xffff) if (opsel & 4) else (s2_raw & 0xffff)
-    # Apply abs/neg modifiers as f16 operations (toggle sign bit 15)
-    if abs_ & 1: s0 &= 0x7fff
-    if abs_ & 2: s1 &= 0x7fff
-    if abs_ & 4: s2 &= 0x7fff
-    if neg & 1: s0 ^= 0x8000
-    if neg & 2: s1 ^= 0x8000
-    if neg & 4: s2 ^= 0x8000
-  elif is_vop2_16bit:
-    # VOP2 16-bit ops: src0 uses f16 inline constants, or VGPR where v128+ = hi half of v0-v127
-    # RDNA3 encoding: for VGPRs, bit 7 of VGPR index (src0-256) selects hi(1) or lo(0) half
-    if src0 >= 256:  # VGPR
-      src0_hi = (src0 - 256) & 0x80 != 0
-      src0_masked = ((src0 - 256) & 0x7f) + 256  # mask out hi bit to get actual VGPR
-      s0_raw = mod_src(st.rsrc(src0_masked, lane), 0)
-      s0 = ((s0_raw >> 16) & 0xffff) if src0_hi else (s0_raw & 0xffff)
-    else:  # SGPR or inline constant
-      s0_raw = mod_src(st.rsrc_f16(src0, lane), 0)
-      s0 = s0_raw & 0xffff
-    # vsrc1: .h suffix encoded in bit 7 of VGPR index (src1 = 256 + vgpr_idx + 0x80 if hi)
-    if src1 is not None:
-      src1_hi = (src1 - 256) & 0x80 != 0
-      src1_masked = ((src1 - 256) & 0x7f) + 256
-      s1_raw = mod_src(st.rsrc(src1_masked, lane), 1)
-      s1 = ((s1_raw >> 16) & 0xffff) if src1_hi else (s1_raw & 0xffff)
-    else:
-      s1 = 0
-    s2 = mod_src(st.rsrc(src2, lane), 2) if src2 is not None else 0
-  elif op_cls is VOP1Op and op in _VOP1_16BIT_SRC_OPS:
-    # VOP1 16-bit source ops: .h encoded in bit 7 of VGPR index (src0 >= 384 means hi half)
-    # For VGPRs: src0 = 256 + vgpr_idx + (0x80 if hi else 0), so bit 7 of (src0-256) is the hi flag
-    src0_hi = src0 >= 256 and ((src0 - 256) & 0x80) != 0
-    src0_masked = ((src0 - 256) & 0x7f) + 256 if src0 >= 256 else src0  # mask out hi bit for VGPR
-    s0_raw = mod_src(st.rsrc(src0_masked, lane), 0)
-    s0 = ((s0_raw >> 16) & 0xffff) if src0_hi else (s0_raw & 0xffff)
-    s1, s2 = 0, 0
-  elif op_cls is VOPCOp and op in _VOPC_16BIT_OPS:
-    # VOPC 16-bit ops: src0 and vsrc1 use same encoding as VOP2 16-bit
-    # For VGPRs, bit 7 of VGPR index selects hi(1) or lo(0) half
-    if src0 >= 256:  # VGPR
-      src0_hi = (src0 - 256) & 0x80 != 0
-      src0_masked = ((src0 - 256) & 0x7f) + 256
-      s0_raw = mod_src(st.rsrc(src0_masked, lane), 0)
-      s0 = ((s0_raw >> 16) & 0xffff) if src0_hi else (s0_raw & 0xffff)
-    else:  # SGPR or inline constant
-      s0_raw = mod_src(st.rsrc_f16(src0, lane), 0)
-      s0 = s0_raw & 0xffff
-    # vsrc1: bit 7 of VGPR index selects hi(1) or lo(0) half
-    if src1 is not None:
-      if src1 >= 256:  # VGPR - use hi/lo encoding
-        src1_hi = (src1 - 256) & 0x80 != 0
-        src1_masked = ((src1 - 256) & 0x7f) + 256
-        s1_raw = mod_src(st.rsrc(src1_masked, lane), 1)
-        s1 = ((s1_raw >> 16) & 0xffff) if src1_hi else (s1_raw & 0xffff)
-      else:  # SGPR or inline constant - read as 32-bit, use low 16 bits
-        s1_raw = mod_src(st.rsrc(src1, lane), 1)
-        s1 = s1_raw & 0xffffffff  # V_CMP_CLASS uses full 32-bit mask
-    else:
-      s1 = 0
-    s2 = 0
-  else:
-    s0 = mod_src(st.rsrc(src0, lane), 0)
-    s1 = mod_src(st.rsrc(src1, lane), 1) if src1 is not None else 0
-    s2 = mod_src(st.rsrc(src2, lane), 2) if src2 is not None else 0
-  # For VOP2 16-bit ops (like V_FMAC_F16), the destination is used as an accumulator.
-  # The pseudocode reads D0.f16 from low 16 bits, so we need to shift hi->lo when vop2_dst_hi is True.
-  if is_vop2_16bit:
-    d0 = ((V[vdst] >> 16) & 0xffff) if vop2_dst_hi else (V[vdst] & 0xffff)
-  else:
-    d0 = V[vdst] if not is_64bit_op else (V[vdst] | (V[vdst + 1] << 32))
+  # 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 op in (VOP3Op.V_CNDMASK_B32, VOP3Op.V_CNDMASK_B16) and inst_type is VOP3 and src2 is not None and src2 < 256 else st.vcc
+  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(s0, s1, s2, d0, st.scc, vcc_for_fn, lane, st.exec_mask, st.literal, st.vgpr, {}, src0_idx, vdst)
+  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
+  # 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_lane' in result:
-    # VOP2 carry instructions (V_ADD_CO_CI_U32, V_SUB_CO_CI_U32, V_SUBREV_CO_CI_U32) write carry to VCC implicitly
-    # VOPC and VOP3-encoded VOPC write to vdst (which is VCC_LO for VOPC, inst.sdst for VOP3)
-    vcc_dst = VCC_LO if op_cls is VOP2Op and op in (VOP2Op.V_ADD_CO_CI_U32, VOP2Op.V_SUB_CO_CI_U32, VOP2Op.V_SUBREV_CO_CI_U32) else vdst
-    st.pend_sgpr_lane(vcc_dst, lane, result['vcc_lane'])
-  if 'exec_lane' in result:
-    # V_CMPX instructions write to EXEC per-lane
-    st.pend_sgpr_lane(EXEC_LO, lane, result['exec_lane'])
-  if 'd0' in result and op_cls not in (VOPCOp,) and 'vgpr_write' not in result:
-    # V_READFIRSTLANE_B32 and V_READLANE_B32 write to SGPR, not VGPR
-    # V_WRITELANE_B32 uses vgpr_write for cross-lane writes, don't overwrite with d0
-    writes_to_sgpr = op in (VOP1Op.V_READFIRSTLANE_B32,) or \
-                     (op_cls is VOP3Op and op in (VOP3Op.V_READFIRSTLANE_B32, VOP3Op.V_READLANE_B32))
-    # Check for 16-bit destination ops (opsel[3] controls hi/lo write)
-    # Must check op_cls to avoid cross-enum value collisions (e.g., VOP1Op.V_MOV_B32=1 vs VOP3Op.V_CMP_LT_F16=1)
-    is_16bit_dst = (op_cls is VOP3Op and op in _VOP3_16BIT_DST_OPS) or (op_cls is VOP1Op and op in _VOP1_16BIT_DST_OPS)
-    if writes_to_sgpr:
-      st.wsgpr(vdst, result['d0'] & 0xffffffff)
-    elif result.get('d0_64'):
-      V[vdst] = result['d0'] & 0xffffffff
-      V[vdst + 1] = (result['d0'] >> 32) & 0xffffffff
-    elif is_16bit_dst and inst_type is VOP3:
-      # VOP3 16-bit ops: opsel[3] (bit 3 of opsel field) controls hi/lo destination
-      if opsel & 8:  # opsel[3] = 1: write to high 16 bits
-        V[vdst] = (V[vdst] & 0x0000ffff) | ((result['d0'] & 0xffff) << 16)
-      else:  # opsel[3] = 0: write to low 16 bits
-        V[vdst] = (V[vdst] & 0xffff0000) | (result['d0'] & 0xffff)
-    elif is_16bit_dst and inst_type is VOP1:
-      # VOP1 16-bit ops: .h suffix encoded in bit 7 of vdst (extracted as vop1_dst_hi)
-      if vop1_dst_hi:  # .h: write to high 16 bits
-        V[vdst] = (V[vdst] & 0x0000ffff) | ((result['d0'] & 0xffff) << 16)
-      else:  # .l: write to low 16 bits
-        V[vdst] = (V[vdst] & 0xffff0000) | (result['d0'] & 0xffff)
-    elif is_vop2_16bit:
-      # VOP2 16-bit ops: .h suffix encoded in bit 7 of vdst (extracted as vop2_dst_hi)
-      if vop2_dst_hi:  # .h: write to high 16 bits
-        V[vdst] = (V[vdst] & 0x0000ffff) | ((result['d0'] & 0xffff) << 16)
-      else:  # .l: write to low 16 bits
-        V[vdst] = (V[vdst] & 0xffff0000) | (result['d0'] & 0xffff)
-    else:
-      V[vdst] = result['d0'] & 0xffffffff
+  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)
@@ -684,82 +404,41 @@ def exec_vector(st: WaveState, inst: Inst, lane: int, lds: bytearray | None = No
 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 matrix A (16x16 f16/bf16) from lanes 0-15, VGPRs src0 to src0+7 (2 f16 per VGPR = 16 values per lane)
-  # Layout: A[row][k] where row = lane (0-15), k comes from 8 VGPRs × 2 halves
-  mat_a = []
-  for lane in range(16):
-    for reg in range(8):
-      val = st.vgpr[lane][src0 - 256 + reg] if src0 >= 256 else st.rsgpr(src0 + reg)
-      mat_a.append(_f16(val & 0xffff))
-      mat_a.append(_f16((val >> 16) & 0xffff))
-  # Read matrix B (16x16 f16/bf16) - same layout, B[col][k] where col comes from lane
-  mat_b = []
-  for lane in range(16):
-    for reg in range(8):
-      val = st.vgpr[lane][src1 - 256 + reg] if src1 >= 256 else st.rsgpr(src1 + reg)
-      mat_b.append(_f16(val & 0xffff))
-      mat_b.append(_f16((val >> 16) & 0xffff))
-
+  # 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
-  # Layout: element i is at lane (i % 32), VGPR (i // 32) + src2
-  mat_c = []
-  for i in range(256):
-    lane, reg = i % 32, i // 32
-    val = st.vgpr[lane][src2 - 256 + reg] if src2 >= 256 else st.rsgpr(src2 + reg)
-    mat_c.append(_f32(val))
-
+  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 = [0.0] * 256
-  for row in range(16):
-    for col in range(16):
-      acc = 0.0
-      for k in range(16):
-        a_val = mat_a[row * 16 + k]
-        b_val = mat_b[col * 16 + k]
-        acc += a_val * b_val
-      mat_d[row * 16 + col] = acc + mat_c[row * 16 + col]
-
-  # Write result matrix D back - same layout as C
+  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:
-    # Output is f16, pack 2 values per VGPR
     for i in range(0, 256, 2):
-      lane, reg = (i // 2) % 32, (i // 2) // 32
-      lo = _i16(mat_d[i]) & 0xffff
-      hi = _i16(mat_d[i + 1]) & 0xffff
-      st.vgpr[lane][vdst + reg] = (hi << 16) | lo
+      st.vgpr[(i//2) % 32][vdst + (i//2)//32] = ((_i16(mat_d[i+1]) & 0xffff) << 16) | (_i16(mat_d[i]) & 0xffff)
   else:
-    # Output is f32
-    for i in range(256):
-      lane, reg = i % 32, i // 32
-      st.vgpr[lane][vdst + reg] = _i32(mat_d[i])
+    for i in range(256): st.vgpr[i % 32][vdst + i//32] = _i32(mat_d[i])
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # MAIN EXECUTION LOOP
 # ═══════════════════════════════════════════════════════════════════════════════
 
-SCALAR_TYPES = {SOP1, SOP2, SOPC, SOPK, SOPP, SMEM}
-VECTOR_TYPES = {VOP1, VOP2, VOP3, VOP3SD, VOPC, FLAT, DS, VOPD, VOP3P}
-
 def step_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int) -> int:
   inst = program.get(st.pc)
   if inst is None: return 1
-  inst_words, st.literal, inst_type = inst._words, getattr(inst, '_literal', None) or 0, type(inst)
+  inst_words, st.literal = inst._words, getattr(inst, '_literal', None) or 0
 
-  if inst_type in SCALAR_TYPES:
+  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_B32 and V_READLANE_B32 write to SGPR, so they should only execute once per wave (lane 0)
-    is_readlane = (inst_type is VOP1 and inst.op == VOP1Op.V_READFIRSTLANE_B32) or \
-                  (inst_type is VOP3 and inst.op in (VOP3Op.V_READFIRSTLANE_B32, VOP3Op.V_READLANE_B32))
-    if is_readlane:
-      exec_vector(st, inst, 0, lds)  # Execute once with lane 0
-    else:
-      exec_mask = st.exec_mask
-      for lane in range(n_lanes):
-        if exec_mask & (1 << lane): exec_vector(st, inst, lane, lds)
+    # 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
@@ -780,31 +459,24 @@ def exec_workgroup(program: Program, workgroup_id: tuple[int, int, int], local_s
     n_lanes, st = min(WAVE_SIZE, total_threads - wave_start), WaveState()
     st.exec_mask = (1 << n_lanes) - 1
     st.wsgpr64(0, args_ptr)
-    gx, gy, gz = workgroup_id
     # Set workgroup IDs in SGPRs based on USER_SGPR_COUNT and enable flags from COMPUTE_PGM_RSRC2
     sgpr_idx = wg_id_sgpr_base
-    if wg_id_enables[0]: st.sgpr[sgpr_idx] = gx; sgpr_idx += 1
-    if wg_id_enables[1]: st.sgpr[sgpr_idx] = gy; sgpr_idx += 1
-    if wg_id_enables[2]: st.sgpr[sgpr_idx] = gz
+    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 if local_size == (lx, 1, 1) else ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
+      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:
-  data = (ctypes.c_char * lib_sz).from_address(lib).raw
-  program = decode_program(data)
+  program = decode_program((ctypes.c_char * lib_sz).from_address(lib).raw)
   if not program: return -1
-  # Parse COMPUTE_PGM_RSRC2 for SGPR layout
-  user_sgpr_count = (rsrc2 >> 1) & 0x1f
-  enable_wg_id_x = bool((rsrc2 >> 7) & 1)
-  enable_wg_id_y = bool((rsrc2 >> 8) & 1)
-  enable_wg_id_z = bool((rsrc2 >> 9) & 1)
-  wg_id_enables = (enable_wg_id_x, enable_wg_id_y, enable_wg_id_z)
+  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, user_sgpr_count, wg_id_enables)
+      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

@@ -1,90 +1,42 @@
 # DSL for RDNA3 pseudocode - makes pseudocode expressions work directly as Python
-import struct, math, re
+import struct, math
+from extra.assembly.amd.dsl import MASK32, MASK64, _f32, _i32, _sext, _f16, _i16, _f64, _i64
 
 # ═══════════════════════════════════════════════════════════════════════════════
-# HELPER FUNCTIONS (previously in helpers.py)
+# HELPER FUNCTIONS
 # ═══════════════════════════════════════════════════════════════════════════════
 
-def _f32(i): return struct.unpack("<f", struct.pack("<I", i & 0xffffffff))[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.unpack("<I", struct.pack("<f", f))[0]
-  except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
 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 _sext(v, b): return v - (1 << b) if v & (1 << (b - 1)) else v
-def _f16(i): return struct.unpack("<e", struct.pack("<H", 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.unpack("<H", struct.pack("<e", f))[0]
-  except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
 def _to_f16_bits(v): return v if isinstance(v, int) else _i16(v)
-def _f64(i): return struct.unpack("<d", struct.pack("<Q", i & 0xffffffffffffffff))[0]
-def _i64(f):
-  if math.isnan(f): return 0x7ff8000000000000
-  if math.isinf(f): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
-  try: return struct.unpack("<Q", struct.pack("<d", f))[0]
-  except (OverflowError, struct.error): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
 def _isnan(x):
   try: return math.isnan(float(x))
   except (TypeError, ValueError): return False
-def _isquietnan(x):
-  """Check if x is a quiet NaN.
-  f16: exponent=31, bit9=1, mantissa!=0
-  f32: exponent=255, bit22=1, mantissa!=0
-  f64: exponent=2047, bit51=1, mantissa!=0
-  """
+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
-    # Get raw bits from TypedView or similar object with _reg attribute
     if hasattr(x, '_reg') and hasattr(x, '_bits'):
       bits = x._reg._val & ((1 << x._bits) - 1)
-      if x._bits == 16:
-        return ((bits >> 10) & 0x1f) == 31 and ((bits >> 9) & 1) == 1 and (bits & 0x3ff) != 0
-      if x._bits == 32:
-        return ((bits >> 23) & 0xff) == 255 and ((bits >> 22) & 1) == 1 and (bits & 0x7fffff) != 0
-      if x._bits == 64:
-        return ((bits >> 52) & 0x7ff) == 0x7ff and ((bits >> 51) & 1) == 1 and (bits & 0xfffffffffffff) != 0
-    return True  # Default to quiet NaN if we can't determine bit pattern
-  except (TypeError, ValueError): return False
-def _issignalnan(x):
-  """Check if x is a signaling NaN.
-  f16: exponent=31, bit9=0, mantissa!=0
-  f32: exponent=255, bit22=0, mantissa!=0
-  f64: exponent=2047, bit51=0, mantissa!=0
-  """
-  try:
-    if not math.isnan(float(x)): return False
-    # Get raw bits from TypedView or similar object with _reg attribute
-    if hasattr(x, '_reg') and hasattr(x, '_bits'):
-      bits = x._reg._val & ((1 << x._bits) - 1)
-      if x._bits == 16:
-        return ((bits >> 10) & 0x1f) == 31 and ((bits >> 9) & 1) == 0 and (bits & 0x3ff) != 0
-      if x._bits == 32:
-        return ((bits >> 23) & 0xff) == 255 and ((bits >> 22) & 1) == 0 and (bits & 0x7fffff) != 0
-      if x._bits == 64:
-        return ((bits >> 52) & 0x7ff) == 0x7ff and ((bits >> 51) & 1) == 0 and (bits & 0xfffffffffffff) != 0
-    return False  # Default to not signaling if we can't determine bit pattern
+      # 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 trunc(x):
-  x = float(x)
-  return x if math.isnan(x) or math.isinf(x) else float(math.trunc(x))
-def floor(x):
-  x = float(x)
-  return x if math.isnan(x) or math.isinf(x) else float(math.floor(x))
-def ceil(x):
-  x = float(x)
-  return x if math.isnan(x) or math.isinf(x) else float(math.ceil(x))
+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))
@@ -92,20 +44,10 @@ class _SafeFloat(float):
 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 f32_to_i32(f):
-  f = float(f)
-  if math.isnan(f): return 0
-  if f >= 2147483647: return 2147483647
-  if f <= -2147483648: return -2147483648
-  return int(f)
-def f32_to_u32(f):
-  f = float(f)
-  if math.isnan(f): return 0
-  if f >= 4294967295: return 4294967295
-  if f <= 0: return 0
-  return int(f)
-f64_to_i32 = f32_to_i32
-f64_to_u32 = f32_to_u32
+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
@@ -129,38 +71,26 @@ def isEven(x):
   return int(x) % 2 == 0
 def fract(x): return x - math.floor(x)
 PI = math.pi
-def sin(x):
-  # V_SIN_F32: pseudocode does sin(input * 2π), but hardware does frac on the input first
-  # So sin(1.0 * 2π) should be sin(frac(1.0) * 2π) = sin(0) = 0
+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")
-  # The input x is already multiplied by 2π in the pseudocode, so we need to
-  # extract the fractional cycle: frac(x / 2π) * 2π
-  cycles = x / (2 * math.pi)
-  frac_cycles = cycles - math.floor(cycles)
-  return math.sin(frac_cycles * 2 * math.pi)
-def cos(x):
-  # V_COS_F32: same as sin, hardware does frac on input cycles
-  if math.isinf(x) or math.isnan(x): return float("nan")
-  cycles = x / (2 * math.pi)
-  frac_cycles = cycles - math.floor(cycles)
-  return math.cos(frac_cycles * 2 * math.pi)
+  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 _brev32(v): return int(bin(v & 0xffffffff)[2:].zfill(32)[::-1], 2)
-def _brev64(v): return int(bin(v & 0xffffffffffffffff)[2:].zfill(64)[::-1], 2)
-def _ctz32(v):
-  v = int(v) & 0xffffffff
-  if v == 0: return 32
-  n = 0
-  while (v & 1) == 0: v >>= 1; n += 1
-  return n
-def _ctz64(v):
-  v = int(v) & 0xffffffffffffffff
-  if v == 0: return 64
-  n = 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:
@@ -184,34 +114,21 @@ def _is_denorm_f64(f):
   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):
-  if math.isnan(b): return a
-  if math.isnan(a): return b
-  return a if _lt_neg_zero(a, b) else b
-def v_max_f32(a, b):
-  if math.isnan(b): return a
-  if math.isnan(a): return b
-  return a if _gt_neg_zero(a, b) else b
-def v_min_i32(a, b): return min(a, b)
-def v_max_i32(a, b): return max(a, b)
-def v_min_u32(a, b): return min(a & 0xffffffff, b & 0xffffffff)
-def v_max_u32(a, b): return max(a & 0xffffffff, b & 0xffffffff)
-v_min_f16 = v_min_f32
-v_max_f16 = v_max_f32
-v_min_i16 = v_min_i32
-v_max_i16 = v_max_i32
+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)
-def v_min3_i32(a, b, c): return min(a, b, c)
-def v_max3_i32(a, b, c): return max(a, b, c)
-def v_min3_u32(a, b, c): return min(a & 0xffffffff, b & 0xffffffff, c & 0xffffffff)
-def v_max3_u32(a, b, c): return max(a & 0xffffffff, b & 0xffffffff, c & 0xffffffff)
-v_min3_f16 = v_min3_f32
-v_max3_f16 = v_max3_f32
-v_min3_i16 = v_min3_i32
-v_max3_i16 = v_max3_i32
+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))
@@ -295,7 +212,7 @@ def signext_from_bit(val, bit):
 
 __all__ = [
   # Classes
-  'Reg', 'SliceProxy', 'TypedView', 'ExecContext', 'compile_pseudocode',
+  'Reg', 'SliceProxy', 'TypedView',
   # Pack functions
   '_pack', '_pack32', 'pack', 'pack32',
   # Constants
@@ -381,8 +298,6 @@ ROUND_MODE = _RoundMode()
 def cvtToQuietNAN(x): return float('nan')
 DST = None  # Placeholder, will be set in context
 
-MASK32, MASK64 = 0xffffffff, 0xffffffffffffffff
-
 # 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
@@ -624,587 +539,4 @@ class Reg:
   def __eq__(s, o): return s._val == int(o)
   def __ne__(s, o): return s._val != int(o)
 
-# ═══════════════════════════════════════════════════════════════════════════════
-# COMPILER: pseudocode -> Python (minimal transforms)
-# ═══════════════════════════════════════════════════════════════════════════════
 
-def compile_pseudocode(pseudocode: str) -> str:
-  """Compile pseudocode to Python. Transforms are minimal - most syntax just works."""
-  # Join continuation lines (lines ending with || or && or open paren)
-  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
-
-    # Control flow - only need pass before outdent (endif/endfor/else/elsif)
-    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(';')
-      # Handle tuple unpacking: { D1.u1, D0.u64 } = expr
-      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)")
-      # Compound assignment
-      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 = lhs.strip(), rhs.strip()
-        stmt = _assign(lhs_s, _expr(rhs_s))
-        # CLZ/CTZ pattern: assignment of loop var to tmp/D0.i32 in first-match loop needs break
-        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 we ended with a control statement that needs a body, add pass
-  if need_pass: lines.append('  ' * indent + "pass")
-  return '\n'.join(lines)
-
-def _assign(lhs: str, rhs: str) -> str:
-  """Generate assignment. Bare tmp/SCC/etc get wrapped in Reg()."""
-  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:
-  """Expression transform: minimal - just fix syntax differences."""
-  e = e.strip()
-  e = e.replace('&&', ' and ').replace('||', ' or ').replace('<>', ' != ')
-  e = re.sub(r'!([^=])', r' not \1', e)
-
-  # Pack: { hi, lo } -> _pack(hi, lo)
-  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)
-
-  # Special constant: 1201'B(2.0 / PI) -> TWO_OVER_PI_1201 (precomputed 1201-bit 2/pi)
-  e = re.sub(r"1201'B\(2\.0\s*/\s*PI\)", "TWO_OVER_PI_1201", e)
-
-  # Literals: 1'0U -> 0, 32'I(x) -> (x), B(x) -> (x)
-  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)  # Bare B( without digit prefix
-  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)
-  # Remove redundant type suffix after lane access: VCC.u64[laneId].u64 -> VCC.u64[laneId]
-  e = re.sub(r'(\[laneId\])\.[uib]\d+', r'\1', e)
-
-  # Constants - INF is defined as an object supporting .f32/.f64 access
-  e = e.replace('+INF', 'INF').replace('-INF', '(-INF)')
-  e = re.sub(r'NAN\.f\d+', 'float("nan")', e)
-
-  # Verilog bit slice syntax: [start +: width] -> extract width bits starting at start
-  # Convert to Python slice: [start + width - 1 : start]
-  def convert_verilog_slice(m):
-    start, width = m.group(1).strip(), m.group(2).strip()
-    # Convert to high:low slice format
-    return f'[({start}) + ({width}) - 1 : ({start})]'
-  e = re.sub(r'\[([^:\[\]]+)\s*\+:\s*([^:\[\]]+)\]', convert_verilog_slice, e)
-
-  # Recursively process bracket contents to handle nested ternaries like S1.u32[x ? a : b]
-  def process_brackets(s):
-    result, i = [], 0
-    while i < len(s):
-      if s[i] == '[':
-        # Find matching ]
-        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])  # Recursively process bracket content
-        result.append('[' + inner + ']')
-        i = j
-      else:
-        result.append(s[i])
-        i += 1
-    return ''.join(result)
-  e = process_brackets(e)
-
-  # Ternary: a ? b : c -> (b if a else c)
-  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
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# EXECUTION CONTEXT
-# ═══════════════════════════════════════════════════════════════════════════════
-
-class ExecContext:
-  """Context for running compiled pseudocode."""
-  def __init__(self, s0=0, s1=0, s2=0, d0=0, scc=0, vcc=0, lane=0, exec_mask=MASK32, literal=0, vgprs=None, src0_idx=0, vdst_idx=0):
-    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."""
-    # Start with module globals (helpers, aliases), then add instance-specific bindings
-    ns = dict(globals())
-    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': SliceProxy(self.EXEC, 31, 0), 'EXEC_HI': 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)
-    # Sync rebinds: if register was reassigned to new Reg or value, copy it back
-    def _sync(ctx_reg, ns_val):
-      if isinstance(ns_val, Reg): ctx_reg._val = ns_val._val
-      else: ctx_reg._val = int(ns_val) & MASK64
-    if ns.get('SCC') is not self.SCC: _sync(self.SCC, ns['SCC'])
-    if ns.get('VCC') is not self.VCC: _sync(self.VCC, ns['VCC'])
-    if ns.get('EXEC') is not self.EXEC: _sync(self.EXEC, ns['EXEC'])
-    if ns.get('D0') is not self.D0: _sync(self.D0, ns['D0'])
-    if ns.get('D1') is not self.D1: _sync(self.D1, ns['D1'])
-    if ns.get('tmp') is not self.tmp: _sync(self.tmp, ns['tmp'])
-    if ns.get('saveexec') is not self.saveexec: _sync(self.saveexec, ns['saveexec'])
-
-  def result(self) -> dict:
-    return {"d0": self.D0._val, "scc": self.SCC._val & 1}
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# PDF EXTRACTION AND CODE GENERATION
-# ═══════════════════════════════════════════════════════════════════════════════
-
-from extra.assembly.amd.dsl import PDF_URLS
-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']  # Malformed pseudocode from PDF
-
-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: continue
-    if re.match(r'^\d+ of \d+$', s): continue
-    if re.match(r'^\d+\.\d+\..*Instructions', s): continue
-    # Skip document headers (RDNA or CDNA)
-    if s.startswith('"RDNA') or s.startswith('AMD ') or s.startswith('CDNA'): continue
-    if s.startswith('Notes') or s.startswith('Functional examples'): break
-    # Track lambda definitions (e.g., BYTE_PERMUTE = lambda(data, sel) (...))
-    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 =']) or
-      any(p in s for p in ['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 _get_op_enums(arch: str) -> list:
-  """Dynamically load op enums from the arch-specific autogen module."""
-  import importlib
-  autogen = importlib.import_module(f"extra.assembly.amd.autogen.{arch}")
-  # Deterministic order: common enums first, then arch-specific
-  enums = []
-  for name in ['SOP1Op', 'SOP2Op', 'SOPCOp', 'SOPKOp', 'SOPPOp', 'VOP1Op', 'VOP2Op', 'VOP3Op', 'VOP3SDOp', 'VOP3POp', 'VOPCOp', 'VOP3AOp', 'VOP3BOp']:
-    if hasattr(autogen, name): enums.append(getattr(autogen, name))
-  return enums
-
-def _parse_pseudocode_from_single_pdf(url: str, defined_ops: dict, OP_ENUMS: list) -> dict:
-  """Parse pseudocode from a single PDF."""
-  import pdfplumber
-  from tinygrad.helpers import fetch
-
-  pdf = pdfplumber.open(fetch(url))
-  total_pages = len(pdf.pages)
-
-  page_cache = {}
-  def get_page_text(i):
-    if i not in page_cache: page_cache[i] = pdf.pages[i].extract_text() or ''
-    return page_cache[i]
-
-  # Find the "Instructions" chapter - typically 10-40% through the document
-  instr_start = None
-  for i in range(int(total_pages * 0.1), int(total_pages * 0.5)):
-    if re.search(r'Chapter \d+\.\s+Instructions\b', get_page_text(i)):
-      instr_start = i
-      break
-  if instr_start is None: instr_start = total_pages // 3  # fallback
-
-  # Find end - stop at "Microcode Formats" chapter (typically 60-70% through)
-  instr_end = total_pages
-  search_starts = [int(total_pages * 0.6), int(total_pages * 0.5), instr_start]
-  for start in search_starts:
-    for i in range(start, min(start + 100, total_pages)):
-      if re.search(r'Chapter \d+\.\s+Microcode Formats', get_page_text(i)):
-        instr_end = i
-        break
-    if instr_end < total_pages: break
-
-  # Extract remaining pages (some already cached from chapter search)
-  all_text = '\n'.join(get_page_text(i) for i in range(instr_start, instr_end))
-  matches = list(INST_PATTERN.finditer(all_text))
-  instructions: dict = {cls: {} for cls in OP_ENUMS}
-
-  for i, match in enumerate(matches):
-    name, opcode = match.group(1), int(match.group(2))
-    key = (name, opcode)
-    if key not in defined_ops: continue
-    start = match.end()
-    end = matches[i + 1].start() if i + 1 < len(matches) else start + 2000
-    snippet = all_text[start:end].strip()
-    if (pseudocode := extract_pseudocode(snippet)):
-      # Assign to all enums that have this op (e.g., both VOPCOp and VOP3AOp)
-      for enum_cls, enum_val in defined_ops[key]:
-        instructions[enum_cls][enum_val] = pseudocode
-
-  return instructions
-
-def parse_pseudocode_from_pdf(arch: str = "rdna3") -> dict:
-  """Parse pseudocode from PDF(s) for all ops. Returns {enum_cls: {op: pseudocode}}."""
-  OP_ENUMS = _get_op_enums(arch)
-  # Build a dict from (name, opcode) -> list of (enum_cls, op) tuples
-  # Multiple enums can have the same op (e.g., VOPCOp and VOP3AOp both have V_CMP_* ops)
-  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))
-
-  urls = PDF_URLS[arch]
-  if isinstance(urls, str): urls = [urls]
-
-  # Parse all PDFs and merge (union of pseudocode)
-  # Reverse order so newer PDFs (RDNA3.5, CDNA4) take priority
-  instructions: dict = {cls: {} for cls in OP_ENUMS}
-  for url in reversed(urls):
-    result = _parse_pseudocode_from_single_pdf(url, defined_ops, OP_ENUMS)
-    for cls, ops in result.items():
-      for op, pseudocode in ops.items():
-        if op in instructions[cls]:
-          if instructions[cls][op] != pseudocode:
-            print(f"  Ignoring {op.name} from older PDF:")
-            print(f"    new: {instructions[cls][op]!r}")
-            print(f"    old: {pseudocode!r}")
-        else:
-          instructions[cls][op] = pseudocode
-
-  return instructions
-
-def generate_gen_pcode(output_path: str = "extra/assembly/amd/autogen/rdna3/gen_pcode.py", arch: str = "rdna3"):
-  """Generate gen_pcode.py - compiled pseudocode functions for the emulator."""
-  from pathlib import Path
-
-  OP_ENUMS = _get_op_enums(arch)
-
-  print("Parsing pseudocode from PDF...")
-  by_cls = parse_pseudocode_from_pdf(arch)
-
-  total_found, total_ops = 0, 0
-  for enum_cls in OP_ENUMS:
-    total = sum(1 for op in enum_cls if op.name.startswith(('S_', 'V_')))
-    found = len(by_cls.get(enum_cls, {}))
-    total_found += found
-    total_ops += total
-    print(f"{enum_cls.__name__}: {found}/{total} ({100*found//total if total else 0}%)")
-  print(f"Total: {total_found}/{total_ops} ({100*total_found//total_ops}%)")
-
-  print("\nCompiling to pseudocode functions...")
-  # Build dynamic import line based on available enums
-  enum_names = [e.__name__ for e in OP_ENUMS]
-  lines = [f'''# autogenerated by pcode.py - do not edit
-# to regenerate: python -m extra.assembly.amd.pcode --arch {arch}
-# ruff: noqa: E501,F405,F403
-# mypy: ignore-errors
-from extra.assembly.amd.autogen.{arch} import {", ".join(enum_names)}
-from extra.assembly.amd.pcode import *
-''']
-
-  compiled_count, skipped_count = 0, 0
-
-  for enum_cls in OP_ENUMS:
-    cls_name = enum_cls.__name__
-    pseudocode_dict = by_cls.get(enum_cls, {})
-    if not pseudocode_dict: continue
-
-    fn_entries = []
-    for op, pc in pseudocode_dict.items():
-      if any(p in pc for p in UNSUPPORTED):
-        skipped_count += 1
-        continue
-
-      try:
-        code = compile_pseudocode(pc)
-        # NOTE: Do NOT add more code.replace() hacks here. Fix issues properly in the DSL
-        # (compile_pseudocode, helper functions, or Reg/TypedView classes) instead.
-        # V_DIV_FMAS_F32/F64: PDF page 449 says 2^32/2^64 but hardware behavior is more complex.
-        # The scale direction depends on S2 (the addend): if exponent(S2) > 127 (i.e., S2 >= 2.0),
-        # scale by 2^+64 (to unscale a numerator that was scaled). Otherwise scale by 2^-64
-        # (to unscale a denominator that was scaled).
-        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)')
-        # V_DIV_SCALE_F32/F64: PDF page 463-464 has several bugs vs hardware behavior:
-        # 1. Zero case: hardware sets VCC=1 (PDF doesn't)
-        # 2. Denorm denom: hardware returns NaN (PDF says scale). VCC is set independently by exp diff check.
-        # 3. Tiny numer (exp<=23): hardware sets VCC=1 (PDF doesn't)
-        # 4. Result would be denorm: hardware doesn't scale, just sets VCC=1
-        if op.name == 'V_DIV_SCALE_F32':
-          # Fix 1: Set VCC=1 when zero operands produce NaN
-          code = code.replace(
-            'D0.f32 = float("nan")',
-            'VCC = Reg(0x1); D0.f32 = float("nan")')
-          # Fix 2: Denorm denom returns NaN. Must check this AFTER all VCC-setting logic runs.
-          # Insert at end of all branches, before the final result is used
-          code = code.replace(
-            'elif S1.f32 == DENORM.f32:\n  D0.f32 = ldexp(S0.f32, 64)',
-            'elif False:\n  pass  # denorm check moved to end')
-          # Add denorm check at the very end - this overrides D0 but preserves VCC
-          code += '\nif S1.f32 == DENORM.f32:\n  D0.f32 = float("nan")'
-          # Fix 3: Tiny numer should set VCC=1
-          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)')
-          # Fix 4: S2/S1 would be denorm - don't scale, just set VCC
-          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':
-          # Same fixes for f64 version
-          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  # denorm check moved to end')
-          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)')
-        # V_DIV_FIXUP_F32/F64: PDF doesn't check isNAN(S0), but hardware returns OVERFLOW if S0 is NaN.
-        # When division fails (e.g., due to denorm denom), S0 becomes NaN, and fixup should return ±inf.
-        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)))')
-        # V_TRIG_PREOP_F64: AMD pseudocode uses (x << shift) & mask but mask needs to extract TOP bits.
-        # The PDF shows: result = 64'F((1201'B(2.0/PI)[1200:0] << shift) & 1201'0x1fffffffffffff)
-        # Issues to fix:
-        # 1. After left shift, the interesting bits are at the top, not bottom - need >> (1201-53)
-        # 2. shift.u32 fails because shift is a plain int after * 53 - use int(shift)
-        # 3. 64'F(...) means convert int to float (not interpret as bit pattern) - use float()
-        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)')
-        # Detect flags for result handling
-        is_64 = any(p in pc for p in ['D0.u64', 'D0.b64', 'D0.f64', 'D0.i64', 'D1.u64', 'D1.b64', 'D1.f64', 'D1.i64'])
-        has_d1 = '{ D1' in pc
-        if has_d1: is_64 = True
-        is_cmp = (cls_name == 'VOPCOp' or cls_name == 'VOP3Op') and 'D0.u64[laneId]' in pc
-        is_cmpx = (cls_name == 'VOPCOp' or cls_name == 'VOP3Op') and 'EXEC.u64[laneId]' in pc  # V_CMPX writes to EXEC per-lane
-        # V_DIV_SCALE passes through S0 if no branch taken
-        is_div_scale = 'DIV_SCALE' in op.name
-        # VOP3SD instructions that write VCC per-lane (either via VCC.u64[laneId] or by setting VCC = 0/1)
-        has_sdst = cls_name == 'VOP3SDOp' and ('VCC.u64[laneId]' in pc or is_div_scale)
-        # Instructions that use/modify PC
-        has_pc = 'PC' in pc
-
-        # Generate function with indented body
-        fn_name = f"_{cls_name}_{op.name}"
-        lines.append(f"def {fn_name}(s0, s1, s2, d0, scc, vcc, lane, exec_mask, literal, VGPR, _vars, src0_idx=0, vdst_idx=0, pc=0):")
-        # Add original pseudocode as comment
-        for pc_line in pc.split('\n'):
-          lines.append(f"  # {pc_line}")
-        # Only create Reg objects for registers actually used in the pseudocode
-        combined = code + pc
-        regs = [('S0', 'Reg(s0)'), ('S1', 'Reg(s1)'), ('S2', 'Reg(s2)'),
-                ('D0', 'Reg(s0)' if is_div_scale else 'Reg(d0)'), ('D1', 'Reg(0)'),
-                ('SCC', 'Reg(scc)'), ('VCC', 'Reg(vcc)'), ('EXEC', 'Reg(exec_mask)'),
-                ('tmp', 'Reg(0)'), ('saveexec', 'Reg(exec_mask)'), ('laneId', 'lane'),
-                ('SIMM16', 'Reg(literal)'), ('SIMM32', 'Reg(literal)'),
-                ('SRC0', 'Reg(src0_idx)'), ('VDST', 'Reg(vdst_idx)'),
-                ('PC', 'Reg(pc)')]  # PC is passed in as byte address
-        used = {name for name, _ in regs if name in combined}
-        # EXEC_LO/EXEC_HI need EXEC
-        if 'EXEC_LO' in combined or 'EXEC_HI' in combined: used.add('EXEC')
-        # VCCZ/EXECZ need VCC/EXEC
-        if 'VCCZ' in combined: used.add('VCC')
-        if 'EXECZ' in combined: used.add('EXEC')
-        for name, init in regs:
-          if name in used: lines.append(f"  {name} = {init}")
-        if 'EXEC_LO' in combined: lines.append("  EXEC_LO = SliceProxy(EXEC, 31, 0)")
-        if 'EXEC_HI' in combined: lines.append("  EXEC_HI = SliceProxy(EXEC, 63, 32)")
-        # VCCZ = 1 if VCC == 0, EXECZ = 1 if EXEC == 0
-        if 'VCCZ' in combined: lines.append("  VCCZ = Reg(1 if VCC._val == 0 else 0)")
-        if 'EXECZ' in combined: lines.append("  EXECZ = Reg(1 if EXEC._val == 0 else 0)")
-        # Add compiled pseudocode with markers
-        lines.append("  # --- compiled pseudocode ---")
-        for line in code.split('\n'):
-          lines.append(f"  {line}")
-        lines.append("  # --- end pseudocode ---")
-        # Generate result dict - use raw params if Reg wasn't created
-        d0_val = "D0._val" if 'D0' in used else "d0"
-        scc_val = "SCC._val & 1" if 'SCC' in used else "scc & 1"
-        lines.append(f"  result = {{'d0': {d0_val}, 'scc': {scc_val}}}")
-        if has_sdst:
-          lines.append("  result['vcc_lane'] = (VCC._val >> lane) & 1")
-        elif 'VCC' in used:
-          lines.append("  if VCC._val != vcc: result['vcc_lane'] = (VCC._val >> lane) & 1")
-        if is_cmpx:
-          lines.append("  result['exec_lane'] = (EXEC._val >> lane) & 1")
-        elif 'EXEC' in used:
-          lines.append("  if EXEC._val != exec_mask: result['exec'] = EXEC._val")
-        if is_cmp:
-          lines.append("  result['vcc_lane'] = (D0._val >> lane) & 1")
-        if is_64:
-          lines.append("  result['d0_64'] = True")
-        if has_d1:
-          lines.append("  result['d1'] = D1._val & 1")
-        if has_pc:
-          # Return new PC as absolute byte address, emulator will compute delta
-          # Handle negative values (backward jumps): PC._val is stored as unsigned, convert to signed
-          lines.append("  _pc = PC._val if PC._val < 0x8000000000000000 else PC._val - 0x10000000000000000")
-          lines.append("  result['new_pc'] = _pc  # absolute byte address")
-        lines.append("  return result")
-        lines.append("")
-
-        fn_entries.append((op, fn_name))
-        compiled_count += 1
-      except Exception as e:
-        print(f"  Warning: Failed to compile {op.name}: {e}")
-        skipped_count += 1
-
-    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('}')
-      lines.append('')
-
-  # Add manually implemented V_WRITELANE_B32 (not in PDF pseudocode, requires special vgpr_write handling)
-  # Only add for architectures that have VOP3Op (RDNA) not VOP3AOp/VOP3BOp (CDNA)
-  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  # lane select (5 bits for wave32)
-  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:
-    cls_name = enum_cls.__name__
-    if by_cls.get(enum_cls): lines.append(f'  {cls_name}: {cls_name}_FUNCTIONS,')
-  lines.append('}')
-  lines.append('')
-  lines.append('def get_compiled_functions(): return COMPILED_FUNCTIONS')
-
-  Path(output_path).write_text('\n'.join(lines))
-  print(f"\nGenerated {output_path}: {compiled_count} compiled, {skipped_count} skipped")
-
-if __name__ == "__main__":
-  import argparse
-  parser = argparse.ArgumentParser(description="Generate pseudocode functions from AMD ISA PDF")
-  parser.add_argument("--arch", choices=list(PDF_URLS.keys()) + ["all"], default="rdna3", help="Target architecture (default: rdna3)")
-  args = parser.parse_args()
-  if args.arch == "all":
-    for arch in PDF_URLS.keys():
-      generate_gen_pcode(output_path=f"extra/assembly/amd/autogen/{arch}/gen_pcode.py", arch=arch)
-  else:
-    generate_gen_pcode(output_path=f"extra/assembly/amd/autogen/{args.arch}/gen_pcode.py", arch=args.arch)

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/external_test_usability.py

@@ -3,7 +3,7 @@
 # Currently many of these tests fail - they document desired behavior
 
 import unittest
-from extra.assembly.amd.autogen.rdna3 import *
+from extra.assembly.amd.autogen.rdna3.ins import *
 from extra.assembly.amd.dsl import Inst, RawImm, SGPR, VGPR
 
 class TestRegisterSliceSyntax(unittest.TestCase):

extra/assembly/amd/test/helpers.py

@@ -22,3 +22,45 @@ def get_llvm_objdump():
   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_emu.py

@@ -6,7 +6,7 @@ Set USE_HW=1 to run on both emulator and real hardware, comparing results.
 """
 
 import ctypes, unittest, os, struct
-from extra.assembly.amd.autogen.rdna3 import *
+from extra.assembly.amd.autogen.rdna3.ins import *
 from extra.assembly.amd.dsl import RawImm
 from extra.assembly.amd.emu import WaveState, run_asm, set_valid_mem_ranges
 from extra.assembly.amd.pcode import _i32, _f32
@@ -3813,3 +3813,446 @@ class TestVTrigPreopF64(unittest.TestCase):
     # Result should still be a valid float (not NaN or inf)
     self.assertFalse(math.isnan(result), "Result should not be NaN")
     self.assertFalse(math.isinf(result), "Result should not be inf")
+
+
+class Test64BitLiterals(unittest.TestCase):
+  """Regression tests for 64-bit instruction literal encoding.
+  Tests verify that Inst.to_bytes() correctly encodes 64-bit literals."""
+
+  def test_64bit_literal_negative_encoding(self):
+    """Verify 64-bit instruction encodes negative literals correctly.
+    Regression test: -33 should encode as 0xffffffdf in the literal field,
+    NOT as 0xffffffff (which would happen with incorrect sign extension)."""
+    neg_val = -33
+    expected_lit = neg_val & 0xffffffff  # 0xffffffdf
+    inst = v_add_f64(v[2], v[0], neg_val)
+    # Check the literal is stored correctly (in high 32 bits for 64-bit ops)
+    self.assertIsNotNone(inst._literal, "Literal should be set")
+    # Literal is stored as (lit32 << 32) for 64-bit ops
+    actual_lit = (inst._literal >> 32) & 0xffffffff
+    self.assertEqual(actual_lit, expected_lit, f"Literal should be {expected_lit:#x}, got {actual_lit:#x}")
+    # Also verify the encoded bytes
+    code = inst.to_bytes()
+    # Literal is last 4 bytes
+    lit_bytes = code[-4:]
+    lit_val = int.from_bytes(lit_bytes, 'little')
+    self.assertEqual(lit_val, expected_lit, f"Encoded literal should be {expected_lit:#x}, got {lit_val:#x}")
+
+  def test_64bit_literal_positive_encoding(self):
+    """Verify 64-bit instruction encodes large positive literals correctly."""
+    large_val = 0x12345678
+    inst = v_add_f64(v[2], v[0], large_val)
+    self.assertIsNotNone(inst._literal, "Literal should be set")
+    actual_lit = (inst._literal >> 32) & 0xffffffff
+    self.assertEqual(actual_lit, large_val, f"Literal should be {large_val:#x}, got {actual_lit:#x}")
+    # Verify encoded bytes
+    code = inst.to_bytes()
+    lit_bytes = code[-4:]
+    lit_val = int.from_bytes(lit_bytes, 'little')
+    self.assertEqual(lit_val, large_val, f"Encoded literal should be {large_val:#x}, got {lit_val:#x}")
+
+
+class TestWave32VCCBranch(unittest.TestCase):
+  """Regression tests for wave32 VCC branch behavior.
+  In wave32 mode, S_CBRANCH_VCCNZ/VCCZ should only check VCC_LO (lower 32 bits),
+  ignoring VCC_HI. Bug: emulator was checking full 64-bit VCC, causing incorrect
+  branches when VCC_LO=0 but VCC_HI!=0."""
+
+  def test_cbranch_vccnz_ignores_vcc_hi(self):
+    """S_CBRANCH_VCCNZ should NOT branch when VCC_LO=0, even if VCC_HI!=0.
+    This is the fix for test_avg_pool3d failure where the emulator incorrectly
+    branched due to stale VCC_HI bits."""
+    instructions = [
+      # Set VCC_HI to non-zero (simulating stale bits from previous ops)
+      s_mov_b32(s[SrcEnum.VCC_HI - 128], 0x80000000),  # VCC_HI = 0x80000000
+      # Set VCC_LO to zero (the condition we're testing)
+      s_mov_b32(s[SrcEnum.VCC_LO - 128], 0),  # VCC_LO = 0
+      # Now S_CBRANCH_VCCNZ should NOT branch since VCC_LO is 0
+      # If it doesn't branch, we'll set v0 = 1; if it branches, v0 stays 0
+      v_mov_b32_e32(v[0], 0),
+      s_cbranch_vccnz(2),  # Skip next instruction if VCC != 0
+      v_mov_b32_e32(v[0], 1),  # This should execute
+      s_nop(0),  # Jump target
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # v0 should be 1 because VCC_LO=0 means no branch
+    self.assertEqual(st.vgpr[0][0], 1, "Should NOT branch when VCC_LO=0 (VCC_HI ignored in wave32)")
+
+  def test_cbranch_vccz_ignores_vcc_hi(self):
+    """S_CBRANCH_VCCZ should branch when VCC_LO=0, regardless of VCC_HI."""
+    instructions = [
+      # Set VCC_HI to non-zero (simulating stale bits)
+      s_mov_b32(s[SrcEnum.VCC_HI - 128], 0x80000000),  # VCC_HI = 0x80000000
+      # Set VCC_LO to zero
+      s_mov_b32(s[SrcEnum.VCC_LO - 128], 0),  # VCC_LO = 0
+      # S_CBRANCH_VCCZ should branch since VCC_LO is 0
+      v_mov_b32_e32(v[0], 0),
+      s_cbranch_vccz(2),  # Skip next instruction if VCC == 0
+      v_mov_b32_e32(v[0], 1),  # This should NOT execute
+      s_nop(0),  # Jump target
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # v0 should be 0 because VCC_LO=0 means branch is taken
+    self.assertEqual(st.vgpr[0][0], 0, "Should branch when VCC_LO=0 (VCC_HI ignored in wave32)")
+
+  def test_cbranch_vccnz_branches_on_vcc_lo(self):
+    """S_CBRANCH_VCCNZ should branch when VCC_LO!=0."""
+    instructions = [
+      # Set VCC_LO to non-zero
+      s_mov_b32(s[SrcEnum.VCC_LO - 128], 1),  # VCC_LO = 1
+      s_mov_b32(s[SrcEnum.VCC_HI - 128], 0),  # VCC_HI = 0
+      v_mov_b32_e32(v[0], 0),
+      s_cbranch_vccnz(2),  # Skip next instruction if VCC != 0
+      v_mov_b32_e32(v[0], 1),  # This should NOT execute
+      s_nop(0),  # Jump target
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # v0 should be 0 because VCC_LO=1 means branch is taken
+    self.assertEqual(st.vgpr[0][0], 0, "Should branch when VCC_LO!=0")
+
+
+class TestVOP3VOPC16Bit(unittest.TestCase):
+  """Regression tests for VOP3-encoded VOPC 16-bit comparison instructions.
+  When VOPC comparisons are encoded in VOP3 format, they use opsel bits to select
+  which 16-bit half of each source to compare.
+  Bug: Emulator was ignoring opsel and using VGPR bit 7 encoding instead."""
+
+  def test_cmp_eq_u16_opsel_lo_lo(self):
+    """V_CMP_EQ_U16 VOP3 with opsel=0 compares lo halves."""
+    # v0 = 0x12340005 (lo=5, hi=0x1234)
+    # v1 = 0x56780005 (lo=5, hi=0x5678)
+    # opsel=0: compare lo halves -> 5 == 5 -> true
+    instructions = [
+      s_mov_b32(s[2], 0x12340005),
+      v_mov_b32_e32(v[0], s[2]),
+      s_mov_b32(s[2], 0x56780005),
+      v_mov_b32_e32(v[1], s[2]),
+      VOP3(VOP3Op.V_CMP_EQ_U16, vdst=v[0], src0=v[0], src1=v[1], opsel=0),  # dst=s0
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # s0 should have bit 0 set (comparison true for lane 0)
+    self.assertEqual(st.sgpr[0] & 1, 1, "lo==lo should be true: 5==5")
+
+  def test_cmp_eq_u16_opsel_hi_hi(self):
+    """V_CMP_EQ_U16 VOP3 with opsel=3 compares hi halves."""
+    # v0 = 0x12340005 (lo=5, hi=0x1234)
+    # v1 = 0x56780005 (lo=5, hi=0x5678)
+    # opsel=3 (bits 0 and 1 set): compare hi halves -> 0x1234 != 0x5678 -> false
+    instructions = [
+      s_mov_b32(s[2], 0x12340005),
+      v_mov_b32_e32(v[0], s[2]),
+      s_mov_b32(s[2], 0x56780005),
+      v_mov_b32_e32(v[1], s[2]),
+      VOP3(VOP3Op.V_CMP_EQ_U16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),  # dst=s0, hi vs hi
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # s0 should have bit 0 clear (comparison false for lane 0)
+    self.assertEqual(st.sgpr[0] & 1, 0, "hi==hi should be false: 0x1234!=0x5678")
+
+  def test_cmp_eq_u16_opsel_hi_hi_equal(self):
+    """V_CMP_EQ_U16 VOP3 with opsel=3 compares hi halves (equal case)."""
+    # v0 = 0x12340005 (lo=5, hi=0x1234)
+    # v1 = 0x12340009 (lo=9, hi=0x1234)
+    # opsel=3: compare hi halves -> 0x1234 == 0x1234 -> true
+    instructions = [
+      s_mov_b32(s[2], 0x12340005),
+      v_mov_b32_e32(v[0], s[2]),
+      s_mov_b32(s[2], 0x12340009),
+      v_mov_b32_e32(v[1], s[2]),
+      VOP3(VOP3Op.V_CMP_EQ_U16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),  # dst=s0, hi vs hi
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # s0 should have bit 0 set (comparison true for lane 0)
+    self.assertEqual(st.sgpr[0] & 1, 1, "hi==hi should be true: 0x1234==0x1234")
+
+  def test_cmp_gt_u16_opsel_hi(self):
+    """V_CMP_GT_U16 VOP3 with opsel=3 compares hi halves."""
+    # v0 = 0x99990005 (lo=5, hi=0x9999)
+    # v1 = 0x12340005 (lo=5, hi=0x1234)
+    # opsel=3: compare hi halves -> 0x9999 > 0x1234 -> true
+    instructions = [
+      s_mov_b32(s[2], 0x99990005),
+      v_mov_b32_e32(v[0], s[2]),
+      s_mov_b32(s[2], 0x12340005),
+      v_mov_b32_e32(v[1], s[2]),
+      VOP3(VOP3Op.V_CMP_GT_U16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),  # dst=s0, hi vs hi
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # s0 should have bit 0 set (comparison true for lane 0)
+    self.assertEqual(st.sgpr[0] & 1, 1, "hi>hi should be true: 0x9999>0x1234")
+
+
+class Test64BitLiteralSources(unittest.TestCase):
+  """Regression tests for 64-bit instruction literal source handling.
+
+  For f64 operations, a 32-bit literal in the instruction stream represents the
+  HIGH 32 bits of the 64-bit value (low 32 bits are implicitly 0).
+
+  Bug: rsrc64() was returning the 32-bit literal as-is instead of shifting it
+  left by 32 bits. This caused V_FMA_F64 and V_LDEXP_F64 to use wrong values
+  when their source is a literal, breaking the f64->i64 conversion sequence.
+
+  The f64->i64 conversion sequence is:
+    v_trunc_f64 -> v_ldexp_f64 (by -32) -> v_floor_f64 -> v_fma_f64 (by -2^32)
+    -> v_cvt_u32_f64 (low bits) -> v_cvt_i32_f64 (high bits)
+
+  The V_FMA_F64 uses literal 0xC1F00000 which is the high 32 bits of f64 -2^32.
+  """
+
+  def test_v_fma_f64_literal_neg_2pow32(self):
+    """V_FMA_F64 with literal encoding of -2^32.
+
+    The f64 value -2^32 (-4294967296.0) has bits 0xC1F0000000000000.
+    The compiler encodes only the high 32 bits (0xC1F00000) as a literal.
+    The emulator must interpret this as 0xC1F00000_00000000.
+    """
+    # v[0:1] = -41.0 (trunc), v[2:3] = -1.0 (floor of -41/2^32)
+    # FMA: result = (-2^32) * (-1.0) + (-41.0) = 4294967296 - 41 = 4294967255.0
+    val_41 = f2i64(-41.0)
+    val_m1 = f2i64(-1.0)
+    # Literal 0xC1F00000 is high 32 bits of f64 -2^32
+    lit = 0xC1F00000
+    instructions = [
+      s_mov_b32(s[0], val_41 & 0xffffffff),
+      s_mov_b32(s[1], (val_41 >> 32) & 0xffffffff),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      s_mov_b32(s[2], val_m1 & 0xffffffff),
+      s_mov_b32(s[3], (val_m1 >> 32) & 0xffffffff),
+      v_mov_b32_e32(v[2], s[2]),
+      v_mov_b32_e32(v[3], s[3]),
+      # V_FMA_F64 v[4:5], literal, v[2:3], v[0:1]
+      # = (-2^32) * (-1.0) + (-41.0) = 4294967255.0
+      VOP3(VOP3Op.V_FMA_F64, vdst=v[4], src0=RawImm(255), src1=v[2], src2=v[0], literal=lit),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = i642f(st.vgpr[0][4] | (st.vgpr[0][5] << 32))
+    expected = 4294967255.0  # 2^32 - 41
+    self.assertAlmostEqual(result, expected, places=0, msg=f"Expected {expected}, got {result}")
+
+  def test_v_ldexp_f64_literal_neg32(self):
+    """V_LDEXP_F64 with literal -32 for exponent.
+
+    V_LDEXP_F64 computes src0 * 2^src1 where src1 is an integer exponent.
+    The literal 0xFFFFFFE0 represents -32 as a 32-bit signed integer.
+    For V_LDEXP_F64, src1 is 32-bit (not 64-bit), so this is correct as-is.
+    """
+    val = f2i64(-41.0)
+    expected = -41.0 * (2.0 ** -32)  # -9.5367431640625e-09
+    instructions = [
+      s_mov_b32(s[0], val & 0xffffffff),
+      s_mov_b32(s[1], (val >> 32) & 0xffffffff),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      # V_LDEXP_F64 v[2:3], v[0:1], -32
+      v_ldexp_f64(v[2:4], v[0:2], 0xFFFFFFE0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = i642f(st.vgpr[0][2] | (st.vgpr[0][3] << 32))
+    self.assertAlmostEqual(result, expected, places=15, msg=f"Expected {expected}, got {result}")
+
+  def test_f64_to_i64_full_sequence(self):
+    """Full f64->i64 conversion sequence with negative value.
+
+    This is the exact sequence generated by the compiler for (long)(-41.0):
+      v_trunc_f64 v[0:1], v[0:1]
+      v_ldexp_f64 v[2:3], v[0:1], -32
+      v_floor_f64 v[2:3], v[2:3]
+      v_fma_f64 v[0:1], 0xc1f00000, v[2:3], v[0:1]  # -2^32
+      v_cvt_u32_f64 v0, v[0:1]
+      v_cvt_i32_f64 v1, v[2:3]
+
+    Result: v1:v0 = 0xFFFFFFFF:0xFFFFFFD7 = -41 as i64
+    """
+    val = f2i64(-41.0)
+    lit = 0xC1F00000  # high 32 bits of f64 -2^32
+    instructions = [
+      s_mov_b32(s[0], val & 0xffffffff),
+      s_mov_b32(s[1], (val >> 32) & 0xffffffff),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_trunc_f64_e32(v[0:2], v[0:2]),
+      v_ldexp_f64(v[2:4], v[0:2], 0xFFFFFFE0),  # -32
+      v_floor_f64_e32(v[2:4], v[2:4]),
+      VOP3(VOP3Op.V_FMA_F64, vdst=v[0], src0=RawImm(255), src1=v[2], src2=v[0], literal=lit),
+      v_cvt_u32_f64_e32(v[4], v[0:2]),
+      v_cvt_i32_f64_e32(v[5], v[2:4]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    lo = st.vgpr[0][4]
+    hi = st.vgpr[0][5]
+    result = struct.unpack('<q', struct.pack('<II', lo, hi))[0]
+    self.assertEqual(result, -41, f"Expected -41, got {result} (lo=0x{lo:08x}, hi=0x{hi:08x})")
+
+  def test_f64_to_i64_large_negative(self):
+    """f64->i64 conversion with larger negative value (-1000000).
+
+    Tests that the conversion sequence works for values that span both
+    high and low 32-bit parts of the result.
+    """
+    val = f2i64(-1000000.0)
+    lit = 0xC1F00000
+    instructions = [
+      s_mov_b32(s[0], val & 0xffffffff),
+      s_mov_b32(s[1], (val >> 32) & 0xffffffff),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_trunc_f64_e32(v[0:2], v[0:2]),
+      v_ldexp_f64(v[2:4], v[0:2], 0xFFFFFFE0),
+      v_floor_f64_e32(v[2:4], v[2:4]),
+      VOP3(VOP3Op.V_FMA_F64, vdst=v[0], src0=RawImm(255), src1=v[2], src2=v[0], literal=lit),
+      v_cvt_u32_f64_e32(v[4], v[0:2]),
+      v_cvt_i32_f64_e32(v[5], v[2:4]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    lo = st.vgpr[0][4]
+    hi = st.vgpr[0][5]
+    result = struct.unpack('<q', struct.pack('<II', lo, hi))[0]
+    self.assertEqual(result, -1000000, f"Expected -1000000, got {result}")
+
+  def test_f64_to_i64_positive(self):
+    """f64->i64 conversion with positive value (1000000)."""
+    val = f2i64(1000000.0)
+    lit = 0xC1F00000
+    instructions = [
+      s_mov_b32(s[0], val & 0xffffffff),
+      s_mov_b32(s[1], (val >> 32) & 0xffffffff),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_trunc_f64_e32(v[0:2], v[0:2]),
+      v_ldexp_f64(v[2:4], v[0:2], 0xFFFFFFE0),
+      v_floor_f64_e32(v[2:4], v[2:4]),
+      VOP3(VOP3Op.V_FMA_F64, vdst=v[0], src0=RawImm(255), src1=v[2], src2=v[0], literal=lit),
+      v_cvt_u32_f64_e32(v[4], v[0:2]),
+      v_cvt_i32_f64_e32(v[5], v[2:4]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    lo = st.vgpr[0][4]
+    hi = st.vgpr[0][5]
+    result = struct.unpack('<q', struct.pack('<II', lo, hi))[0]
+    self.assertEqual(result, 1000000, f"Expected 1000000, got {result}")
+
+  def test_f64_to_i64_large_positive(self):
+    """f64->i64 conversion with value > 2^32 (requires 64-bit result)."""
+    val = f2i64(5000000000.0)  # 5 billion, > 2^32
+    lit = 0xC1F00000
+    instructions = [
+      s_mov_b32(s[0], val & 0xffffffff),
+      s_mov_b32(s[1], (val >> 32) & 0xffffffff),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_trunc_f64_e32(v[0:2], v[0:2]),
+      v_ldexp_f64(v[2:4], v[0:2], 0xFFFFFFE0),
+      v_floor_f64_e32(v[2:4], v[2:4]),
+      VOP3(VOP3Op.V_FMA_F64, vdst=v[0], src0=RawImm(255), src1=v[2], src2=v[0], literal=lit),
+      v_cvt_u32_f64_e32(v[4], v[0:2]),
+      v_cvt_i32_f64_e32(v[5], v[2:4]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    lo = st.vgpr[0][4]
+    hi = st.vgpr[0][5]
+    result = struct.unpack('<q', struct.pack('<II', lo, hi))[0]
+    self.assertEqual(result, 5000000000, f"Expected 5000000000, got {result}")
+
+
+class TestDS2Addr(unittest.TestCase):
+  """Regression tests for DS_LOAD_2ADDR and DS_STORE_2ADDR instructions.
+  These ops use offset scaling: offset * sizeof(data) for address calculation.
+  Bug: Emulator was using offset*4 for both B32 and B64, but B64 needs offset*8."""
+
+  def test_ds_store_load_2addr_b32(self):
+    """DS_STORE_2ADDR_B32 and DS_LOAD_2ADDR_B32 with offset scaling by 4."""
+    # Store 0x12345678 at offset0=0 (*4=0) and 0xDEADBEEF at offset1=1 (*4=4)
+    # Then load them back
+    instructions = [
+      v_mov_b32_e32(v[10], 0),  # addr base = 0
+      s_mov_b32(s[2], 0x12345678),
+      v_mov_b32_e32(v[0], s[2]),  # data0
+      s_mov_b32(s[2], 0xDEADBEEF),
+      v_mov_b32_e32(v[1], s[2]),  # data1
+      DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x12345678, "v2 should have value from offset 0")
+    self.assertEqual(st.vgpr[0][3], 0xDEADBEEF, "v3 should have value from offset 4")
+
+  def test_ds_store_load_2addr_b32_nonzero_offsets(self):
+    """DS_STORE_2ADDR_B32 with non-zero offsets (offset*4 scaling)."""
+    # Store at offset0=2 (*4=8) and offset1=5 (*4=20)
+    instructions = [
+      v_mov_b32_e32(v[10], 0),  # addr base = 0
+      s_mov_b32(s[2], 0x11111111),
+      v_mov_b32_e32(v[0], s[2]),
+      s_mov_b32(s[2], 0x22222222),
+      v_mov_b32_e32(v[1], s[2]),
+      DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=2, offset1=5),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2], offset0=2, offset1=5),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should have value from offset 8 (2*4)")
+    self.assertEqual(st.vgpr[0][3], 0x22222222, "v3 should have value from offset 20 (5*4)")
+
+  def test_ds_store_load_2addr_b64(self):
+    """DS_STORE_2ADDR_B64 and DS_LOAD_2ADDR_B64 with offset scaling by 8."""
+    # For B64: each value is 8 bytes (2 dwords), offsets scaled by 8
+    # Store 64-bit value at offset0=0 (*8=0) and another at offset1=1 (*8=8)
+    instructions = [
+      v_mov_b32_e32(v[10], 0),  # addr base = 0
+      # First 64-bit value: 0x123456789ABCDEF0
+      s_mov_b32(s[2], 0x9ABCDEF0),
+      v_mov_b32_e32(v[0], s[2]),  # low dword
+      s_mov_b32(s[2], 0x12345678),
+      v_mov_b32_e32(v[1], s[2]),  # high dword
+      # Second 64-bit value: 0xDEADBEEFCAFEBABE
+      s_mov_b32(s[2], 0xCAFEBABE),
+      v_mov_b32_e32(v[2], s[2]),  # low dword
+      s_mov_b32(s[2], 0xDEADBEEF),
+      v_mov_b32_e32(v[3], s[2]),  # high dword
+      DS(DSOp.DS_STORE_2ADDR_B64, addr=v[10], data0=v[0], data1=v[2], vdst=v[0], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # v4,v5 = first 64-bit value from offset 0
+    self.assertEqual(st.vgpr[0][4], 0x9ABCDEF0, "v4 should have low dword of first value")
+    self.assertEqual(st.vgpr[0][5], 0x12345678, "v5 should have high dword of first value")
+    # v6,v7 = second 64-bit value from offset 8 (1*8)
+    self.assertEqual(st.vgpr[0][6], 0xCAFEBABE, "v6 should have low dword of second value")
+    self.assertEqual(st.vgpr[0][7], 0xDEADBEEF, "v7 should have high dword of second value")
+
+  def test_ds_2addr_b64_no_overlap(self):
+    """DS_LOAD_2ADDR_B64 with adjacent offsets should not overlap.
+    Regression test: offset1=1 should access bytes 8-15, not overlap with offset0=0 (bytes 0-7)."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      # Store 4 distinct dwords at addresses 0,4,8,12 using regular DS_STORE
+      s_mov_b32(s[2], 0x11111111),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_mov_b32(s[2], 0x22222222),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=4),
+      s_mov_b32(s[2], 0x33333333),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=8),
+      s_mov_b32(s[2], 0x44444444),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=12),
+      s_waitcnt(lgkmcnt=0),
+      # Load with DS_LOAD_2ADDR_B64: offset0=0 should get 0-7, offset1=1 should get 8-15
+      DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # v4,v5 from addr 0-7: 0x11111111, 0x22222222
+    self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should be 0x11111111")
+    self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should be 0x22222222")
+    # v6,v7 from addr 8-15: 0x33333333, 0x44444444
+    self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should be 0x33333333")
+    self.assertEqual(st.vgpr[0][7], 0x44444444, "v7 should be 0x44444444")

extra/assembly/amd/test/test_formats.py

@@ -1,8 +1,9 @@
 #!/usr/bin/env python3
 """Test MUBUF, MTBUF, MIMG, EXP, DS formats against LLVM."""
 import unittest
-from extra.assembly.amd.autogen.rdna3 import *
-from extra.assembly.amd.dsl import encode_src
+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."""
@@ -328,5 +329,79 @@ class TestVOP3Literal(unittest.TestCase):
     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

@@ -2,7 +2,7 @@
 # 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 import *
+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

extra/assembly/amd/test/test_integration.py

@@ -1,7 +1,7 @@
 #!/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 import *
+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
 

extra/assembly/amd/test/test_llvm.py

@@ -1,98 +1,76 @@
 #!/usr/bin/env python3
-"""Test RDNA3 assembler/disassembler against LLVM test vectors."""
+"""Test RDNA3/RDNA4 assembler/disassembler against LLVM test vectors."""
 import unittest, re, subprocess
 from tinygrad.helpers import fetch
-from extra.assembly.amd.autogen.rdna3 import *
-from extra.assembly.amd.asm import asm
 from extra.assembly.amd.test.helpers import get_llvm_mc
 
 LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/main/llvm/test/MC/AMDGPU"
 
-# Format info: (filename, format_class, op_enum)
-LLVM_TEST_FILES = {
-  # Scalar ALU
-  'sop1': ('gfx11_asm_sop1.s', SOP1, SOP1Op),
-  'sop2': ('gfx11_asm_sop2.s', SOP2, SOP2Op),
-  'sopp': ('gfx11_asm_sopp.s', SOPP, SOPPOp),
-  'sopk': ('gfx11_asm_sopk.s', SOPK, SOPKOp),
-  'sopc': ('gfx11_asm_sopc.s', SOPC, SOPCOp),
-  # Vector ALU
-  'vop1': ('gfx11_asm_vop1.s', VOP1, VOP1Op),
-  'vop2': ('gfx11_asm_vop2.s', VOP2, VOP2Op),
-  'vopc': ('gfx11_asm_vopc.s', VOPC, VOPCOp),
-  'vop3': ('gfx11_asm_vop3.s', VOP3, VOP3Op),
-  'vop3p': ('gfx11_asm_vop3p.s', VOP3P, VOP3POp),
-  'vop3sd': ('gfx11_asm_vop3.s', VOP3SD, VOP3SDOp),  # VOP3SD shares file with VOP3
-  'vinterp': ('gfx11_asm_vinterp.s', VINTERP, VINTERPOp),
-  'vopd': ('gfx11_asm_vopd.s', VOPD, VOPDOp),
-  'vopcx': ('gfx11_asm_vopcx.s', VOPC, VOPCOp),  # VOPCX uses VOPC format
-  # VOP3 promotions (VOP1/VOP2/VOPC promoted to VOP3 encoding)
-  'vop3_from_vop1': ('gfx11_asm_vop3_from_vop1.s', VOP3, VOP3Op),
-  'vop3_from_vop2': ('gfx11_asm_vop3_from_vop2.s', VOP3, VOP3Op),
-  'vop3_from_vopc': ('gfx11_asm_vop3_from_vopc.s', VOP3, VOP3Op),
-  'vop3_from_vopcx': ('gfx11_asm_vop3_from_vopcx.s', VOP3, VOP3Op),
-  # Memory
-  'ds': ('gfx11_asm_ds.s', DS, DSOp),
-  'smem': ('gfx11_asm_smem.s', SMEM, SMEMOp),
-  'flat': ('gfx11_asm_flat.s', FLAT, FLATOp),
-  'mubuf': ('gfx11_asm_mubuf.s', MUBUF, MUBUFOp),
-  'mtbuf': ('gfx11_asm_mtbuf.s', MTBUF, MTBUFOp),
-  'mimg': ('gfx11_asm_mimg.s', MIMG, MIMGOp),
-  # WMMA (matrix multiply)
-  'wmma': ('gfx11_asm_wmma.s', VOP3P, VOP3POp),
-  # Additional features
-  'vop3_features': ('gfx11_asm_vop3_features.s', VOP3, VOP3Op),
-  'vop3p_features': ('gfx11_asm_vop3p_features.s', VOP3P, VOP3POp),
-  'vopd_features': ('gfx11_asm_vopd_features.s', VOPD, VOPDOp),
-  # Alias files (alternative mnemonics)
-  'vop3_alias': ('gfx11_asm_vop3_alias.s', VOP3, VOP3Op),
-  'vop3p_alias': ('gfx11_asm_vop3p_alias.s', VOP3P, VOP3POp),
-  'vopc_alias': ('gfx11_asm_vopc_alias.s', VOPC, VOPCOp),
-  'vopcx_alias': ('gfx11_asm_vopcx_alias.s', VOPC, VOPCOp),
-  'vinterp_alias': ('gfx11_asm_vinterp_alias.s', VINTERP, VINTERPOp),
-  'smem_alias': ('gfx11_asm_smem_alias.s', SMEM, SMEMOp),
-  'mubuf_alias': ('gfx11_asm_mubuf_alias.s', MUBUF, MUBUFOp),
-  'mtbuf_alias': ('gfx11_asm_mtbuf_alias.s', MTBUF, MTBUFOp),
+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',
 }
 
-def parse_llvm_tests(text: str) -> list[tuple[str, bytes]]:
+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))):
-      # Match GFX11, W32, or W64 encodings (all valid for gfx11)
-      # Format 1: "// GFX11: v_foo ... ; encoding: [0x01,0x02,...]"
-      # Format 2: "// GFX11: [0x01,0x02,...]" (used by DS, older files)
-      if m := re.search(r'(?:GFX11|W32|W64)[^:]*:.*?encoding:\s*\[(.*?)\]', lines[j]):
+      if m := re.search(pattern, lines[j]):
         hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
-      elif m := re.search(r'(?:GFX11|W32|W64)[^:]*:\s*\[(0x[0-9a-fA-F,x\s]+)\]', lines[j]):
+      elif m := re.search(pattern2, lines[j]):
         hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
-      else:
-        continue
+      else: continue
       if hex_bytes:
         try: tests.append((asm_text, bytes.fromhex(hex_bytes)))
         except ValueError: pass
       break
   return tests
 
-def try_assemble(text: str):
-  """Try to assemble instruction text, return bytes or None on failure."""
-  try: return asm(text).to_bytes()
-  except: return None
-
-def compile_asm_batch(instrs: list[str]) -> list[bytes]:
+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 []
-  asm_text = ".text\n" + "\n".join(instrs) + "\n"
-  result = subprocess.run(
-    [get_llvm_mc(), '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
-    input=asm_text, capture_output=True, text=True, timeout=30)
-  if result.returncode != 0: raise RuntimeError(f"llvm-mc batch failed: {result.stderr.strip()}")
-  # Parse all encodings from output
+  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
@@ -102,100 +80,127 @@ def compile_asm_batch(instrs: list[str]) -> list[bytes]:
   if len(results) != len(instrs): raise RuntimeError(f"expected {len(instrs)} encodings, got {len(results)}")
   return results
 
-class TestLLVM(unittest.TestCase):
-  """Test assembler and disassembler against all LLVM test vectors."""
+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 setUpClass(cls):
-    for name, (filename, _, _) in LLVM_TEST_FILES.items():
+  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.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] = []
 
-# Generate test methods dynamically for each format
-def _make_asm_test(name):
-  def test(self):
-    passed, failed, skipped = 0, 0, 0
-    for asm_text, expected in self.tests.get(name, []):
-      result = try_assemble(asm_text)
-      if result is None: skipped += 1
-      elif result == expected: passed += 1
-      else: failed += 1
-    print(f"{name.upper()} asm: {passed} passed, {failed} failed, {skipped} skipped")
-    self.assertEqual(failed, 0)
-  return test
+  def _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}")
 
-def _make_disasm_test(name):
-  def test(self):
-    _, fmt_cls, op_enum = LLVM_TEST_FILES[name]
-    # VOP3SD opcodes that share encoding with VOP3 (only for vop3sd test, not vopc promotions)
-    vop3sd_opcodes = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
-    is_vopc_promotion = name in ('vop3_from_vopc', 'vop3_from_vopcx')
-    undocumented = {'smem': {34, 35}, 'sopk': {22, 23}, 'sopp': {8, 58, 59}}
+    # 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}'
 
-    # First pass: decode all instructions and collect disasm strings
-    to_test: list[tuple[str, bytes, str | None, str | None]] = []  # (asm_text, data, disasm_str, error)
-    skipped = 0
+    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
-      temp_inst = fmt_cls.from_bytes(data)
-      temp_op = temp_inst._values.get('op', 0)
-      temp_op = temp_op.val if hasattr(temp_op, 'val') else temp_op
-      if temp_op in undocumented.get(name, set()): skipped += 1; continue
-      if name == 'sopp':
-        simm16 = temp_inst._values.get('simm16', 0)
-        simm16 = simm16.val if hasattr(simm16, 'val') else simm16
-        sopp_no_imm = {48, 54, 53, 55, 60, 61, 62}
-        if temp_op in sopp_no_imm and simm16 != 0: skipped += 1; continue
+      if not matches_encoding(data, fmt_cls): continue
       try:
-        if fmt_cls.__name__ in ('VOP3', 'VOP3SD'):
-          temp = VOP3.from_bytes(data)
-          op_val = temp._values.get('op', 0)
-          op_val = op_val.val if hasattr(op_val, 'val') else op_val
-          is_vop3sd = (op_val in vop3sd_opcodes) and not is_vopc_promotion
-          decoded = VOP3SD.from_bytes(data) if is_vop3sd else VOP3.from_bytes(data)
-          if is_vop3sd: VOP3SDOp(op_val)
-          else: VOP3Op(op_val)
-        else:
-          decoded = fmt_cls.from_bytes(data)
-          op_val = decoded._values.get('op', 0)
-          op_val = op_val.val if hasattr(op_val, 'val') else op_val
-          op_enum(op_val)
+        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(), None))
+        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}"))
 
-    # Batch compile all disasm strings with single llvm-mc call
     disasm_strs = [(i, t[2]) for i, t in enumerate(to_test) if t[2] is not None]
-    llvm_results = compile_asm_batch([s for _, s in disasm_strs]) if disasm_strs else []
-    llvm_map = {i: llvm_results[j] for j, (i, _) in enumerate(disasm_strs)}
+    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)}
 
-    # Match results back
-    passed, failed = 0, 0
-    failures: list[str] = []
+    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 is not None and llvm_bytes == data: passed += 1
-        elif llvm_bytes is not None: failed += 1; failures.append(f"'{disasm_str}': expected={data.hex()} got={llvm_bytes.hex()}")
+        if llvm_bytes == data: passed += 1
+        else: failed += 1; failures.append(f"'{disasm_str}': expected={data.hex()} got={llvm_bytes.hex()}")
 
-    print(f"{name.upper()} disasm: {passed} passed, {failed} failed" + (f", {skipped} skipped" if skipped else ""))
-    if failures[:10]: print("  " + "\n  ".join(failures[:10]))
-    self.assertEqual(failed, 0)
+    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 LLVM_TEST_FILES:
-  setattr(TestLLVM, f'test_{name}_asm', _make_asm_test(name))
-  setattr(TestLLVM, f'test_{name}_disasm', _make_disasm_test(name))
+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()
+if __name__ == "__main__": unittest.main()

extra/assembly/amd/test/test_pcode.py

@@ -1,9 +1,11 @@
 #!/usr/bin/env python3
 """Tests for the RDNA3 pseudocode DSL."""
 import unittest
-from extra.assembly.amd.pcode import (Reg, TypedView, SliceProxy, ExecContext, compile_pseudocode, _expr, MASK32, MASK64,
+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):
@@ -227,17 +229,18 @@ 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_lane, even when VCC=0 (no scaling needed).
-    Bug: when VCC._val == vcc (both 0), vcc_lane wasn't returned, so VCC bits weren't written.
+    """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 = 0x3f800000  # 1.0
-    s1 = 0x40400000  # 3.0
-    s2 = 0x3f800000  # 1.0 (numerator)
-    result = _VOP3SDOp_V_DIV_SCALE_F32(s0, s1, s2, 0, 0, 0, 0, 0xffffffff, 0, None, {})
-    # Must always have vcc_lane in result
-    self.assertIn('vcc_lane', result, "V_DIV_SCALE_F32 must always return vcc_lane")
-    self.assertEqual(result['vcc_lane'], 0, "vcc_lane should be 0 when no scaling needed")
+    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.
@@ -246,18 +249,22 @@ class TestPseudocodeRegressions(unittest.TestCase):
     signal_nan = 0x7f800001  # signaling NaN: exponent=255, bit22=0
     # Test quiet NaN detection (bit 1 in mask)
     s1_quiet = 0b0000000010  # bit 1 = quiet NaN
-    result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None, {})
-    self.assertEqual(result['vcc_lane'], 1, "Should detect quiet NaN with quiet NaN mask")
+    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
-    result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None, {})
-    self.assertEqual(result['vcc_lane'], 1, "Should detect signaling NaN with signaling NaN mask")
+    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
-    result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None, {})
-    self.assertEqual(result['vcc_lane'], 0, "Quiet NaN should not 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
-    result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None, {})
-    self.assertEqual(result['vcc_lane'], 0, "Signaling NaN should not 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.

extra/assembly/amd/test/test_pdf_parser.py

@@ -1,10 +1,7 @@
 #!/usr/bin/env python3
 """Test that PDF parser correctly extracts format fields."""
 import unittest, os
-from extra.assembly.amd.autogen.rdna3 import (
-  SOP1, SOP2, SOPK, SOPP, VOP1, VOP2, VOP3SD, VOPC, FLAT, VOPD,
-  SOP1Op, SOP2Op, VOP1Op, VOP3Op
-)
+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 = {

extra/assembly/amd/test/test_rdna3_asm.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import unittest, subprocess
-from extra.assembly.amd.autogen.rdna3 import *
+from extra.assembly.amd.autogen.rdna3.ins import *
 from extra.assembly.amd.test.helpers import get_llvm_mc
 
 def llvm_assemble(asm: str) -> bytes:

extra/assembly/amd/test/test_roundtrip.py

@@ -1,90 +1,37 @@
 #!/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.autogen.rdna3 import *
 from extra.assembly.amd.dsl import Inst
-from extra.assembly.amd.asm import asm
-from extra.assembly.amd.asm import detect_format
 from extra.assembly.amd.test.helpers import get_llvm_mc, get_llvm_objdump
 
-def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
-  """Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
-  old_stdout = sys.stdout
-  sys.stdout = io.StringIO()
-  compiler.disassemble(lib)
-  output = sys.stdout.getvalue()
-  sys.stdout = old_stdout
-
-  results = []
-  for line in output.splitlines():
-    if '//' not in line: continue
-    instr = line.split('//')[0].strip()
-    if not instr: continue
-    comment = line.split('//')[1].strip()
-    if ':' not in comment: continue
-    hex_str = comment.split(':')[1].strip().split()[0]
-    try:
-      machine_bytes = bytes.fromhex(hex_str)[::-1]  # big-endian to little-endian
-      results.append((instr, machine_bytes))
-    except ValueError:
-      continue
-  return results
-
-def compile_asm(instr: str, compiler=None) -> bytes:
-  """Compile a single instruction with llvm-mc and return the machine code bytes."""
-  llvm_mc = get_llvm_mc()
-  result = subprocess.run(
-    [llvm_mc, '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
-    input=f".text\n{instr}\n", capture_output=True, text=True)
-  if result.returncode != 0: raise RuntimeError(f"llvm-mc failed for '{instr}': {result.stderr.strip()}")
-  # Parse encoding: [0x01,0x39,0x0a,0x7e]
-  for line in result.stdout.split('\n'):
-    if 'encoding:' in line:
-      enc = line.split('encoding:')[1].strip()
-      if enc.startswith('[') and enc.endswith(']'):
-        hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
-        return bytes.fromhex(hex_vals)
-  raise RuntimeError(f"no encoding found in llvm-mc output for: {instr}")
-
-def compile_asm_batch(instrs: list[str]) -> list[bytes]:
+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 []
-  llvm_mc = get_llvm_mc()
-  src = ".text\n" + "\n".join(instrs) + "\n"
-  result = subprocess.run(
-    [llvm_mc, '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
-    input=src, capture_output=True, text=True)
+  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()}")
-  # Parse all encodings in order
   encodings = []
   for line in result.stdout.split('\n'):
     if 'encoding:' in line:
       enc = line.split('encoding:')[1].strip()
       if enc.startswith('[') and enc.endswith(']'):
-        hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
-        encodings.append(bytes.fromhex(hex_vals))
+        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], compiler) -> list[str]:
+def compile_and_disasm_batch(instrs: list[str], mcpu: str = 'gfx1100') -> list[str]:
   """Compile instructions with LLVM and get LLVM's disassembly."""
-  import tempfile, os
+  import tempfile
   if not instrs: return []
-  # Build assembly source with all instructions
-  src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n"
-  src += "\n".join(f"  {instr}" for instr in instrs) + "\n"
-  # Use llvm-mc to assemble to object file
+  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', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-filetype=obj', '-o', obj_path],
-      input=src, capture_output=True, text=True)
+    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()}")
-    # Disassemble with llvm-objdump
-    result = subprocess.run([get_llvm_objdump(), '-d', '--mcpu=gfx1100', obj_path], capture_output=True, text=True)
+    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()}")
-    # Parse disassembly output
     results: list[str] = []
     for line in result.stdout.splitlines():
       if '//' not in line: continue
@@ -94,127 +41,143 @@ def compile_and_disasm_batch(instrs: list[str], compiler) -> list[str]:
   finally:
     os.unlink(obj_path)
 
-class TestTinygradKernelRoundtrip(unittest.TestCase):
-  """Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
+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:
-    1. decode -> reencode matches original bytes
-    2. asm(disasm()) matches LLVM output
-    3. our disasm() matches LLVM's disassembly string exactly
-    """
+    """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('gfx1100')
+    compiler = HIPCompiler(self.mcpu)
 
-    # First pass: decode all instructions and collect info
-    decoded_instrs: list[tuple] = []  # list of (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err)
+    # 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:]
-        fmt = detect_format(remaining)
-        if fmt is None:
-          decoded_instrs.append((ki, offset, None, None, None, False, "no format"))
-          offset += 4
-          continue
+        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
+        if len(remaining) < base_size: break
 
         try:
-          decoded = fmt.from_bytes(remaining)  # pass all remaining bytes so from_bytes can read literal
-          size = decoded.size()  # actual size including literal
+          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: str | None = None if decode_ok else f"orig={orig_bytes.hex()} reenc={reencoded.hex()}"
+          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 - skip unknown opcodes (op_X) that LLVM can't compile
-    asm_test_instrs: list[tuple[int, str]] = []  # (idx, our_disasm) for asm test
-    disasm_test_instrs: list[tuple[int, str]] = []  # (idx, our_disasm) for disasm comparison test
-
+    # 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
-      # Skip unknown opcodes and malformed instructions for both tests
       if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm): continue
       asm_test_instrs.append((idx, our_disasm))
-      disasm_test_instrs.append((idx, our_disasm))
 
     # Batch compile for asm test
-    asm_llvm_results = compile_asm_batch([d for _, d in asm_test_instrs])
+    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 disasm_test_instrs], compiler)
-    disasm_llvm_map = {idx: result for (idx, _), result in zip(disasm_test_instrs, disasm_llvm_results)}
+    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)}
 
-    # Now evaluate 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: list[str] = []
-    asm_failures: list[str] = []
-    disasm_failures: list[str] = []
+    decode_failures, asm_failures, disasm_failures = [], [], []
 
     for idx, (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err) in enumerate(decoded_instrs):
-      # Decode test
-      if decode_ok:
-        decode_passed += 1
-      elif decode_err == "no format":
-        decode_skipped += 1
+      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}")
 
-      # Asm test
       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
+          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 comparison test
-      if our_disasm is None:
-        disasm_skipped += 1
-      elif idx in disasm_llvm_map:
-        llvm_disasm = disasm_llvm_map[idx]
-        if our_disasm == llvm_disasm:
-          disasm_passed += 1
-        else:
-          disasm_failed += 1
-          disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm}'")
-      else:
         disasm_skipped += 1
 
-    print(f"decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
-    print(f"asm vs llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
-    print(f"disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
+    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]))
-    # Note: disasm string comparison is informational only - formatting differences between LLVM versions are expected
 
-  # Basic unary ops
+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())
@@ -222,42 +185,62 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
   def test_sin(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).sin())
   def test_sqrt(self): self._test_kernel_roundtrip(lambda T: T([1.0, 4.0, 9.0]).sqrt())
   def test_recip(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
-
-  # Binary ops
   def test_add(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
   def test_sub(self): self._test_kernel_roundtrip(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
   def test_mul(self): self._test_kernel_roundtrip(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
   def test_div(self): self._test_kernel_roundtrip(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
   def test_max_binary(self): self._test_kernel_roundtrip(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
-
-  # Reductions
   def test_sum_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).sum())
   def test_max_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).max())
   def test_mean_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(32).mean())
-
-  # Matmul
   def test_gemm_4x4(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4) @ T.empty(4, 4))
   def test_gemv(self): self._test_kernel_roundtrip(lambda T: T.empty(1, 16) @ T.empty(16, 16))
-
-  # Complex ops
   def test_softmax(self): self._test_kernel_roundtrip(lambda T: T.empty(16).softmax())
   def test_layernorm(self): self._test_kernel_roundtrip(lambda T: T.empty(8, 8).layernorm())
-
-  # Memory patterns
   def test_contiguous(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
   def test_reshape(self): self._test_kernel_roundtrip(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
   def test_expand(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
-
-  # Cast ops
   def test_cast_int(self): self._test_kernel_roundtrip(lambda T: T.empty(16).int().float())
   def test_cast_half(self): self._test_kernel_roundtrip(lambda T: T.empty(16).half().float())
-
-  # Comparison ops
   def test_cmp_lt(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
   def test_where(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
-
-  # Fused ops
   def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
 
+@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/gemm/asm/cdna/gemm.s

@@ -1,13 +1,7 @@
-.text
-.section	.text.
-.global	gemm
-.p2align	8
-.type	gemm,@function
-
-gemm:
 	// ** global buffers
 	s_load_dwordx2  s[28:29], s[0:1], 0x0    // C
-	s_load_dwordx4  s[32:35], s[0:1], 0x8    // A, B
+	s_load_dwordx2  s[34:35], s[0:1], 0x08   // A
+	s_load_dwordx2  s[32:33], s[0:1], 0x10   // B
 	// ** others kernel args
 	s_load_dword    s24, s[0:1], 0x18        // N
 	s_load_dword    s54, s[0:1], 0x1C        // num work groups
@@ -221,11 +215,11 @@ gemm:
 	s_mul_hi_u32 s87, s86, s40                                 // 000000003288: 96572856
 	s_mul_i32 s86, s86, s40                                    // 00000000328C: 92562856
 	s_and_b32 s84, s50, 0x8000                                 // 000000003290: 8654FF32 00008000
-	s_cbranch_scc1 label_GSUC_A                                // 000000003298: BF850003
+	s_cbranch_scc1 skip_offset_A                                // 000000003298: BF850003
 	s_mul_hi_u32 s85, 64, s6                                   // 00000000329C: 965506C0
 	s_mul_i32 s84, 64, s6                                      // 0000000032A0: 925406C0
 
-label_GSUC_A:
+skip_offset_A:
 	s_add_u32 s86, s86, s84                                    // 000000003330: 80565456
 	s_addc_u32 s87, s87, s85                                   // 000000003334: 82575557
 	s_mov_b64 s[60:61], 1                                      // 000000003338: BEBC0181
@@ -259,11 +253,11 @@ label_GSUC_A:
 	s_mul_hi_u32 s87, s86, s42                                 // 0000000033B4: 96572A56
 	s_mul_i32 s86, s86, s42                                    // 0000000033B8: 92562A56
 	s_and_b32 s84, s50, 0x8000                                 // 0000000033BC: 8654FF32 00008000
-	s_cbranch_scc1 label_GSUC_B                                // 0000000033C4: BF850003
+	s_cbranch_scc1 skip_offset_B                                // 0000000033C4: BF850003
 	s_mul_hi_u32 s85, 64, s6                                   // 0000000033C8: 965506C0
 	s_mul_i32 s84, 64, s6                                      // 0000000033CC: 925406C0
 
-label_GSUC_B:
+skip_offset_B:
 	s_add_u32 s86, s86, s84                                    // 00000000345C: 80565456
 	s_addc_u32 s87, s87, s85                                   // 000000003460: 82575557
 	s_mov_b64 s[62:63], 1                                      // 000000003464: BEBE0181
@@ -308,8 +302,6 @@ label_GSUC_B:
 	s_and_b32 s87, s10, 0xe000                                 // 0000000035A4: 8657FF0A 0000E000
 	s_and_b32 s10, s10, 0xff                                   // 0000000035AC: 860AFF0A 000000FF
 	s_mov_b32 s84, s10                                         // 0000000035B4: BED4000A
-
-label_beginStaggerUIter:
 	s_lshl_b32 s85, s84, s86                                   // 0000000035B8: 8E555654
 	s_cmp_ge_u32 s13, s85                                      // 0000000035BC: BF09550D
 	s_sub_u32 s85, s84, 1                                      // 0000000035CC: 80D58154
@@ -344,7 +336,7 @@ label_beginStaggerUIter:
 	s_cselect_b32 s58, s62, -1                                 // 0000000036A4: 853AC13E
 	s_add_u32 s51, s51, 2                                      // 0000000036A8: 80338233
 	s_cmp_eq_u32 s12, 0                                        // 0000000036AC: BF06800C
-	s_cbranch_scc1 label_ShadowInitStart                       // 0000000036B0: BF850092
+	s_cbranch_scc1 init_output_buffers                       // 0000000036B0: BF850092
 	s_mov_b32 m0, s46                                          // 0000000036B4: BEFC002E
 	buffer_load_dwordx4 v0, s[52:55], 0 offen lds              // 0000000036B8: E05D1000 800D0000
 	s_add_u32 m0, m0, 0x1040                                   // 0000000036C0: 807CFF7C 00001040
@@ -431,7 +423,7 @@ label_beginStaggerUIter:
 	s_cmp_eq_u32 s63, 0                                        // 0000000038F4: BF06803F
 	s_cselect_b32 s58, s62, -1                                 // 0000000038F8: 853AC13E
 
-label_ShadowInitStart:
+init_output_buffers:
 	s_mov_b64 s[16:17], s[28:29]                               // 0000000038FC: BE90011C
 	s_mov_b32 s18, 0x80000000                                  // 000000003900: BE9200FF 80000000
 	s_mov_b32 s19, 0x20000                                     // 000000003908: BE9300FF 00020000
@@ -476,12 +468,10 @@ label_ShadowInitStart:
 	s_lshl_b64 s[84:85], s[84:85], 2                           // 0000000039C4: 8ED48254
 	s_add_u32 s16, s16, s84                                    // 0000000039C8: 80105410
 	s_addc_u32 s17, s17, s85                                   // 0000000039CC: 82115511
-
-label_NoBranch_T8JHFHKM7BO5OHXW:
 	s_xor_b32 s46, s48, s46                                    // 0000000039F0: 882E2E30
 	s_xor_b32 s47, s49, s47                                    // 0000000039F4: 882F2F31
 	s_cmp_eq_u32 s12, 1                                        // 0000000039F8: BF06810C
-	s_cbranch_scc1 label_skipPGR2                              // 0000000039FC: BF850040
+	s_cbranch_scc1 after_prefetch                              // 0000000039FC: BF850040
 	s_mov_b32 m0, s46                                          // 000000003A00: BEFC002E
 	buffer_load_dwordx4 v0, s[52:55], 0 offen lds              // 000000003A04: E05D1000 800D0000
 	s_add_u32 m0, m0, 0x1040                                   // 000000003A0C: 807CFF7C 00001040
@@ -517,7 +507,7 @@ label_NoBranch_T8JHFHKM7BO5OHXW:
 	s_xor_b32 s46, s48, s46                                    // 000000003AF8: 882E2E30
 	s_xor_b32 s47, s49, s47                                    // 000000003AFC: 882F2F31
 
-label_skipPGR2:
+after_prefetch:
 	s_waitcnt vmcnt(24)                                        // 000000003B00: BF8C4F78
 	s_barrier                                                  // 000000003B04: BF8A0000
 	ds_read_b128 v[4:7], v2                                    // 000000003B08: D9FE0000 04000002
@@ -539,14 +529,12 @@ label_skipPGR2:
 	ds_read_b128 v[92:95], v3 offset:768                       // 000000003B80: D9FE0300 5C000003
 	ds_read_b128 v[96:99], v3 offset:896                       // 000000003B88: D9FE0380 60000003
 	s_waitcnt lgkmcnt(0)                                       // 000000003B90: BF8CC07F
-
-label_openLoopL:
 	s_cmp_eq_u32 s12, 1                                        // 000000003B94: BF06810C
-	s_cbranch_scc1 label_toPGR1                                // 000000003B98: BF8502E5
+	s_cbranch_scc1 final_compute                                // 000000003B98: BF8502E5
 	s_cmp_le_u32 s12, 2                                        // 000000003B9C: BF0B820C
-	s_cbranch_scc1 label_LoopEndL                              // 000000003BA0: BF85019E
+	s_cbranch_scc1 loop_epilogue                              // 000000003BA0: BF85019E
 
-label_LoopBeginL:
+main_loop:
 	v_mfma_f32_16x16x32_bf16 a[0:3], v[68:71], v[4:7], a[0:3]  // 000000003BA4: D3B58000 04020944
 	ds_read_b128 v[36:39], v2 offset:64                        // 000000003BAC: D9FE0040 24000002
 	v_mfma_f32_16x16x32_bf16 a[4:7], v[68:71], v[8:11], a[4:7] // 000000003BB4: D3B58004 04121144
@@ -770,9 +758,9 @@ label_LoopBeginL:
 	s_cmp_eq_i32 s12, 2                                        // 000000004208: BF00820C
 	s_waitcnt lgkmcnt(0)                                       // 00000000420C: BF8CC07F
 	v_mfma_f32_16x16x32_bf16 a[252:255], v[128:131], v[64:67], a[252:255]// 000000004210: D3B580FC 07F28180
-	s_cbranch_scc0 label_LoopBeginL                            // 000000004218: BF84FE62
+	s_cbranch_scc0 main_loop                            // 000000004218: BF84FE62
 
-label_LoopEndL:
+loop_epilogue:
 	v_mfma_f32_16x16x32_bf16 a[0:3], v[68:71], v[4:7], a[0:3]  // 00000000421C: D3B58000 04020944
 	ds_read_b128 v[36:39], v2 offset:64                        // 000000004224: D9FE0040 24000002
 	v_mfma_f32_16x16x32_bf16 a[4:7], v[68:71], v[8:11], a[4:7] // 00000000422C: D3B58004 04121144
@@ -939,7 +927,7 @@ label_LoopEndL:
 	v_mfma_f32_16x16x32_bf16 a[248:251], v[128:131], v[60:63], a[248:251]// 000000004720: D3B580F8 07E27980
 	v_mfma_f32_16x16x32_bf16 a[252:255], v[128:131], v[64:67], a[252:255]// 000000004728: D3B580FC 07F28180
 
-label_toPGR1:
+final_compute:
 	s_and_b32 s8, s50, 0x3fff                                  // 000000004730: 8608FF32 00003FFF
 	s_and_b32 s84, 0xff, s24                                   // 000000004750: 865418FF 000000FF
 	s_add_u32 s85, -1, s14                                     // 000000004758: 80550EC1
@@ -1095,7 +1083,6 @@ label_toPGR1:
 	v_mfma_f32_16x16x32_bf16 a[248:251], v[128:131], v[60:63], a[248:251]// 000000004BFC: D3B580F8 07E27980
 	v_mfma_f32_16x16x32_bf16 a[252:255], v[128:131], v[64:67], a[252:255]// 000000004C04: D3B580FC 07F28180
 
-label_toPGR1end_OptNLL:
 	v_lshrrev_b32_e32 v4, 6, v134                              // 000000004C0C: 20090C86
 	v_lshrrev_b32_e32 v5, 1, v4                                // 000000004C10: 200A0881
 	v_mul_lo_u32 v5, 16, v5                                    // 000000004C14: D2850005 00020A90
@@ -1114,7 +1101,6 @@ label_toPGR1end_OptNLL:
 	s_mul_i32 s8, 0x100, s3                                    // 000000004C64: 920803FF 00000100
 	v_add_u32_e32 v1, s8, v1                                   // 000000004C6C: 68020208
 
-label_GW_B0_E0:
 	v_add_lshl_u32 v11, v3, v0, 1                              // 000000004C70: D1FE000B 02060103
 	v_accvgpr_read_b32 v16, a0                                 // 000000004C78: D3D84010 18000100
 	v_accvgpr_read_b32 v17, a4                                 // 000000004C80: D3D84011 18000104
@@ -1633,86 +1619,4 @@ label_GW_B0_E0:
 	s_addc_u32 s17, s17, 0                                     // 000000005B14: 82118011
 	buffer_store_dwordx4 v[40:43], v11, s[16:19], 0 offen nt   // 000000005B18: E07E1000 8004280B
 	s_nop 0                                                    // 000000005B20: BF800000
-
-end:
 	s_endpgm                                                   // 00000001F5D0: BF810000
-
-.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 SGPR enables (descriptor bits >448) ----
-  .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 specific (RSRC3) ----
-  .amdhsa_accum_offset 248
-  .amdhsa_uses_dynamic_stack 0
-  .amdhsa_tg_split 0
-
-.end_amdhsa_kernel
-
-.amdgpu_metadata
----
-amdhsa.kernels:
-  - .args:
-      - .address_space:  global
-        .name:           C
-        .offset:         0
-        .size:           8
-        .value_kind:     global_buffer
-        .value_type:     bf16
-      - .address_space:  global
-        .name:           B
-        .offset:         8
-        .size:           8
-        .value_kind:     global_buffer
-        .value_type:     bf16
-      - .address_space:  global
-        .name:           A
-        .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
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 32
-    .max_flat_workgroup_size: 256
-    .name:           gemm
-    .private_segment_fixed_size: 0
-    .sgpr_count:     88
-    .sgpr_spill_count: 0
-    .symbol:         gemm.kd
-    .vgpr_count:     248
-    .vgpr_spill_count: 0
-    .wavefront_size: 64
-amdhsa.version:
-  - 1
-  - 0
-...
-    .end_amdgpu_metadata

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

@@ -48,11 +48,11 @@ ast = sched[-1].ast
 # assembly gemm
 @track_rewrites(name=lambda ret: TracingKey(ret.name, (ret.function_name,), ret))
 def get_asm_prg() -> ProgramSpec:
-  src = fp.read_text()
+  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(B).uop.buffer, from_torch(A).uop.buffer], fixedvars={"SZ":N, "NUM_WG":NUM_WG},
+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):

extra/gemm/asm/unpack_kd.py

@@ -1,12 +1,12 @@
-# unpack the complete kernel descriptor of an amdgpu ELF of for gfx950
+# 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
+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(fp:=pathlib.Path(__file__).parent/"lib", "rb") as f:
+with open(sys.argv[1], "rb") as f:
   lib = f.read()
 
 image, sections, relocs = elf_loader(lib)
@@ -49,7 +49,7 @@ 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
+# rsrc 3 (gfx950)
 
 accum_offset_raw = bits(pgm_rsrc3, 0, 5)
 assert_zero(pgm_rsrc3, 6, 15)
@@ -169,10 +169,10 @@ 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)
 

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/test/hwtest.py

@@ -3,18 +3,21 @@
 
 import numpy as np
 import unittest
-import subprocess, struct, math, textwrap
+import subprocess, struct, math, textwrap, functools
 from tinygrad import Tensor, dtypes, Device, UOp
-from tinygrad.uop.ops import Ops
+from tinygrad.uop.ops import Ops, KernelInfo
 from tinygrad.helpers import getenv
-from tinygrad.runtime.support.compiler_amd import amdgpu_disassemble
-from tinygrad.renderer import ProgramSpec
-from tinygrad.engine.realize import CompiledRunner
 
-from extra.assembly.amd.autogen.rdna3 import *
+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 [
@@ -26,11 +29,9 @@ def get_output(asm:list, n_threads:int=1, vdst:VGPR=v[1]):
     global_store_b32(addr=v[0], data=vdst, saddr=s[0:1]),
     s_endpgm()
   ])
-  prg = ProgramSpec("test", template.replace("fn_name", "test").replace("INSTRUCTION", textwrap.dedent(src)), Device.DEFAULT, UOp(Ops.SINK),
-                    global_size=[1, 1, 1], local_size=[n_threads, 1, 1], globals=[0])
-  car = CompiledRunner(prg)
-  if getenv("PRINT_ASM"): amdgpu_disassemble(car.lib)
-  car([out.uop.buffer], {}, wait=True)
+  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]

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_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/mockgpu/amd/amdgpu.py

@@ -187,7 +187,7 @@ class PM4Executor(AMDQueue):
       if st <= prg_addr < st+sz: prg_sz = sz - (prg_addr - st)
 
     assert prg_sz > 0, "Invalid prg ptr (not found in mapped ranges)"
-    # Pass valid memory ranges and rsrc2 to Python emulator for bounds checking and SGPR layout
+    # 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)

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_multitensor.py

@@ -256,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_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)

test/testextra/test_cfg_viz.py

@@ -10,7 +10,7 @@ from tinygrad.renderer import ProgramSpec
 from tinygrad.helpers import TracingKey, getenv
 from tinygrad.engine.realize import ExecItem, CompiledRunner
 
-from extra.assembly.amd.autogen.rdna3 import *
+from extra.assembly.amd.autogen.rdna3.ins import *
 
 # TODO: use the RDNA3 renderer when it's in master
 template = """.text

test/unit/test_helpers.py

@@ -2,6 +2,7 @@ import ctypes, gzip, unittest, timeit, pickle
 from tinygrad import Variable
 from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, get_contraction, count
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
+from tinygrad.helpers import ceildiv
 from tinygrad.tensor import Tensor, get_shape
 import numpy as np
 
@@ -120,6 +121,25 @@ class TestRoundUp(unittest.TestCase):
     self.assertEqual(round_up(232, 24984), 24984)
     self.assertEqual(round_up(24984, 232), 25056)
 
+class TestCeilDiv(unittest.TestCase):
+  def test_int(self):
+    self.assertEqual(ceildiv(10, 3), 4)
+    self.assertEqual(ceildiv(9, 3), 3)
+    self.assertEqual(ceildiv(0, 5), 0)
+    self.assertEqual(ceildiv(1, 5), 1)
+  def test_symbolic(self):
+    # tests that ceildiv with UOp uses (num + amt - 1) // amt formula for non-negative num
+    v = Variable('v', 0, 100)
+    result = ceildiv(v, 6)
+    self.assertEqual(result.render(), "((v+5)//6)")
+  def test_symbolic_negative_offset(self):
+    # tests ceildiv(v-5, 6) which is used in conv2d output shape
+    # old implementation incorrectly simplified -(x//-y) to ((v+1)//6-1) for v-5
+    # new implementation uses (v-5+5)//6 = v//6 which is correct
+    v = Variable('v', 11, 100)
+    result = ceildiv(v - 5, 6)
+    self.assertEqual(result.render(), "(v//6)")
+
 class TestCount(unittest.TestCase):
   def test_count_basic(self):
     c = count(3)

test/unit/test_linalg.py

@@ -65,6 +65,30 @@ class TestLinAlg(unittest.TestCase):
       orthogonality_helper(Q)
       reconstruction_helper([Q,R],a)
 
+  def test_qr_zero_column(self):
+    a = Tensor([[0.0, 1.0], [0.0, 2.0]]).realize()
+    Q,R = a.qr()
+    assert not np.isnan(Q.numpy()).any()
+    assert not np.isnan(R.numpy()).any()
+    orthogonality_helper(Q)
+    reconstruction_helper([Q,R], a)
+
+  def test_svd_identity(self):
+    for a in (Tensor.eye(2), Tensor.zeros(2, 2)):
+      a = a.realize()
+      U,S,V = a.svd()
+      assert not np.isnan(U.numpy()).any()
+      assert not np.isnan(S.numpy()).any()
+      assert not np.isnan(V.numpy()).any()
+      s_diag = (S.unsqueeze(-2) * Tensor.eye(2))
+      reconstruction_helper([U, s_diag, V], a)
+
+  def test_svd_rank1(self):
+    a = Tensor([[1.0, 1.0], [2.0, 2.0]]).realize()
+    U, S, V = a.svd()
+    np.testing.assert_allclose(S.numpy(), [np.sqrt(10), 0.0], atol=1e-4, rtol=1e-4)
+    reconstruction_helper([U, S.unsqueeze(-2) * Tensor.eye(2), V], a)
+
   def test_newton_schulz(self):
     coefficients = [(2, -1.5, 0.5), (2.0, -1.4, 0.2, 0.2)]#these params map to the sign function
     sizes = [(2,2), (3,2), (2,3), (2,2,2)]

test/unit/test_masked_tensor.py

@@ -1,13 +1,12 @@
 import unittest
 from tinygrad.tensor import Tensor
 
-class TestMaskedShapeTracker(unittest.TestCase):
+class TestMaskedTensor(unittest.TestCase):
   def test_mul_masked(self):
     a = Tensor([1,1,1,1,1])
     b = Tensor([1,1]).pad(((0,3),))
     c = a*b
     assert c.shape == a.shape
-    #assert c.uop.st.views[0].mask is not None
     ret = c.data()
     assert ret.tolist() == [1.0, 1.0, 0.0, 0.0, 0.0]
 
@@ -16,7 +15,6 @@ class TestMaskedShapeTracker(unittest.TestCase):
     b = Tensor([1,1]).pad(((0,3),))
     c = a*b
     assert c.shape == a.shape
-    #assert c.uop.st.views[0].mask is not None
     ret = c.data()
     assert ret.tolist() == [1.0, 1.0, 0.0, 0.0, 0.0]
 
@@ -24,7 +22,6 @@ class TestMaskedShapeTracker(unittest.TestCase):
     a = Tensor([1,1]).pad(((0,2),))
     b = Tensor([1,1]).pad(((0,2),))
     c = a+b
-    #assert c.uop.st.views[0].mask is not None
     ret = c.data()
     assert ret.tolist() == [2.0, 2.0, 0.0, 0.0]
 

test/unit/test_tqdm.py

@@ -128,6 +128,26 @@ class TestProgressBar(unittest.TestCase):
         self._compare_bars(tinytqdm_output, tqdm_output)
         if n > 5: break
 
+  @patch('sys.stderr', new_callable=StringIO)
+  @patch('shutil.get_terminal_size')
+  def test_si_boundary(self, mock_terminal_size, mock_stderr):
+    """Test SI formatting at boundaries (e.g., 999.5 -> 1.00k, not 1000)"""
+    ncols = 80
+    mock_terminal_size.return_value = namedtuple(field_names='columns', typename='terminal_size')(ncols)
+
+    # Test rates at the boundary: 999 stays as "999", 999.5+ becomes "1.00k"
+    for rate in [999, 999.4, 999.5, 1000, 1001]:
+      mock_stderr.truncate(0)
+      mock_stderr.seek(0)
+      elapsed = 1.0 / rate
+      # Need 3 perf_counter calls: init st, init update, final update
+      with patch('time.perf_counter', side_effect=[0, 0, elapsed]):
+        bar = tinytqdm(desc="Test", total=1, unit_scale=True, rate=10**9)
+        bar.update(1, close=True)
+      tinytqdm_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
+      tqdm_output = tqdm.format_meter(n=1, total=1, elapsed=elapsed, ncols=ncols, prefix="Test", unit_scale=True)
+      self._compare_bars(tinytqdm_output, tqdm_output)
+
   @unittest.skip("this is flaky")
   @patch('sys.stderr', new_callable=StringIO)
   @patch('shutil.get_terminal_size')

test/unit/test_validate_oob.py

@@ -0,0 +1,179 @@
+import unittest
+from tinygrad import dtypes, Variable
+from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import Context
+from tinygrad.uop.ops import Ops, UOp, AxisType
+from test.test_uops import to_uops_list
+
+class TestValidateOOB(unittest.TestCase):
+  """Test z3 validation of index bounds for different ALU ops and patterns."""
+
+  # basic index patterns
+  def test_const_index(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      to_uops_list([buf.index(UOp.const(dtypes.int, 0), ptr=True).load(dtype=dtypes.int)])  # valid
+      to_uops_list([buf.index(UOp.const(dtypes.int, 15), ptr=True).load(dtype=dtypes.int)])  # valid (last element)
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(UOp.const(dtypes.int, 16), ptr=True).load(dtype=dtypes.int)])  # off by one
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(UOp.const(dtypes.int, 42), ptr=True).load(dtype=dtypes.int)])  # way out
+
+  def test_variable_index(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      to_uops_list([buf.index(Variable("i", 0, 15), ptr=True).load(dtype=dtypes.int)])  # valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(Variable("i", 0, 20), ptr=True).load(dtype=dtypes.int)])  # oob
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(Variable("i", -5, 10), ptr=True).load(dtype=dtypes.int)])  # negative
+
+  def test_range_with_mask(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      r = UOp.range(42, 0, AxisType.GLOBAL)
+      to_uops_list([buf.index(r.valid(r < 16), ptr=True).load(dtype=dtypes.int)])  # valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r.valid(r < 17), ptr=True).load(dtype=dtypes.int)])  # oob
+
+  def test_variable_with_mask(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      v = Variable("v", -5, 80)
+      to_uops_list([buf.index(v.valid((v >= 0) & (v < 16)), ptr=True).load(dtype=dtypes.int)])  # valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(v.valid(v < 20), ptr=True).load(dtype=dtypes.int)])  # negative not masked
+
+  def test_gated_store(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      v = Variable("v", 0, 20)
+      to_uops_list([buf.index(v.valid(v < 16)).store(0)])  # valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(v.valid(v < 20)).store(0)])  # oob
+
+  # ALU ops in index
+  def test_idiv(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)])  # 0..16 oob
+
+  def test_mod(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      r = UOp.range(100, 0, AxisType.GLOBAL)
+      to_uops_list([buf.index(r % 16, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r % 20, ptr=True).load(dtype=dtypes.int)])  # 0..19 oob
+
+  def test_shr(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)])  # 0..31 oob
+
+  def test_shl(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 0)
+      r = UOp.range(8, 0, AxisType.GLOBAL)
+      to_uops_list([buf.index(r << 2, ptr=True).load(dtype=dtypes.int)])  # 0..28 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r << 4, ptr=True).load(dtype=dtypes.int)])  # 0..112 oob
+
+  def test_and(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      r = UOp.range(100, 0, AxisType.GLOBAL)
+      to_uops_list([buf.index(r & 15, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r & 31, ptr=True).load(dtype=dtypes.int)])  # 0..31 oob
+
+  def test_max(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      to_uops_list([buf.index(Variable("v", -10, 15).maximum(0), ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0), ptr=True).load(dtype=dtypes.int)])  # 0..20 oob
+
+  def test_xor_in_mask(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      r = UOp.range(32, 0, AxisType.GLOBAL)
+      to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16))), ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r.valid((r < 10) ^ (r >= 20)), ptr=True).load(dtype=dtypes.int)])  # 0..9,20..31 oob
+
+  # cast patterns
+  def test_float_cast_in_index(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      r = UOp.range(20, 0)
+      i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
+      to_uops_list([buf.index(i.valid((i >= 0) & (i < 16)), ptr=True).load(dtype=dtypes.int)])
+
+  def test_bool_cast_in_mask(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), (), 0)
+      r = UOp.range(20, 0)
+      to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not()), ptr=True).load(dtype=dtypes.int)])  # only r=0 valid
+
+  # load result as index/mask
+  def test_load_as_index(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
+      buf1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 1)
+      r = UOp.range(42, 0, AxisType.GLOBAL)
+      ld0 = buf0.index(r.valid(r < 8), ptr=True).load(dtype=dtypes.int).cast(dtypes.index)
+      to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32)), ptr=True).load(dtype=dtypes.int)])  # valid
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64)), ptr=True).load(dtype=dtypes.int)])  # oob
+
+  def test_load_bool_as_mask(self):
+    with Context(IGNORE_OOB=0, SPEC=2):
+      buf_bool = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
+      buf_int = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(8), (), 1)
+      gidx = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 16),), "gidx0")
+      ld_bool = buf_bool.index(gidx, ptr=True).load()
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf_int.index(gidx.valid(ld_bool), ptr=True).load()])  # gidx 0..15, buf_int size 8
+
+  # skipped tests (moved from test_uop_graph.py)
+  @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])
+
+  @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])
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad/codegen/__init__.py

@@ -166,8 +166,10 @@ def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> Program
   if ast.arg is None: ast = ast.replace(arg=KernelInfo())
 
   # rewrite to prg
-  full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
-  prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.device)))
+  if ast.op is Ops.PROGRAM: prg = ast
+  else:
+    full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
+    prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.device)))
   prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
 
   # create the ProgramSpec

tinygrad/codegen/opt/postrange.py

@@ -45,6 +45,7 @@ class Scheduler:
     ret = Scheduler(self.ast, self.ren)
     ret.dont_use_locals = self.dont_use_locals
     ret.applied_opts = self.applied_opts[:]
+    if hasattr(self, 'tensor_core'): ret.tensor_core = self.tensor_core
     return ret
 
   kernel_cnt: Final[defaultdict[str, int]] = defaultdict(int)
@@ -307,6 +308,7 @@ class Scheduler:
             reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes]
             if len(reduce_ranges): tc_uop = UOp(Ops.REDUCE, tc_uop.dtype, (tc_uop,)+tuple(reduce_ranges), Ops.ADD)
             self.ast = self.ast.substitute({reduceop: tc_uop})
+          self.tensor_core = tc
           return axes
     return None
 

tinygrad/codegen/opt/search.py

@@ -93,8 +93,8 @@ def _ensure_buffer_alloc(bufs:list[Buffer]) -> list[Buffer]: return [buf.ensure_
 # *** external API ***
 
 # get dictionary of all possible actions
-def get_kernel_actions(s:Scheduler, include_0=True) -> dict[int, Scheduler]:
-  acted, max_up, max_lcl = {0:s} if include_0 else {}, getenv("BEAM_UPCAST_MAX", 256), getenv("BEAM_LOCAL_MAX", 1024)
+def get_kernel_actions(s:Scheduler, include_0=True, max_up:int|None=None) -> dict[int, Scheduler]:
+  acted, max_up, max_lcl = {0:s} if include_0 else {}, getenv("BEAM_UPCAST_MAX", 256) if max_up is None else max_up, getenv("BEAM_LOCAL_MAX", 1024)
   kernel_actions = actions.copy()
 
   for i,a in enumerate(kernel_actions):

tinygrad/dtype.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 from typing import Final, ClassVar, Callable, Literal
 import math, struct, ctypes, functools
 from dataclasses import dataclass, fields
-from tinygrad.helpers import getenv, prod
+from tinygrad.helpers import getenv, prod, round_up, next_power2
 from enum import Enum, auto
 
 class InvalidTypeMetaClass(type):
@@ -101,6 +101,15 @@ class ImageDType(PtrDType):
     assert addrspace == AddrSpace.GLOBAL, "images can't be local"
     return self
   def __repr__(self): return f"dtypes.{self.name}({self.shape})" + (f'.vec({self.v})' if self.v != 1 else '')
+  @property
+  def pitch(self):
+    imgw, imgh, itemsize_log = self.shape[1], self.shape[0], int(math.log2(self.itemsize))
+    pitchalign = max(6, 11 - int(math.log2(imgh))) if imgh > 1 else 6
+    align_up = max(1, (8 // itemsize_log + 1) - imgh // 32) if pitchalign == 6 else (2 ** (pitchalign - itemsize_log - 2))
+
+    granularity = 128 if self.itemsize == 4 else 256
+    pitch_add = (1 << pitchalign) if min(next_power2(imgw), round_up(imgw, granularity)) - align_up + 1 <= imgw and imgw > granularity//2 else 0
+    return round_up(imgw * 4 * self.itemsize, 1 << pitchalign) + pitch_add
 
 class dtypes:
   @staticmethod

tinygrad/engine/realize.py

@@ -125,7 +125,7 @@ def get_runner(device:str, ast:UOp) -> CompiledRunner:
 
 # NOTE: ctx is the buffers
 si_lowerer = PatternMatcher([
-  (UPat(Ops.SINK, name="sink"), lambda ctx,sink: get_runner(ctx[0].device, sink)),
+  (UPat((Ops.SINK, Ops.PROGRAM), name="sink"), lambda ctx,sink: get_runner(ctx[0].device, sink)),
   (UPat(Ops.BUFFER_VIEW), lambda ctx: ViewOp(ctx[0])),
   (UPat(Ops.COPY, name="copy"), lambda ctx,copy: (BufferXfer(ctx[0].nbytes, ctx[0].device, ctx[1].device) \
       if hasattr(Device[ctx[0].device].allocator, '_transfer') and all_same([x.device.split(":")[0] for x in ctx]) \

tinygrad/helpers.py

@@ -38,18 +38,18 @@ def ansilen(s:str): return len(ansistrip(s))
 def make_tuple(x:int|Sequence[int], cnt:int) -> tuple[int, ...]: return (x,)*cnt if isinstance(x, int) else tuple(x)
 def flatten(l:Iterable[Iterable[T]]): return [item for sublist in l for item in sublist]
 def fully_flatten(l):
-  if hasattr(l, "__len__") and hasattr(l, "__getitem__") and not isinstance(l, str):
-    if hasattr(l, "shape") and l.shape == (): return [l[()]]
-    flattened = []
-    for li in l: flattened.extend(fully_flatten(li))
-    return flattened
-  return [l]
+  if not (hasattr(l, "__len__") and hasattr(l, "__getitem__")) or isinstance(l, str): return [l]
+  return [l[()]] if hasattr(l, "shape") and l.shape == () else [x for li in l for x in fully_flatten(li)]
 def fromimport(mod, frm): return getattr(__import__(mod, fromlist=[frm]), frm)
 def _is_balanced(s:str) -> bool: return (d := 0, all((d := d + (c == '(') - (c == ')')) >= 0 for c in s))[1] and d == 0
-def strip_parens(fst:str) -> str: return fst[1:-1] if fst and fst[0]=='(' and fst[-1] == ')' and _is_balanced(fst[1:-1]) else fst
-def ceildiv(num, amt): return int(ret) if isinstance((ret:=-(num//-amt)), float) else ret
+def strip_parens(fst:str) -> str: return fst[1:-1] if fst[:1]=='(' and fst[-1:]==')' and _is_balanced(fst[1:-1]) else fst
+def ceildiv(num, amt):
+  # use (num + amt - 1) // amt when num is a UOp and non-negative to avoid C/Python division mismatch
+  if hasattr(num, 'vmin') and num.vmin >= 0 and (amt > 0 if isinstance(amt, int) else amt.vmin > 0): return (num + amt - 1) // amt
+  return int(ret) if isinstance((ret:=-(num//-amt)), float) else ret
 def round_up(num:int, amt:int) -> int: return (num+amt-1)//amt * amt
 def round_down(num:int, amt:int) -> int: return -round_up(-num, amt)
+def next_power2(x): return 1 if x == 0 else 1 << (x - 1).bit_length()
 # cstyle div and mod
 def cdiv(x:int, y:int) -> int: return abs(x)//abs(y)*(1,-1)[x*y<0] if y != 0 else 0
 def cmod(x:int, y:int) -> int: return x-cdiv(x,y)*y
@@ -87,9 +87,7 @@ def word_wrap(x, wrap=80):
   while len(ansistrip(x[:i])) < wrap and i < len(x): i += 1
   return x[:i] + "\n" + word_wrap(x[i:], wrap)
 def pad_bytes(b:bytes, align:int) -> bytes: return b + b'\x00' * ((align - (len(b) % align)) % align)
-def panic(e:Exception|None=None):
-  if e is None: raise RuntimeError("PANIC!")
-  raise e
+def panic(e:Exception|None=None): raise e if e is not None else RuntimeError("PANIC!")
 
 @functools.cache
 def canonicalize_strides(shape:tuple[T, ...], strides:tuple[T, ...]) -> tuple[T, ...]:
@@ -149,9 +147,7 @@ def getenv(key:str, default:Any=0): return type(default)(os.getenv(key, default)
 def temp(x:str, append_user:bool=False) -> str:
   return (pathlib.Path(tempfile.gettempdir()) / (f"{x}.{getpass.getuser()}" if append_user else x)).as_posix()
 
-def stderr_log(msg):
-  sys.stderr.write(msg)
-  sys.stderr.flush()
+def stderr_log(msg:str): print(msg, end='', file=sys.stderr, flush=True)
 
 class Context(contextlib.ContextDecorator):
   def __init__(self, **kwargs): self.kwargs = kwargs
@@ -181,8 +177,8 @@ JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVa
 WINO, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1)
 USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0)
 TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS", 0)
-SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1)
-LRU = ContextVar("LRU", 1)
+SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("LRU", 1)
+RING, ALL2ALL = ContextVar("RING", 1), ContextVar("ALL2ALL", 0)
 CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
 VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
@@ -512,7 +508,9 @@ class tqdm(Generic[T]):
     if elapsed and self.i/elapsed > self.rate and self.i: self.skip = max(int(self.i/elapsed)//self.rate,1)
     def HMS(t): return ':'.join(f'{x:02d}' if i else str(x) for i,x in enumerate([int(t)//3600,int(t)%3600//60,int(t)%60]) if i or x)
     def SI(x):
-      return (f"{x/1000**int(g:=round(math.log(x,1000),6)):.{int(3-3*math.fmod(g,1))}f}"[:4].rstrip('.')+' kMGTPEZY'[int(g)].strip()) if x else '0.00'
+      if not x: return '0.00'
+      v = f"{x/1000**int(g:=round(math.log(x,1000),6)):.{int(3-3*math.fmod(g,1))}f}"[:4].rstrip('.')
+      return (f"{x/1000**(int(g)+1):.3f}"[:4].rstrip('.')+' kMGTPEZY'[int(g)+1]) if v == "1000" else v+' kMGTPEZY'[int(g)].strip()
     prog_text = f'{SI(self.n)}{f"/{SI(self.t)}" if self.t else self.unit}' if self.unit_scale else f'{self.n}{f"/{self.t}" if self.t else self.unit}'
     est_text = f'<{HMS(elapsed/prog-elapsed) if self.n else "?"}' if self.t else ''
     it_text = (SI(self.n/elapsed) if self.unit_scale else f"{self.n/elapsed:5.2f}") if self.n else "?"

tinygrad/renderer/ptx.py

@@ -133,7 +133,7 @@ string_rewrite = PatternMatcher([
   (UPat(Ops.IF, name="x"), lambda ctx, x: f"@!{ctx.r[x.src[0]]} bra IF_{ctx.r[x.src[0]][1:]}_{ctx.uops.index(x)};"),
   (UPat(Ops.ENDIF, name="x"), lambda ctx, x: f"IF_{ctx.r[x.src[0].src[0]][1:]}_{ctx.uops.index(x.src[0])}:"),
   (UPat(Ops.WMMA, name="x"), lambda ctx, x: list(render_wmma(ctx, x))),
-  (UPat(Ops.BARRIER, name="x"), lambda ctx, x: ctx.barrier),
+  (UPat(Ops.BARRIER), lambda ctx: ctx.barrier),
   (UPat(Ops.DEFINE_VAR, name="x"), lambda ctx, x: f"ld.param.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{x.arg[0]}+0];"),
 ])
 
@@ -180,7 +180,7 @@ class PTXRenderer(Renderer):
     self.uops = uops
 
     def ssa(prefix:str, u:UOp|None=None, dtype:str|None=None) -> str:
-      nonlocal c, r
+      nonlocal c
       prefix += f"_{dtype if dtype is not None else self.types[unwrap(u).dtype.base]}_"
       c[prefix] += 1
       return f"%{prefix}{c[prefix]-1}"
@@ -230,7 +230,7 @@ class PTXRenderer(Renderer):
                        [ssa("wmma_acc", dtype="b32") for _ in range(0, len(r[u.src[2]]), 4 // u.dtype.scalar().itemsize)]]
         r[u] = [ssa("wmma", dtype=self.types[u.dtype.scalar()]) for _ in range(u.dtype.count)]
       prefix, dtype = {Ops.CAST: ("cast", None), Ops.BITCAST: ("cast", None), Ops.END: ("pred", "pred"), Ops.RANGE: ("ridx", None),
-        Ops.DEFINE_VAR: ("dat", None), Ops.CONST: ("const", None), Ops.DEFINE_LOCAL: ("local",self.types[dtypes.ulong]),
+        Ops.DEFINE_VAR: ("dat", None), Ops.CONST: ("const", None), Ops.DEFINE_LOCAL: ("local", self.types[dtypes.ulong]),
         Ops.DEFINE_GLOBAL: ("dat", self.types[dtypes.ulong]), **{op: ("alu", None) for op in GroupOp.ALU}}.get(u.op, (None, None))
       if prefix: r[u] = ssa(prefix, u, dtype)
 

tinygrad/runtime/graph/hcq.py

@@ -1,6 +1,6 @@
-import collections, time
+import collections, itertools, time
 from typing import Any, cast
-from tinygrad.helpers import round_up, PROFILE, merge_dicts, getenv, dedup, suppress_finalizing
+from tinygrad.helpers import round_up, PROFILE, ALL2ALL, merge_dicts, getenv, dedup, suppress_finalizing
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWQueue, HCQArgsState, BumpAllocator, MMIOInterface
 from tinygrad.device import Buffer, BufferSpec, Compiled, Device, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.dtype import dtypes
@@ -22,7 +22,7 @@ class HCQGraph(MultiGraphRunner):
 
     for (j,i), input_idx in self.input_replace.items():
       x = self.input_replace_to_var.setdefault((j,i), UOp.variable(f"input_{input_idx}", 0, 0xffffffffffffffff, dtype=dtypes.uint64))
-      self.hcq_bufs[j][i] = HCQBuffer(x, self.hcq_bufs[j][i].size, texture_info=self.hcq_bufs[j][i].texture_info) # Create fake buffer with variable
+      self.hcq_bufs[j][i] = HCQBuffer(x, self.hcq_bufs[j][i].size, image=self.hcq_bufs[j][i].image) # Create fake buffer with variable
 
     # Allocate kernel args.
     kernargs_size: dict[Compiled, int] = collections.defaultdict(int)
@@ -49,7 +49,9 @@ class HCQGraph(MultiGraphRunner):
     self.ji_schedule: dict[int, tuple[HCQCompiled, HWQueue, list, list, HCQSignal, int|None]] = {}
 
     self.comp_queues: dict[HCQCompiled, HWQueue] = {dev: dev.hw_compute_queue_t() for dev in self.devices}
-    self.copy_queues: dict[HCQCompiled, HWQueue] = {} # lazy allocation
+    self.copy_queues: dict[tuple[HCQCompiled, int], HWQueue] = {} # lazy allocation, keyed by (device, queue_idx)
+    self.num_copy_queues: int = getenv("HCQ_NUM_SDMA", 2 if ALL2ALL >= 1 else 1)
+    self.copy_queue_cnt: collections.defaultdict[HCQCompiled, itertools.count] = collections.defaultdict(itertools.count)
 
     self.signals: dict[Any, HCQSignal] = {**{dev: dev.new_signal(value=0) for dev in self.devices if not dev._is_cpu()},
       **{"KICK": self.devices[0].new_signal(value=0)}, **{dev: self.devices[0].new_signal(value=0) for dev in self.devices if dev._is_cpu()}}
@@ -85,7 +87,8 @@ class HCQGraph(MultiGraphRunner):
         enqueue_queue = self.comp_queues[enqueue_dev]
       else:
         assert (enqueue_dev.hw_copy_queue_t is not None), "device must implement a copy queue"
-        enqueue_queue = self.copy_queues.setdefault(enqueue_dev, enqueue_dev.hw_copy_queue_t())
+        queue_idx = next(self.copy_queue_cnt[enqueue_dev]) % self.num_copy_queues
+        enqueue_queue = self.copy_queues.setdefault((enqueue_dev, queue_idx), enqueue_dev.hw_copy_queue_t(queue_idx=queue_idx))
 
       out_signal = self.signals.setdefault(enqueue_queue, self.devices[0].new_signal(value=0))
 
@@ -175,14 +178,17 @@ class HCQGraph(MultiGraphRunner):
 
     for dev in self.devices:
       for dep_dev in list(self.copy_to_devs[dev]) + [dev]:
-        if dep_dev in self.copy_queues: self.comp_queues[dev].wait(self.signals[(copy_q:=self.copy_queues[dep_dev])], cast(int, last_j[copy_q]) + 1)
+        for copy_q in self._dev_copy_queues(dep_dev):
+          if copy_q in self.signals: self.comp_queues[dev].wait(self.signals[copy_q], cast(int, last_j[copy_q]) + 1)
 
       self.comp_queues[dev].signal(self.virt_timeline_signals[dev], self.virt_timeline_vals[dev] + 1).bind(dev)
-      if dev in self.copy_queues: self.copy_queues[dev].bind(dev)
+      for copy_q in self._dev_copy_queues(dev): copy_q.bind(dev)
 
     self.last_timeline: dict[HCQCompiled, tuple[HCQSignal, int]] = {dev: (dev.timeline_signal, 0) for dev in self.devices}
     self.queue_signals_to_reset = [self.signals[q] for q in list(self.comp_queues.values()) + list(self.copy_queues.values()) if q in self.signals]
 
+  def _dev_copy_queues(self, dev): return [q for (d, _), q in self.copy_queues.items() if d == dev]
+
   def __call__(self, input_rawbuffers: list[Buffer], var_vals: dict[str, int], wait=False) -> float|None:
     # Wait and restore signals
     self.kickoff_value += 1
@@ -205,8 +211,7 @@ class HCQGraph(MultiGraphRunner):
 
     for dev in self.devices:
       self.comp_queues[dev].submit(dev, hcq_var_vals_local:=hcq_var_vals|self.fixedvars.get(dev, {}))
-      if (copy_queue:=self.copy_queues.get(dev, None)) is not None: copy_queue.submit(dev, hcq_var_vals_local)
-
+      for copy_queue in self._dev_copy_queues(dev): copy_queue.submit(dev, hcq_var_vals_local)
       self.last_timeline[dev] = (dev.timeline_signal, dev.next_timeline())
 
     if wait:

tinygrad/runtime/ops_amd.py

@@ -446,8 +446,8 @@ class AMDComputeAQLQueue(AMDComputeQueue):
     dev.compute_queue.signal_doorbell(dev, doorbell_value=dev.compute_queue.put_value-1)
 
 class AMDCopyQueue(HWQueue):
-  def __init__(self, dev, max_copy_size=0x40000000):
-    self.dev, self.sdma, self.internal_cmd_sizes, self.max_copy_size = dev, dev.sdma, [], max_copy_size
+  def __init__(self, dev, max_copy_size=0x40000000, queue_idx=0):
+    self.dev, self.sdma, self.internal_cmd_sizes, self.max_copy_size, self.queue_idx = dev, dev.sdma, [], max_copy_size, queue_idx
     super().__init__()
 
   def q(self, *arr):
@@ -501,41 +501,42 @@ class AMDCopyQueue(HWQueue):
     self._q, self.cmd_sizes = hw_view, [len(self.indirect_cmd)]
 
   def _submit(self, dev:AMDDevice):
+    sdma_queue = dev.sdma_queue(self.queue_idx)
     if self.binded_device == dev:
       # An IB packet must end on a 8 DW boundary.
-      add = (8 - (((dev.sdma_queue.put_value % 32) // 4) + len(self.indirect_cmd) % 8)) % 8
+      add = (8 - (((sdma_queue.put_value % 32) // 4) + len(self.indirect_cmd) % 8)) % 8
       cmds, cmd_sizes = ([0] * add) + self.indirect_cmd, [len(self.indirect_cmd) + add]
 
-      if len(cmds) * 4 >= (dev.sdma_queue.ring.nbytes - dev.sdma_queue.put_value % dev.sdma_queue.ring.nbytes):
+      if len(cmds) * 4 >= (sdma_queue.ring.nbytes - sdma_queue.put_value % sdma_queue.ring.nbytes):
         cmds, cmd_sizes = [0, 0] + self.indirect_cmd, [8]
     else: cmds, cmd_sizes = self._q, self.internal_cmd_sizes
 
     tail_blit_dword = 0
     for cmdsz in cmd_sizes:
-      if (tail_blit_dword + cmdsz) * 4 >= dev.sdma_queue.ring.nbytes - dev.sdma_queue.put_value % dev.sdma_queue.ring.nbytes: break
+      if (tail_blit_dword + cmdsz) * 4 >= sdma_queue.ring.nbytes - sdma_queue.put_value % sdma_queue.ring.nbytes: break
       tail_blit_dword += cmdsz
 
     # Force align of submits to hit our usb layer write cache.
     if (rem_packet_cnt := len(cmds) - tail_blit_dword) > 0 and dev.is_usb(): tail_blit_dword = 0
 
     # USB devices run in single-step mode, so they can't overrun the queue.
-    total_bytes = (tail_blit_dword * 4 if rem_packet_cnt == 0 else -dev.sdma_queue.put_value % dev.sdma_queue.ring.nbytes) + rem_packet_cnt * 4
-    assert total_bytes < dev.sdma_queue.ring.nbytes, "SDMA queue overrun"
-    while not dev.is_usb() and dev.sdma_queue.put_value + total_bytes - dev.sdma_queue.read_ptr > dev.sdma_queue.ring.nbytes: pass
+    total_bytes = (tail_blit_dword * 4 if rem_packet_cnt == 0 else -sdma_queue.put_value % sdma_queue.ring.nbytes) + rem_packet_cnt * 4
+    assert total_bytes < sdma_queue.ring.nbytes, "SDMA queue overrun"
+    while not dev.is_usb() and sdma_queue.put_value + total_bytes - sdma_queue.read_ptr > sdma_queue.ring.nbytes: pass
 
-    start_idx = (dev.sdma_queue.put_value % dev.sdma_queue.ring.nbytes) // 4
-    dev.sdma_queue.ring[start_idx : start_idx + tail_blit_dword] = array.array('I', cmds[:tail_blit_dword])
-    dev.sdma_queue.put_value += tail_blit_dword * 4
+    start_idx = (sdma_queue.put_value % sdma_queue.ring.nbytes) // 4
+    sdma_queue.ring[start_idx : start_idx + tail_blit_dword] = array.array('I', cmds[:tail_blit_dword])
+    sdma_queue.put_value += tail_blit_dword * 4
 
     if (rem_packet_cnt := len(cmds) - tail_blit_dword) > 0:
-      zero_fill = dev.sdma_queue.ring.nbytes - dev.sdma_queue.put_value % dev.sdma_queue.ring.nbytes
-      dev.sdma_queue.ring.view(dev.sdma_queue.put_value % dev.sdma_queue.ring.nbytes, zero_fill, fmt='B')[:] = bytes(zero_fill)
-      dev.sdma_queue.put_value += zero_fill
+      zero_fill = sdma_queue.ring.nbytes - sdma_queue.put_value % sdma_queue.ring.nbytes
+      sdma_queue.ring.view(sdma_queue.put_value % sdma_queue.ring.nbytes, zero_fill, fmt='B')[:] = bytes(zero_fill)
+      sdma_queue.put_value += zero_fill
 
-      dev.sdma_queue.ring[0:rem_packet_cnt] = array.array('I', cmds[tail_blit_dword:])
-      dev.sdma_queue.put_value += rem_packet_cnt * 4
+      sdma_queue.ring[0:rem_packet_cnt] = array.array('I', cmds[tail_blit_dword:])
+      sdma_queue.put_value += rem_packet_cnt * 4
 
-    dev.sdma_queue.signal_doorbell(dev)
+    sdma_queue.signal_doorbell(dev)
 
 class AMDProgram(HCQProgram):
   def __init__(self, dev:AMDDevice, name:str, lib:bytes):
@@ -756,7 +757,8 @@ class KFDIface:
     stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
     assert stm.n_success == 1
 
-  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
+  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
+                   xcc_id=0, idx=0):
     queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring.va_addr, ring_size=ring.size, gpu_id=self.gpu_id,
       queue_type=queue_type, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE|(xcc_id<<8), queue_priority=getenv("AMD_KFD_QUEUE_PRIORITY", 7),
       eop_buffer_address=eop_buffer.va_addr if eop_buffer else 0, eop_buffer_size=eop_buffer.size if eop_buffer else 0, ctl_stack_size=ctl_stack_size,
@@ -826,15 +828,17 @@ class PCIIface(PCIIfaceBase):
       'simd_arrays_per_engine': max_sh_per_se, 'lds_size_in_kb': self.dev_impl.gc_info.gc_lds_size, 'num_xcc': self.dev_impl.gfx.xccs,
       'gfx_target_version': {90403: 90402}.get(gfxver, gfxver)}
 
-  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
+  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
+                   xcc_id=0, idx=0):
     assert cwsr_buffer is None, "no cwsr buffer for am"
 
     if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
-      pv = self.dev_impl.sdma.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr, wptr_addr=gart.va_addr+wptr,
-                                    doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0), pipe=0, queue=0)
+      assert idx <= 3, "only 4 SDMA queues supported in am"
+      pv, doorbell_index = self.dev_impl.sdma.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr,
+                                                         wptr_addr=gart.va_addr+wptr, pipe=0, queue=idx)
     else:
-      pv = self.dev_impl.gfx.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr, wptr_addr=gart.va_addr+wptr,
-        eop_addr=eop_buffer.va_addr, eop_size=eop_buffer.size, doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_MEC_RING0), pipe=0,
+      pv, doorbell_index = self.dev_impl.gfx.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr,
+        wptr_addr=gart.va_addr+wptr, eop_addr=eop_buffer.va_addr, eop_size=eop_buffer.size, pipe=0,
         queue=int(is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL)), aql=is_aql)
 
     return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbells=[self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q')],
@@ -875,9 +879,10 @@ class USBIface(PCIIface):
     barview = self.pci_dev.map_bar(bar=0, off=mapping.paddrs[0][0], size=mapping.size) if cpu_access else None
     return HCQBuffer(mapping.va_addr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=False), view=barview, owner=self.dev)
 
-  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
+  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
+                   xcc_id=0, idx=0):
     if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE: self.pci_dev.usb._pci_cacheable += [(ring.cpu_view().addr, ring.size)]
-    return super().create_queue(queue_type, ring, gart, rptr, wptr, eop_buffer, cwsr_buffer, ctl_stack_size, ctx_save_restore_size, xcc_id)
+    return super().create_queue(queue_type, ring, gart, rptr, wptr, eop_buffer, cwsr_buffer, ctl_stack_size, ctx_save_restore_size, xcc_id, idx)
 
   def sleep(self, timeout): pass
 
@@ -931,8 +936,7 @@ class AMDDevice(HCQCompiled):
       0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
       ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size, debug_memory_size=debug_memory_size)
 
-    max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
-    self.sdma_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20))
+    self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
 
     compilers = CompilerSet([CompilerPair(functools.partial(AMDHIPRenderer, self.arch), None),
                              CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), None, AMD_LLVM),
@@ -940,7 +944,7 @@ class AMDDevice(HCQCompiled):
 
     super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
                      functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),
-                     functools.partial(AMDCopyQueue, self, max_copy_size=max_copy_size),
+                     functools.partial(AMDCopyQueue, self, max_copy_size=self.max_copy_size),
                      kernargs_size=(8 << 10) if self.is_usb() else (16 << 20), sigalloc_size=0x100 if self.is_usb() else 0x1000)
 
     # Scratch setup
@@ -976,7 +980,7 @@ class AMDDevice(HCQCompiled):
       self.sqtt_wptrs = self.allocator.alloc(round_up(self.se_cnt * 4, 0x1000), BufferSpec(cpu_access=True, nolru=True))
       self.sqtt_next_cmd_id = itertools.count(0)
 
-  def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0):
+  def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0, idx=0):
     ring = self.iface.alloc(ring_size, uncached=True, cpu_access=True)
     gart = self.iface.alloc(0x100, uncached=True, cpu_access=True)
 
@@ -993,7 +997,10 @@ class AMDDevice(HCQCompiled):
 
     return (self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
             wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer,
-            ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size))
+            ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx))
+
+  @functools.lru_cache(None)
+  def sdma_queue(self, idx:int=0): return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
 
   def _ensure_has_local_memory(self, private_segment_size):
     if self.max_private_segment_size >= private_segment_size: return

tinygrad/runtime/ops_nv.py

@@ -165,6 +165,10 @@ class NVComputeQueue(NVCommandQueue):
   def _submit(self, dev:NVDevice): self._submit_to_gpfifo(dev, dev.compute_gpfifo)
 
 class NVCopyQueue(NVCommandQueue):
+  def __init__(self, queue_idx=0):
+    self.queue_idx = queue_idx
+    super().__init__()
+
   def copy(self, dest:sint, src:sint, copy_size:int):
     for off in range(0, copy_size, step:=(1 << 31)):
       self.nvm(4, nv_gpu.NVC6B5_OFFSET_IN_UPPER, *data64(src+off), *data64(dest+off))

tinygrad/runtime/ops_qcom.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 import os, ctypes, functools, mmap, struct, array, math, sys, weakref, contextlib
 assert sys.platform != 'win32'
-from types import SimpleNamespace
 from typing import Any, cast
 from tinygrad.device import BufferSpec, CompilerSet, CompilerPair
 from tinygrad.runtime.support.hcq import HCQBuffer, HWQueue, HCQProgram, HCQCompiled, HCQAllocatorBase, HCQSignal, HCQArgsState, BumpAllocator
@@ -11,7 +10,7 @@ from tinygrad.runtime.ops_cl import CLCompiler, CLDevice
 from tinygrad.renderer.cstyle import QCOMRenderer
 from tinygrad.renderer.nir import IR3Renderer
 from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, prod, fromimport, cpu_profile, lo32, PROFILE, suppress_finalizing
-from tinygrad.helpers import flatten, QCOM_IR3, QCOM_CC
+from tinygrad.helpers import next_power2, flatten, QCOM_IR3, QCOM_CC
 from tinygrad.runtime.support.system import System
 if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl  # noqa: F401  # pylint: disable=unused-import
 
@@ -26,7 +25,7 @@ def _qreg_exec(__reg, __val=0, **kwargs):
   return __val
 qreg: Any = type("QREG", (object,), {name[4:].lower(): functools.partial(_qreg_exec, name) for name in mesa.__dict__.keys() if name[:4] == 'REG_'})
 
-def next_power2(x): return 1 if x == 0 else 1 << (x - 1).bit_length()
+def ctz(v): return (v & -v).bit_length() - 1
 
 def parity(val: int):
   for i in range(4,1,-1): val ^= val >> (1 << i)
@@ -191,37 +190,29 @@ class QCOMComputeQueue(HWQueue):
 class QCOMArgsState(HCQArgsState):
   def __init__(self, buf:HCQBuffer, prg:QCOMProgram, bufs:tuple[HCQBuffer, ...], vals:tuple[int, ...]=()):
     super().__init__(buf, prg, bufs, vals=vals)
-
-    if len(bufs) + len(vals) != len(prg.buf_info): raise RuntimeError(f'incorrect args size given={len(bufs)+len(vals)} != want={len(prg.buf_info)}')
-
-    self.buf_info, self.args_info = prg.buf_info[:len(bufs)], prg.buf_info[len(bufs):]
-
     ctypes.memset(cast(int, self.buf.va_addr), 0, prg.kernargs_alloc_size)
+
+    ubos, uavs = [b for b in bufs if b.image is None], [b for b in bufs if b.image is not None]
+    ibos, texs = (uavs, []) if prg.tex_cnt == 0 else (uavs[:-prg.tex_cnt], uavs[-prg.tex_cnt:])
     for cnst_val,cnst_off,cnst_sz in prg.consts_info: to_mv(self.buf.va_addr + cnst_off, cnst_sz)[:] = cnst_val.to_bytes(cnst_sz, byteorder='little')
 
     if prg.samp_cnt > 0: to_mv(self.buf.va_addr + prg.samp_off, len(prg.samplers) * 4).cast('I')[:] = array.array('I', prg.samplers)
-    for i, b in enumerate(bufs):
-      if prg.buf_info[i].type in {BUFTYPE_TEX, BUFTYPE_IBO}:
-        obj = b.texture_info.desc if prg.buf_info[i].type is BUFTYPE_TEX else b.texture_info.ibo
-        to_mv(self.buf.va_addr + prg.buf_info[i].offset, len(obj) * 4).cast('I')[:] = array.array('I', obj)
-      self.bind_sints_to_buf(b.va_addr, buf=self.buf, fmt='Q', offset=self.buf_info[i].offset+(0 if self.buf_info[i].type is BUFTYPE_BUF else 16))
+    if prg.NIR:
+      self.bind_sints_to_buf(*[b.va_addr for b in ubos], buf=self.buf, fmt='Q', offset=prg.buf_off)
+      self.bind_sints_to_buf(*vals, buf=self.buf, fmt='I', offset=prg.buf_off + len(ubos) * 8)
+    else:
+      for i, b in enumerate(ubos): self.bind_sints_to_buf(b.va_addr, buf=self.buf, fmt='Q', offset=prg.buf_offs[i])
+      for i, v in enumerate(vals): self.bind_sints_to_buf(v, buf=self.buf, fmt='I', offset=prg.buf_offs[i+len(ubos)])
 
-    for i, v in enumerate(vals): self.bind_sints_to_buf(v, buf=self.buf, fmt='I', offset=self.args_info[i].offset)
+    def _tex(b, ibo=False):
+      fmt = mesa.FMT6_32_32_32_32_FLOAT if b.image.itemsize == 4 else mesa.FMT6_16_16_16_16_FLOAT
+      return [qreg.a6xx_tex_const_0(fmt=fmt) if ibo else qreg.a6xx_tex_const_0(0x8, swiz_x=0, swiz_y=1, swiz_z=2, swiz_w=3, fmt=fmt),
+              qreg.a6xx_tex_const_1(width=b.image.shape[1], height=b.image.shape[0]),
+              qreg.a6xx_tex_const_2(type=mesa.A6XX_TEX_2D, pitch=b.image.pitch, pitchalign=ctz(b.image.pitch)-6), 0, *data64_le(b.va_addr),
+              qreg.a6xx_tex_const_6(plane_pitch=0x400000), qreg.a6xx_tex_const_7(13), 0, 0, 0, 0, 0, 0, 0, 0]
 
-class IR3ArgsState(HCQArgsState):
-  def __init__(self, buf:HCQBuffer, prg:QCOMProgram, bufs:tuple[HCQBuffer, ...], vals:tuple[int, ...]=()):
-    super().__init__(buf, prg, bufs, vals=vals)
-    ctypes.memset(cast(int, self.buf.va_addr), 0, prg.kernargs_alloc_size)
-    to_mv(self.buf.va_addr + prg.imm_off, len(prg.imm_vals))[:] = prg.imm_vals
-
-    ubos, uavs = [b for b in bufs if b.texture_info is None], [b for b in bufs if b.texture_info is not None]
-    ibos, texs = (uavs, []) if prg.tex_cnt == 0 else (uavs[:-prg.tex_cnt], uavs[-prg.tex_cnt:]) # textures are at the end
-
-    if prg.samp_cnt > 0: to_mv(self.buf.va_addr + prg.samp_off, len(prg.samplers) * 4).cast('I')[:] = array.array('I', prg.samplers)
-    self.bind_sints_to_buf(*[b.va_addr for b in ubos], buf=self.buf, fmt='Q', offset=prg.buf_off)
-    self.bind_sints_to_buf(*vals, buf=self.buf, fmt='I', offset=prg.buf_off + len(ubos) * 8)
-    self.bind_sints_to_buf(*flatten([b.texture_info.desc + ([0] * 8) for b in texs]), buf=self.buf, fmt='I', offset=prg.tex_off)
-    self.bind_sints_to_buf(*flatten([b.texture_info.ibo + ([0] * 8) for b in ibos]), buf=self.buf, fmt='I', offset=prg.ibo_off)
+    self.bind_sints_to_buf(*flatten(map(_tex, texs)), buf=self.buf, fmt='I', offset=prg.tex_off)
+    self.bind_sints_to_buf(*flatten(map(functools.partial(_tex, ibo=True), ibos)), buf=self.buf, fmt='I', offset=prg.ibo_off)
 
 class QCOMProgram(HCQProgram):
   def __init__(self, dev: QCOMDevice, name: str, lib: bytes):
@@ -246,10 +237,11 @@ class QCOMProgram(HCQProgram):
 
       self.tex_off, self.ibo_off, self.samp_off = 2048, 2048 + 0x40 * self.tex_cnt, 2048 + 0x40 * (self.tex_cnt + self.ibo_cnt)
       self.fregs, self.hregs = v.info.max_reg + 1, v.info.max_half_reg + 1
+      self.consts_info:list[tuple] = []
     else: self._parse_lib()
 
     self.lib_gpu: HCQBuffer = self.dev.allocator.alloc(self.image_size, buf_spec:=BufferSpec(cpu_access=True, nolru=True))
-    to_mv(cast(int, self.lib_gpu.va_addr), self.image_size)[:] = self.image
+    to_mv(self.lib_gpu.va_addr, self.image_size)[:] = self.image
 
     self.pvtmem_size_per_item: int = round_up(self.pvtmem, 512) >> 9
     self.pvtmem_size_total: int = self.pvtmem_size_per_item * 128 * 2
@@ -259,7 +251,7 @@ class QCOMProgram(HCQProgram):
     dev._ensure_stack_size(self.hw_stack_offset * 4)
 
     kernargs_alloc_size = round_up(2048 + (self.tex_cnt + self.ibo_cnt) * 0x40 + len(self.samplers) * 4, 0x100)
-    super().__init__(IR3ArgsState if self.NIR else QCOMArgsState, self.dev, self.name, kernargs_alloc_size=kernargs_alloc_size)
+    super().__init__(QCOMArgsState, self.dev, self.name, kernargs_alloc_size=kernargs_alloc_size)
     weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)
 
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
@@ -279,7 +271,7 @@ class QCOMProgram(HCQProgram):
     self.pvtmem, self.shmem = _read_lib(self.lib, image_desc_off+0xc8), _read_lib(self.lib, image_desc_off+0xd8)
 
     # Fill up constants and buffers info
-    self.buf_info, self.consts_info = [], []
+    self.consts_info = []
 
     # Collect sampler info.
     self.samp_cnt = samp_cnt_in_file = _read_lib(self.lib, image_desc_off + 0xdc)
@@ -291,20 +283,17 @@ class QCOMProgram(HCQProgram):
     else: self.samplers = []
 
     # Collect kernel arguments (buffers) info.
-    bdoff = round_up(image_desc_off + 0x158 + len(self.name), 4) + 8 * samp_cnt_in_file
+    bdoff, binfos = round_up(image_desc_off + 0x158 + len(self.name), 4) + 8 * samp_cnt_in_file, []
     while bdoff + 32 <= len(self.lib):
-      length, _, _, offset_words, _, _, _, typ = struct.unpack("IIIIIIII", self.lib[bdoff:bdoff+32])
+      length, _, _, offset_words, _, _, _, typ = struct.unpack("8I", self.lib[bdoff:bdoff+32])
       if length == 0: break
-      self.buf_info.append(SimpleNamespace(offset=offset_words * 4, type=typ))
+      binfos.append((offset_words * 4, typ))
       bdoff += length
+    self.buf_offs = [off for off,typ in binfos if typ not in {BUFTYPE_TEX, BUFTYPE_IBO}]
 
     # Setting correct offsets to textures/ibos.
-    self.tex_cnt, self.ibo_cnt = sum(x.type is BUFTYPE_TEX for x in self.buf_info), sum(x.type is BUFTYPE_IBO for x in self.buf_info)
+    self.tex_cnt, self.ibo_cnt = sum(typ is BUFTYPE_TEX for _,typ in binfos), sum(typ is BUFTYPE_IBO for _,typ in binfos)
     self.ibo_off, self.tex_off, self.samp_off = 2048, 2048 + 0x40 * self.ibo_cnt, 2048 + 0x40 * self.tex_cnt + 0x40 * self.ibo_cnt
-    cur_ibo_off, cur_tex_off = self.ibo_off, self.tex_off
-    for x in self.buf_info:
-      if x.type is BUFTYPE_IBO: x.offset, cur_ibo_off = cur_ibo_off, cur_ibo_off + 0x40
-      elif x.type is BUFTYPE_TEX: x.offset, cur_tex_off = cur_tex_off, cur_tex_off + 0x40
 
     if _read_lib(self.lib, 0xb0) != 0: # check if we have constants.
       cdoff = _read_lib(self.lib, 0xac)
@@ -322,28 +311,10 @@ class QCOMTextureInfo:
     self.pitch, self.real_stride, self.desc, self.ibo = pitch, real_stride, desc, ibo
 
 class QCOMAllocator(HCQAllocatorBase):
-  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
+  def _alloc(self, size:int, opts:BufferSpec) -> HCQBuffer:
     # Recalculate real size for texture
-    if options.image is not None:
-      imgw, imgh, itemsize_log = options.image.shape[1], options.image.shape[0], int(math.log2(options.image.itemsize))
-      pitchalign = max(6, 11 - int(math.log2(imgh))) if imgh > 1 else 6
-      align_up = max(1, (8 // itemsize_log + 1) - imgh // 32) if pitchalign == 6 else (2 ** (pitchalign - itemsize_log - 2))
-
-      granularity = 128 if options.image.itemsize == 4 else 256
-      pitch_add = (1 << pitchalign) if min(next_power2(imgw), round_up(imgw, granularity)) - align_up + 1 <= imgw and imgw > granularity//2 else 0
-      pitch = round_up((real_stride:=imgw * 4 * options.image.itemsize), 1 << pitchalign) + pitch_add
-      size = pitch * imgh
-
-    buf = self.dev._gpu_map(options.external_ptr, size) if options.external_ptr else self.dev._gpu_alloc(size)
-
-    if options.image is not None:
-      tex_fmt = mesa.FMT6_32_32_32_32_FLOAT if options.image.itemsize == 4 else mesa.FMT6_16_16_16_16_FLOAT
-      desc = [qreg.a6xx_tex_const_0(0x8, swiz_x=0, swiz_y=1, swiz_z=2, swiz_w=3, fmt=tex_fmt), qreg.a6xx_tex_const_1(width=imgw, height=imgh),
-              qreg.a6xx_tex_const_2(type=mesa.A6XX_TEX_2D, pitch=pitch, pitchalign=pitchalign-6), 0,
-              *data64_le(buf.va_addr), qreg.a6xx_tex_const_6(plane_pitch=0x400000), qreg.a6xx_tex_const_7(13)]
-
-      buf.texture_info = QCOMTextureInfo(pitch, real_stride, desc, [desc[0] & (~0xffff), *desc[1:len(desc)]])
-    return buf
+    if opts.image is not None: size = opts.image.pitch* opts.image.shape[0]
+    return self.dev._gpu_map(opts.external_ptr, size, image=opts.image) if opts.external_ptr else self.dev._gpu_alloc(size, image=opts.image)
 
   def _do_copy(self, src_addr, dest_addr, src_size, real_size, src_stride, dest_stride, prof_text, dest_off=0, src_off=0):
     with cpu_profile(prof_text, self.dev.device, is_copy=True):
@@ -352,13 +323,13 @@ class QCOMAllocator(HCQAllocatorBase):
         src_off, dest_off = src_off+src_stride, dest_off+dest_stride
 
   def _copyin(self, dest:HCQBuffer, src:memoryview):
-    stride, pitch = (src.nbytes, src.nbytes) if (ti:=cast(QCOMTextureInfo, dest.texture_info)) is None else (ti.real_stride, ti.pitch)
+    stride, pitch = (dest.image.shape[1] * 4 * dest.image.itemsize, dest.image.pitch) if dest.image else (src.nbytes, src.nbytes)
     self._do_copy(mv_address(src), dest.cpu_view().addr, src.nbytes, stride, stride, pitch, f"TINY -> {self.dev.device}")
 
   def _copyout(self, dest:memoryview, src:HCQBuffer):
     self.dev.synchronize()
 
-    stride, pitch = (src.size, src.size) if (ti:=cast(QCOMTextureInfo, src.texture_info)) is None else (ti.real_stride, ti.pitch)
+    stride, pitch = (src.image.shape[1] * 4 * src.image.itemsize, src.image.pitch) if src.image else (src.size, src.size)
     self._do_copy(src.cpu_view().addr, mv_address(dest), src.size, stride, pitch, stride, f"{self.dev.device} -> TINY")
 
   def _as_buffer(self, src:HCQBuffer) -> memoryview:
@@ -405,7 +376,7 @@ class QCOMDevice(HCQCompiled):
     super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal,
                      functools.partial(QCOMComputeQueue, self), None)
 
-  def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> HCQBuffer:
+  def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False, **kwargs) -> HCQBuffer:
     flags |= flag("KGSL_MEMALIGN", alignment_hint:=12) | kgsl.KGSL_MEMFLAGS_USE_CPU_MAP
     if uncached: flags |= flag("KGSL_CACHEMODE", kgsl.KGSL_CACHEMODE_UNCACHED)
 
@@ -413,15 +384,15 @@ class QCOMDevice(HCQCompiled):
     va_addr = self.fd.mmap(0, bosz, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, alloc.id * 0x1000)
 
     if fill_zeroes: ctypes.memset(va_addr, 0, size)
-    return HCQBuffer(va_addr=va_addr, size=size, meta=(alloc, True), view=MMIOInterface(va_addr, size, fmt='B'), owner=self)
+    return HCQBuffer(va_addr=va_addr, size=size, meta=(alloc, True), view=MMIOInterface(va_addr, size, fmt='B'), owner=self, **kwargs)
 
-  def _gpu_map(self, ptr:int, size:int) -> HCQBuffer:
+  def _gpu_map(self, ptr:int, size:int, **kwargs) -> HCQBuffer:
     ptr_aligned, size_aligned = (ptr & ~0xfff), round_up(size + (ptr & 0xfff), 0x1000)
     try:
-      mapinfo = kgsl.IOCTL_KGSL_MAP_USER_MEM(self.fd, hostptr=ptr_aligned, len=size_aligned, memtype=kgsl.KGSL_USER_MEM_TYPE_ADDR)
-      return HCQBuffer(mapinfo.gpuaddr + (ptr - ptr_aligned), size=size, meta=(mapinfo, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self)
+      mi = kgsl.IOCTL_KGSL_MAP_USER_MEM(self.fd, hostptr=ptr_aligned, len=size_aligned, memtype=kgsl.KGSL_USER_MEM_TYPE_ADDR)
+      return HCQBuffer(mi.gpuaddr + (ptr - ptr_aligned), size=size, meta=(mi, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self, **kwargs)
     except OSError as e:
-      if e.errno == 14: return HCQBuffer(va_addr=ptr, size=size, meta=(None, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self)
+      if e.errno == 14: return HCQBuffer(va_addr=ptr, size=size, meta=(None, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self, **kwargs)
       raise RuntimeError("Failed to map external pointer to GPU memory") from e
 
   def _gpu_free(self, mem:HCQBuffer):

tinygrad/runtime/support/am/ip.py

@@ -281,9 +281,10 @@ class AM_GFX(AM_IP):
         self._grbm_select(inst=xcc)
     for xcc in range(self.xccs): self.adev.regGCVM_CONTEXT0_CNTL.write(0, inst=xcc)
 
-  def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, eop_addr:int, eop_size:int, doorbell:int, pipe:int, queue:int,
-                 aql:bool) -> int:
+  def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, eop_addr:int, eop_size:int, pipe:int, queue:int,
+                 aql:bool) -> tuple[int, int]:
     self._grbm_select(me=1, pipe=pipe, queue=queue, inst=0)
+    doorbell = am.AMDGPU_NAVI10_DOORBELL_MEC_RING0
     restore_queue = aql and self.xccs > 1 and self.adev.partial_boot and (self.adev.regCP_HQD_ACTIVE.read(inst=0) & 1)
     restore_ptr = (self.adev.regCP_HQD_PQ_WPTR_LO.read(inst=0) | (self.adev.regCP_HQD_PQ_WPTR_HI.read(inst=0) << 32)) if restore_queue else 0
     if DEBUG >= 2 and restore_queue: print(f"am {self.adev.devfmt}: GFX queue already active, continuing from saved state {restore_ptr=:#x}.")
@@ -327,7 +328,7 @@ class AM_GFX(AM_IP):
       self._grbm_select(inst=xcc)
 
       self.adev.reg(f"regCP_ME1_PIPE{pipe}_INT_CNTL").update(time_stamp_int_enable=1, generic0_int_enable=1, inst=xcc)
-    return restore_ptr // 16
+    return restore_ptr // 16, doorbell
 
   def set_clockgating_state(self):
     if hasattr(self.adev, 'regMM_ATC_L2_MISC_CG'): self.adev.regMM_ATC_L2_MISC_CG.write(enable=1, mem_ls_enable=1)
@@ -414,41 +415,43 @@ class AM_IH(AM_IP):
     self.adev.regIH_RB_RPTR.write(wptr % self.ring_size)
 
 class AM_SDMA(AM_IP):
-  def init_sw(self): self.sdma_name = "F32" if self.adev.ip_ver[am.SDMA0_HWIP] < (7,0,0) else "MCU"
+  def init_sw(self): self.sdma_reginst, self.sdma_name = [], "F32" if self.adev.ip_ver[am.SDMA0_HWIP] < (7,0,0) else "MCU"
   def init_hw(self):
     for pipe_id in range(1):
-      pipe = "" if self.adev.ip_ver[am.SDMA0_HWIP] < (5,0,0) else str(pipe_id)
+      pipe, inst = ("", pipe_id) if self.adev.ip_ver[am.SDMA0_HWIP] < (5,0,0) else (str(pipe_id), 0)
 
       if self.adev.ip_ver[am.SDMA0_HWIP] >= (6,0,0):
-        self.adev.reg(f"regSDMA{pipe}_WATCHDOG_CNTL").update(queue_hang_count=100) # 10s, 100ms per unit
-        self.adev.reg(f"regSDMA{pipe}_UTCL1_CNTL").update(resp_mode=3, redo_delay=9)
+        self.adev.reg(f"regSDMA{pipe}_WATCHDOG_CNTL").update(queue_hang_count=100, inst=inst) # 10s, 100ms per unit
+        self.adev.reg(f"regSDMA{pipe}_UTCL1_CNTL").update(resp_mode=3, redo_delay=9, inst=inst)
 
         # rd=noa, wr=bypass
-        self.adev.reg(f"regSDMA{pipe}_UTCL1_PAGE").update(rd_l2_policy=2, wr_l2_policy=3, **({'llc_noalloc':1} if self.sdma_name == "F32" else {}))
-        self.adev.reg(f"regSDMA{pipe}_{self.sdma_name}_CNTL").update(halt=0, **{f"{'th1_' if self.sdma_name == 'F32' else ''}reset":0})
+        self.adev.reg(f"regSDMA{pipe}_UTCL1_PAGE").update(rd_l2_policy=2, wr_l2_policy=3, **({'llc_noalloc':1} if self.sdma_name == "F32" else {}),
+                                                          inst=inst)
+        self.adev.reg(f"regSDMA{pipe}_{self.sdma_name}_CNTL").update(halt=0, **{f"{'th1_' if self.sdma_name == 'F32' else ''}reset":0}, inst=inst)
 
       self.adev.reg(f"regSDMA{pipe}_CNTL").update(ctxempty_int_enable=1, trap_enable=1,
-        **({'utc_l1_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP] <= (5,2,0) else {}))
+        **({'utc_l1_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP] <= (5,2,0) else {}), inst=inst)
 
     if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
-      self.adev.regDOORBELL0_CTRL_ENTRY_1.write(bif_doorbell1_range_offset_entry=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0*2,
-        bif_doorbell1_range_size_entry=4)
+      for i in range(16): self.adev.reg(f"regDOORBELL0_CTRL_ENTRY_{i+1}").write(**{f"bif_doorbell{i+1}_range_size_entry":4,
+                                        f"bif_doorbell{i+1}_range_offset_entry":(am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0 + i * 0xA) * 2})
       self.adev.soc.doorbell_enable(port=2, awid=0xe, awaddr_31_28_value=0x1, offset=0xe, size=4)
     else: self.adev.soc.doorbell_enable(port=2, awid=0xe, awaddr_31_28_value=0x3, offset=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0*2, size=4)
 
   def fini_hw(self):
-    reg, inst = ("regSDMA_GFX", 0) if self.adev.ip_ver[am.SDMA0_HWIP][:2] == (4,4) else ("regSDMA0_QUEUE0", 0)
+    for reg, inst in self.sdma_reginst:
+      self.adev.reg(f"{reg}_RB_CNTL").update(rb_enable=0, inst=inst)
+      self.adev.reg(f"{reg}_IB_CNTL").update(ib_enable=0, inst=inst)
 
-    self.adev.reg(f"{reg}_RB_CNTL").update(rb_enable=0, inst=inst)
-    self.adev.reg(f"{reg}_IB_CNTL").update(ib_enable=0, inst=inst)
     if self.adev.ip_ver[am.SDMA0_HWIP] >= (6,0,0):
       self.adev.regGRBM_SOFT_RESET.write(soft_reset_sdma0=1)
       time.sleep(0.01)
       self.adev.regGRBM_SOFT_RESET.write(0x0)
 
-  def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, doorbell:int, pipe:int, queue:int) -> int:
-    # Setup the ring
-    reg, inst = ("regSDMA_GFX", pipe*4+queue) if self.adev.ip_ver[am.SDMA0_HWIP][:2] == (4,4) else (f"regSDMA{pipe}_QUEUE{queue}", 0)
+  def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, pipe:int, queue:int) -> tuple[int, int]:
+    reg, inst = ("regSDMA_GFX", pipe+queue*4) if self.adev.ip_ver[am.SDMA0_HWIP][:2] == (4,4) else (f"regSDMA{pipe}_QUEUE{queue}", 0)
+    doorbell = am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0 + (pipe+queue*4) * 0xA
+    self.sdma_reginst.append((reg, inst))
 
     self.adev.reg(f"{reg}_MINOR_PTR_UPDATE").write(0x1, inst=inst)
     if not self.adev.partial_boot: self.adev.wreg_pair(f"{reg}_RB_RPTR", "", "_HI", 0, inst=inst)
@@ -462,7 +465,7 @@ class AM_SDMA(AM_IP):
     self.adev.reg(f"{reg}_RB_CNTL").write(**({f'{self.sdma_name.lower()}_wptr_poll_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP][:2]!=(4,4) else {}),
       rb_vmid=0, rptr_writeback_enable=1, rptr_writeback_timer=4, rb_enable=1, rb_priv=1, rb_size=(ring_size//4).bit_length()-1, inst=inst)
     self.adev.reg(f"{reg}_IB_CNTL").update(ib_enable=1, inst=inst)
-    return self.adev.reg(f"{reg}_RB_WPTR").read() | (self.adev.reg(f"{reg}_RB_WPTR_HI").read() << 32)
+    return self.adev.reg(f"{reg}_RB_WPTR").read(inst=inst) | (self.adev.reg(f"{reg}_RB_WPTR_HI").read(inst=inst) << 32), doorbell
 
 class AM_PSP(AM_IP):
   def init_sw(self):

tinygrad/runtime/support/hcq.py

@@ -8,6 +8,7 @@ from tinygrad.device import BufferSpec, Compiled, LRUAllocator, ProfileDeviceEve
 from tinygrad.uop.ops import sym_infer, sint, UOp
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.memory import BumpAllocator
+from tinygrad.dtype import ImageDType
 
 class MMIOInterface:
   def __init__(self, addr:int, nbytes:int, fmt='B'): self.mv, self.addr, self.nbytes, self.fmt = to_mv(addr, nbytes).cast(fmt), addr, nbytes, fmt
@@ -354,7 +355,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
   cpu_devices: list[HCQCompiled] = []
 
   def __init__(self, device:str, allocator:HCQAllocatorBase, compilers:CompilerSet, runtime, signal_t:Type[SignalType],
-               comp_queue_t:Callable[[], HWQueue], copy_queue_t:Callable[[], HWQueue]|None=None, kernargs_size=(16 << 20), sigalloc_size=0x1000):
+               comp_queue_t:Callable[..., HWQueue], copy_queue_t:Callable[..., HWQueue]|None=None, kernargs_size=(16 << 20), sigalloc_size=0x1000):
     self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
 
     from tinygrad.runtime.graph.hcq import HCQGraph
@@ -455,14 +456,14 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
 
 class HCQBuffer:
-  def __init__(self, va_addr:sint, size:int, texture_info:Any=None, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None,
+  def __init__(self, va_addr:sint, size:int, image:ImageDType|None=None, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None,
                owner:HCQCompiled|None=None):
-    self.va_addr, self.size, self.texture_info, self.meta, self._base, self.view = va_addr, size, texture_info, meta, _base, view
+    self.va_addr, self.size, self.image, self.meta, self._base, self.view = va_addr, size, image, meta, _base, view
     self._devs, self.owner = ([owner] if owner is not None else []), owner
     self._mappings:dict[HCQCompiled, HCQBuffer] = {} # mapping to the other devices
 
   def offset(self, offset:int=0, size:int|None=None) -> HCQBuffer:
-    return HCQBuffer(self.va_addr+offset, size or (self.size - offset), owner=self.owner, texture_info=self.texture_info, meta=self.meta,
+    return HCQBuffer(self.va_addr+offset, size or (self.size - offset), owner=self.owner, image=self.image, meta=self.meta,
       _base=self._base or self, view=(self.view.view(offset=offset, size=size) if self.view is not None else None))
 
   def cpu_view(self) -> MMIOInterface:

tinygrad/schedule/multi.py

@@ -1,6 +1,6 @@
 from typing import cast
 import functools, itertools, operator
-from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, getenv
+from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, ALL2ALL, getenv
 from tinygrad.uop.ops import Ops, UOp, sint, PatternMatcher, UPat, GroupOp, graph_rewrite_map, graph_rewrite
 from tinygrad.device import Device
 
@@ -35,45 +35,49 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
   if not isinstance(buf.device, tuple): return None
   assert all_int(buf.shape), f"does not support symbolic shape {buf.shape}"
   n_lbs, shape, numel = len(buf.device), buf.shape, prod(buf.shape)
+
   # ring allreduce doesn't provide a benefit with only 2 nodes or where number of elements is less than 256k (empirically)
   # fallback to naive allreduce to save on kernel dispatch, chunking and reassembling chunks.
-  use_ring = (RING >= 2 or (n_lbs > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1))
-  if DEBUG >= 2: print(f"{'RING ALLREDUCE' if use_ring else 'NAIVE ALLREDUCE'} {n_lbs}x{numel} | {buf.dtype}")
+  use_all2all = (ALL2ALL >= 2 or (n_lbs > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and ALL2ALL >= 1))
+  use_ring = not use_all2all and (RING >= 2 or (n_lbs > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1))
+  if DEBUG >= 2: print(f"{'ALL2ALL' if use_all2all else 'RING' if use_ring else 'NAIVE'} ALLREDUCE {n_lbs}x{numel} | {buf.dtype}")
 
   # contiguous before we copy it
   buf = buf.contiguous()
 
-  # copy to all devices. if you shrink later, that'll be handled
-  if not use_ring: return functools.reduce(lambda x,y: x.alu(red.arg, y),
-                                           [UOp(Ops.COPY, buf.dtype, (buf.mselect(i), red.src[1])) for i in range(len(buf.device))])
+  # naive: copy to all devices. if you shrink later, that'll be handled
+  if not use_ring and not use_all2all:
+    return functools.reduce(lambda x,y: x.alu(red.arg, y), [UOp(Ops.COPY, buf.dtype, (buf.mselect(i), red.src[1])) for i in range(n_lbs)])
 
-  # new ring reduce
+  # chunk data into n_lbs pieces
   factor = next((f for f in [32, 16, 8, 4, 2] if numel % f == 0), 1)
   base, left = (numel // factor) // n_lbs, (numel // factor) % n_lbs
-  chunk_sizes = [(base + 1) * factor] * left + [base * factor] * (n_lbs - left)
-  chunks = list(itertools.pairwise(itertools.accumulate(chunk_sizes, initial=0)))
+  chunks = list(itertools.pairwise(itertools.accumulate([(base + 1) * factor] * left + [base * factor] * (n_lbs - left), initial=0)))
 
-  # extract chunks and scatter-reduce
+  # reduce-scatter
   reduced_chunks = []
   for i,(s,e) in enumerate(chunks):
-    chunk = buf.reshape((numel,)).shrink(((s,e),))
-    reduced_chunk = chunk
-    for step in range(n_lbs-1):
-      src, dest = (i+step)%n_lbs, (i+step+1)%n_lbs
-      # copy the chunk from the src device to the dest (operating device), and select the chunk on the dest device
-      reduced_chunk = reduced_chunk.copy_to_device(buf.device[dest], src if isinstance(reduced_chunk.device, tuple) else None) \
-        .alu(red.arg, chunk.copy_to_device(buf.device[dest], dest))
-    reduced_chunks.append(reduced_chunk)
+    if use_all2all:
+      chunks_on_i = [buf.mselect(j).reshape((numel,)).shrink(((s,e),)).copy_to_device(buf.device[i]) for j in range(n_lbs)]
+      reduced_chunks.append(functools.reduce(lambda x,y: x.alu(red.arg, y), chunks_on_i))
+    else:
+      chunk, reduced = buf.reshape((numel,)).shrink(((s,e),)), buf.reshape((numel,)).shrink(((s,e),))
+      for step in range(n_lbs-1):
+        src, dest = (i+step)%n_lbs, (i+step+1)%n_lbs
+        cp = reduced.copy_to_device(buf.device[dest], src if isinstance(reduced.device, tuple) else None)
+        reduced = cp.alu(red.arg, chunk.copy_to_device(buf.device[dest], dest))
+      reduced_chunks.append(reduced)
 
   # allgather
   copied_chunks = []
-  for i,c in enumerate(reduced_chunks):
-    this_chunk: list[UOp|None] = [None] * len(buf.device)
-    this_chunk[(i+len(buf.device)-1)%n_lbs] = c
-    for step in range(n_lbs-1):
-      dest = (i+step)%n_lbs
-      this_chunk[dest] = c = c.copy_to_device(buf.device[dest])
-    copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(cast(list[UOp], this_chunk))))
+  for i,rc in enumerate(reduced_chunks):
+    if use_all2all: copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(rc.copy_to_device(buf.device[j]) for j in range(n_lbs))))
+    else:
+      this_chunk: list[UOp|None] = [None] * n_lbs
+      this_chunk[(i+n_lbs-1)%n_lbs] = rc
+      for step in range(n_lbs-1):
+        this_chunk[(i+step)%n_lbs] = rc = rc.copy_to_device(buf.device[(i+step)%n_lbs])
+      copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(cast(list[UOp], this_chunk))))
 
   # reassemble
   pads = [((s,numel-e),) for s,e in chunks]

tinygrad/tensor.py

@@ -127,7 +127,7 @@ class Tensor(OpMixin):
 
     # create a UOp from the different types of inputs
     if isinstance(data, UOp):
-      assert _dtype is None or _dtype==data.dtype, f"dtype doesn't match ({_dtype} vs {data.dtype}), and casting isn't supported"
+      assert _dtype is None or _dtype==data.dtype or data.dtype==dtypes.index, f"dtype mismatch: {_dtype} vs {data.dtype}"
       # if data is dtype.index that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of
       if data.dtype==dtypes.index: data = _index_to_concrete_int(data)
       if data.op is Ops.BIND:  # type: ignore  # mypy type narrowing is bugged here
@@ -3637,11 +3637,13 @@ class Tensor(OpMixin):
     Q = Tensor.eye(m, dtype=self.dtype).reshape((1,) * len(b_shape) + (m, m)).expand(b_shape + (m, m)).contiguous()
     for i in range(min(m, n)):
       x = R[..., i:m, i].contiguous()  # TODO: without contigous this can silently be wrong, should at least assert
-      s = -x[..., 0].sign()
-      u1 = x[..., 0] - s * x.square().sum(-1).sqrt()
-      w = x.unsqueeze(-1) / u1.reshape(b_shape + (1, 1))
+      norm = x.square().sum(-1).sqrt()
+      s = (x[..., 0] != 0).where(-x[..., 0].sign(), -1)
+      u1 = x[..., 0] - s * norm
+      w = x.unsqueeze(-1) / (norm != 0).where(u1, 1).reshape(b_shape + (1, 1))
       w[..., 0, 0] = 1
-      tau = (-s * u1 / x.square().sum(-1).sqrt()).reshape(b_shape + (1, 1))
+      tau = (-s * u1 / (norm != 0).where(norm, 1)).reshape(b_shape + (1, 1))
+      tau = (norm != 0).reshape(b_shape + (1, 1)).where(tau, 0)
       R[..., i:m, :] = R[..., i:m, :] - (w * tau) @ (w.transpose(-2, -1) @ R[..., i:m, :])
       Q[..., :, i:m] = Q[..., :, i:m] - (Q[..., :, i:m] @ w) @ (tau * w).transpose(-2, -1)
     return Q,R
@@ -3668,8 +3670,10 @@ class Tensor(OpMixin):
       #compute the jacobi rotations for each pairing
       gamma = (U_left * U_right).sum(-2).reshape(b_shape + (1, num//2))
       alpha, beta = U_permuted.square().sum(-2).unsqueeze(-2).split(num//2, -1)
-      tau = (beta - alpha) / (2 * gamma)
-      t = tau.sign() / (tau.abs() + (1 + tau.square()).sqrt())
+      rot = gamma != 0
+      tau = (beta - alpha) / (2 * rot.where(gamma, 1))
+      t = (tau != 0).where(tau.sign(), 1) / (tau.abs() + (1 + tau.square()).sqrt())
+      t = rot.where(t, 0)
       c = 1 / (1 + t.square()).sqrt()
       s = c * t
       #apply the rotations
@@ -3686,9 +3690,9 @@ class Tensor(OpMixin):
     for _ in range(max_iterations * iterations_per_round): U, V, permute, inverse_permute = one_round_jacobi(U, V, permute, inverse_permute)
     #extract singular values and sort. construct U from Q
     S, indices = U.square().sum(-2).sqrt().sort(dim = -1, descending=True)
-    new_indices = Tensor.arange(num).reshape((1,) * (self.ndim - 1) + (num,)).expand(b_shape + (num, num)).contiguous()
-    new_indices[..., :num] = indices.reshape(b_shape + (1, num)).expand(b_shape + (num, num))
-    U, V = U.gather(-1, new_indices[...,0:num,0:num]) / S.unsqueeze(-2), V.gather(-1, new_indices[..., 0:num, 0:num]).realize()
+    new_indices = indices.reshape(b_shape + (1, num)).expand(b_shape + (num, num))
+    U = U.gather(-1, new_indices) / (S != 0).where(S, 1).unsqueeze(-2)
+    V = V.gather(-1, new_indices).realize()
 
     padded_u = Tensor.eye(q_num, dtype=U.dtype).reshape((1,) * len(b_shape) + (q_num, q_num)).expand(b_shape + (q_num, q_num)).contiguous()
     padded_u[..., 0:num, 0:num] = U

tinygrad/uop/decompositions.py

@@ -223,26 +223,26 @@ def xlog2(d:UOp) -> UOp:
   Paper: https://arxiv.org/pdf/2001.09258 5.5
   """
   assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES
-  # TODO: float16 denormal need float32 to achieve precision
-  if d.dtype.scalar() == dtypes.float16: return xlog2(d.cast(dtypes.float32)).cast(dtypes.float16)
-  FLT_MIN = d.const_like(1e-6 if d.dtype.scalar() == dtypes.float16 else 1e-4)
+  # float16 uses 2^10 for denormal scaling (2^64 overflows), float32/64 use 2^64
+  denormal_exp = 10 if d.dtype.scalar() == dtypes.float16 else 64
+  FLT_MIN = d.const_like({dtypes.float16: 6.1e-5, dtypes.float32: 1e-4, dtypes.float64: 1e-4}[d.dtype.scalar()])
   is_denormal = d<FLT_MIN
-  a = is_denormal.where(d * (2 ** 64), d)
+  a = is_denormal.where(d * (2 ** denormal_exp), d)
 
   e = ilogb2k(a * (1.0 / 0.75)).cast(a.dtype)
   m = ldexp3k(a, -e)
-  e = is_denormal.where(e - 64, e)
+  e = is_denormal.where(e - denormal_exp, e)
 
   x = (m - 1.0) / (m + 1.0)
   x2 = x * x
   if d.dtype.scalar() == dtypes.float64:
     t = polyN(x2, [0.2211941750456081490e+0, 0.2200768693152277689e+0, 0.2623708057488514656e+0, 0.3205977477944495502e+0,
                    0.4121985945485324709e+0, 0.5770780162997058982e+0, 0.96179669392608091449])
-    s_hi, s_lo = e+x*2.885390081777926774, e.const_like(0)
+    r = t * (x * x2) + e + x * 2.885390081777926774
   else:
     t = polyN(x2, [0.4374550283e+0, 0.5764790177e+0, 0.9618012905120])
-    s_hi, s_lo = e+x*2.8853900432586669922, x*3.2734474483568488616e-08
-  r = t * (x * x2) + (s_hi + s_lo)
+    # s_lo term (x*3.27e-08) only for float32 - underflows in float16
+    r = t * (x * x2) + e + x * 2.8853900432586669922 + (x * 3.2734474483568488616e-08 if d.dtype.scalar() == dtypes.float32 else 0)
 
   # log2(Inf) = Inf
   r = d.ne(math.inf).where(r, r.const_like(math.inf))

tinygrad/uop/validate.py

@@ -1,6 +1,6 @@
 from typing import Callable, cast
 from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, python_alu
-from tinygrad.dtype import ImageDType, dtypes, Invalid
+from tinygrad.dtype import ImageDType, dtypes, Invalid, PtrDType
 from tinygrad.helpers import IGNORE_OOB, cpu_profile
 
 try:
@@ -11,9 +11,8 @@ try:
   # IDIV is truncated division but z3 does euclidian division (floor if b>0 ceil otherwise); mod by power of two sometimes uses Ops.AND
   def z3_cdiv(a, b):return z3.If((a<0), z3.If(0<b, (a+(b-1))/b, (a-(b+1))/b), a/b)
   def z3_xor(a,b):
-    if isinstance(a, z3.BoolRef): return a^b
-    assert a==-1 or b==-1, "xor can only be used in indexing if one of the arguments is -1"
-    return -a-1 if b==-1 else -b-1
+    assert isinstance(a, z3.BoolRef), f"{type(a)=}, {a=}"
+    return a^b
   z3_alu: dict[Ops, Callable] = python_alu | {Ops.MOD: lambda a,b: a-z3_cdiv(a,b)*b, Ops.IDIV: z3_cdiv, Ops.SHR: lambda a,b: a/(2**b.as_long()),
     Ops.SHL: lambda a,b: a*(2**b.as_long()), Ops.AND: lambda a,b: a%(b+1) if isinstance(b, z3.ArithRef) else a&b, Ops.WHERE: z3.If, Ops.XOR: z3_xor,
     Ops.MAX: lambda a,b: z3.If(a<b, b, a),}
@@ -34,7 +33,6 @@ try:
     (UPat(Ops.CONST, dtypes.ints+(dtypes.index,), name="x"), lambda x,ctx: (z3.IntVal(x.arg, ctx=ctx[0].ctx), None)),
     (UPat(Ops.CONST, dtypes.bool, name="x"), lambda x,ctx: (z3.BoolVal(x.arg, ctx=ctx[0].ctx), None)),
     # casts from floats create new variables
-    (UPat(Ops.CAST, dtypes.bool, src=(UPat(dtype=dtypes.floats),), name="x"), lambda x,ctx: (z3.Bool(f"cast{len(ctx[1])}",ctx=ctx[0].ctx), None)),
     (UPat(Ops.CAST, dtypes.ints+(dtypes.index,), src=(UPat(dtype=dtypes.floats),), name="x"), lambda x,ctx:
       create_bounded(f"cast{len(ctx[1])}", x.dtype.min, x.dtype.max, ctx[0])),
     # A comparison between floats introduces a new bool variable
@@ -67,8 +65,10 @@ def validate_index(buf:UOp, idx:UOp, gate:UOp|None=None):
   # We can use UOp min/max to do a faster check, but it can give false positive since its not an exact bound and doesn't consider the mask
   if 0<=idx.vmin and idx.vmax<sz: return True
 
-  # WEBGPU has a BITCAST in the index. TODO: fix
-  if any(x.op is Ops.BITCAST for x in idx.toposort()): return True
+  # TODO: validate these
+  # WEBGPU has a BITCAST in the index, PTX casts pointer to long
+  for x in idx.toposort() | gate.toposort():
+    if x.op is Ops.BITCAST or (x.op is Ops.CAST and isinstance(x.src[0].dtype, PtrDType)): return True
 
   if not z3_imported: raise ImportError("bounds checking requires z3 >= 4.12.4, use IGNORE_OOB=1 to disable, or \"pip install 'z3-solver>=4.12.4\"")
   solver = z3.Solver(ctx=z3.Context())