[Mlir-commits] [mlir] 4341c0c - [mlir][memref] Implement ValueBoundsOpInterface for memref.extract_strided_metadata/assume_alignment (#206466)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jul 6 00:53:43 PDT 2026
Author: Hagai Lev Hacohen
Date: 2026-07-06T07:53:39Z
New Revision: 4341c0c9b8b88aac576ce55211ea9db032b74124
URL: https://github.com/llvm/llvm-project/commit/4341c0c9b8b88aac576ce55211ea9db032b74124
DIFF: https://github.com/llvm/llvm-project/commit/4341c0c9b8b88aac576ce55211ea9db032b74124.diff
LOG: [mlir][memref] Implement ValueBoundsOpInterface for memref.extract_strided_metadata/assume_alignment (#206466)
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.
Added:
Modified:
mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp
mlir/test/Dialect/MemRef/value-bounds-op-interface-impl.mlir
Removed:
################################################################################
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..8408b166a6b2a 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,43 @@ 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);
+ auto result = llvm::cast<OpResult>(value);
+ assert(result.getOwner() == op && "invalid value");
+ int64_t resultNumber = result.getResultNumber();
+
+ if (resultNumber == 1) {
+ cstr.bound(value) == metadataOp.getConstifiedMixedOffset();
+ return;
+ }
+
+ int64_t sourceRank = metadataOp.getSource().getType().getRank();
+ int64_t sizeStart = 2;
+ int64_t strideStart = sizeStart + sourceRank;
+ if (resultNumber >= sizeStart && resultNumber < strideStart) {
+ int64_t idx = resultNumber - sizeStart;
+ cstr.bound(value) >= 0;
+ cstr.bound(value) == cstr.getExpr(metadataOp.getSource(), idx);
+ return;
+ }
+
+ int64_t strideEnd = strideStart + sourceRank;
+ if (resultNumber >= strideStart && resultNumber < strideEnd) {
+ int64_t idx = resultNumber - strideStart;
+ SmallVector<OpFoldResult> strides =
+ metadataOp.getConstifiedMixedStrides();
+ cstr.bound(value) == strides[idx];
+ return;
+ }
+ llvm_unreachable("unexpected index value from extract_strided_metadata");
+ }
+};
+
struct GetGlobalOpInterface
: public ValueBoundsOpInterface::ExternalModel<GetGlobalOpInterface,
GetGlobalOp> {
@@ -162,5 +212,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..32af830aa1ea7 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,29 @@ func.func @memref_expand(%m: memref<?xf32>, %sz: index) -> (index, index) {
// -----
+// CHECK-LABEL: func @memref_extract_strided_metadata_static_metadata(
+// CHECK-SAME: %[[m:.*]]: memref<4x?xf32, strided<[11, 7], offset: 5>>
+// CHECK: %[[c5:.*]] = arith.constant 5 : index
+// CHECK: %[[c4:.*]] = arith.constant 4 : index
+// CHECK: %[[c1:.*]] = arith.constant 1 : index
+// CHECK: %[[dim:.*]] = memref.dim %[[m]], %[[c1]]
+// CHECK: %[[c11:.*]] = arith.constant 11 : index
+// CHECK: %[[c7:.*]] = arith.constant 7 : index
+// CHECK: return %[[c5]], %[[c4]], %[[dim]], %[[c11]], %[[c7]]
+func.func @memref_extract_strided_metadata_static_metadata(
+ %m: memref<4x?xf32, strided<[11, 7], offset: 5>>) -> (index, index, index, index, index) {
+ %base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %m
+ : memref<4x?xf32, strided<[11, 7], offset: 5>> -> memref<f32>, index, index, index, index, index
+ %0 = "test.reify_bound"(%offset) : (index) -> (index)
+ %1 = "test.reify_bound"(%sizes#0) : (index) -> (index)
+ %2 = "test.reify_bound"(%sizes#1) : (index) -> (index)
+ %3 = "test.reify_bound"(%strides#0) : (index) -> (index)
+ %4 = "test.reify_bound"(%strides#1) : (index) -> (index)
+ return %0, %1, %2, %3, %4 : index, index, index, index, index
+}
+
+// -----
+
// CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 * 2)>
// CHECK-LABEL: func @memref_collapse(
// CHECK-SAME: %[[sz0:.*]]: index
More information about the Mlir-commits
mailing list