[mlir][MemRef] Bail out for unsupported cases in FoldMemRefAliasOps pass

[hanchung]

Summary

The pass uses computeSuffixProduct method which only allows static
shapes. This revision adds an early-exit for dynamic cases to avoid
crash.

[mravishankar]

Accepting since it is needed to fix downstream failures, but a proper fix would be to handle dynamic cases as well

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp

65–71

Could you add some comments for the early exit? I think computeSuffixProduct can be made to handle dynamic shapes, but isn't part of this patch

[hanchung]

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp

65–71

The other version of computeSuffixProduct can handle dynamic cases using AffineExpr ops, but that needs more work..

[hanchung]

address comments and fix typos in the lit test.

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp

 /// could be folded into
 ///
 /// %2 = load %0[6 * i1 + i2, %i3] :
 ///          memref<12x42xf32>
 static LogicalResult
 resolveSourceIndicesExpandShape(Location loc, PatternRewriter &rewriter,
                                 memref::ExpandShapeOp expandShapeOp,
                                 ValueRange indices,
                                 SmallVectorImpl<Value> &sourceIndices) {
+  // The below implementation uses computeSuffixProduct method, which only
+  // allows int64_t values (i.e., static shape). Bail out if it has dynamic
+  // shapes.
+  if (!expandShapeOp.getResultType().hasStaticShape())
+    return failure();
+
   MLIRContext *ctx = rewriter.getContext();
   for (ArrayRef<int64_t> groups : expandShapeOp.getReassociationIndices()) {
     assert(!groups.empty() && "association indices groups cannot be empty");
     int64_t groupSize = groups.size();
 
     // Construct the expression for the index value w.r.t to expand shape op
     // source corresponding the indices wrt to expand shape op result.
     SmallVector<int64_t> sizes(groupSize);

mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir

func.func @fold_static_stride_subview_with_affine_load_store_expand_shape_3d(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4: index) -> f32 {
   return %1 : f32
 }
 // CHECK: %[[INDEX:.*]] = affine.apply #[[$MAP]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]]]
 // CHECK-NEXT: %[[RESULT:.*]] = affine.load %[[ARG0]][%[[INDEX]], %[[ARG4]]] : memref<12x32xf32>
 // CHECK-NEXT: return %[[RESULT]] : f32
 
 // -----
 
+// CHECK-LABEL: fold_dynamic_subview_with_memref_load_store_expand_shape
+func.func @fold_dynamic_subview_with_memref_load_store_expand_shape(%arg0 : memref<16x?xf32, strided<[16, 1]>>, %arg1 : index, %arg2 : index) -> f32 {
+  %c0 = arith.constant 0 : index
+  %expand_shape = memref.expand_shape %arg0 [[0, 1], [2, 3]] : memref<16x?xf32, strided<[16, 1]>> into memref<1x16x?x1xf32, strided<[256, 16, 1, 1]>>
+  %0 = memref.load %expand_shape[%c0, %arg1, %arg2, %c0] : memref<1x16x?x1xf32, strided<[256, 16, 1, 1]>>
+  return %0 : f32
+}
+// CHECK: %[[EXPAND_SHAPE:.+]] = memref.expand_shape {{.+}} : memref<16x?xf32, strided<[16, 1]>> into memref<1x16x?x1xf32, strided<[256, 16, 1, 1]>>
+// CHECK: %[[LOAD:.+]] = memref.load %[[EXPAND_SHAPE]]
+// CHECK: return %[[LOAD]]
+
+// -----
+
 // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0 * 1024 + d1)>
 // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0 + d1)>
 // CHECK-LABEL: fold_static_stride_subview_with_affine_load_store_expand_shape
 // CHECK-SAME: (%[[ARG0:.*]]: memref<1024x1024xf32>, %[[ARG1:.*]]: memref<1xf32>, %[[ARG2:.*]]: index)
 func.func @fold_static_stride_subview_with_affine_load_store_expand_shape(%arg0: memref<1024x1024xf32>, %arg1: memref<1xf32>, %arg2: index) -> f32 {
   %0 = memref.expand_shape %arg0 [[0, 1], [2, 3]] : memref<1024x1024xf32> into memref<1x1024x1024x1xf32>
   affine.for %arg3 = 0 to 1 {
     affine.for %arg4 = 0 to 1024 {