dt.count

count
bitsize in nir
2026-08-20 05:26:07 +00:00 · 2026-06-02 13:19:04 -07:00 · 2026-06-02 13:16:42 -07:00 · 2026-06-02 13:11:31 -07:00 · 2026-06-02 13:02:28 -07:00 · 2026-06-02 22:34:25 +03:00
123 changed files with 1699 additions and 1595 deletions
@@ -5,6 +5,7 @@ runs:
  steps:
    - name: Run process replay tests
      shell: bash
+      if: env.CAPTURE_PROCESS_REPLAY == '1'
      run: |
        export PR_TITLE=$(jq -r .pull_request.title "$GITHUB_EVENT_PATH")
        export CURRENT_SHA=${{ github.event.pull_request && github.event.pull_request.head.sha || github.sha }}
@@ -228,6 +228,11 @@ runs:

        sudo chown -R $USER:$USER /var/cache/apt/archives/

+    - name: Add clang to PATH (Linux)
+      if: inputs.llvm == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: echo "/usr/lib/llvm-20/bin" >> "$GITHUB_PATH"
+
    # **** AMD ****
    - name: Setup AMD (Linux)
      if: inputs.amd == 'true' && runner.os == 'Linux'
@@ -281,7 +286,7 @@ runs:
        cache-name: cache-gpuocelot-build-1
      with:
        path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-gpuocelot-f463259669c69abce7b3a0567b6c284f348d0f32-rebuild-${{ env.CACHE_VERSION }}
    - name: Cache gpuocelot
      if: inputs.ocelot == 'true' && github.event_name != 'pull_request'
      id: cache-build
@@ -290,26 +295,22 @@ runs:
        cache-name: cache-gpuocelot-build-1
      with:
        path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-gpuocelot-f463259669c69abce7b3a0567b6c284f348d0f32-rebuild-${{ env.CACHE_VERSION }}
    - name: Clone/compile gpuocelot
      if: inputs.ocelot == 'true' && steps.cache-build-pr.outputs.cache-hit != 'true' && steps.cache-build.outputs.cache-hit != 'true'
      shell: bash
      run: |
-        git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
+        git clone --recurse-submodules https://github.com/tinygrad/gpuocelot.git ${{ github.workspace }}/gpuocelot
        cd ${{ github.workspace }}/gpuocelot/ocelot
-        git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
+        git checkout f463259669c69abce7b3a0567b6c284f348d0f32
        mkdir build
        cd build

        CMAKE_ARGS="-Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5"
        if [[ "${{ runner.os }}" == "macOS" ]]; then
-          sudo xcode-select -s /Applications/Xcode_16.2.app/Contents/Developer
          CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib"
        else
-          curl -fL https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/linux-x86_64/cuda_nvcc-linux-x86_64-11.5.119-archive.tar.xz \
-            | sudo tar -xJ -C /usr/ --strip-components=1
-          curl -fL https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/linux-x86_64/cuda_cudart-linux-x86_64-11.5.117-archive.tar.xz \
-            | sudo tar -xJ -C /usr/ --strip-components=1
+          CMAKE_ARGS="$CMAKE_ARGS -DLLVM_DIR=$(llvm-config-15 --cmakedir)"
        fi

        cmake .. $CMAKE_ARGS
@@ -806,3 +806,16 @@ jobs:
        pkill -f 'extra/remote/serve.py' || true
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
+
+  llvmspeed:
+    name: LLVM Speed
+    runs-on: [self-hosted, Linux, tinyboxrandom]
+    timeout-minutes: 20
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Speed Test
+      run: DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
+    - name: Speed Test (BEAM=2)
+      run: BEAM=2 DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
@@ -2,7 +2,7 @@ name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
  CACHE_VERSION: '19'
-  CAPTURE_PROCESS_REPLAY: 1
+  CAPTURE_PROCESS_REPLAY: ${{ github.event_name == 'pull_request' && contains(github.event.pull_request.title, '[pr]') && '1' || '0' }}
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
  CHECK_OOB: 1
@@ -14,28 +14,14 @@ on:
  pull_request:
  workflow_dispatch:

-jobs:
-  llvmspeed:
-    name: LLVM Speed
-    runs-on: ubuntu-24.04
-    timeout-minutes: 20
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: llvm-speed
-        deps: testing_unit
-        llvm: 'true'
-    - name: Speed Test
-      run: DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
-    - name: Speed Test (BEAM=2)
-      run: BEAM=2 DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
+concurrency:
+  group: test-${{ github.event_name }}-${{ github.event_name == 'pull_request' && github.event.pull_request.number || github.run_id }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}

+jobs:
  docs:
    name: Docs
-    runs-on: ubuntu-22.04
+    runs-on: &linux ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 10
    env:
      CHECK_OOB: 0
@@ -89,7 +75,7 @@ jobs:

  torchbackend:
    name: Torch Backend Tests
-    runs-on: ubuntu-latest
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -125,7 +111,7 @@ jobs:

  torchbackendmore:
    name: Torch Backend Tests More
-    runs-on: ubuntu-latest
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -147,7 +133,7 @@ jobs:

  bepython:
    name: Python Backend
-    runs-on: ubuntu-latest
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -157,65 +143,65 @@ jobs:
      with:
        key: be-minimal
        deps: testing_unit
-    - name: Test dtype with Python emulator
-      run: DEBUG=1 DEV=PYTHON python3 -m pytest -n=auto test/backend/test_dtype.py test/backend/test_dtype_alu.py
-    - name: Test ops with Python emulator
-      run: DEBUG=2 SKIP_SLOW_TEST=1 DEV=PYTHON python3 -m pytest -n=auto test/backend/test_ops.py --durations=20
-    - name: Test uops with Python emulator
-      run: DEV=PYTHON python3 -m pytest test/backend/test_uops.py --durations=20
-    - name: Test symbolic with Python emulator
-      run: DEV=PYTHON python3 test/backend/test_symbolic_ops.py
-    - name: test_renderer_failures with Python emulator
-      run: DEV=PYTHON python3 -m pytest -rA test/backend/test_renderer_failures.py::TestRendererFailures
+    - name: Run backend tests
+      run: SKIP_SLOW_TEST=1 DEV=PYTHON python3 -m pytest -n=auto test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_ops.py test/backend/test_uops.py test/backend/test_symbolic_ops.py test/backend/test_renderer_failures.py::TestRendererFailures --durations=20
    - name: Test IMAGE support
-      run: |
-        IMAGE=1 DEV=PYTHON python3 test/backend/test_ops.py TestOps.test_gemm
-        IMAGE=1 DEV=PYTHON python3 test/backend/test_ops.py TestOps.test_simple_conv2d
+      run: IMAGE=1 DEV=PYTHON python3 test/backend/test_ops.py TestOps.test_gemm TestOps.test_simple_conv2d
    - name: Test emulated METAL tensor cores
+      env:
+        DEV: 'PYTHON::METAL'
      run: |
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::METAL python3 test/backend/test_ops.py TestOps.test_big_gemm
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::METAL python3 test/opt/test_tensor_cores.py
-    - name: Test emulated AMX tensor cores
-      run: DEBUG=2 AMX=1 FORWARD_ONLY=1 DEV=PYTHON::AMX python3 test/backend/test_ops.py TestOps.test_gemm
+        DEBUG=2 python3 test/backend/test_ops.py TestOps.test_big_gemm
+        python3 -m pytest -nauto test/opt/test_tensor_cores.py
    - name: Test emulated AMD tensor cores
+      env:
+        DEV: 'PYTHON::gfx1100'
      run: |
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::gfx1100 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::gfx1100 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::gfx1100 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::gfx1100 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::gfx1100 python3 test/opt/test_tensor_cores.py
+        DEBUG=2 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        python3 -m pytest -nauto test/opt/test_tensor_cores.py
    - name: Test emulated AMD MFMA tensor cores
+      env:
+        DEV: 'PYTHON::gfx950'
      run: |
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::gfx950 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::gfx950 python3 test/opt/test_tensor_cores.py
+        DEBUG=2 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        python3 -m pytest -nauto test/opt/test_tensor_cores.py
    - name: Test emulated AMD RDNA4 tensor cores
+      env:
+        DEV: 'PYTHON::gfx1201'
      run: |
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::gfx1201 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::gfx1201 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::gfx1201 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::gfx1201 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::gfx1201 python3 test/opt/test_tensor_cores.py
+        DEBUG=2 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        python3 -m pytest -nauto test/opt/test_tensor_cores.py
    - name: Test emulated CUDA tensor cores
      run: |
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::sm_80 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 ALLOW_TF32=1 FORWARD_ONLY=1 DEV=PYTHON::sm_80 python3 test/backend/test_ops.py TestOps.test_gemm
-        DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::sm_75 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 ALLOW_TF32=1 FORWARD_ONLY=1 DEV=PYTHON::sm_89 python3 test/opt/test_tensor_cores.py
+        DEBUG=2 DEV=PYTHON::sm_80 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 ALLOW_TF32=1 DEV=PYTHON::sm_80 python3 test/backend/test_ops.py TestOps.test_gemm
+        DEBUG=2 DEV=PYTHON::sm_75 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
+        ALLOW_TF32=1 DEV=PYTHON::sm_89 python3 -m pytest -nauto test/opt/test_tensor_cores.py
    - name: Test emulated INTEL OpenCL tensor cores
-      run: DEBUG=2 FORWARD_ONLY=1 DEV=PYTHON::INTEL HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
+      run: DEBUG=2 DEV=PYTHON::INTEL HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
    - name: Test emulated AMX tensor cores
-      run: DEBUG=2 AMX=1 FORWARD_ONLY=1 DEV=PYTHON::AMX python3 test/opt/test_tensor_cores.py
+      env:
+        DEV: 'PYTHON::AMX'
+      run: |
+        DEBUG=2 python3 test/backend/test_ops.py TestOps.test_gemm
+        python3 -m pytest -nauto test/opt/test_tensor_cores.py
    - name: Test device flop counts
      run: |
        DEBUG=2 DEV=PYTHON::METAL python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
        DEBUG=2 DEV=PYTHON::gfx1100 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
        DEBUG=2 DEV=PYTHON::sm_80 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
        DEBUG=2 DEV=PYTHON::INTEL python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 AMX=1 DEV=PYTHON::AMX python3 ./test/null/test_uops_stats.py TestUOpsStats.test_simple_matmul
+        DEBUG=2 DEV=PYTHON::AMX python3 ./test/null/test_uops_stats.py TestUOpsStats.test_simple_matmul

  linter:
    name: Linters
-    runs-on: ubuntu-latest
+    runs-on: *linux
    timeout-minutes: 10

    steps:
@@ -246,7 +232,7 @@ jobs:

  nulltest:
    name: Null Tests
-    runs-on: ubuntu-latest
+    runs-on: *linux
    timeout-minutes: 15

    steps:
@@ -277,7 +263,7 @@ jobs:

  unittest:
    name: Unit Tests
-    runs-on: ubuntu-latest
+    runs-on: *linux
    timeout-minutes: 15

    steps:
@@ -320,7 +306,7 @@ jobs:
      matrix:
        group: [1, 2]
    name: SPEC=2 (${{ matrix.group }})
-    runs-on: ubuntu-latest
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -331,12 +317,13 @@ jobs:
        key: spec-unit
        deps: testing_unit
        python-version: '3.14'
+        llvm: 'true'
    - name: Test SPEC=2
      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" -k "not test_conv2d_ceildiv_edge_case" --splits 2 --group ${{ matrix.group }}

  fuzzing:
    name: Fuzzing
-    runs-on: ubuntu-latest
+    runs-on: *linux
    timeout-minutes: 10
    steps:
    - name: Checkout Code
@@ -357,7 +344,7 @@ jobs:

  testopenclimage:
    name: CL IMAGE Tests
-    runs-on: ubuntu-22.04
+    runs-on: *linux
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -377,7 +364,7 @@ jobs:

  testgpumisc:
    name: CL Misc tests
-    runs-on: ubuntu-22.04
+    runs-on: *linux
    timeout-minutes: 10
    steps:
      - name: Checkout Code
@@ -402,7 +389,7 @@ jobs:

  testopenpilot:
    name: openpilot Compile Tests
-    runs-on: ubuntu-22.04
+    runs-on: *linux
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -417,8 +404,6 @@ jobs:
      - name: Test openpilot model kernel count and gate usage
        run: |
          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1468 ALLOWED_GATED_READ_IMAGE=18 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
-      - name: Test openpilot CL compile fp16
-        run: FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Test openpilot CL compile fp32 (test correctness)
        run: |
          DEV=CL IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
@@ -432,7 +417,7 @@ jobs:

  testonnxcpu:
    name: ONNX (CPU) Tests
-    runs-on: ubuntu-22.04
+    runs-on: *linux
    timeout-minutes: 20

    steps:
@@ -446,21 +431,13 @@ jobs:
          python-version: '3.12'
          llvm: 'true'
      - name: Test ONNX (CPU)
-        run: DEV=CPU python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
-      - name: Test ONNX (LLVM)
-        run: DEV=CPU:LLVM python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
-      - name: Test ONNX Runner (CPU)
-        run: DEV=CPU python3 test/external/external_test_onnx_runner.py
-      - name: Test Additional ONNX Ops (CPU)
-        run: DEV=CPU python3 test/external/external_test_onnx_ops.py
-      - name: Test Quantize ONNX
-        run: DEV=CPU python3 test/backend/test_quantize_onnx.py
+        run: DEV=CPU python -m pytest -n=auto test/external/external_test_onnx_backend.py test/external/external_test_onnx_runner.py test/external/external_test_onnx_ops.py test/backend/test_quantize_onnx.py --durations=20
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-  testopencl:
-    name: ONNX (CL)+Optimization Tests
-    runs-on: ubuntu-22.04
+  testoptim:
+    name: Optimization Tests
+    runs-on: *linux
    timeout-minutes: 20
    steps:
      - name: Checkout Code
@@ -468,13 +445,11 @@ jobs:
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
-          key: onnxoptl
+          key: optim
          deps: testing
          pydeps: "tensorflow==2.19"
          python-version: '3.12'
          opencl: 'true'
-      - name: Test ONNX (CL)
-        run: DEV=CL python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
      #- name: Test Optimization Helpers
      #  run: DEBUG=1 python3 extra/optimization/test_helpers.py
      #- name: Test Action Space
@@ -494,7 +469,7 @@ jobs:

  testllm:
    name: Test LLM
-    runs-on: ubuntu-24.04
+    runs-on: *linux
    timeout-minutes: 15
    env:
      CHECK_OOB: 0
@@ -519,7 +494,7 @@ jobs:

  testmodels:
    name: Models (llvm+cpu+gpu)
-    runs-on: ubuntu-22.04
+    runs-on: *linux
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -542,7 +517,7 @@ jobs:

  testmetalmodels:
    name: Models (metal)
-    runs-on: macos-14
+    runs-on: &macos macos-26
    timeout-minutes: 20
    steps:
    - name: Checkout Code
@@ -562,7 +537,7 @@ jobs:

  testdsp:
    name: Linux (DSP)
-    runs-on: ubuntu-24.04
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -571,8 +546,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: dsp-minimal
-        deps: testing_unit
-        pydeps: "onnx==1.18.0 onnxruntime ml_dtypes"
+        deps: testing
        llvm: "true"
        qemu: "true"
    - name: Set MOCKDSP env
@@ -580,13 +554,24 @@ jobs:
    - name: Run test_tiny on DSP
      run: DEBUG=2 DEV=DSP python test/test_tiny.py
    - name: Test transcendentals
-      run: CC=clang-20 DEBUG=2 DEV=DSP python test/backend/test_transcendental.py TestTranscendentalVectorized
+      run: DEBUG=2 DEV=DSP python test/backend/test_transcendental.py TestTranscendentalVectorized
    - name: Test quantize onnx
      run: DEBUG=2 DEV=DSP python3 test/backend/test_quantize_onnx.py

-  testwebgpu:
-    name: Linux (WebGPU)
-    runs-on: ubuntu-22.04
+  testlinux:
+    strategy:
+      fail-fast: false
+      matrix:
+        dev:
+          - 'CPU:CLANG'
+          - 'CPU:LLVM'
+          - 'CPU:LVP'
+          - 'CPU:X86'
+          - 'CL'
+          - 'WEBGPU'
+
+    name: Linux (DEV=${{ matrix.dev }})
+    runs-on: *linux
    timeout-minutes: 20
    steps:
    - name: Checkout Code
@@ -594,23 +579,27 @@ jobs:
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: webgpu-minimal
+        key: linux-${{ matrix.dev }}
        deps: testing_unit
        python-version: '3.12'
-        webgpu: 'true'
-    - name: Check Device.DEFAULT (WEBGPU) and print some source
+        llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') || contains(matrix.dev, 'CLANG') }}
+        mesa: ${{ contains(matrix.dev, 'LVP') && 'cpu' || 'false' }}
+        webgpu: ${{ matrix.dev == 'WEBGPU' }}
+        opencl: ${{ matrix.dev == 'CL' }}
+    - name: Set env
+      run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
+    - name: Check Device.DEFAULT and print some source
      run: |
-        DEV=WEBGPU python -c "from tinygrad import Device; assert Device.DEFAULT == 'WEBGPU', Device.DEFAULT"
-        DEV=WEBGPU DEBUG=4 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
-    - name: Run selected webgpu tests
-      run: |
-          DEV=WEBGPU WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/backend --durations=20
+        python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
+        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
+    - name: Run backend tests
+      run: python -m pytest -n=auto test/backend --durations=20
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

  testamdasm:
    name: AMD ASM IDE
-    runs-on: ubuntu-24.04
+    runs-on: *linux
    timeout-minutes: 20
    env:
      DEV: MOCKKFD+AMD
@@ -657,7 +646,7 @@ jobs:

  testmockam:
    name: Linux (am)
-    runs-on: ubuntu-24.04
+    runs-on: *linux
    timeout-minutes: 15
    env:
      DEV: MOCKPCI+AMD
@@ -693,7 +682,7 @@ jobs:
        arch: [gfx1100, gfx1201, gfx950]

    name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
-    runs-on: ubuntu-22.04
+    runs-on: *linux
    timeout-minutes: 15
    env:
      DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
@@ -728,7 +717,7 @@ jobs:
        backend: [ptx, nv]

    name: Linux (${{ matrix.backend }})
-    runs-on: ubuntu-22.04
+    runs-on: *linux
    timeout-minutes: 20
    env:
      FORWARD_ONLY: 1
@@ -756,44 +745,11 @@ jobs:
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-  testcpuopencl:
-    strategy:
-      fail-fast: false
-      matrix:
-        backend: [llvm, cpu, opencl, lvp, x86]
-
-    name: Linux (${{ matrix.backend }})
-    runs-on: ubuntu-22.04
-    timeout-minutes: 20
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v6
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: ${{ matrix.backend }}-minimal
-          deps: testing_unit
-          opencl: ${{ matrix.backend == 'opencl' && 'true' }}
-          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'cpu' || matrix.backend == 'lvp' }}
-          mesa: ${{ matrix.backend == 'lvp' && 'cpu' }}
-      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'DEV=CPU:LLVM' || matrix.backend == 'cpu' && 'CC=clang-20\nDEV=CPU\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'DEV=CL' || matrix.backend == 'lvp' && 'DEV=CPU:LVP' || matrix.backend == 'x86' && 'DEV=CPU:X86' }}" >> $GITHUB_ENV
-      - name: Check Device.DEFAULT and print some source
-        run: |
-          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
-          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
-      - name: Run pytest (${{ matrix.backend }})
-        run: python -m pytest -n=auto test/backend --durations=20
-      - name: Run TRANSCENDENTAL math
-        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
-      - name: Run process replay tests
-        uses: ./.github/actions/process-replay
-
 # ****** OSX Tests ******

-  testmetal:
+  unittestmacos:
    name: MacOS (unit)
-    runs-on: macos-14
+    runs-on: *macos
    timeout-minutes: 20
    steps:
    - name: Checkout Code
@@ -801,18 +757,15 @@ jobs:
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: metal
-        deps: testing
+        key: unittest-macos
+        deps: testing_unit
        python-version: '3.12'
        amd: 'true'
        ocelot: 'true'
-        llvm: 'true'
    - name: Run unit tests
      run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
    - name: Run NULL backend tests
      run: DEV=NULL python -m pytest -n=auto test/null/ --durations=20
-    - name: Run ONNX
-      run: DEV=METAL python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
    - name: Test tensor core ops (fake)
      run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
    - name: Test tensor core ops (real)
@@ -823,20 +776,12 @@ jobs:
      run: DEV=METAL python3 -m pytest test/device/test_metal.py
    #- name: Fuzz Test linearizer
    #  run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
-    - name: Run TRANSCENDENTAL math
-      run: DEV=METAL TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
    - name: Run pytest (amd)
      env:
        DEV: MOCKKFD+AMD
        FORWARD_ONLY: 1
      run: |
        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
-    - name: Run pytest (amd with llvm backend)
-      env:
-        DEV: "MOCKKFD+AMD:LLVM"
-        FORWARD_ONLY: 1
-      run: |
-        python -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py test/device/test_amd_llvm.py --durations=20
    - name: Run pytest (ptx)
      env:
        DEV: "MOCK+NV:PTX"
@@ -848,85 +793,57 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

-  osxwebgpu:
-    name: MacOS (WebGPU)
-    runs-on: macos-14
-    timeout-minutes: 10
+  testmacos:
+    strategy:
+      fail-fast: false
+      matrix:
+        dev:
+          - 'CPU:CLANG'
+          - 'CPU:LLVM'
+          - 'CPU:LVP'
+          - 'METAL'
+          - 'WEBGPU'
+
+    name: MacOS (DEV=${{ matrix.dev }})
+    runs-on: *macos
+    timeout-minutes: 20
    steps:
    - name: Checkout Code
      uses: actions/checkout@v6
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: osx-webgpu
-        deps: testing
-        webgpu: 'true'
-    - name: Build WEBGPU Efficientnet
-      run: DEV=WEBGPU WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m examples.compile_efficientnet
-    - name: Run selected webgpu tests
-      run: DEV=WEBGPU WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m pytest -n=auto test/backend --durations=20
-    #- name: Clean npm cache
-    #  run: npm cache clean --force
-    #- name: Install Puppeteer
-    #  run: npm install puppeteer
-    # this is also flaky
-    #- name: Run WEBGPU Efficientnet
-    #  run: node test/web/test_webgpu.js
-    # this is flaky
-    #- name: Run VIZ tests as external package
-    #  run: |
-    #    mkdir $GITHUB_WORKSPACE/test_dir
-    #    cd $GITHUB_WORKSPACE/test_dir
-    #    python -m venv venv
-    #    source venv/bin/activate
-    #    pip install $GITHUB_WORKSPACE
-    #    cp $GITHUB_WORKSPACE/test/web/test_viz.js .
-    #    node test_viz.js
-    - name: Test ONNX Runner (WEBGPU)
-      run: DEV=WEBGPU python3 test/external/external_test_onnx_runner.py
-
-  osxtests:
-    strategy:
-      fail-fast: false
-      matrix:
-        backend: [metal, llvm, cpu, lvp]
-    name: MacOS (${{ matrix.backend }})
-    runs-on: macos-15
-    timeout-minutes: 20
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v6
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: macos-${{ matrix.backend }}-minimal
-          deps: testing_unit
-          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
-          mesa: ${{ matrix.backend == 'lvp' && 'cpu' }}
-      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'DEV=CPU:LLVM' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_COUNT=2' || matrix.backend == 'metal' && 'DEV=METAL' || matrix.backend == 'lvp' && 'DEV=CPU:LVP' }}" >> $GITHUB_ENV
-      - name: Check Device.DEFAULT and print some source
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
-          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-      - name: Run pytest (${{ matrix.backend }})
-        run: python3 -m pytest -n=auto test/backend --durations=20
-      - name: Run process replay tests
-        uses: ./.github/actions/process-replay
-      - name: Run macOS-specific unit test
-        if: matrix.backend == 'llvm'
-        run: python3 -m pytest test/unit/test_disk_tensor.py::TestDiskTensor::test_copy_to_cpu_not_truncated test/unit/test_cpu.py
+        key: macos-${{ matrix.dev }}
+        deps: testing_unit
+        python-version: '3.12'
+        llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }}
+        mesa: ${{ contains(matrix.dev, 'LVP') && 'cpu' || 'false' }}
+        webgpu: ${{ matrix.dev == 'WEBGPU' }}
+    - name: Set env
+      run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
+    - name: Check Device.DEFAULT and print some source
+      run: |
+        python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
+        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
+    - name: Run backend tests
+      run: python -m pytest -n=auto test/backend --durations=20
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay

 # ****** Windows Tests ******

-  wintests:
+  testwindows:
    strategy:
      fail-fast: false
      matrix:
-        backend: [llvm, cpu, webgpu, x86]
+        dev:
+          - 'CPU:CLANG'
+          - 'CPU:LLVM'
+          - 'CPU:X86'
+          - 'WEBGPU'

-    name: Windows (${{ matrix.backend }})
-    runs-on: windows-latest
+    name: Windows (DEV=${{ matrix.dev }})
+    runs-on: windows-2025
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -934,25 +851,20 @@ jobs:
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
-          key: windows-${{ matrix.backend }}-minimal
+          key: windows-${{ matrix.dev }}-minimal
          deps: testing_unit
-          pydeps: ${{ matrix.backend == 'webgpu' && 'dawn-python' || '' }}
+          pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
      - name: Set env
        shell: bash
-        run:  printf "${{ matrix.backend == 'llvm' && 'DEV=CPU:LLVM' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_COUNT=2' || matrix.backend == 'webgpu' && 'DEV=WEBGPU' || matrix.backend == 'x86' && 'DEV=CPU:X86' }}" >> $GITHUB_ENV
-      - name: Run unit tests
-        if: matrix.backend=='llvm'
-        # test_newton_schulz hits RecursionError
-        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
-      - name: Run NULL backend tests
-        if: matrix.backend=='llvm'
-        shell: bash
-        run: DEV=NULL python -m pytest -n=auto test/null/ --ignore=test/null/test_elf.py --durations=20
-      - name: Run pytest (${{ matrix.backend }})
+        run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
+      - name: Check Device.DEFAULT and print some source
        shell: bash
        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU', 'X86':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
-          python -m pytest -n=auto test/test_tiny.py test/backend/test_ops.py --durations=20
+          python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
+          DEBUG=4 python test/test_tiny.py TestTiny.test_plus
+      - name: Run test_tiny
+        shell: bash
+        run: python -m pytest -n=auto test/test_tiny.py --durations=20

 # ****** Compile-only Tests ******

@@ -962,7 +874,7 @@ jobs:
      matrix:
        backend: [ir3, nak]
    name: Compile-only (${{ matrix.backend }})
-    runs-on: ubuntu-24.04
+    runs-on: *linux
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -1419,10 +1419,7 @@ def train_llama3():

  for p in optim.params:
    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
-    if isinstance(p.device, tuple) and p.uop.axis is not None:
-      p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device[0]).shard_(p.device, axis=p.uop.axis).contiguous()
-    else:
-      p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device).contiguous()
+    p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
  grads = [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)
@@ -1438,13 +1435,14 @@ def train_llama3():

  fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts] if hasattr(model, "_fp8_grad_amax") else []
-  fp8_inv_scales = list(model._fp8_inv_scale.values())
+  fp8_inv_scales = list(model._fp8_inv_scale.values()) + list(model._fp8_next_inv_scale.values())

  from tinygrad.nn.state import get_state_dict
  model_state = get_state_dict(model)
  for wname in model._fp8_inv_scale:
    w = model_state[wname]
    w._inv_scale = model._fp8_inv_scale[wname]
+    w._next_inv_scale = model._fp8_next_inv_scale[wname]
    if optim.master_params:
      idx = next(j for j, p in enumerate(optim.params) if p is w)
      master = optim.master_params[idx]
@@ -1478,7 +1476,7 @@ def train_llama3():
    grad_norm = optim.fstep(grads)
    scheduler.step()

-    for g in grads: g.assign(g.zeros_like())
+    for g in grads: g.assign(g.const_like(0))

    lr_cpu = optim.lr.float().to("CPU")
    grad_norm_cpu = grad_norm.float().to("CPU")
@@ -1500,7 +1498,7 @@ def train_llama3():
  def fake_data(bs, samples):
    import numpy as np
    for _ in range(samples // bs):
-      fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
+      fake_data_np = np.random.randint(0, real_vocab_size, size=(bs, SEQLEN + 1), dtype=np.int32)
      yield Tensor(fake_data_np, device="NPY")

  def get_train_iter():
@@ -136,6 +136,7 @@ class FlatTransformer:
    w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)]
    w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)]
    self._fp8_inv_scale = {name: s.float().contiguous().is_param_(False) for name, s in w_scales}
+    self._fp8_next_inv_scale = {name: s.float().contiguous().is_param_(False) for name, s in w_scales}

  def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02):
    if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
@@ -221,14 +222,19 @@ class FlatTransformer:
      for v in get_parameters(self): v.shard_(device, axis=None)
    else:
      # flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
-      self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
-      self.wo.shard_(device, axis=2).realize()            # (n_layers, dim, in) shard in
+      def _shard_fp8(name:str, axis:int):
+        getattr(self, name).shard_(device, axis=axis)
+        self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().is_param_(False)
+        self._fp8_next_inv_scale[name] = self._fp8_next_inv_scale[name].to(device).contiguous().is_param_(False)
+        Tensor.realize(getattr(self, name), self._fp8_inv_scale[name], self._fp8_next_inv_scale[name])
+      _shard_fp8("wqkv", 1)          # (n_layers, out, dim) shard out
+      _shard_fp8("wo", 2)            # (n_layers, dim, in) shard in
      if SPLIT_W13:
-        self.w1.shard_(device, axis=1).realize()
-        self.w3.shard_(device, axis=1).realize()
+        _shard_fp8("w1", 1)
+        _shard_fp8("w3", 1)
      else:
-        self.w13.shard_(device, axis=1).realize()           # (n_layers, hidden*2, dim) shard out
-      self.w2.shard_(device, axis=2).realize()            # (n_layers, dim, hidden) shard in
+        _shard_fp8("w13", 1)         # (n_layers, hidden*2, dim) shard out
+      _shard_fp8("w2", 2)            # (n_layers, dim, hidden) shard in
      self.attention_norm.shard_(device, axis=None).realize()
      self.ffn_norm.shard_(device, axis=None).realize()
      self.norm.weight.shard_(device, axis=None).realize()
@@ -239,8 +245,6 @@ class FlatTransformer:
        for name in amax_dict:
          for i in range(len(amax_dict[name])):
            amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().is_param_(False)
-      for name in self._fp8_inv_scale:
-        self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().is_param_(False)

  def __call__(self, tokens:Tensor, save:bool=True):
    h = self.tok_embeddings(tokens)
@@ -322,11 +326,7 @@ if __name__ == "__main__":

  # preallocate all the grad buffers and zero them out
  grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
-  def _make_grad(x):
-    if isinstance(x.device, tuple) and x.uop.axis is not None:
-      return Tensor.zeros(x.shape, dtype=grad_dtype(x), device=x.device[0]).shard_(x.device, axis=x.uop.axis).contiguous()
-    return Tensor.zeros(x.shape, dtype=grad_dtype(x), device=x.device).contiguous()
-  grads = {x:_make_grad(x) for x in state.values() if x.is_param}
+  grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}

  fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts]
@@ -6,6 +6,7 @@ from tinygrad.uop.ops import UOp, Ops

 STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
 MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
+FP8_AMAX_MARGIN = getenv("FP8_AMAX_MARGIN", 1.1)

 def stochastic_round_bf16(x:Tensor) -> Tensor:
  bits = x.bitcast(dtypes.uint32)
@@ -39,7 +40,8 @@ class GradAccClipAdamW(Optimizer):
    for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
    # collect inv_scale tensors attached to fp8 params (set by _apply_update)
    fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')]
-    to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales
+    fp8_next_inv_scales = [tt._next_inv_scale for tt in self.params if hasattr(tt, '_next_inv_scale')]
+    to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales+fp8_next_inv_scales

    Tensor.realize(*to_realize)
    return extra[-1]
@@ -88,12 +90,16 @@ class GradAccClipAdamW(Optimizer):
      return out.shard_like(t) if offloaded else out
    if t.dtype in dtypes.fp8s:
      from examples.mlperf.models.flat_llama import FP8_MAX
-      amax = new_w.float().abs().max(axis=tuple(range(1, new_w.ndim))).detach()  # per-layer amax for (n_layers, out, in)
-      scale = FP8_MAX / (amax + 1e-8)
-      fp8_w = (new_w * scale.reshape(-1, *([1]*(new_w.ndim-1)))).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype)
-      if hasattr(t, '_inv_scale'):
-        inv = ((amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
-        t._inv_scale.assign(inv.shard_like(t._inv_scale) if offloaded else inv)
-      return fp8_w.shard_like(t) if offloaded else fp8_w
+      # delayed scaling: reuse previous step's inv_scale
+      t._inv_scale.assign(t._next_inv_scale)
+      inv_scale = t._inv_scale.to(new_w.device) if offloaded else t._inv_scale
+      scale = inv_scale.reciprocal().reshape(-1, *([1]*(new_w.ndim-1)))
+      scaled = (new_w * scale).clamp(-FP8_MAX, FP8_MAX)
+      ret = scaled.cast(t.dtype)
+      # update inv_scale for next step from quantized result
+      new_amax = (ret.float().abs().max(axis=tuple(range(1, ret.ndim))) * inv_scale * FP8_AMAX_MARGIN).detach()
+      new_inv = ((new_amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
+      t._next_inv_scale.assign(new_inv.shard_like(t._next_inv_scale) if offloaded else new_inv)
+      return ret.shard_like(t) if offloaded else ret
    out = new_w.cast(t.dtype)
    return out.shard_like(t) if offloaded else out
@@ -11,6 +11,7 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
@@ -22,7 +23,6 @@ export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
 export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
 export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
 export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0}
 export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

@@ -43,7 +43,7 @@ export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGR

 export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=2
+  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
 fi

 python3 examples/mlperf/model_train.py
@@ -11,6 +11,7 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
@@ -22,7 +23,6 @@ export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
 export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
 export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
 export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0}
 export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

@@ -24,6 +24,7 @@ export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
 export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
 export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export SPLIT_W13=${SPLIT_W13:-0}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
@@ -47,7 +48,7 @@ export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGR

 export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=2
+  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
 fi

 python3 examples/mlperf/model_train.py
@@ -11,6 +11,7 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
@@ -47,7 +48,7 @@ export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGR

 export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=2
+  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
 fi

 python3 examples/mlperf/model_train.py
@@ -24,6 +24,7 @@ export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
 export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
 export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export SPLIT_W13=${SPLIT_W13:-0}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
@@ -11,6 +11,7 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
@@ -23,7 +23,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li

  def name_of(bu:UOp, is_out:bool) -> str:
    nonlocal n
-    if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg), f"input{bu.arg}", prod(bu.shape)*bu.dtype.itemsize
+    if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg.slot), f"input{bu.arg.slot}", prod(bu.shape)*bu.dtype.itemsize
    else:
      b = bu.buffer
      key, size = (id(b.base), b.offset, b.size, b.dtype), b.size*b.dtype.itemsize
@@ -1,40 +1,37 @@
 from __future__ import annotations
 from typing import cast, Callable, TypeVar, Generic, Any, TYPE_CHECKING
-import struct, functools, time, collections
+import struct, functools, time, collections, importlib, itertools, weakref
 from dataclasses import replace
 if TYPE_CHECKING: from tinygrad.engine.realize import ExecContext
-from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, mv_address, round_up, DEBUG, dedup
+from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, mv_address, round_up, DEBUG, dedup, pluralize
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer
-from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites
+from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp
 from tinygrad.uop.symbolic import symbolic_simple, symbolic
 from tinygrad.dtype import dtypes, DType
 from dataclasses import dataclass, field
 from tinygrad.runtime.support.memory import BumpAllocator
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.renderer import Renderer, Estimates
-from tinygrad.engine.realize import to_program, track_stats, get_call_arg_uops, resolve_params
+from tinygrad.engine.realize import to_program, track_stats, get_call_arg_uops, resolve_params, pm_flatten_linear
+from tinygrad.engine.jit import DepsTracker

 HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')

 class HCQ2Compiled(Compiled):
-  """
-  A base class for devices compatible with the HCQ (Hardware Command Queue) API.
-  """
-  timestamp_divider: float = 1000.0  # GPU timestamp counter ticks per microsecond; override per device
+  timestamp_divider: float = 1000.0 # GPU timestamp counter ticks per microsecond; override per device

-  def __init__(self, device:str, allocator:'HCQAllocator', compilers:list[type[Renderer]], runtime,
-               kernargs_size=(16 << 20), can_recover:bool=False, arch=None):
+  def __init__(self, device:str, allocator:'HCQAllocator', compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None):
    self.device_id:int = int(device.split(":")[1]) if ":" in device else 0

-    from extra.hcq2.graph.hcq import HCQ2Graph
-    super().__init__(device, allocator, compilers, lambda *a, **kw: None, HCQ2Graph, arch=arch)
+    # default pm bufferize
+    self.pm_bufferize = PatternMatcher([
+      (UPat(Ops.BUFFER, tag="timeline_signal"), lambda ctx: ctx.timeline_signal),
+      (UPat(Ops.BUFFER, tag="timeline_value"), lambda ctx: ctx.timeline_value),
+      (UPat(Ops.BUFFER, name="b"), lambda ctx, b:
+        Buffer(ctx.device, b.arg, b.dtype, options=BufferSpec(host=True, uncached=True, cpu_access=True, nolru=True))), # TODO: remove nolru
+    ])

-    self.kernargs_size = kernargs_size
-    self.kernargs_offset_allocator:BumpAllocator = BumpAllocator(kernargs_size, wrap=True)
-
-  @functools.cached_property
-  def kernargs_buf(self) -> Buffer:
-    return Buffer(self.device, self.kernargs_size, dtypes.uint8, options=BufferSpec(cpu_access=True), preallocate=True)
+    super().__init__(device, allocator, compilers, lambda *a, **kw: None, None, arch=arch)

  @functools.cached_property
  def timeline_signal(self) -> Buffer:
@@ -50,6 +47,16 @@ class HCQ2Compiled(Compiled):
    buf.as_memoryview(force_zero_copy=True).cast('Q')[0] = 1
    return buf

+  @functools.cache
+  def queue_timeline_signal(self, queue:str) -> Buffer:
+    return Buffer(self.device, 0x100, dtypes.uint8, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
+
+  @functools.cache
+  def queue_timeline_value(self, queue:str) -> Buffer:
+    buf = Buffer("CPU", 1, dtypes.uint64, preallocate=True)
+    buf.as_memoryview(force_zero_copy=True).cast('Q')[0] = 1
+    return buf
+
  def synchronize(self, timeout:int|None=None):
    if not hasattr(self, 'iface'): return
    sig = self.timeline_signal._buf.cpu_view().mv.cast('Q')
@@ -60,14 +67,6 @@ class HCQ2Compiled(Compiled):

  def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.

-  def _realloc(self, oldbuf:HCQ2Buffer|None, new_size:int, options:BufferSpec|None=None, force=False) -> tuple[HCQ2Buffer, bool]:
-    if oldbuf is not None: self.allocator.free(oldbuf, oldbuf.size, options=options)
-    try: buf, realloced = self.allocator.alloc(new_size, options=options), True
-    except MemoryError:
-      if force: raise
-      buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False
-    return buf, realloced
-
  def count(self) -> int: return self.iface.count if hasattr(self, 'iface') else 1

  def _select_iface(self):
@@ -111,6 +110,7 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):

  @suppress_finalizing
  def _free(self, buf:HCQ2Buffer, options:BufferSpec|None=None):
+    self.dev.synchronize()
    if options is not None and options.external_ptr is not None: return
    if hasattr(self, '_do_free'): self._do_free(buf, options)

@@ -126,7 +126,7 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
  def _copy(self, dst:Buffer, src:Buffer):
    from tinygrad.engine.realize import run_linear
    su = UOp.from_buffer(src)
-    run_linear(UOp(Ops.LINEAR, dtypes.void, (su.copy_to_device(dst.device).call(UOp.from_buffer(dst), su),)), jit=True, update_stats=False)
+    run_linear(UOp(Ops.LINEAR, dtypes.void, (su.copy_to_device(dst.device).call(UOp.from_buffer(dst), su),)), update_stats=False)

  def _copyin(self, dest:HCQ2Buffer, src:memoryview):
    s = Buffer(self.dev.device, len(src), dtypes.uint8, options=BufferSpec(host=True), preallocate=True)
@@ -141,33 +141,18 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):

  # def _as_buffer(self, buf): return buf.cpu_view().mv

-# **************** lower context ****************
-
 def unwrap_after(uop):
  while uop.op is Ops.AFTER: uop = uop.src[0]
  return uop

-@dataclass
-class HCQ2DeviceCtx:
-  device:str                       # device name; resolve to instance via Device[device]
-  kernargs_host:UOp                # UOp whose .buffer is dev.kernargs_buf (BUFFER UOp in runtime, PARAM in graph)
-  kernargs_gpu:UOp                 # va_addr const of dev.kernargs_buf
-  kernargs_allocator:BumpAllocator = field(default_factory=lambda: BumpAllocator(2 << 20, wrap=False))
-
-@dataclass
-class HCQ2LowerCtx:
-  name:str
-  inputs:list[Buffer|MultiBuffer] = field(default_factory=list)
-  holds:list[UOp] = field(default_factory=list)
-  dev_ctx:dict[str, HCQ2DeviceCtx] = field(default_factory=dict)
-  addr_table:UOp|None = None
-  next_slot:int = 0
+def make_mstack(uops): return uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, tuple(uops))

 class HCQEncoder:
  def __init__(self): self.blob, self.patches = b'', []

  def get_dev_addr(self, uop:UOp) -> UOp:
-    return UOp(Ops.GETADDR, dtypes.uint64, src=(uop,)) if unwrap_after(uop).op in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT) else uop
+    if unwrap_after(uop).op not in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT): return uop
+    return UOp(Ops.GETADDR, dtypes.uint64, src=(uop, UOp(Ops.DEVICE, arg=self.dev.device)))

  def append(self, *data, dtype=dtypes.uint32):
    for d in data:
@@ -182,302 +167,314 @@ class HCQEncoder:
    buf = UOp.new_buffer(dev, len(self.blob), dtypes.uint8)
    if tag: buf = buf.rtag(tag)
    blob_uop = UOp(Ops.BINARY, dtypes.void, src=(), arg=self.blob)
-    stores = [buf.index(UOp.const(dtypes.int, off)).cast(dt.ptr()).store(val.cast(dt)) for off, val, dt in self.patches]
+    stores = [buf.index(UOp.const(dtypes.int, off), dtype=buf.dtype.ptr()).cast(dt.ptr()).store(val.cast(dt)) for off, val, dt in self.patches]
    return buf.after(buf.store(blob_uop), *stores)

-# **************** prepare runtime ****************
+# *****************
+# 0. helpers

-def _devices(x) -> tuple[str, ...]:
-  return tuple(b.device for b in x.bufs) if isinstance(x, MultiBuffer) else (x.device,) if isinstance(x, Buffer) else x if isinstance(x, tuple) else (x,)
+HCQ_DEVS = frozenset(("AMD",))
+HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",))

-def rebind_program_dev(c:UOp, p:UOp) -> UOp:
-  devs = _devices(c.src[1].buffer)
-  p = p.replace(src=p.src[:1] + (UOp(Ops.DEVICE, arg=devs),) + p.src[2:])
-  return c.replace(src=(Device[devs[0]].pm_lower.rewrite(p),) + c.src[1:])
+def to_tuple(d): return d if isinstance(d, tuple) else (d,)

-def lower_kernargs(call:UOp, prg:UOp) -> UOp:
+def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
+
+# *****************
+# 1.1. prep runtimes: staging copies
+
+def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_devices_in(b.device, HCQ_P2P_DEVS)
+
+def stage_copy(dst:UOp, src:UOp) -> UOp|None:
+  if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
+
+  stage = UOp.new_buffer("CPU", src.buffer.nbytes, dtypes.uint8)
+  return UOp(Ops.LINEAR, dtypes.void, (src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage)))
+pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])
+
+# *****************
+# 1.2. prep runtimes: programs/kernargs
+
+@functools.cache
+def get_pm_prep_program(name:str) -> PatternMatcher|None:
+  try:
+    importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that
+    return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_prep_program
+  except ImportError: return None
+
+def prep_program(call:UOp, prg:UOp) -> UOp|None:
+  dev = call.src[1].device
+  if (pm:=get_pm_prep_program(to_tuple(dev)[0].split(":")[0])) is None or (lowered:=pm.rewrite(prg)) is None: return None
+
+  data, image_bytes = lowered
+  buf = UOp.new_buffer(dev, len(image_bytes), dtypes.uint8).rtag("program")
+  blob = UOp(Ops.BINARY, dtypes.void, src=(), arg=image_bytes)
+  return call.replace(src=(prg.replace(src=(buf.after(buf.store(blob)),), arg=(data, prg.arg)),) + call.src[1:])
+
+def prep_kernargs(call:UOp, prg:UOp) -> UOp:
  data, info = prg.arg
-  enc = HCQEncoder()
-  for gi in info.globals: enc.append(call.src[1+gi], dtype=dtypes.uint64)
-  for v in info.vars: enc.append(v, dtype=dtypes.uint32)
+  patches = [(i*dtypes.uint64.itemsize, UOp(Ops.GETADDR, dtypes.uint64, src=(call.src[1+gi], UOp(Ops.DEVICE, arg=call.src[1+gi].device))),
+              dtypes.uint64) for i,gi in enumerate(info.globals)] \
+          + [(len(info.globals)*dtypes.uint64.itemsize + i*dtypes.uint32.itemsize, v, dtypes.uint32) for i,v in enumerate(info.vars)]
+
+  buf = UOp.new_buffer(call.src[1].device, data.kernargs_alloc_size, dtypes.uint8).rtag("kernargs")
+  kernargs = buf.after(*tuple(buf.index(UOp.const(dtypes.int, o), dtype=buf.dtype.ptr()).cast(dt.ptr()).store(val.cast(dt)) for o, val, dt in patches))

-  enc.blob += b'\x00' * (data.kernargs_alloc_size - len(enc.blob)) # pad blob
-  kernargs = enc.uop(_devices(call.src[1].buffer), tag="kernargs")
  return call.replace(src=(prg.replace(src=prg.src + (kernargs,), arg=(data, info)),) + call.src[1:])

 pm_prep_runtime = PatternMatcher([
  # bind generic PROGRAM device to the call's actual dev(s), then run device-specific lowering
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(), UPat(), UPat(Ops.BINARY)), name="p"),), name="c", allow_any_len=True),
-    rebind_program_dev),
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"),),
+    name="call", allow_any_len=True), prep_program),

  # lower kernargs (PROGRAM.src[0] is now AFTER(BUFFER, COPY) — the lowered program image)
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER),), name="prg"),), name="call", allow_any_len=True), lower_kernargs),
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER),), name="prg"),), name="call", allow_any_len=True), prep_kernargs),
 ])

-# **************** lower ops ****************
+# *****************
+# 2.1. lowering to hcq ir

 def lower_program(call:UOp, prg:UOp) -> UOp:
-  q = UOp(Ops.LINEAR, dtypes.void, (prg,), arg=(_devices(call.src[1].buffer), "COMPUTE"))
-  return UOp(Ops.LINEAR, dtypes.void, (q,), tag=call.tag)
+  q = UOp(Ops.LINEAR, dtypes.void, (prg,), arg=(call.src[1].device, "COMPUTE:0"))
+  return call.replace(src=(q,) + call.src[1:]).rtag('hcq')

-def lower_copy(call:UOp, copy:UOp) -> UOp:
+def lower_copy(call:UOp, copy:UOp) -> UOp|None:
  dst, src = call.src[1], call.src[2]
-  q = UOp(Ops.LINEAR, dtypes.void, (UOp(Ops.COPY, dtypes.void, src=(dst, src), arg=src.buffer.nbytes),), arg=(_devices(dst.buffer), "COPY"))
-  return UOp(Ops.LINEAR, dtypes.void, (q,), tag=call.tag)
+  if (hcq_dev:=next((b.device for b in (dst, src) if b.device.split(":")[0] in HCQ_DEVS), None)) is None: return None
+
+  q = UOp(Ops.LINEAR, dtypes.void, (UOp(Ops.COPY, dtypes.void, src=(dst, src), arg=src.buffer.nbytes),), arg=(hcq_dev, "COPY:0"))
+  return call.replace(src=(q,) + call.src[1:]).rtag('hcq')

 pm_lower_ops = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER), UPat()), name="prg"),), name="call", allow_any_len=True), lower_program),
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER), UPat(Ops.AFTER)), name="prg"),), name="call", allow_any_len=True), lower_program),
  (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), lower_copy),
 ])

-def split_into_queues(outer:UOp) -> UOp:
-  groups:dict[tuple, list[UOp]] = collections.defaultdict(list)
-  for child in outer.src:
-    wrapper = child.src[0] if child.op is Ops.AFTER else child
-    for q in wrapper.src: groups[q.arg].extend(q.src)
-  return outer.replace(src=tuple(UOp(Ops.LINEAR, dtypes.void, tuple(cmds), arg=k) for k, cmds in groups.items()))
-pm_split_into_queues = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR, src=UPat(Ops.LINEAR)).or_after(), name="outer"), split_into_queues)])
+# *****************
+# 2.2. deps tracking
+# device.timeline_signal/value are the per-device schedule epoch. Before a schedule queue accesses memory owned by device N for the first time,
+# it waits for device[N].timeline_signal >= device[N].timeline_value - 1. This orders the schedule after all prior schedules that touched device N.
+#
+# queue.timeline_signal/value are per-queue progress counters used only inside a schedule.
+# Only the owner queue signals its queue.timeline_signal. Values are monotonic.
+#
+# At schedule end, one finalizer queue per touched device[N] waits for every active queue on device[N] to reach its schedule-local
+# final queue.timeline value, then signals device[N].timeline_signal with the schedule's reserved device epoch. After that, buffers/transients
+# for device N from this schedule are safe for the next schedule
+#
+# C programs reserve and bump timeline values, then patch command buffers with the concrete wait/signal values.

-def add_signals(q:UOp) -> UOp:
-  sig = UOp.new_buffer(q.arg[0], 0x100, dtypes.uint8).rtag("timeline_signal")
-  tl = UOp.new_buffer(q.arg[0], 1, dtypes.uint64).rtag("timeline_value").index(UOp.const(dtypes.int, 0))
-  return q.replace(src=(sig.wait(tl-1), *q.src, sig.store(tl)), arg=q.arg)
-pm_add_signals = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"),
-  lambda outer: outer.replace(src=tuple(add_signals(q) for q in outer.src)))])
+@dataclass
+class DepsCtx:
+  deps:DepsTracker = field(default_factory=DepsTracker)
+  evid:itertools.count = field(default_factory=lambda: itertools.count(0))
+  last_per_queue:weakref.WeakValueDictionary[tuple[Any, str], UOp] = field(default_factory=weakref.WeakValueDictionary)

-pm_add_barriers = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"),
-  lambda outer: outer.replace(src=tuple(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), *q.src)) for q in outer.src)))])
+def get_writes_ids(call:UOp) -> tuple[int, ...]:
+  ast, writes = call.src[0].src[0], set()
+  for ast in call.src[0].src:
+    if ast.op is Ops.PROGRAM: writes.update(ast.arg[1].outs)
+    elif ast.op in (Ops.COPY, Ops.SLICE, Ops.CUSTOM_FUNCTION): writes.add(0)
+  return tuple(writes)

-def add_timeline_inc(q:UOp) -> UOp:
-  tl = UOp.new_buffer(q.arg[0], 1, dtypes.uint64).rtag("timeline_value")
-  done = tl.after(UOp(Ops.BARRIER, dtypes.void, src=(q,)))
-  return done.index(UOp.const(dtypes.int, 0), dtype=tl.dtype.ptr()).store(tl.index(UOp.const(dtypes.int, 0)) + 1)
-pm_add_timeline_inc = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"),
-  lambda outer: outer.replace(src=tuple(add_timeline_inc(q) for q in outer.src)))])
+def insert_deps(ctx:DepsCtx, call:UOp) -> UOp|None:
+  q, refs, write = call.src[0].rtag(next(ctx.evid)), [b.buffer for b in get_call_arg_uops(call)], get_writes_ids(call)
+  if q.arg not in ctx.last_per_queue:
+    sig = UOp.new_buffer(q.arg[0], 0x100, dtypes.uint8).rtag("timeline_signal")
+    tl = UOp.new_buffer(q.arg[0], 1, dtypes.uint64).rtag("timeline_value").index(UOp.const(dtypes.int, 0))
+    q = q.replace(src=(sig.wait(tl - 1), *q.src))
+  ctx.last_per_queue[q.arg] = q

-# **************** build host program ****************
+  deps = []
+  for lane in range(len(refs[0].bufs) if isinstance(refs[0], MultiBuffer) else 1):
+    deps += ctx.deps.access_resources([b.bufs[lane] if isinstance(b, MultiBuffer) else b for b in refs], write, q)
+  return call.replace(src=(q.after(*dps).rtag("deps") if (dps:=dedup(deps)) else q,) + call.src[1:])
+pm_insert_deps = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call", allow_any_len=True), insert_deps)])

-def calc_kernargs_sizes(ctx:dict[str,int], u:UOp) -> None:
-  if u.tag != "kernargs": return
-  for d in _devices(u.src[1].arg): ctx[d] = ctx.get(d, 0) + round_up(u.arg, 16)
-pm_calc_kernargs_sizes = PatternMatcher([(UPat(Ops.BUFFER, name="u"), calc_kernargs_sizes)])
+def make_finalizer(devs:tuple[str, ...], queues:list[UOp], nbump:int) -> UOp:
+  sig = UOp.new_buffer(devs, 0x100, dtypes.uint8).rtag("timeline_signal")
+  tl = UOp.new_buffer(devs, 1, dtypes.uint64).rtag("timeline_value")
+  q = UOp(Ops.LINEAR, dtypes.void, (sig.store(tl.index(UOp.const(dtypes.int, 0))),), arg=(devs, "COMPUTE:0"), tag="finalizer")

-# bufferize
+  def bump(b, by): return b.index(UOp.const(dtypes.int, 0), dtype=b.dtype.ptr()).store(b.index(UOp.const(dtypes.int, 0)) + by)
+  bumps = (bump(tl, 1),) + tuple(bump(UOp.new_buffer(devs, 1, dtypes.uint64).rtag((ty, "timeline_value")), nbump) for ty in dedup([q.arg[1] for q in queues]))
+  return UOp(Ops.CALL, dtypes.void, (q.after(*bumps).after(*queues).rtag("deps"),), tag="hcq")

-def _maybe_mstack(srcs:tuple[UOp, ...], tag=None) -> UOp: return srcs[0] if len(srcs) == 1 else UOp(Ops.MSTACK, srcs[0].dtype, srcs, tag=tag)
+def add_finalizer(ctx:DepsCtx, linear:UOp) -> UOp:
+  fams:dict[str, list[UOp]] = collections.defaultdict(list)
+  for q in ctx.last_per_queue.values(): fams[to_tuple(q.arg[0])[0].split(":")[0]].append(q)

-def _lower_stores(host_buf:UOp, buf_node:UOp, stores:tuple[UOp, ...]) -> list[UOp]:
-  def lower(s:UOp) -> UOp:
-    if s.src[1].op is Ops.BINARY: return s.substitute({buf_node: host_buf})
-    idx = s.src[0].src[0]
-    return s.substitute({idx: host_buf.index(UOp.const(dtypes.int, idx.src[1].arg // host_buf.dtype.base.itemsize), dtype=host_buf.dtype.ptr())})
-  return [lower(s) for s in stores]
+  nbump = next(ctx.evid)
+  finalizers = []
+  for queues in fams.values():
+    devs = tuple(sorted(dedup(d for q in queues for d in to_tuple(q.arg[0]))))
+    finalizers.append(make_finalizer(devs, queues, nbump))
+  return linear.replace(src=linear.src + tuple(finalizers))

-_program_uop_cache:dict[tuple[bytes,str], tuple[UOp,UOp]] = {}
-def bufferize_program(ctx:HCQ2LowerCtx, target:UOp, buf_node:UOp) -> UOp:
-  blob, addrs = target.src[1].src[1].arg, []
-  for dev in _devices(buf_node.src[1].arg):
-    if (cached:=_program_uop_cache.get((blob, dev))) is None:
-      lib_gpu = Buffer(dev, round_up(len(blob), 0x1000), dtypes.uint8, options=BufferSpec(nolru=True, cpu_access=True), preallocate=True)
-      lib_gpu._buf.cpu_view()[:len(blob)] = memoryview(blob)
-      cached = _program_uop_cache[(blob, dev)] = (UOp.from_buffer(lib_gpu, dev), UOp.const(dtypes.uint64, lib_gpu._buf.va_addr, device=dev))
-    if cached[0] not in ctx.holds: ctx.holds.append(cached[0])
-    addrs.append(cached[1])
-  return _maybe_mstack(tuple(addrs))
+def add_loads(ctx:set[int], call:UOp, after:UOp) -> UOp:
+  q = unwrap_after(after.src[0])
+  cur_devs = to_tuple(q.arg[0])

-def bufferize_kernargs(ctx:HCQ2LowerCtx, target:UOp, buf_node:UOp) -> UOp:
-  hbufs, addrs = [], []
-  for dev in _devices(buf_node.src[1].arg):
-    dctx = ctx.dev_ctx[dev]
-    isz = dctx.kernargs_host.dtype.base.itemsize
-    off = dctx.kernargs_allocator.alloc(buf_node.arg, 16)
-    hbufs.append(UOp(Ops.SLICE, dctx.kernargs_host.dtype,
-                     src=(dctx.kernargs_host, UOp.const(dtypes.weakint, off // isz)), arg=buf_node.arg // isz))
-    addrs.append(dctx.kernargs_gpu + off)
-  return _maybe_mstack(tuple(addrs)).after(*_lower_stores(_maybe_mstack(tuple(hbufs)), buf_node, target.src[1:]))
+  waits = []
+  for dq in [unwrap_after(dep) for dep in after.src[1:]]:
+    ctx.add(dq.tag)
+    dq_devs = to_tuple(dq.arg[0])
+    sigs = [UOp.new_buffer(d, 0x100, dtypes.uint8).rtag((dq.arg[1], "timeline_signal") if d in dq_devs else "max_sentinel_signal") for d in cur_devs]
+    orig_val = UOp.new_buffer(cur_devs, 1, dtypes.uint64).rtag((dq.arg[1], "timeline_value")).index(UOp.const(dtypes.int, 0))
+    waits.append(make_mstack(sigs).wait(orig_val + dq.tag))
+  return call.replace(src=(after.src[0].substitute({q: q.replace(src=(*waits, *q.src))}),) + call.src[1:])
+pm_add_loads = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.AFTER, tag="deps", name="after"),), name="call", allow_any_len=True), add_loads)])

-def bufferize_cmdbuf(ctx:HCQ2LowerCtx, target:UOp, buf_node:UOp) -> UOp:
-  hbufs = tuple(UOp.from_buffer(Buffer("CPU", buf_node.arg // dtypes.uint32.itemsize, dtypes.uint32,
-                                       options=BufferSpec(cpu_access=True, nolru=True), preallocate=True), dev)
-                for dev in _devices(buf_node.src[1].arg))
-  hbuf = _maybe_mstack(hbufs)
-  return hbuf.after(*_lower_stores(hbuf, buf_node, target.src[1:]), tag=buf_node.tag)
+def add_stores(ctx:set[int], call:UOp) -> UOp|None:
+  if (q:=unwrap_after(call.src[0])).tag not in ctx: return None
+  sig = UOp.new_buffer(q.arg[0], 0x100, dtypes.uint8).rtag((q.arg[1], "timeline_signal"))
+  val = UOp.new_buffer(q.arg[0], 1, dtypes.uint64).rtag((q.arg[1], "timeline_value")).index(UOp.const(dtypes.int, 0))
+  newq = q.replace(src=q.src + (sig.store(val + q.tag),))
+  return call.replace(src=(call.src[0].substitute({q: newq}),) + call.src[1:])
+pm_add_stores = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call", allow_any_len=True), add_stores)])

-def bufferize_binary(ctx:HCQ2LowerCtx, target:UOp, buf_node:UOp) -> UOp|None:
-  if buf_node.tag == "program": return bufferize_program(ctx, target, buf_node)
-  if buf_node.tag == "kernargs": return bufferize_kernargs(ctx, target, buf_node)
-  if buf_node.tag in ("compute", "copy"): return bufferize_cmdbuf(ctx, target, buf_node)
-  return None
+# *****************
+# 2.3. barriers / signals / timeline inc

-# TODO: merge with bufferize_binary
-def resolve_buffer(b:UOp) -> UOp|None:
-  devs = _devices(b.src[1].arg)
-  if b.tag in ("timeline_signal", "timeline_value"):
-    return _maybe_mstack(tuple(UOp.from_buffer(getattr(Device[d], b.tag), d) for d in devs), b.tag)
-  if b.tag == "scratch":
-    return _maybe_mstack(tuple(UOp.from_buffer(Buffer(d, (s:=Device[d].scratch).size, dtypes.uint8, opaque=s, options=BufferSpec(external_ptr=1)), d)
-                               for d in devs), b.tag)
-  if isinstance(b.tag, tuple): # (compute_queue|sdma_queue, ring|write_ptr|doorbell|put_value)
-    return _maybe_mstack(tuple(UOp.from_buffer(getattr(Device[d].compute_queue if b.tag[0] == "compute_queue" else Device[d].sdma_queue(0), b.tag[1]), d)
-                            for d in devs), b.tag)
-  if isinstance(b.device, tuple): return _maybe_mstack(tuple(UOp.from_buffer(buf, buf.device) for buf in b.buffer.bufs))
-  return None
+def add_barriers(call:UOp, q:UOp) -> UOp:
+  return call.replace(src=(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), *q.src)),) + call.src[1:])
+pm_add_barriers = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), add_barriers)])

-pm_bufferize = PatternMatcher([
-  (UPat(Ops.AFTER, src=(UPat(Ops.BUFFER, name="buf_node"),), allow_any_len=True, name="target"), bufferize_binary),
-  (UPat(Ops.BUFFER, name="b"), resolve_buffer), # TODO: cleanup
+# *****************
+# 3.1. encode cmdbufs
+
+@functools.cache
+def get_pm_lower(name:str) -> PatternMatcher|None:
+  try:
+    importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that
+    return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_lower
+  except ImportError: return None
+
+def encode_cmdbuf(call:UOp) -> UOp|None:
+  if (q:=unwrap_after(call.src[0])).op is not Ops.LINEAR: return None
+  if (pm:=get_pm_lower(to_tuple(q.arg[0])[0].split(":")[0])) is None or (encoded:=pm.rewrite(call.src[0])) is None: return None
+  return call.replace(src=(encoded,) + call.src[1:])
+pm_encode_cmdbufs = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call", allow_any_len=True), encode_cmdbuf)])
+
+pm_compose_submit = PatternMatcher([
+  (UPat(Ops.CALL, tag="hcq", src=(UPat(Ops.CUSTOM_FUNCTION, arg="submit", name="sub"),), allow_any_len=True, name="call"),
+    lambda call, sub: call.replace(src=(UOp.group(*sub.src),) + call.src[1:])),
 ])

-def lift_patches_to_cmdbuf(ctx:HCQ2LowerCtx, cmdbuf:UOp) -> UOp|None:
-  if cmdbuf.tag not in ("compute", "copy"): return None
-  patches = dedup(u for store in cmdbuf.src[1:] for u in store.toposort() if u.op is Ops.AFTER)
+# *****************
+# 3.2. add timeline inc
+
+def add_timeline_inc(call:UOp, s:UOp) -> UOp:
+  tl = UOp.new_buffer(s.device, 1, dtypes.uint64).rtag("timeline_value")
+  return call.replace(src=(tl.after(s).index(UOp.const(dtypes.int, 0), dtype=tl.dtype.ptr()).store(tl.index(UOp.const(dtypes.int, 0)) + 1),) + call.src[1:])
+pm_add_timeline_inc = PatternMatcher([(UPat(Ops.CALL, tag="hcq", src=(UPat(name="s"),), name="call", allow_any_len=True), add_timeline_inc)])
+
+# *****************
+# 3.3. lift patches to the command buffer (root)
+
+def lift_patches_to_cmdbuf(cmdbuf:UOp) -> UOp|None:
+  if not (patches:=dedup(u for store in cmdbuf.src[1:] for u in store.toposort() if u.op is Ops.AFTER)): return None
  deps = tuple(d for p in patches for d in p.src[1:])
-  return cmdbuf.replace(src=cmdbuf.src+deps, tag=None).substitute({p:p.src[0] for p in patches})
-pm_lift_patches_to_cmdbuf = PatternMatcher([(UPat(Ops.AFTER, name="cmdbuf", allow_any_len=True), lift_patches_to_cmdbuf)])
+  return cmdbuf.replace(src=cmdbuf.src + deps).substitute({p: p.src[0] for p in patches})
+pm_lift_patches_to_cmdbuf = PatternMatcher([
+  (UPat(Ops.AFTER, src=(UPat(Ops.BUFFER, tag={"compute", "copy"}),), allow_any_len=True, name="cmdbuf"), lift_patches_to_cmdbuf),
+])

-# resolve patches
+# *****************
+# 4. bufferize placeholders: replace placeholders with real buffers.

-def fold_const_store(ctx:HCQ2LowerCtx, buf:UOp, off:UOp, val:UOp) -> UOp:
-  bufs = buf.src if buf.op is Ops.MSTACK else (buf,)
-  vals = val.src if val.op is Ops.MSTACK else (val,) * len(bufs)
-  for b, v in zip(bufs, vals):
+def bufferize_buf(buf:UOp) -> UOp|None:
+  if buf.tag is None: return None
+  uops = tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=dv), dev) for dev in to_tuple(buf.src[1].arg))
+  return make_mstack(uops)
+pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, name="buf"), bufferize_buf)])
+
+# *****************
+# 5.1. capture buffers reachable from each hcq call as BIND, so we don't drop their refs
+
+def hold_call_buffers(call:UOp) -> UOp|None:
+  if not (bufs:=tuple(dedup(u for u in call.src[0].toposort() if u.op is Ops.BUFFER and u not in call.src))): return None
+  return call.replace(src=call.src + (UOp(Ops.BIND, dtypes.void, src=bufs),))
+pm_hold_call_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), hold_call_buffers)])
+
+# *****************
+# 5.2. resolve patches
+
+def push_stack(op, s): return UOp(Ops.STACK, op.dtype.scalar().vec(len(s.src)),
+  tuple(op.replace(dtype=op.dtype.scalar(), src=tuple(x if y is s else y for y in op.src)) for x in s.src))
+
+def fold_blob_store(buf:UOp, blob:UOp) -> UOp:
+  for b in (buf.src if buf.op is Ops.MSTACK else (buf,)): b.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B')[:len(blob.arg)] = blob.arg
+  return UOp(Ops.NOOP)
+
+def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
+  for b, v in zip((buf.src if buf.op is Ops.MSTACK else (buf,)), (val.src if val.op is Ops.STACK else (val,))):
    struct.pack_into(f'<{v.dtype.fmt}', b.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * b.dtype.base.itemsize, v.arg)
  return UOp(Ops.NOOP)

-def fold_blob_store(ctx:HCQ2LowerCtx, buf:UOp, blob:UOp) -> UOp:
-  for b in (buf.src if buf.op is Ops.MSTACK else (buf,)):
-    b.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B')[:len(blob.arg)] = blob.arg
-  return UOp(Ops.NOOP)
+def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
+  if isinstance(b:=buf.buffer, Buffer): return UOp.const(dtypes.uint64, b.get_buf(g.src[1].arg).va_addr)
+  return UOp(Ops.STACK, dtypes.uint64.vec(len(b.bufs)), tuple(UOp.const(dtypes.uint64, x.ensure_allocated()._buf.va_addr) for x in b.bufs))

-def resolve_getaddr(ctx:HCQ2LowerCtx, m:UOp) -> UOp:
-  srcs = m.src if m.op is Ops.MSTACK else (m,)
-  for s in srcs:
-    if s.op in (Ops.BUFFER, Ops.SLICE) and s not in ctx.holds: ctx.holds.append(s)
-  addrs = [s.arg if s.op is Ops.CONST else s.buffer.get_buf(s.device).va_addr for s in srcs]
+pm_resolve_patches = PatternMatcher([
+  # multi
+  (UPat(GroupOp.ALU, src=[UPat(Ops.STACK, name="s"), UPat(Ops.CONST)], name="op"), push_stack),
+  (UPat(Ops.CAST, src=(UPat(Ops.STACK, name="s"),), name="op"), push_stack),

-  # fast-path: all per-dev VAs equal -> just a const
-  if all(v == addrs[0] for v in addrs): return UOp.const(dtypes.uint64, addrs[0])
+  # getaddr
+  (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"), UPat(Ops.DEVICE, name="dev"))), # getaddr(slice(x)) -> offset+getaddr(x)
+    lambda bv, dev: UOp(Ops.GETADDR, dtypes.uint64, src=(bv.src[0], dev)) + UOp.const(dtypes.uint64, bv.src[1].arg * bv.src[0].dtype.itemsize)),
+  (UPat(Ops.GETADDR, src=(UPat({Ops.BUFFER, Ops.MSTACK, Ops.MSELECT}, name="buf"), UPat(Ops.DEVICE)), name="g"), resolve_getaddr),

-  table = _maybe_mstack(tuple(UOp.from_buffer(Buffer("CPU", 1, dtypes.uint64, preallocate=True), "CPU") for _ in range(len(srcs))))
-  vas = _maybe_mstack(tuple(UOp.const(dtypes.uint64, va) for va in addrs))
-  patch = table.index(slot_const:=UOp.const(dtypes.int, 0), dtype=table.dtype.ptr()).store(vas)
-  return table.after(patch).index(slot_const, dtype=table.dtype.ptr()).load(dtype=dtypes.uint64)
+  # folders
+  (UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").store(UPat(Ops.BINARY, name="blob")), fold_blob_store),
+  (UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").index(UPat.cvar("off")).or_casted().store(UPat.any(UPat.cvar("val"), UPat(Ops.STACK, name="val"))),
+    fold_const_store),
+]) + symbolic_simple

-pm_resolve_patches = symbolic + PatternMatcher([
-  # resolve getaddrs
-  (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"),)), # getaddr(buffer_view(x)) -> offset+getaddr(x)
-    lambda ctx, bv: UOp(Ops.GETADDR, dtypes.uint64, src=(bv.src[0],)) + UOp.const(dtypes.uint64, bv.src[1].arg * bv.src[0].dtype.itemsize)),
-  (UPat(Ops.GETADDR, src=(UPat((Ops.BUFFER, Ops.MSTACK), name="m"),)), resolve_getaddr), # getaddr(buffer|mstack) -> addr_table load|const
-  (UPat(Ops.GETADDR, src=(UPat.cvar("const"),)), lambda ctx, const: const), # getaddr(const) -> const
+# *****************
+# 6. callify hcq programs

-  # write consts and binaries directly into the buffer (BUFFER or MSTACK of BUFFERs)
-  (UPat((Ops.BUFFER, Ops.SLICE, Ops.MSTACK), name="buf").store(UPat(Ops.BINARY, name="blob")), fold_blob_store),
-  (UPat((Ops.BUFFER, Ops.SLICE, Ops.MSTACK), name="buf").index(UPat.cvar("off")).or_casted()
-    .store(UPat.any(UPat.cvar("val"), UPat(Ops.MSTACK, src=UPat.cvar(), name="val"))), fold_const_store),
-])
-
-def parametrize_host_buffer(ctx:HCQ2LowerCtx, buf:UOp) -> UOp:
-  # register a host buffer as a launcher input and return its placeholder
-  if buf.op is Ops.AFTER:
-    p = parametrize_host_buffer(ctx, buf.src[0])
-    return p.after(*(s.substitute({buf.src[0]: p}) for s in buf.src[1:]))
-  if (b:=buf.buffer) not in ctx.inputs: ctx.inputs.append(b)
-  return UOp.placeholder((b.size,), b.dtype, ctx.inputs.index(b))
-
-pm_parametrize_host_buffers = PatternMatcher([
-  # resolve buffer views to parametrize only root buffers
-  (UPat(Ops.INDEX, src=(UPat(Ops.SLICE, name="bv"), UPat.var("idx")), name="bi"),
-    lambda bv, idx, bi: bi.replace(src=(bv.src[0], idx * bv.dtype.itemsize // bv.src[0].dtype.itemsize + bv.src[1].arg))),
-
-  # parametrize host buffers
-  (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.SLICE, Ops.MSTACK)),), allow_any_len=True, name="buf"), parametrize_host_buffer),
-  (UPat((Ops.BUFFER, Ops.SLICE, Ops.MSTACK), name="buf"), parametrize_host_buffer),
-
-  # remove UNIQUE/DEVICE to dedup CONST
+pm_fixup = PatternMatcher([ # TODO: this should gone?
  (UPat(Ops.CONST, name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
 ])

-def hcq_callify(ctx:HCQ2LowerCtx, l:UOp) -> UOp:
-  sink = UOp.sink(*l.src, arg=KernelInfo(name=ctx.name, estimates=Estimates()), tag=1)
-  inputs = [UOp.from_buffer(b, tuple(x.device for x in b.bufs) if isinstance(b, MultiBuffer) else "CPU") for b in ctx.inputs]
-  call = to_program(sink, Device["CPU"].renderer).call(*inputs)
-  return call.replace(src=call.src + (UOp(Ops.BIND, dtypes.void, src=tuple(ctx.holds)),)) if ctx.holds else call
-pm_callify = PatternMatcher([(UPat(Ops.LINEAR, name="l", allow_any_len=True), hcq_callify)])
+def to_param(bufs:list[UOp], ref:UOp) -> UOp:
+  bufs.append(ref)
+  return UOp.placeholder((ref.buffer.size,), ref.dtype, len(bufs)-1)
+pm_to_param = PatternMatcher([(UPat({Ops.MSELECT, Ops.MSTACK, Ops.BUFFER}, name="r"), lambda ctx, r: to_param(ctx, r))])

-# **************** schedule ****************
+def parametrize_host_buffers(call:UOp) -> UOp:
+  body = graph_rewrite(call.src[0], pm_to_param, ctx=(bufs:=[]), bottom_up=True, name="parametrize host buffers")
+  return call.replace(src=(body, *bufs) + call.src[1:], tag="hcq_param")
+pm_parametrize_host_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), parametrize_host_buffers)])

-@track_rewrites(name=lambda linear,ast,**kw: f"hcq schedule {getattr(ast.arg, 'name', ast.op.name.lower())}")
-def hcq_schedule(linear:UOp, ast:UOp) -> UOp:
-  # runtime preparation: device-specific program, kernargs for each program
-  linear = graph_rewrite(linear, pm_prep_runtime, name="hcq: prepare runtime")
+def callify_hcq(call:UOp) -> UOp:
+  sink = UOp.sink(call.src[0], arg=KernelInfo(name="hcq_submit", estimates=Estimates()), tag=1)
+  return to_program(sink, Device["CPU"].renderer).call(*call.src[1:])
+pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, tag="hcq_param", name="call"), callify_hcq)])

-  # lower ops into hcq style per-device operations
-  linear = graph_rewrite(linear, pm_lower_ops, name="hcq: lower ops")
+@track_rewrites(lambda _,ret: f"HCQ Schedule {pluralize('Kernel', len(ret.src))}")
+def hcq_schedule(linear:UOp) -> UOp:
+  linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging")
+  linear = graph_rewrite(linear, pm_prep_runtime, name="prepare runtime")

-  # split ops into logical queues
-  linear = graph_rewrite(linear, pm_split_into_queues, name="hcq: split into queues")
+  linear = graph_rewrite(linear, pm_lower_ops, name="lower ops into hcq ir")
+  linear = graph_rewrite(linear, pm_insert_deps, ctx=(deps_ctx:=DepsCtx()), walk=True, name="insert deps")
+  linear = add_finalizer(deps_ctx, linear)
+  linear = graph_rewrite(linear, pm_add_loads, ctx=(waited:=set()), walk=True, name="add loads")
+  linear = graph_rewrite(linear, pm_add_stores, ctx=waited, walk=True, name="add stores")
+  linear = graph_rewrite(linear, pm_add_barriers, walk=True, name="add barriers")
+  linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs")
+  linear = graph_rewrite(linear, pm_compose_submit, walk=True, name="compose submit")
+  linear = graph_rewrite(linear, pm_lift_patches_to_cmdbuf, name="lift patches to cmdbuf", enter_calls=True)

-  # runtime-specific lowering
-  linear = graph_rewrite(linear, pm_add_barriers, walk=True, name="hcq: add barriers")
-  linear = graph_rewrite(linear, pm_add_signals, walk=True, name="hcq: add signals")
-  linear = graph_rewrite(linear, pm_add_timeline_inc, walk=True, name="hcq: add submit")
+  # realize starts from here
+  linear = graph_rewrite(linear, pm_bufferize, bottom_up=True, name="bufferize placeholders", enter_calls=True)
+  linear = graph_rewrite(linear, pm_hold_call_buffers, walk=True, name="hold call buffers")
+  linear = graph_rewrite(linear, pm_resolve_patches, bottom_up=False, name="simplify patches", enter_calls=True)
+  linear = graph_rewrite(linear, pm_fixup, bottom_up=False, name="fixup", enter_calls=True)
+  linear = graph_rewrite(linear, pm_parametrize_host_buffers, name="parametrize host buffers")
+  linear = graph_rewrite(linear, pm_callify_hcq, name="callify hcq")

-  # encode cmdbuffers + submits
-  # TODO: remove dev
-  dev = Device["AMD"]
-  return graph_rewrite(linear, dev.pm_lower, walk=True, name="hcq: encode cmdbuf")
-
-@track_rewrites(name=lambda ctx,linear,ast,**kw: f"hcq realize {getattr(ast.arg, 'name', ast.op.name.lower())}")
-def hcq_realize(ctx:HCQ2LowerCtx, linear:UOp, ast:UOp) -> UOp:
-  # allocate lowering structs
-  graph_rewrite(linear, pm_calc_kernargs_sizes, ctx=(sizes:={}), name=None)
-
-  for dev_name, sz in sizes.items():
-    dev = Device[dev_name]
-    off = dev.kernargs_offset_allocator.alloc(sz, 16)
-    ctx.dev_ctx[dev_name] = HCQ2DeviceCtx(dev_name, UOp.from_buffer(dev.kernargs_buf.view(sz, dtypes.uint8, off), dev_name),
-                                          UOp.const(dtypes.uint64, dev.kernargs_buf.get_buf(dev_name).va_addr + off, device=dev_name))
-
-  linear = graph_rewrite(linear, pm_bufferize, ctx=ctx, bottom_up=True, name="realize binaries")
-  linear = graph_rewrite(linear, pm_lift_patches_to_cmdbuf, ctx=ctx, bottom_up=False, name="lift patches to cmdbuf")
-  linear = graph_rewrite(linear, pm_resolve_patches, ctx=ctx, bottom_up=False, name="simplify patches")
-  linear = graph_rewrite(linear, pm_parametrize_host_buffers, ctx=ctx, bottom_up=True, name="parametrize host buffers")
-  return graph_rewrite(linear, pm_callify, ctx=ctx, name="hcq: callify")
-
-def ensure_accessible(ctx:HCQ2LowerCtx, call:UOp, copy:UOp) -> UOp|None:
-  src_buf = call.src[2].buffer # TODO: cleanup
-  dev = call.src[1].buffer.device
-  try: src_buf.get_buf(dev)
-  except Exception:
-    (cpubuf := Buffer("CPU", src_buf.nbytes, dtypes.uint8, preallocate=True)).copyin(src_buf.ensure_allocated().as_memoryview())
-    ctx.holds.append(buf_uop:=UOp.from_buffer(cpubuf, dev))
-    return call.replace(src=call.src[:2] + (buf_uop,) + call.src[3:])
-pm_ensure_bufs_accessible = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), ensure_accessible)])
-
-def hcq_exec(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  from tinygrad.engine.realize import run_linear
-
-  if ast.src[1].arg.split(":")[0] != "AMD": return None
-
-  # TODO: this mess should gone
-  resolved_call = call.replace(src=(ast,) + tuple(resolve_params(call, ctx.input_uops)) + tuple(s for s in call.src[1:] if s.op is Ops.BIND))
-  bufs = [cast(Buffer, resolved_call.src[1+gi].buffer) for gi in ast.arg.globals] if ast.op is Ops.PROGRAM \
-    else [cast(Buffer, resolved_call.src[i].buffer) for i in range(1, len(resolved_call.src))]
-  hcq_ctx = HCQ2LowerCtx(name="submit")
-  linear = graph_rewrite(UOp(Ops.LINEAR, dtypes.void, (resolved_call,)), pm_ensure_bufs_accessible, ctx=hcq_ctx)
-
-  linear = hcq_schedule(linear, ast)
-  host_call = hcq_realize(hcq_ctx, linear, ast)
-
-  dev = Device["AMD"]
-  with track_stats(ctx, call, dev.device, bufs, ctx.var_vals) as tm:
-    st = time.perf_counter() if ctx.wait else 0.0
-    run_linear(UOp(Ops.LINEAR, dtypes.void, (host_call,)), var_vals=ctx.var_vals, jit=True, update_stats=DEBUG>=3)
-    if ctx.wait:
-      dev.synchronize()
-      tm[0] = time.perf_counter() - st
-  return tm[0] if tm[0] is not None else 0.0
-
-pm_hcq_exec = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat({Ops.PROGRAM, Ops.COPY}, name="ast"),), name="call", allow_any_len=True), hcq_exec),
-])
+  return linear
@@ -3,7 +3,7 @@ from typing import cast
 import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-from extra.hcq2.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, HCQEncoder
+from extra.hcq2.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, HCQEncoder, to_tuple
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
 from tinygrad.dtype import dtypes
@@ -105,10 +105,9 @@ class AMDComputeQueue(HCQEncoder):

    self.acquire_mem(gli=0, gl2=0)

-    scratch_buf = UOp.new_buffer(self.devs if len(self.devs) > 1 else self.devs[0], self.dev.scratch.size, dtypes.uint8).rtag("scratch")
-    scratch_addr = self.get_dev_addr(scratch_buf)
-
+    scratch_addr = self.get_dev_addr(UOp.new_buffer(self.devs, data.private_segment_size, dtypes.uint8).rtag("scratch"))
    args_addr = self.get_dev_addr(args)
+
    user_regs = []
    if data.enable_private_segment_sgpr:
      scratch_hilo = data64_le(scratch_addr)
@@ -119,10 +118,10 @@ class AMDComputeQueue(HCQEncoder):
    self.wreg(self.gc.regCOMPUTE_PGM_LO, *data64_le(prog_addr >> 8))
    self.wreg(self.gc.regCOMPUTE_PGM_RSRC1, data.rsrc1, data.rsrc2)
    self.wreg(self.gc.regCOMPUTE_PGM_RSRC3, data.rsrc3)
-    self.wreg(self.gc.regCOMPUTE_TMPRING_SIZE, self.dev.tmpring_size)
+    self.wreg(self.gc.regCOMPUTE_TMPRING_SIZE, self.dev.tmpring_size(data.private_segment_size))

    for xcc_id in range(self.dev.xccs):
-      scratch_base = scratch_addr + (self.dev.scratch.size // self.dev.xccs * xcc_id)
+      scratch_base = scratch_addr + (data.private_segment_size // self.dev.xccs * xcc_id)
      self.wreg(self.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, *data64_le(scratch_base >> 8))

    self.wreg(self.gc.regCOMPUTE_RESTART_X, 0, 0, 0)
@@ -153,7 +152,7 @@ def amd_submit_pm4(cmdbuf, devs):

  # the compute queue's ring and its host-side ring/write/put pointers (placeholders, resolved in pm_bufferize)
  q = Device['AMD'].compute_queue
-  ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("compute_queue", name))
+  ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("COMPUTE:0", name))
    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))

  # place the cmdbuf at the ring's write offset, wrapping the ring
@@ -163,7 +162,8 @@ def amd_submit_pm4(cmdbuf, devs):
  ring_idx = ((put + i.cast(put.dtype)) % q.ring.size).cast(dtypes.int)

  # copy the cmdbuf into the ring and advance the put/write pointers
-  copy_to_ring = ring.index(ring_idx, dtype=ring.dtype.ptr()).store(cmdbuf.index(i, dtype=dtypes.uint32)).end(i)
+  copy_to_ring = ring.index(ring_idx, dtype=ring.dtype.ptr()).store(
+    cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i)
  bump_put_ptr = put_ptr.index(zero, dtype=put_ptr.dtype.ptr()).store(next_put)
  bump_wptr = wptr.index(zero, dtype=wptr.dtype.ptr()).store(next_put)

@@ -219,7 +219,7 @@ def amd_submit_sdma(cmdbuf, devs):

  # the sdma queue's ring and its host-side ring/write/put pointers
  q = Device['AMD'].sdma_queue(0)
-  ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("sdma_queue", name))
+  ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("COPY:0", name))
    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))

  # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
@@ -233,7 +233,8 @@ def amd_submit_sdma(cmdbuf, devs):
  zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,))
  zero_tail = ring.index(tail_off_dw + zi, dtype=ring.dtype.ptr()).store(UOp.const(dtypes.uint32, 0)).end(zi)
  i = UOp.range(UOp.const(dtypes.int, size_dw), 0, dtype=dtypes.int, src=(cmdbuf,))
-  copy_to_ring = ring.index(start_dw + i, dtype=ring.dtype.ptr()).store(cmdbuf.index(i, dtype=dtypes.uint32)).end(i)
+  copy_to_ring = ring.index(start_dw + i, dtype=ring.dtype.ptr()).store(
+    cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i)

  # advance the put/write pointers past the zeroed tail and the cmdbuf
  next_put_b = put_b + ((zero_amt_dw + size_dw) * 4).cast(put_b.dtype)
@@ -247,14 +248,13 @@ def amd_submit_sdma(cmdbuf, devs):
@dataclass(frozen=True)
 class AMDProgramData:
  entry_point_offset:int; rsrc1:int; rsrc2:int; rsrc3:int; wave32:bool
-  kernargs_segment_size:int; kernargs_alloc_size:int
+  private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
  enable_dispatch_ptr:int; enable_private_segment_sgpr:int

 _amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {}

 def amd_build_program(prg:UOp) -> UOp:
-  devs = prg.src[1].arg  # tuple[str, ...] from rebind_program_dev
-  dev = Device[devs[0]]
+  dev = Device[prg.src[1].arg] # TODO: rm this
  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[4].arg, dev.device))) is None:
    image, sections, relocs = elf_loader(lib)
    rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
@@ -270,14 +270,16 @@ def amd_build_program(prg:UOp) -> UOp:
      rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0),  # priv=1 on gfx11 for cwsr
      rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3,
      wave32=bool(desc.kernel_code_properties & 0x400),
+      private_segment_size=desc.private_segment_fixed_size,
      kernargs_segment_size=desc.kernarg_size,
      kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0),
      enable_dispatch_ptr=edp,
      enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER), bytes(image))
-  data, image_bytes = cached
-  buf_uop = UOp.new_buffer(devs, len(image_bytes), dtypes.uint8).rtag("program")
-  blob_uop = UOp(Ops.BINARY, dtypes.void, src=(), arg=image_bytes)
-  return prg.replace(src=(buf_uop.after(buf_uop.store(blob_uop)),), arg=(data, prg.arg))
+  return cached
+
+pm_prep_program = PatternMatcher([
+  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE, arg="AMD"), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
+])

 class AMDAllocator(HCQAllocator['AMDDevice']):
  def __init__(self, dev:AMDDevice):
@@ -372,18 +374,19 @@ class PCIIface(PCIIfaceBase):
 def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})

 def encode_queue(q:UOp) -> UOp|None:
-  if not (isinstance(q.arg, tuple) and len(q.arg) == 2 and q.arg[1] in ("COMPUTE", "COPY")): return None
-  devs = q.arg[0]
-  return amd_submit_pm4(amd_lower_pm4(q, devs), devs) if q.arg[1] == "COMPUTE" else amd_submit_sdma(amd_lower_sdma(q, devs), devs)
+  q, post = (q.src[0], q.src[1:]) if q.op is Ops.AFTER else (q, ())
+  if not (isinstance(q.arg, tuple) and len(q.arg) == 2 and isinstance(q.arg[1], str) and q.arg[1].startswith(("COMPUTE", "COPY"))): return None
+  devs = to_tuple(q.arg[0])
+  ring = amd_submit_pm4(amd_lower_pm4(q, devs), devs) if q.arg[1].startswith("COMPUTE") else amd_submit_sdma(amd_lower_sdma(q, devs), devs)
+  return UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(ring, *post), arg="submit")
+
+pm_lower = PatternMatcher([
+  (UPat({Ops.LINEAR, Ops.AFTER}, name="q"), encode_queue),
+])

 class AMDDevice(HCQ2Compiled):
  timestamp_divider = 100.0  # AMD GPU clock: ticks/us

-  pm_lower = PatternMatcher([
-    (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
-    (UPat(Ops.LINEAR, name="q"), encode_queue),
-  ])
-
  ifaces = [PCIIface]

  def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
@@ -422,14 +425,13 @@ class AMDDevice(HCQ2Compiled):

    self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
    self.sdma_queues:dict = {}
-    self.has_sdma_queue = self.sdma_queue(0) is not None
+    self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this

-    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None,
-                     kernargs_size=16 << 20, can_recover=self.is_am(), arch=self.arch)
+    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)

    # Scratch setup
    self.max_private_segment_size = 0
-    self._ensure_has_local_memory(4096) # set default scratch size to 128 bytes per thread
+    self.pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.arg))]) + self.pm_bufferize

    self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
    if self.pmc_enabled:
@@ -456,19 +458,6 @@ class AMDDevice(HCQ2Compiled):
      self.sqtt_wptrs = self.allocator.alloc(round_up(self.se_cnt * self.xccs * 4, 0x1000), BufferSpec(cpu_access=True, nolru=True))
      self.sqtt_next_cmd_id = itertools.count(0)

-  @functools.cached_property
-  def compute_queue(self) -> AMDQueueDesc:
-    # https://gitlab.freedesktop.org/agd5f/linux/-/blob/a1fc9f584c4aaf8bc1ebfa459fc57a3f26a290d8/drivers/gpu/drm/amd/amdkfd/kfd_queue.c#L391
-    sgrp_size_per_cu, hwreg_size_per_cu = 0x4000, 0x1000
-    lds_size_per_cu = self.iface.props["lds_size_in_kb"] << 10 if self.target[:2] == (9,5) else 0x10000
-    vgpr_size_per_cu = 0x60000 if self.target in {(11,0,0), (11,0,1), (11,5,1), (12,0,0), (12,0,1)} else 0x80000 if self.target[0] == 9 else 0x40000
-    wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
-    ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
-    return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
-      0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
-      ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
-      debug_memory_size=round_up(self.wave_cnt * 32, 64))
-
  def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0, idx=0):
    ring = self.iface.alloc(ring_size, uncached=True, cpu_access=True)
    gart = self.iface.alloc(0x100, uncached=True, cpu_access=True)
@@ -484,9 +473,32 @@ class AMDDevice(HCQ2Compiled):
    cwsr_buffer = self.iface.alloc(cwsr_buffer_size) if ctx_save_restore_size else None
    eop_buffer = self.iface.alloc(eop_buffer_size) if eop_buffer_size else None

-    return (self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
-            wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer,
-            ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx))
+    queue = (self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
+             wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer,
+             ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx))
+
+    qname = f"{'COPY' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}"
+    self.pm_bufferize = PatternMatcher([
+      (UPat(Ops.BUFFER, tag={(qname, name)}), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
+    ] + [
+      (UPat(Ops.BUFFER, tag={(qname, "timeline_signal")}), lambda ctx, q=qname: ctx.queue_timeline_signal(q)),
+      (UPat(Ops.BUFFER, tag={(qname, "timeline_value")}), lambda ctx, q=qname: ctx.queue_timeline_value(q)),
+    ]) + self.pm_bufferize
+
+    return queue
+
+  @functools.cached_property
+  def compute_queue(self) -> AMDQueueDesc:
+    # https://gitlab.freedesktop.org/agd5f/linux/-/blob/a1fc9f584c4aaf8bc1ebfa459fc57a3f26a290d8/drivers/gpu/drm/amd/amdkfd/kfd_queue.c#L391
+    sgrp_size_per_cu, hwreg_size_per_cu = 0x4000, 0x1000
+    lds_size_per_cu = self.iface.props["lds_size_in_kb"] << 10 if self.target[:2] == (9,5) else 0x10000
+    vgpr_size_per_cu = 0x60000 if self.target in {(11,0,0), (11,0,1), (11,5,1), (12,0,0), (12,0,1)} else 0x80000 if self.target[0] == 9 else 0x40000
+    wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
+    ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
+    return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
+      0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
+      ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
+      debug_memory_size=round_up(self.wave_cnt * 32, 64))

  def sdma_queue(self, idx:int):
    if getenv("AMD_DISABLE_SDMA"): return None
@@ -495,38 +507,49 @@ class AMDDevice(HCQ2Compiled):
      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
    return self.sdma_queues.get(idx, None)

-  def _ensure_has_local_memory(self, private_segment_size):
-    if self.max_private_segment_size >= private_segment_size: return
+  def tmpring_size(self, private_segment_size):
+    private_segment_size = max(private_segment_size, 128)

    lanes_per_wave = 64 # wave64
    mem_alignment_size = 256 if self.target[0] != 9 else 1024
    size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave)
    size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
-    self.scratch, ok = self._realloc(getattr(self, 'scratch', None), size_per_xcc * self.xccs)
-    if ok:
-      # NOTE: xcc logic is correct only for GFX9.
-      max_scratch_waves = self.cu_cnt * self.iface.props['max_slots_scratch_cu'] * self.xccs
-      wave_scratch = ceildiv(lanes_per_wave * size_per_thread, mem_alignment_size)
-      num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target[0] != 9 else 1)

-      tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
-      self.tmpring_size = int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little')
+    # NOTE: xcc logic is correct only for GFX9.
+    max_scratch_waves = self.cu_cnt * self.iface.props['max_slots_scratch_cu'] * self.xccs
+    wave_scratch = ceildiv(lanes_per_wave * size_per_thread, mem_alignment_size)
+    num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target[0] != 9 else 1)
+
+    tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
+    tmpring = int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little')
+
+    if hasattr(self, 'aql_desc'):
+      gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3}
+      rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1,
+              'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})}
+      rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields')
+      rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
+
+      self.aql_desc.scratch_backing_memory_location = int(self.scratch.get_buf().va_addr)
+      self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size * lanes_per_wave // 64
+      self.aql_desc.scratch_resource_descriptor[:] = [lo32(self.scratch.get_buf().va_addr),
+        int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(self.scratch.get_buf().va_addr), SWIZZLE_ENABLE=1), 'little'),
+        lo32(size_per_xcc), int.from_bytes(bytes(rsrc3_t(**rsrc)), 'little')]
+      self.aql_desc.compute_tmpring_size = tmpring
+      self.aql_gart.cpu_view()[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
+
+    return tmpring
+
+  def scratch_buffer(self, private_segment_size):
+    private_segment_size = max(private_segment_size, 128)
+    if self.max_private_segment_size < private_segment_size:
+      lanes_per_wave = 64 # wave64
+      mem_alignment_size = 256 if self.target[0] != 9 else 1024
+      size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave)
+      size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
+      self.scratch = Buffer(self.device, size_per_xcc * self.xccs, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
      self.max_private_segment_size = private_segment_size
-
-      if hasattr(self, 'aql_desc'):
-        gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3}
-        rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1,
-                'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})}
-        rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields')
-        rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
-
-        self.aql_desc.scratch_backing_memory_location = int(self.scratch.va_addr)
-        self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size * lanes_per_wave // 64
-        self.aql_desc.scratch_resource_descriptor[:] = [lo32(self.scratch.va_addr),
-          int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(self.scratch.va_addr), SWIZZLE_ENABLE=1), 'little'),
-          lo32(size_per_xcc), int.from_bytes(bytes(rsrc3_t(**rsrc)), 'little')]
-        self.aql_desc.compute_tmpring_size = self.tmpring_size
-        self.aql_gart.cpu_view()[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
+    return self.scratch

  def on_device_hang(self): self.iface.on_device_hang()

@@ -1,7 +1,8 @@
 from __future__ import annotations
 import functools, pathlib
+from dataclasses import replace
 from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import Ops
+from tinygrad.uop.ops import shape_to_shape_arg
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler

 FP8_MAX = 448.0
@@ -11,7 +12,7 @@ NUM_WG, THREADS_PER_WG = 1024, 256
@functools.cache
 def _local_abs_max_fxn(x_p, device):
  x = Tensor(x_p, device=device)
-  inner = Tensor(x.uop.src[0]) if x.uop.op is Ops.MULTI else x
+  inner = Tensor(x.uop.replace(src=(shape_to_shape_arg(x.uop.shard_shape),), arg=replace(x.uop.arg, axis=None))) if x.uop.axis is not None else x
  return (inner.abs().max(),)

 def local_abs_max(x:Tensor) -> Tensor:
@@ -11,13 +11,13 @@ from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, compile_cpp, al
 _grad_fp8_mailbox:dict[UOp, tuple[UOp, UOp]] = {}

@functools.cache
-def _custom_fused_bwd_w13(grad_xw13:UOp, grad_xw13_fp8:UOp, grad_amax_buf:UOp,
+def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp,
                          xw13:UOp, grad_x2:UOp, amax_state:UOp, grad_amax_state:UOp, dname:str) -> UOp:
  hidden = xw13.shape[2] // 2
  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 5 + n_elems * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(grad_xw13.base, grad_xw13_fp8.base, grad_amax_buf.base,
+  mem = n_elems * 2 * 3 + n_elems * 2 + NUM_WG * 4 + 4
+  sink = UOp.sink(grad_xw13_fp8.base, grad_amax_buf.base,
                  xw13.base, grad_x2.base, amax_state.base, grad_amax_state.base, threads, workgroups,
                  arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=mem)))
  src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden)
@@ -41,23 +41,23 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
  _, _, xw13, amax_state, grad_amax_state = kernel.src[1:]
  device = xw13.device
  axis = xw13.axis if isinstance(device, tuple) else None
-  grad_xw13     = alloc_like(xw13.shape, dtypes.bfloat16, device, axis)
  grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3,  device, axis)
  grad_amax_buf = alloc_local((NUM_WG,), dtypes.float32,  device, axis)
  grad_amax_state_t = Tensor(grad_amax_state, device=device)
  fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device))
-  grad_xw13, grad_xw13_fp8, grad_amax_buf, *_ = Tensor.custom_kernel(
-    grad_xw13, grad_xw13_fp8, grad_amax_buf,
+  grad_xw13_fp8, grad_amax_buf, *_ = Tensor.custom_kernel(
+    grad_xw13_fp8, grad_amax_buf,
    Tensor(xw13, device=device), Tensor(gradient, device=device).cast(dtypes.bfloat16),
    Tensor(amax_state, device=device), grad_amax_state_t, fxn=fxn)
+  grad_xw13_uop = grad_xw13_fp8.uop.cast(dtypes.bfloat16)
  inv_scale = (grad_amax_state_t.float() + 1e-8) / FP8_MAX
  new_grad_amax = scalar_amax(grad_amax_buf)
  store_effect = grad_amax_state_t.uop.store(new_grad_amax.uop)
  assert grad_xw13_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {grad_xw13_fp8.uop.op}"
  grad_xw13_fp8_uop = grad_xw13_fp8.uop.replace(src=grad_xw13_fp8.uop.src + (store_effect,))
  # Stash fp8 companion for cdna_asm_gemm's bwd to attach to grad_a.
-  _grad_fp8_mailbox[grad_xw13.uop] = (grad_xw13_fp8_uop, inv_scale.uop)
-  return (None, None, grad_xw13.uop, None, None)
+  _grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8_uop, inv_scale.uop)
+  return (None, None, grad_xw13_uop, None, None)

 def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor) -> tuple[Tensor, Tensor, Tensor]:
  # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, inv_scale, new_amax)
@@ -21,15 +21,13 @@ constexpr float FP8_MAX = 448.0f;
 static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
 static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC");

-// fused silu*mul backward, three outputs in a single HBM pass:
-//   1) bf16 grad_xw13      — consumed by downstream bf16 autograd chain
-//   2) fp8  grad_xw13_fp8  — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd)
-//   3) fp32 grad_amax_buf  — per-WG partial |grad_xw13|, reduced into next step's grad_amax_state
+// fused silu*mul backward, two outputs in a single HBM pass:
+//   1) fp8  grad_xw13_fp8  — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd)
+//   2) fp32 grad_amax_buf  — per-WG partial |grad_xw13|, reduced into next step's grad_amax_state
 // grad_amax_state is read for the fp8 scale. The store of new_grad_amax into grad_amax_state's
 // buffer is built in Python as a separate effect and threaded into grad_a via .after(store).
 extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
 fused_silu_mul_bwd_w13(
-    __hip_bfloat16*       __restrict__ grad_xw13_out,        // bf16, 2*N_ELEMS
    __hip_fp8_storage_t*  __restrict__ grad_xw13_fp8_out,    // fp8,  2*N_ELEMS
    float*                __restrict__ grad_amax_buf,        // fp32, NUM_WG per-WG partials
    const __hip_bfloat16* __restrict__ xw13,                 // bf16, 2*N_ELEMS
@@ -62,7 +60,6 @@ fused_silu_mul_bwd_w13(
    const __hip_bfloat16 *x3 = reinterpret_cast<const __hip_bfloat16*>(&x3_raw);
    const __hip_bfloat16 *gv = reinterpret_cast<const __hip_bfloat16*>(&g_raw);

-    __hip_bfloat16 out1[VEC], out3[VEC];
    __hip_fp8_storage_t fp8_1[VEC], fp8_3[VEC];
    #pragma unroll
    for (int i = 0; i < VEC; i++) {
@@ -75,15 +72,11 @@ fused_silu_mul_bwd_w13(
      const float gs = fg * scale;
      const float g1 = gs * silu_prime * f3;
      const float g3 = gs * silu;
-      out1[i] = static_cast<__hip_bfloat16>(g1);
-      out3[i] = static_cast<__hip_bfloat16>(g3);
      local_max = fmaxf(local_max, fmaxf(fabsf(g1), fabsf(g3)));
      fp8_1[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g1 * g_scale)), __HIP_SATFINITE, __HIP_E4M3);
      fp8_3[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g3 * g_scale)), __HIP_SATFINITE, __HIP_E4M3);
    }

-    *reinterpret_cast<float4*>(&grad_xw13_out[xw1_off]) = *reinterpret_cast<float4*>(out1);
-    *reinterpret_cast<float4*>(&grad_xw13_out[xw3_off]) = *reinterpret_cast<float4*>(out3);
    *reinterpret_cast<uint64_t*>(&grad_xw13_fp8_out[xw1_off]) = *reinterpret_cast<uint64_t*>(fp8_1);
    *reinterpret_cast<uint64_t*>(&grad_xw13_fp8_out[xw3_off]) = *reinterpret_cast<uint64_t*>(fp8_3);
  }
@@ -4,18 +4,20 @@ from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType

@functools.cache
 def _custom_fused_ce_loss_fwd(loss_out:UOp, max_out:UOp, lse_out:UOp, logits:UOp, targets:UOp,
-                              vocab:int, rows:int, label_smoothing:float) -> UOp:
+                              vocab:int, rows:int, seq:int, label_smoothing:float) -> UOp:
  row = UOp.range(rows, 0)
+  b = row // seq
+  s = row % seq

  v_max = UOp.range(vocab, 1, axis_type=AxisType.REDUCE)
-  row_max = logits[row, v_max].cast(dtypes.float).reduce(v_max, arg=Ops.MAX)
+  row_max = logits[b, s, v_max].cast(dtypes.float).reduce(v_max, arg=Ops.MAX)

  v_lse = UOp.range(vocab, 2, axis_type=AxisType.REDUCE)
-  row_lse = (logits[row, v_lse].cast(dtypes.float) - row_max).exp().reduce(v_lse, arg=Ops.ADD).log() + row_max
+  row_lse = (logits[b, s, v_lse].cast(dtypes.float) - row_max).exp().reduce(v_lse, arg=Ops.ADD).log() + row_max

  v_smooth = UOp.range(vocab, 3, axis_type=AxisType.REDUCE)
-  target = logits[row, targets[row].cast(dtypes.weakint)].cast(dtypes.float)
-  mean_logits = logits[row, v_smooth].cast(dtypes.float).reduce(v_smooth, arg=Ops.ADD) / vocab
+  target = logits[b, s, targets[row].cast(dtypes.weakint)].cast(dtypes.float)
+  mean_logits = logits[b, s, v_smooth].cast(dtypes.float).reduce(v_smooth, arg=Ops.ADD) / vocab
  loss = row_lse - (1.0 - label_smoothing) * target - label_smoothing * mean_logits
  stores = UOp.group(loss_out[row].store(loss), max_out[row].store(row_max), lse_out[row].store(row_lse))

@@ -23,37 +25,42 @@ def _custom_fused_ce_loss_fwd(loss_out:UOp, max_out:UOp, lse_out:UOp, logits:UOp

@functools.cache
 def _custom_fused_ce_loss_bwd(d_logits:UOp, logits:UOp, lse:UOp, targets:UOp, scale:UOp,
-                              vocab:int, rows:int, label_smoothing:float) -> UOp:
+                              vocab:int, rows:int, seq:int, label_smoothing:float) -> UOp:
  row = UOp.range(rows, 0)
  v = UOp.range(vocab, 1)
+  b = row // seq
+  s = row % seq

-  prob = (logits[row, v].cast(dtypes.float) - lse[row]).exp()
+  prob = (logits[b, s, v].cast(dtypes.float) - lse[row]).exp()
  target = v.eq(targets[row].cast(dtypes.weakint)).where(1.0 - label_smoothing, 0.0)
  smooth = label_smoothing / vocab
  grad = (prob - target - smooth) * scale[0]

-  return d_logits[row, v].store(grad.cast(d_logits.dtype.base)).end(v, row).sink(arg=KernelInfo(f"fused_ce_loss_bwd_{rows}_{vocab}"))
+  return d_logits[b, s, v].store(grad.cast(d_logits.dtype.base)).end(v, row).sink(arg=KernelInfo(f"fused_ce_loss_bwd_{rows}_{vocab}"))

 def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float):
  # NOTE: forward inputs are (loss_out, max_out, lse_out, logits, targets)
  # gradient is the upstream grad w.r.t. per-row loss (shape: (rows,) fp32)
  _, _, lse_u, logits_u, targets_u = kernel.src[1:]
  device = logits_u.device
-  rows, VOCAB = logits_u.shape  # (rows, VOCAB) after reshape
+  MBS, SEQ, VOCAB = logits_u.shape
  if isinstance(device, tuple):
    axis = logits_u.axis
    ndev = len(device)
-    d_logits = Tensor(Tensor.invalids(rows // ndev, VOCAB, dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device)
-    rows_per_dev = rows // ndev
+    local_shape = tuple(s//ndev if i == axis else s for i,s in enumerate((MBS, SEQ, VOCAB)))
+    d_logits = Tensor(Tensor.invalids(*local_shape, dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device)
+    rows_per_dev = local_shape[0] * local_shape[1]
+    seq_per_dev = local_shape[1]
  else:
-    d_logits = Tensor.invalids(rows, VOCAB, dtype=dtypes.bfloat16, device=device)
-    rows_per_dev = rows
+    d_logits = Tensor.invalids(MBS, SEQ, VOCAB, dtype=dtypes.bfloat16, device=device)
+    rows_per_dev = MBS * SEQ
+    seq_per_dev = SEQ
  # NOTE: .mean() backward gives same grad per row (1/N), so broadcast is safe; take scalar
  scale = Tensor(gradient, device=device).float().reshape(-1)[0:1].contiguous()
  logits_t = Tensor(logits_u.after(kernel), device=device)
  lse_t = Tensor(lse_u.after(kernel), device=device)
  targets_t = Tensor(targets_u, device=device)
-  fxn = functools.partial(_custom_fused_ce_loss_bwd, vocab=VOCAB, rows=rows_per_dev, label_smoothing=label_smoothing)
+  fxn = functools.partial(_custom_fused_ce_loss_bwd, vocab=VOCAB, rows=rows_per_dev, seq=seq_per_dev, label_smoothing=label_smoothing)
  d_logits, *_ = Tensor.custom_kernel(d_logits, logits_t, lse_t, targets_t, scale, fxn=fxn)
  return (None, None, None, d_logits.uop, None)

@@ -73,17 +80,19 @@ def fused_ce_loss(logits:Tensor, targets:Tensor, label_smoothing:float=0.1) -> T
                      device=logits.device)
    lse_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
                      device=logits.device)
-    rows_per_dev = rows // ndev
+    local_shape = tuple(s//ndev if i == axis else s for i,s in enumerate(logits.shape))
+    rows_per_dev = local_shape[0] * local_shape[1]
+    seq_per_dev = local_shape[1]
  else:
    loss_out = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device)
    max_out  = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device)
    lse_out  = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device)
    rows_per_dev = rows
-  logits_flat = logits.reshape(rows, VOCAB)
+    seq_per_dev = SEQ
  targets_flat = targets.reshape(-1).cast(dtypes.int32)
-  fxn = functools.partial(_custom_fused_ce_loss_fwd, vocab=VOCAB, rows=rows_per_dev,
+  fxn = functools.partial(_custom_fused_ce_loss_fwd, vocab=VOCAB, rows=rows_per_dev, seq=seq_per_dev,
                          label_smoothing=label_smoothing)
  loss_out, max_out, lse_out, *_ = Tensor.custom_kernel(
-    loss_out, max_out, lse_out, logits_flat, targets_flat,
+    loss_out, max_out, lse_out, logits, targets_flat,
    fxn=fxn, grad_fxn=functools.partial(_fused_ce_loss_bwd, label_smoothing=label_smoothing))
  return loss_out.mean()
@@ -25,10 +25,12 @@ def calculate_storage_offset(x: Tensor) -> int:
      u_strides = strides_for_shape(u.src[0].shape)
      for i, (start, _) in enumerate(u.marg): offset += start * u_strides[i]
  return offset
-def wrap(x: Tensor) -> torch.Tensor:
+def wrap(x: Tensor, dev: torch.device|None=None) -> torch.Tensor:
  x._strides = strides_for_shape(x.shape) # always recalculate
  if (not hasattr(x, '_storage_offset')) or (not x.uop.is_realized): x._storage_offset = calculate_storage_offset(x)
-  return mod.wrap(x, _to_torch_dtype(x.dtype), _to_torch_device(x.device).index)
+  # a deviceless tinygrad value takes the device from the op context
+  idx = _to_torch_device(x.device).index if x.device is not None else (dev.index if dev is not None else 0)
+  return mod.wrap(x, _to_torch_dtype(x.dtype), idx)
 def _update_torch_metadata(tensor: torch.Tensor, tiny: Tensor) -> None:
  tiny._strides = strides_for_shape(tiny.shape)
  tiny._storage_offset = calculate_storage_offset(tiny)
@@ -545,14 +547,17 @@ def wrap_out(f):
    assigned = f(*args, **kwargs)
    if getenv("ALLOW_DTYPE_MISMATCH", 1): assigned = assigned.cast(out.dtype)
    assert out.shape == assigned.shape, f"shape mismatch: {assigned.shape} -> {out.shape}"
-    assert out.device == assigned.device, f"device mismatch: {assigned.device} -> {out.device}"
+    assert out.device == assigned.device or out.device is None or assigned.device is None, f"device mismatch: {assigned.device} -> {out.device}"
    assert out.dtype == assigned.dtype, f"dtype mismatch: {assigned.dtype} -> {out.dtype}"
+    if out.device is None and assigned.device is not None: out.replace(out.empty_like(device=assigned.device))
    return out.assign(assigned)
  return _wrap_out

 def _inplace_op(t, new_value):
  if not hasattr(t, "_view_base") and not getattr(canonical_base(t), "_views", set()): t.replace(new_value)
-  else: _apply_inplace(t, new_value)
+  else:
+    if (base:=canonical_base(t)).device is None and new_value.device is not None: base.replace(base.empty_like(device=new_value.device))
+    _apply_inplace(t, new_value)
  return t

 tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
@@ -679,10 +684,11 @@ def wrap_fxn(k,f):
    if TORCH_DEBUG:
      print(k, len(args), [x.shape if isinstance(x, torch.Tensor) else x for x in args],
                          {k:v.shape if isinstance(v, torch.Tensor) else v for k,v in kwargs.items()})
+    dev = next((a.device for a in args if isinstance(a, torch.Tensor) and a.device.type == "tiny"), None)
    args, kwargs = unwrap_args(args, kwargs)
    out = f(*args, **kwargs)
-    if isinstance(out, Tensor): return wrap(out)
-    elif isinstance(out, tuple): return tuple(wrap(x) for x in out)
+    if isinstance(out, Tensor): return wrap(out, dev)
+    elif isinstance(out, tuple): return tuple(wrap(x, dev) for x in out)
    else: raise RuntimeError(f"unknown output type {type(out)}")
  return nf

@@ -240,9 +240,9 @@ class TestTorchBackend(unittest.TestCase):
    np.testing.assert_equal(result.cpu().numpy(), [3., 3., 2.])

  def test_mnist_index(self):
+    # from tinygrad.nn.datasets import mnist
+    X_train, Y_train = Tensor.randint(60000, 1, 28, 28, dtype='uchar').realize(), Tensor.randint(60000, dtype='uchar').realize()
    GlobalCounters.reset()
-    from tinygrad.nn.datasets import mnist
-    X_train, Y_train, _, _ = mnist()
    X_train = torch.tensor(X_train.float().numpy(), device=device)
    Y_train = torch.tensor(Y_train.cast('int64').numpy(), device=device)
    samples = torch.randint(0, X_train.shape[0], (32,))
@@ -74,12 +74,13 @@ A \op{Buffer}'s \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \textt
 \op{Flip} & $(T,)$ & bools $\mathbf{f}$ & Reverse along flagged axes. \\
 \op{Reshape} & $(T, \mathbf{s'})$ & --- & Reinterpret in row-major order. $\prod s_k = \prod s'_k$. \\
 \op{Expand} & $(T, \mathbf{s'})$ & --- & Broadcast size-1 axes. $s_k \in \{1, s'_k\}$. \\
-\op{Pad} & $(T, \mathbf{b}, \mathbf{e})$ & --- & Pad with $0$s: $b_k$ before, $e_k$ after each axis. \\
-\op{Shrink} & $(T, \mathbf{b}, \mathbf{e})$ & --- & Keep $[b_k, e_k)$ per axis. Inverse of \op{Pad}. \\
+\op{Pad} & $(T, \mathbf{o}, \mathbf{s'})$ & --- & Place $T$ at offset $o_k$ in a zero-filled output of shape $s'_k$. \\
+\op{Shrink} & $(T, \mathbf{o}, \mathbf{s'})$ & --- & Keep $s'_k$ elements starting at offset $o_k$ per axis. Inverse of \op{Pad}. \\
 \op{Index} & $(T, i_0, i_1, \ldots)$ & --- & Index from left. $()$-shaped $i$ removes dim; $(k,)$-shaped makes it $k$. \\
 \op{Stack} & $(T_0, T_1, \ldots)$ & --- & Join along a newly created leading axis. All shapes must match. \\
 \op{Replicated} & $(T,)$ & axes & Mark $T$ as replicated along axes. Collapse axes to $1$. \\
 \op{Slice} & $(T, \mathrm{offset})$ & size, dtype & Zero-copy \textit{size} elems of dtype; offset is elems of $T$ dtype. \\
+\op{Bitcast} & $(T,)$ & dtype & Reinterpret storage as target dtype; preserve total bytes. \\
 \bottomrule
 \end{tabular}

@@ -165,8 +166,7 @@ Unary & $(T,)$
  & $\mathrm{trunc}(x)$: round toward zero. \\
 & & \op{Cast}
  & Convert to target dtype (specified in arg). \\
-& & \op{Bitcast}
-  & Reinterpret bits as target dtype. Must be same size. \\[4pt]
+  \\[4pt]
 Binary & $(A, B)$
  & \op{Add}, \op{Mul}, \op{Max}, \op{Mod}, \op{Idiv}
  & $a+b$, $a \cdot b$, $\max(a,b)$, $a \bmod b$, $\lfloor a/b \rfloor$ \\
@@ -167,7 +167,7 @@ class TestDSPcodePatterns(unittest.TestCase):

  def test_global_atomic_add_f32_parsing(self):
    """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype."""
-    vmem = UOp(Ops.PARAM, dtypes.uint32.ptr(1024), arg=2)
+    vmem = UOp.param(2, dtypes.uint32.ptr(1024))
    srcs = {
      'ADDR': UOp.const(dtypes.uint64, 0),
      'DATA': UOp.const(dtypes.uint32, 0x3f800000),
@@ -198,7 +198,7 @@ class TestDSPcodePatterns(unittest.TestCase):
  def test_mem_read_parsing(self):
    """Test MEM[addr].type read expression parsing."""
    # Create a mock LDS buffer
-    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp.param(3, dtypes.uint32.ptr(16384))
    addr = UOp.const(dtypes.uint32, 0)
    vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}

@@ -233,7 +233,7 @@ class TestDSPcodePatterns(unittest.TestCase):
    pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
    self.assertIsNotNone(pcode)
    assert pcode is not None
-    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp.param(3, dtypes.uint32.ptr(16384))
    srcs = {
      'ADDR': UOp.const(dtypes.uint32, 0),
      'OFFSET0': UOp.const(dtypes.uint32, 0),
@@ -314,7 +314,7 @@ class TestConcatWidthParsing(unittest.TestCase):
      self.assertEqual(parsed.simplify().arg, expected)

  def test_permlane64_wave64_pcode_indices(self):
-    vgpr = UOp(Ops.PARAM, dtypes.uint32.ptr(256), arg=0)
+    vgpr = UOp.param(0, dtypes.uint32.ptr(256))
    srcs = {
      'SRC0': UOp.const(dtypes.uint32, 0),
      'VDST': UOp.const(dtypes.uint32, 1),
@@ -347,7 +347,7 @@ class TestAllPcode(unittest.TestCase):
  def _make_srcs(self):
    """Create dummy source variables for pcode parsing."""
    u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v)
-    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp.param(3, dtypes.uint32.ptr(16384))
    return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
            'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
            'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),
@@ -125,8 +125,8 @@ class TestIndexing(unittest.TestCase):
  def test_index_mnist(self, noopt=1, op_limit=512*784*13, split_reduceop=0):
    # WEBGPU generates more ops due to bitpacking of < 4-byte dtypes
    if Device.DEFAULT == "WEBGPU": op_limit *= 15
-    from tinygrad.nn.datasets import mnist
-    X_train, Y_train, _, _ = mnist()
+    # from tinygrad.nn.datasets import mnist
+    X_train, Y_train = Tensor.randint(DSET, 1, 28, 28, dtype='uchar').realize(), Tensor.randint(DSET, dtype='uchar').realize()
    with Context(NOOPT=noopt, SPLIT_REDUCEOP=split_reduceop):
      samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0]).realize()
      GlobalCounters.reset()
@@ -155,7 +155,7 @@ class TestCustomKernel(unittest.TestCase):
    self.assertTrue((ref == tst).all().item())

  def test_eye(self):
-    ref = Tensor.eye(1024).contiguous().realize()
+    ref = Tensor.eye(1024).clone().realize()
    tst = Tensor.empty_like(ref)
    tst = tst.custom_kernel(fxn=custom_eye_kernel)[0]
    self.assertTrue((ref == tst).all().item())
@@ -335,7 +335,7 @@ class TestCustomKernel(unittest.TestCase):
    assert all(x == expected for x in result), f"expected all {expected}, got {result}"

  def test_custom_kernel_sched(self, use_custom=False):
-    x = Tensor.arange(32).reshape(8, 4).realize()
+    x = Tensor.arange(32).reshape(8, 4).clone().realize()
    y = Tensor.empty_like(x)
    y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
    if use_custom:
@@ -352,7 +352,7 @@ class TestCustomKernel(unittest.TestCase):

  @unittest.expectedFailure
  def test_sliced_buffer_function(self):
-    x = Tensor.arange(32).reshape(8, 4).realize()
+    x = Tensor.arange(32).reshape(8, 4).clone().realize()
    from tinygrad import function
    @function(precompile=True)
    def run(x:Tensor) -> Tensor:
@@ -111,6 +111,7 @@ def universal_test_cast(a, in_dtype, dtype):
 def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType):
  if not isinstance(op1, tuple): op1 = (op1, op1)
  if not isinstance(op2, tuple): op2 = (op2, op2)
+  if op1[0] == operator.mod and b == 0: return
  # lt and max with nan is undefined in tinygrad
  if op1[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
  if op2[0] in (operator.lt, Tensor.maximum) and math.isnan(c): return
@@ -219,7 +219,6 @@ class TestUOpValidationIssue(unittest.TestCase):
 class TestEdgeCases(unittest.TestCase):
  # add tests exposing new and diverse kinds of bugs that might impact real users here

-  @unittest.expectedFailure
  def test_circular_pad_negative(self):
    # negative pads with circular mode should wrap like PyTorch
    arr = np.arange(9).reshape(1, 1, 3, 3).astype(np.float32)
@@ -57,7 +57,7 @@ class TestIselX86(unittest.TestCase):
    # need to move src from gpr to xmm before broadcasting
    self.assertTrue(n.arg is X86Ops.VPBROADCASTD and n.src[0].arg is X86Ops.VMOVD)
    # if we can fuse a load we can skip the move and access memory directly
-    load = UOp(Ops.PARAM, dtypes.int32.ptr(), arg=0).index(UOp.const(dtypes.int32, 0), ptr=True).load()
+    load = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load()
    n = self.isel_rewrite(load.broadcast(4))
    self.assertTrue(n.arg is X86Ops.VPBROADCASTD and len(n.src) == 3)

@@ -122,20 +122,20 @@ class TestIselX86(unittest.TestCase):
  # complex address is [base + index*scale + displacement]
  def test_complex_address(self):
    a = UOp.variable("a", 0, 0, dtypes.int32)
-    load = UOp(Ops.PARAM, dtypes.int32.ptr(), arg=0).index(a + 1, ptr=True).load()
+    load = UOp.param(0, dtypes.int32.ptr()).index(a + 1, ptr=True).load()
    n = self.isel_rewrite(load)
    # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].arg == 4)

  def test_fold_load(self):
-    load1 = UOp(Ops.PARAM, dtypes.int32.ptr(), arg=0).index(UOp.const(dtypes.int32, 0), ptr=True).load()
-    load2 = UOp(Ops.PARAM, dtypes.int32.ptr(), arg=0).index(UOp.const(dtypes.int32, 1), ptr=True).load()
+    load1 = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load()
+    load2 = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 1), ptr=True).load()
    n = self.isel_rewrite(load1 + load2)
    self.assertTrue(len(n.src) == 4)

  # don't fold when used multiple times
  def test_dont_fold_load(self):
-    load = UOp(Ops.PARAM, dtypes.int32.ptr(), arg=0).index(UOp.const(dtypes.int32, 0), ptr=True).load()
+    load = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load()
    # used by multiple users
    n = self.isel_rewrite(load + 1 + load)
    self.assertTrue(len(n.src) == 2)
@@ -144,4 +144,4 @@ class TestIselX86(unittest.TestCase):
    self.assertTrue(len(n.src) == 2)

 if __name__ == "__main__":
-  unittest.main()
+  unittest.main()
@@ -49,7 +49,7 @@ class TestJit(unittest.TestCase):
      y = (x + 1).contiguous().realize()
      z = x.shrink(((st, st + N),)).contiguous().realize()
      return y, z
-    x = Tensor.arange(2*N).contiguous().realize()
+    x = Tensor.arange(2*N).clone().realize()
    for _ in range(3): y, z = f(x, Variable("a", 0, N).bind(0))
    self.assertEqual(y.shape, (2*N,))
    self.assertEqual(z.shape, (N,))
@@ -92,7 +92,7 @@ class TestJit(unittest.TestCase):
    @TinyJit
    def f(x): return (x[2:5].contiguous() + 1).realize()
    for i in range(5):
-      x = (Tensor.arange(10).float() + i * 10).contiguous().realize()
+      x = (Tensor.arange(10).float() + i * 10).clone().realize()
      np.testing.assert_allclose(f(x).numpy(), x.numpy()[2:5] + 1)

  def test_jit_multiple_outputs(self):
@@ -8,10 +8,10 @@ class TestKernelCache(unittest.TestCase):
    if Device.DEFAULT not in ["CPU"]:
      self.skipTest("No custom kernel cache is implemented")

-    unique_const = 0.6765677269
+    const_value = 0.6765677269
    a = Tensor.rand(4,4).realize()
    b = Tensor.rand(4,4).realize()
-    x = a + b + unique_const
+    x = a + b + const_value
    x.realize()

    a1 = Tensor.rand(4,4).realize()
@@ -20,7 +20,7 @@ class TestKernelCache(unittest.TestCase):
    Device['CPU'].compiler.compile_cached = None # making it not callable

    try:
-      x1 = a1 + b1 + unique_const
+      x1 = a1 + b1 + const_value
      x1.realize() # Same kernel should be from cache.
    finally:
      Device['CPU'].compiler.compile_cached = orig_compile_func
@@ -70,9 +70,9 @@ class TestLinearizer(unittest.TestCase):
    uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
    # only valid test if outermost range is the reduce
    if uops[uslice].src[-1].arg[-1] == AxisType.REDUCE:
-      load_types = [u.src[0].dtype for u in uops[uslice+1:] if u.op == Ops.LOAD]
+      load_idxs = [u.src[0] for u in uops[uslice+1:] if u.op == Ops.LOAD]
      # assert that there is a global load after the reduce ends
-      assert any(dt.addrspace == AddrSpace.GLOBAL for dt in load_types)
+      assert any(u.addrspace == AddrSpace.GLOBAL for u in load_idxs)

  def _test_no_nested_ranges(self, lins, skip=None):
    for l in lins:
@@ -165,7 +165,6 @@ class TestLinearizer(unittest.TestCase):
    stores = [u for u in uops if u.op is Ops.STORE]
    assert len(accs) == 0  # it's removed now
    assert len(stores) == 1
-    assert stores[0].src[1].dtype == dtypes.float.vec(4)

  # NOTE: can reenable, it does work. it just makes BEAM slow
  @unittest.expectedFailure
@@ -186,12 +185,13 @@ class TestLinearizer(unittest.TestCase):
    opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
    program = to_program(replace_opts(r.schedule_linear().src[-1].src[0], opts_to_apply), renderer=Device[Device.DEFAULT].renderer)

-    stores = [u for u in tuple(program.src[2].src) if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
+    stores = [u for u in tuple(program.src[2].src) if u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG]

    # the first store is to lds and can be upcasted
-    assert stores[0].src[1].dtype == dtypes.float.vec(4)
+    assert stores[0].src[1].max_numel() == 4
    assert any(x.op is Ops.DEFINE_LOCAL for x in stores[0].toposort())
    # the second store is to gds with no upcasts
+    assert stores[1].src[1].max_numel() == 1
    assert stores[1].src[1].dtype == dtypes.float
    assert any(x.op is Ops.PARAM for x in stores[1].toposort())

@@ -367,7 +367,7 @@ class TestLinearizer(unittest.TestCase):
    assert len(barrier) == 1
    # check that the float4 cast collapses for all stores
    for store in local_stores+global_stores:
-      assert store.src[1].dtype.count > 1 # and store.src[2].op is not Ops.VECTORIZE
+      assert store.src[1].max_numel() > 1 # and store.src[2].op is not Ops.VECTORIZE
    # # check the children's vins
    # TODO: src ALU are not the same, should it?
    # assert barrier.src == tuple(local_stores)
@@ -11,16 +11,16 @@ from tinygrad.codegen import to_program
 class TestLinearizerFailure(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
  def test_failure_beam_mnist(self):
-    c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(4014080), arg=0, src=())
+    c0 = UOp.param(0, dtypes.uchar.ptr(4014080))
    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 0, AxisType.GLOBAL)
    c2 = UOp.range(UOp.const(dtypes.weakint, 784), 1, AxisType.GLOBAL)
    c3 = UOp.range(UOp.const(dtypes.weakint, 10), 3, AxisType.GLOBAL)
-    c4 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
+    c4 = UOp.param(1, dtypes.int.ptr(512))
    c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True)))
    c6 = UOp.range(UOp.const(dtypes.weakint, 6000), 1004, AxisType.REDUCE)
    c7 = UOp.range(UOp.const(dtypes.weakint, 3750), 2006, AxisType.REDUCE)
    c8 = UOp.range(UOp.const(dtypes.weakint, 16), 2007, AxisType.GROUP_REDUCE)
-    c9 = UOp(Ops.PARAM, dtypes.uchar.ptr(47040000), arg=2, src=())
+    c9 = UOp.param(2, dtypes.uchar.ptr(47040000))
    c10 = c9.index((((c3*UOp.const(dtypes.weakint, 4704000))+c2)+(c6*UOp.const(dtypes.weakint, 784))).valid(UOp.const(dtypes.bool, True)))
    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.weakint, 6000))+c6)+((c7*UOp.const(dtypes.weakint, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.weakint, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
    c12 = c0.index((((c1*UOp.const(dtypes.weakint, 7840))+(c2*UOp.const(dtypes.weakint, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
@@ -1,8 +1,9 @@
 import unittest
-from tinygrad import Tensor, Device, dtypes, Context
+from tinygrad import Tensor, Device, dtypes, Context, GlobalCounters
 from tinygrad.helpers import getenv
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
 from extra.llama_kernels.fused_ce import fused_ce_loss
+from extra.llama_kernels import local_abs_max
 from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed, quantize_fp8_scalar
 from test.helpers import needs_second_gpu

@@ -82,5 +83,14 @@ class TestQuantizeFP8(unittest.TestCase):
    assert fp8.uop.shape == x.uop.shape
    assert new_amax.shape == ()

+class TestLocalAmax(unittest.TestCase):
+  def test_multi_tensor_local_shard_amax(self):
+    devices = ("CPU:0", "CPU:1")
+    x = Tensor.arange(16, device=devices[0]).reshape(4, 4).cast(dtypes.float).contiguous().realize().shard(devices, axis=0).realize()
+    GlobalCounters.reset()
+    out = (x * local_abs_max(x)).contiguous().realize()
+    self.assertEqual(GlobalCounters.kernel_count, 4)
+    self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])
+
 if __name__ == '__main__':
  unittest.main()
@@ -187,7 +187,7 @@ class TestMultiTensor(unittest.TestCase):
    np.testing.assert_allclose(O.numpy(), X.numpy()[0:2]*W.numpy()[0:2] < 2)

  def test_shrink_on_shard_axis(self):
-    X = Tensor.arange(4*4).reshape(4,4).realize()
+    X = Tensor.arange(4*4).reshape(4,4).clone().realize()
    X_np = X.numpy()
    X.shard_(devices_2, 0)
    # only shrink on the device that owns the shard, this is enabled by the mselect simplifier
@@ -293,7 +293,7 @@ class TestMultiTensor(unittest.TestCase):
    @TinyJit
    def f(x): return (x+1).contiguous().sum()
    for _ in range(5):
-      tt = Tensor.arange(0, 4).contiguous().realize().shard((d1,d2), 0).realize()
+      tt = Tensor.arange(0, 4).clone().realize().shard((d1,d2), 0).realize()
      out = f(tt)
      assert out.item() == 1+2+3+4

@@ -309,7 +309,7 @@ class TestMultiTensor(unittest.TestCase):
    @TinyJit
    def f(x): return (x.shard((d1,d2), 0)+1).contiguous().sum()
    for _ in range(5):
-      tt = Tensor.arange(0, 4).contiguous().realize()
+      tt = Tensor.arange(0, 4).clone().realize()
      out = f(tt)
      assert out.item() == 1+2+3+4

@@ -865,7 +865,7 @@ class TestMultiTensor(unittest.TestCase):
  @unittest.skip("RANGEIFY doesn't support multi const folding")
  def test_multi_const_folding(self):
    with Context(TRACK_MATCH_STATS=0):
-      a = Tensor.arange(3).realize()
+      a = Tensor.arange(3).clone().realize()
      zeros = Tensor.zeros(3).realize()
    b = a.to(devices_2)*zeros.to(devices_2)
    sched = b.schedule_linear().src
@@ -904,7 +904,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):

  # shrink a multitensor on sharded axis
  def test_shrink_bad_args(self):
-    t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
+    t = Tensor.arange(64).reshape(8, 8).clone().realize()
    t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)

    with self.assertRaises(AssertionError):
@@ -927,7 +927,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
  @given(strat.sampled_from([dtypes.float, dtypes.int, dtypes.int64, dtypes.int16]))
  def test_ops(self, dtype):
    if dtype not in Device[Device.DEFAULT].renderer.supported_dtypes(): return
-    t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
+    t = Tensor.arange(64).reshape(8, 8).clone().realize()
    t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
    for i in range(4):
      print(f"{i=}")
@@ -971,7 +971,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
      np.testing.assert_allclose(a.flip(-1).numpy(), b.flip(-1).numpy(), rtol=1e-7, atol=1e-3)

  def test_add_two_partitions(self):
-    t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
+    t = Tensor.arange(64).reshape(8, 8).clone().realize()
    t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)

    a = t.shrink(((2, 4), None))
@@ -988,7 +988,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):

  def test_add_different_tensors(self):
    devices = [f"{Device.DEFAULT}:{i}" for i in range(4)]
-    x = Tensor.arange(64).reshape(8, 8).contiguous().realize().shard(devices, axis=0)
+    x = Tensor.arange(64).reshape(8, 8).clone().realize().shard(devices, axis=0)

    to_add = []
    for i in range(len(devices)):
@@ -1098,7 +1098,7 @@ class TestBatchNorm(unittest.TestCase):
  @given(strat.sampled_from((False, True)))
  def test_batchnorm(self, is_training):
    devices = [f"{Device.DEFAULT}:{i}" for i in range(4)]
-    x = Tensor.arange(4096).reshape(8, 8, 8, 8).contiguous().realize().shard(devices, axis=0)
+    x = Tensor.arange(4096).reshape(8, 8, 8, 8).clone().realize().shard(devices, axis=0)

    with Tensor.train(is_training):
      bns = []
@@ -1184,7 +1184,7 @@ class TestMultiBufferView(unittest.TestCase):

  @unittest.skip("flaky on LLVM")
  def test_shrink_non_shard_axis(self):
-    ref = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().realize()
+    ref = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().realize()
    a = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().shard(devices_2, axis=1).realize()
    self._check(ref, a, lambda t: t[3])

@@ -1296,7 +1296,7 @@ class TestMultiSetitem(unittest.TestCase):
  @needs_second_gpu
  def setUp(self): pass

-  def _t(self, axis): return Tensor.arange(16).contiguous().realize().shard(self.device, axis=axis)
+  def _t(self, axis): return Tensor.arange(16).clone().realize().shard(self.device, axis=axis)

  def test_setitem_scalar_axis0(self):
    t = self._t(0)
@@ -1989,9 +1989,7 @@ class TestOps(unittest.TestCase):
    self.helper_test_exception([(1,1,5,5)],
                                lambda x: torch.nn.functional.pad(x, (3,6,0,0), mode="circular"), lambda x: x.pad((3,6,0,0), mode="circular"),
                                expected=(RuntimeError, ValueError))
-    with self.assertRaises(NotImplementedError):
-      # negative pads with circular pads is not supported
-      Tensor.randn(1,1,5,5).pad((3,-5,1,-5), mode="circular")
+    helper_test_op([(1,1,5,5)], lambda x: torch.nn.functional.pad(x, (1,-2,2,-1), mode="circular"), lambda x: x.pad((1,-2,2,-1), mode="circular"))

  def test_pad_reshape(self):
    helper_test_op([(1, 2)],
@@ -67,7 +67,7 @@ class TestPickle(unittest.TestCase):

  # NOTE: currently Buffer exists on the uop, not tensor
  def test_pickle_buffer_uop(self):
-    t = Tensor.arange(4).realize()
+    t = Tensor.arange(4).clone().realize()
    a = t.uop
    assert a.is_realized
    self.assertIsNotNone(buffer:=a.base.realized)
@@ -95,7 +95,7 @@ class TestPickle(unittest.TestCase):
    np.testing.assert_equal(vt2.numpy(), 20)

  def test_pickle_buffer_view(self):
-    t = Tensor.arange(10, device="CPU").contiguous().realize()
+    t = Tensor.arange(10).clone(device="CPU").realize()
    vt = t[3:5].contiguous().realize()
    assert hasattr(vt.uop.buffer, 'base')
    ref_value = vt.tolist()
@@ -22,30 +22,30 @@ def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None):

 def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
  dtype = alu_src_uops[0].dtype
-  a = UOp(Ops.PARAM, dtype.ptr(), (), 0)
-  b = UOp(Ops.PARAM, dtype.ptr(), (), 1)
+  a = UOp.param(0, dtype.ptr(1))
+  b = UOp.param(1, dtype.ptr(1))
  idx = UOp.const(dtypes.int, 0)
-  ld = b.index(idx)
+  ld = b.index(idx, ptr=True).load()
  alu = ld.alu(alu_op, *alu_src_uops)
-  store = UOp.store(a.index(idx), alu)
+  store = UOp.store(a.index(idx, ptr=True), alu)
  return _test_uop_result([Tensor([input_val])], UOp(Ops.SINK, dtypes.void, (store,), arg=KernelInfo()))[0]

 class TestRendererFailures(unittest.TestCase):
  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
  def test_gated_store_with_alu(self):
-    a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    a = UOp.param(0, dtypes.int.ptr(4))
    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
-    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
+    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu), ptr=True), UOp.const(dtypes.int, 1)))
    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
    ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
    np.testing.assert_equal(ret, [0, 1, 1, 1])

  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
  def test_gated_store_with_alu_2d(self):
-    a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    a = UOp.param(0, dtypes.int.ptr(8))
    gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
    gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0)
-    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
+    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1), ptr=True), UOp.const(dtypes.int, 1)))
    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
    ret = _test_uop_result([], sink, local_size=[4, 2, 1])[0]
    np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1])
@@ -87,7 +87,7 @@ class TestWGSLFailures(unittest.TestCase):
 class TestPTXFailures(unittest.TestCase):
  @unittest.skip("INDEX can only have a gate ALU parent, not an IF")
  def test_gated_store_with_if(self):
-    a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    a = UOp.param(0, dtypes.int.ptr())
    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
    val = UOp.const(dtypes.int, 1)
    if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,))
@@ -221,8 +221,8 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_allclose(out.numpy(), (x.numpy() - x.numpy().max(keepdims=True)).max())

  def test_example_matmul_contig(self):
-    x = Tensor.eye(64).contiguous().realize()
-    y = Tensor.eye(64).contiguous().realize()
+    x = Tensor.eye(64).clone().realize()
+    y = Tensor.eye(64).clone().realize()
    z = y.matmul(x).sum()
    z.backward()
    out = x.grad.contiguous()
@@ -826,7 +826,7 @@ class TestSchedule(unittest.TestCase):
    self._test_fusion([(32, 32)], lambda a:a-a.sum(1), 2)

  def test_cast_padded_view(self):
-    a = Tensor.arange(4).reshape(1, 4)
+    a = Tensor.arange(4).reshape(1, 4).clone().realize()
    casted_view = a.pad(((0, 1), (0, 0))).cast(dtypes.float)
    casted_view.realize()
    self.assertEqual(casted_view.uop.base.realized.size, 8)
@@ -836,7 +836,7 @@ class TestSchedule(unittest.TestCase):

  # NOTE: we only reorder CAST if it's an EXPAND
  def test_cast_after_shrink(self):
-    a = Tensor.arange(4).reshape(1, 4)
+    a = Tensor.arange(4).reshape(1, 4).clone().realize()
    casted_view = a.shrink(((0, 1), (0, 2))).cast(dtypes.float)
    casted_view.realize()
    self.assertEqual(casted_view.uop.base.realized.size, 2)
@@ -991,7 +991,7 @@ class TestSchedule(unittest.TestCase):

  def test_assign_non_contiguous_alt(self): self.test_assign_non_contiguous(alt=True)
  def test_assign_non_contiguous(self, alt=False):
-    x = (Tensor.arange(16)-100).reshape(4,4).contiguous().realize()
+    x = (Tensor.arange(16)-100).reshape(4,4).clone().realize()
    xref = x.numpy()
    if alt:
      y = Tensor.randint(2, 4).contiguous().realize()
@@ -1007,7 +1007,7 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_equal(tst.numpy(), a.numpy())

  def test_setitem_sched(self, mop=lambda x:x, expected_kcount=1):
-    a = Tensor.arange(16, device="CPU").reshape(4, 4).contiguous().realize()
+    a = Tensor.arange(16).reshape(4, 4).clone(device="CPU").realize()
    a2 = mop(a)
    expected = (a+a2).tolist()
    a.assign(a+a2)
@@ -1021,7 +1021,7 @@ class TestSchedule(unittest.TestCase):

  def test_setitem_const_fused(self):
    # https://github.com/tinygrad/tinygrad/issues/10690
-    a = Tensor.arange(16).contiguous().realize()
+    a = Tensor.arange(16).clone().realize()
    GlobalCounters.reset()
    a[4] = 3
    self.assertEqual(GlobalCounters.kernel_count, 0)
@@ -1059,9 +1059,9 @@ class TestSchedule(unittest.TestCase):
    self.assertEqual(b.tolist(), [False, False])

  def test_mnist_val(self):
-    from tinygrad.nn.datasets import mnist
+    # from tinygrad.nn.datasets import mnist
    import torch
-    _, Y_train, _, _ = mnist()
+    Y_train = Tensor.randint(60000, dtype='uchar').realize()
    samples = Tensor.randint(BS:=getenv("BS", 512), high=cast(int,Y_train.shape[-1])).realize()
    yt = Tensor.randn(BS, 10).realize()
    loss = yt.sparse_categorical_crossentropy(Y_train[samples])
@@ -1278,7 +1278,7 @@ class TestCopyFolding(unittest.TestCase):
    self.assertEqual(x.item(), 2.0)

  def test_late_const_copy_folding(self):
-    a = Tensor.arange(3).realize()
+    a = Tensor.arange(3).clone().realize()
    zeros = Tensor.zeros(3, buffer=False).realize()
    b = (a*zeros).to("CPU") + 1
    run_linear(*check_schedule(b, 1, filter_sink=False))
@@ -1353,14 +1353,14 @@ class TestCopyFolding(unittest.TestCase):
    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])

  def test_permute_on_disk(self):
-    with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_memoryview())
+    with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).clone().realize().uop.base.buffer.as_memoryview())
    a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute')}")
    b = a.reshape(2, 2).permute(1, 0).to("CPU")
    b.realize()
    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])

  def test_permute_on_disk_contiguous(self):
-    with open(temp('dt_arange_4_permute_contig'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_memoryview())
+    with open(temp('dt_arange_4_permute_contig'), "wb") as f: f.write(Tensor.arange(4).clone().realize().uop.base.buffer.as_memoryview())
    a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute_contig')}")
    b = a.reshape(2, 2).permute(1, 0).contiguous().to("CPU")
    b.realize()
@@ -1374,7 +1374,7 @@ class TestCopyFolding(unittest.TestCase):

  # NOTE: disk permute must come after COPY
  def test_permute_after_shrink_on_disk(self):
-    with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).realize().uop.base.buffer.as_memoryview())
+    with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).clone().realize().uop.base.buffer.as_memoryview())
    a = Tensor.empty(5, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_5_permute')}")
    b = a.shrink(((0, 4),)).reshape(2, 2).permute(1, 0).to("CPU")
    b.realize()
@@ -32,13 +32,14 @@ class TestSetitem(unittest.TestCase):
    self.assertListEqual(t.tolist(), [0, 1, 11, 3, 11, 5, 6, 7, 8, 9])

  def test_setitem_inplace_mul(self):
-    t = Tensor.arange(10).realize()
+    t = Tensor.arange(10).clone().realize()
    t[:3] *= 10
    self.assertListEqual(t.tolist(), [0, 10, 20, 3, 4, 5, 6, 7, 8, 9])

+  @unittest.skip("crashed in LLVM CI")
  def test_setitem_fancy_on_unrealized_view(self):
    # fancy indexing setitem on unrealized SHRINK view (triggered infinite loop in graph_rewrite)
-    base = Tensor.arange(20, dtype=dtypes.float).reshape(4, 5)
+    base = Tensor.arange(20, dtype=dtypes.float).reshape(4, 5).clone().realize()
    sub = base[1:3]
    flat = sub.reshape(sub.numel()).contiguous()
    idx = Tensor([0, 3, 7, 9])
@@ -229,7 +230,7 @@ class TestSetitem(unittest.TestCase):
    np.testing.assert_equal(t.numpy(), n)

  def test_setitem_swap_rows(self):
-    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).contiguous().realize()
+    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
    tmp = t[0]
    t[0] = t[1]
    t[2] = tmp
@@ -237,7 +238,7 @@ class TestSetitem(unittest.TestCase):
    np.testing.assert_allclose(t.numpy(), [[2, 3], [2, 3], [2, 3]])

    # eager version
-    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).contiguous().realize()
+    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
    tmp = t[0].realize()
    t[0] = t[1].realize()
    t[2] = tmp.realize()
@@ -269,8 +270,8 @@ class TestSetitem(unittest.TestCase):
  def test_cross_assign_independence(self):
    # when assigning to two tensors using computations from both,
    # both assigns should see the OLD values of both tensors
-    a = Tensor.arange(4, dtype=dtypes.float).contiguous().realize()
-    b = Tensor.arange(4, 8, dtype=dtypes.float).contiguous().realize()
+    a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
+    b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
    new_a = a + b    # [4, 6, 8, 10]
    new_b = a * 2    # [0, 2, 4, 6] -- should use OLD a
    a.assign(new_a)
@@ -283,8 +284,8 @@ class TestSetitem(unittest.TestCase):
      np.testing.assert_allclose(b.numpy(), [8, 12, 16, 20])

    # eager version
-    a = Tensor.arange(4, dtype=dtypes.float).contiguous().realize()
-    b = Tensor.arange(4, 8, dtype=dtypes.float).contiguous().realize()
+    a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
+    b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
    new_a = (a + b).realize()
    new_b = (a * 2).realize()
    a.assign(new_a).realize()
@@ -323,7 +324,7 @@ class TestWithGrad(unittest.TestCase):
  def test_set_overlapping_backward(self):
    z = Tensor.zeros(6)
    x = Tensor.ones(4).contiguous()
-    y = Tensor.ones(4).contiguous() * 2
+    y = Tensor.ones(4) * 2
    z[:4] = x
    z[2:] = y
    z.sum().backward()
@@ -36,7 +36,7 @@ class TestSubBuffer(unittest.TestCase):
    assert len(mv) == 5

  def test_subbuffer_used(self):
-    t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
+    t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
    vt = t[2:4].realize()
    out = (vt + 100).tolist()
    assert out == [102, 103]
@@ -44,7 +44,7 @@ class TestSubBuffer(unittest.TestCase):
  @needs_second_gpu
  @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or DEV.interface.startswith("MOCK"), "only NV, AMD, CUDA")
  def test_subbuffer_transfer(self):
-    t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
+    t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
    vt = t[2:5].contiguous().realize()
    out = vt.to(f"{Device.DEFAULT}:1").realize().tolist()
    assert out == [2, 3, 4]
@@ -551,7 +551,7 @@ class TestTinygrad(unittest.TestCase):
      Tensor.zeros(2, 2).realize()

  def test_shrink(self):
-    t = Tensor.arange(32).contiguous().realize()
+    t = Tensor.arange(32).clone().realize()
    self.assertListEqual(t[16:20].tolist(), [16,17,18,19])
    self.assertListEqual(t.shrink_to(16).tolist(), list(range(16)))
    t = t.reshape(4, 8).contiguous().realize()
@@ -20,14 +20,15 @@ def run_uops(uops_list:list[UOp], bufs:list[Buffer]):

 def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
  if op is Ops.CONST: uops.append(UOp.const(dtype, arg))
+  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype).replace(src=()))
  else: uops.append(UOp(op, dtype, tuple(src), arg))
  return uops[-1]

 def _test_single_value(vals, op, dts):
  uops = []
  output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
-  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
-  buf_loads = [uop(uops, Ops.PARAM, dtype.ptr(), (), i+1) for i,dtype in enumerate(dts)]
+  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(1), (), 0)
+  buf_loads = [uop(uops, Ops.PARAM, dtype.ptr(1), (), i+1) for i,dtype in enumerate(dts)]
  loads = (buf_loads[i].index(uop(uops, Ops.CONST, dtypes.int32, (), 0)) for i, dtype in enumerate(dts))
  alu = uop(uops, op, output_dtype, loads)
  out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0), ptr=True), alu))
@@ -41,7 +42,7 @@ def _test_single_value(vals, op, dts):
 def _test_single_value_const(vals, op, dts):
  uops = []
  output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
-  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
+  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(1), (), 0)
  loads = (uop(uops, Ops.CONST, dtype, [], a) for a,dtype in zip(vals, dts))
  alu = uop(uops, op, output_dtype, loads)
  out = buf_store[UOp.const(dtypes.int32, 0)].store(alu)
@@ -53,7 +54,7 @@ def _test_single_value_const(vals, op, dts):

 def _test_uops_result(output_dtype, uops, res):
  # uops = []
-  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
+  buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(1), (), 0)
  # res = output_fn(uops)
  out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), res))
  buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
@@ -220,8 +221,8 @@ class TestLocalAccess(unittest.TestCase):
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends")
 class TestAssembly(unittest.TestCase):
  def test_bitshift_left(self):
-    g1 = UOp(Ops.PARAM, dtypes.int32.ptr(), (), 0)
-    out = UOp(Ops.PARAM, dtypes.int32.ptr(), (), 1)
+    g1 = UOp.param(0, dtypes.int32.ptr(3))
+    out = UOp.param(1, dtypes.int32.ptr(2))
    c1 = UOp.const(dtypes.int, 2)
    c2 = UOp.const(dtypes.int, 3)
    l1 = g1.index(c1)
@@ -248,7 +249,7 @@ class TestAssembly(unittest.TestCase):
    self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)

  def test_mulacc_shl(self):
-    g1 = UOp(Ops.PARAM, dtypes.int32.ptr(), (), 0)
+    g1 = UOp.param(0, dtypes.int32.ptr(2))
    c1 = UOp.const(dtypes.int, 0)
    c2 = UOp.const(dtypes.int, 1)
    expr = g1.index(c1) * UOp.const(dtypes.int, 4096) + g1.index(c2)
@@ -257,7 +258,7 @@ class TestAssembly(unittest.TestCase):
    self.assertIn(Ops.MULACC, [x.op for x in uops])

  def test_use_cmpeq(self):
-    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
+    g = UOp.param(0, dtypes.uint32.ptr(8))
    c = UOp.const(dtypes.uint, 7)
    comp = g.index(c).ne(c).ne(True)
    uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
@@ -12,7 +12,7 @@ from tinygrad.dtype import ImageDType, Invalid
 # PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx

 def vision_conv_143():
-  c0 = UOp(Ops.PARAM, dtypes.imageh((16, 1024, 4)), (), 0)
+  c0 = UOp.param(0, dtypes.imageh((16, 1024, 4)))
  c2 = UOp.range(32, 3, AxisType.LOOP)
  c5 = UOp.range(128, 4, AxisType.LOOP)
  c8 = UOp.range(16, 2, AxisType.LOOP)
@@ -22,13 +22,13 @@ def vision_conv_143():
  c26 = UOp.range(7, 1, AxisType.REDUCE)
  c27 = c2*2+c26
  c32 = ((c27<3)!=True)&(c27<67)
-  c34 = UOp(Ops.PARAM, dtypes.imageh((32, 1024, 4)), (), 1)
+  c34 = UOp.param(1, dtypes.imageh((32, 1024, 4)))
  c38 = c5//2
  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.weakint, Invalid))
  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
-  c49 = UOp(Ops.PARAM, dtypes.imageh((64, 49, 4)), (), 2)
+  c49 = UOp.param(2, dtypes.imageh((64, 49, 4)))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp(Ops.PARAM, dtypes.float.ptr(128), (), 3)
+  c63 = UOp.param(3, dtypes.float.ptr(128))
  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
  c67 = c0.index((c2*128+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)

@@ -38,7 +38,7 @@ def vision_conv_143():
  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))

 def vision_conv_153():
-  c0 = UOp(Ops.PARAM, dtypes.imageh((8, 1024, 4)), (), 0)
+  c0 = UOp.param(0, dtypes.imageh((8, 1024, 4)))
  c2 = UOp.range(16, 3, AxisType.LOOP)
  c5 = UOp.range(256, 4, AxisType.LOOP)
  c8 = UOp.range(8, 2, AxisType.LOOP)
@@ -48,13 +48,13 @@ def vision_conv_153():
  c26 = UOp.range(7, 1, AxisType.REDUCE)
  c27 = c2*2+c26
  c32 = ((c27<3)!=True)&(c27<35)
-  c34 = UOp(Ops.PARAM, dtypes.imageh((16, 1024, 4)), (), 1)
+  c34 = UOp.param(1, dtypes.imageh((16, 1024, 4)))
  c38 = c5//2
  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.weakint, Invalid))
  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
-  c49 = UOp(Ops.PARAM, dtypes.imageh((128, 49, 4)), (), 2)
+  c49 = UOp.param(2, dtypes.imageh((128, 49, 4)))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp(Ops.PARAM, dtypes.float.ptr(256), (), 3)
+  c63 = UOp.param(3, dtypes.float.ptr(256))
  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
  c67 = c0.index((c2*256+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)

@@ -64,16 +64,16 @@ def vision_conv_153():
  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))

 def dm_conv_172():
-  c0 = UOp(Ops.PARAM, dtypes.imageh((1, 240, 4)), (), 0)
+  c0 = UOp.param(0, dtypes.imageh((1, 240, 4)))
  c2 = UOp.range(960, 4, AxisType.LOOP)
-  c5 = UOp(Ops.PARAM, dtypes.imageh((8, 384, 4)), (), 1)
+  c5 = UOp.param(1, dtypes.imageh((8, 384, 4)))
  c7 = UOp.range(32, 0, AxisType.REDUCE)
  c10 = UOp.range(4, 1, AxisType.REDUCE)
  c13 = UOp.range(12, 3, AxisType.REDUCE)
  c18 = UOp.range(8, 2, AxisType.REDUCE)
-  c23 = UOp(Ops.PARAM, dtypes.imageh((240, 128, 4)), (), 2)
+  c23 = UOp.param(2, dtypes.imageh((240, 128, 4)))
  c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
-  c37 = UOp(Ops.PARAM, dtypes.float.ptr(960), (), 3)
+  c37 = UOp.param(3, dtypes.float.ptr(960))
  c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
  c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
  c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666
@@ -99,4 +99,4 @@ bufs = [Buffer(ps.arg.device, g.size, g.dtype if isinstance(g.dtype, ImageDType)

 gsize, lsize = ps.arg.launch_dims({})
 t = rt(*[b._buf for b in bufs], global_size=gsize, local_size=lsize, vals=ps.arg.vals({}), wait=True)
-print(f"{t*1e6:.2f} us")
+print(f"{t*1e6:.2f} us")
@@ -285,7 +285,7 @@ class TestMainOnnxOps(TestOnnxOps):
    for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
      for b in (np.ones([32], dtype=np.int32), np.zeros([32], dtype=np.int32)):
        for channel_shape in [(), (32,)]:
-          with self.subTest(dtype=dtype, zero_point=zero_point, channel_shape=channel_shape):
+          with self.subTest(dtype=dtype.__name__, zero_point=zero_point, channel_shape=channel_shape):
            dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
            inputs = {
              "x": np.random.randint(dtype_min, dtype_max + 1, [1, 3, 224, 224], dtype=dtype),
@@ -304,7 +304,7 @@ class TestMainOnnxOps(TestOnnxOps):

  def test_qlinear_matmul(self):
    for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
-      with self.subTest(dtype=dtype, zero_point=zero_point):
+      with self.subTest(dtype=dtype.__name__, zero_point=zero_point):
        dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
        inputs = {
          "A": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
@@ -512,7 +512,7 @@ class TestContribOnnxOps(TestOnnxOps):

  def test_qlinear_add(self):
    for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
-      with self.subTest(dtype=dtype, zero_point=zero_point):
+      with self.subTest(dtype=dtype.__name__, zero_point=zero_point):
        dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
        inputs = {
          "A": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
@@ -546,7 +546,7 @@ class TestContribOnnxOps(TestOnnxOps):

  def test_qlinear_mul(self):
    for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
-      with self.subTest(dtype=dtype, zero_point=zero_point):
+      with self.subTest(dtype=dtype.__name__, zero_point=zero_point):
        dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
        inputs = {
          "A": np.random.randint(dtype_min, dtype_max + 1, [10, 10], dtype=dtype),
@@ -580,7 +580,7 @@ class TestContribOnnxOps(TestOnnxOps):
  def test_qlinear_global_average_pool(self):
    for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
      for channels_last in [0, 1]:
-        with self.subTest(dtype=dtype, zero_point=zero_point, channels_last=channels_last):
+        with self.subTest(dtype=dtype.__name__, zero_point=zero_point, channels_last=channels_last):
          dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
          # NCHW for channels_last=0, NHWC for channels_last=1
          shape = [1, 3, 32, 32] if channels_last == 0 else [1, 32, 32, 3]
@@ -29,7 +29,7 @@ def gradient_test():
  z = y.matmul(x).sum()
  z.backward()
 def realized_eye():
-  Tensor.eye(3).realize()
+  Tensor.eye(3).clone().realize()
 def realized_list():
  Tensor([[2.0,0,-2.0]]).realize()
 def kernel_matmul():
@@ -4,7 +4,7 @@ import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, functoo
 from dataclasses import replace
 from typing import Callable, Any

-ASSERT_DIFF = int((flag:="[pr]") in os.getenv("COMMIT_MESSAGE", flag) or flag in os.getenv("PR_TITLE", flag))
+ASSERT_DIFF = int((flag:="[PR]") in os.getenv("COMMIT_MESSAGE", flag) or flag in os.getenv("PR_TITLE", flag))
 if not int(os.getenv("ASSERT_PROCESS_REPLAY", "1")): ASSERT_DIFF = 0

 try:
@@ -82,8 +82,8 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
  for buf_dt, data in inputs or []:
    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
-  g = UOp(Ops.PARAM, uop.dtype.ptr(), arg=0, src=())
-  prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
+  g = UOp.param(0, uop.dtype.ptr(1))
+  prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 0), ptr=True), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
  prog = PythonProgram("run", PythonCompiler().compile(prg.src[3].arg))
  prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)
  return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
@@ -423,10 +423,10 @@ def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode
 class _Ctx:
  """Context for instruction compilation - holds buffers and helpers."""
  __slots__ = ('inst_size', 'dyn_fields', '_axis_id', 'wave_size', 'vgpr', 'accvgpr')
-  sgpr = UOp(Ops.PARAM, dtypes.uint32.ptr(SGPR_COUNT), arg=0)
-  vmem = UOp(Ops.PARAM, dtypes.uint32.ptr(1 << 46), arg=2)
-  lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
-  scratch = UOp(Ops.PARAM, dtypes.uint8.ptr(1 << 30), arg=4)
+  sgpr = UOp.param(0, dtypes.uint32.ptr(SGPR_COUNT))
+  vmem = UOp.param(2, dtypes.uint32.ptr(1 << 46))
+  lds = UOp.param(3, dtypes.uint32.ptr(16384))
+  scratch = UOp.param(4, dtypes.uint8.ptr(1 << 30))
  # Cache PARAM UOps by wave_size so all _Ctx instances with same wave_size share identical UOp references
  _vgpr_cache: dict[int, UOp] = {}
  _accvgpr_cache: dict[int, UOp] = {}
@@ -434,10 +434,10 @@ class _Ctx:
  def __init__(self, inst_size: int, wave_size: int = 32):
    self.inst_size, self._axis_id, self.wave_size = inst_size, 0, wave_size
    self.dyn_fields: list[tuple[int, int]] = []  # (lo, hi) of fields read dynamically
-    if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp(Ops.PARAM, dtypes.uint32.ptr(256 * wave_size), arg=1)
+    if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32.ptr(256 * wave_size))
    self.vgpr = _Ctx._vgpr_cache[wave_size]
    if wave_size == 64:
-      if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp(Ops.PARAM, dtypes.uint32.ptr(256 * wave_size), arg=5)
+      if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32.ptr(256 * wave_size))
      self.accvgpr = _Ctx._accvgpr_cache[wave_size]
    else:
      self.accvgpr = self.vgpr
@@ -598,12 +598,13 @@ class _Ctx:
  def rpc(self) -> UOp:
    """Read PC as 64-bit byte address."""
    # Index at PC_LO, then cast to uint64 ptr and load
-    return self.sgpr.index(_c(PC_LO_IDX, dtypes.int), ptr=True).cast(dtypes.uint64.ptr(SGPR_COUNT // 2)).load()
+    return _u64(self.rsgpr_dyn(_c(PC_LO_IDX)), self.rsgpr_dyn(_c(PC_HI_IDX)))

  def inc_pc(self) -> list[UOp]:
    """Increment PC by instruction size in bytes. Returns [store]."""
    new_pc = self.rpc() + UOp.const(dtypes.uint64, self.inst_size)
-    return [self.sgpr.index(_c(PC_LO_IDX, dtypes.int), ptr=True).cast(dtypes.uint64.ptr(SGPR_COUNT // 2)).store(new_pc)]
+    lo, hi = _split64(new_pc)
+    return [self.wsgpr_dyn(_c(PC_LO_IDX), lo), self.wsgpr_dyn(_c(PC_HI_IDX), hi)]

  def scalar_stores(self, assigns: list[tuple[str, UOp]], sdst_reg: UOp, sdst_size: int = 1) -> list[UOp]:
    """Generate stores for scalar assigns with dynamic destination register (D0, SCC, EXEC, VCC)."""
@@ -1343,6 +1344,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
  src0_r = src0_off - _c(256)  # VGPR-relative index (only valid when src is VGPR)
  src1_r = src1_off - _c(256)
  src2_off = ctx.inst_field(type(inst).src2)
+  use_acc = bool(getattr(inst, 'acc_cd', 0))
  # Check if sources are VGPRs (offset >= 256) vs inline constants/SGPRs
  src0_is_vgpr = src0_off >= _c(256)
  src1_is_vgpr = src1_off >= _c(256)
@@ -1503,7 +1505,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
      # So: m_base = half * 16 + (out_reg // 4) * 4 + (out_reg % 4)
      m_base = c_half * UOp.const(dtypes.int, 16) + UOp.const(dtypes.int, (out_reg // 4) * 4 + (out_reg % 4))

-      acc_v = ctx.raccvgpr_dyn(src2_r + _c(out_reg), compute_lane, src2_is_vgpr)
+      acc_v = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(src2_r + _c(out_reg), compute_lane, src2_is_vgpr)
      if is_int_out: acc_v = acc_v.cast(dtypes.int32)
      else: acc_v = acc_v.bitcast(dtypes.float32)
      acc = src2_is_vgpr.where(acc_v, acc_scalar)
@@ -1514,16 +1516,18 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
        acc = acc + a_val * b_val

      if is_int_out:
-        compute_stores.append(ctx.waccvgpr_dyn(vdst_reg + _c(out_reg), compute_lane, acc.cast(dtypes.uint32), exec_mask))
+        compute_stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(
+          vdst_reg + _c(out_reg), compute_lane, acc.cast(dtypes.uint32), exec_mask))
      else:
-        compute_stores.append(ctx.waccvgpr_dyn(vdst_reg + _c(out_reg), compute_lane, acc.bitcast(dtypes.uint32), exec_mask))
+        compute_stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(
+          vdst_reg + _c(out_reg), compute_lane, acc.bitcast(dtypes.uint32), exec_mask))
  else:
    # 16x16 and 4x4: each lane computes out_per_lane outputs
    n_idx = compute_lane % UOp.const(dtypes.int, grp_sub)
    c_grp = compute_lane // UOp.const(dtypes.int, grp_sub)

    for out_reg in range(out_per_lane):
-      acc_v = ctx.raccvgpr_dyn(src2_r + _c(out_reg), compute_lane, src2_is_vgpr)
+      acc_v = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(src2_r + _c(out_reg), compute_lane, src2_is_vgpr)
      if is_int_out: acc_v = acc_v.cast(dtypes.int32)
      else: acc_v = acc_v.bitcast(dtypes.float32)
      acc = src2_is_vgpr.where(acc_v, acc_scalar)
@@ -1544,9 +1548,11 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
          acc = acc + a_val * b_val

      if is_int_out:
-        compute_stores.append(ctx.waccvgpr_dyn(vdst_reg + _c(out_reg), compute_lane, acc.cast(dtypes.uint32), exec_mask))
+        compute_stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(
+          vdst_reg + _c(out_reg), compute_lane, acc.cast(dtypes.uint32), exec_mask))
      else:
-        compute_stores.append(ctx.waccvgpr_dyn(vdst_reg + _c(out_reg), compute_lane, acc.bitcast(dtypes.uint32), exec_mask))
+        compute_stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(
+          vdst_reg + _c(out_reg), compute_lane, acc.bitcast(dtypes.uint32), exec_mask))

  compute_phase = UOp.group(*compute_stores).end(compute_lane)
  return UOp.sink(read_phase, compute_phase, *ctx.inc_pc())
@@ -134,14 +134,14 @@ class TestBitcastConstFolding(unittest.TestCase):
 # folds advance indexing into basic indexing
 class TestIndexingConstFolding(unittest.TestCase):
  def test_scalar_index(self):
-    t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
+    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
    _check_ast_count(1, t[:,:,Tensor(1),:])
    _check_ast_count(1, t[:,:,Tensor(1)+2,:])
    _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])

  def test_const_tensor_index(self):
    # TODO: these can be 0, implement const tensor folded indexing
-    t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
+    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
    _check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:])
    _check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:])
    _check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)])
@@ -28,8 +28,8 @@ class TestDevice(unittest.TestCase):
  def test_nonexistent_renderer(self):
    with self.assertRaisesRegex(RuntimeError, "has no renderer"):
      with Context(DEV="CPU:TYPO"): Device[Device.DEFAULT].renderer
-    with self.assertRaisesRegex(RuntimeError, "did you mean: 'CLANGJIT'"):
-      with Context(DEV="CPU:CLANG"): Device[Device.DEFAULT].renderer
+    with self.assertRaisesRegex(RuntimeError, "did you mean: 'CLANG'"):
+      with Context(DEV="CPU:CLANGJIT"): Device[Device.DEFAULT].renderer

  @unittest.skipIf(Device.DEFAULT != "AMD", "only run on AMD")
  def test_nonexistent_iface(self):
@@ -69,17 +69,17 @@ class TestDevice(unittest.TestCase):
  @unittest.skipIf(WIN, "skipping windows test") # TODO: subprocess causes memory violation?
  def test_env_overwrite_default_compiler(self):
    if Device.DEFAULT == "CPU":
-      from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler
-      try: _, _ = CPULLVMCompiler(), ClangJITCompiler()
+      from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangCompiler
+      try: _, _ = CPULLVMCompiler(), ClangCompiler()
      except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")

-      imports = "from tinygrad import Device; from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler"
+      imports = "from tinygrad import Device; from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangCompiler"
      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, CPULLVMCompiler)"'],
                        shell=True, check=True, env={**os.environ, "DEV": "CPU:LLVM"})
-      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangJITCompiler)"'],
+      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangCompiler)"'],
                        shell=True, check=True, env={**os.environ, "DEV": "CPU"})
-      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangJITCompiler)"'],
-                        shell=True, check=True, env={**os.environ, "DEV": "CPU:CLANGJIT"})
+      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangCompiler)"'],
+                        shell=True, check=True, env={**os.environ, "DEV": "CPU:CLANG"})
    elif Device.DEFAULT == "AMD":
      from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
      try: _, _ = HIPCompiler(Device[Device.DEFAULT].arch), AMDLLVMCompiler(Device[Device.DEFAULT].arch)
@@ -96,15 +96,15 @@ class TestDevice(unittest.TestCase):

  @unittest.skipIf(WIN, "skipping windows test")
  def test_env_online(self):
-    from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler
-    try: _, _ = CPULLVMCompiler(), ClangJITCompiler()
+    from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangCompiler
+    try: _, _ = CPULLVMCompiler(), ClangCompiler()
    except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")

    with Context(DEV="CPU:LLVM"):
      inst = Device["CPU"].compiler
      self.assertIsInstance(Device["CPU"].compiler, CPULLVMCompiler)
    with Context(DEV="CPU"):
-      self.assertIsInstance(Device["CPU"].compiler, ClangJITCompiler)
+      self.assertIsInstance(Device["CPU"].compiler, ClangCompiler)
    with Context(DEV="CPU:LLVM"):
      self.assertIsInstance(Device["CPU"].compiler, CPULLVMCompiler)
      assert inst is Device["CPU"].compiler  # cached
@@ -118,7 +118,7 @@ class TestDevice(unittest.TestCase):

    dev = Device["CPU"]
    dev.cached_renderer.clear()
-    with patch("tinygrad.renderer.cstyle.ClangJITRenderer.__init__", side_effect=RuntimeError("broken")):
+    with patch("tinygrad.renderer.cstyle.ClangRenderer.__init__", side_effect=RuntimeError("broken")):
      self.assertIsInstance(dev.renderer.compiler, CPULLVMCompiler)

  def test_dev_contextvar(self):
@@ -1,5 +1,5 @@
 import unittest, subprocess, platform
-from tinygrad.runtime.support.compiler_cpu import ClangJITCompiler
+from tinygrad.runtime.support.compiler_cpu import ClangCompiler
 from tinygrad.runtime.support.elf import elf_loader

 class TestElfLoader(unittest.TestCase):
@@ -23,7 +23,7 @@ class TestElfLoader(unittest.TestCase):
      }
    '''
    with self.assertRaisesRegex(RuntimeError, 'evil_external_function'):
-      ClangJITCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src)
+      ClangCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src)
  def test_link(self):
    src = '''
      float powf(float, float); // from libm
@@ -96,7 +96,7 @@ class TestGroupedDims(unittest.TestCase):

  def test_global_prod_max(self):
    g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL)
-    sink = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0).index(g + l).store(UOp.const(dtypes.float, 1.0)).end(g, l).sink(arg=KernelInfo())
+    sink = UOp.param(0, dtypes.float.ptr()).index(g + l).store(UOp.const(dtypes.float, 1.0)).end(g, l).sink(arg=KernelInfo())
    class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128)
    specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL]
    self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1)
@@ -2,7 +2,7 @@ import ctypes, gzip, unittest, timeit, pickle
 from tinygrad import Variable
 from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, get_contraction, count, all_same
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
-from tinygrad.helpers import ceildiv
+from tinygrad.helpers import ceildiv, ansistrip
 from tinygrad.tensor import Tensor, get_shape
 import numpy as np

@@ -445,6 +445,13 @@ class TestWordWrap(unittest.TestCase):
    st2 = word_wrap(st, wrap=wrap)
    self.assertEqual(len(st2.splitlines()), 2)

+  def test_wrap_colored_at_boundary(self):
+    wrap = 10
+    st = "x"*(wrap-2) + colored("yyy", "red")
+    st2 = word_wrap(st, wrap=wrap)
+    self.assertEqual(ansistrip(st2), "x"*(wrap-2)+"yy\ny")
+    self.assertNotIn("\x1b[\n", st2)
+
  def test_wrap_explicit_newline(self):
    wrap = 10
    st = "\n".join(["x"*wrap, "x"*wrap, "x"*wrap])
@@ -7,14 +7,14 @@ from tinygrad.codegen import to_program

 class TestLinearizerFailures(unittest.TestCase):
  def test_fail_1(self):
-    c0 = UOp(Ops.PARAM, dtypes.float.ptr(64), arg=0, src=())
+    c0 = UOp.param(0, dtypes.float.ptr(64))
    c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.LOOP)
    c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.LOOP)
    c3 = ((c1*UOp.const(dtypes.weakint, 32))+c2)
-    c4 = UOp(Ops.PARAM, dtypes.float.ptr(163840), arg=1, src=())
+    c4 = UOp.param(1, dtypes.float.ptr(163840))
    c5 = UOp.range(UOp.const(dtypes.weakint, 2560), 0, AxisType.REDUCE)
    c6 = c4.index(((((((c5//UOp.const(dtypes.weakint, 8))%UOp.const(dtypes.weakint, 8))*UOp.const(dtypes.weakint, 8))+(c5%UOp.const(dtypes.weakint, 8)))+(((c2*UOp.const(dtypes.weakint, 40))+(c5//UOp.const(dtypes.weakint, 64)))*UOp.const(dtypes.weakint, 64)))+(c1*UOp.const(dtypes.weakint, 81920))))
-    c7 = UOp(Ops.PARAM, dtypes.float.ptr(64), arg=2, src=())
+    c7 = UOp.param(2, dtypes.float.ptr(64))
    c8 = c7.index(c3)
    c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
    c10 = c0.index(c3).store(c9).end(c1, c2)
@@ -458,12 +458,12 @@ class TestSchedule(unittest.TestCase):
  def test_fold_conv_batchnorm_optim(self, adam=False):
    optim, cnt = (nn.optim.Adam, 29) if adam else (nn.optim.SGD, 15)
    with Tensor.train():
-      img = Tensor.ones(1,3,4,4).realize()
+      img = Tensor.ones(1,3,4,4)
      c1 = nn.Conv2d(3,32,3)
      bn = nn.BatchNorm2d(32, track_running_stats=False)
      _realize_weights([c1, bn])
      opt = optim(nn.state.get_parameters([c1, bn]))
-      Tensor.realize(*nn.state.get_parameters(opt))
+      Tensor.realize(img, *nn.state.get_parameters(opt))
      img_bn = bn(c1(img)).elu().sum()
      opt.zero_grad()
      img_bn.backward()
@@ -968,7 +968,7 @@ class TestSchedule(unittest.TestCase):

  def test_const_schedule_contig(self):
    constv = Tensor.empty(2, 2).uop.const_like(10).contiguous()
-    check_schedule(constv, 1)
+    check_schedule(constv, 0)

  def test_advanced_simple_indexing_combined(self):
    X = Tensor.arange(16).reshape(4, 4)
@@ -15,13 +15,13 @@ def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move

 def get_gated_load_uop(valid:UOp, idx:UOp):
  return UOp(Ops.LOAD, dtypes.float, (
-    UOp(Ops.PARAM, dtypes.float.ptr(), arg=0).index(idx.valid(valid), ptr=True),
+    UOp.param(0, dtypes.float.ptr()).index(idx.valid(valid), ptr=True),
    UOp.const(dtypes.float, 0.0)
  ))

 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
  return UOp(Ops.LOAD, dtypes.float.vec(4), (
-    UOp(Ops.PARAM, dtypes.imagef(image_shape), arg=0).index(idx[1].valid(valid), idx[0].valid(valid), ptr=True),
+    UOp.param(0, dtypes.imagef(image_shape)).index(idx[1].valid(valid), idx[0].valid(valid), ptr=True),
    UOp(Ops.STACK, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4)
  ))

@@ -292,8 +292,7 @@ class TestImageSimplification(unittest.TestCase):
    load = get_load_image_uop(shape, (gidx0<8) & (gidx0<8).ne(True), idx)
    with Context(NOOPT=1, SPEC=0):
      load = full_rewrite(load.sink()).src[0]
-    self.assertEqual(load.op, Ops.STACK)
-    self.assertEqual(load.dtype.count, 4)
+    self.assertFalse(load.op_in_backward_slice_with_self(Ops.LOAD))

  def test_openpilot_conv1(self):
    # first conv in openpilot
@@ -513,7 +512,7 @@ class TestDropTrueGate(unittest.TestCase):
    from tinygrad.codegen.late.devectorizer import load_store_indexing
    from tinygrad.uop.ops import graph_rewrite
    from tinygrad.uop.symbolic import sym
-    buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
+    buf = UOp.param(0, dtypes.int.ptr())
    idx = UOp.const(dtypes.weakint, 0)
    true_gate = UOp.const(dtypes.bool, True)
    index_with_gate = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx.valid(true_gate)))
@@ -557,7 +556,7 @@ class TestRangeShrink(unittest.TestCase):
    # one load guards r < 4, but another load uses r without a gate -> no shrink
    r = Range(0, 204)
    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
-    load2 = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.float.ptr(), arg=1).index(r, ptr=True),))
+    load2 = UOp(Ops.LOAD, dtypes.float, (UOp.param(1, dtypes.float.ptr()).index(r, ptr=True),))
    ranges = self.get_ranges(UOp.sink(load1, load2))
    self.assertEqual(len(ranges), 1)
    self.assertEqual(ranges[0].src[0].arg, 204)
@@ -583,7 +582,7 @@ class TestRangeShrink(unittest.TestCase):
    from tinygrad.dtype import Invalid
    r = Range(0, 204)
    x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
-    ranges = self.get_ranges(UOp(Ops.PARAM, dtypes.float.ptr(), arg=0).index(r).store((r < 4).where(x, 0)).sink())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float.ptr()).index(r).store((r < 4).where(x, 0)).sink())
    self.assertEqual(len(ranges), 1)
    self.assertEqual(ranges[0].src[0].arg, 4)

@@ -592,7 +591,7 @@ class TestRangeShrink(unittest.TestCase):
    from tinygrad.dtype import Invalid
    r = Range(0, 204)
    x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
-    ranges = self.get_ranges(UOp(Ops.PARAM, dtypes.float.ptr(), arg=0).index(r).store((r < 4).where(0, x)).sink())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float.ptr()).index(r).store((r < 4).where(0, x)).sink())
    self.assertEqual(len(ranges), 1)
    self.assertEqual(ranges[0].src[0].arg, 4)

@@ -73,12 +73,12 @@ class TestIdxUpcast(unittest.TestCase):
    # Assert the dtype of the INDEX value, This will need be updated if UOp spec changes
    store = next(uop for uop in uops if uop.op is Ops.STORE)
    assert store.op is Ops.STORE
-    idx = self._find_op(store, Ops.SLICE)
-    # PTX and NIR turn Ops.SLICE into pointer arithmetic earlier than cstyle, plus it's already cast to int64
+    idx = self._find_op(store, Ops.INDEX)
+    # PTX and NIR turn Ops.INDEX into pointer arithmetic earlier than cstyle, plus it's already cast to int64
    if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)):
-      assert idx.op is Ops.SLICE
+      assert idx.op is Ops.INDEX
      idx_val = idx.src[1]
-      self.assertIs(idx_val.dtype, dtype)
+      self.assertFalse(idx_val.overflows(idx_val.dtype.base.scalar()))

  # use expand to generate kernel that uses large idx
  def do_op_then_assert(self, dtype: DType, dim1, dim2, dim3):
@@ -3,7 +3,7 @@ from tinygrad import Tensor, dtypes
 from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite

 _strip_unique_pm = PatternMatcher([
-  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE, name="d")), name="b"), lambda b,d: b.replace(src=(UOp.unique(0), d))),
+  (UPat((Ops.UNIQUE, Ops.LUNIQUE), name="u"), lambda u: u.replace(arg=0) if u.arg != 0 else None),
 ])
 def _strip_unique(u: UOp) -> UOp: return graph_rewrite(u, _strip_unique_pm)

@@ -1,7 +1,7 @@
 import unittest, math
 import numpy as np
 from tinygrad import dtypes
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.uop.ops import UOp
 from tinygrad.uop.decompositions import TRANSCENDENTAL_DTYPES, payne_hanek_reduction, cody_waite_reduction
 from tinygrad.uop.decompositions import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if
 from test.helpers import eval_uop
@@ -10,8 +10,8 @@ class TestTranscendentalFunctions(unittest.TestCase):
  def test_payne_hanek_reduction(self):
    # TODO: Test constant input when constant folding is fixed (or maybe test both variants)
    # Load input value from a buffer to prevent constant folding
-    input_buf = UOp(Ops.PARAM, dtypes.double.ptr(), arg=1, src=())
-    loaded_value = input_buf.index(UOp.const(dtypes.int, 0))
+    input_buf = UOp.param(1, dtypes.double.ptr(1))
+    loaded_value = input_buf.index(UOp.const(dtypes.int, 0), ptr=True).load()
    def eval_payne_hanek_reduction(v:float) -> tuple[float, int]:
      return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value))

@@ -105,8 +105,8 @@ class TestModularWraparound(unittest.TestCase):

 class TestGraphRewrite(unittest.TestCase):
  def test_dedup(self):
-    v1 = UOp(Ops.DEFINE_VAR, dtypes.float)
-    v2 = UOp(Ops.DEFINE_VAR, dtypes.float)
+    v1 = UOp.variable("v", 0, 1, dtypes.float)
+    v2 = UOp.variable("v", 0, 1, dtypes.float)
    nout = graph_rewrite(v1+v2, PatternMatcher([]))
    self.assertIs(nout.src[0], nout.src[1])

@@ -166,7 +166,7 @@ class TestGraphRewrite(unittest.TestCase):
    self.assertEqual(nout.arg, 3.0)

  def test_depth_2_fold(self):
-    v = UOp(Ops.DEFINE_VAR, dtypes.float)
+    v = UOp.variable("v", 0, 1, dtypes.float)
    c1 = UOp.const(dtypes.float, 1.0)
    c2 = UOp.const(dtypes.float, 2.0)
    nout = graph_rewrite(v+c1+c2, simple_pm)
@@ -191,7 +191,7 @@ class TestGraphRewrite(unittest.TestCase):
    b = UOp.variable('b', 0, 1)
    c = UOp.variable('c', 0, 1)
    d = UOp.variable('d', 0, 1)
-    outs = [2+a, 2+a+d+3+b+c+4, UOp(Ops.ADD, a.dtype, src=(a.const_like(2), a)), (4+d)+c+(2+a)+b]
+    outs = [2+a, 2+a+d+3+b+c+4, a.const_like(2)+a, (4+d)+c+(2+a)+b]
    for out in outs:
      sink = graph_rewrite(out, sym)
      print(sink.render())
@@ -203,7 +203,7 @@ class TestUOpGraph(unittest.TestCase):
  def test_add_constant_fold(self):
    c1 = UOp.const(dtypes.float, 1.0)
    c2 = UOp.const(dtypes.float, 2.0)
-    out = UOp(Ops.ADD, dtypes.float, (c1, c2))
+    out = c1+c2
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
@@ -213,9 +213,9 @@ class TestUOpGraph(unittest.TestCase):
  def test_where_same_fold(self):
    v = UOp.variable('tmp', 0, 1)
    c0 = UOp.const(dtypes.weakint, 0)
-    vc = UOp(Ops.CMPNE, dtypes.bool, (v, c0))
+    vc = v != c0
    c1 = UOp.const(dtypes.float, 1.0)
-    out = UOp(Ops.WHERE, dtypes.float, (vc, c1, c1))
+    out = vc.where(c1, c1)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
@@ -226,7 +226,7 @@ class TestUOpGraph(unittest.TestCase):
    bf = UOp.const(dtypes.bool, False)
    c1 = UOp.const(dtypes.float, 1.0)
    c2 = UOp.const(dtypes.float, 2.0)
-    out = UOp(Ops.WHERE, dtypes.float, (bf, c1, c2))
+    out = bf.where(c1, c2)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
@@ -235,7 +235,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_const_cast(self):
    bf = UOp.const(dtypes.bool, False)
-    out = UOp(Ops.CAST, dtypes.int, (bf,))
+    out = bf.cast(dtypes.int)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
@@ -244,7 +244,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_const_bitcast(self):
    bf = UOp.const(dtypes.float, 1.0)
-    out = UOp(Ops.BITCAST, dtypes.uint32, (bf,))
+    out = bf.bitcast(dtypes.uint32)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
@@ -254,15 +254,15 @@ class TestUOpGraph(unittest.TestCase):
  @unittest.expectedFailure
  def test_const_shape_change_bitcast(self):
    bf = UOp.const(dtypes.uint8, 0x3F)
-    out = UOp(Ops.BITCAST, dtypes.half, (bf,))
+    out = bf.bitcast(dtypes.half)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK

  @unittest.skip("this test isn't valid uops")
  def test_noop_vectorize_fold(self):
-    d0 = UOp(Ops.PARAM, dtypes.float.ptr(), arg=0)
+    d0 = UOp.param(0, dtypes.float.ptr())
    idx = UOp.const(dtypes.int, 0)
-    ld = UOp(Ops.LOAD, dtypes.float.vec(2), (d0, idx))
+    ld = d0.load(idx, dtype=dtypes.float.vec(2))
    vec = UOp(Ops.STACK, dtypes.float.vec(2), (ld,))
    x = UOp(Ops.GEP, dtypes.float, (vec, ), arg=0)
    alu = UOp(Ops.SQRT, dtypes.float, (x, ))
@@ -272,13 +272,13 @@ class TestUOpGraph(unittest.TestCase):

  @unittest.skip("this test isn't valid uops")
  def test_gep_vec_fold(self):
-    d0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
-    d1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
-    d2 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 2)
+    d0 = UOp.param(0, dtypes.float.ptr())
+    d1 = UOp.param(1, dtypes.float.ptr())
+    d2 = UOp.param(2, dtypes.float.ptr())
    idx = UOp.const(dtypes.int, 0)
    def _test_vec(geps, count=4):
      vec = UOp(Ops.STACK, dtypes.float.vec(count), geps)
-      out = UOp(Ops.STORE, dtypes.void, (d0.index(idx), vec))
+      out = d0.index(idx).store(vec)
      uops = to_uops_list([out])
      if DEBUG >= 4:
        from tinygrad import Device
@@ -286,28 +286,28 @@ class TestUOpGraph(unittest.TestCase):
      return uops[-2].src[-1]  # -2 to skip SINK

    # possible
-    val = UOp(Ops.LOAD, dtypes.float.vec(4), (d1.index(idx),))
-    xyzw = tuple(UOp(Ops.GEP, dtypes.float, (val,), (i,)) for i in range(4))
+    val = d1.index(idx).load(dtype=dtypes.float.vec(4))
+    xyzw = tuple(val.gep(i) for i in range(4))
    self.assertIs(_test_vec(xyzw).op, Ops.LOAD)

    # unaligned
-    val = UOp(Ops.LOAD, dtypes.float.vec(4), (d1.index(idx),))
-    wzyx = tuple(UOp(Ops.GEP, dtypes.float, (val,), (i,)) for i in reversed(range(4)))
+    val = d1.index(idx).load(dtype=dtypes.float.vec(4))
+    wzyx = tuple(val.gep(i) for i in reversed(range(4)))
    self.assertIs(_test_vec(wzyx).op, Ops.STACK)

    # different_size
-    val = UOp(Ops.LOAD, dtypes.float.vec(2), (d1.index(idx),))
-    xy = tuple(UOp(Ops.GEP, dtypes.float, (val, ), (i,)) for i in range(2))
+    val = d1.index(idx).load(dtype=dtypes.float.vec(2))
+    xy = tuple(val.gep(i) for i in range(2))
    self.assertIs(_test_vec(xy+xy).op, Ops.STACK)
-    val = UOp(Ops.LOAD, dtypes.float.vec(4), (d1.index(idx),))
-    xy = tuple(UOp(Ops.GEP, dtypes.float, (val, ), (i,)) for i in range(2))
+    val = d1.index(idx).load(dtype=dtypes.float.vec(4))
+    xy = tuple(val.gep(i) for i in range(2))
    self.assertIs(_test_vec(xy, count=2).op, Ops.STACK)

    # different vals
-    val1 = UOp(Ops.LOAD, dtypes.float.vec(2), (d1.index(idx),))
-    val2 = UOp(Ops.LOAD, dtypes.float.vec(2), (d2.index(idx),))
-    xy1 = tuple(UOp(Ops.GEP, dtypes.float, (val1, ), (i,)) for i in range(2))
-    xy2 = tuple(UOp(Ops.GEP, dtypes.float, (val2, ), (i,)) for i in range(2))
+    val1 = d1.index(idx).load(dtype=dtypes.float.vec(2))
+    val2 = d2.index(idx).load(dtype=dtypes.float.vec(2))
+    xy1 = tuple(val1.gep(i) for i in range(2))
+    xy2 = tuple(val2.gep(i) for i in range(2))
    self.assertIs(_test_vec(xy1+xy2).op, Ops.STACK)

  def test_gep_vec_const_fold(self):
@@ -323,7 +323,7 @@ class TestUOpGraph(unittest.TestCase):
  def test_wmma_vectorize_fold(self):
    for i in [2, 4, 8]:
      vec = UOp(Ops.STACK, dtypes.half.vec(i), tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
-      var = UOp(Ops.DEFINE_VAR, dtypes.half.vec(i))
+      var = UOp.variable("var", 0, 1, dtypes.half.vec(i))
      acc = UOp.variable('acc', 0, 1, dtypes.half.vec(i))
      wmma = UOp(Ops.WMMA, dtypes.half.vec(i), (vec, var, acc))
      uops = to_uops_list([wmma])
@@ -331,7 +331,7 @@ class TestUOpGraph(unittest.TestCase):
      self.assertEqual(len(uops), 2)  # +1 for SINK

    for i in [2, 4, 8]:
-      var = UOp(Ops.DEFINE_VAR, dtypes.half.vec(i))
+      var = UOp.variable("var", 0, 1, dtypes.half.vec(i))
      vec = UOp(Ops.STACK, dtypes.half.vec(i), tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
      acc = UOp.variable('acc', 0, 1, dtypes.half.vec(i))
      wmma = UOp(Ops.WMMA, dtypes.half.vec(i), (var, vec, acc))
@@ -380,22 +380,22 @@ class TestUOpGraph(unittest.TestCase):
      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK

  def test_cast_alu_fold(self):
-    d0 = UOp(Ops.PARAM, dtypes.bool.ptr(), arg=0)
-    d1 = UOp(Ops.PARAM, dtypes.int.ptr(), arg=1)
+    d0 = UOp.param(0, dtypes.bool.ptr(1))
+    d1 = UOp.param(1, dtypes.int.ptr(1))
    idx = UOp.const(dtypes.int, 0)
    ld = d1.index(idx)
    alu = (ld<1).cast(dtypes.bool)
-    out = UOp(Ops.STORE, dtypes.void, (d0.index(idx), alu))
+    out = d0.index(idx, ptr=True).store(alu)
    uops = to_uops_list([out])
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)

  def test_double_cast_fold(self):
-    d0 = UOp(Ops.PARAM, dtypes.float.ptr(), arg=0)
-    d1 = UOp(Ops.PARAM, dtypes.int.ptr(), arg=1)
+    d0 = UOp.param(0, dtypes.float.ptr(1))
+    d1 = UOp.param(1, dtypes.int.ptr(1))
    idx = UOp.const(dtypes.int, 0)
    ld = d1.index(idx)
    alu = ld.cast(dtypes.float).cast(dtypes.float)
-    out = UOp(Ops.STORE, dtypes.void, (d0.index(idx), alu))
+    out = d0.index(idx, ptr=True).store(alu)
    uops = to_uops_list([out])
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)

@@ -403,8 +403,8 @@ class TestUOpGraph(unittest.TestCase):
    v = UOp.variable("tmp", 0, 1, dtypes.int)
    c2 = UOp.const(dtypes.int, 2)
    c4 = UOp.const(dtypes.int, 4)
-    vc = UOp(Ops.ADD, dtypes.int, (v, c2))
-    out = UOp(Ops.ADD, dtypes.int, (vc, c4))
+    vc = v+c2
+    out = vc+c4
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 4)  # +1 for SINK
    out = uops[-2]  # -2 to skip SINK
@@ -414,7 +414,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_bitcast_to_same_dtype_fold(self):
    for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
-      d0 = UOp(Ops.PARAM, dt.ptr(), arg=0)
+      d0 = UOp.param(0, dt.ptr(1))
      v = d0.index(UOp.const(dtypes.int, 0))
      uops = to_uops_list([v.bitcast(dt)])
      self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}")
@@ -427,18 +427,18 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_fold(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.PARAM, dtypes.long.ptr(), (), 0)
+    d0 = UOp.param(0, dtypes.long.ptr(100))
    ld = d0.index(ridx0.valid(ridx0<50))
    w = (ridx0<50).where(ld, 5)
-    out = UOp(Ops.PARAM, dtypes.long.ptr(), (), 1)
-    uops = to_uops_list([out.index(ridx0).store(w)])
+    out = UOp.param(1, dtypes.long.ptr(100))
+    uops = to_uops_list([out.index(ridx0, ptr=True).store(w)])
    for u in uops:
      assert u.op is not Ops.WHERE
      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg==5

  def test_where_on_gated_load_folds_swapped_branches(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.PARAM, dtypes.long.ptr(), (), 0)
+    d0 = UOp.param(0, dtypes.long.ptr(100))
    ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
    w = (ridx0<50).where(5, ld)
    uops = to_uops_list([w])
@@ -448,40 +448,40 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_with_cast(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    d0 = UOp.param(0, dtypes.int.ptr(100))
    gate_idx = ridx0.valid((ridx0<50))
    ld = d0.index(gate_idx).cast(dtypes.float)
    w = (ridx0<50).where(ld, 5.0)
-    out = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
-    uops = to_uops_list([out.index(ridx0).store(w)])
+    out = UOp.param(1, dtypes.float.ptr(100))
+    uops = to_uops_list([out.index(ridx0, ptr=True).store(w)])
    for u in uops:
      assert u.op is not Ops.WHERE
      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg == 5

  def test_where_on_casted_gated_load_extra_cond(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
+    d0 = UOp.param(0, dtypes.float.ptr(100))
    ld = d0.index(ridx0.valid(ridx0<50))
    w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(dtypes.float, 0)).cast(dtypes.half)
-    out = UOp(Ops.PARAM, dtypes.half.ptr(), (), 1)
-    uops = to_uops_list([out.index(ridx0).store(w)])
+    out = UOp.param(1, dtypes.half.ptr(100))
+    uops = to_uops_list([out.index(ridx0, ptr=True).store(w)])
    for u in uops:
      assert u.op is not Ops.WHERE

  def test_where_on_casted_gated_load_extra_cond_swapped(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
+    d0 = UOp.param(0, dtypes.float.ptr(100))
    ld = d0.index(ridx0.valid(ridx0<50))
    w = ((ridx0<50) & (ridx0>30)).where(UOp.const(dtypes.float, 0), ld).cast(dtypes.half)
-    out = UOp(Ops.PARAM, dtypes.half.ptr(), (), 1)
-    uops = to_uops_list([out.index(ridx0).store(w)])
+    out = UOp.param(1, dtypes.half.ptr(100))
+    uops = to_uops_list([out.index(ridx0, ptr=True).store(w)])
    for u in uops:
      assert u.op is not Ops.WHERE

  def test_where_in_store_becomes_gate(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp(Ops.PARAM, dtypes.long.ptr(), (), 0)
-    idx = d0.index(ridx0)
+    d0 = UOp.param(0, dtypes.long.ptr(100))
+    idx = d0.index(ridx0, ptr=True)
    ld = idx.load()
    val = (ridx0<50).where(5, ld)
    st = idx.store(val).end(ridx0)
@@ -493,14 +493,14 @@ class TestUOpGraph(unittest.TestCase):
  def test_load_idx_becomes_int(self):
    # mnist indexing with split reduceop
    # Make sure we are not doign math on the loaded index, which would promote it to long
-    c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(128000), arg=0, src=())
+    c0 = UOp.param(0, dtypes.uchar.ptr(128000))
    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
-    c3 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
+    c3 = UOp.param(1, dtypes.int.ptr(512))
    c4 = c3.index(c1)
    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
-    c7 = UOp(Ops.PARAM, dtypes.uchar.ptr(60000), arg=2, src=())
+    c7 = UOp.param(2, dtypes.uchar.ptr(60000))
    c8 = c7.index(c6)
    c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
@@ -510,14 +510,14 @@ class TestUOpGraph(unittest.TestCase):

  def test_load_idx_no_math_on_loaded(self):
    # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
-    c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(128000), arg=0, src=())
+    c0 = UOp.param(0, dtypes.uchar.ptr(128000))
    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
-    c3 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
+    c3 = UOp.param(1, dtypes.int.ptr(512))
    c4 = c3.index(c1)  # c4 is a load
    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
-    c7 = UOp(Ops.PARAM, dtypes.uchar.ptr(60000), arg=2, src=())
+    c7 = UOp.param(2, dtypes.uchar.ptr(60000))
    c8 = c7.index(c6)
    # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
    loaded_idx = c4.cast(dtypes.weakint)
@@ -529,48 +529,42 @@ class TestUOpGraph(unittest.TestCase):
      self.assertNotEqual(u.dtype, dtypes.long)

  def test_fold_gated_load(self):
-    glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
-    glbl1 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 1)
-    glbl2 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 2)
+    glbl0 = UOp.param(0, dtypes.int.ptr(1))
+    glbl1 = UOp.param(1, dtypes.int.ptr(1))
+    glbl2 = UOp.param(2, dtypes.int.ptr(1))
    idx = UOp.const(dtypes.int, 0)
    ld0 = glbl1.index(UOp.invalid())
    ld1 = glbl2.index(idx.valid(UOp.const(dtypes.bool, True)))
-    uops = to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0.index(idx), ld1+ld0))])
-    ld0 = uops[-2].src[-1]  # -2 to skip SINK
+    uops = to_uops_list([glbl0.index(idx, ptr=True).store(ld1+ld0)])
    # the gate and invalid value are deleted from ld1
-    self.assertEqual(ld0, UOp.load(glbl2.slice(idx), dtype=dtypes.int))
+    self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 1)

  def test_fold_gated_load_local(self):
-    glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    glbl0 = UOp.param(0, dtypes.int.ptr(16))
    smem = UOp(Ops.DEFINE_LOCAL, dtypes.int.ptr(size=18, addrspace=AddrSpace.LOCAL), (), "temp")
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 16),), "lidx0")
-    st = UOp(Ops.STORE, dtypes.void, (smem.index(lidx, ptr=True), glbl0.index(lidx, ptr=True).load()))
-    barrier = UOp(Ops.BARRIER, dtypes.void, (st, ))
+    lidx = UOp.special(16, "lidx0", dtypes.int)
+    st = smem.index(lidx, ptr=True).store(glbl0.index(lidx, ptr=True).load())
+    barrier = st.barrier()
    ld0 = smem.after(barrier).index(UOp.invalid())
    ld1 = smem.after(barrier).index((lidx+2).valid(UOp.const(dtypes.bool, True)))
-    uops = to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0.index(lidx), ld1+ld0))])
+    uops = to_uops_list([glbl0.index(lidx, ptr=True).store(ld1+ld0)])

-    ld0 = uops[-2].src[-1]  # -2 to skip SINK
    # the gate and invalid value are deleted from ld1
-    new_barrier = ld0.src[0].src[0].src[1]
-    assert new_barrier.op is Ops.BARRIER
-    self.assertEqual(ld0.src[0], smem.after(new_barrier).slice(lidx+2))
+    self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 2)

  def test_fold_gated_store(self):
-    glbl = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    glbl = UOp.param(0, dtypes.int.ptr(1))
    idx0 = UOp.const(dtypes.int, 0)
-    idx1 = UOp.const(dtypes.int, 0)
    val = UOp.const(dtypes.int, 42)
    st0 = glbl.index(UOp.invalid(), ptr=True).store(val)
    st1 = glbl.index(idx0.valid(UOp.const(dtypes.bool, True)), ptr=True).store(val)
    uops = to_uops_list([st0, st1])
    # only the second store happens
-    self.assertEqual(len(uops), 6)  # +1 for SINK
-    self.assertEqual(uops[-2], glbl.slice(idx1).store(val))  # -2 to skip SINK
+    self.assertEqual(len([u for u in uops if u.op is Ops.STORE]), 1)

  @unittest.skip("this is a uop type error")
  def test_asserts_bad_gate(self):
-    glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
+    glbl0 = UOp.param(0, dtypes.int.ptr())
    idx = UOp.const(dtypes.int, 0)
    bad_gate = UOp.const(dtypes.int, 1)
    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0, idx, UOp.const(dtypes.int, 42), bad_gate))])
@@ -781,7 +775,7 @@ class TestLoadStoreFolding(unittest.TestCase):
  def test_gated_load_gep_preserves_alt(self):
    """Test that LOAD(GEP, alt) preserves alt value after rewrite"""
    from tinygrad.codegen.late.devectorizer import load_store_folding
-    buf = UOp(Ops.PARAM, dtypes.float.vec(4).ptr(), (), 0)
+    buf = UOp.param(0, dtypes.float.vec(4).ptr())
    idx = UOp.const(dtypes.int, 0)
    gate = UOp.const(dtypes.bool, True)
    gated_index = buf.index(idx.valid(gate))
@@ -799,8 +793,8 @@ class TestLoadStoreFolding(unittest.TestCase):
  def test_gated_load_ptrcat_preserves_alt(self):
    """Test that LOAD(PTRCAT, alt) preserves alt value after rewrite"""
    from tinygrad.codegen.late.devectorizer import load_store_folding
-    buf1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
-    buf2 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
+    buf1 = UOp.param(0, dtypes.float.ptr())
+    buf2 = UOp.param(1, dtypes.float.ptr())
    idx = UOp.const(dtypes.int, 0)
    idx1 = buf1.index(idx)
    idx2 = buf2.index(idx)
@@ -951,8 +951,8 @@ class TestSymbolic(unittest.TestCase):
    expr = cond.where(a, b).cast(dtypes.half)

    # TODO: copied from render, render does not support cast
-    glbl = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
-    uops = get_uops(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
+    glbl = UOp.param(0, dtypes.int.ptr(1))
+    uops = get_uops(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0), ptr=True), expr)).sink())
    rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]

    self.assertEqual(rewritten_uop, cond.where(a.cast(dtypes.half), b.cast(dtypes.half)))
@@ -1270,7 +1270,7 @@ class TestStoreLoadFolding(unittest.TestCase):
  """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
  def test_store_load_folding(self):
    # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0)
-    buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
+    buf = UOp.param(0, dtypes.int.ptr())
    index = buf.index(UOp.const(dtypes.weakint, 0))
    # Direct: store(idx, load(idx)) -> NOOP
    self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
@@ -1340,7 +1340,7 @@ class TestRangeSplitting(unittest.TestCase):
    from tinygrad.codegen.simplify import pm_split_ranges, pm_flatten_range
    r0 = UOp.range(uconst(8), 0)
    # create a simple expression using the range with mod: store range%2 to a buffer
-    buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
+    buf = UOp.param(0, dtypes.int.ptr())
    val = (r0 % uconst(2)).cast(dtypes.int)
    store = UOp(Ops.STORE, dtypes.void, (buf.index(uconst(0)), val))
    sink = UOp(Ops.SINK, dtypes.void, (UOp(Ops.END, dtypes.void, (store, r0)),))
@@ -82,7 +82,7 @@ class TestVminVmaxProperties(unittest.TestCase):
  def test_vmin_vmax_multiplication_0_inf(self):
    # vmin and vmax for multiplication with a variable
    x = UOp.const(dtypes.float, 0.0)
-    y = UOp.load(UOp(Ops.PARAM, dtypes.float.ptr(1), (), 0), UOp.const(dtypes.int, 0), dtype=dtypes.float)
+    y = UOp.load(UOp.param(0, dtypes.float.ptr(1)), UOp.const(dtypes.int, 0), dtype=dtypes.float)
    uop = x * y
    # TODO: these should be 0, but definitely should not be nan
    self.assertEqual(uop.vmin, -math.inf)
@@ -316,7 +316,7 @@ class TestVminVmaxVConst(unittest.TestCase):

  def test_vmin_vmax_vector_with_gep(self):
    # vmin and vmax for a vector constant of bool values
-    d1 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 1)
+    d1 = UOp.param(1, dtypes.int.ptr())
    idx = UOp.const(dtypes.int, 0)
    val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx).cast(dtypes.int.vec(2).ptr()),))
    uop = (val // 32).gep(0)
@@ -110,10 +110,10 @@ class TestExecALU(unittest.TestCase):

 class TestGatedStoreRewrite(unittest.TestCase):
  def test_tiny_gate_store(self):
-    gmem = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
+    gmem = UOp.param(0, dtypes.float.ptr(8))
    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
    gate = gidx0<UOp.const(dtypes.int, 1)
-    idx = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem, (gidx0 * UOp.const(dtypes.int, 2)).valid(gate)))
+    idx = UOp(Ops.INDEX, dtypes.float.ptr(8), (gmem, (gidx0 * UOp.const(dtypes.int, 2)).valid(gate)))
    val = UOp.const(dtypes.float, 42.0)
    store = UOp(Ops.STORE, dtypes.void, (idx, val))
    uops = to_uops_list([store])
@@ -126,12 +126,12 @@ class TestGatedStoreRewrite(unittest.TestCase):
    self.assertEqual(len(gated_uops[-1].src), 2)

  def test_gate_some_stores(self):
-    gmem0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
-    gmem1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
+    gmem0 = UOp.param(0, dtypes.float.ptr(8))
+    gmem1 = UOp.param(1, dtypes.float.ptr(8))
    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
    idx = gidx0 * UOp.const(dtypes.int, 2)
-    idx0 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem0, idx.valid(gidx0<UOp.const(dtypes.int, 1))))
-    idx1 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem1, idx))
+    idx0 = UOp(Ops.INDEX, dtypes.float.ptr(8), (gmem0, idx.valid(gidx0<UOp.const(dtypes.int, 1))))
+    idx1 = UOp(Ops.INDEX, dtypes.float.ptr(8), (gmem1, idx))
    val = UOp.const(dtypes.float, 42.0)
    stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
    uops = to_uops_list(stores)
@@ -146,13 +146,13 @@ class TestGatedStoreRewrite(unittest.TestCase):
  # scaled down version of TestLinearizerDumb.test_unmerged_ifs
  @unittest.skip("we don't merge ifs anymore")
  def test_merge_ifs_alt(self):
-    gmem0 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0)
-    gmem1 = UOp(Ops.PARAM, dtypes.float.ptr(), (), 1)
+    gmem0 = UOp.param(0, dtypes.float.ptr(8))
+    gmem1 = UOp.param(1, dtypes.float.ptr(8))
    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
    idx = gidx0*UOp.const(dtypes.int, 2)
    gate = gidx0<UOp.const(dtypes.int, 1)
-    idx0 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem0, idx.valid(gate)))
-    idx1 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem1, idx.valid(gate)))
+    idx0 = UOp(Ops.INDEX, dtypes.float.ptr(8), (gmem0, idx.valid(gate)))
+    idx1 = UOp(Ops.INDEX, dtypes.float.ptr(8), (gmem1, idx.valid(gate)))
    val = UOp.const(dtypes.float, 42.0)
    stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
    uops = to_uops_list(stores)
@@ -170,7 +170,7 @@ class TestGatedStoreRewrite(unittest.TestCase):
 class TestFastIdiv(unittest.TestCase):
  def test_division_power_of_two(self):
    for dt in (dtypes.int32, dtypes.uint32):
-      g = UOp(Ops.PARAM, dt.ptr(), (), 0)
+      g = UOp.param(0, dt.ptr(3))
      c = UOp.const(dt, 2)
      l = g.index(c)
      a = UOp(Ops.CDIV, dt, (l, c))
@@ -183,7 +183,7 @@ class TestFastIdiv(unittest.TestCase):
  def test_floormod_power_of_two(self):
    # FLOORMOD by a power of two lowers to AND (correct floor mod for any sign in two's complement)
    for dt in (dtypes.int32, dtypes.uint32):
-      g = UOp(Ops.PARAM, dt.ptr(), (), 0)
+      g = UOp.param(0, dt.ptr(9))
      c = UOp.const(dt, 8)
      a = UOp(Ops.FLOORMOD, dt, (g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -195,7 +195,7 @@ class TestFastIdiv(unittest.TestCase):
  def test_floordiv_power_of_two_uint(self):
    # uint FLOORDIV by a power of two lowers to a shift, leaving no IDIV/FLOORDIV in the kernel
    for dt in (dtypes.uint32, dtypes.uint64):
-      g = UOp(Ops.PARAM, dt.ptr(), (), 0)
+      g = UOp.param(0, dt.ptr(3))
      c = UOp.const(dt, 2)
      a = UOp(Ops.FLOORDIV, dt, (g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -207,7 +207,7 @@ class TestFastIdiv(unittest.TestCase):
  @Context(DISABLE_FAST_IDIV=0)
  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
  def test_fast_idiv_and_mod(self):
-    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
+    g = UOp.param(0, dtypes.uint32.ptr(4))
    c = UOp.const(dtypes.uint, 3)
    l = g.index(c)
    a = UOp(Ops.CDIV, dtypes.uint, (l, c))
@@ -242,7 +242,7 @@ class TestFastIdiv(unittest.TestCase):
  @unittest.expectedFailure
  def test_fast_idiv_overflow(self):
    # This will be possible with a slightly different method for fast_idiv
-    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
+    g = UOp.param(0, dtypes.uint32.ptr(8))
    c = UOp.const(dtypes.uint, 7)
    l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
    a = UOp(Ops.CDIV, dtypes.uint, (l, c))
@@ -253,7 +253,7 @@ class TestFastIdiv(unittest.TestCase):
    self.assertNotIn(Ops.CDIV, ops)

  def test_disable_fast_idiv(self):
-    g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
+    g = UOp.param(0, dtypes.uint32.ptr(4))
    c = UOp.const(dtypes.uint, 3)
    l = g.index(c)
    a = UOp(Ops.CDIV, dtypes.uint, (l, c))
@@ -290,8 +290,8 @@ class TestUOpMethod(unittest.TestCase):
    self.assertEqual((gidx0*3+1).const_factor(), 1)

  def test_replace(self):
-    x = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
-    self.assertIs(x.replace(arg=None).arg, None)
+    x = UOp.param(0, dtypes.int.ptr())
+    self.assertEqual(x.replace(arg=UOp.param(1, dtypes.int.ptr()).arg).arg.slot, 1)
    with self.assertRaises(AssertionError): x.replace(field="a")

  def test_const_zero_neg_zero_different(self):
@@ -139,7 +139,7 @@ class TestUOpsStats(unittest.TestCase):

  #MULACC should have the same stats as MUL + ADD
  def test_mulacc(self):
-    globl = UOp(Ops.PARAM, dtypes.int.ptr(), tuple())
+    globl = UOp.param(0, dtypes.int.ptr())
    o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
    o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
    u1 = globl.index(o1)
@@ -149,7 +149,7 @@ class TestUOpsStats(unittest.TestCase):
    u5 = UOp(Ops.ADD, dtypes.int, (u4,u3))
    uops = tuple(u5.toposort())

-    globl = UOp(Ops.PARAM, dtypes.int.ptr(), tuple())
+    globl = UOp.param(0, dtypes.int.ptr())
    o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
    o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
    u1 = globl.index(o1)
@@ -11,7 +11,7 @@ class TestValidateOOB(unittest.TestCase):
  # basic index patterns
  def test_const_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(16))
      to_uops_list([buf.index(UOp.const(dtypes.int, 0), ptr=True).load(dtype=dtypes.int)])  # valid
      to_uops_list([buf.index(UOp.const(dtypes.int, 15), ptr=True).load(dtype=dtypes.int)])  # valid (last element)
      with self.assertRaises(RuntimeError):
@@ -21,7 +21,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_variable_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(16))
      to_uops_list([buf.index(Variable("i", 0, 15), ptr=True).load(dtype=dtypes.int)])  # valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(Variable("i", 0, 20), ptr=True).load(dtype=dtypes.int)])  # oob
@@ -30,7 +30,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_range_with_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(16))
      r = UOp.range(42, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r.valid(r < 16), ptr=True).load(dtype=dtypes.int)])  # valid
      with self.assertRaises(RuntimeError):
@@ -38,7 +38,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_variable_with_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(16))
      v = Variable("v", -5, 80)
      to_uops_list([buf.index(v.valid((v >= 0) & (v < 16)), ptr=True).load(dtype=dtypes.int)])  # valid
      with self.assertRaises(RuntimeError):
@@ -46,7 +46,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_gated_store(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(16))
      v = Variable("v", 0, 20)
      to_uops_list([buf.index(v.valid(v < 16), ptr=True).store(0)])  # valid
      with self.assertRaises(RuntimeError):
@@ -55,14 +55,14 @@ class TestValidateOOB(unittest.TestCase):
  # ALU ops in index
  def test_floordiv(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(16))
      to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)])  # 0..16 oob

  def test_mod(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(16))
      r = UOp.range(100, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r % 16, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
      with self.assertRaises(RuntimeError):
@@ -70,14 +70,14 @@ class TestValidateOOB(unittest.TestCase):

  def test_shr(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(16))
      to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)])  # 0..31 oob

  def test_shl(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(64), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(64))
      r = UOp.range(8, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r << 2, ptr=True).load(dtype=dtypes.int)])  # 0..28 valid
      with self.assertRaises(RuntimeError):
@@ -85,7 +85,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_and(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(16))
      r = UOp.range(100, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r & 15, ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
      with self.assertRaises(RuntimeError):
@@ -93,14 +93,14 @@ class TestValidateOOB(unittest.TestCase):

  def test_max(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(16))
      to_uops_list([buf.index(Variable("v", -10, 15).maximum(0), ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0), ptr=True).load(dtype=dtypes.int)])  # 0..20 oob

  def test_xor_in_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(16))
      r = UOp.range(32, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16))), ptr=True).load(dtype=dtypes.int)])  # 0..15 valid
      with self.assertRaises(RuntimeError):
@@ -109,22 +109,22 @@ class TestValidateOOB(unittest.TestCase):
  # cast patterns
  def test_float_cast_in_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(16))
      r = UOp.range(20, 0)
      i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
      to_uops_list([buf.index(i.valid((i >= 0) & (i < 16)), ptr=True).load(dtype=dtypes.int)])

  def test_bool_cast_in_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp(Ops.PARAM, dtypes.int.ptr(1), (), 0)
+      buf = UOp.param(0, dtypes.int.ptr(1))
      r = UOp.range(20, 0)
      to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not()), ptr=True).load(dtype=dtypes.int)])  # only r=0 valid

  # load result as index/mask
  def test_load_as_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf0 = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
-      buf1 = UOp(Ops.PARAM, dtypes.int.ptr(64), (), 1)
+      buf0 = UOp.param(0, dtypes.int.ptr(16))
+      buf1 = UOp.param(1, dtypes.int.ptr(64))
      r = UOp.range(42, 0, AxisType.GLOBAL)
      ld0 = buf0.index(r.valid(r < 8), ptr=True).load(dtype=dtypes.int).cast(dtypes.weakint)
      to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32)), ptr=True).load(dtype=dtypes.int)])  # valid
@@ -133,8 +133,8 @@ class TestValidateOOB(unittest.TestCase):

  def test_load_bool_as_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf_bool = UOp(Ops.PARAM, dtypes.bool.ptr(16), (), 0)
-      buf_int = UOp(Ops.PARAM, dtypes.int.ptr(8), (), 1)
+      buf_bool = UOp.param(0, dtypes.bool.ptr(16))
+      buf_int = UOp.param(1, dtypes.int.ptr(8))
      gidx = UOp(Ops.SPECIAL, dtypes.weakint, (UOp.const(dtypes.weakint, 16),), "gidx0")
      ld_bool = buf_bool.index(gidx, ptr=True).load()
      with self.assertRaises(RuntimeError):
@@ -145,7 +145,7 @@ class TestValidateOOB(unittest.TestCase):
  def test_in_bounds_access_gated_local(self):
    with Context(CHECK_OOB=1):
      # Define buffers
-      gbuf = UOp(Ops.PARAM, dtypes.uint.ptr(400), (), 0)
+      gbuf = UOp.param(0, dtypes.uint.ptr(400))
      sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.uint.ptr(8, addrspace=AddrSpace.LOCAL), (), "temp0")

      # Define indices, valids and barrier
@@ -169,8 +169,8 @@ class TestValidateOOB(unittest.TestCase):
  @unittest.skip("Bool load is not supported yet")
  def test_load_mask(self):
    with Context(CHECK_OOB=1):
-      glbl0 = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
-      mask = UOp(Ops.PARAM, dtypes.bool.ptr(16), (), 0)
+      glbl0 = UOp.param(0, dtypes.int.ptr(16))
+      mask = UOp.param(0, dtypes.bool.ptr(16))
      ridx = UOp.range(20, 0)
      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask), ptr=True)))
      to_uops_list([ld0])
@@ -5,15 +5,14 @@ from typing import Generator

 from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, profile_matches
 from tinygrad.uop.symbolic import sym
-from tinygrad.dtype import dtypes
+from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.helpers import colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
 from tinygrad.helpers import cpu_profile, ProfilePointEvent, unwrap
 from tinygrad.device import Buffer

 from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, active_group, _name_cnt, RewriteTrace
-from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData, get_render
-from tinygrad.codegen import to_program_cache
-from tinygrad.codegen import to_program
+from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData, get_render, addrspace_colors
+from tinygrad.codegen import do_to_program

@track_rewrites(name=True)
 def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=None) -> UOp:
@@ -41,7 +40,6 @@ class VizTrace:
@contextlib.contextmanager
 def save_viz():
  for lst in [tracked_keys, tracked_ctxs, active_rewrites, active_group, _name_cnt]: lst.clear()
-  to_program_cache.clear()
  Buffer.profile_events.clear()
  cpu_events.clear()
  viz = VizTrace()
@@ -248,6 +246,19 @@ class TestViz(unittest.TestCase):
    self.assertIn("EXPAND", excluded_nodes)
    self.assertIn("CONST1 1 Ops.DEVICE", graph[id(alu)]["label"])

+  def test_stack_movement_not_folded_unless_all_const(self):
+    a = UOp.variable("a", 0, 10, dtype=dtypes.int)
+    c = UOp.const(dtypes.int, 1)
+    stack = a.vectorize(c)
+    reshaped = stack.reshape((1, 2))
+    graph = uop_to_json(VizData(), reshaped)
+    self.assertFalse(graph[id(stack)]["exclude"])
+
+    const_stack = c.vectorize(UOp.const(dtypes.int, 2))
+    const_reshaped = const_stack.reshape((1, 2))
+    const_graph = uop_to_json(VizData(), const_reshaped)
+    self.assertTrue(const_graph[id(const_stack)]["exclude"])
+
 # VIZ displays nested graph_rewrites in a tree view

 def leaf_rewrite(x:UOp): return x.rtag(1) if x.tag is None else None
@@ -329,12 +340,15 @@ class TestVizIntegration(unittest.TestCase):
  def test_codegen_tracing(self):
    with save_viz() as viz:
      ast = (Tensor.empty(4)+Tensor.empty(4)).schedule_linear().src[0].src[0]
-      prg = to_program(ast, Device[Device.DEFAULT].renderer)
+      prg = do_to_program(ast, Device[Device.DEFAULT].renderer)
    lst = viz.list_items()
    self.assertEqual(len(lst), 3)
    self.assertEqual(lst[0]["name"], "Callify 1 Buffer n1")
    self.assertEqual(lst[1]["name"], "Schedule 1 Kernel n1")
    self.assertEqual(lst[2]["name"], prg.arg.name)
+    input_ast = next(viz.get_details(2, 0))["graph"].values()
+    for u in input_ast:
+      if u["label"].startswith("PARAM\n"): self.assertEqual(u["addrspace"], addrspace_colors[AddrSpace.GLOBAL])

  # schedule graph CALL nodes have a link to jump to codegen
  def test_link_sched_codegen(self):
@@ -346,7 +360,7 @@ class TestVizIntegration(unittest.TestCase):
      from tinygrad.engine.realize import compile_linear
      sched = compile_linear(sched)
      with Context(NO_COLOR=0):
-        prgs = [to_program(si.src[0], Device[c1.device].renderer).arg.name for si in sched.src]
+        prgs = [do_to_program(si.src[0], Device[c1.device].renderer).arg.name for si in sched.src]
    lst = viz.list_items()
    sched_idx = next(i for i,l in enumerate(lst) if l["name"].startswith("Schedule"))
    viz_kernel = next(i for i,s in enumerate(lst[sched_idx]["steps"]) if s["name"] == "View Kernel Graph")
@@ -753,7 +767,7 @@ class TestCfg(unittest.TestCase):
    with save_viz() as viz:
      with Context(DEV=f"NULL::{self.arch}"):
        out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
-        _ = to_program(out.schedule_linear().src[-1].src[0], Device[out.device].renderer)
+        _ = do_to_program(out.schedule_linear().src[-1].src[0], Device[out.device].renderer)
    codegen_rewrites = next(s for s in viz.list_items() if s["name"] == name)
    disasm = next(s for s in codegen_rewrites["steps"] if s["name"] == "View Disassembly")
    return get_render(viz.data, disasm["query"])
@@ -990,8 +1004,9 @@ class TestCLI(unittest.TestCase):
  def test_dedup(self):
    with save_viz() as viz:
      for _ in range(CNT:=4):
-        Tensor.empty(4, device="NULL").add(1).realize()
-        Tensor.empty(8, device="NULL").add(1).realize()
+        # use kernel names unique to this test
+        Tensor.custom_kernel(Tensor.empty(4, device="NULL"), fxn=lambda _: UOp.sink(arg=KernelInfo("k1_test_viz_dedup")))[0].realize()
+        Tensor.custom_kernel(Tensor.empty(8, device="NULL"), fxn=lambda _: UOp.sink(arg=KernelInfo("k2_test_viz_dedup")))[0].realize()
    with write_files(viz) as files, Context(NO_COLOR=1):
      name = run_cli(*files, "-s", "NULL")[0]["name"]
      with Context(DEBUG=3):
@@ -20,7 +20,7 @@ class TestWinograd(unittest.TestCase):
    out = Tensor.conv2d(x,w, padding=1)
    out.mean().backward()
    backward_schedule = x.grad.schedule_linear(w.grad)
-    self.assertEqual(len(backward_schedule.src), 4)
+    self.assertEqual(len(backward_schedule.src), 2)

  @unittest.skip("this requires optimizations")
  def test_counters(self):
@@ -13,12 +13,12 @@ AMX = "AMX" in DEV.arch
 class TestFloat4(unittest.TestCase):
  @staticmethod
  def count_float4(uops: list[UOp], n=4):
-    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.float.vec(n)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.float.vec(n)]))
+    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype.scalar() == dtypes.float and uop.shape == (4,)]),
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype.scalar() == dtypes.float and uop.shape == (4,)]))
  @staticmethod
  def count_half4(uops: list[UOp]):
-    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.half.vec(4)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.half.vec(4)]))
+    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype.scalar() == dtypes.half and uop.shape == (4,)]),
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype.scalar() == dtypes.half and uop.shape == (4,)]))

  def test_float4_basic(self):
    a = Tensor.empty(2, 8).realize()
@@ -281,15 +281,15 @@ class TestAssign(unittest.TestCase):
    np.testing.assert_equal(t.numpy(), [[100, 104, 108, 112], [101, 105, 109, 113], [102, 106, 110, 114], [103, 107, 111, 115]])

  def test_assign_contiguous(self):
-    b = Tensor.arange(16).reshape(4,4).contiguous().realize()
-    a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1)
+    b = Tensor.arange(16).reshape(4,4).clone().realize()
+    a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1)
    GlobalCounters.reset()
    b.assign(a.contiguous()).realize()
    self.assertEqual(GlobalCounters.kernel_count, 2)

  def test_assign_contiguous_permute(self):
-    b = Tensor.arange(16).reshape(4,4).contiguous().realize()
-    a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1).permute((1,0))
+    b = Tensor.arange(16).reshape(4,4).clone().realize()
+    a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1).permute((1,0))
    GlobalCounters.reset()
    b.assign(a.contiguous()).realize()
    self.assertEqual(GlobalCounters.kernel_count, 2)
@@ -325,29 +325,29 @@ class TestAssign(unittest.TestCase):
    np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))

  def test_post_permuted_assignment_alt(self):
-    a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
-    b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
+    a = Tensor.arange(N*N).reshape(N,N).clone().realize()
+    b = Tensor.arange(N*N).reshape(N,N).clone().realize()
    new_a = (a.T+b).numpy()
    a.assign(a.T+b)
    np.testing.assert_allclose(a.numpy(), new_a)

  def test_post_flipped_assignment(self):
-    a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
-    b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
+    a = Tensor.arange(N*N).reshape(N,N).clone().realize()
+    b = Tensor.arange(N*N).reshape(N,N).clone().realize()
    new_a = (a.flip(0)+b).numpy()
    a.assign(a.flip(0)+b)
    np.testing.assert_allclose(a.numpy(), new_a)

  def test_post_flipped_assignment_axis1(self):
-    a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
-    b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
+    a = Tensor.arange(N*N).reshape(N,N).clone().realize()
+    b = Tensor.arange(N*N).reshape(N,N).clone().realize()
    new_a = (a.flip(1)+b).numpy()
    a.assign(a.flip(1)+b)
    np.testing.assert_allclose(a.numpy(), new_a)

  def test_post_reshape_assignment_fine(self):
-    a = Tensor.arange(N*N).reshape(N, N).contiguous().realize()
-    b = Tensor.arange(N*N).reshape(N, N).contiguous().realize()
+    a = Tensor.arange(N*N).reshape(N, N).clone().realize()
+    b = Tensor.arange(N*N).reshape(N, N).clone().realize()
    rhs = a.reshape(-1).reshape(N, N)
    new_a = (rhs+b).numpy()
    a.assign(rhs+b)  # self-assign with reshape view is fine
@@ -355,7 +355,7 @@ class TestAssign(unittest.TestCase):

  @unittest.skip("multi output not supported anymore")
  def test_simple_assignment_multioutput(self):
-    a = Tensor.arange(32*32).reshape(32, 32).contiguous().realize()
+    a = Tensor.arange(32*32).reshape(32, 32).clone().realize()
    b = Tensor.full((32, ), 1.).contiguous().realize()
    c = Tensor.full((32, ), 2.).contiguous().realize()
    d = Tensor.full((32, ), 3.).contiguous().realize()
@@ -375,15 +375,15 @@ class TestAssign(unittest.TestCase):
  # NOTE: if the assign target is read/write in a single kernel, it should be contiguous

  def test_permuted_assignment_correct(self):
-    a = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
-    b = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
+    a = Tensor.arange(4 * 4).reshape(4, 4).clone().realize()
+    b = Tensor.arange(4 * 4).reshape(4, 4).clone().realize()
    a = a.permute(1, 0)
    new_val = a + b
    a.assign(new_val)
    np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4))

  def test_permuted_reduceop_child_dual_use(self):
-    a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
+    a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize()
    b = Tensor.ones(32, 32, dtype=dtypes.int).contiguous().realize()
    r = a.sum(axis=1)
    b.assign(r + b.permute(1, 0))
@@ -392,7 +392,7 @@ class TestAssign(unittest.TestCase):

  @unittest.skip("multi output not supported anymore")
  def test_permuted_reduceop_multioutput_dual_use(self):
-    a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
+    a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize()
    b = Tensor.full((32, 32), 1.).contiguous().realize()
    c = Tensor.full((32, 32), 2.).contiguous().realize()

@@ -405,9 +405,9 @@ class TestAssign(unittest.TestCase):

  @unittest.skip("multi output not supported anymore")
  def test_permuted_reduceop_multioutput_dual_use_possible(self):
-    a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
-    b = Tensor.arange(32 * 32).reshape(32, 32).realize()
-    c = Tensor.arange(32 * 32).reshape(32, 32).realize()
+    a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize()
+    b = Tensor.arange(32 * 32).reshape(32, 32).clone().realize()
+    c = Tensor.arange(32 * 32).reshape(32, 32).clone().realize()

    GlobalCounters.reset()
    r = a.sum(axis=1)
@@ -441,7 +441,7 @@ class TestAssign(unittest.TestCase):
    # Forward shift: read index > write index in overlap
    N = 100000
    shift = 1000
-    a = Tensor.arange(N).float().contiguous().realize()
+    a = Tensor.arange(N).float().clone().realize()
    expected = np.arange(N, dtype=np.float32)
    expected[:N-shift] = expected[shift:].copy()
    with Context(NOOPT=1): a[0:N-shift].assign(a[shift:N]).realize()
@@ -451,7 +451,7 @@ class TestAssign(unittest.TestCase):
    # Reverse shift: write index > read index in overlap
    N = 100000
    shift = 1000
-    a = Tensor.arange(N).float().contiguous().realize()
+    a = Tensor.arange(N).float().clone().realize()
    expected = np.arange(N, dtype=np.float32)
    expected[shift:] = expected[:N-shift].copy()
    with Context(NOOPT=1): a[shift:N].assign(a[0:N-shift]).realize()
@@ -459,7 +459,7 @@ class TestAssign(unittest.TestCase):

  def test_nonoverlapping_shrink_assignment(self):
    # TODO: non-overlapping shrinks don't actually need contiguous, could be 1 kernel with smarter range analysis
-    a = Tensor.arange(100).float().contiguous().realize()
+    a = Tensor.arange(100).float().clone().realize()
    expected = np.arange(100, dtype=np.float32)
    expected[0:10] = expected[50:60].copy()
    GlobalCounters.reset()
@@ -222,7 +222,7 @@ class TestCallSchedule(unittest.TestCase):
    # find the FUNCTION nodes
    c0 = next(u for u in r0.uop.toposort() if u.op is Ops.FUNCTION)
    c1 = next(u for u in r1.uop.toposort() if u.op is Ops.FUNCTION)
-    # the function bodies (src[0]) should have identical keys — local buffer identity must not leak through
+    # the function bodies (src[0]) should have identical keys — unique consts must not leak through
    self.assertEqual(c0.src[0].key, c1.src[0].key)

  def test_precompile_symbolic_2d(self):
@@ -495,6 +495,25 @@ class TestFunctionTuple(unittest.TestCase):
    Tensor.realize(a.grad)
    np.testing.assert_allclose(a.grad.numpy(), [2., 2., 2., 2.])

+  def test_custom_kernel_precompile_multidevice(self):
+    # a custom_kernel output placeholder (invalids) under multi-device @function(precompile=True) must return the
+    # kernel's computed result. read it back through .numpy() so the cross-device gather reads the output buffer
+    devs = ("CPU:0", "CPU:1")
+    def double_kernel(C:UOp, A:UOp) -> UOp:
+      C, A = C.flatten(), A.flatten()
+      i = UOp.range(A.numel(), 0)
+      return C[i].store(A[i] * 2.0).end(i).sink(arg=KernelInfo(name="double_kernel"))
+    def double_grad(d_c:UOp, call:UOp): return (None, (Tensor(d_c) * 2.0).uop)
+
+    @function(precompile=True, precompile_backward=True)
+    def f(a:Tensor):
+      c = Tensor(Tensor.invalids(a.shape[0]//len(devs), a.shape[1], dtype=a.dtype, device=devs).uop.multi(0), device=devs)
+      return Tensor.custom_kernel(c, a, fxn=double_kernel, grad_fxn=double_grad)[0]
+
+    a = Tensor.full((4, 4), 7.0).contiguous().shard(devs, axis=0)
+    Tensor.realize(a)
+    np.testing.assert_allclose(f(a).numpy(), 14.0)
+
  def test_custom_kernel_precompile_further_compute(self):
    def my_kernel(C:UOp, A:UOp) -> UOp:
      i = UOp.range(A.shape[0], 0)
@@ -65,6 +65,10 @@ class TestTensorData(unittest.TestCase):
    assert dat.tolist() == 3
    assert dat.shape == ()

+  def test_const_dtype_for_uop(self):
+    self.assertEqual(Tensor.const(dtypes.int8, UOp.const(dtypes.float32, 1.0)).dtype, dtypes.int8)
+    self.assertEqual(Tensor.const(dtypes.int32, UOp.variable("x", 1, 10).bind(5)).item(), 5)
+
  def test_data_float32(self):
    a = Tensor([[1,2.5],[3,4]], dtype=dtypes.float32)
    dat = a.data()
@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from tinygrad.dtype import dtypes
+from tinygrad.dtype import dtypes, AddrSpace, PtrDType, ImageDType
 from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, graph_rewrite, track_rewrites
 from tinygrad.helpers import VIZ, pluralize, all_int

@@ -176,11 +176,14 @@ def finalize_after(ctx:AllocCtx, x:UOp):
 def replace_input_buffer(ctx:AllocCtx, b:UOp):
  ctx.replacements.append(b)
  return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
-                   b._min_max if b.op is Ops.BIND else None, b.src[0].arg[0] if b.op is Ops.BIND else None)
+                   b._min_max if b.op is Ops.BIND else None, b.src[0].arg[0] if b.op is Ops.BIND else None,
+                   b.addrspace if isinstance(b.dtype, (PtrDType, ImageDType)) else AddrSpace.GLOBAL)

 pm_finalize_call = PatternMatcher([
  (UPat(Ops.AFTER, name="x"), finalize_after),
  (UPat(Ops.COPY, name="x"), lambda ctx,x: ctx.assigns.append(x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None),
+  # remove unique from const. TODO: this is copied in function.py
+  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE, name="d")), name="b"), lambda b,d: b.replace(src=(d,))),
 ])

 pm_replace_buf = PatternMatcher([
@@ -115,7 +115,7 @@ pm_linearize_cleanups = PatternMatcher([
  # if statements are not allowed in the graph
  (UPat((Ops.IF, Ops.ENDIF)), lambda: panic(RuntimeError, "if not allowed in graph")),
  # gated STORE becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF
-  (UPat(Ops.STORE, name="u", src=(UPat(), UPat(), UPat(name="gate", dtype=dtypes.bool))),
+  (UPat(Ops.STORE, name="u", src=(UPat((Ops.INDEX, Ops.SHRINK)).or_casted(), UPat(), UPat(name="gate", dtype=dtypes.bool))),
   lambda u, gate: ((st:=u.replace(src=u.src[0:2])), [mif:=UOp(Ops.IF, src=(gate, u.src[0])), st, UOp(Ops.ENDIF, src=(mif,))]))
 ])

@@ -91,7 +91,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
  subs = {}
  for r in s_topo:
    # look for local INDEXes that are not used in the GLOBAL store, then add them as an INVALID
-    if r.op is Ops.STORE and (idx := r.src[0]).src[0].ptrdtype.addrspace == AddrSpace.GLOBAL:
+    if r.op is Ops.STORE and (idx := r.src[0]).src[0].addrspace == AddrSpace.GLOBAL:
      missing_locals = [all_ranges[rng] for rng in local_dims if all_ranges[rng] not in idx.ranges]
      if len(missing_locals):
        assert len(idx.src) == 2, "index has 2 sources"
@@ -104,7 +104,7 @@ def fold_expanded_index(midx:UOp):
    for grp in grouped_offsets:
      # get the index offset for this element. using [0] is okay, because they are the same
      lidx = midx.src[offsets[grp[0]][0]]
-      if len(grp) > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(len(grp)).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace))
+      if len(grp) > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(len(grp)).ptr(size=buf.max_numel(), addrspace=buf.addrspace))
      # set the idxs of the output
      for i,g in enumerate(grp):
        for oo in offsets[g]: idxs[oo] = global_offset+i
@@ -113,7 +113,7 @@ def fold_expanded_index(midx:UOp):
      global_offset += len(grp)
  assert None not in idxs, f"some idxs are missing {idxs}"
  # this base thing is for image, we want the CAT to be a normal pointer
-  post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace).vec(global_offset), tuple(ret))
+  post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.max_numel(), addrspace=buf.addrspace).vec(global_offset), tuple(ret))
  return post_cat.gep(tuple(cast(list[int], idxs)))

 def cat_after_store(cat:UOp, data:UOp):
@@ -165,7 +165,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
    must_divide = False
  elif buf.dtype.base not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not isinstance(buf.dtype, ImageDType):
    pass
-  elif buf.ptrdtype.addrspace == AddrSpace.REG:
+  elif buf.addrspace == AddrSpace.REG:
    pass
  elif isinstance(buf.dtype, ImageDType):
    lengths = [4]
@@ -186,7 +186,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
    for fold_length in lengths:
      if global_offset+fold_length > sz: continue
      lidx = buf.index((offset + global_offset).valid(mask), ptr=True)
-      if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace))
+      if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.max_numel(), addrspace=buf.addrspace))
      if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length))))))
      else: ret.append(ls.replace(src=(lidx,)+ls.src[1:], dtype=ls.dtype.scalar().vec(fold_length)))
      global_offset += fold_length
@@ -243,7 +243,9 @@ def no_vectorized_alu(alu:UOp):
  return UOp(Ops.STACK, alu.dtype, alus)

 def no_vectorized_buf(buf:UOp):
-  return buf.replace(dtype=buf.ptrdtype.base.scalar().ptr(buf.ptrdtype.size*buf.ptrdtype.count, buf.ptrdtype.addrspace)).cast(buf.dtype)
+  # TODO: this fails on regs
+  #assert buf.max_numel() == buf.ptrdtype.size
+  return buf.replace(dtype=buf.ptrdtype.base.scalar().ptr(buf.ptrdtype.size*buf.ptrdtype.count, buf.addrspace)).cast(buf.dtype)

 def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp, bcast:UOp|None=None):
  cnt = cast.dtype.count
@@ -283,12 +285,6 @@ pm_render = PatternMatcher([
  (UPat(Ops.GEP, name='gep'), lambda gep: UOp(Ops.STACK, gep.dtype, tuple(gep.src[0].gep(x) for x in gep.arg)) if len(gep.arg) > 1 else None),
  (UPat(Ops.GEP, name='gep'), lambda gep: gep.src[0] if gep.src[0].dtype.vcount == 1 and gep.arg == (0,) else None),
  (UPat(Ops.STACK, src=(UPat(name='x'),)), lambda x: x),
-  # rewrite non-image INDEX to SLICE
-  (UPat(Ops.INDEX, name="x"), lambda x: None if isinstance(x.src[0].dtype, ImageDType) else \
-    UOp(Ops.SLICE, dtype=x.dtype, src=x.src, arg=0 if x.dtype.count == 1 else x.dtype.count)),
-  # rewrite CAST on SLICE to just SLICE
-  (UPat(Ops.SLICE, name="bv").cast(name="x"),
-   lambda bv,x: bv.replace(dtype=x.dtype, arg=0 if x.dtype.count == 1 else x.dtype.count))
 ])

 # *** Ops.REDUCE -> Ops.DEFINE_ACC ***
@@ -22,7 +22,7 @@ def linearize(sink:UOp) -> list[UOp]:
    extra = None
    match u.op:
      # the order and placement of these defines is important
-      case Ops.PARAM: priority, extra = -20, u.arg
+      case Ops.PARAM: priority, extra = -20, u.arg.slot
      case Ops.DEFINE_VAR: priority, extra = -19, u.arg
      case Ops.DEFINE_REG: priority = -18
      case Ops.DEFINE_LOCAL: priority = -17
@@ -93,4 +93,4 @@ def do_split_ends(e:UOp):
 pm_split_ends = PatternMatcher([
  # split the ends
  (UPat(Ops.END, name="e"), do_split_ends),
-])
+])
@@ -96,10 +96,11 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:

  # if there are small dims with lots of valid masks, upcast them (they might be from Tensor.stack)
  to_upcast: list[int] = []
+  where_gate_rngs = {r for u in k.ast.backward_slice if u.op is Ops.WHERE for r in u.src[0].ranges}
  # upcast leading axes first (hack-ish for winograd; we actually want to upcast masked axes with low stride first)
  for axis in k.upcastable_dims:
    # for Schedule, we check if the range is used in INDEX gates or WHERE gates
-    is_masked = any(any(o is k.rngs[axis] for o in u.src[0].backward_slice) for u in k.ast.backward_slice if u.op is Ops.WHERE)
+    is_masked = k.rngs[axis] in where_gate_rngs
    if k.full_shape[axis] <= 7 and is_masked and prod(k.full_shape[j] for j in to_upcast) * k.full_shape[axis] <= 7 * 7:
      if DEBUG >= 4: print(f"upcasting masked axis : {axis}")
      to_upcast.append(axis)
@@ -95,7 +95,7 @@ class Scheduler:
    if (old_sz:=rng.src[0].divides(amount)) is None:
      raise KernelOptError(f"{amount} can't divide {rng.src[0]} in {self.colored_shape()}")
    new_rng = UOp.range(amount, next(self.opt_range), new_type) if input_new_rng is None else input_new_rng
-    replaced_rng = rng.replace(src=(UOp.const(dtypes.int, old_sz),))
+    replaced_rng = rng.replace(src=(old_sz,))
    sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
    self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[:-1]} {amount} {str(new_type).split('.')[1].lower()}")
    return replaced_rng, new_rng
@@ -329,8 +329,8 @@ class Scheduler:
  def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE))

 def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
-  glbls = sorted([x for x in ast.backward_slice if x.op is Ops.PARAM], key=lambda x: x.arg)
-  return [Buffer(dname, x.ptrdtype.size, x.dtype.base) for x in glbls]
+  glbls = sorted([x for x in ast.backward_slice if x.op is Ops.PARAM], key=lambda x: x.arg.slot)
+  return [Buffer(dname, x.max_numel(), x.dtype.base) for x in glbls]

 def apply_opts(ast:UOp, ren:Renderer, beam:int=0) -> UOp:
  if ast.tag is not None: return ast
@@ -51,12 +51,12 @@ class DTypeMetaClass(type):

 class AddrSpace(Enum):
  def __repr__(self): return str(self)
-  GLOBAL = auto(); LOCAL = auto(); REG = auto()  # noqa: E702
+  GLOBAL = auto(); LOCAL = auto(); REG = auto(); ANON = auto()  # noqa: E702

@dataclass(frozen=True, eq=False)
 class DType(metaclass=DTypeMetaClass):
  priority: int  # this determines when things get upcasted
-  bitsize: int
+  bitsize: int   # this is the bitsize of the base dtype
  name: str
  fmt: FmtStr|None
  count: int
@@ -76,7 +76,7 @@ class DType(metaclass=DTypeMetaClass):
  def vec(self, sz:int) -> DType:
    assert self.count == 1, f"can't vectorize {self} with size {sz}"
    if sz == 1 or self == dtypes.void: return self  # void doesn't vectorize, and sz=1 is scalar
-    return DType(self.priority, self.bitsize*sz, f"{INVERSE_DTYPES_DICT[self.name]}{sz}", None, sz, self)
+    return DType(self.priority, self.bitsize, f"{INVERSE_DTYPES_DICT[self.name]}{sz}", None, sz, self)
  def ptr(self, size=-1, addrspace=AddrSpace.GLOBAL) -> PtrDType:
    return PtrDType(self.priority, self.bitsize, self.name, self.fmt, self.count, None, self, addrspace, 1, size)
  def scalar(self) -> DType: return self._scalar if self._scalar is not None else self
@@ -87,6 +87,34 @@ def _check_no_non_tensor_return(ret):

 def graph_class(dev): return dev.graph.func if isinstance(dev.graph, functools.partial) else dev.graph

+class DepsTracker:
+  def __init__(self):
+    # tracks (offset, end, dep) ranges per base buffer id to handle suballocated buffers correctly.
+    self.w_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
+    self.r_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
+
+  @staticmethod
+  def _buf_key(buf:Buffer) -> int: return id(buf.base)
+
+  def access_resources(self, bufs:list[Buffer], write:list[int], new_dependency:Any):
+    wait_nodes = []
+    for i,buf in enumerate(bufs):
+      key, s, e = self._buf_key(buf), buf.offset, buf.offset + buf.nbytes
+      wait_nodes += [dep for st,en,dep in self.w_dependency_map[key] if st < e and s < en]
+      if i in write: wait_nodes += [dep for st,en,dep in self.r_dependency_map[key] if st < e and s < en]
+    for i,buf in enumerate(bufs):
+      key, s, e = self._buf_key(buf), buf.offset, buf.offset + buf.nbytes
+      if i in write:
+        for dmap in [self.w_dependency_map, self.r_dependency_map]:
+          kept = []
+          for st,en,dep in dmap[key]:
+            if st < min(s, en): kept.append((st, min(s, en), dep))
+            if max(e, st) < en: kept.append((max(e, st), en, dep))
+          dmap[key] = kept
+        self.w_dependency_map[key].append((s, e, new_dependency))
+      else: self.r_dependency_map[key].append((s, e, new_dependency))
+    return list({id(x):x for x in wait_nodes}.values())
+
 class GraphRunner:
  def __init__(self, linear:UOp, input_uops:tuple[UOp, ...]=()):
    self.linear = linear.src[0]
@@ -94,7 +122,7 @@ class GraphRunner:
    self.runtimes: list[Any|None] = []
    self.uop_replace: list[list[tuple[int, int]]] = []
    for call in self.linear.src:
-      replace = [(p, b.arg) for p, b in enumerate(get_call_arg_uops(call)) if b.op is Ops.PARAM]
+      replace = [(p, b.arg.slot) for p, b in enumerate(get_call_arg_uops(call)) if b.op is Ops.PARAM]
      for dev_idx, (bufs, device_vars) in enumerate(unwrap_multi(call, resolve_params(call, input_uops))):
        self.calls.append((dev_idx, call.src[0], [b.ensure_allocated() for b in bufs], device_vars))
        self.runtimes.append(get_runtime(bufs[0].device, call.src[0]) if call.src[0].op is Ops.PROGRAM else None)
@@ -123,9 +151,8 @@ class GraphRunner:

    estimates = sum((estimate_uop(call) for call in self.linear.src), Estimates())

-    # used in MultiGraphRunner. tracks (offset, end, dep) ranges per base buffer id to handle suballocated buffers correctly.
-    self.w_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
-    self.r_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
+    # used in MultiGraphRunner
+    self.deps = DepsTracker()

    self.device, self.estimates = self.calls[0][2][0].device.split(":")[0], estimates.simplify()

@@ -142,23 +169,7 @@ class GraphRunner:
      yield j, (dims[gl] if gl is not None else self.launch_dims_base[j][0]), (dims[lc] if lc is not None else self.launch_dims_base[j][1])

  def _access_resources(self, bufs:list[Buffer], write:list[int], new_dependency:Any):
-    wait_nodes = []
-    for i,buf in enumerate(bufs):
-      key, s, e = id(buf.base._buf), buf.offset, buf.offset + buf.nbytes
-      wait_nodes += [dep for st,en,dep in self.w_dependency_map[key] if st < e and s < en]
-      if i in write: wait_nodes += [dep for st,en,dep in self.r_dependency_map[key] if st < e and s < en]
-    for i,buf in enumerate(bufs):
-      key, s, e = id(buf.base._buf), buf.offset, buf.offset + buf.nbytes
-      if i in write:
-        for dmap in [self.w_dependency_map, self.r_dependency_map]:
-          kept = []
-          for st,en,dep in dmap[key]:
-            if st < min(s, en): kept.append((st, min(s, en), dep))
-            if max(e, st) < en: kept.append((max(e, st), en, dep))
-          dmap[key] = kept
-        self.w_dependency_map[key].append((s, e, new_dependency))
-      else: self.r_dependency_map[key].append((s, e, new_dependency))
-    return list({id(x):x for x in wait_nodes}.values())
+    return self.deps.access_resources(bufs, write, new_dependency)

  @staticmethod
  def _all_devs(batch_devs:list[Compiled], new_call:UOp) -> list[Compiled]:
@@ -137,8 +137,8 @@ class ExecContext:
  cache: bool = True

 def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
-  if b.op in (Ops.SLICE, Ops.MSELECT) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg], *b.src[1:]))
-  return inputs[b.arg] if b.op is Ops.PARAM else b
+  if b.op in (Ops.SLICE, Ops.MSELECT) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
+  return inputs[b.arg.slot] if b.op is Ops.PARAM else b
 def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)]

 def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], dict[str, int]]]:
@@ -241,14 +241,13 @@ pm_exec = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate),
 ])

-if getenv("HCQ2"):
-  from extra.hcq2.hcq2 import pm_hcq_exec
-  pm_exec = pm_hcq_exec + pm_exec
-
 def compile_linear(linear:UOp, beam:int|None=None, validate=False) -> UOp:
  if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
  if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
  linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
+  if getenv("HCQ2"):
+    from extra.hcq2.hcq2 import hcq_schedule
+    linear = hcq_schedule(linear)
  return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)

 def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:tuple[UOp, ...]=(), update_stats=True, jit=False, wait=False):
@@ -10,13 +10,17 @@ def add_to_ctx(ctx, x:UOp):
  ctx[0].append(x)
  return ret

+pm_transform_unique_const = PatternMatcher([
+  # transform unique consts to LUNIQUE
+  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="x"),
+   lambda ctx,x: x.replace(src=(UOp(Ops.LUNIQUE, arg=next(ctx[1])), x.src[1]))),
+])
+
 pm_ctx = PatternMatcher([
-  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="x"),
-   lambda ctx,x: x.replace(src=(UOp(Ops.LUNIQUE, arg=next(ctx[1])), x.src[1])) if x.src[0].arg > ctx[2] else add_to_ctx(ctx,x)),
-  (UPat(Ops.BIND, name="x"), add_to_ctx),
+  (UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
  (UPat((Ops.AFTER, Ops.CONTIGUOUS), name="x"),
   lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) and x.op_in_backward_slice_with_self(Ops.BUFFER) else None),
-])
+])+pm_transform_unique_const

 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):
@@ -42,7 +46,6 @@ class _function(Generic[ReturnType]):
    # run it and do surgery later
    with Context(ALLOW_DEVICE_USAGE=getenv("DEVICE_IN_FUNCTION_BUG", 0)):
      _function.depth += 1
-      unique_start = next(UOp.unique_num)
      ret = self.fxn(*args, **kwargs)
      _function.depth -= 1
    if isinstance(ret, Tensor):
@@ -62,7 +65,7 @@ class _function(Generic[ReturnType]):

    # the BUFFERs that are left are the implicit inputs
    num_explicit = len(call_uops)
-    uret = graph_rewrite(uret, pm_ctx, (call_uops, itertools.count(0), unique_start), bottom_up=True, name="get_implicit_inputs")
+    uret = graph_rewrite(uret, pm_ctx, (call_uops, itertools.count(0)), bottom_up=True, name="get_implicit_inputs")
    name = getattr(self.fxn, '__qualname__', None) or type(self.fxn).__qualname__
    if not self.allow_implicit:
      implicit_buffers = [x for x in call_uops[num_explicit:] if x.op is Ops.BUFFER]
@@ -1,6 +1,6 @@
 from typing import cast
-import math, dataclasses
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata
+import math, dataclasses, itertools
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, graph_rewrite
 from tinygrad.helpers import argsort
 from tinygrad.dtype import sum_acc_dtype

@@ -16,8 +16,9 @@ def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):

 def _compact_params(body:UOp, all_args:tuple[UOp, ...]) -> tuple[UOp, tuple[UOp, ...]]:
  """Remove unused PARAMs from body and return compacted (body, args)."""
-  used = sorted({p.arg: p for p in body.toposort() if p.op is Ops.PARAM}.items())
-  return body.substitute({p: p.replace(arg=j) for j,(_, p) in enumerate(used)}, walk=True), tuple(all_args[i] for i,_ in used)
+  used = sorted({p.arg.slot: p for p in body.toposort() if p.op is Ops.PARAM}.items())
+  body = body.substitute({p: p.replace(arg=dataclasses.replace(p.arg, slot=j)) for j,(_, p) in enumerate(used)}, walk=True)
+  return body, tuple(all_args[i] for i,_ in used)

 def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
  fxn, args = k.src[0], k.src[1:]
@@ -29,7 +30,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
      return (None,) + (k.arg.grad_fxn(*real, call=k) if len(real) > 1 else k.arg.grad_fxn(real[0], k))
    return (None,) + k.arg.grad_fxn(on_dev(ctx, 0), k)
  assert fxn.op is Ops.TUPLE, f"expected TUPLE body for gradient, got {fxn.op}"
-  params = {x.arg:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
+  params = {x.arg.slot:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
  grad_args = ctx.src
  root_grad = UOp(Ops.TUPLE, src=tuple(UOp(Ops.NOOP) if g.op is Ops.NOOP else
    g if g.base.op is Ops.CONST and g.device is None else g.param_like(len(args)+i) for i,g in enumerate(grad_args)))
@@ -41,6 +42,9 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
  grad_bodies = [(i, grads[p]) for i in needed if (p:=params.get(i)) is not None and p in grads]
  bwd_body = UOp.maketuple(*(gb for _, gb in grad_bodies)).substitute(fwd_subs, walk=True)
  bwd_body, compact_args = _compact_params(bwd_body, (*args, *grad_args, *fwd_outs))
+  # TODO: is this okay here?
+  from tinygrad.function import pm_transform_unique_const
+  bwd_body = graph_rewrite(bwd_body, pm_transform_unique_const, ctx=(None, itertools.count(0)))
  bwd_call = bwd_body.call(*compact_args, name=(k.arg.name or "")+"_backward", precompile=k.arg.precompile_backward)
  gb_map = {i: idx for idx, (i, _) in enumerate(grad_bodies)}
  return (None,) + tuple(bwd_call.gettuple(gb_map[i]) if i in gb_map else None for i in range(len(args)))
@@ -70,7 +74,7 @@ pm_gradient = PatternMatcher([
    (ctx.cast(sum_acc_dtype(ctx.dtype))._rop(Ops.ADD, tuple(i for i,(s,n) in enumerate(zip(ret.src[0].shape, ret.shape)) if s!=n))
     .cast(ctx.dtype), None)),
  (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[0], s+p[0]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
-  (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
+  (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[0]-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)),
  (UPat(Ops.PERMUTE, name="ret"), lambda ctx, ret: (ctx.permute(argsort(ret.marg)),)),
  (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip([i for i,x in enumerate(ret.marg) if x]),)),
  (UPat(Ops.COPY, name="ret"), lambda ctx, ret: (ctx.copy_to_device(ret.src[0].device), None)),
@@ -98,8 +98,8 @@ def get_child(obj, key):
 def word_wrap(x, wrap=80):
  if len(ansistrip(x)) <= wrap: return x
  if len(lines:=x.splitlines()) > 1: return "\n".join(word_wrap(line, wrap) for line in lines)
-  i = 0
-  while len(ansistrip(x[:i])) < wrap and i < len(x): i += 1
+  i = vis = 0
+  while vis < wrap and i < len(x): i, vis = (i + m.end(), vis) if (m:=re.match('\x1b\\[(K|.*?m)', x[i:])) is not None else (i+1, vis+1)
  return x[:i] + "\n" + word_wrap(x[i:], wrap)
 def pad_bytes(b:bytes, align:int) -> bytes: return b + b'\x00' * ((align - (len(b) % align)) % align)

@@ -7,7 +7,7 @@ from tinygrad.mixin.reduce import ReduceMixin
 from tinygrad.uop import Ops
 from tinygrad.uop.ops import _broadcast_shape, resolve, smax, smin, identity_element
 from tinygrad.device import canonicalize_device
-from tinygrad.dtype import ConstType, DType, DTypeLike, InvalidType, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype
+from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, InvalidType, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype
 from tinygrad.helpers import all_int, argfix, ceildiv, flatten, flat_to_grouped, make_tuple, prod, resolve_pool_pads, round_up

 if TYPE_CHECKING:
@@ -18,7 +18,7 @@ ReductionStr = Literal["mean", "sum", "none"]

 class OpMixin(ElementwiseMixin, ReduceMixin):
  @staticmethod
-  def empty(*shape, **kwargs): raise NotImplementedError("creation helpers are only supported on Tensor and UOp")
+  def unique_const(fill_value:ConstType, **kwargs): raise NotImplementedError("creation helpers are only supported on Tensor and UOp")
  @staticmethod
  def const(dtype, b, device=None): raise NotImplementedError("creation helpers are only supported on Tensor and UOp")

@@ -38,25 +38,25 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
    print(Tensor.full((2, 3), False).numpy())
    ```
    """
+    from tinygrad.uop.ops import UOp
    new_shape = argfix(shape)
    dt = to_dtype(dtype) if dtype is not None else None
-    # build the broadcast const value (deviceless for a buffer, device-placed for a value), then clone into storage iff buffer
-    if isinstance(fill_value, get_args(ConstType)):
-      val = cls.const(dt or dtypes.from_py(fill_value), fill_value, None if buffer else canonicalize_device(device))
-    else:  # symbolic UOp fill: keep the value's own dtype, cast only when one is requested
-      val = cls.const(dt, fill_value)
-      if dt is not None: val = val.cast(dt)
+    if isinstance(fill_value, UOp): val = cls.const(dt or fill_value.dtype, fill_value)
+    else: val = cls.const(dt or dtypes.from_py(fill_value), fill_value, None if buffer else canonicalize_device(device))
    val = val.reshape((1,)*len(new_shape)).expand(new_shape)
    return val.clone(device=device) if buffer else val

  @classmethod
  def invalids(cls, *shape, **kwargs) -> Self:
    """
-    Creates an anonymous uninitialized buffer with the given shape.
+    Creates a tensor with the given shape, filled with Invalid.

-    You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
+    This is an alternative to Tensor.empty when you want an "anonymous" buffer.
+
+    Eventually Tensor.empty will be replaced by this.
    """
-    return cls.empty(argfix(*shape), **kwargs)
+    new_shape = argfix(*shape)
+    return cls.unique_const(Invalid, **kwargs).reshape((1,)*len(new_shape)).expand(new_shape)

  @classmethod
  def zeros(cls, *shape, **kwargs) -> Self:
@@ -257,9 +257,11 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
    return MovementMixin.pad(X.const_like(1).cast(dtypes.bool), pads).where(base, base.const_like(value))

  def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self:
-    if any(pB>sh or pA>sh for (pB,pA),sh in zip(pX, self.shape)): raise ValueError('Padding value causes wrapping around more than once.')
-    if any(pB<0 or pA<0 for pB,pA in pX): raise NotImplementedError("Negative pads with circular pads is not supported")
-    orig_shape, X = self.shape, self.repeat(tuple(1 + bool(pB) + bool(pA) for pB,pA in pX))
+    # shrink first for negative pads, then wrap the non-negative remainder
+    X = self.shrink(tuple((-smin(pB,0), smin(pA+sh,sh)) for (pB,pA),sh in zip(pX, self.shape)))
+    pX = tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX)
+    if any(pB>sh or pA>sh for (pB,pA),sh in zip(pX, X.shape)): raise ValueError('Padding value causes wrapping around more than once.')
+    orig_shape, X = X.shape, X.repeat(tuple(1 + bool(pB) + bool(pA) for pB,pA in pX))
    return X.shrink(tuple((0 if pB == 0 else osh-pB, xsh if pA == 0 else xsh-osh+pA) for (pB,pA),osh,xsh in zip(pX, orig_shape, X.shape)))

  def _pad_reflect_replicate(self, pX:tuple[tuple[sint, sint], ...], mode:str) -> Self:
@@ -178,7 +178,7 @@ class MovementMixin:
  def pad(self, arg:tuple[tuple[sint, sint] | None, ...]) -> Self:
    if self.ndim != len(arg):
      raise ValueError(f"{self.ndim=} != {len(arg)=}")
-    ret = self._mop(Ops.PAD, tuple(x if x is not None else (0, 0) for x in arg))
+    ret = self._mop(Ops.PAD, tuple((x[0], s+x[0]+x[1]) if x is not None else (0, s) for x, s in zip(arg, self.shape)))
    return self if ret.shape == self.shape else ret

  def shrink(self, arg: tuple[tuple[sint, sint] | None, ...]) -> Self:
@@ -200,7 +200,7 @@ class MovementMixin:
    """
    if self.ndim != len(arg):
      raise ValueError(f"{self.ndim=} != {len(arg)=}")
-    ret = self._mop(Ops.SHRINK, arg=[x if x is not None else (0, s) for x, s in zip(arg, self.shape)])
+    ret = self._mop(Ops.SHRINK, arg=[(x[0], x[1]-x[0]) if x is not None else (0, s) for x, s in zip(arg, self.shape)])
    return self if ret.shape == self.shape else ret

  def permute(self, order, *args) -> Self:
@@ -251,7 +251,7 @@ class MovementMixin:
    return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)]))

  def pad_to(self, shape, *args) -> Self:
-    return self._mop(Ops.PAD, tuple([(0, 0 if ns is None else ns-s) for s,ns in zip(self.shape, argfix(shape, *args), strict=True)]))
+    return self._mop(Ops.PAD, tuple((0, s if ns is None else ns) for s,ns in zip(self.shape, argfix(shape, *args), strict=True)))

  def view(self, shape, *args) -> Self:
    """`.view` is an alias for `.reshape`."""
@@ -1000,7 +1000,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
      if align_corners: return Tensor.linspace(-1, 1, steps, device=theta.device)
      return Tensor.linspace(-1+1/steps, 1-1/steps, steps, device=theta.device)
    grids = Tensor.meshgrid(*(generate_grid(d) for d in spatial_dims))
-    base_grid = Tensor.stack(*reversed(grids), Tensor.ones_like(grids[0], device=theta.device), dim=-1)
+    base_grid = Tensor.stack(*reversed(grids), grids[0].const_like(1), dim=-1)
    base_grid = base_grid.reshape(1, prod(spatial_dims), len(grids)+1).expand(N, -1, -1)
    return (base_grid @ theta.transpose(1, 2)).reshape(N, *spatial_dims, -1)

@@ -3,7 +3,7 @@ from typing import Callable, cast
 from dataclasses import dataclass
 from tinygrad.helpers import prod, Target, EMULATED_DTYPES
 from tinygrad.uop.ops import Ops, UOp, sint, ssimplify, smin, GroupOp, PatternMatcher
-from tinygrad.dtype import AddrSpace, PtrDType, DType, dtypes
+from tinygrad.dtype import AddrSpace, DType, dtypes
 from tinygrad.codegen.opt.tc import TensorCore
 from tinygrad.device import Compiler

@@ -29,8 +29,8 @@ class Estimates:
      def range_gate(x): return x.op is not Ops.RANGE
      for u in uops:
        if u.op in {Ops.LOAD, Ops.STORE}:
-          # if u.src[0] is SLICE, we have to include the buffer since it might be an AFTER
-          dont_count = dont_count.union((UOp.sink(*u.src[0].src[1:]) if u.src[0].op is Ops.SLICE else u.src[0]).toposort(range_gate))
+          # if u.src[0] is INDEX, we have to include the buffer since it might be an AFTER
+          dont_count = dont_count.union((UOp.sink(*u.src[0].src[1:]) if u.src[0].op is Ops.INDEX else u.src[0]).toposort(range_gate))
          # TODO: is this correct? this all needs to be cleaned up
          if len(u.src) > 2: dont_count = dont_count.union(u.src[2].toposort())
        elif u.op is Ops.IF:
@@ -38,11 +38,11 @@ class Estimates:
    for u in uops:
      if u.op in {Ops.LOAD, Ops.STORE}:
        buf = u
-        while len(buf.src): buf = buf.src[0]
+        while len(buf.src) and buf.op is not Ops.PARAM: buf = buf.src[0]
        if buf.op is Ops.PARAM:
-          # u.src[0] is SLICE, cap at buffer size for re-reads (e.g. matmul)
-          accessed = mem.get((buf, u.op), 0) + u.src[0].dtype.base.itemsize * mults
-          mem[(buf, u.op)] = smin(accessed, buf.ptrdtype.nbytes()) if buf.ptrdtype.size != -1 else accessed
+          # u.src[0] is INDEX, cap at buffer size for re-reads (e.g. matmul)
+          accessed = mem.get((buf, u.op), 0) + u.max_numel() * u.src[0].dtype.itemsize * mults
+          mem[(buf, u.op)] = smin(accessed, buf.max_numel() * buf.dtype.itemsize)
      if u.op is Ops.RANGE:
        mult_stack.append(mults)
        mults *= cast(sint, u.src[0].ssimplify())
@@ -51,10 +51,10 @@ class Estimates:
      elif u.op is Ops.END: mults = mult_stack.pop(-1)
      elif u.op is Ops.SPECIAL: mults *= cast(sint, u.src[0].ssimplify()) # NOTE: we don't push to the mult_stack here, you can't end these
      elif u.op is Ops.DEFINE_VAR and u.arg[0] == 'core_id': mults *= u.arg[2] + 1
-      elif u.op is Ops.LOAD and (not isinstance(u.src[0].dtype, PtrDType) or u.src[0].dtype.addrspace != AddrSpace.REG):
-        lds += u.dtype.itemsize * mults
-      elif u.op is Ops.STORE and (not isinstance(u.src[0].dtype, PtrDType) or u.src[0].dtype.addrspace != AddrSpace.REG):
-        lds += u.src[1].dtype.itemsize * mults
+      elif u.op is Ops.LOAD and u.src[0].addrspace != AddrSpace.REG:
+        lds += u.max_numel() * u.dtype.itemsize * mults
+      elif u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG:
+        lds += u.max_numel() * u.src[1].dtype.itemsize * mults
      elif u.op in GroupOp.ALU and u not in dont_count: flops += (mults * (2 if u.op is Ops.MULACC else 1)) * u.dtype.count
      elif u.op is Ops.WMMA and u not in dont_count: flops += 2 * prod(u.arg[1]) // u.arg[5] * mults
    return Estimates(flops, lds, sum(mem.values()))
@@ -43,10 +43,8 @@ base_rewrite = PatternMatcher([
  (UPat(Ops.CONST, (dtypes.int8, dtypes.int16), name="x"), lambda ctx,x: f"({ctx.render_cast(x.dtype, str(x.arg))})"),
  # default const render
  (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.arg)),
-  # slice is ptr arithmetic
-  (UPat(Ops.SLICE, src=(UPat.var("buf"), UPat.var('idx')), name="x"),
-   lambda ctx,buf,idx,x: ctx.render_cast(x.dtype, f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})")),
  # new load/store
+  (UPat.var("buf").index(UPat.var('idx')), lambda ctx,buf,idx: f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"),
  (UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"(*{ctx[bidx]})"),
  (UPat(Ops.LOAD, src=(UPat.var("bidx"), UPat.var("var"), UPat.var("gate"))), lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"),
  (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
@@ -133,12 +131,12 @@ class CStyleLanguage(Renderer):
  string_rewrite = base_rewrite
  extra_matcher = extra_pm

-  def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[DType,bool]]], uops:list[UOp], prefix=None) -> str:
+  def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[UOp,bool]]], uops:list[UOp], prefix=None) -> str:
    tmp = ""
-    if any(isinstance(dtype, ImageDType) for _,(dtype,_) in bufs):
+    if any(isinstance(u.dtype, ImageDType) for _,(u,_) in bufs):
      tmp = "const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n"
-    buftypes = [(name, self.render_dtype(dtype, mutable)+self.buffer_suffix if isinstance(dtype, (ImageDType, PtrDType)) else
-                self.arg_int_prefix if dtype == dtypes.int else None) for name,(dtype,mutable) in bufs]
+    buftypes = [(name, self.render_dtype(u.dtype, mutable)+self.buffer_suffix if isinstance(u.dtype, (ImageDType, PtrDType)) else
+                self.arg_int_prefix if u.dtype == dtypes.int else None) for name,(u,mutable) in bufs]
    local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]
    launch_bounds = prod([d.vmax for d in local_dims])
    prg = ''.join([f"{self.kernel_typedef.format(launch_bounds=launch_bounds)} {function_name}(",] +
@@ -158,14 +156,14 @@ class CStyleLanguage(Renderer):
    return self.type_map.get(scalar:=dt.scalar(), scalar.name)

  def __getitem__(self, key): return self.r[key]  # hacky helper
-  def _render(self, uops:list[UOp]) -> tuple[str, list[str], list[tuple[str,tuple[DType,bool]]]]:
+  def _render(self, uops:list[UOp]) -> tuple[str, list[str], list[tuple[str,tuple[UOp,bool]]]]:
    r: dict[UOp, str] = {}
    self.r = r

    child_count = Counter(v for ru in uops for v in ru.src)
    # find which PARAMs are stored to with a single toposort
    writable_params = {u for u in UOp.sink(*[u.src[0] for u in uops if u.op is Ops.STORE]).toposort(lambda u: u.op != Ops.END) if u.op is Ops.PARAM}
-    bufs: dict[UOp, tuple[str, tuple[DType, bool]]] = {}
+    bufs: dict[UOp, tuple[str, tuple[UOp, bool]]] = {}
    kernel = []
    depth = 1
    c: defaultdict[str, int] = defaultdict(int)
@@ -180,9 +178,9 @@ class CStyleLanguage(Renderer):
        continue
      if u.op in (Ops.PARAM, Ops.DEFINE_VAR):
        if u.op is not Ops.PARAM: r[u] = u.arg[0]
-        elif isinstance(u.dtype, ImageDType): r[u] = f"data{u.arg}_{u.dtype.shape[0]}x{u.dtype.shape[1]}"
-        else: r[u] = f"data{u.arg}_{sz}" if (sz:=u.ptrdtype.size) > 0 else f"data{u.arg}"
-        bufs[u] = (r[u], (u.dtype, u in writable_params))
+        elif isinstance(u.dtype, ImageDType): r[u] = f"data{u.arg.slot}_{u.dtype.shape[0]}x{u.dtype.shape[1]}"
+        else: r[u] = f"data{u.arg.slot}_{sz}" if (sz:=u.max_numel()) > 0 else f"data{u.arg.slot}"
+        bufs[u] = (r[u], (u, u in writable_params))
        continue

      # naming
@@ -192,15 +190,15 @@ class CStyleLanguage(Renderer):
      else:
        prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const",
                  Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.STACK: "cast",
-                  Ops.SLICE: "bidx", Ops.DEFINE_REG: "acc", Ops.LOAD: "val"}.get(u.op, "alu")
+                  Ops.INDEX: "bidx", Ops.DEFINE_REG: "acc", Ops.LOAD: "val"}.get(u.op, "alu")
        r[u] = f"{prefix}{c[prefix]}"

      l = cast(str, self.string_rewrite.rewrite(u, ctx=self))
      assert l is not None, f"failed to render {u.op} {u.dtype} {[(x.op,x.dtype) for x in u.src]} {u.arg}"

      if u.op in {Ops.ENDIF, Ops.END}: depth -= 1
-      if (u.op is not Ops.CAST or u.dtype.vcount == 1) and (u.op in {Ops.CONST, Ops.GEP, Ops.SLICE, Ops.CUSTOMI} or \
-        (u.op is Ops.LOAD and u.src[0].ptrdtype.addrspace == AddrSpace.REG) or \
+      if (u.op is not Ops.CAST or u.dtype.vcount == 1) and (u.op in {Ops.CONST, Ops.GEP, Ops.INDEX, Ops.CUSTOMI} or \
+        (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG) or \
        (u.op is Ops.CAST and isinstance(u.dtype, PtrDType)) or \
        (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):
        r[u] = l
@@ -246,7 +244,7 @@ class ClangRenderer(CStyleLanguage):
    kernel_typedef = "__attribute__((ms_abi)) void"
  def render_vector_prefix(self, dt:DType) -> str:
    # round (down) to power of two (this is actually the default clang behavior)
-    alignment = 2**int(math.log2(dt.itemsize)) if getenv("ALIGNED", 1) and not dtypes.is_bool(dt) else 1
+    alignment = 2**int(math.log2(dt.itemsize*dt.count)) if getenv("ALIGNED", 1) and not dtypes.is_bool(dt) else 1
    return f"typedef {self.render_dtype(dt.scalar())} {self.render_dtype(dt)} __attribute__((aligned({alignment}),ext_vector_type({dt.count})));"

  def _render_defines(self, uops) -> list[str]:
@@ -268,7 +266,7 @@ class ClangRenderer(CStyleLanguage):
  AMX_SET(1);\n  return data0;\n}}"""]
    return prefix
  def _render_body(self, function_name, kernel, bufs, uops, pref=None) -> str: return super().render_kernel(function_name, kernel, bufs, uops, pref)
-  def _render_entry(self, function_name:str, bufs:list[tuple[str,tuple[DType,bool]]]) -> str: return ""
+  def _render_entry(self, function_name:str, bufs:list[tuple[str,tuple[UOp,bool]]]) -> str: return ""

  def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:
    defines = '\n'.join(self._render_defines(uops))
@@ -277,12 +275,11 @@ class ClangRenderer(CStyleLanguage):
  def supported_dtypes(self):
    return {d for d in super().supported_dtypes() if (d != dtypes.bfloat16 or self.target.arch.startswith(("x86", "arm"))) and d not in dtypes.fp8s}

-class ClangJITRenderer(ClangRenderer):
  def __init__(self, target:Target):
    super().__init__(target)
-    from tinygrad.runtime.support.compiler_cpu import ClangJITCompiler
+    from tinygrad.runtime.support.compiler_cpu import ClangCompiler
    if "AMX" in target.arch: self.tensor_cores = tc.amx
-    self.compiler = ClangJITCompiler([x for x in target.arch.split(",") if x != "AMX"])
+    self.compiler = ClangCompiler([x for x in target.arch.split(",") if x != "AMX"])

 class OpenCLRenderer(CStyleLanguage):
  has_aux = True
@@ -306,12 +303,11 @@ class OpenCLRenderer(CStyleLanguage):
      lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.arg)))[0] >> 16)}u"),
    # load/store image (OpenCL)
    (UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), lambda ctx,buf,idx_y,idx_x: f"IMAGE<{ctx[buf]}, {ctx[idx_y]}, {ctx[idx_x]}>"),
-    (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var"), UPat.var("gate"))),
+    (UPat(Ops.LOAD, dtype=dtypes.float, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var"), UPat.var("gate"))),
      lambda ctx,buf,idx_y,idx_x,var,gate: f"({ctx[gate]}?read_imagef({ctx[buf]}, smp, (int2)({ctx[idx_x]},{ctx[idx_y]})):{ctx[var]})"),
-    (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')),)),
+    (UPat(Ops.LOAD, dtype=dtypes.float, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')),)),
      lambda ctx,buf,idx_y,idx_x: f"read_imagef({ctx[buf]}, smp, (int2)({ctx[idx_x]},{ctx[idx_y]}))"),
-    (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')),
-                          UPat.var("var", dtypes.float.vec(4)))),
+    (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var", dtypes.float))),
      lambda ctx,buf,idx_y,idx_x,var: f"write_imagef({ctx[buf]}, (int2)({ctx[idx_x]},{ctx[idx_y]}), {ctx[var]});"),
  ]) + base_rewrite

@@ -322,8 +318,8 @@ class OpenCLRenderer(CStyleLanguage):
  def aux(self, uops:list[UOp]):
    arg_dtypes:list[list[tuple[int, DType]]] = []
    for i,u in enumerate(u for u in uops if u.op is Ops.PARAM):
-      if len(arg_dtypes) >= u.arg: arg_dtypes.append([])
-      arg_dtypes[u.arg].append((i, u.dtype))
+      while len(arg_dtypes) <= u.arg.slot: arg_dtypes.append([])
+      arg_dtypes[u.arg.slot].append((i, u.dtype))
    return tuple(tuple(a) for a in arg_dtypes),

  def supported_dtypes(self): return {d for d in super().supported_dtypes()
@@ -437,7 +433,8 @@ class CUDARenderer(CStyleLanguage):
  def render_vector_prefix(self, dt:DType) -> str:
    vec, scal = self.render_dtype(dt), self.render_dtype(dt.scalar()),
    elems, header = ', '.join(_nms[:dt.count]), ', '.join([f"{scal} {x}" for x in _nms[:dt.count]])
-    return f"struct __align__({dt.itemsize}) {vec} {{ {scal} {elems}; }}; __device__ {vec} make_{vec}({header}) {{ {vec} r={{{elems}}}; return r; }}"
+    return f"struct __align__({dt.itemsize*dt.count}) {vec} {{ {scal} {elems}; }};" + \
+           f"__device__ {vec} make_{vec}({header}) {{ {vec} r={{{elems}}}; return r; }}"

  def render_kernel(self, function_name, kernel, bufs, uops, prefix=None):
    # TODO: why is dtypes.bfloat16.name == "__bf16"? would be easier not override dtypes.name
@@ -454,7 +451,8 @@ class CUDARenderer(CStyleLanguage):
    for name, (N, M, K), dtype_in, dtype_out, _, _, upcast_axes, _ in wmma_args(uops):
      upcast_sizes = [prod(size for _, size in upcast) for upcast in upcast_axes]
      wmma_dtypes = [self.render_dtype(dtype.vec(size)) for dtype, size in zip([dtype_in, dtype_in, dtype_out], upcast_sizes)]
-      n_operands = [size*dtype.itemsize//4 for dtype, size in zip([dtype_in, dtype_in, dtype_out], upcast_sizes)] # 4 => CUDA reg size in bytes
+      # 4 => CUDA reg size in bytes
+      n_operands = [size*dtype.itemsize*dtype.count//4 for dtype, size in zip([dtype_in, dtype_in, dtype_out], upcast_sizes)]
      operands = [f"%{i}" for i in range(sum(n_operands))]

      # mma operands => {c}, {a}, {b}, {c}
@@ -633,8 +633,8 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
    reg = cast(int, cast(Register, reg_uop.reg).index if reg_uop is not None else reg)
    rm = cast(Register, rm_uop.reg).index
    idx = cast(Register, idx_uop.reg).index if idx_uop is not None and idx_uop.reg is not None else 4
-    rm_sz = 8 if isinstance(rm_uop.dtype, PtrDType) and disp_uop is None else rm_uop.dtype.itemsize
-    reg_sz = (reg_uop.dtype.itemsize if not isinstance(reg_uop.dtype, PtrDType) else 8) if reg_uop is not None else 0
+    rm_sz = 8 if isinstance(rm_uop.dtype, PtrDType) and disp_uop is None else (rm_uop.dtype.itemsize*rm_uop.dtype.count)
+    reg_sz = ((reg_uop.dtype.itemsize*reg_uop.dtype.count) if not isinstance(reg_uop.dtype, PtrDType) else 8) if reg_uop is not None else 0
    sz = reg_sz or rm_sz

    # encode instruction
@@ -74,9 +74,8 @@ lop = {**{x:unsigned_lop for x in (dtypes.bool,)+dtypes.uints}, **{x:signed_lop

 base_rewrite = PatternMatcher([
  # memory load/store
-  (UPat(Ops.SLICE, name="x"), lambda ctx,x:
-   f"  {ctx[x]}_o = getelementptr inbounds {ldt(x.src[0].dtype.base)}, {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}\n"
-   f"  {ctx[x]} = bitcast {ldt(x.src[0].dtype)} {ctx[x]}_o to {ldt(x.dtype)}"),
+  (UPat(Ops.INDEX, name="x"), lambda ctx,x:
+   f"  {ctx[x]} = getelementptr inbounds {ldt(x.dtype.base)}, {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}"),
  (UPat(Ops.LOAD, src=(UPat.var("idx"), UPat.var("alt"), UPat.var("mask")), name="x"),
   lambda ctx,x,idx,alt,mask:
   f"  br label {ctx[x]}_entry\n{ctx[x][1:]}_entry:\n"
@@ -136,13 +135,13 @@ class LLVMRenderer(Renderer):
  code_for_op = {k:lambda:None for v in lop.values() for k in v.keys()}

  extra_matcher = create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast
-  def _render_fn(self, name:str, args:list[tuple[str,DType]], kernel:list[str], prefix:list[str]|None=None) -> str:
+  def _render_fn(self, name:str, args:list[tuple[str,UOp]], kernel:list[str], prefix:list[str]|None=None) -> str:
    # NOTE: CPUAllocator promises 0x20 alignment
-    sargs = ", ".join([f"{ldt(dt)}{' noalias align 32' if isinstance(dt, PtrDType) else ''} {name}" for name,dt in args])
+    sargs = ", ".join([f"{ldt(u.dtype)}{' noalias align 32' if isinstance(u.dtype, PtrDType) else ''} {name}" for name,u in args])
    return "\n".join((prefix or []) + [f"define{' ' + self.abi if self.abi else ''} void @{name}({sargs}) #0", "{"] + kernel + ["  ret void\n}"])
  def _render_kernel(self, uops: list[UOp], prefix:list[str]|None=None) -> tuple[tuple[str, ...], str]:
    r: dict[UOp, str] = {}
-    args: list[tuple[str, DType]] = []
+    args: list[tuple[str, UOp]] = []
    kernel: list[str] = []
    vc = -1

@@ -165,18 +164,18 @@ class LLVMRenderer(Renderer):
        if u.arg is not None: name = u.arg.function_name
        continue
      if u.op in (Ops.PARAM, Ops.DEFINE_VAR):
-        r[u] = f"%data{u.arg}" if u.op is Ops.PARAM else f"%{u.expr}"
-        args.append((r[u], u.dtype))
+        r[u] = f"%data{u.arg.slot}" if u.op is Ops.PARAM else f"%{u.expr}"
+        args.append((r[u], u))
      elif u.op in (Ops.DEFINE_LOCAL, Ops.DEFINE_REG):
        r[u] = f"%{'local' if u.op is Ops.DEFINE_LOCAL else 'reg'}_{str(u.arg).replace('(', '').replace(')', '').replace(',', '_').replace(' ', '')}"
-        assert isinstance(u.dtype, PtrDType)
+        size = u.max_numel()
        if u.op is Ops.DEFINE_REG:
-          kernel.append(f"  {r[u]} = alloca [{u.dtype.size} x {ldt(u.dtype.base)}]")
+          kernel.append(f"  {r[u]} = alloca [{size} x {ldt(u.dtype.base)}]")
        elif self.has_local:
-          local_args.append(f"@{r[u][1:]} = internal unnamed_addr addrspace(3) global [{u.dtype.size} x {ldt(u.dtype)}] undef, align 16")
-          kernel.append(f"  {r[u]} = addrspacecast [{u.dtype.size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{u.dtype.size} x {ldt(u.dtype)}]*")
+          local_args.append(f"@{r[u][1:]} = internal unnamed_addr addrspace(3) global [{size} x {ldt(u.dtype)}] undef, align 16")
+          kernel.append(f"  {r[u]} = addrspacecast [{size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{size} x {ldt(u.dtype)}]*")
        else:
-          kernel.append(f"  {r[u]} = alloca [{u.dtype.size} x {ldt(u.dtype.base)}], align 16")
+          kernel.append(f"  {r[u]} = alloca [{size} x {ldt(u.dtype.base)}], align 16")
      elif u.op is Ops.CONST: r[u] = lconst(u.arg, u.dtype)
      elif u.op is Ops.CAST and (ldt(u.dtype) == ldt(u.src[0].dtype) or isinstance(u.dtype, PtrDType)):
        r[u] = r[u.src[0]] # cast from signed to unsigned of the same size is a noop, or pointer cast
@@ -71,12 +71,12 @@ def nimm_set(imm:mesa.nir_def, x, dtype:DType):
  instr = ctypes.cast(imm.parent_instr, ctypes.POINTER(mesa.nir_load_const_instr))
  struct.pack_into(unwrap(dtype.fmt), (ctypes.c_ubyte * dtype.itemsize).from_address(ctypes.addressof(instr.contents.value)), 0, truncate[dtype](x))

-@nir_instr(nc=1, bs=lambda dtype: dtype.bitsize)
+@nir_instr(nc=1, bs=lambda dtype: dtype.bitsize*dtype.count)
 def nimm(b:mesa.nir_builder, x, dtype:DType) -> mesa.nir_def:
-  nimm_set((instr:=mesa.nir_load_const_instr_create(b.shader, 1, dtype.bitsize)).contents._def, x, dtype)
+  nimm_set((instr:=mesa.nir_load_const_instr_create(b.shader, 1, dtype.bitsize*dtype.count)).contents._def, x, dtype)
  return instr
-@nir_instr(nc=1, bs=lambda dtype: dtype.bitsize)
-def nundef(b, dtype): return mesa.nir_undef_instr_create(b.shader, 1, dtype.bitsize)
+@nir_instr(nc=1, bs=lambda dtype: dtype.bitsize*dtype.count)
+def nundef(b, dtype): return mesa.nir_undef_instr_create(b.shader, 1, dtype.bitsize*dtype.count)

 deref_var = nir_instr(nc=1, bs=32, modes=lambda var:var.data.mode, type=lambda var:var.type, var=lambda var:ctypes.pointer(var))( # pylint: disable=W0108
  lambda b, var: mesa.nir_deref_instr_create(b.shader, mesa.nir_deref_type_var))
@@ -86,7 +86,7 @@ def scope(space): return 'global' if space == AddrSpace.GLOBAL else ('shared' if
 nstore = nir_instr(has_def=False, df=lambda addr:addr, intrins=lambda space,val: {"WRITE_MASK":(1<<val.num_components)-1, **iointr(space)},
  num_components=lambda val:val.num_components, srcs=lambda space, addr, val: [nsrc(val), nsrc(addr)][::1 if space != AddrSpace.REG else -1])(
    lambda b, space, addr, val, dtype: mesa.nir_intrinsic_instr_create(b.shader, g(f"nir_intrinsic_store_{scope(space)}")))
-nload = nir_instr(nc=lambda dtype:dtype.count, bs=lambda dtype:dtype.bitsize//dtype.count, num_components=lambda dtype:dtype.count,
+nload = nir_instr(nc=lambda dtype:dtype.count, bs=lambda dtype:dtype.bitsize, num_components=lambda dtype:dtype.count,
  intrins=lambda space:{**({"ACCESS":mesa.ACCESS_CAN_REORDER} if space==AddrSpace.GLOBAL else {}), **iointr(space)}, srcs=lambda addr: [nsrc(addr)])(
    lambda b, space, addr, dtype: mesa.nir_intrinsic_instr_create(b.shader, g(f"nir_intrinsic_load_{scope(space)}")))

@@ -149,12 +149,12 @@ class NIRRenderer(Renderer):
    (UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 4)),
    (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))),
    (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off"))).or_casted(), UPat.var("val"))),
-     lambda ctx,buf,off,val: nstore(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), ctx.r[val], val.dtype)),
+     lambda ctx,buf,off,val: nstore(ctx.b, buf.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), ctx.r[val], val.dtype)),
    (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))).or_casted(), UPat.var("alt"), UPat.var("gate")), name="x"),
     lambda ctx,x,buf,off,alt,gate: if_phi(ctx.b, ctx.r[gate],
-      lambda: nload(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype, ctx.r[gate]), x.dtype), lambda: ctx.r[alt])),
+      lambda: nload(ctx.b, buf.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype, ctx.r[gate]), x.dtype), lambda: ctx.r[alt])),
    (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))).or_casted(),), name="x"),
-     lambda ctx,x,buf,off: nload(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), x.dtype)),
+     lambda ctx,x,buf,off: nload(ctx.b, buf.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), x.dtype)),
    (UPat(Ops.STACK, name="x"), lambda ctx,x: nalu(ctx.b, f"vec{x.dtype.count}", *[ctx.r[src] for src in x.src])),
    (UPat(GroupOp.ALU, name="x"), lambda ctx,x: nalu(ctx.b, aop[x.src[0].dtype.scalar()][x.op], *[ctx.r[src] for src in x.src])),
    (UPat(Ops.CAST, name="x"), lambda ctx,x: ncast(ctx.b, ctx.r[x.src[0]], x.src[0].dtype, x.dtype)),
@@ -63,7 +63,7 @@ def mem_type(x:UOp) -> str:
  match x.op:
    case Ops.AFTER: return mem_type(x.src[0])
    case Ops.DEFINE_LOCAL: return 'shared'
-    case Ops.PARAM: return 'global'
+    case Ops.PARAM: return 'shared' if x.addrspace == AddrSpace.LOCAL else 'global'
    case _: raise RuntimeError(f"{x.op} needs to be memory")

 def render_wmma(ctx: "PTXRenderer", wmma: UOp):
@@ -90,7 +90,7 @@ string_rewrite = PatternMatcher([
  (UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.arg, x.dtype)}, 0;"),
  (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.arg, x.dtype)};"),
  (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"mov.u32 %{x.arg}, %{'ctaid' if x.arg[0] == 'g' else 'tid'}.{chr(120+int(x.arg[-1]))};"),
-  (UPat(Ops.PARAM, name="x"), lambda ctx, x: f"ld.param.{ctx.types[dtypes.ulong]} {ctx.r[x]}, [data{x.arg}+0];"),
+  (UPat(Ops.PARAM, name="x"), lambda ctx, x: f"ld.param.{ctx.types[dtypes.ulong]} {ctx.r[x]}, [data{x.arg.slot}+0];"),
  (UPat((Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ), name="x", allow_any_len=True, src=(UPat.var("src0"),)),
    lambda ctx, x, src0: ctx.code_for_op[x.op](ctx.r[x], *[ctx.r[v] for v in x.src], src0.dtype, ctx.types[src0.dtype])),
  (UPat(GroupOp.ALU, name="x"), lambda ctx, x: ctx.code_for_op[x.op](ctx.r[x], *[ctx.r[v] for v in x.src], x.dtype, ctx.types[x.dtype])),
@@ -202,7 +202,7 @@ class PTXRenderer(Renderer):
        r[u] = r[u.src[0]]
        continue
      if u.op is Ops.DEFINE_REG:
-        r[u] = [ssa("reg", u, self.types[u.dtype.base.scalar()]) for _ in range(u.ptrdtype.size)]
+        r[u] = [ssa("reg", u, self.types[u.dtype.base.scalar()]) for _ in range(u.max_numel())]
        continue
      if u.op in {Ops.INDEX, Ops.LOAD, Ops.STORE} and isinstance(u.src[0].dtype, PtrDType) and u.src[0].dtype.addrspace == AddrSpace.REG:
        if u.op is Ops.INDEX:
@@ -219,7 +219,7 @@ class PTXRenderer(Renderer):
      elif u.op is Ops.DEFINE_VAR: bufs.append((u.expr, u.dtype))
      elif u.op is Ops.LOAD:
        r[u] = [ssa('val', dtype=self.types[u.dtype.scalar()]) for _ in range(u.dtype.count)] if u.dtype.count > 1 else ssa('val', u)
-      elif u.op is Ops.PARAM: bufs.append((f"data{u.arg}", u.dtype))
+      elif u.op is Ops.PARAM: bufs.append((f"data{u.arg.slot}", u.dtype))
      elif u.op is Ops.WMMA:
        # registers for packing/unpacking input and acc
        self.wmma_r = [[ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[0]]), 4 // u.src[0].dtype.scalar().itemsize)],
@@ -27,10 +27,8 @@ def packed_load(root:UOp, bidx:UOp, dtype:DType, var:UOp|None=None, gate:UOp|Non
  val = (load.cast(dtypes.uint32) >> shift_am) & mask
  return sign_extend(val, 8*dtype.itemsize).cast(dtype) if dtype in [dtypes.char, dtypes.short] else val.cast(dtype)

-def is_packed(dt:DType, odt:DType|None = None) -> bool:
-  if odt is None: odt = dt
-  return dt.itemsize < 4 and dt.base != dtypes.half and (not isinstance(odt, PtrDType) or odt.addrspace != AddrSpace.REG)
-def _packed_size(dt:PtrDType): return dt.size // (4//dt.itemsize) if is_packed(dt) else dt.size
+def is_packed(u:UOp) -> bool: return u.dtype.itemsize < 4 and u.dtype.base != dtypes.half and u.addrspace != AddrSpace.REG
+def _packed_size(u:UOp): return u.max_numel() // (4//u.dtype.itemsize) if is_packed(u) else u.max_numel()

 def is_nan(a):
  bs, (exp, mant) = a.dtype.bitsize, dtypes.finfo(a.dtype)
@@ -41,12 +39,12 @@ wgsl_matcher = PatternMatcher([
   lambda a,b,c: a.cast(dtypes.int).alu(c.op, b.cast(dtypes.int)).cast(dtypes.bool)),
  # TODO: load alt value doesnt have to be a const
  (UPat.load(UPat.var("b"), UPat.cvar("c"), UPat.var("gate"), name="l"),
-   lambda l,b,c,gate: packed_load(l,b,l.dtype,c.cast(dtypes.uint32),gate) if is_packed(l.dtype, b.dtype) else None),
-  (UPat.load(UPat.var("b"), name='l'), lambda l,b: packed_load(l, b, l.dtype) if is_packed(l.dtype, b.dtype) else None),
+   lambda l,b,c,gate: packed_load(l,b,l.dtype,c.cast(dtypes.uint32),gate) if is_packed(b) else None),
+  (UPat.load(UPat.var("b"), name='l'), lambda l,b: packed_load(l, b, l.dtype) if is_packed(b) else None),
  (UPat.store(UPat.var("bidx"), UPat.var("var"), UPat.var("gate")),
-   lambda bidx,var,gate: packed_store(bidx,var,gate) if is_packed(var.dtype, bidx.dtype) else None),
+   lambda bidx,var,gate: packed_store(bidx,var,gate) if is_packed(bidx) else None),
  (UPat.store(UPat.var("bidx"), UPat.var("var")),
-   lambda bidx,var: packed_store(bidx,var) if is_packed(var.dtype, bidx.dtype) else None),
+   lambda bidx,var: packed_store(bidx,var) if is_packed(bidx) else None),
  (UPat.var("a") << UPat.var("b"),lambda a,b:(a.bitcast(dtypes.uint32)<<b.cast(dtypes.uint32)).bitcast(a.dtype) if b.dtype!=dtypes.uint32 else None),
  (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
  # fix nan check: 'a != a -> is_nan()'
@@ -71,8 +69,8 @@ class WGSLRenderer(CStyleLanguage):
    (UPat(Ops.CONST, dtype=(dtypes.uchar, dtypes.ushort, dtypes.uint32), name="x"),
     lambda x: f"bitcast<u32>({x.arg})" if x.arg < 0 else f"{x.arg&0xFFFFFFFF}u"),
    (UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}"),
-    (UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x: f"var<workgroup> {ctx[x]}: array<{ctx.buf_map(x.dtype.base)},{_packed_size(x.dtype)}>;"),
-    (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"var {ctx[x]}: array<{ctx.buf_map(x.dtype)},{_packed_size(x.dtype)}>;"),
+    (UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x: f"var<workgroup> {ctx[x]}: array<{ctx.buf_map(x)},{_packed_size(x)}>;"),
+    (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"var {ctx[x]}: array<{ctx.buf_map(x)},{_packed_size(x)}>;"),
    (UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)),
     lambda ctx,x: f"bitcast<vec2<f16>>({ctx[x.src[0]]})[0]"),
    (UPat(Ops.BITCAST, dtype=dtypes.uchar, name="x"), lambda ctx,x: f"bitcast<u32>({ctx[x.src[0]]}&0xFF)"),
@@ -84,21 +82,21 @@ class WGSLRenderer(CStyleLanguage):
    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast<{ctx.type_map[x.dtype]}>({ctx[x.src[0]]})"),
    # TODO: load alt value doesnt have to be a const
    (UPat.load(UPat.var("b"), UPat.cvar("v"), UPat.var("gate")),
-      lambda ctx,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b],b.src[0].dtype)}, {ctx[gate]})"),
-    (UPat.load(UPat.var("b")), lambda ctx, b: ctx.render_load(ctx[b], b.dtype)),
+      lambda ctx,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b], b.src[0])}, {ctx[gate]})"),
+    (UPat.load(UPat.var("b")), lambda ctx, b: ctx.render_load(ctx[b], b)),
    (UPat.store(UPat.var("b"), UPat.var("v")), lambda ctx,b,v:\
     # (load & mask) | var -> mask = v.src[0].src[1], var = v.src[1]
-     f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n  atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b.src[0].dtype) \
+     f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n  atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b.src[0]) \
      else f"{ctx[b]} = {ctx[v]};"),
-    (UPat(Ops.SLICE, src=(UPat.var("b"), UPat.var("idx"))),
+    (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"))),
     lambda ctx,b,idx: f"{ctx[b]}[{strip_parens(ctx[idx]) if idx.arg is Ops.ADD else ctx[idx]}]"),
  ]) + base_rewrite

  def render_cast(self, dt:DType, val: str) -> str: return f"{self.type_map[dt]}({val})"
  def render_dtype(self, dt:DType, mutable=True) -> str: return "var"
-  def render_load(self, x:str, dt:DType) -> str: return f"atomicLoad(&{x})" if is_packed(dt) else x
-  def buf_map(self, dt:DType) -> str: return "atomic<u32>" if is_packed(dt) else self.type_map[dt.base]
-  def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[DType,bool]]], uops:list[UOp], prefix=None) -> str:
+  def render_load(self, x:str, u:UOp) -> str: return f"atomicLoad(&{x})" if is_packed(u) else x
+  def buf_map(self, u:UOp) -> str: return "atomic<u32>" if is_packed(u) else self.type_map[u.dtype.base]
+  def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[UOp,bool]]], uops:list[UOp], prefix=None) -> str:
    local_size = [u.src[0].ssimplify() for u in sorted([u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == 'l'], key=lambda u: u.arg)]
    if not local_size: local_size = [1]
    bind_it = iter(range(len(bufs)))
@@ -108,8 +106,8 @@ class WGSLRenderer(CStyleLanguage):
    prg += "fn nan() -> f32 { let bits = 0xffffffffu; return bitcast<f32>(bits); }\n"
    prg += "@group(0) @binding(0)\nvar<uniform> INFINITY : f32;\n"
    prg += "\n".join((external_local_bufs or [])+[f"@group(0) @binding({next(bind_it)+1})" +
-      f"{'var<storage,read_write>' if isinstance(dtype, PtrDType) else 'var<uniform>'}" +
-      f"{name}:{f'array<{self.buf_map(dtype.base)}>' if isinstance(dtype,PtrDType) else self.buf_map(dtype)};" for name,(dtype,_) in bufs])
+      f"{'var<storage,read_write>' if isinstance(u.dtype, PtrDType) else 'var<uniform>'}" +
+      f"{name}:{f'array<{self.buf_map(u)}>' if isinstance(u.dtype,PtrDType) else self.buf_map(u)};" for name,(u,_) in bufs])
    prg += f"\n@compute @workgroup_size({','.join([str(x) for x in local_size])}) fn {function_name}(@builtin(workgroup_id) gindex: vec3<u32>,"
    return prg + "@builtin(local_invocation_id) lindex: vec3<u32>) {\n" + "\n".join(kernel) + "\n}"

@@ -4,7 +4,7 @@ from tinygrad.helpers import to_mv, OSX, WIN, mv_address, suppress_finalizing, u
 from tinygrad.device import BufferSpec
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
 from tinygrad.runtime.support.hcq import CLikeArgsState
-from tinygrad.renderer.cstyle import ClangJITRenderer
+from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.renderer.llvmir import CPULLVMRenderer
 from tinygrad.renderer.nir import LVPRenderer
 from tinygrad.renderer.isa.x86 import X86Renderer
@@ -138,5 +138,5 @@ class CPUDevice(HCQCompiled):
  def __init__(self, device:str=""):
    self.tasks:queue.Queue = queue.Queue()
    CPUWorker(self, self.tasks, thread_id=0).start()
-    super().__init__(device, CPUAllocator(self), [ClangJITRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], functools.partial(CPUProgram, self),
+    super().__init__(device, CPUAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], functools.partial(CPUProgram, self),
                     CPUSignal, CPUComputeQueue, arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native")
@@ -2,7 +2,7 @@ from __future__ import annotations
 import ctypes, os, mmap, tempfile, pathlib, array, functools, threading, contextlib, sys, subprocess, struct
 assert sys.platform != 'win32'
 from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler
-from tinygrad.dtype import dtypes, DType, PtrDType
+from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import getenv, round_up, mv_address, to_mv, cpu_objdump, system, DEBUG, suppress_finalizing, Target
 from tinygrad.renderer.cstyle import ClangRenderer
@@ -53,18 +53,19 @@ class DSPRenderer(ClangRenderer):
      'void* HAP_mmap(void *addr, int len, int prot, int flags, int fd, long offset);', 'int HAP_munmap(void *addr, int len);',
      'unsigned long long HAP_perf_get_time_us(void);'] + super()._render_defines(uops)

-  def _render_entry(self, function_name:str, bufs:list[tuple[str,tuple[DType,bool]]]) -> str:
+  def _render_entry(self, function_name:str, bufs:list[tuple[str,tuple[UOp,bool]]]) -> str:
    msrc = ['int entry(unsigned long long handle, unsigned int sc, remote_arg* pra) {',
            'struct dcvs_v2_req req = {.type=7, .dcvs_enable=0, .set_latency=1, .latency=100, .set_dcvs_params=1, .target_corner = 6 /* TURBO */};',
            'HAP_power_set((void*)handle, (void*)&req);']
    msrc += ['if ((sc>>24) != 2) return 0;']
    msrc += [f'int sz_or_val_{i} = ((int*)pra[0].buf.pv)[{i}];' for i,b in enumerate(bufs)]
-    msrc += [f'int off{i} = ((int*)pra[1].buf.pv)[{i}];' for i,b in enumerate(bufs) if isinstance(b[1][0], PtrDType)]
-    msrc += [f'void *buf_{i} = HAP_mmap(0,sz_or_val_{i},3,0,pra[{i+3}].dma.fd,0)+off{i};' for i,b in enumerate(bufs) if isinstance(b[1][0], PtrDType)]
+    msrc += [f'int off{i} = ((int*)pra[1].buf.pv)[{i}];' for i,b in enumerate(bufs) if isinstance(b[1][0].dtype, PtrDType)]
+    msrc += [f'void *buf_{i} = HAP_mmap(0,sz_or_val_{i},3,0,pra[{i+3}].dma.fd,0)+off{i};' for i,b in enumerate(bufs)
+             if isinstance(b[1][0].dtype, PtrDType)]
    msrc += ["unsigned long long start = HAP_perf_get_time_us();"]
-    msrc += [f"{function_name}({', '.join([(f'buf_{i}' if isinstance(b[1][0], PtrDType) else f'sz_or_val_{i}') for i,b in enumerate(bufs)])});"]
+    msrc += [f"{function_name}({', '.join([(f'buf_{i}' if isinstance(b[1][0].dtype, PtrDType) else f'sz_or_val_{i}') for i,b in enumerate(bufs)])});"]
    msrc += ["*(unsigned long long *)(pra[2].buf.pv) = HAP_perf_get_time_us() - start;"]
-    msrc += [f'HAP_munmap(buf_{i}, sz_or_val_{i});' for i,b in enumerate(bufs) if isinstance(b[1][0], PtrDType)]
+    msrc += [f'HAP_munmap(buf_{i}, sz_or_val_{i});' for i,b in enumerate(bufs) if isinstance(b[1][0].dtype, PtrDType)]
    msrc += ["return 0; }"]
    return '\n'.join(msrc)

@@ -273,22 +274,23 @@ return (void*)syscall((long)addr, length, prot, flags, fd, offset, 222); }}'''
 class MockDSPRenderer(DSPRenderer):
  def __init__(self, target:Target): self.target, self.compiler = target, DSPCompiler(mock=True)
  def _render_defines(self, uops) -> list[str]: return ClangRenderer._render_defines(self, uops)
-  def _render_entry(self, function_name:str, bufs:list[tuple[str,tuple[DType,bool]]]) -> str:
+  def _render_entry(self, function_name:str, bufs:list[tuple[str,tuple[UOp,bool]]]) -> str:
    # https://gpages.juszkiewicz.com.pl/syscalls-table/syscalls.html
    # control register 21 is HEX_REG_QEMU_INSN_CNT, 0x6a15c000 loads it
    msrc = [mockdsp_boilerplate, 'void _start(void) {']
    for i,b in enumerate(bufs):
-      if isinstance(b[1][0], PtrDType):
-        sz = b[1][0].size*b[1][0].itemsize
+      if isinstance(b[1][0].dtype, PtrDType):
+        sz = b[1][0].dtype.size*b[1][0].dtype.itemsize
        # for loop for big reads
        msrc.append(f"void *buf{i} = mmap2(0, {sz}, 3, 0x21, -1, 0); for(int rd = 0; rd < {sz}; rd += read(0, buf{i}+rd, {sz}-rd));")
      else:
        msrc.append(f"unsigned int val{i}; read(0, &val{i}, 4);")
    msrc.append("unsigned int st = inscount();")
-    msrc.append(f"{function_name}({', '.join([(f'(void*)buf{i}' if isinstance(b[1][0], PtrDType) else f'val{i}') for i,b in enumerate(bufs)])});")
+    params = [(f'(void*)buf{i}' if isinstance(b[1][0].dtype, PtrDType) else f'val{i}') for i,b in enumerate(bufs)]
+    msrc.append(f"{function_name}({', '.join(params)});")
    msrc.append("unsigned int et = inscount() - st; write(1, &et, sizeof(et));")
    for i,b in enumerate(bufs):
-      if isinstance(b[1][0], PtrDType): msrc.append(f"write(1, buf{i}, {b[1][0].size*b[1][0].itemsize});")
+      if isinstance(b[1][0].dtype, PtrDType): msrc.append(f"write(1, buf{i}, {b[1][0].dtype.size*b[1][0].dtype.itemsize});")
    msrc.append('exit(0); }')
    return '\n'.join(msrc)

@@ -41,105 +41,104 @@ def generic_wmma_helper(inp, warp_size, WARP_THREADS, K, NUM_A, NUM_B, NUM_C, a_

 class PythonProgram:
  def __init__(self, name:str, lib:bytes, **kwargs):
-    self.uops: list[tuple[Ops, DType, list[int], Any]] = pickle.loads(lib)
+    self.uops: list[UOp] = pickle.loads(lib)
+    self.uop_to_index: dict[UOp, int] = {u:i for i,u in enumerate(self.uops)}
+    self.loop_ends: dict[UOp, int] = {u.src[1]:i for i, u in enumerate(self.uops) if u.op == Ops.END}
  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
    st = time.perf_counter()
    warp = list(itertools.product(*[range(x) for x in local_size[::-1]]))
    warp_size = len(warp)
    void_ops = {Ops.END, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP, Ops.GROUP, Ops.STORE}
-    loop_ends: dict[int, int] = {srcs[1]:i for i, (uop, _, srcs, _) in enumerate(self.uops) if uop == Ops.END}
    for idxs in itertools.product(*[range(x) for x in global_size[::-1]]):
-      values: dict[int, Any] = {}
+      values: dict[UOp, Any] = {}
      pbufs: list[memoryview] = list(bufs)
      pvals: list[int] = list(vals)
      exec_masks = [[True] * warp_size]
      i = 0
      while i < len(self.uops):
-        uop, dtype, srcs, arg = self.uops[i]
-        src_values = [values[v] for v in srcs if self.uops[v][0] not in void_ops]
-        src_dtypes = [self.uops[v][1] for v in srcs if self.uops[v][0] not in void_ops]
-        if getenv("TRACE"): print(i, uop, dtype, arg, src_values, src_dtypes)
-        if uop is Ops.END:
-          i = srcs[1]
+        u = self.uops[i]
+        src_values = [values[v] for v in u.src if v.op not in void_ops]
+        src_dtypes = [v.dtype for v in u.src if v.op not in void_ops]
+        if getenv("TRACE"): print(i, u.op, u.dtype, u.arg, src_values, src_dtypes)
+        if u.op is Ops.END:
+          i = self.uop_to_index[u.src[1]]
          continue
-        if uop is Ops.IF:
+        if u.op is Ops.IF:
          exec_masks.append([x and y for x,y in zip(exec_masks[-1], src_values[0])])
          i += 1
          continue
-        if uop is Ops.ENDIF:
+        if u.op is Ops.ENDIF:
          exec_masks.pop()
          i += 1
          continue
-        if uop in (Ops.BARRIER, Ops.SINK, Ops.NOOP, Ops.GROUP):
+        if u.op in (Ops.BARRIER, Ops.SINK, Ops.NOOP, Ops.GROUP):
          # in the python emulator, the warp is always in sync
          i += 1
          continue
-        assert dtype is not None, f"{uop} is missing a dtype"
-        if uop is Ops.STORE:
+        assert u.dtype is not None, f"{u.op} is missing a dtype"
+        if u.op is Ops.STORE:
          assert len(src_values) == 2, f"STORE must be lowered to 2 srcs, got {len(src_values)}"
          store_gate = exec_masks[-1]
-          for j,val in enumerate(src_values[1] if src_dtypes[1].count > 1 else [src_values[1]]):
+          for j,val in enumerate(src_values[1] if u.max_numel() > 1 else [src_values[1]]):
            for (m,o),v,g in zip(src_values[0], val, store_gate):
              if g: _store(m, o+j, v, src_dtypes[1].scalar())
          i += 1
          continue
-        if uop is Ops.AFTER: values[i] = src_values[0]
-        elif uop in {Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}:
-          assert isinstance(dtype, PtrDType), dtype
-          storage_fmt = storage_fmt_for_dtype(dtype.base.scalar())
-          if storage_fmt is None: raise RuntimeError(f"{dtype=} is not supported")
+        if u.op is Ops.AFTER: values[u] = src_values[0]
+        elif u.op in {Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}:
+          storage_fmt = storage_fmt_for_dtype(u.dtype.base.scalar())
+          if storage_fmt is None: raise RuntimeError(f"dtype={u.dtype} is not supported")
          if TYPE_CHECKING or sys.version_info < (3, 12): assert storage_fmt != "e"
-          if uop is Ops.DEFINE_REG:
+          if u.op is Ops.DEFINE_REG:
            # REGs are per thread
-            values[i] = [memoryview(bytearray(dtype.size*dtype.itemsize)).cast(storage_fmt) for _ in range(warp_size)]
+            values[u] = [memoryview(bytearray(u.max_numel()*u.dtype.itemsize)).cast(storage_fmt) for _ in range(warp_size)]
          else:
-            buf = memoryview(bytearray(dtype.size*dtype.itemsize)) if uop is not Ops.PARAM else pbufs.pop(0)
-            values[i] = [buf.cast(storage_fmt)] * warp_size
-        elif uop is Ops.DEFINE_VAR:
-          values[i] = [pvals.pop(0)] * warp_size
-        elif uop is Ops.SPECIAL:
-          if arg[0] == 'g': values[i] = [idxs[2-int(arg[-1])]] * warp_size
-          elif arg[0] == 'l': values[i] = [x[2-int(arg[-1])] for x in warp]
-        elif uop is Ops.CONST: values[i] = [arg] * warp_size
-        elif uop is Ops.SLICE:
-          assert len(src_values) == 2, "non-image index must be 2 srcs"
+            buf = memoryview(bytearray(u.max_numel()*u.dtype.itemsize)) if u.op is not Ops.PARAM else pbufs.pop(0)
+            values[u] = [buf.cast(storage_fmt)] * warp_size
+        elif u.op is Ops.DEFINE_VAR:
+          values[u] = [pvals.pop(0)] * warp_size
+        elif u.op is Ops.SPECIAL:
+          if u.arg[0] == 'g': values[u] = [idxs[2-int(u.arg[-1])]] * warp_size
+          elif u.arg[0] == 'l': values[u] = [x[2-int(u.arg[-1])] for x in warp]
+        elif u.op is Ops.CONST: values[u] = [u.arg] * warp_size
+        elif u.op is Ops.INDEX:
          ret:list = []
-          for m,o in zip(*src_values): ret.append((m,o))
-          values[i] = ret
-        elif uop is Ops.INDEX:
-          assert isinstance(src_dtypes[0], ImageDType), "only image INDEX is supported"
-          ret = []
-          assert len(src_values) == 3, "image index must be 3 srcs"
-          for m,oy,ox in zip(*src_values):
-            if ox < 0 or ox >= src_dtypes[0].shape[1] or oy < 0 or oy >= src_dtypes[0].shape[0]: ret.append((m, None))
-            else: ret.append((m, ox*4 + oy*src_dtypes[0].shape[1]*4))
-          values[i] = ret
-        elif uop is Ops.CAST and isinstance(dtype, PtrDType):
-          values[i] = src_values[0]
-        elif uop is Ops.RANGE:
-          if i not in values: values[i] = [0] * warp_size
+          if isinstance(src_dtypes[0], ImageDType):
+            assert len(src_values) == 3, "image index must be 3 srcs"
+            for m,oy,ox in zip(*src_values):
+              if ox < 0 or ox >= src_dtypes[0].shape[1] or oy < 0 or oy >= src_dtypes[0].shape[0]: ret.append((m, None))
+              else: ret.append((m, ox*4 + oy*src_dtypes[0].shape[1]*4))
          else:
-            for j in range(len(values[i])):
-              values[i][j] += 1
-          if values[i][0] == src_values[0][0]:
-            del values[i]
-            i = loop_ends[i] + 1
+            assert len(src_values) == 2, "non-image index must be 2 srcs"
+            for m,o in zip(*src_values): ret.append((m,o))
+          values[u] = ret
+        elif u.op is Ops.CAST and isinstance(u.dtype, PtrDType):
+          values[u] = src_values[0]
+        elif u.op is Ops.RANGE:
+          if u not in values: values[u] = [0] * warp_size
+          else:
+            for j in range(len(values[u])):
+              values[u][j] += 1
+          if values[u][0] == src_values[0][0]:
+            del values[u]
+            i = self.loop_ends[u] + 1
            continue
-        elif uop is Ops.STACK: values[i] = src_values
-        elif uop is Ops.BITCAST: values[i] = [bitcast(x, src_dtypes[0], dtype) for x in src_values[0]]
-        elif uop is Ops.CAST:
-          values[i] = [truncate.get(dtype, lambda dt: dt)(dtype.const(x)) for x in src_values[0]]
-        elif uop is Ops.LOAD:
-          if dtype.count > 1:
-            values[i] = [load([src_values[i][j] if i != 0 and src_dtypes[i].count > 1 else src_values[i] \
-                               for i in range(len(src_values))], j, dtype.scalar()) for j in range(dtype.count)]
+        elif u.op is Ops.STACK: values[u] = src_values
+        elif u.op is Ops.BITCAST: values[u] = [bitcast(x, src_dtypes[0], u.dtype) for x in src_values[0]]
+        elif u.op is Ops.CAST:
+          values[u] = [truncate.get(u.dtype, lambda dt: dt)(u.dtype.const(x)) for x in src_values[0]]
+        elif u.op is Ops.LOAD:
+          if (load_sz := u.max_numel()) > 1:
+            # buf and gate are not vecs
+            values[u] = [load([src_values[k] if k in [0,2] else src_values[k][j] \
+                               for k in range(len(src_values))], j, u.dtype.scalar()) for j in range(load_sz)]
          else:
-            values[i] = load(src_values, 0, dtype)
-        elif uop is Ops.GEP: values[i] = src_values[0][get_single_element(arg)]
-        elif uop is Ops.WMMA:
-          first_src_dtype = self.uops[srcs[0]][1]
+            values[u] = load(src_values, 0, u.dtype)
+        elif u.op is Ops.GEP: values[u] = src_values[0][get_single_element(u.arg)]
+        elif u.op is Ops.WMMA:
+          first_src_dtype = u.src[0].dtype
          assert isinstance(first_src_dtype, DType) # mypy
-          dims, dtype_in, device, threads = arg[1], first_src_dtype.scalar(), arg[4], arg[5]
+          dims, dtype_in, device, threads = u.arg[1], first_src_dtype.scalar(), u.arg[4], u.arg[5]
          wmma_helper = functools.partial(generic_wmma_helper, src_values, warp_size)
          # TODO: refactor these to a shared TensorCoreLayout
          if device == "METAL":
@@ -147,17 +146,17 @@ class PythonProgram:
            def a_b_elem(x, i, j, goff): return x[(i%2)][goff+(i//2)%2+(j%4)*2+(i//4)*8+(j//4)*16]
            # (i, j), C, D (2 elements on 32 threads): row major same as A/B
            def c_map(lane, elem): return (elem + ((lane%2)*2) + ((lane//8)%2)*4, ((lane//2)%4) + (lane//16)*4)
-            values[i] = wmma_helper(32, 8, 2, 2, 2, a_b_elem, a_b_elem, c_map)
+            values[u] = wmma_helper(32, 8, 2, 2, 2, a_b_elem, a_b_elem, c_map)
          elif device == "AMD" and threads == 64:
            def a_elem(x, k, row, goff): return x[k%(dims[2]//4)][goff + (k//(dims[2]//4))*16 + row]
            def b_elem(x, col, k, goff): return a_elem(x, k, col, goff)  # pylint: disable=arguments-out-of-order
            def c_map(lane, elem): return (lane%16, (lane//16)*4 + elem)
-            values[i] = wmma_helper(64, dims[2], len(src_values[0]), len(src_values[1]), len(src_values[2]), a_elem, b_elem, c_map)
+            values[u] = wmma_helper(64, dims[2], len(src_values[0]), len(src_values[1]), len(src_values[2]), a_elem, b_elem, c_map)
          elif device == "AMD" and len(src_values[0]) == 8: # RDNA4
            def a_elem(x, k, row, goff): return x[k - [0, 4, 4, 8][k//4]][goff + row + [0, 16, 0, 16][k//4]]
            def b_elem(x, col, k, goff): return a_elem(x, k, col, goff)
            def c_map(lane, elem): return (lane%16, (lane//16)*8 + elem)
-            values[i] = wmma_helper(32, 16, 8, 8, 8, a_elem, b_elem, c_map)
+            values[u] = wmma_helper(32, 16, 8, 8, 8, a_elem, b_elem, c_map)
          elif device == "AMD":
            # A (16 elements on 32 threads): col major, lane 16-32 == lane 0-15
            def a_elem(x, k, row, goff):
@@ -166,7 +165,7 @@ class PythonProgram:
            # B (16 elements on 32 threads): row major, lane 16-32 == lane 0-15
            def b_elem(x, col, k, goff): return a_elem(x, k, col, goff)  # pylint: disable=arguments-out-of-order
            def c_map(lane, elem): return (lane%16, lane//16+elem*2) # (i, j), C, D (8 elements on 32 threads): row major
-            values[i] = wmma_helper(32, 16, 16, 16, 8, a_elem, b_elem, c_map)
+            values[u] = wmma_helper(32, 16, 16, 16, 8, a_elem, b_elem, c_map)
          elif device == "CUDA":
            # (col, row) given (lane, elem) for C & D (4 elements on 32 threads); shared by all tc shapes with M=16 N=8
            def c_map(lane, elem): return (elem%2 + (lane%4)*2, lane//4 + (elem//2)*8)
@@ -174,24 +173,24 @@ class PythonProgram:
            if dims == (8,16,16):
              def a_elem(x, k, row, goff): return x[k%2 + (row//8)*2 + (k//8)*4][goff + (k//2)%4 + (row%8)*4]
              def b_elem(x, col, k, goff): return x[k%2 + (k//8)*2][goff + (k//2)%4 + col*4]
-              values[i] = wmma_helper(32, 16, 8, 4, 4, a_elem, b_elem, c_map)
+              values[u] = wmma_helper(32, 16, 8, 4, 4, a_elem, b_elem, c_map)

            elif dims == (8,16,32):
              def a_elem(x, k, row, goff): return x[k%4 + (row//8)*4 + (k//16)*8][goff + (k//4)%4 + (row%8)*4]
              def b_elem(x, col, k, goff): return x[k%4 + (k//16)*4][goff + (k//4)%4  + col*4]
-              values[i] = wmma_helper(32, 32, 16, 8, 4, a_elem, b_elem, c_map)
+              values[u] = wmma_helper(32, 32, 16, 8, 4, a_elem, b_elem, c_map)

            elif dims == (8,16,8) and dtype_in == dtypes.half:
              def a_elem(x, k, row, goff): return x[k%2 + (row//8)*2][goff + k//2 + (row%8)*4]
              def b_elem(x, col, k, goff): return x[k%2][goff + k//2 + col*4]
-              values[i] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map)
+              values[u] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map)

            elif dims == (8,16,8) and dtype_in == dtypes.float:
              def a_elem(x, k, row, goff): return x[(k//4)*2 + row//8][goff + k%4 + (row%8)*4]
              def b_elem(x, col, k, goff): return x[k//4][goff + k%4 + col*4]
-              values[i] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map)
+              values[u] = wmma_helper(32, 8, 4, 2, 4, a_elem, b_elem, c_map)

-            else: raise NotImplementedError(f"unimplemented tensor core {arg}")
+            else: raise NotImplementedError(f"unimplemented tensor core {u.arg}")
          elif device == "INTEL":
            # A (16 elements on 8 threads)
            def a_elem(x, k, row, goff): return x[k%2+row*2][goff+k//2]
@@ -199,17 +198,17 @@ class PythonProgram:
            def b_elem(x, col, k, goff): return x[k][goff+col]
            # C, D (8 elements on 8 threads)
            def c_map(lane, elem): return (lane, elem)
-            values[i] = wmma_helper(8, 16, 16, 16, 8, a_elem, b_elem, c_map)
+            values[u] = wmma_helper(8, 16, 16, 16, 8, a_elem, b_elem, c_map)
          elif device == "CPU":
            def elem(x, col, row, _): return x[col+row][0] # k is always 0
            def c_map(lane, elem): return (elem%16, elem//16)
-            values[i] = wmma_helper(1, 1, 16, 16, 256, elem, elem, c_map)
-          else: raise NotImplementedError(f"unimplemented tensor core {arg}")
-        elif uop in GroupOp.ALU:
-          assert all_same([len(x) for x in src_values]), f"{[len(x) for x in src_values]} doesn't match on {uop}"
-          assert all_same([dtype] + src_dtypes) or uop in {*GroupOp.Comparison, Ops.WHERE}, f"dtype mismatch on {uop}"
-          values[i] = [exec_alu(uop, dtype, p) for p in zip(*src_values)]
-        assert i in values, (uop, dtype, srcs, arg)
+            values[u] = wmma_helper(1, 1, 16, 16, 256, elem, elem, c_map)
+          else: raise NotImplementedError(f"unimplemented tensor core {u.arg}")
+        elif u.op in GroupOp.ALU:
+          assert all_same([len(x) for x in src_values]), f"{[len(x) for x in src_values]} doesn't match on {u.op}"
+          assert all_same([u.dtype] + src_dtypes) or u.op in {*GroupOp.Comparison, Ops.WHERE}, f"dtype mismatch on {u.op}"
+          values[u] = [exec_alu(u.op, u.dtype, p) for p in zip(*src_values)]
+        assert u in values, u
        i += 1
    return time.perf_counter() - st

@@ -236,10 +235,7 @@ class PythonRenderer(Renderer):
    elif IMAGE and not target.arch: self.target = replace(target, arch="IMAGE_PITCH_ALIGNMENT=1")
    else: self.target = target

-  def render(self, uops:list[UOp]) -> str:
-    # the value of SPECIAL comes from local/global_size, not form its source
-    lops = [(u.op, u.dtype, [uops.index(v) for v in u.src if u.op is not Ops.SPECIAL], u.arg) for u in uops]
-    return base64.b64encode(pickle.dumps(lops)).decode()
+  def render(self, uops:list[UOp]) -> str: return base64.b64encode(pickle.dumps(uops)).decode()

  def supported_dtypes(self): return {d for d in super().supported_dtypes() if d != dtypes.half or sys.version_info >= (3, 12)}

--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	c7b6ee0c7d	dt.count	2026-06-02 13:19:04 -07:00
geohot	13f5d39fcf	count	2026-06-02 13:16:42 -07:00
geohot	431accc9b7	bitsize in nir	2026-06-02 13:11:31 -07:00
geohot	df000116ea	more renderer cleanups	2026-06-02 13:02:28 -07:00
nimlgenandGitHub	99e37b1ee3	hcq2: deps (#16459 ) * start * sin * f	2026-06-02 22:34:25 +03:00
George HotzandGitHub	82f1c983d4	clean renderer migrations [pr] (#16472 ) * clean renderer migrations * minor webgpu * use PARAM UOp as API * make linter happy	2026-06-02 11:19:00 -07:00
sirhcmandGitHub	9897658895	ci: fix ocelot compilation on macos (#16471 )	2026-06-02 12:43:31 -04:00
chenyuandGitHub	6b7d2b91df	update test_uop_graph (#16470 ) use UOp methods instead of constructing UOp directly, some of it violated spec	2026-06-02 08:53:54 -04:00
qazalandGitHub	854eac09c6	llama: no E_ copy after bf16 GEMM (#16458 )	2026-06-02 14:14:13 +09:00
George HotzandGitHub	7d8ed8d4d7	add store to buffer's addrspace (#16468 )	2026-06-01 22:07:43 -07:00
George HotzandGitHub	20242fdf1d	update test + spec from shrink_in_render (#16467 ) * update test + spec from shrink_in_render * cast	2026-06-01 19:24:43 -07:00
sirhcmandGitHub	c6cad1ad67	ci: standardize runs-on (#16466 ) * ci: use macos 26 * ugh github * stick with github for arm	2026-06-01 21:39:58 -04:00
sirhcmandGitHub	b0ecbb34d9	ci: cleanup python backend tests (#16465 )	2026-06-01 20:08:05 -04:00
sirhcmandGitHub	2d0f132a3b	ci: cleanup more duplicate tests (#16462 )	2026-06-01 18:56:29 -04:00
wozeparrotandGitHub	aab9a5a8a3	llama: allow specifying layer count (#16464 )	2026-06-01 15:36:04 -07:00
chenyuandGitHub	0167401fa2	minor hcopt WHERE cleanup [PR] (#16463 )	2026-06-01 17:58:38 -04:00
George HotzandGitHub	124d2f8227	anon addrspace from new renderer (#16461 ) * anon addrspace from new renderer * use max_numel in python renderer * add sizes to ptrs in tests * more * correct fix	2026-06-01 14:42:02 -07:00
chenyuandGitHub	517eea5985	no CONST(DEVICE) in create_allreduce_function (#16460 )	2026-06-01 17:12:34 -04:00
chenyuandGitHub	7e7b481ba7	less CONST(DEVICE) (#16452 ) * less CONST(DEVICE) no DEVICE for single device in const_like, multi has other issues * maybe * that?	2026-06-01 15:55:12 -04:00
George HotzandGitHub	556defa0f7	minor updates from vec removal (#16456 )	2026-05-31 09:48:51 -07:00
Javier De JesusandGitHub	989f713c1b	support negative pads in circular pad mode (#16448 )	2026-05-31 09:28:45 -07:00
nimlgenandGitHub	2c2cb339e0	fix word wrap (#16450 )	2026-05-30 23:21:24 +03:00
qazalandGitHub	29b47a0057	llama: update local amax implementation after ParamArgs change (#16446 ) * local amax failing test * update _local_abs_max_fxn	2026-05-30 16:55:43 +09:00
wozeparrotandGitHub	6795c2d5c9	llama: zero grad this way (#16445 )	2026-05-29 20:25:21 -07:00
George HotzandGitHub	cf55aaf01f	python prg is pkl uops (#16443 ) * python prg is pkl uops * refactor to use uop * refactor to u.	2026-05-29 19:13:51 -07:00
sirhcmandGitHub	c377d01491	ci: run dsp on tinygrad[testing] (#16442 )	2026-05-29 21:16:56 -04:00
wozeparrotandGitHub	c23652e486	llama: minimize peak init mem (#16440 )	2026-05-29 18:00:37 -07:00
sirhcmandGitHub	d943493b79	ci: remove duplicate op compile test (#16441 )	2026-05-29 19:20:31 -04:00
chenyuandGitHub	8ac62b28e5	fix AffineGrid fusion (#16439 )	2026-05-29 17:59:47 -04:00
sirhcmandGitHub	ef50a49693	ci: macos dev matrix (#16436 )	2026-05-29 17:40:32 -04:00
sirhcmandGitHub	434cfa96a3	ci: no fetch in backend tests (#16438 ) should make for less actions cache thrashing	2026-05-29 17:11:16 -04:00
chenyuandGitHub	b7280705a7	limit CONST(UNIQUE) to invalids only (#16432 )	2026-05-29 16:02:06 -04:00
George HotzandGitHub	9506b78d73	fix viz addrspace (#16437 ) * fix viz addrspace * revert that	2026-05-29 12:58:05 -07:00
nimlgenandGitHub	d69aca41a9	hcq2: rework pm_bufferize (#16431 )	2026-05-29 22:09:52 +03:00
George HotzandGitHub	e2a0434403	full derivation of addrspace (#16433 ) * full derivation of addrspace * w/e, it fixes it	2026-05-29 11:39:31 -07:00
wozeparrotandGitHub	6787de9f52	llama: fix mp (#16434 )	2026-05-29 11:21:43 -07:00
chenyuandGitHub	2d7e5baab4	remove vec= from UPat.cvar [PR] (#16430 )	2026-05-29 10:52:30 -04:00
chenyuandGitHub	fa666cefe8	remove dead branch in UOp [PR] (#16429 )	2026-05-29 10:38:49 -04:00
qazalandGitHub	81bc00c006	do not require clearing method_cache in viz tests (#16428 ) * update * update test_dedup	2026-05-29 18:12:34 +09:00
qazalandGitHub	54cfb794b8	viz: addrspace little colored box (#16427 ) * return addrspace * layout * render * addrspace encodes color * update colors * in input_ast all are params are green * update stroke	2026-05-29 17:25:07 +09:00
qazalandGitHub	814d414f41	viz: set label offset for asm (#16426 )	2026-05-29 13:16:34 +09:00
wozeparrotandGitHub	f86966af56	llama: optim amax margin (#16425 )	2026-05-28 20:18:11 -07:00
sirhcmandGitHub	6e0d5262dc	ci: autocancel outdated pr jobs (#16424 )	2026-05-28 23:14:35 -04:00
sirhcmandGitHub	69aa2054f6	rename clangjit to clang (#16423 )	2026-05-28 22:41:58 -04:00
sirhcmandGitHub	a909acb882	move llvmspeed to benchmarks (#16422 )	2026-05-28 22:26:22 -04:00
George HotzandGitHub	1e7f1dcf49	add ParamArgs [pr] (#16421 ) * add ParamArgs * fix export * cleanups * fixes * simpler	2026-05-28 19:17:17 -07:00
sirhcmandGitHub	7d38edffdb	ci: dev matrix (#16420 ) windows just runs test_tiny	2026-05-28 22:04:04 -04:00
wozeparrotandGitHub	36c8ff70c1	llama: use old scale for dequant in optim (#16417 )	2026-05-28 15:21:19 -07:00
George Hotz GitHub sirhcm	c87f3433d1	use namespace runners (#16387 ) Co-authored-by: Christopher Milan <[email protected]>	2026-05-28 18:05:46 -04:00
George HotzandGitHub	c9adde72c1	addrspace property (#16418 ) * addrspace property * movement addrspace * regs	2026-05-28 14:39:25 -07:00
sirhcmandGitHub	c8af163d2b	disable process replay by default (#16419 ) enable process replay with [pr] and assert with [PR] process replay no longer captures on master	2026-05-28 17:36:28 -04:00
nimlgenandGitHub	b0e49afaf1	hcq2: new multi (#16413 ) * hcq2: new multi * op	2026-05-28 22:16:10 +03:00
George HotzandGitHub	edca5df25a	flip offset and shape in pad and shrink (#16414 ) * flip offset and shape in pad and shrink * dumb test	2026-05-28 11:58:19 -07:00
chenyuandGitHub	d72d8ee065	.const() should not ignore dtype (#16412 ) fixed a bug in postrange, also cleaner	2026-05-28 10:49:15 -04:00
sirhcmandGitHub	0ae957bb0a	refactor webgpu (#16406 )	2026-05-27 23:13:08 -04:00
qazalandGitHub	202adc644e	viz: make call toggle easier to click on (#16411 ) * call tag is a rect * details * colors * simplify, better comment	2026-05-28 11:53:36 +09:00
George HotzandGitHub	5ee6b6b79e	fix slice store to remove the index (#16410 ) * fix slice store to remove the index * fix spec	2026-05-27 19:17:53 -07:00
qazalandGitHub	88e88d63d6	viz: click on +- toggles sources (#16409 )	2026-05-28 09:12:43 +09:00
George HotzandGitHub	b21afb4883	marg line cleanup (#16408 ) * marg line cleanup * bitcast is a mop	2026-05-27 16:41:04 -07:00
wozeparrotandGitHub	dac3743d75	llama: delayed scaling in optim (#16407 )	2026-05-27 15:40:03 -07:00
George HotzandGitHub	8ee3a37524	shrink/pad use (new_shape, offset) (#16405 ) * shrink uses offset and shape * pad does too * fix	2026-05-27 15:13:08 -07:00
sirhcmandGitHub	171401e8df	skip modulo by zero in test_dtype_alu (#16404 )	2026-05-27 17:09:05 -04:00
qazalandGitHub	452c7d4230	llama: don't allocate grad_xw13 in bf16 (#16359 )	2026-05-28 04:33:07 +09:00
nimlgenandGitHub	0c385e31c6	hcq2 rewrite (#16375 ) * hcq2 rewrite * fi * x * simpler	2026-05-27 22:25:35 +03:00
chenyuandGitHub	c33b767407	bring back test and torch backend change for unique const (#16403 )	2026-05-27 15:16:08 -04:00
sirhcmandGitHub	bacabf0866	webgpu: fix enums (#16402 )	2026-05-27 13:09:50 -04:00
chenyuandGitHub	6da785562b	test_custom_kernel_precompile_multidevice (#16401 ) add a test to show what invalids need	2026-05-27 11:19:16 -04:00
chenyuandGitHub	3e80f375ee	skip test_setitem_fancy_on_unrealized_view (#16400 ) crashes in linux llvm ci	2026-05-27 09:50:26 -04:00
chenyuandGitHub	945ed4f689	revert const unique changes (#16395 )	2026-05-27 00:06:41 -04:00
sirhcmandGitHub	aacc8addf4	ci: use ubuntu 24.04 (#16393 )	2026-05-26 23:22:01 -04:00
chenyuandGitHub	fa14cde05c	test update for arange and eye (#16394 ) these will need explicit clone to make a buffer	2026-05-26 22:48:34 -04:00
wozeparrotandGitHub	3a7a6da7d5	llama: fakedata uses real vocab size (#16389 )	2026-05-26 18:58:55 -07:00