gemv: add optional fused GELU epilogue (NPU2)

aie_kernels/aie2p/gelu.cc

@@ -8,56 +8,57 @@
 
 using namespace aie;
 
-void gelu_tanh_approx_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict output_vector, const int32_t vector_size)
+// One 16-lane GELU (tanh approximation): 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))).
+static inline aie::vector<bfloat16, 16> gelu_tanh_approx_v16(aie::vector<bfloat16, 16> x)
 {
-    event0();
-
-    auto it_in = aie::begin_restrict_vector<16>((bfloat16 *)input_vector);
-    auto it_out = aie::begin_restrict_vector<16>((bfloat16 *)output_vector);
-
-    aie::vector<bfloat16, 16> input;
-
-    // Constants
     const bfloat16 k0_5 = 0.5f;
     const bfloat16 k1 = 1.0f;
-    const bfloat16 sqrt_2_over_pi = 0.79788456f; // ≈ sqrt(2/π)
+    const bfloat16 sqrt_2_over_pi = 0.79788456f; // sqrt(2/pi)
     const bfloat16 kBeta = 0.044715f;
 
     auto v05 = aie::broadcast<bfloat16, 16>(k0_5);
     auto v1 = aie::broadcast<bfloat16, 16>(k1);
     auto vs2opi = aie::broadcast<bfloat16, 16>(sqrt_2_over_pi);
     auto vBeta = aie::broadcast<bfloat16, 16>(kBeta);
 
-    AIE_PREPARE_FOR_PIPELINING
-    AIE_LOOP_MIN_ITERATION_COUNT(64)
-    for (int i = 0; i < vector_size; i += 16) {
-        input = *it_in++;
-        auto x = input;
-
-        // Compute x^3
-        aie::vector<bfloat16, 16> x2 = aie::mul(x, x);  // x^2
-        aie::vector<bfloat16, 16> x3 = aie::mul(x, x2); // x^3
-
-        // inner = sqrt(2/pi) * (x + 0.044715 * x^3)
-        aie::vector<bfloat16, 16> x3_beta = aie::mul(x3, vBeta);
-        aie::vector<bfloat16, 16> inner = aie::add(x, x3_beta);
-        auto inner1 = aie::mul(inner, vs2opi);
-
-        // tanh_out = tanh(inner)
-        auto tanh_out = aie::tanh<bfloat16>(inner1.to_vector<float>());
+    aie::vector<bfloat16, 16> x2 = aie::mul(x, x);
+    aie::vector<bfloat16, 16> x3 = aie::mul(x, x2);
+    aie::vector<bfloat16, 16> x3_beta = aie::mul(x3, vBeta);
+    aie::vector<bfloat16, 16> inner = aie::add(x, x3_beta);
+    auto inner1 = aie::mul(inner, vs2opi);
+    auto tanh_out = aie::tanh<bfloat16>(inner1.to_vector<float>());
+    aie::vector<bfloat16, 16> one_plus_tanh = aie::add(tanh_out, v1);
+    aie::vector<bfloat16, 16> mul_v05 = aie::mul(v05, one_plus_tanh);
+    return aie::mul(x, mul_v05).to_vector<bfloat16>();
+}
 
-        // result = 0.5 * x * (1 + tanh_out)
-        aie::vector<bfloat16, 16> one_plus_tanh = aie::add(tanh_out, v1);
-        // Multiply by x and 0.5
-        aie::vector<bfloat16, 16> mul_v05 = aie::mul(v05, one_plus_tanh);
-        auto result = aie::mul(x, mul_v05);
+// Out-of-place GELU: output_vector = gelu(input_vector). input and output must not alias.
+void gelu_tanh_approx_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict output_vector, const int32_t vector_size)
+{
+    event0();
+    auto it_in = aie::begin_restrict_vector<16>((bfloat16 *)input_vector);
+    auto it_out = aie::begin_restrict_vector<16>((bfloat16 *)output_vector);
 
-        *it_out++ = result.to_vector<bfloat16>();
+    AIE_PREPARE_FOR_PIPELINING
+    AIE_LOOP_MIN_ITERATION_COUNT(1)
+    for (int i = 0; i < vector_size; i += 16) {
+        *it_out++ = gelu_tanh_approx_v16(*it_in++);
     }
-
     event1();
+}
 
-    return;
+// In-place GELU: v = gelu(v). Single pointer, so aliasing-correct (each 16-lane slot is read then written).
+static inline void gelu_tanh_approx_inplace_bf16(bfloat16 *restrict v, const int32_t vector_size)
+{
+    event0();
+    auto it = aie::begin_restrict_vector<16>(v);
+    AIE_PREPARE_FOR_PIPELINING
+    AIE_LOOP_MIN_ITERATION_COUNT(1)
+    for (int i = 0; i < vector_size; i += 16) {
+        aie::vector<bfloat16, 16> x = *it;
+        *it++ = gelu_tanh_approx_v16(x);
+    }
+    event1();
 }
 
 extern "C" {
@@ -67,4 +68,11 @@ void gelu_bf16(bfloat16 *restrict input, bfloat16 *restrict output, int input_si
     gelu_tanh_approx_bf16(input, output, input_size);
 }
 
+// In-place GELU over n bf16 elements (n a multiple of 16). Intended as a fused epilogue over a compute
+// tile (e.g. a GEMV output tile), applied once per tile in the producing core.
+void gelu_tile_bf16(uint32_t n, bfloat16 *restrict c)
+{
+    gelu_tanh_approx_inplace_bf16(c, (int32_t)n);
+}
+
 } // extern "C"

iron/operators/gemv/design.py

@@ -36,6 +36,7 @@ def my_matvec(
     kernel_object="mv.o",
     func_prefix="",
     verbose=False,
+    epilogue="none",
 ):
     if m_output is None:
         m_output = m_input
@@ -85,6 +86,20 @@ def my_matvec(
         f"{func_prefix}{kernel_object}",
         [np.int32, np.int32, L1_A_ty, L1_B_ty, L1_C_ty],
     )
+    # Optional fused activation over the full m_output C-tile, applied once per tile in core_body
+    # (after the matvec inner-loop has filled all rows) rather than per matvec call, whose m_input
+    # tile can be smaller than the 16-wide activation vector.
+    assert epilogue in ("none", "gelu")
+    gelu_kernel = None
+    if epilogue == "gelu":
+        assert (
+            m_output % 16 == 0
+        ), f"gelu epilogue needs m_output % 16 == 0 (got {m_output})"
+        gelu_kernel = Kernel(
+            f"{func_prefix}gelu_tile_bf16",
+            f"{func_prefix}{kernel_object}",
+            [np.int32, L1_C_ty],
+        )
 
     A_L3L1_fifos = [
         ObjectFifo(L1_A_ty, name=f"A_L3L1_{i}", depth=2) for i in range(cols)
@@ -96,7 +111,7 @@ 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):
+    def core_body(A_L3L1_fifo, B_L3L1_fifo, C_L1L3_fifo, matvec, gelu_kernel=None):
         one_idx = index.constant(1)
         for _ in range_(0xFFFFFFFF):  # batch dim handled as part of this loop
             b = B_L3L1_fifo.acquire(1)
@@ -110,6 +125,8 @@ def core_body(A_L3L1_fifo, B_L3L1_fifo, C_L1L3_fifo, matvec):
                     a = A_L3L1_fifo.acquire(1)
                     matvec(m_input, output_row_offset, a, b, c)
                     A_L3L1_fifo.release(1)
+                if gelu_kernel is not None:
+                    gelu_kernel(m_output, c)
                 C_L1L3_fifo.release(1)
             B_L3L1_fifo.release(1)
 
@@ -121,7 +138,8 @@ def core_body(A_L3L1_fifo, B_L3L1_fifo, C_L1L3_fifo, matvec):
                 B_L3L1_fifos[i].cons(),
                 C_L1L3_fifos[i].prod(),
                 matvec,
-            ],
+            ]
+            + ([gelu_kernel] if epilogue == "gelu" else []),
         )
         for i in range(cols)
     ]

iron/operators/gemv/op.py

@@ -8,11 +8,13 @@
     MLIROperator,
     AIERuntimeArgSpec,
     KernelObjectArtifact,
+    KernelArchiveArtifact,
     SourceArtifact,
     PythonGeneratedMLIRArtifact,
     DesignGenerator,
 )
 import aie.utils as aie_utils
+from iron.common.device_utils import get_kernel_dir
 
 
 @dataclass
@@ -26,6 +28,10 @@ class GEMV(MLIROperator):
     tile_size_output: int | None = None
     num_batches: int = 1
     kernel_vector_size: int = field(default=64, repr=False)
+    # Optional fused activation applied to each output tile in the producing core.
+    # "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)
     context: object = field(default=None, repr=False)
 
     _name_aliases: ClassVar[Dict[str, str]] = {
@@ -49,9 +55,36 @@ def __post_init__(self):
             self.K >= self.kernel_vector_size and self.K % self.kernel_vector_size == 0
         ):
             raise ValueError("K must be multiple of kernel_vector_size")
+        if self.epilogue not in ("none", "gelu"):
+            raise ValueError(
+                f"unknown epilogue {self.epilogue!r} (expected 'none' or 'gelu')"
+            )
+        if self.epilogue == "gelu" and self.tile_size_output % 16 != 0:
+            raise ValueError(
+                f"gelu epilogue needs tile_size_output % 16 == 0 (got {self.tile_size_output})"
+            )
 
         MLIROperator.__init__(self, context=self.context)
 
+    @property
+    def name(self) -> str:
+        # epilogue is repr=False so the default path keeps a stable name, but the fused
+        # variant must not share an artifact name with the plain GEMV of the same shape:
+        # 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}"
+
+    @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"
+
     def get_mlir_artifact(self):
         mlir_verbose = getattr(self.context, "mlir_verbose", False)
 
@@ -71,26 +104,46 @@ def get_mlir_artifact(self):
                 ),
                 {
                     "verbose": mlir_verbose,
-                    "kernel_object": f"gemv_{self.K}k_{self.kernel_vector_size}vs.o",
+                    "kernel_object": self._kernel_link_file,
+                    "epilogue": self.epilogue,
                 },
             ),
         )
 
     def get_kernel_artifacts(self):
-        return [
-            KernelObjectArtifact(
-                f"gemv_{self.K}k_{self.kernel_vector_size}vs.o",
+        matvec_obj = KernelObjectArtifact(
+            f"gemv_{self.K}k_{self.kernel_vector_size}vs.o",
+            dependencies=[
+                SourceArtifact(
+                    self.context.base_dir / "aie_kernels" / "generic" / "mv.cc"
+                )
+            ],
+            extra_flags=[
+                f"-DDIM_K={self.K}",
+                f"-DVEC_SIZE={self.kernel_vector_size}",
+            ],
+        )
+        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" / "generic" / "mv.cc"
+                        self.context.base_dir / "aie_kernels" / "aie2p" / "gelu.cc"
                     )
                 ],
-                extra_flags=[
-                    f"-DDIM_K={self.K}",
-                    f"-DVEC_SIZE={self.kernel_vector_size}",
-                ],
-            ),
-        ]
+            )
+            return [
+                KernelArchiveArtifact(
+                    self._kernel_link_file, dependencies=[matvec_obj, gelu_obj]
+                )
+            ]
+        return [matvec_obj]
 
     def get_arg_spec(self):
         batch_dim = (self.num_batches,) if self.num_batches > 1 else ()

iron/operators/gemv/reference.py

@@ -64,3 +64,13 @@ def generate_golden_reference_batched(M=128, K=128, num_batches=2, seed=42):
     for b in range(num_batches):
         C[b] = A[b] @ B[b]
     return {"A": A, "B": B, "C": C}
+
+
+def gelu_tanh_approx(x):
+    """Tanh-approximation GELU, matching aie_kernels/aie2p/gelu.cc.
+
+    0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))). Computed in float32.
+    """
+    xf = np.asarray(x, dtype=np.float32)
+    inner = 0.79788456 * (xf + 0.044715 * xf**3)
+    return 0.5 * xf * (1.0 + np.tanh(inner))

iron/operators/gemv/test.py

@@ -9,7 +9,11 @@
 from iron.operators.gemv.reference import (
     generate_golden_reference,
     generate_golden_reference_batched,
+    gelu_tanh_approx,
 )
+from iron.common.device_utils import get_kernel_dir
+import numpy as np
+import torch
 from iron.common.test_utils import run_test
 
 
@@ -131,3 +135,54 @@ def test_gemv_batched(
     print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n")
 
     assert not errors, f"batched GEMV failed: {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",
+    [
+        pytest.param(128, 128, 1, 32, 128),
+        pytest.param(2048, 8192, 1, 1, 2048),
+        pytest.param(8192, 2048, 1, 4, 1024),
+    ],
+)
+def test_gemv_gelu(
+    M, K, num_aie_columns, tile_size_input, tile_size_output, aie_context
+):
+    """GEMV with the fused GELU epilogue (NPU2-only) vs a gelu(A @ B) golden."""
+    if get_kernel_dir() != "aie2p":
+        pytest.skip("gemv gelu epilogue is only available on NPU2 (aie2p)")
+
+    golden_ref = generate_golden_reference(M=M, K=K)
+    c_ref = golden_ref["C"].to(torch.float32).numpy()
+    c_gelu = torch.from_numpy(gelu_tanh_approx(c_ref).astype(np.float32)).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,
+        epilogue="gelu",
+        context=aie_context,
+    )
+
+    input_buffers = {"matrix": golden_ref["A"].flatten(), "vector": golden_ref["B"]}
+    output_buffers = {"output": c_gelu}
+
+    errors, latency_us, bandwidth_gbps = run_test(
+        operator, input_buffers, output_buffers, rel_tol=0.06, abs_tol=2e-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}"