[mlir] Move linalg::PadTensorOp to tensor::PadOp.

RFC: https://llvm.discourse.group/t/rfc-move-linalg-padtensorop-to-tensor-padop/5785

Differential Revision: https://reviews.llvm.org/D117892

Author: pifon2a

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

@@ -18,7 +18,6 @@ include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
-include "mlir/Interfaces/TilingInterface.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 
 // Base class for Linalg dialect ops that do not correspond to library calls.
@@ -130,207 +129,6 @@ def Linalg_InitTensorOp : Linalg_Op<"init_tensor",
   let hasCanonicalizer = 1;
 }
 
-def Linalg_PadTensorOp : Linalg_Op<"pad_tensor",
-    [AttrSizedOperandSegments, NoSideEffect,
-     DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
-     DeclareOpInterfaceMethods<TilingInterface,
-         ["getDestinationOperands", "getLoopIteratorTypes", "getIterationDomain",
-          "getTiledImplementation"]>]> {
-  let summary = "tensor pad operation";
-  let description = [{
-    `linalg.pad_tensor` is an operation that pads the `source` tensor
-    with given `low` and `high` padding config.
-
-    The PadTensor operation supports the following arguments:
-
-    * source: the "base" tensor on which to pad.
-    * low: A list contains the padding along the start of each
-           dimension, i.e `low`.
-    * high: A list contains the padding along the end of each
-            dimension, i.e. `high`.
-    * nofold: indicates that the operation should not be folded when source and
-              result types are equal.
-
-    The result tensor dimensions are `low` + `dim` + `high` along that
-    dimension. The number of elements of `low` and `high` must match
-    the rank of the input tensor. They can be either a constant or a
-    dynamic value.
-
-    The region of the `pad_tensor` operation returns the value to use
-    for the padding. The arguments of the region represent the index
-    of the source being accessed. There should be as many arguments as
-    the rank of the `source` tensor. The value `yield`-ed by the
-    region is used as the value of the view at the given position.
-
-    If `nofold` is set, the padding operation will not be folded away even
-    if the source type and the padded type have the same static shape. This can
-    be used, e.g., for packing or promotion to faster memory.
-
-    Example 1:
-
-    ```mlir
-      %pad_value = ... : f32
-      %0 = linalg.pad_tensor %0 low[1, 2] high[2, 3] {
-      ^bb0(%arg0 : index, %arg1 : index):
-        linalg.yield %pad_value : f32
-      } : tensor<?x?xf32> to tensor<?x?xf32>
-    ```
-
-    Example 2:
-
-    ```mlir
-      %pad_value = ... : f32
-      %0 = linalg.pad_tensor %arg0 low[2, %arg1, 3, 3] high[3, 3, %arg1, 2] {
-      ^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index):
-          linalg.yield %pad_value : f32
-      } : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32>
-    ```
-
-    Example 3:
-
-    ```mlir
-      %pad_value = ... : f32
-      %0 = linalg.pad_tensor %arg0 low[0, 0] high[%ub0, %ub1] {
-      ^bb0(%arg1: index, %arg2: index):
-        linalg.yield %pad_value : f32
-      } : tensor<2x3xf32> to tensor<?x?xf32>
-    ```
-
-    Example 4:
-
-    ```mlir
-      // Force a padded value to be always exist with `nofold`.
-      %pad_value = ... : f32
-      %0 = linalg.pad_tensor %arg0 nofold low[0, 0] high[0, 0] {
-      ^bb0(%arg1: index, %arg2: index):
-        linalg.yield %pad_value : f32
-      } : tensor<2x3xf32> to tensor<2x3xf32>
-    ```
-  }];
-
-  let arguments = (ins
-    AnyTensor:$source,
-    Variadic<Index>:$low,
-    Variadic<Index>:$high,
-    I64ArrayAttr:$static_low,
-    I64ArrayAttr:$static_high,
-    UnitAttr:$nofold);
-
-  let regions = (region SizedRegion<1>:$region);
-
-  let results = (outs AnyTensor:$result);
-
-  // TODO: Remove custom<InferType> when AllTypesMatch supports opt. operands.
-  let assemblyFormat = [{
-    $source
-    (`nofold` $nofold^)?
-    `low` `` custom<OperandsOrIntegersSizesList>($low, $static_low)
-    `high` `` custom<OperandsOrIntegersSizesList>($high, $static_high)
-    $region attr-dict `:` type($source) `to` type($result)
-  }];
-
-  let extraClassDeclaration = [{
-    static StringRef getStaticLowAttrName() {
-      return "static_low";
-    }
-
-    static StringRef getStaticHighAttrName() {
-      return "static_high";
-    }
-
-    RankedTensorType getSourceType() {
-      return source().getType().cast<RankedTensorType>();
-    }
-    RankedTensorType getResultType() {
-      return getResult().getType().cast<RankedTensorType>();
-    }
-
-    // Infer the shape of the result tensor given the type of the source tensor
-    // and paddings. Known result dimensions that cannot necessarily be inferred
-    // from low/high padding sizes can be optionally specified. Those will be
-    // considered when computing the result type.
-    static RankedTensorType inferResultType(
-                                RankedTensorType sourceType,
-                                ArrayRef<int64_t> staticLow,
-                                ArrayRef<int64_t> staticHigh,
-                                ArrayRef<int64_t> resultShape = {});
-
-    // Return a PadTensorOp that pads `source` to `type` size where the static
-    // sizes are assumed to be greater than the dynamic sizes. The op performs
-    // "high" padding (i.e. it adds trailing padding values until the desired
-    // size is met).
-    static linalg::PadTensorOp createPadHighOp(
-        Type type, Value source, Value pad, bool nofold, Location loc,
-        OpBuilder & builder);
-
-    // Return a PadTensorOp that pads `source to `type` size with `pad` value.
-    // I.e., a block will be created and the `pad` value will be yielded
-    // directly. If the type passed is nullptr, it is inferred.
-    static linalg::PadTensorOp createPadScalarOp(
-        Type type, Value source, Value pad, ArrayRef<OpFoldResult> low,
-        ArrayRef<OpFoldResult> high, bool nofold, Location loc,
-        OpBuilder & builder);
-
-    // Return the pad value if it is a constant. Return null value otherwise.
-    Value getConstantPaddingValue();
-
-    // Return a vector of all the static or dynamic values (low/high padding) of
-    // the op.
-    inline SmallVector<OpFoldResult> getMixedPadImpl(ArrayAttr staticAttrs,
-                                                     ValueRange values) {
-      SmallVector<OpFoldResult> res;
-      unsigned numDynamic = 0;
-      unsigned count = staticAttrs.size();
-      for (unsigned idx = 0; idx < count; ++idx) {
-        if (ShapedType::isDynamic(staticAttrs[idx].cast<IntegerAttr>().getInt()))
-          res.push_back(values[numDynamic++]);
-        else
-          res.push_back(staticAttrs[idx]);
-      }
-      return res;
-    }
-    SmallVector<OpFoldResult> getMixedLowPad() {
-      return getMixedPadImpl(static_low(), low());
-    }
-    SmallVector<OpFoldResult> getMixedHighPad() {
-      return getMixedPadImpl(static_high(), high());
-    }
-    // Return true if low padding is guaranteed to be 0.
-    bool hasZeroLowPad() {
-      return llvm::all_of(getMixedLowPad(), [](OpFoldResult ofr) {
-        return getConstantIntValue(ofr) == static_cast<int64_t>(0);
-      });
-    }
-    // Return true if high padding is guaranteed to be 0.
-    bool hasZeroHighPad() {
-      return llvm::all_of(getMixedHighPad(), [](OpFoldResult ofr) {
-        return getConstantIntValue(ofr) == static_cast<int64_t>(0);
-      });
-    }
-  }];
-
-  let builders = [
-    // Build a PadTensorOp with mixed static and dynamic entries.
-    OpBuilder<(ins "Value":$source, "ArrayRef<int64_t>":$staticLow,
-      "ArrayRef<int64_t>":$staticHigh, "ValueRange":$low, "ValueRange":$high,
-      CArg<"bool", "false">:$nofold,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
-    // Build a PadTensorOp with all dynamic entries.
-    OpBuilder<(ins "Value":$source, "ValueRange":$low, "ValueRange":$high,
-      CArg<"bool", "false">:$nofold,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
-    // Build a PadTensorOp with mixed static and dynamic entries and custom
-    // result type. If the type passed is nullptr, it is inferred.
-    OpBuilder<(ins "Type":$resultType, "Value":$source,
-      "ArrayRef<OpFoldResult>":$low, "ArrayRef<OpFoldResult>":$high,
-      CArg<"bool", "false">:$nofold,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
-  ];
-
-  let hasCanonicalizer = 1;
-  let hasFolder = 1;
-}
-
 def Linalg_YieldOp : Linalg_Op<"yield", [NoSideEffect, ReturnLike, Terminator]>,
     Arguments<(ins Variadic<AnyType>:$values)> {
   let summary = "Linalg yield operation";

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

@@ -1,4 +1,4 @@
-//===- HoistPadding.h - Hoisting transformation for PadTensorOp -*- C++ -*-===//
+//===- HoistPadding.h - Hoisting for tensor::PadOp -*- C++ --------------*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -14,8 +14,11 @@
 namespace mlir {
 class Value;
 
+namespace tensor {
+class PadOp;
+} // namespace tensor
+
 namespace linalg {
-class PadTensorOp;
 
 /// Mechanically hoist padding operations on tensors by `numLoops` into a new,
 /// generally larger tensor. This achieves packing of multiple padding ops into
@@ -59,8 +62,8 @@ class PadTensorOp;
 ///      }
 ///    }
 /// ```
-FailureOr<Value> hoistPaddingOnTensors(PadTensorOp opToHoist, int numLoops,
-                                       PadTensorOp &hoistedOp);
+FailureOr<Value> hoistPaddingOnTensors(tensor::PadOp opToHoist, int numLoops,
+                                       tensor::PadOp &hoistedOp);
 
 } // namespace linalg
 } // namespace mlir

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

@@ -1132,18 +1132,18 @@ void populateLinalgDistributeTiledLoopPattern(
 // Op-specific patterns.
 //===----------------------------------------------------------------------===//
 
-/// PadTensorOp is not canonicalized away yet, so we provide a transformation to
-/// `linalg.generic`.
-struct PadTensorOpTransformationPattern : public OpRewritePattern<PadTensorOp> {
-  using OpRewritePattern<PadTensorOp>::OpRewritePattern;
+/// tensor::PadOp is not canonicalized away yet, so we provide a transformation
+/// to `linalg.generic`.
+struct PadOpTransformationPattern : public OpRewritePattern<tensor::PadOp> {
+  using OpRewritePattern<tensor::PadOp>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(PadTensorOp padOp,
+  LogicalResult matchAndRewrite(tensor::PadOp padOp,
                                 PatternRewriter &rewriter) const override;
 };
 
 /// Pad the operands of `opToPad` to a static bounding box. Use `paddingFunc`
 /// and `nofoldFunc` to set the padding value and the nofold attribute of the
-/// introduced PadTensorOps, respectively. Update `paddedOp` to the cloned
+/// introduced tensor::PadOps, respectively. Update `paddedOp` to the cloned
 /// statically shaped operation and return the extracted dynamically shaped
 /// results. If padding fails, return failure.
 FailureOr<SmallVector<Value>>
@@ -1153,23 +1153,23 @@ rewriteAsPaddedOp(OpBuilder &b, LinalgOp opToPad,
                   LinalgOp &paddedOp);
 
 using OptimizeCopyFn =
-    std::function<LogicalResult(PatternRewriter &, PadTensorOp, Value)>;
+    std::function<LogicalResult(PatternRewriter &, tensor::PadOp, Value)>;
 
-/// Rewrite a PadTensorOp into a sequence of InitTensorOp, FillOp and
+/// Rewrite a tensor::PadOp into a sequence of InitTensorOp, FillOp and
 /// InsertSliceOp. For now, only constant padding values are supported.
 /// `OptimizeCopyFn` can be used to customize copying step optimization.
-struct GeneralizePadTensorOpPattern : public OpRewritePattern<PadTensorOp> {
-  GeneralizePadTensorOpPattern(MLIRContext *context,
-                               OptimizeCopyFn optimizeCopyFn = nullptr,
-                               PatternBenefit benefit = 1)
-      : OpRewritePattern<PadTensorOp>(context, benefit),
+struct GeneralizePadOpPattern : public OpRewritePattern<tensor::PadOp> {
+  GeneralizePadOpPattern(MLIRContext *context,
+                         OptimizeCopyFn optimizeCopyFn = nullptr,
+                         PatternBenefit benefit = 1)
+      : OpRewritePattern<tensor::PadOp>(context, benefit),
         optimizeCopyFn(std::move(optimizeCopyFn)) {}
-  LogicalResult matchAndRewrite(PadTensorOp padOp,
+  LogicalResult matchAndRewrite(tensor::PadOp padOp,
                                 PatternRewriter &rewriter) const override;
 
 protected:
   OptimizeCopyFn optimizeCopyFn;
-  Value createFillOrGenerateOp(PatternRewriter &rewriter, PadTensorOp padOp,
+  Value createFillOrGenerateOp(PatternRewriter &rewriter, tensor::PadOp padOp,
                                Value dest,
                                const SmallVector<Value> &dynSizes) const;
 };
@@ -1179,9 +1179,9 @@ struct GeneralizePadTensorOpPattern : public OpRewritePattern<PadTensorOp> {
 /// are used to encode a certain ordering of pattern application. To avoid
 /// scattering magic constants throughout the code base, the patterns must be
 /// added with this function. `baseBenefit` can be used to offset the benefit
-/// of all PadTensorOp vectorization patterns by a certain value.
-void populatePadTensorOpVectorizationPatterns(RewritePatternSet &patterns,
-                                              PatternBenefit baseBenefit = 1);
+/// of all tensor::PadOp vectorization patterns by a certain value.
+void populatePadOpVectorizationPatterns(RewritePatternSet &patterns,
+                                        PatternBenefit baseBenefit = 1);
 
 /// Match and rewrite for the pattern:
 /// ```

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

@@ -107,12 +107,12 @@ tensor::ExtractSliceOp makeComposedExtractSliceOp(
     OpBuilder &b, Location loc, Value source, ArrayRef<OpFoldResult> offsets,
     ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides);
 
-/// Create a PadTensorOp that pads `source` to the size of the statically sized
-/// `type` whose static sizes are assumed to be greater than the dynamic
+/// Create a tensor::PadOp that pads `source` to the size of the statically
+/// sized `type` whose static sizes are assumed to be greater than the dynamic
 /// `source` size. The padding introduces trailing `pad` values until the target
 /// size is met. If `source` is defined by one or more LinalgOps that have been
 /// padded with the same value and sizes, return their padded result instead of
-/// creating a PadTensorOp.
+/// creating a tensor::PadOp.
 ///
 /// Example:
 /// ```

mlir/include/mlir/Dialect/Tensor/IR/Tensor.h

@@ -19,6 +19,7 @@
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
 //===----------------------------------------------------------------------===//