Skip to content

[MLIR][SROA][Mem2Reg] Add data layout to interface methods #85644

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 20, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 22 additions & 21 deletions mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
Original file line number Diff line number Diff line change
Expand Up @@ -83,11 +83,10 @@ 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.
to memory slots.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Doesn't this remain true in the context of the PromotableMemOpInterface used by the Mem2Reg pass?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's not yet fully clear if this requirement will remain valid for LLVM dialect. I can remove this change and adapt it when it shows to no longer hold.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After some more though, this is really no longer needed. Changing the LLVM load and store interface implementations to support loading an smaller type from a larger slot should be easy enough.
We are in fact considering such changes, as the type consistency pass showed to produce substantial performance degradations for parts of our benchmark set.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What about loading larger types from smaller slots?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not supported, as this would break things on multiple levels. See the stacked PR for the constrains added specifically in SROA. The additional steps for Mem2Reg will be added in subsequent patches.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don’t understand, I would say it is obviously not supported, hence why the original documentation was there. I’m not sure where else the constraints on bounding are documented if you remove this.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can clarify in the comment that only in-bounds accesses should be considered valid.


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*.
For a memory operation on a slot to be valid, it must strictly operate
within the bounds of the slot.

If the same operation does both loads and stores on the same slot, the
load must semantically happen first.
Expand Down Expand Up @@ -142,7 +141,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
Expand Down Expand Up @@ -197,7 +197,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
Expand Down Expand Up @@ -285,29 +286,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)
>
];
}
Expand All @@ -323,21 +323,21 @@ 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.
}],
"bool",
"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
Expand All @@ -351,7 +351,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)
>
];
}
Expand Down
2 changes: 1 addition & 1 deletion mlir/include/mlir/Transforms/Mem2Reg.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion mlir/include/mlir/Transforms/SROA.h
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
Loading