MDescriptor 0.3.2__tar.gz → 0.3.3__tar.gz

{mdescriptor-0.3.2 → mdescriptor-0.3.3}/CMakeLists.txt

@@ -5,9 +5,6 @@ set(PYBIND11_FINDPYTHON ON)
 find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)
 find_package(pybind11 CONFIG REQUIRED)
 
-option(MDESCRIPTOR_BUILD_NATIVE "Build the CPU native extension" ON)
-
-if(MDESCRIPTOR_BUILD_NATIVE)
 # OpenBLAS is a build-time dependency and is vendored into the resulting
 # wheel.  scipy-openblas32 deliberately exports a prefixed CBLAS ABI so that
 # this extension cannot accidentally bind to NumPy/SciPy's copy of OpenBLAS.
@@ -198,7 +195,6 @@ install(
     LIBRARY DESTINATION mdescriptor
     RUNTIME DESTINATION mdescriptor
 )
-endif()
 
 # The CPU and CUDA backends ship inside the same wheel; CUDA availability is
 # a runtime choice (execution.device), never a separate distribution.

{mdescriptor-0.3.2 → mdescriptor-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: MDescriptor
-Version: 0.3.2
+Version: 0.3.3
 Summary: Batch-oriented periodic atomic descriptors with C++17 kernels
 License:                     GNU GENERAL PUBLIC LICENSE
                                 Version 3, 29 June 2007
@@ -716,9 +716,6 @@ Requires-Dist: pytest>=8; extra == "reference-nepadapters"
 Requires-Dist: ase==3.26.0; extra == "reference-nepadapters"
 Requires-Dist: scipy>=1.11; extra == "reference-nepadapters"
 Requires-Dist: nep-adapters==1.0.2; extra == "reference-nepadapters"
-Provides-Extra: all
-Requires-Dist: ase>=3.22; extra == "all"
-Requires-Dist: scipy>=1.11; extra == "all"
 Provides-Extra: dev
 Requires-Dist: pytest>=8; extra == "dev"
 Requires-Dist: pytest-cov>=7; extra == "dev"
@@ -1062,10 +1059,8 @@ vendor adapters. The default inference graphs are lowered into the private
 ```
 
 Coverage excludes the isolated vendored DPA implementation and enforces 75%
-branch coverage for project-owned Python code.  Tests marked `timing` record
-repeatable timing samples but intentionally do not impose a runner-dependent
-speedup threshold; controlled performance reports are produced by
-`scripts/run_benchmark.py`.
+branch coverage for project-owned Python code.  Controlled performance
+reports are produced by `scripts/benchmarking/run_descriptor_benchmark.py`.
 
 The extension is private (`mdescriptor._native`). C++ shared math and batch
 helpers live in named headers under `cpp/include/mdescriptor/detail/`.

{mdescriptor-0.3.2 → mdescriptor-0.3.3}/README.md

@@ -332,10 +332,8 @@ vendor adapters. The default inference graphs are lowered into the private
 ```
 
 Coverage excludes the isolated vendored DPA implementation and enforces 75%
-branch coverage for project-owned Python code.  Tests marked `timing` record
-repeatable timing samples but intentionally do not impose a runner-dependent
-speedup threshold; controlled performance reports are produced by
-`scripts/run_benchmark.py`.
+branch coverage for project-owned Python code.  Controlled performance
+reports are produced by `scripts/benchmarking/run_descriptor_benchmark.py`.
 
 The extension is private (`mdescriptor._native`). C++ shared math and batch
 helpers live in named headers under `cpp/include/mdescriptor/detail/`.

{mdescriptor-0.3.2 → mdescriptor-0.3.3}/cpp/cuda/include/mdescriptor/cuda/batch.hpp

@@ -57,9 +57,6 @@ public:
     }
 
 private:
-    template <typename Value>
-    static void release(Value*& pointer) noexcept;
-
     std::int32_t* numbers_ = nullptr;
     double* positions_ = nullptr;
     double* positions_soa_ = nullptr;

{mdescriptor-0.3.2 → mdescriptor-0.3.3}/cpp/cuda/include/mdescriptor/cuda/dpa4.hpp

@@ -26,7 +26,6 @@ public:
     DeviceDpa4Model& operator=(const DeviceDpa4Model&) = delete;
     ~DeviceDpa4Model() noexcept;
 
-    std::int64_t feature_count() const noexcept;
     double cutoff() const noexcept;
 
     std::vector<double> compute(

{mdescriptor-0.3.2 → mdescriptor-0.3.3}/cpp/cuda/include/mdescriptor/cuda/dpa4c.hpp

@@ -24,16 +24,11 @@ public:
     DeviceDpa4cModel& operator=(const DeviceDpa4cModel&) = delete;
     ~DeviceDpa4cModel() noexcept;
 
-    std::int64_t feature_count() const noexcept { return feature_count_; }
     double cutoff() const noexcept { return rcut_; }
 
-    // Returns the checkpoint type index for an atomic number, or -1 when the
-    // payload did not provide a type_numbers entry or the number is absent.
-    int type_index_for_number(std::int32_t number) const noexcept;
-
     // type_indices is intentionally host-owned.  The implementation uploads
     // it on the model's stream and returns host float64 values with shape
-    // (batch.atoms, feature_count()).
+    // (batch.atoms, feature_count_).
     std::vector<double> compute(
         CudaExecutionContext& context,
         const DeviceBatch& batch,
@@ -64,11 +59,8 @@ private:
     std::vector<int> degree_channels_;
     std::vector<int> bispectrum_ranks_;
     std::vector<std::int64_t> degree_offsets_;
-    std::vector<std::int64_t> gram_offsets_;
     std::vector<std::int32_t> gram_index_;
     std::vector<float> gram_scale_;
-    std::vector<std::int32_t> type_numbers_;
-    std::vector<std::int32_t> host_type_lookup_;
 
     // Every tensor is contiguous and device-resident.  Keeping the ownership
     // in this object makes model construction exception safe and avoids a
@@ -95,6 +87,10 @@ private:
     std::unique_ptr<DeviceArray> probe_scale_;
     std::unique_ptr<DeviceArray> output_mean_;
     std::unique_ptr<DeviceArray> output_stddev_;
+    // The compact gram metadata is constant per model; it is uploaded once at
+    // construction instead of being copied into the workspace tail per call.
+    std::unique_ptr<DeviceArray> gram_index_device_;
+    std::unique_ptr<DeviceArray> gram_scale_device_;
 };
 
 } // namespace mdescriptor::cuda

{mdescriptor-0.3.2 → mdescriptor-0.3.3}/cpp/cuda/include/mdescriptor/cuda/neighbor_graph.hpp

@@ -4,6 +4,8 @@
 #include "mdescriptor/cuda/context.hpp"
 #include "mdescriptor/detail/batch.hpp"
 
+#include <cuda_runtime.h>
+
 #include <cstddef>
 #include <cstdint>
 #include <vector>
@@ -15,6 +17,30 @@ enum class NeighborGraphOrdering {
     Canonical,
 };
 
+// An edge is the exact self pair when it connects an atom to itself in the
+// home image.  Shifts may be enumerated locally (cell images) or read from a
+// CSR shift array; the array form stays null-safe because graphs built without
+// periodic images carry no shifts.
+__device__ __forceinline__ bool exact_self_edge(
+    std::int64_t center,
+    std::int64_t atom,
+    std::int64_t shift_x,
+    std::int64_t shift_y,
+    std::int64_t shift_z) {
+    return atom == center && shift_x == 0 && shift_y == 0 && shift_z == 0;
+}
+
+__device__ __forceinline__ bool exact_self_edge(
+    std::int64_t center,
+    std::int64_t atom,
+    const std::int32_t* shifts,
+    std::int64_t edge) {
+    return shifts != nullptr
+        && exact_self_edge(
+            center, atom, shifts[edge * 3], shifts[edge * 3 + 1],
+            shifts[edge * 3 + 2]);
+}
+
 // CSR + SoA representation shared by the CUDA descriptor family.  The graph
 // can either be uploaded by an adapter that already owns a host graph or be
 // constructed from a device batch.  DPA4/DPA4C use the latter path: only
@@ -28,13 +54,6 @@ public:
     DeviceNeighborGraph& operator=(const DeviceNeighborGraph&) = delete;
     ~DeviceNeighborGraph() noexcept;
 
-    void upload(
-        CudaExecutionContext& context,
-        const std::vector<std::int64_t>& offsets,
-        const std::vector<std::int32_t>& atoms,
-        const std::vector<std::int32_t>& shifts,
-        const std::vector<double>& displacements,
-        const std::vector<double>& distance2);
     // Build the DPA graph from device-resident coordinates.  The host batch
     // is metadata used to validate PBC and upload compact inverse cells/image
     // bounds; it is never used to enumerate or materialize pairs.
@@ -61,7 +80,6 @@ public:
     void clear() noexcept;
 
     std::size_t pairs() const noexcept { return pairs_; }
-    std::int64_t max_neighbors() const noexcept { return max_neighbors_; }
     bool slot_major() const noexcept { return slot_major_; }
     std::int64_t neighbor_stride() const noexcept { return neighbor_stride_; }
     const std::int32_t* neighbor_counts() const noexcept { return neighbor_counts_; }
@@ -83,16 +101,12 @@ private:
         const std::vector<double>& grid_spacing,
         const std::vector<std::int32_t>& grid_dimensions);
 
-    template <typename Value>
-    void ensure_capacity(Value** pointer, std::size_t* capacity, std::size_t count);
-
     std::int64_t* offsets_ = nullptr;
     std::int32_t* atoms_ = nullptr;
     std::int32_t* shifts_ = nullptr;
     double* displacements_ = nullptr;
     double* distance2_ = nullptr;
     std::size_t pairs_ = 0;
-    std::int64_t max_neighbors_ = 0;
     bool slot_major_ = false;
     std::int64_t neighbor_stride_ = 0;
     std::size_t offsets_capacity_ = 0;
@@ -124,7 +138,6 @@ private:
     std::int32_t* atom_to_structure_ = nullptr;
     std::int32_t* cell_counts_ = nullptr;
     std::int32_t* cell_offsets_ = nullptr;
-    std::int32_t* cell_fill_ = nullptr;
     std::int32_t* cell_atoms_ = nullptr;
     std::int32_t* cell_sort_keys_ = nullptr;
     std::int32_t* atom_cells_ = nullptr;
@@ -146,7 +159,6 @@ private:
     std::size_t reference_cell_inverses_capacity_ = 0;
     std::size_t cell_counts_capacity_ = 0;
     std::size_t cell_offsets_capacity_ = 0;
-    std::size_t cell_fill_capacity_ = 0;
 };
 
 } // namespace mdescriptor::cuda

{mdescriptor-0.3.2 → mdescriptor-0.3.3}/cpp/cuda/include/mdescriptor/cuda/nep.hpp

@@ -7,10 +7,15 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <memory>
 #include <vector>
 
 namespace mdescriptor::cuda {
 
+// Defined in device_memory.cuh; opaque here so the plugin-internal header
+// stays free of CUDA headers.
+struct DeviceArray;
+
 // Device-resident descriptor portion of one parsed NEP model.  The host
 // parser remains the single source of truth for the model protocol; this
 // object only owns the arrays needed by the descriptor kernel.
@@ -23,14 +28,12 @@ public:
     DeviceNepModel& operator=(const DeviceNepModel&) = delete;
     ~DeviceNepModel() noexcept;
 
-    int version() const noexcept { return version_; }
     int num_types() const noexcept { return num_types_; }
     int n_max_radial() const noexcept { return n_max_radial_; }
     int n_max_angular() const noexcept { return n_max_angular_; }
     int basis_size_radial() const noexcept { return basis_size_radial_; }
     int basis_size_angular() const noexcept { return basis_size_angular_; }
     int l_max() const noexcept { return l_max_; }
-    int num_l() const noexcept { return num_l_; }
     int dimension() const noexcept { return dimension_; }
     double radial_cutoff_max() const noexcept { return radial_cutoff_max_; }
     double angular_cutoff_max() const noexcept { return angular_cutoff_max_; }
@@ -46,28 +49,20 @@ public:
     // semantics instead of silently writing an all-zero GPU row.
     bool supports_atomic_number(std::int32_t number) const noexcept;
 
-    const std::int32_t* type_lookup() const noexcept { return type_lookup_; }
-    const float* radial_cutoff_pair() const noexcept { return radial_cutoff_pair_; }
-    const float* angular_cutoff_pair() const noexcept { return angular_cutoff_pair_; }
-    const float* radial_pair_coefficients() const noexcept {
-        return radial_pair_coefficients_;
-    }
-    const float* angular_pair_coefficients() const noexcept {
-        return angular_pair_coefficients_;
-    }
-    const float* scalers() const noexcept { return scalers_; }
+    const std::int32_t* type_lookup() const noexcept;
+    const float* radial_cutoff_pair() const noexcept;
+    const float* angular_cutoff_pair() const noexcept;
+    const float* radial_pair_coefficients() const noexcept;
+    const float* angular_pair_coefficients() const noexcept;
+    const float* scalers() const noexcept;
 
 private:
-    void release() noexcept;
-
-    int version_ = 0;
     int num_types_ = 0;
     int n_max_radial_ = 0;
     int n_max_angular_ = 0;
     int basis_size_radial_ = 0;
     int basis_size_angular_ = 0;
     int l_max_ = 0;
-    int num_l_ = 0;
     int dimension_ = 0;
     double radial_cutoff_max_ = 0.0;
     double angular_cutoff_max_ = 0.0;
@@ -79,19 +74,18 @@ private:
     bool has_q_134_ = false;
     std::vector<std::int32_t> host_type_lookup_;
 
-    std::int32_t* type_lookup_ = nullptr;
-    float* radial_cutoff_pair_ = nullptr;
-    float* angular_cutoff_pair_ = nullptr;
-    float* radial_pair_coefficients_ = nullptr;
-    float* angular_pair_coefficients_ = nullptr;
-    float* scalers_ = nullptr;
+    std::unique_ptr<DeviceArray> type_lookup_;
+    std::unique_ptr<DeviceArray> radial_cutoff_pair_;
+    std::unique_ptr<DeviceArray> angular_cutoff_pair_;
+    std::unique_ptr<DeviceArray> radial_pair_coefficients_;
+    std::unique_ptr<DeviceArray> angular_pair_coefficients_;
+    std::unique_ptr<DeviceArray> scalers_;
 };
 
 std::vector<double> compute_nep(
     CudaExecutionContext& context,
     const DeviceBatch& batch,
     const DeviceNeighborGraph& graph,
-    const DeviceNepModel& model,
-    bool reference_radial_accumulation = false);
+    const DeviceNepModel& model);
 
 } // namespace mdescriptor::cuda

{mdescriptor-0.3.2 → mdescriptor-0.3.3}/cpp/cuda/src/backend.cu

@@ -31,7 +31,11 @@
 
 namespace py = pybind11;
 
-namespace mdescriptor::cuda {
+// Shared extended-descriptor glue.  This header leaves namespace
+// mdescriptor::cuda open for its includers, so it must stay the last include:
+// everything below is already inside that namespace and the file closes it.
+#include "extended_descriptors_common.cuh"
+
 namespace {
 
 using I32Array = py::array_t<std::int32_t, py::array::c_style | py::array::forcecast>;
@@ -75,23 +79,6 @@ BatchArrays arrays_from_batch(const py::object& value) {
             std::move(pbc), std::move(offsets), view};
 }
 
-template <typename Value>
-Value option(const py::dict& options, const char* name, Value fallback) {
-    const py::str key(name);
-    if (!options.contains(key) || options[key].is_none()) {
-        return fallback;
-    }
-    return py::cast<Value>(options[key]);
-}
-
-std::vector<std::int32_t> species_option(const py::dict& options) {
-    const py::str key("species");
-    if (!options.contains(key) || options[key].is_none()) {
-        return {};
-    }
-    return py::cast<std::vector<std::int32_t>>(options[key]);
-}
-
 py::dict dpa_payload_option(const py::dict& options, const std::string& name) {
     const py::str key("_cuda_payload");
     if (!options.contains(key) || options[key].is_none()) {
@@ -148,13 +135,9 @@ std::int64_t feature_count_for(
     // that value private lets the CUDA plugin add a family without duplicating
     // every label/layout formula in this dispatch seam (matrix descriptors
     // may legitimately resolve their width only after seeing a batch).
-    const py::str feature_key("_cuda_feature_count");
-    if (options.contains(feature_key) && !options[feature_key].is_none()) {
-        const auto value = py::cast<std::int64_t>(options[feature_key]);
-        return value > 0 ? value : 0;
-    }
-    if (is_extended_descriptor(name)) {
-        return 0;
+    const std::int64_t configured = feature_count_option(options, 0);
+    if (is_extended_descriptor(name) || configured > 0) {
+        return configured;
     }
     throw std::invalid_argument("CUDA backend does not support this descriptor");
 }
@@ -190,8 +173,6 @@ std::vector<std::int32_t> dpa_type_indices(
     return result;
 }
 
-py::list generic_labels(const std::string& name, std::int64_t features);
-
 py::list dpa_labels(const py::dict& options, const std::string& name, std::int64_t features) {
     const py::dict payload = dpa_payload_option(options, name);
     const py::str key("labels");
@@ -202,7 +183,7 @@ py::list dpa_labels(const py::dict& options, const std::string& name, std::int64
             throw std::invalid_argument(name + " CUDA model payload has invalid labels");
         }
     }
-    return generic_labels(name, features);
+    return labels_option(options, name, features);
 }
 
 bool cancelled(const py::object& control) {
@@ -227,24 +208,6 @@ void check_cancelled(const py::object& control) {
     }
 }
 
-py::array double_array(const std::vector<double>& values, std::int64_t rows, std::int64_t columns) {
-    py::array_t<double> result({
-        static_cast<py::ssize_t>(rows), static_cast<py::ssize_t>(columns),
-    });
-    if (!values.empty()) {
-        std::copy(values.begin(), values.end(), result.mutable_data());
-    }
-    return result;
-}
-
-py::list generic_labels(const std::string& name, std::int64_t features) {
-    py::list labels;
-    for (std::int64_t index = 0; index < features; ++index) {
-        labels.append(name + ":" + std::to_string(index));
-    }
-    return labels;
-}
-
 py::list nep_labels(std::int64_t features) {
     py::list labels;
     for (std::int64_t index = 0; index < features; ++index) {
@@ -253,30 +216,6 @@ py::list nep_labels(std::int64_t features) {
     return labels;
 }
 
-py::array offsets_array(const std::vector<std::int64_t>& offsets) {
-    py::array_t<std::int64_t> result(offsets.size());
-    std::copy(offsets.begin(), offsets.end(), result.mutable_data());
-    return result;
-}
-
-py::dict result_metadata(const std::string& name, const py::dict& options) {
-    py::dict metadata;
-    metadata["descriptor"] = name;
-    metadata["backend"] = "mdescriptor-cuda";
-    py::dict execution;
-    execution["device"] = "cuda";
-    const py::str execution_key("execution");
-    if (options.contains(execution_key) && !options[execution_key].is_none()) {
-        const py::dict options_execution = py::cast<py::dict>(options[execution_key]);
-        execution["num_threads"] = options_execution.contains("num_threads")
-            ? options_execution["num_threads"] : py::none();
-    } else {
-        execution["num_threads"] = py::none();
-    }
-    metadata["execution"] = execution;
-    return metadata;
-}
-
 __global__ void write_neighbor_records(
     const std::int64_t* graph_offsets,
     const std::int32_t* graph_atoms,
@@ -300,7 +239,7 @@ __global__ void write_neighbor_records(
         const std::int32_t shift_x = graph_shifts[index * 3 + 0];
         const std::int32_t shift_y = graph_shifts[index * 3 + 1];
         const std::int32_t shift_z = graph_shifts[index * 3 + 2];
-        if (!self_pairs && atom == center && shift_x == 0 && shift_y == 0 && shift_z == 0) {
+        if (!self_pairs && exact_self_edge(center, atom, shift_x, shift_y, shift_z)) {
             continue;
         }
         if (!full_neighbor_list && !detail::keep_half_neighbor(
@@ -340,8 +279,7 @@ __global__ void count_filtered_neighbor_records(
     for (std::int64_t index = begin; index < end; ++index) {
         const std::int32_t atom = graph_atoms[index];
         const std::int32_t* shift = graph_shifts + index * 3;
-        if (!self_pairs && atom == center
-            && shift[0] == 0 && shift[1] == 0 && shift[2] == 0) {
+        if (!self_pairs && exact_self_edge(center, atom, graph_shifts, index)) {
             continue;
         }
         if (!full_neighbor_list && !detail::keep_half_neighbor(
@@ -502,17 +440,17 @@ py::object Backend::compute(py::object batch_object, py::object control) {
             }
         }
         py::dict result;
-        result["values"] = double_array(values, rows, 4);
+        result["values"] = values_array(values, rows, 4);
         result["level"] = "pair";
-        result["row_offsets"] = offsets_array(row_offsets);
-        result["pair_records"] = double_array(pair_records, rows, 5);
+        result["row_offsets"] = i64_array(row_offsets);
+        result["pair_records"] = values_array(pair_records, rows, 5);
         py::list labels;
         labels.append("dx");
         labels.append("dy");
         labels.append("dz");
         labels.append("distance");
         result["labels"] = labels;
-        result["metadata"] = result_metadata(name_, options_);
+        result["metadata"] = mdescriptor::cuda::metadata(options_, name_);
         return std::move(result);
     }
 
@@ -522,11 +460,11 @@ py::object Backend::compute(py::object batch_object, py::object control) {
                 mark_completed(control);
             }
             py::dict result;
-            result["values"] = double_array({}, 0, feature_count_);
+            result["values"] = values_array({}, 0, feature_count_);
             result["level"] = "atom";
             result["row_offsets"] = arrays.offsets;
             result["labels"] = nep_labels(feature_count_);
-            result["metadata"] = result_metadata(name_, options_);
+            result["metadata"] = mdescriptor::cuda::metadata(options_, name_);
             return std::move(result);
         }
         for (std::int64_t atom = 0; atom < arrays.view.atoms; ++atom) {
@@ -557,7 +495,7 @@ py::object Backend::compute(py::object batch_object, py::object control) {
             device_graph_.build_nep(
                 *context_, *compute_batch, compute_view, cutoff);
             const auto computed = compute_nep(
-                *context_, *compute_batch, device_graph_, *nep_model_, true);
+                *context_, *compute_batch, device_graph_, *nep_model_);
             values = computed;
         }
         check_cancelled(control);
@@ -565,11 +503,11 @@ py::object Backend::compute(py::object batch_object, py::object control) {
             mark_completed(control);
         }
         py::dict result;
-        result["values"] = double_array(values, arrays.view.atoms, feature_count_);
+        result["values"] = values_array(values, arrays.view.atoms, feature_count_);
         result["level"] = "atom";
         result["row_offsets"] = arrays.offsets;
         result["labels"] = nep_labels(feature_count_);
-        result["metadata"] = result_metadata(name_, options_);
+        result["metadata"] = mdescriptor::cuda::metadata(options_, name_);
         return std::move(result);
     }
 
@@ -592,11 +530,11 @@ py::object Backend::compute(py::object batch_object, py::object control) {
             mark_completed(control);
         }
         py::dict result;
-        result["values"] = double_array(values, arrays.view.atoms, feature_count_);
+        result["values"] = values_array(values, arrays.view.atoms, feature_count_);
         result["level"] = "atom";
         result["row_offsets"] = arrays.offsets;
         result["labels"] = dpa_labels(options_, name_, feature_count_);
-        result["metadata"] = result_metadata(name_, options_);
+        result["metadata"] = mdescriptor::cuda::metadata(options_, name_);
         return std::move(result);
     }
 
@@ -649,16 +587,16 @@ py::object Backend::compute(py::object batch_object, py::object control) {
         mark_completed(control);
     }
     py::dict result;
-    result["values"] = double_array(values, arrays.view.atoms, features);
+    result["values"] = values_array(values, arrays.view.atoms, features);
     result["level"] = "atom";
     result["row_offsets"] = arrays.offsets;
-    result["labels"] = generic_labels(name_, features);
-    result["metadata"] = result_metadata(name_, options_);
+    result["labels"] = labels_option(options_, name_, features);
+    result["metadata"] = mdescriptor::cuda::metadata(options_, name_);
     return std::move(result);
 }
 
 py::dict Backend::metadata() const {
-    return result_metadata(name_, options_);
+    return mdescriptor::cuda::metadata(options_, name_);
 }
 
 void Backend::close() noexcept {