[Mlir-commits] [mlir] [MLIR][SROA][Mem2Reg] Add data layout to interface methods (PR #85644)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Mar 18 07:10:10 PDT 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-core
Author: Christian Ulmann (Dinistro)
<details>
<summary>Changes</summary>
This commit expends the Mem2Reg and SROA interface methods with passed in handles to a `DataLayout` structure. This is done to avoid superfluous retreiving of data layouts during each conversion of intrinsics.
This change, additionally, enables subsequent changes to make the LLVM dialect implementation of these interfaces type agnostic.
---
Patch is 38.11 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/85644.diff
7 Files Affected:
- (modified) mlir/include/mlir/Interfaces/MemorySlotInterfaces.td (+20-22)
- (modified) mlir/include/mlir/Transforms/Mem2Reg.h (+1-1)
- (modified) mlir/include/mlir/Transforms/SROA.h (+2-1)
- (modified) mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp (+92-61)
- (modified) mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp (+12-6)
- (modified) mlir/lib/Transforms/Mem2Reg.cpp (+17-7)
- (modified) mlir/lib/Transforms/SROA.cpp (+19-11)
``````````diff
diff --git a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
index 9ffa709cc5bfd4..adf324bf559791 100644
--- a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
+++ b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
@@ -83,11 +83,7 @@ def PromotableAllocationOpInterface
def PromotableMemOpInterface : OpInterface<"PromotableMemOpInterface"> {
let description = [{
Describes an operation that can load from memory slots and/or store
- to memory slots. Loads and stores must be of whole values of the same
- type as the slot itself.
-
- For a memory operation on a slot to be valid, it must operate on the slot
- pointer *only as a pointer to an element of the type of the slot*.
+ to memory slots.
If the same operation does both loads and stores on the same slot, the
load must semantically happen first.
@@ -142,7 +138,8 @@ def PromotableMemOpInterface : OpInterface<"PromotableMemOpInterface"> {
}], "bool", "canUsesBeRemoved",
(ins "const ::mlir::MemorySlot &":$slot,
"const ::llvm::SmallPtrSetImpl<::mlir::OpOperand *> &":$blockingUses,
- "::llvm::SmallVectorImpl<::mlir::OpOperand *> &":$newBlockingUses)
+ "::llvm::SmallVectorImpl<::mlir::OpOperand *> &":$newBlockingUses,
+ "const ::mlir::DataLayout &":$datalayout)
>,
InterfaceMethod<[{
Transforms IR to ensure that the current operation does not use the
@@ -197,7 +194,8 @@ def PromotableOpInterface : OpInterface<"PromotableOpInterface"> {
No IR mutation is allowed in this method.
}], "bool", "canUsesBeRemoved",
(ins "const ::llvm::SmallPtrSetImpl<::mlir::OpOperand *> &":$blockingUses,
- "::llvm::SmallVectorImpl<::mlir::OpOperand *> &":$newBlockingUses)
+ "::llvm::SmallVectorImpl<::mlir::OpOperand *> &":$newBlockingUses,
+ "const ::mlir::DataLayout &":$datalayout)
>,
InterfaceMethod<[{
Transforms IR to ensure that the current operation does not use the
@@ -285,29 +283,28 @@ def DestructurableAllocationOpInterface
def SafeMemorySlotAccessOpInterface
: OpInterface<"SafeMemorySlotAccessOpInterface"> {
let description = [{
- Describes operations using memory slots in a type-safe manner.
+ Describes operations using memory slots in a safe manner.
}];
let cppNamespace = "::mlir";
let methods = [
InterfaceMethod<[{
Returns whether all accesses in this operation to the provided slot are
- done in a type-safe manner. To be type-safe, the access must only load
- the value in this type as the type of the slot, and without assuming any
- context around the slot. For example, a type-safe load must not load
- outside the bounds of the slot.
+ done in a safe manner. To be safe, the access most only access the slot
+ inside the bounds that its type implies.
- If the type-safety of the accesses depends on the type-safety of the
- accesses to further memory slots, the result of this method will be
- conditioned to the type-safety of the accesses to the slots added by
- this method to `mustBeSafelyUsed`.
+ If the safety of the accesses depends on the safety of the accesses to
+ further memory slots, the result of this method will be conditioned to
+ the safety of the accesses to the slots added by this method to
+ `mustBeSafelyUsed`.
No IR mutation is allowed in this method.
}],
"::mlir::LogicalResult",
"ensureOnlySafeAccesses",
(ins "const ::mlir::MemorySlot &":$slot,
- "::mlir::SmallVectorImpl<::mlir::MemorySlot> &":$mustBeSafelyUsed)
+ "::mlir::SmallVectorImpl<::mlir::MemorySlot> &":$mustBeSafelyUsed,
+ "const ::mlir::DataLayout &":$dataLayout)
>
];
}
@@ -323,13 +320,12 @@ def DestructurableAccessorOpInterface
InterfaceMethod<[{
For a given destructurable memory slot, returns whether this operation can
rewire its uses of the slot to use the slots generated after
- destructuring. This may involve creating new operations, and usually
- amounts to checking if the pointer types match.
+ destructuring. This may involve creating new operations.
This method must also register the indices it will access within the
`usedIndices` set. If the accessor generates new slots mapping to
subelements, they must be registered in `mustBeSafelyUsed` to ensure
- they are used in a locally type-safe manner.
+ they are used in a safe manner.
No IR mutation is allowed in this method.
}],
@@ -337,7 +333,8 @@ def DestructurableAccessorOpInterface
"canRewire",
(ins "const ::mlir::DestructurableMemorySlot &":$slot,
"::llvm::SmallPtrSetImpl<::mlir::Attribute> &":$usedIndices,
- "::mlir::SmallVectorImpl<::mlir::MemorySlot> &":$mustBeSafelyUsed)
+ "::mlir::SmallVectorImpl<::mlir::MemorySlot> &":$mustBeSafelyUsed,
+ "const ::mlir::DataLayout &":$dataLayout)
>,
InterfaceMethod<[{
Rewires the use of a slot to the generated subslots, without deleting
@@ -351,7 +348,8 @@ def DestructurableAccessorOpInterface
"rewire",
(ins "const ::mlir::DestructurableMemorySlot &":$slot,
"::llvm::DenseMap<::mlir::Attribute, ::mlir::MemorySlot> &":$subslots,
- "::mlir::RewriterBase &":$rewriter)
+ "::mlir::RewriterBase &":$rewriter,
+ "const ::mlir::DataLayout &":$dataLayout)
>
];
}
diff --git a/mlir/include/mlir/Transforms/Mem2Reg.h b/mlir/include/mlir/Transforms/Mem2Reg.h
index d145f7ed437582..ed10644e26a519 100644
--- a/mlir/include/mlir/Transforms/Mem2Reg.h
+++ b/mlir/include/mlir/Transforms/Mem2Reg.h
@@ -27,7 +27,7 @@ struct Mem2RegStatistics {
/// at least one memory slot was promoted.
LogicalResult
tryToPromoteMemorySlots(ArrayRef<PromotableAllocationOpInterface> allocators,
- RewriterBase &rewriter,
+ RewriterBase &rewriter, const DataLayout &dataLayout,
Mem2RegStatistics statistics = {});
} // namespace mlir
diff --git a/mlir/include/mlir/Transforms/SROA.h b/mlir/include/mlir/Transforms/SROA.h
index 1af1fe930723f1..d09a7989edeab2 100644
--- a/mlir/include/mlir/Transforms/SROA.h
+++ b/mlir/include/mlir/Transforms/SROA.h
@@ -31,7 +31,8 @@ struct SROAStatistics {
/// failure if no slot was destructured.
LogicalResult tryToDestructureMemorySlots(
ArrayRef<DestructurableAllocationOpInterface> allocators,
- RewriterBase &rewriter, SROAStatistics statistics = {});
+ RewriterBase &rewriter, const DataLayout &dataLayout,
+ SROAStatistics statistics = {});
} // namespace mlir
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index bf24194d03ddb2..a0eabc2250f7e5 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
@@ -132,7 +132,8 @@ Value LLVM::StoreOp::getStored(const MemorySlot &slot, RewriterBase &rewriter) {
bool LLVM::LoadOp::canUsesBeRemoved(
const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses) {
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
if (blockingUses.size() != 1)
return false;
Value blockingUse = (*blockingUses.begin())->get();
@@ -155,7 +156,8 @@ DeletionKind LLVM::LoadOp::removeBlockingUses(
bool LLVM::StoreOp::canUsesBeRemoved(
const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses) {
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
if (blockingUses.size() != 1)
return false;
Value blockingUse = (*blockingUses.begin())->get();
@@ -181,12 +183,14 @@ DeletionKind LLVM::StoreOp::removeBlockingUses(
}
LogicalResult LLVM::LoadOp::ensureOnlySafeAccesses(
- const MemorySlot &slot, SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
+ const MemorySlot &slot, SmallVectorImpl<MemorySlot> &mustBeSafelyUsed,
+ const DataLayout &dataLayout) {
return success(getAddr() != slot.ptr || getType() == slot.elemType);
}
LogicalResult LLVM::StoreOp::ensureOnlySafeAccesses(
- const MemorySlot &slot, SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
+ const MemorySlot &slot, SmallVectorImpl<MemorySlot> &mustBeSafelyUsed,
+ const DataLayout &dataLayout) {
return success(getAddr() != slot.ptr ||
getValue().getType() == slot.elemType);
}
@@ -206,7 +210,8 @@ static bool forwardToUsers(Operation *op,
bool LLVM::BitcastOp::canUsesBeRemoved(
const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses) {
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
return forwardToUsers(*this, newBlockingUses);
}
@@ -217,7 +222,8 @@ DeletionKind LLVM::BitcastOp::removeBlockingUses(
bool LLVM::AddrSpaceCastOp::canUsesBeRemoved(
const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses) {
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
return forwardToUsers(*this, newBlockingUses);
}
@@ -228,7 +234,8 @@ DeletionKind LLVM::AddrSpaceCastOp::removeBlockingUses(
bool LLVM::LifetimeStartOp::canUsesBeRemoved(
const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses) {
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
return true;
}
@@ -239,7 +246,8 @@ DeletionKind LLVM::LifetimeStartOp::removeBlockingUses(
bool LLVM::LifetimeEndOp::canUsesBeRemoved(
const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses) {
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
return true;
}
@@ -250,7 +258,8 @@ DeletionKind LLVM::LifetimeEndOp::removeBlockingUses(
bool LLVM::InvariantStartOp::canUsesBeRemoved(
const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses) {
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
return true;
}
@@ -261,7 +270,8 @@ DeletionKind LLVM::InvariantStartOp::removeBlockingUses(
bool LLVM::InvariantEndOp::canUsesBeRemoved(
const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses) {
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
return true;
}
@@ -272,7 +282,8 @@ DeletionKind LLVM::InvariantEndOp::removeBlockingUses(
bool LLVM::DbgDeclareOp::canUsesBeRemoved(
const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses) {
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
return true;
}
@@ -283,7 +294,8 @@ DeletionKind LLVM::DbgDeclareOp::removeBlockingUses(
bool LLVM::DbgValueOp::canUsesBeRemoved(
const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses) {
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
// There is only one operand that we can remove the use of.
if (blockingUses.size() != 1)
return false;
@@ -318,7 +330,8 @@ static bool hasAllZeroIndices(LLVM::GEPOp gepOp) {
bool LLVM::GEPOp::canUsesBeRemoved(
const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses) {
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
// GEP can be removed as long as it is a no-op and its users can be removed.
if (!hasAllZeroIndices(*this))
return false;
@@ -337,7 +350,8 @@ static bool isFirstIndexZero(LLVM::GEPOp gep) {
}
LogicalResult LLVM::GEPOp::ensureOnlySafeAccesses(
- const MemorySlot &slot, SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
+ const MemorySlot &slot, SmallVectorImpl<MemorySlot> &mustBeSafelyUsed,
+ const DataLayout &dataLayout) {
if (getBase() != slot.ptr)
return success();
if (slot.elemType != getElemType())
@@ -353,9 +367,9 @@ LogicalResult LLVM::GEPOp::ensureOnlySafeAccesses(
bool LLVM::GEPOp::canRewire(const DestructurableMemorySlot &slot,
SmallPtrSetImpl<Attribute> &usedIndices,
- SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
- auto basePtrType = llvm::dyn_cast<LLVM::LLVMPointerType>(getBase().getType());
- if (!basePtrType)
+ SmallVectorImpl<MemorySlot> &mustBeSafelyUsed,
+ const DataLayout &dataLayout) {
+ if (!isa<LLVM::LLVMPointerType>(getBase().getType()))
return false;
if (getBase() != slot.ptr || slot.elemType != getElemType())
@@ -378,7 +392,8 @@ bool LLVM::GEPOp::canRewire(const DestructurableMemorySlot &slot,
DeletionKind LLVM::GEPOp::rewire(const DestructurableMemorySlot &slot,
DenseMap<Attribute, MemorySlot> &subslots,
- RewriterBase &rewriter) {
+ RewriterBase &rewriter,
+ const DataLayout &dataLayout) {
IntegerAttr firstLevelIndex =
llvm::dyn_cast_if_present<IntegerAttr>(getIndices()[1]);
const MemorySlot &newSlot = subslots.at(firstLevelIndex);
@@ -447,7 +462,7 @@ std::optional<uint64_t> getStaticMemIntrLen(LLVM::MemcpyInlineOp op) {
/// of the bounds of the given slot, on a best-effort basis.
template <class MemIntr>
static bool definitelyWritesOnlyWithinSlot(MemIntr op, const MemorySlot &slot,
- DataLayout &dataLayout) {
+ const DataLayout &dataLayout) {
if (!isa<LLVM::LLVMPointerType>(slot.ptr.getType()) ||
op.getDst() != slot.ptr)
return false;
@@ -512,7 +527,8 @@ Value LLVM::MemsetOp::getStored(const MemorySlot &slot,
bool LLVM::MemsetOp::canUsesBeRemoved(
const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses) {
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
// TODO: Support non-integer types.
bool canConvertType =
TypeSwitch<Type, bool>(slot.elemType)
@@ -526,8 +542,7 @@ bool LLVM::MemsetOp::canUsesBeRemoved(
if (getIsVolatile())
return false;
- DataLayout layout = DataLayout::closest(*this);
- return getStaticMemIntrLen(*this) == layout.getTypeSize(slot.elemType);
+ return getStaticMemIntrLen(*this) == dataLayout.getTypeSize(slot.elemType);
}
DeletionKind LLVM::MemsetOp::removeBlockingUses(
@@ -537,14 +552,15 @@ DeletionKind LLVM::MemsetOp::removeBlockingUses(
}
LogicalResult LLVM::MemsetOp::ensureOnlySafeAccesses(
- const MemorySlot &slot, SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
- DataLayout dataLayout = DataLayout::closest(*this);
+ const MemorySlot &slot, SmallVectorImpl<MemorySlot> &mustBeSafelyUsed,
+ const DataLayout &dataLayout) {
return success(definitelyWritesOnlyWithinSlot(*this, slot, dataLayout));
}
bool LLVM::MemsetOp::canRewire(const DestructurableMemorySlot &slot,
SmallPtrSetImpl<Attribute> &usedIndices,
- SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
+ SmallVectorImpl<MemorySlot> &mustBeSafelyUsed,
+ const DataLayout &dataLayout) {
if (&slot.elemType.getDialect() != getOperation()->getDialect())
return false;
@@ -558,13 +574,13 @@ bool LLVM::MemsetOp::canRewire(const DestructurableMemorySlot &slot,
if (!areAllIndicesI32(slot))
return false;
- DataLayout dataLayout = DataLayout::closest(*this);
return definitelyWritesOnlyWithinSlot(*this, slot, dataLayout);
}
DeletionKind LLVM::MemsetOp::rewire(const DestructurableMemorySlot &slot,
DenseMap<Attribute, MemorySlot> &subslots,
- RewriterBase &rewriter) {
+ RewriterBase &rewriter,
+ const DataLayout &dataLayout) {
std::optional<DenseMap<Attribute, Type>> types =
slot.elemType.cast<DestructurableTypeInterface>().getSubelementIndexMap();
@@ -579,7 +595,6 @@ DeletionKind LLVM::MemsetOp::rewire(const DestructurableMemorySlot &slot,
packed = structType.isPacked();
Type i32 = IntegerType::get(getContext(), 32);
- DataLayout dataLayout = DataLayout::closest(*this);
uint64_t memsetLen = memsetLenAttr.getValue().getZExtValue();
uint64_t covered = 0;
for (size_t i = 0; i < types->size(); i++) {
@@ -642,7 +657,8 @@ template <class MemcpyLike>
static bool
memcpyCanUsesBeRemoved(MemcpyLike op, const MemorySlot &slot,
const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses) {
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
// If source and destination are the same, memcpy behavior is undefined and
// memmove is a no-op. Because there is no memory change happening here,
// simplifying such operations is left to canonicalization.
@@ -652,8 +668,7 @@ memcpyCanUsesBeRemoved(MemcpyLike op, const MemorySlot &slot,
if (op.getIsVolatile())
return false;
- DataLayout layout = DataLayout::closest(op);
- return getStaticMemIntrLen(op) == layout.getTypeSize(slot.elemType);
+ return getStaticMemIntrLen(op) == dataLayout.getTypeSize(slot.elemType);
}
template <class MemcpyLike>
@@ -681,7 +696,8 @@ memcpyEnsureOnlySafeAccesses(MemcpyLike op, const MemorySlot &slot,
template <class MemcpyLike>
static bool memcpyCanRewire(MemcpyLike op, const DestructurableMemorySlot &slot,
SmallPtrSetImpl<Attribute> &usedIndices,
- SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
+ SmallVectorImpl<MemorySlot> &mustBeSafelyUsed,
+ const DataLayout &dataLayout) {
if (op.getIsVolatile())
return false;
@@ -693,7 +709,6 @@ static bool memcpyCanRewire(MemcpyLike op, const DestructurableMemorySlot &slot,
return false;
// Only full copies are supported.
- DataLayout dataLayout = DataLayout::closest(op);
if (getStaticMemIntrLen(op) != dataLayout.getTypeSize(slot.elemType))
return false;
@@ -733,15 +748,13 @@ void createMemcpyLikeToReplace(RewriterBase &rewriter, const DataLayout &layout,
/// Rewires a memcpy-like operation. Only copies to or from the full slot are
/// supported.
template <class MemcpyLike>
-static DeletionKind memcpyRewire(MemcpyLike op,
- const DestructurableMemorySlot &slot,
- DenseMap<Attribute, MemorySlot> &subslots,
- RewriterBase &rewriter) {
+static DeletionKind
+memcpyRewire(MemcpyLike op, const DestructurableMemorySlot &slot,
+ DenseMap<Attribute, MemorySlot> &subslots, RewriterBase &rewriter,
+ const DataLayout &dataLayout) {
if (subslots.empty())
return DeletionKind::Delete;
- DataLayout layout = DataLayout::closest(op);
-
assert((slot.ptr == op.getDst()) != (slot.ptr == op.getSrc()));
bool isDst = slot.ptr == op.getDst();
@@ -772,7 +785,7 @@ static DeletionKind memcpyRewire(MemcpyLike op,
isDst ? op.getSrc() : op.getDst(), gepIndices);
// Then create a new memcpy out of this source pointer.
- createMemcpyLikeToReplace(rewriter, layout, op,
+ createMemcpyLikeToReplace(rewriter, dataLayout, op,
isDst ? subslot.ptr : subslotPtrInOther,
isDst ? subslotPtrInOther : subslot.ptr,
subslot.elemType, op.getIsVolatile());
@@ -798,8 +811,10 @@ Value LLVM::MemcpyOp::getStored(const MemorySlot &slot,
bool LLVM::MemcpyOp::canUsesBeRemov...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/85644
More information about the Mlir-commits
mailing list