diff --git a/src/xdsl_quantum/dialects/__init__.py b/src/xdsl_quantum/dialects/__init__.py index 9cb4132..7b3d8aa 100644 --- a/src/xdsl_quantum/dialects/__init__.py +++ b/src/xdsl_quantum/dialects/__init__.py @@ -6,5 +6,11 @@ def get_all_dialects() -> dict[str, Callable[[], Dialect]]: """Returns all available dialects.""" - # Add your dialects here to be discovered by `xdsl-opt` - return {} + def get_qir(): + from xdsl_quantum.dialects.qir import QIR + + return QIR + + return { + "qir": get_qir, + } diff --git a/src/xdsl_quantum/dialects/qir.py b/src/xdsl_quantum/dialects/qir.py new file mode 100644 index 0000000..6188b5f --- /dev/null +++ b/src/xdsl_quantum/dialects/qir.py @@ -0,0 +1,643 @@ +from abc import ABC, abstractmethod + +from xdsl.dialects import llvm +from xdsl.dialects.builtin import I32, I64, Float64Type, i1, i32, i64 +from xdsl.ir import Dialect, Operation, ParametrizedAttribute, SSAValue, TypeAttribute +from xdsl.irdl import ( + IRDLOperation, + irdl_attr_definition, + irdl_op_definition, + operand_def, + result_def, +) + + +@irdl_attr_definition +class QubitType(ParametrizedAttribute, TypeAttribute): + """ + QIR qubit type. Lowers to opaque pointer. + """ + + name = "qir.qubit" + + +@irdl_attr_definition +class ResultType(ParametrizedAttribute, TypeAttribute): + """ + QIR result type. Lowers to opaque pointer. + """ + + name = "qir.result" + + +@irdl_attr_definition +class ArrayType(ParametrizedAttribute, TypeAttribute): + """ + QIR array + """ + + name = "qir.array" + + +class QIROperation(IRDLOperation, ABC): + """ + QIR operations are defined using opaque function. + """ + + @staticmethod + @abstractmethod + def get_func_name() -> str: + """ + Return the name of the opaque function. + """ + ... + + @classmethod + def get_func_type(cls) -> llvm.LLVMFunctionType: + """ + Return the llvm function type for this function. + If this function has inputs/outputs which are not llvm pointers, + then this function must be overriden. + """ + irdl_def = cls.get_irdl_definition() + operands = len(irdl_def.operands) + results = len(irdl_def.results) + + return llvm.LLVMFunctionType( + (llvm.LLVMPointerType(),) * operands, + llvm.LLVMPointerType() if results else None, + ) + + +@irdl_op_definition +class ArrayCreate1D(QIROperation): + """ + MLIR equivalent of __quantum__rt__array_create_1d + We restrict to arrays of qubits for now. + """ + + name = "qir.array_create_1d" + + elem_size = operand_def(I32) + + count = operand_def(I64) + + out = result_def(ArrayType) + + assembly_format = "$elem_size `,` $count attr-dict" + + @staticmethod + def get_func_name() -> str: + return "__quantum__rt__array_create_1d" + + @classmethod + def get_func_type(cls) -> llvm.LLVMFunctionType: + return llvm.LLVMFunctionType((i32, i64), llvm.LLVMPointerType()) + + def __init__(self, elem_size: SSAValue | Operation, count: SSAValue | Operation): + super().__init__(operands=(elem_size, count), result_types=((ArrayType(),))) + + +@irdl_op_definition +class ArrayGetElementPtr(QIROperation): + """ + MLIR equivalent of __quantum__rt__array_get_element_ptr_1d + """ + + name = "qir.get_element_ptr" + + arr = operand_def(ArrayType) + + index = operand_def(I64) + + out = result_def(llvm.LLVMPointerType()) + + assembly_format = "$arr `[` $index `]` attr-dict" + + @staticmethod + def get_func_name() -> str: + return "__quantum__rt__array_get_element_ptr_1d" + + @classmethod + def get_func_type(cls) -> llvm.LLVMFunctionType: + return llvm.LLVMFunctionType( + (llvm.LLVMPointerType(), i64), llvm.LLVMPointerType() + ) + + def __init__(self, arr: SSAValue | Operation, index: SSAValue | Operation): + super().__init__( + operands=(arr, index), result_types=((llvm.LLVMPointerType(),)) + ) + + +@irdl_op_definition +class ResultGetOneOp(QIROperation): + """ + MLIR equivalent of __quantum__rt__result_get_one + """ + + name = "qir.result_get_one" + + out = result_def(ResultType) + + assembly_format = "attr-dict" + + @staticmethod + def get_func_name() -> str: + return "__quantum__rt__result_get_one" + + def __init__(self): + super().__init__(result_types=(ResultType(),)) + + +@irdl_op_definition +class ResultEqualOp(QIROperation): + """ + MLIR equivalent of __quantum__rt__result_equal + """ + + name = "qir.result_equal" + + lhs = operand_def(ResultType) + rhs = operand_def(ResultType) + + out = result_def(i1) + + assembly_format = "$lhs `,` $rhs attr-dict" + + @staticmethod + def get_func_name() -> str: + return "__quantum__rt__result_equal" + + @classmethod + def get_func_type(cls) -> llvm.LLVMFunctionType: + return llvm.LLVMFunctionType( + (llvm.LLVMPointerType(), llvm.LLVMPointerType()), i1 + ) + + def __init__(self, lhs: SSAValue | Operation, rhs: SSAValue | Operation): + super().__init__(operands=(lhs, rhs), result_types=(i1,)) + + +@irdl_op_definition +class ReadResultOp(QIROperation): + """ + MLIR equivalent of __quantum__rt__read_result__body + """ + + name = "qir.read_result" + + result = operand_def(ResultType) + + out = result_def(i1) + + assembly_format = "$result attr-dict" + + @staticmethod + def get_func_name() -> str: + return "__quantum__rt__read_result__body" + + @classmethod + def get_func_type(cls) -> llvm.LLVMFunctionType: + return llvm.LLVMFunctionType((llvm.LLVMPointerType(),), i1) + + def __init__(self, result: SSAValue | Operation): + super().__init__(operands=(result,), result_types=(i1,)) + + +@irdl_op_definition +class QubitAllocateOp(QIROperation): + """ + MLIR equivalent of __quantum__rt__qubit_allocate + """ + + name = "qir.qubit_allocate" + + out = result_def(QubitType) + + assembly_format = "attr-dict" + + @staticmethod + def get_func_name() -> str: + return "__quantum__rt__qubit_allocate" + + def __init__(self): + super().__init__(result_types=(QubitType(),)) + + +@irdl_op_definition +class MeasureOp(QIROperation): + """ + MLIR equivalent of __quantum__qis__m__body + """ + + name = "qir.m" + + qubit = operand_def(QubitType) + res = result_def(ResultType) + + assembly_format = "$qubit attr-dict" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__m__body" + + def __init__(self, qubit: SSAValue | Operation): + super().__init__(operands=(qubit,), result_types=(ResultType(),)) + + +@irdl_op_definition +class ReleaseOp(QIROperation): + """ + MLIR equivalent of __quantum__rt__qubit_release + """ + + name = "qir.qubit_release" + + qubit = operand_def(QubitType) + + assembly_format = "$qubit attr-dict" + + @staticmethod + def get_func_name() -> str: + return "__quantum__rt__qubit_release" + + def __init__(self, qubit: SSAValue | Operation): + super().__init__(operands=(qubit,)) + + +class ControlledOperation(QIROperation, ABC): + """ + Base class for operations with control and target qubits. + """ + + control = operand_def(QubitType) + target = operand_def(QubitType) + + assembly_format = "$control `,` $target attr-dict" + + def __init__(self, control: SSAValue | Operation, target: SSAValue | Operation): + super().__init__(operands=(control, target)) + + +@irdl_op_definition +class CNotOp(ControlledOperation): + """ + MLIR equivalent of __quantum__qis__cnot__body + """ + + name = "qir.cnot" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__cnot__body" + + +@irdl_op_definition +class CZOp(ControlledOperation): + """ + MLIR equivalent of __quantum__qis__cz__body + """ + + name = "qir.cz" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__cz__body" + + +class RotationOperation(QIROperation, ABC): + """ + Base class for rotation gates + """ + + angle = operand_def(Float64Type) + qubit = operand_def(QubitType) + + assembly_format = "`` `<` $angle `>` $qubit attr-dict" + + def __init__(self, angle: SSAValue | Operation, qubit: SSAValue | Operation): + super().__init__(operands=(angle, qubit)) + + @classmethod + def get_func_type(cls) -> llvm.LLVMFunctionType: + return llvm.LLVMFunctionType((Float64Type(), llvm.LLVMPointerType())) + + +@irdl_op_definition +class RXOp(RotationOperation): + """ + MLIR equivalent of __quantum__qis__rx__body + """ + + name = "qir.rx" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__rx__body" + + +@irdl_op_definition +class RYOp(RotationOperation): + """ + MLIR equivalent of __quantum__qis__ry__body + """ + + name = "qir.ry" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__ry__body" + + +@irdl_op_definition +class RZOp(RotationOperation): + """ + MLIR equivalent of __quantum__qis__rz__body + """ + + name = "qir.rz" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__rz__body" + + +@irdl_op_definition +class RZZOp(QIROperation): + """ + MLIR equivalent of __quantum__qis__rzz__body + """ + + name = "qir.rzz" + + angle = operand_def(Float64Type) + qubit1 = operand_def(QubitType) + qubit2 = operand_def(QubitType) + + assembly_format = "`` `<` $angle `>` $qubit1 `,` $qubit2 attr-dict" + + def __init__( + self, + angle: SSAValue | Operation, + qubit1: SSAValue | Operation, + qubit2: SSAValue | Operation, + ): + super().__init__(operands=(angle, qubit1, qubit2)) + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__rzz__body" + + @classmethod + def get_func_type(cls) -> llvm.LLVMFunctionType: + return llvm.LLVMFunctionType( + (Float64Type(), llvm.LLVMPointerType(), llvm.LLVMPointerType()) + ) + + +class ControlledRotationOperation(QIROperation, ABC): + """ + Base class for controlled rotation gates + """ + + angle = operand_def(Float64Type) + control = operand_def(ArrayType) + target = operand_def(QubitType) + + assembly_format = "`` `<` $angle `>` $control `,` $target attr-dict" + + def __init__( + self, + angle: SSAValue | Operation, + control: SSAValue | Operation, + target: SSAValue | Operation, + ): + super().__init__(operands=(angle, control, target)) + + @classmethod + def get_func_type(cls) -> llvm.LLVMFunctionType: + return llvm.LLVMFunctionType( + (Float64Type(), llvm.LLVMPointerType(), llvm.LLVMPointerType()) + ) + + +@irdl_op_definition +class CRXOp(ControlledRotationOperation): + """ + MLIR equivalent of __quantum__qis__rx__ctl + """ + + name = "qir.crx" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__rx__ctl" + + +@irdl_op_definition +class CRYOp(ControlledRotationOperation): + """ + MLIR equivalent of __quantum__qis__ry__ctl + """ + + name = "qir.cry" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__ry__ctl" + + +@irdl_op_definition +class CRZOp(ControlledRotationOperation): + """ + MLIR equivalent of __quantum__qis__rz__ctl + """ + + name = "qir.crz" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__rz__ctl" + + +class SingleQubitOperation(QIROperation, ABC): + """ + Base class for operations on a single qubit and no other operands. + """ + + qubit = operand_def(QubitType) + + assembly_format = "$qubit attr-dict" + + def __init__(self, qubit: SSAValue | Operation): + super().__init__(operands=(qubit,)) + + +@irdl_op_definition +class HOp(SingleQubitOperation): + """ + MLIR equivalent of __quantum__qis__h__body + """ + + name = "qir.h" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__h__body" + + +@irdl_op_definition +class SOp(SingleQubitOperation): + """ + MLIR equivalent of __quantum__qis__s__body + """ + + name = "qir.s" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__s__body" + + +@irdl_op_definition +class SAdjOp(SingleQubitOperation): + """ + MLIR equivalent of __quantum__qis__s__adj + """ + + name = "qir.s_adj" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__s__adj" + + +@irdl_op_definition +class TOp(SingleQubitOperation): + """ + MLIR equivalent of __quantum__qis__t__body + """ + + name = "qir.t" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__t__body" + + +@irdl_op_definition +class TAdjOp(SingleQubitOperation): + """ + MLIR equivalent of __quantum__qis__t__adj + """ + + name = "qir.t_adj" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__t__adj" + + +@irdl_op_definition +class XOp(SingleQubitOperation): + """ + MLIR equivalent of __quantum__qis__x__body + """ + + name = "qir.x" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__x__body" + + +@irdl_op_definition +class YOp(SingleQubitOperation): + """ + MLIR equivalent of __quantum__qis__y__body + """ + + name = "qir.y" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__y__body" + + +@irdl_op_definition +class ZOp(SingleQubitOperation): + """ + MLIR equivalent of __quantum__qis__z__body + """ + + name = "qir.z" + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__z__body" + + +@irdl_op_definition +class CCXOp(QIROperation): + """ + MLIR equivalent of __quantum__qis__ccx__body + """ + + name = "qir.ccx" + + c1 = operand_def(QubitType) + c2 = operand_def(QubitType) + targ = operand_def(QubitType) + + assembly_format = "$c1 `,` $c2 `,` $targ attr-dict" + + def __init__( + self, + c1: SSAValue | Operation, + c2: SSAValue | Operation, + targ: SSAValue | Operation, + ): + super().__init__(operands=(c1, c2, targ)) + + @staticmethod + def get_func_name() -> str: + return "__quantum__qis__ccx__body" + + +QIR = Dialect( + "qir", + [ + ArrayCreate1D, + ArrayGetElementPtr, + ResultGetOneOp, + ResultEqualOp, + ReadResultOp, + QubitAllocateOp, + MeasureOp, + ReleaseOp, + CNotOp, + CZOp, + HOp, + RXOp, + RYOp, + RZOp, + RZZOp, + CRXOp, + CRYOp, + CRZOp, + SOp, + SAdjOp, + TOp, + TAdjOp, + XOp, + YOp, + ZOp, + CCXOp, + ], + [ + QubitType, + ResultType, + ArrayType, + ], +) diff --git a/tests/filecheck/dialecs/qir/ops.mlir b/tests/filecheck/dialecs/qir/ops.mlir new file mode 100644 index 0000000..f038c2f --- /dev/null +++ b/tests/filecheck/dialecs/qir/ops.mlir @@ -0,0 +1,114 @@ +// RUN: XDSL_ROUNDTRIP +// RUN: XDSL_GENERIC_ROUNDTRIP + +// CHECK: %{{.*}} = qir.result_get_one +// CHECK-GENERIC: %{{.*}} = "qir.result_get_one"() : () -> !qir.result +%0 = qir.result_get_one + +// CHECK-NEXT: %{{.*}} = qir.result_equal %{{.*}}, %{{.*}} +// CHECK-GENERIC-NEXT: %{{.*}} = "qir.result_equal"(%{{.*}}, %{{.*}}) : (!qir.result, !qir.result) -> i1 +%1 = qir.result_equal %0, %0 + +// CHECK-NEXT: %{{.*}} = qir.read_result %{{.*}} +// CHECK-GENERIC-NEXT: %{{.*}} = "qir.read_result"(%{{.*}}) : (!qir.result) -> i1 +%2 = qir.read_result %0 + +// CHECK-NEXT: %{{.*}} = qir.qubit_allocate +// CHECK-GENERIC-NEXT: %{{.*}} = "qir.qubit_allocate"() : () -> !qir.qubit +%3 = qir.qubit_allocate + +// CHECK-NEXT: %{{.*}} = qir.m %{{.*}} +// CHECK-GENERIC-NEXT: %{{.*}} = "qir.m"(%{{.*}}) : (!qir.qubit) -> !qir.result +%4 = qir.m %3 + +// CHECK-NEXT: qir.qubit_release %{{.*}} +// CHECK-GENERIC-NEXT: "qir.qubit_release"(%{{.*}}) : (!qir.qubit) -> () +qir.qubit_release %3 + +%5, %6, %7 = "test.op"() : () -> (!qir.qubit, !qir.qubit, !qir.qubit) + +// CHECK: qir.cnot %{{.*}}, %{{.*}} +// CHECK-GENERIC: "qir.cnot"(%{{.*}}, %{{.*}}) : (!qir.qubit, !qir.qubit) -> () +qir.cnot %5, %6 + +// CHECK-NEXT: qir.cz %{{.*}}, %{{.*}} +// CHECK-GENERIC-NEXT: "qir.cz"(%{{.*}}, %{{.*}}) : (!qir.qubit, !qir.qubit) -> () +qir.cz %5, %6 + +// CHECK-NEXT: qir.h %{{.*}} +// CHECK-GENERIC-NEXT: "qir.h"(%{{.*}}) : (!qir.qubit) -> () +qir.h %5 + +%8 = arith.constant 1.000000e-01 : f64 + +// CHECK: qir.rx<%{{.*}}> %{{.*}} +// CHECK-GENERIC: "qir.rx"(%{{.*}}, %{{.*}}) : (f64, !qir.qubit) -> () +qir.rx<%8> %5 + +// CHECK-NEXT: qir.ry<%{{.*}}> %{{.*}} +// CHECK-GENERIC-NEXT: "qir.ry"(%{{.*}}, %{{.*}}) : (f64, !qir.qubit) -> () +qir.ry<%8> %5 + +// CHECK-NEXT: qir.rz<%{{.*}}> %{{.*}} +// CHECK-GENERIC-NEXT: "qir.rz"(%{{.*}}, %{{.*}}) : (f64, !qir.qubit) -> () +qir.rz<%8> %5 + +// CHECK-NEXT: qir.rzz<%{{.*}}> %{{.*}}, %{{.*}} +// CHECK-GENERIC-NEXT: "qir.rzz"(%{{.*}}, %{{.*}}, %{{.*}}) : (f64, !qir.qubit, !qir.qubit) -> () +qir.rzz<%8> %5, %6 + +// CHECK-NEXT: qir.s %{{.*}} +// CHECK-GENERIC-NEXT: "qir.s"(%{{.*}}) : (!qir.qubit) -> () +qir.s %5 + +// CHECK-NEXT: qir.s_adj %{{.*}} +// CHECK-GENERIC-NEXT: "qir.s_adj"(%{{.*}}) : (!qir.qubit) -> () +qir.s_adj %5 + +// CHECK-NEXT: qir.t %{{.*}} +// CHECK-GENERIC-NEXT: "qir.t"(%{{.*}}) : (!qir.qubit) -> () +qir.t %5 + +// CHECK-NEXT: qir.t_adj %{{.*}} +// CHECK-GENERIC-NEXT: "qir.t_adj"(%{{.*}}) : (!qir.qubit) -> () +qir.t_adj %5 + +// CHECK-NEXT: qir.x %{{.*}} +// CHECK-GENERIC-NEXT: "qir.x"(%{{.*}}) : (!qir.qubit) -> () +qir.x %5 + +// CHECK-NEXT: qir.y %{{.*}} +// CHECK-GENERIC-NEXT: "qir.y"(%{{.*}}) : (!qir.qubit) -> () +qir.y %5 + +// CHECK-NEXT: qir.z %{{.*}} +// CHECK-GENERIC-NEXT: "qir.z"(%{{.*}}) : (!qir.qubit) -> () +qir.z %5 + +// CHECK-NEXT: qir.ccx %{{.*}}, %{{.*}}, %{{.*}} +// CHECK-GENERIC-NEXT: "qir.ccx"(%{{.*}}, %{{.*}}, %{{.*}}) : (!qir.qubit, !qir.qubit, !qir.qubit) -> () +qir.ccx %5, %6, %7 + +%c8 = arith.constant 8 : i32 +%c1 = arith.constant 1 : i64 +%c0 = arith.constant 0 : i64 + +// CHECK: %9 = qir.array_create_1d %c8, %c1 +// CHECK-GENERIC: %9 = "qir.array_create_1d"(%c8, %c1) : (i32, i64) -> !qir.array +%9 = qir.array_create_1d %c8, %c1 + +// CHECK-NEXT: %10 = qir.get_element_ptr %9[%c0] +// CHECK-GENERIC-NEXT: "qir.get_element_ptr"(%9, %c0) : (!qir.array, i64) -> !llvm.ptr +%10 = qir.get_element_ptr %9[%c0] + +// CHECK: qir.crx<%{{.*}}> %{{.*}}, %{{.*}} +// CHECK-GENERIC: "qir.crx"(%{{.*}}, %{{.*}}, %{{.*}}) : (f64, !qir.array, !qir.qubit) -> () +qir.crx<%8> %9, %6 + +// CHECK-NEXT: qir.cry<%{{.*}}> %{{.*}}, %{{.*}} +// CHECK-GENERIC-NEXT: "qir.cry"(%{{.*}}, %{{.*}}, %{{.*}}) : (f64, !qir.array, !qir.qubit) -> () +qir.cry<%8> %9, %6 + +// CHECK-NEXT: qir.crz<%{{.*}}> %{{.*}}, %{{.*}} +// CHECK-GENERIC-NEXT: "qir.crz"(%{{.*}}, %{{.*}}, %{{.*}}) : (f64, !qir.array, !qir.qubit) -> () +qir.crz<%8> %9, %6