Domain-specific Algorithm
Auto-vectorization
Title Slide with
2 MLIR
Meng Sun, Advay Singh, Linus Wang, Panagiotis Papanikolaou
Apr. 22, 2025
Needleman-Wunsch Algorithm
● The algorithm uses a grid (matrix) to compare two sequences (like DNA or
proteins) character by character
● Start by filling the first row and column with gap penalties (usually negative
numbers)
● For each remaining cell, calculate a score by taking the maximum of three
options:
○ Diagonal move (match/mismatch): Add match score or subtract
mismatch penalty from diagonal cell
○ Horizontal move (gap in sequence 1): Add gap penalty to left cell
○ Vertical move (gap in sequence 2): Add gap penalty to cell above
● After filling the entire grid, trace backward from bottom-right to top-left to find
the optimal alignment
● When tracing back, diagonal moves represent aligned pairs, horizontal
moves represent gaps in sequence 1, and vertical moves represent gaps in
sequence 2
Match: +2
Mismatch: -1
Gap: -2
Parallelising Sequence Alignment
● By default, the Needleman-Wunsch
algorithm is highly sequential and
has O(mn) time complexity
● Inter-sequence parallelization cannot
be done under the existing LLVM
infrastructure
● Vectron adds a compiler pass that
allows for vectorising operations
between different instances of the
sequence alignment problem
(batch processing)
● We replicate that in MLIR
MLIR / Polygeist
● MLIR offers a rapid way to create higher-level
IR, allowing auto-vectorization to occur at a
higher level of abstraction.
● MLIR can provide different IRs for different
languages or applications, making
application-specific optimizations possible.
● Polygeist serves as a frontend translator,
converting high-level languages like C/C++
into MLIR, enabling gradual lowering and
optimization of source code.
MLIR Pass Infrastructure mlir::PassRegistration<mlir::Vectorize>();
...
1. Inherit the CRTP class Vectorize :
public PassWrapper<Vectorize,
“PassWrapper” utility OperationPass<mlir::func::FuncOp>> {
class
private:
void runOnOperation() override;
2. Implement the virtual StringRef getArgument() const final { return
"vectorize"; }
function
“runOnOperation()”, StringRef getDescription() const final {
which handles the real return "Vectorize Needleman-Wunsch
Algorithm across multiple instances";
pass }
};
3. “PassRegistration”
Vectorization
[Link] %arg3 = 0 to 512 step 8 {
[Link] %arg3 = 0 to 512 {
[Link] %arg4 = 1 to 51 {
[Link] %arg4 = 1 to 51 {
[Link] %arg5 = 1 to 51 {
[Link] %arg5 = 1 to 51 {
%3 = [Link] %arg2[%arg4 - 1, %arg5,
%0 = [Link] %arg2[%arg4 - 1, %arg5, %arg3] :
%arg3] : memref<?x51x512xi32>, vector<8xi32>
memref<?x51x512xi32>
%4 = [Link] %3, %0 : vector<8xi32>
%1 = [Link] %0, %c-5_i32 : i32
%5 = [Link] %arg2[%arg4, %arg5 - 1,
%2 = [Link] %arg2[%arg4, %arg5 - 1, %arg3] :
%arg3] : memref<?x51x512xi32>, vector<8xi32>
memref<?x51x512xi32>
%6 = [Link] %5, %0 : vector<8xi32>
%3 = [Link] %2, %c-5_i32 : i32
%7 = [Link] %arg2[%arg4 - 1, %arg5 - 1,
%4 = [Link] %arg2[%arg4 - 1, %arg5 - 1, %arg3] :
%arg3] : memref<?x51x512xi32>, vector<8xi32>
memref<?x51x512xi32>
%8 = [Link] %arg0[%arg4 - 1, %arg3] :
%5 = [Link] %arg0[%arg4 - 1, %arg3] :
memref<?x512xi32>, vector<8xi32>
memref<?x512xi32>
%9 = [Link] %arg1[%arg4 - 1, %arg3] :
%6 = [Link] %arg1[%arg4 - 1, %arg3] :
memref<?x512xi32>, vector<8xi32>
memref<?x512xi32>
%10 = [Link] eq, %8, %9 : vector<8xi32>
%7 = [Link] eq, %5, %6 : i32
...
...
[Link] %18, %arg2[%arg4, %arg5, %arg3] :
[Link] %15, %arg2[%arg4, %arg5, %arg3] :
memref<?x51x512xi32>, vector<8xi32>
memref<?x51x512xi32>
}
}
}
}
}
}
MLIR
Dialect to
LLVM IR