[Mlir-commits] [mlir] [mlir][memref] Add ValueBoundsOpInterface external models for the memref operations extract_strided_metadata and assume_alignment (PR #206466)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Hagai Lev Hacohen (HagaiLevHacohen)

<details>
<summary>Changes</summary>

Add ValueBoundsOpInterface external models for memref.assume_alignment and memref.extract_strided_metadata, so dimension and metadata bounds can be propagated through these view-like ops during value-bounds analysis.

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


3 Files Affected:

- (modified) mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp (+3-2) 
- (modified) mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp (+48) 
- (modified) mlir/test/Dialect/MemRef/value-bounds-op-interface-impl.mlir (+37) 


``````````diff
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
index a1e3f10a871c1..8cecf5a4897e9 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
@@ -55,8 +55,9 @@ void mlir::memref::MemRefDialect::initialize() {
   declarePromisedInterfaces<RuntimeVerifiableOpInterface, AssumeAlignmentOp,
                             AtomicRMWOp, CastOp, CopyOp, DimOp, ExpandShapeOp,
                             GenericAtomicRMWOp, LoadOp, StoreOp, SubViewOp>();
-  declarePromisedInterfaces<ValueBoundsOpInterface, AllocOp, AllocaOp, CastOp,
-                            DimOp, GetGlobalOp, RankOp, SubViewOp>();
+  declarePromisedInterfaces<
+      ValueBoundsOpInterface, AllocOp, AllocaOp, AssumeAlignmentOp, CastOp,
+      DimOp, ExtractStridedMetadataOp, GetGlobalOp, RankOp, SubViewOp>();
   declarePromisedInterface<DestructurableTypeInterface, MemRefType>();
 }
 
diff --git a/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp
index 69afbcadb0b07..042c8920fadd8 100644
--- a/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -30,6 +30,19 @@ struct AllocOpInterface
   }
 };
 
+struct AssumeAlignmentOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<AssumeAlignmentOpInterface,
+                                                   memref::AssumeAlignmentOp> {
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto assumeAlignmentOp = cast<memref::AssumeAlignmentOp>(op);
+    assert(value == assumeAlignmentOp.getResult() && "invalid value");
+
+    cstr.bound(value)[dim] ==
+        cstr.getExpr(assumeAlignmentOp.getViewSource(), dim);
+  }
+};
+
 struct CastOpInterface
     : public ValueBoundsOpInterface::ExternalModel<CastOpInterface, CastOp> {
   void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
@@ -70,6 +83,37 @@ struct ExpandShapeOpInterface
   }
 };
 
+struct ExtractStridedMetadataOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<
+          ExtractStridedMetadataOpInterface, memref::ExtractStridedMetadataOp> {
+  void populateBoundsForIndexValue(Operation *op, Value value,
+                                   ValueBoundsConstraintSet &cstr) const {
+    auto metadataOp = cast<memref::ExtractStridedMetadataOp>(op);
+
+    if (value == metadataOp.getOffset()) {
+      cstr.bound(value) == metadataOp.getConstifiedMixedOffset();
+      return;
+    }
+
+    for (auto [idx, size] : llvm::enumerate(metadataOp.getSizes())) {
+      if (value != size)
+        continue;
+      cstr.bound(value) >= 0;
+      cstr.bound(value) == cstr.getExpr(metadataOp.getSource(), idx);
+      return;
+    }
+
+    SmallVector<OpFoldResult> strides = metadataOp.getConstifiedMixedStrides();
+    for (auto [idx, stride] : llvm::enumerate(metadataOp.getStrides())) {
+      if (value != stride)
+        continue;
+      cstr.bound(value) == strides[idx];
+      return;
+    }
+    llvm_unreachable("unexpected index value from extract_strided_metadata");
+  }
+};
+
 struct GetGlobalOpInterface
     : public ValueBoundsOpInterface::ExternalModel<GetGlobalOpInterface,
                                                    GetGlobalOp> {
@@ -162,5 +206,9 @@ void mlir::memref::registerValueBoundsOpInterfaceExternalModels(
     memref::GetGlobalOp::attachInterface<memref::GetGlobalOpInterface>(*ctx);
     memref::RankOp::attachInterface<memref::RankOpInterface>(*ctx);
     memref::SubViewOp::attachInterface<memref::SubViewOpInterface>(*ctx);
+    memref::AssumeAlignmentOp::attachInterface<
+        memref::AssumeAlignmentOpInterface>(*ctx);
+    memref::ExtractStridedMetadataOp::attachInterface<
+        memref::ExtractStridedMetadataOpInterface>(*ctx);
   });
 }
diff --git a/mlir/test/Dialect/MemRef/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/MemRef/value-bounds-op-interface-impl.mlir
index d0aec68d54988..53a73a77b99c3 100644
--- a/mlir/test/Dialect/MemRef/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/MemRef/value-bounds-op-interface-impl.mlir
@@ -27,6 +27,22 @@ func.func @memref_alloca(%sz: index) -> (index, index) {
 
 // -----
 
+// CHECK-LABEL: func @memref_assume_alignment(
+//  CHECK-SAME:     %[[sz:.*]]: index
+//       CHECK:   %[[c6:.*]] = arith.constant 6 : index
+//       CHECK:   %[[c1:.*]] = arith.constant 1 : index
+//       CHECK:   %[[dim:.*]] = memref.dim %{{.*}}, %[[c1]] : memref<6x?xf32>
+//       CHECK:   return %[[c6]], %[[dim]]
+func.func @memref_assume_alignment(%sz: index) -> (index, index) {
+  %0 = memref.alloc(%sz) : memref<6x?xf32>
+  %1 = memref.assume_alignment %0, 16 : memref<6x?xf32>
+  %2 = "test.reify_bound"(%1) {dim = 0} : (memref<6x?xf32>) -> (index)
+  %3 = "test.reify_bound"(%1) {dim = 1} : (memref<6x?xf32>) -> (index)
+  return %2, %3 : index, index
+}
+
+// -----
+
 // CHECK-LABEL: func @memref_cast(
 //       CHECK:   %[[c10:.*]] = arith.constant 10 : index
 //       CHECK:   return %[[c10]]
@@ -77,6 +93,27 @@ func.func @memref_expand(%m: memref<?xf32>, %sz: index) -> (index, index) {
 
 // -----
 
+// CHECK-LABEL: func @memref_extract_strided_metadata(
+//  CHECK-SAME:     %[[sz:.*]]: index
+//       CHECK:   %[[c0:.*]] = arith.constant 0 : index
+//       CHECK:   %[[c10:.*]] = arith.constant 10 : index
+//       CHECK:   %[[c1:.*]] = arith.constant 1 : index
+//       CHECK:   %[[dim:.*]] = memref.dim %{{.*}}, %[[c1]] : memref<10x?xf32>
+//       CHECK:   %[[c1_0:.*]] = arith.constant 1 : index
+//       CHECK:   return %[[c0]], %[[c10]], %[[dim]], %[[c1_0]]
+func.func @memref_extract_strided_metadata(%sz: index) -> (index, index, index, index) {
+  %0 = memref.alloc(%sz) : memref<10x?xf32>
+  %base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %0
+    : memref<10x?xf32> -> memref<f32>, index, index, index, index, index
+  %1 = "test.reify_bound"(%offset) : (index) -> (index)
+  %2 = "test.reify_bound"(%sizes#0) : (index) -> (index)
+  %3 = "test.reify_bound"(%sizes#1) : (index) -> (index)
+  %4 = "test.reify_bound"(%strides#1) : (index) -> (index)
+  return %1, %2, %3, %4 : index, index, index, index
+}
+
+// -----
+
 //       CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 * 2)>
 // CHECK-LABEL: func @memref_collapse(
 //  CHECK-SAME:     %[[sz0:.*]]: index

``````````

</details>


https://github.com/llvm/llvm-project/pull/206466