[PATCH] D135836: [mlir][MemRef] Add a extract_strided_metadata(extract_strided_metadata) pattern

[Quentin Colombet via Phabricator via llvm-commits]

This revision was landed with ongoing or failed builds.
This revision was automatically updated to reflect the committed changes.
Closed by commit rG3a33c146edd2: [mlir][MemRef] Add a extract_strided_metadata(extract_strided_metadata) pattern (authored by qcolombet).

Changed prior to commit:
  https://reviews.llvm.org/D135836?vs=467323&id=467878#toc

Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D135836/new/

https://reviews.llvm.org/D135836

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


Index: mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir
===================================================================
--- mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir
+++ mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir
@@ -841,3 +841,29 @@
   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
+}
Index: mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp
===================================================================
--- mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp
+++ mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp
@@ -718,6 +718,34 @@
     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(
@@ -731,7 +759,8 @@
            ForwardStaticMetadata,
            ExtractStridedMetadataOpAllocFolder<memref::AllocOp>,
            ExtractStridedMetadataOpAllocFolder<memref::AllocaOp>,
-           RewriteExtractAlignedPointerAsIndexOfViewLikeOp>(
+           RewriteExtractAlignedPointerAsIndexOfViewLikeOp,
+           ExtractStridedMetadataOpExtractStridedMetadataFolder>(
           patterns.getContext());
 }
 


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D135836.467878.patch
Type: text/x-patch
Size: 3088 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20221014/9ee70fb3/attachment.bin>