diff --git a/aie_kernels/aie2p/gelu.cc b/aie_kernels/aie2p/gelu.cc index c964feb9..8ddfbd56 100644 --- a/aie_kernels/aie2p/gelu.cc +++ b/aie_kernels/aie2p/gelu.cc @@ -8,19 +8,13 @@ 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 +gelu_tanh_approx_v16(aie::vector 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 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(k0_5); @@ -28,36 +22,44 @@ void gelu_tanh_approx_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict o auto vs2opi = aie::broadcast(sqrt_2_over_pi); auto vBeta = aie::broadcast(kBeta); + aie::vector x2 = aie::mul(x, x); + aie::vector x3 = aie::mul(x, x2); + aie::vector x3_beta = aie::mul(x3, vBeta); + aie::vector inner = aie::add(x, x3_beta); + auto inner1 = aie::mul(inner, vs2opi); + auto tanh_out = aie::tanh(inner1.to_vector()); + aie::vector one_plus_tanh = aie::add(tanh_out, v1); + aie::vector mul_v05 = aie::mul(v05, one_plus_tanh); + return aie::mul(x, mul_v05).to_vector(); +} + +// 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); + 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 x2 = aie::mul(x, x); // x^2 - aie::vector x3 = aie::mul(x, x2); // x^3 - - // inner = sqrt(2/pi) * (x + 0.044715 * x^3) - aie::vector x3_beta = aie::mul(x3, vBeta); - aie::vector inner = aie::add(x, x3_beta); - auto inner1 = aie::mul(inner, vs2opi); - - // tanh_out = tanh(inner) - auto tanh_out = aie::tanh(inner1.to_vector()); - - // result = 0.5 * x * (1 + tanh_out) - aie::vector one_plus_tanh = aie::add(tanh_out, v1); - // Multiply by x and 0.5 - aie::vector mul_v05 = aie::mul(v05, one_plus_tanh); - auto result = aie::mul(x, mul_v05); - - *it_out++ = result.to_vector(); + *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(64) + for (int i = 0; i < vector_size; i += 16) { + aie::vector x = *it; + *it++ = gelu_tanh_approx_v16(x); + } + event1(); } extern "C" { @@ -67,4 +69,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" diff --git a/iron/operators/gemv/design.py b/iron/operators/gemv/design.py index 654dc7f6..6712e065 100644 --- a/iron/operators/gemv/design.py +++ b/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,18 @@ 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 +109,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 +123,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 +136,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) ] diff --git a/iron/operators/gemv/op.py b/iron/operators/gemv/op.py index 01739bf7..b7c4df7f 100644 --- a/iron/operators/gemv/op.py +++ b/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,23 @@ 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 _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 +91,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 () diff --git a/iron/operators/gemv/reference.py b/iron/operators/gemv/reference.py index dc72ae96..0f70dae3 100644 --- a/iron/operators/gemv/reference.py +++ b/iron/operators/gemv/reference.py @@ -40,3 +40,13 @@ def generate_golden_reference( "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)) diff --git a/iron/operators/gemv/test.py b/iron/operators/gemv/test.py index 2abb42c0..f7f1ba95 100755 --- a/iron/operators/gemv/test.py +++ b/iron/operators/gemv/test.py @@ -6,7 +6,10 @@ import aie.utils as aie_utils from iron.operators.gemv.op import GEMV -from iron.operators.gemv.reference import generate_golden_reference +from iron.operators.gemv.reference import generate_golden_reference, 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 @@ -69,3 +72,50 @@ def test_gemv(M, K, num_aie_columns, tile_size_input, tile_size_output, aie_cont 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[\d\.]+)", + Bandwidth=r"Effective Bandwidth: (?P[\d\.e\+-]+) GB/s", + Throughput=r"Throughput: (?P[\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}"