Page MenuHomePhabricator

[mlir][memref] Add realloc op.
ClosedPublic

Authored by bixia on Sep 7 2022, 7:12 AM.

Details

Summary

Add memref.realloc and canonicalization of the op. Add conversion patterns for
lowering the op to LLVM using unaligned alloc or aligned alloc based on the
conversion option.

Add filecheck tests for parsing and converting the op. Add an integration test.

Diff Detail

Event Timeline

bixia created this revision.Sep 7 2022, 7:12 AM
bixia requested review of this revision.Sep 7 2022, 7:12 AM
bixia retitled this revision from [mlir][memref] Add realloc op. to WIP [mlir][memref] Add realloc op..Sep 7 2022, 7:15 AM
bixia retitled this revision from WIP [mlir][memref] Add realloc op. to [mlir][memref] Add realloc op..Sep 7 2022, 7:33 AM
aartbik added inline comments.Sep 7 2022, 9:51 AM
mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
200

maybe add a comment that referencing memory through the "old" SSA value after the realloc is undefined behavior, perhaps with an example:

%new = memref.realloc(%old) ....
 
%0 = memref.load %new[%c]   // ok
%1 = memref.load %old[%c]   // UB
Hardcode84 added inline comments.
mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
2078

Please, do not add Dealloc here, it was specifically moved to these ifs, so it can be skipped completely

bixia updated this revision to Diff 458765.Sep 8 2022, 8:53 AM

Add description about referencing memref via the old SSA value.
Revert the change that hoist the adding DeallocOpLowering.

bixia added inline comments.Sep 8 2022, 8:55 AM
mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
200

Thanks for the suggestion! Add this to the end of the description.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
2078

Thanks for catching this! I reverted this change.

bixia updated this revision to Diff 459813.Sep 13 2022, 11:03 AM

Fix code format.

bixia updated this revision to Diff 459922.Sep 13 2022, 5:06 PM

Avoid reallocation when new_size <= old_size.

bixia updated this revision to Diff 460532.Sep 15 2022, 4:12 PM

Restore the description about referencing memref via the old SSA value.

bixia updated this revision to Diff 460834.Sep 16 2022, 11:30 AM

Add a build method.

ftynse requested changes to this revision.Sep 16 2022, 11:37 PM

High-level comment: since the lowering goes to alloc/memcpy/free, can't you have a pattern that rewrites realloc to memref.alloc/copy/dealloc instead of going directly to LLVM with a lot of code. The rewriter infrastructure should be able to handle that.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
266

Could we please make the syntax of this op consistent with that of memref.alloc(a). Both take the dynamic size in parentheses.

mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp
53

Nit: this name confused me several times. I read it as emitting the IR that will produce a non-aligned allocation. But in fact the allocation is still aligned! It's just modified manually.

64

Nit: please expand auto unless the type is obvious from context (cast on the RHS) or difficult to spell (iterators).

65

Nit: I'd rather get the first parent with the SymbolTable trait. Allocations may well be in other module-like operations than the built-in module.

164

Please don't remove the newline.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
103–257

///-document here what this does, specifically that it is not intended to be used directly + that derived classes must implement a method with a specific contract.

104

Nit: I'd rather suffix this with Base than with Impl.

141

Pass the callback by function_ref instead of std::function if it is not stored inside. The latter allocates on heap.

165

Use cast instead of dyn_cast if you never check the result for being null.

173

Direct IR modifications are not allowed in conversion patterns. This argument has to be passed to createBlock of the rewriter. Since the block is originally obtained by splitting, you can create a new block with the argument before the split-out pattern, and then merge that part into the newly created block with the argument.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
259

Please don't refer to layout map anymore, layouts have not been only maps for over a year. It's just a layout.

This revision now requires changes to proceed.Sep 16 2022, 11:37 PM
bixia updated this revision to Diff 461231.Sep 19 2022, 9:23 AM
bixia marked 10 inline comments as done.

Address review comments.

bixia added inline comments.Sep 19 2022, 9:29 AM
mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp
53

Rename this to allocateBufferManuallyAlign.

64

Thanks! Fix here and other places.

65

Sorry, I don't quite understand what you said. Can you give an example for how the code may look like?
We are trying to create a forward declaration of the allocation function so that we can use the function in the module.
Here is the relevant part of the original code we are trying to rewrite in this PR.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
103–257

Add document, please take another look.

173

Thanks!

ftynse added inline comments.Sep 20 2022, 6:14 AM
mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp
65

Instead of doing getParentOfType<ModuleOp>, it should be doing getParentWithTrait<SymbolTable>. MLIR offers no guarantees that functions are always contained in a ModuleOp. FWIW, in a GPU-targeting compilation, they will be contained in a GPUModuleOp.

Have you seen my top-level comment in the original review? That one is a blocker, the rest are tidying up.

bixia added a comment.Sep 20 2022, 7:38 AM

Have you seen my top-level comment in the original review? That one is a blocker, the rest are tidying up.

Sorry I forgot to respond to that. I can't just use alloc/copy/dealloc because copy requires the two memref to have the same element type and same number of elements.
I may use "view", to create a "smaller view of the bigger memref" (if the realloc buffer is bigger, create a view of it. If the old buffer is bigger, create a view of it). However,
"view" requires the source memref to have i8 type, and memref.cast is not for this kind of cases. So I am kind of stuck here.

bixia added a comment.Sep 20 2022, 7:46 AM

Have you seen my top-level comment in the original review? That one is a blocker, the rest are tidying up.

Sorry I forgot to respond to that. I can't just use alloc/copy/dealloc because copy requires the two memref to have the same element type and same number of elements.
I may use "view", to create a "smaller view of the bigger memref" (if the realloc buffer is bigger, create a view of it. If the old buffer is bigger, create a view of it). However,
"view" requires the source memref to have i8 type, and memref.cast is not for this kind of cases. So I am kind of stuck here.

Also the direct lowering implementation avoids the alloc/copy/dealloc when the new size is smaller.

ftynse accepted this revision.Sep 21 2022, 3:44 AM
ftynse added inline comments.
mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
177

Nit: cast here instead of dyn_cast.

mlir/test/Conversion/MemRefToLLVM/convert-dynamic-memref-ops.mlir
634–659

Overall suggestion for tests: drop the trailing types from operations except when they are meaningful for comprehension. For example, bitcast and the likes should keep the types because it's important to see what is being casted to what, but mul/add/extract value don't really need types. This will make the tests more robust to eventual syntactic changes that are in fact irrelevant to what is being tested.

This revision is now accepted and ready to land.Sep 21 2022, 3:44 AM
bixia updated this revision to Diff 461887.Sep 21 2022, 7:05 AM
bixia marked 2 inline comments as done.

Replace a dyn_cast with cast. Simplify FileCheck tests by removing unnecessary checking of trailing types.

bixia added inline comments.Sep 21 2022, 7:06 AM
mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp
65

Per offline conversation, add a TODO to fix these utility functions before we can use getParentWithTrait here.

mlir/test/Conversion/MemRefToLLVM/convert-dynamic-memref-ops.mlir
634–659

Thanks for the suggestion, it also help simplify the tests beside making them more robust.

bixia updated this revision to Diff 461889.Sep 21 2022, 7:25 AM

Revert and unintentional change to a test.

This revision was landed with ongoing or failed builds.Sep 21 2022, 8:04 AM
This revision was automatically updated to reflect the committed changes.