Merge branch 'master' into pyremu_try

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

@@ -10,7 +10,6 @@ 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,19 +402,17 @@ 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 cross-lane ops (READLANE/WRITELANE/PERMLANE) using pcode parser."""
+    """Compile READLANE/READFIRSTLANE/WRITELANE 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 = []
@@ -429,8 +427,7 @@ 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,
-                 'ROUND_MODE': _c(0), 'ROUND_TOWARD_ZERO': _c(0)})  # rounding mode: 0=RNE, RTZ constant
+    srcs.update({'EXEC': exec_mask, 'SCC': self.rsgpr_dyn(_c(SCC.offset)), 'laneId': lane})
     _, assigns = parse_pcode(pcode, srcs)
 
     raw_stores: list = []
@@ -487,7 +484,6 @@ 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)
@@ -555,7 +551,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 in ('V_READFIRSTLANE_B32_E32', 'V_PERMLANE64_B32_E32'): return ctx.compile_lane_pcode(inst.op, inst)
+  if op_name == 'V_READFIRSTLANE_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)
@@ -648,10 +644,6 @@ 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)
@@ -795,7 +787,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 = {'S@0': s0_mod, 'S@1': s1_mod, 'S@2': s2_mod,
+    srcs = {'S0': s0_mod, 'S1': s1_mod, 'S2': 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:
@@ -804,13 +796,10 @@ 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))
-    # 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)
+    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}
   return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask)
 
 def _compile_vopd(inst: VOPD, ctx: _Ctx) -> UOp:
@@ -879,15 +868,6 @@ 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,19 +94,13 @@ 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 [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
+  for bits, mask, ext in [(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)
@@ -200,17 +194,6 @@ 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),
@@ -244,50 +227,11 @@ _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
@@ -295,7 +239,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, 'u4': dtypes.uint8, 'i4': dtypes.int8, 'u1': dtypes.uint32}
+          'u8': dtypes.uint8, 'i8': dtypes.int8, 'b8': dtypes.uint8, '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]:
@@ -452,7 +396,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' and self.at('LBRACKET'):
+      if name == 'VGPR':
         self.eat('LBRACKET')
         lane = self.parse()
         self.eat('RBRACKET')
@@ -479,21 +423,7 @@ class Parser:
       if self.try_eat('LBRACE'):
         idx = self.eat('NUM').val
         self.eat('RBRACE')
-        # 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)
+        elem = self.vars.get(f'{name}{idx}', _u32(0))
         if self.try_eat('DOT'):
           dt_name = self.eat('IDENT').val
           return _cast_to(elem, DTYPES.get(dt_name, dtypes.uint32))
@@ -502,17 +432,27 @@ 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]
-        assert isinstance(v, UOp), f"expected UOp for {name}, got {type(v)}"
-        return v
+        return v if isinstance(v, UOp) else _u32(0) if isinstance(v, dict) else _u32(0)
       raise RuntimeError(f"unknown variable: {name}")
     raise RuntimeError(f"unexpected token in primary: {self.peek()}")
 
   def _handle_dot(self, base, field: str) -> UOp:
-    assert isinstance(base, UOp), f"expected UOp for dot access, got {type(base)}"
+    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)
     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')
@@ -527,7 +467,6 @@ 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:
@@ -570,18 +509,16 @@ class Parser:
       var_name = self._find_var_name(base)
     if first.op == Ops.CONST:
       idx = int(first.arg)
-      # Check for array element (var@idx)
-      if var_name and f'{var_name}@{idx}' in self.vars:
-        v = self.vars[f'{var_name}@{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]):
@@ -600,7 +537,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))  # use @ to avoid collision with temps like A4
+      elem = self.vars.get(f'{var_name}{idx}', _u32(0))
       if self.try_eat('DOT'):
         dt_name = self.eat('IDENT').val
         return _cast_to(elem, DTYPES.get(dt_name, dtypes.uint32))
@@ -662,14 +599,13 @@ 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'):
-      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 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 'U' in suffix: return _const(dtypes.uint32, val)
     return _const(dtypes.int if val < 0 else dtypes.uint32, val)
 
@@ -687,8 +623,7 @@ 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)
-        assert result is not None, f"lambda {name} must return a value"
-        return result
+        return result if result is not None else _u32(0)
       return parse_expr(body, lv, self.funcs)
     if name in self.funcs:
       return self.funcs[name](*args)
@@ -696,7 +631,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')
-    assert mem is not None, "memory load requires _vmem or _lds"
+    if mem is None: return _const(dt, 0)
     adt = dtypes.uint64 if addr.dtype == dtypes.uint64 else dtypes.uint32
     active = self.vars.get('_active')
     gate = (active,) if active is not None else ()
@@ -758,9 +693,29 @@ 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 with its value."""
+  """Substitute loop variable and evaluate bracket expressions.
+  Converts var[loop_var] to var{val} for array element access (like the old regex parser)."""
   toks = tokenize(line)
-  return ' '.join(str(val) if t.type == 'IDENT' and t.val == loop_var else t.val for t in toks if t.type != 'EOF')
+  # 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)
 
 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."""
@@ -810,9 +765,8 @@ 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().simplify()
-        assert expr.op == Ops.CONST, f"loop bound must be constant, got {expr}"
-        return int(expr.arg)
+        expr = p.parse()
+        return int(expr.arg) if expr.op == Ops.CONST else 0
       start_val = parse_bound()
       p.eat('COLON')
       end_val = parse_bound()
@@ -833,7 +787,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, **block_assigns}, funcs, assigns)
+        _, iter_assigns, _ = parse_block(subst_lines, 0, vars, funcs, assigns)
         if has_break:
           assert found_var is not None
           found = block_assigns.get(found_var, vars.get(found_var))
@@ -858,9 +812,7 @@ def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: di
 
     # declare
     if first == 'declare':
-      # 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))
+      if '[' not in line and len(toks) >= 2 and toks[1].type == 'IDENT': vars[toks[1].val] = _u32(0)
       i += 1; continue
 
     # lambda definition
@@ -918,7 +870,6 @@ 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)))
@@ -982,32 +933,19 @@ 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 (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
+    # 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
 
     # Compound assignment: var += or var -=
     assign_op = next((j for j, t in enumerate(toks) if t.type == 'ASSIGN_OP'), None)
@@ -1054,14 +992,13 @@ 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 is_const(c, v): return c.op == Ops.CONST and c.arg is v
+      def not_static_false(c): return c.op != Ops.CONST or c.arg is not False
       cond = parse_cond(line, 'if')
-      conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not is_const(cond, False) else []
+      conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not_static_false(cond) 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 if not is_const(cond, False) else None)
+      i, branch, ret = parse_block(lines, i, vars, funcs, assigns)
       if conditions: conditions[0] = (cond, ret if ret is not None else branch)
       vars.clear(); vars.update(vars_snap)
       while i < len(lines):
@@ -1070,16 +1007,12 @@ 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')
-          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
+          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))
           vars.clear(); vars.update(vars_snap)
         elif lf == 'else':
-          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)
+          i += 1; i, branch, ret = parse_block(lines, i, vars, funcs, assigns)
+          else_branch = (ret, branch)
           vars.clear(); vars.update(vars_snap)
         elif lf == 'endif': i += 1; break
         else: break
@@ -1091,21 +1024,17 @@ 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
-      # 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
+      # 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
       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,23 +43,6 @@ 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
@@ -229,12 +212,8 @@ 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, 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)
-  """
+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."""
   import math
   diffs = []
   for i in range(n_vgprs):
@@ -245,11 +224,6 @@ 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]
@@ -262,19 +236,16 @@ 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, ulp_tolerance: int = 0) -> WaveState:
+def run_program(instructions: list, n_lanes: int = 1) -> 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, ulp_tolerance=ulp_tolerance)
+    diffs = compare_wave_states(emu_st, hw_st, n_lanes)
     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,47 +719,5 @@ 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,7 +62,6 @@ 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
@@ -74,7 +73,6 @@ 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,6 +255,7 @@ 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]),
@@ -292,6 +293,7 @@ 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]),
@@ -300,7 +302,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."""
@@ -694,6 +696,7 @@ 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),
@@ -706,6 +709,7 @@ 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),
@@ -718,6 +722,7 @@ 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),
@@ -730,6 +735,7 @@ 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),
@@ -742,6 +748,7 @@ 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)),
@@ -750,8 +757,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)
 
@@ -797,6 +804,7 @@ 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),
@@ -808,6 +816,7 @@ 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),
@@ -1503,82 +1512,5 @@ 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,6 +857,7 @@ 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 = [
@@ -867,12 +868,13 @@ 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 = [
@@ -883,12 +885,13 @@ 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 = [
@@ -899,12 +902,13 @@ 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 = [
@@ -915,11 +919,12 @@ 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 = [
@@ -930,7 +935,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):
@@ -938,6 +943,7 @@ 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]),
@@ -948,8 +954,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}")
 
@@ -2949,394 +2955,5 @@ 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,6 +149,7 @@ 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)),
@@ -165,6 +166,7 @@ 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 = [
@@ -197,6 +199,7 @@ 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)
@@ -214,6 +217,7 @@ 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)
@@ -231,6 +235,7 @@ 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 = [
@@ -249,6 +254,7 @@ 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]),
@@ -261,13 +267,14 @@ 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]),
@@ -279,12 +286,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)
     # 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]),
@@ -297,7 +305,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)
 
 
@@ -306,6 +314,7 @@ 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
@@ -315,13 +324,14 @@ 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
@@ -331,13 +341,14 @@ 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
@@ -349,8 +360,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
 
@@ -359,6 +370,7 @@ 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]),
@@ -366,8 +378,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)
@@ -376,6 +388,7 @@ 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 = [
@@ -385,8 +398,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)
 
@@ -400,6 +413,7 @@ 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)
@@ -418,12 +432,13 @@ 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
@@ -443,7 +458,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")
 
@@ -453,6 +468,7 @@ 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)
@@ -474,7 +490,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")
 
@@ -697,6 +713,7 @@ 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
@@ -706,13 +723,14 @@ 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
@@ -725,277 +743,5 @@ 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(), 'ROUND_TOWARD_ZERO': u32(), 'ROUND_NEAREST_EVEN': u32(), 'WAVE_STATUS': u32(),
+            'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': 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.renderer.arch}")
+  print(f"Device arch: {dev.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.renderer.arch)
+    asm = build_kernel(dev.arch)
 
   binary = dev.compiler.compile(asm)
   print(f"Compiled! Binary size: {len(binary)} bytes")

extra/gemm/asm/cdna/gemm.py

@@ -1,95 +0,0 @@
-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

@@ -0,0 +1,78 @@
+.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

@@ -0,0 +1,73 @@
+# 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

@@ -1,46 +0,0 @@
-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 os, zlib, struct, hashlib
-from tinygrad.helpers import getenv
+import sys, os, zlib, struct, hashlib
+from tinygrad.helpers import DEBUG, getenv, fetch
 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, use_bot=True)
+    dev = USB3(vendor, device, 0x81, 0x83, 0x02, 0x04)
     break
   except RuntimeError: pass
 if dev is None:

test/test_edgecases.py

@@ -195,10 +195,8 @@ 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,21 +1,11 @@
 # 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_memory_planner.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import dtypes
+from tinygrad import dtypes, Device
 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("NULL", size, dtypes.int8, base=global_map[base] if base is not None else None, offset=offset)
+  global_map[i] = Buffer(Device.DEFAULT, 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/test_multitensor.py

@@ -256,18 +256,6 @@ 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))
@@ -1285,20 +1273,19 @@ 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, dev_count:int):
+  def _test_matmul_half(self, devs):
     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, device=devs[0]).shard(devs, axis=0)
-      b = Tensor.empty((N, N), dtype=dtype, device=devs[0]).shard(devs, axis=None)
+      a = Tensor.empty((N, N), dtype=dtype).shard(devs, axis=0)
+      b = Tensor.empty((N, N), dtype=dtype).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(dev_count=2)
-  def test_matmul_half_alt(self): self._test_matmul_half(dev_count=4)
+  def test_matmul_half(self): self._test_matmul_half(devices_2)
+  def test_matmul_half_alt(self): self._test_matmul_half(devices_4)
 
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiFromUnrenderable(unittest.TestCase):

test/test_uop_graph.py

@@ -479,26 +479,6 @@ 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)
@@ -706,97 +686,6 @@ 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/testextra/test_cfg_viz.py

@@ -2,7 +2,7 @@
 # allow define from star imports
 
 import unittest
-import functools
+import textwrap, functools
 
 from tinygrad import Device, Tensor
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
@@ -11,27 +11,69 @@ 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 s
+from extra.assembly.amd.dsl import Inst
 
-# TODO: this belongs to the dsl infrastructure
-from extra.gemm.amd_asm_matmul import Kernel
+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: shouldn't need compiler once we can output ELF
 # outputs a text disassembly for humans and a machine readable binary
-def assemble(name:str, k:Kernel, compiler:Compiler) -> tuple[str, bytes]:
-  src = k.to_asm()
+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))
   return (src, compiler.compile(src))
 
-def asm_kernel(out:UOp, k:Kernel, name:str, device:str, compiler:Compiler, n_threads:int=1, n_workgroups:int=1) -> UOp:
+def asm_kernel(out:UOp, insts:list[str|Inst], 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, k, compiler)
+  src, lib = assemble(name, insts, 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, k:Kernel) -> None:
-  fxn = functools.partial(asm_kernel, k=k, name=name, device=Device.DEFAULT, compiler=HIPCompiler(Device[Device.DEFAULT].renderer.arch))
+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))
   out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
   out.realize()
 
@@ -43,32 +85,32 @@ class TestCfg(unittest.TestCase):
       self.skipTest(f"tests written for RDNA, got arch {arch}")
 
   def test_simple(self):
-    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)
+    run_asm("simple", [
+      "entry:",
+        "s_branch bb1",
+      "bb1:",
+        s_endpgm(),
+        s_code_end(),
+    ])
 
   def test_diamond(self):
-    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))
+    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))
     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())
@@ -82,138 +124,133 @@ class TestCfg(unittest.TestCase):
     self.assertEqual(insts, ['s_mov_b32', 's_cmp_eq_u64'])
 
   def test_loop(self):
-    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)
+    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(),
+    ])
 
   def test_loop_branch(self):
-    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)
+    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(),
+    ])
 
   def test_loop_break(self):
-    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)
+    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(),
+    ])
 
   def test_switch(self):
-    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)
+    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(),
+    ])
 
   def test_ping_pong(self):
-    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)
+    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(),
+    ])
 
   def test_colored_blocks(self):
     N = 10
-    k = Kernel(arch=Device["AMD"].arch)
-    k.label("entry")
-    k.emit(s_branch(), target="init0")
+    asm = ["entry:", "s_branch init0"]
     for i in range(N):
-      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)
+      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)
 
   def test_jump_back_to_end(self):
-    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)
+    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(),
+    ])
 
   def test_hit_count(self):
-    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)
+    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(),
+    ])
 
 if __name__ == "__main__":
   unittest.main()

test/unit/test_assign.py

@@ -461,27 +461,6 @@ 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)
@@ -493,38 +472,6 @@ 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))
@@ -619,12 +566,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, dtype=dtypes.int32).contiguous().realize()
+    buf = Tensor.zeros(4).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, dtype=dtypes.int32).contiguous().realize()
+    buf = Tensor.zeros(4).contiguous().realize()
     buf[1:3].assign(Tensor([5, 6])).realize()
     np.testing.assert_equal(buf.numpy(), [0, 5, 6, 0])
 
@@ -706,7 +653,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, dtype=dtypes.int32).contiguous().realize()
+    a = Tensor.zeros(4).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
@@ -718,8 +665,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, dtype=dtypes.int32).contiguous().realize()
-    b = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
+    a = Tensor.zeros(4).contiguous().realize()
+    b = Tensor.zeros(4).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,13 +5,24 @@ from tinygrad.runtime.support.c import DLL, record, init_records
 from tinygrad.runtime.support import c
 from tinygrad.runtime.support.autogen import gen
 
-@unittest.skipIf(WIN, "doesn't compile on windows")
-class TestC(unittest.TestCase):
+class TestAutogen(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:
@@ -34,6 +45,7 @@ class TestC(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:
@@ -63,6 +75,7 @@ class TestC(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:
@@ -90,6 +103,7 @@ class TestC(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:
@@ -117,6 +131,7 @@ class TestC(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:
@@ -136,6 +151,7 @@ class TestC(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:
@@ -157,6 +173,7 @@ class TestC(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:
@@ -174,6 +191,7 @@ class TestC(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:
@@ -199,6 +217,7 @@ class TestC(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:
@@ -223,88 +242,7 @@ class TestC(unittest.TestCase):
     self.assertEqual(out.contents.a, 20)
     self.assertEqual(out.contents.b, 10)
 
-  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
-
+  @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_packed_structs(self):
     ns = self.run_gen("""
 typedef unsigned NvU32;
@@ -354,6 +292,47 @@ 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 {
@@ -399,6 +378,7 @@ typedef struct
     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;
@@ -428,6 +408,77 @@ typedef struct
     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 {
@@ -440,6 +491,7 @@ typedef struct
     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 };
@@ -459,43 +511,4 @@ typedef struct
     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, atol=1e-6)
+    np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5)
 
   @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, atol=1e-6)
+    np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5)
 
 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.assign(Tensor.full((128, 64), 0x4248, dtype=dtypes.uint16).bitcast(dtypes.uint8)).realize()
+    t.bitcast(dtypes.uint16).assign(Tensor.full((128, 64), 0x4248, dtype=dtypes.uint16)).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.assign(Tensor.full((128, 32), 3.1415927, dtype=dtypes.float32).bitcast(dtypes.uint8)).realize()
+    t.bitcast(dtypes.float32).assign(Tensor.full((128, 32), 3.1415927, dtype=dtypes.float32)).realize()
     _test_bitcasted(t, dtypes.float32, 3.1415927)
     _test_bitcasted(t, dtypes.uint32, 0x40490FDB)
     # doesn't suport normal cast
@@ -250,24 +250,6 @@ 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)
 
@@ -348,25 +330,13 @@ class TestDiskTensor(unittest.TestCase):
 
   def test_assign_with_bitcast(self):
     # bitcast assign is used in safe_save for writing header length
-    # bitcast on source side works, bitcast on target side raises
+    # this tests the synchronous disk assign hack handles bitcast correctly
     pathlib.Path(temp(fn:="dt_assign_bitcast")).unlink(missing_ok=True)
     t = Tensor.empty(16, device=f"disk:{temp(fn)}", dtype=dtypes.uint8)
-    # correct way: bitcast the source to match target dtype
-    t[0:8].assign(Tensor([12345], dtype=dtypes.int64, device="CPU").bitcast(dtypes.uint8))
+    t[0:8].bitcast(dtypes.int64).assign([12345])
+    # verify the data was written correctly
     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_schedule_cache.py

@@ -35,6 +35,7 @@ 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)
@@ -42,6 +43,7 @@ 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,19 +470,5 @@ 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_symbolic.py

@@ -1074,24 +1074,6 @@ 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())),
-   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()),
+        allow_any_len=True), 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")), name="sto"), gep_on_store),
+  (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st")), allow_any_len=True, 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")), name="sto"), cat_after_store),
+  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), allow_any_len=True, 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), UPat(name="val")), name="s"), lambda s,val: s.replace(src=(s.src[0].src[0], val))),
+  (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:])),
 ])
 

tinygrad/helpers.py

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

tinygrad/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].assign(Tensor([len(j)], dtype=dtypes.int64, device="CPU").bitcast(dtypes.uint8))
+  t[0:8].bitcast(dtypes.int64).assign([len(j)])
   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"))),
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat.var("gate"))).or_casted("bidx"), UPat.var("var")), allow_any_len=True),
    lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"),
-  (UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"(*{ctx[bidx]})"),
-  (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
+  (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]};"),
   # 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,13 +4,12 @@ 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
-DLL("CoreGraphics", "CoreGraphics")
+ctypes.CDLL("/System/Library/Frameworks/CoreGraphics.framework/CoreGraphics")
 
 # FIXME: these need autogen to support objc categories
 # https://developer.apple.com/library/archive/documentation/Cocoa/Conceptual/ObjectiveC/Chapters/ocCategories.html
@@ -68,7 +67,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 = DLL("MTLCompiler", "MTLCompiler")
+  support = ctypes.CDLL("/System/Library/PrivateFrameworks/MTLCompiler.framework/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 = 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")
+    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")
 
   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,15 +131,13 @@ 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", f"/System/Library/PrivateFrameworks/{p}.framework"],
+                  "darwin": ["/opt/homebrew/lib", f"/System/Library/Frameworks/{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
@@ -156,12 +154,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 = nm, emsg
+    self.nm, self.emsg, self.loaded = nm, emsg, False
     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_.add(self.nm)
+        self.loaded = True
       except OSError as e:
         self.emsg = str(e)
         if DEBUG >= 3: print(f"loading {nm} failed: {e}")
@@ -177,6 +175,5 @@ class DLL(ctypes.CDLL):
     return wrapper
 
   def __getattr__(self, nm):
-    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'}?"))
+    if not 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, use_bot=False):
+  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):
     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, self.use_bot = max_streams, use_bot
+    self.max_streams = max_streams
     self.ctx = ctypes.POINTER(libusb.struct_libusb_context)()
 
     if libusb.libusb_init(ctypes.byref(self.ctx)): raise RuntimeError("libusb_init failed")
@@ -25,34 +25,30 @@ 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")
 
-    if use_bot:
-      self._tag = 0
-    else:
-      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)
 
-      # 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)
+    # 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}")
 
-      # 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}")
+    # Base cmd
+    cmd_template = bytes([0x01, 0x00, 0x00, 0x01, *([0] * 12), 0xE4, 0x24, 0x00, 0xB2, 0x1A, 0x00, 0x00, 0x00, *([0] * 8)])
 
-      # Base cmd
-      cmd_template = bytes([0x01, 0x00, 0x00, 0x01, *([0] * 12), 0xE4, 0x24, 0x00, 0xB2, 0x1A, 0x00, 0x00, 0x00, *([0] * 8)])
+    # Init pools
+    self.tr = {ep: [libusb.libusb_alloc_transfer(0) for _ in range(self.max_streams)] for ep in all_eps}
 
-      # Init pools
-      self.tr = {ep: [libusb.libusb_alloc_transfer(0) for _ in range(self.max_streams)] for ep in all_eps}
+    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)]
 
-      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)
+    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
@@ -72,90 +68,38 @@ 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:list[bytes|None] = []
-    tr_window, op_window = [], []
+    results, tr_window, op_window = [], [], []
 
     for idx, (cdb, rlen, send_data) in enumerate(zip(cdbs, idata, odata)):
-      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"
+      # allocate slot and stream. stream is 1-based
+      slot, stream = idx % self.max_streams, (idx % self.max_streams) + 1
 
-        # 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)
+      # build cmd packet
+      self.buf_cmd[slot][16:16+len(cdb)] = list(cdb)
 
-        # 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)
+      # 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))
 
-        # 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 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))
 
-        # build cmd packet
-        self.buf_cmd[slot][16:16+len(cdb)] = list(cdb)
+      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))
 
-        # 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))
+        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)))
 
-        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 = []
+      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,8 +71,7 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
   # allgather
   copied_chunks = []
   for i,rc in enumerate(reduced_chunks):
-    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))))
+    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))))
     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, src=(UPat(Ops.RESHAPE, name="x2"), UPat()), name="x"),
+  (UPat(Ops.RESHAPE, name="x2").f(Ops.RESHAPE, allow_any_len=True, 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.RESHAPE, src=(UPat(Ops.BUFFER), UPat()), name="r").f(Ops.CONTIGUOUS, name="c"), lambda r,c: r.replace(tag=c.tag)),
+  (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)),
 
   # split_reduceop
   (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)), split_reduceop),
@@ -126,10 +126,6 @@ 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 \
@@ -251,7 +247,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 (ranges don't matter for const - it's the same value everywhere)
+  # no buffers for const
   (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),
@@ -275,20 +271,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 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"
+  while not any(u.op is Ops.INDEX for u in x.src):
+    assert x.op not in GroupOp.Elementwise, "can't buffer view elementwise"
     x = x.src[0]
+  x = next(u for u in x.src if u.op is Ops.INDEX)
 
   if len(shape) == 0: offset = x.src[1].arg
-  else: offset = sum(idx.vmin for idx in x.src[1:])
-  if offset < 0: raise RuntimeError(f"negative offset {offset} in buffer view")
+  else: offset = max(sum(idx.vmin for idx in x.src[1:]), 0)
 
-  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.BUFFERIZE, src=(UPat((Ops.BITCAST, Ops.CONTIGUOUS), name="t"), UPat()), name="b"), late_buffer_view),
+  (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)),
 ])
 
 DEVICE_MAX_BUFS = {"METAL": 31, "WEBGPU": 8} # TODO: get from device?
@@ -455,6 +451,9 @@ 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),
 ])
@@ -470,6 +469,9 @@ 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, ASM_GEMM, ceildiv, fetch, polyN, is_numpy_ndarray, TracingKey, cpu_profile
+from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, 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 is_disk:
+    if isinstance(self.device, str) and self.device.startswith("DISK"):
+      if not isinstance(x, Tensor): x = Tensor(x, device="CPU", dtype=self.dtype)
       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,9 +2431,6 @@ 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}")
@@ -3868,10 +3865,7 @@ class Tensor(OpMixin):
 
   def bitcast(self, dtype:DTypeLike) -> Tensor:
     """
-    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.
+    Bitcasts `self` to the given `dtype` of the same itemsize.
 
     `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_VAR = auto(); BIND = auto()
+  DEFINE_GLOBAL = auto(); DEFINE_VAR = auto(); BIND = auto()
 
   # this is a RANGE for GPU dimensions, similar to symbolic shapes but not exactly
   SPECIAL = auto()
@@ -28,9 +28,6 @@ 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.SINK | \
+           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.CUSTOM_KERNEL | \
            Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY:
         return None
 
@@ -224,9 +224,8 @@ 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_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)
+      case Ops.DEFINE_GLOBAL | 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
@@ -1437,7 +1436,6 @@ 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
@@ -1446,13 +1444,7 @@ 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,11 +84,14 @@ _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)), name="buf"),
+  (UPat(Ops.BUFFER, src=(UPat((Ops.LUNIQUE, Ops.UNIQUE)), UPat(Ops.DEVICE)), allow_any_len=True, 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.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), lambda: True),
+  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.BUFFER_VIEW, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND, Ops.CONTIGUOUS))), 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:])),
@@ -313,8 +316,7 @@ def eval_pyrender(code:str) -> UOp:
   return lcls['ast']
 
 def test_pyrender(test_ast:UOp, assert_parents=True):
-  try: code = pyrender(test_ast)
-  except NotImplementedError: return None  # this is okay, not all ops can be pyrendered
+  code = pyrender(test_ast)
   ast:UOp = eval_pyrender(code)
   if ast is not test_ast:
     if assert_parents:

tinygrad/viz/serve.py

@@ -422,14 +422,16 @@ 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
-  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]
+  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]
 
 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
@@ -556,7 +558,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, "type":"", "hits":{"cols":inst_columns, "rows":[]}}
+      rows[pc] = {"pc":pc-start_pc, "inst":inst, "hit_count":0, "dur":0, "stall":0, "hits":{"cols":inst_columns, "rows":[]}, "type":""}
     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