[mlir][MemRef] Add a extract_strided_metadata(extract_strided_metadata) pattern

[qcolombet]

Summary

This pattern will be useful to get cleaner code when lowering view like operations.

Note: This patch is part of a series that replaces the lowering of view like operations in MemRefToLLVM with composable patterns.

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

matchAndRewrite(memref::ExtractAlignedPointerAsIndexOp extractOp,
     if (!viewLikeOp)
       return rewriter.notifyMatchFailure(extractOp, "not a ViewLike source");
     rewriter.updateRootInPlace(extractOp, [&]() {
       extractOp.getSourceMutable().assign(viewLikeOp.getViewSource());
     });
     return success();
   }
 };
+
+/// Replace `base, offset =
+///            extract_strided_metadata(extract_strided_metadata(src)#0)`
+/// With
+/// ```
+/// base, ... = extract_strided_metadata(src)
+/// offset = 0
+/// ```
+class ExtractStridedMetadataOpExtractStridedMetadataFolder
+    : public OpRewritePattern<memref::ExtractStridedMetadataOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult
+  matchAndRewrite(memref::ExtractStridedMetadataOp extractStridedMetadataOp,
+                  PatternRewriter &rewriter) const override {
+    auto sourceExtractStridedMetadataOp =
+        extractStridedMetadataOp.getSource()
+            .getDefiningOp<memref::ExtractStridedMetadataOp>();
+    if (!sourceExtractStridedMetadataOp)
+      return failure();
+    Location loc = extractStridedMetadataOp.getLoc();
+    rewriter.replaceOp(extractStridedMetadataOp,
+                       {sourceExtractStridedMetadataOp.getBaseBuffer(),
+                        getValueOrCreateConstantIndexOp(
+                            rewriter, loc, rewriter.getIndexAttr(0))});
+    return success();
+  }
+};
 } // namespace
 
 void memref::populateSimplifyExtractStridedMetadataOpPatterns(
     RewritePatternSet &patterns) {
   patterns
       .add<ExtractStridedMetadataOpSubviewFolder,
            ExtractStridedMetadataOpReshapeFolder<
                memref::ExpandShapeOp, getExpandedSizes, getExpandedStrides>,
            ExtractStridedMetadataOpReshapeFolder<
                memref::CollapseShapeOp, getCollapsedSize, getCollapsedStride>,
            ForwardStaticMetadata,
            ExtractStridedMetadataOpAllocFolder<memref::AllocOp>,
            ExtractStridedMetadataOpAllocFolder<memref::AllocaOp>,
-           RewriteExtractAlignedPointerAsIndexOfViewLikeOp>(
+           RewriteExtractAlignedPointerAsIndexOfViewLikeOp,
+           ExtractStridedMetadataOpExtractStridedMetadataFolder>(
           patterns.getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // Pass registration
 //===----------------------------------------------------------------------===//
 
 namespace {

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

func.func @extract_strided_metadata_of_collapse_to_rank0(%arg : memref<1x1x1x1x1x1xi32>)
 
   %base, %offset =
     memref.extract_strided_metadata %collapsed_view : memref<i32>
     -> memref<i32>, index
 
   return %base, %offset :
       memref<i32>, index
 }
+
+// -----
+
+// Check that we simplify extract_strided_metadata of
+// extract_strided_metadata.
+//
+// CHECK-LABEL: func @extract_strided_metadata_of_extract_strided_metadata(
+//  CHECK-SAME: %[[ARG:.*]]: memref<i32>)
+//
+//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[BASE:.*]], %[[OFFSET:.*]] = memref.extract_strided_metadata %[[ARG]]
+//
+//       CHECK: return %[[BASE]], %[[C0]]
+func.func @extract_strided_metadata_of_extract_strided_metadata(%arg : memref<i32>)
+  -> (memref<i32>, index) {
+
+  %base, %offset =
+    memref.extract_strided_metadata %arg:memref<i32>
+    -> memref<i32>, index
+  %base2, %offset2 =
+    memref.extract_strided_metadata %base:memref<i32>
+    -> memref<i32>, index
+
+  return %base2, %offset2 :
+      memref<i32>, index
+}