Skip to content

Commit

Permalink
Implement get_c_interface method for Kokkos plugin (#837)
Browse files Browse the repository at this point in the history
**Context:** Catalyst is removing the Kokkos plugin from its code base
(PennyLaneAI/catalyst#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 <[email protected]>
  • Loading branch information
rauletorresc and ringo-but-quantum authored Aug 7, 2024
1 parent 619b807 commit 8f517d2
Show file tree
Hide file tree
Showing 5 changed files with 95 additions and 3 deletions.
1 change: 1 addition & 0 deletions .github/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
5 changes: 3 additions & 2 deletions .github/workflows/tests_lkcpu_python.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -297,4 +298,4 @@ jobs:
with:
fail_ci_if_error: true
verbose: true
token: ${{ secrets.CODECOV_TOKEN }}
token: ${{ secrets.CODECOV_TOKEN }}
2 changes: 1 addition & 1 deletion pennylane_lightning/core/_version.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,4 +16,4 @@
Version number (major.minor.patch[-label])
"""

__version__ = "0.38.0-dev25"
__version__ = "0.38.0-dev26"
44 changes: 44 additions & 0 deletions pennylane_lightning/lightning_kokkos/lightning_kokkos.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.<system>-<architecture>-<python-id>"
# 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")
46 changes: 46 additions & 0 deletions tests/test_device.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 8f517d2

Please sign in to comment.