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

[llvmlistbot at llvm.org]

================
@@ -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()) {
----------------
MaheshRavishankar wrote:

I agree this check is needed, but it might be better to do this as a follow up. You actually dont need identity, you only need to check that the collapsed dimensions are contiguous. Since this isnt in the existing behavior, lets do this in a follow up.

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