Migrating to One-Shot Bufferization

[tzunghanjuang]

Context: MLIR removed old bufferization (Conversion-Based bufferization) passes and introduced new One-Shot-Bufferization pass. [Link] Note that Ownership-based Buffer Deallocation will be handled in a different PR.

• ~Only update llvm and mhlo.~ Jax sticks with 0.4.28.

Description of the Change: Newly Added Bufferization Interface: Comparing implemented interface methods with BufferizationOpInterface.td. Operation List: https://github.com/PennyLaneAI/catalyst/blob/8d0324db672d507d651b46fd65e9058455a9bb03/mlir/lib/Bufferization.md

Major Updates:

• Create one-shot-bufferize interface for catalyst-bufferize, gradient-bufferize, and quantum-bufferize.
• The operations to be bufferized should be registered with declarePromisedInterfaces.
• Bufferize FunctionOPInterface like ForwardOP and ReverseOp need to access OneShotModuleBufferize.
• Gradient.ReturnOP is ReturnLike now to support bufferization.
• Set function-boundary-type-conversion to identity-layout-map to avoid invalid type conversion between Memref<T> and Memref<T, Stride<...>>
• Use allow-return-allocs-from-loops flag to avoid scp.for and scp.while to avoid arg and yield types mismatch
• Async tests need copy-before-write flag

Required LLVM and MHLO Patches:

• Make OneShotModuleBufferize support FunctionOpInterface instead of FuncOP. https://github.com/llvm/llvm-project/pull/110322
• FunctionOpInterface patch to mhlo. https://github.com/openxla/xla/pull/17935
• Update MemRefType conversion in CallOp bufferization. https://github.com/llvm/llvm-project/pull/105929

Related GitHub Issues: [sc-71487] [Discourse discussion about migration] [Discord discussion about FunctionOpInterface and FuncOp] [Discord discussion about memory layout] [LLVM Upstream]

[LLVM PR1] [LLVM PR2]

https://github.com/llvm/llvm-project/commit/7f04a8ad131881b5a58b97c8191733ed42d18e20

youtube link with discussion about one-shot bufferizer pdf with info