PennyLaneAI / catalyst

A JIT compiler for hybrid quantum programs in PennyLane
https://docs.pennylane.ai/projects/catalyst
Apache License 2.0
138 stars 35 forks source link

[MLIR] Dialect `complex' not found for custom op 'complex.constant' #64

Closed sergei-mironov closed 1 year ago

sergei-mironov commented 1 year ago

Consider the following program running Catalyst 0.1.0 and PennyLane 0.28:

import pennylane as qml
import jax.numpy as jnp
from catalyst import qjit

@qjit
def workflow(params):
    @qml.qnode(qml.device("lightning.qubit", wires=1),)
    def circuit(phi):
        qml.PhaseShift(phi=phi, wires=[0])
        return qml.state()

    phi = params[0]
    for _ in range(10):
        state = circuit(phi)[0]
        phi = jnp.mean(state).real
        # phi = (state[0] + state[1])/2).real       # In contrast, this approach works
    return state

print(workflow(jnp.array([jnp.pi], dtype=jnp.float64)))

An attempt to run the program results in the following compilation error. Other expressions involving complex/float computations seem to cause similar errors.

/tmp/tmpkjcy955a/workflow.nohlo.mlir:6:12: error: Dialect `complex' not found for custom op 'complex.constant'
    %cst = complex.constant [0.000000e+00, 0.000000e+00] : complex<f64>
           ^
/tmp/tmpkjcy955a/workflow.nohlo.mlir:6:12: note: Registered dialects: arith, builtin, func, gradient, index, linalg, llvm, memref, quantum, scf, tensor ; for more info on dialect registration see https://mlir.llvm.org/getting_started/Faq/#registered-loaded-dependent-whats-up-with-dialects-management
Traceback (most recent call last):
  File "/home/sergei/complex_dialiect_issue.py", line 18, in <module>
    print(workflow(jnp.array([jnp.pi], dtype=jnp.float64)))
  File "/home/sergei/.local/lib/python3.8/site-packages/catalyst/compilation_pipelines.py", line 611, in __call__
    self.compiled_function = self.compile()
  File "/home/sergei/.local/lib/python3.8/site-packages/catalyst/compilation_pipelines.py", line 585, in compile
    shared_object, self._llvmir = compiler.compile(
  File "/home/sergei/.local/lib/python3.8/site-packages/catalyst/compiler.py", line 318, in compile
    buff = bufferize_tensors(mlir)
  File "/home/sergei/.local/lib/python3.8/site-packages/catalyst/compiler.py", line 195, in bufferize_tensors
    subprocess.run(command, stdout=file, check=True)
  File "/usr/lib/python3.8/subprocess.py", line 516, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['/home/sergei/.local/lib/python3.8/site-packages/catalyst/bin/quantum-opt', '/tmp/tmpkjcy955a/workflow.nohlo.mlir', '--lower-gradients', '--gradient-bufferize', '--scf-bufferize', '--convert-tensor-to-linalg', '--convert-elementwise-to-linalg', '--arith-bufferize', '--empty-tensor-to-alloc-tensor', '--bufferization-bufferize', '--tensor-bufferize', '--linalg-bufferize', '--tensor-bufferize', '--quantum-bufferize', '--func-bufferize', '--finalizing-bufferize', '--buffer-hoisting', '--convert-bufferization-to-memref', '--canonicalize', '--cse']' returned non-zero exit status 1.
dime10 commented 1 year ago

Thanks for reporting this one, should be an easy fix by registering the complex dialect in the quantum-opt tool. We could also consider registering all built-in dialects to avoid such issues in the future.