Reduce redundant FlyDSL numeric wrappers (#448)

Author: coderfeli

kernels/blockscale_preshuffle_gemm.py

@@ -18,8 +18,7 @@
 from flydsl._mlir import ir
 
 from flydsl.expr import arith, gpu, buffer_ops, vector, rocdl
-from flydsl.expr.arith import ArithValue
-from flydsl.expr.typing import T
+from flydsl.expr.typing import T, Vector as Vec
 
 
 from kernels.mfma_preshuffle_pipeline import (
@@ -124,6 +123,9 @@ def compile_blockscale_preshuffle_gemm(
     def _out_elem_type():
         return T.bf16 if is_bf16_out else T.f16
 
+    def _out_elem_dtype():
+        return fx.BFloat16 if is_bf16_out else fx.Float16
+
     epilog_tag = "cshuffle" if use_cshuffle_epilog else "direct"
 
     module_name = (
@@ -168,10 +170,9 @@ def kernel_gemm(
         i32_m: fx.Int32,
         i32_n: fx.Int32,
     ):
-        c_m = arith.index_cast(T.index, i32_m)
-        c_n = arith.index_cast(T.index, i32_n)
+        c_n = fx.Index(i32_n)
 
-        acc_init = arith.constant_vector(0.0, T.f32x4)
+        acc_init = fx.full(4, 0.0, fx.Float32)
 
         # ---- B layout (compile-time since N, K are known) ----
         n0_val = N // 16
@@ -189,7 +190,7 @@ def kernel_gemm(
         stride_lds = fx.make_stride(tile_k, 1)
         layout_lds = fx.make_layout(shape_lds, stride_lds)
 
-        k_blocks16 = arith.index(tile_k_bytes // 16)
+        k_blocks16 = tile_k_bytes // 16
 
         tx = gpu.thread_id("x")
         bx = gpu.block_id("x")
@@ -215,14 +216,14 @@ def kernel_gemm(
             lds_out = None
 
         # ---- Buffer resources (explicit num_records_bytes for correct OOB on M tail) ----
-        rt_M = arith.index_cast(T.index, i32_m)
-        rt_N = arith.index_cast(T.index, i32_n)
-        a_nbytes = rt_M * fx.Index(K)  # fp8: 1 byte/elem
+        rt_M = fx.Index(i32_m)
+        rt_N = fx.Index(i32_n)
+        a_nbytes = rt_M * K  # fp8: 1 byte/elem
         a_rsrc = buffer_ops.create_buffer_resource(arg_a, max_size=False, num_records_bytes=a_nbytes)
         out_elem_bytes = 2  # bf16/fp16
-        c_nbytes = rt_M * rt_N * fx.Index(out_elem_bytes)
+        c_nbytes = rt_M * rt_N * out_elem_bytes
         c_rsrc = buffer_ops.create_buffer_resource(arg_c, max_size=False, num_records_bytes=c_nbytes)
-        sa_nbytes = arith.index(K // 128) * rt_M * fx.Index(4)  # [scale_k, M] f32
+        sa_nbytes = (K // 128) * rt_M * 4  # [scale_k, M] f32
         scale_a_rsrc = buffer_ops.create_buffer_resource(arg_scale_a, max_size=False, num_records_bytes=sa_nbytes)
 
         b_rsrc = buffer_ops.create_buffer_resource(arg_b, max_size=True)
@@ -268,7 +269,7 @@ def load_b_pack(base_k, ki_step, ni):
             )
 
         c64_b = 64
-        _lds_k_dim_c = fx.Index(tile_k)
+        _lds_k_dim_c = tile_k
 
         def load_b_packs_k64(base_k, ku: int, ni: int):
             base_k_bytes = base_k
@@ -282,9 +283,9 @@ def load_b_packs_k64(base_k, ku: int, ni: int):
                 elem_type=T.f8, vec_elems=16, elem_bytes=elem_bytes,
                 offset_in_bytes=True,
             )
-            b_i64x2 = vector.bitcast(T.i64x2, b16)
-            b0_i64 = vector.extract(b_i64x2, static_position=[0], dynamic_position=[])
-            b1_i64 = vector.extract(b_i64x2, static_position=[1], dynamic_position=[])
+            b_i64x2 = Vec(b16).bitcast(fx.Int64)
+            b0_i64 = b_i64x2[0]
+            b1_i64 = b_i64x2[1]
             return b0_i64, b1_i64
 
         def load_b_tile(base_k):
@@ -307,9 +308,9 @@ def lds_load_16b(curr_row_a_lds, col_base, lds_buffer):
 
         def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer):
             loaded_a16 = lds_load_16b(curr_row_a_lds, col_base, lds_buffer)
-            a_i64x2 = vector.bitcast(T.i64x2, loaded_a16)
-            a0_i64 = vector.extract(a_i64x2, static_position=[0], dynamic_position=[])
-            a1_i64 = vector.extract(a_i64x2, static_position=[1], dynamic_position=[])
+            a_i64x2 = Vec(loaded_a16).bitcast(fx.Int64)
+            a0_i64 = a_i64x2[0]
+            a1_i64 = a_i64x2[1]
             return a0_i64, a1_i64
 
         # ── A global→reg load ─────────────────────────────────────────────
@@ -318,8 +319,7 @@ def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer):
             (tile_m, tile_k_dwords), (tile_k_dwords, 1)
         )
         chunk_i32_a = a_load_bytes // 4
-        c_chunk_a = fx.Index(chunk_i32_a)
-        tx_i32_base = tx * c_chunk_a
+        tx_i32_base = tx * chunk_i32_a
 
         def load_a(idx_i32, a_load_bytes_v):
             if const_expr(a_load_bytes_v == 16):
@@ -348,12 +348,12 @@ def load_a_tile(base_k_div4, a_load_bytes_v, tx_i32_base_v, chunk_i32_a_v):
                 idx_i32 = row_a_global * _k_div4_factor + (base_k_div4 + col_a_local_i32)
                 a_vec = load_a(idx_i32, a_load_bytes_v)
                 if const_expr(a_load_bytes_v == 16):
-                    parts.append(vector.bitcast(T.i32x4, a_vec))
+                    parts.append(Vec(a_vec).bitcast(fx.Int32))
                 else:
                     parts.append(a_vec)
             return parts
 
-        c4_bytes = fx.Index(4)  # bytes per dword (always 4, used for LDS byte addressing)
+        c4_bytes = 4  # bytes per dword (always 4, used for LDS byte addressing)
 
         def store_a_tile_to_lds(vec_a_parts, lds_buffer, a_load_bytes_v, tx_i32_base_v, chunk_i32_a_v):
             for i in range_constexpr(num_a_loads):
@@ -365,7 +365,7 @@ def store_a_tile_to_lds(vec_a_parts, lds_buffer, a_load_bytes_v, tx_i32_base_v,
                         layout_lds=layout_lds,
                         row_local=row_a_local, col_local_i32=col_a_local_i32,
                         tx_c4=c4_bytes, k_blocks16=k_blocks16,
-                        lds_base=fx.Index(0),
+                        lds_base=0,
                         vec_part_i32x4=vec_a_parts[i], elem_bytes=elem_bytes,
                     )
                 elif const_expr(a_load_bytes_v == 8):
@@ -375,7 +375,7 @@ def store_a_tile_to_lds(vec_a_parts, lds_buffer, a_load_bytes_v, tx_i32_base_v,
                         layout_lds=layout_lds,
                         row_local=row_a_local, col_local_i32=col_a_local_i32,
                         tx_c4=c4_bytes, k_blocks16=k_blocks16,
-                        lds_base=fx.Index(0),
+                        lds_base=0,
                         vec_part_i32x2=vec_a_parts[i],
                     )
 
@@ -402,20 +402,20 @@ def dma_a_tile_to_lds(base_k_div4, lds_buffer):
                 col_a_local_sw = swizzle_xor16(row_a_local, col_a_local_i32 * c4_bytes, k_blocks16)
                 row_a_global = bx_m + row_a_local
                 global_byte_idx = row_a_global * k_bytes_factor + (base_k_div4 * c4_bytes + col_a_local_sw)
-                global_offset = arith.index_cast(T.i32, global_byte_idx)
+                global_offset = fx.Int32(global_byte_idx)
 
                 if const_expr(i == 0):
                     lds_addr = memref_dialect.extract_aligned_pointer_as_index(lds_buffer) + wave_id * 64 * dma_bytes
-                    lds_ptr_i64_lane0 = rocdl.readfirstlane(T.i64, arith.index_cast(T.i64, lds_addr))
+                    lds_ptr_i64_lane0 = rocdl.readfirstlane(T.i64, fx.Int64(lds_addr))
                 else:
                     lds_ptr_i64_lane0 += total_threads * dma_bytes
                 lds_ptr_type = ir.Type.parse("!llvm.ptr<3>")
                 lds_ptr = llvm.inttoptr(lds_ptr_type, lds_ptr_i64_lane0)
 
-                size_i32 = arith.constant(dma_bytes, type=T.i32)
-                soffset = arith.constant(0, type=T.i32)
-                offset_imm = arith.constant(0, type=T.i32)
-                aux = arith.constant(1, type=T.i32)
+                size_i32 = fx.Int32(dma_bytes)
+                soffset = fx.Int32(0)
+                offset_imm = fx.Int32(0)
+                aux = fx.Int32(1)
 
                 rocdl.raw_ptr_buffer_load_lds(
                     a_rsrc,
@@ -447,11 +447,8 @@ def _mfma_fn_placeholder(*args, **kwargs):
         mfma_fn = _mfma_fn_placeholder
 
         if const_expr(_is_gfx950):
-            c0_i64 = arith.constant(0, type=T.i64)
-
             def pack_i64x4_to_i32x8(x0, x1, x2, x3):
-                v4 = vector.from_elements(T.vec(4, T.i64), [x0, x1, x2, x3])
-                return vector.bitcast(T.vec(8, T.i32), v4)
+                return Vec.from_elements([x0, x1, x2, x3], fx.Int64).bitcast(fx.Int32)
         else:
             mfma_fn = rocdl.mfma_f32_16x16x32_fp8_fp8
 
@@ -479,17 +476,17 @@ def load_scales_for_tile(k_base):
                 sa_base_offset = kb * c_M
                 s_a_vecs = []
                 for mi in range_constexpr(m_repeat):
-                    row_base_m = bx_m + arith.index(mi * 16)
+                    row_base_m = bx_m + mi * 16
                     row_g_base = row_base_m + row_off_base
                     sa_idx = sa_base_offset + row_g_base
                     s_a_vec = buffer_ops.buffer_load(
                         scale_a_rsrc, sa_idx, vec_width=4, dtype=T.f32
                     )
-                    s_a_vecs.append(vector.bitcast(T.f32x4, s_a_vec))
+                    s_a_vecs.append(Vec(s_a_vec).bitcast(fx.Float32))
 
                 s_b_vals = []
                 for ni in range_constexpr(num_acc_n):
-                    col_base_ni = by_n + n_tile_base + arith.index(ni * 16)
+                    col_base_ni = by_n + n_tile_base + ni * 16
                     n_block = col_base_ni // c_128
                     sb_idx = n_block * c_scale_k + kb
                     s_b_val = buffer_ops.buffer_load(
@@ -499,7 +496,7 @@ def load_scales_for_tile(k_base):
 
                 s_b_vecs = []
                 for ni in range_constexpr(num_acc_n):
-                    s_b_vecs.append(vector.broadcast(T.f32x4, s_b_vals[ni]))
+                    s_b_vecs.append(Vec.filled(4, fx.Float32(s_b_vals[ni]), fx.Float32))
 
                 combined_scales = []
                 for mi in range_constexpr(m_repeat):
@@ -611,28 +608,23 @@ def write_row_to_lds(
                         col_local = col_base_local + (ni * 16)
                         acc_idx = mi * num_acc_n + ni
                         acc = final_accs[acc_idx]
-                        val = vector.extract(
-                            acc, static_position=[ii], dynamic_position=[]
-                        )
-                        v_out = arith.trunc_f(_out_elem_type(), val)
+                        val = Vec(acc)[ii]
+                        v_out = val.to(_out_elem_dtype())
                         lds_idx = row_base_lds + col_local
-                        v1 = vector.from_elements(vec1_out, [v_out])
+                        v1 = Vec.from_elements([v_out], _out_elem_dtype())
                         vector.store(v1, lds_out, [lds_idx], alignment=2)
 
                 def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag):
                     idx_out = row * c_n + col_g0
                     byte_off = idx_out * 2
                     e_vec = 4 if (int(tile_n) % (32 * 4)) == 0 else 2
                     if const_expr(e_vec == 4):
-                        frag_i32x2 = vector.bitcast(T.vec(2, T.i32), frag)
+                        frag_i32x2 = Vec(frag).bitcast(fx.Int32)
                         buffer_ops.buffer_store(
                             frag_i32x2, c_rsrc, byte_off, offset_is_bytes=True
                         )
                     else:
-                        frag_i32x1 = vector.bitcast(T.vec(1, T.i32), frag)
-                        frag_i32 = vector.extract(
-                            frag_i32x1, static_position=[0], dynamic_position=[]
-                        )
+                        frag_i32 = Vec(frag).bitcast(fx.Int32)[0]
                         buffer_ops.buffer_store(
                             frag_i32, c_rsrc, byte_off, offset_is_bytes=True
                         )
@@ -660,10 +652,8 @@ def body_row(*, mi, ii, row_in_tile, row):
                 for ni in range_constexpr(num_acc_n):
                     acc_idx = mi * num_acc_n + ni
                     acc = final_accs[acc_idx]
-                    val = vector.extract(
-                        acc, static_position=[ii], dynamic_position=[]
-                    )
-                    val_out = arith.trunc_f(_out_elem_type(), val)
+                    val = Vec(acc)[ii]
+                    val_out = val.to(_out_elem_dtype())
                     idx_out = idx_base + (ni * 16)
                     buffer_ops.buffer_store(val_out, c_rsrc, idx_out)
 
@@ -784,7 +774,7 @@ def _load_a_to_lds(base_k, lds_buffer, a_load_bytes_v, tx_i32_base_v, chunk_i32_
                 gpu.barrier()
                 a0_prefetch_pong = prefetch_a0_pack(lds_a_pong)
 
-            last_k = arith.index(K - tile_k)
+            last_k = K - tile_k
             final_accs = compute_tile_blockscale(
                 global_accs, b_tile_pong, lds_a_pong, scales_pong,
                 a0_prefetch=a0_prefetch_pong,
@@ -826,8 +816,7 @@ def _load_a_to_lds(base_k, lds_buffer, a_load_bytes_v, tx_i32_base_v, chunk_i32_
                 gpu.barrier()
                 a0_prefetch_pong = prefetch_a0_pack(lds_a_pong)
 
-            last_k = arith.index(K - tile_k)
-            second_last_k = arith.index(K - tile_k * 2)
+            last_k = K - tile_k
 
             _load_a_to_lds(last_k, lds_a_ping, a_load_bytes, tx_i32_base, chunk_i32_a)
             b_tile_ping = prefetch_b_tile(last_k)

kernels/fused_rope_cache_kernel.py

@@ -40,7 +40,7 @@
 import flydsl.expr as fx
 from flydsl.expr import arith, buffer_ops, const_expr, range_constexpr, vector
 from flydsl.expr.arith import ArithValue
-from flydsl.expr.typing import T
+from flydsl.expr.typing import T, Vector as Vec
 from kernels.kernels_common import get_warp_size
 
 # WARP_SIZE is 32 on RDNA (wave32: gfx10xx/gfx11xx/gfx12xx) and 64 on CDNA (wave64: gfx9xx).
@@ -124,6 +124,7 @@ def fused_qk_rope_reshape_and_cache(
         tid = fx.thread_idx.x
 
         elem_type = T.bf16 if dtype_str == "bf16" else T.f16
+        elem_dtype = fx.BFloat16 if dtype_str == "bf16" else fx.Float16
 
         # --- Layout API setup ---
         copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy(copy_bits), elem_bits)
@@ -168,18 +169,18 @@ def ds_bpermute_pair(vec_val, pair_byte_addr):
                 # Truncate back to i16, bitcast to elem_type, reconstruct vector<1xelem_type>
                 peer_i16 = ArithValue(peer_i32).trunci(T.i16)
                 peer_elem = ArithValue(peer_i16).bitcast(elem_type)
-                return vector.from_elements(T.vec(1, elem_type), [peer_elem])
+                return Vec.from_elements([peer_elem], elem_dtype)
             else:
                 # VEC_WIDTH>=2: VEC_WIDTH bf16/f16 elements → n_i32 x i32, one ds_bpermute per chunk.
                 # VEC_WIDTH=2 → n_i32=1 (32 bits); VEC_WIDTH=4 → n_i32=2 (64 bits), etc.
                 n_i32 = VEC_WIDTH // 2
-                v_i32 = vector.bitcast(T.vec(n_i32, T.i32), vec_val)
+                v_i32 = Vec(vec_val).bitcast(fx.Int32)
                 peer_chunks = []
                 for ci in range_constexpr(n_i32):
-                    chunk = vector.extract(v_i32, static_position=[ci], dynamic_position=[])
+                    chunk = v_i32[ci]
                     peer_chunks.append(fx.rocdl.ds_bpermute(T.i32, pair_byte_addr, chunk))
-                peer_v_i32 = vector.from_elements(T.vec(n_i32, T.i32), peer_chunks)
-                return vector.bitcast(T.vec(VEC_WIDTH, elem_type), peer_v_i32)
+                peer_v_i32 = Vec.from_elements(peer_chunks, fx.Int32)
+                return peer_v_i32.bitcast(elem_dtype)
 
         if tid < vecs_per_head:
             # --- Load position (scalar i32) ---
@@ -288,11 +289,13 @@ def ds_bpermute_pair(vec_val, pair_byte_addr):
 
                         k_scaled = []
                         v_scaled = []
+                        k_rot_vec = Vec(k_rot_e.ir_value())
+                        v_vec = Vec(v_e)
                         for i in range_constexpr(VEC_WIDTH):
                             # Always use vector.extract; works for VEC_WIDTH=1 (vector<1xbf16>)
                             # and VEC_WIDTH>1 equally.
-                            ke = ArithValue(vector.extract(k_rot_e.ir_value(), static_position=[i], dynamic_position=[])).extf(T.f32) * k_rcp
-                            ve = ArithValue(vector.extract(v_e, static_position=[i], dynamic_position=[])).extf(T.f32) * v_rcp
+                            ke = k_rot_vec[i].to(fx.Float32) * k_rcp
+                            ve = v_vec[i].to(fx.Float32) * v_rcp
                             k_scaled.append(ke)
                             v_scaled.append(ve)
 
@@ -417,7 +420,7 @@ def pack_fp8(vals):
                                     + pid_b * x_size
                                     + sub_o
                                 )
-                                k_elem = vector.extract(k_rot_e.ir_value(), static_position=[vi], dynamic_position=[])
+                                k_elem = Vec(k_rot_e.ir_value())[vi]
                                 buffer_ops.buffer_store(k_elem, kc_rsrc, kc_nf_off)
 
                             for vi in range_constexpr(VEC_WIDTH):
@@ -428,7 +431,7 @@ def pack_fp8(vals):
                                     + d_idx * block_size
                                     + pid_b
                                 )
-                                v_elem = vector.extract(v_e, static_position=[vi], dynamic_position=[])
+                                v_elem = Vec(v_e)[vi]
                                 buffer_ops.buffer_store(v_elem, vc_rsrc, vc_nf_off)
 
     def _mark_token_layout_dynamic(tensor):

kernels/gemm_fp8fp4_gfx1250.py

@@ -1243,10 +1243,10 @@ def _l2_prefetch(k_base):
         desc_as_init = make_desc_as(stages_as_mem[0], split_k_base)
         desc_bs_init = make_desc_bs(stages_bs_mem[0], split_k_base)
 
-        adv_a_i32 = tile_k // PACK_FACTOR_A
-        adv_b_i32 = packed_tile_k_b * 16
-        adv_as_i32 = tile_k // SCALE_BLOCK * wmma_m_rep
-        adv_bs_i32 = tile_k // SCALE_BLOCK * b_scale_load_rep
+        adv_a_i32 = fx.Int32(tile_k // PACK_FACTOR_A)
+        adv_b_i32 = fx.Int32(packed_tile_k_b * 16)
+        adv_as_i32 = fx.Int32(tile_k // SCALE_BLOCK * wmma_m_rep)
+        adv_bs_i32 = fx.Int32(tile_k // SCALE_BLOCK * b_scale_load_rep)
 
         pred_const = fx.Int32(1)
 

kernels/silu_and_mul_fq.py

@@ -44,7 +44,7 @@
 from flydsl.expr import arith, buffer_ops, const_expr, range_constexpr, rocdl, vector
 from flydsl.expr import math as fx_math
 from flydsl.expr.arith import ArithValue
-from flydsl.expr.typing import T
+from flydsl.expr.typing import T, Vector as Vec
 from kernels.kernels_common import get_warp_size
 
 BLOCK_THREADS = 256
@@ -356,18 +356,14 @@ def _f32_to_e2m1(qx_f32):
                     else:
                         # quant_mode == "none": write bf16 out directly.
                         # out row stride = inter_dim * 2 bytes.
-                        act_f32_vec = vector.from_elements(vec_f32_ty, act_vals)
-                        act_f32_av = ArithValue(act_f32_vec)
-                        act_bf16_vec = act_f32_av.truncf(T.vec(VEC, T.bf16))
+                        act_bf16_vec = Vec.from_elements(act_vals, fx.Float32).to(fx.BFloat16)
                         # Write as packed i32 (VEC/2 dwords).
                         vec_dw = VEC // 2  # each dword = 2 bf16 elems
                         if const_expr(vec_dw >= 1):
-                            act_i32 = vector.bitcast(T.vec(vec_dw, T.i32), act_bf16_vec)
+                            act_i32 = act_bf16_vec.bitcast(fx.Int32)
                             bf16_byte_off = in_row * (inter_dim * 2) + col0 * 2
                             if const_expr(vec_dw == 1):
-                                store_val = vector.extract(
-                                    act_i32, static_position=[0], dynamic_position=[]
-                                )
+                                store_val = act_i32[0]
                                 buffer_ops.buffer_store(
                                     store_val,
                                     out_rsrc,