[mlir][Linalg] Add a pattern to decompose `linalg.generic` ops.

[mravishankar]

Summary

This patch adds a pattern to decompose a linalg.generic operations
that

• has only parallel iterator types
• has more than 2 statements (including the yield)

into multiple linalg.generic operation such that each operation has
a single statement and a yield.
The pattern added here just splits the matching linalg.generic into
two linalg.generics, one containing the first statement, and the
other containing the remaining. The same pattern can be applied
repeatedly on the second op to ultimately fully decompose the generic
op.

[mravishankar]

Fix bug related to output indexing map handling.

[mravishankar]

Add lit tests and finish corner cases.

[mravishankar]

Fix lit test.

[stellaraccident]

A few small comments. I need to look at the lit test on a bigger screen and will follow-up. Nothing obvious, but looking to see if it can be simplified.

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

704

Thanks - I've open-coded this a few times.

mlir/include/mlir/IR/Types.h

97 ↗

(On Diff #444915)

I'm trying to spot the diff here. Is it just formatting fallout?

132 ↗

(On Diff #444915)

I'm not sure these helpers are in the "this was a great idea" bucket. I usually just say t.isa<IntegerType>. A few extra characters but less core coupling.

Minimally: this new function is now not consistent with the comment.

mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp

139

River would say: Remove trivial braces (here and below)

(if River is watching: see, I can be taught)

152

I'm trying to understand why this restriction needs to be this strict. Isn't the restriction actually that the peeled operation cannot use the out operand (and let us guarantee that we have not violated ordering semantics)?

I'm fine with the stricter form as a starting point but making sure I understand (and if we think it can be loosened, to note that in a comment).

164

Nit: "operation has less than 2 statements"

179

I would have expected that a BlockAndValueMapper would be needed to remap the operands? I'm a bit surprised this works without leaving the uses pointing at the original block (which will be deleted and should assert).

(edit: I see you being clever below. disregard)

241

Stale comment?

[mravishankar]

Address comments and remove all changes to Types.h/Types.cpp

[mravishankar]

Actually drop the changes to Types.h

[mravishankar]

Addressed the comments, but they were not in the place they were initially intended to be, so hopefully I read the comments correctly.

mlir/include/mlir/IR/Types.h

97 ↗

(On Diff #444915)

Yeah I think so. Ill revert these changes.

132 ↗

(On Diff #444915)

Ah thats right. Ill clean up this and all related changes.

mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp

152

Left a comment saying it might be relaxed. It is definitely being conservative. Mostly the outs values are only read when the op has a reduction iterator type. If you have a case where this is done, worth looking into. IIUC we should be able to "just drop the condition" and still not violate any dependency.

164

its less than 3? It is meant to abort on only "yield" or "op + yield"

[mravishankar]

Rebase

[stellaraccident]

mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp

164

Oh right. Had my original impl in mind.

mlir/test/Dialect/Linalg/decompose-ops.mlir

2

I wish this test were able to write workout being so extensive. But this is also a tricky/detailed pattern and I can't see an immediate simplification to how you have it.

[mravishankar]

Rebase to retrigger pre-merge check

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

InterfaceMethod<
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         assert(opOperand->getOwner() == this->getOperation());
         return getBlock()->getArgument(opOperand->getOperandNumber());
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
+        Return the operand for a `blockArgument`.
+      }],
+      /*retTy=*/"OpOperand *",
+      /*methodName=*/"getTiedOpOperand",
+      /*args=*/(ins "BlockArgument":$blockArgument),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        assert(blockArgument.getOwner() == getBlock());
+        return &this->getOperation()->getOpOperand(
+            blockArgument.getArgNumber());
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the input or output indexing map for `opOperand`.
       }],
       /*retTy=*/"AffineMap",
       /*methodName=*/"getTiedIndexingMap",
       /*args=*/(ins "OpOperand*":$opOperand),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         assert(opOperand->getOwner() == this->getOperation());

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 //===----------------------------------------------------------------------===//
 // Transformations exposed as function calls.
 //===----------------------------------------------------------------------===//
 using LinalgLoops = SmallVector<Operation *, 4>;
 
 void populatePadTensorTilingPatterns(RewritePatternSet &patterns,
                                      const LinalgTilingOptions &options);
 
+/// Populate patterns for splitting a `LinalgOp` with multiple statements within
+/// its payload into multiple `GenericOp` that have a single statement.
+void populateDecomposeLinalgOpsPattern(RewritePatternSet &patterns);
+
 /// Populate patterns for vectorizing low-D convolution ops. This is a step in
 /// progressive lowering for convolution ops, it assume high-D convolution ops
 /// were decomposed previously.
 void populateConvolutionVectorizationPatterns(RewritePatternSet &patterns,
                                               PatternBenefit benefit = 1);
 
 /// Populate patterns that convert `ElementwiseMappable` ops to linalg
 /// parallel loops.

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRLinalgTransforms
   BubbleUpExtractSlice.cpp
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   CodegenStrategy.cpp
   ConstantFold.cpp
+  DecomposeLinalgOps.cpp
   Detensorize.cpp
   DropUnitDims.cpp
   ElementwiseOpFusion.cpp
   ElementwiseToLinalg.cpp
   FusePadOpWithLinalgProducer.cpp
   Fusion.cpp
   FusionOnTensors.cpp
   Generalization.cpp

mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp

+//===- DecomposeLinalgOps.cpp - Pattern to break up Linalg ops ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+
+using namespace mlir;
+using namespace mlir::linalg;
+
+namespace {
+
+/// Pattern to decompose a GenericOp that has more than two statements
+/// into one GenericOp with the first statement (i.e. peeled operation), and
+/// a second GenericOp with the remaining statements (i.e. residual operations).
+
+/// - The result of the first GenericOp has the same shape as the iteration
+///   space of the GenericOp. The number of results is same as the number of
+///   results of the peeled operation. The element type of the results of the
+///   first operation is same as the type of the results of the peeled
+///   operation.
+/// - The residual operation gets new `ins` operands that are the result of the
+///   peeled operation. If any operand of the residual operation becomes dead it
+///   is expected to be dropped by further canonicalization.
+/// - If the result of the peeled operation was yielded by the original
+///   GenericOp the uses of the corresponding results will be replaced with the
+///   result of the first GenericOp created.
+struct DecomposeLinalgOp : public OpRewritePattern<GenericOp> {
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GenericOp genericOp,
+                                PatternRewriter &rewriter) const override;
+
+private:
+  /// Helper method to create a generic op for the peeled scalar operation. The
+  /// created op has an empty region.
+  GenericOp createPeeledGenericOp(GenericOp genericOp,
+                                  PatternRewriter &rewriter) const;
+
+  /// Helper method to create a generic op for the residual scalar operation.
+  /// The created op has the same region as the original op.
+  GenericOp createResidualGenericOp(GenericOp genericOp,
+                                    GenericOp peeledGenericOp,
+                                    PatternRewriter &rewriter) const;
+};
+} // namespace
+
+/// Helper method to compute the range of a generic op.
+static SmallVector<OpFoldResult> getGenericOpLoopRange(OpBuilder &b,
+                                                       GenericOp op) {
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPoint(op);
+  Location loc = op.getLoc();
+  auto allShapesSizes =
+      cast<LinalgOp>(op.getOperation()).createFlatListOfOperandDims(b, loc);
+  AffineMap map = op.getShapesToLoopsMap();
+  return getAsOpFoldResult(applyMapToValues(b, loc, map, allShapesSizes));
+}
+
+GenericOp
+DecomposeLinalgOp::createPeeledGenericOp(GenericOp genericOp,
+                                         PatternRewriter &rewriter) const {
+  Block *body = genericOp.getBody();
+  Operation *peeledScalarOperation = &(*body->begin());
+  SmallVector<AffineMap> peeledGenericOpIndexingMaps;
+  for (auto operand : genericOp.getInputOperands()) {
+    peeledGenericOpIndexingMaps.push_back(
+        genericOp.getTiedIndexingMap(operand));
+  }
+
+  /// Compute the loop ranges for operation. This is the shape of the result of
+  /// the generic op for the peeled operation.
+  Location loc = genericOp.getLoc();
+  SmallVector<OpFoldResult> resultShape =
+      getGenericOpLoopRange(rewriter, genericOp);
+  SmallVector<Value> initValues;
+  SmallVector<Type> resultTypes;
+  AffineMap identityMap = rewriter.getMultiDimIdentityMap(resultShape.size());
+  for (auto scalarResult : peeledScalarOperation->getResults()) {
+    Value initTensor = rewriter.create<linalg::InitTensorOp>(
+        loc, resultShape, scalarResult.getType());
+    initValues.push_back(initTensor);
+    resultTypes.push_back(initTensor.getType());
+    peeledGenericOpIndexingMaps.push_back(identityMap);
+  }
+
+  /// Create the peeled generic op with an empty body.
+  auto indexingMapAttr =
+      rewriter.getAffineMapArrayAttr(peeledGenericOpIndexingMaps);
+  return rewriter.create<GenericOp>(
+      loc, resultTypes, genericOp.inputs(), initValues, indexingMapAttr,
+      genericOp.iterator_types(), /*doc=*/nullptr, /*libraryCall=*/nullptr,
+      [](OpBuilder, Location, ValueRange) {});
+}
+
+GenericOp
+DecomposeLinalgOp::createResidualGenericOp(GenericOp genericOp,
+                                           GenericOp peeledGenericOp,
+                                           PatternRewriter &rewriter) const {
+  /// Append all results from the peeledGenericOps as `ins` operand for the
+  /// residual generic op.
+  SmallVector<Value> residualGenericOpOperands = llvm::to_vector(
+      llvm::map_range(genericOp.getInputOperands(),
+                      [](OpOperand *operand) { return operand->get(); }));
+  auto peeledGenericOpResults = llvm::to_vector(llvm::map_range(
+      peeledGenericOp.getResults(), [](OpResult r) -> Value { return r; }));
+  residualGenericOpOperands.append(peeledGenericOpResults);
+
+  /// Add identity maps for the newly added operands.
+  AffineMap identityMap =
+      rewriter.getMultiDimIdentityMap(genericOp.getNumLoops());
+  auto indexingMaps = llvm::to_vector(
+      llvm::map_range(genericOp.getInputOperands(), [&](OpOperand *operand) {
+        return genericOp.getTiedIndexingMap(operand);
+      }));
+  indexingMaps.resize(indexingMaps.size() + peeledGenericOpResults.size(),
+                      identityMap);
+  for (OpOperand *outOperand : genericOp.getOutputOperands())
+    indexingMaps.push_back(genericOp.getTiedIndexingMap(outOperand));
+
+  auto indexingMapAttr = rewriter.getAffineMapArrayAttr(indexingMaps);
+  return rewriter.create<GenericOp>(
+      genericOp->getLoc(), genericOp->getResultTypes(),
+      residualGenericOpOperands, genericOp.outputs(), indexingMapAttr,
+      genericOp.iterator_types(), /*doc=*/nullptr, /*libraryCall=*/nullptr,
+      [](OpBuilder, Location, ValueRange) {});
+}
+
+LogicalResult
+DecomposeLinalgOp::matchAndRewrite(GenericOp genericOp,
+                                   PatternRewriter &rewriter) const {
+  /// For now only match on operations where the iterator types are all parallel
+  if (genericOp.getNumParallelLoops() != genericOp.getNumLoops()) {
+    return rewriter.notifyMatchFailure(genericOp,
+                                       "unhandled decomposition of operation "
+                                       "with non-parallel iterator types");
+  }
+  // TODO: this could be generalized to handle `linalg.generic` with buffer
+  // operands too but requires allocation for intermediates. Punt on this for
+  // now.
+  if (!genericOp.hasTensorSemantics()) {
+    return rewriter.notifyMatchFailure(
+        genericOp, "only operations with tensor semantics are handled");
+  }
+
+  if (llvm::any_of(genericOp.getOutputOperands(), [&](OpOperand *outOperand) {
+        return genericOp.payloadUsesValueFromOperand(outOperand);
+      })) {
+    return rewriter.notifyMatchFailure(
+        genericOp, "unhandled decomposition of generic op with use of out "
+                   "operand value in payload");
+  }
+
+  /// If the op has only a single statement (apart from the yield), do nothing.
+  Block *body = genericOp.getBody();
+  if (body->getOperations().size() <= 2) {
+    return rewriter.notifyMatchFailure(genericOp,
+                                       "body has less than 3 statements");
+  }
+
+  GenericOp peeledGenericOp = createPeeledGenericOp(genericOp, rewriter);
+  GenericOp residualGenericOp =
+      createResidualGenericOp(genericOp, peeledGenericOp, rewriter);
+
+  /// Move the first statement of the original operation into the body of the
+  /// generic op for the peeled operation.
+  Block *peeledGenericOpBody = peeledGenericOp.getBody();
+  Block *residualGenericOpBody = residualGenericOp.getBody();
+  assert(peeledGenericOpBody->empty() && residualGenericOpBody->empty() &&
+         "expected split generic ops to have empty region");
+  peeledGenericOpBody->getOperations().splice(
+      peeledGenericOpBody->begin(), body->getOperations(), body->begin());
+  residualGenericOpBody->getOperations().splice(residualGenericOpBody->begin(),
+                                                body->getOperations());
+
+  Operation *peeledScalarOperation = &(*peeledGenericOpBody->begin());
+  {
+    // Yield all the result of the peeled scalar operation.
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPointToEnd(peeledGenericOpBody);
+    rewriter.create<YieldOp>(genericOp.getLoc(),
+                             peeledScalarOperation->getResults());
+  }
+
+  /// In the split operations, replace block arguments uses that refer to
+  /// original operation to the block arguments of the newly created operation.
+  unsigned origNumInputs = genericOp.getNumInputs();
+  for (auto inputBlockArg :
+       llvm::enumerate(genericOp.getBody()->getArguments())) {
+    Value residualOpReplacementArg =
+        residualGenericOpBody->getArgument(inputBlockArg.index());
+    inputBlockArg.value().replaceUsesWithIf(
+        residualOpReplacementArg, [&](OpOperand &use) {
+          return use.getOwner()->getBlock() == residualGenericOpBody;
+        });
+
+    Value peeledOpReplacementArg =
+        peeledGenericOpBody->getArgument(inputBlockArg.index());
+    inputBlockArg.value().replaceUsesWithIf(
+        peeledOpReplacementArg, [&](OpOperand &use) {
+          return use.getOwner()->getBlock() == peeledGenericOpBody;
+        });
+  }
+
+  /// Before fixing up the residual operation, track what values are yielded. If
+  /// any of those are from the peeled scalar operation, the uses of the
+  /// corresponding result have to be remapped to result of the generic op for
+  /// the peeled operation.
+  auto yieldOp = residualGenericOpBody->getTerminator();
+  SmallVector<Value> replacements;
+  for (auto yieldValue : llvm::enumerate(yieldOp->getOperands())) {
+    OpResult opr = yieldValue.value().dyn_cast<OpResult>();
+    if (!opr || opr.getOwner() != peeledScalarOperation) {
+      replacements.push_back(residualGenericOp.getResult(yieldValue.index()));
+    } else {
+      replacements.push_back(peeledGenericOp->getResult(opr.getResultNumber()));
+    }
+  }
+
+  /// Update all uses of the peeled scalar operation results in the residual op
+  /// with the newly added arguments.
+  {
+    SmallVector<Value> scalarReplacements;
+    unsigned peeledGenericOpNumResults = peeledGenericOp->getNumResults();
+    scalarReplacements.reserve(peeledGenericOpNumResults);
+    for (auto num : llvm::seq<unsigned>(0, peeledGenericOpNumResults))
+      scalarReplacements.push_back(
+          residualGenericOpBody->getArgument(num + origNumInputs));
+    bool allUsesReplaced = false;
+    rewriter.replaceOpWithinBlock(peeledScalarOperation, scalarReplacements,
+                                  residualGenericOpBody, &allUsesReplaced);
+    assert(!allUsesReplaced &&
+           "peeled scalar operation is erased when it wasnt expected to be");
+
+    // Replace
+  }
+
+  // Replace the original operation
+  rewriter.replaceOp(genericOp, replacements);
+  return success();
+}
+
+void mlir::linalg::populateDecomposeLinalgOpsPattern(
+    RewritePatternSet &patterns) {
+  patterns.insert<DecomposeLinalgOp>(patterns.getContext());
+}

mlir/test/Dialect/Linalg/decompose-ops.mlir

+func.func @simple_op(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>)
+    -> (tensor<?x?xf32>, tensor<?x?xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+  %d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
+  %init1 = linalg.init_tensor [%d1, %d0] : tensor<?x?xf32>
+  %init2 = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32>
+  %result:2 = linalg.generic {
+    indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, 
+                     affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1, d0)>,
+                     affine_map<(d0, d1) -> (d0, d1)>],
+    iterator_types = ["parallel", "parallel"]}
+    ins(%arg0, %arg1, %arg2 : tensor<?x?xf32>, tensor<?xf32>, tensor<?xf32>)
+    outs(%init1, %init2 : tensor<?x?xf32>, tensor<?x?xf32>) {
+    ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32, %b4 : f32) :
+      %0 = arith.addf %b0, %b1 : f32
+      %1 = arith.mulf %0, %b2 : f32
+      linalg.yield %0, %1 : f32, f32
+    } -> (tensor<?x?xf32>, tensor<?x?xf32>)
+  return %result#0, %result#1 : tensor<?x?xf32>, tensor<?x?xf32>
+}

mlir/test/lib/Dialect/Linalg/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRLinalgTestPasses
   TestLinalgCodegenStrategy.cpp
+  TestLinalgDecomposeOps.cpp
   TestLinalgElementwiseFusion.cpp
   TestLinalgFusionTransforms.cpp
   TestLinalgHoisting.cpp
   TestLinalgTransforms.cpp
   TestPadFusion.cpp
 
   EXCLUDE_FROM_LIBMLIR
 

mlir/test/lib/Dialect/Linalg/TestLinalgDecomposeOps.cpp

+//===- TestLinalgDecomposeOps.cpp - Test Linalg decomposition  ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass for testing decomposition of Linalg ops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace {
+struct TestLinalgDecomposeOps
+    : public PassWrapper<TestLinalgDecomposeOps, OperationPass<>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgDecomposeOps)
+
+  TestLinalgDecomposeOps() = default;
+  TestLinalgDecomposeOps(const TestLinalgDecomposeOps &pass)
+      : PassWrapper(pass) {}
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AffineDialect, linalg::LinalgDialect>();
+  }
+  StringRef getArgument() const final {
+    return "test-linalg-decompose-ops-pattern";
+  }
+  StringRef getDescription() const final {
+    return "Test Linalg decomposition patterns";
+  }
+
+  void runOnOperation() override {
+    MLIRContext *context = &this->getContext();
+    RewritePatternSet decompositionPatterns(context);
+    linalg::populateDecomposeLinalgOpsPattern(decompositionPatterns);
+    if (failed(applyPatternsAndFoldGreedily(
+            getOperation(), std::move(decompositionPatterns)))) {
+      return signalPassFailure();
+    }
+  }
+};
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestLinalgDecomposeOps() {
+  PassRegistration<TestLinalgDecomposeOps>();
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestMultiBuffering();
 void registerTestIntRangeInference();
 void registerTestIRVisitorsPass();
 void registerTestGenericIRVisitorsPass();
 void registerTestGenericIRVisitorsInterruptPass();
 void registerTestInterfaces();
 void registerTestLastModifiedPass();
 void registerTestLinalgCodegenStrategy();
+void registerTestLinalgDecomposeOps();
 void registerTestLinalgElementwiseFusion();
 void registerTestLinalgFusionTransforms();
 void registerTestLinalgTensorFusionTransforms();
 void registerTestLinalgTiledLoopFusionTransforms();
 void registerTestLinalgGreedyFusion();
 void registerTestLinalgHoisting();
 void registerTestLinalgTileAndFuseSequencePass();
 void registerTestLinalgTransforms();
#endif
   mlir::test::registerTestComposeSubView();
   mlir::test::registerTestMultiBuffering();
   mlir::test::registerTestIntRangeInference();
   mlir::test::registerTestIRVisitorsPass();
   mlir::test::registerTestGenericIRVisitorsPass();
   mlir::test::registerTestInterfaces();
   mlir::test::registerTestLastModifiedPass();
   mlir::test::registerTestLinalgCodegenStrategy();
+  mlir::test::registerTestLinalgDecomposeOps();
   mlir::test::registerTestLinalgElementwiseFusion();
   mlir::test::registerTestLinalgFusionTransforms();
   mlir::test::registerTestLinalgTensorFusionTransforms();
   mlir::test::registerTestLinalgTiledLoopFusionTransforms();
   mlir::test::registerTestLinalgGreedyFusion();
   mlir::test::registerTestLinalgHoisting();
   mlir::test::registerTestLinalgTileAndFuseSequencePass();
   mlir::test::registerTestLinalgTransforms();