cleanups

* viz/cli: add pmc printer

* cli work

* s

* linter

* pack workgroups

* add : to wgp

* counter name

.github/workflows/test.yml

@@ -793,7 +793,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        backend: [llvm, cpu, opencl, lvp]
+        backend: [llvm, cpu, opencl, lvp, x86]
 
     name: Linux (${{ matrix.backend }})
     runs-on: ubuntu-22.04
@@ -810,7 +810,7 @@ jobs:
           llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
           mesa: ${{ matrix.backend == 'lvp' && 'true' }}
       - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'DEV=CPU:LLVM' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'DEV=CL' || matrix.backend == 'lvp' && 'DEV=CPU:LVP' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'DEV=CPU:LLVM' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'DEV=CL' || matrix.backend == 'lvp' && 'DEV=CPU:LVP' || matrix.backend == 'x86' && 'DEV=CPU:X86' }}" >> $GITHUB_ENV
       - name: Check Device.DEFAULT and print some source
         run: |
           python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
@@ -960,7 +960,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        backend: [llvm, cpu, webgpu]
+        backend: [llvm, cpu, webgpu, x86]
 
     name: Windows (${{ matrix.backend }})
     runs-on: windows-latest
@@ -976,7 +976,7 @@ jobs:
           pydeps: ${{ matrix.backend == 'webgpu' && 'dawn-python' || '' }}
       - name: Set env
         shell: bash
-        run:  printf "${{ matrix.backend == 'llvm' && 'DEV=CPU:LLVM' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_COUNT=2' || matrix.backend == 'webgpu' && 'DEV=WEBGPU'}}" >> $GITHUB_ENV
+        run:  printf "${{ matrix.backend == 'llvm' && 'DEV=CPU:LLVM' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_COUNT=2' || matrix.backend == 'webgpu' && 'DEV=WEBGPU' || matrix.backend == 'x86' && 'DEV=CPU:X86' }}" >> $GITHUB_ENV
       - name: Run unit tests
         if: matrix.backend=='llvm'
         # test_newton_schulz hits RecursionError
@@ -988,7 +988,7 @@ jobs:
       - name: Run pytest (${{ matrix.backend }})
         shell: bash
         run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU', 'X86':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
           python -m pytest -n=auto test/test_tiny.py test/backend/test_ops.py --durations=20
 
 # ****** Compile-only Tests ******

examples/mlperf/model_train.py

@@ -246,7 +246,7 @@ def train_resnet():
 
       if i == BENCHMARK:
         assert not math.isnan(loss)
-        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
+        median_step_time = sorted(step_times)[BENCHMARK // 2]  # in seconds
         estimated_total_minutes = int(median_step_time * steps_in_train_epoch * epochs / 60)
         print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
         print(f"epoch global_ops: {steps_in_train_epoch * GlobalCounters.global_ops:_}, "
@@ -593,7 +593,7 @@ def train_retinanet():
 
       if i == BENCHMARK:
         assert not math.isnan(loss)
-        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
+        median_step_time = sorted(step_times)[BENCHMARK // 2]  # in seconds
         estimated_total_minutes = int(median_step_time * steps_in_train_epoch * EPOCHS / 60)
         print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
         print(f"epoch global_ops: {steps_in_train_epoch * GlobalCounters.global_ops:_}, "
@@ -868,7 +868,7 @@ def train_unet3d():
         i += 1
 
         if i == BENCHMARK:
-          median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
+          median_step_time = sorted(step_times)[BENCHMARK // 2]  # in seconds
           estimated_total_minutes = int(median_step_time * SAMPLES_PER_EPOCH * NUM_EPOCHS / 60)
           print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
           if (TRAIN_BEAM or EVAL_BEAM) and epoch == start_epoch: break
@@ -1167,7 +1167,7 @@ def train_bert():
       i += 1
 
       if i == BENCHMARK:
-        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
+        median_step_time = sorted(step_times)[BENCHMARK // 2]  # in seconds
         estimated_total_minutes = int(median_step_time * train_steps / 60)
         print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
         print(f"epoch global_ops: {train_steps * GlobalCounters.global_ops:_}, "
@@ -1544,7 +1544,7 @@ def train_llama3():
 
       mem_gb = GlobalCounters.mem_used / 1e9
       gflops = GlobalCounters.global_ops / 1e9 / dev_time
-      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 2.3e15)) * 100
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * (4.6e15 if FP8 else 2.3e15))) * 100
       tqdm.write(
           f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
           f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
@@ -1577,8 +1577,9 @@ def train_llama3():
         safe_save(get_state_dict(scheduler), fn)
 
       if i == BENCHMARK:
-        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]
-        estimated_total_minutes = int(median_step_time * (SAMPLES // GBS) / 60)
+        median_step_time = sorted(step_times)[BENCHMARK // 2]
+        estimated_steps = 200_000 // GBS if getenv("LLAMA3_SIZE", "8B") == "8B" else MAX_STEPS
+        estimated_total_minutes = int(median_step_time * estimated_steps / 60)
         print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
         print(f"epoch global_ops: {GlobalCounters.global_ops:_}, "
               f"epoch global_mem: {GlobalCounters.global_mem:_}")

examples/mlperf/models/flat_llama.py

@@ -35,12 +35,15 @@ def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
   return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal()
 
 def matmul(x:Tensor, w:Tensor, fp8=FP8, amax_x:Tensor|None=None, amax_w:Tensor|None=None) -> Tensor:
-  if not fp8: return x @ w.T
-  from tinygrad.helpers import ASM_GEMM
+  if not fp8:
+    if getenv("ASM_GEMM"):
+      from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
+      if can_use_asm_gemm(x, w.T): return asm_gemm(x, w.T)
+    return x @ w.T
   x_fp8, x_scale = quantize_fp8(x, amax_state=amax_x)
   w_fp8, w_scale = quantize_fp8(w, amax_state=amax_w)
   combined_scale = x_scale * w_scale
-  if ASM_GEMM:
+  if getenv("ASM_GEMM"):
     from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
     if can_use_asm_gemm(x_fp8, w_fp8.T): return asm_gemm(x_fp8, w_fp8.T, combined_scale=combined_scale)
   return x_fp8.dot(w_fp8.T, dtype=dtypes.float) * combined_scale
@@ -122,7 +125,11 @@ class FlatTransformer:
     xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
     if FP8: 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)
-    attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
+    if getenv("HK_FLASH_ATTENTION"):
+      from extra.thunder.amd.fa import flash_attention
+      attn = flash_attention(xq, xk, xv, is_causal=True)
+    else:
+      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
     attn = attn.reshape(bsz, seqlen, -1)
     return matmul(attn, wo, amax_x=amax_xo, amax_w=amax_wo)
 
@@ -193,7 +200,7 @@ class FlatTransformer:
                          self.attention_norm[i], self.wo[i],
                          self.ffn_norm[i], self.w1[i], self.w2[i], self.w3[i],
                          **attn_kwargs, **amax_attn, **amax_layer)
-    logits = (self.norm(h).contiguous().contiguous_backward() @ self.output[0].T).contiguous_backward()
+    logits = matmul(self.norm(h).contiguous().contiguous_backward(), self.output[0], fp8=False).contiguous_backward()
     return logits
 
 def _get_pads(uop:UOp) -> list[UOp]:

examples/mlperf/models/llama.py

@@ -1,80 +0,0 @@
-from tinygrad import Tensor, nn
-from tinygrad.helpers import getenv
-from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
-
-class Attention:
-  def __init__(self, dim:int, n_heads:int, n_kv_heads:int|None=None, linear=nn.Linear):
-    self.n_heads = n_heads
-    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
-    self.head_dim = dim // n_heads
-    self.n_rep = self.n_heads // self.n_kv_heads
-
-    if getenv("WQKV"):
-      self.wqkv = linear(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2, bias=False)
-    else:
-      self.wq = linear(dim, self.n_heads * self.head_dim, bias=False)
-      self.wk = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-      self.wv = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-
-    self.wo = linear(self.n_heads * self.head_dim, dim, bias=False)
-
-  def __call__(self, x:Tensor, freqs_cis:Tensor) -> Tensor:
-    if getenv("WQKV"):
-      xqkv = self.wqkv(x)
-      xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
-      xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-      xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-      xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-    else:
-      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
-
-    xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
-    xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
-    xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
-
-    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-    bsz, seqlen, _, _ = xq.shape
-
-    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)
-    return self.wo(attn)
-
-class FeedForward:
-  def __init__(self, dim:int, hidden_dim:int, linear=nn.Linear):
-    self.w1 = linear(dim, hidden_dim, bias=False)
-    self.w2 = linear(hidden_dim, dim, bias=False)
-    self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit
-
-  def __call__(self, x:Tensor) -> Tensor:
-    w1 = self.w1(x).silu()
-    w3 = self.w3(x)
-    return self.w2(w1 * w3)
-
-class TransformerBlock:
-  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int|None, norm_eps:float, linear=nn.Linear):
-    self.attention = Attention(dim, n_heads, n_kv_heads, linear)
-    self.feed_forward = FeedForward(dim, hidden_dim, linear)
-    self.attention_norm = nn.RMSNorm(dim, norm_eps)
-    self.ffn_norm = nn.RMSNorm(dim, norm_eps)
-
-  def __call__(self, x:Tensor, freqs_cis:Tensor):
-    h = x + self.attention(self.attention_norm(x), freqs_cis)
-    return h + self.feed_forward(self.ffn_norm(h))
-
-class Transformer:
-  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
-               rope_theta:int=10000, max_context:int=1024, linear=nn.Linear, embedding=nn.Embedding):
-    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, linear) for _ in range(n_layers)]
-    self.norm = nn.RMSNorm(dim, norm_eps)
-    self.tok_embeddings = embedding(vocab_size, dim)
-    self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
-
-  def __call__(self, tokens:Tensor):
-    h = self.tok_embeddings(tokens)
-    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
-    for layer in self.layers: h = layer(h, freqs_cis)
-    logits = self.output(self.norm(h))
-    return logits

extra/gemm/cdna_asm_gemm.py

@@ -2716,8 +2716,11 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
     assert all_same([gradient.device, a.device, b.device, out.device])
     a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
     g_t = g_t[:a.shape[0]]
-    grad_a = (g_t @ b_t.T).uop
-    grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
+    if can_use_asm_gemm(g_t, b_t.T): grad_a = asm_gemm(g_t, b_t.T).uop
+    else: grad_a = (g_t @ b_t.T).uop
+    a_t_flat, g_t_flat = a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1), g_t.reshape(-1, g_t.shape[-1])
+    if can_use_asm_gemm(a_t_flat, g_t_flat): grad_b = asm_gemm(a_t_flat, g_t_flat).uop
+    else: grad_b = (a_t_flat @ g_t_flat).uop
     return (None, grad_a, grad_b)
 
 # ** main gemm function

extra/gemm/rdna4_asm_matmul.py

@@ -0,0 +1,245 @@
+# RDNA4 128x128 GEMM using WMMA — optimized DS scheduling
+import numpy as np
+from tinygrad import Tensor, Device, Context, GlobalCounters
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.helpers import getenv, colored
+from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.engine.realize import Estimates
+from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL, src, ttmp
+from tinygrad.runtime.autogen.amd.rdna4.ins import *
+
+BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 16
+TILES_M, TILES_N = 4, 4
+THREADS, ELEM = 128, 2
+LDS_A_ROW = BLOCK_K*ELEM   # 32
+LDS_B_ROW = BLOCK_N*ELEM  # 256
+LDS_A_SIZE = BLOCK_M * LDS_A_ROW  # 4096
+LDS_B_SIZE = BLOCK_K * LDS_B_ROW  # 4096
+LDS_SIZE = LDS_A_SIZE + LDS_B_SIZE  # 8192
+LDS_B_OFF = LDS_A_SIZE
+ACC, DA, DB, FA, FB, ET = 60, 188, 196, 204, 44, 10
+
+def build_kernel(N, arch='gfx1200'):
+  assert N % BLOCK_M == 0 and N >= 256
+  NO_ALU, NO_DS, NO_GLOBAL = getenv("NO_ALU", 0), getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
+  I, L, B = [], {}, []
+  def e(i): I.append(i); return i
+  def label(n): L[n] = sum(i.size() for i in I)
+  def br(i, t): B.append((len(I)-1, t))
+
+  e(s_load_b128(sdata=s[4:7], sbase=s[0:1], ioffset=0, soffset=NULL))
+  e(s_load_b64(sdata=s[8:9], sbase=s[0:1], ioffset=0x10, soffset=NULL))
+  e(s_wait_kmcnt(simm16=0))
+  e(s_mov_b32(s[10], ttmp[9])); e(s_and_b32(s[11], ttmp[7], 0xFFFF))
+  e(s_lshl_b32(s[10], s[10], 7)); e(s_lshl_b32(s[11], s[11], 7))
+  e(s_mov_b32(s[12], N)); e(s_lshl_b32(s[13], s[12], 1))
+  e(s_mul_i32(s[14], s[12], BLOCK_K*ELEM))
+  e(s_add_co_i32(s[17], s[12], -2*BLOCK_K))  # loop bound
+
+  e(v_and_b32_e32(v[1], 31, v[0])); e(v_lshrrev_b32_e32(v[2], 5, v[0]))
+  e(v_and_b32_e32(v[3], 1, v[2])); e(v_lshrrev_b32_e32(v[2], 1, v[2]))
+
+  e(v_lshlrev_b32_e32(v[4], 5, v[0]))
+  # B store: transposed layout for stride-32 reads. addr = LDS_B_OFF + (tid%8)*512 + (tid/8)*32
+  e(v_and_b32_e32(v[48], 7, v[0])); e(v_lshlrev_b32_e32(v[5], 9, v[48]))   # (tid%8)*512
+  e(v_lshrrev_b32_e32(v[48], 3, v[0])); e(v_lshlrev_b32_e32(v[48], 5, v[48]))  # (tid/8)*32
+  e(v_add_nc_u32_e32(v[5], v[5], v[48])); e(v_add_nc_u32_e32(v[5], LDS_B_OFF, v[5]))
+
+  e(v_add_nc_u32_e32(v[48], s[11], v[0]))
+  e(v_mul_lo_u32(v[6], v[48], N*ELEM)); e(v_mov_b32_e32(v[7], 0))
+  e(v_lshrrev_b32_e32(v[48], 3, v[0])); e(v_mul_lo_u32(v[8], v[48], N*ELEM))
+  e(v_and_b32_e32(v[48], 7, v[0])); e(v_lshlrev_b32_e32(v[48], 5, v[48]))
+  e(v_add_nc_u32_e32(v[8], v[8], v[48]))
+  e(s_mul_i32(s[15], s[10], ELEM)); e(v_add_nc_u32_e32(v[8], s[15], v[8]))
+  e(v_mov_b32_e32(v[9], 0))
+
+  # LDS read addrs with padded strides (eliminates bank conflicts)
+  # A: (lane%16)*LDS_A_ROW + (lane/16)*16 + wave_m*64*LDS_A_ROW
+  # B: (lane%16)*LDS_B_ROW + (lane/16)*16 + wave_n*64*ELEM + LDS_B_OFF
+  LLA, LLB = 40, 43
+  e(v_and_b32_e32(v[50], 15, v[1])); e(v_lshrrev_b32_e32(v[51], 4, v[1]))
+  e(v_lshlrev_b32_e32(v[LLA], 5, v[50]))        # (lane%16) * 32
+  e(v_lshlrev_b32_e32(v[51], 4, v[51]))        # (lane/16) * 16
+  e(v_add_nc_u32_e32(v[LLA], v[LLA], v[51]))
+  e(v_lshlrev_b32_e32(v[52], 11, v[2]))         # wave_m * 2048
+  e(v_add_nc_u32_e32(v[LLA], v[LLA], v[52]))
+  # B read: transposed layout. addr = LDS_B_OFF + (lane%16)*32 + (lane/16)*16 + wave_n*2*512
+  # wave_n selects column panels: wave_n*2 panels (each panel=16 cols, wave_n covers 64 cols = 4 panels)
+  # But wave_n*2*512 = wave_n*1024. Hmm, wave_n covers cols [wave_n*64 : (wave_n+1)*64].
+  # Each panel = 16 cols = 512 bytes. wave_n*64/16 = wave_n*4 panels. Offset = wave_n*4*512 = wave_n*2048.
+  e(v_lshlrev_b32_e32(v[LLB], 5, v[50]))         # (lane%16) * 32 (stride 32!)
+  e(v_add_nc_u32_e32(v[LLB], v[LLB], v[51]))     # + (lane/16)*16
+  e(v_lshlrev_b32_e32(v[52], 11, v[3]))           # wave_n * 2048
+  e(v_add_nc_u32_e32(v[LLB], v[LLB], v[52]))
+  e(v_add_nc_u32_e32(v[LLB], LDS_B_OFF, v[LLB]))
+
+  for i in range(0, 128, 2):
+    e(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[ACC+i], vdsty=v[ACC+i+1], srcx0=0, srcy0=0))
+  e(s_mov_b32(s[16], 0))
+
+  if not NO_GLOBAL:
+    for i in range(2): e(global_load_b128(vdst=v[DA+i*4:DA+i*4+3], vaddr=v[6:7], saddr=s[4:5], ioffset=i*16))
+    for i in range(2): e(global_load_b128(vdst=v[DB+i*4:DB+i*4+3], vaddr=v[8:9], saddr=s[6:7], ioffset=i*16))
+    e(s_wait_loadcnt(simm16=0))
+  if not NO_DS:
+    for i in range(2): e(ds_store_b128(addr=v[4], data0=v[DA+i*4:DA+i*4+3], offset0=(i*16)&0xFF, offset1=(i*16)>>8))
+    for i in range(2): e(ds_store_b128(addr=v[5], data0=v[DB+i*4:DB+i*4+3], offset0=(i*16)&0xFF, offset1=(i*16)>>8))
+  if not NO_GLOBAL:
+    e(v_add_nc_u32_e32(v[6], BLOCK_K*ELEM, v[6]))
+    e(v_add_nc_u32_e32(v[8], s[14], v[8]))
+
+  # =============================================================================
+  def emit_iter_body(load_set='AB'):
+    if not NO_DS:
+      e(s_wait_dscnt(simm16=0))
+      e(s_barrier_signal(ssrc0=src[193])); e(s_barrier_wait(simm16=0xFFFF))
+    if not NO_GLOBAL:
+      if 'A' in load_set:
+        for i in range(2): e(global_load_b128(vdst=v[DA+i*4:DA+i*4+3], vaddr=v[6:7], saddr=s[4:5], ioffset=i*16))
+        e(v_add_nc_u32_e32(v[6], BLOCK_K*ELEM, v[6]))
+      if 'B' in load_set:
+        for i in range(2): e(global_load_b128(vdst=v[DB+i*4:DB+i*4+3], vaddr=v[8:9], saddr=s[6:7], ioffset=i*16))
+        e(v_add_nc_u32_e32(v[8], s[14], v[8]))
+    if not NO_DS:
+      # Issue 6 loads: A[0:3] + B[0] + B[1]. B[2:3] interleaved with WMMAs.
+      for tm in range(TILES_M):
+        aoff = tm * 16 * LDS_A_ROW
+        e(ds_load_b128(vdst=v[FA+tm*4:FA+tm*4+3], addr=v[LLA], offset0=aoff&0xFF, offset1=aoff>>8))
+      e(ds_load_b128(vdst=v[FB:FB+3], addr=v[LLB], offset0=0, offset1=0))
+      e(ds_load_b128(vdst=v[FB+4:FB+7], addr=v[LLB], offset0=0, offset1=2))
+      e(s_wait_dscnt(simm16=0))  # wait for 6 loads (no stall!)
+    if not NO_ALU:
+      # B[0] WMMAs — issue B[2] during compute
+      if not NO_DS: e(ds_load_b128(vdst=v[FB+8:FB+11], addr=v[LLB], offset0=0, offset1=4))
+      for tm in range(TILES_M):
+        ac = ACC + (tm*TILES_N+0)*8
+        e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB:FB+3], src2=v[ac:ac+7]))
+      # B[1] WMMAs — issue B[3] during compute
+      if not NO_DS:
+        e(ds_load_b128(vdst=v[FB+12:FB+15], addr=v[LLB], offset0=0, offset1=6))
+      for tm in range(TILES_M):
+        ac = ACC + (tm*TILES_N+1)*8
+        e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+4:FB+7], src2=v[ac:ac+7]))
+      # B[2] WMMAs — B[2] loaded during B[0] WMMAs (~100 cycles ago)
+      if not NO_DS: e(s_wait_dscnt(simm16=1))  # B[2] done, B[3] may still be loading
+      for tm in range(TILES_M):
+        ac = ACC + (tm*TILES_N+2)*8
+        e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+8:FB+11], src2=v[ac:ac+7]))
+      # B[3] WMMAs
+      if not NO_DS: e(s_wait_dscnt(simm16=0))
+      for tm in range(TILES_M):
+        ac = ACC + (tm*TILES_N+3)*8
+        e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+12:FB+15], src2=v[ac:ac+7]))
+    if not NO_GLOBAL and not NO_DS: e(s_wait_loadcnt(simm16=0))
+    if not NO_DS:
+      for i in range(2): e(ds_store_b128(addr=v[4], data0=v[DA+i*4:DA+i*4+3], offset0=(i*16)&0xFF, offset1=(i*16)>>8))
+      for i in range(2): e(ds_store_b128(addr=v[5], data0=v[DB+i*4:DB+i*4+3], offset0=(i*16)&0xFF, offset1=(i*16)>>8))
+    e(s_add_co_i32(s[16], s[16], BLOCK_K))
+
+  label('LOOP')
+  emit_iter_body(load_set='A')
+  emit_iter_body(load_set='B')
+  e(s_cmp_lt_i32(s[16], s[17])); e(s_cbranch_scc1(simm16=0)); br(I[-1], 'LOOP')
+
+  emit_iter_body(load_set='AB')  # tail with prefetch
+
+  # Final iteration: no prefetch, no ds_store needed
+  if not NO_DS:
+    e(s_wait_dscnt(simm16=0))
+    e(s_barrier_signal(ssrc0=src[193])); e(s_barrier_wait(simm16=0xFFFF))
+  if not NO_DS:
+    for tm in range(TILES_M):
+      aoff = tm * 16 * LDS_A_ROW
+      e(ds_load_b128(vdst=v[FA+tm*4:FA+tm*4+3], addr=v[LLA], offset0=aoff&0xFF, offset1=aoff>>8))
+    e(ds_load_b128(vdst=v[FB:FB+3], addr=v[LLB], offset0=0, offset1=0))
+    e(ds_load_b128(vdst=v[FB+4:FB+7], addr=v[LLB], offset0=0, offset1=2))
+    e(s_wait_dscnt(simm16=0))
+  if not NO_ALU:
+    if not NO_DS: e(ds_load_b128(vdst=v[FB+8:FB+11], addr=v[LLB], offset0=0, offset1=4))
+    for tm in range(TILES_M):
+      ac = ACC + (tm*TILES_N+0)*8
+      e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB:FB+3], src2=v[ac:ac+7]))
+    if not NO_DS: e(ds_load_b128(vdst=v[FB+12:FB+15], addr=v[LLB], offset0=0, offset1=6))
+    for tm in range(TILES_M):
+      ac = ACC + (tm*TILES_N+1)*8
+      e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+4:FB+7], src2=v[ac:ac+7]))
+    if not NO_DS: e(s_wait_dscnt(simm16=1))
+    for tm in range(TILES_M):
+      ac = ACC + (tm*TILES_N+2)*8
+      e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+8:FB+11], src2=v[ac:ac+7]))
+    if not NO_DS: e(s_wait_dscnt(simm16=0))
+    for tm in range(TILES_M):
+      ac = ACC + (tm*TILES_N+3)*8
+      e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+12:FB+15], src2=v[ac:ac+7]))
+
+  label('EPILOGUE')
+  e(v_and_b32_e32(v[ET], 15, v[1]))
+  e(v_lshrrev_b32_e32(v[ET+1], 4, v[1])); e(v_lshlrev_b32_e32(v[ET+1], 3, v[ET+1]))
+  e(v_lshlrev_b32_e32(v[ET+2], 6, v[2])); e(v_add_nc_u32_e32(v[ET+2], s[11], v[ET+2]))
+  e(v_lshlrev_b32_e32(v[ET+3], 6, v[3])); e(v_add_nc_u32_e32(v[ET+3], s[10], v[ET+3]))
+  e(v_add_nc_u32_e32(v[ET+3], v[ET+3], v[ET])); e(v_mov_b32_e32(v[ET+5], 0))
+
+  for tm in range(TILES_M):
+    for tn in range(TILES_N):
+      ac = ACC + (tm*TILES_N+tn)*8; r_off, c_off = tm*16, tn*16
+      e(v_add_nc_u32_e32(v[ET+6], r_off, v[ET+2])); e(v_add_nc_u32_e32(v[ET+6], v[ET+1], v[ET+6]))
+      e(v_mul_lo_u32(v[ET+4], v[ET+6], s[12])); e(v_add_nc_u32_e32(v[ET+4], v[ET+4], v[ET+3]))
+      if c_off: e(v_add_nc_u32_e32(v[ET+4], c_off, v[ET+4]))
+      e(v_lshlrev_b32_e32(v[ET+4], 1, v[ET+4]))
+      for elem in range(8):
+        e(v_cvt_f16_f32_e32(v[ET+7], v[ac+elem]))
+        e(global_store_b16(vaddr=v[ET+4:ET+5], vsrc=v[ET+7], saddr=s[8:9]))
+        if elem < 7: e(v_add_nc_u32_e32(v[ET+4], s[13], v[ET+4]))
+
+  e(s_wait_storecnt(simm16=0)); e(s_sendmsg(simm16=3)); e(s_endpgm())
+
+  for idx, target in B:
+    off = (L[target] - sum(i.size() for i in I[:idx+1])) // 4
+    assert -32768 <= off <= 32767; I[idx].simm16 = off
+  return I
+
+N = getenv("N", 4096)
+
+def test_matmul():
+  dev = Device[Device.DEFAULT]
+  arch = getattr(dev.renderer, 'arch', 'gfx1200')
+  print(f"Device arch: {arch}")
+  insts = build_kernel(N, arch)
+
+  rng = np.random.default_rng(42)
+  a = Tensor(rng.random((N, N), dtype=np.float32).astype(np.float16))
+  b = Tensor(rng.random((N, N), dtype=np.float32).astype(np.float16))
+  c = Tensor.empty(N, N, dtype=dtypes.half)
+  Tensor.realize(a, b, c)
+
+  grid, local = (N//BLOCK_N, N//BLOCK_M, 1), (THREADS, 1, 1)
+  print(f"Grid: {grid}, Local: {local}")
+
+  dname = Device.DEFAULT
+  def asm_kernel(A, B, C):
+    gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
+    lidxs = [UOp.special(THREADS, "lidx0")]
+    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=max(LDS_SIZE, 65536//getenv("LIMIT_OCC",2)), addrspace=AddrSpace.LOCAL), (), 'lds')
+    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs,
+                    arg=KernelInfo(name=colored("kernel","cyan"), estimates=Estimates(ops=N*N*N*2, mem=N*N*2*3)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+
+  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
+  ei = c.schedule()[0].lower()
+
+  ets = []
+  with Context(DEBUG=2):
+    for _ in range(getenv("CNT", 5)): ets.append(ei.run(wait=True))
+  print(f"REAL TFLOPS {N*N*N*2 / min(ets) * 1e-12:.2f}")
+
+  if getenv("VERIFY", 1):
+    GlobalCounters.reset()
+    c_np = c.float().numpy()
+    a_np, b_np = a.float().numpy(), b.float().numpy()
+    ref = a_np @ b_np
+    err = np.sqrt(np.mean((c_np - ref)**2)) / np.sqrt(np.mean(ref**2))
+    print(f"relative RMSE {err:.6f}")
+    if err != err or err > 0.05: raise RuntimeError(f"matmul is wrong! RMSE={err}")
+
+if __name__ == "__main__":
+  test_matmul()

extra/mlx_driver/loopback.py

@@ -6,7 +6,8 @@ from tinygrad.device import Device, BufferSpec
 from tinygrad.runtime.support.system import PCIDevice
 from tinygrad.runtime.support.memory import AddrSpace
 from tinygrad.runtime.ops_amd import AMDComputeQueue
-from extra.mlx_driver.mlxdev import MLXDev, MLXQP, to_be
+from tinygrad.helpers import to_be32, to_be64
+from extra.mlx_driver.mlxdev import MLXDev, MLXQP
 
 BUF_SIZE = 0x1000
 MLX_PCI = getenv("MLX_PCI", "0000:41:00.0")
@@ -49,7 +50,7 @@ rq_wqe = qp.qp_buf.view((qp.rq_head & rq_mask) * 16, 16)
 rq_wqe[:] = struct.pack('>IIQ', len(test_msg), dev.mkey, dst_paddr)
 qp.rq_head += 1
 # ring recv doorbell from CPU (DBR offset 0 = recv counter)
-dev.dbr[qp.qp_dbr // 4] = to_be('I', qp.rq_head)
+dev.dbr[qp.qp_dbr // 4] = to_be32(qp.rq_head)
 
 # build send WQE in SQ from CPU (opcode 0x0a = SEND, ds_count=2)
 sq_head = qp.sq_head
@@ -60,7 +61,7 @@ wqe[0:8] = struct.pack('>II', (sq_head << 8) | 0x0a, (qp.qp_info['qpn'] << 8) |
 wqe[11] = 0x08  # CE: signal completion
 wqe[16:32] = struct.pack('>IIQ', len(test_msg), dev.mkey, src_paddr)
 qp.sq_head += 1
-doorbell_val = to_be('Q', int.from_bytes(bytes(wqe[0:8]), 'big'))
+doorbell_val = to_be64(int.from_bytes(bytes(wqe[0:8]), 'big'))
 
 # map MLX5 UAR and DBR into GPU VA
 uar_paddr = dev.pci_dev.bar_info(0)[0] + dev.uar * 0x1000
@@ -72,7 +73,7 @@ print(f"UAR gpu_va=0x{uar_gpu_va:x} DBR gpu_va=0x{dbr_gpu_va:x}")
 q = AMDComputeQueue(gpu)
 q.wait(gpu.timeline_signal, gpu.timeline_value - 1)
 # write DBR (32-bit sq_head) - send doorbell at qp_dbr + 4
-q.release_mem(dbr_gpu_va + qp.qp_dbr + 4, to_be('I', qp.sq_head), q.pm4.data_sel__mec_release_mem__send_32_bit_low,
+q.release_mem(dbr_gpu_va + qp.qp_dbr + 4, to_be32(qp.sq_head), q.pm4.data_sel__mec_release_mem__send_32_bit_low,
               q.pm4.int_sel__mec_release_mem__none)
 # write UAR doorbell (64-bit)
 q.release_mem(uar_gpu_va + 0x800, doorbell_val, q.pm4.data_sel__mec_release_mem__send_64_bit_data,

extra/sqtt/examples/generate_examples.py

@@ -10,6 +10,7 @@ EXAMPLES = {
   "plus":"test/test_tiny.py TestTiny.test_plus",
   "gemm":"-c \"from tinygrad import Tensor; (Tensor.empty(N:=32, N)@Tensor.empty(N, N)).realize()\"",
   "sync":"test/amd/test_custom_kernel.py TestCustomKernel.test_lds_sync",
+  "handwritten":"test/amd/test_custom_kernel.py TestCustomKernel.test_handwritten",
 }
 
 if __name__ == "__main__":
@@ -21,6 +22,6 @@ if __name__ == "__main__":
     for i in range(2):
       # AM_RESET=1 gets a clear trace, does not work on mi300 machines
       subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
-                     env={**os.environ, "AMD":"1", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
+                     env={**os.environ, "DEV":"AMD", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
       PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
       print(f"saved SQTT trace to {dest}")

extra/viz/README.md

@@ -24,7 +24,7 @@ List all codegen steps for a kernel: `--rewrites -s E_3`
 Get source code: `--rewrites -s E_3 -i "View Source"`
 Inspect a graph rewrite: `--rewrites -s E_3 -i "initial symbolic"`
 
-# SQTT tracing
+## SQTT tracing
 
 Supported on AMD for RDNA3 and RDNA4 (best) and CDNA (developing).
 
@@ -38,8 +38,12 @@ You can select a specific trace with --source, Example workflow:
 VIZ=-2 python extra/gemm/amd_asm_matmul.py
 
 # View barriers
-extra/viz/cli.py --profile -s "SQTT kernel PKTS SE:0" | rg BARRIER | head -10
+extra/viz/cli.py --profile -s "kernel SQTT SE:0 PKTS" | rg BARRIER | head -10
+
+# Get bank conflicts from performance counters
+
+python extra/viz/cli.py -p -s "kernel PMC" -i "SQC_LDS_BANK_CONFLICT"
 
 # Find the EXEC corresponding to a DISPATCH at cycle 410
-extra/viz/cli.py --profile -s "SQTT kernel PKTS SE:0" | awk '/EXEC/ && $1 - $5 == 410'
+extra/viz/cli.py --profile -s "kernel SQTT SE:0 PKTS" | awk '/EXEC/ && $1 - $5 == 410'
 ```

extra/viz/cli.py

@@ -47,7 +47,9 @@ def decode_profile(data:bytes) -> dict:
 def get(data:dict, key:str):
   for k,v in data.items():
     if ansistrip(k) == key: return v
-  raise RuntimeError(f'item "{key}" not found in list')
+  import difflib
+  match = difflib.get_close_matches(key, [ansistrip(k) for k in data], n=1, cutoff=0.6)
+  raise RuntimeError(f'item "{key}" not found in list'+(f", did you mean {match[0]!r}?" if match else ''))
 
 def main(args) -> None:
   viz.trace = viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))
@@ -59,8 +61,8 @@ def main(args) -> None:
     events:list = viz.load_pickle(args.profile_path, default=[])
     if (profile_bytes:=viz.get_profile(events)) is None: raise RuntimeError(f"empty profile in {args.profile_path}")
     profile = decode_profile(profile_bytes)
-    profile["layout"].update([(f'{c["name"]} {s["name"]}', s["data"]) for c in viz.ctxs if c["name"].startswith("SQTT") for s in c["steps"]
-                              if "PKTS" in s["name"]])
+    profile["layout"].update([(f'{c["name"][5:]}{" SQTT" if s["name"].endswith("PKTS") else ""} {s["name"]}', s["data"]) for c in viz.ctxs
+                              if c["name"].startswith("SQTT") for s in c["steps"] if s["name"].endswith(("PMC", "PKTS"))])
     if args.src is None:
       for k in profile["layout"]:
         print(f"  {format_colored(k)}")
@@ -87,7 +89,7 @@ def main(args) -> None:
         op_str = hex_colored(op_name, color) if color and not args.no_color else op_name
         phase, delay = None, 0
         idx = next(pkt_idxs.setdefault(e.device, itertools.count()))
-        if e.device.startswith("WAVE") or e.device == "OTHER_SIMD":
+        if e.device.startswith("WAVE"):
           inst = f"0x{(pc:=int(info.replace('PC:', ''))):05x} {pc_map[pc]}" if info else f"{'':7} {op_name}"
           dispatch_to_inst[f"{e.device}-{idx}"] = (inst, int(e.st))
           phase = "DISPATCH"
@@ -99,6 +101,20 @@ def main(args) -> None:
         print(f"{int(e.st)-inst_st:<12} {unit:<20} {op_str}{' '*(22-ansilen(op_str))} {int(unwrap(e.en)-e.st):<4} {str(delay or ''):<4} {info}")
       return None
 
+    # ** PMC printer
+    if "PMC" in args.src:
+      table = viz.unpack_pmc(data[0])
+      cols = table["cols"]
+      rows:list = []
+      for r in table["rows"]:
+        if args.item is None: rows.append(r[:2])
+        elif args.item == r[0]:
+          rows = r[2]["rows"] if len(r) > 2 else [r[:2]]
+          cols = r[2]["cols"] if len(r) > 2 else cols
+      from tabulate import tabulate
+      print(tabulate(rows, headers=cols, tablefmt="github"))
+      return None
+
     # ** Profiler printer
     agg:dict[str, tuple[float, int]] = {}
     total = 0

test/amd/test_custom_kernel.py

@@ -104,20 +104,21 @@ def custom_handwritten(A:UOp, arch:str) -> UOp:
   lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=512, addrspace=AddrSpace.LOCAL), (), 'lds')  # 128 * 4 bytes
   k = Kernel(arch)
   k.emit(r4.s_nop(0))
-  k.emit(r4.v_mov_b32_e32(v[1], 10))
+  k.emit(r4.v_mov_b32_e32(v[1], 4))
   def emit_alt():
-    for i in range(4):
+    for i in range(2):
       k.emit(r4.v_mov_b32_e32(v[20+i], 4.0))
       k.emit(r4.v_rcp_f32_e32(v[22+i], v[20+i]))
       k.emit(r4.s_mov_b32(s[20+i], i))
+      k.emit(r4.s_mul_i32(s[14+i], s[12+i], 32))
   def emit_wmma():
-    for _ in range(4):
+    for _ in range(2):
       k.emit(r4.v_wmma_f32_16x16x16_f16(v[0:7], v[8:11], v[8:11], 1))
   k.label("start")
   k.emit(s_mov_b32(s[1], 10))
   k.label("loop")
   # wmma should've overlapped here if it was a different unit?
-  for _ in range(4):
+  for _ in range(2):
     emit_wmma()
     emit_alt()
   for _ in range(8): k.emit(s_nop(1))

test/amd/test_sqttmap.py

@@ -100,9 +100,7 @@ class TestSQTTMapBase(unittest.TestCase):
             elif "WAVE" in e.device:
               # sopk/immediates don't get ALU/MEM EXEC
               if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE", "BARRIER", "BARRIER_SIGNAL",
-                                             "WAVEEND", "WAVERDY"}: insts += 1
-            # OTHER_ is its own stream, it's the INST from other SIMDs that share the same EXEC.
-            elif e.device.startswith("OTHER"): continue
+                                             "WAVEEND", "WAVERDY"} and not e.name.display_name.startswith("OTHER_"): insts += 1
             else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
           self.assertEqual(execs, insts)
 
@@ -131,7 +129,18 @@ class TestSQTTMapBase(unittest.TestCase):
 
 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"
 
-class TestSQTTMapRDNA4(TestSQTTMapBase): target = "gfx1200"
+class TestSQTTMapRDNA4(TestSQTTMapBase):
+  target = "gfx1200"
+
+  @unittest.expectedFailure
+  def test_rdna4_wmma(self):
+    events, kernels, target = self.examples["profile_handwritten_run_0"]
+    row_ends = {}
+    for e in sqtt_timeline(events[0].blob, list(kernels.values())[0].lib, target):
+      if type(e).__name__ != "ProfileRangeEvent" or e.device != "ALUEXEC:0 WMMA": continue
+      if (et:=row_ends.get(e.device)) is not None and e.st < et:
+        raise RuntimeError(f"WMMA exec overlaps in {e.device}: {e.st} {et}.")
+      row_ends[e.device] = e.en
 
 class TestSQTTMapCDNA(TestSQTTMapBase):
   target = "gfx950"

test/backend/test_arange.py

@@ -24,6 +24,10 @@ class TestArange(unittest.TestCase):
     self.assertEqual(self._get_flops(Tensor.arange(256), np.arange(256)), 0)
     self.assertEqual(self._get_flops(Tensor.arange(2560), np.arange(2560)), 0)
 
+  @unittest.skipIf(Device.DEFAULT == "CL", "TODO: fails on CI CL")
+  def test_arange_cumsum(self):
+    np.testing.assert_equal(Tensor.arange(513).cumsum(0).numpy(), np.arange(513).cumsum())
+
   def test_arange_cat(self):
     t = Tensor.arange(2, dtype=dtypes.int)+Tensor([3])
     self.assertEqual(t.cat(t).tolist(), [3, 4, 3, 4])

test/backend/test_asm_gemm.py

@@ -21,9 +21,8 @@ def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=N
 
   a, b = a_rand.clone().requires_grad_(), b_rand.clone().requires_grad_()
   if multi: a, b = a.shard(devs, axis=a_shard), b.shard(devs, axis=b_shard)
-  with Context(ASM_GEMM=1):
-    tst = asm_gemm(a, b)
-    tst.sum().backward()
+  tst = asm_gemm(a, b)
+  tst.sum().backward()
   Tensor.realize(tst, a.grad, b.grad)
 
   a_ref, b_ref = a_rand.clone().requires_grad_(), b_rand.clone().requires_grad_()
@@ -32,9 +31,8 @@ def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=N
     a_ref = a_ref.cast(dtypes.bfloat16)
     b_ref = b_ref.cast(dtypes.bfloat16)
   if multi: a_ref, b_ref = a_ref.shard(devs, axis=a_shard), b_ref.shard(devs, axis=b_shard)
-  with Context(ASM_GEMM=0):
-    ref = asm_gemm(a_ref, b_ref)
-    ref.sum().backward()
+  ref = a_ref @ b_ref
+  ref.sum().backward()
   Tensor.realize(ref, a_ref.grad, b_ref.grad)
 
   # no validation on the NULL device
@@ -136,14 +134,12 @@ class TestGemmLlama(unittest.TestCase):
     if not is_cdna4() or getenv("MOCKGPU"):
       self.skipTest("very slow on non mi350x")
 
-  @Context(ASM_GEMM=1)
-  def test_empty(self): (Tensor.empty(N:=getenv("N", 4096), N, dtype=self.dtype)@Tensor.empty(N, N, dtype=self.dtype)).realize()
+  def test_empty(self): asm_gemm(Tensor.empty(N:=getenv("N", 4096), N, dtype=self.dtype), Tensor.empty(N, N, dtype=self.dtype)).realize()
 
-  @Context(ASM_GEMM=1)
   def test_empty_bw(self):
     x = Tensor.empty(1, N:=getenv("N", 4096), N, dtype=self.dtype, requires_grad=True)
     y = Tensor.empty((N, N), dtype=self.dtype, requires_grad=True)
-    z = x @ y
+    z = asm_gemm(x, y)
     z.sum().backward()
     Tensor.realize(z, x.grad, y.grad)
     # FP8 forward output is bf16, gradients use fp8e5m2 (aka bf8)

test/backend/test_encodings.py

@@ -0,0 +1,150 @@
+import unittest
+from typing import cast
+from tinygrad import Device
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.dtype import dtypes
+from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm, def_reg
+
+def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=op, dtype=dt, src=src, tag=tag)
+
+@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only on x86")
+class TestEncodingsX86(unittest.TestCase):
+  # NOTE: x86 supports a single displacement as memory address and index without base memory address
+  # these have no use cases so they aren't supported
+  def encode(self, u:UOp): return cast(X86Renderer, Device[Device.DEFAULT].renderer).render([u])
+
+  # displacement of 0 isn't emitted
+  def test_base_address(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RDI)
+    # mov edi, dword ptr [rdi]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 3F"))
+
+  # rsp/r12 require a sib byte when used as base memory address
+  def test_rsp_base_address(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RSP), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RSP)
+    # mov esp, dword ptr [rsp]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 24 24"))
+
+  # rbp/r13 require a displacement when used as base memory address
+  def test_rbp_base_address(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RBP), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RBP)
+    # mov ebp, dword ptr [rbp + 0]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 6D 00"))
+
+  # test [base + index*scale]
+  def test_base_index_address(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, RDX), imm(dtypes.int8, 0)), RAX)
+    # mov eax, dword ptr [rax + rdx*4]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 04 90"))
+
+  # rsp as index means no index
+  def test_rsp_index_address(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, RSP), imm(dtypes.int8, 0)), RAX)
+    # mov eax, dword ptr [rax]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 00"))
+
+  # however r12 is a valid index
+  def test_r12_index_address(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, GPR[12]), imm(dtypes.int8, 0)), RAX)
+    # mov eax, dword ptr [rax + r12*4]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("42 8B 04 A0"))
+
+  # test [base + index*scale + 8bit disp]
+  def test_complex_address_8bit_disp(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10)), RDI)
+    # mov edi, dword ptr [rdi + rsi*4 + 0xa]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 7C B7 0A"))
+
+  # test [base + index*scale + 32bit disp]
+  def test_complex_address_32bit_disp(self):
+    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10000)), RDI)
+    # mov edi, dword ptr [rdi + rsi*4 + 0x2710]
+    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B BC B7 10 27 00 00"))
+
+  # 8bit variants of legacy instructions subtract 1 from opcode
+  def test_8bit_legacy_encoding(self):
+    cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(dtypes.int8, RDX),), RAX)
+    # movsx eax, dl
+    self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("0F BE C2"))
+
+  # accessing lower 8 bits of rsp, rbp, rsi, rdi requires rex prefix
+  def test_lower_8bits_reg(self):
+    cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(dtypes.int8, RDI),), RAX)
+    # movsx eax, dil
+    self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("40 0F BE C7"))
+
+  # test 16 bit variant of legacy instruction
+  def test_16bit_legacy_encoding(self):
+    cast = ins(X86Ops.MOVSX, dtypes.int16, (def_reg(dtypes.int8, RDX),), RAX)
+    # movsx ax, dl
+    self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("66 0F BE C2"))
+
+  # test 64 bit variant of legacy instruction
+  def test_64bit_legacy_encoding(self):
+    cast = ins(X86Ops.MOVSX, dtypes.int64, (def_reg(dtypes.int8, RDX),), RAX)
+    # movsx rax, dl
+    self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("48 0F BE C2"))
+
+  # test compact vex encoding
+  def test_compact_vex_encoding(self):
+    xmm0, xmm1 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1])
+    add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm1), XMM[0])
+    # vaddss xmm0, xmm0, xmm1
+    self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FA 58 C1"))
+
+  # test long vex encoding
+  def test_long_vex_encoding(self):
+    xmm0, xmm8 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[8])
+    add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm8), XMM[0])
+    # vaddss xmm0, xmm0, xmm8
+    self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C4 C1 7A 58 C0"))
+
+  # test ymm encoding
+  def test_ymm_encoding(self):
+    xmm0, xmm1 = def_reg(dtypes.float32.vec(8), XMM[0]), def_reg(dtypes.float32.vec(8), XMM[1])
+    add = ins(X86Ops.VADDPS, dtypes.float32.vec(8), (xmm0, xmm1), XMM[0])
+    # vaddps ymm0, ymm0, ymm1
+    self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FC 58 C1"))
+
+  # test encoding where register is in the immediate field
+  def test_reg_in_imm_field(self):
+    xmm0, xmm1, xmm2 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1]), def_reg(dtypes.float32, XMM[2])
+    blend = ins(X86Ops.VBLENDVPS, dtypes.float32, (xmm0, xmm1, xmm2), XMM[0])
+    # vblendvps xmm0, xmm0, xmm1, xmm2
+    self.assertEqual(bytes.fromhex(self.encode(blend)), bytes.fromhex("C4 E3 79 4A C1 20"))
+
+  # when writting to mem the uop takes the store form where dtype is void and there's no definition
+  def test_write_mem(self):
+    base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10)
+    xmm0 = def_reg(dtypes.float32, XMM[0])
+    extr = ins(X86Ops.VPEXTRD, dtypes.void, (base, index, disp, xmm0, imm(dtypes.uint8, 0)))
+    # vpextrd dword ptr [rdi + rsi*4 + 0xa], xmm0, 0
+    self.assertEqual(bytes.fromhex(self.encode(extr)), bytes.fromhex("C4 E3 79 16 44 B7 0A 00"))
+
+  # test two address instruction with fused load works
+  def test_two_address_load(self):
+    base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10)
+    cmove = ins(X86Ops.CMOVE, dtypes.int32,  (base, index, disp), RAX)
+    # cmove eax, dword ptr [rdi + rsi*4 + 0xa]
+    self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 44 B7 0A"))
+
+  # test instruction where displacement and imm have the same value
+  def test_disp_imm_same_value(self):
+    base, index, disp = def_reg(dtypes.int8.ptr(), RDI), def_reg(dtypes.int8, RSI), imm(dtypes.int8, 10)
+    mov = ins(X86Ops.MOVi, dtypes.void, (base, index, disp, disp))
+    # mov byte ptr [rdi + rsi + 0xa], 0xa
+    self.assertEqual(bytes.fromhex(self.encode(mov)), bytes.fromhex("40 C6 44 37 0A 0A"))
+
+    base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10)
+    imul = ins(X86Ops.IMULi, dtypes.int32, (base, index, disp) + (imm(dtypes.int32, 10),), RDI)
+    # imul edi, dword ptr [rdi + rsi*4 + 0xa], 0xa
+    self.assertEqual(bytes.fromhex(self.encode(imul)), bytes.fromhex("69 BC B7 0A 00 00 00 0A 00 00 00"))
+
+  # cmoves have the cmp as the last src even though it is not explicitly used, the cmp doesn't define a reg and is ignored in the encoding
+  def test_cmove_ignore_cmp(self):
+    cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(dtypes.int32, RAX), UOp(Ops.INS, arg=X86Ops.CMP)), RDX)
+    # cmove edx, eax
+    self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 D0"))
+
+if __name__ == "__main__":
+  unittest.main()

test/backend/test_isel.py

@@ -0,0 +1,152 @@
+import unittest
+from typing import cast
+from tinygrad import Device
+from tinygrad.uop import Ops
+from tinygrad.uop.ops import UOp, dtypes, graph_rewrite
+from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
+from tinygrad.renderer.isa import IselContext
+
+# these tests are to catch changes that don't cause incorrect codegen but cause worse codegen
+@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
+class TestIselX86(unittest.TestCase):
+  def isel_rewrite(self, x:UOp):
+    return graph_rewrite(x, cast(X86Renderer, Device[Device.DEFAULT].renderer).isel_matcher, IselContext(x), bottom_up=True)
+
+  def _check_op(self, dt_op, expr):
+    nargs = expr.__code__.co_argcount
+    for dt,op in dt_op:
+      with self.subTest(dtype=dt):
+        v = [UOp.variable(str(i), 0, 0, dt) for i in range(nargs)]
+        n = self.isel_rewrite(expr(*v))
+        self.assertIs(n.arg, op)
+
+  def test_cmove(self):
+    a = UOp.variable("a", 0, 0, dtypes.int32)
+    b = UOp.variable("b", 0, 0, dtypes.int32)
+    c = (a < b).where(a, b)
+    d = (a != b).where(a, b)
+    f = c + d
+    n = self.isel_rewrite(f)
+    self.assertTrue(n.src[0].arg is X86Ops.CMOVL and n.src[1].arg is X86Ops.CMOVNE)
+    # both comparisons become the same instruction
+    self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg is X86Ops.CMP)
+
+  def test_vmax(self):
+    dt_op = [(dtypes.float32, X86Ops.VMAXSS), (dtypes.float64, X86Ops.VMAXSD),
+             (dtypes.float32.vec(4), X86Ops.VMAXPS), (dtypes.float64.vec(4), X86Ops.VMAXPD)]
+    self._check_op(dt_op, lambda a,b: (a < b).where(b, a))
+
+  def test_vmin(self):
+    dt_op = [(dtypes.float32, X86Ops.VMINSS), (dtypes.float64, X86Ops.VMINSD),
+             (dtypes.float32.vec(4), X86Ops.VMINPS), (dtypes.float64.vec(4), X86Ops.VMINPD)]
+    self._check_op(dt_op, lambda a,b: (a < b).where(a, b))
+
+  def test_vfmadd(self):
+    dt_op = [(dtypes.float32, X86Ops.VFMADD213SS), (dtypes.float64, X86Ops.VFMADD213SD),
+             (dtypes.float32.vec(4), X86Ops.VFMADD213PS), (dtypes.float64.vec(4), X86Ops.VFMADD213PD)]
+    self._check_op(dt_op, lambda a,b,c: a * b + c)
+
+  # TODO: shouldn't match fmadd if var is used multiple times
+  @unittest.expectedFailure
+  def test_vfmadd_fail(self):
+    dt_op = [(dtypes.float32, X86Ops.VADDSS), (dtypes.float64, X86Ops.VADDSD),
+             (dtypes.float32.vec(4), X86Ops.VADDPS), (dtypes.float64.vec(4), X86Ops.VADDPD)]
+    self._check_op(dt_op, lambda a,b: a * b + b)
+
+  def test_vpbroadcast(self):
+    a = UOp.variable("a", 0, 0, dtypes.int32)
+    n = self.isel_rewrite(a.broadcast(4))
+    # need to move src from gpr to xmm before broadcasting
+    self.assertTrue(n.arg is X86Ops.VPBROADCASTD and n.src[0].arg is X86Ops.VMOVD)
+    # if we can fuse a load we can skip the move and access memory directly
+    load = UOp(Ops.PARAM, dtypes.int32.ptr(), arg=0).index(UOp.const(dtypes.int32, 0), ptr=True).load()
+    n = self.isel_rewrite(load.broadcast(4))
+    self.assertTrue(n.arg is X86Ops.VPBROADCASTD and len(n.src) == 3)
+
+  def test_vbroadcastss(self):
+    a = UOp.variable("a", 0, 0, dtypes.float32)
+    valid = [UOp.vectorize(a, a, a, a), UOp.vectorize(a, a, a, a, a, a, a, a)]
+    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VBROADCASTSS)
+
+  def test_vshufps(self):
+    a = UOp.variable("a", 0, 0, dtypes.float32.vec(8))
+    b = UOp.variable("b", 0, 0, dtypes.float32.vec(8))
+    c = UOp.variable("c", 0, 0, dtypes.float32)
+    d = UOp.variable("d", 0, 0, dtypes.float32)
+
+    valid = [UOp.vectorize(c, c, d, d),
+             UOp.vectorize(a.gep(0), a.gep(1), c, c),
+             UOp.vectorize(a.gep(0), a.gep(1), b.gep(2), b.gep(3)),
+             UOp.vectorize(a.gep(1), a.gep(2), a.gep(3), a.gep(0)),
+             UOp.vectorize(a.gep(3), a.gep(2), a.gep(1), a.gep(0), a.gep(7), a.gep(6), a.gep(5), a.gep(4)),
+             UOp.vectorize(a.gep(0), a.gep(0), b.gep(1), b.gep(1), a.gep(4), a.gep(4), b.gep(5), b.gep(5))]
+    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPS)
+
+    invalid = [UOp.vectorize(a.gep(0), a.gep(1), b.gep(4), b.gep(5)),
+               UOp.vectorize(a.gep(0), a.gep(5), b.gep(2), b.gep(3)),
+               UOp.vectorize(a.gep(0), a.gep(0), a.gep(0), a.gep(0), a.gep(4), a.gep(4), a.gep(4), a.gep(5)),
+               UOp.vectorize(a.gep(0), a.gep(0), b.gep(0), b.gep(0), a.gep(4), a.gep(4), b.gep(4), a.gep(4))]
+    for shuf in invalid: self.assertIsNot(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPS)
+
+  def test_vshufpd(self):
+    a = UOp.variable("a", 0, 0, dtypes.float64.vec(4))
+    b = UOp.variable("b", 0, 0, dtypes.float64.vec(4))
+    c = UOp.variable("c", 0, 0, dtypes.float64)
+    d = UOp.variable("d", 0, 0, dtypes.float64)
+
+    valid = [UOp.vectorize(c, d),
+             UOp.vectorize(a.gep(0), c),
+             UOp.vectorize(a.gep(1), b.gep(1)),
+             UOp.vectorize(a.gep(0), b.gep(1), a.gep(2), b.gep(3)),
+             UOp.vectorize(a.gep(1), a.gep(1), a.gep(3), a.gep(3))]
+    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPD)
+
+    invalid = [UOp.vectorize(c, c, c, c),
+               UOp.vectorize(a.gep(0), a.gep(1), b.gep(2), b.gep(3)),
+               UOp.vectorize(a.gep(2), b.gep(3), a.gep(2), b.gep(3)),
+               UOp.vectorize(a.gep(0), b.gep(1), a.gep(0), b.gep(1))]
+    for shuf in invalid: self.assertIsNot(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPD)
+
+  # this is the fallback slow VECTORIZE, 1 vinsertps per src in VECTORIZE
+  def test_vinsertps(self):
+    a = UOp.variable("a", 0, 0, dtypes.float32.vec(4))
+    b = UOp.variable("b", 0, 0, dtypes.float32.vec(4))
+    c = UOp.variable("c", 0, 0, dtypes.float32.vec(4))
+    d = UOp.variable("e", 0, 0, dtypes.float32)
+    # pack 1 from vector and 1 from scalar, moving 0th element to position 0 does nothing so only 1 vinsertps is generated
+    n = self.isel_rewrite(UOp.vectorize(a.gep(0), d))
+    self.assertIs(n.arg, X86Ops.VINSERTPS)
+    self.assertIsNot(n.src[0].arg, X86Ops.VINSERTPS)
+
+    valid = [UOp.vectorize(a.gep(0), b.gep(1), a.gep(2), b.gep(3)), # TODO: this should be vunpck
+             UOp.vectorize(a.gep(3), b.gep(2), c.gep(1), d)]
+    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VINSERTPS)
+
+  # complex address is [base + index*scale + displacement]
+  def test_complex_address(self):
+    a = UOp.variable("a", 0, 0, dtypes.int32)
+    load = UOp(Ops.PARAM, dtypes.int32.ptr(), arg=0).index(a + 1, ptr=True).load()
+    n = self.isel_rewrite(load)
+    # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
+    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].arg == 4)
+
+  def test_fold_load(self):
+    load1 = UOp(Ops.PARAM, dtypes.int32.ptr(), arg=0).index(UOp.const(dtypes.int32, 0), ptr=True).load()
+    load2 = UOp(Ops.PARAM, dtypes.int32.ptr(), arg=0).index(UOp.const(dtypes.int32, 1), ptr=True).load()
+    n = self.isel_rewrite(load1 + load2)
+    self.assertTrue(len(n.src) == 4)
+
+  # don't fold when used multiple times
+  def test_dont_fold_load(self):
+    load = UOp(Ops.PARAM, dtypes.int32.ptr(), arg=0).index(UOp.const(dtypes.int32, 0), ptr=True).load()
+    # used by multiple users
+    n = self.isel_rewrite(load + 1 + load)
+    self.assertTrue(len(n.src) == 2)
+    # used mutiple times by same user
+    n = self.isel_rewrite(load * load)
+    self.assertTrue(len(n.src) == 2)
+
+  # TODO: might want to check that load isn't part of another range when fusing
+
+if __name__ == "__main__":
+  unittest.main()

test/backend/test_schedule.py

@@ -510,7 +510,7 @@ class TestSchedule(unittest.TestCase):
     np.testing.assert_allclose(out[1].numpy(), np.sqrt(np.square(y.numpy() - np_mu).sum(-1)/y.shape[-1]), atol=1e-4, rtol=1e-4)
 
   def test_cumsum_parallel_reduce_fused(self):
-    # two-stage cumsum + ops triggers parallel REDUCEs in one kernel that must share an END
+    # two-stage cumsum + ops triggers parallel REDUCEs in one kernel that must share an END (same nesting context = should merge)
     step, num_steps = 513, 10
     t = Tensor.arange(step).float().realize()
     phase = t.cumsum()
@@ -521,6 +521,12 @@ class TestSchedule(unittest.TestCase):
     expected = (expected * np.array([1,0,0,1,0,0,0,0,1,0]).reshape(num_steps, 1)).flatten()
     np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)
 
+  @unittest.skipIf(Device.DEFAULT == "CL", "TODO: fails on CI CL")
+  def test_reduce_different_nesting_depth(self):
+    # two REDUCEs sharing the same RANGE at different nesting depths must NOT merge
+    x = Tensor.arange(768).reshape(3, 256).float()
+    np.testing.assert_allclose((x.sum(axis=1) + x.sum(axis=1).sum()).numpy(), x.numpy().sum(axis=1) + x.numpy().sum(axis=1).sum())
+
   def test_multimatmul_fusion(self):
     Tensor.manual_seed(0)
     a,b = Tensor.randn(4, 64).realize(), Tensor.rand(64,8).realize()

test/external/external_llm_eval.py

@@ -8,7 +8,7 @@ LABEL = ["A", "B", "C", "D"]
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
-  parser.add_argument("--port", "-p", type=int, default=11434)
+  parser.add_argument("--port", "-p", type=int, default=8000)
   parser.add_argument("--limit", "-L", type=int, default=None)
   parser.add_argument("--max_tokens", "-T", type=int, default=4096)
   parser.add_argument("--offset", "-O", type=int, default=0)

test/null/test_autogen.py

@@ -12,6 +12,51 @@ class TestC(unittest.TestCase):
       subprocess.check_output(('clang', '-x', 'c', '-fPIC', '-shared', '-', '-o', f.name), input=src.encode())
       return DLL("test", f.name)
 
+  def test_struct_array_init(self):
+    @record
+    class Foo:
+      SIZE = 12
+      a: Annotated[ctypes.c_int * 3, 0]
+    init_records()
+
+    f = Foo((1,2,3))
+    assert f.a[0] == 1
+    assert f.a[1] == 2
+    assert f.a[2] == 3
+    f = Foo((ctypes.c_int * 3)(1,2,3))
+    assert f.a[0] == 1
+    assert f.a[1] == 2
+    assert f.a[2] == 3
+
+  def test_field_ranges(self):
+    @record
+    class Foo:
+      SIZE = 2
+      s: Annotated[ctypes.c_int8, 0]
+      u: Annotated[ctypes.c_uint8, 1]
+    init_records()
+
+    f = Foo()
+    f.s = -1
+    f.u = -1
+    assert f.s == -1
+    assert f.u == 255
+
+  # this syntax is inherited from ctypes, but it seems a bit nonsensical?
+  def test_voidp_none(self):
+    @record
+    class Foo:
+      SIZE = 8
+      p: Annotated[ctypes.c_void_p, 0]
+    init_records()
+
+    f = Foo(None)
+    assert f.p is None
+    f.p = ctypes.c_void_p(0xDEADBEEF)
+    assert f.p == 0xDEADBEEF
+    f.p = None
+    assert f.p is None
+
   def test_packed_struct(self):
     @record
     class Baz:

test/null/test_viz.py

@@ -470,7 +470,7 @@ class TestVizProfiler(unittest.TestCase):
     j = load_profile(prof)
     event = j['layout']['NV:SDMA:0']['events'][0]
     gbs = sz/(dur*1e-6)*1e-9
-    self.assertEqual(event['fmt'], f"{gbs:.0f} GB/s")
+    self.assertEqual(event['fmt'], f"{gbs:.0f} GB/s\n{sz/1e6:.0f} MB")
 
   def test_graph(self):
     prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)),
@@ -512,7 +512,7 @@ class TestVizProfiler(unittest.TestCase):
     j = load_profile(prof)
     sdma_events = j['layout']['NV:1:SDMA:0']['events']
     gbs = sz/(dur*1e-6)*1e-9
-    self.assertEqual(sdma_events[0]['fmt'], f"{gbs:.0f} GB/s")
+    self.assertEqual(sdma_events[0]["fmt"], f"{gbs:.0f} GB/s\n{sz/1e6:.0f} MB")
 
   def test_block_ordering(self):
     prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)),

test/unit/test_gguf.py

@@ -30,11 +30,15 @@ class TestGGUF(unittest.TestCase):
 
   def test_dequantization_q4_0(self): self._test_dequantization(GGMLQuantizationType.Q4_0)
   def test_dequantization_q4_1(self): self._test_dequantization(GGMLQuantizationType.Q4_1)
+  def test_dequantization_q5_0(self): self._test_dequantization(GGMLQuantizationType.Q5_0)
+  def test_dequantization_q5_1(self): self._test_dequantization(GGMLQuantizationType.Q5_1)
   def test_dequantization_q8_0(self): self._test_dequantization(GGMLQuantizationType.Q8_0)
   def test_dequantization_q4_k(self): self._test_dequantization(GGMLQuantizationType.Q4_K)
   def test_dequantization_q5_k(self): self._test_dequantization(GGMLQuantizationType.Q5_K)
   def test_dequantization_q6_k(self): self._test_dequantization(GGMLQuantizationType.Q6_K)
   def test_dequantization_mxfp4(self): self._test_dequantization(GGMLQuantizationType.MXFP4)
+  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "Backend must support bfloat16")
+  def test_dequantization_bf16(self): self._test_dequantization(GGMLQuantizationType.BF16)
   def test_dequantization_mxfp4_old(self):
     def encode(nibbles, E):
       packed = [(low & 0xF) | ((high & 0xF) << 4) for low, high in zip(nibbles[:16], nibbles[16:])]
@@ -120,17 +124,21 @@ class TestGGUF(unittest.TestCase):
 class TestGGUFGEMV(unittest.TestCase):
   def _test_gguf_gemv(self, qtype: GGMLQuantizationType):
     block_size, type_size = GGML_QUANT_SIZES[qtype]
-    rows, cols = 8192, 2048
+    rows, cols = (1024, 512) if qtype == GGMLQuantizationType.BF16 else (8192, 2048)
     n_blocks = rows * cols // block_size
     rng = np.random.default_rng(42)
-    # generate random quantized blocks with valid fp16 scale fields (random bytes can produce NaN scales)
-    q_data = rng.integers(0, 256, size=n_blocks * type_size, dtype=np.uint8).reshape(n_blocks, type_size)
-    scales = np.float16(rng.standard_normal(n_blocks * 4)).view(np.uint8).reshape(n_blocks, -1)
-    if qtype == GGMLQuantizationType.Q8_0: q_data[:, :2] = scales[:, :2]                    # d at offset 0
-    elif qtype in (GGMLQuantizationType.Q4_K, GGMLQuantizationType.Q5_K): q_data[:, :4] = scales[:, :4] # d, dmin at offset 0
-    elif qtype == GGMLQuantizationType.Q6_K: q_data[:, -2:] = scales[:, :2]                 # d at end
-    elif qtype == GGMLQuantizationType.MXFP4: q_data[:, 0] = rng.integers(120, 136, size=n_blocks, dtype=np.uint8) # constrain byte0
-    q_data = q_data.flatten()
+    if qtype == GGMLQuantizationType.BF16:
+      q_data = (rng.standard_normal(rows * cols).astype(np.float32).view(np.uint32) >> 16).astype(np.uint16).view(np.uint8)
+    else:
+      # generate random quantized blocks with valid fp16 scale fields (random bytes can produce NaN scales)
+      q_data = rng.integers(0, 256, size=n_blocks * type_size, dtype=np.uint8).reshape(n_blocks, type_size)
+      scales = np.float16(rng.standard_normal(n_blocks * 4)).view(np.uint8).reshape(n_blocks, -1)
+      if qtype in (GGMLQuantizationType.Q5_0, GGMLQuantizationType.Q8_0): q_data[:, :2] = scales[:, :2]  # d at offset 0
+      elif qtype in (GGMLQuantizationType.Q5_1, GGMLQuantizationType.Q4_K, GGMLQuantizationType.Q5_K):
+        q_data[:, :4] = scales[:, :4]  # d, m/dmin at offset 0
+      elif qtype == GGMLQuantizationType.Q6_K: q_data[:, -2:] = scales[:, :2]               # d at end
+      elif qtype == GGMLQuantizationType.MXFP4: q_data[:, 0] = rng.integers(120, 136, size=n_blocks, dtype=np.uint8) # constrain byte0
+      q_data = q_data.flatten()
     ref = dequantize(q_data, qtype).reshape(rows, cols)
 
     # build a minimal gguf in memory: header + 1 tensor info + aligned data
@@ -148,15 +156,18 @@ class TestGGUFGEMV(unittest.TestCase):
 
     x = rng.standard_normal(cols).astype(np.float32)
     np.testing.assert_allclose((tensors["weight"] @ Tensor(x)).numpy(), ref @ x, atol=1e-2, rtol=1e-2)
-    # can only expect the weights to be identical if we really support float16 (ie. not decompositions)
-    if is_dtype_supported(dtypes.half): np.testing.assert_equal(tensors["weight"].numpy(), ref)
+    if qtype == GGMLQuantizationType.BF16 or is_dtype_supported(dtypes.half): np.testing.assert_equal(tensors["weight"].numpy(), ref)
     assert np.isfinite(ref).all() and np.isfinite(tensors["weight"].numpy()).all(), f"{qtype.name} has NaN/Inf"
 
   def test_gguf_gemv_q8_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q8_0)
+  def test_gguf_gemv_q5_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_0)
+  def test_gguf_gemv_q5_1(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_1)
   def test_gguf_gemv_q4_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q4_K)
   def test_gguf_gemv_q5_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_K)
   def test_gguf_gemv_q6_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q6_K)
   def test_gguf_gemv_mxfp4(self): self._test_gguf_gemv(GGMLQuantizationType.MXFP4)
+  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "Backend must support bfloat16")
+  def test_gguf_gemv_bf16(self): self._test_gguf_gemv(GGMLQuantizationType.BF16)
 
 if __name__ == '__main__':
   unittest.main()

tinygrad/apps/llm.py

@@ -287,8 +287,7 @@ models = {
   "olmoe": "https://huggingface.co/allenai/OLMoE-1B-7B-0924-Instruct-GGUF/resolve/main/olmoe-1b-7b-0924-instruct-q4_k_m.gguf",
 }
 
-# *** simple OpenAI compatible server on 11434 to match ollama ***
-# OPENAI_BASE_URL=http://localhost:11434/v1 OPENAI_API_KEY=ollama uvx --from gpt-command-line gpt
+# *** simple OpenAI API compatible server with web interface on http://localhost:8000/ ***
 
 CHAT_HTML = b'''<!DOCTYPE html><html><head><title>tinygrad chat</title><style>
   * { margin: 0 }
@@ -397,7 +396,7 @@ if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument("--model", "-m", default=list(models.keys())[0], help=f"Model choice ({', '.join(models.keys())}) or path to a local GGUF file")
   parser.add_argument("--max_context", type=int, default=4096, help="Max Context Length")
-  parser.add_argument("--serve", nargs='?', type=int, const=11434, metavar="PORT", help="Run OpenAI compatible API (optional port, default 11434)")
+  parser.add_argument("--serve", nargs='?', type=int, const=8000, metavar="PORT", help="Run OpenAI compatible API (optional port, default 8000)")
   parser.add_argument("--warmup", action="store_true", help="warmup the JIT")
   parser.add_argument("--benchmark", nargs='?', type=int, const=20, metavar="COUNT", help="Benchmark tok/s (optional count, default 20)")
   args = parser.parse_args()

tinygrad/codegen/__init__.py

@@ -101,6 +101,13 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
   # this was the linearizer
   sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True)
 
+  from tinygrad.renderer.isa import ISARenderer, IselContext
+  if isinstance(ren, ISARenderer):
+    linear_sink = graph_rewrite(sink, ren.pre_isel_matcher, name="pre instruction selection", bottom_up=True)
+    isel_ctx = IselContext(linear_sink)
+    linear_sink = graph_rewrite(linear_sink, ren.isel_matcher, ctx=isel_ctx, name="instruction selection", bottom_up=True)
+    sink = linear_sink
+
   # return the rewritten sink
   return sink
 
@@ -114,20 +121,35 @@ pm_linearize_cleanups = PatternMatcher([
 ])
 
 # requires lst be toposorted. like graph rewrite, but for lines
-def line_rewrite(lst:list[UOp], pm:PatternMatcher) -> list[UOp]:
+def line_rewrite(lst:list[UOp], pm:PatternMatcher, ctx=None) -> list[UOp]:
   newlst = []
   replaced: dict[UOp, UOp] = {}
   for u in lst:
-    nu = u.replace(src=tuple([replaced[x] for x in u.src]))
-    ret: tuple[UOp, list[UOp]] = cast(tuple[UOp, list[UOp]]|None, pm.rewrite(nu)) or (nu, [nu])
+    nu = u.replace(src=tuple([replaced.get(x, x) for x in u.src]))
+    ret: tuple[UOp, list[UOp]] = cast(tuple[UOp, list[UOp]]|None, pm.rewrite(nu, ctx)) or (nu, [nu])
     replaced[u] = ret[0]
     newlst.extend(ret[1])
   return newlst
 
-def do_linearize(prg:UOp, sink:UOp) -> UOp:
-  lst = line_rewrite(linearize(sink), pm_linearize_cleanups)
-  if SPEC: type_verify(lst, program_spec)
-  return prg.replace(src=prg.src + (UOp(Ops.LINEAR, src=tuple(lst)),))
+def do_linearize(ctx:Renderer, prg:UOp, sink:UOp) -> UOp:
+  from tinygrad.renderer.isa import ISARenderer
+  generic_lst = line_rewrite(linearize(sink), pm_linearize_cleanups) if sink.arg.estimates is None and not isinstance(ctx, ISARenderer) else None
+  if isinstance(ctx, ISARenderer):
+    from tinygrad.renderer.isa import PreRegAllocContext
+    from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
+    lst = linearize(sink)
+    if ctx.pre_regalloc_matcher is not None: lst = line_rewrite(lst, ctx.pre_regalloc_matcher, PreRegAllocContext())
+    regalloc_ctx = LinearScanRegallocContext(lst, ctx)
+    lst = line_rewrite(lst, pm_regalloc_rewrite, regalloc_ctx)
+    if ctx.late_regalloc_matcher is not None: lst = line_rewrite(lst, ctx.late_regalloc_matcher, regalloc_ctx)
+    lst = line_rewrite(lst, ctx.post_regalloc_matcher, regalloc_ctx)
+    if DEBUG >= 4: print(ctx.asm(lst, sink.arg.function_name))
+    if SPEC: type_verify(lst, ctx.isa_spec)
+  else:
+    assert generic_lst is not None
+    lst = generic_lst
+    if SPEC: type_verify(lst, program_spec)
+  return prg.replace(src=(sink,)+prg.src[1:] + (UOp(Ops.LINEAR, src=tuple(lst)),))
 
 def do_estimates(prg:UOp, sink:UOp, lin:UOp) -> UOp|None:
   if sink.arg.estimates is not None: return None

tinygrad/codegen/late/devectorizer.py

@@ -328,11 +328,23 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
   return acc.after(end).index(UOp.const(dtypes.int, 0))
 
 def merge_reduce_ends(ctx:ReduceContext, sink:UOp):
-  # merge ENDs that share the same range (only those created by reduce_to_acc)
+  # merge ENDs that share the same range and nesting context (only those created by reduce_to_acc)
+  # ENDs at different nesting depths get cloned RANGEs so each RANGE maps to one END
   range_to_ends: dict[tuple[UOp, ...], list[UOp]] = {}
   for u in sink.backward_slice:
     if u.op is Ops.END and u.tag == "mergeable": range_to_ends.setdefault(u.src[1:], []).append(u)
-  subs = {e: UOp.group(*(e.src[0] for e in ends)).end(*r) for r, ends in range_to_ends.items() if len(ends) > 1 for e in ends}
+  subs: dict[UOp, UOp] = {}
+  next_axis = max((u.arg[0] for u in sink.backward_slice if u.op is Ops.RANGE), default=-1) + 1
+  for r, ends in range_to_ends.items():
+    if len(ends) <= 1: continue
+    by_ctx: dict[frozenset[UOp], list[UOp]] = {}
+    for e in ends: by_ctx.setdefault(frozenset(e.ranges), []).append(e)
+    for i, group in enumerate(by_ctx.values()):
+      tr = r if i == 0 else tuple(rr.replace(arg=(next_axis + j, *rr.arg[1:])) for j, rr in enumerate(r))
+      if i > 0: next_axis += len(r)
+      mapped = [e.substitute(dict(zip(r, tr))) if i > 0 else e for e in group]
+      merged = mapped[0] if len(mapped) == 1 else UOp.group(*(e.src[0] for e in mapped)).end(*tr)
+      for e in group: subs[e] = merged
   return sink.substitute(subs) if subs else None
 
 pm_reduce = PatternMatcher([

tinygrad/codegen/late/linearizer.py

@@ -2,7 +2,7 @@ import heapq
 from typing import Any
 from collections import defaultdict
 from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, multirange_str
-from tinygrad.helpers import prod, getenv, TUPLE_ORDER
+from tinygrad.helpers import prod, getenv, TUPLE_ORDER, DEV
 
 def linearize(sink:UOp) -> list[UOp]:
   # this is a toposort with priority
@@ -31,10 +31,18 @@ def linearize(sink:UOp) -> list[UOp]:
       case Ops.RANGE: priority = 5    # placing RANGE is good
       case Ops.END: priority = -5     # placing END is bad
       case _: priority = 0            # everything else has priority 0
+
+    # stack pointer needs to be scheduled at the top of the kernel
+    # TODO: remove once there's a proper isa scheduler
+    if u.op is Ops.INS:
+      from tinygrad.renderer.isa.x86 import X86Ops, RSP
+      match u.arg:
+        case X86Ops.DEFINE_REG: priority, extra = (-21 if u.tag[0] == RSP else -20), u.tag[0].index
+
     priorities[u] = (run_count, priority, extra)
 
   # number the uops in "ideal" order
-  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if TUPLE_ORDER else ())))}
+  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if TUPLE_ORDER and DEV.value.renderer != "X86" else ())))}
 
   # then force them to be toposorted in as close to the ideal order as possible
   heap = [(-nkey[sink], sink)]
@@ -93,4 +101,4 @@ def do_split_ends(e:UOp):
 pm_split_ends = PatternMatcher([
   # split the ends
   (UPat(Ops.END, name="e"), do_split_ends),
-])
+])

tinygrad/codegen/late/regalloc.py

@@ -0,0 +1,132 @@
+import itertools
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
+from tinygrad.renderer.isa import ISARenderer, Register
+from tinygrad.dtype import dtypes, PtrDType
+
+PSEUDO_OPS = {Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP}
+def _uop_key(u:UOp): return (u.op, u.dtype, u.arg)
+
+# loosely based on: https://bernsteinbear.com/assets/img/register-spilling-range-splitting-ssa.pdf
+class LinearScanRegallocContext:
+  def __init__(self, uops:list[UOp], ren:ISARenderer):
+    if saved:=ren.callee_saved():
+      ret_i = next(i for i,u in reversed(tuple(enumerate(uops))) if u.op is Ops.INS and getattr(u.arg, "name", None) == "RET")
+      uops[0:0] = saved
+      uops[ret_i+len(saved)] = uops[ret_i+len(saved)].replace(src=uops[ret_i+len(saved)].src + saved)
+    live_range: dict[Register, list[int]] = {}
+    live: dict[Register, Register] = {}
+    live_ins: list[dict[Register, Register]] = []
+    self.defs: dict[Register, UOp] = {} # mapping from virtual to uop that defines it
+    self.real_defs: dict[Register, Register] = {} # mapping from virtual to real at definition
+    self.spills: dict[Register, UOp] = {} # mapping from virtual to stack slot
+    self.fills: dict[int, dict[int, tuple[Register, Register]]] = {} # mapping from program point to mapping from idx to virtual and real to fill to
+    self.insert_before: dict[int, list[tuple[Register, Register]]] = {} # mapping from program point to fills to be inserted
+    self.idx = itertools.count()
+    self.ren = ren
+    self.stack_size = 0
+    # the label associated with each loop NOTE: this is only used post regalloc and should be removed
+    self.loop_label: dict[UOp, str] = {}
+    arg_order = {Ops.PARAM: 0, Ops.DEFINE_VAR: 1, Ops.SPECIAL: 2}
+    self.func_arg_idxs = {_uop_key(u): i for i,u in enumerate(sorted({u for u in uops if u.op in arg_order}, key=lambda k: (arg_order[k.op], k.arg)))}
+    self.local_offsets: dict[tuple, int] = {}
+    for u in uops:
+      if u.op not in (Ops.DEFINE_LOCAL, Ops.DEFINE_REG): continue
+      self.local_offsets.setdefault(_uop_key(u), self.stack_size)
+      self.stack_size += u.dtype.nbytes()
+    # compute live ranges
+    lr = live_range
+    ranges: list[Register] = []
+    for i,u in enumerate(reversed(uops)):
+      if u.op in PSEUDO_OPS: continue
+      defs = u.tag if isinstance(u.tag, tuple) else ()
+      for v in defs + tuple(s.reg for s in set(u.src)):
+        if isinstance(v, Register): lr.setdefault(v, []).insert(0, len(uops) - 1 - i)
+      for v in defs:
+        if isinstance(v, Register): self.defs[v] = u
+        if v in lr and (n:=max((lr[rng][-1] for rng in ranges if lr[rng][0] <= lr[v][-1] < lr[rng][-1]), default=None)): lr[v].append(n)
+      if u.op is Ops.RANGE: ranges.append(u.reg)
+
+    def alloc(cons:tuple[Register, ...], i:int) -> Register:
+      live_inv = {v:k for k,v in live.items()}
+      # allocate the best register. Registers not in live or not used again are free and have priority,
+      # otherwise pick the one with the furthest next use. Regs that appear first in cons have priority in case of a tie
+      reg,vreg = max(((r,live_inv.get(r)) for r in cons),
+                    key=lambda rv: next((j-i for j in ([] if rv[1] is None else live_range[rv[1]]) if j >= i), len(uops)))
+      return live.pop(vreg) if vreg is not None else reg
+
+    # assign register to spilled virtual and record load to be emitted before current uop, also assign it a stack slot
+    def fill(v:Register, i:int, cons:tuple[Register, ...]|None=None) -> Register:
+      if v not in self.spills:
+        dt = self.defs[v].dtype
+        sz = dt.scalar().itemsize * dt.count if not isinstance(dt, PtrDType) else 8
+        assert sz > 0
+        offset = self.stack_size + (sz - self.stack_size % sz) % sz
+        self.spills[v] = UOp.const(dtypes.int32, offset)
+        self.stack_size = offset + sz
+      r = alloc(cons if cons is not None else v.cons, i)
+      self.insert_before.setdefault(i, []).append((v, r))
+      return r
+
+    for i,u in enumerate(uops):
+      if u.op in PSEUDO_OPS: continue
+      # allocate uses
+      for j,s in enumerate(u.src):
+        # HACK: cause of later hacks to lower range
+        if u.op is Ops.END: continue
+        if not isinstance(v:=s.reg, Register): continue
+        if v not in live: live[v] = fill(v, i)
+        if v in self.spills: self.fills.setdefault(i, {})[j] = (v, live[v])
+
+      # allocate defs
+      if isinstance(u.tag, tuple):
+        for j,v in enumerate(u.tag):
+          assert isinstance(v, Register) and v not in live
+          cons = v.cons
+          # two address instructions (src is reused by def) can only coalesce reused src. reused src goes first to get priority in case of a tiebreak
+          if ren.is_two_address(u) and j == 0:
+            ins = tuple(live.get(s.reg) for s in u.src)
+            cons = ((ins[0],) if ins[0] in cons else ()) + tuple(r for r in cons if r not in ins)
+            assert cons
+          # HACK: cause the range is missing the comparison
+          self.real_defs[v] = live[v] = alloc(cons, i+1 if u.op is not Ops.RANGE else i)
+
+      # loop prologue, avoid loading inside the loop
+      if u.op is Ops.RANGE:
+        # we move to registers vars used in the loop sorted by next use, vars not used in the loop will not be reloaded in the epilogue
+        used_in_loop = [v for v in live.keys() | self.spills.keys() if any(i <= l < live_range[u.reg][-1] for l in live_range[v])]
+        sorted_uses = sorted(used_in_loop, key=lambda k: next(l-i for l in live_range[k] if l >= i))
+        live_in: dict[Register, Register] = {}
+        for v in sorted_uses:
+          # if all the possible registers are already in live_in there's no space for this var
+          if set(v.cons).issubset(live_in.values()): continue
+          if v not in live: live[v] = fill(v, i)
+          assert live[v] not in live_in.values()
+          live_in[v] = live[v]
+        live_ins.append(live_in)
+
+      # loop epilogue, reload registers that were live at loop entry
+      if u.op is Ops.END:
+        # TODO: if a uop is in a different reg in live out vs live in move between registers instead of loading
+        # TODO: don't reload if first use in loop is a load
+        for v,r in live_ins.pop().items():
+          if v not in live or live[v] != r: live[v] = fill(v, i, (r,))
+
+def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp):
+  i = next(ctx.idx)
+  if x.op in PSEUDO_OPS: return None
+  nsrc = []
+  for j,s in enumerate(x.src):
+    if i in ctx.fills and j in ctx.fills[i]:
+      v,r = ctx.fills[i][j]
+      nsrc.append(ctx.ren.fill(ctx.spills[v], ctx.defs[v], r))
+    else: nsrc.append(s)
+  ndefs = tuple(ctx.real_defs[v] for v in x.tag) if isinstance(x.tag, tuple) else x.tag
+  nx = x.replace(src=tuple(nsrc), tag=ndefs)
+  fills = [ctx.ren.fill(ctx.spills[v], ctx.defs[v], r) for v,r in ctx.insert_before.get(i, [])]
+  spills = [ctx.ren.spill(ctx.spills[v], nx) for v in x.tag if v in ctx.spills] if isinstance(x.tag, tuple) else []
+  return nx, fills + [nx] + spills
+
+pm_regalloc_rewrite = PatternMatcher([
+  (UPat({Ops.INS, Ops.CONST, Ops.RANGE, Ops.END, Ops.NOOP, Ops.GROUP, Ops.AFTER, Ops.BARRIER,
+         Ops.PARAM, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}, name="x"), regalloc_rewrite),
+])

tinygrad/device.py

@@ -319,7 +319,8 @@ def is_dtype_supported(dtype:DType, target:Target|None=None) -> bool:
       case "METAL": return not CI or BENCHMARKS
       case "CUDA": return (not CI or BENCHMARKS) and target.renderer != "PTX"
       case "NV": return (not CI or BENCHMARKS) and target.renderer not in ("PTX", "NAK")
-      case "CPU": return (not CI or BENCHMARKS) and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and target.renderer != "LVP"
+      case "CPU": return (not CI or BENCHMARKS) and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and \
+          target.renderer not in ("LVP", "X86")
       case "AMD" | "CL" | "PYTHON" | "NULL": return True
       case _: return False
   if dtype in dtypes.fp8_ocp:

tinygrad/engine/realize.py

@@ -9,6 +9,27 @@ from tinygrad.device import Device, Buffer
 from tinygrad.renderer import ProgramSpec, Estimates
 from tinygrad.codegen import get_program
 
+# **************** Stat ****************
+
+def update_stats(display_name:str, device:str, estimates:Estimates, var_vals:dict[str, int], et:float|None, buf_count:int,
+                 jit=False, metadata:tuple[Metadata, ...]=(), first_run=False):
+  GlobalCounters.kernel_count += 1
+  GlobalCounters.global_ops += (op_est:=sym_infer(estimates.ops, var_vals))
+  GlobalCounters.global_mem += (mem_est:=sym_infer(estimates.mem, var_vals))
+  if et is not None: GlobalCounters.time_sum_s += et
+  if DEBUG >= 2:
+    lds_est = sym_infer(estimates.lds, var_vals)
+    header_color = 'magenta' if jit else ('green' if first_run else None)
+    ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
+    flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
+    flops_str = f"{flops*1e-9:7.0f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:7.0f} TFLOPS", 'green')
+    mem_str = f"{membw*1e-9:4.0f}|{ldsbw*1e-9:<6.0f} GB/s" if membw < 1e13 and ldsbw < 1e15 else \
+      colored(f"{membw*1e-12:4.0f}|{ldsbw*1e-12:<6.0f} TB/s", 'green')
+    print(f"{colored(f'*** {device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
+      f" {display_name+' '*(46-ansilen(display_name))} arg {buf_count:2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
+      ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})")+
+      f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in metadata] if metadata else ''}")
+
 # **************** Runners ****************
 
 class Runner:
@@ -166,22 +187,7 @@ class ExecItem:
       cpu_events.append(ProfilePointEvent(self.prg.device, "exec", len(cpu_events), payload))
     et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2)
     if do_update_stats:
-      GlobalCounters.kernel_count += 1
-      GlobalCounters.global_ops += (op_est:=sym_infer(self.prg.estimates.ops, var_vals))
-      GlobalCounters.global_mem += (mem_est:=sym_infer(self.prg.estimates.mem, var_vals))
-      if et is not None: GlobalCounters.time_sum_s += et
-      if DEBUG >= 2:
-        lds_est = sym_infer(self.prg.estimates.lds, var_vals)
-        header_color = 'magenta' if jit else ('green' if self.prg.first_run else None)
-        ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
-        flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
-        flops_str = f"{flops*1e-9:7.0f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:7.0f} TFLOPS", 'green')
-        mem_str = f"{membw*1e-9:4.0f}|{ldsbw*1e-9:<6.0f} GB/s" if membw < 1e13 and ldsbw < 1e15 else \
-          colored(f"{membw*1e-12:4.0f}|{ldsbw*1e-12:<6.0f} TB/s", 'green')
-        print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
-          f" {self.prg.display_name+' '*(46-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
-          ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})")+
-          f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in self.metadata] if self.metadata else ''}")
+      update_stats(self.prg.display_name, self.prg.device, self.prg.estimates, var_vals, et, len(bufs), jit, self.metadata, self.prg.first_run)
       self.prg.first_run = False
     return et
 

tinygrad/helpers.py

@@ -251,8 +251,6 @@ ALLOW_TF32 = ContextVar("ALLOW_TF32", 0)
 SCACHE = ContextVar("SCACHE", 1)
 # allow use of atomics for embedding backward
 USE_ATOMICS = ContextVar("USE_ATOMICS", 0)
-# allow use of assembly for gemm
-ASM_GEMM = ContextVar("ASM_GEMM", 0)
 
 @dataclass(frozen=True)
 class Metadata:

tinygrad/mixin/movement.py

@@ -107,6 +107,12 @@ class MovementMixin:
     ret = self._mop(Ops.RESHAPE, arg=new_shape)
     return self if ret.shape == self.shape else ret
 
+  def pad(self, arg:tuple[tuple[sint, sint] | None, ...]) -> Self:
+    if self.ndim != len(arg):
+      raise ValueError(f"{self.ndim=} != {len(arg)=}")
+    ret = self._mop(Ops.PAD, tuple(x if x is not None else (0, 0) for x in arg))
+    return self if ret.shape == self.shape else ret
+
   def shrink(self, arg: tuple[tuple[sint, sint] | None, ...]) -> Self:
     """
     Returns a tensor that shrinks the each axis based on input arg.

tinygrad/nn/state.py

@@ -297,13 +297,19 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
   """
   Converts ggml tensor data to a tinygrad tensor.
 
-  Supported native types: float32 (id: 0), float16 (id: 1), int8 (id: 16), int16 (id: 17), int32 (id: 18)
-  Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q8_0 (id: 8), Q4_K (id: 12), Q5_K (id: 13), Q6_K (id: 14), MXFP4 (id: 39)
+  Supported native types: float32 (id: 0), float16 (id: 1), bfloat16 (id: 30),
+  int8 (id: 16), int16 (id: 17), int32 (id: 18)
+  Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q5_0 (id: 6),
+  Q5_1 (id: 7), Q8_0 (id: 8), Q4_K (id: 12), Q5_K (id: 13),
+  Q6_K (id: 14), MXFP4 (id: 39)
   """
   # https://github.com/ggerganov/ggml/blob/323951f1bdcdfbd5b5ff3a9a7c3770e63b1a560e/include/ggml.h#L356
 
   # native types
-  if (dtype := { 0: dtypes.float32, 1: dtypes.float16, 16: dtypes.int8, 17: dtypes.int16, 18: dtypes.int32 }.get(ggml_type)) is not None:
+  if (dtype := {
+    0: dtypes.float32, 1: dtypes.float16, 16: dtypes.int8,
+    17: dtypes.int16, 18: dtypes.int32, 30: dtypes.bfloat16,
+  }.get(ggml_type)) is not None:
     return t[:dtype.itemsize * n].contiguous().bitcast(dtype)
 
   def q_to_uint8(t: Tensor, b: int) -> Tensor:
@@ -312,12 +318,21 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
     return t.unsqueeze(-1).expand((*t.shape,8//b)).idiv(shift_tensor).bitwise_and(bitmask).transpose(-1, -2).flatten(-2)
 
   # map to (number of elements, number of bytes)
-  if (nelements_nbytes := { 2:(32,18), 3:(32,20), 8:(32,34), 12:(256,144), 13:(256,176), 14:(256,210), 39:(32,17) }.get(ggml_type)) is not None:
+  if (nelements_nbytes := {
+    2:(32,18), 3:(32,20), 6:(32,22), 7:(32,24), 8:(32,34),
+    12:(256,144), 13:(256,176), 14:(256,210), 39:(32,17),
+  }.get(ggml_type)) is not None:
     blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1])).contiguous()
     if ggml_type == 2: return (q_to_uint8(blocks[:,2:], 4).bitcast(dtypes.int8) - 8) * blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
     if ggml_type == 3:
       d, m = (blocks[:,s:s+2].bitcast(dtypes.float16).cast(dtypes.float32) for s in [ 0, 2 ])
       return q_to_uint8(blocks[:,4:], 4).bitcast(dtypes.int8) * d + m
+    if ggml_type in (6, 7):
+      d = blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
+      qh_off = 2 if ggml_type == 6 else 4
+      qh = q_to_uint8(blocks[:,qh_off:qh_off+4], 1).reshape((-1, 8, 4)).transpose(-1, -2).flatten(-2).bitcast(dtypes.int8)
+      q = q_to_uint8(blocks[:,qh_off+4:], 4).bitcast(dtypes.int8) + qh * 16
+      return q * d + (blocks[:,2:4].bitcast(dtypes.float16).cast(dtypes.float32) if ggml_type == 7 else -16 * d)
     if ggml_type == 8: return blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32) * blocks[:,2:].bitcast(dtypes.int8)
      # Q4_K: 256 elements per 144-byte block (d:2, dmin:2, scales:12, qs:128)
      # Q5_K: 256 elements per 176-byte block (d:2, dmin:2, scales:12, qh:32, qs:128)

tinygrad/renderer/amd/elf.py

@@ -3,7 +3,7 @@ import ctypes
 from tinygrad.helpers import ceildiv, round_up
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.runtime.autogen import amdgpu_kd, hsa, libc
-from tinygrad.renderer.amd.dsl import Reg, FixedBitField
+from tinygrad.renderer.amd.dsl import Inst, Reg, FixedBitField
 from tinygrad.runtime.autogen.amd.common import OpType
 
 # instructions used for padding
@@ -11,8 +11,9 @@ from tinygrad.runtime.autogen.amd.rdna3.ins import s_code_end # same encoding as
 from tinygrad.runtime.autogen.amd.cdna.ins import s_nop as s_nop_cdna
 
 _arch_map = {"gfx9": "cdna", "gfx10": "rdna3", "gfx11": "rdna3", "gfx12": "rdna4"}
-def do_assemble_amd(ctx, prg:UOp, lin:UOp) -> UOp:
+def do_assemble_amd(ctx, prg:UOp, lin:UOp) -> UOp|None:
   insts = [u.arg for u in lin.src]
+  if not all(isinstance(inst, Inst) for inst in insts): return None
 
   # ** scan for max vgpr/sgpr/accvgpr
   max_vgpr, max_sgpr, max_accvgpr = 0, 0, 0

tinygrad/renderer/isa/__init__.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+import itertools
+from dataclasses import dataclass, field
+from tinygrad.renderer import Renderer
+from tinygrad.uop.ops import PatternMatcher, UOp, Ops
+
+@dataclass(frozen=True)
+class Register:
+  name: str
+  index: int
+  _cons: tuple[Register, ...] = field(default_factory=tuple)
+  @property
+  def cons(self): return self._cons or (self,)
+  def __repr__(self): return self.name
+
+class IselContext:
+  def __init__(self, sink:UOp):
+    self.uses = sink.get_consumer_map()
+    self.reg_n = itertools.count()
+    arg_order = {Ops.PARAM: 0, Ops.DEFINE_VAR: 1, Ops.SPECIAL: 2}
+    self.func_args = sorted([u for u in self.uses if u.op in arg_order], key=lambda k: (arg_order[k.op], k.arg))
+
+  def vreg(self, cons:tuple[Register, ...]|Register):
+    return Register(f"v{next(self.reg_n)}", 0, _cons=cons if isinstance(cons, tuple) else (cons,))
+
+@dataclass
+class PreRegAllocContext:
+  lock: UOp|None = None
+  clobbered: set[UOp] = field(default_factory=set)
+
+class ISARenderer(Renderer):
+  isa_spec: PatternMatcher
+  pre_isel_matcher: PatternMatcher
+  isel_matcher: PatternMatcher
+  pre_regalloc_matcher: PatternMatcher|None = None
+  late_regalloc_matcher: PatternMatcher|None = None
+  post_regalloc_matcher: PatternMatcher
+
+  def callee_saved(self) -> tuple[UOp, ...]: return tuple()
+  def is_two_address(self, x:UOp) -> bool: return False
+  def copy(self, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific")
+  def spill(self, disp:UOp, x:UOp) -> UOp: raise NotImplementedError("arch specific")
+  def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific")
+  def asm(self, uops:list[UOp], function_name:str) -> str: raise NotImplementedError("arch specific")

tinygrad/renderer/isa/x86.py

@@ -0,0 +1,936 @@
+# flake8: noqa: E702
+# allow semicolons to put multiple ops on one line
+import sys, struct, functools
+from typing import cast
+from tinygrad.dtype import dtypes, PtrDType, DType, truncate
+from tinygrad.uop import FastEnum, auto, Ops, GroupOp
+from tinygrad.uop.ops import UOp, UPat, PatternMatcher
+from tinygrad.renderer.isa import ISARenderer, IselContext, Register, PreRegAllocContext
+from tinygrad.helpers import getenv, CPU_COUNT, unwrap, Target
+
+# ***** X86 Ops *****
+
+class X86Ops(FastEnum):
+  # NOTE: X86Ops with i suffix are variants that take an immediate, m suffix are variants that can write to memory instead of read from
+  # these aren't real instructions
+  DEFINE_REG = auto(); FRAME_INDEX = auto(); LABEL = auto()
+  # index
+  LEA = auto()
+  # register / memory / immediate moves
+  MOV = auto(); MOVm = auto(); MOVi = auto(); MOVABS = auto()
+  VMOVSS = auto(); VMOVSD = auto(); VMOVUPS = auto()
+  VMOVSSm = auto(); VMOVSDm = auto(); VMOVUPSm = auto()
+  # casts
+  MOVZX = auto(); MOVSX = auto(); MOVSXD = auto()
+  VPMOVZXBW = auto(); VPMOVZXBD = auto(); VPMOVZXBQ = auto()
+  VPMOVZXWD = auto(); VPMOVZXWQ = auto(); VPMOVZXDQ = auto()
+  VPMOVSXBW = auto(); VPMOVSXBD = auto(); VPMOVSXBQ = auto()
+  VPMOVSXWD = auto(); VPMOVSXWQ = auto(); VPMOVSXDQ = auto()
+  VCVTDQ2PS = auto(); VCVTDQ2PD = auto(); VCVTTPS2DQ = auto(); VCVTTPD2DQ = auto()
+  VCVTPH2PS = auto(); VCVTPS2PH = auto(); VCVTPS2PD = auto(); VCVTPD2PS = auto()
+  VCVTSS2SD = auto(); VCVTSD2SS = auto(); VCVTSI2SS = auto(); VCVTSI2SD = auto()
+  VCVTTSS2SI = auto(); VCVTTSD2SI = auto()
+  # bitcasts
+  VMOVD = auto(); VMOVQ = auto(); VMOVDm = auto(); VMOVQm = auto()
+  # comparisons
+  VUCOMISS = auto(); VUCOMISD = auto()
+  VCMPSS = auto(); VCMPSD = auto(); VCMPPS = auto(); VCMPPD = auto()
+  VPCMPGTB = auto(); VPCMPGTW = auto(); VPCMPGTD = auto(); VPCMPGTQ = auto()
+  VPCMPEQB = auto(); VPCMPEQW = auto(); VPCMPEQD = auto(); VPCMPEQQ = auto()
+  SETNE = auto(); SETE = auto(); SETL = auto(); SETB = auto()
+  # where
+  CMOVNE = auto(); CMOVE = auto(); CMOVL = auto(); CMOVB = auto()
+  VPBLENDVB = auto(); VBLENDVPS = auto(); VBLENDVPD = auto()
+  # jumps
+  JNE = auto(); JE = auto(); JL = auto(); JB = auto(); JGE = auto(); JMP = auto()
+  # vectorize / gep
+  VSHUFPS = auto(); VSHUFPD = auto(); VINSERTPS = auto(); VPSRLDQ = auto()
+  VPEXTRB = auto(); VPEXTRW = auto(); VPEXTRD = auto(); VPEXTRQ = auto()
+  VPINSRB = auto(); VPINSRW = auto(); VPINSRD = auto(); VPINSRQ = auto()
+  VPBROADCASTB = auto(); VPBROADCASTW = auto(); VPBROADCASTD = auto(); VPBROADCASTQ = auto()
+  VBROADCASTSS = auto()
+  # int binary
+  IDIV = auto(); DIV = auto()
+  ADD = auto(); ADDi = auto(); SUB = auto(); SUBi = auto(); IMUL = auto(); IMULi = auto()
+  AND = auto(); ANDi = auto(); XOR = auto(); XORi = auto(); OR = auto(); ORi = auto()
+  SHL = auto(); SHLi = auto(); SHR = auto(); SHRi = auto(); SAR = auto(); SARi = auto(); CMP = auto(); CMPi = auto()
+  # float unary (sometimes not unary)
+  VROUNDSS = auto(); VROUNDSD = auto(); VROUNDPS = auto(); VROUNDPD = auto()
+  VSQRTSS = auto(); VSQRTSD = auto(); VSQRTPS = auto(); VSQRTPD = auto()
+  # float scalar / vector binary
+  VADDSS = auto(); VADDSD = auto(); VADDPS = auto(); VADDPD = auto()
+  VSUBSS = auto(); VSUBSD = auto(); VSUBPS = auto(); VSUBPD = auto()
+  VMULSS = auto(); VMULSD = auto(); VMULPS = auto(); VMULPD = auto()
+  VDIVSS = auto(); VDIVSD = auto(); VDIVPS = auto(); VDIVPD = auto()
+  VMAXSS = auto(); VMAXSD = auto(); VMAXPS = auto(); VMAXPD = auto()
+  VMINSS = auto(); VMINSD = auto(); VMINPS = auto(); VMINPD = auto()
+  # int vector binary
+  VPADDB = auto(); VPADDW = auto(); VPADDD = auto(); VPADDQ = auto()
+  VPSUBB = auto(); VPSUBW = auto(); VPSUBD = auto(); VPSUBQ = auto()
+  VPMULLW = auto(); VPMULLD = auto()
+  # packed bitwise
+  VPAND = auto(); VPOR = auto(); VPXOR = auto()
+  # packed variable shifts
+  VPSLLVD = auto(); VPSLLVQ = auto(); VPSRLVD = auto(); VPSRLVQ = auto(); VPSRAVD = auto()
+  # fused multiply add TODO: add other variants to fuse more loads
+  VFMADD213SS = auto(); VFMADD213SD = auto(); VFMADD213PS = auto(); VFMADD213PD = auto()
+  # return
+  RET = auto()
+
+# TODO: add commutative groupop to fuse more loads
+class X86GroupOp:
+  # X86Ops whose first src is also the destination
+  TwoAddress = {X86Ops.ADD, X86Ops.ADDi, X86Ops.AND, X86Ops.ANDi, X86Ops.XOR, X86Ops.XORi, X86Ops.OR, X86Ops.ORi, X86Ops.IMUL,
+                X86Ops.SUB, X86Ops.SUBi, X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR, X86Ops.SARi,
+                X86Ops.IDIV, X86Ops.DIV, X86Ops.VFMADD213SS, X86Ops.VFMADD213SD, X86Ops.VFMADD213PS, X86Ops.VFMADD213PD,
+                X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB}
+
+  # X86Ops whose first src can read from memory
+  ReadMem1st = {X86Ops.MOV, X86Ops.VMOVSS, X86Ops.VMOVSD, X86Ops.VMOVUPS, X86Ops.MOVZX, X86Ops.MOVSX, X86Ops.MOVSXD, X86Ops.VMOVD, X86Ops.VMOVQ,
+                X86Ops.VPMOVZXBW, X86Ops.VPMOVZXBD, X86Ops.VPMOVZXBQ, X86Ops.VPMOVZXWD, X86Ops.VPMOVZXWQ, X86Ops.VPMOVZXDQ,
+                X86Ops.VPMOVSXBW, X86Ops.VPMOVSXBD, X86Ops.VPMOVSXBQ, X86Ops.VPMOVSXWD, X86Ops.VPMOVSXWQ, X86Ops.VPMOVSXDQ,
+                X86Ops.VCVTDQ2PS, X86Ops.VCVTDQ2PD, X86Ops.VCVTTPS2DQ, X86Ops.VCVTTPD2DQ, X86Ops.VCVTTSS2SI, X86Ops.VCVTTSD2SI,
+                X86Ops.VCVTPH2PS, X86Ops.VCVTPS2PD, X86Ops.VCVTPD2PS, X86Ops.VROUNDPS, X86Ops.VROUNDPD, X86Ops.VSQRTPS, X86Ops.VSQRTPD,
+                X86Ops.VPBROADCASTB, X86Ops.VPBROADCASTW, X86Ops.VPBROADCASTD, X86Ops.VPBROADCASTQ, X86Ops.VBROADCASTSS,
+                X86Ops.CMPi, X86Ops.IMULi, X86Ops.LEA}
+
+  # X86Ops whose second src can read from memory NOTE: some of these are TwoAddress so the second src is actually the first
+  ReadMem2nd = {X86Ops.ADD, X86Ops.SUB, X86Ops.AND, X86Ops.OR, X86Ops.XOR, X86Ops.SHL, X86Ops.SHR, X86Ops.SAR, X86Ops.IMUL, X86Ops.CMP,
+                X86Ops.VADDSS, X86Ops.VADDSD, X86Ops.VADDPS, X86Ops.VADDPD, X86Ops.VSUBSS, X86Ops.VSUBSD, X86Ops.VSUBPS, X86Ops.VSUBPD,
+                X86Ops.VMULSS, X86Ops.VMULSD, X86Ops.VMULPS, X86Ops.VMULPD, X86Ops.VDIVSS, X86Ops.VDIVSD, X86Ops.VDIVPS, X86Ops.VDIVPD,
+                X86Ops.VPADDB, X86Ops.VPADDW, X86Ops.VPADDD, X86Ops.VPADDQ, X86Ops.VPSUBB, X86Ops.VPSUBW, X86Ops.VPSUBD, X86Ops.VPSUBQ,
+                X86Ops.VPCMPEQB, X86Ops.VPCMPEQW, X86Ops.VPCMPEQD, X86Ops.VPCMPEQQ, X86Ops.VPBLENDVB, X86Ops.VBLENDVPS, X86Ops.VBLENDVPD,
+                X86Ops.VPCMPGTB, X86Ops.VPCMPGTW, X86Ops.VPCMPGTD, X86Ops.VPCMPGTQ, X86Ops.VCMPSS, X86Ops.VCMPSD, X86Ops.VCMPPS, X86Ops.VCMPPD,
+                X86Ops.VPMULLW, X86Ops.VPMULLD, X86Ops.VROUNDSS, X86Ops.VROUNDSD, X86Ops.VSQRTSS, X86Ops.VSQRTSD, X86Ops.VSHUFPS, X86Ops.VINSERTPS,
+                X86Ops.VPINSRB, X86Ops.VPINSRW, X86Ops.VPINSRD, X86Ops.VPINSRQ, X86Ops.VPAND, X86Ops.VPOR, X86Ops.VPXOR, X86Ops.VPSLLVD,
+                X86Ops.VPSLLVQ, X86Ops.VPSRLVD, X86Ops.VPSRLVQ, X86Ops.VPSRAVD, X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB,
+                X86Ops.VMAXSS, X86Ops.VMAXSD, X86Ops.VMAXPS, X86Ops.VMAXPD, X86Ops.VMINSS, X86Ops.VMINSD, X86Ops.VMINPS, X86Ops.VMINPD,
+                X86Ops.VCVTSI2SS, X86Ops.VCVTSI2SD, X86Ops.VCVTSS2SD, X86Ops.VCVTSD2SS, X86Ops.VUCOMISS, X86Ops.VUCOMISD, X86Ops.IDIV, X86Ops.DIV,
+                X86Ops.VSHUFPD}
+
+  # X86Ops whose third src can read from memory NOTE: these are TwoAddress so the third src is actually the second
+  ReadMem3rd = {X86Ops.VFMADD213SS, X86Ops.VFMADD213SD, X86Ops.VFMADD213PS, X86Ops.VFMADD213PD}
+
+  # X86Ops that can write to memory
+  WriteMem = {X86Ops.MOVm, X86Ops.MOVi, X86Ops.VMOVSSm, X86Ops.VMOVSDm, X86Ops.VMOVUPSm, X86Ops.VMOVDm, X86Ops.VMOVQm,
+              X86Ops.ADDi, X86Ops.SUBi, X86Ops.ANDi, X86Ops.ORi, X86Ops.XORi, X86Ops.SHLi, X86Ops.SHRi, X86Ops.SARi, X86Ops.SETNE,
+              X86Ops.SETE, X86Ops.SETL, X86Ops.SETB, X86Ops.VCVTPS2PH, X86Ops.VPEXTRB, X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VPEXTRQ}
+
+  # X86Ops that read flags
+  ReadFlags = {X86Ops.CMOVB, X86Ops.CMOVL, X86Ops.CMOVE, X86Ops.CMOVNE, X86Ops.SETB, X86Ops.SETL, X86Ops.SETE, X86Ops.SETNE, X86Ops.JB, X86Ops.JL,
+               X86Ops.JE, X86Ops.JNE, X86Ops.JGE}
+
+  # X86Ops that write flags or can modify flags to undefined values
+  WriteFlags = {X86Ops.CMP, X86Ops.CMPi, X86Ops.ADD, X86Ops.ADDi, X86Ops.SUB, X86Ops.SUBi, X86Ops.IMUL, X86Ops.IMULi, X86Ops.IDIV, X86Ops.DIV,
+                X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR, X86Ops.SARi, X86Ops.AND, X86Ops.ANDi, X86Ops.XOR, X86Ops.XORi,
+                X86Ops.OR, X86Ops.ORi, X86Ops.VUCOMISS, X86Ops.VUCOMISD}
+
+  # X86Ops whose first src is the rm field
+  Rm1st = ReadMem1st | (ReadMem2nd & TwoAddress) | {X86Ops.VPSRLDQ}
+
+  # X86Ops whose second src is the rm field
+  Rm2nd = ReadMem2nd | (ReadMem3rd & TwoAddress)
+
+  All = set(X86Ops)
+
+# ***** X86 legalization *****
+
+extra_matcher = PatternMatcher([
+  # bool CMPNE is XOR, bool CMPEQ is XOR+XOR, bool CMPLT is XOR+AND
+  # TODO: how does this work for vector dtypes?
+  (UPat.var('x', dtypes.bool).ne(UPat.var('y')), lambda x,y: x^y),
+  (UPat.var('x', dtypes.bool).alu(Ops.CMPEQ, UPat.var('y')), lambda x,y: (x^y)^True),
+  (UPat.var('x', dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
+  # cast to pointer is a noop
+  (UPat.var("y").cast(name="x"), lambda y,x: y if isinstance(x.dtype, PtrDType) or y.dtype == dtypes.void else None),
+  # can't cast from float16 to ints/float64 directly and vice versa
+  (UPat.var("y", dtypes.float16).cast((dtypes.float64,)+dtypes.ints, name="x"), lambda y,x: y.cast(dtypes.float32).cast(x.dtype)),
+  (UPat.var("y", (dtypes.float64,)+dtypes.ints).cast(dtypes.float16, name="x"), lambda y,x: y.cast(dtypes.float32).cast(x.dtype)),
+  # can't cast from float to int8/16 directly and vice versa
+  (UPat.var("y", dtypes.floats).cast(dtypes.int8s+dtypes.int16s, name="x"), lambda y,x: y.cast(dtypes.int32).cast(x.dtype)),
+  (UPat.var("y", (dtypes.bool,)+dtypes.int8s+dtypes.int16s).cast(dtypes.floats, name="x"), lambda y,x: y.cast(dtypes.int32).cast(x.dtype)),
+  # int/float casts only for signed int
+  (UPat.var("y", dtypes.uint32).cast(dtypes.floats, name="x"), lambda y,x: y.cast(dtypes.int64).cast(x.dtype)),
+  # casting uint64 to float requires special handling
+  (UPat.var("y", dtypes.uint64).cast(dtypes.floats, name="x"), lambda y,x:
+   (y >> 1).cast(dtypes.int64).cast(x.dtype) * 2 + (y & 1).cast(dtypes.int64).cast(x.dtype)),
+  # no int8 mul or cmove, cast to int16
+  (UPat.var("a", dtypes.int8s) * UPat.var("b"), lambda a,b: (a.cast(dtypes.int16) * b.cast(dtypes.int16)).cast(a.dtype)),
+  (UPat.var("m").where(UPat.var("a", (dtypes.bool,)+dtypes.int8s), UPat.var("b")),
+   lambda m,a,b: m.where(a.cast(dtypes.int16), b.cast(dtypes.int16)).cast(a.dtype) if a.dtype.count == 1 else None),
+  # float16 alus are done in float32
+  (UPat(GroupOp.ALU, dtypes.float16, name="x"), lambda x: UOp(x.op, dtypes.float.vec(x.dtype.count),
+   tuple(s.cast(dtypes.float) if s.dtype != dtypes.bool else s for s in x.src)).cast(x.dtype)),
+  (UPat(GroupOp.Comparison, src=(UPat.var("a", dtypes.float16), UPat.var("b")), name="x"),
+   lambda x,a,b: UOp(x.op, x.dtype, (a.cast(dtypes.float32), b.cast(dtypes.float32))).cast(x.dtype)),
+  # no cmpne for packed ints, y != x => !(y==x)
+  (UPat(Ops.CMPNE, src=(UPat.var("y", dtypes.ints), UPat.var("x")), name="cmp"),
+   lambda y,x,cmp: UOp(Ops.CMPEQ, cmp.dtype, (y,x))^True if y.dtype.count > 1 else None),
+  # float where expects a mask TODO: handle float64 cmp to float32 where
+  (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats), UPat.var("b")),
+   lambda m,a,b: m.cast(a.dtype).ne(0).where(a, b) if m.src[0].dtype not in dtypes.floats else None),
+  # TODO: do we want this? If yes make it general
+  #(UPat(Ops.VECTORIZE, dtypes.float16, name="x"), lambda x: x.replace(dtype=dtypes.float32.vec(x.dtype.count),
+  #  src=tuple(s.src[0] for s in x.src)).cast(x.dtype) if all(s.op is Ops.CAST for s in x.src) else None),
+  # rewrite -x -> 0 - x
+  (UPat(Ops.NEG, name="x"), lambda x: UOp(Ops.SUB, x.dtype, (x.const_like(0),) + x.src)),
+])
+
+# ***** X86 pre instruction selection *****
+
+# these must be done in a separate matcher because they violate the spec
+pre_isel_matcher = PatternMatcher([
+  # zero extending scalar 32bit int is a noop
+  (UPat.var("y", dtypes.uint32).cast(dtypes.int64s, name="x"), lambda y,x: x.replace(op=Ops.NOOP) if y.dtype.count == 1 else None),
+  # cast between signed and unsigned int is a noop
+  (UPat.var("y", dtypes.ints+(dtypes.bool,)).cast(dtypes.ints, name="x"),
+   lambda y,x: x.replace(op=Ops.NOOP) if x.dtype.itemsize == y.dtype.itemsize else None),
+  # cast to < scalar int is a noop
+  (UPat.var("y", dtypes.ints).cast(dtypes.ints, name="x"),
+   lambda y,x: x.replace(op=Ops.NOOP) if x.dtype.itemsize < y.dtype.itemsize and y.dtype.count == 1 else None),
+  # bitcasts between scalar floats and ints are real, rest are noops
+  (UPat.var("y").bitcast().named("x"), lambda y,x: None if y.dtype in dtypes.floats and x.dtype in dtypes.ints or \
+   y.dtype in dtypes.ints and x.dtype in dtypes.floats else x.replace(op=Ops.NOOP)),
+  # noop of a noop is removed
+  (UPat(Ops.NOOP, src=(UPat(Ops.NOOP),), name="x"), lambda x: x.replace(src=x.src[0].src)),
+  # moving elements of a single register to another without shuffling is a noop
+  (UPat(Ops.VECTORIZE, src=(UPat.var("y"),), allow_any_len=True, name="x"),
+   lambda y,x: UOp(Ops.NOOP, x.dtype, y.src) if all(s.op is Ops.GEP and s.src == y.src and s.arg[0] == i for i,s in enumerate(x.src)) else None),
+  # gated index becomes a conditional move on the index, the load/store are unconditional
+  (UPat.var("base").index(UPat.var("idx"), UPat.var("gate")).load(UPat.var("alt"), name="x"), lambda base,idx,gate,alt,x:
+   gate.where(base.index(idx, ptr=True), (l:=UOp(Ops.DEFINE_LOCAL, base.dtype.base.ptr(x.dtype.count), arg=0)
+              .index(UOp.const(dtypes.int32, 0), ptr=True)).after(l.store(alt))).load(dtype=x.dtype)),
+  (UPat.var("base").index(UPat.var("idx"), UPat.var("gate")).store(UPat.var("val")), lambda base,idx,gate,val:
+   gate.where(base.index(idx, ptr=True), UOp(Ops.DEFINE_LOCAL, base.dtype.base.ptr(val.dtype.count), arg=0)
+              .index(UOp.const(dtypes.int32, 0), ptr=True)).store(val)),
+  # TODO: remove this once we allow all flag producing ops in cmove
+  # if gate in scalar int cmove is not a comparison need to add one to set the flag
+  (UPat.var("m", dtypes.bool).where(UPat.var("a"), UPat.var("b")),
+   lambda m,a,b: m.ne(0).where(a,b) if m.op not in GroupOp.Comparison and a.dtype.count == 1 else None),
+])
+
+# ***** X86 registers *****
+
+RAX = Register("rax", 0)
+RCX = Register("rcx", 1)
+RDX = Register("rdx", 2)
+RBX = Register("rbx", 3)
+RSP = Register("rsp", 4)
+RBP = Register("rbp", 5)
+RSI = Register("rsi", 6)
+RDI = Register("rdi", 7)
+GPR = (RAX, RCX, RDX, RBX, RSP, RBP, RSI, RDI) + tuple(Register(f"r{i}", i) for i in range(8, 16))
+XMM = tuple(Register(f"xmm{i}", i) for i in range(16))
+# gprs you can write to
+WGPR = tuple(r for r in GPR if r != RSP)
+
+CALLEE_SAVED = (RBX, RSP, RBP, GPR[12], GPR[13], GPR[14], GPR[15]) + ((RSI, RDI) + XMM[6:16] if sys.platform == "win32" else ())
+
+reg_strs = {"rax": {4:"eax", 2:"ax", 1:"al"}, "rcx": {4:"ecx", 2:"cx", 1:"cl"}, "rdx": {4:"edx", 2:"dx", 1:"dl"}, "rbx": {4:"ebx", 2:"bx", 1:"bl"},
+        "rsp": {4:"esp", 2:"sp", 1:"spl"}, "rbp": {4:"ebp", 2:"bp", 1:"bpl"}, "rsi": {4:"esi", 2:"si", 1:"sil"}, "rdi": {4:"edi", 2:"di", 1:"dil"},
+        **{f"r{i}": {4:f"r{i}d", 2:f"r{i}w", 1:f"r{i}b"} for i in range(8, 16)}, **{f"xmm{i}": {64:f"zmm{i}", 32:f"ymm{i}"} for i in range(16)}}
+
+# ***** X86 instruction selection *****
+# if the load is used multiple times we don't fold
+def is_foldable_load(ctx:IselContext, x:UOp, s:UOp) -> bool: return s.op is Ops.LOAD and len(ctx.uses[s]) == x.src.count(s) == 1
+def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.GEP else s
+def lane(x:UOp, i:int) -> int: return s.arg[0] if (s:=x.src[i]).op is Ops.GEP else 0
+def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt]
+def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, arg=X86Ops.DEFINE_REG, dtype=dt, tag=None if reg is None else (reg,))
+def imm(dt:DType, v:int) -> UOp: return UOp(Ops.CONST, dt, arg=truncate[dt](v), tag="__x86_imm__")
+def _uop_key(u:UOp): return (u.op, u.dtype, u.arg)
+def to_imm(c:UOp) -> UOp|None:
+  if c.op is not Ops.CONST: return None
+  if c.dtype is dtypes.int64: return imm(dtypes.int32, c.arg) if not c.overflows(dtypes.int32) else None
+  if c.dtype is dtypes.uint64: return imm(dtypes.uint32, c.arg) if not c.overflows(dtypes.uint32) else None
+  if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, c.arg)
+  return None
+def cmp(x:UOp) -> UOp:
+  if x.src[0].dtype is dtypes.float32: return x.ins(X86Ops.VUCOMISS, dtype=dtypes.void)
+  if x.src[0].dtype is dtypes.float64: return x.ins(X86Ops.VUCOMISD, dtype=dtypes.void)
+  return x.ins(X86Ops.CMP, dtype=dtypes.void) if (i:=to_imm(x.src[1])) is None else x.ins(X86Ops.CMPi, dtype=dtypes.void, src=(x.src[0], i))
+def vcmp(x:UOp) -> UOp:
+  v = imm(dtypes.uint8, {Ops.CMPLT: 1, Ops.CMPNE: 4, Ops.CMPEQ: 0}[x.op])
+  if x.dtype.scalar() is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.dtype.count == 1 else X86Ops.VCMPPS, src=x.src + (v,))
+  return x.ins(X86Ops.VCMPSD if x.dtype.count == 1 else X86Ops.VCMPPD, src=x.src + (v,))
+
+# vshufps xmm2, xmm0, xmm1, imm
+# for 128 bit xmm2 selects its lower 2 32 bits from xmm0 and its upper 2 32 bits from xmm1 according to imm
+# for 256 bit ymm2 repeats the shuffle for its upper 128 bits selecting from the upper 128 bits of ymm0 and ymm1
+def vshufps(x:UOp) -> UOp|None:
+  a, b = base(x, 0), base(x, 2)
+  if not (a is base(x, 1) and b is base(x, 3)) or any(lane(x, i) > 3 for i in range(4)): return None
+  if len(x.src) == 8:
+    if not (a is base(x, 4) is base(x, 5) and b is base(x, 6) is base(x, 7)) or any(lane(x, i+4) != lane(x, i)+4 for i in range(4)): return None
+  return x.ins(X86Ops.VSHUFPS, src=(a, b, imm(dtypes.uint8, sum(lane(x, i) << 2*i for i in range(4)))))
+
+# vshufpd xmm2, xmm0, xmm1, imm
+# for 128 bit xmm2 selects its lower 64 bits from xmm0 and its upper 64 bits from xmm1 according to imm
+# for 256 bit ymm2 also selects its upper 128 bits from the upper 128 bits of ymm0 and ymm1 following the same constraint
+def vshufpd(x:UOp) -> UOp|None:
+  a, b = base(x, 0), base(x, 1)
+  if lane(x, 0) > 1 or lane(x, 1) > 1: return None
+  if len(x.src) == 4 and not (a is base(x, 2) and b is base(x, 3) and lane(x, 2) > 1 and lane(x, 3) > 1): return None
+  return x.ins(X86Ops.VSHUFPD, src=(a, b, imm(dtypes.uint8, sum(lane(x, i) << i for i in range(len(x.src))))))
+
+# vinsertps xmm2, xmm0, xmm1, imm
+# inserts any 32 bit element in xmm1 into any position in xmm0 according to immm, result is written to xmm2
+# this is the fallback slow case for when you can't match more a powerful shuffle
+def vinsertps(x:UOp) -> UOp:
+  def _insert(ret:UOp, i:int) -> UOp:
+    s, v = base(x, i), lane(x, i)
+    # moving the 0th element into the 0th position does nothing
+    return s if i == v == 0 else x.ins(X86Ops.VINSERTPS, src=(ret, s, imm(dtypes.uint8, v << 6 | i << 4)))
+  return functools.reduce(_insert, range(len(x.src)), def_reg(x.dtype))
+
+# vpinsq xmm2, xmm0, rax, imm
+# inserts element in rax into any position in xmm0, result is written to xmm2 according to imm
+def vpins(x:UOp) -> UOp:
+  op = {1: X86Ops.VPINSRB, 2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD, 8: X86Ops.VPINSRQ}[x.dtype.scalar().itemsize]
+  return functools.reduce(lambda ret,i: x.ins(op, src=(ret, x.src[i], imm(dtypes.uint8, i))), range(len(x.src)), def_reg(x.dtype))
+
+# vpbroadcastd xmm1, xmm0
+# inserts scalar int in xmm0 into all lanes of xmm1
+def vpbroadcast(ctx:IselContext, x:UOp, y:UOp) -> UOp:
+  n = x.ins({1: X86Ops.VPBROADCASTB, 2: X86Ops.VPBROADCASTW, 4: X86Ops.VPBROADCASTD, 8: X86Ops.VPBROADCASTQ}[y.dtype.itemsize], src=(y,))
+  if is_foldable_load(ctx, n, y): return n
+  # if there isn't a load we can fold we need to move y from gpr to xmm
+  # this is hacky but required because int.vec(1) isn't supported
+  y = y if y.dtype.itemsize > 1 else y.cast(dtypes.int16)
+  return n.replace(src=(y.bitcast({2:dtypes.float16, 4:dtypes.float32, 8:dtypes.float64}[y.dtype.itemsize]),))
+
+# we don't call ctx.vreg on the srcs to avoid duplicates, a rewrite will assign the tuple of valid registers to a vreg
+def idiv(ctx:IselContext, x:UOp) -> UOp:
+  op = X86Ops.DIV if x.dtype in dtypes.uints else X86Ops.IDIV
+  # for >8bit need to zero/sign extend rax to rdx
+  if x.dtype in dtypes.int8s: ext = []
+  elif x.dtype in dtypes.uints: ext = [x.ins(X86Ops.MOVi, src=(imm(min(dtypes.uint32, x.dtype), 0),), tag=(RDX,))]
+  else: ext = [x.ins(X86Ops.SARi, src=(x.src[0], imm(dtypes.uint8, x.dtype.itemsize * 8 - 1)), tag=(RDX,))]
+  # for 8bit need to zero/sign extend al to ah
+  if x.dtype is dtypes.uint8: dividend = UOp(Ops.INS, arg=X86Ops.MOVZX, dtype=dtypes.int16, src=(x.src[0],), tag=(RAX,))
+  elif x.dtype is dtypes.int8: dividend = UOp(Ops.INS, arg=X86Ops.MOVSX, dtype=dtypes.int16, src=(x.src[0],), tag=(RAX,))
+  else: dividend = x.ins(X86Ops.MOV, src=(x.src[0],), tag=(RAX,))
+  # divisor can't be in rax or rdx
+  divisor = x.ins(X86Ops.MOV, src=(x.src[1],), tag=tuple(r for r in WGPR if r not in (RAX, RDX)))
+  # for >8bit both rax and rdx are written to
+  defs = (ctx.vreg(RAX),) if x.dtype in dtypes.int8s else (ctx.vreg(RAX), ctx.vreg(RDX))
+  idiv = x.ins(op, src=(dividend, divisor) + tuple(ext), tag=defs)
+  # this move "cleanses" the register constraint (rax) of idiv as it only applies on definition and not on the uses of idiv
+  return x.ins(X86Ops.MOV, src=(idiv,))
+
+def fold_address(x:UOp) -> tuple[UOp, UOp, UOp]:
+  def _disp(v:int) -> UOp: return imm(dtypes.int32 if abs(v) > dtypes.int8.max else dtypes.int8, v)
+  def _cast(v:UOp) -> UOp: return v.cast(dtypes.int64) if v.vmin < 0 else v
+  if x.op is not Ops.INDEX: return (x, UOp(Ops.NOOP), _disp(0))
+  base, idx = x.src
+  disp_scale = base.dtype.itemsize if isinstance(base.dtype, PtrDType) else 1
+  if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: return (base, _cast(idx.src[0]), _disp(idx.src[1].arg * disp_scale))
+  if idx.op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.arg * disp_scale))
+  return (base, _cast(idx), _disp(0))
+
+def alloc_defs(ctx:IselContext, x:UOp) -> UOp|None:
+  if x.dtype is dtypes.void or isinstance(x.tag, tuple): return None
+  if x.op in {Ops.PARAM, Ops.DEFINE_VAR, Ops.SPECIAL}:
+    i = ctx.func_args.index(x)
+    regs = (RCX, RDX, GPR[8], GPR[9]) if sys.platform == "win32" else (RDI, RSI, RDX, RCX, GPR[8], GPR[9])
+    if i < len(regs): return x.replace(tag=(ctx.vreg(regs[i]),))
+  defs = [ctx.vreg(WGPR)] if x.dtype in dtypes.ints+(dtypes.bool,) or isinstance(x.dtype, PtrDType) else [ctx.vreg(XMM)]
+  return x.replace(tag=tuple(defs))
+
+def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None:
+  # immediates and real registers
+  if x.op is Ops.CONST: return None
+  if x.arg in (X86Ops.FRAME_INDEX, X86Ops.DEFINE_REG) and x.tag is not None: return None
+  # no register definition
+  if x.dtype is dtypes.void: return None
+  # already allocated vregs
+  if isinstance(x.tag, tuple) and x.tag[0]._cons: return None
+  # allocate vreg definitions
+  defs = []
+  if isinstance(x.tag, tuple): defs = [ctx.vreg(x.tag)]
+  elif x.dtype in dtypes.ints+(dtypes.bool,) or isinstance(x.dtype, PtrDType): defs = [ctx.vreg(WGPR)]
+  elif x.dtype in dtypes.floats or x.dtype.count > 1: defs = [ctx.vreg(XMM)]
+  # TODO: add this once the scheduler can track register pressure
+  # if x.arg in X86GroupOp.WriteFlags: defs.append(ctx.vreg(RFLAGS))
+  return x.replace(tag=tuple(defs))
+
+def lower_abi(ctx, x:UOp):
+  i = ctx.func_arg_idxs[_uop_key(x)]
+  if sys.platform == "win32": regs, stack_base = (RCX, RDX, GPR[8], GPR[9]), 32
+  else: regs, stack_base = (RDI, RSI, RDX, RCX, GPR[8], GPR[9]), 0
+  if i < len(regs):
+    src = def_reg(x.dtype, regs[i])
+    if x.reg == src.reg: return src, [src]
+    return (nx:=x.ins(X86Ops.MOV, src=(src,)), [src, nx])
+  fi = UOp(Ops.INS, arg=X86Ops.FRAME_INDEX, dtype=dtypes.int32, tag=(i-len(regs)+1)*8+stack_base)
+  nx = x.ins(X86Ops.MOV, src=(def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), fi))
+  return nx, [fi, nx]
+
+def lower_stack_define(ctx, x:UOp):
+  disp = imm(dtypes.int32, ctx.local_offsets[_uop_key(x)])
+  nx = x.ins(X86Ops.LEA, src=(def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), disp))
+  return nx, [disp, nx]
+
+dts = dtypes.ints + (dtypes.bool, dtypes.float16, dtypes.float32, dtypes.float64)
+dt_16bit = tuple(dt.vec(l) for dt in dts for l in [2,1] if l*dt.itemsize == 2 and dt not in dtypes.int16s)
+dt_32bit = tuple(dt.vec(l) for dt in dts for l in [4,2,1] if l*dt.itemsize == 4 and dt not in dtypes.int32s)
+dt_64bit = tuple(dt.vec(l) for dt in dts for l in [8,4,2,1] if l*dt.itemsize == 8 and dt not in dtypes.int64s)
+dt_128bit = tuple(dt.vec(l) for dt in dts for l in [16,8,4,2,1] if l*dt.itemsize == 16)
+
+isel_matcher = PatternMatcher([
+  # **** Op -> Op ****
+  # float gep(0) is a noop as it just moves the 0th element from one xmm register to another
+  # this is done here to not interfere with shuffles
+  (UPat(dtype=dtypes.floats).gep(0, name="x"), lambda x: x.replace(op=Ops.NOOP, arg=None)),
+  # range is lowered to acc, cmp, jmp after regalloc
+  (UPat(Ops.RANGE, src=(UPat.cvar("c"),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(c.dtype, c.arg),) + x.src[1:])),
+  (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None),
+  # **** Op -> X86Op ****
+  # append return, callee-saved live ranges are inserted in regalloc
+  (UPat(Ops.SINK, name="x"), lambda x:
+   x.replace(src=(x.ins(X86Ops.RET, src=x.src),))
+     if not (len(x.src) == 1 and x.src[0].op is Ops.INS and x.src[0].arg is X86Ops.RET) else None),
+  # late lowered function args and stack backed locals still need virtual registers
+  (UPat((Ops.PARAM, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.DEFINE_REG, Ops.DEFINE_LOCAL), name="x"), alloc_defs),
+  # constants that can't be immediates, move them to registers
+  (UPat.cvar("x", dtypes.int64s), lambda x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, x.arg),)) if x.tag is None else None),
+  (UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, x.arg),)) if x.tag is None else None),
+  (UPat.cvar("x", dtypes.floats), lambda x:
+   UOp.const(dt:=to_int(x.dtype), struct.unpack(dt.fmt, struct.pack(x.dtype.fmt, x.arg))[0]).bitcast(x.dtype) if x.tag is None else None),
+  # TODO: these should use a.maximum(b) / a.minimum(b)
+  ((UPat.var("a") < UPat.var("b")).where(UPat.var("b", dtypes.float32), UPat.var("a")), lambda a,b:
+   a.ins(X86Ops.VMAXSS if a.dtype.count == 1 else X86Ops.VMAXPS, src=(a, b))),
+  ((UPat.var("a") < UPat.var("b")).where(UPat.var("b", dtypes.float64), UPat.var("a")), lambda a,b:
+   a.ins(X86Ops.VMAXSD if a.dtype.count == 1 else X86Ops.VMAXPD, src=(a, b))),
+  ((UPat.var("a") < UPat.var("b")).where(UPat.var("a", dtypes.float32), UPat.var("b")), lambda a,b:
+   a.ins(X86Ops.VMINSS if a.dtype.count == 1 else X86Ops.VMINPS, src=(a, b))),
+  ((UPat.var("a") < UPat.var("b")).where(UPat.var("a", dtypes.float64), UPat.var("b")), lambda a,b:
+   a.ins(X86Ops.VMINSD if a.dtype.count == 1 else X86Ops.VMINPD, src=(a, b))),
+  # conditional moves that use masks NOTE: these currently assume a mask producing cmp exists
+  (UPat.var("m").where(UPat.var("a", dtypes.ints), UPat.var("b")), lambda m,a,b:
+   a.ins(X86Ops.VPBLENDVB, src=(b, a, m.replace(dtype=m.src[0].dtype))) if a.dtype.count > 1 else None),
+  (UPat.var("m").where(UPat.var("a", dtypes.float32), UPat.var("b")), lambda m,a,b:
+   a.ins(X86Ops.VBLENDVPS, src=(b, a, m.replace(dtype=m.src[0].dtype)))),
+  (UPat.var("m").where(UPat.var("a", dtypes.float64), UPat.var("b")), lambda m,a,b:
+   a.ins(X86Ops.VBLENDVPD, src=(b, a, m.replace(dtype=m.src[0].dtype)))),
+  # in this case we have a mask producing comparison whose user expects a bool, so we convert to bool
+  (UPat(GroupOp.Comparison, dtypes.bool, (UPat.var("y", (dtypes.float32, dtypes.float64)), UPat()), name="x"), lambda y,x:
+   x.replace(dtype=y.dtype).bitcast(to_int(y.dtype)).bitwise_and(1).f(Ops.NOOP, dtype=dtypes.bool)),
+  # conditional moves that use flags
+  (UPat(Ops.CMPLT, src=(UPat(dtype=dtypes.sints), UPat()), name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b:
+   a.ins(X86Ops.CMOVL, src=(b, a, cmp(m)))),
+  (UPat(Ops.CMPLT, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.CMOVB, src=(b, a, cmp(m)))),
+  (UPat(Ops.CMPEQ, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.CMOVE, src=(b, a, cmp(m)))),
+  (UPat(Ops.CMPNE, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.CMOVNE, src=(b, a, cmp(m)))),
+  # jumps, use flags
+  (UPat(Ops.IF, src=(UPat(Ops.CMPLT, src=(UPat(dtype=dtypes.uints), UPat()), name="y"),), name="x"), lambda y,x: x.ins(X86Ops.JB, src=(cmp(y),))),
+  (UPat(Ops.IF, src=(UPat(Ops.CMPLT, name="y"),), name="x"), lambda y,x: x.ins(X86Ops.JL, src=(cmp(y),))),
+  (UPat(Ops.IF, src=(UPat(Ops.CMPEQ, name="y"),), name="x"), lambda y,x: x.ins(X86Ops.JE, src=(cmp(y),))),
+  (UPat(Ops.IF, src=(UPat(Ops.CMPNE, name="y"),), name="x"), lambda y,x: x.ins(X86Ops.JNE, src=(cmp(y),))),
+  # comparisons whose user doesn't use the flag, move flag result to register
+  (UPat(Ops.CMPLT, dtypes.bool, (UPat(dtype=dtypes.uints), UPat()), name="x"), lambda x: x.ins(X86Ops.SETB, src=(cmp(x),))),
+  (UPat(Ops.CMPLT, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETL, src=(cmp(x),))),
+  (UPat(Ops.CMPEQ, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETE, src=(cmp(x),))),
+  (UPat(Ops.CMPNE, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETNE, src=(cmp(x),))),
+  # comparisons that produce masks (these aren't bool dtype)
+  (UPat(GroupOp.Comparison, src=(UPat(dtype=(dtypes.float32, dtypes.float64)), UPat()), name="x"), vcmp),
+  (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int8s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQB)),
+  (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int16s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQW)),
+  (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int32s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQD)),
+  (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int64s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQQ)),
+  (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int8s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTB, src=(b, a))),
+  (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int16s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTW, src=(b, a))),
+  (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int32s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTD, src=(b, a))),
+  (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int64s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTQ, src=(b, a))),
+  # float unary
+  (UPat.var("y", dtypes.float32).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSS, src=(y, y)) if x.dtype.count == 1 else x.ins(X86Ops.VSQRTPS)),
+  (UPat.var("y", dtypes.float64).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSD, src=(y, y)) if x.dtype.count == 1 else x.ins(X86Ops.VSQRTPD)),
+  (UPat.var("y", dtypes.float32).trunc().named("x"), lambda y,x:
+   x.ins(X86Ops.VROUNDSS, src=(y, y, imm(dtypes.uint8, 3))) if x.dtype.count == 1 else x.ins(X86Ops.VROUNDPS, src=(y, imm(dtypes.uint8, 3)))),
+  (UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x:
+   x.ins(X86Ops.VROUNDSD, src=(y, y, imm(dtypes.uint8, 3))) if x.dtype.count == 1 else x.ins(X86Ops.VROUNDPD, src=(y, imm(dtypes.uint8, 3)))),
+  # shufles
+  (UPat.var("y", dtypes.float32).broadcast(name="x"), lambda y,x: x.ins(X86Ops.VBROADCASTSS, src=(y,))),
+  # for float16 we route the srcs through gprs unless we can fold them, this is suboptimal for values in xmms, in that case we want vpunpcklwd
+  (UPat(Ops.VECTORIZE, dtypes.float16, name="x"), lambda ctx,x:
+   vpins(x.replace(src=tuple(s if is_foldable_load(ctx, x, s) else s.bitcast(dtypes.int16) for s in x.src)))),
+  (UPat(Ops.VECTORIZE, (dtypes.float32.vec(4), dtypes.float32.vec(8)), name="x"), vshufps),
+  (UPat(Ops.VECTORIZE, (dtypes.float64.vec(2), dtypes.float64.vec(4)), name="x"), vshufpd),
+  (UPat(Ops.VECTORIZE, dtypes.float32, name="x"), vinsertps),
+  (UPat.var("y", dtypes.ints+(dtypes.bool,)).broadcast(name="x"), vpbroadcast),
+  (UPat(Ops.VECTORIZE, dtypes.ints+(dtypes.bool,), name="x"), vpins),
+  # gep
+  (UPat.var("y", dtypes.int8s+(dtypes.bool,)).gep(name="x"), lambda y,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, x.arg[0])))),
+  (UPat.var("y", dtypes.int16s).gep(name="x"), lambda y,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, x.arg[0])))),
+  (UPat.var("y", dtypes.int32s).gep(name="x"), lambda y,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, x.arg[0])))),
+  (UPat.var("y", dtypes.int64s).gep(name="x"), lambda y,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, x.arg[0])))),
+  (UPat.var("y", dtypes.floats).gep(name="x"), lambda y,x: x.ins(X86Ops.VPSRLDQ, src=(y, imm(dtypes.uint8, x.arg[0] * x.dtype.itemsize)))),
+  # fused multiply add TODO: don't fuse if mul used several times
+  (UPat.var('a', dtypes.float32) * UPat.var('b') + UPat.var('c'), lambda a,b,c:
+   a.ins(X86Ops.VFMADD213SS if a.dtype.count == 1 else X86Ops.VFMADD213PS, src=(a, b, c))),
+  (UPat.var('a', dtypes.float64) * UPat.var('b') + UPat.var('c'), lambda a,b,c:
+   a.ins(X86Ops.VFMADD213SD if a.dtype.count == 1 else X86Ops.VFMADD213PD, src=(a, b, c))),
+  # packed bitwise
+  ((UPat() & UPat()).named("x"), lambda x: x.ins(X86Ops.VPAND) if x.dtype.count > 1 else None),
+  ((UPat() | UPat()).named("x"), lambda x: x.ins(X86Ops.VPOR) if x.dtype.count > 1 else None),
+  ((UPat() ^ UPat()).named("x"), lambda x: x.ins(X86Ops.VPXOR) if x.dtype.count > 1 else None),
+  # packed int binary
+  ((UPat(dtype=dtypes.int32s) << UPat()).named("x"), lambda x: x.ins(X86Ops.VPSLLVD) if x.dtype.count > 1 else None),
+  ((UPat(dtype=dtypes.int64s) << UPat()).named("x"), lambda x: x.ins(X86Ops.VPSLLVQ) if x.dtype.count > 1 else None),
+  ((UPat(dtype=dtypes.uint32) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRLVD) if x.dtype.count > 1 else None),
+  ((UPat(dtype=dtypes.uint64) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRLVQ) if x.dtype.count > 1 else None),
+  ((UPat(dtype=dtypes.int32) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRAVD) if x.dtype.count > 1 else None),
+  ((UPat(dtype=dtypes.int8s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDB) if x.dtype.count > 1 else None),
+  ((UPat(dtype=dtypes.int16s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDW) if x.dtype.count > 1 else None),
+  ((UPat(dtype=dtypes.int32s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDD) if x.dtype.count > 1 else None),
+  ((UPat(dtype=dtypes.int64s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDQ) if x.dtype.count > 1 else None),
+  (UPat(Ops.SUB, dtypes.int8s, name="x"), lambda x: x.ins(X86Ops.VPSUBB) if x.dtype.count > 1 else None),
+  (UPat(Ops.SUB, dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPSUBW) if x.dtype.count > 1 else None),
+  (UPat(Ops.SUB, dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPSUBD) if x.dtype.count > 1 else None),
+  (UPat(Ops.SUB, dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPSUBQ) if x.dtype.count > 1 else None),
+  (UPat(Ops.MUL, dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMULLW) if x.dtype.count > 1 else None),
+  (UPat(Ops.MUL, dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMULLD) if x.dtype.count > 1 else None),
+  # scalar int binary
+  ((UPat(dtype=dtypes.ints) // UPat()).named("x"), idiv),
+  # scalar int binary with immediate
+  (UPat.var("a", dtypes.ints) << UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.arg)))),
+  (UPat.var("a", dtypes.uints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.arg)))),
+  (UPat.var("a", dtypes.sints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.arg)))),
+  (UPat.var("a", dtypes.ints) + UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ADDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints) * UPat.cvar("c"), lambda a,c: a.ins(X86Ops.IMULi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ANDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints+(dtypes.bool,)) | UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.cvar("c"), lambda a,c: a.ins(X86Ops.XORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.cvar("c"))), lambda a,c: a.ins(X86Ops.SUBi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  # scalar int binary with register
+  (UPat.var("a", dtypes.ints) << UPat.var("b"), lambda a,b: a.ins(X86Ops.SHL, src=(a, b))),
+  (UPat.var("a", dtypes.uints) >> UPat.var("b"), lambda a,b: a.ins(X86Ops.SHR, src=(a, b))),
+  (UPat.var("a", dtypes.sints) >> UPat.var("b"), lambda a,b: a.ins(X86Ops.SAR, src=(a, b))),
+  (UPat.var("a", dtypes.ints) + UPat.var("b"), lambda a,b: a.ins(X86Ops.ADD, src=(a, b))),
+  (UPat.var("a", dtypes.ints) * UPat.var("b"), lambda a,b: a.ins(X86Ops.IMUL, src=(a, b))),
+  (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.var("b"), lambda a,b: a.ins(X86Ops.AND, src=(a, b))),
+  (UPat.var("a", dtypes.ints+(dtypes.bool,)) | UPat.var("b"), lambda a,b: a.ins(X86Ops.OR, src=(a, b))),
+  (UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.var("b"), lambda a,b: a.ins(X86Ops.XOR, src=(a, b))),
+  (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.var("b"))), lambda a,b: a.ins(X86Ops.SUB, src=(a, b))),
+  # float binary
+  ((UPat(dtype=dtypes.float32) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSS if x.dtype.count == 1 else X86Ops.VADDPS)),
+  ((UPat(dtype=dtypes.float64) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSD if x.dtype.count == 1 else X86Ops.VADDPD)),
+  ((UPat(dtype=dtypes.float32) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSS if x.dtype.count == 1 else X86Ops.VMULPS)),
+  ((UPat(dtype=dtypes.float64) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSD if x.dtype.count == 1 else X86Ops.VMULPD)),
+  (UPat(Ops.SUB, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VSUBSS if x.dtype.count == 1 else X86Ops.VSUBPS)),
+  (UPat(Ops.SUB, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VSUBSD if x.dtype.count == 1 else X86Ops.VSUBPD)),
+  (UPat(Ops.FDIV, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VDIVSS if x.dtype.count == 1 else X86Ops.VDIVPS)),
+  (UPat(Ops.FDIV, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VDIVSD if x.dtype.count == 1 else X86Ops.VDIVPD)),
+  # casts
+  (UPat(dtype=dtypes.int32).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTDQ2PS) if x.dtype.count > 1 else None),
+  (UPat(dtype=dtypes.int32).cast(dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VCVTDQ2PD) if x.dtype.count > 1 else None),
+  (UPat(dtype=dtypes.float32).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VCVTTPS2DQ) if x.dtype.count > 1 else None),
+  (UPat(dtype=dtypes.float64).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VCVTTPD2DQ) if x.dtype.count > 1 else None),
+  (UPat(dtype=dtypes.float32).cast(dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VCVTPS2PD) if x.dtype.count > 1 else None),
+  (UPat(dtype=dtypes.float64).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTPD2PS) if x.dtype.count > 1 else None),
+  (UPat(dtype=dtypes.float32).cast(dtypes.float16, name="x"), lambda x: x.ins(X86Ops.VCVTPS2PH, src=x.src + (imm(dtypes.uint8, 4),))),
+  (UPat(dtype=dtypes.float16).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTPH2PS)),
+  (UPat(dtype=dtypes.float32).cast(dtypes.int32s+dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VCVTTSS2SI)),
+  (UPat(dtype=dtypes.float64).cast(dtypes.int32s+dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VCVTTSD2SI)),
+  (UPat.var("y", dtypes.float32).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSS2SD, src=(y, y))),
+  (UPat.var("y", dtypes.float64).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSD2SS, src=(y, y))),
+  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SS, src=(def_reg(x.dtype), y))),
+  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SD, src=(def_reg(x.dtype), y))),
+  (UPat(dtype=dtypes.uints+(dtypes.bool,)).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOVZX) if x.dtype.count == 1 else None),
+  (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.MOVSXD) if x.dtype.count == 1 else None),
+  (UPat(dtype=dtypes.sints).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOVSX) if x.dtype.count == 1 else None),
+  (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBW)),
+  (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBD)),
+  (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBQ)),
+  (UPat(dtype=dtypes.uint16).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXWD)),
+  (UPat(dtype=dtypes.uint16).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXWQ)),
+  (UPat(dtype=dtypes.uint32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXDQ)),
+  (UPat(dtype=dtypes.int8).cast(dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBW)),
+  (UPat(dtype=dtypes.int8).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBD)),
+  (UPat(dtype=dtypes.int8).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBQ)),
+  (UPat(dtype=dtypes.int16).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXWD)),
+  (UPat(dtype=dtypes.int16).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXWQ)),
+  (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXDQ)),
+  # bitcasts
+  (UPat.var("y", dtypes.float16).bitcast(dtypes.int16s).named("x"), lambda y,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, 0)))),
+  (UPat(dtype=dtypes.int16s).bitcast(dtypes.float16).named("x"), vpins),
+  (UPat(dtype=dtypes.int32s).bitcast(dtypes.float32).named("x"), lambda x: x.ins(X86Ops.VMOVD)),
+  (UPat(dtype=dtypes.int64s).bitcast(dtypes.float64).named("x"), lambda x: x.ins(X86Ops.VMOVQ)),
+  (UPat(dtype=dtypes.float32).bitcast(dtypes.int32s).named("x"), lambda x: x.ins(X86Ops.VMOVDm)),
+  (UPat(dtype=dtypes.float64).bitcast(dtypes.int64s).named("x"), lambda x: x.ins(X86Ops.VMOVQm)),
+  # index
+  (UPat(Ops.INDEX, name="x"), lambda x: x.ins(X86Ops.LEA, src=fold_address(x))),
+  # TODO: fuse stores, very few cases -- store cmp becomes setcc, store gep int becomes vpextr, store bitcast to int becomes vmovd/q
+  # copy, load, store
+  # NOTE: copy here violates the spec, it only happens post register allocation when a reg to reg move needs to be inserted
+  (UPat(Ops.COPY, dt_128bit, name="x"), lambda x: x.ins(X86Ops.VMOVUPS)),
+  (UPat(Ops.COPY, dt_64bit, name="x"), lambda x: x.ins(X86Ops.VMOVSD)),
+  (UPat(Ops.COPY, dt_32bit+dt_16bit, name="x"), lambda x: x.ins(X86Ops.VMOVSS)),
+  (UPat(Ops.COPY, dtypes.ints+(dtypes.bool,), name="x"), lambda x: x.ins(X86Ops.MOV)),
+  (UPat(Ops.LOAD, dt_128bit, name="x"), lambda x: x.ins(X86Ops.VMOVUPS, src=fold_address(x.src[0]))),
+  (UPat(Ops.LOAD, dt_64bit, name="x"), lambda x: x.ins(X86Ops.VMOVSD, src=fold_address(x.src[0]))),
+  (UPat(Ops.LOAD, dt_32bit, name="x"), lambda x: x.ins(X86Ops.VMOVSS, src=fold_address(x.src[0]))),
+  (UPat(Ops.LOAD, dt_16bit, name="x"), lambda x:
+   x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(x.src[0]) + (imm(dtypes.uint8, 0),))),
+  (UPat(Ops.LOAD, dtypes.ints+(dtypes.bool,), name="x"), lambda x: x.ins(X86Ops.MOV, src=fold_address(x.src[0]))),
+  (UPat.var("a").store(UPat.var("b", dt_128bit), name="x"), lambda a,b,x: x.ins(X86Ops.VMOVUPSm, src=fold_address(a) + (b,))),
+  (UPat.var("a").store(UPat.var("b", dt_64bit), name="x"), lambda a,b,x: x.ins(X86Ops.VMOVSDm, src=fold_address(a) + (b,))),
+  (UPat.var("a").store(UPat.var("b", dt_32bit), name="x"), lambda a,b,x: x.ins(X86Ops.VMOVSSm, src=fold_address(a) + (b,))),
+  (UPat.var("a").store(UPat.var("b", dt_16bit), name="x"), lambda a,b,x: x.ins(X86Ops.VPEXTRW, src=fold_address(a) + (b, imm(dtypes.uint8, 0)))),
+  (UPat.var("a").store(UPat.var("b", dtypes.ints+(dtypes.bool,)), name="x"), lambda a,b,x:
+   x.ins(X86Ops.MOVm, src=fold_address(a) + (b,)) if (i:=to_imm(b)) is None else x.ins(X86Ops.MOVi, src=fold_address(a) + (i,))),
+  # **** X86Op -> X86Op ****
+  # fold loads into X86Ops that allow it, if beneficial
+  (UPat(Ops.INS, src=(UPat(Ops.LOAD, name="y"),), allow_any_len=True, name="x"), lambda ctx,y,x:
+   x.replace(src=fold_address(y.src[0]) + x.src[1:]) if x.arg in X86GroupOp.ReadMem1st and is_foldable_load(ctx, x, y) else None),
+  (UPat(Ops.INS, src=(UPat(), UPat(Ops.LOAD, name="y")), allow_any_len=True, name="x"), lambda ctx,y,x:
+   x.replace(src=x.src[:1] + fold_address(y.src[0]) + x.src[2:]) if x.arg in X86GroupOp.ReadMem2nd and is_foldable_load(ctx, x, y) else None),
+  (UPat(Ops.INS, src=(UPat(), UPat(), UPat(Ops.LOAD, name="y")), allow_any_len=True, name="x"), lambda ctx,y,x:
+   x.replace(src=x.src[:2] + fold_address(y.src[0]) + x.src[3:]) if x.arg in X86GroupOp.ReadMem3rd and is_foldable_load(ctx, x, y) else None),
+  # allocate virtual registers
+  (UPat(Ops.INS, name="x"), alloc_vregs),
+])
+
+# ***** pre register allocation *****
+# this handles flag clobbers. Unfortunately x86 doesn't have a good way to store/restore the flag register (then regalloc would handle it)
+# so we rematerialize. This is different from rematerialization you might want to do in regalloc because it is not optional,
+# regalloc shouldn't rematerialize if a src of the instruction is dead, but here you need to as there's no fallback load from stack
+def flag_rematerialize(ctx:PreRegAllocContext, x:UOp):
+  flag_def = x if x.arg in X86GroupOp.WriteFlags or x.op is Ops.RANGE else x.src[-1] if x.arg in X86GroupOp.ReadFlags else None
+  if flag_def is None: return None
+  if ctx.lock is not None and ctx.lock is not flag_def: ctx.clobbered.add(ctx.lock)
+  ctx.lock = flag_def
+  if flag_def not in ctx.clobbered: return None
+  ctx.clobbered.remove(flag_def)
+  return (x, [flag_def, x])
+
+pre_regalloc_matcher = PatternMatcher([
+  (UPat((Ops.INS, Ops.RANGE), name="x"), flag_rematerialize),
+])
+
+late_regalloc_matcher = PatternMatcher([
+  (UPat((Ops.PARAM, Ops.DEFINE_VAR, Ops.SPECIAL), name="x"), lower_abi),
+  (UPat((Ops.DEFINE_REG, Ops.DEFINE_LOCAL), name="x"), lower_stack_define),
+])
+
+# ***** post register allocation *****
+# TODO: control flow should be overhauled so that this isn't necessary
+def lower_range(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
+  loop_label = "_".join(str(i) for i in x.arg[:-1])
+  acc = x.ins(X86Ops.MOVi, src=(imm(x.dtype, 0),) + x.src[1:])
+  label = UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_{loop_label}")
+  cmp = UOp(Ops.INS, arg=X86Ops.CMPi if x.src[0].op is Ops.CONST else X86Ops.CMP, src=(acc, x.src[0]))
+  jump_out = UOp(Ops.INS, arg=X86Ops.JGE, src=(cmp,), tag=f".LOOP_OUT_{loop_label}")
+  ctx.loop_label[acc] = loop_label
+  return (acc, [acc, label, cmp, jump_out])
+
+# final rewrite to match the isa spec
+post_regalloc_matcher = PatternMatcher([
+  # alloc stack space
+  (UPat(Ops.INS, arg=X86Ops.DEFINE_REG, dtype=dtypes.uint64, name="x"), lambda ctx,x:
+   (x, [x, x.ins(X86Ops.SUBi, src=(imm(dtypes.uint32, ctx.stack_size),), tag=(RSP,))]) if ctx.stack_size > 0 and x.reg is RSP else None),
+  # dealloc stack space
+  (UPat(Ops.INS, arg=X86Ops.RET, name="x"), lambda ctx,x: (x, [UOp(Ops.INS, arg=X86Ops.ADDi, dtype=dtypes.uint64,
+                      src=(imm(dtypes.uint32, ctx.stack_size),), tag=(RSP,)), x]) if ctx.stack_size > 0 else None),
+  # rewrite FRAME_INDEX to CONST now that the stack size is known
+  (UPat(Ops.INS, arg=X86Ops.FRAME_INDEX, name="x"), lambda ctx,x: (nx:=UOp.const(x.dtype, ctx.stack_size + x.tag), [nx])),
+  # rewrite RANGE to ACC = 0 -> LABEL -> JUMP if ACC >= loop bound
+  (UPat(Ops.RANGE, name="x"), lambda ctx,x: lower_range(ctx, x)),
+  # rewrite END to ACC + 1 -> JUMP -> LABEL, also add the out of loop JUMP to the src so this becomes the jump target
+  (UPat(Ops.END, name="x"), lambda ctx,x: (jmp:=UOp(Ops.INS, arg=X86Ops.JMP, tag=f".LOOP_{ctx.loop_label[x.src[1]]}"),
+   [x.src[1].ins(X86Ops.ADDi, src=(imm(x.src[1].dtype, 1),)), jmp, UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_OUT_{ctx.loop_label[x.src[1]]}")])),
+  # rewrite two address instructions to two address form, if reused src wasn't coalesced insert a move
+  (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=x.replace(src=x.src[1:]),
+   [ctx.ren.copy(x.src[0], x.reg), nx] if x.reg != x.src[0].reg else [nx]) if x.arg in X86GroupOp.TwoAddress else None),
+])
+
+# ***** X86 spec *****
+# TODO: do we even want this?
+isa_spec = PatternMatcher([
+  # these are the only non X86Ops allowed
+  (UPat(Ops.CONST), lambda: True),
+  (UPat((Ops.NOOP, Ops.GROUP, Ops.AFTER, Ops.BARRIER, Ops.SINK)), lambda: True),
+  (UPat(Ops.INS, name="x"), lambda x: x.arg in X86GroupOp.All),
+])
+
+# ***** X86 instruction encoding *****
+
+def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) -> bytes|None:
+  def _encode(reg_uop:UOp|None, rm_uop:UOp, idx_uop:UOp|None=None, disp_uop:UOp|None=None, vvvv_uop:UOp|None=None, imm_uop:UOp|None=None) -> bytes:
+    nonlocal reg, opc
+    # get the encoding values of the different fields
+    reg = cast(int, cast(Register, reg_uop.reg).index if reg_uop is not None else reg)
+    rm = cast(Register, rm_uop.reg).index
+    idx = cast(Register, idx_uop.reg).index if idx_uop is not None and idx_uop.reg is not None else 4
+    rm_sz = 8 if isinstance(rm_uop.dtype, PtrDType) and disp_uop is None else rm_uop.dtype.itemsize
+    reg_sz = (reg_uop.dtype.itemsize if not isinstance(reg_uop.dtype, PtrDType) else 8) if reg_uop is not None else 0
+    sz = reg_sz or rm_sz
+
+    # encode instruction
+    inst = bytes([])
+    assert 0 <= reg <= 15 and 0 <= idx <= 15 and 0 <= rm <= 15
+    # r extends reg field, x extends index field, b extends rm or base field
+    r, _x, b = reg >> 3, idx >> 3, rm >> 3
+    if sel: # VEX bytes
+      vvvv = cast(Register, vvvv_uop.reg).index if vvvv_uop is not None else 0
+      l = (max(reg_sz, rm_sz) > 16) & 0b1
+      if sel == 1 and _x == b == we == 0: inst += bytes([0xC5, (~r & 0b1) << 7 | (~vvvv & 0b1111) << 3 | l << 2 | pp])
+      else: inst += bytes([0xC4, (~r & 0b1) << 7 | (~_x & 0b1) << 6 | (~b & 0b1) << 5 | sel, we << 7 | (~vvvv & 0b1111) << 3 | l << 2 | pp])
+    else: # optional PREFIX and REX bytes
+      # PREFIX byte signaling 16 bit variant of instruction
+      if sz == 2: inst += bytes([0x66])
+      # bit signaling 64 bit variant of instruction
+      w = sz == 8
+      # REX byte is required when 64 bit or an extended reg is used (index 8 - 15) or lower 8 bits of (rsp, rbp, rsi, rdi) are accessed
+      if w | r | _x | b | (reg_sz == 1 & reg >> 2) | (rm_sz == 1 & rm >> 2): inst += bytes([0b0100 << 4 | w << 3 | r << 2 | _x << 1 | b])
+      # legacy 8bit opcode is 1 less than 16-64bit variants
+      if (rm_sz == 1 or reg_sz == 1) and x.arg not in X86GroupOp.ReadFlags | {X86Ops.LEA}: opc -= 1
+    # OPCODE byte
+    inst += opc.to_bytes((opc.bit_length() + 7) // 8, 'big')
+    # MODRM byte
+    # now we only care about the lower 3 bits
+    idx, rm, reg = idx & 0b111, rm & 0b111, reg & 0b111
+    # 0b00 -- signals memory access with no displacement
+    # 0b01 -- signals memory access with 8bit displacement
+    # 0b10 -- signals memory access with 32bit displacement
+    # 0b11 -- signals no memory access
+    if disp_uop is not None:
+      assert disp_uop.dtype in (dtypes.int8, dtypes.int32), "displacement can only be 1 or 4 byte signed int"
+      # rbp/r13 always require a displacement
+      if disp_uop.arg != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10
+      else: mod = 0b00
+    else: mod = 0b11
+    # x 0b0 and idx 0b100 means rsp which means no index exists
+    # rm 0b100 (rsp/r12) signals a sib byte is required, rm then is encoded in the base field of SIB
+    _rm = rm if idx == 0b100 and _x == 0b0 else 0b100
+    inst += bytes([mod << 6 | reg << 3 | _rm])
+    # SIB byte
+    if _rm == 0b100 and mod != 0b11:
+      scale = {1: 0b00, 2: 0b01, 4: 0b10, 8: 0b11}[1 if idx == 0b100 and _x == 0b0 else rm_sz]
+      inst += bytes([scale << 6 | idx << 3 | rm])
+    # DISP byte
+    if mod == 0b01 or mod == 0b10:
+      assert disp_uop is not None
+      inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.arg)
+    # IMM byte
+    if imm_uop is not None:
+      if imm_uop.op is Ops.CONST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.arg)
+      elif isinstance(imm_uop.reg, Register): inst += bytes([(imm_uop.reg.index & 0b1111) << 4 | 0b0000])
+    return inst
+
+  # get the encoding structure of the uop
+  # when a uop writes to memory it takes the form of a store, dtype is void, no definition
+  address:tuple[UOp|None, ...]
+  if x.arg in X86GroupOp.WriteMem:
+    if len(x.src) > 3: address, rest = x.src[:3], x.src[3:]
+    else: address, rest = (x, None, None), x.src
+    return _encode(rest[0], *address, *(None, *rest[1:])) if reg is None else _encode(None, *address, *(None, *rest[:1]))
+
+  if x.arg in X86GroupOp.Rm1st:
+    if len(x.src) > 2: address, rest = x.src[:3], x.src[3:]
+    else: address, rest = (x.src[0], None, None), x.src[1:]
+    imm_uop = rest[:1] if rest and (rest[0].op is Ops.CONST or isinstance(rest[0].reg, Register)) else (None,)
+    return _encode(x, *address, *(None, *imm_uop)) if reg is None else _encode(None, *address, *(x if sel else None, *imm_uop))
+
+  if x.arg in X86GroupOp.Rm2nd:
+    if len(x.src) > 3: address, rest = x.src[1:4], x.src[:1] + x.src[4:]
+    else: address, rest = (x.src[1], None, None), x.src[:1] + x.src[2:]
+    # cmp/vucomiss reg, rm don't define a new register
+    return _encode(x, *address, *rest) if x.dtype is not dtypes.void else _encode(rest[0], *address)
+
+  return None
+
+# https://www.felixcloutier.com/x86/
+# NOTE: LEGACY prefix == VEX prefix
+# pp field: None == 0, 66 == 1, F3 == 2, F2 == 3
+# map select: 0F == 1, 0F38 == 2, 0F3A == 3
+encodings = {
+  # moves
+  X86Ops.MOVABS: lambda x:
+   bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | x.tag[0].index >> 3, 0xB8 + (x.tag[0].index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].arg),
+  X86Ops.MOV: lambda x: encode(x, 0x8B), X86Ops.MOVi: lambda x: encode(x, 0xC7, reg=0),
+  X86Ops.MOVm: lambda x: encode(x, 0x89), X86Ops.LEA: lambda x: encode(x, 0x8D),
+  X86Ops.VMOVSS: lambda x: encode(x, 0x10, pp=2, sel=1), X86Ops.VMOVSSm: lambda x: encode(x, 0x11, pp=2, sel=1),
+  X86Ops.VMOVSD: lambda x: encode(x, 0x10, pp=3, sel=1), X86Ops.VMOVSDm: lambda x: encode(x, 0x11, pp=3, sel=1),
+  X86Ops.VMOVUPS: lambda x: encode(x, 0x10, pp=0, sel=1), X86Ops.VMOVUPSm: lambda x: encode(x, 0x11, pp=0, sel=1),
+  X86Ops.VMOVD: lambda x: encode(x, 0x6E, pp=1, sel=1), X86Ops.VMOVQ: lambda x: encode(x, 0x6E, pp=1, sel=1, we=1),
+  X86Ops.VMOVDm: lambda x: encode(x, 0x7E, pp=1, sel=1), X86Ops.VMOVQm: lambda x: encode(x, 0x7E, pp=1, sel=1, we=1),
+  # casts
+  X86Ops.MOVZX: lambda x: encode(x, 0x0FB7),
+  X86Ops.MOVSX: lambda x: encode(x, 0x0FBF), X86Ops.MOVSXD: lambda x: encode(x, 0x63),
+  X86Ops.VPMOVZXBW: lambda x: encode(x, 0x30, pp=1, sel=2), X86Ops.VPMOVZXBD: lambda x: encode(x, 0x31, pp=1, sel=2),
+  X86Ops.VPMOVZXBQ: lambda x: encode(x, 0x32, pp=1, sel=2), X86Ops.VPMOVZXWD: lambda x: encode(x, 0x33, pp=1, sel=2),
+  X86Ops.VPMOVZXWQ: lambda x: encode(x, 0x34, pp=1, sel=2), X86Ops.VPMOVZXDQ: lambda x: encode(x, 0x35, pp=1, sel=2),
+  X86Ops.VPMOVSXBW: lambda x: encode(x, 0x20, pp=1, sel=2), X86Ops.VPMOVSXBD: lambda x: encode(x, 0x21, pp=1, sel=2),
+  X86Ops.VPMOVSXBQ: lambda x: encode(x, 0x22, pp=1, sel=2), X86Ops.VPMOVSXWD: lambda x: encode(x, 0x23, pp=1, sel=2),
+  X86Ops.VPMOVSXWQ: lambda x: encode(x, 0x24, pp=1, sel=2), X86Ops.VPMOVSXDQ: lambda x: encode(x, 0x25, pp=1, sel=2),
+  X86Ops.VCVTSS2SD: lambda x: encode(x, 0x5A, pp=2, sel=1), X86Ops.VCVTSD2SS: lambda x: encode(x, 0x5A, pp=3, sel=1),
+  X86Ops.VCVTPH2PS: lambda x: encode(x, 0x13, pp=1, sel=2), X86Ops.VCVTPS2PH: lambda x: encode(x, 0x1D, pp=1, sel=3),
+  X86Ops.VCVTDQ2PS: lambda x: encode(x, 0x5B, pp=0, sel=1), X86Ops.VCVTDQ2PD: lambda x: encode(x, 0xE6, pp=2, sel=1),
+  X86Ops.VCVTPS2PD: lambda x: encode(x, 0x5A, pp=0, sel=1), X86Ops.VCVTPD2PS: lambda x: encode(x, 0x5A, pp=1, sel=1),
+  X86Ops.VCVTTPS2DQ: lambda x: encode(x, 0x5B, pp=2, sel=1), X86Ops.VCVTTPD2DQ: lambda x: encode(x, 0xE6, pp=1, sel=1),
+  X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=x.src[1].dtype.itemsize == 8),
+  X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=x.src[1].dtype.itemsize == 8),
+  X86Ops.VCVTTSS2SI: lambda x: encode(x, 0x2C, pp=2, sel=1, we=x.dtype.itemsize == 8),
+  X86Ops.VCVTTSD2SI: lambda x: encode(x, 0x2C, pp=3, sel=1, we=x.dtype.itemsize == 8),
+  # int division
+  X86Ops.IDIV: lambda x: encode(x, 0xF7, reg=7), X86Ops.DIV: lambda x: encode(x, 0xF7, reg=6),
+  # scalar int binary
+  X86Ops.SHLi: lambda x: encode(x, 0xC1, reg=4),
+  X86Ops.SHRi: lambda x: encode(x, 0xC1, reg=5), X86Ops.SARi: lambda x: encode(x, 0xC1, reg=7),
+  X86Ops.ADD: lambda x: encode(x, 0x03), X86Ops.ADDi: lambda x: encode(x, 0x81, reg=0),
+  X86Ops.SUB: lambda x: encode(x, 0x2B), X86Ops.SUBi: lambda x: encode(x, 0x81, reg=5),
+  X86Ops.AND: lambda x: encode(x, 0x23), X86Ops.ANDi: lambda x: encode(x, 0x81, reg=4),
+  X86Ops.XOR: lambda x: encode(x, 0x33), X86Ops.XORi: lambda x: encode(x, 0x81, reg=6),
+  X86Ops.OR: lambda x: encode(x, 0x0B), X86Ops.ORi: lambda x: encode(x, 0x81, reg=1),
+  X86Ops.CMP: lambda x: encode(x, 0x3B), X86Ops.CMPi: lambda x: encode(x, 0x81, reg=7),
+  X86Ops.IMUL: lambda x: encode(x, 0x0FAF), X86Ops.IMULi: lambda x: encode(x, 0x69),
+  X86Ops.SETB: lambda x: encode(x, 0x0F92, reg=0), X86Ops.SETL: lambda x: encode(x, 0x0F9C, reg=0),
+  X86Ops.SETE: lambda x: encode(x, 0x0F94, reg=0), X86Ops.SETNE: lambda x: encode(x, 0x0F95, reg=0),
+  # packed bitwise NOTE: only bitwise and packed
+  X86Ops.VPAND: lambda x: encode(x, 0xDB, pp=1, sel=1), X86Ops.VPXOR: lambda x: encode(x, 0xEF, pp=1, sel=1),
+  X86Ops.VPOR: lambda x: encode(x, 0xEB, pp=1, sel=1),
+  # unary
+  X86Ops.VSQRTSS: lambda x: encode(x, 0x51, pp=2, sel=1), X86Ops.VSQRTPS: lambda x: encode(x, 0x51, pp=0, sel=1),
+  X86Ops.VSQRTSD: lambda x: encode(x, 0x51, pp=3, sel=1), X86Ops.VSQRTPD: lambda x: encode(x, 0x51, pp=1, sel=1),
+  X86Ops.VROUNDSS: lambda x: encode(x, 0x0A, pp=1, sel=3), X86Ops.VROUNDPS: lambda x: encode(x, 0x08, pp=1, sel=3),
+  X86Ops.VROUNDSD: lambda x: encode(x, 0x0B, pp=1, sel=3), X86Ops.VROUNDPD: lambda x: encode(x, 0x09, pp=1, sel=3),
+  # packed int binary
+  X86Ops.VPSLLVD: lambda x: encode(x, 0x47, pp=1, sel=2), X86Ops.VPSLLVQ: lambda x: encode(x, 0x47, pp=1, sel=2, we=1),
+  X86Ops.VPSRLVD: lambda x: encode(x, 0x45, pp=1, sel=2), X86Ops.VPSRLVQ: lambda x: encode(x, 0x45, pp=1, sel=2, we=1),
+  X86Ops.VPCMPGTB: lambda x: encode(x, 0x64, pp=1, sel=1), X86Ops.VPCMPGTW: lambda x: encode(x, 0x65, pp=1, sel=1),
+  X86Ops.VPCMPGTD: lambda x: encode(x, 0x66, pp=1, sel=1), X86Ops.VPCMPGTQ: lambda x: encode(x, 0x37, pp=1, sel=2),
+  X86Ops.VPCMPEQB: lambda x: encode(x, 0x74, pp=1, sel=1), X86Ops.VPCMPEQW: lambda x: encode(x, 0x75, pp=1, sel=1),
+  X86Ops.VPCMPEQD: lambda x: encode(x, 0x76, pp=1, sel=1), X86Ops.VPCMPEQQ: lambda x: encode(x, 0x29, pp=1, sel=2),
+  X86Ops.VPMULLW: lambda x: encode(x, 0xD5, pp=1, sel=1), X86Ops.VPMULLD: lambda x: encode(x, 0x40, pp=1, sel=2),
+  X86Ops.VPADDB: lambda x: encode(x, 0xFC, pp=1, sel=1), X86Ops.VPADDW: lambda x: encode(x, 0xFD, pp=1, sel=1),
+  X86Ops.VPADDD: lambda x: encode(x, 0xFE, pp=1, sel=1), X86Ops.VPADDQ: lambda x: encode(x, 0xD4, pp=1, sel=1),
+  X86Ops.VPSUBB: lambda x: encode(x, 0xF8, pp=1, sel=1), X86Ops.VPSUBW: lambda x: encode(x, 0xF9, pp=1, sel=1),
+  X86Ops.VPSUBD: lambda x: encode(x, 0xFA, pp=1, sel=1), X86Ops.VPSUBQ: lambda x: encode(x, 0xFB, pp=1, sel=1),
+  X86Ops.VPSRAVD: lambda x: encode(x, 0x46, pp=1, sel=2),
+  # float cmp
+  X86Ops.VUCOMISS: lambda x: encode(x, 0x2E, pp=0, sel=1), X86Ops.VUCOMISD: lambda x: encode(x, 0x2E, pp=1, sel=1),
+  # scalar / packed float binary
+  X86Ops.VADDSS: lambda x: encode(x, 0x58, pp=2, sel=1), X86Ops.VADDPS: lambda x: encode(x, 0x58, pp=0, sel=1),
+  X86Ops.VADDSD: lambda x: encode(x, 0x58, pp=3, sel=1), X86Ops.VADDPD: lambda x: encode(x, 0x58, pp=1, sel=1),
+  X86Ops.VSUBSS: lambda x: encode(x, 0x5C, pp=2, sel=1), X86Ops.VSUBPS: lambda x: encode(x, 0x5C, pp=0, sel=1),
+  X86Ops.VSUBSD: lambda x: encode(x, 0x5C, pp=3, sel=1), X86Ops.VSUBPD: lambda x: encode(x, 0x5C, pp=1, sel=1),
+  X86Ops.VMULSS: lambda x: encode(x, 0x59, pp=2, sel=1), X86Ops.VMULPS: lambda x: encode(x, 0x59, pp=0, sel=1),
+  X86Ops.VMULSD: lambda x: encode(x, 0x59, pp=3, sel=1), X86Ops.VMULPD: lambda x: encode(x, 0x59, pp=1, sel=1),
+  X86Ops.VDIVSS: lambda x: encode(x, 0x5E, pp=2, sel=1), X86Ops.VDIVPS: lambda x: encode(x, 0x5E, pp=0, sel=1),
+  X86Ops.VDIVSD: lambda x: encode(x, 0x5E, pp=3, sel=1), X86Ops.VDIVPD: lambda x: encode(x, 0x5E, pp=1, sel=1),
+  X86Ops.VCMPSS: lambda x: encode(x, 0xC2, pp=2, sel=1), X86Ops.VCMPPS: lambda x: encode(x, 0xC2, pp=0, sel=1),
+  X86Ops.VCMPSD: lambda x: encode(x, 0xC2, pp=3, sel=1), X86Ops.VCMPPD: lambda x: encode(x, 0xC2, pp=1, sel=1),
+  X86Ops.VMAXSS: lambda x: encode(x, 0x5F, pp=2, sel=1), X86Ops.VMAXPS: lambda x: encode(x, 0x5F, pp=0, sel=1),
+  X86Ops.VMAXSD: lambda x: encode(x, 0x5F, pp=3, sel=1), X86Ops.VMAXPD: lambda x: encode(x, 0x5F, pp=1, sel=1),
+  X86Ops.VMINSS: lambda x: encode(x, 0x5D, pp=2, sel=1), X86Ops.VMINPS: lambda x: encode(x, 0x5D, pp=0, sel=1),
+  X86Ops.VMINSD: lambda x: encode(x, 0x5D, pp=3, sel=1), X86Ops.VMINPD: lambda x: encode(x, 0x5D, pp=1, sel=1),
+  # ternary
+  X86Ops.CMOVB: lambda x: encode(x, 0x0F42), X86Ops.CMOVL: lambda x: encode(x, 0x0F4C),
+  X86Ops.CMOVE: lambda x: encode(x, 0x0F44), X86Ops.CMOVNE: lambda x: encode(x, 0x0F45),
+  X86Ops.VFMADD213SS: lambda x: encode(x, 0xA9, pp=1, sel=2), X86Ops.VFMADD213SD: lambda x: encode(x, 0xA9, pp=1, sel=2, we=1),
+  X86Ops.VFMADD213PS: lambda x: encode(x, 0xA8, pp=1, sel=2), X86Ops.VFMADD213PD: lambda x: encode(x, 0xA8, pp=1, sel=2, we=1),
+  X86Ops.VBLENDVPS: lambda x: encode(x, 0x4A, pp=1, sel=3), X86Ops.VBLENDVPD: lambda x: encode(x, 0x4B, pp=1, sel=3),
+  X86Ops.VPBLENDVB: lambda x: encode(x, 0x4C, pp=1, sel=3),
+  # shuffles
+  X86Ops.VPBROADCASTB: lambda x: encode(x, 0x78, pp=1, sel=2), X86Ops.VPBROADCASTW: lambda x: encode(x, 0x79, pp=1, sel=2),
+  X86Ops.VPBROADCASTD: lambda x: encode(x, 0x58, pp=1, sel=2), X86Ops.VPBROADCASTQ: lambda x: encode(x, 0x59, pp=1, sel=2),
+  X86Ops.VBROADCASTSS: lambda x: encode(x, 0x18, pp=1, sel=2), X86Ops.VPSRLDQ: lambda x: encode(x, 0x73, reg=3, pp=1, sel=1),
+  X86Ops.VPINSRB: lambda x: encode(x, 0x20, pp=1, sel=3), X86Ops.VPINSRW: lambda x: encode(x, 0xC4, pp=1, sel=1),
+  X86Ops.VPINSRD: lambda x: encode(x, 0x22, pp=1, sel=3), X86Ops.VPINSRQ: lambda x: encode(x, 0x22, pp=1, sel=3, we=1),
+  X86Ops.VSHUFPS: lambda x: encode(x, 0xC6, pp=0, sel=1), X86Ops.VSHUFPD: lambda x: encode(x, 0xC6, pp=1, sel=1),
+  X86Ops.VINSERTPS: lambda x: encode(x, 0x21, pp=1, sel=3),
+  # extract
+  X86Ops.VPEXTRB: lambda x: encode(x, 0x14, pp=1, sel=3), X86Ops.VPEXTRW: lambda x: encode(x, 0x15, pp=1, sel=3),
+  X86Ops.VPEXTRD: lambda x: encode(x, 0x16, pp=1, sel=3), X86Ops.VPEXTRQ: lambda x: encode(x, 0x16, pp=1, sel=3, we=1),
+  # jumps are encoded with a placeholder which gets patched later once the real offset is known
+  X86Ops.JE: lambda x: bytes([0x0F, 0x84]) + int(0).to_bytes(4),
+  X86Ops.JNE: lambda x: bytes([0x0F, 0x85]) + int(0).to_bytes(4),
+  X86Ops.JL: lambda x: bytes([0x0F, 0x8C]) + int(0).to_bytes(4),
+  X86Ops.JB: lambda x: bytes([0x0F, 0x82]) + int(0).to_bytes(4),
+  X86Ops.JGE: lambda x: bytes([0x0F, 0x8D]) + int(0).to_bytes(4),
+  X86Ops.JMP: lambda x: bytes([0xE9]) + int(0).to_bytes(4),
+  X86Ops.RET: lambda x: bytes([0xC3]),
+}
+
+class X86Renderer(ISARenderer):
+  device = "CPU"
+  has_local = False
+  has_threads = bool(getenv("THREADS", 1))
+  global_max = (CPU_COUNT.value, 0, 0)
+  extra_matcher = extra_matcher
+  pre_isel_matcher = pre_isel_matcher
+  isel_matcher = isel_matcher
+  pre_regalloc_matcher = pre_regalloc_matcher
+  late_regalloc_matcher = late_regalloc_matcher
+  post_regalloc_matcher = post_regalloc_matcher
+  isa_spec = isa_spec
+  code_for_op = {x: lambda: None for x in (Ops.SQRT, Ops.AND, Ops.OR, Ops.SHL, Ops.SHR, Ops.NEG, Ops.SUB, Ops.FDIV, Ops.CMPLT, Ops.CMPEQ)}
+  def __init__(self, target:Target):
+    super().__init__(target)
+    from tinygrad.runtime.support.compiler_cpu import X86Compiler
+    self.compiler = X86Compiler()
+  def callee_saved(self):
+    ordered = (RSP,) + tuple(r for r in CALLEE_SAVED if r is not RSP)
+    return tuple(def_reg(dtypes.uint64 if r in GPR else dtypes.float64.vec(2), r) for r in ordered)
+  def is_two_address(self, x:UOp) -> bool: return x.arg in X86GroupOp.TwoAddress
+  # nasty hacks to deal with pointers TODO: rm pointers
+  def copy(self, x:UOp, reg:Register):
+    dt = dtypes.uint64 if isinstance(x.dtype, PtrDType) else x.dtype
+    ret = isel_matcher.rewrite(UOp(Ops.COPY, dt, (x,), tag=reg))
+    assert ret is not None
+    return ret.replace(dtype=x.dtype)
+
+  def spill(self, disp:UOp, x:UOp) -> UOp:
+    nx = x.replace(dtype=dtypes.uint64 if isinstance(x.dtype, PtrDType) else x.dtype)
+    ret = isel_matcher.rewrite(def_reg(dtypes.uint64, RSP).index(disp).store(nx))
+    assert ret is not None
+    return ret.replace(src=(s if s is not nx else x for s in ret.src))
+
+  def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp:
+    ndt = dtypes.uint64 if isinstance(x.dtype, PtrDType) else x.dtype
+    ret = isel_matcher.rewrite(def_reg(dtypes.uint64, RSP).index(disp).load(dtype=ndt, tag=reg))
+    assert ret is not None
+    return ret.replace(dtype=x.dtype)
+
+  def asm(self, uops:list[UOp], function_name:str) -> str:
+    def _format_op(x:UOp) -> str: return f"    {(o[7:-1] if (o:=str(x.arg))[-1] in ('i', 'm') else o[7:]).lower():7s}"
+    def _format_operands(x:UOp) -> str:
+      def _format(src:tuple[UOp, ...]) -> list[str]:
+        return [str(s.arg) if s.op is Ops.CONST else reg_strs[o].get(s.dtype.itemsize if not isinstance(s.dtype, PtrDType) else 8, o) if \
+                (o:=str(s.reg)) in reg_strs else o for s in src if s.op is Ops.CONST or s.reg is not None]
+      def _mem_adress(base:UOp, idx:UOp, disp:UOp) -> str:
+        return f"[{base.reg}" + (f" + {idx.reg}*{base.dtype.itemsize}" if idx.reg else "") + (f" + {disp.arg}" if disp.arg else "") + "]"
+
+      if len(x.src) > 3 and x.arg in X86GroupOp.WriteMem: return ", ".join([_mem_adress(*x.src[:3])] + _format(x.src[3:]))
+      elif len(x.src) > 2 and x.arg in X86GroupOp.Rm1st: return ", ".join(_format((x,)) + [_mem_adress(*x.src[:3])] + _format(x.src[3:]))
+      elif len(x.src) > 3 and x.arg in X86GroupOp.Rm2nd: return ", ".join(_format((x, x.src[0])) + [_mem_adress(*x.src[1:4])] + _format(x.src[4:]))
+      return ", ".join(_format((x,) + x.src))
+
+    asm = [f".{function_name}:"]
+    for u in uops:
+      if u.op is not Ops.INS: continue
+      if u.arg is X86Ops.DEFINE_REG: continue
+      if u.arg is X86Ops.LABEL: asm.append(f"{str(u.tag)}:")
+      elif u.arg is X86Ops.RET: asm.append(_format_op(u))
+      else: asm.append(_format_op(u) + " " + _format_operands(u))
+    return "\n".join(asm)
+
+  def render(self, uops:list[UOp]) -> str:
+    targets: dict[str, int] = {}
+    jumps: dict[UOp, int] = {}
+    binary = bytearray()
+    for u in uops:
+      if u.op is not Ops.INS: continue
+      if u.arg is X86Ops.DEFINE_REG: continue
+      if u.arg is X86Ops.LABEL:
+        targets[u.tag] = len(binary)
+        continue
+      if u.arg not in encodings or (l:=encodings[u.arg](u)) is None:
+        raise RuntimeError(f"failed to encode {u.arg} with {u.dtype} srcs {[x.dtype for x in u.src]}")
+      binary.extend(l)
+      if u.arg in (X86Ops.JL, X86Ops.JB, X86Ops.JE, X86Ops.JNE, X86Ops.JGE, X86Ops.JMP): jumps[u] = len(binary)
+    # fixup jump targets now that encoding size is known
+    for u in uops:
+      if (t:=jumps.get(u)) is not None: binary[t-4:t] = (targets[u.tag] - t).to_bytes(4, 'little', signed=True)
+    return binary.hex()

tinygrad/runtime/graph/hcq.py

@@ -144,7 +144,7 @@ class HCQGraph(MultiGraphRunner):
       self.last_j[enqueue_queue] = j
 
     # Check which signals are used in the profile graph.
-    self.prof_signal_is_used = [any(ent.st_id == j or ent.en_id == j for ent in self.prof_graph_entries) for j in range(len(self.jit_cache) * 2)]
+    self.prof_signal_is_used: set[int] = {sid for ent in self.prof_graph_entries for sid in (ent.st_id, ent.en_id)}
 
     # Build hardware queues.
     self.copy_to_devs: dict[HCQCompiled, set[HCQCompiled]] = {dev: set() for dev in self.devices}
@@ -167,7 +167,7 @@ class HCQGraph(MultiGraphRunner):
       for sig, val in sync_signals + deps: enqueue_queue.wait(sig, val)
 
       # Encode waits and start profile timestamp (if needed).
-      if PROFILE and self.prof_signal_is_used[j * 2]: enqueue_queue.timestamp(self.prof_signals[j * 2])
+      if PROFILE and j * 2 in self.prof_signal_is_used: enqueue_queue.timestamp(self.prof_signals[j * 2])
 
       # Encode main commands based on ji type.
       if isinstance(ji.prg, CompiledRunner):
@@ -203,7 +203,7 @@ class HCQGraph(MultiGraphRunner):
         self.copy_to_devs[cast(HCQCompiled, Device[dest.device])].add(cast(HCQCompiled, Device[src.device]))
 
       # Encode finish profile timestamp (if needed).
-      if PROFILE and self.prof_signal_is_used[j * 2 + 1]: enqueue_queue.timestamp(self.prof_signals[j * 2 + 1])
+      if PROFILE and j * 2 + 1 in self.prof_signal_is_used: enqueue_queue.timestamp(self.prof_signals[j * 2 + 1])
 
       if signal_val is not None: enqueue_queue.signal(signal, signal_val)
 

tinygrad/runtime/ops_cpu.py

@@ -8,6 +8,7 @@ from tinygrad.runtime.support.hcq import CLikeArgsState
 from tinygrad.renderer.cstyle import ClangJITRenderer
 from tinygrad.renderer.llvmir import CPULLVMRenderer
 from tinygrad.renderer.nir import LVPRenderer
+from tinygrad.renderer.isa.x86 import X86Renderer
 from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.uop.ops import sint
 
@@ -136,5 +137,5 @@ class CPUDevice(HCQCompiled):
   def __init__(self, device:str=""):
     self.tasks:queue.Queue = queue.Queue()
     CPUWorker(self, self.tasks, thread_id=0).start()
-    renderers:list[type[Renderer]] = [ClangJITRenderer, CPULLVMRenderer, LVPRenderer]
+    renderers:list[type[Renderer]] = [ClangJITRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer]
     super().__init__(device, CPUAllocator(self), renderers, functools.partial(CPUProgram, self), CPUSignal, CPUComputeQueue)

tinygrad/runtime/ops_rdma.py

@@ -5,8 +5,8 @@ from tinygrad.uop.ops import sint
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocatorBase, HCQAllocator, HWQueue, HCQBuffer, FileIOInterface
 from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta
 from tinygrad.runtime.support.memory import VirtMapping, AddrSpace
-from tinygrad.runtime.support.mlx.mlxdev import MLXDev, MLXQP, to_be
-from tinygrad.helpers import unwrap
+from tinygrad.runtime.support.mlx.mlxdev import MLXDev, MLXQP
+from tinygrad.helpers import unwrap, to_be32, to_be64
 
 class RDMACopyQueue(HWQueue):
   def __init__(self, dev:RDMADevice):
@@ -19,10 +19,10 @@ class RDMACopyQueue(HWQueue):
 
   def encode_ring(self, hwq:HWQueue, dev:HCQCompiled, iface:MLXIface, qp:MLXQP, cq_buf:HCQBuffer, head:sint, ring_uar:bool=False):
     for buf in [iface.dbr_buf, cq_buf] + ([iface.uar_buf] if ring_uar else []): cast(HCQAllocator, dev.allocator).map(buf)
-    hwq.write(iface.dbr_buf.offset(qp.qp_dbr + (4 if ring_uar else 0)), to_be('I', head + 1))
-    if ring_uar: hwq.write(iface.uar_buf.offset(0x800), to_be('Q', ((head << 8) | 0x0a) << 32 | ((qp.qp_info['qpn'] << 8) | 2)), b64=True)
+    hwq.write(iface.dbr_buf.offset(qp.qp_dbr + (4 if ring_uar else 0)), to_be32(head + 1))
+    if ring_uar: hwq.write(iface.uar_buf.offset(0x800), to_be64(((head << 8) | 0x0a) << 32 | ((qp.qp_info['qpn'] << 8) | 2)), b64=True)
     hwq.poll_bit(cq_buf.offset((head & (qp.cq_size - 1)) * 64 + 60, 4), ((head >> (qp.cq_size.bit_length() - 1)) & 1) << 24, mask=0x01000000)
-    hwq.write(iface.dbr_buf.offset(qp.cq_dbr), to_be('I', (head + 1) & 0xFFFFFF))
+    hwq.write(iface.dbr_buf.offset(qp.cq_dbr), to_be32((head + 1) & 0xFFFFFF))
     return self
 
   def copy(self, dest:HCQBuffer, src:HCQBuffer, sz:int):
@@ -39,8 +39,8 @@ class RDMACopyQueue(HWQueue):
   def _submit(self, dev:RDMADevice):
     for remote_nic, sq_wqe, rq_wqe in zip(self._q[0::3], self._q[1::3], self._q[2::3]):
       src_qp, dest_qp, _, _ = dev.iface.connect(remote_nic)
-      assert src_qp.head + 1 - to_be('I', src_qp.dev.dbr[src_qp.qp_dbr // 4 + 1]) <= (1 << src_qp.log_sq_size), "SQ ring full"
-      assert src_qp.head + 1 - to_be('I', dest_qp.dev.dbr[dest_qp.qp_dbr // 4]) <= (1 << dest_qp.log_rq_size), "RQ ring full"
+      assert src_qp.head + 1 - to_be32(src_qp.dev.dbr[src_qp.qp_dbr // 4 + 1]) <= (1 << src_qp.log_sq_size), "SQ ring full"
+      assert src_qp.head + 1 - to_be32(dest_qp.dev.dbr[dest_qp.qp_dbr // 4]) <= (1 << dest_qp.log_rq_size), "RQ ring full"
       dest_qp.qp_buf.view((src_qp.head & ((1 << dest_qp.log_rq_size) - 1)) * 16, 16)[:] = rq_wqe
       sq_view = src_qp.qp_buf.view(src_qp.sq_offset + (src_qp.head & ((1 << src_qp.log_sq_size) - 1)) * 64, 64)
       sq_view[:] = struct.pack('>I', (src_qp.head << 8) | 0x0a) + sq_wqe[4:]

tinygrad/runtime/support/c.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 import ctypes, functools, os, pathlib, re, sys, sysconfig
 from tinygrad.helpers import ceildiv, getenv, unwrap, DEBUG, OSX, WIN
-from _ctypes import Array as _CArray, _SimpleCData, _Pointer
 from typing import TYPE_CHECKING, get_type_hints, get_args, get_origin, overload, Annotated, Any, Generic, Iterable, ParamSpec, TypeVar
 
 def _do_ioctl(__idir, __base, __nr, __struct, __fd, *args, __payload=None, **kwargs):
@@ -34,22 +33,22 @@ if TYPE_CHECKING:
   from _ctypes import _CData
   class Array(Generic[T, U], _CData):
     @overload
-    def __getitem__(self: Array[_SimpleCData[V], Any], key: int) -> V: ...
+    def __getitem__(self: Array[ctypes._SimpleCData[V], Any], key: int) -> V: ...
     @overload
     def __getitem__(self: Array[T, Any], key: slice) -> list[T]: ...
     @overload
     def __getitem__(self: Array[T, Any], key: int) -> T: ...
     def __getitem__(self, key) -> Any: ...
     @overload
-    def __setitem__(self: Array[_SimpleCData[V], Any], key: int, val: V): ...
+    def __setitem__(self: Array[ctypes._SimpleCData[V], Any], key: int, val: V): ...
     @overload
     def __setitem__(self: Array[T, Any], key: int, val: T): ...
     @overload
     def __setitem__(self: Array[T, Any], key: slice, val: Iterable[T]): ...
     def __setitem__(self, key, val): ...
-  class POINTER(Generic[T], _Pointer): ...
+  class POINTER(Generic[T], ctypes._Pointer): ...
   class CFUNCTYPE(Generic[T, P], _CFunctionType): ...
-  class Enum(_SimpleCData):
+  class Enum(ctypes._SimpleCData):
     @classmethod
     def get(cls, val:int, default="unknown") -> str: ...
     @classmethod
@@ -80,14 +79,9 @@ else:
       return val
   def pointer(obj): return ctypes.pointer(obj)
 
-def i2b(i:int, sz:int) -> bytes: return i.to_bytes(sz, sys.byteorder)
-def b2i(b:bytes) -> int: return int.from_bytes(b, sys.byteorder)
-def mv(st) -> memoryview: return memoryview(st).cast('B')
-
 class Struct(ctypes.Structure):
   def __init__(self, *args, **kwargs):
     ctypes.Structure.__init__(self)
-    self._objects_ = {}
     for f,v in [*zip((rf[0] for rf in self._real_fields_), args), *kwargs.items()]: setattr(self, f, v)
 
 def record(cls) -> type[Struct]:
@@ -98,38 +92,38 @@ def record(cls) -> type[Struct]:
 def init_records() -> None:
   for cls, struct, ns in _pending_records:
     setattr(struct, '_real_fields_', [])
-    for nm, t in get_type_hints(cls, globalns=ns, include_extras=True).items():
-      if t.__origin__ in (bool, bytes, str, int, float): setattr(struct, nm, Field(*(f:=t.__metadata__)))
-      else: setattr(struct, nm, Field(*(f:=(del_an(t.__origin__), *t.__metadata__))))
-      struct._real_fields_.append((nm,) + f) # type: ignore
+    for i, (nm, t) in enumerate(get_type_hints(cls, globalns=ns, include_extras=True).items()):
+      struct._real_fields_.append((nm, *(f:=(del_an(t.__origin__), *t.__metadata__) if isinstance(t.__metadata__[0], int) else t.__metadata__))) # type: ignore
+      setattr(struct, nm, Field(nm, i, *f))
   _pending_records.clear()
 
-class Field(property):
-  def __init__(self, typ, off:int, bit_width=None, bit_off=0):
-    if bit_width is not None:
-      sl, set_mask = slice(off,off+(sz:=ceildiv(bit_width+bit_off, 8))), ~((mask:=(1 << bit_width) - 1) << bit_off)
+class Field:
+  def __init__(self, nm, idx, typ, off, bit_width=None, bit_off=0):
+    self.nm, self.idx, self.typ, self.off, self.bit_width, self.bit_off = nm, idx, typ, off, bit_width, bit_off
+
+  # lazily resolve field descriptors
+  def _resolve(self, cls):
+    if self.bit_width: # handle bitfields ourselves
+      sl, set_mask = slice(self.off, self.off+(sz:=ceildiv(self.bit_width+self.bit_off, 8))), ~((mask:=(1 << self.bit_width) - 1) << self.bit_off)
+      def b2i(obj): return int.from_bytes(memoryview(obj).cast("B")[sl], sys.byteorder)
+      def bset(obj, v): memoryview(obj).cast("B")[sl] = ((b2i(obj) & set_mask) | v << self.bit_off).to_bytes(sz, sys.byteorder)
       # FIXME: signedness
-      super().__init__(lambda self: (b2i(mv(self)[sl]) >> bit_off) & mask,
-                       lambda self,v: mv(self).__setitem__(sl, i2b((b2i(mv(self)[sl]) & set_mask) | (v << bit_off), sz)))
-    else:
-      sl = slice(off, off + ctypes.sizeof(typ))
-      def set_with_objs(f):
-        def wrapper(self, v):
-          if hasattr(v, '_objects') and hasattr(self, '_objects_'): self._objects_[off] = {'_self_': v, **(v._objects or {})}
-          mv(self).__setitem__(sl, bytes(v if isinstance(v, typ) else f(v)))
-        return wrapper
-      if issubclass(typ, _CArray):
-        getter = (lambda self: typ.from_buffer(mv(self)[sl]).value) if typ._type_ is ctypes.c_char else (lambda self: typ.from_buffer(mv(self)[sl]))
-        super().__init__(getter, set_with_objs(lambda v: typ(*v)))
-      else: super().__init__(lambda self: v.value if isinstance(v:=typ.from_buffer(mv(self)[sl]), _SimpleCData) else v, set_with_objs(typ))
-    self.offset = off
+      cf = property(lambda obj: b2i(obj) >> self.bit_off & mask, bset)
+    # pull the CField descriptor from a dummy class, zero length arrays are so ctypes manages references to child objects for us
+    else: cf = type(self.nm, (ctypes.Structure,), {"_layout_": "ms", "_pack_": 1, "_fields_": [(str(i), ctypes.c_byte * 0) for i in range(self.idx)] +
+                                                                                              [("_", ctypes.c_byte * self.off), ("v", self.typ)]}).v # type: ignore
+    setattr(cls, self.nm, cf)
+    return cf
+
+  def __get__(self, obj, objtype=None): return self._resolve(objtype).__get__(obj, objtype) if objtype else self
+  def __set__(self, obj, value): self._resolve(obj.__class__).__set__(obj, value)
 
 @functools.cache
 def init_c_struct_t(sz:int, fields: tuple[tuple, ...]):
   CStruct = type("CStruct", (Struct,), {'_fields_': [('_mem_', ctypes.c_byte * sz)], '_real_fields_': []})
-  for nm,ty,*args in fields:
-    setattr(CStruct, nm, Field(*(f:=(del_an(ty), *args))))
-    CStruct._real_fields_.append((nm,) + f) # type: ignore
+  for i,(nm,ty,*args) in enumerate(fields):
+    CStruct._real_fields_.append((nm, *(f:=(del_an(ty), *args)))) # type: ignore
+    setattr(CStruct, nm, Field(nm, i, *f))
   return CStruct
 def init_c_var(ty, creat_cb): return (creat_cb(v:=del_an(ty)()), v)[1]
 

tinygrad/runtime/support/compiler_cpu.py

@@ -91,3 +91,8 @@ class CPULLVMCompiler(LLVMCompiler):
     # +reserve-x18 here does the same thing as -ffixed-x18 in ops_cpu.py, see comments there for why it's needed on arm osx
     cpu, feats = ctypes.string_at(llvm.LLVMGetHostCPUName()), (b'+reserve-x18,' if OSX else b'') + ctypes.string_at(llvm.LLVMGetHostCPUFeatures())
     super().__init__(cpu.decode(), feats.decode(), cache_key)
+
+class X86Compiler(Compiler):
+  def __init__(self): super().__init__(None)
+  def compile(self, src:str) -> bytes: return bytes.fromhex(src)
+  def disassemble(self, lib:bytes): return capstone_flatdump(lib)

tinygrad/runtime/support/mlx/mlxdev.py

@@ -11,7 +11,6 @@ MLX5_CMD_STRUCTS = {v: (getattr(mlx5, f"struct_mlx5_ifc_{n[12:].lower()}_in_bits
   getattr(mlx5, f"struct_mlx5_ifc_{n[12:].lower()}_out_bits", None)) for n, v in mlx5.__dict__.items() if n.startswith("MLX5_CMD_OP_")}
 MLX5_CMD_STRUCTS[mlx5.MLX5_CMD_OP_ACCESS_REG] = (mlx5.struct_mlx5_ifc_access_register_in_bits, mlx5.struct_mlx5_ifc_access_register_out_bits)
 
-def to_be(fmt, val): return to_be32(val) if fmt == 'I' else to_be64(val)
 def ipv4_to_gid(ip): return bytes(10) + b'\xff\xff' + socket.inet_aton(ip)
 
 def udp_sport(lqpn, rqpn):
@@ -64,7 +63,7 @@ class MLXCmdQueue:
     for i in range(n):
       off, _ = self.mboxes[base + i]
       blk = mlx5.struct_mlx5_cmd_prot_block(data=list(data[i*chunk_sz:(i+1)*chunk_sz].ljust(chunk_sz, b'\x00')),
-        next=to_be('Q', self.mboxes[base+i+1][1]) if i < n-1 else 0, block_num=to_be('I', i), token=tok)
+        next=to_be64(self.mboxes[base+i+1][1]) if i < n-1 else 0, block_num=to_be32(i), token=tok)
       self.queue[off:off + ctypes.sizeof(mlx5.struct_mlx5_cmd_prot_block)] = bytes(blk)
     return (self.mboxes[base][0], self.mboxes[base][1], n)
 
@@ -81,9 +80,9 @@ class MLXCmdQueue:
     # prepare mailboxes and build command layout
     _, in_ptr, n_in = self.create_mbox_chain(0, tok, inp[16:])
     _, out_ptr, n_out = self.create_mbox_chain(n_in, tok, bytes(out_sz))
-    cmd = mlx5.struct_mlx5_cmd_layout(type=mlx5.MLX5_PCI_CMD_XPORT, inlen=to_be('I', len(inp)), in_ptr=to_be('Q', in_ptr),
+    cmd = mlx5.struct_mlx5_cmd_layout(type=mlx5.MLX5_PCI_CMD_XPORT, inlen=to_be32(len(inp)), in_ptr=to_be64(in_ptr),
       _in=[int.from_bytes(inp[i:i+4], 'little') for i in range(0, 16, 4)],
-      out_ptr=to_be('Q', out_ptr), outlen=to_be('I', 16 + out_sz), token=tok, status_own=mlx5.CMD_OWNER_HW)
+      out_ptr=to_be64(out_ptr), outlen=to_be32(16 + out_sz), token=tok, status_own=mlx5.CMD_OWNER_HW)
     cmd_bytes = bytearray(bytes(cmd))
     cmd_bytes[mlx5.struct_mlx5_cmd_layout.sig.offset] = (~functools.reduce(lambda a, b: a ^ b, cmd_bytes)) & 0xFF  # type: ignore[attr-defined]
 
@@ -115,8 +114,8 @@ class MLXDev:
 
     self.init_hw(ip)
 
-  def rreg(self, off): return to_be('I',self.bar[off // 4])
-  def wreg(self, off, val): self.bar[off // 4] = to_be('I',val)
+  def rreg(self, off): return to_be32(self.bar[off // 4])
+  def wreg(self, off, val): self.bar[off // 4] = to_be32(val)
   def iseg_r(self, field): return self.rreg(getattr(mlx5.struct_mlx5_init_seg, field).offset)
   def iseg_w(self, field, val): self.wreg(getattr(mlx5.struct_mlx5_init_seg, field).offset, val)
 

tinygrad/runtime/support/system.py

@@ -66,7 +66,9 @@ class _System:
       iokit.IOServiceGetMatchingServices(0, iokit.IOServiceMatching(b"IOPCIDevice"), ctypes.byref(iterator:=ctypes.c_uint()))
       while svc:=iokit.IOIteratorNext(iterator): all_devs.append((v:=read_prop(svc, "vendor-id"), d:=read_prop(svc, "device-id"), f"{v:x}:{d:x}"))
     else:
-      for pcibus in FileIOInterface("/sys/bus/pci/devices").listdir():
+      try: devs = FileIOInterface("/sys/bus/pci/devices")
+      except FileNotFoundError: raise RuntimeError("no pcie")
+      for pcibus in devs.listdir():
         if base_class is not None and int(FileIOInterface(f"/sys/bus/pci/devices/{pcibus}/class").read(), 16) >> 16 != base_class: continue
         all_devs.append((int(FileIOInterface(f"/sys/bus/pci/devices/{pcibus}/vendor").read(), 16),
                          int(FileIOInterface(f"/sys/bus/pci/devices/{pcibus}/device").read(), 16), pcibus))

tinygrad/tensor.py

@@ -7,7 +7,7 @@ if TYPE_CHECKING: import numpy
 from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_float, least_upper_dtype, to_dtype, truncate
 from tinygrad.dtype import _from_np_dtype, _to_np_dtype, PyConst, Invalid, InvalidType
 from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten
-from tinygrad.helpers import IMAGE, FLOAT16, WINO, Metadata, TRACEMETA, ASM_GEMM, ceildiv, fetch, is_numpy_ndarray, TracingKey, cpu_profile
+from tinygrad.helpers import IMAGE, FLOAT16, WINO, Metadata, TRACEMETA, ceildiv, fetch, is_numpy_ndarray, TracingKey, cpu_profile
 from tinygrad.helpers import suppress_finalizing, disable_gc
 from tinygrad.gradient import compute_gradient
 from tinygrad.mixin import OpMixin
@@ -832,7 +832,7 @@ class Tensor(OpMixin):
     if isinstance(self.device, tuple): return self._multi_like(Tensor.rand, **kwargs)
     return Tensor.rand(*self.shape, device=kwargs.pop("device", self.device), dtype=kwargs.pop("dtype", self.dtype), **kwargs)
 
-  # ***** rng hlops *****
+  # ***** random functions *****
 
   def randn_like(self, dtype:DTypeLike|None=None, requires_grad:bool|None=None, **kwargs) -> Tensor:
     """
@@ -880,9 +880,8 @@ class Tensor(OpMixin):
     print(Tensor.randint(2, 3, low=5, high=10).numpy())
     ```
     """
-    if not isinstance(low, int) or not isinstance(high, int): raise TypeError(f"{low=} and {high=} must be integers")
-    dtype = to_dtype(dtype)
-    if not dtypes.is_int(dtype): raise TypeError(f"{dtype=} must be int")
+    if not all_int([low, high]): raise TypeError(f"{low=} and {high=} must be integers")
+    if not dtypes.is_int(dtype := to_dtype(dtype)): raise TypeError(f"{dtype=} must be int")
     return Tensor.uniform(*shape, low=low, high=high, dtype=dtype, **kwargs)
 
   @staticmethod
@@ -898,7 +897,7 @@ class Tensor(OpMixin):
     print(Tensor.normal(2, 3, mean=10, std=2).numpy())
     ```
     """
-    return ((std * Tensor.randn(*shape, **kwargs)) + mean).requires_grad_(requires_grad)
+    return (std * Tensor.randn(*shape, **kwargs) + mean).requires_grad_(requires_grad)
 
   @staticmethod
   def uniform(*shape, low=0.0, high=1.0, dtype:DTypeLike|None=None, requires_grad:bool|None=None, **kwargs) -> Tensor:
@@ -931,7 +930,6 @@ class Tensor(OpMixin):
     """
     return Tensor.uniform(*shape, low=-1.0, high=1.0, **kwargs).mul(prod(argfix(*shape))**-0.5)
 
-  # https://www.tensorflow.org/api_docs/python/tf/keras/initializers/GlorotUniform
   @staticmethod
   def glorot_uniform(*shape, **kwargs) -> Tensor:
     """
@@ -945,9 +943,9 @@ class Tensor(OpMixin):
     print(Tensor.glorot_uniform(2, 3).numpy())
     ```
     """
-    return Tensor.uniform(*shape, low=-1.0, high=1.0, **kwargs).mul((6/(argfix(*shape)[0]+prod(argfix(*shape)[1:])))**0.5)
+    bound = (6 / (argfix(*shape)[0]+prod(argfix(*shape)[1:]))) ** 0.5
+    return Tensor.uniform(*shape, low=-bound, high=bound, **kwargs)
 
-  # https://pytorch.org/docs/stable/_modules/torch/nn/init.html#kaiming_uniform_
   @staticmethod
   def kaiming_uniform(*shape, a:float = 0.01, **kwargs) -> Tensor:
     """
@@ -961,10 +959,9 @@ class Tensor(OpMixin):
     print(Tensor.kaiming_uniform(2, 3).numpy())
     ```
     """
-    bound = math.sqrt(3.0) * math.sqrt(2.0 / (1 + a ** 2)) / math.sqrt(prod(argfix(*shape)[1:]))
+    bound = (6 / (1 + a ** 2) / prod(argfix(*shape)[1:])) ** 0.5
     return Tensor.uniform(*shape, low=-bound, high=bound, **kwargs)
 
-  # https://pytorch.org/docs/stable/_modules/torch/nn/init.html#kaiming_normal_
   @staticmethod
   def kaiming_normal(*shape, a:float = 0.01, **kwargs) -> Tensor:
     """
@@ -978,7 +975,7 @@ class Tensor(OpMixin):
     print(Tensor.kaiming_normal(2, 3).numpy())
     ```
     """
-    std = math.sqrt(2.0 / (1 + a ** 2)) / math.sqrt(prod(argfix(*shape)[1:]))
+    std = (2 / (1 + a ** 2) / prod(argfix(*shape)[1:])) ** 0.5
     return Tensor.normal(*shape, mean=0.0, std=std, **kwargs)
 
   @staticmethod
@@ -1059,7 +1056,7 @@ class Tensor(OpMixin):
       else: t.grad.assign(t.grad + g.to(t.grad.device))
     return self
 
-  # ***** movement low level ops *****
+  # ***** movement ops *****
 
   def _mop(self, op:Ops, arg) -> Tensor: return self._apply_uop(UOp._mop, extra_args=(op,), arg=arg)
   def _rop(self, op:Ops, axis:tuple[int, ...]) -> Tensor: return self._apply_uop(UOp._rop, op=op, axis=axis)
@@ -1137,8 +1134,6 @@ class Tensor(OpMixin):
     if mode in {"reflect", "replicate"}: return self._pad_reflect_replicate(pX, mode)
     raise NotImplementedError(f"{mode=} is not supported")
 
-  # ***** movement high level ops *****
-
   def _getitem(self, indices, v: Tensor|None = None) -> Tensor:
     # wrap single index into a list
     if (isinstance(indices, list) and all_int(indices)) or not isinstance(indices, (tuple, list)): indices = [indices]
@@ -2032,9 +2027,6 @@ class Tensor(OpMixin):
 
   def dot(self, w:Tensor, dtype:DTypeLike|None=None) -> Tensor:
     if IMAGE: return self.image_dot(w, dtype)
-    if ASM_GEMM:
-      from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
-      if can_use_asm_gemm(self, w): return asm_gemm(self, w)
     return super().dot(w, dtype=dtype)
 
   def _cumalu(self, axis:int, op:Ops, _include_initial=False) -> Tensor:
@@ -2629,14 +2621,6 @@ class Tensor(OpMixin):
     print(q.scaled_dot_product_attention(k, v).numpy())
     ```
     """
-    if getenv("FLASH_ATTENTION"):
-      from extra.thunder.tiny.fa import flash_attention
-      return flash_attention(self, key, value, attn_mask=attn_mask, is_causal=is_causal)
-
-    if getenv("HK_FLASH_ATTENTION"):
-      from extra.thunder.amd.fa import flash_attention
-      return flash_attention(self, key, value, attn_mask=attn_mask, is_causal=is_causal)
-
     # GQA: https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
     if enable_gqa:
       key = key.repeat_interleave(int(self.shape[-3] // key.shape[-3]), dim=-3)

tinygrad/uop/ops.py

@@ -473,6 +473,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def end(self, *src:UOp): return UOp(Ops.END, src=(self,)+src) if len(src) else self
   def after(self, *src:UOp, **kwargs): return UOp(Ops.AFTER, self.dtype, (self,)+src, **kwargs) if len(src) else self
   def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
+  def ins(self, arg, **kwargs): return UOp(Ops.INS, kwargs.pop("dtype", self.dtype), kwargs.pop("src", self.src), arg, kwargs.pop("tag", self.tag))
   def contract(self, *rngs:UOp):
     assert all(x.arg[-1] == AxisType.UPCAST for x in rngs), "all contract ranges must be upcast"
     return UOp(Ops.CONTRACT, dtype=self.dtype.vec(prod([x.vmax+1 for x in rngs])), src=(self,), arg=tuple((x.arg[0], x.vmax+1) for x in rngs))
@@ -537,6 +538,13 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       for s in self.src: yield from s.split_uop(sep)
     else: yield self
 
+  @property
+  def reg(self:UOp):
+    # TODO: add a way to access the nth element in src, sea of nodes call this a projection
+    if self.op in (Ops.NOOP, Ops.AFTER) and self.src: return self.src[0].reg
+    if isinstance(self.tag, tuple): return self.tag[0]
+    return self.tag
+
   # *** multi-device helpers ***
 
   def multi(self, axis:int|None):
@@ -650,16 +658,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     if ret.shape == self.shape and same_shape_noop: return self
     return ret
 
-  # in these four, if the shape doesn't change we can return self
-  #def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=True)
-  #def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg, same_shape_noop=True)
-  #def shrink(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.SHRINK, arg, same_shape_noop=True)
-  def pad(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.PAD, arg, same_shape_noop=True)
-
-  # in these two, we have custom logic to check if they are a no-op
-  #def permute(self, arg:tuple[int, ...]): return self._mop(Ops.PERMUTE, arg, same_shape_noop=False) if arg != tuple(range(len(self.shape))) else self
-  #def flip(self, arg:tuple[bool, ...]): return self._mop(Ops.FLIP, arg, same_shape_noop=False) if any(arg) and len(arg) == len(self.shape) else self
-
   # *** uop UNIQUE ***
 
   # TODO: use this in Buffer
@@ -1020,10 +1018,13 @@ def exec_alu(op:Ops, dtype:DType, operands, truncate_output=True):
 # ***** uop helpers *****
 
 def print_uops(uops:list[UOp]):
+  def format_tag(u:UOp) -> str: return "" if u.tag is None else str(u.tag)
   uops_index = {u:i for i,u in enumerate(uops)}
+  tag_width = max((len(format_tag(u)) for u in uops), default=0)
   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}")
+    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} {format_tag(u):{tag_width}s} {u.arg}")
 
 # ***** pattern matcher *****
 

tinygrad/viz/serve.py

@@ -45,7 +45,7 @@ from tinygrad.dtype import dtypes
 
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
                **{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, Ops.REDUCE_AXIS: "#FF6B6B", Ops.SHAPED_WMMA: "#FF5B5B",
-               Ops.RANGE: "#c8a0e0", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
+               Ops.RANGE: "#76349c", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
                Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.INS: "#eec4ff",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
                Ops.BUFFER_VIEW: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
@@ -104,7 +104,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
     if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE, Ops.LUNIQUE} and u is not x: excluded.add(u)
     if u.op is Ops.CONST and len(u.src) and u.src[0].op in {Ops.UNIQUE, Ops.LUNIQUE}: excluded.remove(u)
     if u.op is Ops.VCONST and u.dtype.scalar() == dtypes.weakint and u is not x: excluded.add(u)
-    if u.op is Ops.VECTORIZE and len(u.src) == 0: excluded.add(u)
+    if u.op in {Ops.VECTORIZE, Ops.NOOP} and len(u.src) == 0: excluded.add(u)
     # exclude RESHAPE/EXPAND that only serve to broadcast a CONST
     if u.op in {Ops.RESHAPE, Ops.EXPAND} and len(u.src) >= 1 and u.src[0] in excluded and u is not x: excluded.add(u)
   for u in toposort:
@@ -225,8 +225,9 @@ def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:
       ref = next((v for k in e.name.keys if (v:=ref_map.get(k)) is not None), None)
       if isinstance(e.name.ret, str): fmt.append(e.name.ret)
       elif isinstance(e.name.ret, int):
-        membw = e.name.ret / (dur * 1e-6)
+        membw = (nbytes:=e.name.ret) / (dur * 1e-6)
         fmt.append(f"{membw*1e-9:.0f} GB/s" if membw < 1e13 else f"{membw*1e-12:.0f} TB/s")
+        fmt.append(next((f"{nbytes/s:.0f} {u}" for s,u in ((1e9,"GB"),(1e6,"MB"),(1e3,"KB")) if nbytes>=s), f"{nbytes:.0f} B"))
       elif e.name.tb: fmt.append("TB:"+json.dumps(e.name.tb))
     events.append(struct.pack("<IIIIfI", enum_str(name, scache), option(ref), option(key), rel_ts(st,start_ts, f"'{name}' on {e.device}"),
                               dur, enum_str("\n".join(fmt),scache)))
@@ -287,18 +288,22 @@ metrics:dict[str, Callable[[dict[str, tuple[int, int, int]]], str]] = {
 
 def unpack_pmc(e) -> dict:
   agg_cols = ["Name", "Sum"]
-  sample_cols = ["XCC", "INST", "SE", "SA", "WGP", "Value"]
   rows:list[list] = []
   stats:dict[str, tuple[int, int, int]] = {}  # name -> (sum, max, count)
   view, ptr = memoryview(e.blob).cast('Q'), 0
   for s in e.sched:
+    sample_cols = ["XCC", "INST", "SE", "SA"] + [f"WGP:{i}" for i in range(s.wgp)]
     row:list = [s.name, 0, {"cols":sample_cols, "rows":[]}]
     max_val, cnt = 0, 0
-    for sample in itertools.product(range(s.xcc), range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
-      row[1] += (val:=int(view[ptr]))
-      max_val, cnt = max(max_val, val), cnt + 1
-      row[2]["rows"].append(sample+(val,))
-      ptr += 1
+    for sample in itertools.product(range(s.xcc), range(s.inst), range(s.se), range(s.sa)):
+      vals:list[int] = []
+      # pack work group processors on the same se
+      for _ in range(s.wgp):
+        row[1] += (val:=int(view[ptr]))
+        max_val, cnt = max(max_val, val), cnt + 1
+        vals.append(val)
+        ptr += 1
+      row[2]["rows"].append(sample+tuple(vals))
     stats[s.name] = (row[1], max_val, cnt)
     rows.append(row)
   for name, fn in metrics.items():
@@ -351,12 +356,10 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> Generator[ProfileEvent,
   row_counts:dict[str, itertools.count] = {}
   curr_barrier:dict[int, ProfileRangeEvent] = {}
   exec_pending:dict[str, list[tuple[str, str]]] = {}
-  is_cdna = target.startswith("gfx9")
   dispatch_to_exec = {"WMMA":"VALU", "VALU":"VALU", "VALU1":"VALU", "VALUT":"VALU", "VALUB":"VALU", "VALUINST":"VALU", "VINTERP":"VALU",
                       "SGMEM":"VMEM", "FLAT":"VMEM", "LDS":"LDS", "SALU":"SALU", "SMEM":"SALU", "VMEM":"VMEM"}
   def add(name:str, p:PacketType, wave:int|None=None, info:InstructionInfo|None=None) -> Generator[ProfileEvent, None, None]:
-    row = "OTHER_SIMD" if name.startswith("OTHER_") else f"WAVE:{wave}" if (wave:=getattr(p, "wave", wave)) is not None \
-        else f"{p.__class__.__name__}:0 {name.replace('_ALT', '')}"
+    row = f"WAVE:{wave}" if (wave:=getattr(p, "wave", wave)) is not None else f"{p.__class__.__name__}:0 {name.replace('_ALT', '')}"
     # by default we extend the packet to one cycle after timestamp
     start_time, end_time = p._time, p._time+1
     # exec links to dispatch, dispatch links to PC
@@ -365,11 +368,11 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> Generator[ProfileEvent,
       dispatch_id, op_type = exec_pending[name].pop(0)
       # get the number of cycles from the op type
       duration = int(dur_match.group(1)) if (dur_match:=re.match(r".*_(\d+)$", op_type)) else 1
-      # for execs, timestamp is the completion time
-      start_time, end_time = p._time-duration, p._time
+      # for execs, extend end time by the duration
+      start_time, end_time = p._time, p._time+duration
       link = f"LINK:{dispatch_id}"
-      # add wmma in the exec name for coloring
-      if op_type.startswith("WMMA"): name += "_WMMA"
+      # wmma exec gets its own row and color
+      if op_type.startswith("WMMA"): name, row = name+"_WMMA", "ALUEXEC:0 WMMA"
     # queue inst dispatches
     idx = next(row_counts.setdefault(row, itertools.count(0)))
     if isinstance(p, (VALUINST, INST, INST_RDNA4)) and (exec_type:=dispatch_to_exec.get(name.replace("OTHER_", "").split("_")[0])) is not None:
@@ -381,9 +384,6 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> Generator[ProfileEvent,
     # construct and yield the event for this packet
     if row not in row_ends: yield ProfilePointEvent(row, "JSON", "pcMap", pc_map, ts=Decimal(0))
     yield (e:=ProfileRangeEvent(row, TracingKey(name, ret=link), Decimal(start_time), Decimal(end_time)))
-    # allow CDNA packets to overlap, NOT allowed on RDNA.
-    if (et:=row_ends.get(row)) is not None and e.st < et and not is_cdna and not isinstance(p, (ALUEXEC, VMEMEXEC)):
-      RuntimeError(f"packet {row}-{idx} overlaps: {e.st} {et}.")
     row_ends[row] = unwrap(e.en)
     # barrier on this wave extends to fill the time it was waiting
     if wave is not None: