[mlir][tensor] Fix transpose permutation in tensor.pack generalization pattern

[qedawkins]

Summary

The generalization pattern for tensor.pack was inverting the
innerDimsPos permutation when normalizing. Thus, the transpose op
produced when generalizing wouldn't do the correct transpose. This can
be observed with the following example by comparing the IR generated
with and without data layout op (pack/unpack) propagation.

#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
builtin.module {
  func.func @forward(%arg0: tensor<3x5x7xf32>, %arg1: tensor<3x5x7xf32>) -> tensor<1x1x1x5x7x3xf32> {
    %empty = tensor.empty() : tensor<3x5x7xf32>
    %elementwise = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg0, %arg1 : tensor<3x5x7xf32>, tensor<3x5x7xf32>) outs(%empty : tensor<3x5x7xf32>) {
    ^bb0(%in: f32, %in_1: f32, %out: f32):
      %add = arith.addf %in, %in_1 : f32
      linalg.yield %add : f32
    } -> tensor<3x5x7xf32>
    %pack_empty = tensor.empty() : tensor<1x1x1x5x7x3xf32>
    %pack = tensor.pack %elementwise inner_dims_pos = [1, 2, 0] inner_tiles = [5, 7, 3] into %pack_empty : tensor<3x5x7xf32> -> tensor<1x1x1x5x7x3xf32>
    return %pack : tensor<1x1x1x5x7x3xf32>
  }
}

With the data layout propagation patterns through elementwise ops:

#map = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>
module {
  func.func @forward(%arg0: tensor<3x5x7xf32>, %arg1: tensor<3x5x7xf32>) -> tensor<1x1x1x5x7x3xf32> {
    %0 = tensor.empty() : tensor<1x1x1x5x7x3xf32>
    %1 = tensor.empty() : tensor<1x1x1x5x7x3xf32>
    %2 = tensor.empty() : tensor<7x3x5xf32>
    %transposed = linalg.transpose ins(%arg0 : tensor<3x5x7xf32>) outs(%2 : tensor<7x3x5xf32>) permutation = [2, 0, 1]
    %inserted_slice = tensor.insert_slice %transposed into %1[0, 0, 0, 0, 0, 0] [1, 1, 1, 7, 3, 5] [1, 1, 1, 1, 1, 1] : tensor<7x3x5xf32> into tensor<1x1x1x5x7x3xf32>
    %3 = tensor.empty() : tensor<1x1x1x5x7x3xf32>
    %4 = tensor.empty() : tensor<7x3x5xf32>
    %transposed_0 = linalg.transpose ins(%arg1 : tensor<3x5x7xf32>) outs(%4 : tensor<7x3x5xf32>) permutation = [2, 0, 1]
    %inserted_slice_1 = tensor.insert_slice %transposed_0 into %3[0, 0, 0, 0, 0, 0] [1, 1, 1, 7, 3, 5] [1, 1, 1, 1, 1, 1] : tensor<7x3x5xf32> into tensor<1x1x1x5x7x3xf32>
    %5 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%inserted_slice, %inserted_slice_1 : tensor<1x1x1x5x7x3xf32>, tensor<1x1x1x5x7x3xf32>) outs(%0 : tensor<1x1x1x5x7x3xf32>) {
    ^bb0(%in: f32, %in_2: f32, %out: f32):
      %6 = arith.addf %in, %in_2 : f32
      linalg.yield %6 : f32
    } -> tensor<1x1x1x5x7x3xf32>
    return %5 : tensor<1x1x1x5x7x3xf32>
  }
}

Without propagation:

#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
module {
  func.func @forward(%arg0: tensor<3x5x7xf32>, %arg1: tensor<3x5x7xf32>) -> tensor<1x1x1x5x7x3xf32> {
    %0 = tensor.empty() : tensor<3x5x7xf32>
    %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg0, %arg1 : tensor<3x5x7xf32>, tensor<3x5x7xf32>) outs(%0 : tensor<3x5x7xf32>) {
    ^bb0(%in: f32, %in_0: f32, %out: f32):
      %4 = arith.addf %in, %in_0 : f32
      linalg.yield %4 : f32
    } -> tensor<3x5x7xf32>
    %2 = tensor.empty() : tensor<1x1x1x5x7x3xf32>
    %3 = tensor.empty() : tensor<7x3x5xf32>
    %transposed = linalg.transpose ins(%1 : tensor<3x5x7xf32>) outs(%3 : tensor<7x3x5xf32>) permutation = [2, 0, 1]
    %inserted_slice = tensor.insert_slice %transposed into %2[0, 0, 0, 0, 0, 0] [1, 1, 1, 7, 3, 5] [1, 1, 1, 1, 1, 1] : tensor<7x3x5xf32> into tensor<1x1x1x5x7x3xf32>
    return %inserted_slice : tensor<1x1x1x5x7x3xf32>
  }
}

Where data layout propagation is doing a different transpose than what
generalization comes up with.

[hanchung]

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

588–595

The implementation of getPackNormalizedInnerPerm is weird to me.. it's calling getPackUnpackNormalizedInnerPerm and generates the inverse permutation.

Why not just use the same method and call invertPermutationVector here?

[qedawkins]

Use invert helper instead of specialized function

[qedawkins]

lint

[qedawkins]

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

588–595

Thanks, this is much cleaner :)

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

getPackUnpackNormalizedInnerPerm(int rank, ArrayRef<int64_t> innerDimsPos) {
   SmallVector<int64_t> vec(rank, kNonTiledMarker);
   for (auto [index, value] : llvm::enumerate(innerDimsPos))
     vec[value] = index;
   SmallVector<int64_t> perm = llvm::to_vector(llvm::make_filter_range(
       vec, [&](int64_t v) { return v != kNonTiledMarker; }));
   return perm;
 }
 
+static SmallVector<int64_t>
+getPackNormalizedInnerPerm(int rank, ArrayRef<int64_t> innerDimsPos) {
+  auto perm = getPackUnpackNormalizedInnerPerm(rank, innerDimsPos);
+  SmallVector<int64_t> invPerm(perm.size());
+  for (auto [index, value] : llvm::enumerate(perm))
+    invPerm[value] = index;
+  return invPerm;
+}
+
 LogicalResult GeneralizeOuterUnitDimsPackOpPattern::matchAndRewrite(
     tensor::PackOp packOp, PatternRewriter &rewriter) const {
   // TODO: support the case that outer dimensions are not all 1s A
   // tensor.expand_shape will be generated in this case.
   int64_t srcRank = packOp.getSourceRank();
   if (llvm::any_of(packOp.getDestType().getShape().take_front(srcRank),
                    [](int64_t val) { return val != 1; })) {
     return rewriter.notifyMatchFailure(
LogicalResult GeneralizeOuterUnitDimsPackOpPattern::matchAndRewrite(
   Type elemType = packOp.getSourceType().getElementType();
   auto readType = RankedTensorType::get(readShape, elemType);
 
   Value input = getPackOpSourceOrPaddedSource(rewriter, packOp);
   Value tile = rewriter.create<tensor::ExtractSliceOp>(
       loc, readType, input, readOffsets, readSizes, readStrides);
 
   // 2. Transpose the tile to match the inner tile order.
   SmallVector<int64_t> perm =
-      getPackUnpackNormalizedInnerPerm(srcRank, packOp.getInnerDimsPos());
+      getPackNormalizedInnerPerm(srcRank, packOp.getInnerDimsPos());
+
+  LLVM_DEBUG(DBGS() << "Pack permutation: " << packOp << "\n";
+             llvm::interleaveComma(perm, DBGS() << "perm: "); DBGSNL(););
+
   SmallVector<int64_t> transpShape = readShape;
   applyPermutationToVector<int64_t>(transpShape, perm);
 
   Value empty = rewriter.create<tensor::EmptyOp>(loc, transpShape, elemType);
   auto transposedOp =
       rewriter.create<linalg::TransposeOp>(loc, tile, empty, perm);
 
   // 3. Insert the inner tile to the destination.
   int64_t destRank = packOp.getDestRank();
   SmallVector<OpFoldResult> writeStrides(destRank, oneIdxAttr);

mlir/test/Dialect/Linalg/generalize-tensor-pack.mlir

 // CHECK:         %[[EMPTY:.+]] = tensor.empty() : tensor<32x8xf32>
 // CHECK:         %[[TRANSP:.+]] =  linalg.transpose
 // CHECK-SAME:      ins(%[[SRC]] : tensor<32x8xf32>)
 // CHECK-SAME:      outs(%[[EMPTY]] : tensor<32x8xf32>)
 // CHECK-SAME:      permutation = [0, 1]
 // CHECK:         %[[INSERT:.+]] = tensor.insert_slice %[[TRANSP]] into %[[DEST]]
 // CHECK-SAME:      [0, 0, 0, 0] [1, 1, 32, 8] [1, 1, 1, 1]
 // CHECK:         return %[[INSERT]]
+
+// -----
+
+func.func @simple_CHW_to_CHWhwc(%arg0: tensor<3x5x7xf32>, %arg1: tensor<1x1x1x5x7x3xf32>) -> tensor<1x1x1x5x7x3xf32> {
+  %0 = tensor.pack %arg0 inner_dims_pos = [1, 2, 0] inner_tiles = [5, 7, 3] into %arg1 : tensor<3x5x7xf32> -> tensor<1x1x1x5x7x3xf32>
+  return %0 : tensor<1x1x1x5x7x3xf32>
+}
+// CHECK-LABEL: func.func @simple_CHW_to_CHWhwc
+// CHECK-SAME:    %[[SRC:[a-zA-Z0-9]+]]
+// CHECK-SAME:    %[[DEST:[a-zA-Z0-9]+]]
+// CHECK:         %[[EMPTY:.+]] = tensor.empty() : tensor<5x7x3xf32>
+// CHECK:         %[[TRANSP:.+]] =  linalg.transpose
+// CHECK-SAME:      ins(%[[SRC]] : tensor<3x5x7xf32>)
+// CHECK-SAME:      outs(%[[EMPTY]] : tensor<5x7x3xf32>)
+// CHECK-SAME:      permutation = [1, 2, 0]
+// CHECK:         %[[INSERT:.+]] = tensor.insert_slice %[[TRANSP]] into %[[DEST]]
+// CHECK-SAME:      [0, 0, 0, 0, 0, 0] [1, 1, 1, 5, 7, 3] [1, 1, 1, 1, 1, 1]
+// CHECK:         return %[[INSERT]]