Clean up postrange.py by removing comments

does that pass?
one call to hc opt
2025-09-10 09:40:20 +08:00 · 2025-09-08 11:12:46 -07:00 · 2025-09-08 10:30:04 -07:00 · 2025-09-08 18:28:09 +03:00 · 2025-09-08 17:37:17 +03:00 · 2025-09-08 10:17:03 -04:00
60 changed files with 612 additions and 1007 deletions
@@ -0,0 +1,91 @@
+name: Autogen
+env:
+  # increment this when downloads substantially change to avoid the internet
+  DOWNLOAD_CACHE_VERSION: '12'
+  PYTHON_CACHE_VERSION: '3'
+  APT_CACHE_VERSION: '1'
+  BUILD_CACHE_VERSION: '1'
+  CAPTURE_PROCESS_REPLAY: 1
+  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  PYTHONPATH: ${{ github.workspace }}
+
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+    paths: 
+    - 'tinygrad/runtime/autogen/**/*'
+  workflow_dispatch:
+    paths: 
+    - 'tinygrad/runtime/autogen/**/*'
+
+jobs:
+  autogen:
+    name: Autogen
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        opencl: 'true'
+        amd: 'true'
+        cuda: 'true'
+        webgpu: 'true'
+        llvm: 'true'
+    - name: Install autogen support packages
+      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev
+    - name: Verify OpenCL autogen
+      run: |
+        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
+        ./autogen_stubs.sh opencl
+        diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
+    - name: Verify CUDA autogen
+      run: |
+        cp tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
+        cp tinygrad/runtime/autogen/nv_gpu.py /tmp/nv_gpu.py.bak
+        ./autogen_stubs.sh cuda
+        ./autogen_stubs.sh nv
+        diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
+        diff /tmp/nv_gpu.py.bak tinygrad/runtime/autogen/nv_gpu.py
+    - name: Verify AMD autogen
+      run: |
+        cp tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
+        cp tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
+        cp tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
+        cp tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
+        cp tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
+        ./autogen_stubs.sh hsa
+        ./autogen_stubs.sh kfd
+        ./autogen_stubs.sh comgr
+        ./autogen_stubs.sh amd
+        ./autogen_stubs.sh sqtt
+        diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
+        diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
+        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
+        diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
+        diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
+    - name: Verify Linux autogen
+      run: |
+        cp tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
+        cp tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
+        cp tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
+        ./autogen_stubs.sh libc
+        ./autogen_stubs.sh io_uring
+        ./autogen_stubs.sh ib
+        diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
+        diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
+        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
+    - name: Verify WebGPU autogen
+      run: |
+        cp tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
+        ./autogen_stubs.sh webgpu
+        diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
+    - name: Verify LLVM autogen
+      run: |
+        cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
+        ./autogen_stubs.sh llvm
+        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
@@ -89,75 +89,6 @@ jobs:
        clang -O2 recognize.c -lm -o recognize
        cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock

-  autogen:
-    name: Autogen
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        opencl: 'true'
-        amd: 'true'
-        cuda: 'true'
-        webgpu: 'true'
-        llvm: 'true'
-    - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev
-    - name: Verify OpenCL autogen
-      run: |
-        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
-        ./autogen_stubs.sh opencl
-        diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
-    - name: Verify CUDA autogen
-      run: |
-        cp tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
-        cp tinygrad/runtime/autogen/nv_gpu.py /tmp/nv_gpu.py.bak
-        ./autogen_stubs.sh cuda
-        ./autogen_stubs.sh nv
-        diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
-        diff /tmp/nv_gpu.py.bak tinygrad/runtime/autogen/nv_gpu.py
-    - name: Verify AMD autogen
-      run: |
-        cp tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
-        cp tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
-        cp tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
-        cp tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
-        cp tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
-        ./autogen_stubs.sh hsa
-        ./autogen_stubs.sh kfd
-        ./autogen_stubs.sh comgr
-        ./autogen_stubs.sh amd
-        ./autogen_stubs.sh sqtt
-        diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
-        diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
-        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
-        diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
-        diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
-    - name: Verify Linux autogen
-      run: |
-        cp tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
-        cp tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
-        cp tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
-        ./autogen_stubs.sh libc
-        ./autogen_stubs.sh io_uring
-        ./autogen_stubs.sh ib
-        diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
-        diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
-        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
-    - name: Verify WebGPU autogen
-      run: |
-        cp tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
-        ./autogen_stubs.sh webgpu
-        diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
-    - name: Verify LLVM autogen
-      run: |
-        cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
-        ./autogen_stubs.sh llvm
-        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
-
  torchbackend:
    name: Torch Backend Tests
    runs-on: ubuntu-latest
@@ -372,7 +303,7 @@ jobs:
        CAPTURE_PROCESS_REPLAY=1 python test/test_tiny.py TestTiny.test_plus
        python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
-        DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
+        #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
    - name: Repo line count < 18000 lines
      run: MAX_LINE_COUNT=18000 python sz.py

@@ -532,8 +463,8 @@ jobs:
          opencl: 'true'
      - name: Test ONNX (GPU)
        run: GPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
-      - name: Test Optimization Helpers
-        run: DEBUG=1 python3 extra/optimization/test_helpers.py
+      #- name: Test Optimization Helpers
+      #  run: DEBUG=1 python3 extra/optimization/test_helpers.py
      #- name: Test Action Space
      #  run: DEBUG=1 GPU=1 python3 extra/optimization/get_action_space.py
      - name: Test Beam Search
@@ -608,8 +539,9 @@ jobs:
          test/test_outerworld_range.py test/test_sample.py test/test_randomness.py
    - name: Test CPU=1 RANGEIFY=2
      run: CPU=1 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
-    - name: Test LLVM=1 RANGEIFY=1 (slow tests)
-      run: LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20
+    # slow (and still wrong on beautiful_mnist)
+    #- name: Test LLVM=1 RANGEIFY=1 (slow tests)
+    #  run: LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20

  testdevectorize:
    name: Linux (devectorize)
@@ -798,7 +730,7 @@ jobs:
          opencl: ${{ matrix.backend == 'gpu' && 'true' }}
          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'gpu' && 'GPU=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_COUNT=2' || matrix.backend == 'gpu' && 'GPU=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['LLVM','CPU','GPU'], Device.DEFAULT"
@@ -1022,7 +954,7 @@ jobs:
          pydeps: "capstone"
          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == '${{ matrix.backend }}'.upper(), Device.DEFAULT"
@@ -1059,7 +991,7 @@ jobs:
          pydeps: ${{ matrix.backend == 'webgpu' && 'dawn-python' || '' }}
      - name: Set env
        shell: bash
-        run:  printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'webgpu' && 'WEBGPU=1'}}" >> $GITHUB_ENV
+        run:  printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_COUNT=2' || matrix.backend == 'webgpu' && 'WEBGPU=1'}}" >> $GITHUB_ENV
      - name: Run unit tests
        if: matrix.backend=='llvm'
        # test_newton_schulz hits RecursionError
@@ -10,13 +10,13 @@ from tinygrad.renderer.cstyle import ClangRenderer
 render_dtype = ClangRenderer().render_dtype

 class ClangGraph(GraphRunner):
-  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, int]):
    super().__init__(jit_cache, input_rawbuffers, var_vals)
    if not all(isinstance(ji.prg, CompiledRunner) for ji in jit_cache): raise GraphException

    prgs = '\n'.join(dedup([cast(CompiledRunner, ji.prg).p.src for ji in jit_cache]))
    args = [f"{render_dtype(x.dtype)}* arg{i}" for i,x in enumerate(input_rawbuffers)]
-    args += sorted([f"int {v.expr}" for v in var_vals])
+    args += sorted([f"int {v}" for v in var_vals])
    code = ["void batched("+','.join(args)+") {"]
    for ji in jit_cache:
      args = []
@@ -34,6 +34,6 @@ class ClangGraph(GraphRunner):
    assert compiler is not None
    self._prg = ClangProgram("batched", compiler.compile(prgs+"\n"+"\n".join(code))) # no point in caching the pointers

-  def __call__(self, rawbufs: List[Buffer], var_vals: Dict[Variable, int], wait=False):
+  def __call__(self, rawbufs: List[Buffer], var_vals: Dict[str, int], wait=False):
    return cpu_time_execution(
-    lambda: self._prg(*[x._buf for x in rawbufs], *[x[1] for x in sorted(var_vals.items(), key=lambda x: x[0].expr)]), enable=wait)
+    lambda: self._prg(*[x._buf for x in rawbufs], *[x[1] for x in sorted(var_vals.items(), key=lambda x: x[0])]), enable=wait)
@@ -26,7 +26,7 @@ class VirtAQLQueue(AQLQueue):
    self.available_packet_slots -= 1

 class HSAGraph(MultiGraphRunner):
-  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, int]):
    super().__init__(jit_cache, input_rawbuffers, var_vals)

    # Check all jit items are compatible.
@@ -53,7 +53,7 @@ class HSAGraph(MultiGraphRunner):
      self.ji_kargs_structs[j] = ji.prg._prg.args_struct_t.from_address(kernargs_ptrs[ji.prg.dev])
      kernargs_ptrs[ji.prg.dev] += round_up(ctypes.sizeof(ji.prg._prg.args_struct_t), 16)
      for i in range(len(ji.bufs)): self.ji_kargs_structs[j].__setattr__(f'f{i}', cast(Buffer, ji.bufs[i])._buf)
-      for i in range(len(ji.prg.p.vars)): self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[ji.prg.p.vars[i]])
+      for i in range(len(ji.prg.p.vars)): self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[ji.prg.p.vars[i].expr])

    # Build queues.
    self.virt_aql_queues: Dict[Compiled, VirtAQLQueue] = {dev:VirtAQLQueue(dev, 2*len(self.jit_cache)+16) for dev in self.devices}
@@ -106,7 +106,7 @@ class HSAGraph(MultiGraphRunner):
    for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 0)
    hsa.hsa_signal_silent_store_relaxed(self.finish_signal, 0)

-  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int], wait=False) -> Optional[float]:
+  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[str, int], wait=False) -> Optional[float]:
    # Wait and restore signals
    hsa.hsa_signal_wait_scacquire(self.finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
    for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 1)
@@ -123,7 +123,7 @@ class HSAGraph(MultiGraphRunner):
    # Update var_vals
    for j in self.jc_idx_with_updatable_var_vals:
      for i,v in enumerate(cast(CompiledRunner, self.jit_cache[j].prg).p.vars):
-        self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v])
+        self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v.expr])

    # Update launch dims
    for j in self.jc_idx_with_updatable_launch_dims:
@@ -67,11 +67,12 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
  forward_args = ",".join(f"{dtype}{'*' if name not in symbolic_vars.values() else ''} {name}" for name,dtype,_ in (outputs+inputs if wasm else inputs+outputs))

  if not wasm:
+    thread_id = 0 # NOTE: export does not support threading, thread_id is always 0
    for name,cl in bufs_to_save.items():
      weight = ''.join(["\\x%02X"%x for x in bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
      cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
    cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
-    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
+    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)}, {thread_id});" for (name, args, _global_size, _local_size) in statements] + ["}"]
    return '\n'.join(headers + cprog)
  else:
    if bufs_to_save:
@@ -239,7 +240,9 @@ 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"
-  with Context(JIT=2): run,special_names = jit_model(model, *inputs)
+
+  # NOTE: CPU_COUNT=1, since export does not support threading
+  with Context(JIT=2, CPU_COUNT=1): run,special_names = jit_model(model, *inputs)
  functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
  state = get_state_dict(model)
  weight_names = {id(x.uop.base.realized): name for name, x in state.items()}
@@ -2,7 +2,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
-from tinygrad.codegen.opt.kernel import OptOps
+from tinygrad.codegen.opt import OptOps
 from tinygrad.engine.realize import lower_schedule

 dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
@@ -88,7 +88,7 @@ def mcts_search(lin:Kernel, rawbufs:List[Buffer], amt:int) -> Kernel:
    return ret

  rawbufs = _ensure_buffer_alloc(rawbufs)
-  var_vals = {k:(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
+  var_vals = {k.expr:(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
  dev = Device[lin.opts.device]
  root = MCTSNode(lin)

@@ -1,6 +1,6 @@
 # stuff needed to unpack a kernel
 from tinygrad import Variable
-from tinygrad.codegen.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -115,7 +115,7 @@ def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_
  assert dev.compiler is not None

  rawbufs = _ensure_buffer_alloc(rawbufs)
-  var_vals: dict[Variable, int] = {k:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
+  var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
  p = get_program(lin.get_optimized_ast(), lin.opts)
  tms = _time_program(p, dev.compiler.compile(p.src), var_vals, rawbufs,
                      max_global_size=max_global_size if allow_test_size else None, clear_l2=clear_l2, cnt=cnt, name=to_function_name(lin.name))
@@ -2,7 +2,6 @@ import itertools
 from enum import Enum, auto
 from collections import defaultdict
 from typing import List, Tuple, DefaultDict
-from extra.optimization.helpers import load_worlds, ast_str_to_ast
 from tinygrad.helpers import prod, tqdm
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -121,7 +120,7 @@ def st_equivalent(st1: ShapeTracker, st2: ShapeTracker):
    if i > 1000:
      print("WARNING: did not search all possible combinations")
      break
-    var_vals = {k:v for k,v in zip(vs, ranges)}
+    var_vals = {k.expr:v for k,v in zip(vs, ranges)}
    r1 = sym_infer(idx1, var_vals) if sym_infer(valid1, var_vals) else 0
    r2 = sym_infer(idx2, var_vals) if sym_infer(valid2, var_vals) else 0
    if r1 != r2: return False
@@ -147,6 +146,7 @@ def test_rebuild_bufferop_st(ast:UOp):
  for src in ast.src: test_rebuild_bufferop_st(src)

 if __name__ == "__main__":
+  from extra.optimization.helpers import load_worlds, ast_str_to_ast
  ast_strs = load_worlds(False, False, True)[:2000]
  for ast_str in tqdm(ast_strs):
    test_rebuild_bufferop_st(ast_str_to_ast(ast_str))
@@ -52,12 +52,12 @@ class TestHCQ(unittest.TestCase):
      with self.subTest(name=str(queue_type)):
        q = queue_type().signal(virt_signal, virt_val)

-        var_vals = {virt_signal.base_buf.va_addr: TestHCQ.d0.timeline_signal.base_buf.va_addr, virt_val: TestHCQ.d0.timeline_value}
+        var_vals = {virt_signal.base_buf.va_addr.expr: TestHCQ.d0.timeline_signal.base_buf.va_addr, virt_val.expr: TestHCQ.d0.timeline_value}
        q.submit(TestHCQ.d0, var_vals)
        TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
        TestHCQ.d0.timeline_value += 1

-        var_vals = {virt_signal.base_buf.va_addr: TestHCQ.d0.timeline_signal.base_buf.va_addr, virt_val: TestHCQ.d0.timeline_value}
+        var_vals = {virt_signal.base_buf.va_addr.expr: TestHCQ.d0.timeline_signal.base_buf.va_addr, virt_val.expr: TestHCQ.d0.timeline_value}
        q.submit(TestHCQ.d0, var_vals)
        TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
        TestHCQ.d0.timeline_value += 1
@@ -106,7 +106,7 @@ class TestHCQ(unittest.TestCase):

        fake_signal.value = 0x30

-        q.submit(TestHCQ.d0, {virt_signal.base_buf.va_addr: fake_signal.base_buf.va_addr, virt_val: fake_signal.value})
+        q.submit(TestHCQ.d0, {virt_signal.base_buf.va_addr.expr: fake_signal.base_buf.va_addr, virt_val.expr: fake_signal.value})
        TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
        TestHCQ.d0.timeline_value += 1

@@ -131,7 +131,7 @@ class TestHCQ(unittest.TestCase):
     .signal(TestHCQ.d0.timeline_signal, virt_val)

    for _ in range(100):
-      q.submit(TestHCQ.d0, {virt_val: TestHCQ.d0.timeline_value})
+      q.submit(TestHCQ.d0, {virt_val.expr: TestHCQ.d0.timeline_value})
      TestHCQ.d0.timeline_value += 1

    val = TestHCQ.a.uop.buffer.as_buffer().cast("f")[0]
@@ -146,7 +146,7 @@ class TestHCQ(unittest.TestCase):
    q.exec(TestHCQ.runner._prg, TestHCQ.kernargs_ba_ptr, sint_global, sint_local) \
     .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)

-    q.submit(TestHCQ.d0, {sint_global[0]: 1, sint_local[0]: 1})
+    q.submit(TestHCQ.d0, {sint_global[0].expr: 1, sint_local[0].expr: 1})
    TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
    TestHCQ.d0.timeline_value += 1

@@ -181,7 +181,7 @@ class TestHCQ(unittest.TestCase):
        for z in range(1, 4):
          ctypes.memset(zt._buf.va_addr, 0, zb.nbytes)

-          q.submit(TestHCQ.d0, {virt_val: TestHCQ.d0.timeline_value, virt_local[0]: x, virt_local[1]: y, virt_local[2]: z})
+          q.submit(TestHCQ.d0, {virt_val.expr: TestHCQ.d0.timeline_value, virt_local[0].expr: x, virt_local[1].expr: y, virt_local[2].expr: z})
          TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
          TestHCQ.d0.timeline_value += 1

@@ -253,7 +253,7 @@ class TestHCQ(unittest.TestCase):
                                    .copy(virt_dest_addr, virt_src_addr, 8) \
                                    .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)

-    q.submit(TestHCQ.d0, {virt_src_addr: TestHCQ.a.uop.buffer._buf.va_addr, virt_dest_addr: TestHCQ.b.uop.buffer._buf.va_addr})
+    q.submit(TestHCQ.d0, {virt_src_addr.expr: TestHCQ.a.uop.buffer._buf.va_addr, virt_dest_addr.expr: TestHCQ.b.uop.buffer._buf.va_addr})

    TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
    TestHCQ.d0.timeline_value += 1
@@ -276,7 +276,7 @@ class TestHCQ(unittest.TestCase):
                                    .copy(virt_dest_addr, virt_src_addr, sz) \
                                    .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)

-    q.submit(TestHCQ.d0, {virt_src_addr: buf2._buf.va_addr, virt_dest_addr: buf1._buf.va_addr})
+    q.submit(TestHCQ.d0, {virt_src_addr.expr: buf2._buf.va_addr, virt_dest_addr.expr: buf1._buf.va_addr})

    TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
    TestHCQ.d0.timeline_value += 1
@@ -299,7 +299,7 @@ class TestHCQ(unittest.TestCase):

        fake_signal.value = 0x30

-        q.submit(TestHCQ.d0, {virt_signal.base_buf.va_addr: fake_signal.base_buf.va_addr, virt_val: fake_signal.value})
+        q.submit(TestHCQ.d0, {virt_signal.base_buf.va_addr.expr: fake_signal.base_buf.va_addr, virt_val.expr: fake_signal.value})
        TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
        TestHCQ.d0.timeline_value += 1

@@ -2,7 +2,7 @@
 import unittest

 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.codegen.opt.search import Opt, OptOps
+from .search import Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
@@ -22,7 +22,7 @@ if os.getenv("VALIDATE_HCQ", 0) != 0:
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.codegen.opt.search import get_kernel_actions, bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG, Timing
@@ -90,7 +90,7 @@ def get_fuzz_rawbuf_like(old_rawbuf, zero=False, copy=False, size=None, force_de

 def run_linearizer(lin: Kernel, rawbufs=None, var_vals=None) -> tuple[str, Any]: # (error msg, run state)
  if rawbufs is None: rawbufs = bufs_from_lin(lin)
-  if var_vals is None: var_vals = {v: v.min for v in lin.vars}
+  if var_vals is None: var_vals = {v.expr: v.min for v in lin.vars}

  # TODO: images needs required_optimization
  try:
@@ -129,7 +129,7 @@ def compare_linearizer(lin: Kernel, rawbufs=None, var_vals=None, ground_truth=No

  if var_vals is None:
    # TODO: handle symbolic max case
-    var_vals = {v: random.randint(v.vmin, v.vmax) for v in lin.ast.variables()}
+    var_vals = {v.expr: random.randint(v.vmin, v.vmax) for v in lin.ast.variables()}

  if ground_truth is None and not has_bf16:
    unoptimized = Kernel(lin.ast)
@@ -12,7 +12,7 @@ try:
  from tinygrad.renderer import Renderer, ProgramSpec
  from tinygrad.engine.realize import get_program
  from tinygrad.uop.ops import UOp, Ops, KernelInfo
-  from tinygrad.codegen.opt.kernel import Opt
+  from tinygrad.codegen.opt import Opt
  from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
  from tinygrad.device import Device
 except ImportError as e:
@@ -336,5 +336,19 @@ class TestKernelOpts(unittest.TestCase):
      [Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.UNROLL, axis=1, arg=4)], # noqa: E501
    ])

+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_threads, "test requires threads")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.global_max is not None and
+                       Device[Device.DEFAULT].renderer.global_max[0] > 1, "test requires multicore")
+  def test_thread_opts(self):
+    a = Tensor.rand(4, 4, 4, 4)
+    b = Tensor.rand(4, 4, 4)
+    r = (b.sqrt() + ((a+1).sum(axis=3).exp()))
+    helper_linearizer_opt(r, [
+      [Opt(OptOps.THREAD, 0, 2)],
+      [Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.THREAD, 0, 2)],
+      [Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.THREAD, 0, 2), Opt(OptOps.UNROLL, 0, 2)],
+    ] + [[Opt(OptOps.THREAD, 0, 4)] if Device[Device.DEFAULT].renderer.global_max[0] >= 4 else []]
+      + [[Opt(OptOps.THREAD, 0, 8)] if Device[Device.DEFAULT].renderer.global_max[0] >= 8 else []])
+
 if __name__ == '__main__':
  unittest.main()
@@ -26,6 +26,10 @@ import unittest
 import numpy as np
 import torch
 from tinygrad import Tensor, dtypes, nn
+from tinygrad.device import is_dtype_supported
+from tinygrad.helpers import getenv
+
+MOCKGPU = getenv("MOCKGPU")

 class TestNaNEdgeCases(unittest.TestCase):
  # we don't need more of these. it's unclear if torch's behavior is desired here
@@ -167,34 +171,6 @@ class TestZeroFolding(unittest.TestCase):
    with self.assertRaises(RuntimeError):
      (x % x).numpy()

-class TestArangeUOpValidationIssue(unittest.TestCase):
-  # these fail with UOp verification error.
-  # we don't need more of these involving arange
-
-  @unittest.expectedFailure
-  def test_large_arange_sum(self):
-    # Summing a huge arange should either succeed or raise a MemoryError.
-    n = 2**31 + 3
-    expected = (n - 1) * n // 2
-    out = Tensor.arange(n).sum().item()
-    self.assertEqual(out, expected)
-
-  @unittest.expectedFailure
-  def test_large_arange_index(self):
-    # Indexing a huge arange should return the correct value instead of failing
-    # with a UOp verification error.
-    n = 2**31 + 3
-    out = Tensor.arange(n)[0].item()
-    self.assertEqual(out, 0)
-
-  @unittest.expectedFailure
-  def test_large_arange_permute(self):
-    # Permuting a huge tensor should not trigger UOp verification failures.
-    n = 2**31 + 3
-    out = Tensor.arange(n).reshape(n, 1).permute(1, 0)
-    self.assertEqual(out.shape, (1, n))
-    out.realize()
-
 class TestAssignIssues(unittest.TestCase):
  # these are good failures. i'm not sure we need more, but we need to fix these.

@@ -230,10 +206,8 @@ class TestUOpValidationIssue(unittest.TestCase):
  # these fail with UOp verification error.
  # we want more of these with diverse errors!

-  @unittest.expectedFailure
+  @unittest.skipIf((not is_dtype_supported(dtypes.long)) or MOCKGPU, "hangs gpuocelot")
  def test_tensor_index_overflow(self):
-    # Advanced indexing on tensors expanded past int32 should not error, but
-    # tinygrad fails with a UOp verification error.
    val = Tensor([1])
    big = val.expand(2**31 + 3)
    idx = Tensor([0, 2**31 + 2])
@@ -273,4 +247,4 @@ class TestEdgeCases(unittest.TestCase):


 if __name__ == "__main__":
-  unittest.main()
+  unittest.main()
@@ -928,6 +928,7 @@ class TestIdxUpcast(unittest.TestCase):
      self.do_op_then_assert(dtypes.long, 2048, 2048, UOp.variable("dim3", 1, 2048).bind(32))

  @unittest.skipIf(is_dtype_supported(dtypes.long), "int64 is supported")
+  @unittest.expectedFailure  # bug in gpu dims limiting
  def test_int64_unsupported_overflow(self):
    with self.assertRaises(KeyError):
      self.do_op_then_assert(dtypes.long, 2048, 2048, 2048)
@@ -99,7 +99,7 @@ class TestTiny(unittest.TestCase):
  # *** a model ***

  # TODO: this is failing because of how swizzling rewrites the ShapeTracker of the final STORE
-  @unittest.skipIf(IMAGE>0 or (CI and Device.DEFAULT == "DSP"), "failing because of make things that can't be images not images")
+  @unittest.skipIf(CI and Device.DEFAULT == "DSP", "failing because of make things that can't be images not images")
  def test_mnist(self):
    layers = [
      nn.Conv2d(1, 32, 5), Tensor.relu,
@@ -118,7 +118,7 @@ class TestTiny(unittest.TestCase):
    self.assertEqual(len(probs[0]), 10)

  # TODO: this is failing because of how swizzling rewrites the ShapeTracker of the final STORE
-  @unittest.skipIf(IMAGE>0 or (CI and Device.DEFAULT == "DSP"), "failing because of make things that can't be images not images")
+  @unittest.skipIf(CI and Device.DEFAULT == "DSP", "failing because of make things that can't be images not images")
  def test_mnist_backward(self):
    # NOTE: we don't have the whole model here for speed
    layers = [
@@ -212,7 +212,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_same_fold(self):
    v = UOp.variable('tmp', 0, 1)
-    c0 = UOp(Ops.CONST, dtypes.int, arg=0)
+    c0 = UOp(Ops.CONST, dtypes.index, arg=0)
    vc = UOp(Ops.CMPNE, dtypes.bool, (v, c0))
    c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
    out = UOp(Ops.WHERE, dtypes.float, (vc, c1, c1))
@@ -398,7 +398,7 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)

  def test_depth_2_const_fold(self):
-    v = UOp.variable("tmp", 0, 1)
+    v = UOp.variable("tmp", 0, 1, dtypes.int)
    c2 = UOp(Ops.CONST, dtypes.int, arg=2)
    c4 = UOp(Ops.CONST, dtypes.int, arg=4)
    vc = UOp(Ops.ADD, dtypes.int, (v, c2))
@@ -417,6 +417,17 @@ class TestUOpGraph(unittest.TestCase):
      uops = to_uops_list([v.bitcast(dt)])
      self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST]), 0, f"dtype = {dt}")

+  def test_load_idx_becomes_int(self):
+    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
+    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 1)
+    l0 = UOp(Ops.LOAD, dtypes.long, (d0.index(UOp.const(dtypes.int, 0)),))
+    idx = l0 * 600
+    valid = (l0<-1).ne(True)&(l0<3000)
+    l1 = UOp(Ops.LOAD, dtypes.long, (d1.index(idx, valid),))
+    uops = to_uops_list([l1])
+    for u in uops:
+      if u.op is Ops.INDEX: self.assertEqual(u.src[1].dtype, dtypes.int)
+
  def test_in_out_of_bounds_access(self):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
@@ -512,7 +523,7 @@ class TestUOpGraph(unittest.TestCase):
  def test_in_out_bounds_access_with_mask(self):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 42),), "gidx0")
+      gidx0 = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 42),), "gidx0")
      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, (5<gidx0)&(gidx0<16)),))
      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, gidx0<16),))
      to_uops_list([ld0, ld1])
@@ -536,9 +547,9 @@ class TestUOpGraph(unittest.TestCase):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
      glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 0)
-      gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 42),), "gidx0")
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, gidx0<8),))
-      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(ld0*2, (ld0>=0)&(ld0<32)),))
+      gidx0 = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 42),), "gidx0")
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, gidx0<8),)).cast(dtypes.index)
+      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(ld0*2, (ld0>=0)&(ld0<32)),)).cast(dtypes.index)
      to_uops_list([ld1])

      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(ld0*2, (ld0>=0)&(ld0<64)),))
@@ -476,7 +476,7 @@ class TestUOpMethod(unittest.TestCase):
    st_var = Tensor.empty((2, 10))[:, :a.bind(1)]
    _, var_vals = (uop_var+st_var).schedule_with_vars()
    self.assertEqual(len(var_vals), 1)
-    self.assertEqual(list(var_vals)[0], a)
+    self.assertEqual(list(var_vals)[0], a.expr)

  def test_const_factor(self):
    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 8),), 'gidx0')
@@ -97,19 +97,22 @@ class TestFoldingAndReduction(unittest.TestCase):

 class TestModuloAndDivisionFolding(unittest.TestCase):
  def test_full_graph_rewrite_modulo_folding_with_define_var(self):
-    x_var_uop = UOp.variable('x', 0, 100)
+    # index dtype because div-mod rules only work on index
+    x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.index)
    optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4)
    self.assertEqual(optimized_mod_uop.op, Ops.CONST)
    self.assertEqual(optimized_mod_uop.arg, 2)

  def test_full_graph_rewrite_division_folding_with_define_var(self):
-    n_var_uop = UOp.variable('n', 1, 1000)
+    # index dtype because div-mod rules only work on index
+    n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.index)
    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].arg, 2)

  def test_full_graph_rewrite_complex_mod_div_folding(self):
-    k_var_uop = UOp.variable('k', 0, 50)
+    # index dtype because div-mod rules only work on index
+    k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.index)
    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.arg, 1)
@@ -126,8 +129,9 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
    if opt.op is Ops.VECTORIZE: self.assertFalse(all_same(opt.src))

  def test_full_graph_rewrite_modulo_large_divisor(self):
+    # index dtype because div-mod rules only work on index
    x_var_uop = UOp.variable('x', 1, 5)
-    self.assertIs(apply_rewrite(x_var_uop % 10), x_var_uop)
+    self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.index) % 10).render(simplify=False), x_var_uop.render(simplify=False))

  def test_full_graph_rewrite_division_with_remainder(self):
    x_var_uop = UOp.variable('x', 7, 9)
@@ -46,8 +46,8 @@ class TestRewriteMap(unittest.TestCase):

  def test_add_zero(self):
    # Build a small graph: add(0, add(const=0, const=5))
-    zero_node = UOp.const(dtypes.int, 0)
-    five_node = UOp.const(dtypes.int, 5)
+    zero_node = UOp.const(dtypes.index, 0)
+    five_node = UOp.const(dtypes.index, 5)
    inner_add = zero_node + five_node
    root_add = zero_node + inner_add

@@ -67,7 +67,7 @@ class TestRewriteMap(unittest.TestCase):
    Test rewriting neg(neg(5)) => 5 using symbolic.
    """
    # In some versions of TinyGrad, you might do: (-(-five_node))
-    five_node = UOp.const(dtypes.int, 5)
+    five_node = UOp.const(dtypes.index, 5)
    # If your code allows UOp(...), do that; else you might do something like:
    # double_neg_five = -(-five_node)
    # But let's be explicit:
@@ -85,8 +85,8 @@ class TestRewriteMap(unittest.TestCase):
    """
    Combine both rewrites: add(0, neg(neg(5))) => add(0, 5) => 5
    """
-    zero_node = UOp.const(dtypes.int, 0)
-    five_node = UOp.const(dtypes.int, 5)
+    zero_node = UOp.const(dtypes.index, 0)
+    five_node = UOp.const(dtypes.index, 5)
    neg_five = -five_node
    double_neg_five = -neg_five
    root_add = zero_node + double_neg_five
@@ -103,7 +103,7 @@ class TestRewriteMap(unittest.TestCase):
  def test_multi_var_rewrites(self):
    x_var = UOp.variable('x', 0, 10)
    y_var = UOp.variable('y', -5, 5)
-    zero_node = UOp.const(dtypes.int, 0)
+    zero_node = UOp.const(dtypes.index, 0)

    sum_with_zero = y_var + zero_node    # (y + 0)
    combined = x_var + sum_with_zero     # x + (y + 0)
@@ -155,8 +155,8 @@ class TestRewriteMap(unittest.TestCase):
    x_var = UOp.variable('x', 1, 10)
    y_var = UOp.variable('y', -5, 5)
    z_var = UOp.variable('z', 0, 5)
-    zero_node = UOp.const(dtypes.int, 0)
-    one_node = UOp.const(dtypes.int, 1)
+    zero_node = UOp.const(dtypes.index, 0)
+    one_node = UOp.const(dtypes.index, 1)

    # Build sub-expressions
    yz_sum = y_var + z_var              # (y + z)
@@ -18,7 +18,7 @@ def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UO
    UOp(Ops.VECTORIZE, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4)
  ))

-def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.int32, (UOp.const(dtypes.int, nmax),), expr)
+def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, nmax),), expr)
 def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
 def Range(n, nmax): return UOp.range(nmax, n)

@@ -73,19 +73,19 @@ class TestSymbolicVarVals(unittest.TestCase):

  def test_var_vals_shape(self):
    x = Variable("x", 1, 100).bind(3)
-    assert ShapeTracker.from_shape((x, 3)).var_vals == {Variable("x", 1, 100): 3}
+    assert ShapeTracker.from_shape((x, 3)).var_vals == {"x": 3}

  def test_var_vals_offset(self):
    x = Variable("x", 1, 100).bind(3)
    st = ShapeTracker.from_shape((4, 3)).shrink(((x, x+1), (0, 3)))
    self.assert_equal(st.views[-1].offset, x * 3)
-    assert st.var_vals == {Variable("x", 1, 100): 3}
+    assert st.var_vals == {"x": 3}

  def test_var_vals_mask(self):
    x = Variable("x", 1, 100).bind(3)
    view = View.create(shape=(3,4), strides=(4,1), offset=0, mask=((0, x), (0, 4)))
    st = ShapeTracker(views=(view,))
-    assert st.var_vals == {Variable("x", 1, 100): 3}
+    assert st.var_vals == {"x": 3}

  def test_var_vals_complex(self):
    x = Variable("x", 1, 100).bind(3)
@@ -93,13 +93,13 @@ class TestSymbolicVarVals(unittest.TestCase):
    z = Variable("z", 1, 100).bind(5)
    st = ShapeTracker.from_shape((x, 5, y)).shrink(((0, x), (z, z+1), (0, 3)))
    self.assert_equal(st.views[-1].offset, y * z)
-    assert st.var_vals == {Variable("x", 1, 100): 3, Variable("y", 1, 100):4, Variable("z", 1, 100): 5}
+    assert st.var_vals == {"x": 3, "y": 4, "z": 5}

  def test_shrink_reshape(self):
    x = Variable("x", 1, 100).bind(3)
    st = ShapeTracker.from_shape((10, 10, 10)).shrink(((x, x+3), (3, 7), (2, 5)))
    st = st.reshape((3*4*3,))
-    assert st.var_vals == {Variable("x", 1, 100): 3}
+    assert st.var_vals == {"x": 3}

 class TestShapeTrackerUnbind(unittest.TestCase):
  def test_view_unbind(self):
@@ -2,22 +2,20 @@
 import unittest, pickle, functools, math
 import z3

-from tinygrad.dtype import dtypes, ConstType
+from tinygrad.dtype import dtypes, ConstType, DType
 from tinygrad.codegen import full_rewrite
-from tinygrad.codegen.late.devectorizer import sym
 from tinygrad.helpers import Context
-from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
-from tinygrad import Variable
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer, track_rewrites
+from tinygrad.uop.symbolic import sym
 from tinygrad.uop.spec import uops_to_z3

-def render(self) -> tuple[str, ConstType, ConstType]:
-  # NOTE: we need STORE so the ALU op has children
-  glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
-  uops = full_rewrite(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), self)).sink())
-  rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
-  return rewritten_uop.render(simplify=False), rewritten_uop.vmin, rewritten_uop.vmax
+@track_rewrites(name="simplify symbolic uop")
+def render(v) -> UOp:
+  v_simplified = graph_rewrite(v, sym)
+  return v_simplified

-def uconst(val): return UOp.const(dtypes.int, val)
+def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.index): return UOp.variable(name,min_val,max_val,dtype)
+def uconst(val): return UOp.const(dtypes.index, 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)

@@ -30,11 +28,12 @@ class TestSymbolicPickle(unittest.TestCase):

 class TestSymbolic(unittest.TestCase):
  def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
+    v_simplified = render(v)
    if test_z3:
      solver = z3.Solver()
-      expr, expr_simplified = uops_to_z3(solver, v, v.simplify())
+      expr, expr_simplified = uops_to_z3(solver, v, v_simplified)
      self.assertEqual(solver.check(expr != expr_simplified), z3.unsat, "simplified expression not equal to original")
-    rendered, nmin, nmax = render(v)
+    rendered, nmin, nmax = v_simplified.render(simplify=False), v_simplified.vmin, v_simplified.vmax
    if isinstance(s, tuple): self.assertIn(rendered, s)
    else: self.assertEqual(rendered, s)
    self.assertEqual(nmin, n)
@@ -111,7 +110,7 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(-Variable("a", 0, 8), -8, 0, "(a*-1)")

  def test_xor_0(self):
-    self.helper_test_variable(Variable("a", 0, 8) ^ 0, 0, 8, "a")
+    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) ^ 0, 0, 8, "a")

  def test_add_1(self):
    self.helper_test_variable(Variable("a", 0, 8)+1, 1, 9, "(a+1)")
@@ -209,12 +208,12 @@ class TestSymbolic(unittest.TestCase):
    self.assertEqual((Variable("x", -10, 0)%Variable("y", 1, 10))._min_max, (-9, 0))

  def test_range_div_its_symbolic_bound(self):
-    a = Variable("a", 1, 10)
+    a = Variable("a", 1, 10, dtypes.index)
    ridx0 = UOp.range(a+2, 0)
    self.helper_test_variable(ridx0//(a+2), 0, 0, "0")

  def test_range_mod_its_symbolic_bound(self):
-    a = Variable("a", 1, 10)
+    a = Variable("a", 1, 10, dtypes.index)
    ridx = UOp.range(a+2, 0)
    self.helper_test_variable(ridx%(a+2), 0, 11, "ridx0")

@@ -463,8 +462,8 @@ class TestSymbolic(unittest.TestCase):
      self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -8, 0, "(((idx*10)//11)*-1)")

  def test_nest_div_negative_factor(self):
-    ridx0=UOp.variable("ridx0", 0, 9)
-    ridx1=UOp.variable("ridx1", 0, 6)
+    ridx0=Variable("ridx0", 0, 9)
+    ridx1=Variable("ridx1", 0, 6)
    self.helper_test_variable(((((ridx0*-7)+ridx1)+63)//35), 0, 1, "(((ridx0//5)*-1)+1)")

  def test_div_into_mod(self):
@@ -533,8 +532,8 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(x//y, 2, 2, "2")
    self.helper_test_variable(x%y, 0, 7, "(x+(y*-2))")
    # ensure all 4 corners are checked
-    x = Variable("x", -10, 10)
-    y = Variable("y", -8, 9)
+    x = Variable("x", -10, 10, dtypes.int)
+    y = Variable("y", -8, 9, dtypes.int)
    self.helper_test_variable(x//y, -2147483648, 2147483647, "(x//y)")
    self.helper_test_variable(x%y, -2147483648, 2147483647, "(x%y)")

@@ -587,6 +586,12 @@ class TestSymbolic(unittest.TestCase):
    unrolled_div = (gidx+2561)//4+(gidx+2562)//4+(gidx+2560)//4+(gidx+2559)//4
    self.helper_test_variable(unrolled_div, 2559, 5118, "(gidx+2559)")

+  def test_arange_unrolled4_with_cast(self):
+    gidx = Variable("gidx", 0, 2559, dtypes.index)
+    dt = dtypes.int
+    unrolled_div = ((gidx+2561)//4 + 2).cast(dt)+((gidx+2562)//4).cast(dt)+((gidx+2560)//4).cast(dt)+((gidx+2559)//4).cast(dt)
+    self.helper_test_variable(unrolled_div, 2561, 5120, "((int)(gidx)+2561)")
+
  def test_arange_unrolled4_mul(self):
    gidx = Variable("gidx", 0, 2559)
    unrolled_div = 2*((gidx+2561)//4)+2*((gidx+2562)//4)+2*((gidx+2560)//4)+2*((gidx+2559)//4)
@@ -688,10 +693,10 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(-a<-b, False, True, "(b<a)")

  def test_where_cast(self):
-    s = Variable("s", 0, 3)
+    s = Variable("s", 0, 3, dtypes.int)
    cond = s < 2
-    a = Variable("a", 0, 3)
-    b = Variable("b", 0, 3)
+    a = Variable("a", 0, 3, dtypes.int)
+    b = Variable("b", 0, 3, dtypes.int)
    expr = cond.where(a, b).cast(dtypes.half)

    # TODO: copied from render, render does not support cast
@@ -738,8 +743,8 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(a.trunc(), 1, 10, "a", test_z3=False)

  def test_do_math_in_int32(self):
-    a = Variable("a", 1, 10)
-    b = Variable("b", 1, 10)
+    a = Variable("a", 1, 10, dtypes.int)
+    b = Variable("b", 1, 10, dtypes.int)
    self.helper_test_variable(a.cast(dtypes.long)+b.cast(dtypes.long), 2, 20, "(long)((a+b))")
    self.helper_test_variable(a.cast(dtypes.long)*b.cast(dtypes.long), 1, 100, "(long)((a*b))")

@@ -804,7 +809,7 @@ class TestSymInfer(unittest.TestCase):
    a = Variable("a", 0, 10)
    b = Variable("b", 0, 10)
    c = Variable("c", 0, 10)
-    var_vals = {a: 2, b: 3, c: 4}
+    var_vals = {a.expr: 2, b.expr: 3, c.expr: 4}
    assert sym_infer(5, var_vals) == 5
    assert sym_infer(4.2, var_vals) == 4.2
    assert sym_infer(a, var_vals) == 2
@@ -817,7 +822,7 @@ class TestSymInfer(unittest.TestCase):
  def test_sym_infer_cdiv_cmod(self):
    a = Variable("a", -1000, 1)
    b = Variable("b", -1000, 1)
-    var_vals = {a: 1, b: -1000}
+    var_vals = {a.expr: 1, b.expr: -1000}
    assert sym_infer(a%b, var_vals) == 1
    assert sym_infer(a//b, var_vals) == 0

@@ -142,7 +142,7 @@ class TestViz(BaseTestViz):

  def test_const_node_visibility(self):
    a = UOp.variable("a", 0, 10)
-    z = UOp.const(dtypes.int, 0)
+    z = UOp.const(dtypes.index, 0)
    alu = a*z
    exec_rewrite(alu, [sym])
    lst = get_viz_list()
@@ -2,7 +2,7 @@ from typing import Any, Callable
 import functools
 from dataclasses import dataclass
 from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL
-from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp
+from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype
 from tinygrad.uop.spec import type_verify
 from tinygrad.renderer import Renderer

@@ -10,7 +10,7 @@ from tinygrad.renderer import Renderer
 from tinygrad.codegen.lowerer import pm_lowerer, get_index
 from tinygrad.codegen.quantize import pm_quant
 from tinygrad.codegen.gpudims import pm_add_gpudims
-from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing
+from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, cast_folding
 from tinygrad.uop.decompositions import get_late_rewrite_patterns
 from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_expander
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
@@ -18,6 +18,7 @@ from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_in
 from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
 from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
 from tinygrad.codegen.opt.postrange import pm_postrange_opt
+from tinygrad.codegen.simplify import pm_simplify_ranges
 from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen

@dataclass
@@ -43,27 +44,32 @@ rewrites_for_linearizer = [
  RewriteStep(block_merge, name="Linearizer: Merge Blocks"),
  RewriteStep(pm_finalize, name="Linearizer: Finalize")]

-def get_rewrites_for_renderer(opts:Renderer, linearizer:bool=True) -> list[RewriteStep]:
+def get_rewrites_for_renderer(opts:Renderer, optimize:bool=True, linearizer:bool=True) -> list[RewriteStep]:
  # cache with the values of the context vars
-  return _get_rewrites_for_renderer(opts, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value)
+  return _get_rewrites_for_renderer(opts, optimize, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value)

@functools.cache
-def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
+def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
  # ** lowerer (rewrite_shapetracker_with_index) **
  ret: list[RewriteStep] = []

-  # view pushing
-  ret.extend(rewrites_for_views)
+  if optimize:
+    # view pushing
+    ret.extend(rewrites_for_views)

-  # lowerer first
-  if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
-  ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True))
+    # lowerer first
+    if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
+    ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True))

-  # optimize (schedule) the AST
-  ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
+    # symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
+    ret.append(RewriteStep(sym, name="initial symbolic"))
+
+    # optimize (schedule) the AST
+    ret.append(RewriteStep(pm_simplify_ranges, name="simplify ranges"))
+    ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))

  # ** expander (expand_rewrite) **
-  ret.append(RewriteStep(sym+migrate_indexing, name="initial symbolic"))
+  ret.append(RewriteStep(sym+migrate_indexing, name="postopt symbolic"))

  # expand
  ret.append(RewriteStep(sym+pm_pre_expander+expander, name="expander"))
@@ -87,6 +93,9 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
  supported_ops = tuple(opts.code_for_op.keys())
  extra_matcher = opts.extra_matcher if opts.extra_matcher is not None else PatternMatcher([])

+  # lower the index dtype to a concrete int
+  ret.append(RewriteStep(pm_lower_index_dtype+cast_folding+load_store_indexing, lambda _: opts.device, name="lower all index dtypes"))
+
  # optional pre matcher
  if opts.pre_matcher is not None: ret.append(RewriteStep(opts.pre_matcher, name="pre_matcher"))

@@ -101,8 +110,8 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
  # return the list (with optional linearizer)
  return ret + (rewrites_for_linearizer if linearizer else [])

-def full_rewrite_to_sink(sink:UOp, opts:Renderer|None=None, linearizer:bool=False) -> UOp:
-  return apply_rewrites(sink, get_rewrites_for_renderer(opts if opts is not None else Renderer(), linearizer))
+def full_rewrite_to_sink(sink:UOp, opts:Renderer|None=None, optimize:bool=True, linearizer:bool=False) -> UOp:
+  return apply_rewrites(sink, get_rewrites_for_renderer(opts if opts is not None else Renderer(), optimize, linearizer))

 def full_rewrite(sink:UOp, opts:Renderer|None=None) -> list[UOp]:
  """
@@ -116,6 +125,6 @@ def full_rewrite(sink:UOp, opts:Renderer|None=None) -> list[UOp]:
    Linear program in UOps.
  """

-  lst = list(full_rewrite_to_sink(sink, opts, linearizer=True).arg.lst)
+  lst = list(full_rewrite_to_sink(sink, opts, optimize=sink.tag is None, linearizer=True).arg.lst)
  if __debug__: type_verify(lst)
  return lst
@@ -34,7 +34,7 @@ def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|No
  if max_sizes is not None and len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
  # try to split up dims: (a,) -> (b, c)
  if limited == dims: limited = _split_dims(dims, max_sizes) if max_sizes is not None else dims
-  ret = raw_idxs = [UOp(Ops.SPECIAL, dtypes.int, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
+  ret = raw_idxs = [UOp(Ops.SPECIAL, dtypes.index, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
  if len(limited) < len(dims):
    ret = []
    if (contraction:=get_contraction(dims, limited)) is None: raise AssertionError(f"get_contraction should not be None {dims=} {limited=}")
@@ -59,7 +59,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
  all_ranges = {x.arg[0:-1]:x for x in s_topo if x.op is Ops.RANGE}

  # extract global/local dims
-  global_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] is AxisType.GLOBAL]))
+  global_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.GLOBAL, AxisType.THREAD)]))
  local_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)]))
  if not global_dims and not local_dims: return None

@@ -4,7 +4,7 @@ from collections import defaultdict
 from dataclasses import dataclass
 from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, graph_rewrite, GroupOp, identity_element
-from tinygrad.uop.symbolic import uop_given_valid, parse_valid, simplify_valid, sym, symbolic_flat
+from tinygrad.uop.symbolic import uop_given_valid, parse_valid, sym, symbolic_flat, cast_folding
 from tinygrad.helpers import getenv, flatten, AMX, prod, partition
 from tinygrad.renderer import Renderer

@@ -51,12 +51,13 @@ def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp,
  return UOp.store(buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx), val, *store.src[2:])

 load_store_indexing = PatternMatcher([
-  # simplify valid
-  (UPat(Ops.AND, name="valid"), simplify_valid),
  # image load valid idx simplification
  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("start_idx"), UPat.var("valid"))), simplify_valid_load),
  # index True is just Index
  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("start_idx"), UPat(Ops.CONST, arg=True))), lambda buf,start_idx: buf.index(start_idx)),
+  # remove hanging cast
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.int).cast()),), lambda buf,idx: buf.index(idx)),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.int).cast(), UPat.var("valid"))), lambda buf,idx,valid: buf.index(idx, valid)),
  # delete_redundant_gates (after expand)
  (UPat(Ops.STORE, src=(UPat.any(stidx:=UPat.var("buf").index(UPat.var("idx"), UPat.var("store_gate")), stidx.cast().named("cast")),
                                  UPat.var("val")), name="store", allow_any_len=True), delete_redundant_gates),
@@ -316,12 +317,12 @@ def no_range(u:UOp) -> bool: return not any(x.op is Ops.RANGE for x in u.sparent

 pm_reduce_collapse = PatternMatcher([
  # lift x+y out of reduce on lt
-  ((UPat.var("x")+UPat.var("y")) < UPat.var("c"), lambda x,y,c: (x < (c-y)) if no_range(y) and no_range(c) else None),
+  ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
  # lift x*y out of reduce
  ((UPat.var("x")*UPat.var("y")) < UPat.var("c"),
   lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and y.vmin > 0 else None),
  # lift x+y out of reduce on ne
-  ((UPat.var("x")+UPat.var("y")) != UPat.var("c"), lambda x,y,c: (x != (c-y)) if no_range(y) and no_range(c) else None),
+  ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"), lambda x,y,c: (x != (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
  # fold the range
  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.cvar("val")).reduce(arg=Ops.ADD, allow_any_len=True),
   lambda r,cut,val: (r.src[0]-cut).maximum(0).minimum(r.src[0]).cast(val.dtype) * val),
@@ -351,7 +352,7 @@ pm_reduce_collapse = PatternMatcher([
  (UPat(Ops.VECTORIZE, name="inp").reduce(name="red", allow_any_len=True), no_vectorized_reduce),
  # index/load/where. TODO: this is more aggressive than needed
  (UPat((Ops.INDEX, Ops.LOAD, Ops.WHERE), name="alu"), no_vectorized_alu),
-])+sym
+])+sym+cast_folding

 def reduce_collapse(red:UOp):
  included, not_included = partition(red.parents, lambda x: any(y in x.sparents for y in red.src[1:]))
@@ -1,5 +1,4 @@
 # this converts a lowerer program into a vectorized program
-
 import functools, itertools, operator
 from tinygrad.dtype import dtypes, PtrDType, AddrSpace
 from tinygrad.helpers import AMX, dedup, flatten, all_same, prod, partition
@@ -152,7 +151,7 @@ def fix_group_for_reduce(x:UOp):
 pm_pre_expander = PatternMatcher([
  # rewrite UPCAST/UNROLL range to something to be expanded
  (UPat(Ops.RANGE, name="r"),
-   lambda r: UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0],s),)) \
+   lambda r: UOp(Ops.UNROLL, r.dtype, (UOp.const(r.dtype.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0],s),)) \
    if r.arg[1] in {AxisType.UNROLL, AxisType.UPCAST} else None),
  # fix REDUCEs with UNROLLs
  (UPat(Ops.REDUCE, name="x"), fix_reduce_unroll),
@@ -5,7 +5,7 @@ from dataclasses import dataclass
 from tinygrad.uop.ops import AxisType

 class OptOps(Enum):
-  TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto() # noqa: E702
+  TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto(); THREAD = auto() # noqa: E702
  GROUP = auto(); GROUPTOP = auto(); NOLOCALS = auto(); PADTO = auto(); SWAP = auto() # noqa: E702
  def __lt__(self, x:OptOps): return self.value < x.value

@@ -16,10 +16,10 @@ class Opt:
  arg: int|tuple|None = None
  def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg})"

-axis_letters = {AxisType.GLOBAL: "g", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
+axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
-axis_colors = {AxisType.GLOBAL: "blue", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE", AxisType.UPCAST: "yellow",
-               AxisType.GROUP_REDUCE: "green", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
+axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
+               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "green", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}

 class KernelOptError(Exception): pass
 def check(cond:bool, msg:str=""):
@@ -3,12 +3,9 @@ from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS, TC_OPT, TC_SELECT, USE_TC, AMX
 from tinygrad.dtype import ImageDType
 from tinygrad.uop.ops import Ops, resolve, AxisType
-
-# both versions
-from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.codegen.opt.postrange import Scheduler

-def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
+def hand_coded_optimizations(k:Scheduler) -> Scheduler:
  # first try the tensor cores
  """ Attempts to apply a tensor core optimization to the kernel. If one exists and applies properly, return true, otherwise return false.
  Tensor cores are optimized instructions that matrix multiply-accumulate across a wave of threads: D(M, N) = A(M, K) * B(K, N) + C(M, N).
@@ -38,15 +35,7 @@ def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
      pass
    if good_tc_opt:
      # skip hand-coded TC opts if AMX, upcasting will make kernel slower
-      if isinstance(k, Kernel) and (tc_opts:=tk.tensor_core_opts) is not None and not AMX:
-        # hand-coded TC opts
-        for tc_dim in [tc_dim for tc_dim in [1,0] if tc_opts.axes_exist[tc_dim]]: # attempt to upcast M and N
-          szs = [sz for sz in [5,4,3,2] if tk.full_shape[tc_opts.axes[tc_dim]] % sz == 0]
-          if szs: tk.apply_opt(Opt(OptOps.UPCAST, tc_opts.axes[tc_dim], szs[0]))
-
-        if tc_opts.axes_exist[0] and (szs := [sz for sz in [4,2] if tk.full_shape[tc_opts.axes[0]] % sz == 0]): # attempt to local N
-          tk.apply_opt(Opt(OptOps.LOCAL, tc_opts.axes[0], szs[0]))
-      elif isinstance(k, Scheduler) and rngs is not None and not AMX:
+      if rngs is not None and not AMX:
        for tc_dim in [1,0]: # attempt to upcast M and N
          szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None]
          if szs:
@@ -54,7 +43,7 @@ def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
            rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0]
        if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N
          tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0]))
-      return tk.applied_opts
+      return tk

  # make a copy so it does not mutate the input
  k = k.copy()
@@ -64,32 +53,17 @@ def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
  if k.opts.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and  \
    k.reduceop is not None and k.reduceop.arg[0] is Ops.ADD and len(k.full_shape) >= 2 and k.opts.has_shared and \
    (mulop:=k.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.LOAD and mulop.src[1].op is Ops.LOAD:
-    if isinstance(k, Kernel):
-      st0, st1 = k.sts[k.bufs.index(mulop.src[0])], k.sts[k.bufs.index(mulop.src[1])]
-      strides0, strides1 = st0.real_strides(), st1.real_strides()
-      def has_expanded_axis(shape, strides): return any(resolve(s > 1) and not resolve(st != 0) for s,st in zip(shape,strides))
-      if strides0[first_reduce:=(k.axes_of(AxisType.REDUCE)[0])] == 1 and \
-        not (has_expanded_axis(st0.shape, strides0) and has_expanded_axis(st1.shape, strides1)):
-        for global_idx in k.axes_of(AxisType.GLOBAL):
-          if k.full_shape[first_reduce]%MV_THREADS_PER_ROW == 0 and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
-            if DEBUG >= 3:
-              print(f"MATVEC: {k.full_shape=} {first_reduce=} {strides0=} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
-            if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
-            if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
-            if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
-            return k.applied_opts
-    else:
-      idx0, idx1 = mulop.src[0].src[0].src[1], mulop.src[1].src[0].src[1]
-      first_reduce_rng = k.ranges_of(AxisType.REDUCE)[0]
-      if any(u is first_reduce_rng for u in idx0.split_uop(Ops.ADD)) and all(r in idx1.ranges for r in idx0.ranges):
-        for global_idx in k.axes_of(AxisType.GLOBAL):
-          if first_reduce_rng.src[0].divides(MV_THREADS_PER_ROW) is not None and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
-            if DEBUG >= 3:
-              print(f"MATVEC: {k.full_shape=} {first_reduce_rng.render()} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
-            if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
-            if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
-            if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
-            return k.applied_opts
+    idx0, idx1 = mulop.src[0].src[0].src[1], mulop.src[1].src[0].src[1]
+    first_reduce_rng = k.ranges_of(AxisType.REDUCE)[0]
+    if any(u is first_reduce_rng for u in idx0.split_uop(Ops.ADD)) and all(r in idx1.ranges for r in idx0.ranges):
+      for global_idx in k.axes_of(AxisType.GLOBAL):
+        if first_reduce_rng.src[0].divides(MV_THREADS_PER_ROW) is not None and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
+          if DEBUG >= 3:
+            print(f"MATVEC: {k.full_shape=} {first_reduce_rng.render()} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
+          if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
+          if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
+          if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
+          return k

  # are we grouping? (requires local shape support)
  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False):
@@ -102,11 +76,8 @@ def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
  # upcast float4 images
  for buf_index,buf in enumerate(k.bufs):
    if isinstance(buf.src[0].dtype, ImageDType):
-      if hasattr(k, "sts"):
-        unit_stride_axes_mul_4 = [i for i in k.sts[buf_index].unit_stride_axes(ignore_valid=True) if k.sts[buf_index].shape[i]%4 == 0]
-      else:
-        # part of real_strides
-        unit_stride_axes_mul_4 = [k.rngs.index(c) for c in k.bufs[buf_index].src[1].split_uop(Ops.ADD) if c.op is Ops.RANGE and (c.vmax+1)%4 == 0]
+      # part of real_strides
+      unit_stride_axes_mul_4 = [k.rngs.index(c) for c in k.bufs[buf_index].src[1].split_uop(Ops.ADD) if c.op is Ops.RANGE and (c.vmax+1)%4 == 0]
      if len(unit_stride_axes_mul_4):
        if (axis:=unit_stride_axes_mul_4[0]) in k.upcastable_dims:
          k.apply_opt(Opt(OptOps.UPCAST, axis, 4))
@@ -114,7 +85,7 @@ def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
          k.apply_opt(Opt(OptOps.UNROLL, k.unrollable_dims.index(axis), 4))

  # no more opt if we are grouping
-  if k.group_for_reduces: return k.applied_opts
+  if k.group_for_reduces: return k

  # **** below this line need to be optional and benchmarked ****

@@ -122,11 +93,9 @@ def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
  to_upcast: list[int] = []
  # upcast leading axes first (hack-ish for winograd; we actually want to upcast masked axes with low stride first)
  for axis in k.upcastable_dims:
-    if isinstance(k, Kernel): is_masked = any(st.axis_is_masked(axis) for st in k.sts)
-    else:
-      # for Schedule, we check if the range is used in INDEX gates or WHERE gates
-      is_masked = any(len(st.src) > 2 and k.rngs[axis] in st.src[2].parents for st in k.bufs) or \
-                  any(any(o is k.rngs[axis] for o in u.src[0].parents) for u in k.ast.parents if u.op is Ops.WHERE)
+    # for Schedule, we check if the range is used in INDEX gates or WHERE gates
+    is_masked = any(len(st.src) > 2 and k.rngs[axis] in st.src[2].parents for st in k.bufs) or \
+                any(any(o is k.rngs[axis] for o in u.src[0].parents) for u in k.ast.parents if u.op is Ops.WHERE)
    if k.full_shape[axis] <= 7 and is_masked and prod(k.full_shape[j] for j in to_upcast) * k.full_shape[axis] <= 7 * 7:
      if DEBUG >= 4: print(f"upcasting masked axis : {axis}")
      to_upcast.append(axis)
@@ -141,24 +110,16 @@ def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
    for axis, upcast_amount in itertools.product(k.upcastable_dims, ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]):
      # if we haven't upcasted it, it mods, and buffer has stride 0 on axis while having no stride 0 in the upcasted axis already
      if axis in upcasted_axis or k.full_shape[axis]%upcast_amount != 0: continue
-      if isinstance(k, Kernel):
-        # must have stride 0 on a view
-        # must have all non stride 0 on what's upcasted before
-        if any(st.views[-1].strides[axis] == 0 and \
-              all(x != 0 for t,x in zip(k.axis_types, st.real_strides()) if t in (AxisType.UPCAST, AxisType.UNROLL)) for st in k.sts):
-          xb_choices.append((sum(st.views[-1].strides[axis]>0 for st in k.sts),
-                            sum(st.views[-1].strides[axis] for st in k.sts), axis, upcast_amount))
-      else:
-        rng = k.rngs[axis]
-        if any(rng not in b.src[1].parents and all(r2 in b.src[1].parents for r2 in k.ranges_of(AxisType.UPCAST, AxisType.UNROLL)) for b in k.bufs):
-          num_strides, sum_strides = 0, 0
-          for b in k.bufs:
-            if rng in b.src[1].parents: num_strides += 1
-            for c in b.src[1].split_uop(Ops.ADD):
-              if c is rng: sum_strides += 1
-              if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].arg
-              if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].arg
-          xb_choices.append((num_strides, sum_strides, axis, upcast_amount))
+      rng = k.rngs[axis]
+      if any(rng not in b.src[1].parents and all(r2 in b.src[1].parents for r2 in k.ranges_of(AxisType.UPCAST, AxisType.UNROLL)) for b in k.bufs):
+        num_strides, sum_strides = 0, 0
+        for b in k.bufs:
+          if rng in b.src[1].parents: num_strides += 1
+          for c in b.src[1].split_uop(Ops.ADD):
+            if c is rng: sum_strides += 1
+            if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].arg
+            if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].arg
+        xb_choices.append((num_strides, sum_strides, axis, upcast_amount))
    if xb_choices:
      xb_choices = sorted(xb_choices)
      if DEBUG >= 4: print(f"more upcast axis : {xb_choices}")
@@ -196,11 +157,8 @@ def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
      k.apply_opt(Opt(OptOps.NOLOCALS))
    else:
      # prioritize making expand axes local
-      if isinstance(k, Kernel):
-        local_axis_ranking = [(any(st.views[-1].strides[axis] == 0 for st in k.sts), axis) for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP)]
-      else:
-        local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].parents for b in k.bufs), axis) \
-                                for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
+      local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].parents for b in k.bufs), axis) \
+                              for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
      to_local: list[tuple[int, int]] = []
      for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])):
        local_size = prod(sz for _, sz in to_local)
@@ -213,4 +171,16 @@ def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
        k.apply_opt(Opt(OptOps.LOCAL, axis, local_sz))
        if will_delete_shape: deleted_shape += 1

-  return k.applied_opts
+  # **** threading ****
+
+  if k.opts.has_threads and k.opts.global_max is not None:
+    for threads in [32,16,12,8,6,5,4,3,2]:
+      # Skip is too many threads. Heuristic: use about 128K ops per thread
+      if threads > k.opts.global_max[0] or resolve(prod(k.full_shape) // (128 << 10) < threads): continue
+      for axis in k.axes_of(AxisType.LOOP):
+        if k.full_shape[axis] % threads == 0:
+          k.apply_opt(Opt(OptOps.THREAD, axis, threads))
+          break
+      if k.applied_opts and k.applied_opts[-1].op is OptOps.THREAD: break
+
+  return k
@@ -1,435 +0,0 @@
-from __future__ import annotations
-import itertools, functools, math
-from dataclasses import dataclass
-from collections import defaultdict
-from typing import cast, Final, Callable, Sequence
-from tinygrad.codegen.opt import OptOps, Opt, KernelOptError, check, axis_letters, axis_colors
-from tinygrad.uop.ops import GroupOp, KernelInfo, UOp, Ops, can_pad, resolve, Variable, sint, graph_rewrite, AxisType
-from tinygrad.uop.spec import type_verify, ast_spec
-from tinygrad.device import Device
-from tinygrad.codegen.opt.tc import TensorCore
-from tinygrad.renderer import Renderer
-from tinygrad.dtype import ImageDType
-from tinygrad.helpers import all_same, colored, ansilen, dedup, prod, round_up, to_function_name, unwrap, argfix, DEBUG
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import strides_for_shape, get_contraction
-from tinygrad.codegen.opt.swizzler import view_left, view_left_through_load
-
-@dataclass
-class TensorCoreOptions:
-  axes: tuple[int, ...] # the location of the original N and M axes if still in the shape
-  axes_exist: tuple[bool, ...] # true if the original N and M axes are still in the shape
-  axis_pads: tuple[tuple[int, int], ...]
-  def fix_axes(self, removed_axis:int): # adjust the TC axes if necessary when a dimension is removed
-    axes, axes_exist = list(self.axes), list(self.axes_exist)
-    for tc_dim in [i for i in range(2) if axes_exist[i]]:
-      if removed_axis < axes[tc_dim]: axes[tc_dim] -= 1
-      elif removed_axis == axes[tc_dim]: axes_exist[tc_dim] = False
-    self.axes, self.axes_exist = tuple(axes), tuple(axes_exist)
-
-class Kernel:
-  def __init__(self, ast:UOp, opts:Renderer|None=None):
-    assert ast.op is Ops.SINK, ast.op
-    self.ast = ast
-
-    self.opts = opts if opts is not None else Device[Device.DEFAULT].renderer
-    # verify AST matches the spec
-    if __debug__: type_verify(list(self.ast.toposort()), ast_spec)
-
-    self.vars: list[Variable] = self.ast.variables()
-    # NOTE: this requires a specific order with the [::-1], this is likely a bug
-    self.bufs: list[UOp] = [x for x in self.ast.toposort() if x.op in GroupOp.Buffer and x.st is not None][::-1]
-
-    # create new shapetrackers inside this kernel, we will permute them
-    self.sts: list[ShapeTracker] = [x.st_arg for x in self.bufs]
-
-    # add the shapetrackers for each reduce
-    # we use this to track which axes are reduced in each reduce
-    self.reduceops = [x for x in self.ast.toposort() if x.op is Ops.REDUCE_AXIS]
-    for x in self.reduceops:
-      self.sts.append(unwrap(x.st))
-      self.sts.append(unwrap(x.src[0].st))
-
-    # add a shapetracker to the end to track the full shape, with 0 strides so it can merge
-    full_shape = ast.full_shape
-    self.sts.append(ShapeTracker.from_shape(full_shape, (0,)*len(full_shape)))
-
-    # parameters for optimization
-    self.tensor_core: TensorCore|None = None
-    self.tensor_core_opts: TensorCoreOptions|None = None
-    self.use_tensor_cores: int = 0
-    self.applied_opts: list[Opt] = []
-    self.dont_use_locals = False
-    self.finalized: bool = False
-
-    # group simplifies
-    self.simplify_ones()
-    self.simplify_merge_adjacent()
-
-    # axis types
-    global_loops = AxisType.GLOBAL if self.opts.has_local else AxisType.LOOP
-    self.axis_types: list[AxisType] = [AxisType.REDUCE if resolve(x!=y) else global_loops for x,y in zip(self.output_shape, self.full_shape)]
-
-    # confirm all reduce axes are at the end
-    if (final_reduces := [x for x in self.axis_types if x == AxisType.REDUCE]) and final_reduces != self.axis_types[-len(final_reduces):]:
-      raise RuntimeError(f"reduces are not at the end of the shape {self.full_shape} -> {self.output_shape}")
-
-  def copy(self):
-    ret = type(self).__new__(type(self))
-
-    # base linearizer params
-    ret.opts, ret.ast = self.opts, self.ast
-
-    # things downstream of the AST
-    ret.reduceops, ret.vars, ret.bufs = self.reduceops, self.vars, self.bufs
-    ret.sts = self.sts[:]
-    ret.axis_types = self.axis_types[:]
-
-    # parameters for optimizations
-    ret.applied_opts, ret.dont_use_locals = self.applied_opts[:], self.dont_use_locals
-    ret.tensor_core, ret.tensor_core_opts, ret.use_tensor_cores = self.tensor_core, self.tensor_core_opts, self.use_tensor_cores
-    ret.finalized = self.finalized
-
-    return ret
-
-  @property
-  def reduceop(self) -> UOp|None: return self.reduceops[0] if len(self.reduceops) > 0 else None
-  @property
-  def full_shape(self) -> tuple[sint, ...]: return self.sts[-1].shape
-
-  @property
-  def output_shape(self) -> tuple[sint, ...]: return self.sts[0].shape
-  @property
-  def shape_len(self) -> int: return len(self.full_shape)
-
-  def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in argfix(axis_type)]
-  @property
-  def upcasted(self) -> int: return len(self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
-  @property
-  def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE))
-
-  # heuristic helpers
-  @property
-  def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP) \
-                                                  if isinstance(s:=self.full_shape[i], int) and s > 1]
-  @property
-  def unrollable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE) \
-                                                  if isinstance(s:=self.full_shape[i], int) and s > 1]
-
-  # ******************** colors and names ********************
-
-  def colors(self) -> list[str]:
-    assert len(self.axis_types) == self.shape_len, "colors size mismatch"
-    return [axis_colors[x] if not self.dont_use_locals or not x == AxisType.GLOBAL else "BLUE" for x in self.axis_types]
-
-  def colored_shape(self, pad:int|None=None, dense=False) -> str:
-    shape_strs = [(s if dense else f"{s:4d}") if isinstance(s, int) else s.render() for s in self.full_shape]
-    ret = ' '.join(colored(s, color) for s,color in zip(shape_strs, self.colors()))
-    if pad: ret += ' '*(pad-ansilen(ret))
-    return ret
-
-  kernel_cnt: Final[defaultdict[str, int]] = defaultdict(int)
-  @functools.cached_property
-  def name(self) -> str:
-    # kernel name (before late upcast)
-    kernel_type = "r" if self.reduceop is not None else ("C" if all(x.op is Ops.SINK or x.op in GroupOp.Buffer for x in self.ast.toposort()) else "E")
-    suffix = colored('_', 'BLACK').join([colored(x.render() if isinstance(x, UOp) else str(x), c) for x,c in zip(self.full_shape, self.colors())])
-    name = kernel_type + (f"{len(self.ast.src)}" if len(self.ast.src) > 1 else "") + "_" + suffix
-
-    # name the function something unique
-    Kernel.kernel_cnt[(function_name := to_function_name(name))] += 1
-    num = f"n{Kernel.kernel_cnt[function_name]-1}" if Kernel.kernel_cnt[function_name] > 1 else ""
-    return name + colored(num, 'BLACK')
-
-  # ******************** base simplifiers ********************
-
-  # apply reshape and permute to all shapetrackers
-  def reshape(self, new_shape_fxn:Callable[[tuple[sint, ...]], Sequence[sint]]):
-    self.sts = [st.reshape(tuple(new_shape_fxn(st.shape))) for st in self.sts]
-  def permute(self, new_axes:Sequence[int]): self.sts = [st.permute(tuple(new_axes)) for st in self.sts]
-
-  # axis : the axis to pull from
-  # amount : the amount to take
-  # top : if you want to pull that amount from the top
-  # insert_at : place to insert the new stuff
-  def shift_to(self, axis:int, amount:int, new_type:AxisType, top:bool=False, insert_at:int|None=None) -> int:
-    if insert_at is None: insert_at = self.shape_len
-    self.axis_types.insert(insert_at, new_type)
-    move_axis = axis if top else axis+1
-    if move_axis < insert_at: insert_at += 1
-    def new_shape_fxn(x): return x[0:axis] + (((amount,x[axis]//amount) if top else (x[axis]//amount,amount)) if x[axis] > 1 else (1,1)) + x[axis+1:]
-    new_axes = [i for i in range(insert_at) if i != move_axis]+[move_axis]+[i for i in range(insert_at, self.shape_len+1) if i != move_axis]
-    self.reshape(new_shape_fxn)
-    self.permute(new_axes)
-    return insert_at
-
-  # ******************** complex simplifiers ********************
-
-  def simplify_ones(self) -> bool:
-    # remove places where the shape is all ones
-    if any(all_ones:=[s==1 for s in self.full_shape]):
-      if hasattr(self, 'axis_types'):
-        self.axis_types = [x for i,x in enumerate(self.axis_types) if not all_ones[i]]
-      self.reshape(lambda shape: [x for i,x in enumerate(shape) if not all_ones[i]])
-      return True
-    return False
-
-  def simplify_merge_adjacent(self):
-    assert not hasattr(self, 'axis_types'), "don't call this after init"
-    if self.shape_len == 0: return
-    shapes, strides = [x.shape for x in self.sts], [x.real_strides() for x in self.sts]
-    # NOTE: we can't use self.first_reduce yet
-    first_reduce = [resolve(x!=y) for x,y in zip(self.output_shape+(0,), self.full_shape+(1,))].index(True)
-
-    # if it's an image, insert fake strides such that this fusion doesn't happen across image axes
-    # TODO: remove membufs
-    membufs = dedup([x.src[0].base for x in self.bufs if x.op in {Ops.LOAD, Ops.STORE}])
-    if isinstance(membufs[0].base.dtype, ImageDType):
-      base_shape = membufs[0].base.dtype.shape
-      if shape_idx_groups := get_contraction(self.output_shape, base_shape):
-        special_strides: tuple[sint, ...] = tuple()
-        for i,g in enumerate(shape_idx_groups):
-          shape_piece = tuple(self.output_shape[x] for x in g)
-          assert prod(shape_piece) == base_shape[i], f"get_contraction was wrong? {shape_piece} != {base_shape[i]}"
-          special_strides += strides_for_shape(shape_piece)
-        # adding the fake image shape
-        shapes.append(self.output_shape)
-        strides.append(special_strides)
-
-    # merge dimensions if we can, multi _merge_dims
-    # NOTE: this does not always preserve the reduce dimension
-    # TODO: move this into shapetracker, with tests!
-    # TODO: how does this work with multi-reduce?
-    rets = [[(s[0], st[0])] for s,st in zip(shapes, strides)]
-    for i in range(1, len(shapes[0])):
-      can_merge = []
-      for s,st,ret in zip(shapes, strides, rets):
-        # TODO: added the always mergeability of 1s, is this right? if so, add to shapetracker in the 1 case
-        si, sti, last_st = s[i], st[i], ret[-1][1]
-        can_merge.append((sti is not None) and ((sti != 0 and last_st == si*sti) or (sti == 0 and last_st == 0)))
-      # more can merge than this
-      mergeable = all(can_merge) and i != first_reduce
-      for j,(s,st) in enumerate(zip(shapes, strides)):
-        if mergeable: rets[j][-1] = (rets[j][-1][0] * s[i], st[i])
-        else: rets[j].append((s[i], st[i]))
-
-    # do the reshapes
-    for i,x in enumerate(rets[:len(self.sts)]): self.sts[i] = self.sts[i].reshape(tuple([y[0] for y in x]))
-
-  # ******************** apply optimizations ********************
-
-  def real_axis(self, op:OptOps, axis:int|None):
-    try:
-      if axis is None: return -1
-      if op is OptOps.UNROLL: return self.unrollable_dims[axis]
-      if op in {OptOps.GROUP, OptOps.GROUPTOP}: return self.axes_of(AxisType.REDUCE)[axis]
-      check(axis < self.shape_len, f"invalid axis on {axis=} {op=} {self.shape_len=}")
-      return axis
-    except IndexError as e: raise KernelOptError from e
-
-  def apply_opt(self, opt:Opt, append_opt:bool=True) -> int|None:
-    if self.finalized: raise RuntimeError("can't optimize Kernel after it's finalized")
-    if self.dont_use_locals: check(opt.op not in {OptOps.LOCAL, OptOps.GROUP, OptOps.GROUPTOP}, "not using locals")
-
-    if opt.op is OptOps.TC:
-      check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: things like PADTO might be fine
-      check(len(self.opts.tensor_cores) > 0, "must have tensor cores")
-      check(opt.axis is not None, "tensor core opts must have an axis")
-      check(opt.arg is not None and isinstance(opt.arg, tuple) and len(opt.arg) == 3, "tensor core opts must have valid arg")
-      check(-1 <= (tc_select:=cast(tuple, opt.arg)[0]) < len(self.opts.tensor_cores), "tensor core opts must have valid tc_select")
-      check(0 <= (tc_opt:=cast(tuple, opt.arg)[1]) <= 2, "tensor core opts must have valid tc_opt")
-      check(0 < (use_tensor_cores:=cast(tuple, opt.arg)[2]) <= 2, "use_tensor_cores value is not valid")
-      check(self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt), "no tensor core available")
-      self.applied_opts.append(opt)
-      return None
-
-    axis = self.real_axis(opt.op, opt.axis)
-
-    if opt.op is OptOps.SWAP: amt = self.real_axis(opt.op, cast(int, opt.arg))  # arg is an axis in the SWAPs
-    elif opt.arg is not None:
-      check(isinstance(opt.arg, int), "arg should be int")
-      amt = arg if (arg:=cast(int, opt.arg)) != 0 else self.full_shape[axis]
-      check(isinstance(amt, int) and amt != 1, f"shift/padto of {amt=}, 1 or symbolic amount is meaningless")
-      if opt.op is not OptOps.PADTO:
-        # we check both the full_shape and each shape
-        check(self.full_shape[axis] % amt == 0, f"no longer valid shift {self.full_shape[axis]=}, {amt=}")
-        for st in self.sts: check(st.shape[axis] == 1 or st.shape[axis] % amt == 0, f"no longer valid shift {st.shape[axis]=}, {amt=}")
-    else: amt = -1
-
-    if self.reduceop is not None and (opt.op in {OptOps.GROUP, OptOps.GROUPTOP} or \
-                                      (self.group_for_reduces and opt.op not in {OptOps.NOLOCALS, OptOps.PADTO})):
-      acc_sz = self.reduceop.dtype.itemsize
-      upcast_sz = prod([self.full_shape[a] for a in self.axes_of(AxisType.UPCAST)])
-      local_sz = prod([self.full_shape[a] for a in self.axes_of(AxisType.LOCAL)])
-      smem_sz = amt*acc_sz*upcast_sz*local_sz
-      check(smem_sz <= self.opts.shared_max, f"exceeds maximum shared memory size: needs {smem_sz}, max {self.opts.shared_max}")
-
-    new_axis = None
-    if opt.op is OptOps.LOCAL:    # cyan
-      # NOTE: LLVM/CPU can use locals too, but they are treated the same as globals (still helpful for L1 cache)
-      # it's disabled for now since it makes BEAM slow for little gain
-      check(self.opts.has_local, "target does not support local")
-      check(self.axis_types[axis] is AxisType.GLOBAL, "local is for globals")
-      new_axis = self.shift_to(axis, amt, AxisType.LOCAL, insert_at=max(self.axes_of(AxisType.GLOBAL, AxisType.LOCAL))+1)
-    elif opt.op in {OptOps.GROUP, OptOps.GROUPTOP}:   # green
-      check(self.opts.has_local and self.opts.has_shared, "target does not support local or shared mem")
-      check(self.axis_types[axis] is AxisType.REDUCE, "must be reduce axis to group")
-      check(not self.tensor_core, "can't group with tensor cores")
-      check(len(reduce_axes:=[i for r in self.reduceops for i in r.axis_arg]) == len(set(reduce_axes)), "can't group with parallel reduces")
-      new_axis = self.shift_to(axis, amt, AxisType.GROUP_REDUCE, top=(opt.op is OptOps.GROUPTOP), insert_at=min(self.axes_of(AxisType.REDUCE)))
-    elif opt.op is OptOps.UNROLL:                     # purple
-      check(self.axis_types[axis] not in (AxisType.UPCAST, AxisType.UNROLL), "can't upcasted already upcasted")
-      check(amt <= 32, "don't unroll more than 32")
-      new_axis = self.shift_to(axis, amt, AxisType.UNROLL, insert_at=None)
-    elif opt.op is OptOps.UPCAST:                     # yellow
-      check(axis in self.upcastable_dims, f"{axis=} not in {self.upcastable_dims=}")
-      # NOTE: assume the first get_local_axes() LOCAL are for TC
-      check(not (self.tensor_core and axis in self.axes_of(AxisType.LOCAL)[:len(self.tensor_core.get_local_axes())]), "can't upcast TC locals")
-      check((self.opts is not None and self.opts.device == "DSP") or amt <= 16, "don't upcast more than 16")
-      new_axis = self.shift_to(axis, amt, AxisType.UPCAST,
-                               insert_at=max(self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP, AxisType.UPCAST))+1)
-    elif opt.op is OptOps.NOLOCALS:
-      check(self.opts.has_local and not self.dont_use_locals, "NOLOCALS is meaningless if target does not support local or already not using locals")
-      check(AxisType.LOCAL not in self.axis_types and self.group_for_reduces == 0, "can't have no locals with locals")
-      self.dont_use_locals = True
-    elif opt.op is OptOps.SWAP:
-      check(axis < amt, f"swap is only for axis < amt, getting {amt=}, {axis=}")
-      check(self.axis_types[axis]==self.axis_types[amt]==AxisType.GLOBAL, f"swap is for globals {self.axis_types[axis]=}, {self.axis_types[amt]=}")
-      permute = list(range(self.shape_len))
-      permute[axis], permute[amt] = permute[amt], permute[axis]
-      self.permute(tuple(permute))
-    elif opt.op is OptOps.PADTO:
-      check(not self.vars, "does not work with symbolic shape")
-      check(self.axis_types[axis] not in (AxisType.UPCAST, AxisType.UNROLL), "cannot pad upcasted")
-      # ok to pad SUM if all parent ALU ops have f(0) = 0
-      if (r:=self.reduceop) is not None and self.axis_types[axis] in (AxisType.GROUP_REDUCE, AxisType.REDUCE):
-        check(r.arg[0] is Ops.ADD and can_pad(r, {}), f"cannot pad {r}")
-      padded = False
-      for i,st in enumerate(self.sts):
-        if (s:=st.shape[axis]) == 1: continue  # reduced
-        check(s > amt//4, f"pad adds more than quadruple the work {st.shape[axis]=} > {amt//4=}")
-        if (ru := round_up(cast(int, s), amt) - s):
-          # pad right seems to be faster
-          self.sts[i] = st.pad(((0,0),) * axis + ((0,ru),) + ((0,0),) * (len(st.shape)-axis-1))
-          padded = True
-      check(padded, "nothing was padded")
-
-    if append_opt: self.applied_opts.append(opt)
-    if self.simplify_ones() and self.tensor_core_opts:
-      self.tensor_core_opts.fix_axes(axis) # fix up axes in TC opts if required after simplify_ones()
-    return new_axis
-
-  def apply_opts(self, opts:Sequence[Opt]) -> Kernel:
-    for opt in opts: self.apply_opt(opt)
-    return self
-
-  # **** kernel outputs, mostly tensor cores ****
-
-  def _create_tc_opts(self, reduceop:UOp, tc:TensorCore, axis:int, opt_level:int) -> TensorCoreOptions|None:
-    has_cast = tc.dtype_in != tc.dtype_out
-    if has_cast and not (reduceop.src[0].op is Ops.CAST and reduceop.src[0].dtype == tc.dtype_out): return None
-
-    mul_op = reduceop.src[0].src[0] if has_cast else reduceop.src[0]
-    if mul_op.op is not Ops.MUL: return None
-
-    def buf_index(src:UOp) -> int|None:
-      # TODO: apply tc even if the sources are not from LOAD
-      if src.op is Ops.LOAD and src.dtype == tc.dtype_in: return self.bufs.index(src)
-      try:
-        if opt_level >= 1 and src.op is Ops.CAST and src.dtype == tc.dtype_in: return self.bufs.index(src.src[0])
-      except ValueError: return None
-      return None
-    if (buf0:=buf_index(mul_op.src[0])) is None or (buf1:=buf_index(mul_op.src[1])) is None: return None
-
-    buf0_strides, buf1_strides = self.sts[buf0].real_strides(), self.sts[buf1].real_strides()
-    axis_buf0 = [(i,self.full_shape[i],buf1_strides[i]) for i in self.upcastable_dims if buf0_strides[i] == 0]
-    axis_buf1 = [(i,self.full_shape[i],buf0_strides[i]) for i in self.upcastable_dims if buf1_strides[i] == 0]
-    if not (axis_buf0 and axis_buf1 and (len(self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE)) == 1 or (opt_level >= 1))): return None
-
-    axis_choices = list(itertools.product(axis_buf0, axis_buf1, self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE)))
-    if not (axis < len(axis_choices)): return None
-
-    s0, s1, s2 = axis_choices[-(axis+1)][0][0], axis_choices[-(axis+1)][1][0], axis_choices[-(axis+1)][2]  # s0 is n, s1 is m, s2 is k
-    axis_pads = tuple((x, tc.dims[i]) for i, x in enumerate([s0, s1, s2]) if resolve(self.full_shape[x]%tc.dims[i] != 0))
-    if axis_pads and (opt_level < 2): return None
-    if DEBUG >= 3: print("TENSOR CORES", axis_buf0, axis_buf1, tc)
-    return TensorCoreOptions(axes=(s0, s1, s2), axes_exist=(True, True), axis_pads=axis_pads)
-
-  def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> bool:
-    if use_tensor_cores and self.reduceop is not None and self.reduceop.arg[0] is Ops.ADD:
-      tensor_cores = self.opts.tensor_cores if tc_select == -1 else [self.opts.tensor_cores[tc_select]]
-      for tc in tensor_cores:
-        tensor_core_opts = [self._create_tc_opts(reduceop, tc, axis, opt_level) for reduceop in self.reduceops]
-        if tensor_core_opts[0] is None: continue
-        # can only fuse reduces with the same tc options
-        assert all_same(tensor_core_opts)
-        self.tensor_core_opts = tc_opts = tensor_core_opts[0]
-
-        # attempt to pad the tensor axes that require it
-        try:
-          for axis, dim in tc_opts.axis_pads: self.apply_opt(Opt(OptOps.PADTO, axis, dim), append_opt=False) # PADTO might fail
-        except KernelOptError: continue
-        # tensor core -- unroll the reduce dim (K), upcast and local the inner and outer dims (N, M)
-        for opt in tc.opts: self.apply_opt(Opt({"u":OptOps.UPCAST, "l":OptOps.LOCAL}[opt[0]], tc_opts.axes[int(opt[1])], 2), append_opt=False)
-        for dim, amt in tc.get_reduce_axes(): self.apply_opt(Opt(OptOps.UNROLL, 0, amt), append_opt=False) # TODO: this should be the reduce, not 0
-        self.tensor_core = tc
-        self.use_tensor_cores = use_tensor_cores  # TC=2 will do the shape ops without the WMMA
-        return True
-    return False
-
-  # strings like ['g0', 'g1', 'l0', 'l1', 'l2', 'l3', 'l4', 'l5', 'R0', 'r0', 'r1', 'r2', 'u0', 'u1', 'u2']
-  def shape_str(self) -> list[str]:
-    ret: list[str] = []
-    cnt: dict[AxisType, int] = {}
-    for x in self.axis_types:
-      cnt[x] = (cnt[x] + 1) if x in cnt else 0
-      ret.append(f"{axis_letters[x]}{cnt[x]}")
-    return ret
-  def shape_str_to_axis(self, nms:list[str]) -> tuple[int, ...]: return tuple([self.shape_str().index(x) for x in nms])
-
-  def get_optimized_ast(self, name_override:str|None=None) -> UOp:
-    @functools.cache
-    def fixup_ast(op:UOp) -> UOp:
-      ret = op.replace(src=tuple(fixup_ast(x) for x in op.src)) # noqa: F821
-      if op.op in GroupOp.Buffer and op in self.bufs:
-        st = self.sts[self.bufs.index(op)]
-        # replace the VIEW source
-        return ret.replace(src=(ret.src[0].replace(arg=st),)+ret.src[1:])
-      if op.op is Ops.SINK:
-        # NOTE: should group_for_reduces be added to the local_dims?
-        # TODO: arg.name should be able to be None
-        kernel_name = ret.arg.name if ret.arg is not None and ret.arg.name != "test" else self.name if name_override is None else name_override
-        return ret.replace(arg=KernelInfo(kernel_name, tuple(self.axis_types), self.dont_use_locals, tuple(self.applied_opts)))
-      if op.op is Ops.REDUCE_AXIS:
-        reduce_idx = len(self.bufs) + self.reduceops.index(op) * 2
-        changed = tuple(i for i in range(self.shape_len) if resolve(self.sts[reduce_idx].shape[i] != self.sts[reduce_idx + 1].shape[i]))
-        axes = tuple(i for i in self.axes_of(AxisType.REDUCE, AxisType.GROUP_REDUCE, AxisType.UNROLL) if i in changed)
-        if (tc := self.tensor_core) and self.use_tensor_cores == 1:
-          # get reduce/upcast axes for the tensor cores
-          tc_reduce_axes = self.shape_str_to_axis([f"r{i}" for i in range(len(tc.get_reduce_axes()))])
-          base_upcast_axes = tuple([(s,2) for s in self.shape_str_to_axis(tc.base_upcast_axes())])
-          tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
-
-          # permute the srcs
-          srcs = list((ret.src[0] if ret.src[0].op is not Ops.CAST else ret.src[0].src[0]).src)
-          for i, (src, permaxis) in enumerate(zip(srcs, tc.permutes_for_shape_str(self.shape_str()))):
-            src_st = (src if src.op is Ops.LOAD else src.src[0]).st_arg
-            srcs[i] = src.view(ShapeTracker.from_shape(src_st.shape).permute(permaxis))
-
-          # construct the op
-          wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.opts.device, tc.threads, tc_upcast_axes, tc_reduce_axes)
-          wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
-            UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0]),
-            UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1]),
-            UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg)
-          tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2])
-
-          # preserve any other reduce
-          return ret.replace(src=(tc_uop,), arg=(Ops.ADD, new_axes)) if (new_axes := tuple(i for i in axes if i not in tc_reduce_axes)) else tc_uop
-
-        ret = ret.replace(arg = (op.arg[0], axes))
-      return ret
-    self.finalized = True
-    fixed_ast = fixup_ast(self.ast)
-    del fixup_ast
-    return graph_rewrite(fixed_ast, view_left+view_left_through_load, name="fixup optimized AST")
@@ -1,34 +1,20 @@
 from __future__ import annotations
 import math, itertools
 from collections import defaultdict
-from typing import cast, Final, Sequence
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, _substitute, AxisType, ssimplify, can_pad
-from tinygrad.uop.symbolic import symbolic_flat
+from typing import cast, Final
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, can_pad
 from tinygrad.device import Buffer
 from tinygrad.dtype import AddrSpace, dtypes, ImageDType
 from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod
 from tinygrad.codegen.opt import axis_colors, Opt, OptOps, KernelOptError, check, axis_letters
+from tinygrad.codegen.simplify import pm_flatten_range
 from tinygrad.renderer import Renderer
 from tinygrad.schedule.rangeify import remove_tags

 # NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
-axis_to_pos = {AxisType.LOOP: -1, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
+axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
               AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}

-def flatten_range(r:UOp):
-  off = 2 if r.op is Ops.STORE else 1
-  rngs = r.src[off:]
-  if not len(rngs): return None
-  new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
-  return r.replace(src=r.src[:off]+tuple(new_rngs))
-
-pm_flatten_range = PatternMatcher([
-  # real ranges only
-  (UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range),
-])
-
-def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
-
 class Scheduler:
  def __init__(self, ast:UOp, opts:Renderer):
    self.ast, self.opts = ast, opts
@@ -58,17 +44,11 @@ class Scheduler:
    return ret
  def shape_str_to_axis(self, nms:list[str]) -> tuple[int, ...]: return tuple([self.shape_str().index(x) for x in nms])

-  @property
-  def termination(self):
-    terminators = [u for u in self.ast.parents if u.op in {Ops.REDUCE, Ops.STORE}]
-    termination = {}
-    for t in terminators:
-      # works without pm_flatten_range
-      for u in UOp.sink(*t.src[1 if t.op is Ops.REDUCE else 2:]).parents:
-        if u.op is Ops.RANGE: termination[u] = t
-    return termination
-
-  def copy(self): return Scheduler(self.get_optimized_ast(), self.opts)
+  def copy(self):
+    ret = Scheduler(self.ast, self.opts)
+    ret.dont_use_locals = self.dont_use_locals
+    ret.applied_opts = self.applied_opts[:]
+    return ret

  kernel_cnt: Final[defaultdict[str, int]] = defaultdict(int)
  def get_optimized_ast(self, name_override:str|None=None):
@@ -82,8 +62,7 @@ class Scheduler:
    self.ast = graph_rewrite(self.ast, pm_flatten_range, name="flatten range")
    return self.ast.replace(arg=KernelInfo(name=name, applied_opts=tuple(self.applied_opts), dont_use_locals=self.dont_use_locals), tag=1)

-  def convert_loop_to_global(self):
-    if not self.opts.has_local: return None
+  def _globalizable_rngs(self) -> list[UOp]:
    store_rngs = self.ast.src[0].src[2:]

    # filter any not in local stores
@@ -91,29 +70,16 @@ class Scheduler:
                        or (x.op is Ops.BUFFERIZE and x.arg == AddrSpace.LOCAL)]
    for ls in local_store_rngs: store_rngs = tuple([x for x in store_rngs if x in ls])

-    store_rng = [x for x in UOp.sink(*store_rngs).toposort() if x.op is Ops.RANGE] if store_rngs else []
-    rng = [x.replace(arg=(x.arg[0], AxisType.GLOBAL)) if x.arg[1] == AxisType.LOOP and x in store_rng else x for x in self.rngs]
+    return [x for x in UOp.sink(*store_rngs).toposort() if x.op is Ops.RANGE and x.arg[1] == AxisType.LOOP] if store_rngs else []
+
+  def convert_loop_to_global(self):
+    if not self.opts.has_local: return None
+
+    globalizible_rngs = self._globalizable_rngs()
+    rng = [x.replace(arg=(x.arg[0], AxisType.GLOBAL)) if x in globalizible_rngs else x for x in self.rngs]

    self.ast = self.ast.substitute(dict(zip(self.rngs, rng)))

-  def simplify_merge_adjacent(self):
-    i = 0
-    while i < len(self.rngs)-1:
-      r0, r1 = self.rngs[i], self.rngs[i+1]
-      # same axistype and same termination
-      termination = self.termination
-      if r0.arg[1] == r1.arg[1] and r0 in termination and r1 in termination and termination[r0] == termination[r1]:
-        s0, s1 = r0.src[0], r1.src[0]
-        # do the merge
-        oidx = self.ast.simplify()
-        new_range = r0.replace(src=(s0*s1,))
-        nidx = graph_rewrite(oidx, _substitute+symbolic_flat+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1}, name=f"check_merge_{i}_{i+1}")
-        # check if it simplifies
-        if count_divmod(nidx) <= count_divmod(oidx):
-          self.ast = nidx
-          continue
-      i += 1
-
  def colors(self) -> list[str]: return [axis_colors[x] if not self.dont_use_locals or not x == AxisType.GLOBAL else "BLUE" for x in self.axis_types]
  def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():>4s}', color) for x,color in zip(self.rngs, self.colors())])

@@ -123,7 +89,7 @@ class Scheduler:
    new_rng = UOp.range(amount, self.maxarg+1, new_type) if input_new_rng is None else input_new_rng
    replaced_rng = rng.replace(src=(UOp.const(dtypes.int, old_sz),))
    sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
-    self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[0]} {amount}")
+    self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[0]} {amount} {str(new_type).split('.')[1].lower()}")
    return replaced_rng, new_rng

  def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
@@ -146,10 +112,6 @@ class Scheduler:
      return axis
    except IndexError as e: raise KernelOptError from e

-  def apply_opts(self, opts:Sequence[Opt]) -> Scheduler:
-    for opt in opts: self.apply_opt(opt)
-    return self
-
  def apply_opt(self, opt:Opt, append_opt:bool=True):
    if opt.op is OptOps.NOLOCALS:
      check(all(x not in {AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE} for x in self.axis_types), "no locals can't have locals")
@@ -165,7 +127,7 @@ class Scheduler:
    opt_to_at = {
      OptOps.LOCAL: AxisType.LOCAL, OptOps.UPCAST: AxisType.UPCAST,
      OptOps.UNROLL: AxisType.UNROLL, OptOps.GROUP: AxisType.GROUP_REDUCE,
-      OptOps.GROUPTOP: AxisType.GROUP_REDUCE}
+      OptOps.GROUPTOP: AxisType.GROUP_REDUCE, OptOps.THREAD: AxisType.THREAD}

    ret = None
    if opt.op in opt_to_at:
@@ -187,11 +149,16 @@ class Scheduler:
      if opt.op is OptOps.LOCAL:
        check(not self.dont_use_locals, "can't use locals")
        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOOP}, "local is for globals")
+      if opt.op is OptOps.THREAD:
+        check(self.opts is not None and self.opts.has_threads, "target does not support threads")
+        check(self.opts is not None and self.opts.global_max is not None and amt <= self.opts.global_max[0], "too many threads")
+        check(all(x is not AxisType.THREAD for x in self.axis_types), "already threaded")
+        check(rng in self._globalizable_rngs(), "can't apply range to this dim")
      if opt.op in {OptOps.GROUP, OptOps.GROUPTOP}:
        check(all(x.op is not OptOps.TC for x in self.applied_opts), "no grouping with tensor cores")  # TODO: why is this wrong?
        check(not self.dont_use_locals, "can't use locals")
        check(rng.arg[-1] == AxisType.REDUCE, "group is for reduce")
-      ret = self.shift_to(rng, amt, opt_to_at[opt.op], top=opt.op==OptOps.GROUPTOP)
+      ret = self.shift_to(rng, amt, opt_to_at[opt.op], top=opt.op in {OptOps.GROUPTOP, OptOps.THREAD})
    elif opt.op is OptOps.TC:
      check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: remove the need for this by having warps
      check(opt.axis is not None, "tensor core opts must have an axis")
@@ -204,6 +171,7 @@ class Scheduler:
    elif opt.op is OptOps.PADTO:
      check(rng.src[0].op is Ops.CONST, "only pad const axes")
      check(rng.arg[-1] not in {AxisType.UPCAST, AxisType.UNROLL}, "cannot pad upcasted") # TODO: why is this wrong?
+      check(rng.arg[-1] is not AxisType.THREAD, "cannot pad thread")
      # ok to pad SUM if all parent ALU ops have f(0) = 0
      if (r:=self.reduceop) is not None and rng.arg[-1] in (AxisType.GROUP_REDUCE, AxisType.REDUCE):
        check(r.arg[0] is Ops.ADD and can_pad(r, {}), f"cannot pad {r}")
@@ -347,7 +315,6 @@ def apply_opts(ctx:Renderer, ast:UOp):
  if ast.tag is not None: return None
  k = Scheduler(ast, ctx)
  k.convert_loop_to_global()
-  k.simplify_merge_adjacent()
  if BEAM >= 1:
    from tinygrad.codegen.opt.search import beam_search
    rawbufs = bufs_from_ast(ast, ctx.device)
@@ -358,7 +325,7 @@ def apply_opts(ctx:Renderer, ast:UOp):
    from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
    # NOTE: hand_coded_optimizations doesn't support multiblock opts yet
    if all(len(u.src) == 1 for u in ast.parents if u.op is Ops.LOAD):
-      for opt in hand_coded_optimizations(k): k.apply_opt(opt)
+      k = hand_coded_optimizations(k)
  return k.get_optimized_ast(name_override=ast.arg.name if ast.arg is not None and ast.arg.name != "test" else None)

 pm_postrange_opt = PatternMatcher([
@@ -1,19 +1,14 @@
 from typing import cast
 import functools, math, time, multiprocessing, traceback, signal, atexit
-from collections import defaultdict
 from dataclasses import replace
-from tinygrad.uop.ops import UOp, Ops, Variable, sym_infer, AxisType, pyrender
+from tinygrad.uop.ops import sym_infer, AxisType, pyrender
 from tinygrad.device import Device, Buffer, Compiler
 from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str
 from tinygrad.helpers import IGNORE_BEAM_CACHE
-from tinygrad.dtype import ImageDType, PtrDType
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.tensor import Tensor
 from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.renderer import ProgramSpec
-
-# both versions
-from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.codegen.opt.postrange import Scheduler

 actions = [Opt(op=OptOps.UPCAST, axis=axis, arg=amt) for amt in [0,2,3,4,5,7] for axis in range(8)]
@@ -21,12 +16,13 @@ actions += [Opt(op=OptOps.UNROLL, axis=axis, arg=amt) for amt in [0,4,7] for axi
 actions += [Opt(op=OptOps.LOCAL, axis=axis, arg=amt) for amt in [2,3,4,8,13,16,29] for axis in range(6)]
 actions += [Opt(op=OptOps.GROUPTOP, axis=axis, arg=amt) for amt in [13,16,28,29,32,49,64,256] for axis in range(3)]
 actions += [Opt(op=OptOps.GROUP, axis=axis, arg=amt) for amt in [0,4,8,16] for axis in range(3)]
-if getenv("BEAM_PADTO", 1): actions += [Opt(op=OptOps.PADTO, axis=axis, arg=amt) for amt in [32] for axis in range(7)]
+if getenv("BEAM_PADTO", 0): actions += [Opt(op=OptOps.PADTO, axis=axis, arg=amt) for amt in [32] for axis in range(7)]
 actions += [Opt(op=OptOps.LOCAL, axis=0, arg=32), Opt(op=OptOps.LOCAL, axis=6, arg=2)]
 actions += [Opt(op=OptOps.TC, axis=0, arg=(-1, 0, getenv("TC", 1)))]
 # covers resnet kernels (3 global * 3 reduce)
 actions += [Opt(op=OptOps.TC, axis=axis, arg=(-1, getenv("TC_OPT", 2), getenv("TC", 1))) for axis in range(9)]
 actions += [Opt(op=OptOps.SWAP, axis=axis_0, arg=axis_1) for axis_0 in range(5) for axis_1 in range(axis_0+1, 5)]
+actions += [Opt(op=OptOps.THREAD, axis=axis, arg=amt) for amt in [2,3,4,5,8,12,16,24,32,64] for axis in range(3)]
 if getenv("NOLOCALS"): actions += [Opt(op=OptOps.NOLOCALS)]

 def get_test_global_size(global_size, max_global_size, var_vals):
@@ -39,7 +35,7 @@ def get_test_global_size(global_size, max_global_size, var_vals):
        break
  return test_global_size, input_size / prod(test_global_size)

-def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[Variable, int], rawbufs:list[Buffer], early_stop:float|None=None,
+def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[str, int], rawbufs:list[Buffer], early_stop:float|None=None,
                  allow_test_size:int=True, max_global_size:int|None=65536, clear_l2=False, cnt=3, name="test") -> list[float]:
  factor = 1
  if allow_test_size and p.global_size is not None and max_global_size is not None:
@@ -63,7 +59,7 @@ def timeout_handler(signum, frame):
  if DEBUG >= 2: print("*** BEAM COMPILE TIMEOUT")
  raise TimeoutException()

-def _try_compile_linearized_w_idx(x:tuple[int,Kernel], compiler:Compiler) -> tuple[int, tuple[ProgramSpec, bytes, float]|None]:
+def _try_compile_linearized_w_idx(x:tuple[int,Scheduler], compiler:Compiler) -> tuple[int, tuple[ProgramSpec, bytes, float]|None]:
  if hasattr(signal, "alarm"):
    signal.signal(getattr(signal, 'SIGALRM'), timeout_handler)
    # set timeout
@@ -96,26 +92,8 @@ def _ensure_buffer_alloc(bufs:list[Buffer]) -> list[Buffer]: return [buf.ensure_

 # *** external API ***

-# get (scrap) buffers for timing the linearizer
-# NOTE: there's also bufs_from_ast in postrange
-def bufs_from_lin(lin:Kernel, allocate:bool=True) -> list[Buffer]:
-  bufsts: defaultdict[int, list[UOp]] = defaultdict(list)
-  for x in lin.bufs:
-    if x.src[0].base.op is Ops.DEFINE_GLOBAL: bufsts[x.src[0].base.arg].append(x)
-  # TODO: Nones are staying in here if buffers are optimized out!
-  # TODO: add a test for this
-  rawbufs: list[Buffer|None] = [None]*(max(bufsts)+1)
-  for k,lx in bufsts.items():
-    buf_size = prod(dtype.shape) if isinstance(dtype:=lx[0].src[0].dtype, ImageDType) else max(y.st_arg.real_size() for y in lx)
-    assert isinstance(dtype, (PtrDType, ImageDType))
-    if buf_size == 0: buf_size = 1  # create a size 1 buffer if no cell is accessed in kernel. # TODO: remove from kernel input in this case.
-    buf_dtype = dtype if isinstance(dtype, ImageDType) else dtype.base
-    rawbufs[k] = Buffer(lin.opts.device, buf_size, buf_dtype).allocate() if allocate else Buffer(lin.opts.device, buf_size, buf_dtype)
-  #assert all(r is not None for r in rawbufs)
-  return cast(list[Buffer], rawbufs)
-
 # get dictionary of all possible actions
-def get_kernel_actions(lin:Kernel|Scheduler, include_0=True, candidates:list[Opt]|None=None) -> dict[int, Kernel|Scheduler]:
+def get_kernel_actions(lin:Scheduler, include_0=True, candidates:list[Opt]|None=None) -> dict[int, Scheduler]:
  acted_lins, max_up, max_lcl = {0:lin} if include_0 else {}, getenv("BEAM_UPCAST_MAX", 256), getenv("BEAM_LOCAL_MAX", 1024)
  kernel_actions = (actions if candidates is None else candidates).copy()

@@ -139,7 +117,7 @@ def get_kernel_actions(lin:Kernel|Scheduler, include_0=True, candidates:list[Opt
  return acted_lins

 beam_pool, BEAM_DEBUG = None, getenv("BEAM_DEBUG")
-def beam_search(lin:Kernel|Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value):
+def beam_search(lin:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value):
  global beam_pool
  key = {"ast": lin.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": lin.opts.device, "suffix": lin.opts.suffix}
  if not disable_cache and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
@@ -147,7 +125,7 @@ def beam_search(lin:Kernel|Scheduler, rawbufs:list[Buffer], amt:int, allow_test_
    for o in val[len(lin.applied_opts):]: ret.apply_opt(o)
    return ret

-  beam: list[tuple[Kernel|Scheduler, float]] = [(lin, float("inf"))]
+  beam: list[tuple[Scheduler, float]] = [(lin, float("inf"))]
  seen_libs = set()

  default_parallel = multiprocessing.cpu_count() if lin.opts.device in {"CUDA", "AMD", "NV", "METAL", "HIP"} else 0
@@ -164,12 +142,12 @@ def beam_search(lin:Kernel|Scheduler, rawbufs:list[Buffer], amt:int, allow_test_

  try:
    rawbufs = _ensure_buffer_alloc(rawbufs)
-    var_vals: dict[Variable, int] = {k:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
+    var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
    exiting, st = False, time.perf_counter()
    dev = Device[lin.opts.device]
    while not exiting:
-      acted_lins: list[Kernel|Scheduler] = flatten([get_kernel_actions(lin, include_0=False).values() for lin,_ in beam])
-      timed_lins: list[tuple[Kernel|Scheduler, float]] = []
+      acted_lins: list[Scheduler] = flatten([get_kernel_actions(lin, include_0=False).values() for lin,_ in beam])
+      timed_lins: list[tuple[Scheduler, float]] = []
      _compile_fn = functools.partial(_try_compile_linearized_w_idx, compiler=dev.compiler)
      least_compute_ops = math.inf
      for i,proc in (map(_compile_fn, enumerate(acted_lins)) if beam_pool is None else beam_pool.imap_unordered(_compile_fn, enumerate(acted_lins))):
@@ -0,0 +1,37 @@
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute
+from tinygrad.uop.symbolic import symbolic_flat
+
+def flatten_range(r:UOp):
+  off = 2 if r.op is Ops.STORE else 1
+  rngs = r.src[off:]
+  if not len(rngs): return None
+  new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
+  return r.replace(src=r.src[:off]+tuple(new_rngs))
+
+pm_flatten_range = PatternMatcher([
+  # real ranges only
+  (UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range),
+])
+
+def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
+def simplify_merge_adjacent(u:UOp) -> UOp|None:
+  i = 2 if u.op is Ops.STORE else 1
+  while i < len(u.src)-1:
+    r0, r1 = u.src[i], u.src[i+1]
+    # check same type
+    if r0.arg[-1] == r1.arg[-1]:
+      s0, s1 = r0.src[0], r1.src[0]
+      # do the merge
+      new_range = r0.replace(src=(s0*s1,))
+      nidx = graph_rewrite(u, _substitute+symbolic_flat+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
+                           name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
+      # check if it simplifies
+      if count_divmod(nidx) <= count_divmod(u):
+        u = nidx
+        continue
+    i += 1
+  return u
+
+pm_simplify_ranges = PatternMatcher([
+  (UPat((Ops.STORE, Ops.REDUCE), name="u"), simplify_merge_adjacent),
+])
@@ -298,6 +298,7 @@ class Compiled:

 # TODO: move this to each Device
 def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
+  if dtype == dtypes.index: return False
  if device is None: device = Device.DEFAULT
  if dtype == dtypes.bfloat16:
    if device == "METAL": return not CI
@@ -89,7 +89,7 @@ class dtypes:
  def is_float(x: DType) -> bool: return x.scalar() in dtypes.floats or isinstance(x, ImageDType)
  @staticmethod # static methods on top, or bool in the type info will refer to dtypes.bool
  @functools.cache
-  def is_int(x: DType) -> bool: return x.scalar() in dtypes.ints
+  def is_int(x: DType) -> bool: return x.scalar() in dtypes.ints + (dtypes.index,)
  @staticmethod
  @functools.cache
  def is_unsigned(x: DType) -> bool: return x.scalar() in dtypes.uints
@@ -128,6 +128,7 @@ class dtypes:
  @staticmethod
  def fields() -> dict[str, DType]: return DTYPES_DICT
  void: Final[DType] = DType.new(-1, 0, "void", None)
+  index: Final[DType] = DType.new(-1,100, "index", None)
  bool: Final[DType] = DType.new(0, 1, "bool", '?')
  int8: Final[DType] = DType.new(1, 1, "signed char", 'b')
  uint8: Final[DType] = DType.new(2, 1, "unsigned char", 'B')
@@ -164,7 +165,7 @@ class dtypes:
  uints = (uint8, uint16, uint32, uint64)
  sints = (int8, int16, int32, int64)
  ints = uints + sints
-  all = floats + ints + (bool,)
+  all = floats + ints + (bool, index)

 if (env_default_float := getenv("DEFAULT_FLOAT", "")):
  dtypes.default_float = getattr(dtypes, env_default_float.lower())
@@ -186,11 +187,12 @@ def _get_recursive_parents(dtype:DType) -> set[DType]:
  return set.union(*[_get_recursive_parents(d) for d in promo_lattice[dtype]], {dtype}) if dtype != dtypes.float64 else {dtypes.float64}
@functools.cache
 def least_upper_dtype(*ds:DType) -> DType:
-  return min(set.intersection(*[_get_recursive_parents(d) for d in ds])) if not (images:=[d for d in ds if isinstance(d, ImageDType)]) else images[0]
+  return min(set.intersection(*[_get_recursive_parents(d.scalar()) for d in ds])) \
+      if not (images:=[d for d in ds if isinstance(d, ImageDType)]) else images[0]
 def least_upper_float(dt:DType) -> DType: return dt if dtypes.is_float(dt) else least_upper_dtype(dt, dtypes.default_float)

-DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void"))}
-INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void"}
+DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "index"))}
+INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void", "index":"index"}

@functools.cache
 def can_safe_cast(dt0:DType, dt1:DType) -> bool:
@@ -198,6 +200,7 @@ def can_safe_cast(dt0:DType, dt1:DType) -> bool:
  # https://numpy.org/doc/stable/reference/generated/numpy.can_cast.html
  if dt0 == dt1 or dt0 == dtypes.bool: return True
  match dt1:
+    case dtypes.index: return dt0 in dtypes.ints
    case dtypes.double: return dt0 in (dtypes.float, dtypes.half, dtypes.bfloat16,
      dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8)
    case dtypes.float: return dt0 in (dtypes.half, dtypes.bfloat16, dtypes.uint16, dtypes.uint8, dtypes.int16, dtypes.int8)
@@ -315,4 +318,4 @@ def _to_torch_dtype(dtype:DType) -> 'torch.dtype'|None:  # type: ignore [name-de
  except TypeError: return None
@functools.cache
 def _from_torch_dtype(torchdtype:'torch.dtype') -> DType: # type: ignore [name-defined] # noqa: F821
-  return {v:k for k in dtypes.all if (v:=_to_torch_dtype(k)) is not None}[torchdtype]
+  return {v:k for k in DTYPES_DICT.values() if (v:=_to_torch_dtype(k)) is not None}[torchdtype]
@@ -16,7 +16,7 @@ class GraphException(Exception): pass

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

-def apply_graph_to_jit(jit_cache: list[ExecItem], input_rawbuffers: list[Buffer], var_vals: dict[Variable, int], max_batch_size=0) -> list[ExecItem]:
+def apply_graph_to_jit(jit_cache: list[ExecItem], input_rawbuffers: list[Buffer], var_vals: dict[str, int], max_batch_size=0) -> list[ExecItem]:
  # Split JIT cache into batches for faster graph execution.
  # This allows the accelerator to run some batches while subsequent graphs are still being updated.
  graphed_jit_cache: list[ExecItem] = []
@@ -73,7 +73,7 @@ def get_input_replace(jit_cache: list[ExecItem], input_rawbuffers:list[Buffer])
  return input_replace

 class GraphRunner(Runner):
-  def __init__(self, jit_cache: list[ExecItem], input_rawbuffers: list[Buffer], var_vals: dict[Variable, int]):
+  def __init__(self, jit_cache: list[ExecItem], input_rawbuffers: list[Buffer], var_vals: dict[str, int]):
    self.jit_cache = jit_cache  # NOTE: this is not used, but you have to keep these objects alive for the Graph
    self.input_replace:dict[tuple[int, int], int] = get_input_replace(jit_cache, input_rawbuffers)
    self.var_vals_replace:dict[int, list[tuple[int, int]]] = {}
@@ -82,7 +82,7 @@ class GraphRunner(Runner):

    def is_sym_dim(dim) -> bool: return not all(isinstance(d, (int, float)) for d in dim)

-    self.vars = sorted(var_vals.keys(), key=lambda v: v.expr)
+    self.vars = sorted(var_vals.keys())
    self.symbolic_dims = dedup([tuple(d) for ji in jit_cache if isinstance(ji.prg, CompiledRunner) and (d:=ji.prg.p.local_size) and is_sym_dim(d)] +
                               [tuple(d) for ji in jit_cache if isinstance(ji.prg, CompiledRunner) and (d:=ji.prg.p.global_size) and is_sym_dim(d)])
    def find_symbolic_dim(dim): return self.symbolic_dims.index(tuple(dim)) if dim is not None and tuple(dim) in self.symbolic_dims else None
@@ -91,7 +91,7 @@ class GraphRunner(Runner):
    for j,ji in enumerate(jit_cache):
      estimates += ji.prg.estimates
      if isinstance(ji.prg, CompiledRunner):
-        if ji.prg.p.vars: self.var_vals_replace[j] = [(i, self.vars.index(v)) for i, v in enumerate(ji.prg.p.vars) if v not in ji.fixedvars]
+        if ji.prg.p.vars: self.var_vals_replace[j] = [(i, self.vars.index(v.expr)) for i, v in enumerate(ji.prg.p.vars) if v.expr not in ji.fixedvars]

        global_dim_idx, local_dim_idx = find_symbolic_dim(ji.prg.p.global_size), find_symbolic_dim(ji.prg.p.local_size)
        if global_dim_idx is not None or local_dim_idx is not None:
@@ -105,12 +105,12 @@ class GraphRunner(Runner):

    super().__init__(colored(f"<batched {len(jit_cache)}>", "cyan"), jit_cache[0].prg.device.split(":")[0], estimates.simplify())

-  def updated_vars(self, var_vals: dict[Variable, int]):
+  def updated_vars(self, var_vals: dict[str, int]):
    vals = [var_vals[v] for v in self.vars]
    for j, vidxs in self.var_vals_replace.items():
      for i, v in vidxs: yield j, i, vals[v]

-  def updated_launch_dims(self, var_vals: dict[Variable, int]):
+  def updated_launch_dims(self, var_vals: dict[str, int]):
    dims = [tuple(sym_infer(s, var_vals) for s in dim) for dim in self.symbolic_dims]
    for j, (gl, lc) in self.launch_dims_replace.items():
      yield j, (dims[gl] if gl is not None else self.launch_dims_base[j][0]), (dims[lc] if lc is not None else self.launch_dims_base[j][1])
@@ -192,7 +192,7 @@ class CapturedJit(Generic[ReturnType]):
    self.__post_init__()

  # jit exec
-  def __call__(self, input_buffers:list[Buffer], var_vals:dict[Variable, int]) -> ReturnType:
+  def __call__(self, input_buffers:list[Buffer], var_vals:dict[str, int]) -> ReturnType:
    # assign inputs
    for idx, offset, device, size, dtype in self.extra_view_inputs:
      input_buffers.append(Buffer(device, size, dtype, base=input_buffers[idx], offset=offset).ensure_allocated())
@@ -225,7 +225,8 @@ def _prepare_jit_inputs(args, kwargs):
                                         for lb in lbs if lb.base.realized is not None])
  assert len(set(input_buffers)) == len(input_buffers), "duplicate inputs to JIT"
  st_varval_dtype_device = [(*unwrap(lb.st).unbind(), lb.dtype, lb.device) for lb in lbs]
-  var_vals = merge_dicts([x[1] for x in st_varval_dtype_device] + [dict(v.unbind() for v in (args + tuple(kwargs.values())) if isinstance(v, UOp))])
+  _var_vals = merge_dicts([x[1] for x in st_varval_dtype_device] + [dict(v.unbind() for v in (args + tuple(kwargs.values())) if isinstance(v, UOp))])
+  var_vals = {k.expr:v for k,v in _var_vals.items()}
  st_vars_dtype_device = [(x[0], tuple(sorted(x[1].keys(), key=lambda v: v.expr)), x[2], x[3]) for x in st_varval_dtype_device]
  return input_buffers, var_vals, names, st_vars_dtype_device

@@ -3,12 +3,12 @@ import time, pprint, random, itertools, math
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, Variable, sym_infer, graph_rewrite, print_uops, track_rewrites, KernelInfo, pyrender
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, graph_rewrite, print_uops, track_rewrites, KernelInfo, pyrender
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
 from tinygrad.engine.schedule import ScheduleItem
 from tinygrad.codegen import full_rewrite
-from tinygrad.codegen.opt.kernel import Opt
+from tinygrad.codegen.opt import Opt

 # **************** Program Creation ****************

@@ -47,7 +47,8 @@ def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None)
  src = renderer.render(uops)

  return ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test", src, renderer.device, ast, uops,
-                     global_size=[1,1,1] if renderer.has_local else None, local_size=[1,1,1] if renderer.has_local else None)
+                     global_size=[1,1,1] if renderer.has_local or renderer.has_threads else None,
+                     local_size=[1,1,1] if renderer.has_local else None)

 # **************** Runners ****************

@@ -56,9 +57,9 @@ class Runner:
    self.first_run, self.display_name, self.device, self.estimates = True, display_name, device, estimates
  @property
  def dev(self): return Device[self.device]
-  def exec(self, rawbufs:list[Buffer], var_vals:dict[Variable, int]|None=None) -> float|None:
+  def exec(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None) -> float|None:
    return self(rawbufs, {} if var_vals is None else var_vals)
-  def __call__(self, rawbufs:list[Buffer], var_vals:dict[Variable, int], wait=False) -> float|None:
+  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False) -> float|None:
    raise NotImplementedError("override this")

 def optimize_local_size(_prg:Callable, global_size:list[int], rawbufs:list[Buffer]) -> list[int]:
@@ -89,9 +90,10 @@ class CompiledRunner(Runner):

  def __reduce__(self): return self.__class__, (self.p, self.lib)

-  def __call__(self, rawbufs:list[Buffer], var_vals:dict[Variable, int], wait=False) -> float|None:
+  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False) -> float|None:
+    has_local = Device[self.p.device].renderer.has_local
    global_size, local_size = self.p.launch_dims(var_vals)
-    if global_size is not None and local_size is None and all_int(self.p.global_size): # type: ignore[arg-type]
+    if has_local and global_size is not None and local_size is None and all_int(self.p.global_size): # type: ignore[arg-type]
      local_size = optimize_local_size(self._prg, global_size, rawbufs)
      global_size = [g//l if g%l == 0 else g/l for g,l in zip(global_size, local_size)]
      self.p = replace(self.p, global_size=global_size, local_size=local_size)
@@ -102,11 +104,11 @@ class CompiledRunner(Runner):
    if local_size:
      lra['local_size'] = tuple(local_size)
      assert len(local_size) == 3, "local size must have len 3"
-    return self._prg(*[x._buf for x in rawbufs], **lra, vals=tuple(var_vals[k] for k in self.p.vars), wait=wait)
+    return self._prg(*[x._buf for x in rawbufs], **lra, vals=tuple(var_vals[k.expr] for k in self.p.vars), wait=wait)

 class ViewOp(Runner):
  def __init__(self, buf:Buffer): super().__init__(colored(f"view {buf.nbytes:8d} @ {buf.offset:<10d}", "yellow"), buf.device)
-  def __call__(self, rawbufs:list[Buffer], var_vals:dict[Variable, int], wait=False):
+  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False):
    assert rawbufs[0]._base is not None and rawbufs[0]._base == rawbufs[1].base, f"must be base {rawbufs}"

 class BufferCopy(Runner):
@@ -124,7 +126,7 @@ class BufferCopy(Runner):
      src.allocator._copyout(dest.allocator._as_buffer(dest._buf), src._buf)
    else:
      dest.copyin(src.as_buffer(allow_zero_copy=True))  # may allocate a CPU buffer depending on allow_zero_copy
-  def __call__(self, rawbufs:list[Buffer], var_vals:dict[Variable, int], wait=False):
+  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False):
    dest, src = rawbufs[0:2]
    assert dest.size == src.size and dest.dtype == src.dtype, f"buffer copy mismatch, {dest.size} != {src.size}, {dest.dtype} != {src.dtype}"
    st = time.perf_counter()
@@ -159,8 +161,8 @@ class ExecItem:
  prg: Runner
  bufs: list[Buffer|None]
  metadata: tuple[Metadata, ...]|None = None
-  fixedvars: dict[Variable, int] = field(default_factory=dict)
-  def run(self, _var_vals:dict[Variable, int]|None=None, wait=False, jit=False, do_update_stats=True) -> float|None:
+  fixedvars: dict[str, int] = field(default_factory=dict)
+  def run(self, _var_vals:dict[str, int]|None=None, wait=False, jit=False, do_update_stats=True) -> float|None:
    var_vals = self.fixedvars if _var_vals is None else (_var_vals|self.fixedvars)
    bufs = [cast(Buffer, x) for x in self.bufs] if jit else [cast(Buffer, x).ensure_allocated() for x in self.bufs]
    if PROFILE: cpu_events.append(ProfilePointEvent(self.prg.device, "exec", self.prg.display_name, {"metadata":self.metadata, "var_vals":var_vals}))
@@ -211,7 +213,7 @@ def lower_schedule(schedule:list[ScheduleItem]) -> Generator[tuple[ScheduleItem,

 capturing: list = []  # put classes with an add method in here

-def run_schedule(schedule:list[ScheduleItem], var_vals:dict[Variable, int]|None=None, do_update_stats=True):
+def run_schedule(schedule:list[ScheduleItem], var_vals:dict[str, int]|None=None, do_update_stats=True):
  for si, ei in lower_schedule(schedule):
    if len(capturing) and CAPTURING: capturing[0].add(ei)
    if VALIDATE_WITH_CPU and si.ast.op is Ops.SINK:
@@ -229,4 +231,3 @@ def run_schedule(schedule:list[ScheduleItem], var_vals:dict[Variable, int]|None=
      np.testing.assert_allclose(si.bufs[0].numpy(), nb[0].numpy(), rtol=1e-3, atol=1e-3)
    else:
      ei.run(var_vals, do_update_stats=do_update_stats)
-
@@ -1,7 +1,7 @@
 from typing import cast
 from dataclasses import dataclass, field
 from collections import deque, defaultdict
-from tinygrad.uop.ops import UOp, Variable, Ops, buffers
+from tinygrad.uop.ops import UOp, Ops, buffers
 from tinygrad.device import Device, Buffer, MultiBuffer
 from tinygrad.helpers import Metadata, all_same

@@ -12,15 +12,15 @@ class ScheduleItem:
  ast: UOp
  bufs: tuple[Buffer, ...]
  metadata: tuple[Metadata, ...] = ()
-  fixedvars: dict[Variable, int] = field(default_factory=dict)
+  fixedvars: dict[str, int] = field(default_factory=dict)

 # **** schedule linearizer

-def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[Variable, int]]:
+def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[str, int]]:
  # construct the KERNEL children graph based on assigns
  children: defaultdict[UOp, list[UOp]] = defaultdict(list)
  in_degree: dict[UOp, int] = {}
-  var_vals: dict[Variable, int] = {}
+  var_vals: dict[str, int] = {}
  for u in sched_sink.toposort():
    if u.op is not Ops.ASSIGN: continue  # anything that's not an ASSIGN doesn't write a kernel, so we can skip
    k = u.src[1]
@@ -40,8 +40,8 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
        pass  # a BUFFER is already realized, nothing to do here
      elif s.op is Ops.BIND:
        var, val = s.unbind()
-        assert var not in var_vals or var_vals[var] == val, f"bind mismatch on {var}, {var_vals[var]} != {val}"
-        var_vals[var] = val
+        assert var.expr not in var_vals or var_vals[var.expr] == val, f"bind mismatch on {var}, {var_vals[var.expr]} != {val}"
+        var_vals[var.expr] = val
      else:
        raise RuntimeError(f"input to kernel must be ASSIGN or BUFFER, not {s.op}")

@@ -72,7 +72,7 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
      assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
      dnums = [x for x in ast.variables() if x.arg[0] == '_device_num']
      for i,bufs in enumerate(zip(*[x.bufs for x in cast(tuple[MultiBuffer, ...], ubufs)])):
-        schedule.append(ScheduleItem(ast, bufs, k.arg.metadata, {dnums[0]:i} if len(dnums) else {}))
+        schedule.append(ScheduleItem(ast, bufs, k.arg.metadata, {dnums[0].expr:i} if len(dnums) else {}))
    else:
      # ONE -> ONE
      schedule.append(ScheduleItem(ast, cast(tuple[Buffer, ...], ubufs), k.arg.metadata))
@@ -10,11 +10,12 @@ U = TypeVar("U")
 def prod(x:Iterable[T]) -> T|int: return functools.reduce(operator.mul, x, 1)

 # NOTE: helpers is not allowed to import from anything else in tinygrad
-OSX = platform.system() == "Darwin"
+OSX, WIN = platform.system() == "Darwin", sys.platform == "win32"
 CI = os.getenv("CI", "") != ""
+ARCH_X86 = any(x in platform.processor() for x in ("Intel", "i386", "x86_64"))

 # fix colors on Windows, https://stackoverflow.com/questions/12492810/python-how-can-i-make-the-ansi-escape-codes-to-work-also-in-windows
-if sys.platform == "win32": os.system("")
+if WIN: os.system("")

 def dedup(x:Iterable[T]): return list(dict.fromkeys(x))   # retains list order
 def argfix(*x):
@@ -126,8 +127,7 @@ class ContextVar:
  def __lt__(self, x): return self.value < x

 DEBUG, IMAGE, BEAM, NOOPT = ContextVar("DEBUG", 0), ContextVar("IMAGE", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
-JIT = ContextVar("JIT", 2 if platform.system() == 'Darwin' and ('Intel' in platform.processor() or 'i386' in platform.processor()) else 1)
-JIT_BATCH_SIZE = ContextVar("JIT_BATCH_SIZE", 32)
+JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32)
 WINO, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1)
 USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0)
 TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS", 0)
@@ -142,6 +142,7 @@ CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), Co
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, AMD_LLVM = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0), ContextVar("AMD_LLVM", 1)
 RANGEIFY, FUSE_ATTENTION = ContextVar("RANGEIFY", 0), ContextVar("FUSE_ATTENTION", 0)
 EMULATE = ContextVar("EMULATE", "")
+CPU_COUNT = ContextVar("CPU_COUNT", max(1, (os.cpu_count() or 1) // (4 if ARCH_X86 else 2))) # take 1/2 of the cores, accounting HT

@dataclass(frozen=True)
 class Metadata:
@@ -101,7 +101,7 @@ class ProgramSpec:
  def applied_opts(self) -> tuple[Opt, ...]|None: return self.uops[-1].arg.applied_opts if \
    self.uops is not None and self.uops[-1].op is Ops.SINK and self.uops[-1].arg is not None else None

-  def launch_dims(self, var_vals:dict[Variable, int]):
+  def launch_dims(self, var_vals:dict[str, int]):
    global_size = [sym_infer(sz, var_vals) for sz in self.global_size] if self.global_size is not None else None
    local_size = [sym_infer(sz, var_vals) for sz in self.local_size] if self.local_size is not None else None
    return global_size, local_size
@@ -112,6 +112,7 @@ class Renderer:
  # TODO: make this generic with a list of supported types
  supports_float4: bool = True
  has_local: bool = True
+  has_threads: bool = False
  has_shared: bool = True
  # NOTE: these two should be in (x,y,z) order to match the max_sizes argument in get_grouped_dims
  global_max: tuple[int, ...]|None = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now
@@ -3,7 +3,7 @@ import os, math, sys
 from collections import defaultdict, Counter
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, sint_to_uop
-from tinygrad.helpers import strip_parens, getenv, prod, dedup, AMX
+from tinygrad.helpers import strip_parens, getenv, prod, dedup, AMX, CPU_COUNT
 from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate
 from tinygrad.renderer import Renderer
 from tinygrad.codegen.late.devectorizer import no_vectorized_alu
@@ -192,9 +192,12 @@ class ClangRenderer(CStyleLanguage):
  float4_style = ('{', '}')
  gep_arr_threshold = 0
  has_local = False
-  global_max = None
+  has_threads = True
+  global_max = (CPU_COUNT.value, 0, 0)
  infinity = "__builtin_inff()"
  nan = '__builtin_nanf("")'
+  code_for_workitem = {"g": lambda _: "core_id"}
+  extra_args = ['int core_id']
  if AMX: tensor_cores = tc.amx

  # language options
@@ -4,12 +4,11 @@ import tinygrad.runtime.autogen.cuda as cuda
 from tinygrad.helpers import init_c_var, dedup
 from tinygrad.device import Buffer, Device
 from tinygrad.runtime.ops_cuda import CUDADevice, check, encode_args, cu_time_execution
-from tinygrad.uop.ops import Variable
 from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner
 from tinygrad.engine.jit import MultiGraphRunner, GraphException

 class CUDAGraph(MultiGraphRunner):
-  def __init__(self, jit_cache: list[ExecItem], input_rawbuffers: list[Buffer], var_vals: dict[Variable, int]):
+  def __init__(self, jit_cache: list[ExecItem], input_rawbuffers: list[Buffer], var_vals: dict[str, int]):
    super().__init__(jit_cache, input_rawbuffers, var_vals)

    # Check all jit items are compatible.
@@ -28,7 +27,7 @@ class CUDAGraph(MultiGraphRunner):
        deps = self._access_resources([x.base for x in ji.bufs if x is not None], ji.prg.p.outs, new_dependency=new_node)
        c_deps = (cuda.CUgraphNode*len(deps))(*deps) if deps else None

-        c_args, vargs = encode_args([cast(Buffer, x)._buf for x in ji.bufs], [var_vals.get(x, ji.fixedvars.get(x)) for x in ji.prg.p.vars])
+        c_args, vargs = encode_args([cast(Buffer, x)._buf for x in ji.bufs], [var_vals.get(x.expr, ji.fixedvars.get(x.expr)) for x in ji.prg.p.vars])
        kern_params = cuda.CUDA_KERNEL_NODE_PARAMS_v1(ji.prg._prg.prg, *global_size, *local_size, 0, None, vargs)
        check(cuda.cuGraphAddKernelNode(ctypes.byref(new_node), self.graph, c_deps, len(deps), ctypes.byref(kern_params)))

@@ -48,7 +47,7 @@ class CUDAGraph(MultiGraphRunner):

    self.instance = init_c_var(cuda.CUgraphExec(), lambda x: check(cuda.cuGraphInstantiate_v2(ctypes.byref(x), self.graph, None, None, 0)))

-  def __call__(self, input_rawbuffers: list[Buffer], var_vals: dict[Variable, int], wait=False) -> float|None:
+  def __call__(self, input_rawbuffers: list[Buffer], var_vals: dict[str, int], wait=False) -> float|None:
    # Update rawbuffers in the c_args struct.
    for (j,i),input_idx in self.input_replace.items():
      if not self.updatable_nodes[j][3]: setattr(self.updatable_nodes[j][2], f'f{i}', input_rawbuffers[input_idx]._buf)
@@ -9,7 +9,7 @@ from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner, Buffer
 from tinygrad.engine.jit import MultiGraphRunner

 class HCQGraph(MultiGraphRunner):
-  def __init__(self, jit_cache: list[ExecItem], input_rawbuffers: list[Buffer], var_vals: dict[Variable, int]):
+  def __init__(self, jit_cache: list[ExecItem], input_rawbuffers: list[Buffer], var_vals: dict[str, int]):
    super().__init__(jit_cache, input_rawbuffers, var_vals)
    self.devices = list(set(cast(HCQCompiled, d) for ji in jit_cache for d in [Device[cast(Buffer, x).device] for x in ji.bufs]))

@@ -69,7 +69,7 @@ class HCQGraph(MultiGraphRunner):
    for dev, queue in self.comp_queues.items(): dev_access[queue].add(dev)

    self.input_replace_map: dict[HCQCompiled, set[int]] = collections.defaultdict(set)
-    self.fixedvars: dict[HCQCompiled, dict[Variable, int]] = {}
+    self.fixedvars: dict[HCQCompiled, dict[str, int]] = {}

    for j,ji in enumerate(jit_cache):
      if is_exec_prg:=isinstance(ji.prg, CompiledRunner): enqueue_dev: HCQCompiled = ji.prg.dev
@@ -183,7 +183,7 @@ class HCQGraph(MultiGraphRunner):
    self.last_timeline: dict[HCQCompiled, tuple[HCQSignal, int]] = {dev: (dev.timeline_signal, 0) for dev in self.devices}
    self.queue_signals_to_reset = [self.signals[q] for q in list(self.comp_queues.values()) + list(self.copy_queues.values()) if q in self.signals]

-  def __call__(self, input_rawbuffers: list[Buffer], var_vals: dict[Variable, int], wait=False) -> float|None:
+  def __call__(self, input_rawbuffers: list[Buffer], var_vals: dict[str, int], wait=False) -> float|None:
    # Wait and restore signals
    self.kickoff_value += 1
    for dev in self.devices: self.last_timeline[dev][0].wait(self.last_timeline[dev][1])
@@ -195,12 +195,13 @@ class HCQGraph(MultiGraphRunner):

    if PROFILE and self.kickoff_value > 1: self.collect_timestamps()

-    hcq_var_vals = {self.kickoff_var: self.kickoff_value, **var_vals,
-                    **{var: dev.timeline_value - 1 for dev, var in self.virt_timeline_vals.items()},
-                    **{sig.base_buf.va_addr: dev.timeline_signal.base_buf.va_addr for dev, sig in self.virt_timeline_signals.items()}}
+    hcq_var_vals = {self.kickoff_var.expr: self.kickoff_value, **var_vals,
+                    **{var.expr: dev.timeline_value - 1 for dev, var in self.virt_timeline_vals.items()},
+                    **{sig.base_buf.va_addr.expr: dev.timeline_signal.base_buf.va_addr for dev, sig in self.virt_timeline_signals.items()}}

    # Update rawbuffers
-    for (j,i),input_idx in self.input_replace.items(): hcq_var_vals[self.input_replace_to_var.get((j,i))] = input_rawbuffers[input_idx]._buf.va_addr
+    for (j,i),input_idx in self.input_replace.items():
+      hcq_var_vals[self.input_replace_to_var[(j,i)].expr] = input_rawbuffers[input_idx]._buf.va_addr

    for dev in self.devices:
      self.comp_queues[dev].submit(dev, hcq_var_vals_local:=hcq_var_vals|self.fixedvars.get(dev, {}))
@@ -5,7 +5,6 @@ from tinygrad.helpers import dedup, getenv, merge_dicts, PROFILE
 from tinygrad.device import Buffer, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.engine.realize import ExecItem, CompiledRunner
 from tinygrad.engine.jit import GraphRunner, GraphException
-from tinygrad.uop.ops import Variable
 from tinygrad.runtime.ops_metal import wait_check, msg, libobjc, to_struct, objc_instance,\
  MTLResourceOptions, cmdbuf_st_time, cmdbuf_en_time, objc_id, to_ns_str

@@ -17,7 +16,7 @@ class MTLResourceUsage:
  MTLResourceUsageWrite = 0b10

 class MetalGraph(GraphRunner):
-  def __init__(self, jit_cache: list[ExecItem], input_rawbuffers: list[Buffer], var_vals: dict[Variable, int]):
+  def __init__(self, jit_cache: list[ExecItem], input_rawbuffers: list[Buffer], var_vals: dict[str, int]):
    super().__init__(jit_cache, input_rawbuffers, var_vals)
    if not all(isinstance(ji.prg, CompiledRunner) for ji in jit_cache): raise GraphException

@@ -48,7 +47,8 @@ class MetalGraph(GraphRunner):
        if b is not None and b not in input_rawbuffers:
          msg("setKernelBuffer:offset:atIndex:")(icb_command, b._buf.buf, b._buf.offset, i)
          all_resources.append(b._buf.buf)
-      for i,v in enumerate(prg.p.vars): msg("setKernelBuffer:offset:atIndex:")(icb_command, self.int_buf.buf, self.varlist.index(v)*4, len(ji.bufs)+i)
+      for i,v in enumerate(prg.p.vars):
+        msg("setKernelBuffer:offset:atIndex:")(icb_command, self.int_buf.buf, self.varlist.index(v.expr)*4, len(ji.bufs)+i)

      global_size, local_size = prg.p.launch_dims(var_vals)
      msg("concurrentDispatchThreadgroups:threadsPerThreadgroup:")(icb_command, to_struct(*global_size), to_struct(*local_size))
@@ -61,7 +61,7 @@ class MetalGraph(GraphRunner):
    for var in self.fixedvars: self.int_buf_view[self.varlist.index(var)] = self.fixedvars[var]
    self.range = to_struct(0, len(jit_cache))

-  def __call__(self, input_rawbuffers: list[Buffer], var_vals: dict[Variable, int], wait=False) -> float|None:
+  def __call__(self, input_rawbuffers: list[Buffer], var_vals: dict[str, int], wait=False) -> float|None:
    if self.command_buffer is not None and self.command_buffer in self.dev.mtl_buffers_in_flight: wait_check(self.command_buffer)
    # NOTE: old command buffer may not be inflight anymore
    if self.command_buffer is not None and PROFILE: self.collect_timestamps()
@@ -1,5 +1,4 @@
 import time, itertools
-from tinygrad.uop.ops import Variable
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.engine.realize import CompiledRunner, BufferXfer, ExecItem
 from tinygrad.device import Device, Compiled, Buffer
@@ -18,7 +17,7 @@ def dev_key(dev:RemoteDevice): return dev.conn if dev.properties.graph_supports_
 def map_rawbuf(rawbuf:Buffer): return (cast(RemoteDevice, Device[rawbuf.device]).session, rawbuf._buf)

 class RemoteGraph(MultiGraphRunner):
-  def __init__(self, jit_cache: list[ExecItem], rawbufs: list[Buffer], var_vals: dict[Variable, int]):
+  def __init__(self, jit_cache: list[ExecItem], rawbufs: list[Buffer], var_vals: dict[str, int]):
    super().__init__(jit_cache, rawbufs, var_vals)
    devices = dedup(flatten([[Device[unwrap(buf).device] for buf in ji.bufs] for ji in jit_cache]))
    c2d = {device.conn: device for device in devices}
@@ -93,7 +92,7 @@ class RemoteGraph(MultiGraphRunner):
    for req in self.template:
      match req:
        case GraphExec(): RemoteConnection(unwrap(req.session).host).q(GraphFree(req.graph_num, session=req.session))
-  def __call__(self, rawbufs: list[Buffer], var_vals: dict[Variable, int], wait=False):
+  def __call__(self, rawbufs: list[Buffer], var_vals: dict[str, int], wait=False):
    if wait: st = time.perf_counter()
    rmap = {orig: map_rawbuf(rawbufs[replace_idx]) for orig,replace_idx in self.handle_indexes.items()}
    for req in self.template:
@@ -1,6 +1,6 @@
 from __future__ import annotations
 import platform, subprocess, sys, ctypes, functools, time, mmap, threading, queue
-from tinygrad.helpers import capstone_flatdump, getenv, from_mv, to_mv, OSX, mv_address, wait_cond, cpu_profile
+from tinygrad.helpers import capstone_flatdump, getenv, from_mv, to_mv, OSX, WIN, mv_address, wait_cond, cpu_profile, suppress_finalizing
 from tinygrad.device import Compiler, BufferSpec, DMACPURef
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocatorBase, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
 from tinygrad.runtime.support.elf import jit_loader
@@ -30,29 +30,38 @@ class ClangJITCompiler(Compiler):
 class CPUWorker(threading.Thread):
  def __init__(self, dev, tasks, thread_id):
    super().__init__()
-    self.dev, self.tasks, self.thread_id, self.daemon = dev, tasks, thread_id, True
+    self.dev, self.tasks, self.thread_id, self.pool, self.daemon = dev, tasks, thread_id, [], True
+
+  def push_task(self, tid, cmd, args):
+    if len(self.pool) <= tid:
+      self.pool.append(queue.Queue())
+      CPUWorker(self, self.pool[tid], thread_id=tid+1).start()
+    self.pool[tid].put([cmd, 1, len(args)] + args)

  def run(self):
    while True:
      cmd_iter = iter(self.tasks.get())
      for cmd in cmd_iter:
-        args_cnt = next(cmd_iter)
-        cmd(*[next(cmd_iter) for _ in range(args_cnt)])
+        threads, args_cnt = next(cmd_iter), next(cmd_iter)
+        args = [next(cmd_iter) for _ in range(args_cnt)]
+        for th in range(threads - 1): self.push_task(th, cmd, args)
+        cmd(self.thread_id, *args)
+        for th in range(threads - 1): self.pool[th].join()
      self.tasks.task_done()

 class CPUComputeQueue(HWQueue):
-  def _exec(self, prg, bufs, *args):
-    prg.fxn(*map(ctypes.c_uint64, args[:bufs]), *map(ctypes.c_int64 if platform.machine() == "arm64" else ctypes.c_int32, args[bufs:]))
-  def _signal(self, signal_addr, value): to_mv(signal_addr, 4).cast('I')[0] = value
-  def _wait(self, signal_addr, value): wait_cond(lambda: to_mv(signal_addr, 4).cast('I')[0] >= value, timeout_ms=60000)
-  def _timestamp(self, timestamp_addr): to_mv(timestamp_addr, 8).cast('Q')[0] = time.perf_counter_ns()
-  def cmd(self, cmd, *args):
-    self.q(cmd, len(args), *args)
+  def _exec(self, tid, prg, bufs, *args):
+    prg.fxn(*map(ctypes.c_uint64, args[:bufs]), *map(ctypes.c_int64 if platform.machine() == "arm64" else ctypes.c_int32, args[bufs:]), tid)
+  def _signal(self, tid, signal_addr, value): to_mv(signal_addr, 4).cast('I')[0] = value
+  def _wait(self, tid, signal_addr, value): wait_cond(lambda: to_mv(signal_addr, 4).cast('I')[0] >= value, timeout_ms=60000)
+  def _timestamp(self, tid, timestamp_addr): to_mv(timestamp_addr, 8).cast('Q')[0] = time.perf_counter_ns()
+  def cmd(self, cmd, *args, threads=1):
+    self.q(cmd, threads, len(args), *args)
    return self

  def memory_barrier(self): return self
  def exec(self, prg:CPUProgram, args_state:HCQArgsState, global_size, local_size):
-    return self.cmd(self._exec, prg, len(args_state.bufs), *[x.va_addr for x in args_state.bufs], *args_state.vals)
+    return self.cmd(self._exec, prg, len(args_state.bufs), *[x.va_addr for x in args_state.bufs], *args_state.vals, threads=(global_size or (1,))[0])
  def wait(self, signal, value=0): return self.cmd(self._wait, signal.value_addr, value)
  def timestamp(self, signal): return self.cmd(self._timestamp, signal.timestamp_addr)
  def signal(self, signal, value:sint=0): return self.cmd(self._signal, signal.value_addr, value)
@@ -62,10 +71,10 @@ class CPUComputeQueue(HWQueue):
 MAP_JIT = 0x0800

 class CPUProgram(HCQProgram):
-  rt_lib = ctypes.CDLL(ctypes.util.find_library('System' if OSX else 'kernel32') if OSX or sys.platform == "win32" else 'libgcc_s.so.1')
+  rt_lib = ctypes.CDLL(ctypes.util.find_library('System' if OSX else 'kernel32') if OSX or WIN else 'libgcc_s.so.1')

  def __init__(self, dev, name:str, lib:bytes):
-    if sys.platform == "win32":
+    if sys.platform == "win32": # mypy doesn't understand when WIN is used here
      PAGE_EXECUTE_READWRITE, MEM_COMMIT, MEM_RESERVE = 0x40, 0x1000, 0x2000
      ctypes.windll.kernel32.VirtualAlloc.restype = ctypes.c_void_p
      self.mem = ctypes.windll.kernel32.VirtualAlloc(ctypes.c_void_p(0), ctypes.c_size_t(len(lib)), MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE)
@@ -93,14 +102,14 @@ class CPUProgram(HCQProgram):

    super().__init__(HCQArgsState, dev, name, kernargs_alloc_size=0)

+  @suppress_finalizing
  def __del__(self):
-    if getattr(sys, 'is_finalizing', lambda: True)(): return
    if sys.platform == 'win32': ctypes.windll.kernel32.VirtualFree(ctypes.c_void_p(self.mem), ctypes.c_size_t(0), 0x8000) #0x8000 - MEM_RELEASE

 class CPUAllocator(HCQAllocatorBase):
  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
    if options.external_ptr: addr, buf = options.external_ptr, None
-    elif sys.platform == "win32": addr = mv_address(buf:=mmap.mmap(-1, size, access=mmap.ACCESS_WRITE))
+    elif WIN: addr = mv_address(buf:=mmap.mmap(-1, size, access=mmap.ACCESS_WRITE))
    else: addr = mv_address(buf:=mmap.mmap(-1, size, mmap.MAP_ANON | mmap.MAP_PRIVATE, mmap.PROT_READ | mmap.PROT_WRITE))
    return HCQBuffer(va:=addr, sz:=size, meta=buf, view=MMIOInterface(va, sz, fmt='B'), owner=self.dev)
  def _as_buffer(self, src) -> memoryview:
@@ -36,8 +36,10 @@ dsp_string = PatternMatcher([
 class DSPRenderer(ClangRenderer):
  device = "DSP"
  supports_float4 = True
+  has_threads = False
  buffer_suffix = " restrict __attribute__((align_value(128)))"
  kernel_typedef = "__attribute__((noinline)) void"
+  extra_args = []
  pre_matcher = dsp_pm
  extra_matcher = dsp_pm_late+ClangRenderer.extra_matcher
  string_rewrite = dsp_string+ClangRenderer.string_rewrite
@@ -100,7 +100,7 @@ class GraphComputeItem:
  datahash: str
  bufs: tuple[int, ...]
  vars: tuple[Variable, ...]
-  fixedvars: dict[Variable, int]
+  fixedvars: dict[str, int]
  ins: tuple[int, ...]
  outs: tuple[int, ...]
  global_size: tuple[sint, ...]|None
@@ -111,7 +111,7 @@ class GraphAlloc(RemoteRequest):
  graph_num: int
  jit_cache: tuple[GraphComputeItem|Transfer, ...]
  bufs: tuple[tuple[SessionKey, int], ...]
-  var_vals: dict[Variable, int]
+  var_vals: dict[str, int]

@dataclass(frozen=True)
 class GraphFree(RemoteRequest):
@@ -121,7 +121,7 @@ class GraphFree(RemoteRequest):
 class GraphExec(RemoteRequest):
  graph_num: int
  bufs: tuple[tuple[SessionKey, int], ...]
-  var_vals: dict[Variable, int]
+  var_vals: dict[str, int]
  wait: bool

 # for safe deserialization
@@ -6,7 +6,7 @@ except ImportError: fcntl = None #type:ignore[assignment]
 from tinygrad.helpers import PROFILE, getenv, to_mv, round_up, ProfileRangeEvent
 from tinygrad.renderer import Renderer
 from tinygrad.device import BufferSpec, Compiler, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent
-from tinygrad.uop.ops import sym_infer, sint, Variable, UOp
+from tinygrad.uop.ops import sym_infer, sint, UOp
 from tinygrad.runtime.autogen import libc

 class MMIOInterface:
@@ -28,10 +28,7 @@ class FileIOInterface:
  def __del__(self):
    if hasattr(self, 'fd'): os.close(self.fd)
  def ioctl(self, request, arg): return fcntl.ioctl(self.fd, request, arg)
-  def mmap(self, start, sz, prot, flags, offset):
-    x = libc.mmap(start, sz, prot, flags, self.fd, offset)
-    if x == 0xffffffffffffffff: raise OSError(f"Failed to mmap {sz} bytes at {hex(start)}: {os.strerror(ctypes.get_errno())}")
-    return x
+  def mmap(self, start, sz, prot, flags, offset): return FileIOInterface._mmap(start, sz, prot, flags, self.fd, offset)
  def read(self, size=None, binary=False, offset=None):
    if offset is not None: self.seek(offset)
    with open(self.fd, "rb" if binary else "r", closefd=False) as file: return file.read(size)
@@ -41,11 +38,13 @@ class FileIOInterface:
  def listdir(self): return os.listdir(self.path)
  def seek(self, offset): os.lseek(self.fd, offset, os.SEEK_SET)
  @staticmethod
-  def anon_mmap(start, sz, prot, flags, offset):
-    x = libc.mmap(start, sz, prot, flags, -1, offset)
+  def _mmap(start, sz, prot, flags, fd, offset):
+    x = libc.mmap(start, sz, prot, flags, fd, offset)
    if x == 0xffffffffffffffff: raise OSError(f"Failed to mmap {sz} bytes at {hex(start)}: {os.strerror(ctypes.get_errno())}")
    return x
  @staticmethod
+  def anon_mmap(start, sz, prot, flags, offset): return FileIOInterface._mmap(start, sz, prot, flags, -1, offset)
+  @staticmethod
  def munmap(buf, sz): return libc.munmap(buf, sz)
  @staticmethod
  def exists(path): return os.path.exists(path)
@@ -192,7 +191,7 @@ class HWQueue(Generic[SignalType, HCQDeviceType, ProgramType, ArgsStateType]):
      if isinstance(val, int): mv[i] = val if mask is None else ((mv[i] & ~mask) | val)
      else: self.mv_sints.append((mv, i, self._new_sym(val), mask))

-  def _apply_var_vals(self, var_vals:dict[Variable, int]):
+  def _apply_var_vals(self, var_vals:dict[str, int]):
    resolved_syms = [sym_infer(sym, var_vals) for sym in self.syms]

    for off, sym_idx in self.q_sints:
@@ -205,7 +204,7 @@ class HWQueue(Generic[SignalType, HCQDeviceType, ProgramType, ArgsStateType]):

    self._prev_resolved_syms = cast(list[int|None], resolved_syms)

-  def submit(self, dev:HCQDeviceType, var_vals:dict[Variable, int]|None=None):
+  def submit(self, dev:HCQDeviceType, var_vals:dict[str, int]|None=None):
    """
    Submits the command queue to a specific device for execution.

@@ -119,7 +119,7 @@ def map_reshape(idx:UOp, r:UOp):
  for s,src in list(zip(idx.shape, idx.src[1:]))[::-1]:
    to_sum.append(acc*src)
    acc *= s
-  mish = sum(to_sum, start=UOp.const(dtypes.int, 0))
+  mish = sum(to_sum, start=UOp.const(dtypes.index, 0))
  ret:list[UOp] = []
  for s in r.src[0].shape[::-1]:
    ret.append(mish % s) # NOTE: simplify will turn this to CONST
@@ -186,7 +186,7 @@ def map_partial_contiguous(ctx:RangeifyContext, x:UOp, idx:UOp):
      ranges.append(idx.src[1+i])
      continue
    passthrough_idx.append(idx.src[1+i])
-    ranges.append(ctx.new_range(s) if resolve(s!=1) else UOp.const(dtypes.int, 0))
+    ranges.append(ctx.new_range(s) if resolve(s!=1) else UOp.const(dtypes.index, 0))
    new_ranges.append(ranges[-1])
  ret = x.src[0].index(*ranges).bufferize(*[x for x in new_ranges if x.op is not Ops.CONST], arg=x.device)
  return ret.index(*passthrough_idx)
@@ -195,7 +195,7 @@ def map_contiguous(ctx:RangeifyContext, x:UOp):
  if x.arg is not None: return None
  ranges = []
  for s in x.shape[len(x.src)-1:]:
-    ranges.append(ctx.new_range(s) if resolve(s!=1) else UOp.const(dtypes.int, 0))
+    ranges.append(ctx.new_range(s) if resolve(s!=1) else UOp.const(dtypes.index, 0))
  ret = x.src[0].index(*ranges).bufferize(*x.src[1:], *[x for x in ranges if x.op is not Ops.CONST], arg=x.device)
  return ret.shrink(((0, prod(x.shape)),)).forced_reshape(x.shape)

@@ -5,22 +5,8 @@ import functools
 from typing import Callable
 from tinygrad.helpers import merge_dicts, getenv
 from tinygrad.shape.view import View, unravel
-from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import UOp, Ops, graph_rewrite, Variable, sint, sint_to_uop, Context, PatternMatcher, UPat, GroupOp
 from tinygrad.uop.symbolic import symbolic_flat, uop_given_valid, simplify_valid
-
-# If a node overflow, its srcs need to be checked to see if this overflow is the result of an ALU operation,
-# or that the node simply inherits the dtype from srcs. Upcast is either `Ops.CAST`+`replace` or just `replace`.
-def handle_upcast(u: UOp) -> UOp|None:
-  dtype = dtypes.int64.vec(u.dtype.count) if u.dtype.count > 1 else dtypes.int64
-  # check for overflow, upcast this to int64
-  if u.vmax > dtypes.max(dtypes.int) or u.vmin < dtypes.min(dtypes.int):
-    return u.replace(dtype=dtype, src=tuple([x.cast(dtype) for x in u.src]))
-  # if any inputs are int64 and this *doesn't* overflow, cast back to int
-  if any(x.dtype == dtypes.int64 for x in u.src):
-    return u.replace(dtype=dtype, src=tuple([x.cast(dtype) for x in u.src])).cast(u.dtype)
-  return None
-pm_upcast = PatternMatcher([(UPat(GroupOp.ALU, dtype=dtypes.int, name="u"), handle_upcast),])
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite, Variable, sint, sint_to_uop, Context

@functools.cache
 def views_to_indexed_uops(views: tuple[View, ...], _idxs:tuple[UOp, ...]|None=None) -> tuple[UOp, UOp]:
@@ -34,8 +20,8 @@ def views_to_indexed_uops(views: tuple[View, ...], _idxs:tuple[UOp, ...]|None=No
    # simplify
    if (newvalid:=simplify_valid(valid)) is not None: valid = newvalid
    if (newidx:=uop_given_valid(valid, idx)) is not None: idx = newidx
-    # symbolic again, upcast if needed
-    return graph_rewrite(UOp.sink(idx, valid), symbolic_flat+pm_upcast, name="indexing sym @ 2").src
+    # symbolic again
+    return graph_rewrite(UOp.sink(idx, valid), symbolic_flat, name="indexing sym @ 2").src

@functools.cache
 def views_to_real_strides(views: tuple[View, ...], ignore_valid=False) -> tuple[sint|None, ...]:
@@ -97,7 +83,7 @@ class ShapeTracker:
  def vars(self) -> set[Variable]: return set().union(*[v.vars() for v in self.views])

  @property
-  def var_vals(self) -> dict[Variable, int]: return merge_dicts([dict([v.unbind()]) for v in self.vars()])
+  def var_vals(self) -> dict[str, int]: return merge_dicts([{(vu:=v.unbind())[0].expr:vu[1]} for v in self.vars()])

  def unbind(self) -> tuple[ShapeTracker, dict[Variable, int]]:
    unbound_views, var_vals = zip(*[v.unbind() for v in self.views])
@@ -204,15 +204,15 @@ class View:

    # Merge dimensions in vm2 if required.
    # NB: Merging too many dimensions can make it difficult to project vm2's mask, hence only combining when required.
-    idxs: list[UOp] = [UOp.variable(f"idx{i}", 0, s-1) for i,s in enumerate(vm1.shape)]
-    merged_size, merged_term = 1, UOp.const(dtypes.int, 0)
+    idxs: list[UOp] = [UOp.variable(f"idx{i}", 0, s-1, dtypes.index) for i,s in enumerate(vm1.shape)]
+    merged_size, merged_term = 1, UOp.const(dtypes.index, 0)
    extents: list[tuple[sint, UOp]] = []
    for term, s, o in zip(reversed(terms), reversed(vm2.shape), reversed(origin)):
      merged_term += (sum([idxs[d1] * s1 for d1, s1 in term]) + o) * merged_size
      merged_size *= s
      if resolve(merged_term < merged_size, False) and resolve(0 <= merged_term, False):
        extents.append((merged_size, merged_term))
-        merged_size, merged_term = 1, UOp.const(dtypes.int, 0)
+        merged_size, merged_term = 1, UOp.const(dtypes.index, 0)
    if resolve(merged_term != 0): return None
    if (vm2_shape := tuple(s for s,_ in reversed(extents))) != vm2.shape:
      if (reshaped_vm2 := vm2.reshape(vm2_shape)) is None: return None
@@ -8,7 +8,7 @@ from tinygrad.dtype import _from_np_dtype, _to_np_dtype
 from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten, dedup
 from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG, is_numpy_ndarray, RANGEIFY, FUSE_ATTENTION
 from tinygrad.gradient import compute_gradient
-from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, Variable, MathTrait, identity_element, all_metadata
+from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, MathTrait, identity_element, all_metadata, index_to_concrete_int
 from tinygrad.uop.spec import tensor_uop_spec, type_verify
 from tinygrad.device import Device, Buffer
 from tinygrad.engine.realize import run_schedule
@@ -139,11 +139,13 @@ class Tensor(MathTrait):
    # create a UOp from the different types of inputs
    if isinstance(data, UOp):
      assert dtype is None or dtype==data.dtype, "dtype doesn't match, and casting isn't supported"
-      if data.op is Ops.BIND:
-        var, val = data.unbind()
+      # if data is dtype.index that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of
+      if data.dtype==dtypes.index: data = index_to_concrete_int(data)
+      if data.op is Ops.BIND:  # type: ignore  # mypy type narrowing is bugged here
+        var, val = data.unbind()  # type: ignore
        # give the bound constant a device
        const = UOp.const(var.dtype, val, device, ())
-        data = data.replace(src=(var.replace(src=const.src), const))
+        data = data.replace(src=(var.replace(src=const.src), const))  # type: ignore
    elif data is None: data = UOp.const(dtype or dtypes.default_float, 0, device, ())
    elif isinstance(data, get_args(ConstType)): data = UOp.const(dtype or dtypes.from_py(data), data, device, ())
    elif isinstance(data, bytes): data = _frompy(data, dtypes.uint8 if dtype is None else dtype)
@@ -241,7 +243,7 @@ class Tensor(MathTrait):
    _apply_map_to_tensors(becomes_map, name="Apply Kernelize Map")
    return self

-  def schedule_with_vars(self, *lst:Tensor) -> tuple[list[ScheduleItem], dict[Variable, int]]:
+  def schedule_with_vars(self, *lst:Tensor) -> tuple[list[ScheduleItem], dict[str, int]]:
    """
    Creates the schedule needed to realize these Tensor(s), with Variables.

@@ -5,7 +5,7 @@ from dataclasses import dataclass, field
 from enum import Enum, auto
 from tinygrad.uop import Ops, GroupOp
 from tinygrad.uop.mathtraits import MathTrait
-from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType
+from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
 from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey
 if TYPE_CHECKING:
@@ -14,6 +14,7 @@ if TYPE_CHECKING:

 class AxisType(Enum):
  GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
+  THREAD = auto()

 # https://en.wikipedia.org/wiki/Identity_element
 def identity_element(op:Ops, dt:DType) -> ConstType: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dtypes.min(dt)}[op], dt)
@@ -37,7 +38,7 @@ def smin(*lst): return _suop(argfix(*lst), UOp.minimum, min)
 def srender(x) -> str: return x.render() if isinstance(x, UOp) else str(x)

 def ssimplify(uop): return uop.ssimplify() if isinstance(uop, UOp) else uop
-def sym_infer(uop: UOp|int, var_vals: dict[UOp, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop
+def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop

 # used for UOp and UPat
 def pretty_print(x:Any, rep:Callable, srcfn=lambda x: x.src, cache=None, d=0)->str:
@@ -307,7 +308,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
  def range(end:sint, *arg):
    if len(arg) == 0: raise RuntimeError("range needs an arg")
    if len(arg) == 1: arg = arg+(AxisType.LOOP,)
-    return UOp(Ops.RANGE, dtype=dtypes.int, src=(sint_to_uop(end),), arg=arg)
+    return UOp(Ops.RANGE, dtype=dtypes.index, src=(sint_to_uop(end),), arg=arg)
  def r(self, op:Ops, axis:tuple[int, ...]):
    axis = tuple(sorted([x for x in axis if resolve(self.shape[x] != 1)]))
    if len(axis) == 0: return self
@@ -482,11 +483,11 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
  # *** uop Variable stuff ***

  @staticmethod
-  def variable(name:str, min_val:ConstType, max_val:ConstType, dtype:DType=dtypes.int) -> UOp:
+  def variable(name:str, min_val:ConstType, max_val:ConstType, dtype:DType=dtypes.index) -> UOp:
    assert not isinstance(min_val, UOp) and not isinstance(max_val, UOp), f"can't create Variable {name} with {min_val}/{max_val}"
    return UOp(Ops.DEFINE_VAR, dtype, arg=(name, min_val, max_val))
  @property
-  def expr(self):
+  def expr(self) -> str:
    assert self.op is Ops.DEFINE_VAR, f"op is {self.op}, need DEFINE_VAR"
    return self.arg[0]
  def bind(self, val:int|UOp):
@@ -573,7 +574,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    if self.op is Ops.VCONST: return (min(self.arg), max(self.arg))
    if self.op is Ops.GEP: return self.src[0]._min_max
    # TODO: CAST to bool/unsigned is not monotone, still some case can be simplified
-    if self.op is Ops.CAST and self.dtype in (dtypes.floats+dtypes.sints):
+    if self.op is Ops.CAST and self.dtype in dtypes.floats+dtypes.sints+(dtypes.index,):
      return max(dtypes.min(self.dtype), self.src[0].vmin), min(self.src[0].vmax, dtypes.max(self.dtype))
    return dtypes.min(self.dtype), dtypes.max(self.dtype)

@@ -584,9 +585,9 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    # TODO: sanitize varnames, or don't use naked eval while staying fast
    return eval("lambda "+','.join(varnames)+": "+sself.render(pm=renderer_infer)), varnames  # pylint: disable=eval-used

-  def sym_infer(self, var_vals:dict[UOp, int]):
+  def sym_infer(self, var_vals:dict[str, int]):
    fxn, varnames = self._sym_fxn
-    return fxn(**{k.arg[0]:v for k,v in var_vals.items() if k.arg[0] in varnames})
+    return fxn(**{k:v for k,v in var_vals.items() if k in varnames})

  def render(self, simplify=True, pm:PatternMatcher|None=None) -> str:
    with Context(TRACK_MATCH_STATS=0):
@@ -1025,7 +1026,30 @@ def graph_rewrite_map(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, na
    for k,v in input_map.items(): new_map[k] = new_map.get(v,v)
  return new_map

-def sint_to_uop(x:sint, dtype:DType=dtypes.int) -> UOp: return UOp.const(dtype, x) if isinstance(x, int) else x
+def sint_to_uop(x:sint) -> UOp: return UOp.const(dtypes.index, x) if isinstance(x, int) else x.cast(dtypes.index)
+
+def select_dtype(u): return (dtypes.long if u.overflows(dtypes.int32) else dtypes.int).vec(u.dtype.count)
+pm_lower_index_dtype = PatternMatcher([
+  # There are no Unary ops at this point in symbolic, those are introduced later
+  (UPat(GroupOp.Binary, dtypes.index, name="u", src=(UPat.var("x"), UPat.var("y"))), lambda u,x,y:
+    x.cast(dt:=least_upper_dtype(select_dtype(u), x.dtype, y.dtype)).alu(u.op, y.cast(dt))),
+  # comparison ops might now have different dtypes in their sources
+  (UPat(GroupOp.Comparison, name="u", src=(UPat.var("x",dtypes.ints), UPat.var("y", dtypes.ints))), lambda u,x,y:
+    x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)).alu(u.op, y.cast(dt)) if x.dtype!=y.dtype else None),
+  (UPat(Ops.WHERE, dtype=dtypes.index, src=(UPat.var("cond"), UPat.var("x"), UPat.var("y")), name="u"), lambda cond,u,x,y:
+    cond.where(x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)), y.cast(dt))),
+  (UPat((Ops.CONST, Ops.VCONST), dtype=dtypes.index, name="u"), lambda u: u.replace(dtype=select_dtype(u))),
+  (UPat((Ops.RANGE,), dtype=dtypes.index, src=(UPat.var("end")), name="r"), lambda ctx,r,end:
+    r.replace(dtype=(dt:=select_dtype(r)), src=(end.cast(dt),))),
+  (UPat(Ops.CAST, dtype=dtypes.index, src=(UPat.var("x", dtypes.ints),), name="u"), lambda u,x: x),
+  (UPat(Ops.VECTORIZE, dtype=dtypes.index, name="u"), lambda u: u.replace(
+    dtype=(dt:=least_upper_dtype(*[x.dtype for x in u.src])).vec(u.dtype.count), src=tuple(x.cast(dt) for x in u.src))),
+  (UPat(Ops.VECTORIZE, dtype=dtypes.index, name="u"), lambda u: u.replace(dtype=(dt:=(dtypes.long if any(v.overflows(dtypes.int) for v in u.src)
+    else dtypes.long)).vec(u.dtype.count),src=tuple(x.cast(dt) for x in u.src))),
+  (UPat((Ops.SPECIAL,Ops.DEFINE_VAR), dtypes.index, name="u"), lambda u: u.replace(dtype=dtypes.int)),
+  (UPat((Ops.BIND), dtypes.index, name="u"), lambda u: u.replace(dtype=u.src[0].dtype)),
+])
+def index_to_concrete_int(u:UOp): return graph_rewrite(u, pm_lower_index_dtype)

 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])

@@ -28,17 +28,17 @@ try:
    # float loads only become a variable when they get cast to int/bool
    (UPat(Ops.LOAD, dtypes.ints, name="x"),
      lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],create_bounded(f"load{ctx[1].setdefault(x, len(ctx[1]))}", x.dtype.min, x.dtype.max, ctx[0])))),
-    (UPat(Ops.CONST, dtype=dtypes.ints+(dtypes.bool,), name="x"),
+    (UPat(Ops.CONST, dtype=dtypes.ints+(dtypes.bool,dtypes.index), name="x"),
      lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],(z3.BoolVal if dtypes.is_bool(x.dtype) else z3.IntVal)(x.arg, ctx=ctx[0].ctx)))),
    # z3 can cast from bool to int automatically
-    (UPat(Ops.CAST, dtype=dtypes.ints, src=UPat(Ops.NOOP), name="x"), lambda x: x.src[0]),
+    (UPat(Ops.CAST, dtype=dtypes.ints+(dtypes.index,), src=UPat(Ops.NOOP), name="x"), lambda x: x.src[0]),
    (UPat(Ops.CAST, dtype=dtypes.bool, src=UPat(Ops.NOOP), name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0], x.src[0].arg[1]!=0))),
    # if the source of the cast is not a noop it means that it is a float and so we create a new variable
-    (UPat(Ops.CAST, dtype=dtypes.ints, name="x"), lambda x,ctx:
+    (UPat(Ops.CAST, dtype=dtypes.ints+(dtypes.index,), name="x"), lambda x,ctx:
      UOp(Ops.NOOP, arg=(ctx[0], create_bounded(f"cast{ctx[1].setdefault(x, len(ctx[1]))}", x.dtype.min, x.dtype.max, ctx[0])))),
    (UPat(Ops.CAST, dtype=dtypes.bool, name="x"), lambda x,ctx:
      UOp(Ops.NOOP, arg=(ctx[0], z3.Bool(f"cast{ctx[1].setdefault(x, len(ctx[1]))}",ctx=ctx[0].ctx)))),
-    (UPat(Ops.XOR, src=UPat(Ops.NOOP), name="x"),
+    (UPat(Ops.XOR, dtype=dtypes.ints+(dtypes.bool, ), src=UPat(Ops.NOOP), name="x"),
      lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0], z3.BV2Int(z3_alu[x.op](*(z3.Int2BV(s.arg[1], x.dtype.itemsize*8) for s in x.src)))))),
    (UPat(GroupOp.ALU, src=UPat(Ops.NOOP), name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0], z3_alu[x.op](*(s.arg[1] for s in x.src))))),
    # A comparison between floats introduces a new bool variable
@@ -95,7 +95,7 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([
  (UPat(Ops.VIEW, src=(UPat.var("x"),)), lambda x: x.base.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.ASSIGN, Ops.CONST, Ops.DEVICE}),

  # Tensor variable bindings
-  (UPat(Ops.BIND, dtypes.int, (UPat(Ops.DEFINE_VAR), UPat.cvar(dtype=dtypes.int)), arg=None), lambda: True),
+  (UPat(Ops.BIND, (dtypes.int,dtypes.index,), (UPat(Ops.DEFINE_VAR), UPat.cvar(dtype=(dtypes.int,dtypes.index,))), arg=None), lambda: True),

  # Tensor const has a device and an unmasked ShapeTracker of stride 0
  # NOTE: variables in shape can cause multiple views in this ShapeTracker and other issues, see TestSymbolicJit.test_ones_sum
@@ -120,6 +120,7 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([
 # ***** uop type spec *****

 def validate_index(idx:UOp, gate:UOp=UOp.const(dtypes.bool, True)):
+  # TODO: check for overflow
  if IGNORE_OOB or isinstance(idx.dtype, ImageDType) or (sz := idx.src[0].ptrdtype.size) == -1: return True
  # We can use UOp min/max to do a faster check, but it can give false positive since its not an exact bound and doesn't consider the mask
  if 0<=idx.src[1].vmin and idx.src[1].vmax<sz: return True
@@ -175,6 +176,9 @@ spec = PatternMatcher([

  # **** new style load/store ****

+  # make sure all index dtypes have been lowered
+  (UPat(GroupOp.All, dtype=dtypes.index), lambda: False),
+
  # INDEX is used in new style load/store
  # INDEX takes a <buf, alu, gate?>
  (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines), UPat())), lambda: True),
@@ -26,7 +26,7 @@ symbolic_simple = PatternMatcher([
  # ** self folding **
  (UPat.var("x") + 0, lambda x: x),    # x+0 -> x
  (UPat.var("x") * 1, lambda x: x),    # x*1 -> x
-  (UPat.var("x", dtype=dtypes.ints) ^ 0, lambda x: x), # x^0 -> x
+  (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.index)) ^ 0, lambda x: x), # x^0 -> x
  (UPat.var("x") // UPat.var("x"), lambda x: x.const_like(1)), # x//x -> 1
  (UPat.var("x") // 1, lambda x: x),   # x//1 -> x
  (UPat.var("x") // -1, lambda x: -x), # x//-1 -> -x
@@ -49,11 +49,11 @@ symbolic_simple = PatternMatcher([
  (UPat.var("x", dtype=dtypes.bool).logical_not().logical_not(), lambda x: x),
  (UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, True), UPat.const(dtypes.bool, False)), lambda x: x),
  (UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, False), UPat.const(dtypes.bool, True)), lambda x: x.logical_not()),
-  (UPat.var("x", dtype=dtypes.ints+(dtypes.bool,)).trunc(), lambda x: x),
+  (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.index)).trunc(), lambda x: x),
  # ** zero folding **
  (UPat.var("x") < UPat.var("x"), lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x < x -> False
  (UPat.var("x") % UPat.var("x"), lambda x: x.const_like(0)), # x%x -> 0
-  (UPat.var("x", dtype=dtypes.ints) != UPat.var("x", dtype=dtypes.ints),
+  (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.index)) != UPat.var("x"),
   lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x != x -> False (only ints)
  # x*0 -> 0 or 0*x -> 0
  # if x is nan or inf it should render the nan value.
@@ -108,6 +108,7 @@ def fold_unrolled_divs(divs:UOp, denominator: int, fac=1) -> UOp|None:
    if fac!=1:
      if u.op is not Ops.MUL or u.src[1].op is not Ops.CONST or u.src[1].arg != fac: return None
      u = u.src[0]
+    if u.op is Ops.CAST and u.src[0].dtype == dtypes.index: u = u.src[0]
    if not (u.op is Ops.IDIV and u.src[1].op is Ops.CONST): return None
    if denominator != u.src[1].arg: return None
    if (s0:=u.src[0]).vmin < 0: return None
@@ -123,7 +124,7 @@ def fold_unrolled_divs(divs:UOp, denominator: int, fac=1) -> UOp|None:
  for i in range(denominator-len(seen_const)):
    if ans is not None and 0 <= ans.vmin and ans.vmax + i < denominator: seen_const.append(i)
  if sorted(seen_const)==list(range(denominator)):
-    return fac*ans
+    return (fac*ans).cast(divs.dtype)
  return None

 def lt_folding(x:UOp, c:int) -> UOp|None:
@@ -270,10 +271,19 @@ gep_pushing = PatternMatcher([
  (UPat(Ops.WMMA, name="wmma").f(Ops.GEP, name="gep"), gep_through_wmma),
 ])

+cast_folding = PatternMatcher([
+  (UPat.var('x', dtypes.ints+(dtypes.index,)).cast(dtypes.ints+(dtypes.index,), name="a").cast(name="b"),
+    lambda x,a,b: x.cast(b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None),
+  # try to do math in int instead of long
+  (UPat(GroupOp.Binary, src=(UPat.var("x", dtypes.long), UPat.var("y", dtypes.long)), name="u"), lambda u,x,y:
+    x.cast(dtypes.int).alu(u.op, y.cast(dtypes.int)).cast(u.dtype) if not any(v.overflows(dtypes.int) for v in (u,x,y)) else None),
+  ((UPat.var("x", dtypes.index) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)),
+])
+
 commutative = PatternMatcher([
-  # ** COMMUTATIVE flipping (only for ints) **
+  # ** COMMUTATIVE flipping (only for index) **
  # NOTE: this can break merging vector math by only flipping some of them
-  (UPat(GroupOp.Commutative, dtype=dtypes.int, name='x'), lambda x: x.replace(src=x.src[::-1]) if x.src[1].tuplize < x.src[0].tuplize else None),
+  (UPat(GroupOp.Commutative, dtype=dtypes.index, name='x'), lambda x: x.replace(src=x.src[::-1]) if x.src[1].tuplize < x.src[0].tuplize else None),
 ])

 symbolic = symbolic_simple+commutative+PatternMatcher([
@@ -289,10 +299,6 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
  ((UPat.var("y") + UPat.var("x")) + UPat.var("x"), lambda y,x: y+x*2),
  ((UPat.var("x") / UPat.var("x2")) / UPat.var("x3"), lambda x,x2,x3: x/(x2*x3) if x2 is not x3 else None), # (x/x2)/x3 -> x/(x2*x3)
  (-1 * (UPat.var("x") + UPat.cvar("c")), lambda x,c: (-x)+(-c)),  # -(x+c) -> -x + -c
-  (UPat.var('x', dtypes.ints).cast(dtypes.ints, name="a").cast(name="b"),
-    lambda x,a,b: x.cast(b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None),
-  (UPat(GroupOp.Binary, src=(UPat.var("x",dtypes.long), UPat.var("y", dtypes.long)), name="u"), lambda u,x,y:
-    x.cast(dtypes.int).alu(u.op, y.cast(dtypes.int)).cast(u.dtype) if not any(v.overflows(dtypes.int) for v in (u,x,y)) else None),
  # a conditional with the same results either way is a noop, also fold const conditionals
  (UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val),
  (UPat.cvar("gate", vec=False).where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),
@@ -317,27 +323,27 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
  ((UPat.var("x") // UPat.cvar("c1")) // UPat.cvar("c2"), lambda x,c1,c2: x//(c1*c2)), # (x//c1)//c2 -> x//(c1*c2)
  # ** lt **
  # c0*x<c1 for positive int c0,c1
-  ((UPat.cvar("c0", vec=False)*UPat.var("x", dtype=dtypes.ints))<UPat.cvar("c1", vec=False),
+  ((UPat.cvar("c0", vec=False)*UPat.var("x", dtype=dtypes.index))<UPat.cvar("c1", vec=False),
   lambda x,c0,c1: x<math.ceil(c1.arg/c0.arg) if c0.arg > 0 and c1.arg > 0 else None),
  # c0*x<c1 for negative int c0 and non-positive c1
-  ((UPat.cvar("c0", vec=False)*UPat.var("x", dtype=dtypes.ints))<UPat.cvar("c1", vec=False),
+  ((UPat.cvar("c0", vec=False)*UPat.var("x", dtype=dtypes.index))<UPat.cvar("c1", vec=False),
   lambda x,c0,c1: (-x)<(-(math.floor(-c1.arg/-c0.arg))) if c0.arg < 0 and c0.arg != -1 and c1.arg <= 0 else None),
  # x//d<c
-  ((UPat.var("x", dtype=dtypes.ints)//UPat.cvar("d", vec=False))<UPat.cvar("c", vec=False),
+  ((UPat.var("x", dtype=dtypes.index)//UPat.cvar("d", vec=False))<UPat.cvar("c", vec=False),
   lambda x,d,c: (x<(c.arg*d.arg) if c.arg > 0 else x<(c.arg*d.arg-(d.arg-1))) if d.arg > 0 else None),
  # ** move add/mul consts to end (NOTE: this is still happening before constant folding) **
  ((UPat.var("x") + UPat.cvar("c1")) + UPat.var("y"), lambda x,c1,y: (x+y)+c1),
  ((UPat.var("x") * UPat.cvar("c1")) * UPat.var("y"), lambda x,c1,y: (x*y)*c1),
  # *** rules from symbolic ***
  # unrolled arange div folding
-  ((UPat() + UPat()//UPat.cvar("d", vec=False)).named("divs"), lambda divs,d: fold_unrolled_divs(divs, d.arg)),
-  ((UPat() + (UPat()//UPat.cvar("d", vec=False))*UPat.cvar("c")).named("divs"), lambda divs,d,c: fold_unrolled_divs(divs, d.arg, c.arg)),
+  ((UPat()+(UPat()//UPat.cvar("d", vec=False)).or_casted()).named("divs"), lambda divs,d: fold_unrolled_divs(divs, d.arg)),
+  ((UPat()+((UPat()//UPat.cvar("d", vec=False)).or_casted()*UPat.cvar("c"))).named("divs"), lambda divs,d,c: fold_unrolled_divs(divs, d.arg, c.arg)),
  # generic lt folding
-  (UPat.var("x", dtypes.sints)<UPat.cvar("c", vec=False), lambda x,c: lt_folding(x, c.arg) if 0 < c.arg else None),
-  (UPat.var("x", dtypes.sints)*-1 < UPat.var("y", dtypes.sints)*-1, lambda x,y: y<x),
+  (UPat.var("x", dtypes.index)<UPat.cvar("c", vec=False), lambda x,c: lt_folding(x, c.arg) if 0 < c.arg else None),
+  (UPat.var("x", dtypes.index)*-1 < UPat.var("y")*-1, lambda x,y: y<x),
  # canonicalize a simplex with positive coefficients > 0
  # not x < 1 -> X > 0
-  ((UPat.var("x", dtypes.ints)<1).ne(True), lambda x: (newx<1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None),
+  ((UPat.var("x", dtypes.index)<1).ne(True), lambda x: (newx<1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None),
  # ** div **
  # div folding
  ((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d)
@@ -345,30 +351,28 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
  # a range mod its own upper bound is just the range
  (UPat(Ops.RANGE, src=UPat.var("end"), name="r")%UPat.var("end"), lambda r,end: r),
  (UPat(Ops.RANGE, src=UPat.var("end"), name="r")//UPat.var("end"), lambda r,end: r.const_like(0)),
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.var("y"))), cancel_divmod),
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), fold_binary_numerator),
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), fold_divmod_congruence),
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), divide_by_gcd),
-  (UPat(Ops.MOD, dtypes.sints, name="m", src=(UPat.var("x"), UPat.cvar("y", vec=False))), remove_nested_mod),
-  (UPat((Ops.IDIV), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), nest_div_by_smallest_factor),
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), simplify_remainder),
+  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d", src=(UPat.var("x"), UPat.var("y"))), cancel_divmod),
+  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), fold_binary_numerator),
+  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), fold_divmod_congruence),
+  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), divide_by_gcd),
+  (UPat(Ops.MOD, dtypes.index, name="m", src=(UPat.var("x"), UPat.cvar("y", vec=False))), remove_nested_mod),
+  (UPat((Ops.IDIV), dtypes.index, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), nest_div_by_smallest_factor),
+  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), simplify_remainder),
  (UPat.var("x") // UPat.var("d"), lambda x,d: -(x//(-d)) if d.vmax < 0 else None),
  (UPat.var("x") // UPat.var("d"), lambda x,d: -((-x)//d) if x.vmax <=0 else None),
-  ((UPat.var("x", dtypes.sints)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
+  ((UPat.var("x", dtypes.index)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
    lambda x,c,n,d: (-(-(c.arg%d.arg + x - (d.arg-1))//d) + c.arg//d.arg) if x.vmax<=0 and n.vmin>=0 and d.arg>0 else None),
  # ** mod **
  # mod folding
  (UPat.var("x") % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <= 0 else None),
  (UPat.var("x") % UPat.var("d"), lambda x,d: (x%(-d)) if d.vmax <  0 else None),
-  # up + x//c*c + x%c
-  (UPat.var("up") + UPat.var("x", dtypes.ints)//UPat.cvar("c")*UPat.cvar("c") + UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda up,x,c: up+x),
-])+gep_pushing
+])+gep_pushing+cast_folding

 symbolic_flat = symbolic+PatternMatcher([
  # ** combine terms (opinionated) **
  (-1 * (UPat.var("x") + UPat.var("y")), lambda x,y: (-x)+(-y)),  # -(x+y) -> -x + -y
  # (x+y)*c -> x*c+y*c. only for int, float has inf*0=nan issue
-  ((UPat.var("x", dtypes.ints) + UPat.var("y")) * UPat.cvar("c"), lambda x,y,c: x*c+y*c),
+  ((UPat.var("x", dtypes.index) + UPat.var("y")) * UPat.cvar("c"), lambda x,y,c: x*c+y*c),
 ])

 # ******** we take a small aside to "simplify_valid" to rewrite valids ********
@@ -400,6 +404,7 @@ def uop_given_valid(valid:UOp, uop:UOp) -> UOp|None:
  # simplify uop given that valid is True
  for expr,v in bounds.items():
    v0, v1 = (expr.vmin if v[0] is None else v[0], expr.vmax if v[1] is None else v[1])
+    expr = expr.substitute(load_subs)  # make sure expr appears in same form in the uop
    # some expr has lower bound > upper bound -> valid is an empty set and we return None
    if v0 > v1: return None
    # whole node became a const
@@ -459,6 +464,8 @@ def reduce_mul_chain(r:UOp):
 REMOVE_FROM_SINK = {Ops.SINK, Ops.UNROLL, Ops.PTRCAT, Ops.CAT, Ops.NOOP}
 REMOVE_FROM_BARRIER = {Ops.VECTORIZE, Ops.SINK, Ops.CAT, Ops.PTRCAT, Ops.NOOP}
 sym = symbolic_flat+PatternMatcher([
+  # simplify valid
+  (UPat(Ops.AND, name="valid"), simplify_valid),
  # LOAD/STORE -> NOOP
  (UPat.var('x').store(UPat.var('x').load(), allow_any_len=True), lambda x: None if x.dtype.addrspace != AddrSpace.REG else x.src[0].src[0]),
  (UPat(Ops.LOAD, src=(UPat.cvar('c'))), lambda c: c),
@@ -207,7 +207,7 @@ async function renderProfiler() {
    if (eventType === EventTypes.TIMELINE) {
      const levelHeight = baseHeight-padding;
      const levels = [];
-      data.tracks.set(k, { shapes, offsetY });
+      data.tracks.set(k, { shapes, visible:[], offsetY });
      let colorKey, ref;
      for (let j=0; j<eventsLen; j++) {
        const e = {name:strings[u32()], ref:optional(u32()), st:u32(), dur:f32(), info:strings[u32()] || null};
@@ -236,8 +236,6 @@ async function renderProfiler() {
      div.style("height", levelHeight*levels.length+padding+"px").style("pointerEvents", "none");
    } else {
      const peak = u64();
-      const height = heightScale(peak);
-      const yscale = d3.scaleLinear().domain([0, peak]).range([height, 0]);
      let x = 0, y = 0;
      const buf_shapes = new Map(), temp = new Map();
      const timestamps = [];
@@ -268,12 +266,14 @@ async function renderProfiler() {
        v.y.push(v.y.at(-1));
      }
      timestamps.push(dur);
+      const height = heightScale(peak);
+      const yscale = d3.scaleLinear().domain([0, peak]).range([height, 0]);
      for (const [_, {dtype, sz, nbytes, y, x:steps}] of buf_shapes) {
        const x = steps.map(s => timestamps[s]);
        const arg = {tooltipText:`${dtype} len:${formatUnit(sz)}\n${formatUnit(nbytes, "B")}`};
        shapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, shapes.length) });
      }
-      data.tracks.set(k, { shapes, offsetY, height, peak, scaleFactor:maxheight*4/height });
+      data.tracks.set(k, { shapes, visible:[], offsetY, height, peak, scaleFactor:maxheight*4/height });
      div.style("height", height+padding+"px").style("cursor", "pointer").on("click", (e) => {
        const newFocus = e.currentTarget.id === focusedDevice ? null : e.currentTarget.id;
        let offset = 0;
@@ -292,18 +292,16 @@ async function renderProfiler() {
  // draw events on a timeline
  const dpr = window.devicePixelRatio || 1;
  const ellipsisWidth = ctx.measureText("...").width;
-  const rectLst = [];
  function render(transform) {
    zoomLevel = transform;
-    rectLst.length = 0;
    ctx.clearRect(0, 0, canvas.clientWidth, canvas.clientHeight);
    // rescale to match current zoom
    const xscale = d3.scaleLinear().domain([0, dur]).range([0, canvas.clientWidth]);
    const visibleX = xscale.range().map(zoomLevel.invertX, zoomLevel).map(xscale.invert, xscale);
    xscale.domain(visibleX);
-    const yscale = data.axes.y != null ? d3.scaleLinear().domain(data.axes.y.domain).range(data.axes.y.range) : null;
    // draw shapes
-    for (const [_, { offsetY, shapes }] of data.tracks) {
+    for (const [_, { offsetY, shapes, visible }] of data.tracks) {
+      visible.length = 0;
      for (const e of shapes) {
        if (e.width == null) { start = e.x[0]; end = end = e.x[e.x.length-1]; }
        else { start = e.x; end = e.x+e.width; }
@@ -319,12 +317,8 @@ async function renderProfiler() {
          p.closePath();
          ctx.fill(p);
          // NOTE: y coordinates are in reverse order
-          for (let i = 0; i < x.length - 1; i++) {
-            let tooltipText = e.arg.tooltipText;
-            if (yscale != null && ((yaxisVal=yscale.invert(offsetY+e.y1[i]))>0)) {
-              tooltipText += `\nTotal: ${formatUnit(yaxisVal, data.axes.y.fmt)}`;
-            }
-            rectLst.push({ x0:x[i], x1:x[i+1], y0:offsetY+e.y1[i], y1:offsetY+e.y0[i], arg:{...e.arg, tooltipText} });
+          for (let i = 0; i<x.length-1; i++) {
+            visible.push({ x0:x[i], x1:x[i+1], y0:offsetY+e.y1[i], y1:offsetY+e.y0[i], arg:e.arg });
          }
          continue;
        }
@@ -332,7 +326,7 @@ async function renderProfiler() {
        const x = xscale(start);
        const width = xscale(end)-x;
        ctx.fillRect(x, offsetY+e.y, width, e.height);
-        rectLst.push({ y0:offsetY+e.y, y1:offsetY+e.y+e.height, x0:x, x1:x+width, arg:e.arg });
+        visible.push({ y0:offsetY+e.y, y1:offsetY+e.y+e.height, x0:x, x1:x+width, arg:e.arg });
        // add label
        if (e.label == null) continue;
        ctx.textAlign = "left";
@@ -361,8 +355,9 @@ async function renderProfiler() {
      ctx.textBaseline = "top";
      ctx.fillText(formatTime(tick, dur), x+ctx.lineWidth+2, tickSize);
    }
-    if (yscale != null) {
-      drawLine(ctx, [0, 0], yscale.range());
+    if (data.axes.y != null) {
+      drawLine(ctx, [0, 0], data.axes.y.range);
+      const yscale = d3.scaleLinear().domain(data.axes.y.domain).range(data.axes.y.range);
      for (const tick of yscale.ticks()) {
        const y = yscale(tick);
        drawLine(ctx, [0, tickSize], [y, y]);
@@ -398,10 +393,16 @@ async function renderProfiler() {
  new ResizeObserver(([e]) => e.contentRect.width > 0 && resize()).observe(profiler.node());

  function findRectAtPosition(x, y) {
+    let tid = null;
+    for (const k of data.tracks.keys()) {
+      const r = rect(document.getElementById(k));
+      if (y >= r.y && y <= r.y+r.height) { tid = k; break; }
+    }
+    if (tid == null) return;
    const { top, left, width, height } = rect(canvas);
    const X = ((x-left) * (canvas.width/width))/dpr;
    const Y = ((y-top) * (canvas.height/height))/dpr;
-    for (const r of rectLst) {
+    for (const r of data.tracks.get(tid).visible) {
      if (Y>=r.y0 && Y<=r.y1 && X>=r.x0 && X<=r.x1) return r.arg;
    }
  }
@@ -11,7 +11,7 @@ from tinygrad.uop.ops import TrackedGraphRewrite, UOp, Ops, printable, GroupOp,
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, Device
 from tinygrad.renderer import ProgramSpec
 from tinygrad.dtype import dtypes
-from tinygrad.codegen.opt.kernel import axis_colors
+from tinygrad.codegen.opt import axis_colors

 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
               Ops.DEFINE_GLOBAL: "#ffe0b0", Ops.DEFINE_LOCAL: "#ffe0d0", Ops.DEFINE_REG: "#f0ffe0", Ops.REDUCE_AXIS: "#FF6B6B",
Author	SHA1	Message	Date
George HotzandGitHub	8006595cde	Clean up postrange.py by removing comments	2025-09-10 09:40:20 +08:00
geohot	18da4cdcbd	does that pass?	2025-09-08 11:12:46 -07:00
geohot	8b944a32eb	one call to hc opt	2025-09-08 10:30:04 -07:00
nimlgenandGitHub	ebbcdd6577	cpu: use suppress_finalizing (#12071 )	2025-09-08 18:28:09 +03:00
qazalandGitHub	73ca0e870c	viz: index visible rects (#12070 )	2025-09-08 17:37:17 +03:00
chenyuandGitHub	d40f5b766b	default BEAM_PADTO to 0 (#12069 ) seems incorrect, disable by default now	2025-09-08 10:17:03 -04:00
Sieds LyklesandGitHub	75b58fe2d3	move simplify_valid pat to sym (#12065 ) * move simplify_valid pat to sym * fix expectedfailure	2025-09-08 07:01:26 +02:00
chenyuandGitHub	56861852be	enable IMAGE for test_mnist and test_mnist_backward (#12064 ) passes now	2025-09-07 09:06:39 -04:00
nimlgenandGitHub	ef71acc88a	hcq: cleanup fileio iface (#12063 ) * hcq: cleanup fileio iface * typo * _	2025-09-07 15:43:27 +03:00
nimlgenandGitHub	35ddfc3d39	change default cpu_count (#12062 )	2025-09-06 23:30:20 +03:00
nimlgenandGitHub	97187bf8b6	cleanup win and arch checks (#12060 ) * cleanup win and arch checks * stupid mypy	2025-09-06 23:08:46 +03:00
Sieds LyklesandGitHub	f326df8ae8	add type: ignore (#12059 )	2025-09-06 21:17:35 +02:00
George HotzandGitHub	c66935f7b9	only run hcopts once (#12053 ) * only run hcopts once * same?	2025-09-06 11:14:52 -07:00
qazalandGitHub	801be5f7b9	viz: memory graph cleanups (#12057 ) * delete the total nbytes tooltip * split pixel rescaling from layout	2025-09-06 19:44:53 +03:00
nimlgenandGitHub	10ac427aaa	cpu threading (#11951 ) * start cpu threading * fix * fix2 * fix * hacks? * threads * minor * no dsp * dsp 2 * n * more * test * xm * cleaner * readable * f * reorder * when no threads * rangeify * typos * not needed * reapply * remoev this * linter * fixed cpu count in ci * fix * fixes * rm * typo * sort based on speed * test if test works in ci * Revert "test if test works in ci" This reverts commit `1f05edb531`. * do not pad thread	2025-09-06 16:13:43 +03:00
nimlgenandGitHub	2b1844da27	cpu: support several threads in runtime (#12055 )	2025-09-06 13:29:31 +03:00
nimlgenandGitHub	f37b836618	factor out _globalizable_rngs (#12054 )	2025-09-06 13:29:23 +03:00
nimlgenandGitHub	1630c87d0e	run optimize_local_size only when locals supported (#12056 )	2025-09-06 13:29:09 +03:00
Jordan ChalupkaandGitHub	48ec5efad9	only run autogen tests on change (#12049 ) * only run autogen tests on change * example change * rm example change	2025-09-05 23:53:01 -07:00
Sieds LyklesandGitHub	581b2388c2	add dtypes.index (#12015 ) * add dtypes.index * cast shape, stride and mask to dtypes.index in view.create * move pm_lower_index_dtype to ops * DEFINE_VAR is dtype.index by default * merge var_val_using_str * remove int from commutative * fix test_rewrite_map * change that to dtypes.index * change some int to index * shorten those * remove old cast in renderer * cleanup * change that back * add comment * delete comment * just delete those * view doesnt have to cast anymore * adjust comment	2025-09-06 06:03:44 +02:00
Sieds LyklesandGitHub	c6c16b2946	`var_vals` uses str for var (#12011 ) * var_vals is str,int * remove imports * remove print * fix test * change var_vals in hcq * update test_hcq * fix multitensor _device_num var * fix syminfer test * shorten line * p.vars stays list[Variable] * shorten line * vars is back to tuple[Variable, ...] * change var_vals in extra * change var_vals from shapetracker * var_vals is str:int * fix signature	2025-09-06 04:16:12 +02:00
geohot	8658a97197	hotfix: name the shift rewrite better + no ctx there	2025-09-05 19:01:59 -07:00
George HotzandGitHub	6ef3270fc8	fix opt gate (#12050 )	2025-09-05 18:59:54 -07:00
geohot	66c5206b42	hotfix: minimal scheduler copy	2025-09-05 18:24:00 -07:00
geohot	478e758755	Revert "fix scheduler copy (#12048 )" This reverts commit `51b7c40788`.	2025-09-05 18:21:55 -07:00
George HotzandGitHub	51b7c40788	fix scheduler copy (#12048 ) * fix scheduler copy * hand coded opt only runs once	2025-09-05 17:17:49 -07:00
George HotzandGitHub	0123c394e5	early simplfy_merge_adjacent (#12045 ) * do simplify_merge_adjacent before schedule * do simplify_merge_adjacent before schedule * disable that slow test	2025-09-05 16:39:20 -07:00
George HotzandGitHub	8423c06144	delete unused bufs_from_lin (#12044 )	2025-09-05 16:08:28 -07:00
George HotzandGitHub	38dcadf07b	delete kernel.py (#12040 ) * delete kernel.py * delete that file * rip and tear * don't test search * imports * fix torch frontend * not a part of regen	2025-09-05 15:52:07 -07:00