gemv: add optional fused RMSNorm/LayerNorm prologue

iron/operators/gemv/design.py

@@ -8,7 +8,7 @@
 from aie.dialects.aie import T
 from aie.helpers.dialects.scf import _for as range_
 from aie.helpers.taplib import TensorAccessPattern
-from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker
+from aie.iron import Buffer, Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker
 
 """
 Matrix-vector design
@@ -37,6 +37,7 @@ def my_matvec(
     func_prefix="",
     verbose=False,
     epilogue="none",
+    prologue="none",
 ):
     if m_output is None:
         m_output = m_input
@@ -101,6 +102,26 @@ def my_matvec(
             [np.int32, L1_C_ty],
         )
 
+    # Optional norm applied to the shared B vector once per acquire, before it feeds every
+    # matvec call in this batch -- so the cost is O(K) per batch, not O(M) as re-normalizing
+    # per output tile would be. Both kernels take separate in/out pointers (restrict on
+    # both), so normalizing needs a second L1 buffer rather than an in-place call.
+    assert prologue in ("none", "rms", "ln")
+    norm_kernel = None
+    b_norm_bufs = [None] * cols
+    PROLOGUE_EPSILON = 1e-5  # matches layer_norm.cc's hardcoded epsilon
+    if prologue != "none":
+        norm_fn = "rms_norm_bf16_vector" if prologue == "rms" else "layer_norm"
+        norm_args = (
+            [L1_B_ty, L1_B_ty, np.int32, np.float32]
+            if prologue == "rms"
+            else [L1_B_ty, L1_B_ty, np.int32]
+        )
+        norm_kernel = Kernel(
+            f"{func_prefix}{norm_fn}", f"{func_prefix}{kernel_object}", norm_args
+        )
+        b_norm_bufs = [Buffer(type=L1_B_ty, name=f"B_norm_{i}") for i in range(cols)]
+
     A_L3L1_fifos = [
         ObjectFifo(L1_A_ty, name=f"A_L3L1_{i}", depth=2) for i in range(cols)
     ]
@@ -111,10 +132,20 @@ def my_matvec(
         ObjectFifo(L1_C_ty, name=f"C_L1L3_{i}", depth=2) for i in range(cols)
     ]
 
-    def core_body(A_L3L1_fifo, B_L3L1_fifo, C_L1L3_fifo, matvec, gelu_kernel=None):
+    def core_body(
+        A_L3L1_fifo, B_L3L1_fifo, C_L1L3_fifo, matvec, norm_kernel, b_norm, gelu_kernel
+    ):
         one_idx = index.constant(1)
         for _ in range_(0xFFFFFFFF):  # batch dim handled as part of this loop
-            b = B_L3L1_fifo.acquire(1)
+            b_raw = B_L3L1_fifo.acquire(1)
+            if norm_kernel is not None:
+                if prologue == "rms":
+                    norm_kernel(b_raw, b_norm, K, PROLOGUE_EPSILON)
+                else:
+                    norm_kernel(b_raw, b_norm, K)
+                b = b_norm
+            else:
+                b = b_raw
             # The kernel function computes m output rows; each core is responsible for (M/cols) output rows, so we need to call the kernel (M/cols)/m times.
             for i_idx in range_(M // m_output // cols):
                 c = C_L1L3_fifo.acquire(1)
@@ -138,8 +169,10 @@ def core_body(A_L3L1_fifo, B_L3L1_fifo, C_L1L3_fifo, matvec, gelu_kernel=None):
                 B_L3L1_fifos[i].cons(),
                 C_L1L3_fifos[i].prod(),
                 matvec,
-            ]
-            + ([gelu_kernel] if epilogue == "gelu" else []),
+                norm_kernel,
+                b_norm_bufs[i],
+                gelu_kernel,
+            ],
         )
         for i in range(cols)
     ]

iron/operators/gemv/op.py

@@ -32,6 +32,10 @@ class GEMV(MLIROperator):
     # "none" (default) leaves the output unchanged; "gelu" applies GELU(tanh approx).
     # repr=False keeps operator/artifact names stable for the default path.
     epilogue: str = field(default="none", repr=False)
+    # Optional norm applied to the shared vector once before the matvec, fusing a
+    # RMSNorm/LayerNorm decode prologue into the same dispatch. Affine-free (gamma/beta
+    # fold into the weight matrix host-side, same convention the norm kernels already use).
+    prologue: str = field(default="none", repr=False)
     context: object = field(default=None, repr=False)
 
     _name_aliases: ClassVar[Dict[str, str]] = {
@@ -63,6 +67,10 @@ def __post_init__(self):
             raise ValueError(
                 f"gelu epilogue needs tile_size_output % 16 == 0 (got {self.tile_size_output})"
             )
+        if self.prologue not in ("none", "rms", "ln"):
+            raise ValueError(
+                f"unknown prologue {self.prologue!r} (expected 'none', 'rms' or 'ln')"
+            )
 
         MLIROperator.__init__(self, context=self.context)
 
@@ -73,17 +81,25 @@ def name(self) -> str:
         # both would emit the same .mlir/.xclbin, and in a shared build dir a cached unfused
         # build can then satisfy the fused op (running the raw matvec with no activation).
         base = super().name
-        if self.epilogue == "none":
-            return base
-        return f"{base}_epi{self.epilogue}"
+        suffix = ""
+        if self.prologue != "none":
+            suffix += f"_pro{self.prologue}"
+        if self.epilogue != "none":
+            suffix += f"_epi{self.epilogue}"
+        return base + suffix
 
     @property
     def _kernel_link_file(self):
-        # With the gelu epilogue the core also links the gelu kernel, so the object becomes an
-        # archive of (matvec, gelu); the plain matvec stays a single object.
-        if self.epilogue == "gelu":
-            return f"gemv_{self.K}k_{self.kernel_vector_size}vs_gelu_kernels.a"
-        return f"gemv_{self.K}k_{self.kernel_vector_size}vs.o"
+        # With a prologue and/or the gelu epilogue the core links extra kernels, so the
+        # object becomes an archive; with neither, the plain matvec stays a single object.
+        if self.prologue == "none" and self.epilogue == "none":
+            return f"gemv_{self.K}k_{self.kernel_vector_size}vs.o"
+        suffix = ""
+        if self.prologue != "none":
+            suffix += f"_pro{self.prologue}"
+        if self.epilogue != "none":
+            suffix += f"_epi{self.epilogue}"
+        return f"gemv_{self.K}k_{self.kernel_vector_size}vs{suffix}_kernels.a"
 
     def get_mlir_artifact(self):
         mlir_verbose = getattr(self.context, "mlir_verbose", False)
@@ -106,6 +122,7 @@ def get_mlir_artifact(self):
                     "verbose": mlir_verbose,
                     "kernel_object": self._kernel_link_file,
                     "epilogue": self.epilogue,
+                    "prologue": self.prologue,
                 },
             ),
         )
@@ -123,27 +140,48 @@ def get_kernel_artifacts(self):
                 f"-DVEC_SIZE={self.kernel_vector_size}",
             ],
         )
+        extra_objs = []
+        if self.prologue != "none":
+            # rms_norm.cc / layer_norm.cc exist under both aie2 and aie2p with the same
+            # extern-C signature, so the prologue works on NPU1 and NPU2 alike.
+            norm_src = "rms_norm.cc" if self.prologue == "rms" else "layer_norm.cc"
+            extra_objs.append(
+                KernelObjectArtifact(
+                    norm_src.replace(".cc", ".o"),
+                    dependencies=[
+                        SourceArtifact(
+                            self.context.base_dir
+                            / "aie_kernels"
+                            / get_kernel_dir()
+                            / norm_src
+                        )
+                    ],
+                )
+            )
         if self.epilogue == "gelu":
             # The gelu kernel lives in aie2p/gelu.cc, so the fused epilogue is NPU2-only.
             if get_kernel_dir() != "aie2p":
                 raise NotImplementedError(
                     "gemv gelu epilogue is only available on NPU2 (aie2p); "
                     f"current kernel dir is {get_kernel_dir()!r}"
                 )
-            gelu_obj = KernelObjectArtifact(
-                "gelu.o",
-                dependencies=[
-                    SourceArtifact(
-                        self.context.base_dir / "aie_kernels" / "aie2p" / "gelu.cc"
-                    )
-                ],
-            )
-            return [
-                KernelArchiveArtifact(
-                    self._kernel_link_file, dependencies=[matvec_obj, gelu_obj]
+            extra_objs.append(
+                KernelObjectArtifact(
+                    "gelu.o",
+                    dependencies=[
+                        SourceArtifact(
+                            self.context.base_dir / "aie_kernels" / "aie2p" / "gelu.cc"
+                        )
+                    ],
                 )
-            ]
-        return [matvec_obj]
+            )
+        if not extra_objs:
+            return [matvec_obj]
+        return [
+            KernelArchiveArtifact(
+                self._kernel_link_file, dependencies=[matvec_obj] + extra_objs
+            )
+        ]
 
     def get_arg_spec(self):
         batch_dim = (self.num_batches,) if self.num_batches > 1 else ()

iron/operators/gemv/reference.py

@@ -74,3 +74,22 @@ def gelu_tanh_approx(x):
     xf = np.asarray(x, dtype=np.float32)
     inner = 0.79788456 * (xf + 0.044715 * xf**3)
     return 0.5 * xf * (1.0 + np.tanh(inner))
+
+
+def rms_norm_ref(x, epsilon=1e-5):
+    """RMSNorm, matching aie_kernels/{aie2,aie2p}/rms_norm.cc's rms_norm_bf16_vector:
+    f32 sum-of-squares reduction, affine-free (gamma=1), computed in float32.
+    """
+    xf = np.asarray(x, dtype=np.float32)
+    inv_rms = 1.0 / np.sqrt(np.mean(xf * xf) + epsilon)
+    return xf * inv_rms
+
+
+def layer_norm_ref(x, epsilon=1e-5):
+    """LayerNorm, matching aie_kernels/{aie2,aie2p}/layer_norm.cc's layer_norm:
+    f32 mean/var reduction, affine-free (gamma=1, beta=0), computed in float32.
+    """
+    xf = np.asarray(x, dtype=np.float32)
+    mean = np.mean(xf)
+    var = np.mean(xf * xf) - mean * mean
+    return (xf - mean) / np.sqrt(var + epsilon)

iron/operators/gemv/test.py

@@ -10,6 +10,8 @@
     generate_golden_reference,
     generate_golden_reference_batched,
     gelu_tanh_approx,
+    rms_norm_ref,
+    layer_norm_ref,
 )
 from iron.common.device_utils import get_kernel_dir
 import numpy as np
@@ -186,3 +188,60 @@ def test_gemv_gelu(
     print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n")
 
     assert not errors, f"Test failed with errors: {errors}"
+
+
+@pytest.mark.metrics(
+    Latency=r"Latency \(us\): (?P<value>[\d\.]+)",
+    Bandwidth=r"Effective Bandwidth: (?P<value>[\d\.e\+-]+) GB/s",
+    Throughput=r"Throughput: (?P<value>[\d\.e\+-]+) GFLOP/s",
+)
+@pytest.mark.parametrize(
+    "M,K,num_aie_columns,tile_size_input,tile_size_output,prologue,norm_ref",
+    [
+        pytest.param(128, 128, 1, 32, 128, "rms", rms_norm_ref),
+        pytest.param(2048, 8192, 1, 1, 2048, "rms", rms_norm_ref),
+        pytest.param(8192, 2048, 1, 4, 1024, "rms", rms_norm_ref),
+        pytest.param(128, 128, 1, 32, 128, "ln", layer_norm_ref),
+        pytest.param(2048, 8192, 1, 1, 2048, "ln", layer_norm_ref),
+        pytest.param(8192, 2048, 1, 4, 1024, "ln", layer_norm_ref),
+    ],
+)
+def test_gemv_norm_prologue(
+    M,
+    K,
+    num_aie_columns,
+    tile_size_input,
+    tile_size_output,
+    prologue,
+    norm_ref,
+    aie_context,
+):
+    """GEMV with a fused RMSNorm/LayerNorm prologue vs an A @ norm(B) golden."""
+    golden_ref = generate_golden_reference(M=M, K=K)
+    b_ref = golden_ref["B"].to(torch.float32).numpy()
+    b_norm = torch.from_numpy(norm_ref(b_ref).astype(np.float32)).to(torch.bfloat16)
+    c_pro = (golden_ref["A"] @ b_norm).to(torch.bfloat16)
+
+    operator = GEMV(
+        M=M,
+        K=K,
+        num_aie_columns=num_aie_columns,
+        tile_size_input=tile_size_input,
+        tile_size_output=tile_size_output,
+        prologue=prologue,
+        context=aie_context,
+    )
+
+    input_buffers = {"matrix": golden_ref["A"].flatten(), "vector": golden_ref["B"]}
+    output_buffers = {"output": c_pro}
+
+    errors, latency_us, bandwidth_gbps = run_test(
+        operator, input_buffers, output_buffers, rel_tol=0.05, abs_tol=1e-2
+    )
+
+    print(f"\nLatency: {latency_us:.1f} us")
+    gflops = (2.0 * M * K) / (latency_us * 1e-6) / 1e9
+    print(f"Throughput: {gflops:.6e} GFLOP/s")
+    print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n")
+
+    assert not errors, f"Test failed with errors: {errors}"