From 8f517d24c71c6c7765f5c1bf29b0264b951de96a Mon Sep 17 00:00:00 2001 From: Raul Torres <138264735+rauletorresc@users.noreply.github.com> Date: Wed, 7 Aug 2024 12:45:21 -0500 Subject: [PATCH] Implement get_c_interface method for Kokkos plugin (#837) **Context:** Catalyst is removing the Kokkos plugin from its code base (https://github.com/PennyLaneAI/catalyst/pull/974) and is attempting to use it directly from Lightning. However, the C interface was not provided via Python. **Description of the Change:** Implement the C interface provider method in Python. **Benefits:** Catalyst can use the Kokkos plugin. --------- Co-authored-by: ringo-but-quantum --- .github/CHANGELOG.md | 1 + .github/workflows/tests_lkcpu_python.yml | 5 +- pennylane_lightning/core/_version.py | 2 +- .../lightning_kokkos/lightning_kokkos.py | 44 ++++++++++++++++++ tests/test_device.py | 46 +++++++++++++++++++ 5 files changed, 95 insertions(+), 3 deletions(-) diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md index c0e20d79f8..79e3ef5f9e 100644 --- a/.github/CHANGELOG.md +++ b/.github/CHANGELOG.md @@ -62,6 +62,7 @@ * Add a Catalyst-specific wrapping class for Lightning Kokkos. [(#770)](https://github.com/PennyLaneAI/pennylane-lightning/pull/770) + [(#837)](https://github.com/PennyLaneAI/pennylane-lightning/pull/837) * Add `initial_state_prep` option to Catalyst TOML file. [(#826)](https://github.com/PennyLaneAI/pennylane-lightning/pull/826) diff --git a/.github/workflows/tests_lkcpu_python.yml b/.github/workflows/tests_lkcpu_python.yml index 0e211f2e91..1dbad00706 100644 --- a/.github/workflows/tests_lkcpu_python.yml +++ b/.github/workflows/tests_lkcpu_python.yml @@ -246,7 +246,8 @@ jobs: run: | cd main/ DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"` - PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS --splits 7 --group ${{ matrix.group }} \ + # Remove `python -m` to avoid running tests with relative modules + PL_DEVICE=${DEVICENAME} pytest tests/ $COVERAGE_FLAGS --splits 7 --group ${{ matrix.group }} \ --store-durations --durations-path='.github/workflows/python_lightning_kokkos_test_durations.json' --splitting-algorithm=least_duration mv .github/workflows/python_lightning_kokkos_test_durations.json ${{ github.workspace }}/.test_durations-${{ matrix.exec_model }}-${{ matrix.group }} pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append @@ -297,4 +298,4 @@ jobs: with: fail_ci_if_error: true verbose: true - token: ${{ secrets.CODECOV_TOKEN }} \ No newline at end of file + token: ${{ secrets.CODECOV_TOKEN }} diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py index 3961045c01..444de99a58 100644 --- a/pennylane_lightning/core/_version.py +++ b/pennylane_lightning/core/_version.py @@ -16,4 +16,4 @@ Version number (major.minor.patch[-label]) """ -__version__ = "0.38.0-dev25" +__version__ = "0.38.0-dev26" diff --git a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py index c3dae23ea8..6528253264 100644 --- a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py +++ b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py @@ -17,6 +17,8 @@ interfaces with C++ for fast linear algebra calculations. """ +import os +import sys from os import getenv from pathlib import Path from typing import List @@ -834,3 +836,45 @@ def processing_fn(tape): return self.adjoint_jacobian(new_tape, starting_state, use_device_state) return processing_fn + + @staticmethod + def get_c_interface(): + """Returns a tuple consisting of the device name, and + the location to the shared object with the C/C++ device implementation. + """ + + # The shared object file extension varies depending on the underlying operating system + file_extension = "" + OS = sys.platform + if OS == "linux": + file_extension = ".so" + elif OS == "darwin": + file_extension = ".dylib" + else: + raise RuntimeError( + f"'LightningKokkosSimulator' shared library not available for '{OS}' platform" + ) + + lib_name = "liblightning_kokkos_catalyst" + file_extension + package_root = Path(__file__).parent + + # The absolute path of the plugin shared object varies according to the installation mode. + + # Wheel mode: + # Fixed location at the root of the project + wheel_mode_location = package_root.parent / lib_name + if wheel_mode_location.is_file(): + return "LightningKokkosSimulator", wheel_mode_location.as_posix() + + # Editable mode: + # The build directory contains a folder which varies according to the platform: + # lib.--" + # To avoid mismatching the folder name, we search for the shared object instead. + # TODO: locate where the naming convention of the folder is decided and replicate it here. + editable_mode_path = package_root.parent.parent / "build" + for path, _, files in os.walk(editable_mode_path): + if lib_name in files: + lib_location = (Path(path) / lib_name).as_posix() + return "LightningKokkosSimulator", lib_location + + raise RuntimeError("'LightningKokkosSimulator' shared library not found") diff --git a/tests/test_device.py b/tests/test_device.py index 0619546aed..a64e4d51d1 100644 --- a/tests/test_device.py +++ b/tests/test_device.py @@ -107,3 +107,49 @@ def circuit(): return qml.state() assert np.allclose(circuit(), np.array([1.0])) + + +@pytest.mark.skipif( + (device_name != "lightning.kokkos" or sys.platform != "win32"), + reason="This test is for Kokkos under Windows only.", +) +def test_unsupported_windows_platform_kokkos(): + """Test unsupported Windows platform for Kokkos.""" + + dev = qml.device(device_name, wires=0) + + with pytest.raises( + RuntimeError, + match="'LightningKokkosSimulator' shared library not available for 'win32' platform", + ): + dev.get_c_interface() + + +@pytest.mark.skipif( + (device_name != "lightning.kokkos" or sys.platform != "linux"), + reason="This test is for Kokkos under Linux only.", +) +def test_supported_linux_platform_kokkos(): + """Test supported Linux platform for Kokkos.""" + + dev = qml.device(device_name, wires=0) + + dev_name, shared_lib_name = dev.get_c_interface() + + assert dev_name == "LightningKokkosSimulator" + assert "liblightning_kokkos_catalyst.so" in shared_lib_name + + +@pytest.mark.skipif( + (device_name != "lightning.kokkos" or sys.platform != "darwin"), + reason="This test is for Kokkos under MacOS only.", +) +def test_supported_macos_platform_kokkos(): + """Test supported MacOS platform for Kokkos.""" + + dev = qml.device(device_name, wires=0) + + dev_name, shared_lib_name = dev.get_c_interface() + + assert dev_name == "LightningKokkosSimulator" + assert "liblightning_kokkos_catalyst.dylib" in shared_lib_name