[Expr] Add type-aware extrema and integer ceiling division (#995)

Author: Phil-amd

.claude/skills/flydsl-kernel-authoring/SKILL.md

@@ -433,6 +433,9 @@ mask = fx.Int32(0xFF)                            # i32 constant (preferred)
 result = a + b
 result = a * scale
 result = cond.select(true_val, false_val)
+largest = fx.max(a, b)
+smallest = fx.min(a, b)
+tiles = fx.ceildiv(count, tile_size)  # signed/unsigned dispatch from typed operands
 
 # Keep direct arith.*FOp only when explicit fastmath flags are required.
 ```
@@ -484,7 +487,8 @@ Use `Vec.filled(...)` for splats and `Vec.from_elements(...)` for vectors from s
 | Add | `a + b` | Yes | Use direct FOp only for explicit fastmath |
 | Multiply | `a * b` | Yes | Use direct FOp only for explicit fastmath |
 | Negate | `-a` | Yes | |
-| Max | `a.maximumf(b)` | Yes | Good for ReLU |
+| Max / Min | `fx.max(a, b)` / `fx.min(a, b)` | Yes | Float forms propagate NaN; `fx.maxnumf` does not |
+| Integer ceil-div | `fx.ceildiv(a, b)` | Yes | Direct signed/unsigned op; distinct from layout `fx.ceil_div` |
 | Compare | `arith.cmpf(a, b, pred)` | Yes | Returns i1/vec<i1> |
 | Select | `cond.select(t, f)` | Yes | |
 | Abs | no direct helper | Use `-v`, comparison, and `cond.select(...)` |
@@ -783,7 +787,7 @@ vD = vAB + Vec(fx.memref_load_vec(rC))
 # --- ReLU: C = max(A, 0) ---
 vA = Vec(fx.memref_load_vec(rA))
 zero_vec = Vec.filled(vec_width, 0.0, fx.Float32)
-vC = vA.maximumf(zero_vec)
+vC = fx.max(vA, zero_vec)
 
 # --- Abs: C = |A| (arith.absf does NOT exist) ---
 vA = fx.memref_load_vec(rA)

.claude/skills/kernel-code-cleanup/SKILL.md

@@ -132,7 +132,10 @@ fx.copy(copy, fx.slice(tA, (None, tid)), rA)   # after partitioning tA (§7b: pr
 | `arith.index_cast(T.i32, v)` | `fx.Int32(v)` |
 | `arith.select(cond, t, f)` | `cond.select(t, f)` |
 | `arith.cmpi(slt, a, b)` | `a < b` |
-| `arith.maxnumf(a,b)` | `a.maximumf(b)` |
+| `arith.maximumf/minimumf(a,b)` | `fx.max(a, b)` / `fx.min(a, b)` |
+| `arith.maxsi/maxui/minsi/minui(a,b)` | `fx.max(a, b)` / `fx.min(a, b)` |
+| `arith.maxnumf(a,b)` | `fx.maxnumf(a, b)` — different NaN semantics from `fx.max` |
+| `arith.ceildivsi/ceildivui(a,b)` | `fx.ceildiv(a, b)` |
 
 Keep `arith.cmpf` / explicit `*FOp` only where no operator exists or fastmath is
 needed.
@@ -471,6 +474,7 @@ def _run_compiled(exe, *args):             # in-tree
 | `arith.unwrap(v)` / `_to_raw(v)` | `v.ir_value()` (boundary only) |
 | `fx.Index(n)` / `arith.index` / `arith.index_cast` | explicit `fx.Int64/Int32(...)` |
 | `arith.mulf/addf/trunc_f/select` | `*`, `+`, `.to(ty)`, `.select(...)` |
+| raw integer min/max or ceil-div | `fx.max` / `fx.min` / `fx.ceildiv` |
 | `vector.extract/bitcast/splat` | `fx.Vector(v)[i]` / `.bitcast(ty)` / `.filled(...)` |
 | `scf.ForOp` / `scf.IfOp` | `range_constexpr` / `range(..., init=)` / Python `if` / `const_expr` |
 | `buffer_ops.*` + offsets | `fx.rocdl.make_buffer_tensor` + layout + `fx.copy` |

.github/workflows/pre-checks.yaml

@@ -46,6 +46,12 @@ jobs:
           REVIEWDOG_REPORTER: ${{ github.event_name == 'pull_request' && 'github-pr-review' || 'github-check' }}
           REVIEWDOG_FILTER_MODE: ${{ github.event_name == 'pull_request' && 'diff_context' || 'nofilter' }}
 
+      - name: Reject new legacy typed arithmetic
+        run: python3 scripts/check_typed_arithmetic_usage.py
+        env:
+          BASE_SHA: ${{ github.event.pull_request.base.sha || github.event.before || '' }}
+          HEAD_SHA: ${{ github.sha }}
+
   cpp-style:
     name: Check C++ Code Style
     runs-on: ubuntu-24.04

docs/api_stability.md

@@ -162,7 +162,7 @@ excludes it from the catalog above.
 | `fx.index_cast` | `fx.Index(x)` | v0.4 |
 | `fx.constant_vector` | `Numeric` and `Vector` member functions | v0.4 |
 | `fx.tdm_ops` and content reached through this alias | `fx.rocdl.tdm_ops`; the latter is a target-specific unstable path | v0.4 |
-| `fx.Numeric.maximumf`, `fx.Numeric.minimumf` | `fx.arith.maximumf(x, y)`, `fx.arith.minimumf(x, y)` | v0.4 |
+| `fx.Numeric.maximumf`, `fx.Numeric.minimumf` | `fx.max(x, y)`, `fx.min(x, y)` | v0.4 |
 | `fx.Numeric.shrui`, `fx.Numeric.addf` | `fx.arith.shrui(x, amount)`, `x + y` with a `fastmath` context | v0.4 |
 | `fx.Numeric.exp2` | `fx.math.exp2(x)` | v0.4 |
 | `fx.Numeric.shuffle_xor` | `fx.gpu.shuffle_xor(x, offset, width)` | v0.4 |

docs/kernel_authoring_guide.md

@@ -201,6 +201,9 @@ result = a + b
 result = a * 2
 result = a // 4
 result = a % 16
+largest = fx.max(a, b)
+smallest = fx.min(a, b)
+tiles = fx.ceildiv(count, tile_size)
 
 # Cast (prefer DSL numeric constructors)
 i64_val = fx.Int64(int_val) # cast to 64-bit integer (fx.Index is deprecated)
@@ -216,6 +219,9 @@ result = a << 4
 ```
 
 Use direct `arith.*FOp(..., fastmath=...)` only where explicit fastmath flags are performance-critical.
+Use `fx.max` / `fx.min` for type-dispatched extrema and `fx.ceildiv` for
+overflow-safe integer ceil division. `fx.maxnumf` intentionally retains
+non-NaN-wins semantics, and `fx.ceil_div` remains the layout/int-tuple API.
 
 ### 4.2 Vector values (`Vector`)
 

docs/language/arithmetic_types.md

@@ -47,6 +47,33 @@ The following methods apply to both `Numeric` and `Vector`, and are elementwise
 | `as_numeric` / `Numeric.from_python_value(value)` | build a `Numeric` from a Python value | `as_numeric(5)` → `Int32(5)` |
 | `Numeric.from_ir_type(ir_type)` | the `Numeric` type for an MLIR type | `Numeric.from_ir_type(T.f32())` → `Float32` |
 
+### Function-level typed arithmetic
+
+`fx.max`, `fx.min`, and `fx.ceildiv` normalize Python literals and DSL operands,
+resolve one common type, broadcast scalar operands to `Vector` shapes, and emit
+the dedicated MLIR operation for that type:
+
+| API | Float | Signed integer | Unsigned integer |
+|---|---|---|---|
+| `fx.max` | `arith.maximumf` | `arith.maxsi` | `arith.maxui` |
+| `fx.min` | `arith.minimumf` | `arith.minsi` | `arith.minui` |
+| `fx.ceildiv` | unsupported | `arith.ceildivsi` | `arith.ceildivui` |
+
+`fx.max` and `fx.min` are variadic and accept nested lists/tuples. Their float
+forms propagate NaN and order signed zero as `-0.0 < +0.0`. This is deliberately
+different from `fx.maxnumf`, which returns the non-NaN input when exactly one
+operand is NaN.
+
+`fx.ceildiv` rounds integer division toward positive infinity. It does not use
+`(a + b - 1) // b`, whose intermediate addition can overflow at run time, and
+it does not change the floor-division meaning of `//`. The similarly named
+`fx.ceil_div` remains the layout/int-tuple operation.
+
+Boolean inputs to `fx.max` / `fx.min` widen to `Int32`. Boolean, `Index`, float,
+and narrow storage-float inputs to `fx.ceildiv` are rejected. `Index` and narrow
+storage floats are also rejected by `fx.max` / `fx.min`; cast them to an
+explicit supported arithmetic type first.
+
 ### Vector
 
 `Vector` is a fixed-length sequence of `N` elements of a single `Numeric`

kernels/attention/pa_metadata.py

@@ -971,7 +971,7 @@ def _num_part(batch_idx):
             # number of partition_size-token partitions for this batch =
             # ceil(context_len[batch_idx] / kv_granularity)
             ctxv = _load(ctx_rsrc, batch_idx)
-            return fx.Int32(arith.ceildivui(ctxv.ir_value(), c_kvg.ir_value()))
+            return fx.Int32(fx.ceildiv(fx.Uint32(ctxv), fx.Uint32(c_kvg)))
 
         def _store(rsrc, off, val):
             # NOTE: no masked stores — masked buffer_store sets OOB offset

kernels/common/utils.py

@@ -41,8 +41,11 @@ def exp2_f32_fast(value):
     return exp2_amdgcn_scalar(raw)
 
 
-def cdiv(numer: int, denom: int) -> int:
-    return (numer + denom - 1) // denom
+def cdiv(numer, denom):
+    """Ceiling division for host integers and typed DSL integer values."""
+    if isinstance(numer, (fx.Numeric, fx.Vector)) or isinstance(denom, (fx.Numeric, fx.Vector)):
+        return fx.ceildiv(numer, denom)
+    return -(-numer // denom)
 
 
 # Alias: several kernels historically spelled this ``ceildiv``.

kernels/moe/topk_gating_softmax_kernel.py

@@ -205,7 +205,7 @@ def group_reduce(x, mode):
             off = fx.Int32(THREADS_PER_TOKEN // (2 << _sh))
             peer = w.shuffle_xor(off, width_i32)
             if mode == "max":
-                w = w.maximumf(peer)
+                w = fx.max(w, peer)
             else:
                 w = w.addf(peer, fastmath=fm_fast)
         return w
@@ -314,7 +314,7 @@ def _store_scalar_i32(divided, index, val):
             val_e = vector.extract(as_ir_value(atom_vec), dynamic_position=[], static_position=[v])
             xv = val_e if dtype_str == "f32" else val_e.extf(compute_type)
             x_list.append(xv)
-            thread_max = thread_max.maximumf(xv)
+            thread_max = fx.max(thread_max, xv)
 
     group_max = group_reduce(thread_max, "max")
 
@@ -364,7 +364,7 @@ def _store_scalar_i32(divided, index, val):
 
     # Pass 5: leader writes weights/indices/tei (with optional renorm).
     c_eps = fx.Float32(1e-20)
-    denom = selected_sum.maximumf(c_eps)
+    denom = fx.max(selected_sum, c_eps)
     inv_denom = c_one_f / denom
 
     if (expert_lane == fx.Int32(0)) & (global_token < i32_num_tokens):
@@ -467,7 +467,7 @@ def group_reduce(x, mode):
                 off = fx.Int32(THREADS_PER_TOKEN // (2 << _sh))
                 peer = w.shuffle_xor(off, width_i32)
                 if mode == "max":
-                    w = w.maximumf(peer)
+                    w = fx.max(w, peer)
                 else:
                     w = w.addf(peer, fastmath=fm_fast)
             return w
@@ -568,7 +568,7 @@ def _store_scalar_i32(divided, index, val):
                 val_e = vector.extract(as_ir_value(atom_vec), dynamic_position=[], static_position=[v])
                 xv = val_e if dtype_str == "f32" else val_e.extf(compute_type)
                 x_list.append(xv)
-                thread_max = thread_max.maximumf(xv)
+                thread_max = fx.max(thread_max, xv)
 
         group_max = group_reduce(thread_max, "max")
 
@@ -632,7 +632,7 @@ def _store_scalar_i32(divided, index, val):
         # Pass 5: Leader writes weights/indices/tei (with optional renorm)
         # ==================================================================
         c_eps = fx.Float32(1e-20)
-        denom = selected_sum.maximumf(c_eps)
+        denom = fx.max(selected_sum, c_eps)
         inv_denom = c_one_f / denom
 
         # Inline the leader-active predicate so the AST rewriter recognises it

kernels/norm/layernorm_kernel.py

@@ -837,7 +837,7 @@ def wave_reduce_max(x):
             for _sh_exp in range_constexpr(int(math.log2(WARP_SIZE))):
                 off = WARP_SIZE // (2 << _sh_exp)
                 peer = w.shuffle_xor(off, WARP_SIZE)
-                w = w.maximumf(peer)
+                w = fx.max(w, peer)
             return w
 
         def block_reduce_add2(val0, val1):
@@ -960,7 +960,7 @@ def block_reduce_max(val):
                 y_local.append(y)
                 y_abs = (y.bitcast(fx.Uint32) & abs_mask).bitcast(fx.Float32)
                 tile_max = y_abs.reduce(ReductionOp.MAX)
-                thread_row_max = thread_row_max.maximumf(tile_max)
+                thread_row_max = fx.max(thread_row_max, tile_max)
 
             row_max = block_reduce_max(thread_row_max)
             scale = row_max / c_dtype_max
@@ -1050,7 +1050,7 @@ def _abs_scalar(val):
                     s = s_e if dtype_str == "f32" else s_e.to(fx.Float32)
                     y = y * s
                 y_abs = _abs_scalar(y)
-                thread_row_max = thread_row_max.maximumf(is_valid.select(y_abs, c_zero_f))
+                thread_row_max = fx.max(thread_row_max, is_valid.select(y_abs, c_zero_f))
 
             row_max = block_reduce_max(thread_row_max)
             scale = row_max / c_dtype_max
@@ -1186,7 +1186,7 @@ def wave_reduce_max(x):
             for _sh_exp in range_constexpr(int(math.log2(WARP_SIZE))):
                 off = WARP_SIZE // (2 << _sh_exp)
                 peer = w.shuffle_xor(off, WARP_SIZE)
-                w = w.maximumf(peer)
+                w = fx.max(w, peer)
             return w
 
         def block_reduce_add2(val0, val1):
@@ -1318,7 +1318,7 @@ def block_reduce_max(val):
                 y_local.append(y)
                 y_abs = (y.bitcast(fx.Uint32) & abs_mask).bitcast(fx.Float32)
                 tile_max = y_abs.reduce(ReductionOp.MAX)
-                thread_row_max = thread_row_max.maximumf(tile_max)
+                thread_row_max = fx.max(thread_row_max, tile_max)
 
             row_max = block_reduce_max(thread_row_max)
             scale = row_max / c_dtype_max
@@ -1420,7 +1420,7 @@ def _abs_scalar(val):
                     s = s_e if dtype_str == "f32" else s_e.to(fx.Float32)
                     y = y * s
                 y_abs = _abs_scalar(y)
-                thread_row_max = thread_row_max.maximumf(is_valid.select(y_abs, c_zero_f))
+                thread_row_max = fx.max(thread_row_max, is_valid.select(y_abs, c_zero_f))
 
             row_max = block_reduce_max(thread_row_max)
             scale = row_max / c_dtype_max