[Mlir-commits] [mlir] [mlir][Linalg] Allow more control in drop unit dims (PR #171796)

[Lukas Sommer]

================
@@ -264,23 +268,25 @@ expandValue(RewriterBase &rewriter, Location loc, Value result, Value origDest,
         loc, result, origDest, offsets, sizes, strides);
   }
 
-  assert(rankReductionStrategy ==
+  assert(control.rankReductionStrategy ==
              ControlDropUnitDims::RankReductionStrategy::ReassociativeReshape &&
          "unknown rank reduction strategy");
   return tensor::ExpandShapeOp::create(rewriter, loc, origResultType, result,
                                        reassociation)
       .getResult();
 }
 
-/// Collapse the given `value` so that the type matches the type of
-/// `origOutput`. The `reassociation` is used when `rankReductionStrategy` is
-/// set to `RankReductionStrategy::ReassociativeReshape`.
-static Value collapseValue(
-    RewriterBase &rewriter, Location loc, Value operand,
-    ArrayRef<int64_t> targetShape, ArrayRef<ReassociationIndices> reassociation,
-    ControlDropUnitDims::RankReductionStrategy rankReductionStrategy) {
+FailureOr<Value>
+linalg::collapseValue(RewriterBase &rewriter, Location loc, Value operand,
+                      ArrayRef<int64_t> targetShape,
+                      ArrayRef<ReassociationIndices> reassociation,
+                      const ControlDropUnitDims &control) {
   if (auto memrefType = dyn_cast<MemRefType>(operand.getType())) {
-    if (rankReductionStrategy ==
+    if (!memrefType.getLayout().isIdentity()) {
----------------
sommerlukas wrote:

The previous implementation also used a check for identity layout [here](https://github.com/llvm/llvm-project/blob/b880428bf0fca4eef7460c440df19412863001c2/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp#L462-L464) to treat `memref` with non-identify layout different from identity layout. That's where I took this check from, but I can remove it for now if you prefer.

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