[mlir][memref]Add pattern to forward memref.extract_aligned_pointer_as_index(view_like_op) to its source

[qcolombet]

LGTM.

Not for this patch, but we may want to rename the pass. E.g., SimplifyExtractOpsPass, not a great name though :P.

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

LogicalResult matchAndRewrite(memref::ExtractStridedMetadataOp op,
 
     for (OpFoldResult stride : strides)
       results.push_back(getValueOrCreateConstantIndexOp(rewriter, loc, stride));
 
     rewriter.replaceOp(op, results);
     return success();
   }
 };
+
+/// Rewrite memref.extract_aligned_pointer_as_index of a ViewLikeOp to the
+/// source of the ViewLikeOp.
+class RewriteExtractAlignedPointerAsIndexOfViewLikeOp
+    : public OpRewritePattern<memref::ExtractAlignedPointerAsIndexOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult
+  matchAndRewrite(memref::ExtractAlignedPointerAsIndexOp extractOp,
+                  PatternRewriter &rewriter) const override {
+    auto viewLikeOp =
+        extractOp.getSource().getDefiningOp<ViewLikeOpInterface>();
+    if (!viewLikeOp)
+      return rewriter.notifyMatchFailure(extractOp, "not a ViewLike source");
+    rewriter.updateRootInPlace(extractOp, [&]() {
+      extractOp.sourceMutable().assign(viewLikeOp.getViewSource());
+    });
+    return success();
+  }
+};
 } // namespace
 
 void memref::populateSimplifyExtractStridedMetadataOpPatterns(
     RewritePatternSet &patterns) {
   patterns.add<ExtractStridedMetadataOpSubviewFolder,
                ExtractStridedMetadataOpExpandShapeFolder, ForwardStaticMetadata,
                ExtractStridedMetadataOpAllocFolder<memref::AllocOp>,
-               ExtractStridedMetadataOpAllocFolder<memref::AllocaOp>>(
+               ExtractStridedMetadataOpAllocFolder<memref::AllocaOp>,
+               RewriteExtractAlignedPointerAsIndexOfViewLikeOp>(
       patterns.getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // Pass registration
 //===----------------------------------------------------------------------===//
 
 namespace {

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

func.func @extract_strided_metadata_of_alloc_with_strided(%arg : index)
   %A = memref.alloc() : memref<4xi16, strided<[12]>>
   %base, %offset, %size, %stride = memref.extract_strided_metadata %A :
     memref<4xi16, strided<[12]>>
     -> memref<i16>, index, index, index
 
   return %base, %offset, %size, %stride :
       memref<i16>, index, index, index
 }
+
+// -----
+
+// CHECK-LABEL: extract_aligned_pointer_as_index
+//  CHECK-SAME: (%[[ARG0:.*]]: memref<f32>
+func.func @extract_aligned_pointer_as_index(%arg0: memref<f32>) -> index {
+  // CHECK-NOT: memref.subview
+  //     CHECK: memref.extract_aligned_pointer_as_index %[[ARG0]]
+  %c = memref.subview %arg0[] [] [] : memref<f32> to memref<f32>
+  %r = memref.extract_aligned_pointer_as_index %arg0: memref<f32> -> index
+  return %r : index
+}