cleanups

* STORE doesn't have ranges anymore

* fix

.github/workflows/autogen.yml

@@ -48,7 +48,7 @@ jobs:
         python3 -c "from tinygrad.runtime.autogen import opencl"
         python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
         python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
-        python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v13_0_12, smu_v14_0_2, fw"
+        python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v13_0_12, smu_v14_0_2, fw, navi_offsets, vega_offsets"
         python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
         python3 -c "from tinygrad.runtime.autogen import llvm"
         python3 -c "from tinygrad.runtime.autogen import webgpu"

docs/tensor/elementwise.md

@@ -68,6 +68,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.div
 ::: tinygrad.Tensor.idiv
 ::: tinygrad.Tensor.mod
+::: tinygrad.Tensor.fmod
 ::: tinygrad.Tensor.bitwise_xor
 ::: tinygrad.Tensor.bitwise_and
 ::: tinygrad.Tensor.bitwise_or

examples/mlperf/model_train.py

@@ -1446,6 +1446,10 @@ def train_llama3():
       idx = next(j for j, p in enumerate(optim.params) if p is w)
       optim.master_params[idx].assign((optim.master_params[idx] * w._inv_scale.reshape(-1, *([1]*(w.ndim-1)))).contiguous())
 
+  # realize everything here
+  if optim.master_params: Tensor.realize(*optim.master_params)
+  Tensor.realize(*optim.params, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax)
+
   @TinyJit
   def minibatch(tokens:Tensor):
     if is_dp: tokens = tokens.to(None).shard(device, 0)

examples/mlperf/models/flat_llama.py

@@ -158,14 +158,14 @@ class FlatTransformer:
 
     xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
     xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
-    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
     if getenv("HK_FLASH_ATTENTION"):
       from extra.thunder.amd.fa import flash_attention
       attn, *save = flash_attention(xq, xk, xv, is_causal=True)
       saves.extend(save)
     else:
-      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True)
-    attn = attn.transpose(1, 2).reshape(bsz, seqlen, -1)
+      xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
+      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
+    attn = attn.reshape(bsz, seqlen, -1)
 
     out, *ret = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo)
     new_amaxs.extend(ret[:1])

extra/amdpci/am_smi.py

@@ -64,7 +64,7 @@ def get_bar0_size(pcibus):
 
 class AMSMI(AMDev):
   def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
-    self.pcibus = pcibus
+    self.pcibus, self.devfmt = pcibus, pcibus
     self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
     self.pci_state = self.read_pci_state()
     if self.pci_state == "D0": self._init_from_d0()
@@ -91,6 +91,7 @@ class SMICtx:
     self.prev_lines_cnt = 0
     self.prev_terminal_width = 0
     self.prev_terminal_height = 0
+    self.prev_metrics = {}
 
     remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:", "Processing accelerators:"]
     lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
@@ -235,6 +236,29 @@ class SMICtx:
       case (13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.SocketPowerLimit)
       case _: return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX
 
+  def get_throttle_info(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12):
+        throttle_fields = [('ProchotResidencyAcc', 'Prochot'), ('PptResidencyAcc', 'PPT'),
+                           ('SocketThmResidencyAcc', 'Socket Thm'), ('VrThmResidencyAcc', 'VR Thm'), ('HbmThmResidencyAcc', 'HBM Thm')]
+        prev = self.prev_metrics.get(dev.pcibus)
+        active = []
+        if prev is not None:
+          acc_delta = metrics.AccumulationCounter - prev.AccumulationCounter
+          if acc_delta > 0:
+            for field, name in throttle_fields:
+              delta = getattr(metrics, field) - getattr(prev, field)
+              if delta > 0 and (pct := min(100, (delta * 100 + acc_delta // 2) // acc_delta)) > 0: active.append((name, pct))
+        return active
+      case _:
+        smu_mod = dev.smu.smu_mod
+        throttler_names = {getattr(smu_mod, a): a[len('THROTTLER_'):-len('_BIT')]
+                           for a in dir(smu_mod) if a.startswith('THROTTLER_') and a.endswith('_BIT')}
+        active = []
+        for i, pct in enumerate(metrics.SmuMetrics.ThrottlingPercentage):
+          if pct > 0: active.append((throttler_names.get(i, f"UNK_{i}"), int(pct)))
+        return active
+
   def get_mem_usage(self, dev):
     usage = 0
     pt_stack = [dev.mm.root_page_table]
@@ -281,6 +305,13 @@ class SMICtx:
                     + [f"MEM Activity {draw_bar(self.get_mem_activity(dev, metrics) / 100, activity_line_width)}"] \
                     + [f"MEM Usage    {draw_bar(mem_used / mem_total, activity_line_width, opt_text=mem_fmt)}"] \
 
+      throttle_info = self.get_throttle_info(dev, metrics)
+      if throttle_info:
+        throttle_text = colored(', '.join(f"{name} {pct}%" for name, pct in throttle_info), "red")
+      else:
+        throttle_text = colored("None", "green")
+      activity_line += [f"Throttle     {throttle_text}" + " " * (activity_line_width + 2)]
+
       temps_data, temps_data_compact = self.get_temps(dev, metrics), self.get_temps(dev, metrics, compact=True)
       temps_table = ["=== Temps (°C) ==="] + [f"{name:<16}: {color_temp(val)}" for name, val in temps_data.items()]
       temps_table_compact = ["Temps (°C):" + '/'.join([f"{color_temp(val)} {name}" for name, val in temps_data_compact.items()])]
@@ -324,6 +355,8 @@ class SMICtx:
 
       dev_content.append(device_line + activity_line + same_line([temps_table, power_table, frequency_table]))
 
+    self.prev_metrics = {dev.pcibus: m for dev, m in dev_metrics.items() if m is not None}
+
     raw_text = 'AM Monitor'.center(terminal_width) + "\n" + "=" * terminal_width + "\n\n"
     for i in range(0, len(dev_content), 2):
       if i + 1 < len(dev_content): raw_text += '\n'.join(same_line([dev_content[i], dev_content[i+1]], split=padding))

extra/gemm/cdna_asm_gemm.py

@@ -2628,21 +2628,24 @@ def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
 # ** FP8 GEMM custom kernel
 
 @functools.cache
-def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, X_s:UOp, W_s:UOp, *extra:UOp, dname:str) -> UOp:
-  # A is (batch, M, K), B is (N, K) transposed, X_s is x_scale, W_s is w_scale — kernel multiplies by both.
-  # extra is unused fwd inputs (e.g. grad_amax_state) plumbed through so the bwd can read them via kernel.src.
+def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str, scale_mode:int=3) -> UOp:
+  # scale_mode: 0=no scale, 1=x only, 2=w only, 3=both
+  n_scales = (1 if scale_mode & 1 else 0) + (1 if scale_mode & 2 else 0)
+  scales, extra = args[:n_scales], args[n_scales:]
   M, K = A.shape[0]*A.shape[1], A.shape[2]
   N, K2 = B.shape[(1 if B.ndim == 3 else 0):]
   assert K == K2, f"{A.shape} {B.shape}"
   block_size = 256
   threads = UOp.special(64 * 8, "lidx0")
   workgroups = UOp.special((M // block_size) * (N // block_size), "gidx0")
-  sink = UOp.sink(C.base, A.base, B.base, X_s.base, W_s.base, threads, workgroups,
+  sink_inputs = (C.base, A.base, B.base) + tuple(s.base for s in scales) + (threads, workgroups)
+  sink = UOp.sink(*sink_inputs,
                   arg=KernelInfo(f"hk_fp8_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
   kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
   src = (kittens_path/"gemm_fp8.cpp").read_text()
   lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
-                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}"]).compile_cached(src)
+                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
+                                 f"-DSCALE_MODE={scale_mode}"]).compile_cached(src)
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
                                UOp(Ops.BINARY, arg=lib)))
 
@@ -2699,8 +2702,7 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
 
 def custom_gemm_bw(gradient:UOp, kernel:UOp):
   inputs = kernel.src[1:]
-  # fp8 scaled gemm has 5 inputs (out, a, b, x_scale, w_scale) optionally plus grad_amax_state (6 total); plain gemm has 3
-  if len(inputs) >= 5:
+  if inputs[1].dtype == FP8_DTYPE:
     grad_amax_state = inputs[5] if len(inputs) == 6 else None
     out, a, b, s_x, s_w = inputs[:5]
     a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
@@ -2720,8 +2722,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
     # dgrad: uses g_scale * x_scale * w_scale
     grad_a = asm_gemm(g_fp8, b_t, x_scale=g_scale * s_x_t, w_scale=s_w_t)
     # wgrad: no w_scale
-    _one = Tensor(1.0, dtype=dtypes.float, device=a.device)
-    grad_b = asm_gemm(g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1), a_t.reshape(-1, a_t.shape[-1]), x_scale=g_scale * s_x_t, w_scale=_one)
+    grad_b = asm_gemm(g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1), a_t.reshape(-1, a_t.shape[-1]), x_scale=g_scale * s_x_t)
     # Attach the delayed-amax store effect (if any) to grad_a so realizing grads commits the amax update.
     ret = (None, grad_a.uop.after(store_effect), grad_b.uop, None, None)
     if len(inputs) == 6: ret = ret + (None,)
@@ -2774,11 +2775,11 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
   if arch.startswith("gfx950") and getenv("USE_ASM", 1):
     # fp8 gemm computes [email protected], kernel multiplies output by x_scale * w_scale before bf16 store
     if a.dtype == FP8_DTYPE:
-      _one = lambda: Tensor(1.0, dtype=dtypes.float, device=a.device)
-      xs = x_scale if x_scale is not None else _one()
-      ws = w_scale if w_scale is not None else _one()
+      scales = tuple(s for s in (x_scale, w_scale) if s is not None)
+      scale_mode = (1 if x_scale is not None else 0) | (2 if w_scale is not None else 0)
       extra = [grad_amax_state] if grad_amax_state is not None else []
-      out = Tensor.custom_kernel(out, a, b.T, xs, ws, *extra, fxn=functools.partial(custom_hk_fp8_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
+      fxn = functools.partial(custom_hk_fp8_gemm, dname=dname, scale_mode=scale_mode)
+      out = Tensor.custom_kernel(out, a, b.T, *scales, *extra, fxn=fxn, grad_fxn=custom_gemm_bw)[0]
     else:
       out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
   else:

extra/thunder/amd/fa.py

@@ -55,8 +55,6 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
   assert attn_mask is None, "attn_mask not supported"
   assert is_causal, "only causal attention supported"
 
-  xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
-
   B, N, H, D = xq.shape
   H_KV = xk.shape[2]
   assert D == 128, "only D=128 supported"
@@ -81,7 +79,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
   attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D), grad_fxn=grad)[:2]
 
-  return attn.transpose(1, 2), attn, l_vec
+  return attn, attn, l_vec
 
 @functools.cache
 def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):

extra/thunder/amd/gemm_fp8.cpp

@@ -93,7 +93,20 @@ constexpr int NUM_WARPS = 8;
 
 using G = kittens::group<NUM_WARPS>;
 
-__global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr, float *x_scale_ptr, float *w_scale_ptr) {
+// scale_mode: 0=no scale, 1=x only, 2=w only, 3=both
+#ifndef SCALE_MODE
+#define SCALE_MODE 3
+#endif
+
+__global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr
+#if SCALE_MODE == 1
+    , float *x_scale_ptr
+#elif SCALE_MODE == 2
+    , float *w_scale_ptr
+#elif SCALE_MODE == 3
+    , float *x_scale_ptr, float *w_scale_ptr
+#endif
+) {
     constexpr int M = GEMM_M, N = GEMM_N, K = GEMM_K;
 
     kittens::gl<fp8e4m3, 1, 1, M, K> A{A_ptr, nullptr, nullptr, nullptr, nullptr};
@@ -333,11 +346,25 @@ __global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_pt
     }
 
     // apply x_scale * w_scale before bf16 store to prevent overflow
+#if SCALE_MODE == 1
+    float scale = *x_scale_ptr;
+    mul(cA, cA, scale);
+    mul(cB, cB, scale);
+    mul(cC, cC, scale);
+    mul(cD, cD, scale);
+#elif SCALE_MODE == 2
+    float scale = *w_scale_ptr;
+    mul(cA, cA, scale);
+    mul(cB, cB, scale);
+    mul(cC, cC, scale);
+    mul(cD, cD, scale);
+#elif SCALE_MODE == 3
     float scale = *x_scale_ptr * *w_scale_ptr;
     mul(cA, cA, scale);
     mul(cB, cB, scale);
     mul(cC, cC, scale);
     mul(cD, cD, scale);
+#endif
 
     store(C, cA, {0, 0, block_row * WARPS_ROW * 2 + warp_m, block_col * WARPS_COL * 2 + warp_n});
     store(C, cB, {0, 0, block_row * WARPS_ROW * 2 + warp_m, block_col * WARPS_COL * 2 + WARPS_COL + warp_n});

test/backend/test_custom_kernel.py

@@ -321,8 +321,23 @@ class TestCustomKernel(unittest.TestCase):
     self.assertEqual(GlobalCounters.kernel_count, 2)
     self.assertEqual(z.tolist(), x.add(2).tolist())
 
+  @unittest.expectedFailure
   def test_custom_kernel_sched_copy(self): self.test_custom_kernel_sched(use_custom=True)
 
+  @unittest.expectedFailure
+  def test_sliced_buffer_function(self):
+    x = Tensor.arange(32).reshape(8, 4).realize()
+    from tinygrad import function
+    @function(precompile=True)
+    def run(x:Tensor) -> Tensor:
+      y = Tensor.invalids(*x.shape, dtype=x.dtype)
+      return Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
+    GlobalCounters.reset()
+    y = run(x[0]).realize()
+    # it's copying the input and the output
+    self.assertEqual(GlobalCounters.kernel_count, 1)
+    self.assertEqual(y.tolist(), [1, 2, 3, 4])
+
 class TestUOpReduce(unittest.TestCase):
   def test_uop_sum(self):
     a = Tensor([1.0, 2, 3, 4, 5])

test/backend/test_linearizer.py

@@ -14,7 +14,7 @@ from tinygrad.renderer.cstyle import CUDARenderer
 from test.helpers import replace_opts
 MOCKGPU = DEV.interface.startswith("MOCK")
 
-from tinygrad.uop.ops import print_uops # noqa: F401 # pylint: disable=unused-import
+from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
 
 class TestLinearizer(unittest.TestCase):
   def test_arg_dedup(self):

test/backend/test_ops.py

@@ -636,6 +636,17 @@ class TestOps(unittest.TestCase):
         helper_test_op(None, lambda x: 100%x, forward_only=True, vals=[va])
         helper_test_op(None, lambda x: 100.5%x, forward_only=True, vals=[va])
 
+  def test_fmod(self):
+    a = [-4, 7, 5, 4, -7, 8, -9]
+    b = [2, -3, 8, -2, 3, 5, -5]
+    for float_a in [True, False]:
+      for float_b in [True, False]:
+        va = [float(ai) for ai in a] if float_a else a
+        vb = [float(bi) for bi in b] if float_b else b
+        helper_test_op(None, lambda x,y: x.fmod(y), forward_only=True, vals=[va, vb])
+        helper_test_op(None, lambda x: x.fmod(2), forward_only=True, vals=[va])
+        helper_test_op(None, lambda x: x.fmod(3.5), forward_only=True, vals=[va])
+
   def test_mul_naninf(self):
     helper_test_op([(45,65)], lambda x: x*math.inf)
     helper_test_op([(45,65)], lambda x: x*-math.inf)

test/backend/test_profiler.py

@@ -51,11 +51,11 @@ class TestProfiler(unittest.TestCase):
     TestProfiler.runtime = get_runtime(TestProfiler.d0.device, TestProfiler.prg)
     TestProfiler.b.uop.buffer.allocate()
 
-  def test_profile_kernel_run(self):
+  def test_profile_kernel_run(self, wait=False):
     runner_name = TestProfiler.runtime.name
     with helper_collect_profile(TestProfiler.d0) as profile:
       gs, ls = TestProfiler.prg.arg.launch_dims({})
-      TestProfiler.runtime(TestProfiler.b.uop.buffer._buf, TestProfiler.a.uop.buffer._buf, global_size=gs, local_size=ls)
+      TestProfiler.runtime(TestProfiler.b.uop.buffer._buf, TestProfiler.a.uop.buffer._buf, global_size=gs, local_size=ls, wait=wait)
 
     profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
     kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent)]
@@ -63,6 +63,9 @@ class TestProfiler(unittest.TestCase):
     assert kernel_runs[0].name == runner_name, "kernel name is not correct"
     assert _dev_base(kernel_runs[0].device) == kernel_runs[0].device, "kernel should not be on a sub-device"
 
+  def test_profile_kernel_run_wait(self):
+    self.test_profile_kernel_run(wait=True)
+
   def test_profile_copyin(self):
     buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
 

test/backend/test_randomness.py

@@ -361,7 +361,7 @@ class TestRandomness(unittest.TestCase):
     _check_with_torch(w=[0.231, 0., 1., 0.5], num_samples=300, replacement=True)
     _check_with_torch(w=[[0.2, 0.8]], num_samples=300, replacement=True)  # 2D but only 1 row
     _check_with_torch(w=[[0.453, 0., 1., 0.81], [0.1, 0.8, 0., 0.1]], num_samples=300, replacement=True)
-    # no-replacement isn't supported, unless taking only one sample
+    # no-replacement
     w = [0.1, 0.9]
     self.assertRaises(AssertionError, lambda: Tensor(w).multinomial(100, replacement=False))
 
@@ -372,6 +372,23 @@ class TestRandomness(unittest.TestCase):
     torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(1000)]
     self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_samples), lambda _: torch.tensor(torch_samples)))
 
+    w = list(range(32))
+    s1 = Tensor(w).multinomial(5, replacement=False).numpy()
+    self.assertEqual(len(set(s1.tolist())), 5)
+    s2 = Tensor(w).multinomial(5, replacement=False).numpy()
+    self.assertFalse(np.array_equal(s1, s2))
+    full = Tensor(w).multinomial(len(w), replacement=False).numpy()
+    self.assertEqual(sorted(full.tolist()), w)
+
+    w = [0.1, 0.2, 0.3, 0.4]
+    @TinyJit
+    def sample_three(): return Tensor(w).multinomial(3, replacement=False).realize()
+
+    tiny_draws = np.array([sample_three().numpy() for _ in range(1000)])
+    torch_draws = np.array([torch.tensor(w).multinomial(3, replacement=False).numpy() for _ in range(1000)])
+    for pos in range(3):
+      self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_draws[:, pos]), lambda _: torch.tensor(torch_draws[:, pos])))
+
   @unittest.skip("this test is flaky")
   def test_multinomial_counterexample(self):
     tiny_res = Tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True)

test/device/test_metal.py

@@ -50,19 +50,6 @@ kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgr
       compiled = compiled[:40] # corrupt the compiled program
       MetalProgram(device, "r_5", compiled)
 
-  def test_wait_skips_in_flight(self):
-    device = MetalDevice("metal")
-    compiled = MetalCompiler().compile("""
-#include <metal_stdlib>
-kernel void noop(uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {}
-""")
-    prg = MetalProgram(device, "noop", compiled)
-    self.assertIsInstance(prg(wait=True), float)
-    self.assertEqual(device.mtl_buffers_in_flight, [])
-    self.assertIsNone(prg(wait=False))
-    self.assertEqual(len(device.mtl_buffers_in_flight), 1)
-    device.synchronize()
-
   def test_free(self):
     size = 2**16
     device = Device['METAL']

test/external/external_fuzz_beam_timeout_recovery.py

@@ -3,7 +3,7 @@
 Stress test for beam timeout + device recovery on AM devices.
 
 Usage:
-  DEV=AMD python test/external/external_test_beam_timeout_recovery.py
+  DEV=AMD python test/external/external_fuzz_beam_timeout_recovery.py
 """
 from tinygrad import Tensor, Device
 from tinygrad.helpers import Context

test/mockgpu/amd/emu.py

@@ -375,22 +375,24 @@ def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32,
   """Conditional memory store with sub-word support. Returns list of store UOps."""
   adt = dtypes.uint64 if addr_bits == 64 else dtypes.uint32
   word_addr = addr >> UOp.const(adt, 2)
-  idx = mem.index(word_addr.cast(dtypes.int), active)
-  if data_bits == 32: return [idx.store(active.where(_to_u32(val), idx))]
+  bidx = mem.index(word_addr.cast(dtypes.int), ptr=True)
+  if data_bits == 32: return [UOp(Ops.STORE, dtypes.void, (bidx, _to_u32(val), active))]
   # Sub-word store: read-modify-write with mask
+  cur = bidx.load(active, _c(0, dtypes.uint32))
   byte_pos = addr.cast(dtypes.uint32) & _c(3)
   byte_shift = byte_pos * _c(8)
   val_u32, size_mask = val.cast(dtypes.uint32), _c(0xFF if data_bits == 8 else 0xFFFF)
   mask = size_mask << byte_shift
-  new_word = (idx & (mask ^ _c(0xFFFFFFFF))) | ((val_u32 & size_mask) << byte_shift)
-  if data_bits == 8: return [idx.store(active.where(new_word, idx))]
+  new_word = (cur & (mask ^ _c(0xFFFFFFFF))) | ((val_u32 & size_mask) << byte_shift)
+  if data_bits == 8: return [UOp(Ops.STORE, dtypes.void, (bidx, new_word, active))]
   # 16-bit cross-word case: byte_pos == 3 means value spans two words
   is_cross = byte_pos.eq(_c(3))
-  cross_word0 = (idx & _c(0x00FFFFFF)) | ((val_u32 & _c(0xFF)) << _c(24))
-  store0 = idx.store(active.where(is_cross.where(cross_word0, new_word), idx))
-  next_idx = mem.index((word_addr + UOp.const(adt, 1)).cast(dtypes.int), active & is_cross)
-  cross_word1 = (next_idx & _c(0xFFFFFF00)) | ((val_u32 >> _c(8)) & _c(0xFF))
-  return [store0, next_idx.store((active & is_cross).where(cross_word1, next_idx))]
+  cross_word0 = (cur & _c(0x00FFFFFF)) | ((val_u32 & _c(0xFF)) << _c(24))
+  store0 = UOp(Ops.STORE, dtypes.void, (bidx, is_cross.where(cross_word0, new_word), active))
+  next_bidx = mem.index((word_addr + UOp.const(adt, 1)).cast(dtypes.int), ptr=True)
+  next_cur = next_bidx.load(active & is_cross, _c(0, dtypes.uint32))
+  cross_word1 = (next_cur & _c(0xFFFFFF00)) | ((val_u32 >> _c(8)) & _c(0xFF))
+  return [store0, UOp(Ops.STORE, dtypes.void, (next_bidx, cross_word1, active & is_cross))]
 
 def _mem_store_bytes(mem: UOp, addr: UOp, val: UOp, active: UOp, data_bits: int = 32) -> list[UOp]:
   """Store to byte-addressable memory (scratch). addr is byte offset, mem is uint8 buffer."""
@@ -398,7 +400,8 @@ def _mem_store_bytes(mem: UOp, addr: UOp, val: UOp, active: UOp, data_bits: int
   val_u32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
   for i in range(data_bits // 8):
     byte_val = (val_u32 >> UOp.const(dtypes.uint32, i * 8)) & UOp.const(dtypes.uint32, 0xFF)
-    stores.append(mem.index((addr + UOp.const(dtypes.uint64, i)).cast(dtypes.int), active).store(byte_val.cast(dtypes.uint8)))
+    bidx = mem.index((addr + UOp.const(dtypes.uint64, i)).cast(dtypes.int), ptr=True)
+    stores.append(UOp(Ops.STORE, dtypes.void, (bidx, byte_val.cast(dtypes.uint8), active)))
   return stores
 
 def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode_vars: dict | None = None, op_name: str = "") -> dict[int, UOp]:
@@ -516,14 +519,15 @@ class _Ctx:
   # Dynamic register access (takes UOp index instead of int)
   def rsgpr_dyn(self, reg: UOp, valid: UOp | None = None) -> UOp:
     """Read SGPR with dynamic register index."""
-    if valid is not None: return self.sgpr.index(reg.cast(dtypes.int), valid, ptr=True).load()
+    if valid is not None: return self.sgpr.index(reg.cast(dtypes.int), ptr=True).load(valid, _c(0, dtypes.uint32))
     return self.sgpr.index(reg.cast(dtypes.int), ptr=True).load()
 
   def wsgpr_dyn(self, reg: UOp, val: UOp) -> UOp:
     """Write SGPR with dynamic register index. On RDNA, index 124 = NULL (writes discarded). On CDNA, index 124 = M0 (read/write)."""
     # RDNA: NULL (124) discards writes. CDNA: M0 (124) is writable.
     valid = None if self.wave_size == 64 else reg.ne(_c(124))
-    return self.sgpr.index(reg.cast(dtypes.int), valid).store(val.cast(dtypes.uint32))
+    bidx = self.sgpr.index(reg.cast(dtypes.int), ptr=True)
+    return UOp(Ops.STORE, dtypes.void, (bidx, val.cast(dtypes.uint32))+((valid,) if valid is not None else ()))
 
   def wmask(self, reg: UOp, val: UOp) -> list[UOp]:
     """Write a lane mask (VCC/EXEC). Splits into lo/hi for wave64."""
@@ -540,24 +544,26 @@ class _Ctx:
   def rvgpr_dyn(self, reg: UOp, lane: UOp, valid: UOp | None = None) -> UOp:
     """Read VGPR with dynamic register index."""
     idx = reg.cast(dtypes.int) * _c(self.wave_size, dtypes.int) + lane.cast(dtypes.int)
-    return self.vgpr.index(idx, valid, ptr=True).load() if valid is not None else self.vgpr.index(idx, ptr=True).load()
+    if valid is not None: return self.vgpr.index(idx, ptr=True).load(valid, _c(0, dtypes.uint32))
+    return self.vgpr.index(idx, ptr=True).load()
 
   def wvgpr_dyn(self, reg: UOp, lane: UOp, val: UOp, exec_mask: UOp, after: UOp | None = None) -> UOp:
     """Write VGPR with dynamic register index."""
     buf = self.vgpr.after(after) if after is not None else self.vgpr
     offset = reg.cast(dtypes.int) * _c(self.wave_size, dtypes.int) + lane.cast(dtypes.int)
-    return buf.index(offset, _lane_active(exec_mask, lane)).store(val.cast(dtypes.uint32))
+    return UOp(Ops.STORE, dtypes.void, (buf.index(offset, ptr=True), val.cast(dtypes.uint32), _lane_active(exec_mask, lane)))
 
   def raccvgpr_dyn(self, reg: UOp, lane: UOp, valid: UOp | None = None) -> UOp:
     """Read ACCVGPR with dynamic register index (CDNA only)."""
     idx = reg.cast(dtypes.int) * _c(self.wave_size, dtypes.int) + lane.cast(dtypes.int)
-    return self.accvgpr.index(idx, valid, ptr=True).load() if valid is not None else self.accvgpr.index(idx, ptr=True).load()
+    if valid is not None: return self.accvgpr.index(idx, ptr=True).load(valid, _c(0, dtypes.uint32))
+    return self.accvgpr.index(idx, ptr=True).load()
 
   def waccvgpr_dyn(self, reg: UOp, lane: UOp, val: UOp, exec_mask: UOp, after: UOp | None = None) -> UOp:
     """Write ACCVGPR with dynamic register index (CDNA only)."""
     buf = self.accvgpr.after(after) if after is not None else self.accvgpr
     offset = reg.cast(dtypes.int) * _c(self.wave_size, dtypes.int) + lane.cast(dtypes.int)
-    return buf.index(offset, _lane_active(exec_mask, lane)).store(val.cast(dtypes.uint32))
+    return UOp(Ops.STORE, dtypes.void, (buf.index(offset, ptr=True), val.cast(dtypes.uint32), _lane_active(exec_mask, lane)))
 
   def rsrc_dyn(self, off: UOp, lane: UOp | None, bits: int = 32, literal: UOp | None = None, is_f64: bool = False, do_cast: bool = True) -> UOp:
     """Read source operand with dynamic offset. Handles SGPR/inline constants (<256), VGPR (>=256).
@@ -713,7 +719,7 @@ class _Ctx:
         old = self.vgpr.index(val[0].cast(dtypes.int), ptr=True).load()
         new_val = _set_bits(old, _val_to_bits(val[1]), width, lo_bit).cast(dtypes.uint32)
         active = _lane_active(exec_mask, lane)
-        raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].cast(dtypes.int), active).store(new_val)))
+        raw_stores.append(('vgpr_direct', UOp(Ops.STORE, dtypes.void, (self.vgpr.index(val[0].cast(dtypes.int), ptr=True), new_val, active))))
         continue
       if 'D0' in dest and '[laneId]' in dest:
         old_vcc = self.rmask(_c(VCC_LO.offset))
@@ -1847,15 +1853,16 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
     if data_bits < 32:
       # Sub-dword LDS write: read-modify-write within the uint32 slot
       word_addr = (addr >> addr_shift).cast(dtypes.int)
-      idx = mem.index(word_addr, active)
+      bidx = mem.index(word_addr, ptr=True)
+      cur = bidx.load(active, _c(0, dtypes.uint32))
       byte_pos = addr.cast(dtypes.uint32) & _c(3)
       byte_shift = byte_pos * _c(8)
       size_mask = _c(0xFF if data_bits == 8 else 0xFFFF)
       mask = size_mask << byte_shift
-      new_word = (idx & (mask ^ _c(0xFFFFFFFF))) | ((val.cast(dtypes.uint32) & size_mask) << byte_shift)
-      return idx.store(active.where(new_word, idx))
-    idx = mem.index((addr >> addr_shift).cast(dtypes.int))
-    return idx.store(active.where(val, idx.load()))
+      new_word = (cur & (mask ^ _c(0xFFFFFFFF))) | ((val.cast(dtypes.uint32) & size_mask) << byte_shift)
+      return UOp(Ops.STORE, dtypes.void, (bidx, new_word, active))
+    bidx = mem.index((addr >> addr_shift).cast(dtypes.int), ptr=True)
+    return UOp(Ops.STORE, dtypes.void, (bidx, val, active))
 
   def make_srcs(lane: UOp) -> dict:
     addr = make_addr(lane)
@@ -2005,17 +2012,18 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
       word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
       val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64), ptr=True).load(), _c(0))
       lds_idx = ((lds_addr + _c(i * 4)) >> _c(2)).cast(dtypes.int)
-      stores.append(ctx.lds.index(lds_idx, active).store(active.where(val, ctx.lds.index(lds_idx, active))))
+      bidx = ctx.lds.index(lds_idx, ptr=True)
+      stores.append(UOp(Ops.STORE, dtypes.void, (bidx, val, active)))
   elif is_store:
     for i in range(n_dwords):
       word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
-      idx = mem.index(word_addr.cast(dtypes.int64), in_bounds)
+      idx = mem.index(word_addr.cast(dtypes.int64), ptr=True)
       val = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(vdata + _c(i), lane)
-      stores.append(idx.store(in_bounds.where(_to_u32(val), idx)))
+      stores.append(UOp(Ops.STORE, dtypes.void, (idx, _to_u32(val), in_bounds)))
   else:
     for i in range(n_dwords):
       word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
-      val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64), in_bounds, ptr=True).load(), _c(0))
+      val = mem.index(word_addr.cast(dtypes.int64), ptr=True).load(in_bounds, _c(0, dtypes.uint32))
       stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(vdata + _c(i), lane, val, exec_mask))
   return UOp.sink(UOp.group(*stores).end(lane), *ctx.inc_pc())
 

test/mockgpu/amd/pcode.py

@@ -828,28 +828,30 @@ class Parser:
     assert mem is not None, "memory load requires _vmem or _lds"
     adt = dtypes.uint64 if addr.dtype == dtypes.uint64 else dtypes.uint32
     active = self.vars.get('_active')
-    gate = (active,) if active is not None else ()
+    # gate now lives on LOAD; helper to construct gated load with 0 alt
+    def _gload(bidx, dtype):
+      return bidx.load(active, _const(dtype.base, 0)) if active is not None else bidx.load()
     byte_mem = mem.dtype.base == dtypes.uint8
     if byte_mem:
       idx = addr.cast(dtypes.int)
       if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
         val = _u32(0).cast(dtypes.uint64)
-        for i in range(8): val = val | (mem.index(idx + _const(dtypes.int, i), *gate, ptr=True).load().cast(dtypes.uint64) << _u64(i * 8))
+        for i in range(8): val = val | (_gload(mem.index(idx + _const(dtypes.int, i), ptr=True), mem.dtype).cast(dtypes.uint64) << _u64(i * 8))
       elif dt in (dtypes.uint8, dtypes.int8):
-        val = mem.index(idx, *gate, ptr=True).load().cast(dt)
+        val = _gload(mem.index(idx, ptr=True), mem.dtype).cast(dt)
       elif dt in (dtypes.uint16, dtypes.int16, dtypes.short):
-        lo = mem.index(idx, *gate, ptr=True).load().cast(dtypes.uint32)
-        hi = mem.index(idx + _const(dtypes.int, 1), *gate, ptr=True).load().cast(dtypes.uint32)
+        lo = _gload(mem.index(idx, ptr=True), mem.dtype).cast(dtypes.uint32)
+        hi = _gload(mem.index(idx + _const(dtypes.int, 1), ptr=True), mem.dtype).cast(dtypes.uint32)
         val = (lo | (hi << _u32(8))).cast(dt)
       else:
         val = _u32(0)
-        for i in range(4): val = val | (mem.index(idx + _const(dtypes.int, i), *gate, ptr=True).load().cast(dtypes.uint32) << _u32(i * 8))
+        for i in range(4): val = val | (_gload(mem.index(idx + _const(dtypes.int, i), ptr=True), mem.dtype).cast(dtypes.uint32) << _u32(i * 8))
     else:
       idx = (addr >> _const(addr.dtype, 2)).cast(dtypes.int)
-      val = mem.index(idx, *gate)
+      val = _gload(mem.index(idx, ptr=True), mem.dtype)
       if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
         idx2 = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int)
-        val = val.cast(dtypes.uint64) | (mem.index(idx2, *gate).cast(dtypes.uint64) << _u64(32))
+        val = val.cast(dtypes.uint64) | (_gload(mem.index(idx2, ptr=True), mem.dtype).cast(dtypes.uint64) << _u64(32))
       elif dt in (dtypes.uint8, dtypes.int8): val = (val >> ((addr & _const(adt, 3)).cast(dtypes.uint32) * _u32(8))) & _u32(0xFF)
       elif dt in (dtypes.uint16, dtypes.int16):
         val = (val >> (((addr >> _const(adt, 1)) & _const(adt, 1)).cast(dtypes.uint32) * _u32(16))) & _u32(0xFFFF)
@@ -862,7 +864,7 @@ class Parser:
         idx_native = (addr >> _const(adt, 2)).cast(dtypes.int64)
         idx_hi_native = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int64)
         safe_idx_hi = is_unaligned.where(idx_hi_native, idx_native)
-        hi = mem.index(safe_idx_hi, *gate)
+        hi = _gload(mem.index(safe_idx_hi, ptr=True), mem.dtype)
         combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
         val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(dtypes.uint64, 8))).cast(dtypes.uint32), val)
     return _cast_to(val, dt)

test/null/test_simplify_valid_idx.py

@@ -49,7 +49,7 @@ class TestValidIdxSimplification(unittest.TestCase):
   def check(self, load, sidx, svalid, extra=()):
     with Context(NOOPT=1, SPEC=0):
       load = full_rewrite_to_sink(UOp.sink(load, *extra)).src[0]
-    idx, valid = load.src[0].src[1], load.src[0].src[2]
+    idx, valid = load.src[0].src[1], load.src[1]
     check_uop_against_string(self, idx, sidx)
     check_uop_against_string(self, valid, svalid)
 
@@ -225,9 +225,11 @@ class TestImageSimplification(unittest.TestCase):
     check_uop_against_string(self, idx0, sidx0)
     check_uop_against_string(self, idx1, sidx1)
     if svalid is not None:
-      check_uop_against_string(self, load.src[0].src[2], svalid)
+      check_uop_against_string(self, load.src[1], svalid)
     else:
-      self.assertEqual(len(load.src[0].src), 2, "svalid is None but load still has a valid")
+      # gate is at LOAD.src[1] when present; if simplified away, src[1] should not be bool
+      self.assertFalse(len(load.src) >= 2 and load.src[1].dtype.scalar() == dtypes.bool,
+                       "svalid is None but load still has a valid")
 
   def test_idx_gt_c(self):
     # (idx1 < c+1).ne(True) ? (..., idx1-1+c) : 0 can drop the valid
@@ -512,18 +514,34 @@ class TestUnfoldableImage(unittest.TestCase):
       self.assertEqual(res.src[0].src[0].dtype, dtypes.float.ptr(400))
 
 class TestDropTrueGate(unittest.TestCase):
-  def test_drop_true_gate_on_index(self):
-    # test that INDEX with a constant True gate gets simplified to drop the gate
+  def test_drop_true_gate_on_load(self):
+    # test that LOAD with a constant True gate gets simplified to drop the gate
     from tinygrad.codegen.late.devectorizer import load_store_indexing
     from tinygrad.uop.ops import graph_rewrite
     buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
     idx = UOp.const(dtypes.weakint, 0)
     true_gate = UOp.const(dtypes.bool, True)
-    index_with_gate = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx, true_gate))
+    bidx = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx))
+    load = UOp(Ops.LOAD, dtypes.int, (bidx, true_gate))
     # apply the optimization
-    result = graph_rewrite(index_with_gate, load_store_indexing)
-    # the True gate should be dropped (INDEX should only have 2 sources)
-    self.assertEqual(len(result.src), 2, "True gate should be dropped from INDEX")
+    result = graph_rewrite(load, load_store_indexing)
+    # the True gate should be dropped (LOAD should only have 1 source)
+    self.assertEqual(len(result.src), 1, "True gate should be dropped from LOAD")
+
+  def test_drop_true_gate_on_store(self):
+    # test that STORE with a constant True gate gets simplified to drop the gate
+    from tinygrad.codegen.late.devectorizer import load_store_indexing
+    from tinygrad.uop.ops import graph_rewrite
+    buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
+    idx = UOp.const(dtypes.weakint, 0)
+    val = UOp.const(dtypes.int, 42)
+    true_gate = UOp.const(dtypes.bool, True)
+    bidx = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx))
+    store = UOp(Ops.STORE, dtypes.void, (bidx, val, true_gate))
+    # apply the optimization
+    result = graph_rewrite(store, load_store_indexing)
+    # the True gate should be dropped (STORE should only have 2 sources)
+    self.assertEqual(len(result.src), 2, "True gate should be dropped from STORE")
 
 class TestRangeShrink(unittest.TestCase):
   def get_ranges(self, sink):

test/null/test_uop_graph.py

@@ -428,7 +428,8 @@ class TestUOpGraph(unittest.TestCase):
     uops = to_uops_list([w, red])
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg==5
+      # alt is at src[2] in new gated LOAD shape (idx, gate, alt)
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[2].arg==5
 
   def test_where_on_gated_load_folds_swapped_branches(self):
     ridx0 = UOp.range(100, 0)
@@ -438,7 +439,7 @@ class TestUOpGraph(unittest.TestCase):
     uops = to_uops_list([w])
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD: assert u.src[1].arg==5
+      if u.op is Ops.LOAD: assert u.src[2].arg==5
 
   def test_where_on_gated_load_with_cast(self):
     ridx0 = UOp.range(100, 0)
@@ -451,7 +452,7 @@ class TestUOpGraph(unittest.TestCase):
     uops = to_uops_list([w, red])
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg == 5
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[2].arg == 5
 
   def test_where_on_casted_gated_load_extra_cond(self):
     ridx0 = UOp.range(100, 0)

tinygrad/codegen/__init__.py

@@ -3,7 +3,8 @@ from dataclasses import replace
 import itertools
 from tinygrad.helpers import DISABLE_FAST_IDIV, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
 from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, Target, panic
-from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, pyrender
+from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo
+from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.dtype import dtypes
@@ -107,9 +108,9 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True, b
 pm_linearize_cleanups = PatternMatcher([
   # if statements are not allowed in the graph
   (UPat((Ops.IF, Ops.ENDIF)), lambda: panic(RuntimeError, "if not allowed in graph")),
-  # gated INDEX becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF
-  (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat(name="gate", dtype=dtypes.bool))).or_casted(), UPat())),
-   lambda u, gate: (u, [mif:=UOp(Ops.IF, src=(gate, u.src[0])), u, UOp(Ops.ENDIF, src=(mif,))]))
+  # gated STORE becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF
+  (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX).or_casted(), UPat(), UPat(name="gate", dtype=dtypes.bool))),
+   lambda u, gate: (u.replace(src=u.src[:2]), [mif:=UOp(Ops.IF, src=(gate, u.src[0])), u.replace(src=u.src[:2]), UOp(Ops.ENDIF, src=(mif,))]))
 ])
 
 # requires lst be toposorted. like graph rewrite, but for lines

tinygrad/codegen/late/devectorizer.py

@@ -53,10 +53,15 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
 load_store_indexing = PatternMatcher([
   # image load valid idx simplification
   (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
-  # simplify away long after index has been lowered
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("x", dtypes.long), UPat.var("c", dtypes.bool))), lambda buf,x,c: simplify_valid_load(buf, x, c)),
-  # drop true gate
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("x"), UPat.const(dtypes.bool, True)),), lambda buf,x: buf.index(x, ptr=True)),
+  # drop true gate from gated LOAD with alt: also drop the now-unused alt
+  (UPat(Ops.LOAD, src=(UPat.var("idx"), UPat.const(dtypes.bool, True), UPat()), allow_any_len=True, name="ld"),
+   lambda ld,idx: ld.replace(src=(idx,)+ld.src[3:])),
+  # drop true gate from gated LOAD without alt
+  (UPat(Ops.LOAD, src=(UPat.var("idx"), UPat.const(dtypes.bool, True)), allow_any_len=True, name="ld"),
+   lambda ld,idx: ld.replace(src=(idx,)+ld.src[2:])),
+  # drop true gate from STORE
+  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat.var("val"), UPat.const(dtypes.bool, True))),
+   lambda idx,val: idx.store(val)),
 ])
 
 # ***** load/store grouping *****
@@ -116,22 +121,22 @@ def fold_expanded_index(midx:UOp):
   post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace).vec(global_offset), tuple(ret))
   return post_cat.gep(tuple(cast(list[int], idxs)))
 
-def cat_after_store(cat:UOp, data:UOp, sto:UOp):
+def cat_after_store(cat:UOp, data:UOp):
   # TODO: this is written in many places
   offset = 0
   ret: list[UOp] = []
   for s in cat.src:
-    ret.append(s.store(data.gep(tuple(range(offset, offset+s.dtype.count))), *sto.src[2:]))
+    ret.append(s.store(data.gep(tuple(range(offset, offset+s.dtype.count)))))
     offset += s.dtype.count
   return UOp.group(*ret)
 
-def gep_on_store(gep:UOp, st:UOp, sto:UOp):
+def gep_on_store(gep:UOp, st:UOp):
   # NOTE: we need to invert the gep here, but it may be an expanding gep
   # fake argsort. TODO: handle duplicates
   a = {}
   for i,x in enumerate(gep.arg): a[x] = i
   new_arg = tuple(x[1] for x in sorted(a.items()))
-  return gep.src[0].store(st.gep(new_arg), *sto.src[2:])
+  return gep.src[0].store(st.gep(new_arg))
 
 load_store_folding = PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat(Ops.STACK, src=UPat(GroupOp.Defines).or_after(name="buf")), UPat.var("vec"))), expand_index),
@@ -140,12 +145,12 @@ load_store_folding = PatternMatcher([
   (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)),
   # GEP on data of STORE
-  (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st")), name="sto"), gep_on_store),
+  (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st"))), gep_on_store),
   # put PTRCAT after LOAD
   (UPat(Ops.LOAD, src=(UPat(Ops.PTRCAT, name="cat"),), name="ld", allow_any_len=True),
    lambda cat,ld: UOp(Ops.VCAT, cat.dtype.base.vec(cat.dtype.vcount), tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))),
   # put PTRCAT after STORE
-  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), name="sto"), cat_after_store),
+  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data"))), cat_after_store),
 ])
 
 # *** correct load/store ***
@@ -187,7 +192,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
       if global_offset+fold_length > sz: continue
       lidx = buf.index((offset + global_offset).valid(mask), ptr=True)
       if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace))
-      if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length))))+ls.src[2:]))
+      if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length))))))
       else: ret.append(ls.replace(src=(lidx,)+ls.src[1:], dtype=ls.dtype.scalar().vec(fold_length)))
       global_offset += fold_length
       break
@@ -197,8 +202,9 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
   return UOp(Ops.VCAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)
 
 def get_image_idx(idx:UOp, width:int):
-  oidx = UOp(Ops.STACK, dtypes.weakint.vec(2), (((x:=idx.src[1].get_idx()) // 4) % width, (x // (4*width))))
-  return idx.replace(src=(idx.src[0], oidx.valid(idx.src[1].get_valid())))
+  x, valid = idx.src[1].get_idx(), idx.src[1].get_valid()
+  idx_x, idx_y = (x // 4) % width, x // (4*width)
+  return idx.replace(src=(idx.src[0], UOp.vectorize(idx_x, idx_y).valid(valid)))
 
 def image_fixup(ls:UOp):
   # normal image load or store, with the CAST from expand_index
@@ -280,18 +286,18 @@ pm_render = PatternMatcher([
   (UPat(Ops.GEP, name='gep'), lambda gep: UOp(Ops.STACK, gep.dtype, tuple(gep.src[0].gep(x) for x in gep.arg)) if len(gep.arg) > 1 else None),
   (UPat(Ops.GEP, name='gep'), lambda gep: gep.src[0] if gep.src[0].dtype.vcount == 1 and gep.arg == (0,) else None),
   (UPat(Ops.STACK, src=(UPat(name='x'),)), lambda x: x),
-  # give any loads that are masked an alt value
-  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted(),), allow_any_len=True, name="x"),
-    lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:])
-      if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE, Ops.BARRIER) else None),
+  # give any gated loads (gate at src[1]) an alt value at src[2]
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX).or_casted(), UPat(dtype=dtypes.bool)), allow_any_len=True, name="x"),
+    lambda x: x.replace(src=(x.src[0], x.src[1], x.const_like(0))+x.src[2:])
+      if len(x.src) == 2 or x.src[2].op in (Ops.CUSTOM, Ops.STORE, Ops.BARRIER) else None),
   # Where after gated load becomes alt value
   # NOTE: if a is CAST and a.src[0].dtype == l.dtype, use a.src[0] to avoid roundtrip cast (e.g. uint->float->uint)
-  (UPat.var("c").where(UPat(Ops.LOAD, src=(UPat().index(UPat(), UPat.var("c")).or_casted(),), allow_any_len=True, name="l").or_casted(),
-    UPat.var("a")), lambda c,l,a: l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype))+
-                                            l.src[2:]).cast(a.dtype)),
-  (UPat.var("c").where(UPat.var("a"), UPat(Ops.LOAD, src=(UPat().index(UPat(), UPat.var("c", dtype=dtypes.bool).logical_not()).or_casted(),),
-    allow_any_len=True, name="l").or_casted()), lambda c,l,a: l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype
-                                                                             else a.cast(l.dtype))+l.src[2:]).cast(a.dtype)),
+  (UPat.var("c").where(UPat(Ops.LOAD, src=(UPat(Ops.INDEX).or_casted(), UPat.var("c")), allow_any_len=True, name="l").or_casted(),
+    UPat.var("a")), lambda c,l,a: l.replace(src=(l.src[0], l.src[1], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype))+
+                                            l.src[3:]).cast(a.dtype)),
+  (UPat.var("c").where(UPat.var("a"), UPat(Ops.LOAD, src=(UPat(Ops.INDEX).or_casted(), UPat.var("c", dtype=dtypes.bool).logical_not()),
+    allow_any_len=True, name="l").or_casted()), lambda c,l,a: l.replace(src=(l.src[0], l.src[1], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype
+                                                                             else a.cast(l.dtype))+l.src[3:]).cast(a.dtype)),
 ])
 
 # *** Ops.REDUCE -> Ops.DEFINE_ACC ***

tinygrad/codegen/late/expander.py

@@ -45,7 +45,7 @@ def do_expand(root:UOp):
     else:
       # non-UNROLL input
       if root.op in range_start and i >= range_start[root.op]:
-        # for any range args of STORE/REDUCE, pass them through
+        # for any range args of REDUCE/WMMA/END/etc., pass them through
         new_srcs.append(src)
       elif root.op is Ops.INDEX and i >= 1 and not isinstance(root.dtype, PtrDType):
         new_srcs.append(src)

tinygrad/codegen/opt/search.py

@@ -1,6 +1,7 @@
 import math, time, multiprocessing, traceback, signal, atexit
 from dataclasses import replace
-from tinygrad.uop.ops import sym_infer, AxisType, pyrender, UOp
+from tinygrad.uop.ops import sym_infer, AxisType, UOp
+from tinygrad.uop.render import pyrender
 from tinygrad.device import Device, Buffer
 from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str, unwrap
 from tinygrad.helpers import IGNORE_BEAM_CACHE

tinygrad/codegen/simplify.py

@@ -14,7 +14,7 @@ def flatten_range(r:UOp) -> UOp|None:
 
 pm_flatten_range = PatternMatcher([
   # real ranges only
-  (UPat((Ops.REDUCE, Ops.STORE, Ops.END), name="r"), flatten_range),
+  (UPat((Ops.REDUCE, Ops.END), name="r"), flatten_range),
 ])
 
 def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.IDIV, Ops.MOD} for u in x.backward_slice)

tinygrad/nn/onnx.py

@@ -644,7 +644,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   def BitwiseOr(x:Tensor,y:Tensor): return x | y
   def BitwiseXor(x:Tensor,y:Tensor): return x ^ y
   def BitwiseNot(x:Tensor): return ~x
-  def Mod(x:Tensor,y:Tensor,fmod=0): return x - x.div(y, rounding_mode="trunc") * y if fmod else x % y
+  def Mod(x:Tensor,y:Tensor,fmod=0): return x.fmod(y) if fmod else x % y
 
   # ***** Casting Ops *****
   # NOTE: saturate only applies to FP8 types

tinygrad/renderer/__init__.py

@@ -3,7 +3,7 @@ from typing import Callable, cast
 from dataclasses import dataclass
 from tinygrad.helpers import prod, Target
 from tinygrad.uop.ops import Ops, UOp, sint, ssimplify, smin, GroupOp, PatternMatcher
-from tinygrad.dtype import AddrSpace, PtrDType
+from tinygrad.dtype import AddrSpace, PtrDType, dtypes
 from tinygrad.codegen.opt.tc import TensorCore
 from tinygrad.device import Compiler
 
@@ -31,8 +31,11 @@ class Estimates:
         if u.op in {Ops.LOAD, Ops.STORE}:
           # if u.src[0] is INDEX, we have to include the buffer since it might be an AFTER
           dont_count = dont_count.union((UOp.sink(*u.src[0].src[1:]) if u.src[0].op is Ops.INDEX else u.src[0]).toposort(range_gate))
-          # TODO: is this correct? this all needs to be cleaned up
-          if len(u.src) > 2: dont_count = dont_count.union(u.src[2].toposort())
+          # gate (bool-typed src) is part of indexing/predication, exclude its computation
+          # LOAD: gate at src[1]; STORE: gate at src[2]
+          gate_pos = 1 if u.op is Ops.LOAD else 2
+          if len(u.src) > gate_pos and u.src[gate_pos].dtype.scalar() == dtypes.bool:
+            dont_count = dont_count.union(u.src[gate_pos].toposort(range_gate))
         elif u.op is Ops.IF:
           dont_count = dont_count.union(u.src[0].toposort())
     for u in uops:

tinygrad/renderer/cstyle.py

@@ -44,9 +44,9 @@ base_rewrite = PatternMatcher([
   # default const render
   (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.arg)),
   # new load/store
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), allow_any_len=True),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx'))),
    lambda ctx,buf,idx: f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"),
-  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat.var("gate"))).or_casted("bidx"), UPat.var("var"))),
+  (UPat(Ops.LOAD, src=(UPat.var('bidx'), UPat.var("gate", dtype=dtypes.bool), UPat.var("var"))),
    lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"),
   (UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"(*{ctx[bidx]})"),
   (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
@@ -302,12 +302,13 @@ class OpenCLRenderer(CStyleLanguage):
     (UPat(Ops.CONST, dtypes.bfloat16, name="x"),
       lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.arg)))[0] >> 16)}u"),
     # load/store image (OpenCL)
-    (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2)), UPat.var("gate")), UPat.var("var"))),
+    (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2))), UPat.var("gate", dtype=dtypes.bool),
+                                                    UPat.var("var"))),
       lambda ctx,buf,idx,var,gate: f"({ctx[gate]}?read_imagef({ctx[buf]}, smp, {ctx[idx]}):{ctx[var]})"),
     (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2))),)),
       lambda ctx,buf,idx: f"read_imagef({ctx[buf]}, smp, {ctx[idx]})"),
-    (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2)), allow_any_len=True),
-                          UPat.var("var", dtypes.float.vec(4))), allow_any_len=True),
+    (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2))),
+                          UPat.var("var", dtypes.float.vec(4)))),
       lambda ctx,buf,idx,var: f"write_imagef({ctx[buf]}, {ctx[idx]}, {ctx[var]});"),
   ]) + base_rewrite
 

tinygrad/renderer/llvmir.py

@@ -76,7 +76,7 @@ base_rewrite = PatternMatcher([
   # memory load/store
   (UPat(Ops.INDEX, name="x"), lambda ctx,x:
    f"  {ctx[x]} = getelementptr inbounds {ldt(x.dtype.base)}, {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}"),
-  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat.var("mask"))).or_casted("idx"), UPat.var("alt")), allow_any_len=True, name="x"),
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX).or_casted("idx"), UPat.var("mask", dtype=dtypes.bool), UPat.var("alt")), allow_any_len=True, name="x"),
    lambda ctx,x,idx,alt,mask:
    f"  br label {ctx[x]}_entry\n{ctx[x][1:]}_entry:\n"
    f"  br i1 {ctx[mask]}, label {ctx[x]}_load, label {ctx[x]}_exit\n{ctx[x][1:]}_load:\n"

tinygrad/renderer/nir.py

@@ -125,19 +125,20 @@ class NIRRenderer(Renderer):
     (UPat.cvar("x", dtypes.uints), lambda x: UOp.const(x.dtype, x.dtype.max+x.arg+1) if x.arg < 0 else None),
     # from ptx
     (UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
-    # load/store bool -> uint8
+    # load/store bool -> uint8 (alt at src[2] in new gated shape; preserve gate at src[1])
     (UPat(Ops.LOAD, dtypes.bool, name="x"),
-     lambda x: x.replace(dtype=dtypes.uint8, src=x.src[0:1]+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())+x.src[2:]).cast(dtypes.bool)),
+     lambda x: x.replace(dtype=dtypes.uint8, src=tuple(s.cast(dtypes.uint8) if i == 2 and s.dtype.scalar() == dtypes.bool else s
+                                                       for i,s in enumerate(x.src))).cast(dtypes.bool)),
     (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
-     lambda x: x.replace(src=x.src[0:1] + (x.src[1].cast(dtypes.uint8),) + x.src[2:])),
+     lambda x: x.replace(src=(x.src[0], x.src[1].cast(dtypes.uint8))+x.src[2:])),
     # NIR requires shift amount to be 32 bit: https://docs.mesa3d.org/nir/alu.html#nir-alu-op-ishl
     (UPat((Ops.SHL, Ops.SHR), name="x"), lambda x: x.replace(src=(x.src[0], x.src[1].cast(dtypes.uint))) if x.src[1].dtype.bitsize != 32 else None),
     # OpConvertFToU is undefined if Result Type is not wide enough, cast through int32
     # ref: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpConvertFToU
     (UPat(Ops.CAST, (dtypes.uchar, dtypes.ushort), src=(UPat.var("x", dtypes.floats),), name="c"), lambda x,c: x.cast(dtypes.int32).cast(c.dtype)),
     # load/store use pointer arithmetic, and the cast does nothing
-    (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"), lambda x,buf,off: x.replace(
-      src=(buf,off.cast(dtypes.long))+x.src[2:]) if buf.dtype.addrspace != AddrSpace.REG and off.op not in (Ops.CAST, Ops.STACK) else None),
+    (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off")), name="x"), lambda x,buf,off: x.replace(
+      src=(buf,off.cast(dtypes.long))) if buf.dtype.addrspace != AddrSpace.REG and off.op not in (Ops.CAST, Ops.STACK) else None),
     (UPat(Ops.CAST, name="x"), lambda x: x.src[0] if isinstance(x.dtype, PtrDType) or x.src[0].dtype == dtypes.void else None),
   ])
 
@@ -146,9 +147,10 @@ class NIRRenderer(Renderer):
     (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 8)),
     (UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 4)),
     (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))),
-    (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val")), allow_any_len=True),
+    (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off"))), UPat.var("val"))),
      lambda ctx,buf,off,val: nstore(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), ctx.r[val], val.dtype)),
-    (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"), UPat.var("gate"))), UPat.var("alt")), allow_any_len=True, name="x"),
+    (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))), UPat.var("gate", dtype=dtypes.bool), UPat.var("alt")),
+          allow_any_len=True, name="x"),
      lambda ctx,x,buf,off,alt,gate: if_phi(ctx.b, ctx.r[gate],
       lambda: nload(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype, ctx.r[gate]), x.dtype), lambda: ctx.r[alt])),
     (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))),), allow_any_len=True, name="x"),
@@ -268,9 +270,9 @@ class IR3Renderer(NIRRenderer, OpenCLRenderer):
     return _nload_img(ctx.b, ctx.r[img], ctx.r[coord], img.dtype)
 
   def_rewrite = PatternMatcher([
-    (UPat(Ops.STORE, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2)), allow_any_len=True), UPat.var("val")),
-          allow_any_len=True), lambda ctx,img,coord,val: nstore_img(ctx.b, ctx.r[img], ctx.r[coord], ctx.r[val], val.dtype)),
-    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2)), UPat.var("gate")), UPat.var("alt"))),
+    (UPat(Ops.STORE, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2))), UPat.var("val"))),
+     lambda ctx,img,coord,val: nstore_img(ctx.b, ctx.r[img], ctx.r[coord], ctx.r[val], val.dtype)),
+    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2))), UPat.var("gate", dtype=dtypes.bool), UPat.var("alt"))),
      lambda ctx,img,coord,alt,gate: if_phi(ctx.b, ctx.r[gate], lambda: ctx.nload_img(img, coord), lambda: ctx.r[alt])),
     (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2))),)), nload_img),
   ]) + NIRRenderer.def_rewrite

tinygrad/renderer/ptx.py

@@ -45,14 +45,15 @@ ptx_matcher = PatternMatcher([
   # upcast to float32 all the ops that don't support half
   (UPat(doesnt_support_half, dtype=dtypes.half, name="x"),
     lambda x: (UOp(x.op, dtypes.float32, tuple(vv.cast(dtypes.float32) for vv in x.src), x.arg).cast(dtypes.half))),
-  # load/store bool -> uint8
+  # load/store bool -> uint8 (alt at src[2] in new gated shape; preserve gate at src[1])
   (UPat(Ops.LOAD, dtypes.bool, src=(UPat(dtype=dtypes.int64),), name="x", allow_any_len=True),
-   lambda x: UOp(x.op, dtypes.uint8, x.src[0:1] + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ()) + x.src[2:]).cast(dtypes.bool)),
+   lambda x: UOp(x.op, dtypes.uint8, tuple(s.cast(dtypes.uint8) if i == 2 and s.dtype.scalar() == dtypes.bool else s
+                                           for i,s in enumerate(x.src))).cast(dtypes.bool)),
   (UPat(Ops.STORE, src=(UPat(dtype=dtypes.int64), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
-   lambda x: UOp(x.op, dtypes.void, x.src[0:1] + (x.src[1].cast(dtypes.uint8),) + x.src[2:])),
+   lambda x: UOp(x.op, dtypes.void, (x.src[0], x.src[1].cast(dtypes.uint8))+x.src[2:])),
   # indexing on PTX is in uint64, we do the math while it's still in the graph
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx")), name="op", allow_any_len=True), lambda buf,idx,op:
-    UOp(Ops.INDEX, dtype=dtypes.int64, src=(buf, buf.cast(dtypes.int64)+idx.cast(dtypes.int64)*buf.dtype.itemsize)+op.src[2:]) \
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx")), name="op"), lambda buf,idx,op:
+    UOp(Ops.INDEX, dtype=dtypes.int64, src=(buf, buf.cast(dtypes.int64)+idx.cast(dtypes.int64)*buf.dtype.itemsize)) \
       if op.dtype != dtypes.int64 and buf.dtype.addrspace != AddrSpace.REG else None),
   # load/store use pointer arithmetic, and the cast does nothing
   (UPat(Ops.CAST, name="x"), lambda x: x.src[0] if isinstance(x.dtype, PtrDType) or x.src[0].dtype == dtypes.void else None),
@@ -102,18 +103,18 @@ string_rewrite = PatternMatcher([
   (UPat(Ops.CAST, name="x", src=(UPat.var("a"),)),
    lambda ctx, x, a: f"cvt{modifier(x.dtype, a.dtype)}.{ctx.cast_types[x.dtype]}.{ctx.cast_types[a.dtype]} {ctx.r[x]}, {ctx.r[a]};"),
   # store / gated load / load
-  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc")), allow_any_len=True), UPat.var("var"))),
+  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))), UPat.var("var"))),
    lambda ctx, loc, var, buf: f"st.{mem_type(buf)}" + \
     f"{f'.v{cnt}' if ((cnt:=var.dtype.count)>1) else ''}.{ctx.mem_types[var.dtype.scalar()]} " + \
     f"[{ctx.r[loc]}+0], {('{' + ', '.join(ctx.r[var]) + '}') if var.dtype.count > 1 else ctx.r[var]};"),
-  (UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"), UPat.var("gate"))), UPat.var("alt")), allow_any_len=True),
+  (UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))), UPat.var("gate", dtype=dtypes.bool), UPat.var("alt"))),
     lambda ctx, x, loc, alt, gate, buf: flatten([
     [f"mov.{ctx.mem_types[x.dtype.scalar()]} {v}, {render_val(0, x.dtype.scalar())};" for v in ctx.r[x]],
     [f"@{ctx.r[gate]} ld.{mem_type(buf)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];"]
   ]) if alt.dtype.count > 1 else [
     f"@{ctx.r[gate]} ld.{mem_type(buf)}.{ctx.mem_types[x.dtype.scalar()]} {ctx.r[x]}, [{ctx.r[loc]}+0];",
     f"@!{ctx.r[gate]} mov.b{ctx.types[x.dtype.scalar()][1:]} {ctx.r[x]}, {ctx.r[alt]};"]),
-  (UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))),), allow_any_len=True),
+  (UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))),)),
     lambda ctx, x, loc, buf: f"ld.{mem_type(buf)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];" \
      if x.dtype.count > 1 else f"ld.{mem_type(buf)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];"),
   # simple

tinygrad/renderer/wgsl.py

@@ -10,21 +10,27 @@ def sign_extend(val:UOp, sext_am:int):
         | val.bitcast(dtypes.uint32)).bitcast(dtypes.int)
 
 # store for char: buf[idx/4] <- (var << (idx%4)*8))
-def packed_store(bidx:UOp, var:UOp):
+def packed_store(bidx:UOp, var:UOp, *extra:UOp):
   elems, mask = 4//var.dtype.itemsize, _mask(var.dtype)
   shift_am, div_idx = (bidx.src[1].cast(dtypes.uint32) % elems) * (8*var.dtype.itemsize), bidx.src[1] // elems
   new_v, wmask = (var & mask).cast(dtypes.uint32) << shift_am, ((mask << shift_am) ^ 0xFFFFFFFF).cast(dtypes.uint32)
-  # preserve valid condition (bidx.src[2]) if it exists for gated stores
-  idx_src = (bidx.src[0], div_idx) if len(bidx.src) == 2 else (bidx.src[0], div_idx, bidx.src[2])
-  buf = UOp.load(UOp(Ops.INDEX, bidx.dtype, idx_src), dtype=dtypes.uint32)
-  return UOp.store(UOp(Ops.INDEX, bidx.dtype, idx_src), (buf & wmask) | new_v)
+  new_idx = UOp(Ops.INDEX, bidx.dtype, (bidx.src[0], div_idx))
+  buf = UOp.load(new_idx, dtype=dtypes.uint32)
+  # preserve trailing srcs (e.g. gate at src[2] for gated stores)
+  return UOp(Ops.STORE, dtypes.void, (new_idx, (buf & wmask) | new_v) + extra)
 
 # load for char: sign_extend(buf[idx/4] >> ((idx%4)*8))
 def packed_load(root:UOp, bidx:UOp, dtype:DType, var:UOp|None=None):
   elems, mask = 4//dtype.itemsize, _mask(dtype)
   shift_am, div_idx = (bidx.src[1].cast(dtypes.uint32) % elems) * (8*dtype.itemsize), bidx.src[1] // elems
-  idx = UOp(Ops.INDEX, bidx.dtype, (bidx.src[0], div_idx, bidx.src[2]) if var is not None else (bidx.src[0], div_idx))
-  load = UOp.load(idx, *([var] if var is not None else root.src[1:]), dtype=dtypes.uint32, arg=root.arg)
+  new_idx = UOp(Ops.INDEX, bidx.dtype, (bidx.src[0], div_idx))
+  # rebuild LOAD srcs preserving gate at src[1] (if bool) and replacing alt with var if provided
+  other_srcs = list(root.src[1:])
+  if var is not None:
+    alt_pos = 1 if (len(other_srcs) >= 1 and other_srcs[0].dtype.scalar() == dtypes.bool) else 0
+    if alt_pos < len(other_srcs): other_srcs[alt_pos] = var
+    else: other_srcs.append(var)
+  load = UOp.load(new_idx, *other_srcs, dtype=dtypes.uint32, arg=root.arg)
   val = (load.cast(dtypes.uint32) >> shift_am) & mask
   return sign_extend(val, 8*dtype.itemsize).cast(dtype) if dtype in [dtypes.char, dtypes.short] else val.cast(dtype)
 
@@ -40,12 +46,12 @@ def is_nan(a):
 wgsl_matcher = PatternMatcher([
   (UPat((Ops.CMPLT, Ops.XOR), src=(UPat(name="a", dtype=dtypes.bool), UPat.var("b")), name="c"),
    lambda a,b,c: a.cast(dtypes.int).alu(c.op, b.cast(dtypes.int)).cast(dtypes.bool)),
-  # TODO: load alt value doesnt have to be a const
-  (UPat.load(UPat.var("b"), UPat.cvar("c"), allow_any_len=True, name="l"),
-   lambda l,b,c: packed_load(l,b,l.dtype,c.cast(dtypes.uint32)) if is_packed(l.dtype, b.dtype) else None),
+  # TODO: load alt value doesnt have to be a const (alt is at src[2] in gated LOAD)
+  (UPat.load(UPat.var("b"), UPat.var("g", dtype=dtypes.bool), UPat.cvar("c"), name="l"),
+   lambda l,b,g,c: packed_load(l,b,l.dtype,c.cast(dtypes.uint32)) if is_packed(l.dtype, b.dtype) else None),
   (UPat.load(UPat.var("b"), name='l', allow_any_len=True), lambda l,b: packed_load(l, b, l.dtype) if is_packed(l.dtype, b.dtype) else None),
-  (UPat.store(UPat.var("bidx"), UPat.var("var"), allow_any_len=True),
-   lambda bidx,var: packed_store(bidx,var) if is_packed(var.dtype, bidx.dtype) else None),
+  (UPat.store(UPat.var("bidx"), UPat.var("var"), allow_any_len=True, name="sto"),
+   lambda bidx,var,sto: packed_store(bidx,var,*sto.src[2:]) if is_packed(var.dtype, bidx.dtype) else None),
   (UPat.var("a") << UPat.var("b"),lambda a,b:(a.bitcast(dtypes.uint32)<<b.cast(dtypes.uint32)).bitcast(a.dtype) if b.dtype!=dtypes.uint32 else None),
   (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
   # fix nan check: 'a != a -> is_nan()'
@@ -81,15 +87,15 @@ class WGSLRenderer(CStyleLanguage):
     (UPat(Ops.BITCAST, dtype=dtypes.short, name="x"), lambda ctx,x: f"bitcast<i32>(vec2<f16>({ctx[x.src[0]]},0))" \
      if x.src[0].dtype == dtypes.half else f"((i32({ctx[x.src[0]]}&0xFFFF)<<16)>>16)"),
     (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast<{ctx.type_map[x.dtype]}>({ctx[x.src[0]]})"),
-    # TODO: load alt value doesnt have to be a const
-    (UPat.load(UPat.var("b"), UPat.cvar("v"), allow_any_len=True),
-      lambda ctx,b,v: f"select({ctx[v]}, {ctx.render_load(ctx[b],b.src[0].dtype)}, {ctx[b.src[2]]})"),
+    # TODO: load alt value doesnt have to be a const (gated load: src[1]=gate, src[2]=alt)
+    (UPat.load(UPat.var("b"), UPat.var("g", dtype=dtypes.bool), UPat.cvar("v")),
+      lambda ctx,b,g,v: f"select({ctx[v]}, {ctx.render_load(ctx[b],b.src[0].dtype)}, {ctx[g]})"),
     (UPat.load(UPat.var("b"), allow_any_len=True), lambda ctx, b: ctx.render_load(ctx[b], b.dtype)),
     (UPat.store(UPat.var("b"), UPat.var("v"), allow_any_len=True),lambda ctx,b,v:\
      # (load & mask) | var -> mask = v.src[0].src[1], var = v.src[1]
      f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n  atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b.src[0].dtype) \
       else f"{ctx[b]} = {ctx[v]};"),
-    (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx")), allow_any_len=True),
+    (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"))),
      lambda ctx,b,idx: f"{ctx[b]}[{strip_parens(ctx[idx]) if idx.arg is Ops.ADD else ctx[idx]}]"),
   ]) + base_rewrite
 

tinygrad/runtime/autogen/am/__init__.py

@@ -34,4 +34,6 @@ def __getattr__(nm):
                                                                            for f in files if (p:=pathlib.Path(f)).is_file()] + ["}"])
       return load("am/fw", ["{}/amdgpu/psp_*_sos.bin", "{}/amdgpu/smu_*.bin", "{}/amdgpu/sdma_*.bin"] +
                            [f"{{}}/amdgpu/gc_*_{x}.bin" for x in ["pfp", "me", "mec", "imu", "rlc"]], srcs=fw_src, gen=genfw)
+    case "navi_offsets": return load("am/navi_offsets", [f"{AMD}/include/sienna_cichlid_ip_offset.h"], srcs=am_src)
+    case "vega_offsets": return load("am/vega_offsets", [f"{AMD}/include/vega20_ip_offset.h"], srcs=am_src)
     case _: raise AttributeError(f"no such autogen: {nm}")

tinygrad/runtime/autogen/am/navi_offsets.py

@@ -0,0 +1,823 @@
+# mypy: disable-error-code="empty-body"
+from __future__ import annotations
+import ctypes
+from typing import Literal, TypeAlias
+from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support import c
+@c.record
+class struct_IP_BASE_INSTANCE(c.Struct):
+  SIZE = 20
+  segment: c.Array[ctypes.c_uint32, Literal[5]]
+struct_IP_BASE_INSTANCE.register_fields([('segment', c.Array[ctypes.c_uint32, Literal[5]], 0)])
+@c.record
+class struct_IP_BASE(c.Struct):
+  SIZE = 140
+  instance: c.Array[struct_IP_BASE_INSTANCE, Literal[7]]
+struct_IP_BASE.register_fields([('instance', c.Array[struct_IP_BASE_INSTANCE, Literal[7]], 0)])
+MAX_INSTANCE = 7
+MAX_SEGMENT = 5
+ATHUB_BASE__INST0_SEG0 = 0x00000C00
+ATHUB_BASE__INST0_SEG1 = 0x02408C00
+ATHUB_BASE__INST0_SEG2 = 0
+ATHUB_BASE__INST0_SEG3 = 0
+ATHUB_BASE__INST0_SEG4 = 0
+ATHUB_BASE__INST1_SEG0 = 0
+ATHUB_BASE__INST1_SEG1 = 0
+ATHUB_BASE__INST1_SEG2 = 0
+ATHUB_BASE__INST1_SEG3 = 0
+ATHUB_BASE__INST1_SEG4 = 0
+ATHUB_BASE__INST2_SEG0 = 0
+ATHUB_BASE__INST2_SEG1 = 0
+ATHUB_BASE__INST2_SEG2 = 0
+ATHUB_BASE__INST2_SEG3 = 0
+ATHUB_BASE__INST2_SEG4 = 0
+ATHUB_BASE__INST3_SEG0 = 0
+ATHUB_BASE__INST3_SEG1 = 0
+ATHUB_BASE__INST3_SEG2 = 0
+ATHUB_BASE__INST3_SEG3 = 0
+ATHUB_BASE__INST3_SEG4 = 0
+ATHUB_BASE__INST4_SEG0 = 0
+ATHUB_BASE__INST4_SEG1 = 0
+ATHUB_BASE__INST4_SEG2 = 0
+ATHUB_BASE__INST4_SEG3 = 0
+ATHUB_BASE__INST4_SEG4 = 0
+ATHUB_BASE__INST5_SEG0 = 0
+ATHUB_BASE__INST5_SEG1 = 0
+ATHUB_BASE__INST5_SEG2 = 0
+ATHUB_BASE__INST5_SEG3 = 0
+ATHUB_BASE__INST5_SEG4 = 0
+ATHUB_BASE__INST6_SEG0 = 0
+ATHUB_BASE__INST6_SEG1 = 0
+ATHUB_BASE__INST6_SEG2 = 0
+ATHUB_BASE__INST6_SEG3 = 0
+ATHUB_BASE__INST6_SEG4 = 0
+CLK_BASE__INST0_SEG0 = 0x00016C00
+CLK_BASE__INST0_SEG1 = 0x02401800
+CLK_BASE__INST0_SEG2 = 0
+CLK_BASE__INST0_SEG3 = 0
+CLK_BASE__INST0_SEG4 = 0
+CLK_BASE__INST1_SEG0 = 0x00016E00
+CLK_BASE__INST1_SEG1 = 0x02401C00
+CLK_BASE__INST1_SEG2 = 0
+CLK_BASE__INST1_SEG3 = 0
+CLK_BASE__INST1_SEG4 = 0
+CLK_BASE__INST2_SEG0 = 0x00017000
+CLK_BASE__INST2_SEG1 = 0x02402000
+CLK_BASE__INST2_SEG2 = 0
+CLK_BASE__INST2_SEG3 = 0
+CLK_BASE__INST2_SEG4 = 0
+CLK_BASE__INST3_SEG0 = 0x00017200
+CLK_BASE__INST3_SEG1 = 0x02402400
+CLK_BASE__INST3_SEG2 = 0
+CLK_BASE__INST3_SEG3 = 0
+CLK_BASE__INST3_SEG4 = 0
+CLK_BASE__INST4_SEG0 = 0x0001B000
+CLK_BASE__INST4_SEG1 = 0x0242D800
+CLK_BASE__INST4_SEG2 = 0
+CLK_BASE__INST4_SEG3 = 0
+CLK_BASE__INST4_SEG4 = 0
+CLK_BASE__INST5_SEG0 = 0x0001B200
+CLK_BASE__INST5_SEG1 = 0x0242DC00
+CLK_BASE__INST5_SEG2 = 0
+CLK_BASE__INST5_SEG3 = 0
+CLK_BASE__INST5_SEG4 = 0
+CLK_BASE__INST6_SEG0 = 0x0001B400
+CLK_BASE__INST6_SEG1 = 0x0242E000
+CLK_BASE__INST6_SEG2 = 0
+CLK_BASE__INST6_SEG3 = 0
+CLK_BASE__INST6_SEG4 = 0
+DF_BASE__INST0_SEG0 = 0x00007000
+DF_BASE__INST0_SEG1 = 0x0240B800
+DF_BASE__INST0_SEG2 = 0
+DF_BASE__INST0_SEG3 = 0
+DF_BASE__INST0_SEG4 = 0
+DF_BASE__INST1_SEG0 = 0
+DF_BASE__INST1_SEG1 = 0
+DF_BASE__INST1_SEG2 = 0
+DF_BASE__INST1_SEG3 = 0
+DF_BASE__INST1_SEG4 = 0
+DF_BASE__INST2_SEG0 = 0
+DF_BASE__INST2_SEG1 = 0
+DF_BASE__INST2_SEG2 = 0
+DF_BASE__INST2_SEG3 = 0
+DF_BASE__INST2_SEG4 = 0
+DF_BASE__INST3_SEG0 = 0
+DF_BASE__INST3_SEG1 = 0
+DF_BASE__INST3_SEG2 = 0
+DF_BASE__INST3_SEG3 = 0
+DF_BASE__INST3_SEG4 = 0
+DF_BASE__INST4_SEG0 = 0
+DF_BASE__INST4_SEG1 = 0
+DF_BASE__INST4_SEG2 = 0
+DF_BASE__INST4_SEG3 = 0
+DF_BASE__INST4_SEG4 = 0
+DF_BASE__INST5_SEG0 = 0
+DF_BASE__INST5_SEG1 = 0
+DF_BASE__INST5_SEG2 = 0
+DF_BASE__INST5_SEG3 = 0
+DF_BASE__INST5_SEG4 = 0
+DF_BASE__INST6_SEG0 = 0
+DF_BASE__INST6_SEG1 = 0
+DF_BASE__INST6_SEG2 = 0
+DF_BASE__INST6_SEG3 = 0
+DF_BASE__INST6_SEG4 = 0
+DIO_BASE__INST0_SEG0 = 0x02404000
+DIO_BASE__INST0_SEG1 = 0
+DIO_BASE__INST0_SEG2 = 0
+DIO_BASE__INST0_SEG3 = 0
+DIO_BASE__INST0_SEG4 = 0
+DIO_BASE__INST1_SEG0 = 0
+DIO_BASE__INST1_SEG1 = 0
+DIO_BASE__INST1_SEG2 = 0
+DIO_BASE__INST1_SEG3 = 0
+DIO_BASE__INST1_SEG4 = 0
+DIO_BASE__INST2_SEG0 = 0
+DIO_BASE__INST2_SEG1 = 0
+DIO_BASE__INST2_SEG2 = 0
+DIO_BASE__INST2_SEG3 = 0
+DIO_BASE__INST2_SEG4 = 0
+DIO_BASE__INST3_SEG0 = 0
+DIO_BASE__INST3_SEG1 = 0
+DIO_BASE__INST3_SEG2 = 0
+DIO_BASE__INST3_SEG3 = 0
+DIO_BASE__INST3_SEG4 = 0
+DIO_BASE__INST4_SEG0 = 0
+DIO_BASE__INST4_SEG1 = 0
+DIO_BASE__INST4_SEG2 = 0
+DIO_BASE__INST4_SEG3 = 0
+DIO_BASE__INST4_SEG4 = 0
+DIO_BASE__INST5_SEG0 = 0
+DIO_BASE__INST5_SEG1 = 0
+DIO_BASE__INST5_SEG2 = 0
+DIO_BASE__INST5_SEG3 = 0
+DIO_BASE__INST5_SEG4 = 0
+DIO_BASE__INST6_SEG0 = 0
+DIO_BASE__INST6_SEG1 = 0
+DIO_BASE__INST6_SEG2 = 0
+DIO_BASE__INST6_SEG3 = 0
+DIO_BASE__INST6_SEG4 = 0
+DCN_BASE__INST0_SEG0 = 0x00000012
+DCN_BASE__INST0_SEG1 = 0x000000C0
+DCN_BASE__INST0_SEG2 = 0x000034C0
+DCN_BASE__INST0_SEG3 = 0x00009000
+DCN_BASE__INST0_SEG4 = 0x02403C00
+DCN_BASE__INST1_SEG0 = 0
+DCN_BASE__INST1_SEG1 = 0
+DCN_BASE__INST1_SEG2 = 0
+DCN_BASE__INST1_SEG3 = 0
+DCN_BASE__INST1_SEG4 = 0
+DCN_BASE__INST2_SEG0 = 0
+DCN_BASE__INST2_SEG1 = 0
+DCN_BASE__INST2_SEG2 = 0
+DCN_BASE__INST2_SEG3 = 0
+DCN_BASE__INST2_SEG4 = 0
+DCN_BASE__INST3_SEG0 = 0
+DCN_BASE__INST3_SEG1 = 0
+DCN_BASE__INST3_SEG2 = 0
+DCN_BASE__INST3_SEG3 = 0
+DCN_BASE__INST3_SEG4 = 0
+DCN_BASE__INST4_SEG0 = 0
+DCN_BASE__INST4_SEG1 = 0
+DCN_BASE__INST4_SEG2 = 0
+DCN_BASE__INST4_SEG3 = 0
+DCN_BASE__INST4_SEG4 = 0
+DCN_BASE__INST5_SEG0 = 0
+DCN_BASE__INST5_SEG1 = 0
+DCN_BASE__INST5_SEG2 = 0
+DCN_BASE__INST5_SEG3 = 0
+DCN_BASE__INST5_SEG4 = 0
+DCN_BASE__INST6_SEG0 = 0
+DCN_BASE__INST6_SEG1 = 0
+DCN_BASE__INST6_SEG2 = 0
+DCN_BASE__INST6_SEG3 = 0
+DCN_BASE__INST6_SEG4 = 0
+DPCS_BASE__INST0_SEG0 = 0x00000012
+DPCS_BASE__INST0_SEG1 = 0x000000C0
+DPCS_BASE__INST0_SEG2 = 0x000034C0
+DPCS_BASE__INST0_SEG3 = 0x00009000
+DPCS_BASE__INST0_SEG4 = 0x02403C00
+DPCS_BASE__INST1_SEG0 = 0
+DPCS_BASE__INST1_SEG1 = 0
+DPCS_BASE__INST1_SEG2 = 0
+DPCS_BASE__INST1_SEG3 = 0
+DPCS_BASE__INST1_SEG4 = 0
+DPCS_BASE__INST2_SEG0 = 0
+DPCS_BASE__INST2_SEG1 = 0
+DPCS_BASE__INST2_SEG2 = 0
+DPCS_BASE__INST2_SEG3 = 0
+DPCS_BASE__INST2_SEG4 = 0
+DPCS_BASE__INST3_SEG0 = 0
+DPCS_BASE__INST3_SEG1 = 0
+DPCS_BASE__INST3_SEG2 = 0
+DPCS_BASE__INST3_SEG3 = 0
+DPCS_BASE__INST3_SEG4 = 0
+DPCS_BASE__INST4_SEG0 = 0
+DPCS_BASE__INST4_SEG1 = 0
+DPCS_BASE__INST4_SEG2 = 0
+DPCS_BASE__INST4_SEG3 = 0
+DPCS_BASE__INST4_SEG4 = 0
+DPCS_BASE__INST5_SEG0 = 0
+DPCS_BASE__INST5_SEG1 = 0
+DPCS_BASE__INST5_SEG2 = 0
+DPCS_BASE__INST5_SEG3 = 0
+DPCS_BASE__INST5_SEG4 = 0
+DPCS_BASE__INST6_SEG0 = 0
+DPCS_BASE__INST6_SEG1 = 0
+DPCS_BASE__INST6_SEG2 = 0
+DPCS_BASE__INST6_SEG3 = 0
+DPCS_BASE__INST6_SEG4 = 0
+FUSE_BASE__INST0_SEG0 = 0x00017400
+FUSE_BASE__INST0_SEG1 = 0x02401400
+FUSE_BASE__INST0_SEG2 = 0
+FUSE_BASE__INST0_SEG3 = 0
+FUSE_BASE__INST0_SEG4 = 0
+FUSE_BASE__INST1_SEG0 = 0
+FUSE_BASE__INST1_SEG1 = 0
+FUSE_BASE__INST1_SEG2 = 0
+FUSE_BASE__INST1_SEG3 = 0
+FUSE_BASE__INST1_SEG4 = 0
+FUSE_BASE__INST2_SEG0 = 0
+FUSE_BASE__INST2_SEG1 = 0
+FUSE_BASE__INST2_SEG2 = 0
+FUSE_BASE__INST2_SEG3 = 0
+FUSE_BASE__INST2_SEG4 = 0
+FUSE_BASE__INST3_SEG0 = 0
+FUSE_BASE__INST3_SEG1 = 0
+FUSE_BASE__INST3_SEG2 = 0
+FUSE_BASE__INST3_SEG3 = 0
+FUSE_BASE__INST3_SEG4 = 0
+FUSE_BASE__INST4_SEG0 = 0
+FUSE_BASE__INST4_SEG1 = 0
+FUSE_BASE__INST4_SEG2 = 0
+FUSE_BASE__INST4_SEG3 = 0
+FUSE_BASE__INST4_SEG4 = 0
+FUSE_BASE__INST5_SEG0 = 0
+FUSE_BASE__INST5_SEG1 = 0
+FUSE_BASE__INST5_SEG2 = 0
+FUSE_BASE__INST5_SEG3 = 0
+FUSE_BASE__INST5_SEG4 = 0
+FUSE_BASE__INST6_SEG0 = 0
+FUSE_BASE__INST6_SEG1 = 0
+FUSE_BASE__INST6_SEG2 = 0
+FUSE_BASE__INST6_SEG3 = 0
+FUSE_BASE__INST6_SEG4 = 0
+GC_BASE__INST0_SEG0 = 0x00001260
+GC_BASE__INST0_SEG1 = 0x0000A000
+GC_BASE__INST0_SEG2 = 0x0001C000
+GC_BASE__INST0_SEG3 = 0x02402C00
+GC_BASE__INST0_SEG4 = 0
+GC_BASE__INST1_SEG0 = 0
+GC_BASE__INST1_SEG1 = 0
+GC_BASE__INST1_SEG2 = 0
+GC_BASE__INST1_SEG3 = 0
+GC_BASE__INST1_SEG4 = 0
+GC_BASE__INST2_SEG0 = 0
+GC_BASE__INST2_SEG1 = 0
+GC_BASE__INST2_SEG2 = 0
+GC_BASE__INST2_SEG3 = 0
+GC_BASE__INST2_SEG4 = 0
+GC_BASE__INST3_SEG0 = 0
+GC_BASE__INST3_SEG1 = 0
+GC_BASE__INST3_SEG2 = 0
+GC_BASE__INST3_SEG3 = 0
+GC_BASE__INST3_SEG4 = 0
+GC_BASE__INST4_SEG0 = 0
+GC_BASE__INST4_SEG1 = 0
+GC_BASE__INST4_SEG2 = 0
+GC_BASE__INST4_SEG3 = 0
+GC_BASE__INST4_SEG4 = 0
+GC_BASE__INST5_SEG0 = 0
+GC_BASE__INST5_SEG1 = 0
+GC_BASE__INST5_SEG2 = 0
+GC_BASE__INST5_SEG3 = 0
+GC_BASE__INST5_SEG4 = 0
+GC_BASE__INST6_SEG0 = 0
+GC_BASE__INST6_SEG1 = 0
+GC_BASE__INST6_SEG2 = 0
+GC_BASE__INST6_SEG3 = 0
+GC_BASE__INST6_SEG4 = 0
+HDA_BASE__INST0_SEG0 = 0x004C0000
+HDA_BASE__INST0_SEG1 = 0x02404800
+HDA_BASE__INST0_SEG2 = 0
+HDA_BASE__INST0_SEG3 = 0
+HDA_BASE__INST0_SEG4 = 0
+HDA_BASE__INST1_SEG0 = 0
+HDA_BASE__INST1_SEG1 = 0
+HDA_BASE__INST1_SEG2 = 0
+HDA_BASE__INST1_SEG3 = 0
+HDA_BASE__INST1_SEG4 = 0
+HDA_BASE__INST2_SEG0 = 0
+HDA_BASE__INST2_SEG1 = 0
+HDA_BASE__INST2_SEG2 = 0
+HDA_BASE__INST2_SEG3 = 0
+HDA_BASE__INST2_SEG4 = 0
+HDA_BASE__INST3_SEG0 = 0
+HDA_BASE__INST3_SEG1 = 0
+HDA_BASE__INST3_SEG2 = 0
+HDA_BASE__INST3_SEG3 = 0
+HDA_BASE__INST3_SEG4 = 0
+HDA_BASE__INST4_SEG0 = 0
+HDA_BASE__INST4_SEG1 = 0
+HDA_BASE__INST4_SEG2 = 0
+HDA_BASE__INST4_SEG3 = 0
+HDA_BASE__INST4_SEG4 = 0
+HDA_BASE__INST5_SEG0 = 0
+HDA_BASE__INST5_SEG1 = 0
+HDA_BASE__INST5_SEG2 = 0
+HDA_BASE__INST5_SEG3 = 0
+HDA_BASE__INST5_SEG4 = 0
+HDA_BASE__INST6_SEG0 = 0
+HDA_BASE__INST6_SEG1 = 0
+HDA_BASE__INST6_SEG2 = 0
+HDA_BASE__INST6_SEG3 = 0
+HDA_BASE__INST6_SEG4 = 0
+HDP_BASE__INST0_SEG0 = 0x00000F20
+HDP_BASE__INST0_SEG1 = 0x0240A400
+HDP_BASE__INST0_SEG2 = 0
+HDP_BASE__INST0_SEG3 = 0
+HDP_BASE__INST0_SEG4 = 0
+HDP_BASE__INST1_SEG0 = 0
+HDP_BASE__INST1_SEG1 = 0
+HDP_BASE__INST1_SEG2 = 0
+HDP_BASE__INST1_SEG3 = 0
+HDP_BASE__INST1_SEG4 = 0
+HDP_BASE__INST2_SEG0 = 0
+HDP_BASE__INST2_SEG1 = 0
+HDP_BASE__INST2_SEG2 = 0
+HDP_BASE__INST2_SEG3 = 0
+HDP_BASE__INST2_SEG4 = 0
+HDP_BASE__INST3_SEG0 = 0
+HDP_BASE__INST3_SEG1 = 0
+HDP_BASE__INST3_SEG2 = 0
+HDP_BASE__INST3_SEG3 = 0
+HDP_BASE__INST3_SEG4 = 0
+HDP_BASE__INST4_SEG0 = 0
+HDP_BASE__INST4_SEG1 = 0
+HDP_BASE__INST4_SEG2 = 0
+HDP_BASE__INST4_SEG3 = 0
+HDP_BASE__INST4_SEG4 = 0
+HDP_BASE__INST5_SEG0 = 0
+HDP_BASE__INST5_SEG1 = 0
+HDP_BASE__INST5_SEG2 = 0
+HDP_BASE__INST5_SEG3 = 0
+HDP_BASE__INST5_SEG4 = 0
+HDP_BASE__INST6_SEG0 = 0
+HDP_BASE__INST6_SEG1 = 0
+HDP_BASE__INST6_SEG2 = 0
+HDP_BASE__INST6_SEG3 = 0
+HDP_BASE__INST6_SEG4 = 0
+MMHUB_BASE__INST0_SEG0 = 0x0001A000
+MMHUB_BASE__INST0_SEG1 = 0x02408800
+MMHUB_BASE__INST0_SEG2 = 0
+MMHUB_BASE__INST0_SEG3 = 0
+MMHUB_BASE__INST0_SEG4 = 0
+MMHUB_BASE__INST1_SEG0 = 0
+MMHUB_BASE__INST1_SEG1 = 0
+MMHUB_BASE__INST1_SEG2 = 0
+MMHUB_BASE__INST1_SEG3 = 0
+MMHUB_BASE__INST1_SEG4 = 0
+MMHUB_BASE__INST2_SEG0 = 0
+MMHUB_BASE__INST2_SEG1 = 0
+MMHUB_BASE__INST2_SEG2 = 0
+MMHUB_BASE__INST2_SEG3 = 0
+MMHUB_BASE__INST2_SEG4 = 0
+MMHUB_BASE__INST3_SEG0 = 0
+MMHUB_BASE__INST3_SEG1 = 0
+MMHUB_BASE__INST3_SEG2 = 0
+MMHUB_BASE__INST3_SEG3 = 0
+MMHUB_BASE__INST3_SEG4 = 0
+MMHUB_BASE__INST4_SEG0 = 0
+MMHUB_BASE__INST4_SEG1 = 0
+MMHUB_BASE__INST4_SEG2 = 0
+MMHUB_BASE__INST4_SEG3 = 0
+MMHUB_BASE__INST4_SEG4 = 0
+MMHUB_BASE__INST5_SEG0 = 0
+MMHUB_BASE__INST5_SEG1 = 0
+MMHUB_BASE__INST5_SEG2 = 0
+MMHUB_BASE__INST5_SEG3 = 0
+MMHUB_BASE__INST5_SEG4 = 0
+MMHUB_BASE__INST6_SEG0 = 0
+MMHUB_BASE__INST6_SEG1 = 0
+MMHUB_BASE__INST6_SEG2 = 0
+MMHUB_BASE__INST6_SEG3 = 0
+MMHUB_BASE__INST6_SEG4 = 0
+MP0_BASE__INST0_SEG0 = 0x00016000
+MP0_BASE__INST0_SEG1 = 0x00DC0000
+MP0_BASE__INST0_SEG2 = 0x00E00000
+MP0_BASE__INST0_SEG3 = 0x00E40000
+MP0_BASE__INST0_SEG4 = 0x0243FC00
+MP0_BASE__INST1_SEG0 = 0
+MP0_BASE__INST1_SEG1 = 0
+MP0_BASE__INST1_SEG2 = 0
+MP0_BASE__INST1_SEG3 = 0
+MP0_BASE__INST1_SEG4 = 0
+MP0_BASE__INST2_SEG0 = 0
+MP0_BASE__INST2_SEG1 = 0
+MP0_BASE__INST2_SEG2 = 0
+MP0_BASE__INST2_SEG3 = 0
+MP0_BASE__INST2_SEG4 = 0
+MP0_BASE__INST3_SEG0 = 0
+MP0_BASE__INST3_SEG1 = 0
+MP0_BASE__INST3_SEG2 = 0
+MP0_BASE__INST3_SEG3 = 0
+MP0_BASE__INST3_SEG4 = 0
+MP0_BASE__INST4_SEG0 = 0
+MP0_BASE__INST4_SEG1 = 0
+MP0_BASE__INST4_SEG2 = 0
+MP0_BASE__INST4_SEG3 = 0
+MP0_BASE__INST4_SEG4 = 0
+MP0_BASE__INST5_SEG0 = 0
+MP0_BASE__INST5_SEG1 = 0
+MP0_BASE__INST5_SEG2 = 0
+MP0_BASE__INST5_SEG3 = 0
+MP0_BASE__INST5_SEG4 = 0
+MP0_BASE__INST6_SEG0 = 0
+MP0_BASE__INST6_SEG1 = 0
+MP0_BASE__INST6_SEG2 = 0
+MP0_BASE__INST6_SEG3 = 0
+MP0_BASE__INST6_SEG4 = 0
+MP1_BASE__INST0_SEG0 = 0x00016000
+MP1_BASE__INST0_SEG1 = 0x00DC0000
+MP1_BASE__INST0_SEG2 = 0x00E00000
+MP1_BASE__INST0_SEG3 = 0x00E40000
+MP1_BASE__INST0_SEG4 = 0x0243FC00
+MP1_BASE__INST1_SEG0 = 0
+MP1_BASE__INST1_SEG1 = 0
+MP1_BASE__INST1_SEG2 = 0
+MP1_BASE__INST1_SEG3 = 0
+MP1_BASE__INST1_SEG4 = 0
+MP1_BASE__INST2_SEG0 = 0
+MP1_BASE__INST2_SEG1 = 0
+MP1_BASE__INST2_SEG2 = 0
+MP1_BASE__INST2_SEG3 = 0
+MP1_BASE__INST2_SEG4 = 0
+MP1_BASE__INST3_SEG0 = 0
+MP1_BASE__INST3_SEG1 = 0
+MP1_BASE__INST3_SEG2 = 0
+MP1_BASE__INST3_SEG3 = 0
+MP1_BASE__INST3_SEG4 = 0
+MP1_BASE__INST4_SEG0 = 0
+MP1_BASE__INST4_SEG1 = 0
+MP1_BASE__INST4_SEG2 = 0
+MP1_BASE__INST4_SEG3 = 0
+MP1_BASE__INST4_SEG4 = 0
+MP1_BASE__INST5_SEG0 = 0
+MP1_BASE__INST5_SEG1 = 0
+MP1_BASE__INST5_SEG2 = 0
+MP1_BASE__INST5_SEG3 = 0
+MP1_BASE__INST5_SEG4 = 0
+MP1_BASE__INST6_SEG0 = 0
+MP1_BASE__INST6_SEG1 = 0
+MP1_BASE__INST6_SEG2 = 0
+MP1_BASE__INST6_SEG3 = 0
+MP1_BASE__INST6_SEG4 = 0
+NBIO_BASE__INST0_SEG0 = 0x00000000
+NBIO_BASE__INST0_SEG1 = 0x00000014
+NBIO_BASE__INST0_SEG2 = 0x00000D20
+NBIO_BASE__INST0_SEG3 = 0x00010400
+NBIO_BASE__INST0_SEG4 = 0x0241B000
+NBIO_BASE__INST1_SEG0 = 0
+NBIO_BASE__INST1_SEG1 = 0
+NBIO_BASE__INST1_SEG2 = 0
+NBIO_BASE__INST1_SEG3 = 0
+NBIO_BASE__INST1_SEG4 = 0
+NBIO_BASE__INST2_SEG0 = 0
+NBIO_BASE__INST2_SEG1 = 0
+NBIO_BASE__INST2_SEG2 = 0
+NBIO_BASE__INST2_SEG3 = 0
+NBIO_BASE__INST2_SEG4 = 0
+NBIO_BASE__INST3_SEG0 = 0
+NBIO_BASE__INST3_SEG1 = 0
+NBIO_BASE__INST3_SEG2 = 0
+NBIO_BASE__INST3_SEG3 = 0
+NBIO_BASE__INST3_SEG4 = 0
+NBIO_BASE__INST4_SEG0 = 0
+NBIO_BASE__INST4_SEG1 = 0
+NBIO_BASE__INST4_SEG2 = 0
+NBIO_BASE__INST4_SEG3 = 0
+NBIO_BASE__INST4_SEG4 = 0
+NBIO_BASE__INST5_SEG0 = 0
+NBIO_BASE__INST5_SEG1 = 0
+NBIO_BASE__INST5_SEG2 = 0
+NBIO_BASE__INST5_SEG3 = 0
+NBIO_BASE__INST5_SEG4 = 0
+NBIO_BASE__INST6_SEG0 = 0
+NBIO_BASE__INST6_SEG1 = 0
+NBIO_BASE__INST6_SEG2 = 0
+NBIO_BASE__INST6_SEG3 = 0
+NBIO_BASE__INST6_SEG4 = 0
+OSSSYS_BASE__INST0_SEG0 = 0x000010A0
+OSSSYS_BASE__INST0_SEG1 = 0x0240A000
+OSSSYS_BASE__INST0_SEG2 = 0
+OSSSYS_BASE__INST0_SEG3 = 0
+OSSSYS_BASE__INST0_SEG4 = 0
+OSSSYS_BASE__INST1_SEG0 = 0
+OSSSYS_BASE__INST1_SEG1 = 0
+OSSSYS_BASE__INST1_SEG2 = 0
+OSSSYS_BASE__INST1_SEG3 = 0
+OSSSYS_BASE__INST1_SEG4 = 0
+OSSSYS_BASE__INST2_SEG0 = 0
+OSSSYS_BASE__INST2_SEG1 = 0
+OSSSYS_BASE__INST2_SEG2 = 0
+OSSSYS_BASE__INST2_SEG3 = 0
+OSSSYS_BASE__INST2_SEG4 = 0
+OSSSYS_BASE__INST3_SEG0 = 0
+OSSSYS_BASE__INST3_SEG1 = 0
+OSSSYS_BASE__INST3_SEG2 = 0
+OSSSYS_BASE__INST3_SEG3 = 0
+OSSSYS_BASE__INST3_SEG4 = 0
+OSSSYS_BASE__INST4_SEG0 = 0
+OSSSYS_BASE__INST4_SEG1 = 0
+OSSSYS_BASE__INST4_SEG2 = 0
+OSSSYS_BASE__INST4_SEG3 = 0
+OSSSYS_BASE__INST4_SEG4 = 0
+OSSSYS_BASE__INST5_SEG0 = 0
+OSSSYS_BASE__INST5_SEG1 = 0
+OSSSYS_BASE__INST5_SEG2 = 0
+OSSSYS_BASE__INST5_SEG3 = 0
+OSSSYS_BASE__INST5_SEG4 = 0
+OSSSYS_BASE__INST6_SEG0 = 0
+OSSSYS_BASE__INST6_SEG1 = 0
+OSSSYS_BASE__INST6_SEG2 = 0
+OSSSYS_BASE__INST6_SEG3 = 0
+OSSSYS_BASE__INST6_SEG4 = 0
+PCIE0_BASE__INST0_SEG0 = 0x00000000
+PCIE0_BASE__INST0_SEG1 = 0x00000014
+PCIE0_BASE__INST0_SEG2 = 0x00000D20
+PCIE0_BASE__INST0_SEG3 = 0x00010400
+PCIE0_BASE__INST0_SEG4 = 0x0241B000
+PCIE0_BASE__INST1_SEG0 = 0
+PCIE0_BASE__INST1_SEG1 = 0
+PCIE0_BASE__INST1_SEG2 = 0
+PCIE0_BASE__INST1_SEG3 = 0
+PCIE0_BASE__INST1_SEG4 = 0
+PCIE0_BASE__INST2_SEG0 = 0
+PCIE0_BASE__INST2_SEG1 = 0
+PCIE0_BASE__INST2_SEG2 = 0
+PCIE0_BASE__INST2_SEG3 = 0
+PCIE0_BASE__INST2_SEG4 = 0
+PCIE0_BASE__INST3_SEG0 = 0
+PCIE0_BASE__INST3_SEG1 = 0
+PCIE0_BASE__INST3_SEG2 = 0
+PCIE0_BASE__INST3_SEG3 = 0
+PCIE0_BASE__INST3_SEG4 = 0
+PCIE0_BASE__INST4_SEG0 = 0
+PCIE0_BASE__INST4_SEG1 = 0
+PCIE0_BASE__INST4_SEG2 = 0
+PCIE0_BASE__INST4_SEG3 = 0
+PCIE0_BASE__INST4_SEG4 = 0
+PCIE0_BASE__INST5_SEG0 = 0
+PCIE0_BASE__INST5_SEG1 = 0
+PCIE0_BASE__INST5_SEG2 = 0
+PCIE0_BASE__INST5_SEG3 = 0
+PCIE0_BASE__INST5_SEG4 = 0
+PCIE0_BASE__INST6_SEG0 = 0
+PCIE0_BASE__INST6_SEG1 = 0
+PCIE0_BASE__INST6_SEG2 = 0
+PCIE0_BASE__INST6_SEG3 = 0
+PCIE0_BASE__INST6_SEG4 = 0
+SDMA0_BASE__INST0_SEG0 = 0x00001260
+SDMA0_BASE__INST0_SEG1 = 0x0000A000
+SDMA0_BASE__INST0_SEG2 = 0x0001C000
+SDMA0_BASE__INST0_SEG3 = 0x02402C00
+SDMA0_BASE__INST0_SEG4 = 0
+SDMA0_BASE__INST1_SEG0 = 0
+SDMA0_BASE__INST1_SEG1 = 0
+SDMA0_BASE__INST1_SEG2 = 0
+SDMA0_BASE__INST1_SEG3 = 0
+SDMA0_BASE__INST1_SEG4 = 0
+SDMA0_BASE__INST2_SEG0 = 0
+SDMA0_BASE__INST2_SEG1 = 0
+SDMA0_BASE__INST2_SEG2 = 0
+SDMA0_BASE__INST2_SEG3 = 0
+SDMA0_BASE__INST2_SEG4 = 0
+SDMA0_BASE__INST3_SEG0 = 0
+SDMA0_BASE__INST3_SEG1 = 0
+SDMA0_BASE__INST3_SEG2 = 0
+SDMA0_BASE__INST3_SEG3 = 0
+SDMA0_BASE__INST3_SEG4 = 0
+SDMA0_BASE__INST4_SEG0 = 0
+SDMA0_BASE__INST4_SEG1 = 0
+SDMA0_BASE__INST4_SEG2 = 0
+SDMA0_BASE__INST4_SEG3 = 0
+SDMA0_BASE__INST4_SEG4 = 0
+SDMA0_BASE__INST5_SEG0 = 0
+SDMA0_BASE__INST5_SEG1 = 0
+SDMA0_BASE__INST5_SEG2 = 0
+SDMA0_BASE__INST5_SEG3 = 0
+SDMA0_BASE__INST5_SEG4 = 0
+SDMA0_BASE__INST6_SEG0 = 0
+SDMA0_BASE__INST6_SEG1 = 0
+SDMA0_BASE__INST6_SEG2 = 0
+SDMA0_BASE__INST6_SEG3 = 0
+SDMA0_BASE__INST6_SEG4 = 0
+SDMA1_BASE__INST0_SEG0 = 0x00001260
+SDMA1_BASE__INST0_SEG1 = 0x0000A000
+SDMA1_BASE__INST0_SEG2 = 0x0001C000
+SDMA1_BASE__INST0_SEG3 = 0x02402C00
+SDMA1_BASE__INST0_SEG4 = 0
+SDMA1_BASE__INST1_SEG0 = 0
+SDMA1_BASE__INST1_SEG1 = 0
+SDMA1_BASE__INST1_SEG2 = 0
+SDMA1_BASE__INST1_SEG3 = 0
+SDMA1_BASE__INST1_SEG4 = 0
+SDMA1_BASE__INST2_SEG0 = 0
+SDMA1_BASE__INST2_SEG1 = 0
+SDMA1_BASE__INST2_SEG2 = 0
+SDMA1_BASE__INST2_SEG3 = 0
+SDMA1_BASE__INST2_SEG4 = 0
+SDMA1_BASE__INST3_SEG0 = 0
+SDMA1_BASE__INST3_SEG1 = 0
+SDMA1_BASE__INST3_SEG2 = 0
+SDMA1_BASE__INST3_SEG3 = 0
+SDMA1_BASE__INST3_SEG4 = 0
+SDMA1_BASE__INST4_SEG0 = 0
+SDMA1_BASE__INST4_SEG1 = 0
+SDMA1_BASE__INST4_SEG2 = 0
+SDMA1_BASE__INST4_SEG3 = 0
+SDMA1_BASE__INST4_SEG4 = 0
+SDMA1_BASE__INST5_SEG0 = 0
+SDMA1_BASE__INST5_SEG1 = 0
+SDMA1_BASE__INST5_SEG2 = 0
+SDMA1_BASE__INST5_SEG3 = 0
+SDMA1_BASE__INST5_SEG4 = 0
+SDMA1_BASE__INST6_SEG0 = 0
+SDMA1_BASE__INST6_SEG1 = 0
+SDMA1_BASE__INST6_SEG2 = 0
+SDMA1_BASE__INST6_SEG3 = 0
+SDMA1_BASE__INST6_SEG4 = 0
+SMUIO_BASE__INST0_SEG0 = 0x00016800
+SMUIO_BASE__INST0_SEG1 = 0x00016A00
+SMUIO_BASE__INST0_SEG2 = 0x00440000
+SMUIO_BASE__INST0_SEG3 = 0x02401000
+SMUIO_BASE__INST0_SEG4 = 0
+SMUIO_BASE__INST1_SEG0 = 0
+SMUIO_BASE__INST1_SEG1 = 0
+SMUIO_BASE__INST1_SEG2 = 0
+SMUIO_BASE__INST1_SEG3 = 0
+SMUIO_BASE__INST1_SEG4 = 0
+SMUIO_BASE__INST2_SEG0 = 0
+SMUIO_BASE__INST2_SEG1 = 0
+SMUIO_BASE__INST2_SEG2 = 0
+SMUIO_BASE__INST2_SEG3 = 0
+SMUIO_BASE__INST2_SEG4 = 0
+SMUIO_BASE__INST3_SEG0 = 0
+SMUIO_BASE__INST3_SEG1 = 0
+SMUIO_BASE__INST3_SEG2 = 0
+SMUIO_BASE__INST3_SEG3 = 0
+SMUIO_BASE__INST3_SEG4 = 0
+SMUIO_BASE__INST4_SEG0 = 0
+SMUIO_BASE__INST4_SEG1 = 0
+SMUIO_BASE__INST4_SEG2 = 0
+SMUIO_BASE__INST4_SEG3 = 0
+SMUIO_BASE__INST4_SEG4 = 0
+SMUIO_BASE__INST5_SEG0 = 0
+SMUIO_BASE__INST5_SEG1 = 0
+SMUIO_BASE__INST5_SEG2 = 0
+SMUIO_BASE__INST5_SEG3 = 0
+SMUIO_BASE__INST5_SEG4 = 0
+SMUIO_BASE__INST6_SEG0 = 0
+SMUIO_BASE__INST6_SEG1 = 0
+SMUIO_BASE__INST6_SEG2 = 0
+SMUIO_BASE__INST6_SEG3 = 0
+SMUIO_BASE__INST6_SEG4 = 0
+THM_BASE__INST0_SEG0 = 0x00016600
+THM_BASE__INST0_SEG1 = 0x02400C00
+THM_BASE__INST0_SEG2 = 0
+THM_BASE__INST0_SEG3 = 0
+THM_BASE__INST0_SEG4 = 0
+THM_BASE__INST1_SEG0 = 0
+THM_BASE__INST1_SEG1 = 0
+THM_BASE__INST1_SEG2 = 0
+THM_BASE__INST1_SEG3 = 0
+THM_BASE__INST1_SEG4 = 0
+THM_BASE__INST2_SEG0 = 0
+THM_BASE__INST2_SEG1 = 0
+THM_BASE__INST2_SEG2 = 0
+THM_BASE__INST2_SEG3 = 0
+THM_BASE__INST2_SEG4 = 0
+THM_BASE__INST3_SEG0 = 0
+THM_BASE__INST3_SEG1 = 0
+THM_BASE__INST3_SEG2 = 0
+THM_BASE__INST3_SEG3 = 0
+THM_BASE__INST3_SEG4 = 0
+THM_BASE__INST4_SEG0 = 0
+THM_BASE__INST4_SEG1 = 0
+THM_BASE__INST4_SEG2 = 0
+THM_BASE__INST4_SEG3 = 0
+THM_BASE__INST4_SEG4 = 0
+THM_BASE__INST5_SEG0 = 0
+THM_BASE__INST5_SEG1 = 0
+THM_BASE__INST5_SEG2 = 0
+THM_BASE__INST5_SEG3 = 0
+THM_BASE__INST5_SEG4 = 0
+THM_BASE__INST6_SEG0 = 0
+THM_BASE__INST6_SEG1 = 0
+THM_BASE__INST6_SEG2 = 0
+THM_BASE__INST6_SEG3 = 0
+THM_BASE__INST6_SEG4 = 0
+UMC_BASE__INST0_SEG0 = 0x00014000
+UMC_BASE__INST0_SEG1 = 0x02425800
+UMC_BASE__INST0_SEG2 = 0
+UMC_BASE__INST0_SEG3 = 0
+UMC_BASE__INST0_SEG4 = 0
+UMC_BASE__INST1_SEG0 = 0x00054000
+UMC_BASE__INST1_SEG1 = 0x02425C00
+UMC_BASE__INST1_SEG2 = 0
+UMC_BASE__INST1_SEG3 = 0
+UMC_BASE__INST1_SEG4 = 0
+UMC_BASE__INST2_SEG0 = 0x00094000
+UMC_BASE__INST2_SEG1 = 0x02426000
+UMC_BASE__INST2_SEG2 = 0
+UMC_BASE__INST2_SEG3 = 0
+UMC_BASE__INST2_SEG4 = 0
+UMC_BASE__INST3_SEG0 = 0x000D4000
+UMC_BASE__INST3_SEG1 = 0x02426400
+UMC_BASE__INST3_SEG2 = 0
+UMC_BASE__INST3_SEG3 = 0
+UMC_BASE__INST3_SEG4 = 0
+UMC_BASE__INST4_SEG0 = 0x00114000
+UMC_BASE__INST4_SEG1 = 0x02426800
+UMC_BASE__INST4_SEG2 = 0
+UMC_BASE__INST4_SEG3 = 0
+UMC_BASE__INST4_SEG4 = 0
+UMC_BASE__INST5_SEG0 = 0x00154000
+UMC_BASE__INST5_SEG1 = 0x02426C00
+UMC_BASE__INST5_SEG2 = 0
+UMC_BASE__INST5_SEG3 = 0
+UMC_BASE__INST5_SEG4 = 0
+UMC_BASE__INST6_SEG0 = 0x00194000
+UMC_BASE__INST6_SEG1 = 0x02427000
+UMC_BASE__INST6_SEG2 = 0
+UMC_BASE__INST6_SEG3 = 0
+UMC_BASE__INST6_SEG4 = 0
+USB0_BASE__INST0_SEG0 = 0x0242A800
+USB0_BASE__INST0_SEG1 = 0x05B00000
+USB0_BASE__INST0_SEG2 = 0
+USB0_BASE__INST0_SEG3 = 0
+USB0_BASE__INST0_SEG4 = 0
+USB0_BASE__INST1_SEG0 = 0
+USB0_BASE__INST1_SEG1 = 0
+USB0_BASE__INST1_SEG2 = 0
+USB0_BASE__INST1_SEG3 = 0
+USB0_BASE__INST1_SEG4 = 0
+USB0_BASE__INST2_SEG0 = 0
+USB0_BASE__INST2_SEG1 = 0
+USB0_BASE__INST2_SEG2 = 0
+USB0_BASE__INST2_SEG3 = 0
+USB0_BASE__INST2_SEG4 = 0
+USB0_BASE__INST3_SEG0 = 0
+USB0_BASE__INST3_SEG1 = 0
+USB0_BASE__INST3_SEG2 = 0
+USB0_BASE__INST3_SEG3 = 0
+USB0_BASE__INST3_SEG4 = 0
+USB0_BASE__INST4_SEG0 = 0
+USB0_BASE__INST4_SEG1 = 0
+USB0_BASE__INST4_SEG2 = 0
+USB0_BASE__INST4_SEG3 = 0
+USB0_BASE__INST4_SEG4 = 0
+USB0_BASE__INST5_SEG0 = 0
+USB0_BASE__INST5_SEG1 = 0
+USB0_BASE__INST5_SEG2 = 0
+USB0_BASE__INST5_SEG3 = 0
+USB0_BASE__INST5_SEG4 = 0
+USB0_BASE__INST6_SEG0 = 0
+USB0_BASE__INST6_SEG1 = 0
+USB0_BASE__INST6_SEG2 = 0
+USB0_BASE__INST6_SEG3 = 0
+USB0_BASE__INST6_SEG4 = 0
+VCN_BASE__INST0_SEG0 = 0x00007800
+VCN_BASE__INST0_SEG1 = 0x00007E00
+VCN_BASE__INST0_SEG2 = 0x02403000
+VCN_BASE__INST0_SEG3 = 0
+VCN_BASE__INST0_SEG4 = 0
+VCN_BASE__INST1_SEG0 = 0x00007B00
+VCN_BASE__INST1_SEG1 = 0x00012000
+VCN_BASE__INST1_SEG2 = 0x02445000
+VCN_BASE__INST1_SEG3 = 0
+VCN_BASE__INST1_SEG4 = 0
+VCN_BASE__INST2_SEG0 = 0
+VCN_BASE__INST2_SEG1 = 0
+VCN_BASE__INST2_SEG2 = 0
+VCN_BASE__INST2_SEG3 = 0
+VCN_BASE__INST2_SEG4 = 0
+VCN_BASE__INST3_SEG0 = 0
+VCN_BASE__INST3_SEG1 = 0
+VCN_BASE__INST3_SEG2 = 0
+VCN_BASE__INST3_SEG3 = 0
+VCN_BASE__INST3_SEG4 = 0
+VCN_BASE__INST4_SEG0 = 0
+VCN_BASE__INST4_SEG1 = 0
+VCN_BASE__INST4_SEG2 = 0
+VCN_BASE__INST4_SEG3 = 0
+VCN_BASE__INST4_SEG4 = 0
+VCN_BASE__INST5_SEG0 = 0
+VCN_BASE__INST5_SEG1 = 0
+VCN_BASE__INST5_SEG2 = 0
+VCN_BASE__INST5_SEG3 = 0
+VCN_BASE__INST5_SEG4 = 0
+VCN_BASE__INST6_SEG0 = 0
+VCN_BASE__INST6_SEG1 = 0
+VCN_BASE__INST6_SEG2 = 0
+VCN_BASE__INST6_SEG3 = 0
+VCN_BASE__INST6_SEG4 = 0

tinygrad/runtime/autogen/am/vega_offsets.py

@@ -0,0 +1,774 @@
+# mypy: disable-error-code="empty-body"
+from __future__ import annotations
+import ctypes
+from typing import Literal, TypeAlias
+from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support import c
+@c.record
+class struct_IP_BASE_INSTANCE(c.Struct):
+  SIZE = 24
+  segment: c.Array[ctypes.c_uint32, Literal[6]]
+struct_IP_BASE_INSTANCE.register_fields([('segment', c.Array[ctypes.c_uint32, Literal[6]], 0)])
+@c.record
+class struct_IP_BASE(c.Struct):
+  SIZE = 144
+  instance: c.Array[struct_IP_BASE_INSTANCE, Literal[6]]
+struct_IP_BASE.register_fields([('instance', c.Array[struct_IP_BASE_INSTANCE, Literal[6]], 0)])
+MAX_INSTANCE = 6
+MAX_SEGMENT = 6
+ATHUB_BASE__INST0_SEG0 = 0x00000C20
+ATHUB_BASE__INST0_SEG1 = 0
+ATHUB_BASE__INST0_SEG2 = 0
+ATHUB_BASE__INST0_SEG3 = 0
+ATHUB_BASE__INST0_SEG4 = 0
+ATHUB_BASE__INST0_SEG5 = 0
+ATHUB_BASE__INST1_SEG0 = 0
+ATHUB_BASE__INST1_SEG1 = 0
+ATHUB_BASE__INST1_SEG2 = 0
+ATHUB_BASE__INST1_SEG3 = 0
+ATHUB_BASE__INST1_SEG4 = 0
+ATHUB_BASE__INST1_SEG5 = 0
+ATHUB_BASE__INST2_SEG0 = 0
+ATHUB_BASE__INST2_SEG1 = 0
+ATHUB_BASE__INST2_SEG2 = 0
+ATHUB_BASE__INST2_SEG3 = 0
+ATHUB_BASE__INST2_SEG4 = 0
+ATHUB_BASE__INST2_SEG5 = 0
+ATHUB_BASE__INST3_SEG0 = 0
+ATHUB_BASE__INST3_SEG1 = 0
+ATHUB_BASE__INST3_SEG2 = 0
+ATHUB_BASE__INST3_SEG3 = 0
+ATHUB_BASE__INST3_SEG4 = 0
+ATHUB_BASE__INST3_SEG5 = 0
+ATHUB_BASE__INST4_SEG0 = 0
+ATHUB_BASE__INST4_SEG1 = 0
+ATHUB_BASE__INST4_SEG2 = 0
+ATHUB_BASE__INST4_SEG3 = 0
+ATHUB_BASE__INST4_SEG4 = 0
+ATHUB_BASE__INST4_SEG5 = 0
+ATHUB_BASE__INST5_SEG0 = 0
+ATHUB_BASE__INST5_SEG1 = 0
+ATHUB_BASE__INST5_SEG2 = 0
+ATHUB_BASE__INST5_SEG3 = 0
+ATHUB_BASE__INST5_SEG4 = 0
+ATHUB_BASE__INST5_SEG5 = 0
+CLK_BASE__INST0_SEG0 = 0x00016C00
+CLK_BASE__INST0_SEG1 = 0x00016E00
+CLK_BASE__INST0_SEG2 = 0x00017000
+CLK_BASE__INST0_SEG3 = 0x00017200
+CLK_BASE__INST0_SEG4 = 0x0001B000
+CLK_BASE__INST0_SEG5 = 0x0001B200
+CLK_BASE__INST1_SEG0 = 0
+CLK_BASE__INST1_SEG1 = 0
+CLK_BASE__INST1_SEG2 = 0
+CLK_BASE__INST1_SEG3 = 0
+CLK_BASE__INST1_SEG4 = 0
+CLK_BASE__INST1_SEG5 = 0
+CLK_BASE__INST2_SEG0 = 0
+CLK_BASE__INST2_SEG1 = 0
+CLK_BASE__INST2_SEG2 = 0
+CLK_BASE__INST2_SEG3 = 0
+CLK_BASE__INST2_SEG4 = 0
+CLK_BASE__INST2_SEG5 = 0
+CLK_BASE__INST3_SEG0 = 0
+CLK_BASE__INST3_SEG1 = 0
+CLK_BASE__INST3_SEG2 = 0
+CLK_BASE__INST3_SEG3 = 0
+CLK_BASE__INST3_SEG4 = 0
+CLK_BASE__INST3_SEG5 = 0
+CLK_BASE__INST4_SEG0 = 0
+CLK_BASE__INST4_SEG1 = 0
+CLK_BASE__INST4_SEG2 = 0
+CLK_BASE__INST4_SEG3 = 0
+CLK_BASE__INST4_SEG4 = 0
+CLK_BASE__INST4_SEG5 = 0
+CLK_BASE__INST5_SEG0 = 0
+CLK_BASE__INST5_SEG1 = 0
+CLK_BASE__INST5_SEG2 = 0
+CLK_BASE__INST5_SEG3 = 0
+CLK_BASE__INST5_SEG4 = 0
+CLK_BASE__INST5_SEG5 = 0
+DCE_BASE__INST0_SEG0 = 0x00000012
+DCE_BASE__INST0_SEG1 = 0x000000C0
+DCE_BASE__INST0_SEG2 = 0x000034C0
+DCE_BASE__INST0_SEG3 = 0
+DCE_BASE__INST0_SEG4 = 0
+DCE_BASE__INST0_SEG5 = 0
+DCE_BASE__INST1_SEG0 = 0
+DCE_BASE__INST1_SEG1 = 0
+DCE_BASE__INST1_SEG2 = 0
+DCE_BASE__INST1_SEG3 = 0
+DCE_BASE__INST1_SEG4 = 0
+DCE_BASE__INST1_SEG5 = 0
+DCE_BASE__INST2_SEG0 = 0
+DCE_BASE__INST2_SEG1 = 0
+DCE_BASE__INST2_SEG2 = 0
+DCE_BASE__INST2_SEG3 = 0
+DCE_BASE__INST2_SEG4 = 0
+DCE_BASE__INST2_SEG5 = 0
+DCE_BASE__INST3_SEG0 = 0
+DCE_BASE__INST3_SEG1 = 0
+DCE_BASE__INST3_SEG2 = 0
+DCE_BASE__INST3_SEG3 = 0
+DCE_BASE__INST3_SEG4 = 0
+DCE_BASE__INST3_SEG5 = 0
+DCE_BASE__INST4_SEG0 = 0
+DCE_BASE__INST4_SEG1 = 0
+DCE_BASE__INST4_SEG2 = 0
+DCE_BASE__INST4_SEG3 = 0
+DCE_BASE__INST4_SEG4 = 0
+DCE_BASE__INST4_SEG5 = 0
+DCE_BASE__INST5_SEG0 = 0
+DCE_BASE__INST5_SEG1 = 0
+DCE_BASE__INST5_SEG2 = 0
+DCE_BASE__INST5_SEG3 = 0
+DCE_BASE__INST5_SEG4 = 0
+DCE_BASE__INST5_SEG5 = 0
+DF_BASE__INST0_SEG0 = 0x00007000
+DF_BASE__INST0_SEG1 = 0
+DF_BASE__INST0_SEG2 = 0
+DF_BASE__INST0_SEG3 = 0
+DF_BASE__INST0_SEG4 = 0
+DF_BASE__INST0_SEG5 = 0
+DF_BASE__INST1_SEG0 = 0
+DF_BASE__INST1_SEG1 = 0
+DF_BASE__INST1_SEG2 = 0
+DF_BASE__INST1_SEG3 = 0
+DF_BASE__INST1_SEG4 = 0
+DF_BASE__INST1_SEG5 = 0
+DF_BASE__INST2_SEG0 = 0
+DF_BASE__INST2_SEG1 = 0
+DF_BASE__INST2_SEG2 = 0
+DF_BASE__INST2_SEG3 = 0
+DF_BASE__INST2_SEG4 = 0
+DF_BASE__INST2_SEG5 = 0
+DF_BASE__INST3_SEG0 = 0
+DF_BASE__INST3_SEG1 = 0
+DF_BASE__INST3_SEG2 = 0
+DF_BASE__INST3_SEG3 = 0
+DF_BASE__INST3_SEG4 = 0
+DF_BASE__INST3_SEG5 = 0
+DF_BASE__INST4_SEG0 = 0
+DF_BASE__INST4_SEG1 = 0
+DF_BASE__INST4_SEG2 = 0
+DF_BASE__INST4_SEG3 = 0
+DF_BASE__INST4_SEG4 = 0
+DF_BASE__INST4_SEG5 = 0
+DF_BASE__INST5_SEG0 = 0
+DF_BASE__INST5_SEG1 = 0
+DF_BASE__INST5_SEG2 = 0
+DF_BASE__INST5_SEG3 = 0
+DF_BASE__INST5_SEG4 = 0
+DF_BASE__INST5_SEG5 = 0
+FUSE_BASE__INST0_SEG0 = 0x00017400
+FUSE_BASE__INST0_SEG1 = 0
+FUSE_BASE__INST0_SEG2 = 0
+FUSE_BASE__INST0_SEG3 = 0
+FUSE_BASE__INST0_SEG4 = 0
+FUSE_BASE__INST0_SEG5 = 0
+FUSE_BASE__INST1_SEG0 = 0
+FUSE_BASE__INST1_SEG1 = 0
+FUSE_BASE__INST1_SEG2 = 0
+FUSE_BASE__INST1_SEG3 = 0
+FUSE_BASE__INST1_SEG4 = 0
+FUSE_BASE__INST1_SEG5 = 0
+FUSE_BASE__INST2_SEG0 = 0
+FUSE_BASE__INST2_SEG1 = 0
+FUSE_BASE__INST2_SEG2 = 0
+FUSE_BASE__INST2_SEG3 = 0
+FUSE_BASE__INST2_SEG4 = 0
+FUSE_BASE__INST2_SEG5 = 0
+FUSE_BASE__INST3_SEG0 = 0
+FUSE_BASE__INST3_SEG1 = 0
+FUSE_BASE__INST3_SEG2 = 0
+FUSE_BASE__INST3_SEG3 = 0
+FUSE_BASE__INST3_SEG4 = 0
+FUSE_BASE__INST3_SEG5 = 0
+FUSE_BASE__INST4_SEG0 = 0
+FUSE_BASE__INST4_SEG1 = 0
+FUSE_BASE__INST4_SEG2 = 0
+FUSE_BASE__INST4_SEG3 = 0
+FUSE_BASE__INST4_SEG4 = 0
+FUSE_BASE__INST4_SEG5 = 0
+FUSE_BASE__INST5_SEG0 = 0
+FUSE_BASE__INST5_SEG1 = 0
+FUSE_BASE__INST5_SEG2 = 0
+FUSE_BASE__INST5_SEG3 = 0
+FUSE_BASE__INST5_SEG4 = 0
+FUSE_BASE__INST5_SEG5 = 0
+GC_BASE__INST0_SEG0 = 0x00002000
+GC_BASE__INST0_SEG1 = 0x0000A000
+GC_BASE__INST0_SEG2 = 0
+GC_BASE__INST0_SEG3 = 0
+GC_BASE__INST0_SEG4 = 0
+GC_BASE__INST0_SEG5 = 0
+GC_BASE__INST1_SEG0 = 0
+GC_BASE__INST1_SEG1 = 0
+GC_BASE__INST1_SEG2 = 0
+GC_BASE__INST1_SEG3 = 0
+GC_BASE__INST1_SEG4 = 0
+GC_BASE__INST1_SEG5 = 0
+GC_BASE__INST2_SEG0 = 0
+GC_BASE__INST2_SEG1 = 0
+GC_BASE__INST2_SEG2 = 0
+GC_BASE__INST2_SEG3 = 0
+GC_BASE__INST2_SEG4 = 0
+GC_BASE__INST2_SEG5 = 0
+GC_BASE__INST3_SEG0 = 0
+GC_BASE__INST3_SEG1 = 0
+GC_BASE__INST3_SEG2 = 0
+GC_BASE__INST3_SEG3 = 0
+GC_BASE__INST3_SEG4 = 0
+GC_BASE__INST3_SEG5 = 0
+GC_BASE__INST4_SEG0 = 0
+GC_BASE__INST4_SEG1 = 0
+GC_BASE__INST4_SEG2 = 0
+GC_BASE__INST4_SEG3 = 0
+GC_BASE__INST4_SEG4 = 0
+GC_BASE__INST4_SEG5 = 0
+GC_BASE__INST5_SEG0 = 0
+GC_BASE__INST5_SEG1 = 0
+GC_BASE__INST5_SEG2 = 0
+GC_BASE__INST5_SEG3 = 0
+GC_BASE__INST5_SEG4 = 0
+GC_BASE__INST5_SEG5 = 0
+HDP_BASE__INST0_SEG0 = 0x00000F20
+HDP_BASE__INST0_SEG1 = 0
+HDP_BASE__INST0_SEG2 = 0
+HDP_BASE__INST0_SEG3 = 0
+HDP_BASE__INST0_SEG4 = 0
+HDP_BASE__INST0_SEG5 = 0
+HDP_BASE__INST1_SEG0 = 0
+HDP_BASE__INST1_SEG1 = 0
+HDP_BASE__INST1_SEG2 = 0
+HDP_BASE__INST1_SEG3 = 0
+HDP_BASE__INST1_SEG4 = 0
+HDP_BASE__INST1_SEG5 = 0
+HDP_BASE__INST2_SEG0 = 0
+HDP_BASE__INST2_SEG1 = 0
+HDP_BASE__INST2_SEG2 = 0
+HDP_BASE__INST2_SEG3 = 0
+HDP_BASE__INST2_SEG4 = 0
+HDP_BASE__INST2_SEG5 = 0
+HDP_BASE__INST3_SEG0 = 0
+HDP_BASE__INST3_SEG1 = 0
+HDP_BASE__INST3_SEG2 = 0
+HDP_BASE__INST3_SEG3 = 0
+HDP_BASE__INST3_SEG4 = 0
+HDP_BASE__INST3_SEG5 = 0
+HDP_BASE__INST4_SEG0 = 0
+HDP_BASE__INST4_SEG1 = 0
+HDP_BASE__INST4_SEG2 = 0
+HDP_BASE__INST4_SEG3 = 0
+HDP_BASE__INST4_SEG4 = 0
+HDP_BASE__INST4_SEG5 = 0
+HDP_BASE__INST5_SEG0 = 0
+HDP_BASE__INST5_SEG1 = 0
+HDP_BASE__INST5_SEG2 = 0
+HDP_BASE__INST5_SEG3 = 0
+HDP_BASE__INST5_SEG4 = 0
+HDP_BASE__INST5_SEG5 = 0
+MMHUB_BASE__INST0_SEG0 = 0x0001A000
+MMHUB_BASE__INST0_SEG1 = 0
+MMHUB_BASE__INST0_SEG2 = 0
+MMHUB_BASE__INST0_SEG3 = 0
+MMHUB_BASE__INST0_SEG4 = 0
+MMHUB_BASE__INST0_SEG5 = 0
+MMHUB_BASE__INST1_SEG0 = 0
+MMHUB_BASE__INST1_SEG1 = 0
+MMHUB_BASE__INST1_SEG2 = 0
+MMHUB_BASE__INST1_SEG3 = 0
+MMHUB_BASE__INST1_SEG4 = 0
+MMHUB_BASE__INST1_SEG5 = 0
+MMHUB_BASE__INST2_SEG0 = 0
+MMHUB_BASE__INST2_SEG1 = 0
+MMHUB_BASE__INST2_SEG2 = 0
+MMHUB_BASE__INST2_SEG3 = 0
+MMHUB_BASE__INST2_SEG4 = 0
+MMHUB_BASE__INST2_SEG5 = 0
+MMHUB_BASE__INST3_SEG0 = 0
+MMHUB_BASE__INST3_SEG1 = 0
+MMHUB_BASE__INST3_SEG2 = 0
+MMHUB_BASE__INST3_SEG3 = 0
+MMHUB_BASE__INST3_SEG4 = 0
+MMHUB_BASE__INST3_SEG5 = 0
+MMHUB_BASE__INST4_SEG0 = 0
+MMHUB_BASE__INST4_SEG1 = 0
+MMHUB_BASE__INST4_SEG2 = 0
+MMHUB_BASE__INST4_SEG3 = 0
+MMHUB_BASE__INST4_SEG4 = 0
+MMHUB_BASE__INST4_SEG5 = 0
+MMHUB_BASE__INST5_SEG0 = 0
+MMHUB_BASE__INST5_SEG1 = 0
+MMHUB_BASE__INST5_SEG2 = 0
+MMHUB_BASE__INST5_SEG3 = 0
+MMHUB_BASE__INST5_SEG4 = 0
+MMHUB_BASE__INST5_SEG5 = 0
+MP0_BASE__INST0_SEG0 = 0x00016000
+MP0_BASE__INST0_SEG1 = 0
+MP0_BASE__INST0_SEG2 = 0
+MP0_BASE__INST0_SEG3 = 0
+MP0_BASE__INST0_SEG4 = 0
+MP0_BASE__INST0_SEG5 = 0
+MP0_BASE__INST1_SEG0 = 0
+MP0_BASE__INST1_SEG1 = 0
+MP0_BASE__INST1_SEG2 = 0
+MP0_BASE__INST1_SEG3 = 0
+MP0_BASE__INST1_SEG4 = 0
+MP0_BASE__INST1_SEG5 = 0
+MP0_BASE__INST2_SEG0 = 0
+MP0_BASE__INST2_SEG1 = 0
+MP0_BASE__INST2_SEG2 = 0
+MP0_BASE__INST2_SEG3 = 0
+MP0_BASE__INST2_SEG4 = 0
+MP0_BASE__INST2_SEG5 = 0
+MP0_BASE__INST3_SEG0 = 0
+MP0_BASE__INST3_SEG1 = 0
+MP0_BASE__INST3_SEG2 = 0
+MP0_BASE__INST3_SEG3 = 0
+MP0_BASE__INST3_SEG4 = 0
+MP0_BASE__INST3_SEG5 = 0
+MP0_BASE__INST4_SEG0 = 0
+MP0_BASE__INST4_SEG1 = 0
+MP0_BASE__INST4_SEG2 = 0
+MP0_BASE__INST4_SEG3 = 0
+MP0_BASE__INST4_SEG4 = 0
+MP0_BASE__INST4_SEG5 = 0
+MP0_BASE__INST5_SEG0 = 0
+MP0_BASE__INST5_SEG1 = 0
+MP0_BASE__INST5_SEG2 = 0
+MP0_BASE__INST5_SEG3 = 0
+MP0_BASE__INST5_SEG4 = 0
+MP0_BASE__INST5_SEG5 = 0
+MP1_BASE__INST0_SEG0 = 0x00016000
+MP1_BASE__INST0_SEG1 = 0
+MP1_BASE__INST0_SEG2 = 0
+MP1_BASE__INST0_SEG3 = 0
+MP1_BASE__INST0_SEG4 = 0
+MP1_BASE__INST0_SEG5 = 0
+MP1_BASE__INST1_SEG0 = 0
+MP1_BASE__INST1_SEG1 = 0
+MP1_BASE__INST1_SEG2 = 0
+MP1_BASE__INST1_SEG3 = 0
+MP1_BASE__INST1_SEG4 = 0
+MP1_BASE__INST1_SEG5 = 0
+MP1_BASE__INST2_SEG0 = 0
+MP1_BASE__INST2_SEG1 = 0
+MP1_BASE__INST2_SEG2 = 0
+MP1_BASE__INST2_SEG3 = 0
+MP1_BASE__INST2_SEG4 = 0
+MP1_BASE__INST2_SEG5 = 0
+MP1_BASE__INST3_SEG0 = 0
+MP1_BASE__INST3_SEG1 = 0
+MP1_BASE__INST3_SEG2 = 0
+MP1_BASE__INST3_SEG3 = 0
+MP1_BASE__INST3_SEG4 = 0
+MP1_BASE__INST3_SEG5 = 0
+MP1_BASE__INST4_SEG0 = 0
+MP1_BASE__INST4_SEG1 = 0
+MP1_BASE__INST4_SEG2 = 0
+MP1_BASE__INST4_SEG3 = 0
+MP1_BASE__INST4_SEG4 = 0
+MP1_BASE__INST4_SEG5 = 0
+MP1_BASE__INST5_SEG0 = 0
+MP1_BASE__INST5_SEG1 = 0
+MP1_BASE__INST5_SEG2 = 0
+MP1_BASE__INST5_SEG3 = 0
+MP1_BASE__INST5_SEG4 = 0
+MP1_BASE__INST5_SEG5 = 0
+NBIO_BASE__INST0_SEG0 = 0x00000000
+NBIO_BASE__INST0_SEG1 = 0x00000014
+NBIO_BASE__INST0_SEG2 = 0x00000D20
+NBIO_BASE__INST0_SEG3 = 0x00010400
+NBIO_BASE__INST0_SEG4 = 0
+NBIO_BASE__INST0_SEG5 = 0
+NBIO_BASE__INST1_SEG0 = 0
+NBIO_BASE__INST1_SEG1 = 0
+NBIO_BASE__INST1_SEG2 = 0
+NBIO_BASE__INST1_SEG3 = 0
+NBIO_BASE__INST1_SEG4 = 0
+NBIO_BASE__INST1_SEG5 = 0
+NBIO_BASE__INST2_SEG0 = 0
+NBIO_BASE__INST2_SEG1 = 0
+NBIO_BASE__INST2_SEG2 = 0
+NBIO_BASE__INST2_SEG3 = 0
+NBIO_BASE__INST2_SEG4 = 0
+NBIO_BASE__INST2_SEG5 = 0
+NBIO_BASE__INST3_SEG0 = 0
+NBIO_BASE__INST3_SEG1 = 0
+NBIO_BASE__INST3_SEG2 = 0
+NBIO_BASE__INST3_SEG3 = 0
+NBIO_BASE__INST3_SEG4 = 0
+NBIO_BASE__INST3_SEG5 = 0
+NBIO_BASE__INST4_SEG0 = 0
+NBIO_BASE__INST4_SEG1 = 0
+NBIO_BASE__INST4_SEG2 = 0
+NBIO_BASE__INST4_SEG3 = 0
+NBIO_BASE__INST4_SEG4 = 0
+NBIO_BASE__INST4_SEG5 = 0
+NBIO_BASE__INST5_SEG0 = 0
+NBIO_BASE__INST5_SEG1 = 0
+NBIO_BASE__INST5_SEG2 = 0
+NBIO_BASE__INST5_SEG3 = 0
+NBIO_BASE__INST5_SEG4 = 0
+NBIO_BASE__INST5_SEG5 = 0
+OSSSYS_BASE__INST0_SEG0 = 0x000010A0
+OSSSYS_BASE__INST0_SEG1 = 0
+OSSSYS_BASE__INST0_SEG2 = 0
+OSSSYS_BASE__INST0_SEG3 = 0
+OSSSYS_BASE__INST0_SEG4 = 0
+OSSSYS_BASE__INST0_SEG5 = 0
+OSSSYS_BASE__INST1_SEG0 = 0
+OSSSYS_BASE__INST1_SEG1 = 0
+OSSSYS_BASE__INST1_SEG2 = 0
+OSSSYS_BASE__INST1_SEG3 = 0
+OSSSYS_BASE__INST1_SEG4 = 0
+OSSSYS_BASE__INST1_SEG5 = 0
+OSSSYS_BASE__INST2_SEG0 = 0
+OSSSYS_BASE__INST2_SEG1 = 0
+OSSSYS_BASE__INST2_SEG2 = 0
+OSSSYS_BASE__INST2_SEG3 = 0
+OSSSYS_BASE__INST2_SEG4 = 0
+OSSSYS_BASE__INST2_SEG5 = 0
+OSSSYS_BASE__INST3_SEG0 = 0
+OSSSYS_BASE__INST3_SEG1 = 0
+OSSSYS_BASE__INST3_SEG2 = 0
+OSSSYS_BASE__INST3_SEG3 = 0
+OSSSYS_BASE__INST3_SEG4 = 0
+OSSSYS_BASE__INST3_SEG5 = 0
+OSSSYS_BASE__INST4_SEG0 = 0
+OSSSYS_BASE__INST4_SEG1 = 0
+OSSSYS_BASE__INST4_SEG2 = 0
+OSSSYS_BASE__INST4_SEG3 = 0
+OSSSYS_BASE__INST4_SEG4 = 0
+OSSSYS_BASE__INST4_SEG5 = 0
+OSSSYS_BASE__INST5_SEG0 = 0
+OSSSYS_BASE__INST5_SEG1 = 0
+OSSSYS_BASE__INST5_SEG2 = 0
+OSSSYS_BASE__INST5_SEG3 = 0
+OSSSYS_BASE__INST5_SEG4 = 0
+OSSSYS_BASE__INST5_SEG5 = 0
+SDMA0_BASE__INST0_SEG0 = 0x00001260
+SDMA0_BASE__INST0_SEG1 = 0
+SDMA0_BASE__INST0_SEG2 = 0
+SDMA0_BASE__INST0_SEG3 = 0
+SDMA0_BASE__INST0_SEG4 = 0
+SDMA0_BASE__INST0_SEG5 = 0
+SDMA0_BASE__INST1_SEG0 = 0
+SDMA0_BASE__INST1_SEG1 = 0
+SDMA0_BASE__INST1_SEG2 = 0
+SDMA0_BASE__INST1_SEG3 = 0
+SDMA0_BASE__INST1_SEG4 = 0
+SDMA0_BASE__INST1_SEG5 = 0
+SDMA0_BASE__INST2_SEG0 = 0
+SDMA0_BASE__INST2_SEG1 = 0
+SDMA0_BASE__INST2_SEG2 = 0
+SDMA0_BASE__INST2_SEG3 = 0
+SDMA0_BASE__INST2_SEG4 = 0
+SDMA0_BASE__INST2_SEG5 = 0
+SDMA0_BASE__INST3_SEG0 = 0
+SDMA0_BASE__INST3_SEG1 = 0
+SDMA0_BASE__INST3_SEG2 = 0
+SDMA0_BASE__INST3_SEG3 = 0
+SDMA0_BASE__INST3_SEG4 = 0
+SDMA0_BASE__INST3_SEG5 = 0
+SDMA0_BASE__INST4_SEG0 = 0
+SDMA0_BASE__INST4_SEG1 = 0
+SDMA0_BASE__INST4_SEG2 = 0
+SDMA0_BASE__INST4_SEG3 = 0
+SDMA0_BASE__INST4_SEG4 = 0
+SDMA0_BASE__INST4_SEG5 = 0
+SDMA0_BASE__INST5_SEG0 = 0
+SDMA0_BASE__INST5_SEG1 = 0
+SDMA0_BASE__INST5_SEG2 = 0
+SDMA0_BASE__INST5_SEG3 = 0
+SDMA0_BASE__INST5_SEG4 = 0
+SDMA0_BASE__INST5_SEG5 = 0
+SDMA1_BASE__INST0_SEG0 = 0x00001860
+SDMA1_BASE__INST0_SEG1 = 0
+SDMA1_BASE__INST0_SEG2 = 0
+SDMA1_BASE__INST0_SEG3 = 0
+SDMA1_BASE__INST0_SEG4 = 0
+SDMA1_BASE__INST0_SEG5 = 0
+SDMA1_BASE__INST1_SEG0 = 0
+SDMA1_BASE__INST1_SEG1 = 0
+SDMA1_BASE__INST1_SEG2 = 0
+SDMA1_BASE__INST1_SEG3 = 0
+SDMA1_BASE__INST1_SEG4 = 0
+SDMA1_BASE__INST1_SEG5 = 0
+SDMA1_BASE__INST2_SEG0 = 0
+SDMA1_BASE__INST2_SEG1 = 0
+SDMA1_BASE__INST2_SEG2 = 0
+SDMA1_BASE__INST2_SEG3 = 0
+SDMA1_BASE__INST2_SEG4 = 0
+SDMA1_BASE__INST2_SEG5 = 0
+SDMA1_BASE__INST3_SEG0 = 0
+SDMA1_BASE__INST3_SEG1 = 0
+SDMA1_BASE__INST3_SEG2 = 0
+SDMA1_BASE__INST3_SEG3 = 0
+SDMA1_BASE__INST3_SEG4 = 0
+SDMA1_BASE__INST3_SEG5 = 0
+SDMA1_BASE__INST4_SEG0 = 0
+SDMA1_BASE__INST4_SEG1 = 0
+SDMA1_BASE__INST4_SEG2 = 0
+SDMA1_BASE__INST4_SEG3 = 0
+SDMA1_BASE__INST4_SEG4 = 0
+SDMA1_BASE__INST4_SEG5 = 0
+SDMA1_BASE__INST5_SEG0 = 0
+SDMA1_BASE__INST5_SEG1 = 0
+SDMA1_BASE__INST5_SEG2 = 0
+SDMA1_BASE__INST5_SEG3 = 0
+SDMA1_BASE__INST5_SEG4 = 0
+SDMA1_BASE__INST5_SEG5 = 0
+SMUIO_BASE__INST0_SEG0 = 0x00016800
+SMUIO_BASE__INST0_SEG1 = 0x00016A00
+SMUIO_BASE__INST0_SEG2 = 0
+SMUIO_BASE__INST0_SEG3 = 0
+SMUIO_BASE__INST0_SEG4 = 0
+SMUIO_BASE__INST0_SEG5 = 0
+SMUIO_BASE__INST1_SEG0 = 0
+SMUIO_BASE__INST1_SEG1 = 0
+SMUIO_BASE__INST1_SEG2 = 0
+SMUIO_BASE__INST1_SEG3 = 0
+SMUIO_BASE__INST1_SEG4 = 0
+SMUIO_BASE__INST1_SEG5 = 0
+SMUIO_BASE__INST2_SEG0 = 0
+SMUIO_BASE__INST2_SEG1 = 0
+SMUIO_BASE__INST2_SEG2 = 0
+SMUIO_BASE__INST2_SEG3 = 0
+SMUIO_BASE__INST2_SEG4 = 0
+SMUIO_BASE__INST2_SEG5 = 0
+SMUIO_BASE__INST3_SEG0 = 0
+SMUIO_BASE__INST3_SEG1 = 0
+SMUIO_BASE__INST3_SEG2 = 0
+SMUIO_BASE__INST3_SEG3 = 0
+SMUIO_BASE__INST3_SEG4 = 0
+SMUIO_BASE__INST3_SEG5 = 0
+SMUIO_BASE__INST4_SEG0 = 0
+SMUIO_BASE__INST4_SEG1 = 0
+SMUIO_BASE__INST4_SEG2 = 0
+SMUIO_BASE__INST4_SEG3 = 0
+SMUIO_BASE__INST4_SEG4 = 0
+SMUIO_BASE__INST4_SEG5 = 0
+SMUIO_BASE__INST5_SEG0 = 0
+SMUIO_BASE__INST5_SEG1 = 0
+SMUIO_BASE__INST5_SEG2 = 0
+SMUIO_BASE__INST5_SEG3 = 0
+SMUIO_BASE__INST5_SEG4 = 0
+SMUIO_BASE__INST5_SEG5 = 0
+THM_BASE__INST0_SEG0 = 0x00016600
+THM_BASE__INST0_SEG1 = 0
+THM_BASE__INST0_SEG2 = 0
+THM_BASE__INST0_SEG3 = 0
+THM_BASE__INST0_SEG4 = 0
+THM_BASE__INST0_SEG5 = 0
+THM_BASE__INST1_SEG0 = 0
+THM_BASE__INST1_SEG1 = 0
+THM_BASE__INST1_SEG2 = 0
+THM_BASE__INST1_SEG3 = 0
+THM_BASE__INST1_SEG4 = 0
+THM_BASE__INST1_SEG5 = 0
+THM_BASE__INST2_SEG0 = 0
+THM_BASE__INST2_SEG1 = 0
+THM_BASE__INST2_SEG2 = 0
+THM_BASE__INST2_SEG3 = 0
+THM_BASE__INST2_SEG4 = 0
+THM_BASE__INST2_SEG5 = 0
+THM_BASE__INST3_SEG0 = 0
+THM_BASE__INST3_SEG1 = 0
+THM_BASE__INST3_SEG2 = 0
+THM_BASE__INST3_SEG3 = 0
+THM_BASE__INST3_SEG4 = 0
+THM_BASE__INST3_SEG5 = 0
+THM_BASE__INST4_SEG0 = 0
+THM_BASE__INST4_SEG1 = 0
+THM_BASE__INST4_SEG2 = 0
+THM_BASE__INST4_SEG3 = 0
+THM_BASE__INST4_SEG4 = 0
+THM_BASE__INST4_SEG5 = 0
+THM_BASE__INST5_SEG0 = 0
+THM_BASE__INST5_SEG1 = 0
+THM_BASE__INST5_SEG2 = 0
+THM_BASE__INST5_SEG3 = 0
+THM_BASE__INST5_SEG4 = 0
+THM_BASE__INST5_SEG5 = 0
+UMC_BASE__INST0_SEG0 = 0x00014000
+UMC_BASE__INST0_SEG1 = 0
+UMC_BASE__INST0_SEG2 = 0
+UMC_BASE__INST0_SEG3 = 0
+UMC_BASE__INST0_SEG4 = 0
+UMC_BASE__INST0_SEG5 = 0
+UMC_BASE__INST1_SEG0 = 0
+UMC_BASE__INST1_SEG1 = 0
+UMC_BASE__INST1_SEG2 = 0
+UMC_BASE__INST1_SEG3 = 0
+UMC_BASE__INST1_SEG4 = 0
+UMC_BASE__INST1_SEG5 = 0
+UMC_BASE__INST2_SEG0 = 0
+UMC_BASE__INST2_SEG1 = 0
+UMC_BASE__INST2_SEG2 = 0
+UMC_BASE__INST2_SEG3 = 0
+UMC_BASE__INST2_SEG4 = 0
+UMC_BASE__INST2_SEG5 = 0
+UMC_BASE__INST3_SEG0 = 0
+UMC_BASE__INST3_SEG1 = 0
+UMC_BASE__INST3_SEG2 = 0
+UMC_BASE__INST3_SEG3 = 0
+UMC_BASE__INST3_SEG4 = 0
+UMC_BASE__INST3_SEG5 = 0
+UMC_BASE__INST4_SEG0 = 0
+UMC_BASE__INST4_SEG1 = 0
+UMC_BASE__INST4_SEG2 = 0
+UMC_BASE__INST4_SEG3 = 0
+UMC_BASE__INST4_SEG4 = 0
+UMC_BASE__INST4_SEG5 = 0
+UMC_BASE__INST5_SEG0 = 0
+UMC_BASE__INST5_SEG1 = 0
+UMC_BASE__INST5_SEG2 = 0
+UMC_BASE__INST5_SEG3 = 0
+UMC_BASE__INST5_SEG4 = 0
+UMC_BASE__INST5_SEG5 = 0
+UVD_BASE__INST0_SEG0 = 0x00007800
+UVD_BASE__INST0_SEG1 = 0x00007E00
+UVD_BASE__INST0_SEG2 = 0
+UVD_BASE__INST0_SEG3 = 0
+UVD_BASE__INST0_SEG4 = 0
+UVD_BASE__INST0_SEG5 = 0
+UVD_BASE__INST1_SEG0 = 0
+UVD_BASE__INST1_SEG1 = 0x00009000
+UVD_BASE__INST1_SEG2 = 0
+UVD_BASE__INST1_SEG3 = 0
+UVD_BASE__INST1_SEG4 = 0
+UVD_BASE__INST1_SEG5 = 0
+UVD_BASE__INST2_SEG0 = 0
+UVD_BASE__INST2_SEG1 = 0
+UVD_BASE__INST2_SEG2 = 0
+UVD_BASE__INST2_SEG3 = 0
+UVD_BASE__INST2_SEG4 = 0
+UVD_BASE__INST2_SEG5 = 0
+UVD_BASE__INST3_SEG0 = 0
+UVD_BASE__INST3_SEG1 = 0
+UVD_BASE__INST3_SEG2 = 0
+UVD_BASE__INST3_SEG3 = 0
+UVD_BASE__INST3_SEG4 = 0
+UVD_BASE__INST3_SEG5 = 0
+UVD_BASE__INST4_SEG0 = 0
+UVD_BASE__INST4_SEG1 = 0
+UVD_BASE__INST4_SEG2 = 0
+UVD_BASE__INST4_SEG3 = 0
+UVD_BASE__INST4_SEG4 = 0
+UVD_BASE__INST4_SEG5 = 0
+UVD_BASE__INST5_SEG0 = 0
+UVD_BASE__INST5_SEG1 = 0
+UVD_BASE__INST5_SEG2 = 0
+UVD_BASE__INST5_SEG3 = 0
+UVD_BASE__INST5_SEG4 = 0
+UVD_BASE__INST5_SEG5 = 0
+VCE_BASE__INST0_SEG0 = 0x00008800
+VCE_BASE__INST0_SEG1 = 0
+VCE_BASE__INST0_SEG2 = 0
+VCE_BASE__INST0_SEG3 = 0
+VCE_BASE__INST0_SEG4 = 0
+VCE_BASE__INST0_SEG5 = 0
+VCE_BASE__INST1_SEG0 = 0
+VCE_BASE__INST1_SEG1 = 0
+VCE_BASE__INST1_SEG2 = 0
+VCE_BASE__INST1_SEG3 = 0
+VCE_BASE__INST1_SEG4 = 0
+VCE_BASE__INST1_SEG5 = 0
+VCE_BASE__INST2_SEG0 = 0
+VCE_BASE__INST2_SEG1 = 0
+VCE_BASE__INST2_SEG2 = 0
+VCE_BASE__INST2_SEG3 = 0
+VCE_BASE__INST2_SEG4 = 0
+VCE_BASE__INST2_SEG5 = 0
+VCE_BASE__INST3_SEG0 = 0
+VCE_BASE__INST3_SEG1 = 0
+VCE_BASE__INST3_SEG2 = 0
+VCE_BASE__INST3_SEG3 = 0
+VCE_BASE__INST3_SEG4 = 0
+VCE_BASE__INST3_SEG5 = 0
+VCE_BASE__INST4_SEG0 = 0
+VCE_BASE__INST4_SEG1 = 0
+VCE_BASE__INST4_SEG2 = 0
+VCE_BASE__INST4_SEG3 = 0
+VCE_BASE__INST4_SEG4 = 0
+VCE_BASE__INST4_SEG5 = 0
+VCE_BASE__INST5_SEG0 = 0
+VCE_BASE__INST5_SEG1 = 0
+VCE_BASE__INST5_SEG2 = 0
+VCE_BASE__INST5_SEG3 = 0
+VCE_BASE__INST5_SEG4 = 0
+VCE_BASE__INST5_SEG5 = 0
+XDMA_BASE__INST0_SEG0 = 0x00003400
+XDMA_BASE__INST0_SEG1 = 0
+XDMA_BASE__INST0_SEG2 = 0
+XDMA_BASE__INST0_SEG3 = 0
+XDMA_BASE__INST0_SEG4 = 0
+XDMA_BASE__INST0_SEG5 = 0
+XDMA_BASE__INST1_SEG0 = 0
+XDMA_BASE__INST1_SEG1 = 0
+XDMA_BASE__INST1_SEG2 = 0
+XDMA_BASE__INST1_SEG3 = 0
+XDMA_BASE__INST1_SEG4 = 0
+XDMA_BASE__INST1_SEG5 = 0
+XDMA_BASE__INST2_SEG0 = 0
+XDMA_BASE__INST2_SEG1 = 0
+XDMA_BASE__INST2_SEG2 = 0
+XDMA_BASE__INST2_SEG3 = 0
+XDMA_BASE__INST2_SEG4 = 0
+XDMA_BASE__INST2_SEG5 = 0
+XDMA_BASE__INST3_SEG0 = 0
+XDMA_BASE__INST3_SEG1 = 0
+XDMA_BASE__INST3_SEG2 = 0
+XDMA_BASE__INST3_SEG3 = 0
+XDMA_BASE__INST3_SEG4 = 0
+XDMA_BASE__INST3_SEG5 = 0
+XDMA_BASE__INST4_SEG0 = 0
+XDMA_BASE__INST4_SEG1 = 0
+XDMA_BASE__INST4_SEG2 = 0
+XDMA_BASE__INST4_SEG3 = 0
+XDMA_BASE__INST4_SEG4 = 0
+XDMA_BASE__INST4_SEG5 = 0
+XDMA_BASE__INST5_SEG0 = 0
+XDMA_BASE__INST5_SEG1 = 0
+XDMA_BASE__INST5_SEG2 = 0
+XDMA_BASE__INST5_SEG3 = 0
+XDMA_BASE__INST5_SEG4 = 0
+XDMA_BASE__INST5_SEG5 = 0
+RSMU_BASE__INST0_SEG0 = 0x00012000
+RSMU_BASE__INST0_SEG1 = 0
+RSMU_BASE__INST0_SEG2 = 0
+RSMU_BASE__INST0_SEG3 = 0
+RSMU_BASE__INST0_SEG4 = 0
+RSMU_BASE__INST0_SEG5 = 0
+RSMU_BASE__INST1_SEG0 = 0
+RSMU_BASE__INST1_SEG1 = 0
+RSMU_BASE__INST1_SEG2 = 0
+RSMU_BASE__INST1_SEG3 = 0
+RSMU_BASE__INST1_SEG4 = 0
+RSMU_BASE__INST1_SEG5 = 0
+RSMU_BASE__INST2_SEG0 = 0
+RSMU_BASE__INST2_SEG1 = 0
+RSMU_BASE__INST2_SEG2 = 0
+RSMU_BASE__INST2_SEG3 = 0
+RSMU_BASE__INST2_SEG4 = 0
+RSMU_BASE__INST2_SEG5 = 0
+RSMU_BASE__INST3_SEG0 = 0
+RSMU_BASE__INST3_SEG1 = 0
+RSMU_BASE__INST3_SEG2 = 0
+RSMU_BASE__INST3_SEG3 = 0
+RSMU_BASE__INST3_SEG4 = 0
+RSMU_BASE__INST3_SEG5 = 0
+RSMU_BASE__INST4_SEG0 = 0
+RSMU_BASE__INST4_SEG1 = 0
+RSMU_BASE__INST4_SEG2 = 0
+RSMU_BASE__INST4_SEG3 = 0
+RSMU_BASE__INST4_SEG4 = 0
+RSMU_BASE__INST4_SEG5 = 0
+RSMU_BASE__INST5_SEG0 = 0
+RSMU_BASE__INST5_SEG1 = 0
+RSMU_BASE__INST5_SEG2 = 0
+RSMU_BASE__INST5_SEG3 = 0
+RSMU_BASE__INST5_SEG4 = 0
+RSMU_BASE__INST5_SEG5 = 0

tinygrad/runtime/ops_metal.py

@@ -147,10 +147,10 @@ class MetalProgram:
     encoder.endEncoding()
     command_buffer.setLabel(to_ns_str(self.name)) # TODO: is this always needed?
     command_buffer.commit()
+    self.dev.mtl_buffers_in_flight.append(command_buffer)
     if wait:
       wait_check(command_buffer)
       return command_buffer.GPUEndTime() - command_buffer.GPUStartTime()
-    self.dev.mtl_buffers_in_flight.append(command_buffer)
 
 class MetalBuffer:
   def __init__(self, buf:metal.MTLBuffer, size:int, offset=0): self.buf, self.size, self.offset = buf, size, offset

tinygrad/runtime/ops_python.py

@@ -18,8 +18,10 @@ def _load(m, i, dtype: DType):
   return from_storage_scalar(m[i], dtype)
 
 def load(inp, j, dtype: DType):
-  if len(inp) == 2: return [_load(m, x+j if x is not None else None, dtype) if gate else default for (m,x,gate),default in zip(*inp)]
-  return [_load(m, x+j if x is not None else None, dtype) for m,x,_ in inp[0]]
+  # inp is [index_values, gates, alts] (gated load with alt) or [index_values] (plain load)
+  if len(inp) == 3: return [_load(m, x+j if x is not None else None, dtype) if g else default
+                            for (m,x),g,default in zip(inp[0], inp[1], inp[2])]
+  return [_load(m, x+j if x is not None else None, dtype) for m,x in inp[0]]
 
 def _store(m, i, v, dtype: DType):
   if i < 0 or i >= len(m): raise IndexError(f"store out of bounds, size is {len(m)}, access is {i}, value is {v}")
@@ -67,8 +69,10 @@ class PythonProgram:
           continue
         assert dtype is not None, f"{uop} is missing a dtype"
         if uop is Ops.STORE:
+          # gate is at src[2] for gated stores; default to all-True for plain stores
+          gates = src_values[2] if len(src_values) >= 3 else [True]*len(src_values[0])
           for j,val in enumerate(src_values[1] if src_dtypes[1].count > 1 else [src_values[1]]):
-            for (m,o,g),v in zip(src_values[0], val):
+            for (m,o),g,v in zip(src_values[0], gates, val):
               if g: _store(m, o+j, v, src_dtypes[1].scalar())
           i += 1
           continue
@@ -98,7 +102,7 @@ class PythonProgram:
               else: ret.append((m, ox*4 + oy*src_dtypes[0].shape[1]*4))
           else:
             for m,o in zip(src_values[0], src_values[1]): ret.append((m,o))
-          values[i] = [(m,o,g) for (m,o),g in zip(ret, src_values[2] if len(src_values) == 3 else [True]*len(ret))] # set the gate last
+          values[i] = ret
         elif uop is Ops.CAST and isinstance(dtype, PtrDType):
           values[i] = src_values[0]
         elif uop is Ops.RANGE:

tinygrad/runtime/support/am/amdev.py

@@ -177,10 +177,12 @@ class AMDev:
     # Init hw for IP blocks where it is needed
     if not self.partial_boot:
       if self.psp.is_sos_alive() and self.smu.is_smu_alive():
+        self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) & ~pci.PCI_COMMAND_MASTER, 2)
         if self.is_hive():
           if reset_mode: return # in reset mode, do not raise
           raise RuntimeError("Malformed state. Use extra/amdpci/hive_reset.py to reset the hive")
         self.smu.mode1_reset()
+      self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
       self.init_hw(self.soc, self.gmc, self.ih, self.psp, self.smu)
 
     # Booting done
@@ -188,7 +190,6 @@ class AMDev:
 
     # Re-initialize main blocks
     self.init_hw(self.gfx, self.sdma)
-    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
 
     if (max_power:=getenv("AM_POWER_LIMIT", 0.0)) > 0:
       self.smu.set_power_limit(max_power)

tinygrad/runtime/support/am/ip.py

@@ -304,9 +304,10 @@ class AM_GFX(AM_IP):
   def reset_mec(self):
     self._dequeue_hqds()
 
-    for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(soft_reset_cp=1, soft_reset_cpc=1, inst=xcc)
-    time.sleep(0.05)
-    for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(0x0, inst=xcc)
+    if self.adev.ip_ver[am.GC_HWIP] < (10,0,0): # gfx10+ uses mec_pipe0_reset
+      for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(soft_reset_cp=1, soft_reset_cpc=1, inst=xcc)
+      time.sleep(0.05)
+      for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(0x0, inst=xcc)
 
     self._config_mec()
     self._enable_mec()

tinygrad/runtime/support/amd.py

@@ -62,7 +62,7 @@ def import_soc(ip):
   # rocm soc headers have more profiling enums than upstream linux
   return type("SOC", (object,), import_header(f"aqlprofile/linux/{({9: 'vega10', 10: 'navi10', 11: 'soc21', 12: 'soc24'}[ip[0]])}_enum.h", ROCM_URL))
 
-def import_ip_offsets(ip): return type("IPOFF", (object,), import_header(f"include/{('sienna_cichlid' if ip[0] > 9 else 'vega20')}_ip_offset.h"))
+def import_ip_offsets(ip): return getattr(tinygrad.runtime.autogen.am, f"{'navi' if ip[0] > 9 else 'vega'}_offsets")
 
 def import_pmc(ip) -> dict[str, tuple[str, int]]:
   res:dict[str, tuple[str, int]] = {}

tinygrad/runtime/support/nv/nvdev.py

@@ -73,7 +73,6 @@ class NVMemoryManager(MemoryManager):
 class NVDev:
   def __init__(self, pci_dev:PCIDevice):
     self.pci_dev, self.devfmt, self.mmio = pci_dev, pci_dev.pcibus, pci_dev.map_bar(0, fmt='I')
-    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
 
     self.smi_dev, self.is_booting, self.is_err_state = False, True, False
     self._early_ip_init()
@@ -104,10 +103,12 @@ class NVDev:
     self.include("src/common/inc/swref/published/ampere/ga102/dev_gc6_island_addendum.h")
 
     if (needs_reset:=self.reg("NV_PFB_PRI_MMU_WPR2_ADDR_HI").read() != 0):
+      self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) & ~pci.PCI_COMMAND_MASTER, 2)
       if DEBUG >= 2: print(f"nv {self.devfmt}: WPR2 is up. Issuing a full reset.", flush=True)
       self.pci_dev.reset()
       time.sleep(0.1) # wait until device can respond again
 
+    self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
     self.chip_id = self.reg("NV_PMC_BOOT_0").read()
     self.chip_details = self.reg("NV_PMC_BOOT_42").read_bitfields()
     self.chip_name = {0x17: "GA1", 0x19: "AD1", 0x1b: "GB2"}[self.chip_details['architecture']] + f"{self.chip_details['implementation']:02d}"

tinygrad/runtime/support/system.py

@@ -196,6 +196,9 @@ class PCIDevice:
   def reset(self): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{self.pcibus}/reset'")
   def read_config(self, offset:int, size:int): return int.from_bytes(self.cfg_fd.read(size, binary=True, offset=offset), byteorder='little')
   def write_config(self, offset:int, value:int, size:int): self.cfg_fd.write(value.to_bytes(size, byteorder='little'), binary=True, offset=offset)
+  def write_config_flush(self, offset:int, value:int, size:int):
+    self.write_config(offset, value, size)
+    self.read_config(offset, size)
 
   @functools.cache
   def bar_fd(self, bar_idx:int) -> FileIOInterface:

tinygrad/schedule/indexing.py

@@ -5,7 +5,7 @@ from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches
 from tinygrad.uop.ops import consumer_map_from_toposort, gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
-from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored
+from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC
 
 ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
                      Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
@@ -265,7 +265,9 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     # 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)
 
-  tsink = graph_rewrite(tsink, pm_apply_rangeify, ctx=rctx, bottom_up=True, name="apply rangeify")
+  # NOTE: SPEC=3 is broken here with shape
+  with Context(SPEC=min(SPEC.value, 2)):
+    tsink = graph_rewrite(tsink, pm_apply_rangeify, ctx=rctx, bottom_up=True, name="apply rangeify")
   return tsink, rctx
 
 def render_ranges(*rngs_list, realized) -> str:

tinygrad/schedule/rangeify.py

@@ -442,7 +442,7 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
   (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.AFTER, name="y"))), lambda x,y: x.after(*y.src[1:])),
 
   # remove invalid writes
-  (UPat(Ops.STORE, src=(UPat(), UPat(Ops.CONTIGUOUS, src=(UPat(Ops.CONST, arg=Invalid),))), allow_any_len=True), lambda: UOp(Ops.NOOP)),
+  (UPat(Ops.STORE, src=(UPat(), UPat(Ops.CONTIGUOUS, src=(UPat(Ops.CONST, arg=Invalid),)))), lambda: UOp(Ops.NOOP)),
   (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.NOOP, src=()))), lambda x: x),
   (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.END, src=(UPat(Ops.NOOP, src=()),), allow_any_len=True))), lambda x: x),
 ])

tinygrad/tensor.py

@@ -822,19 +822,27 @@ class Tensor(OpMixin):
     """
     Returns a tensor with `num_samples` indices sampled from a multinomial distribution weighted by `self`.
 
-    NOTE: `replacement=False` for `num_samples > 1` is not supported yet.
     ```python exec="true" source="above" session="tensor" result="python"
     Tensor.manual_seed(42)
     t = Tensor([1, 2, 3, 4])
     print(t.multinomial(20, replacement=True).numpy())
     ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    Tensor.manual_seed(42)
+    t = Tensor([1, 2, 3, 4])
+    print(t.multinomial(3, replacement=False).numpy())
+    ```
     """
     assert 1 <= self.ndim <= 2 and num_samples > 0, f"{self.ndim=} must be 1 or 2 dim, {num_samples=} must be positive"
-    assert replacement or num_samples == 1, "no replacement only supports num_samples = 1"
     weight = self.unsqueeze(0) if self.ndim == 1 else self
-    cdf = (cw := weight.cumsum(1).float()) / cw[:, -1].unsqueeze(1)
-    unif_samples = Tensor.rand(num_samples, cdf.shape[0], 1).to(self.device)
-    indices = (unif_samples.expand((-1, -1, cdf.shape[1])) >= cdf).sum(2).permute((1, 0))
+    assert replacement or num_samples <= weight.shape[1], "no replacement samples must not exceed population size"
+    if replacement or num_samples == 1:
+      cdf = (cw := weight.cumsum(1).float()) / cw[:, -1].unsqueeze(1)
+      unif_samples = Tensor.rand(num_samples, cdf.shape[0], 1).to(self.device)
+      indices = (unif_samples.expand((-1, -1, cdf.shape[1])) >= cdf).sum(2).permute((1, 0))
+    else:
+      # Efraimidis–Spirakis
+      indices = (weight.rand_like(dtype=dtypes.float32).log2() / weight).topk(num_samples, dim=1)[1]
     return (indices.squeeze(0) if self.ndim == 1 else indices).cast(dtypes.int32)
 
   # ***** toposort and backward pass *****
@@ -1323,6 +1331,18 @@ class Tensor(OpMixin):
     a, b = self._broadcasted(x, reverse)
     return a - a.div(b, rounding_mode="floor") * b
 
+  def fmod(self, x:Tensor|ConstType) -> Tensor:
+    """
+    C-style remainder of `self` divided by `x` (sign follows the dividend), using truncating division.
+    Differs from `mod`/`%`, which uses Python floor remainder.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-4, 7, 5, 4, -7, 8]).fmod(Tensor([2, -3, 8, -2, 3, 5])).numpy())
+    ```
+    """
+    a, b = self._broadcasted(x)
+    return a - a.div(b, rounding_mode="trunc") * b
+
   def where(self:Tensor, x:Tensor|ConstType|sint, y:Tensor|ConstType|sint) -> Tensor:
     """
     Returns a tensor of elements selected from either `x` or `y`, depending on `self`.

tinygrad/uop/ops.py

@@ -9,7 +9,7 @@ from tinygrad.dtype import ConstFloat, PyConst, storage_fmt_for_dtype, to_storag
 from tinygrad.device import Buffer, MultiBuffer, canonicalize_device
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
 from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, VIZ, SPEC, CAPTURE_PROCESS_REPLAY
-from tinygrad.helpers import strip_parens, colored, ansilen, printable
+from tinygrad.helpers import colored, ansilen, printable
 if TYPE_CHECKING:
   from tinygrad.renderer import Estimates
 
@@ -26,7 +26,7 @@ axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL:
 axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
                AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
 
-range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1,
+range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1,
                Ops.COPY: 2, Ops.BUFFER_VIEW: 1, Ops.LINEAR: 0}
 
 # https://en.wikipedia.org/wiki/Identity_element
@@ -77,16 +77,6 @@ def consumer_map_from_toposort(lst:Iterable[UOp]):
       if s in ret: ret[s][u] = None
   return ret
 
-def pretty_print(x:UOp, cache=None, d=0)->str:
-  def dfs(x:UOp, cache:dict):
-    for s in x.src:
-      cache.setdefault(s, [len(cache), 0, False])[1] += 1
-      if cache[s][1] == 1: dfs(s, cache)
-  if cache is None: dfs(x, cache:={})
-  if (cx:=cache.setdefault(x, [0,0,False]))[2]: return f"{' '*d}x{cx[0]}"
-  cx[2], srcs = True, (''.join(f'\n{pretty_print(s, cache, d+2)},' for s in x.src))
-  return f"{' '*d}{f'x{cx[0]}:=' * (cx[1]>1)}{type(x).__name__}({x.op}, {x.dtype}, arg={x.argstr()}{x.tagstr()}, src=({srcs}))"
-
 class UOpMetaClass(type):
   ucache:dict[tuple, weakref.ReferenceType[UOp]] = {}
   def __call__(cls, op:Ops, dtype:DType=dtypes.void, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None,
@@ -100,7 +90,11 @@ class UOpMetaClass(type):
       buffers[created] = _buffer
     if SPEC > 1:
       from tinygrad.uop.spec import full_spec, test_pyrender
-      if SPEC > 2: test_pyrender(created)
+      if SPEC > 2:
+        # SPEC=3 checks the shape
+        _ = created._shape
+        if SPEC > 3:
+          test_pyrender(created)
       with Context(CHECK_OOB=0): fret = cast(bool|None, full_spec.rewrite(created))
       if fret is not True: raise RuntimeError(f"SPEC ISSUE {fret}: {created}")
     return created
@@ -150,7 +144,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   @recursive_property
   def key(self) -> bytes:
     return hashlib.sha256(str((self.op, self.dtype, self.arg)).encode() + b"".join([s.key for s in self.src])).digest()
-  def __repr__(self): return pretty_print(self)
+  def __repr__(self):
+    from tinygrad.uop.render import pretty_print
+    return pretty_print(self)
   def argstr(self):
     if self.op is Ops.REDUCE: return f'({", ".join(map(str, self.arg))})'
     return f"ConstFloat({float.__repr__(self.arg)})" if isinstance(self.arg, ConstFloat) else repr(self.arg)
@@ -212,7 +208,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     match self.op:
       # late ops don't have shape
       case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
-           Ops.STACK | Ops.GEP | Ops.UNROLL | Ops.CONTRACT | Ops.SINK | Ops.END | Ops.REWRITE_ERROR | \
+           Ops.CONTRACT | Ops.SINK | Ops.END | Ops.REWRITE_ERROR | Ops.PTRCAT | Ops.ENDIF | \
            Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS | Ops.TUPLE | Ops.CALL | Ops.FUNCTION:
         return None
 
@@ -228,24 +224,26 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
         return inner_shape
 
       case Ops.CAST:
-        # when PTX casts from ptr to non ptr, remove the shape
+        # when PTX casts from ptr to non ptr, remove the shape of the buffer
         if isinstance(self.src[0].dtype, PtrDType) and not isinstance(self.src[0].dtype, ImageDType) and not isinstance(self.dtype, PtrDType):
-          return None
+          return ()
 
       case Ops.INDEX:
-        # non pointer index doesn't have a shape
-        if not isinstance(self.dtype, PtrDType): return None
-        # fully indexed doesn't have a shape. TODO: remove this
-        if self.src[0]._shape is None or len(self.src[1:]) == len(self.src[0].shape): return None
-        # pointer index
-        return self.src[0].shape[len(self.src[1:]):]
+        shp:list[sint] = []
+        for s in self.src[1:]: shp.extend(list(s.shape))
+        return tuple(shp) + self.src[0].shape[len(self.src[1:]):]
+
+      # TODO: these should have the shape of the dtype.count
+      case Ops.CONST | Ops.DEFINE_VAR: return ()
+      case Ops.GEP | Ops.STACK | Ops.VCONST | Ops.VCAT: return ()
 
       # some ops init the shape
-      case Ops.CONST | Ops.DEFINE_VAR | Ops.BIND | Ops.RANGE | Ops.SPECIAL: return ()
-      # TODO: VCONST should have the shape of the arg
-      case Ops.VCONST: return ()
+      case Ops.BIND | Ops.RANGE | Ops.SPECIAL | Ops.UNROLL: return ()
       case Ops.BUFFER: return (self.arg,)
-      case Ops.BUFFER_VIEW: return (self.arg[0],)
+      case Ops.BUFFER_VIEW:
+        # HACK: BUFFER_VIEW is used inside kernels, so we set the shape to () if it's on an INDEX
+        if self.src[0].op is Ops.INDEX: return ()
+        return (self.arg[0],)
       case Ops.CUSTOM_FUNCTION: return None
       case Ops.BUFFERIZE: return tuple([int(r.vmax+1) for r in self.src[1:]])
       case Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)
@@ -280,7 +278,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     # NOTE: ssimplify is required because the shape needs to be canonical for broadcasting and same shape checking
     if self.op in GroupOp.Movement.union({Ops.MULTI, Ops.REDUCE}):
       ps = self.src[0]._shape
-      if ps is None: raise RuntimeError(f"movement op {self.op} requires shape")
+      if ps is None: raise RuntimeError(f"movement op {self.op} requires shape, {self.src[0].op} doesn't have one")
       match self.op:
         case Ops.RESHAPE:
           if not all(x >= 0 for x in self.marg): raise ValueError(f"shape can't contain negative numbers {self.marg}")
@@ -316,7 +314,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     if self.op in GroupOp.ALU.union({Ops.CAST, Ops.COPY, Ops.NOOP, Ops.GROUP, Ops.SINK, Ops.ALLREDUCE, Ops.STORE}):
       input_shapes = [x._shape for x in self.src if x._shape is not None]
       if len(input_shapes) == 0: return None
-      if not all_same(input_shapes): raise RuntimeError(f"shape mismatch at {self.op}: {input_shapes}")
+      if not all_same(input_shapes): raise RuntimeError(f"shape mismatch at {self.op}: {input_shapes} {[x.op for x in self.src]}")
       return input_shapes[0]
 
     # all Ops must be explicitly handled
@@ -889,6 +887,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
 
   @functools.cached_property
   def _sym_fxn(self):
+    from tinygrad.uop.render import _render_with_splits, renderer_infer
     sself = self.simplify()
     varnames = tuple(x.expr for x in sself.toposort() if x.op is Ops.DEFINE_VAR)
     # TODO: sanitize varnames, or don't use naked eval while staying fast
@@ -904,12 +903,15 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
 
   def render(self, simplify=True, pm:PatternMatcher|None=None) -> str:
     ctx: dict[UOp, str] = {}
+    from tinygrad.uop.render import renderer
     pm = renderer if pm is None else pm
     for u in (s:=self.simplify() if simplify else self).toposort():
       ctx[u] = cast(str, pm.rewrite(u, ctx=ctx))
     return ctx[s]
 
-  def pyrender(self): return pyrender(self)
+  def pyrender(self):
+    from tinygrad.uop.render import pyrender
+    return pyrender(self)
 
   # *** uop high level syntactic sugar ***
 
@@ -953,9 +955,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     body = self if self.op is Ops.TUPLE else UOp.maketuple(self)
     return UOp(Ops.FUNCTION, dtypes.void, (body,)+srcs, CallInfo(grad_fxn, metadata, name, precompile, precompile_backward))
   def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
-    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(srcs)]
-    kernel = fxn(*placeholders).call(*srcs, grad_fxn=grad_fxn)
-    return [s.after(kernel) for s in srcs]
+    contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs)
+    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)]
+    kernel = fxn(*placeholders).call(*contig_srcs, grad_fxn=grad_fxn)
+    return [s.after(kernel) for s in contig_srcs]
 
 @dataclass(frozen=True)
 class KernelInfo:
@@ -1055,13 +1058,15 @@ def exec_alu(op:Ops, dtype:DType, operands, truncate_output=True):
   alu = python_alu[op](*operands)
   return truncate.get(dtype, lambda x: x)(alu) if truncate_output else alu
 
-# ***** uop helpers *****
-
-def print_uops(uops:list[UOp]):
-  uops_index = {u:i for i,u in enumerate(uops)}
-  for i,u in enumerate(uops):
-    formatted_srcs = [(uops_index[x] if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src]
-    print(f"{i:4d} {str(u.op):20s}: {multirange_str(u.ranges, color=True, pad=10)} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")
+def bitcast(x, in_dtype:DType, out_dtype:DType):
+  assert in_dtype.itemsize == out_dtype.itemsize, "bitcast itemsize mismatch"
+  in_count, out_count = in_dtype.count, out_dtype.count
+  in_vals = (x,) if in_count == 1 else tuple(x)
+  assert len(in_vals) == in_count, f"bitcast expected {in_count} values, got {len(in_vals)}"
+  packed = struct.pack(f"{in_count}{storage_fmt_for_dtype(in_dtype.scalar())}", *[to_storage_scalar(v, in_dtype.scalar()) for v in in_vals])
+  out_vals = struct.unpack(f"{out_count}{storage_fmt_for_dtype(out_dtype.scalar())}", packed)
+  ret = tuple(from_storage_scalar(v, out_dtype.scalar()) for v in out_vals)
+  return ret[0] if out_count == 1 else ret
 
 # ***** pattern matcher *****
 
@@ -1135,8 +1140,8 @@ class UPat(OpMixin):
 
   # copied from UOp
   def sink(self, *srcs:UPat|None, **kwargs): return UPat(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
-  def index(self, idx:UPat, valid:UPat|None=None, **kwargs):
-    return UPat(Ops.INDEX, self.match_dtype, (self,idx,valid) if valid is not None else (self,idx), **kwargs)
+  def index(self, *srcs:UPat|None, **kwargs):
+    return UPat(Ops.INDEX, self.match_dtype, (self,)+tuple(x for x in srcs if x is not None), **kwargs)
   def cast(self, dtype=None, **kwargs):
     if dtype is not None and self.match_dtype == (dtype,): return self
     return UPat(Ops.CAST, dtype, (self,), **kwargs)
@@ -1522,17 +1527,22 @@ pm_lower_index_dtype = PatternMatcher([
   (UPat(Ops.DEFINE_VAR, dtype=dtypes.weakint, name="u"), lambda u: u.replace(dtype=dtypes.int).cast(dtypes.weakint)),
   (UPat(Ops.BIND, src=(UPat.var("var").cast(dtypes.weakint), UPat.cvar("val").cast(dtypes.weakint))),
     lambda var,val: var.bind(val).cast(dtypes.weakint)),
-  # lower Invalid
-  (UPat.var("buf").index(UPat.var("cond").where(UPat.var("idx"), UPat(Ops.CONST, arg=Invalid))), lambda buf,idx,cond: buf.index(idx, cond, ptr=True)),
+  # lower Invalid: lift gate from INDEX up to the parent LOAD/STORE
+  (UPat(Ops.LOAD, src=(UPat.var("buf").index(UPat.var("cond").where(UPat.var("idx"),
+        UPat(Ops.CONST, arg=Invalid))).or_casted("bidx"),), allow_any_len=True, name="ld"),
+   lambda ld,buf,cond,idx,bidx: ld.replace(src=((nbidx:=buf.index(idx, ptr=True)) if bidx.op is Ops.INDEX
+                                                else bidx.replace(src=(buf.index(idx, ptr=True),)), cond) + ld.src[1:])),
+  (UPat(Ops.STORE, src=(UPat.var("buf").index(UPat.var("cond").where(UPat.var("idx"),
+        UPat(Ops.CONST, arg=Invalid))).or_casted("bidx"), UPat.var("val")), name="st"),
+   lambda st,buf,cond,idx,bidx,val: st.replace(src=(buf.index(idx, ptr=True) if bidx.op is Ops.INDEX
+                                                    else bidx.replace(src=(buf.index(idx, ptr=True),)), val, cond))),
   # remove hanging casts
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast()),), lambda buf,idx: buf.index(idx, ptr=True)),
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast(), UPat.var("valid"))),
-   lambda buf,idx,valid: buf.index(idx, valid, ptr=True)),
   (UPat((Ops.SINK, Ops.NOOP, Ops.END), name="n"),
    lambda n: n.replace(src=tuple(s.src[0] if s.op is Ops.CAST and s.dtype == dtypes.weakint else s for s in n.src))),
   # vectorized indexes (ie. images) must be int
-  (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.STACK, dtypes.long, name="vec")), allow_any_len=True, name="idx"),
-   lambda idx,vec: idx.replace(src=(idx.src[0], UOp.vectorize(*(u.cast(dtypes.int) for u in vec.src)), *idx.src[2:])))
+  (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.STACK, dtypes.long, name="vec")), name="idx"),
+   lambda idx,vec: idx.replace(src=(idx.src[0], UOp.vectorize(*(u.cast(dtypes.int) for u in vec.src)))))
 ])
 def _index_to_concrete_int(u:UOp) -> UOp: return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]
 
@@ -1551,152 +1561,6 @@ def do_unbind(ctx:dict[Variable, int], x:UOp):
   return v
 pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)])
 
-# for debug
-syms = { Ops.ADD: "+", Ops.SUB: "-", Ops.IDIV: "//", Ops.MOD: "%", Ops.SHL: "<<", Ops.SHR: ">>",
-         Ops.MUL: "*", Ops.CMPLT: "<", Ops.CMPNE: "!=", Ops.AND: "&", Ops.OR: "|", Ops.XOR: "^"}
-# comparison operators are not in here because they are chained in python, not left-associative
-precedence = {Ops.MUL:1, Ops.IDIV:1, Ops.MOD:1, Ops.ADD:2, Ops.SUB:2, Ops.SHL:3, Ops.SHR:3, Ops.AND:4, Ops.XOR:5, Ops.OR:6}
-def strip_binary_parens(x:UOp, left:str, right:str, code_for_op) -> str:
-  if x.op not in precedence: return code_for_op(left, right)
-  return code_for_op(strip_parens(left) if precedence.get(x.src[0].op,99)<=precedence[x.op] else left, strip_parens(right) if
-    precedence.get(x.src[1].op,99)<precedence[x.op] else right)
-
-def bitcast(x, in_dtype:DType, out_dtype:DType):
-  assert in_dtype.itemsize == out_dtype.itemsize, "bitcast itemsize mismatch"
-  in_count, out_count = in_dtype.count, out_dtype.count
-  in_vals = (x,) if in_count == 1 else tuple(x)
-  assert len(in_vals) == in_count, f"bitcast expected {in_count} values, got {len(in_vals)}"
-  packed = struct.pack(f"{in_count}{storage_fmt_for_dtype(in_dtype.scalar())}", *[to_storage_scalar(v, in_dtype.scalar()) for v in in_vals])
-  out_vals = struct.unpack(f"{out_count}{storage_fmt_for_dtype(out_dtype.scalar())}", packed)
-  ret = tuple(from_storage_scalar(v, out_dtype.scalar()) for v in out_vals)
-  return ret[0] if out_count == 1 else ret
-
-renderer = PatternMatcher([
-  (UPat((Ops.DEFINE_VAR,), name="x"), lambda x: x.expr),
-  (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.NOOP, name="x"))), lambda x: x.arg),
-  (UPat((Ops.SPECIAL), name="x"), lambda x: x.arg),
-  (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"),
-  (UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: str(x.arg)),
-  (UPat(Ops.UNROLL, name="x"), lambda ctx,x,u: f"UNROLL({ctx[x.src[0]]}, {u.arg})"),
-  (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({str(x.dtype)[7:]})({ctx[x.src[0]]})"),
-  (UPat(Ops.BIND, name="x"), lambda ctx,x: ctx[x.src[0]]),
-  (UPat(Ops.NEG, name="x"), lambda ctx,x: f"(-{ctx[x.src[0]]})"),
-  (UPat(Ops.RECIPROCAL, name="x"), lambda ctx,x: f"(1/{ctx[x.src[0]]})"),
-  (UPat(Ops.MAX, name="x"), lambda ctx,x: f"max({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
-  (UPat(Ops.MULACC, name="x"), lambda ctx,x: f"({ctx[x.src[0]]}*{ctx[x.src[1]]}+{ctx[x.src[2]]})"),
-  (UPat(Ops.WHERE, name="x"), lambda ctx,x: f"({ctx[x.src[1]]} if {ctx[x.src[0]]} else {ctx[x.src[2]]})"),
-  (UPat(set(syms.keys()), name="x"), lambda ctx,x: strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
-  (UPat((Ops.INDEX, Ops.BUFFERIZE), name="x"), lambda x, ctx: ''.join([f"[{strip_parens(ctx[y])}]" for y in x.src[1:]])),
-  (UPat(Ops.STACK, name="x"),
-   lambda ctx,x: f"{{{','.join([ctx[y] for y in x.src])}}}" if not x.src or not all_same(x.src) else f"{{{ctx[x.src[0]]}, ...}}"),
-  (UPat(GroupOp.All, name="x"), lambda x: str(x)),
-])
-
-renderer_infer = PatternMatcher([
-  (UPat(Ops.MOD, name="x"), lambda ctx,x: f"cmod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
-  (UPat(Ops.IDIV, name="x"), lambda ctx,x: f"cdiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
-  (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast({ctx[x.src[0]]}, {x.src[0].dtype!r}, {x.dtype!r})"),
-]) + renderer
-
-# *** pyrender ***
-
-def srcs(ctx, src): return f"({ctx[src[0]]},)" if len(src) == 1 else f"({', '.join([ctx[x] for x in src])})"
-def render_marg(ctx,x:UOp):
-  if x.op is Ops.PERMUTE: return str(x.marg)
-  if x.op is Ops.FLIP: return str(tuple([i for i,x in enumerate(x.marg) if x]))
-  pieces = []
-  if x.op in {Ops.RESHAPE, Ops.EXPAND}:
-    pieces = [f"{ctx[a] if isinstance(a, UOp) else str(a)}" for a in x.marg]
-  if x.op in {Ops.PAD, Ops.SHRINK}:
-    pieces = [f"({ctx[a[0]] if isinstance(a[0], UOp) else str(a[0])}, {ctx[a[1]] if isinstance(a[1], UOp) else str(a[1])})" for a in x.marg]
-  return f"({','.join(pieces)})" if len(pieces) != 1 else f"({pieces[0]},)"
-
-sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.UNIQUE, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY,
-         Ops.WHERE, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH}
-pm_pyrender_extra = PatternMatcher([
-  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"),
-   lambda x,u,d: f"UOp.unique_const({x.arg}, dtype={x.dtype}, device={repr(d.arg)}, unique={u.arg})"),
-  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE, name="d"),), name="x"), lambda x,d: f"UOp.const({x.dtype}, {x.arg}, device={repr(d.arg)})"),
-  (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.dtype}, {x.arg})"),
-  (UPat(Ops.DEFINE_VAR, src=(), name="x"), lambda x:
-    f"UOp.variable(\"{x.arg[0]}\", {x.arg[1]}, {x.arg[2]}{', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else ''})"),
-  (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"),
-  (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].arg}, {repr(x.arg)}, dtype={x.dtype})"),
-  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"), lambda x,u,d:
-    f"UOp.new_buffer({repr(d.arg)}, {x.arg}, {x.dtype}, {u.arg})"),
-  (UPat(Ops.COPY, src=(UPat(name="x"), UPat(Ops.DEVICE, name="d"))), lambda ctx,x,d: f"{ctx[x]}.copy_to_device({repr(d.arg)})"),
-  (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda ctx,x: f"UOp(Ops.CUSTOM_FUNCTION, {x.dtype}, src={srcs(ctx, x.src)}, arg={x.arg!r})"),
-  (UPat(Ops.REDUCE, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}._rop({r.arg[0]}, {r.arg[1]})" if len(r.arg[1]) else None),
-  # NOTE: range has srcs sometimes after control flow
-  (UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
-    "UOp.range("+', '.join([str(c.arg)] + [repr(y) for y in x.arg])+
-      (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else '')+")"),
-  # TODO: index shouldn't mismatch dtype
-  (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:
-   f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, "+''.join([f"{ctx[xx]}, " for xx in x.src[2:]])+
-    (f"dtype={x.dtype})" if x.src[0].dtype != x.dtype else "ptr=True)") if x.src[0].dtype.base != x.dtype else None),
-  # TODO: movement ops simplify stuff, this can break SPEC=2
-  #(UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
-  # NOTE: CMPNE doesn't work cause there's no __rne__
-  # NOTE: only match CONSTs without UNIQUE (len(src)==1), unique_const needs explicit rendering
-  (UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE}, src=(UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="y"), UPat(name="z")), name="x"),
-   lambda ctx,x,y,z: strip_binary_parens(x, str(y.arg), ctx[z], lambda a,b: f"({a}{syms[x.op]}{b})")),
-  # NOTE: sub doesn't work cause it's written as add/mul
-  (UPat(set(syms.keys())-{Ops.SUB}, src=(UPat(name="y"), UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="z")), name="x"), lambda ctx,x,y,z:
-    strip_binary_parens(x, ctx[y], str(z.arg), lambda a,b: f"({a}{syms[x.op]}{b})")),
-  (UPat(set(syms.keys())-{Ops.SUB}, name="x"), lambda ctx,x:
-    strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
-  (UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
-  (UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \
-    ([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
-])
-
-# NOTE: you can remove pm_pyrender_extra and it'll still be correct
-pm_pyrender = pm_pyrender_extra+PatternMatcher([
-  (UPat(GroupOp.All, name="u"), lambda ctx,u: f"UOp({u.op}, {u.dtype}, {srcs(ctx,u.src)}"+(f", {repr(u.arg)})" if u.arg is not None else ")")),
-])
-
-def _render_with_splits(lst:list[UOp], pm:PatternMatcher, to_render:set[UOp], split_depth:int=100) -> dict[str, str]:
-  r: dict[UOp, str] = {}
-  ret: dict[str, str] = {}
-  depth: dict[UOp, int] = {}
-  for i,u in enumerate(lst):
-    # limit inline depth to avoid "too many nested parentheses" in Python parser
-    op_depth = 1 + max([depth.get(s, 0) for s in u.src], default=0)
-    if op_depth > split_depth: to_render.add(u)
-    depth[u] = 0 if u in to_render else op_depth
-    ren = cast(str, pm.rewrite(u, ctx=r))
-    assert isinstance(ren, str)
-    if u.tag is not None: ren += f".rtag({repr(u.tag)})"
-    if u not in to_render: r[u] = ren
-    else:
-      r[u] = f"c{i}" if u is not lst[-1] else "ast"
-      ret[r[u]] = ren
-  return ret
-
-def pyrender(ast:UOp) -> str:
-  lst = list(ast.toposort())
-
-  cmap = consumer_map_from_toposort(lst)
-  not_rendered = {Ops.CONST, Ops.VCONST, Ops.DEVICE}
-  always_rendered = {Ops.PARAM, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.STACK,
-                     Ops.BUFFER, Ops.COPY, Ops.CALL, Ops.FUNCTION, Ops.WHERE, Ops.END}
-
-  to_render: set[UOp] = {ast}
-  for u in lst:
-    if u.op in {Ops.SINK}:
-      for s in u.src: to_render.add(s)
-    if u.op is Ops.STORE: to_render.add(u.src[1])
-    if u.op is Ops.REDUCE: to_render.add(u.src[0])
-    if u.op in {Ops.CALL, Ops.FUNCTION}: raise NotImplementedError("call can't be pyrendered")
-    if u.op in not_rendered: continue
-    # checking the consumers is not enough, you have to make sure it's not used twice by the one consumer
-    if len(cmap[u]) == 1 and len([x for x in list(cmap[u].keys())[0].src if x is u]) == 1 and u.op not in always_rendered: continue
-    to_render.add(u)
-
-  ret = _render_with_splits(lst, pm_pyrender, to_render)
-  return '\n'.join([f"{k} = {strip_parens(v)}" for k,v in ret.items()])
-
 # *** what was symbolic.py ***
 
 sint = int|UOp

tinygrad/uop/render.py

@@ -0,0 +1,159 @@
+from typing import cast
+from tinygrad.dtype import dtypes
+from tinygrad.uop import Ops, GroupOp
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, multirange_str, range_str, consumer_map_from_toposort
+from tinygrad.helpers import strip_parens, all_same
+
+def pretty_print(x:UOp, cache=None, d=0)->str:
+  def dfs(x:UOp, cache:dict):
+    for s in x.src:
+      cache.setdefault(s, [len(cache), 0, False])[1] += 1
+      if cache[s][1] == 1: dfs(s, cache)
+  if cache is None: dfs(x, cache:={})
+  if (cx:=cache.setdefault(x, [0,0,False]))[2]: return f"{' '*d}x{cx[0]}"
+  cx[2], srcs = True, (''.join(f'\n{pretty_print(s, cache, d+2)},' for s in x.src))
+  return f"{' '*d}{f'x{cx[0]}:=' * (cx[1]>1)}{type(x).__name__}({x.op}, {x.dtype}, arg={x.argstr()}{x.tagstr()}, src=({srcs}))"
+
+# ***** uop helpers *****
+
+def print_uops(uops:list[UOp]):
+  uops_index = {u:i for i,u in enumerate(uops)}
+  for i,u in enumerate(uops):
+    formatted_srcs = [(uops_index[x] if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src]
+    print(f"{i:4d} {str(u.op):20s}: {multirange_str(u.ranges, color=True, pad=10)} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")
+
+# for debug
+syms = { Ops.ADD: "+", Ops.SUB: "-", Ops.IDIV: "//", Ops.MOD: "%", Ops.SHL: "<<", Ops.SHR: ">>",
+         Ops.MUL: "*", Ops.CMPLT: "<", Ops.CMPNE: "!=", Ops.AND: "&", Ops.OR: "|", Ops.XOR: "^"}
+# comparison operators are not in here because they are chained in python, not left-associative
+precedence = {Ops.MUL:1, Ops.IDIV:1, Ops.MOD:1, Ops.ADD:2, Ops.SUB:2, Ops.SHL:3, Ops.SHR:3, Ops.AND:4, Ops.XOR:5, Ops.OR:6}
+def strip_binary_parens(x:UOp, left:str, right:str, code_for_op) -> str:
+  if x.op not in precedence: return code_for_op(left, right)
+  return code_for_op(strip_parens(left) if precedence.get(x.src[0].op,99)<=precedence[x.op] else left, strip_parens(right) if
+    precedence.get(x.src[1].op,99)<precedence[x.op] else right)
+
+renderer = PatternMatcher([
+  (UPat((Ops.DEFINE_VAR,), name="x"), lambda x: x.expr),
+  (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.NOOP, name="x"))), lambda x: x.arg),
+  (UPat((Ops.SPECIAL), name="x"), lambda x: x.arg),
+  (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"),
+  (UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: str(x.arg)),
+  (UPat(Ops.UNROLL, name="x"), lambda ctx,x,u: f"UNROLL({ctx[x.src[0]]}, {u.arg})"),
+  (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({str(x.dtype)[7:]})({ctx[x.src[0]]})"),
+  (UPat(Ops.BIND, name="x"), lambda ctx,x: ctx[x.src[0]]),
+  (UPat(Ops.NEG, name="x"), lambda ctx,x: f"(-{ctx[x.src[0]]})"),
+  (UPat(Ops.RECIPROCAL, name="x"), lambda ctx,x: f"(1/{ctx[x.src[0]]})"),
+  (UPat(Ops.MAX, name="x"), lambda ctx,x: f"max({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
+  (UPat(Ops.MULACC, name="x"), lambda ctx,x: f"({ctx[x.src[0]]}*{ctx[x.src[1]]}+{ctx[x.src[2]]})"),
+  (UPat(Ops.WHERE, name="x"), lambda ctx,x: f"({ctx[x.src[1]]} if {ctx[x.src[0]]} else {ctx[x.src[2]]})"),
+  (UPat(set(syms.keys()), name="x"), lambda ctx,x: strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
+  (UPat((Ops.INDEX, Ops.BUFFERIZE), name="x"), lambda x, ctx: ''.join([f"[{strip_parens(ctx[y])}]" for y in x.src[1:]])),
+  (UPat(Ops.STACK, name="x"),
+   lambda ctx,x: f"{{{','.join([ctx[y] for y in x.src])}}}" if not x.src or not all_same(x.src) else f"{{{ctx[x.src[0]]}, ...}}"),
+  (UPat(GroupOp.All, name="x"), lambda x: str(x)),
+])
+
+renderer_infer = PatternMatcher([
+  (UPat(Ops.MOD, name="x"), lambda ctx,x: f"cmod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
+  (UPat(Ops.IDIV, name="x"), lambda ctx,x: f"cdiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
+  (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast({ctx[x.src[0]]}, {x.src[0].dtype!r}, {x.dtype!r})"),
+]) + renderer
+
+# *** pyrender ***
+
+def srcs(ctx, src): return f"({ctx[src[0]]},)" if len(src) == 1 else f"({', '.join([ctx[x] for x in src])})"
+def render_marg(ctx,x:UOp):
+  if x.op is Ops.PERMUTE: return str(x.marg)
+  if x.op is Ops.FLIP: return str(tuple([i for i,x in enumerate(x.marg) if x]))
+  pieces = []
+  if x.op in {Ops.RESHAPE, Ops.EXPAND}:
+    pieces = [f"{ctx[a] if isinstance(a, UOp) else str(a)}" for a in x.marg]
+  if x.op in {Ops.PAD, Ops.SHRINK}:
+    pieces = [f"({ctx[a[0]] if isinstance(a[0], UOp) else str(a[0])}, {ctx[a[1]] if isinstance(a[1], UOp) else str(a[1])})" for a in x.marg]
+  return f"({','.join(pieces)})" if len(pieces) != 1 else f"({pieces[0]},)"
+
+sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.UNIQUE, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY,
+         Ops.WHERE, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH}
+pm_pyrender_extra = PatternMatcher([
+  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"),
+   lambda x,u,d: f"UOp.unique_const({x.arg}, dtype={x.dtype}, device={repr(d.arg)}, unique={u.arg})"),
+  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE, name="d"),), name="x"), lambda x,d: f"UOp.const({x.dtype}, {x.arg}, device={repr(d.arg)})"),
+  (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.dtype}, {x.arg})"),
+  (UPat(Ops.DEFINE_VAR, src=(), name="x"), lambda x:
+    f"UOp.variable(\"{x.arg[0]}\", {x.arg[1]}, {x.arg[2]}{', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else ''})"),
+  (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"),
+  (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].arg}, {repr(x.arg)}, dtype={x.dtype})"),
+  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"), lambda x,u,d:
+    f"UOp.new_buffer({repr(d.arg)}, {x.arg}, {x.dtype}, {u.arg})"),
+  (UPat(Ops.COPY, src=(UPat(name="x"), UPat(Ops.DEVICE, name="d"))), lambda ctx,x,d: f"{ctx[x]}.copy_to_device({repr(d.arg)})"),
+  (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda ctx,x: f"UOp(Ops.CUSTOM_FUNCTION, {x.dtype}, src={srcs(ctx, x.src)}, arg={x.arg!r})"),
+  (UPat(Ops.REDUCE, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}._rop({r.arg[0]}, {r.arg[1]})" if len(r.arg[1]) else None),
+  # NOTE: range has srcs sometimes after control flow
+  (UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
+    "UOp.range("+', '.join([str(c.arg)] + [repr(y) for y in x.arg])+
+      (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else '')+")"),
+  # TODO: index shouldn't mismatch dtype
+  (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:
+   f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, "+''.join([f"{ctx[xx]}, " for xx in x.src[2:]])+
+    (f"dtype={x.dtype})" if x.src[0].dtype != x.dtype else "ptr=True)") if x.src[0].dtype.base != x.dtype else None),
+  # TODO: movement ops simplify stuff, this can break SPEC=2
+  #(UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
+  # NOTE: CMPNE doesn't work cause there's no __rne__
+  # NOTE: only match CONSTs without UNIQUE (len(src)==1), unique_const needs explicit rendering
+  (UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE}, src=(UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="y"), UPat(name="z")), name="x"),
+   lambda ctx,x,y,z: strip_binary_parens(x, str(y.arg), ctx[z], lambda a,b: f"({a}{syms[x.op]}{b})")),
+  # NOTE: sub doesn't work cause it's written as add/mul
+  (UPat(set(syms.keys())-{Ops.SUB}, src=(UPat(name="y"), UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="z")), name="x"), lambda ctx,x,y,z:
+    strip_binary_parens(x, ctx[y], str(z.arg), lambda a,b: f"({a}{syms[x.op]}{b})")),
+  (UPat(set(syms.keys())-{Ops.SUB}, name="x"), lambda ctx,x:
+    strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
+  (UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
+  (UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \
+    ([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
+])
+
+# NOTE: you can remove pm_pyrender_extra and it'll still be correct
+pm_pyrender = pm_pyrender_extra+PatternMatcher([
+  (UPat(GroupOp.All, name="u"), lambda ctx,u: f"UOp({u.op}, {u.dtype}, {srcs(ctx,u.src)}"+(f", {repr(u.arg)})" if u.arg is not None else ")")),
+])
+
+def _render_with_splits(lst:list[UOp], pm:PatternMatcher, to_render:set[UOp], split_depth:int=100) -> dict[str, str]:
+  r: dict[UOp, str] = {}
+  ret: dict[str, str] = {}
+  depth: dict[UOp, int] = {}
+  for i,u in enumerate(lst):
+    # limit inline depth to avoid "too many nested parentheses" in Python parser
+    op_depth = 1 + max([depth.get(s, 0) for s in u.src], default=0)
+    if op_depth > split_depth: to_render.add(u)
+    depth[u] = 0 if u in to_render else op_depth
+    ren = cast(str, pm.rewrite(u, ctx=r))
+    assert isinstance(ren, str)
+    if u.tag is not None: ren += f".rtag({repr(u.tag)})"
+    if u not in to_render: r[u] = ren
+    else:
+      r[u] = f"c{i}" if u is not lst[-1] else "ast"
+      ret[r[u]] = ren
+  return ret
+
+def pyrender(ast:UOp) -> str:
+  lst = list(ast.toposort())
+
+  cmap = consumer_map_from_toposort(lst)
+  not_rendered = {Ops.CONST, Ops.VCONST, Ops.DEVICE}
+  always_rendered = {Ops.PARAM, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.STACK,
+                     Ops.BUFFER, Ops.COPY, Ops.CALL, Ops.FUNCTION, Ops.WHERE, Ops.END}
+
+  to_render: set[UOp] = {ast}
+  for u in lst:
+    if u.op in {Ops.SINK}:
+      for s in u.src: to_render.add(s)
+    if u.op is Ops.STORE: to_render.add(u.src[1])
+    if u.op is Ops.REDUCE: to_render.add(u.src[0])
+    if u.op in {Ops.CALL, Ops.FUNCTION}: raise NotImplementedError("call can't be pyrendered")
+    if u.op in not_rendered: continue
+    # checking the consumers is not enough, you have to make sure it's not used twice by the one consumer
+    if len(cmap[u]) == 1 and len([x for x in list(cmap[u].keys())[0].src if x is u]) == 1 and u.op not in always_rendered: continue
+    to_render.add(u)
+
+  ret = _render_with_splits(lst, pm_pyrender, to_render)
+  return '\n'.join([f"{k} = {strip_parens(v)}" for k,v in ret.items()])

tinygrad/uop/spec.py

@@ -1,12 +1,13 @@
 import math
 from typing import cast, Any
-from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, AxisType, KernelInfo, pyrender
+from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, AxisType, KernelInfo
+from tinygrad.uop.render import print_uops, pyrender
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid, ConstFloat
 from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata, panic, CHECK_OOB
 
-def validate_index(buf:UOp, idx:UOp, gate:UOp|None=None):
+def validate_index(buf:UOp, idx:UOp):
+  # gate now lives on LOAD/STORE; INDEX is always 2-src (buf, idx)
   if idx.op is Ops.CONST and idx.arg is Invalid: return True
-  if gate is None: gate = UOp.const(dtypes.bool, True)
   # TODO: check for overflow
   if not CHECK_OOB or isinstance(buf.dtype, ImageDType) or (sz := buf.ptrdtype.size) == -1: return True
 
@@ -16,12 +17,12 @@ def validate_index(buf:UOp, idx:UOp, gate:UOp|None=None):
   # TODO: validate these
   # WEBGPU has a BITCAST in the index, PTX casts pointer to long
   # VECTORIZE/GEP can't be properly modeled in z3 since it doesn't support vectors
-  for x in idx.toposort() | gate.toposort():
+  for x in idx.toposort():
     if x.op in {Ops.BITCAST, Ops.STACK, Ops.GEP} or (x.op is Ops.CAST and isinstance(x.src[0].dtype, PtrDType)): return True
 
   # if all is good and CHECK_OOB=1, validate with z3
   from tinygrad.uop.validate import validate_index_with_z3
-  return validate_index_with_z3(sz, idx, gate)
+  return validate_index_with_z3(sz, idx, UOp.const(dtypes.bool, True))
 
 # four specs:
 #   shared_spec  -- usable anywhere
@@ -173,10 +174,12 @@ shared_codegen_spec = PatternMatcher([
   (UPat(Ops.STACK, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.vcount and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)),
   (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()),
 
-  # LOAD(idx) / STORE(idx, val)
+  # LOAD(idx) / STORE(idx, val) / LOAD(idx, gate, alt?) gated / STORE(idx, val, gate) gated
   (UPat().index(UPat()).or_casted().load(), lambda: True),
-  (UPat().index(UPat(), UPat(dtype=dtypes.bool)).or_casted().load(), lambda: True),  # gated load (alt added in program_spec)
-  (UPat(Ops.INDEX).or_casted().store(UPat()), lambda: True),
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX).or_casted(), UPat(dtype=dtypes.bool))), lambda: True),  # gated load
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX).or_casted(), UPat(dtype=dtypes.bool), UPat())), lambda: True),  # gated load with alt
+  (UPat(Ops.STORE, dtypes.void, src=(UPat(Ops.INDEX).or_casted(), UPat())), lambda: True),
+  (UPat(Ops.STORE, dtypes.void, src=(UPat(Ops.INDEX).or_casted(), UPat(), UPat(dtype=dtypes.bool))), lambda: True),  # gated store
 
   # CUSTOM (inline and non inline)
   (UPat((Ops.CUSTOMI, Ops.CUSTOM)), lambda: True),
@@ -184,9 +187,8 @@ shared_codegen_spec = PatternMatcher([
   # assembly instruction
   (UPat(Ops.INS), lambda: True),
 
-  # INDEX (2-arg and 3-arg with bool gate)
+  # INDEX (always 2-arg, no gate; gate now lives on LOAD/STORE)
   (UPat(GroupOp.Defines|{Ops.AFTER}, name="buf").index(UPat.var("idx")), validate_index),
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines|{Ops.AFTER}, name="buf"), UPat.var("idx"), UPat.var("gate", dtype=dtypes.bool))), validate_index),
 
   # SPECIAL
   (UPat(Ops.SPECIAL, src=(UPat.var("x", (dtypes.weakint, dtypes.int32)),), name="s"), lambda s,x: s.dtype == x.dtype and isinstance(s.arg, str)),
@@ -235,8 +237,8 @@ tensor_spec = PatternMatcher([
 # ***** UOp spec in linearized programs *****
 
 program_spec = PatternMatcher([
-  # LOAD (idx, alt_value), LOAD can have an alt value, but only if the index has a gate
-  (UPat().index(UPat(), UPat(dtype=dtypes.bool)).or_casted().load(UPat()), lambda: True),
+  # LOAD (idx, gate, alt_value), LOAD can have an alt value, but only if there's a gate
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX).or_casted(), UPat(dtype=dtypes.bool), UPat())), lambda: True),
 
   # END closes ranges
   (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE)), dtype=dtypes.void), lambda: True),

tinygrad/uop/symbolic.py

@@ -442,16 +442,16 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([
   # ** load/store folding **
   (UPat.store(UPat(Ops.INDEX, name="index"), UPat.load(UPat(Ops.INDEX, name="index"))), lambda index: UOp(Ops.NOOP)),
   (UPat.store(UPat(Ops.INDEX, name="index"), UPat.var("gate").where(UPat.var("alt"),
-                                                                    UPat.load(UPat(Ops.INDEX, name="index"))), allow_any_len=True, name="store"),
-   lambda index, gate, alt, store: UOp.store(index.src[0].index(gate.where(index.src[1], UOp.invalid())), alt, *store.src[2:])),
+                                                                    UPat.load(UPat(Ops.INDEX, name="index")))),
+   lambda index, gate, alt: UOp.store(index.src[0].index(gate.where(index.src[1], UOp.invalid())), alt)),
   # fold gated LOAD/STORE
-  (UPat(Ops.STORE, src=(UPat().index(UPat.const(dtypes.weakint, Invalid)).or_casted(),), allow_any_len=True, name="x"), lambda x: UOp(Ops.NOOP)),
+  (UPat(Ops.STORE, src=(UPat().index(UPat.const(dtypes.weakint, Invalid)).or_casted(), UPat())), lambda: UOp(Ops.NOOP)),
   (UPat(Ops.LOAD, src=(UPat().index(UPat.const(dtypes.weakint, Invalid)).or_casted(),), allow_any_len=True, name="x"),
     lambda x: x.src[1] if len(x.src) > 1 else x.const_like(0)), # invalid load produces 0, or the alt value if we have one
-  (UPat(Ops.STORE, src=(UPat(), invalid_pat), allow_any_len=True), lambda i: UOp(Ops.NOOP)),
+  (UPat(Ops.STORE, src=(UPat(), invalid_pat)), lambda i: UOp(Ops.NOOP)),
   # store of where with invalid -> gated store
-  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, name="index"), UPat.var("cond").where(UPat.var("val"), invalid_pat)), allow_any_len=True, name="store"),
-   lambda index, cond, val, store, i: UOp.store(index.src[0].index(cond.where(index.src[1], UOp.invalid())), val, *store.src[2:])),
+  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, name="index"), UPat.var("cond").where(UPat.var("val"), invalid_pat))),
+   lambda index, cond, val, i: UOp.store(index.src[0].index(cond.where(index.src[1], UOp.invalid())), val)),
   ((UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()),  # 1/(x^c) -> (1/x)^c
   ((UPat.var("x") * UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()*x.reciprocal()),
   ((UPat.var("x") * UPat.cvar("c")).reciprocal(), lambda x,c: x.reciprocal()*c.reciprocal()), # 1/(x*c) -> (1/c)*(1/x)

tinygrad/viz/serve.py

@@ -39,8 +39,8 @@ class HTTPRequestHandler(BaseHTTPRequestHandler):
     # pass if client closed connection
     except (BrokenPipeError, ConnectionResetError): return
 
-from tinygrad.uop.ops import TrackedGraphRewrite, RewriteTrace, UOp, Ops, GroupOp, srender, sint, sym_infer, range_str, pyrender
-from tinygrad.uop.ops import print_uops, range_start, multirange_str
+from tinygrad.uop.ops import TrackedGraphRewrite, RewriteTrace, UOp, Ops, GroupOp, srender, sint, sym_infer, range_str, range_start, multirange_str
+from tinygrad.uop.render import print_uops, pyrender
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, ProfileProgramEvent
 from tinygrad.dtype import dtypes