[LinalgExt][Attention] Specialize generics to matmuls (#24668)
Specialize generic matmul to named one in decompose-attention to prevent
BubbleUpExpandShapes from propagating reshapes through contractions,
preserving codegen-friendly structure for microkernels.
Signed-off-by: Kamil Karwacki <karwacki@roofline.ai>
Co-authored-by: Thomas Ziereis <ziereis@roofline.ai>
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp
index b12f5a1..1b76613 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp
@@ -16,6 +16,7 @@
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -153,7 +154,7 @@
static Value computeMatmul(OpBuilder &builder, Location loc, AffineMap lhsMap,
AffineMap rhsMap, AffineMap accMap, Value lhs,
- Value rhs, Value acc) {
+ Value rhs, Value acc, bool specialize = false) {
SmallVector<AffineMap> compressedMaps =
compressUnusedDims(SmallVector<AffineMap>{lhsMap, rhsMap, accMap});
@@ -183,6 +184,19 @@
linalg::YieldOp::create(b, loc, add);
});
+ if (specialize) {
+ // Try to specialize the generic into a named op (e.g. batch_matmul).
+ // Named ops prevent BubbleUpExpandShapes from propagating reshapes through
+ // contractions, preserving codegen-friendly structure for microkernels.
+ IRRewriter rewriter(builder);
+ rewriter.setInsertionPoint(genericOp);
+ FailureOr<linalg::LinalgOp> specializedOp =
+ linalg::specializeGenericOp(rewriter, genericOp);
+ if (succeeded(specializedOp)) {
+ return specializedOp->getOperation()->getResult(0);
+ }
+ }
+
return genericOp.getResult(0);
}
@@ -313,7 +327,7 @@
SmallVector<OpFoldResult> iterationDomain,
Type sElementType, Region &elementwiseRegion,
DictionaryAttr qkAttrs, bool lowPrecision,
- bool useExp2) {
+ bool useExp2, bool specialize = false) {
MLIRContext *ctx = b.getContext();
// If using exp2 for attention instead of the original exp, we have to
// multiply the scale by log2(e). We use exp2 instead of exp as most platforms
@@ -358,7 +372,7 @@
Value sZero = arith::ConstantOp::create(b, loc, b.getZeroAttr(sElementType));
Value s = linalg::FillOp::create(b, loc, sZero, emptyS).getResult(0);
- s = computeMatmul(b, loc, qMap, kMap, sMap, query, key, s);
+ s = computeMatmul(b, loc, qMap, kMap, sMap, query, key, s, specialize);
if (qkAttrs) {
s.getDefiningOp()->setAttrs(qkAttrs);
}
@@ -437,7 +451,8 @@
// ---- QK Matmul + elementwise math ----
Value s = computeQKAndElementwise(
loc, b, query, key, getScale(), mask, qMap, kMap, sMap, getMaskMap(),
- sizes, f32Type, getRegion(), qkAttrs, lowPrecision, /*useExp2=*/true);
+ sizes, f32Type, getRegion(), qkAttrs, lowPrecision, /*useExp2=*/true,
+ /*specialize=*/true);
// ---- Softmax ----
@@ -502,8 +517,8 @@
}
// result = P @ V + acc
- Value result =
- computeMatmul(b, loc, pMap, getValueMap(), accMap, p, value, accFill);
+ Value result = computeMatmul(b, loc, pMap, getValueMap(), accMap, p, value,
+ accFill, /*specialize=*/true);
if (pvAttrs) {
result.getDefiningOp()->setAttrs(pvAttrs);
}
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/BUILD.bazel b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/BUILD.bazel
index d30d0d7..9480824 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/BUILD.bazel
@@ -96,6 +96,7 @@
"@llvm-project//mlir:InliningUtils",
"@llvm-project//mlir:LinalgDialect",
"@llvm-project//mlir:LinalgStructuredOpsIncGen",
+ "@llvm-project//mlir:LinalgTransforms",
"@llvm-project//mlir:LinalgUtils",
"@llvm-project//mlir:MathDialect",
"@llvm-project//mlir:MemRefDialect",
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/CMakeLists.txt
index 50a477c..2621f05 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/CMakeLists.txt
@@ -57,6 +57,7 @@
MLIRInferTypeOpInterface
MLIRLinalgDialect
MLIRLinalgStructuredOpsIncGenLib
+ MLIRLinalgTransforms
MLIRLinalgUtils
MLIRMathDialect
MLIRMemRefDialect
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/test/decompose_aggregate_op.mlir b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/test/decompose_aggregate_op.mlir
index 85c2181..ff12285 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/test/decompose_aggregate_op.mlir
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/test/decompose_aggregate_op.mlir
@@ -107,12 +107,7 @@
// CHECK: linalg.generic
// CHECK: arith.mulf
// S = Q @ K
-// CHECK: linalg.generic
-// CHECK: arith.extf
-// CHECK: arith.extf
-// CHECK: arith.mulf
-// CHECK: arith.addf
-// CHECK: linalg.yield
+// CHECK: linalg.batch_matmul
// max = rowMax(S)
// CHECK: linalg.generic
// CHECK-NOT: arith.extf
@@ -139,13 +134,8 @@
// CHECK-NOT: arith.extf
// CHECK: arith.truncf
// CHECK: linalg.yield
-// newAcc = P @ V
-// CHECK: linalg.generic
-// CHECK: arith.extf
-// CHECK: arith.extf
-// CHECK: arith.mulf
-// CHECK: arith.addf
-// CHECK: linalg.yield
+// result = P @ V
+// CHECK: linalg.batch_matmul
// -----
