[mlir][Linalg] Rewrite DownscaleSizeOneWindowed2DConvolution to use rank-reducing insert/extract slices.

[nicolasvasilache]

Summary

This rewriting enables better bufferization and canonicalizations.

[gysit]

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

96

nit: tensor to dest?

mlir/include/mlir/IR/BuiltinTypes.h

242

Is it OK to let this->shape reference newShape here? I would expect the lifetime of newShape ends immediately after exiting dropDim? Or is it prolonged to the lifetime of the parent statement?

[nicolasvasilache]

mlir/include/mlir/IR/BuiltinTypes.h

242

A new RankedTensorType is returned which has its own allocated shape for this.
I'll add a comment to avoid accidental return of Builder & which would be wrong indeed.

[nicolasvasilache]

Address.

[gysit]

mlir/include/mlir/IR/BuiltinTypes.h

242

Ah that makes sense. I missed the thing is casted!

[bondhugula]

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

96

to dest -> to that of dest.

96

It isn't clear from the comment here why a Value has been passed instead of a target type.

mlir/include/mlir/IR/BuiltinTypes.h

206–208

Please add in sorted order.

211

out-live -> outlive

mlir/include/mlir/IR/BuiltinTypes.td

706

Likewise.

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

 /// ```mlir
 ///   %2 = consumer %0 ... : tensor<8x16xf32> ...
 /// ```
 bool canFoldIntoConsumerOp(CastOp castOp);
 
 /// Performs folding of any operand of `op` if it comes from a tensor::CastOp
 /// that can be folded.
 LogicalResult foldTensorCast(Operation *op);
+
+/// Create a rank-reducing ExtractSliceOp @[0 .. 0] with strides [1 .. 1] and
+/// appropriate sizes to reduce the rank of `tensor` to `targetType`.
+Value createCanonicalRankReducingExtractSliceOp(OpBuilder &b, Location loc,
+                                                Value tensor,
+                                                RankedTensorType targetType);
+
+/// Create a rank-reducing InsertSliceOp @[0 .. 0] with strides [1 .. 1] and
+/// appropriate sizes to increase the rank of `tensor` to `dest`.
+Value createCanonicalRankReducingInsertSliceOp(OpBuilder &b, Location loc,
+                                               Value tensor, Value dest);
+
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_IR_TENSOR_H_

mlir/include/mlir/IR/BuiltinTypes.h

 
 //===----------------------------------------------------------------------===//
 // Tablegen Interface Declarations
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/BuiltinTypeInterfaces.h.inc"
 
 namespace mlir {
+
+//===----------------------------------------------------------------------===//
+// RankedTensorType
+//===----------------------------------------------------------------------===//
+
+/// This is a builder type that keeps local references to arguments. Arguments
+/// that are passed into the builder must out-live the builder.
+class RankedTensorType::Builder {
+public:
+  /// Build from another RankedTensorType.
+  explicit Builder(RankedTensorType other)
+      : shape(other.getShape()), elementType(other.getElementType()),
+        encoding(other.getEncoding()) {}
+
+  /// Build from scratch.
+  Builder(ArrayRef<int64_t> shape, Type elementType, Attribute encoding)
+      : shape(shape), elementType(elementType), encoding(encoding) {}
+
+  Builder &setShape(ArrayRef<int64_t> newShape) {
+    shape = newShape;
+    return *this;
+  }
+
+  Builder &setElementType(Type newElementType) {
+    elementType = newElementType;
+    return *this;
+  }
+
+  Builder &setEncoding(Attribute newEncoding) {
+    encoding = newEncoding;
+    return *this;
+  }
+
+  /// Create a new RankedTensor by erasing a dim from shape.
+  // Note: the newly created type has ownership of a new shape vector.
+  RankedTensorType dropDim(unsigned dim) {
+    SmallVector<int64_t, 4> newShape(shape.begin(), shape.end());
+    newShape.erase(newShape.begin() + dim);
+    return setShape(newShape);
+  }
+
+  operator RankedTensorType() {
+    return RankedTensorType::get(shape, elementType, encoding);
+  }
+
+private:
+  ArrayRef<int64_t> shape;
+  Type elementType;
+  Attribute encoding;
+};
+
 //===----------------------------------------------------------------------===//
 // MemRefType
 //===----------------------------------------------------------------------===//
 
 /// This is a builder type that keeps local references to arguments. Arguments
 /// that are passed into the builder must out-live the builder.
 class MemRefType::Builder {
 public:

mlir/include/mlir/IR/BuiltinTypes.td

let builders = [
     TypeBuilderWithInferredContext<(ins
       "ArrayRef<int64_t>":$shape,
       "Type":$elementType,
       CArg<"Attribute", "{}">:$encoding
     ), [{
       return $_get(elementType.getContext(), shape, elementType, encoding);
     }]>
   ];
+  let extraClassDeclaration = [{
+    /// This is a builder type that keeps local references to arguments.
+    /// Arguments that are passed into the builder must out-live the builder.
+    class Builder;
+  }];
   let skipDefaultBuilders = 1;
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TupleType
 //===----------------------------------------------------------------------===//
 

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

LogicalResult matchAndRewrite(linalg::Conv2DNhwcHwcfOp convOp,
     Value input = convOp.inputs().front();
     Value filter = convOp.inputs().back();
     Value output = convOp.outputs().front();
 
     auto inputType = input.getType().dyn_cast<RankedTensorType>();
     auto filterType = filter.getType().dyn_cast<RankedTensorType>();
     auto outputType = output.getType().dyn_cast<RankedTensorType>();
 
-    auto inputShape = inputType.getShape();
     auto filterShape = filterType.getShape();
     auto outputShape = outputType.getShape();
 
     // Only handle the case where at least one of the window dimensions is
     // of size 1. Other cases can rely on tiling to reduce to such cases.
     int64_t fhSize = filterShape[0], fwSize = filterShape[1];
     int64_t ohSize = outputShape[1], owSize = outputShape[2];
-    if (!(fhSize == 1 && ohSize == 1) && !(fwSize == 1 && owSize == 1))
+    bool removeH = (fhSize == 1 && ohSize == 1);
+    bool removeW = (fwSize == 1 && owSize == 1);
+    if (!removeH && !removeW)
       return failure();
-    bool removeH = ohSize == 1;
 
     // Get new shapes and types for all operands by removing the size-1
     // dimension.
+    using RTTBuilder = RankedTensorType::Builder;
+    auto newInputType = RTTBuilder(inputType).dropDim((removeH ? 1 : 2));
+    auto newFilterType = RTTBuilder(filterType).dropDim((removeH ? 0 : 1));
+    auto newOutputType = RTTBuilder(outputType).dropDim(removeH ? 1 : 2);
 
-    SmallVector<int64_t, 3> newInputShape{
-        inputShape[0], inputShape[removeH ? 2 : 1], inputShape[3]};
-    auto newInputType = RankedTensorType::get(
-        newInputShape, inputType.getElementType(), inputType.getEncoding());
-
-    SmallVector<int64_t, 3> newFilterShape{filterShape[removeH ? 1 : 0],
-                                           filterShape[2], filterShape[3]};
-    auto newFilterType = RankedTensorType::get(
-        newFilterShape, filterType.getElementType(), filterType.getEncoding());
-
-    SmallVector<int64_t, 3> newOutputShape{
-        outputShape[0], outputShape[removeH ? 2 : 1], outputShape[3]};
-    auto newOutputType = RankedTensorType::get(
-        newOutputShape, outputType.getElementType(), outputType.getEncoding());
-
-    SmallVector<ReassociationIndices, 3> ioReshapeIndices = {{0}, {1, 2}, {3}};
-    SmallVector<ReassociationIndices, 3> fReshapeIndices = {{0, 1}, {2}, {3}};
-
-    // Reshape all operands for 1-D convolution.
+    // Rank-reduce operands.
     Location loc = convOp.getLoc();
-    Value newInput = rewriter.create<linalg::TensorCollapseShapeOp>(
-        loc, newInputType, input, ioReshapeIndices);
-    Value newFilter = rewriter.create<linalg::TensorCollapseShapeOp>(
-        loc, newFilterType, filter, fReshapeIndices);
-    Value newOutput = rewriter.create<linalg::TensorCollapseShapeOp>(
-        loc, newOutputType, output, ioReshapeIndices);
-
-    // We need to shrink the strides and dilations too.
-    auto stride = convOp.strides().getValues<int64_t>()[removeH ? 1 : 0];
-    auto stridesAttr = rewriter.getI64VectorAttr(stride);
-    auto dilation = convOp.dilations().getValues<int64_t>()[removeH ? 1 : 0];
-    auto dilationsAttr = rewriter.getI64VectorAttr(dilation);
+    Value newInput = tensor::createCanonicalRankReducingExtractSliceOp(
+        rewriter, loc, input, newInputType);
+    Value newFilter = tensor::createCanonicalRankReducingExtractSliceOp(
+        rewriter, loc, filter, newFilterType);
+    Value newOutput = tensor::createCanonicalRankReducingExtractSliceOp(
+        rewriter, loc, output, newOutputType);
+
+    // Rank-reduce strides and dilations too.
+    // TODO: dropDim 1-liner helper.
+    auto strides = llvm::to_vector<4>(convOp.strides().getValues<int64_t>());
+    strides.erase(strides.begin() + (removeH ? 0 : 1));
+    auto stridesAttr = rewriter.getI64VectorAttr(strides);
+
+    auto dilations =
+        llvm::to_vector<4>(convOp.dilations().getValues<int64_t>());
+    dilations.erase(dilations.begin() + (removeH ? 0 : 1));
+    auto dilationsAttr = rewriter.getI64VectorAttr(dilations);
 
     auto conv1DOp = rewriter.create<linalg::Conv1DNwcWcfOp>(
         loc, newOutputType, ValueRange{newInput, newFilter},
         ValueRange{newOutput}, stridesAttr, dilationsAttr);
 
-    rewriter.replaceOpWithNewOp<linalg::TensorExpandShapeOp>(
-        convOp, outputType, conv1DOp.getResult(0), ioReshapeIndices);
+    // Insert back.
+    Value inserted = tensor::createCanonicalRankReducingInsertSliceOp(
+        rewriter, loc, conv1DOp.getResult(0), output);
+    rewriter.replaceOp(convOp, inserted);
+
     return success();
   };
 };
 
 } // namespace
 
 void linalg::populateDecomposeConvolutionPatterns(RewritePatternSet &patterns,
                                                   PatternBenefit benefit) {
   patterns.add<DownscaleSizeOneWindowed2DConvolution>(patterns.getContext(),
                                                       benefit);
 }

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

OpFoldResult ExtractSliceOp::fold(ArrayRef<Attribute>) {
   if (getSourceType() == getType() &&
       succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType())))
     return this->source();
   if (Value slice = foldExtractAfterInsertSlice(*this))
     return slice;
   return OpFoldResult();
 }
 
+Value mlir::tensor::createCanonicalRankReducingExtractSliceOp(
+    OpBuilder &b, Location loc, Value tensor, RankedTensorType targetType) {
+  auto rankedTensorType = tensor.getType().cast<RankedTensorType>();
+  unsigned rank = rankedTensorType.getRank();
+  auto shape = rankedTensorType.getShape();
+  SmallVector<OpFoldResult> offsets(rank, b.getIndexAttr(0));
+  SmallVector<OpFoldResult> sizes;
+  for (unsigned i = 0, e = rank; i < e; ++i) {
+    OpFoldResult dim;
+    if (rankedTensorType.isDynamicDim(i))
+      dim = b.createOrFold<tensor::DimOp>(
+          loc, tensor, b.create<arith::ConstantIndexOp>(loc, i));
+    else
+      dim = b.getIndexAttr(shape[i]);
+    sizes.push_back(dim);
+  }
+  SmallVector<OpFoldResult> strides(rank, b.getIndexAttr(1));
+  return b.createOrFold<tensor::ExtractSliceOp>(loc, targetType, tensor,
+                                                offsets, sizes, strides);
+}
+
 //===----------------------------------------------------------------------===//
 // InsertSliceOp
 //===----------------------------------------------------------------------===//
 
 // Build a InsertSliceOp with mixed static and dynamic entries.
 void InsertSliceOp::build(OpBuilder &b, OperationState &result, Value source,
                           Value dest, ArrayRef<OpFoldResult> offsets,
                           ArrayRef<OpFoldResult> sizes,

 } // namespace
 
 void InsertSliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
   results.add<InsertSliceOpConstantArgumentFolder, InsertSliceOpCastFolder,
               InsertSliceOpSourceCastInserter>(context);
 }
 
+Value mlir::tensor::createCanonicalRankReducingInsertSliceOp(OpBuilder &b,
+                                                             Location loc,
+                                                             Value tensor,
+                                                             Value dest) {
+  auto rankedTensorType = dest.getType().cast<RankedTensorType>();
+  unsigned rank = rankedTensorType.getRank();
+  auto shape = rankedTensorType.getShape();
+  SmallVector<OpFoldResult> offsets(rank, b.getIndexAttr(0));
+  SmallVector<OpFoldResult> sizes;
+  for (unsigned i = 0, e = rank; i < e; ++i) {
+    OpFoldResult dim;
+    if (rankedTensorType.isDynamicDim(i))
+      dim = b.createOrFold<tensor::DimOp>(
+          loc, dest, b.create<arith::ConstantIndexOp>(loc, i));
+    else
+      dim = b.getIndexAttr(shape[i]);
+    sizes.push_back(dim);
+  }
+  SmallVector<OpFoldResult> strides(rank, b.getIndexAttr(1));
+  return b.createOrFold<tensor::InsertSliceOp>(loc, tensor, dest, offsets,
+                                               sizes, strides);
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tensor/IR/TensorOps.cpp.inc"

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

 // RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-decompose-convolution-patterns %s | FileCheck %s
 
 // CHECK-LABEL: func @conv2d_nhwc_4x1x2x8_tensor
 //  CHECK-SAME: (%[[INPUT:.+]]: tensor<4x1x6x3xf32>, %[[FILTER:.+]]: tensor<1x2x3x8xf32>, %[[INIT:.+]]: tensor<4x1x2x8xf32>)
 func @conv2d_nhwc_4x1x2x8_tensor(%input: tensor<4x1x6x3xf32>, %filter: tensor<1x2x3x8xf32>, %init: tensor<4x1x2x8xf32>) -> tensor<4x1x2x8xf32> {
   %0 = linalg.conv_2d_nhwc_hwcf
     {dilations = dense<[2, 3]> : tensor<2xi64>, strides = dense<[3, 2]> : tensor<2xi64>}
     ins(%input, %filter : tensor<4x1x6x3xf32>, tensor<1x2x3x8xf32>)
     outs(%init : tensor<4x1x2x8xf32>) -> tensor<4x1x2x8xf32>
   return %0 : tensor<4x1x2x8xf32>
 }
 
-//               CHECK: %[[INPUT_1D:.+]] = linalg.tensor_collapse_shape %[[INPUT]]
-// CHECK-SAME{LITERAL}:   [[0], [1, 2], [3]] : tensor<4x1x6x3xf32> into tensor<4x6x3xf32>
-//               CHECK: %[[FILTER_1D:.+]] = linalg.tensor_collapse_shape %[[FILTER]]
-// CHECK-SAME{LITERAL}:   [[0, 1], [2], [3]] : tensor<1x2x3x8xf32> into tensor<2x3x8xf32>
-//               CHECK: %[[INIT_1D:.+]] = linalg.tensor_collapse_shape %[[INIT]]
-// CHECK-SAME{LITERAL}:   [[0], [1, 2], [3]] : tensor<4x1x2x8xf32> into tensor<4x2x8xf32>
+//               CHECK: %[[INPUT_1D:.+]] = tensor.extract_slice %[[INPUT]]
+// CHECK-SAME{LITERAL}:   [0, 0, 0, 0] [4, 1, 6, 3] [1, 1, 1, 1] : tensor<4x1x6x3xf32> to tensor<4x6x3xf32>
+//               CHECK: %[[FILTER_1D:.+]] = tensor.extract_slice %[[FILTER]]
+// CHECK-SAME{LITERAL}:   [0, 0, 0, 0] [1, 2, 3, 8] [1, 1, 1, 1] : tensor<1x2x3x8xf32> to tensor<2x3x8xf32>
+//               CHECK: %[[INIT_1D:.+]] = tensor.extract_slice %[[INIT]]
+// CHECK-SAME{LITERAL}:   [0, 0, 0, 0] [4, 1, 2, 8] [1, 1, 1, 1] : tensor<4x1x2x8xf32> to tensor<4x2x8xf32>
 //               CHECK: %[[CONV_1D:.+]] = linalg.conv_1d_nwc_wcf
 //          CHECK-SAME:     dilations = dense<3> : vector<1xi64>
 //          CHECK-SAME:     strides = dense<2> : vector<1xi64>
 //          CHECK-SAME:   ins(%[[INPUT_1D]], %[[FILTER_1D]] : tensor<4x6x3xf32>, tensor<2x3x8xf32>)
 //          CHECK-SAME:   outs(%[[INIT_1D]] : tensor<4x2x8xf32>)
-//               CHECK: %[[CONV_2D:.+]] = linalg.tensor_expand_shape %[[CONV_1D]]
-// CHECK-SAME{LITERAL}:   [[0], [1, 2], [3]] : tensor<4x2x8xf32> into tensor<4x1x2x8xf32>
+//               CHECK: %[[CONV_2D:.+]] = tensor.insert_slice %[[CONV_1D]] into %[[INIT]]
+// CHECK-SAME{LITERAL}:   [0, 0, 0, 0] [4, 1, 2, 8] [1, 1, 1, 1] : tensor<4x2x8xf32> into tensor<4x1x2x8xf32>
 //               CHECK: return %[[CONV_2D]]
-
 // -----
 
 // CHECK-LABEL: func @conv2d_nhwc_qxqx1xq_tensor
 //  CHECK-SAME: (%[[INPUT:.+]]: tensor<?x?x1x?xf32>, %[[FILTER:.+]]: tensor<?x1x?x?xf32>, %[[INIT:.+]]: tensor<?x?x1x?xf32>)
 func @conv2d_nhwc_qxqx1xq_tensor(%input: tensor<?x?x1x?xf32>, %filter: tensor<?x1x?x?xf32>, %init: tensor<?x?x1x?xf32>) -> tensor<?x?x1x?xf32> {
   %0 = linalg.conv_2d_nhwc_hwcf
     {dilations = dense<[2, 3]> : tensor<2xi64>, strides = dense<[3, 2]> : tensor<2xi64>}
     ins(%input, %filter : tensor<?x?x1x?xf32>, tensor<?x1x?x?xf32>)
     outs(%init : tensor<?x?x1x?xf32>) -> tensor<?x?x1x?xf32>
   return %0 : tensor<?x?x1x?xf32>
 }
 
-//               CHECK: %[[INPUT_1D:.+]] = linalg.tensor_collapse_shape %[[INPUT]]
-// CHECK-SAME{LITERAL}:   [[0], [1, 2], [3]] : tensor<?x?x1x?xf32> into tensor<?x?x?xf32>
-//               CHECK: %[[FILTER_1D:.+]] = linalg.tensor_collapse_shape %[[FILTER]]
-// CHECK-SAME{LITERAL}:   [[0, 1], [2], [3]] : tensor<?x1x?x?xf32> into tensor<?x?x?xf32>
-//               CHECK: %[[INIT_1D:.+]] = linalg.tensor_collapse_shape %[[INIT]]
-// CHECK-SAME{LITERAL}:   [[0], [1, 2], [3]] : tensor<?x?x1x?xf32> into tensor<?x?x?xf32>
+//               CHECK: %[[INPUT_1D:.+]] = tensor.extract_slice %[[INPUT]]
+//          CHECK-SAME:   [0, 0, 0, 0] [%{{.*}}, %{{.*}}, 1, %{{.*}}] [1, 1, 1, 1] :
+//          CHECK-SAME:     tensor<?x?x1x?xf32> to tensor<?x?x?xf32>
+//               CHECK: %[[FILTER_1D:.+]] = tensor.extract_slice %[[FILTER]]
+//          CHECK-SAME:   [0, 0, 0, 0] [%{{.*}}, 1, %{{.*}}, %{{.*}}] [1, 1, 1, 1] :
+//          CHECK-SAME:     tensor<?x1x?x?xf32> to tensor<?x?x?xf32>
+//               CHECK: %[[INIT_1D:.+]] = tensor.extract_slice %[[INIT]]
+//          CHECK-SAME:   [0, 0, 0, 0] [%{{.*}}, %{{.*}}, 1, %{{.*}}] [1, 1, 1, 1] :
+//          CHECK-SAME:     tensor<?x?x1x?xf32> to tensor<?x?x?xf32>
 //               CHECK: %[[CONV_1D:.+]] = linalg.conv_1d_nwc_wcf
 //          CHECK-SAME:     dilations = dense<2> : vector<1xi64>
 //          CHECK-SAME:     strides = dense<3> : vector<1xi64>
 //          CHECK-SAME:   ins(%[[INPUT_1D]], %[[FILTER_1D]] : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
 //          CHECK-SAME:   outs(%[[INIT_1D]] : tensor<?x?x?xf32>)
-//               CHECK: %[[CONV_2D:.+]] = linalg.tensor_expand_shape %[[CONV_1D]]
-// CHECK-SAME{LITERAL}:   [[0], [1, 2], [3]] : tensor<?x?x?xf32> into tensor<?x?x1x?xf32>
+//               CHECK: %[[CONV_2D:.+]] = tensor.insert_slice %[[CONV_1D]] into %[[INIT]]
+//          CHECK-SAME:   [0, 0, 0, 0] [%{{.*}}, %{{.*}}, 1, %{{.*}}] [1, 1, 1, 1] :
+//          CHECK-SAME:     tensor<?x?x?xf32> into tensor<?x?x1x?xf32>
 //               CHECK: return %[[CONV_2D]]
 
 // -----
 
 // Do not convert convolution ops whose window dimensions are not ones.
 
 // CHECK-LABEL: func @conv2d_nhwc_4x1x2x8_tensor
 func @conv2d_nhwc_4x1x2x8_tensor(%input: tensor<4x3x5x3xf32>, %filter: tensor<2x2x3x8xf32>, %init: tensor<4x1x2x8xf32>) -> tensor<4x1x2x8xf32> {
   // CHECK: linalg.conv_2d_nhwc_hwcf
   %0 = linalg.conv_2d_nhwc_hwcf
     {dilations = dense<[2, 3]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
     ins(%input, %filter : tensor<4x3x5x3xf32>, tensor<2x2x3x8xf32>)
     outs(%init : tensor<4x1x2x8xf32>) -> tensor<4x1x2x8xf32>
   return %0 : tensor<4x1x2x8xf32>
 }