fixes

fix after merge
Merge branch 'master' into shrink_in_render
2026-08-15 15:18:26 +00:00 · 2026-06-02 12:53:00 -07:00 · 2026-06-02 12:45:51 -07:00 · 2026-06-02 11:20:24 -07:00 · 2026-06-02 11:19:00 -07:00 · 2026-06-02 12:43:31 -04:00
105 changed files with 1404 additions and 1225 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'
@@ -303,7 +308,6 @@ runs:

        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
          CMAKE_ARGS="$CMAKE_ARGS -DLLVM_DIR=$(llvm-config-15 --cmakedir)"
@@ -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-24.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-24.04
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -125,7 +111,7 @@ jobs:

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

  bepython:
    name: Python Backend
-    runs-on: ubuntu-24.04
+    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-24.04
+    runs-on: *linux
    timeout-minutes: 10

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

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

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

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

    steps:
@@ -320,7 +306,7 @@ jobs:
      matrix:
        group: [1, 2]
    name: SPEC=2 (${{ matrix.group }})
-    runs-on: ubuntu-24.04
+    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-24.04
+    runs-on: *linux
    timeout-minutes: 10
    steps:
    - name: Checkout Code
@@ -357,7 +344,7 @@ jobs:

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

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

  testopenpilot:
    name: openpilot Compile Tests
-    runs-on: ubuntu-24.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-24.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-24.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-24.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-24.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-24.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-24.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-24.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")
@@ -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)
@@ -276,9 +280,9 @@ def apply_grad(grad_buf:Tensor, new_grad:UOp):
    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
    return
  cur = grad_buf.uop
-  for pad in sorted(pads, key=lambda p: p.marg[0][1] if p.op == Ops.PAD else 0, reverse=True):
+  for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
    if pad.op == Ops.PAD:
-      grad_shrink = tuple([(p[1], s+p[1]) for s,p in zip(pad.src[0].shape, pad.marg)])
+      grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
      buf_slice = cur.shrink(grad_shrink)
      cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
    else:
@@ -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]
@@ -89,13 +91,15 @@ class GradAccClipAdamW(Optimizer):
    if t.dtype in dtypes.fp8s:
      from examples.mlperf.models.flat_llama import FP8_MAX
      # delayed scaling: reuse previous step's inv_scale
-      scale = t._inv_scale.reciprocal().reshape(-1, *([1]*(new_w.ndim-1)))
+      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))) * t._inv_scale).detach()
-      inv = ((new_amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
-      t._inv_scale.assign(inv.shard_like(t._inv_scale) if offloaded else inv)
+      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
@@ -5,7 +5,7 @@ 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, 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
@@ -17,14 +17,18 @@ from tinygrad.engine.realize import to_program, track_stats, get_call_arg_uops,
 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, can_recover:bool=False, arch=None):
    self.device_id:int = int(device.split(":")[1]) if ":" in device else 0

+    # 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))),
+    ])
+
    super().__init__(device, allocator, compilers, lambda *a, **kw: None, None, arch=arch)

  @functools.cached_property
@@ -41,14 +45,6 @@ class HCQ2Compiled(Compiled):
    buf.as_memoryview(force_zero_copy=True).cast('Q')[0] = 1
    return buf

-  @functools.cached_property
-  def pm_bufferize(self) -> PatternMatcher:
-    return 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))),
-    ])
-
  def synchronize(self, timeout:int|None=None):
    if not hasattr(self, 'iface'): return
    sig = self.timeline_signal._buf.cpu_view().mv.cast('Q')
@@ -59,14 +55,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):
@@ -167,15 +155,24 @@ class HCQEncoder:
    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)

+# *****************
+# 0. helpers
+
+HCQ_DEVS = frozenset(("AMD",))
+HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",))
+
+def to_tuple(d): return d if isinstance(d, tuple) else (d,)
+
+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

-HCQ_DEVS = ("AMD",)
-HCQ_P2P_DEVS = HCQ_DEVS + ("CPU",)
+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:
-  dst_dev, src_dev = dst.device.split(":")[0], src.device.split(":")[0]
-  if not ((src_dev in HCQ_DEVS and dst_dev not in HCQ_P2P_DEVS) or (dst_dev in HCQ_DEVS and src_dev not in HCQ_P2P_DEVS)): return 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)])
@@ -192,7 +189,8 @@ def get_pm_prep_program(name:str) -> PatternMatcher|None:

 def prep_program(call:UOp, prg:UOp) -> UOp|None:
  dev = call.src[1].device
-  if (pm:=get_pm_prep_program(dev.split(":")[0])) is None or (lowered:=pm.rewrite(prg)) is None: return None
+  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)
@@ -276,7 +274,7 @@ def get_pm_lower(name:str) -> PatternMatcher|None:
  except ImportError: return None

 def encode_cmdbuf(call:UOp, q:UOp) -> UOp|None:
-  if (pm:=get_pm_lower(call.src[1].device.split(":")[0])) is None or (encoded:=pm.rewrite(q)) is None: return None
+  if (pm:=get_pm_lower(to_tuple(q.arg[0])[0].split(":")[0])) is None or (encoded:=pm.rewrite(q)) is None: return None
  return call.replace(src=(encoded,) + call.src[1:])
 pm_encode_cmdbufs = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), encode_cmdbuf)])

@@ -304,40 +302,52 @@ pm_lift_patches_to_cmdbuf = PatternMatcher([

 def bufferize_buf(buf:UOp) -> UOp|None:
  if buf.tag is None: return None
-  d = buf.src[1].arg
-  uops = tuple(UOp.from_buffer((d:=Device[dev]).pm_bufferize.rewrite(buf, ctx=d), dev) for dev in (d if isinstance(d, tuple) else (d,)))
+  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 uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, uops)
 pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, name="buf"), bufferize_buf)])

 # *****************
-# 5.1. capture buffers reachable from each hcq call as BIND, so resolve_patches doesn't drop their refs
+# 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))): return 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:
-  buf.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B')[:len(blob.arg)] = blob.arg
+  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:
-  struct.pack_into(f'<{val.dtype.fmt}', buf.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * buf.dtype.base.itemsize, val.arg)
+  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)

-pm_resolve_patches = symbolic + PatternMatcher([
+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))
+
+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),
+
+  # 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, name="buf"), UPat(Ops.DEVICE)), name="g"),
-    lambda buf, g: UOp.const(dtypes.uint64, buf.buffer.get_buf(g.src[1].arg).va_addr)),
-  (UPat(Ops.GETADDR, src=(UPat.cvar("const"), UPat())), lambda const: const),
+  (UPat(Ops.GETADDR, src=(UPat({Ops.BUFFER, Ops.MSTACK, Ops.MSELECT}, name="buf"), UPat(Ops.DEVICE)), name="g"), resolve_getaddr),

-  (UPat(Ops.BUFFER, name="buf").store(UPat(Ops.BINARY, name="blob")), fold_blob_store),
-  (UPat(Ops.BUFFER, name="buf").index(UPat.cvar("off")).or_casted().store(UPat.cvar("val")), fold_const_store),
-])
+  # 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

 # *****************
 # 6. callify hcq programs
@@ -346,12 +356,20 @@ pm_fixup = PatternMatcher([ # TODO: this should gone?
  (UPat(Ops.CONST, name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
 ])

+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))])
+
+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)])
+
 def callify_hcq(call:UOp) -> UOp:
-  bufs = dedup([u for u in call.src[0].toposort() if u.op is Ops.BUFFER])
-  subs = {b: UOp.placeholder((b.arg,), b.dtype, i) for i, b in enumerate(bufs)}
-  sink = UOp.sink(call.src[0].substitute(subs), arg=KernelInfo(name="hcq_submit", estimates=Estimates()), tag=1)
-  return to_program(sink, Device["CPU"].renderer).call(*bufs, UOp(Ops.BIND, dtypes.void, src=call.src[1:]))
-pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), callify_hcq)])
+  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)])

@track_rewrites(lambda _,ret: f"HCQ Schedule {pluralize('Kernel', len(ret.src))}")
 def hcq_schedule(linear:UOp) -> UOp:
@@ -371,6 +389,7 @@ def hcq_schedule(linear:UOp) -> UOp:
  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_callify_hcq, name="parametrize host buffers")
+  linear = graph_rewrite(linear, pm_parametrize_host_buffers, name="parametrize host buffers")
+  linear = graph_rewrite(linear, pm_callify_hcq, name="callify hcq")

  return linear
@@ -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
@@ -220,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(("SDMA: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
@@ -249,7 +248,7 @@ 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]] = {}
@@ -271,6 +270,7 @@ 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,
@@ -423,13 +423,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, 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 +456,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 +471,29 @@ 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"{'SDMA' 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"]
+    ]) + 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,50 +502,50 @@ 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()

-  @functools.cached_property
-  def pm_bufferize(self) -> PatternMatcher:
-    return PatternMatcher([
-      (UPat(Ops.BUFFER, tag="scratch"),
-        lambda ctx: Buffer(ctx.device, ctx.scratch.size, dtypes.uint8, opaque=ctx.scratch, options=BufferSpec(external_ptr=1))),
-      (UPat(Ops.BUFFER, tag={("compute_queue", n) for n in ("ring", "write_ptr", "doorbell", "put_value")}, name="b"),
-        lambda ctx, b: getattr(ctx.compute_queue, b.tag[1])),
-      (UPat(Ops.BUFFER, tag={("sdma_queue", n) for n in ("ring", "write_ptr", "doorbell", "put_value")}, name="b"),
-        lambda ctx, b: getattr(ctx.sdma_queue(0), b.tag[1])),
-    ]) + super().pm_bufferize
-
  def device_props(self): return self.iface.props
@@ -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:
@@ -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()
@@ -23,12 +23,14 @@ def calculate_storage_offset(x: Tensor) -> int:
  for u in x.uop.toposort():
    if u.op == Ops.SHRINK:
      u_strides = strides_for_shape(u.src[0].shape)
-      for i, (_, start) in enumerate(u.marg): offset += start * u_strides[i]
+      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,8 +74,8 @@ 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{s'}, \mathbf{o})$ & --- & Place $T$ at offset $o_k$ in a zero-filled output of shape $s'_k$. \\
-\op{Shrink} & $(T, \mathbf{s'}, \mathbf{o})$ & --- & Keep $s'_k$ elements starting at offset $o_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$. \\
@@ -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()
@@ -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()
@@ -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()
@@ -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)],
@@ -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,))
@@ -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])
@@ -323,7 +323,7 @@ class TestWithGrad(unittest.TestCase):

  def test_set_overlapping_backward(self):
    z = Tensor.zeros(6)
-    x = Tensor.ones(4)
+    x = Tensor.ones(4).contiguous()
    y = Tensor.ones(4) * 2
    z[:4] = x
    z[2:] = y
@@ -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]
@@ -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())
@@ -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)
@@ -456,14 +456,14 @@ class TestSchedule(unittest.TestCase):
      check_schedule(out, 4, [c1.weight, c1.bias, *nn.state.get_parameters(bn)])

  def test_fold_conv_batchnorm_optim(self, adam=False):
-    # 2 is too low?
-    optim, cnt = (nn.optim.Adam, 16) if adam else (nn.optim.SGD, 2)
+    optim, cnt = (nn.optim.Adam, 29) if adam else (nn.optim.SGD, 15)
    with Tensor.train():
      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(img, *nn.state.get_parameters(opt))
      img_bn = bn(c1(img)).elu().sum()
      opt.zero_grad()
      img_bn.backward()
@@ -477,7 +477,7 @@ class TestSchedule(unittest.TestCase):
      fw = bn(x).contiguous_backward().relu().contiguous()
      fw.sum().backward()
      # TODO: this is too many
-      check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 9)
+      check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 10, nn.state.get_parameters(bn))

  def test_fold_conv_relu(self):
    c1 = nn.Conv2d(3,16,3)
@@ -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)

@@ -78,7 +78,7 @@ class TestIdxUpcast(unittest.TestCase):
    if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)):
      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):
@@ -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,46 +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.index(idx, ptr=True), 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
-    self.assertEqual(ld0.src[0], smem.after(barrier).index(lidx+2, ptr=True))
+    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.index(idx1, ptr=True).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))])
@@ -779,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))
@@ -797,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()
@@ -15,6 +15,10 @@ class TestRealizeIsRealized(unittest.TestCase):
    t = Tensor.zeros(10).contiguous().realize()
    assert t.uop.is_realized

+  def test_ones(self):
+    t = Tensor.ones(4, 4).realize()
+    assert t.uop.is_realized
+
  def test_bytes(self):
    t = Tensor(b'\x01\x02\x03').realize()
    assert t.uop.is_realized
@@ -51,10 +55,6 @@ class TestRealizeIsRealized(unittest.TestCase):
    t = Tensor(3.14).realize()
    assert not t.uop.is_realized

-  def test_ones_not_realized(self):
-    t = Tensor.ones(4, 4).realize()
-    assert not t.uop.is_realized
-
  def test_none_not_realized(self):
    t = Tensor(None).realize()
    assert not t.uop.is_realized
@@ -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,7 +176,8 @@ 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),
@@ -3,16 +3,16 @@ from dataclasses import replace
 import itertools
 from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
 from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic
-from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo
+from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext
-from tinygrad.dtype import dtypes
+from tinygrad.dtype import dtypes, PtrDType, ImageDType

 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
-from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load
+from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load, pm_clean_up_group_sink
 from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_transcendental_patterns, pm_dtype_decomps
 from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize_buf_and_index, devectorize_alu, pm_reduce, \
@@ -24,6 +24,25 @@ from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, p
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite

+# NOTE: this is temporary until we fix the devectorizer
+pm_index_is_shrink = PatternMatcher([
+  # rewrite non-image INDEX to SHRINK
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).cast(name="x"), lambda buf,idx,x:
+    UOp(Ops.SHRINK, dtype=buf.dtype.base, src=(buf, idx, UOp.const(dtypes.int, x.dtype.count))) if isinstance(buf.dtype, PtrDType) else None),
+  # rewrite GEP to INDEX
+  (UPat(Ops.GEP, name="x"), lambda x: x.replace(op=Ops.INDEX, src=x.src+(UOp.const(dtypes.int, x.arg),), arg=None)),
+])
+
+pm_remove_vec_dtypes = PatternMatcher([
+  # rewrite PARAM to non pointer
+  (UPat((Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf"), lambda buf:
+   buf.replace(dtype=buf.dtype.base, src=(UOp.const(dtypes.int, buf.ptrdtype.size),)) \
+    if isinstance(buf.dtype, PtrDType) and not isinstance(buf.dtype, ImageDType) else None),
+  # remove all vec dtypes
+  (UPat(GroupOp.All-{Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}, name="x"),
+   lambda x: x.replace(dtype=x.dtype.base.scalar().base)),
+])+pm_clean_up_group_sink
+
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
  if DEBUG >= 5: print(pyrender(ast))
@@ -100,6 +119,8 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
  pm_final_rewrite = pm_decomp+pm_render+extra_matcher+pm_split_ends
  sink = graph_rewrite(sink, pm_final_rewrite, ctx=ren, name="final rewrite")
+  sink = graph_rewrite(sink, pm_index_is_shrink, name="index is shrink")
+  sink = graph_rewrite(sink, pm_remove_vec_dtypes, name="remove vec dtypes")

  # this was the linearizer
  sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True)
@@ -115,7 +136,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(Ops.INDEX).or_casted(), 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
@@ -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,7 +51,7 @@ 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):
@@ -94,7 +94,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)
@@ -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]]]:
@@ -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)))
@@ -72,8 +73,8 @@ pm_gradient = PatternMatcher([
  (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret:
    (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[1], 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[1], s-p[0]-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, 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[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)

@@ -11,7 +11,7 @@ from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, InvalidType, Pt
 from tinygrad.helpers import all_int, argfix, ceildiv, flatten, flat_to_grouped, make_tuple, prod, resolve_pool_pads, round_up

 if TYPE_CHECKING:
-  from tinygrad.uop.ops import sint
+  from tinygrad.uop.ops import sint, UOp

 ReductionStr = Literal["mean", "sum", "none"]

@@ -23,12 +23,13 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
  def const(dtype, b, device=None): raise NotImplementedError("creation helpers are only supported on Tensor and UOp")

  @classmethod
-  def full(cls, shape:tuple[sint, ...], fill_value:ConstType, **kwargs) -> Self:
+  def full(cls, shape:tuple[sint, ...], fill_value:ConstType|UOp, dtype:DTypeLike|None=None,
+           device:str|tuple[str, ...]|None=None, buffer=True) -> Self:
    """
    Creates a tensor with the given shape, filled with the given value.

    You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
-    Additionally, all other keyword arguments are passed to the constructor of the tensor.
+    Pass `buffer=False` to get a broadcast const value instead of a materialized buffer.

    ```python exec="true" source="above" session="tensor" result="python"
    print(Tensor.full((2, 3), 42).numpy())
@@ -37,11 +38,13 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
    print(Tensor.full((2, 3), False).numpy())
    ```
    """
+    from tinygrad.uop.ops import UOp
    new_shape = argfix(shape)
-    if not kwargs.pop("buffer", True):
-      dt = to_dtype(kwargs.pop("dtype", None) or dtypes.from_py(fill_value))
-      return cls.const(dt, fill_value, canonicalize_device(kwargs.pop("device", None))).reshape((1,)*len(new_shape)).expand(new_shape)
-    return cls.unique_const(fill_value, **kwargs).reshape((1,)*len(new_shape)).expand(new_shape)
+    dt = to_dtype(dtype) if dtype is not None else None
+    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:
@@ -52,7 +55,8 @@ class OpMixin(ElementwiseMixin, ReduceMixin):

    Eventually Tensor.empty will be replaced by this.
    """
-    return cls.full(argfix(*shape), Invalid, **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:
@@ -253,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((s+x[0]+x[1], x[0]) if x is not None else (s, 0) for x, s in zip(arg, self.shape)))
+    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[1]-x[0], x[0]) if x is not None else (s, 0) 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((s if ns is None else ns, 0) 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

@@ -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 INDEX, 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
+          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()))
@@ -8,9 +8,17 @@ from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, trunc
 from tinygrad.renderer import Renderer
 from tinygrad.codegen.late.devectorizer import no_vectorized_alu

+def render_index(ctx,buf,idx):
+  base = buf
+  while base.op is Ops.AFTER: base = base.src[0]
+  if base.addrspace == AddrSpace.ANON:
+    assert idx.op is Ops.CONST, f"{idx.op} must be CONST"
+    return f"{ctx[buf]}[{idx.arg}]"
+  else:
+    return f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"

 base_rewrite = PatternMatcher([
-  (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),
+  (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{ctx[x.src[0]]}];"),
  (UPat(Ops.IF, name="x"), lambda ctx,x: f"if ({ctx[x.src[0]]}) {{"),
  (UPat((Ops.ENDIF, Ops.END)), lambda ctx: "}"),
  (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]})"),
@@ -18,14 +26,14 @@ base_rewrite = PatternMatcher([
  (UPat(Ops.RANGE, name="x"),
   lambda ctx,x: f"for ({ctx.render_dtype(x.dtype)} {ctx[x]} = 0; {ctx[x]} < {ctx[x.src[0]]}; {ctx[x]}++) {{"),
  (UPat(Ops.STACK, name="x"),
-   lambda ctx,x: f"{ctx.float4.replace('float4', ctx.render_dtype(x.dtype))}" + \
-    f"{ctx.float4_style[0]}{','.join([ctx[y] for y in x.src])}{ctx.float4_style[1]}"),
+   lambda ctx,x: f"{ctx.float4.replace('float4', ctx.render_dtype(ctx.render_dtype_with_shape(x)))}" + \
+     f"{ctx.float4_style[0]}{','.join([ctx[y] for y in x.src])}{ctx.float4_style[1]}"),
  (UPat(Ops.CAST, name="x"), lambda ctx,x:
    f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_dtype(x.dtype)})" if x.dtype.count > 1 and not isinstance(x.dtype, PtrDType) else None),
  (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({ctx.render_cast(x.dtype, ctx[x.src[0]])})"),
  (UPat(Ops.BITCAST, name="x"), lambda ctx,x:
    f"__builtin_bit_cast({ctx.render_dtype(x.dtype)}, ({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
-  (UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x: f"{ctx.smem_align}{ctx.smem_prefix}{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),
+  (UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x: f"{ctx.smem_align}{ctx.smem_prefix}{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{ctx[x.src[0]]}];"),
  (UPat(Ops.BARRIER), lambda ctx: ctx.barrier),
  (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"{ctx.code_for_workitem[x.arg[0]](x.arg[-1])}; /* {(x.src[0]).render()} */"),
  # const
@@ -43,24 +51,26 @@ 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)),
+  # SHRINK/INDEX
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx'))), render_index),
+  (UPat(Ops.SHRINK, src=(UPat.var("buf"), UPat.var('idx'), UPat.cvar())), render_index),
  # 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]};"),
+  (UPat(Ops.LOAD, src=(UPat.var('bidx'),), name="x"), lambda ctx,x,bidx: ctx.render_access(bidx, ctx.render_dtype_with_shape(x))),
+  (UPat(Ops.LOAD, src=(UPat.var("bidx"), UPat.var("var"), UPat.var("gate")), name="x"),
+   lambda ctx,x,bidx,var,gate: f"({ctx[gate]}?{ctx.render_access(bidx, ctx.render_dtype_with_shape(x))}:{ctx[var]})"),
+  (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))),
+   lambda ctx,bidx,var: f"{ctx.render_access(bidx, ctx.render_dtype_with_shape(var))} = {ctx[var]};"),
  # alu/gep
  # TODO: look for left-associative
  (UPat(GroupOp.ALU, name="x"), lambda ctx,x: ctx.code_for_op[x.op](
    *([strip_parens(ctx[v]) if v.op == x.op and x.op in {Ops.ADD, Ops.MUL, Ops.XOR, Ops.OR, Ops.AND} else ctx[v] for v in x.src]), x.dtype)),
-  (UPat(Ops.GEP, name="x"), lambda ctx,x: ctx[x.src[0]] + \
-    (f"[{x.arg[0]}]" if x.src[0].dtype.count > ctx.gep_arr_threshold else f".{'xyzwabcd'[x.arg[0]]}")),
  # custom passes through with format
  (UPat((Ops.CUSTOM, Ops.CUSTOMI), name="x"), lambda ctx,x: x.arg.format(*[ctx[y] for y in x.src])),
 ])

 extra_pm = PatternMatcher([
  # devectorize any bools
-  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.INDEX), dtype=dtypes.bool, name="alu"), no_vectorized_alu),
+  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.SHRINK), dtype=dtypes.bool, name="alu"), no_vectorized_alu),
  # CAST (from bool) can't be vectorized
  (UPat(Ops.CAST, src=(UPat(dtype=dtypes.bool),), name="alu"), no_vectorized_alu),
  # WHERE can't be vectorized
@@ -95,7 +105,11 @@ pm_manual_bf16_cast = PatternMatcher([
  (UPat(Ops.CAST, dtype=dtypes.bfloat16, src=(UPat.var("x", dtype=dtypes.float),)), cast_float_to_bf16),
 ])

-def uops_to_dtypes(uops:list[UOp]) -> list[DType]: return dedup(u.dtype for u in uops if not isinstance(u.dtype, (ImageDType, PtrDType)))
+def dtype_with_shape(dtype:DType, shape:tuple) -> DType:
+  return dtype.scalar().vec(prod(shape)) if dtype.count == 1 and len(shape) == 1 and isinstance(shape[0], int) and shape[0] > 1 else dtype
+def uops_to_dtypes(uops:list[UOp]) -> list[DType]:
+  return dedup(dtype_with_shape(u.dtype, u._shape or ()) if u.addrspace not in {AddrSpace.GLOBAL, AddrSpace.LOCAL} else u.dtype
+               for u in uops if not isinstance(u.dtype, (ImageDType, PtrDType)))

 # (name, dims, dtype_in, dtype_out, device, threads, upcast_axes, reduce_axes)
 def wmma_args(uops:list[UOp]):
@@ -131,12 +145,14 @@ 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):
+    def arg_dtype(u:UOp) -> DType:
+      return u.dtype if isinstance(u.dtype, (ImageDType, PtrDType)) or u.op is not Ops.PARAM else u.dtype.ptr(u.max_numel(), u.addrspace)
+    if any(isinstance(arg_dtype(u), 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(dt, mutable)+self.buffer_suffix if isinstance(dt, (ImageDType, PtrDType)) else
+                self.arg_int_prefix if dt == dtypes.int else None) for name,(u,mutable) in bufs for dt in (arg_dtype(u),)]
    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}(",] +
@@ -145,6 +161,9 @@ class CStyleLanguage(Renderer):
    return prg if prefix is None else "\n".join(prefix)+f"\n{prg}"

  def render_cast(self, dt:DType, val: str) -> str: return f"({self.render_dtype(dt)})({val})"
+  def render_dtype_with_shape(self, u:UOp) -> DType: return dtype_with_shape(u.dtype, u.shape)
+  def render_access(self, bidx:UOp, dtype:DType) -> str:
+    return f"(*(({self.render_dtype(dtype.ptr(addrspace=bidx.addrspace))})({self[bidx]})))" if dtype.count > 1 else f"(*{self[bidx]})"
  def render_dtype(self, dt:DType, mutable=True) -> str:
    if isinstance(dt, ImageDType): return f"{'write_only' if mutable else 'read_only'} image2d_t"
    if isinstance(dt, PtrDType):
@@ -156,14 +175,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)
@@ -178,9 +197,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
@@ -190,21 +209,24 @@ 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.INDEX: "bidx", Ops.DEFINE_REG: "acc", Ops.LOAD: "val"}.get(u.op, "alu")
+                  Ops.INDEX: "bidx", Ops.SHRINK: "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.INDEX, 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.SHRINK, 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
      else:
+        if u.op is Ops.SHRINK or u._shape is None: u_dtype = u.src[0].dtype
+        else: u_dtype = self.render_dtype_with_shape(u)
        if u.op not in {Ops.RANGE, Ops.DEFINE_LOCAL, Ops.STORE, Ops.DEFINE_REG} and u.dtype != dtypes.void:
-          l = f"{self.render_dtype(u.dtype)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "")
+          l = f"{self.render_dtype(u_dtype)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "")
        kernel.append("  "*depth + l)
        if prefix: c[prefix] += 1  # if it was used, increment
      if u.op in {Ops.IF, Ops.RANGE}: depth += 1
@@ -266,7 +288,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))
@@ -275,12 +297,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
@@ -304,12 +325,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

@@ -320,8 +340,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()
@@ -5,12 +5,14 @@ from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import HIPRenderer, create_non_native_float_pats, pm_manual_bf16_cast
 from tinygrad.uop.decompositions import xexp2, xlog2
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, range_str
-from tinygrad.dtype import dtypes, float_to_fp8, DType, PtrDType, truncate
+from tinygrad.dtype import dtypes, float_to_fp8, DType, PtrDType, truncate, AddrSpace
 from tinygrad.helpers import prod, Target, CPU_COUNT, getenv, OSX

-def ldt(dt:DType):
-  if dt.vcount > 1: return f"<{dt.vcount} x {ldt(dt.scalar())}>"
-  if isinstance(dt, PtrDType): return ldt(dt.base) + "*"
+def ldt(dt:DType, count=1, ptr=False):
+  #if dt.vcount > 1: return f"<{dt.vcount} x {ldt(dt.scalar())}>"
+  #if isinstance(dt, PtrDType): return ldt(dt.base) + "*"
+  if ptr: return ldt(dt, count) + "*"
+  if count > 1: return f"<{count} x {ldt(dt, 1, ptr)}>"
  return {dtypes.void: "void", dtypes.bool: "i1", dtypes.int8: "i8", dtypes.int16: "i16", dtypes.int32: "i32", dtypes.int64: "i64",
          dtypes.uint8: "i8", dtypes.uint16: "i16", dtypes.uint32: "i32", dtypes.uint64: "i64", dtypes.fp8e4m3: "i8", dtypes.fp8e5m2: "i8",
          dtypes.float16: "half", dtypes.bfloat16: "bfloat", dtypes.float32: "float", dtypes.float64: "double"}[dt]
@@ -54,12 +56,13 @@ def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str:
  N,M,K = wmma.arg[1]
  if cdna:
    if K == 32: dt_map.update({dtypes.half: ".f16", dtypes.bfloat16: ".bf16"})
-    return f"  {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype.scalar()]}" + \
-           f".{N}x{M}x{K}{dt_map[wmma.arg[2]]}(" + ", ".join([f"{ldt(w.dtype)} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
+    return f"  {ctx[wmma]} = call {ldt(wmma.dtype, count=wmma.max_numel())} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype.scalar()]}" + \
+           f".{N}x{M}x{K}{dt_map[wmma.arg[2]]}(" + ", ".join([f"{ldt(w.dtype, count=w.max_numel())} {ctx[w]}" for w in wmma.src]) + \
+            ", i32 0, i32 0, i32 0)"
  # https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.wmma_32.ll
  # example: %wmma0 = call <8 x float> @llvm.amdgcn.wmma.f32.16x16x16.f16(<16 x half> %v99,<16 x half> %v100,<8 x float> %v101)
-  return f"  {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype.scalar()]}.16x16x16." + \
-    f"{dt_map[wmma.src[0].dtype.scalar()]}(" + ", ".join([f"{ldt(w.dtype)} {ctx[w]}" for w in wmma.src]) + (", i1 false)" \
+  return f"  {ctx[wmma]} = call {ldt(wmma.dtype, count=wmma.max_numel())} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype.scalar()]}.16x16x16." + \
+    f"{dt_map[wmma.src[0].dtype.scalar()]}(" + ", ".join([f"{ldt(w.dtype, count=w.max_numel())} {ctx[w]}" for w in wmma.src]) + (", i1 false)" \
      if wmma.dtype.scalar() != dtypes.float else ")")

 # llvm ops, lop[<dtype>][<op>]
@@ -75,25 +78,31 @@ lop = {**{x:unsigned_lop for x in (dtypes.bool,)+dtypes.uints}, **{x:signed_lop
 base_rewrite = PatternMatcher([
  # memory load/store
  (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]]}"),
+   f"  {ctx[x]} = extractelement {ldt(x.src[0].dtype, x.src[0].max_numel())} {ctx[x.src[0]]}, i32 {x.src[1].arg}" \
+    if x.addrspace == AddrSpace.ANON else None),
+  (UPat((Ops.INDEX, Ops.SHRINK), name="x"), lambda ctx,x:
+   f"  {ctx[x]} = getelementptr inbounds {ldt(x.dtype)}, {ldt(x.dtype, ptr=True)} {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"
   f"  br i1 {ctx[mask]}, label {ctx[x]}_load, label {ctx[x]}_exit\n{ctx[x][1:]}_load:\n"
-   f"  {ctx[x]}_yes = load {ldt(x.dtype)}, {ldt(idx.dtype)} {ctx[idx]}\n"
+   f"  {ctx[x]}_yes = load {ldt(x.dtype, idx.max_numel())}, {ldt(idx.dtype, x.max_numel(), True)} {ctx[idx]}\n"
   f"  br label {ctx[x]}_exit\n{ctx[x][1:]}_exit:\n"
-   f"  {ctx[x]} = phi {ldt(x.dtype)} [{ctx[x]}_yes, {ctx[x]}_load], [{ctx[alt]}, {ctx[x]}_entry]"),
-  (UPat.var('idx').load(name="x"), lambda ctx,x,idx: f"  {ctx[x]} = load {ldt(x.dtype)}, {ldt(idx.dtype)} {ctx[idx]}"),
-  (UPat.var('idx').store(UPat.var("var")), lambda ctx,idx,var: f"  store {ldt(var.dtype)} {ctx[var]}, {ldt(idx.dtype)} {ctx[idx]}"),
-
+   f"  {ctx[x]} = phi {ldt(x.dtype, idx.max_numel())} [{ctx[x]}_yes, {ctx[x]}_load], [{ctx[alt]}, {ctx[x]}_entry]"),
+  (UPat.var('idx').load(name="x"),
+   lambda ctx,x,idx: f"  {ctx[x]} = load {ldt(idx.dtype, idx.max_numel())}, {ldt(idx.dtype, idx.max_numel(), True)} {ctx[idx]}"),
+  (UPat.var('idx').store(UPat.var("var")),
+   lambda ctx,idx,var:
+   f"  store {ldt(var.dtype, idx.max_numel())} {ctx[var]}, {ldt(idx.dtype, idx.max_numel(), True)} {ctx[idx]}"),
  # GEP/VECTORIZE/CAST for float4 support
-  (UPat(Ops.GEP, name="x"), lambda ctx,x: f"  {ctx[x]} = extractelement {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, i32 {x.arg[0]}"),
+  #(UPat(Ops.GEP, name="x"), lambda ctx,x: f"  {ctx[x]} = extractelement {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, i32 {x.arg[0]}"),
  (UPat(Ops.STACK, src=UPat.var('y'), name="x"), lambda ctx,x,y:
   f"  {ctx[x]}_z = insertelement <1 x {ldt(y.dtype)}> poison, {ldt(y.dtype)} {ctx[y]}, i32 0\n"
-   f"  {ctx[x]} = shufflevector <1 x {ldt(y.dtype)}> {ctx[x]}_z, <1 x {ldt(y.dtype)}> poison, <{x.dtype.count} x i32> zeroinitializer"),
-  (UPat(Ops.STACK, name="x"), lambda ctx,x: "\n".join([(f"  {ctx[x]}_{i}" if i+1 != len(x.src) else f"  {ctx[x]}")+
-                                                            f" = insertelement {ldt(x.dtype)} "+(f"{ctx[x]}_{i-1}" if i != 0 else "poison")+
-                                                            f", {ldt(u.dtype)} {ctx[u]}, i32 {i}" for i,u in enumerate(x.src)])),
+   f"  {ctx[x]} = shufflevector <1 x {ldt(y.dtype)}> {ctx[x]}_z, <1 x {ldt(y.dtype)}> poison, <{x.max_numel()} x i32> zeroinitializer"),
+  (UPat(Ops.STACK, name="x"), lambda ctx,x: "\n".join([(
+   f"  {ctx[x]}_{i}" if i+1 != len(x.src) else f"  {ctx[x]}")+
+   f" = insertelement {ldt(x.dtype, x.max_numel())} "+(f"{ctx[x]}_{i-1}" if i != 0 else "poison")+
+   f", {ldt(u.dtype)} {ctx[u]}, i32 {i}" for i,u in enumerate(x.src)])),
  # unary/binary/ternary ops
  (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"  {ctx[x]} = bitcast {ldt(x.src[0].dtype)} {ctx[x.src[0]]} to {ldt(x.dtype)}"),
  (UPat(Ops.CAST, name="x"), lambda ctx,x: f"  {ctx[x]} = {lcast(x.src[0].dtype, x.dtype)} {ldt(x.src[0].dtype)} {ctx[x.src[0]]} to {ldt(x.dtype)}"),
@@ -135,13 +144,14 @@ 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, ptr=u.addrspace == AddrSpace.GLOBAL)}{' noalias align 32' if u.addrspace == AddrSpace.GLOBAL 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

@@ -164,18 +174,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
@@ -1,17 +1,17 @@
 from typing import Callable, cast, Any
-from tinygrad.dtype import AddrSpace, DType, PtrDType, ImageDType, dtypes, truncate
+from tinygrad.dtype import AddrSpace, DType, ImageDType, dtypes, truncate
 from tinygrad.helpers import DEBUG, OSX, unwrap, fromimport, Target
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import CUDARenderer, OpenCLRenderer
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
 from tinygrad.runtime.autogen import mesa
 from tinygrad.runtime.support.c import POINTER
-import base64, ctypes, ctypes.util, struct, functools, inspect, itertools
+import base64, ctypes, ctypes.util, struct, functools, inspect, itertools, os, warnings

 def g(s:str): return getattr(mesa, s)
 def nsrc(d:mesa.nir_def) -> mesa.nir_src: return mesa.nir_src(ssa=ctypes.pointer(d))

-def glsl_type(t:DType): return mesa.glsl_array_type(glsl_type(t.base), t.size, 0).contents if isinstance(t, PtrDType) else {
+def glsl_type(t:DType): return {
  **{getattr(dtypes,k):g(f"glsl_type_builtin_{v}") for k,v in [('double','double'),('float','float'),('float16','float16_t'),('bool','uint8_t')]},
  **{d:g(f"glsl_type_builtin_{'u' * (d in dtypes.uints)}int{str(d.bitsize)+'_t' if d.itemsize != 4 else ''}") for d in dtypes.ints}}[t]

@@ -25,7 +25,6 @@ aop = {**{x:u_aop for x in (dtypes.bool,)+dtypes.uints}, **{x:s_aop for x in dty

 def c(t:DType, u:bool=True) -> str: return "u" if t in dtypes.uints and u else ("i" if t in dtypes.ints else ("f" if t in dtypes.floats else "b"))
 def ncast(b:mesa.nir_builder, src:mesa.nir_def, it:DType, ot:DType) -> mesa.nir_def:
-  if isinstance(it, PtrDType) and ot == dtypes.long: return src
  return nalu(b, f"{c(it)}2{c(it) if it in dtypes.ints and ot in dtypes.ints else c(ot, ot == dtypes.bool)}{ot.bitsize}", src)

 def nif(b:mesa.nir_builder, cond:mesa.nir_def, then_fn:Callable, else_fn:Callable):
@@ -86,9 +85,9 @@ 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 count:count, bs=lambda dtype:dtype.bitsize, num_components=lambda count: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)}")))
+    lambda b, space, addr, dtype, count=1: mesa.nir_intrinsic_instr_create(b.shader, g(f"nir_intrinsic_load_{scope(space)}")))

 ngid = nir_instr(nc=3, bs=32)(lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_workgroup_id))
 nlid = nir_instr(nc=3, bs=32)(lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_local_invocation_id))
@@ -104,16 +103,31 @@ def njump(b:mesa.nir_builder, typ, tgt=None, cond=None, else_tgt=None): return m

 def if_phi(b:mesa.nir_builder, cond, then_fn, else_fn): return mesa.nir_if_phi(b, *nif(b, cond, then_fn, else_fn)).contents

-def nidx(b:mesa.nir_builder, buf, off, dtype, gate=None) -> mesa.nir_def:
+def _load_count(x:UOp) -> int: return x.max_numel() if 1 < x.max_numel() <= 4 else 1
+def _pad_count(b:mesa.nir_builder, dtype:DType, count:int, val):
+  return val if val.num_components == count else nalu(b, f"vec{count}", val, *[nundef(b, dtype) for _ in range(count-1)])
+
+def nidx(b:mesa.nir_builder, buf, off, dtype, addrspace, gate=None) -> mesa.nir_def:
  @nir_instr(nc=1, bs=32, modes=lambda buf: buf.data.mode, type=lambda buf: mesa.glsl_get_array_element(buf.type))
  def reg(b, buf):
    deref = mesa.nir_deref_instr_create(b.shader, mesa.nir_deref_type_array)
    deref.contents.parent, deref.contents.arr.index = nsrc(deref_var(b, buf)), nsrc(off)
    return deref
-  f = (functools.partial(reg, b, buf) if dtype.addrspace == AddrSpace.REG else
+  f = (functools.partial(reg, b, buf) if addrspace == AddrSpace.REG else
       lambda: nalu(b, "iadd", buf, nalu(b, "imul", off, nimm(b, dtype.itemsize, dtypes.long))))
  return if_phi(b, gate, f, lambda: buf) if gate is not None else f()

+def ngated_load_index(ctx, x, buf, off, alt, gate):
+  cnt = _load_count(x)
+  return if_phi(ctx.b, ctx.r[gate],
+    lambda: nload(ctx.b, buf.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype, buf.addrspace, ctx.r[gate]), x.dtype, cnt),
+    lambda: _pad_count(ctx.b, x.dtype, cnt, ctx.r[alt]))
+
+def ngated_load_shrink(ctx, x, idx, alt, gate):
+  cnt = _load_count(idx)
+  return if_phi(ctx.b, ctx.r[gate], lambda: nload(ctx.b, idx.addrspace, ctx.r[idx], x.dtype, cnt),
+    lambda: _pad_count(ctx.b, x.dtype, cnt, ctx.r[alt]))
+
 class NIRRenderer(Renderer):
  suffix = "NIR"
  nir_options: bytes
@@ -137,7 +151,7 @@ class NIRRenderer(Renderer):
    (UPat(Ops.CAST, (dtypes.uchar, dtypes.ushort), src=(UPat.var("x", dtypes.floats),), name="c"), lambda x,c: x.cast(dtypes.int32).cast(c.dtype)),
    # load/store use pointer arithmetic, and the cast does nothing. NOTE: this doesn't apply to image indexing cause it's 1-D
    (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off")), name="x"), lambda x,buf,off: x.replace(
-      src=(buf,off.cast(dtypes.long))) if buf.dtype.addrspace != AddrSpace.REG and off.op not in (Ops.CAST, Ops.STACK) else None),
+      src=(buf,off.cast(dtypes.long))) if buf.addrspace != AddrSpace.REG and off.op not in (Ops.CAST, Ops.STACK) else None),
    # images need index to be int for nir
    (UPat.var("buf").index(UPat.var("idx_y"), UPat.var("idx_x")),
     lambda buf,idx_y,idx_x: buf.index(idx_y.cast(dtypes.int), idx_x.cast(dtypes.int))),
@@ -149,18 +163,28 @@ 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, buf.addrspace), ctx.r[val], val.dtype)),
+    (UPat(Ops.STORE, src=(UPat(Ops.SHRINK, name="idx"), UPat.var("val"))),
+     lambda ctx,idx,val: nstore(ctx.b, idx.addrspace, ctx.r[idx], 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])),
+     ngated_load_index),
+    (UPat(Ops.LOAD, src=(UPat(Ops.SHRINK, name="idx"), UPat.var("alt"), UPat.var("gate")), name="x"),
+     ngated_load_shrink),
    (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)),
-    (UPat(Ops.STACK, name="x"), lambda ctx,x: nalu(ctx.b, f"vec{x.dtype.count}", *[ctx.r[src] for src in x.src])),
+     lambda ctx,x,buf,off: nload(ctx.b, buf.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype, buf.addrspace), x.dtype, _load_count(x))),
+    (UPat(Ops.LOAD, src=(UPat(Ops.SHRINK, name="idx"),), name="x"),
+     lambda ctx,x,idx: nload(ctx.b, idx.addrspace, ctx.r[idx], x.dtype, _load_count(idx))),
+    (UPat(Ops.SHRINK, src=(UPat.var("buf"), UPat.var("off"), UPat.cvar()), name="x"),
+     lambda ctx,x,buf,off: nidx(ctx.b, ctx.r[buf], ctx.r[off], x.dtype, x.addrspace)),
+    (UPat(Ops.STACK, name="x"), lambda ctx,x: ctx.r[x.src[0]] if len(x.src) == 1 else
+     nalu(ctx.b, f"vec{len(x.src)}", *[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)),
    (UPat(Ops.BITCAST, src=(UPat.var("a"),), allow_any_len=True), lambda ctx,a: ctx.r[a]),
-    (UPat(Ops.GEP, src=(UPat.var("a"),), name="x"), lambda ctx,x,a: nchannel(ctx.b, ctx.r[a], x.arg[0])),
-    (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x:mesa.nir_local_variable_create(ctx.b.impl, glsl_type(x.dtype), f"acc{x.arg}".encode()).contents),
+    (UPat(Ops.INDEX, src=(UPat.var("a"), UPat.cvar("idx"))),
+     lambda ctx,a,idx: nchannel(ctx.b, ctx.r[a], idx.arg) if a.addrspace == AddrSpace.ANON else None),
+    (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: mesa.nir_local_variable_create(ctx.b.impl,
+      mesa.glsl_array_type(glsl_type(x.dtype), x.src[0].arg, 0), f"acc{x.arg}".encode()).contents),
    (UPat(Ops.BARRIER), lambda ctx: nbarrier(ctx.b)),
    (UPat(Ops.IF, name="x"), lambda ctx,x: mesa.nir_push_if(ctx.b, ctx.r[x.src[0]])),
    (UPat(Ops.ENDIF, name="x"), lambda ctx,x: (lambda _: mesa.nir_def())(mesa.nir_pop_if(ctx.b, ctx.r[x.src[0]])))
@@ -189,8 +213,7 @@ class NIRRenderer(Renderer):
    self.param_idx, ranges = 0, []

    for u in uops:
-      if u.op in {Ops.NOOP, Ops.GROUP, Ops.INDEX}: pass
-      elif u.op is Ops.CAST and isinstance(u.dtype, PtrDType): pass
+      if u.op in {Ops.NOOP, Ops.GROUP} or (u.op is Ops.INDEX and u.src[0].addrspace != AddrSpace.ANON): pass
      elif u.op is Ops.AFTER:
        self.r[u] = self.r[u.src[0]]
      elif u.op == Ops.SINK:
@@ -198,7 +221,7 @@ class NIRRenderer(Renderer):
          self.b.shader.contents.info.name = ctypes.cast(ctypes.create_string_buffer(u.arg.function_name.encode()), POINTER[ctypes.c_char])
      elif u.op == Ops.DEFINE_LOCAL:
        self.r[u] = nimm(self.b, self.b.shader.contents.info.shared_size, dtypes.long)
-        self.b.shader.contents.info.shared_size += u.dtype.nbytes()
+        self.b.shader.contents.info.shared_size += u.src[0].arg * u.dtype.itemsize
      elif u.op == Ops.RANGE:
        ranges.append(i:=deref_var(self.b, mesa.nir_local_variable_create(self.b.impl, glsl_type(u.dtype), f"idx{range_str(u)}".encode()).contents))
        nstore(self.b, AddrSpace.REG, i, nimm(self.b, 0, u.dtype), u.dtype)
@@ -217,7 +240,17 @@ class NIRRenderer(Renderer):
        self.r[u] = cast(mesa.nir_def, d)
    self.postrender(uops)

-    mesa.nir_validate_shader(self.b.shader, b"after render")
+    if DEBUG >= 2 and hasattr(os, "fork"):
+      with warnings.catch_warnings():
+        warnings.simplefilter("ignore", DeprecationWarning)
+        pid = os.fork()
+      if pid == 0:
+        mesa.nir_validate_shader(self.b.shader, b"after render")
+        os._exit(0)
+      _, status = os.waitpid(pid, 0)
+      if os.WIFSIGNALED(status): raise RuntimeError(f"NIR validation failed after render with signal {os.WTERMSIG(status)}")
+      if os.WEXITSTATUS(status) != 0: raise RuntimeError(f"NIR validation failed after render with exit code {os.WEXITSTATUS(status)}")
+    else: mesa.nir_validate_shader(self.b.shader, b"after render")
    if DEBUG >= 4: mesa.nir_print_shader(self.b.shader, ctypes.POINTER(mesa.struct__IO_FILE).in_dll(ctypes.CDLL(ctypes.util.find_library('c')),
                                                                                                    "__stdoutp" if OSX else "stdout"))
    mesa.nir_serialize(blob:=mesa.struct_blob(), self.b.shader, False)
@@ -257,9 +290,10 @@ class LVPRenderer(NIRRenderer):

 def tovec(b, idx_y, idx_x): return nalu(b, "vec4", idx_x, idx_y, nundef(b, dtypes.int), nundef(b, dtypes.int))
 def nfloat(dtype): return mesa.nir_type_float16 if dtype == dtypes.half else mesa.nir_type_float32
-nstore_img = nir_instr(has_def=False, df=lambda img:img, num_components=lambda val:val.num_components,
+nstore_img = nir_instr(has_def=False, df=lambda img:img, num_components=4,
  intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2D, 'ACCESS':mesa.ACCESS_CAN_REORDER, 'SRC_TYPE':nfloat(dtype)},
-  srcs=lambda b,img,idx_y,idx_x,val:[nsrc(x) for x in [img, tovec(b, idx_y, idx_x), nundef(b, dtypes.int), val, nimm(b, 0, dtypes.int)]])(
+  srcs=lambda b,img,idx_y,idx_x,val,dtype:[nsrc(x) for x in [img, tovec(b, idx_y, idx_x), nundef(b, dtypes.int),
+    val if val.num_components == 4 else nalu(b, "vec4", val, nundef(b, dtype), nundef(b, dtype), nundef(b, dtype)), nimm(b, 0, dtypes.int)]])(
    lambda b,img,idx_y,idx_x,val,dtype:mesa.nir_intrinsic_instr_create(b.shader,g("nir_intrinsic_image_store")))

 _nload_img = nir_instr(intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2D, 'ACCESS':mesa.ACCESS_CAN_REORDER, 'DEST_TYPE':nfloat(dtype)},
@@ -277,8 +311,11 @@ class IR3Renderer(NIRRenderer, OpenCLRenderer):
  def_rewrite = PatternMatcher([
    (UPat(Ops.STORE, src=(UPat.var('img').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("val")), allow_any_len=True),
     lambda ctx,img,idx_y,idx_x,val: nstore_img(ctx.b, ctx.r[img], ctx.r[idx_y], ctx.r[idx_x], ctx.r[val], val.dtype)),
-    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("alt"), UPat.var("gate"))),
-     lambda ctx,img,idx_y,idx_x,alt,gate: if_phi(ctx.b, ctx.r[gate], lambda: ctx.nload_img(img, idx_y, idx_x), lambda: ctx.r[alt])),
+    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("alt"), UPat.var("gate")), name="x"),
+     lambda ctx,x,img,idx_y,idx_x,alt,gate: if_phi(ctx.b, ctx.r[gate],
+       lambda: ctx.nload_img(img, idx_y, idx_x) if len(x.shape) > 0 and x.shape[-1] == 4 else nchannel(ctx.b, ctx.nload_img(img, idx_y, idx_x), 0),
+       lambda: ctx.r[alt] if len(x.shape) == 0 or x.shape[-1] != 4 or ctx.r[alt].num_components == 4 else
+         nalu(ctx.b, "vec4", ctx.r[alt], nundef(ctx.b, x.dtype), nundef(ctx.b, x.dtype), nundef(ctx.b, x.dtype)))),
    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('idx_y'), UPat.var('idx_x')),)), nload_img),
  ]) + NIRRenderer.def_rewrite

@@ -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,19 @@ 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(x:UOp) -> bool:
+  if x.op is Ops.LOAD: dt, addrspace = x.dtype, x.src[0].addrspace
+  elif x.op is Ops.STORE: dt, addrspace = x.src[1].dtype, x.src[0].addrspace
+  else: dt, addrspace = x.dtype.base, x.dtype.addrspace if isinstance(x.dtype, PtrDType) else x.addrspace
+  return dt.itemsize < 4 and dt.base != dtypes.half and addrspace != AddrSpace.REG
+
+def _packed_size(ctx, x:UOp):
+  size = ctx[x.src[0]]
+  if not is_packed(x): return size
+  elems = 4 // x.dtype.base.itemsize
+  return str((x.src[0].arg + elems - 1) // elems) if x.src[0].op is Ops.CONST else f"(({size}+{elems-1})/{elems})"
+
+def _buf_map(ctx, x:UOp): return ctx.type_map[x.dtype.base] if x.addrspace == AddrSpace.REG else ctx.buf_map(x.dtype.base)

 def is_nan(a):
  bs, (exp, mant) = a.dtype.bitsize, dtypes.finfo(a.dtype)
@@ -41,12 +50,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),
-  (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),
-  (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 l,b,c,gate: packed_load(l,b,l.dtype,c.cast(dtypes.uint32),gate) if is_packed(l) else None),
+  (UPat.load(UPat.var("b"), name='l'), lambda l,b: packed_load(l, b, l.dtype) if is_packed(l) else None),
+  (UPat.store(UPat.var("bidx"), UPat.var("var"), UPat.var("gate"), name="s"),
+   lambda s,bidx,var,gate: packed_store(bidx,var,gate) if is_packed(s) else None),
+  (UPat.store(UPat.var("bidx"), UPat.var("var"), name="s"),
+   lambda s,bidx,var: packed_store(bidx,var) if is_packed(s) 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 +80,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<{_buf_map(ctx,x)},{_packed_size(ctx,x)}>;"),
+    (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"var {ctx[x]}: array<{_buf_map(ctx,x)},{_packed_size(ctx,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)"),
@@ -83,12 +92,12 @@ class WGSLRenderer(CStyleLanguage):
     if x.src[0].dtype == dtypes.half else f"((i32({ctx[x.src[0]]}&0xFFFF)<<16)>>16)"),
    (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)),
-    (UPat.store(UPat.var("b"), UPat.var("v")), lambda ctx,b,v:\
+    (UPat.load(UPat.var("b"), UPat.cvar("v"), UPat.var("gate"), name="l"),
+      lambda ctx,l,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b],b)}, {ctx[gate]})"),
+    (UPat.load(UPat.var("b"), name="l"), lambda ctx,l,b: ctx.render_load(ctx[b], b)),
+    (UPat.store(UPat.var("b"), UPat.var("v"), name="s"), lambda ctx,s,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) and b.addrspace != AddrSpace.REG \
      else f"{ctx[b]} = {ctx[v]};"),
    (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]}]"),
@@ -96,9 +105,11 @@ class WGSLRenderer(CStyleLanguage):

  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, uop:UOp) -> str: return f"atomicLoad(&{x})" if is_packed(uop) and uop.addrspace != AddrSpace.REG else x
+  def buf_map(self, dt:DType) -> str: return "atomic<u32>" if dt.itemsize < 4 and dt != dtypes.half else self.type_map[dt.base]
+  def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[UOp,bool]]], uops:list[UOp], prefix=None) -> str:
+    def arg_dtype(u:UOp) -> DType:
+      return u.dtype if isinstance(u.dtype, PtrDType) or u.op is not Ops.PARAM else u.dtype.ptr(u.max_numel(), u.addrspace)
    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 +119,9 @@ 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(dt, PtrDType) else 'var<uniform>'}" +
+      f"{name}:{f'array<{self.buf_map(dt.base)}>' if isinstance(dt,PtrDType) else self.buf_map(dt)};"
+      for name,(u,_) in bufs for dt in (arg_dtype(u),)])
    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,21 @@ 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:
+    def arg_dtype(u:UOp): return u.dtype if isinstance(u.dtype, PtrDType) or u.op is not Ops.PARAM else u.dtype.ptr(u.max_numel(), u.addrspace)
    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(arg_dtype(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(arg_dtype(b[1][0]), 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)])});"]
+    params = [(f'buf_{i}' if isinstance(arg_dtype(b[1][0]), PtrDType) else f'sz_or_val_{i}') for i,b in enumerate(bufs)]
+    msrc += [f"{function_name}({', '.join(params)});"]
    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(arg_dtype(b[1][0]), PtrDType)]
    msrc += ["return 0; }"]
    return '\n'.join(msrc)

@@ -273,22 +276,24 @@ 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:
+    def arg_dtype(u:UOp): return u.dtype if isinstance(u.dtype, PtrDType) or u.op is not Ops.PARAM else u.dtype.ptr(u.max_numel(), u.addrspace)
    # 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(dt:=arg_dtype(b[1][0]), PtrDType):
+        sz = dt.size*dt.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(arg_dtype(b[1][0]), 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(dt:=arg_dtype(b[1][0]), PtrDType): msrc.append(f"write(1, buf{i}, {dt.size*dt.itemsize});")
    msrc.append('exit(0); }')
    return '\n'.join(msrc)

@@ -5,8 +5,8 @@
 from typing import Any, TYPE_CHECKING
 import pickle, base64, itertools, time, sys, functools
 from dataclasses import replace
-from tinygrad.dtype import DType, dtypes, ImageDType, PtrDType, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
-from tinygrad.helpers import all_same, getenv, flatten, get_single_element, Target, IMAGE
+from tinygrad.dtype import DType, dtypes, ImageDType, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar, AddrSpace
+from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE
 from tinygrad.device import Compiled, Compiler, Allocator
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp, bitcast
@@ -41,103 +41,105 @@ 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.INDEX:
+            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.SHRINK or (u.op is Ops.INDEX and len(src_values) == 2):
+          if u.addrspace == AddrSpace.ANON:
+            # old GEP
+            assert all_same(src_values[1]), "all index must be the same"
+            values[u] = src_values[0][src_values[1][0]]
+          else:
+            # normal index
+            values[u] = [(m,o) for m,o in zip(src_values[0], src_values[1])]
+        elif u.op is Ops.INDEX and len(src_values) == 3:
+          assert isinstance(src_dtypes[0], ImageDType), "3 src index is only for Image"
          ret:list = []
-          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))
+          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[u] = ret
+        elif u.op is Ops.RANGE:
+          if u not in values: values[u] = [0] * warp_size
          else:
-            assert len(src_values) == 2, "non-image index must be 2 srcs"
-            for m,o in zip(*src_values): ret.append((m,o))
-          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
-          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
+            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.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":
@@ -145,17 +147,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):
@@ -164,7 +166,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)
@@ -172,24 +174,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]
@@ -197,17 +199,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

@@ -234,10 +236,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)}

@@ -20,7 +20,7 @@ BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2
 def dcache_flush():
  from tinygrad.uop.ops import UOp, Ops, KernelInfo
  from tinygrad.codegen import to_program
-  buf, n = UOp(Ops.PARAM, dtypes.uint8.ptr(), arg=0), UOp(Ops.PARAM, dtypes.uint8.ptr(), arg=1)
+  buf, n = UOp.param(0, dtypes.uint8.ptr()), UOp.param(1, dtypes.uint8.ptr())
  i = UOp.range(n.cast(dtypes.int), 0, dtype=dtypes.int)
  flush = UOp(Ops.CUSTOM, dtypes.void, (buf.cast(dtypes.ulong) + i.cast(dtypes.ulong) * UOp.const(dtypes.ulong, 64),),
              arg='__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");')
@@ -1,157 +1,120 @@
 import functools, struct
 from tinygrad.device import Compiled, Allocator, BufferSpec
 from tinygrad.renderer.wgsl import WGSLRenderer
-from tinygrad.helpers import round_up, suppress_finalizing
+from tinygrad.helpers import round_up, suppress_finalizing, getenv, to_mv
 from tinygrad.runtime.autogen import webgpu
 from tinygrad.runtime.support import c
-from typing import cast, List, Any, TypeAlias
+from typing import Callable
 import ctypes
-import os

-WGPUDevPtr: TypeAlias = webgpu.WGPUDevice
-WGPUBufPtr: TypeAlias = webgpu.WGPUBuffer
+backend_types = {v: k for k, v in webgpu.enum_WGPUBackendType.items()}
+instance = webgpu.wgpuCreateInstance(webgpu.WGPUInstanceDescriptor(features=webgpu.WGPUInstanceFeatures(timedWaitAnyEnable=True)))

-backend_types = {v: k for k, v in webgpu.enum_WGPUBackendType.items() }
+def from_wgpu_str(string_view:webgpu.WGPUStringView) -> str: return ctypes.string_at(string_view.data, string_view.length).decode()
+def to_wgpu_str(_str:str) -> webgpu.WGPUStringView: return webgpu.WGPUStringView(data=ctypes.create_string_buffer(_str.encode()), length=len(_str))

-instance = webgpu.wgpuCreateInstance(webgpu.WGPUInstanceDescriptor(features = webgpu.WGPUInstanceFeatures(timedWaitAnyEnable = True)))
+# gets a memoryview from a buffer, which is assumed to have MAP_READ (see _readable_buffer)
+def buf_to_mv(buf:webgpu.WGPUBuffer) -> memoryview:
+  BufferMapAsync(buf, webgpu.WGPUMapMode_Read, 0, size:=webgpu.wgpuBufferGetSize(buf))
+  return to_mv(webgpu.wgpuBufferGetConstMappedRange(buf, 0, size), size)

-def to_c_string(_str:str) -> ctypes.Array: return ctypes.create_string_buffer(_str.encode('utf-8'))
+# turns a webgpu function returning a future into python-synchronous function
+# the new function handles the status code and optional error message, returning the other callback arguments
+def synchronous(status_enum:dict[int, str], has_emsg:bool=False):
+  def wrap(fn:Callable[..., webgpu.WGPUFuture]) -> Callable:
+    @functools.wraps(fn)
+    def wrapper(*args):
+      status, payload, emsg = 0, [], None

-def from_wgpu_str(string_view:webgpu.struct_WGPUStringView) -> str: return ctypes.string_at(string_view.data, string_view.length).decode("utf-8")
+      @next(ty for nm, ty, *_ in fn.argtypes[-1]._real_fields_ if nm == "callback") # type: ignore
+      def cb(s:int, *args):
+        nonlocal status, payload, emsg
+        # the last two arguments are "userdata1" and "userdata2", which we drop
+        # we must process wgpu strings in this callback, as they will be freed after we return
+        status, (*payload, emsg) = s, [from_wgpu_str(a) if type(a) is webgpu.WGPUStringView else a for a in args[:-2]] + ([] if has_emsg else [None])

-def to_wgpu_str(_str:str) -> webgpu.struct_WGPUStringView:
-  return webgpu.WGPUStringView(data=ctypes.cast(ctypes.pointer(to_c_string(_str)), ctypes.POINTER(ctypes.c_char)), length=len(_str))
+      future = fn(*args, fn.argtypes[-1](mode=webgpu.WGPUCallbackMode_WaitAnyOnly, callback=cb)) # type: ignore
+      if (future_status:=webgpu.wgpuInstanceWaitAny(instance, 1, webgpu.WGPUFutureWaitInfo(future), 2**64-1)) != webgpu.WGPUWaitStatus_Success:
+        raise RuntimeError(f"error while waiting for future ({fn.__name__}): {webgpu.enum_WGPUWaitStatus.get(future_status)}")

-def _wait(future:webgpu.struct_WGPUFuture):
-  assert webgpu.wgpuInstanceWaitAny(instance, 1, webgpu.WGPUFutureWaitInfo(future=future), 2**64-1) == webgpu.WGPUWaitStatus_Success, "Future failed"
+      if status != 1: raise RuntimeError(f"[{status_enum.get(status)}]{emsg or ''}")
+      return payload if len(payload) > 1 else payload[0] if len(payload) == 1 else None
+    return wrapper
+  return wrap

-def write_buffer(device:WGPUDevPtr, buf:WGPUBufPtr, offset:int, src:memoryview|bytearray|bytes):
-  src = bytearray(src)
-  webgpu.wgpuQueueWriteBuffer(webgpu.wgpuDeviceGetQueue(device), buf, offset, (ctypes.c_uint8 * len(src)).from_buffer(src), len(src))
-
-def _run(async_fun, cb_info_type, cb_type, status_enum:dict|None, res_idx:int|None, msg_idx:int|None, *params):
-  result: List[Any] = []
-
-  def cb(*params):
-    result[:] = params
-    if msg_idx: result[msg_idx] = from_wgpu_str(result[msg_idx])
-
-  cb_info = cb_info_type(mode=webgpu.WGPUCallbackMode_WaitAnyOnly, callback=cb_type(cb))
-  _wait(async_fun(*params, cb_info))
-
-  if result[0] != 1: raise RuntimeError(f"[{status_enum.get(result[0]) if status_enum else 'ERROR'}]{result[msg_idx] if msg_idx else ''}")
-  return result[res_idx] if res_idx else None
-
-def copy_buffer_to_buffer(dev:WGPUDevPtr, src:WGPUBufPtr, src_offset:int, dst:WGPUBufPtr, dst_offset:int, size:int):
-  encoder = webgpu.wgpuDeviceCreateCommandEncoder(dev, webgpu.WGPUCommandEncoderDescriptor())
-  webgpu.wgpuCommandEncoderCopyBufferToBuffer(encoder, src, src_offset, dst, dst_offset, size)
-  cb = webgpu.wgpuCommandEncoderFinish(encoder, webgpu.WGPUCommandBufferDescriptor())
-  webgpu.wgpuQueueSubmit(webgpu.wgpuDeviceGetQueue(dev), 1, (webgpu.WGPUCommandBuffer*1)(cb))
-  webgpu.wgpuCommandBufferRelease(cb)
-  webgpu.wgpuCommandEncoderRelease(encoder)
-
-def read_buffer(dev:WGPUDevPtr, buf:WGPUBufPtr) -> memoryview:
-  size = webgpu.wgpuBufferGetSize(buf)
-  tmp_buffer = webgpu.wgpuDeviceCreateBuffer(dev, webgpu.WGPUBufferDescriptor(size=size,
-    usage=webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_MapRead, mappedAtCreation=False))
-  copy_buffer_to_buffer(dev, buf, 0, tmp_buffer, 0, size)
-  _run(webgpu.wgpuBufferMapAsync2, webgpu.WGPUBufferMapCallbackInfo2, webgpu.WGPUBufferMapCallback2, webgpu.enum_WGPUBufferMapAsyncStatus, None, 0,
-       tmp_buffer, webgpu.WGPUMapMode_Read, 0, size)
-  void_ptr = ctypes.cast(webgpu.wgpuBufferGetConstMappedRange(tmp_buffer, 0, size), ctypes.c_void_p)
-  buf_copy = bytearray((ctypes.c_uint8 * size).from_address(void_ptr.value))
-  webgpu.wgpuBufferUnmap(tmp_buffer)
-  webgpu.wgpuBufferDestroy(tmp_buffer)
-  return memoryview(buf_copy).cast("B")
-
-def pop_error(device:WGPUDevPtr) -> str:
-  return _run(webgpu.wgpuDevicePopErrorScopeF, webgpu.WGPUPopErrorScopeCallbackInfo, webgpu.WGPUPopErrorScopeCallback, None, 2, 2, device)
-
-def create_uniform(wgpu_device:WGPUDevPtr, val:int|float) -> WGPUBufPtr:
-  buf = webgpu.wgpuDeviceCreateBuffer(wgpu_device,
-    webgpu.WGPUBufferDescriptor(size=4, usage=webgpu.WGPUBufferUsage_Uniform | webgpu.WGPUBufferUsage_CopyDst))
-  write_buffer(wgpu_device, buf, 0, val.to_bytes(4, "little") if isinstance(val, int) else struct.pack('<f', val))
-  return buf
+BufferMapAsync = synchronous(webgpu.enum_WGPUBufferMapAsyncStatus, True)(webgpu.wgpuBufferMapAsync2)
+DevicePopErrorScope = synchronous(webgpu.enum_WGPUPopErrorScopeStatus)(webgpu.wgpuDevicePopErrorScope2)
+DeviceCreateComputePipeline = synchronous(webgpu.enum_WGPUCreatePipelineAsyncStatus, True)(webgpu.wgpuDeviceCreateComputePipelineAsync2)
+InstanceRequestAdapter = synchronous(webgpu.enum_WGPURequestAdapterStatus, True)(webgpu.wgpuInstanceRequestAdapter2)
+AdapterRequestDevice = synchronous(webgpu.enum_WGPURequestDeviceStatus, True)(webgpu.wgpuAdapterRequestDevice2)
+QueueOnSubmittedWorkDone = synchronous(webgpu.enum_WGPUQueueWorkDoneStatus)(webgpu.wgpuQueueOnSubmittedWorkDone2)

 class WebGPUProgram:
-  def __init__(self, dev:tuple[WGPUDevPtr, bool], name:str, lib:bytes, **kwargs):
-    (self.dev, self.timestamp_supported) = dev
+  def __init__(self, dev:'WebGpuDevice', name:str, lib:bytes, **kwargs):
+    self.dev, self.name = dev, to_wgpu_str(name)

    # Creating shader module
    shader = webgpu.WGPUShaderModuleWGSLDescriptor(code=to_wgpu_str(lib.decode()),
-      chain=webgpu.WGPUChainedStruct(sType=webgpu.WGPUSType_ShaderSourceWGSL))
-    module = webgpu.WGPUShaderModuleDescriptor()
-    module.nextInChain = ctypes.cast(ctypes.pointer(shader), c.POINTER[webgpu.struct_WGPUChainedStruct])
+                                                   chain=webgpu.WGPUChainedStruct(sType=webgpu.WGPUSType_ShaderSourceWGSL))
+    module = webgpu.WGPUShaderModuleDescriptor(nextInChain=ctypes.cast(ctypes.pointer(shader), ctypes.POINTER(webgpu.struct_WGPUChainedStruct)))

    # Check compiler error
-    webgpu.wgpuDevicePushErrorScope(self.dev, webgpu.WGPUErrorFilter_Validation)
-    shader_module = webgpu.wgpuDeviceCreateShaderModule(self.dev, module)
+    webgpu.wgpuDevicePushErrorScope(self.dev.device_res, webgpu.WGPUErrorFilter_Validation)
+    self.prg = webgpu.wgpuDeviceCreateShaderModule(self.dev.device_res, module)
+    if err := self.dev.pop_error(): raise RuntimeError(f"Shader compilation failed: {err}")

-    if err := pop_error(self.dev): raise RuntimeError(f"Shader compilation failed: {err}")
+  @suppress_finalizing
+  def __del__(self): webgpu.wgpuShaderModuleRelease(self.prg)

-    self.name, self.lib, self.prg = name, lib, shader_module
-  def __call__(self, *bufs:WGPUBufPtr, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
+  def __call__(self, *bufs:webgpu.WGPUBuffer, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
               vals:tuple[int, ...]=(), wait=False, **kw) -> float|None:
-    wait = wait and self.timestamp_supported
-    tmp_bufs = [*bufs]
-    buf_patch = False
-
-    # WebGPU does not allow using the same buffer for input and output
-    for i in range(1, len(bufs)):
-      if ctypes.addressof(bufs[i]) == ctypes.addressof(bufs[0]):
-        tmp_bufs[0] = webgpu.wgpuDeviceCreateBuffer(self.dev,
-          webgpu.WGPUBufferDescriptor(size=webgpu.wgpuBufferGetSize(bufs[0]), usage=webgpu.wgpuBufferGetUsage(bufs[0])))
-        buf_patch = True
+    wait = wait and webgpu.WGPUFeatureName_TimestampQuery in self.dev.features

    # Creating bind group layout
-    binding_layouts = [webgpu.WGPUBindGroupLayoutEntry(binding=0, visibility= webgpu.WGPUShaderStage_Compute,
-      buffer=webgpu.WGPUBufferBindingLayout(type=webgpu.WGPUBufferBindingType_Uniform))]
-    binding_layouts += [webgpu.WGPUBindGroupLayoutEntry(binding=i+1, visibility=webgpu.WGPUShaderStage_Compute,
-      buffer=webgpu.WGPUBufferBindingLayout(type=webgpu.WGPUBufferBindingType_Uniform if i >= len(tmp_bufs)
-      else webgpu.WGPUBufferBindingType_Storage)) for i in range(len(tmp_bufs)+len(vals))]
+    def bgl_entry(n:int, ty:str):
+      return webgpu.WGPUBindGroupLayoutEntry(binding=n, visibility=webgpu.WGPUShaderStage_Compute,
+                                             buffer=webgpu.WGPUBufferBindingLayout(type=getattr(webgpu, f'WGPUBufferBindingType_{ty}')))
+    bind_entries = (webgpu.WGPUBindGroupLayoutEntry * (1+len(bufs)+len(vals)))(
+      bgl_entry(0, 'Uniform'), *(bgl_entry(i+1, 'Uniform' if i >= len(bufs) else 'Storage') for i in range(len(bufs)+len(vals))))

-    bl_arr_type = webgpu.WGPUBindGroupLayoutEntry * len(binding_layouts)
-    webgpu.wgpuDevicePushErrorScope(self.dev, webgpu.WGPUErrorFilter_Validation)
-    bind_group_layouts = [webgpu.wgpuDeviceCreateBindGroupLayout(self.dev, webgpu.WGPUBindGroupLayoutDescriptor(
-      entryCount=len(binding_layouts), entries=ctypes.cast(bl_arr_type(*binding_layouts), ctypes.POINTER(webgpu.WGPUBindGroupLayoutEntry))))]
+    webgpu.wgpuDevicePushErrorScope(self.dev.device_res, webgpu.WGPUErrorFilter_Validation)
+    bind_layout = webgpu.wgpuDeviceCreateBindGroupLayout(self.dev.device_res,
+                                                         webgpu.WGPUBindGroupLayoutDescriptor(entryCount=len(bind_entries), entries=bind_entries))

-    if bg_layout_err := pop_error(self.dev): raise RuntimeError(f"Error creating bind group layout: {bg_layout_err}")
+    if err := self.dev.pop_error(): raise RuntimeError(f"Error creating bind group layout: {err}")

    # Creating pipeline layout
-    pipeline_layout_desc = webgpu.WGPUPipelineLayoutDescriptor(bindGroupLayoutCount=len(bind_group_layouts),
-      bindGroupLayouts = (webgpu.WGPUBindGroupLayout * len(bind_group_layouts))(*bind_group_layouts))
+    pipeline_layout_desc = webgpu.WGPUPipelineLayoutDescriptor(bindGroupLayoutCount=1, bindGroupLayouts=(webgpu.WGPUBindGroupLayout*1)(bind_layout))

-    webgpu.wgpuDevicePushErrorScope(self.dev, webgpu.WGPUErrorFilter_Validation)
-    pipeline_layout = webgpu.wgpuDeviceCreatePipelineLayout(self.dev, pipeline_layout_desc)
-
-    if pipe_err := pop_error(self.dev): raise RuntimeError(f"Error creating pipeline layout: {pipe_err}")
+    webgpu.wgpuDevicePushErrorScope(self.dev.device_res, webgpu.WGPUErrorFilter_Validation)
+    pipeline_layout = webgpu.wgpuDeviceCreatePipelineLayout(self.dev.device_res, pipeline_layout_desc)
+    if err := self.dev.pop_error(): raise RuntimeError(f"Error creating pipeline layout: {err}")

    # Creating bind group
-    bindings = [webgpu.WGPUBindGroupEntry(binding=0, buffer=create_uniform(self.dev, float('inf')), offset=0, size=4)]
-    bindings += [webgpu.WGPUBindGroupEntry(binding=i+1, buffer=create_uniform(self.dev, cast(int, x)) if i >= len(tmp_bufs) else x, offset=0,
-      size=4 if i >= len(tmp_bufs) else webgpu.wgpuBufferGetSize(x)) for i,x in enumerate(tuple(tmp_bufs)+vals)]
+    def bg_entry(n:int, x:webgpu.WGPUBuffer|int|float):
+      buf = x if isinstance(x, webgpu.WGPUBuffer) else self.dev.create_uniform(x)
+      return webgpu.WGPUBindGroupEntry(binding=n, buffer=buf, offset=0, size=webgpu.wgpuBufferGetSize(buf))
+    bindings = (webgpu.WGPUBindGroupEntry * (1+len(bufs)+len(vals)))(bg_entry(0, float('inf')), *(bg_entry(i+1, x) for i,x in enumerate(bufs+vals)))

-    bg_arr_type = webgpu.WGPUBindGroupEntry * len(bindings)
-    bind_group_desc = webgpu.WGPUBindGroupDescriptor(layout=bind_group_layouts[0], entryCount=len(bindings), entries=bg_arr_type(*bindings))
-    webgpu.wgpuDevicePushErrorScope(self.dev, webgpu.WGPUErrorFilter_Validation)
-    bind_group = webgpu.wgpuDeviceCreateBindGroup(self.dev, bind_group_desc)
-
-    if bind_err := pop_error(self.dev): raise RuntimeError(f"Error creating bind group: {bind_err}")
+    bind_group_desc = webgpu.WGPUBindGroupDescriptor(layout=bind_layout, entryCount=len(bindings), entries=bindings)
+    webgpu.wgpuDevicePushErrorScope(self.dev.device_res, webgpu.WGPUErrorFilter_Validation)
+    bind_group = webgpu.wgpuDeviceCreateBindGroup(self.dev.device_res, bind_group_desc)
+    if err := self.dev.pop_error(): raise RuntimeError(f"Error creating bind group: {err}")

    # Creating compute pipeline
    compute_desc = webgpu.WGPUComputePipelineDescriptor(layout=pipeline_layout,
-      compute=webgpu.WGPUComputeState(module=self.prg, entryPoint=to_wgpu_str(self.name)))
-    pipeline_result = _run(webgpu.wgpuDeviceCreateComputePipelineAsync2, webgpu.WGPUCreateComputePipelineAsyncCallbackInfo2,
-    webgpu.WGPUCreateComputePipelineAsyncCallback2, webgpu.enum_WGPUCreatePipelineAsyncStatus, 1, None, self.dev, compute_desc)
+                                                        compute=webgpu.WGPUComputeState(module=self.prg, entryPoint=self.name))
+    pipeline_result = DeviceCreateComputePipeline(self.dev.device_res, compute_desc)

-    command_encoder = webgpu.wgpuDeviceCreateCommandEncoder(self.dev, webgpu.WGPUCommandEncoderDescriptor())
+    command_encoder = webgpu.wgpuDeviceCreateCommandEncoder(self.dev.device_res, webgpu.WGPUCommandEncoderDescriptor())
    comp_pass_desc = webgpu.WGPUComputePassDescriptor()

    if wait:
-      query_set = webgpu.wgpuDeviceCreateQuerySet(self.dev, webgpu.WGPUQuerySetDescriptor(type=webgpu.WGPUQueryType_Timestamp, count=2))
-      query_buf = webgpu.wgpuDeviceCreateBuffer(self.dev,
-        webgpu.WGPUBufferDescriptor(size=16, usage=webgpu.WGPUBufferUsage_QueryResolve | webgpu.WGPUBufferUsage_CopySrc))
-      comp_pass_desc.timestampWrites = c.pointer(webgpu.WGPUComputePassTimestampWrites(
-        querySet=query_set, beginningOfPassWriteIndex=0, endOfPassWriteIndex=1))
+      query_set = webgpu.wgpuDeviceCreateQuerySet(self.dev.device_res, webgpu.WGPUQuerySetDescriptor(type=webgpu.WGPUQueryType_Timestamp, count=2))
+      query_buf = webgpu.wgpuDeviceCreateBuffer(
+        self.dev.device_res, webgpu.WGPUBufferDescriptor(size=16, usage=webgpu.WGPUBufferUsage_QueryResolve | webgpu.WGPUBufferUsage_CopySrc))
+      comp_pass_desc.timestampWrites = c.pointer(webgpu.WGPUComputePassTimestampWrites(querySet=query_set, beginningOfPassWriteIndex=0,
+                                                                                       endOfPassWriteIndex=1))

    # Begin compute pass
    compute_pass = webgpu.wgpuCommandEncoderBeginComputePass(command_encoder, comp_pass_desc)
@@ -163,63 +126,96 @@ class WebGPUProgram:
    if wait: webgpu.wgpuCommandEncoderResolveQuerySet(command_encoder, query_set, 0, 2, query_buf, 0)

    cmd_buf = webgpu.wgpuCommandEncoderFinish(command_encoder, webgpu.WGPUCommandBufferDescriptor())
-    webgpu.wgpuQueueSubmit(webgpu.wgpuDeviceGetQueue(self.dev), 1, (webgpu.WGPUCommandBuffer*1)(cmd_buf))
+    webgpu.wgpuQueueSubmit(self.dev.queue, 1, (webgpu.WGPUCommandBuffer*1)(cmd_buf))

-    if buf_patch:
-      copy_buffer_to_buffer(self.dev, tmp_bufs[0], 0, bufs[0], 0, webgpu.wgpuBufferGetSize(bufs[0]))
-      webgpu.wgpuBufferDestroy(tmp_bufs[0])
+    # release created objects
+    webgpu.wgpuBindGroupLayoutRelease(bind_layout)
+    webgpu.wgpuPipelineLayoutRelease(pipeline_layout)
+    webgpu.wgpuBindGroupRelease(bind_group)
+    webgpu.wgpuComputePipelineRelease(pipeline_result)
+    webgpu.wgpuCommandEncoderRelease(command_encoder)
+    webgpu.wgpuComputePassEncoderRelease(compute_pass)
+    webgpu.wgpuCommandBufferRelease(cmd_buf)

    if wait:
-      time = ((timestamps:=read_buffer(self.dev, query_buf).cast("Q").tolist())[1] - timestamps[0]) / 1e9
-      webgpu.wgpuBufferDestroy(query_buf)
+      time = ((timestamps:=buf_to_mv(tmp_buf:=self.dev._readable_buffer(query_buf)).cast("Q").tolist())[1] - timestamps[0]) / 1e9
+      self.dev.free(query_buf)
+      self.dev.free(tmp_buf)
      webgpu.wgpuQuerySetDestroy(query_set)
+      webgpu.wgpuQuerySetRelease(query_set)
      return time
    return None

 class WebGpuAllocator(Allocator['WebGpuDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> WGPUBufPtr:
+  def _alloc(self, size:int, options:BufferSpec) -> webgpu.WGPUBuffer:
    # WebGPU buffers have to be 4-byte aligned
    return webgpu.wgpuDeviceCreateBuffer(self.dev.device_res, webgpu.WGPUBufferDescriptor(size=round_up(size, 4),
      usage=webgpu.WGPUBufferUsage_Storage | webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_CopySrc))
-  def _copyin(self, dest:WGPUBufPtr, src:memoryview):
+  def _copyin(self, dest:webgpu.WGPUBuffer, src:memoryview):
    if src.nbytes % 4:
      padded_src = bytearray(round_up(src.nbytes, 4))
      padded_src[:src.nbytes] = src
-    write_buffer(self.dev.device_res, dest, 0, padded_src if src.nbytes % 4 else src)
-  def _copyout(self, dest:memoryview, src:WGPUBufPtr):
-    buffer_data = read_buffer(self.dev.device_res, src)
-    dest[:] = buffer_data[:dest.nbytes] if webgpu.wgpuBufferGetSize(src)  > dest.nbytes else buffer_data
-  @suppress_finalizing
-  def _free(self, opaque:WGPUBufPtr, options:BufferSpec): webgpu.wgpuBufferDestroy(opaque)
+    self.dev.write_buffer(dest, padded_src if src.nbytes % 4 else src)
+  def _copyout(self, dest:memoryview, src:webgpu.WGPUBuffer):
+    dest[:] = buf_to_mv(tmp_buf:=self.dev._readable_buffer(src))[:dest.nbytes]
+    self.dev.free(tmp_buf)
+
+  def _free(self, opaque:webgpu.WGPUBuffer, options:BufferSpec): self.dev.free(opaque)

 class WebGpuDevice(Compiled):
  def __init__(self, device:str):
    # Requesting an adapter
-    adapter_res = _run(webgpu.wgpuInstanceRequestAdapterF, webgpu.WGPURequestAdapterCallbackInfo, webgpu.WGPURequestAdapterCallback,
-      webgpu.enum_WGPUCreatePipelineAsyncStatus, 1, 2, instance, webgpu.WGPURequestAdapterOptions(
-        powerPreference=webgpu.WGPUPowerPreference_HighPerformance, backendType=backend_types.get(os.getenv("WEBGPU_BACKEND", ""), 0)))
+    adapter_res = InstanceRequestAdapter(instance, webgpu.WGPURequestAdapterOptions(
+      powerPreference=webgpu.WGPUPowerPreference_HighPerformance, backendType=backend_types.get(getenv("WEBGPU_BACKEND", ""), 0)))

    # Get supported features
-    supported_features = webgpu.WGPUSupportedFeatures()
-    webgpu.wgpuAdapterGetFeatures(adapter_res, supported_features)
-    supported = [supported_features.features[i] for i in range(supported_features.featureCount)]
-    features = [feat for feat in [webgpu.WGPUFeatureName_TimestampQuery, webgpu.WGPUFeatureName_ShaderF16] if feat in supported]
-    dev_desc = webgpu.WGPUDeviceDescriptor(requiredFeatureCount=len(features),
-                                           requiredFeatures=c.Array(webgpu.WGPUFeatureName, len(features))(*features)) # type: ignore
+    webgpu.wgpuAdapterGetFeatures(adapter_res, supported_features:=webgpu.WGPUSupportedFeatures())
+    self.features = [feat for i in range(supported_features.featureCount)
+                     if (feat:=supported_features.features[i]) in [webgpu.WGPUFeatureName_TimestampQuery, webgpu.WGPUFeatureName_ShaderF16]]
+    webgpu.wgpuSupportedFeaturesFreeMembers(supported_features)
+    dev_desc = webgpu.WGPUDeviceDescriptor(requiredFeatureCount=len(self.features),
+                                           requiredFeatures=(webgpu.WGPUFeatureName * len(self.features))(*self.features))

    # Limits
-    supported_limits = webgpu.WGPUSupportedLimits()
-    webgpu.wgpuAdapterGetLimits(adapter_res, ctypes.cast(ctypes.pointer(supported_limits),ctypes.POINTER(webgpu.struct_WGPUSupportedLimits)))
-    limits = webgpu.WGPURequiredLimits(limits=supported_limits.limits)
-    dev_desc.requiredLimits = c.pointer(limits)
+    webgpu.wgpuAdapterGetLimits(adapter_res, supported_limits:=webgpu.WGPUSupportedLimits())
+    dev_desc.requiredLimits = c.pointer(webgpu.WGPURequiredLimits(limits=supported_limits.limits))

    # Requesting a device
-    self.device_res = _run(webgpu.wgpuAdapterRequestDeviceF, webgpu.WGPURequestDeviceCallbackInfo, webgpu.WGPURequestDeviceCallback,
-      webgpu.enum_WGPURequestDeviceStatus, 1, 2, adapter_res, dev_desc)
+    self.device_res = AdapterRequestDevice(adapter_res, dev_desc)
+    self.queue = webgpu.wgpuDeviceGetQueue(self.device_res)

-    program = functools.partial(WebGPUProgram, (self.device_res, webgpu.WGPUFeatureName_TimestampQuery in supported))
-    super().__init__(device, WebGpuAllocator(self), [WGSLRenderer], program, arch="shader-f16" * (webgpu.WGPUFeatureName_ShaderF16 in supported))
+    webgpu.wgpuAdapterRelease(adapter_res)

-  def synchronize(self):
-    _run(webgpu.wgpuQueueOnSubmittedWorkDone2, webgpu.WGPUQueueWorkDoneCallbackInfo2, webgpu.WGPUQueueWorkDoneCallback2,
-    webgpu.enum_WGPUQueueWorkDoneStatus, None, None, webgpu.wgpuDeviceGetQueue(self.device_res))
+    super().__init__(device, WebGpuAllocator(self), [WGSLRenderer], functools.partial(WebGPUProgram, self),
+                     arch="shader-f16" * (webgpu.WGPUFeatureName_ShaderF16 in self.features))
+
+  def synchronize(self): QueueOnSubmittedWorkDone(self.queue)
+
+  @suppress_finalizing
+  def free(self, buf:webgpu.WGPUBuffer):
+    if webgpu.wgpuBufferGetMapState(buf) == webgpu.WGPUBufferMapState_Mapped: webgpu.wgpuBufferUnmap(buf)
+    webgpu.wgpuBufferDestroy(buf)
+    webgpu.wgpuBufferRelease(buf)
+
+  def pop_error(self) -> str: return DevicePopErrorScope(self.device_res)[1]
+  def create_uniform(self, val:int|float) -> webgpu.WGPUBuffer:
+    buf = webgpu.wgpuDeviceCreateBuffer(self.device_res,
+                                        webgpu.WGPUBufferDescriptor(size=4, usage=webgpu.WGPUBufferUsage_Uniform | webgpu.WGPUBufferUsage_CopyDst))
+    self.write_buffer(buf, val.to_bytes(4, "little") if isinstance(val, int) else struct.pack('<f', val))
+    return buf
+  def _readable_buffer(self, buf:webgpu.WGPUBuffer) -> webgpu.WGPUBuffer:
+    size = webgpu.wgpuBufferGetSize(buf)
+    ret = webgpu.wgpuDeviceCreateBuffer(self.device_res,
+      webgpu.WGPUBufferDescriptor(size=size, usage=webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_MapRead, mappedAtCreation=False))
+
+    # copy_buffer_to_buffer
+    encoder = webgpu.wgpuDeviceCreateCommandEncoder(self.device_res, webgpu.WGPUCommandEncoderDescriptor())
+    webgpu.wgpuCommandEncoderCopyBufferToBuffer(encoder, buf, 0, ret, 0, size)
+    cmd_buf = webgpu.wgpuCommandEncoderFinish(encoder, webgpu.WGPUCommandBufferDescriptor())
+    webgpu.wgpuQueueSubmit(self.queue, 1, (webgpu.WGPUCommandBuffer*1)(cmd_buf))
+    webgpu.wgpuCommandBufferRelease(cmd_buf)
+    webgpu.wgpuCommandEncoderRelease(encoder)
+
+    return ret
+  def write_buffer(self, buf:webgpu.WGPUBuffer, src:memoryview|bytearray|bytes):
+    webgpu.wgpuQueueWriteBuffer(self.queue, buf, 0, (ctypes.c_uint8 * len(src)).from_buffer_copy(src), len(src))
@@ -114,6 +114,8 @@ class DLL(ctypes.CDLL):

  def __init__(self, nm:str, paths:str|list[str], extra_paths=[], emsg="", **kwargs):
    self.nm, self.emsg = nm, emsg or f"try setting {nm.upper()+'_PATH'}?"
+    if nm == 'llvm' and (ver:=getenv("LLVM_VERSION", "")):
+      paths = ([f"/opt/homebrew/opt/llvm@{ver}/lib/libLLVM.dylib"] if OSX else [f"LLVM-{ver}"]) + (paths if isinstance(paths, list) else [paths])
    if (path:= DLL.findlib(nm, paths if isinstance(paths, list) else [paths], extra_paths if isinstance(extra_paths, list) else [extra_paths])):
      if DEBUG >= 3: print(f"loading {nm} from {path}")
      try:
@@ -132,6 +134,7 @@ class DLL(ctypes.CDLL):
        nonlocal cfunc
        if cfunc is None: (cfunc:=getattr(self, fn.__name__)).argtypes, cfunc.restype = argtypes, restype
        return cfunc(*args)
+      wrapper.restype, wrapper.argtypes = restype, argtypes # type: ignore
      return wrapper
    return wrap

@@ -4,7 +4,7 @@ from tinygrad.helpers import getenv, capstone_flatdump, DEBUG, unwrap
 from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.runtime.autogen import llvm

-class ClangJITCompiler(Compiler):
+class ClangCompiler(Compiler):
  def __init__(self, arch:list[str], cachekey="compile_clang_jit"):
    assert len(arch) >= 2, f"invalid arch string: {','.join(arch)!r}, expected '<arch>,<cpu>,[<feats>]' (eg. 'x86_64,znver2')"
    self.arch, cpu, *feats = arch
@@ -98,7 +98,7 @@ class CPULLVMCompiler(LLVMCompiler):
    if cpu == "native":
      cpu = ctypes.string_at(llvm.LLVMGetHostCPUName()).decode()
      featstr = (featstr + "," if featstr else "") + ctypes.string_at(llvm.LLVMGetHostCPUFeatures()).decode()
-    # +reserve-x18 here does the same thing as -ffixed-x18 in ClangJITCompiler, see comments there for why it's needed on arm osx
+    # +reserve-x18 here does the same thing as -ffixed-x18 in ClangCompiler, see comments there for why it's needed on arm osx
    super().__init__(self.arch, cpu, ('+reserve-x18,' if self.arch == "arm64" else '') + featstr, cache_key)

  def disassemble(self, lib:bytes): capstone_flatdump(lib, self.arch)
@@ -77,7 +77,7 @@ def create_new_buffer(ctx:tuple[dict[UOp, UOp], tuple[UOp, ...]], b:UOp):
  return ret

 pm_post_sched_cache = PatternMatcher([
-  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg]),
+  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg.slot]),
  # create new BUFFERs for LUNIQUE BUFFERs from rangeify
  (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), create_new_buffer),
 ])
@@ -55,7 +55,7 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
  return UOp.usum(*[c.pad(((s,numel-e),)) for (s,e),c in zip(chunks, copied_chunks)]).reshape(shape)

 def create_allreduce_function(buf:UOp, red:UOp, output:UOp|None=None) -> UOp|None:
-  if output is None: output = UOp.const(red.dtype, Invalid, red.device, red.shape).clone()
+  if output is None: output = UOp.const(red.dtype, Invalid, shape=red.shape).clone(device=red.device)
  to = red.param_like(0)
  src = buf.param_like(1)
  red = src.allreduce(red.arg, red.src[1])
@@ -128,7 +128,7 @@ def _apply_reshape(in_shape:tuple[sint,...], out_shape:tuple[sint, ...], urngs:U
@functools.cache
 def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
  match op:
-    case Ops.SHRINK:  rngs = tuple(a if off == 0 else a+off for a,(_,off) in zip(rngs, arg))
+    case Ops.SHRINK:  rngs = tuple(a if off == 0 else a+off for a,(off,_) in zip(rngs, arg))
    case Ops.PERMUTE: rngs = tuple(rngs[p] for p in argsort(arg))
    case Ops.FLIP:    rngs = tuple(((s-1)-a) if f else a for a,s,f in zip(rngs, in_shape, arg))
    case Ops.EXPAND:  rngs = tuple(a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, in_shape, arg))
@@ -136,7 +136,7 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
      # NOTE: the .where(r-s, i) is not inside the graph_rewrite so that `convert_pad_to_where_to_keep_behavior_local`
      #       wraps the pad with only the newly added valid
      rngs = tuple(r if (sz == sh and off == 0) else graph_rewrite((r >= off) & (r < (sh+off)),
-        symbolic+pm_simplify_valid, name="pad").where(r-off, UOp.invalid()) for r,sh,(sz,off) in zip(rngs, in_shape, arg))
+        symbolic+pm_simplify_valid, name="pad").where(r-off, UOp.invalid()) for r,sh,(off,sz) in zip(rngs, in_shape, arg))
    case Ops.RESHAPE:
      sink = UOp.sink(*rngs).simplify() # NOTE: this applies any commutative flips to the rngs early
      sub_array = {r:UOp.range(r.src[0], i, AxisType.PLACEHOLDER) for i,r in enumerate(sink.ranges)}
@@ -88,8 +88,8 @@ def expand_multi(root:UOp, multi:UOp):
  return multi.src[0].expand(new_shape).multi(multi.axis)

 def pad_multi(root:UOp, multi:UOp):
-  assert multi.axis is None or root.marg[multi.axis] == (multi.shape[multi.axis], 0), f"padding not supported for {root.marg=}"
-  local_pad = tuple((multi.src[0].shape[multi.axis], 0) if a == multi.axis else s for a,s in enumerate(root.marg))
+  assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]), f"padding not supported for {root.marg=}"
+  local_pad = tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.marg))
  return multi.src[0]._mop(Ops.PAD, local_pad).multi(multi.axis)

 def permute_multi(root:UOp, multi:UOp):
@@ -97,15 +97,15 @@ def permute_multi(root:UOp, multi:UOp):
  return multi.src[0].permute(root.marg).multi(root.axis)

 def shrink_multi(root:UOp, multi:UOp):
-  shard_bounds = tuple((e-s,s) for s,e in multi.bounds) if multi.axis is not None else ()
-  assert multi.axis is None or root.marg[multi.axis] == (multi.shape[multi.axis], 0) or root.marg[multi.axis] in shard_bounds, \
+  shard_bounds = tuple((s,e-s) for s,e in multi.bounds) if multi.axis is not None else ()
+  assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]) or root.marg[multi.axis] in shard_bounds, \
    f"shrinking not supported for {root.marg=}"
-  if multi.axis is not None and root.marg[multi.axis] in shard_bounds and root.marg[multi.axis] != (multi.shape[multi.axis], 0):
+  if multi.axis is not None and root.marg[multi.axis] in shard_bounds and root.marg[multi.axis] != (0, multi.shape[multi.axis]):
    # NOTE: shrink on the shard axis is only allowed when result is a single partition, denoted by the new real
    # we just copy it to all the devices, no real. this will be optimized out later
-    non_shard_shrink = tuple((multi.src[0].shape[i], 0) if i == multi.axis else s for i, s in enumerate(root.marg))
+    non_shard_shrink = tuple((0, multi.src[0].shape[i]) if i == multi.axis else s for i, s in enumerate(root.marg))
    return multi.src[0].copy_to_device(multi.device, arg=shard_bounds.index(root.marg[multi.axis]))._mop(Ops.SHRINK, non_shard_shrink)
-  local_shrink = tuple((multi.src[0].shape[multi.axis], 0) if a == multi.axis else s for a,s in enumerate(root.marg))
+  local_shrink = tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.marg))
  return multi.src[0]._mop(Ops.SHRINK, local_shrink).multi(multi.axis)

 def flip_multi(root:UOp, multi:UOp):
@@ -137,7 +137,7 @@ def rewrite_into_function(call:UOp):

 def param_to_multi(p:UOp):
  if p.axis is None: return None
-  return UOp.param(p.arg, p.dtype, p.shard_shape, p.device).multi(p.axis)
+  return UOp.param(p.arg.slot, p.dtype, p.shard_shape, p.device, p.arg.vmin_vmax, p.arg.name, p.arg.addrspace).multi(p.axis)

 # NOTE: this is the same pattern as Ops.UNROLL
 multi_pm = PatternMatcher([
@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field, replace
 import itertools
 from tinygrad.dtype import dtypes, PtrDType, AddrSpace, Invalid
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo, ParamArg
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, identity_element
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import prod, all_same, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
@@ -52,7 +52,7 @@ def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp):
    if x.op is Ops.PERMUTE: x, after = x.src[0], after.permute(argsort(x.marg))
    elif x.op is Ops.RESHAPE: x, after = x.src[0], after.reshape(x.src[0].shape)
    elif x.op is Ops.WHERE and x.src[2].base.arg == Invalid and x.src[1].op is Ops.PAD:
-      x, after = x.src[1].src[0], after.shrink(tuple((o, s+o) for (_,o),s in zip(x.src[1].marg, x.src[1].src[0].shape)))
+      x, after = x.src[1].src[0], after.shrink(tuple((o, s+o) for (o,_),s in zip(x.src[1].marg, x.src[1].src[0].shape)))
    else: break
  ctx[x] = after

@@ -130,15 +130,15 @@ def resolve_function(c:UOp, allow_param_mismatch=True) -> UOp|None:
  if c.arg.precompile: return None
  params: list[UOp] = []
  graph_rewrite(c.src[0], pm_gather_params, bottom_up=True, ctx=params, name="gather params")
-  params = sorted(params, key=lambda x: x.arg)
+  params = sorted(params, key=lambda x: x.arg.slot)
  args = c.src[1:]

  # NOTE: this isn't really needed. it's okay if there's unused args in the function
  if not allow_param_mismatch:
-    if [x.arg for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg for x in params]}")
+    if [x.arg.slot for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg.slot for x in params]}")
    if len(params) != len(args): raise TypeError(f"expected {len(params)} args, got {len(args)}")

-  dict_map = {x:args[x.arg] for x in params}
+  dict_map = {x:args[x.arg.slot] for x in params}
  for i, (p, a) in enumerate(dict_map.items()):
    if p.axis != a.axis: raise TypeError(f"arg {i} axis mismatch: expected {p.axis}, got {a.axis}")
    if p.max_shape != a.max_shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
@@ -347,11 +347,12 @@ def late_buffer_view(t:UOp, b:UOp):
    assert x.op not in GroupOp.Elementwise, "can't buffer view elementwise"
    x = x.src[0]
  x = next(u for u in x.src if u.op is Ops.INDEX)
+  assert x.op is Ops.INDEX, "must be INDEX"

  if len(shape) == 0: offset = x.src[1].arg
  else: offset = max(sum(idx.vmin for idx in x.src[1:]), 0)

-  return b.replace(src=(UOp(Ops.SLICE, t.dtype, (x.base, UOp.const(dtypes.weakint, offset)), size), b.src[1]))
+  return b.replace(src=(UOp(Ops.SLICE, t.dtype, (x.src[0], UOp.const(dtypes.weakint, offset)), size),))

 to_bufferview = PatternMatcher([
  (UPat(Ops.STAGE, src=(UPat((Ops.BITCAST, Ops.CONTIGUOUS), name="t"), UPat()), name="b"), late_buffer_view),
@@ -413,7 +414,11 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
  # NOTE: the DEFINE_LOCAL needs to be disambiguated here
  if sdtype.addrspace == AddrSpace.GLOBAL:
    buf = UOp(Ops.BUFFER, x.dtype, (UOp(Ops.LUNIQUE, arg=next(ctx)), UOp(Ops.DEVICE, arg=x.arg.device)), size)
-    do_store = buf.index(idx, dtype=sdtype).store(x.src[0]).end(*rngs)
+    if x.src[0].op is Ops.SLICE:
+      # no INDEX on SLICE, this could be cleaner
+      do_store = buf.store(x.src[0]).end(*rngs)
+    else:
+      do_store = buf.index(idx, dtype=sdtype).store(x.src[0]).end(*rngs)
    return buf.after(do_store)

  if allow_locals:
@@ -472,7 +477,7 @@ class LocalAddBufferContext:
  opts:tuple|None = None

 def debuf(ctx:LocalAddBufferContext, buf:UOp):
-  ret = UOp(Ops.PARAM, buf.dtype.ptr(prod(buf.max_shape)), arg=ctx.dg).reshape(buf.max_shape)
+  ret = UOp(Ops.PARAM, buf.dtype.ptr(prod(buf.max_shape), buf.addrspace), arg=ParamArg(ctx.dg, addrspace=buf.addrspace)).reshape(buf.max_shape)
  # if the buffer has symbolic shape, shrink the max-sized view to the actual shape
  if buf.max_shape != buf.shape: ret = ret.shrink(tuple((0, s) for s in buf.shape))
  if buf not in ctx.map: ctx.map[buf] = buf
@@ -484,7 +489,7 @@ def unbind_kernel(ctx:LocalAddBufferContext, b:UOp):
  return b.src[0]

 def handle_after(ctx:LocalAddBufferContext, after:UOp):
-  if isinstance(after.dtype, PtrDType) and after.ptrdtype.addrspace == AddrSpace.LOCAL: return None
+  if isinstance(after.dtype, PtrDType) and after.addrspace == AddrSpace.LOCAL: return None
  buf = after.buf_uop
  # HACK to put the buffer in the MAP instead of MSTACK/MSELECT
  if buf.op in {Ops.MSTACK, Ops.MSELECT}: buf = buf.src[0]
@@ -507,9 +512,11 @@ def find_bufs(x:UOp):
 to_define_global = PatternMatcher([
  (UPat(Ops.STORE, name="x"), find_bufs),
  (UPat(Ops.BUFFER, name="buf"), debuf),
-  (UPat(Ops.PARAM, src=(UPat(), UPat(Ops.DEVICE)), name="buf"), debuf),
-  (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat.cvar('vmin'), UPat.cvar('vmax'), UPat.var("nm")), name="v"),
-   lambda v, vmin, vmax, nm: UOp.variable(nm.arg, vmin.arg, vmax.arg, v.dtype)),
+  (UPat(Ops.PARAM, name="v"), lambda v:
+   UOp.variable(v.arg.name, v.arg.vmin_vmax[0], v.arg.vmin_vmax[1], v.dtype)
+   if v.arg.name is not None and v.arg.vmin_vmax is not None else None),
+  (UPat(Ops.PARAM, name="buf"), lambda ctx, buf:
+   None if isinstance(buf.dtype, PtrDType) or buf.arg.name is not None or buf._shape is None else debuf(ctx, buf)),
  (UPat(Ops.INDEX, src=(UPat(Ops.DEFINE_VAR, name="v"),)), lambda v: v),

  (UPat(Ops.BIND, name="b"), unbind_kernel),
@@ -160,7 +160,7 @@ class Tensor(OpMixin):
  def const_like(self, b:ConstType) -> Tensor: return Tensor(self.uop.const_like(b))
  @staticmethod
  def const(dtype:DType, b:ConstType|UOp, device:str|tuple[str, ...]|None=None) -> Tensor:
-    return Tensor(b if isinstance(b, UOp) else UOp.const(dtype, b, device))
+    return Tensor(UOp.const(dtype, b, device))
  @staticmethod
  def unique_const(fill_value:ConstType|UOp, **kwargs) -> Tensor:
    if isinstance(fill_value, UOp): return Tensor(fill_value, **kwargs)
@@ -203,7 +203,7 @@ class Tensor(OpMixin):

  def as_param(self, slot:int):
    if self.uop.axis is not None:
-      param = UOp.param(slot, self.dtype, self.uop.shard_shape, self.device).multi(self.uop.axis)
+      param = UOp.param(slot, self.dtype, self.uop.shard_shape, self.device, axis=self.uop.axis)
    else:
      param = UOp.param(slot, self.dtype, self.shape, self.device)
    return Tensor(param)
@@ -580,7 +580,7 @@ class Tensor(OpMixin):

  def _multi_like(self, fxn, *args, **kwargs) -> Tensor:
    dtype = kwargs.pop("dtype", self.dtype)
-    if kwargs.get("device") is not None: raise RuntimeError("cannot specify `device` on `*_like` of a multi device tensor")
+    if kwargs.pop("device", None) is not None: raise RuntimeError("cannot specify `device` on `*_like` of a multi device tensor")
    assert isinstance(self.device, tuple), f"_multi_like needs a multi device tensor, got {self.device}"
    if self.uop.axis is None: return fxn(self.shape, *args, dtype=dtype, **kwargs).shard(self.device)
    stacked = UOp.mstack(*[fxn(self.uop.shard_shape, *args, device=d, dtype=dtype, **kwargs).uop for d in self.device])
@@ -598,9 +598,8 @@ class Tensor(OpMixin):
    print(Tensor.full_like(t, 42).numpy())
    ```
    """
-    if device is None: return super().full_like(fill_value, dtype)
-    if isinstance(self.device, tuple): raise RuntimeError("cannot specify `device` on `full_like` of a multi device tensor")
-    return Tensor.full(self.shape, fill_value, dtype=dtype or self.dtype, device=device)
+    if isinstance(self.device, tuple): return self._multi_like(Tensor.full, fill_value, dtype=dtype or self.dtype, device=device)
+    return Tensor.full(self.shape, fill_value, dtype=dtype or self.dtype, device=self.device if device is None else device)

  def rand_like(self, **kwargs) -> Tensor:
    """
@@ -852,6 +851,7 @@ class Tensor(OpMixin):
    # clear contexts
    for t,g in zip(tensors_need_grad, self.gradient(*tensors_need_grad, gradient=gradient)):
      assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
+      if g.device is None and t.device is not None: g = g.clone(device=t.device)
      if t.grad is None: t.grad = g
      else: t.grad.assign(t.grad + g.to(t.grad.device))
    return self
@@ -117,7 +117,7 @@ class GroupOp:
            Ops.XOR, Ops.SHL, Ops.SHR, Ops.OR, Ops.AND, Ops.THREEFRY, Ops.SUB, Ops.FDIV, Ops.POW, Ops.FLOORDIV, Ops.FLOORMOD}
  Ternary = {Ops.WHERE, Ops.MULACC}
  ALU = set.union(Unary, Binary, Ternary)
-  Broadcastable = set.union(Binary, Ternary, {Ops.GROUP, Ops.STORE})
+  Broadcastable = set.union(Binary, Ternary, {Ops.GROUP})

  # TODO: is BITCAST always Elementwise if it's shape changing?
  Elementwise = set.union(ALU, {Ops.CAST, Ops.BITCAST})
@@ -477,7 +477,7 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa
        if (v:=powers_of_two.get(c.arg, 0)) else None)]
    if not disable_fast_idiv:
      # fast_idiv handles non-pow2: only fire on non-negative inputs (signed magic-mul is unreliable for x<0)
-      pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d", vec=False))),
+      pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))),
        lambda ctx, x, d: fast_idiv(ctx, x, d.arg) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
      # rewrite raw CMOD -> x - d*CDIV(x,d) so fast_idiv can pick up the CDIV. only on non-negative inputs;
      # avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.CMOD as an implementation detail)
@@ -493,7 +493,7 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa
      ((UPat.cvar("c", dtypes.sints) < UPat.var("x", dtypes.sints)).logical_not(), lambda x,c: x<c+1),
      (UPat.var("x", dtypes.sints)*-1 < UPat.var("y", dtypes.sints)*UPat.cvar("c"), lambda x,y,c: y*(-c)<x),
      (UPat.var("x", dtypes.sints)*-1 < UPat.cvar("c"), lambda x,c:-c<x),
-      ((UPat.cvar("c1",vec=False)<UPat.var("x", dtypes.sints)) & (UPat.var("x", dtypes.sints)<UPat.cvar("c2",vec=False)),
+      ((UPat.cvar("c1")<UPat.var("x", dtypes.sints)) & (UPat.var("x", dtypes.sints)<UPat.cvar("c2")),
        lambda x,c1,c2: x.eq(c1+1) if c1.arg+1==c2.arg-1 else None),  # (c-1)<x & x<(c+1) -> x==c
    ]
  if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))]
@@ -108,7 +108,7 @@ div_and_mod_symbolic = PatternMatcher([
  # (x//c+a)//d -> (x+a*c)//(c*d) for c>0, d>0
  ((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d) if c.vmin>0 and d.vmin>0 else None),
  # (x+c)//d -> (x+c%d)//d + c//d for d>0 (split out the multiple of d in the constant)
-  ((UPat.var("x", dtypes.weakint)+UPat.cvar("c", vec=False))//UPat.cvar("d", vec=False),
+  ((UPat.var("x", dtypes.weakint)+UPat.cvar("c"))//UPat.cvar("d"),
    lambda x,c,d: (x+c.arg%d.arg)//d + c.arg//d.arg if c.arg%d.arg!=c.arg and d.arg>0 else None),

  # ** 2. Slow Rules **
@@ -18,6 +18,19 @@ class AxisType(Enum):
  def __repr__(self): return str(self)
  GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
  THREAD = auto(); PLACEHOLDER = auto() # noqa: E702
+
+@dataclass(frozen=True, order=True)
+class ParamArg:
+  slot: int
+  vmin_vmax: tuple[PyConst, PyConst]|None = None
+  name: str|None = None
+  addrspace: AddrSpace = AddrSpace.GLOBAL
+  axis: int|None = None
+  device: str|tuple[str, ...]|None = None
+  def __repr__(self):
+    fields = (("vmin_vmax", None), ("name", None), ("addrspace", AddrSpace.GLOBAL), ("axis", None), ("device", None))
+    args = [str(self.slot)] + [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default]
+    return f"ParamArg({', '.join(args)})"
 axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
 axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
@@ -275,22 +288,23 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
        # STAGE adds the existing shape to the front, opposite of INDEX
        return tuple([int(r.vmax+1) for r in self.src[1:]])+self.src[0].shape
      case Ops.DEFINE_LOCAL | Ops.DEFINE_REG:
+        if len(self.src) >= 1:
+          # NOTE: this is the same as PARAM
+          return tuple(self.src[0].sgep(i) for i in range(self.src[0].dtype.count))
        if isinstance(self.dtype, PtrDType):
          return (self.ptrdtype.size, self.dtype.count) if self.dtype.count > 1 else (self.ptrdtype.size,)
        return (self.dtype.count,) if self.dtype.count > 1 else ()
      case Ops.PARAM:
        if isinstance(self.dtype, ImageDType): return self.dtype.shape
        if isinstance(self.dtype, PtrDType): return (self.ptrdtype.size,)
-        # NOTE: copied from marg
-        if len(self.src) >= 1: return tuple(self.src[0].sgep(i) for i in range(self.src[0].dtype.count))
-        return None
+        return tuple(self.src[0].sgep(i) for i in range(self.src[0].dtype.count)) if len(self.src) >= 1 else None

      # wmma output shape = accumulator shape (src[2])
      case Ops.WMMA | Ops.SHAPED_WMMA: return self.src[2]._shape

      # passthrough ops
      case Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.LOAD | \
-           Ops.COPY | Ops.ALLREDUCE:
+           Ops.COPY | Ops.ALLREDUCE | Ops.STORE:
        return self.src[0]._shape
      # REDUCE with empty axis is passthrough (lowered form)
      case Ops.REDUCE if len(self.arg[1]) == 0:
@@ -305,7 +319,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
          return ps[:-1]+(ssimplify((ps[-1]*input_sz) // output_sz),) if len(ps) > 0 else ps
        return ps

-      # MULTI marker (axis info in PARAM sources) has no shape
+      # MULTI marker has no shape
      case Ops.MULTI if len(self.src) == 0: return None

    # movement ops change the shape
@@ -327,20 +341,20 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
          return tuple(ps[i] for i in self.marg)
        case Ops.PAD:
          # TODO: why do i need resolve here?
-          if len(ps) != len(self.marg) or not all(resolve(sz>=0) and resolve(0<=o) and resolve(o+s<=sz) for s,(sz,o) in zip(ps, self.marg)):
+          if len(ps) != len(self.marg) or not all(resolve(sz>=0) and resolve(0<=o) and resolve(o+s<=sz) for s,(o,sz) in zip(ps, self.marg)):
            raise ValueError(f"invalid pad {self.marg} for {ps}")
-          return tuple(ssimplify(sz) for sz,_ in self.marg)
+          return tuple(ssimplify(sz) for _,sz in self.marg)
        case Ops.SHRINK:
          # TODO: why do i need resolve here?
-          if len(ps) != len(self.marg) or not all(resolve(0<=b) and resolve(sz>=0) and resolve(b+sz<=s) for s,(sz,b) in zip(ps, self.marg)):
+          if len(ps) != len(self.marg) or not all(resolve(0<=o) and resolve(sz>=0) and resolve(o+sz<=s) for s,(o,sz) in zip(ps, self.marg)):
            raise ValueError(f"invalid shrink {self.marg} for {ps}")
-          return tuple(ssimplify(sz) for sz,_ in self.marg)
+          return tuple(ssimplify(sz) for _,sz in self.marg)
        case Ops.FLIP:
          if len(ps) != len(self.marg) or not all(isinstance(x, bool) for x in self.marg): raise ValueError(f"bad flip on {ps}, {self.marg}")
          return ps
        case Ops.MULTI: return tuple(s*len(self.device) if a == self.axis else s for a,s in enumerate(ps))
        case Ops.REDUCE:
-          axis_arg = self.arg[1] if self.op is Ops.REDUCE else self.arg[7]
+          axis_arg = self.arg[1]
          if not isinstance(axis_arg, tuple) or not all(isinstance(x, int) and x>=0 and x<len(ps) for x in axis_arg):
            raise ValueError(f"invalid type for axis: {axis_arg}")
          return tuple(1 if i in axis_arg else s for i,s in enumerate(ps))
@@ -368,6 +382,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):

  @property
  def max_shape(self) -> tuple[int, ...]: return to_max_shape(self.shape)
+  def max_numel(self) -> int: return prod(self.max_shape)

  @property
  def shard_shape(self) -> tuple[sint, ...]:
@@ -477,8 +492,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
      return perm.index(*non_slice_args, ptr=True)
    return self.index(*[UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in idx])
  def const_like(self, b:ConstLike, dtype:DType|None=None):
-    # constants can optionally have a DEVICE source
-    ret = UOp.const(dtype or self.dtype.base, b, device=self.device, shape=self.shard_shape if self.axis is not None else self._shape)
+    # multi constants can optionally have a DEVICE source  # TODO: no const with DEVICE
+    dev = self.device if isinstance(self.device, tuple) else None
+    ret = UOp.const(dtype or self.dtype.base, b, device=dev, shape=self.shard_shape if self.axis is not None else self._shape)
    return ret.multi(self.axis) if self.axis is not None else ret
  def ufix(self, x):
    if isinstance(x, UOp): return x
@@ -525,7 +541,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
    return UOp(op, out_dtype, all_srcs, **kwargs)
  @staticmethod
  def const(dtype:DType, b:ConstLike, device:str|tuple[str, ...]|None=None, shape:tuple[sint, ...]|None=None):
-    if isinstance(b, UOp): return b.unbind()[0] if b.op is Ops.BIND else b
+    if isinstance(b, UOp): return b.cast(dtype)
    if isinstance(b, tuple) and all_same(b):
      assert len(b) > 0, "can't create const from empty tuple"
      b = b[0]  # doesn't have to be a STACK if they are all the same
@@ -554,7 +570,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
  @staticmethod
  def invalid(count=1): return UOp(Ops.CONST, dtypes.weakint.vec(count), src=(), arg=Invalid)
  def valid(self, cond):
-    return self if cond.op is Ops.WHERE and cond.arg else cond.where(self.cast(dtypes.weakint), UOp.invalid(self.dtype.count))
+    return cond.where(self.cast(dtypes.weakint), UOp.invalid(self.dtype.count))
  def get_idx(self) -> UOp:
    assert self.dtype.scalar() is dtypes.weakint, "Can only call get_idx on index dtype"
    return self.src[1] if self.op is Ops.WHERE and self.src[2].arg is Invalid else self
@@ -610,16 +626,12 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
    if self.op is Ops.GETTUPLE:
      in_tuple = self.src[0].src[0] if self.src[0].op is Ops.FUNCTION else self.src[0]
      return in_tuple.src[self.arg].axis if in_tuple.op is Ops.TUPLE else None
-    # PARAM: axis is stored as a MULTI source
-    if self.op is Ops.PARAM:
-      for s in self.src:
-        if s.op is Ops.MULTI: return s.arg
-      return None
+    if self.op is Ops.PARAM: return self.arg.axis
    # NOTE: they all have to share an axis, we always choose [-1]
    if self.op in GroupOp.ALU: return axes[-1] if (axes := dedup([x.axis for x in self.src if x.axis is not None])) else None
    if len(self.src) == 0: return None
    src_axis = self.src[0].axis
-    if self.op is Ops.SHRINK and src_axis is not None and self.marg[src_axis] != (self.src[0].shape[src_axis], 0):
+    if self.op is Ops.SHRINK and src_axis is not None and self.marg[src_axis] != (0, self.src[0].shape[src_axis]):
      return None # SHRINK will remove the sharding if it's on axis
    if self.op is Ops.REDUCE: return None if src_axis is not None and src_axis in self.arg[1] else src_axis
    if self.op is Ops.RESHAPE:
@@ -680,7 +692,11 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
  def marg(self):
    match self.op:
      case Ops.RESHAPE | Ops.EXPAND: return tuple(ssimplify(self.src[1].sgep(i)) for i in range(self.src[1].dtype.count))
-      case Ops.PAD | Ops.SHRINK: return tuple((self.src[1].sgep(i), self.src[2].sgep(i)) for i in range(self.src[1].dtype.count))
+      case Ops.PAD | Ops.SHRINK:
+        # this is like broadcasting for shapes
+        return tuple(((ssimplify(self.src[1]) if self.src[1].shape == () else self.src[1].sgep(i)),
+                      (ssimplify(self.src[2]) if self.src[2].shape == () else self.src[2].sgep(i)))
+                     for i in range(max(self.src[1].dtype.count, self.src[2].dtype.count)))
      case Ops.PERMUTE | Ops.FLIP: return self.arg
      case _: raise RuntimeError(f"{self.op} is not a MovementOp")

@@ -731,6 +747,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
    return ret.after(ret.store(src))
  @recursive_property
  def device(self) -> str|tuple[str, ...]|None:
+    if self.op is Ops.PARAM: return self.arg.device
    if self.op is Ops.DEVICE: return self.arg
    if self.op is Ops.STAGE: return self.arg.device
    if self.op is Ops.AFTER: return self.src[0].device
@@ -742,6 +759,21 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
    for x in self.src:
      if x.device is not None: return x.device
    return None
+  @recursive_property
+  def addrspace(self) -> AddrSpace|None:
+    if self.op is Ops.PARAM: return self.arg.addrspace
+    if self.op is Ops.BUFFER: return AddrSpace.GLOBAL
+    if self.op is Ops.DEFINE_LOCAL: return AddrSpace.LOCAL
+    if self.op is Ops.DEFINE_REG: return AddrSpace.REG
+    if self.op is Ops.LOAD: return AddrSpace.ANON # LOAD brings things into anonymous registers
+    if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.GEP, Ops.STORE}:
+      return self.src[0].addrspace
+    if self.op in GroupOp.Movement: return self.src[0].addrspace
+    if self.op in {Ops.STACK, Ops.WMMA} or self.op in GroupOp.Elementwise:
+      ad = [x.addrspace for x in self.src if x.addrspace is not None]
+      if not len(ad) or not all_same(ad): return None
+      return ad[0]
+    return None
  @property
  def buf_uop(self) -> UOp:
    if self.op in {Ops.BUFFER, Ops.PARAM}: return self
@@ -949,7 +981,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
    # float has NAN issue and we use explicit NAN in transcendental
    if self.op is Ops.WHERE and dtypes.is_int(self.dtype): return min(self.src[1].vmin, self.src[2].vmin), max(self.src[1].vmax, self.src[2].vmax)
    # NOTE: returned UOp is assumed to be CONST
-    if self.op is Ops.PARAM and len(self.src) >= 4: return self.src[2].arg, self.src[3].arg
+    if self.op is Ops.PARAM and self.arg.vmin_vmax is not None: return self.arg.vmin_vmax
    if self.op is Ops.DEFINE_VAR and self.arg: return self.arg[1], self.arg[2]
    if self.op in (Ops.RANGE, Ops.SPECIAL): return 0, (self.src[0]-1).vmax
    if self.op is Ops.BIND: return self.src[0]._min_max # ignore the bound value
@@ -994,7 +1026,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
  @staticmethod
  def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL):
    lookup = {AddrSpace.GLOBAL: Ops.PARAM, AddrSpace.LOCAL: Ops.DEFINE_LOCAL, AddrSpace.REG: Ops.DEFINE_REG}
-    ret = UOp(lookup[addrspace], dtype.ptr(prod(shape), addrspace), arg=slot)
+    arg = ParamArg(slot, addrspace=addrspace) if addrspace is AddrSpace.GLOBAL else slot
+    ret = UOp(lookup[addrspace], dtype.ptr(prod(shape), addrspace), arg=arg)
    if len(shape) > 1: ret = ret.reshape(shape)
    return ret
  def placeholder_like(self, slot:int):
@@ -1007,18 +1040,17 @@ class UOp(OpMixin, metaclass=UOpMetaClass):

  # TODO: this should replace placeholder
  @staticmethod
-  def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None, name=None):
-    src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),) + \
-                           (UOp(Ops.NOOP) if device is None else UOp(Ops.DEVICE, arg=device),)
-    if vmin_vmax is not None: src += (UOp.const(dtype, vmin_vmax[0]), UOp.const(dtype.scalar(), vmin_vmax[1]))
-    if name is not None: src += (UOp(Ops.NOOP, arg=name),)
-    return UOp(Ops.PARAM, dtype, src, arg=slot)
+  def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None, name=None,
+            addrspace=AddrSpace.GLOBAL, axis:int|None=None):
+    if shape is not None and axis is not None and isinstance(device, tuple):
+      shape = tuple(s*len(device) if i == axis else s for i,s in enumerate(shape))
+    src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),)
+    return UOp(Ops.PARAM, dtype, src, arg=ParamArg(slot, vmin_vmax, name, addrspace, axis, device))
  def param_like(self, slot:int):
+    addrspace = self.addrspace if isinstance(self.dtype, (PtrDType, ImageDType)) else AddrSpace.GLOBAL
    if self.op is Ops.BIND:
-      return UOp.param(slot, self.dtype, self._shape, self.device, self._min_max, self.src[0].arg[0])
-    p = UOp.param(slot, self.dtype, self._shape, self.device)
-    if self.axis is not None: p = p.replace(src=p.src + (UOp(Ops.MULTI, arg=self.axis),))
-    return p
+      return UOp.param(slot, self.dtype, self._shape, self.device, cast(tuple[int, int], self._min_max), self.src[0].arg[0], addrspace)
+    return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis)

  # opaque bodies stay as Ops.CALL; value-producing bodies become Ops.FUNCTION (wrapped in TUPLE)
  _OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.SLICE, Ops.CUSTOM_FUNCTION}
@@ -1082,10 +1114,10 @@ class ProgramInfo:
    local_size: list[int]|None = [1, 1, 1]
    for u in sink.toposort():
      if u.op is Ops.DEFINE_VAR: _vars.append(u)
-      if u.op is Ops.PARAM: _globals.append(u.arg)
+      if u.op is Ops.PARAM: _globals.append(u.arg.slot)
      if u.op in (Ops.STORE, Ops.LOAD):
-        if (idx:=u.src[0]).op is Ops.INDEX or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX):
-          if (buf:=idx.src[0]).op is Ops.PARAM: (outs if u.op is Ops.STORE else ins).append(buf.arg)
+        if (idx:=u.src[0]).op in (Ops.INDEX, Ops.SHRINK) or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX):
+          if (buf:=idx.src[0]).op is Ops.PARAM: (outs if u.op is Ops.STORE else ins).append(buf.arg.slot)
      if u.op is Ops.SPECIAL:
        if u.arg[0] == 'i': local_size = None
        special_size = local_size if u.arg[0] == 'l' else global_size
@@ -1211,8 +1243,7 @@ class UPat(OpMixin):
  def var(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None): return UPat(dtype=dtype, name=name)
  @staticmethod
  @functools.cache
-  def cvar(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None, vec=True, arg=None):
-    return UPat(Ops.CONST, dtype, name=name, arg=arg)
+  def cvar(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None, arg=None): return UPat(Ops.CONST, dtype, name=name, arg=arg)
  @staticmethod
  def const(dtype:DType|tuple[DType, ...]|None, b:ConstType, device=None): return UPat(Ops.CONST, dtype=dtype, arg=b)

@@ -1630,7 +1661,7 @@ pm_lower_index_dtype = PatternMatcher([
 def _index_to_concrete_int(u:UOp) -> UOp: return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]

 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])
-_pm_resolve_params = PatternMatcher([(UPat(Ops.PARAM, name="p"), lambda ctx,p: ctx[p.arg])])
+_pm_resolve_params = PatternMatcher([(UPat(Ops.PARAM, name="p"), lambda ctx,p: ctx[p.arg.slot])])
 remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])

 def gate_kernel_sink(x:UOp) -> bool:
@@ -34,7 +34,7 @@ def strip_binary_parens(x:UOp, left:str, right:str, code_for_op) -> str:

 renderer = PatternMatcher([
  (UPat((Ops.DEFINE_VAR,), name="x"), lambda x: x.expr),
-  (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.NOOP, name="x"))), lambda x: x.arg),
+  (UPat(Ops.PARAM, name="x"), lambda x: x.arg.name if x.arg.name is not None else f"p{x.arg.slot}"),
  (UPat((Ops.SPECIAL), name="x"), lambda x: x.arg),
  (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"),
  (UPat(Ops.CONST, name="x"), lambda x: str(x.arg)),
@@ -1,9 +1,9 @@
 import math
 from typing import cast, Any
-from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, AxisType, KernelInfo
+from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, AxisType, KernelInfo, ParamArg
 from tinygrad.uop.render import print_uops, pyrender
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid, ConstFloat
-from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata, panic, CHECK_OOB
+from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata, panic, CHECK_OOB, all_same

 # ***** uop helpers *****

@@ -13,7 +13,10 @@ def validate_index(uidx:UOp, gate:UOp|None=None):
  if idx.op is Ops.CONST and idx.arg is Invalid: return True
  if gate is None: gate = UOp.const(dtypes.bool, True)
  # TODO: check for overflow
-  if not CHECK_OOB or isinstance(buf.dtype, ImageDType) or (sz := buf.ptrdtype.size) == -1: return True
+  if not CHECK_OOB or isinstance(buf.dtype, ImageDType): return True
+
+  # buffer size
+  sz = buf.max_numel()

  # We can use UOp min/max to do a faster check, but it can give false positive since its not an exact bound and doesn't consider the mask
  if 0<=idx.vmin and idx.vmax<sz: return True
@@ -70,16 +73,15 @@ spec_shared = PatternMatcher([
  (UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(dtypes.is_int(y.dtype) for y in x.src[1:]) or None),
  (UPat(Ops.END, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(u.op is Ops.RANGE for u in x.src[1:])),

-  # PARAM (that's really a DEFINE_GLOBAL)
-  (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.dtype, (PtrDType, ImageDType)) and x.dtype.addrspace == AddrSpace.GLOBAL),
+  # PARAM
+  (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)),

  # GROUP of stores (or groups, or NOOPs)
  # TODO: remove UNROLL here, it's for SPEC=2
  (UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.UNROLL, Ops.INS))), lambda: True),

  # TOOD: these should be buffer with different addrspace
-  (UPat(Ops.DEFINE_LOCAL, name="x"), lambda x: isinstance(x.dtype, PtrDType) and x.dtype.addrspace == AddrSpace.LOCAL),
-  (UPat(Ops.DEFINE_REG, src=()), lambda: True),
+  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)), lambda: True),

  # AFTER on Movement Op, PARAM, BUFFER, CONTIGUOUS, or another AFTER
  (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.BUFFER, Ops.CONTIGUOUS, Ops.DEFINE_REG, Ops.DEFINE_LOCAL, Ops.AFTER, Ops.MULTI,
@@ -99,11 +101,11 @@ spec_shared = PatternMatcher([
  (UPat(Ops.INS), lambda: True),

  # LOAD(idx) / STORE(idx, val) with gates on the LOAD/STORE
-  (UPat(Ops.INDEX, name="uidx").or_casted().load(), validate_index),
-  (UPat(Ops.INDEX, name="uidx").or_casted().load(UPat.var("alt"), UPat.var("gate", dtype=dtypes.bool), name="load"),
+  (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().load(), validate_index),
+  (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().load(UPat.var("alt"), UPat.var("gate", dtype=dtypes.bool), name="load"),
   lambda uidx,gate,alt,load: validate_index(uidx, gate) if alt.dtype == load.dtype else False),
-  (UPat(Ops.INDEX, name="uidx").or_casted().store(UPat()), validate_index),
-  (UPat(Ops.INDEX, name="uidx").or_casted().store(UPat(), UPat.var("gate", dtype=dtypes.bool)), validate_index),
+  (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().store(UPat()), validate_index),
+  (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().store(UPat(), UPat.var("gate", dtype=dtypes.bool)), validate_index),

  # STORE in tensor graph: store a value into a target
  (UPat(Ops.STORE, dtypes.void, (UPat(name="x"), UPat())), lambda x: True),
@@ -124,15 +126,12 @@ spec_tensor = PatternMatcher([

  # CONST with a UNIQUE or DEVICE
  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE),)), lambda: True),
-  (UPat(Ops.CONST, src=(UPat((Ops.UNIQUE, Ops.LUNIQUE)), UPat(Ops.DEVICE))), lambda: True),
+  (UPat(Ops.CONST, src=(UPat((Ops.UNIQUE, Ops.LUNIQUE)), UPat(Ops.DEVICE)), name="c"), lambda c: c.arg is Invalid),

  # BUFFER
  (UPat(Ops.BUFFER, src=(UPat((Ops.UNIQUE, Ops.LUNIQUE)), UPat(Ops.DEVICE)), name="buf"),
   lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, DType)),

-  # PARAM (that's really a variable)
-  (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat(), UPat(), UPat()), name="x"), lambda x: True),
-
  # Tensor variable bindings
  (UPat(Ops.BIND, (dtypes.int, dtypes.weakint,), (UPat(Ops.DEFINE_VAR), UPat.cvar(dtype=(dtypes.int,dtypes.weakint,))), arg=None), lambda: True),

@@ -147,11 +146,6 @@ spec_tensor = PatternMatcher([
  (UPat(Ops.TUPLE, dtypes.void), lambda: True),
  (UPat(Ops.GETTUPLE, src=(UPat((Ops.FUNCTION, Ops.TUPLE)),), name="g"), lambda g: isinstance(g.arg, int)),

-  # PARAM
-  (UPat(Ops.PARAM, src=(UPat(), UPat(Ops.NOOP)), name="x"), lambda x: True),  # TODO: why does this have NOOP?
-  (UPat(Ops.PARAM, src=(UPat(), UPat(Ops.DEVICE)), name="x"), lambda x: True),
-  (UPat(Ops.PARAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.MULTI)), name="x"), lambda x: True),
-
  # inputs to movement ops
  (UPat(Ops.STACK), lambda: True),
  (UPat({Ops.ADD, Ops.MUL, Ops.CDIV, Ops.FLOORDIV}, dtype=dtypes.weakint), lambda: True),
@@ -174,7 +168,7 @@ spec_tensor = PatternMatcher([
  # MULTI/MSELECT/MSTACK
  (UPat(Ops.MULTI, name="multi"), lambda multi: all(x.dtype == multi.dtype for x in multi.src) and isinstance(multi.arg, int)),
  (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)),
-  (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(x.device, str) for x in x.src)),
+  (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)),

  # CONTIGUOUS ensures the source UOp realizes
  (UPat((Ops.DETACH, Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD), name="root", src=(UPat.var("x"),), arg=None),
@@ -202,6 +196,9 @@ spec_program = PatternMatcher([
  # weakint is not allowed in programs
  (UPat(GroupOp.All, dtypes.weakint), lambda: False),

+  # allow special SHRINK
+  (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.AFTER)), UPat(), UPat(Ops.CONST))), lambda: True),
+
  # movement ops are not allowed in programs
  (UPat(GroupOp.Movement), lambda: False),

@@ -213,23 +210,21 @@ spec_program = PatternMatcher([
   lambda x: False if x.dtype.count > 1 and (x.dtype.count,) != x.shape else None),

  # STACK/GEP in program. TODO: this should match Tensor
-  (UPat(Ops.STACK, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.vcount and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)),
+  (UPat(Ops.STACK, name="x"), lambda x: len(x.src)>1),
  (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()),

  # if has a <gate, index_for_dedup>
-  (UPat(Ops.IF, dtype=dtypes.void, src=(UPat(dtype=dtypes.bool), UPat((Ops.CAST, Ops.INDEX)))), lambda: True),
+  (UPat(Ops.IF, dtype=dtypes.void, src=(UPat(dtype=dtypes.bool), UPat((Ops.CAST, Ops.INDEX, Ops.SHRINK)))), lambda: True),
  (UPat(Ops.ENDIF, dtype=dtypes.void, src=(UPat(Ops.IF),)), lambda: True),
 ])+spec_shared

 # these are intermediate ops. everything should be deleted from here
 spec_full = PatternMatcher([
  # SLICE on BUFFER is allowed if BUFFER is
-  (UPat(Ops.SLICE, src=(UPat((Ops.BUFFER, Ops.PARAM)), UPat(Ops.CONST, dtype=dtypes.weakint)), allow_any_len=True, name="bv"),
+  (UPat(Ops.SLICE, src=(UPat(GroupOp.Movement.union({Ops.BUFFER, Ops.PARAM, Ops.STAGE, Ops.AFTER})),
+                        UPat(Ops.CONST, dtype=dtypes.weakint)), allow_any_len=True, name="bv"),
   lambda bv: isinstance(bv.arg, int)),

-  # TODO: SLICE shouldn't go on INDEX. why is this allowed? remove these both
-  (UPat(Ops.SLICE, src=(UPat((Ops.INDEX,)), UPat(Ops.CONST, dtype=dtypes.weakint)), allow_any_len=True, name="bv"),
-   lambda bv: isinstance(bv.arg, int)),
  (UPat(Ops.CALL, src=(UPat((Ops.SLICE,)),), allow_any_len=True), lambda: True),

  # codegen may end ranges after gpudims has replaced RANGE with SPECIAL.
@@ -266,7 +261,7 @@ from tinygrad.schedule.rangeify import BufferizeOpts
 glbls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Metadata": Metadata,
                        "UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid,
                        "Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace, "panic": panic,
-                        "ConstFloat": ConstFloat}
+                        "ConstFloat": ConstFloat, "ParamArg": ParamArg}
 def eval_pyrender(code:str) -> UOp:
  lcls:dict[str, Any] = {}
  exec(code, glbls, lcls)
@@ -98,14 +98,14 @@ symbolic_simple = propagate_invalid + PatternMatcher([
  ((UPat.var() % UPat.var("y")).named("base") % UPat.var("y"), lambda base,y: base),  # (x%y)%y = -> x%y (rewritten with base for speed)
  # variations of (x%c)+(x//c)*c = x
  (UPat(Ops.ADD, dtype=dtypes.weakint, name="x"), fold_add_divmod_recombine),
-  (UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c", vec=False), lambda x,c: x if c.arg else c),
-  (UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c", vec=False), lambda x,c: c if c.arg else x),
+  (UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c"), lambda x,c: x if c.arg else c),
+  (UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c"), lambda x,c: c if c.arg else x),
  (UPat(GroupOp.Idempotent, src=(UPat.var("x"), UPat.var("x"))), lambda x: x),
  (UPat.var("x", dtype=dtypes.bool).logical_not().logical_not(), lambda x: x),
  (UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, True), UPat.const(dtypes.bool, False)), lambda x: x),
  (UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, False), UPat.const(dtypes.bool, True)), lambda x: x.logical_not()),
  # CAST(bool -> int) != const — CAST(True)=1, CAST(False)=0, so fold based on const value
-  (UPat.var("x", dtype=dtypes.bool).cast(dtypes.ints+(dtypes.weakint,)) != UPat.cvar("c", vec=False),
+  (UPat.var("x", dtype=dtypes.bool).cast(dtypes.ints+(dtypes.weakint,)) != UPat.cvar("c"),
   lambda x,c: x if c.arg == 0 else x.logical_not() if c.arg == 1 else x.const_like(True)),
  (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)).trunc(), lambda x: x),
  # ** zero folding **
@@ -115,7 +115,7 @@ symbolic_simple = propagate_invalid + PatternMatcher([
  (UPat.var("x") & 0, lambda x: x.const_like(0)), # x&0 -> 0
  # (x&mask)>>k -> x>>k when mask only clears bits below k
  # TODO: combine this with "# rules for threefry" below
-  ((UPat.var("x") & UPat.cvar("mask", vec=False)) >> UPat.cvar("k", vec=False),
+  ((UPat.var("x") & UPat.cvar("mask")) >> UPat.cvar("k"),
   lambda x,mask,k: x >> k.arg if mask.arg | ((1 << k.arg) - 1) == -1 else None),
  (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) != UPat.var("x"),
   lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x != x -> False (only ints)
@@ -148,9 +148,9 @@ symbolic_simple = propagate_invalid + PatternMatcher([
  (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x if x.dtype == b.dtype and can_lossless_cast(b.dtype, a.dtype) else None),
  (UPat.var("x").cast(dtypes.bool), lambda x: x != 0),
  # ** pow **
-  (UPat.var("x").alu(Ops.POW, UPat.cvar("c", vec=False)), simplify_pow),
+  (UPat.var("x").alu(Ops.POW, UPat.cvar("c")), simplify_pow),
  # positive const ** x
-  (UPat.cvar("c", vec=False).alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.arg == 1 else (x*math.log2(c.arg)).exp2() if c.arg > 0 else None),
+  (UPat.cvar("c").alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.arg == 1 else (x*math.log2(c.arg)).exp2() if c.arg > 0 else None),
  # rules for threefry
  ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)),
  (((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y',  dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
@@ -160,7 +160,7 @@ symbolic_simple = propagate_invalid + PatternMatcher([
  # ** simple where folding **
  # a conditional with the same results either way is a noop, also fold const conditionals
  (UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val),
-  (UPat.cvar("gate", vec=False).where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),
+  (UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),
  # a.where(b.where(c, d), d) -> (a & b).where(c, d)
  (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
 ])
@@ -205,7 +205,7 @@ gep_pushing = PatternMatcher([
   lambda g1, g2: g2.src[0].gep(tuple(g2.arg[g1.arg[i]] for i in range(len(g1.arg))))),
  (UPat(Ops.STACK, name='vec').f(Ops.GEP, name='gep'),
   lambda gep, vec: UOp(Ops.STACK, gep.dtype, tuple(vec.src[i] for i in gep.arg)) if len(gep.arg) > 1 else vec.src[gep.arg[0]]),
-  (UPat.cvar("c", vec=False).f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(c.arg)),
+  (UPat.cvar("c").f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(c.arg)),
  # GEP on void is skipped
  (UPat(Ops.GEP, src=(UPat(dtype=dtypes.void, name="x"),)), lambda x: x),
  # GEP in order is removed
@@ -270,20 +270,20 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
  ((UPat.var("x") // UPat.cvar("c1")) // UPat.cvar("c2"), lambda x,c1,c2: x//(c1*c2) if c2.vmin>0 else None),
  # ** lt **
  # c0*x<c1 for positive int c0,c1
-  ((UPat.cvar("c0", vec=False)*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1", vec=False),
+  ((UPat.cvar("c0")*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1"),
   lambda x,c0,c1: x<math.ceil(c1.arg/c0.arg) if c0.arg > 0 and c1.arg > 0 else None),
  # c0*x<c1 for negative int c0 and non-positive c1
-  ((UPat.cvar("c0", vec=False)*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1", vec=False),
+  ((UPat.cvar("c0")*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1"),
   lambda x,c0,c1: (-x)<(-(math.floor(-c1.arg/-c0.arg))) if c0.arg < 0 and c0.arg != -1 and c1.arg <= 0 else None),
  # x//d<c -> x<c*d for d>0
-  ((UPat.var("x", dtype=dtypes.weakint)//UPat.cvar("d", vec=False))<UPat.cvar("c", vec=False),
+  ((UPat.var("x", dtype=dtypes.weakint)//UPat.cvar("d"))<UPat.cvar("c"),
   lambda x,d,c: x<(c.arg*d.arg) if d.arg > 0 else None),
  # ** move add/mul consts to end (NOTE: this is still happening before constant folding) **
  ((UPat.var("x") + UPat.cvar("c1")) + UPat.var("y"), lambda x,c1,y: (x+y)+c1),
  ((UPat.var("x") * UPat.cvar("c1")) * UPat.var("y"), lambda x,c1,y: (x*y)*c1),
  # *** rules from symbolic ***
  # generic lt folding
-  (UPat.var("x", dtypes.weakint)<UPat.cvar("c", vec=False), lambda x,c: lt_folding(x, c.arg) if 0 < c.arg else None),
+  (UPat.var("x", dtypes.weakint)<UPat.cvar("c"), lambda x,c: lt_folding(x, c.arg) if 0 < c.arg else None),
  (UPat.var("x", dtypes.weakint)*-1 < UPat.var("y")*-1, lambda x,y: y<x),
  # canonicalize a simplex with positive coefficients > 0. NOTE: not x < 1 means x > 0
  ((UPat.var("x", dtypes.weakint)<1).ne(True), lambda x: (newx<1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None),
@@ -438,7 +438,15 @@ pm_simplify_valid = PatternMatcher([
 ])

 # this is symbolic 2.0
-REMOVE_FROM_SINK_LIKE = {Ops.UNROLL, Ops.NOOP, Ops.STACK, Ops.SINK}
+REMOVE_FROM_SINK_LIKE = {Ops.UNROLL, Ops.NOOP, Ops.STACK, Ops.SINK, Ops.GROUP}
+pm_clean_up_group_sink = PatternMatcher([
+  # clean up GROUP/SINK
+  (UPat(Ops.GROUP, src=(UPat.var("x"),)), lambda x: x),
+  (UPat((Ops.SINK, Ops.GROUP), name="root"),
+    lambda root: UOp(root.op, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_SINK_LIKE else (x,) for x in root.src)), root.arg)
+      if any(x.op in REMOVE_FROM_SINK_LIKE for x in root.src) else None),
+])
+
 sym = symbolic+pm_simplify_valid+PatternMatcher([
  # reorder ALU/VECTORIZE
  (UPat(GroupOp.ALU, src=(UPat(Ops.STACK, src=UPat(name='x')), UPat(Ops.STACK, src=UPat(name='y'))), name='alu'),
@@ -467,16 +475,11 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([
  (UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")*UPat.var("y")), lambda x,y,d: y*(1-d)),
  (UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")+UPat.var("y")), lambda x,y,d: (1-d)+x*y),
  # move const multiply after REDUCE (NOTE: the mul chain can do this, but only if it's a same dtype reduce)
-  ((UPat.var("x")*UPat.cvar("c", vec=False)).reduce(arg=Ops.ADD, name="r", allow_any_len=True), lambda x,c,r: r.replace(src=(x,)+r.src[1:])*c.arg),
+  ((UPat.var("x")*UPat.cvar("c")).reduce(arg=Ops.ADD, name="r", allow_any_len=True), lambda x,c,r: r.replace(src=(x,)+r.src[1:])*c.arg),
  # reduce mul chain, move muls after the reduce
  (UPat(Ops.MUL).reduce(name="r", allow_any_len=True), reduce_mul_chain),
-  # clean up GROUP/SINK
-  (UPat(Ops.GROUP, src=(UPat.var("x"),)), lambda x: x),
-  (UPat((Ops.SINK, Ops.GROUP), name="root"),
-    lambda root: UOp(root.op, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_SINK_LIKE else (x,) for x in root.src)), root.arg)
-      if any(x.op in REMOVE_FROM_SINK_LIKE for x in root.src) else None),
  # ** combine terms (opinionated) **
  (-1 * (UPat.var("x") + UPat.var("y")), lambda x,y: (-x)+(-y)),  # -(x+y) -> -x + -y
  # (x+y)*c -> x*c+y*c. only for int, float has inf*0=nan issue
  ((UPat.var("x", dtypes.weakint) + UPat.var("y")) * UPat.cvar("c"), lambda x,y,c: x*c+y*c),
-])
+])+pm_clean_up_group_sink
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	e76de41110	fixes	2026-06-02 12:53:00 -07:00
geohot	5768042e3f	fix after merge	2026-06-02 12:45:51 -07:00
George HotzandGitHub	ef9c60238e	Merge branch 'master' into shrink_in_render	2026-06-02 11:20:24 -07: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	5825d4d833	Merge branch 'master' into shrink_in_render	2026-06-01 22:08:00 -07:00
George HotzandGitHub	7d8ed8d4d7	add store to buffer's addrspace (#16468 )	2026-06-01 22:07:43 -07:00
geohot	394afe40c5	fix for nir	2026-06-01 21:46:45 -07:00
geohot	774847a54d	is github broken	2026-06-01 19:25:48 -07:00
George HotzandGitHub	a35c1c9c38	Merge branch 'master' into shrink_in_render	2026-06-01 19:25:15 -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
George HotzandGitHub	2bcd46d946	Merge branch 'master' into shrink_in_render	2026-06-01 19:10:41 -07:00
geohot	1ffe08bab9	webgpu	2026-06-01 18:58:06 -07:00
geohot	894f1221c2	webgpu fixes	2026-06-01 18:44:09 -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
geohot	f3ecb4f8e8	amd emu fix	2026-06-01 18:16:52 -07:00
geohot	3dbaa526fa	update for ANON	2026-06-01 17:49:12 -07:00
sirhcmandGitHub	b0ecbb34d9	ci: cleanup python backend tests (#16465 )	2026-06-01 20:08:05 -04:00
geohot	31a87addca	renderer fixes	2026-06-01 16:58:01 -07:00
George HotzandGitHub	c64a37fa7d	Merge branch 'master' into shrink_in_render	2026-06-01 16:48:22 -07:00
geohot	d6f1aadeb7	test fixes	2026-06-01 16:47:26 -07:00
geohot	beb6c3ab3e	fix llvm test_tiny	2026-06-01 16:28:28 -07:00
geohot	08fe658f74	fix tests	2026-06-01 16:22:16 -07:00
geohot	f532e1b2d0	test updates	2026-06-01 16:12:37 -07:00
geohot	1c18e1bae8	fix examples	2026-06-01 16:06:52 -07:00
sirhcmandGitHub	2d0f132a3b	ci: cleanup more duplicate tests (#16462 )	2026-06-01 18:56:29 -04:00
geohot	50ac2872b3	fixes	2026-06-01 15:50:40 -07:00
wozeparrotandGitHub	aab9a5a8a3	llama: allow specifying layer count (#16464 )	2026-06-01 15:36:04 -07:00
geohot	8d327d4877	index gate	2026-06-01 15:19:00 -07:00
chenyuandGitHub	0167401fa2	minor hcopt WHERE cleanup [PR] (#16463 )	2026-06-01 17:58:38 -04:00
geohot	bdbee57f34	fix dl/dr shape	2026-06-01 14:53:41 -07:00
geohot	7b00120d92	oops, remove that	2026-06-01 14:45:29 -07:00
George HotzandGitHub	46541d70f4	Merge branch 'master' into shrink_in_render	2026-06-01 14:43:17 -07: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
geohot	8850ce9380	reg/local	2026-06-01 13:50:40 -07:00
geohot	4571b0d98a	llvm work	2026-06-01 13:22:54 -07: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
geohot	9f78877d14	some llvm fixes	2026-06-01 11:47:20 -07:00
geohot	6f506dc55e	fix python	2026-06-01 11:38:40 -07:00
geohot	12752b8a44	fix wmma	2026-05-31 17:59:48 -07:00
geohot	e808f698bc	fix uops stats	2026-05-31 14:58:03 -07:00
geohot	27835b5a31	fix CHECK_OOB	2026-05-31 14:21:35 -07:00
George HotzandGitHub	9ccee6aae7	Merge branch 'master' into shrink_in_render	2026-05-31 09:49:24 -07:00
George HotzandGitHub	556defa0f7	minor updates from vec removal (#16456 )	2026-05-31 09:48:51 -07:00
geohot	fdc7d4c0af	work	2026-05-31 09:29:58 -07:00
George HotzandGitHub	ea70715344	Merge branch 'master' into shrink_in_render	2026-05-31 09:29:52 -07:00
Javier De JesusandGitHub	989f713c1b	support negative pads in circular pad mode (#16448 )	2026-05-31 09:28:45 -07:00
George HotzandGitHub	8ba3ee138e	Merge branch 'master' into shrink_in_render	2026-05-31 09:27:40 -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	604b35aa67	Merge branch 'master' into shrink_in_render	2026-05-29 19:17:11 -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
geohot	7754025f2a	just global/local	2026-05-29 18:20:48 -07:00
geohot	5ac25ba991	that	2026-05-29 18:18:36 -07:00
sirhcmandGitHub	c377d01491	ci: run dsp on tinygrad[testing] (#16442 )	2026-05-29 21:16:56 -04:00
geohot	507c68dbc8	more passing	2026-05-29 18:02:25 -07:00
wozeparrotandGitHub	c23652e486	llama: minimize peak init mem (#16440 )	2026-05-29 18:00:37 -07:00
geohot	3e0335a4a0	work	2026-05-29 17:57:18 -07:00
sirhcmandGitHub	d943493b79	ci: remove duplicate op compile test (#16441 )	2026-05-29 19:20:31 -04:00
geohot	cb755bded6	test tiny passes	2026-05-29 15:28:53 -07:00
geohot	10c2a50e79	no GEP in program	2026-05-29 15:03:01 -07:00
chenyuandGitHub	8ac62b28e5	fix AffineGrid fusion (#16439 )	2026-05-29 17:59:47 -04:00
geohot	7b951e691e	both	2026-05-29 14:51:56 -07: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
geohot	90b2c7e115	do it as renderer hacks for now	2026-05-29 13:42:52 -07:00
geohot	14394eb97d	stack spec	2026-05-29 13:24:16 -07:00
geohot	213cc5b6b0	test plus works	2026-05-29 13:18:12 -07:00
George HotzandGitHub	1670dbfacd	Merge branch 'master' into shrink_in_render	2026-05-29 13:06:31 -07:00
chenyuandGitHub	b7280705a7	limit CONST(UNIQUE) to invalids only (#16432 )	2026-05-29 16:02:06 -04:00
George HotzandGitHub	927e16fbdb	Merge branch 'master' into shrink_in_render	2026-05-29 12:58:31 -07: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	f165839386	Merge branch 'master' into shrink_in_render	2026-05-29 11:40:00 -07: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
George HotzandGitHub	3baba3f23f	Merge branch 'master' into shrink_in_render	2026-05-29 01:37:25 -07: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
geohot	a75ad9fbaa	no vec dtype	2026-05-28 20:34:40 -07: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
George HotzandGitHub	4115d330ab	Merge branch 'master' into shrink_in_render	2026-05-28 19:27:45 -07:00
sirhcmandGitHub	a909acb882	move llvmspeed to benchmarks (#16422 )	2026-05-28 22:26:22 -04:00
George HotzandGitHub	7da2c151be	Merge branch 'master' into shrink_in_render	2026-05-28 19:18:31 -07: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 HotzandGitHub	19246184d3	Merge branch 'master' into shrink_in_render	2026-05-28 15:06:43 -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	f17eb03634	Merge branch 'master' into shrink_in_render	2026-05-28 14:40:06 -07: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
geohot	1c10882bf0	something	2026-05-28 13:31:48 -07:00
nimlgenandGitHub	b0e49afaf1	hcq2: new multi (#16413 ) * hcq2: new multi * op	2026-05-28 22:16:10 +03:00
geohot	6881b32e84	use shrink in renderers	2026-05-28 11:58:57 -07: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