[mlir][linalg] Convert tensor.generate to destination style

[springerm]

Summary

This can be a pre-processing for bufferization and allows for more efficient lowerings without an alloc.

[mravishankar]

The pattern itself looks fine, but I think this has a dependency issue. There is already a Tensor -> Linalg dependence. This is adding a Linalg -> Tensor dependence.... That could be problematic... is there a way to do this as part of bufferization or something, or move these patterns to Linalg?

[springerm]

In D142205#4069807, @mravishankar wrote:

The pattern itself looks fine, but I think this has a dependency issue. There is already a Tensor -> Linalg dependence. This is adding a Linalg -> Tensor dependence.... That could be problematic... is there a way to do this as part of bufferization or something, or move these patterns to Linalg?

Note it is a dependency from TensorTransforms -> Linalg and we already have it without this change. The implementation is actually largely copied from Tensor/Transforms/BufferizableOpInterfaceImpl.cpp. I'm still on the fence on whether we should completely delete the tensor.generate bufferization and make ConvertToDestinationStyle a mandatory pre-processing to bufferize tensor.generate. It does solve a few issues such as "inferring the correct memory space". Whether the pattern lives in LinalgTransforms or TensorTransforms does not really matter in the end, but it conceptually fits into TensorTransforms as all bufferization-related transforms/interface impls are currently in the dialect of the respective ops.

In general, I found it better to do less work in bufferization and do more in pre/post processing passes. EmptyTensorElimination, BufferDeallocation and BufferHoisting are three examples. ConvertToDestinationPassingStyle is another one that is introduced with this revision. Otherwise, One-Shot Bufferize becomes really large and complicated.

[mravishankar]

In D142205#4069971, @springerm wrote:

In D142205#4069807, @mravishankar wrote:

The pattern itself looks fine, but I think this has a dependency issue. There is already a Tensor -> Linalg dependence. This is adding a Linalg -> Tensor dependence.... That could be problematic... is there a way to do this as part of bufferization or something, or move these patterns to Linalg?

Note it is a dependency from TensorTransforms -> Linalg and we already have it without this change. The implementation is actually largely copied from Tensor/Transforms/BufferizableOpInterfaceImpl.cpp.

This is adding a dependency from Linalg to Tensor.... Doesnt that lead to circular dependency. Would be good to see we dont break the layering check in bazel.... Maybe specific Tensor:IR -> Linalg:IR -> Tensor:Transforms -> Linalg:Transforms doesnt create a circular dependency, but thats strange layering. I am just saying maybe move this to Linalg since you are generating Linalg op.s

I'm still on the fence on whether we should completely delete the tensor.generate bufferization and make ConvertToDestinationStyle a mandatory pre-processing to bufferize tensor.generate. It does solve a few issues such as "inferring the correct memory space". Whether the pattern lives in LinalgTransforms or TensorTransforms does not really matter in the end, but it conceptually fits into TensorTransforms as all bufferization-related transforms/interface impls are currently in the dialect of the respective ops.

In general, I found it better to do less work in bufferization and do more in pre/post processing passes. EmptyTensorElimination, BufferDeallocation and BufferHoisting are three examples. ConvertToDestinationPassingStyle is another one that is introduced with this revision. Otherwise, One-Shot Bufferize becomes really large and complicated.

I totally agree with this. Doing pre-processing this way is definitely better. I am just thinking of dialect dependency.

[springerm]

address comments

[springerm]

In D142205#4070313, @mravishankar wrote:

This is adding a dependency from Linalg to Tensor.... Doesnt that lead to circular dependency.

It's was adding a dependency from TensorTransforms (not Tensor); now LinalgTransforms (not Linalg). So either one if fine.

[springerm]

Moved everything to Linalg.

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

 /// its payload into multiple `GenericOp` that have a single statement.
 /// The option `removeDeadArgsAndResults` adds patterns to remove dead arguments
 /// and results from the generated decomposed ops. This is default `true` since
 /// the core decomposition patterns relies on these clean up patterns. It is set
 /// to false only for testing purposes.
 void populateDecomposeLinalgOpsPattern(RewritePatternSet &patterns,
                                        bool removeDeadArgsAndResults = true);
 
+/// Populate patterns that convert non-destination-style ops to destination
+/// style ops.
+void populateConvertToDestinationStylePatterns(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
   ConstantFold.cpp
+  ConvertToDestinationStyle.cpp
   DataLayoutPropagation.cpp
   DecomposeLinalgOps.cpp
   Detensorize.cpp
   DropUnitDims.cpp
   ElementwiseOpFusion.cpp
   ElementwiseToLinalg.cpp
   EraseUnusedOperandsAndResults.cpp
   FusePadOpWithLinalgProducer.cpp

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

+//===- ConvertToDestinationStyle.cpp - Convert non-DPS to DPS 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains patterns to convert non-DPS ops to DPS ops. New
+// tensor.empty ops are inserted as a destination. Such tensor.empty can be
+// eliminated with "empty tensor elimination", allowing them to bufferize
+// without an allocation (assuming there are no further conflicts).
+//
+//===----------------------------------------------------------------------===//
+//
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/PatternMatch.h"
+#include "llvm/Support/Debug.h"
+
+using namespace mlir;
+using namespace mlir::tensor;
+
+namespace {
+
+/// Lower tensor.generate to linalg.generic.
+struct GenerateOpConverter : public OpRewritePattern<GenerateOp> {
+  using OpRewritePattern<GenerateOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GenerateOp generateOp,
+                                PatternRewriter &rewriter) const override {
+    // Only ops with exactly one block are supported.
+    if (!generateOp.getBody().hasOneBlock())
+      return failure();
+
+    Location loc = generateOp.getLoc();
+    RankedTensorType tensorType = generateOp.getType().cast<RankedTensorType>();
+
+    // Create tensor.empty.
+    auto emptyOp = rewriter.create<EmptyOp>(loc, tensorType,
+                                            generateOp.getDynamicExtents());
+
+    // Create linalg.generic.
+    SmallVector<utils::IteratorType> iteratorTypes(
+        tensorType.getRank(), utils::IteratorType::parallel);
+    SmallVector<AffineMap> indexingMaps(
+        1, rewriter.getMultiDimIdentityMap(tensorType.getRank()));
+    auto genericOp = rewriter.create<linalg::GenericOp>(
+        loc, tensorType, /*inputs=*/ValueRange(),
+        /*outputs=*/ValueRange{emptyOp.getResult()}, /*indexingMaps=*/
+        indexingMaps, iteratorTypes);
+    Block *body = rewriter.createBlock(&genericOp->getRegion(0), {},
+                                       tensorType.getElementType(), loc);
+    rewriter.setInsertionPointToStart(body);
+    SmallVector<Value> bbArgReplacements;
+    for (int64_t i = 0; i < tensorType.getRank(); ++i)
+      bbArgReplacements.push_back(rewriter.create<linalg::IndexOp>(loc, i));
+    rewriter.mergeBlocks(&generateOp.getBody().front(), body,
+                         bbArgReplacements);
+
+    // Update terminator.
+    auto yieldOp = cast<tensor::YieldOp>(body->getTerminator());
+    rewriter.replaceOpWithNewOp<linalg::YieldOp>(yieldOp, yieldOp.getValue());
+
+    // Replace tensor.generate.
+    rewriter.replaceOp(generateOp, genericOp->getResult(0));
+    return success();
+  }
+};
+
+} // namespace
+
+void linalg::populateConvertToDestinationStylePatterns(
+    RewritePatternSet &patterns) {
+  patterns.insert<GenerateOpConverter>(patterns.getContext());
+}

mlir/test/Dialect/Linalg/convert-to-destination-style.mlir

+// RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-convert-to-destination-style-patterns %s | FileCheck %s
+
+// CHECK: #[[$map:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func @tensor_generate(
+//  CHECK-SAME:     %[[s1:.*]]: index, %[[s2:.*]]: index
+//       CHECK:   %[[empty:.*]] = tensor.empty(%[[s1]], %[[s2]]) : tensor<?x?xindex>
+//       CHECK:   %[[generic:.*]] = linalg.generic
+//  CHECK-SAME:       {indexing_maps = [#[[$map]]], iterator_types = ["parallel", "parallel"]}
+//  CHECK-SAME:       outs(%[[empty]] : tensor<?x?xindex>) {
+//       CHECK:     %[[i0:.*]] = linalg.index 0
+//       CHECK:     %[[i1:.*]] = linalg.index 1
+//       CHECK:     %[[added:.*]] = arith.addi %[[i0]], %[[i1]]
+//       CHECK:     linalg.yield %[[added]]
+//       CHECK:   }
+//       CHECK:   return %[[generic]]
+func.func @tensor_generate(%s1: index, %s2: index) -> tensor<?x?xindex> {
+  %0 = tensor.generate %s1, %s2 {
+    ^bb0(%arg0: index, %arg1: index):
+    %1 = arith.addi %arg0, %arg1 : index
+    tensor.yield %1 : index
+  } : tensor<?x?xindex>
+  return %0 : tensor<?x?xindex>
+}

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

struct TestLinalgTransforms
   Option<bool> testEraseUnusedOperandsAndResults{
       *this, "test-erase-unused-operands-and-results",
       llvm::cl::desc("Test patterns to erase unused operands and results"),
       llvm::cl::init(false)};
   Option<bool> testEraseUnnecessaryInputs{
       *this, "test-erase-unnecessary-inputs",
       llvm::cl::desc("Test patterns to erase unnecessary inputs"),
       llvm::cl::init(false)};
+  Option<bool> testConvertToDestinationStylePatterns{
+      *this, "test-convert-to-destination-style-patterns",
+      llvm::cl::desc("Test patterns that convert ops to destination style"),
+      llvm::cl::init(false)};
 };
 } // namespace
 
 static void applyPatterns(func::FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   RewritePatternSet patterns(ctx);
 
   //===--------------------------------------------------------------------===//

 }
 
 static void applyEraseUnnecessaryInputs(func::FuncOp funcOp) {
   RewritePatternSet patterns(funcOp.getContext());
   populateEraseUnnecessaryInputsPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
+static void applyConvertToDestinationStylePatterns(Operation *rootOp) {
+  RewritePatternSet patterns(rootOp->getContext());
+  populateConvertToDestinationStylePatterns(patterns);
+  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
+}
+
 /// Apply transformations specified as patterns.
 void TestLinalgTransforms::runOnOperation() {
   if (testPatterns)
     return applyPatterns(getOperation());
   if (testVectorTransferForwardingPatterns)
     return applyVectorTransferForwardingPatterns(getOperation());
   if (testGenericToVectorPattern)
     return applyLinalgToVectorPatterns(getOperation());
void TestLinalgTransforms::runOnOperation() {
   if (testBubbleUpExtractSliceOpPattern)
     return applyBubbleUpExtractSliceOpPattern(getOperation());
   if (testSwapExtractSliceWithFill)
     return applySwapExtractSliceWithFillPattern(getOperation());
   if (testEraseUnusedOperandsAndResults)
     return applyEraseUnusedOperandsAndResultsPatterns(getOperation());
   if (testEraseUnnecessaryInputs)
     return applyEraseUnnecessaryInputs(getOperation());
+  if (testConvertToDestinationStylePatterns)
+    applyConvertToDestinationStylePatterns(getOperation());
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms>();
 }
 } // namespace test
 } // namespace mlir