[Mlir-commits] [mlir] df63870 - [mlir][memref]Add pattern to forward memref.extract_aligned_pointer_as_index(view_like_op) to its source

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2022-09-29T02:27:01-07:00
New Revision: df6387079e531bffe385d8c6aafa6d9aeb6a385c

URL: https://github.com/llvm/llvm-project/commit/df6387079e531bffe385d8c6aafa6d9aeb6a385c
DIFF: https://github.com/llvm/llvm-project/commit/df6387079e531bffe385d8c6aafa6d9aeb6a385c.diff

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

Differential Revision: https://reviews.llvm.org/D134835

Added: 
    

Modified: 
    mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp
    mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp b/mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp
index 1d42c9109447a..5a6f72ea56398 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp
@@ -575,6 +575,26 @@ struct ExtractStridedMetadataOpAllocFolder
     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(
@@ -582,7 +602,8 @@ void memref::populateSimplifyExtractStridedMetadataOpPatterns(
   patterns.add<ExtractStridedMetadataOpSubviewFolder,
                ExtractStridedMetadataOpExpandShapeFolder, ForwardStaticMetadata,
                ExtractStridedMetadataOpAllocFolder<memref::AllocOp>,
-               ExtractStridedMetadataOpAllocFolder<memref::AllocaOp>>(
+               ExtractStridedMetadataOpAllocFolder<memref::AllocaOp>,
+               RewriteExtractAlignedPointerAsIndexOfViewLikeOp>(
       patterns.getContext());
 }
 

diff  --git a/mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir b/mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir
index 449b35b9ed667..2862ef96a53c8 100644
--- a/mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir
+++ b/mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir
@@ -750,3 +750,15 @@ func.func @extract_strided_metadata_of_alloc_with_strided(%arg : 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
+}