[mlir][memref] Remove invalid extract_aligned_pointer_as_index folding in ExpandStridedMetadata

[Hardcode84]

RewriteExtractAlignedPointerAsIndexOfViewLikeOp tries to propagate extract_aligned_pointer_as_index through the view ops.

ViewLikeOpInterface by itself doesn't guarantee to preserve the base pointer and memref.view is one such example, so limit pattern to a few specific ops.

[llvmbot]

@llvm/pr-subscribers-mlir-memref

@llvm/pr-subscribers-mlir

Author: Ivan Butygin (Hardcode84)

Changes

ViewLikeOpInterface by itself doesn't guarantee to preserve the base pointer and memref.view is one such example, so limit pattern to a few specific ops.

---

Full diff: https://github.com/llvm/llvm-project/pull/167615.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp (+5-1)
• (modified) mlir/test/Dialect/MemRef/expand-strided-metadata.mlir (+17)

diff --git a/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp b/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp
index bd02516d5b527..c9352e8f700d7 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp
@@ -959,7 +959,11 @@ class RewriteExtractAlignedPointerAsIndexOfViewLikeOp
                   PatternRewriter &rewriter) const override {
     auto viewLikeOp =
         extractOp.getSource().getDefiningOp<ViewLikeOpInterface>();
-    if (!viewLikeOp || extractOp.getSource() != viewLikeOp.getViewDest())
+    // ViewLikeOpInterface by itself doesn't guarantee to preserve the base
+    // pointer in general and `memref.view` is one such example, so just check
+    // for a few specific cases.
+    if (!viewLikeOp || extractOp.getSource() != viewLikeOp.getViewDest() ||
+        !isa<memref::SubViewOp, memref::ReinterpretCastOp>(viewLikeOp))
       return rewriter.notifyMatchFailure(extractOp, "not a ViewLike source");
     rewriter.modifyOpInPlace(extractOp, [&]() {
       extractOp.getSourceMutable().assign(viewLikeOp.getViewSource());
diff --git a/mlir/test/Dialect/MemRef/expand-strided-metadata.mlir b/mlir/test/Dialect/MemRef/expand-strided-metadata.mlir
index 18cdfb73f6ba8..ec17282ed20ec 100644
--- a/mlir/test/Dialect/MemRef/expand-strided-metadata.mlir
+++ b/mlir/test/Dialect/MemRef/expand-strided-metadata.mlir
@@ -1455,3 +1455,20 @@ func.func @extract_strided_metadata_of_memory_space_cast_no_base(%base: memref<2
 
 // CHECK-LABEL:  func @extract_strided_metadata_of_memory_space_cast_no_base
 //   CHECK-NOT:  memref.memory_space_cast
+
+// -----
+
+func.func @negative_memref_view_extract_aligned_pointer(%arg0: memref<?xi8>) -> index {
+  // `extract_aligned_pointer_as_index` must not be folded as `memref.view` can change the base pointer
+  // CHECK-LABEL: func @negative_memref_view_extract_aligned_pointer
+  // CHECK-SAME: (%[[ARG0:.*]]: memref<?xi8>)
+  // CHECK: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK: %[[VIEW:.*]] = memref.view %[[ARG0]][%[[C0]]][] : memref<?xi8> to memref<f32>
+  // CHECK: %[[PTR:.*]] = memref.extract_aligned_pointer_as_index %[[VIEW]] : memref<f32> -> index
+  // CHECK: return %[[PTR]] : index
+
+  %c0 = arith.constant 0 : index
+  %0 = memref.view %arg0[%c0][] : memref<?xi8> to memref<f32>
+  %1 = memref.extract_aligned_pointer_as_index %0: memref<f32> -> index
+  return %1 : index
+}

[krzysz00]

Approved modulo test updates