[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/Tensor/Transforms/Transforms.h

 /// tensor.[extract_slice|cast|expand_shape|collapse_shape].
 void populateFoldTensorEmptyPatterns(RewritePatternSet &patterns);
 
 /// Populates `patterns` with patterns that fold operations like `tensor.pad`
 /// and `tensor.extract_slice` into `tensor.pack` and `tensor.unpack` operations
 /// respectively.
 void populateFoldIntoPackAndUnpackPatterns(RewritePatternSet &patterns);
 
+/// Populates `patterns` with patterns what convert non-destination-style ops
+/// to destination style ops.
+void populateConvertToDestinationStylePatterns(RewritePatternSet &patterns);
+
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMS_H

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

 add_mlir_dialect_library(MLIRTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
+  ConvertToDestinationStyle.cpp
   EmptyOpPatterns.cpp
   ExtractSliceFromReshapeUtils.cpp
   FoldIntoPackAndUnpackPatterns.cpp
   MergeConsecutiveInsertExtractSlicePatterns.cpp
   ReshapePatterns.cpp
   SplitPaddingPatterns.cpp
   SwapExtractSliceWithProducerPatterns.cpp
 

mlir/lib/Dialect/Tensor/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/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Transforms/Transforms.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 tensor::populateConvertToDestinationStylePatterns(
+    RewritePatternSet &patterns) {
+  patterns.insert<GenerateOpConverter>(patterns.getContext());
+}

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

+// RUN: mlir-opt -split-input-file -test-tensor-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/Tensor/TestTensorTransforms.cpp

StringRef getArgument() const final {
     return "test-tensor-transform-patterns";
   }
   StringRef getDescription() const final {
     return "Test Tensor transformation patterns by applying them greedily.";
   }
 
   void runOnOperation() override;
 
+  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)};
+
   Option<bool> testSplitPaddingPatterns{
       *this, "test-split-padding-patterns",
       llvm::cl::desc("Test patterns to split tensor.pad ops"),
       llvm::cl::init(false)};
 
   Option<bool> testFoldConstantExtractSlice{
       *this, "test-fold-constant-extract-slice",
       llvm::cl::desc("Test folding arith.constant and tensor.extract_slice"),
struct TestTensorTransforms
 
   Option<bool> testSimplifyPackPatterns{
       *this, "test-simplify-pack-patterns",
       llvm::cl::desc("Test patterns to simplify tensor.pack"),
       llvm::cl::init(false)};
 };
 } // namespace
 
+static void applyConvertToDestinationStylePatterns(Operation *rootOp) {
+  RewritePatternSet patterns(rootOp->getContext());
+  tensor::populateConvertToDestinationStylePatterns(patterns);
+  (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
+}
+
 static void applyReassociativeReshapeFoldingPatterns(Operation *rootOp) {
   RewritePatternSet patterns(rootOp->getContext());
   tensor::populateReassociativeReshapeFoldingPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
 
 static void applyEmptyOpFoldingPatterns(Operation *rootOp) {
   RewritePatternSet patterns(rootOp->getContext());
applyRewriteExtractFromCollapseShapePatterns(Operation *rootOp,
   else
     patterns.add<RewriteExtractSliceFromCollapseShapeUsingScfFor>(
         rootOp->getContext());
   return applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
 
 void TestTensorTransforms::runOnOperation() {
   Operation *rootOp = getOperation();
+  if (testConvertToDestinationStylePatterns)
+    applyConvertToDestinationStylePatterns(rootOp);
   if (testSimplifyPackPatterns)
     applySimplifyPackPatterns(rootOp);
   if (testSplitPaddingPatterns)
     applySplitPaddingPatterns(rootOp);
   if (testFoldConstantExtractSlice)
     applyFoldConstantExtractSlicePatterns(rootOp);
   if (testFoldConsecutiveInsertExtractSlice)
     applyFoldConsecutiveInsertExtractSlicePatterns(rootOp);