Add RMSNorm kernel to latency-demo roofline notebook (#165)

notebooks/demo_gen_mlir.py

@@ -355,6 +355,118 @@ def gen_rope_mlir(num_heads, seq_len, head_dim, grid_s, grid_h, tile_seq=256):
   }}
 }}"""
 
+def gen_rmsnorm_mlir(n_rows, hidden_dim, num_cores=4, block_size=1024):
+    """Generate RMSNorm kernel: y = x * rsqrt(mean(x^2) + eps) * w.
+
+    Two-pass embarrassingly parallel over rows. Grid = [num_cores, 1].
+    Weight W is 1D [hidden_dim] — broadcast to match tile shape.
+
+    Design decisions:
+    - Unfused standalone kernel (no predecessor/successor fusion) — models the
+      HBM-resident pass-through as it appears between matmul stages in prefill.
+    - 1D grid [N,1] row-partition: prefill has abundant row parallelism
+      (seq_len >> num_cores), so hidden-dim sharding (2x2) adds allreduce
+      communication for zero benefit.  Matches adjacent matmul/softmax/SDPA grids.
+    - W is 1D [hidden_dim] per PyTorch convention (nn.Parameter(torch.ones(H))).
+      Loaded per-block in pass 2; small enough (8 KB) to be LX-resident on real
+      hardware, though the charge-model estimator costs every load from HBM.
+    """
+    assert hidden_dim % block_size == 0, "hidden_dim must be divisible by block_size"
+    hd = hidden_dim
+    bs = block_size
+    x_view = _mem_view("x_view", "x_ptr", [n_rows, hd], [hd, 1])
+    y_view = _mem_view("y_view", "y_ptr", [n_rows, hd], [hd, 1])
+    w_view = _mem_view("w_view", "w_ptr", [hd], [1])
+    x_acc = _access_tile("x_acc", "x_view", ["%row", "%col"], [1, bs], [n_rows, hd])
+    x_acc2 = _access_tile("x_acc2", "x_view", ["%row", "%col"], [1, bs], [n_rows, hd])
+    w_acc = _access_tile("w_acc", "w_view", ["%col"], [bs], [hd])
+    y_acc = _access_tile("y_acc", "y_view", ["%row", "%col"], [1, bs], [n_rows, hd])
+    return f"""module {{
+  func.func @rmsnorm_kernel(
+      %x_ptr: index, %y_ptr: index, %w_ptr: index,
+      %n_rows: index, %N: index, %eps: f16, %BLOCK_SIZE: index
+  ) attributes {{grid = [{num_cores}, 1]}} {{
+    %core_id = ktdp.get_compute_tile_id : index
+    %step = arith.constant {num_cores} : index
+    %c0 = arith.constant 0 : index
+    %c_hd = arith.constant {hd} : index
+
+{x_view}
+
+{y_view}
+
+{w_view}
+
+    scf.for %row = %core_id to %n_rows step %step : index {{
+
+        // === Pass 1: sum of squares over hidden dim (f32 accumulator) ===
+        %zero_f32 = arith.constant 0.0 : f32
+        %zero_block = tensor.splat %zero_f32 : tensor<1x{bs}xf32>
+
+        %sq_acc = scf.for %col = %c0 to %c_hd step %BLOCK_SIZE
+            iter_args(%acc = %zero_block) -> tensor<1x{bs}xf32> {{
+
+      {x_acc}
+
+            %x_blk = ktdp.load %x_acc : !ktdp.access_tile<1x{bs}xindex> -> tensor<1x{bs}xf16>
+
+            %x_sq = arith.mulf %x_blk, %x_blk : tensor<1x{bs}xf16>
+            %x_sq_f32 = arith.extf %x_sq : tensor<1x{bs}xf16> to tensor<1x{bs}xf32>
+            %acc_next = arith.addf %acc, %x_sq_f32 : tensor<1x{bs}xf32>
+
+            scf.yield %acc_next : tensor<1x{bs}xf32>
+        }}
+
+        // Reduce f32 accumulator
+        %reduce_init = tensor.splat %zero_f32 : tensor<1xf32>
+        %sum_sq = linalg.reduce {{ arith.addf }}
+            ins(%sq_acc : tensor<1x{bs}xf32>)
+            outs(%reduce_init : tensor<1xf32>)
+            dimensions = [1]
+
+        // === Compute rstd = rsqrt(sum_sq / N + eps) on tensor<1xf32> ===
+        %N_i32 = arith.index_cast %N : index to i32
+        %N_f32 = arith.sitofp %N_i32 : i32 to f32
+        %N_t = tensor.splat %N_f32 : tensor<1xf32>
+        %eps_f32 = arith.extf %eps : f16 to f32
+        %eps_t = tensor.splat %eps_f32 : tensor<1xf32>
+
+        %mean_sq = arith.divf %sum_sq, %N_t : tensor<1xf32>
+        %mean_sq_plus_eps = arith.addf %mean_sq, %eps_t : tensor<1xf32>
+        %rstd_f32 = math.rsqrt %mean_sq_plus_eps : tensor<1xf32>
+        %rstd_f16 = arith.truncf %rstd_f32 : tensor<1xf32> to tensor<1xf16>
+        %c0_idx = arith.constant 0 : index
+        %rstd_scalar = tensor.extract %rstd_f16[%c0_idx] : tensor<1xf16>
+        %rstd_block = tensor.splat %rstd_scalar : tensor<1x{bs}xf16>
+
+        // === Pass 2: normalize and scale ===
+        scf.for %col = %c0 to %c_hd step %BLOCK_SIZE {{
+
+      {x_acc2}
+
+      {w_acc}
+
+            %x2 = ktdp.load %x_acc2 : !ktdp.access_tile<1x{bs}xindex> -> tensor<1x{bs}xf16>
+            %w_1d = ktdp.load %w_acc : !ktdp.access_tile<{bs}xindex> -> tensor<{bs}xf16>
+
+            %w_init = arith.constant dense<0.0> : tensor<1x{bs}xf16>
+            %w = linalg.broadcast ins(%w_1d : tensor<{bs}xf16>) outs(%w_init : tensor<1x{bs}xf16>) dimensions = [0]
+
+            %x_norm = arith.mulf %x2, %rstd_block : tensor<1x{bs}xf16>
+            %y_blk = arith.mulf %x_norm, %w : tensor<1x{bs}xf16>
+
+      {y_acc}
+
+            ktdp.store %y_blk, %y_acc : tensor<1x{bs}xf16>, !ktdp.access_tile<1x{bs}xindex>
+
+            scf.yield
+        }}
+        scf.yield
+    }}
+    return
+  }}
+}}"""
+
 
 def gen_sdpa_mlir(seq_len, head_dim, block_m):
     """Generate a naive fused SDPA kernel (single-pass, no K-tiling).

notebooks/demo_helpers.py

@@ -18,6 +18,7 @@
     gen_matmul_mlir,
     gen_softmax_mlir,
     gen_rope_mlir,
+    gen_rmsnorm_mlir,
     gen_sdpa_mlir,
     gen_sdpa_decode_pv_mlir,
     gen_paged_attention_mlir,
@@ -81,7 +82,7 @@ def print_hw_config(hw: HardwareConfig):
     print(f"  {'cores_active':<14} {'systolic':>10} {'SIMD':>10}")
     for nc in [4, 32]:
         bw = chip_bw_per_cycle / nc
-        print(f"  {nc:<14} {hw.systolic_flops_per_cycle / bw:>7.0f} F/B {hw.simd_elements_per_cycle / bw:>7.0f} F/B")
+        print(f"  {nc:<14} {hw.systolic_flops_per_cycle / bw:>7.2f} F/B {hw.simd_elements_per_cycle / bw:>7.2f} F/B")
 
 
 def print_core_roofline(report, hw: HardwareConfig):
@@ -248,7 +249,6 @@ def run_kernel_softmax(hw, n_rows, row_width, num_cores, rng=None):
              input_ptr=rng.standard_normal((n_rows, row_width)).astype(np.float16)),
         dict(n_rows=n_rows))
 
-
 def run_kernel_rope(hw, num_heads, seq_len, head_dim, grid_s, grid_h, tile_seq=256, rng=None):
     """Generate RoPE MLIR, create tensors, run, return LatencyReport."""
     if rng is None:
@@ -266,6 +266,16 @@ def run_kernel_rope(hw, num_heads, seq_len, head_dim, grid_s, grid_h, tile_seq=2
              sin_ptr=sin_table,
              out_ptr=np.zeros((num_heads * seq_len, head_dim), dtype=np.float16)))
 
+def run_kernel_rmsnorm(hw, n_rows, hidden_dim, num_cores, rng=None, block_size=1024):
+    """Generate RMSNorm MLIR, create tensors, run, return LatencyReport."""
+    if rng is None:
+        rng = np.random.default_rng(0)
+    mlir = gen_rmsnorm_mlir(n_rows, hidden_dim, num_cores, block_size)
+    return run_kernel(hw, mlir, "rmsnorm_kernel",
+        dict(x_ptr=rng.standard_normal((n_rows, hidden_dim)).astype(np.float16),
+             y_ptr=np.zeros((n_rows, hidden_dim), dtype=np.float16),
+             w_ptr=rng.uniform(0.5, 1.5, (hidden_dim,)).astype(np.float16)),
+        dict(n_rows=n_rows, N=hidden_dim, eps=np.float16(1e-5), BLOCK_SIZE=block_size))
 
 def run_kernel_sdpa(hw, seq_len, head_dim, block_m, rng=None):
     """Generate SDPA MLIR, create tensors, run, return LatencyReport."""