fixups

two stage removal
replace NOOP with AFTER in reg init (#12804 )
2025-10-20 17:30:07 +08:00 · 2025-10-20 16:21:39 +08:00 · 2025-10-20 15:34:32 +08:00 · 2025-10-20 13:11:57 +08:00 · 2025-10-20 12:37:03 +08:00 · 2025-10-20 12:27:56 +08:00
46 changed files with 922 additions and 639 deletions
@@ -131,7 +131,7 @@ jobs:
    - name: UsbGPU copy speeds
      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
    #- name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
@@ -625,16 +625,12 @@ jobs:
      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
    - name: benchmark openpilot 0.9.9 dmonitoring
      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
-    - name: openpilot compile3 0.9.9 driving_vision
-      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.9.9 driving_policy
-      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=7 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.9.9 dmonitoring
-      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
-    - name: openpilot compile3 Space Lab policy + vision
-      run: |
-        PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
-        PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
+    - name: openpilot compile3 0.10.1 driving_vision
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=25 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+    - name: openpilot compile3 0.10.1 driving_policy
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=7 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/refs/heads/master/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.10.1 dmonitoring
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/refs/heads/master/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
@@ -374,15 +374,13 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2081 ALLOWED_GATED_READ_IMAGE=28 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
-      - name: Test openpilot alt model correctness (float32)
-        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
-      - name: Test openpilot fastvits model correctness (float32)
-        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
-      # - name: Test openpilot simple_plan vision model correctness (float32)
-      #   run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
-      - name: Test openpilot LLVM compile
-        run: CPU=1 CPU_LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1452 ALLOWED_GATED_READ_IMAGE=122 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+      - name: Test openpilot CL compile fp16
+        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+      - name: Test openpilot CL compile fp32 (test correctness)
+        run: DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
+      - name: Test openpilot LLVM compile fp16
+        run: FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -1,9 +1,5 @@
 import os, sys, pickle, time, re
 import numpy as np
-if "FLOAT16" not in os.environ: os.environ["FLOAT16"] = "1"
-if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
-if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
-if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"

 from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
@@ -21,11 +17,14 @@ def compile(onnx_file):

  input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
  input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
+
  # Float inputs and outputs to tinyjits for openpilot are always float32
+  # TODO this seems dumb
  input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()}
  Tensor.manual_seed(100)
-  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in sorted(input_shapes.items())}
-  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
+  inputs = {k:Tensor(Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize().numpy(), device='NPY') for k,shp in sorted(input_shapes.items())}
+  if not getenv("NPY_IMG"):
+    inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()}
  print("created tensors")

  run_onnx_jit = TinyJit(lambda **kwargs:
@@ -33,8 +32,6 @@ def compile(onnx_file):
  for i in range(3):
    GlobalCounters.reset()
    print(f"run {i}")
-    inputs = {**{k:v.clone() for k,v in new_inputs.items() if 'img' in k},
-              **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1)):
      ret = run_onnx_jit(**inputs).numpy()
    # copy i == 1 so use of JITBEAM is okay
@@ -69,14 +66,9 @@ def compile(onnx_file):
  print(f"mdl size is {mdl_sz/1e6:.2f}M")
  print(f"pkl size is {pkl_sz/1e6:.2f}M")
  print("**** compile done ****")
-  return test_val
+  return inputs, test_val

-def test_vs_compile(run, new_inputs, test_val=None):
-  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
-
-  # create fake "from_blob" tensors for the inputs, and wrapped NPY tensors for the numpy inputs (these have the same underlying memory)
-  inputs = {**{k:v for k,v in new_inputs.items() if 'img' in k},
-            **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
+def test_vs_compile(run, inputs, test_val=None):

  # run 20 times
  step_times = []
@@ -93,68 +85,49 @@ def test_vs_compile(run, new_inputs, test_val=None):
    min_time = min(step_times)
    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"

-  print(out, val.shape, val.dtype)
  if test_val is not None: np.testing.assert_equal(test_val, val)
  print("**** test done ****")

  # test that changing the numpy changes the model outputs
-  if any([x.device == 'NPY' for x in inputs.values()]):
-    for v in new_inputs_numpy.values(): v *= 2
-    out = run(**inputs)
-    changed_val = out.numpy()
-    np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
+  inputs_2x = {k: Tensor(v.numpy()*2, device=v.device) for k,v in inputs.items()}
+  out = run(**inputs_2x)
+  changed_val = out.numpy()
+  np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
  return val

-def test_vs_onnx(new_inputs, test_val, onnx_file, ort=False):
-  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
+def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
+  import onnxruntime as ort
+  
+  onnx_inputs = {k:v.numpy() for k,v in new_inputs.items()}
  onnx_model = onnx.load(onnx_file)

-  timings = []
-  if ort:
-    # test with onnxruntime
-    import onnxruntime as ort
-    onnx_session = ort.InferenceSession(onnx_file)
-    for _ in range(1 if test_val is not None else 5):
-      st = time.perf_counter()
-      onnx_output = onnx_session.run([onnx_model.graph.output[0].name], {k:v.astype(np.float16) for k,v in new_inputs_numpy.items()})
-      timings.append(time.perf_counter() - st)
-    new_torch_out = onnx_output[0]
-  else:
-    # test with torch
-    import torch
-    from onnx2torch import convert
-    inputs = {k.name:new_inputs_numpy[k.name] for k in onnx_model.graph.input}
-    torch_model = convert(onnx_model).float()
-    with torch.no_grad():
-      for _ in range(1 if test_val is not None else 5):
-        st = time.perf_counter()
-        torch_out = torch_model(*[torch.tensor(x) for x in inputs.values()])
-        timings.append(time.perf_counter() - st)
-      new_torch_out = torch_out.numpy()
+  ORT_TO_NP_DTYPES: dict[str, np.dtype] = {
+    'tensor(float)': np.dtype('float32'),
+    'tensor(float16)': np.dtype('float16'),
+    'tensor(uint8)': np.dtype('uint8'),
+  }

-  if test_val is not None:
-    np.testing.assert_allclose(new_torch_out.reshape(test_val.shape), test_val, atol=1e-4, rtol=1e-2)
-    print("test vs onnx passed")
+  timings = []
+  onnx_session = ort.InferenceSession(onnx_file)
+  onnx_types = {x.name: ORT_TO_NP_DTYPES[x.type] for x in onnx_session.get_inputs()}
+  onnx_inputs = {k:onnx_inputs[k].astype(onnx_types[k]) for k in onnx_inputs}
+
+  for _ in range(1 if test_val is not None else 5):
+    st = time.perf_counter()
+    onnx_output = onnx_session.run([onnx_model.graph.output[0].name], onnx_inputs)
+    timings.append(time.perf_counter() - st)
+
+  np.testing.assert_allclose(onnx_output[0].reshape(test_val.shape), test_val, atol=tol, rtol=tol)
+  print("test vs onnx passed")
  return timings

 if __name__ == "__main__":
  onnx_file = fetch(OPENPILOT_MODEL)
-  test_val = compile(onnx_file) if not getenv("RUN") else None
+  inputs, outputs = compile(onnx_file)

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

-  # same randomness as compile
-  Tensor.manual_seed(100)
-  new_inputs = {nm:Tensor.randn(*st.shape, dtype=dtype).mul(8).realize() for nm, (st, _, dtype, _) in
-                sorted(zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_st_vars_dtype_device))}
-
-  test_val = test_vs_compile(pickle_loaded, new_inputs, test_val)
-  if getenv("BENCHMARK"):
-    for be in ["torch", "ort"]:
-      try:
-        timings = test_vs_onnx(new_inputs, None, onnx_file, be=="ort")
-        print(f"timing {be}: {min(timings)*1000:.2f} ms")
-      except Exception as e:
-        print(f"{be} fail with {e}")
-  if not getenv("FLOAT16"): test_vs_onnx(new_inputs, test_val, onnx_file, getenv("ORT"))
+  test_vs_compile(pickle_loaded, inputs, outputs)
+  if not getenv("FLOAT16"):
+    test_vs_onnx(inputs, outputs, onnx_file, 1e-4)

@@ -7,31 +7,34 @@ import os
 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
 NUM_WAVES = 2
-FLOPS_PER_MATMUL =  16*16*16*2
-INTERNAL_LOOP  = 1_000_000
-INSTRUCTIONS_PER_LOOP = 1_000
+FLOPS_PER_MATMUL = 16*16*16*2
+INTERNAL_LOOP = 1_000_00
+INSTRUCTIONS_PER_LOOP = 200
+DIRECTIVE = ".amdhsa_wavefront_size32 1"

 assemblyTemplate = (pathlib.Path(__file__).parent / "template.s").read_text()

-def launchBenchmark(instruction, vgprIndices, dense = True):
-  if dense:
+def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, extra=""):
+  if accum:
+    instructions = "{} a[0:{}], v[{}:{}], v[{}:{}], 1{}\n".format(instruction, vgprIndices[0],
+                                                                vgprIndices[1], vgprIndices[2],
+                                                                vgprIndices[1], vgprIndices[2], extra)
+  elif dense:
    instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], 1\n".format(instruction, vgprIndices[0],
                                                                vgprIndices[1], vgprIndices[2],
-                                                                vgprIndices[1], vgprIndices[2]) * INSTRUCTIONS_PER_LOOP
+                                                                vgprIndices[1], vgprIndices[2])
  else:
    instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], v{}\n".format(instruction, vgprIndices[0],
-                                                                vgprIndices[1], vgprIndices[2],
-                                                                vgprIndices[3], vgprIndices[4],
-                                                                vgprIndices[5]) * INSTRUCTIONS_PER_LOOP
-  src = assemblyTemplate.replace("INSTRUCTION", instructions)
+                                                                  vgprIndices[1], vgprIndices[2],
+                                                                  vgprIndices[3], vgprIndices[4],
+                                                                  vgprIndices[5])
+  src = assemblyTemplate.replace("INTERNAL_LOOP", str(INTERNAL_LOOP)).replace("INSTRUCTION", instructions*INSTRUCTIONS_PER_LOOP)
+  src = src.replace("DIRECTIVE", DIRECTIVE)
  lib = COMPILER.compile(src)
  fxn = AMDProgram(DEV, "matmul", lib)
-  start = time.perf_counter()
-  fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True) #For some reason the returned time is very small after the first kernel execution
-  end = time.perf_counter()
-  elapsed = end-start
+  elapsed = fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True)
  FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
-  print("{:<29} : {} T(FL)OPS".format(instruction, round(FLOPs/elapsed/10**12, 2)))
+  print(f"{instruction:<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")

 if __name__=="__main__":
  DEVICENUM = os.getenv("DEVICENUM", "0")
@@ -40,18 +43,17 @@ if __name__=="__main__":
  except:
    raise RuntimeError("Error while initiating AMD device")

-  if (ARCH := DEV.arch) not in ['gfx1100', 'gfx1201']:
-    raise RuntimeError("only gfx1100 and gfx1201 supported")
-  COMPILER = HIPCompiler(ARCH)
-
-  if ARCH == 'gfx1100':
+  COMPILER = HIPCompiler(DEV.arch)
+  if DEV.arch in {'gfx1100', 'gfx1103'}:
+    if DEV.arch == 'gfx1103':
+      NUM_WORKGROUPS = 8
    launchBenchmark("v_wmma_bf16_16x16x16_bf16", (7,8,15))
    launchBenchmark("v_wmma_f16_16x16x16_f16", (7,8,15))
    launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,15))
    launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,15))
    launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,9))
    launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,11))
-  if ARCH == 'gfx1201':
+  elif DEV.arch == 'gfx1201':
    NUM_WORKGROUPS = 64
    launchBenchmark("v_wmma_bf16_16x16x16_bf16", (3,4,7))
    launchBenchmark("v_wmma_f16_16x16x16_f16", (3,4,7))
@@ -76,4 +78,18 @@ if __name__=="__main__":
    launchBenchmark("v_swmmac_f32_16x16x32_bf8_fp8", (7,8,9,10,13,14), False)
    launchBenchmark("v_swmmac_f32_16x16x32_bf8_bf8", (7,8,9,10,13,14), False)
    FLOPS_PER_MATMUL = 16*16*64*2
-    launchBenchmark("v_swmmac_i32_16x16x64_iu4", (7,8,9,10,13,14), False)
+    launchBenchmark("v_swmmac_i32_16x16x64_iu4", (7,8,9,10,13,14), False)
+  elif DEV.arch == 'gfx950':
+    DIRECTIVE = ".amdhsa_accum_offset 4"
+    NUM_WORKGROUPS = 256
+    WAVE_SIZE = 64
+    NUM_WAVES = 4
+    launchBenchmark("v_mfma_f32_16x16x16_bf16", (3,0,1), accum=True)
+    FLOPS_PER_MATMUL = 16*16*32*2
+    launchBenchmark("v_mfma_f32_16x16x32_bf16", (3,0,3), accum=True)
+    FLOPS_PER_MATMUL = 16*16*128*2
+    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,7), accum=True) # fp8
+    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,5), accum=True, extra=", cbsz:2 blgp:2") # fp6
+    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,3), accum=True, extra=", cbsz:4 blgp:4") # fp4
+  else:
+    raise RuntimeError(f"arch {DEV.arch} not supported.")
@@ -1,9 +1,9 @@
    .text
    .globl matmul
-    .p2align 8 
+    .p2align 8
    .type matmul,@function
 matmul:
-    s_mov_b32 s1, 1000000
+    s_mov_b32 s1, INTERNAL_LOOP
    s_mov_b32 s2, 0
    inner_loop:
        INSTRUCTION
@@ -17,7 +17,7 @@ matmul:
 .amdhsa_kernel matmul
  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
-  .amdhsa_wavefront_size32 1
+  DIRECTIVE
 .end_amdhsa_kernel

 .amdgpu_metadata
@@ -28,7 +28,7 @@ amdhsa.version:
 amdhsa.kernels:
  - .name: matmul
    .symbol: matmul.kd
-    .kernarg_segment_size:  0
+    .kernarg_segment_size: 0
    .group_segment_fixed_size: 0
    .private_segment_fixed_size: 0
    .kernarg_segment_align: 4
@@ -36,6 +36,5 @@ amdhsa.kernels:
    .sgpr_count: 8
    .vgpr_count: 32
    .max_flat_workgroup_size: 1024
-    .args:
 ...
 .end_amdgpu_metadata
@@ -1,7 +1,8 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <database xmlns="http://nouveau.freedesktop.org/"
 xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
-xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
+xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
+<import file="freedreno_copyright.xml"/>
 <import file="adreno/adreno_common.xml"/>

 <enum name="vgt_event_type" varset="chip">
@@ -20,9 +21,9 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="HLSQ_FLUSH" value="7" variants="A3XX-A4XX"/>
 	<value name="VIZQUERY_END" value="8" variants="A2XX"/>
 	<value name="SC_WAIT_WC" value="9" variants="A2XX"/>
-	<value name="WRITE_PRIMITIVE_COUNTS" value="9" variants="A6XX"/>
-	<value name="START_PRIMITIVE_CTRS" value="11" variants="A6XX"/>
-	<value name="STOP_PRIMITIVE_CTRS" value="12" variants="A6XX"/>
+	<value name="WRITE_PRIMITIVE_COUNTS" value="9" variants="A6XX-"/>
+	<value name="START_PRIMITIVE_CTRS" value="11" variants="A6XX-"/>
+	<value name="STOP_PRIMITIVE_CTRS" value="12" variants="A6XX-"/>
 	<!-- Not sure that these 4 events don't have the same meaning as on A5XX+ -->
 	<value name="RST_PIX_CNT" value="13" variants="A2XX-A4XX"/>
 	<value name="RST_VTX_CNT" value="14" variants="A2XX-A4XX"/>
@@ -30,8 +31,8 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="STAT_EVENT" value="16" variants="A2XX-A4XX"/>
 	<value name="CACHE_FLUSH_AND_INV_TS_EVENT" value="20" variants="A2XX-A4XX"/>
 	<doc>
-		If A6XX_RB_SAMPLE_COUNT_CONTROL.copy is true, writes OQ Z passed
-		sample counts to RB_SAMPLE_COUNT_ADDR.  This writes to main
+		If A6XX_RB_SAMPLE_COUNTER_CNTL.copy is true, writes OQ Z passed
+		sample counts to RB_SAMPLE_COUNTER_BASE.  This writes to main
 		memory, skipping UCHE.
 	</doc>
 	<value name="ZPASS_DONE" value="21"/>
@@ -96,6 +97,13 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	</doc>
 	<value name="BLIT" value="30" variants="A5XX-"/>

+	<doc>
+	Flip between the primary and secondary LRZ buffers. This is used
+	for concurrent binning, so that BV can write to one buffer while
+	BR reads from the other.
+	</doc>
+	<value name="LRZ_FLIP_BUFFER" value="36" variants="A7XX-"/>
+
 	<doc>
 		Clears based on GRAS_LRZ_CNTL configuration, could clear
 		fast-clear buffer or LRZ direction.
@@ -112,11 +120,12 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="LRZ_FLUSH" value="38" variants="A5XX-"/>
 	<value name="BLIT_OP_FILL_2D" value="39" variants="A5XX-"/>
 	<value name="BLIT_OP_COPY_2D" value="40" variants="A5XX-A6XX"/>
-	<value name="UNK_40" value="40" variants="A7XX"/>
+	<value name="LRZ_CACHE_INVALIDATE" value="40" variants="A7XX-"/>
+	<value name="LRZ_Q_CACHE_INVALIDATE" value="41" variants="A7XX-"/>
 	<value name="BLIT_OP_SCALE_2D" value="42" variants="A5XX-"/>
 	<value name="CONTEXT_DONE_2D" value="43" variants="A5XX-"/>
-	<value name="UNK_2C" value="44" variants="A5XX-"/>
-	<value name="UNK_2D" value="45" variants="A5XX-"/>
+	<value name="VSC_BINNING_START" value="44" variants="A5XX-"/>
+	<value name="VSC_BINNING_END" value="45" variants="A5XX-"/>

 	<!-- a6xx events -->
 	<doc>
@@ -129,21 +138,22 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<!-- note, some of these are the same as a6xx, just named differently -->

 	<doc> Doesn't seem to do anything </doc>
-	<value name="DUMMY_EVENT" value="1" variants="A7XX"/>
-	<value name="CCU_INVALIDATE_DEPTH" value="24" variants="A7XX"/>
-	<value name="CCU_INVALIDATE_COLOR" value="25" variants="A7XX"/>
-	<value name="CCU_RESOLVE_CLEAN" value="26" variants="A7XX"/>
-	<value name="CCU_FLUSH_DEPTH" value="28" variants="A7XX"/>
-	<value name="CCU_FLUSH_COLOR" value="29" variants="A7XX"/>
-	<value name="CCU_RESOLVE" value="30" variants="A7XX"/>
-	<value name="CCU_END_RESOLVE_GROUP" value="31" variants="A7XX"/>
-	<value name="CCU_CLEAN_DEPTH" value="32" variants="A7XX"/>
-	<value name="CCU_CLEAN_COLOR" value="33" variants="A7XX"/>
-	<value name="CACHE_RESET" value="48" variants="A7XX"/>
-	<value name="CACHE_CLEAN" value="49" variants="A7XX"/>
+	<value name="DUMMY_EVENT" value="1" variants="A7XX-"/>
+	<value name="CCU_INVALIDATE_DEPTH" value="24" variants="A7XX-"/>
+	<value name="CCU_INVALIDATE_COLOR" value="25" variants="A7XX-"/>
+	<value name="CCU_RESOLVE_CLEAN" value="26" variants="A7XX-"/>
+	<value name="CCU_FLUSH_DEPTH" value="28" variants="A7XX-"/>
+	<value name="CCU_FLUSH_COLOR" value="29" variants="A7XX-"/>
+	<value name="CCU_RESOLVE" value="30" variants="A7XX-"/>
+	<value name="CCU_END_RESOLVE_GROUP" value="31" variants="A7XX-"/>
+	<value name="CCU_CLEAN_DEPTH" value="32" variants="A7XX-"/>
+	<value name="CCU_CLEAN_COLOR" value="33" variants="A7XX-"/>
+	<value name="CACHE_RESET" value="48" variants="A7XX-"/>
+	<value name="CACHE_CLEAN" value="49" variants="A7XX-"/>
 	<!-- TODO: deal with name conflicts with other gens -->
-	<value name="CACHE_FLUSH7" value="50" variants="A7XX"/>
-	<value name="CACHE_INVALIDATE7" value="51" variants="A7XX"/>
+	<value name="CACHE_FLUSH7" value="50" variants="A7XX-"/>
+	<value name="CACHE_INVALIDATE7" value="51" variants="A7XX-"/>
+	<value name="DEPTH_BUFFER_FLIP" value="0x3d" variants="A8XX-"/>
 </enum>

 <enum name="pc_di_primtype">
@@ -324,7 +334,7 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<doc>fetch state sub-blocks and initiate shader code DMAs</doc>
 	<value name="CP_SET_STATE" value="0x25"/>
 	<doc>load constant into chip and to memory</doc>
-	<value name="CP_SET_CONSTANT" value="0x2d"/>
+	<value name="CP_SET_CONSTANT" value="0x2d" variants="A2XX"/>
 	<doc>load sequencer instruction memory (pointer-based)</doc>
 	<value name="CP_IM_LOAD" value="0x27"/>
 	<doc>load sequencer instruction memory (code embedded in packet)</doc>
@@ -371,7 +381,7 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="CP_LOAD_STATE" value="0x30" variants="A3XX"/>
 	<value name="CP_LOAD_STATE4" value="0x30" variants="A4XX-A5XX"/>
 	<doc>Conditionally load a IB based on a flag, prefetch enabled</doc>
-	<value name="CP_COND_INDIRECT_BUFFER_PFE" value="0x3a"/>
+	<value name="CP_COND_INDIRECT_BUFFER_PFE" value="0x3a" variants="A3XX-A5XX"/>
 	<doc>Conditionally load a IB based on a flag, prefetch disabled</doc>
 	<value name="CP_COND_INDIRECT_BUFFER_PFD" value="0x32" variants="A3XX"/>
 	<doc>Load a buffer with pre-fetch enabled</doc>
@@ -514,7 +524,7 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<!--
 	Seems to set the mode flags which control which CP_SET_DRAW_STATE
 	packets are executed, based on their ENABLE_MASK values
-	
+
 	CP_SET_MODE w/ payload of 0x1 seems to cause CP_SET_DRAW_STATE
 	packets w/ ENABLE_MASK & 0x6 to execute immediately
 	 -->
@@ -537,7 +547,7 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="CP_LOAD_STATE6_GEOM" value="0x32" variants="A6XX-"/>
 	<value name="CP_LOAD_STATE6_FRAG" value="0x34" variants="A6XX-"/>
 	<!--
-	Note: For IBO state (Image/SSBOs) which have shared state across
+	Note: For UAV state (Image/SSBOs) which have shared state across
 	shader stages, for 3d pipeline CP_LOAD_STATE6 is used.  But for
 	compute shaders, CP_LOAD_STATE6_FRAG is used.  Possibly they are
 	interchangable.
@@ -566,20 +576,21 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="IN_PREEMPT" value="0x0f" variants="A6XX-"/>

 	<!-- TODO do these exist on A5xx? -->
-	<value name="CP_SCRATCH_WRITE" value="0x4c" variants="A6XX"/>
+	<value name="CP_SCRATCH_WRITE" value="0x4c" variants="A6XX-"/>
 	<value name="CP_REG_TO_MEM_OFFSET_MEM" value="0x74" variants="A6XX-"/>
 	<value name="CP_REG_TO_MEM_OFFSET_REG" value="0x72" variants="A6XX-"/>
 	<value name="CP_WAIT_MEM_GTE" value="0x14" variants="A6XX"/>
 	<value name="CP_WAIT_TWO_REGS" value="0x70" variants="A6XX"/>
 	<value name="CP_MEMCPY" value="0x75" variants="A6XX-"/>
 	<value name="CP_SET_BIN_DATA5_OFFSET" value="0x2e" variants="A6XX-"/>
+	<!-- A750+, set in place of CP_SET_BIN_DATA5_OFFSET but has different values -->
+	<value name="CP_SET_UNK_BIN_DATA" value="0x2d" variants="A7XX-"/>
 	<doc>
                Write CP_CONTEXT_SWITCH_*_INFO from CP to the following dwords,
                and forcibly switch to the indicated context.
 	</doc>
 	<value name="CP_CONTEXT_SWITCH" value="0x54" variants="A6XX"/>
-	<!-- Note, kgsl calls this CP_SET_AMBLE: -->
-	<value name="CP_SET_CTXSWITCH_IB" value="0x55" variants="A6XX-"/>
+	<value name="CP_SET_AMBLE" value="0x55" variants="A6XX-"/>

 	<!--
 	Seems to always have the payload:
@@ -630,8 +641,7 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="CP_BV_BR_COUNT_OPS" value="0x1b" variants="A7XX-"/>
 	<doc> Clears, adds to local, or adds to global timestamp </doc>
 	<value name="CP_MODIFY_TIMESTAMP" value="0x1c" variants="A7XX-"/>
-	<!-- similar to CP_CONTEXT_REG_BUNCH, but discards first two dwords?? -->
-	<value name="CP_CONTEXT_REG_BUNCH2" value="0x5d" variants="A7XX-"/>
+	<value name="CP_NON_CONTEXT_REG_BUNCH" value="0x5d" variants="A7XX-"/>
 	<doc>
 		Write to a scratch memory that is read by CP_REG_TEST with
 		SOURCE_SCRATCH_MEM set. It's not the same scratch as scratch registers.
@@ -648,6 +658,11 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">

 	<doc>Reset various on-chip state used for synchronization</doc>
 	<value name="CP_RESET_CONTEXT_STATE" value="0x1f" variants="A7XX-"/>
+
+	<doc>Invalidates the "CCHE" introduced on a740</doc>
+	<value name="CP_CCHE_INVALIDATE" value="0x3a" variants="A7XX-"/>
+
+	<value name="CP_SCOPE_CNTL" value="0x6c" variants="A7XX-"/>
 </enum>


@@ -790,14 +805,14 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<value name="SB6_GS_SHADER" value="0xb"/>
 		<value name="SB6_FS_SHADER" value="0xc"/>
 		<value name="SB6_CS_SHADER" value="0xd"/>
-		<value name="SB6_IBO"       value="0xe"/>
-		<value name="SB6_CS_IBO"    value="0xf"/>
+		<value name="SB6_UAV"       value="0xe"/>
+		<value name="SB6_CS_UAV"    value="0xf"/>
 	</enum>
 	<enum name="a6xx_state_type">
 		<value name="ST6_SHADER" value="0"/>
 		<value name="ST6_CONSTANTS" value="1"/>
 		<value name="ST6_UBO" value="2"/>
-		<value name="ST6_IBO" value="3"/>
+		<value name="ST6_UAV" value="3"/>
 	</enum>
 	<enum name="a6xx_state_src">
 		<value name="SS6_DIRECT" value="0"/>
@@ -903,12 +918,6 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>

 	<stripe varset="chip" variants="A5XX-">
-		<reg32 offset="4" name="4">
-			<bitfield name="INDX_BASE_LO" low="0" high="31"/>
-		</reg32>
-		<reg32 offset="5" name="5">
-			<bitfield name="INDX_BASE_HI" low="0" high="31"/>
-		</reg32>
 		<reg64 offset="4" name="INDX_BASE" type="address"/>
 		<reg32 offset="6" name="6">
 			<!-- max # of elements in index buffer -->
@@ -1084,8 +1093,10 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 			<bitfield name="BINNING" pos="20" varset="chip" variants="A6XX-" type="boolean"/>
 			<bitfield name="GMEM" pos="21" varset="chip" variants="A6XX-" type="boolean"/>
 			<bitfield name="SYSMEM" pos="22" varset="chip" variants="A6XX-" type="boolean"/>
-			<bitfield name="GROUP_ID" low="24" high="28" type="uint"/>
+			<!-- high bit is 28 until a750: -->
+			<bitfield name="GROUP_ID" low="24" high="29" type="uint"/>
 		</reg32>
+		<reg64 offset="1" name="ADDR" type="address"/>
 		<reg32 offset="1" name="1">
 			<bitfield name="ADDR_LO" low="0" high="31" type="hex"/>
 		</reg32>
@@ -1119,39 +1130,63 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>
 </domain>

+<enum name="a7xx_abs_mask_mode">
+	<value name="ABS_MASK" value="0x1"/>
+	<value name="NO_ABS_MASK" value="0x0"/>
+</enum>
+
 <domain name="CP_SET_BIN_DATA5" width="32">
 	<reg32 offset="0" name="0">
+		<bitfield name="VSC_MASK" low="0" high="15" type="hex">
+			<doc>
+				A mask of bins, starting at VSC_N, whose
+				visibility is OR'd together. A value of 0 is
+				interpreted as 1 (i.e. just use VSC_N for
+				visbility) for backwards compatibility. Only
+				exists on a7xx.
+			</doc>
+		</bitfield>
 		<!-- equiv to PC_VSTREAM_CONTROL.SIZE on a3xx/a4xx: -->
 		<bitfield name="VSC_SIZE" low="16" high="21" type="uint"/>
 		<!-- equiv to PC_VSTREAM_CONTROL.N on a3xx/a4xx: -->
 		<bitfield name="VSC_N" low="22" high="26" type="uint"/>
+		<bitfield name="ABS_MASK" pos="28" type="a7xx_abs_mask_mode" addvariant="yes">
+			<doc>
+				If this field is 1, VSC_MASK and VSC_N are
+				ignored and instead a new ordinal immediately
+				after specifies the full 32-bit mask of bins
+				to use. The mask is "absolute" instead of
+				relative to VSC_N.
+			</doc>
+		</bitfield>
 	</reg32>
-	<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
-	<reg32 offset="1" name="1">
-		<bitfield name="BIN_DATA_ADDR_LO" low="0" high="31" type="hex"/>
-	</reg32>
-	<reg32 offset="2" name="2">
-		<bitfield name="BIN_DATA_ADDR_HI" low="0" high="31" type="hex"/>
-	</reg32>
-	<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
-	<reg32 offset="3" name="3">
-		<bitfield name="BIN_SIZE_ADDRESS_LO" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="4" name="4">
-		<bitfield name="BIN_SIZE_ADDRESS_HI" low="0" high="31"/>
-	</reg32>
-	<!-- new on a6xx, where BIN_DATA_ADDR is the DRAW_STRM: -->
-	<reg32 offset="5" name="5">
-		<bitfield name="BIN_PRIM_STRM_LO" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="6" name="6">
-		<bitfield name="BIN_PRIM_STRM_HI" low="0" high="31"/>
-	</reg32>
-	<!--
-		a7xx adds a few more addresses to the end of the pkt
-	 -->
-	<reg64 offset="7" name="7"/>
-	<reg64 offset="9" name="9"/>
+	<stripe varset="a7xx_abs_mask_mode" variants="NO_ABS_MASK">
+		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
+		<reg64 offset="1" name="BIN_DATA_ADDR" type="address"/>
+		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
+		<reg64 offset="3" name="BIN_SIZE_ADDR" type="address"/>
+		<!-- new on a6xx, where BIN_DATA_ADDR is the DRAW_STRM: -->
+		<reg64 offset="5" name="BIN_PRIM_STRM" type="address"/>
+		<!--
+			a7xx adds a few more addresses to the end of the pkt
+		 -->
+		<reg64 offset="7" name="7"/>
+		<reg64 offset="9" name="9"/>
+	</stripe>
+	<stripe varset="a7xx_abs_mask_mode" variants="ABS_MASK">
+		<reg32 offset="1" name="ABS_MASK"/>
+		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
+		<reg64 offset="2" name="BIN_DATA_ADDR" type="address"/>
+		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
+		<reg64 offset="4" name="BIN_SIZE_ADDR" type="address"/>
+		<!-- new on a6xx, where BIN_DATA_ADDR is the DRAW_STRM: -->
+		<reg64 offset="6" name="BIN_PRIM_STRM" type="address"/>
+		<!--
+			a7xx adds a few more addresses to the end of the pkt
+		 -->
+		<reg64 offset="8" name="8"/>
+		<reg64 offset="10" name="10"/>
+	</stripe>
 </domain>

 <domain name="CP_SET_BIN_DATA5_OFFSET" width="32">
@@ -1162,23 +1197,42 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
                stream is recorded.
 	</doc>
 	<reg32 offset="0" name="0">
+		<bitfield name="VSC_MASK" low="0" high="15" type="hex"/>
 		<!-- equiv to PC_VSTREAM_CONTROL.SIZE on a3xx/a4xx: -->
 		<bitfield name="VSC_SIZE" low="16" high="21" type="uint"/>
 		<!-- equiv to PC_VSTREAM_CONTROL.N on a3xx/a4xx: -->
 		<bitfield name="VSC_N" low="22" high="26" type="uint"/>
+		<bitfield name="ABS_MASK" pos="28" type="a7xx_abs_mask_mode" addvariant="yes"/>
 	</reg32>
-	<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
-	<reg32 offset="1" name="1">
-		<bitfield name="BIN_DATA_OFFSET" low="0" high="31" type="uint"/>
-	</reg32>
-	<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
-	<reg32 offset="2" name="2">
-		<bitfield name="BIN_SIZE_OFFSET" low="0" high="31" type="uint"/>
-	</reg32>
-	<!-- BIN_DATA2_ADDR -> VSC_PIPE[p].DATA2_ADDRESS -->
-	<reg32 offset="3" name="3">
-		<bitfield name="BIN_DATA2_OFFSET" low="0" high="31" type="uint"/>
-	</reg32>
+	<stripe varset="a7xx_abs_mask_mode" variants="NO_ABS_MASK">
+		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
+		<reg32 offset="1" name="1">
+			<bitfield name="BIN_DATA_OFFSET" low="0" high="31" type="uint"/>
+		</reg32>
+		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
+		<reg32 offset="2" name="2">
+			<bitfield name="BIN_SIZE_OFFSET" low="0" high="31" type="uint"/>
+		</reg32>
+		<!-- BIN_DATA2_ADDR -> VSC_PIPE[p].DATA2_ADDRESS -->
+		<reg32 offset="3" name="3">
+			<bitfield name="BIN_DATA2_OFFSET" low="0" high="31" type="uint"/>
+		</reg32>
+	</stripe>
+	<stripe varset="a7xx_abs_mask_mode" variants="ABS_MASK">
+		<reg32 offset="1" name="ABS_MASK"/>
+		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
+		<reg32 offset="2" name="2">
+			<bitfield name="BIN_DATA_OFFSET" low="0" high="31" type="uint"/>
+		</reg32>
+		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
+		<reg32 offset="3" name="3">
+			<bitfield name="BIN_SIZE_OFFSET" low="0" high="31" type="uint"/>
+		</reg32>
+		<!-- BIN_DATA2_ADDR -> VSC_PIPE[p].DATA2_ADDRESS -->
+		<reg32 offset="4" name="4">
+			<bitfield name="BIN_DATA2_OFFSET" low="0" high="31" type="uint"/>
+		</reg32>
+	</stripe>
 </domain>

 <domain name="CP_REG_RMW" width="32">
@@ -1196,6 +1250,9 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</doc>
 	<reg32 offset="0" name="0">
 		<bitfield name="DST_REG" low="0" high="17" type="hex"/>
+		<bitfield name="DST_SCRATCH" pos="19" type="boolean" varset="chip" variants="A7XX-"/>
+		<!-- skip implied CP_WAIT_FOR_IDLE + CP_WAIT_FOR_ME -->
+		<bitfield name="SKIP_WAIT_FOR_ME" pos="23" type="boolean" varset="chip" variants="A7XX-"/>
 		<bitfield name="ROTATE" low="24" high="28" type="uint"/>
 		<bitfield name="SRC1_ADD" pos="29" type="boolean"/>
 		<bitfield name="SRC1_IS_REG" pos="30" type="boolean"/>
@@ -1209,7 +1266,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>
 </domain>

-<domain name="CP_REG_TO_MEM" width="32">
+<domain name="CP_REG_TO_MEM" width="32" prefix="chip">
 	<reg32 offset="0" name="0">
 		<bitfield name="REG" low="0" high="17" type="hex"/>
 		<!-- number of registers/dwords copied is max(CNT, 1). -->
@@ -1217,12 +1274,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="64B" pos="30" type="boolean"/>
 		<bitfield name="ACCUMULATE" pos="31" type="boolean"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="DEST" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
-		<bitfield name="DEST_HI" low="0" high="31"/>
-	</reg32>
+	<stripe varset="chip" variants="A2XX-A4XX">
+		<reg32 offset="1" name="DEST" type="address"/>
+	</stripe>
+	<stripe varset="chip" variants="A5XX-">
+		<reg64 offset="1" name="DEST" type="address"/>
+	</stripe>
 </domain>

 <domain name="CP_REG_TO_MEM_OFFSET_REG" width="32">
@@ -1238,12 +1295,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="64B" pos="30" type="boolean"/>
 		<bitfield name="ACCUMULATE" pos="31" type="boolean"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="DEST" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
-		<bitfield name="DEST_HI" low="0" high="31"/>
-	</reg32>
+	<reg64 offset="1" name="DEST" type="waddress"/>
 	<reg32 offset="3" name="3">
 		<bitfield name="OFFSET0" low="0" high="17" type="hex"/>
 		<bitfield name="OFFSET0_SCRATCH" pos="19" type="boolean"/>
@@ -1263,18 +1315,8 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="64B" pos="30" type="boolean"/>
 		<bitfield name="ACCUMULATE" pos="31" type="boolean"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="DEST" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
-		<bitfield name="DEST_HI" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="3" name="3">
-		<bitfield name="OFFSET_LO" low="0" high="31" type="hex"/>
-	</reg32>
-	<reg32 offset="4" name="4">
-		<bitfield name="OFFSET_HI" low="0" high="31" type="hex"/>
-	</reg32>
+	<reg64 offset="1" name="DEST" type="waddress"/>
+	<reg64 offset="3" name="OFFSET" type="waddress"/>
 </domain>

 <domain name="CP_MEM_TO_REG" width="32">
@@ -1287,12 +1329,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<!-- does the same thing as CP_MEM_TO_MEM::UNK31 -->
 		<bitfield name="UNK31" pos="31" type="boolean"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="SRC" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
-		<bitfield name="SRC_HI" low="0" high="31"/>
-	</reg32>
+	<stripe varset="chip" variants="A2XX-A4XX">
+		<reg32 offset="1" name="SRC" type="address"/>
+	</stripe>
+	<stripe varset="chip" variants="A5XX-">
+		<reg64 offset="1" name="SRC" type="address"/>
+	</stripe>
 </domain>

 <domain name="CP_MEM_TO_MEM" width="32">
@@ -1312,6 +1354,10 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<!-- some other kind of wait -->
 		<bitfield name="UNK31" pos="31" type="boolean"/>
 	</reg32>
+	<reg64 offset="1" name="DST" type="waddress"/>
+	<reg64 offset="3" name="SRC_A" type="address"/>
+	<reg64 offset="5" name="SRC_B" type="address"/>
+	<reg64 offset="7" name="SRC_C" type="address"/>
 	<!--
 	followed by sequence of addresses.. the first is the
 	destination and the rest are N src addresses which are
@@ -1346,6 +1392,8 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="SCRATCH" low="20" high="22" type="uint"/>
 		<!-- number of registers/dwords copied is CNT + 1. -->
 		<bitfield name="CNT" low="24" high="26" type="uint"/>
+		<!-- skip implied CP_WAIT_FOR_IDLE + CP_WAIT_FOR_ME -->
+		<bitfield name="SKIP_WAIT_FOR_ME" pos="27" type="boolean" varset="chip" variants="A7XX-"/>
 	</reg32>
 </domain>

@@ -1368,12 +1416,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 </domain>

 <domain name="CP_MEM_WRITE" width="32">
-	<reg32 offset="0" name="0">
-		<bitfield name="ADDR_LO" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="ADDR_HI" low="0" high="31"/>
-	</reg32>
+	<stripe varset="chip" variants="A2XX-A4XX">
+		<reg32 offset="0" name="ADDR" type="address"/>
+	</stripe>
+	<stripe varset="chip" variants="A5XX-">
+		<reg64 offset="0" name="ADDR" type="address"/>
+	</stripe>
 	<!-- followed by the DWORDs to write -->
 </domain>

@@ -1425,24 +1473,14 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="POLL" low="4" high="5" type="poll_memory_type"/>
 		<bitfield name="WRITE_MEMORY" pos="8" type="boolean"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="POLL_ADDR_LO" low="0" high="31" type="hex"/>
-	</reg32>
-	<reg32 offset="2" name="2">
-		<bitfield name="POLL_ADDR_HI" low="0" high="31" type="hex"/>
-	</reg32>
+	<reg64 offset="1" name="POLL_ADDR" type="address"/>
 	<reg32 offset="3" name="3">
 		<bitfield name="REF" low="0" high="31"/>
 	</reg32>
 	<reg32 offset="4" name="4">
 		<bitfield name="MASK" low="0" high="31"/>
 	</reg32>
-	<reg32 offset="5" name="5">
-		<bitfield name="WRITE_ADDR_LO" low="0" high="31" type="hex"/>
-	</reg32>
-	<reg32 offset="6" name="6">
-		<bitfield name="WRITE_ADDR_HI" low="0" high="31" type="hex"/>
-	</reg32>
+	<reg64 offset="5" name="WRITE_ADDR" type="waddress"/>
 	<reg32 offset="7" name="7">
 		<bitfield name="WRITE_DATA" low="0" high="31"/>
 	</reg32>
@@ -1457,12 +1495,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<!-- Reserved for flags, presumably? Unused in FW -->
 		<bitfield name="RESERVED" low="0" high="31" type="hex"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="POLL_ADDR_LO" low="0" high="31" type="hex"/>
-	</reg32>
-	<reg32 offset="2" name="2">
-		<bitfield name="POLL_ADDR_HI" low="0" high="31" type="hex"/>
-	</reg32>
+	<reg64 offset="1" name="POLL_ADDR" type="address"/>
 	<reg32 offset="3" name="3">
 		<bitfield name="REF" low="0" high="31"/>
 	</reg32>
@@ -1480,12 +1513,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="POLL" low="4" high="5" type="poll_memory_type"/>
 		<bitfield name="WRITE_MEMORY" pos="8" type="boolean"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="POLL_ADDR_LO" low="0" high="31" type="hex"/>
-	</reg32>
-	<reg32 offset="2" name="2">
-		<bitfield name="POLL_ADDR_HI" low="0" high="31" type="hex"/>
-	</reg32>
+	<reg64 offset="1" name="POLL_ADDR" type="address"/>
 	<reg32 offset="3" name="3">
 		<bitfield name="REF" low="0" high="31"/>
 	</reg32>
@@ -1619,12 +1647,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	TODO what is gpuaddr for, seems to be all 0's.. maybe needed for
 	context switch?
 	 -->
-	<reg32 offset="1" name="1">
-		<bitfield name="ADDR_0_LO" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="2" name="2">
-		<bitfield name="ADDR_0_HI" low="0" high="31"/>
-	</reg32>
+	<reg64 offset="1" name="ADDR" type="waddress"/>
 	<reg32 offset="3" name="3">
 		<!-- ??? -->
 	</reg32>
@@ -1653,8 +1676,8 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="WRITE_SAMPLE_COUNT" pos="12" type="boolean"/>
 		<!-- Write sample count at (iova + 16) -->
 		<bitfield name="SAMPLE_COUNT_END_OFFSET" pos="13" type="boolean"/>
-		<!-- *(iova + 8) = *(iova + 16) - *iova -->
-		<bitfield name="WRITE_SAMPLE_COUNT_DIFF" pos="14" type="boolean"/>
+		<!-- *(iova + 8) += *(iova + 16) - *iova -->
+		<bitfield name="WRITE_ACCUM_SAMPLE_COUNT_DIFF" pos="14" type="boolean"/>

 		<!-- Next 4 flags are valid to set only when concurrent binning is enabled -->
 		<!-- Increment 16b BV counter. Valid only in BV pipe -->
@@ -1668,15 +1691,11 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="WRITE_DST" pos="24" type="event_write_dst" addvariant="yes"/>
 		<!-- Writes into WRITE_DST from WRITE_SRC. RB_DONE_TS requires WRITE_ENABLED. -->
 		<bitfield name="WRITE_ENABLED" pos="27" type="boolean"/>
+		<bitfield name="IRQ" pos="31" type="boolean"/>
 	</reg32>

 	<stripe varset="event_write_dst" variants="EV_DST_RAM">
-		<reg32 offset="1" name="1">
-			<bitfield name="ADDR_0_LO" low="0" high="31"/>
-		</reg32>
-		<reg32 offset="2" name="2">
-			<bitfield name="ADDR_0_HI" low="0" high="31"/>
-		</reg32>
+		<reg64 offset="1" name="1" type="waddress"/>
 		<reg32 offset="3" name="3">
 			<bitfield name="PAYLOAD_0" low="0" high="31"/>
 		</reg32>
@@ -1743,9 +1762,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	<reg32 offset="0" name="0">
 	</reg32>
 	<stripe varset="chip" variants="A4XX">
-		<reg32 offset="1" name="1">
-			<bitfield name="ADDR" low="0" high="31"/>
-		</reg32>
+		<reg32 offset="1" name="ADDR" type="address"/>
 		<reg32 offset="2" name="2">
 			<!-- localsize is value minus one: -->
 			<bitfield name="LOCALSIZEX" low="2" high="11" type="uint"/>
@@ -1754,12 +1771,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		</reg32>
 	</stripe>
 	<stripe varset="chip" variants="A5XX-">
-		<reg32 offset="1" name="1">
-			<bitfield name="ADDR_LO" low="0" high="31"/>
-		</reg32>
-		<reg32 offset="2" name="2">
-			<bitfield name="ADDR_HI" low="0" high="31"/>
-		</reg32>
+		<reg64 offset="1" name="ADDR" type="address"/>
 		<reg32 offset="3" name="3">
 			<!-- localsize is value minus one: -->
 			<bitfield name="LOCALSIZEX" low="2" high="11" type="uint"/>
@@ -1771,40 +1783,88 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)

 <domain name="CP_SET_MARKER" width="32" varset="chip" prefix="chip" variants="A6XX-">
 	<doc>Tell CP the current operation mode, indicates save and restore procedure</doc>
+	<enum name="set_marker_mode">
+		<value value="0" name="SET_RENDER_MODE"/>
+		<!-- IFPC - inter-frame power collapse -->
+		<value value="1" name="SET_IFPC_MODE"/>
+	</enum>
+	<enum name="a6xx_ifpc_mode">
+		<value value="0" name="IFPC_ENABLE"/>
+		<value value="1" name="IFPC_DISABLE"/>
+	</enum>
 	<enum name="a6xx_marker">
-		<value value="1" name="RM6_BYPASS"/>
-		<value value="2" name="RM6_BINNING"/>
-		<value value="4" name="RM6_GMEM"/>
-		<value value="5" name="RM6_ENDVIS"/>
-		<value value="6" name="RM6_RESOLVE"/>
-		<value value="7" name="RM6_YIELD"/>
+		<value value="1" name="RM6_DIRECT_RENDER"/>
+		<value value="2" name="RM6_BIN_VISIBILITY"/>
+		<value value="3" name="RM6_BIN_DIRECT"/>
+		<value value="4" name="RM6_BIN_RENDER_START"/>
+		<value value="5" name="RM6_BIN_END_OF_DRAWS"/>
+		<value value="6" name="RM6_BIN_RESOLVE"/>
+		<value value="7" name="RM6_BIN_RENDER_END"/>
 		<value value="8" name="RM6_COMPUTE"/>
-		<value value="0xc" name="RM6_BLIT2DSCALE"/>  <!-- no-op (at least on current sqe fw) -->
+		<value value="12" name="RM6_BLIT2DSCALE"/>  <!-- no-op (at least on current sqe fw) -->

 		<!--
 			These values come from a6xx_set_marker() in the
 			downstream kernel, and they can only be set by the kernel
 		-->
-		<value value="0xd" name="RM6_IB1LIST_START"/>
-		<value value="0xe" name="RM6_IB1LIST_END"/>
-		<!-- IFPC - inter-frame power collapse -->
-		<value value="0x100" name="RM6_IFPC_ENABLE"/>
-		<value value="0x101" name="RM6_IFPC_DISABLE"/>
+		<value value="13" name="RM6_IB1LIST_START"/>
+		<value value="14" name="RM6_IB1LIST_END"/>
+		<value value="15" name="RM7_BIN_VISIBILITY_END"/>
+
+		<!-- new in a8xx: -->
+		<value value="32" name="RM8_DEPTH_PASS_START"/>
+		<value value="33" name="RM8_DEPTH_PASS_END"/>
 	</enum>
-	<reg32 offset="0" name="0">
-		<!--
-			NOTE: blob driver and some versions of freedreno/turnip set
-			b4, which is unused (at least by current sqe fw), but interferes
-			with parsing if we extend the size of the bitfield to include
-			b8 (only sent by kernel mode driver).  Really, the way the
-			parsing works in the firmware, only b0-b3 are considered, but
-			if b8 is set, the low bits are interpreted differently.  To
-			model this, without getting confused by spurious b4, this is
-			described as two overlapping bitfields:
-		 -->
-		<bitfield name="MODE" low="0" high="8" type="a6xx_marker"/>
-		<bitfield name="MARKER" low="0" high="3" type="a6xx_marker"/>
-	</reg32>
+	<stripe varset="chip" variants="A6XX-A7XX">
+		<reg32 offset="0" name="0">
+			<!-- if b8 is set, the low bits are interpreted differently (and b4 ignored) -->
+			<bitfield name="MARKER_MODE" pos="8" type="set_marker_mode" addvariant="yes"/>
+
+
+			<bitfield name="MODE" low="0" high="3" type="a6xx_marker" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
+			<!-- used by preemption to determine if GMEM needs to be saved or not -->
+			<bitfield name="USES_GMEM" pos="4" type="boolean" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
+
+
+			<bitfield name="IFPC_MODE" pos="0" type="a6xx_ifpc_mode" varset="set_marker_mode" variants="SET_IFPC_MODE"/>
+
+
+			<!--
+				CP_SET_MARKER is used with these bits to create a
+				critical section around a workaround for ray tracing.
+				The workaround happens after BVH building, and appears
+				to invalidate the RTU's BVH node cache. It makes sure
+				that only one of BR/BV/LPAC is executing the
+				workaround at a time, and no draws using RT on BV/LPAC
+				are executing while the workaround is executed on BR (or
+				vice versa, that no draws on BV/BR using RT are executed
+				while the workaround executes on LPAC), by
+				hooking subsequent CP_EVENT_WRITE/CP_DRAW_*/CP_EXEC_CS.
+				The blob usage is:
+
+
+				CP_SET_MARKER(RT_WA_START)
+				... workaround here ...
+				CP_SET_MARKER(RT_WA_END)
+				...
+				CP_SET_MARKER(SHADER_USES_RT)
+				CP_DRAW_INDX(...) or CP_EXEC_CS(...)
+			-->
+			<bitfield name="SHADER_USES_RT" pos="9" type="boolean" variants="A7XX-"/>
+			<bitfield name="RT_WA_START" pos="10" type="boolean" variants="A7XX-"/>
+			<bitfield name="RT_WA_END" pos="11" type="boolean" variants="A7XX-"/>
+		</reg32>
+	</stripe>
+	<stripe varset="chip" variants="A8XX-">
+		<reg32 offset="0" name="0">
+			<!-- if b8 is set, the low bits are interpreted differently (and b4 ignored) -->
+			<bitfield name="MARKER_MODE" pos="8" type="set_marker_mode" addvariant="yes"/>
+			<bitfield name="USES_GMEM" pos="7" type="boolean" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
+			<bitfield name="MODE" low="0" high="6" type="a6xx_marker" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
+			<bitfield name="IFPC_MODE" pos="0" type="a6xx_ifpc_mode" varset="set_marker_mode" variants="SET_IFPC_MODE"/>
+			<!-- idk if the RT w/a fields apply to a8xx as well -->
+		</reg32>
+	</stripe>
 </domain>

 <domain name="CP_SET_PSEUDO_REG" width="32" varset="chip" prefix="chip" variants="A6XX-">
@@ -1830,9 +1890,9 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 			If concurrent binning is disabled then BR also does binning so it will also
 			write the "real" registers in BR.
 		-->
-		<value value="8" name="DRAW_STRM_ADDRESS"/>
-		<value value="9" name="DRAW_STRM_SIZE_ADDRESS"/>
-		<value value="10" name="PRIM_STRM_ADDRESS"/>
+		<value value="8" name="VSC_PIPE_DATA_DRAW_BASE"/>
+		<value value="9" name="VSC_SIZE_BASE"/>
+		<value value="10" name="VSC_PIPE_DATA_PRIM_BASE"/>
 		<value value="11" name="UNK_STRM_ADDRESS"/>
 		<value value="12" name="UNK_STRM_SIZE_ADDRESS"/>

@@ -1933,11 +1993,11 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 			a bitmask of which modes pass the test.
 		-->

-		<!-- RM6_BINNING -->
+		<!-- RM6_BIN_VISIBILITY -->
 		<bitfield name="BINNING" pos="25" variants="RENDER_MODE" type="boolean"/>
 		<!-- all others -->
 		<bitfield name="GMEM" pos="26" variants="RENDER_MODE" type="boolean"/>
-		<!-- RM6_BYPASS -->
+		<!-- RM6_DIRECT_RENDER -->
 		<bitfield name="SYSMEM" pos="27" variants="RENDER_MODE" type="boolean"/>

 		<bitfield name="BV" pos="25" variants="THREAD_MODE" type="boolean"/>
@@ -2010,54 +2070,45 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>
 </domain>

-<domain name="CP_SET_CTXSWITCH_IB" width="32">
+<domain name="CP_SET_AMBLE" width="32">
 	<doc>
-                Used by the userspace driver to set various IB's which are
-                executed during context save/restore for handling
-                state that isn't restored by the
-                context switch routine itself.
+		Used by the userspace and kernel drivers to set various IB's
+		which are executed during context save/restore for handling
+		state that isn't restored by the context switch routine itself.
 	</doc>
-	<enum name="ctxswitch_ib">
-		<value name="RESTORE_IB" value="0">
+	<enum name="amble_type">
+		<value name="PREAMBLE_AMBLE_TYPE" value="0">
 			<doc>Executed unconditionally when switching back to the context.</doc>
 		</value>
-		<value name="YIELD_RESTORE_IB" value="1">
+		<value name="BIN_PREAMBLE_AMBLE_TYPE" value="1">
                        <doc>
 				Executed when switching back after switching
 				away during execution of
-				a CP_SET_MARKER packet with RM6_YIELD as the
-				payload *and* the normal save routine was
-				bypassed for a shorter one. I think this is
-				connected to the "skipsaverestore" bit set by
-				the kernel when preempting.
+				a CP_SET_MARKER packet with RM6_BIN_RENDER_END as the
+				payload *and* skipsaverestore is set. This is
+				expected to restore static register values not
+				saved when skipsaverestore is set.
 			</doc>
 		</value>
-		<value name="SAVE_IB" value="2">
+		<value name="POSTAMBLE_AMBLE_TYPE" value="2">
                        <doc>
 				Executed when switching away from the context,
 				except for context switches initiated via
 				CP_YIELD.
                        </doc>
 		</value>
-		<value name="RB_SAVE_IB" value="3">
+		<value name="KMD_AMBLE_TYPE" value="3">
 			<doc>
 				This can only be set by the RB (i.e. the kernel)
 				and executes with protected mode off, but
-				is otherwise similar to SAVE_IB.
-
-				Note, kgsl calls this CP_KMD_AMBLE_TYPE
+				is otherwise similar to POSTAMBLE_AMBLE_TYPE.
 			</doc>
 		</value>
 	</enum>
-	<reg32 offset="0" name="0">
-		<bitfield name="ADDR_LO" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="ADDR_HI" low="0" high="31"/>
-	</reg32>
+	<reg64 offset="0" name="ADDR" type="address"/>
 	<reg32 offset="2" name="2">
 		<bitfield name="DWORDS" low="0" high="19" type="uint"/>
-		<bitfield name="TYPE" low="20" high="21" type="ctxswitch_ib"/>
+		<bitfield name="TYPE" low="20" high="21" type="amble_type"/>
 	</reg32>
 </domain>

@@ -2089,12 +2140,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<value name="UNK_EVENT_WRITE" value="0x4"/>
 		<doc>
 			Tracks GRAS_LRZ_CNTL::GREATER, GRAS_LRZ_CNTL::DIR, and
-			GRAS_LRZ_DEPTH_VIEW with previous values, and if one of
+			GRAS_LRZ_VIEW_INFO with previous values, and if one of
 			the following is true:
 			- GRAS_LRZ_CNTL::GREATER has changed
 			- GRAS_LRZ_CNTL::DIR has changed, the old value is not
 			  CUR_DIR_GE, and the new value is not CUR_DIR_DISABLED
-			- GRAS_LRZ_DEPTH_VIEW has changed
+			- GRAS_LRZ_VIEW_INFO has changed
 			then it does a LRZ_FLUSH with GRAS_LRZ_CNTL::ENABLE
 			forced to 1.
 			Only exists in a650_sqe.fw.
@@ -2209,7 +2260,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)

 <domain name="CP_MEM_TO_SCRATCH_MEM" width="32">
 	<doc>
-		Best guess is that it is a faster way to fetch all the VSC_STATE registers
+		Best guess is that it is a faster way to fetch all the VSC_CHANNEL_VISIBILITY registers
 		and keep them in a local scratch memory instead of fetching every time
 		when skipping IBs.
 	</doc>
@@ -2257,7 +2308,25 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	<reg32 offset="0" name="0">
 		<bitfield name="CLEAR_ON_CHIP_TS" pos="0" type="boolean"/>
 		<bitfield name="CLEAR_RESOURCE_TABLE" pos="1" type="boolean"/>
-		<bitfield name="CLEAR_GLOBAL_LOCAL_TS" pos="2" type="boolean"/>
+		<bitfield name="CLEAR_BV_BR_COUNTER" pos="2" type="boolean"/>
+		<bitfield name="RESET_GLOBAL_LOCAL_TS" pos="3" type="boolean"/>
+	</reg32>
+</domain>
+
+<domain name="CP_SCOPE_CNTL" width="32">
+	<enum name="cp_scope">
+		<value value="0" name="INTERRUPTS"/>
+	</enum>
+	<reg32 offset="0" name="0">
+		<bitfield name="DISABLE_PREEMPTION" pos="0" type="boolean"/>
+		<bitfield low="28" high="31" name="SCOPE" type="cp_scope"/>
+	</reg32>
+</domain>
+
+<domain name="CP_INDIRECT_BUFFER" width="32" varset="chip" prefix="chip" variants="A5XX-">
+	<reg64 offset="0" name="IB_BASE" type="address"/>
+	<reg32 offset="2" name="2">
+		<bitfield name="IB_SIZE" low="0" high="19"/>
 	</reg32>
 </domain>

@@ -97,7 +97,7 @@ def parse_cmd_buf(dat):

        if state_block == SB6_CS_SHADER:
          from extra.disassemblers.adreno import disasm_raw
-          if state_type == ST6_SHADER and IOCTL > 2:
+          if state_type == ST6_SHADER and IOCTL > 3:
            disasm_raw(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
          if state_type == ST6_CONSTANTS:
            x = get_mem(((vals[2] << 32) | vals[1]), num_unit*4)
@@ -106,25 +106,30 @@ def parse_cmd_buf(dat):
              print('constants')
              hexdump(x)
          if state_type == ST6_IBO:
-            ibos_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 16 * 4)
+            if state_src == 0x1:
+              ibos_bytes = get_mem(CAPTURED_STATE['bindless_base'] + ((vals[2] << 32) | vals[1]) * 4, num_unit * 64)
+            else: ibos_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 16 * 4)
            CAPTURED_STATE['ibos'] = ibos_bytes[:]
            if IOCTL > 1:
              print('texture ibos')
              hexdump(ibos_bytes)
        elif state_block == SB6_CS_TEX:
          if state_type == ST6_SHADER:
-            samplers_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 4 * 4)
+            if state_src == 0x1:
+              samplers_bytes = get_mem(CAPTURED_STATE['bindless_base'] + ((vals[2] << 32) | vals[1]) * 4, num_unit * 64)
+            else: samplers_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 4 * 4)
            CAPTURED_STATE['samplers'] = samplers_bytes[:]
            if IOCTL > 1:
              print('texture samplers')
              hexdump(samplers_bytes)
          if state_type == ST6_CONSTANTS:
-            descriptors_bytes = get_mem((vals[2] << 32) | vals[1], 1600)
+            if state_src == 0x1:
+              descriptors_bytes = get_mem(CAPTURED_STATE['bindless_base'] + ((vals[2] << 32) | vals[1]) * 4, num_unit * 64)
+            else: descriptors_bytes = get_mem((vals[2] << 32) | vals[1], 1600)
            CAPTURED_STATE['descriptors'] = descriptors_bytes[:]
            if IOCTL > 1:
              print('texture descriptors')
              hexdump(descriptors_bytes)
-
      elif ops[opcode] == "CP_REG_TO_MEM":
        reg, cnt, b64, accum = vals[0] & 0x3FFFF, (vals[0] >> 18) & 0xFFF, (vals[0] >> 30) & 0x1, (vals[0] >> 31) & 0x1
        dest = vals[1] | (vals[2] << 32)
@@ -152,6 +157,10 @@ def parse_cmd_buf(dat):
        if IOCTL > 0:
          print(f'THREADSIZE-{(vals[0] >> 20)&0x1}\nEARLYPREAMBLE-{(vals[0] >> 23) & 0x1}\nMERGEDREGS-{(vals[0] >> 3) & 0x1}\nTHREADMODE-{vals[0] & 0x1}\nHALFREGFOOTPRINT-{(vals[0] >> 1) & 0x3f}\nFULLREGFOOTPRINT-{(vals[0] >> 7) & 0x3f}\nBRANCHSTACK-{(vals[0] >> 14) & 0x3f}\n')
          print(f'SP_CS_UNKNOWN_A9B1-{vals[1]}\nSP_CS_BRANCH_COND-{vals[2]}\nSP_CS_OBJ_FIRST_EXEC_OFFSET-{vals[3]}\nSP_CS_OBJ_START-{vals[4] | (vals[5] << 32)}\nSP_CS_PVT_MEM_PARAM-{vals[6]}\nSP_CS_PVT_MEM_ADDR-{vals[7] | (vals[8] << 32)}\nSP_CS_PVT_MEM_SIZE-{vals[9]}')
+      if offset == 0xa9e8:
+        CAPTURED_STATE['bindless_base'] = (vals[0] | (vals[1] << 32)) & ~0b11
+        # print(hex(CAPTURED_STATE['bindless_base']))
+        # hexdump(get_mem(CAPTURED_STATE['bindless_base'], 0x200))
      if offset == 0xb180:
        if IOCTL > 0:
          print('border color offset', hex(vals[1] << 32 | vals[0]))
@@ -171,8 +180,8 @@ def ioctl(fd, request, argp):
    name, stype = nrs[nr]
    s = get_struct(argp, stype)
    if IOCTL > 0: print(f"{ret:2d} = {name:40s}", ' '.join(format_struct(s)))
-    if name == "IOCTL_KGSL_GPUOBJ_INFO": pass
-      # mmaped[s.gpuaddr] = mmap.mmap(fd, s.size, offset=s.id*0x1000)
+    if name == "IOCTL_KGSL_GPUOBJ_INFO":
+      mmaped[s.gpuaddr] = mmap.mmap(fd, s.size, offset=s.id*0x1000)
    if name == "IOCTL_KGSL_GPU_COMMAND":
      for i in range(s.numcmds):
        cmd = get_struct(s.cmdlist+ctypes.sizeof(msm_kgsl.struct_kgsl_command_object)*i, msm_kgsl.struct_kgsl_command_object)
@@ -321,7 +321,7 @@
 				CLANG_ENABLE_MODULES = YES;
 				CODE_SIGN_ENTITLEMENTS = macOS/macOS.entitlements;
 				CODE_SIGN_IDENTITY = "-";
-				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "Apple Development";
+				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "-";
 				CODE_SIGN_STYLE = Automatic;
 				COMBINE_HIDPI_IMAGES = YES;
 				CURRENT_PROJECT_VERSION = 1;
@@ -357,7 +357,7 @@
 				CLANG_ENABLE_MODULES = YES;
 				CODE_SIGN_ENTITLEMENTS = macOS/macOS.entitlements;
 				CODE_SIGN_IDENTITY = "-";
-				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "Apple Development";
+				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "-";
 				CODE_SIGN_STYLE = Automatic;
 				COMBINE_HIDPI_IMAGES = YES;
 				CURRENT_PROJECT_VERSION = 1;
@@ -502,7 +502,7 @@
 			buildSettings = {
 				AD_HOC_CODE_SIGNING_ALLOWED = YES;
 				CODE_SIGN_ENTITLEMENTS = TinyGPUDriverExtension/TinyGPUDriver.entitlements;
-				CODE_SIGN_IDENTITY = "Apple Development";
+				CODE_SIGN_IDENTITY = "-";
 				CODE_SIGN_STYLE = Automatic;
 				CURRENT_PROJECT_VERSION = 1;
 				DEVELOPMENT_TEAM = 9YG3G8543N;
@@ -530,7 +530,7 @@
 			buildSettings = {
 				AD_HOC_CODE_SIGNING_ALLOWED = YES;
 				CODE_SIGN_ENTITLEMENTS = TinyGPUDriverExtension/TinyGPUDriver.entitlements;
-				CODE_SIGN_IDENTITY = "Apple Development";
+				CODE_SIGN_IDENTITY = "-";
 				CODE_SIGN_STYLE = Automatic;
 				CURRENT_PROJECT_VERSION = 1;
 				DEVELOPMENT_TEAM = 9YG3G8543N;
@@ -282,11 +282,11 @@ class TestTrainingOnnxOps(TestOnnxOps):
    tiny_out = runner(inps)
    onnx_out = onnx_fxn(**inps, **opts)
    for (nm, t_out), o_out in  zip(tiny_out.items(), onnx_out):
-      np.testing.assert_allclose(t_out.numpy(), o_out, rtol=1e-3, atol=1e-6, err_msg=f"{nm} failed")
+      np.testing.assert_allclose(t_out.numpy(), o_out, rtol=1e-6, atol=1e-6, err_msg=f"{nm} failed")

-  def test_adagrad_t_greater_than_zero(self):
+  def test_adagrad_t(self):
    from onnx.backend.test.case.node.adagrad import apply_adagrad
-    for t in [1, 3, 100]:
+    for t in [0, 1, 3, 100]:
      inputs = {
        "r": np.array(0.01, dtype=np.float32),
        "t": np.array(t, dtype=np.int32),
@@ -298,10 +298,10 @@ class TestTrainingOnnxOps(TestOnnxOps):
      outputs = ["X_out", "H_out"]
      self._validate_training("Adagrad", apply_adagrad, inputs, attributes, outputs)

-  def test_momentum_t_greater_than_zero(self):
+  def test_momentum(self):
    from onnx.backend.test.case.node.momentum import apply_momentum, apply_nesterov
    for onnx_fxn, mode in ((apply_momentum, "standard"), (apply_nesterov, "nesterov")):
-      for t in [1, 3, 100]:
+      for t in [0, 1, 3, 100]:
        inputs = {
          "r": np.array(0.01, dtype=np.float32),
          "t": np.array(t, dtype=np.int32),
@@ -313,9 +313,9 @@ class TestTrainingOnnxOps(TestOnnxOps):
        outputs = ["X_out", "V_out"]
        self._validate_training("Momentum", onnx_fxn, inputs, attributes, outputs)

-  def test_adam_t_greater_than_zero(self):
+  def test_adam(self):
    from onnx.backend.test.case.node.adam import apply_adam
-    for t in [1, 3, 100]:
+    for t in [0, 1, 3, 100]:
      inputs = {
        "r": np.array(0.01, dtype=np.float32),
        "t": np.array(t, dtype=np.int32),
@@ -422,6 +422,7 @@ class TestContribOnnxOps(TestOnnxOps):
        outputs = ["C"]
        self.helper_test_single_op("QLinearAdd", inputs, attributes, outputs, atol=1) # TODO: look into why this is inaccurate

+  def test_qlinear_add_round_half_to_even(self):
    with self.subTest(test_case="round_half_to_even"):
      inputs = {
        "A": np.array([1, 1, 1, 1], dtype=np.int8),
@@ -435,7 +436,7 @@ class TestContribOnnxOps(TestOnnxOps):
      }
      attributes = {}
      outputs = ["C"]
-      self.helper_test_single_op("QLinearAdd", inputs, attributes, outputs)
+      self.helper_test_single_op("QLinearAdd", inputs, attributes, outputs, atol=1) # TODO: look into why this is inaccurate

  def test_qlinear_mul(self):
    for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
@@ -24,7 +24,6 @@ class TestUnaryOpsConstFolding(unittest.TestCase):
    _check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
    _check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))

-  @unittest.expectedFailure  # no two level fold
  def test_neg_folding(self):
    _check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
    _check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
@@ -83,10 +82,8 @@ class TestBinaryOpsConstFolding(unittest.TestCase):
  def test_div_tensor_one(self):
    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / Tensor.ones(4))

-  @unittest.expectedFailure  # TODO: fix
  def test_idiv_literal_one(self):
    _check_ast_count(0, Tensor([1, 2, 3, 4]) // 1)
-  @unittest.expectedFailure  # TODO: fix
  def test_idiv_tensor_one(self):
    _check_ast_count(0, Tensor([1, 2, 3, 4]) // Tensor.ones(4, dtype=dtypes.int32))

@@ -92,7 +92,9 @@ class TestProfiler(unittest.TestCase):
    #  assert evs[i].st > evs[i-1].en, "timestamp not aranged"

  def test_profile_multidev(self):
-    d1 = Device[f"{Device.DEFAULT}:1"]
+    try: d1 = Device[f"{Device.DEFAULT}:1"]
+    except Exception as e: self.skipTest(f"second device not available {e}")
+
    buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
    buf2 = Buffer(f"{Device.DEFAULT}:1", 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()

@@ -109,7 +111,8 @@ class TestProfiler(unittest.TestCase):
      assert evs[0].is_copy, "kernel should be copy"

  def test_profile_multidev_transfer(self):
-    d1 = Device[f"{Device.DEFAULT}:1"]
+    try: d1 = Device[f"{Device.DEFAULT}:1"]
+    except Exception as e: self.skipTest(f"second device not available {e}")

    buf1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:0").realize()
    with helper_collect_profile(TestProfiler.d0, d1) as profile:
@@ -122,7 +125,8 @@ class TestProfiler(unittest.TestCase):

  @unittest.skipIf(Device.DEFAULT in "METAL" or (MOCKGPU and Device.DEFAULT == "AMD"), "AMD mockgpu does not support queue wait interrupts")
  def test_profile_graph(self):
-    d1 = Device[f"{Device.DEFAULT}:1"]
+    try: d1 = Device[f"{Device.DEFAULT}:1"]
+    except Exception as e: self.skipTest(f"second device not available {e}")

    def f(a):
      x = (a + 1).realize()
@@ -145,7 +149,9 @@ class TestProfiler(unittest.TestCase):
  @unittest.skipIf(CI or not issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "skip CI")
  def test_dev_jitter_matrix(self):
    dev_cnt = 6
-    devs = [Device[f"{Device.DEFAULT}:{i}"] for i in range(dev_cnt)]
+    try: devs = [Device[f"{Device.DEFAULT}:{i}"] for i in range(dev_cnt)]
+    except Exception as e: self.skipTest(f"multiple devices not available {e}")
+
    for dev in devs: dev.synchronize()
    for dev in devs: dev._at_profile_finalize()

@@ -219,7 +225,7 @@ class TestProfiler(unittest.TestCase):
    exec_points = [e for e in profile if isinstance(e, ProfilePointEvent) and e.name == "exec"]
    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent) and not e.is_copy]
    self.assertEqual(len(exec_points), len(range_events), 2)
-    self.assertEqual(len(dedup(e.key for e in exec_points)), 1)
+    self.assertEqual(len(dedup(e.arg['name'] for e in exec_points)), 1)
    self.assertEqual(len(dedup(e.arg['metadata'] for e in exec_points)), 1)

 if __name__ == "__main__":
@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, nn
-from tinygrad.helpers import Context, GlobalCounters, CI
+from tinygrad.helpers import Context, GlobalCounters, CI, CPU_LVP, getenv
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops

 class TestRangeifyAssign(unittest.TestCase):
@@ -28,6 +28,73 @@ class TestRangeifyEdgeCase(unittest.TestCase):
    res = Tensor.cat(a, c, dim=0)
    self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)

+if getenv("BIG") > 2:
+  # llama 8B (8192)
+  BS, HEADS, SEQLEN, EMB = 4, 32, 8192, 128
+elif getenv("BIG") > 1:
+  # llama 8B
+  BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
+elif getenv("BIG") > 0:
+  # bigger
+  BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
+else:
+  BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
+
+@unittest.skipIf(CPU_LVP, "broken in LVP")
+class TestPcontig(unittest.TestCase):
+  def test_flash_attention_bw(self):
+    def fa_bw():
+      Tensor.manual_seed(1337)
+      with Context(DEBUG=0):
+        q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize().requires_grad_() for _ in range(3)]
+        attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
+        attn_output.weight.requires_grad_().realize()
+        target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()
+
+      GlobalCounters.reset()
+      attn = q.scaled_dot_product_attention(k, v).contiguous().contiguous_backward()
+      attn = attn.transpose(1, 2).reshape(BS, SEQLEN, -1)
+      out = attn_output(attn)
+      loss = (out - target).square().mean()
+      loss.backward()
+      #ret = [out, Tensor.stack(q.grad, k.grad, v.grad)]
+      ret = [out, q.grad, k.grad, v.grad]
+      Tensor.realize(*ret)
+      return ret
+
+    with Context(PCONTIG=2, REAL_SUBSTITUTE=1, DEBUG=2):
+      grads = fa_bw()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+
+    with Context(DEBUG=2):
+      cmp_grads = fa_bw()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+
+    with Context(DEBUG=0):
+      mses = [((x-y)**2).sum().item() for x,y in zip(grads, cmp_grads)]
+    mse = sum(mses)
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-6)
+
+  def test_flash_attention(self):
+    def fa():
+      Tensor.manual_seed(1337)
+      with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
+      GlobalCounters.reset()
+      return q.scaled_dot_product_attention(k, v).realize()
+
+    with Context(PCONTIG=2, DEBUG=2):
+      ret = fa()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+    with Context(DEBUG=2):
+      cmp = fa()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+    with Context(DEBUG=0):
+      mse = ((cmp-ret)**2).sum().item()
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-6)
+
+
 # *** non CI rangeify tests below this line ***

 N = 256
@@ -215,33 +282,6 @@ class TestRangeify(unittest.TestCase):
    out = blk._feed_forward(x)
    out.realize()

-  @unittest.skip("RANGEIFY=0 does nothing")
-  def test_flash_attention(self):
-    BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
-
-    # bigger
-    #BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
-
-    # llama 8B
-    #BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
-
-    def fa():
-      Tensor.manual_seed(1337)
-      with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
-      return q.scaled_dot_product_attention(k, v).realize()
-
-    with Context(DEBUG=4):
-      GlobalCounters.reset()
-      ret = fa()
-    with Context(RANGEIFY=0):
-      with Context(DEBUG=2):
-        GlobalCounters.reset()
-        cmp = fa()
-      with Context(DEBUG=0):
-        mse = ((cmp-ret)**2).sum().item()
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-6)
-
 # contiguous + reduce can support ranges?

@unittest.skip("pm_rangeify no longer exists. test this in a different way")
@@ -333,7 +333,7 @@ class TestSchedule(unittest.TestCase):
    r1 = (x - r0).sum(axis=0).div(2)
    out0 = r0 + y
    out1 = r1 + y
-    schedule = check_schedule([out0, out1], 4)
+    schedule = check_schedule([out0, out1], 3)
    reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
    self.assertEqual(len(reduceops), 2) # why is RANGEIFY different?

@@ -1274,7 +1274,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters([c1, c2]), nesterov=True, momentum=0.9, weight_decay=0.1)
      opt.zero_grad()
      c2(c1(img).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 15)
+      check_schedule(opt.schedule_step(), 13)

  def test_sgd_4convs_fuse(self):
    with Tensor.train():
@@ -287,7 +287,6 @@ class TestSymbolicOps(unittest.TestCase):
        symbolic = symbolic_result[:].numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)

-  @unittest.expectedFailure
  def test_conv2d_ceildiv_edge_case(self):
    v = Variable('v', 11, 50_000)
    val = 39601
@@ -295,9 +294,10 @@ class TestSymbolicOps(unittest.TestCase):
    weight = Tensor.randn(256, 22, 12)

    result = x.conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))
-    var_val = {v: val}
+    var_val = {v.expr: val}
    shape = tuple(sym_infer(s, var_val) for s in result.shape)
-    self.assertEqual(shape, (1, 256, 6600))  # TODO: fails if ceildiv is incorrect
+    with self.assertRaises(AssertionError):
+      self.assertEqual(shape, (1, 256, 6600))  # TODO: fails if ceildiv is incorrect
    # TODO: test output is correct

 if __name__ == '__main__':
@@ -433,17 +433,15 @@ class TestDiskTensorMovement(unittest.TestCase):
    t = Tensor(self.fn)
    self.assertListEqual(t[16:18].tolist(), [16,17])

-  # TODO: fix this! at least assert on it
-  @unittest.expectedFailure
  def test_slice_read_cat(self):
    t = Tensor(self.fn)
-    self.assertListEqual(Tensor.cat(t[16:18], t[20:22]).tolist(), [16,17,20,21])
+    with self.assertRaises(AssertionError):
+      self.assertListEqual(Tensor.cat(t[16:18], t[20:22]).tolist(), [16,17,20,21])

-  # TODO: fix this! at least assert on it
-  @unittest.expectedFailure
  def test_slice_sum(self):
    t = Tensor(self.fn)
-    self.assertListEqual((t[16:18]+t[20:22]).tolist(), [16+20,17+21])
+    with self.assertRaises(AssertionError):
+      self.assertListEqual((t[16:18]+t[20:22]).tolist(), [16+20,17+21])

 if __name__ == "__main__":
  unittest.main()
@@ -20,8 +20,8 @@ class TestKernelize(unittest.TestCase):
    self.assertEqual(len([s for s in a0.uop.toposort() if s.op is Ops.KERNEL]), 2)
    self.assertIs(a1.uop.base.op, Ops.REDUCE_AXIS)
    # input Tensor and user contiguous kernelize
-    self.assertIs(a0.uop.base.op, Ops.ASSIGN)
-    self.assertIs(a.uop.base.op, Ops.ASSIGN)
+    self.assertIs(a0.uop.base.op, Ops.AFTER)
+    self.assertIs(a.uop.base.op, Ops.AFTER)

  def test_two_reduce_w_add(self):
    a = Tensor.ones(16,16).contiguous()
@@ -31,7 +31,7 @@ class TestKernelize(unittest.TestCase):
    # NOTE: the +1 is fused with a1, so a1 is not kernelized
    self.assertIs(a1.uop.base.op, Ops.REDUCE_AXIS)
    # the input to the REDUCE_AXIS is an ASSIGN though
-    self.assertIs(a1.uop.base.src[0].base.op, Ops.ASSIGN)
+    self.assertIs(a1.uop.base.src[0].base.op, Ops.AFTER)

 if __name__ == '__main__':
  unittest.main()
@@ -2,7 +2,7 @@ import unittest, decimal, json, struct
 from dataclasses import dataclass
 from typing import Generator

-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, TRACK_MATCH_STATS
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, TRACK_MATCH_STATS, profile_matches
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
@@ -117,6 +117,28 @@ class TestViz(BaseTestViz):
    # NOTE: names from TracingKey do not get deduped
    self.assertEqual(lst[0]["name"], "custom_name")

+  def test_profile_matches(self):
+    @profile_matches
+    def nested_function(u:UOp):
+      for i in range(2): graph_rewrite(u, PatternMatcher([]), name=f"step {i+1}")
+
+    @track_rewrites()
+    def main_rewrite(u:UOp):
+      graph_rewrite(u, PatternMatcher([]), name="init")
+      nested_function(u)
+
+    main_rewrite(UOp.variable("a", 1, 10)+UOp.variable("b", 1, 10))
+    steps = get_viz_list()[0]["steps"]
+    self.assertEqual(steps[0]["name"], "init")
+    self.assertEqual(steps[1]["name"], "nested_function")
+    self.assertEqual(len(steps), 4)
+
+  def test_profile_matches_invalid_arg(self):
+    @profile_matches
+    def invalid_fxn(arg:str): return graph_rewrite(UOp(Ops.SINK), PatternMatcher([]))
+    with self.assertRaisesRegex(AssertionError, "invalid match tracing input"):
+      invalid_fxn("test")
+
  def test_colored_label(self):
    # NOTE: dataclass repr prints literal escape codes instead of unicode chars
    @dataclass(frozen=True)
@@ -127,7 +149,7 @@ class TestViz(BaseTestViz):
    self.assertEqual(ansistrip(a2["label"]), f"CUSTOM\n{TestStruct.__qualname__}(colored_field='xyz12345')")

  def test_inf_loop(self):
-    a = UOp.variable('a', 0, 10)
+    a = UOp.variable('a', 0, 10, dtype=dtypes.int)
    b = a.replace(op=Ops.CONST)
    pm = PatternMatcher([
      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
@@ -142,8 +164,8 @@ class TestViz(BaseTestViz):
    self.assertEqual(graphs[2], uop_to_json(nop)[id(nop)])

  def test_const_node_visibility(self):
-    a = UOp.variable("a", 0, 10)
-    z = UOp.const(dtypes.index, 0)
+    a = UOp.variable("a", 0, 10, dtype=dtypes.int)
+    z = UOp.const(a.dtype, 0)
    alu = a*z
    exec_rewrite(alu, [sym])
    lst = get_viz_list()
@@ -326,14 +348,14 @@ def load_profile(lst:list[ProfileEvent]) -> dict:
    event_type, event_count = u("<BI")
    if event_type == 0:
      for _ in range(event_count):
-        name, ref, st, dur, _ = u("<IIIfI")
-        v["events"].append({"name":strings[name], "ref":option(ref), "st":st, "dur":dur})
+        name, ref, key, st, dur, _ = u("<IIIIfI")
+        v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur})
    else:
      v["peak"] = u("<Q")[0]
      for _ in range(event_count):
        alloc, ts, key = u("<BII")
        if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-        else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg":{"users":[u("<I")[0] for _ in range(u("<I")[0])]}})
+        else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIBB") for _ in range(u("<I")[0])]}})
  return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}

 class TestVizProfiler(unittest.TestCase):
@@ -498,6 +520,27 @@ class TestVizMemoryLayout(BaseTestViz):
    buffers = profile["layout"]["NULL Memory"]["events"]
    user_cnt = [len(b["arg"]["users"]) for b in buffers if b["arg"].get("users")]
    self.assertEqual(max(user_cnt), n)
+    input_buf = buffers.pop()
+    assert all(u[3] == 0 for u in input_buf["arg"]["users"])
+
+  def test_annotate_read_write(self):
+    a = Tensor.ones(4, device="NULL").contiguous().realize()
+    b = a.assign(a+2)
+    c = a+1
+    Tensor.realize(b, c)
+    buf_events = load_profile(cpu_events+Buffer.profile_events)["layout"]["NULL Memory"]["events"]
+    users = next((b["arg"]["users"] for b in buf_events if len(b["arg"].get("users",[])) == 3))
+    self.assertEqual(users[0][3], 1) # write Tensor.ones
+    self.assertEqual(users[1][3], 2) # read+write Tensor.assign
+    self.assertEqual(users[2][3], 0) # readonly
+
+  def test_dedup_users(self):
+    a = Tensor.empty(1, device="NULL")
+    for _ in range(n:=4): a.add(1).realize()
+    profile = load_profile(cpu_events+Buffer.profile_events)
+    programs = profile["layout"][a.device]["events"]
+    users = profile["layout"][f"{a.device} Memory"]["events"].pop()["arg"]["users"]
+    self.assertEqual(len(programs), len(set(users)), n)

 if __name__ == "__main__":
  unittest.main()
@@ -123,7 +123,7 @@ def gep_on_store(gep:UOp, st:UOp, sto:UOp):
  return gep.src[0].store(st.gep(new_arg), *sto.src[2:])

 load_store_folding = PatternMatcher([
-  (UPat(Ops.INDEX, src=(UPat(Ops.VECTORIZE, src=UPat(GroupOp.Defines, name="buf")), UPat.var("vec"))), expand_index),
+  (UPat(Ops.INDEX, src=(UPat(Ops.VECTORIZE, src=UPat(GroupOp.Defines).or_after(name="buf")), UPat.var("vec"))), expand_index),
  # GEP after LOAD
  (UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True),
   lambda gep, ld: ld.replace(dtype=ld.dtype.scalar().vec(gep.dtype.count), src=(gep.src[0],)+ld.src[1:]).gep(gep.arg)),
@@ -242,11 +242,13 @@ def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp):
  return buf.broadcast(cnt).index(idx.broadcast(cnt)*cnt+UOp.const(dtypes.index.vec(cnt), tuple(range(cnt))))

 devectorize = PatternMatcher([
+  # CAST after AFTER
+  (UPat(Ops.CAST, name="c").f(Ops.AFTER, allow_any_len=True, name="a"), lambda c,a: c.src[0].after(*a.src[1:]).cast(c.dtype)),
  # no ALU on vectorized dtypes
  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name="alu"), no_vectorized_alu),
  (UPat(Ops.WMMA, name="wmma"), no_vectorized_wmma),
  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf"), no_vectorized_buf),
-  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf").cast(name="cast").index(UPat.var("idx")), no_vectorized_index),
+  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)).or_after(name="buf").cast(name="cast").index(UPat.var("idx")), no_vectorized_index),
 ])

 pm_render = PatternMatcher([
@@ -296,12 +298,14 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
    stored_ranges = flatten([x.src[2:] for x in topo if x.op is Ops.STORE])
    input_ranges = tuple([x for x in topo if x.op is Ops.RANGE and x not in reduce_range and x not in stored_ranges])
    identity = red.const(red.dtype, identity_element(red.arg, red.dtype.scalar()))
-    acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), arg=(ctx.acc_num,)).index(UOp.const(dtypes.int, 0))
-    do_store = acc.store(identity, UOp(Ops.NOOP, src=input_ranges)) if len(input_ranges) else acc.store(identity)
-    lst = [acc.load(do_store, *reduce_range)] + lst  # put acc as the first element
+    acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), arg=(ctx.acc_num,))
+    acc_init = acc.after(*input_ranges).index(UOp.const(dtypes.int, 0)).store(identity) if len(input_ranges) else \
+               acc.index(UOp.const(dtypes.int, 0)).store(identity)
+    lst = [acc.after(acc_init, *reduce_range).index(UOp.const(dtypes.int, 0)).load()] + lst  # put acc as the first element
    ctx.acc_num += 1
  ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst)
-  return acc.load(acc.store(ret, *reduce_range)) if len(reduce_range) != 0 else ret
+  if len(reduce_range) == 0: return ret
+  return acc.after(acc.index(UOp.const(dtypes.int, 0)).store(ret, *reduce_range)).index(UOp.const(dtypes.int, 0)).load()

 pm_reduce = PatternMatcher([
  # REDUCE -> DEFINE_ACC+ASSIGN
@@ -5,7 +5,7 @@ from typing import cast, Final
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp
 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, merge_dicts, get_single_element
+from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
 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
@@ -88,7 +88,14 @@ class Scheduler:

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

-  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 colors(self) -> list[str]:
+    store_rngs = flatten([x.src[2:] for x in self.ast.src])
+    ret = []
+    for x,r in zip(self.axis_types, self.rngs):
+      if self.dont_use_locals and x == AxisType.GLOBAL: ret.append("BLUE")
+      elif r not in store_rngs and x == AxisType.LOOP: ret.append("BLACK")
+      else: ret.append(axis_colors[x])
+    return ret
  def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():>4s}', color) for x,color in zip(self.rngs, self.colors())])

  def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False, input_new_rng=None):
@@ -1,5 +1,5 @@
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, ImageDType
-from tinygrad.uop.symbolic import symbolic_flat, sym
+from tinygrad.uop.symbolic import symbolic_flat
 from tinygrad.helpers import partition
 from tinygrad.dtype import dtypes

@@ -69,18 +69,24 @@ pm_split_ranges = PatternMatcher([

 def no_range(u:UOp) -> bool: return not any(x.op is Ops.RANGE for x in u.backward_slice_with_self)

-def reduce_rangeless(red:UOp):
-  # TODO: share code with reduce_unparented
-  if red.arg not in {Ops.ADD, Ops.MAX}: return None
-  if red.src[0].dtype != red.dtype: return None
-  if not no_range(red.src[0]): return None
-  ret = red.src[0]
+def reduce_unparented(red:UOp):
+  if red.arg not in {Ops.ADD, Ops.MAX, Ops.MUL}: return None
+  assert all(x.op is Ops.RANGE for x in red.src[1:]), "some reduce srcs aren't ranges"
+  reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].ranges)
+  if len(reduce_unparented) == 0: return None
+  ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0]
  if red.arg is Ops.ADD:
-    for r in red.src[1:]:
-      ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
+    for r in reduce_unparented: ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
+  if red.arg is Ops.MUL:
+    for r in reduce_unparented: ret = ret ** r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
  return ret

-pm_reduce_collapse = PatternMatcher([
+pm_reduce_unparented = PatternMatcher([
+  # remove any ranges from a REDUCE that aren't referenced in the reduce source
+  (UPat(Ops.REDUCE, name="red"), reduce_unparented),
+])
+
+pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
  # lift x+y out of reduce on lt
  ((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
@@ -106,9 +112,7 @@ pm_reduce_collapse = PatternMatcher([
  # AND on WHERE
  ((UPat(Ops.DEFINE_VAR, name="x") & UPat.var("y")).where(UPat.cvar("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
    lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)),
-  # remove REDUCEs that no longer have a RANGE in the src
-  (UPat(Ops.REDUCE, name="red"), reduce_rangeless),
-])+sym
+])+symbolic_flat

 def reduce_collapse(red:UOp):
  included, not_included = partition(red.backward_slice, lambda x: any(y in x.backward_slice_with_self for y in red.src[1:]))
@@ -122,23 +126,6 @@ def reduce_collapse(red:UOp):
  sink = graph_rewrite(collapse_fxn, pm_reduce_collapse, name="reduce_collapse")
  return sink.substitute({v:k for k,v in replaces.items()}) if no_range(sink) else None

-def reduce_unparented(red:UOp):
-  if red.arg not in {Ops.ADD, Ops.MAX, Ops.MUL}: return None
-  assert all(x.op is Ops.RANGE for x in red.src[1:]), "some reduce srcs aren't ranges"
-  reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].ranges)
-  if len(reduce_unparented) == 0: return None
-  ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0]
-  if red.arg is Ops.ADD:
-    for r in reduce_unparented: ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
-  if red.arg is Ops.MUL:
-    for r in reduce_unparented: ret = ret ** r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
-  return ret
-
-pm_reduce_unparented = PatternMatcher([
-  # remove any ranges from a REDUCE that aren't referenced in the reduce source
-  (UPat(Ops.REDUCE, name="red"), reduce_unparented),
-])
-
 pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([
  # remove REDUCE without loads (generic arange opt / indexing). TODO: support multi range
  (UPat(Ops.REDUCE, src=(UPat(), UPat()), name="red"), reduce_collapse),
@@ -166,8 +166,9 @@ class ExecItem:
    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:
-      payload = {"metadata":self.metadata, "var_vals":var_vals, "bufs":[b.trace_num for b in bufs]}
-      cpu_events.append(ProfilePointEvent(self.prg.device, "exec", self.prg.display_name, payload))
+      payload = {"metadata":self.metadata, "var_vals":var_vals, "bufs":[b.trace_num for b in bufs], "name":self.prg.display_name}
+      payload["outputs"], payload["inputs"] = (self.prg.p.outs, self.prg.p.ins) if isinstance(self.prg, CompiledRunner) else ([0], [1])
+      cpu_events.append(ProfilePointEvent(self.prg.device, "exec", len(cpu_events), payload))
    et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2)
    if do_update_stats:
      GlobalCounters.kernel_count += 1
@@ -180,10 +181,11 @@ class ExecItem:
        header_color = 'magenta' if jit else ('green' if self.prg.first_run else None)
        ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
        flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
-        flops_str = f"{flops*1e-9:9.2f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:9.2f} TFLOPS", 'green')
-        mem_str = f"{membw*1e-9:6.1f}|{ldsbw*1e-9:<7.1f} GB/s" if membw < 1e13 else colored(f"{membw*1e-12:6.1f}|{ldsbw*1e-12:<7.1f} TB/s", 'green')
+        flops_str = f"{flops*1e-9:7.0f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:7.0f} TFLOPS", 'green')
+        mem_str = f"{membw*1e-9:4.0f}|{ldsbw*1e-9:<6.0f} GB/s" if membw < 1e13 and ldsbw < 1e15 else \
+          colored(f"{membw*1e-12:4.0f}|{ldsbw*1e-12:<6.0f} TB/s", 'green')
        print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
-          f" {self.prg.display_name+' '*(44-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB"+
+          f" {self.prg.display_name+' '*(46-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
          ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})")+
          f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in self.metadata] if self.metadata else ''}")
      self.prg.first_run = False
@@ -22,18 +22,18 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
  in_degree: dict[UOp, 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
+    if u.op is not Ops.AFTER: continue  # anything that's not an ASSIGN doesn't write a kernel, so we can skip
    k = u.src[1]
    in_degree.setdefault(k, 0)
    for s in k.src:
-      if s.op is Ops.ASSIGN:
+      if s.op is Ops.AFTER:
        children[s.src[1]].append(k)
        in_degree[k] += 1
      elif s.op in {Ops.MSELECT, Ops.MSTACK}:
        for ss in s.src:
          if ss.op is Ops.MSELECT: ss = ss.src[0]
          if ss.op is not Ops.BUFFER:
-            assert ss.op is Ops.ASSIGN, f"ss.op is not ASSIGN, it's {ss.op}"
+            assert ss.op is Ops.AFTER, f"ss.op is not AFTER, it's {ss.op}"
            children[ss.src[1]].append(k)
            in_degree[k] += 1
      elif s.op is Ops.BUFFER:
@@ -43,7 +43,7 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
        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}")
+        raise RuntimeError(f"input to kernel must be AFTER or BUFFER, not {s.op}")

  # linearize KERNEL UOps into ScheduleItems in BFS order

@@ -169,6 +169,8 @@ VIZ = PROFILE = ContextVar("VIZ", 0)
 SPEC = ContextVar("SPEC", 0)
 # TODO: disable by default due to speed
 IGNORE_OOB = ContextVar("IGNORE_OOB", 1)
+PCONTIG = ContextVar("PCONTIG", 0)  # partial contiguous in rangeify
+REAL_SUBSTITUTE = ContextVar("REAL_SUBSTITUTE", 0)

@dataclass(frozen=True)
 class Metadata:
@@ -1242,7 +1242,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    G, V, H = G.detach(), V.detach(), H.detach()
    X.grad = norm_coefficient * X.detach() + G
    opt = TinyAdam([X], b1=alpha, b2=beta, eps=epsilon)
-    opt.m, opt.v, opt.lr = [V], [H], R
+    # NOTE: FUSE_OPTIM can change shapes of m and v
+    opt.m, opt.v, opt.lr = [V.reshape(opt.m[0].shape)], [H.reshape(opt.v[0].shape)], R
    # need no-op for m_hat and v_hat if T == 0
    if T == 0: opt.b1_t, opt.b2_t = opt.b1_t.zeros_like(), opt.b2_t.zeros_like()
    else:
@@ -116,9 +116,7 @@ class LARS(Optimizer):
      # classic momentum does post learning rate update
      if self.classic: g = g * r * self.lr
      if self.momentum:
-        # TODO: this contiguous is required for correctness because self.b[i] becomes a non contiguous view
-        # the scheduler should detect this and just insert contiguous
-        self.b[i].assign(self.momentum * self.b[i].contiguous() + g)  # NOTE: self.b[i] is zero on the first run, no if required
+        self.b[i].assign(self.momentum * self.b[i] + g)  # NOTE: self.b[i] is zero on the first run, no if required
        g = (g + self.momentum * self.b[i]) if self.nesterov else self.b[i]
      if self.ns_params: g = g.reshape(g.shape[0], -1).newton_schulz(self.ns_steps, self.ns_params).reshape(g.shape)
      # muon does post momentum weight decay
@@ -1,7 +1,7 @@
 from typing import Literal, Callable, cast
 import os, math, sys
 from collections import defaultdict, Counter
-from tinygrad.codegen.opt import tc
+from tinygrad.codegen.opt import tc, axis_letters
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
 from tinygrad.helpers import strip_parens, getenv, prod, dedup, AMX, CPU_COUNT
 from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate
@@ -144,6 +144,9 @@ class CStyleLanguage(Renderer):
    name = "test"
    for u in uops:
      if u.op is Ops.NOOP: continue
+      if u.op is Ops.AFTER:
+        r[u] = r[u.src[0]]
+        continue
      if u.op is Ops.SINK:
        if u.arg is not None: name = u.arg.function_name
        continue
@@ -160,7 +163,7 @@ class CStyleLanguage(Renderer):
      # naming
      prefix = None
      if u.op is Ops.SPECIAL: r[u] = u.arg
-      elif u.op is Ops.RANGE: r[u] = "ridx"+range_str(u)
+      elif u.op is Ops.RANGE: r[u] = f"{axis_letters[u.arg[-1]]}idx"+range_str(u)
      else:
        prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const",
                  Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.VECTORIZE: "cast", Ops.PRECAST: "precast",
@@ -167,6 +167,9 @@ class LLVMRenderer(Renderer):
    name = "test"
    for u in uops:
      if u.op is Ops.NOOP: continue
+      if u.op is Ops.AFTER:
+        r[u] = r[u.src[0]]
+        continue
      if u.op is Ops.SINK:
        if u.arg is not None: name = u.arg.function_name
        continue
@@ -1,4 +1,4 @@
-from typing import Callable, cast
+from typing import Callable, cast, Any
 from tinygrad.dtype import AddrSpace, DType, PtrDType, dtypes
 from tinygrad.helpers import DEBUG, OSX, unwrap
 from tinygrad.renderer import Renderer
@@ -169,10 +169,13 @@ class NIRRenderer(Renderer):
  def render(self, uops:list[UOp]):
    self.prerender(uops)
    for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]: self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].arg
-    self.r, self.param_idx, ranges = {}, 0, []
+    self.r: dict[UOp, Any] = {}
+    self.param_idx, ranges = 0, []

    for u in uops:
      if u.op == Ops.NOOP or u.op == Ops.INDEX: pass
+      elif u.op is Ops.AFTER:
+        self.r[u] = self.r[u.src[0]]
      elif u.op == Ops.SINK:
        if u.arg is not None: self.b.shader.contents.info.name = mesa.char_pointer_cast(u.arg.function_name)
      elif u.op == Ops.DEFINE_LOCAL:
@@ -179,6 +179,9 @@ class PTXRenderer(Renderer):
    name = "test"
    for u in uops:
      if u.op is Ops.NOOP: continue
+      if u.op is Ops.AFTER:
+        self.r[u] = self.r[u.src[0]]
+        continue
      if u.op is Ops.SINK:
        if u.arg is not None: name = u.arg.function_name
        continue
@@ -278,6 +278,7 @@ class AMDComputeQueue(HWQueue):
  def timestamp(self, signal:AMDSignal):
    with self.pred_exec(xcc_mask=0b1):
      self.release_mem(signal.timestamp_addr, 0, self.pm4.data_sel__mec_release_mem__send_gpu_clock_counter, self.pm4.int_sel__mec_release_mem__none)
+      self.acquire_mem() # ensure timestamp is written
    return self

  def signal(self, signal:AMDSignal, value:sint=0):
@@ -334,7 +335,7 @@ class AMDComputeAQLQueue(AMDComputeQueue):
    def flush_pm4_batch():
      nonlocal pm4_batch
      if not pm4_batch: return bytes()
-      dev.pm4_ibs.cpu_view().view(off:=dev.pm4_ib_alloc.alloc(len(pm4_batch) * 4), fmt='I')[:len(pm4_batch)] = array.array('I', pm4_batch)
+      dev.pm4_ibs.cpu_view().view(off:=dev.pm4_ib_alloc.alloc(len(pm4_batch) * 4, 16), fmt='I')[:len(pm4_batch)] = array.array('I', pm4_batch)
      pkt = [AQL_HDR | (hsa.HSA_PACKET_TYPE_VENDOR_SPECIFIC << hsa.HSA_PACKET_HEADER_TYPE) | (1 << 16),
        self.pm4.PACKET3(self.pm4.PACKET3_INDIRECT_BUFFER, 2), *data64_le(dev.pm4_ibs.va_addr+off), len(pm4_batch)|self.pm4.INDIRECT_BUFFER_VALID, 10]
      pm4_batch.clear()
@@ -701,7 +702,7 @@ class PCIIface(PCIIfaceBase):
      read_ptrs=[gart.cpu_view().view(offset=rptr, size=8, fmt='Q')], write_ptrs=[gart.cpu_view().view(offset=wptr, size=8, fmt='Q')])

  def sleep(self, timeout):
-    if self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
+    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
      self.pci_dev.irq_fd.read(8 * events_cnt)
      self.dev_impl.ih.interrupt_handler()

@@ -833,7 +834,7 @@ class AMDDevice(HCQCompiled):
        read_dispatch_id_field_base_byte_offset=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
        max_cu_id=self.max_cu_id, max_wave_id=self.max_wave_id)
      gart.cpu_view().view(fmt='B')[:ctypes.sizeof(aql_desc)] = bytes(aql_desc)
-      self.aql_desc = hsa.amd_queue_t.from_address(gart.va_addr)
+      self.aql_desc = hsa.amd_queue_t.from_address(gart.cpu_view().addr)

    cwsr_buffer_size = round_up((ctx_save_restore_size + debug_memory_size) * self.iface.props.get('num_xcc', 1), mmap.PAGESIZE)
    cwsr_buffer = self.iface.alloc(cwsr_buffer_size) if ctx_save_restore_size else None
@@ -17,7 +17,7 @@ class NullRenderer(CStyleLanguage):
 class NullProgram:
  def __init__(self, device:str, name:str, lib:bytes): self.device, self.name = device, name
  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
-    with cpu_profile(self.name, self.device): return 1e-4
+    with cpu_profile(self.name, self.device): return 1e-3

 class NullAllocator(Allocator['NullDevice']):
  def _alloc(self, size, options): pass
@@ -28,7 +28,7 @@ class NullAllocator(Allocator['NullDevice']):
  def _offset(self, buf, offset:int, size:int): pass

 class NullGraph(MultiGraphRunner):
-  def __call__(self, input_rawbuffers, var_vals, wait=False) -> float|None: return 1e-3
+  def __call__(self, input_rawbuffers, var_vals, wait=False) -> float|None: return 1e-1

 class NullDevice(Compiled):
  def __init__(self, device:str):
@@ -72,7 +72,8 @@ class PythonProgram:
              if g: _store(m, o+j, v, dtp[1].scalar())
          i += 1
          continue
-        if uop in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}:
+        if uop is Ops.AFTER: ul[i] = inp[0]
+        elif uop in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}:
          assert isinstance(dtype, PtrDType), dtype
          storage_fmt = storage_fmt_for_dtype(dtype.base.scalar())
          if storage_fmt is None: raise RuntimeError(f"{dtype=} is not supported")
@@ -9,7 +9,7 @@ from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
 from tinygrad.runtime.autogen import kgsl, adreno
 from tinygrad.runtime.ops_cl import CLCompiler, CLDevice
 from tinygrad.renderer.cstyle import QCOMRenderer
-from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, prod, fromimport
+from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, prod, fromimport, cpu_profile
 if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl  # noqa: F401  # pylint: disable=unused-import

 BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2
@@ -289,20 +289,21 @@ class QCOMAllocator(HCQAllocatorBase):
      buf.texture_info = QCOMTextureInfo(pitch, real_stride, desc, [desc[0] & (~0xffff), *desc[1:len(desc)]])
    return buf

-  def _do_copy(self, src_addr, dest_addr, src_size, real_size, src_stride, dest_stride, dest_off=0, src_off=0):
-    while src_off < src_size:
-      ctypes.memmove(dest_addr+dest_off, src_addr+src_off, real_size)
-      src_off, dest_off = src_off+src_stride, dest_off+dest_stride
+  def _do_copy(self, src_addr, dest_addr, src_size, real_size, src_stride, dest_stride, prof_text, dest_off=0, src_off=0):
+    with cpu_profile(prof_text, self.dev.device, is_copy=True):
+      while src_off < src_size:
+        ctypes.memmove(dest_addr+dest_off, src_addr+src_off, real_size)
+        src_off, dest_off = src_off+src_stride, dest_off+dest_stride

  def _copyin(self, dest:HCQBuffer, src:memoryview):
    stride, pitch = (src.nbytes, src.nbytes) if (ti:=cast(QCOMTextureInfo, dest.texture_info)) is None else (ti.real_stride, ti.pitch)
-    self._do_copy(mv_address(src), dest.va_addr, src.nbytes, stride, stride, pitch)
+    self._do_copy(mv_address(src), dest.va_addr, src.nbytes, stride, stride, pitch, f"TINY -> {self.dev.device}")

  def _copyout(self, dest:memoryview, src:HCQBuffer):
    self.dev.synchronize()

    stride, pitch = (src.size, src.size) if (ti:=cast(QCOMTextureInfo, src.texture_info)) is None else (ti.real_stride, ti.pitch)
-    self._do_copy(src.va_addr, mv_address(dest), src.size, stride, pitch, stride)
+    self._do_copy(src.va_addr, mv_address(dest), src.size, stride, pitch, stride, f"{self.dev.device} -> TINY")

  def _as_buffer(self, src:HCQBuffer) -> memoryview:
    self.dev.synchronize()
@@ -261,20 +261,20 @@ class HCQSignal(Generic[HCQDeviceType]):
@contextlib.contextmanager
 def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Callable[[], HWQueue]|None=None, queue:HWQueue|None=None):
  st, en = (dev.new_signal(), dev.new_signal()) if enabled else (None, None)
+  assert queue is not None or queue_type is not None, "Either queue or queue_type must be provided"

  if enabled and queue is not None: queue.timestamp(st)
-  elif enabled:
-    assert queue_type is not None
+  elif enabled and queue_type is not None:
    queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(st).signal(dev.timeline_signal, dev.next_timeline()).submit(dev)

  try: yield (st, en)
  finally:
    if enabled and queue is not None: queue.timestamp(en)
-    elif enabled:
-      assert queue_type is not None
+    elif enabled and queue_type is not None:
      queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(en).signal(dev.timeline_signal, dev.next_timeline()).submit(dev)

-    if enabled and PROFILE: dev.sig_prof_records.append((cast(HCQSignal, st), cast(HCQSignal, en), desc, queue_type is dev.hw_copy_queue_t))
+    if enabled and PROFILE:
+      dev.sig_prof_records.append((cast(HCQSignal, st), cast(HCQSignal, en), desc, (queue_type or type(queue)) is dev.hw_copy_queue_t))

 class HCQArgsState(Generic[ProgramType]):
  def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:tuple[HCQBuffer, ...], vals:tuple[sint, ...]=()):
@@ -1,6 +1,6 @@
 import os, mmap, array, functools, ctypes, select, contextlib, dataclasses, sys, errno, itertools
 from typing import cast, ClassVar
-from tinygrad.helpers import round_up, getenv, OSX, temp
+from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv
 from tinygrad.runtime.autogen import libc, vfio
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, HCQBuffer
 from tinygrad.runtime.support.memory import MemoryManager, VirtMapping
@@ -77,15 +77,14 @@ class _System:
      sysmem_view = System.iokit_pci_memmap(round_up(size, mmap.PAGESIZE))
      paddrs = list(itertools.takewhile(lambda p: p[1] != 0, zip(sysmem_view.view(fmt='Q')[0::2], sysmem_view.view(fmt='Q')[1::2])))
      assert not contiguous or len(paddrs) == 1, "not contiguous, but required"
-      paged_paddrs = [p + i for p, sz in paddrs for i in range(0, sz, 0x1000)][:round_up(size, 0x1000)//0x1000]
    else:
      assert not contiguous or size <= (2 << 20), "Contiguous allocation is only supported for sizes up to 2MB"
-      flags = (libc.MAP_HUGETLB if contiguous and (size:=round_up(size, mmap.PAGESIZE)) > 0x1000 else 0) | (MAP_FIXED if vaddr else 0)
+      flags = (libc.MAP_HUGETLB if contiguous and (size:=round_up(size, mmap.PAGESIZE)) > mmap.PAGESIZE else 0) | (MAP_FIXED if vaddr else 0)
      va = FileIOInterface.anon_mmap(vaddr, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED|mmap.MAP_ANONYMOUS|MAP_POPULATE|MAP_LOCKED|flags, 0)
-      sysmem_view, paged_paddrs = MMIOInterface(va, size), self.system_paddrs(va, size)
+      sysmem_view, paddrs = MMIOInterface(va, size), [(x, mmap.PAGESIZE) for x in self.system_paddrs(va, size)]

    if data is not None: sysmem_view[:len(data)] = data
-    return sysmem_view, paged_paddrs
+    return sysmem_view, [p + i for p, sz in paddrs for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)]

  def pci_reset(self, gpu):
    if OSX: System.iokit_pci_rpc(__TinyGPURPCReset:=2)
@@ -202,7 +201,7 @@ class LNXPCIIfaceBase:
      mapping = self.dev_impl.mm.map_range(vaddr, size, [(paddr, 0x1000) for paddr in paddrs], system=True, snooped=True, uncached=True)
      return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=True, hMemory=paddrs[0]), view=memview, owner=self.dev)

-    mapping = self.dev_impl.mm.valloc(size:=round_up(size, 4 << 10), uncached=uncached, contiguous=cpu_access)
+    mapping = self.dev_impl.mm.valloc(size:=round_up(size, 0x1000), uncached=uncached, contiguous=cpu_access)
    barview = self.pci_dev.map_bar(bar=self.vram_bar, off=mapping.paddrs[0][0], size=mapping.size) if cpu_access else None
    return HCQBuffer(mapping.va_addr, size, view=barview, meta=PCIAllocationMeta(mapping, cpu_access, hMemory=mapping.paddrs[0][0]), owner=self.dev)

@@ -2,9 +2,9 @@ from typing import Iterator
 import functools, operator, itertools
 from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, AddrSpace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
-from tinygrad.helpers import argsort, all_same, cpu_profile, TracingKey
+from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored

 ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
                     Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
@@ -40,25 +40,32 @@ class BufferizeOpts:

@dataclass
 class IndexingContext:
-  realize_map: dict[UOp, None] = field(default_factory=dict)
+  realize_map: dict[UOp, None|list[int]] = field(default_factory=dict)
  range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[UOp, ...]]] = field(default_factory=dict)

  # create ranges
  range_idx: Iterator[int] = field(default_factory=itertools.count)
-  def new_range(self, s:sint, axistype:AxisType=AxisType.LOOP):
+  def new_range(self, s:sint, axistype:AxisType=AxisType.LOOP) -> UOp:
+    if isinstance(s, UOp) and s.op is Ops.RANGE: return s
+    # if a range has a 1 src, it's the same as UOp.const(dtypes.index, 0)
    return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0)

 def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
  if x.op in {Ops.BUFFERIZE, Ops.INDEX, Ops.KERNEL}: return None
-  if x.op is Ops.ASSIGN and x.src[1].op is Ops.KERNEL: return None
+  if x.op is Ops.AFTER and x.src[1].op is Ops.KERNEL: return None
  new_srcs = []
  for s in x.src:
    new_src = s
-    if s.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT} or (s.op is Ops.ASSIGN and s.src[1].op is Ops.KERNEL):
+    if s.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT} or (s.op is Ops.AFTER and s.src[1].op is Ops.KERNEL):
      if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0])
    elif s in ctx.realize_map:
-      new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(new_src,)+tuple(ctx.range_map[s][1]), arg=BufferizeOpts(device=s.device), tag=s.tag)
-      if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0])
+      realized_ranges = ctx.realize_map[s]
+      assert isinstance(realized_ranges, list), "realize map must contain range list"
+      closed_ranges = tuple([r for i,r in enumerate(ctx.range_map[s][1]) if i in realized_ranges])
+      # None in the device assigns it a number later
+      opts = BufferizeOpts(device=s.device) if len(ctx.range_map[s][1]) == len(realized_ranges) else BufferizeOpts(None, AddrSpace.LOCAL)
+      new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(new_src,)+closed_ranges, arg=opts, tag=s.tag if opts.addrspace == AddrSpace.GLOBAL else None)
+      if x in ctx.range_map: new_src = new_src.index(*[r for i,r in enumerate(ctx.range_map[x][0]) if i in realized_ranges])
    new_srcs.append(new_src)
  # NOTE: do we need this?
  return x.replace(src=tns) if x.src != (tns:=tuple(new_srcs)) else None
@@ -132,26 +139,23 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
    case _: raise RuntimeError(f"{op} is not a MovementOp")
  return rngs

-@cpu_profile(TracingKey("run_rangeify"), "TINY")
+@profile_matches
 def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
  rctx = IndexingContext()

  # get ops to realize
-  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, name="Input Graph")
+  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, name="get realize")

  # get the traversal order
-  with cpu_profile(TracingKey("reverse toposort"), "TINY"):
+  with cpu_profile("reverse toposort", "TINY"):
    tsink_reverse_toposort = tsink.reverse_toposort(consumer_map:=tsink.get_consumer_map())

  # explicit rangeify
-  ending_ranges: dict[UOp, bool] = {}
+  ending_ranges: dict[UOp, list[UOp]] = {}
  for x in tsink_reverse_toposort:
    if x.op in {Ops.DEVICE, Ops.UNIQUE}: continue
    if x.dtype.scalar() == dtypes.index: continue  # TODO: why do I need this?
-    ending_ranges[x] = any(ending_ranges[u] for u in consumer_map[x])
-
-    # if this element has weight and it's ending a range, we (force) realize it
-    if ending_ranges[x] and x.op in GroupOp.Elementwise.union({Ops.REDUCE_AXIS}): rctx.realize_map[x] = None
+    ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])

    # *** the ranges on the output are
    #  1. new if this op is realized
@@ -161,9 +165,12 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
    consumer_rngs = [rctx.range_map[c][0] for c in consumer_map[x] if c in rctx.range_map]
    if x in rctx.realize_map:
      # if this is in the realize_map, we create new ranges (at the output)
-      out_rngs = tuple(rctx.new_range(s) if not isinstance(s, UOp) or s.op is not Ops.RANGE else s for s in x.shape)
+      out_rngs = tuple(rctx.new_range(s) for s in x.shape)
      # all ranges are ended now
-      ending_ranges[x] = False
+      ending_ranges[x] = []
+      # mark all ranges as ended
+      assert rctx.realize_map[x] is None
+      rctx.realize_map[x] = list(range(len(x.shape)))
    elif x.op in {Ops.MSTACK, Ops.MSELECT}:
      # treat MSTACK/MSELECT like SINK
      continue
@@ -175,29 +182,41 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
      out_rngs = consumer_rngs[0]
    elif len(consumer_rngs) > 1:
      # if this has two consumers, we have to merge the ranges and might create new ones
-      all_rngs = list(zip(*consumer_rngs))
+      all_rngs: list[tuple[UOp, ...]] = list(zip(*consumer_rngs))
      rngs_valids = []
      for valid_rngs in all_rngs:
        local_rngs, valids = zip(*[(r.get_idx(), r.get_valid()) for r in valid_rngs])
-        # if a range has a 1 src, it's the same as UOp.const(dtypes.index, 0)
-        same_rngs = [x if x.op is not Ops.RANGE or resolve(x.src[0] != 1) else UOp.const(dtypes.index, 0) for x in local_rngs]
-        rngs_valids.append((local_rngs, valids, all_same(same_rngs)))
+        rngs_valids.append((local_rngs, valids))

      # TODO: in RANGEIFY > 1 all_all_same isn't required
-      all_all_same = all(same_rngs for _,_,same_rngs in rngs_valids)
+      all_all_same = all(all_same(local_rngs) for local_rngs,_ in rngs_valids)
      _out_rngs = []
-      for i,(local_rngs,valids,same_rngs) in enumerate(rngs_valids):
+      _realize_axis = []
+      for i,(local_rngs,valids) in enumerate(rngs_valids):
        # we compare the ranges without their valids
-        if all_all_same:
+        if all_all_same or (PCONTIG and all_same(local_rngs)):
          # the new valid is the OR of all the children valids
          minimum_valid = functools.reduce(operator.or_, valids, UOp.const(dtypes.bool, False))
          _out_rngs.append(graph_rewrite(minimum_valid.where(local_rngs[0], UOp.invalid()), symbolic, name="minimum_valid"))
        else:
          _out_rngs.append(rctx.new_range(x.shape[i]))
+          _realize_axis.append(i)
      out_rngs = tuple(_out_rngs)

-      # we have to realize here if there's new ranges
-      if not all_all_same: rctx.realize_map[x] = None
+      # we have to (partially) realize here if there's new ranges
+      if len(_realize_axis): rctx.realize_map[x] = _realize_axis
+
+    # if this element is a reduce and there's ended ranges, we might have to end some other ranges
+    if len(ending_ranges[x]) and x.op in GroupOp.Elementwise.union({Ops.REDUCE_AXIS}):
+      _realize_axis = rctx.realize_map.get(x, []) or []
+      for i,r in enumerate(out_rngs):
+        if i in _realize_axis: continue
+        if not (PCONTIG > 1) or any(any(rr.arg > e.arg for e in ending_ranges[x]) for rr in r.ranges):
+          _realize_axis.append(i)
+      ending_ranges[x] = []
+      if len(_realize_axis):
+        rctx.realize_map[x] = _realize_axis
+        out_rngs = tuple([(rctx.new_range(x.shape[i]) if i in _realize_axis else r) for i,r in enumerate(out_rngs)])

    # TODO: some ops don't have shape, enable this after the `.st` property is removed
    #assert len(out_rngs) == len(x.shape), \
@@ -213,15 +232,22 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
    # apply movement ops
    if x.op in GroupOp.Movement: rngs = apply_movement_op(x.op, x.src[0].shape, x.marg, rngs)
    # if the EXPAND is used to inject a range, we don't mark it as ending_ranges. otherwise we do.
-    if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape): ending_ranges[x] = True
+    # NOTE: this doesn't actually always end a range, but this is why convs are realized, so for now we need it
+    if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape):
+      ending_ranges[x] = list(UOp.sink(*[ro for ri, ro in zip(rngs, out_rngs) if ri is not ro]).ranges.keys())

    # REDUCE_AXIS creates ranges for the axes it is reducing
    if x.op is Ops.REDUCE_AXIS:
      rngs = tuple(rctx.new_range(s, axistype=AxisType.REDUCE) if i in x.arg[1] else r for i,(r,s) in enumerate(zip(rngs, x.src[0].shape)))

    if debug:
-      print("***" if x in rctx.realize_map else "   ", len(consumer_map[x]), f"{str(x.op):20s}",
-            UOp.sink().index(*rngs).render(), " -> ", UOp.sink().index(*out_rngs).render())
+      realized_ranges = rctx.realize_map.get(x, None)
+      disp = []
+      for i, (ri, ro) in enumerate(zip([r.render() for r in rngs], [r.render() for r in out_rngs])):
+        rng = f"{ri}" if ri == ro else f"{ri} -> {ro}"
+        if realized_ranges is not None and i in realized_ranges: rng = colored(rng, "yellow")
+        disp.append("["+rng+"]")
+      print("***" if x in rctx.realize_map else "   ", len(consumer_map[x]), f"{str(x.op):20s}", ''.join(disp))

    # assign to the range map. rngs are the input ranges, out_rngs are the output ranges, from the x op.
    rctx.range_map[x] = (rngs, out_rngs)
@@ -1,7 +1,7 @@
 from typing import cast
 import functools, itertools, operator
 from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, getenv
-from tinygrad.uop.ops import Ops, UOp, sint, PatternMatcher, UPat, GroupOp, track_rewrites, graph_rewrite_map
+from tinygrad.uop.ops import Ops, UOp, sint, PatternMatcher, UPat, GroupOp, track_rewrites, graph_rewrite_map, graph_rewrite
 from tinygrad.device import Device

 # *** allreduce implementation ***
@@ -219,4 +219,8 @@ multi_pm = PatternMatcher([
 ])+replace_allreduce

@track_rewrites()
-def get_multi_map(big_sink:UOp) -> dict[UOp, UOp]: return graph_rewrite_map(big_sink, multi_pm, name="multi_pm")
+def get_multi_map(big_sink:UOp) -> dict[UOp, UOp]:
+  if getenv("VIZ"): graph_rewrite(big_sink, PatternMatcher([]), name="View Multi AST")
+  ret = graph_rewrite_map(big_sink, multi_pm, name="multi_pm")
+  if getenv("VIZ"): graph_rewrite(ret[big_sink], PatternMatcher([]), name="View Post Multi AST")
+  return ret
@@ -3,8 +3,8 @@ from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, ssimplify, KernelInfo
 from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType
-from tinygrad.uop.symbolic import symbolic_simple
-from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, Metadata
+from tinygrad.uop.symbolic import symbolic_flat
+from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, Metadata, REAL_SUBSTITUTE
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_unparented
 from tinygrad.codegen.opt import Opt
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext, apply_movement_op
@@ -178,8 +178,11 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
  # if it makes it here, the bufferize is removed
  # this is the ranges replaced
  # NOTE: if buf src is a const, we don't replace it
-  replaces = flatten([(k,v) for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST])
-  return UOp(Ops.SUBSTITUTE, dtype=src.dtype, src=(src, UOp(Ops.NOOP, src=tuple(replaces[0::2])), UOp(Ops.NOOP, src=tuple(replaces[1::2]))))
+  if REAL_SUBSTITUTE:
+    return src.substitute({k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST})
+  else:
+    replaces = flatten([(k,v) for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST])
+    return UOp(Ops.SUBSTITUTE, dtype=src.dtype, src=(src, UOp(Ops.NOOP, src=tuple(replaces[0::2])), UOp(Ops.NOOP, src=tuple(replaces[1::2]))))

 def pre_bufferize(b:UOp, x:UOp, copy:UOp):
  nb = b.replace(src=(b.src[0].contiguous(),)+b.src[1:])
@@ -208,6 +211,38 @@ pm_cleanups = pm_mops+PatternMatcher([
   lambda s: UOp.const(c.dtype, c.arg) if (c:=s.base).op is Ops.CONST else None),
 ])

+def second_stage_removal(src:UOp, buf:UOp):
+  if buf.arg.addrspace != AddrSpace.GLOBAL: return None
+  # if it's user contiguous, we never remove it
+  if src.op in ALWAYS_RUN_OPS: return None
+  accessed_buffers: list[UOp] = []
+  def red_gate(x:UOp):
+    if x.op is Ops.BUFFERIZE and x.arg.addrspace == AddrSpace.GLOBAL:
+      accessed_buffers.append(x)
+      return False
+    if x.op is Ops.BUFFER:
+      accessed_buffers.append(x)
+    return True
+  src.toposort(gate=red_gate)
+  del red_gate
+  accessed_buffers = dedup(accessed_buffers)
+  out_in_ratio = prod(buf.shape) / sum([x.size for x in accessed_buffers])
+  print(f"ratio {out_in_ratio:.2f} {buf.shape} from {[x.size for x in accessed_buffers]}")
+  if out_in_ratio > 10: return buf.replace(op=Ops.REMOVE, arg=None)
+
+pm_cleanups_2 = PatternMatcher([
+  (UPat.var("src").f(Ops.BUFFERIZE, allow_any_len=True, name="buf"), second_stage_removal),
+])
+
+def do_remove(rem:UOp, idx:UOp):
+  assert len(rem.src) == len(idx.src), f"remove on wrong bufferize, {len(rem.src)} != {len(idx.src)}"
+  return rem.src[0].substitute({k:v for k,v in zip(rem.src[1:], idx.src[1:]) if k.op is not Ops.CONST})
+
+pm_cleanup_remove = PatternMatcher([
+  # NOTE: this is better than substitute
+  (UPat(Ops.REMOVE, name="rem").f(Ops.INDEX, allow_any_len=True, name="idx"), do_remove),
+])
+
 def late_buffer_view(t:UOp, b:UOp):
  if isinstance(b.device, str) and (b.device.startswith("DISK") or b.device.startswith("TINYFS")):
    rngs = b.src[1:]
@@ -215,7 +250,9 @@ def late_buffer_view(t:UOp, b:UOp):

    # walk up for the INDEX
    x = t
-    while not any(u.op is Ops.INDEX for u in x.src): x = x.src[0]
+    while not any(u.op is Ops.INDEX for u in x.src):
+      assert x.op not in GroupOp.Elementwise, "can't buffer view elementwise"
+      x = x.src[0]
    x = next(u for u in x.src if u.op is Ops.INDEX)

    if len(shape) == 0: offset = x.src[1].arg
@@ -237,7 +274,7 @@ def limit_bufs(ctx:IndexingContext, root:UOp):
  bufs: set[UOp] = set()
  def gate_input(u:UOp):
    # TODO: add cache to fix n^2
-    if is_load:=(u.op in {Ops.BUFFERIZE, Ops.ASSIGN, Ops.BUFFER, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_VAR}): bufs.add(u)
+    if is_load:=(u.op in {Ops.BUFFERIZE, Ops.AFTER, Ops.BUFFER, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_VAR}): bufs.add(u)
    return not is_load
  root.toposort(gate=gate_input)

@@ -272,7 +309,8 @@ def bufferize_to_store(x:UOp):
    assert assign_target.op is Ops.INDEX, f"{assign_target.op} is not index"
    # in assign, this is the buffer size, not the bufferize size
    # TODO: assign_mops here
-    ret = assign_target.replace(dtype=sdtype).store(assign_src, *rngs, dtype=x.dtype).replace(tag=x.tag)
+    do_store = assign_target.replace(dtype=sdtype).store(assign_src, *rngs).replace(tag=x.tag)
+    ret = assign_target.src[0].after(do_store)
    mops = []
    walk = assign_mops
    while walk is not assign_mops.base:
@@ -284,8 +322,8 @@ def bufferize_to_store(x:UOp):
  # NOTE: the DEFINE_LOCAL needs to be disambiguated here
  if sdtype.addrspace == AddrSpace.GLOBAL:
    buf = UOp.new_buffer(x.arg.device, size, x.dtype)
-    ret = buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=x.dtype).replace(tag=x.tag)
-    ret = ret.forced_reshape(shape)
+    do_store = buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs).replace(tag=x.tag)
+    ret = buf.after(do_store).forced_reshape(shape)
    # TODO: is this right? what if it's offset
    if any(r.op is Ops.RANGE and r.src[0].op is not Ops.CONST for r in rngs):
      sym_shape = tuple([ssimplify(r.src[0]) if r.op is not Ops.CONST else 1 for r in rngs])
@@ -297,7 +335,7 @@ def bufferize_to_store(x:UOp):
  if tag is None: tag = UOp.unique().arg # TODO: hack
  buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=tag)
  # store has the other dtype here
-  # TODO: how is this unified?
+  # TODO: use after here?
  return buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=sdtype).reshape(shape)

 pm_add_buffers = pm_mops+to_bufferview+PatternMatcher([
@@ -330,12 +368,12 @@ def unbind_kernel(ctx:LocalAddBufferContext, b:UOp):
  ctx.vars[b] = None
  return b.src[0]

-def handle_assign(ctx:LocalAddBufferContext, assign:UOp):
-  buf = assign.as_buf()
+def handle_after(ctx:LocalAddBufferContext, after:UOp):
+  buf = after.as_buf()
  # HACK to put the buffer in the MAP instead of MSTACK/MSELECT
  if buf.op in {Ops.MSTACK, Ops.MSELECT}: buf = buf.src[0]
  assert buf not in ctx.map
-  ctx.map[buf] = assign
+  ctx.map[buf] = after
  return buf

 def renumber_range(ctx:LocalAddBufferContext, r:UOp):
@@ -345,7 +383,7 @@ def renumber_range(ctx:LocalAddBufferContext, r:UOp):
  return ret

 def find_bufs(x:UOp):
-  idxs = [s for s in x.toposort(gate=lambda x: x.op is not Ops.ASSIGN) if s.op is Ops.INDEX]
+  idxs = [s for s in x.toposort(gate=lambda x: x.op is not Ops.AFTER) if s.op is Ops.INDEX]
  read_from: dict[UOp, Ops] = {}
  if any((buf:=idx.as_buf()).op is Ops.BUFFER and read_from.setdefault(buf, op:=idx.src[0].op) is not op for idx in idxs):
    raise RuntimeError(f"cycle detected while indexing {buf}")
@@ -354,7 +392,7 @@ to_define_global = PatternMatcher([
  (UPat(Ops.STORE, name="x"), find_bufs),
  (UPat(Ops.BUFFER, name="buf"), debuf),
  (UPat(Ops.BIND, name="b"), unbind_kernel),
-  (UPat((Ops.ASSIGN, Ops.MSTACK, Ops.MSELECT), name="assign"), handle_assign),
+  (UPat((Ops.MSTACK, Ops.MSELECT, Ops.AFTER), name="after"), handle_after),

  # HACK in case any CONSTs were replaced
  # this is only needed if you are using symbolic
@@ -413,7 +451,7 @@ class Kernel:
    ast_rep = f"SINK{tuple(s.op for s in self.ast.src)}" if self.ast.op is Ops.SINK else repr(self.ast.op)
    return f"<Kernel {len(list(self.ast.toposort()))} {ast_rep} {self.metadata}>"

-def split_store(ctx:list[UOp], x:UOp):
+def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
  if len(x.ranges): return None
  if x.src[0].ptrdtype.addrspace is AddrSpace.LOCAL: return None

@@ -431,7 +469,7 @@ def split_store(ctx:list[UOp], x:UOp):
  kernel = UOp(Ops.KERNEL, src=tuple(lctx.map.values())+tuple(lctx.vars.keys()), arg=kernel_arg)
  if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src if x.op is not Ops.BIND]):
    raise RuntimeError(f"all buffers must be on the same device: {tuple(b.buf_uop.buffer for b in kernel.src)}")
-  return x.as_buf().assign(kernel)
+  return kernel

 split_kernels = PatternMatcher([
  (UPat(Ops.STORE, name="x"), split_store),
@@ -486,6 +524,7 @@ pm_substitute_recurse = PatternMatcher([(UPat(Ops.SUBSTITUTE, src=(UPat(), UPat(

@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len([u for u in UOp.sink(*ret.values()).toposort() if u.op is Ops.KERNEL]))}", True)
 def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
+  if getenv("VIZ"): graph_rewrite(sink, PatternMatcher([]), name="View Input Graph")
  uop_list: list[UOp] = []
  tsink = graph_rewrite(sink, add_tags, ctx=uop_list, bottom_up=True, name="number the uops")

@@ -495,10 +534,12 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
  tsink, rctx = run_rangeify(tsink, getenv("DEBUG_RANGEIFY", 0))

  # NOTE: sym (vs symbolic_simple) breaks things here because ranges with len 1 aren't handled right
-  tsink = graph_rewrite(tsink, symbolic_simple+pm_reduce_unparented, name="symbolic")  # this supports const folding
+  tsink = graph_rewrite(tsink, symbolic_flat+pm_reduce_unparented, name="symbolic")  # this supports const folding
  tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers")
  # TODO: can you substitute and remove costly buffers at the same time?
  tsink = graph_rewrite(tsink, pm_substitute_recurse, bottom_up=True, name="run substitutes")
+  tsink = graph_rewrite(tsink, pm_cleanups_2, name="remove (more) costly buffers")
+  tsink = graph_rewrite(tsink, pm_cleanup_remove, name="actually remove")
  tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")

  # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph
@@ -517,13 +558,13 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
  kernel_assign: dict[UOp, UOp] = {}
  assign_rep: dict[UOp, UOp] = {}
  for u in tsink.toposort():
-    if u.op is not Ops.ASSIGN: continue
+    if u.op is not Ops.AFTER: continue
    kernel_assign[u.buf_uop] = u
    for s in u.src[1].src:
      # TODO: this is probably broken for MSELECT/MSTACK
      if s.op is not Ops.BUFFER or s is u.buf_uop or (a:=kernel_assign.get(s)) is None: continue
-      if any(x.op is Ops.ASSIGN and x.buf_uop is s for x in u.toposort()):
-        raise RuntimeError(f"cycle detected in graph, kernel for {u.buf_uop} must either depend on ASSIGN or BUFFER")
+      if any(x.op is Ops.AFTER and x.buf_uop is s for x in u.toposort()):
+        raise RuntimeError(f"cycle detected in graph, kernel for {u.buf_uop} must either depend on AFTER or BUFFER")
      assign_rep[a] = kernel_assign[s] = a.replace(src=a.src+(u,))
  if assign_rep: tsink = graph_rewrite(tsink, _substitute, ctx=assign_rep, bottom_up=True, name="fix_assign")

@@ -249,9 +249,9 @@ class Tensor(MathTrait):
    self.kernelize(*lst)
    sink = UOp.sink(*[x.uop for x in (self,)+lst])

-    # remove all ASSIGNs, after scheduling, the tensors are just buffers
-    remove_assign_map = {u:u.buf_uop for u in sink.toposort() if u.op is Ops.ASSIGN}
-    _apply_map_to_tensors(remove_assign_map, name="Remove Assigns")
+    # remove all AFTERs, after scheduling, the tensors are just buffers
+    remove_assign_map = {u:u.buf_uop for u in sink.toposort() if u.op is Ops.AFTER}
+    _apply_map_to_tensors(remove_assign_map, name="Remove After")

    # create the schedule
    schedule, var_vals = create_schedule_with_vars(sink)
@@ -12,11 +12,15 @@ class Ops(FastEnum):
  NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); PRECAST = auto(); REWRITE_ERROR = auto()  # noqa: E702
  SENTINEL = auto()

+  # AFTER passes src[0] through and promises in the toposort that any consumers of the AFTER run after src[1:]
+  AFTER = auto()
+
  # buffer ops
  COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto() # noqa: E702

  # create buffer
  BUFFERIZE = auto()
+  REMOVE = auto()
  SUBSTITUTE = auto()

  # ops that adjust the behavior of the scheduler
@@ -17,7 +17,7 @@ class AxisType(Enum):
  GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
  THREAD = auto()

-range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3}
+range_start = {Ops.BUFFERIZE: 1, Ops.REMOVE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3}

 # 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)
@@ -186,11 +186,12 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
      case Ops.CONST | Ops.DEFINE_VAR | Ops.BIND: return () if self._device is not None else None
      case Ops.BUFFER: return (self.arg,)
      case Ops.BUFFER_VIEW: return (self.arg[0],)
-      case Ops.BUFFERIZE: return tuple([int(r.vmax+1) for r in self.src[1:]])
+      case Ops.BUFFERIZE | Ops.REMOVE: return tuple([int(r.vmax+1) for r in self.src[1:]])
      case Ops.DEFINE_GLOBAL | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)

      # passthrough ops
-      case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.FUSE: return self.src[0]._shape
+      case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.FUSE | Ops.AFTER:
+        return self.src[0]._shape

      # ops with custom handling
      case Ops.KERNEL: return self.arg.ast._shape
@@ -324,7 +325,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
  def detach(self): return UOp(Ops.DETACH, self.dtype, (self,))
  def index(self, *srcs:UOp|None, **kwargs):
    return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
-  def __getitem__(self, idx): return self.index(idx)
+  def __getitem__(self, *idx): return self.index(*idx)
  def const_like(self, b:ConstLike):
    # constants can optionally have a DEVICE source
    return UOp.const(self.dtype, b, device=self._device, shape=self._shape)
@@ -349,6 +350,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    return UOp(Ops.GEP, self.dtype.scalar().vec(len(i)) if len(i) > 1 else self.dtype.scalar(), (self,), i)
  def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype.base), src=(self,)+src, **kwargs)
  def store(self, *src:UOp, **kwargs): return UOp(Ops.STORE, kwargs.pop("dtype", dtypes.void), (self,)+src, **kwargs)
+  def after(self, *src:UOp): return UOp(Ops.AFTER, self.dtype, (self,)+src)
  def assign(self, x:UOp): return UOp(Ops.ASSIGN, self.dtype, (self, x))
  def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
  def alu(self, op, *src:UOp, **kwargs):
@@ -525,6 +527,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
  def _device(self) -> str|tuple[str, ...]|None:
    if self.op is Ops.DEVICE: return self.arg
    if self.op is Ops.BUFFERIZE: return self.arg.device
+    if self.op is Ops.AFTER: return self.src[0]._device
    if self.op is Ops.MSELECT:
      assert isinstance(self.src[0].device, tuple), "mselect must be on tuple device"
      return self.src[0].device[self.arg]
@@ -538,8 +541,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    if self.op is Ops.BUFFER: return self
    if self.op is Ops.MSELECT: return self.src[0].buf_uop.mselect(self.arg)
    if self.op is Ops.MSTACK: return UOp(Ops.MSTACK, self.dtype, src=tuple(x.buf_uop for x in self.src))
-    assert self.op is Ops.ASSIGN, f"must be ASSIGN {self.op}"
-    return self.src[0].base
+    assert self.op is Ops.AFTER, f"must be AFTER {self.op}"
+    return self.src[0].buf_uop.base

  def as_buf(self) -> UOp:
    if self.op is Ops.MSELECT: return self.src[0].as_buf().mselect(self.arg)
@@ -652,6 +655,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
      new_count.subtract(div_fac.split_uop(Ops.MUL))
      if const%div_const==0 and all(v>=0 for v in new_count.values()): return math.prod([*new_count.elements(), self.const_like(const//div_const)])
    return None # generic None if we aren't sure
+  def sum(self:UOp, *uops:UOp) -> UOp: return functools.reduce(operator.or_ if self.dtype is dtypes.bool else operator.add, uops, self)
+  def prod(self:UOp, *uops:UOp) -> UOp: return functools.reduce(operator.and_ if self.dtype is dtypes.bool else operator.mul, uops, self)
  @property
  def vmin(self) -> ConstType: return self._min_max[0]
  @property
@@ -808,6 +813,8 @@ class UPat(MathTrait):
  @staticmethod
  def any(*src): return UPatAny(src=src)
  def or_casted(self, name:str|None=None): return UPat.any(self if name is None else self.named(name), UPat(Ops.CAST, name=name, src=(self,)))
+  def or_after(self, name:str|None=None):
+    return UPat.any(self if name is None else self.named(name), UPat(Ops.AFTER, name=name, src=(self,), allow_any_len=True))

  @staticmethod
  @functools.cache
@@ -928,7 +935,7 @@ class TrackedGraphRewrite:
  loc:tuple[str, int]                           # location that called graph_rewrite
  sink:int                                      # the sink input to graph_rewrite
  matches:list[tuple[int, int, tuple, float]]   # before/after UOp, UPat location and time
-  name:str|None                                 # optional name of the rewrite
+  name:str                                      # name of the rewrite
  depth:int                                     # depth if it's a subrewrite
  bottom_up:bool

@@ -973,19 +980,21 @@ def track_rewrites(name:Callable[..., str|TracingKey]|bool=True, replay:bool=Fal
  return _decorator

 active_rewrites:list[TrackedGraphRewrite] = []
-def track_matches(func):
-  def _track_func(*args, **kwargs):
+def profile_matches(fxn:Callable):
+  def wrap(*args, **kwargs):
+    name = str(kwargs.get("name", None) or fxn.__name__)
+    assert args and isinstance(args[0], UOp), f"invalid match tracing inputs for {name} with {args}"
    if tracking:=(TRACK_MATCH_STATS >= 2):
      loc = ((frm:=sys._getframe(1)).f_code.co_filename, frm.f_lineno)
      depth = len(active_rewrites)
-      if not tracked_ctxs: add_trace_group(TracingKey(f"default {func.__name__}"))
-      tracked_ctxs[-1].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], kwargs.get("name", None), depth, kwargs.get("bottom_up", False)))
+      if not tracked_ctxs: add_trace_group(TracingKey(f"default {fxn.__name__}"))
+      tracked_ctxs[-1].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], name, depth, kwargs.get("bottom_up", False)))
      active_rewrites.append(ctx)
-    with cpu_profile(kwargs.get("name", "<unnamed>"), "TINY", display=tracking):
-      ret = func(*args, **kwargs)
+    with cpu_profile(name, "TINY", display=tracking):
+      ret = fxn(*args, **kwargs)
    if tracking: active_rewrites.pop()
    return ret
-  return _track_func
+  return wrap

 class TrackedPatternMatcher(PatternMatcher):
  def rewrite(self, uop:UOp, ctx=None) -> UOp|None:
@@ -1125,12 +1134,12 @@ class RewriteContext:
          self.replace[n] = replaced_new_n
    return self.replace[root]

-@track_matches
+@profile_matches
 def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None) -> UOp:
  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx)
  return rewrite_ctx.unified_rewrite(sink)

-@track_matches
+@profile_matches
 def graph_rewrite_map(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None,
                      input_map:dict[UOp, UOp]|None=None, ) -> dict[UOp, UOp]:
  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx)
@@ -1167,7 +1176,10 @@ pm_lower_index_dtype = PatternMatcher([
  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast(), UPat.var("valid"))), lambda buf,idx,valid: buf.index(idx, valid)),
  (UPat((Ops.STORE, Ops.LOAD), src=(UPat(), UPat(), UPat().cast(dtypes.index)), allow_any_len=True, name="s"),
    lambda s: s.replace(src=s.src[:2]+tuple(u.src[0] for u in s.src[2:]))),
+  # TODO: this is only triggering if they are all casts, correct?
  (UPat((Ops.SINK, Ops.NOOP), src=UPat().cast(dtypes.index), name="n"), lambda n: n.replace(src=tuple(s.src[0] for s in n.src))),
+  # TODO: this should be more general
+  (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=tuple(y.src[0] if y.op is Ops.CAST and y.dtype.scalar()==dtypes.index else y for y in x.src))),
 ])
 def _index_to_concrete_int(u:UOp): return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]

@@ -66,8 +66,8 @@ buffer_spec = PatternMatcher([
 ])

 assign_spec = PatternMatcher([
-  # KERNEL can attach to an ASSIGN to describe the compute required to realize a BUFFER
-  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.BUFFER_VIEW, Ops.ASSIGN, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), lambda: True),
+  # KERNEL can attach to an AFTER to describe the compute required to realize a BUFFER
+  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.BUFFER_VIEW, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), lambda: True),

  # ASSIGN has a target and a value. It can also optionally depend on other assigns
  (UPat(Ops.ASSIGN, name="x"), lambda x: len(x.src) >= 2 and all(s.op is Ops.ASSIGN for s in x.src[2:])),
@@ -111,6 +111,9 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([

  # REDUCE with an outerworld range
  (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])),
+
+  # AFTER if things were kernelized
+  (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.AFTER)),), allow_any_len=True), lambda: True)
 ])

 # ***** uop type spec *****
@@ -162,6 +165,9 @@ spec = PatternMatcher([

  (UPat(Ops.CONST, src=(), name="x"), lambda x: type(x.arg) is type(dtypes.as_const(x.arg, x.dtype))),

+  # allow AFTER on buffers
+  (UPat(Ops.AFTER, src=(UPat(GroupOp.Defines),), allow_any_len=True), lambda: True),
+
  # **** new style load/store ****

  # make sure all index dtypes have been lowered
@@ -171,8 +177,8 @@ spec = PatternMatcher([

  # INDEX is used in new style load/store
  # INDEX takes a <buf, alu, gate?>
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines), UPat())), lambda: True),
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines), UPat(), UPat(dtype=dtypes.bool))), lambda: True),
+  (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines).or_after(), UPat())), lambda: True),
+  (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines).or_after(), UPat(), UPat(dtype=dtypes.bool))), lambda: True),

  # LOAD on STORE
  (UPat(Ops.LOAD, src=(UPat(Ops.STORE),), allow_any_len=True), lambda: True),
@@ -283,6 +289,8 @@ full_spec = PatternMatcher([
  (UPat(Ops.DEFINE_VAR), lambda: True),
  # reshape on STORE
  (UPat(Ops.RESHAPE, src=(UPat(Ops.STORE),)), lambda: True),
+  # allow any AFTER
+  (UPat(Ops.AFTER, src=(UPat(),), allow_any_len=True), lambda: True),
 ])+tensor_uop_spec+spec

 # ***** uop helpers *****
@@ -28,6 +28,7 @@ invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat)
 propagate_invalid = PatternMatcher([
  # this needs to be before symbolic so that 0*something_that_might_be_invalid doesnt become 0
  # propagate invalid, push it past children
+  (invalid_gate.cast(name="cast"), lambda i,x,cond,cast: x.cast(cast.dtype) if cast.dtype is not dtypes.index else None),
  *((invalid_gate.alu(op, UPat.var("y")).named("alu"), lambda cond,x,y,alu,i: cond.where(x.alu(alu.op,y), i))
    for op in GroupOp.Binary-GroupOp.Comparison),
  *((invalid_gate.alu(op, UPat.var("y")).named("alu"), lambda cond,x,y,alu,i: x.alu(alu.op,y)) for op in GroupOp.Comparison),
@@ -130,7 +131,7 @@ symbolic_simple = propagate_invalid + PatternMatcher([
 def lt_folding(x:UOp, c:int) -> UOp|None:
  p, np = partition(x.split_uop(Ops.ADD), lambda u: u.const_factor() == 1)
  if np and (d:=math.gcd(*[u.const_factor() for u in np], c)) > 1 and 0 <= sum(u.vmin for u in p) and sum(u.vmax for u in p) < d:
-    return cast(UOp, functools.reduce(operator.add, np).divides(d))<(c//d)
+    return cast(UOp, UOp.sum(*np).divides(d))<(c//d)
  return None

 def canonicalize_simplex(X:UOp) -> UOp|None:
@@ -144,7 +145,7 @@ def canonicalize_simplex(X:UOp) -> UOp|None:
      u = u.src[0]
    if not (u.op in GroupOp.Irreducible and u.vmin >= 0): return None
    ret.append(u)
-  return functools.reduce(operator.add, ret) if changed else None
+  return UOp.sum(*ret) if changed else None

 def cancel_divmod(d: UOp, x: UOp, y: UOp) -> UOp|None:
  # simple cancel div/mod case when the range of the numerator lies within a single denominator interval
@@ -167,7 +168,7 @@ def remove_nested_mod(m: UOp, x: UOp, y: UOp) -> UOp|None:
        something_changed = True
        u = u.src[0]
    new_xs.append(u)
-  new_x: UOp = functools.reduce(operator.add, new_xs)
+  new_x: UOp = UOp.sum(*new_xs)
  if something_changed and new_x.vmin>=0: return new_x % y
  return None

@@ -300,6 +301,7 @@ 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.cvar("y") * (UPat.var("x", dtype=dtypes.index) + UPat.cvar("c")), lambda x,y,c: (y*x)+(y*c)),  # -(x+c) -> -x + -c
  # ** where folding **
  (UPat.var("cond", dtype=dtypes.bool).logical_not().where(UPat.var("t"), UPat.var("f")), lambda cond, t, f: cond.where(f,t)
    if f.arg is not Invalid else None),
@@ -452,9 +454,9 @@ def simplify_valid(valid:UOp) -> UOp|None:
  something_changed = False
  valids = list(valid.split_uop(Ops.AND))
  for stmt in sorted(valids, key=lambda v: _valid_priority(v, valids)):
-    ret.append(uop_given_valid(functools.reduce(operator.and_, ret), stmt) if ret else stmt)
+    ret.append(uop_given_valid(UOp.prod(*ret), stmt) if ret else stmt)
    if ret[-1] is not stmt: something_changed = True
-  return functools.reduce(operator.and_, ret) if something_changed else None
+  return UOp.prod(*ret) if something_changed else None

 # ******** phase 3 is the complete symbolic, and deals with very complex things like loop rewriting and threefry transform ********

@@ -471,7 +473,7 @@ def reduce_mul_chain(r:UOp):

 def drop_and_clauses(cond:UOp, x:UOp, i:UOp) -> UOp|None:
  if not (dropped_clauses:=[c for c in cond.split_uop(Ops.AND) if not any(r in x.ranges for r in c.ranges)]): return None
-  return functools.reduce(operator.and_, [c for c in cond.split_uop(Ops.AND) if c not in dropped_clauses], UOp.const(dtypes.bool, True)).where(x, i)
+  return UOp.const(dtypes.bool, True).prod(*[c for c in cond.split_uop(Ops.AND) if c not in dropped_clauses]).where(x, i)
 pm_drop_and_clauses = PatternMatcher([(UPat.var("cond").where(UPat.var("x", dtype=dtypes.index), invalid_pat), drop_and_clauses)])

 def where_on_load(l, c1, buf, x):
@@ -483,7 +485,7 @@ def where_on_load(l, c1, buf, x):
    and not c.op_in_backward_slice_with_self(Ops.LOAD)]
  if not (removed:=moved_clauses+duplicate_clauses): return None
  # aditionally we can drop the clause on the where if it already exists in the load
-  remaining_clause = functools.reduce(operator.and_, [c for c in c1.split_uop(Ops.AND) if c not in removed], UOp.const(dtypes.bool, True))
+  remaining_clause = UOp.const(dtypes.bool, True).prod(*[c for c in c1.split_uop(Ops.AND) if c not in removed])
  return remaining_clause.where(UOp.load(buf.index(x.get_idx().valid(functools.reduce(operator.and_, moved_clauses, c2)), *l.src[1:])), 0)
 pm_move_where_on_load = PatternMatcher([
  (UPat.var("c1").where(UPat(Ops.LOAD, src=(UPat.var("buf").index(UPat.var("x")),), name="l"), 0), where_on_load),
@@ -41,11 +41,13 @@
  }
  ul {
    padding: 0;
-    opacity: 0.6;
    white-space: nowrap;
    cursor: pointer;
  }
-  ul.active {
+  ul > p {
+    opacity: 0.6;
+  }
+  ul.active > p {
    opacity: 1;
  }
  ul > ul {
@@ -54,8 +56,10 @@
  ul.expanded > ul {
    display: block;
  }
-  ul.disabled {
+  ul.disabled > p {
    opacity: 0.4;
+  }
+  ul.disabled {
    pointer-events: none;
  }
  label {
@@ -51,7 +51,7 @@ function addTags(root) {
  root.selectAll("text").data(d => [d]).join("text").text(d => d).attr("dy", "0.35em");
 }

-let [workerUrl, worker] = [null, null];
+let workerUrl = null, worker = null;
 async function initWorker() {
  const resp = await Promise.all(["/assets/dagrejs.github.io/project/dagre/latest/dagre.min.js","/js/worker.js"].map(u => fetch(u)));
  workerUrl = URL.createObjectURL(new Blob([(await Promise.all(resp.map((r) => r.text()))).join("\n")], { type: "application/javascript" }));
@@ -202,6 +202,8 @@ async function renderProfiler() {
  const canvasTop = rect(canvas).top;
  // color by key (name/device)
  const colorMap = new Map();
+  // map shapes by event key
+  const shapeMap = new Map();
  data = {tracks:new Map(), axes:{}};
  const heightScale = d3.scaleLinear().domain([0, tracePeak]).range([4,maxheight=100]);
  for (let i=0; i<layoutsLen; i++) {
@@ -216,10 +218,10 @@ async function renderProfiler() {
    if (eventType === EventTypes.TIMELINE) {
      const levelHeight = baseHeight-padding;
      const levels = [];
-      data.tracks.set(k, { shapes, visible, offsetY });
+      data.tracks.set(k, { shapes, visible, offsetY, pcolor:"#9ea2ad" });
      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};
+        const e = {name:strings[u32()], ref:optional(u32()), key:optional(u32()), st:u32(), dur:f32(), info:strings[u32()] || null};
        // find a free level to put the event
        let depth = levels.findIndex(levelEt => e.st >= levelEt);
        const et = e.st+Math.trunc(e.dur);
@@ -239,7 +241,18 @@ async function renderProfiler() {
          const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
          if (stepIdx !== -1) { ref.step = stepIdx; shapeRef = ref; }
        }
-        const arg = { tooltipText:colored(e.name).outerHTML+"\n"+formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...shapeRef };
+        const html = document.createElement("div");
+        html.appendChild(tabulate([["Name", colored(e.name)], ["Duration", formatTime(e.dur)], ["Start Time", formatTime(e.st)]]).node());
+        if (e.info != null) html.appendChild(document.createElement("p")).innerText = "\n"+e.info;
+        if (shapeRef != null) {
+          const p = html.appendChild(document.createElement("p"));
+          p.innerText = "\nView Codegen Rewrite"; p.style.cursor = "pointer";
+          p.onclick = () => setCtxWithHistory(shapeRef.ctx, shapeRef.step);
+        }
+        // tiny device events go straight to the rewrite rule
+        const key = k.startsWith("TINY") ? null : `${k}-${j}`;
+        const arg = { tooltipText:colored(e.name).outerHTML+"\n"+formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), html, key, ...shapeRef };
+        if (e.key != null) shapeMap.set(e.key, arg);
        // offset y by depth
        shapes.push({x:e.st, y:levelHeight*depth, width:e.dur, height:levelHeight, arg, label, fillColor });
      }
@@ -259,7 +272,7 @@ async function renderProfiler() {
          x += 1; y += nbytes; valueMap.set(ts, y);
        } else {
          const free = buf_shapes.get(key);
-          free.users = Array.from({ length: u32() }, () => strings[u32()]);
+          free.users = Array.from({ length: u32() }, () => ({shape:shapeMap.get(u32()), repr:strings[u32()], num:u8(), mode:u8()}));
          timestamps.push(ts); valueMap.set(ts, y);
          x += 1; y -= free.nbytes;
          free.x.push(x);
@@ -283,11 +296,13 @@ async function renderProfiler() {
        if (users != null) rows.push(["Users", users.length]);
        const info = html.appendChild(tabulate(rows).node());
        for (let u=0; u<users?.length; u++) {
-          const p = html.appendChild(document.createElement("p")); p.style.marginTop = "4px"; p.style.cursor = "pointer";
-          const name = users[u]; p.appendChild(colored(`[${u}] ${name}`));
-          p.onclick = () => {
-            const cid = ctxs.findIndex(c => c.name === name);
-            if (cid != null) setCtxWithHistory(cid-1);
+          const p = html.appendChild(document.createElement("p")); p.style.marginTop = "4px";
+          const { repr, num, mode, shape } = users[u]; p.appendChild(colored(`[${u}] ${repr} ${mode == 2 ? 'read+write' : mode == 1 ? 'write' : 'read'}@data${num}`));
+          const metadata = shape?.tooltipText?.split("\n").at(-1);
+          if (metadata != null) p.appendChild(document.createElement("span")).innerText = "\n"+metadata;
+          if (shape != null) {
+            p.style.cursor = "pointer";
+            p.onclick = () => focusShape(shape);
          }
        }
        const arg = {tooltipText:info.outerHTML, html, key:`${k}-${num}`};
@@ -313,7 +328,7 @@ async function renderProfiler() {
        sum.x.push(allX[i], allX[i+1]);
        const y = maxY.get(allX[i]); sum.y1.push(y, y); sum.y0.push(base0, base0);
      }
-      data.tracks.set(k, { shapes:[sum], visible, offsetY, height, peak, scaleFactor:maxheight*4/height, views:[[sum], shapes], valueMap });
+      data.tracks.set(k, { shapes:[sum], visible, offsetY, pcolor:"#c9a8ff", height, peak, scaleFactor:maxheight*4/height, views:[[sum], shapes], valueMap });
      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;
@@ -342,14 +357,13 @@ async function renderProfiler() {
    xscale.domain(visibleX);
    // draw shapes
    const paths = [];
-    for (const [_, { offsetY, shapes, visible, valueMap }] of data.tracks) {
+    for (const [_, { offsetY, shapes, visible, valueMap, pcolor }] of data.tracks) {
      visible.length = 0;
      for (const e of shapes) {
-        // generic polygon
-        if (e.width == null) {
+        const p = new Path2D();
+        if (e.width == null) { // generic polygon
          if (e.x[0]>et || e.x.at(-1)<st) continue;
          const x = e.x.map(xscale);
-          const p = new Path2D();
          p.moveTo(x[0], offsetY+e.y0[0]);
          for (let i=1; i<x.length; i++) {
            p.lineTo(x[i], offsetY+e.y0[i]);
@@ -360,32 +374,29 @@ async function renderProfiler() {
          for (let i=x.length-1; i>=0; i--) p.lineTo(x[i], offsetY+e.y1[i]);
          p.closePath();
          ctx.fillStyle = e.fillColor; ctx.fill(p);
-          if (focusedShape?.key && e.arg?.key === focusedShape.key) { paths.push(p); }
-          continue;
-        }
-        // contiguous rect
-        if (e.x>et || e.x+e.width<st) continue;
-        const x = xscale(e.x);
-        const y = offsetY+e.y;
-        const width = xscale(e.x+e.width)-x;
-        ctx.fillStyle = e.fillColor; ctx.fillRect(x, y, width, e.height);
-        visible.push({ y0:y, y1:y+e.height, x0:x, x1:x+width, arg:e.arg });
-        // add label
-        if (e.label == null) continue;
-        ctx.textAlign = "left";
-        ctx.textBaseline = "middle";
-        let labelX = x+2, labelWidth = 0;
-        const labelY = y+e.height/2;
-        for (const [i,l] of e.label.entries()) {
-          if (labelWidth+l.width+(i===e.label.length-1 ? 0 : ellipsisWidth)+2 > width) {
-            if (labelWidth !== 0) ctx.fillText("...", labelX, labelY);
-            break;
+        } else { // contiguous rect
+          if (e.x>et || e.x+e.width<st) continue;
+          const x = xscale(e.x);
+          const y = offsetY+e.y;
+          const width = xscale(e.x+e.width)-x;
+          p.rect(x, y, width, e.height);
+          visible.push({ y0:y, y1:y+e.height, x0:x, x1:x+width, arg:e.arg });
+          ctx.fillStyle = e.fillColor; ctx.fill(p);
+          // add label
+          let lw = 0;
+          const lx = x+2, ly = y+e.height/2;
+          for (let li=0; li<e.label?.length; li++) {
+            if (lw+e.label[li].width+(li===e.label.length-1 ? 0 : ellipsisWidth)+2 > width) {
+              if (lw>0) ctx.fillText("...", lx+lw, ly);
+              break;
+            }
+            ctx.textBaseline = "middle";
+            ctx.fillStyle = e.label[li].color;
+            ctx.fillText(e.label[li].st, lx+lw, ly);
+            lw += e.label[li].width;
          }
-          ctx.fillStyle = l.color;
-          ctx.fillText(l.st, labelX, labelY);
-          labelWidth += l.width;
-          labelX += l.width;
        }
+        if (focusedShape?.key && e.arg?.key === focusedShape.key) { paths.push([p, pcolor]); }
      }
    }
    // draw axes
@@ -415,7 +426,7 @@ async function renderProfiler() {
      drawLine(ctx, [x, x], [0, canvas.clientHeight], { color:m.color });
      ctx.fillText(m.name, x+2, 1);
    }
-    for (const p of paths) { ctx.lineWidth = 1.4; ctx.strokeStyle = "#c9a8ff"; ctx.stroke(p); }
+    for (const [p, color] of paths) { ctx.lineWidth = 1.4; ctx.strokeStyle = color; ctx.stroke(p); }
  }

  function resize() {
@@ -452,12 +463,15 @@ async function renderProfiler() {
    }
  }

+  function focusShape(shape) {
+    focusedShape = shape; render(zoomLevel);
+    return document.querySelector(".metadata").replaceChildren(shape?.html ?? "");
+  }
  canvas.addEventListener("click", e => {
    e.preventDefault();
    const foundRect = findRectAtPosition(e.clientX, e.clientY);
-    if (foundRect?.step != null) return setCtxWithHistory(foundRect.ctx, foundRect.step);
-    if (foundRect?.key != focusedShape?.key) { focusedShape = foundRect; render(zoomLevel); }
-    return document.querySelector(".metadata").replaceChildren(foundRect?.html ?? "");
+    if (foundRect?.step != null && foundRect?.key == null) { return setCtxWithHistory(foundRect.ctx, foundRect.step); }
+    if (foundRect?.key != focusedShape?.key) { focusShape(foundRect); }
  });

  canvas.addEventListener("mousemove", e => {
@@ -516,11 +530,6 @@ function codeBlock(st, language, { loc, wrap }={}) {
  return ret;
 }

-function appendTd(tr, value, unit=null) {
-  const fmt = (typeof value === "number" && !Number.isInteger(value)) ? value.toFixed(2) : value;
-  tr.appendChild(document.createElement("td")).innerText = unit == "us" ? formatTime(value) : fmt+(unit ?? "");
-}
-
 function setActive(e) {
  if (e == null) return;
  e.classList.add("active");
@@ -599,7 +608,8 @@ async function main() {
      for (const [j,u] of steps.entries()) {
        const inner = ul.appendChild(document.createElement("ul"));
        inner.id = `step-${i}-${j}`;
-        inner.innerText = `${u.name ?? u.loc[0].replaceAll("\\", "/").split("/").pop()+':'+u.loc[1]}`+(u.match_count ? ` - ${u.match_count}` : '');
+        const p = inner.appendChild(document.createElement("p"));
+        p.innerText = `${u.name}`+(u.match_count ? ` - ${u.match_count}` : '');
        inner.style.marginLeft = `${8*u.depth}px`;
        inner.onclick = (e) => {
          e.stopPropagation();
@@ -645,7 +655,7 @@ async function main() {
        tr.className = "main-row code-row";
        for (const [i,value] of r.entries()) {
          // string format scalar values
-          if (!Array.isArray(value)) appendTd(tr, value);
+          if (!Array.isArray(value)) tr.appendChild(document.createElement("td")).innerText = value;
          // display arrays in a bar graph
          else {
            const segmentsTd = tr.appendChild(document.createElement("td"));
@@ -697,8 +707,9 @@ async function main() {
    rewriteList.className = "rewrite-list";
    for (let s=0; s<=step.match_count; s++) {
      const ul = rewriteList.appendChild(document.createElement("ul"));
-      ul.innerText = s;
      ul.id = `rewrite-${s}`;
+      const p = ul.appendChild(document.createElement("p"));
+      p.innerText = s;
      ul.onclick = () => setState({ currentRewrite:s });
      ul.className = s > ret.length-1 ? "disabled" : s === currentRewrite ? "active" : "";
      if (s > 0 && s === currentRewrite) {
@@ -20,7 +20,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
               **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER_VIEW: "#E5EAFF",
               Ops.BLOCK: "#C4A484", Ops.BLOCKEND: "#C4A4A4", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.FUSE: "#FFa500",
               Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
-               Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e", Ops.SUBSTITUTE: "#ffff00"}
+               Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e", Ops.SUBSTITUTE: "#ffff00", Ops.AFTER: "#8A7866"}

 # VIZ API

@@ -56,7 +56,7 @@ def pystr(u:UOp, i:int) -> str:
    except Exception: pass
  return str(u)

-def uop_to_json(x:UOp) -> dict[int, dict]:
+def uop_to_json(x:UOp, ignore_indexing=False) -> dict[int, dict]:
  assert isinstance(x, UOp)
  graph: dict[int, dict] = {}
  excluded: set[UOp] = set()
@@ -64,13 +64,14 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
    # always exclude DEVICE/CONST/UNIQUE
    if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE} and u is not x: excluded.add(u)
    if u.op is Ops.VCONST and u.dtype.scalar() == dtypes.index and u is not x: excluded.add(u)
+    if u.dtype.scalar() is dtypes.index and ignore_indexing: excluded.update(u.backward_slice_with_self)
  for u in toposort:
    if u in excluded: continue
    argst = codecs.decode(str(u.arg), "unicode_escape")
    if u.op in GroupOp.Movement: argst = (mask_to_str if u.op in {Ops.SHRINK, Ops.PAD} else shape_to_str)(u.marg)
    label = f"{str(u.op).split('.')[1]}{(chr(10)+word_wrap(argst.replace(':', ''))) if u.arg is not None else ''}"
    if u.dtype != dtypes.void: label += f"\n{u.dtype}"
-    for idx,x in enumerate(u.src):
+    for idx,x in enumerate(u.src[:1] if u.op in {Ops.BUFFERIZE, Ops.INDEX} else u.src):
      if x in excluded:
        arg = f"{x.arg:g}" if x.op is Ops.CONST and dtypes.is_float(x.dtype) else f"{x.arg}"
        label += f"\n{x.op.name}{idx} {arg}" + (f" {x.src[0].op}" if len(x.src) else "")
@@ -97,14 +98,19 @@ def _reconstruct(a:int):
  return UOp(op, dtype, tuple(_reconstruct(s) for s in src), arg, *rest)

 def get_full_rewrite(ctx:TrackedGraphRewrite, i:int=0) -> Generator[GraphRewriteDetails, None, None]:
-  yield {"graph":uop_to_json(next_sink:=_reconstruct(ctx.sink)), "uop":pystr(next_sink,i), "changed_nodes":None, "diff":None, "upat":None}
+  next_sink = _reconstruct(ctx.sink)
+  # in the schedule graph we don't show indexing ops (unless it's in a kernel AST or rewriting dtypes.index sink)
+  ignore_indexing = trace.keys[i].display_name.startswith("Schedule") and not (ctx.name in {"kernel split"} or \
+      any(s.dtype is dtypes.index for s in next_sink.src+(next_sink,)))
+  yield {"graph":uop_to_json(next_sink, ignore_indexing), "uop":pystr(next_sink,i), "changed_nodes":None, "diff":None, "upat":None}
  replaces: dict[UOp, UOp] = {}
  for u0_num,u1_num,upat_loc,dur in tqdm(ctx.matches):
    replaces[u0:=_reconstruct(u0_num)] = u1 = _reconstruct(u1_num)
    try: new_sink = next_sink.substitute(replaces)
    except RuntimeError as e: new_sink = UOp(Ops.NOOP, arg=str(e))
    match_repr = f"# {dur*1e6:.2f} us\n"+printable(upat_loc)
-    yield {"graph":(sink_json:=uop_to_json(new_sink)), "uop":pystr(new_sink,i), "changed_nodes":[id(x) for x in u1.toposort() if id(x) in sink_json],
+    yield {"graph":(sink_json:=uop_to_json(new_sink, ignore_indexing)), "uop":pystr(new_sink,i),
+           "changed_nodes":[id(x) for x in u1.toposort() if id(x) in sink_json],
           "diff":list(difflib.unified_diff(pystr(u0,i).splitlines(),pystr(u1,i).splitlines())), "upat":(upat_loc, match_repr)}
    if not ctx.bottom_up: next_sink = new_sink

@@ -136,25 +142,30 @@ def flatten_events(profile:list[ProfileEvent]) -> Generator[tuple[Decimal, Decim
 # normalize event timestamps and attach kernel metadata
 def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int, scache:dict[str, int]) -> bytes|None:
  events:list[bytes] = []
-  exec_points:dict[str, dict] = {}
+  exec_points:dict[str, ProfilePointEvent] = {}
  for st,et,dur,e in dev_events:
-    if isinstance(e, ProfilePointEvent) and e.name == "exec": exec_points[e.key] = e.arg
+    if isinstance(e, ProfilePointEvent) and e.name == "exec": exec_points[e.arg["name"]] = e
    if dur == 0: continue
-    name, info = e.name, None
+    name, info, key = e.name, None, None
    if (ref:=ref_map.get(name)) is not None:
      name = ctxs[ref]["name"]
      if isinstance(p:=trace.keys[ref].ret, ProgramSpec) and (ei:=exec_points.get(p.name)) is not None:
-        info = f"{sym_infer(p.estimates.ops, ei['var_vals'])/(t:=dur*1e3):.2f} GFLOPS {sym_infer(p.estimates.mem, ei['var_vals'])/t:4.1f}"+ \
-               f"|{sym_infer(p.estimates.lds,ei['var_vals'])/t:.1f} GB/s\n{ei['metadata']}"
+        info = f"{sym_infer(p.estimates.ops, ei.arg['var_vals'])/(t:=dur*1e3):.2f} GFLOPS {sym_infer(p.estimates.mem, ei.arg['var_vals'])/t:4.1f}"+ \
+               f"|{sym_infer(p.estimates.lds,ei.arg['var_vals'])/t:.1f} GB/s\n{ei.arg['metadata']}"
+        key = ei.key
    elif isinstance(e.name, TracingKey):
      name = e.name.display_name
      ref = next((v for k in e.name.keys if (v:=ref_map.get(k)) is not None), None)
-    events.append(struct.pack("<IIIfI", enum_str(name, scache), option(ref), st-start_ts, dur, enum_str(info or "", scache)))
+    events.append(struct.pack("<IIIIfI", enum_str(name, scache), option(ref), option(key), st-start_ts, dur, enum_str(info or "", scache)))
  return struct.pack("<BI", 0, len(events))+b"".join(events) if events else None

 def encode_mem_free(key:int, ts:int, execs:list[ProfilePointEvent], scache:dict) -> bytes:
-  kernel_names = [enum_str(ei.key, scache) for ei in execs]
-  return struct.pack(f"<BIII{len(kernel_names)}I", 0, ts, key, len(kernel_names), *kernel_names)
+  ei_encoding:list[tuple[int, int, int, int]] = [] # <[u32, u32, u8, u8] [run id, display name, buffer number and mode (2 = r/w, 1 = w, 0 = r)]
+  for e in execs:
+    num = next(i for i,k in enumerate(e.arg["bufs"]) if k == key)
+    mode = 2 if (num in e.arg["inputs"] and num in e.arg["outputs"]) else 1 if (num in e.arg["outputs"]) else 0
+    ei_encoding.append((e.key, enum_str(e.arg["name"], scache), num, mode))
+  return struct.pack("<BIII", 0, ts, key, len(ei_encoding))+b"".join(struct.pack("<IIBB", *t) for t in ei_encoding)

 def mem_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int], dtype_size:dict[str, int],
               scache:dict[str, int]) -> bytes|None:
Author	SHA1	Message	Date
geohot	45032b5c96	fixups	2025-10-20 17:30:07 +08:00
geohot	6efc381825	two stage removal	2025-10-20 16:21:39 +08:00
George HotzandGitHub	b8a9cce783	replace NOOP with AFTER in reg init (#12804 ) * after op * fix tests * replace NOOP with AFTER in reg init * closer * or_after there * fix device * fix all renderers * better spec for after	2025-10-20 15:34:32 +08:00
qazalandGitHub	12fd2c9c7b	explicitly set ignore_indexing for schedule only (#12803 )	2025-10-20 13:11:57 +08:00
qazalandGitHub	734c99f722	viz: show indexing rewrites during run_rangeify (#12802 ) * viz: show indexing rewrites during run_rangeify * sinking index	2025-10-20 12:37:03 +08:00
George HotzandGitHub	2e9082e0bc	after op (#12801 ) * after op * fix tests	2025-10-20 12:27:56 +08:00
qazalandGitHub	339e6edb7d	viz: ui prereqs for hierarchical rewrites (#12799 )	2025-10-20 12:15:15 +08:00
wozeparrotandGitHub	357dac8425	feat: allow tuple indexing on uops (#12797 )	2025-10-19 19:11:05 -07:00
George Hotz GitHub qazal	ba593f7b98	don't render index (#12796 ) * don't render index * update to ignore_indexing --------- Co-authored-by: qazal <[email protected]>	2025-10-20 09:48:36 +08:00
geohot	cad3ada909	tinygpu: build with SIP off works	2025-10-20 09:11:09 +08:00
nimlgenandGitHub	9cd35deae7	amd: fix alignment + pointers for aql over usb (#12793 )	2025-10-19 23:55:57 +08:00
nimlgenandGitHub	59784a5972	amd: ensure ts is written (#12794 )	2025-10-19 23:55:49 +08:00
chenyuandGitHub	63a23dfe80	test step 0 in TestTrainingOnnxOps (#12790 ) and tighter rtol	2025-10-19 09:15:49 -04:00
chenyuandGitHub	e8158afd4b	update test_qlinear_add_round_half_to_even (#12789 ) this does not pass locally	2025-10-19 08:47:27 -04:00
Sieds LyklesandGitHub	1df9c7d7e7	reduce_collapse uses symbolic_flat (#12766 ) * sym->symbolic_flat * cast invalid drops invalid	2025-10-19 12:27:47 +02:00
Sieds LyklesandGitHub	fd6ef4801c	rangeify uses symbolic_flat (#12786 ) * symbolic_simple -> symbolic_flat * remove expected failures	2025-10-19 12:27:14 +02:00
geohot	89e7f2fa00	mmapeak: gfx1103 support	2025-10-19 16:57:28 +08:00
George HotzandGitHub	617614beb7	add mi350x support to mmapeak (#12784 )	2025-10-19 16:11:07 +08:00
qazalandGitHub	c8ef4b60f6	viz: share match tracing and TINY device profiler (#12783 ) * set a default name for the traces * set profile_matches + renames * profile_matches test * traces 4 steps total	2025-10-19 14:30:07 +08:00
chenyu GitHub Bruce Wayne	350a4754a9	Update openpilot models (#12780 ) * Update openpilot models * Update slower model * fix that --------- Co-authored-by: Bruce Wayne <[email protected]>	2025-10-18 20:32:35 -04:00
chenyuandGitHub	30ff84d050	update test_conv2d_ceildiv_edge_case (#12779 )	2025-10-18 16:43:32 -04:00
nimlgenandGitHub	442218266d	qcom: fix profiler (#12778 ) * qcom: fix profiler * this way	2025-10-19 01:27:59 +08:00
Harald SchäferandGitHub	addc54b96c	Simplify openpilot compile3.py (#12748 ) * Simpler compile3 * tests * remove default args * onnx file is still fp16 * self-test FP16 too * allow test disable * absurd tolerance * Just do latest * Try simplest * use later models * kernel count not relevant if speed is good * dead improts * Revert "dead improts" This reverts commit `f68c2cd15d`. * Revert "kernel count not relevant if speed is good" This reverts commit `0955ca4ee0`. * add back kernal count check on latest model	2025-10-18 10:12:22 -04:00
nimlgenandGitHub	037f6e8fa0	qcom: ioctl for 7xx (#12777 )	2025-10-18 20:33:14 +08:00
wozeparrotandGitHub	82f10cfe2e	feat: assert on bufferview math (#12772 )	2025-10-17 14:20:08 -07:00
chenyuandGitHub	fcdf4ab37e	remove a contiguous in LARS (#12770 )	2025-10-17 17:07:30 -04:00
nimlgenandGitHub	910d698b78	system: cleanup page sizes (#12771 ) * system: cleanup page sizes * ooops	2025-10-18 02:06:42 +08:00
George HotzandGitHub	062a6d68d7	test flash attention backward (#12762 ) * test flash attention backward * TODO: fix pcontig * end ranges * render colors * very big * multiout at every level * reset ending ranges * fix tests * ugh	2025-10-17 23:15:59 +08:00
George HotzandGitHub	33025b99f6	small changes from fa backward (#12769 )	2025-10-17 22:41:18 +08:00
chenyuandGitHub	e0d0d4372d	fix shape of m and v in onnx Adam with FUSE_OPTIM (#12768 ) value is still slightly off but that's not onnx specific	2025-10-17 10:32:41 -04:00
qazalandGitHub	bd662bea67	viz: light up program runs (#12764 ) * basics work * fix the color * light up program events * swap a with p * better	2025-10-17 19:33:18 +08:00
George HotzandGitHub	c9a3464f76	those decimals never mattered (#12760 ) * those decimals never mattered * this * improve debug * real substitute fixes pcontig * locals are different buffers	2025-10-17 17:16:24 +08:00
qazalandGitHub	0160f034d6	viz: show display name for copy runners (#12761 ) * viz: show display name for copy runners * more u32	2025-10-17 16:59:51 +08:00
qazalandGitHub	253d32b065	viz: add metadata to buffer user list (#12758 ) * simple failing test * encodings * test passing * key is deduped	2025-10-17 16:28:54 +08:00
George HotzandGitHub	935a60db72	bring back partial contig and flash attention (#12756 ) * bring back partial contig and flash attention * why not 2 * work * that * fix pcontig	2025-10-17 16:19:05 +08:00
Sieds LyklesandGitHub	f6bc620169	UOp.prod and UOp.sum methods (#12755 )	2025-10-17 10:02:01 +02:00
Sieds LyklesandGitHub	d1bb5c0426	slightly flatter symbolic (#12757 )	2025-10-17 09:58:45 +02:00
qazalandGitHub	5417e4b099	viz helper cleanups (#12754 )	2025-10-17 15:20:24 +08:00
qazalandGitHub	3196a7aae3	viz: pre reqs for lighting up programs (#12753 )	2025-10-17 15:03:21 +08:00
qazalandGitHub	dfb8f9fc9e	viz: annotate buffer mutability in the memory graph (#12750 )	2025-10-17 11:53:02 +08:00
Sieds LyklesandGitHub	79c2f1ae26	remove reduce_rangless and replace with reduce_unparented (#12749 )	2025-10-17 04:46:05 +02:00