[Runtime][Frontend] Add support for custom QuantumDevices. (#327)

**Context:** Changes needed to support third party devices.

Third party devices must implement the
`Catalyst::Runtime::QuantumDevice` interface and add a `getCustomDevice`
method which returns an instance of their device.

```cpp
extern "C" Catalyst::Runtime::QuantumDevice*
getCustomDevice() { return new CustomDevice(); }
```

PennyLane plugin devices should also implement

```python
    @staticmethod
    def get_c_interface():
        return path_to_shared_lib
```

Then the following works:

```python

@qjit
@qml.qnode(qml.device("custom.device", wires=wires))
def circuit():
    ...
```

**Description of the Change:** Changes in the frontend to allow any
devices that implement a `get_c_interface` to be loaded into the
runtime. Changes to the runtime to load the shared library and call
`getCustomDevice`.

**Benefits:** Third party developers can implement their own devices as
long as they follow the spec.

**Possible Drawbacks:** Benchmarks showed no impact.

**Related GitHub Issues:**

- [x] Tests
- [x] Documentation

**Note**: `RTLD_DEEPBIND` is needed for allowing custom QuantumDevices
to work. However, `RTLD_DEEPBIND` is incompatible with sanitizers.

[sc-46413]

---------

Co-authored-by: Josh Izaac <[email protected]>
Co-authored-by: Ali Asadi <[email protected]>
Co-authored-by: David Ittah <[email protected]>

.github/workflows/check-catalyst.yaml

@@ -36,6 +36,16 @@ jobs:
           ENABLE_OPENQASM=ON \
           make runtime
 
+          # This is needed in the artifact for the pytests
+          # Note the lack of sanitizers.
+          # Left other flags the same.
+          COMPILER_LAUNCHER="" \
+          RT_BUILD_DIR="$(pwd)/runtime-build" \
+          ENABLE_LIGHTNING_KOKKOS=ON \
+          ENABLE_OPENQASM=ON \
+          make dummy_device
+
+
       - name: Upload Catalyst-Runtime Artifact
         uses: actions/upload-artifact@v3
         with:

Makefile

@@ -22,6 +22,7 @@ help:
 	@echo "  frontend           to install Catalyst Frontend"
 	@echo "  mlir               to build MLIR and custom Catalyst dialects"
 	@echo "  runtime            to build Catalyst Runtime with PennyLane-Lightning"
+	@echo "  dummy_device       needed for frontend tests"
 	@echo "  test               to run the Catalyst test suites"
 	@echo "  docs               to build the documentation for Catalyst"
 	@echo "  clean              to uninstall Catalyst and delete all temporary and cache files"
@@ -31,6 +32,7 @@ help:
 	@echo "  format [check=1]   to apply C++ and Python formatter; use with 'check=1' to check instead of modify (requires black, pylint and clang-format)"
 	@echo "  format [version=?] to apply C++ and Python formatter; use with 'version={version}' to run clang-format-{version} instead of clang-format"
 
+
 .PHONY: all
 all: runtime mlir frontend
 
@@ -58,6 +60,9 @@ dialects:
 runtime:
 	$(MAKE) -C runtime all
 
+dummy_device:
+	$(MAKE) -C runtime dummy_device
+
 .PHONY: test test-runtime test-frontend lit pytest test-demos
 test: test-runtime test-frontend test-demos
 

doc/changelog.md

@@ -8,6 +8,10 @@
   messages and includes a verbose trace if verbose mode is enabled.
   [(#303)](https://github.com/PennyLaneAI/catalyst/pull/303)
 
+* Add support for third party devices.
+  Third party `QuantumDevice` implementations can now be loaded into the runtime.
+  [(#327)](https://github.com/PennyLaneAI/catalyst/pull/327)
+
 <h3>Breaking changes</h3>
 
 * The axis ordering for `catalyst.jacobian` is updated to match `jax.jacobian`. Assume we have parameters of shape 

doc/dev/custom_devices.rst

@@ -0,0 +1,159 @@
+
+Custom Devices
+##############
+
+Differences between PennyLane and Catalyst
+==========================================
+
+PennyLane and Catalyst treat devices a bit differently.
+In PennyLane, one is able to `define devices <https://docs.pennylane.ai/en/stable/development/plugins.html>`_ in Python.
+Catalyst cannot interface with Python devices yet.
+Instead, Catalyst can only interact with devices that implement the `QuantumDevice <../api/file_runtime_include_QuantumDevice.hpp.html>`_ class.
+
+Here is an example of a custom ``QuantumDevice`` in which every single quantum operation is implemented as a no-operation.
+Additionally, all measurements will always return ``true``.
+
+.. code-block:: c++
+
+        #include <QuantumDevice.hpp>
+
+        struct DummyDevice : public Catalyst::Runtime::QuantumDevice {
+            DummyDevice() = default;          // LCOV_EXCL_LINE
+            virtual ~DummyDevice() = default; // LCOV_EXCL_LINE
+
+            DummyDevice &operator=(const QuantumDevice &) = delete;
+            DummyDevice(const DummyDevice &) = delete;
+            DummyDevice(DummyDevice &&) = delete;
+            DummyDevice &operator=(QuantumDevice &&) = delete;
+
+            virtual std::string getName(void) { return "DummyDevice"; }
+
+            auto AllocateQubit() -> QubitIdType { return 0; }
+            virtual auto AllocateQubits(size_t num_qubits) -> std::vector<QubitIdType>
+            {
+                return std::vector<QubitIdType>(num_qubits);
+            }
+            [[nodiscard]] virtual auto Zero() const -> Result { return NULL; }
+            [[nodiscard]] virtual auto One() const -> Result { return NULL; }
+            virtual auto Observable(ObsId, const std::vector<std::complex<double>> &,
+                                    const std::vector<QubitIdType> &) -> ObsIdType
+            {
+                return 0;
+            }
+            virtual auto TensorObservable(const std::vector<ObsIdType> &) -> ObsIdType { return 0; }
+            virtual auto HamiltonianObservable(const std::vector<double> &, const std::vector<ObsIdType> &)
+                -> ObsIdType
+            {
+                return 0;
+            }
+            virtual auto Measure(QubitIdType) -> Result
+            {
+                bool *ret = (bool *)malloc(sizeof(bool));
+                *ret = true;
+                return ret;
+            }
+
+            virtual void ReleaseQubit(QubitIdType) {}
+            virtual void ReleaseAllQubits() {}
+            [[nodiscard]] virtual auto GetNumQubits() const -> size_t { return 0; }
+            virtual void SetDeviceShots(size_t shots) {}
+            [[nodiscard]] virtual auto GetDeviceShots() const -> size_t { return 0; }
+            virtual void StartTapeRecording() {}
+            virtual void StopTapeRecording() {}
+            virtual void PrintState() {}
+            virtual void NamedOperation(const std::string &, const std::vector<double> &,
+                                        const std::vector<QubitIdType> &, bool)
+            {
+            }
+
+            virtual void MatrixOperation(const std::vector<std::complex<double>> &,
+                                         const std::vector<QubitIdType> &, bool)
+            {
+            }
+
+            virtual auto Expval(ObsIdType) -> double { return 0.0; }
+            virtual auto Var(ObsIdType) -> double { return 0.0; }
+            virtual void State(DataView<std::complex<double>, 1> &) {}
+            virtual void Probs(DataView<double, 1> &) {}
+            virtual void PartialProbs(DataView<double, 1> &, const std::vector<QubitIdType> &) {}
+            virtual void Sample(DataView<double, 2> &, size_t) {}
+            virtual void PartialSample(DataView<double, 2> &, const std::vector<QubitIdType> &, size_t) {}
+            virtual void Counts(DataView<double, 1> &, DataView<int64_t, 1> &, size_t) {}
+
+            virtual void PartialCounts(DataView<double, 1> &, DataView<int64_t, 1> &,
+                                       const std::vector<QubitIdType> &, size_t)
+            {
+            }
+
+            virtual void Gradient(std::vector<DataView<double, 1>> &, const std::vector<size_t> &) {}
+        };
+
+In addition to implementing the ``QuantumDevice`` class, one must implement the following method:
+
+.. code-block:: c++
+
+    extern "C" Catalyst::Runtime::QuantumDevice*
+    getCustomDevice() { return new CustomDevice(); }
+
+where ``CustomDevice()`` is a constructor for your custom device.
+``CustomDevice``'s destructor will be called by the runtime.
+
+.. note::
+
+    This interface might change quickly in the near future.
+    Please check back regularly for updates and to ensure your device is compatible with a specific version of Catalyst.
+
+How to compile custom devices
+=============================
+
+One can follow the ``catalyst/runtime/tests/third_party/CMakeLists.txt`` `as an example. <https://github.com/PennyLaneAI/catalyst/blob/26b412b298f22565fea529d2019554e7ad9b9624/runtime/tests/third_party/CMakeLists.txt>`_
+
+.. code-block:: cmake
+
+        cmake_minimum_required(VERSION 3.20)
+
+        project(third_party_device)
+
+        set(CMAKE_CXX_STANDARD 20)
+        set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+        add_library(dummy_device SHARED dummy_device.cpp)
+        target_include_directories(dummy_device PUBLIC ${runtime_includes})
+        set_property(TARGET dummy_device PROPERTY POSITION_INDEPENDENT_CODE ON)
+
+
+
+Integration with Python devices
+===============================
+
+If you already have a custom PennyLane device defined in Python and have added a shared object that corresponds to your implementation of the ``QuantumDevice`` class, then all you need to do is to add a ``get_c_interface`` method to your PennyLane device.
+The ``get_c_interface`` method should be a static method that takes no parameters and returns the complete path to your shared library with the ``QuantumDevice`` implementation.
+After doing so, Catalyst should be able to interface with your custom device.
+
+.. code-block:: python
+
+    class DummyDevice(qml.QubitDevice):
+        """Dummy Device"""
+
+        name = "Dummy Device"
+        short_name = "dummy.device"
+        pennylane_requires = "0.32.0"
+        version = "0.0.1"
+        author = "Erick Ochoa"
+
+        def __init__(self, shots=None, wires=None):
+            super().__init__(wires=wires, shots=shots)
+
+        def apply(self, operations, **kwargs):
+            """Your normal definitions"""
+
+        @staticmethod
+        def get_c_interface():
+            """Location to shared object with C/C++ implementation"""
+            return "full/path/to/libdummy_device.so"
+
+    @qjit
+    @qml.qnode(DummyDevice(wires=1))
+    def f():
+        return measure(0)
+

doc/index.rst

@@ -75,6 +75,7 @@ Catalyst
    modules/mlir
    modules/runtime
    dev/debugging
+   dev/custom_devices
    dev/roadmap
 
 .. toctree::

frontend/catalyst/pennylane_extensions.py

@@ -19,6 +19,7 @@
 # pylint: disable=too-many-lines
 
 import numbers
+import pathlib
 from functools import update_wrapper
 from typing import Any, Callable, Iterable, List, Optional, Union
 
@@ -126,11 +127,25 @@ def __call__(self, *args, **kwargs):
                 name = self.device.short_name
             else:
                 name = self.device.name
-            if name not in QFunc.RUNTIME_DEVICES:
+
+            is_known_device = name in QFunc.RUNTIME_DEVICES
+            implements_c_interface = hasattr(self.device, "get_c_interface")
+            is_valid_device = is_known_device or implements_c_interface
+            if not is_valid_device:
                 raise CompileError(
-                    f"The {name} device is not " "supported for compilation at the moment."
+                    f"The {name} device is not supported for compilation at the moment."
                 )
 
+            # TODO:
+            # Once all devices get converted to shared libraries this name should just be the path.
+            backend_path_or_name = name
+            if implements_c_interface:
+                impl = self.device.get_c_interface()
+                if not pathlib.Path(impl).is_file():
+                    raise CompileError(f"Device at {impl} cannot be found!")
+
+                backend_path_or_name = self.device.get_c_interface()
+
             backend_kwargs = {}
             if hasattr(self.device, "shots"):
                 backend_kwargs["shots"] = self.device.shots if self.device.shots else 0
@@ -142,7 +157,7 @@ def __call__(self, *args, **kwargs):
                     backend_kwargs["s3_destination_folder"] = str(self.device._s3_folder)
 
             device = QJITDevice(
-                self.device.shots, self.device.wires, self.device.short_name, backend_kwargs
+                self.device.shots, self.device.wires, backend_path_or_name, backend_kwargs
             )
         else:
             # Allow QFunc to still be used by itself for internal testing.

frontend/test/pytest/test_custom_devices.py

@@ -0,0 +1,102 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit test for custom device integration with Catalyst.
+"""
+import pathlib
+
+import pennylane as qml
+import pytest
+
+from catalyst import measure, qjit
+from catalyst.compiler import get_lib_path
+from catalyst.utils.exceptions import CompileError
+
+
+@pytest.mark.skipif(
+    not pathlib.Path(get_lib_path("runtime", "RUNTIME_LIB_DIR") + "/libdummy_device.so").is_file(),
+    reason="lib_dummydevice.so was not found.",
+)
+def test_custom_device():
+    """Test that custom device can run using Catalyst."""
+
+    class DummyDevice(qml.QubitDevice):
+        """Dummy Device"""
+
+        name = "Dummy Device"
+        short_name = "dummy.device"
+        pennylane_requires = "0.32.0"
+        version = "0.0.1"
+        author = "Dummy"
+
+        # Doesn't matter as at the moment it is dictated by QJITDevice
+        operations = []
+        observables = []
+
+        def __init__(self, shots=None, wires=None):
+            super().__init__(wires=wires, shots=shots)
+
+        def apply(self, operations, **kwargs):
+            """Unused"""
+            raise RuntimeError("Only C/C++ interface is defined")
+
+        @staticmethod
+        def get_c_interface():
+            """Location to shared object with C/C++ implementation"""
+            return get_lib_path("runtime", "RUNTIME_LIB_DIR") + "/libdummy_device.so"
+
+    @qjit
+    @qml.qnode(DummyDevice(wires=1))
+    def f():
+        """This function would normally return False.
+        However, DummyDevice as defined in libdummy_device.so
+        has been implemented to always return True."""
+        return measure(0)
+
+    assert True == f()
+
+
+def test_custom_device_bad_directory():
+    """Test that custom device error."""
+
+    class DummyDevice(qml.QubitDevice):
+        """Dummy Device"""
+
+        name = "Dummy Device"
+        short_name = "dummy.device"
+        pennylane_requires = "0.32.0"
+        version = "0.0.1"
+        author = "Dummy"
+
+        # Doesn't matter as at the moment it is dictated by QJITDevice
+        operations = []
+        observables = []
+
+        def __init__(self, shots=None, wires=None):
+            super().__init__(wires=wires, shots=shots)
+
+        def apply(self, operations, **kwargs):
+            """Unused."""
+            raise RuntimeError("Only C/C++ interface is defined")
+
+        @staticmethod
+        def get_c_interface():
+            """Location to shared object with C/C++ implementation"""
+            return "this-file-does-not-exist.so"
+
+    with pytest.raises(CompileError, match="Device .* cannot be found"):
+
+        @qjit
+        @qml.qnode(DummyDevice(wires=1))
+        def f():
+            return measure(0)

runtime/Makefile

@@ -64,6 +64,10 @@ runtime: $(RT_BUILD_DIR)/lib/backend/librt_backend.so $(RT_BUILD_DIR)/lib/librt_
 $(RT_BUILD_DIR)/tests/runner_tests: configure
 	cmake --build $(RT_BUILD_DIR) --target runner_tests -j$(NPROC)
 
+.PHONY: dummy_device
+dummy_device:
+	cmake --build $(RT_BUILD_DIR) --target dummy_device -j$(NPROC)
+
 .PHONY: test
 test: $(RT_BUILD_DIR)/tests/runner_tests
 	@echo "test the Catalyst runtime test suite"