[MLIR][Tensor] Add canonicalization patterns for `tensor.pack`

[chelini]

Summary

• Fold an unpack(pack(x)) to x.

• Rewrite a tensor.pack to an tensor.expand_shape if only one dimension is packed.

[chelini]

Take into account padding semantics.

[chelini]

Add test for padding

[rengolin]

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

3402

Not clear why you outlined this code. Seems simple enough to be inside matchAndRewrite.

[chelini]

Rename

[mravishankar]

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

3386

Does this need to check for padding semantics? Maybe not, but I am not sure.

3399

This doesnt look like a canonicalization to me. I'd rather add a separate "populate*" method to pull this in.

[chelini]

Address comments:

• Add a populate* method to simplify tensor.pack.

[rengolin]

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

3386

I think we do. Both padding and inner tiles.

If the pack has padding, than the packed tensor may be "unpackable" unless we define the operations on the padding irrelevant. But if that's true, then we need to let unpack know that this is what we've done, so that it can match its output shape to be the same as the original (not sure how, since pack doesn't have an amount to pad, just a constant to use as padding).

If the unpack has different tiles as the pack, then the order of elements will be different and will create a different output tensor, so we can't collapse the pair.

The size of the output tensor has to validate from the shape of the type. In both pad and uneven tiling cases, we have to make sure that the pre-packed shape is identical to the post-unpacked shape. We can do that via use-def chains or by writing a strict shape policy and carry padding sizes around.

None of those things are simple, so I'd just avoid them for now. Just make sure there's no padding and the inner tiles are the same, too.

[hanchung]

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

3037–3038

I think using RankedTensorType is a bit more accurate.

3053

I looked into other file and think this should be // namespace.

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

294–296 ↗

(On Diff #487447)

llvm style nit: we don't use braces for single statements.

https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

[chelini]

• Bail out if tensor.pack has padding

• Check tiles for pack and unpack to be the same when canonicalize chain

• Use single-braces in TestTensorTransforms.cpp

• Fix namespace: end namespace -> namespace

• Add more tests with dynamic tiles/dims

• Use RankedTensorType instead if ShapedType in SimplifyPackToExandShape

• Improve code reuse between pack::canonicalize and unpack::canonicalize

[rengolin]

Thanks @chelini, LGTM now, but I'll let others make the final approval.

[hanchung]

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

3028–3030

[optional] The operandType is only used once; it does not improve readability. The method itself describes it. I think we can save one line here.

3384

My experience tells me that we seldom add braces to it (though I could not find style guide about it).

[tyb0807]

Thank you, LGTM

[chelini]

• Remove braces

• Address comment on operandType

[chelini]

Fix compile error

[mravishankar]

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

3386

Agreed. I didnt see that the padding semantics is checked in the pattern. If the pack has padding, then we cant simply fold the pack and unpack away. Cause pack is doing more than what unpack is doing. THats what my concern was. I dont fully follow everything stated above, but the code itself looks fine.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

let extraClassDeclaration = commonExtraClassDeclaration # [{
     static ShapedType inferPackedType(ShapedType sourceType,
         ArrayRef<int64_t> innerTileSizes, ArrayRef<int64_t> innerDimsPos,
         ArrayRef<int64_t> outerDimsPerm = {});
 
     static Value createDestinationTensor(OpBuilder &b, Location loc,
         Value source, ArrayRef<OpFoldResult> innerTileSizes,
         ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> outerDimsPerm);
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // UnPackOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_UnPackOp : Tensor_RelayoutOp<"unpack"> {
   let summary = "tensor unpack operation";

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

 template <typename OpTy>
 static LogicalResult
 reifyResultShapesImpl(OpTy op, OpBuilder &builder,
                       ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   static_assert(llvm::is_one_of<OpTy, PackOp, UnPackOp>::value,
                 "applies to only pack or unpack operations");
   int64_t destRank = op.getDestRank();
   reifiedReturnShapes.resize(1, SmallVector<Value>(destRank));
   for (auto dim : llvm::seq<int64_t>(0, destRank)) {
     reifiedReturnShapes[0][dim] =
         builder.createOrFold<tensor::DimOp>(op.getLoc(), op.getDest(), dim);
   }
   return success();
 }
 
 template <typename OpTy>
 static DenseMap<int64_t, OpFoldResult> getDimAndTileMappingImpl(OpTy op) {
   static_assert(llvm::is_one_of<OpTy, PackOp, UnPackOp>::value,
                 "applies to only pack or unpack operations");
   DenseMap<int64_t, OpFoldResult> dimAndTileMapping;
   ArrayRef<int64_t> dimsToTile = op.getInnerDimsPos();
   SmallVector<OpFoldResult> tiles = op.getMixedTiles();
   assert(tiles.size() == dimsToTile.size() &&
          "tiles must match indices of dimension to block");
   // bind the dimension `i` with the tile factor.
   for (auto i : llvm::seq<int64_t>(0, dimsToTile.size()))
     dimAndTileMapping[dimsToTile[i]] = tiles[i];
   return dimAndTileMapping;
 }
 
 template <typename OpTy>
 static SmallVector<OpFoldResult> getMixedTilesImpl(OpTy op) {
   static_assert(llvm::is_one_of<OpTy, PackOp, UnPackOp>::value,
                 "applies to only pack or unpack operations");
   Builder builder(op);
   SmallVector<OpFoldResult> mixedInnerTiles;
   unsigned dynamicValIndex = 0;
   for (int64_t staticTile : op.getStaticInnerTiles()) {
     if (!ShapedType::isDynamic(staticTile))
       mixedInnerTiles.push_back(builder.getI64IntegerAttr(staticTile));
     else
       mixedInnerTiles.push_back(op.getInnerTiles()[dynamicValIndex++]);
Speculation::Speculatability PackOp::getSpeculatability() {
   // sizes of the tiles do not divide perfectly the dimension; thus, check only
   // to have constant tiles and tiled inner dimensions.
   if (!areTilesAndTiledDimsAllConstant(*this))
     return Speculation::NotSpeculatable;
 
   return Speculation::Speculatable;
 }
 
+namespace {
+
+/// Fold an unpack(pack(x)) to x.
+struct PackOfUnPack : public OpRewritePattern<PackOp> {
+  using OpRewritePattern<PackOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(PackOp packOp,
+                                PatternRewriter &rewriter) const override {
+    UnPackOp unPackOp = packOp.getSource().getDefiningOp<UnPackOp>();
+    if (!unPackOp || unPackOp.getSourceType() != packOp.getDestType())
+      return failure();
+    if (packOp.getInnerDimsPos() != unPackOp.getInnerDimsPos())
+      return failure();
+    if (packOp.getOuterDimsPerm() != unPackOp.getOuterDimsPerm())
+      return failure();
+    rewriter.replaceOp(packOp, unPackOp.getSource());
+    return success();
+  }
+};
+
+/// Packing one-dimensional tensor can be expressed as an expand shape op.
+struct PackToExpandShape : public OpRewritePattern<PackOp> {
+  using OpRewritePattern<PackOp>::OpRewritePattern;
+
+  Value insertExpand(RewriterBase &rewriter, Location loc, Value operand,
+                     Type newOperandType, ArrayAttr reassociation) const {
+    Type operandType = operand.getType();
+    if (operandType == newOperandType)
+      return operand;
+    return rewriter.create<tensor::ExpandShapeOp>(loc, newOperandType, operand,
+                                                  reassociation);
+  }
+
+  LogicalResult matchAndRewrite(PackOp packOp,
+                                PatternRewriter &rewriter) const override {
+    ShapedType sourceType = packOp.getSourceType();
+    ShapedType destType = packOp.getDestType();
+    if (sourceType.getRank() != 1 || packOp.getPaddingValue())
+      return failure();
+    auto reassociation =
+        getReassociationIndicesForReshape(sourceType, destType);
+    if (!reassociation)
+      return failure();
+    Value expanded = insertExpand(
+        rewriter, packOp.getLoc(), packOp.getSource(), destType,
+        getReassociationIndicesAttribute(rewriter, *reassociation));
+    rewriter.replaceOp(packOp, expanded);
+    return success();
+  }
+};
+
+} // end namespace
+
+void PackOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                         MLIRContext *ctx) {
+  results.add<PackOfUnPack, PackToExpandShape>(ctx);
+}
+
 //===----------------------------------------------------------------------===//
 // UnPackOp
 //===----------------------------------------------------------------------===//
 
 void UnPackOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "unpack");
 }

mlir/test/Dialect/Tensor/canonicalize.mlir

 func.func @unpack_pack(%t: tensor<128x128xf32>, %tile1: index, %tile2: index) -> tensor<128x128xf32> {
   %tensor_empty = tensor.empty(%tile1, %tile2) : tensor<16x16x?x?xf32>
   %packed = tensor.pack %t inner_dims_pos = [0, 1] inner_tiles = [%tile1, %tile2] into %tensor_empty : tensor<128x128xf32> -> tensor<16x16x?x?xf32>
   %tensor_empty1 = tensor.empty() : tensor<128x128xf32>
   %unpacked = tensor.unpack %packed inner_dims_pos = [0, 1] inner_tiles = [%tile1, %tile2] into %tensor_empty1 : tensor<16x16x?x?xf32> -> tensor
 <128x128xf32>
   return %unpacked : tensor<128x128xf32>
 }
+
+// -----
+
+// CHECK: func.func @single_dim_packing(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<256xf32>)
+// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1]] : tensor<256xf32> into tensor<8x32xf32>
+// CHECK: return %[[EXPANDED]] : tensor<8x32xf32>
+func.func @single_dim_packing(%arg0: tensor<256xf32>) -> tensor<8x32xf32> {
+  %empty = tensor.empty() : tensor<8x32xf32>
+  %0 = tensor.pack %arg0 inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<256xf32> -> tensor<8x32xf32>
+  return %0 : tensor<8x32xf32>
+}
+
+// -----
+
+// CHECK: func.func @single_dim_packing_with_padding(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<255xf32>)
+// CHECK-NOT: tensor.expand_shape
+// CHECK: tensor.pack
+func.func @single_dim_packing_with_padding(%arg0: tensor<255xf32>) -> tensor<8x32xf32> {
+  %empty = tensor.empty() : tensor<8x32xf32>
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.pack %arg0 padding_value(%cst : f32) inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<255xf32> -> tensor<8x32xf32>
+  return %0 : tensor<8x32xf32>
+}
+
+// -----
+
+// Chain NCnc -> NC -> NC -> NCnc
+// CHECK: func.func @pack_unpack(
+// CHECK-SAME: %[[T:.+]]: tensor<16x16x?x?xf32>,
+// CHECK: return %[[T]] : tensor<16x16x?x?xf32>
+func.func @pack_unpack(%t: tensor<16x16x?x?xf32>, %tile1: index, %tile2: index) -> tensor<16x16x?x?xf32> {
+  %tensor_empty = tensor.empty() : tensor<128x128xf32>
+  %unpacked = tensor.unpack %t inner_dims_pos = [0, 1] inner_tiles = [%tile1, %tile2] into %tensor_empty : tensor<16x16x?x?xf32> -> tensor<128x128xf32>
+  %tensor_empty1 = tensor.empty(%tile1, %tile2) : tensor<16x16x?x?xf32>
+  %packed = tensor.pack %unpacked inner_dims_pos = [0, 1] inner_tiles = [%tile1, %tile2] into %tensor_empty1 : tensor<128x128xf32> -> tensor<16x16x?x?xf32>
+  return %packed : tensor<16x16x?x?xf32>
+}
+
+// -----
+
+// Chain NCnc -> NC -> NC -> NCnc
+// CHECK: func.func @pack_unpack(
+// CHECK-SAME: %[[T:.+]]: tensor<16x16x8x8xf32>
+// CHECK: return %[[T]] : tensor<16x16x8x8xf32>
+func.func @pack_unpack(%t: tensor<16x16x8x8xf32>) -> tensor<16x16x8x8xf32> {
+  %tensor_empty = tensor.empty() : tensor<128x128xf32>
+  %unpacked = tensor.unpack %t inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %tensor_empty : tensor<16x16x8x8xf32> -> tensor<128x128xf32>
+  %tensor_empty1 = tensor.empty() : tensor<16x16x8x8xf32>
+  %packed = tensor.pack %unpacked inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %tensor_empty1 : tensor<128x128xf32> -> tensor<16x16x8x8xf32>
+  return %packed : tensor<16x16x8x8xf32>
+}