diff --git a/frontend/catalyst/pipelines.py b/frontend/catalyst/pipelines.py index 67a6e4ef45..6348e1eb87 100644 --- a/frontend/catalyst/pipelines.py +++ b/frontend/catalyst/pipelines.py @@ -240,8 +240,7 @@ def get_bufferization_stage(options: CompileOptions) -> List[str]: "func.func(buffer-hoisting)", "func.func(buffer-loop-hoisting)", "func.func(promote-buffers-to-stack)", - # TODO: migrate to new buffer deallocation "buffer-deallocation-pipeline" - "func.func(buffer-deallocation)", + "buffer-deallocation-pipeline", "convert-arraylist-to-memref", "convert-bufferization-to-memref", "canonicalize", # Must be after convert-bufferization-to-memref diff --git a/mlir/include/Catalyst/IR/CatalystOps.td b/mlir/include/Catalyst/IR/CatalystOps.td index 69ab0f54e5..99e387bf99 100644 --- a/mlir/include/Catalyst/IR/CatalystOps.td +++ b/mlir/include/Catalyst/IR/CatalystOps.td @@ -26,13 +26,13 @@ include "mlir/IR/OpBase.td" include "Catalyst/IR/CatalystDialect.td" -def ListInitOp : Catalyst_Op<"list_init"> { +def ListInitOp : Catalyst_Op<"list_init", [MemoryEffects<[MemAlloc]>]> { let summary = "Initialize a dynamically resizable arraylist."; let results = (outs ArrayListType:$list); let assemblyFormat = [{ attr-dict `:` type($list) }]; } -def ListDeallocOp : Catalyst_Op<"list_dealloc"> { +def ListDeallocOp : Catalyst_Op<"list_dealloc", [MemoryEffects<[MemFree]>]> { let summary = "Deallocate the underlying memory of an arraylist."; let arguments = (ins ArrayListType:$list); let assemblyFormat = [{ $list attr-dict `:` type($list) }]; @@ -41,7 +41,8 @@ def ListDeallocOp : Catalyst_Op<"list_dealloc"> { def ListPushOp : Catalyst_Op<"list_push", [TypesMatchWith<"type of 'value' matches element type of 'list'", "list", "value", - "mlir::cast($_self).getElementType()">]> { + "mlir::cast($_self).getElementType()">, + MemoryEffects<[MemWrite, MemAlloc]>]> { let summary = "Append an element to the end of an array list."; let arguments = (ins AnyType:$value, ArrayListType:$list); let assemblyFormat = [{ $value `,` $list attr-dict `:` type($list) }]; @@ -50,14 +51,15 @@ def ListPushOp : Catalyst_Op<"list_push", def ListPopOp : Catalyst_Op<"list_pop", [TypesMatchWith<"type of 'result' matches element type of 'list'", "list", "result", - "mlir::cast($_self).getElementType()">]> { + "mlir::cast($_self).getElementType()">, + MemoryEffects<[MemRead, MemFree]>]> { let summary = "Remove an element from the end of an array list and return it."; let arguments = (ins ArrayListType:$list); let results = (outs AnyType:$result); let assemblyFormat = [{ $list attr-dict `:` type($list) }]; } -def ListLoadDataOp : Catalyst_Op<"list_load_data"> { +def ListLoadDataOp : Catalyst_Op<"list_load_data", [MemoryEffects<[MemRead]>]> { let summary = "Get the underlying memref storing the data of an array list."; let arguments = (ins ArrayListType:$list); let results = (outs AnyMemRef:$data); @@ -71,7 +73,7 @@ def ListLoadDataOp : Catalyst_Op<"list_load_data"> { let assemblyFormat = [{ $list attr-dict `:` type($list) `->` type($data) }]; } -def PrintOp : Catalyst_Op<"print"> { +def PrintOp : Catalyst_Op<"print", [MemoryEffects<[MemRead, MemWrite]>]> { let summary = "Prints numeric values or constant strings at runtime."; let arguments = (ins @@ -81,7 +83,7 @@ def PrintOp : Catalyst_Op<"print"> { ); } -def AssertionOp : Catalyst_Op<"assert"> { +def AssertionOp : Catalyst_Op<"assert", [MemoryEffects<[MemWrite]>]> { let summary = "Asserts condition at runtime."; let arguments = (ins diff --git a/mlir/include/Gradient/IR/GradientOps.td b/mlir/include/Gradient/IR/GradientOps.td index 8931e109e0..cbb985ab23 100644 --- a/mlir/include/Gradient/IR/GradientOps.td +++ b/mlir/include/Gradient/IR/GradientOps.td @@ -105,7 +105,8 @@ def ValueAndGradOp : Gradient_Op<"value_and_grad", [ } -def AdjointOp : Gradient_Op<"adjoint", [AttrSizedOperandSegments]> { +def AdjointOp : Gradient_Op<"adjoint", [AttrSizedOperandSegments, + MemoryEffects<[MemRead, MemWrite, MemAlloc, MemFree]>]> { let summary = "Perform quantum AD using the adjoint method on a device."; let arguments = (ins @@ -134,7 +135,8 @@ def AdjointOp : Gradient_Op<"adjoint", [AttrSizedOperandSegments]> { def BackpropOp : Gradient_Op<"backprop", [ AttrSizedOperandSegments, AttrSizedResultSegments, - DeclareOpInterfaceMethods + DeclareOpInterfaceMethods, + MemoryEffects<[MemRead, MemWrite, MemAlloc, MemFree]> ]> { let summary = "Perform classic automatic differentiation using Enzyme AD."; @@ -179,7 +181,8 @@ def JVPOp : Gradient_Op<"jvp", [ SameVariadicResultSize, DeclareOpInterfaceMethods, DeclareOpInterfaceMethods, - GradientOpInterface + GradientOpInterface, + MemoryEffects<[MemRead, MemWrite, MemAlloc, MemFree]> ]> { let summary = "Compute the jvp of a function."; @@ -214,7 +217,8 @@ def VJPOp : Gradient_Op<"vjp", [ AttrSizedResultSegments, DeclareOpInterfaceMethods, DeclareOpInterfaceMethods, - GradientOpInterface + GradientOpInterface, + MemoryEffects<[MemRead, MemWrite, MemAlloc, MemFree]> ]> { let summary = "Compute the vjp of a function."; diff --git a/mlir/include/MBQC/IR/MBQCOps.td b/mlir/include/MBQC/IR/MBQCOps.td index 33ed1b2243..26f93017a0 100644 --- a/mlir/include/MBQC/IR/MBQCOps.td +++ b/mlir/include/MBQC/IR/MBQCOps.td @@ -54,7 +54,11 @@ def MeasurementPlaneAttr : EnumAttr { +def MeasureInBasisOp : MBQC_Op<"measure_in_basis", + // Measurement both read and write since it collapses states + // In fact, MBQC's entire point is to alter the state via measurement. + [MemoryEffects<[MemRead, MemWrite]>] + > { let summary = "A parametric single-qubit projective measurement in an arbitrary basis."; let description = [{ A parametric single-qubit projective measurement is equivalent to the `quantum.measure` diff --git a/mlir/include/Quantum/IR/QuantumOps.td b/mlir/include/Quantum/IR/QuantumOps.td index f1fcf50c8e..50e85fe39f 100644 --- a/mlir/include/Quantum/IR/QuantumOps.td +++ b/mlir/include/Quantum/IR/QuantumOps.td @@ -60,7 +60,7 @@ def NamedObservableAttr : EnumAttr traits = []> : Op; -def InitializeOp : Quantum_Op<"init"> { +def InitializeOp : Quantum_Op<"init", [MemoryEffects<[MemAlloc]>]> { let summary = "Initialize the quantum runtime."; let assemblyFormat = [{ @@ -68,7 +68,7 @@ def InitializeOp : Quantum_Op<"init"> { }]; } -def FinalizeOp : Quantum_Op<"finalize"> { +def FinalizeOp : Quantum_Op<"finalize", [MemoryEffects<[MemFree]>]> { let summary = "Teardown the quantum runtime."; let assemblyFormat = [{ @@ -76,7 +76,7 @@ def FinalizeOp : Quantum_Op<"finalize"> { }]; } -def DeviceInitOp : Quantum_Op<"device"> { +def DeviceInitOp : Quantum_Op<"device", [MemoryEffects<[MemAlloc]>]> { let summary = "Initialize a quantum device."; let arguments = (ins @@ -92,7 +92,7 @@ def DeviceInitOp : Quantum_Op<"device"> { } -def DeviceReleaseOp : Quantum_Op<"device_release"> { +def DeviceReleaseOp : Quantum_Op<"device_release", [MemoryEffects<[MemFree]>]> { let summary = "Release the active quantum device."; let assemblyFormat = [{ @@ -169,7 +169,7 @@ def DeallocQubitOp : Memory_Op<"dealloc_qb"> { }]; } -def DeallocOp : Memory_Op<"dealloc"> { +def DeallocOp : Memory_Op<"dealloc", [MemoryEffects<[MemFree]>]> { let summary = "Deallocate a quantum register."; let description = [{ }]; @@ -338,7 +338,8 @@ class UnitaryGate_Op traits = []> : let extraClassDeclaration = extraBaseClassDeclaration; } -def SetStateOp : Gate_Op<"set_state"> { +def SetStateOp : Gate_Op<"set_state", + [MemoryEffects<[MemWrite]>]> { let summary = "Set state to a complex vector."; let description = [{ This operation is useful for simulators implementing state preparation. @@ -369,7 +370,8 @@ def SetStateOp : Gate_Op<"set_state"> { } -def SetBasisStateOp : Gate_Op<"set_basis_state"> { +def SetBasisStateOp : Gate_Op<"set_basis_state", + [MemoryEffects<[MemWrite]>]> { let summary = "Set basis state."; let description = [{ This operation is useful for simulators implementing set basis state. @@ -402,7 +404,9 @@ def SetBasisStateOp : Gate_Op<"set_basis_state"> { } def CustomOp : UnitaryGate_Op<"custom", [DifferentiableGate, NoMemoryEffect, - AttrSizedOperandSegments, AttrSizedResultSegments]> { + AttrSizedOperandSegments, AttrSizedResultSegments, + MemoryEffects<[MemRead]> + ]> { let summary = "A generic quantum gate on n qubits with m floating point parameters."; let description = [{ }]; @@ -525,7 +529,8 @@ def CustomOp : UnitaryGate_Op<"custom", [DifferentiableGate, NoMemoryEffect, let hasCanonicalizeMethod = 1; } -def GlobalPhaseOp : UnitaryGate_Op<"gphase", [DifferentiableGate, AttrSizedOperandSegments]> { +def GlobalPhaseOp : UnitaryGate_Op<"gphase", [DifferentiableGate, AttrSizedOperandSegments, + MemoryEffects<[MemRead, MemWrite]>]> { let summary = "Global Phase."; let description = [{ }]; @@ -564,7 +569,8 @@ def GlobalPhaseOp : UnitaryGate_Op<"gphase", [DifferentiableGate, AttrSizedOpera } def MultiRZOp : UnitaryGate_Op<"multirz", [DifferentiableGate, NoMemoryEffect, - AttrSizedOperandSegments, AttrSizedResultSegments]> { + AttrSizedOperandSegments, AttrSizedResultSegments, + MemoryEffects<[MemRead]>]> { let summary = "Apply an arbitrary multi Z rotation"; let description = [{ The `quantum.multirz` operation applies an arbitrary multi Z rotation to the state-vector. @@ -598,7 +604,8 @@ def MultiRZOp : UnitaryGate_Op<"multirz", [DifferentiableGate, NoMemoryEffect, } def QubitUnitaryOp : UnitaryGate_Op<"unitary", [ParametrizedGate, NoMemoryEffect, - AttrSizedOperandSegments, AttrSizedResultSegments]> { + AttrSizedOperandSegments, AttrSizedResultSegments, + MemoryEffects<[MemRead]>]> { let summary = "Apply an arbitrary fixed unitary matrix"; let description = [{ The `quantum.unitary` operation applies an arbitrary fixed unitary matrix to the @@ -639,7 +646,8 @@ def QubitUnitaryOp : UnitaryGate_Op<"unitary", [ParametrizedGate, NoMemoryEffect class Region_Op traits = []> : Quantum_Op; -def AdjointOp : Region_Op<"adjoint", [QuantumRegion, SingleBlockImplicitTerminator<"YieldOp">]> { +def AdjointOp : Region_Op<"adjoint", [QuantumRegion, SingleBlockImplicitTerminator<"YieldOp">, + MemoryEffects<[MemRead, MemWrite]>]> { let summary = "Calculate the adjoint of the enclosed operations"; let regions = (region SizedRegion<1>:$region); @@ -680,7 +688,8 @@ def YieldOp : Quantum_Op<"yield", [Pure, ReturnLike, Terminator, ParentOneOf<["A class Observable_Op traits = []> : Quantum_Op; -def ComputationalBasisOp : Observable_Op<"compbasis", [AttrSizedOperandSegments]> { +def ComputationalBasisOp : Observable_Op<"compbasis", [AttrSizedOperandSegments, + MemoryEffects<[MemRead]>]> { let summary = "Define a pseudo-obeservable of the computational basis for use in measurements"; let description = [{ The `quantum.compbasis` operation defines a quantum observable to be used by other @@ -718,7 +727,7 @@ def ComputationalBasisOp : Observable_Op<"compbasis", [AttrSizedOperandSegments] let hasVerifier = 1; } -def NamedObsOp : Observable_Op<"namedobs"> { +def NamedObsOp : Observable_Op<"namedobs", [MemoryEffects<[MemRead]>]> { let summary = "Define a Named observable for use in measurements"; let description = [{ The `quantum.namedobs` operation defines a quantum observable to be used by measurement @@ -853,7 +862,12 @@ def HamiltonianOp : Observable_Op<"hamiltonian"> { class Measurement_Op traits = []> : Quantum_Op; -def MeasureOp : Quantum_Op<"measure"> { +def MeasureOp : Quantum_Op<"measure", + // Measurement both read and write since it collapses states + // This is even more true for the mbqc variant, since that one's entire point + // is to alter the state via measurement. + [MemoryEffects<[MemRead, MemWrite]>] + > { let summary = "A single-qubit projective measurement in the computational basis."; let description = [{ }]; @@ -873,7 +887,8 @@ def MeasureOp : Quantum_Op<"measure"> { }]; } -def SampleOp : Measurement_Op<"sample", [AttrSizedOperandSegments]> { +def SampleOp : Measurement_Op<"sample", [AttrSizedOperandSegments, + MemoryEffects<[MemAlloc, MemRead]>]> { let summary = "Sample eigenvalues from the given observable for the current state"; let description = [{ The `quantum.sample` operation represents the measurement process of sampling eigenvalues @@ -939,7 +954,8 @@ def SampleOp : Measurement_Op<"sample", [AttrSizedOperandSegments]> { let hasVerifier = 1; } -def CountsOp : Measurement_Op<"counts", [AttrSizedOperandSegments, SameVariadicResultSize]> { +def CountsOp : Measurement_Op<"counts", [AttrSizedOperandSegments, SameVariadicResultSize, + MemoryEffects<[MemAlloc, MemRead]>]> { let summary = "Compute sample counts for the given observable for the current state"; let description = [{ The `quantum.counts` operation represents the measurement process of sampling eigenvalues @@ -998,7 +1014,7 @@ def CountsOp : Measurement_Op<"counts", [AttrSizedOperandSegments, SameVariadicR let hasVerifier = 1; } -def ExpvalOp : Measurement_Op<"expval"> { +def ExpvalOp : Measurement_Op<"expval", [MemoryEffects<[MemRead]>]> { let summary = "Compute the expectation value of the given observable for the current state"; let description = [{ The `quantum.expval` operation represents the measurement process of computing the @@ -1035,7 +1051,7 @@ def ExpvalOp : Measurement_Op<"expval"> { }]; } -def VarianceOp : Measurement_Op<"var"> { +def VarianceOp : Measurement_Op<"var", [MemoryEffects<[MemRead]>]> { let summary = "Compute the variance of the given observable for the current state"; let description = [{ The `quantum.var` operation represents the measurement process of computing the variance of @@ -1071,7 +1087,8 @@ def VarianceOp : Measurement_Op<"var"> { }]; } -def ProbsOp : Measurement_Op<"probs", [AttrSizedOperandSegments]> { +def ProbsOp : Measurement_Op<"probs", [AttrSizedOperandSegments, + MemoryEffects<[MemAlloc, MemRead]>]> { let summary = "Compute computational basis probabilities for the current state"; let description = [{ The `quantum.probs` operation represents the measurement process of computing probabilities @@ -1117,7 +1134,8 @@ def ProbsOp : Measurement_Op<"probs", [AttrSizedOperandSegments]> { let hasVerifier = 1; } -def StateOp : Measurement_Op<"state", [AttrSizedOperandSegments]> { +def StateOp : Measurement_Op<"state", [AttrSizedOperandSegments, + MemoryEffects<[MemAlloc, MemRead]>]> { let summary = "Return the current statevector"; let description = [{ The `quantum.state` operation represents the measurement process of returning the current diff --git a/mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp index c5b93d1602..239bbcacb8 100644 --- a/mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp +++ b/mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp @@ -75,8 +75,6 @@ struct PrintOpInterface struct CustomCallOpInterface : public bufferization::BufferizableOpInterface::ExternalModel { - bool bufferizesToAllocation(Operation *op, Value value) const { return true; } - bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, const bufferization::AnalysisState &state) const { @@ -251,8 +249,6 @@ void convertTypes(SmallVector inTypes, SmallVector &convertedResults struct CallbackCallOpInterface : public bufferization::BufferizableOpInterface::ExternalModel { - bool bufferizesToAllocation(Operation *op, Value value) const { return true; } - bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, const bufferization::AnalysisState &state) const { diff --git a/mlir/lib/Driver/Pipelines.cpp b/mlir/lib/Driver/Pipelines.cpp index 46f6905a66..63ce765476 100644 --- a/mlir/lib/Driver/Pipelines.cpp +++ b/mlir/lib/Driver/Pipelines.cpp @@ -86,8 +86,8 @@ void createBufferizationPipeline(OpPassManager &pm) pm.addNestedPass(mlir::bufferization::createBufferHoistingPass()); pm.addNestedPass(mlir::bufferization::createBufferLoopHoistingPass()); pm.addNestedPass(mlir::bufferization::createPromoteBuffersToStackPass()); - // TODO: migrate to new buffer deallocation "buffer-deallocation-pipeline" - pm.addNestedPass(catalyst::createBufferDeallocationPass()); + mlir::bufferization::BufferDeallocationPipelineOptions bufferDeallocOptions; + mlir::bufferization::buildBufferDeallocationPipeline(pm, bufferDeallocOptions); pm.addPass(catalyst::createArrayListToMemRefPass()); pm.addPass(mlir::createConvertBufferizationToMemRefPass()); pm.addPass(mlir::createCanonicalizerPass()); diff --git a/mlir/lib/Quantum/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Quantum/Transforms/BufferizableOpInterfaceImpl.cpp index 4a0312478e..bf150b0cc6 100644 --- a/mlir/lib/Quantum/Transforms/BufferizableOpInterfaceImpl.cpp +++ b/mlir/lib/Quantum/Transforms/BufferizableOpInterfaceImpl.cpp @@ -177,8 +177,6 @@ struct SampleOpInterface return false; } - bool bufferizesToAllocation(Operation *op, Value value) const { return true; } - bufferization::AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand, const bufferization::AnalysisState &state) const @@ -227,8 +225,6 @@ struct CountsOpInterface return false; } - bool bufferizesToAllocation(Operation *op, Value value) const { return true; } - bufferization::AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand, const bufferization::AnalysisState &state) const @@ -287,8 +283,6 @@ struct ProbsOpInterface return false; } - bool bufferizesToAllocation(Operation *op, Value value) const { return true; } - bufferization::AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand, const bufferization::AnalysisState &state) const @@ -340,8 +334,6 @@ struct StateOpInterface return false; } - bool bufferizesToAllocation(Operation *op, Value value) const { return true; } - bufferization::AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand, const bufferization::AnalysisState &state) const diff --git a/mlir/test/Quantum/AllocationTest.mlir b/mlir/test/Quantum/AllocationTest.mlir index 0d930954ba..3ee02385da 100644 --- a/mlir/test/Quantum/AllocationTest.mlir +++ b/mlir/test/Quantum/AllocationTest.mlir @@ -12,25 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. -// RUN: quantum-opt --buffer-deallocation --split-input-file %s | FileCheck %s +// RUN: quantum-opt --buffer-deallocation-pipeline --split-input-file %s | FileCheck %s // CHECK-LABEL: @existing_deallocation -func.func @existing_deallocation() { +func.func @existing_deallocation() -> !quantum.reg { // CHECK: [[reg:%.+]] = quantum.alloc %0 = quantum.alloc( 1) : !quantum.reg // CHECK: quantum.dealloc [[reg]] quantum.dealloc %0 : !quantum.reg - return + return %0 : !quantum.reg } // ----- // CHECK-LABEL: @missing_deallocation -func.func @missing_deallocation() { +func.func @missing_deallocation() -> !quantum.reg { // CHECK: [[reg:%.+]] = quantum.alloc %0 = quantum.alloc( 1) : !quantum.reg // CHECK-NOT: quantum.dealloc [[reg]] - return + return %0 : !quantum.reg } // -----