[Mlir-commits] [mlir] [MLIR][linalg] Share external interface implementations to reduce build time (NFC) (PR #221814)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/221814

The Linalg tiling and partial-reduction external models instantiate their operation-independent implementations for every registered structured op. Keep thin operation-specific external models, but inherit the implementations from non-template bases so the compiler emits and optimizes them only once. 

Assisted-by: Codex

>From 922fb041e493ba15331422ae6afa9cf69ff16542 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Mon, 7 Sep 2026 05:53:18 -0700
Subject: [PATCH] [mlir][linalg] Share external interface implementations

The Linalg tiling and partial-reduction external models instantiate their
operation-independent implementations for every registered structured op.
Keep thin operation-specific external models, but inherit the implementations
from non-template bases so the compiler emits and optimizes them only once.

The baseline was captured during a clean all build of 5,525 Ninja steps. That
build took 7m17.17s wall time, 11,387.70s user time, and 694.49s system time at
2,763% effective CPU, with 3,007,068 KiB peak RSS and zero swaps. It produced
1,377 MLIR C++ TU traces totaling 9,710.73s of ExecuteCompiler time.

Before this change, TilingInterfaceImpl.cpp took 86.197s ExecuteCompiler,
31.278s frontend, 54.444s backend, 41.046s optimizer, 13.363s code generation,
21.486s function instantiation, 20.940s pending instantiations, 11.571s class
instantiation, 5.283s source processing, and 4.865s class parsing. The single
registerAll instantiation for 90 generated operation types took 19.205s.

The final source took 7.286s ExecuteCompiler, 3.811s frontend, 3.413s backend,
2.211s optimizer, 1.199s code generation, 1.480s function instantiation,
1.226s pending instantiations, 1.174s class instantiation, 2.141s source
processing, and 2.030s class parsing. The registerAll instantiation fell to
0.503s.

That is a 91.5% reduction in ExecuteCompiler time, 87.8% in frontend time,
93.7% in backend time, 94.6% in optimizer time, 91.0% in code generation, and
97.4% in registerAll instantiation time. The baseline came from a contended
full build while the final trace came from an isolated compile, so wall time is
not a controlled A/B; the phase and object-size reductions independently
confirm the improvement.

The final one-object compile took 7.33s wall time, 6.90s user time, and 0.30s
system time at 98% CPU, with 638,176 KiB peak RSS and zero swaps. The trace
shrank from 35,940,437 to 2,094,717 bytes, a 94.2% reduction.

Object text shrank from 2,605,544 to 572,391 bytes, a 78.0% reduction, and the
object file shrank from 5,870,360 to 2,137,088 bytes, a 63.6% reduction. Text
symbols belonging to the tiling and partial-reduction model families shrank
from 2,201,961 to 279,183 bytes, an 87.3% reduction. The straight-line
registerAll function remained 78,059 bytes.

The initial validation rebuild of MLIRLinalgTransforms and mlir-opt took 6.71s
wall time, 6.53s user time, and 3.00s system time at 141% CPU, with 1,629,384
KiB peak RSS and zero swaps. The 174-test Linalg dialect suite ran in 0.74s and
passed completely. The full check-mlir invocation took 7.38s wall time, 1.62s
user time, and 0.38s system time at 27% CPU, with 174,632 KiB peak RSS and zero
swaps; its lit suite ran in 6.44s with 3,937 passed, 619 unsupported, and 1
expected failure.

Assisted-by: Codex
---
 .../Linalg/Transforms/TilingInterfaceImpl.cpp | 121 +++++++++++++-----
 1 file changed, 87 insertions(+), 34 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
index c84e5e50cb0e2d..bbb23c00e5e193 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
@@ -200,31 +200,12 @@ validateTilingSemiAffineMaps(LinalgOp linalgOp, ArrayRef<OpFoldResult> sizes) {
 //===----------------------------------------------------------------------===//
 
 namespace {
-/// External model implementation of TilingInterface for LinalgOps. An external
-/// model implementation is used for now till the use of `TilingInterface` is
-/// on-par with the current Linalg tiling + fusion patterns. Once it is
-/// maybe possible to move this into the op-definition (though there are
-/// advantages to leaving it as an external model)
-template <typename LinalgOpTy>
-struct LinalgOpTilingInterface
-    : public TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
-                                            LinalgOpTy> {
-  using Base =
-      TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
-                                     LinalgOpTy>;
-  // Inherit the defaulted hint-bearing overloads; these ops do not require the
-  // hint (no inner tiles).
-  using Base::generateResultTileValue;
-  using Base::getIterationDomainTileFromOperandTiles;
-  using Base::getTiledImplementation;
-  using Base::getTiledImplementationFromOperandTiles;
-
-  /// Return the loop iterator type.
-  SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
-    LinalgOpTy concreteOp = cast<LinalgOpTy>(op);
-    return concreteOp.getIteratorTypesArray();
-  }
-
+/// Operation-independent implementation shared by the external models for
+/// LinalgOps. External models are used for now until `TilingInterface` is
+/// on-par with the current Linalg tiling and fusion patterns. It may then be
+/// possible to move this into the op definitions, though there are advantages
+/// to leaving it as an external model.
+struct LinalgOpTilingInterfaceImpl {
   /// Return the iteration domain range.
   SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
     OpBuilder::InsertionGuard g(b);
@@ -517,6 +498,69 @@ struct LinalgOpTilingInterface
   }
 };
 
+template <typename LinalgOpTy>
+struct LinalgOpTilingInterfaceModel
+    : public TilingInterface::ExternalModel<
+          LinalgOpTilingInterfaceModel<LinalgOpTy>, LinalgOpTy>,
+      public LinalgOpTilingInterfaceImpl {
+  using ExternalModel =
+      TilingInterface::ExternalModel<LinalgOpTilingInterfaceModel<LinalgOpTy>,
+                                     LinalgOpTy>;
+
+  using LinalgOpTilingInterfaceImpl::generateScalarImplementation;
+  using LinalgOpTilingInterfaceImpl::getIterationDomain;
+  using LinalgOpTilingInterfaceImpl::getIterationDomainTileFromResultTile;
+  using LinalgOpTilingInterfaceImpl::getResultTilePosition;
+  using LinalgOpTilingInterfaceImpl::isOpFusableWithConsumerSlice;
+  using LinalgOpTilingInterfaceImpl::isOpFusableWithProducerSlices;
+
+  /// Return the loop iterator type without a dynamic LinalgOp interface lookup.
+  SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
+    return cast<LinalgOpTy>(op).getIteratorTypesArray();
+  }
+
+  // Preserve the hint-bearing ExternalModel defaults while routing the
+  // no-hint overloads through the shared implementation.
+  using ExternalModel::generateResultTileValue;
+  FailureOr<TilingResult>
+  generateResultTileValue(Operation *op, OpBuilder &b, unsigned resultNumber,
+                          ArrayRef<OpFoldResult> offsets,
+                          ArrayRef<OpFoldResult> sizes) const {
+    return LinalgOpTilingInterfaceImpl::generateResultTileValue(
+        op, b, resultNumber, offsets, sizes);
+  }
+
+  using ExternalModel::getIterationDomainTileFromOperandTiles;
+  LogicalResult getIterationDomainTileFromOperandTiles(
+      Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
+      ArrayRef<SmallVector<OpFoldResult>> allOffsets,
+      ArrayRef<SmallVector<OpFoldResult>> allSizes,
+      SmallVectorImpl<OpFoldResult> &iterDomainOffsets,
+      SmallVectorImpl<OpFoldResult> &iterDomainSizes) const {
+    return LinalgOpTilingInterfaceImpl::getIterationDomainTileFromOperandTiles(
+        op, b, operandNumbers, allOffsets, allSizes, iterDomainOffsets,
+        iterDomainSizes);
+  }
+
+  using ExternalModel::getTiledImplementation;
+  FailureOr<TilingResult>
+  getTiledImplementation(Operation *op, OpBuilder &b,
+                         ArrayRef<OpFoldResult> offsets,
+                         ArrayRef<OpFoldResult> sizes) const {
+    return LinalgOpTilingInterfaceImpl::getTiledImplementation(op, b, offsets,
+                                                               sizes);
+  }
+
+  using ExternalModel::getTiledImplementationFromOperandTiles;
+  FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
+      Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
+      ArrayRef<SmallVector<OpFoldResult>> allOffsets,
+      ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
+    return LinalgOpTilingInterfaceImpl::getTiledImplementationFromOperandTiles(
+        op, b, operandNumbers, allOffsets, allSizes);
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // External Model for implementing `PartialReductionInterface` for `LinalgOp`s.
 //===----------------------------------------------------------------------===//
@@ -655,12 +699,9 @@ static InitSliceInfo getInitSliceInfo(MLIRContext *context,
       partialReductionMap, initOperandShape);
 }
 
-/// External model implementation of PartialReductionInterface for
-/// LinalgOps.
-template <typename LinalgOpTy>
-struct LinalgOpPartialReductionInterface
-    : public PartialReductionOpInterface::ExternalModel<
-          LinalgOpPartialReductionInterface<LinalgOpTy>, LinalgOpTy> {
+/// Operation-independent implementation shared by the
+/// PartialReductionInterface external models for LinalgOps.
+struct LinalgOpPartialReductionInterfaceImpl {
   FailureOr<SmallVector<Value>> generateInitialTensorForPartialReduction(
       Operation *op, OpBuilder &b, Location loc, ArrayRef<OpFoldResult> sizes,
       const SetVector<unsigned> &reductionDims) const {
@@ -892,6 +933,18 @@ struct LinalgOpPartialReductionInterface
   }
 };
 
+template <typename LinalgOpTy>
+struct LinalgOpPartialReductionInterfaceModel
+    : public PartialReductionOpInterface::ExternalModel<
+          LinalgOpPartialReductionInterfaceModel<LinalgOpTy>, LinalgOpTy>,
+      public LinalgOpPartialReductionInterfaceImpl {
+  using LinalgOpPartialReductionInterfaceImpl::
+      generateInitialTensorForPartialReduction;
+  using LinalgOpPartialReductionInterfaceImpl::getPartialResultTilePosition;
+  using LinalgOpPartialReductionInterfaceImpl::mergeReductions;
+  using LinalgOpPartialReductionInterfaceImpl::tileToPartialReduction;
+};
+
 template <typename OpTy>
 static SmallVector<Range> getPackUnPackIterationDomain(OpTy op,
                                                        OpBuilder &builder) {
@@ -2008,9 +2061,9 @@ struct UnPackOpTiling
 
 template <typename OpType>
 static void registerOne(MLIRContext *ctx) {
-  OpType::template attachInterface<LinalgOpTilingInterface<OpType>>(*ctx);
-  OpType::template attachInterface<LinalgOpPartialReductionInterface<OpType>>(
-      *ctx);
+  OpType::template attachInterface<LinalgOpTilingInterfaceModel<OpType>>(*ctx);
+  OpType::template attachInterface<
+      LinalgOpPartialReductionInterfaceModel<OpType>>(*ctx);
 }
 
 /// Variadic helper function.