diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index 282f915d..584fd669 100644 --- a/.github/workflows/python-tests.yml +++ b/.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 diff --git a/.github/workflows/rust-tests.yml b/.github/workflows/rust-tests.yml index f9c2a1c4..36eb59bf 100644 --- a/.github/workflows/rust-tests.yml +++ b/.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: | diff --git a/.github/workflows/torch-tests.yml b/.github/workflows/torch-tests.yml index 93661740..d9dd09dc 100644 --- a/.github/workflows/torch-tests.yml +++ b/.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: | diff --git a/metatomic-core/CMakeLists.txt b/metatomic-core/CMakeLists.txt index 128d52f9..e7dd396a 100644 --- a/metatomic-core/CMakeLists.txt +++ b/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() diff --git a/metatomic-core/Cargo.toml b/metatomic-core/Cargo.toml index 934b6fe2..db3482f6 100644 --- a/metatomic-core/Cargo.toml +++ b/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 } diff --git a/metatomic-core/cmake/metatomic-config.in.cmake b/metatomic-core/cmake/metatomic-config.in.cmake index 15652cba..4bccc34b 100644 --- a/metatomic-core/cmake/metatomic-config.in.cmake +++ b/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 diff --git a/metatomic-core/src/kernels/cpu.rs b/metatomic-core/src/kernels/cpu.rs new file mode 100644 index 00000000..1225a904 --- /dev/null +++ b/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, +) -> Result { + let values: ArrayViewD = tensor.try_into()?; + return Ok(values == reference.cpu.view()); +} + +macro_rules! validate_cell { + ($T: ty, $pbc: expr, $cell: expr) => { + let pbc_array: ArrayView1 = $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(()); +} diff --git a/metatomic-core/src/kernels/cuda.rs b/metatomic-core/src/kernels/cuda.rs new file mode 100644 index 00000000..b1fcb833 --- /dev/null +++ b/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, + module: Arc, + is_equal_i32: CudaFunction, + validate_cell_pbc_f32: CudaFunction, + validate_cell_pbc_f64: CudaFunction, +} + +impl CudaKernelCache { + fn new(device_id: usize) -> Result { + 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>> = LazyLock::new(|| Mutex::new(HashMap::new())); + +fn get_or_init(device_id: usize) -> Result, 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) -> Result { + 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::(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::(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(()) +} diff --git a/metatomic-core/src/kernels/cuda_kernels.cu b/metatomic-core/src/kernels/cuda_kernels.cu new file mode 100644 index 00000000..3fc47ad1 --- /dev/null +++ b/metatomic-core/src/kernels/cuda_kernels.cu @@ -0,0 +1,95 @@ +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +#define MAX_NDIM 7 + +/// Multi-dimensional strided index (up to MAX_NDIM dimensions). +/// Decomposes a flat linear index into multi-dimensional coordinates from the +/// shape, then computes the strided memory offset using the stride array. +/// +/// WARNING: any change here needs to be reflected in the Rust and Metal sources. +struct StridedNDIndex { + int64_t ndim; + int64_t shape[MAX_NDIM]; + int64_t strides[MAX_NDIM]; + + /// Get the offset from the start of the array for a given flat index + __device__ int64_t offset(int64_t flat_idx) const { + int64_t off = 0; + for (int d = this->ndim - 1; d >= 0; d--) { + int64_t coord = flat_idx % this->shape[d]; + flat_idx /= this->shape[d]; + off += coord * this->strides[d]; + } + return off; + } +}; + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +extern "C" __global__ void is_equal_i32( + const int* values, + StridedNDIndex values_idx, + const int* reference, + StridedNDIndex reference_idx, + int64_t n, + int* mismatch +) { + int64_t i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + int64_t value_offset = values_idx.offset(i); + int64_t reference_offset = reference_idx.offset(i); + if (values[value_offset] != reference[reference_offset]) { + atomicMax(mismatch, 1); + } + } +} + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template +__device__ void validate_cell_pbc_impl( + const bool* pbc, + StridedNDIndex pbc_idx, + const T* cell, + StridedNDIndex cell_idx, + int* mismatch_idx +) { + int i = threadIdx.x; + if (i < 3) { + if (!pbc[pbc_idx.offset(i)]) { + if ( + cell[cell_idx.offset(i * 3 + 0)] != T(0) || + cell[cell_idx.offset(i * 3 + 1)] != T(0) || + cell[cell_idx.offset(i * 3 + 2)] != T(0) + ) { + atomicMax(mismatch_idx, i + 1); + } + } + } +} + +extern "C" __global__ void validate_cell_pbc_f32( + const bool* pbc, + StridedNDIndex pbc_idx, + const float* cell, + StridedNDIndex cell_idx, + int* mismatch_idx +) { + validate_cell_pbc_impl(pbc, pbc_idx, cell, cell_idx, mismatch_idx); +} + +extern "C" __global__ void validate_cell_pbc_f64( + const bool* pbc, + StridedNDIndex pbc_idx, + const double* cell, + StridedNDIndex cell_idx, + int* mismatch_idx +) { + validate_cell_pbc_impl(pbc, pbc_idx, cell, cell_idx, mismatch_idx); +} + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// diff --git a/metatomic-core/src/kernels/metal.rs b/metatomic-core/src/kernels/metal.rs new file mode 100644 index 00000000..67339d85 --- /dev/null +++ b/metatomic-core/src/kernels/metal.rs @@ -0,0 +1,311 @@ +use std::collections::{HashMap, hash_map::Entry}; +use std::ptr::NonNull; +use std::sync::Mutex; +use std::sync::LazyLock; + +use objc2::rc::Retained; +use objc2::runtime::ProtocolObject; +use objc2_foundation::ns_string; + +use objc2_metal::{ + MTLBuffer, MTLCommandBuffer, MTLCommandEncoder, MTLCommandQueue, + MTLComputeCommandEncoder, MTLComputePipelineState, + MTLCreateSystemDefaultDevice, MTLCompileOptions, + MTLDevice, MTLLibrary, MTLResourceOptions, MTLSize, +}; + +use dlpk::DLPackTensorRef; + +use crate::Error; +use super::{ReferenceValue, StridedNDIndex}; + +// Small wrapper around MTLBuffer to implement Send and Sync, since the data is +// read-only after initialization. +pub(crate) struct MetalBuffer(Retained>); + +unsafe impl Send for MetalBuffer {} +unsafe impl Sync for MetalBuffer {} + +impl std::ops::Deref for MetalBuffer { + type Target = ProtocolObject; + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +const KERNEL_SRC: &str = include_str!("metal_kernels.metal"); + +/// Cached metal ressources: device, command queue, and pipeline states for kernels. +struct MetalKernelCache { + device: Retained>, + queue: Retained>, + is_equal_i32: Retained>, + validate_cell_pbc_f32: Retained>, +} + +impl MetalKernelCache { + fn new(device_id: usize) -> Result { + let device = MTLCreateSystemDefaultDevice() + .ok_or_else(|| Error::Internal(format!("no Metal device found for id {device_id}")))?; + + let library = device + .newLibraryWithSource_options_error( + ns_string!(KERNEL_SRC), + Some(&MTLCompileOptions::new()), + ) + .map_err(|e| Error::Internal(format!("MSL compile failed: {e}")))?; + + let is_equal_i32 = make_pipeline(&device, &library, "is_equal_i32")?; + let validate_cell_pbc_f32 = make_pipeline(&device, &library, "validate_cell_pbc_f32")?; + + let queue = device + .newCommandQueue() + .ok_or_else(|| Error::Internal("failed to create command queue".into()))?; + + Ok(Self { + device, + queue, + is_equal_i32, + validate_cell_pbc_f32, + }) + } +} + +fn make_pipeline( + device: &ProtocolObject, + library: &ProtocolObject, + name: &str, +) -> Result>, Error> { + use objc2_foundation::NSString; + + let ns_name = NSString::from_str(name); + let function = library + .newFunctionWithName(&ns_name) + .ok_or_else(|| Error::Internal(format!("get_function({name}): not found")))?; + + device + .newComputePipelineStateWithFunction_error(&function) + .map_err(|e| Error::Internal(format!("pipeline state ({name}): {e}"))) +} + +static METAL_CACHE: LazyLock>> = LazyLock::new(|| Mutex::new(HashMap::new())); + +fn get_or_init(cache: &mut HashMap, device_id: usize) -> Result<&MetalKernelCache, Error> { + let entry = match cache.entry(device_id) { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => entry.insert(MetalKernelCache::new(device_id)?), + }; + Ok(entry) +} + +/// Compute the byte span of a DLPack tensor's data (including gaps from +/// strides). +#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] +fn tensor_num_bytes(tensor: &DLPackTensorRef<'_>) -> usize { + let elem_size = tensor.dtype().bits as usize / 8; + let shape = tensor.shape(); + match tensor.strides() { + None => shape.iter().map(|&s| s as usize).product::() * elem_size, + Some(strides) => { + let max_idx: i64 = shape.iter() + .zip(strides.iter()) + .map(|(&s, &st)| (s - 1) * st) + .sum(); + (max_idx as usize + 1) * elem_size + } + } +} + +/// Extract a raw pointer to the tensor's data, accounting for byte_offset. +/// +/// # Safety +/// +/// The returned pointer is only valid as long as the DLPack tensor's backing +/// memory is alive. +#[allow(clippy::cast_possible_truncation)] +fn dlpack_data_ptr(tensor: &DLPackTensorRef<'_>) -> *const std::ffi::c_void { + unsafe { + tensor.raw.data.cast::().add(tensor.raw.byte_offset as usize).cast() + } +} + +/// Check that the values of a Metal-resident i32 DLPack tensor match an expected +/// reference array. +#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)] +pub(crate) fn is_equal_i32(tensor: DLPackTensorRef<'_>, reference: &ReferenceValue) -> Result { + let device_id = tensor.device().device_id as usize; + let mut lock = METAL_CACHE.lock().expect("failed to lock METAL_CACHE"); + let cache = get_or_init(&mut lock, device_id)?; + + let n_elements: usize = tensor.shape().iter().map(|&s| s as usize).product(); + let ref_bytes = n_elements * std::mem::size_of::(); + + // Build strided index for the values + let values_idx = StridedNDIndex::from_dlpack(&tensor); + + // Upload reference values to Metal (cached after first call) + let (ref_buf, reference_idx) = reference.metal.get_or_init(|| { + let ref_bytes = reference.cpu.len() * std::mem::size_of::(); + let ref_ptr: *const std::ffi::c_void = reference.cpu.as_slice() + .expect("reference should be contiguous") + .as_ptr() + .cast(); + let buf = unsafe { + cache.device.newBufferWithBytes_length_options( + NonNull::new(ref_ptr.cast_mut()).expect("reference pointer must not be null"), + ref_bytes, + MTLResourceOptions::empty(), + ).expect("failed to create reference buffer") + }; + let idx = StridedNDIndex::from_ndarray(&reference.cpu.view()); + (MetalBuffer(buf), idx) + }); + + let values_buf = unsafe { + cache.device.newBufferWithBytes_length_options( + NonNull::new(dlpack_data_ptr(&tensor).cast_mut()).expect("values pointer must not be null"), + tensor_num_bytes(&tensor), + MTLResourceOptions::empty(), + ).expect("failed to create values buffer") + }; + let result_buf = unsafe { + cache.device.newBufferWithBytes_length_options( + NonNull::from(&0i32).cast(), + std::mem::size_of::(), + MTLResourceOptions::empty(), + ).expect("failed to create result buffer") + }; + + objc2::rc::autoreleasepool(|_| { + let cmd_buf = cache.queue.commandBuffer().expect("failed to create command buffer"); + let encoder = cmd_buf.computeCommandEncoder().expect("failed to create compute encoder"); + + encoder.setComputePipelineState(&cache.is_equal_i32); + unsafe { + encoder.setBuffer_offset_atIndex(Some(&*values_buf), 0, 0); + + encoder.setBytes_length_atIndex( + NonNull::from(&values_idx).cast(), + std::mem::size_of::(), + 1, + ); + + encoder.setBuffer_offset_atIndex(Some(&*ref_buf), 0, 2); + + encoder.setBytes_length_atIndex( + NonNull::from(&reference_idx).cast(), + std::mem::size_of::(), + 3, + ); + + encoder.setBytes_length_atIndex( + NonNull::from(&(n_elements as u32)).cast(), + std::mem::size_of::(), + 4, + ); + + encoder.setBuffer_offset_atIndex(Some(&*result_buf), 0, 5); + } + + let tg_size = 32; + let tg_count = n_elements.div_ceil(tg_size); + encoder.dispatchThreadgroups_threadsPerThreadgroup( + MTLSize { width: tg_count, height: 1, depth: 1 }, + MTLSize { width: tg_size, height: 1, depth: 1 }, + ); + encoder.endEncoding(); + cmd_buf.commit(); + cmd_buf.waitUntilCompleted(); + }); + + let result = unsafe { + *result_buf.contents().as_ptr().cast::() + }; + return Ok(result == 0); +} + +/// Validate that cell vectors are zero for non-periodic dimensions on Metal. +#[allow(clippy::cast_sign_loss)] +pub(crate) fn validate_cell_pbc( + pbc: DLPackTensorRef<'_>, + cell: DLPackTensorRef<'_>, +) -> Result<(), Error> { + let device_id = pbc.device().device_id as usize; + let mut lock = METAL_CACHE.lock().expect("failed to lock METAL_CACHE"); + let cache = get_or_init(&mut lock, device_id)?; + + let pbc_idx = StridedNDIndex::from_dlpack(&pbc); + let cell_idx = StridedNDIndex::from_dlpack(&cell); + + let pbc_buf = unsafe { + cache.device.newBufferWithBytes_length_options( + NonNull::new(dlpack_data_ptr(&pbc).cast_mut()).expect("pbc pointer must not be null"), + tensor_num_bytes(&pbc), + MTLResourceOptions::empty(), + ).expect("failed to create pbc buffer") + }; + let cell_buf = unsafe { + cache.device.newBufferWithBytes_length_options( + NonNull::new(dlpack_data_ptr(&cell).cast_mut()).expect("cell pointer must not be null"), + tensor_num_bytes(&cell), + MTLResourceOptions::empty(), + ).expect("failed to create cell buffer") + }; + let result_buf = unsafe { + cache.device.newBufferWithBytes_length_options( + NonNull::from(&0i32).cast(), + std::mem::size_of::(), + MTLResourceOptions::empty(), + ).expect("failed to create result buffer") + }; + + objc2::rc::autoreleasepool(|_| { + let cmd_buf = cache.queue.commandBuffer().expect("failed to create command buffer"); + let encoder = cmd_buf.computeCommandEncoder().expect("failed to create compute encoder"); + + assert!(cell.dtype().bits == 32, "only float32 is supported on Metal"); + + encoder.setComputePipelineState(&cache.validate_cell_pbc_f32); + unsafe { + encoder.setBuffer_offset_atIndex(Some(&*pbc_buf), 0, 0); + + encoder.setBytes_length_atIndex( + NonNull::from(&pbc_idx).cast(), + std::mem::size_of::(), + 1, + ); + + encoder.setBuffer_offset_atIndex(Some(&*cell_buf), 0, 2); + + encoder.setBytes_length_atIndex( + NonNull::from(&cell_idx).cast(), + std::mem::size_of::(), + 3, + ); + + encoder.setBuffer_offset_atIndex(Some(&*result_buf), 0, 4); + } + + encoder.dispatchThreadgroups_threadsPerThreadgroup( + MTLSize { width: 1, height: 1, depth: 1 }, + MTLSize { width: 3, height: 1, depth: 1 }, + ); + encoder.endEncoding(); + cmd_buf.commit(); + cmd_buf.waitUntilCompleted(); + }); + + let result = unsafe { + *result_buf.contents().as_ptr().cast::() + }; + + if result != 0 { + let dim = result - 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(()) +} diff --git a/metatomic-core/src/kernels/metal_kernels.metal b/metatomic-core/src/kernels/metal_kernels.metal new file mode 100644 index 00000000..aff14a43 --- /dev/null +++ b/metatomic-core/src/kernels/metal_kernels.metal @@ -0,0 +1,76 @@ +#include +using namespace metal; + +// --------------------------------------------------------------------------- +// Multi-dimensional strided index helper (up to MAX_NDIM dimensions). +// +// Decomposes a flat linear index into multi-dimensional coordinates based on +// the shape and then computes the strided memory offset using the stride +// array. +// +// WARNING: the layout of this struct must match both the CUDA +// (cuda_kernels.cu) and Rust (kernels/mod.rs) definitions. +// --------------------------------------------------------------------------- +constant long MAX_NDIM [[maybe_unused]] = 7; + +struct StridedNDIndex { + long ndim; + long shape[MAX_NDIM]; + long strides[MAX_NDIM]; + + long offset(long flat_idx) const { + long off = 0; + for (int d = ndim - 1; d >= 0; d--) { + long coord = flat_idx % shape[d]; + flat_idx /= shape[d]; + off += coord * strides[d]; + } + return off; + } +}; + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +kernel void is_equal_i32( + [[buffer(0)]] device const int* values, + [[buffer(1)]] constant StridedNDIndex& values_idx, + [[buffer(2)]] device const int* reference, + [[buffer(3)]] constant StridedNDIndex& reference_idx, + [[buffer(4)]] constant uint& n, + [[buffer(5)]] device atomic_int* mismatch, + [[thread_position_in_grid]] uint gid +) { + if (gid < n) { + long v_off = values_idx.offset(gid); + long r_off = reference_idx.offset(gid); + if (values[v_off] != reference[r_off]) { + atomic_fetch_max_explicit(mismatch, 1, memory_order_relaxed); + } + } +} + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +/// Validate cell vectors against PBC flags (f32 only on Metal). +kernel void validate_cell_pbc_f32( + [[buffer(0)]] device const bool* pbc, + [[buffer(1)]] constant StridedNDIndex& pbc_idx, + [[buffer(2)]] device const float* cell, + [[buffer(3)]] constant StridedNDIndex& cell_idx, + [[buffer(4)]] device atomic_int* mismatch_idx, + [[thread_position_in_threadgroup]] uint tid +) { + if (tid < 3) { + if (!pbc[pbc_idx.offset(tid)]) { + if ( + cell[cell_idx.offset(tid * 3 + 0)] != 0.0f || + cell[cell_idx.offset(tid * 3 + 1)] != 0.0f || + cell[cell_idx.offset(tid * 3 + 2)] != 0.0f + ) { + atomic_fetch_max_explicit(mismatch_idx, int(tid + 1), memory_order_relaxed); + } + } + } +} diff --git a/metatomic-core/src/kernels/mod.rs b/metatomic-core/src/kernels/mod.rs new file mode 100644 index 00000000..6c586d8c --- /dev/null +++ b/metatomic-core/src/kernels/mod.rs @@ -0,0 +1,168 @@ +use std::sync::OnceLock; + +use cudarc::driver::CudaSlice; +use dlpk::sys::DLDeviceType; +use dlpk::DLPackTensorRef; +use ndarray::{ArrayD, ArrayViewD}; + +use crate::Error; + +mod cpu; +mod cuda; + +#[cfg(target_os = "macos")] +mod metal; + +const MAX_NDIM: usize = 7; + +/// Multi-dimensional strided index (up to MAX_NDIM dimensions). +/// +/// Decomposes a flat linear index into multi-dimensional coordinates from the +/// shape, then computes the strided memory offset using the stride array. +/// +/// WARNING: any change here needs to be reflected in the CUDA and Metal sources. +#[repr(C)] +pub(crate) struct StridedNDIndex { + pub(crate) ndim: i64, + pub(crate) shape: [i64; MAX_NDIM], + pub(crate) strides: [i64; MAX_NDIM], +} + +#[allow(clippy::cast_possible_wrap)] +impl StridedNDIndex { + /// Create a `StridedNDIndex` from a DLPack tensor's shape and strides. + pub(crate) fn from_dlpack(tensor: &DLPackTensorRef<'_>) -> Self { + Self::from_shape_strides(tensor.shape(), tensor.strides()) + } + + /// Create a `StridedNDIndex` from an ndarray view's shape and strides. + pub(crate) fn from_ndarray(array: &ArrayViewD<'_, T>) -> Self { + let shape: Vec = array.shape().iter().map(|&s| s as i64).collect(); + let strides: Vec = array.strides().iter().map(|&s| s as i64).collect(); + Self::from_shape_strides(&shape, Some(&strides)) + } + + /// Create a `StridedNDIndex` from shape and optional strides. + /// + /// If strides is `None`, the strides are computed as if the array were + /// contiguous (row-major / C-contiguous). + pub(crate) fn from_shape_strides(shape: &[i64], strides: Option<&[i64]>) -> Self { + let ndim = shape.len(); + assert!( + ndim <= MAX_NDIM, + "StridedNDIndex only supports up to {MAX_NDIM} dimensions, got {ndim}" + ); + let mut shape_arr = [0i64; MAX_NDIM]; + let mut strides_arr = [0i64; MAX_NDIM]; + + // Contiguous fallback strides (row-major / C-contiguous) + let mut acc: i64 = 1; + for i in (0..ndim).rev() { + shape_arr[i] = shape[i]; + strides_arr[i] = acc; + acc *= shape[i]; + } + + if let Some(strides) = strides { + strides_arr[..ndim].copy_from_slice(&strides[..ndim]); + } + StridedNDIndex { ndim: ndim as i64, shape: shape_arr, strides: strides_arr } + } +} + +/// Store and cache reference values for different backends (CPU, CUDA, Metal). +pub struct ReferenceValue { + /// The reference values stored on the CPU, always there + pub(crate) cpu: ArrayD, + /// Reference values stored on CUDA, intialized on first use from the CPU values + pub(crate) cuda: OnceLock<(CudaSlice, StridedNDIndex)>, + #[cfg(target_os = "macos")] + /// Reference values stored on Metal, intialized on first use from the CPU values + pub(crate) metal: OnceLock<(metal::MetalBuffer, StridedNDIndex)>, +} + +impl ReferenceValue { + pub(crate) fn new(cpu: ArrayD) -> Self { + Self { + cpu, + cuda: OnceLock::new(), + #[cfg(target_os = "macos")] + metal: OnceLock::new(), + } + } +} + +/// Check that the values of an i32 DLPack tensor match the expected reference. +/// +/// This dispatches to the appropriate backend based on the device of `tensor`. +/// +/// # 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) -> Result { + match tensor.device().device_type { + DLDeviceType::kDLCPU | DLDeviceType::kDLCUDAHost | DLDeviceType::kDLROCMHost => { + cpu::is_equal_i32(tensor, reference) + } + DLDeviceType::kDLCUDA | DLDeviceType::kDLCUDAManaged => { + cuda::is_equal_i32(tensor, reference) + } + DLDeviceType::kDLMetal => { + #[cfg(target_os = "macos")] { + metal::is_equal_i32(tensor, reference) + } + #[cfg(not(target_os = "macos"))] { + Err(Error::Internal( + "Metal backend is only available on macOS".into(), + )) + } + } + _ => { + eprintln!( + "is_equal_i32 for device {:?} is not implemented", + tensor.device() + ); + Ok(true) + } + } +} + +/// Validate that cell vectors are zero for non-periodic dimensions. +/// +/// This dispatches to the appropriate backend based on the device of `pbc`. +/// +/// # 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> { + debug_assert!( + pbc.device() == cell.device(), + "pbc and cell must be on the same device" + ); + + match pbc.device().device_type { + DLDeviceType::kDLCPU | DLDeviceType::kDLCUDAHost | DLDeviceType::kDLROCMHost => { + cpu::validate_cell_pbc(pbc, cell) + } + DLDeviceType::kDLCUDA | DLDeviceType::kDLCUDAManaged => { + cuda::validate_cell_pbc(pbc, cell) + } + DLDeviceType::kDLMetal => { + #[cfg(target_os = "macos")] { + metal::validate_cell_pbc(pbc, cell) + } + #[cfg(not(target_os = "macos"))] { + Err(Error::Internal( + "Metal backend is only available on macOS".into(), + )) + } + } + _ => { + eprintln!( + "Cell/PBC validation for device {:?} is not implemented", + pbc.device() + ); + Ok(()) + } + } +} diff --git a/metatomic-core/src/lib.rs b/metatomic-core/src/lib.rs index 18dd99ca..954693f5 100644 --- a/metatomic-core/src/lib.rs +++ b/metatomic-core/src/lib.rs @@ -25,6 +25,8 @@ pub use self::metadata::{Device, DType, ModelCapabilities, ModelMetadata, PairLi mod quantities; pub use self::quantities::{Quantity, SampleKind, Gradients}; +mod kernels; + mod system; pub use self::system::System; diff --git a/metatomic-core/src/system.rs b/metatomic-core/src/system.rs index b3aeb8e3..740cece0 100644 --- a/metatomic-core/src/system.rs +++ b/metatomic-core/src/system.rs @@ -1,17 +1,36 @@ use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::LazyLock; -use dlpk::sys::{DLDataType, DLDevice, DLDeviceType}; +use dlpk::sys::{DLDataType, DLDevice}; use dlpk::{DLPackTensor, DLPackTensorRef}; use metatensor::{TensorBlock, TensorMap}; use crate::{Error, PairListOptions}; +use crate::kernels::ReferenceValue; /// Names that can never be used as custom data in a system static INVALID_DATA_NAMES: LazyLock> = LazyLock::new(|| { HashSet::from(["types", "type", "positions", "position", "cell", "neighbors", "neighbor", "pair", "pairs"]) }); +static XYZ_REFERENCE: LazyLock> = LazyLock::new(|| { + ReferenceValue::new( + ndarray::ArrayD::from_shape_vec( + ndarray::IxDyn(&[3usize, 1]), + vec![0i32, 1, 2], + ).unwrap() + ) +}); + +static DISTANCE_REFERENCE: LazyLock> = LazyLock::new(|| { + ReferenceValue::new( + ndarray::ArrayD::from_shape_vec( + ndarray::IxDyn(&[1usize, 1]), + vec![0i32], + ).unwrap() + ) +}); + /// Storage for an atomistic system. /// /// This owns the raw DLPack tensors and metatensor objects used at FFI @@ -51,9 +70,7 @@ impl System { custom_data: HashMap::new(), }; - if system.device().device_type == DLDeviceType::kDLCPU { - validate_cpu_system_data(&system)?; - } + crate::kernels::validate_cell_pbc(system.pbc(), system.cell())?; return Ok(system); } @@ -121,12 +138,18 @@ impl System { )); } - #[allow(clippy::collapsible_if)] - if components[0].device().device_type == DLDeviceType::kDLCPU { - if components[0][0] != [0] || components[0][1] != [1] || components[0][2] != [2] { + { + let mts_array = components[0].values(); + let dl_tensor = mts_array.as_dlpack( + components[0].device(), + None, + dlpk::sys::DLPackVersion::current(), + )?; + + if !crate::kernels::is_equal_i32(dl_tensor.as_ref(), &XYZ_REFERENCE)? { return Err(Error::InvalidParameter( - "invalid components for `pairs`: the 'xyz' \ - component should contain [0, 1, 2]".into() + "invalid components for `pairs`: the 'xyz' component should \ + contain [[0], [1], [2]]".into() )); } } @@ -139,9 +162,15 @@ impl System { )); } - #[allow(clippy::collapsible_if)] - if properties.device().device_type == DLDeviceType::kDLCPU { - if properties[0] != [0] { + { + let mts_array = properties.values(); + let dl_tensor = mts_array.as_dlpack( + properties.device(), + None, + dlpk::sys::DLPackVersion::current(), + )?; + + if !crate::kernels::is_equal_i32(dl_tensor.as_ref(), &DISTANCE_REFERENCE)? { return Err(Error::InvalidParameter( "invalid properties for `pairs`: the 'distance' property \ should contain [0]".into() @@ -344,36 +373,6 @@ fn validate_system_tensors( return Ok(()); } -fn validate_cpu_system_data(system: &System) -> Result<(), Error> { - let pbc_array: ndarray::ArrayView1 = system.pbc().try_into()?; - - if system.dtype().bits == 32 { - let cell_array: ndarray::ArrayView2 = system.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 - ))); - } - } - } else { - let cell_array: ndarray::ArrayView2 = system.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 - ))); - } - } - } - - return Ok(()); -} - #[cfg(test)] pub(crate) use tests::test_system;