Merge branch 'master' into invalid_try_3

invalid try 3 (kimi)
2026-08-15 19:18:27 +00:00 · 2026-07-25 14:07:16 -07:00 · 2026-07-25 14:06:29 -07:00
190 changed files with 2343 additions and 7723 deletions
@@ -137,7 +137,6 @@ runs:
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
      shell: bash
      run: |
-        sudo mkdir -p /var/cache/apt/archives
        sudo chown -R $USER:$USER /var/cache/apt/archives

        echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
@@ -215,7 +214,6 @@ runs:
          sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
        fi

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

    - name: Add clang to PATH (Linux)
@@ -117,10 +117,10 @@ jobs:
      run: python3 test/external/process_replay/reset.py
    - name: Run llama3.2
      run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.6
-      # qwen3.6:35b-a3b doesn't fit on mac
+    - name: Run qwen3.5
+      # qwen3.5:35b-a3b doesn't fit on mac
      if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen36_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.6:35b-a3b --benchmark --warmup
+      run: BENCHMARK_LOG=qwen35_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.5:35b-a3b --benchmark --warmup
    - name: Run olmoe
      # just metal for now
      if: ${{ matrix.dev == 'METAL' }}
@@ -402,35 +402,6 @@ jobs:
      run: PYTHONPATH=. DEV=PCI+NV:NAK python3.11 test/test_tiny.py

  testcommalatest:
-    name: comma Benchmark (0.11.2)
-    runs-on: [self-hosted, Linux, comma]
-    timeout-minutes: 12
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
-    - name: openpilot compile3 0.11.2 supercombo
-      run: BENCHMARK_LOG=openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
-    - name: openpilot compile3 0.11.2 supercombo (from pickle)
-      run: BENCHMARK_LOG=openpilot_0_11_2_supercombo_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
-    - name: IR3 openpilot compile3 0.11.2 supercombo
-      run: BENCHMARK_LOG=ir3_openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=41 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
-    - name: openpilot compile3 0.11.2 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_11_2_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testcommaold:
    name: comma Benchmark (0.11.0)
    runs-on: [self-hosted, Linux, comma]
    timeout-minutes: 12
@@ -461,6 +432,35 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

+  testcommaold:
+    name: comma Benchmark (0.10.1)
+    runs-on: [self-hosted, Linux, comma]
+    timeout-minutes: 12
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
+    - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
+      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.10.1 driving_vision
+      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.10.1 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3.2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.10.1 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
  testqualcommdsp:
    name: DSP Benchmark
    runs-on: [self-hosted, Linux, comma4]
@@ -515,12 +515,12 @@ jobs:
      run: |
        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: openpilot compile3 big_driving_supercombo
-      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo PICKLE_OOB=1 PYTHONPATH="." TC_OPT=2 GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/10926f2c0911821ca0e72439c1c3bf3ec11f0a08789aa14b7ee8f25379b2afa4 openpilot.pkl
-    - name: openpilot load_pickle big_driving_supercombo
-      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_load_pickle PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=25 python3 examples/openpilot/load_pickle.py openpilot.pkl
-    - name: openpilot run_pickle big_driving_supercombo
-      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
+    - name: openpilot compile3 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot load_pickle 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
+    - name: openpilot run_pickle 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py
    - name: Test copy speeds
      run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds

@@ -1,34 +0,0 @@
-#!/usr/bin/env python3
-# Sticky PR comment via the REST API: find an existing comment containing MARKER and PATCH it, or POST a new one.
-# Works on GitHub and Gitea (stdlib only, replaces marocchino/sticky-pull-request-comment which needs GraphQL).
-# Env vars: GITHUB_TOKEN, GITHUB_API_URL, GITHUB_REPOSITORY (set by the runner), PR_NUMBER, MARKER, and BODY_FILE or MESSAGE.
-import json, os, sys, urllib.request
-
-api, repo = os.environ["GITHUB_API_URL"], os.environ["GITHUB_REPOSITORY"]
-pr, marker = os.environ["PR_NUMBER"], os.environ["MARKER"]
-body = open(os.environ["BODY_FILE"]).read() if os.environ.get("BODY_FILE") else os.environ["MESSAGE"]
-
-if not body.strip():
-  print("comment body is empty, not posting")
-  sys.exit(0)
-
-def req(url, method="GET", payload=None):
-  r = urllib.request.Request(url, data=None if payload is None else json.dumps(payload).encode(), method=method,
-    headers={"Authorization": f"token {os.environ['GITHUB_TOKEN']}", "Accept": "application/json", "Content-Type": "application/json"})
-  return json.load(urllib.request.urlopen(r))
-
-# find the latest sticky comment (paginate, 100 comments per page)
-existing, page = None, 1
-while True:
-  comments = req(f"{api}/repos/{repo}/issues/{pr}/comments?per_page=100&page={page}")
-  stickies = [c for c in comments if marker in (c.get("body") or "")]
-  if stickies: existing = stickies[-1]
-  if not comments or len(comments) < 100: break
-  page += 1
-
-if existing is not None and existing["body"] == body:
-  print("comment is already up to date")
-  sys.exit(0)
-url = f"{api}/repos/{repo}/issues/comments/{existing['id']}" if existing is not None else f"{api}/repos/{repo}/issues/{pr}/comments"
-resp = req(url, 'PATCH' if existing is not None else 'POST', {'body': body})
-print(f"{'updated' if existing is not None else 'created'} comment {resp['id']}")
@@ -26,14 +26,14 @@ jobs:
      - name: Check whether branch is up-to-date
        id: brstat
        run: |
-          # fetch master from the base repo (tinygrad/tinygrad on GitHub, the mirror on Gitea), not the PR head remote
-          git fetch "${{ github.event.pull_request.base.repo.clone_url }}" master
+          git remote add tinygrad https://github.com/tinygrad/tinygrad
+          git fetch tinygrad master
          echo "${{ github.event.pull_request.head.sha }}"
-          git rev-list --left-right --count FETCH_HEAD...${{ github.event.pull_request.head.sha }} | awk '{print "Behind "$1" - Ahead "$2""}'
-          count=$(git rev-list --left-right --count FETCH_HEAD...${{ github.event.pull_request.head.sha }} | awk '{print $1}')
+          git rev-list --left-right --count  tinygrad/master...${{ github.event.pull_request.head.sha }} | awk '{print "Behind "$1" - Ahead "$2""}'
+          count=$(git rev-list --left-right --count  tinygrad/master...${{ github.event.pull_request.head.sha }} | awk '{print $1}')
          if [ $count -gt 0 ]
          then
-            echo "Current branch is behind ${{ github.event.pull_request.base.repo.full_name }} master branch!"
+            echo "Current branch is behind tinygrad master branch!"
            echo "stat=true" >> "$GITHUB_OUTPUT"
          else
            echo "stat=false" >> "$GITHUB_OUTPUT"
@@ -75,13 +75,13 @@ jobs:
          python sz.py "$BASE" "$PR" > loc_content.txt
      - name: Comment Code Line Diff
        continue-on-error: false
-        env:
+        uses: marocchino/sticky-pull-request-comment@v3
+        with:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          PR_NUMBER: ${{ github.event.pull_request.number }}
-          MARKER: "### Changes"
-          BODY_FILE: loc_content.txt
-        # note: run the script from the base checkout, never from the PR checkout
-        run: python3 "$GITHUB_WORKSPACE/base/.github/workflows/sticky_comment.py"
+          ignore_empty: true
+          skip_unchanged: true
+          recreate: true
+          path: loc_content.txt

  rebase:
    name: Core Library Line Difference
@@ -91,14 +91,12 @@ jobs:
    needs: checkbranch
    if: needs.checkbranch.outputs.branchstat == 'true'
    steps:
-      # pull_request_target: a plain checkout gets the base repo, so no PR code is executed
-      - uses: actions/checkout@v6
      - name: Comment Rebase
        continue-on-error: false
-        env:
+        uses: marocchino/sticky-pull-request-comment@v3
+        with:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          PR_NUMBER: ${{ github.event.pull_request.number }}
-          MARKER: "line count difference bot is disabled"
-          MESSAGE: |
-            This branch currently is behind ${{ github.event.pull_request.base.repo.full_name }} master. The line count difference bot is disabled.
-        run: python3 .github/workflows/sticky_comment.py
+          skip_unchanged: true
+          recreate: true
+          message: |
+            This branch currently is behind tinygrad/master. The line count difference bot is disabled.
@@ -167,7 +167,6 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: unittest-13
-        python-version: '3.11'
        deps: testing_unit
        llvm: 'true'
        amd: 'true'
@@ -177,14 +176,13 @@ jobs:
      run: |
        DEV=NULL python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
        DEV=NULL VIZ=1 python3 -m pytest -n=auto test/null/test_viz.py
-        DEBUG=7 python -m tinygrad.viz.cli --json | jq empty
    # TODO: too slow
    # - name: Run SDXL on NULL backend
    #   run: DEV=NULL DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
    - name: Run Clip tests for SD MLPerf on NULL backend
      run: DEV=NULL python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
    - name: Run AMD emulated BERT training on NULL backend
-      run: DEV=NULL::gfx1201 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: DEV=NULL::gfx1201 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    # TODO: support fake weights
    #- name: Run LLaMA 7B on 4 fake devices
    #  run: DEV=NULL python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -201,7 +199,6 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: unittest-13
-        python-version: '3.11'
        pydeps: "pre-commit"
        deps: testing_unit
        llvm: 'true'
@@ -353,8 +350,6 @@ jobs:
        run: DEV=NULL NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
      - name: Test llama 3 training
        run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
-      - name: Test gpt-oss training
-        run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=32 BS=2 SEQLEN=128 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 MXFP8=1 VOCAB_SIZE=32000 LAYERS=2 EXPERTS=4 MODEL=gptoss PYTHONPATH=. python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -483,7 +478,7 @@ jobs:
          wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
          echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-21 main" | sudo tee /etc/apt/sources.list.d/llvm.list
          sudo apt-get update
-          sudo apt-get install -y llvm-21 llvm-21-tools cloc
+          sudo apt-get install llvm-21 llvm-21-tools cloc
      - name: Install rocprof-trace-decoder
        run: sudo PYTHONPATH="." ./extra/sqtt/install_rocprof_decoder.py
      - name: Run AMD renderer tests
@@ -633,7 +628,7 @@ jobs:

  unittestmacos:
    name: MacOS (unit)
-    runs-on: macos-26
+    runs-on: &macos macos-26
    timeout-minutes: 20
    steps:
    - name: Checkout Code
@@ -643,8 +638,12 @@ jobs:
      with:
        key: unittest-macos
        deps: testing_unit
+        amd: 'true'
+        ocelot: 'true'
    - name: Run unit tests
      run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --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)
@@ -655,25 +654,6 @@ 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 process replay tests
-      uses: ./.github/actions/process-replay
-
-  unittestmacosmock:
-    name: MacOS (unit, mock)
-    runs-on: macos-26
-    timeout-minutes: 20
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: unittest-macos-mock
-        deps: testing_unit
-        amd: 'true'
-        ocelot: 'true'
-    - name: Run NULL backend tests
-      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
    - name: Run pytest (amd)
      env:
        DEV: MOCKKFD+AMD
@@ -697,7 +677,7 @@ jobs:
      matrix:
        group: [1, 2]
    name: MacOS (DEV=METAL) (${{ matrix.group }})
-    runs-on: macos-26
+    runs-on: *macos
    timeout-minutes: 20
    env:
      DEV: METAL
@@ -729,7 +709,7 @@ jobs:
          - 'WEBGPU'

    name: MacOS (DEV=${{ matrix.dev }})
-    runs-on: macos-26
+    runs-on: *macos
    timeout-minutes: 20
    steps:
    - name: Checkout Code
@@ -3,4 +3,4 @@
 - Run tests with `-n12` for speed (e.g. `python -m pytest test/null/test_dtype.py -x -q -n12`)
 - Run `python -m mypy tinygrad/` to typecheck
 - Run `python -m ruff check .` to lint
- Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules
+- Read `./tinygrad/viz/README` for profiling
@@ -88,7 +88,7 @@ def example_3_custom_uop(a:Tensor, correct):

    # store all the per lane accumulators to LOCAL
    local_accs = UOp.placeholder((LCLS,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-    local_accs = local_accs.after(local_accs[lane].store(acc[0]))
+    local_accs = local_accs.after(local_accs[lane].store(acc[0]).barrier())

    # accumulate LOCALs into a single per CU accumulator
    late_reduce_loop = UOp.range(LCLS, 3, AxisType.REDUCE)
@@ -1,8 +1,7 @@
 ::: tinygrad.dtype.DType

-::: tinygrad.dtype.DTypes
+::: tinygrad.dtype.dtypes
    options:
-        heading: dtypes
        members: true
        members_order: source
        show_labels: false
@@ -9,7 +9,8 @@ from extra.lr_scheduler import OneCycleLR
 GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]

 # override tinygrad defaults
-Context(DEFAULT_FLOAT=dtypes.half, FUSE_OPTIM=1).__enter__()
+dtypes.default_float = dtypes.half
+Context(FUSE_OPTIM=1).__enter__()

 # from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 batchsize = getenv("BS", 1024)
@@ -22,6 +22,10 @@ class Attention:
    self.head_dim = dim // n_heads

  def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]) -> Tensor:
+    if mask is not None or start_pos.val == 0:
+      # no symbolic shape qkv when consuming prompts
+      start_pos = start_pos.val
+
    if HALF: x = x.half()
    xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
    xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
@@ -34,8 +38,12 @@ class Attention:
    # update the cache
    self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()

-    keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
-    values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
+    if start_pos > 0:
+      keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
+      values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
+    else:
+      keys = xk
+      values = xv

    xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
    return self.c_proj(xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2).reshape(bsz, seqlen, self.dim))
@@ -78,14 +86,15 @@ class Transformer:
      seqlen = tokens.shape[1]
      tok_emb = self.wte(tokens)

-    # start_pos is a bound Variable, so everything below it stays symbolic
-    pos_emb = self.wpe(self.allpos.shrink((None, (start_pos, start_pos+seqlen))))
+    # not symbolic when consuming the prompt
+    selected_pos = (0, seqlen) if start_pos.val == 0 else (start_pos, start_pos+1)
+    pos_emb = self.wpe(self.allpos.shrink((None, selected_pos)))

    h = tok_emb + pos_emb

    if HALF: h = h.half()

-    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos+1) if seqlen > 1 else None
+    mask = Tensor.full((1, 1, seqlen, start_pos.val+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos.val+1) if seqlen > 1 else None

    for hi in self.h: h = hi(h, start_pos, mask)

@@ -1,11 +1,11 @@
 import os, random, pickle, queue, struct, math, functools, hashlib, time
 from typing import List
 from pathlib import Path
-from multiprocessing import Queue, Process, shared_memory, connection, Lock
+from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count

 import numpy as np
 from tinygrad import dtypes, Tensor
-from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, NUM_CPU_THREADS
+from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX
 from tinygrad.nn.state import TensorIO

 ### ResNet
@@ -131,7 +131,7 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir
    else: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name}")
    Y = [None] * (batch_size*BATCH_COUNT)

-    for _ in range(NUM_CPU_THREADS.value):
+    for _ in range(cpu_count()):
      p = Process(target=loader_process, args=(q_in, q_out, X, seed))
      p.daemon = True
      p.start()
@@ -212,7 +212,7 @@ def batch_load_train_bert(BS:int, seed:int|None=None):
    rng.shuffle(fs)
    train_files.append(fs.pop(0))

-  cycle_length = min(NUM_CPU_THREADS.value, len(train_files))
+  cycle_length = min(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), len(train_files))
  assert cycle_length > 0, "cycle_length must be greater than 0"

  dataset = InterleavedDataset(train_files, cycle_length)
@@ -301,7 +301,7 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=
    X = Tensor.empty(*sz, dtype=dtypes.float32, device=f"disk:/dev/shm/{shm_name_x}")
    Y = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name_y}")

-    for _ in range(NUM_CPU_THREADS.value):
+    for _ in range(cpu_count()):
      proc = Process(target=load_unet3d_data, args=(preprocessed_dataset_dir, seed, queue_in, queue_out, X, Y))
      proc.daemon = True
      proc.start()
@@ -437,7 +437,7 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
  dataset_iter = iter(image_ids)

  try:
-    for _ in range(NUM_CPU_THREADS.value):
+    for _ in range(cpu_count()):
      proc = Process(
        target=load_retinanet_data,
        args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),
@@ -1668,7 +1668,7 @@ def train_llama3():
 def train_gptoss():
  from examples.mlperf.models.gpt_oss import GPTOSS, GPT_OSS_20B, apply_grad, FP8_DTYPE
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, clip_grads
+  from examples.mlperf.optim import GradAccClipAdamW, clip_grads

  BENCHMARK = getenv("BENCHMARK")

@@ -1711,10 +1711,9 @@ def train_gptoss():
    wandb.init(config=config, **wandb_args, project="MLPerf-gpt-oss")

  model_params = GPT_OSS_20B
-  model_params['vocab_size'] = getenv("VOCAB_SIZE", 128256)
+  model_params['vocab_size'] = 128256
  real_vocab_size = model_params['vocab_size']
  if (layers:=getenv("LAYERS")) != 0: model_params['n_layers'] = layers
-  if (experts:=getenv("EXPERTS")) != 0: model_params['n_experts'] = experts
  print(f"model parameters: {model_params}")

  model = GPTOSS(**model_params, max_context=SEQLEN)
@@ -1735,12 +1734,7 @@ def train_gptoss():
  is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
  is_fake_offload = Device.DEFAULT == "NULL"
  optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
-  params_wd = [p for p in params if p.ndim >= 3]
-  params_no_wd = [p for p in params if p.ndim < 3]
-  optim = GradAccClipAdamWGroup(
-    GradAccClipAdamW(params_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device),
-    GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=grad_acc, device=optim_device),
-  )
+  optim = GradAccClipAdamW(params, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)

  for p in optim.params:
    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
@@ -1749,10 +1743,7 @@ def train_gptoss():

  from extra.gemm.cdna_asm_gemm import _mx_block_scale
  model_state = get_state_dict(model)
-  def _scale_key(n):
-    if "." in n and (c:=f"{(b:=n.rsplit('.',1))[0]}_scale.{b[1]}") in model_state: return c
-    return f"{n}_scale"
-  fp8_scale_names = {n: _scale_key(n) for n, t in model_state.items() if t.dtype == FP8_DTYPE}
+  fp8_scale_names = {n: f"{n}_scale" for n, t in model_state.items() if t.dtype == FP8_DTYPE}
  fp8_inv_scales = [model_state[sname] for sname in fp8_scale_names.values()]
  for wname, sname in fp8_scale_names.items():
    w, scale = model_state[wname], model_state[sname]
@@ -25,7 +25,6 @@ FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
 SPLIT_W13 = getenv("SPLIT_W13", 0)
 COLUMNWISE_WEIGHT_SCALE = getenv("COLUMNWISE_WEIGHT_SCALE", 0)
 MXFP8 = getenv("MXFP8", 0)
-MXFP4 = getenv("MXFP4", 0)

 FP8_DTYPE = dtypes.fp8e4m3
 FP8_GRAD_DTYPE = dtypes.fp8e5m2
@@ -45,11 +44,6 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
      from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
      if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
    return (x @ w.T,)
-  if MXFP4:
-    assert x is not None, "MXFP4 matmul requires an unquantized input"
-    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm
-    if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T, mxfp4=True),)
-    return (x @ w.T,)
  assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
  if MXFP8:
    from extra.gemm.cdna_asm_gemm import asm_gemm, quantize_mxfp8, mx_pack, can_use_asm_gemm, _mx_block_scale
@@ -85,7 +79,7 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca

 def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
                         next_amax_x:Tensor, grad_amax_state:Tensor, next_grad_amax_state:Tensor):
-  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
+  if FUSED_ADD_NORM_MUL_QUANTIZE:
    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
    x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
@@ -98,7 +92,7 @@ def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, ep

 def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
                             next_amax_x:Tensor, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
-  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
+  if FUSED_ADD_NORM_MUL_QUANTIZE:
    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
    x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
@@ -114,7 +108,7 @@ def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
                             amax_x2:Tensor, next_amax_x2:Tensor,
                             grad_amax_xw13:Tensor, next_grad_amax_xw13:Tensor,
                             grad_amax_xout:Tensor, next_grad_amax_xout:Tensor):
-  if FUSED_SILU_W13 and not MXFP4:
+  if FUSED_SILU_W13:
    from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
    x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
                                                 next_grad_amax_state=next_grad_amax_xw13, amax_out=next_amax_x2)
@@ -185,9 +179,6 @@ class FlatTransformer:
      from extra.gemm.cdna_asm_gemm import quantize_mxfp8
      w_q, w_e8, _ = quantize_mxfp8(w.reshape(self.n_layers * out_features, in_features))
      return w_q.reshape(self.n_layers, out_features, in_features), w_e8.reshape(self.n_layers, out_features, in_features // 32)
-    if MXFP4:
-      # FP4 is produced dynamically so optimizer updates always start from the current BF16 weight.
-      return w.cast(dtypes.bfloat16), Tensor.ones(self.n_layers)
    amax = (w.abs().max(axis=2) if COLUMNWISE_WEIGHT_SCALE else w.abs().flatten(1).max(1)).detach()
    scale = FP8_MAX / (amax + 1e-8)
    inv_scale = (amax + 1e-8) / FP8_MAX
@@ -13,14 +13,10 @@ from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
 from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
-from extra.gemm.moe_gemm import grouped_mx_gemm
-from extra.gemm.moe_routing import route, dispatch, combine

 FP8_DTYPE = dtypes.fp8e4m3
 FP8_MAX = 448.0
-INIT_STD = 0.02
-ASM_GEMM = getenv("ASM_GEMM", 0)
-
+INIT_STD = 0.008

 def _quant_dequant_fwd(x:Tensor) -> Tensor:
  # x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
@@ -54,7 +50,8 @@ def _dequant_fwd_fxn(wq_p, ws_p, device):
  return _dequant_fwd(Tensor(wq_p, device=device), Tensor(ws_p, device=device))

 def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
-  return (Tensor(grad).cast(dtypes.bfloat16).uop, None)
+  w_scale = Tensor(call.src[2])
+  return ((Tensor(grad).cast(dtypes.bfloat16) * _mx_scale(w_scale).cast(dtypes.bfloat16)).uop, None)

 def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
  fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)
@@ -63,34 +60,10 @@ def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:

 def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
  l_shape = x.shape[:-1]
-  if ASM_GEMM:
-    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
-    x2, K, N = x.reshape(-1, x.shape[-1]), x.shape[-1], w_q.shape[0]
-    wq, ws = w_q, w_scale
-    if (pad := (-K) % 256):
-      x2 = x2.pad(((0, 0), (0, pad)))
-      wq = wq.pad(((0, 0), (0, pad)))
-      ws = ws.pad(((0, 0), (0, pad // 32)), value=127).cast(dtypes.uint8)
-    if (npad := (-N) % 256):
-      wq = wq.pad(((0, npad), (0, 0)))
-      ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8)
-    x_q, x_e8, x_si = quantize_mxfp8(x2)
-    if x_si is not None and can_use_asm_gemm(x_q, wq.T):
-      out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(ws), ws), mx_w_stored=True)
-      return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
  x_phys = quant_dequant_mx(x.reshape(-1, x.shape[-1])).reshape(*l_shape, x.shape[-1])
  w_phys = dequant_weight(w_q, w_scale)
  return (x_phys @ w_phys.T).cast(dtypes.bfloat16)

-def _pad_to_mult(t:Tensor, axis:int, mult:int=256) -> Tensor:
-  if (r := (-t.shape[axis]) % mult) == 0: return t
-  pads = [(0, 0)] * t.ndim
-  pads[axis] = (0, r)
-  return t.pad(tuple(pads))
-
-def _pad_cols(t:Tensor) -> Tensor: return _pad_to_mult(t, -1)
-def _pad_rows(t:Tensor) -> Tensor: return _pad_to_mult(t, -2)
-
 def swiglu(x:Tensor, limit:float=7.0, alpha:float=1.702) -> Tensor:
  x_glu, x_linear = x[..., ::2], x[..., 1::2]
  x_glu = x_glu.clamp(max_=limit)
@@ -127,9 +100,9 @@ class GPTOSS:
    self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
    self.gate = Tensor.normal(n_layers, n_experts, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
    self.gate_bias = Tensor.zeros(n_layers, n_experts, dtype=dtypes.bfloat16).contiguous()
-    self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim, moe=True)
+    self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim)
    self.w_gate_up_bias = Tensor.zeros(n_layers, n_experts, intermediate_size * 2, dtype=dtypes.bfloat16).contiguous()
-    self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std, moe=True)
+    self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std)
    self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous()

    # output
@@ -139,15 +112,10 @@ class GPTOSS:
    self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
    self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False)

-  def _quant_weight(self, *shape:int, std:float=INIT_STD, moe:bool=False):
-    def _one(*s:int):
-      w = Tensor.zeros(*s) if getenv("ZEROS") else Tensor.normal(*s, mean=0.0, std=std)
-      w_q, w_e8, _ = quantize_mxfp8(_pad_cols(_pad_rows(w)) if moe else w)
-      return w_q, w_e8.is_param_(False)
-    if moe:
-      qs = [_one(*shape[1:]) for _ in range(shape[0])]
-      return [q[0] for q in qs], [q[1] for q in qs]
-    return _one(*shape)
+  def _quant_weight(self, *shape:int, std:float=INIT_STD):
+    w = Tensor.zeros(*shape) if getenv("ZEROS") else Tensor.normal(*shape, mean=0.0, std=std)
+    w_q, w_e8, _ = quantize_mxfp8(w)
+    return w_q, w_e8.is_param_(False)

  def _attn_mask(self, seqlen:int, dtype) -> Tensor:
    i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
@@ -206,32 +174,17 @@ class GPTOSS:
                   w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor):
    x_normed, rrms = rmsnorm(x, self.norm_eps)
    inp = x_normed * ffn_norm
+
    logits = inp.float() @ gate.float().T + gate_bias.float()
-    dim, inter = self.dim, self.intermediate_size
+    thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
+    weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)

-    if getenv("GROUPED_MOE", 0):
-      bsz, seqlen = x.shape[:2]
-      inp, logits = inp.reshape(-1, dim), logits.reshape(-1, self.n_experts)
-      r = route(logits, self.experts_per_tok, self.n_experts)
-      onehot = r.rows_e.one_hot(self.n_experts).float()
-      xg = dispatch(_pad_cols(inp.cast(dtypes.bfloat16)), r)
-      h = grouped_mx_gemm(xg, (w_gate_up, w_gate_up_scale), r.off)[:, :2*inter] + (onehot @ w_gate_up_bias.float()).cast(dtypes.bfloat16)
-      y = swiglu(h, self.swiglu_limit)
-      z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
-          + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
-      out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
-    else:
-      thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
-      weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
-
-      out = None
-      for e in range(self.n_experts):
-        gu_q, gu_s = w_gate_up[e][:2*inter, :dim].contiguous(), w_gate_up_scale[e][:2*inter, :dim//32].contiguous()
-        dn_q, dn_s = w_down[e][:dim, :inter].contiguous(), w_down_scale[e][:dim, :inter//32].contiguous()
-        gate_up = matmul_mx(inp, gu_q, gu_s) + w_gate_up_bias[e]
-        y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), dn_q, dn_s) + w_down_bias[e]).contiguous()
-        contrib = weights[..., e:e+1].cast(y.dtype) * y
-        out = contrib if out is None else out + contrib
+    out = None
+    for e in range(self.n_experts):
+      gate_up = matmul_mx(inp, w_gate_up[e], w_gate_up_scale[e]) + w_gate_up_bias[e]
+      y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), w_down[e], w_down_scale[e]) + w_down_bias[e]).contiguous()
+      contrib = weights[..., e:e+1].cast(y.dtype) * y
+      out = contrib if out is None else out + contrib
    return out, [x_normed, rrms]

  @function(precompile=True, precompile_backward=True)
@@ -1,8 +1,8 @@
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.nn.optim import Optimizer, OptimizerGroup
+from tinygrad.nn.optim import Optimizer
 from tinygrad.helpers import FUSE_OPTIM, getenv
-from tinygrad.uop.ops import UOp, Ops, AxisType
+from tinygrad.uop.ops import UOp, Ops

 STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
 MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
@@ -42,8 +42,8 @@ class GradAccClipAdamW(Optimizer):
      self.master_params = None

  def _zero_shard(self, t:Tensor) -> Tensor:
-    if not self.zero or t.ndim < 2 or (t.shape[0] % len(self.device)) != 0: return t
-    return Tensor(t.uop._shard(0, UOp.range(len(self.device), -1, AxisType.DEVICE)).unshard(0)).clone()
+    if not self.zero or (t.shape[0] % len(self.device)) != 0: return t
+    return Tensor(t.uop._shard(0, len(self.device)).multi(0)).clone()

  def _zero_gather(self, t:Tensor) -> Tensor:
    if not isinstance(t.device, tuple) or t.uop.axis != 0: return t
@@ -121,24 +121,3 @@ class GradAccClipAdamW(Optimizer):
      return ret.shard_like(t) if offloaded else ret
    out = new_w.cast(t.dtype)
    return out.shard_like(t) if offloaded else out
-
-class GradAccClipAdamWGroup(OptimizerGroup):
-  def __init__(self, *optimizers:GradAccClipAdamW):
-    super().__init__(*optimizers)
-    for o in self.optimizers[1:]: o.lr = self.optimizers[0].lr
-  def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
-    offset = 0
-    to_realize = []
-    for o in self.optimizers:
-      n = len(o.params)
-      to_realize += o.fschedule_step(grads[offset:offset+n])
-      offset += n
-    Tensor.realize(*to_realize, *([grad_norm] if grad_norm is not None else []))
-  @property
-  def lr(self): return self.optimizers[0].lr
-  @property
-  def device(self): return self.optimizers[0].device
-  @property
-  def master_params(self):
-    mp = [mp for o in self.optimizers for mp in (o.master_params or [])]
-    return mp if mp else None
@@ -1,4 +1,4 @@
-import os, sys, pickle, time, re, tempfile, struct, shutil, io
+import os, sys, pickle, time, re
 import numpy as np
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"

@@ -9,39 +9,6 @@ from tinygrad.nn.onnx import OnnxRunner

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
-PICKLE_OOB = getenv("PICKLE_OOB")
-
-def dump_pickle(obj, f):
-  if PICKLE_OOB:
-    # allows pickling when buffers don't fit in (CPU) RAM
-    # from openpilot/selfdrive/modeld/helpers.py
-    with tempfile.TemporaryFile(dir=".") as tmp:
-      def buffer_callback(pb: pickle.PickleBuffer):
-        m = pb.raw()
-        tmp.write(struct.pack('<q', m.nbytes))
-        tmp.write(m)
-        pb.release() # keep peak ram at ~1 buffer
-      stream = io.BytesIO()
-      pickle.Pickler(stream, protocol=5, buffer_callback=buffer_callback).dump(obj)
-      opcodes = stream.getvalue()
-      f.write(struct.pack('<q', len(opcodes)))
-      f.write(opcodes)
-      tmp.seek(0)
-      shutil.copyfileobj(tmp, f)
-  else: pickle.dump(obj, f)
-
-def load_pickle(f):
-  if PICKLE_OOB:
-    # allows unpickling when buffers don't fit in (CPU) RAM
-    # from openpilot/selfdrive/modeld/helpers.py
-    opcodes = f.read(struct.unpack('<q', f.read(8))[0])
-    def buffers():
-      while (h := f.read(8)):
-        pb = pickle.PickleBuffer(bytearray(struct.unpack('<q', h)[0]))
-        f.readinto(pb)
-        yield pb
-    return pickle.load(io.BytesIO(opcodes), buffers=buffers())
-  else: return pickle.load(f)

 def compile(onnx_file):
  run_onnx = OnnxRunner(onnx_file)
@@ -61,8 +28,8 @@ def compile(onnx_file):
    inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()}
  print("created tensors")

-  @TinyJit(prune=True)
-  def run_onnx_jit(**kwargs): return next(iter(run_onnx({k:v.to(Device.DEFAULT) for k,v in kwargs.items()}).values())).cast('float32')
+  run_onnx_jit = TinyJit(lambda **kwargs:
+                         next(iter(run_onnx({k:v.to(Device.DEFAULT) for k,v in kwargs.items()}).values())).cast('float32'), prune=True)
  for i in range(3):
    GlobalCounters.reset()
    print(f"run {i}")
@@ -98,7 +65,8 @@ def compile(onnx_file):
  if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
    assert gated_read_image_count == allowed_gated_read_image, f"different gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"

-  with open(OUTPUT, "wb") as f: dump_pickle(run_onnx_jit, f)
+  with open(OUTPUT, "wb") as f:
+    pickle.dump(run_onnx_jit, f)
  mdl_sz = os.path.getsize(onnx_file)
  pkl_sz = os.path.getsize(OUTPUT)
  print(f"mdl size is {mdl_sz/1e6:.2f}M")
@@ -168,7 +136,7 @@ def bench(run, inputs):

 if __name__ == "__main__":
  if getenv("RUN_PICKLE"):
-    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
    inputs = {name: Tensor(Tensor.randn(*view.shape, dtype=dtype).numpy(), device=device)
              for name, (view, _vars, dtype, device) in zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_input_info)}
    test_vs_compile(pickle_loaded, inputs)
@@ -176,7 +144,7 @@ if __name__ == "__main__":
    onnx_file = fetch(OPENPILOT_MODEL)
    inputs, outputs = compile(onnx_file)

-    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)

    test_vs_compile(pickle_loaded, inputs, outputs)
    if getenv("SELFTEST"):
@@ -1,6 +1,5 @@
-import sys
+import sys, pickle
 from extra.bench_log import WallTimeEvent, BenchEvent
-from examples.openpilot.compile3 import load_pickle
 from tinygrad.helpers import getenv

 PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
@@ -8,7 +7,7 @@ PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
 load_times = []

 for _ in range(10):
-  with WallTimeEvent(BenchEvent.STEP) as wte: load_pickle(open(PKL, 'rb'))
+  with WallTimeEvent(BenchEvent.STEP) as wte: pickle.load(open(PKL, 'rb'))
  load_times.append(wte.time)
  print(f"pickle load: {wte.time:6.2f} s")

@@ -241,8 +241,8 @@ export default {model_name};
 def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False):
  assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"

-  # NOTE: NUM_CPU_THREADS=1, since export does not support threading
-  with Context(JIT=2, NUM_CPU_THREADS=1): linear, output_bufs = jit_model(model, *inputs)
+  # NOTE: CPU_COUNT=1, since export does not support threading
+  with Context(JIT=2, CPU_COUNT=1): linear, output_bufs = jit_model(model, *inputs)
  functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
  state = get_state_dict(model)
  weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
@@ -264,7 +264,7 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model"
        if getattr(dim, "op", None) is Ops.ADD and len(dim.src) == 2 and \
           any(s.op is Ops.PARAM and s.addrspace is AddrSpace.ALU for s in dim.src) and any(s.op is Ops.CONST for s in dim.src):
          name, val = dim.src if dim.src[1].op is Ops.CONST else reversed(dim.src)
-          global_size[j] = f"_{name.expr}[0] + {val.val}"
+          global_size[j] = f"_{name.expr}[0] + {val.arg}"

  prg = ""
  if target == "clang":
@@ -18,9 +18,9 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
  SEQ = inp.shape[1]
  OUT = weight.shape[0]
  IN = weight.shape[-1]
-  seq_idx = UOp.range(SEQ, 2)
-  out_idx = UOp.range(OUT, 3)
-  batch_idx = UOp.range(output.size//SEQ//OUT, 1)
+  seq_idx = UOp.range(SEQ, 2, AxisType.LOOP)
+  out_idx = UOp.range(OUT, 3, AxisType.LOOP)
+  batch_idx = UOp.range(output.size//SEQ//OUT, 1, AxisType.LOOP)
  reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
  product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
  reduced = product.reduce(reduce_idx, arg=Ops.ADD)
@@ -53,7 +53,7 @@ class FP8Linear:
    x_fp8, x_scale = quantize_to_fp8(x)
    GPUS = self.weight.device
    if isinstance(GPUS, tuple) and len(GPUS) > 1:
-      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.unshard(0), device=GPUS)
+      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.multi(0), device=GPUS)
    else:
      y = Tensor.empty((batch, seq, self.weight.shape[0]), dtype=dtypes.float)
    y = Tensor.custom_kernel(y, x_fp8, w_fp8, fxn=custom_matmul, grad_fxn=custom_matmul_backward)[0]
@@ -58,8 +58,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
  B_copy = B_local.permute((1,0)) if use_wmma else B_local
  A_store = A_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
  B_store = B_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
-  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
-  A_local, B_local = A_local.after(A_store, B_store), B_local.after(A_store, B_store)
+  barrier = UOp.barrier(A_store, B_store)
+  A_local, B_local = A_local.after(barrier), B_local.after(barrier)

  # -- COMPUTE --
  lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
@@ -70,8 +70,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:

  if use_wmma:
    k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
-    tile_m = UOp.range(TM // WMMA_ACC, 200)
-    tile_n = UOp.range(TN, 201)
+    tile_m = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
+    tile_n = UOp.range(TN, 201, AxisType.LOOP)

    acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0,2,1)[tile_m, tile_n]
    a_frag = A_local.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_K // WMMA_K, WMMA_K)[wave_m, tile_m, lane_n, k]
@@ -96,8 +96,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
    b_frag = b_frag.reshape(1, TN).expand(TM, TN)
    acc_store = acc.store(acc.after(k) + (a_frag * b_frag))

-  # store accumulator and loop (the barrier at the end of the loop is implied by the LOCAL buffers stored and loaded in the loop)
-  acc = acc.after(acc_store.end(k).end(k_tile))
+  # store accumulator and loop
+  acc = acc.after(acc_store.end(k).barrier().end(k_tile))

  # store accumulator to output (unified)
  c = c.reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M,
@@ -84,16 +84,16 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
    q.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
  K_store = KV_lds.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
    k[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
-  Q_lds = Q_lds.after(UOp.group(Q_store, K_store))
-  KV_lds_k = KV_lds.after(UOp.group(Q_store, K_store))
+  qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store))
+  Q_lds = Q_lds.after(qk_load_barrier)
+  KV_lds_k = KV_lds.after(qk_load_barrier)

  # -- S = Q @ K^T via WMMA (re-init each n_tile) --
  S_reg = UOp.placeholder((TM, TN), dtypes.float, slot=6, addrspace=AddrSpace.REG)
  S_reg = S_reg.after(S_reg.after(n_tile).store(S_reg.const_like(0)))
  k_qk = UOp.range(D // WMMA_K, 101, AxisType.REDUCE)
-  tm1 = UOp.range(TM // WMMA_ACC, 200)
-  tn1 = UOp.range(TN, 201)
+  tm1 = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
+  tn1 = UOp.range(TN, 201, AxisType.LOOP)
  S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
  q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
  k_frag = KV_lds_k.reshape(WAVES_N, TN, WMMA_N, D // WMMA_K, WMMA_K)[wave_n, tn1, lane_n, k_qk]
@@ -110,7 +110,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
  rm2 = UOp.range(TN, 261, AxisType.REDUCE)
  m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
  # warp reduce max (in-place)
-  ri_w = UOp.range(TM, 270)
+  ri_w = UOp.range(TM, 270, AxisType.LOOP)
  m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce_max(m_ij[ri_w], lane)).end(ri_w))

  # compute P = exp(S - m_ij) in S_reg
@@ -120,7 +120,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
  p_local = p_local.after(p_local.after(n_tile).store(p_local.const_like(0)))
  rp2 = UOp.range(TN, 291, AxisType.REDUCE)
  p_local = p_local.after(p_local.store(p_local.after(rp2) + S_reg[:, rp2]).end(rp2))
-  ri_ws = UOp.range(TM, 295)
+  ri_ws = UOp.range(TM, 295, AxisType.LOOP)
  p_sum = p_local.after(p_local[ri_ws].store(warp_reduce_sum(p_local[ri_ws], lane)).end(ri_ws))

  # write P = exp(S - m_ij) to P_lds (reuses slot 0, Q no longer needed)
@@ -130,11 +130,11 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
  P_store = P_write[tid].store(S_reg.cast(dtypes.half))

  # -- online softmax correction --
-  ri4 = UOp.range(TM, 330)
+  ri4 = UOp.range(TM, 330, AxisType.LOOP)
  m_new_val = m_i[ri4].maximum(m_ij[ri4])
  alpha_val = ((m_i[ri4] - m_new_val) * LOG2E).exp2()
  beta_val = ((m_ij[ri4] - m_new_val) * LOG2E).exp2()
-  rj4 = UOp.range(TD, 331)
+  rj4 = UOp.range(TD, 331, AxisType.LOOP)
  correction = UOp.group(
    acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4),
    l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]),
@@ -147,21 +147,21 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
  # load V into KV_lds (must wait for QK WMMA to finish reading K from KV_lds)
  V_store = KV_lds.after(qk_done).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
    v[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
-  P_lds = P_lds.after(UOp.group(P_store, V_store))
-  KV_lds_v = KV_lds.after(UOp.group(P_store, V_store))
+  pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
+  P_lds = P_lds.after(pv_barrier)
+  KV_lds_v = KV_lds.after(pv_barrier)

  # -- acc += P @ V via WMMA --
  k_pv = UOp.range(BLOCK_N // WMMA_K, 400, AxisType.REDUCE)
-  tm2 = UOp.range(TM // WMMA_ACC, 401)
-  tn2 = UOp.range(TD, 402)
+  tm2 = UOp.range(TM // WMMA_ACC, 401, AxisType.LOOP)
+  tn2 = UOp.range(TD, 402, AxisType.LOOP)
  acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
  p_frag = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
  v_frag = KV_lds_v.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
  pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), *WMMA_ARG)

  # end KV tile loop
-  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).end(n_tile)
+  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).barrier().end(n_tile)
  acc = acc.after(n_tile_end)
  l_i = l_i.after(n_tile_end)
  m_i = m_i.after(n_tile_end)
@@ -28,10 +28,10 @@ REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
 assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
 assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"

-def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.WEAK): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
+def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
 def copy(dest:UOp, src:UOp, rng:int, upcast=False):
  assert dest.shape == src.shape
-  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.WEAK)
+  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
  return dest[*rngs].store(src[*rngs]).end(*rngs)

 def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
@@ -66,8 +66,9 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
  B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
  B_local_store = copy(B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)

-  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
-  A_local, B_local = A_local.after(A_local_store, B_local_store), B_local.after(A_local_store, B_local_store)
+  # TODO: can we automate barrier?
+  barrier = UOp.barrier(A_local_store, B_local_store)
+  A_local, B_local = A_local.after(barrier), B_local.after(barrier)

  # open inner k range
  k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
@@ -101,7 +102,7 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iter_m, t_m] * B_row[iter_n, t_n]).end(iter_m, iter_n, t_m, t_n)

  # Close k, sync, and close K tiles
-  sink = sink.end(k).end(k_tile_range)
+  sink = sink.end(k).barrier().end(k_tile_range)

  # ---------------------------
  # REG -> GLOBAL (epilogue)
@@ -1,9 +1,8 @@
-import atexit, functools, math, pathlib
+import atexit, functools, pathlib
 from tinygrad import Tensor, Device, dtypes
-from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
 from tinygrad.renderer import Estimates
-from tinygrad.helpers import getenv, all_same, DEBUG, ceildiv
+from tinygrad.helpers import getenv, all_same, DEBUG
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8

@@ -73,7 +72,7 @@ def hk_fp8_atb_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, g_amax:Tensor|
    elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=dtypes.bfloat16, device=a.device), 2
    elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=dtypes.bfloat16, device=a.device), 1
    else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device), 0, True
-    out = Tensor(inv.uop.unshard(out_axis), device=a.device)
+    out = Tensor(inv.uop.multi(out_axis), device=a.device)
    dname = a.device[0]
  else:
    out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device)
@@ -108,23 +107,6 @@ def custom_hk_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:U
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
                               UOp(Ops.BINARY, arg=lib)))

-# ** MXFP4 GEMM custom kernel
-
-@functools.cache
-def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp, tile_m:int, tile_n:int) -> UOp:
-  from extra.gemm.gemm_mxfp4 import build_kernel
-  M, half_k = math.prod(A.shape[:-1]), A.shape[-1]
-  N, half_k_b = math.prod(B.shape[:-1]), B.shape[-1]
-  K = half_k * 2
-  assert half_k == half_k_b and math.prod(C.shape[:-1]) == M and C.shape[-1] == N
-  threads = UOp.special(256, "lidx0")
-  groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
-  lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
-  sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
-                  arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
-  insts = build_kernel(M, N, K, tile_m, tile_n)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))
-
 def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
  # 1x32 block scaling along the last axis
  *batch, K = x.shape
@@ -137,50 +119,6 @@ def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
  packed = mx_pack(e8) if len(batch) == 1 and scale_K % 4 == 0 else None
  return x_clamped.cast(FP8_DTYPE), e8, packed

-def _mxfp4_shuffle_weight(x:Tensor) -> Tensor:
-  # shuffle_weight(x, layout=(16, 16)) on the packed uint8 buffer.
-  if x.ndim == 3:
-    ndev, rows, half_k = x.shape
-    return x.reshape(ndev, rows//16, 16, half_k//32, 2, 16).permute(0, 1, 3, 4, 2, 5).reshape(ndev, rows, half_k).contiguous()
-  rows, half_k = x.shape
-  return x.reshape(rows//16, 16, half_k//32, 2, 16).permute(0, 2, 3, 1, 4).reshape(rows, half_k).contiguous()
-
-def _mxfp4_shuffle_scales(x:Tensor) -> Tensor:
-  # e8m0_shuffle: each 256x8 scale tile is arranged for the raw MFMA scale loads.
-  if x.ndim == 3:
-    ndev, rows, scale_k = x.shape
-    return x.reshape(ndev, rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 1, 4, 6, 3, 5, 2).reshape(ndev, rows, scale_k).contiguous()
-  rows, scale_k = x.shape
-  return x.reshape(rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 3, 5, 2, 4, 1).reshape(rows, scale_k).contiguous()
-
-def quantize_mxfp4(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
-  # OCP MXFP4: 1x32 blocks, E2M1 values packed low-nibble first, and E8M0 scales.
-  *batch, K = x.shape
-  rows = math.prod(batch)
-  assert x.ndim >= 2 and K % 256 == 0 and rows % 32 == 0, \
-    f"mxfp4 quantization needs rows%32 and K%256, got {x.shape}"
-  xb = x.float().reshape(*batch, K//32, 32)
-  amax = xb.abs().max(axis=-1)
-
-  # even scale rounding: round the fp32 significand before choosing 2^(floor(log2)-2).
-  amax_rounded = ((amax.bitcast(dtypes.uint32) + 0x200000) & 0xFF800000).bitcast(dtypes.float32)
-  scale_exp = (amax_rounded.maximum(2**-126).log2().floor() - 2).clamp(-127, 127)
-  e8 = (scale_exp + 127).cast(dtypes.uint8)
-  scaled = xb * (-scale_exp).exp2().reshape(*batch, K//32, 1)
-
-  mag = scaled.abs()
-  code = sum(x.cast(dtypes.uint8) for x in
-             (mag > .25, mag >= .75, mag > 1.25, mag >= 1.75, mag > 2.5, mag >= 3.5, mag > 5.0))
-  code = code | ((scaled < 0).cast(dtypes.uint8) << 3)
-  code = code.reshape(*batch, K)
-  packed = code[..., 0::2] | (code[..., 1::2] << 4)
-  if isinstance(x.device, tuple) and x.uop.axis == x.ndim-2 and x.shape[x.uop.axis] == len(x.device):
-    axis = x.uop.axis
-    order = (axis, *range(axis), *range(axis+1, e8.ndim))
-    e8_local = e8.permute(order)
-    return packed, e8, _mxfp4_shuffle_scales(e8_local.reshape(e8_local.shape[0], -1, K//32))
-  return packed, e8, _mxfp4_shuffle_scales(e8.reshape(rows, K//32))
-
 def mx_pack(e8:Tensor) -> Tensor:
  rows, scale_K = e8.shape
  return e8.reshape(rows, scale_K // 4, 4).bitcast(dtypes.uint32).reshape(rows, scale_K // 4).permute(1, 0).contiguous()
@@ -233,13 +171,13 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
  M, K = A.shape[0]*A.shape[1], A.shape[2]
  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
  assert K == K2
-  m = UOp.range(M, 1)
-  n = UOp.range(N, 2)
+  m = UOp.range(M, 1, AxisType.LOOP)
+  n = UOp.range(N, 2, AxisType.LOOP)
  k = UOp.range(K, 0, AxisType.REDUCE)
-  mul = (A.flatten().index((m*UOp.const(K)+k))*
-         B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
+  mul = (A.flatten().index((m*UOp.const(dtypes.weakint, K)+k))*
+         B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)
  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
-  store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n)
+  store = C.flatten().index((m*UOp.const(dtypes.weakint, N)+n)).store(red).end(m, n)
  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))

 # ** bf16 A @ B.T kernel in C
@@ -296,7 +234,7 @@ def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
    elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=a.dtype, device=a.device), 2
    elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=a.dtype, device=a.device), 1
    else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device), 0, True
-    out = Tensor(inv.uop.unshard(out_axis), device=a.device)
+    out = Tensor(inv.uop.multi(out_axis), device=a.device)
    dname = a.device[0]
  else:
    out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device)
@@ -306,6 +244,7 @@ def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
  if reduce_out: out = out.sum(0)
  return out.squeeze(0) if out.ndim == 3 else out

+
 # ** backward gemm, might use the asm gemm

 def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=False, has_w_post:bool=False):
@@ -402,29 +341,13 @@ def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=F
  if wp is not None: grad_b = grad_b / wp.reshape(-1, 1)
  return (None, grad_a.uop, grad_b.uop) + tuple(None for _ in inputs[3:])

-# ** mxfp4 gemm backward
-
-def custom_mxfp4_gemm_bw(gradient:UOp, kernel:UOp):
-  # The raw kernel consumes quantized buffers, while the final two inputs retain the BF16 operands for STE gradients.
-  inputs = kernel.src[1:]  # (out, a_q, b_q, scale_a, scale_b, a, w)
-  assert len(inputs) == 7
-  a, w = Tensor(inputs[5], device=inputs[5].device), Tensor(inputs[6], device=inputs[6].device)
-  g = Tensor(gradient, device=a.device)[:a.shape[0]].cast(dtypes.bfloat16)
-  grad_a = asm_gemm(g, w, mxfp4=True)
-  a_flat, g_flat = a.reshape(-1, a.shape[-1]), g.reshape(-1, g.shape[-1])
-  grad_w = asm_gemm(g_flat.T, a_flat, mxfp4=True)
-  return (None, None, None, None, None, grad_a.uop, grad_w.uop)
-
 # ** main gemm function

 def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None, grad_amax_state:Tensor|None=None,
             next_grad_amax_state:Tensor|None=None,
             w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_amax:Tensor|None=None,
-             a_pretranspose:Tensor|None=None, mxfp4:bool=False) -> Tensor:
+             a_pretranspose:Tensor|None=None) -> Tensor:
  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
-  if mxfp4:
-    assert not mx and mx_scales is None, "mxfp4 owns quantization; mx/mx_scales are for mxfp8"
-    assert a.dtype == dtypes.bfloat16, f"cannot quantize {a.dtype} to mxfp4"
  counters["used"] += 1
  unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
  if unfold_batch:
@@ -432,7 +355,7 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
    a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
  squeeze = a.ndim == 2
  if squeeze: a = a.unsqueeze(0)
-  out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE or mxfp4 else a.dtype
+  out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE else a.dtype

  batch, M, K = a.shape
  N = b.shape[1]
@@ -443,11 +366,11 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N

  if is_multi:
    if n_sharded:
-      out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.unshard(2), device=a.device)
+      out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.multi(2), device=a.device)
    elif m_sharded:
-      out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(1), device=a.device)
+      out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.multi(1), device=a.device)
    else:
-      out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(0),
+      out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.multi(0),
                   device=a.device)
  else:
    out = Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device)
@@ -455,21 +378,7 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
  renderer = Device[dname:=(a.device[0] if is_multi else a.device)].renderer
  dname, arch = dname.split(":")[0], renderer.target.arch
  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    if mxfp4:
-      tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if (batch*M) % tm == N % tn == 0)
-      fxn = functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n)
-      w = b.T
-      if k_sharded:
-        ndev = len(a.device)
-        a_q, _, scale_a = quantize_mxfp4(a.reshape(batch, M, ndev, K))
-        b_q, _, scale_b = quantize_mxfp4(w.reshape(w.shape[0], ndev, K))
-        b_q = _mxfp4_shuffle_weight(b_q.permute(1, 0, 2))
-      else:
-        a_q, _, scale_a = quantize_mxfp4(a.reshape(batch*M, K))
-        b_q, _, scale_b = quantize_mxfp4(w)
-        a_q, b_q = a_q.reshape(batch, M, K//2).contiguous(), _mxfp4_shuffle_weight(b_q)
-      out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, a, w, fxn=fxn, grad_fxn=custom_mxfp4_gemm_bw)[0]
-    elif mx:
+    if mx:
      # mxfp8 1x32 block scaling
      if mx_scales is not None:
        a_si, a_e8, b_si, b_e8 = mx_scales
@@ -79,7 +79,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
  # this is the big accumulator
  acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float, 0, AddrSpace.REG)
  assert acc.size*WARP_SIZE*WARPGROUP_SIZE*4 == BLOCK_M*BLOCK_N
-  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const((0.0,)*4, dtypes.float), end=init_l)
+  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float, (0.0,)*4), end=init_l)

  # create locals (note A is permuted, and the stride is changed to avoid bank conflicts)
  def make_locals(slot) -> tuple[UOp, UOp]:
@@ -29,7 +29,7 @@ TID_SIZE = WARPGROUP_SIZE*WARP_SIZE

 def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=()):
  assert dest.shape == src.shape
-  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.WEAK) for i,s in enumerate(src.shape)]
+  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.LOOP) for i,s in enumerate(src.shape)]
  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
  return dest.after(copy) if set else copy

@@ -55,7 +55,7 @@ def grouped_mx_wgrad(g:Tensor, xg:Tensor, expert_off:Tensor, n_experts:int) -> T
  dname = (g.device[0] if isinstance(g.device, tuple) else g.device).split(":")[0]
  is_multi = isinstance(g.device, tuple)
  inv = Tensor.invalids(1, n_experts * N, K, dtype=dtypes.bfloat16, device=g.device)
-  out = Tensor(inv.uop.unshard(0), device=g.device) if is_multi else inv
+  out = Tensor(inv.uop.multi(0), device=g.device) if is_multi else inv
  out = Tensor.custom_kernel(out, gT, xT, g_si, x_si, expert_off,
                             fxn=functools.partial(custom_hk_grouped_mxfp8_wgrad, dname=dname, n_experts=n_experts))[0]
  out = out.sum(0) if is_multi else out.squeeze(0)
@@ -103,7 +103,7 @@ def grouped_mx_gemm(x:Tensor, w:Tensor|tuple[Tensor, Tensor], expert_off:Tensor)
  if isinstance(x.device, tuple) and (row_axis := x.uop.axis) is not None:
    ndev = len(x.device)
    out = Tensor(Tensor.invalids(*(s // ndev if i == row_axis else s for i, s in enumerate(out_shape)),
-                                 dtype=dtypes.bfloat16, device=x.device).uop.unshard(row_axis), device=x.device)
+                                 dtype=dtypes.bfloat16, device=x.device).uop.multi(row_axis), device=x.device)
  else:
    out = Tensor.invalids(*out_shape, dtype=dtypes.bfloat16, device=x.device)
  return Tensor.custom_kernel(out, x_q, w_q, x_si, w_si, xe_in, w_e8, expert_off,
@@ -1,130 +0,0 @@
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-
-BLOCK_ROW = 256
-
-def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor:
-  if isinstance(device, tuple):
-    return Tensor(Tensor.invalids(shape[0] // len(device), *shape[1:], dtype=dtype, device=device).uop.multi(0), device=device)
-  return Tensor.invalids(*shape, dtype=dtype, device=device)
-
-def _atomic_add(device:str) -> str:
-  return "__hip_atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);" if device == "AMD" \
-    else "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"
-
-def _blk_for(D:int) -> int:
-  blk = 64
-  while D % blk: blk //= 2
-  return blk
-
-def _kv_ranges(G, N, D, BLK):
-  g = UOp.range(G, 0)
-  m = UOp.range(N, 1)
-  jo = UOp.range(D // BLK, 2)
-  ji = UOp.range(BLK, 3, AxisType.LOCAL)
-  return g, m, jo * BLK + ji, jo, ji
-
-def _ggather_fwd_kernel(out:UOp, table:UOp, idx:UOp) -> UOp:
-  G, M, D = out.shape
-  g, m, j, jo, ji = _kv_ranges(G, M, D, _blk_for(D))
-  row = idx.index(g, m).cast(dtypes.weakint)
-  val = table.index(g, row, j).load()
-  return out.index(g, m, j).store(val).end(g, m, jo, ji).sink(
-    arg=KernelInfo(name=f"ggather_fwd_{M}_{D}", opts_to_apply=()))
-
-def _ggather_zero_kernel(out:UOp) -> UOp:
-  i = UOp.range(out.numel(), 0)
-  return out.flatten().index(i).store(UOp.const(0.0, out.dtype)).end(i).sink(arg=KernelInfo(name="ggather_zero"))
-
-def _sharded_zeros(shape:tuple[int, ...], dtype, device) -> Tensor:
-  return Tensor.custom_kernel(_sharded_invalids(shape, dtype, device), fxn=_ggather_zero_kernel)[0]
-
-def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
-  _, table_u, idx_u = kernel.src[1:4]
-  dev = table_u.device
-  device = (dev[0] if isinstance(dev, tuple) else dev).split(":")[0]
-  G, R, D = table_u.shape
-  gt = _sharded_zeros((G, R, D), dtypes.float32, dev)
-  go = Tensor(gradient, device=dev)
-  atomic_str = _atomic_add(device)
-  def _bwd_kernel(gtab:UOp, gout:UOp, idx:UOp) -> UOp:
-    Gk, M, Dk = gout.shape
-    g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
-    row = idx.index(g, m).cast(dtypes.weakint)
-    val = gout.index(g, m, j).load().cast(dtypes.float32)
-    atomic = UOp(Ops.CUSTOM, dtypes.void, (gtab.index(g, row, j), val), arg=atomic_str)
-    return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=()))
-  grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0]
-  return (None, grad_table.cast(table_u.dtype).uop, None)
-
-def grouped_gather_rows(table:Tensor, idx:Tensor, n_groups:int) -> Tensor:
-  G, R, D = table.shape
-  M = idx.shape[1]
-  out = _sharded_invalids((G, M, D), table.dtype, table.device)
-  return Tensor.custom_kernel(out, table, idx, fxn=_ggather_fwd_kernel, grad_fxn=_ggather_bwd)[0]
-
-def _gscatter_fwd_kernel(out:UOp, src:UOp, idx:UOp) -> UOp:
-  G, M, D = out.shape
-  k = idx.shape[1] // src.shape[1]
-  g, m, j, jo, ji = _kv_ranges(G, idx.shape[1], D, _blk_for(D))
-  row = idx.index(g, m).cast(dtypes.weakint)
-  val = src.index(g, (m // k).cast(dtypes.weakint), j).load()
-  return out.index(g, row, j).store(val).end(g, m, jo, ji).sink(
-    arg=KernelInfo(name=f"gscatter_fwd_{idx.shape[1]}_{D}", opts_to_apply=()))
-
-def _gscatter_bwd(gradient:UOp, kernel:UOp) -> tuple:
-  _, src_u, idx_u = kernel.src[1:4]
-  dev = src_u.device
-  G, T_l, D = src_u.shape
-  k = idx_u.shape[1] // T_l
-  sel = grouped_gather_rows(Tensor(gradient, device=dev), Tensor(idx_u, device=dev), G)
-  return (None, sel.reshape(G, T_l, k, D).sum(2).cast(src_u.dtype).uop, None)
-
-def grouped_scatter_rows(src:Tensor, idx:Tensor, m_l:int) -> Tensor:
-  G, T_l, D = src.shape
-  zero = _sharded_zeros((G, m_l, D), src.dtype, src.device)
-  return Tensor.custom_kernel(zero, src, idx, fxn=_gscatter_fwd_kernel, grad_fxn=_gscatter_bwd)[0]
-
-def m_max_for(t_local:int, experts_per_tok:int, n_experts:int) -> int:
-  return (-(-t_local * experts_per_tok // BLOCK_ROW) + n_experts) * BLOCK_ROW
-
-class Routing:
-  def __init__(self, weights:Tensor, dest_row:Tensor, off:Tensor, m_l:int, n_groups:int, t_local:int):
-    self.weights, self.dest_row = weights, dest_row
-    self.off = off
-    self.m_l, self.n_groups, self.t_local = m_l, n_groups, t_local
-
-  @property
-  def rows_e(self) -> Tensor:
-    G, E = self.off.shape[0], self.off.shape[1] - 1
-    tr = Tensor.arange(self.m_l // BLOCK_ROW, dtype=dtypes.int32).reshape(1, -1, 1) * BLOCK_ROW
-    tr = tr.shard(self.off.device) if isinstance(self.off.device, tuple) else tr.to(self.off.device)
-    tile_e = ((tr >= self.off[:, :E].reshape(G, 1, E)).sum(-1) - 1).cast(dtypes.int32)
-    return tile_e.reshape(-1, 1).expand(-1, BLOCK_ROW).reshape(-1)
-
-def n_groups_of(t:Tensor) -> int:
-  return len(t.device) if isinstance(t.device, tuple) else 1
-
-def route(logits:Tensor, experts_per_tok:int, n_experts:int) -> Routing:
-  T, E = logits.shape
-  k, G = experts_per_tok, n_groups_of(logits)
-  assert T % G == 0, f"tokens {T} must split across {G} devices"
-  T_l, m_l = T // G, m_max_for(T // G, k, n_experts)
-
-  topv, topi = logits.reshape(G, T_l, E).topk(k)
-  weights = topv.softmax(-1)
-  m = topi.reshape(G, T_l * k).cast(dtypes.int32).one_hot(E).cast(dtypes.int32)
-
-  pad = ((m.sum(1) + (BLOCK_ROW - 1)) // BLOCK_ROW) * BLOCK_ROW
-  off = pad.cumsum(1).pad(((0, 0), (1, 0)))
-  dest_row = ((m.cumsum(1) + off[:, :E].reshape(G, 1, E)) * m).sum(-1).sub(1).cast(dtypes.int32)
-  return Routing(weights, dest_row, off, m_l, G, T_l)
-
-def dispatch(x:Tensor, r:Routing) -> Tensor:
-  G, D = r.n_groups, x.shape[-1]
-  return grouped_scatter_rows(x.reshape(G, r.t_local, D), r.dest_row, r.m_l).reshape(G * r.m_l, D)
-
-def combine(y:Tensor, r:Routing, n_tokens:int, experts_per_tok:int) -> Tensor:
-  G, D, k = r.n_groups, y.shape[-1], experts_per_tok
-  sel = grouped_gather_rows(y.reshape(G, r.m_l, D), r.dest_row, G).reshape(G, r.t_local, k, D)
-  return (sel * r.weights.reshape(G, r.t_local, k, 1).cast(sel.dtype)).sum(2).reshape(n_tokens, D).cast(y.dtype)
@@ -1,139 +0,0 @@
-"""
-tilelang-style matmul_relu written with tinygrad UOp APIs.
-
-Reference tilelang kernel:
-  @tilelang.jit
-  def matmul_relu(A, B, block_M=64, block_N=64, block_K=64,
-                  dtype=T.float16, accum_dtype=T.float32):
-    M, N, K = T.const('M, N, K')
-    C = T.empty([M, N], dtype)
-    with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
-      A_shared = T.alloc_shared((block_M, block_K), dtype)
-      B_shared = T.alloc_shared((block_K, block_N), dtype)
-      C_local  = T.alloc_fragment((block_M, block_N), accum_dtype)
-      T.clear(C_local)
-      for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
-        T.copy(A[by * block_M, ko * block_K], A_shared)
-        T.copy(B[ko * block_K, bx * block_N], B_shared)
-        T.gemm(A_shared, B_shared, C_local)
-      for i, j in T.Parallel(block_M, block_N):
-        C_local[i, j] = T.max(C_local[i, j], 0)
-      T.copy(C_local, C[by * block_M, bx * block_N])
-    return C
-"""
-
-from tinygrad.dtype import dtypes, AddrSpace, DType
-from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
-from tinygrad.helpers import cdiv, getenv
-from tinygrad.tensor import Tensor
-
-# ---------------------------------------------------------------------------
-# tilelang builtins, expressed with tinygrad UOp APIs
-# ---------------------------------------------------------------------------
-
-def alloc_shared(shape:tuple[int, ...], dtype:DType, slot:int) -> UOp:
-  """T.alloc_shared: one LOCAL buffer shared by all threads in the block."""
-  return UOp.placeholder(tuple(shape), dtype, slot, AddrSpace.LOCAL)
-
-def alloc_fragment(shape:tuple[int, ...], dtype:DType, slot:int, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp:
-  """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid."""
-  assert len(axes) == len(rngs)
-  assert all(tnum.op is Ops.RANGE and tnum.arg[-1] is AxisType.LOCAL for tnum in rngs), "fragments shard over LOCAL ranges"
-  by_axis = dict(zip(axes, rngs))
-  shard_shape = tuple(s // (int(by_axis[i].vmax)+1) if i in by_axis else s for i, s in enumerate(shape))
-  fragment = UOp.placeholder(shard_shape, dtype, slot, AddrSpace.REG)
-  return fragment.unshard(axes, rngs)
-
-# ---------------------------------------------------------------------------
-# GEMM kernel: C = relu(A @ B), float inputs (fp16 or fp32), fp32 fragment accumulator, no WMMA
-# ---------------------------------------------------------------------------
-
-# 64x64 output tile per block, 128 threads as an 8x16 grid; each thread owns an 8x4 fragment sub-tile
-# (the 2-D per-thread layout tilelang infers for this GEMM). The 4 contiguous columns (TN=4) are what
-# let codegen vectorize loads/stores to float4, matching tilelang's lowering exactly.
-BLOCK_M = BLOCK_N = BLOCK_K = 64
-TY = 8
-TX = 16
-THREADS = TY * TX
-TM = BLOCK_M // TY   # fragment rows per thread  (8)
-TN = BLOCK_N // TX   # fragment columns per thread (4)
-
-def matmul_relu_kernel(c:UOp, a:UOp, b:UOp) -> UOp:
-  """C[M, N] = relu(A[M, K] @ B[K, N]) -- one 64x64 tile per block, locals + a 2-D fragment."""
-  M, K = a.shape
-  K2, N = b.shape
-  assert K == K2 and a.dtype == b.dtype == c.dtype and not dtypes.is_int(a.dtype)
-  assert not (K % BLOCK_K or M % BLOCK_M or N % BLOCK_N), "test sizes must be multiples of the block sizes"
-
-  # with T.Kernel(T.ceildiv(N, BLOCK_N), T.ceildiv(M, BLOCK_M), threads=128) as (bx, by):
-  bx = UOp.range(cdiv(N, BLOCK_N), 0, AxisType.GLOBAL)
-  by = UOp.range(cdiv(M, BLOCK_M), 1, AxisType.GLOBAL)
-
-  # 16*8 threads = 128 threads
-  tx = UOp.range(TX, 2, AxisType.LOCAL)
-  ty = UOp.range(TY, 3, AxisType.LOCAL)
-
-  # shared + fragment (regs)
-  A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype, 0)
-  B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype, 1)
-  C_local = alloc_fragment((TM, TY, TX, TN), dtypes.float32, 0, (1, 2), (ty, tx))
-
-  # zero out the regs to start. this is expanded by the devectorizer
-  C_local = C_local.after(C_local.store(0.0))
-
-  # for ko in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3):
-  ko = UOp.range(cdiv(K, BLOCK_K), 6, AxisType.LOOP)
-
-  # index the outer matrices
-  a = a.rearrange("(m bm) (k bk) -> m k bm bk", bm=BLOCK_M, bk=BLOCK_K)[by, ko]
-  b = b.rearrange("(k bk) (n bn) -> k n bk bn", bk=BLOCK_K, bn=BLOCK_N)[ko, bx]
-  c = c.rearrange("(m bm) (n bn) -> m n bm bn", bm=BLOCK_M, bn=BLOCK_N)[by, bx]
-
-  # T.copy: A_shared <- a, B_shared <- b
-  def with_threads(x:UOp): return x.rearrange("(tm ty) (tx tn) -> ty tx tm tn", tm=TM, tn=TN)[ty, tx]
-  A_shared = A_shared.after(with_threads(A_shared).store(with_threads(a)))
-  B_shared = B_shared.after(with_threads(B_shared).store(with_threads(b)))
-
-  # T.gemm(A_shared, B_shared, C_local), no WMMA
-  kk = UOp.range(BLOCK_K, 11, AxisType.LOOP)
-  ir = UOp.range(TM, 12, AxisType.LOOP)
-  jj = UOp.range(TN, 13, AxisType.UPCAST)
-  acc = C_local.after(kk)[ir, ty, tx, jj] + A_shared[ir*TM + ty, kk].cast(dtypes.float32) * B_shared[kk, tx*TN + jj].cast(dtypes.float32)
-  # closing the ko loop here too; codegen adds the barrier so no thread overwrites the tiles while others still read them
-  C_local = C_local[ir, ty, tx, jj].set(acc, end=(kk, ir, jj, ko))
-
-  # c <- C_local (with relu and cast): every thread stores its shard's sub-view of the output tile
-  c_st = c.reshape(C_local.shape).store(C_local.relu().cast(c.dtype))
-
-  # close the locals and globals
-  return c_st.end(tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=()))
-
-# ---------------------------------------------------------------------------
-# python wrapper: same signature as the tilelang function
-# ---------------------------------------------------------------------------
-
-def matmul_relu(a:Tensor, b:Tensor) -> Tensor:
-  """C = relu(A @ B), fp16 in/out with an fp32 fragment accumulator."""
-  c = Tensor.empty(a.shape[0], b.shape[1], dtype=a.dtype, device=a.device)
-  return c.custom_kernel(a, b, fxn=matmul_relu_kernel)[0]
-
-# ---------------------------------------------------------------------------
-# test
-# ---------------------------------------------------------------------------
-
-if __name__ == "__main__":
-  from tinygrad import Device
-  assert Device[Device.DEFAULT].renderer.has_local, "this GPU-style kernel needs a backend with local memory (LOCAL ranges + barriers)"
-  M = K = N = getenv("N", 256)  # 4x4 grid of 64x64 tiles, 4 K chunks
-  dtype_in = dtypes.half if getenv("HALF") else dtypes.float
-
-  a = Tensor.randn(M, K, dtype=dtype_in).contiguous()
-  b = Tensor.randn(K, N, dtype=dtype_in).contiguous()
-  ref = (a @ b).relu().realize()
-
-  for _ in range(10):
-    out = matmul_relu(a, b).realize()
-
-  import numpy as np
-  np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=1e-1, rtol=1e-2)
-  print("matmul_relu passed!")
@@ -1,18 +1,18 @@
 from __future__ import annotations
-from typing import cast, Callable, TypeVar, Generic, Any, Sequence
+from typing import cast, Callable, TypeVar, Generic, Any
 import struct, functools, time, collections, itertools
 from dataclasses import replace, dataclass
 from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap
 from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar
-from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer, DepsTracker
+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, GroupOp
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.dtype import dtypes, truncate
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.memory import BumpAllocator
 from tinygrad.renderer import Renderer, Estimates
-from tinygrad.engine.realize import to_program, get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop
-from tinygrad.engine.realize import pm_flatten_linear
+from tinygrad.engine.realize import to_program, get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop, pm_flatten_linear
+from tinygrad.engine.jit import DepsTracker

 # *****************
 # 0. helpers
@@ -38,31 +38,28 @@ class HCQInfo:
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c

 def unwrap_mstack(u):
-  if u.op is Ops.MSTACK: return tuple(x for s in u.src for x in unwrap_mstack(s))
-  return unwrap_mstack(u.src[0]) if u.op in {Ops.MSELECT, Ops.SLICE} else (u,)
+  return tuple(x for s in u.src for x in unwrap_mstack(s)) if u.op is Ops.MSTACK else (unwrap_mstack(u.src[0]) if u.op in {Ops.MSELECT, Ops.SLICE} else (u,))

-def make_patches(buf:UOp, patches:Sequence[tuple[sint, UOp]]) -> UOp:
-  offsets = UOp(Ops.STACK, dtypes.int, tuple(UOp.const(off // buf.dtype.itemsize, dtypes.int) for off,_ in patches))
-  values = UOp(Ops.STACK, buf.dtype, tuple(val.cast(buf.dtype) for _,val in patches))
-  return buf.index(offsets).store(values)
+def make_patch(buf:UOp, off:sint, val:UOp) -> UOp:
+  return buf.index(UOp.const(dtypes.int, off // buf.dtype.itemsize)).store(val.simplify().cast(buf.dtype))

 def make_binary_patch(buf:UOp, blob:bytes) -> UOp:
  data = UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype)
  r = UOp.range(len(blob) // buf.dtype.itemsize, 0, dtype=dtypes.int, src=(buf, data))
  return buf.index(r).store(data.index(r).load()).end(r)

-def make_cmdbuf(lin, devs, buf:UOp|None=None, dep:UOp|None=None):
-  blob, patches = bytearray(), []
+def make_cmdbuf(lin, devs):
+  blob, patches = b'', []
  for s in (s for ins in lin.src for s in ins.src):
-    if s.op is not Ops.CONST: patches.append((len(blob), s))
-    blob.extend(struct.pack(f'<{s.dtype.fmt}', s.val if s.op is Ops.CONST else 0x0))
-  cmdbuf = buf if buf is not None else UOp.placeholder((len(blob) // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("cmdbuf")
-  writable = cmdbuf.after(dep) if dep is not None else cmdbuf
-  return cmdbuf.after(make_binary_patch(writable, bytes(blob)), *((make_patches(writable, patches),) if patches else ()))
+    if (ssimp:=s.simplify()).op is not Ops.CONST: patches.append((len(blob), ssimp))
+    blob += struct.pack(f'<{ssimp.dtype.fmt}', ssimp.arg if ssimp.op is Ops.CONST else 0x0)
+  cmdbuf = UOp.placeholder((len(blob) // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("cmdbuf")
+  return cmdbuf.after(make_binary_patch(cmdbuf, blob), *[make_patch(cmdbuf, off, s) for off, s in patches])

 def make_signal(devs, queue="COMPUTE:0", sentinel=False):
-  return UOp.placeholder((1,), dtypes.uint64, 0, device=devs, volatile=True).rtag("sentinel_signal" if sentinel else f"{queue}_timeline_signal")
-def make_signal_value(devs, queue="COMPUTE:0"): return UOp.placeholder((1,), dtypes.uint64, 0, device=devs).rtag(f"{queue}_timeline_value")
+  return UOp.placeholder((1,), dtypes.uint64, 0, device=devs).rtag("sentinel_signal" if sentinel else f"{queue}_timeline_signal")
+def make_signal_value(devs, queue="COMPUTE:0"):
+  return UOp.placeholder((1,), dtypes.uint64, 0, device=devs).rtag(f"{queue}_timeline_value")

 def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
  return UOp.custom_function("submit_cmdbuf", UOp(Ops.LINEAR, src=tuple(cmds), arg=(to_tuple(devs), queue)))
@@ -72,16 +69,14 @@ def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp:
  data, info = prg.arg
  buf = UOp.placeholder((data.kernargs_alloc_size // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("kernargs")
  words = [w for gi in info.globals for w in data64_le(get_call_arg_uops(call)[gi].getaddr(devs))] + list(info.vars)
-  return buf.after(*((make_patches(buf, [(i * 4, w) for i, w in enumerate(words)]),) if words else ()))
+  return buf.after(*[make_patch(buf, i * 4, w) for i, w in enumerate(words)])

 # *****************
 # 0.1. prep: replace buffers with params

-def replace_call_buffers(ctx:tuple[list[UOp], dict[UOp, int]], call:UOp) -> UOp|None:
-  bufs, slots = ctx
-  for s in call.src[1:]:
-    if s.op not in (Ops.PARAM, Ops.BIND) and slots.setdefault(s, len(bufs)) == len(bufs): bufs.append(s)
-  return call.replace(src=call.src[:1] + tuple(s if s.op in (Ops.PARAM, Ops.BIND) else s.param_like(slots[s]) for s in call.src[1:]))
+def replace_call_buffers(ctx:list[UOp], call:UOp) -> UOp|None:
+  ctx += [s for s in call.src[1:] if s not in ctx and s.op not in (Ops.PARAM, Ops.BIND)]
+  return call.replace(src=call.src[:1] + tuple(s if s.op in (Ops.PARAM, Ops.BIND) else s.param_like(ctx.index(s)) for s in call.src[1:]))
 pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_buffers)])

 # *****************
@@ -111,8 +106,7 @@ def _get_call_bufs_by_lane(call:UOp, devices:tuple[str, ...]) -> list[list[Any]]
 def _get_deps(ctx:DepsTracker, bufs_by_lane:list[list[Any]], write, key:tuple[tuple[str, ...], str, int]) -> list[tuple[tuple, int, int]]:
  dep_lanes:list[tuple[tuple, int, int]] = []
  for lane, bufs in enumerate(bufs_by_lane):
-    written = write if write is not None else list(range(len(bufs)))
-    dep_lanes += [(dep, dlane, lane) for dep, dlane in ctx.access_resources(bufs, written, (key, lane))]
+    dep_lanes += [(dep, dlane, lane) for dep, dlane in ctx.access_resources(bufs, write if write is not None else range(len(bufs)), (key, lane))]
  return dep_lanes

 def _build_wait_cmds(dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, ...], queue:str) -> tuple[list[UOp], set[int]]:
@@ -129,7 +123,7 @@ def _build_wait_cmds(dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str,
  for (ddevs, dqueue, dtag), lanes in deps.items():
    sig = UOp.mstack(*[make_signal(d if dl is None else ddevs[dl], queue=dqueue, sentinel=dl is None) for dl, d in zip(lanes, devices)])
    val = UOp.mstack(*[make_signal_value(d if dl is None else ddevs[dl], queue=dqueue) for dl, d in zip(lanes, devices)])
-    waits.append(UOp(Ops.INS, arg="wait", src=(sig, val.index(UOp.const(0, dtypes.int)) + dtag)))
+    waits.append(UOp(Ops.INS, arg="wait", src=(sig, val.index(UOp.const(dtypes.int, 0)) + dtag)))
  return waits, {dtag for _, _, dtag in deps}

 def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]],
@@ -140,8 +134,8 @@ def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[t
    for b in itertools.chain.from_iterable(_get_call_bufs_by_lane(call, devices)):
      for bd in to_tuple(b.device): dev_bufs[bd][id(b)] = b

-  zero, n, submits, bumps, waited = UOp.const(0, dtypes.int), len(batch_info), [], [], set()
-  for _, devgroup in itertools.groupby(sorted(dev_bufs), key=lambda d: d.split(":")[0]):
+  zero, n, finalizers, waited = UOp.const(dtypes.int, 0), len(batch_info), [], set()
+  for _, devgroup in itertools.groupby(sorted(dedup([d for devs, _ in batch_info for d in devs])), key=lambda d: d.split(":")[0]):
    devs = tuple(devgroup)

    # to finalize the batch, sync all accesses from other devices to buffers that belong to this device
@@ -149,16 +143,14 @@ def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[t
    waits, cur_waited = _build_wait_cmds(fin_deps, devs, "COMPUTE:0")
    waited |= cur_waited

-    # wait the syncs, store the device epoch
+    # wait the syncs, store the device epoch; value bumps are a separate call: no lane may bump until every lane has patched its waits
    store = UOp(Ops.INS, arg="store", src=(make_signal(devs), (tl:=make_signal_value(devs)).index(zero) + n))
-    submits.append((devs, make_submit(*waits, store, devs=devs, queue="COMPUTE:0")))
+    submit = make_submit(*waits, store, devs=devs, queue="COMPUTE:0")
    upd = [(tl, n + 1)] + [(make_signal_value(devs, queue=qn), n)
                           for qn in dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)]) if qn != "COMPUTE:0"]
-    bumps.append((devs, UOp.barrier(*[s.index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd])))
-
-  # NOTE: submit before bumps
-  fins = [UOp.custom_function("hcq", b.sink()).call(aux=HCQInfo("hcq_finalizer", Estimates(), devs, "COMPUTE:0")) for devs, b in submits + bumps]
-  return fins, waited
+    bump = UOp.barrier(*[s.index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd])
+    finalizers += [UOp.custom_function("hcq", b.sink()).call(aux=HCQInfo("hcq_finalizer", Estimates(), devs, "COMPUTE:0")) for b in (submit, bump)]
+  return finalizers, waited

 def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]:
  batch_info = [(devices, "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0") for call, devices in batch]
@@ -199,6 +191,7 @@ def sched_hcq_batches(l:UOp) -> UOp:
    if (devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is not None: batch.append((call, to_tuple(devs)))
    else: srcs, batch = srcs + _finalize_batch(batch) + [call], []
  return l.replace(src=tuple(srcs + _finalize_batch(batch)))
+pm_sched_hcq_batches = PatternMatcher([(UPat(Ops.LINEAR, name="l"), sched_hcq_batches)])

 # *****************
 # 3. merge into queues
@@ -213,7 +206,7 @@ def _merged_hcq_call(calls:list[UOp]) -> UOp: # TODO: simplify?
 def merge_queues(linear:UOp) -> UOp:
  new_src:list[UOp] = []
  opened_qs:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> list of hcq calls, kept in submit order
-  limits:dict[tuple[tuple[str, ...], str], int] = collections.defaultdict(lambda: JIT_BATCH_SIZE.value)
+  limits = collections.defaultdict(lambda: JIT_BATCH_SIZE.value)

  for call in linear.src:
    if not isinstance(info:=call.arg.aux, HCQInfo) or info.name == "hcq_finalizer": # non-hcq call or finalizer: close all open queues
@@ -230,10 +223,7 @@ def merge_queues(linear:UOp) -> UOp:
      new_rec = [call]
    opened_qs[(info.device, info.queue)] = new_rec
  return linear.replace(src=tuple(new_src + [_merged_hcq_call(c) for c in opened_qs.values()]))
-
-def schedule_and_merge(ctx:dict[UOp, UOp], linear:UOp) -> UOp:
-  return merge_queues(sched_hcq_batches(linear).substitute(ctx, walk=True, enter_calls=True))
-pm_schedule_and_merge = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), schedule_and_merge)])
+pm_merge_queues = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), merge_queues)])

 # *****************
 # 4.2. hcq lowering: ops to ir
@@ -241,8 +231,7 @@ pm_schedule_and_merge = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), schedu
 def encode_cmdbuf(submit:UOp, lin:UOp) -> UOp|None:
  if (pm:=Device.get_class(lin.arg[0][0]).pm_lower) is None: return None
  return graph_rewrite(submit, pm, name=f"encode {lin.arg[0]}", enter_calls=True)
-pm_encode_cmdbufs = PatternMatcher([
-  (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="lin"),), name="submit"), encode_cmdbuf)])
+pm_encode_cmdbufs = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="lin"),), name="submit"), encode_cmdbuf)])

 # *****************

@@ -254,8 +243,7 @@ def is_value_known_at_link(val:UOp) -> bool:
  return not val.variables() and not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs)

 def is_link_patch(p:UOp, jit:bool) -> bool:
-  if p.tag == "link": return True
-  store = p.src[0] if (is_binary_patch:=(p.op is Ops.END and p.src[0].op is Ops.STORE)) else p
+  store = p.src[0] if (is_binary_patch:=p.op is Ops.END) else p
  if not jit: return store.buf_uop.tag == "program"
  return is_binary_patch or (store.op is Ops.STORE and is_value_known_at_link(store.src[1]))

@@ -269,8 +257,7 @@ def trim_link_patches(ctx:tuple[bool, list[UOp]], a:UOp) -> UOp|None:
 pm_trim_link_patches = PatternMatcher([(UPat(Ops.AFTER, src=(UPat((Ops.PARAM, Ops.MSTACK)),), allow_any_len=True, name="a"), trim_link_patches)])

 def split_patches(ctx:bool, call:UOp) -> UOp|None:
-  lt_patches:list[UOp] = []
-  body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(ctx, lt_patches), name=f"trim link-time patches ({call.arg.aux.name})")
+  body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(ctx, lt_patches:=[]), name=f"trim link-time patches ({call.arg.aux.name})")

  lt_srcs = collections.defaultdict(list)
  for p in lt_patches: lt_srcs[p.buf_uop].append(p)
@@ -286,8 +273,8 @@ def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[dict[UOp, UOp
  slots = {g:i for i,g in enumerate(order)}
  table = UOp.placeholder((len(order),), dtypes.uint64, next(UOp.unique_num), device=call.arg.aux.device).rtag(name)

-  reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(slots[bare[g]], dtypes.int)).load() for g in gaddrs}
-  return reads, (table.after(make_patches(table, [(i * table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots else ()
+  reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, slots[bare[g]])).load() for g in gaddrs}
+  return reads, (table.after(*[make_patch(table, i * table.dtype.itemsize, addr) for addr, i in slots.items()]),) if slots else ()

 def make_blob_bufs(call:UOp, blobs:list[UOp]) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]:
  bufs = {b: UOp.placeholder((b.max_numel(),), b.dtype, next(UOp.unique_num), device=call.arg.aux.device).rtag("template") for b in blobs}
@@ -305,8 +292,7 @@ def rm_rt_uops(call:UOp) -> UOp|None:
  reads, fills = reads | {k:v for r,_ in tables for k,v in r.items()}, [f for _,fs in tables for f in fs]
  return call.replace(src=(call.src[0].substitute(reads), *call.src[1:], *fills),
                      arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(sorted(dedup(g.buf_uop.arg.slot for g in inputs))))))
-pm_rm_rt_uops = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_uops)])
+pm_rm_rt_uops = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_uops)])

 # *****************

@@ -318,25 +304,18 @@ def replace_params(call:UOp) -> UOp|None:
  by_root = {p.src[0]: p for p in patched}
  c_args = [by_root.get(a, a) for a in args]

-  # keep buffers whose addresses become link-time constants alive and mapped
-  held = args + [r.without_after for r in refhold]
-  addrs = dedup([g.src[0].without_after for x in call.src for g in x.toposort() if g.op is Ops.GETADDR])
-  refhold += [a for a in addrs if a not in held and all(b.op is not Ops.PARAM or b.tag is not None for b in unwrap_mstack(a))]
-
-  sub = {(b:=u.without_after): UOp.param(i, u.dtype, shape=b.shape, device=HCQ_RUNTIME_DEV.value, volatile=b.op is Ops.PARAM and b.arg.volatile)
-         for i,u in enumerate(c_args)} | {v: v.replace(arg=replace(v.arg, slot=-1)) for v in variables if v.op is Ops.PARAM}
+  sub = {(b:=u.without_after): UOp.param(i, u.dtype, shape=b.shape, device=u.device) for i,u in enumerate(c_args)} | \
+        {v: v.replace(arg=replace(v.arg, slot=-1)) for v in variables if v.op is Ops.PARAM}
  info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args) if u.tag == "inputs"), None))
-  return call.replace(src=(body.substitute(sub).replace(arg="hcq_args"), *c_args, *refhold),
-                      arg=replace(call.arg, aux=info)) # TODO: call.after(*refhold)?
-pm_replace_params = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), replace_params)])
+  return call.replace(src=(body.substitute(sub), *c_args, *refhold), arg=replace(call.arg, aux=info)) # TODO: call.after(*refhold)?
+pm_replace_params = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), replace_params)])

 # *****************

 def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp:
  base = bv.src[0].after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ())
  itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
-  return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].val * itemsize, dtypes.uint64)
+  return UOp(Ops.GETADDR, dtypes.uint64, src=(base,), arg=g.arg) + UOp.const(dtypes.uint64, bv.src[1].arg * itemsize)

 pm_early_simplify = PatternMatcher([
  (UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat(Ops.SLICE, name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice),
@@ -349,8 +328,7 @@ pm_early_simplify = PatternMatcher([

 def pack_hcq_placeholders(call:UOp) -> UOp|None:
  bufs = [b for b in call.src[0].toposort() if b.op is Ops.PARAM and b.tag in {"scratch", "kernargs"}]
-  offs:dict[UOp, int] = {}
-  sizes:dict[Any, int] = {}
+  offs, sizes = {}, {}
  for b in bufs:
    if b.tag == "scratch": sizes[b.tag] = max(sizes.get(b.tag, 0), b.max_numel())
    else:
@@ -358,44 +336,43 @@ def pack_hcq_placeholders(call:UOp) -> UOp|None:
      sizes[b.tag] = offs[b] + b.max_numel()
  counts = collections.Counter(b.tag for b in bufs)
  bases = {b.tag:UOp.placeholder((sizes[b.tag],), b.dtype, next(UOp.unique_num), device=b.device).rtag(b.tag) for b in bufs if counts[b.tag] > 1}
-  subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
+  subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.weakint, offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
  return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None
-pm_pack_placeholders = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])
+pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])

 # *****************
 # 8. callify hcq programs

-def callify_hcq(call:UOp, cf:UOp) -> UOp:
-  prg = to_program(cf.src[0].replace(arg=KernelInfo("hcq_submit"), tag=1), Device[HCQ_RUNTIME_DEV.value].renderer)
-  return call.replace(src=(cf.replace(src=(prg,), arg="hcq"), *call.src[1:]))
-pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, src=(
-  UPat(Ops.CUSTOM_FUNCTION, arg="hcq_args", src=(UPat(Ops.SINK),), name="cf"),), name="call", allow_any_len=True), callify_hcq)])
+pm_callify_hcq = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="hcq", src=(UPat(Ops.SINK),), name="cf"),
+  lambda cf: cf.replace(src=(to_program(cf.src[0].replace(arg=KernelInfo("hcq_submit"), tag=1), Device[HCQ_RUNTIME_DEV.value].renderer),)))])

 hcq_compile_cache:dict[tuple[bytes, bool], UOp] = {}

@track_rewrites(lambda linear,input_uops,jit,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
 def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp:
-  if input_uops is not None:
-    slots = {u:i for i,u in reversed(tuple(enumerate(input_uops)))}
-    linear = graph_rewrite(linear, pm_replace_buffers, ctx=(input_uops, slots), walk=True, name="replace buffer")
+  if input_uops is not None: linear = graph_rewrite(linear, pm_replace_buffers, ctx=input_uops, walk=True, enter_calls=True, name="replace buffer")

  if (final_linear:=(hcq_compile_cache.get(cache_key:=(linear.key, jit)))) is None:
    # prep
    linear = linear.substitute(back_map:={s.param_like(i): s for i,s in enumerate(input_uops)} if input_uops is not None else {}, walk=True)
-    linear = graph_rewrite(linear, pm_insert_copy_staging+pm_flatten_linear, name="insert copy staging")
+    linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging")

    # schedule
-    linear = graph_rewrite(linear, pm_schedule_and_merge, ctx={s:p for p,s in back_map.items()}, walk=True, name="schedule and merge hcq")
+    linear = graph_rewrite(linear, pm_sched_hcq_batches, walk=True, name="schedule hcq batches")
+    linear = linear.substitute({s: p for p, s in back_map.items()}, walk=True, enter_calls=True)
+    linear = graph_rewrite(linear, pm_merge_queues, walk=True, name="merge queues")

    # lowering to hcq ir
-    linear = graph_rewrite(linear, pm_encode_cmdbufs+pm_pack_placeholders, walk=True, name="encode and pack", enter_calls=True)
+    linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs", enter_calls=True)
+    linear = graph_rewrite(linear, pm_pack_placeholders, walk=True, name="pack placeholders")

-    # patches
-    linear = graph_rewrite(linear, pm_split_patches+pm_early_simplify+symbolic, ctx=jit, bottom_up=False, name="simplify patches", enter_calls=True)
+    # pie
+    linear = graph_rewrite(linear, pm_split_patches, ctx=jit, walk=True, name="split rt/lt patches")
+    linear = graph_rewrite(linear, pm_early_simplify + symbolic, bottom_up=False, name="simplify packed placeholders", enter_calls=True)
+    linear = graph_rewrite(linear, pm_rm_rt_uops, walk=True, name="replace rt uops")
+    linear = graph_rewrite(linear, pm_replace_params, walk=True, name="replace with args")

    # and compile it
-    linear = graph_rewrite(linear, pm_replace_params, bpm=pm_rm_rt_uops, name="replace rt uops and params")
    final_linear = hcq_compile_cache[cache_key] = graph_rewrite(linear, pm_callify_hcq, name="callify hcq", enter_calls=True)

  return final_linear
@@ -412,8 +389,8 @@ pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, name="buf"), bufferize_buf)])
 # *****************
 # 7. resolve patches

-def push_stack(op, s): return UOp(Ops.STACK,
-  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 push_stack(op, s): return UOp(Ops.STACK, op.dtype.scalar(),
+  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_binary(buf:UOp, blob:UOp) -> UOp:
  for b in (m.bufs if isinstance(m:=buf.buffer, MultiBuffer) else (m,)):
@@ -421,10 +398,9 @@ def fold_binary(buf:UOp, blob:UOp) -> UOp:
  return UOp(Ops.NOOP)

 def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
-  for off,val in zip(off.src, val.src):
-    for b,v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)):
-      data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype](v.val))
-      b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[(bo:=off.val*buf.dtype.itemsize):bo+len(data)] = data
+  for b, v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)):
+    data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype](v.arg))
+    b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[(byte_off:=off.arg*buf.dtype.itemsize):byte_off+len(data)] = data
  return UOp(Ops.NOOP)

 def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
@@ -433,8 +409,8 @@ def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
  devs, b = g.arg, buf.buffer
  bufs = tuple(cast(Buffer, x.buffer) for x in buf.src) if buf.op is Ops.MSTACK else tuple(b.bufs if isinstance(b, MultiBuffer) else (b,)*len(devs))
  assert len(bufs) == len(devs), f"can't resolve {len(bufs)} buffers on {len(devs)} devices"
-  addrs = tuple(UOp.const(x.get_buf(d).va_addr, dtypes.uint64) for x, d in zip(bufs, devs))
-  return addrs[0] if len(addrs) == 1 else UOp(Ops.STACK, src=addrs)
+  addrs = tuple(UOp.const(dtypes.uint64, x.get_buf(d).va_addr) for x, d in zip(bufs, devs))
+  return addrs[0] if len(addrs) == 1 else UOp(Ops.STACK, dtypes.uint64, addrs)

 pm_resolve_patches = PatternMatcher([
  # multi
@@ -448,26 +424,27 @@ pm_resolve_patches = PatternMatcher([
  (UPat(name="buf").index(UPat(Ops.RANGE), allow_any_len=True)
    .store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast()).index(UPat(Ops.RANGE), allow_any_len=True).load())
    .end(UPat(Ops.RANGE)), fold_binary),
-  (UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
+  (UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").index(UPat.cvar("off"))
+    .store(UPat.any(UPat.cvar("val"), UPat(Ops.STACK, name="val"))), fold_const_store),
 ])

 pm_assert_no_afters = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: panic(RuntimeError, f"AFTER left at hcq_link: {a.src[0].op}"))])

-def link_buf_key(a:UOp): return a.key, to_tuple(a.device)
-link_buf_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {}
-link_linear_cache:dict[tuple[bytes, bool], UOp] = {}
+hcq_link_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {}

-@track_rewrites(lambda _,jit,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
-def hcq_link(linear:UOp, jit=False, cache=True) -> UOp:
-  if (linked:=link_linear_cache.get(linear_key:=(linear.key, jit))) is not None: return linked
+def link_cache_key(a:UOp): return a.key, to_tuple(a.device)
+pm_link_cache = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: hcq_link_cache.get(link_cache_key(a)))])

-  bufs = {(j,i):a for j,c in enumerate(linear.src) for i,a in enumerate(c.src[1:], 1)
-          if a.op is Ops.AFTER and unwrap_mstack(a.src[0])[0].tag in HCQ_CACHE_TAGS}
-  linear = linear.substitute({x:link_buf_cache[k] for a in bufs.values() if (k:=link_buf_key(a)) in link_buf_cache for x in (a, a.src[0])}, walk=True)
-  linear = graph_rewrite(linear, pm_resolve_patches+symbolic+pm_assert_no_afters, bpm=pm_bufferize, ctx=cache, bottom_up=False,
-                         name="resolve patches")
-  for (j,i),a in bufs.items(): link_buf_cache.setdefault(link_buf_key(a), linear.src[j].src[i])
-  if cache: link_linear_cache[linear_key] = linear
+@track_rewrites(lambda _,jit,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
+def hcq_link(linear:UOp, jit=False) -> UOp:
+  cacheable = {(j,i):a for j,c in enumerate(linear.src) for i,a in enumerate(c.src[1:], 1)
+               if a.op is Ops.AFTER and unwrap_mstack(a.src[0])[0].tag in HCQ_CACHE_TAGS}
+  hits = {a.src[0]:hcq_link_cache[key] for a in cacheable.values() if (key:=link_cache_key(a)) in hcq_link_cache}
+  linear = graph_rewrite(linear, pm_link_cache, name="apply link cache").substitute(hits, walk=True)
+  linear = graph_rewrite(linear, pm_bufferize, ctx=jit, bottom_up=True, walk=True, name="bufferize placeholders")
+  linear = graph_rewrite(linear, pm_resolve_patches + symbolic, bottom_up=False, name="simplify patches")
+  linear = graph_rewrite(linear, pm_assert_no_afters, name="assert no afters")
+  for (j,i),a in cacheable.items(): hcq_link_cache.setdefault(link_cache_key(a), linear.src[j].src[i])
  return linear

 # *****************
@@ -481,41 +458,39 @@ class HCQ2Compiled(Compiled):

    self.pm_bufferize = PatternMatcher([
      (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].timeline_signal("sentinel", (1 << 64) - 1)),
-      (UPat(Ops.PARAM, name="b"), lambda ctx, b: None if b.tag is None else ctx[0].new_buffer(b, cache=ctx[1]))
+      (UPat(Ops.PARAM, name="b"), lambda ctx, b: None if b.tag is None else ctx[0].new_buffer(b, jit=ctx[1]))
    ])

-    super().__init__(device, allocator, compilers, runtime, None, arch=arch)
+    super().__init__(device, allocator, compilers, lambda *a, **kw: None, None, arch=arch)

    self.rt_buffer = Buffer(self.device, 64 << 20, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True))
-    self.rt_allocator = BumpAllocator(64 << 20)
+    self.rt_allocator = BumpAllocator(64 << 20, wrap=False)

-  def new_buffer(self, b:UOp, cache:bool) -> Buffer:
-    if cache or b.tag in HCQ_CACHE_TAGS:
+  def new_buffer(self, b:UOp, jit:bool) -> Buffer:
+    if jit or b.tag in HCQ_CACHE_TAGS:
      return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(uncached=True, cpu_access=True, nolru=True))
    return self.rt_buffer.view(b.max_numel(), b.dtype, self.rt_allocator.alloc(b.max_numel() * b.dtype.itemsize, alignment=128))

  @functools.cache
-  def timeline_signal(self, queue:str, init_value:int=0) -> Buffer:
+  def timeline_signal(self, queue:str="COMPUTE:0", init_value:int=0) -> Buffer:
    buf = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
    buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
    return buf

  @functools.cache
-  def timeline_value(self, queue:str, init_value:int=1) -> Buffer:
+  def timeline_value(self, queue:str="COMPUTE:0", init_value:int=1) -> Buffer:
    buf = Buffer("CPU", 1, dtypes.uint64, preallocate=True)
    buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
    return buf

  def synchronize(self, timeout:int|None=None):
    if not hasattr(self, 'iface'): return
-    sig = self.timeline_signal("COMPUTE:0").as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
-    tl = self.timeline_value("COMPUTE:0").as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
+    sig = self.timeline_signal().as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
+    tl = self.timeline_value().as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
    st = time.perf_counter()
    while sig[0] < tl[0] - 1:
      if time.perf_counter() - st > (timeout or 3000) / 1000: self.on_device_hang()

-  def on_device_hang(self): raise RuntimeError(f"{self.device} hang detected")
-
  def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.

  def count(self) -> int: return self.iface.count if hasattr(self, 'iface') else 1
@@ -3,8 +3,8 @@ from typing import cast, Any, Callable
 import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
-from tinygrad.runtime.support.hcq2 import make_binary_patch, make_patches
+from extra.hcq2.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
+from extra.hcq2.hcq2 import make_binary_patch
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
 from tinygrad.dtype import dtypes
@@ -37,7 +37,7 @@ class PM4Ops(FastEnum):
  RELEASE_MEM = auto(); DISPATCH_DIRECT = auto(); EVENT_WRITE = auto()  # noqa: E702

 def pkt3(ctx, op:PM4Ops, *vals):
-  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(x, dtypes.uint32)
+  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(dtypes.uint32, x)
    for x in (ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), len(vals) - 1), *vals)))

 def wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int):
@@ -146,38 +146,31 @@ pm_pm4_opsel = PatternMatcher([
  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
 ])

-def pm4_submit(ctx, lin):
-  # ensure compute queues are allocated
-  for d in (devs:=ctx.devs): q = Device[d].compute_queue
+def pm4_submit(cmdbuf, devs):
+  size, zero = UOp.const(dtypes.uint32, cmdbuf.nbytes() // dtypes.uint32.itemsize), UOp.const(dtypes.int, 0)
+
+  # the compute queue's ring and its host-side ring/write/put pointers (placeholders, resolved in pm_bufferize)
+  for d in devs: q = Device[d].compute_queue
  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}")
    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))

-  # two tail dwords coordinate safe IB reuse: GPU completions and host submits
-  size_dw = sum(len(ins.src) for ins in lin.src) + len(release_mem(ctx, 0, 0).src)
-  assert size_dw < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet"
+  # place the cmdbuf at the ring's write offset, wrapping the ring
+  put = put_ptr.index(zero)
+  next_put = put + size.cast(put.dtype)
+  i = UOp.range(size, 0, dtype=dtypes.int, src=(cmdbuf,))
+  ring_idx = ((put + i.cast(put.dtype)) % q.ring.size).cast(dtypes.int)

-  ib = UOp.placeholder((size_dw + 2,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
-  done_idx, submit_idx = UOp.const(size_dw + 0, dtypes.int), UOp.const(size_dw + 1, dtypes.int)
-  init_counters = make_patches(ib, [((size_dw + i) * 4, UOp.const(0, dtypes.uint32)) for i in range(2)]).rtag("link")
-  submitted = (counter:=ib.after(init_counters).index(submit_idx)).load()
-  completed = ib.after(loop:=UOp.loop(0)).index(done_idx).load()
-  ib_free = completed.end(loop, completed != submitted)
+  # copy the cmdbuf into the ring and advance the put/write pointers
+  copy_to_ring = ring.index(ring_idx).store(cmdbuf.index(i).load()).end(i)
+  bump_put_ptr = put_ptr.index(zero).store(next_put)
+  bump_wptr = wptr.index(zero).store(next_put)

-  bump_fence = pm4_store(ctx, UOp(Ops.SLICE, dtypes.uint32, (ib, UOp.const(size_dw)), 2), (submitted + 1).cast(dtypes.uint64))
-  cmdbuf = make_cmdbuf(lin.replace(src=lin.src + (bump_fence,)), devs, buf=ib, dep=ib_free)
+  # ring the doorbell once the copy and pointer bumps have landed
+  flush = UOp.barrier(copy_to_ring, bump_put_ptr, bump_wptr)
+  return doorbell.after(flush).index(zero).store(next_put)

-  # the ring itself only carries a packet pointing at the ib, wrapping the ring
-  put = put_ptr.index(zero:=UOp.const(0, dtypes.int))
-  pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER, 2), *data64_le(cmdbuf.getaddr(devs)), size_dw | ctx.pm4.INDIRECT_BUFFER_VALID)
-  write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(x, dtypes.uint32)) for off,x in enumerate(pkt)])
-
-  # advance the put/write pointers past the packet
-  bump_put_ptr = put_ptr.index(zero).store(put + len(pkt))
-  bump_wptr = wptr.index(zero).store(put + len(pkt))
-  flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr, counter.store(submitted + 1))
-  return doorbell.after(flush).index(zero).store(put + len(pkt))
-
-pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), pm4_submit)])
+pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
+  lambda lin: pm4_submit(make_cmdbuf(lin, to_tuple(lin.arg[0])), to_tuple(lin.arg[0])))])

 # *****************
 # SDMA
@@ -187,26 +180,26 @@ class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TR
 def sdma_copy(ctx, call):
  sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
  src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs)
-  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(x, dtypes.uint32) for off in range(0, sz, ctx.max_copy_size) for x in (
+  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(dtypes.uint32, x) for off in range(0, sz, ctx.max_copy_size) for x in (
    ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
    ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0, *data64_le(src_addr+off), *data64_le(dst_addr+off))))

 def sdma_wait(ctx, ins, dst, val):
  op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
     | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
-  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(x, dtypes.uint32) for x in (
+  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(dtypes.uint32, x) for x in (
    op, *data64_le(dst.getaddr(ctx.devs)), val, 0xffffffff,
    ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))))

 def sdma_store(ctx, ins, dst, val):
  op = ctx.sdma.SDMA_OP_FENCE | (ctx.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if ctx.target[0] != 9 else 0)
  return UOp(Ops.LINEAR, src=(
-    ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
-    ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(x, dtypes.uint32) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))
+    ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
+    ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(dtypes.uint32, x) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))

 def sdma_timestamp(ctx, ins, dst):
  op = ctx.sdma.SDMA_OP_TIMESTAMP | ctx.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL)
-  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))
+  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))

 pm_sdma_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
@@ -219,7 +212,7 @@ pm_sdma_opsel = PatternMatcher([

 def sdma_submit(cmdbuf, devs):
  # the cmdbuf to submit + the patch writes that fill it
-  size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(0, dtypes.int)
+  size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(dtypes.int, 0)

  # the sdma queue's ring and its host-side ring/write/put pointers
  for d in devs: q = Device[d].sdma_queue(0)
@@ -235,8 +228,8 @@ def sdma_submit(cmdbuf, devs):

  # zero the wrapped tail, then copy the cmdbuf into the ring
  zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,))
-  zero_tail = ring.index(tail_off_dw + zi).store(UOp.const(0, dtypes.uint32)).end(zi)
-  i = UOp.range(UOp.const(size_dw, dtypes.int), 0, dtype=dtypes.int, src=(cmdbuf,))
+  zero_tail = ring.index(tail_off_dw + zi).store(UOp.const(dtypes.uint32, 0)).end(zi)
+  i = UOp.range(UOp.const(dtypes.int, size_dw), 0, dtype=dtypes.int, src=(cmdbuf,))
  copy_to_ring = ring.index(start_dw + i).store(cmdbuf.index(i).load()).end(i)

  # advance the put/write pointers past the zeroed tail and the cmdbuf
@@ -249,7 +242,7 @@ def sdma_submit(cmdbuf, devs):
  return doorbell.after(flush).index(zero).store(next_put_b)

 pm_sdma_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
-  lambda ctx, lin: sdma_submit(make_cmdbuf(lin, ctx.devs), ctx.devs))])
+  lambda lin: sdma_submit(make_cmdbuf(lin, to_tuple(lin.arg[0])), to_tuple(lin.arg[0])))])

@dataclass(frozen=True)
 class AMDEncodeCtx:  # encode-time constants for one queue: devs (every cmdbuf address resolves into these) + gfx version + packet/ip modules
@@ -260,7 +253,7 @@ def encode_queue(q:UOp) -> UOp|None:
  d = Device[(devs:=to_tuple(q.arg[0]))[0]]
  ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size)
  opsel, submit = (pm_pm4_opsel, pm_pm4_submit) if q.arg[1].startswith("COMPUTE") else (pm_sdma_opsel, pm_sdma_submit)
-  return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"), ctx)
+  return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"))

@dataclass(frozen=True)
 class AMDProgramData:
@@ -518,8 +511,7 @@ class PCIIface(PCIIfaceBase):
      cq = d.compute_queue
      for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
      d.iface.dev_impl.gfx.setup_ring(*cq.params)
-      d.timeline_signal('COMPUTE:0')._buf.cpu_view().mv.cast('Q')[0] = \
-        d.timeline_value('COMPUTE:0').as_memoryview(force_zero_copy=True).cast('Q')[0] - 1
+      d.timeline_signal()._buf.cpu_view().mv.cast('Q')[0] = d.timeline_value().as_memoryview(force_zero_copy=True).cast('Q')[0] - 1

  def sleep(self, timeout):
    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -31,12 +31,12 @@ def dname_of(device) -> str:

 def alloc_like(shape, dtype, device, axis=None) -> Tensor:
  if isinstance(device, tuple) and axis is not None:
-    return Tensor(Tensor.invalids(*shard_shape(shape, axis, len(device)), dtype=dtype, device=device).uop.unshard(axis), device=device)
+    return Tensor(Tensor.invalids(*shard_shape(shape, axis, len(device)), dtype=dtype, device=device).uop.multi(axis), device=device)
  return Tensor.invalids(*shape, dtype=dtype, device=device)

 def alloc_local(shape, dtype, device, axis=None) -> Tensor:
  if isinstance(device, tuple) and axis is not None:
-    return Tensor(Tensor.invalids(*shape, dtype=dtype, device=device).uop.unshard(0), device=device)
+    return Tensor(Tensor.invalids(*shape, dtype=dtype, device=device).uop.multi(0), device=device)
  return Tensor.invalids(*shape, dtype=dtype, device=device)

 def compile_hip(src:str, defines:list[str]):
@@ -48,7 +48,7 @@ def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float):
    axis = logits_u.axis
    ndev = len(device)
    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.unshard(axis), device=device)
+    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:
@@ -74,11 +74,11 @@ def fused_ce_loss(logits:Tensor, targets:Tensor, label_smoothing:float=0.1) -> T
    axis = logits.uop.axis
    assert axis in (0, 1), f"unsupported sharding axis={axis} for CE loss"
    ndev = len(logits.device)
-    loss_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.unshard(0),
+    loss_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
                      device=logits.device)
-    max_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.unshard(0),
+    max_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
                      device=logits.device)
-    lse_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.unshard(0),
+    lse_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
                      device=logits.device)
    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]
@@ -16,7 +16,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:

  wg = UOp.range(NUM_WG, 0, AxisType.GLOBAL)
  tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
-  it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.WEAK)
+  it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.LOOP)
  lane = UOp.range(VEC, 3, AxisType.UNROLL)

  idx = (((it * NUM_WG + wg) * THREADS_PER_WG + tid) * VEC) + lane
@@ -36,19 +36,19 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
  lmax_val = lmax.after(lmax_store.end(it))[0]

  lds = UOp.placeholder((THREADS_PER_WG,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-  lds = lds.after(lds[tid].store(lmax_val))
+  lds = lds.after(lds[tid].store(lmax_val).barrier())

  step = THREADS_PER_WG // 2
  while step:
    active = tid < step
    other = lds[(tid + step).valid(active)].load()
-    lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)))
+    lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)).barrier())
    step //= 2

  device = device[0].split(":")[0] if isinstance(device, tuple) else device.split(":")[0]
  if device in {"AMD", "NULL"}: atomic_arg = "if ({2} > {3}) __hip_atomic_fetch_max((int*){0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);"
  else: raise NotImplementedError(f"no atomic max for device {device}")
-  amax_idx = amax_out.reshape((1,)).index(UOp.const(0))
+  amax_idx = amax_out.reshape((1,)).index(UOp.const(dtypes.weakint, 0))
  max_val = lds[0].load()
  atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
  return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))
@@ -4,7 +4,7 @@ from hexdump import hexdump
 from copy import deepcopy
 import pathlib, sys
 from tinygrad.helpers import to_mv, getenv
-from tinygrad.runtime.autogen import mesa
+from tinygrad.runtime.autogen import adreno
 sys.path.append(pathlib.Path(__file__).parent.parent.parent.as_posix())

 IOCTL = getenv("IOCTL", 0)
@@ -23,7 +23,7 @@ for child in xml.getroot():
 CAPTURED_STATE = {}

 REGS = {}
-for k, v in mesa.__dict__.items():
+for k, v in adreno.__dict__.items():
  if k.startswith("REG_") and isinstance(v, int) and v > 1024: REGS[v] = k

 from extra.qcom_gpu_driver import msm_kgsl
@@ -42,7 +42,7 @@ def get_struct(argp, stype):

 def format_struct(s):
  sdats = []
-  for field_name, *_ in s._fields_:
+  for field_name, *_ in s._real_fields_:
    if field_name in {"__pad", "PADDING_0"}: continue
    dat = getattr(s, field_name)
    if isinstance(dat, int): sdats.append(f"{field_name}:0x{dat:X}")
@@ -96,9 +96,9 @@ def parse_cmd_buf(dat):
        CAPTURED_STATE['LOAD_FRAGS'].append((state_block, state_type, num_unit, dst_off))

        if state_block == SB6_CS_SHADER:
-          from tinygrad.runtime.support.compiler_mesa import disas_adreno
+          from extra.disassemblers.adreno import disasm_raw
          if state_type == ST6_SHADER and IOCTL > 3:
-            disas_adreno(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
+            disasm_raw(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
          if state_type == ST6_CONSTANTS:
            x = get_mem(((vals[2] << 32) | vals[1]), num_unit*4)
            CAPTURED_STATE['constants'] = x[:]
@@ -142,7 +142,7 @@ def parse_cmd_buf(dat):
      vals = struct.unpack("I"*size, dat[ptr+4:ptr+4+4*size])
      if IOCTL > 0: print(f"{ptr:3X} -- typ 4: {size=:3d}, {reg_name}", hprint(vals))
      for vi,v in enumerate(vals): CAPTURED_STATE[offset+vi] = v
-      if offset == mesa.REG_A6XX_SP_CS_CONFIG:
+      if offset == adreno.REG_A6XX_SP_CS_CONFIG:
        val = vals[0]
        if IOCTL > 0:
          print(f"\tBINDLESS_TEX={(val >> 0) & 0b1}")
@@ -215,3 +215,79 @@ def install_hook(c_function, python_function):

 libc = ctypes.CDLL(ctypes.util.find_library("libc"))
 install_hook(libc.ioctl, ioctl)
+
+def before_launch():
+  global CAPTURED_STATE
+  CAPTURED_STATE.clear()
+def collect_last_launch_state():
+  global CAPTURED_STATE
+  return deepcopy(CAPTURED_STATE)
+def compare_launch_state(state, good_state):
+  cmp = [
+    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NTEX__MASK),
+    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NSAMP__MASK),
+    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NIBO__MASK),
+    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_ENABLED),
+    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_TEX),
+    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_SAMP),
+    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_IBO),
+    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_UBO),
+
+    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_HALFREGFOOTPRINT__MASK),
+    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_FULLREGFOOTPRINT__MASK),
+    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_BRANCHSTACK__MASK),
+    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_FULLREGFOOTPRINT__MASK),
+    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_THREADMODE__MASK),
+    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_EARLYPREAMBLE),
+    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_MERGEDREGS),
+
+    (adreno.REG_A6XX_SP_CS_PVT_MEM_PARAM, adreno.A6XX_SP_CS_PVT_MEM_PARAM_MEMSIZEPERITEM__MASK),
+    (adreno.REG_A6XX_SP_CS_PVT_MEM_PARAM, adreno.A6XX_SP_CS_PVT_MEM_PARAM_HWSTACKSIZEPERTHREAD__MASK),
+
+    (adreno.REG_A6XX_SP_CS_UNKNOWN_A9B1, adreno.A6XX_SP_CS_UNKNOWN_A9B1_UNK5),
+    (adreno.REG_A6XX_SP_CS_UNKNOWN_A9B1, adreno.A6XX_SP_CS_UNKNOWN_A9B1_UNK6),
+
+    (adreno.REG_A6XX_SP_CS_BRANCH_COND, 0xffffffff),
+
+    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_KERNELDIM__MASK),
+    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEX__MASK),
+    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEY__MASK),
+    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEZ__MASK),
+
+    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_1, 0xffffffff),
+    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_2, 0xffffffff),
+    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_3, 0xffffffff),
+    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_4, 0xffffffff),
+    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_5, 0xffffffff),
+    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_6, 0xffffffff),
+
+    (adreno.REG_A6XX_HLSQ_CS_CNTL_0, 0xffffffff),
+    (adreno.REG_A6XX_HLSQ_CS_CNTL_1, 0xffffffff),
+    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_X, 0xffffffff),
+    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_Y, 0xffffffff),
+    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_Z, 0xffffffff),
+  ]
+
+  for x,m in cmp:
+    print(f"Field {REGS[x]}, mask: 0x{m:X} cmp: {state.get(x, 0) & m} vs {good_state.get(x, 0) & m}")
+    if state.get(x, 0) & m != good_state.get(x, 0) & m:
+      return False, f"Field {REGS[x]}, mask: 0x{m:X} mismatch: {state.get(x, 0) & m} vs {good_state.get(x, 0) & m}"
+
+  for n in ['descriptors', 'ibos']:
+    if n not in good_state: continue
+    mv1, mv2 = state.get(n), good_state.get(n)
+
+    if len(mv1) != len(mv2): return False, f"{n}: len mismatch {len(mv1)} != {len(mv2)}"
+    mv1 = memoryview(bytearray(mv1)).cast('I')
+    mv2 = memoryview(bytearray(mv2)).cast('I')
+    for i in range(len(mv2)):
+      if i % 8 == 5 or i % 8 == 4: continue # addresses
+      if mv1[i]!=mv2[i]: return False, f"{n}: content mismatch {i} {mv1[i]} {mv2[i]}"
+
+  for n in ['samplers']:
+    if n not in good_state: continue
+    mv1, mv2 = state.get(n), good_state.get(n)
+    if len(mv1) != len(mv2): return False, f"{n}: len mismatch {len(mv1)} != {len(mv2)}"
+    if any(mv1[i]!=mv2[i] for i in range(len(mv1))): return False, f"{n}: content mismatch"
+
+  return True, "PASS"
@@ -14,7 +14,7 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
  shard_axis = ref.uop.axis if axis is None else axis
  shape = tuple(s // len(ref.device) if i == shard_axis else s for i, s in enumerate(shape))
  axis = ref.uop.axis if axis is None else axis
-  return Tensor(Tensor.invalids(*shape, dtype=dtype, device=ref.device).uop.unshard(axis), dtype=dtype, device=ref.device)
+  return Tensor(Tensor.invalids(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)

@functools.cache
 def custom_fused_qkv_rope_forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs_cis:UOp,
@@ -16,7 +16,7 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None) -> Tensor:
  if not isinstance(ref.device, tuple): return Tensor.empty(*shape, dtype=ref.dtype, device=ref.device)
  shape = tuple(s // len(ref.device) if i == ref.uop.axis else s for i, s in enumerate(shape))
  axis = ref.uop.axis if axis is None else axis
-  return Tensor(Tensor.empty(*shape, dtype=ref.dtype, device=ref.device).uop.unshard(axis), dtype=ref.dtype, device=ref.device)
+  return Tensor(Tensor.empty(*shape, dtype=ref.dtype, device=ref.device).uop.multi(axis), dtype=ref.dtype, device=ref.device)

 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
  return _sharded_empty(ref.shape, ref, axis)
@@ -224,7 +224,7 @@ class Group:

      # store to shared memory
      red_local_store = red_local[self.laneid].store(red_reg[0])
-      red_local = red_local.after(red_local_store).reshape(red_local.shape)
+      red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)

      # reduce from shared memory
      for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
@@ -258,7 +258,7 @@ class Group:

      # store to shared memory
      red_local_store = red_local[self.laneid].store(red_reg[0])
-      red_local = red_local.after(red_local_store).reshape(red_local.shape)
+      red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)

      # reduce from shared memory
      for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
@@ -342,7 +342,7 @@ class Group:
          if src.dtype != dst.dtype:
            src_load = src_load.cast(dst.dtype)
          dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load)
-          dst_store = dst_store.end(height, width, outer, inner)
+          dst_store = dst_store.end(height, width, outer, inner).barrier()
    elif dst.addrspace == AddrSpace.REG and src.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT):
      srcf = src.flatten()
      row_stride = prod(src.shape[axis+1:])
@@ -48,14 +48,14 @@ class Kernel(AbstractContextManager):
  @property
  def warpgroup(self): return self.group(4)

-  def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.WEAK, track:bool=True):
+  def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.LOOP, track:bool=True):
    if end == 0: start, end = 0, start
    rng = _tk_range(start, end, step, axis_type, self.range_id)
    self.range_id += 1
    if track: self.range_stack.append(rng)
    return rng

-  def raw_range(self, end:int=0, axis_type:AxisType=AxisType.WEAK):
+  def raw_range(self, end:int=0, axis_type:AxisType=AxisType.LOOP):
    rng = UOp.range(end, self.range_id, axis_type=axis_type)
    self.range_id += 1
    return rng
@@ -555,8 +555,6 @@ def wrap_out(f):
    assert out.shape == assigned.shape, f"shape mismatch: {assigned.shape} -> {out.shape}"
    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}"
-    # an out= that is a view has to be written through its base, and _apply_inplace gives a deviceless base its buffer first
-    if canonical_base(out) is not out: return _apply_inplace(out, assigned) or out
    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
@@ -98,21 +98,6 @@ class TestTorchBackend(unittest.TestCase):
  def test_empty_strided_default_dtype(self):
    self.assertEqual(torch.empty_strided((2,3), (1,2), device=device).dtype, torch.get_default_dtype())

-  @unittest.expectedFailure  # TODO: empty_strided ignores the requested strides, the backend treats everything as contiguous
-  def test_empty_strided_honors_strides(self):
-    self.assertEqual(tuple(torch.empty_strided((2,3), (1,2), device=device).stride()), (1,2))
-
-  @unittest.expectedFailure  # TODO: torch refuses an out= that overlaps an input, we compute silently
-  def test_out_overlapping_input_is_rejected(self):
-    x = torch.arange(6., device=device)
-    with self.assertRaises(RuntimeError): torch.add(x[:-1], 10, out=x[1:])
-
-  def test_out_disjoint_input_is_allowed(self):
-    # torch permits an out= that shares a base with an input as long as they do not overlap
-    x, xc = torch.arange(6., device=device), torch.arange(6.)
-    torch.add(x[:3], 10, out=x[3:]); torch.add(xc[:3], 10, out=xc[3:])
-    np.testing.assert_equal(x.cpu().numpy(), xc.numpy())
-
  def test_plus_inplace(self):
    a = torch.ones(4, device=device)
    b = torch.ones(4, device=device)
@@ -0,0 +1,86 @@
+# Multi-device op migration: MULTI/MSELECT/MSTACK → PAD / WHERE / STACK+INDEX
+
+## Status (updated)
+
+- **Stage 0 — DONE.** Internal `Ops.PAD` fills **Invalid** (`schedule/indexing.py:104`, bool keeps 0-fill); external `Tensor.pad`/`pad_to` always emit an explicit fill mask (`mixin/op.py:289`, `mixin/movement.py:267`) — required because a bare Invalid-pad leaks through elementwise ALU (`pad(x)+1` would read 0 instead of 1 in pad regions). REDUCE inputs with Invalid contribute the reduce identity (`pm_invalid_reduce_identity` in `uop/symbolic.py`, run in `get_kernel_graph` after gate lifting in `schedule/rangeify.py`) — only WHERE-alt gates whose condition involves a reduce range are rewritten, so gather-with-Invalid-index still poisons whole lanes. Same-condition nested where collapse rule added (`c?(c?t:f):f2 -> c?t:f2`) so the mask form folds to a single gate. All suites green (`test/unit`, `test/null`, `test/backend`, `test/external/external_test_schedule_scaling.py`, mypy, ruff).
+- **Stage 1 — representation in place behind `SYMBOLIC_MULTI`.** `symbolic_multi_pm` (`schedule/multi.py`) converts `MULTI→_unshard` (raw Invalid pad), `MSELECT→dnum.eq(i).where(x, Invalid)`, `MSTACK→STACK.index(dnum)` + INDEX(STACK,var)→nested-where lowering. `_unshard` uses the raw Invalid pad; `_unshard_fill` (0-fill) is used for the ALU allreduce in `copy_multi` because gated stores leave stale pad regions (the ALU-sum path can't use the identity rule). Basic shard ops work; full parity is Stage 2.
+- **Stage 2 — remaining.** Buffer level, reduce/allreduce split for shard-axis reduces, API surface.
+- **Stage 3 — remaining.**
+
+Notes: `test_schedule.py:test_pad_reduce_unsafe_multiview_st` went 4→5 kernels (pad now materializes an explicit mask; the mask form is also what makes the previously-wrong masked-pad+hazard case correct). `test_jit_footguns.py:test_symbolic_pad_view_frozen` went 2→4: the explicit mask recomputes from the symbolic shape, fixing the frozen-pad footgun. Also fixed a latent infinite loop: `(x+y) !=/< c → x !=/< c-y` collapse rules in `codegen/simplify.py` now only fire when the remaining side still contains the range (they previously shuffled constants forever when both sides were range-free).
+
+## Goal
+
+Replace the three multi-device UOps with a symbolic `_device_num` representation:
+
+| Old op | New form |
+|---|---|
+| `MULTI(x, axis)` | `x._unshard(axis)` — PAD with `_device_num`-dependent bounds back to full shape (helper already exists at `tinygrad/uop/ops.py:704-707`) |
+| `MSELECT(x, i)` | `dnum.eq(i).where(x, x.const_like(Invalid))` |
+| `MSTACK(s0..sn)` | `UOp(Ops.STACK, src=srcs).index(dnum)` — leading device axis, indexed per-device |
+
+where `dnum = UOp.variable("_device_num", 0, ndev-1)`. The per-device specialization mechanism already exists: `unwrap_multi` (`tinygrad/engine/realize.py:148-153`) binds `_device_num` per device at exec time.
+
+**Key semantic decision (approved):** internal `Ops.PAD` produces **Invalid** in padded regions; external `Tensor.pad` API still pads with 0. Staged migration: introduce the new representation first, keep old ops working, migrate call sites incrementally, delete old ops last.
+
+## Background: current design
+
+- `Ops.MULTI(src, axis)` (`tinygrad/uop/__init__.py:100`) — per-shard graph marker. Eliminated by `multi_pm` (`tinygrad/schedule/multi.py:162-195`) as the first step of `get_kernel_graph` (`tinygrad/schedule/rangeify.py:548`). Shape/axis tracking: `UOp.axis`/`UOp.bounds` (`tinygrad/uop/ops.py:667-702`).
+- `Ops.MSELECT(x, i)` / `Ops.MSTACK(srcs)` (`__init__.py:96`) — buffer-level ops. Spec at `tinygrad/uop/spec.py:181-184`; device prop `ops.py:816-819`; per-kernel PARAMs via debuf (`rangeify.py:474`); per-device dependency states (`tinygrad/schedule/__init__.py:11-17`); `MultiBuffer` (`ops.py:904-930`, `tinygrad/device.py:88-99`); only MSTACK can be `realized` (`ops.py:920-930`).
+- `_shard`/`_unshard` (`ops.py:704-714`) already emit symbolic SHRINK/PAD bounds with `_device_num`.
+- Naive allreduce already uses the target pattern: `dnum.eq(i).where(buf, state)` (`tinygrad/schedule/allreduce.py:27-33`).
+
+## Existing Invalid machinery (rely on this)
+
+- `pm_data_invalid` (`tinygrad/uop/symbolic.py:71-92`): Invalid poisons ALU (ops move inside the gate); gated LOAD folds to alt/0, gated STORE folds to NOOP.
+- `pm_remove_invalid` (`symbolic.py:94-96`): leftover Invalid → 0 in final codegen (`codegen/__init__.py:345`). Spec forbids Invalid in final programs (`spec.py:217`), so materialized Invalid regions read as 0.
+- STORE of CONST(Invalid) → NOOP (`rangeify.py:423-424`).
+- `identity_element(op, dtype)` exists (`ops.py:51`): ADD→0, MUL→1, MAX→dtype.min.
+- `found_after` (`rangeify.py:26`) already matches `WHERE(cond, PAD(x), Invalid)`.
+
+## Stage 0 — internal PAD = Invalid; external pad = explicit 0
+
+1. `tinygrad/schedule/indexing.py:100-104` (`convert_pad_to_where_to_keep_behavior_local`): fill value `0` → `UOp.const(x.dtype, Invalid)`, **except `dtypes.bool` keeps 0-fill** (False is the bool-reduce identity, and the external-pad mask below needs it).
+2. `tinygrad/mixin/op.py:282-290` (`_pad_constant`): **remove the `if value == 0: return base` shortcut** — always emit `pad(bool_ones).where(base, value)`. Required because bare Invalid-pad leaks through elementwise ALU: `pad(x)+1` gate-lifts to `where(valid, x+1, Invalid)` and reads 0 instead of 1 in pad regions. The mask lowers to a pure index expression (`valid.where(1,0)`), no extra kernel. External behavior unchanged for all `value`.
+3. **New rule**: `REDUCE(where(c, x, Invalid), op)` → `REDUCE(where(c, x, identity_element(op, dtype)), op)`. Must fire in rangeify/symbolic *before* codegen builds the accumulator loop — otherwise `pm_data_invalid` gate-lifts `acc + where(c,x,Invalid)` into `where(c, acc+x, Invalid)` and one invalid lane poisons the whole reduction. Placement (symbolic.py vs the reduce path in indexing.py) TBD at implementation; verify with `Tensor.pad(...).sum()/max()` tests.
+4. Audit: schedule tests with kernel counts involving pads; circular/reflect/replicate pads don't use PAD fill (verified, `op.py:292-312`) — unaffected; `allreduce.py:59,76` usum-of-padded-chunks gets *more* correct (disjoint regions).
+
+## Stage 1 — new representation behind env flag
+
+New `symbolic_multi_pm` PatternMatcher (in `schedule/multi.py` or new file), gated by env (e.g. `SYMBOLIC_MULTI`), run in `get_kernel_graph` right after `multi_pm`:
+
+- `MULTI(x, axis)` → `x._unshard(axis)`
+- `MSELECT(x, i)` → `dnum.eq(i).where(x, x.const_like(Invalid))` (Invalid from `tinygrad.dtype`)
+- `MSTACK(srcs)` → `STACK(*srcs).index(dnum)`, plus new lowering `INDEX(STACK(vals), var)` → nested `var.eq(k).where(src_k, Invalid)` (analogous to `convert_stack_to_where`, `indexing.py:113-121`; must fire before `validate_index` spec, `spec.py:118-122`)
+
+Flag off = zero behavior change; flag on = new forms flow through rangeify and specialize per device at exec.
+
+## Stage 2 — migrate producers/consumers (one commit each, independently testable)
+
+1. `UOp.shard` (`ops.py:715-717`): emit symbolic `_shard`+`_unshard` full-shape form directly instead of `.multi(axis)`; delete movement-op `multi_pm` rules that PAD subsumes (`pad_multi`, `permute_multi`, `expand_multi`, `reshape_multi`, `flip_multi`, `shrink_multi` — `multi.py:93-125`).
+2. ALU/STACK: `alu_multi`/`shard_srcs`/`stack_multi` (`multi.py:55-78,127-131`) become plain elementwise on full-shape padded tensors. `reduce_multi` (`multi.py:80-91`) keeps the shard-axis → local-reduce + ALLREDUCE split; Invalid-pad + identity rule replaces neutral-pad-value reasoning.
+3. allreduce (`schedule/allreduce.py`): naive path already matches; migrate ring/all2all MSELECT/MSTACK scratch-buffer assembly (lines 35-76) to WHERE/STACK+INDEX forms.
+4. Buffer level: debuf (`rangeify.py:474`), `_states`/`_unwrap_src` (`schedule/__init__.py:11-17`), `_collect_bufs` (`schedule/memory.py:9`), `unwrap_multi` (`realize.py:148-153`), JIT (`jit.py:127-130, 237`), callify (`callify.py:52-95`), `buffer`/`realized`/`buf_uop`/`has_buffer_identity` (`ops.py:841-930`).
+5. API surface: `UOp.multi/mselect/mstack` (`ops.py:662-725`), `Tensor.shard` (`tensor.py:333-347`), gradient (`mixin/gradient.py:72`), `_multi_like` (`mixin/creation.py:16-20`), embedding backward (`nn/__init__.py:309-354`), `copy_to_device(arg=)` MSELECT path (`ops.py:719-723`).
+
+## Stage 3 — removal
+
+Delete `Ops.MULTI/MSELECT/MSTACK` from the enum (`uop/__init__.py:96,100`), spec rules, viz colors (`viz/serve.py:51,56`), `UOp.axis`/`bounds` machinery (`ops.py:667-702`), remaining `multi_pm` rules, and `MultiBuffer` if fully subsumed. Flip flag default-on, then delete the flag.
+
+## Open implementation details
+
+- REDUCE-identity rule placement (must precede codegen accumulator construction).
+- INDEX(STACK, var) spec timing — the value-STACK INDEX violates the pointer-INDEX spec until lowered.
+- Whether `MultiBuffer`/tuple-`device` survives as the runtime container, or buffers become single-device with the device axis explicit in shape — decides how much of Stage 2.4 is rewrite vs delete.
+- Bool carve-out in Stage 0.1: verify no internal consumer needs Invalid-filled bool pads.
+
+## Verification (run at each stage)
+
+```bash
+python -m pytest test/unit/test_multitensor.py test/unit/test_allreduce.py test/null/test_multitensor.py test/unit/test_call.py -x -q -n12
+python -m pytest test/external/external_test_schedule_scaling.py -x -q   # test_concat_scaling
+python -m mypy tinygrad/
+python -m ruff check .
+```
+
+Also pad/reduce numeric tests after Stage 0 (`test_ops` pad tests, `Tensor.pad(...).sum()/max()`).
@@ -29,7 +29,6 @@
 \definecolor{axbrred}{HTML}{E53935}      % GROUP_REDUCE
 \definecolor{axyellow}{HTML}{F9A825}     % UPCAST
 \definecolor{axmagenta}{HTML}{7B1FA2}    % UNROLL
-\definecolor{axgreen}{HTML}{2E7D32}      % DEVICE

 \title{tinygrad: a single dialect from Tensor programs to Command Buffers}
 \author{tinygrad, Corp. \\ \texttt{[email protected]}}
@@ -80,7 +79,6 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \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{Bitcast} & $(T,)$ & dtype & Reinterpret storage as target dtype; preserve total bytes. \\
-\op{Unshard} & $(T, R_0, R_1, \ldots)$ & axes $(a_0, a_1, \ldots)$ & Concatenate shards of \op{Range} $R_k$ along axis $a_k$; $R_k$ is outer. \\
 \bottomrule
 \end{tabular}

@@ -260,7 +258,6 @@ Every UOp has a \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{addrspa
 \op{Const} & from arg & $()$ & \textsc{null} & $[v, v]$ \\
 \op{Param} & from arg & from $\mathrm{src}[0]$ & from arg & from src or dtype range \\[3pt]
 Movement ops & $\mathrm{src}[0].\mathrm{dtype}$ & (see op) & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
-\op{Unshard} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0]$, each $a_k \times n_k$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Reduce} & $\mathrm{src}[0].\mathrm{dtype}$ & remove first $n$ axes & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\[3pt]
 \op{Cast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & clamped to dtype \\
 \op{Bitcast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
@@ -286,9 +283,9 @@ $[a,A]$, $[b,B]$, $[c,C]$ denote min\_max of $\mathrm{src}[0]$, $\mathrm{src}[1]
 Default \emph{dtype range}: $[\mathrm{dtype\_min},\, \mathrm{dtype\_max}]$.

 \medskip
-\textbf{sharding} tracks multi-device sharding as a set of (axis, \op{Range}) pairs. \op{Unshard} defines it: arg is the tuple of sharded axes, one \op{Range} in src per axis (positional: the $k$-th \op{Range} shards the $k$-th axis). \op{Buffer} with $n$-tuple device: sharded on axis $0$ (device dim). The single-axis convenience \textbf{axis} is \textsc{null} unless exactly one axis is sharded.
-\op{Reshape} remaps each sharded axis to preserve its shard boundary. \op{Permute} follows the permutation. \op{Expand} shifts all sharded axes right by $|\mathbf{n}|$.
-\op{Reduce} on a sharded axis drops it. \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}.
+\textbf{axis} tracks the multi-device sharding dimension. \op{Buffer} with $n$-tuple device: axis $= 0$ (device dim).
+\op{Reshape} remaps axis to preserve the shard boundary. \op{Permute} follows the permutation. \op{Expand} shifts axis right by $|\mathbf{n}|$.
+\op{Reduce} on the shard axis $\to$ \textsc{null} (shard axis is among the first $n$ axes). \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}.

 %% ============================================================
 \subsection*{Kernel Optimizations (OptOps) \normalfont\small--- schedule-level transforms on kernel ranges}
@@ -300,7 +297,6 @@ Each kernel's iteration space is a set of \op{Range} axes. Every range has an \t
 \toprule
 \textbf{AxisType} & \textbf{Letter} & \textbf{Split from} & \textbf{Direction} & \textbf{Semantics} \\
 \midrule
-{\color{axgreen}\texttt{DEVICE}} & \texttt{d} & --- & --- & Multi-device sharding dimension. \\
 {\color{axblue}\texttt{GLOBAL}} & \texttt{g} & --- & --- & GPU global workgroup dimension. \\
 {\color{axcyan}\texttt{LOCAL}} & \texttt{l} & g, L & inner & Workgroup local dimension (shared memory). \\
 {\color{axbrcyan}\texttt{WARP}} & \texttt{w} & \multicolumn{2}{l}{(created by \op{TC})} & Warp-level lanes for tensor cores. \\
@@ -382,8 +378,6 @@ def scatter_add(T, idx, val):
 Let $D = (d_0, \ldots, d_{n-1})$ be an $n$-tuple device.
 \op{Copy} to an $n$-tuple device reshards with axis $= 0$. \op{Copy} never changes shape.

-\textbf{Sharding} splits a tensor along an axis across $n$ devices. It opens a \op{Range} of type \texttt{DEVICE} (a symbolic per-device index $d$), shrinks each device's view to its piece, then closes the range with \op{Unshard}$(T, R, a)$. The result is a logical tensor whose shape along axis $a$ is the full size; each device holds $1/n$ of it. \op{Unshard} is the inverse of sharding --- it marks the boundary between per-device computation and the logical multi-device tensor. The range need not be \texttt{DEVICE}; e.g.\ a \texttt{WARP} range closes the same way, concatenating per-lane shards along $a$ with the range as the outer factor. A tensor may be sharded along several axes at once: \op{Unshard}$(T, R_0, R_1, \ldots;\; a_0, a_1, \ldots)$ carries one \op{Range} per sharded axis, and every movement op maps each sharded axis independently.
-
 \begin{lstlisting}
 # T has shape (s,) on a single device.

@@ -11,8 +11,8 @@ from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op, SOP2Op, DSOp

 def _srcs():
  """Create minimal source variables for pcode parsing."""
-  def u32(v=0): return UOp.const(v, dtypes.uint32)
-  return {'S0': u32(), 'S1': u32(), 'S2': u32(), 'SCC': u32(), 'VCC': UOp.const(0, dtypes.uint64), 'laneId': u32()}
+  def u32(v=0): return UOp.const(dtypes.uint32, v)
+  return {'S0': u32(), 'S1': u32(), 'S2': u32(), 'SCC': u32(), 'VCC': UOp.const(dtypes.uint64, 0), 'laneId': u32()}

 class TestBasicParsing(unittest.TestCase):
  """Test basic pcode parsing for common instruction patterns."""
@@ -44,8 +44,8 @@ class TestWithSources(unittest.TestCase):

  def test_v_add_f32_with_sources(self):
    """Test V_ADD_F32 with actual float constants."""
-    s0 = UOp.const(0x3f800000, dtypes.uint32)  # 1.0f
-    s1 = UOp.const(0x40000000, dtypes.uint32)  # 2.0f
+    s0 = UOp.const(dtypes.uint32, 0x3f800000)  # 1.0f
+    s1 = UOp.const(dtypes.uint32, 0x40000000)  # 2.0f
    _, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_F32_E32], {'S0': s0, 'S1': s1})
    self.assertEqual(len(assigns), 1)
    dest, val = assigns[0]
@@ -55,8 +55,8 @@ class TestWithSources(unittest.TestCase):

  def test_v_mul_f32_with_sources(self):
    """Test V_MUL_F32 with actual float constants."""
-    s0 = UOp.const(0x40000000, dtypes.uint32)  # 2.0f
-    s1 = UOp.const(0x40400000, dtypes.uint32)  # 3.0f
+    s0 = UOp.const(dtypes.uint32, 0x40000000)  # 2.0f
+    s1 = UOp.const(dtypes.uint32, 0x40400000)  # 3.0f
    _, assigns = parse_pcode(PCODE[VOP2Op.V_MUL_F32_E32], {'S0': s0, 'S1': s1})
    self.assertEqual(len(assigns), 1)
    dest, val = assigns[0]
@@ -67,36 +67,36 @@ class TestParseExpr(unittest.TestCase):

  def test_integer_literals(self):
    """Test parsing integer literals."""
-    self.assertEqual(parse_expr('0', {}).val, 0)
-    self.assertEqual(parse_expr('42', {}).val, 42)
-    self.assertEqual(parse_expr('42U', {}).val, 42)
+    self.assertEqual(parse_expr('0', {}).arg, 0)
+    self.assertEqual(parse_expr('42', {}).arg, 42)
+    self.assertEqual(parse_expr('42U', {}).arg, 42)

  def test_negative_integers(self):
    """Test parsing negative integer literals."""
    result = parse_expr('-1', {})
-    self.assertEqual(result.val, -1)
+    self.assertEqual(result.arg, -1)
    self.assertEqual(result.dtype, dtypes.int)

  def test_float_literals(self):
    """Test parsing float literals."""
    result = parse_expr('1.0F', {})
-    self.assertEqual(result.val, 1.0)
+    self.assertEqual(result.arg, 1.0)
    self.assertEqual(result.dtype, dtypes.float32)

  def test_hex_literals(self):
    """Test parsing hex literals."""
    result = parse_expr('0xFF', {})
-    self.assertEqual(result.val, 255)
+    self.assertEqual(result.arg, 255)

  def test_variable_lookup(self):
    """Test variable lookup in parse_expr."""
-    vrs = {'x': UOp.const(42, dtypes.uint32)}
+    vrs = {'x': UOp.const(dtypes.uint32, 42)}
    result = parse_expr('x', vrs)
-    self.assertEqual(result.val, 42)
+    self.assertEqual(result.arg, 42)

  def test_binary_ops(self):
    """Test parsing binary operations."""
-    vrs = {'a': UOp.const(10, dtypes.uint32), 'b': UOp.const(5, dtypes.uint32)}
+    vrs = {'a': UOp.const(dtypes.uint32, 10), 'b': UOp.const(dtypes.uint32, 5)}

    # Addition
    result = parse_expr('a + b', vrs)
@@ -105,11 +105,11 @@ class TestParseExpr(unittest.TestCase):
    # Subtraction with constant folding
    result = parse_expr('10 - 5', {})
    self.assertEqual(result.op, Ops.CONST)
-    self.assertEqual(result.val, 5)
+    self.assertEqual(result.arg, 5)

  def test_ternary(self):
    """Test parsing ternary expressions."""
-    vrs = {'cond': UOp.const(True), 'a': UOp.const(1, dtypes.uint32), 'b': UOp.const(0, dtypes.uint32)}
+    vrs = {'cond': UOp.const(dtypes.bool, True), 'a': UOp.const(dtypes.uint32, 1), 'b': UOp.const(dtypes.uint32, 0)}
    result = parse_expr('cond ? a : b', vrs)
    self.assertEqual(result.op, Ops.WHERE)

@@ -127,7 +127,7 @@ class TestForLoopParsing(unittest.TestCase):
  def test_clz_parsing(self):
    """Test CLZ pcode parsing produces correct structure."""
    pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
-    S0 = UOp.const(0xFFFFFFFF, dtypes.uint32)  # All ones - CLZ should be 0
+    S0 = UOp.const(dtypes.uint32, 0xFFFFFFFF)  # All ones - CLZ should be 0
    _vrs, assigns = parse_pcode(pcode, {'S0': S0})

    self.assertEqual(len(assigns), 1)
@@ -139,7 +139,7 @@ class TestForLoopParsing(unittest.TestCase):
  def test_clz_with_zero(self):
    """Test CLZ with input 0 - should return -1."""
    pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
-    S0 = UOp.const(0, dtypes.uint32)
+    S0 = UOp.const(dtypes.uint32, 0)
    _vrs, assigns = parse_pcode(pcode, {'S0': S0})

    # Check that the innermost value (default) is -1 (may be wrapped in CAST)
@@ -150,7 +150,7 @@ class TestForLoopParsing(unittest.TestCase):
    # Unwrap CAST if present
    while val.op == Ops.CAST:
      val = val.src[0]
-    self.assertEqual(val.val, -1)
+    self.assertEqual(val.arg, -1)

  def test_ctz_parsing(self):
    """Test CTZ pcode parsing."""
@@ -158,7 +158,7 @@ class TestForLoopParsing(unittest.TestCase):
    if pcode is None:
      self.skipTest("V_CTZ_I32_B32_E32 pcode not available")

-    S0 = UOp.const(1, dtypes.uint32)  # LSB set - CTZ should be 0
+    S0 = UOp.const(dtypes.uint32, 1)  # LSB set - CTZ should be 0
    _vrs, assigns = parse_pcode(pcode, {'S0': S0})
    self.assertEqual(len(assigns), 1)

@@ -169,8 +169,8 @@ class TestDSPcodePatterns(unittest.TestCase):
    """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype."""
    vmem = UOp.param(2, dtypes.uint32, (1024,))
    srcs = {
-      'ADDR': UOp.const(0, dtypes.uint64),
-      'DATA': UOp.const(0x3f800000, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint64, 0),
+      'DATA': UOp.const(dtypes.uint32, 0x3f800000),
      '_vmem': vmem,
    }

@@ -199,8 +199,8 @@ class TestDSPcodePatterns(unittest.TestCase):
    """Test MEM[addr].type read expression parsing."""
    # Create a mock LDS buffer
    lds = UOp.param(3, dtypes.uint32, (16384,))
-    addr = UOp.const(0, dtypes.uint32)
-    vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(0, dtypes.uint32)}
+    addr = UOp.const(dtypes.uint32, 0)
+    vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}

    result = parse_expr('MEM[ADDR + OFFSET].b32', vrs)
    # Should be an INDEX operation into LDS
@@ -212,13 +212,13 @@ class TestDSPcodePatterns(unittest.TestCase):
    self.assertIsNotNone(pcode)
    assert pcode is not None
    srcs = {
-      'ADDR': UOp.const(0, dtypes.uint32),
-      'OFFSET0': UOp.const(0, dtypes.uint32),
-      'OFFSET1': UOp.const(1, dtypes.uint32),
-      'DATA': UOp.const(0xAAAAAAAA, dtypes.uint32),
-      'DATA2': UOp.const(0xBBBBBBBB, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint32, 0),
+      'OFFSET0': UOp.const(dtypes.uint32, 0),
+      'OFFSET1': UOp.const(dtypes.uint32, 1),
+      'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
+      'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
    }
-    srcs['laneId'] = UOp.const(0, dtypes.uint32)
+    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
    _, assigns = parse_pcode(pcode, srcs)
    # Should have 2 MEM write assignments
    self.assertEqual(len(assigns), 2)
@@ -235,12 +235,12 @@ class TestDSPcodePatterns(unittest.TestCase):
    assert pcode is not None
    lds = UOp.param(3, dtypes.uint32, (16384,))
    srcs = {
-      'ADDR': UOp.const(0, dtypes.uint32),
-      'OFFSET0': UOp.const(0, dtypes.uint32),
-      'OFFSET1': UOp.const(1, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint32, 0),
+      'OFFSET0': UOp.const(dtypes.uint32, 0),
+      'OFFSET1': UOp.const(dtypes.uint32, 1),
      '_lds': lds,
    }
-    srcs['laneId'] = UOp.const(0, dtypes.uint32)
+    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
    _, assigns = parse_pcode(pcode, srcs)
    # Should have 2 RETURN_DATA assignments
    self.assertEqual(len(assigns), 2)
@@ -252,36 +252,36 @@ class TestDSPcodePatterns(unittest.TestCase):
    pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
    assert pcode is not None
    srcs = {
-      'ADDR': UOp.const(100, dtypes.uint32),
-      'OFFSET0': UOp.const(2, dtypes.uint32),
-      'OFFSET1': UOp.const(5, dtypes.uint32),
-      'DATA': UOp.const(0xAAAAAAAA, dtypes.uint32),
-      'DATA2': UOp.const(0xBBBBBBBB, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint32, 100),
+      'OFFSET0': UOp.const(dtypes.uint32, 2),
+      'OFFSET1': UOp.const(dtypes.uint32, 5),
+      'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
+      'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
    }
-    srcs['laneId'] = UOp.const(0, dtypes.uint32)
+    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
    _, assigns = parse_pcode(pcode, srcs)
    # Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120
    # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
-    self.assertEqual(assigns[0][1][0].simplify().val, 108)  # type: ignore[index]
-    self.assertEqual(assigns[1][1][0].simplify().val, 120)  # type: ignore[index]
+    self.assertEqual(assigns[0][1][0].simplify().arg, 108)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][0].simplify().arg, 120)  # type: ignore[index]

  def test_ds_store_data_values(self):
    """Test DS_STORE_2ADDR_B32 uses correct data values."""
    pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
    assert pcode is not None
    srcs = {
-      'ADDR': UOp.const(0, dtypes.uint32),
-      'OFFSET0': UOp.const(0, dtypes.uint32),
-      'OFFSET1': UOp.const(1, dtypes.uint32),
-      'DATA': UOp.const(0xAAAAAAAA, dtypes.uint32),
-      'DATA2': UOp.const(0xBBBBBBBB, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint32, 0),
+      'OFFSET0': UOp.const(dtypes.uint32, 0),
+      'OFFSET1': UOp.const(dtypes.uint32, 1),
+      'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
+      'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
    }
-    srcs['laneId'] = UOp.const(0, dtypes.uint32)
+    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
    _, assigns = parse_pcode(pcode, srcs)
    # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
    # DATA[31:0] should preserve the value
-    self.assertEqual(assigns[0][1][1].simplify().val, 0xAAAAAAAA)  # type: ignore[index]
-    self.assertEqual(assigns[1][1][1].simplify().val, 0xBBBBBBBB)  # type: ignore[index]
+    self.assertEqual(assigns[0][1][1].simplify().arg, 0xAAAAAAAA)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][1].simplify().arg, 0xBBBBBBBB)  # type: ignore[index]

 class TestConditionalParsing(unittest.TestCase):
  """Test conditional (if/elsif/else) pcode parsing."""
@@ -290,9 +290,9 @@ class TestConditionalParsing(unittest.TestCase):
    """Test parsing ternary expression (which becomes WHERE)."""
    # S_CSELECT_B32: D0.u32 = SCC ? S0.u32 : S1.u32
    pcode = PCODE[SOP2Op.S_CSELECT_B32]
-    s0 = UOp.const(10, dtypes.uint32)
-    s1 = UOp.const(20, dtypes.uint32)
-    scc = UOp.const(1, dtypes.uint32)
+    s0 = UOp.const(dtypes.uint32, 10)
+    s1 = UOp.const(dtypes.uint32, 20)
+    scc = UOp.const(dtypes.uint32, 1)
    _vrs, assigns = parse_pcode(pcode, {'S0': s0, 'S1': s1, 'SCC': scc})
    self.assertEqual(len(assigns), 1)
    dest, val = assigns[0]
@@ -305,26 +305,26 @@ class TestConcatWidthParsing(unittest.TestCase):

  def test_permlanex16_altrow_concat(self):
    for row, expected in [(0, 1), (1, 0), (2, 3), (3, 2)]:
-      parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(row, dtypes.uint32)})
-      self.assertEqual(parsed.simplify().val, expected)
+      parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(dtypes.uint32, row)})
+      self.assertEqual(parsed.simplify().arg, expected)

  def test_permlane64_altlane_concat(self):
    for lane, expected in [(0, 32), (1, 33), (31, 63), (32, 0), (63, 31)]:
-      parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(lane, dtypes.uint32)})
-      self.assertEqual(parsed.simplify().val, expected)
+      parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(dtypes.uint32, lane)})
+      self.assertEqual(parsed.simplify().arg, expected)

  def test_permlane64_wave64_pcode_indices(self):
    vgpr = UOp.param(0, dtypes.uint32, (256,))
    srcs = {
-      'SRC0': UOp.const(0, dtypes.uint32),
-      'VDST': UOp.const(1, dtypes.uint32),
-      'EXEC_LO': UOp.const(0xFFFFFFFF, dtypes.uint32),
-      'EXEC': UOp.const(0xFFFFFFFFFFFFFFFF, dtypes.uint64),
+      'SRC0': UOp.const(dtypes.uint32, 0),
+      'VDST': UOp.const(dtypes.uint32, 1),
+      'EXEC_LO': UOp.const(dtypes.uint32, 0xFFFFFFFF),
+      'EXEC': UOp.const(dtypes.uint64, 0xFFFFFFFFFFFFFFFF),
      '_vgpr': vgpr,
      '_wave_size': 64,
-      'S0': UOp.const(0, dtypes.uint32),
-      'S1': UOp.const(0, dtypes.uint32),
-      'S2': UOp.const(0, dtypes.uint32),
+      'S0': UOp.const(dtypes.uint32, 0),
+      'S1': UOp.const(dtypes.uint32, 0),
+      'S2': UOp.const(dtypes.uint32, 0),
    }

    def load_idx(v: UOp) -> int:
@@ -333,12 +333,12 @@ class TestConcatWidthParsing(unittest.TestCase):
      self.assertEqual(simp.src[0].op, Ops.INDEX)
      idx = simp.src[0].src[1].simplify()
      self.assertEqual(idx.op, Ops.CONST)
-      return idx.val
+      return idx.arg

    _, assigns = parse_pcode(PCODE[VOP1Op.V_PERMLANE64_B32_E32], srcs)
    self.assertEqual(len(assigns), 64)
    for lane, (dst_idx, src_idx) in {0: (64, 32), 31: (95, 63), 32: (96, 0), 63: (127, 31)}.items():
-      self.assertEqual(assigns[lane][1][0].simplify().val, dst_idx)  # type: ignore[index]
+      self.assertEqual(assigns[lane][1][0].simplify().arg, dst_idx)  # type: ignore[index]
      self.assertEqual(load_idx(assigns[lane][1][1]), src_idx)  # type: ignore[index]

 class TestAllPcode(unittest.TestCase):
@@ -346,7 +346,7 @@ class TestAllPcode(unittest.TestCase):

  def _make_srcs(self):
    """Create dummy source variables for pcode parsing."""
-    u32, u64 = lambda v=0: UOp.const(v, dtypes.uint32), lambda v=0: UOp.const(v, dtypes.uint64)
+    u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v)
    lds = UOp.param(3, dtypes.uint32, (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(),
@@ -358,7 +358,7 @@ class TestAllPcode(unittest.TestCase):
            'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'ROUND_TOWARD_ZERO': u32(),
            'ROUND_NEAREST_EVEN': u32(), 'WAVE_STATUS': u32(),
            'MAX_FLOAT_F32': u32(0x7f7fffff), 'Unsigned': u32(1), 'clampedLOD': u32(),
-            '_lds': lds, '_vmem': lds, '_active': UOp.const(True)}
+            '_lds': lds, '_vmem': lds, '_active': UOp.const(dtypes.bool, True)}

  def _parse_all_pcode(self, pcode_dict, arch: str, min_pct: float):
    """Parse all pcode. RuntimeError = parser limitation (ok), other exceptions = real bugs."""
@@ -19,12 +19,6 @@ class TestArange(unittest.TestCase):
    self.assertLess(self._get_flops(Tensor.arange(256).clone(), np.arange(256)), 256*4)
    self.assertLess(self._get_flops(Tensor.arange(2560).clone(), np.arange(2560)), 2560*4)

-  def test_cat_complexity(self):
-    x = Tensor.arange(2**10) + Tensor.empty((), dtype=dtypes.uint32)
-    out = x.cat(x).cat(Tensor.empty(1, dtype=dtypes.uint32))
-    linear = compile_linear(out.schedule_linear())
-    self.assertLessEqual(estimate_uop(linear.src[-1]).ops, out.numel()*20)
-
  @unittest.skipIf(Device.DEFAULT == "CL", "flaky in CI")
  def test_arange_cumsum(self):
    np.testing.assert_equal(Tensor.arange(513).cumsum(0).numpy(), np.arange(513).cumsum())
@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.helpers import getenv, system, DEV
-from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm, quantize_mxfp4
+from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm
 from test.helpers import needs_second_gpu
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX

@@ -150,45 +150,6 @@ class TestAsmGEMM(unittest.TestCase):
    with self.assertRaisesRegex(AssertionError, "not a multiple"):
      verify_asm_gemm(1, 256, 1000, 256)

-class TestMXFP4(unittest.TestCase):
-  def setUp(self):
-    if not is_cdna4() or DEV.interface.startswith("MOCK"):
-      self.skipTest("requires real amd machine")
-
-  def test_quantize(self):
-    import numpy as np
-    block = np.array([0, .26, .74, .75, 1.26, 1.75, 2.51, 3.5, 5.1, 6, -6] + [0] * 21, dtype=np.float32)
-    x = Tensor(np.tile(block, (32, 8)), dtype=dtypes.bfloat16)
-    packed, scale, _ = quantize_mxfp4(x)
-    p = packed.numpy()
-    codes = np.stack((p & 0xF, p >> 4), axis=-1).reshape(32, 256)
-    np.testing.assert_array_equal(codes[0, :11], [0, 1, 1, 2, 3, 4, 5, 6, 7, 7, 15])
-    np.testing.assert_array_equal(scale.numpy(), np.full((32, 8), 127, dtype=np.uint8))
-
-  def test_correctness(self):
-    import numpy as np
-    M = N = K = 256
-    rng = np.random.default_rng(1)
-    a = Tensor(rng.standard_normal((M, K), dtype=np.float32), dtype=dtypes.bfloat16)
-    b = Tensor(rng.standard_normal((N, K), dtype=np.float32), dtype=dtypes.bfloat16)
-    out = asm_gemm(a, b.T, mxfp4=True).realize()
-    # reference gemm
-    a_packed, scale_a, _ = quantize_mxfp4(a)
-    b_packed, scale_b, _ = quantize_mxfp4(b)
-    def unpack(x): return np.stack((x & 0xF, x >> 4), axis=-1).reshape(x.shape[0], -1)
-    code_a, code_b = unpack(a_packed.numpy()), unpack(b_packed.numpy())
-    lut = np.array([0, .5, 1, 1.5, 2, 3, 4, 6, -0., -.5, -1, -1.5, -2, -3, -4, -6], dtype=np.float32)
-    a_dequant = lut[code_a] * np.repeat(np.exp2(scale_a.numpy().astype(np.int16)-127), 32, axis=1)
-    b_dequant = lut[code_b] * np.repeat(np.exp2(scale_b.numpy().astype(np.int16)-127), 32, axis=1)
-    ref = Tensor(a_dequant @ b_dequant.T, dtype=dtypes.bfloat16).realize().numpy()
-    np.testing.assert_array_equal(out.numpy(), ref)
-
-  def test_empty(self):
-    M, N, K = getenv("M", 16384), getenv("N", 4096), getenv("K", 14336)
-    a = Tensor.empty(M, K, dtype=dtypes.bfloat16)
-    b = Tensor.empty(N, K, dtype=dtypes.bfloat16)
-    asm_gemm(a, b.T, mxfp4=True).realize()
-
 # test the Asm GEMM with Llama shapes, only run on the real machine for speed

@unittest.skipUnless(has_hipcc(), "requires hipcc to compile")
@@ -6,11 +6,11 @@ from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.uop.ops import KernelInfo

 def call_out_kernel(F:UOp, C:UOp) -> UOp:
-  call = F[0].load().call(UOp.const(3).cast(dtypes.int), C[0], ret_dtype=dtypes.void)
+  call = F[0].load().call(UOp.const(dtypes.int, 3), C[0], ret_dtype=dtypes.void)
  return C.after(call)[1].store(C.after(call)[0].load() + 1).sink(arg=KernelInfo(name="call_out"))

 def call_ret_kernel(F:UOp, C:UOp) -> UOp:
-  val = F[0].load().call(UOp.const(21).cast(dtypes.int), ret_dtype=dtypes.int)
+  val = F[0].load().call(UOp.const(dtypes.int, 21), ret_dtype=dtypes.int)
  return C[0].store(val * 2).sink(arg=KernelInfo(name="call_ret"))

@unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "TODO: CALL is rendered in C style only")
@@ -16,7 +16,7 @@ def _check_ast_count(desired_count:int, t:Tensor):

 class TestMovedConstFolding(unittest.TestCase):
  def test_contiguous_deviceless_const(self):
-    t = Tensor(UOp.const(2.0, dtypes.float)).contiguous()
+    t = Tensor(UOp.const(dtypes.float, 2.0)).contiguous()
    self.assertIs(t.uop.op, Ops.CONST)
    self.assertIsNone(t.uop.device)

@@ -169,7 +169,7 @@ class TestMultiConstFolding(unittest.TestCase):
 class TestThreefryConstFolding(unittest.TestCase):
  def test_threefry(self):
    # THREEFRY(const,const) folds to a const once decomposed
-    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64))
+    x = threefry2x32(UOp.const(dtypes.uint64, 5), UOp.const(dtypes.uint64, 10))
    self.assertIs(x.simplify().op, Ops.CONST)

 class TestTautologicalCompare(unittest.TestCase):
@@ -1,9 +1,7 @@
 import unittest
 from tinygrad import Tensor, UOp, GlobalCounters, Context, Device
-import numpy as np
 from tinygrad.dtype import AddrSpace, dtypes, Invalid
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops
-from tinygrad.renderer.ptx import PTXRenderer

 # **** kernels ****

@@ -57,7 +55,7 @@ def flip_contract_kernel(dest:UOp, src:UOp):
  return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.numel()}", opts_to_apply=()))

 def slice_sum_kernel(dest:UOp, src:UOp):
-  G = UOp.range(src.shape[0], 0, dtype=dtypes.int)
+  G = UOp.range(src.shape[0], 0)
  slice_src = src[G, :]
  reg = UOp.placeholder((1,), dest.dtype, 0, addrspace=AddrSpace.REG)
  reg = reg.after(G)[0].set(0)
@@ -119,20 +117,13 @@ class TestCustomKernel(unittest.TestCase):
    out = c.flatten().tolist()
    assert all(x == 2 for x in out), "all 2"

-  def test_duplicate_call_arg(self):
-    x = Tensor.arange(4).clone().realize()
-    x = Tensor.custom_kernel(x, x, fxn=custom_add_one_kernel)[0]
-    # webgpu silently errors when a kernel has duplicate buffer args, so the list stays the same.
-    # https://gpuweb.github.io/gpuweb/#abstract-opdef-encoder-bind-groups-alias-a-writable-resource
-    self.assertEqual(x.tolist(), [1, 2, 3, 4] if Device.DEFAULT != "WEBGPU" else [0, 1, 2, 3])
-
  def test_simple_sharded(self):
    devs = ("CPU:0", "CPU:1")

    a = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
    b = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
    # ugly construction to get a sharded empty tensor
-    c = Tensor(Tensor.empty(8, 16, device=devs).uop.unshard(0), device=devs)
+    c = Tensor(Tensor.empty(8, 16, device=devs).uop.multi(0), device=devs)
    c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
    out = c.flatten().tolist()
    assert all(x == 2 for x in out), "all 2"
@@ -141,7 +132,7 @@ class TestCustomKernel(unittest.TestCase):
    # PYTHON backend explicitly checks for OOB access for wrong multi shape regression
    devs = ("PYTHON:0", "PYTHON:1")
    a = Tensor.ones(4, 4).contiguous().shard(devs, axis=0)
-    c = Tensor(Tensor.empty(2, 4, device=devs).uop.unshard(0), device=devs)
+    c = Tensor(Tensor.empty(2, 4, device=devs).uop.multi(0), device=devs)
    c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
    assert (c == 2).all().item()

@@ -215,7 +206,7 @@ class TestCustomKernel(unittest.TestCase):
    N = 16
    a = Tensor.randn(N, N).shard_(devs, axis=0)
    b = Tensor.randn(N, N).to(devs)
-    c = Tensor(Tensor.empty(N//2, N, device=devs).uop.unshard(0), device=devs)
+    c = Tensor(Tensor.empty(N//2, N, device=devs).uop.multi(0), device=devs)
    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
    self.assertTrue(tst.allclose(a@b, atol=1e-3).item())

@@ -332,7 +323,7 @@ class TestCustomKernel(unittest.TestCase):
  def test_multi_invalids_custom_kernel_no_copy(self):
    devs = ("CPU:0", "CPU:1")
    a = Tensor.ones(4, 4).shard(devs, axis=0).realize()
-    c = Tensor(Tensor.invalids(2, 4, dtype=dtypes.float, device=devs).uop.unshard(0), device=devs)
+    c = Tensor(Tensor.invalids(2, 4, dtype=dtypes.float, device=devs).uop.multi(0), device=devs)
    c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
    GlobalCounters.reset()
    c.realize()
@@ -341,7 +332,7 @@ class TestCustomKernel(unittest.TestCase):

  def test_partial_invalid_store_keeps_uncovered_reads(self):
    x = Tensor([10., 20., 30., 40.])
-    after = x.uop.after(x.uop.shrink(((0, 2),)).store(Invalid))
+    after = x.uop.after(x.uop.shrink(((0, 2),)).store(UOp.const(dtypes.float, Invalid, shape=(2,))))
    self.assertEqual(Tensor(after).contiguous().tolist(), [10., 20., 30., 40.])

  def test_multi_after_invalid_store_dep_removed(self):
@@ -397,7 +388,7 @@ class TestCustomKernel(unittest.TestCase):
    y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
    if use_custom:
      z = Tensor.empty_like(x)
-      z = Tensor.custom_kernel(z, y.T.T, fxn=custom_add_one_kernel)[0]
+      z = Tensor.custom_kernel(y, y.T.T, fxn=custom_add_one_kernel)[0]
    else: z = y.T.T+1
    GlobalCounters.reset()
    z.realize()
@@ -423,181 +414,19 @@ class TestCustomKernel(unittest.TestCase):

  @Context(DEV="CPU")
  def test_simple_from_source(self):
-    a = Tensor.arange(4).clone().realize()
-    src = "void test_src(int* restrict a) { a[0] = 1; }"
+    a = Tensor([0., 1., 2.]).realize()
+
+    src = "void test_src(float* restrict a) { a[0] = 1.0; }"
    # TODO: it currently requires a compiler for Ops.BINARY
    from tinygrad.device import Device
    binary = Device[a.device].renderer.compiler.compile(src)
    def custom_src_kernel(A:UOp) -> UOp:
      sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
-    a = Tensor.custom_kernel(a.reshape(2, 2).T, fxn=custom_src_kernel)[0]
-    self.assertEqual(a.tolist(), [[1, 2], [1, 3]])
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())),
+                                   UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))

-class TestCustomKernelInput(unittest.TestCase):
-  def _test_mop(self, mop_fxn, max_kernels):
-    # default: input is BUFFER
-    x = mop_fxn(Tensor.arange(32).clone("CPU").realize())
-    y = Tensor.custom_kernel(Tensor.empty_like(x), x, fxn=custom_add_one_kernel)[0]
-    GlobalCounters.reset()
-    y.realize()
-    kernel_count = GlobalCounters.kernel_count
-    self.assertEqual(y.tolist(), x.add(1).tolist())
-    self.assertLessEqual(kernel_count, max_kernels)
-    # same test with @function, input is PARAM
-    from tinygrad import function
-    x0 = Tensor.arange(32).clone("CPU").realize()
-    @function(precompile=True)
-    def run(a:Tensor) -> Tensor:
-      xv = mop_fxn(a)
-      y = Tensor.invalids(*xv.shape, dtype=xv.dtype, device=a.device)
-      return Tensor.custom_kernel(y, xv, fxn=custom_add_one_kernel)[0]
-    GlobalCounters.reset()
-    y = run(x0).realize()
-    kernel_count = GlobalCounters.kernel_count
-    self.assertEqual(y.tolist(), mop_fxn(x0).add(1).tolist())
-    self.assertLessEqual(kernel_count, max_kernels)
-
-  def test_reshape(self): self._test_mop(lambda x: x.reshape(16, 2), max_kernels=2)
-  def test_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T, max_kernels=3)
-  def test_double_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T.T, max_kernels=3)
-  def test_shrink(self): self._test_mop(lambda x: x[:4], max_kernels=2)
-  def test_pad(self): self._test_mop(lambda x: x[:4].pad(((0, 4),)), max_kernels=2)
-  def test_flip(self): self._test_mop(lambda x: x.flip(0), max_kernels=2)
-  def test_offset_shrink(self): self._test_mop(lambda x: x[4:8], max_kernels=2)
-  def test_2d_shrink(self): self._test_mop(lambda x: x.reshape(4, 8)[:, 2:6], max_kernels=3)
-  def test_expand(self): self._test_mop(lambda x: x.reshape(16, 2)[:, :1].expand(16, 2), max_kernels=3)
-
-class TestUnshardIndex(unittest.TestCase):
-  """Regression tests for INDEX on UNSHARD (fragment) resolution in schedule/multi.py.
-
-  A fragment is a per-thread REG buffer wrapped in UNSHARD over LOCAL thread ranges.
-  index_multi must resolve an INDEX on the UNSHARD view into an INDEX on the per-thread
-  shard. Two ownership patterns must work:
-    contiguous: idx = rng*shard_sz + local   (thread rng owns [rng*shard_sz, ...))
-    strided:    idx = rng + ir*shard_sz      (thread rng owns {rng, rng+shard_sz, ...})
-  """
-  def _run(self, kernel, shape=(8, 8)):
-    c = Tensor.empty(*shape)
-    out = Tensor.custom_kernel(c, fxn=kernel)[0]
-    try: return out.numpy()
-    except RuntimeError as e:
-      if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e):
-        self.skipTest("PTX does not support dynamic register indexing")
-      raise
-
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_contiguous_fragment_index(self):
-    # thread ty owns rows [ty*8, ty*8+8) of a 64-row fragment -- contiguous ownership.
-    # This is the pre-existing case that index_multi always handled.
-    def kernel(C:UOp) -> UOp:
-      ty = UOp.range(8, 0, AxisType.LOCAL)
-      ir = UOp.range(8, 1, AxisType.LOOP)
-      j = UOp.range(8, 2, AxisType.LOOP)
-      # 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows [ty*8, ty*8+8).
-      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
-      return C[ty*8 + ir, j].store(frag[ty*8 + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="contig_frag"))
-    out = self._run(kernel, (64, 8))
-    assert out.shape == (64, 8)
-
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_strided_fragment_index(self):
-    # thread ty owns rows {ty, ty+8, ty+16, ty+24, ..., ty+56} of a 64-row fragment --
-    # strided ownership. idx = ty + ir*8 where shard_sz=8 (8 threads, shard rows=8).
-    # The contiguous check (idx - rng*shard_sz) fails; the strided check
-    # (idx-rng) % shard_sz == 0 must succeed. This is the pattern the index_multi fix adds.
-    def kernel(C:UOp) -> UOp:
-      ty = UOp.range(8, 0, AxisType.LOCAL)
-      ir = UOp.range(8, 1, AxisType.LOOP)
-      j = UOp.range(8, 2, AxisType.LOOP)
-      # 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows {ty, ty+8, ..., ty+56}.
-      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
-      return C[ty + ir*8, j].store(frag[ty + ir*8, j]).end(j, ir, ty).sink(arg=KernelInfo(name="strided_frag"))
-    out = self._run(kernel, (64, 8))
-    assert out.shape == (64, 8)
-
-  def test_fragment_index_cannot_shard(self):
-    # thread ty indexing rows [ty, ty+8) overlaps with other threads' rows -- this matches neither
-    # the contiguous nor the strided ownership pattern, so index_multi must raise.
-    def kernel(C:UOp) -> UOp:
-      ty = UOp.range(8, 0, AxisType.LOCAL)
-      ir = UOp.range(8, 1, AxisType.LOOP)
-      j = UOp.range(8, 2, AxisType.LOOP)
-      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
-      return C[ty + ir, j].store(frag[ty + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="bad_frag"))
-    with self.assertRaisesRegex(RuntimeError, "cannot shard index"):
-      self._run(kernel, (64, 8))
-
-def _run_fragment_kernel(testcase, kernel, out_shape, inputs=()):
-  c = Tensor.empty(*out_shape)
-  out = Tensor.custom_kernel(c, *inputs, fxn=kernel)[0]
-  try: return out.numpy()
-  except RuntimeError as e:
-    if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e):
-      testcase.skipTest("PTX does not support dynamic register indexing")
-    raise
-
-class TestUnshardAlu(unittest.TestCase):
-  """Tests for ALU on (fragment) UNSHARD values in schedule/multi.py's alu_multi.
-
-  An ALU with UNSHARD srcs lowers to per-shard ops when every src is one of:
-    same sharding:  peel the UNSHARD, keep the layout
-    scalar:         broadcast to every shard
-    whole unsharded same-shape value: takes its per-shard sub-view (shard_subview)
-  """
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_alu_scalar_broadcast(self):
-    # scalar srcs broadcast to every shard: frag*2.0 where frag is 1.5 per thread -> 3.0 everywhere
-    def kernel(C:UOp) -> UOp:
-      ty = UOp.range(8, 0, AxisType.LOCAL)
-      # 8 values per thread, 8 threads -> 64-value full view
-      frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
-      v = frag.after(frag.store(1.5)) * 2.0
-      return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_scalar", opts_to_apply=()))
-    out = _run_fragment_kernel(self, kernel, (64,))
-    np.testing.assert_allclose(out, 3.0)
-
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_alu_whole_value_subview(self):
-    # UNSHARD + whole unsharded same-shape value: each shard adds its own sub-view of A.
-    def kernel(C:UOp, A:UOp) -> UOp:
-      ty = UOp.range(8, 0, AxisType.LOCAL)
-      frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
-      v = frag.after(frag.store(0.0)) + A
-      return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_subview", opts_to_apply=()))
-    a = Tensor(np.arange(64, dtype=np.float32))
-    out = _run_fragment_kernel(self, kernel, (64,), inputs=(a,))
-    np.testing.assert_allclose(out, a.numpy(), atol=1e-4)
-
-class TestUnshardStore(unittest.TestCase):
-  """Tests for STORE of a sharded value into an unsharded dest (store_value_multi in schedule/multi.py).
-
-  Every shard stores its value into its own contiguous sub-view of the dest, one SHRINK per sharded axis.
-  """
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_store_unshard_value(self):
-    # single-axis: 8 threads each own 8 values of the 64-value output tile
-    def kernel(C:UOp) -> UOp:
-      ty = UOp.range(8, 0, AxisType.LOCAL)
-      frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
-      v = frag.after(frag.store(0.0)) + 2.5
-      return C.store(v).end(ty).sink(arg=KernelInfo(name="store_unshard", opts_to_apply=()))
-    out = _run_fragment_kernel(self, kernel, (64,))
-    np.testing.assert_allclose(out, 2.5)
-
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_store_unshard_value_2axis(self):
-    # two sharded axes (the gemm fragment layout): thread (ty, tx) owns the (2, 1, 1, 2) sub-view of the
-    # (2, 4, 2, 2) output tile; the store must SHRINK dest on both sharded axes
-    def kernel(C:UOp, A:UOp) -> UOp:
-      ty = UOp.range(4, 0, AxisType.LOCAL)
-      tx = UOp.range(2, 1, AxisType.LOCAL)
-      frag = UOp.placeholder((2, 1, 1, 2), dtypes.float32, 0, AddrSpace.REG).unshard((1, 2), (ty, tx))
-      v = frag.after(frag.store(0.0)) + A
-      return C.store(v).end(tx, ty).sink(arg=KernelInfo(name="store_unshard_2axis", opts_to_apply=()))
-    a = Tensor(np.arange(32, dtype=np.float32).reshape(2, 4, 2, 2))
-    out = _run_fragment_kernel(self, kernel, (2, 4, 2, 2), inputs=(a,))
-    np.testing.assert_allclose(out, a.numpy(), atol=1e-4)
+    a = Tensor.custom_kernel(a, fxn=custom_src_kernel)[0]
+    self.assertEqual(a.tolist(), [1., 1., 2.])

 class TestUOpReduce(unittest.TestCase):
  def test_uop_sum(self):
@@ -124,7 +124,6 @@ class TestFp8sConversions(unittest.TestCase):
  def test_float_to_fp8e4m3(self, x):
    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())

-  @unittest.skip("fp8 overflow semantics are inconsistent")
  def test_float_to_fp8e4m3_extreme_values(self):
    for x in [FP8E4M3_MAX, FP8E4M3_MAX*1.01, -FP8E4M3_MAX, -FP8E4M3_MAX*1.01, math.inf, -math.inf, math.nan, -math.nan]:
      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())
@@ -7,7 +7,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
 from tinygrad.renderer.isa import IselContext

 # INDEX on a register value with a constant index extracts a single element (the old GEP)
-def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), dtype=y.dtype.scalar())
+def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(dtypes.int, i), dtype=y.dtype.scalar())

@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
 class TestIselX86(unittest.TestCase):
@@ -49,7 +49,7 @@ class TestIselX86(unittest.TestCase):
    load = UOp.param(0, dtypes.int32, (16,)).index(a + 1).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].val == 4)
+    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].arg == 4)

 if __name__ == "__main__":
  unittest.main()
@@ -13,7 +13,6 @@ from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.renderer.isa import ISARenderer
 from test.helpers import replace_opts
-from test.backend.test_softmax_fusion import single_kernel_softmax
 MOCKGPU = DEV.interface.startswith("MOCK")

 from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
@@ -268,9 +267,9 @@ class TestLinearizer(unittest.TestCase):
    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
    idxs = sorted(idxs, key=lambda uop: uop.arg)
-    assert (idxs[0].arg, idxs[0].src[0].val) == ('gidx0', 6), idxs[0]
-    assert (idxs[1].arg, idxs[1].src[0].val) == ('gidx1', 5), idxs[1].arg
-    assert (idxs[2].arg, idxs[2].src[0].val) == ('gidx2', 4), idxs[2].arg
+    assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0]
+    assert (idxs[1].arg, idxs[1].src[0].arg) == ('gidx1', 5), idxs[1].arg
+    assert (idxs[2].arg, idxs[2].src[0].arg) == ('gidx2', 4), idxs[2].arg

  def test_sum_collapse(self):
    t = Tensor([2]).reshape(1, 1).expand(256, 256).sum()
@@ -393,16 +392,6 @@ class TestLinearizer(unittest.TestCase):
    # the global store doesn't change
    assert stores[1].src[1].dtype == dtypes.float

-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
-  def test_two_grouped_stores_local(self):
-    # GROUP on both reduces puts two LOCAL buffers in one kernel, and the store to each needs its own barrier
-    a = Tensor.rand(32, 32).realize()
-    opts = [Opt(OptOps.GROUP, 1, 4), Opt(OptOps.GROUP, 2, 4)]
-    ast = helper_linearizer_opt(single_kernel_softmax(a), [opts])
-    uops = to_program(replace_opts(ast, opts), renderer=Device[Device.DEFAULT].renderer).src[1].src
-    self.assertEqual(len([u for u in uops if u.op is Ops.BARRIER]), 2)
-
 # *** helpers ***

 def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
@@ -12,18 +12,18 @@ class TestLinearizerFailure(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
  def test_failure_beam_mnist(self):
    c0 = UOp.param(0, dtypes.uchar, (4014080,))
-    c1 = UOp.range(UOp.const(512), 0, AxisType.GLOBAL)
-    c2 = UOp.range(UOp.const(784), 1, AxisType.GLOBAL)
-    c3 = UOp.range(UOp.const(10), 3, AxisType.GLOBAL)
+    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.param(1, dtypes.int, (512,))
-    c5 = c4.index(c1.valid(UOp.const(True)))
-    c6 = UOp.range(UOp.const(6000), 1004, AxisType.REDUCE)
-    c7 = UOp.range(UOp.const(3750), 2006, AxisType.REDUCE)
-    c8 = UOp.range(UOp.const(16), 2007, AxisType.GROUP_REDUCE)
+    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.param(2, dtypes.uchar, (47040000,))
-    c10 = c9.index((((c3*UOp.const(4704000))+c2)+(c6*UOp.const(784))).valid(UOp.const(True)))
-    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(6000))+c6)+((c7*UOp.const(16))+c8)).alu(Ops.CMPLT, UOp.const(59999)).where(UOp.const(0).cast(dtypes.int), UOp.const(1).cast(dtypes.int)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(-1).cast(dtypes.int))).where(UOp.const(0).cast(dtypes.uchar), c10).reduce(c6, arg=Ops.ADD)
-    c12 = c0.index((((c1*UOp.const(7840))+(c2*UOp.const(10)))+c3).valid(UOp.const(True))).store(c11).end(c1, c2, c3)
+    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)
    ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
    _ = to_program(ast, Device["METAL"].renderer)

@@ -80,7 +80,7 @@ class TestQuantizeFP8(unittest.TestCase):
  @needs_second_gpu
  def test_multi(self):
    devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(8))
-    x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.unshard(0)
+    x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.multi(0)
    x = Tensor(x, device=devs)
    amax_state = Tensor.full((), 2.0, dtype=dtypes.float32, device=devs).contiguous()
    amax_out = Tensor.zeros((), dtype=dtypes.float32, device=devs).realize()
@@ -1,6 +1,6 @@
 import unittest, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
-from tinygrad.uop.ops import Ops, UOp, AxisType
+from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters
 from tinygrad.engine.realize import run_linear, compile_linear
@@ -52,10 +52,8 @@ class TestMultiTensor(unittest.TestCase):
  def test_shard(self):
    X = Tensor.ones(256).contiguous().realize()
    X.shard_(devices_2, 0)
-    assert X.uop.src[0].shape == (128,)
-    # the MULTI carries and ends the DEVICE range as its second src
-    assert X.uop.src[1].op is Ops.RANGE and X.uop.src[1].arg[-1] is AxisType.DEVICE
-    assert X.uop.ended_ranges == X.uop.src[1:]
+    for lb in X.uop.src:
+      assert lb.shape == (128,)
    (X + X).realize()

  @unittest.expectedFailure # TODO: fix
@@ -76,13 +74,6 @@ class TestMultiTensor(unittest.TestCase):
    run_linear(linear)
    self.assertEqual(len(set(names)), 1, "function was relinearized")

-  def test_shard_beam(self):
-    cpu_2 = ("CPU:1", "CPU:2")
-    src = Tensor.ones(16).shard(cpu_2, 0).realize()
-    pad = src.to(cpu_2[::-1]).schedule_linear().src[0]
-    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): prg = compile_linear(UOp(Ops.LINEAR, src=(pad,))).src[0].src[0]
-    self.assertNotEqual(prg.src[0].arg.applied_opts, ())
-
  def test_shard_same_device(self):
    X = Tensor.ones(256).contiguous().realize()
    X.shard_((d1, X.device), 0)
@@ -432,49 +423,6 @@ class TestMultiBufferView(unittest.TestCase):
    run_linear(linear, var_vals)
    np.testing.assert_equal(out.numpy(), ref[5].numpy())

-@unittest.skipIf(not_support_multi_device(), "need multi")
-class Test2DShard(unittest.TestCase):
-  def setUp(self):
-    self.devices_4 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
-    self.rng = UOp.range(4, -1, AxisType.DEVICE)
-    self.rng0, self.rng1 = self.rng // 2, self.rng % 2
-
-  def _shard_2d(self, t:Tensor) -> Tensor:
-    u = t.uop.copy_to_device(self.devices_4)._shard(0, self.rng0)._shard(1, self.rng1).unshard((0, 1), (self.rng0, self.rng1))
-    return Tensor(u)
-
-  def test_2d_shard_basic(self):
-    ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
-    t = self._shard_2d(ref)
-    out = t.contiguous().realize()
-    np.testing.assert_equal(out.numpy(), ref.numpy())
-
-  def test_2d_shard_elementwise(self):
-    ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
-    t = self._shard_2d(ref)
-    out = (t + 1).contiguous().realize()
-    np.testing.assert_equal(out.numpy(), ref.numpy() + 1)
-
-  def test_2d_shard_sum_all(self):
-    ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
-    t = self._shard_2d(ref)
-    out = t.sum().contiguous().realize()
-    np.testing.assert_equal(out.numpy(), np.array(ref.numpy().sum()))
-
-  def test_2d_shard_sum_non_sharded_axis(self):
-    ref = Tensor.arange(4*4*2).reshape(4, 4, 2).contiguous().realize()
-    t = self._shard_2d(ref)
-    out = t.sum(axis=2).contiguous().realize()
-    np.testing.assert_equal(out.numpy(), ref.numpy().sum(axis=2))
-
-  def test_2d_shard_matmul(self):
-    a = Tensor.arange(16).reshape(4, 4).contiguous().realize()
-    b = Tensor.arange(16).reshape(4, 4).contiguous().realize()
-    a_s = self._shard_2d(a)
-    b_s = self._shard_2d(b)
-    out = (a_s @ b_s).contiguous().realize()
-    np.testing.assert_equal(out.numpy(), a.numpy() @ b.numpy())
-
@unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiTransformer(unittest.TestCase):
  @needs_second_gpu
@@ -340,11 +340,11 @@ class TestOps(unittest.TestCase):

  def test_where(self):
    helper_test_op([], lambda: torch.where(torch.tensor([True, False]), 1, 3).type(torch.int32),
-                   lambda: Tensor([True, False]).where(1, 3).clone(), forward_only=True)
+                   lambda: Tensor([True, False]).where(1, 3), forward_only=True)
    helper_test_op(
      [(100,)],
      lambda x: torch.where(x > 0.5, 4, 2).type(torch.int32),
-      lambda x: (x > 0.5).where(4, 2).clone(), forward_only=True)
+      lambda x: (x > 0.5).where(4, 2), forward_only=True)

    for shps in [[(8,),(1,),(1,)], [(10,10),(10,),(10,)], [(100,)]*3, [(10,10)]*3]:
      helper_test_op(
@@ -356,7 +356,7 @@ class TestOps(unittest.TestCase):
    helper_test_op(
      [(5, 5)],
      lambda x: torch.where(x > 0.5, 4, 2).type(torch.int32).permute((1, 0)),
-      lambda x: (x > 0.5).where(4, 2).clone().permute((1, 0)), forward_only=True)
+      lambda x: (x > 0.5).where(4, 2).permute((1, 0)), forward_only=True)

  def _test_cmp(self, fxn, reverse=True):
    # test different dtypes
@@ -636,9 +636,9 @@ class TestOps(unittest.TestCase):
        helper_test_op(None, lambda x,y: x%y, forward_only=True, vals=[va, vb])
        helper_test_op(None, lambda x: x%2, forward_only=True, vals=[va])
        helper_test_op(None, lambda x: x%3, forward_only=True, vals=[va])
-        helper_test_op(None, lambda x: x%3.5, lambda x: (x%3.5).clone(), forward_only=True, vals=[va])
+        helper_test_op(None, lambda x: x%3.5, forward_only=True, vals=[va])
        helper_test_op(None, lambda x: 100%x, forward_only=True, vals=[va])
-        helper_test_op(None, lambda x: 100.5%x, lambda x: (100.5%x).clone(), forward_only=True, vals=[va])
+        helper_test_op(None, lambda x: 100.5%x, forward_only=True, vals=[va])

  def test_fmod(self):
    a = [-4, 7, 5, 4, -7, 8, -9]
@@ -649,7 +649,7 @@ class TestOps(unittest.TestCase):
        vb = [float(bi) for bi in b] if float_b else b
        helper_test_op(None, lambda x,y: x.fmod(y), forward_only=True, vals=[va, vb])
        helper_test_op(None, lambda x: x.fmod(2), forward_only=True, vals=[va])
-        helper_test_op(None, lambda x: x.fmod(3.5), lambda x: x.fmod(3.5).clone(), forward_only=True, vals=[va])
+        helper_test_op(None, lambda x: x.fmod(3.5), forward_only=True, vals=[va])

  def test_mul_naninf(self):
    helper_test_op([(45,65)], lambda x: x*math.inf)
@@ -706,7 +706,7 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, lambda x: 0.7**x, vals=[[-2.,-1,0,1,2,3]])
    helper_test_op(None, lambda x: (-2)**x, vals=[[-2.,-1,0,1,2,3]])
    # float to power of int
-    helper_test_op(None, lambda x: 0.7**x, lambda x: (0.7**x).clone(), vals=[[-2,-1,0,1,2,3]], forward_only=True)
+    helper_test_op(None, lambda x: 0.7**x, vals=[[-2,-1,0,1,2,3]], forward_only=True)

  @unittest.skipIf(COMPILE_ONLY, "test requires runtime")
  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, NIRRenderer), "TODO: broken in LVP")
@@ -775,7 +775,7 @@ class TestOps(unittest.TestCase):
  def test_pow_int_base_float_exponent(self):
    for exponent in (0.5, 1.5, 2.0, -1.0, 0.0):
      helper_test_op([], lambda: torch.tensor([1, 2, 3, 4], dtype=torch.int) ** exponent,
-                         lambda: (Tensor([1, 2, 3, 4], dtype=dtypes.int32) ** exponent).clone(), forward_only=True)
+                         lambda: Tensor([1, 2, 3, 4], dtype=dtypes.int32) ** exponent, forward_only=True)

  def test_sqrt(self):
    helper_test_op([(45,65)], lambda x: x.sqrt())
@@ -1728,8 +1728,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([(10,10,10)], lambda x: x.log_softmax(0), atol=1e-7, grad_atol=1e-7)
    helper_test_op([(10,10,10)], lambda x: x.log_softmax(1), atol=1e-7, grad_atol=1e-7)
    helper_test_op([(10,10,10)], lambda x: x.log_softmax(2), atol=1e-7, grad_atol=1e-7)
-  def test_softmin(self):
-    helper_test_op([(45,65)], torch.nn.Softmin(dim=1), Tensor.softmin, atol=1e-7, grad_atol=1e-7)

  def test_normalize(self):
    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x), lambda x: x.normalize(), atol=1e-7, grad_atol=1e-7)
@@ -143,11 +143,12 @@ class TestOptim(unittest.TestCase):

  @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
  def test_mixed_precision(self):
-    self.enterContext(Context(DEFAULT_FLOAT=dtypes.half))
+    old_default_float, dtypes.default_float = dtypes.default_float, dtypes.half
    # weight update would overflow without upcasting
    self._test_sgd(10, {'lr': 1e10}, 1e-6, 3e-4)
    self._test_adam(1, {'lr': 1e10}, 1e-4, 1e-4)
    self._test_adamw(1, {'lr': 1e10}, 1e-4, 1e-4)
+    dtypes.default_float = old_default_float

  def test_assert_tensor_train(self):
    t = Tensor.ones((1,1))
@@ -13,7 +13,7 @@ class TestPickle(unittest.TestCase):

  def test_pickle_pattern_matcher(self):
    pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)])
-    sink = UOp.const(2)
+    sink = UOp.const(dtypes.int, 2)
    tt = pm.rewrite(sink)
    pm_str = pickle.dumps(pm)
    pm2 = pickle.loads(pm_str)
@@ -1,7 +1,7 @@
 import unittest, math

 from tinygrad import dtypes, Tensor, Device
-from tinygrad.helpers import getenv, DEV, Context
+from tinygrad.helpers import getenv, DEV
 from tinygrad.codegen import to_program

 from tinygrad.uop.ops import Ops
@@ -232,14 +232,16 @@ class TestRandomness(unittest.TestCase):
  @given(strat.sampled_from([dtypes.float, dtypes.float16, dtypes.bfloat16]))
  def test_randn_finite(self, default_float):
    if default_float not in Device[Device.DEFAULT].renderer.supported_dtypes(): return
+    old_default_float = dtypes.default_float
    # low precision can result in inf from randn
-    self.enterContext(Context(DEFAULT_FLOAT=default_float))
+    dtypes.default_float = default_float
    t = Tensor.randn(64, 64)
    mx = t.max().numpy().item()
    mn = t.min().numpy().item()
    print(f"testing with {default_float=}")
    assert math.isfinite(mx), mx
    assert math.isfinite(mn), mn
+    dtypes.default_float = old_default_float

  def test_random_counter_overflow(self):
    device = Device.DEFAULT
@@ -24,7 +24,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
  dtype = alu_src_uops[0].dtype
  a = UOp.param(0, dtype, (1,))
  b = UOp.param(1, dtype, (1,))
-  idx = UOp.const(0)
+  idx = UOp.const(dtypes.int, 0)
  ld = b.index(idx).load()
  alu = ld.alu(alu_op, *alu_src_uops)
  store = UOp.store(a.index(idx), alu)
@@ -35,7 +35,7 @@ class TestRendererFailures(unittest.TestCase):
  def test_gated_store_with_alu(self):
    a = UOp.param(0, dtypes.int, (4,))
    gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
-    gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(1).cast(dtypes.int)))
+    gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
    sink = UOp(Ops.SINK, src=(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])
@@ -45,7 +45,7 @@ class TestRendererFailures(unittest.TestCase):
    a = UOp.param(0, dtypes.int, (8,))
    gate_alu_0 = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
    gate_alu_1 = (lidx1:=UOp.special(2, 'lidx1')).ne(0)
-    gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(1).cast(dtypes.int)))
+    gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
    sink = UOp(Ops.SINK, src=(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])
@@ -54,7 +54,7 @@ class TestRendererFailures(unittest.TestCase):
 class TestCStyleFailures(unittest.TestCase):
  def test_inline_const_alu(self):
    # CPU doesn't use the max function
-    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int.min+1).cast(dtypes.int))
+    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.int.min+1))
    self.assertEqual(ret[0], 1)

  def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True):
@@ -80,7 +80,7 @@ class TestWGSLFailures(unittest.TestCase):
  def test_multiply_infinity(self):
    # multiplying a positive constant by infinity should return infinity
    # WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer
-    ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(float("inf")).cast(dtypes.float32))
+    ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(dtypes.float32, float("inf")))
    self.assertEqual(ret[0], float("inf"))

  # WGSL has a specific select(alt, val, gate) ternary operator instead of gate?val:alt
@@ -104,7 +104,7 @@ class TestPTXFailures(unittest.TestCase):
  def test_gated_store_with_if(self):
    a = UOp.param(0, dtypes.int, (4,))
    gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
-    val = UOp.const(1).cast(dtypes.int)
+    val = UOp.const(dtypes.int, 1)
    if_uop = UOp(Ops.IF, src=(gate_alu,))
    gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0, if_uop), val))
    sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
@@ -363,7 +363,7 @@ class TestCopyFolding(unittest.TestCase):

  def test_one_hot_with_copy(self):
    y = Tensor([1, 2, 3]).to("CPU")
-    x = y.one_hot(10).int()
+    x = y.one_hot(10)
    check_schedule(x, 3, filter_sink=False)

  @unittest.skip("no longer supported")
@@ -24,13 +24,13 @@ class TestTinygrad(unittest.TestCase):
    self.assertEqual(Tensor(3.14).shape, ())

  def test_deviceless_const_construct_device_repr(self):
-    t = Tensor(UOp.const(2.0).cast(dtypes.float))
+    t = Tensor(UOp.const(dtypes.float, 2.0))
    self.assertIsNone(t.uop.device)
    self.assertIsNone(t.device)
    self.assertIn("<UOp None", repr(t))

  def test_deviceless_const_realize_noop(self):
-    t = Tensor(UOp.const(2.0).cast(dtypes.float))
+    t = Tensor(UOp.const(dtypes.float, 2.0))
    uop = t.uop
    t.realize()
    self.assertIs(t.uop, uop)
@@ -728,12 +728,12 @@ class TestZeroShapeTensor(unittest.TestCase):
    self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)

  def test_clone_deviceless_const(self):
-    t = Tensor(UOp.const(2.0).cast(dtypes.float)).clone()
+    t = Tensor(UOp.const(dtypes.float, 2.0)).clone()
    np.testing.assert_equal(t.numpy(), 2.0)
    self.assertTrue(t.uop.has_buffer_identity())

  def test_numpy_deviceless_const(self):
-    np.testing.assert_equal(Tensor(UOp.const(2.0).cast(dtypes.float)).numpy(), 2.0)
+    np.testing.assert_equal(Tensor(UOp.const(dtypes.float, 2.0)).numpy(), 2.0)

  def test_clone_with_shrink(self):
    a = Tensor.rand(16, 16)
@@ -756,7 +756,7 @@ class TestZeroShapeTensor(unittest.TestCase):
    np.testing.assert_allclose(a.grad.numpy(), b.grad.numpy())

  def test_clone_deviceless_const_to_cpu(self):
-    t = Tensor(UOp.const(2.0).cast(dtypes.float)).clone(device="CPU")
+    t = Tensor(UOp.const(dtypes.float, 2.0)).clone(device="CPU")
    self.assertEqual(t.device, "CPU")
    np.testing.assert_equal(t.numpy(), 2.0)

@@ -1,6 +1,6 @@
 import unittest
 import numpy as np
-from tinygrad import Device, Tensor, Variable, TinyJit, dtypes
+from tinygrad import Tensor, Variable, dtypes
 from tinygrad.helpers import CHECK_OOB

 class TestTensorVariable(unittest.TestCase):
@@ -18,18 +18,10 @@ class TestTensorVariable(unittest.TestCase):
      self.assertListEqual((vv * t).tolist(), [2, 2, 2])
    except RuntimeError: pass

-  @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires long support")
+  # TODO: a Variable PARAM lowers to int32, so a bound value that doesn't fit int32 truncates or fails to bind
+  @unittest.expectedFailure
  def test_large_range_variable(self):
-    self.assertEqual(Tensor(Variable("b", 0, 2**40, dtype=dtypes.long).bind(2**35)).clone(Device.DEFAULT).item(), 2**35)
-
-  @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires long support")
-  def test_large_range_variable_jit(self):
-    @TinyJit
-    def f(a,b): return (Tensor(a+b).clone(Device.DEFAULT) * 2).realize()
-    for i in range(3):
-      a = Variable("a", 0, 2**10, dtype=dtypes.int).bind(i)
-      b = Variable("b", 0, 2**40, dtype=dtypes.long).bind(2**35)
-      self.assertEqual(f(a,b).item(), (2**35 + i) * 2)
+    self.assertEqual(Tensor(Variable("b", 0, 2**40).bind(2**35)).item(), 2**35)

  def test_variable_defers_like_a_literal(self):
    vv = Variable("a", 1, 10).bind(2)
@@ -19,7 +19,7 @@ def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
  run_linear(UOp(Ops.LINEAR, src=(UOp.sink(*uops_list, arg=KernelInfo()).call(*buf_uops),)))

 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(arg).cast(dtype))
+  if op is Ops.CONST: uops.append(UOp.const(dtype, arg))
  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,)))
  else: uops.append(UOp(op, dtype, tuple(src), arg))
  return uops[-1]
@@ -43,7 +43,7 @@ def _test_single_value_const(vals, op, dts):
  buf_store = uop(uops, Ops.PARAM, output_dtype, (), 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(0).cast(dtypes.int32)].store(alu)
+  out = buf_store[UOp.const(dtypes.int32, 0)].store(alu)
  buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
  run_uops([out], [buf])
  return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
@@ -221,12 +221,12 @@ class TestAssembly(unittest.TestCase):
  def test_bitshift_left(self):
    g1 = UOp.param(0, dtypes.int32, shape=(3,))
    out = UOp.param(1, dtypes.int32, shape=(2,))
-    c1 = UOp.const(2)
-    c2 = UOp.const(3)
+    c1 = UOp.const(dtypes.int, 2)
+    c2 = UOp.const(dtypes.int, 3)
    l1 = g1.index(c1)
    a1 = UOp(Ops.MUL, src=(l1, c1))
    a2 = UOp(Ops.MUL, src=(l1, c2))
-    uops = to_uops_list([out.index(UOp.const(0)).store(a1), out.index(UOp.const(1)).store(a2)],
+    uops = to_uops_list([out.index(UOp.const(dtypes.int, 0)).store(a1), out.index(UOp.const(dtypes.int, 1)).store(a2)],
                        ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    ops = [x.op for x in uops]
@@ -249,16 +249,16 @@ class TestAssembly(unittest.TestCase):

  def test_mulacc_shl(self):
    g1 = UOp.param(0, dtypes.int32, shape=(2,))
-    c1 = UOp.const(0)
-    c2 = UOp.const(1)
-    expr = g1.index(c1) * UOp.const(4096) + g1.index(c2)
+    c1 = UOp.const(dtypes.int, 0)
+    c2 = UOp.const(dtypes.int, 1)
+    expr = g1.index(c1) * UOp.const(dtypes.int, 4096) + g1.index(c2)
    uops = to_uops_list([expr], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    self.assertIn(Ops.MULACC, [x.op for x in uops])

  def test_use_cmpeq(self):
    g = UOp.param(0, dtypes.uint32, shape=(8,))
-    c = UOp.const(7)
+    c = UOp.const(dtypes.uint, 7)
    comp = g.index(c).ne(c).ne(True)
    uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
@@ -1,8 +1,7 @@
 from tinygrad import Tensor, dtypes
+dtypes.default_float = dtypes.float16
 from tinygrad.dtype import to_dtype
-from tinygrad.helpers import getenv, Context
-
-Context(DEFAULT_FLOAT=dtypes.float16).__enter__()
+from tinygrad.helpers import getenv

 if __name__ == "__main__":
  # matmuls in bert layers
@@ -1,10 +1,9 @@
 from tinygrad import Tensor, dtypes, GlobalCounters
+dtypes.default_float = dtypes.float16
 from tinygrad.dtype import to_dtype
-from tinygrad.helpers import getenv, Context
+from tinygrad.helpers import getenv
 from test.backend.test_softmax_fusion import single_kernel_softmax

-Context(DEFAULT_FLOAT=dtypes.float16).__enter__()
-
 if __name__ == "__main__":
  # softmax in bert layers
  BS = getenv("BS", 96//6)
@@ -13,9 +13,9 @@ from tinygrad.dtype import Invalid

 def vision_conv_143():
  c0 = UOp.param(0, dtypes.half, shape=(16, 1024, 4))
-  c2 = UOp.range(32, 3)
-  c5 = UOp.range(128, 4)
-  c8 = UOp.range(16, 2)
+  c2 = UOp.range(32, 3, AxisType.LOOP)
+  c5 = UOp.range(128, 4, AxisType.LOOP)
+  c8 = UOp.range(16, 2, AxisType.LOOP)
  c16 = UOp.range(7, 0, AxisType.REDUCE)
  c17 = c8*2+c16
  c24 = ((c17<3)!=True)&(c17<35)
@@ -24,8 +24,8 @@ def vision_conv_143():
  c32 = ((c27<3)!=True)&(c27<67)
  c34 = UOp.param(1, dtypes.half, shape=(32, 1024, 4))
  c38 = c5//2
-  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(Invalid, dtypes.weakint))
-  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
+  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.param(2, dtypes.half, shape=(64, 49, 4))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
  c63 = UOp.param(3, dtypes.float, (128,))
@@ -39,9 +39,9 @@ def vision_conv_143():

 def vision_conv_153():
  c0 = UOp.param(0, dtypes.half, shape=(8, 1024, 4))
-  c2 = UOp.range(16, 3)
-  c5 = UOp.range(256, 4)
-  c8 = UOp.range(8, 2)
+  c2 = UOp.range(16, 3, AxisType.LOOP)
+  c5 = UOp.range(256, 4, AxisType.LOOP)
+  c8 = UOp.range(8, 2, AxisType.LOOP)
  c16 = UOp.range(7, 0, AxisType.REDUCE)
  c17 = c8*2+c16
  c24 = ((c17<3)!=True)&(c17<19)
@@ -50,8 +50,8 @@ def vision_conv_153():
  c32 = ((c27<3)!=True)&(c27<35)
  c34 = UOp.param(1, dtypes.half, shape=(16, 1024, 4))
  c38 = c5//2
-  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(Invalid, dtypes.weakint))
-  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
+  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.param(2, dtypes.half, shape=(128, 49, 4))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
  c63 = UOp.param(3, dtypes.float, (256,))
@@ -65,7 +65,7 @@ def vision_conv_153():

 def dm_conv_172():
  c0 = UOp.param(0, dtypes.half, shape=(1, 240, 4))
-  c2 = UOp.range(960, 4)
+  c2 = UOp.range(960, 4, AxisType.LOOP)
  c5 = UOp.param(1, dtypes.half, shape=(8, 384, 4))
  c7 = UOp.range(32, 0, AxisType.REDUCE)
  c10 = UOp.range(4, 1, AxisType.REDUCE)
@@ -5,9 +5,8 @@ These tests intentionally cause GPU faults to verify error handling.
 Run with: DEV=AMD python -m pytest test/external/external_test_gpu_crash.py -v
 """
 import unittest, re, importlib
-from tinygrad import Device, UOp, dtypes
-from tinygrad.engine.realize import run_linear
-from tinygrad.uop.ops import Ops, KernelInfo
+from tinygrad.device import Device, TinyELF
+from tinygrad.helpers import Target
 from tinygrad.renderer.amd.dsl import s, v, Inst, NULL

 RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "s_waitcnt": "s_waitcnt", "s_endpgm": "s_endpgm",
@@ -15,11 +14,20 @@ RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "
                   "global_atomic_add_u32": "global_atomic_add", "flat_load_b32": "flat_load_dword",
                   "flat_store_b32": "flat_store_dword", "flat_atomic_add_u32": "flat_atomic_add", "s_load_b32": "s_load_dword"}

+def assemble(code:str, name:str="test", is_cdna:bool=False) -> str:
+  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
+  if is_cdna: kd["accum_offset"] = 8
+  else: kd["wavefront_size32"] = 1
+  return f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{code}\n.rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + \
+         "\n".join(f".amdhsa_{k} {v}" for k,v in kd.items()) + "\n.end_amdhsa_kernel"
+
@unittest.skipIf(Device.DEFAULT != "AMD", "AMD required")
 class TestGPUCrash(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
    cls.dev = Device["AMD"]
+    cls.compiler = HIPCompiler(cls.dev.arch)
    cls.is_cdna = cls.dev.target[0] < 10
    ins = importlib.import_module('tinygrad.runtime.autogen.amd.' + ('cdna' if cls.is_cdna else 'rdna3') + '.ins')
    for rdna3_name, cdna3_name in RDNA3_CDNA3_MAP.items():
@@ -34,11 +42,14 @@ class TestGPUCrash(unittest.TestCase):
    except Exception:
      self.fail("Device not working before test")

+  def _run(self, code: str):
+    prg = self.dev.runtime(TinyELF(self.compiler.compile(assemble(code, is_cdna=self.is_cdna)), "test",
+                                   Target("AMD", arch=self.dev.arch), ()))
+    prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)
+
  def _run_insts(self, insts: list[Inst]):
-    buf = UOp.new_buffer("AMD", 64, dtypes.uint8)
-    sink = UOp.sink(UOp.param(0, dtypes.uint8, (64,), device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
-    prg = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=i) for i in insts))))
-    run_linear(UOp(Ops.LINEAR, src=(prg.call(buf),)), wait=True)
+    from test.amd.disasm import disasm
+    self._run("\n".join(disasm(i) for i in insts))

  def _assert_gpu_fault(self, func):
    """Assert that func raises a RuntimeError indicating a GPU fault (not a setup error)."""
@@ -3,7 +3,6 @@ import unittest
 import numpy as np
 from tinygrad import Tensor, dtypes
 from tinygrad.engine.jit import TinyJit
-from tinygrad.helpers import Context
 from test.helpers import derandomize_model

 from examples.llama import Transformer
@@ -15,7 +14,8 @@ def helper_test_jitted_correctness(gen, train, train_jit):

 class TestJittedModels(unittest.TestCase):
  def test_jitted_tiny_llama(self):
-    self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
+    old_float = dtypes.default_float
+    dtypes.default_float = dtypes.float16

    args_tiny = {"dim": 1024, "hidden_dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000}
    model = Transformer(**args_tiny)
@@ -25,6 +25,7 @@ class TestJittedModels(unittest.TestCase):
    @TinyJit
    def test_jit(t): return model(t, 0).realize()
    helper_test_jitted_correctness(lambda: (Tensor([[1,]]),), test, test_jit)
+    dtypes.default_float = old_float

  def test_jitted_stable_diffusion(self):
    from examples.stable_diffusion import UNetModel, unet_params
@@ -40,7 +40,7 @@ def random_int_expr(depth=10):
 def random_bool_expr(depth=10, expr1=None):
  if depth == 0: return True
  if expr1 is None: expr1 = random_int_expr(depth-1)
-  expr2 = random.choice([random_or_sub_expression_int(depth-1, expr1), UOp.const(random.randint(-10, 10))])
+  expr2 = random.choice([random_or_sub_expression_int(depth-1, expr1), UOp.const(dtypes.weakint, random.randint(-10, 10))])
  return random.choice(comp_ops)(expr1, expr2)


@@ -82,8 +82,8 @@ if __name__ == "__main__":
            f"v2=Variable(\"{u2.arg[0]}\", {u2.arg[1]}, {u2.arg[2]})\n" +\
            f"v3=Variable(\"{u3.arg[0]}\", {u3.arg[1]}, {u3.arg[2]})\n" +\
            f"expr = {expr}\n" +\
-            f"v1_val, v2_val, v3_val = UOp.const({n1.as_long()}), UOp.const({n2.as_long()})," +\
-                f"UOp.const({n3.as_long()})\n" +\
+            f"v1_val, v2_val, v3_val = UOp.const(dtypes.weakint, {n1.as_long()}), UOp.const(dtypes.weakint, {n2.as_long()})," +\
+                f"UOp.const(dtypes.weakint, {n3.as_long()})\n" +\
            "num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\
            "rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\
            "assert num==rn, f\"{num} != {rn}\"\n"
@@ -51,7 +51,7 @@ class TestEval(unittest.TestCase):
      vae = AutoencoderKL()
      for p in get_parameters(vae): p.to_(GPUS).realize()
      x = Tensor.zeros(48,4,64,64).contiguous().to(GPUS).realize()
-      x.uop = x.uop.unshard(0)
+      x.uop = x.uop.multi(0)
      for _ in range(2): vae_decode(x, vae)

 if __name__=="__main__":
@@ -90,7 +90,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
    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.param(0, uop.dtype, (1,))
-  prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
+  prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
  prog = dev.runtime(prg.to_elf())
  prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)
  return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
@@ -131,8 +131,7 @@ class PM4Executor(AMDQueue):
    _ = self._next_dword() # ev

    ptr = to_mv(self.gpu.translate_addr(addr_lo + (addr_hi << 32)), 8)
-    if mem_data_sel == 1: ptr.cast('I')[0] = val & 0xffffffff
-    elif mem_data_sel == 2: ptr.cast('Q')[0] = val
+    if mem_data_sel == 1 or mem_data_sel == 2: ptr.cast('Q')[0] = val
    elif mem_data_sel == 3:
      if mem_event_type == CACHE_FLUSH_AND_INV_TS_EVENT: ptr.cast('Q')[0] = int(time.perf_counter() * 1e8)
      else: raise RuntimeError(f"Unknown {mem_data_sel=} {mem_event_type=}")
@@ -51,7 +51,7 @@ class _MXCSRContext:
    if lib is None or not hasattr(self, '_saved'): return
    lib.set_fpcr(self._saved)

-from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.device import Buffer, BufferSpec, Device
 from tinygrad.runtime.autogen import hsa
@@ -192,16 +192,16 @@ def _init_sqtt_encoder():

  return emit, finish, finalize

-def _c(val, dtype=dtypes.uint32): return UOp.const(val, dtype)
+def _c(val, dtype=dtypes.uint32): return UOp.const(dtype, val)

 def _u64(lo: UOp, hi: UOp) -> UOp:
  """Combine two 32-bit UOps into a 64-bit UOp."""
-  return lo.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
+  return lo.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))

 def _split64(val: UOp) -> tuple[UOp, UOp]:
  """Split a 64-bit value into (lo, hi) 32-bit values."""
  v64 = val.bitcast(dtypes.uint64) if val.dtype == dtypes.float64 else val.cast(dtypes.uint64) if val.dtype != dtypes.uint64 else val
-  return v64.cast(dtypes.uint32), (v64 >> UOp.const(32, dtypes.uint64)).cast(dtypes.uint32)
+  return v64.cast(dtypes.uint32), (v64 >> UOp.const(dtypes.uint64, 32)).cast(dtypes.uint32)

 _SRC_MOD_TYPES = {16: (dtypes.uint16, dtypes.half, 0x7FFF), 32: (dtypes.uint32, dtypes.float32, 0x7FFFFFFF),
                  64: (dtypes.uint64, dtypes.float64, 0x7FFFFFFFFFFFFFFF)}
@@ -210,7 +210,7 @@ def _apply_src_mods(val: UOp, mod_bit: int, abs_bits: int, neg_bits: int, bits:
  if not (abs_bits & (1 << mod_bit)) and not (neg_bits & (1 << mod_bit)): return val
  ut, ft, mask = _SRC_MOD_TYPES[bits]
  fv = val.cast(ut).bitcast(ft) if bits == 16 else val.bitcast(ft) if val.dtype == ut else val
-  if abs_bits & (1 << mod_bit): fv = (fv.bitcast(ut) & UOp.const(mask, ut)).bitcast(ft)
+  if abs_bits & (1 << mod_bit): fv = (fv.bitcast(ut) & UOp.const(ut, mask)).bitcast(ft)
  if neg_bits & (1 << mod_bit): fv = fv.neg()
  return fv.bitcast(ut).cast(dtypes.uint32) if bits == 16 else fv.bitcast(ut)

@@ -251,7 +251,7 @@ def _to_u32(val: UOp) -> UOp:
  if val.dtype.itemsize == 4: return val.bitcast(dtypes.uint32)  # same size: bitcast (float32->uint32)
  return val.cast(dtypes.uint32)  # different size: cast (bool, int16, etc)
 def _lane_active(exec_mask: UOp, lane: UOp) -> UOp:
-  if exec_mask.dtype == dtypes.uint64: return ((exec_mask >> lane.cast(dtypes.uint64)) & UOp.const(1, dtypes.uint64)).ne(UOp.const(0, dtypes.uint64))
+  if exec_mask.dtype == dtypes.uint64: return ((exec_mask >> lane.cast(dtypes.uint64)) & UOp.const(dtypes.uint64, 1)).ne(UOp.const(dtypes.uint64, 0))
  return ((exec_mask >> lane.cast(dtypes.uint32)) & _c(1)).ne(_c(0))
 def _hi16(v: UOp) -> UOp: return (v >> _c(16)) & _c(0xFFFF)
 def _cond(cond, if_true, if_false):
@@ -264,9 +264,9 @@ def _set_lane_bit(old: UOp, lane: UOp, val: UOp, exec_mask: UOp) -> UOp:
  """Set/clear a single bit in a mask based on lane index, respecting exec mask."""
  if old.dtype in (dtypes.uint64, dtypes.int64):
    dt = dtypes.uint64
-    mask = UOp.const(1, dt) << lane.cast(dt)
+    mask = UOp.const(dt, 1) << lane.cast(dt)
    new_bit = _to_u32(val).cast(dt) << lane.cast(dt)
-    cleared = old.cast(dt) & (mask ^ UOp.const(0xFFFFFFFFFFFFFFFF, dt))
+    cleared = old.cast(dt) & (mask ^ UOp.const(dt, 0xFFFFFFFFFFFFFFFF))
    return _lane_active(exec_mask, lane).where(cleared | new_bit, old.cast(dt))
  mask = _c(1) << lane.cast(dtypes.uint32)
  new_bit = _to_u32(val) << lane.cast(dtypes.uint32)
@@ -365,7 +365,7 @@ def _write_64bit(val: UOp, wfn, reg_or_addr, is_mem: bool, *args) -> list[UOp]:
  """Write a 64-bit value as two 32-bit writes. args passed to wfn after reg/addr and lo/hi value."""
  lo, hi = _split64(val)
  incr = 4 if is_mem else 1  # 4 bytes for memory addresses, 1 for register indices
-  return [wfn(reg_or_addr, lo, *args), wfn(reg_or_addr + (UOp.const(incr, reg_or_addr.dtype) if isinstance(reg_or_addr, UOp) else incr), hi, *args)]
+  return [wfn(reg_or_addr, lo, *args), wfn(reg_or_addr + (UOp.const(reg_or_addr.dtype, incr) if isinstance(reg_or_addr, UOp) else incr), hi, *args)]

 def _write_val(bits: int, val: UOp, wfn, reg_or_addr, *args, is_mem: bool = False) -> list[UOp]:
  """Write value, splitting 64-bit if needed. bits=64 for 64-bit writes, otherwise 32-bit."""
@@ -374,7 +374,7 @@ def _write_val(bits: int, val: UOp, wfn, reg_or_addr, *args, is_mem: bool = Fals
 def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32, data_bits: int = 32) -> list[UOp]:
  """Conditional memory store with sub-word support. Returns list of store UOps."""
  adt = dtypes.uint64 if addr_bits == 64 else dtypes.uint32
-  word_addr = addr >> UOp.const(2, adt)
+  word_addr = addr >> UOp.const(adt, 2)
  idx = mem.index(word_addr.valid(active))
  if data_bits == 32: return [idx.store(active.where(_to_u32(val), idx))]
  # Sub-word store: read-modify-write with mask
@@ -388,7 +388,7 @@ def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32,
  is_cross = byte_pos.eq(_c(3))
  cross_word0 = (idx & _c(0x00FFFFFF)) | ((val_u32 & _c(0xFF)) << _c(24))
  store0 = idx.store(active.where(is_cross.where(cross_word0, new_word), idx))
-  next_idx = mem.index((word_addr + UOp.const(1, adt)).valid(active & is_cross))
+  next_idx = mem.index((word_addr + UOp.const(adt, 1)).valid(active & is_cross))
  cross_word1 = (next_idx & _c(0xFFFFFF00)) | ((val_u32 >> _c(8)) & _c(0xFF))
  return [store0, next_idx.store((active & is_cross).where(cross_word1, next_idx))]

@@ -397,8 +397,8 @@ def _mem_store_bytes(mem: UOp, addr: UOp, val: UOp, active: UOp, data_bits: int
  stores = []
  val_u32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
  for i in range(data_bits // 8):
-    byte_val = (val_u32 >> UOp.const(i * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)
-    stores.append(mem.index((addr + UOp.const(i, dtypes.uint64)).valid(active)).store(byte_val.cast(dtypes.uint8)))
+    byte_val = (val_u32 >> UOp.const(dtypes.uint32, i * 8)) & UOp.const(dtypes.uint32, 0xFF)
+    stores.append(mem.index((addr + UOp.const(dtypes.uint64, i)).valid(active)).store(byte_val.cast(dtypes.uint8)))
  return stores

 def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode_vars: dict | None = None, op_name: str = "") -> dict[int, UOp]:
@@ -446,7 +446,7 @@ class _Ctx:
    """Create a lane range UOp with unique axis ID."""
    if n is None: n = self.wave_size
    self._axis_id += 1
-    return UOp.range(n, self._axis_id, dtype=dtypes.int)
+    return UOp.range(n, self._axis_id, AxisType.LOOP, dtype=dtypes.int)

  def unroll_lanes(self, get_lane_bit, exec_mask: UOp, apply_exec: bool = True) -> UOp:
    """Combine lane bits into a mask using RANGE+REDUCE (32-bit for RDNA, 64-bit for CDNA)."""
@@ -462,8 +462,8 @@ class _Ctx:
  def inst_word(self, dword_idx: int) -> UOp:
    """Read instruction dword from vmem at PC + dword_idx*4."""
    pc = self.rpc()
-    addr = pc if dword_idx == 0 else pc + UOp.const(dword_idx * 4, dtypes.uint64)
-    return self.vmem.index(addr >> UOp.const(2, dtypes.uint64)).load()
+    addr = pc if dword_idx == 0 else pc + UOp.const(dtypes.uint64, dword_idx * 4)
+    return self.vmem.index(addr >> UOp.const(dtypes.uint64, 2)).load()

  def inst_field(self, field) -> UOp:
    """Extract field bits from instruction encoding. Tracks field for canonical key computation."""
@@ -475,15 +475,15 @@ class _Ctx:
    word = self.inst_word(dword_idx)
    if lo // 32 == hi // 32:  # Same dword
      mask = (1 << (hi - lo + 1)) - 1
-      shifted = word if lo_in_dword == 0 else word >> UOp.const(lo_in_dword, dtypes.uint32)
-      return shifted & UOp.const(mask, dtypes.uint32)
+      shifted = word if lo_in_dword == 0 else word >> UOp.const(dtypes.uint32, lo_in_dword)
+      return shifted & UOp.const(dtypes.uint32, mask)
    else:  # Spans two dwords
      lo_bits = 32 - lo_in_dword
      lo_mask = (1 << lo_bits) - 1
      hi_mask = (1 << (hi_in_dword + 1)) - 1
-      lo_part = (word >> UOp.const(lo_in_dword, dtypes.uint32)) & UOp.const(lo_mask, dtypes.uint32)
-      hi_part = self.inst_word(dword_idx + 1) & UOp.const(hi_mask, dtypes.uint32)
-      return lo_part | (hi_part << UOp.const(lo_bits, dtypes.uint32))
+      lo_part = (word >> UOp.const(dtypes.uint32, lo_in_dword)) & UOp.const(dtypes.uint32, lo_mask)
+      hi_part = self.inst_word(dword_idx + 1) & UOp.const(dtypes.uint32, hi_mask)
+      return lo_part | (hi_part << UOp.const(dtypes.uint32, lo_bits))

  def inst_field_signed(self, field) -> UOp:
    """Extract field and sign-extend based on field width."""
@@ -584,7 +584,7 @@ class _Ctx:
      inline = is_float_const.where(float_inline.bitcast(dtypes.uint64), int_inline.bitcast(dtypes.uint64))
      # Literal handling: F64 VOP puts literal in high 32 bits; B64/I64/U64 VOP and SOP zero-extend
      if literal is not None:
-        lit_val = literal.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64) if is_f64 else literal.cast(dtypes.uint64)
+        lit_val = literal.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32) if is_f64 else literal.cast(dtypes.uint64)
        inline = off.eq(_c(255)).where(lit_val, inline)
      scalar_val = (off < _c(128)).where(sgpr_val, inline)
    else:
@@ -602,7 +602,7 @@ class _Ctx:

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

@@ -668,7 +668,7 @@ class _Ctx:
    if 'VCC' not in srcs: srcs['VCC'] = self.rmask(_c(vcc_reg))
    srcs.update({'EXEC': exec_mask, 'SCC': self.rsgpr_dyn(_c(SCC.offset)), 'laneId': lane, 'VDST': vdst_reg,
                  'ROUND_MODE': _c(0), 'ROUND_TOWARD_ZERO': _c(0), 'ROUND_NEAREST_EVEN': _c(0), '_vgpr': self.vgpr, '_wave_size': self.wave_size,
-                  'MAX_FLOAT_F32': UOp.const(3.4028234663852886e38, dtypes.float32),
+                  'MAX_FLOAT_F32': UOp.const(dtypes.float32, 3.4028234663852886e38),
                  # CDNA SDWA byte/word select constants (E32 always uses BYTE0/WORD0 defaults)
                  'SDWA_SRC0_SEL': _c(0), 'BYTE0': _c(0), 'BYTE1': _c(1), 'BYTE2': _c(2), 'BYTE3': _c(3),
                  'WORD0': _c(0), 'WORD1': _c(1)})  # rounding mode and SDWA constants
@@ -709,7 +709,7 @@ class _Ctx:
      # VGPR bit-slice assignment: VGPR[lane][reg][hi:lo] = (vgpr_idx, rhs_val, hi, lo[, cond]) -> read-modify-write
      if dest.startswith('VGPR[') and re.search(r'\[\d+:\d+\]', dest):
        # VGPR bit-slice: (vgpr_idx, rhs_val, hi_bit, lo_bit) - hi/lo are UOp constants
-        hi_bit, lo_bit = int(val[2].val), int(val[3].val)
+        hi_bit, lo_bit = int(val[2].arg), int(val[3].arg)
        width = hi_bit - lo_bit + 1
        old = self.vgpr.index(val[0]).load()
        new_val = _set_bits(old, _val_to_bits(val[1]), width, lo_bit).cast(dtypes.uint32)
@@ -732,24 +732,24 @@ class _Ctx:
            d0_width, slice_mask = d0_hi_bit - d0_lo_bit + 1, (1 << (d0_hi_bit - d0_lo_bit + 1)) - 1
            val_bits = val.bitcast(dtypes.uint16).cast(dtypes.uint32) if val.dtype == dtypes.half else \
                       val.cast(dtypes.uint32) if val.dtype in (dtypes.uint16, dtypes.int16) else \
-                       val.cast(dtypes.uint32) & UOp.const(slice_mask, dtypes.uint32)
+                       val.cast(dtypes.uint32) & UOp.const(dtypes.uint32, slice_mask)
            raw_stores.append(('vgpr_slice', (d0_lo_bit, d0_width, val_bits)))
            continue
        # For integer ops with clamp, use pre-computed saturated value; for floats, clamp to [0,1]
        if int_saturate is not None: val = int_saturate
        elif clmp and val.dtype in (dtypes.float32, dtypes.half, dtypes.float64):
-          clamped = val.maximum(UOp.const(0.0, val.dtype)).minimum(UOp.const(1.0, val.dtype))
-          val = _FUNCS['isNAN'](val).where(UOp.const(0.0, val.dtype), clamped)
+          clamped = val.maximum(UOp.const(val.dtype, 0.0)).minimum(UOp.const(val.dtype, 1.0))
+          val = _FUNCS['isNAN'](val).where(UOp.const(val.dtype, 0.0), clamped)
        if val.dtype in (dtypes.uint64, dtypes.int64, dtypes.float64):
          lo, hi = _split64(val)
          raw_stores.extend([('vgpr', self.wvgpr_dyn(vdst_reg, lane, lo, exec_mask)),
                             ('vgpr', self.wvgpr_dyn(vdst_reg + _c(1), lane, hi, exec_mask))])
        elif val.dtype in (dtypes.half, dtypes.uint16, dtypes.int16):
          result, old_val = _val_to_u32(val), self.rvgpr_dyn(vdst_reg, lane)
-          hi_result = (old_val & UOp.const(0xFFFF, dtypes.uint32)) | (result << UOp.const(16, dtypes.uint32))
+          hi_result = (old_val & UOp.const(dtypes.uint32, 0xFFFF)) | (result << UOp.const(dtypes.uint32, 16))
          # GFX9/CDNA zeroes upper 16 bits on lo-half write; RDNA preserves them
-          lo_result = (result & UOp.const(0xFFFF, dtypes.uint32)) if self.wave_size == 64 else \
-                      (old_val & UOp.const(0xFFFF0000, dtypes.uint32)) | (result & UOp.const(0xFFFF, dtypes.uint32))
+          lo_result = (result & UOp.const(dtypes.uint32, 0xFFFF)) if self.wave_size == 64 else \
+                      (old_val & UOp.const(dtypes.uint32, 0xFFFF0000)) | (result & UOp.const(dtypes.uint32, 0xFFFF))
          result = opsel_dst_hi.where(hi_result, lo_result) if isinstance(opsel_dst_hi, UOp) else hi_result if opsel_dst_hi else lo_result
          raw_stores.append(('vgpr', self.wvgpr_dyn(vdst_reg, lane, result, exec_mask)))
        else: raw_stores.append(('vgpr', self.wvgpr_dyn(vdst_reg, lane, _val_to_u32(val), exec_mask)))
@@ -767,8 +767,8 @@ class _Ctx:
    if slice_stores:
      result = self.rvgpr_dyn(vdst_reg, lane)
      for lo_bit, width, val_bits in slice_stores:
-        mask = UOp.const(((1 << width) - 1) << lo_bit, dtypes.uint32)
-        result = (result & (mask ^ UOp.const(0xFFFFFFFF, dtypes.uint32))) | (val_bits << UOp.const(lo_bit, dtypes.uint32))
+        mask = UOp.const(dtypes.uint32, ((1 << width) - 1) << lo_bit)
+        result = (result & (mask ^ UOp.const(dtypes.uint32, 0xFFFFFFFF))) | (val_bits << UOp.const(dtypes.uint32, lo_bit))
      lane_stores.append(self.wvgpr_dyn(vdst_reg, lane, result, exec_mask))
    # VCC/EXEC mask writes must be computed BEFORE VGPR stores to avoid reading modified VGPRs.
    # When vdst overlaps with src operands (e.g. v_add_co_u32 v[0], vcc, s[8], v[0]), the carry
@@ -790,8 +790,8 @@ class _Ctx:
 def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
  simm16 = ctx.inst_field_signed(type(inst).simm16).cast(dtypes.int16)
  if inst.op in (ir3.SOPPOp.S_ENDPGM, ir4.SOPPOp.S_ENDPGM, irc.SOPPOp.S_ENDPGM):
-    return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), UOp.const(0xFFFFFFFF, dtypes.uint32)),
-                          ctx.wsgpr_dyn(_c(PC_HI_IDX), UOp.const(0xFFFFFFFF, dtypes.uint32)))
+    return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)),
+                          ctx.wsgpr_dyn(_c(PC_HI_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)))
  # S_BARRIER: advance PC past the barrier instruction. The execution loop detects barriers before executing and handles synchronization.
  barrier_ops = {ir3.SOPPOp.S_BARRIER, irc.SOPPOp.S_BARRIER}
  if hasattr(ir4.SOPPOp, 'S_BARRIER_WAIT'): barrier_ops.add(ir4.SOPPOp.S_BARRIER_WAIT)
@@ -804,8 +804,8 @@ def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
    pc_bytes = ctx.rpc()  # PC is already 64-bit byte address
    vcc, exec_val = ctx.rmask(_c(VCC_LO.offset)), ctx.rexec()
    srcs: dict[str, UOp|int] = {'PC': pc_bytes.cast(dtypes.int64), 'SIMM16': simm16, 'SCC': ctx.rsgpr_dyn(_c(SCC.offset)), 'VCC': vcc,
-            'VCCZ': vcc.eq(UOp.const(0, vcc.dtype)).cast(dtypes.uint32),
-            'EXECZ': exec_val.eq(UOp.const(0, exec_val.dtype)).cast(dtypes.uint32)}
+            'VCCZ': vcc.eq(UOp.const(vcc.dtype, 0)).cast(dtypes.uint32),
+            'EXECZ': exec_val.eq(UOp.const(exec_val.dtype, 0)).cast(dtypes.uint32)}
    for dest, val in parse_pcode(pcode, srcs)[1]:
      if dest == 'PC' or dest.startswith('PC.'):
        lo, hi = _split64(val.cast(dtypes.uint64))
@@ -833,12 +833,12 @@ def _compile_smem(inst: ir3.SMEM | ir4.SMEM, ctx: _Ctx) -> UOp:
  part = op_name.rsplit('_', 1)[1]  # B32, DWORD, DWORDX2, U8, I8, etc.
  nval = int(part.removeprefix('DWORD').removeprefix('X') or '1') if 'DWORD' in part else int(part[1:]) / 32 * (-1 if part[0] == 'I' else 1)
  ndwords = max(1, int(abs(nval)))
-  dword_base = addr >> UOp.const(2, dtypes.uint64)
-  vals = [ctx.vmem.index(dword_base + UOp.const(i, dtypes.uint64)) for i in range(ndwords)]
+  dword_base = addr >> UOp.const(dtypes.uint64, 2)
+  vals = [ctx.vmem.index(dword_base + UOp.const(dtypes.uint64, i)) for i in range(ndwords)]
  if abs(nval) < 1:
    nbits = int(abs(nval) * 32)
-    byte_off = (addr & UOp.const(3, dtypes.uint64)).cast(dtypes.uint32) * UOp.const(8, dtypes.uint32)
-    extracted = (vals[0] >> byte_off) & UOp.const((1 << nbits) - 1, dtypes.uint32)
+    byte_off = (addr & UOp.const(dtypes.uint64, 3)).cast(dtypes.uint32) * UOp.const(dtypes.uint32, 8)
+    extracted = (vals[0] >> byte_off) & UOp.const(dtypes.uint32, (1 << nbits) - 1)
    vals[0] = extracted.cast({8: dtypes.int8, 16: dtypes.int16}[nbits]).cast(dtypes.int32).bitcast(dtypes.uint32) if nval < 0 else extracted
  stores = [ctx.wsgpr_dyn(sdata_reg + _c(i), vals[i]) for i in range(ndwords)]
  return UOp.sink(*stores, *ctx.inc_pc())
@@ -941,7 +941,7 @@ def _dpp16_ctrl(lane: UOp, dpp: int, row_mask: int, bank_mask: int, wave_size: i
  enabled = (((_c(row_mask) >> row.cast(dtypes.uint32)) & _c(1)).ne(_c(0)) &
             (((_c(bank_mask) >> bank.cast(dtypes.uint32)) & _c(1)).ne(_c(0))))
  op, arg = decode_dpp16(dpp)
-  src_lane, valid = lane_i, UOp.const(True)
+  src_lane, valid = lane_i, UOp.const(dtypes.bool, True)

  if op == 'quad_perm':
    assert isinstance(arg, tuple)
@@ -967,7 +967,7 @@ def _load_dpp16_src0(ctx: _Ctx, inst, lane: UOp, fallback: UOp) -> UOp:
                                         getattr(inst, 'bank_mask', 0xf) or 0xf, ctx.wave_size)
  safe_src_lane = (enabled & valid).where(src_lane, _c(0, dtypes.int))
  swizzled = ctx.rvgpr_dyn(ctx.inst_field(type(inst).vsrc0), safe_src_lane)
-  invalid = UOp.const(0, fallback.dtype) if getattr(inst, 'bc', 0) else fallback
+  invalid = UOp.const(fallback.dtype, 0) if getattr(inst, 'bc', 0) else fallback
  return enabled.where(valid.where(swizzled, invalid), fallback)

 def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc.VOPC_SDWA_SDST, ctx: _Ctx) -> UOp:
@@ -1159,7 +1159,7 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16
        s1 = _apply_src_mods(s1, 0, 1 if getattr(inst, 'src1_abs', 0) else 0, 1 if getattr(inst, 'src1_neg', 0) else 0, bits['s1'])
      s0 = _apply_src_mods(s0, 0, abs_bits, neg_bits, bits['s0'])
      s1 = _apply_src_mods(s1, 1, abs_bits, neg_bits, bits['s1'])
-    for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': UOp.const(0, dtypes.uint64)})[1]:
+    for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': UOp.const(dtypes.uint64, 0)})[1]:
      if '[laneId]' in dest and ('D0' in dest or 'EXEC' in dest): return val.cast(dtypes.uint32)
    return _c(0)

@@ -1189,8 +1189,8 @@ def _compile_bitop3(inst, ctx: _Ctx, exec_mask: UOp, bits: dict, op_name: str) -
  is_16 = 'B16' in op_name
  dt, mask = (dtypes.uint16, 0xFFFF) if is_16 else (dtypes.uint32, 0xFFFFFFFF)
  s0, s1, s2 = src0.cast(dt), src1.cast(dt), src2.cast(dt)
-  def bnot(v): return v ^ UOp.const(mask, dt)
-  result = UOp.const(0, dt)
+  def bnot(v): return v ^ UOp.const(dt, mask)
+  result = UOp.const(dt, 0)
  for i in range(8):
    if not (ttbl & (1 << i)): continue
    result = result | ((s0 if i & 4 else bnot(s0)) & (s1 if i & 2 else bnot(s1)) & (s2 if i & 1 else bnot(s2)))
@@ -1244,7 +1244,7 @@ def _compile_vop3(inst: ir3.VOP3 | ir4.VOP3 | irc.VOP3, ctx: _Ctx) -> UOp:
  src0 = _apply_src_mods(src0, 0, abs_bits, neg_bits, bits['s0'])
  src1 = _apply_src_mods(src1, 1, abs_bits, neg_bits, bits['s1'])
  src2 = _apply_src_mods(src2, 2, abs_bits, neg_bits, bits['s2'])
-  srcs = {'S0': src0, 'S1': src1, 'S2': src2, 'OPSEL': UOp.const(opsel, dtypes.uint32)}
+  srcs = {'S0': src0, 'S1': src1, 'S2': src2, 'OPSEL': UOp.const(dtypes.uint32, opsel)}
  if 'CNDMASK' in op_name and src2 is not None: srcs['VCC'] = src2
  # FMAC instructions need D0 (accumulator) from destination register
  if 'FMAC' in op_name: srcs['D0'] = ctx.rvgpr_dyn(vdst_reg, lane)
@@ -1398,7 +1398,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
  # Layout: tmp[0..n_a_elems-1] = A[m][k], tmp[n_a_elems..n_a_elems+n_b_elems-1] = B[n][k]
  # Within each group of lanes, lane%grp_sub gives M/N index, lane//grp_sub gives sub-block
  grp_sub = min(M, 16)  # lanes within group mapped to M/N dimension
-  b_off = UOp.const(n_a_elems, dtypes.int)
+  b_off = UOp.const(dtypes.int, n_a_elems)
  acc_dt = dtypes.int32 if is_int_out else dtypes.float32
  # Use uint32 temp array to prevent optimizer from eliminating f16→f32 bitcast chains.
  # The optimizer folds bitcast(uint32→float32) stores to float32 arrays, losing the conversion.
@@ -1407,44 +1407,44 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
  def cvt_elem(raw: UOp, sub_idx: int) -> UOp:
    if is_i8:
      # Extract i8, sign-extend to i32
-      byte_val = (raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)
-      return (byte_val.cast(dtypes.int32) ^ UOp.const(0x80, dtypes.int32)) - UOp.const(0x80, dtypes.int32)
+      byte_val = (raw >> UOp.const(dtypes.uint32, sub_idx * 8)) & UOp.const(dtypes.uint32, 0xFF)
+      return (byte_val.cast(dtypes.int32) ^ UOp.const(dtypes.int32, 0x80)) - UOp.const(dtypes.int32, 0x80)
    elif is_f32_src:
      return raw  # already uint32 (f32 bit pattern)
    elif is_fp8:
-      return ((raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)).cast(dtypes.uint32)
+      return ((raw >> UOp.const(dtypes.uint32, sub_idx * 8)) & UOp.const(dtypes.uint32, 0xFF)).cast(dtypes.uint32)
    elif is_bf16:
      # bf16→f32 bits: just shift left by 16 (bf16 is upper 16 bits of f32)
-      return ((raw >> UOp.const(sub_idx * 16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)) << UOp.const(16, dtypes.uint32)
+      return ((raw >> UOp.const(dtypes.uint32, sub_idx * 16)) & UOp.const(dtypes.uint32, 0xFFFF)) << UOp.const(dtypes.uint32, 16)
    else:
      # f16→f32 conversion using float arithmetic to avoid UOp optimizer eliminating the conversion.
      # The optimizer folds bitcast(uint32→float32) chains, so we compute the float value directly.
-      h = (raw >> UOp.const(sub_idx * 16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)
-      sign = (h >> UOp.const(15, dtypes.uint32)) & UOp.const(1, dtypes.uint32)
-      exp = (h >> UOp.const(10, dtypes.uint32)) & UOp.const(0x1F, dtypes.uint32)
-      mant = h & UOp.const(0x3FF, dtypes.uint32)
+      h = (raw >> UOp.const(dtypes.uint32, sub_idx * 16)) & UOp.const(dtypes.uint32, 0xFFFF)
+      sign = (h >> UOp.const(dtypes.uint32, 15)) & UOp.const(dtypes.uint32, 1)
+      exp = (h >> UOp.const(dtypes.uint32, 10)) & UOp.const(dtypes.uint32, 0x1F)
+      mant = h & UOp.const(dtypes.uint32, 0x3FF)
      # Use bf16 path: shift left by 16 to create bf16 bits, then shift mantissa and adjust exponent in float domain
      # bf16 bits = (sign << 15) | (exp_bf16 << 7) | mant_bf16 -- but f16 and bf16 have different formats
      # Instead: construct f32 bits properly, use a local uint32 array to force materialization
-      f32_bits = (sign << UOp.const(31, dtypes.uint32)) | \
-                 ((exp + UOp.const(112, dtypes.uint32)) << UOp.const(23, dtypes.uint32)) | \
-                 (mant << UOp.const(13, dtypes.uint32))
-      is_zero = exp.eq(UOp.const(0, dtypes.uint32))
+      f32_bits = (sign << UOp.const(dtypes.uint32, 31)) | \
+                 ((exp + UOp.const(dtypes.uint32, 112)) << UOp.const(dtypes.uint32, 23)) | \
+                 (mant << UOp.const(dtypes.uint32, 13))
+      is_zero = exp.eq(UOp.const(dtypes.uint32, 0))
      # Return uint32 (f32 bit pattern) — stored directly to uint32 temp array, bitcast to float on read
-      return is_zero.where(UOp.const(0, dtypes.uint32), f32_bits)
+      return is_zero.where(UOp.const(dtypes.uint32, 0), f32_bits)

  read_lane = ctx.range()
  # For 32x32: lane%16 = M/N index within 16-wide block, lane//16 = which of 4 quarter-waves
  # Groups: lanes 0-31 = group 0, lanes 32-63 = group 1
  # Within group: (lane%32)%16 = M/N[0-15], (lane%32)//16 selects M/N[0-15] or [16-31]
-  lane_in_grp = read_lane % UOp.const(grp_size, dtypes.int)
-  grp_idx = read_lane // UOp.const(grp_size, dtypes.int)
+  lane_in_grp = read_lane % UOp.const(dtypes.int, grp_size)
+  grp_idx = read_lane // UOp.const(dtypes.int, grp_size)

  if M == 32:
    # 32x32: lane_in_grp%16 = sub-row/col (0-15), lane_in_grp//16 = block (0=rows 0-15, 1=rows 16-31)
-    sub_mn = lane_in_grp % UOp.const(16, dtypes.int)
-    block_mn = lane_in_grp // UOp.const(16, dtypes.int)
-    mn_idx = block_mn * UOp.const(16, dtypes.int) + sub_mn  # actual M/N index (0-31)
+    sub_mn = lane_in_grp % UOp.const(dtypes.int, 16)
+    block_mn = lane_in_grp // UOp.const(dtypes.int, 16)
+    mn_idx = block_mn * UOp.const(dtypes.int, 16) + sub_mn  # actual M/N index (0-31)
  else:
    mn_idx = lane_in_grp  # for 16x16 and 4x4

@@ -1456,18 +1456,18 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
                                ctx.rsrc_dyn(src0_off, _c(0, dtypes.int), 32))
    a_val = cvt_elem(a_raw, sub_idx)
    if M == 4:
-      a_idx = grp_idx * UOp.const(M * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int)
+      a_idx = grp_idx * UOp.const(dtypes.int, M * K) + mn_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, kl)
    else:
-      a_idx = mn_idx * UOp.const(K, dtypes.int) + grp_idx * UOp.const(k_per_grp, dtypes.int) + UOp.const(kl, dtypes.int)
+      a_idx = mn_idx * UOp.const(dtypes.int, K) + grp_idx * UOp.const(dtypes.int, k_per_grp) + UOp.const(dtypes.int, kl)
    read_stores.append(tmp.index(a_idx).store(a_val))

    b_raw = src1_is_vgpr.where(ctx.rvgpr_dyn(src1_r + _c(reg_idx), read_lane),
                                ctx.rsrc_dyn(src1_off, _c(0, dtypes.int), 32))
    b_val = cvt_elem(b_raw, sub_idx)
    if M == 4:
-      b_idx = b_off + grp_idx * UOp.const(N * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int)
+      b_idx = b_off + grp_idx * UOp.const(dtypes.int, N * K) + mn_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, kl)
    else:
-      b_idx = b_off + mn_idx * UOp.const(K, dtypes.int) + grp_idx * UOp.const(k_per_grp, dtypes.int) + UOp.const(kl, dtypes.int)
+      b_idx = b_off + mn_idx * UOp.const(dtypes.int, K) + grp_idx * UOp.const(dtypes.int, k_per_grp) + UOp.const(dtypes.int, kl)
    read_stores.append(tmp.index(b_idx).store(b_val))

  read_phase = UOp.group(*read_stores).end(read_lane)
@@ -1488,11 +1488,11 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
    # Lane mapping: n = (lane%32)%16 + ((lane%32)//16)*16, gives column 0-31
    # Row groups: 4 groups of 4, covering rows 0-31. Group g (0-3): rows g*4 .. g*4+3
    # group assignment: lane//16 gives quarter (0-3), each quarter maps to 4 rows
-    c_lane_in_32 = compute_lane % UOp.const(32, dtypes.int)
-    c_sub = c_lane_in_32 % UOp.const(16, dtypes.int)
-    c_block = c_lane_in_32 // UOp.const(16, dtypes.int)
-    n_idx = c_block * UOp.const(16, dtypes.int) + c_sub
-    c_half = compute_lane // UOp.const(32, dtypes.int)  # 0 or 1
+    c_lane_in_32 = compute_lane % UOp.const(dtypes.int, 32)
+    c_sub = c_lane_in_32 % UOp.const(dtypes.int, 16)
+    c_block = c_lane_in_32 // UOp.const(dtypes.int, 16)
+    n_idx = c_block * UOp.const(dtypes.int, 16) + c_sub
+    c_half = compute_lane // UOp.const(dtypes.int, 32)  # 0 or 1

    for out_reg in range(16):
      # Each half covers 8 rows. out_reg 0-3: rows 0-3 (half0) or 16-19 (half1)
@@ -1503,7 +1503,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
      # acc[r] at lane l -> D[m][n] where n = (l%32)%16 + ((l%32)//16)*16
      # m = (l//32)*16 + (r//4)*4 + (r%4)  ... giving rows in blocks of 4
      # So: m_base = half * 16 + (out_reg // 4) * 4 + (out_reg % 4)
-      m_base = c_half * UOp.const(16, dtypes.int) + UOp.const((out_reg // 4) * 4 + (out_reg % 4), dtypes.int)
+      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 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)
@@ -1511,8 +1511,8 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
      acc = src2_is_vgpr.where(acc_v, acc_scalar)

      for k in range(K):
-        a_val = tmp2.index(m_base * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt)
-        b_val = tmp2.index(b_off + n_idx * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt)
+        a_val = tmp2.index(m_base * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt)
+        b_val = tmp2.index(b_off + n_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt)
        acc = acc + a_val * b_val

      if is_int_out:
@@ -1523,8 +1523,8 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
          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(grp_sub, dtypes.int)
-    c_grp = compute_lane // UOp.const(grp_sub, dtypes.int)
+    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 if use_acc else ctx.rvgpr_dyn)(src2_r + _c(out_reg), compute_lane, src2_is_vgpr)
@@ -1534,17 +1534,17 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:

      if M == 4:
        # 4x4: each group is independent. A/B indexed per-group.
-        m_base = c_grp * UOp.const(M * K, dtypes.int) + UOp.const(out_reg * K, dtypes.int)
+        m_base = c_grp * UOp.const(dtypes.int, M * K) + UOp.const(dtypes.int, out_reg * K)
        for k in range(K):
-          a_val = tmp2.index(m_base + UOp.const(k, dtypes.int)).bitcast(acc_dt)
-          b_val = tmp2.index(b_off + c_grp * UOp.const(N*K, dtypes.int) + n_idx * UOp.const(K, dtypes.int)+UOp.const(k, dtypes.int)).bitcast(acc_dt)
+          a_val = tmp2.index(m_base + UOp.const(dtypes.int, k)).bitcast(acc_dt)
+          b_val = tmp2.index(b_off + c_grp * UOp.const(dtypes.int, N*K) + n_idx * UOp.const(dtypes.int, K)+UOp.const(dtypes.int, k)).bitcast(acc_dt)
          acc = acc + a_val * b_val
      else:
        # 16x16: K is split across groups. Shared MxK/NxK arrays.
-        m_base = c_grp * UOp.const(out_per_lane, dtypes.int) + UOp.const(out_reg, dtypes.int)
+        m_base = c_grp * UOp.const(dtypes.int, out_per_lane) + UOp.const(dtypes.int, out_reg)
        for k in range(K):
-          a_val = tmp2.index(m_base * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt)
-          b_val = tmp2.index(b_off + n_idx * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt)
+          a_val = tmp2.index(m_base * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt)
+          b_val = tmp2.index(b_off + n_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt)
          acc = acc + a_val * b_val

      if is_int_out:
@@ -1570,8 +1570,8 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
  is_rdna4 = isinstance(inst, ir4.VOP3P)
  # read 16x16 F16/BF16 matrix from VGPRs → flat f32 array[row*16+k]
  def read_f16_val(src, lane, vgpr, half):
-    v = ctx.rvgpr_dyn(src + _c(vgpr), UOp.const(lane, dtypes.int))
-    return cvt((v >> UOp.const(16, dtypes.uint32)) if half else (v & UOp.const(0xFFFF, dtypes.uint32)))
+    v = ctx.rvgpr_dyn(src + _c(vgpr), UOp.const(dtypes.int, lane))
+    return cvt((v >> UOp.const(dtypes.uint32, 16)) if half else (v & UOp.const(dtypes.uint32, 0xFFFF)))

  # RDNA3: 16 lanes × 8 VGPRs × 2 halves, k maps linearly
  # RDNA4: 32 lanes × 4 VGPRs × 2 halves, k bits are scrambled (k[2] goes to lane bit 4)
@@ -1593,20 +1593,20 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
             for m in range(16) for n in range(16) for lane, vgpr in [d_map(m, n)]]
    mat_d = [sum(mat_a[r*16+k] * mat_b[c*16+k] for k in range(16)) + mat_c[r*16+c] for r in range(16) for c in range(16)]
    def f32_to_f16_bits(v: UOp) -> UOp: return v.cast(dtypes.half).bitcast(dtypes.uint16).cast(dtypes.uint32)
-    def f32_to_bf16_bits(v: UOp) -> UOp: return (v.bitcast(dtypes.uint32) >> UOp.const(16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)
+    def f32_to_bf16_bits(v: UOp) -> UOp: return (v.bitcast(dtypes.uint32) >> UOp.const(dtypes.uint32, 16)) & UOp.const(dtypes.uint32, 0xFFFF)
    out_cvt = f32_to_bf16_bits if is_bf16 else f32_to_f16_bits
    if is_rdna4:  # pack 2 f16 per VGPR: adjacent m values share (lane, vgpr) since vgpr=m&7, half=m&1
-      stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1] // 2), UOp.const(d_map(m, n)[0], dtypes.int),
-                out_cvt(mat_d[m*16+n]) | (out_cvt(mat_d[(m+1)*16+n]) << UOp.const(16, dtypes.uint32)), exec_mask)
+      stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1] // 2), UOp.const(dtypes.int, d_map(m, n)[0]),
+                out_cvt(mat_d[m*16+n]) | (out_cvt(mat_d[(m+1)*16+n]) << UOp.const(dtypes.uint32, 16)), exec_mask)
                for n in range(16) for m in range(0, 16, 2)]
    else:  # (rdna3) 1 f16 per VGPR (lo half only)
-      stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int), out_cvt(mat_d[m*16+n]), exec_mask)
+      stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(dtypes.int, d_map(m, n)[0]), out_cvt(mat_d[m*16+n]), exec_mask)
                for m in range(16) for n in range(16)]
  else: # f32
-    mat_c = [ctx.rvgpr_dyn(src2_r + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int)).bitcast(dtypes.float32)
+    mat_c = [ctx.rvgpr_dyn(src2_r + _c(d_map(m, n)[1]), UOp.const(dtypes.int, d_map(m, n)[0])).bitcast(dtypes.float32)
             for m in range(16) for n in range(16)]
    mat_d = [sum(mat_a[r*16+k] * mat_b[c*16+k] for k in range(16)) + mat_c[r*16+c] for r in range(16) for c in range(16)]
-    stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int), mat_d[m*16+n].bitcast(dtypes.uint32), exec_mask)
+    stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(dtypes.int, d_map(m, n)[0]), mat_d[m*16+n].bitcast(dtypes.uint32), exec_mask)
              for m in range(16) for n in range(16)]
  return UOp.sink(*stores, *ctx.inc_pc())

@@ -1688,8 +1688,8 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
      scalar_hi_sel = src_lo if not opsel_hi_bit else is_sgpr_pair.where(sgpr_hi, src_lo)
      lo = is_vgpr.where(vgpr_hi if opsel_lo else vgpr_lo, scalar_lo_sel)
      hi = is_vgpr.where(vgpr_hi if opsel_hi_bit else vgpr_lo, scalar_hi_sel)
-      if neg_lo: lo = lo ^ UOp.const(0x80000000, dtypes.uint32)
-      if neg_hi_bit: hi = hi ^ UOp.const(0x80000000, dtypes.uint32)
+      if neg_lo: lo = lo ^ UOp.const(dtypes.uint32, 0x80000000)
+      if neg_hi_bit: hi = hi ^ UOp.const(dtypes.uint32, 0x80000000)
      return _u64(lo, hi)
    srcs = {'S0': build_pk_f32(src0, src_offs[0], opsel & 1, opsel_hi & 1, neg & 1, neg_hi & 1),
            'S1': build_pk_f32(src1, src_offs[1], opsel & 2, opsel_hi & 2, neg & 2, neg_hi & 2),
@@ -1700,35 +1700,35 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
    def apply_abs(v, bit, opsel_hi_bit, opsel_bit):
      if not (neg_hi & bit): return v
      # Apply abs based on whether source is f32 or f16
-      if not (combined_opsel_hi & opsel_hi_bit): return v & UOp.const(0x7FFFFFFF, dtypes.uint32)  # f32 abs
-      if opsel & opsel_bit: return v & UOp.const(0x7FFF0000, dtypes.uint32)  # f16 hi abs (preserve lo)
-      return v & UOp.const(0xFFFF7FFF, dtypes.uint32)  # f16 lo abs (preserve hi)
+      if not (combined_opsel_hi & opsel_hi_bit): return v & UOp.const(dtypes.uint32, 0x7FFFFFFF)  # f32 abs
+      if opsel & opsel_bit: return v & UOp.const(dtypes.uint32, 0x7FFF0000)  # f16 hi abs (preserve lo)
+      return v & UOp.const(dtypes.uint32, 0xFFFF7FFF)  # f16 lo abs (preserve hi)
    def apply_neg_mix(v, bit, opsel_hi_bit, opsel_bit):
      if not (neg & bit): return v
-      if not (combined_opsel_hi & opsel_hi_bit): return v ^ UOp.const(0x80000000, dtypes.uint32)  # f32 neg
-      if opsel & opsel_bit: return v ^ UOp.const(0x80000000, dtypes.uint32)  # f16 hi neg
-      return v ^ UOp.const(0x00008000, dtypes.uint32)  # f16 lo neg
+      if not (combined_opsel_hi & opsel_hi_bit): return v ^ UOp.const(dtypes.uint32, 0x80000000)  # f32 neg
+      if opsel & opsel_bit: return v ^ UOp.const(dtypes.uint32, 0x80000000)  # f16 hi neg
+      return v ^ UOp.const(dtypes.uint32, 0x00008000)  # f16 lo neg
    s0_mod = apply_neg_mix(apply_abs(src0, 1, 1, 1), 1, 1, 1)
    s1_mod = apply_neg_mix(apply_abs(src1, 2, 2, 2), 2, 2, 2)
    s2_mod = apply_neg_mix(apply_abs(src2, 4, 4, 4), 4, 4, 4)
    srcs = {'S@0': s0_mod, 'S@1': s1_mod, 'S@2': s2_mod,
-            'OPSEL_HI': UOp.const(combined_opsel_hi, dtypes.uint32), 'OPSEL': UOp.const(opsel, dtypes.uint32)}
+            'OPSEL_HI': UOp.const(dtypes.uint32, combined_opsel_hi), 'OPSEL': UOp.const(dtypes.uint32, opsel)}
  else:
    def get_half_bits(val: UOp, use_hi: bool, apply_neg: bool = False) -> UOp:
-      bits = ((val >> UOp.const(16, dtypes.uint32)) if use_hi else val) & UOp.const(0xFFFF, dtypes.uint32)
+      bits = ((val >> UOp.const(dtypes.uint32, 16)) if use_hi else val) & UOp.const(dtypes.uint32, 0xFFFF)
      if apply_neg: bits = bits.cast(dtypes.uint16).bitcast(dtypes.half).neg().bitcast(dtypes.uint16).cast(dtypes.uint32)
      return bits
    def build_remapped_src(src: UOp, opsel_lo_bit: int, opsel_hi_bit: int, neg_lo_bit: int, neg_hi_bit: int) -> UOp:
      lo = get_half_bits(src, bool(opsel_lo_bit), bool(neg_lo_bit))
      hi = get_half_bits(src, bool(opsel_hi_bit), bool(neg_hi_bit))
-      return lo | (hi << UOp.const(16, dtypes.uint32))
+      return lo | (hi << UOp.const(dtypes.uint32, 16))
    # DOT IU instructions use NEG bits for signed/unsigned selection, not fp16 negation
    is_dot_iu = 'DOT' in op_name and 'IU' in op_name
    n0, n1, n2, nh0, nh1, nh2 = (0, 0, 0, 0, 0, 0) if is_dot_iu else (neg & 1, neg & 2, neg & 4, neg_hi & 1, neg_hi & 2, neg_hi & 4)
    srcs = {'S0': build_remapped_src(src0, opsel & 1, opsel_hi & 1, n0, nh0),
            'S1': build_remapped_src(src1, opsel & 2, opsel_hi & 2, n1, nh1),
            'S2': build_remapped_src(src2, opsel & 4, 1 if opsel_hi2 else 0, n2, nh2)}
-    if is_dot_iu: srcs['NEG'] = UOp.const(neg, dtypes.uint32)
+    if is_dot_iu: srcs['NEG'] = UOp.const(dtypes.uint32, neg)
  return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask)

 def _compile_vopd(inst: ir3.VOPD | ir4.VOPD, ctx: _Ctx) -> UOp:
@@ -1781,7 +1781,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
  # CDNA acc bit: when set, VGPR operands (vdst/vdata) target ACCVGPR file instead of VGPR
  use_acc = bool(getattr(inst, 'acc', 0))
  mem = ctx.lds if is_lds else ctx.scratch if is_scratch else ctx.vmem
-  addr_shift = UOp.const(2, dtypes.uint32 if is_lds else dtypes.uint64)
+  addr_shift = UOp.const(dtypes.uint32 if is_lds else dtypes.uint64, 2)

  # Extract register info - all dynamic for deduplication
  if is_lds:
@@ -1831,20 +1831,20 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
      return addr
    offset64 = offset.cast(dtypes.uint64)
    # Dynamic saddr check: saddr < 124 means valid SGPR, otherwise use VGPR pair for address
-    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(False)
+    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(dtypes.bool, False)
    if is_scratch:
      scratch_stride = ctx.rsgpr_dyn(_c(SCRATCH_STRIDE_IDX)).cast(dtypes.uint64)
      base = lane.cast(dtypes.uint64) * scratch_stride
      # SVE (Scratch VGPR Enable): when SVE=1, VADDR is used as offset; when SVE=0, VADDR is ignored
      sve = getattr(inst, 'sve', 0)
      vaddr = ctx.rvgpr_dyn(addr_reg, lane).cast(dtypes.uint64)
-      addr_offset = vaddr if sve == 1 else UOp.const(0, dtypes.uint64)
+      addr_offset = vaddr if sve == 1 else UOp.const(dtypes.uint64, 0)
      # Add saddr value only if use_saddr is true (saddr < 124)
-      saddr_contrib = use_saddr.where(ctx.rsgpr_dyn(saddr_reg).cast(dtypes.uint64), UOp.const(0, dtypes.uint64)) \
-        if saddr_reg is not None else UOp.const(0, dtypes.uint64)
+      saddr_contrib = use_saddr.where(ctx.rsgpr_dyn(saddr_reg).cast(dtypes.uint64), UOp.const(dtypes.uint64, 0)) \
+        if saddr_reg is not None else UOp.const(dtypes.uint64, 0)
      return base + addr_offset + saddr_contrib + offset64
    # FLAT/GLOBAL: choose between SGPR base (saddr) or VGPR pair (addr) based on saddr validity
-    saddr_base = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(0, dtypes.uint64)
+    saddr_base = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(dtypes.uint64, 0)
    vaddr_base = _u64(ctx.rvgpr_dyn(addr_reg, lane), ctx.rvgpr_dyn(addr_reg + _c(1), lane))
    # When saddr is valid: base = saddr pair, vaddr is 32-bit offset; otherwise: base = 0, vaddr is 64-bit address
    base_addr = use_saddr.where(saddr_base + ctx.rvgpr_dyn(addr_reg, lane).cast(dtypes.uint64), vaddr_base)
@@ -1874,18 +1874,18 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
        data = {'DATA': ctx.rvgpr_dyn(vdata_reg, lane), 'DATA1': ctx.rvgpr_dyn(vdata_reg + _c(1), lane),
                'DATA2': ctx.rvgpr_dyn(vdata_reg + _c(2), lane)}
      elif data_bits_mem <= 32:
-        data = {'DATA': ctx.rvgpr_dyn(vdata_reg, lane), 'DATA2': ctx.rvgpr_dyn(data1_reg, lane) if has_data1 else UOp.const(0, dtypes.uint32)}
+        data = {'DATA': ctx.rvgpr_dyn(vdata_reg, lane), 'DATA2': ctx.rvgpr_dyn(data1_reg, lane) if has_data1 else UOp.const(dtypes.uint32, 0)}
      else:
        data = {'DATA': _u64(ctx.rvgpr_dyn(vdata_reg, lane), ctx.rvgpr_dyn(vdata_reg + _c(1), lane)),
-                'DATA2': _u64(ctx.rvgpr_dyn(data1_reg, lane), ctx.rvgpr_dyn(data1_reg + _c(1), lane)) if has_data1 else UOp.const(0, dtypes.uint64)}
+                'DATA2': _u64(ctx.rvgpr_dyn(data1_reg, lane), ctx.rvgpr_dyn(data1_reg + _c(1), lane)) if has_data1 else UOp.const(dtypes.uint64, 0)}
      # RDNA3 uses ADDR/OFFSET, RDNA4 uses vgpr_a/offset (lowercase) + CalcDsAddr function
      return {'ADDR': addr, 'ADDR_BASE': addr, 'OFFSET': offset, 'OFFSET0': offset0, 'OFFSET1': offset1, '_lds': mem, 'laneId': lane,
              'vgpr_a': ctx.rvgpr_dyn(addr_reg, lane), 'offset': offset, 'offset0': offset0, 'offset1': offset1, **data}
    active = _lane_active(exec_mask, lane)
    # saddr < 124 means valid SGPR pair, otherwise use 0 (NULL means no saddr contribution)
-    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(False)
-    saddr_raw = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(0, dtypes.uint64)
-    saddr_base = use_saddr.where(saddr_raw, UOp.const(0, dtypes.uint64))
+    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(dtypes.bool, False)
+    saddr_raw = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(dtypes.uint64, 0)
+    saddr_base = use_saddr.where(saddr_raw, UOp.const(dtypes.uint64, 0))
    # Sign-extend offset to 64-bit for the final address calculation
    ioffset64 = offset.cast(dtypes.int64).cast(dtypes.uint64)
    # v_addr for CalcGlobalAddr: when saddr valid, use low 32 bits as offset; otherwise full 64-bit address. Include ioffset.
@@ -1900,13 +1900,13 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
    # acc bit: read/write ACCVGPR instead of VGPR for data operands
    _rvdata = (lambda r, l, *a: ctx.raccvgpr_dyn(r, l)) if use_acc else ctx.rvgpr_dyn
    vdata = _rvdata(vdata_reg, lane).cast(dtypes.uint64) if 'STORE' in op_name \
-      else _rvdata(vdst_reg, lane) if 'D16' in op_name else UOp.const(0, dtypes.uint32)
+      else _rvdata(vdst_reg, lane) if 'D16' in op_name else UOp.const(dtypes.uint32, 0)
    if 'STORE' in op_name and data_bits_mem >= 64:
-      vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
+      vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
    srcs = {'ADDR': addr, 'VDATA': vdata, '_vmem': mem, '_active': active,
            'laneId': lane, 'v_addr': vaddr_base, 's_saddr': saddr_base, 'SADDR': saddr_base, 'OFFSET': offset}
    for i in range(data_bits_mem // 32):
-      srcs[f'VDATA{i}'] = _rvdata(vdata_reg + _c(i), lane) if 'STORE' in op_name else UOp.const(0, dtypes.uint32)
+      srcs[f'VDATA{i}'] = _rvdata(vdata_reg + _c(i), lane) if 'STORE' in op_name else UOp.const(dtypes.uint32, 0)
    return srcs

  def make_stores(dest: str, val: UOp, lane: UOp, active: UOp, writes_return_data: bool) -> list[UOp]:
@@ -1983,7 +1983,7 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
  offset, offen, idxen = ctx.inst_field(type(inst).offset), ctx.inst_field(type(inst).offen), ctx.inst_field(type(inst).idxen)

  # V# descriptor: base[0:1], num_records[2], stride=word3[13:0]
-  base = _u64(ctx.rsgpr_dyn(srsrc), ctx.rsgpr_dyn(srsrc + _c(1))) & UOp.const(0xFFFFFFFFFFFF, dtypes.uint64)
+  base = _u64(ctx.rsgpr_dyn(srsrc), ctx.rsgpr_dyn(srsrc + _c(1))) & UOp.const(dtypes.uint64, 0xFFFFFFFFFFFF)
  num_records = ctx.rsgpr_dyn(srsrc + _c(2))
  stride = (ctx.rsgpr_dyn(srsrc + _c(3)) & _c(0x3FFF)).cast(dtypes.uint64)

@@ -2000,7 +2000,7 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
  buffer_offset = (stride * index + voff + offset.cast(dtypes.uint64)).cast(dtypes.uint32)
  in_bounds = active & buffer_offset.__lt__(num_records)
  addr = base + soff + buffer_offset.cast(dtypes.uint64)
-  addr = in_bounds.where(addr, UOp.const(0, dtypes.uint64))  # safe address when OOB
+  addr = in_bounds.where(addr, UOp.const(dtypes.uint64, 0))  # safe address when OOB
  mem = ctx.vmem

  stores: list[UOp] = []
@@ -2009,20 +2009,20 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
    lds_base = ctx.rsgpr_dyn(_c(124)) & _c(0x3FFFF)
    lds_addr = lds_base + lane.cast(dtypes.uint32) * _c(n_dwords * 4)
    for i in range(n_dwords):
-      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
+      word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
      val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64)).load(), _c(0))
      lds_idx = (lds_addr + _c(i * 4)) >> _c(2)
      lds_slot = ctx.lds.index(lds_idx.valid(active))
      stores.append(lds_slot.store(active.where(val, lds_slot)))
  elif is_store:
    for i in range(n_dwords):
-      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
+      word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
      idx = mem.index(word_addr.cast(dtypes.int64).valid(in_bounds))
      val = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(vdata + _c(i), lane)
      stores.append(idx.store(in_bounds.where(_to_u32(val), idx)))
  else:
    for i in range(n_dwords):
-      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
+      word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
      val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64).valid(in_bounds)).load(), _c(0))
      stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(vdata + _c(i), lane, val, exec_mask))
  return UOp.sink(UOp.group(*stores).end(lane), *ctx.inc_pc())
@@ -7,7 +7,7 @@ from tinygrad.codegen.decomp.dtype import f2f
 # Type alias for vars dict: stores UOps and tuples for lambda definitions
 VarVal = UOp | tuple[str, list[str], str]

-def _const(dt, v): return UOp.const(v, dt)
+def _const(dt, v): return UOp.const(dt, v)
 def _u32(v): return _const(dtypes.uint32, v)
 def _u64(v): return _const(dtypes.uint64, v)
 def _to_u32(v): return v if v.dtype == dtypes.uint32 else v.bitcast(dtypes.uint32) if v.dtype.itemsize == 4 else v.cast(dtypes.uint32)
@@ -55,8 +55,8 @@ def _expr_bits(v: UOp) -> int:
  if v.op in (Ops.AND, Ops.XOR):
    widths: list[int] = []
    for src in v.src:
-      if src.op == Ops.CONST and isinstance(src.val, int) and src.val > 0 and (src.val & (src.val + 1)) == 0:
-        widths.append(src.val.bit_length())
+      if src.op == Ops.CONST and isinstance(src.arg, int) and src.arg > 0 and (src.arg & (src.arg + 1)) == 0:
+        widths.append(src.arg.bit_length())
    if widths: return max(widths)
  return v.dtype.bitsize

@@ -144,9 +144,9 @@ def _minmax_reduce(is_max: bool, dt, *args: UOp) -> UOp:
 def _find_two_pi_mul(x):
  if x.op != Ops.MUL or len(x.src) != 2: return None
  for i, s in enumerate(x.src):
-    if s.op == Ops.CONST and abs(s.val - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
+    if s.op == Ops.CONST and abs(s.arg - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
    if s.op == Ops.MUL and len(s.src) == 2:
-      vals = [ss.val for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].val for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST]
+      vals = [ss.arg for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].arg for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST]
      if len(vals) == 2 and abs(vals[0] * vals[1] - 6.283185307179586) < 1e-5: return (x.src[1-i], vals[0] * vals[1])
  return None

@@ -163,7 +163,7 @@ def _trig_reduce(x, phase=0.0):

 def _signext(val: UOp) -> UOp:
  for bits, mask, ext in [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
-    if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].val == mask) or val.dtype.itemsize == bits // 8:
+    if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].arg == mask) or val.dtype.itemsize == bits // 8:
      v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
      sb = (v32 >> _u32(bits - 1)) & _u32(1)
      return sb.ne(_u32(0)).where(v32 | _u32(ext), v32).cast(dtypes.int)
@@ -497,7 +497,7 @@ class Parser:
        if not dtypes.is_int(right.dtype): right = right.cast(dtypes.uint32)
        return (left >> right) if op == '>>' else (left << right)
      case '+' | '-':
-        if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.val - right.val)
+        if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.arg - right.arg)
        return (left + right) if op == '+' else (left - right)
      case '*' | '/':
        # Integer promotion: promote 16-bit integers to 32-bit before multiply to avoid overflow
@@ -507,7 +507,7 @@ class Parser:
          left, right = left.cast(pdt), right.cast(pdt)
        if op == '*': return left * right
        return (left // right) if dtypes.is_int(left.dtype) else (left / right)
-      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.val == 2.0 else left
+      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.arg == 2.0 else left

  _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)]

@@ -530,7 +530,7 @@ class Parser:
    if self.try_eat_val('-', 'OP'):
      inner = self.unary()
      if inner.op == Ops.CONST:
-        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.val)
+        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.arg)
      return inner.neg()
    if self.try_eat_val('+', 'OP'): return self.unary()
    return self.postfix()
@@ -670,14 +670,14 @@ class Parser:
      width = self.parse()
      self.eat('RBRACKET')
      if width.op == Ops.CONST:
-        w = int(width.val)
+        w = int(width.arg)
        return (base >> _to_u32(first)) & _const(base.dtype, (1 << w) - 1)
      return base
    if self.try_eat('COLON'):
      second = self.parse()
      self.eat('RBRACKET')
      if first.op == Ops.CONST and second.op == Ops.CONST:
-        a, b = int(first.val), int(second.val)
+        a, b = int(first.arg), int(second.arg)
        if a < b: return _bitreverse(base, b - a + 1)
        hi, lo = a, b
        if lo >= base.dtype.itemsize * 8:
@@ -699,7 +699,7 @@ class Parser:
    if var_name is None:
      var_name = self._find_var_name(base)
    if first.op == Ops.CONST:
-      idx = int(first.val)
+      idx = int(first.arg)
      # Check for array element (var@idx)
      if var_name and f'{var_name}@{idx}' in self.vars:
        v = self.vars[f'{var_name}@{idx}']
@@ -866,13 +866,13 @@ class Parser:
        idx_hi_native = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int64)
        safe_idx_hi = is_unaligned.where(idx_hi_native, idx_native)
        hi = mindex(safe_idx_hi)
-        combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
-        val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(8, dtypes.uint64))).cast(dtypes.uint32), val)
+        combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
+        val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(dtypes.uint64, 8))).cast(dtypes.uint32), val)
    return _cast_to(val, dt)

  def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
    if l.dtype != r.dtype:
-      if r.dtype == dtypes.int and r.op == Ops.CONST and r.val < 0: l = l.cast(dtypes.int)
+      if r.dtype == dtypes.int and r.op == Ops.CONST and r.arg < 0: l = l.cast(dtypes.int)
      else: r = r.cast(l.dtype)
    return l, r

@@ -970,7 +970,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
        if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
        expr = p.parse().simplify()
        assert expr.op == Ops.CONST, f"loop bound must be constant, got {expr}"
-        return int(expr.val)
+        return int(expr.arg)
      start_val = parse_bound()
      p.eat('COLON')
      end_val = parse_bound()
@@ -1258,7 +1258,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
      def parse_cond(s, kw):
        ll = s.lower()
        return _to_bool(parse_expr(s[ll.find(kw) + len(kw):ll.rfind('then')].strip(), env, funcs))
-      def is_const(c, v): return c.op == Ops.CONST and c.val is v
+      def is_const(c, v): return c.op == Ops.CONST and c.arg is v
      cond = parse_cond(line, 'if')
      conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not is_const(cond, False) else []
      branch_assigns: list[tuple[UOp, list]] = []  # (cond, assigns_list) for side-effect merging
@@ -1339,7 +1339,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
              return (dest, (val[0], cnd.where(val[1], val[1])))
            return (dest, val)
          # Build combined condition: each branch fires when its cond is true AND no earlier cond was true
-          remaining = UOp.const(True)
+          remaining = UOp.const(dtypes.bool, True)
          for bc, bse in branch_assigns:
            effective = remaining & bc if remaining.op != Ops.CONST else bc
            for dest, val in bse: assigns.append(_cond_side_effect(effective, dest, val))
@@ -91,7 +91,7 @@ class GPFIFO:
    args_cnt, vals_cnt = const0[80], const0[81]
    args_addr = qmd.constant_buffer_addr_lower_0 + (qmd.constant_buffer_addr_upper_0 << 32) + 0x160
    args = to_mv(args_addr, args_cnt*8).cast('Q')
-    vals = to_mv(args_addr + args_cnt*8, vals_cnt*8).cast('Q')
+    vals = to_mv(args_addr + args_cnt*8, vals_cnt*4).cast('I')
    cargs = [ctypes.cast(args[i], ctypes.c_void_p) for i in range(args_cnt)] + [ctypes.cast(vals[i], ctypes.c_void_p) for i in range(vals_cnt)]
    gx, gy, gz = qmd.cta_raster_width, qmd.cta_raster_height, qmd.cta_raster_depth
    lx, ly, lz = qmd.cta_thread_dimension0, qmd.cta_thread_dimension1, qmd.cta_thread_dimension2
@@ -1,6 +1,6 @@
 import unittest, itertools, math
 from tinygrad import Tensor, dtypes, Context
-from tinygrad.dtype import DType, ConstType
+from tinygrad.dtype import DType, ConstType, Invalid
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np
@@ -36,20 +36,21 @@ class TestUnaryOpsConstFolding(unittest.TestCase):

 class TestWeakConstFolding(unittest.TestCase):
  def test_weakint_math(self):
-    out = (UOp.const(2**40) + UOp.const(2**40)).simplify()
-    self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakint, 2**41))
+    out = (UOp.const(dtypes.weakint, 2**40) + UOp.const(dtypes.weakint, 2**40)).simplify()
+    self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakint, 2**41))

  def test_float_unaries(self):
-    for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL):
-      out = UOp.const(4.0).alu(op).simplify()
-      self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat))
+    for dtype in (dtypes.weakfloat,):
+      for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL):
+        out = UOp.const(dtype, 4).alu(op).simplify()
+        self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat))

  def test_weakfloat_math(self):
-    out = (UOp.const(1.25) + UOp.const(2.5)).simplify()
-    self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakfloat, 3.75))
+    out = (UOp.const(dtypes.weakfloat, 1.25) + UOp.const(dtypes.weakfloat, 2.5)).simplify()
+    self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakfloat, 3.75))

  def test_invalid_poison(self):
-    self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)
+    self.assertIs(UOp.invalid().alu(Ops.CDIV, UOp.const(dtypes.weakint, 0)).simplify().arg, Invalid)

 class TestBinaryOpsConstFolding(unittest.TestCase):
  def test_add_literal_zero(self):
@@ -120,10 +121,10 @@ class TestBitcastConstFolding(unittest.TestCase):
    def t(cases: dict[DType, ConstType]):
      for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
        if not math.isnan(from_v):
-          r = full_rewrite(UOp.const(from_v, from_dt).bitcast(to_dt).sink()).src[0]
+          r = full_rewrite(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0]
          self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
          self.assertEqual(r.dtype, to_dt, msg)
-          np.testing.assert_equal(r.val, to_v, msg)
+          np.testing.assert_equal(r.arg, to_v, msg)

    t({dtypes.int8: 0, dtypes.uint8: 0, dtypes.bool: False})
    t({dtypes.int8: 1, dtypes.uint8: 1, dtypes.bool: True})
@@ -144,9 +145,9 @@ class TestBitcastConstFolding(unittest.TestCase):

  def test_vec_bitcast(self):
    with Context(SPEC=0):
-      srcs = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink()).src
+      srcs = full_rewrite(UOp.const(dtypes.int32, (-1, -2**31, 75)).bitcast(dtypes.uint32).sink()).src
    self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs))
-    self.assertEqual(tuple(x.val for x in srcs), (2**32-1, 2**31, 75))
+    self.assertEqual(tuple(x.arg for x in srcs), (2**32-1, 2**31, 75))

 # folds advance indexing into basic indexing
 class TestIndexingConstFolding(unittest.TestCase):
@@ -2,7 +2,7 @@ import unittest, math, struct, operator
 from tinygrad import Tensor, Device
 from tinygrad.dtype import DTYPES_DICT, dtypes, Invalid, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float

-from tinygrad.helpers import getenv, Context
+from tinygrad.helpers import getenv
 from hypothesis import given, settings, strategies as strat
 import numpy as np
 import torch
@@ -58,8 +58,8 @@ class TestHelpers(unittest.TestCase):
  def test_from_py(self):
    assert dtypes.from_py(True) == dtypes.bool
    assert dtypes.from_py(Invalid) == dtypes.bool
-    assert dtypes.from_py(2) == dtypes.weakint
-    assert dtypes.from_py(3.0) == dtypes.weakfloat
+    assert dtypes.from_py(2) == dtypes.default_int
+    assert dtypes.from_py(3.0) == dtypes.default_float
    assert dtypes.from_py([]) == dtypes.default_float
    assert dtypes.from_py(()) == dtypes.default_float
    assert dtypes.from_py([True]) == dtypes.bool
@@ -245,14 +245,19 @@ class TestTypePromotion(unittest.TestCase):
    assert least_upper_dtype(dtypes.weakfloat, dtypes.float64) == dtypes.float64

 class TestTypeSpec(unittest.TestCase):
+  def setUp(self):
+    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
+  def tearDown(self):
+    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
+
  def test_set_dtype_default(self):
    for default_int in [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64]:
-      with Context(DEFAULT_INT=default_int):
-        assert dtypes.default_int == default_int
+      dtypes.default_int = default_int
+      assert dtypes.default_int == default_int

    for default_float in [*dtypes.fp8s, dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
-      with Context(DEFAULT_FLOAT=default_float):
-        assert dtypes.default_float == default_float
+      dtypes.default_float = default_float
+      assert dtypes.default_float == default_float

  @given(strat.sampled_from(core_dtypes), strat.sampled_from([operator.gt, operator.ge, operator.le, operator.lt, operator.eq, operator.ne]))
  def test_bool_ops(self, dtype, op):
@@ -260,7 +265,7 @@ class TestTypeSpec(unittest.TestCase):

  @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
  def test_functions_return_index(self, dtype, default_int, default_float):
-    self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float))
+    dtypes.default_int, dtypes.default_float = default_int, default_float
    assert Tensor([0, 1], dtype=dtype).argmax().dtype == dtypes.int32
    assert Tensor([0, 1], dtype=dtype).argmin().dtype == dtypes.int32
    assert Tensor([0, 1], dtype=dtype).multinomial().dtype == dtypes.int32
@@ -280,7 +285,7 @@ class TestTypeSpec(unittest.TestCase):

  @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
  def test_attention_returns_same_dtype(self, data_dtype, default_float):
-    self.enterContext(Context(DEFAULT_FLOAT=default_float))
+    dtypes.default_float = default_float
    query = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
    key = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
    value = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
@@ -291,27 +296,32 @@ class TestTypeSpec(unittest.TestCase):
    assert query.scaled_dot_product_attention(key, value, attn_mask=mask).dtype == data_dtype

 class TestAutoCastType(unittest.TestCase):
+  def setUp(self):
+    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
+  def tearDown(self):
+    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
+
  @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
  def test_least_upper_float_input_is_float(self, input_dtype, default_float):
-    self.enterContext(Context(DEFAULT_FLOAT=default_float))
+    dtypes.default_float = default_float
    self.assertEqual(least_upper_float(input_dtype), input_dtype)

  @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
  def test_least_upper_float_input_is_int(self, input_dtype, default_float):
-    self.enterContext(Context(DEFAULT_FLOAT=default_float))
+    dtypes.default_float = default_float
    self.assertEqual(least_upper_float(input_dtype), default_float)

  @given(strat.sampled_from(core_dtypes))
  def test_broadcast_scalar(self, dt):
-    assert (Tensor.ones(4, 4, dtype=dt) + 2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.weakfloat)
-    assert (Tensor.ones(4, 4, dtype=dt) + 2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint)
+    assert (Tensor.ones(4, 4, dtype=dt) + 2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
+    assert (Tensor.ones(4, 4, dtype=dt) + 2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
    assert (Tensor.ones(4, 4, dtype=dt) + True).dtype == dt

  @given(strat.sampled_from(core_dtypes))
  def test_pad_scalar(self, dt):
    t = Tensor.ones(4, dtype=dt)
-    assert t.pad(((1, 1),), value=2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.weakfloat)
-    assert t.pad(((1, 1),), value=2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint)
+    assert t.pad(((1, 1),), value=2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
+    assert t.pad(((1, 1),), value=2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
    assert t.pad(((1, 1),), value=True).dtype == dt

  @given(strat.sampled_from(core_dtypes))
@@ -322,7 +332,7 @@ class TestAutoCastType(unittest.TestCase):

  @given(strat.sampled_from(dtype_floats))
  def test_int_div_int(self, default_float):
-    self.enterContext(Context(DEFAULT_FLOAT=default_float))
+    dtypes.default_float = default_float
    self.assertEqual(Tensor([1]).div(Tensor([2])).dtype, default_float)

  def test_sum(self):
@@ -410,16 +420,16 @@ class TestAutoCastType(unittest.TestCase):
  @given(strat.sampled_from(core_dtypes))
  def test_where_one_scalar(self, dt):
    t = Tensor(2, dtype=dt)
-    self.check_where_alternate_input_other(t, 3.2, (dt if dtypes.is_float(dt) else dtypes.weakfloat))
-    self.check_where_alternate_input_other(t, 3, (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint))
+    self.check_where_alternate_input_other(t, 3.2, (dt if dtypes.is_float(dt) else dtypes.default_float))
+    self.check_where_alternate_input_other(t, 3, (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int))
    self.check_where_alternate_input_other(t, True, dt)

  def test_where_two_scalars(self):
-    self.check_where_alternate_input_other(3.1, 3.2, dtypes.weakfloat)
-    self.check_where_alternate_input_other(3.1, 3, dtypes.weakfloat)
-    self.check_where_alternate_input_other(3.1, True, dtypes.weakfloat)
-    self.check_where_alternate_input_other(3, 2, dtypes.weakint)
-    self.check_where_alternate_input_other(3, True, dtypes.weakint)
+    self.check_where_alternate_input_other(3.1, 3.2, dtypes.default_float)
+    self.check_where_alternate_input_other(3.1, 3, dtypes.default_float)
+    self.check_where_alternate_input_other(3.1, True, dtypes.default_float)
+    self.check_where_alternate_input_other(3, 2, dtypes.default_int)
+    self.check_where_alternate_input_other(3, True, dtypes.default_int)

  def test_where_non_bool_cond_raises(self):
    with self.assertRaises(RuntimeError): Tensor([1, 0, 2]).where(1, 0)
@@ -431,8 +441,8 @@ class TestAutoCastType(unittest.TestCase):

  @given(strat.sampled_from(core_dtypes))
  def test_maximum_const(self, dt):
-    assert Tensor([1, 2], dtype=dt).maximum(3.1).dtype == (dt if dtypes.is_float(dt) else dtypes.weakfloat)
-    assert Tensor([1, 2], dtype=dt).maximum(3).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint)
+    assert Tensor([1, 2], dtype=dt).maximum(3.1).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
+    assert Tensor([1, 2], dtype=dt).maximum(3).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
    assert Tensor([1, 2], dtype=dt).maximum(True).dtype == dt

  def test_div(self):
@@ -443,7 +453,7 @@ class TestAutoCastType(unittest.TestCase):

  def test_div_const(self):
    assert (Tensor([1, 2], dtype=dtypes.int32) / 2).dtype == dtypes.default_float
-    assert (Tensor([1, 2], dtype=dtypes.int32) / 2.0).dtype == dtypes.weakfloat
+    assert (Tensor([1, 2], dtype=dtypes.int32) / 2.0).dtype == dtypes.default_float
    assert (Tensor([1, 2], dtype=dtypes.float16) / 2).dtype == dtypes.float16
    assert (Tensor([1, 2], dtype=dtypes.float16) / 2.0).dtype == dtypes.float16

@@ -12,7 +12,7 @@ class TestGroupedDims(unittest.TestCase):
    idxs = get_grouped_dims(prefix, dims, max_sizes, reverse)
    loop_idxs = dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs]))
    loop_idxs = sorted(loop_idxs, key=lambda uop: uop.arg)
-    sizes = [x.src[0].val for x in loop_idxs]
+    sizes = [x.src[0].arg for x in loop_idxs]
    assert len(idxs) == len(dims), f"expected idxs to have same length as dims {len(dims)}, got {len(idxs)}"
    if assert_same_length:
      assert len(loop_idxs) == min(len(sizes), len(dims)), f"expected idxs to have length {min(len(sizes), len(dims))}, got {len(loop_idxs)}"
@@ -24,7 +24,7 @@ class TestGroupedDims(unittest.TestCase):
    total = math.prod(dims)
    specials = sorted(dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs])), key=lambda u: u.arg)
    # build flat index and primed flat (same expression with renamed SPECIALs)
-    flat = UOp.const(0)
+    flat = UOp.const(dtypes.weakint, 0)
    for i, idx in enumerate(idxs):
      flat = flat + idx * int(math.prod(dims[i+1:]))
    flat_p = flat.substitute({s: UOp(Ops.SPECIAL, src=s.src, arg=s.arg+"_p") for s in specials})
@@ -107,7 +107,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.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
+    sink = UOp.param(0, dtypes.float, (512,)).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)
@@ -14,10 +14,10 @@ class TestGradient(unittest.TestCase):
  def _test_one_input_function(self, f:Callable, jf:Callable|None=None):
    if jf is None: jf = f
    x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
-    gx = compute_gradient(f(x), UOp.const(1.0), set([x]))[x]
+    gx = compute_gradient(f(x), UOp.const(dtypes.float, 1.0), set([x]))[x]

    for val in [-5., -2.0, 0.0, 2.0, 5.]:
-      tg_out = gx.substitute({x: UOp.const(val)}).ssimplify()
+      tg_out = gx.substitute({x: x.const_like(val)}).ssimplify()
      tx = torch.tensor([val], dtype=torch.float, requires_grad=True)
      torch_out = torch.autograd.grad(jf(tx), tx)[0].item()
      self._cmp_nan_okay(tg_out, torch_out)
@@ -26,13 +26,13 @@ class TestGradient(unittest.TestCase):
    if jf is None: jf = f
    x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
    y = UOp.variable('y', -math.inf, math.inf, dtype=dtypes.float)
-    grads = compute_gradient(f(x, y), UOp.const(1.0), set([x, y]))
+    grads = compute_gradient(f(x, y), UOp.const(dtypes.float, 1.0), set([x, y]))
    gx, gy = grads[x], grads[y]

    for valx in [-5., -2.0, 0.0, 2.0, 5.]:
      for valy in [-5., -2.0, 0.0, 2.0, 5.]:
        # Substitute the values into the gradient expressions
-        substitutions = {x: UOp.const(valx), y: UOp.const(valy)}
+        substitutions = {x: x.const_like(valx), y: y.const_like(valy)}
        tg_out_x = gx.substitute(substitutions).ssimplify()
        tg_out_y = gy.substitute(substitutions).ssimplify()

@@ -15,13 +15,13 @@ def apply_rewrite(expr):
 def apply_rewrite_values(expr):
  srcs = full_rewrite(expr.sink()).src
  if len(srcs) == 1:
-    if srcs[0].op is Ops.CONST: return (srcs[0].val,)
-    if srcs[0].op is Ops.STACK: return tuple(s.val for s in srcs[0].src)
-  return tuple(s.val for s in srcs)
+    if srcs[0].op is Ops.CONST: return (srcs[0].arg,) if not isinstance(srcs[0].arg, tuple) else srcs[0].arg
+    if srcs[0].op is Ops.STACK: return tuple(s.arg for s in srcs[0].src)
+  return tuple(s.arg for s in srcs)

 def evaluate_uop(uop, variables):
  if uop.op == Ops.CONST:
-    return uop.val
+    return uop.arg
  elif uop.op == Ops.PARAM and uop.arg.addrspace is AddrSpace.ALU:
    return variables[uop.expr]
  elif uop.op in GroupOp.ALU:
@@ -32,67 +32,67 @@ def evaluate_uop(uop, variables):

 class TestArithmeticSimplifications(unittest.TestCase):
  def test_full_graph_rewrite_division_by_zero(self):
-    optimized_div_uop = apply_rewrite(UOp.const(10.0) / UOp.const(0.0))
+    optimized_div_uop = apply_rewrite(UOp.const(dtypes.float32, 10.0) / UOp.const(dtypes.float32, 0.0))
    self.assertEqual(optimized_div_uop.op, Ops.CONST)
-    self.assertTrue(math.isinf(optimized_div_uop.val) or math.isnan(optimized_div_uop.val))
+    self.assertTrue(math.isinf(optimized_div_uop.arg) or math.isnan(optimized_div_uop.arg))

  def test_full_graph_rewrite_redundant_operations(self):
-    optimized_uop = apply_rewrite((UOp.const(10.0) + UOp.const(0.0)) * UOp.const(1.0))
+    optimized_uop = apply_rewrite((UOp.const(dtypes.float32, 10.0) + UOp.const(dtypes.float32, 0.0)) * UOp.const(dtypes.float32, 1.0))
    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 10.0)
+    self.assertEqual(optimized_uop.arg, 10.0)

  def test_full_graph_rewrite_large_graph(self):
-    prev_uop = UOp.const(0)
+    prev_uop = UOp.const(dtypes.int32, 0)
    for i in range(1, 101):
-      prev_uop += UOp.const(i)
+      prev_uop += UOp.const(dtypes.int32, i)
    optimized_uop = apply_rewrite(prev_uop)
    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, sum(range(1, 101)))
+    self.assertEqual(optimized_uop.arg, sum(range(1, 101)))

  def test_full_graph_rewrite_division_by_one(self):
-    optimized_uop = apply_rewrite(UOp.const(42.0) / UOp.const(1.0))
+    optimized_uop = apply_rewrite(UOp.const(dtypes.float32, 42.0) / UOp.const(dtypes.float32, 1.0))
    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 42.0)
+    self.assertEqual(optimized_uop.arg, 42.0)

  def test_full_graph_rewrite_modulo_by_one(self):
-    optimized_uop = apply_rewrite(UOp.const(42) % UOp.const(1))
+    optimized_uop = apply_rewrite(UOp.const(dtypes.int32, 42) % UOp.const(dtypes.int32, 1))
    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 0)
+    self.assertEqual(optimized_uop.arg, 0)


 class TestFoldingAndReduction(unittest.TestCase):
  @unittest.skip("reduce is removed now")
  def test_full_graph_rewrite_constant_reduction_folding(self):
-    const1 = UOp.const(5)
-    const2 = UOp.const(10)
-    const3 = UOp.const(20)
+    const1 = UOp.const(dtypes.int32, 5)
+    const2 = UOp.const(dtypes.int32, 10)
+    const3 = UOp.const(dtypes.int32, 20)
    optimized_sink = apply_rewrite((const1 + const2 + const3).reduce(Ops.ADD))
    expected_sum = 5 + 10 + 20
-    self.assertEqual(optimized_sink.val, expected_sum)
+    self.assertEqual(optimized_sink.arg, expected_sum)

  @unittest.skip("reduce is removed now")
  def test_full_graph_rewrite_reduction_with_unused_range(self):
-    const1 = UOp.const(15)
-    const2 = UOp.const(25)
+    const1 = UOp.const(dtypes.int32, 15)
+    const2 = UOp.const(dtypes.int32, 25)
    rng = UOp.range(10, idx=0)
    optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng))
    expected_sum = 10 * (15 + 25)
-    self.assertEqual(optimized_sink.val, expected_sum)
+    self.assertEqual(optimized_sink.arg, expected_sum)

  @unittest.skip("currently failing")
  def test_full_graph_rewrite_range_reduction(self):
    simple_range = UOp.range(5, idx=0)
    optimized_sink = apply_rewrite(simple_range.reduce(Ops.ADD, simple_range))
    expected_sum = sum(range(5))
-    self.assertEqual(optimized_sink.val, expected_sum)
+    self.assertEqual(optimized_sink.arg, expected_sum)

  @unittest.skip("currently failing")
  def test_full_graph_rewrite_simple_reduction_folding(self):
    simple_range = UOp.range(4, idx=0)
-    add_uop = simple_range + UOp.const(1)
+    add_uop = simple_range + UOp.const(dtypes.int32, 1)
    optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range))
    expected_sum = sum(i + 1 for i in range(4))
-    self.assertEqual(optimized_sink.val, expected_sum)
+    self.assertEqual(optimized_sink.arg, expected_sum)

  @unittest.skip("currently failing")
  def test_full_graph_rewrite_nested_loop_collapse(self):
@@ -101,7 +101,7 @@ class TestFoldingAndReduction(unittest.TestCase):
    expr = (outer_range * 10) + inner_range
    optimized_reduce_uop = apply_rewrite(expr.reduce(Ops.ADD, outer_range, inner_range))
    self.assertEqual(optimized_reduce_uop.op, Ops.CONST)
-    self.assertEqual(optimized_reduce_uop.val, sum((i * 10) + j for i in range(8) for j in range(4)))
+    self.assertEqual(optimized_reduce_uop.arg, sum((i * 10) + j for i in range(8) for j in range(4)))


 class TestModuloAndDivisionFolding(unittest.TestCase):
@@ -110,27 +110,27 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
    x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint)
    optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4)
    self.assertEqual(optimized_mod_uop.op, Ops.CONST)
-    self.assertEqual(optimized_mod_uop.val, 2)
+    self.assertEqual(optimized_mod_uop.arg, 2)

  def test_full_graph_rewrite_division_folding_with_define_var(self):
    # index dtype because div-mod rules only work on index
    n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint)
    optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3)
    self.assertEqual(optimized_div_uop.op, Ops.MUL)
-    self.assertEqual(optimized_div_uop.src[1].val, 2)
+    self.assertEqual(optimized_div_uop.src[1].arg, 2)

  def test_full_graph_rewrite_complex_mod_div_folding(self):
    # index dtype because div-mod rules only work on index
    k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint)
    optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2)
    self.assertEqual(optimized_div_uop.op, Ops.CONST)
-    self.assertEqual(optimized_div_uop.val, 1)
+    self.assertEqual(optimized_div_uop.arg, 1)

  def test_graph_rewrite_div_folding_bug(self):
    lhs = UOp(Ops.ADD, src=(
-      UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(32),), arg='lidx0'),)*4),
-      UOp.const((0, 256, 512, 768))))
-    rhs = UOp.const((2,)*4)
+      UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 32),), arg='lidx0'),)*4),
+      UOp.const(dtypes.int, (0, 256, 512, 768))))
+    rhs = UOp.const(dtypes.int, (2,)*4)
    unopt = lhs<rhs
    opt = apply_rewrite(unopt)
    print(unopt)
@@ -159,11 +159,11 @@ class TestModuloAndDivisionFolding(unittest.TestCase):

 class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
  def test_full_graph_rewrite_transcendental_edge_cases(self):
-    optimized_sink = full_rewrite(UOp.const(-1.0).log2().sink(UOp.const(0.0).reciprocal()))
+    optimized_sink = full_rewrite(UOp.const(dtypes.float32, -1.0).log2().sink(UOp.const(dtypes.float32, 0.0).reciprocal()))
    optimized_log2_neg, optimized_recip_zero = optimized_sink.src
-    self.assertTrue(math.isnan(optimized_log2_neg.val), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.val}")
-    self.assertTrue(math.isinf(optimized_recip_zero.val) and optimized_recip_zero.val > 0,
-                    f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.val}")
+    self.assertTrue(math.isnan(optimized_log2_neg.arg), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.arg}")
+    self.assertTrue(math.isinf(optimized_recip_zero.arg) and optimized_recip_zero.arg > 0,
+                    f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.arg}")

  @unittest.skip("broken")
  def test_full_graph_rewrite_modulo_negative_dividend(self):
@@ -182,27 +182,27 @@ class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
 class TestGEPAndVectorizeRewrite(unittest.TestCase):
  def test_gep_single_element_extraction(self):
    # GEP on a vector dtype to extract a single element
-    base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(base_vector.index(2)).val, 3.0)
+    base_vector = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
+    self.assertEqual(apply_rewrite(base_vector.index(2)).arg, 3.0)

  def test_gep_tuple_extraction(self):
    # GEP on a vector dtype to extract multiple elements as a vector
-    base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
+    base_vector = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
    self.assertEqual(list(apply_rewrite_values(UOp.stack(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0])

  def test_gep_on_const_stack(self):
    # GEP on a const STACK to extract a single element
-    const_stack = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(const_stack.index(2)).val, 3.0)
+    const_stack = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
+    self.assertEqual(apply_rewrite(const_stack.index(2)).arg, 3.0)

  def test_gep_tuple_on_const_stack(self):
    # GEP on a const STACK using a tuple to extract multiple elements
-    const_stack = UOp.const((7.0, 8.0, 9.0, 10.0))
+    const_stack = UOp.const(dtypes.float32, (7.0, 8.0, 9.0, 10.0))
    self.assertEqual(list(apply_rewrite_values(UOp.stack(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])

  def test_vectorize_multiple_elements(self):
    # Vectorizing multiple elements using GEP
-    base_vector = UOp.const((5.0, 10.0, 15.0, 20.0))
+    base_vector = UOp.const(dtypes.float32, (5.0, 10.0, 15.0, 20.0))
    vectorized_uop = UOp(Ops.STACK, src=tuple(base_vector.index(i) for i in range(4)))
    self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0])

@@ -213,7 +213,7 @@ from tinygrad.uop.symbolic import symbolic_simple

 class TestBottomUpRewrite(unittest.TestCase):
  def test_const_folding(self):
-    a = UOp.const(5)
+    a = UOp.const(dtypes.int, 5)
    ret = (a*3) + (a*7)
    gt = graph_rewrite(ret, symbolic_simple)
    ret = graph_rewrite(ret, symbolic_simple, bottom_up=True)
@@ -305,7 +305,7 @@ class TestRecurse(unittest.TestCase):
    graph_rewrite(a, pm, bottom_up=True)

  def test_inf_loop(self):
-    a = UOp.const(3)
+    a = UOp.const(dtypes.int, 3)
    pm = PatternMatcher([
      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
@@ -314,7 +314,7 @@ class TestRecurse(unittest.TestCase):
      graph_rewrite(a, pm)

  def test_inf_loop_bottom_up(self):
-    a = UOp.const(3)
+    a = UOp.const(dtypes.int, 3)
    pm = PatternMatcher([
      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
@@ -322,11 +322,11 @@ class TestRecurse(unittest.TestCase):
    with self.assertRaises(RuntimeError):
      graph_rewrite(a, pm, bottom_up=True)

-def bidir_append(ctx, x, b): ctx.append((x.val if x.op is Ops.CONST else "+", b))
+def bidir_append(ctx, x, b): ctx.append((x.arg if x.op is Ops.CONST else "+", b))
 class TestBidirectional(unittest.TestCase):
  def test_simple(self):
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
    c = a + b
    pm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, False)) ])
    bpm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, True)) ])
@@ -336,14 +336,14 @@ class TestBidirectional(unittest.TestCase):

 class TestStopEarly(unittest.TestCase):
  def test_stop_early(self):
-    a = UOp.const(3)
-    b = UOp.const(4)
+    a = UOp.const(dtypes.int, 3)
+    b = UOp.const(dtypes.int, 4)
    c = a+b
-    cn = UOp.const(7)
-    d = UOp.const(2)
+    cn = UOp.const(dtypes.int, 7)
+    d = UOp.const(dtypes.int, 2)
    def visit_const(c:UOp):
-      print(f"visit {c.val}")
-      assert c.val not in (3,4)
+      print(f"visit {c.arg}")
+      assert c.arg not in (3,4)
    pm_cvisit = PatternMatcher([(UPat(Ops.CONST, name="c"), visit_const),])
    ret = (c+d).substitute({c:cn}, extra_pm=pm_cvisit)
    assert ret == cn+d
@@ -376,7 +376,7 @@ class TestWalkRewrite(unittest.TestCase):

  def test_walk_topdown_no_fixed_point(self):
    """A bouncing pattern applies once and stops instead of looping."""
-    a = UOp.const(3)
+    a = UOp.const(dtypes.int, 3)
    pm = PatternMatcher([
      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
@@ -384,7 +384,7 @@ class TestWalkRewrite(unittest.TestCase):
    with self.assertRaises(RuntimeError):
      graph_rewrite(a, pm, bottom_up=True)
    ret = graph_rewrite(a, pm, walk=True)
-    self.assertIs(ret, UOp.const(4))
+    self.assertIs(ret, UOp.const(dtypes.int, 4))

  def test_walk_topdown_rewrites_children(self):
    a = UOp.variable('a', 0, 10)
@@ -418,11 +418,11 @@ class TestWalkRewrite(unittest.TestCase):
    """Top-down walk fires pm after children are processed (post-order)."""
    visited = []
    def track_visit(ctx, x):
-      ctx.append(x.val if x.op is Ops.CONST else x.op)
+      ctx.append(x.arg if x.op is Ops.CONST else x.op)
      return None
    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
    graph_rewrite(a + b, pm, ctx=visited, walk=True)
    self.assertEqual(visited, [1, 2, Ops.ADD])

@@ -454,23 +454,23 @@ class TestWalkRewrite(unittest.TestCase):

  def test_walk_bottomup_no_fixed_point(self):
    """Bottom-up walk also applies once per node, no fixed-point iteration."""
-    a = UOp.const(3)
+    a = UOp.const(dtypes.int, 3)
    pm = PatternMatcher([
      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
    ])
    ret = graph_rewrite(a, pm, bottom_up=True, walk=True)
-    self.assertIs(ret, UOp.const(4))
+    self.assertIs(ret, UOp.const(dtypes.int, 4))

  def test_walk_bottomup_visit_order(self):
    """Bottom-up walk fires bpm before descending (pre-order)."""
    visited = []
    def track_visit(ctx, x):
-      ctx.append(x.val if x.op is Ops.CONST else x.op)
+      ctx.append(x.arg if x.op is Ops.CONST else x.op)
      return None
    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
    graph_rewrite(a + b, pm, ctx=visited, bottom_up=True, walk=True)
    # bpm fires on each node before children: +, 1, 2
    self.assertEqual(visited, [Ops.ADD, 1, 2])
@@ -490,15 +490,15 @@ class TestWalkRewrite(unittest.TestCase):
    """Bidirectional walk: bpm fires pre-order, pm fires post-order."""
    visited = []
    def bpm_visit(ctx, x):
-      ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm"))
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm"))
      return None
    def pm_visit(ctx, x):
-      ctx.append((x.val if x.op is Ops.CONST else x.op, "pm"))
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
      return None
    bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_visit)])
    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_visit)])
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
    graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
    # bpm fires pre-order, pm fires post-order
    self.assertEqual(visited, [
@@ -509,23 +509,23 @@ class TestWalkRewrite(unittest.TestCase):
    """If bpm matches, children are skipped and pm never fires on that node."""
    visited = []
    def bpm_match(ctx, x):
-      ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm"))
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm"))
      # rewrite const(1) -> const(10), short-circuiting its subtree
-      if x.op is Ops.CONST and x.val == 1: return x.replace(arg=10)
+      if x.op is Ops.CONST and x.arg == 1: return x.replace(arg=10)
      return None
    def pm_match(ctx, x):
-      ctx.append((x.val if x.op is Ops.CONST else x.op, "pm"))
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
      return None
    bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_match)])
    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_match)])
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
    ret = graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
    # bpm matches const(1) and short-circuits it, so pm never fires on const(1)
    self.assertNotIn((1, "pm"), visited)
    # but pm still fires on const(2) and the rebuilt ADD
    self.assertIn((2, "pm"), visited)
-    self.assertIs(ret, UOp.const(10) + b)
+    self.assertIs(ret, UOp.const(dtypes.int, 10) + b)

 if __name__ == '__main__':
  unittest.main()
@@ -297,10 +297,10 @@ class TestPolyN(unittest.TestCase):
    from tinygrad.dtype import dtypes
    from tinygrad.uop.ops import UOp
    from test.helpers import eval_uop
-    np.testing.assert_allclose(eval_uop(polyN(UOp.const(1.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 0.0)
-    np.testing.assert_allclose(eval_uop(polyN(UOp.const(2.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 1.0)
-    np.testing.assert_allclose(eval_uop(polyN(UOp.const(3.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 4.0)
-    np.testing.assert_allclose(eval_uop(polyN(UOp.const(4.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 9.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 1.0), [1.0, -2.0, 1.0])), 0.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 2.0), [1.0, -2.0, 1.0])), 1.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 3.0), [1.0, -2.0, 1.0])), 4.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 4.0), [1.0, -2.0, 1.0])), 9.0)

 class TestTimeToStr(unittest.TestCase):
  def test_seconds(self):           self.assertEqual("   10.01s ", time_to_str(10.01))
@@ -8,15 +8,15 @@ from tinygrad.codegen import to_program
 class TestLinearizerFailures(unittest.TestCase):
  def test_fail_1(self):
    c0 = UOp.param(0, dtypes.float, (64,))
-    c1 = UOp.range(UOp.const(2), 1, AxisType.WEAK)
-    c2 = UOp.range(UOp.const(32), 2, AxisType.WEAK)
-    c3 = ((c1*UOp.const(32))+c2)
+    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.param(1, dtypes.float, (163840,))
-    c5 = UOp.range(UOp.const(2560), 0, AxisType.REDUCE)
-    c6 = c4.index(((((((c5//UOp.const(8))%UOp.const(8))*UOp.const(8))+(c5%UOp.const(8)))+(((c2*UOp.const(40))+(c5//UOp.const(64)))*UOp.const(64)))+(c1*UOp.const(81920))))
+    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.param(2, dtypes.float, (64,))
    c8 = c7.index(c3)
-    c9 = ((((c6+(c8*UOp.const(-1.0)))*(c6+(c8*UOp.const(-1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(0.000390625))+UOp.const(1e-05)).sqrt().reciprocal()
+    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)
    ast = c10.sink(arg=KernelInfo())
    to_program(ast, renderer=Device[Device.DEFAULT].renderer)
@@ -42,7 +42,7 @@ def check_assign(buffer_lists, copies=None):
  for orig_si, new_si in zip(linear.src, result.src):
    for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
      if new.op is Ops.SLICE and id(orig) not in replace_map:
-        replace_map[id(orig)] = (new.src[0], new.src[1].val * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)
+        replace_map[id(orig)] = (new.src[0], new.src[1].arg * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)

  # verify pinned buffers are not planned
  for buf in held_bufs:
@@ -1,5 +1,5 @@
 import unittest, time
-from tinygrad import Tensor, UOp, getenv
+from tinygrad import dtypes, Tensor, UOp, getenv
 from tinygrad.helpers import Profiling

 PYPROFILE = getenv("PYPROFILE")
@@ -27,29 +27,29 @@ class TestBench(unittest.TestCase):
    print(f"{self._testMethodName:30s} {et*1e6/self.N:.2f} us")

  def test_uop_instant_creation(self):
-    for i in range(self.N): UOp.const(100+i)
+    for i in range(self.N): UOp.const(dtypes.int, 100+i)

  def test_uop_list_creation(self):
-    [UOp.const(100+i) for i in range(self.N)]
+    [UOp.const(dtypes.int, 100+i) for i in range(self.N)]

  def test_uop_add_2n(self):
-    a = UOp.const(2)
+    a = UOp.const(dtypes.int, 2)
    for _ in range(self.N): a = a + a

  def test_uop_toposort(self):
-    a = UOp.const(0)
-    for i in range(self.N): a = a + UOp.const(100+i)
+    a = UOp.const(dtypes.int, 0)
+    for i in range(self.N): a = a + UOp.const(dtypes.int, 100+i)
    self.start_time()
    self.assertEqual(len(a.toposort()), 2*self.N+1)

  def test_uop_toposort_2n(self):
-    a = UOp.const(0)
+    a = UOp.const(dtypes.int, 0)
    for _ in range(self.N): a = a + a
    self.start_time()
    self.assertEqual(len(a.toposort()), self.N+1)

  def test_uop_simplify(self):
-    a = UOp.const(2)
+    a = UOp.const(dtypes.int, 2)
    for _ in range(self.N): (a+a).simplify()

  def test_uop_simplify_complex(self):
@@ -68,7 +68,7 @@ class TestBench(unittest.TestCase):
    for _ in range(self.N): expr.simplify()

  def test_uop_chain_free(self):
-    a = UOp.const(2)
+    a = UOp.const(dtypes.int, 2)
    for _ in range(self.N): a = a + a
    self.start_time()
    del a
@@ -1,5 +1,5 @@
 import gc, unittest
-from tinygrad import Tensor, UOp, GlobalCounters, dtypes
+from tinygrad import Tensor, GlobalCounters, dtypes
 from tinygrad.engine.jit import TinyJit
 from tinygrad.helpers import Context

@@ -181,7 +181,7 @@ class TestMultiScalarALU(unittest.TestCase):
    @functools.cache
    def _fxn(x_p, device):
      t = Tensor(x_p, device=device)
-      inner = Tensor(t.uop.src[0]) if t.uop.op is Ops.UNSHARD else t
+      inner = Tensor(t.uop.src[0]) if t.uop.op is Ops.MULTI else t
      return (inner.sum(),)
    param = x.as_param(0)
    fxn = _fxn(param.uop, x.device)
@@ -218,10 +218,5 @@ class TestMultiAxis(unittest.TestCase):
    self.assertEqual(e.uop.axis, 0)
    self.assertTrue(e.uop.has_buffer_identity())

-  def test_symbolic_reshape_shard_axis(self):
-    rows = UOp.variable("rows", 1, 4).bind(3)
-    x = Tensor.empty(4, 2).shard(("NULL:1", "NULL:2"), axis=1)[:rows]
-    self.assertEqual(x.reshape(rows, 1, 2).uop.axis, 2)
-
 if __name__ == '__main__':
  unittest.main()
@@ -5,9 +5,9 @@ from tinygrad.uop.ops import PatternMatcher, UPat

 class TestPatternMatcher(unittest.TestCase):
  def test_simple_match(self):
-    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.weakfloat), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(1)
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.int, 1)
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), None)

@@ -61,16 +61,16 @@ class TestPatternMatcher(unittest.TestCase):

  def test_uop(self):
    matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
+    c1 = UOp.const(dtypes.float, 1.0)
    c2 = UOp(Ops.ADD, src=(c1, c1))
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), None)

  def test_uop_set(self):
    matcher = PatternMatcher([(UPat((Ops.CONST, Ops.CAST), name="x"), lambda x: x.rtag())])
-    c1 = UOp.const(False)
+    c1 = UOp.const(dtypes.bool, False)
    c2 = UOp(Ops.CAST, arg=dtypes.int, src=(c1,))
-    c3 = UOp.const(1.0)
+    c3 = UOp.const(dtypes.float, 1.0)
    c4 = UOp(Ops.ADD, src=(c3, c3))
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), c2.rtag())
@@ -82,11 +82,11 @@ class TestPatternMatcher(unittest.TestCase):
      (UPat(Ops.CONST, arg=False, name="x"), lambda x: x.rtag()),
      (UPat(Ops.MAX, name="x"), lambda x: x.rtag()),
    ])
-    c1 = UOp.const(0.0)
-    c2 = UOp.const(False)
+    c1 = UOp.const(dtypes.float, 0.0)
+    c2 = UOp.const(dtypes.bool, False)
    c3 = UOp(Ops.MAX, src=(c1, c1))
    c4 = UOp(Ops.MUL, src=(c1, c1))
-    c5 = UOp.const(-1)
+    c5 = UOp.const(dtypes.int, -1)
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), c2.rtag())
    self.assertEqual(matcher.rewrite(c3), c3.rtag())
@@ -96,11 +96,11 @@ class TestPatternMatcher(unittest.TestCase):
  def test_filter_arg(self):
    matcher = PatternMatcher([
      (UPat(Ops.MUL, src=[UPat(Ops.CONST, name="c"), UPat(Ops.CONST, arg=2)], name="x"),
-       lambda x,c: x.rtag() if c.val in {1, -1} else None)
+       lambda x,c: x.rtag() if c.arg in {1, -1} else None)
    ])
-    y1 = UOp.const(1)
-    y2 = UOp.const(2)
-    y3 = UOp.const(-1)
+    y1 = UOp.const(dtypes.int, 1)
+    y2 = UOp.const(dtypes.int, 2)
+    y3 = UOp.const(dtypes.int, -1)
    c1 = UOp(Ops.MUL, src=(y1, y2))
    c2 = UOp(Ops.MUL, src=(y2, y2))
    c3 = UOp(Ops.MUL, src=(y3, y2))
@@ -114,27 +114,26 @@ class TestPatternMatcher(unittest.TestCase):

  def test_dup_name(self):
    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST, name="y"), UPat(Ops.CONST, name="y"))), lambda x, y: x.rtag())])
-    y1 = UOp.const(1.0)
-    y2 = UOp.const(1.0)
+    y1 = UOp.const(dtypes.float, 1.0)
+    y2 = UOp.const(dtypes.float, 1.0)
    c1 = UOp(Ops.ADD, src=(y1, y1))
    c2 = UOp(Ops.ADD, src=(y1, y2))
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), c1.rtag())

  def test_dtype(self):
-    # a concrete const dtype lives on the pair's CAST
-    matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype=dtypes.float32), lambda x: x.rtag())])
-    c1 = UOp.const(1.0).cast(dtypes.float32)
-    c2 = UOp.const(1.0).cast(dtypes.float64)
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float32), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float64, 1.0)
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), None)

  def test_dtype_set(self):
-    matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())])
-    c1 = UOp.const(1.0).cast(dtypes.float32)
-    c2 = UOp.const(1.0).cast(dtypes.float64)
-    c3 = UOp.const(1.0).cast(dtypes.float16)
-    c4 = UOp.const(1).cast(dtypes.int)
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float64, 1.0)
+    c3 = UOp.const(dtypes.float16, 1.0)
+    c4 = UOp.const(dtypes.int, 1)
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), c2.rtag())
    self.assertEqual(matcher.rewrite(c3), None)
@@ -142,8 +141,8 @@ class TestPatternMatcher(unittest.TestCase):

  def test_src_one(self):
    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(Ops.CONST))), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    c3 = UOp(Ops.ADD, src=(c1,c2))
    self.assertEqual(matcher.rewrite(c3), c3.rtag())
    self.assertEqual(matcher.rewrite(c2), None)
@@ -159,8 +158,8 @@ class TestPatternMatcher(unittest.TestCase):

  def test_src_permutations(self):
    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    c3 = UOp(Ops.ADD, src=(c1,c2))
    c4 = UOp(Ops.ADD, src=(c3,c2))
    c5 = UOp(Ops.ADD, src=(c2,c3))
@@ -172,8 +171,8 @@ class TestPatternMatcher(unittest.TestCase):

  def test_src_repeat(self):
    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=UPat(Ops.CONST)), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    c3 = UOp(Ops.ADD, src=(c1,c2))
    c4 = UOp(Ops.ADD, src=(c2,c3))
    self.assertEqual(matcher.rewrite(c3), c3.rtag())
@@ -181,9 +180,9 @@ class TestPatternMatcher(unittest.TestCase):

  def test_allow_len(self):
    matcher = PatternMatcher([(UPat(Ops.MULACC, name="x", src=(UPat(Ops.CONST),), allow_any_len=True), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
    c4 = UOp(Ops.EXP2, src=(c1,))
    c5 = UOp(Ops.ADD, src=(c1,c2))
    c6 = UOp(Ops.MULACC, src=(c1,c2,c3))
@@ -192,8 +191,8 @@ class TestPatternMatcher(unittest.TestCase):
    self.assertEqual(matcher.rewrite(c6), c6.rtag())

  def test_deep_src_permutations(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    u1 = (c1 + c2) + c1
    u2 = (c2 + c1) + c1
    matcher = PatternMatcher([
@@ -47,8 +47,12 @@ class TestRealWorld(unittest.TestCase):
    gc.collect()
    global global_mem_used
    global_mem_used = GlobalCounters.mem_used
+    self.old_float = dtypes.default_float
    np.random.seed(2002)

+  def tearDown(self):
+    dtypes.default_float = self.old_float
+
  @slow
  @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
  def test_stable_diffusion(self):
@@ -77,7 +81,7 @@ class TestRealWorld(unittest.TestCase):

  @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
  def test_llama(self):
-    self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
+    dtypes.default_float = dtypes.float16

    args_tiny = {"dim": 1024, "hidden_dim": 2048, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000}
    model = LLaMaTransformer(**args_tiny)
@@ -89,7 +93,7 @@ class TestRealWorld(unittest.TestCase):

  @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
  def test_gpt2(self):
-    self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
+    dtypes.default_float = dtypes.float16

    args_tiny = {"dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-5, "vocab_size": 1000}
    model = GPT2Transformer(**args_tiny)
@@ -146,7 +150,7 @@ class TestRealWorld(unittest.TestCase):

  @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
  def test_train_cifar_hyp(self):
-    self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
+    dtypes.default_float = dtypes.float16
    with Context(TRAINING=1):
      model = SpeedyResNet(Tensor.ones((12,3,2,2)))
      optimizer = optim.SGD(get_parameters(model), lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['bias_decay'])
@@ -600,15 +600,18 @@ class TestSchedule(unittest.TestCase):
    p = p.pad(((1, 0), ))
    p = p.repeat([2])
    # TODO: this should be 3 if fix store hazard worked correctly
-    check_schedule(p, 4)
+    # NOTE: pad now always has an explicit fill mask (internal PAD is Invalid-filled), which materializes here
+    check_schedule(p, 5)

  def test_conv2d(self, allowed=4, dtype=dtypes.float):
-    self.enterContext(Context(DEFAULT_FLOAT=dtype))
+    old_default_float, dtypes.default_float = dtypes.default_float, dtype
+    dtypes.default_float = dtype
    Tensor.manual_seed(0)
    BS, CIN = 2, 3
    img = Tensor.randn(BS, CIN, 64, 64).realize()
    w = Tensor.uniform(16, CIN, 3, 3).realize()
    ret = Tensor.conv2d(img, w).relu().mean().backward()
+    dtypes.default_float = old_default_float
    linear, var_vals = Tensor.linear_with_vars(ret, img.grad, w.grad)
    cnt = len([call for call in linear.src if call.src[0].op is Ops.SINK])
    assert cnt == allowed, f"expected {allowed} kernels, got {cnt}"
@@ -1640,11 +1643,11 @@ class TestSchedule(unittest.TestCase):
    self.assertEqual(GlobalCounters.mem_used-base, 0)

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

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

  def test_advanced_simple_indexing_combined(self):
@@ -2,7 +2,7 @@ import unittest, itertools

 from tinygrad.codegen.late.coalesce import indexing_simplify
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import UOp, Ops, graph_rewrite, pm_lower_index_dtype
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite
 from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
 from tinygrad.helpers import Context
 from test.helpers import full_rewrite
@@ -23,7 +23,7 @@ def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UO
    UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)),
  ))

-def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(nmax),), arg=expr)
+def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(dtypes.weakint, nmax),), arg=expr)
 def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
 def Range(n, nmax): return UOp.range(nmax, n)

@@ -207,7 +207,7 @@ class TestImageSimplification(unittest.TestCase):
    if svalid is not None:
      check_uop_against_string(self, off.src[1].get_valid(), svalid)
    else:
-      self.assertEqual(off.src[1].get_valid(), UOp.const(True), "svalid is None but valid is not True")
+      self.assertEqual(off.src[1].get_valid(), UOp.const(dtypes.bool, True), "svalid is None but valid is not True")

  def test_idx_gt_c(self):
    # (idx1 < c+1).ne(True) ? (..., idx1-1+c) : 0 can drop the valid
@@ -455,7 +455,7 @@ class TestImageSimplification(unittest.TestCase):
    A1 = lidx0*32 + r0*32 + lidx1*4 - 99
    valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 19)
    alu0 = gidx0 + (A1 % 32)*32 + (A1 // 32 % 16)*1024
-    load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(0)))
+    load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(dtypes.weakint, 0)))
    try:
      self.check(load, None, "(gidx0+lidx0*1024+r0*1024+lidx1*128+-3168)", "0")
    except AssertionError:
@@ -474,7 +474,7 @@ class TestImageSimplification(unittest.TestCase):
    A1 = lidx0*16 + r0*16 + lidx1*4 - 51
    valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 11)
    alu0 = lidx2 + gidx0*4 + (A1 % 16)*64 + (A1 // 16 % 8)*1024
-    load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(0)))
+    load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(dtypes.weakint, 0)))
    try:
      self.check(load, None, "(lidx2+gidx0*4+lidx0*1024+r0*1024+lidx1*256+-3264)", "0")
    except AssertionError:
@@ -488,20 +488,10 @@ class TestImageSimplification(unittest.TestCase):
    gidx0 = Special("gidx0", 1064)
    r12 = Range(12, 3)
    valid = ((gidx0 < 645).ne(True)) & (gidx0 < 653)
-    idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(0))
+    idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(dtypes.weakint, 0))
    load = get_load_image_uop((1, 48, 4), valid, idx)
    self.check(load, None, "(r12*4+(gidx0+3)%4+(gidx0+3)//4*24+-3888)", "0")

-  def test_drop_gate_committed_in_the_index_pass(self):
-    # the fused index pass runs without symbolic, so committing a weak src must not leave a CAST that
-    # symbolic later folds inside the index only: the gate's copy of the expression has to stay the same node
-    f = UOp.variable("f", 0.0, 9.0, dtypes.float)
-    idx_y = (f + UOp.const(1.0)).cast(dtypes.int)
-    load = get_load_image_uop((10, 10, 4), (UOp.const(-1) < idx_y) & (idx_y < UOp.const(10)),
-                              (Special("gidx0", 10), idx_y))
-    off = graph_rewrite(load.sink(), pm_lower_index_dtype+indexing_simplify, ctx={}).src[0].src[0]
-    self.assertEqual(off.src[1].get_valid(), UOp.const(True))
-
 class TestDropTrueGate(unittest.TestCase):
  def test_drop_true_gate_on_index(self):
    # test that INDEX with a constant True valid gets simplified to drop the valid
@@ -509,8 +499,8 @@ class TestDropTrueGate(unittest.TestCase):
    from tinygrad.uop.ops import graph_rewrite
    from tinygrad.uop.symbolic import sym
    buf = UOp.param(0, dtypes.int, (1,))
-    idx = UOp.const(0)
-    true_gate = UOp.const(True)
+    idx = UOp.const(dtypes.weakint, 0)
+    true_gate = UOp.const(dtypes.bool, True)
    index_with_gate = UOp(Ops.INDEX, src=(buf, idx.valid(true_gate)))
    # apply the optimization
    result = graph_rewrite(index_with_gate, sym+indexing_simplify)
@@ -526,50 +516,50 @@ class TestRangeShrink(unittest.TestCase):
  def test_range_shrink_single_guard(self):
    # range 0..203 guarded by r < 4 everywhere -> shrink to 0..3
    r = Range(0, 204)
-    load = get_gated_load_uop(r < UOp.const(4), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
    ranges = self.get_ranges(load.sink())
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assertEqual(ranges[0].src[0].arg, 4)

  def test_range_shrink_picks_max_guard(self):
    # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8
    r = Range(0, 204)
-    load1 = get_gated_load_uop(r < UOp.const(4), r)
-    load2 = get_gated_load_uop(r < UOp.const(8), r)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
+    load2 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 8), r)
    ranges = self.get_ranges(UOp.sink(load1, load2))
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 8)
+    self.assertEqual(ranges[0].src[0].arg, 8)

  def test_range_no_shrink_guard_ge_max(self):
    # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain)
    r = Range(0, 204)
-    load = get_gated_load_uop(r < UOp.const(300), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 300), r)
    ranges = self.get_ranges(load.sink())
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assertEqual(ranges[0].src[0].arg, 204)

  def test_range_no_shrink_when_unguarded_elsewhere(self):
    # 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(4), r)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
    load2 = UOp(Ops.LOAD, src=(UOp.param(1, dtypes.float, (204,)).index(r),))
    ranges = self.get_ranges(UOp.sink(load1, load2))
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assertEqual(ranges[0].src[0].arg, 204)

  def test_range_no_shrink_when_used_in_reduce(self):
    # range used in both a gated load AND directly in the reduce expression -> no shrink
    r = Range(0, 204)
-    gated_load = get_gated_load_uop(r < UOp.const(4), r)
+    gated_load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
    red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD)
    ranges = self.get_ranges(red.sink())
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assertEqual(ranges[0].src[0].arg, 204)

  def test_range_shrink_to_single_iteration(self):
    # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely
    r = Range(0, 204)
-    load = get_gated_load_uop(r < UOp.const(1), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 1), r)
    ranges = self.get_ranges(load.sink())
    self.assertEqual(len(ranges), 0)

@@ -577,19 +567,19 @@ class TestRangeShrink(unittest.TestCase):
    # emulates mask.where(x.pad_to(mask.shape), Invalid): range should shrink accordingly
    from tinygrad.dtype import Invalid
    r = Range(0, 204)
-    x = (r < 4).where(UOp.const(1.0), Invalid)
+    x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assertEqual(ranges[0].src[0].arg, 4)

  def test_range_shrink_store_where_invalid_flipped(self):
    # above, but flipped
    from tinygrad.dtype import Invalid
    r = Range(0, 204)
-    x = (r < 4).where(UOp.const(1.0), Invalid)
+    x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assertEqual(ranges[0].src[0].arg, 4)

 if __name__ == '__main__':
  unittest.main()
@@ -1,6 +1,5 @@
 import unittest
 from tinygrad import Variable
-from tinygrad.uop.ops import UOp


 class TestFuzzFailure(unittest.TestCase):
@@ -9,7 +8,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 2)
    v3=Variable('v3', 0, 1)
    expr = (((((((((((((((((((((((0//4)%2)//8)+-2)+-4)+-3)+v1)+-4)+v2)+-2)+v3)+v2)//3)%7)*1)//2)+v2)*-1)+2)+1)+0)+-3)+v3)
-    v1_val, v2_val, v3_val = UOp.const(8), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -20,7 +19,7 @@ class TestFuzzFailure(unittest.TestCase):
    v3=Variable('v3', 0, 3)
    expr = (((((((((((((((((((((((((0*4)//5)*2)*-1)*-2)+-4)*4)*2)*3)*4)+-4)*4)+v2)+v2)+v3)//3)+v2)+v1)//9)+3)+1)//1)+-4)//4)*2)
    expr = (((((v1+(v2+(((v3+(v2*2))+1)//3)))+4)//9)+-57)//(9*4))
-    v1_val, v2_val, v3_val = UOp.const(6), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(6), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -30,7 +29,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 1)
    v3=Variable('v3', 0, 2)
    expr = (((((((((((((((((((0//2)//3)+v3)+0)+-4)*-2)*-2)+-1)+2)+3)+v3)+0)//8)*-3)+0)*-2)*-4)*-2)//5)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -40,7 +39,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 3)
    v3=Variable('v3', 0, 4)
    expr = (((((((((((((((((((((((((((((0*-2)+0)*-1)//9)//6)//8)+v1)*-4)+v2)//4)//8)+4)*3)+v1)+v3)//8)//7)+4)+v3)*-4)+1)+v1)*3)+4)*2)//5)//2)//3)*-4)
-    v1_val, v2_val, v3_val = UOp.const(2), UOp.const(0), UOp.const(2)
+    v1_val, v2_val, v3_val = v1.const_like(2), v2.const_like(0), v3.const_like(2)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -50,7 +49,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 1)
    v3=Variable('v3', 0, 3)
    expr = ((((((((((((((0+v2)+v1)*0)+v2)//1)//7)+-2)+v2)+v1)*4)+-3)//5)+v2)+1)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -61,7 +60,7 @@ class TestFuzzFailure(unittest.TestCase):
    v3=Variable('v3', 0, 128)
    expr = (((((((((((((((((((((((((((((0//3)+4)+v1)//2)+-1)//1)*1)*-1)*4)//5)+v1)//6)+v1)*-1)+-4)+v2)+-2)*-3)+v3)+-4)+-2)*-1)//8)//4)*-4)+3)+v3)*
            -2)+v2)
-    v1_val, v2_val, v3_val = UOp.const(8), UOp.const(3), UOp.const(2)
+    v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(3), v3.const_like(2)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -71,7 +70,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 5)
    v3=Variable('v3', 0, 128)
    expr = (((((((((((((((((((((((((((((0+v2)*-4)+0)//9)+-4)*-2)*3)*4)//9)+v3)+v1)//4)+v1)+v3)+-1)*4)//4)+v2)//7)//3)+v1)+v2)+v3)+1)*2)//4)*3)+-1)*1)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(2), UOp.const(65)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(2), v3.const_like(65)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -81,7 +80,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 8)
    v3=Variable('v3', 0, 9)
    expr = (((((((0+-1)+2)+v1)*-2)//3)+v1)*-4)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -91,7 +90,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 1)
    v3=Variable('v3', 0, 8)
    expr = (((((((((((((((((((((((((((((0*-2)//1)+3)*-2)+-3)*-4)*1)+v1)+0)%2)%8)%9)+v2)%9)+-4)//4)+-1)*-2)+0)+v1)+v1)+3)+v1)+4)+-4)+0)*2)+-3)%6)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(1), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(1), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -102,7 +101,7 @@ class TestFuzzFailure(unittest.TestCase):
    v3=Variable("v3", 0, 32)
    x5 = (v1 <= 9).where(v1 * -4 - 4, v1 // 9) // 9
    expr = ((x5 >= -4).where(x5, (v2 % 3 + v2) // 5) * -1).maximum(((v1 * -2) % 6 + v3 % 1) * -1) * -1
-    v1_val, v2_val, v3_val = UOp.const(9), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(9), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -112,7 +111,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable("v2", 0, 128)
    v3=Variable("v3", 0, 5)
    expr = (((v2 * 0).maximum(8) - v2 * 2) % 5 + v1 // 6 + v1 + 5) % 5
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(7), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(7), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -181,7 +181,7 @@ class TestTensorPad(unittest.TestCase):
    t = Tensor.arange(9).reshape(1, 1, 3, 3)
    self.assertEqual(t.dtype, dtypes.int)
    r = t.pad((1, 2, 0, -1), value=-float('inf'))
-    self.assertEqual(r.dtype, dtypes.weakfloat)
+    self.assertEqual(r.dtype, dtypes.float)
    self.assertEqual(r.shape, (1, 1, 2, 6))

 class TestTensorDeviceMismatch(unittest.TestCase):
@@ -60,7 +60,7 @@ class TestTensorUOpClone(unittest.TestCase):
    t = _t(3, 4).float()
    self.assertIs(_strip_unique(t.clone().uop), _strip_unique(t.uop.clone()))
  def test_clone_deviceless_const(self):
-    u = UOp.const(2.0)
+    u = UOp.const(dtypes.float, 2.0)
    self.assertIs(_strip_unique(Tensor(u).clone().uop), _strip_unique(u.clone()))

 class TestTensorUOpGradient(unittest.TestCase):
@@ -382,7 +382,7 @@ class TestTensorUOpStack(unittest.TestCase):
    self.assertIs(_t(2, 3).uop.stack(w.uop).dtype, dtypes.float32)
  def test_stack_index_dtype(self):
    # index is outside the promotion lattice, equal dtypes bypass promotion
-    self.assertEqual(UOp.const(1).stack(UOp.const(2)).shape, (2,))
+    self.assertEqual(UOp.const(dtypes.weakint, 1).stack(UOp.const(dtypes.weakint, 2)).shape, (2,))

 class TestTensorUOpConv2d(unittest.TestCase):
  def test_conv2d_basic(self):
@@ -10,7 +10,7 @@ class TestTranscendentalFunctions(unittest.TestCase):
    # 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.param(1, dtypes.double, (1,))
-    loaded_value = input_buf.index(UOp.const(0)).load()
+    loaded_value = input_buf.index(UOp.const(dtypes.int, 0)).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))

@@ -27,48 +27,48 @@ class TestTranscendentalFunctions(unittest.TestCase):
    np.testing.assert_equal(q, 4)

  def test_cody_waite_reduction(self):
-    r, q = (eval_uop(u) for u in cody_waite_reduction(UOp.const(12 * math.pi + 0.1).cast(dtypes.float64)))
+    r, q = (eval_uop(u) for u in cody_waite_reduction(UOp.const(dtypes.float64, 12 * math.pi + 0.1)))
    np.testing.assert_allclose(r, 0.1)
    np.testing.assert_equal(q, 12)

  def test_frexp(self):
    for x in (1, -1):
-      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(float(x)).cast(dtypes.float64)))
+      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, x)))
      np.testing.assert_equal(mantissa, 0.5)
      np.testing.assert_equal(exponent, 1)

    for x in (2, -2):
-      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(2.0).cast(dtypes.float64)))
+      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 2.0)))
      np.testing.assert_equal(mantissa, 0.5)
      np.testing.assert_equal(exponent, 2)

-    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(5.0).cast(dtypes.float64)))
+    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 5.0)))
    np.testing.assert_equal(mantissa, 0.625)
    np.testing.assert_equal(exponent, 3)

-    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(1000.0).cast(dtypes.float64)))
+    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 1000.0)))
    np.testing.assert_allclose(mantissa, 0.9765625)
    np.testing.assert_equal(exponent, 10)

  def test_rintk(self):
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(0.0).cast(dtypes.float))), 0)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.0).cast(dtypes.float))), 5)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.5).cast(dtypes.float))), 6)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.999).cast(dtypes.float))), 6)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.0).cast(dtypes.float))), -5)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.5).cast(dtypes.float))), -6)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.999).cast(dtypes.float))), -6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 0.0))), 0)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.0))), 5)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.5))), 6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.999))), 6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.0))), -5)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.5))), -6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.999))), -6)

  def test_pow2if(self):
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(0).cast(dtypes.int), dtypes.float)), 1.0)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(1).cast(dtypes.int), dtypes.float)), 2.0)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(2).cast(dtypes.int), dtypes.float)), 4.0)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(10).cast(dtypes.int), dtypes.float)), 1024.0)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(63).cast(dtypes.int), dtypes.float)), 2**63)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-1).cast(dtypes.int), dtypes.float)), 0.5)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-2).cast(dtypes.int), dtypes.float)), 0.25)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-10).cast(dtypes.int), dtypes.float)), 2**-10)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-63).cast(dtypes.int), dtypes.float)), 2**-63)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 0), dtypes.float)), 1.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 1), dtypes.float)), 2.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 2), dtypes.float)), 4.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 10), dtypes.float)), 1024.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 63), dtypes.float)), 2**63)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -1), dtypes.float)), 0.5)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -2), dtypes.float)), 0.25)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -10), dtypes.float)), 2**-10)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -63), dtypes.float)), 2**-63)

 if __name__ == '__main__':
  unittest.main()
@@ -7,35 +7,35 @@ from tinygrad.uop.symbolic import sym
 from test.helpers import to_uops_list

 simple_pm = PatternMatcher([
-  (UPat.cvar('x', dtypes.weakint), lambda x: UOp.const(1.0) + UOp.const(2.0)),
-  (UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(x.val+y.val)),
-  (UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(x.val*y.val*z.val)),
-  ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.val+c2.val)),
+  (UPat.cvar('x', dtypes.int), lambda x: UOp.const(dtypes.float, 1.0) + UOp.const(dtypes.float, 2.0)),
+  (UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(dtypes.float, x.arg+y.arg)),
+  (UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(dtypes.float, x.arg*y.arg*z.arg)),
+  ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.arg+c2.arg)),
 ])

 def const_values(u:UOp):
-  if u.op is Ops.CONST: return (u.val,)
-  if u.op is Ops.STACK: return tuple(x.val for x in u.src)
+  if u.op is Ops.CONST: return (u.arg,) if not isinstance(u.arg, tuple) else u.arg
+  if u.op is Ops.STACK: return tuple(x.arg for x in u.src)
  raise AssertionError(f"expected const-like UOp, got {u.op}")

 class TestGraphRewriteConst(unittest.TestCase):
  def test_gep_const(self):
-    v1 = UOp.const((0,1,2), dtypes.int)
+    v1 = UOp.const(dtypes.int, (0,1,2))
    v2 = v1.index(1)
    ret = graph_rewrite(v2, sym)
    self.assertEqual(ret.dtype, dtypes.int)
-    self.assertEqual(ret.val, 1)
+    self.assertEqual(ret.arg, 1)

  def test_add_const(self):
-    v1 = UOp.const((0,1,2))
-    v2 = UOp.const((5,6,7))
+    v1 = UOp.const(dtypes.int, (0,1,2))
+    v2 = UOp.const(dtypes.int, (5,6,7))
    ret = graph_rewrite(v1+v2, sym)
    self.assertEqual(ret.op, Ops.STACK)
    self.assertEqual(const_values(ret), (5,7,9))

  def test_add_const_lose_v(self):
-    v1 = UOp.const((0,1,2))
-    v2 = UOp.const((2,1,0))
+    v1 = UOp.const(dtypes.int, (0,1,2))
+    v2 = UOp.const(dtypes.int, (2,1,0))
    ret = graph_rewrite(v1+v2, sym)
    self.assertEqual(ret.op, Ops.STACK)
    self.assertEqual(const_values(ret), (2,2,2))
@@ -49,47 +49,47 @@ class TestModularWraparound(unittest.TestCase):
    self.assertEqual(len(results), 2)  # +1 for SINK
    self.assertEqual(results[0].op, Ops.CONST)
    self.assertEqual(results[0].dtype, uop.dtype)
-    self.assertEqual(results[0].val, expected)
+    self.assertEqual(results[0].arg, expected)

  @xfail_broken_const_wraparound
  def test_cast(self):
    t = self._test
-    t(UOp.const(0xABCD17D6, dtypes.uint).cast(dtypes.uint8), 0xD6)
-    t(UOp.const(0xABCD17D6, dtypes.uint).cast(dtypes.uint8).cast(dtypes.uint), 0xD6)
+    t(UOp.const(dtypes.uint, 0xABCD17D6).cast(dtypes.uint8), 0xD6)
+    t(UOp.const(dtypes.uint, 0xABCD17D6).cast(dtypes.uint8).cast(dtypes.uint), 0xD6)

  @xfail_broken_const_wraparound
  def test_mul(self):
    t = self._test
-    t(UOp.const(0xABCD17D6, dtypes.uint) * 0xAABBCCDD, 1147018174)
-    t(UOp.const(0xABCD17D6, dtypes.int) * 10, -1241321892)
+    t(UOp.const(dtypes.uint, 0xABCD17D6) * 0xAABBCCDD, 1147018174)
+    t(UOp.const(dtypes.int, 0xABCD17D6) * 10, -1241321892)

  @xfail_broken_const_wraparound
  def test_div(self):
    t = self._test
-    t(UOp.const(0xABCD17D6, dtypes.uint) * 0xAABBCCDD // 11, 104274379)
-    t(UOp.const(0xABCD17D6, dtypes.int) * 10 // 11, -112847444)
+    t(UOp.const(dtypes.uint, 0xABCD17D6) * 0xAABBCCDD // 11, 104274379)
+    t(UOp.const(dtypes.int, 0xABCD17D6) * 10 // 11, -112847444)

  @xfail_broken_const_wraparound
  def test_neg(self):
    t = self._test
-    t(-UOp.const(1, dtypes.uint8), 0xFF)
-    t(-UOp.const(1, dtypes.uint16), 0xFFFF)
-    t(-UOp.const(1, dtypes.uint32), 0xFFFFFFFF)
-    t(-UOp.const(1, dtypes.uint64), 0xFFFFFFFFFFFFFFFF)
+    t(-UOp.const(dtypes.uint8, 1), 0xFF)
+    t(-UOp.const(dtypes.uint16, 1), 0xFFFF)
+    t(-UOp.const(dtypes.uint32, 1), 0xFFFFFFFF)
+    t(-UOp.const(dtypes.uint64, 1), 0xFFFFFFFFFFFFFFFF)

  @xfail_broken_const_wraparound
  def test_neg_min_int(self):
    t = self._test
-    t(-UOp.const(-2**7, dtypes.int8), -2**7)
-    t(-UOp.const(-2**15, dtypes.int16), -2**15)
-    t(-UOp.const(-2**31, dtypes.int32), -2**31)
-    t(-UOp.const(-2**63, dtypes.int64), -2**63)
+    t(-UOp.const(dtypes.int8, -2**7), -2**7)
+    t(-UOp.const(dtypes.int16, -2**15), -2**15)
+    t(-UOp.const(dtypes.int32, -2**31), -2**31)
+    t(-UOp.const(dtypes.int64, -2**63), -2**63)

  @xfail_broken_const_wraparound
  def test_payne_hanek_reduction_bug(self):
    t = self._test
-    a = (UOp.const(43748177600, dtypes.uint).cast(dtypes.uint) | 36).cast(dtypes.ulong)
-    b = 2536655455 * a + 4294967296 * UOp.const(25366554550, dtypes.ulong)
+    a = (UOp.const(dtypes.uint, 43748177600).cast(dtypes.uint) | 36).cast(dtypes.ulong)
+    b = 2536655455 * a + 4294967296 * UOp.const(dtypes.ulong, 25366554550)
    c = (b + 2261737165) // 4611686018427387904
    t(c, 0)

@@ -103,67 +103,67 @@ class TestGraphRewrite(unittest.TestCase):
  # NOTE: this shows why we can't have a UOp in arg
  @unittest.expectedFailure
  def test_no_dedup_args(self):
-    a1 = UOp.variable("a1", UOp.const(0), UOp.const(11), dtypes.int)
-    a2 = UOp.variable("a2", UOp.const(0), UOp.const(11), dtypes.int)
+    a1 = UOp.variable("a1", UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 11), dtypes.int)
+    a2 = UOp.variable("a2", UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 11), dtypes.int)
    sink = a1.sink(a2)
    variables = [x for x in graph_rewrite(sink, PatternMatcher([])).toposort() if x.op is Ops.PARAM and x.addrspace is AddrSpace.ALU]
    self.assertEqual(len(variables), 1)

  def test_simple(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    nout = graph_rewrite(c1+c2, simple_pm)
    self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 3.0)
+    self.assertEqual(nout.arg, 3.0)

  def test_depth_2_late(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
    nout = graph_rewrite(c1*c2*(c3+c3), simple_pm)
    self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 12.0)
+    self.assertEqual(nout.arg, 12.0)

  def test_double(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
    nout = graph_rewrite(c1+c2+c3, simple_pm)
    self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 6.0)
+    self.assertEqual(nout.arg, 6.0)

  def test_triple(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
-    c4 = UOp.const(4.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
+    c4 = UOp.const(dtypes.float, 4.0)
    nout = graph_rewrite(c1+c2+c3+c4, simple_pm)
    self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 10.0)
+    self.assertEqual(nout.arg, 10.0)

  def test_diamond(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
    nout = graph_rewrite((c1+c2)+(c1+c3), simple_pm)
    self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 7.0)
+    self.assertEqual(nout.arg, 7.0)

  def test_magic_4(self):
-    c1 = UOp.const(4)
+    c1 = UOp.const(dtypes.int, 4.0)
    nout = graph_rewrite(c1, simple_pm)
    self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 3.0)
+    self.assertEqual(nout.arg, 3.0)

  def test_depth_2_fold(self):
    v = UOp.variable("v", 0, 1, dtypes.float)
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    nout = graph_rewrite(v+c1+c2, simple_pm)
    self.assertEqual(nout.op, Ops.ADD)
    self.assertEqual(nout.src[0].op, Ops.PARAM)
    self.assertEqual(nout.src[1].op, Ops.CONST)
-    self.assertEqual(nout.src[1].val, 3.0)
+    self.assertEqual(nout.src[1].arg, 3.0)

  def test_commutative_work(self):
    a = UOp.variable('a', 0, 1)
@@ -174,14 +174,14 @@ class TestGraphRewrite(unittest.TestCase):
    a = UOp.variable('a', 0, 1)
    tst = (2+a).simplify()
    self.assertIs(tst.src[0], a)
-    self.assertIs(tst.src[1], UOp.const(2))
+    self.assertIs(tst.src[1], a.const_like(2))

  def test_consts_go_last(self):
    a = UOp.variable('a', 0, 1)
    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.const(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())
@@ -191,75 +191,67 @@ class TestGraphRewrite(unittest.TestCase):

 class TestUOpGraph(unittest.TestCase):
  def test_add_constant_fold(self):
-    c1 = UOp.const(1.0, dtypes.float)
-    c2 = UOp.const(2.0, dtypes.float)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    out = c1+c2
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 3.0)
+    self.assertEqual(out.arg, 3.0)

  def test_where_same_fold(self):
    v = UOp.variable('tmp', 0, 1)
-    c0 = UOp.const(0)
+    c0 = UOp.const(dtypes.weakint, 0)
    vc = v != c0
-    c1 = UOp.const(1.0, dtypes.float)
+    c1 = UOp.const(dtypes.float, 1.0)
    out = vc.where(c1, c1)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 1.0)
+    self.assertEqual(out.arg, 1.0)

  def test_where_const_fold(self):
-    bf = UOp.const(False)
-    c1 = UOp.const(1.0, dtypes.float)
-    c2 = UOp.const(2.0, dtypes.float)
+    bf = UOp.const(dtypes.bool, False)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    out = bf.where(c1, c2)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 2.0)
+    self.assertEqual(out.arg, 2.0)

  def test_const_cast(self):
-    bf = UOp.const(False)
+    bf = UOp.const(dtypes.bool, False)
    out = bf.cast(dtypes.int)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 0)
+    self.assertEqual(out.arg, 0)

  def test_const_bitcast(self):
-    bf = UOp.const(1.0, dtypes.float)
+    bf = UOp.const(dtypes.float, 1.0)
    out = bf.bitcast(dtypes.uint32)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 0x3F800000)
+    self.assertEqual(out.arg, 0x3F800000)

  @unittest.expectedFailure
  def test_const_shape_change_bitcast(self):
-    bf = UOp.const(0x3F).cast(dtypes.uint8)
+    bf = UOp.const(dtypes.uint8, 0x3F)
    out = bf.bitcast(dtypes.half)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK

-  def test_devectorize_derives_lane_dtype(self):
-    from tinygrad.codegen import do_devectorize
-    # an Invalid lane derives bool while the value lane derives float: the lane rebuild must derive, not inherit
-    lhs = UOp.stack(UOp.invalid(), UOp.const(1.0).cast(dtypes.float))
-    out = do_devectorize(lhs * lhs)
-    invalid_lane_mul = next(u for u in out.src[0].toposort() if u.op is Ops.MUL)
-    self.assertIs(invalid_lane_mul.dtype, dtypes.bool)
-
  @unittest.skip("this test isn't valid uops")
  def test_noop_vectorize_fold(self):
    d0 = UOp.param(0, dtypes.float, (1,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
    ld = d0.load(idx, dtype=dtypes.float)
    vec = UOp(Ops.STACK, dtypes.float, (ld,))
    x = vec.index(0)
@@ -273,7 +265,7 @@ class TestUOpGraph(unittest.TestCase):
    d0 = UOp.param(0, dtypes.float, (1,))
    d1 = UOp.param(1, dtypes.float, (1,))
    d2 = UOp.param(2, dtypes.float, (1,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
    def _test_vec(geps, count=4):
      vec = UOp(Ops.STACK, dtypes.float, geps)
      out = d0.index(idx).store(vec)
@@ -310,7 +302,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_gep_vec_const_fold(self):
    for vec_size in [2, 4, 8]:
-      consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)]
+      consts = [UOp.const(dtypes.float, float(i)) for i in range(vec_size)]
      vec = UOp(Ops.STACK, src=tuple(consts))
      with Context(SPEC=0):
        uops = to_uops_list([vec.index(i) for i in range(vec_size)])
@@ -320,7 +312,7 @@ class TestUOpGraph(unittest.TestCase):
  def test_cast_alu_fold(self):
    d0 = UOp.param(0, dtypes.bool, (1,))
    d1 = UOp.param(1, dtypes.int, (1,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
    ld = d1.index(idx)
    alu = (ld<1).cast(dtypes.bool)
    out = d0.index(idx).store(alu)
@@ -330,7 +322,7 @@ class TestUOpGraph(unittest.TestCase):
  def test_double_cast_fold(self):
    d0 = UOp.param(0, dtypes.float, (1,))
    d1 = UOp.param(1, dtypes.int, (1,))
-    idx = UOp.const(0, dtypes.int)
+    idx = UOp.const(dtypes.int, 0)
    ld = d1.index(idx)
    alu = ld.cast(dtypes.float).cast(dtypes.float)
    out = d0.index(idx).store(alu)
@@ -339,8 +331,8 @@ class TestUOpGraph(unittest.TestCase):

  def test_depth_2_const_fold(self):
    v = UOp.variable("tmp", 0, 1, dtypes.int)
-    c2 = UOp.const(2, dtypes.int)
-    c4 = UOp.const(4, dtypes.int)
+    c2 = UOp.const(dtypes.int, 2)
+    c4 = UOp.const(dtypes.int, 4)
    vc = v+c2
    out = vc+c4
    uops = to_uops_list([out])
@@ -348,19 +340,19 @@ class TestUOpGraph(unittest.TestCase):
    out = uops[-2]  # -2 to skip SINK
    self.assertEqual(out.op, Ops.ADD)
    self.assertEqual(out.src[1].op, Ops.CONST)
-    self.assertEqual(out.src[1].val, 6)
+    self.assertEqual(out.src[1].arg, 6)

  def test_bitcast_to_same_dtype_fold(self):
    for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
      d0 = UOp.param(0, dt, (1,))
-      v = d0.index(UOp.const(0))
+      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}")

  def test_sub_with_cast_folds(self):
    a = Variable("a", 0, 5)
    uops = to_uops_list([a.cast(dtypes.int)+(-a).cast(dtypes.int)])
-    assert uops[0] == UOp.const(0, dtypes.int)
+    assert uops[0] == UOp.const(dtypes.int, 0)
    assert uops[-1].op == Ops.SINK

  def test_where_on_gated_load_fold(self):
@@ -372,7 +364,7 @@ class TestUOpGraph(unittest.TestCase):
    uops = to_uops_list([out.index(ridx0).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].val==5
+      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)
@@ -382,7 +374,7 @@ class TestUOpGraph(unittest.TestCase):
    uops = to_uops_list([w])
    for u in uops:
      assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD: assert u.src[1].val==5
+      if u.op is Ops.LOAD: assert u.src[1].arg==5

  def test_where_on_gated_load_with_cast(self):
    ridx0 = UOp.range(100, 0)
@@ -394,13 +386,13 @@ class TestUOpGraph(unittest.TestCase):
    uops = to_uops_list([out.index(ridx0).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].val == 5
+      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.param(0, dtypes.float, (100,))
    ld = d0.index(ridx0.valid(ridx0<50))
-    w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(0.0)).cast(dtypes.half)
+    w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(dtypes.float, 0)).cast(dtypes.half)
    out = UOp.param(1, dtypes.half, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
    for u in uops:
@@ -410,7 +402,7 @@ class TestUOpGraph(unittest.TestCase):
    ridx0 = UOp.range(100, 0)
    d0 = UOp.param(0, dtypes.float, (100,))
    ld = d0.index(ridx0.valid(ridx0<50))
-    w = ((ridx0<50) & (ridx0>30)).where(UOp.const(0.0), ld).cast(dtypes.half)
+    w = ((ridx0<50) & (ridx0>30)).where(UOp.const(dtypes.float, 0), ld).cast(dtypes.half)
    out = UOp.param(1, dtypes.half, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
    for u in uops:
@@ -426,22 +418,22 @@ class TestUOpGraph(unittest.TestCase):
    uops = to_uops_list([st])
    for u in uops:
      assert u.op is not Ops.WHERE
-      if u.op is Ops.STORE: assert u.src[1].val==5
+      if u.op is Ops.STORE: assert u.src[1].arg==5

  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.param(0, dtypes.uchar, (128000,))
-    c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
-    c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
    c3 = UOp.param(1, dtypes.int, (512,))
    c4 = c3.index(c1)
-    c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
-    c6 = ((c2*UOp.const(240))+c5)
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
    c7 = UOp.param(2, dtypes.uchar, (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(250))+c2)).store(c9).end(c1, c2)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
    uops = to_uops_list([c10])
    for u in uops:
      self.assertNotEqual(u.dtype, dtypes.long)
@@ -449,19 +441,19 @@ 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.param(0, dtypes.uchar, (128000,))
-    c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
-    c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
    c3 = UOp.param(1, dtypes.int, (512,))
    c4 = c3.index(c1)  # c4 is a load
-    c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
-    c6 = ((c2*UOp.const(240))+c5)
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
    c7 = UOp.param(2, dtypes.uchar, (60000,))
    c8 = c7.index(c6)
    # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
    loaded_idx = c4.cast(dtypes.weakint)
-    comparison = (loaded_idx + c5) < UOp.const(60000)
+    comparison = (loaded_idx + c5) < UOp.const(dtypes.weakint, 60000)
    c9 = comparison.where(c8.cast(dtypes.uint).cast(dtypes.uchar), 0).reduce(c5, arg=Ops.ADD)
-    c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
    uops = to_uops_list([c10])
    for u in uops:
      self.assertNotEqual(u.dtype, dtypes.long)
@@ -470,9 +462,9 @@ class TestUOpGraph(unittest.TestCase):
    glbl0 = UOp.param(0, dtypes.int, (1,))
    glbl1 = UOp.param(1, dtypes.int, (1,))
    glbl2 = UOp.param(2, dtypes.int, (1,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
    ld0 = glbl1.index(UOp.invalid())
-    ld1 = glbl2.index(idx.valid(UOp.const(True)))
+    ld1 = glbl2.index(idx.valid(UOp.const(dtypes.bool, True)))
    uops = to_uops_list([glbl0.index(idx).store(ld1+ld0)])
    # the gate and invalid value are deleted from ld1
    self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 1)
@@ -484,7 +476,7 @@ class TestUOpGraph(unittest.TestCase):
    st = smem.index(lidx).store(glbl0.index(lidx).load())
    barrier = st.barrier()
    ld0 = smem.after(barrier).index(UOp.invalid())
-    ld1 = smem.after(barrier).index((lidx+2).valid(UOp.const(True)))
+    ld1 = smem.after(barrier).index((lidx+2).valid(UOp.const(dtypes.bool, True)))
    uops = to_uops_list([glbl0.index(lidx).store(ld1+ld0)])

    # the gate and invalid value are deleted from ld1
@@ -492,10 +484,10 @@ class TestUOpGraph(unittest.TestCase):

  def test_fold_gated_store(self):
    glbl = UOp.param(0, dtypes.int, (1,))
-    idx0 = UOp.const(0)
-    val = UOp.const(42)
+    idx0 = UOp.const(dtypes.int, 0)
+    val = UOp.const(dtypes.int, 42)
    st0 = glbl.index(UOp.invalid()).store(val)
-    st1 = glbl.index(idx0.valid(UOp.const(True))).store(val)
+    st1 = glbl.index(idx0.valid(UOp.const(dtypes.bool, True))).store(val)
    uops = to_uops_list([st0, st1])
    # only the second store happens
    self.assertEqual(len([u for u in uops if u.op is Ops.STORE]), 1)
@@ -503,9 +495,9 @@ class TestUOpGraph(unittest.TestCase):
  @unittest.skip("this is a uop type error")
  def test_asserts_bad_gate(self):
    glbl0 = UOp.param(0, dtypes.int, (1,))
-    idx = UOp.const(0)
-    bad_gate = UOp.const(1)
-    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(42), bad_gate))])
+    idx = UOp.const(dtypes.int, 0)
+    bad_gate = UOp.const(dtypes.int, 1)
+    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(dtypes.int, 42), bad_gate))])

  def test_after_end(self):
    r = UOp.range(10, 0)
@@ -513,7 +505,7 @@ class TestUOpGraph(unittest.TestCase):
    c = r + 1
    self.assertIn(r, c.ranges)

-    e = UOp.const(1).end(r)
+    e = UOp.const(dtypes.int, 1).end(r)
    self.assertNotIn(r, e.ranges)

    a = c.after(e)
@@ -563,7 +555,7 @@ class TestConstBufferize(unittest.TestCase):
    CONST doesn't depend on ranges (constant is same value everywhere).
    """
    from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
-    c = UOp.const(42.0)
+    c = UOp.const(dtypes.float, 42.0)
    r1 = UOp.range(3, 0)
    bufferize_with_range = UOp(Ops.STAGE, src=(c, r1), arg=BufferizeOpts(device="CPU"))
    self.assertEqual(len(bufferize_with_range.src), 2)  # const + 1 range
@@ -571,13 +563,13 @@ class TestConstBufferize(unittest.TestCase):
    result = graph_rewrite(bufferize_with_range, pm_const_buffer_folding, name='test')
    # BUFFERIZE should be removed, result is const broadcast to shape
    self.assertNotEqual(result.op, Ops.STAGE)
-    const_vals = [u.val for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat]
+    const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
    self.assertIn(42.0, const_vals)

  def test_const_bufferize_with_multiple_ranges(self):
    """Test CONST.BUFFERIZE with multiple ranges is also folded."""
    from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
-    c = UOp.const(3.14)
+    c = UOp.const(dtypes.float, 3.14)
    r1 = UOp.range(3, 0)
    r2 = UOp.range(4, 1)
    bufferize_with_ranges = UOp(Ops.STAGE, src=(c, r1, r2), arg=BufferizeOpts(device="CPU"))
@@ -586,14 +578,14 @@ class TestConstBufferize(unittest.TestCase):
    result = graph_rewrite(bufferize_with_ranges, pm_const_buffer_folding, name='test')
    # BUFFERIZE should be removed
    self.assertNotEqual(result.op, Ops.STAGE)
-    const_vals = [u.val for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat]
+    const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
    self.assertIn(3.14, const_vals)

 class TestUOpTags(unittest.TestCase):
  def test_inc_by_one(self):
-    g = UOp.const(1) + UOp.const(1)
+    g = UOp.const(dtypes.int, 1) + UOp.const(dtypes.int, 1)
    assert g.ssimplify() == 2
-    pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.replace(arg=x.val+1, tag=1) if x.tag is None else None)])
+    pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.replace(arg=x.arg+1, tag=1) if x.tag is None else None)])
    pm_strip_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
    g = graph_rewrite(g, pm_plus_1)
    assert g.ssimplify() == 4
@@ -682,36 +674,36 @@ class TestUOpGetItem(unittest.TestCase):

 class TestUOpBroadcast(unittest.TestCase):
  def test_broadcast_row(self):
-    a = UOp.const(1.0).expand((4, 8))
-    b = UOp.const(2.0).expand((4, 1))
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    b = UOp.const(dtypes.float, 2, shape=(4, 1))
    c = a + b
    self.assertEqual(c.shape, (4, 8))
    self.assertEqual(c.op, Ops.ADD)

  def test_broadcast_col(self):
-    a = UOp.const(1.0).expand((4, 8))
-    b = UOp.const(2.0).expand((1, 8))
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    b = UOp.const(dtypes.float, 2, shape=(1, 8))
    c = a + b
    self.assertEqual(c.shape, (4, 8))
    self.assertEqual(c.op, Ops.ADD)

  def test_broadcast_lower_dim(self):
-    a = UOp.const(1.0).expand((4, 8))
-    b = UOp.const(2.0).expand((8,))
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    b = UOp.const(dtypes.float, 2, shape=(8,))
    c = a * b
    self.assertEqual(c.shape, (4, 8))
    self.assertEqual(c.op, Ops.MUL)

  def test_broadcast_scalar(self):
-    a = UOp.const(1.0).expand((4, 8))
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
    c = a * 2
    self.assertEqual(c.shape, (4, 8))
    self.assertEqual(c.op, Ops.MUL)

  def test_broadcast_symbolic_same_shape(self):
    t = Variable("t", 1, 10)
-    a = UOp.const(1.0).expand((1, 1, t))
-    b = UOp.const(2.0).expand((1, 1, t))
+    a = UOp.const(dtypes.float, 1, shape=(1, 1, t))
+    b = UOp.const(dtypes.float, 2, shape=(1, 1, t))
    c = a + b
    self.assertEqual(c.op, Ops.ADD)

@@ -1,34 +1,34 @@
 import unittest
-from tinygrad import UOp
+from tinygrad import UOp, dtypes

 class TestUOpRepr(unittest.TestCase):
  def test_simple_const(self):
-    a = UOp.const(42)
-    self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.weakint, arg=42, src=())")
+    a = UOp.const(dtypes.int, 42)
+    self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.int, arg=42, src=())")
  def test_different_consts(self):
-    a, b = UOp.const(42), UOp.const(3)
+    a, b = UOp.const(dtypes.int, 42), UOp.const(dtypes.int, 3)
    expected = (
-      "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
-      "  UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
-      "  UOp(Ops.CONST, dtypes.weakint, arg=3, src=()),))"
+      "UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
+      "  UOp(Ops.CONST, dtypes.int, arg=42, src=()),\n" +
+      "  UOp(Ops.CONST, dtypes.int, arg=3, src=()),))"
    )
    self.assertEqual(repr(a+b), expected)
  def test_walrus_operator_indentation(self):
    # The reference should have the same indentation as the definition
-    a = UOp.const(42)
+    a = UOp.const(dtypes.int, 42)
    expected = (
-      "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
-      "  x0:=UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
+      "UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.CONST, dtypes.int, arg=42, src=()),\n" +
      "  x0,))"
    )
    self.assertEqual(repr(a+a), expected)
  def test_nested_walrus_indentation(self):
    # Ensure indentation is consistent at multiple levels
-    b = (a:=UOp.const(1)) + a
+    b = (a:=UOp.const(dtypes.int, 1)) + a
    expected = (
-      "UOp(Ops.MUL, dtypes.weakint, arg=None, src=(\n" +
-      "  x0:=UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
-      "    x1:=UOp(Ops.CONST, dtypes.weakint, arg=1, src=()),\n" +
+      "UOp(Ops.MUL, dtypes.int, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
+      "    x1:=UOp(Ops.CONST, dtypes.int, arg=1, src=()),\n" +
      "    x1,)),\n" +
      "  x0,))"
    )
@@ -4,48 +4,44 @@ from tinygrad.uop.ops import UOp, resolve

 class TestUOpResolve(unittest.TestCase):
  def test_simple_int(self):
-    u = UOp.const(4, dtypes.int)
+    u = UOp.const(dtypes.int, 4)
    self.assertEqual(int(u), 4)

-  def test_weak_const(self):
-    self.assertEqual(int(UOp.const(5)), 5)
-    self.assertEqual(float(UOp.const(1.5)), 1.5)
-
  def test_int_add(self):
-    u = UOp.const(4, dtypes.int) + 7
+    u = UOp.const(dtypes.int, 4) + 7
    self.assertEqual(int(u), 11)

  def test_lt(self):
-    u = UOp.const(4) < 7
+    u = UOp.const(dtypes.int, 4) < 7
    self.assertTrue(u)

  def test_rfloordiv(self):
-    u = 8 // UOp.const(4, dtypes.int)
+    u = 8 // UOp.const(dtypes.int, 4)
    self.assertEqual(int(u), 2)

  def test_rtruediv(self):
-    u = 9 / UOp.const(4, dtypes.float)
+    u = 9 / UOp.const(dtypes.float, 4)
    self.assertEqual(float(u), 2.25)

  def test_leq(self):
-    u = UOp.const(4) <= 4
+    u = UOp.const(dtypes.int, 4) <= 4
    self.assertTrue(u)

  def test_ne(self):
-    u = UOp.const(4) != 7
+    u = UOp.const(dtypes.int, 4) != 7
    self.assertTrue(u)

  def test_ne_f(self):
-    u = UOp.const(4) != 4
+    u = UOp.const(dtypes.int, 4) != 4
    self.assertFalse(u)

  def test_ngt(self):
-    u = UOp.const(4) > 7
+    u = UOp.const(dtypes.int, 4) > 7
    self.assertFalse(u)

  def test_ssimplify(self):
-    self.assertEqual((8 % UOp.const(4)).ssimplify(), 0)
-    self.assertEqual((8 * UOp.const(4)).ssimplify(), 32)
+    self.assertEqual((8 % UOp.const(dtypes.int, 4)).ssimplify(), 0)
+    self.assertEqual((8 * UOp.const(dtypes.int, 4)).ssimplify(), 32)

  def test_ambiguous_less_than(self):
    u = UOp.variable("i", 1, 10)
@@ -56,7 +52,7 @@ class TestUOpResolve(unittest.TestCase):
    self.assertFalse(resolve(u < -1, True))

  def test_float_direct(self):
-    u = UOp.const(4.5, dtypes.float) + 7
+    u = UOp.const(dtypes.float, 4.5) + 7
    self.assertEqual(float(u), 11.5)

  def test_var_cmp_t(self):
@@ -12,12 +12,13 @@ from tinygrad.uop.validate import uops_to_z3
 def check_uop_against_string(self, v:UOp, s:str):
  sym_vars = {v.render():v for v in v.toposort() if v.op in (Ops.RANGE, Ops.SPECIAL, Ops.PARAM)}
  s_eval = eval(s, sym_vars)
-  if isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(s_eval)
+  if isinstance(s_eval, int) and v.dtype==dtypes.weakint: s_eval = UOp.const(dtypes.weakint, s_eval)
+  elif isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(dtypes.from_py(s_eval), s_eval)
  s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval")
  self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v.render()} for {s}")

 def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint): return UOp.variable(name,min_val,max_val,dtype)
-def uconst(val): return UOp.const(val)
+def uconst(val): return UOp.const(dtypes.weakint, val)
 def usum(ops): return functools.reduce(lambda x,y: x+y, ops)
 def uand(ops): return functools.reduce(lambda x,y: x*y, ops)

@@ -103,16 +104,16 @@ class TestSymbolic(unittest.TestCase):
    self.assertEqual(UOp.gcd(a, a*b, a*3).simplify(), a)
    self.assertEqual(UOp.gcd(a*a*a, a*b*a, a*3*a).simplify(), a*a)
    self.assertEqual(UOp.gcd(a*a*10, b*a*5, a*a*5).simplify(), a*5)
-    self.assertEqual(UOp.gcd(a*10, b*5, a*5).simplify(), uconst(5))
-    self.assertEqual(UOp.gcd(a, b*5, a*5).simplify(), uconst(1))
+    self.assertEqual(UOp.gcd(a*10, b*5, a*5).simplify(), a.const_like(5))
+    self.assertEqual(UOp.gcd(a, b*5, a*5).simplify(), a.const_like(1))

  def test_divides_exact(self):
    a = Variable("a", 1, 8)
    b = Variable("b", 1, 8)
    self.assertEqual((a*a*3).divide_exact(a).simplify(), a*3)
-    self.assertEqual((a*a*3).divide_exact(a*a*3).simplify(), uconst(1))
-    self.assertEqual((a*a*6).divide_exact(a*a*3).simplify(), uconst(2))
-    self.assertEqual((a*b*3).divide_exact(uconst(3)).simplify(), a*b)
+    self.assertEqual((a*a*3).divide_exact(a*a*3).simplify(), a.const_like(1))
+    self.assertEqual((a*a*6).divide_exact(a*a*3).simplify(), a.const_like(2))
+    self.assertEqual((a*b*3).divide_exact(a.const_like(3)).simplify(), a*b)
    self.assertEqual((a*a*3).divide_exact(a*(-3)).simplify(), a*-1)
    self.assertEqual((a*a*b*3).divide_exact(a*b).simplify(), a*3)
    self.assertEqual((a*3+a*b).divide_exact(a).simplify(), b+3)
@@ -346,12 +347,6 @@ class TestSymbolic(unittest.TestCase):
  def test_mul_lt(self):
    self.helper_test_variable(Variable("a", 0, 5)*4 < 13, 0, 1, "(a<4)")
    self.helper_test_variable(Variable("a", 0, 5)*4 < 16, 0, 1, "(a<4)")
-    self.helper_test_variable(Variable("a", -5, 5)*4 < -13, 0, 1, "(a<-3)")
-    self.helper_test_variable(Variable("a", -5, 5)*-4 < 13, 0, 1, "((a*-1)<4)")
-    c0, c1 = 2, 2**54+1
-    self.helper_test_variable(Variable("a", 0, c1)*c0 < c1, 0, 1, f"(a<{2**53+1})")
-    c0, c1 = -2, -(2**54-1)
-    self.helper_test_variable(Variable("a", 0, -c1)*c0 < c1, 0, 1, f"((a*-1)<{-(2**53-1)})")
    self.helper_test_variable(Variable("a", 0, 5)*(-2) < 0, 0, 1, "((a*-1)<0)")
    self.helper_test_variable(Variable("a", 0, 5)*4 >= 12, 0, 1, "((a<3)!=True)")
    self.helper_test_variable(Variable("a", 0, 5)*4 >= 13, 0, 1, "((a<4)!=True)")
@@ -882,7 +877,7 @@ class TestSymbolic(unittest.TestCase):
    idx = Variable("idx", 0, 24)
    self.helper_test_variable(idx//4, 0, 6, "(idx//4)")
    # TODO: simplify the true branch
-    self.helper_test_variable((idx<4).where(idx//4, uconst(-1)), -1, 6, "(idx<4).where((idx//4), -1)")
+    self.helper_test_variable((idx<4).where(idx//4, idx.const_like(-1)), -1, 6, "(idx<4).where((idx//4), -1)")

  def test_floordiv_lt(self):
    # x//d<c <=> x<c*d for d>0, and <=> c*d<x for d<0
@@ -941,7 +936,7 @@ class TestSymbolic(unittest.TestCase):

  def test_where_removal(self):
    cond = Variable("a", 0, 3) < 2
-    u1, u0 = UOp.const(True), UOp.const(False)
+    u1, u0 = cond.const_like(True), cond.const_like(False)
    self.helper_test_variable(cond, 0, 1, "(a<2)")
    self.helper_test_variable(cond.where(u1, u0), 0, 1, "(a<2)")
    self.helper_test_variable(cond.where(u1, u0).where(u1, u0), 0, 1, "(a<2)")
@@ -1016,12 +1011,12 @@ class TestSymbolic(unittest.TestCase):

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

    # the vars are now scalar PARAMs
    pvar = {u.expr: u for u in rewritten_uop.toposort() if u.op is Ops.PARAM}
-    self.assertEqual(rewritten_uop, (pvar['s']<UOp.const(2, dtypes.int)).where(pvar['a'].cast(dtypes.half), pvar['b'].cast(dtypes.half)))
+    self.assertEqual(rewritten_uop, (pvar['s']<2).where(pvar['a'].cast(dtypes.half), pvar['b'].cast(dtypes.half)))

  def test_where_merge_branches(self):
    cond1 = Variable("s", 0, 10) < 6
@@ -1178,7 +1173,7 @@ class TestSymbolicVariables(unittest.TestCase):
    b = Variable("x", 1, 1).bind(1)
    s = b.simplify()
    self.assertEqual(s.op, Ops.CONST)
-    self.assertEqual(s.val, 1)
+    self.assertEqual(s.arg, 1)

 class TestSymInfer(unittest.TestCase):
  def test_sym_infer(self):
@@ -1202,13 +1197,9 @@ class TestSymInfer(unittest.TestCase):
    # floor: 1 % -1000 = -999, 1 // -1000 = -1
    assert sym_infer(a%b, var_vals) == -999
    assert sym_infer(a//b, var_vals) == -1
-  def test_sym_infer_with_cast(self):
-    a = Variable("a", 0, 100, dtypes.int)
-    assert sym_infer(a.cast(dtypes.long) + 1, {a.expr: 5}) == 6
-    assert sym_infer(a.cast(dtypes.float) * 0.5, {a.expr: 5}) == 2.5
  def test_sym_infer_with_bitcast(self):
    a = Variable("a", 1, 10, dtypes.int)
-    expr = ((a.bitcast(dtypes.uint) << UOp.const(1)).bitcast(dtypes.int) + 2)
+    expr = ((a.bitcast(dtypes.uint) << UOp.const(dtypes.uint, 1)).bitcast(dtypes.int) + 2)
    ret = sym_infer(expr, {a.expr: 2})
    assert isinstance(ret, int)
    assert ret == 6
@@ -1219,7 +1210,7 @@ class TestSymInfer(unittest.TestCase):
    c = Variable("c", 0, 0xFFFFFFFF, dtypes.uint)
    assert sym_infer(c.bitcast(dtypes.int), {c.expr: 0xFFFFFFFF}) == -1

-    assert sym_infer(UOp.const(1.5).cast(dtypes.float).bitcast(dtypes.uint), {}) == 1069547520
+    assert sym_infer(UOp.const(dtypes.float, 1.5).bitcast(dtypes.uint), {}) == 1069547520

  def test_sym_infer_deeply_nested(self):
    # build an expression that exceeds Python's nested parentheses limit for eval
@@ -1340,7 +1331,7 @@ class TestInvalidIndex(unittest.TestCase):
  def test_invalid_times_0(self):
    ridx = Variable("ridx", 0, 10)
    idx = (ridx<5).where(ridx, UOp.invalid())*0
-    self.assertIs(idx.simplify(), (ridx<5).where(uconst(0), UOp.invalid()),
+    self.assertIs(idx.simplify(), (ridx<5).where(UOp.const(dtypes.weakint, 0), UOp.invalid()),
                  "multiplying an index by 0 should preserve the invalid")

  def test_alu_moves_inside_invalid(self):
@@ -1365,24 +1356,24 @@ class TestInvalidIndex(unittest.TestCase):
    self.assertIs((UOp.invalid()<Variable("a",0,10)).simplify().dtype, dtypes.bool)

  def test_alu_invalid_vconst(self):
-    c1 = UOp.const((1, 1, Invalid, Invalid))
-    c2 = UOp.const((1, Invalid, 1, 1))
-    self.assertIs((c1+c2).simplify(), UOp.const((2, Invalid, Invalid, Invalid)))
+    c1 = UOp.const(dtypes.weakint, (1, 1, Invalid, Invalid))
+    c2 = UOp.const(dtypes.weakint, (1, Invalid, 1, 1))
+    self.assertIs((c1+c2).simplify(), UOp.const(dtypes.weakint, (2, Invalid, Invalid, Invalid)))

 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.param(0, dtypes.int, (1,))
-    index = buf.index(UOp.const(0))
+    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)
    # Emergent: store(idx, load(idx) + 0) -> store(idx, load(idx)) -> NOOP
-    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(0)), sym).op, Ops.NOOP)
+    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(dtypes.int, 0)), sym).op, Ops.NOOP)
    # Emergent: store(idx, load(idx) * 1) -> store(idx, load(idx)) -> NOOP
-    self.assertEqual(graph_rewrite(index.store(index.load() * UOp.const(1)), sym).op, Ops.NOOP)
+    self.assertEqual(graph_rewrite(index.store(index.load() * UOp.const(dtypes.int, 1)), sym).op, Ops.NOOP)
    # Negative: store(idx, load(idx) + 1) should NOT fold
-    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(1)), sym).op, Ops.STORE)
+    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(dtypes.int, 1)), sym).op, Ops.STORE)

 class TestMoveWhereOnLoad(unittest.TestCase):
  def test_bool_index_preserves_dtype(self):
@@ -1393,7 +1384,7 @@ class TestMoveWhereOnLoad(unittest.TestCase):
    cond = (a < 4) & (r < 2)
    valid = (a < 2)  # pre-existing valid on the load (to pass can_move check for the r-only clause)
    idx = buf.index(a.valid(valid))
-    expr = cond.where(idx, UOp.const(0))
+    expr = cond.where(idx, idx.const_like(0))
    out = graph_rewrite(expr, pm_move_where_on_load)
    type_verify(out, spec_shared)  # Invalid matches any dtype

@@ -1478,7 +1469,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 2)
    v3=Variable('v3', 0, 1)
    expr = (((((((((((((((((((((((0//4)%2)//8)+-2)+-4)+-3)+v1)+-4)+v2)+-2)+v3)+v2)//3)%7)*1)//2)+v2)*-1)+2)+1)+0)+-3)+v3)
-    v1_val, v2_val, v3_val = UOp.const(8), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    assert num==rn, f"{num} != {rn}"
@@ -5,12 +5,12 @@ from tinygrad.dtype import dtypes, Invalid
 class TestVminVmaxProperties(unittest.TestCase):
  def test_vmin_vmax_constant(self):
    # vmin and vmax for a constant
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
    self.assertEqual(uop.vmin, 42)
    self.assertEqual(uop.vmax, 42)

  def test_vmin_vmax_cmpne(self):
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
    def test_bool(u, x):
      self.assertEqual(u.vmin, x)
      self.assertEqual(u.vmax, x)
@@ -81,8 +81,8 @@ class TestVminVmaxProperties(unittest.TestCase):

  def test_vmin_vmax_multiplication_0_inf(self):
    # vmin and vmax for multiplication with a variable
-    x = UOp.const(0.0)
-    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0), dtype=dtypes.float)
+    x = UOp.const(dtypes.float, 0.0)
+    y = UOp.load(UOp.param(0, dtypes.float, (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)
@@ -167,7 +167,7 @@ class TestVminVmaxProperties(unittest.TestCase):
    self.assertNotEqual(i.vmin, i.vmax)

  def test_vmin_vmax_invalid_vconst(self):
-    x = UOp.const((0, 4, Invalid, Invalid))
+    x = UOp.const(dtypes.weakint, (0, 4, Invalid, Invalid))
    self.assertEqual((x.vmin, x.vmax), (0, 4))

 class TestVminVmaxDivMod(unittest.TestCase):
@@ -198,14 +198,14 @@ class TestVminVmaxDivMod(unittest.TestCase):

  def test_vmin_vmax_floordiv_floormod(self):
    x = UOp.variable('x', -7, 7)
-    floordiv = x.alu(Ops.FLOORDIV, UOp.const(3))
+    floordiv = x.alu(Ops.FLOORDIV, x.const_like(3))
    self.assertEqual(floordiv.vmin, -3)
    self.assertEqual(floordiv.vmax, 2)
-    floormod = x.alu(Ops.FLOORMOD, UOp.const(3))
+    floormod = x.alu(Ops.FLOORMOD, x.const_like(3))
    self.assertEqual(floormod.vmin, 0)
    self.assertEqual(floormod.vmax, 2)
    # negative const divisor: floormod range is [c+1, 0]
-    floormod_neg = x.alu(Ops.FLOORMOD, UOp.const(-3))
+    floormod_neg = x.alu(Ops.FLOORMOD, x.const_like(-3))
    self.assertEqual(floormod_neg.vmin, -2)
    self.assertEqual(floormod_neg.vmax, 0)

@@ -286,44 +286,44 @@ class TestVminVmaxDivMod(unittest.TestCase):
 class TestVminVmaxVConst(unittest.TestCase):
  def test_vmin_vmax_vconst_single_element(self):
    # vmin and vmax for a single-element vector constant
-    uop = UOp.const((42,))
+    uop = UOp.const(dtypes.int32, (42,))
    self.assertEqual(uop.vmin, 42)
    self.assertEqual(uop.vmax, 42)

  def test_vmin_vmax_vconst_multiple_elements(self):
    # vmin and vmax for a multi-element vector constant
-    uop = UOp.const((10, 20, -5, 7))
+    uop = UOp.const(dtypes.int32, (10, 20, -5, 7))
    self.assertEqual(uop.vmin, -5)
    self.assertEqual(uop.vmax, 20)

  def test_vmin_vmax_vconst_all_equal(self):
    # vmin and vmax for a vector where all elements are equal
-    uop = UOp.const((7, 7, 7))
+    uop = UOp.const(dtypes.int32, (7, 7, 7))
    self.assertEqual(uop.vmin, 7)
    self.assertEqual(uop.vmax, 7)

  def test_vmin_vmax_vconst_with_negative_values(self):
    # vmin and vmax for a vector constant containing negative values
-    uop = UOp.const((-10, -20, -5, -15))
+    uop = UOp.const(dtypes.int32, (-10, -20, -5, -15))
    self.assertEqual(uop.vmin, -20)
    self.assertEqual(uop.vmax, -5)

  def test_vmin_vmax_vconst_with_floats(self):
    # vmin and vmax for a vector constant of float values
-    uop = UOp.const((1.5, -3.2, 0.0))
+    uop = UOp.const(dtypes.float32, (1.5, -3.2, 0.0))
    self.assertEqual(uop.vmin, -3.2)
    self.assertEqual(uop.vmax, 1.5)

  def test_vmin_vmax_vconst_with_bools(self):
    # vmin and vmax for a vector constant of bool values
-    uop = UOp.const((True, False, False))
+    uop = UOp.const(dtypes.bool, (True, False, False))
    self.assertIs(uop.vmin, False)
    self.assertIs(uop.vmax, True)

  def test_vmin_vmax_vector_with_gep(self):
    # vmin and vmax for a vector constant of bool values
    d1 = UOp.param(1, dtypes.int, (1,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
    val = UOp(Ops.LOAD, src=(d1.index(idx),))
    uop = (val // 32)
    self.assertEqual(uop.vmin, -67108864)
@@ -332,17 +332,17 @@ class TestVminVmaxVConst(unittest.TestCase):
 class TestConstFactor(unittest.TestCase):
  def test_const_factor_constant(self):
    # const_factor for a constant
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
    self.assertEqual(uop.const_factor(), 42)

  def test_const_factor_addition(self):
    # const_factor for an addition of constants
-    uop = UOp.const(30) + UOp.const(12)
+    uop = UOp.const(dtypes.int32, 30) + UOp.const(dtypes.int32, 12)
    self.assertEqual(uop.const_factor(), 6)  # GCD(30, 12) = 6

  def test_const_factor_multiplication(self):
    # const_factor for a multiplication of constants
-    uop = UOp.const(5) * UOp.const(7)
+    uop = UOp.const(dtypes.int32, 5) * UOp.const(dtypes.int32, 7)
    self.assertEqual(uop.const_factor(), 5)  # For multiplication, it's one of the factors

  def test_const_factor_with_variable(self):
@@ -377,14 +377,14 @@ class TestConstFactor(unittest.TestCase):
 class TestDivides(unittest.TestCase):
  def test_divides_constant_exact(self):
    # Divides a constant by an exact divisor
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
    result = uop.divides(7)
    self.assertIsNotNone(result)
    self.assertEqual(result.const_factor(), 6)  # 42 / 7 = 6

  def test_divides_constant_inexact(self):
    # Try to divide a constant by a non-exact divisor
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
    result = uop.divides(5)
    self.assertIsNone(result)  # 42 is not divisible by 5

--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
George HotzandGitHub	1a523193dd	Merge branch 'master' into invalid_try_3	2026-07-25 14:07:16 -07:00
geohot	e70a79465e	invalid try 3 (kimi)	2026-07-25 14:06:29 -07:00