Merge pull request #18 from qognitive/feature/nanobind

Feature/nanobind

nanobind_test/CMakeLists.txt

@@ -0,0 +1,41 @@
+cmake_minimum_required(VERSION 3.27)
+
+set(CMAKE_EXPORT_COMPILE_COMMANDS
+    TRUE
+    CACHE BOOL "Export compile commands to build directory" FORCE)
+
+include(cmake/CPM.cmake)
+
+project(nanobind_test LANGUAGES CXX)
+
+# Set C++ standard
+set(CMAKE_CXX_STANDARD 23)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CXX_EXTENSIONS OFF)
+# Need to enforce -fPIC across whole project to build shared libraries
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -Werror -g3")
+
+find_package(
+  Python 3.12
+  COMPONENTS Interpreter Development
+  REQUIRED)
+
+# TODO different compilation flags for Debug/Release modes; also pull in
+# target_compile_options for fast_pauli target defined below
+if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
+  set(CMAKE_BUILD_TYPE
+      Release
+      CACHE STRING "Choose the type of build." FORCE)
+  set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release"
+                                               "MinSizeRel" "RelWithDebInfo")
+endif()
+
+# Dependencies
+cpmaddpackage("gh:wjakob/nanobind#v2.0.0")
+cpmaddpackage("gh:fmtlib/fmt#10.2.1")
+cpmaddpackage("gh:kokkos/mdspan#b885a2c60ad42f9e1aaa0d317a38105b950cbed0")
+
+# add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/${nanobind_BUILD_DIR})
+nanobind_add_module(mdspan_wrapper mdspan_wrapper.cpp)
+target_link_libraries(mdspan_wrapper PRIVATE mdspan fmt::fmt)

nanobind_test/README.md

@@ -0,0 +1,13 @@
+# `nb::ndarray` :handshake: `std::mdspan`
+
+How to connect `nanobind`'s `nb::ndarray` and `std::mdspan`.
+
+Properties of `nb::ndarray`
+- work with numpy, torch, tensorflow, jax, cupy, and anything that supports DLPack
+- Supports zero-copy exchange using two protocols: 1) buffer protocol and 2) DLPack
+- Can return data that is owned by a Python object
+
+## Useful guides from `nanobind`
+- https://github.com/wjakob/nanobind/blob/b1531b9397c448c6b784520a4f052608f28a5e8d/include/nanobind/eigen/dense.h#L93C7-L93C24
+- https://github.com/wjakob/nanobind/blob/b1531b9397c448c6b784520a4f052608f28a5e8d/include/nanobind/eigen/dense.h#L124
+- https://github.com/wjakob/nanobind/blob/b1531b9397c448c6b784520a4f052608f28a5e8d/include/nanobind/eigen/dense.h#L242-L245

nanobind_test/cmake/CPM.cmake

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: MIT
+#
+# SPDX-FileCopyrightText: Copyright (c) 2019-2023 Lars Melchior and contributors
+
+set(CPM_DOWNLOAD_VERSION 0.39.0)
+set(CPM_HASH_SUM
+    "66639bcac9dd2907b2918de466783554c1334446b9874e90d38e3778d404c2ef")
+
+if(CPM_SOURCE_CACHE)
+  set(CPM_DOWNLOAD_LOCATION
+      "${CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
+elseif(DEFINED ENV{CPM_SOURCE_CACHE})
+  set(CPM_DOWNLOAD_LOCATION
+      "$ENV{CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
+else()
+  set(CPM_DOWNLOAD_LOCATION
+      "${CMAKE_BINARY_DIR}/cmake/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
+endif()
+
+# Expand relative path. This is important if the provided path contains a tilde
+# (~)
+get_filename_component(CPM_DOWNLOAD_LOCATION ${CPM_DOWNLOAD_LOCATION} ABSOLUTE)
+
+file(
+  DOWNLOAD
+  https://github.com/cpm-cmake/CPM.cmake/releases/download/v${CPM_DOWNLOAD_VERSION}/CPM.cmake
+  ${CPM_DOWNLOAD_LOCATION}
+  EXPECTED_HASH SHA256=${CPM_HASH_SUM})
+
+include(${CPM_DOWNLOAD_LOCATION})

nanobind_test/mdspan_wrapper.cpp

@@ -0,0 +1,183 @@
+#include <cstdint>
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+
+#include <algorithm>
+#include <experimental/mdspan>
+#include <fmt/format.h>
+#include <fmt/ranges.h>
+#include <iostream>
+#include <span>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+namespace nb = nanobind;
+using namespace nb::literals;
+
+template <typename T, size_t ndim>
+std::mdspan<T, std::dextents<size_t, ndim>>
+ndarray_cast_from_py(nb::ndarray<T> &a) {
+
+  std::array<size_t, ndim> shape;
+  for (size_t i = 0; i < ndim; ++i) {
+    shape[i] = a.shape(i);
+  }
+
+  return std::mdspan<T, std::dextents<size_t, ndim>>(a.data(), shape);
+}
+
+template <typename T, size_t ndim>
+std::pair<std::vector<T>, std::array<size_t, ndim>>
+cast_ndarray_to_blob(nb::ndarray<T> &a) {
+  // Shape info
+  size_t size = 1;
+  std::array<size_t, ndim> shape;
+  for (size_t i = 0; i < a.ndim(); ++i) {
+    shape[i] = a.shape(i);
+    size *= a.shape(i);
+  }
+
+  // Copy the raw data
+  std::vector<T> _data(size);
+  std::memcpy(_data.data(), a.data(), size * sizeof(T));
+  return std::make_pair(_data, shape);
+}
+template <typename T> struct PauliOp {
+  // As n_pauli_strings x n_qubits
+  std::mdspan<int, std::dextents<size_t, 2>> pauli_strings;
+  std::mdspan<T, std::dextents<size_t, 1>> coeffs;
+  std::vector<T> _coeffs;
+
+  PauliOp(nb::ndarray<int> &ps, nb::ndarray<T> &c) {
+    pauli_strings = ndarray_cast_from_py<int, 2>(ps);
+    std::array<size_t, 1> shape;
+    std::tie(_coeffs, shape) = cast_ndarray_to_blob<T, 1>(c);
+    coeffs = std::mdspan<T, std::dextents<size_t, 1>>(_coeffs.data(), shape);
+  }
+
+  bool operator==(const PauliOp &other) const {
+    if (pauli_strings.extent(0) != other.pauli_strings.extent(0) ||
+        pauli_strings.extent(1) != other.pauli_strings.extent(1)) {
+      return false;
+    }
+
+    if (coeffs.extent(0) != other.coeffs.extent(0)) {
+      return false;
+    }
+    // fmt::println("Dimensions match");
+
+    for (size_t i = 0; i < pauli_strings.extent(0); ++i) {
+      for (size_t j = 0; j < pauli_strings.extent(1); ++j) {
+        if (pauli_strings[i, j] != other.pauli_strings[i, j]) {
+          return false;
+        }
+      }
+    }
+
+    // fmt::println("Pauli strings match");
+    for (size_t i = 0; i < coeffs.extent(0); ++i) {
+      if (coeffs[i] != other.coeffs[i]) {
+        return false;
+      }
+    }
+
+    // fmt::println("Coeffs match");
+
+    return true;
+  }
+
+  void scale(T scale) {
+    for (size_t i = 0; i < coeffs.extent(0); ++i) {
+      coeffs[i] *= scale;
+    }
+  }
+
+  void multiply_coeff(std::mdspan<T, std::dextents<size_t, 1>> other) {
+    for (size_t i = 0; i < coeffs.extent(0); ++i) {
+      coeffs[i] *= other[i];
+    }
+  }
+
+  void print() {
+    fmt::print("Coeffs[");
+    for (size_t i = 0; i < coeffs.extent(0); ++i) {
+      fmt::print("{}, ", coeffs[i]);
+    }
+    fmt::print("]\n");
+  }
+
+  void return_coeffs(std::mdspan<T, std::dextents<size_t, 1>> &out) {
+    for (size_t i = 0; i < coeffs.extent(0); ++i) {
+      out[i] = coeffs[i];
+    }
+  }
+};
+
+NB_MODULE(mdspan_wrapper, m) {
+  //
+  nb::class_<PauliOp<double>>(m, "PauliOp")
+      .def(nb::init<nb::ndarray<int> &, nb::ndarray<double> &>())
+      .def("scale", &PauliOp<double>::scale, "scale"_a)
+      .def("print", &PauliOp<double>::print)
+      .def("__eq__", &PauliOp<double>::operator==)
+      .def("multiply_coeff",
+           [](PauliOp<double> &op, nb::ndarray<double> &c) {
+             op.multiply_coeff(ndarray_cast_from_py<double, 1>(c));
+           })
+      .def("return_coeffs",
+           [](PauliOp<double> &op, nb::ndarray<double> &out) {
+             auto out_mdspan = ndarray_cast_from_py<double, 1>(out);
+             op.return_coeffs(out_mdspan);
+           })
+      .def("return_coeffs_owning", [](PauliOp<double> &op) {
+        struct Temp {
+          std::vector<double> data;
+        };
+
+        Temp *tmp = new Temp{op._coeffs};
+
+        fmt::println("copied data: [{}]", fmt::join(tmp->data, ", "));
+        std::cout << std::flush;
+
+        nb::capsule deleter(
+            tmp, [](void *data) noexcept { delete static_cast<Temp *>(data); });
+
+        return nb::ndarray<nb::numpy, double>(
+            /*data*/ tmp->data.data(),
+            /*shape */ {tmp->data.size()},
+            /*deleter*/ deleter);
+      });
+
+  m.def("return_coeffs",
+        [](size_t n) {
+          std::vector<double> data(n);
+          for (size_t i = 0; i < n; ++i) {
+            data[i] = i;
+          }
+
+          struct Temp {
+            std::vector<double> data;
+          };
+
+          Temp *tmp = new Temp{data};
+
+          nb::capsule deleter(tmp, [](void *data) noexcept {
+            delete static_cast<Temp *>(data);
+          });
+
+          return nb::ndarray<nb::numpy, double>(
+              /*data*/ tmp->data.data(),
+              /*shape */ {tmp->data.size()},
+              /*deleter*/ deleter);
+
+          // nb::ndarray<nb::numpy> out{data.data(), {data.size()}, deleter};
+
+          // nb::object res = nb::cast(out, nb::rv_policy::copy);
+          // return res;
+        }
+
+        /**/
+  );
+}

nanobind_test/test.py

@@ -0,0 +1,87 @@
+import sys
+
+import numpy as np
+import pytest
+
+sys.path.append("build")
+
+import mdspan_wrapper
+
+
+def test_fake_pauli_op() -> None:
+    pauli_strings = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32)
+    coeffs = np.array([1.0, 2.0, 3.0])
+    op = mdspan_wrapper.PauliOp(pauli_strings, coeffs)
+
+    print(coeffs)
+
+    op.scale(2.0)
+    print(coeffs)
+
+
+@pytest.fixture
+def my_op() -> mdspan_wrapper.PauliOp:
+    pauli_strings = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32)
+    coeffs = np.array([1.0, 2.0, 3.0])
+    return mdspan_wrapper.PauliOp(pauli_strings, coeffs)
+
+
+def test_fake_pauli_op_scope(my_op: mdspan_wrapper.PauliOp) -> None:
+    my_op.scale(2.0)
+    op2 = mdspan_wrapper.PauliOp(
+        np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32),
+        np.array([1.0, 2.0, 3.0]) * 2,
+    )
+
+    # my_op.scale(2.0)
+    # print("After scaling, after op2")
+    # my_op.print()
+
+    # op3 = mdspan_wrapper.PauliOp(
+    #     np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32),
+    #     np.array([1.0, 2.0, 3.0]) * 4,
+    # )
+    # print("op3")
+    # op3.print()
+
+    # print("my_op")
+    # my_op.print()
+
+    # print("op2")
+    # op2.print()
+    assert my_op == op2
+
+
+def test_multiply_coeff(my_op: mdspan_wrapper.PauliOp) -> None:
+    coeffs = np.array([1.0, 2.0, 3.0])
+    my_op.multiply_coeff(coeffs)
+    op2 = mdspan_wrapper.PauliOp(
+        np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.int32),
+        np.array([1.0, 4.0, 9.0]),
+    )
+    assert my_op == op2
+
+
+def test_multiply_coeff_diff_type(my_op: mdspan_wrapper.PauliOp) -> None:
+    coeffs = np.array([1.0, 2.0, 3.0], dtype=np.complex128)
+    with pytest.raises(TypeError):
+        my_op.multiply_coeff(coeffs)
+
+
+def test_return_coeffs_non_owning(my_op: mdspan_wrapper.PauliOp) -> None:
+    c = np.zeros(3)
+    print(c)
+    my_op.return_coeffs(c)
+
+    print(c)
+    np.testing.assert_allclose(c, np.array([1.0, 2.0, 3.0]))
+
+
+def test_return_coeffs_owning(my_op: mdspan_wrapper.PauliOp) -> None:
+    c = my_op.return_coeffs_owning()
+    print("returned array", c)
+    np.testing.assert_allclose(c, np.array([1.0, 2.0, 3.0]))
+
+
+def test_standalone() -> None:
+    np.testing.assert_allclose(mdspan_wrapper.return_coeffs(3), np.array([0, 1.0, 2.0]))

pyproject.toml

@@ -94,6 +94,7 @@ exclude = [
     "venv",
     ".venv",
     "fast_pauli/__version__.py",
+    "nanobind_test"
 ]
 
 line-length = 88