🤖 machine learning

LP-GEMM: Integrating Layout Propagation into GEMM Operations

This paper introduces LP-GEMM, a technique that eliminates redundant data packing and unpacking in sequential matrix multiplications by propagating memory layouts across operations, achieving significant speedups on x86 and RISC-V architectures while maintaining full BLAS semantic correctness.

Original authors: César Guedes Carneiro, Lucas Alvarenga, Guido Araujo, Sandro Rigo

Published 2026-04-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: César Guedes Carneiro, Lucas Alvarenga, Guido Araujo, Sandro Rigo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a massive, high-speed bakery that specializes in making complex, multi-layered cakes (these are your Machine Learning models). The most critical step in your bakery is mixing ingredients. In the world of computers, this mixing process is called GEMM (General Matrix Multiplication).

For decades, the standard way to run this bakery (using libraries like OpenBLAS) has been incredibly efficient at mixing one cake at a time. However, it follows a rigid rule: Every time you finish mixing one cake, you must immediately take it apart, put it back into a standard box, and then take it out of that box again to start the next cake.

The Problem: The "Re-packing" Bottleneck

Let's say you are making a 10-layer cake.

Layer 1: You mix the batter. The standard rule says, "Stop! Put this batter into a square box, seal it, and label it."
Layer 2: You need that batter to make the next layer. But it's in a box! So, you have to unbox it, pour it into a new mixing bowl, and then start mixing.
Layer 3: You finish mixing. Again, you must re-box it, seal it, and label it.
Layer 4: You have to unbox it again.

In a real computer, this "boxing and unboxing" (called packing and unpacking) happens millions of times. It doesn't actually mix the cake; it just moves it around. It wastes time, energy, and the baker's (CPU's) attention.

The Solution: LP-GEMM (The "Assembly Line" Approach)

The paper introduces LP-GEMM (Layout Propagation GEMM). Think of this as redesigning the bakery into a seamless assembly line.

Instead of stopping to box the batter after every layer, the bakers agree on a special, custom shape for the mixing bowls.

The First Baker (Initial Kernel): Takes the raw ingredients and puts them into this special custom shape.
The Middle Bakers (Intermediate Kernels): They receive the batter already in the custom shape. They mix the next layer directly into that shape. No unboxing. No re-boxing. They just keep the flow moving.
The Last Baker (Ending Kernel): Once the final cake is done, only then does someone take the custom-shaped batter and put it into the standard box so the customer can take it home.

Why This Matters

The paper tested this idea on two types of "bakers" (computer chips):

Intel x86 (The Powerhouse): A very fast, complex processor.
RISC-V (The Efficient Newcomer): A newer, simpler, and very energy-efficient processor.

The Results:

On the Intel chip: The new method was about 2.25 times faster for sequences of operations. It's like going from a 10-minute commute to a 4-minute commute.
On the RISC-V chip: The speedup was even more dramatic, up to 5 times faster. Because the RISC-V chip is more sensitive to "wasted movement," removing the packing/unpacking steps made a huge difference.

Real-World Application: The Llama 3.2 Test

To prove this wasn't just a theory, the authors built a version of Llama 3.2 (a famous AI chatbot) using this new method.

The Challenge: AI models like Llama don't just do one mix; they do a long chain of them (like the "Attention" mechanism in Transformers).
The Fix: They adapted the "middle bakers" to handle the custom shape even when the recipe called for a quick pause (like adding a spice or a temperature check).
The Outcome: The AI ran significantly faster, proving that keeping the data in a "flowing" state rather than a "boxed" state is a game-changer.

The Analogy Summary

Old Way (OpenBLAS): Like a courier who delivers a package, unpacks it, hands it to the next person, who then immediately repacks it, labels it, and sends it to the next person. Lots of wasted time with the tape gun and boxes.
New Way (LP-GEMM): Like a conveyor belt where the package stays on a specialized tray. The first person puts it on the tray, the next person works on it while it's still on the tray, and the next person does the same. Only at the very end is it put into a standard shipping box.

The Bottom Line

LP-GEMM is a clever trick that stops computers from doing unnecessary "cleaning up" between math problems. By letting data flow in a format that is ready for the next step, rather than forcing it into a standard format after every step, we can make AI and scientific computing much faster and more efficient, without needing to invent new hardware.

1. Problem Statement

In modern Machine Learning (ML) and Scientific Computing workloads, execution time is often dominated by sequences of dependent General Matrix Multiplications (GEMMs), such as those found in Multi-Layer Perceptrons (MLPs) and Transformer attention mechanisms.

The Bottleneck: State-of-the-art BLAS libraries (e.g., OpenBLAS, Intel MKL) optimize individual GEMM calls aggressively. However, they are constrained by the standard BLAS API, which mandates that every operation must independently pack input matrices into a specific internal layout for computation and unpack the output back to a canonical memory layout (e.g., row-major or column-major).
Redundancy: In sequential GEMM chains, the output of one operation immediately becomes the input of the next. Consequently, the library performs a redundant unpacking step (to restore the canonical layout) followed immediately by a repacking step (to prepare for the next operation). This cycle wastes significant computational resources and memory bandwidth, particularly in data movement rather than arithmetic.

2. Methodology: LP-GEMM

The authors propose LP-GEMM (Layout Propagation GEMM), a decomposition of the standard GEMM kernel designed to eliminate redundant packing/unpacking by propagating the internal data layout across sequential operations.

Core Concept

Instead of treating each GEMM as an isolated kernel that resets the data layout, LP-GEMM aligns three critical aspects of the computation:

Packed Layout: The internal memory arrangement used for reading inputs.
Iteration Order: The order in which the micro-kernel computes partial results.
Final Stored Layout: The layout in which the result is written to memory.

By unifying these three ( $Layout_{packed} = Layout_{compute} = Layout_{stored}$ ), the output of one GEMM is immediately consumable by the next without intermediate reordering.

Kernel Decomposition

To maintain correctness at the boundaries of a computation chain while enabling propagation internally, LP-GEMM splits the GEMM operation into three specialized kernels:

Initial GEMM (ini-GEMM):
- Performs the initial packing of input matrices.
- Executes the computation using a Propagate-Layout Micro-kernel.
- Stores the output in the propagated (non-canonical) layout.
Intermediate GEMM (mid-GEMM):
- Assumes the input multiplier is already packed in the propagated layout.
- Skips the packing step entirely.
- Continues the layout propagation, storing results in the same propagated format.
- Note: This kernel is used for the bulk of operations in a sequence.
Ending GEMM (end-GEMM):
- Performs the final computation.
- Uses a Default Micro-kernel that resets the data back to the canonical BLAS layout.
- Unpacks the result for the user or subsequent non-GEMM operations.

Integration with ML Workloads

The paper addresses the challenge of integrating layout propagation with non-linear operations (e.g., Softmax, RoPE) found in Transformer attention layers:

RoPE (Rotary Position Embedding): Naturally aligns with the propagated layout as it operates row-wise, allowing for efficient SIMD execution.
Softmax: Requires adaptation. Since the propagated layout interleaves rows for outer-product processing, Softmax is reorganized to operate over multiple rows simultaneously to minimize address-jump overhead.

3. Key Contributions

LP-GEMM Decomposition: A novel architectural decomposition of the OpenBLAS GEMM kernel into ini-GEMM, mid-GEMM, and end-GEMM to enable data layout propagation.
Standalone Implementation: A complete C++ implementation of the Llama-3.2 (1B) inference path using exclusively BLAS-level GEMM calls, demonstrating the practicality of the approach outside of microbenchmarks.
Cross-Architecture Validation: Implementation and evaluation on both x86 (AVX-512) and RISC-V (RVV 1.0) architectures, proving the approach is not tied to specific hardware vector extensions.
Additional Optimizations: Implementation of other common ML matrix operations (Softmax, RoPE, RMSNorm) optimized for the propagated layout.

4. Experimental Results

The authors evaluated LP-GEMM against OpenBLAS, BLIS, Intel MKL, and FlashGEMM on Intel Xeon Gold 6252 (x86) and SpacemiT X60 (RISC-V).

Single GEMM Performance:
- The ini-GEMM kernel performs similarly to OpenBLAS.
- The mid-GEMM and end-GEMM kernels achieve a median speedup of ~1.5x over OpenBLAS, with some cases exceeding 2x. This confirms that packing/unpacking is often the dominant cost, not the micro-kernel arithmetic itself.
Sequential GEMMs (Attention Layer):
- x86: Achieved average speedups of 2.25x over OpenBLAS for sequential operations.
- RISC-V: Showed even more dramatic gains, with speedups reaching 5x for Attention-like workloads. This is attributed to the inefficiency of the baseline OpenBLAS RISC-V kernel, which suffers heavily from out-of-order memory accesses during unpacking.
Consecutive GEMM Benchmarks:
- In benchmarks consisting of three consecutive GEMMs (simulating DNN bottleneck blocks), LP-GEMM outperformed both OpenBLAS and the fusion-based FlashGEMM on most test cases.
Scalability:
- On x86, speedups scale inversely with matrix size (packing costs become negligible for very large matrices compared to $O(n^3)$ compute).
- On RISC-V, speedups scale linearly with problem size because the penalty of scattered memory accesses in the baseline grows with matrix size, while LP-GEMM maintains contiguous access.

5. Significance and Conclusion

Paradigm Shift: LP-GEMM shifts the optimization focus from "isolated kernel tuning" to "pipeline-aware optimization." It demonstrates that understanding the data flow between operations can yield significant performance gains without requiring deep microarchitectural re-tuning.
Practicality: The approach is highly practical for Edge AI and inference workloads where latency is critical. The implementation of Llama-3.2 proves that these optimizations can be integrated into real-world models with modest code changes.
Portability: The technique is orthogonal to the underlying micro-kernel design. It can be combined with highly optimized vendor kernels (like MKL) to further close the performance gap, suggesting a future where layout propagation is a standard feature in linear algebra libraries.
Future Work: The authors suggest extending this approach to other operations like convolutions and deeper integration with compiler-driven optimization frameworks to automate layout propagation decisions.

In summary, LP-GEMM effectively eliminates the "packing tax" in sequential matrix multiplications, delivering substantial performance improvements across diverse hardware architectures and real-world ML models.