This is an archive of the discontinued LLVM Phabricator instance.

[mlir] Add a generic SROA implementation.
ClosedPublic

Authored by Moxinilian on May 9 2023, 2:28 AM.

Details

Summary

This revision introduces a generic implementation of Scalar Replacement
Of Aggregates. In contrast to the implementation in LLVM, this focuses
on the core of SROA: destructuring aggregates. By implementing
interfaces on allocators and accessors, memory allocators can be
destructured into smaller allocators, through the MemorySlot
abstraction.

This pass only works on aggregates that are accessed in a "type-safe"
way, that is within the bounds and respecting the type of a given memory
slot. The destructuring pattern and functions only peel off the first
layer of aggregates and can safely be applied repeatedly. For
convenience, the transformation is also available as a pass that will
apply the pattern repeatedly.

Depends on D149958

Diff Detail

Event Timeline

Moxinilian created this revision.May 9 2023, 2:28 AM
Herald added a reviewer: dcaballe. · View Herald Transcript
Herald added a project: Restricted Project. · View Herald Transcript
Moxinilian requested review of this revision.May 9 2023, 2:29 AM

Thanks. I am interested in the pass and of course mem2reg as a pass. My only annoyance is that I have to run the pass several times.

I added a bunch of comments. The pass already looks great.

Make sure to update the commit message, as the pass applies the pattern until convergence, not just once.

mlir/include/mlir/Interfaces/MemorySlotInterfaces.h
28

Ultra NIT: I prefer explicit public inheritance

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
263

This is a very general name for an interface. We should consider to rename this to something that ties it to the MemorySlot abstraction.

351

Only allowing to use attributes to index into the list of types seems a bit cumbersome for other usages.
Adding a second InterfaceMethod that allows you to access this with an unsigned might be nice to have.

The function with an attribute parameter could then even have a default implementation that dispatches to the other one.

mlir/include/mlir/Transforms/SROA.h
19

NIT: Top-level comment.

28

NIT: llvm::Optional is deprecated.

mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
72

NIT: destructuredTypes or destructuredTypeMap

282

NIT: Top level comment

294

std::optional

300

NIT: Drop auto, as the type is not directly clear.

343

According to the GEPOp definition, the type of base can also be a vector. In that case, the cast would assert.
Maybe adding a test for this case would also be nice.

395

std::optional

399

Ultra NIT: Usually, const is used in front of the type of a reference, i.e., const auto &. There is no semantic difference, but it is what is mainly used in the MLIR codebase.

mlir/lib/Transforms/SROA.cpp
26

std::optional

35

NIT: Can you use the SmallPtrSetImpl type here?

157

NIT: use notifyMatchFailure() instead.

177

IIRC, the PatternRewriter class inherits from OpBuilder. If there is a reason why you cannot use the rewriter, comment it.

Moxinilian marked 15 inline comments as done.May 10 2023, 2:39 AM

Thank you for the review!

@tschuett Sorry, that was not clear. While the pattern only peels off one layer of aggregate (as it would typically be used recursively in pattern application), the pass applies it repeatedly. I updated the summary accordingly.

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
263

I went with SafeMemorySlotAccessOpInterface, does that sound good?

351

I'm worried about this. LLVM is so far the only use case I need to implement that uses integers directly. In MemRef you want to use multi-dimensional indices, not just integers. And then in an hypothetical hashmap dialect you may want to use strings or any attribute really. Blessing the "integer-only" use case feels very ad hoc to me, especially considering implementers will have to handle Attribute anyway because they may accidentally receive or want to process Attribute indices no matter what. Could you be more specific of the approach you want for this?

mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
72

So the intent was that this represents the type in its destructured form. I went for subelementIndexMap then.

343

Great catch. For the test, I am not sure, considering this only gets invoked currently when the base is the result of an alloca, which must be a pointer. Any idea?

Moxinilian marked 3 inline comments as done.
Moxinilian edited the summary of this revision. (Show Details)

Address comments.

Thanks. Then I will take the pass.

Improve destructuring algorithm in AllocaOp.

Very nice work!

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h
205

There are still at least two uses of Optional. Please search&replace with std::optional.

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
231
249

nit: I assume usedIndices an input only? Could it be const and or documented?

mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
424

Should size limit be a pass parameter and also apply to structs? This maybe part of a followup revision.

mlir/lib/Transforms/SROA.cpp
17

nit: Can you check the includes? I don't think there TypeSwitch is used in this file?

125

nit: Drop llvm:: here and in llvm::reverse below (I believe).

146

I would define a SROA DEBUG_TYPE as in the mem2reg revision.

172

nit: Let's use *info to access optional values.

Dinistro added inline comments.May 10 2023, 11:23 PM
mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
263

Sounds good to me.

351

Oh, I didn't realize that the index can be multidimensional. In that case it can probably stay an attribute. Just be aware that attributes are always uniqued and aren't cleaned up, even if they are no longer used.

mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
343

In that case you cannot add a test, I guess.

Rebase, add statistics and improve pattern.

Moxinilian marked 12 inline comments as done.

Address comments.

mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
424

Yes probably. The ideal maximum size would be +inf, the limit being only for practical reasons. I think it is fine that the limit does not apply to structs for example, as you cannot create a struct with a trillion fields as easily as an array with a trillion elements. Let's save that for a follow up if it proves useful, though.

mlir/lib/Transforms/SROA.cpp
125

So I kind of prefer having llvm:: in front of this sort of utilities, like topological sort, to signify that this comes from the infra and is not a local concept (so the reader does not assume enumerate has some form of contextual local logic). Does that sound good to you? Otherwise I can drop them, it's fine.

Nice we are almost there! I have some nit comment but am otherwise happy!

The build bot seems to be unable to apply the patch. I assume a rebase should fix this.

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
243–244

nit: I would try to make the last sentence a bit more direct, e.g., "Only generate subslots for the indices found in usedIndices since all other subslots are unused"

mlir/include/mlir/Transforms/Passes.td
251

nit: Total number of memory slots (since they seem countable?)

Note this appears in several places!

256

And if possible try to make the lines shorter to stay within the 80 columns. I expect [{}] should work.

261

In think I would also use number / Maximal number?

mlir/lib/Transforms/SROA.cpp
192–193

nit: Would it make sense to move this check into computeDestructuringInfo to make it more self contained.

197

nit: I would suggest to inverse the condition to save on indent in the body of the current if, i.e.:

if (!info)
  continue;
destructureSlot(slot, allocator, rewriter, *info, statistics);
destructuredAny = true;
mlir/test/Dialect/LLVMIR/sroa.mlir
184

nit: I would stick to no_dynamic_indexing. The double negation is a bit hard to parse.

Moxinilian marked 7 inline comments as done.

Address comments.

mlir/include/mlir/Transforms/Passes.td
256

I tried! Unfortunately it does not escape carriage returns and breaks the generated code entirely...

Fix tests + improve code layout.

Fix forgotten s. Sorry!

Fix GEPIndidicesAdaptor bug in release mode.

gysit accepted this revision.May 22 2023, 12:46 AM

LGTM!

Do you want to try another rebase. Maybe the CI then turns green assuming the underlying problem has been fixed in meantime?

This revision is now accepted and ready to land.May 22 2023, 12:46 AM

Refresh commit for CI.

This revision was automatically updated to reflect the committed changes.