[MLIR][Linalg] Change insertion point for `bubbleUpPackOpThroughElemGenericOp`

[chelini]

Summary

Currently, the insertion point for bubbleUpPackOpThroughElemGenericOp
is after the tensor.pack this means that the new generic will be created
right after the tensor.pack. This is inconvenient because we are moving
the position of the generic; the idea is to move pack/unpack around, not
linalg.generics. This PR changes the insertion point to preserve the
position of the generic.

Additionally, it restricts the pattern to fire if the generic has a
single user (tensor.pack) to avoid introducing recomputation.

[hanchung]

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

318–326

style nit: we need braces if {} is used in else branch.

323–325

[optional] Is it a heavy check? If so, maybe we should have simpler check and bail-out. We can maybe introduce an aggressive mode of this is needed.

[chelini]

Braces

[chelini]

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

323–325

Difficult to say; it depends on the situation. I never hit this branch in my codebase as the destination of a tensor.pack (by construction) is always a tensor.empty. I added this branch to preserve the tests we already had, where the destination is a function argument. For this specific case, dominance is just a pointer comparison.

[rengolin]

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

323–325

If you don't exercise this code then perhaps you should craft a test showing it is valid, or emit an assert on !emptyOp.

[chelini]

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

323–325

Are you referring to the "else" branch? All the code is exercised, and most of the tests (but not all) we have in data-layout-propagtion.mlir takes the "else" branch. I also added another test, would_break_dominance, where the destination is not a function argument, so we should be covered.

[rengolin]

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

323–325

Perfect, thanks!

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

 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/IR/Dominance.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/Support/Debug.h"
 #include <optional>
 
 namespace mlir {
 #define GEN_PASS_DEF_LINALGDATALAYOUTPROPAGATION
 #include "mlir/Dialect/Linalg/Passes.h.inc"
 } // namespace mlir

 ///     ^bb0(%in: f32, %out: f32):
 ///       %4 = arith.addf %in, %in : f32
 ///       linalg.yield %4 : f32
 ///     } -> tensor<?x?x8x2xf32>
 static FailureOr<GenericOp>
 bubbleUpPackOpThroughElemGenericOp(RewriterBase &rewriter,
                                    tensor::PackOp packOp) {
   auto genericOp = packOp.getSource().getDefiningOp<GenericOp>();
-  if (!genericOp)
-    return failure();
-
-  if (!isElementwise(genericOp))
+  if (!genericOp || !isElementwise(genericOp))
     return failure();
 
   // TODO: Relax the restriction. We are able to bubble up the pack op through
   // multi-result generic op. It just needs more work.
   if (genericOp.getNumResults() != 1)
     return failure();
 
+  // Bail-out if the result of the generic has multiple uses, as bubbling up
+  // creates recomputation if the generic has multiple users.
+  if (!genericOp->getResult(0).hasOneUse())
+    return failure();
+
+  // We want to move the pack not the generic.
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(genericOp);
+
+  // We need to handle two cases:
+  // 1) The tensor.pack destination is a tensor.empty. If this is the case, we
+  // create a new tensor.empty to avoid breaking dominance, as we are moving the
+  // tensor.pack above the linalg.generic.
+  // 2) The destination is not a tensor.empty. In this case we can replace only
+  // if the destination of the tensor.pack dominates the linalg.generic.
+  Value packOpDest = packOp.getDest();
+  if (!packOpDest.hasOneUse())
+    return failure();
+  if (auto emptyOp = packOpDest.getDefiningOp<tensor::EmptyOp>())
+    packOpDest = rewriter.create<tensor::EmptyOp>(
+        genericOp->getLoc(), emptyOp.getMixedSizes(),
+        emptyOp.getType().getElementType());
+  else {
+    DominanceInfo dom(genericOp);
+    if (!dom.properlyDominates(packOpDest, genericOp))
+      return failure();
+  }
+
   // TODO: Add an option for allowing padding values. It could introduce
   // undefined behavior if we unconditionally propagate pack op through all
   // the ops. E.g., if the padding value is zero and there are division ops in
   // a generic op. Some values of padding area could be NaN (0/0).
   if (packOp.getPaddingValue())
     return failure();
 
   OpOperand *opOperand = genericOp.getDpsInitOperand(0);
   auto packInfo = getPackingInfoFromOperand(
       genericOp.getMatchingIndexingMap(opOperand), packOp);
 
   // Rebuild the indexing map for the corresponding init operand.
   auto [packedOutOperand, packedOutIndexingMap] =
       getOrCreatePackedViewOfOperand(rewriter, genericOp.getLoc(), packInfo,
                                      genericOp, opOperand);
 
   // We'll replace the init operand with the destination of pack op if the init
   // operand has not users in the body of the linalg.generic (pure elementwise).
   // If it has users we need to pack the init operand too and replace the init
   // with the packing result.
   Value dest = (genericOp.getRegionOutputArgs()[0].use_empty())
-                   ? packOp.getDest()
+                   ? packOpDest
                    : packedOutOperand;
 
   return packElementWiseOp(rewriter, genericOp, dest, packedOutIndexingMap,
                            packInfo);
 }
 
 /// Wrapper pattern that applies bubbleUpPackOpThroughElemGenericOp method.
 struct BubbleUpPackOpThroughElemGenericOpPattern

mlir/test/Dialect/Linalg/data-layout-propagation.mlir

 // CHECK: func.func @pad_along_unpacked_dim(
 // CHECK: %[[ARG0:.+]]: tensor<1x2x56x56x32xf32>)
 // CHECK: %[[CST:.+]] = arith.constant 0.000000e+00 : f32
 // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x56x56x64xf32>
 // CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]] 
 // CHECK-SAME:  outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] 
 // CHECK-SAME:  into %[[EMPTY]] : tensor<1x2x56x56x32xf32> -> tensor<1x56x56x64xf32>
 // CHECK: %[[PADDED:.+]] = tensor.pad %[[UNPACK]] low[0, 1, 1, 1] high[0, 1, 1, 1]
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+func.func @would_break_dominance(%arg0: tensor<128x256xi32>) -> tensor<4x16x16x32xi32>{
+  %init = tensor.empty() : tensor<128x256xi32>
+  %elem = linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel", "parallel"]}
+      ins(%arg0 : tensor<128x256xi32>)
+      outs(%init : tensor<128x256xi32>) {
+    ^bb0(%arg3: i32, %arg4: i32):
+      %4 = arith.addi %arg3, %arg3 : i32
+      linalg.yield %4 : i32
+  } -> tensor<128x256xi32>
+  %dest = bufferization.alloc_tensor() : tensor<4x16x16x32xi32>
+  %pack = tensor.pack %elem
+    inner_dims_pos = [1, 0]
+    inner_tiles = [16, 32]
+    into %dest : tensor<128x256xi32> -> tensor<4x16x16x32xi32>
+  return %pack : tensor<4x16x16x32xi32>
+}
+
+// CHECK: func.func @would_break_dominance(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<128x256xi32>)
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<128x256xi32>
+// CHECK-NEXT: %[[GEN:.+]] = linalg.generic
+// CHECK-SAME:  ins(%[[ARG0]]
+// CHECK-SAME:  outs(%[[EMPTY]]
+// CHECK: %[[ALLOC:.+]] = bufferization.alloc_tensor() : tensor<4x16x16x32xi32>
+// CHECK-NEXT: %{{.+}} = tensor.pack %[[GEN]]
+// CHECK-SAME:  inner_dims_pos = [1, 0] inner_tiles = [16, 32] 
+// CHECK-SAME:  into %[[ALLOC]]