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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/221814.diff


1 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp (+87-34) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
index c84e5e50cb0e2..bbb23c00e5e19 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.

``````````

</details>


https://github.com/llvm/llvm-project/pull/221814