[mlir][memref]: Collapse strided unit dim even if strides are dynamic

[amrami]

…e dynamic

[llvmbot]

@llvm/pr-subscribers-mlir-memref

@llvm/pr-subscribers-mlir

Author: Maya Amrami (amrami)

Changes

…e dynamic

---

Full diff: https://github.com/llvm/llvm-project/pull/157330.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp (+5-5)
• (modified) mlir/test/Dialect/MemRef/ops.mlir (+6-1)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index b59d73d1291c8..3bdeaea300659 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -2401,6 +2401,11 @@ computeCollapsedLayoutMap(MemRefType srcType,
     auto trailingReassocs = ArrayRef<int64_t>(reassoc).drop_front();
     auto stride = SaturatedInteger::wrap(resultStrides[resultStrideIndex--]);
     for (int64_t idx : llvm::reverse(trailingReassocs)) {
+      // Dimensions of size 1 should be skipped, because their strides are
+      // meaningless and could have any arbitrary value.
+      if (srcShape[idx - 1] == 1)
+        continue;
+
       stride = stride * SaturatedInteger::wrap(srcShape[idx]);
 
       // Both source and result stride must have the same static value. In that
@@ -2415,11 +2420,6 @@ computeCollapsedLayoutMap(MemRefType srcType,
       if (strict && (stride.saturated || srcStride.saturated))
         return failure();
 
-      // Dimensions of size 1 should be skipped, because their strides are
-      // meaningless and could have any arbitrary value.
-      if (srcShape[idx - 1] == 1)
-        continue;
-
       if (!stride.saturated && !srcStride.saturated && stride != srcStride)
         return failure();
     }
diff --git a/mlir/test/Dialect/MemRef/ops.mlir b/mlir/test/Dialect/MemRef/ops.mlir
index 6c2298a3f8acb..50683761db5bf 100644
--- a/mlir/test/Dialect/MemRef/ops.mlir
+++ b/mlir/test/Dialect/MemRef/ops.mlir
@@ -431,7 +431,8 @@ func.func @expand_collapse_shape_dynamic(%arg0: memref<?x?x?xf32>,
          %arg4: index,
          %arg5: index,
          %arg6: index,
-         %arg7: memref<4x?x4xf32>) {
+         %arg7: memref<4x?x4xf32>,
+         %arg8: memref<1x1x18x?xsi8, strided<[?, ?, ?, 1], offset: ?>>) {
 //       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
 //  CHECK-SAME:     memref<?x?x?xf32> into memref<?x?xf32>
   %0 = memref.collapse_shape %arg0 [[0, 1], [2]] :
@@ -480,6 +481,10 @@ func.func @expand_collapse_shape_dynamic(%arg0: memref<?x?x?xf32>,
 //       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]]
   %4 = memref.expand_shape %arg7 [[0, 1], [2], [3, 4]] output_shape [2, 2, %arg4, 2, 2]
         : memref<4x?x4xf32> into memref<2x2x?x2x2xf32>
+
+//       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2], [3]]
+//  CHECK-SAME:     memref<1x1x18x?xsi8, strided<[?, ?, ?, 1], offset: ?>> into memref<1x18x?xsi8, strided<[?, ?, 1], offset: ?>>
+  %5 = memref.collapse_shape %arg8 [[0, 1], [2], [3]] : memref<1x1x18x?xsi8, strided<[?, ?, ?, 1], offset: ?>> into memref<1x18x?xsi8, strided<[?, ?, 1], offset: ?>>
   return
 }
 

[amrami]

@matthias-springer Can you take a look? 😊

[matthias-springer]

I don't see a problem with this PR, but I'm not working with collapse_shape/expand_shape anymore. @MaheshRavishankar is this safe to merge?

[amrami]

@MaheshRavishankar @Groverkss Can you take a look? :)

[krzysz00]

LGTM here

[krzysz00]

Do we need something further before landing this?

[hanhanW]

Hi, I'd like to signal that this breaks IREE downstream lit tests.

With the change, the below memref is no longer collapsable, but it should be okay to collapse?

memref<3x3x1x96xf32, strided<[288, 96, 96, 1], offset: 864>>

The below method starts returning false with the change.

llvm-project/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

Lines 2597 to 2605 in 92a1eb3

	bool CollapseShapeOp::isGuaranteedCollapsible(
	MemRefType srcType, ArrayRef<ReassociationIndices> reassociation) {
	// MemRefs with identity layout are always collapsible.
	if (srcType.getLayout().isIdentity())
	return true;
	return succeeded(computeCollapsedLayoutMap(srcType, reassociation,
	/*strict=*/true));
	}

which does not compute stride properly in the change.

[hanhanW]

It fails in the below check. stride=1 mismatch srcStride=96. I think stride should be computed as 96 like the old behavior.

llvm-project/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

Lines 2590 to 2591 in 8f683c3

	if (!stride.saturated && !srcStride.saturated && stride != srcStride)
	return failure();

[hanhanW]

IMHO, the element type does not matter. We should probably just use f32 or i8 for consistency in this file.