[mlir][MemRef] Change the anchor point of a subview pattern

[qcolombet]

Summary

Essentially, this patches changes the anchor point of the extract_strided_metadata(subview) pattern from extract_strided_metadata to subview.

In details, this means that instead of replacing:

base, offset, sizes, strides = extract_strided_metadata(subview(src))

With

base, ... = extract_strided_metadata(src)
offset = <some math>
sizes = subSizes
strides = <some math>

We replace only the subview part and connect it back with a reinterpret_cast:

val = subview(src)

>

base, ... = extract_strided_metadata(src)
offset = <some math>
sizes = subSizes
strides = <some math>
val = reinterpret_cast base, offset, sizes, strides

Note: This patch is part of a series that replaces the lowering of view like operations in MemRefToLLVM with composable patterns.
Note2: This patch depends on other patches in that series but mostly only because of the output of the tests.

[qcolombet]

• Rebase the patch after the fix at https://reviews.llvm.org/D136107

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

 namespace memref {
 #define GEN_PASS_DEF_SIMPLIFYEXTRACTSTRIDEDMETADATA
 #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
 } // namespace memref
 } // namespace mlir
 using namespace mlir;
 
 namespace {
-/// Replace `baseBuffer, offset, sizes, strides =
-///              extract_strided_metadata(subview(memref, subOffset,
-///                                               subSizes, subStrides))`
+/// Replace `dst = subview(memref, subOffset, subSizes, subStrides))`
 /// With
 ///
 /// \verbatim
 /// baseBuffer, baseOffset, baseSizes, baseStrides =
 ///     extract_strided_metadata(memref)
 /// strides#i = baseStrides#i * subSizes#i
 /// offset = baseOffset + sum(subOffset#i * strides#i)
 /// sizes = subSizes
+/// dst = reinterpret_cast baseBuffer, offset, sizes, strides
 /// \endverbatim
 ///
 /// In other words, get rid of the subview in that expression and canonicalize
 /// on its effects on the offset, the sizes, and the strides using affine.apply.
-struct ExtractStridedMetadataOpSubviewFolder
-    : public OpRewritePattern<memref::ExtractStridedMetadataOp> {
+struct SubviewFolder : public OpRewritePattern<memref::SubViewOp> {
 public:
-  using OpRewritePattern<memref::ExtractStridedMetadataOp>::OpRewritePattern;
+  using OpRewritePattern<memref::SubViewOp>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(memref::ExtractStridedMetadataOp op,
+  LogicalResult matchAndRewrite(memref::SubViewOp subview,
                                 PatternRewriter &rewriter) const override {
-    auto subview = op.getSource().getDefiningOp<memref::SubViewOp>();
-    if (!subview)
-      return failure();
-
-    // Build a plain extract_strided_metadata(memref) from
-    // extract_strided_metadata(subview(memref)).
-    Location origLoc = op.getLoc();
+    // Build a plain extract_strided_metadata(memref) from subview(memref).
+    Location origLoc = subview.getLoc();
     Value source = subview.getSource();
     auto sourceType = source.getType().cast<MemRefType>();
     unsigned sourceRank = sourceType.getRank();
 
     auto newExtractStridedMetadata =
         rewriter.create<memref::ExtractStridedMetadataOp>(origLoc, source);
 
     SmallVector<int64_t> sourceStrides;
for (unsigned i = 0; i < sourceRank; ++i) {
       values[subStrideForDim] = subStrides[i];
       values[origStrideForDim] = origStride;
     }
 
     // Compute the offset.
     OpFoldResult finalOffset =
         makeComposedFoldedAffineApply(rewriter, origLoc, expr, values);
 
-    SmallVector<Value> results;
     // The final result is  <baseBuffer, offset, sizes, strides>.
     // Thus we need 1 + 1 + subview.getRank() + subview.getRank(), to hold all
     // the values.
     auto subType = subview.getType().cast<MemRefType>();
     unsigned subRank = subType.getRank();
-    // Properly size the array so that we can do random insertions
-    // at the right indices.
-    // We do that to populate the non-dropped sizes and strides in one go.
-    results.resize_for_overwrite(subRank * 2 + 2);
-
-    results[0] = newExtractStridedMetadata.getBaseBuffer();
-    results[1] =
-        getValueOrCreateConstantIndexOp(rewriter, origLoc, finalOffset);
 
     // The sizes of the final type are defined directly by the input sizes of
     // the subview.
     // Moreover subviews can drop some dimensions, some strides and sizes may
     // not end up in the final <base, offset, sizes, strides> value that we are
     // replacing.
     // Do the filtering here.
     SmallVector<OpFoldResult> subSizes = subview.getMixedSizes();
-    const unsigned sizeStartIdx = 2;
-    const unsigned strideStartIdx = sizeStartIdx + subRank;
-    unsigned insertedDims = 0;
     llvm::SmallBitVector droppedDims = subview.getDroppedDims();
+
+    SmallVector<OpFoldResult> finalSizes;
+    finalSizes.reserve(subRank);
+
+    SmallVector<OpFoldResult> finalStrides;
+    finalStrides.reserve(subRank);
+
     for (unsigned i = 0; i < sourceRank; ++i) {
       if (droppedDims.test(i))
         continue;
 
-      results[sizeStartIdx + insertedDims] =
-          getValueOrCreateConstantIndexOp(rewriter, origLoc, subSizes[i]);
-      results[strideStartIdx + insertedDims] =
-          getValueOrCreateConstantIndexOp(rewriter, origLoc, strides[i]);
-      ++insertedDims;
+      finalSizes.push_back(subSizes[i]);
+      finalStrides.push_back(strides[i]);
     }
-    assert(insertedDims == subRank &&
+    assert(finalSizes.size() == subRank &&
            "Should have populated all the values at this point");
 
-    rewriter.replaceOp(op, results);
+    auto memrefDesc = rewriter.create<memref::ReinterpretCastOp>(
+        origLoc, subType, newExtractStridedMetadata.getBaseBuffer(),
+        finalOffset,
+        /*sizes=*/finalSizes,
+        /*strides=*/finalStrides);
+    rewriter.replaceOp(subview, memrefDesc.getResult());
     return success();
   }
 };
 
 /// Compute the expanded sizes of the given \p expandShape for the
 /// \p groupId-th reassociation group.
 /// \p origSizes hold the sizes of the source shape as values.
 /// This is used to compute the new sizes in cases of dynamic shapes.
matchAndRewrite(memref::ExtractStridedMetadataOp extractStridedMetadataOp,
     return success();
   }
 };
 } // namespace
 
 void memref::populateSimplifyExtractStridedMetadataOpPatterns(
     RewritePatternSet &patterns) {
   patterns
-      .add<ExtractStridedMetadataOpSubviewFolder,
+      .add<SubviewFolder,
            ExtractStridedMetadataOpReshapeFolder<
                memref::ExpandShapeOp, getExpandedSizes, getExpandedStrides>,
            ExtractStridedMetadataOpReshapeFolder<
                memref::CollapseShapeOp, getCollapsedSize, getCollapsedStride>,
            ExtractStridedMetadataOpAllocFolder<memref::AllocOp>,
            ExtractStridedMetadataOpAllocFolder<memref::AllocaOp>,
            RewriteExtractAlignedPointerAsIndexOfViewLikeOp,
            ExtractStridedMetadataOpReinterpretCastFolder,

mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir

func.func @extract_strided_metadata_constants(%base: memref<5x4xf32, strided<[4, 1], offset: 2>>)
 
   // CHECK: %[[BASE]], %[[C2]], %[[C5]], %[[C4]], %[[C4]], %[[C1]]
   return %base_buffer, %offset, %sizes#0, %sizes#1, %strides#0, %strides#1 :
     memref<f32>, index, index, index, index, index
 }
 
 // -----
 
+// Check that we simplify subview(src) into:
+// base, offset, sizes, strides xtract_strided_metadata src
+// final_sizes = subSizes
+// final_strides = <some math> strides
+// final_offset = <some math> offset
+// reinterpret_cast base to final_offset, final_sizes, final_ strides
+//
+// Orig strides: [s0, s1, s2]
+// Sub strides: [subS0, subS1, subS2]
+// => New strides: [s0 * subS0, s1 * subS1, s2 * subS2]
+// ==> 1 affine map (used for each stride) with two values.
+//
+// Orig offset: origOff
+// Sub offsets: [subO0, subO1, subO2]
+// => Final offset: s0 * subS0 * subO0 + ... + s2 * subS2 * subO2 + origOff
+// ==> 1 affine map with (rank * 3 + 1) symbols
+//
+// CHECK-DAG: #[[$STRIDE_MAP:.*]] = affine_map<()[s0, s1] -> (s0 * s1)>
+// CHECK-DAG: #[[$OFFSET_MAP:.*]] = affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0 + (s1 * s2) * s3 + (s4 * s5) * s6 + (s7 * s8) * s9)>
+// CHECK-LABEL: func @simplify_subview_all_dynamic
+//  CHECK-SAME: (%[[ARG:.*]]: memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>, %[[DYN_OFFSET0:.*]]: index, %[[DYN_OFFSET1:.*]]: index, %[[DYN_OFFSET2:.*]]: index, %[[DYN_SIZE0:.*]]: index, %[[DYN_SIZE1:.*]]: index, %[[DYN_SIZE2:.*]]: index, %[[DYN_STRIDE0:.*]]: index, %[[DYN_STRIDE1:.*]]: index, %[[DYN_STRIDE2:.*]]: index)
+//
+//   CHECK-DAG: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:3, %[[STRIDES:.*]]:3 = memref.extract_strided_metadata %[[ARG]]
+//
+//  CHECK-DAG: %[[FINAL_STRIDE0:.*]] = affine.apply #[[$STRIDE_MAP]]()[%[[DYN_STRIDE0]], %[[STRIDES]]#0]
+//  CHECK-DAG: %[[FINAL_STRIDE1:.*]] = affine.apply #[[$STRIDE_MAP]]()[%[[DYN_STRIDE1]], %[[STRIDES]]#1]
+//  CHECK-DAG: %[[FINAL_STRIDE2:.*]] = affine.apply #[[$STRIDE_MAP]]()[%[[DYN_STRIDE2]], %[[STRIDES]]#2]
+//
+//  CHECK-DAG: %[[FINAL_OFFSET:.*]] = affine.apply #[[$OFFSET_MAP]]()[%[[OFFSET]], %[[DYN_OFFSET0]], %[[DYN_STRIDE0]], %[[STRIDES]]#0, %[[DYN_OFFSET1]], %[[DYN_STRIDE1]], %[[STRIDES]]#1, %[[DYN_OFFSET2]], %[[DYN_STRIDE2]], %[[STRIDES]]#2]
+//
+//      CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[FINAL_OFFSET]]], sizes: [%[[DYN_SIZE0]], %[[DYN_SIZE1]], %[[DYN_SIZE2]]], strides: [%[[FINAL_STRIDE0]], %[[FINAL_STRIDE1]], %[[FINAL_STRIDE2]]]
+//
+//       CHECK: return %[[RES]]
+func.func @simplify_subview_all_dynamic(
+    %base: memref<?x?x?xf32, strided<[?,?,?], offset:?>>,
+    %offset0: index, %offset1: index, %offset2: index,
+    %size0: index, %size1: index, %size2: index,
+    %stride0: index, %stride1: index, %stride2: index)
+    -> memref<?x?x?xf32, strided<[?,?,?], offset:?>> {
+
+  %subview = memref.subview %base[%offset0, %offset1, %offset2]
+                                 [%size0, %size1, %size2]
+                                 [%stride0, %stride1, %stride2] :
+    memref<?x?x?xf32, strided<[?,?,?], offset: ?>> to
+      memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
+
+  return %subview : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
+}
+
+// -----
+
 // Check that we simplify extract_strided_metadata of subview to
 // base_buf, base_offset, base_sizes, base_strides = extract_strided_metadata
 // strides = base_stride_i * subview_stride_i
 // offset = base_offset + sum(subview_offsets_i * strides_i).
 //
 // This test also checks that we don't create useless arith operations
 // when subview_offsets_i is 0.
 //