[Mlir-commits] [mlir] [MLIR] Folding unpack and pack sequence in data layout propagation from padded domain (PR #138332)

[Han-Chung Wang]

================
@@ -298,20 +298,56 @@ getOrCreatePackedViewOfOperand(OpBuilder &b, Location loc, PackInfo packInfo,
   return std::make_tuple(packedOperand, indexingMap);
 }
 
+static bool isGenericOutsNotUsed(linalg::GenericOp genericOp) {
+  int numDpsOuts = genericOp.getNumDpsInits();
+  for (int i = 0; i < numDpsOuts; ++i) {
+    Block *block = genericOp.getBody();
+    int numBlockArgs = block->getNumArguments();
+    int matchingInitArgIndex = numBlockArgs - numDpsOuts + i;
+    return block->getArgument(matchingInitArgIndex).use_empty();
+  }
+  return true;
+}
+
 /// Pack a genericOp and return it.
 static GenericOp packGenericOp(RewriterBase &rewriter, GenericOp genericOp,
                                Value dest, AffineMap packedOutIndexingMap,
                                const PackInfo &packInfo) {
   Location loc = genericOp.getLoc();
   SmallVector<Value> inputOperands;
+  SmallVector<Value> inputOperandsFromUnpackedSource;
   SmallVector<AffineMap> indexingMaps;
+
+  // Note: canUnpackPackFold needs to also guarantee the generic body
+  // doesn't have gather semantics. Since such scenarios has been
+  // rejected by both BubbleUpPackOpThroughGenericOp and
+  // PushDownUnPackOpThroughGenericOp, we can safely assume
+  // canUnpackPackFold is as long as init is not used.
+  bool canUnpackPackFold = isGenericOutsNotUsed(genericOp);
----------------
hanhanW wrote:

> Thanks for the reminder. I'd adopt the && !hasGatherSemantics(genericOp) instead of relying on the comment to stay up to date. I was attempting to save an redundant check if possible :-p.

In this case, maybe we can put `hasGatherSemantics` in an assertion and document the requirement to the method. It is not trivial to support the op that has gather semantics, but I think it is doable. We need either an assertion or the actual check, so future contributors won't miss updating the method, IMO.

> I don't think isElementwise(genericOp) is necessary here. Per a past discussion with @Max191 offline, I think we are good as long as outs are unused. For example, the below IR isn't elementwise but we'd be okay to push unpack down or bubble the pack up.

After reviewing the `isElementwise` implementation and definition, I realized that I had wrong understanding about it. I thought that it requires `outs` is not used, but the implementation says no -- I can see the reason, but I'm still not fully convinced. Anyway, my initial idea is to only handle the cases you understand, and my assumption is that you only want to support elementwise operations when all the outs are not used. I'm being conservative here because people have different uses for linalg dialect. They could have a creative generic op that not uses outs but accidentally meets the requirement, and it would open up a can of worms. It prevents the divergence of the expectation of the pass between users and authors.

I don't follow the example, the computation body looks like an elementwise operation to me. Did you miss indexing maps or something else? My understanding is that it sums up the `in_0` and  `in_1` and yield the result? It is a generic op form of `arith.addf in0, in1 : tensor<...xf32>`, IIUC.

https://github.com/llvm/llvm-project/pull/138332