Merge branch 'master' into precompile_backward

combine no_vectorized_index and no_vectorized_index_broadcast

.github/actions/setup-tinygrad/action.yml

@@ -45,6 +45,10 @@ inputs:
     description: "Install mesa"
     required: false
     default: 'false'
+  tinydreno:
+    description: "Install tinydreno"
+    required: false
+    default: 'false'
 runs:
   using: "composite"
   steps:
@@ -62,14 +66,14 @@ runs:
       uses: actions/cache/restore@v4
       with:
         path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
     - name: Cache Python packages
       if: github.event_name != 'pull_request'
       id: restore-venv
       uses: actions/cache@v4
       with:
         path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
 
     # **** Caching downloads ****
 
@@ -195,13 +199,13 @@ runs:
       uses: actions/cache/restore@v4
       with:
         path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
     - name: Cache apt
       if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name != 'pull_request'
       uses: actions/cache@v4
       with:
         path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
 
     - name: Run apt Update + Install
       if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
@@ -233,7 +237,7 @@ runs:
       shell: bash
       run: |
         sudo mkdir -p /usr/local/lib
-        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/nimlgen/amdcomgr_dylib/releases/latest | \
+        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/tinygrad/amdcomgr_dylib/releases/latest | \
           jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
           sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
         cargo build --release --manifest-path ./extra/remu/Cargo.toml
@@ -326,3 +330,9 @@ runs:
       if: inputs.mesa == 'true' && runner.os == 'macOS'
       shell: bash
       run: brew install sirhcm/tinymesa/tinymesa_cpu
+
+    # *** tinydreno ***
+    - name: Install tinydreno (linux)
+      if: inputs.tinydreno == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so

.github/workflows/autogen.yml

@@ -32,6 +32,7 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
+        key: 'autogen'
         opencl: 'true'
         amd: 'true'
         cuda: 'true'
@@ -81,6 +82,7 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
+        key: 'autogen-mac'
         llvm: 'true'
     - name: Regenerate autogen files
       run: |
@@ -110,6 +112,8 @@ jobs:
       uses: actions/checkout@v4
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
+      with:
+        key: 'autogen-comgr'
     - name: Install autogen support packages
       run: |
         wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null

.github/workflows/benchmark.yml

@@ -332,7 +332,7 @@ jobs:
     # - name: Fuzz Padded Tensor Core GEMM (PTX)
     #   run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
     - name: HEVC Decode Benchmark
-      run: VALIDATE=1 MAX_FRAMES=100 JITBEAM=1 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
+      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 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
     - name: Run 10 CIFAR training steps
@@ -520,7 +520,7 @@ jobs:
       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
-    # this needs to be mocked and testable on a local machine
+    # TODO: broken on some of the machines
     #- name: Test full tinyfs load
     #  run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
     - name: Run process replay tests
@@ -617,6 +617,27 @@ jobs:
     - 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
 
+  testcommausbgpubenchmark:
+    name: UsbGPU Benchmark (comma)
+    runs-on: [self-hosted, Linux, comma4]
+    timeout-minutes: 20
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: openpilot compile3 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot load_pickle 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
+
   testreddriverbenchmark:
     name: AM Benchmark
     runs-on: [self-hosted, Linux, tinyboxrandom]

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '17'
+  CACHE_VERSION: '18'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
@@ -244,6 +244,37 @@ jobs:
     - name: Run TYPED=1
       run: CHECK_OOB=0 DEV=CPU TYPED=1 python test/test_tiny.py
 
+  nulltest:
+    name: Null Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: unittest-13
+        pydeps: "pillow ftfy regex pre-commit"
+        deps: testing_unit
+        llvm: 'true'
+        amd: 'true'
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
+    - name: Run targetted tests on NULL backend
+      run: NULL=1 python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
+    # TODO: too slow
+    # - name: Run SDXL on NULL backend
+    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
+    - name: Run Clip tests for SD MLPerf on NULL backend
+      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
+    - name: Run AMD emulated BERT training on NULL backend
+      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=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
+    # TODO: support fake weights
+    #- name: Run LLaMA 7B on 4 fake devices
+    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
+
   unittest:
     name: Unit Tests
     runs-on: ubuntu-latest
@@ -268,20 +299,6 @@ jobs:
       run: |
         CPU=1 python test/null/test_device.py TestRunAsModule.test_module_runs
         CPU=1 python -m pytest -n=auto test/unit/ --durations=20
-    - name: Run NULL backend tests
-      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
-    - name: Run targetted tests on NULL backend
-      run: NULL=1 python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
-    # TODO: too slow
-    # - name: Run SDXL on NULL backend
-    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
-    - name: Run Clip tests for SD MLPerf on NULL backend
-      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
-    - name: Run AMD emulated BERT training on NULL backend
-      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=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
-    # TODO: support fake weights
-    #- name: Run LLaMA 7B on 4 fake devices
-    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
     - name: Run GC tests
       run: python test/external/external_uop_gc.py
     - name: External Benchmark Schedule
@@ -644,7 +661,7 @@ jobs:
           sudo apt-get update
           sudo apt-get install llvm-21 llvm-21-tools cloc
       - name: Install rocprof-trace-decoder
-        run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
+        run: sudo PYTHONPATH="." ./extra/sqtt/install_rocprof_decoder.py
       - name: Run AMD renderer tests
         run: AMD_LLVM=0 python -m pytest -n=auto test/amd/ --durations 20
       - name: Run AMD renderer tests (AMD_LLVM=1)
@@ -994,3 +1011,26 @@ jobs:
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
           DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
           python -m pytest -n=auto test/backend/test_ops.py --durations=20
+  qcomclcompiletests:
+    name: Compile-only (QCOM CL)
+    runs-on: ubuntu-24.04-arm
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: compile-qcomcl
+          deps: testing_unit
+          tinydreno: 'true'
+          python-version: '3.12'
+      - name: Set env
+        shell: bash
+        run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\nNULL_QCOMCL=1" >> $GITHUB_ENV
+      - name: Run test_ops
+        shell: bash
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
+          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20

examples/mlperf/dataloader.py

@@ -396,6 +396,7 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
       queue_in.put((idx, img, tgt))
 
   def _setup_shared_mem(shm_name:str, size:tuple[int, ...], dtype:dtypes) -> tuple[shared_memory.SharedMemory, Tensor]:
+    shm_name = f"{shm_name}_{os.getpid()}"
     if os.path.exists(f"/dev/shm/{shm_name}"): os.unlink(f"/dev/shm/{shm_name}")
     shm = shared_memory.SharedMemory(name=shm_name, create=True, size=prod(size))
     shm_tensor = Tensor.empty(*size, dtype=dtype, device=f"disk:/dev/shm/{shm_name}")

examples/mlperf/model_train.py

@@ -3,7 +3,7 @@ from pathlib import Path
 import multiprocessing
 
 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
-from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker, DEBUG
 from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
 
@@ -1282,7 +1282,7 @@ def train_bert():
         previous_step = i
 
 def train_llama3():
-  from extra.models.llama import Transformer
+  from examples.mlperf.models.llama import Transformer
   from examples.llama3 import MODEL_PARAMS
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
   from examples.mlperf.optim import GradAccClipAdamW
@@ -1335,10 +1335,16 @@ def train_llama3():
   model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
   # vocab_size from the mixtral tokenizer
   if not SMALL: model_params |= {"vocab_size": 32000}
+  real_vocab_size = model_params['vocab_size']
   if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
   print(f"model parameters: {model_params}")
 
-  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
+  # pad vocab
+  if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
+  vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
+
+  model = Transformer(**model_params, max_context=SEQLEN)
+
   params = get_parameters(model)
   # weights are all bfloat16 for now
   assert params and all(p.dtype == dtypes.bfloat16 for p in params)
@@ -1352,6 +1358,8 @@ def train_llama3():
     for v in get_parameters(model):
       v.shard_(device, axis=None)
 
+    vocab_mask.shard_(device, axis=None)
+
   if (MP := getenv("MP", 1)) > 1:
     device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
     for k,v in get_state_dict(model).items():
@@ -1359,6 +1367,7 @@ def train_llama3():
       elif '.attention.wq' in k: v.shard_(device, axis=0)
       elif '.attention.wk' in k: v.shard_(device, axis=0)
       elif '.attention.wv' in k: v.shard_(device, axis=0)
+      elif '.attention.wqkv' in k: v.shard_(device, axis=0)
       elif '.attention.wo' in k: v.shard_(device, axis=1)
       elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
       elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
@@ -1371,13 +1380,21 @@ def train_llama3():
       # prevents memory spike on device 0
       v.realize()
 
-  optim_device = "CPU" if getenv("OFFLOAD_OPTIM") else None
+    vocab_mask.shard_(device, axis=2).realize()
+
+  is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
+  is_fake_offload = Device.DEFAULT == "NULL"
+  optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
   optim = GradAccClipAdamW(get_parameters(model), lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
                            eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
 
   # init grads
-  for p in optim.params:
-    p.grad = p.zeros_like().contiguous().realize()
+  if is_offload_optim:
+    for p in optim.params:
+      p.grad = Tensor.zeros(p.shape, dtype=p.dtype, device=optim_device, requires_grad=False).contiguous().realize()
+  else:
+    for p in optim.params:
+      p.grad = p.zeros_like().contiguous().realize()
   grads: list[Tensor] = [p.grad for p in optim.params]
 
   scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
@@ -1393,51 +1410,55 @@ def train_llama3():
 
   @TinyJit
   def minibatch(tokens:Tensor):
-    tokens = tokens.to(None)
     if (DP := getenv("DP", 1)) > 1:
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.shard(device, 0)
+      tokens = tokens.to(None).shard(device, 0)
     if (MP := getenv("MP", 1)) > 1:
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
       tokens = tokens.shard(device)
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    if DP == 1 and MP == 1: tokens = tokens.to(None)
+    logits:Tensor = model(tokens[:, :-1])
+    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
     loss.backward()
     assert all(p.grad is g for p,g in zip(optim.params, grads))
-    Tensor.realize(loss, *grads)
-    return loss.flatten().float().to("CPU")
+    loss_cpu = loss.flatten().float().to("CPU")
+    Tensor.realize(loss_cpu, *grads)
+    return loss_cpu
 
   @TinyJit
   def optim_step():
-    optim.step()
+    grad_norm = optim.fstep(grads)
     scheduler.step()
 
     for g in grads:
       g.assign(g.zeros_like())
 
-    lr = optim.lr
-    Tensor.realize(lr, *grads)
+    lr_cpu = optim.lr.float().to("CPU")
+    grad_norm_cpu = grad_norm.float().to("CPU")
+    Tensor.realize(lr_cpu, grad_norm_cpu, *grads)
 
-    return lr.float().to("CPU")
+    return lr_cpu, grad_norm_cpu
 
   @TinyJit
   @Tensor.train(False)
   def eval_step(tokens:Tensor):
-    tokens = tokens.to(None)
     if (DP := getenv("DP", 1)) > 1:
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.shard(device, 0)
+      tokens = tokens.to(None).shard(device, 0)
     if (MP := getenv("MP", 1)) > 1:
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
       tokens = tokens.shard(device)
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    if DP == 1 and MP == 1: tokens = tokens.to(None)
+    logits:Tensor = model(tokens[:, :-1])
+    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
     return loss.flatten().float().to("CPU")
 
   # ** data iters **
   def fake_data(bs, samples):
+    import numpy as np
     for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
+      fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
+      yield Tensor(fake_data_np, device="NPY")
 
   def get_train_iter():
     if getenv("FAKEDATA", 0):
@@ -1464,49 +1485,53 @@ def train_llama3():
   step_times = []
   while i < MAX_STEPS:
     GlobalCounters.reset()
+    actual_gbs = GBS if i >= 2 else BS
     if getenv("TRAIN", 1):
       profile_marker(f"train @ {i}")
       st = time.perf_counter()
 
       stopped = False
-      for _ in range(grad_acc):
+      losses, data_time, dev_time = [], 0, 0
+      for _ in range(grad_acc if i >= 2 else 1):
         ist = time.perf_counter()
         try: tokens = next(train_iter)
         except StopIteration:
           stopped = True
           break
-        dt = time.perf_counter()
-        loss = minibatch(tokens)
+        mst = time.perf_counter()
+        data_time += mst - ist
+        losses.append(minibatch(tokens).item())
+        dev_time += time.perf_counter() - mst
       if stopped: break
 
       gt = time.perf_counter()
-      lr = optim_step()
-      ot = time.perf_counter()
-
-      loss = loss.float().item()
-      lr = lr.item()
-
+      ret = optim_step()
+      lr, grad_norm = ret[0].item(), ret[1].item()
       et = time.perf_counter()
+
+      loss = sum(losses) / len(losses)
+      optim_time = et - gt
+      dev_time += optim_time
       step_time = et - st
       gbs_time = gt - st
-      optim_time = ot - gt
-      data_time = dt - ist
-      dev_time = step_time - data_time * grad_acc
       if BENCHMARK: step_times.append(step_time)
 
       i += 1
-      sequences_seen += GBS
+      sequences_seen += actual_gbs
 
       mem_gb = GlobalCounters.mem_used / 1e9
       gflops = GlobalCounters.global_ops / 1e9 / dev_time
       mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * max(getenv("DP", 1), getenv("MP", 1)) * 2.3e15)) * 100
       tqdm.write(
           f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
-          f"{lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
+          f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
+      if DEBUG >= 1: tqdm.write("  mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
 
       if WANDB:
         wandb.log({
-          "lr": lr, "train/loss": loss,
+          "train/loss": loss,
+          "train/lr": lr,
+          "train/grad_norm": grad_norm,
           "train/step_time": step_time,
           "train/gbs_time": gbs_time,
           "train/optim_time": optim_time,
@@ -1535,7 +1560,7 @@ def train_llama3():
         print(f"epoch global_ops: {GlobalCounters.global_ops:_}, "
               f"epoch global_mem: {GlobalCounters.global_mem:_}")
 
-    if (sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
+    if (sequences_seen // EVAL_FREQ != (sequences_seen - actual_gbs) // EVAL_FREQ and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
       if EVAL_BS == 0: return
       tqdm.write(f"evaluating after {sequences_seen} sequences")
       profile_marker(f"eval @ {i}")
@@ -1543,7 +1568,7 @@ def train_llama3():
       # run eval
       eval_losses = []
       eval_iter = get_eval_iter()
-      tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
+      tqdm.write(f"evaluating {EVAL_SAMPLES//EVAL_BS} batches of {EVAL_BS} sequences")
 
       for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
         eval_losses += eval_step(tokens).tolist()

examples/mlperf/models/llama.py

@@ -0,0 +1,80 @@
+from tinygrad import Tensor, nn
+from tinygrad.helpers import getenv
+from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
+
+class Attention:
+  def __init__(self, dim:int, n_heads:int, n_kv_heads:int|None=None, linear=nn.Linear):
+    self.n_heads = n_heads
+    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
+    self.head_dim = dim // n_heads
+    self.n_rep = self.n_heads // self.n_kv_heads
+
+    if getenv("WQKV"):
+      self.wqkv = linear(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2, bias=False)
+    else:
+      self.wq = linear(dim, self.n_heads * self.head_dim, bias=False)
+      self.wk = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+      self.wv = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+
+    self.wo = linear(self.n_heads * self.head_dim, dim, bias=False)
+
+  def __call__(self, x:Tensor, freqs_cis:Tensor) -> Tensor:
+    if getenv("WQKV"):
+      xqkv = self.wqkv(x)
+      xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
+      xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+    else:
+      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+
+    xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
+    xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
+    xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
+
+    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+    bsz, seqlen, _, _ = xq.shape
+
+    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
+    attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
+
+    attn = attn.reshape(bsz, seqlen, -1)
+    return self.wo(attn)
+
+class FeedForward:
+  def __init__(self, dim:int, hidden_dim:int, linear=nn.Linear):
+    self.w1 = linear(dim, hidden_dim, bias=False)
+    self.w2 = linear(hidden_dim, dim, bias=False)
+    self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit
+
+  def __call__(self, x:Tensor) -> Tensor:
+    w1 = self.w1(x).silu()
+    w3 = self.w3(x)
+    return self.w2(w1 * w3)
+
+class TransformerBlock:
+  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int|None, norm_eps:float, linear=nn.Linear):
+    self.attention = Attention(dim, n_heads, n_kv_heads, linear)
+    self.feed_forward = FeedForward(dim, hidden_dim, linear)
+    self.attention_norm = nn.RMSNorm(dim, norm_eps)
+    self.ffn_norm = nn.RMSNorm(dim, norm_eps)
+
+  def __call__(self, x:Tensor, freqs_cis:Tensor):
+    h = x + self.attention(self.attention_norm(x), freqs_cis)
+    return h + self.feed_forward(self.ffn_norm(h))
+
+class Transformer:
+  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
+               rope_theta:int=10000, max_context:int=1024, linear=nn.Linear, embedding=nn.Embedding):
+    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, linear) for _ in range(n_layers)]
+    self.norm = nn.RMSNorm(dim, norm_eps)
+    self.tok_embeddings = embedding(vocab_size, dim)
+    self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
+
+  def __call__(self, tokens:Tensor):
+    h = self.tok_embeddings(tokens)
+    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
+    for layer in self.layers: h = layer(h, freqs_cis)
+    logits = self.output(self.norm(h))
+    return logits

examples/mlperf/optim.py

@@ -7,40 +7,51 @@ class GradAccClipAdamW(Optimizer):
   def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
     super().__init__(params, lr, device, fused)
     self.b1, self.b2, self.eps, self.wd = b1, b2, eps, weight_decay
-    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False).contiguous() for _ in [b1, b2])
+    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False) for _ in [b1, b2])
     self.m = self._new_optim_param()
     self.v = self._new_optim_param()
     self.grad_acc, self.clip_norm = grad_acc, clip_norm
 
+  def fstep(self, grads:list[Tensor]):
+    if self.fused:
+      out, extra = self._step([], grads)
+      updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
+    else:
+      updates, extra = self._step([], grads)
+    for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i]))
+    to_realize = extra+self.params+self.buffers
+
+    Tensor.realize(*to_realize)
+    return extra[-1]
+
   def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
     for i in range(len(grads)):
-      if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
+      if grads[i].device != self.m[i].device: grads[i].assign(grads[i].to(self.m[i].device))
 
     if self.fused:
-      grads[0] = grads[0] / self.grad_acc
+      grads[0].assign(grads[0] / self.grad_acc)
       total_norm = grads[0].float().square().sum().sqrt()
-      grads[0] = (grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype)
+      grads[0].assign((grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype))
     else:
-      total_norm = Tensor.zeros((), dtype=dtypes.float32, device=self.device)
-      for g in grads:
-        total_norm += g.float().square().sum()
-      total_norm = total_norm.sqrt()
       for i in range(len(grads)):
-        grads[i] = grads[i] / self.grad_acc
-        grads[i] = (grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)
+        grads[i].assign(grads[i] / self.grad_acc)
+      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
+      for i in range(len(grads)):
+        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype))
 
     ret = []
     self.b1_t *= self.b1
     self.b2_t *= self.b2
-    for i, (t, g) in enumerate(zip(params, grads)):
+    for i, g in enumerate(grads):
       self.m[i].assign((self.b1 * self.m[i] + (1.0 - self.b1) * g).cast(self.m[i].dtype))
       self.v[i].assign((self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).cast(self.v[i].dtype))
-      m_hat = self.m[i] / (1.0 - self.b1_t)
-      v_hat = self.v[i] / (1.0 - self.b2_t)
+      m_hat = (self.m[i] / (1.0 - self.b1_t)).cast(self.m[i].dtype)
+      v_hat = (self.v[i] / (1.0 - self.b2_t)).cast(self.v[i].dtype)
       up = m_hat / (v_hat.sqrt() + self.eps)
-      ret.append((self.lr * up).cast(t.dtype))
-    return ret, [self.b1_t, self.b2_t] + self.m + self.v
+      ret.append((self.lr * up).cast(g.dtype))
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
 
   def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
-    up = up.shard_like(t) + self.lr.to(t.device) * self.wd * t.detach()
+    wd = self.wd if t.ndim >= 2 else 0.0
+    up = up.shard_like(t) + self.lr.to(t.device) * wd * t.detach()
     return t.detach() - up.cast(t.dtype)

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama405b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -0,0 +1,37 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+
+export DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8}
+export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4/"
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"405B"}
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+export FAKEDATA=1 BENCHMARK=10
+if [ -z "$FULL_LAYERS" ]; then
+  export LLAMA_LAYERS=2
+fi
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama405b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8}
+export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1152}
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4/"
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"405B"}
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -5,6 +5,7 @@ export DEV=${DEV:-AMD}
 export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
 
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
@@ -14,7 +15,7 @@ export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-0}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"
@@ -22,7 +23,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh

@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+export FAKEDATA=1 BENCHMARK=10
+if [ -z "$FULL_LAYERS" ]; then
+  export LLAMA_LAYERS=2
+fi
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -5,6 +5,7 @@ export DEV=${DEV:-AMD}
 export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
 
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
@@ -14,7 +15,7 @@ export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-0}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"
@@ -22,7 +23,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run_mp.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-32}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh

@@ -3,4 +3,4 @@ export BENCHMARK=5
 export EVAL_BS=0
 export VIZ=${VIZ:--1}
 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
-PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20
+extra/viz/cli.py --profile --device "AMD" --top 20

examples/openpilot/compile3.py

@@ -31,7 +31,7 @@ def compile(onnx_file):
   for i in range(3):
     GlobalCounters.reset()
     print(f"run {i}")
-    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1)):
+    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1), OPENPILOT_HACKS=1):
       ret = run_onnx_jit(**inputs).numpy()
     # copy i == 1 so use of JITBEAM is okay
     if i == 1: test_val = np.copy(ret)

examples/openpilot/load_pickle.py

@@ -0,0 +1,16 @@
+import sys, pickle
+from extra.bench_log import WallTimeEvent, BenchEvent
+from tinygrad.helpers import getenv
+
+PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
+
+load_times = []
+
+for _ in range(10):
+  with WallTimeEvent(BenchEvent.STEP) as wte: pickle.load(open(PKL, 'rb'))
+  load_times.append(wte.time)
+  print(f"pickle load: {wte.time:6.2f} s")
+
+if (assert_time:=getenv("ASSERT_MIN_LOAD_TIME")):
+  min_time = min(load_times)
+  assert min_time < assert_time, f"Speed regression, expected min load time of < {assert_time} s but took: {min_time} s"

extra/amdpci/am_smi.py

@@ -65,7 +65,7 @@ def get_bar0_size(pcibus):
 class AMSMI(AMDev):
   def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
     self.pcibus = pcibus
-    self.vram, self.doorbell64, self.mmio, self.dma_regions = vram_bar, doorbell_bar, mmio_bar, None
+    self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
     self.pci_state = self.read_pci_state()
     if self.pci_state == "D0": self._init_from_d0()
 

extra/amdpci/hive_reset.py

@@ -2,7 +2,7 @@
 
 import os
 from tinygrad.helpers import Context
-from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
+from tinygrad.runtime.support.system import System, PCIDevice
 from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev
 
@@ -12,7 +12,7 @@ if __name__ == "__main__":
     drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
     if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":
       raise RuntimeError(f"amdgpu is bound to {gpu}. Stopping...")
-  pcidevs = [PCIDevice("AM", gpu, bars=[0, 2, 5]) for gpu in gpus]
+  pcidevs = [PCIDevice("AM", gpu) for gpu in gpus]
   amdevs = []
   with Context(DEBUG=2):
     for pcidev in pcidevs:

extra/amdpci/proclogs.py

@@ -7,8 +7,8 @@ class GFXFake:
   def __init__(self): self.xccs = 8
 
 class AMDFake(AMDev):
-  def __init__(self, pci_dev, dma_regions=None):
-    self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
+  def __init__(self, pci_dev):
+    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
     self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
     self._run_discovery()
     self._build_regs()

extra/bench_log.py

@@ -34,7 +34,8 @@ class WallTimeEvent:
     self.start = time.monotonic()
     return self
   def __exit__(self, *_):
-    _events[self.event]["wall"].append(time.monotonic() - self.start)
+    self.time = time.monotonic() - self.start
+    _events[self.event]["wall"].append(self.time)
     return False
 
 class KernelTimeEvent:

extra/gemm/amd_copy_matmul.py

@@ -0,0 +1,84 @@
+from tinygrad import UOp, getenv
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+from tinygrad.dtype import AddrSpace, dtypes
+
+N = getenv("N", 4096)
+M = getenv("M", N)
+K = getenv("K", N)
+
+WARP_SIZE = 32
+BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 8
+TM, TN = 4, 4
+LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
+assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0
+
+is_kernel5 = getenv("K5", 0)
+THREADS_PER_BLOCK = 128 if is_kernel5 else 256
+WAVES_PER_BLOCK_N = 1 if is_kernel5 else 2
+WAVES_PER_BLOCK_M = THREADS_PER_BLOCK // WARP_SIZE // WAVES_PER_BLOCK_N
+REG_TILES_PER_WAVE_N = BLOCK_N // (WAVES_PER_BLOCK_N * LANES_PER_WAVE_N * TN)
+REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
+
+assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
+assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
+
+consts = {"wpb_m":WAVES_PER_BLOCK_M, "lpw_m":LANES_PER_WAVE_M, "rt_m":REG_TILES_PER_WAVE_M, "t_m": TM,
+          "wpb_n":WAVES_PER_BLOCK_N, "lpw_n":LANES_PER_WAVE_N, "rt_n":REG_TILES_PER_WAVE_N, "t_n": TN}
+
+# 128x128 out, kx128, kx128 in
+def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
+  tid = UOp.range(THREADS_PER_BLOCK, 2, AxisType.LOCAL)
+
+  #tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
+  warp, lane = tid // WARP_SIZE, tid % WARP_SIZE
+  wave_n, wave_m = warp % WAVES_PER_BLOCK_N, warp // WAVES_PER_BLOCK_N
+  lane_n, lane_m = lane % LANES_PER_WAVE_N, lane // LANES_PER_WAVE_N
+
+  # define locals
+  A_local = UOp.placeholder((BLOCK_K, BLOCK_M), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
+  B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+
+  # open the main reduction range and copy in GLOBAL -> LOCAL
+  a = a.reshape(K // BLOCK_K, BLOCK_K, BLOCK_M)
+  b = b.reshape(K // BLOCK_K, BLOCK_K, BLOCK_N)
+  k_tile_range = UOp.range(K // BLOCK_K, 3, AxisType.REDUCE)
+  A_store = A_local.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile_range].reshape(-1, THREADS_PER_BLOCK)[:, tid])
+  B_store = B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile_range].reshape(-1, THREADS_PER_BLOCK)[:, tid])
+  barrier = UOp.barrier(A_store, B_store)
+  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+
+  # define accumulator (128x128), but broadcast across tid
+  c_regs = UOp.placeholder((REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  c_regs = c_regs.after(c_regs.store(UOp.const(dtypes.float, 0).reshape((1,)*len(c_regs.shape)).expand(c_regs.shape)))
+
+  # define registers (NOTE: the thread count is the device count for this multi, it's sharded across the THREADS_PER_BLOCK)
+  A_col = UOp.placeholder((REG_TILES_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  B_row = UOp.placeholder((REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+
+  # LOCAL -> REGS
+  k = UOp.range(BLOCK_K, 4, AxisType.REDUCE)
+  A_col = A_col.after(A_col.store(A_local[k].reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM)[wave_m, :, lane_m, :]))
+  B_row = B_row.after(B_row.store(B_local[k].reshape(WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)[wave_n, :, lane_n, :]))
+
+  # do FMA
+  A_col = A_col.reshape(REG_TILES_PER_WAVE_M*TM, 1).expand(REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN)
+  B_row = B_row.reshape(1, REG_TILES_PER_WAVE_N*TN).expand(REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN)
+  c_regs = c_regs.after(c_regs.store(c_regs.after(k) + (A_col * B_row)).end(k).barrier().end(k_tile_range))
+
+  # store back to c
+  c_store = c.rearrange("(wpb_m rt_m lpw_m t_m) (wpb_n rt_n lpw_n t_n) -> (wpb_m wpb_n lpw_m lpw_n) (rt_m t_m) (rt_n t_n)", **consts)
+  return c_store[tid].store(c_regs).end(tid)
+
+def amd_copy_matmul(c:UOp, a:UOp, b:UOp) -> UOp:
+  block_id_n = UOp.range(N // BLOCK_N, 0, AxisType.GLOBAL)
+  block_id_m = UOp.range(M // BLOCK_M, 1, AxisType.GLOBAL)
+
+  # index the output with the globals
+  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]
+  a = a.T.reshape(K, M // BLOCK_M, BLOCK_M)[:, block_id_m, :]
+  b = b.reshape(K, N // BLOCK_N, BLOCK_N)[:, block_id_n, :]
+  return block_128x128_gemm(c, a, b).end(block_id_n, block_id_m).sink(arg=KernelInfo(opts_to_apply=()))
+
+if __name__ == "__main__":
+  from amd_uop_matmul import eval_custom_matmul
+  eval_custom_matmul(amd_copy_matmul)

extra/gemm/amd_uop_matmul.py

@@ -1,98 +1,74 @@
-import numpy as np
-from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
+from tinygrad import Tensor, Context, GlobalCounters, dtypes
 from tinygrad.uop.ops import UOp, KernelInfo, sint, AxisType
-from tinygrad.engine.realize import ExecItem, get_runner
 from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import getenv
+from tinygrad.helpers import DEBUG, getenv
 
 N = getenv("N", 4096)
-M = K = N
-run_count = getenv("CNT", 5)
+M = getenv("M", N)
+K = getenv("K", N)
+NUM_RUNS = getenv("CNT", 5)
 
 # ---------------------------
 # launch/config constants
 # ---------------------------
 
 WARP_SIZE = 32
-
-# Threadblock tile sizes (block-level tile of C that a block computes)
-BLOCK_N = 128   # columns of C (N-dim) per block
-BLOCK_M = 128   # rows of C (M-dim) per block
-BLOCK_K = 8     # K-slice per block iteration
-
-# Register tile sizes (per-thread accumulator tile of C)
-TN = 4     # columns per thread
-TM = 4     # rows per thread
+BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 8
+TM, TN = 4, 4
+LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
+assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0
 
 is_kernel5 = getenv("K5", 0)
 THREADS_PER_BLOCK = 128 if is_kernel5 else 256
-assert THREADS_PER_BLOCK % BLOCK_N == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_N"
-assert THREADS_PER_BLOCK % BLOCK_K == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_K"
-assert (BLOCK_N * BLOCK_K) % THREADS_PER_BLOCK == 0
-assert (BLOCK_M * BLOCK_K) % THREADS_PER_BLOCK == 0
+WAVES_PER_BLOCK_N = 1 if is_kernel5 else 2
+WAVES_PER_BLOCK_M = THREADS_PER_BLOCK // WARP_SIZE // WAVES_PER_BLOCK_N
+REG_TILES_PER_WAVE_N = BLOCK_N // (WAVES_PER_BLOCK_N * LANES_PER_WAVE_N * TN)
+REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
 
-WARPS_PER_BLOCK = THREADS_PER_BLOCK // WARP_SIZE
-WAVE_TILE_N = 128 if is_kernel5 else 64
-WAVE_TILE_M = BLOCK_N * BLOCK_M // WARPS_PER_BLOCK // WAVE_TILE_N
-assert BLOCK_N % WAVE_TILE_N == 0, "BN must be a multiple of WN"
-assert BLOCK_M % WAVE_TILE_M == 0, "BM must be a multiple of WM"
-WAVES_IN_BLOCK_X = BLOCK_N // WAVE_TILE_N
-WAVES_IN_BLOCK_Y = BLOCK_M // WAVE_TILE_M
-assert WAVES_IN_BLOCK_X * WAVES_IN_BLOCK_Y == WARPS_PER_BLOCK, "wave grid must match warps/block"
-
-LANES_PER_WAVE_X = 8
-LANES_PER_WAVE_Y = 4
-ITERS_PER_WAVE_N = WAVE_TILE_N // (LANES_PER_WAVE_X * TN)
-ITERS_PER_WAVE_M = WAVE_TILE_M // (LANES_PER_WAVE_Y * TM)
-assert WAVE_TILE_N % (LANES_PER_WAVE_X * TN) == 0, "WAVE_TILE_N must be divisible by LANES_PER_WAVE_X*TN"
-assert WAVE_TILE_M % (LANES_PER_WAVE_Y * TM) == 0, "WAVE_TILE_M must be divisible by LANES_PER_WAVE_Y*TM"
+assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
+assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
 
 def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
-def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=False):
+def copy(dest:UOp, src:UOp, rng:int, upcast=False):
   assert dest.shape == src.shape
   rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
-  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
-  return dest.after(copy) if set else copy
+  return dest[*rngs].store(src[*rngs]).end(*rngs)
 
-def hand_spec_kernel3():
+def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
   # ---------------------------
-  # block indices & placeholders
+  # block indices
   # ---------------------------
-  blockIdx_x = UOp.special(N // BLOCK_N, "gidx0")
-  blockIdx_y = UOp.special(N // BLOCK_M, "gidx1")
-
-  a = UOp.placeholder((N, N), dtypes.float, slot=1)
-  b = UOp.placeholder((N, N), dtypes.float, slot=2)
-  c = UOp.placeholder((N, N), dtypes.float, slot=0)
+  block_id_n = UOp.special(N // BLOCK_N, "gidx0")
+  block_id_m = UOp.special(M // BLOCK_M, "gidx1")
 
   # index the output with the globals
-  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[blockIdx_y, :, blockIdx_x, :]
+  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]
 
   # open the main reduction range
-  k_tile_range = UOp.range(N // BLOCK_K, 0, AxisType.REDUCE)
-  a = a.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_K, BLOCK_K)[blockIdx_y, :, k_tile_range, :]
-  b = b.reshape(N // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, blockIdx_x, :]
+  k_tile_range = UOp.range(K // BLOCK_K, 0, AxisType.REDUCE)
+  a = a.reshape(M // BLOCK_M, BLOCK_M, K // BLOCK_K, BLOCK_K)[block_id_m, :, k_tile_range, :]
+  b = b.reshape(K // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, block_id_n, :]
 
   # globals are no longer used, they are already in the indexes
-  del blockIdx_y, blockIdx_x
+  del block_id_m, block_id_n
 
   # ---------------------------
-  # GLOBAL -> LOCAL (As, Bs)
+  # GLOBAL -> LOCAL (A_local, B_local)
   # ---------------------------
   tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
 
   # A: read BM x BK tiles (permute on store into locals)
-  BM_As_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
-  As = UOp.placeholder((BLOCK_K, BM_As_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
-  As_store = copy(As.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
+  BM_A_local_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
+  A_local = UOp.placeholder((BLOCK_K, BM_A_local_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
+  A_local_store = copy(A_local.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
 
   # B: read BK x BN tiles
-  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
-  Bs_store = copy(Bs.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
+  B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+  B_local_store = copy(B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
 
   # TODO: can we automate barrier?
-  barrier = UOp.barrier(As_store, Bs_store)
-  As, Bs = As.after(barrier), Bs.after(barrier)
+  barrier = UOp.barrier(A_local_store, B_local_store)
+  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
 
   # open inner k range
   k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
@@ -100,31 +76,30 @@ def hand_spec_kernel3():
   # ---------------------------
   # LOCAL -> REG (per-wave tiles)
   # ---------------------------
-  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
-  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
-  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
+  warp, lane = tid // WARP_SIZE, tid % WARP_SIZE
+  waveIdx, waveIdy = warp % WAVES_PER_BLOCK_N, warp // WAVES_PER_BLOCK_N
+  laneIdx, laneIdy = lane % LANES_PER_WAVE_N, lane // LANES_PER_WAVE_N
+  assert waveIdy.vmax+1 == WAVES_PER_BLOCK_M and laneIdy.vmax+1 == LANES_PER_WAVE_M
 
-  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
-  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
-  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y
+  A_col = UOp.placeholder((REG_TILES_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  A_local_slice = A_local[k, :].reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM)[waveIdy, :, laneIdy, :]
+  A_col = A_col.after(copy(A_col, A_local_slice, 300, upcast=True))
 
-  A_col = UOp.placeholder((ITERS_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  A_col = copy(A_col, As[k, :].reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)[waveIdy, :, laneIdy, :], 300, set=True, upcast=True)
-
-  B_row = UOp.placeholder((ITERS_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
-  B_row = copy(B_row, Bs[k, :].reshape(WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)[waveIdx, :, laneIdx, :], 400, set=True, upcast=True)
+  B_row = UOp.placeholder((REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+  B_local_slice = B_local[k, :].reshape(WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)[waveIdx, :, laneIdx, :]
+  B_row = B_row.after(copy(B_row, B_local_slice, 400, upcast=True))
 
   # ---------------------------
   # FMA: c_regs += A_col * B_row
   # ---------------------------
-  c_regs = UOp.placeholder((ITERS_PER_WAVE_M, TM, ITERS_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  c_regs = UOp.placeholder((REG_TILES_PER_WAVE_M, TM, REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
   i = UOp.range(c_regs.size, 16)
   c_regs = c_regs.after(c_regs.flatten()[i].store(0.0).end(i))
 
   # TODO: why don't these work as upcast?
   # why if the ranges merge is it slow?!? (if you change the order on end, they will merge. big slowdown on METAL)
-  iterWaveM, yt, iterWaveN, xt = rngs = rngs_for_shape(c_regs.shape, 500)
-  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iterWaveM, yt] * B_row[iterWaveN, xt]).end(iterWaveM, iterWaveN, yt, xt)
+  iter_m, t_m, iter_n, t_n = rngs = rngs_for_shape(c_regs.shape, 500)
+  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iter_m, t_m] * B_row[iter_n, t_n]).end(iter_m, iter_n, t_m, t_n)
 
   # Close k, sync, and close K tiles
   sink = sink.end(k).barrier().end(k_tile_range)
@@ -132,38 +107,37 @@ def hand_spec_kernel3():
   # ---------------------------
   # REG -> GLOBAL (epilogue)
   # ---------------------------
-  c = c.reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
-                WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
+  c = c.reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM,
+                WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)
   c = c[waveIdy, :, laneIdy, :,
         waveIdx, :, laneIdx, :]
   sink = copy(c, c_regs.after(sink), rng=600)
 
   return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
 
-def test_matmul(sink:UOp, dtype=dtypes.float32, N=N):
-  rng = np.random.default_rng()
-  a = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
-  b = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
-  hc = Tensor.empty(N, N, dtype=dtype)
-  Tensor.realize(a, b, hc)
-
-  ei = ExecItem(sink, [t.uop.buffer for t in [hc, a, b]], prg=get_runner(Device.DEFAULT, sink))
+def eval_custom_matmul(fxn):
+  a = Tensor.randn(M, K, dtype=dtypes.float)
+  b = Tensor.randn(K, N, dtype=dtypes.float)
+  c = Tensor.empty(M, N, dtype=dtypes.float)
+  with Context(DEBUG=0): Tensor.realize(a, b)
 
   ets = []
-  with Context(DEBUG=2):
-    for _ in range(run_count):
-      ets.append(ei.run(wait=True))
-  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
+  with Context(DEBUG=max(2, DEBUG.value)):
+    for _ in range(NUM_RUNS):
+      GlobalCounters.reset()
+      tst = Tensor.custom_kernel(c, a, b, fxn=fxn)[0].realize()
+      ets.append(GlobalCounters.time_sum_s)
+  print(f"REAL TFLOPS {M * N * K * 2 / min(ets) * 1e-12:.2f}")
 
   if getenv("VERIFY", 1):
     GlobalCounters.reset()
     with Context(DEBUG=2):
       tc = (a @ b).realize()
     with Context(DEBUG=0):
-      err = (hc - tc).square().mean().item()
+      err = (tc - tst).square().mean().item()
     print(f"mean squared error {err}")
     if err > 1e-06:
       raise RuntimeError("matmul is wrong!")
 
 if __name__ == "__main__":
-  test_matmul(hand_spec_kernel3(), N=N)
+  eval_custom_matmul(hand_spec_kernel3)

extra/gemm/asm/cdna/gemm.py

@@ -1,101 +0,0 @@
-import atexit, functools
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from tinygrad.renderer import Estimates
-from tinygrad.helpers import getenv, all_same, dedup
-from extra.gemm.asm.cdna.asm import build_kernel, TILE_M, TILE_N, TILE_K, NUM_WG
-
-# ** CDNA4 assembly gemm
-
-WORKGROUP_SIZE = 256
-
-@functools.cache
-def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str, arch:str, wg:int) -> UOp:
-  batch, M, K = A.shape
-  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
-  assert K == K2
-  lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
-  gidx = UOp.special(wg, "gidx0")
-  insts = build_kernel(batch, M, N, K, A.dtype.base)
-  lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=133_120, addrspace=AddrSpace.LOCAL), (), 'lds')
-  sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
-                  arg=KernelInfo(name=f"gemm_{batch}_{M}_{N}_{K}", estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname),
-                                UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
-
-counters = {"used":0, "todos":[]}
-def todo(msg:str) -> bool: counters["todos"].append(msg); return False
-atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used'))
-
-def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
-  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
-  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
-  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
-  N = b.shape[1]
-  # only sharding on the batch or K is tested, others might work too
-  if isinstance(a.device, tuple):
-    if a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
-    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
-    else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
-    dname = a.device[0]
-  else: dname = a.device
-  arch = getattr(Device[dname].renderer, "arch", "")
-  if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
-  if (M % TILE_M != 0 or N % TILE_N != 0 or K % TILE_K != 0) and arch == "gfx950":
-    return todo(f"GEMM shape ({M},{N},{K}) not a multiple of ({TILE_M},{TILE_N},{TILE_K})")
-  return True
-
-# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
-# note: this can be removed after we have GEMM on mixins
-
-def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
-  M, K = A.shape[0]*A.shape[1], A.shape[2]
-  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
-  assert K == K2
-  m = UOp.range(M, 1, AxisType.LOOP)
-  n = UOp.range(N, 2, AxisType.LOOP)
-  k = UOp.range(K, 0, AxisType.REDUCE)
-  mul = (A.index((m*UOp.const(dtypes.index, K)+k))*B.index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
-  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
-  store = C.index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
-  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
-
-# ** backward gemm, might use the asm gemm
-
-def custom_gemm_bw(gradient:UOp, kernel:UOp):
-  out, a, b = kernel.src[1:]
-  assert all_same([gradient.device, a.device, b.device, out.device])
-  a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
-  # TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
-  g_t = g_t[:a.shape[0]]
-  grad_a = (g_t @ b_t.T).uop
-  grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
-  return (None, grad_a, grad_b)
-
-# ** main gemm function
-
-def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
-  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
-  counters["used"] += 1
-  squeeze = a.ndim == 2
-  if squeeze: a = a.unsqueeze(0)
-
-  batch, M, K = a.shape
-  N = b.shape[1]
-  is_multi = isinstance(a.device, tuple)
-  if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
-
-  if is_multi:
-    out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
-  else:
-    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
-
-  renderer = Device[a.device[0] if is_multi else a.device].renderer
-  dname, arch = renderer.device, getattr(renderer, "arch", "")
-  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=NUM_WG, arch=arch), grad_fxn=custom_gemm_bw)[0]
-  else:
-    out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
-  if k_sharded: out = out.sum(0)
-  return out.squeeze(0) if squeeze else out

extra/gemm/asm/rdna3/template.s

@@ -1,76 +0,0 @@
-.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 30336
-  .amdhsa_private_segment_fixed_size 0
-  .amdhsa_kernarg_size 32
-  # register usage (RSRC1)
-  .amdhsa_next_free_vgpr 256
-  .amdhsa_next_free_sgpr 100
-  # 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: kernarg ptr in s[0:1]
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_user_sgpr_count 2
-  # gfx10+ / gfx11 specifics (RSRC1[29..31])
-  .amdhsa_wavefront_size32 1
-  .amdhsa_workgroup_processor_mode 1
-  .amdhsa_memory_ordered 1
-  .amdhsa_forward_progress 1
-  # misc for gfx11
-  .amdhsa_dx10_clamp 1
-  .amdhsa_ieee_mode 1
-  .amdhsa_uses_dynamic_stack 0
-.end_amdhsa_kernel
-
-.amdgpu_metadata
----
-amdhsa.kernels:
-  - .args:
-      - .address_space:  generic
-        .name:           C
-        .offset:         0
-        .size:           8
-        .value_kind:     global_buffer
-        .value_type:     f16
-      - .address_space:  generic
-        .name:           A
-        .offset:         8
-        .size:           8
-        .value_kind:     global_buffer
-        .value_type:     f16
-      - .address_space:  generic
-        .name:           B
-        .offset:         16
-        .size:           8
-        .value_kind:     global_buffer
-        .value_type:     f16
-    .group_segment_fixed_size: 30336
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 32
-    .max_flat_workgroup_size: 128
-    .name:           gemm
-    .private_segment_fixed_size: 0
-    .sgpr_count:     70
-    .sgpr_spill_count: 0
-    .symbol:         gemm.kd
-    .vgpr_count:     256
-    .vgpr_spill_count: 0
-    .wavefront_size: 32
-amdhsa.version:
-  - 1
-  - 1
-...
-.end_amdgpu_metadata

extra/gemm/asm/rdna3/test.py

@@ -1,30 +0,0 @@
-import math, pathlib
-
-from tinygrad import Device, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-
-from extra.gemm.amd_uop_matmul import test_matmul
-
-N = 4096
-TN = 96
-THREADS_PER_WG = 128
-NUM_WG = math.ceil(N / TN) * math.ceil(N / TN)
-
-dname:str = Device.DEFAULT
-template:str = (pathlib.Path(__file__).parent/"template.s").read_text()
-
-def asm_kernel() -> UOp:
-  lidx = UOp.special(THREADS_PER_WG, "lidx0")
-  gidx = UOp.special(NUM_WG, "gidx0")
-
-  a = UOp.placeholder((N*N,), dtypes.half, slot=1)
-  b = UOp.placeholder((N*N,), dtypes.half, slot=2)
-  c = UOp.placeholder((N*N,), dtypes.half, slot=0)
-
-  src = template.replace("INSTRUCTIONS", (pathlib.Path(__file__).parent/"gemm.s").read_text())
-
-  sink = UOp.sink(a, b, c, lidx, gidx, arg=KernelInfo(name="gemm"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
-
-if __name__ == "__main__":
-  test_matmul(asm_kernel(), dtype=dtypes.half, N=N)

extra/gemm/asm/unpack_kd.py

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

extra/hevc/decode.py

@@ -10,9 +10,9 @@ HEVC_ROUNDUP = getenv("DATA_ROUNDUP", 32)
 @functools.cache
 def _hevc_jitted_decoder(out_image_size:tuple[int, int], max_hist:int, inplace:bool):
   def hevc_decode_frame(pos:Variable, hevc_tensor:Tensor, offset:Variable, sz:Variable, opaque:Tensor, i:Variable, *hist:Tensor, outbuf:Tensor|None=None):
-    x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist)
+    x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist).realize()
     if outbuf is not None: outbuf.assign(x).realize()
-    return x.realize()
+    return x
   return TinyJit(hevc_decode_frame)
 
 def hevc_decode(hevc_tensor:Tensor, opaque:Tensor, frame_info:list, luma_h:int, luma_w:int,
@@ -74,10 +74,14 @@ if __name__ == "__main__":
   Device.default.synchronize()
 
   # decode all frames using the iterator
-  with Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps")):
+  tm = Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps"))
+  with tm:
     images = list(hevc_decode(hevc_tensor, opaque_nv, frame_info, luma_h, luma_w, history=hist, preallocated_outputs=out_images))
     Device.default.synchronize()
 
+  fps = len(frame_info)/(tm.et/1e9)
+  assert fps >= getenv("ASSERT_FPS", 0), f"HEVC decode too slow: {fps:.2f} fps"
+
   # validation
   if getenv("VALIDATE", 0):
     import pickle

extra/models/llama.py

@@ -56,7 +56,10 @@ class Attention:
   def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]=None) -> Tensor:
     if getenv("WQKV"):
       xqkv = self.wqkv(x)
-      xq, xk, xv = xqkv.split([self.n_heads * self.head_dim, self.n_kv_heads * self.head_dim, self.n_kv_heads * self.head_dim], dim=2)
+      xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
+      xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
     else:
       xq, xk, xv = self.wq(x), self.wk(x.contiguous_backward()), self.wv(x)
 

extra/sqtt/README.md

@@ -2,30 +2,13 @@
 
 ## Getting SQ Thread Trace
 
-SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.
+`VIZ=2` to enable SQTT profiling.
+
+`SQTT_ITRACE_SE_MASK=X` to select shader engines for instruction tracing, -1 = all, 0 = disabled, >0 = SE bitmask, default 0b11.
 
 `SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.
 
-`SQTT_ITRACE_SE_MASK=X` to select for which shader engines instruction tracing will be enabled, -1 is all, 0 is none (instruction tracing disabled), >0 is
-bitfield/mask for SEs to enable instruction tracing on. Masking shader engines will give smaller file sizes at a cost of less hits and kernels that
-don't have any wavefront on first simd of shader engine with instruction tracing enabled will not have instruction timings.
-The default is 2 (second shader engine only), only one for file size reasons, second instead of first because dispatch starts from it so there is
-greater chance that kernels with small global size will have instruction tracing data.
- 
-Note that instruction tracing might not be available for kernels with small global dims, this is not a bug, but it can be improved with various hacks
-to the point where it can reliably trace a kernel consisting of a single wavefront (am only, not quite reliable under amdgpu due to waves sometimes
-being dispatched starting from different simds). More info in comments in ops_amd.py
+## Viewing the traces
 
-## Converting pickled profile with SQTT data into RGP file
-
-```bash
-extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
-```
-
-Then load gpu0.rgp into Radeon GPU Profiler. It works just fine both in wine (macos, native version available for linux) and via ssh X forwarding
-
-If multiple gpus are used you can select which one to export with `-d` like this:
-
-```bash
-extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -d 'AMD:5' -o /tmp/gpu5.rgp
-```
+- Web UI: `tinygrad/viz/serve.py`
+- Command line: `python -m tinygrad.renderer.amd.sqtt`

extra/sqtt/active_sqtt_parse.py

@@ -1,152 +0,0 @@
-import os
-os.environ["PYTHONPATH"] = "."
-os.environ["SQTT"] = "1"
-if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
-os.environ["PROFILE"] = "1"
-os.environ["AMD_LLVM"] = "0"
-
-from dataclasses import replace
-import atexit, contextlib
-from tinygrad import Tensor
-from tinygrad.helpers import system, OSX
-from tinygrad.runtime.ops_amd import AMDProgram
-from extra.sqtt.roc import decode, WaveExec, ProfileSQTTEvent
-from tinygrad.device import Device
-
-from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
-
-dev = Device["AMD"]
-
-@contextlib.contextmanager
-def save_sqtt():
-  # clear the old traces
-  dev.profile_events.clear()
-  sqtt:dict[str, list[WaveExec]] = {}
-  yield sqtt
-  events = dev.profile_events
-
-  #rctx = decode(events)
-  #assert len(rctx.inst_execs) > 0, "empty sqtt output"
-  #sqtt.update(rctx.inst_execs)
-
-  for e in events:
-    if isinstance(e, ProfileSQTTEvent):
-      print(replace(e, blob=b''))
-      if e.se == 0:
-        parse_sqtt_print_packets(e.blob)
-
-template = """.text
-.globl matmul
-.p2align 8
-.type matmul,@function
-matmul:
-  INSTRUCTION
-
-.rodata
-.p2align 6
-.amdhsa_kernel matmul
-  .amdhsa_kernarg_size 8
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
-  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
-  .amdhsa_wavefront_size32 1
-.end_amdhsa_kernel
-
-.amdgpu_metadata
----
-amdhsa.version:
-  - 1
-  - 0
-amdhsa.kernels:
-  - .name: matmul
-    .symbol: matmul.kd
-    .group_segment_fixed_size: 0
-    .private_segment_fixed_size: 0
-    .wavefront_size: 32
-    .sgpr_count: 8
-    .vgpr_count: 8
-    .max_flat_workgroup_size: 1024
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 8
-    .args:
-      - .address_space:  global
-        .name:           a
-        .offset:         0
-        .size:           8
-        .type_name:      'float*'
-        .value_kind:     global_buffer
-...
-.end_amdgpu_metadata
-"""
-
-def run_asm(src, num_workgroups=1, num_waves=1):
-  WAVE_SIZE = 32
-  t = Tensor.empty(0x1000).realize()
-  buf = t.uop.buffer.ensure_allocated()
-  lib = dev.compiler.compile(template.replace("INSTRUCTION", '\n'.join(src)))
-  dev.compiler.disassemble(lib)
-  fxn = AMDProgram(dev, "matmul", lib)
-  fxn(buf._buf, global_size=(num_workgroups,1,1), local_size=(WAVE_SIZE*num_waves,1,1), wait=True)
-
-if __name__ == "__main__":
-  with save_sqtt() as sqtt:
-    run_asm([
-      "s_nop 100",
-      "s_nop 100",
-      "s_load_b64 s[0:1], s[0:1], null",
-      "s_waitcnt lgkmcnt(0)",
-      "s_nop 100",
-      "s_nop 100",
-      "s_add_i32 s2, s2, 10",
-      "s_add_i32 s2, s2, 10",
-      "s_nop 100",
-      "s_nop 100",
-      "v_mov_b32_e32 v0, 0",
-      "v_mov_b32_e32 v0, 0",
-      "s_nop 100",
-      "s_nop 100",
-      "v_dual_fmac_f32 v2, v48, v24 :: v_dual_fmac_f32 v9, v37, v51",
-      "v_dual_fmac_f32 v2, v48, v24 :: v_dual_fmac_f32 v9, v37, v51",
-      "s_nop 100",
-      "s_nop 100",
-      "global_load_b128 v[2:5], v0, s[0:1]",
-      "global_load_b128 v[2:5], v0, s[0:1]",
-      "s_nop 100",
-      "s_nop 100",
-      "s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)",
-      "s_endpgm",
-    ], num_workgroups=1, num_waves=1)
-  exit(0)
-
-  with save_sqtt() as sqtt:
-    #(Tensor.empty(16,16) @ Tensor.empty(16,16)).elu().realize()
-    #Tensor.empty(1, 64).sum(axis=1).realize()
-    Tensor.empty(1).log2().realize()
-  exit(0)
-
-  with save_sqtt() as sqtt:
-    # what's in v0?
-    run_asm([
-      "v_mov_b32_e32 v0, 0",
-      "v_mov_b32_e32 v1, 0",
-      "s_clause 0x1",
-      "s_load_b64 s[0:1], s[0:1], null",
-      "s_waitcnt lgkmcnt(0)",
-    ]+[
-      "global_load_b32 v1, v0, s[0:1]",
-    ]*10+[
-      "global_load_b32 v10, v1, s[0:1]",
-      "s_waitcnt vmcnt(0)",
-
-      #"v_rcp_f32 v1, v0"
-      #"v_add_f32_e32 v1 v0 v0",
-      #"v_add_f32_e32 v5 v4 v4",
-      #"v_add_f32_e32 v7 v6 v6",
-      #"v_add_f32_e32 v1 v0 v0",
-      #"v_add_f32_e32 v2 v1 v1",
-      #"s_nop 1"
-    ]*5+[
-      "v_add_f32_e32 v3 v2 v2",
-    ]*5+[
-      "v_mul_f32_e32 v3 v2 v2",
-    ]*7)

extra/sqtt/attempt_sqtt_parse.py

@@ -1,548 +0,0 @@
-import pickle, sys
-from tinygrad.helpers import getenv, Timing, colored
-from extra.sqtt.roc import decode, ProfileSQTTEvent
-
-# do these enums match fields in the packets?
-#from tinygrad.runtime.support.amd import import_soc
-#soc = import_soc([11])
-#perf_sel = {getattr(soc, k):k for k in dir(soc) if k.startswith("SQ_PERF_")}
-
-# Instruction packets (one per ISA op)
-# NOTE: these are bad guesses and may be wrong! feel free to update if you know better
-# some names were taken from SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT
-
-# we see 18 opcodes
-# opcodes(18):  1  2  3  4  5  6  8  9  F 10 11 12 14 15 16 17 18 19
-# if you exclude everything, you are left with 6
-# opcodes( 6): 10 11 14 15 16 17
-# sometimes we see a lot of B, but not repeatable
-
-# not seen
-# 7 A C
-
-# NOTE: INST runs before EXEC
-
-OPCODE_COLORS = {
-  # dispatches are BLACK
-  0x1: "BLACK",
-  0x18: "BLACK",
-
-  # execs are yellow
-  0x2: "yellow",
-  0x3: "yellow",
-  0x4: "YELLOW",
-  0x5: "YELLOW",
-
-  # waves are blue
-  0x8: "blue",
-  0x9: "blue",
-  0x6: "cyan",
-  0xb: "cyan",
-}
-
-OPCODE_NAMES = {
-  # gated by SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT (but others must be enabled for it to show)
-  0x01: "VALUINST",
-  # gated by SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT
-  0x02: "VMEMEXEC",
-  # gated by SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT
-  0x03: "ALUEXEC",
-  # gated by SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT
-  0x04: "IMMEDIATE",
-  0x05: "IMMEDIATE_MASK",
-
-  # gated by SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT
-  0x06: "WAVERDY",
-  # gated by SQ_TT_TOKEN_EXCLUDE_WAVESTARTEND_SHIFT
-  0x08: "WAVEEND",
-  0x09: "WAVESTART",
-  # gated by SQ_TT_TOKEN_EXCLUDE_WAVEALLOC_SHIFT
-  0x0B: "WAVEALLOC",  # FFF00
-
-  # gated by NOT SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT
-  0x0D: "PERF",
-  # gated by SQ_TT_TOKEN_EXCLUDE_EVENT_SHIFT
-  0x12: "EVENT",
-  0x13: "EVENT_BIG",  # FFFFF800
-  # some gated by SQ_TT_TOKEN_EXCLUDE_REG_SHIFT, some always there. something is broken with the timing on this
-  0x14: "REG",
-  # gated by SQ_TT_TOKEN_EXCLUDE_INST_SHIFT
-  0x18: "INST",
-  # gated by SQ_TT_TOKEN_EXCLUDE_UTILCTR_SHIFT
-  0x19: "UTILCTR",
-
-  # this is the first (8 byte) packet in the bitstream
-  0x17: "LAYOUT_HEADER",       # layout/mode/group + selectors A/B (reversed)
-
-  # pure time (no extra bits)
-  0x0F: "TS_DELTA_SHORT",
-  0x10: "NOP",
-  0x11: "TS_WAVE_STATE",     # almost pure time, has a small flag
-
-  # not a good name, but seen and understood mostly
-  0x15: "SNAPSHOT",          # small delta + 50-ish bits of snapshot
-  0x16: "TS_DELTA_OR_MARK",  # 36-bit long delta or 36-bit marker
-
-  # packets we haven't seen / rarely see 0x0b
-  0x07: "TS_DELTA_S8_W3_7",         # shift=8,  width=3  (small delta)
-  0x0A: "TS_DELTA_S5_W2_A",         # shift=5,  width=2
-  0x0C: "TS_DELTA_S5_W3_B",         # shift=5,  width=3 (different consumer)
-}
-
-#  SALU    =  0x0 / s_mov_b32
-#  SMEM    =  0x1 / s_load_b*
-#  JUMP    =  0x3 / s_cbranch_scc0
-#  NEXT    =  0x4 / s_cbranch_execz
-#  MESSAGE =  0x9 / s_sendmsg
-#  VALU    =  0xb / v_(exp,log)_f32_e32
-#  VALU    =  0xd / v_lshlrev_b64
-#  VALU    =  0xe / v_mad_u64_u32
-#  VMEM    = 0x21 / global_load_b32
-#  VMEM    = 0x22 / global_load_b32
-#  VMEM    = 0x24 / global_store_b32
-#  VMEM    = 0x25 / global_store_b64
-#  VMEM    = 0x27 / global_store
-#  VMEM    = 0x28 / global_store_b64
-#  LDS     = 0x29 / ds_load_b128
-#  LDS     = 0x2b / ds_store_b32
-#  LDS     = 0x2e / ds_store_b128
-#  ????    = 0x5a / hidden global_load  instruction
-#  ????    = 0x5b / hidden global_load  instruction
-#  ????    = 0x5c / hidden global_store instruction
-#  VALU    = 0x73 / v_cmpx_eq_u32_e32 (not normal VALUINST)
-OPNAME = {
-  0x0: "SALU",
-  0x1: "SMEM",
-  0x3: "JUMP",
-  0x4: "NEXT",
-  0x9: "MESSAGE",
-  0xb: "VALU",
-  0xd: "VALU",
-  0xe: "VALU",
-  0x21: "VMEM_LOAD",
-  0x22: "VMEM_LOAD",
-  0x24: "VMEM_STORE",
-  0x25: "VMEM_STORE",
-  0x26: "VMEM_STORE",
-  0x27: "VMEM_STORE",
-  0x28: "VMEM_STORE",
-  0x29: "LDS_LOAD",
-  0x2b: "LDS_STORE",
-  0x2e: "LDS_STORE",
-  0x50: "__SIMD_LDS_LOAD",
-  0x51: "__SIMD_LDS_LOAD",
-  0x54: "__SIMD_LDS_STORE",
-  0x5a: "__SIMD_VMEM_LOAD",
-  0x5b: "__SIMD_VMEM_LOAD",
-  0x5c: "__SIMD_VMEM_STORE",
-  0x5d: "__SIMD_VMEM_STORE",
-  0x5e: "__SIMD_VMEM_STORE",
-  0x5f: "__SIMD_VMEM_STORE",
-  0x72: "SALU_OR",
-  0x73: "VALU_CMPX",
-}
-
-ALUSRC = {
-  1: "SALU",
-  2: "VALU",
-  3: "VALU_SALU",
-}
-
-MEMSRC = {
-  0: "LDS",
-  1: "__LDS",
-  2: "VMEM",
-  3: "__VMEM",
-}
-
-
-# these tables are from rocprof trace decoder
-# rocprof_trace_decoder_parse_data-0x11c6a0
-# parse_sqtt_180 = b *rocprof_trace_decoder_parse_data-0x11c6a0+0x110040
-
-# ---------- 1. local_138: 256-byte state->opcode table ----------
-
-STATE_TO_OPCODE: bytes = bytes([
-  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x12, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x13, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-])
-
-# opcode mask (the bits used to determine the opcode, worked out by looking at the repeats in STATE_TO_OPCODE)
-
-opcode_mask = {
-  0x10: 0b1111,
-
-  0x16: 0b1111111,
-  0x17: 0b1111111,
-  0x07: 0b1111111,
-  0x19: 0b1111111,
-  0x11: 0b1111111,
-  0x12: 0b11111111,
-  0x13: 0b11111111,
-  0x15: 0b1111111,
-
-  0x18: 0b111,
-  0x1: 0b111,
-
-  0x5: 0b11111,
-  0x6: 0b11111,
-  0xb: 0b11111,
-  0x8: 0b11111,
-  0xc: 0b11111,
-  0xd: 0b11111,
-
-  0xf: 0b1111,
-  0x14: 0b1111,
-
-  0x9: 0b11111,
-  0xa: 0b11111,
-
-  0x4: 0b1111,
-  0x3: 0b1111,
-  0x2: 0b1111,
-}
-
-# ---------- 2. DAT_0012e280: nibble budget per opcode&0x1F ----------
-
-NIBBLE_BUDGET = [
-  0x08, 0x0C, 0x08, 0x08, 0x0C, 0x18, 0x18, 0x40, 0x14, 0x20, 0x30, 0x14, 0x34, 0x1C, 0x30, 0x08,
-  0x04, 0x18, 0x18, 0x20, 0x40, 0x40, 0x30, 0x40, 0x14, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-]
-
-# ---------- 3. delta_map from your hash nodes ----------
-
-# opcode -> (shift, width)
-DELTA_MAP_DEFAULT = {
-  0x01: (3,  3),   # shift=3,  end=6
-  0x02: (4,  2),   # shift=4,  end=6
-  0x03: (4,  2),   # shift=4,  end=6
-  0x04: (4,  3),   # shift=4,  end=7
-  0x05: (5,  3),   # shift=5,  end=8
-  0x06: (5,  3),   # shift=5,  end=8
-  0x07: (8,  3),   # shift=8,  end=11
-  0x08: (5,  3),   # shift=5,  end=8
-  0x09: (5,  2),   # shift=5,  end=7
-  0x0A: (5,  2),   # shift=5,  end=7
-  0x0B: (5,  3),   # shift=5,  end=8
-  0x0C: (5,  3),   # shift=5,  end=8
-  0x0D: (5,  3),   # shift=5,  end=8
-  # NOTE: 0x0e can never be decoded, it's not in the STATE_TO_OPCODE table
-  #0x0E: (7,  2),   # shift=7,  end=9
-  0x0F: (4,  4),   # shift=4,  end=8
-  0x10: (0,  0),   # shift=0,  end=0  (no delta)
-  0x11: (7,  9),   # shift=7,  end=16
-  0x12: (8,  3),   # shift=8,  end=11
-  0x13: (8,  3),   # shift=8,  end=11
-  0x14: (4,  3),   # shift=4,  end=7
-  0x15: (7,  3),   # shift=7,  end=10
-  0x16: (12, 36),  # shift=12, end=48 (36-bit field, matches the 0x16 special-case)
-  0x17: (0,  0),   # shift=0,  end=0  (no delta)
-  0x18: (4,  3),   # shift=4,  end=7
-  0x19: (7,  2),   # shift=7,  end=9
-}
-
-# ---------- 4. One-line-per-packet parser ----------
-
-def reg_mask(opcode):
-  nb_bits = NIBBLE_BUDGET[opcode & 0x1F]
-  shift, width = DELTA_MAP_DEFAULT[opcode]
-  delta_mask = ((1 << width) - 1) << shift
-  assert delta_mask & opcode_mask[opcode] == 0, "masks shouldn't overlap"
-  return ((1 << nb_bits) - 1) & ~(delta_mask | opcode_mask[opcode])
-
-def decode_packet_fields(opcode: int, reg: int) -> str:
-  """
-  Decode packet payloads conservatively, using:
-    - NIBBLE_BUDGET[opcode & 0x1F] to mask reg down to true width.
-    - DELTA_MAP_DEFAULT[opcode] to expose the "primary" field (often delta).
-    - Per-opcode layouts derived from rocprof's decompiled consumers.
-  """
-  # --- 0. Restrict to real packet bits not used in delta ---------------------------------
-  pkt = reg & reg_mask(opcode)
-  fields: list[str] = []
-
-  match opcode:
-    case 0x01: # VALUINST
-      # 6 bit field
-      flag = (pkt >> 6) & 1
-      wave = pkt >> 7
-      fields.append(f"wave={wave:x}")
-      if flag: fields.append("flag")
-    case 0x02: # VMEMEXEC
-      # 2 bit field (pipe is a guess)
-      src = pkt>>6
-      fields.append(f"src={src} [{MEMSRC.get(src, '')}]")
-    case 0x03: # ALUEXEC
-      # 2 bit field
-      src = pkt>>6
-      fields.append(f"src={src} [{ALUSRC.get(src, '')}]")
-    case 0x04: # IMMEDIATE_4
-      # 5 bit field (actually 4)
-      wave = pkt >> 7
-      fields.append(f"wave={wave:x}")
-    case 0x05: # IMMEDIATE_5
-      # 16 bit field
-      # 1 bit per wave
-      fields.append(f"mask={pkt>>8:016b}")
-    case 0x6:
-      # wave ready FFFF00
-      # 16 bit field
-      # 1 bit per wave
-      fields.append(f"mask={pkt>>8:016b}")
-    case 0x0d:
-      # 20 bit field
-      fields.append(f"arg = {pkt>>8:X}")
-    case 0x12:
-      fields.append(f"event = {pkt>>11:X}")
-    case 0x15:
-      fields.append(f"snap = {pkt>>10:X}")
-    case 0x19:
-      # wave end
-      fields.append(f"ctr = {pkt>>9:X}")
-    case 0xf:
-      extracted_delta = (reg >> 4) & 0xF
-      fields.append(f"strange_delta=0x{extracted_delta:x}")
-    case 0x11:
-      # DELTA_MAP_DEFAULT: shift=7, width=9 -> small delta.
-      # FF0000 is the mask
-      coarse    = pkt >> 16
-      fields.append(f"coarse=0x{coarse:02x}")
-      # From decomp:
-      #  - when layout<3 and coarse&1, it sets a "has interesting wave" flag
-      #  - when coarse&8, it marks all live waves as "terminated"
-      if coarse & 0x01:
-        fields.append("flag_wave_interest=1")
-      if coarse & 0x08:
-        fields.append("flag_terminate_all=1")
-    case 0x8:
-      # wave end, this is 20 bits (FFF00)
-      flag7   = (pkt >> 8) & 1
-      simd    = (pkt >> 9) & 3
-      cu      = ((pkt >> 11) & 0x7) | (flag7 << 3)
-      wave    = (pkt >> 15) & 0x1f
-      fields.append(f"wave={wave:x}")
-      fields.append(f"simd={simd}")
-      fields.append(f"cu={cu}")
-    case 0x9:
-      # From case 9 (WAVESTART) in multiple consumers:
-      #   flag7  = (w >> 7) & 1        (low bit of uVar41)
-      #   cls2   = (w >> 8) & 3        (class / group)
-      #   slot4  = (w >> 10) & 0xf     (slot / group index)
-      #   idx_lo = (w >> 0xd) & 0x1f   (low index, layout<4 path)
-      #   idx_hi = (w >> 0xf) & 0x1f   (high index, layout>=4 path)
-      #   id7    = (w >> 0x19) & 0x7f  (7-bit id)
-      flag7   = (pkt >> 7) & 1
-      simd    = (pkt >> 8) & 3
-      cu      = ((pkt >> 10) & 0x7) | (flag7 << 3)
-      wave    = (pkt >> 13) & 0x1F
-      id7     = (pkt >> 17)
-      fields.append(f"wave={wave:x}")
-      fields.append(f"simd={simd}")
-      fields.append(f"cu={cu}")
-      fields.append(f"id7=0x{id7:x}")
-    case 0x18:
-      # FFF88 is the mask
-      # From case 0x18:
-      #   low3   = w & 7
-      #   grp3   = (w >> 3) or (w >> 4) & 7   (layout-dependent)
-      #   flags  = bits 6 (B6) and 7 (B7)
-      #   hi8    = (w >> 0xc) & 0xff   (layout 4 path)
-      #   hi7    = (w >> 0xd) & 0x7f   (other layouts)
-      #   idx5   = (w >> 7) or (w >> 8) & 0x1f, used as wave index
-      flag1 = (pkt >> 3) & 1
-      flag2 = (pkt >> 7) & 1
-      wave = (pkt >> 8) & 0x1F
-      op = (pkt >> 13)
-      fields.append(f"wave={wave:x}")
-      fields.append(f"op=0x{op:02x} [{OPNAME.get(op, '')}]")
-      if flag1: fields.append("flag1")
-      if flag2: fields.append("flag2")
-    case 0x14:
-      subop   = (pkt >> 16) & 0xFFFF       # (short)(w >> 0x10)
-      val32   = (pkt >> 32) & 0xFFFFFFFF   # (uint)(w >> 0x20)
-      slot    = (pkt >> 7) & 0x7           # index in local_168[...] tables
-      hi_byte = (pkt >> 8) & 0xFF          # determines config vs marker
-
-      fields.append(f"subop=0x{subop:04x}")
-      fields.append(f"slot={slot}")
-      fields.append(f"val32=0x{val32:08x}")
-
-      if hi_byte & 0x80:
-        # Config flavour: writes config words into per-slot state arrays.
-        fields.append("kind=config")
-        if subop == 0x000C:
-          fields.append("slot=lo")
-        elif subop == 0x000D:
-          fields.append("slot=hi")
-      else:
-        # COR marker: subop 0xC342, payload "COR\0" → start of a COR region.
-        if subop == 0xC342:
-          fields.append("kind=cor_stream")
-          if val32 == 0x434F5200:
-            fields.append("cor_magic='COR\\0'")
-    case 0x16:
-      # Bits:
-      #   bit8  -> 0x100
-      #   bit9  -> 0x200
-      #   bits 12..47 -> 36-bit field used as delta or marker
-      bit8 = bool(pkt & 0x100)
-      bit9 = bool(pkt & 0x200)
-      if not bit9:
-        mode = "delta"
-      elif not bit8:
-        mode = "marker"
-      else:
-        mode = "other"
-      # need to use reg here
-      val36 = (reg >> 12) & ((1 << 36) - 1)
-      fields.append(f"mode={mode}")
-      if mode != "delta":
-        fields.append(f"val36=0x{val36:x}")
-    case 0x17:
-      # From decomp (two sites with identical logic):
-      #   layout = (w >> 7) & 0x3f
-      #   mode   = (w >> 0xd) & 3
-      #   group  = (w >> 0xf) & 7
-      #   sel_a  = (w >> 0x1c) & 0xf
-      #   sel_b  = (w >> 0x21) & 7
-      #   flag4  = (w >> 0x3b) & 1  (only meaningful when layout == 4)
-      layout = (pkt >> 7)  & 0x3F
-      simd   = (pkt >> 13) & 0x3  # you can change this by changing traced simd
-      group  = (pkt >> 15) & 0x7
-      sel_a  = (pkt >> 0x1C) & 0xF
-      sel_b  = (pkt >> 0x21) & 0x7
-      flag4  = (pkt >> 0x3B) & 0x1
-
-      fields.append(f"layout={layout}")
-      fields.append(f"group={group}")
-      fields.append(f"simd={simd}")
-      fields.append(f"sel_a={sel_a}")
-      fields.append(f"sel_b={sel_b}")
-      if layout == 4:
-        fields.append(f"layout4_flag={flag4}")
-    case _:
-      fields.append(f"{pkt:X} & {reg_mask(opcode):X}")
-  return ",".join(fields)
-
-FILTER_LEVEL = getenv("FILTER", 1)
-
-DEFAULT_FILTER: tuple[int, ...] = tuple()
-# NOP + pure time + "sample"
-if FILTER_LEVEL >= 0: DEFAULT_FILTER += (0x10, 0xf, 0x11)
-# reg + event + sample + marker
-# TODO: events are probably good
-if FILTER_LEVEL >= 1: DEFAULT_FILTER += (0x14, 0x12, 0x16)
-# instruction runs + valuinst
-if FILTER_LEVEL >= 2: DEFAULT_FILTER += (0x01, 0x02, 0x03)
-# instructions dispatch (inst, immed)
-if FILTER_LEVEL >= 3: DEFAULT_FILTER += (0x4, 0x5, 0x18)
-# waves
-if FILTER_LEVEL >= 4: DEFAULT_FILTER += (0x6, 0x8, 0x9)
-
-def parse_sqtt_print_packets(data: bytes, filter=DEFAULT_FILTER, verbose=True) -> None:
-  """
-  Minimal debug: print ONE LINE per decoded token (packet).
-
-  Now prints only the actual nibbles that belong to each packet, instead of
-  the full 64-bit shift register.
-  """
-  n = len(data)
-  time = 0
-  last_printed_time = 0
-  reg = 0          # shift register
-  offset = 0       # bit offset, in steps of 4 (one nibble)
-  nib_budget = 0x40
-  flags = 0
-  token_index = 0
-  opcodes_seen = set()
-
-  while (offset >> 3) < n:
-    # 1) Fill register with nibbles according to nib_budget
-    if nib_budget != 0:
-      target = offset + 4 + ((nib_budget - 1) & ~3)
-      while offset != target and (offset >> 3) < n:
-        byte = data[offset >> 3]
-        nib = (byte >> (offset & 4)) & 0xF
-        reg = ((reg >> 4) | (nib << 60)) & ((1 << 64) - 1)
-        offset += 4
-      if offset != target: break  # don't parse past the end
-
-    # 2) Decode token from low 8 bits
-    opcode = STATE_TO_OPCODE[reg & 0xFF]
-    opcodes_seen.add(opcode)
-
-    # 4) Set next nibble budget based on opcode
-    nib_budget = NIBBLE_BUDGET[opcode & 0x1F]
-
-    # 5) Get delta
-    shift, width = DELTA_MAP_DEFAULT[opcode]
-    delta = (reg >> shift) & ((1 << width) - 1)
-
-    # 6) Update time and handle special opcodes 0xF/0x16
-    if opcode == 0x16:
-      two_bits = (reg >> 8) & 0x3
-      if two_bits == 1:
-        flags |= 0x01
-
-      # Common 36-bit field at bits [12..47]
-      if (reg & 0x200) == 0:
-        # delta mode: add 36-bit delta to time
-        pass
-      elif (reg & 0x100) == 0:
-        # marker / other modes: no time advance
-        # real marker: bit9=1, bit8=0, non-zero payload
-        # "other" 0x16 variants, ignored for timing
-        delta = 0
-      else:
-        raise RuntimeError("unknown 0x16 delta")
-    elif opcode == 0x0F:
-      # opcode 0x0F has an offset of 4 to the delta
-      # update: it's actually computed to be 8 to match WAVESTART
-      delta = delta + 8
-
-    # Append extra decoded fields into the note string
-    note = decode_packet_fields(opcode, reg)
-
-    # this delta happens before the instruction
-    time += delta
-    token_index += 1
-
-    if verbose and (filter is None or opcode not in filter):
-      print(f"{time:8d} +{time-last_printed_time:8d} : "+colored(f"{OPCODE_NAMES[opcode]:18s} ", OPCODE_COLORS.get(opcode, "white"))+f"{note}")
-      last_printed_time = time
-
-  # Optional summary at the end
-  print(f"# done: tokens={token_index:_}, final_time={time}, flags=0x{flags:02x}")
-  if verbose:
-    print(f"opcodes({len(opcodes_seen):2d}):",
-          ' '.join([colored(f"{op:2X}", "WHITE" if op in opcodes_seen else "BLACK") for op in sorted(opcode_mask)]))
-
-
-def parse(fn:str):
-  with Timing(f"unpickle {fn}: "): dat = pickle.load(open(fn, "rb"))
-  #if getenv("ROCM", 0):
-  #  with Timing(f"decode {fn}: "): ctx = decode(dat)
-  dat_sqtt = [x for x in dat if isinstance(x, ProfileSQTTEvent)]
-  print(f"got {len(dat_sqtt)} SQTT events in {fn}")
-  return dat_sqtt
-
-if __name__ == "__main__":
-  fn = "extra/sqtt/examples/profile_gemm_run_0.pkl"
-  dat_sqtt = parse(sys.argv[1] if len(sys.argv) > 1 else fn)
-  for i,dat in enumerate(dat_sqtt):
-    with Timing(f"decode pkt {i} with len {len(dat.blob):_}: "):
-      parse_sqtt_print_packets(dat.blob, verbose=getenv("V", 1))

extra/sqtt/examples/generate_examples.py

@@ -1,24 +1,25 @@
-import os, subprocess, sys
+import os, subprocess, sys, shlex
 from pathlib import Path
 from tinygrad.helpers import temp
 
 EXAMPLES_DIR = Path(__file__).parent
 PROFILE_PATH = Path(temp("profile.pkl", append_user=True))
 
-EXAMPLES = [
-  "test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
-  "test/test_tiny.py TestTiny.test_plus",
-  "test/test_tiny.py TestTiny.test_gemm",
-]
+EXAMPLES = {
+  "empty":"test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
+  "plus":"test/test_tiny.py TestTiny.test_plus",
+  "gemm":"-c \"from tinygrad import Tensor; (Tensor.empty(N:=32, N)@Tensor.empty(N, N)).realize()\"",
+  "sync":"test/amd/test_custom_kernel.py TestCustomKernel.test_wave_sync",
+}
 
 if __name__ == "__main__":
   arch = subprocess.check_output(["python", "-c", "from tinygrad import Device; print(Device['AMD'].arch)"], text=True,
                                  env={**os.environ, "DEBUG":"0"}).rstrip()
   (EXAMPLES_DIR/arch).mkdir(exist_ok=True)
-  for test in EXAMPLES:
+  for name,test in EXAMPLES.items():
     for i in range(2):
       # AM_RESET=1 gets a clear trace, does not work on mi300 machines
-      subprocess.run([sys.executable, *test.split()], cwd=EXAMPLES_DIR.parent.parent.parent,
+      subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
                      env={**os.environ, "AMD":"1", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
-      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{test.split('.')[-1].replace('test_', '')}_run_{i}.pkl")
+      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
       print(f"saved SQTT trace to {dest}")

extra/sqtt/roc.py

@@ -118,7 +118,7 @@ def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]])
     nonlocal exc
     try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
     except AttributeError as e:
-      exc = RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install")
+      exc = RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_rocprof_decoder.py to install")
       exc.__cause__ = e
   (t:=threading.Thread(target=worker, daemon=True)).start()
   t.join()

extra/thunder/amd/fa.py

@@ -19,6 +19,38 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
   return _sharded_empty(ref.shape, ref, axis)
 
+@functools.cache
+def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch):
+  def grad(dou:UOp, ker:UOp) -> tuple[None, None, UOp, UOp, UOp]:
+    do = Tensor(dou, device=dou.device)
+    attn = Tensor(ker.src[1].after(ker), device=ker.src[1].device)
+    l_vec = Tensor(ker.src[2].after(ker), device=ker.src[2].device)
+    xq = Tensor(ker.src[3], device=ker.src[3].device)
+    xk = Tensor(ker.src[4], device=ker.src[4].device)
+    xv = Tensor(ker.src[5], device=ker.src[5].device)
+
+    dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
+    GROUP_SIZE = H_local // H_KV_local
+    dk_partial = _sharded_empty((B * GROUP_SIZE, N, H_KV, D), xk, axis=shard_axis)
+    dv_partial = _sharded_empty((B * GROUP_SIZE, N, H_KV, D), xv, axis=shard_axis)
+
+    # delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
+    delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
+    delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
+
+    dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
+
+    # unshuffle dq: atomic_pk_add_bf16_with_warpid creates a shuffled layout within each 16x128 tile
+    # decompose each tile into (j=4, a=2, b=2, d=4, e=4, k=4, c=2) and permute to (e, k, j, a, d, b, c) = standard row-major
+    dq = dq.reshape(B, H, N//16, 4, 2, 2, 4, 4, 4, 2).permute(0, 1, 2, 7, 8, 3, 4, 6, 5, 9).reshape(B, H, N, D).transpose(1, 2)
+
+    # reduce partial dK/dV across GROUP_SIZE query heads
+    dk = dk_partial.reshape(B, GROUP_SIZE, N, H_KV, D).sum(1)
+    dv = dv_partial.reshape(B, GROUP_SIZE, N, H_KV, D).sum(1)
+
+    return None, None, dq.uop, dk.uop, dv.uop
+  return grad
+
 def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False):
   assert attn_mask is None, "attn_mask not supported"
   assert is_causal, "only causal attention supported"
@@ -45,23 +77,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
   attn = _sharded_empty_like(xq, axis=shard_axis)
   l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
 
-  def grad(dou:UOp, _) -> tuple[None, None, UOp, UOp, UOp]:
-    do = Tensor(dou, device=dou.device)
-    dq_in = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
-    dq = _sharded_empty_like(xq, axis=shard_axis)
-    dk = _sharded_empty_like(xk, axis=shard_axis)
-    dv = _sharded_empty_like(xv, axis=shard_axis)
-
-    # delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
-    delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
-    delta_vec, dq_in = Tensor.custom_kernel(delta_vec, dq_in, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
-
-    dq_in, dk, dv = Tensor.custom_kernel(dq_in, dk, dv, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
-
-    # unshuffle dq
-    dq = Tensor.custom_kernel(dq, dq_in, fxn=functools.partial(custom_fa_backward_post, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[0]
-
-    return None, None, dq.uop, dk.uop, dv.uop
+  grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch)
 
   attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D), grad_fxn=grad)[:2]
 
@@ -89,7 +105,6 @@ def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, device:str, arch:st
                   arg=KernelInfo(name="custom_fa_forward", estimates=estimates))
 
   lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-
   lib = bytearray(lib)
   rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
   struct.pack_into('<I', lib, rodata_off, 160000)
@@ -120,7 +135,6 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
                   arg=KernelInfo(name="custom_fa_backward_pre", estimates=estimates))
 
   lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-
   lib = bytearray(lib)
   rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
   struct.pack_into('<I', lib, rodata_off, 160000)
@@ -138,7 +152,7 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
   BLOCK_SIZE_KV = 256
   NUM_WARPS = 4
   NUM_THREADS = 64 * NUM_WARPS
-  gsz = (H_KV, N // BLOCK_SIZE_KV, B)
+  gsz = (H, N // BLOCK_SIZE_KV, B)
   lsz = (NUM_THREADS, 1, 1)
   threadIdx_x = UOp.special(lsz[0], "lidx0")
   blockIdx_x, blockIdx_y, blockIdx_z = UOp.special(gsz[0], "gidx0"), UOp.special(gsz[1], "gidx1"), UOp.special(gsz[2], "gidx2")
@@ -151,7 +165,6 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
                   arg=KernelInfo(name="custom_fa_backward", estimates=estimates))
 
   lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-
   lib = bytearray(lib)
   rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
   struct.pack_into('<I', lib, rodata_off, 160000)
@@ -182,7 +195,6 @@ def custom_fa_backward_post(dq_out:UOp, dq_in:UOp, device:str, arch:str, B:int,
                   arg=KernelInfo(name="custom_fa_backward_post", estimates=estimates))
 
   lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-
   lib = bytearray(lib)
   rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
   struct.pack_into('<I', lib, rodata_off, 160000)

extra/thunder/amd/fa_bwd_causal.cpp

@@ -37,7 +37,7 @@ using namespace kittens;
 using _gl_QdO  = gl<bf16, ATTN_B, ATTN_N, ATTN_H, ATTN_D>;
 using _gl_KV   = gl<bf16, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D>;
 using _gl_dQ   = gl<bf16, ATTN_B, ATTN_H, ATTN_N, ATTN_D>;
-using _gl_dKV  = gl<bf16, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D>;
+using _gl_dKV  = gl<bf16, ATTN_B * GROUP_SIZE, ATTN_N, ATTN_H_KV, ATTN_D>;
 using _gl_Lvec = gl<float, ATTN_B, ATTN_H, 1, ATTN_N>;
 
 template<int D> struct attn_bwd_combined_globals {
@@ -47,7 +47,7 @@ template<int D> struct attn_bwd_combined_globals {
   _gl_dQ dQg;
   _gl_dKV dKg, dVg;
   _gl_Lvec L_vec, delta_vec;
-  dim3 grid() { return dim3(ATTN_H_KV, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
+  dim3 grid() { return dim3(ATTN_H, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
   dim3 block() { return dim3(NUM_THREADS); }
   size_t dynamic_shared_memory() { return MAX_SHARED_MEMORY; }
 };
@@ -55,10 +55,12 @@ template<int D> struct attn_bwd_combined_globals {
 template<int D> __launch_bounds__(NUM_THREADS, 1)
 __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr, bf16 *dO_ptr, bf16 *Q_ptr, bf16 *K_ptr, bf16 *V_ptr, float *L_vec_ptr, float *delta_vec_ptr) {
 
-  const int kv_head_idx = blockIdx.x;  // This is the KV head index
+  const int q_head_idx_fixed = blockIdx.x;  // This is the query head index [0, ATTN_H)
+  const int kv_head_idx = q_head_idx_fixed / GROUP_SIZE;
+  const int q_head_in_group = q_head_idx_fixed % GROUP_SIZE;
   const int seq_idx = blockIdx.y;
   const int batch_idx = blockIdx.z;
-  const int first_q_head = kv_head_idx * GROUP_SIZE;
+  const int first_q_head = q_head_idx_fixed;
 
   const int warpid = kittens::warpid();
   const int j = seq_idx * NUM_WARPS + warpid;
@@ -70,7 +72,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
   // first Q step that can overlap this K_span:
   const int first_step = max(0, k_start_min / STEP_QO);
   const int num_steps_per_head = total_steps_per_head - first_step;
-  const int num_steps = num_steps_per_head * GROUP_SIZE;
+  const int num_steps = num_steps_per_head;
   const int k_pos = j * WARP_SIZE_KV;
 
   constexpr float L_SCALE_FACTOR = 1.44269504089f;
@@ -3355,14 +3357,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
     }
   }
 
-  store<1>(g.dVg, dV_j, {batch_idx, 0, kv_head_idx, 0}, {0, j, 0, 0});
+  store<1>(g.dVg, dV_j, {batch_idx * GROUP_SIZE + q_head_in_group, 0, kv_head_idx, 0}, {0, j, 0, 0});
   __builtin_amdgcn_s_waitcnt(0);
   __builtin_amdgcn_s_barrier();
 
   // We first copy dV_j_T from accumulator GPRs to vector GPRs and then perform the store
   accvgpr_read(dV_j_T, dK_j_T);
   mul(dV_j_T, dV_j_T, dP_SCALE_FACTOR);
-  store<1>(g.dKg, dV_j, {batch_idx, 0, kv_head_idx, 0}, {0, j, 0, 0});
+  store<1>(g.dKg, dV_j, {batch_idx * GROUP_SIZE + q_head_in_group, 0, kv_head_idx, 0}, {0, j, 0, 0});
 
   // Write out final dQ_i slice
   mul(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);

extra/thunder/amd/fa_fwd_causal.cpp

@@ -66,7 +66,7 @@ template<int D, typename T=bf16, typename L=row_l, typename S=rt_32x16_s> using
 template<int D, typename T=bf16, typename L=col_l, typename S=rt_16x32_s> using qo_tile_transposed = rt<T, D, Q_BLOCK_SIZE, L, S>;
 template<int D, typename T=bf16, typename L=row_l, typename S=rt_32x16_s> using kv_tile = rt<T, KV_BLOCK_SIZE, D, L, S>;
 template<int D, typename T=bf16, typename L=col_l, typename S=rt_16x32_s> using kv_tile_transposed = rt<T, D, KV_BLOCK_SIZE, L, S>;
-template<int D, typename T=float, typename L=col_l, typename S=rt_16x32_4_s> using attn_tile = rt<T, KV_BLOCK_SIZE, Q_BLOCK_SIZE, L, S>;
+template<typename T=float, typename L=col_l, typename S=rt_16x32_4_s> using attn_tile = rt<T, KV_BLOCK_SIZE, Q_BLOCK_SIZE, L, S>;
 
 /**********************************************************/
 template<int THR_X, int THR_Y>
@@ -103,7 +103,7 @@ __device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_
 
     #pragma unroll
     for (int i = 0; i < dst.height; ++i) {
-        // Row base of the 32x* chunk produced by MFMA 
+        // Row base of the 32x* chunk produced by MFMA
         const int row_base = (i * 32) + ((lane >> 5) << 2); // multiplesof 4
 
         // Relative index of the FIRST element in this row-chunk w.r.t. q_pos
@@ -148,7 +148,7 @@ __device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_
 /**********************************************************/
 
 template<int D> struct attn_globals {
-    _gl_QKVO Qg, Kg, Vg, Og; 
+    _gl_QKVO Qg, Kg, Vg, Og;
     gl<float, -1, -1, -1, -1> L_vec;
     dim3 grid() { return dim3(ATTN_H, ((ATTN_N / Q_BLOCK_SIZE + NUM_WARPS - 1) / NUM_WARPS), ATTN_B); }
     dim3 block() { return dim3(NUM_THREADS); }
@@ -196,10 +196,10 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
 
     kv_tile<D, bf16, col_l, rt_16x32_4_s> v_reg;
     qo_tile_transposed<D, float, col_l, rt_32x32_s> o_reg; // Output tile.
-    attn_tile<D, float, col_l, rt_32x32_s> att_block[2]; // attention tile, in float.
-    attn_tile<D, bf16, col_l, rt_32x32_s> att_block_bf16;
-    attn_tile<D, bf16, col_l, rt_16x32_4_s> att_block_bf16_in;
-    typename attn_tile<D, float, col_l, rt_32x32_s>::row_vec max_vec, norm_vec, max_vec_prev, scale_vec;
+    attn_tile<float, col_l, rt_32x32_s> att_block[2]; // attention tile, in float.
+    attn_tile<bf16, col_l, rt_32x32_s> att_block_bf16;
+    attn_tile<bf16, col_l, rt_16x32_4_s> att_block_bf16_in;
+    typename attn_tile<float, col_l, rt_32x32_s>::row_vec max_vec, norm_vec, max_vec_prev, scale_vec;
 
     zero(o_reg);
     zero(norm_vec);
@@ -241,8 +241,8 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
     zero(att_block[0]);
     transpose(k_reg_transposed, k_reg);
     mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
-    __builtin_amdgcn_sched_barrier(0); 
-    if constexpr (causal) { 
+    __builtin_amdgcn_sched_barrier(0);
+    if constexpr (causal) {
         const int kv_end_pos = (1) * KV_BLOCK_SIZE;
         if (__builtin_expect(q_start_pos < kv_end_pos, 0)) {  // Only mask if needed
             mask_kv_tile(att_block[0], tile_idx, 0, neg_inf_v, lane);
@@ -269,7 +269,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
     load(k_reg, k_smem[1]);
     // All warps then collaboratively load in the third slice of K (K2) into shared memory
     G::load<1, false>(k_smem[0], g.Kg, {batch_idx, 2, head_idx_kv, 0}, swizzled_offsets_K);
-    // All warps then collaboratively load in the second slice of V (V1) into shared memory 
+    // All warps then collaboratively load in the second slice of V (V1) into shared memory
     G::load<1, false>(v_smem[1], g.Vg, {batch_idx, 1, head_idx_kv, 0}, swizzled_offsets_V);
     asm volatile("s_waitcnt lgkmcnt(0)");
     asm volatile("s_waitcnt vmcnt(4)");
@@ -288,7 +288,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
         mul(norm_vec, norm_vec, scale_vec);
         col_sum(norm_vec, att_block[0], norm_vec);
         copy(att_block_bf16, att_block[0]);
-        att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+        att_block_bf16_in = *reinterpret_cast<attn_tile< bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
         sched_barrier_exp_pairs<6, 3, 1>();
         sched_barrier_pairs<10, 5, 1>();
         __builtin_amdgcn_sched_barrier(0);
@@ -296,7 +296,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
         __builtin_amdgcn_sched_barrier(0);
 
         // Cluster 1:
-        //      Load K3 into shared 
+        //      Load K3 into shared
         G::load<1, false>(k_smem[1], g.Kg, {batch_idx, j, head_idx_kv, 0}, swizzled_offsets_K);
         //      Load V0 into registers
         load(v_reg, v_smem[0]);
@@ -348,7 +348,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
         mul(norm_vec, norm_vec, scale_vec);
         col_sum(norm_vec, att_block[1], norm_vec);
         copy(att_block_bf16, att_block[1]);
-        att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+        att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
         sched_barrier_exp_pairs<6, 3, 3>();
         sched_barrier_pairs<10, 5, 3>();
         __builtin_amdgcn_s_setprio(0);
@@ -417,7 +417,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
 
     col_sum(norm_vec, att_block[0], norm_vec);
     copy(att_block_bf16, att_block[0]);
-    att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
     sched_barrier_exp_pairs<6, 3, 5>();
     sched_barrier_pairs<10, 5, 5>();
     __builtin_amdgcn_sched_barrier(0);
@@ -482,7 +482,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
     mul(norm_vec, norm_vec, scale_vec);
     col_sum(norm_vec, att_block[1], norm_vec);
     copy(att_block_bf16, att_block[1]);
-    att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
     sched_barrier_exp_pairs<6, 3, 7>();
     sched_barrier_pairs<10, 5, 7>();
     __builtin_amdgcn_sched_barrier(0);
@@ -544,7 +544,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
     mul(norm_vec, norm_vec, scale_vec);
     col_sum(norm_vec, att_block[0], norm_vec);
     copy(att_block_bf16, att_block[0]);
-    att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
     sched_barrier_exp_pairs<6, 3, 9>();
     sched_barrier_pairs<10, 5, 9>();
     __builtin_amdgcn_sched_barrier(0);
@@ -586,7 +586,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
 
     col_sum(norm_vec, att_block[1], norm_vec);
     copy(att_block_bf16, att_block[1]);
-    att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
 
     __builtin_amdgcn_sched_barrier(0);
     mul_col(o_reg, o_reg, scale_vec);

extra/torch_backend/backend.py

@@ -505,7 +505,9 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
   "aten.lt.Tensor_out": Tensor.__lt__, "aten.lt.Scalar_out": Tensor.__lt__,
   "aten.le.Tensor_out": Tensor.__le__, "aten.le.Scalar_out": Tensor.__le__,
   "aten.clamp_max.Tensor_out": lambda input,max_: input.clamp(max_=max_),
+  "aten.clamp_max.out": lambda input,max_: input.clamp(max_=max_),
   "aten.clamp_min.Tensor_out": lambda input,min_: input.clamp(min_=min_),
+  "aten.clamp_min.out": lambda input,min_: input.clamp(min_=min_),
   "aten.fmod.Tensor_out": lambda input,other: input-input.div(other, rounding_mode="trunc")*other,
   # TODO: this might result in overflow issues
   "aten.round.decimals_out": lambda self,decimals: (self*10**decimals).round()/10**decimals,
@@ -588,7 +590,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.repeat": lambda x,*repeats: Tensor.repeat(x,*repeats).contiguous(), # not a view
   "aten._softmax": lambda self,dim,half_to_float: self.softmax(dim),
   "aten._log_softmax": lambda self,dim,half_to_float: self.log_softmax(dim),
-  "aten.random_": lambda self: Tensor.randint(*self.shape, low=dtypes.min(self.dtype), high=dtypes.max(self.dtype), device=self.device, dtype=self.dtype),
+  "aten.random_": lambda self: Tensor.randint(*self.shape, low=self.dtype.min, high=self.dtype.max, device=self.device, dtype=self.dtype),
   "aten.random_.from": lambda self, from_, to: Tensor.randint(*self.shape, low=from_, high=to, device=self.device, dtype=self.dtype),
   "aten.uniform_": lambda self, low=0, high=1: Tensor.uniform(*self.shape, low=low, high=high, dtype=self.dtype),
   "aten.normal_": lambda self, mean=0, std=1: Tensor.normal(*self.shape, mean=mean, std=std, dtype=self.dtype),

extra/torch_backend/test_kernel_fusion.py

@@ -1,7 +1,6 @@
 # simple tests
 import unittest
 import torch
-import warnings
 from tinygrad.helpers import getenv, GlobalCounters
 if getenv("TINY_BACKEND2"):
   import extra.torch_backend.backend2
@@ -18,9 +17,7 @@ class TestKernelFusionRegression(unittest.TestCase):
     torch.manual_seed(42)
     GlobalCounters.reset()
     fn().detach().cpu().numpy()
-    expectation = f"{GlobalCounters.kernel_count} vs {expected_kernels} expected."
-    if GlobalCounters.kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
-    self.assertLessEqual(GlobalCounters.kernel_count, expected_kernels, f"{expectation}")
+    self.assertEqual(GlobalCounters.kernel_count, expected_kernels)
 
   def test_elementwise_fusion(self):
     def fn():
@@ -34,7 +31,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       conv = torch.nn.Conv2d(3, 16, 3, padding=1).to(device)
       with torch.no_grad():
         return torch.nn.functional.relu(conv(x))
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 6)
 
   def test_batchnorm_fusion(self):
     def fn():
@@ -44,7 +41,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       bn.eval()
       with torch.no_grad():
         return torch.nn.functional.relu(bn(conv(x)))
-    self._check_kernel_count(fn, 16)
+    self._check_kernel_count(fn, 10)
 
   def test_reduce_fusion(self):
     def fn():
@@ -92,7 +89,7 @@ class TestKernelFusionRegression(unittest.TestCase):
         out = bn(conv(x))
         out += identity
         return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 17)
+    self._check_kernel_count(fn, 12)
 
   def test_multiple_inplace_ops_fusion(self):
     def fn():
@@ -117,7 +114,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       bn.train()
       with torch.no_grad():
         return bn(x)
-    self._check_kernel_count(fn, 10)
+    self._check_kernel_count(fn, 8)
 
   # this is a minimal extra/other_mnist/beautiful_mnist_torch.py to cover fusion for training with optimizer
   def test_mnist_training_fusion(self):
@@ -138,7 +135,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       loss.backward()
       optimizer.step()
       return loss
-    self._check_kernel_count(fn, 28)
+    self._check_kernel_count(fn, 24)
 
 if __name__ == "__main__":
   unittest.main()

extra/usbgpu/tbgpu/installer/Shared/server.c

@@ -17,6 +17,7 @@
 // Protocol
 
 enum {
+  CMD_PROBE = 0,          // probe devices, returns count
   CMD_MAP_BAR = 1,        // map PCI BAR, returns size
   CMD_MAP_SYSMEM_FD = 2,  // alloc DMA memory, returns fd via SCM_RIGHTS
   CMD_CFG_READ = 3,       // read PCI config space
@@ -24,11 +25,14 @@ enum {
   CMD_RESET = 5,          // reset device
   CMD_MMIO_READ = 6,      // bulk read from BAR
   CMD_MMIO_WRITE = 7,     // bulk write to BAR
+  CMD_MAP_SYSMEM = 8,     // map system memory
+  CMD_SYSMEM_READ = 9,    // bulk read from system memory
+  CMD_SYSMEM_WRITE = 10,  // bulk write to system memory
   RESP_OK = 0, RESP_ERR = 1,
 };
 
-typedef struct { uint8_t cmd, bar; uint64_t offset, size, value; } __attribute__((packed)) request_t;
-typedef struct { uint8_t status; uint64_t value, addr; } __attribute__((packed)) response_t;
+typedef struct { uint8_t cmd; uint32_t dev_id, bar; uint64_t arg0, arg1, arg2; } __attribute__((packed)) request_t;
+typedef struct { uint8_t status; uint64_t resp0, resp1; } __attribute__((packed)) response_t;
 
 // Constants and state
 
@@ -77,7 +81,7 @@ static int send_response(int fd, response_t *resp, int send_fd) {
 }
 
 static void send_error(int fd, const char *msg) {
-  response_t resp = {.status = RESP_ERR, .value = strlen(msg)};
+  response_t resp = {.status = RESP_ERR, .resp0 = strlen(msg)};
   send_response(fd, &resp, -1);
   send(fd, msg, strlen(msg), 0);
 }
@@ -116,8 +120,8 @@ static int dext_rpc(uint32_t sel, uint64_t *in, uint32_t in_cnt, uint64_t *out_v
 static int map_bar(uint32_t bar, response_t *resp) {
   if (bar >= MAX_BARS) return -1;
   if (!g_bars[bar].addr && IOConnectMapMemory64(g_conn, bar, mach_task_self(), &g_bars[bar].addr, &g_bars[bar].size, kIOMapAnywhere)) return -1;
-  resp->addr = g_bars[bar].addr;
-  resp->value = g_bars[bar].size;
+  resp->resp0 = g_bars[bar].addr;
+  resp->resp1 = g_bars[bar].size;
   return 0;
 }
 
@@ -148,7 +152,7 @@ static int map_sysmem_fd(uint64_t size, response_t *resp, int *out_fd) {
   strncpy(g_sysmem[idx].shm_name, shm_name, sizeof(g_sysmem[idx].shm_name));
   g_sysmem_count++;
 
-  *resp = (response_t){.addr = idx, .value = alloc_sz};
+  *resp = (response_t){.resp0 = alloc_sz, .resp1 = idx};
   *out_fd = fd;
   return 0;
 
@@ -203,19 +207,19 @@ static void handle_client(int fd) {
 
     case CMD_MAP_SYSMEM_FD: {
       int shm_fd = -1;
-      resp.status = map_sysmem_fd(req.size, &resp, &shm_fd) ? 1 : 0;
+      resp.status = map_sysmem_fd(req.arg0, &resp, &shm_fd) ? 1 : 0;
       send_response(fd, &resp, shm_fd);
       continue;
     }
 
     case CMD_CFG_READ: {
-      uint64_t in[2] = {req.offset, req.size};
-      resp.status = dext_rpc(0, in, 2, &resp.value) ? 1 : 0;
+      uint64_t in[2] = {req.arg0, req.arg1};
+      resp.status = dext_rpc(0, in, 2, &resp.resp0) ? 1 : 0;
       break;
     }
 
     case CMD_CFG_WRITE: {
-      uint64_t in[3] = {req.offset, req.size, req.value};
+      uint64_t in[3] = {req.arg0, req.arg1, req.arg2};
       resp.status = dext_rpc(1, in, 3, NULL) ? 1 : 0;
       break;
     }
@@ -225,17 +229,17 @@ static void handle_client(int fd) {
       break;
 
     case CMD_MMIO_READ:
-      if (validate_bar(req.bar, req.offset, req.size)) { resp.status = 1; break; }
-      mmio_copy(g_bulk_buf, (void*)(g_bars[req.bar].addr + req.offset), req.size);
-      resp.value = req.size;
+      if (validate_bar(req.bar, req.arg0, req.arg1)) { resp.status = 1; break; }
+      mmio_copy(g_bulk_buf, (void*)(g_bars[req.bar].addr + req.arg0), req.arg1);
+      resp.resp0 = req.arg1;
       send_response(fd, &resp, -1);
-      send(fd, g_bulk_buf, req.size, 0);
+      send(fd, g_bulk_buf, req.arg1, 0);
       continue;
 
     case CMD_MMIO_WRITE:
-      recvall(fd, g_bulk_buf, req.size);
-      if (!validate_bar(req.bar, req.offset, req.size))
-        mmio_copy((void*)(g_bars[req.bar].addr + req.offset), g_bulk_buf, req.size);
+      recvall(fd, g_bulk_buf, req.arg1);
+      if (!validate_bar(req.bar, req.arg0, req.arg1))
+        mmio_copy((void*)(g_bars[req.bar].addr + req.arg0), g_bulk_buf, req.arg1);
       continue;
 
     default:

extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension/TinyGPUDriver.NV.Release.entitlements

@@ -0,0 +1,17 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>com.apple.application-identifier</key>
+	<string>9YG3G8543N.org.tinygrad.tinygpu.edriver</string>
+	<key>com.apple.developer.driverkit</key>
+	<true/>
+	<key>com.apple.developer.driverkit.transport.pci</key>
+	<array>
+		<dict>
+			<key>IOPCIPrimaryMatch</key>
+			<string>0x000010de&amp;0x0000FFFF</string>
+		</dict>
+	</array>
+</dict>
+</plist>

extra/usbgpu/tbgpu/installer/build_and_sign_nv.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+set -e
+
+xcodebuild clean build CODE_SIGN_IDENTITY="" CODE_SIGNING_REQUIRED=NO -alltargets -configuration Release build
+
+cp "../profiles/edriver_rel_2.provisionprofile" "./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext/embedded.provisionprofile"
+cp "../profiles/installer_provisioning.provisionprofile" "./build/Release/TinyGPU.app/Contents/embedded.provisionprofile"
+
+codesign \
+    --sign "Developer ID Application: tinygrad, Corp. (9YG3G8543N)" \
+    --entitlements ./TinyGPUDriverExtension/TinyGPUDriver.NV.Release.entitlements \
+    --verbose \
+    --options runtime \
+    --timestamp \
+    --force \
+    ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
+
+codesign \
+    --sign "Developer ID Application: tinygrad, Corp. (9YG3G8543N)" \
+    --entitlements ./macOS/macOS.entitlements \
+    --options runtime \
+    --verbose \
+    --timestamp \
+    --force \
+    ./build/Release/TinyGPU.app
+
+codesign --verify --deep --strict --verbose=4 ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
+
+codesign --verify --deep --strict --verbose=4 ./build/Release/TinyGPU.app
+
+spctl -a -vv ./build/Release/TinyGPU.app
+
+spctl -a -vv ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext

extra/usbgpu/tbgpu/installer/notary_tool.sh

@@ -3,3 +3,7 @@ set -e
 
 ditto -c -k --keepParent ./build/Release/TinyGPU.app ./build/Release/TinyGPU.zip
 xcrun notarytool submit ./build/Release/TinyGPU.zip --keychain-profile "hgwJFhdheiIEy82nDN" --wait
+
+rm ./build/Release/TinyGPU.zip
+xcrun stapler staple ./build/Release/TinyGPU.app
+ditto -c -k --keepParent ./build/Release/TinyGPU.app ./build/Release/TinyGPU.zip

extra/viz/README

@@ -1,17 +1,17 @@
 A command line tool for exploring the VIZ trace.
 
-After running with VIZ=-1, use `PYTHONPATH=. extra/viz/cli.py` to explore the saved trace files.
+After running with VIZ=-1, use `extra/viz/cli.py` to explore the saved trace files.
 
 ## Inspect runtime profiling
 
-Use `PYTHONPATH=. extra/viz/cli.py --profile` to list all traced devices.
+Use `extra/viz/cli.py --profile` to list all traced devices.
 
 List top slowest kernels on a device: `--profile --device "AMD"`
 List samples of a kernel on a device: `--profile --device "AMD" --kernel E_3`
 
 ## Inspect codegen and PatternMatcher
 
-Use `PYTHONPATH=. extra/viz/cli.py --rewrites` to list all traced kernels.
+Use `extra/viz/cli.py --rewrites` to list all traced kernels.
 
 List all codegen steps for a kernel: `--rewrites --kernel E_3`
 Get source code: `--rewrites --kernel E_3 --select "View Source"`

extra/viz/cli.py

@@ -1,44 +1,78 @@
 #!/usr/bin/env python3
 import os
 os.environ["VIZ"] = "0"
-import argparse, pathlib
+import argparse, pathlib, sys, struct, json
 from typing import Iterator
 from tinygrad.viz import serve as viz
 from tinygrad.uop.ops import RewriteTrace
 from tinygrad.helpers import temp, ansistrip, colored, time_to_str, ansilen
-from test.null.test_viz import load_profile
+
+# ** generic helpers
 
 def optional_eq(val:dict, arg:str|None) -> bool: return arg is None or ansistrip(val["name"]) == arg
 
 def print_data(data:dict) -> None:
   if isinstance(data.get("value"), Iterator):
     for m in data["value"]:
-      if m.get("uop"):
-        print("Input UOp:")
-        print(m["uop"])
-      if not m["diff"]: continue
-      print("Rewrites:")
-      fp = pathlib.Path(m["upat"][0][0])
-      print(f"{fp.parent.name}/{fp.name}:{m['upat'][0][1]}")
-      print(m["upat"][1])
-      for line in m["diff"]:
-        color = "red" if line.startswith("-") else "green" if line.startswith("+") else None
-        print(colored(line, color))
+      if m.get("uop"): print(f"Input UOp:\n{m['uop']}")
+      if m.get("diff"):
+        loc = pathlib.Path(m["upat"][0][0])
+        print(f"Rewrite at {loc.parent.name}/{loc.name}:{m['upat'][0][1]}\n{m['upat'][1]}")
+        for line in m["diff"]: print(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None))
   if data.get("src") is not None: print(data["src"])
 
+# ** Profiler trace decoder
+
+# 0 means None, otherwise it's an enum value
+def option(i:int) -> int|None: return None if i == 0 else i-1
+
+def decode_profile(data:bytes) -> dict:
+  ret, off = data, 0
+  def u(fmt:str) -> tuple:
+    nonlocal off
+    vals = struct.unpack_from(fmt, ret, off)
+    off += struct.calcsize(fmt)
+    return vals
+  total_dur, global_peak, index_len, layout_len = u("<IQII")
+  strings, dtypes, markers = json.loads(ret[off:off+index_len]).values()
+  off += index_len
+  layout:dict[str, dict] = {}
+  for _ in range(layout_len):
+    klen = u("<B")[0]
+    k = ret[off:off+klen].decode()
+    off += klen
+    layout[k] = v = {"events":[]}
+    event_type, event_count = u("<BI")
+    if event_type == 0:
+      for _ in range(event_count):
+        name, ref, key, st, dur, fmt = u("<IIIIfI")
+        v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur, "fmt":strings[fmt]})
+    else:
+      v["linear"] = u("<B")[0]
+      v["peak"] = u("<Q")[0]
+      for _ in range(event_count):
+        if v["linear"]:
+          ts, value = u("<IQ")
+          v["events"].append({"event":"freq", "ts":ts, "value":value})
+        else:
+          alloc, ts, key = u("<BII")
+          if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
+          else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
+  return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
+
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   g_mode = parser.add_argument_group("mode")
   g_mode.add_argument("--profile", action="store_true", help="View profile trace")
   g_mode.add_argument("--rewrites", action="store_true", help="View rewrites trace")
+  g_common = parser.add_argument_group("common options")
+  g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
   g_profile = parser.add_argument_group("profile options")
   g_profile.add_argument("--device", type=str, default=None, metavar="NAME", help="Select a device (optional name, default: only list names)")
   g_profile.add_argument("--top", type=int, default=10, metavar="N", help="Number of top kernels to show (-1 for all, default: 10)")
   g_rewrites = parser.add_argument_group("rewrites options")
   g_rewrites.add_argument("--select", type=str, default=None, metavar="NAME",
                           help="Select an item within the chosen kernel (optional name, default: only list names)")
-  g_common = parser.add_argument_group("common options")
-  g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
   parser.add_argument("--profile-path", type=pathlib.Path, metavar="PATH", help="Path to profile (optional file, default: latest profile)",
                         default=pathlib.Path(temp("profile.pkl", append_user=True)))
   parser.add_argument("--rewrites-path", type=pathlib.Path, metavar="PATH", help="Path to rewrites (optional file, default: latest rewrites)",
@@ -46,14 +80,14 @@ if __name__ == "__main__":
   args = parser.parse_args()
   if not args.profile and not args.rewrites:
     parser.print_help()
-    exit(0)
+    sys.exit(0)
 
   viz.trace = viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))
   viz.ctxs = viz.get_rewrites(viz.trace)
 
   if args.profile:
     from tabulate import tabulate
-    profile = load_profile(viz.load_pickle(args.profile_path, default=[]))
+    profile = decode_profile(viz.get_profile(viz.load_pickle(args.profile_path, default=[])))
     agg, total, n = {}, 0, 0
     if args.device is None: print("Select a device:")
     for k,v in profile["layout"].items():
@@ -63,7 +97,7 @@ if __name__ == "__main__":
       for e in v.get("events", []):
         et = e["dur"]*1e-6
         if args.kernel is not None:
-          if ansistrip(e["name"]) == args.kernel and n < 10:
+          if optional_eq(e, args.kernel) and n < 10:
             ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
             name = e["name"]+(" " * (46 - ansilen(e["name"])))
             print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms  "+e['fmt'].replace('\n', ' | ')+"  ")
@@ -81,7 +115,7 @@ if __name__ == "__main__":
         other_t = total-sum(t for _, (t, _) in sel)
         table.append([f"Other ({len(other)} unique)", time_to_str(other_t, w=9), sum(c for _,(_,c) in other), f"{other_t/total*100.0:.2f}%"])
       print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github"))
-    exit(0)
+    sys.exit(0)
 
   for k in viz.ctxs:
     if not optional_eq(k, args.kernel): continue

pyproject.toml

@@ -74,7 +74,7 @@ testing_minimal = [
   "hypothesis>=6.148.9",
   "z3-solver<4.15.4",  # 4.15.4 has a segfault when creating many z3.Context()
 ]
-testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "gguf"]
+testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "gguf>=0.18"]
 testing = [
   "tinygrad[testing_unit]",
   "pillow",

test/amd/disasm.py

@@ -324,6 +324,12 @@ def _disasm_smem(inst: SMEM) -> str:
   if name in ('s_memrealtime', 's_memtime'): return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}"
   return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
 
+R4_TH_LOAD = {1: 'TH_LOAD_NT', 2: 'TH_LOAD_HT', 3: 'TH_LOAD_LU', 4: 'TH_LOAD_RT_WB', 5: 'TH_LOAD_NT_WB'}
+R4_TH_STORE = {1: 'TH_STORE_NT', 2: 'TH_STORE_HT', 3: 'TH_STORE_ST', 4: 'TH_STORE_RT_WB', 5: 'TH_STORE_NT_WB'}
+R4_TH_ATOMIC = {1: 'TH_ATOMIC_RETURN', 2: 'TH_ATOMIC_NT', 3: 'TH_ATOMIC_RETURN_NT',
+                4: 'TH_ATOMIC_CASCADE_RT', 5: 'TH_ATOMIC_CASCADE_RETURN', 6: 'TH_ATOMIC_CASCADE_NT', 7: 'TH_ATOMIC_CASCADE_RETURN_NT'}
+R4_SCOPE = {1: 'SCOPE_SE', 2: 'SCOPE_DEV', 3: 'SCOPE_SYS'}
+
 def _disasm_flat(inst: FLAT) -> str:
   name, cdna, r4 = inst.op_name.lower(), _is_cdna(inst), _is_r4(inst)
   acc = getattr(inst, 'acc', 0)
@@ -331,9 +337,10 @@ def _disasm_flat(inst: FLAT) -> str:
   if r4: seg = 'flat' if (cls_name:=inst.__class__.__name__) == 'VFLAT' else ('global' if cls_name == 'VGLOBAL' else 'scratch')
   else: seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
   instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
-  # Global/scratch uses 13-bit signed offset
+  # Global/scratch uses 13-bit signed offset (RDNA3/CDNA), 24-bit signed offset (RDNA4)
   offset = inst.ioffset if r4 else inst.offset  # type: ignore[attr-defined]
-  if seg != 'flat':
+  if r4: off_val = offset if offset < (1 << 23) else offset - (1 << 24)  # sign extend 24-bit
+  elif seg != 'flat':
     if cdna:
       # CDNA: bit 12 is sign bit but not in offset field
       raw = int.from_bytes(inst.to_bytes(), 'little')
@@ -348,7 +355,9 @@ def _disasm_flat(inst: FLAT) -> str:
   w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'atomic' in name else regs.get('d', 1)
   off_s = f" offset:{off_val}" if off_val else ""
   if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}"  # type: ignore[attr-defined]
-  elif r4: mods = f"{off_s}{' scope' if inst.scope else ''}{' th' if inst.th else ''}"  # type: ignore[attr-defined]
+  elif r4:
+    th_names = R4_TH_ATOMIC if 'atomic' in name else (R4_TH_STORE if 'store' in name else R4_TH_LOAD)
+    mods = off_s + (f" th:{th_names[inst.th]}" if inst.th in th_names else "") + (f" scope:{R4_SCOPE[inst.scope]}" if inst.scope in R4_SCOPE else "")
   else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
   if seg == 'flat': saddr_s = ""
   elif _unwrap(inst.saddr) in (0x7F, 124): saddr_s = ", off"
@@ -357,7 +366,7 @@ def _disasm_flat(inst: FLAT) -> str:
     saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[_unwrap(inst.saddr)]}"
   elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
   else: saddr_s = f", {_sreg(inst.saddr, 2) if _unwrap(inst.saddr) < 106 else decode_src(_unwrap(inst.saddr), cdna)}"
-  if 'addtid' in name: return f"{instr} {reg_fn(inst.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
+  if 'addtid' in name: return f"{instr} {reg_fn((inst.vsrc if r4 else inst.data) if 'store' in name else inst.vdst)}{saddr_s}{mods}"
    # RDNA4: vaddr instead of addr, vsrc instead of data
   addr = inst.vaddr if r4 else inst.addr  # type: ignore[attr-defined]
   data = inst.vsrc if r4 else inst.data  # type: ignore[attr-defined]
@@ -372,7 +381,7 @@ def _disasm_flat(inst: FLAT) -> str:
   addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
   data_s, vdst_s = reg_fn(data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
   if 'atomic' in name:
-    glc_or_sc0 = inst.sc0 if cdna else inst.glc  # type: ignore[attr-defined]
+    glc_or_sc0 = inst.sc0 if cdna else (inst.th & 1 if r4 else inst.glc)  # type: ignore[attr-defined]
     sfx = f"{saddr_s if seg != 'flat' else ''}{mods}"
     return f"{instr} {vdst_s}, {addr_s}, {data_s}{sfx}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{sfx}"
   if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"

test/amd/helpers.py

@@ -5,8 +5,8 @@ from tinygrad.runtime.autogen import llvm
 from tinygrad.runtime.support.elf import elf_loader
 
 ARCH_TO_TARGET:dict[str, list[str]] = {
-  "rdna3":["gfx1100"],
-  "rdna4":["gfx1200"],
+  "rdna3":["gfx1100", "gfx1151"],
+  "rdna4":["gfx1200", "gfx1201"],
   "cdna":["gfx950", "gfx942"],
 }
 

test/amd/hw/test_vopc.py

@@ -104,6 +104,34 @@ class TestCmpClass(unittest.TestCase):
     st = run_program(instructions, n_lanes=1)
     self.assertEqual(st.vcc & 1, 0, "Signaling NaN should not match quiet mask")
 
+  def test_v_cmp_lg_f32_nan(self):
+    """v_cmp_lg_f32 is ordered not-equal (<>): NaN <> x should be False per IEEE 754."""
+    quiet_nan = 0x7fc00000
+    one_f32 = 0x3f800000  # 1.0f
+    instructions = [
+      s_mov_b32(s[0], quiet_nan),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[1], one_f32),
+      v_mov_b32_e32(v[1], s[1]),
+      v_cmp_lg_f32_e32(v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc & 1, 0, "v_cmp_lg_f32(NaN, 1.0) should be 0")
+
+  def test_v_cmp_neq_f32_nan(self):
+    """v_cmp_neq_f32 is unordered not-equal (!=): NaN != x should be True per IEEE 754."""
+    quiet_nan = 0x7fc00000
+    one_f32 = 0x3f800000  # 1.0f
+    instructions = [
+      s_mov_b32(s[0], quiet_nan),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[1], one_f32),
+      v_mov_b32_e32(v[1], s[1]),
+      v_cmp_neq_f32_e32(v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc & 1, 1, "v_cmp_neq_f32(NaN, 1.0) should be 1")
+
   def test_v_cmp_sets_vcc_bits(self):
     """V_CMP_EQ sets VCC bits based on per-lane comparison."""
     instructions = [

test/amd/test_custom_kernel.py

@@ -1,10 +1,12 @@
 import unittest
+import functools
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
-
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
+from tinygrad.runtime.autogen.amd.rdna4.ins import s_barrier_wait, s_barrier_signal
 from tinygrad.renderer.amd.dsl import s, v
+from test.amd.helpers import TARGET_TO_ARCH
 
 def custom_add_one(A:UOp) -> UOp:
   A = A.flatten()
@@ -43,9 +45,26 @@ def custom_add_var(A:UOp, B:UOp) -> UOp:
   sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(f"custom_add_var_{A.size}"))
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
 
+def custom_wave_sync(A:UOp, arch:str) -> UOp:
+  # 4 waves across 1024 WG — enough to saturate a SIMD with many concurrent WGs
+  # s_sleep yields the SIMD so waves from different WGs interleave, causing barrier packet reordering
+  threads = UOp.special(128, "lidx0")
+  wg = UOp.special(1024, "gidx0")
+  insts = []
+  for _ in range(4):
+    insts.append(s_sleep(4))
+    insts += [s_barrier()] if arch == "rdna3" else [s_barrier_signal(), s_barrier_wait()]
+    insts += [s_nop(0)]*4
+  insts.append(s_endpgm())
+  sink = UOp.sink(A.base, threads, wg, arg=KernelInfo("custom_wave_sync"))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+
 @unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestCustomKernel(unittest.TestCase):
+  def setUp(self): self.arch = TARGET_TO_ARCH[Device["AMD"].arch]
+
   def test_simple(self):
+    if self.arch != "rdna3": self.skipTest("only rdna3")
     a = Tensor.full((16, 16), 1.).contiguous().realize()
     a = Tensor.custom_kernel(a, fxn=custom_add_one)[0]
     ei = a.schedule()[-1].lower()
@@ -55,6 +74,7 @@ class TestCustomKernel(unittest.TestCase):
     self.assertTrue((a.numpy() == 2.).all())
 
   def test_variable(self):
+    if self.arch != "rdna3": self.skipTest("only rdna3")
     b = Tensor.full((16, 16), 1, dtype=dtypes.uint32).contiguous().realize()
     a = Tensor.zeros_like(b).contiguous().realize()
     a = Tensor.custom_kernel(a, b, fxn=custom_add_var)[0]
@@ -63,5 +83,9 @@ class TestCustomKernel(unittest.TestCase):
       ei.run({"var":i})
       self.assertTrue((a.numpy() == 1+i).all())
 
+  def test_wave_sync(self):
+    if self.arch not in {"rdna3", "rdna4"}: self.skipTest("only rdna3 or rdna4")
+    Tensor.empty(1).custom_kernel(fxn=functools.partial(custom_wave_sync, arch=self.arch))[0].realize()
+
 if __name__ == "__main__":
   unittest.main()

test/amd/test_llvm.py

@@ -40,7 +40,7 @@ RDNA4_FILES = ['gfx12_asm_sop1.s', 'gfx12_asm_sop2.s', 'gfx12_asm_sopp.s', 'gfx1
   'gfx12_asm_vop1.s', 'gfx12_asm_vop2.s', 'gfx12_asm_vopc.s', 'gfx12_asm_vopcx.s', 'gfx12_asm_vop3.s', 'gfx12_asm_vop3c.s',
   'gfx12_asm_vop3cx.s', 'gfx12_asm_vop3p.s', 'gfx12_asm_vop3_from_vop1.s', 'gfx12_asm_vop3_from_vop2.s',
   'gfx12_asm_vop3p_features.s', 'gfx12_asm_vopd.s', 'gfx12_asm_vopd_features.s',
-  'gfx12_asm_ds.s', 'gfx12_asm_smem.s',
+  'gfx12_asm_ds.s', 'gfx12_asm_smem.s', 'gfx12_asm_vflat.s',
   'gfx12_asm_wmma_w32.s']
 
 def _parse_llvm_tests(text: str, pattern: str) -> list[tuple[str, bytes]]:

test/amd/test_sqtt_examples.py

@@ -9,8 +9,8 @@ from tinygrad.renderer.amd import decode_inst
 from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP
 from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp
 from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, INST, INST_RDNA4, VALUINST,
-                                     IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4,
-                                     InstOp, InstOpRDNA4, print_packets)
+                                     IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4, PACKET_TYPES_CDNA, CDNA_WAVESTART,
+                                     InstOp, InstOpRDNA4, print_packets, CDNA_WAVEEND, CDNA_INST)
 from test.amd.helpers import TARGET_TO_ARCH
 
 import tinygrad
@@ -21,7 +21,7 @@ OTHER_SIMD_OPS = {InstOp.OTHER_LDS_LOAD, InstOp.OTHER_LDS_STORE, InstOp.OTHER_LD
                   InstOp.OTHER_FLAT_STORE_128, InstOp.OTHER_GLOBAL_LOAD, InstOp.OTHER_GLOBAL_LOAD_VADDR,
                   InstOp.OTHER_GLOBAL_STORE_64, InstOp.OTHER_GLOBAL_STORE_96, InstOp.OTHER_GLOBAL_STORE_128,
                   InstOp.OTHER_GLOBAL_STORE_VADDR_128}
-OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM, InstOpRDNA4.OTHER_VMEM_STORE}
+OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM, InstOpRDNA4.OTHER_VMEM_5, InstOpRDNA4.OTHER_LDS_1, InstOpRDNA4.OTHER_LDS_2}
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # ROCPROF DECODER
@@ -125,7 +125,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
           self.assertIsInstance(packets[0], LAYOUT_HEADER, f"first packet should be LAYOUT_HEADER in {name}")
 
   def test_packet_types_valid(self):
-    all_classes = set(PACKET_TYPES_RDNA3.values()) | set(PACKET_TYPES_RDNA4.values())
+    all_classes = set(PACKET_TYPES_RDNA3.values()) | set(PACKET_TYPES_RDNA4.values()) | set(PACKET_TYPES_CDNA.values())
     for name, (events, *_) in self.examples.items():
       for i, event in enumerate(events):
         with self.subTest(example=name, event=i):
@@ -138,8 +138,8 @@ class SQTTExamplesTestBase(unittest.TestCase):
       if "empty" in name: continue
       with self.subTest(example=name):
         all_packets = [p for e in events for p in decode(e.blob)]
-        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_RDNA4))]), 0, f"no WAVESTART in {name}")
-        self.assertGreater(len([p for p in all_packets if isinstance(p, WAVEEND)]), 0, f"no WAVEEND in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_RDNA4, CDNA_WAVESTART))]), 0, f"no WAVESTART in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVEEND, CDNA_WAVEEND))]), 0, f"no WAVEEND in {name}")
 
   def test_time_monotonic(self):
     for name, (events, *_) in self.examples.items():
@@ -153,7 +153,10 @@ class SQTTExamplesTestBase(unittest.TestCase):
       if "gemm" not in name: continue
       with self.subTest(example=name):
         all_packets = [p for e in events for p in decode(e.blob)]
-        self.assertGreater(len([p for p in all_packets if isinstance(p, (INST, INST_RDNA4))]), 0, f"no INST packets in {name}")
+        inst_packets = [p for p in all_packets if isinstance(p, (INST, INST_RDNA4, CDNA_INST))]
+        self.assertGreater(len(inst_packets), 0, f"no INST packets in {name}")
+        if isinstance(inst_packets[0], (INST, INST_RDNA4)):
+          self.assertGreater(len([p for p in inst_packets if p.op.name.startswith("JUMP")]), 0, f"no JUMP packets in {name}")
 
   expected: dict[str, list[int]] = {}  # override in subclasses
   def test_packet_counts(self):
@@ -181,8 +184,8 @@ class SQTTExamplesTestBase(unittest.TestCase):
         for event in events:
           wave_starts: dict[tuple[int, int, int], int] = {}
           for p in decode(event.blob):
-            if isinstance(p, (WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
-            elif isinstance(p, WAVEEND) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
+            if isinstance(p, (WAVESTART, CDNA_WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
+            elif isinstance(p, (WAVEEND, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
               our_waves.append((wave_starts[key], p._time))
         self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")
 
@@ -208,17 +211,22 @@ class SQTTExamplesTestBase(unittest.TestCase):
 class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
   target = "gfx1100"
   expected = {
-    "profile_empty_run_0": [1744, 1801, 1854, 1890, 1917, 1822],
-    "profile_empty_run_1": [1744, 1801, 1854, 1886, 1921, 1906],
-    "profile_gemm_run_0": [1800, 1867, 1899, 1898, 1914, 1895, 1694, 1779, 1819, 1872, 1877, 1858, 1750, 1834, 1866, 1834, 1911, 1796],
-    "profile_gemm_run_1": [1806, 1874, 1837, 1885, 1907, 1906, 1694, 1778, 1810, 1873, 1885, 1867, 1750, 1834, 1866, 1856, 1903, 1897],
-    "profile_plus_run_0": [1744, 1878, 1854, 1890, 1878, 1910],
-    "profile_plus_run_1": [1744, 1878, 1854, 1886, 1921, 1909],
+    "profile_empty_run_0": [1880, 1867, 1920, 1971, 1998, 1904],
+    "profile_empty_run_1": [1880, 1867, 1920, 1971, 1998, 1904],
+    "profile_gemm_run_0": [3275, 3278, 2426, 2475, 2511, 2431],
+    "profile_gemm_run_1": [3264, 3268, 2420, 2469, 2504, 2401],
+    "profile_ops_run_0": [1944, 4903, 1984, 2035, 2062, 1968],
+    "profile_ops_run_1": [1944, 4918, 1984, 2035, 2062, 1968],
+    "profile_plus_run_0": [1938, 1932, 1978, 2029, 2056, 1962],
+    "profile_plus_run_1": [1891, 1874, 1931, 1982, 2009, 1915],
   }
 
 class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"
-@unittest.skip("TODO: fix CDNA")
-class TestSQTTExamplesCDNA(SQTTExamplesTestBase): target = "gfx950"
+
+class TestSQTTExamplesCDNA(SQTTExamplesTestBase):
+  target = "gfx950"
+  def test_rocprof_wave_times_match(self): self.skipTest("TODO: requires timestamp patching")
+  def test_rocprof_inst_times_match(self): self.skipTest("TODO: requires timestamp patching")
 
 if __name__ == "__main__":
   unittest.main()

test/amd/test_sqtt_tables.py

@@ -1,183 +0,0 @@
-"""Tests comparing sqtt.py PACKET_TYPES_RDNA3/RDNA4 against AMD's rocprof-trace-decoder binary."""
-import unittest, struct, ctypes, pickle
-from pathlib import Path
-
-ROCPROF_LIB = Path("/usr/lib/librocprof-trace-decoder.so")
-import tinygrad
-EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
-
-# CDNA pkt_fmt -> size in bytes (extracted from rocprof hash table)
-CDNA_PKT_SIZES = {0: 2, 1: 8, 2: 8, 3: 4, 4: 2, 5: 6, 6: 2, 7: 2, 8: 2, 9: 2, 10: 2, 11: 8, 12: 6, 13: 4, 14: 8, 15: 6}
-
-def _find_segment(perms: str):
-  """Find a segment of the loaded library with given permissions (e.g. 'rw-p', 'r--p')."""
-  with open('/proc/self/maps', 'r') as f:
-    for line in f:
-      if 'librocprof-trace-decoder.so' in line and f' {perms} ' in line:
-        parts = line.split()
-        return int(parts[0].split('-')[0], 16), int(parts[2], 16)
-  return None, None
-
-def _read_array(file_offset: int, count: int):
-  """Read an array of uint8 at file_offset from the loaded library."""
-  base, seg_offset = _find_segment('rw-p')
-  if base is None: return None
-  return list((ctypes.c_uint8 * count).from_address(base + (file_offset - seg_offset)))
-
-def _load_lib():
-  if not ROCPROF_LIB.exists(): return False
-  ctypes.CDLL(str(ROCPROF_LIB))
-  return True
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# RDNA EXTRACTION (nibble-based format)
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def extract_bit_tables():
-  """Extract bit budget tables. Returns (layout2, layout3, layout4) or None."""
-  if not _load_lib(): return None
-  return _read_array(0x2d220, 32), _read_array(0x2d280, 32), _read_array(0x2d2c0, 32)
-
-def extract_delta_fields():
-  """Extract delta bitfield tables. Returns (layout2, layout3, layout4) dicts mapping type_id -> (lo, hi)."""
-  if not _load_lib(): return None
-  ro_base, ro_offset = _find_segment('r--p')
-  if ro_base is None: return None
-
-  def read_table(file_offset, num_entries):
-    addr = ro_base + (file_offset - ro_offset)
-    data = bytes((ctypes.c_uint8 * (num_entries * 12)).from_address(addr))
-    return {type_id: (lo, hi) for j in range(0, len(data), 12)
-            for type_id, lo, hi in [struct.unpack('<III', data[j:j+12])] if type_id < 32}
-
-  return read_table(0x26800, 24), read_table(0x26dc0, 25), read_table(0x27300, 27)
-
-def extract_packet_encodings():
-  """Extract packet encodings. Returns (L2, L3, L4) dicts mapping type_id -> (mask, value)."""
-  if not _load_lib(): return None
-  rw_base, rw_offset = _find_segment('rw-p')
-  if rw_base is None: return None
-
-  # Read base encodings from registration vector at 0x2d340
-  vec_start = ctypes.c_void_p.from_address(rw_base + (0x2d340 - rw_offset)).value
-  vec_end = ctypes.c_void_p.from_address(rw_base + (0x2d348 - rw_offset)).value
-  base = {}
-  if vec_start and vec_end:
-    for i in range((vec_end - vec_start) // 32):
-      addr = vec_start + i * 32
-      type_id = ctypes.c_uint8.from_address(addr).value
-      pat_start = ctypes.c_void_p.from_address(addr + 8).value
-      pat_end = ctypes.c_void_p.from_address(addr + 16).value
-      if pat_start and pat_end and 0 < (n := pat_end - pat_start) <= 8:
-        pat = list((ctypes.c_uint8 * n).from_address(pat_start))
-        base[type_id] = (sum(1 << j for j in range(n)), sum(b << j for j, b in enumerate(pat)))
-
-  return {**base, 17: (0x7f, 0x51), 25: (0x7f, 0x31)}, base, {**base}  # L2 has overrides
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# CDNA EXTRACTION (16-bit header format)
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def extract_cdna_packet_sizes():
-  """Extract CDNA pkt_fmt -> size mapping by running rocprof decoder to populate its hash table."""
-  if not _load_lib(): return None
-  from test.amd.test_sqtt_examples import run_rocprof_decoder
-
-  if not (pkl_path := next((EXAMPLES_DIR / "gfx950").glob("*.pkl"), None)): return None
-  with open(pkl_path, "rb") as f: data = pickle.load(f)
-  sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
-  prg = next((e for e in data if type(e).__name__ == "ProfileProgramEvent"), None)
-  if not sqtt_events or not prg: return None
-
-  # Run decoder to trigger hash table initialization
-  run_rocprof_decoder([e.blob for e in sqtt_events], prg.lib, prg.base, "gfx950")
-
-  # Extract hash table: head at 0x2d4f0, nodes are 16 bytes (next[8], key[4], value[4])
-  rw_base, rw_offset = _find_segment('rw-p')
-  if not (head := ctypes.c_void_p.from_address(rw_base + (0x2d4f0 - rw_offset)).value if rw_base else None): return None
-
-  pkt_sizes: dict[int, int] = {}
-  node, seen = head, set()
-  while node and node not in seen and len(pkt_sizes) < 20:
-    seen.add(node)
-    key, val = ctypes.c_uint32.from_address(node + 8).value, ctypes.c_uint32.from_address(node + 12).value
-    if key < 16 and val in (0x10, 0x20, 0x30, 0x40): pkt_sizes[key] = {0x10: 2, 0x20: 4, 0x30: 6, 0x40: 8}[val]
-    node = ctypes.c_void_p.from_address(node).value  # type: ignore[assignment]
-  return pkt_sizes if len(pkt_sizes) == 16 else None
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# TESTS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-class TestSQTTMatchesBinary(unittest.TestCase):
-  def test_bit_counts_match_layout3(self): self._test_bit_counts(3)
-  def test_bit_counts_match_layout4(self): self._test_bit_counts(4)
-  def test_encodings_match_layout3(self): self._test_encodings(3)
-  def test_encodings_match_layout4(self): self._test_encodings(4)
-  def test_delta_fields_match_layout3(self): self._test_delta_fields(3)
-  def test_delta_fields_match_layout4(self): self._test_delta_fields(4)
-
-  def test_cdna_packet_sizes(self):
-    """Extract and verify CDNA pkt_fmt -> size mapping from rocprof's hash table."""
-    if not (EXAMPLES_DIR / "gfx950").exists(): self.skipTest("no CDNA examples")
-    if not (pkt_sizes := extract_cdna_packet_sizes()): self.skipTest("rocprof-trace-decoder not installed")
-    for pkt_fmt, size in CDNA_PKT_SIZES.items():
-      with self.subTest(pkt_fmt=pkt_fmt): self.assertEqual(pkt_sizes.get(pkt_fmt), size)
-
-  def test_cdna_packet_definitions(self):
-    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_CDNA
-    for pkt_fmt, pkt_cls in PACKET_TYPES_CDNA.items():
-      with self.subTest(packet=pkt_cls.__name__):
-        self.assertEqual(pkt_cls.encoding.default, pkt_fmt)
-        self.assertEqual(CDNA_PKT_SIZES[pkt_fmt] * 2, pkt_cls._size_nibbles)  # type: ignore[attr-defined]
-
-  def _test_bit_counts(self, layout: int):
-    if not (tables := extract_bit_tables()): self.skipTest("rocprof-trace-decoder not installed")
-    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4
-    for type_id, pkt_cls in {3: PACKET_TYPES_RDNA3, 4: PACKET_TYPES_RDNA4}[layout].items():
-      with self.subTest(packet=pkt_cls.__name__):
-        self.assertEqual(pkt_cls._size_nibbles * 4, tables[layout - 2][type_id])  # type: ignore[attr-defined]
-
-  def _test_encodings(self, layout: int):
-    if not (encodings := extract_packet_encodings()): self.skipTest("rocprof-trace-decoder not installed")
-    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4
-    for type_id, pkt_cls in {3: PACKET_TYPES_RDNA3, 4: PACKET_TYPES_RDNA4}[layout].items():
-      with self.subTest(packet=pkt_cls.__name__):
-        self.assertEqual((pkt_cls.encoding.mask, pkt_cls.encoding.default), encodings[layout - 2][type_id])
-
-  def _test_delta_fields(self, layout: int):
-    if not (deltas := extract_delta_fields()): self.skipTest("rocprof-trace-decoder not installed")
-    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4
-    for type_id, pkt_cls in {3: PACKET_TYPES_RDNA3, 4: PACKET_TYPES_RDNA4}[layout].items():
-      if type_id not in deltas[layout - 2]: continue
-      delta = getattr(pkt_cls, 'delta', None)
-      actual = (0, 0) if delta is None else (delta.lo, delta.hi + 1)
-      with self.subTest(packet=pkt_cls.__name__): self.assertEqual(actual, deltas[layout - 2][type_id])
-
-if __name__ == "__main__":
-  tables = extract_bit_tables()
-  encodings = extract_packet_encodings()
-  deltas = extract_delta_fields()
-
-  TYPE_NAMES = {1: 'VALUINST', 2: 'VMEMEXEC', 3: 'ALUEXEC', 4: 'IMMEDIATE', 5: 'IMMEDIATE_MASK', 6: 'WAVERDY',
-    7: 'TS_DELTA_S8_W3', 8: 'WAVEEND', 9: 'WAVESTART', 10: 'TS_DELTA_S5_W2', 11: 'WAVEALLOC', 12: 'TS_DELTA_S5_W3',
-    13: 'PERF', 14: 'UTILCTR', 15: 'TS_DELTA_SHORT', 16: 'NOP', 17: 'TS_WAVE_STATE', 18: 'EVENT', 19: 'EVENT_BIG',
-    20: 'REG', 21: 'SNAPSHOT', 22: 'TS_DELTA_OR_MARK', 23: 'LAYOUT_HEADER', 24: 'INST', 25: 'UNK_25'}
-
-  print("L2:", tables[0], "\nL3:", tables[1], "\nL4:", tables[2])
-  if encodings and tables:
-    print(f"\n{'TypeID':>6} {'Name':>18} {'L2 enc':>12} {'L3 enc':>12} {'L4 enc':>12}"
-          f" {'L2':>4} {'L3':>4} {'L4':>4} {'L2 delta':>12} {'L3 delta':>12} {'L4 delta':>12}")
-    print("-" * 140)
-    for type_id in sorted(set(encodings[0]) | set(encodings[1]) | set(encodings[2])):
-      name = TYPE_NAMES.get(type_id, f'UNK_{type_id}')
-      bits = [tables[i][type_id] if type_id < len(tables[i]) else 0 for i in range(3)]
-      enc_strs = [f"0x{encodings[i][type_id][0]:02x}/0x{encodings[i][type_id][1]:02x}" if type_id in encodings[i] else "-" for i in range(3)]
-      delta_strs = [f"[{d[1]-1}:{d[0]}]" if (d := deltas[i].get(type_id, (0, 0)))[1] > d[0] else "-" for i in range(3)]
-      print(f"{type_id:6d} {name:>18} {enc_strs[0]:>12} {enc_strs[1]:>12} {enc_strs[2]:>12}"
-            f" {bits[0]:4d} {bits[1]:4d} {bits[2]:4d} {delta_strs[0]:>12} {delta_strs[1]:>12} {delta_strs[2]:>12}")
-
-  cdna = extract_cdna_packet_sizes()
-  if cdna: print(f"\nCDNA packet sizes: {cdna}")
-
-  unittest.main()

test/amd/test_sqttmap.py

@@ -2,9 +2,10 @@
 import unittest, pickle
 from typing import Iterator
 from pathlib import Path
-from tinygrad.helpers import DEBUG
+from tinygrad.helpers import DEBUG, getenv, temp
 from tinygrad.renderer.amd.sqtt import print_packets, map_insts
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
+from tinygrad.viz.serve import sqtt_timeline
 from test.amd.disasm import disasm
 
 import tinygrad
@@ -24,7 +25,7 @@ def rocprof_inst_traces_match(sqtt, prg, target):
 
   passed_insts = 0
   for pkt, info in map_insts(sqtt.blob, prg.lib, target):
-    if DEBUG >= 2: print_packets([pkt])
+    if DEBUG >= 2: print_packets([(pkt, info)])
     if info is None: continue
     if DEBUG >= 2: print(f"{' '*29}{disasm(info.inst)}")
     rocprof_inst = next(rwaves_iter[info.wave][0])
@@ -53,7 +54,7 @@ class TestSQTTMapBase(unittest.TestCase):
   def setUpClass(cls):
     if cls is TestSQTTMapBase: raise unittest.SkipTest("base class")
     cls.examples = {}
-    for pkl_path in sorted((EXAMPLES_DIR/cls.target).glob("*.pkl")):
+    for pkl_path in ([Path(temp("profile.pkl", append_user=True))] if getenv("LOAD_PROFILE") else sorted((EXAMPLES_DIR/cls.target).glob("*.pkl"))):
       with open(pkl_path, "rb") as f:
         data = pickle.load(f)
       sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
@@ -63,6 +64,8 @@ class TestSQTTMapBase(unittest.TestCase):
 
   def test_rocprof_inst_traces_match(self):
     for name, (events, kern_events, target) in self.examples.items():
+      if "sync" in name and self.target.startswith("gfx12"):
+        self.skipTest("our timestamps are off by a few cycles because rocprof patches timestamps for rdna4 barriers")
       for event in events:
         if not event.itrace: continue
         if event.kern not in kern_events: continue
@@ -70,9 +73,47 @@ class TestSQTTMapBase(unittest.TestCase):
           passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target)
           if n_waves: print(f"{name}: passed for {passed_insts} instructions across {n_waves} waves scheduled on {n_units} wave units")
 
+  def test_sqtt_timeline(self):
+    for name, (events, kern_events, target) in self.examples.items():
+      for event in events:
+        if (p:=kern_events.get(event.kern)) is None: continue
+        with self.subTest(example=name, kern=event.kern):
+          # skip if there's no SQTT frequency data
+          if not (timeline:=sqtt_timeline(event.blob, p.lib, target)): continue
+          if not (frequency:=[e.key for e in timeline if type(e).__name__ == "ProfilePointEvent" and e.name == "freq_hz"]): continue
+          mean = sum(frequency) / len(frequency)
+          variance = sum((v - mean) ** 2 for v in frequency) / len(frequency)
+          self.assertGreater(mean, 0)
+          if DEBUG >= 2: print(f"{name:20s} SE:{event.se} {mean/1e9:.2f} GHz mean, {variance/1e18:.2f} GHz^2 variance")
+          events = [e for e in timeline if type(e).__name__ == "ProfileRangeEvent"]
+          insts, execs = 0, 0
+          for e in events:
+            if "EXEC" in e.device:
+              if "ALT" not in e.name.display_name: execs += 1
+            elif "WAVE" in e.device:
+              # sopk/immediates don't get ALU/MEM EXEC
+              if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE", "BARRIER", "BARRIER_SIGNAL"}: insts += 1
+            else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
+          self.assertEqual(execs, insts)
+
+  def test_wave_sync(self):
+    for name, (events, kern_events, target) in self.examples.items():
+      for event in events:
+        wave_barriers = {}
+        for e in sqtt_timeline(event.blob, kern_events[event.kern].lib, target):
+          if type(e).__name__ == "ProfileRangeEvent" and e.name.display_name == "BARRIER": wave_barriers.setdefault(e.device, []).append(e)
+        if not wave_barriers: continue
+        for row, events in wave_barriers.items():
+          for e in events:
+            assert e.en-e.st > 1, f"all barriers must have a duration greater than 1, got {e}"
+
 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"
 
 class TestSQTTMapRDNA4(TestSQTTMapBase): target = "gfx1200"
 
+class TestSQTTMapCDNA(TestSQTTMapBase):
+  target = "gfx950"
+  def test_rocprof_inst_traces_match(self): self.skipTest("requires timestamp patching to match rocprof, currently it's off by a few cycles")
+
 if __name__ == "__main__":
   unittest.main()

test/backend/test_arange.py

@@ -8,6 +8,7 @@ from tinygrad.engine.schedule import ExecItem
 from tinygrad.uop.ops import Ops
 from tinygrad.renderer import Estimates
 from tinygrad.renderer.ptx import PTXRenderer
+from test.helpers import needs_second_gpu
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, tensor, desired):
@@ -188,6 +189,33 @@ class TestIndexing(unittest.TestCase):
     for i in idx.flatten().numpy(): expected_grad[i] += 2
     np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
 
+  @needs_second_gpu
+  @unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
+  @Context(USE_ATOMICS=1, SPEC=1)
+  def test_embedding_backward_vocab_sharded(self):
+    from tinygrad.renderer.cstyle import CStyleLanguage
+    if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
+    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
+    vocab_size, embed_size = 1000, 128
+    bs, seqlen = 4, 256
+    idx = Tensor.randint(bs, seqlen, high=vocab_size)
+    emb = nn.Embedding(vocab_size, embed_size)
+    emb.weight = Tensor.ones(vocab_size, embed_size, requires_grad=True)
+    gt = Tensor.zeros(bs, seqlen, embed_size)
+    Tensor.realize(idx, emb.weight, gt)
+    # compute expected grad on single device
+    expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
+    for i in idx.flatten().numpy(): expected_grad[i] += 2
+    # now shard the embedding weight on vocab axis and recompute
+    emb.weight = Tensor.ones(vocab_size, embed_size, requires_grad=True)
+    emb.weight.shard_(devices, axis=0)
+    idx = idx.shard(devices, axis=None)
+    gt = gt.shard(devices, axis=None)
+    Tensor.realize(idx, emb.weight, gt)
+    loss = (emb(idx)-gt).square().sum()
+    loss.backward()
+    np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
+
   @unittest.skipUnless(Device.DEFAULT == "AMD" or (Device.DEFAULT == "NULL" and EMULATE.value.startswith("AMD")), "tests AMD bf16 cast overhead")
   def base_test_llama_8b_rope_backward(self, dtype):
     from extra.models.llama import precompute_freqs_cis, apply_rotary_emb

test/backend/test_asm_gemm.py

@@ -2,7 +2,7 @@ import unittest
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv
-from extra.gemm.asm.cdna.gemm import asm_gemm
+from extra.gemm.cdna_asm_gemm import asm_gemm
 from test.helpers import needs_second_gpu
 
 # On non CDNA4 it will only validate the Tensor.custom_kernel integration
@@ -47,6 +47,18 @@ def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:i
 def verify_asm_gemm_k_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=8) -> None:
   run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=1, b_shard=0, gpus=gpus)
 
+def verify_asm_gemm_n_sharded(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
+
+def verify_asm_gemm_m_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
+
+def verify_asm_gemm_n_sharded_2d(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
+
+def verify_asm_gemm_k_sharded_3d(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=2, b_shard=0, gpus=gpus)
+
 # 128x smaller than usual
 # uses the UOp GEMM, runs on non CDNA4 and CI
 @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@@ -60,6 +72,14 @@ class TestGemm(unittest.TestCase):
   def test_gemm_multi(self): verify_asm_gemm(2, 64, 32, 32, gpus=2)
   @needs_second_gpu
   def test_gemm_k_sharded(self): verify_asm_gemm_k_sharded(64, 64, 2*64, gpus=2)
+  @needs_second_gpu
+  def test_gemm_m_sharded(self): verify_asm_gemm_m_sharded(2*64, 64, 32, gpus=2)
+  @needs_second_gpu
+  def test_gemm_n_sharded(self): verify_asm_gemm_n_sharded(1, 64, 64, 32, gpus=2)
+  @needs_second_gpu
+  def test_gemm_n_sharded_2d(self): verify_asm_gemm_n_sharded_2d(64, 2*64, 32, gpus=2)
+  @needs_second_gpu
+  def test_gemm_k_sharded_3d(self): verify_asm_gemm_k_sharded_3d(1, 64, 32, 2*64, gpus=2)
 
 # uses the Asm GEMM on CDNA4 only for speed reasons
 class TestGemmLarge(unittest.TestCase):
@@ -67,6 +87,9 @@ class TestGemmLarge(unittest.TestCase):
     if not is_cdna4():
       self.skipTest("very slow on non mi350x")
 
+  @Context(ASM_GEMM=1)
+  def test_empty(self): (Tensor.empty(N:=getenv("N", 4096), N, dtype=dtypes.half)@Tensor.empty(N, N, dtype=dtypes.half)).realize()
+
   def test_tiny(self): verify_asm_gemm(1, 256, 256, 64)
   def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 4096), N, N, dtype=dtypes.half)
   def test_gemm(self): verify_asm_gemm(1, 8192, 4096, 14336)
@@ -101,6 +124,20 @@ class TestGemmLarge(unittest.TestCase):
       verify_asm_gemm(3, 256, 256, 256)
   def test_gemm_previously_unsupported(self): verify_asm_gemm(8, 1024, 1024, 4096, gpus=8)
 
+  # M-sharded 2D
+  def test_m_sharded_1(self): verify_asm_gemm_m_sharded(8*8192, 4096, 4096, dtype=dtypes.bfloat16, gpus=8)
+  def test_m_sharded_2(self): verify_asm_gemm_m_sharded(8*4096, 14336, 4096, dtype=dtypes.bfloat16, gpus=8)
+
+  # N-sharded 2D
+  def test_n_sharded_2d_1(self): verify_asm_gemm_n_sharded_2d(8192, 8*4096, 4096, dtype=dtypes.bfloat16, gpus=8)
+  def test_n_sharded_2d_2(self): verify_asm_gemm_n_sharded_2d(4096, 8*14336, 4096, dtype=dtypes.bfloat16, gpus=8)
+
+  # tensor parallel shapes (Llama 8B, MP=8)
+  def test_tp_n_sharded_wq(self): verify_asm_gemm_n_sharded(1, 8192, 4096, 4096, dtype=dtypes.bfloat16, gpus=8)
+  def test_tp_n_sharded_w1(self): verify_asm_gemm_n_sharded(1, 8192, 14336, 4096, dtype=dtypes.bfloat16, gpus=8)
+  def test_tp_k_sharded_wo(self): verify_asm_gemm_k_sharded_3d(1, 8192, 4096, 4096, dtype=dtypes.bfloat16, gpus=8)
+  def test_tp_k_sharded_w2(self): verify_asm_gemm_k_sharded_3d(1, 8192, 4096, 14336, dtype=dtypes.bfloat16, gpus=8)
+
   # more shapes: vary M, N, K independently
   def test_shape_small_square(self): verify_asm_gemm(1, 256, 256, 256)
   def test_shape_small_rect_m(self): verify_asm_gemm(1, 512, 256, 256)
@@ -123,7 +160,7 @@ class TestGemmLarge(unittest.TestCase):
 
 class TestMagicGu(unittest.TestCase):
   def test_magicgu_matches_old(self):
-    from extra.gemm.asm.cdna.asm import _magicgu_mulhi, TILE_M, TILE_N, TILE_K
+    from extra.gemm.cdna_asm_gemm import _magicgu_mulhi, TILE_M, TILE_N, TILE_K
     old_iters_args = {64: (67108864, 0), 128: (33554432, 0), 224: (613566757, 2147483656)}
     old_gemm_shapes = [
       (8192, 4096, 4096), (8192, 14336, 4096), (8192, 4096, 14336),

test/backend/test_const_folding.py

@@ -27,6 +27,11 @@ class TestMovedConstFolding(unittest.TestCase):
   def test_add_padded_one(self):
     _check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))
 
+  def test_copy_padded_const(self):
+    schedule = Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").schedule()
+    assert not any(si.ast.op is Ops.COPY for si in schedule), "const copy should be folded"
+    np.testing.assert_equal(Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").numpy(), [0, 1, 1, 1, 1, 0])
+
   def test_cast_padded(self):
     # NOTE: it's always 1 kernel when calling .numpy, limitation of _check_ast_count
     if is_dtype_supported(dtypes.int16):

test/backend/test_dtype.py

@@ -10,7 +10,7 @@ from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Context, Device, Tensor, dtypes
 from hypothesis import given, settings, strategies as strat
 from test.helpers import rand_for_dtype
-from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX
+from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")
 
@@ -101,14 +101,14 @@ class TestDType(unittest.TestCase):
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "skip for now")
   def test_uint_overflow(self):
     if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned")
-    v = dtypes.max(self.DTYPE)
+    v = self.DTYPE.max
     _test_to_np(Tensor(v, dtype=self.DTYPE)+2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))+2)
     _test_to_np(Tensor(v, dtype=self.DTYPE)*2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))*2)
 
   def test_dtypes_DTYPES_DICT(self):
     self.assertIn("float", DTYPES_DICT)
     self.assertIn("float32", DTYPES_DICT)
-    self.assertEqual(len(DTYPES_DICT), 26)
+    self.assertEqual(len(DTYPES_DICT), 28)
     self.assertTrue(all(isinstance(value, DType) for value in DTYPES_DICT.values()))
     self.assertTrue(all(issubclass(_to_np_dtype(value), np.generic) for value in DTYPES_DICT.values() if _to_np_dtype(value) is not None))
 
@@ -143,6 +143,8 @@ def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
 class TestFp8s(unittest.TestCase):
   def test_fp8e4m3_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3).dtype == dtypes.fp8e4m3
   def test_fp8e5m2_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2).dtype == dtypes.fp8e5m2
+  def test_fp8e4m3fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3fnuz).dtype == dtypes.fp8e4m3fnuz
+  def test_fp8e5m2fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2fnuz).dtype == dtypes.fp8e5m2fnuz
 
 class TestFp8sConversions(unittest.TestCase):
   @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX))
@@ -169,6 +171,30 @@ class TestFp8sConversions(unittest.TestCase):
   def test_fp8e5m2_to_float(self, x):
     np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2).float().item())
 
+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3FNUZ_MAX, max_value=FP8E4M3FNUZ_MAX))
+  def test_float_to_fp8e4m3fnuz(self, x):
+    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
+
+  def test_float_to_fp8e4m3fnuz_extreme_values(self):
+    for x in [FP8E4M3FNUZ_MAX, FP8E4M3FNUZ_MAX*1.01, -FP8E4M3FNUZ_MAX, -FP8E4M3FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
+      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
+
+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2FNUZ_MAX, max_value=FP8E5M2FNUZ_MAX))
+  def test_float_to_fp8e5m2fnuz(self, x):
+    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
+
+  def test_float_to_fp8e5m2fnuz_extreme_values(self):
+    for x in [FP8E5M2FNUZ_MAX, FP8E5M2FNUZ_MAX*1.01, -FP8E5M2FNUZ_MAX, -FP8E5M2FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
+      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
+
+  @given(strat.integers(min_value=0, max_value=255))
+  def test_fp8e4m3fnuz_to_float(self, x):
+    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fnuz).float().item())
+
+  @given(strat.integers(min_value=0, max_value=255))
+  def test_fp8e5m2fnuz_to_float(self, x):
+    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2fnuz).float().item())
+
 @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
 class TestBFloat16(unittest.TestCase):
   def test_bf16_creation_numpy(self):
@@ -490,4 +516,3 @@ class TestOpsBFloat16(unittest.TestCase):
 
 if __name__ == '__main__':
   unittest.main()
-

test/backend/test_dtype_alu.py

@@ -52,6 +52,8 @@ class ht:
 ht.bfloat16 = ht.uint16.filter(lambda x: ((x >> 7) & 0xFF) != 0)  # filter subnormal bfloat16
 ht.fp8e4m3 = ht.uint8
 ht.fp8e5m2 = ht.uint8
+ht.fp8e4m3fnuz = ht.uint8
+ht.fp8e5m2fnuz = ht.uint8
 
 def universal_test(a, b, dtype, op):
   if not isinstance(op, tuple): op = (op, op)
@@ -67,7 +69,8 @@ def universal_test(a, b, dtype, op):
     if not is_dtype_supported(dtype) or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero
       fe, fm = dtypes.finfo(dtype)
       atol, rtol = 2 ** (2 - (1 << (fe - 1))), 2 ** (-fm)
-    else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype, (1e-10, 1e-7))
+    else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1),
+                        dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz:(5e-1, 5e-1)}.get(dtype, (1e-10, 1e-7))
     np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
   else: np.testing.assert_equal(tensor_value, numpy_value)
 
@@ -87,7 +90,8 @@ def universal_test_unary(a, dtype, op):
   else: tensor_value, numpy_value = op[0](ta).numpy(), op[1](ta.numpy())
   if dtype in dtypes.floats:
     atol, rtol = { dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2),
-      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1)}.get(dtype, (1e-6, 1e-5))
+      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1),
+      dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz: (5e-1, 5e-1)}.get(dtype, (1e-6, 1e-5))
     np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
   else: np.testing.assert_equal(tensor_value, numpy_value)
 
@@ -155,6 +159,26 @@ class TestDTypeALU(unittest.TestCase):
   def test_emulated_fp8e5m2(self, a, b, op):
     universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)
 
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}")
+  @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
+  def test_fp8e4m3fnuz(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}")
+  @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
+  def test_fp8e5m2fnuz(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
+
+  @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="fp8e4m3fnuz")
+  def test_emulated_fp8e4m3fnuz(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
+
+  @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="fp8e5m2fnuz")
+  def test_emulated_fp8e5m2fnuz(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
+
   @given(ht.float32, strat.sampled_from(unary_operations))
   def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)
 
@@ -198,6 +222,30 @@ class TestDTypeALU(unittest.TestCase):
     if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
     universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)
 
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}")
+  @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
+  def test_fp8e4m3fnuz_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}")
+  @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
+  def test_fp8e5m2fnuz_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
+
+  @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="fp8e4m3fnuz")
+  def test_emulated_fp8e4m3fnuz_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
+
+  @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="fp8e5m2fnuz")
+  def test_emulated_fp8e5m2fnuz_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
+
   @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
   def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)
 
@@ -318,7 +366,7 @@ class TestDTypeALU(unittest.TestCase):
 
   @unittest.expectedFailure
   def test_unsafe_cast_float_to_int_failure(self):
-    val = float(dtypes.max(dtypes.int32) - 1)
+    val = float(dtypes.int32.max - 1)
     t1 = Tensor([val], dtype=dtypes.float32).cast(dtypes.int32)
     t2 = Tensor(val, dtype=dtypes.float32).cast(dtypes.int32)
     np.testing.assert_equal(t1.item(), t2.item())

test/backend/test_graph.py

@@ -15,7 +15,7 @@ from test.helpers import needs_second_gpu
 np.random.seed(1337)
 Tensor.manual_seed(1337)
 BUF_SIZE = 4096
-RUN_CNT = 4
+RUN_CNT = 5
 
 cached_prgs = {}
 def helper_exec_op(device, outbuf, inbufs):
@@ -47,6 +47,17 @@ def helper_create_offset_rawbuffer(base, offset=0):
   x = Buffer(base.device, base.size-offset, base.dtype, base=base, offset=offset)
   return x.ensure_allocated()
 
+def helper_alloc_rawbuffer_sized(device, size, fill=False):
+  rawbuf = Buffer(device, size, dtypes.int).ensure_allocated()
+  if fill:
+    with Context(DEBUG=0):
+      data = np.random.randint(-10000, 10000, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
+      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_memoryview())
+  return rawbuf
+
+def helper_make_view(base, offset_elems, size_elems):
+  return Buffer(base.device, size_elems, base.dtype, base=base, offset=offset_elems * base.dtype.itemsize).ensure_allocated()
+
 def helper_run_jit(jis, bufs, out_buffers):
   for rawbuf in out_buffers:
     mv = memoryview(bytearray(rawbuf.size * rawbuf.dtype.itemsize))
@@ -80,6 +91,14 @@ def helper_test_graphs(graph_impl, graphs, runs=RUN_CNT):
 
 @unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
 class TestGraph(unittest.TestCase):
+  def skip_if_no_offset(self):
+    if not hasattr(Device[Device.DEFAULT].allocator, "_offset"): self.skipTest("device does not support _offset")
+
+  def skip_if_not_multigraph(self):
+    graph = g.func if isinstance(g:=(d:=Device[Device.DEFAULT]).graph, functools.partial) else g
+    if not issubclass(graph, MultiGraphRunner): self.skipTest("graph is not supported (not MultiGraphRunner)")
+    if not hasattr(d.allocator, '_transfer') or not d.allocator.supports_transfer: self.skipTest("device is not supported (no transfers)")
+
   def test_order_2_writes_to_same_buf(self):
     d0 = Device.DEFAULT
     b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(5)]
@@ -110,11 +129,6 @@ class TestGraph(unittest.TestCase):
 
     helper_test_graphs(Device[d0].graph, graphs)
 
-  def skip_if_not_multigraph(self):
-    graph = g.func if isinstance(g:=(d:=Device[Device.DEFAULT]).graph, functools.partial) else g
-    if not issubclass(graph, MultiGraphRunner): self.skipTest("graph is not supported (not MultiGraphRunner)")
-    if not hasattr(d.allocator, '_transfer') or not d.allocator.supports_transfer: self.skipTest("device is not supported (no transfers)")
-
   def test_order_copy_writed(self):
     self.skip_if_not_multigraph()
 
@@ -265,5 +279,58 @@ class TestGraph(unittest.TestCase):
 
     helper_test_graphs(Device[d0].graph, graphs)
 
+  def test_partial_write_preserves_write_dep(self):
+    self.skip_if_not_multigraph()
+    self.skip_if_no_offset()
+    d0 = Device.DEFAULT
+
+    base = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
+    copy_src_full = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
+    copy_src_lo = helper_alloc_rawbuffer(d0, fill=True)
+    v_lo = helper_make_view(base, 0, BUF_SIZE)
+    v_hi = helper_make_view(base, BUF_SIZE, BUF_SIZE)
+    a, c = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(2)]
+
+    graphs = [
+      [helper_copy_op(d0, base, copy_src_full), helper_copy_op(d0, v_lo, copy_src_lo), helper_exec_op(d0, c, [v_hi, a])]
+    ]
+    helper_test_graphs(Device[d0].graph, graphs)
+
+  def test_partial_write_preserves_read_dep(self):
+    self.skip_if_not_multigraph()
+    self.skip_if_no_offset()
+    d0 = Device.DEFAULT
+
+    base = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
+    copy_dst = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
+    copy_src_lo = helper_alloc_rawbuffer(d0, fill=True)
+    v_lo = helper_make_view(base, 0, BUF_SIZE)
+    v_hi = helper_make_view(base, BUF_SIZE, BUF_SIZE)
+    a, b = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(2)]
+
+    graphs = [
+      [helper_copy_op(d0, copy_dst, base), helper_copy_op(d0, v_lo, copy_src_lo), helper_exec_op(d0, v_hi, [a, b])]
+    ]
+    helper_test_graphs(Device[d0].graph, graphs)
+
+  def test_middle_write_splits_write_dep(self):
+    self.skip_if_not_multigraph()
+    self.skip_if_no_offset()
+    d0 = Device.DEFAULT
+
+    base = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 3, fill=True)
+    copy_src_full = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 3, fill=True)
+    copy_src_mid = helper_alloc_rawbuffer(d0, fill=True)
+    v_lo = helper_make_view(base, 0, BUF_SIZE)
+    v_mid = helper_make_view(base, BUF_SIZE, BUF_SIZE)
+    v_hi = helper_make_view(base, BUF_SIZE * 2, BUF_SIZE)
+    a, c, e = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(3)]
+
+    graphs = [
+      [helper_copy_op(d0, base, copy_src_full), helper_copy_op(d0, v_mid, copy_src_mid),
+       helper_exec_op(d0, c, [v_lo, a]), helper_exec_op(d0, e, [v_hi, a])]
+    ]
+    helper_test_graphs(Device[d0].graph, graphs)
+
 if __name__ == '__main__':
   unittest.main()

test/backend/test_multitensor.py

@@ -228,17 +228,17 @@ 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_multiple_to_single_device_naive(self):
-    with Context(RING=0):
-      t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize()
-    self.assertEqual(t.device, Device.DEFAULT)
-    np.testing.assert_equal(t.numpy(), np.arange(32))
-
-  def test_multiple_to_single_device_ring(self):
-    with Context(RING=2):
-      t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize()
-    self.assertEqual(t.device, Device.DEFAULT)
-    np.testing.assert_equal(t.numpy(), np.arange(32))
+  def test_multiple_to_single_device(self):
+    kernel_counts = {}
+    for ring in (0, 2):
+      GlobalCounters.reset()
+      with Context(RING=ring, SCACHE=0):
+        t = Tensor.arange(32).contiguous().shard(devices_4, 0).to(Device.DEFAULT)
+        t.realize()
+      kernel_counts[ring] = GlobalCounters.kernel_count
+      self.assertEqual(t.device, Device.DEFAULT)
+      np.testing.assert_equal(t.numpy(), np.arange(32))
+    self.assertNotEqual(kernel_counts[0], kernel_counts[2])
 
   def test_allreduce_all2all(self):
     with Context(ALL2ALL=2):

test/backend/test_ops.py

@@ -6,6 +6,7 @@ from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, AMD_LL
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
+from tinygrad.renderer.cstyle import QCOMCLRenderer
 from tinygrad.renderer.nir import NIRRenderer
 
 TINY_BACKEND = getenv("TINY_BACKEND")
@@ -436,7 +437,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,35), (45,35), (45,35)], lambda x,y,z: x.lerp(y,z))
     helper_test_op(None, lambda x,y,z: x.lerp(y,z), vals=[[1.,2.,3.], [4.,5.,6.], 0.5])
 
-  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_tril(self):
     helper_test_op([(3,3)], lambda x: x.tril())
     helper_test_op([(3,3)], lambda x: x.tril(1))
@@ -454,7 +455,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(5,3,3)], lambda x: x.tril(1))
     helper_test_op(None, lambda x: x.tril(), vals=[[[True] * 3] * 3], forward_only=True)
 
-  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_triu(self):
     helper_test_op([(3,3)], lambda x: x.triu())
     helper_test_op([(3,3)], lambda x: x.triu(1))
@@ -478,9 +479,9 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[1., 0., 3., -4.], 3.])
     helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[1., 0., 3., -4.], [-1., -2., 3., 0.]])
     helper_test_op(None, torch.maximum, Tensor.maximum,
-                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.max(dtypes.int)], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.max], forward_only=True)
     helper_test_op(None, torch.maximum, Tensor.maximum,
-                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.min(dtypes.int)], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.min], forward_only=True)
     helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[True, False, False], True], forward_only=True)
     helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[True, False, False], [True, True, False]], forward_only=True)
 
@@ -495,9 +496,9 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[1., 0., 3., -4.], 3.])
     helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[1., 0., 3., -4.], [-1., -2., 3., 0.]])
     helper_test_op(None, torch.minimum, Tensor.minimum,
-                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.max(dtypes.int)], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.max], forward_only=True)
     helper_test_op(None, torch.minimum, Tensor.minimum,
-                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.min(dtypes.int)], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.min], forward_only=True)
     helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[True, False, False], True], forward_only=True)
     helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[True, False, False], [True, True, False]], forward_only=True)
 
@@ -765,6 +766,7 @@ class TestOps(unittest.TestCase):
 
     self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_xor(y), expected=RuntimeError)
 
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_and(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)
@@ -782,6 +784,7 @@ class TestOps(unittest.TestCase):
 
     self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_and(y), expected=RuntimeError)
 
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_or(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)
@@ -1170,6 +1173,7 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: x.type(torch.int32).argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[False, True]])
     helper_test_op(None, lambda x: x.type(torch.int32).argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[True, False]])
 
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_argmin(self):
     # check if it returns the first index for multiple occurrences
     helper_test_op(None, lambda x: x.argmin().type(torch.int32), lambda x: x.argmin(), forward_only=True, vals=[[2, 2]])
@@ -1475,6 +1479,7 @@ class TestOps(unittest.TestCase):
   def test_prod_dtype_arg(self):
     with self.assertRaises(AttributeError): Tensor([1.0, 2.0]).prod(dtype="")
 
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_min(self):
     helper_test_op([(3,3)], lambda x: x.min())
     helper_test_op([(45,3)], lambda x: x.min())
@@ -1503,7 +1508,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(3,3)], lambda x: torch.full_like(x, 2).prod(), lambda x: (x.full_like(2)).prod(), forward_only=True)
     helper_test_op([(3,3)], lambda x: torch.full_like(x, 2).max(), lambda x: (x.full_like(2)).max(), forward_only=True)
 
-  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
   def test_any(self):
     helper_test_op([(3,4,5,6)], lambda x: x.any(), forward_only=True)
     helper_test_op(None, lambda x: x.any(), vals=[[True, True]], forward_only=True)
@@ -1515,7 +1519,7 @@ class TestOps(unittest.TestCase):
   def test_any_zero_axis(self):
     helper_test_op([(1,0,3,0,5)], lambda x: x.any(axis=(1,3)), forward_only=True)
 
-  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_all(self):
     helper_test_op([(3,4,5,6)], lambda x: x.all(), forward_only=True)
     helper_test_op(None, lambda x: x.all(), vals=[[True, True]], forward_only=True)
@@ -1665,6 +1669,15 @@ class TestOps(unittest.TestCase):
     helper_test_op([(10,10,10)], lambda x: x.log_softmax(1), atol=1e-7, grad_atol=1e-7)
     helper_test_op([(10,10,10)], lambda x: x.log_softmax(2), atol=1e-7, grad_atol=1e-7)
 
+  def test_normalize(self):
+    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x), lambda x: x.normalize(), atol=1e-7, grad_atol=1e-7)
+    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, dim=0), lambda x: x.normalize(dim=0), atol=1e-7, grad_atol=1e-7)
+    helper_test_op([(10,10,10)], lambda x: torch.nn.functional.normalize(x, dim=2), lambda x: x.normalize(dim=2), atol=1e-7, grad_atol=1e-7)
+    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=1), lambda x: x.normalize(p=1), atol=1e-7, grad_atol=1e-7)
+    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=3, dim=0), lambda x: x.normalize(p=3, dim=0), atol=1e-7, grad_atol=1e-7)
+    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=0), lambda x: x.normalize(p=0), atol=1e-7, grad_atol=1e-7)
+    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=-1), lambda x: x.normalize(p=-1), atol=1e-7, grad_atol=1e-7)
+
   def test_logsumexp(self):
     helper_test_op([(45,65)], lambda x: torch.logsumexp(x, dim=0), lambda x: x.logsumexp(0), atol=1e-7, grad_atol=1e-7)
     helper_test_op([(45,65)], lambda x: torch.logsumexp(x, dim=0, keepdim=True), lambda x: x.logsumexp(0, True), atol=1e-7, grad_atol=1e-7)
@@ -2880,6 +2893,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[...,c,:,e], lambda x: x[...,k,:,p])
 
   @slow_test
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_dim_collapse_int(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # dim collapse from int
@@ -2890,6 +2904,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[1,:,3:11:2,d,0:2], lambda x: x[1,:,3:11:2,o,0:2])
 
   @slow_test
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_dim_inject_none(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # dim injection from None
@@ -2924,6 +2939,7 @@ class TestOps(unittest.TestCase):
                             lambda x: x[Tensor([[0,1,-1],[-1,-2,0]]), Tensor([2,1,-1])])
 
   @slow_test
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_list_indices(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[((0,),)])
@@ -2935,6 +2951,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,(2,1,0),c,(-2,1,0),e], lambda x: x[i,(2,1,0),k,(-2,1,0),p])
 
   @slow_test
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_tuple_indices(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[(((0,),),)], lambda x: x[(((0,),),)])
@@ -3276,7 +3293,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(20,)], lambda x: (x>0.5).nonzero().int(), lambda x: (x>0.5).nonzero(), forward_only=True)
     helper_test_op([(10, 5, 3)], lambda x: (x>0.5).nonzero().int(), lambda x: (x>0.5).nonzero(), forward_only=True)
 
-  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
   def test_cast(self):
     helper_test_op([(3, 3)], lambda x: x.float())
     helper_test_op(None, lambda x: x.float(), vals=[[0, 1, 2, 3]], forward_only=True)

test/backend/test_quantize_onnx.py

@@ -204,13 +204,13 @@ class TestQuantizeOnnx(unittest.TestCase):
     W = Tensor(m2:=(np.random.uniform(0, 255, size=(N,N)).astype(wi))).realize()
     tg_dtype = dtypes.int8 if xi == np.int8 else dtypes.uint8
     out = (X.int().matmul(W.int())//1000)
-    if clip: out = out.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
+    if clip: out = out.clip(tg_dtype.min, tg_dtype.max)
     out = out.cast(tg_dtype)
     opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
     sexec(out, opts, replace_src, run_count=1)
     tout = out.numpy()
     mout = ((m1.astype(np.int32) @ m2.astype(np.int32)) // 1000)
-    if clip: mout = mout.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
+    if clip: mout = mout.clip(tg_dtype.min, tg_dtype.max)
     mout = mout.astype(xi)
     print(tout)
     print(mout)

test/backend/test_renderer_failures.py

@@ -62,7 +62,7 @@ class TestRendererFailures(unittest.TestCase):
 class TestCStyleFailures(unittest.TestCase):
   def test_inline_const_alu(self):
     # CPU doesn't use the max function
-    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.min(dtypes.int)+1))
+    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.int.min+1))
     self.assertEqual(ret[0], 1)
 
   def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True):

test/backend/test_schedule.py

@@ -795,6 +795,59 @@ class TestSchedule(unittest.TestCase):
       self.assertIsNotNone(out.uop.base.realized)
       self.assertIsInstance(out.uop.base.realized.dtype, ImageDType)
 
+  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
+  def test_image_dot_f16_fusion(self):
+    with Context(FLOAT16=1, OPENPILOT_HACKS=1):
+      def cnt():
+        x, y, z = Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float')
+        a = (x @ y).relu()
+        sched = ((a @ z).relu() + a).schedule()
+        for si in sched: si.lower()
+        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
+
+      with Context(IMAGE=1): cnt1 = cnt()
+      with Context(IMAGE=2): cnt2 = cnt()
+
+      self.assertEqual(cnt1, 5)
+      self.assertEqual(cnt2, 5)
+
+  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
+  def test_image_f16_residual_fusion(self):
+    with Context(FLOAT16=1, OPENPILOT_HACKS=1):
+      def cnt():
+        inp = Tensor.empty((512,), dtype='float')
+        b1, b2 = Tensor.empty((512, 1024), dtype='float'), Tensor.empty((1024, 512), dtype='float')
+        c1, c2 = Tensor.empty((1024,), dtype='float'), Tensor.empty((512,), dtype='float')
+        rb = (((((inp @ b1) + c1).relu() @ b2) + c2).relu() + inp).relu()
+        b16, c16 = Tensor.empty((512, 16), dtype='float'), Tensor.empty((16,), dtype='float')
+        b32, c32 = Tensor.empty((512, 32), dtype='float'), Tensor.empty((32,), dtype='float')
+        sched = Tensor.schedule((rb @ b16 + c16).relu(), (rb @ b32 + c32).relu())
+        for si in sched: si.lower()
+        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
+
+      with Context(IMAGE=1): cnt1 = cnt()
+      with Context(IMAGE=2): cnt2 = cnt()
+
+      self.assertEqual(cnt1, 9)
+      self.assertEqual(cnt2, 9)
+
+  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
+  @unittest.expectedFailure
+  def test_image_conv_fusion(self):
+    with Context(OPENPILOT_HACKS=1):
+      def cnt():
+        x, y, z = Tensor.empty((1, 4, 3, 3)), Tensor.empty((4, 1, 3, 3)), Tensor.empty((4, 1, 7, 7))
+        a = x.conv2d(y, Tensor.empty(4), groups=4, padding=1)
+        b = a.conv2d(z, groups=4, padding=3)
+        sched = (a + b).schedule()
+        for si in sched: si.lower()
+        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
+
+      with Context(IMAGE=1): cnt1 = cnt()
+      with Context(IMAGE=2): cnt2 = cnt()
+
+      self.assertEqual(cnt1, cnt2)
+
   def _test_fusion(self, shapes, f, cnt):
     with Context(DEBUG=0, TRACK_MATCH_STATS=0): args = [Tensor.randn(s).realize() for s in shapes]
     run_schedule(check_schedule(compare:=f(*args), cnt))

test/backend/test_setitem.py

@@ -205,6 +205,20 @@ class TestSetitem(unittest.TestCase):
     n[:, ind_1.numpy(), :, ind_2.numpy(), :] = v.numpy()
     np.testing.assert_equal(t.numpy(), n)
 
+  def test_setitem_tensor_int_indexing(self):
+    t = Tensor.zeros(4, 3, dtype=dtypes.int).contiguous()
+    t[Tensor([0, 2]), 0] = Tensor([99, 88], dtype=dtypes.int)
+    n = np.zeros((4, 3), dtype=np.int32)
+    n[[0, 2], 0] = [99, 88]
+    np.testing.assert_equal(t.numpy(), n)
+
+  def test_setitem_tensor_slice_indexing(self):
+    t = Tensor.zeros(4, 3, dtype=dtypes.int).contiguous()
+    t[Tensor([0, 2]), :2] = Tensor([[10, 20], [30, 40]], dtype=dtypes.int)
+    n = np.zeros((4, 3), dtype=np.int32)
+    n[[0, 2], :2] = [[10, 20], [30, 40]]
+    np.testing.assert_equal(t.numpy(), n)
+
   def test_setitem_2d_tensor_indexing(self):
     t = Tensor.zeros(2, dtype=dtypes.int).contiguous()
     index = Tensor([[0, 1], [1,0]])
@@ -237,8 +251,11 @@ class TestSetitem(unittest.TestCase):
     s1 = t.sum()
     t[3:].assign(2.0)
     s2 = t.sum()
-    # TODO: s0 and s1 see final buffer state, should be [0.0, 3.0, 9.0]
-    np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [9.0, 9.0, 9.0])
+    try:
+      np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [0.0, 3.0, 9.0])
+    except AssertionError:
+      # TODO: broken now, lazy sums all see final buffer state
+      np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [9.0, 9.0, 9.0])
 
     # eager version
     t = Tensor.zeros(6).contiguous().realize()
@@ -259,8 +276,11 @@ class TestSetitem(unittest.TestCase):
     a.assign(new_a)
     b.assign(new_b)
     np.testing.assert_allclose(a.numpy(), [4, 6, 8, 10])
-    # TODO: new_b sees mutated a, should be [0, 2, 4, 6]
-    np.testing.assert_allclose(b.numpy(), [8, 12, 16, 20])
+    try:
+      np.testing.assert_allclose(b.numpy(), [0, 2, 4, 6])
+    except AssertionError:
+      # TODO: broken now, new_b sees mutated a
+      np.testing.assert_allclose(b.numpy(), [8, 12, 16, 20])
 
     # eager version
     a = Tensor.arange(4, dtype=dtypes.float).contiguous().realize()
@@ -279,17 +299,43 @@ class TestWithGrad(unittest.TestCase):
     x = Tensor.rand(8)
     z[:3] = x
 
-  def test_set_into_requires_grad(self):
-    z = Tensor.rand(8, 8, requires_grad=True)
-    x = Tensor.rand(8)
-    with self.assertRaises(NotImplementedError):
-      z[:3] = x
-
   def test_set_with_requires_grad(self):
-    z = Tensor.rand(8, 8)
-    x = Tensor.rand(8, requires_grad=True)
-    with self.assertRaises(NotImplementedError):
-      z[:3] = x
+    z = Tensor.ones(8, 8)
+    x = Tensor.rand(8, 8, requires_grad=True)
+    z[:] = x
+    z.sum().backward()
+    np.testing.assert_allclose(x.grad.numpy(), np.ones((8, 8)))
+
+  def test_set_nonleaf_requires_grad(self):
+    x = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
+    z = x * 2
+    z[:2] = Tensor([10.0, 20.0])
+    z.sum().backward()
+    np.testing.assert_allclose(x.grad.numpy(), [0, 0, 2, 2])
+
+  def test_set_overlapping_requires_grad(self):
+    z = Tensor.zeros(6, requires_grad=True)
+    x = Tensor.ones(4, requires_grad=True)
+    y = Tensor.ones(4, requires_grad=True) * 2
+    z[:4] = x
+    z[2:] = y
+    z.sum().backward()
+    np.testing.assert_allclose(x.grad.numpy(), [1, 1, 0, 0])
+    np.testing.assert_allclose(y.grad.numpy(), np.ones(4))
+
+  def test_set_iadd_requires_grad(self):
+    z = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
+    x = Tensor([10.0, 20.0], requires_grad=True)
+    z[:2] += x
+    z.sum().backward()
+    np.testing.assert_allclose(z.grad.numpy(), np.ones(4))
+    np.testing.assert_allclose(x.grad.numpy(), np.ones(2))
+
+  def test_set_used_before_setitem(self):
+    z = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
+    _ = z.sum()
+    with self.assertRaises(RuntimeError):
+      z[:2] = Tensor([0.0, 0.0])
 
 class TestSetitemLoop(unittest.TestCase):
   def test_arange(self):

test/backend/test_stunning.py

@@ -25,7 +25,7 @@ class TestStunning(unittest.TestCase):
     nv = a[12].cat(a[76]).tolist()
 
     vi = Variable('i', 0, a.shape[0]-1)
-    with self.assertRaisesRegex(AssertionError, "bind mismatch on"):
+    with self.assertRaisesRegex(RuntimeError, "bind mismatch on"):
       wv = a[vi.bind(12)].cat(a[vi.bind(76)]).tolist()
       self.assertListEqual(nv, wv)
 

test/backend/test_subbuffer.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.device import Buffer
-from tinygrad.helpers import Context
+from tinygrad.helpers import Context, getenv
 from test.helpers import needs_second_gpu
 
 @unittest.skipUnless(hasattr(Device[Device.DEFAULT].allocator, "_offset"), "subbuffer not supported")
@@ -42,7 +42,7 @@ class TestSubBuffer(unittest.TestCase):
     assert out == [102, 103]
 
   @needs_second_gpu
-  @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"}, "only NV, AMD, CUDA")
+  @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or getenv("MOCKGPU"), "only NV, AMD, CUDA")
   def test_subbuffer_transfer(self):
     t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
     vt = t[2:5].contiguous().realize()

test/backend/test_symbolic_ops.py

@@ -69,6 +69,58 @@ class TestSymbolicOps(unittest.TestCase):
         # symbolic shape dropout is not supported
         self.test_attention(dropout_p=0.5)
 
+  def test_sdpa_symbolic_seq_len(self):
+    # symbolic seq_len on all of q/k/v (dim -2 after transpose)
+    q = Tensor.rand(2, 10, 4, 8)
+    k = Tensor.rand(2, 10, 4, 8)
+    v = Tensor.rand(2, 10, 4, 8)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      Tensor.realize(q, k, v)
+      symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention(
+        k[:, :vi].transpose(1, 2), v[:, :vi].transpose(1, 2)).realize()[:2, :4, :i, :8].numpy()
+      expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention(
+        k[:, :i].transpose(1, 2), v[:, :i].transpose(1, 2)).realize().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_sdpa_symbolic_seq_len_query_only(self):
+    # symbolic seq_len on query only (dim -2 after transpose)
+    q = Tensor.rand(2, 10, 4, 8)
+    k = Tensor.rand(2, 5, 4, 8)
+    v = Tensor.rand(2, 5, 4, 8)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      Tensor.realize(q, k, v)
+      symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention(
+        k.transpose(1, 2), v.transpose(1, 2)).realize()[:2, :4, :i, :8].numpy()
+      expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention(
+        k.transpose(1, 2), v.transpose(1, 2)).realize().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_sdpa_symbolic_batch(self):
+    # symbolic batch dim (dim 0)
+    q = Tensor.rand(10, 4, 3, 8)
+    k = Tensor.rand(10, 4, 3, 8)
+    v = Tensor.rand(10, 4, 3, 8)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      Tensor.realize(q, k, v)
+      symbolic = q[:vi].scaled_dot_product_attention(k[:vi], v[:vi]).realize()[:i, :4, :3, :8].numpy()
+      expected = q[:i].scaled_dot_product_attention(k[:i], v[:i]).realize().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_sdpa_symbolic_heads(self):
+    # symbolic heads dim (dim -3)
+    q = Tensor.rand(2, 10, 3, 8)
+    k = Tensor.rand(2, 10, 3, 8)
+    v = Tensor.rand(2, 10, 3, 8)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      Tensor.realize(q, k, v)
+      symbolic = q[:, :vi].scaled_dot_product_attention(k[:, :vi], v[:, :vi]).realize()[:2, :i, :3, :8].numpy()
+      expected = q[:, :i].scaled_dot_product_attention(k[:, :i], v[:, :i]).realize().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
   def test_attention_pos_0_sz_0(self):
     Attention(128, 8)(Tensor.ones(1, 0, 128), Variable("start_pos", 0, 128).bind(0), None)
 

test/backend/test_tensor_variable.py

@@ -136,6 +136,30 @@ class TestTensorVariable(unittest.TestCase):
     with self.assertRaises(AssertionError):
       t.chunk(2, dim=0)  # can't split along symbolic dim
 
+  def test_symbolic_var_sum(self, var_name="u"):
+    t = Variable("t", 1, 10).bind(4)
+    v = Variable(var_name, 1, 5).bind(1)
+    mask = (Tensor.full((1, 1, t, v+t), 1) + 1).contiguous()
+    mask.shrink(((0, 1), (0, 1), (0, 4), (0, 4))).numpy()
+  def test_symbolic_var_sum_alt_name(self): self.test_symbolic_var_sum("s")
+
+  def test_symbolic_triu(self):
+    t = Variable("t", 1, 10).bind(4)
+    for start_pos in (0, 1, 3):
+      var_start_pos = Variable("start_pos", 0, 5).bind(start_pos)
+      mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).triu(var_start_pos+1)
+      out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy()
+      expected = np.triu(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1)
+      np.testing.assert_equal(out, expected)
+
+  def test_symbolic_tril(self):
+    t = Variable("t", 1, 10).bind(4)
+    for start_pos in (0, 1, 3):
+      var_start_pos = Variable("start_pos", 0, 5).bind(start_pos)
+      mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).tril(var_start_pos+1)
+      out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy()
+      expected = np.tril(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1)
+      np.testing.assert_equal(out, expected)
 
 if __name__ == '__main__':
   unittest.main()

test/external/external_fuzz_beam_timeout_recovery.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+"""
+Stress test for beam timeout + device recovery on AM devices.
+
+Usage:
+  AMD=1 python test/external/external_test_beam_timeout_recovery.py
+"""
+from tinygrad import Tensor, Device
+from tinygrad.helpers import Context
+from tinygrad.runtime.ops_amd import AMDDevice
+
+if __name__ == "__main__":
+  dev = Device["AMD"]
+  assert isinstance(dev, AMDDevice) and dev.is_am(), "not am"
+
+  N = 10000
+  for i in range(N):
+    with Context(DEBUG=0, BEAM=0):
+      a = Tensor.rand(4096, 4096, device="AMD").contiguous().realize()
+      b = Tensor.rand(4096, 4096, device="AMD").contiguous().realize()
+      c = a.matmul(b)
+      c.realize()
+    try: dev.synchronize(timeout=1)
+    except RuntimeError as e: print(e)
+    with Context(DEBUG=0, BEAM=0):
+      a = Tensor.ones(512, 512, device="AMD").contiguous().realize()
+      b = Tensor.ones(512, 512, device="AMD").contiguous().realize()
+      result = a.matmul(b).realize()[0, 0].item()
+    assert result == 512.0, f"iter {i}: got {result}"
+    print(f"  iter {i+1}/{N}: ok")
+  print(f"=== All {N} iterations passed ===")

test/external/external_fuzz_sdma_warm_start.py

@@ -0,0 +1,55 @@
+import subprocess, sys, os, random
+
+CHILD_SCRIPT = """
+import os, random
+import numpy as np
+from tinygrad import Tensor, Device
+from tinygrad.runtime.ops_amd import AMDDevice
+
+dev = Device["AMD"]
+for i in range({N}):
+  sz = random.randint(1, {MAX_SZ})
+  data = np.random.randint(0, 256, sz, dtype=np.uint8)
+  t = Tensor(data, device="AMD").contiguous().realize()
+  dev.synchronize()
+  result = t.numpy()
+  assert (result == data).all(), f"Data mismatch at iter {{i}}"
+""".strip()
+
+def run_child(n_ops, max_sz, timeout):
+  env = os.environ.copy()
+  env.setdefault("SDMA_RING_SIZE", "4096")
+
+  script = CHILD_SCRIPT.format(N=n_ops, MAX_SZ=max_sz)
+  p = subprocess.Popen([sys.executable, "-c", script], stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env)
+
+  try:
+    _, stderr = p.communicate(timeout=timeout)
+    return ("ok" if p.returncode == 0 else "fail"), stderr.decode(errors='replace')
+  except subprocess.TimeoutExpired:
+    p.kill()
+    p.communicate()
+    return "timeout", "TIMEOUT: SDMA ring likely stuck"
+
+if __name__ == "__main__":
+  n_iters = int(os.environ.get("FUZZ_ITERS", "10000"))
+  timeout = int(os.environ.get("FUZZ_TIMEOUT", "10"))
+  max_sz = int(os.environ.get("FUZZ_MAX_SZ", "65536"))
+
+  timeouts = 0
+  failures = 0
+
+  for i in range(n_iters):
+    # Run child with many ops to stress the small sdma ring buffer across warm starts
+    n_ops = random.randint(20, 100)
+    status, stderr = run_child(n_ops=n_ops, max_sz=max_sz, timeout=timeout)
+    if status == "timeout":
+      timeouts += 1
+      print(f"\tstderr: {stderr[:500]}")
+    elif status == "fail":
+      failures += 1
+      print(f"\tstderr: {stderr[:500]}")
+    else:
+      print(f"iter {i}: ok (n_ops={n_ops})")
+
+  print(f"\n=== Results: {n_iters} iterations, {timeouts} timeouts, {failures} failures ===")

test/external/external_test_gpu_crash.py

@@ -4,13 +4,19 @@
 These tests intentionally cause GPU faults to verify error handling.
 Run with: AMD=1 python -m pytest test/external/external_test_gpu_crash.py -v
 """
-import unittest, re
+import unittest, re, importlib
 from tinygrad.device import Device
-from tinygrad.runtime.autogen.amd.rdna3.ins import *  # noqa: F403
 from tinygrad.renderer.amd.dsl import s, v, Inst, NULL
 
-def assemble(code:str, name:str="test") -> str:
-  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "wavefront_size32": 1, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
+RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "s_waitcnt": "s_waitcnt", "s_endpgm": "s_endpgm",
+                   "global_load_b32": "global_load_dword", "global_store_b32": "global_store_dword",
+                   "global_atomic_add_u32": "global_atomic_add", "flat_load_b32": "flat_load_dword",
+                   "flat_store_b32": "flat_store_dword", "flat_atomic_add_u32": "flat_atomic_add", "s_load_b32": "s_load_dword"}
+
+def assemble(code:str, name:str="test", is_cdna:bool=False) -> str:
+  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
+  if is_cdna: kd["accum_offset"] = 8
+  else: kd["wavefront_size32"] = 1
   return f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{code}\n.rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + \
          "\n".join(f".amdhsa_{k} {v}" for k,v in kd.items()) + "\n.end_amdhsa_kernel"
 
@@ -21,6 +27,10 @@ class TestGPUCrash(unittest.TestCase):
     from tinygrad.runtime.support.compiler_amd import HIPCompiler
     cls.dev = Device["AMD"]
     cls.compiler = HIPCompiler(cls.dev.arch)
+    cls.is_cdna = cls.dev.target[0] < 10
+    ins = importlib.import_module('tinygrad.runtime.autogen.amd.' + ('cdna' if cls.is_cdna else 'rdna3') + '.ins')
+    for rdna3_name, cdna3_name in RDNA3_CDNA3_MAP.items():
+      setattr(cls, rdna3_name, getattr(ins, cdna3_name if cls.is_cdna else rdna3_name))
 
   def setUp(self):
     # Verify device works before each test
@@ -33,7 +43,7 @@ class TestGPUCrash(unittest.TestCase):
 
   def _run(self, code: str):
     from tinygrad.runtime.ops_amd import AMDProgram
-    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code)))
+    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code, is_cdna=self.is_cdna)))
     prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)
 
   def _run_insts(self, insts: list[Inst]):
@@ -57,32 +67,32 @@ class TestOutOfBoundsMemoryAccess(TestGPUCrash):
 
   def test_global_load_null_ptr(self):
     """Global load from NULL pointer."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
-             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0),
+             self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
   def test_global_store_null_ptr(self):
     """Global store to NULL pointer."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
-             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF),
+             self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
   def test_global_load_unmapped_high_address(self):
     """Global load from high unmapped address (0xDEAD00000000)."""
-    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD),
-             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD),
+             self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
   def test_global_store_unmapped_high_address(self):
     """Global store to high unmapped address."""
-    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678),
-             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 0x12345678),
+             self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
   def test_global_atomic_unmapped(self):
     """Atomic operation on unmapped memory."""
-    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1),
-             global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0xBEEF0000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 1),
+             self.global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
 
@@ -91,14 +101,14 @@ class TestSMEMFaults(TestGPUCrash):
 
   def test_smem_load_null(self):
     """SMEM load from NULL base."""
-    insts = [s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
-             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
+    insts = [self.s_mov_b32(s[2], 0), self.s_mov_b32(s[3], 0),
+             self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
   def test_smem_load_unmapped(self):
     """SMEM load from unmapped address."""
-    insts = [s_mov_b32(s[2], 0xBEEF0000), s_mov_b32(s[3], 0xDEAD),
-             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
+    insts = [self.s_mov_b32(s[2], 0xBEEF0000), self.s_mov_b32(s[3], 0xDEAD),
+             self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
 
@@ -107,20 +117,20 @@ class TestFlatMemoryFaults(TestGPUCrash):
 
   def test_flat_load_null(self):
     """FLAT load from NULL address."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
-             flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0),
+             self.flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
   def test_flat_store_null(self):
     """FLAT store to NULL address."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
-             flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF),
+             self.flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
   def test_flat_atomic_null(self):
     """FLAT atomic on NULL address."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 1),
-             flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 1),
+             self.flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
 

test/external/external_test_hive_reset.py

@@ -2,8 +2,9 @@
 import subprocess, sys
 from tinygrad.helpers import getenv
 
-LOOPS = getenv("LOOPS", 10)
+LOOPS = getenv("LOOPS", 50)
 BROKEN = getenv("BROKEN", 0)
+ONLY_RESET = getenv("ONLY_RESET", 0)
 
 BROKEN_KERNEL_SCRIPT = """
 from tinygrad.device import Device
@@ -36,7 +37,7 @@ for i in range(LOOPS):
     print(f"=== Running broken kernel ({i+1}/{LOOPS}) ===")
     ret = subprocess.run([sys.executable, "-c", BROKEN_KERNEL_SCRIPT])
     print(f"=== broken kernel exited with code {ret.returncode} ===")
-
-  print(f"=== Running test_tiny.py ({i+1}/{LOOPS}) ===")
-  ret = subprocess.run([sys.executable, "test/test_tiny.py", "TestTiny.test_plus"])
-  print(f"=== test_tiny.py exited with code {ret.returncode} ===")
+  elif not ONLY_RESET:
+    print(f"=== Running test_tiny.py ({i+1}/{LOOPS}) ===")
+    ret = subprocess.run([sys.executable, "test/test_tiny.py", "TestTiny.test_plus"])
+    print(f"=== test_tiny.py exited with code {ret.returncode} ===")