failing test from application

* assembly/amd: add ds perm instructions

* NO SKIP

* fix preexisting RDNA3 issues

* pcode

* assert

* asserts

* unify

* simp

* good fix

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -10,6 +10,7 @@ export DEBUG=${DEBUG:-0}
 export FLASH_ATTENTION=${FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}

extra/assembly/amd/emu.py

@@ -402,17 +402,19 @@ class _Ctx:
     return UOp.sink(*self.scalar_stores(assigns, sdst_reg, sdst_size), *self.inc_pc())
 
   def compile_lane_pcode(self, op, inst) -> UOp:
-    """Compile READLANE/READFIRSTLANE/WRITELANE using pcode parser."""
+    """Compile cross-lane ops (READLANE/WRITELANE/PERMLANE) using pcode parser."""
     pcode = get_pcode(op)
     op_name = op.name if hasattr(op, 'name') else str(op)
     src0_off, vdst_off = self.inst_field(type(inst).src0), self.inst_field(type(inst).vdst)
     src0_reg = (src0_off >= _c(256)).where(src0_off - _c(256), _c(0))  # VGPR index or 0
     src1_off = self.inst_field(type(inst).src1) if hasattr(type(inst), 'src1') else None
+    src2_off = self.inst_field(type(inst).src2) if hasattr(type(inst), 'src2') else None
     exec_lo = self.rsgpr_dyn(_c(EXEC_LO.offset))
     srcs = {
       'SRC0': src0_reg, 'VDST': vdst_off, 'EXEC_LO': exec_lo, 'EXEC': exec_lo.cast(dtypes.uint64), '_vgpr': self.vgpr,
       'S0': self.rsrc_dyn(src0_off, _c(0, dtypes.int)) if 'WRITELANE' in op_name else src0_reg,
       'S1': self.rsrc_dyn(src1_off, _c(0, dtypes.int)) if src1_off is not None else _c(0),
+      'S2': self.rsrc_dyn(src2_off, _c(0, dtypes.int)) if src2_off is not None else _c(0),
     }
     _, assigns = parse_pcode(pcode, srcs)
     stores = []
@@ -427,7 +429,8 @@ class _Ctx:
     pcode = get_pcode(op)
     vcc_reg = sdst_reg if sdst_reg is not None else VCC_LO.offset
     if 'VCC' not in srcs: srcs['VCC'] = self.rsgpr_dyn(_c(vcc_reg))
-    srcs.update({'EXEC': exec_mask, 'SCC': self.rsgpr_dyn(_c(SCC.offset)), 'laneId': lane})
+    srcs.update({'EXEC': exec_mask, 'SCC': self.rsgpr_dyn(_c(SCC.offset)), 'laneId': lane,
+                 'ROUND_MODE': _c(0), 'ROUND_TOWARD_ZERO': _c(0)})  # rounding mode: 0=RNE, RTZ constant
     _, assigns = parse_pcode(pcode, srcs)
 
     raw_stores: list = []
@@ -484,6 +487,7 @@ def _compile_sopp(inst: SOPP, ctx: _Ctx) -> UOp:
   if inst.op == SOPPOp.S_ENDPGM:
     return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)),
                           ctx.wsgpr_dyn(_c(PC_HI_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)))
+  if inst.op == SOPPOp.S_NOP: return UOp.sink(*ctx.inc_pc())  # S_NOP is a no-op
   # NOTE: we ignore SOPPs without PCODE
   if inst.op in PCODE:
     pcode = get_pcode(inst.op)
@@ -551,7 +555,7 @@ def _compile_sop(inst: SOP1 | SOP2 | SOPC | SOPK, ctx: _Ctx) -> UOp:
 
 def _compile_vop12(inst: VOP1 | VOP1_SDST | VOP2, ctx: _Ctx) -> UOp:
   op_name = _op_name(inst)
-  if op_name == 'V_READFIRSTLANE_B32_E32': return ctx.compile_lane_pcode(inst.op, inst)
+  if op_name in ('V_READFIRSTLANE_B32_E32', 'V_PERMLANE64_B32_E32'): return ctx.compile_lane_pcode(inst.op, inst)
   lane, exec_mask, bits = ctx.range(), ctx.rsgpr_dyn(_c(EXEC_LO.offset)), inst.canonical_op_bits
   literal = ctx.inst_field(type(inst).literal) if hasattr(type(inst), 'literal') else None
   vdst_reg = ctx.inst_field(VOP1.vdst)
@@ -644,6 +648,10 @@ def _compile_vop3(inst: VOP3, ctx: _Ctx) -> UOp:
   if op_name in ('V_READLANE_B32', 'V_READFIRSTLANE_B32', 'V_READFIRSTLANE_B32_E64', 'V_WRITELANE_B32'):
     return ctx.compile_lane_pcode(inst.op, inst)
 
+  # V_PERMLANE16_B32 / V_PERMLANEX16_B32: cross-lane swizzle via pcode
+  if 'PERMLANE16' in op_name or 'PERMLANEX16' in op_name:
+    return ctx.compile_lane_pcode(inst.op, inst)
+
   # VOP3 VOPC (v_cmp_*_e64) - delegate to unified VOPC handler
   if 'V_CMP' in op_name or 'V_CMPX' in op_name:
     return _compile_vopc(inst, ctx, opsel=opsel, abs_bits=getattr(inst, 'abs', 0) or 0, neg_bits=getattr(inst, 'neg', 0) or 0)
@@ -787,7 +795,7 @@ def _compile_vop3p(inst: VOP3P, ctx: _Ctx) -> UOp:
     s0_mod = apply_neg_mix(apply_abs(src0, 1, 1, 1), 1, 1, 1)
     s1_mod = apply_neg_mix(apply_abs(src1, 2, 2, 2), 2, 2, 2)
     s2_mod = apply_neg_mix(apply_abs(src2, 4, 4, 4), 4, 4, 4)
-    srcs = {'S0': s0_mod, 'S1': s1_mod, 'S2': s2_mod,
+    srcs = {'S@0': s0_mod, 'S@1': s1_mod, 'S@2': s2_mod,
             'OPSEL_HI': UOp.const(dtypes.uint32, combined_opsel_hi), 'OPSEL': UOp.const(dtypes.uint32, opsel)}
   else:
     def get_half_bits(val: UOp, use_hi: bool, apply_neg: bool = False) -> UOp:
@@ -796,10 +804,13 @@ def _compile_vop3p(inst: VOP3P, ctx: _Ctx) -> UOp:
       return bits
     def build_remapped_src(src: UOp, opsel_lo_bit: int, opsel_hi_bit: int, neg_lo_bit: int, neg_hi_bit: int) -> UOp:
       return get_half_bits(src, bool(opsel_lo_bit), bool(neg_lo_bit)) | (get_half_bits(src, bool(opsel_hi_bit), bool(neg_hi_bit)) << UOp.const(dtypes.uint32, 16))
-    s0_new = build_remapped_src(src0, opsel & 1, opsel_hi & 1, neg & 1, neg_hi & 1)
-    s1_new = build_remapped_src(src1, opsel & 2, opsel_hi & 2, neg & 2, neg_hi & 2)
-    s2_new = build_remapped_src(src2, opsel & 4, 1 if opsel_hi2 else 0, neg & 4, neg_hi & 4)
-    srcs = {'S0': s0_new, 'S1': s1_new, 'S2': s2_new}
+    # DOT IU instructions use NEG bits for signed/unsigned selection, not fp16 negation
+    is_dot_iu = 'DOT' in op_name and 'IU' in op_name
+    n0, n1, n2, nh0, nh1, nh2 = (0, 0, 0, 0, 0, 0) if is_dot_iu else (neg & 1, neg & 2, neg & 4, neg_hi & 1, neg_hi & 2, neg_hi & 4)
+    srcs = {'S0': build_remapped_src(src0, opsel & 1, opsel_hi & 1, n0, nh0),
+            'S1': build_remapped_src(src1, opsel & 2, opsel_hi & 2, n1, nh1),
+            'S2': build_remapped_src(src2, opsel & 4, 1 if opsel_hi2 else 0, n2, nh2)}
+    if is_dot_iu: srcs['NEG'] = UOp.const(dtypes.uint32, neg)
   return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask)
 
 def _compile_vopd(inst: VOPD, ctx: _Ctx) -> UOp:
@@ -868,6 +879,15 @@ def _compile_mem_op(inst: DS | FLAT | GLOBAL | SCRATCH, ctx: _Ctx) -> UOp:
   has_data1 = is_lds and hasattr(inst, 'data1') and inst.data1 is not None
   data1_reg = ctx.inst_field(DS.data1) if is_lds else _c(0)
 
+  # DS_PERMUTE/DS_BPERMUTE: cross-lane VGPR access via pcode
+  if is_lds and 'PERMUTE' in op_name:
+    pcode = get_pcode(inst.op)
+    srcs = {'ADDR': addr_reg, 'DATA0': vdata_reg, 'VDST': vdst_reg, 'OFFSET': offset,
+            'EXEC': exec_mask.cast(dtypes.uint64), '_vgpr': ctx.vgpr}
+    _, assigns = parse_pcode(pcode, srcs)
+    stores = [ctx.vgpr.index(val[0].cast(dtypes.int)).store(val[1].cast(dtypes.uint32)) for dest, val in assigns if dest.startswith('VGPR[')]
+    return UOp.sink(*stores, *ctx.inc_pc())
+
   def make_addr(lane: UOp) -> UOp:
     if is_lds: return ctx.rvgpr_dyn(addr_reg, lane)
     offset64 = offset.cast(dtypes.uint64)

extra/assembly/amd/pcode.py

@@ -94,13 +94,19 @@ def _trig_reduce(x, phase=0.0):
   return UOp(Ops.SIN, x.dtype, (x - n * _const(x.dtype, 6.283185307179586),))
 
 def _signext(val: UOp) -> UOp:
-  for bits, mask, ext in [(8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
+  for bits, mask, ext in [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
     if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].arg == mask) or val.dtype.itemsize == bits // 8:
       v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
       sb = (v32 >> _u32(bits - 1)) & _u32(1)
       return sb.ne(_u32(0)).where(v32 | _u32(ext), v32).cast(dtypes.int)
   return val.cast(dtypes.int64) if val.dtype in (dtypes.int, dtypes.int32) else val
 
+def _signext_4bit(val: UOp) -> UOp:
+  """Sign extend a 4-bit value to 32-bit signed integer."""
+  v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
+  sb = (v32 >> _u32(3)) & _u32(1)  # sign bit at position 3
+  return sb.ne(_u32(0)).where(v32 | _u32(0xFFFFFFF0), v32).bitcast(dtypes.int)
+
 def _abs(val: UOp) -> UOp:
   if val.dtype not in (dtypes.float32, dtypes.float64, dtypes.half): return val
   _, _, _, _, shift = _float_info(val)
@@ -194,6 +200,17 @@ def _ff1(val: UOp, bits: int) -> UOp:
     result = cond.where(_const(dtypes.int, i), result)
   return result
 
+def _sad_u8(a: UOp, b: UOp, acc: UOp, masked: bool = False) -> UOp:
+  """Sum of absolute differences of 4 unsigned bytes + accumulator. If masked, skips bytes where a == 0."""
+  a, b, acc = a.cast(dtypes.uint32), b.cast(dtypes.uint32), acc.cast(dtypes.uint32)
+  result = acc
+  for i in range(4):
+    a_byte = (a >> _u32(i * 8)) & _u32(0xFF)
+    b_byte = (b >> _u32(i * 8)) & _u32(0xFF)
+    diff = (a_byte > b_byte).where(a_byte - b_byte, b_byte - a_byte)
+    result = result + (a_byte.ne(_u32(0)).where(diff, _u32(0)) if masked else diff)
+  return result
+
 _FUNCS: dict[str, Callable[..., UOp]] = {
   'sqrt': lambda a: UOp(Ops.SQRT, a.dtype, (a,)), 'trunc': lambda a: UOp(Ops.TRUNC, a.dtype, (a,)),
   'log2': lambda a: UOp(Ops.LOG2, a.dtype, (a,)), 'sin': lambda a: _trig_reduce(a),
@@ -227,11 +244,50 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
   'signext_from_bit': _signext_from_bit, 'ldexp': _ldexp, 'frexp_mant': _frexp_mant, 'mantissa': _frexp_mant,
   'frexp_exp': _frexp_exp, 'trig_preop_result': _trig_preop,
   's_ff1_i32_b32': lambda a: _ff1(a, 32), 's_ff1_i32_b64': lambda a: _ff1(a, 64),
+  # Normalization conversions: map [-1,1] or [0,1] to integer range
+  # Use floor(x + 0.5) for round-to-nearest
+  # SNORM: round(value * 32767), range is [-32767, 32767] (hardware behavior)
+  'f16_to_snorm': lambda a: _floor(_f16_extract(a).cast(dtypes.float32) * _const(dtypes.float32, 32767) + _const(dtypes.float32, 0.5)).cast(dtypes.int).cast(dtypes.int16),
+  'f16_to_unorm': lambda a: _floor(_f16_extract(a).cast(dtypes.float32) * _const(dtypes.float32, 65535) + _const(dtypes.float32, 0.5)).cast(dtypes.uint16),
+  'f32_to_snorm': lambda a: _floor(a.bitcast(dtypes.float32) * _const(dtypes.float32, 32767) + _const(dtypes.float32, 0.5)).cast(dtypes.int).cast(dtypes.int16),
+  'f32_to_unorm': lambda a: _floor(a.bitcast(dtypes.float32) * _const(dtypes.float32, 65535) + _const(dtypes.float32, 0.5)).cast(dtypes.uint16),
+  'f32_to_u8': lambda a: _f_to_u(a.bitcast(dtypes.float32), dtypes.uint8),
+  # Integer truncation conversions
+  'i32_to_i16': lambda a: a.cast(dtypes.int).cast(dtypes.int16),
+  'u32_to_u16': lambda a: a.cast(dtypes.uint32).cast(dtypes.uint16),
+  'u16_to_u32': lambda a: (a.cast(dtypes.uint32) & _u32(0xFFFF)),
+  'u8_to_u32': lambda a: (a.cast(dtypes.uint32) & _u32(0xFF)),
+  'u4_to_u32': lambda a: (a.cast(dtypes.uint32) & _u32(0xF)),
+  # Signed extraction with sign extension for dot products
+  'i16_to_i32': lambda a: _signext(a.cast(dtypes.uint32) & _u32(0xFFFF)),
+  'i8_to_i32': lambda a: _signext(a.cast(dtypes.uint32) & _u32(0xFF)),
+  'i4_to_i32': lambda a: _signext_4bit(a.cast(dtypes.uint32) & _u32(0xF)),
+  # Float to int16 conversions
+  'v_cvt_i16_f32': lambda a: UOp(Ops.TRUNC, dtypes.float32, (a.bitcast(dtypes.float32),)).cast(dtypes.int16),
+  'v_cvt_u16_f32': lambda a: _f_to_u(a.bitcast(dtypes.float32), dtypes.uint16),
+  # SAD (Sum of Absolute Differences) - sum |a_i - b_i| for 4 bytes + accumulator
+  'v_sad_u8': lambda a, b, c: _sad_u8(a, b, c),
+  'v_msad_u8': lambda a, b, c: _sad_u8(a, b, c, masked=True),
+  # System NOPs - these are scheduling hints, no effect on emulation
+  'MIN': lambda a, b: (a < b).where(a, b),
+  's_nop': lambda a: _u32(0),
 }
 for is_max, name in [(False, 'min'), (True, 'max')]:
   for dt, sfx in [(dtypes.float32, 'f32'), (dtypes.int, 'i32'), (dtypes.uint32, 'u32'), (dtypes.int16, 'i16'), (dtypes.uint16, 'u16')]:
     _FUNCS[f'v_{name}_{sfx}'] = lambda *a, im=is_max, d=dt: _minmax_reduce(im, d, *a)
     _FUNCS[f'v_{name}3_{sfx}'] = lambda *a, im=is_max, d=dt: _minmax_reduce(im, d, *a)
+# f16 min/max/min3/max3/med3
+for is_max, name in [(False, 'min'), (True, 'max')]:
+  _FUNCS[f'v_{name}_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}3_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}_num_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}_num_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+  _FUNCS[f'v_{name}3_num_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}3_num_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+  _FUNCS[f'v_{name}imum_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}imum_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+  _FUNCS[f'v_{name}imum3_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
+  _FUNCS[f'v_{name}imum3_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # TOKENIZER/PARSER
@@ -239,7 +295,7 @@ for is_max, name in [(False, 'min'), (True, 'max')]:
 
 DTYPES = {'u32': dtypes.uint32, 'i32': dtypes.int, 'f32': dtypes.float32, 'b32': dtypes.uint32, 'u64': dtypes.uint64, 'i64': dtypes.int64,
           'f64': dtypes.float64, 'b64': dtypes.uint64, 'u16': dtypes.uint16, 'i16': dtypes.short, 'f16': dtypes.half, 'b16': dtypes.uint16,
-          'u8': dtypes.uint8, 'i8': dtypes.int8, 'b8': dtypes.uint8, 'u1': dtypes.uint32}
+          'u8': dtypes.uint8, 'i8': dtypes.int8, 'b8': dtypes.uint8, 'u4': dtypes.uint8, 'i4': dtypes.int8, 'u1': dtypes.uint32}
 _BITS_DT = {8: dtypes.uint8, 16: dtypes.uint16, 32: dtypes.uint32, 64: dtypes.uint64}
 _NUM_SUFFIXES = ('ULL', 'LL', 'UL', 'U', 'L', 'F', 'f')
 def _strip_suffix(num: str) -> tuple[str, str]:
@@ -396,7 +452,7 @@ class Parser:
         self.eat('DOT')
         dt_name = self.eat('IDENT').val
         return self._handle_mem_load(addr, DTYPES.get(dt_name, dtypes.uint32))
-      if name == 'VGPR':
+      if name == 'VGPR' and self.at('LBRACKET'):
         self.eat('LBRACKET')
         lane = self.parse()
         self.eat('RBRACKET')
@@ -423,7 +479,21 @@ class Parser:
       if self.try_eat('LBRACE'):
         idx = self.eat('NUM').val
         self.eat('RBRACE')
-        elem = self.vars.get(f'{name}{idx}', _u32(0))
+        # Handle VGPR{lane}[reg] - 2D array access after loop unrolling
+        if name == 'VGPR' and self.at('LBRACKET'):
+          self.eat('LBRACKET')
+          reg = self.parse()
+          self.eat('RBRACKET')
+          vgpr = self.vars.get('_vgpr')
+          if vgpr is None: return _u32(0)
+          return vgpr.index(_to_u32(reg) * _u32(32) + _u32(int(idx)), ptr=True).load()
+        elem = self.vars.get(f'{name}@{idx}', self.vars.get(f'{name}{idx}'))
+        if elem is None:
+          # Extract bit idx from base variable (like var[idx])
+          base = self.vars.get(name)
+          assert isinstance(base, UOp), f"unknown variable: {name}{idx}"
+          dt = dtypes.uint64 if base.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
+          elem = (base.cast(dt) >> _const(dt, int(idx))) & _const(dt, 1)
         if self.try_eat('DOT'):
           dt_name = self.eat('IDENT').val
           return _cast_to(elem, DTYPES.get(dt_name, dtypes.uint32))
@@ -432,27 +502,17 @@ class Parser:
         return elem
       if self.at('LBRACKET') and name not in self.vars:
         self.eat('LBRACKET')
-        if self.at('NUM'):
-          idx_num = int(self.peek().val)
-          if f'{name}{idx_num}' in self.vars:
-            self.eat('NUM')
-            self.eat('RBRACKET')
-            elem = self.vars[f'{name}{idx_num}']
-            if self.try_eat('DOT'): return _cast_to(elem, DTYPES.get(self.eat('IDENT').val, dtypes.uint32))
-            return elem
         first = self.parse()
         return self._handle_bracket_rest(first, _u32(0), name)
       if name in self.vars:
         v = self.vars[name]
-        return v if isinstance(v, UOp) else _u32(0) if isinstance(v, dict) else _u32(0)
+        assert isinstance(v, UOp), f"expected UOp for {name}, got {type(v)}"
+        return v
       raise RuntimeError(f"unknown variable: {name}")
     raise RuntimeError(f"unexpected token in primary: {self.peek()}")
 
   def _handle_dot(self, base, field: str) -> UOp:
-    if isinstance(base, str): return _u32(0)
-    if not isinstance(base, UOp):
-      if isinstance(base, dict): return base.get(field, _u32(0))
-      return _u32(0)
+    assert isinstance(base, UOp), f"expected UOp for dot access, got {type(base)}"
     if field == 'u64' and self.at('LBRACKET') and self.peek(1).type == 'IDENT' and self.peek(1).val == 'laneId':
       self.eat('LBRACKET')
       self.eat_val('laneId', 'IDENT')
@@ -467,6 +527,7 @@ class Parser:
     if dt == base.dtype: return base
     if dt.itemsize == 2 and base.dtype.itemsize == 4:
       return (base & _const(base.dtype, 0xFFFF)).cast(dtypes.uint16) if dt == dtypes.uint16 else (base & _const(base.dtype, 0xFFFF)).cast(dtypes.uint16).bitcast(dt)
+    if field == 'i4': return _signext_4bit(base)
     return _cast_to(base, dt)
 
   def _handle_bracket(self, base, var_name: str | None = None) -> UOp:
@@ -509,16 +570,18 @@ class Parser:
       var_name = self._find_var_name(base)
     if first.op == Ops.CONST:
       idx = int(first.arg)
-      if var_name and f'{var_name}{idx}' in self.vars:
-        v = self.vars[f'{var_name}{idx}']
+      # Check for array element (var@idx)
+      if var_name and f'{var_name}@{idx}' in self.vars:
+        v = self.vars[f'{var_name}@{idx}']
         return _cast_to(v, dt_suffix) if dt_suffix else v
+      # Bit extraction
       dt = dtypes.uint64 if base.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
       base_cast = base.cast(dt) if base.dtype != dt else base
       result = ((base_cast >> _const(dt, idx)) & _const(dt, 1))
       return _cast_to(result, dt_suffix) if dt_suffix else result
     if var_name:
       idx_u32 = _to_u32(first)
-      elems = [(i, self.vars[f'{var_name}{i}']) for i in range(256) if f'{var_name}{i}' in self.vars]
+      elems = [(i, self.vars[f'{var_name}@{i}']) for i in range(256) if f'{var_name}@{i}' in self.vars]
       if elems:
         result = elems[-1][1]
         for ei, ev in reversed(elems[:-1]):
@@ -537,7 +600,7 @@ class Parser:
     self.eat('RBRACE')
     var_name = self._find_var_name(base)
     if var_name:
-      elem = self.vars.get(f'{var_name}{idx}', _u32(0))
+      elem = self.vars.get(f'{var_name}@{idx}', _u32(0))  # use @ to avoid collision with temps like A4
       if self.try_eat('DOT'):
         dt_name = self.eat('IDENT').val
         return _cast_to(elem, DTYPES.get(dt_name, dtypes.uint32))
@@ -599,13 +662,14 @@ class Parser:
     raise RuntimeError(f"unexpected token after {bits}': {self.peek()}")
 
   def _parse_number(self, num: str) -> UOp:
-    if num.startswith('0x') or num.startswith('0X'): return _const(dtypes.uint64, int(num.rstrip('ULul'), 16))
-    suffix, num = _strip_suffix(num)
-    if '.' in num or suffix in ('F', 'f'):
-      return _const(dtypes.float32 if suffix in ('F', 'f') else dtypes.float64, float(num))
-    val = int(num)
-    if 'ULL' in suffix: return _const(dtypes.uint64, val)
-    if 'LL' in suffix or 'L' in suffix: return _const(dtypes.uint64, val)
+    if num.startswith('0x') or num.startswith('0X'):
+      is_u64 = num.upper().endswith('ULL') or num.upper().endswith('LL') or num.upper().endswith('UL')
+      return _const(dtypes.uint64 if is_u64 else dtypes.uint32, int(num.rstrip('ULul'), 16))
+    suffix, num_str = _strip_suffix(num)
+    if '.' in num_str or suffix in ('F', 'f'):
+      return _const(dtypes.float32 if suffix in ('F', 'f') else dtypes.float64, float(num_str))
+    val = int(num_str)
+    if 'ULL' in suffix or 'LL' in suffix or 'L' in suffix: return _const(dtypes.uint64, val)
     if 'U' in suffix: return _const(dtypes.uint32, val)
     return _const(dtypes.int if val < 0 else dtypes.uint32, val)
 
@@ -623,7 +687,8 @@ class Parser:
       if ';' in body or '\n' in body or 'return' in body.lower():
         lines = [l.strip() for l in body.replace(';', '\n').split('\n') if l.strip() and not l.strip().startswith('//')]
         _, _, result = parse_block(lines, 0, lv, self.funcs)
-        return result if result is not None else _u32(0)
+        assert result is not None, f"lambda {name} must return a value"
+        return result
       return parse_expr(body, lv, self.funcs)
     if name in self.funcs:
       return self.funcs[name](*args)
@@ -631,7 +696,7 @@ class Parser:
 
   def _handle_mem_load(self, addr: UOp, dt) -> UOp:
     mem = self.vars.get('_vmem') if '_vmem' in self.vars else self.vars.get('_lds')
-    if mem is None: return _const(dt, 0)
+    assert mem is not None, "memory load requires _vmem or _lds"
     adt = dtypes.uint64 if addr.dtype == dtypes.uint64 else dtypes.uint32
     active = self.vars.get('_active')
     gate = (active,) if active is not None else ()
@@ -693,29 +758,9 @@ def parse_tokens(toks: list[Token], vars: dict[str, VarVal], funcs: dict | None
 
 # Unified block parser for pcode
 def _subst_loop_var(line: str, loop_var: str, val: int) -> str:
-  """Substitute loop variable and evaluate bracket expressions.
-  Converts var[loop_var] to var{val} for array element access (like the old regex parser)."""
+  """Substitute loop variable with its value."""
   toks = tokenize(line)
-  # First pass: convert var[loop_var] to var{loop_var} to mark for array element assignment
-  result_toks: list[Token] = []
-  j = 0
-  while j < len(toks):
-    t = toks[j]
-    # Check for pattern: IDENT[loop_var] where it's not preceded by a dot (not .type[...])
-    if t.type == 'IDENT' and j+3 < len(toks) and toks[j+1].type == 'LBRACKET' and toks[j+2].type == 'IDENT' and toks[j+2].val == loop_var and toks[j+3].type == 'RBRACKET':
-      # Check that it's not .type[loop_var]
-      if not result_toks or result_toks[-1].type != 'DOT':
-        result_toks.append(t)
-        result_toks.append(Token('LBRACE', '{'))
-        result_toks.append(Token('NUM', str(val)))
-        result_toks.append(Token('RBRACE', '}'))
-        j += 4
-        continue
-    result_toks.append(t)
-    j += 1
-  # Second pass: substitute loop variable in remaining positions
-  subst_parts = [str(val) if t.type == 'IDENT' and t.val == loop_var else t.val for t in result_toks if t.type != 'EOF']
-  return ' '.join(subst_parts)
+  return ' '.join(str(val) if t.type == 'IDENT' and t.val == loop_var else t.val for t in toks if t.type != 'EOF')
 
 def _set_bits(old: UOp, val: UOp, width: int, offset: int) -> UOp:
   """Set bits [offset:offset+width) in old to val, masking and shifting appropriately."""
@@ -765,8 +810,9 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
       def parse_bound():
         if p.at('NUM') and p.peek(1).type == 'QUOTE': p.eat('NUM'); p.eat('QUOTE')
         if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
-        expr = p.parse()
-        return int(expr.arg) if expr.op == Ops.CONST else 0
+        expr = p.parse().simplify()
+        assert expr.op == Ops.CONST, f"loop bound must be constant, got {expr}"
+        return int(expr.arg)
       start_val = parse_bound()
       p.eat('COLON')
       end_val = parse_bound()
@@ -787,7 +833,7 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
       if found_var: vars[found_var] = block_assigns[found_var] = _const(dtypes.bool, False)
       for loop_i in range(start_val, end_val + 1):
         subst_lines = [_subst_loop_var(bl, loop_var, loop_i) for bl in body_lines if not (has_break and bl.strip().lower() == 'break')]
-        _, iter_assigns, _ = parse_block(subst_lines, 0, vars, funcs, assigns)
+        _, iter_assigns, _ = parse_block(subst_lines, 0, {**vars, **block_assigns}, funcs, assigns)
         if has_break:
           assert found_var is not None
           found = block_assigns.get(found_var, vars.get(found_var))
@@ -812,7 +858,9 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
 
     # declare
     if first == 'declare':
-      if '[' not in line and len(toks) >= 2 and toks[1].type == 'IDENT': vars[toks[1].val] = _u32(0)
+      # Initialize scalar declarations (skip arrays and vars already passed as srcs)
+      if '[' not in line and len(toks) >= 2 and toks[1].type == 'IDENT':
+        vars.setdefault(toks[1].val, _u32(0))
       i += 1; continue
 
     # lambda definition
@@ -870,6 +918,7 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
       j, lane_toks = _match_bracket(toks, 1)
       if j < len(toks) and toks[j].type == 'LBRACKET':
         j, reg_toks = _match_bracket(toks, j)
+        if j < len(toks) and toks[j].type == 'DOT': j += 2  # skip .type suffix
         if j < len(toks) and toks[j].type == 'EQUALS': j += 1
         ln, rg, val = parse_tokens(lane_toks, vars, funcs), parse_tokens(reg_toks, vars, funcs), parse_tokens(toks[j:], vars, funcs)
         if assigns is not None: assigns.append((f'VGPR[{_tok_str(lane_toks)}][{_tok_str(reg_toks)}]', (_to_u32(rg) * _u32(32) + _to_u32(ln), val)))
@@ -933,19 +982,32 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
             block_assigns[var] = vars[var] = _set_bit(existing, _to_u32(parse_tokens(bit_toks, vars, funcs)), parse_tokens(toks[j+1:], vars, funcs))
             i += 1; continue
 
-    # Array element: var{idx} = value
-    if len(toks) >= 5 and toks[0].type == 'IDENT' and toks[1].type == 'LBRACE' and toks[2].type == 'NUM':
-      var, idx = toks[0].val, int(toks[2].val)
-      j = 4
-      while j < len(toks) and toks[j].type != 'EQUALS': j += 1
-      if j < len(toks):
-        val = parse_tokens(toks[j+1:], vars, funcs)
-        existing = block_assigns.get(var, vars.get(var))
-        if existing is not None and isinstance(existing, UOp):
-          block_assigns[var] = vars[var] = _set_bit(existing, _u32(idx), val)
-        else:
-          block_assigns[f'{var}{idx}'] = vars[f'{var}{idx}'] = val
-        i += 1; continue
+    # Array element: var[idx] = value (static index) or var[expr] = value (dynamic)
+    if len(toks) >= 4 and toks[0].type == 'IDENT' and toks[1].type == 'LBRACKET':
+      var = toks[0].val
+      j, idx_toks = _match_bracket(toks, 1)
+      if j < len(toks) and toks[j].type == 'EQUALS':
+        # Static index: var[NUM] = value
+        if len(idx_toks) == 1 and idx_toks[0].type == 'NUM':
+          idx = int(idx_toks[0].val.rstrip('UuLl'))
+          val = parse_tokens(toks[j+1:], vars, funcs)
+          existing = block_assigns.get(var, vars.get(var))
+          if existing is not None and isinstance(existing, UOp):
+            block_assigns[var] = vars[var] = _set_bit(existing, _u32(idx), val)
+          else:
+            block_assigns[f'{var}@{idx}'] = vars[f'{var}@{idx}'] = val
+          i += 1; continue
+        # Dynamic index: var[expr] = value where var has @-elements
+        elems = [(k.split('@')[1], v) for k, v in {**vars, **block_assigns}.items() if k.startswith(f'{var}@') and isinstance(v, UOp)]
+        if elems:
+          idx_expr = parse_tokens(idx_toks, vars, funcs)
+          val = parse_tokens(toks[j+1:], vars, funcs)
+          for elem_idx_str, old_elem in elems:
+            elem_idx = int(elem_idx_str)
+            cond = _to_u32(idx_expr).eq(_u32(elem_idx))
+            new_val = cond.where(val.cast(old_elem.dtype) if val.dtype != old_elem.dtype else val, old_elem)
+            block_assigns[f'{var}@{elem_idx}'] = vars[f'{var}@{elem_idx}'] = new_val
+          i += 1; continue
 
     # Compound assignment: var += or var -=
     assign_op = next((j for j, t in enumerate(toks) if t.type == 'ASSIGN_OP'), None)
@@ -992,13 +1054,14 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
       def parse_cond(s, kw):
         ll = s.lower()
         return _to_bool(parse_expr(s[ll.find(kw) + len(kw):ll.rfind('then')].strip(), vars, funcs))
-      def not_static_false(c): return c.op != Ops.CONST or c.arg is not False
+      def is_const(c, v): return c.op == Ops.CONST and c.arg is v
       cond = parse_cond(line, 'if')
-      conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not_static_false(cond) else []
+      conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not is_const(cond, False) else []
       else_branch: tuple[UOp | None, dict[str, VarVal]] = (None, {})
       vars_snap = dict(vars)
+      static_true = is_const(cond, True)  # track if any condition is statically true
       i += 1
-      i, branch, ret = parse_block(lines, i, vars, funcs, assigns)
+      i, branch, ret = parse_block(lines, i, vars, funcs, assigns if not is_const(cond, False) else None)
       if conditions: conditions[0] = (cond, ret if ret is not None else branch)
       vars.clear(); vars.update(vars_snap)
       while i < len(lines):
@@ -1007,12 +1070,16 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
         lf = ltoks[0].val.lower()
         if lf == 'elsif':
           c = parse_cond(lines[i], 'elsif')
-          i += 1; i, branch, ret = parse_block(lines, i, vars, funcs, assigns)
-          if not_static_false(c): conditions.append((c, ret if ret is not None else branch))
+          take = not static_true and not is_const(c, False)
+          i += 1; i, branch, ret = parse_block(lines, i, vars, funcs, assigns if take else None)
+          if take:
+            conditions.append((c, ret if ret is not None else branch))
+            if is_const(c, True): static_true = True
           vars.clear(); vars.update(vars_snap)
         elif lf == 'else':
-          i += 1; i, branch, ret = parse_block(lines, i, vars, funcs, assigns)
-          else_branch = (ret, branch)
+          i += 1
+          i, branch, ret = parse_block(lines, i, vars, funcs, assigns if not static_true else None)
+          if not static_true: else_branch = (ret, branch)
           vars.clear(); vars.update(vars_snap)
         elif lf == 'endif': i += 1; break
         else: break
@@ -1024,17 +1091,21 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
             if rv.dtype != result.dtype and rv.dtype.itemsize == result.dtype.itemsize: result = result.cast(rv.dtype)
             result = c.where(rv, result)
         return i, block_assigns, result
-      # Main style: merge variable assignments with WHERE
-      else_assigns = else_branch[1]
-      all_vars = set().union(*[ba.keys() for _, ba in conditions if isinstance(ba, dict)], else_assigns.keys())
-      for var in all_vars:
-        res: Any = else_assigns.get(var, block_assigns.get(var, vars.get(var, _u32(0))))
-        for cond, ba in reversed(conditions):
-          if isinstance(ba, dict) and var in ba:
-            tv = ba[var]
-            if isinstance(tv, UOp) and isinstance(res, UOp):
-              res = cond.where(tv, res.cast(tv.dtype) if tv.dtype != res.dtype and tv.dtype.itemsize == res.dtype.itemsize else res)
-        block_assigns[var] = vars[var] = res
+      # If statically true, use that branch directly; otherwise merge with WHERE
+      if static_true:
+        ba = next((b for c, b in conditions if is_const(c, True) and isinstance(b, dict)), {})
+        block_assigns.update(ba); vars.update(ba)
+      else:
+        else_assigns = else_branch[1]
+        all_vars = set().union(*[ba.keys() for _, ba in conditions if isinstance(ba, dict)], else_assigns.keys())
+        for var in all_vars:
+          res: Any = else_assigns.get(var, block_assigns.get(var, vars.get(var, _u32(0))))
+          for cond, ba in reversed(conditions):
+            if isinstance(ba, dict) and var in ba:
+              tv = ba[var]
+              if isinstance(tv, UOp) and isinstance(res, UOp):
+                res = cond.where(tv, res.cast(tv.dtype) if tv.dtype != res.dtype and tv.dtype.itemsize == res.dtype.itemsize else res)
+          block_assigns[var] = vars[var] = res
       continue
 
     # Regular assignment: var = value

extra/assembly/amd/test/hw/helpers.py

@@ -13,7 +13,7 @@ def _i32(f: float) -> int: return struct.unpack('<I', struct.pack('<f', f))[0]
 def _f32(i: int) -> float: return struct.unpack('<f', struct.pack('<I', i & 0xFFFFFFFF))[0]
 
 # f16 conversion helpers
-def _f16(i: int) -> float: return struct.unpack('<e', struct.pack('<H', i & 0xFFFF))[0]
+def f16(i: int) -> float: return struct.unpack('<e', struct.pack('<H', i & 0xFFFF))[0]
 def f32_to_f16(f: float) -> int:
   f = float(f)
   if math.isnan(f): return 0x7e00
@@ -43,6 +43,23 @@ VCC = VCC_LO  # For VOP3SD sdst field (VCC_LO is exported from dsl)
 USE_HW = os.environ.get("USE_HW", "0") == "1"
 FLOAT_TOLERANCE = 1e-5
 
+def get_gpu_target() -> tuple[int, int, int]:
+  """Get the GPU target as (major, minor, stepping) tuple."""
+  if not USE_HW: return (0, 0, 0)
+  from tinygrad.device import Device
+  return Device["AMD"].target
+
+def skip_unless_gfx(min_major: int, min_minor: int = 0, reason: str = ""):
+  """Skip test if GPU target is below the minimum required version."""
+  import unittest
+  def decorator(test_func):
+    if not USE_HW: return test_func
+    target = get_gpu_target()
+    if target[0] < min_major or (target[0] == min_major and target[1] < min_minor):
+      return unittest.skip(reason or f"requires gfx{min_major}{min_minor}0+")(test_func)
+    return test_func
+  return decorator
+
 # Output buffer layout: vgpr[16][32], sgpr[16], vcc, scc, exec
 N_VGPRS, N_SGPRS, WAVE_SIZE = 16, 16, 32
 VGPR_BYTES = N_VGPRS * WAVE_SIZE * 4  # 16 regs * 32 lanes * 4 bytes = 2048
@@ -212,8 +229,12 @@ amdhsa.kernels:
 
   return parse_output(bytes(out_buf), n_lanes)
 
-def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgprs: int = N_VGPRS) -> list[str]:
-  """Compare two WaveStates and return list of differences."""
+def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgprs: int = N_VGPRS, ulp_tolerance: int = 0) -> list[str]:
+  """Compare two WaveStates and return list of differences.
+
+  Args:
+    ulp_tolerance: Allow up to this many ULPs difference for float comparisons (0 = exact match required)
+  """
   import math
   diffs = []
   for i in range(n_vgprs):
@@ -224,6 +245,11 @@ def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgp
         emu_f, hw_f = _f32(emu_val), _f32(hw_val)
         if math.isnan(emu_f) and math.isnan(hw_f):
           continue
+        # Check ULP difference for floats (only for same-sign values)
+        if ulp_tolerance > 0 and (emu_val < 0x80000000) == (hw_val < 0x80000000):
+          ulp_diff = abs(int(emu_val) - int(hw_val))
+          if ulp_diff <= ulp_tolerance:
+            continue
         diffs.append(f"v[{i}] lane {lane}: emu=0x{emu_val:08x} ({emu_f:.6g}) hw=0x{hw_val:08x} ({hw_f:.6g})")
   for i in range(N_SGPRS):
     emu_val = emu_st.sgpr[i]
@@ -236,16 +262,19 @@ def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgp
     diffs.append(f"scc: emu={emu_st.scc} hw={hw_st.scc}")
   return diffs
 
-def run_program(instructions: list, n_lanes: int = 1) -> WaveState:
+def run_program(instructions: list, n_lanes: int = 1, ulp_tolerance: int = 0) -> WaveState:
   """Run instructions and return WaveState.
 
   If USE_HW=1, runs on both emulator and hardware, compares results, and raises if they differ.
   Otherwise, runs only on emulator.
+
+  Args:
+    ulp_tolerance: Allow up to this many ULPs difference for float comparisons (0 = exact match required)
   """
   emu_st = run_program_emu(instructions, n_lanes)
   if USE_HW:
     hw_st = run_program_hw(instructions, n_lanes)
-    diffs = compare_wave_states(emu_st, hw_st, n_lanes)
+    diffs = compare_wave_states(emu_st, hw_st, n_lanes, ulp_tolerance=ulp_tolerance)
     if diffs:
       raise AssertionError(f"Emulator vs Hardware mismatch:\n" + "\n".join(diffs))
     return hw_st

extra/assembly/amd/test/hw/test_ds.py

@@ -719,5 +719,47 @@ class TestAtomicOrdering(unittest.TestCase):
     self.assertEqual(st.vgpr[0][4], 150, "Final value should be 150")
 
 
+class TestDsPermute(unittest.TestCase):
+  """Tests for DS_PERMUTE_B32 and DS_BPERMUTE_B32 instructions."""
+
+  def test_ds_permute_b32_identity(self):
+    """DS_PERMUTE_B32 with identity permutation (lane 0 sends to lane 0)."""
+    # For simplicity, test with single lane
+    instructions = [
+      v_mov_b32_e32(v[0], 0),  # addr = 0 (lane 0)
+      v_mov_b32_e32(v[1], 0xDEADBEEF),  # data
+      ds_permute_b32(v[2], v[0], v[1]),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # Lane 0 sends to lane 0, so lane 0 gets 0xDEADBEEF
+    self.assertEqual(st.vgpr[0][2], 0xDEADBEEF)
+
+  def test_ds_bpermute_b32_identity(self):
+    """DS_BPERMUTE_B32 with identity permutation (each lane reads from itself)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),  # addr = 0 (read from lane 0)
+      v_mov_b32_e32(v[1], 0xCAFEBABE),  # data in lane 0
+      ds_bpermute_b32(v[2], v[0], v[1]),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # Lane 0 reads from lane 0's v[1]
+    self.assertEqual(st.vgpr[0][2], 0xCAFEBABE)
+
+  def test_ds_permute_b32_broadcast(self):
+    """DS_PERMUTE_B32 broadcast - all lanes send to lane 0."""
+    # With 4 lanes, all sending to lane 0, highest lane wins
+    instructions = [
+      v_mov_b32_e32(v[0], 0),  # All lanes send to addr 0 (lane 0)
+      v_mov_b32_e32(v[1], 0x11111111),  # All lanes send same data
+      ds_permute_b32(v[2], v[0], v[1]),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=4)
+    # Lane 0 receives data (highest numbered active lane wins)
+    self.assertEqual(st.vgpr[0][2], 0x11111111)
+
+
 if __name__ == '__main__':
   unittest.main()

extra/assembly/amd/test/hw/test_sop.py

@@ -62,6 +62,7 @@ class TestBasicScalar(unittest.TestCase):
     st = run_program(instructions, n_lanes=1)
     self.assertEqual(st.sgpr[1], 0x80000000)
 
+  @skip_unless_gfx(11, 5, "SALU FP ops require gfx1150+")
   def test_s_fmamk_f32(self):
     """S_FMAMK_F32: D = S0 * literal + S1."""
     # 2.0 * 3.0 + 1.0 = 7.0
@@ -73,6 +74,7 @@ class TestBasicScalar(unittest.TestCase):
     st = run_program(instructions, n_lanes=1)
     self.assertEqual(st.sgpr[2], f2i(7.0))
 
+  @skip_unless_gfx(11, 5, "SALU FP ops require gfx1150+")
   def test_s_fmamk_f32_negative(self):
     """S_FMAMK_F32 with negative values."""
     # -2.0 * 4.0 + 10.0 = 2.0

extra/assembly/amd/test/hw/test_vop1.py

@@ -255,7 +255,6 @@ class TestF16Conversions(unittest.TestCase):
 
   def test_v_cvt_f16_f32_small(self):
     """V_CVT_F16_F32 converts small f32 value."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     instructions = [
       v_mov_b32_e32(v[0], 0.5),
       v_cvt_f16_f32_e32(v[1], v[0]),
@@ -293,7 +292,6 @@ class TestF16Conversions(unittest.TestCase):
 
   def test_v_cvt_f16_f32_reads_full_32bit_source(self):
     """V_CVT_F16_F32 must read full 32-bit f32 source."""
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], 0x3fc00000),  # f32 1.5
       v_mov_b32_e32(v[0], s[0]),
@@ -302,7 +300,7 @@ class TestF16Conversions(unittest.TestCase):
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][1]
     lo_bits = result & 0xffff
-    self.assertEqual(lo_bits, 0x3e00, f"Expected f16(1.5)=0x3e00, got 0x{lo_bits:04x} ({_f16(lo_bits)})")
+    self.assertEqual(lo_bits, 0x3e00, f"Expected f16(1.5)=0x3e00, got 0x{lo_bits:04x} ({f16(lo_bits)})")
 
   def test_v_cvt_i16_f16_zero(self):
     """V_CVT_I16_F16 converts f16 zero to i16 zero."""
@@ -696,7 +694,6 @@ class TestCvtF16Modifiers(unittest.TestCase):
 
   def test_v_cvt_f32_f16_abs_negative(self):
     """V_CVT_F32_F16 with |abs| on negative value."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_neg1 = f32_to_f16(-1.0)  # 0xbc00
     instructions = [
       s_mov_b32(s[0], f16_neg1),
@@ -709,7 +706,6 @@ class TestCvtF16Modifiers(unittest.TestCase):
 
   def test_v_cvt_f32_f16_abs_positive(self):
     """V_CVT_F32_F16 with |abs| on positive value (should stay positive)."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_2 = f32_to_f16(2.0)  # 0x4000
     instructions = [
       s_mov_b32(s[0], f16_2),
@@ -722,7 +718,6 @@ class TestCvtF16Modifiers(unittest.TestCase):
 
   def test_v_cvt_f32_f16_neg_positive(self):
     """V_CVT_F32_F16 with neg on positive value."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_2 = f32_to_f16(2.0)  # 0x4000
     instructions = [
       s_mov_b32(s[0], f16_2),
@@ -735,7 +730,6 @@ class TestCvtF16Modifiers(unittest.TestCase):
 
   def test_v_cvt_f32_f16_neg_negative(self):
     """V_CVT_F32_F16 with neg on negative value (double negative)."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_neg2 = f32_to_f16(-2.0)  # 0xc000
     instructions = [
       s_mov_b32(s[0], f16_neg2),
@@ -748,7 +742,6 @@ class TestCvtF16Modifiers(unittest.TestCase):
 
   def test_v_cvt_f16_f32_then_pack_for_wmma(self):
     """CVT F32->F16 followed by pack (common WMMA pattern)."""
-    from extra.assembly.amd.test.hw.helpers import _f16
     f32_val = 3.5
     instructions = [
       s_mov_b32(s[0], f2i(f32_val)),
@@ -757,8 +750,8 @@ class TestCvtF16Modifiers(unittest.TestCase):
       v_pack_b32_f16(v[2], v[1], v[1]),  # Pack same value
     ]
     st = run_program(instructions, n_lanes=1)
-    lo = _f16(st.vgpr[0][2] & 0xffff)
-    hi = _f16((st.vgpr[0][2] >> 16) & 0xffff)
+    lo = f16(st.vgpr[0][2] & 0xffff)
+    hi = f16((st.vgpr[0][2] >> 16) & 0xffff)
     self.assertAlmostEqual(lo, f32_val, places=1)
     self.assertAlmostEqual(hi, f32_val, places=1)
 
@@ -804,7 +797,6 @@ class TestConversionRounding(unittest.TestCase):
 
   def test_f16_to_f32_precision(self):
     """F16 to F32 conversion precision."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_val = f32_to_f16(1.5)
     instructions = [
       s_mov_b32(s[0], f16_val),
@@ -816,7 +808,6 @@ class TestConversionRounding(unittest.TestCase):
 
   def test_f16_denormal_to_f32(self):
     """F16 denormal converts to small positive f32."""
-    from extra.assembly.amd.test.hw.helpers import _f16
     f16_denorm = 0x0001  # Smallest positive f16 denormal
     instructions = [
       v_mov_b32_e32(v[0], f16_denorm),
@@ -1512,5 +1503,82 @@ class TestReciprocalF16(unittest.TestCase):
     self.assertAlmostEqual(result, 0.25, places=2, msg="1/4.0 should be 0.25")
 
 
+class TestCvtNormF16(unittest.TestCase):
+  """Tests for V_CVT_NORM_I16_F16 and V_CVT_NORM_U16_F16."""
+
+  def test_cvt_norm_i16_f16_positive(self):
+    """V_CVT_NORM_I16_F16: f16 1.0 -> i16 max (32767)."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(1.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_cvt_norm_i16_f16_e32(v[1], v[0]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1] & 0xffff
+    self.assertEqual(result, 32767)
+
+  def test_cvt_norm_i16_f16_negative(self):
+    """V_CVT_NORM_I16_F16: f16 -1.0 -> i16 -32767 (0x8001)."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(-1.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_cvt_norm_i16_f16_e32(v[1], v[0]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1] & 0xffff
+    self.assertEqual(result, 0x8001)  # -32767, hardware uses symmetric range
+
+  def test_cvt_norm_i16_f16_zero(self):
+    """V_CVT_NORM_I16_F16: f16 0.0 -> i16 0."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),
+      v_cvt_norm_i16_f16_e32(v[1], v[0]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1] & 0xffff
+    self.assertEqual(result, 0)
+
+  def test_cvt_norm_u16_f16_one(self):
+    """V_CVT_NORM_U16_F16: f16 1.0 -> u16 max (65535)."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(1.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_cvt_norm_u16_f16_e32(v[1], v[0]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1] & 0xffff
+    self.assertEqual(result, 65535)
+
+  def test_cvt_norm_u16_f16_half(self):
+    """V_CVT_NORM_U16_F16: f16 0.5 -> u16 ~32768."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(0.5)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_cvt_norm_u16_f16_e32(v[1], v[0]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1] & 0xffff
+    self.assertAlmostEqual(result, 32768, delta=1)
+
+
+class TestPermlane64(unittest.TestCase):
+  """Tests for V_PERMLANE64_B32 instruction (wave64 cross-half swap)."""
+
+  def test_v_permlane64_b32_is_nop_in_wave32(self):
+    """V_PERMLANE64_B32 is a NOP in wave32 mode.
+
+    Per AMD pcode: "if WAVE32 then s_nop(...) else ... endif"
+    The emulator runs in wave32 mode, so this instruction should not modify registers.
+    """
+    instructions = [
+      v_mov_b32_e32(v[0], 0xCAFEBABE),  # source
+      v_mov_b32_e32(v[1], 0x12345678),  # dest (should be preserved)
+      v_permlane64_b32_e32(v[1], v[0]),  # NOP in wave32
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # Dest register should be unchanged (NOP behavior in wave32)
+    self.assertEqual(st.vgpr[0][1], 0x12345678)
+
+
 if __name__ == '__main__':
   unittest.main()

extra/assembly/amd/test/hw/test_vop3.py

@@ -857,7 +857,6 @@ class TestF16Modifiers(unittest.TestCase):
 
   def test_v_fma_f16_inline_const_1_0(self):
     """V_FMA_F16: a*b + 1.0 should use f16 inline constant."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
     f16_a = f32_to_f16(0.325928)  # ~0x3537
     f16_b = f32_to_f16(-0.486572)  # ~0xb7c9
     instructions = [
@@ -868,13 +867,12 @@ class TestF16Modifiers(unittest.TestCase):
       v_fma_f16(v[4], v[4], v[6], 1.0),  # 1.0 is inline constant
     ]
     st = run_program(instructions, n_lanes=1)
-    result = _f16(st.vgpr[0][4] & 0xffff)
+    result = f16(st.vgpr[0][4] & 0xffff)
     expected = 0.325928 * (-0.486572) + 1.0
     self.assertAlmostEqual(result, expected, delta=0.01)
 
   def test_v_fma_f16_inline_const_0_5(self):
     """V_FMA_F16: a*b + 0.5 should use f16 inline constant."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
     f16_a = f32_to_f16(2.0)
     f16_b = f32_to_f16(3.0)
     instructions = [
@@ -885,13 +883,12 @@ class TestF16Modifiers(unittest.TestCase):
       v_fma_f16(v[2], v[0], v[1], 0.5),  # 0.5 is inline constant
     ]
     st = run_program(instructions, n_lanes=1)
-    result = _f16(st.vgpr[0][2] & 0xffff)
+    result = f16(st.vgpr[0][2] & 0xffff)
     expected = 2.0 * 3.0 + 0.5
     self.assertAlmostEqual(result, expected, delta=0.01)
 
   def test_v_fma_f16_inline_const_neg_1_0(self):
     """V_FMA_F16: a*b + (-1.0) should use f16 inline constant."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
     f16_a = f32_to_f16(2.0)
     f16_b = f32_to_f16(3.0)
     instructions = [
@@ -902,13 +899,12 @@ class TestF16Modifiers(unittest.TestCase):
       v_fma_f16(v[2], v[0], v[1], -1.0),  # -1.0 is inline constant
     ]
     st = run_program(instructions, n_lanes=1)
-    result = _f16(st.vgpr[0][2] & 0xffff)
+    result = f16(st.vgpr[0][2] & 0xffff)
     expected = 2.0 * 3.0 + (-1.0)
     self.assertAlmostEqual(result, expected, delta=0.01)
 
   def test_v_add_f16_abs_both(self):
     """V_ADD_F16 with abs on both operands."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
     f16_neg2 = f32_to_f16(-2.0)
     f16_neg3 = f32_to_f16(-3.0)
     instructions = [
@@ -919,12 +915,11 @@ class TestF16Modifiers(unittest.TestCase):
       v_add_f16_e64(v[2], abs(v[0]), abs(v[1])),  # |-2| + |-3| = 5
     ]
     st = run_program(instructions, n_lanes=1)
-    result = _f16(st.vgpr[0][2] & 0xffff)
+    result = f16(st.vgpr[0][2] & 0xffff)
     self.assertAlmostEqual(result, 5.0, delta=0.01)
 
   def test_v_mul_f16_neg_abs(self):
     """V_MUL_F16 with neg on one operand and abs on another."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
     f16_2 = f32_to_f16(2.0)
     f16_neg3 = f32_to_f16(-3.0)
     instructions = [
@@ -935,7 +930,7 @@ class TestF16Modifiers(unittest.TestCase):
       v_mul_f16_e64(v[2], -v[0], abs(v[1])),  # -(2) * |-3| = -6
     ]
     st = run_program(instructions, n_lanes=1)
-    result = _f16(st.vgpr[0][2] & 0xffff)
+    result = f16(st.vgpr[0][2] & 0xffff)
     self.assertAlmostEqual(result, -6.0, delta=0.01)
 
   def test_v_fmac_f16_hi_dest(self):
@@ -943,7 +938,6 @@ class TestF16Modifiers(unittest.TestCase):
 
     This tests the case from AMD_LLVM sin(0) where V_FMAC_F16 writes to v0.h.
     """
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], 0x38003c00),  # v0 = {hi=0.5, lo=1.0}
       v_mov_b32_e32(v[0], s[0]),
@@ -954,8 +948,8 @@ class TestF16Modifiers(unittest.TestCase):
     ]
     st = run_program(instructions, n_lanes=1)
     v0 = st.vgpr[0][0]
-    result_hi = _f16((v0 >> 16) & 0xffff)
-    result_lo = _f16(v0 & 0xffff)
+    result_hi = f16((v0 >> 16) & 0xffff)
+    result_lo = f16(v0 & 0xffff)
     self.assertAlmostEqual(result_hi, 0.5, delta=0.01, msg=f"Expected hi=0.5, got {result_hi}")
     self.assertAlmostEqual(result_lo, 1.0, delta=0.01, msg=f"Expected lo=1.0, got {result_lo}")
 
@@ -2955,5 +2949,394 @@ class TestVOP3Clamp(unittest.TestCase):
     self.assertAlmostEqual(i2f(st.vgpr[3][1]), 1.0, places=5, msg="lane 3: 2.5 should clamp to 1.0")
 
 
+class TestCvtPkF16(unittest.TestCase):
+  """Tests for V_CVT_PK_RTZ_F16_F32 - pack two f32 to f16 with round toward zero."""
+
+  def test_cvt_pk_rtz_f16_f32_basic(self):
+    """V_CVT_PK_RTZ_F16_F32: basic pack of two f32 values."""
+    instructions = [
+      v_mov_b32_e32(v[0], 1.0),
+      v_mov_b32_e32(v[1], 2.0),
+      v_cvt_pk_rtz_f16_f32_e64(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo_f16 = f16(result & 0xffff)
+    hi_f16 = f16((result >> 16) & 0xffff)
+    self.assertAlmostEqual(lo_f16, 1.0, delta=0.01)
+    self.assertAlmostEqual(hi_f16, 2.0, delta=0.01)
+
+
+class TestCvtPkNorm(unittest.TestCase):
+  """Tests for V_CVT_PK_NORM_I16_F32 and V_CVT_PK_NORM_U16_F32."""
+
+  def test_cvt_pk_norm_i16_f32_basic(self):
+    """V_CVT_PK_NORM_I16_F32: pack two f32 to normalized i16."""
+    instructions = [
+      v_mov_b32_e32(v[0], 1.0),
+      v_mov_b32_e32(v[1], -1.0),
+      v_cvt_pk_norm_i16_f32(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    self.assertEqual(lo, 32767)
+    self.assertEqual(hi, 0x8001)  # -32767, hardware uses symmetric range
+
+  def test_cvt_pk_norm_u16_f32_basic(self):
+    """V_CVT_PK_NORM_U16_F32: pack two f32 to normalized u16."""
+    instructions = [
+      v_mov_b32_e32(v[0], 1.0),
+      v_mov_b32_e32(v[1], 0.5),
+      v_cvt_pk_norm_u16_f32(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    self.assertEqual(lo, 65535)
+    self.assertAlmostEqual(hi, 32768, delta=1)
+
+
+class TestCvtPkInt(unittest.TestCase):
+  """Tests for V_CVT_PK_I16_I32, V_CVT_PK_U16_U32, V_CVT_PK_I16_F32, V_CVT_PK_U16_F32."""
+
+  def test_cvt_pk_i16_i32_basic(self):
+    """V_CVT_PK_I16_I32: pack two i32 to i16."""
+    instructions = [
+      s_mov_b32(s[0], 100),
+      s_mov_b32(s[1], -100 & 0xffffffff),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_cvt_pk_i16_i32(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    lo_signed = lo if lo < 32768 else lo - 65536
+    hi_signed = hi if hi < 32768 else hi - 65536
+    self.assertEqual(lo_signed, 100)
+    self.assertEqual(hi_signed, -100)
+
+  def test_cvt_pk_u16_u32_basic(self):
+    """V_CVT_PK_U16_U32: pack two u32 to u16."""
+    instructions = [
+      s_mov_b32(s[0], 1000),
+      s_mov_b32(s[1], 2000),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_cvt_pk_u16_u32(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    self.assertEqual(lo, 1000)
+    self.assertEqual(hi, 2000)
+
+  def test_cvt_pk_i16_f32_basic(self):
+    """V_CVT_PK_I16_F32: convert two f32 to packed i16."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100.5),
+      v_mov_b32_e32(v[1], -50.7),
+      v_cvt_pk_i16_f32(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    lo_signed = lo if lo < 32768 else lo - 65536
+    hi_signed = hi if hi < 32768 else hi - 65536
+    self.assertEqual(lo_signed, 100)
+    self.assertEqual(hi_signed, -50)
+
+  def test_cvt_pk_u16_f32_basic(self):
+    """V_CVT_PK_U16_F32: convert two f32 to packed u16."""
+    instructions = [
+      v_mov_b32_e32(v[0], 100.9),
+      v_mov_b32_e32(v[1], 200.1),
+      v_cvt_pk_u16_f32(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    self.assertEqual(lo, 100)
+    self.assertEqual(hi, 200)
+
+  def test_cvt_pk_u8_f32_basic(self):
+    """V_CVT_PK_U8_F32: convert f32 to u8 and pack at byte position."""
+    instructions = [
+      v_mov_b32_e32(v[0], 128.5),
+      v_mov_b32_e32(v[1], 0),
+      v_mov_b32_e32(v[2], 0),
+      v_cvt_pk_u8_f32(v[2], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    byte0 = result & 0xff
+    self.assertEqual(byte0, 128)
+
+
+class TestDotProduct(unittest.TestCase):
+  """Tests for dot product instructions V_DOT4_U32_U8, V_DOT8_U32_U4."""
+
+  def test_v_dot4_u32_u8_basic(self):
+    """V_DOT4_U32_U8: 4-element dot product of u8 vectors."""
+    src0 = 0x04030201  # {4, 3, 2, 1}
+    src1 = 0x01010101  # {1, 1, 1, 1}
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot4_u32_u8(v[2], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    self.assertEqual(result, 10)
+
+  def test_v_dot4_u32_u8_with_accumulator(self):
+    """V_DOT4_U32_U8 with non-zero accumulator."""
+    src0 = 0x02020202  # {2, 2, 2, 2}
+    src1 = 0x03030303  # {3, 3, 3, 3}
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 100),
+      v_dot4_u32_u8(v[2], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    self.assertEqual(result, 124)
+
+  def test_v_dot8_u32_u4_basic(self):
+    """V_DOT8_U32_U4: 8-element dot product of u4 vectors."""
+    # src0 = 8 nibbles: {1,2,3,4,5,6,7,8} packed as 0x87654321
+    # src1 = 8 nibbles: {1,1,1,1,1,1,1,1} packed as 0x11111111
+    # result = 1+2+3+4+5+6+7+8 = 36
+    src0 = 0x87654321
+    src1 = 0x11111111
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot8_u32_u4(v[2], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    self.assertEqual(result, 36)
+
+
+class TestMinMaxF16Vop3(unittest.TestCase):
+  """Tests for V_MIN3_F16, V_MAX3_F16, V_MED3_F16, V_MINMAX_F16, V_MAXMIN_F16."""
+
+  def test_v_min3_f16_basic(self):
+    """V_MIN3_F16: minimum of three f16 values."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(3.0)),
+      s_mov_b32(s[1], f32_to_f16(1.0)),
+      s_mov_b32(s[2], f32_to_f16(2.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_min3_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 1.0, delta=0.01)
+
+  def test_v_max3_f16_basic(self):
+    """V_MAX3_F16: maximum of three f16 values."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(1.0)),
+      s_mov_b32(s[1], f32_to_f16(3.0)),
+      s_mov_b32(s[2], f32_to_f16(2.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_max3_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 3.0, delta=0.01)
+
+  def test_v_med3_f16_basic(self):
+    """V_MED3_F16: median of three f16 values."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(3.0)),
+      s_mov_b32(s[1], f32_to_f16(1.0)),
+      s_mov_b32(s[2], f32_to_f16(2.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_med3_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 2.0, delta=0.01)
+
+  def test_v_minmax_f16_basic(self):
+    """V_MINMAX_F16: clamp(src0, min=src1, max=src2)."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(2.5)),
+      s_mov_b32(s[1], f32_to_f16(1.0)),
+      s_mov_b32(s[2], f32_to_f16(2.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_minmax_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 2.0, delta=0.01)
+
+  def test_v_maxmin_f16_basic(self):
+    """V_MAXMIN_F16: clamp(src0, min=src2, max=src1)."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(0.5)),
+      s_mov_b32(s[1], f32_to_f16(2.0)),
+      s_mov_b32(s[2], f32_to_f16(1.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_maxmin_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 1.0, delta=0.01)
+
+  def test_v_min3_f16_with_neg(self):
+    """V_MIN3_F16 with neg modifier: min(-3, 1, 2) = -3."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(3.0)),
+      s_mov_b32(s[1], f32_to_f16(1.0)),
+      s_mov_b32(s[2], f32_to_f16(2.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_min3_f16(v[3], -v[0], v[1], v[2]),  # neg on first operand
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, -3.0, delta=0.01)
+
+  def test_v_max3_f16_with_abs(self):
+    """V_MAX3_F16 with abs modifier: max(|-3|, 1, 2) = 3."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(-3.0)),
+      s_mov_b32(s[1], f32_to_f16(1.0)),
+      s_mov_b32(s[2], f32_to_f16(2.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      v_max3_f16(v[3], abs(v[0]), v[1], v[2]),  # abs on first operand
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 3.0, delta=0.01)
+
+  def test_v_med3_f16_opsel_hi(self):
+    """V_MED3_F16 with opsel reading from hi half."""
+    # Pack two f16 values: hi=5.0, lo=1.0
+    packed = (f32_to_f16(5.0) << 16) | f32_to_f16(1.0)
+    instructions = [
+      s_mov_b32(s[0], packed),
+      s_mov_b32(s[1], f32_to_f16(3.0)),
+      s_mov_b32(s[2], f32_to_f16(4.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], s[2]),
+      # Read hi half of v[0] (5.0), med3(5, 3, 4) = 4
+      v_med3_f16(v[3], v[0].h, v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 4.0, delta=0.01)
+
+
+class TestSadHi(unittest.TestCase):
+  """Tests for V_SAD_HI_U8 instruction."""
+
+  def test_v_sad_hi_u8_basic(self):
+    """V_SAD_HI_U8: (sad << 16) + acc."""
+    # |1-5| + |2-6| + |3-7| + |4-8| = 16, << 16 = 0x100000, + 100 = 0x100064
+    instructions = [
+      v_mov_b32_e32(v[0], 0x04030201),
+      v_mov_b32_e32(v[1], 0x08070605),
+      v_mov_b32_e32(v[2], 100),
+      v_sad_hi_u8(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], (16 << 16) + 100)
+
+  def test_v_sad_hi_u8_zero_diff(self):
+    """V_SAD_HI_U8: identical inputs gives acc only."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x12345678),
+      v_mov_b32_e32(v[2], 50),
+      v_sad_hi_u8(v[3], v[0], v[0], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], 50)
+
+
+class TestPermlane(unittest.TestCase):
+  """Tests for V_PERMLANE16_B32 and V_PERMLANEX16_B32 instructions."""
+
+  def test_v_permlane16_b32_identity(self):
+    """V_PERMLANE16_B32 with identity permutation (lane i reads from lane i within row)."""
+    # lanesel encodes 4 bits per position: position i gets lanesel[i*4+3:i*4]
+    # Identity: position 0->0, 1->1, ..., 15->15
+    # lanesel = 0xFEDCBA9876543210 (positions 15-0 in nibbles)
+    instructions = [
+      v_mov_b32_e32(v[0], 0xDEADBEEF),  # source data
+      s_mov_b32(s[0], 0x76543210),       # lanesel low (positions 0-7)
+      s_mov_b32(s[1], 0xFEDCBA98),       # lanesel high (positions 8-15)
+      v_permlane16_b32(v[1], v[0], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    # Lane 0 reads from lane 0 (position 0 -> lanesel[3:0] = 0)
+    self.assertEqual(st.vgpr[0][1], 0xDEADBEEF)
+
+  def test_v_permlane16_b32_broadcast(self):
+    """V_PERMLANE16_B32 broadcast lane 0 to all lanes in row."""
+    # lanesel = all zeros -> all positions read from lane 0 within row
+    instructions = [
+      v_mov_b32_e32(v[0], 0xCAFEBABE),  # source data
+      s_mov_b32(s[0], 0),                # lanesel low = 0 (all read lane 0)
+      s_mov_b32(s[1], 0),                # lanesel high = 0
+      v_permlane16_b32(v[1], v[0], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=4)
+    # All lanes read from lane 0 of their row
+    for lane in range(4):
+      self.assertEqual(st.vgpr[lane][1], 0xCAFEBABE)
+
+  def test_v_permlanex16_b32_identity(self):
+    """V_PERMLANEX16_B32 cross-row read with identity selection."""
+    # In wave32: row 0 (lanes 0-15) reads from row 1 (lanes 16-31) and vice versa
+    # With single lane in row 0, it reads from lane 0 of row 1 (lane 16)
+    # But lane 16 doesn't exist in 1-lane test, so use 32 lanes
+    instructions = [
+      v_mov_b32_e32(v[0], 0x11111111),  # All lanes have this initially
+      s_mov_b32(s[0], 0x76543210),       # lanesel low
+      s_mov_b32(s[1], 0xFEDCBA98),       # lanesel high
+      v_permlanex16_b32(v[1], v[0], s[0], s[1]),
+    ]
+    st = run_program(instructions, n_lanes=32)
+    # Lane 0 in row 0 reads from lane 0 of row 1 (lane 16)
+    self.assertEqual(st.vgpr[0][1], 0x11111111)
+    # Lane 16 in row 1 reads from lane 0 of row 0 (lane 0)
+    self.assertEqual(st.vgpr[16][1], 0x11111111)
+
+
 if __name__ == '__main__':
   unittest.main()

extra/assembly/amd/test/hw/test_vop3p.py

@@ -149,7 +149,6 @@ class TestFmaMix(unittest.TestCase):
 
   def test_v_fma_mix_f32_src2_f16_lo(self):
     """V_FMA_MIX_F32 with src2 as f16 from lo bits."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_2 = f32_to_f16(2.0)
     instructions = [
       s_mov_b32(s[0], f2i(1.0)),
@@ -166,7 +165,6 @@ class TestFmaMix(unittest.TestCase):
 
   def test_v_fma_mix_f32_src2_f16_hi(self):
     """V_FMA_MIX_F32 with src2 as f16 from hi bits."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_2 = f32_to_f16(2.0)
     val = (f16_2 << 16) | 0
     instructions = [
@@ -199,7 +197,6 @@ class TestFmaMix(unittest.TestCase):
 
   def test_v_fma_mix_f32_with_abs_f16_src2_lo(self):
     """V_FMA_MIX_F32 with abs modifier on f16 src2 (lo half). Regression test for sin(1.0) bug."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_neg1 = f32_to_f16(-1.0)  # 0xbc00
     instructions = [
       s_mov_b32(s[0], f2i(0.0)),  # src0 = 0.0 (f32)
@@ -217,7 +214,6 @@ class TestFmaMix(unittest.TestCase):
 
   def test_v_fma_mix_f32_with_neg_f16_src2_lo(self):
     """V_FMA_MIX_F32 with neg modifier on f16 src2 (lo half)."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_1 = f32_to_f16(1.0)  # 0x3c00
     instructions = [
       s_mov_b32(s[0], f2i(0.0)),  # src0 = 0.0 (f32)
@@ -235,7 +231,6 @@ class TestFmaMix(unittest.TestCase):
 
   def test_v_fma_mix_f32_with_abs_f16_src2_hi(self):
     """V_FMA_MIX_F32 with abs modifier on f16 src2 (hi half)."""
-    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_neg1 = f32_to_f16(-1.0)  # 0xbc00
     val = (f16_neg1 << 16) | 0  # -1.0 in hi, 0 in lo
     instructions = [
@@ -254,7 +249,6 @@ class TestFmaMix(unittest.TestCase):
 
   def test_v_fma_mixlo_f16(self):
     """V_FMA_MIXLO_F16 writes to low 16 bits of destination."""
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], f2i(2.0)),
       v_mov_b32_e32(v[0], s[0]),
@@ -267,14 +261,13 @@ class TestFmaMix(unittest.TestCase):
       VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
     ]
     st = run_program(instructions, n_lanes=1)
-    lo = _f16(st.vgpr[0][3] & 0xffff)
+    lo = f16(st.vgpr[0][3] & 0xffff)
     hi = (st.vgpr[0][3] >> 16) & 0xffff
     self.assertAlmostEqual(lo, 7.0, places=1)
     self.assertEqual(hi, 0xdead, f"hi should be preserved, got 0x{hi:04x}")
 
   def test_v_fma_mixlo_f16_all_f32_sources(self):
     """V_FMA_MIXLO_F16 with all f32 sources."""
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], f2i(1.0)),
       v_mov_b32_e32(v[0], s[0]),
@@ -286,13 +279,12 @@ class TestFmaMix(unittest.TestCase):
       VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
     ]
     st = run_program(instructions, n_lanes=1)
-    lo = _f16(st.vgpr[0][3] & 0xffff)
+    lo = f16(st.vgpr[0][3] & 0xffff)
     # 1*2+3 = 5
     self.assertAlmostEqual(lo, 5.0, places=1)
 
   def test_v_fma_mixlo_f16_sin_case(self):
     """V_FMA_MIXLO_F16 case from sin kernel."""
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], 0x3f800000),  # f32 1.0
       v_mov_b32_e32(v[3], s[0]),
@@ -305,7 +297,7 @@ class TestFmaMix(unittest.TestCase):
       VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[3], src1=s[6], src2=v[5], opsel=0, opsel_hi=0, opsel_hi2=0),
     ]
     st = run_program(instructions, n_lanes=1)
-    lo = _f16(st.vgpr[0][3] & 0xffff)
+    lo = f16(st.vgpr[0][3] & 0xffff)
     self.assertAlmostEqual(lo, -3.14159, delta=0.01)
 
 
@@ -314,7 +306,6 @@ class TestVOP3P(unittest.TestCase):
 
   def test_v_pk_add_f16_basic(self):
     """V_PK_ADD_F16 adds two packed f16 values."""
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], 0x40003c00),  # hi=2.0, lo=1.0
       s_mov_b32(s[1], 0x44004200),  # hi=4.0, lo=3.0
@@ -324,14 +315,13 @@ class TestVOP3P(unittest.TestCase):
     ]
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][2]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
     self.assertAlmostEqual(lo, 4.0, places=2)
     self.assertAlmostEqual(hi, 6.0, places=2)
 
   def test_v_pk_mul_f16_basic(self):
     """V_PK_MUL_F16 multiplies two packed f16 values."""
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], 0x42004000),  # hi=3.0, lo=2.0
       s_mov_b32(s[1], 0x45004400),  # hi=5.0, lo=4.0
@@ -341,14 +331,13 @@ class TestVOP3P(unittest.TestCase):
     ]
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][2]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
     self.assertAlmostEqual(lo, 8.0, places=1)
     self.assertAlmostEqual(hi, 15.0, places=1)
 
   def test_v_pk_fma_f16_basic(self):
     """V_PK_FMA_F16: D = A * B + C for packed f16."""
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], 0x42004000),  # A: hi=3.0, lo=2.0
       s_mov_b32(s[1], 0x45004400),  # B: hi=5.0, lo=4.0
@@ -360,8 +349,8 @@ class TestVOP3P(unittest.TestCase):
     ]
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][3]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
     self.assertAlmostEqual(lo, 9.0, places=1)   # 2*4+1
     self.assertAlmostEqual(hi, 16.0, places=0)  # 3*5+1
 
@@ -370,7 +359,6 @@ class TestVOP3P(unittest.TestCase):
     Inline constants for VOP3P are f16 values in the low 16 bits only.
     hi half of inline constant is 0, so hi result = v0.hi + 0 = 1.0.
     """
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], 0x3c003c00),  # packed f16: hi=1.0, lo=1.0
       v_mov_b32_e32(v[0], s[0]),
@@ -378,8 +366,8 @@ class TestVOP3P(unittest.TestCase):
     ]
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][1]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
     # lo = 1.0 + 1.0 = 2.0, hi = 1.0 + 0.0 = 1.0 (inline const hi half is 0)
     self.assertAlmostEqual(lo, 2.0, places=2)
     self.assertAlmostEqual(hi, 1.0, places=2)
@@ -388,7 +376,6 @@ class TestVOP3P(unittest.TestCase):
     """V_PK_MUL_F16 with inline constant POS_TWO (2.0).
     Inline constant has value only in low 16 bits, hi is 0.
     """
-    from extra.assembly.amd.test.hw.helpers import _f16
     # v0 = packed (3.0, 4.0), multiply by POS_TWO
     # lo = 3.0 * 2.0 = 6.0, hi = 4.0 * 0.0 = 0.0 (inline const hi is 0)
     instructions = [
@@ -398,8 +385,8 @@ class TestVOP3P(unittest.TestCase):
     ]
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][1]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
     self.assertAlmostEqual(lo, 6.0, places=1)
     self.assertAlmostEqual(hi, 0.0, places=1)
 
@@ -413,7 +400,6 @@ class TestWMMAF16(unittest.TestCase):
 
   def test_v_wmma_f16_16x16x16_f16_all_ones(self):
     """V_WMMA_F16_16X16X16_F16 with all ones produces 16.0 in f16."""
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = []
     instructions.append(s_mov_b32(s[0], 0x3c003c00))  # packed f16 1.0
     # Initialize A matrix in v[16:23] (8 regs)
@@ -432,13 +418,12 @@ class TestWMMAF16(unittest.TestCase):
     for lane in range(32):
       for reg in range(8):
         result = st.vgpr[lane][reg]
-        lo = _f16(result & 0xffff)
+        lo = f16(result & 0xffff)
         self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
         self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
 
   def test_v_wmma_f16_16x16x16_f16_with_accumulator(self):
     """V_WMMA_F16_16X16X16_F16 with non-zero accumulator."""
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = []
     instructions.append(s_mov_b32(s[0], 0x3c003c00))  # packed f16 1.0
     instructions.append(s_mov_b32(s[1], 0x4500))  # f16 5.0 in lo bits only
@@ -458,7 +443,7 @@ class TestWMMAF16(unittest.TestCase):
     for lane in range(32):
       for reg in range(8):
         result = st.vgpr[lane][reg]
-        lo = _f16(result & 0xffff)
+        lo = f16(result & 0xffff)
         self.assertAlmostEqual(lo, 21.0, places=0, msg=f"v[{reg}] lane {lane}: expected 21.0, got {lo}")
         self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
 
@@ -468,7 +453,6 @@ class TestWMMAF16(unittest.TestCase):
     Regression test: WMMA was using static register indices instead of dynamic.
     This test uses v[64:71] for A, v[80:87] for B, v[96:103] for C/D.
     """
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = []
     instructions.append(s_mov_b32(s[0], 0x3c003c00))  # packed f16 1.0
     # Initialize A matrix in v[64:71] (8 regs)
@@ -490,7 +474,7 @@ class TestWMMAF16(unittest.TestCase):
     for lane in range(32):
       for reg in range(8):
         result = st.vgpr[lane][reg]
-        lo = _f16(result & 0xffff)
+        lo = f16(result & 0xffff)
         self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
         self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
 
@@ -713,7 +697,6 @@ class TestPackedMixedSigns(unittest.TestCase):
 
   def test_pk_add_f16_mixed_signs(self):
     """V_PK_ADD_F16 with mixed positive/negative values."""
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], 0xc0003c00),  # packed: hi=-2.0, lo=1.0
       s_mov_b32(s[1], 0x3c003c00),  # packed: hi=1.0, lo=1.0
@@ -723,14 +706,13 @@ class TestPackedMixedSigns(unittest.TestCase):
     ]
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][2]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
     self.assertAlmostEqual(lo, 2.0, places=2)   # 1.0 + 1.0
     self.assertAlmostEqual(hi, -1.0, places=2)  # -2.0 + 1.0
 
   def test_pk_mul_f16_zero(self):
     """V_PK_MUL_F16 with zero."""
-    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], 0x40004000),  # packed: 2.0, 2.0
       s_mov_b32(s[1], 0x00000000),  # packed: 0.0, 0.0
@@ -743,5 +725,277 @@ class TestPackedMixedSigns(unittest.TestCase):
     self.assertEqual(result, 0x00000000, "2.0 * 0.0 should be 0.0")
 
 
+class TestDot2F32F16(unittest.TestCase):
+  """Tests for V_DOT2_F32_F16 - dot product of f16 pairs producing f32."""
+
+  def test_v_dot2_f32_f16_basic(self):
+    """V_DOT2_F32_F16: dot product of two packed f16 pairs -> f32."""
+    # src0 = {hi=2.0, lo=1.0}, src1 = {hi=4.0, lo=3.0}
+    # result = 1.0*3.0 + 2.0*4.0 + 0 = 3 + 8 = 11.0
+    src0 = (f32_to_f16(2.0) << 16) | f32_to_f16(1.0)
+    src1 = (f32_to_f16(4.0) << 16) | f32_to_f16(3.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot2_f32_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = i2f(st.vgpr[0][3])
+    self.assertAlmostEqual(result, 11.0, places=2)
+
+  def test_v_dot2_f32_f16_with_accumulator(self):
+    """V_DOT2_F32_F16 with non-zero f32 accumulator."""
+    # src0 = {hi=1.0, lo=1.0}, src1 = {hi=1.0, lo=1.0}, acc = 5.0
+    # result = 1.0*1.0 + 1.0*1.0 + 5.0 = 7.0
+    src0 = (f32_to_f16(1.0) << 16) | f32_to_f16(1.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], f2i(5.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[0]),  # same as src0
+      v_mov_b32_e32(v[2], s[1]),
+      v_dot2_f32_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = i2f(st.vgpr[0][3])
+    self.assertAlmostEqual(result, 7.0, places=2)
+
+  def test_v_dot2_f32_f16_negative_values(self):
+    """V_DOT2_F32_F16 with negative f16 values."""
+    # src0 = {hi=-2.0, lo=3.0}, src1 = {hi=1.0, lo=2.0}
+    # result = 3.0*2.0 + (-2.0)*1.0 + 0 = 6 - 2 = 4.0
+    # NOTE: Hardware DOT2 may have up to 1 ULP difference due to internal implementation
+    src0 = (f32_to_f16(-2.0) << 16) | f32_to_f16(3.0)
+    src1 = (f32_to_f16(1.0) << 16) | f32_to_f16(2.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot2_f32_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
+    ]
+    st = run_program(instructions, n_lanes=1, ulp_tolerance=1)
+    result = i2f(st.vgpr[0][3])
+    self.assertAlmostEqual(result, 4.0, places=2)
+
+
+class TestDot2F16F16(unittest.TestCase):
+  """Tests for V_DOT2_F16_F16 - dot product of f16 pairs producing f16."""
+
+  def test_v_dot2_f16_f16_basic(self):
+    """V_DOT2_F16_F16: dot product of two packed f16 pairs -> f16."""
+    # src0 = {hi=2.0, lo=1.0}, src1 = {hi=3.0, lo=2.0}
+    # result = 1.0*2.0 + 2.0*3.0 + 0 = 2 + 6 = 8.0 (f16)
+    src0 = (f32_to_f16(2.0) << 16) | f32_to_f16(1.0)
+    src1 = (f32_to_f16(3.0) << 16) | f32_to_f16(2.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot2_f16_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 8.0, places=1)
+
+  def test_v_dot2_f16_f16_with_accumulator(self):
+    """V_DOT2_F16_F16 with non-zero f16 accumulator."""
+    # src0 = {hi=1.0, lo=1.0}, src1 = {hi=1.0, lo=1.0}, acc = 3.0 (f16)
+    # result = 1.0*1.0 + 1.0*1.0 + 3.0 = 5.0 (f16)
+    src0 = (f32_to_f16(1.0) << 16) | f32_to_f16(1.0)
+    acc = f32_to_f16(3.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[2], acc),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[0]),  # same as src0
+      v_mov_b32_e32(v[2], s[2]),
+      v_dot2_f16_f16(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = f16(st.vgpr[0][3] & 0xffff)
+    self.assertAlmostEqual(result, 5.0, places=1)
+
+
+class TestSignedDotProducts(unittest.TestCase):
+  """Tests for V_DOT4_I32_IU8 and V_DOT8_I32_IU4 with signed inputs."""
+
+  def test_v_dot4_i32_iu8_signed_both(self):
+    """V_DOT4_I32_IU8 with both inputs signed (neg=0b011)."""
+    # src0 = {-1, -2, 3, 4} as i8 = {0xff, 0xfe, 0x03, 0x04}
+    # src1 = {1, 1, 1, 1} as i8
+    # result = (-1)*1 + (-2)*1 + 3*1 + 4*1 = -1 - 2 + 3 + 4 = 4
+    src0 = (0xff << 24) | (0xfe << 16) | (0x03 << 8) | 0x04  # -1, -2, 3, 4
+    src1 = 0x01010101  # 1, 1, 1, 1
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b011),  # both signed
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][3]
+    # Result is i32, interpret as signed
+    if result >= 0x80000000:
+      result = result - 0x100000000
+    self.assertEqual(result, 4)
+
+  def test_v_dot4_i32_iu8_src0_signed(self):
+    """V_DOT4_I32_IU8 with only src0 signed (neg=0b001)."""
+    # src0 = {-1, -1, -1, -1} as i8 = {0xff, 0xff, 0xff, 0xff}
+    # src1 = {2, 2, 2, 2} as u8
+    # result = (-1)*2 + (-1)*2 + (-1)*2 + (-1)*2 = -8
+    src0 = 0xffffffff  # -1, -1, -1, -1 (as i8)
+    src1 = 0x02020202  # 2, 2, 2, 2 (as u8)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b001),  # src0 signed
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][3]
+    if result >= 0x80000000:
+      result = result - 0x100000000
+    self.assertEqual(result, -8)
+
+  def test_v_dot4_i32_iu8_src1_signed(self):
+    """V_DOT4_I32_IU8 with only src1 signed (neg=0b010)."""
+    # src0 = {2, 2, 2, 2} as u8
+    # src1 = {-1, -1, -1, -1} as i8 = {0xff, 0xff, 0xff, 0xff}
+    # result = 2*(-1) + 2*(-1) + 2*(-1) + 2*(-1) = -8
+    src0 = 0x02020202  # 2, 2, 2, 2 (as u8)
+    src1 = 0xffffffff  # -1, -1, -1, -1 (as i8)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b010),  # src1 signed
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][3]
+    if result >= 0x80000000:
+      result = result - 0x100000000
+    self.assertEqual(result, -8)
+
+  def test_v_dot4_i32_iu8_unsigned_as_reference(self):
+    """V_DOT4_I32_IU8 with both unsigned (neg=0) - same as V_DOT4_U32_U8."""
+    # src0 = {0xff, 0xff, 0xff, 0xff} = 255 each as u8
+    # src1 = {1, 1, 1, 1}
+    # result = 255*1 + 255*1 + 255*1 + 255*1 = 1020
+    src0 = 0xffffffff
+    src1 = 0x01010101
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0),  # both unsigned
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], 1020)
+
+  def test_v_dot8_i32_iu4_signed_both(self):
+    """V_DOT8_I32_IU4 with both inputs signed (neg=0b011)."""
+    # src0 = 8 nibbles: {-1, -2, 3, 4, -1, -2, 3, 4} as i4
+    # i4 -1 = 0xf, -2 = 0xe, 3 = 0x3, 4 = 0x4
+    # src0 = 0xfe34fe34
+    # src1 = {1, 1, 1, 1, 1, 1, 1, 1} as i4 = 0x11111111
+    # result = 2 * ((-1)*1 + (-2)*1 + 3*1 + 4*1) = 2 * 4 = 8
+    src0 = 0xfe34fe34
+    src1 = 0x11111111
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot8_i32_iu4(v[3], v[0], v[1], v[2], neg=0b011),  # both signed
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][3]
+    if result >= 0x80000000:
+      result = result - 0x100000000
+    self.assertEqual(result, 8)
+
+  def test_v_dot8_i32_iu4_all_negative(self):
+    """V_DOT8_I32_IU4 with all negative signed values."""
+    # src0 = 8 nibbles all -1 (0xf) = 0xffffffff
+    # src1 = 8 nibbles all 1 = 0x11111111
+    # result = 8 * ((-1)*1) = -8
+    src0 = 0xffffffff  # all -1 as i4
+    src1 = 0x11111111  # all 1
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_mov_b32_e32(v[2], 0),
+      v_dot8_i32_iu4(v[3], v[0], v[1], v[2], neg=0b011),  # both signed
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][3]
+    if result >= 0x80000000:
+      result = result - 0x100000000
+    self.assertEqual(result, -8)
+
+
+class TestPkMinMaxF16(unittest.TestCase):
+  """Tests for V_PK_MIN_F16 and V_PK_MAX_F16."""
+
+  def test_v_pk_min_f16_basic(self):
+    """V_PK_MIN_F16: packed min of two f16 pairs."""
+    # src0 = {hi=3.0, lo=1.0}, src1 = {hi=2.0, lo=4.0}
+    # result = {min(3,2)=2, min(1,4)=1}
+    src0 = (f32_to_f16(3.0) << 16) | f32_to_f16(1.0)
+    src1 = (f32_to_f16(2.0) << 16) | f32_to_f16(4.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_pk_min_f16(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
+    self.assertAlmostEqual(lo, 1.0, delta=0.01)
+    self.assertAlmostEqual(hi, 2.0, delta=0.01)
+
+  def test_v_pk_max_f16_basic(self):
+    """V_PK_MAX_F16: packed max of two f16 pairs."""
+    # src0 = {hi=3.0, lo=1.0}, src1 = {hi=2.0, lo=4.0}
+    # result = {max(3,2)=3, max(1,4)=4}
+    src0 = (f32_to_f16(3.0) << 16) | f32_to_f16(1.0)
+    src1 = (f32_to_f16(2.0) << 16) | f32_to_f16(4.0)
+    instructions = [
+      s_mov_b32(s[0], src0),
+      s_mov_b32(s[1], src1),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[1], s[1]),
+      v_pk_max_f16(v[2], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][2]
+    lo = f16(result & 0xffff)
+    hi = f16((result >> 16) & 0xffff)
+    self.assertAlmostEqual(lo, 4.0, delta=0.01)
+    self.assertAlmostEqual(hi, 3.0, delta=0.01)
+
+
 if __name__ == '__main__':
   unittest.main()

extra/assembly/amd/test/test_emu2_pcode.py

@@ -294,7 +294,7 @@ class TestAllPcode(unittest.TestCase):
             'ADDR': u32(), 'ADDR_BASE': u32(), 'TADDR': u32(), 'DATA': u32(), 'DATA0': u32(), 'DATA1': u32(), 'DATA2': u32(),
             'VDATA': u32(), 'VDATA0': u32(), 'VDATA1': u32(), 'VDATA2': u32(), 'VDATA3': u32(),
             'OPSEL': u32(), 'OPSEL_HI': u32(), 'NEG': u32(), 'NEG_HI': u32(), 'CLAMP': u32(),
-            'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'WAVE_STATUS': u32(),
+            'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'ROUND_TOWARD_ZERO': u32(), 'ROUND_NEAREST_EVEN': u32(), 'WAVE_STATUS': u32(),
             'MAX_FLOAT_F32': u32(0x7f7fffff), 'Unsigned': u32(1), 'clampedLOD': u32(),
             '_lds': lds, '_vmem': lds, '_active': UOp.const(dtypes.bool, True)}
 

extra/gemm/amd_asm_matmul.py

@@ -471,7 +471,7 @@ THREADS = 128
 
 def test_matmul():
   dev = Device[Device.DEFAULT]
-  print(f"Device arch: {dev.arch}")
+  print(f"Device arch: {dev.renderer.arch}")
 
   if getenv("STOCK", 0):
     # Load the stock kernel from amd_seb/kernel8_batched_gmem.s
@@ -479,7 +479,7 @@ def test_matmul():
     asm = stock_path.read_text()
     print(f"Loaded stock kernel from {stock_path}")
   else:
-    asm = build_kernel(dev.arch)
+    asm = build_kernel(dev.renderer.arch)
 
   binary = dev.compiler.compile(asm)
   print(f"Compiled! Binary size: {len(binary)} bytes")

extra/gemm/asm/cdna/gemm.py

@@ -0,0 +1,95 @@
+import atexit, functools
+from tinygrad.runtime.support.compiler_amd import HIPCompiler
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.renderer import Estimates
+from tinygrad.helpers import getenv, all_same, dedup
+from extra.gemm.asm.cdna.asm import build_kernel, GEMM_ARGS
+
+# ** CDNA4 assembly gemm
+
+WORKGROUP_SIZE = 256
+
+def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str, arch:str, wg:int) -> UOp:
+  batch, M, K = A.shape
+  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
+  assert K == K2
+  lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
+  gidx = UOp.special(wg, "gidx0")
+  k = build_kernel(batch, M, N, K, A.dtype.base)
+  sink = UOp.sink(C.base, A.base, B.base, lidx, gidx,
+                  arg=KernelInfo(name=k.name, estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
+  binary = HIPCompiler(arch).compile(k.to_asm())
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=k.to_text()), UOp(Ops.BINARY, arg=binary)))
+
+counters = {"used":0, "todos":[]}
+def todo(msg:str) -> bool: counters["todos"].append(msg); return False
+atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used'))
+
+def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
+  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
+  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
+  # only sharding on the batch is tested, others might work too
+  if isinstance(a.device, tuple) and not (a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None):
+    return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
+  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
+  N = b.shape[1]
+  if isinstance(a.device, tuple): batch //= len(a.device)
+  if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
+  if (key:=(M, N, K)) not in GEMM_ARGS: return todo(f"GEMM shape not supported {key}")
+  return True
+
+# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
+# note: this can be removed after we have GEMM on mixins
+
+def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
+  M, K = A.shape[0]*A.shape[1], A.shape[2]
+  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
+  assert K == K2
+  m = UOp.range(M, 1, AxisType.LOOP)
+  n = UOp.range(N, 2, AxisType.LOOP)
+  k = UOp.range(K, 0, AxisType.REDUCE)
+  mul = (A.index((m*UOp.const(dtypes.index, K)+k))*B.index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
+  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
+  store = C.index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
+  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
+
+# ** backward gemm, might use the asm gemm
+
+def custom_gemm_bw(gradient:UOp, kernel:UOp):
+  out, a, b = kernel.src
+  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)
+  grad_a = (g_t @ b_t.T).uop
+  a_T = a_t.transpose(-2, -1)
+  a_T = a_T.reshape(*a_T.shape[:-1], 1, a_T.shape[-1])
+  g_r = g_t.reshape(*g_t.shape[:-2], 1, *g_t.shape[-2:]).transpose(-1, -2)
+  grad_b = (a_T * g_r).sum((-1, 0)).uop
+  return (None, grad_a, grad_b)
+
+# ** main gemm function
+
+def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
+  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
+  counters["used"] += 1
+  squeeze = a.ndim == 2
+  if squeeze: a = a.unsqueeze(0)
+
+  batch, M, K = a.shape
+  N = b.shape[1]
+  is_multi = isinstance(a.device, tuple)
+
+  if is_multi:
+    out = Tensor(Tensor.empty(batch//len(a.device), M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+  else:
+    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
+
+  dname = a.device[0] if is_multi else a.device
+  arch = getattr(Device[dname].renderer, "arch", None)
+  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
+    numWG = GEMM_ARGS[(M, N, K)][0]
+    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=numWG, arch=arch), grad_fxn=custom_gemm_bw)[0]
+  else:
+    out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
+  return out.squeeze(0) if squeeze else out

extra/gemm/asm/cdna/template.s

@@ -1,78 +0,0 @@
-.text
-.section	.text.
-.global	gemm
-.p2align	8
-.type	gemm,@function
-
-gemm:
-INSTRUCTIONS
-
-.section .rodata,"a",@progbits
-.p2align 6, 0x0
-.amdhsa_kernel gemm
-  # basic memory requirements
-  .amdhsa_group_segment_fixed_size 133120
-  .amdhsa_private_segment_fixed_size 0
-  .amdhsa_kernarg_size 28
-  # register usage (RSRC1)
-  .amdhsa_next_free_vgpr 504
-  .amdhsa_next_free_sgpr 96
-  # workgroup / workitem IDs (RSRC2)
-  .amdhsa_system_sgpr_workgroup_id_x 1
-  .amdhsa_system_sgpr_workgroup_id_y 1
-  .amdhsa_system_sgpr_workgroup_id_z 1
-  # user SGPRs, we only specify the kernel args ptr in s[0:1]
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_user_sgpr_count 2
-  .amdhsa_user_sgpr_kernarg_preload_length 0
-  .amdhsa_user_sgpr_kernarg_preload_offset 0
-  # gfx90a / gfx940 specifics (RSRC3)
-  .amdhsa_accum_offset 248
-  .amdhsa_uses_dynamic_stack 0
-  .amdhsa_tg_split 0
-.end_amdhsa_kernel
-
-.amdgpu_metadata
----
-amdhsa.kernels:
-  - .name: gemm
-    .symbol: gemm.kd
-    .args:
-      - .name: C
-        .address_space: global
-        .offset: 0
-        .size: 8
-        .value_kind: global_buffer
-        .value_type: bf16
-      - .name: B
-        .address_space: global
-        .offset: 8
-        .size: 8
-        .value_kind: global_buffer
-        .value_type: bf16
-      - .name: A
-        .address_space: global
-        .offset: 16
-        .size: 8
-        .value_kind: global_buffer
-        .value_type: bf16
-      - .name: sz
-        .offset: 24
-        .size: 4
-        .value_kind: by_value
-        .value_type: u32
-    .group_segment_fixed_size: 133120
-    .private_segment_fixed_size: 0
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 28
-    .max_flat_workgroup_size: 256
-    .sgpr_count: 88
-    .sgpr_spill_count: 0
-    .vgpr_count: 248
-    .vgpr_spill_count: 0
-    .wavefront_size: 64
-amdhsa.version:
-  - 1
-  - 0
-...
-.end_amdgpu_metadata

extra/gemm/asm/cdna/test.py

@@ -1,73 +0,0 @@
-# Run assembly on the AMD runtime and check correctness
-# VIZ=2 to profile
-import pathlib
-from tinygrad import Tensor, Device, dtypes, Context
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.engine.realize import Estimates
-from tinygrad.helpers import getenv
-
-fp = pathlib.Path(__file__).parent/"gemm.s"
-
-N = getenv("N", 8192)
-THREADS_PER_WG = 256
-NUM_WG = N//THREADS_PER_WG * N//THREADS_PER_WG
-
-assert N % THREADS_PER_WG == 0, "N must be divisible by THREADS_PER_WG"
-
-# ** generate inputs on CPU
-
-scale = 10.0
-
-import torch
-torch.manual_seed(0)
-A = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
-B = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
-Bt = B.t().contiguous() # transpose B for the asm gemm
-C_torch = A@B
-
-# ** copy buffers to AMD
-
-# input creation and validation run on the copy engine for simpler tracing
-
-def from_torch(t:torch.Tensor) -> Tensor:
-  return Tensor.from_blob(t.data_ptr(), t.shape, dtype=dtypes.bfloat16, device="cpu").to(Device.DEFAULT).realize()
-
-C_tiny = from_torch(A) @ from_torch(B)
-C_asm = Tensor.empty_like(C_tiny)
-
-# ** assembly custom kernel
-
-def custom_asm_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
-  lidx = UOp.special(THREADS_PER_WG, "lidx0")
-  gidx = UOp.special(NUM_WG, "gidx0")
-
-  src = (pathlib.Path(__file__).parent/"template.s").read_text().replace("INSTRUCTIONS", fp.read_text())
-
-  sz = UOp.variable("SZ", 256, 8192)
-
-  sink = UOp.sink(C.base, A.base, B.base, sz, lidx, gidx, arg=KernelInfo(name="gemm", estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
-
-C_asm = Tensor.custom_kernel(C_asm, from_torch(A), from_torch(Bt), fxn=custom_asm_gemm)[0]
-
-# ** run gemms
-
-sched = Tensor.schedule(C_tiny, C_asm)
-eis = [si.lower() for si in sched]
-
-with Context(DEBUG=2):
-  for ei in eis:
-    et = ei.run({"SZ":N}, wait=True)
-    print(f"{(N*N*N*2 / et)*1e-12:.2f} REAL TFLOPS")
-
-# ** correctness
-
-import ctypes
-
-def torch_bf16(t:Tensor) -> torch.tensor:
-  asm_out = t.to("cpu").realize().uop.buffer._buf
-  buf = (ctypes.c_uint16*C_asm.uop.size).from_address(asm_out.va_addr)
-  return torch.frombuffer(buf, dtype=torch.bfloat16, count=C_asm.uop.size).reshape(C_asm.shape)
-
-assert torch.allclose(torch_bf16(C_asm), C_torch, rtol=1e-2, atol=1e-3)
-assert torch.allclose(torch_bf16(C_tiny), C_torch, rtol=1e-2, atol=1e-3)

extra/gemm/asm/cdna/test_asm_gemm.py

@@ -0,0 +1,46 @@
+import unittest
+from tinygrad import Tensor, Device, dtypes, Context
+from tinygrad.helpers import getenv
+from extra.gemm.asm.cdna.gemm import asm_gemm
+
+def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, multi=False) -> None:
+  Tensor.manual_seed(0)
+  a_rand = Tensor.randn((batch, M, K), dtype=dtypes.float).sub(0.5).cast(dtype)
+  b_rand = Tensor.randn((K, N), dtype=dtypes.float).sub(0.5).cast(dtype)
+  with Context(DEBUG=0):
+    Tensor.realize(a_rand, b_rand)
+
+  devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(8)) if multi else None
+
+  a, b = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
+  if multi: a, b = a.shard(devs, axis=0), b.shard(devs, axis=None)
+  tst = asm_gemm(a, b)
+  tst.sum().backward()
+  Tensor.realize(tst, a.grad, b.grad)
+
+  a_ref, b_ref = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
+  if multi: a_ref, b_ref = a_ref.shard(devs, axis=0), b_ref.shard(devs, axis=None)
+  with Context(ASM_GEMM=0): ref = a_ref @ b_ref
+  ref.sum().backward()
+  Tensor.realize(ref, a_ref.grad, b_ref.grad)
+
+  with Context(DEBUG=0):
+    assert (tst - ref).square().max().float().item() < 1e-6, "forward mismatch"
+    assert (a.grad - a_ref.grad).square().max().float().item() < 1e-3, "grad_a mismatch"
+    assert (b.grad - b_ref.grad).square().max().float().item() < 1e-3, "grad_b mismatch"
+
+class TestGemm(unittest.TestCase):
+  def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 4096), N, N, dtype=dtypes.half)
+
+  def test_gemm1(self): verify_asm_gemm(8, 8192, 4096, 14336, multi=True)
+  def test_gemm2(self): verify_asm_gemm(8, 8192, 128256, 4096, multi=True)
+  def test_gemm3(self): verify_asm_gemm(8, 8192, 14336, 4096, multi=True)
+  def test_gemm4(self): verify_asm_gemm(8, 4096, 14336, 4096, multi=True)
+  def test_gemm5(self): verify_asm_gemm(8, 4096, 4096, 14336, multi=True)
+  def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, multi=True)
+  def test_gemm_unsupported(self):
+    with self.assertRaisesRegex(AssertionError, "shape not supported"):
+      verify_asm_gemm(8, 8192, 1024, 4096, multi=True)
+
+if __name__ == "__main__":
+  unittest.main()

extra/usbgpu/patch.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
-import sys, os, zlib, struct, hashlib
-from tinygrad.helpers import DEBUG, getenv, fetch
+import os, zlib, struct, hashlib
+from tinygrad.helpers import getenv
 from tinygrad.runtime.support.usb import USB3
 
 SUPPORTED_CONTROLLERS = [
@@ -50,7 +50,7 @@ patched_fw = patch(file_path, file_hash, patches)
 dev = None
 for vendor, device in SUPPORTED_CONTROLLERS:
   try:
-    dev = USB3(vendor, device, 0x81, 0x83, 0x02, 0x04)
+    dev = USB3(vendor, device, 0x81, 0x83, 0x02, 0x04, use_bot=True)
     break
   except RuntimeError: pass
 if dev is None:

test/external/external_test_gpu_crash.py

@@ -1,126 +0,0 @@
-# ruff: noqa: F405
-"""Tests for GPU crash scenarios using AMD assembly to trigger invalid operations.
-
-These tests intentionally cause GPU faults to verify error handling.
-Run with: AMD=1 python -m pytest test/external/external_test_gpu_crash.py -v
-"""
-import unittest, re
-from tinygrad.device import Device
-from extra.assembly.amd.autogen.rdna3.ins import *  # noqa: F403
-from extra.assembly.amd.dsl import s, v, Inst, NULL
-
-def assemble(code:str, name:str="test") -> str:
-  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "wavefront_size32": 1, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
-  return f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{code}\n.rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + \
-         "\n".join(f".amdhsa_{k} {v}" for k,v in kd.items()) + "\n.end_amdhsa_kernel"
-
-@unittest.skipIf(Device.DEFAULT != "AMD", "AMD required")
-class TestGPUCrash(unittest.TestCase):
-  @classmethod
-  def setUpClass(cls):
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    cls.dev = Device["AMD"]
-    cls.compiler = HIPCompiler(cls.dev.arch)
-
-  def setUp(self):
-    # Verify device works before each test
-    from tinygrad import Tensor
-    try:
-      t = Tensor([1.0, 2.0], device="AMD").realize()
-      assert (t + 1).numpy().tolist() == [2.0, 3.0]
-    except Exception:
-      self.fail("Device not working before test")
-
-  def _run(self, code: str):
-    from tinygrad.runtime.ops_amd import AMDProgram
-    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code)))
-    prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)
-
-  def _run_insts(self, insts: list[Inst]): self._run("\n".join(i.disasm() for i in insts))
-
-  def _assert_gpu_fault(self, func):
-    """Assert that func raises a RuntimeError indicating a GPU fault (not a setup error)."""
-    with self.assertRaises(RuntimeError) as cm:
-      func()
-    err_msg = str(cm.exception).lower()
-    # Verify it's a GPU fault, not a setup/device initialization error
-    self.assertTrue(
-      re.search(r'fault|hang|timeout|illegal|memviol', err_msg),
-      f"Expected GPU fault error, got: {cm.exception}"
-    )
-
-
-class TestOutOfBoundsMemoryAccess(TestGPUCrash):
-  """Tests for out-of-bounds memory accesses."""
-
-  def test_global_load_null_ptr(self):
-    """Global load from NULL pointer."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
-             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_global_store_null_ptr(self):
-    """Global store to NULL pointer."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
-             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_global_load_unmapped_high_address(self):
-    """Global load from high unmapped address (0xDEAD00000000)."""
-    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD),
-             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_global_store_unmapped_high_address(self):
-    """Global store to high unmapped address."""
-    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678),
-             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_global_atomic_unmapped(self):
-    """Atomic operation on unmapped memory."""
-    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1),
-             global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-
-class TestSMEMFaults(TestGPUCrash):
-  """Tests for scalar memory (SMEM) faults."""
-
-  def test_smem_load_null(self):
-    """SMEM load from NULL base."""
-    insts = [s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
-             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_smem_load_unmapped(self):
-    """SMEM load from unmapped address."""
-    insts = [s_mov_b32(s[2], 0xBEEF0000), s_mov_b32(s[3], 0xDEAD),
-             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-
-class TestFlatMemoryFaults(TestGPUCrash):
-  """Tests for FLAT memory instruction faults."""
-
-  def test_flat_load_null(self):
-    """FLAT load from NULL address."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
-             flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_flat_store_null(self):
-    """FLAT store to NULL address."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
-             flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-  def test_flat_atomic_null(self):
-    """FLAT atomic on NULL address."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 1),
-             flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
-    self._assert_gpu_fault(lambda: self._run_insts(insts))
-
-
-if __name__ == "__main__":
-  unittest.main()

test/test_edgecases.py

@@ -195,8 +195,10 @@ class TestAssignIssues(unittest.TestCase):
     t.shrink(((1, 3), (1, 3))).assign(Tensor.ones(2, 2))
     np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())
 
+  @unittest.expectedFailure
   def test_assign_broadcast(self):
     # broadcasting during assign should behave like PyTorch
+    # NOTE: we don't want implicit dtype casting (int64 -> float32 loses precision), so this fails
     torch_tensor = torch.zeros(3, 5)
     torch_tensor[:] = torch.arange(5)
     t = Tensor.zeros(3, 5)

test/test_linearizer_failures.py

@@ -1,11 +1,21 @@
 # ruff: noqa: E501
 import unittest
+from tinygrad import Tensor
 from tinygrad.uop.ops import UOp, Ops, AxisType
 from tinygrad.dtype import dtypes
 from tinygrad.engine.realize import get_program
 from tinygrad.device import Device
 
 class TestLinearizerFailures(unittest.TestCase):
+  def test_cumsum_repeat_reshape_multiply(self):
+    # cumsum + repeat + reshape + multiply fails when step > 512
+    step, num_steps = 513, 10
+    t = Tensor.arange(step).float()
+    phase = t.cumsum()
+    tiled = phase.repeat((num_steps,)).reshape(num_steps, step)
+    pattern = Tensor([1,0,0,1,0,0,0,0,1,0]).reshape(num_steps, 1)
+    result = (tiled * pattern).flatten()
+    result.numpy()  # should not raise AssertionError in CFGContext
   def test_fail_1(self):
     c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=0, src=())
     c1 = UOp.range(UOp.const(dtypes.index, 2), 1, AxisType.LOOP)

test/test_multitensor.py

@@ -256,6 +256,18 @@ class TestMultiTensor(unittest.TestCase):
       a,b = _test_allreduce(Tensor.rand(256, 256))
       np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
 
+  def test_multiple_to_single_device_naive(self):
+    with Context(RING=0):
+      t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize()
+    self.assertEqual(t.device, Device.DEFAULT)
+    np.testing.assert_equal(t.numpy(), np.arange(32))
+
+  def test_multiple_to_single_device_ring(self):
+    with Context(RING=2):
+      t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize()
+    self.assertEqual(t.device, Device.DEFAULT)
+    np.testing.assert_equal(t.numpy(), np.arange(32))
+
   def test_allreduce_all2all(self):
     with Context(ALL2ALL=2):
       a,b = _test_allreduce(Tensor.rand(256, 256))
@@ -1273,19 +1285,20 @@ class TestMultiRamUsage(unittest.TestCase):
     _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).contiguous().realize()
     self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
 
-  def _test_matmul_half(self, devs):
+  def _test_matmul_half(self, dev_count:int):
     N = 32
     total_mem = {}
+    devs = tuple(f"NULL:{i}" for i in range(dev_count))
     for dtype in {dtypes.float, dtypes.half}:
       GlobalCounters.reset()
-      a = Tensor.empty((N, N), dtype=dtype).shard(devs, axis=0)
-      b = Tensor.empty((N, N), dtype=dtype).shard(devs, axis=None)
+      a = Tensor.empty((N, N), dtype=dtype, device=devs[0]).shard(devs, axis=0)
+      b = Tensor.empty((N, N), dtype=dtype, device=devs[0]).shard(devs, axis=None)
       (a @ b).realize()
       total_mem[dtype] = GlobalCounters.global_mem
     self.assertEqual(total_mem[dtypes.half], total_mem[dtypes.float] // 2)
 
-  def test_matmul_half(self): self._test_matmul_half(devices_2)
-  def test_matmul_half_alt(self): self._test_matmul_half(devices_4)
+  def test_matmul_half(self): self._test_matmul_half(dev_count=2)
+  def test_matmul_half_alt(self): self._test_matmul_half(dev_count=4)
 
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiFromUnrenderable(unittest.TestCase):

test/testextra/test_cfg_viz.py

@@ -2,7 +2,7 @@
 # allow define from star imports
 
 import unittest
-import textwrap, functools
+import functools
 
 from tinygrad import Device, Tensor
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
@@ -11,69 +11,27 @@ from tinygrad.runtime.support.compiler_amd import HIPCompiler
 from tinygrad.viz.serve import amdgpu_cfg
 
 from extra.assembly.amd.autogen.rdna3.ins import *
-from extra.assembly.amd.dsl import Inst
+from extra.assembly.amd.dsl import s
 
-template = """.text
-.globl fn_name
-.p2align 8
-.type fn_name,@function
-fn_name:
-  INSTRUCTION
-
-.rodata
-.p2align 6
-.amdhsa_kernel fn_name
-  .amdhsa_kernarg_size 8
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
-  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
-  .amdhsa_wavefront_size32 1
-.end_amdhsa_kernel
-
-.amdgpu_metadata
----
-amdhsa.version:
-  - 1
-  - 0
-amdhsa.kernels:
-  - .name: fn_name
-    .symbol: fn_name.kd
-    .group_segment_fixed_size: 0
-    .private_segment_fixed_size: 0
-    .wavefront_size: 32
-    .sgpr_count: 8
-    .vgpr_count: 8
-    .max_flat_workgroup_size: 1024
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 8
-    .args:
-      - .address_space:  global
-        .name:           a
-        .offset:         0
-        .size:           8
-        .type_name:      'float*'
-        .value_kind:     global_buffer
-...
-.end_amdgpu_metadata
-"""
+# TODO: this belongs to the dsl infrastructure
+from extra.gemm.amd_asm_matmul import Kernel
 
 # TODO: shouldn't need compiler once we can output ELF
 # outputs a text disassembly for humans and a machine readable binary
-def assemble(name:str, insts:list[str|Inst], compiler:Compiler) -> tuple[str, bytes]:
-  asm = "\n".join([inst if isinstance(inst, str) else inst.disasm() for inst in insts])
-  src = template.replace("fn_name", name).replace("INSTRUCTION", textwrap.dedent(asm))
+def assemble(name:str, k:Kernel, compiler:Compiler) -> tuple[str, bytes]:
+  src = k.to_asm()
   return (src, compiler.compile(src))
 
-def asm_kernel(out:UOp, insts:list[str|Inst], name:str, device:str, compiler:Compiler, n_threads:int=1, n_workgroups:int=1) -> UOp:
+def asm_kernel(out:UOp, k:Kernel, name:str, device:str, compiler:Compiler, n_threads:int=1, n_workgroups:int=1) -> UOp:
   lidx = UOp.special(n_threads, "lidx0")
   gidx = UOp.special(n_workgroups, "gidx0")
   sink = UOp.sink(out, lidx, gidx, arg=KernelInfo(name=name))
-  src, lib = assemble(name, insts, compiler)
+  src, lib = assemble(name, k, compiler)
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=device), UOp(Ops.LINEAR, src=(*sink.src, sink)),
                                UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
 
-def run_asm(name:str, insts:list) -> None:
-  fxn = functools.partial(asm_kernel, insts=insts, name=name, device=Device.DEFAULT, compiler=HIPCompiler(Device[Device.DEFAULT].renderer.arch))
+def run_asm(name:str, k:Kernel) -> None:
+  fxn = functools.partial(asm_kernel, k=k, name=name, device=Device.DEFAULT, compiler=HIPCompiler(Device[Device.DEFAULT].renderer.arch))
   out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
   out.realize()
 
@@ -85,32 +43,32 @@ class TestCfg(unittest.TestCase):
       self.skipTest(f"tests written for RDNA, got arch {arch}")
 
   def test_simple(self):
-    run_asm("simple", [
-      "entry:",
-        "s_branch bb1",
-      "bb1:",
-        s_endpgm(),
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_branch(), target="bb1")
+    k.label("bb1")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("simple", k)
 
   def test_diamond(self):
-    run_asm("diamond", insts:=[
-      "entry:",
-        s_mov_b32(s[0], 0),
-        s_mov_b32(s[1], 0),
-        s_cmp_eq_u64(s[0:1], 0),
-        "s_cbranch_scc1 if",
-        "s_branch else",
-      "if:",
-        s_nop(1),
-        "s_branch end",
-      "else:",
-        s_nop(0),
-      "end:",
-        s_endpgm(),
-        s_code_end(),
-    ])
-    _, lib = assemble("diamond", insts, HIPCompiler(Device[Device.DEFAULT].arch))
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[0], 0))
+    k.emit(s_mov_b32(s[1], 0))
+    k.emit(s_cmp_eq_u64(s[0:1], 0))
+    k.emit(s_cbranch_scc1(), target="if")
+    k.emit(s_branch(), target="else")
+    k.label("if")
+    k.emit(s_nop(1))
+    k.emit(s_branch(), target="end")
+    k.label("else")
+    k.emit(s_nop(0))
+    k.label("end")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("diamond", k)
+    _, lib = assemble("diamond", k, HIPCompiler(Device[Device.DEFAULT].arch))
     cfg = amdgpu_cfg(lib, Device[Device.DEFAULT].device_props()["gfx_target_version"])["data"]
     self.assertEqual(len(cfg["blocks"]), 5)
     edge_count = sum(len(v) for v in cfg["paths"].values())
@@ -124,133 +82,138 @@ class TestCfg(unittest.TestCase):
     self.assertEqual(insts, ['s_mov_b32', 's_cmp_eq_u64'])
 
   def test_loop(self):
-    run_asm("simple_loop", [
-      "entry:",
-        s_mov_b32(s[1], 4),
-      "loop:",
-        s_add_u32(s[1], s[1], -1),
-        s_cmp_eq_i32(s[1], 0),
-        "s_cbranch_scc0 loop",
-        s_endpgm(),
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 4))
+    k.label("loop")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_cbranch_scc0(), target="loop")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("simple_loop", k)
 
   def test_loop_branch(self):
-    run_asm("loop_if", [
-      "entry:",
-        s_mov_b32(s[1], 4),
-      "loop:",
-        s_add_u32(s[1], s[1], -1),
-        s_cmp_eq_i32(s[1], 2),
-        "s_cbranch_scc1 cond",
-        "s_branch cont",
-      "cond:",
-        s_add_u32(s[1], s[1], -2),
-      "cont:",
-        s_cmp_eq_i32(s[1], 0),
-        "s_cbranch_scc0 loop",
-        s_endpgm(),
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 4))
+    k.label("loop")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_cmp_eq_i32(s[1], 2))
+    k.emit(s_cbranch_scc1(), target="cond")
+    k.emit(s_branch(), target="cont")
+    k.label("cond")
+    k.emit(s_add_u32(s[1], s[1], -2))
+    k.label("cont")
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_cbranch_scc0(), target="loop")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("loop_if", k)
 
   def test_loop_break(self):
-    run_asm("loop_break", [
-      "entry:",
-        s_mov_b32(s[1], 8),
-      "loop:",
-        s_add_u32(s[1], s[1], -1),
-        s_cmp_eq_i32(s[1], 5),
-        "s_cbranch_scc1 break",
-        s_cmp_eq_i32(s[1], 0),
-        "s_cbranch_scc0 loop",
-      "break:",
-        s_endpgm(),
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 8))
+    k.label("loop")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_cmp_eq_i32(s[1], 5))
+    k.emit(s_cbranch_scc1(), target="break")
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_cbranch_scc0(), target="loop")
+    k.label("break")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("loop_break", k)
 
   def test_switch(self):
-    run_asm("switch_case", [
-      "entry:",
-        s_cmp_eq_i32(s[0], 0),
-        "s_cbranch_scc1 case0",
-        s_cmp_eq_i32(s[0], 1),
-        "s_cbranch_scc1 case1",
-        "s_branch case2",
-      "case0:",
-        s_nop(0),
-        "s_branch join",
-      "case1:",
-        s_nop(1),
-        "s_branch join",
-      "case2:",
-        s_nop(2),
-        "s_branch join",
-      "join:",
-        s_endpgm(),
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_cmp_eq_i32(s[0], 0))
+    k.emit(s_cbranch_scc1(), target="case0")
+    k.emit(s_cmp_eq_i32(s[0], 1))
+    k.emit(s_cbranch_scc1(), target="case1")
+    k.emit(s_branch(), target="case2")
+    k.label("case0")
+    k.emit(s_nop(0))
+    k.emit(s_branch(), target="join")
+    k.label("case1")
+    k.emit(s_nop(1))
+    k.emit(s_branch(), target="join")
+    k.label("case2")
+    k.emit(s_nop(2))
+    k.emit(s_branch(), target="join")
+    k.label("join")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("switch_case", k)
 
   def test_ping_pong(self):
-    run_asm("ping_pong", [
-      "entry:",
-        s_cmp_eq_i32(s[0], 0),
-        "s_cbranch_scc1 ping",
-        "s_branch pong",
-      "ping:",
-        s_cmp_eq_i32(s[1], 0),
-        "s_cbranch_scc1 pong",
-        "s_branch end",
-      "pong:",
-        s_cmp_eq_i32(s[2], 0),
-        "s_cbranch_scc1 ping",
-      "end:",
-        s_endpgm(),
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_cmp_eq_i32(s[0], 0))
+    k.emit(s_cbranch_scc1(), target="ping")
+    k.emit(s_branch(), target="pong")
+    k.label("ping")
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_cbranch_scc1(), target="pong")
+    k.emit(s_branch(), target="end")
+    k.label("pong")
+    k.emit(s_cmp_eq_i32(s[2], 0))
+    k.emit(s_cbranch_scc1(), target="ping")
+    k.label("end")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("ping_pong", k)
 
   def test_colored_blocks(self):
     N = 10
-    asm = ["entry:", "s_branch init0"]
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_branch(), target="init0")
     for i in range(N):
-      asm += [f"init{i}:", s_mov_b32(s[1], i + 1), f"s_branch {(loop:=f'loop{i}')}"]
-      asm += [
-        f"{loop}:",
-          s_nop(i & 7),
-          s_add_u32(s[1], s[1], -1),
-          s_cmp_eq_i32(s[1], 0),
-          f"s_cbranch_scc0 {loop}",
-          f"s_branch {'init' + str(i+1) if i + 1 < N else 'end'}",
-      ]
-    asm += ["end:", s_endpgm(), s_code_end()]
-    run_asm("test_colored_blocks", asm)
+      loop = f"loop{i}"
+      k.label(f"init{i}")
+      k.emit(s_mov_b32(s[1], i + 1))
+      k.emit(s_branch(), target=loop)
+      k.label(loop)
+      k.emit(s_nop(i & 7))
+      k.emit(s_add_u32(s[1], s[1], -1))
+      k.emit(s_cmp_eq_i32(s[1], 0))
+      k.emit(s_cbranch_scc0(), target=loop)
+      k.emit(s_branch(), target=f"init{i+1}" if i + 1 < N else "end")
+    k.label("end")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("test_colored_blocks", k)
 
   def test_jump_back_to_end(self):
-    run_asm("jump_back_to_end", [
-      "entry:",
-        s_mov_b32(s[1], 2),
-        "s_cbranch_execz loop",
-      "end:",
-        s_endpgm(),
-      "loop:",
-        s_add_u32(s[1], s[1], -1),
-        s_cmp_eq_i32(s[1], 0),
-        "s_branch end",
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 2))
+    k.emit(s_cbranch_execz(), target="loop")
+    k.label("end")
+    k.emit(s_endpgm())
+    k.label("loop")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_branch(), target="end")
+    k.emit(s_code_end())
+    run_asm("jump_back_to_end", k)
 
   def test_hit_count(self):
-    run_asm("test_hit_count", [
-      "entry:",
-        s_mov_b32(s[1], 1),
-        "s_branch alt",
-      "continue:",
-        s_mov_b32(s[2], 2),
-        s_add_u32(s[1], s[1], s[2]),
-      "alt:",
-        s_add_u32(s[1], s[1], -1),
-        s_endpgm(),
-        s_code_end(),
-    ])
+    k = Kernel(arch=Device["AMD"].arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 1))
+    k.emit(s_branch(), target="alt")
+    k.label("continue")
+    k.emit(s_mov_b32(s[2], 2))
+    k.emit(s_add_u32(s[1], s[1], s[2]))
+    k.label("alt")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    run_asm("test_hit_count", k)
 
 if __name__ == "__main__":
   unittest.main()

test/unit/test_assign.py

@@ -461,6 +461,27 @@ class TestAssign(unittest.TestCase):
     a[2:5] = [1, 2, 3]
     np.testing.assert_allclose(a.numpy(), [0., 0., 1., 2., 3., 0., 0., 0.])
 
+  def test_assign_bitcast(self):
+    # assign to a bitcast view should modify the underlying buffer
+    a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    # IEEE 754: 1.0f = 0x3f800000, 2.0f = 0x40000000, 3.0f = 0x40400000, 4.0f = 0x40800000
+    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32)).realize()
+    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
+    # double bitcast
+    b = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    b.bitcast(dtypes.uint32).bitcast(dtypes.int32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.int32)).realize()
+    np.testing.assert_allclose(b.numpy(), [4.0, 3.0, 2.0, 1.0])
+    # shrink then bitcast
+    c = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    c[0:2].bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000], dtype=dtypes.uint32)).realize()
+    np.testing.assert_allclose(c.numpy(), [4.0, 3.0, 3.0, 4.0])
+
+  def test_assign_bitcast_different_size(self):
+    # different-size bitcast creates a new tensor, not a view, so assign doesn't modify the original
+    a = Tensor([0]*8, dtype=dtypes.uint8).realize()
+    a.bitcast(dtypes.int64).assign(Tensor([12345], dtype=dtypes.int64)).realize()
+    np.testing.assert_equal(a.numpy(), [0]*8)
+
   @unittest.skip("don't use output buffer, and mismatch dtype no longer supported")
   def test_cast_assignment(self):
     a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
@@ -472,6 +493,38 @@ class TestAssign(unittest.TestCase):
     assert oba1 is None and oba2 is None
     np.testing.assert_allclose(a.numpy(), np.arange(N*N,dtype=np.int32).reshape((N,N)))
 
+  def test_assign_dtype_mismatch(self):
+    # assign should not implicitly cast dtypes - this can lose precision
+    a = Tensor.zeros(4, dtype=dtypes.float32).contiguous().realize()
+    b = Tensor([1, 2, 3, 4], dtype=dtypes.int32)
+    with self.assertRaisesRegex(RuntimeError, "assign dtype mismatch"):
+      a.assign(b)
+
+  def test_assign_dtype_mismatch_int64_to_float32(self):
+    # int64 -> float32 loses precision for large values, should not be implicit
+    a = Tensor.zeros(1, dtype=dtypes.float32).contiguous().realize()
+    b = Tensor([16777217], dtype=dtypes.int64)  # 2^24 + 1, not exactly representable in float32
+    with self.assertRaisesRegex(RuntimeError, "assign dtype mismatch"):
+      a.assign(b)
+
+  def test_assign_shape_broadcast(self):
+    # shape broadcasting should work when dtypes match
+    a = Tensor.zeros(3, 5, dtype=dtypes.float32).contiguous().realize()
+    b = Tensor([1., 2., 3., 4., 5.], dtype=dtypes.float32)
+    a.assign(b)
+    a.realize()
+    expected = np.array([[1., 2., 3., 4., 5.]] * 3)
+    np.testing.assert_allclose(a.numpy(), expected)
+
+  def test_assign_shape_broadcast_2d(self):
+    # broadcast (1, 5) to (3, 5)
+    a = Tensor.zeros(3, 5, dtype=dtypes.float32).contiguous().realize()
+    b = Tensor([[1., 2., 3., 4., 5.]], dtype=dtypes.float32)
+    a.assign(b)
+    a.realize()
+    expected = np.array([[1., 2., 3., 4., 5.]] * 3)
+    np.testing.assert_allclose(a.numpy(), expected)
+
   def test_disk_assignment(self):
     a = Tensor.empty(5, device=f"disk:{temp('disk_assignment')}").assign(Tensor.ones(5)).numpy()
     np.testing.assert_equal(a, np.ones(5))
@@ -566,12 +619,12 @@ class TestAssignOrdering(unittest.TestCase):
   def test_slice_write_then_full_read(self):
     """Write to slice, then read full buffer."""
     # without .realize(): orphan slice assign not triggered by .numpy()
-    buf = Tensor.zeros(4).contiguous().realize()
+    buf = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
     buf[1:3].assign(Tensor([5, 6]))
     np.testing.assert_equal(buf.numpy(), [0, 0, 0, 0])  # TODO: wrong! should be [0, 5, 6, 0]
 
     # with .realize(): assign executes
-    buf = Tensor.zeros(4).contiguous().realize()
+    buf = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
     buf[1:3].assign(Tensor([5, 6])).realize()
     np.testing.assert_equal(buf.numpy(), [0, 5, 6, 0])
 
@@ -653,7 +706,7 @@ class TestAssignOrdering(unittest.TestCase):
 
   def test_three_buffer_chain(self):
     """Chain: A depends on B, B depends on C - ordering matters."""
-    a = Tensor.zeros(4).contiguous().realize()
+    a = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
     b = Tensor([1, 2, 3, 4]).contiguous().realize()
     c = Tensor([10, 10, 10, 10]).contiguous().realize()
     # b reads from c, a reads from b
@@ -665,8 +718,8 @@ class TestAssignOrdering(unittest.TestCase):
 
   def test_interleaved_assign_read_patterns(self):
     """Complex interleaved pattern: write A, read A into B, write B, read B."""
-    a = Tensor.zeros(4).contiguous().realize()
-    b = Tensor.zeros(4).contiguous().realize()
+    a = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
+    b = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
 
     a.assign(Tensor([1, 2, 3, 4]))
     b.assign(a.contiguous())       # b should get [1,2,3,4]

test/unit/test_autogen.py

@@ -5,24 +5,13 @@ from tinygrad.runtime.support.c import DLL, record, init_records
 from tinygrad.runtime.support import c
 from tinygrad.runtime.support.autogen import gen
 
-class TestAutogen(unittest.TestCase):
+@unittest.skipIf(WIN, "doesn't compile on windows")
+class TestC(unittest.TestCase):
   def compile(self, src):
     with tempfile.NamedTemporaryFile(suffix=".so") as f:
       subprocess.check_output(('clang', '-x', 'c', '-fPIC', '-shared', '-', '-o', f.name), input=src.encode())
       return DLL("test", f.name)
 
-  def run_gen(self, contents):
-    with tempfile.NamedTemporaryFile(mode='w', suffix='.h') as f:
-      f.write(contents)
-      f.flush()
-
-      generated_code = gen(name="test_header", dll=None, files=[f.name])
-
-      namespace = {}
-      exec(generated_code, namespace)
-      return namespace
-
-  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_packed_struct(self):
     @record
     class Baz:
@@ -45,7 +34,6 @@ class TestAutogen(unittest.TestCase):
     assert b.c == 1
     assert b.d == 0
 
-  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_packed_struct_interop(self):
     @record
     class Baz:
@@ -75,7 +63,6 @@ class TestAutogen(unittest.TestCase):
     self.assertEqual(test(b), b.a + b.b + b.c + b.d)
 
   # https://github.com/python/cpython/issues/90914
-  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_bitfield_interop(self):
     @record
     class Baz:
@@ -103,7 +90,6 @@ class TestAutogen(unittest.TestCase):
     def test(x:Baz) -> ctypes.c_int: ...
     for i in range(8): self.assertEqual(test(Baz(*(j==i for j in range(8)))), i==2)
 
-  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_struct_interop(self):
     @record
     class Baz:
@@ -131,7 +117,6 @@ class TestAutogen(unittest.TestCase):
     def test(x:Baz) -> Baz: ...
     self.assertEqual(bytes(test(Baz(*range(8)))), struct.pack("8i", *range(7, -1, -1)))
 
-  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_aos_interop(self):
     @record
     class Item:
@@ -151,7 +136,6 @@ class TestAutogen(unittest.TestCase):
     def test(arr:(Item * 3)) -> ctypes.c_int: ...
     self.assertEqual(test((Item * 3)(Item(10), Item(20), Item(30))), 60)
 
-  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_soa_interop(self):
     @record
     class Row:
@@ -173,7 +157,6 @@ class TestAutogen(unittest.TestCase):
     self.assertEqual(r.data[1], 20)
     self.assertEqual(r.data[2], 10)
 
-  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_soa_ptr_interop(self):
     @record
     class Row:
@@ -191,7 +174,6 @@ class TestAutogen(unittest.TestCase):
     def test(x:Row) -> ctypes.c_int: ...
     assert test(Row((ctypes.c_int * 3)(10, 20, 30))) == 60
 
-  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_nested_struct_interop(self):
     @record
     class Inner:
@@ -217,7 +199,6 @@ class TestAutogen(unittest.TestCase):
     self.assertEqual(o.inner.a, 20)
     self.assertEqual(o.b, 10)
 
-  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_struct_pointer_interop(self):
     @record
     class Foo:
@@ -242,7 +223,88 @@ class TestAutogen(unittest.TestCase):
     self.assertEqual(out.contents.a, 20)
     self.assertEqual(out.contents.b, 10)
 
-  @unittest.skipIf(WIN, "doesn't compile on windows")
+  def test_pointer_field_roundtrip(self):
+    # This tests storing a pointer in a record struct field and passing it to C
+    # Mimics how mesa.struct_lp_build_tgsi_params.mask is used
+    from tinygrad.runtime.support.c import POINTER
+    @record
+    class Inner:
+      SIZE = 8
+      value: Annotated[ctypes.c_int, 0]
+      flag: Annotated[ctypes.c_int, 4]
+    @record
+    class Outer:
+      SIZE = 16
+      x: Annotated[ctypes.c_int, 0]
+      inner_ptr: Annotated[POINTER[Inner], 8]
+    init_records()
+
+    src = """
+      struct inner { int value; int flag; };
+      struct outer { int x; struct inner *inner_ptr; };
+      int test(struct inner *p) {
+        return p->value + p->flag;
+      }
+    """
+    dll = self.compile(src)
+    @dll.bind
+    def test(p:POINTER[Inner]) -> ctypes.c_int: ...
+
+    inner = Inner(value=42, flag=10)
+    outer = Outer(x=1, inner_ptr=ctypes.pointer(inner))
+    # Retrieve pointer from struct field and pass to C
+    self.assertEqual(test(outer.inner_ptr), 52)
+
+  def test_pointer_field_loses_reference(self):
+    # BUG: When a pointer is stored in a record struct field, only the address bytes are saved.
+    # The pointer's _objects dict (which prevents GC of the pointed-to object) is lost.
+    # This causes the pointed-to object to be garbage collected, leading to use-after-free.
+    from tinygrad.runtime.support.c import POINTER
+    @record
+    class MaskContext:
+      SIZE = 16
+      value: Annotated[ctypes.c_int, 0]
+      initialized: Annotated[ctypes.c_int, 4]
+      ptr: Annotated[ctypes.c_void_p, 8]
+    @record
+    class Params:
+      SIZE = 16
+      x: Annotated[ctypes.c_int, 0]
+      mask: Annotated[POINTER[MaskContext], 8]
+    init_records()
+
+    src = """
+      struct mask_ctx { int value; int initialized; void *ptr; };
+      void mask_begin(struct mask_ctx *m, int val) { m->value = val; m->initialized = 1; }
+      int mask_end(struct mask_ctx *m) { return m->value + m->initialized; }
+    """
+    dll = self.compile(src)
+    @dll.bind
+    def mask_begin(m:POINTER[MaskContext], val:ctypes.c_int) -> None: ...
+    @dll.bind
+    def mask_end(m:POINTER[MaskContext]) -> ctypes.c_int: ...
+
+    # When MaskContext() is created inline, it gets garbage collected after the pointer
+    # is stored because only the address bytes are saved, not the _objects reference.
+    params = Params(x=1, mask=ctypes.pointer(MaskContext()))
+    mask_begin(params.mask, 42)
+    result = mask_end(params.mask)
+    self.assertEqual(result, 43)  # 42 + 1
+
+@unittest.skipIf(OSX and ('MTLCompiler' in DLL._loaded_ or 'llvm' in DLL._loaded_), "libclang can't be loaded after MTLCompiler or llvm on OSX")
+@unittest.skipIf(WIN, "doesn't compile on windows")
+class TestAutogen(unittest.TestCase):
+  def run_gen(self, contents):
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.h') as f:
+      f.write(contents)
+      f.flush()
+
+      generated_code = gen(name="test_header", dll=None, files=[f.name])
+
+      namespace = {}
+      exec(generated_code, namespace)
+      return namespace
+
   def test_packed_structs(self):
     ns = self.run_gen("""
 typedef unsigned NvU32;
@@ -292,47 +354,6 @@ typedef struct
     assert frts_cmd.readVbiosDesc.__class__ is FWSECLIC_READ_VBIOS_DESC
     assert frts_cmd.frtsRegionDesc.__class__ is FWSECLIC_FRTS_REGION_DESC
 
-  @unittest.skipIf(WIN, "doesn't compile on windows")
-  @unittest.skipIf(OSX, "can't find stdint?")
-  def test_packed_fields(self):
-    ns = self.run_gen("""#include <stdint.h>
-typedef struct die_info
- {
-	 uint16_t die_id;
-	 uint16_t die_offset; /* Points to the corresponding die_header structure */
- } die_info;
-
-typedef struct ip_discovery_header
- {
-	 uint32_t signature;    /* Table Signature */
-	 uint16_t version;      /* Table Version */
-	 uint16_t size;         /* Table Size */
-	 uint32_t id;           /* Table ID */
-	 uint16_t num_dies;     /* Number of Dies */
-	 die_info die_info[16]; /* list die information for up to 16 dies */
-	 union {
-		 uint16_t padding[1];	/* version <= 3 */
-		 struct {		/* version == 4 */
-			 uint8_t base_addr_64_bit : 1; /* ip structures are using 64 bit base address */
-			 uint8_t reserved : 7;
-			 uint8_t reserved2;
-		 };
-	 };
- } ip_discovery_header;
-""")
-
-    ip_discovery_header = ns['ip_discovery_header']
-
-    hdr = b'IPDS\x04\x00|\x1d\x80\x1a\xffd\x01\x00\x00\x00\x8c\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00' # noqa: E501
-    ihdr = ip_discovery_header.from_buffer_copy(hdr)
-
-    assert ctypes.sizeof(ihdr) == 80
-    assert ihdr.signature == 0x53445049
-    assert ihdr.version == 0x0004
-    assert ihdr.num_dies == 1
-    assert ihdr.base_addr_64_bit == 1
-
-  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_gen_from_header(self):
     namespace = self.run_gen("""
     typedef struct {
@@ -378,7 +399,6 @@ typedef struct ip_discovery_header
     self.assertTrue(hasattr(rect, 'height'))
     self.assertTrue(hasattr(rect, 'color'))
 
-  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_struct_ordering(self):
     namespace = self.run_gen("""
     struct A;
@@ -408,77 +428,6 @@ typedef struct ip_discovery_header
     self.assertTrue(hasattr(b, 'c_ptr'))
     self.assertTrue(hasattr(c, 'a_ptr'))
 
-  @unittest.skipIf(WIN, "doesn't compile on windows")
-  def test_pointer_field_roundtrip(self):
-    # This tests storing a pointer in a record struct field and passing it to C
-    # Mimics how mesa.struct_lp_build_tgsi_params.mask is used
-    from tinygrad.runtime.support.c import POINTER
-    @record
-    class Inner:
-      SIZE = 8
-      value: Annotated[ctypes.c_int, 0]
-      flag: Annotated[ctypes.c_int, 4]
-    @record
-    class Outer:
-      SIZE = 16
-      x: Annotated[ctypes.c_int, 0]
-      inner_ptr: Annotated[POINTER[Inner], 8]
-    init_records()
-
-    src = """
-      struct inner { int value; int flag; };
-      struct outer { int x; struct inner *inner_ptr; };
-      int test(struct inner *p) {
-        return p->value + p->flag;
-      }
-    """
-    dll = self.compile(src)
-    @dll.bind
-    def test(p:POINTER[Inner]) -> ctypes.c_int: ...
-
-    inner = Inner(value=42, flag=10)
-    outer = Outer(x=1, inner_ptr=ctypes.pointer(inner))
-    # Retrieve pointer from struct field and pass to C
-    self.assertEqual(test(outer.inner_ptr), 52)
-
-  @unittest.skipIf(WIN, "doesn't compile on windows")
-  def test_pointer_field_loses_reference(self):
-    # BUG: When a pointer is stored in a record struct field, only the address bytes are saved.
-    # The pointer's _objects dict (which prevents GC of the pointed-to object) is lost.
-    # This causes the pointed-to object to be garbage collected, leading to use-after-free.
-    from tinygrad.runtime.support.c import POINTER
-    @record
-    class MaskContext:
-      SIZE = 16
-      value: Annotated[ctypes.c_int, 0]
-      initialized: Annotated[ctypes.c_int, 4]
-      ptr: Annotated[ctypes.c_void_p, 8]
-    @record
-    class Params:
-      SIZE = 16
-      x: Annotated[ctypes.c_int, 0]
-      mask: Annotated[POINTER[MaskContext], 8]
-    init_records()
-
-    src = """
-      struct mask_ctx { int value; int initialized; void *ptr; };
-      void mask_begin(struct mask_ctx *m, int val) { m->value = val; m->initialized = 1; }
-      int mask_end(struct mask_ctx *m) { return m->value + m->initialized; }
-    """
-    dll = self.compile(src)
-    @dll.bind
-    def mask_begin(m:POINTER[MaskContext], val:ctypes.c_int) -> None: ...
-    @dll.bind
-    def mask_end(m:POINTER[MaskContext]) -> ctypes.c_int: ...
-
-    # When MaskContext() is created inline, it gets garbage collected after the pointer
-    # is stored because only the address bytes are saved, not the _objects reference.
-    params = Params(x=1, mask=ctypes.pointer(MaskContext()))
-    mask_begin(params.mask, 42)
-    result = mask_end(params.mask)
-    self.assertEqual(result, 43)  # 42 + 1
-
-  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_anonymous_children(self):
     namespace = self.run_gen("""
       struct foo {
@@ -491,7 +440,6 @@ typedef struct ip_discovery_header
     self.assertIn('struct_foo', namespace)
     self.assertIn('struct_foo_bar', namespace)
 
-  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_enums(self):
     namespace = self.run_gen("""
       enum Foo { A, B, C };
@@ -511,4 +459,43 @@ typedef struct ip_discovery_header
     assert namespace["enum_Bar"].get(1) == "Y"
     assert namespace["enum_Bar"].get(2) == "Z"
 
+  @unittest.skipIf(OSX, "can't find stdint?")
+  def test_packed_fields(self):
+    ns = self.run_gen("""#include <stdint.h>
+typedef struct die_info
+ {
+	 uint16_t die_id;
+	 uint16_t die_offset; /* Points to the corresponding die_header structure */
+ } die_info;
+
+typedef struct ip_discovery_header
+ {
+	 uint32_t signature;    /* Table Signature */
+	 uint16_t version;      /* Table Version */
+	 uint16_t size;         /* Table Size */
+	 uint32_t id;           /* Table ID */
+	 uint16_t num_dies;     /* Number of Dies */
+	 die_info die_info[16]; /* list die information for up to 16 dies */
+	 union {
+		 uint16_t padding[1];	/* version <= 3 */
+		 struct {		/* version == 4 */
+			 uint8_t base_addr_64_bit : 1; /* ip structures are using 64 bit base address */
+			 uint8_t reserved : 7;
+			 uint8_t reserved2;
+		 };
+	 };
+ } ip_discovery_header;
+""")
+
+    ip_discovery_header = ns['ip_discovery_header']
+
+    hdr = b'IPDS\x04\x00|\x1d\x80\x1a\xffd\x01\x00\x00\x00\x8c\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00' # noqa: E501
+    ihdr = ip_discovery_header.from_buffer_copy(hdr)
+
+    assert ctypes.sizeof(ihdr) == 80
+    assert ihdr.signature == 0x53445049
+    assert ihdr.version == 0x0004
+    assert ihdr.num_dies == 1
+    assert ihdr.base_addr_64_bit == 1
+
 if __name__ == "__main__": unittest.main()

test/unit/test_call.py

@@ -42,7 +42,7 @@ class TestCall(unittest.TestCase):
     b = Tensor.randn(K, N)
     Tensor.realize(a, b)
     c = Tensor.call(a, b, fxn=a.as_param(0) @ b.as_param(1))
-    np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5)
+    np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)
 
   @unittest.skip("needs GEMM on mixins")
   def test_call_gemm_uop(self):
@@ -56,7 +56,7 @@ class TestCall(unittest.TestCase):
     y = UOp.param(1, dtypes.float, shape=(K, N))
     c = Tensor.call(a, b, fxn=x@y)
 
-    np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5)
+    np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_disk_tensor.py

@@ -59,13 +59,13 @@ class TestRawDiskBuffer(unittest.TestCase):
     _test_bitcasted(t, dtypes.float32, 0.0)
     _test_bitcasted(t, dtypes.uint32, 0)
     # pi in float16 stored via int16
-    t.bitcast(dtypes.uint16).assign(Tensor.full((128, 64), 0x4248, dtype=dtypes.uint16)).realize()
+    t.assign(Tensor.full((128, 64), 0x4248, dtype=dtypes.uint16).bitcast(dtypes.uint8)).realize()
     _test_bitcasted(t, dtypes.float16, 3.140625)
     _test_bitcasted(t, dtypes.float32, 50.064727)
     _test_bitcasted(t, dtypes.uint16, 0x4248)
     _test_bitcasted(t, dtypes.uint32, 0x42484248)
     # pi in float32 stored via float32
-    t.bitcast(dtypes.float32).assign(Tensor.full((128, 32), 3.1415927, dtype=dtypes.float32)).realize()
+    t.assign(Tensor.full((128, 32), 3.1415927, dtype=dtypes.float32).bitcast(dtypes.uint8)).realize()
     _test_bitcasted(t, dtypes.float32, 3.1415927)
     _test_bitcasted(t, dtypes.uint32, 0x40490FDB)
     # doesn't suport normal cast
@@ -250,6 +250,24 @@ class TestDiskTensor(unittest.TestCase):
     tout = [(x//256, x%256) for x in out]
     assert tout == list([(x+1,x) for x in range(32,64,2)])
 
+  def test_strided_read(self):
+    # test non-contiguous (strided) read - should read elements at indices 0, 2, 4
+    pathlib.Path(temp(fn:="dt_strided_read")).unlink(missing_ok=True)
+    dt = Tensor([0, 1, 2, 3, 4, 5]).to(f"disk:{temp(fn)}")
+    result = dt[::2].tolist()
+    # TODO: dt[::2] selects indices 0, 2, 4, so result should be [0, 2, 4]
+    # self.assertEqual(result, [0, 2, 4])
+    self.assertEqual(result, [0, 1, 2])  # wrong!
+
+  def test_permuted_read(self):
+    # test non-contiguous (permuted) read - should read transposed
+    pathlib.Path(temp(fn:="dt_permuted_read")).unlink(missing_ok=True)
+    dt = Tensor([[0, 1, 2], [3, 4, 5]]).to(f"disk:{temp(fn)}")
+    result = dt.T.tolist()
+    # TODO: transpose should give [[0, 3], [1, 4], [2, 5]]
+    # self.assertEqual(result, [[0, 3], [1, 4], [2, 5]])
+    self.assertEqual(result, [[0, 1], [2, 3], [4, 5]])  # wrong!
+
   def test_write_ones(self):
     pathlib.Path(temp("dt_write_ones")).unlink(missing_ok=True)
 
@@ -330,13 +348,25 @@ class TestDiskTensor(unittest.TestCase):
 
   def test_assign_with_bitcast(self):
     # bitcast assign is used in safe_save for writing header length
-    # this tests the synchronous disk assign hack handles bitcast correctly
+    # bitcast on source side works, bitcast on target side raises
     pathlib.Path(temp(fn:="dt_assign_bitcast")).unlink(missing_ok=True)
     t = Tensor.empty(16, device=f"disk:{temp(fn)}", dtype=dtypes.uint8)
-    t[0:8].bitcast(dtypes.int64).assign([12345])
-    # verify the data was written correctly
+    # correct way: bitcast the source to match target dtype
+    t[0:8].assign(Tensor([12345], dtype=dtypes.int64, device="CPU").bitcast(dtypes.uint8))
     val = int.from_bytes(t[0:8].data(), 'little')
     self.assertEqual(val, 12345)
+    # bitcast on target with non-broadcastable dtype raises
+    with self.assertRaises(RuntimeError):
+      t[0:4].bitcast(dtypes.int32).assign(Tensor([12345], dtype=dtypes.int64))
+
+  def test_assign_to_bitcast_view(self):
+    # assign float values to a float32 view of a uint8 disk buffer (used by safe_save)
+    pathlib.Path(temp(fn:="dt_bitcast_view_assign")).unlink(missing_ok=True)
+    t = Tensor.empty(32, device=f"disk:{temp(fn)}", dtype=dtypes.uint8)
+    # create float32 view of bytes 8-24 (4 floats)
+    float_view = t[8:24].bitcast(dtypes.float32)
+    float_view.assign(Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32, device="CPU"))
+    np.testing.assert_array_equal(float_view.numpy(), [1.0, 2.0, 3.0, 4.0])
 
   def test_assign_cross_device(self):
     # disk assign allows cross-device (source on GPU/CPU, target on disk)

test/unit/test_memory_planner.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import dtypes, Device
+from tinygrad import dtypes
 from tinygrad.device import Buffer
 from tinygrad.engine.memory import _internal_memory_planner
 
@@ -7,7 +7,7 @@ global_map = {}
 def b(i, base=None, offset=0, pin=False, size=16):
   global global_map
   if i in global_map: return global_map[i]
-  global_map[i] = Buffer(Device.DEFAULT, size, dtypes.int8, base=global_map[base] if base is not None else None, offset=offset)
+  global_map[i] = Buffer("NULL", size, dtypes.int8, base=global_map[base] if base is not None else None, offset=offset)
   if pin: global_map[i].ref(1)
   return global_map[i]
 

test/unit/test_schedule_cache.py

@@ -35,7 +35,6 @@ class TestScheduleCache(unittest.TestCase):
     _, var_vals = t.schedule_with_vars()
     self.assertEqual(var_vals, {'pos': 42})
 
-  @Context(SPEC=0)
   def test_custom_kernel(self):
     for i in range(4):
       a = Tensor.empty(1)
@@ -43,7 +42,6 @@ class TestScheduleCache(unittest.TestCase):
       a.realize()
       self.assertEqual(a.item(), i)
 
-  @Context(SPEC=0)
   def test_same_custom_function_reuses_cache(self):
     schedule_cache.clear()
     fxn = functools.partial(custom_set0_kernel, num=10)

test/unit/test_simplify_valid_idx.py

@@ -470,5 +470,19 @@ class TestUnfoldableImageChannelSelection(unittest.TestCase):
     load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.DEFINE_GLOBAL, dtypes.imagef((10, 10, 4)), arg=0).index(x, ptr=True), UOp.const(dtypes.float, 0)))
     self.assertEqual(self._count_nans(load), 1)
 
+class TestDropTrueGate(unittest.TestCase):
+  def test_drop_true_gate_on_index(self):
+    # test that INDEX with a constant True gate gets simplified to drop the gate
+    from tinygrad.codegen.late.devectorizer import load_store_indexing
+    from tinygrad.uop.ops import graph_rewrite
+    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
+    idx = UOp.const(dtypes.index, 0)
+    true_gate = UOp.const(dtypes.bool, True)
+    index_with_gate = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx, true_gate))
+    # apply the optimization
+    result = graph_rewrite(index_with_gate, load_store_indexing)
+    # the True gate should be dropped (INDEX should only have 2 sources)
+    self.assertEqual(len(result.src), 2, "True gate should be dropped from INDEX")
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_uop_graph.py

@@ -479,6 +479,26 @@ class TestUOpGraph(unittest.TestCase):
     for u in uops:
       self.assertNotEqual(u.dtype, dtypes.long)
 
+  def test_load_idx_no_math_on_loaded(self):
+    # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
+    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(128000), arg=0, src=())
+    c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
+    c3 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
+    c4 = c3.index(c1)  # c4 is a load
+    c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
+    c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(60000), arg=2, src=())
+    c8 = c7.index(c6)
+    # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
+    loaded_idx = c4.cast(dtypes.index)
+    comparison = (loaded_idx + c5) < UOp.const(dtypes.index, 60000)
+    c9 = comparison.where(c8.cast(dtypes.uint).cast(dtypes.uchar), 0).reduce(c5, arg=Ops.ADD)
+    c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
+    uops = to_uops_list([c10])
+    for u in uops:
+      self.assertNotEqual(u.dtype, dtypes.long)
+
   def test_fold_gated_load(self):
     glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
     glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 1)
@@ -686,6 +706,97 @@ class TestExpander(unittest.TestCase):
     sink = expander_rewrite(sink)
     print(sink)
 
+class TestReduceCollapse(unittest.TestCase):
+  def test_multi_range_reduce_add(self):
+    """Test that (x + y).reduce(r1, r2) distributes over multiple ranges"""
+    from tinygrad.codegen.simplify import pm_reduce_collapse
+    # Create two ranges
+    r1 = UOp.range(3, 0)
+    r2 = UOp.range(4, 1)
+    # Create x + y where x and y depend on different ranges
+    x = r1.cast(dtypes.float)
+    y = r2.cast(dtypes.float)
+    # (x + y).reduce(r1, r2) should be rewritten
+    red = (x + y).reduce(r1, r2, arg=Ops.ADD)
+    self.assertEqual(len(red.src), 3)  # value + 2 ranges
+    result = graph_rewrite(red, pm_reduce_collapse, name='test')
+    # Should become add of two separate reduces
+    self.assertEqual(result.op, Ops.ADD)
+
+class TestLoadStoreFolding(unittest.TestCase):
+  def test_gated_load_gep_preserves_alt(self):
+    """Test that LOAD(GEP, alt) preserves alt value after rewrite"""
+    from tinygrad.codegen.late.devectorizer import load_store_folding
+    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.vec(4).ptr(), (), 0)
+    idx = UOp.const(dtypes.int, 0)
+    gate = UOp.const(dtypes.bool, True)
+    gated_index = buf.index(idx, gate)
+    gep = gated_index.gep(0)
+    alt = UOp.const(dtypes.float, 42.0)
+    gated_load = gep.load(alt)
+    self.assertEqual(len(gated_load.src), 2)  # GEP + alt
+    result = graph_rewrite(gated_load, load_store_folding, name='test')
+    # After rewrite, should still have alt value preserved
+    self.assertEqual(result.op, Ops.GEP)
+    inner_load = result.src[0]
+    self.assertEqual(inner_load.op, Ops.LOAD)
+    self.assertEqual(len(inner_load.src), 2)  # INDEX + alt
+
+  def test_gated_load_ptrcat_preserves_alt(self):
+    """Test that LOAD(PTRCAT, alt) preserves alt value after rewrite"""
+    from tinygrad.codegen.late.devectorizer import load_store_folding
+    buf1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
+    buf2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
+    idx = UOp.const(dtypes.int, 0)
+    idx1 = buf1.index(idx)
+    idx2 = buf2.index(idx)
+    ptrcat = UOp(Ops.PTRCAT, dtypes.float.ptr().vec(2), (idx1, idx2))
+    alt = UOp.const(dtypes.float.vec(2), 42.0)
+    gated_load = ptrcat.load(alt)
+    self.assertEqual(len(gated_load.src), 2)  # PTRCAT + alt
+    result = graph_rewrite(gated_load, load_store_folding, name='test')
+    # After rewrite, should be CAT of LOADs, each preserving alt
+    self.assertEqual(result.op, Ops.CAT)
+    for inner_load in result.src:
+      self.assertEqual(inner_load.op, Ops.LOAD)
+      self.assertEqual(len(inner_load.src), 2)  # INDEX + alt
+    self.assertEqual(inner_load.src[1].arg, 42.0)  # alt value preserved
+
+class TestConstBufferize(unittest.TestCase):
+  def test_const_bufferize_with_ranges(self):
+    """Test that CONST.BUFFERIZE with ranges is folded correctly.
+
+    BUFFERIZE can have ranges as additional sources beyond the value.
+    The pattern at rangeify.py uses allow_any_len=True because
+    CONST doesn't depend on ranges (constant is same value everywhere).
+    """
+    from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
+    c = UOp.const(dtypes.float, 42.0)
+    r1 = UOp.range(3, 0)
+    bufferize_with_range = UOp(Ops.BUFFERIZE, dtypes.float, (c, r1), arg=BufferizeOpts(device="CPU"))
+    self.assertEqual(len(bufferize_with_range.src), 2)  # const + 1 range
+
+    result = graph_rewrite(bufferize_with_range, pm_const_buffer_folding, name='test')
+    # BUFFERIZE should be removed, result is const broadcast to shape
+    self.assertNotEqual(result.op, Ops.BUFFERIZE)
+    const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
+    self.assertIn(42.0, const_vals)
+
+  def test_const_bufferize_with_multiple_ranges(self):
+    """Test CONST.BUFFERIZE with multiple ranges is also folded."""
+    from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
+    c = UOp.const(dtypes.float, 3.14)
+    r1 = UOp.range(3, 0)
+    r2 = UOp.range(4, 1)
+    bufferize_with_ranges = UOp(Ops.BUFFERIZE, dtypes.float, (c, r1, r2), arg=BufferizeOpts(device="CPU"))
+    self.assertEqual(len(bufferize_with_ranges.src), 3)  # const + 2 ranges
+
+    result = graph_rewrite(bufferize_with_ranges, pm_const_buffer_folding, name='test')
+    # BUFFERIZE should be removed
+    self.assertNotEqual(result.op, Ops.BUFFERIZE)
+    const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
+    self.assertIn(3.14, const_vals)
+
 class TestUOpTags(unittest.TestCase):
   def test_inc_by_one(self):
     g = UOp.const(dtypes.int, 1) + UOp.const(dtypes.int, 1)

test/unit/test_uop_symbolic.py

@@ -1074,6 +1074,24 @@ class TestGatedUopGivenValid(unittest.TestCase):
     expected_vec = UOp(Ops.VECTORIZE, dtypes.index.vec(2), (uconst(0), r0))
     self.assertEqual(idx, (r0 < 3).where(expected_vec, UOp.invalid()))
 
+class TestRangeSplitting(unittest.TestCase):
+  def test_range_split_on_mod(self):
+    # test that mark_range_mod splits RANGE(8) into RANGE(4)*2 + RANGE(2) when used with %2
+    from tinygrad.codegen.simplify import pm_split_ranges, pm_flatten_range
+    r0 = UOp.range(uconst(8), 0)
+    # create a simple expression using the range with mod: store range%2 to a buffer
+    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
+    val = (r0 % uconst(2)).cast(dtypes.int)
+    store = UOp(Ops.STORE, dtypes.void, (buf.index(uconst(0)), val))
+    sink = UOp(Ops.SINK, dtypes.void, (UOp(Ops.END, dtypes.void, (store, r0)),))
+    # count RANGEs before
+    ranges_before = len([u for u in sink.toposort() if u.op is Ops.RANGE])
+    # apply the range splitting optimization
+    sink_after = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="test split ranges")
+    # count RANGEs after - should have more due to splitting
+    ranges_after = len([u for u in sink_after.toposort() if u.op is Ops.RANGE])
+    self.assertGreater(ranges_after, ranges_before, "RANGE should be split when used with mod of divisible constant")
+
 class TestBounds(unittest.TestCase):
   def test_unrolled_arange(self):
     # #include <metal_stdlib>

tinygrad/codegen/__init__.py

@@ -115,8 +115,8 @@ pm_linearize_cleanups = PatternMatcher([
   # if statements are not allowed in the graph
   (UPat((Ops.IF, Ops.ENDIF)), lambda: panic(RuntimeError("if not allowed in graph"))),
   # gated INDEX becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF
-  (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat(name="gate", dtype=dtypes.bool))).or_casted(), UPat()),
-        allow_any_len=True), lambda u, gate: (u, [mif:=UOp(Ops.IF, src=(gate, u.src[0])), u, UOp(Ops.ENDIF, src=(mif,))]))
+  (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat(name="gate", dtype=dtypes.bool))).or_casted(), UPat())),
+   lambda u, gate: (u, [mif:=UOp(Ops.IF, src=(gate, u.src[0])), u, UOp(Ops.ENDIF, src=(mif,))]))
 ])
 
 # requires lst be toposorted. like graph rewrite, but for lines

tinygrad/codegen/late/devectorizer.py

@@ -123,12 +123,12 @@ load_store_folding = PatternMatcher([
   (UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True),
    lambda gep, ld: ld.replace(dtype=ld.dtype.scalar().vec(gep.dtype.count), src=(gep.src[0],)+ld.src[1:]).gep(gep.arg)),
   # GEP on data of STORE
-  (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st")), allow_any_len=True, name="sto"), gep_on_store),
+  (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st")), name="sto"), gep_on_store),
   # put PTRCAT after LOAD
   (UPat(Ops.LOAD, src=(UPat(Ops.PTRCAT, name="cat"),), name="ld", allow_any_len=True),
    lambda cat,ld: UOp(Ops.CAT, cat.dtype.base.vec(cat.dtype.vcount), tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))),
   # put PTRCAT after STORE
-  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), allow_any_len=True, name="sto"), cat_after_store),
+  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), name="sto"), cat_after_store),
 ])
 
 # *** correct load/store ***
@@ -342,6 +342,6 @@ pm_add_loads = PatternMatcher([
   (UPat(Ops.INDEX, name="idx"), lambda idx: None if isinstance(idx.dtype, (PtrDType, ImageDType)) else
     idx.replace(dtype=idx.src[0].dtype).load(dtype=idx.dtype.base)),
   # remove loads from stores
-  (UPat(Ops.STORE, src=(UPat(Ops.LOAD),), allow_any_len=True, name="s"), lambda s: s.replace(src=(s.src[0].src[0],)+s.src[1:])),
+  (UPat(Ops.STORE, src=(UPat(Ops.LOAD), UPat(name="val")), name="s"), lambda s,val: s.replace(src=(s.src[0].src[0], val))),
 ])
 

tinygrad/helpers.py

@@ -206,6 +206,8 @@ 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/nn/state.py

@@ -78,7 +78,7 @@ def safe_save(tensors:dict[str, Tensor], fn:str, metadata:dict[str, Any]|None=No
   j += "\x20"*(round_up(len(j),8)-len(j))
   pathlib.Path(fn).unlink(missing_ok=True)
   t = Tensor.empty(8+len(j)+offset, dtype=dtypes.uint8, device=f"disk:{fn}")
-  t[0:8].bitcast(dtypes.int64).assign([len(j)])
+  t[0:8].assign(Tensor([len(j)], dtype=dtypes.int64, device="CPU").bitcast(dtypes.uint8))
   t[8:8+len(j)].assign(list(j.encode('utf-8')))
   for k,v in safe_load(t).items(): v.assign(tensors[k])
 

tinygrad/renderer/cstyle.py

@@ -46,10 +46,10 @@ base_rewrite = PatternMatcher([
   # new load/store
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), allow_any_len=True),
    lambda ctx,buf,idx: f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"),
-  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat.var("gate"))).or_casted("bidx"), UPat.var("var")), allow_any_len=True),
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat.var("gate"))).or_casted("bidx"), UPat.var("var"))),
    lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"),
-  (UPat(Ops.LOAD, src=(UPat.var('bidx'),), allow_any_len=True), lambda ctx,bidx: f"(*{ctx[bidx]})"),
-  (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var")), allow_any_len=True), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
+  (UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"(*{ctx[bidx]})"),
+  (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
   # alu/gep
   # TODO: look for left-associative
   (UPat(GroupOp.ALU, name="x"), lambda ctx,x: ctx.code_for_op[x.op](

tinygrad/runtime/ops_metal.py

@@ -4,12 +4,13 @@ import tinygrad.runtime.support.objc as objc
 from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent, CompilerSet, CompilerPair
 from tinygrad.renderer.cstyle import MetalRenderer
 from tinygrad.runtime.autogen import metal
+from tinygrad.runtime.support.c import DLL
 
 # 13 is requestType that metal uses to compile source code into MTLB, there aren't any docs or symbols.
 REQUEST_TYPE_COMPILE = 13
 
 # Must be loaded for default Metal Device: https://developer.apple.com/documentation/metal/1433401-mtlcreatesystemdefaultdevice?language=objc
-ctypes.CDLL("/System/Library/Frameworks/CoreGraphics.framework/CoreGraphics")
+DLL("CoreGraphics", "CoreGraphics")
 
 # FIXME: these need autogen to support objc categories
 # https://developer.apple.com/library/archive/documentation/Cocoa/Conceptual/ObjectiveC/Chapters/ocCategories.html
@@ -67,7 +68,7 @@ class MetalCompiler(Compiler):
   # doesn't seem to be anything we can do.
   with contextlib.suppress(FileNotFoundError, ModuleNotFoundError):
     import tinygrad.runtime.autogen.llvm # noqa: F401
-  support = ctypes.CDLL("/System/Library/PrivateFrameworks/MTLCompiler.framework/MTLCompiler")
+  support = DLL("MTLCompiler", "MTLCompiler")
   support.MTLCodeGenServiceCreate.restype = ctypes.c_void_p
 
   def __init__(self):

tinygrad/runtime/support/amd.py

@@ -88,8 +88,8 @@ def import_asic_regs(prefix:str, version:tuple[int, ...], cls=AMDReg) -> dict[st
     return x
   def _download_file(ver, suff) -> str:
     dir_prefix = {"osssys": "oss"}.get(prefix, prefix)
-    fetch_name, file_name = f"{prefix}_{'_'.join(map(str, ver))}_{suff}.h", f"{prefix}_{'_'.join(map(str, version))}_{suff}.h"
-    return header_download(f"include/asic_reg/{dir_prefix}/{fetch_name}", name=file_name, subdir="asic_regs")
+    fetch_name = f"{prefix}_{'_'.join(map(str, ver))}_{suff}.h"
+    return header_download(f"include/asic_reg/{dir_prefix}/{fetch_name}", name=fetch_name, subdir="asic_regs")
 
   for ver in fixup_ip_version(prefix, version):
     try: offs, sh_masks = _extract_regs(_download_file(ver, "offset")), _extract_regs(_download_file(ver, "sh_mask"))

tinygrad/runtime/support/c.py

@@ -131,13 +131,15 @@ def init_c_struct_t(sz:int, fields: tuple[tuple, ...]):
 def init_c_var(ty, creat_cb): return (creat_cb(v:=del_an(ty)()), v)[1]
 
 class DLL(ctypes.CDLL):
+  _loaded_: set[str] = set()
+
   @staticmethod
   def findlib(nm:str, paths:list[str], extra_paths=[]):
     if nm == 'libc' and OSX: return '/usr/lib/libc.dylib'
     if pathlib.Path(path:=getenv(nm.replace('-', '_').upper()+"_PATH", '')).is_file(): return path
     for p in paths:
       libpaths = {"posix": ["/usr/lib64", "/usr/lib", "/usr/local/lib"], "nt": os.environ['PATH'].split(os.pathsep),
-                  "darwin": ["/opt/homebrew/lib", f"/System/Library/Frameworks/{p}.framework"],
+                  "darwin": ["/opt/homebrew/lib", f"/System/Library/Frameworks/{p}.framework", f"/System/Library/PrivateFrameworks/{p}.framework"],
                   'linux': ['/lib', '/lib64', f"/lib/{sysconfig.get_config_var('MULTIARCH')}", "/usr/lib/wsl/lib/"]}
       if (pth:=pathlib.Path(p)).is_absolute():
         if pth.is_file(): return p
@@ -154,12 +156,12 @@ class DLL(ctypes.CDLL):
               if f.read(4) == b'\x7FELF': return str(l)
 
   def __init__(self, nm:str, paths:str|list[str], extra_paths=[], emsg="", **kwargs):
-    self.nm, self.emsg, self.loaded = nm, emsg, False
+    self.nm, self.emsg = nm, emsg
     if (path:= DLL.findlib(nm, paths if isinstance(paths, list) else [paths], extra_paths if isinstance(extra_paths, list) else [extra_paths])):
       if DEBUG >= 3: print(f"loading {nm} from {path}")
       try:
         super().__init__(path, **kwargs)
-        self.loaded = True
+        self._loaded_.add(self.nm)
       except OSError as e:
         self.emsg = str(e)
         if DEBUG >= 3: print(f"loading {nm} failed: {e}")
@@ -175,5 +177,6 @@ class DLL(ctypes.CDLL):
     return wrapper
 
   def __getattr__(self, nm):
-    if not self.loaded: raise AttributeError(f"failed to load library {self.nm}: " + (self.emsg or f"try setting {self.nm.upper()+'_PATH'}?"))
+    if self.nm not in self._loaded_:
+      raise AttributeError(f"failed to load library {self.nm}: " + (self.emsg or f"try setting {self.nm.upper()+'_PATH'}?"))
     return super().__getattr__(nm)

tinygrad/runtime/support/usb.py

@@ -5,10 +5,10 @@ from tinygrad.helpers import DEBUG, to_mv, round_up, OSX
 from tinygrad.runtime.support.hcq import MMIOInterface
 
 class USB3:
-  def __init__(self, vendor:int, dev:int, ep_data_in:int, ep_stat_in:int, ep_data_out:int, ep_cmd_out:int, max_streams:int=31):
+  def __init__(self, vendor:int, dev:int, ep_data_in:int, ep_stat_in:int, ep_data_out:int, ep_cmd_out:int, max_streams:int=31, use_bot=False):
     self.vendor, self.dev = vendor, dev
     self.ep_data_in, self.ep_stat_in, self.ep_data_out, self.ep_cmd_out = ep_data_in, ep_stat_in, ep_data_out, ep_cmd_out
-    self.max_streams = max_streams
+    self.max_streams, self.use_bot = max_streams, use_bot
     self.ctx = ctypes.POINTER(libusb.struct_libusb_context)()
 
     if libusb.libusb_init(ctypes.byref(self.ctx)): raise RuntimeError("libusb_init failed")
@@ -25,30 +25,34 @@ class USB3:
     # Set configuration and claim interface
     if libusb.libusb_set_configuration(self.handle, 1): raise RuntimeError("set_configuration failed")
     if libusb.libusb_claim_interface(self.handle, 0): raise RuntimeError("claim_interface failed. sudo required?")
-    if libusb.libusb_set_interface_alt_setting(self.handle, 0, 1): raise RuntimeError("alt_setting failed")
 
-    # Clear any stalled endpoints
-    all_eps = (self.ep_data_out, self.ep_data_in, self.ep_stat_in, self.ep_cmd_out)
-    for ep in all_eps: libusb.libusb_clear_halt(self.handle, ep)
+    if use_bot:
+      self._tag = 0
+    else:
+      if libusb.libusb_set_interface_alt_setting(self.handle, 0, 1): raise RuntimeError("alt_setting failed")
 
-    # Allocate streams
-    stream_eps = (ctypes.c_uint8 * 3)(self.ep_data_out, self.ep_data_in, self.ep_stat_in)
-    if (rc:=libusb.libusb_alloc_streams(self.handle, self.max_streams * len(stream_eps), stream_eps, len(stream_eps))) < 0:
-      raise RuntimeError(f"alloc_streams failed: {rc}")
+      # Clear any stalled endpoints
+      all_eps = (self.ep_data_out, self.ep_data_in, self.ep_stat_in, self.ep_cmd_out)
+      for ep in all_eps: libusb.libusb_clear_halt(self.handle, ep)
 
-    # Base cmd
-    cmd_template = bytes([0x01, 0x00, 0x00, 0x01, *([0] * 12), 0xE4, 0x24, 0x00, 0xB2, 0x1A, 0x00, 0x00, 0x00, *([0] * 8)])
+      # Allocate streams
+      stream_eps = (ctypes.c_uint8 * 3)(self.ep_data_out, self.ep_data_in, self.ep_stat_in)
+      if (rc:=libusb.libusb_alloc_streams(self.handle, self.max_streams * len(stream_eps), stream_eps, len(stream_eps))) < 0:
+        raise RuntimeError(f"alloc_streams failed: {rc}")
 
-    # Init pools
-    self.tr = {ep: [libusb.libusb_alloc_transfer(0) for _ in range(self.max_streams)] for ep in all_eps}
+      # Base cmd
+      cmd_template = bytes([0x01, 0x00, 0x00, 0x01, *([0] * 12), 0xE4, 0x24, 0x00, 0xB2, 0x1A, 0x00, 0x00, 0x00, *([0] * 8)])
 
-    self.buf_cmd = [(ctypes.c_uint8 * len(cmd_template))(*cmd_template) for _ in range(self.max_streams)]
-    self.buf_stat = [(ctypes.c_uint8 * 64)() for _ in range(self.max_streams)]
-    self.buf_data_in = [(ctypes.c_uint8 * 0x1000)() for _ in range(self.max_streams)]
-    self.buf_data_out = [(ctypes.c_uint8 * 0x80000)() for _ in range(self.max_streams)]
-    self.buf_data_out_mvs = [to_mv(ctypes.addressof(self.buf_data_out[i]), 0x80000) for i in range(self.max_streams)]
+      # Init pools
+      self.tr = {ep: [libusb.libusb_alloc_transfer(0) for _ in range(self.max_streams)] for ep in all_eps}
 
-    for slot in range(self.max_streams): struct.pack_into(">B", self.buf_cmd[slot], 3, slot + 1)
+      self.buf_cmd = [(ctypes.c_uint8 * len(cmd_template))(*cmd_template) for _ in range(self.max_streams)]
+      self.buf_stat = [(ctypes.c_uint8 * 64)() for _ in range(self.max_streams)]
+      self.buf_data_in = [(ctypes.c_uint8 * 0x1000)() for _ in range(self.max_streams)]
+      self.buf_data_out = [(ctypes.c_uint8 * 0x80000)() for _ in range(self.max_streams)]
+      self.buf_data_out_mvs = [to_mv(ctypes.addressof(self.buf_data_out[i]), 0x80000) for i in range(self.max_streams)]
+
+      for slot in range(self.max_streams): struct.pack_into(">B", self.buf_cmd[slot], 3, slot + 1)
 
   def _prep_transfer(self, tr, ep, stream_id, buf, length):
     tr.contents.dev_handle, tr.contents.endpoint, tr.contents.length, tr.contents.buffer = self.handle, ep, length, buf
@@ -68,38 +72,90 @@ class USB3:
         if tr.contents.status == libusb.LIBUSB_TRANSFER_COMPLETED: running -= 1
         elif tr.contents.status != 0xFF: raise RuntimeError(f"EP 0x{tr.contents.endpoint:02X} error: {tr.contents.status}")
 
+  def _bulk_out(self, ep: int, payload: bytes, timeout: int = 1000):
+    transferred = ctypes.c_int(0)
+    rc = libusb.libusb_bulk_transfer(
+      self.handle,
+      ep,
+      (ctypes.c_ubyte * len(payload))(*payload),
+      len(payload),
+      ctypes.byref(transferred),
+      timeout,
+    )
+    assert rc == 0, f"bulk OUT 0x{ep:02X} failed: {rc}"
+    assert transferred.value == len(payload), f"bulk OUT short write on 0x{ep:02X}: {transferred.value}/{len(payload)} bytes"
+
+  def _bulk_in(self, ep: int, length: int, timeout: int = 1000) -> bytes:
+    buf, transferred = (ctypes.c_ubyte * length)(), ctypes.c_int(0)
+    rc = libusb.libusb_bulk_transfer(
+      self.handle,
+      ep,
+      buf,
+      length,
+      ctypes.byref(transferred),
+      timeout,
+    )
+    assert rc == 0, f"bulk IN 0x{ep:02X} failed: {rc}"
+    return bytes(buf[:transferred.value])
+
   def send_batch(self, cdbs:list[bytes], idata:list[int]|None=None, odata:list[bytes|None]|None=None) -> list[bytes|None]:
     idata, odata = idata or [0] * len(cdbs), odata or [None] * len(cdbs)
-    results, tr_window, op_window = [], [], []
+    results:list[bytes|None] = []
+    tr_window, op_window = [], []
 
     for idx, (cdb, rlen, send_data) in enumerate(zip(cdbs, idata, odata)):
-      # allocate slot and stream. stream is 1-based
-      slot, stream = idx % self.max_streams, (idx % self.max_streams) + 1
+      if self.use_bot:
+        dir_in = rlen > 0
+        data_len = rlen if dir_in else (len(send_data) if send_data is not None else 0)
+        assert (data_len == 0) if dir_in else (rlen == 0), "BOT mode only supports either read or write per command"
 
-      # build cmd packet
-      self.buf_cmd[slot][16:16+len(cdb)] = list(cdb)
+        # CBW
+        self._tag += 1
+        flags = 0x80 if dir_in else 0x00
+        cbw = struct.pack("<IIIBBB", 0x43425355, self._tag, data_len, flags, 0, len(cdb)) + cdb + b"\x00" * (16 - len(cdb))
+        self._bulk_out(self.ep_data_out, cbw)
 
-      # cmd + stat transfers
-      tr_window.append(self._prep_transfer(self.tr[self.ep_cmd_out][slot], self.ep_cmd_out, None, self.buf_cmd[slot], len(self.buf_cmd[slot])))
-      tr_window.append(self._prep_transfer(self.tr[self.ep_stat_in][slot], self.ep_stat_in, stream, self.buf_stat[slot], 64))
+        # DAT
+        if dir_in:
+          results.append(self._bulk_in(self.ep_data_in, rlen))
+        else:
+          if send_data is not None:
+            self._bulk_out(self.ep_data_out, send_data)
+          results.append(None)
 
-      if rlen:
-        if rlen > len(self.buf_data_in[slot]): self.buf_data_in[slot] = (ctypes.c_uint8 * round_up(rlen, 0x1000))()
-        tr_window.append(self._prep_transfer(self.tr[self.ep_data_in][slot], self.ep_data_in, stream, self.buf_data_in[slot], rlen))
+        # CSW
+        sig, rtag, residue, status = struct.unpack("<IIIB", self._bulk_in(self.ep_data_in, 13, timeout=2000))
+        assert sig == 0x53425355, f"Bad CSW signature 0x{sig:08X}, expected 0x53425355"
+        assert rtag == self._tag, f"CSW tag mismatch: got {rtag}, expected {self._tag}"
+        assert status == 0, f"SCSI command failed, CSW status=0x{status:02X}, residue={residue}"
+      else:
+        # allocate slot and stream. stream is 1-based
+        slot, stream = idx % self.max_streams, (idx % self.max_streams) + 1
 
-      if send_data is not None:
-        if len(send_data) > len(self.buf_data_out[slot]):
-          self.buf_data_out[slot] = (ctypes.c_uint8 * len(send_data))()
-          self.buf_data_out_mvs[slot] = to_mv(ctypes.addressof(self.buf_data_out[slot]), len(send_data))
+        # build cmd packet
+        self.buf_cmd[slot][16:16+len(cdb)] = list(cdb)
 
-        self.buf_data_out_mvs[slot][:len(send_data)] = bytes(send_data)
-        tr_window.append(self._prep_transfer(self.tr[self.ep_data_out][slot], self.ep_data_out, stream, self.buf_data_out[slot], len(send_data)))
+        # cmd + stat transfers
+        tr_window.append(self._prep_transfer(self.tr[self.ep_cmd_out][slot], self.ep_cmd_out, None, self.buf_cmd[slot], len(self.buf_cmd[slot])))
+        tr_window.append(self._prep_transfer(self.tr[self.ep_stat_in][slot], self.ep_stat_in, stream, self.buf_stat[slot], 64))
 
-      op_window.append((idx, slot, rlen))
-      if (idx + 1 == len(cdbs)) or len(op_window) >= self.max_streams:
-        self._submit_and_wait(tr_window)
-        for idx, slot, rlen in op_window: results.append(bytes(self.buf_data_in[slot][:rlen]) if rlen else None)
-        tr_window = []
+        if rlen:
+          if rlen > len(self.buf_data_in[slot]): self.buf_data_in[slot] = (ctypes.c_uint8 * round_up(rlen, 0x1000))()
+          tr_window.append(self._prep_transfer(self.tr[self.ep_data_in][slot], self.ep_data_in, stream, self.buf_data_in[slot], rlen))
+
+        if send_data is not None:
+          if len(send_data) > len(self.buf_data_out[slot]):
+            self.buf_data_out[slot] = (ctypes.c_uint8 * len(send_data))()
+            self.buf_data_out_mvs[slot] = to_mv(ctypes.addressof(self.buf_data_out[slot]), len(send_data))
+
+          self.buf_data_out_mvs[slot][:len(send_data)] = bytes(send_data)
+          tr_window.append(self._prep_transfer(self.tr[self.ep_data_out][slot], self.ep_data_out, stream, self.buf_data_out[slot], len(send_data)))
+
+        op_window.append((idx, slot, rlen))
+        if (idx + 1 == len(cdbs)) or len(op_window) >= self.max_streams:
+          self._submit_and_wait(tr_window)
+          for idx, slot, rlen in op_window: results.append(bytes(self.buf_data_in[slot][:rlen]) if rlen else None)
+          tr_window = []
 
     return results
 

tinygrad/schedule/multi.py

@@ -71,7 +71,8 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
   # allgather
   copied_chunks = []
   for i,rc in enumerate(reduced_chunks):
-    if use_all2all: copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(rc.copy_to_device(buf.device[j]) for j in range(n_lbs))))
+    if isinstance(red.src[1].arg, str): copied_chunks.append(rc.copy_to_device(red.src[1].arg))
+    elif use_all2all: copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(rc.copy_to_device(buf.device[j]) for j in range(n_lbs))))
     else:
       this_chunk: list[UOp|None] = [None] * n_lbs
       this_chunk[(i+n_lbs-1)%n_lbs] = rc

tinygrad/schedule/rangeify.py

@@ -60,7 +60,7 @@ def split_reduceop(reduce:UOp, x:UOp):
 
 mop_cleanup = PatternMatcher([
   # merge adjacent RESHAPES, safe because they are not tagged
-  (UPat(Ops.RESHAPE, name="x2").f(Ops.RESHAPE, allow_any_len=True, name="x"),
+  (UPat(Ops.RESHAPE, src=(UPat(Ops.RESHAPE, name="x2"), UPat()), name="x"),
    lambda x,x2: x.replace(src=(x2.src[0], x.src[1])) if x.tag is None and x2.tag is None else None),
 ])
 
@@ -90,7 +90,7 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   (UPat(Ops.CUSTOM_KERNEL, name="ck"), resolve_custom_kernel),
 
   # remove CONTIGUOUS if the BUFFER is already contiguous
-  (UPat(Ops.BUFFER).f(Ops.RESHAPE, allow_any_len=True, name="r").f(Ops.CONTIGUOUS, name="c"), lambda r,c: r.replace(tag=c.tag)),
+  (UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER), UPat()), name="r").f(Ops.CONTIGUOUS, name="c"), lambda r,c: r.replace(tag=c.tag)),
 
   # split_reduceop
   (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)), split_reduceop),
@@ -126,6 +126,10 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
 
   # ** assign rules **
 
+  # move bitcast from assign target to source: a.bitcast(X).assign(src) -> a.assign(src.bitcast(a.dtype))
+  (UPat(Ops.ASSIGN, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(name="src")), name="assign"),
+   lambda target, src, assign: target.assign(src.bitcast(target.dtype)).replace(tag=assign.tag)),
+
   # assign only to buffer, otherwise make it a CONTIGUOUS
   (UPat(Ops.ASSIGN, src=(UPat(GroupOp.All-{Ops.BUFFER}, name="target"), UPat(name="x")), name="assign"),
    lambda x,target,assign: x.f(Ops.CONTIGUOUS, tag=assign.tag) if ((t:=target.base).op is not Ops.BUFFER and \
@@ -247,7 +251,7 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([
   # dont bufferize an arange
   (UPat.any((r:=UPat(dtype=dtypes.index).cast()).named("src"), r.eq(UPat()).named("src")).f(Ops.BUFFERIZE,
     allow_any_len=True, name="buf").f(Ops.INDEX, allow_any_len=True, name="idx"), remove_bufferize),
-  # no buffers for const
+  # no buffers for const (ranges don't matter for const - it's the same value everywhere)
   (UPat(Ops.CONST, name='c').f(Ops.BUFFERIZE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.arg).rtag(b.tag)),
   # indexing a const is a const
   (UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),),), lambda c: c),
@@ -271,20 +275,20 @@ def late_buffer_view(t:UOp, b:UOp):
   size = prod(shape)
 
   # walk up for the INDEX
+  # NOTE: even though we allow RESHAPE and SHRINK, they can combine to form non-contiguous access patterns (e.g. t[::2])
   x = t
-  while not any(u.op is Ops.INDEX for u in x.src):
-    assert x.op not in GroupOp.Elementwise, "can't buffer view elementwise"
+  while x.op is not Ops.INDEX:
+    assert x.op in {Ops.BITCAST, Ops.CONTIGUOUS, Ops.SHRINK, Ops.RESHAPE}, f"unexpected op {x.op} in buffer view walk"
     x = x.src[0]
-  x = next(u for u in x.src if u.op is Ops.INDEX)
 
   if len(shape) == 0: offset = x.src[1].arg
-  else: offset = max(sum(idx.vmin for idx in x.src[1:]), 0)
+  else: offset = sum(idx.vmin for idx in x.src[1:])
+  if offset < 0: raise RuntimeError(f"negative offset {offset} in buffer view")
 
-  return b.replace(src=(UOp(Ops.BUFFER_VIEW, t.dtype, (x.base,), (size, offset), tag=t.tag),) + b.src[1:])
+  return b.replace(src=(UOp(Ops.BUFFER_VIEW, t.dtype, (x.base,), (size, offset), tag=t.tag), b.src[1]))
 
 to_bufferview = PatternMatcher([
-  (UPat((Ops.BITCAST, Ops.CONTIGUOUS), name="t").f(Ops.BUFFERIZE, allow_any_len=True, name="b"), late_buffer_view),
-  (UPat((Ops.BITCAST, Ops.CONTIGUOUS)).f(Ops.BUFFER_VIEW, name="b"), lambda b: b.replace(src=b.src[0].src)),
+  (UPat(Ops.BUFFERIZE, src=(UPat((Ops.BITCAST, Ops.CONTIGUOUS), name="t"), UPat()), name="b"), late_buffer_view),
 ])
 
 DEVICE_MAX_BUFS = {"METAL": 31, "WEBGPU": 8} # TODO: get from device?
@@ -451,9 +455,6 @@ to_define_global = PatternMatcher([
   # this is only needed if you are using symbolic
   (UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
 
-  # remove RANGE with 0 size
-  (UPat(Ops.RANGE, name="r"), lambda r: UOp.const(dtypes.index, 0) if r.vmax == 0 else None),
-
   # renumber the ranges starting with 0 so that kernel deduping works
   (UPat(Ops.RANGE, name="r"), renumber_range),
 ])
@@ -469,9 +470,6 @@ rangeify_codegen = PatternMatcher([
   # TODO: this can be moved into codegen?
   (UPat(Ops.NOOP, name="x"), lambda x: x.src[0]),
 
-  # strip the arg from store
-  (UPat(Ops.STORE, name="x"), lambda x: x.replace(arg=None) if x.arg is not None else None),
-
   # add loads to non ptr indexes
   # TODO: this can be moved into codegen?
   #(UPat.any(UPat(Ops.DEFINE_GLOBAL, name="dg"), UPat(Ops.DEFINE_LOCAL).f(Ops.AFTER, allow_any_len=True, name="dg"))

tinygrad/tensor.py

@@ -7,7 +7,7 @@ if TYPE_CHECKING: import numpy
 from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate
 from tinygrad.dtype import _from_np_dtype, _to_np_dtype, PyConst
 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, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, is_numpy_ndarray, TracingKey, cpu_profile
+from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ASM_GEMM, ceildiv, fetch, polyN, 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
@@ -286,20 +286,20 @@ class Tensor(OpMixin):
     return self
 
   def assign(self, x:Tensor|PyConst|list|tuple) -> Tensor:
+    is_disk = isinstance(self.device, str) and self.device.startswith("DISK")
+    if not isinstance(x, Tensor): x = Tensor(x, device="CPU" if is_disk else self.device, dtype=self.dtype)
+    if self.uop is x.uop: return self  # a self assign is a NOOP
+    # broadcast x (shape only, dtype must match)
+    if self.shape != x.shape: x = x._broadcast_to(self.shape)
+    if self.shape != x.shape: raise RuntimeError(f"assign shape mismatch {self.shape} != {x.shape}")
+    if not is_disk and self.device != x.device: raise RuntimeError(f"assign device mismatch {self.device} != {x.device}")
+    if self.dtype != x.dtype: raise RuntimeError(f"assign dtype mismatch {self.dtype} != {x.dtype}")
+    if isinstance(self.device, tuple) and self.uop.axis != x.uop.axis: raise RuntimeError(f"multi axis mismatch {self.uop.axis} != {x.uop.axis}")
+
     # TODO: this is a hack for writing to DISK. remove with working assign
-    if isinstance(self.device, str) and self.device.startswith("DISK"):
-      if not isinstance(x, Tensor): x = Tensor(x, device="CPU", dtype=self.dtype)
+    if is_disk:
       self._buffer().copyin(x._data())
       return self
-    if not isinstance(x, Tensor): x = Tensor(x, device=self.device, dtype=self.dtype)
-    if self.uop is x.uop: return self  # a self assign is a NOOP
-    # NOTE: we allow cross device assign
-    # broadcast x
-    if least_upper_dtype(self.dtype, x.dtype) == self.dtype: x = x._broadcast_to(self.shape).cast(self.dtype)
-    assert self.shape == x.shape, f"assign shape mismatch {self.shape} != {x.shape}"
-    assert self.device == x.device, f"assign device mismatch {self.device} != {x.device}"
-    assert self.dtype == x.dtype, f"assign dtype mismatch {self.dtype} != {x.dtype}"
-    assert not isinstance(self.device, tuple) or self.uop.axis == x.uop.axis, f"multi assign axis mismatch {self.uop.axis} != {x.uop.axis}"
     return self.replace(self._apply_uop(UOp.assign, x))
 
   def detach(self) -> Tensor:
@@ -2431,6 +2431,9 @@ class Tensor(OpMixin):
     ```
     """
     if IMAGE: return self.image_dot(w, dtype)
+    if ASM_GEMM:
+      from extra.gemm.asm.cdna.gemm import can_use_asm_gemm, asm_gemm
+      if can_use_asm_gemm(self, w): return asm_gemm(self, w)
     x, dx, dw = self, self.ndim, w.ndim
     if not (dx > 0 and dw > 0): raise RuntimeError(f"both tensors need to be at least 1D, got {dx}D and {dw}D")
     if x.shape[-1] != w.shape[axis_w:=-min(w.ndim,2)]: raise RuntimeError(f"cannot dot {x.shape} and {w.shape}")
@@ -3865,7 +3868,10 @@ class Tensor(OpMixin):
 
   def bitcast(self, dtype:DTypeLike) -> Tensor:
     """
-    Bitcasts `self` to the given `dtype` of the same itemsize.
+    Bitcasts `self` to the given `dtype`.
+
+    When the target dtype has the same itemsize, this is a view of the same memory.
+    When itemsizes differ, the last dimension is adjusted and a new Tensor is created.
 
     `self` must not require a gradient.
 

tinygrad/uop/__init__.py

@@ -14,7 +14,7 @@ class Ops(FastEnum):
   # ** 1 -- defines/special **
 
   # define GLOBAL/VAR are ptrs to outside the Kernel
-  DEFINE_GLOBAL = auto(); DEFINE_VAR = auto(); BIND = auto()
+  DEFINE_VAR = auto(); BIND = auto()
 
   # this is a RANGE for GPU dimensions, similar to symbolic shapes but not exactly
   SPECIAL = auto()
@@ -28,6 +28,9 @@ class Ops(FastEnum):
   NOOP = auto(); REWRITE_ERROR = auto()
   PARAM = auto(); CALL = auto()
 
+  # TODO: remove this alias, DEFINE_GLOBAL is PARAM now
+  DEFINE_GLOBAL = PARAM
+
   # renderer
   # LINEAR is a list of UOps, SOURCE has a str arg that's human readable, BINARY has bytes arg that's compiled
   PROGRAM = auto(); LINEAR = auto(); SOURCE = auto(); BINARY = auto()

tinygrad/uop/ops.py

@@ -206,7 +206,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     match self.op:
       # late ops don't have shape
       case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.RANGE | Ops.LOAD | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
-           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.CUSTOM_KERNEL | \
+           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.CUSTOM_KERNEL | Ops.SINK | \
            Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY:
         return None
 
@@ -224,8 +224,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       case Ops.BUFFER_VIEW: return (self.arg[0],)
       case Ops.ENCDEC: return self.arg[0]
       case Ops.BUFFERIZE: return tuple([int(r.vmax+1) for r in self.src[1:]])
-      case Ops.DEFINE_GLOBAL | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)
+      case Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)
       case Ops.PARAM:
+        if isinstance(self.dtype, PtrDType): return (self.ptrdtype.size,)
         # NOTE: copied from marg
         if len(self.src) >= 1: return tuple(self.src[0].sgep(i) for i in range(self.src[0].dtype.count))
         return None
@@ -1436,6 +1437,7 @@ def pyrender(ast:UOp) -> str:
       for s in u.src: to_render.add(s)
     if u.op is Ops.STORE: to_render.add(u.src[1])
     if u.op in {Ops.REDUCE, Ops.REDUCE_AXIS}: to_render.add(u.src[0])
+    if u.op in {Ops.CUSTOM_KERNEL, Ops.CALL}: raise NotImplementedError("custom_kernel / call can't be pyrendered")
     if u.op in not_rendered: continue
     # checking the consumers is not enough, you have to make sure it's not used twice by the one consumer
     if len(cmap[u]) == 1 and len([x for x in list(cmap[u].keys())[0].src if x is u]) == 1 and u.op not in always_rendered: continue
@@ -1444,7 +1446,13 @@ def pyrender(ast:UOp) -> str:
   kernels: dict[UOp, tuple[str, str]] = {}
   r: dict[UOp, str] = {}
   ret: dict[str, str] = {}
+  depth: dict[UOp, int] = {}
   for i,u in enumerate(lst):
+    # limit inline depth to avoid "too many nested parentheses" in Python parser
+    op_depth = 1 + max([depth[s] for s in u.src], default=0)
+    if op_depth > 100: to_render.add(u)
+    depth[u] = 0 if u in to_render else op_depth
+    # do the rendering
     if u.op is Ops.KERNEL:
       if u.arg.ast not in kernels:
         kernels[u.arg.ast] = (f"k{len(kernels)}", f"def k{len(kernels)}():\n  " + pyrender(u.arg.ast).replace('\n', '\n  ') + "\n  return ast\n\n")

tinygrad/uop/spec.py

@@ -84,14 +84,11 @@ _tensor_spec = PatternMatcher([
   (UPat(Ops.LUNIQUE, dtypes.void, ()), lambda: True),
   (UPat(Ops.DEVICE, dtypes.void, (), name="d"), lambda d:
    isinstance(d.arg, str) or (isinstance(d.arg, tuple) and all(isinstance(s, str) for s in d.arg))),
-  (UPat(Ops.BUFFER, src=(UPat((Ops.LUNIQUE, Ops.UNIQUE)), UPat(Ops.DEVICE)), allow_any_len=True, name="buf"),
+  (UPat(Ops.BUFFER, src=(UPat((Ops.LUNIQUE, Ops.UNIQUE)), UPat(Ops.DEVICE)), name="buf"),
    lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, (DType, ImageDType))),
-  (UPat(Ops.BUFFER_VIEW, src=(UPat(Ops.BUFFER),), name="buf_view"),
-   lambda buf_view: isinstance(buf_view.arg, tuple) and len(buf_view.arg) == 2 and all(isinstance(arg, (int, UOp)) for arg in buf_view.arg)),
-  (UPat(Ops.BUFFER_VIEW, src=(UPat(Ops.MSTACK, src=UPat(Ops.BUFFER)),)), lambda: True),
 
   # KERNEL can attach to an AFTER to describe the compute required to realize a BUFFER
-  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.BUFFER_VIEW, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND, Ops.CONTIGUOUS))), lambda: True),
+  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), lambda: True),
 
   # ASSIGN has a target and a value. It can also optionally depend on other assigns
   (UPat(Ops.ASSIGN, name="x"), lambda x: len(x.src) >= 2 and all(s.op is Ops.ASSIGN for s in x.src[2:])),
@@ -316,7 +313,8 @@ def eval_pyrender(code:str) -> UOp:
   return lcls['ast']
 
 def test_pyrender(test_ast:UOp, assert_parents=True):
-  code = pyrender(test_ast)
+  try: code = pyrender(test_ast)
+  except NotImplementedError: return None  # this is okay, not all ops can be pyrendered
   ast:UOp = eval_pyrender(code)
   if ast is not test_ast:
     if assert_parents:

tinygrad/viz/serve.py

@@ -422,16 +422,14 @@ def get_stdout(f: Callable) -> str:
   return buf.getvalue()
 
 def amd_readelf(lib:bytes) -> list[dict]:
+  from tinygrad.runtime.autogen import amdgpu_kd
   from tinygrad.runtime.support.elf import elf_loader
-  import msgpack
-  _, sections, __ = elf_loader(lib)
-  data = next((s for s in sections if s.name.startswith(".note"))).content
-  namesz, descsz, typ = struct.unpack_from(hdr:="<III", data, 0)
-  offset = (struct.calcsize(hdr)+namesz+3) & -4
-  notes = msgpack.unpackb(data[offset:offset+descsz])
-  keys = {".sgpr_count":"SGPRs", ".vgpr_count":"VGPRs", ".max_flat_workgroup_size":"Max WGP size",
-          ".group_segment_fixed_size":"LDS size", ".private_segment_fixed_size":"Scratch size"}
-  return [{"label":label, "value":v} for k,label in keys.items() if (v:=notes["amdhsa.kernels"][0][k]) > 0]
+  image, sections, __ = elf_loader(lib)
+  rodata = next((s for s in sections if s.name == ".rodata")).content
+  kd = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t.from_buffer_copy(bytearray(rodata))
+  vgpr_gran = kd.compute_pgm_rsrc1 & amdgpu_kd.COMPUTE_PGM_RSRC1_GRANULATED_WORKITEM_VGPR_COUNT
+  return [{"label":f"{resource} Alloc", "value":val} for resource,val in [("VGPR", (vgpr_gran+1)*8-7), ("LDS",kd.group_segment_fixed_size),
+                                                                          ("Scratch", kd.private_segment_fixed_size)] if val > 0]
 
 def amd_decode(lib:bytes, target:int) -> dict[int, Any]: # Any is the Inst class from extra.assembly.amd.dsl
   from tinygrad.runtime.support.elf import elf_loader
@@ -558,7 +556,7 @@ def get_render(query:str) -> dict:
     rows:dict[int, dict] = {}
     for pc, (inst,_) in pc_to_inst.items():
       if start_pc is None: start_pc = pc
-      rows[pc] = {"pc":pc-start_pc, "inst":inst, "hit_count":0, "dur":0, "stall":0, "hits":{"cols":inst_columns, "rows":[]}, "type":""}
+      rows[pc] = {"pc":pc-start_pc, "inst":inst, "hit_count":0, "dur":0, "stall":0, "type":"", "hits":{"cols":inst_columns, "rows":[]}}
     for e in w.unpack_insts():
       if not (inst:=rows[e.pc]).get("type"): inst["type"] = str(e.typ).split("_")[-1]
       inst["hit_count"] += 1