catalyst icon indicating copy to clipboard operation
catalyst copied to clipboard

Migrate to new buffer deallocation pipeline

Open paul0403 opened this issue 6 months ago • 1 comments

As part of the llvm update, the old --buffer-deallocation pass is removed.

Intended replacement is --buffer-deallocation-pipeline. llvm/llvm-project#126366, https://discourse.llvm.org/t/psa-bufferization-new-buffer-deallocation-pipeline/73375. However, we encountered a few issues. See detailed discord discussion at https://discord.com/channels/636084430946959380/642426447167881246/1376919538301403276

The heart of the new deallocation pass, --ownership-based-buffer-deallocation, blanket-ly fails for ops with unknown memory effects: https://github.com/llvm/llvm-project/blob/da4958ae2b384c2a027cf20c67b7e211d39fcbfe/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp#L523 To solve this issue, the suggestion was to add memory effects to custom operations. The migration guide also suggests that bufferizable ops no longer implement the bufferizesToAllocation method, so we remove them.

This was supposed to be done alongside the llvm update in https://github.com/PennyLaneAI/catalyst/pull/1752; However, soon it became clear that this migration to the new buffer deallocation is very complicated, and should be its own story.

The llvm update in #1752 thus did not finish this migration. This PR records the work that was already done on this back in #1752.


List of memory effects added to the ops back in #1752 :

quantum dialect: quantum.init: MemAlloc<QuantumMemory> quantum.finalize: MemFree<QuantumMemory> quantum.device: MemAlloc<QuantumMemory> quantum.device_release: MemFree<QuantumMemory> (!) quantum.alloc already has a NoMemoryEffect. I left it alone. quantum.dealloc: MemFree<QuantumMemory> quantum.set_state and set_basis_state: MemWrite<QuantumMemory> quantum.custom, multi_rz and unitary: MemRead<QuantumMemory> Originally I had both read and write here, but this caused passes like cancel inverses to fail to remove some gates without uses after pattern rewrite. quantum.gphase: MemRead<QuantumMemory>, MemWrite<QuantumMemory> quantum.adjoint: MemRead<QuantumMemory>, MemWrite<QuantumMemory> quantum.compbasis and namedobs: MemRead<QuantumMemory> quantum.measure: MemRead<QuantumMemory>, MemWrite<QuantumMemory> quantum.sample, counts, state, probs: MemAlloc, MemRead<QuantumMemory> These need to allocate a classical array to store the results quantum.expval, var: MemRead<QuantumMemory>

mbqc dialect: mbqc.measure_in_basis: MemRead<QuantumMemory>, MemWrite<QuantumMemory>

catalyst dialect: catalyst.list_init: MemAlloc catalyst.list_dealloc: MemFree catalyst.list_push: MemWrite, MemAlloc catalyst.list_pop: MemRead, MemFree catalyst.list_load_data: MemRead catalyst.print: MemRead, MemAlloc (alloc is needed to print strings) catalyst.assert: MemWrite (write assertion error message)

gradient dialect: I am not super sure regarding them, so I just turn on all effects to be safe. gradient.adjiont: MemRead, MemWrite, MemAlloc, MemFree gradient.backprop: MemRead, MemWrite, MemAlloc, MemFree gradient.jvp: MemRead, MemWrite, MemAlloc, MemFree gradient.vjp: MemRead, MemWrite, MemAlloc, MemFree

Note that these are not tested/thoroughly considered. See conversations in #1752 for some of the initial discussions on these memory effects.

paul0403 avatar May 30 '25 22:05 paul0403

I am opening this PR simply to record the WIP progress of what we have so far for the new dealloc. This can be beneficial when we revisit.

The huge diff is just from #1752 itself, and will go away after it's merged.

Please leave this PR in draft to save CI. Thanks!

paul0403 avatar May 30 '25 22:05 paul0403