Custom cuda and metal kernels for system validation

.github/workflows/python-tests.yml

@@ -62,11 +62,6 @@ jobs:
         with:
           toolchain: stable
 
-      - name: Cache Rust dependencies
-        uses: Leafwing-Studios/cargo-cache@v2.6.1
-        with:
-          sweep-cache: true
-
       - name: Setup sccache
         if: ${{ !env.ACT }}
         uses: mozilla-actions/sccache-action@v0.0.10

.github/workflows/rust-tests.yml

@@ -104,12 +104,6 @@ jobs:
         with:
           python-version: "3.14"
 
-      - name: Cache Rust dependencies
-        uses: Leafwing-Studios/cargo-cache@v2.6.1
-        if: matrix.container == null
-        with:
-          sweep-cache: true
-
       - name: install valgrind
         if: matrix.do-valgrind
         run: |

.github/workflows/torch-tests.yml

@@ -70,11 +70,6 @@ jobs:
         with:
           toolchain: stable
 
-      - name: Cache Rust dependencies
-        uses: Leafwing-Studios/cargo-cache@v2.6.1
-        with:
-          sweep-cache: true
-
       - name: install valgrind
         if: matrix.do-valgrind
         run: |

metatomic-core/CMakeLists.txt

@@ -400,6 +400,13 @@ endif()
 target_link_libraries(metatomic::static INTERFACE nlohmann_json::nlohmann_json)
 target_link_libraries(metatomic::shared INTERFACE nlohmann_json::nlohmann_json)
 
+if(APPLE)
+    target_link_libraries(metatomic::static INTERFACE
+        "-framework Metal" "-framework CoreGraphics" "-framework CoreFoundation" "-framework Foundation" objc
+    )
+endif()
+
+
 if (BUILD_SHARED_LIBS)
     add_library(metatomic ALIAS metatomic::shared)
 else()

metatomic-core/Cargo.toml

@@ -19,9 +19,18 @@ dlpk = { version = "0.4", features = ["ndarray"]}
 json = "0.12"
 libloading = "0.9"
 ndarray = "0.17"
+
+# For serialization of the systems
 zip = { version = "8.6.0", default-features = false }
 byteorder = {version = "1"}
 
+# For custom kernels
+cudarc = {version = "0.19", default-features = false, features=["std", "cuda-13030", "driver", "nvrtc", "dynamic-loading"]}
+
+[target.'cfg(target_os = "macos")'.dependencies]
+objc2-metal = "0.3"
+objc2 = "0.6"
+objc2-foundation = "0.3"
 
 [build-dependencies]
 cbindgen = { version = "0.29", default-features = false }

metatomic-core/cmake/metatomic-config.in.cmake

@@ -78,7 +78,15 @@ if (@METATOMIC_INSTALL_BOTH_STATIC_SHARED@ OR NOT @BUILD_SHARED_LIBS@)
     )
 
     target_compile_features(metatomic::static INTERFACE cxx_std_17)
-    target_link_libraries(metatomic::static INTERFACE metatensor nlohmann_json::nlohmann_json)
+
+    target_link_libraries(metatomic::static INTERFACE metatensor)
+    target_link_libraries(metatomic::static INTERFACE nlohmann_json::nlohmann_json)
+
+    if(APPLE)
+        target_link_libraries(metatomic::static INTERFACE
+            "-framework Metal" "-framework CoreGraphics" "-framework CoreFoundation" "-framework Foundation" objc
+        )
+    endif()
 endif()
 
 # Export either the shared or static library as the metatomic target

metatomic-core/src/kernels/cpu.rs

@@ -0,0 +1,61 @@
+use dlpk::DLPackTensorRef;
+use ndarray::{ArrayView1, ArrayView2, ArrayViewD};
+
+use crate::Error;
+use super::ReferenceValue;
+
+/// Check that the values of an i32 DLPack tensor match the expected reference.
+///
+/// The tensor is converted to an ndarray view and compared element-wise and
+/// shape-wise against `reference`. The `description` is used verbatim in the
+/// error message on mismatch.
+///
+/// # Parameters
+/// - `tensor`: DLPack tensor with i32 data type
+/// - `reference`: expected values with the same shape as the tensor
+pub(crate) fn is_equal_i32(
+    tensor: DLPackTensorRef<'_>,
+    reference: &ReferenceValue<i32>,
+) -> Result<bool, Error> {
+    let values: ArrayViewD<i32> = tensor.try_into()?;
+    return Ok(values == reference.cpu.view());
+}
+
+macro_rules! validate_cell {
+    ($T: ty, $pbc: expr, $cell: expr) => {
+        let pbc_array: ArrayView1<bool> = $pbc.try_into()?;
+        let cell_array: ArrayView2<$T> = $cell.try_into()?;
+        for i in 0..3 {
+            if !pbc_array[i] && !cell_array.row(i).iter().all(|&x| x == 0.0) {
+                return Err(Error::InvalidParameter(format!(
+                    "invalid cell: for non-periodic dimensions, the corresponding \
+                    cell vector must be zero, but cell[{}] contains non-zero values",
+                    i
+                )));
+            }
+        }
+    };
+}
+
+/// Validate that cell vectors are zero for non-periodic dimensions on CPU.
+///
+/// Converts the DLPack tensors to ndarray views and checks that for every
+/// dimension where `pbc` is false, the corresponding row of `cell` contains
+/// only zeros.
+///
+/// # Parameters
+/// - `pbc`: 1D boolean tensor of length 3 (periodic boundary condition flags)
+/// - `cell`: 3x3 tensor (unit cell vectors as rows)
+pub(crate) fn validate_cell_pbc(
+    pbc: DLPackTensorRef<'_>,
+    cell: DLPackTensorRef<'_>,
+) -> Result<(), Error> {
+    let dtype = cell.dtype();
+    if dtype.bits == 32 {
+        validate_cell!(f32, pbc, cell);
+    } else {
+        assert_eq!(dtype.bits, 64);
+        validate_cell!(f64, pbc, cell);
+    }
+    return Ok(());
+}

metatomic-core/src/kernels/cuda.rs

@@ -0,0 +1,240 @@
+use std::collections::hash_map::Entry;
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex, LazyLock};
+
+use cudarc::driver::safe::DeviceRepr;
+use cudarc::driver::safe::{
+    CudaContext, CudaFunction, CudaModule, CudaStream, LaunchConfig, PushKernelArg,
+};
+use cudarc::nvrtc::compile_ptx;
+use dlpk::DLPackTensorRef;
+
+use crate::Error;
+use super::{ReferenceValue, StridedNDIndex};
+
+// CUDA kernel source compiled at runtime via NVRTC for the exact GPU
+const KERNEL_SRC: &str = include_str!("cuda_kernels.cu");
+
+unsafe impl DeviceRepr for StridedNDIndex {}
+
+/// Zero-cost wrapper to pass an existing device pointer as a CUDA kernel
+/// argument.
+///
+/// Does NOT own the memory — the caller (DLPack tensor) is responsible for
+/// lifetime and must ensure the pointer remains valid for the duration of the
+/// kernel launch.
+///
+/// The `#[repr(transparent)]` wrapper over `cudarc::driver::sys::CUdeviceptr`
+/// is passed to `PushKernelArg::arg()` which pushes the address of this struct
+/// on the host stack. CUDA reads 8 bytes from that address as the kernel
+/// parameter value, giving the kernel the correct device pointer.
+#[repr(transparent)]
+struct DevicePtrArg {
+    ptr: cudarc::driver::sys::CUdeviceptr,
+}
+
+unsafe impl DeviceRepr for DevicePtrArg {}
+
+/// Per-device cached resources: context, module, and kernel function handles.
+struct CudaKernelCache {
+    ctx: Arc<CudaContext>,
+    module: Arc<CudaModule>,
+    is_equal_i32: CudaFunction,
+    validate_cell_pbc_f32: CudaFunction,
+    validate_cell_pbc_f64: CudaFunction,
+}
+
+impl CudaKernelCache {
+    fn new(device_id: usize) -> Result<Self, Error> {
+        let ctx = CudaContext::new(device_id)
+            .map_err(|e| Error::Internal(format!("CudaContext::new({device_id}): {e}")))?;
+        let ptx = compile_ptx(KERNEL_SRC)
+            .map_err(|e| Error::Internal(format!("NVRTC compile failed: {e}")))?;
+        let module = ctx
+            .load_module(ptx)
+            .map_err(|e| Error::Internal(format!("PTX load failed: {e}")))?;
+        let is_equal_i32 = module
+            .load_function("is_equal_i32")
+            .map_err(|e| Error::Internal(format!("load_function(is_equal_i32): {e}")))?;
+        let validate_cell_pbc_f32 = module
+            .load_function("validate_cell_pbc_f32")
+            .map_err(|e| Error::Internal(format!("load_function(validate_cell_pbc_f32): {e}")))?;
+        let validate_cell_pbc_f64 = module
+            .load_function("validate_cell_pbc_f64")
+            .map_err(|e| Error::Internal(format!("load_function(validate_cell_pbc_f64): {e}")))?;
+        Ok(Self {
+            ctx,
+            module,
+            is_equal_i32,
+            validate_cell_pbc_f32,
+            validate_cell_pbc_f64,
+        })
+    }
+}
+
+static CUDA_CACHE: LazyLock<Mutex<HashMap<usize, CudaKernelCache>>> = LazyLock::new(|| Mutex::new(HashMap::new()));
+
+fn get_or_init(device_id: usize) -> Result<Arc<CudaStream>, Error> {
+    let mut cache = CUDA_CACHE.lock().expect("failed to lock CUDA_CACHE");
+    let entry = match cache.entry(device_id) {
+        Entry::Occupied(entry) => entry.into_mut(),
+        Entry::Vacant(entry) => entry.insert(CudaKernelCache::new(device_id)?),
+    };
+    Ok(entry.ctx.default_stream())
+}
+
+/// Extract a `CUdeviceptr` from a DLPack tensor's raw `data` + `byte_offset`.
+///
+/// # Safety
+///
+/// The returned `CUdeviceptr` is only valid as long as the DLPack tensor's
+/// backing memory is alive. The caller must ensure the tensor is not dropped
+/// before the kernel finishes execution.
+unsafe fn dlpack_to_device_ptr(tensor: &DLPackTensorRef<'_>) -> cudarc::driver::sys::CUdeviceptr {
+    debug_assert!(
+        tensor.device().device_type == dlpk::sys::DLDeviceType::kDLCUDA,
+        "dlpack_to_device_ptr called on non-CUDA tensor"
+    );
+    let raw_ptr = tensor.raw.data as u64;
+    (raw_ptr + tensor.raw.byte_offset) as cudarc::driver::sys::CUdeviceptr
+}
+
+/// Check that the values of a CUDA-resident i32 DLPack tensor match an expected
+/// reference array.
+///
+/// The comparison is performed entirely on-device: the existing GPU pointer
+/// from `tensor` is wrapped as a `DevicePtrArg`, the reference is uploaded to
+/// the GPU (and cached for subsequent calls), and a single-element result flag
+/// (`0` = ok, `1` = mismatch) is read back.
+#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
+pub(crate) fn is_equal_i32(tensor: DLPackTensorRef<'_>, reference: &ReferenceValue<i32>) -> Result<bool, Error> {
+    debug_assert!(
+        tensor.device().device_type == dlpk::sys::DLDeviceType::kDLCUDA,
+        "is_equal_i32 called on non-CUDA tensor"
+    );
+    debug_assert!(tensor.device().device_id >= 0, "is_equal_i32 called on invalid device_id");
+
+    let device_id = tensor.device().device_id as usize;
+    let stream = get_or_init(device_id)?;
+    let cache = CUDA_CACHE.lock().expect("failed to lock CUDA_CACHE");
+    let entry = &cache[&device_id];
+
+    let n_elements: i64 = tensor.shape().iter().product();
+
+    // Build strided index from the DLPack tensor (preserves actual strides)
+    let values_idx = StridedNDIndex::from_dlpack(&tensor);
+
+    // Wrap the existing GPU-allocated tensor pointer
+    let tensor_ptr = unsafe { DevicePtrArg { ptr: dlpack_to_device_ptr(&tensor) } };
+
+    // Upload reference values to GPU (cached after first call)
+    let (ref_dev, reference_idx) = reference.cuda.get_or_init(|| {
+        let slice = stream
+            .clone_htod(reference.cpu.as_slice().expect("reference should be contiguous"))
+            .expect("clone_htod reference failed");
+        let idx = StridedNDIndex::from_ndarray(&reference.cpu.view());
+        (slice, idx)
+    });
+
+    // Allocate result flag (initialized to 0 = no mismatch)
+    let mut result = stream.alloc_zeros::<i32>(1)
+        .map_err(|e| Error::Internal(format!("alloc_zeros: {e}")))?;
+
+    unsafe {
+        stream.launch_builder(&entry.is_equal_i32)
+            .arg(&tensor_ptr)
+            .arg(&values_idx)
+            .arg(ref_dev)
+            .arg(reference_idx)
+            .arg(&n_elements)
+            .arg(&mut result)
+            .launch(LaunchConfig::for_num_elems(u32::try_from(n_elements).expect("tensor too large for CUDA kernel")))
+            .map_err(|e| Error::Internal(format!("kernel launch (is_equal_i32): {e}")))?;
+    }
+
+    stream.synchronize()
+        .map_err(|e| Error::Internal(format!("device sync: {e}")))?;
+
+    let host = stream.clone_dtoh(&result)
+        .map_err(|e| Error::Internal(format!("clone_dtoh result: {e}")))?;
+
+    return Ok(host[0] == 0);
+}
+
+/// Validate that cell vectors are zero for non-periodic dimensions, on CUDA device.
+#[allow(clippy::cast_sign_loss)]
+pub(crate) fn validate_cell_pbc(
+    pbc: DLPackTensorRef<'_>,
+    cell: DLPackTensorRef<'_>,
+) -> Result<(), Error> {
+    debug_assert!(
+        pbc.device().device_type == dlpk::sys::DLDeviceType::kDLCUDA,
+        "validate_cell_pbc called on non-CUDA tensor"
+    );
+    debug_assert!(pbc.device().device_id >= 0, "validate_cell_pbc called on invalid device_id");
+    debug_assert!(cell.device() == pbc.device(), "pbc and cell must be on the same device");
+
+
+    let device_id = pbc.device().device_id as usize;
+    let stream = get_or_init(device_id)?;
+    let cache = CUDA_CACHE.lock().expect("failed to lock CUDA_CACHE");
+    let entry = &cache[&device_id];
+
+    let pbc_ptr = unsafe { DevicePtrArg { ptr: dlpack_to_device_ptr(&pbc) } };
+    let cell_ptr = unsafe { DevicePtrArg { ptr: dlpack_to_device_ptr(&cell) } };
+
+    let pbc_idx = StridedNDIndex::from_dlpack(&pbc);
+    let cell_idx = StridedNDIndex::from_dlpack(&cell);
+
+    let mut result = stream.alloc_zeros::<i32>(1)
+        .map_err(|e| Error::Internal(format!("alloc_zeros: {e}")))?;
+
+    if cell.dtype().bits == 32 {
+        unsafe {
+            stream.launch_builder(&entry.validate_cell_pbc_f32)
+                .arg(&pbc_ptr)
+                .arg(&pbc_idx)
+                .arg(&cell_ptr)
+                .arg(&cell_idx)
+                .arg(&mut result)
+                .launch(LaunchConfig {
+                    grid_dim: (1, 1, 1),
+                    block_dim: (3, 1, 1),
+                    shared_mem_bytes: 0,
+                })
+                .map_err(|e| Error::Internal(format!("kernel launch (f32): {e}")))?;
+        }
+    } else {
+        assert_eq!(cell.dtype().bits, 64, "validate_cell_pbc: unsupported cell dtype");
+        unsafe {
+            stream.launch_builder(&entry.validate_cell_pbc_f64)
+                .arg(&pbc_ptr)
+                .arg(&pbc_idx)
+                .arg(&cell_ptr)
+                .arg(&cell_idx)
+                .arg(&mut result)
+                .launch(LaunchConfig {
+                    grid_dim: (1, 1, 1),
+                    block_dim: (3, 1, 1),
+                    shared_mem_bytes: 0,
+                })
+                .map_err(|e| Error::Internal(format!("kernel launch (f64): {e}")))?;
+        }
+    }
+
+    stream.synchronize()
+        .map_err(|e| Error::Internal(format!("device sync: {e}")))?;
+
+    let host = stream.clone_dtoh(&result)
+        .map_err(|e| Error::Internal(format!("clone_dtoh result: {e}")))?;
+
+    if host[0] != 0 {
+        let dim = host[0] - 1;
+        return Err(Error::InvalidParameter(format!(
+            "invalid cell: for non-periodic dimensions, the corresponding \
+             cell vector must be zero, but cell[{}] contains non-zero values",
+            dim
+        )));
+    }
+    Ok(())
+}