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

[springerm]

Summary

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

Depends On: D142206

[mravishankar]

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

164

This logic is implemented on the pad op using the reifyResultShapes interface...Could we just reuse that?

187

I am not sure I want to go down this path. This linalg.generic isnt vectorizable AFAIK. We could use a slice and insert, etc. but tahts too cumbersome. For now maybe just handle the zero-padded case....

[springerm]

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

187

OK then I'm going to add a default linalg.fill lowering here, so that it can vectorize. In the unlikely case that it is not a constant, it will still lower to linalg,generic.

[springerm]

address comments

[springerm]

rebase

[mravishankar]

Unblocking, but have one comment remaining.

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

178 ↗

(On Diff #491296)

First, I am not sure we need to make this distinction.... constant yielded value is already outside the op.

Even if we need to, could this be made simpler by just having

if (matchPattern(yieldedValue, m_Constant(...)) {
  yieldedValue = ...
}
auto fillOp = rewriter.create<linalg::FillOp>(..., yieldeValue, ...);

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/Affine/IR/AffineOps.h"
 #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;
LogicalResult matchAndRewrite(GenerateOp generateOp,
     rewriter.replaceOpWithNewOp<linalg::YieldOp>(yieldOp, yieldOp.getValue());
 
     // Replace tensor.generate.
     rewriter.replaceOp(generateOp, genericOp->getResult(0));
     return success();
   }
 };
 
+/// Lower tensor.pad to linalg.generic + tensor.insert_slice.
+struct PadOpConverter : public OpRewritePattern<PadOp> {
+  using OpRewritePattern<PadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(PadOp padOp,
+                                PatternRewriter &rewriter) const override {
+    // Only ops with exactly one block are supported.
+    if (!padOp.getBodyRegion().hasOneBlock())
+      return failure();
+
+    Location loc = padOp.getLoc();
+    RankedTensorType resultType = padOp.getResultType();
+    RankedTensorType srcType = padOp.getSourceType();
+
+    auto toValue = [&](OpFoldResult ofr) {
+      if (ofr.is<Value>())
+        return ofr.get<Value>();
+      return rewriter
+          .create<arith::ConstantIndexOp>(loc, *getConstantIntValue(ofr))
+          .getResult();
+    };
+
+    // Compute dynamic result dimensions.
+    SmallVector<OpFoldResult> mixedLowPad = padOp.getMixedLowPad();
+    SmallVector<OpFoldResult> mixedHighPad = padOp.getMixedHighPad();
+    SmallVector<Value> dynamicSizes;
+    for (int64_t i = 0; i < resultType.getRank(); ++i) {
+      if (!resultType.isDynamicDim(i))
+        continue;
+      Value srcDim = rewriter.create<tensor::DimOp>(loc, padOp.getSource(), i);
+      Value lowPad = toValue(mixedLowPad[i]);
+      Value highPad = toValue(mixedHighPad[i]);
+      AffineExpr s0, s1, s2;
+      bindSymbols(padOp->getContext(), s0, s1, s2);
+      AffineExpr sumExpr = s0 + s1 + s2;
+      Value sum = rewriter.create<AffineApplyOp>(
+          loc, sumExpr, ValueRange{srcDim, lowPad, highPad});
+      dynamicSizes.push_back(sum);
+    }
+
+    // Create tensor.empty.
+    auto emptyOp = rewriter.create<EmptyOp>(loc, resultType, dynamicSizes);
+
+    // Create linalg.generic.
+    // TODO: Create linalg.fill in case the padded value is a constant.
+    SmallVector<utils::IteratorType> iteratorTypes(
+        resultType.getRank(), utils::IteratorType::parallel);
+    SmallVector<AffineMap> indexingMaps(
+        1, rewriter.getMultiDimIdentityMap(resultType.getRank()));
+    auto genericOp = rewriter.create<linalg::GenericOp>(
+        loc, resultType, /*inputs=*/ValueRange(),
+        /*outputs=*/ValueRange{emptyOp.getResult()}, /*indexingMaps=*/
+        indexingMaps, iteratorTypes);
+    Block *body = rewriter.createBlock(&genericOp->getRegion(0), {},
+                                       resultType.getElementType(), loc);
+    rewriter.setInsertionPointToStart(body);
+    SmallVector<Value> bbArgReplacements;
+    for (int64_t i = 0; i < resultType.getRank(); ++i)
+      bbArgReplacements.push_back(rewriter.create<linalg::IndexOp>(loc, i));
+    rewriter.mergeBlocks(padOp.getBody(), body, bbArgReplacements);
+
+    // Update terminator.
+    auto yieldOp = cast<tensor::YieldOp>(body->getTerminator());
+    rewriter.replaceOpWithNewOp<linalg::YieldOp>(yieldOp, yieldOp.getValue());
+    rewriter.setInsertionPointAfter(genericOp);
+
+    // Create tensor::InsertSliceOp.
+    SmallVector<OpFoldResult> sliceSizes =
+        getMixedSizes(rewriter, loc, padOp.getSource());
+    SmallVector<OpFoldResult> sliceStrides(srcType.getRank(),
+                                           rewriter.getIndexAttr(1));
+    rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
+        padOp, padOp.getSource(), genericOp->getResult(0),
+        /*offsets=*/padOp.getMixedLowPad(), sliceSizes, sliceStrides);
+
+    return success();
+  }
+};
+
 } // namespace
 
 void tensor::populateConvertToDestinationStylePatterns(
     RewritePatternSet &patterns) {
-  patterns.insert<FromElementsOpConverter, GenerateOpConverter>(
+  patterns.insert<FromElementsOpConverter, GenerateOpConverter, PadOpConverter>(
       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
+// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-convert-to-destination-style-patterns -canonicalize %s | FileCheck %s
 
 // CHECK-LABEL: func @tensor_from_elements_0d(
 //  CHECK-SAME:     %[[arg0:.*]]: index
 //       CHECK:   %[[empty:.*]] = tensor.empty() : tensor<index>
 //       CHECK:   %[[insert:.*]] = tensor.insert %[[arg0]] into %[[empty]][]
 //       CHECK:   return %[[insert]]
 func.func @tensor_from_elements_0d(%arg0: index) -> tensor<index> {
   %0 = tensor.from_elements %arg0 : tensor<index>

 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>
 }
+
+// -----
+
+// CHECK:       #[[$map:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 5)>
+// CHECK:       #[[$map1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 10)>
+// CHECK:       #[[$map2:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func @tensor_pad(
+//  CHECK-SAME:   %[[t1:.*]]: tensor<?x10xindex>, %[[l2:.*]]: index, %[[h1:.*]]: index, %[[h2:.*]]: index
+//   CHECK-DAG:   %[[c0:.*]] = arith.constant 0 : index
+//   CHECK-DAG:   %[[dim0:.*]] = tensor.dim %[[t1]], %[[c0]]
+//   CHECK-DAG:   %[[size0:.*]] = affine.apply #[[$map]]()[%[[dim0]], %[[h1]]]
+//   CHECK-DAG:   %[[size1:.*]] = affine.apply #[[$map1]]()[%[[l2]], %[[h2]]]
+//       CHECK:   %[[empty:.*]] = tensor.empty(%[[size0]], %[[size1]]) : tensor<?x?xindex>
+//       CHECK:   %[[generic:.*]] = linalg.generic
+//  CHECK-SAME:       {indexing_maps = [#[[$map2]]], iterator_types = ["parallel", "parallel"]}
+//  CHECK-SAME:       outs(%[[empty]] : tensor<?x?xindex>) {
+//       CHECK:     %[[i0:.*]] = linalg.index 0
+//       CHECK:     %[[i1:.*]] = linalg.index 1
+//       CHECK:     %[[mul:.*]] = arith.muli %[[i0]], %[[i1]]
+//       CHECK:     linalg.yield %[[mul]]
+//       CHECK:   }
+//   CHECK-DAG:   %[[dim0:.*]] = tensor.dim %[[t1]], %[[c0]]
+//       CHECK:   %[[inserted:.*]] = tensor.insert_slice %[[t1]] into %[[generic]][5, %[[l2]]] [%[[dim0]], 10] [1, 1] : tensor<?x10xindex> into tensor<?x?xindex>
+//       CHECK:   return %[[inserted]]
+func.func @tensor_pad(%t1: tensor<?x10xindex>, %l2: index, %h1: index,
+                      %h2: index) -> tensor<?x?xindex> {
+  %0 = tensor.pad %t1 low[5, %l2] high[%h1, %h2] {
+  ^bb0(%arg0: index, %arg1: index):
+    %m = arith.muli %arg0, %arg1 : index
+    tensor.yield %m : index
+  } : tensor<?x10xindex> to tensor<?x?xindex>
+  return %0 : tensor<?x?xindex>
+}