[FlyDSL] Simplify kernel APIs and remove redundant wrappers

aiter/ops/flydsl/kernels/act.py

@@ -20,9 +20,7 @@
 
 def sigmoid_batch(xs, *, alpha=1.0):
     """Emit all exponentials before their reciprocals to preserve batch scheduling."""
-    e = [
-        fx.Float32(rocdl.exp2(T.f32, _raw(x * fx.Float32(-alpha * LOG2E)))) for x in xs
-    ]
+    e = [fx.exp2(x * fx.Float32(-alpha * LOG2E), fastmath="afn") for x in xs]
     return [fx.Float32(rocdl.rcp(T.f32, _raw(fx.Float32(1.0) + ei))) for ei in e]
 
 
@@ -56,7 +54,7 @@ def tanh_batch(xs):
     es = []
     for x in xs:
         abs_x = fx.max(x, -x)
-        es.append(fx.Float32(rocdl.exp2(T.f32, _raw(abs_x * neg_two_log2e))))
+        es.append(fx.exp2(abs_x * neg_two_log2e, fastmath="afn"))
     recips = [fx.Float32(rocdl.rcp(T.f32, _raw(fx.Float32(1.0) + e))) for e in es]
     zero = fx.Float32(0.0)
     out = []

aiter/ops/flydsl/kernels/bmm_a8w8_mxscale_gfx950.py

@@ -36,8 +36,7 @@
 
 import flydsl.compiler as flyc
 import flydsl.expr as fx
-from flydsl._mlir import ir
-from flydsl._mlir.dialects import rocdl as rocdl_ir
+from flydsl.compiler.mlir_utils import convert_to_mlir_attr
 from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
 from flydsl.expr.typing import BFloat16, Constexpr, Float32, Int8, Int32, T
 from flydsl.expr.typing import Vector as Vec
@@ -513,14 +512,14 @@ def _raw_dma(rsrc, size, byte, lds_off, aux=0):
                 lds_ptr = fx.to_llvm_ptr(
                     fx.add_offset(lds_i8, rocdl.readfirstlane(T.i32, lds_off))
                 )
-                rocdl_ir.raw_ptr_buffer_load_async_lds(
+                rocdl.raw_ptr_buffer_load_async_lds(
                     rsrc,
                     lds_ptr,
                     fx.Int32(size).ir_value(),
                     fx.Int32(byte).ir_value(),
                     fx.Int32(0).ir_value(),
                     fx.Int32(0).ir_value(),
-                    aux=ir.IntegerAttr.get(ir.IntegerType.get_signless(32), aux),
+                    aux=convert_to_mlir_attr(aux),
                 )
 
             def _dma_b(byte, lds_off):

aiter/ops/flydsl/kernels/dpp_utils.py

@@ -20,21 +20,14 @@ def update_dpp_i32(
     280/264 for row xor-8, 276/260 for row xor-4, 78 for xor-2,
     and 177 for xor-1 within a 16-lane row.
     """
-    from flydsl._mlir.dialects import llvm as _llvm
-
-    return _llvm.call_intrinsic(
+    return fx.rocdl.update_dpp(
         fx.Int32.ir_type,
-        "llvm.amdgcn.update.dpp.i32",
-        [
-            fx.Int32(old).ir_value(),
-            fx.Int32(src).ir_value(),
-            fx.Int32(dpp_ctrl).ir_value(),
-            fx.Int32(row_mask).ir_value(),
-            fx.Int32(bank_mask).ir_value(),
-            fx.Boolean(bound_ctrl).ir_value(),
-        ],
-        [],
-        [],
+        fx.Int32(old).ir_value(),
+        fx.Int32(src).ir_value(),
+        dpp_ctrl,
+        row_mask,
+        bank_mask,
+        bound_ctrl,
         **kw,
     )
 

aiter/ops/flydsl/kernels/gather_gemm_8wave.py

@@ -6,16 +6,15 @@
 import flydsl.compiler as flyc
 import flydsl.expr as fx
 
-# Raw dialect, and the only one here. There IS an fx atom for global -> LDS,
+# The ROCDL operation below has no matching synchronous fx copy atom. There is
 # `cdna4.GlobalLoadAsyncLDS128b`, but only the async one: it completes through
 # asyncmark / wait_asyncmark, while this pipeline's `wait_barrier` is a
 # hand-counted `s_waitcnt vmcnt(N)` covering A and B together. Tracking A one
 # way and B the other is how the MFMAs come to read stale LDS -- see BlockScale
 # below. `rocdl.global.load.lds` is the synchronous form and counts in vmcnt,
 # so the existing barrier keeps covering it. Pointers still go through
 # `fx.to_llvm_ptr`, so no address space is hardcoded.
-from flydsl._mlir.dialects import rocdl as _rocdl_d
-from flydsl.expr import T, arith, const_expr, range_constexpr, rocdl
+from flydsl.expr import T, const_expr, range_constexpr, rocdl
 from flydsl.expr.rocdl import cvt_pk_f32_fp8
 from flydsl.expr.typing import Vector as Vec
 
@@ -125,7 +124,7 @@ def load(self, lds_dst, k_offset):
             )
             src = fx.add_offset(self.kv_ptr, self.gl_offsets[step] + fx.Int64(k_offset))
             dst = fx.add_offset(lds_dst.ptr, fx.Int32(step_off) + self.lane_id * 16)
-            _rocdl_d.global_load_lds(
+            rocdl.global_load_lds(
                 _raw(fx.to_llvm_ptr(src)), _raw(fx.to_llvm_ptr(dst)), 16, 0
             )
 
@@ -379,9 +378,7 @@ def _emit(
                     if const_expr(n_valid < self.lds_half):
                         # Dead columns of the widened part: an index past the
                         # descriptor's last record discards the store.
-                        index = arith.select(
-                            col < n_valid, index, self.m_rows * row_stride
-                        )
+                        index = (col < n_valid).select(index, self.m_rows * row_stride)
                     self._store(vals[i], div, index)
 
     def store_k(self, c_frag, base_row, head, col_base):

aiter/ops/flydsl/kernels/gemm_a16w16_gfx950_utils.py

@@ -3,15 +3,13 @@
 
 import flydsl.compiler as flyc
 import flydsl.expr as fx
-from flydsl._mlir import ir
 from flydsl._mlir.dialects import llvm
 from flydsl.expr import (
     const_expr,
     gpu,
     range_constexpr,
     rocdl,
 )
-from flydsl.expr.typing import T
 
 from .tensor_shim import buf_base_i64
 
@@ -25,10 +23,18 @@ def wait_vmcnt_and_barrier(vmcnt=0):
     rocdl.s_barrier()
 
 
-def get_llvm_ptr(ptr, offset, dtype_bytes, ptr_type):
-    byte_offset = fx.Int64(fx.Index(offset) * fx.Index(dtype_bytes))
+def get_llvm_ptr(ptr, offset, dtype_bytes):
+    byte_offset = fx.Int64(offset) * dtype_bytes
     address = buf_base_i64(ptr) + byte_offset
-    return llvm.IntToPtrOp(ptr_type, address.ir_value()).result
+    ptr_type = fx.PointerType.get(fx.Int8.ir_type, fx.AddressSpace.Global, 1)
+    return fx.to_llvm_ptr(fx.inttoptr(ptr_type, address))
+
+
+def _global_i32_ptr(ptr, offset):
+    """Four-byte-aligned global pointer for split-K synchronization."""
+    address = buf_base_i64(ptr) + fx.Int64(offset) * 4
+    ptr_type = fx.PointerType.get(fx.Int32.ir_type, fx.AddressSpace.Global, 4)
+    return fx.inttoptr(ptr_type, address)
 
 
 def store_global_f32_vec(c_ptr, global_offset, vec, vec_size):
@@ -41,7 +47,6 @@ def store_global_f32_vec(c_ptr, global_offset, vec, vec_size):
             c_ptr,
             global_offset + vec_idx * 4,
             4,
-            ir.Type.parse("!llvm.ptr<1>"),
         )
         llvm.InlineAsmOp(
             None,
@@ -184,7 +189,6 @@ def zero_c(self):
                             self.c_ptr,
                             c_offset,
                             self.C_DTYPE_BYTES,
-                            ir.Type.parse("!llvm.ptr<1>"),
                         )
                         llvm.InlineAsmOp(
                             None,
@@ -195,17 +199,10 @@ def zero_c(self):
                         )
             wait_vmcnt_and_barrier(0)
             if self.tid == 0:
-                signal_ptr = get_llvm_ptr(
-                    self.signal_ptr,
-                    self.signal_idx,
-                    4,
-                    ir.Type.parse("!llvm.ptr<1>"),
-                )
-                llvm.StoreOp(
-                    fx.Int32(1).ir_value(),
-                    signal_ptr,
-                    alignment=4,
-                    ordering=llvm.AtomicOrdering.monotonic,
+                fx.generic_store(
+                    _global_i32_ptr(self.signal_ptr, self.signal_idx),
+                    fx.Int32(1),
+                    memory_order=fx.AtomicOrdering.Monotonic,
                     syncscope="agent",
                 )
 
@@ -216,72 +213,41 @@ def wait_until_initialized(self):
             # (cur != 0). cur starts at 0 to force at least one load.
             cur = fx.Int32(0)
             while cur == 0:
-                signal_ptr = get_llvm_ptr(
-                    self.signal_ptr,
-                    self.signal_idx,
-                    4,
-                    ir.Type.parse("!llvm.ptr<1>"),
-                )
                 cur = fx.Int32(
-                    llvm.LoadOp(
-                        T.i32,
-                        signal_ptr,
-                        alignment=4,
-                        ordering=llvm.AtomicOrdering.monotonic,
+                    fx.generic_load(
+                        _global_i32_ptr(self.signal_ptr, self.signal_idx),
+                        memory_order=fx.AtomicOrdering.Monotonic,
                         syncscope="agent",
-                    ).result
+                    )
                 )
         rocdl.sched_barrier(0)
         gpu.barrier()
 
     @flyc.jit
     def reset_sync_state(self):
-        semaphore_ptr = get_llvm_ptr(
-            self.semaphore_ptr,
-            self.signal_idx,
-            4,
-            ir.Type.parse("!llvm.ptr<1>"),
-        )
-        signal_ptr = get_llvm_ptr(
-            self.signal_ptr,
-            self.signal_idx,
-            4,
-            ir.Type.parse("!llvm.ptr<1>"),
-        )
-        zero = fx.Int32(0).ir_value()
-        llvm.StoreOp(
+        zero = fx.Int32(0)
+        fx.generic_store(
+            _global_i32_ptr(self.semaphore_ptr, self.signal_idx),
             zero,
-            semaphore_ptr,
-            alignment=4,
-            ordering=llvm.AtomicOrdering.monotonic,
+            memory_order=fx.AtomicOrdering.Monotonic,
             syncscope="agent",
         )
-        llvm.StoreOp(
+        fx.generic_store(
+            _global_i32_ptr(self.signal_ptr, self.signal_idx),
             zero,
-            signal_ptr,
-            alignment=4,
-            ordering=llvm.AtomicOrdering.monotonic,
+            memory_order=fx.AtomicOrdering.Monotonic,
             syncscope="agent",
         )
 
     @flyc.jit
     def finish_split(self, split_k):
         gpu.barrier()
         if self.tid == 0:
-            semaphore_ptr = get_llvm_ptr(
-                self.semaphore_ptr,
-                self.signal_idx,
-                4,
-                ir.Type.parse("!llvm.ptr<1>"),
-            )
-            arrive_idx = llvm.AtomicRMWOp(
-                llvm.AtomicBinOp.add,
-                semaphore_ptr,
-                fx.Int32(1).ir_value(),
-                llvm.AtomicOrdering.monotonic,
+            arrive_idx = fx.atomic_add(
+                _global_i32_ptr(self.semaphore_ptr, self.signal_idx),
+                fx.Int32(1),
                 syncscope="agent",
-                alignment=4,
-            ).result
+            )
             if arrive_idx == split_k - 1:
                 self.reset_sync_state()
         gpu.barrier()

aiter/ops/flydsl/kernels/gemm_common_gfx1250.py

@@ -3,7 +3,6 @@
 from collections import namedtuple
 
 import flydsl.expr as fx
-from flydsl._mlir import ir
 from flydsl._mlir.dialects import llvm as llvm_dialect
 from flydsl.expr import as_ir_value, gpu, rocdl
 from flydsl.expr.rocdl import cluster, tdm_ops
@@ -56,7 +55,7 @@ def make_sgpr_opaque(val_i32):
     Emits no instruction.
     """
     op = llvm_dialect.InlineAsmOp(
-        res=ir.IntegerType.get_signless(32),
+        res=T.i32,
         operands_=[as_ir_value(val_i32)],
         asm_string="",
         constraints="=s,s",
@@ -66,31 +65,30 @@ def make_sgpr_opaque(val_i32):
     return op.res
 
 
-def _raw_lds_ptr(lds_base_idx, byte_offset):
-    """Materialize an LLVM LDS pointer from a pre-extracted byte base."""
-    lds_ptr_ty = ir.Type.parse("!llvm.ptr<3>")
-    total_byte = fx.Index(lds_base_idx) + fx.Index(byte_offset)
-    return llvm_dialect.inttoptr(lds_ptr_ty, fx.Int32(total_byte).ir_value())
+def _lds_ptr(lds_base_idx, byte_offset, element_type, alignment):
+    """Materialize an aligned LDS pointer from a pre-extracted byte base."""
+    lds_ptr_ty = fx.PointerType.get(element_type, fx.AddressSpace.Shared, alignment)
+    total_byte = fx.Int32(lds_base_idx) + fx.Int32(byte_offset)
+    return fx.inttoptr(lds_ptr_ty, total_byte)
 
 
 def lds_load_b128_raw(lds_base_idx, byte_offset):
-    """Load 16 bytes from LDS using a pre-extracted base index (raw LLVM)."""
-    ptr_val = _raw_lds_ptr(lds_base_idx, byte_offset)
-    return llvm_dialect.load(
-        ir.VectorType.get([4], ir.IntegerType.get_signless(32)), ptr_val
-    )
+    """Load 16 bytes from LDS using a pre-extracted byte base."""
+    vec_type = fx.Vector.make_type(4, fx.Int32)
+    ptr_val = _lds_ptr(lds_base_idx, byte_offset, fx.Int32.ir_type, 16)
+    return as_ir_value(fx.ptr_load(ptr_val, result_type=vec_type))
 
 
 def lds_load_b32_raw(lds_base_idx, byte_offset):
     """Load 4 bytes from LDS as ``i32`` using a pre-extracted base index."""
-    ptr_val = _raw_lds_ptr(lds_base_idx, byte_offset)
-    return llvm_dialect.load(ir.IntegerType.get_signless(32), ptr_val)
+    ptr_val = _lds_ptr(lds_base_idx, byte_offset, fx.Int32.ir_type, 4)
+    return as_ir_value(fx.ptr_load(ptr_val))
 
 
 def lds_store_b128_raw(lds_base_idx, byte_offset, data):
-    """Store 16 bytes to LDS using a pre-extracted base index (raw LLVM)."""
-    ptr_val = _raw_lds_ptr(lds_base_idx, byte_offset)
-    llvm_dialect.store(as_ir_value(data), ptr_val)
+    """Store 16 bytes to LDS using a pre-extracted byte base."""
+    ptr_val = _lds_ptr(lds_base_idx, byte_offset, fx.Int32.ir_type, 16)
+    fx.ptr_store(as_ir_value(data), ptr_val)
 
 
 def workgroup_barrier(use_cluster=False):
@@ -135,9 +133,7 @@ def pipeline_fence_wait(use_cluster=False):
 
 def fmin_f32(a, b):
     """Scalar f32 min (maps to v_min_num_f32)."""
-    from flydsl.expr import arith
-
-    return fx.Float32(arith.minnumf(as_ir_value(a), as_ir_value(b)))
+    return fx.minnumf(fx.Float32(a), fx.Float32(b))
 
 
 def fclamp_f32(x, lo, hi):
@@ -149,15 +145,7 @@ def fclamp_f32(x, lo, hi):
 
 def _fma_f32(a, b, c):
     """Evaluate ``a * b + c`` as one fused f32 operation."""
-    return fx.Float32(
-        llvm_dialect.call_intrinsic(
-            T.f32,
-            "llvm.fma.f32",
-            [as_ir_value(a), as_ir_value(b), as_ir_value(c)],
-            [],
-            [],
-        )
-    )
+    return fx.Float32(fx.fma(fx.Float32(a), fx.Float32(b), fx.Float32(c)))
 
 
 def fused_silu_swiglu_elem(g, u, *, swiglu, limit_f32, neg_limit_f32):