diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 633cb9b..eff6743 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -23,7 +23,7 @@ jobs: runs-on: ubuntu-latest if: github.event_name == 'pull_request' steps: - - uses: christophebedard/dco-check@v0.5.0 + - uses: christophebedard/dco-check@0.5.1 with: token: ${{ secrets.GITHUB_TOKEN }} diff --git a/Cargo.lock b/Cargo.lock index 5e2cd33..8919106 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -338,6 +338,27 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "csv" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde_core", +] + +[[package]] +name = "csv-core" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" +dependencies = [ + "memchr", +] + [[package]] name = "either" version = "1.15.0" @@ -474,6 +495,12 @@ dependencies = [ "rustversion", ] +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + [[package]] name = "js-sys" version = "0.3.91" @@ -653,6 +680,7 @@ dependencies = [ "mohu-buffer", "mohu-dtype", "mohu-error", + "mohu-testing", "num-complex", "num-traits", "rayon", @@ -675,6 +703,7 @@ name = "mohu-io" version = "0.1.0" dependencies = [ "arrow", + "csv", "memmap2", "mohu-core", "serde", diff --git a/Cargo.toml b/Cargo.toml index b658360..d9454e5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -104,6 +104,7 @@ serde = { version = "1", features = ["derive"] } # ── I/O ─────────────────────────────────────────────────────────────────────── memmap2 = "0.9" +csv = "1.4" # ── Python bindings ─────────────────────────────────────────────────────────── pyo3 = { version = "0.23", features = ["extension-module"] } diff --git a/README.md b/README.md index ca96ed0..e0218f9 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,15 @@ Polars proved you can rewrite the data layer in Rust and win. mohu is that same Early. The foundation is being laid. If you believe the Python numerical stack deserves a rewrite, watch this repo or contribute. +## windows + +On Windows, use `scripts/use-llvm-mingw.ps1` to run the `mohu-io` test target +with the LLVM-MinGW linker setup that works in this workspace. + +```powershell +.\scripts\use-llvm-mingw.ps1 +``` + ## built with - [Rust](https://rust-lang.org) diff --git a/crates/mohu-array/src/array.rs b/crates/mohu-array/src/array.rs index e69de29..8b13789 100644 --- a/crates/mohu-array/src/array.rs +++ b/crates/mohu-array/src/array.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-array/src/iter.rs b/crates/mohu-array/src/iter.rs index e69de29..8b13789 100644 --- a/crates/mohu-array/src/iter.rs +++ b/crates/mohu-array/src/iter.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-array/src/shape.rs b/crates/mohu-array/src/shape.rs index e69de29..8b13789 100644 --- a/crates/mohu-array/src/shape.rs +++ b/crates/mohu-array/src/shape.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-array/src/slice.rs b/crates/mohu-array/src/slice.rs index e69de29..8b13789 100644 --- a/crates/mohu-array/src/slice.rs +++ b/crates/mohu-array/src/slice.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-array/src/view.rs b/crates/mohu-array/src/view.rs index e69de29..8b13789 100644 --- a/crates/mohu-array/src/view.rs +++ b/crates/mohu-array/src/view.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-buffer/examples/alloc_and_pool.rs b/crates/mohu-buffer/examples/alloc_and_pool.rs index 8f402ca..4ad6f50 100644 --- a/crates/mohu-buffer/examples/alloc_and_pool.rs +++ b/crates/mohu-buffer/examples/alloc_and_pool.rs @@ -7,10 +7,8 @@ // directly, but is critical for high-performance scientific computing. use mohu_buffer::{ - AllocHandle, AllocStats, Strategy, - BufferPool, PoolStats, GLOBAL_POOL, - SIMD_ALIGN, CACHE_LINE, - Buffer, + AllocHandle, AllocStats, Buffer, CACHE_LINE, GLOBAL_POOL, SIMD_ALIGN, + Strategy, }; use mohu_dtype::DType; @@ -22,22 +20,32 @@ fn main() { // ── Direct allocation with AllocHandle ───────────────────────────────── println!("\n── AllocHandle (raw aligned allocation) ──"); let handle = AllocHandle::alloc(1024, SIMD_ALIGN).unwrap(); - println!("Allocated: len={}, align={}, strategy={:?}", - handle.len(), handle.align(), handle.strategy()); + println!( + "Allocated: len={}, align={}, strategy={:?}", + handle.len(), + handle.align(), + handle.strategy() + ); println!(" Pointer: {:p}", handle.as_ptr()); println!(" Aligned to 64 bytes: {}", handle.is_aligned_to(64)); assert!(matches!(handle.strategy(), Strategy::Heap)); // Zeroed allocation let zeroed = AllocHandle::alloc_zeroed(512, SIMD_ALIGN).unwrap(); - println!("\nZeroed alloc: len={}, all zeros: {}", + println!( + "\nZeroed alloc: len={}, all zeros: {}", zeroed.len(), - zeroed.as_byte_slice().iter().all(|&b| b == 0)); + zeroed.as_byte_slice().iter().all(|&b| b == 0) + ); // Zero-size allocation let empty = AllocHandle::alloc(0, SIMD_ALIGN).unwrap(); - println!("Zero-size: len={}, is_empty={}, strategy={:?}", - empty.len(), empty.is_empty(), empty.strategy()); + println!( + "Zero-size: len={}, is_empty={}, strategy={:?}", + empty.len(), + empty.is_empty(), + empty.strategy() + ); // ── Global allocation stats ──────────────────────────────────────────── println!("\n── Allocation stats (global) ──"); @@ -89,13 +97,17 @@ fn main() { println!("\n── Pool warm-up ──"); pool.warm(&[8192, 32768], 2).unwrap(); let after_warm = pool.stats(); - println!("After warm-up: cached={} bytes, blocks={}", - after_warm.cached_bytes, after_warm.cached_blocks); + println!( + "After warm-up: cached={} bytes, blocks={}", + after_warm.cached_bytes, after_warm.cached_blocks + ); // Per-size-class breakdown for sc in pool.size_class_stats() { - println!(" Class {:>8} bytes: {} handles ({} bytes cached)", - sc.size_class, sc.cached_handles, sc.cached_bytes); + println!( + " Class {:>8} bytes: {} handles ({} bytes cached)", + sc.size_class, sc.cached_handles, sc.cached_bytes + ); } // ── Pool trim (reduce memory footprint) ──────────────────────────────── @@ -106,8 +118,11 @@ fn main() { println!("\n── Buffer reuse pattern ──"); for i in 0..3 { let buf = Buffer::zeros(DType::F32, &[1000]).unwrap(); - println!(" Iteration {i}: allocated {} elements, {} bytes", - buf.len(), buf.nbytes()); + println!( + " Iteration {i}: allocated {} elements, {} bytes", + buf.len(), + buf.nbytes() + ); // buf is dropped here, its allocation returns to the pool } println!(" Pool after loop: {} cached bytes", pool.cached_bytes()); diff --git a/crates/mohu-buffer/examples/buffer_basics.rs b/crates/mohu-buffer/examples/buffer_basics.rs index 131107c..c13b075 100644 --- a/crates/mohu-buffer/examples/buffer_basics.rs +++ b/crates/mohu-buffer/examples/buffer_basics.rs @@ -12,7 +12,7 @@ use mohu_buffer::{ Buffer, SliceArg, - strides::{c_strides, f_strides, NdIndexIter, unravel_index, ravel_multi_index}, + strides::{NdIndexIter, c_strides, f_strides, ravel_multi_index, unravel_index}, }; use mohu_dtype::DType; @@ -20,7 +20,12 @@ fn main() { // ── Buffer creation from a Rust slice ────────────────────────────────── // NumPy: a = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32) let a = Buffer::from_slice::(&[1.0, 2.0, 3.0, 4.0, 5.0]).unwrap(); - println!("from_slice: dtype={}, shape={:?}, len={}", a.dtype(), a.shape(), a.len()); + println!( + "from_slice: dtype={}, shape={:?}, len={}", + a.dtype(), + a.shape(), + a.len() + ); // Read elements back let v: f32 = a.get(&[2]).unwrap(); @@ -30,7 +35,11 @@ fn main() { // NumPy: a = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float64) let rows: &[&[f64]] = &[&[1.0, 2.0, 3.0], &[4.0, 5.0, 6.0]]; let a2d = Buffer::from_slice_2d(rows).unwrap(); - println!("\nfrom_slice_2d: shape={:?}, strides={:?}", a2d.shape(), a2d.strides()); + println!( + "\nfrom_slice_2d: shape={:?}, strides={:?}", + a2d.shape(), + a2d.strides() + ); let v: f64 = a2d.get(&[1, 2]).unwrap(); println!(" a[1,2] = {v}"); // 6.0 @@ -79,7 +88,11 @@ fn main() { // ── Transpose (zero-copy view operation) ─────────────────────────────── // NumPy: a.T let t = a2d.transpose(); - println!("\nTranspose: shape={:?}, strides={:?}", t.shape(), t.strides()); + println!( + "\nTranspose: shape={:?}, strides={:?}", + t.shape(), + t.strides() + ); let v: f64 = t.get(&[2, 1]).unwrap(); // was a2d[1,2] = 6.0 println!(" a.T[2,1] = {v}"); // 6.0 @@ -93,14 +106,25 @@ fn main() { // ── Slicing (zero-copy view with adjusted strides) ───────────────────── // NumPy: a[::2] (every other element) let five = Buffer::from_slice::(&[10.0, 20.0, 30.0, 40.0, 50.0]).unwrap(); - let sliced = five.slice_axis(0, SliceArg { start: Some(0), stop: None, step: Some(2) }).unwrap(); + let sliced = five + .slice_axis( + 0, + SliceArg { + start: Some(0), + stop: None, + step: Some(2), + }, + ) + .unwrap(); let data = sliced.to_vec::().unwrap(); println!("\nslice [::2]: {:?}", data); // [10.0, 30.0, 50.0] // ── Broadcast (zero-copy virtual replication) ────────────────────────── // NumPy: np.broadcast_to(np.array([1, 2, 3]), (3, 3)) - let row = Buffer::from_slice::(&[1.0, 2.0, 3.0]).unwrap() - .expand_dims(0).unwrap(); + let row = Buffer::from_slice::(&[1.0, 2.0, 3.0]) + .unwrap() + .expand_dims(0) + .unwrap(); let bc = row.broadcast_to(&[3, 3]).unwrap(); println!("\nBroadcast to (3,3): shape={:?}", bc.shape()); for i in 0..3 { @@ -115,8 +139,8 @@ fn main() { println!("\n── Layout inspection ──"); println!("a2d C-contiguous: {}", a2d.is_c_contiguous()); println!("a2d F-contiguous: {}", a2d.is_f_contiguous()); - println!("a2d.T C-contiguous: {}", t.is_c_contiguous()); // false after transpose - println!("a2d.T F-contiguous: {}", t.is_f_contiguous()); // true + println!("a2d.T C-contiguous: {}", t.is_c_contiguous()); // false after transpose + println!("a2d.T F-contiguous: {}", t.is_f_contiguous()); // true // ── Stride arithmetic ────────────────────────────────────────────────── println!("\n── Stride arithmetic ──"); diff --git a/crates/mohu-buffer/src/alloc.rs b/crates/mohu-buffer/src/alloc.rs index db56d32..108eef8 100644 --- a/crates/mohu-buffer/src/alloc.rs +++ b/crates/mohu-buffer/src/alloc.rs @@ -44,10 +44,10 @@ pub const POISON_BYTE: u8 = 0xDE; // ─── Global statistics ──────────────────────────────────────────────────────── -static LIVE_BYTES: AtomicI64 = AtomicI64::new(0); -static PEAK_BYTES: AtomicU64 = AtomicU64::new(0); -static ALLOC_COUNT: AtomicU64 = AtomicU64::new(0); -static FREE_COUNT: AtomicU64 = AtomicU64::new(0); +static LIVE_BYTES: AtomicI64 = AtomicI64::new(0); +static PEAK_BYTES: AtomicU64 = AtomicU64::new(0); +static ALLOC_COUNT: AtomicU64 = AtomicU64::new(0); +static FREE_COUNT: AtomicU64 = AtomicU64::new(0); static HEAP_LIVE_BYTES: AtomicI64 = AtomicI64::new(0); static MMAP_LIVE_BYTES: AtomicI64 = AtomicI64::new(0); @@ -55,13 +55,13 @@ static MMAP_LIVE_BYTES: AtomicI64 = AtomicI64::new(0); #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct AllocStats { /// Net bytes currently held by live mohu buffers. - pub live_bytes: i64, + pub live_bytes: i64, /// Maximum `live_bytes` observed since process start. - pub peak_bytes: u64, + pub peak_bytes: u64, /// Total number of successful allocations performed. - pub alloc_count: u64, + pub alloc_count: u64, /// Total number of frees performed. - pub free_count: u64, + pub free_count: u64, /// Bytes currently held by heap (`std::alloc`) allocations. pub heap_live_bytes: i64, /// Bytes currently held by mmap allocations. @@ -72,10 +72,10 @@ impl AllocStats { /// Takes an atomic snapshot of all global counters. pub fn snapshot() -> Self { Self { - live_bytes: LIVE_BYTES.load(Ordering::Relaxed), - peak_bytes: PEAK_BYTES.load(Ordering::Relaxed), - alloc_count: ALLOC_COUNT.load(Ordering::Relaxed), - free_count: FREE_COUNT.load(Ordering::Relaxed), + live_bytes: LIVE_BYTES.load(Ordering::Relaxed), + peak_bytes: PEAK_BYTES.load(Ordering::Relaxed), + alloc_count: ALLOC_COUNT.load(Ordering::Relaxed), + free_count: FREE_COUNT.load(Ordering::Relaxed), heap_live_bytes: HEAP_LIVE_BYTES.load(Ordering::Relaxed), mmap_live_bytes: MMAP_LIVE_BYTES.load(Ordering::Relaxed), } @@ -92,17 +92,24 @@ fn record_alloc(bytes: usize, strategy: Strategy) { let mut peak = PEAK_BYTES.load(Ordering::Relaxed); while (new as u64) > peak { match PEAK_BYTES.compare_exchange_weak( - peak, new as u64, Ordering::Relaxed, Ordering::Relaxed, + peak, + new as u64, + Ordering::Relaxed, + Ordering::Relaxed, ) { - Ok(_) => break, + Ok(_) => break, Err(current) => peak = current, } } ALLOC_COUNT.fetch_add(1, Ordering::Relaxed); match strategy { - Strategy::Heap => { HEAP_LIVE_BYTES.fetch_add(bytes as i64, Ordering::Relaxed); } - Strategy::Mmap => { MMAP_LIVE_BYTES.fetch_add(bytes as i64, Ordering::Relaxed); } - Strategy::ZeroSize => {} + Strategy::Heap => { + HEAP_LIVE_BYTES.fetch_add(bytes as i64, Ordering::Relaxed); + }, + Strategy::Mmap => { + MMAP_LIVE_BYTES.fetch_add(bytes as i64, Ordering::Relaxed); + }, + Strategy::ZeroSize => {}, } } @@ -110,9 +117,13 @@ fn record_free(bytes: usize, strategy: Strategy) { LIVE_BYTES.fetch_sub(bytes as i64, Ordering::Relaxed); FREE_COUNT.fetch_add(1, Ordering::Relaxed); match strategy { - Strategy::Heap => { HEAP_LIVE_BYTES.fetch_sub(bytes as i64, Ordering::Relaxed); } - Strategy::Mmap => { MMAP_LIVE_BYTES.fetch_sub(bytes as i64, Ordering::Relaxed); } - Strategy::ZeroSize => {} + Strategy::Heap => { + HEAP_LIVE_BYTES.fetch_sub(bytes as i64, Ordering::Relaxed); + }, + Strategy::Mmap => { + MMAP_LIVE_BYTES.fetch_sub(bytes as i64, Ordering::Relaxed); + }, + Strategy::ZeroSize => {}, } } @@ -164,27 +175,31 @@ pub enum MmapAdvice { impl MmapAdvice { pub(crate) fn to_libc(self) -> libc::c_int { match self { - MmapAdvice::Normal => libc::MADV_NORMAL, + MmapAdvice::Normal => libc::MADV_NORMAL, MmapAdvice::Sequential => libc::MADV_SEQUENTIAL, - MmapAdvice::Random => libc::MADV_RANDOM, - MmapAdvice::WillNeed => libc::MADV_WILLNEED, - MmapAdvice::DontNeed => libc::MADV_DONTNEED, + MmapAdvice::Random => libc::MADV_RANDOM, + MmapAdvice::WillNeed => libc::MADV_WILLNEED, + MmapAdvice::DontNeed => libc::MADV_DONTNEED, #[cfg(target_os = "linux")] - MmapAdvice::HugePage => libc::MADV_HUGEPAGE, + MmapAdvice::HugePage => libc::MADV_HUGEPAGE, #[cfg(target_os = "linux")] MmapAdvice::NoHugePage => libc::MADV_NOHUGEPAGE, #[cfg(not(target_os = "linux"))] - MmapAdvice::HugePage => libc::MADV_NORMAL, + MmapAdvice::HugePage => libc::MADV_NORMAL, #[cfg(not(target_os = "linux"))] MmapAdvice::NoHugePage => libc::MADV_NORMAL, - MmapAdvice::Free => { + MmapAdvice::Free => { // MADV_FREE: Linux >= 4.5 uses 8, macOS uses 5. // Both libc crates define MADV_FREE if available. #[cfg(any(target_os = "linux", target_os = "macos"))] - { libc::MADV_FREE } + { + libc::MADV_FREE + } #[cfg(not(any(target_os = "linux", target_os = "macos")))] - { libc::MADV_NORMAL } - } + { + libc::MADV_NORMAL + } + }, } } } @@ -193,7 +208,10 @@ impl MmapAdvice { enum AllocInner { ZeroSize, - Heap { ptr: NonNull, layout: StdLayout }, + Heap { + ptr: NonNull, + layout: StdLayout, + }, #[cfg(feature = "mmap")] Mmap(Box), } @@ -216,9 +234,9 @@ unsafe impl Sync for AllocInner {} /// - Alignment is at least [`SIMD_ALIGN`] bytes. /// - `Drop` frees the memory exactly once. pub struct AllocHandle { - inner: AllocInner, - len: usize, - align: usize, + inner: AllocInner, + len: usize, + align: usize, #[cfg(unix)] mlocked: bool, } @@ -233,8 +251,8 @@ impl AllocHandle { let align = align.max(SIMD_ALIGN).next_power_of_two(); if len == 0 { return Ok(Self { - inner: AllocInner::ZeroSize, - len: 0, + inner: AllocInner::ZeroSize, + len: 0, align, #[cfg(unix)] mlocked: false, @@ -243,8 +261,7 @@ impl AllocHandle { #[cfg(feature = "mmap")] let inner = if len >= MMAP_THRESHOLD { - let mmap = memmap2::MmapMut::map_anon(len) - .map_err(|_| MohuError::alloc(len))?; + let mmap = memmap2::MmapMut::map_anon(len).map_err(|_| MohuError::alloc(len))?; AllocInner::Mmap(Box::new(mmap)) } else { heap_alloc_inner(len, align)? @@ -275,10 +292,10 @@ impl AllocHandle { match &handle.inner { AllocInner::Heap { ptr, .. } => { unsafe { ptr.as_ptr().write_bytes(0, len) }; - } + }, #[cfg(feature = "mmap")] - AllocInner::Mmap(_) => { /* mmap pages arrive zero-filled from the OS */ } - AllocInner::ZeroSize => {} + AllocInner::Mmap(_) => { /* mmap pages arrive zero-filled from the OS */ }, + AllocInner::ZeroSize => {}, } } Ok(handle) @@ -290,10 +307,10 @@ impl AllocHandle { #[inline] pub fn as_ptr(&self) -> *const u8 { match &self.inner { - AllocInner::ZeroSize => NonNull::dangling().as_ptr(), - AllocInner::Heap { ptr, .. } => ptr.as_ptr(), + AllocInner::ZeroSize => NonNull::dangling().as_ptr(), + AllocInner::Heap { ptr, .. } => ptr.as_ptr(), #[cfg(feature = "mmap")] - AllocInner::Mmap(mmap) => mmap.as_ptr(), + AllocInner::Mmap(mmap) => mmap.as_ptr(), } } @@ -305,30 +322,39 @@ impl AllocHandle { #[inline] pub fn as_mut_ptr(&mut self) -> *mut u8 { match &mut self.inner { - AllocInner::ZeroSize => NonNull::dangling().as_ptr(), - AllocInner::Heap { ptr, .. } => ptr.as_ptr(), + AllocInner::ZeroSize => NonNull::dangling().as_ptr(), + AllocInner::Heap { ptr, .. } => ptr.as_ptr(), #[cfg(feature = "mmap")] - AllocInner::Mmap(mmap) => mmap.as_mut_ptr(), + AllocInner::Mmap(mmap) => mmap.as_mut_ptr(), } } // ─── Metadata ───────────────────────────────────────────────────────────── /// Byte length of this allocation (0 for zero-size handles). - #[inline] pub fn len(&self) -> usize { self.len } + #[inline] + pub fn len(&self) -> usize { + self.len + } /// Returns `true` if this is a zero-size allocation. - #[inline] pub fn is_empty(&self) -> bool { self.len == 0 } + #[inline] + pub fn is_empty(&self) -> bool { + self.len == 0 + } /// Minimum alignment of the allocation in bytes. - #[inline] pub fn align(&self) -> usize { self.align } + #[inline] + pub fn align(&self) -> usize { + self.align + } /// Returns the allocation strategy used for this handle. #[inline] pub fn strategy(&self) -> Strategy { match &self.inner { - AllocInner::ZeroSize => Strategy::ZeroSize, - AllocInner::Heap { .. } => Strategy::Heap, + AllocInner::ZeroSize => Strategy::ZeroSize, + AllocInner::Heap { .. } => Strategy::Heap, #[cfg(feature = "mmap")] - AllocInner::Mmap(_) => Strategy::Mmap, + AllocInner::Mmap(_) => Strategy::Mmap, } } @@ -341,7 +367,9 @@ impl AllocHandle { /// Returns the start pointer as a `NonNull`, or an error for zero-size. pub fn as_non_null(&self) -> MohuResult> { if self.len == 0 { - return Err(MohuError::bug("as_non_null called on a zero-size AllocHandle")); + return Err(MohuError::bug( + "as_non_null called on a zero-size AllocHandle", + )); } Ok(unsafe { NonNull::new_unchecked(self.as_ptr() as *mut u8) }) } @@ -400,16 +428,20 @@ impl AllocHandle { /// be in cache (e.g., after acquiring a cold buffer from the pool). #[inline] pub fn prefetch_read(&self) { - if self.len == 0 { return; } + if self.len == 0 { + return; + } let ptr = self.as_ptr(); let len = self.len; #[cfg(target_arch = "x86_64")] { - use std::arch::x86_64::{_mm_prefetch, _MM_HINT_T0}; + use std::arch::x86_64::{_MM_HINT_T0, _mm_prefetch}; let mut offset = 0usize; while offset < len { - unsafe { _mm_prefetch(ptr.add(offset) as *const i8, _MM_HINT_T0); } + unsafe { + _mm_prefetch(ptr.add(offset) as *const i8, _MM_HINT_T0); + } offset += CACHE_LINE; } } @@ -442,16 +474,20 @@ impl AllocHandle { /// Reduces RFO (read-for-ownership) stalls in write-heavy loops. #[inline] pub fn prefetch_write(&self) { - if self.len == 0 { return; } + if self.len == 0 { + return; + } let ptr = self.as_ptr(); let len = self.len; #[cfg(target_arch = "x86_64")] { - use std::arch::x86_64::{_mm_prefetch, _MM_HINT_ET0}; + use std::arch::x86_64::{_MM_HINT_ET0, _mm_prefetch}; let mut offset = 0usize; while offset < len { - unsafe { _mm_prefetch(ptr.add(offset) as *const i8, _MM_HINT_ET0); } + unsafe { + _mm_prefetch(ptr.add(offset) as *const i8, _MM_HINT_ET0); + } offset += CACHE_LINE; } } @@ -488,10 +524,10 @@ impl AllocHandle { pub fn mlock(&mut self) -> MohuResult<()> { #[cfg(unix)] { - if self.len == 0 || self.mlocked { return Ok(()); } - let ret = unsafe { - libc::mlock(self.as_ptr() as *const libc::c_void, self.len) - }; + if self.len == 0 || self.mlocked { + return Ok(()); + } + let ret = unsafe { libc::mlock(self.as_ptr() as *const libc::c_void, self.len) }; if ret != 0 { return Err(MohuError::bug(format!( "mlock failed: {}", @@ -522,8 +558,14 @@ impl AllocHandle { /// Returns `true` if this allocation's pages are currently locked in RAM. #[inline] pub fn is_mlocked(&self) -> bool { - #[cfg(unix)] { self.mlocked } - #[cfg(not(unix))] { false } + #[cfg(unix)] + { + self.mlocked + } + #[cfg(not(unix))] + { + false + } } // ─── In-place grow (Linux mremap) ───────────────────────────────────────── @@ -565,7 +607,9 @@ impl AllocHandle { /// For mmap regions, this is a hint to the OS via `MADV_DONTNEED` (Linux) /// or an explicit `memset` on other platforms — the result is zeroed bytes. pub fn zero(&mut self) { - if self.len == 0 { return; } + if self.len == 0 { + return; + } #[cfg(target_os = "linux")] if matches!(self.inner, AllocInner::Mmap(_)) { @@ -576,7 +620,9 @@ impl AllocHandle { } // Heap (or non-Linux mmap): explicit memset - unsafe { self.as_mut_ptr().write_bytes(0, self.len); } + unsafe { + self.as_mut_ptr().write_bytes(0, self.len); + } } // ─── Debug poison ───────────────────────────────────────────────────────── @@ -592,7 +638,9 @@ impl AllocHandle { pub fn poison(&mut self) { #[cfg(debug_assertions)] if self.len > 0 { - unsafe { self.as_mut_ptr().write_bytes(POISON_BYTE, self.len); } + unsafe { + self.as_mut_ptr().write_bytes(POISON_BYTE, self.len); + } } } @@ -601,7 +649,9 @@ impl AllocHandle { /// Useful in tests to assert that a buffer was properly poisoned before /// being placed in the pool, confirming no aliasing occurred. pub fn check_poison(&self) -> bool { - if self.len == 0 { return true; } + if self.len == 0 { + return true; + } let slice = unsafe { std::slice::from_raw_parts(self.as_ptr(), self.len) }; slice.iter().all(|&b| b == POISON_BYTE) } @@ -610,9 +660,11 @@ impl AllocHandle { /// /// Use to detect use-after-free in debug scenarios. pub fn has_poison(&self) -> bool { - if self.len == 0 { return false; } + if self.len == 0 { + return false; + } let slice = unsafe { std::slice::from_raw_parts(self.as_ptr(), self.len) }; - slice.iter().any(|&b| b == POISON_BYTE) + slice.contains(&POISON_BYTE) } // ─── Byte-range view ────────────────────────────────────────────────────── @@ -650,10 +702,10 @@ impl Drop for AllocHandle { match &self.inner { AllocInner::Heap { ptr, layout } => { unsafe { alloc::dealloc(ptr.as_ptr(), *layout) }; - } + }, #[cfg(feature = "mmap")] - AllocInner::Mmap(_) => { /* Box calls munmap on drop */ } - AllocInner::ZeroSize => {} + AllocInner::Mmap(_) => { /* Box calls munmap on drop */ }, + AllocInner::ZeroSize => {}, } } } @@ -661,9 +713,9 @@ impl Drop for AllocHandle { impl std::fmt::Debug for AllocHandle { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("AllocHandle") - .field("ptr", &format_args!("{:p}", self.as_ptr())) - .field("len", &self.len) - .field("align", &self.align) + .field("ptr", &format_args!("{:p}", self.as_ptr())) + .field("len", &self.len) + .field("align", &self.align) .field("strategy", &self.strategy()) .finish() } @@ -672,8 +724,7 @@ impl std::fmt::Debug for AllocHandle { // ─── Internal helpers ───────────────────────────────────────────────────────── fn heap_alloc_inner(len: usize, align: usize) -> MohuResult { - let layout = StdLayout::from_size_align(len, align) - .map_err(|_| MohuError::alloc(len))?; + let layout = StdLayout::from_size_align(len, align).map_err(|_| MohuError::alloc(len))?; let raw = unsafe { alloc::alloc(layout) }; let ptr = NonNull::new(raw).ok_or_else(|| MohuError::alloc(len))?; Ok(AllocInner::Heap { ptr, layout }) diff --git a/crates/mohu-buffer/src/buffer.rs b/crates/mohu-buffer/src/buffer.rs index 3b6bed6..873dc6e 100644 --- a/crates/mohu-buffer/src/buffer.rs +++ b/crates/mohu-buffer/src/buffer.rs @@ -11,13 +11,10 @@ /// `Buffer` can wrap externally-owned memory via [`Buffer::from_dlpack`]. /// In that case the backing `RawBuffer` holds a `*mut DLManagedTensor` /// and calls its deleter when the last Arc reference is dropped. -use std::{ - ptr::NonNull, - sync::Arc, -}; +use std::{ptr::NonNull, sync::Arc}; use mohu_dtype::{ - dlpack::{assert_cpu_device, DLDataType}, + dlpack::{DLDataType, assert_cpu_device}, dtype::DType, promote::CastMode, scalar::Scalar, @@ -38,24 +35,32 @@ pub struct BufferFlags(u8); impl BufferFlags { /// Array is writeable (not read-only / not a broadcast view). - pub const WRITEABLE: Self = Self(1 << 0); + pub const WRITEABLE: Self = Self(1 << 0); /// This `Buffer` is the (sole or shared) owner of the backing bytes. - pub const OWNS_DATA: Self = Self(1 << 1); + pub const OWNS_DATA: Self = Self(1 << 1); /// Array is C-contiguous in the backing buffer. - pub const C_CONTIGUOUS: Self = Self(1 << 2); + pub const C_CONTIGUOUS: Self = Self(1 << 2); /// Array is Fortran-contiguous in the backing buffer. - pub const F_CONTIGUOUS: Self = Self(1 << 3); + pub const F_CONTIGUOUS: Self = Self(1 << 3); /// Backing memory is SIMD-aligned (≥ 64 bytes). - pub const ALIGNED: Self = Self(1 << 4); + pub const ALIGNED: Self = Self(1 << 4); /// Returns an empty flag set with no flags enabled. - pub const fn empty() -> Self { Self(0) } + pub const fn empty() -> Self { + Self(0) + } /// Returns `true` if all flags in `other` are set in `self`. - pub const fn contains(self, other: Self) -> bool { self.0 & other.0 == other.0 } + pub const fn contains(self, other: Self) -> bool { + self.0 & other.0 == other.0 + } /// Returns a new flag set with all flags from both `self` and `other`. - pub const fn insert(self, other: Self) -> Self { Self(self.0 | other.0) } + pub const fn insert(self, other: Self) -> Self { + Self(self.0 | other.0) + } /// Returns a new flag set with `other`'s flags cleared from `self`. - pub const fn remove(self, other: Self) -> Self { Self(self.0 & !other.0) } + pub const fn remove(self, other: Self) -> Self { + Self(self.0 & !other.0) + } } // ─── DLPack C-ABI types ─────────────────────────────────────────────────────── @@ -64,8 +69,8 @@ impl BufferFlags { #[repr(C)] #[derive(Debug, Clone, Copy)] pub struct RawDLDataType { - pub code: u8, - pub bits: u8, + pub code: u8, + pub bits: u8, pub lanes: u16, } @@ -81,40 +86,42 @@ impl From for RawDLDataType { #[derive(Debug, Clone, Copy)] pub struct RawDLDevice { pub device_type: i32, - pub device_id: i32, + pub device_id: i32, } /// C-ABI DLTensor (DLPack v0.8 layout). #[repr(C)] pub struct DLTensor { - pub data: *mut std::ffi::c_void, - pub device: RawDLDevice, - pub ndim: i32, - pub dtype: RawDLDataType, - pub shape: *const i64, - pub strides: *const i64, + pub data: *mut std::ffi::c_void, + pub device: RawDLDevice, + pub ndim: i32, + pub dtype: RawDLDataType, + pub shape: *const i64, + pub strides: *const i64, pub byte_offset: u64, } /// Context kept alive for the lifetime of an exported `DLManagedTensor`. struct DLExportCtx { /// Keeps the backing buffer alive until the DLPack consumer is done. - _raw: Arc, - shape: Vec, + _raw: Arc, + shape: Vec, strides: Vec, } /// C-ABI DLManagedTensor (DLPack v0.8). #[repr(C)] pub struct DLManagedTensor { - pub dl_tensor: DLTensor, + pub dl_tensor: DLTensor, pub manager_ctx: *mut std::ffi::c_void, - pub deleter: Option, + pub deleter: Option, } /// Called by the DLPack consumer when it no longer needs the tensor. unsafe extern "C" fn dlmanaged_deleter(ptr: *mut DLManagedTensor) { - if ptr.is_null() { return; } + if ptr.is_null() { + return; + } unsafe { let managed = &*ptr; if !managed.manager_ctx.is_null() { @@ -133,9 +140,7 @@ enum BufferSource { Owned(AllocHandle), /// Externally owned memory imported via DLPack. /// The deleter is invoked when this `RawBuffer` is dropped. - DLPack { - managed: *mut DLManagedTensor, - }, + DLPack { managed: *mut DLManagedTensor }, } // SAFETY: *mut DLManagedTensor is owned exclusively by this RawBuffer. @@ -163,7 +168,7 @@ pub struct RawBuffer { source: BufferSource, /// Pointer to the usable start of the data (may be offset into the /// DLPack allocation). - ptr: NonNull, + ptr: NonNull, /// Number of usable bytes. nbytes: usize, } @@ -185,7 +190,11 @@ impl RawBuffer { } else { handle.as_non_null()? }; - Ok(Self { source: BufferSource::Owned(handle), ptr, nbytes }) + Ok(Self { + source: BufferSource::Owned(handle), + ptr, + nbytes, + }) } /// Wraps an externally owned DLPack pointer. @@ -197,8 +206,8 @@ impl RawBuffer { /// in `managed` is invoked. unsafe fn from_dlpack_ptr( managed: *mut DLManagedTensor, - ptr: NonNull, - nbytes: usize, + ptr: NonNull, + nbytes: usize, ) -> Self { Self { source: BufferSource::DLPack { managed }, @@ -208,11 +217,20 @@ impl RawBuffer { } /// Returns the data pointer. - #[inline] pub fn as_ptr(&self) -> *const u8 { self.ptr.as_ptr() } + #[inline] + pub fn as_ptr(&self) -> *const u8 { + self.ptr.as_ptr() + } /// Returns a mutable data pointer (caller must ensure exclusive access). - #[inline] pub fn as_mut_ptr(&self) -> *mut u8 { self.ptr.as_ptr() } + #[inline] + pub fn as_mut_ptr(&self) -> *mut u8 { + self.ptr.as_ptr() + } /// Returns the byte capacity of this raw buffer. - #[inline] pub fn nbytes(&self) -> usize { self.nbytes } + #[inline] + pub fn nbytes(&self) -> usize { + self.nbytes + } /// Returns `true` if the backing memory is SIMD-aligned. #[inline] @@ -229,7 +247,7 @@ impl RawBuffer { impl std::fmt::Debug for RawBuffer { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("RawBuffer") - .field("ptr", &format_args!("{:p}", self.ptr)) + .field("ptr", &format_args!("{:p}", self.ptr)) .field("nbytes", &self.nbytes) .field("external", &self.is_external()) .finish() @@ -250,10 +268,10 @@ impl std::fmt::Debug for RawBuffer { /// To get an independent copy of the data, call [`make_unique`](Buffer::make_unique). #[derive(Debug)] pub struct Buffer { - raw: Arc, - dtype: DType, + raw: Arc, + dtype: DType, layout: Layout, - flags: BufferFlags, + flags: BufferFlags, } impl Buffer { @@ -268,7 +286,12 @@ impl Buffer { Order::F => Layout::new_f(shape, dtype.itemsize())?, }; let flags = Self::compute_flags(&raw, &layout); - Ok(Self { raw, dtype, layout, flags }) + Ok(Self { + raw, + dtype, + layout, + flags, + }) } /// Allocates a zeroed buffer with `dtype` and `shape`. @@ -277,7 +300,12 @@ impl Buffer { let raw = Arc::new(RawBuffer::alloc(nbytes, true)?); let layout = Layout::new_c(shape, dtype.itemsize())?; let flags = Self::compute_flags(&raw, &layout); - Ok(Self { raw, dtype, layout, flags }) + Ok(Self { + raw, + dtype, + layout, + flags, + }) } /// Allocates a buffer filled with the one-value for `dtype`. @@ -292,15 +320,12 @@ impl Buffer { /// Allocates a buffer filled with `fill_bytes` repeated for each element. /// /// `fill_bytes.len()` must equal `dtype.itemsize()`. - pub fn full( - dtype: DType, - shape: &[usize], - fill_bytes: &[u8], - ) -> MohuResult { + pub fn full(dtype: DType, shape: &[usize], fill_bytes: &[u8]) -> MohuResult { if fill_bytes.len() != dtype.itemsize() { return Err(MohuError::bug(format!( "Buffer::full: fill_bytes.len()={} != dtype.itemsize()={}", - fill_bytes.len(), dtype.itemsize() + fill_bytes.len(), + dtype.itemsize() ))); } let mut buf = Self::alloc(dtype, shape, Order::C)?; @@ -312,22 +337,22 @@ impl Buffer { /// Copies elements from a typed slice into a new C-contiguous buffer. pub fn from_slice(data: &[T]) -> MohuResult { - let dtype = T::DTYPE; - let shape = [data.len()]; + let dtype = T::DTYPE; + let shape = [data.len()]; let nbytes = data.len() * dtype.itemsize(); - let raw = Arc::new(RawBuffer::alloc(nbytes, false)?); + let raw = Arc::new(RawBuffer::alloc(nbytes, false)?); // SAFETY: raw has exactly nbytes of valid writable memory. unsafe { - std::ptr::copy_nonoverlapping( - data.as_ptr() as *const u8, - raw.as_mut_ptr(), - nbytes, - ); + std::ptr::copy_nonoverlapping(data.as_ptr() as *const u8, raw.as_mut_ptr(), nbytes); } let layout = Layout::new_c(&shape, dtype.itemsize())?; - let flags = Self::compute_flags(&raw, &layout) - .insert(BufferFlags::WRITEABLE); - Ok(Self { raw, dtype, layout, flags }) + let flags = Self::compute_flags(&raw, &layout).insert(BufferFlags::WRITEABLE); + Ok(Self { + raw, + dtype, + layout, + flags, + }) } /// Copies a 2D slice-of-slices into a row-major buffer. @@ -336,35 +361,35 @@ impl Buffer { return Self::zeros(T::DTYPE, &[0, 0]); } let cols = data[0].len(); - for (_row_idx, row) in data.iter().enumerate() { + for row in data.iter() { if row.len() != cols { return Err(MohuError::ShapeMismatch { expected: vec![cols], - got: vec![row.len()], + got: vec![row.len()], }); } } - let rows = data.len(); - let dtype = T::DTYPE; - let shape = [rows, cols]; + let rows = data.len(); + let dtype = T::DTYPE; + let shape = [rows, cols]; let nbytes = rows * cols * dtype.itemsize(); - let raw = Arc::new(RawBuffer::alloc(nbytes, false)?); + let raw = Arc::new(RawBuffer::alloc(nbytes, false)?); unsafe { let mut dst = raw.as_mut_ptr(); for row in data { let row_bytes = row.len() * dtype.itemsize(); - std::ptr::copy_nonoverlapping( - row.as_ptr() as *const u8, - dst, - row_bytes, - ); + std::ptr::copy_nonoverlapping(row.as_ptr() as *const u8, dst, row_bytes); dst = dst.add(row_bytes); } } let layout = Layout::new_c(&shape, dtype.itemsize())?; - let flags = Self::compute_flags(&raw, &layout) - .insert(BufferFlags::WRITEABLE); - Ok(Self { raw, dtype, layout, flags }) + let flags = Self::compute_flags(&raw, &layout).insert(BufferFlags::WRITEABLE); + Ok(Self { + raw, + dtype, + layout, + flags, + }) } /// Wraps a `Vec` by copying it into a mohu buffer. @@ -383,23 +408,27 @@ impl Buffer { /// /// Prefer `from_slice` or DLPack import for safer alternatives. pub unsafe fn from_raw_parts( - ptr: NonNull, + ptr: NonNull, nbytes: usize, - dtype: DType, + dtype: DType, layout: Layout, ) -> Self { // We create a fake AllocHandle that owns nothing — zero-size handle. // Caller is responsible for the actual lifetime. let raw = Arc::new(RawBuffer { source: BufferSource::Owned( - AllocHandle::alloc(0, SIMD_ALIGN) - .expect("zero-size alloc never fails") + AllocHandle::alloc(0, SIMD_ALIGN).expect("zero-size alloc never fails"), ), ptr, nbytes, }); let flags = Self::compute_flags(&raw, &layout); - Self { raw, dtype, layout, flags } + Self { + raw, + dtype, + layout, + flags, + } } // ─── DLPack import ──────────────────────────────────────────────────────── @@ -417,27 +446,37 @@ impl Buffer { if managed.is_null() { return Err(MohuError::DLPackNullPointer); } - let (tensor_device, tensor_dtype, tensor_ndim, tensor_data, - tensor_byte_offset, tensor_shape, tensor_strides) = unsafe { + let ( + tensor_device, + tensor_dtype, + tensor_ndim, + tensor_data, + tensor_byte_offset, + tensor_shape, + tensor_strides, + ) = unsafe { let m = &*managed; let t = &m.dl_tensor; - (t.device, t.dtype, t.ndim, t.data, t.byte_offset, t.shape, t.strides) + ( + t.device, + t.dtype, + t.ndim, + t.data, + t.byte_offset, + t.shape, + t.strides, + ) }; assert_cpu_device(tensor_device.device_type)?; - let dtype = DType::from_dlpack( - tensor_dtype.code, - tensor_dtype.bits, - tensor_dtype.lanes, - )?; + let dtype = DType::from_dlpack(tensor_dtype.code, tensor_dtype.bits, tensor_dtype.lanes)?; let ndim = tensor_ndim as usize; let byte_offset = tensor_byte_offset as usize; let base_ptr = unsafe { (tensor_data as *mut u8).add(byte_offset) }; - let ptr = NonNull::new(base_ptr).ok_or_else(|| { - MohuError::DLPackInvalid("DLTensor.data is null".to_string()) - })?; + let ptr = NonNull::new(base_ptr) + .ok_or_else(|| MohuError::DLPackInvalid("DLTensor.data is null".to_string()))?; let shape: Vec = if ndim == 0 { vec![] @@ -467,7 +506,12 @@ impl Buffer { let mut flags = Self::compute_flags(&raw, &layout); flags = flags.remove(BufferFlags::WRITEABLE); - Ok(Self { raw, dtype, layout, flags }) + Ok(Self { + raw, + dtype, + layout, + flags, + }) } // ─── DLPack export ──────────────────────────────────────────────────────── @@ -484,41 +528,46 @@ impl Buffer { /// Dropping the returned pointer without calling the deleter leaks memory. pub fn to_dlpack(&self) -> MohuResult<*mut DLManagedTensor> { let dl_dtype = RawDLDataType::from(self.dtype.to_dlpack()); - let ndim = self.layout.ndim(); + let ndim = self.layout.ndim(); - let shape: Vec = self.layout.shape().iter() - .map(|&d| d as i64).collect(); + let shape: Vec = self.layout.shape().iter().map(|&d| d as i64).collect(); // Convert byte strides to element strides. let itemsize = self.dtype.itemsize() as isize; - let strides: Vec = self.layout.strides().iter() - .map(|&s| (s / itemsize) as i64).collect(); + let strides: Vec = self + .layout + .strides() + .iter() + .map(|&s| (s / itemsize) as i64) + .collect(); let ctx = Box::new(DLExportCtx { - _raw: Arc::clone(&self.raw), + _raw: Arc::clone(&self.raw), shape, strides, }); let ctx_ptr = Box::into_raw(ctx); - let data_ptr = unsafe { - self.raw.as_mut_ptr().add(self.layout.offset()) as *mut std::ffi::c_void - }; + let data_ptr = + unsafe { self.raw.as_mut_ptr().add(self.layout.offset()) as *mut std::ffi::c_void }; let dl_tensor = DLTensor { - data: data_ptr, - device: RawDLDevice { device_type: 1, device_id: 0 }, - ndim: ndim as i32, - dtype: dl_dtype, - shape: unsafe { (*ctx_ptr).shape.as_ptr() }, - strides: unsafe { (*ctx_ptr).strides.as_ptr() }, + data: data_ptr, + device: RawDLDevice { + device_type: 1, + device_id: 0, + }, + ndim: ndim as i32, + dtype: dl_dtype, + shape: unsafe { (*ctx_ptr).shape.as_ptr() }, + strides: unsafe { (*ctx_ptr).strides.as_ptr() }, byte_offset: 0, }; let managed = Box::new(DLManagedTensor { dl_tensor, manager_ctx: ctx_ptr as *mut std::ffi::c_void, - deleter: Some(dlmanaged_deleter), + deleter: Some(dlmanaged_deleter), }); Ok(Box::into_raw(managed)) @@ -527,40 +576,84 @@ impl Buffer { // ─── Properties ─────────────────────────────────────────────────────────── /// Returns the element data type of this buffer. - #[inline] pub fn dtype(&self) -> DType { self.dtype } + #[inline] + pub fn dtype(&self) -> DType { + self.dtype + } /// Returns a reference to this buffer's layout descriptor. - #[inline] pub fn layout(&self) -> &Layout { &self.layout } + #[inline] + pub fn layout(&self) -> &Layout { + &self.layout + } /// Returns the shape of this buffer as a slice of dimension sizes. - #[inline] pub fn shape(&self) -> &[usize] { self.layout.shape() } + #[inline] + pub fn shape(&self) -> &[usize] { + self.layout.shape() + } /// Returns the byte strides of this buffer. - #[inline] pub fn strides(&self) -> &[isize] { self.layout.strides() } + #[inline] + pub fn strides(&self) -> &[isize] { + self.layout.strides() + } /// Returns the number of dimensions (axes) of this buffer. - #[inline] pub fn ndim(&self) -> usize { self.layout.ndim() } + #[inline] + pub fn ndim(&self) -> usize { + self.layout.ndim() + } /// Returns the total number of elements in this buffer. - #[inline] pub fn len(&self) -> usize { self.layout.size() } + #[inline] + pub fn len(&self) -> usize { + self.layout.size() + } /// Returns the total byte size of this buffer's data (`len * itemsize`). - #[inline] pub fn nbytes(&self) -> usize { self.layout.nbytes() } + #[inline] + pub fn nbytes(&self) -> usize { + self.layout.nbytes() + } /// Returns the byte size of a single element. - #[inline] pub fn itemsize(&self) -> usize { self.dtype.itemsize() } + #[inline] + pub fn itemsize(&self) -> usize { + self.dtype.itemsize() + } /// Returns the byte offset from the backing buffer start to element `[0, …, 0]`. - #[inline] pub fn offset(&self) -> usize { self.layout.offset() } + #[inline] + pub fn offset(&self) -> usize { + self.layout.offset() + } /// Returns the bitfield flags describing this buffer's properties. - #[inline] pub fn flags(&self) -> BufferFlags { self.flags } + #[inline] + pub fn flags(&self) -> BufferFlags { + self.flags + } /// Returns `true` if any dimension is zero (zero-element buffer). - pub fn is_empty(&self) -> bool { self.layout.is_empty() } + pub fn is_empty(&self) -> bool { + self.layout.is_empty() + } /// Returns `true` if this buffer is writeable. - pub fn is_writeable(&self) -> bool { self.flags.contains(BufferFlags::WRITEABLE) } + pub fn is_writeable(&self) -> bool { + self.flags.contains(BufferFlags::WRITEABLE) + } /// Returns `true` if this buffer is C-contiguous (row-major). - pub fn is_c_contiguous(&self) -> bool { self.layout.is_c_contiguous() } + pub fn is_c_contiguous(&self) -> bool { + self.layout.is_c_contiguous() + } /// Returns `true` if this buffer is Fortran-contiguous (column-major). - pub fn is_f_contiguous(&self) -> bool { self.layout.is_f_contiguous() } + pub fn is_f_contiguous(&self) -> bool { + self.layout.is_f_contiguous() + } /// Returns `true` if this buffer is contiguous in either C or F order. - pub fn is_contiguous(&self) -> bool { self.layout.is_contiguous() } + pub fn is_contiguous(&self) -> bool { + self.layout.is_contiguous() + } /// Returns `true` if the backing memory is SIMD-aligned. - pub fn is_aligned(&self) -> bool { self.flags.contains(BufferFlags::ALIGNED) } + pub fn is_aligned(&self) -> bool { + self.flags.contains(BufferFlags::ALIGNED) + } /// Returns `true` if the backing memory is shared with other `Buffer` instances. - pub fn is_shared(&self) -> bool { Arc::strong_count(&self.raw) > 1 } + pub fn is_shared(&self) -> bool { + Arc::strong_count(&self.raw) > 1 + } /// Returns a raw const pointer to element `[0, 0, …, 0]`. #[inline] @@ -588,7 +681,7 @@ impl Buffer { if T::DTYPE != self.dtype { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: self.dtype.to_string(), + got: self.dtype.to_string(), }); } if !self.is_c_contiguous() { @@ -606,7 +699,7 @@ impl Buffer { if T::DTYPE != self.dtype { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: self.dtype.to_string(), + got: self.dtype.to_string(), }); } if !self.is_writeable() { @@ -630,7 +723,7 @@ impl Buffer { if T::DTYPE != self.dtype { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: self.dtype.to_string(), + got: self.dtype.to_string(), }); } let off = self.layout.byte_offset(indices)?; @@ -644,7 +737,7 @@ impl Buffer { if T::DTYPE != self.dtype { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: self.dtype.to_string(), + got: self.dtype.to_string(), }); } if !self.is_writeable() { @@ -666,10 +759,10 @@ impl Buffer { /// To get an independent copy, call [`make_unique`](Self::make_unique). pub fn share(&self) -> Self { Self { - raw: Arc::clone(&self.raw), - dtype: self.dtype, + raw: Arc::clone(&self.raw), + dtype: self.dtype, layout: self.layout.clone(), - flags: self.flags.remove(BufferFlags::WRITEABLE), // shared = read-only + flags: self.flags.remove(BufferFlags::WRITEABLE), // shared = read-only } } @@ -713,11 +806,13 @@ impl Buffer { /// Returns a transposed view (reverses axis order, no copy). pub fn transpose(&self) -> Self { Self { - raw: Arc::clone(&self.raw), - dtype: self.dtype, + raw: Arc::clone(&self.raw), + dtype: self.dtype, layout: self.layout.transpose(), - flags: self.flags.remove(BufferFlags::WRITEABLE) - .remove(BufferFlags::C_CONTIGUOUS), + flags: self + .flags + .remove(BufferFlags::WRITEABLE) + .remove(BufferFlags::C_CONTIGUOUS), } } @@ -725,21 +820,23 @@ impl Buffer { pub fn permute(&self, axes: &[usize]) -> MohuResult { let layout = self.layout.permute(axes)?; Ok(Self { - raw: Arc::clone(&self.raw), - dtype: self.dtype, + raw: Arc::clone(&self.raw), + dtype: self.dtype, layout, - flags: self.flags.remove(BufferFlags::WRITEABLE) - .remove(BufferFlags::C_CONTIGUOUS), + flags: self + .flags + .remove(BufferFlags::WRITEABLE) + .remove(BufferFlags::C_CONTIGUOUS), }) } /// Returns a reshaped view. Requires C-contiguous layout. pub fn reshape(&self, new_shape: &[usize]) -> MohuResult { let layout = self.layout.reshape(new_shape)?; - let flags = Self::compute_flags(&self.raw, &layout); + let flags = Self::compute_flags(&self.raw, &layout); Ok(Self { - raw: Arc::clone(&self.raw), - dtype: self.dtype, + raw: Arc::clone(&self.raw), + dtype: self.dtype, layout, flags, }) @@ -748,21 +845,27 @@ impl Buffer { /// Returns a slice along `axis` with the given `SliceArg`. pub fn slice_axis(&self, axis: usize, arg: SliceArg) -> MohuResult { let layout = self.layout.slice_axis(axis, arg)?; - let flags = self.flags + let flags = self + .flags .remove(BufferFlags::WRITEABLE) .remove(BufferFlags::C_CONTIGUOUS) .remove(BufferFlags::F_CONTIGUOUS); - Ok(Self { raw: Arc::clone(&self.raw), dtype: self.dtype, layout, flags }) + Ok(Self { + raw: Arc::clone(&self.raw), + dtype: self.dtype, + layout, + flags, + }) } /// Returns a broadcast view to `new_shape`. Broadcast axes are read-only. pub fn broadcast_to(&self, new_shape: &[usize]) -> MohuResult { let layout = self.layout.broadcast_to(new_shape)?; Ok(Self { - raw: Arc::clone(&self.raw), - dtype: self.dtype, + raw: Arc::clone(&self.raw), + dtype: self.dtype, layout, - flags: self.flags.remove(BufferFlags::WRITEABLE), + flags: self.flags.remove(BufferFlags::WRITEABLE), }) } @@ -770,20 +873,20 @@ impl Buffer { pub fn expand_dims(&self, axis: usize) -> MohuResult { let layout = self.layout.expand_dims(axis)?; Ok(Self { - raw: Arc::clone(&self.raw), - dtype: self.dtype, + raw: Arc::clone(&self.raw), + dtype: self.dtype, layout, - flags: self.flags, + flags: self.flags, }) } /// Removes all axes of size 1. pub fn squeeze(&self) -> Self { Self { - raw: Arc::clone(&self.raw), - dtype: self.dtype, + raw: Arc::clone(&self.raw), + dtype: self.dtype, layout: self.layout.squeeze(), - flags: self.flags, + flags: self.flags, } } @@ -809,7 +912,7 @@ impl Buffer { if T::DTYPE != self.dtype { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: self.dtype.to_string(), + got: self.dtype.to_string(), }); } if self.is_c_contiguous() { @@ -835,9 +938,15 @@ impl Buffer { let mut f = BufferFlags::empty() .insert(BufferFlags::WRITEABLE) .insert(BufferFlags::OWNS_DATA); - if raw.is_aligned() { f = f.insert(BufferFlags::ALIGNED); } - if layout.is_c_contiguous() { f = f.insert(BufferFlags::C_CONTIGUOUS); } - if layout.is_f_contiguous() { f = f.insert(BufferFlags::F_CONTIGUOUS); } + if raw.is_aligned() { + f = f.insert(BufferFlags::ALIGNED); + } + if layout.is_c_contiguous() { + f = f.insert(BufferFlags::C_CONTIGUOUS); + } + if layout.is_f_contiguous() { + f = f.insert(BufferFlags::F_CONTIGUOUS); + } f } } @@ -847,10 +956,10 @@ impl Clone for Buffer { /// for a deep (independent) copy. fn clone(&self) -> Self { Self { - raw: Arc::clone(&self.raw), - dtype: self.dtype, + raw: Arc::clone(&self.raw), + dtype: self.dtype, layout: self.layout.clone(), - flags: self.flags.remove(BufferFlags::WRITEABLE), + flags: self.flags.remove(BufferFlags::WRITEABLE), } } } @@ -889,7 +998,7 @@ impl Buffer { let buf = Self::alloc(DType::F64, &[n], Order::C)?; if n > 0 { - let ptr = unsafe { buf.as_mut_ptr() as *mut f64 }; + let ptr = unsafe { buf.as_mut_ptr() as *mut f64 }; let slice = unsafe { std::slice::from_raw_parts_mut(ptr, n) }; // Index-based so every element is independent — safe for Rayon. use rayon::prelude::*; @@ -898,28 +1007,36 @@ impl Buffer { }); } - if dtype == DType::F64 { Ok(buf) } else { buf.cast(dtype, CastMode::Unsafe) } + if dtype == DType::F64 { + Ok(buf) + } else { + buf.cast(dtype, CastMode::Unsafe) + } } /// Creates a 1-D buffer of `n` evenly-spaced values from `start` to `stop`. /// /// Equivalent to `np.linspace`. If `endpoint` is `true`, `stop` is included. pub fn linspace( - start: f64, - stop: f64, - n: usize, + start: f64, + stop: f64, + n: usize, endpoint: bool, - dtype: DType, + dtype: DType, ) -> MohuResult { if n == 0 { return Self::alloc(dtype, &[0], Order::C); } - let div = if endpoint && n > 1 { (n - 1) as f64 } else { n as f64 }; + let div = if endpoint && n > 1 { + (n - 1) as f64 + } else { + n as f64 + }; let span = stop - start; let buf = Self::alloc(DType::F64, &[n], Order::C)?; { - let ptr = unsafe { buf.as_mut_ptr() as *mut f64 }; + let ptr = unsafe { buf.as_mut_ptr() as *mut f64 }; let slice = unsafe { std::slice::from_raw_parts_mut(ptr, n) }; use rayon::prelude::*; slice.par_iter_mut().enumerate().for_each(|(i, v)| { @@ -930,7 +1047,11 @@ impl Buffer { } } - if dtype == DType::F64 { Ok(buf) } else { buf.cast(dtype, CastMode::Unsafe) } + if dtype == DType::F64 { + Ok(buf) + } else { + buf.cast(dtype, CastMode::Unsafe) + } } /// Creates an `n × m` identity matrix with ones on diagonal `k`. @@ -939,11 +1060,13 @@ impl Buffer { /// Equivalent to `np.eye`. pub fn eye(n: usize, m: usize, k: i64, dtype: DType) -> MohuResult { let buf = Self::zeros(dtype, &[n, m])?; - if n == 0 || m == 0 { return Ok(buf); } + if n == 0 || m == 0 { + return Ok(buf); + } let one_bytes = dtype_one_bytes(dtype); - let itemsize = dtype.itemsize(); - let raw_ptr = unsafe { buf.as_mut_ptr() }; + let itemsize = dtype.itemsize(); + let raw_ptr = unsafe { buf.as_mut_ptr() }; let (row_start, col_start) = if k >= 0 { (0usize, k as usize) @@ -968,13 +1091,17 @@ impl Buffer { if v.ndim() != 1 { return Err(MohuError::bug("Buffer::diag: input must be 1-D")); } - let n = v.len(); - let size = if k >= 0 { n + k as usize } else { n + (-k) as usize }; + let n = v.len(); + let size = if k >= 0 { + n + k as usize + } else { + n + (-k) as usize + }; let out = Self::zeros(v.dtype(), &[size, size])?; let itemsize = v.dtype().itemsize(); - let src_raw = v.as_ptr(); - let dst_raw = unsafe { out.as_mut_ptr() }; + let src_raw = v.as_ptr(); + let dst_raw = unsafe { out.as_mut_ptr() }; let (row_start, col_start) = if k >= 0 { (0usize, k as usize) @@ -1000,9 +1127,9 @@ impl Buffer { if self.ndim() < 2 { return Err(MohuError::bug("diagonal: requires at least 2 dimensions")); } - let nd = self.ndim(); - let n = self.shape()[nd - 2]; - let m = self.shape()[nd - 1]; + let nd = self.ndim(); + let n = self.shape()[nd - 2]; + let m = self.shape()[nd - 1]; let (row_start, col_start) = if k >= 0 { (0usize, k as usize) @@ -1028,18 +1155,15 @@ impl Buffer { let mut new_strides: Vec = self.strides()[..nd - 2].to_vec(); new_strides.push(diag_stride); - let layout = Layout::new_custom( - &new_shape, - &new_strides, - start_off, - self.layout.itemsize(), - )?; + let layout = + Layout::new_custom(&new_shape, &new_strides, start_off, self.layout.itemsize())?; Ok(Self { - raw: Arc::clone(&self.raw), - dtype: self.dtype, + raw: Arc::clone(&self.raw), + dtype: self.dtype, layout, - flags: self.flags + flags: self + .flags .remove(BufferFlags::WRITEABLE) .remove(BufferFlags::C_CONTIGUOUS) .remove(BufferFlags::F_CONTIGUOUS), @@ -1057,7 +1181,7 @@ impl Buffer { self.ndim() ))); } - let dim = self.shape()[axis]; + let dim = self.shape()[axis]; let mut new_strides: Vec = self.strides().to_vec(); // New offset = old offset + (dim-1) * old_stride[axis] let offset_delta = if dim > 0 { @@ -1080,10 +1204,11 @@ impl Buffer { )?; Ok(Self { - raw: Arc::clone(&self.raw), - dtype: self.dtype, + raw: Arc::clone(&self.raw), + dtype: self.dtype, layout, - flags: self.flags + flags: self + .flags .remove(BufferFlags::WRITEABLE) .remove(BufferFlags::C_CONTIGUOUS) .remove(BufferFlags::F_CONTIGUOUS), @@ -1156,15 +1281,19 @@ impl Buffer { /// Works for non-square matrices. The buffer must be writeable. pub fn fill_diagonal(&mut self, value: T) -> MohuResult<()> { if self.ndim() != 2 { - return Err(MohuError::bug("fill_diagonal: requires exactly 2 dimensions")); + return Err(MohuError::bug( + "fill_diagonal: requires exactly 2 dimensions", + )); } if T::DTYPE != self.dtype { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: self.dtype.to_string(), + got: self.dtype.to_string(), }); } - if !self.is_writeable() { return Err(MohuError::ReadOnly); } + if !self.is_writeable() { + return Err(MohuError::ReadOnly); + } self.make_unique()?; let n = self.shape()[0].min(self.shape()[1]); @@ -1187,7 +1316,9 @@ impl Buffer { /// Computes the arithmetic mean of all elements as f64. pub fn mean_all_f64(&self) -> MohuResult { let n = self.len(); - if n == 0 { return Ok(f64::NAN); } + if n == 0 { + return Ok(f64::NAN); + } Ok(crate::ops::sum_all_f64(self)? / n as f64) } @@ -1213,7 +1344,9 @@ impl Buffer { }); } let n = self.len(); - if n <= ddof { return Ok(f64::NAN); } + if n <= ddof { + return Ok(f64::NAN); + } let mean = self.mean_all_f64()?; // Second pass: sum of squared deviations. @@ -1227,10 +1360,13 @@ impl Buffer { unsafe { std::slice::from_raw_parts(c.as_ptr(), n) } }; use rayon::prelude::*; - let ss: f64 = s.par_iter().map(|&x| { - let d = x as f64 - mean; - d * d - }).sum(); + let ss: f64 = s + .par_iter() + .map(|&x| { + let d = x as f64 - mean; + d * d + }) + .sum(); return Ok(ss / (n - ddof) as f64); } macro_rules! do_var { @@ -1243,26 +1379,29 @@ impl Buffer { unsafe { std::slice::from_raw_parts(c.as_ptr() as *const $T, n) } }; use rayon::prelude::*; - let ss: f64 = s.par_iter().map(|&x| { - let d = num_traits::cast::<$T, f64>(x).unwrap_or(0.0) - mean; - d * d - }).sum(); + let ss: f64 = s + .par_iter() + .map(|&x| { + let d = num_traits::cast::<$T, f64>(x).unwrap_or(0.0) - mean; + d * d + }) + .sum(); Ok(ss / (n - ddof) as f64) }}; } match self.dtype { - DType::I8 => do_var!(i8), - DType::I16 => do_var!(i16), - DType::I32 => do_var!(i32), - DType::I64 => do_var!(i64), - DType::U8 => do_var!(u8), - DType::U16 => do_var!(u16), - DType::U32 => do_var!(u32), - DType::U64 => do_var!(u64), - DType::F16 => do_var!(::half::f16), + DType::I8 => do_var!(i8), + DType::I16 => do_var!(i16), + DType::I32 => do_var!(i32), + DType::I64 => do_var!(i64), + DType::U8 => do_var!(u8), + DType::U16 => do_var!(u16), + DType::U32 => do_var!(u32), + DType::U64 => do_var!(u64), + DType::F16 => do_var!(::half::f16), DType::BF16 => do_var!(::half::bf16), - DType::F32 => do_var!(f32), - DType::F64 => do_var!(f64), + DType::F32 => do_var!(f32), + DType::F64 => do_var!(f64), _ => unreachable!(), } } @@ -1288,13 +1427,18 @@ impl Buffer { pub fn sum_axis(&self, axis: usize, keepdims: bool) -> MohuResult { if axis >= self.ndim() { return Err(MohuError::bug(format!( - "sum_axis: axis {axis} out of bounds for ndim {}", self.ndim() + "sum_axis: axis {axis} out of bounds for ndim {}", + self.ndim() ))); } // Build output shape let mut out_shape: Vec = self.shape().to_vec(); let axis_size = out_shape[axis]; - if keepdims { out_shape[axis] = 1; } else { out_shape.remove(axis); } + if keepdims { + out_shape[axis] = 1; + } else { + out_shape.remove(axis); + } let out = Self::zeros(DType::F64, &out_shape)?; let out_raw = unsafe { out.as_mut_ptr() as *mut f64 }; @@ -1307,7 +1451,7 @@ impl Buffer { // Build src index from out index by inserting the axis. let itemsize = self.dtype.itemsize(); - let src_raw = self.as_ptr(); + let src_raw = self.as_ptr(); for (out_flat, out_idx) in NdIndexIter::new(&out_shape_full).enumerate() { let mut acc = 0.0f64; @@ -1324,7 +1468,9 @@ impl Buffer { let val = read_as_f64(unsafe { src_raw.add(off) }, self.dtype, itemsize); acc += val; } - unsafe { out_raw.add(out_flat).write(acc); } + unsafe { + out_raw.add(out_flat).write(acc); + } } Ok(out) @@ -1348,7 +1494,7 @@ impl Buffer { unsafe { std::slice::from_raw_parts(c.as_ptr(), n) } }; return Ok(s.par_iter().any(|&x| x != 0)); - } + }, DType::C64 => { let c; let s: &[num_complex::Complex] = if self.is_c_contiguous() { @@ -1358,7 +1504,7 @@ impl Buffer { unsafe { std::slice::from_raw_parts(c.as_ptr() as *const _, n) } }; return Ok(s.par_iter().any(|x| x.re != 0.0 || x.im != 0.0)); - } + }, DType::C128 => { let c; let s: &[num_complex::Complex] = if self.is_c_contiguous() { @@ -1368,8 +1514,8 @@ impl Buffer { unsafe { std::slice::from_raw_parts(c.as_ptr() as *const _, n) } }; return Ok(s.par_iter().any(|x| x.re != 0.0 || x.im != 0.0)); - } - _ => {} + }, + _ => {}, } macro_rules! do_any { ($T:ty) => {{ @@ -1381,23 +1527,25 @@ impl Buffer { unsafe { std::slice::from_raw_parts(c.as_ptr() as *const $T, n) } }; Ok(s.par_iter().any(|&x| { - num_traits::cast::<$T, f64>(x).map(|v| v != 0.0).unwrap_or(false) + num_traits::cast::<$T, f64>(x) + .map(|v| v != 0.0) + .unwrap_or(false) })) }}; } match self.dtype { - DType::I8 => do_any!(i8), - DType::I16 => do_any!(i16), - DType::I32 => do_any!(i32), - DType::I64 => do_any!(i64), - DType::U8 => do_any!(u8), - DType::U16 => do_any!(u16), - DType::U32 => do_any!(u32), - DType::U64 => do_any!(u64), - DType::F16 => do_any!(::half::f16), + DType::I8 => do_any!(i8), + DType::I16 => do_any!(i16), + DType::I32 => do_any!(i32), + DType::I64 => do_any!(i64), + DType::U8 => do_any!(u8), + DType::U16 => do_any!(u16), + DType::U32 => do_any!(u32), + DType::U64 => do_any!(u64), + DType::F16 => do_any!(::half::f16), DType::BF16 => do_any!(::half::bf16), - DType::F32 => do_any!(f32), - DType::F64 => do_any!(f64), + DType::F32 => do_any!(f32), + DType::F64 => do_any!(f64), _ => unreachable!(), } } @@ -1417,7 +1565,7 @@ impl Buffer { unsafe { std::slice::from_raw_parts(c.as_ptr(), n) } }; return Ok(s.par_iter().all(|&x| x != 0)); - } + }, DType::C64 => { let c; let s: &[num_complex::Complex] = if self.is_c_contiguous() { @@ -1427,7 +1575,7 @@ impl Buffer { unsafe { std::slice::from_raw_parts(c.as_ptr() as *const _, n) } }; return Ok(s.par_iter().all(|x| x.re != 0.0 || x.im != 0.0)); - } + }, DType::C128 => { let c; let s: &[num_complex::Complex] = if self.is_c_contiguous() { @@ -1437,8 +1585,8 @@ impl Buffer { unsafe { std::slice::from_raw_parts(c.as_ptr() as *const _, n) } }; return Ok(s.par_iter().all(|x| x.re != 0.0 || x.im != 0.0)); - } - _ => {} + }, + _ => {}, } macro_rules! do_all { ($T:ty) => {{ @@ -1450,23 +1598,25 @@ impl Buffer { unsafe { std::slice::from_raw_parts(c.as_ptr() as *const $T, n) } }; Ok(s.par_iter().all(|&x| { - num_traits::cast::<$T, f64>(x).map(|v| v != 0.0).unwrap_or(false) + num_traits::cast::<$T, f64>(x) + .map(|v| v != 0.0) + .unwrap_or(false) })) }}; } match self.dtype { - DType::I8 => do_all!(i8), - DType::I16 => do_all!(i16), - DType::I32 => do_all!(i32), - DType::I64 => do_all!(i64), - DType::U8 => do_all!(u8), - DType::U16 => do_all!(u16), - DType::U32 => do_all!(u32), - DType::U64 => do_all!(u64), - DType::F16 => do_all!(::half::f16), + DType::I8 => do_all!(i8), + DType::I16 => do_all!(i16), + DType::I32 => do_all!(i32), + DType::I64 => do_all!(i64), + DType::U8 => do_all!(u8), + DType::U16 => do_all!(u16), + DType::U32 => do_all!(u32), + DType::U64 => do_all!(u64), + DType::F16 => do_all!(::half::f16), DType::BF16 => do_all!(::half::bf16), - DType::F32 => do_all!(f32), - DType::F64 => do_all!(f64), + DType::F32 => do_all!(f32), + DType::F64 => do_all!(f64), _ => unreachable!(), } } @@ -1486,7 +1636,7 @@ impl Buffer { unsafe { std::slice::from_raw_parts(c.as_ptr(), n) } }; return Ok(s.par_iter().filter(|&&x| x != 0).count()); - } + }, DType::C64 => { let c; let s: &[num_complex::Complex] = if self.is_c_contiguous() { @@ -1496,7 +1646,7 @@ impl Buffer { unsafe { std::slice::from_raw_parts(c.as_ptr() as *const _, n) } }; return Ok(s.par_iter().filter(|x| x.re != 0.0 || x.im != 0.0).count()); - } + }, DType::C128 => { let c; let s: &[num_complex::Complex] = if self.is_c_contiguous() { @@ -1506,8 +1656,8 @@ impl Buffer { unsafe { std::slice::from_raw_parts(c.as_ptr() as *const _, n) } }; return Ok(s.par_iter().filter(|x| x.re != 0.0 || x.im != 0.0).count()); - } - _ => {} + }, + _ => {}, } macro_rules! do_cnz { ($T:ty) => {{ @@ -1518,24 +1668,28 @@ impl Buffer { c = self.to_contiguous()?; unsafe { std::slice::from_raw_parts(c.as_ptr() as *const $T, n) } }; - Ok(s.par_iter().filter(|&&x| { - num_traits::cast::<$T, f64>(x).map(|v| v != 0.0).unwrap_or(false) - }).count()) + Ok(s.par_iter() + .filter(|&&x| { + num_traits::cast::<$T, f64>(x) + .map(|v| v != 0.0) + .unwrap_or(false) + }) + .count()) }}; } match self.dtype { - DType::I8 => do_cnz!(i8), - DType::I16 => do_cnz!(i16), - DType::I32 => do_cnz!(i32), - DType::I64 => do_cnz!(i64), - DType::U8 => do_cnz!(u8), - DType::U16 => do_cnz!(u16), - DType::U32 => do_cnz!(u32), - DType::U64 => do_cnz!(u64), - DType::F16 => do_cnz!(::half::f16), + DType::I8 => do_cnz!(i8), + DType::I16 => do_cnz!(i16), + DType::I32 => do_cnz!(i32), + DType::I64 => do_cnz!(i64), + DType::U8 => do_cnz!(u8), + DType::U16 => do_cnz!(u16), + DType::U32 => do_cnz!(u32), + DType::U64 => do_cnz!(u64), + DType::F16 => do_cnz!(::half::f16), DType::BF16 => do_cnz!(::half::bf16), - DType::F32 => do_cnz!(f32), - DType::F64 => do_cnz!(f64), + DType::F32 => do_cnz!(f32), + DType::F64 => do_cnz!(f64), _ => unreachable!(), } } @@ -1552,10 +1706,12 @@ impl Buffer { let a = self.cast(DType::F64, CastMode::Unsafe)?; let b = other.cast(DType::F64, CastMode::Unsafe)?; let as_ = a.as_slice::()?; - let bs = b.as_slice::()?; + let bs = b.as_slice::()?; use rayon::prelude::*; Ok(as_.par_iter().zip(bs.par_iter()).all(|(a, b)| { - if a.is_nan() && b.is_nan() { return true; } + if a.is_nan() && b.is_nan() { + return true; + } (a - b).abs() <= atol + rtol * b.abs() })) } @@ -1684,11 +1840,17 @@ impl Buffer { if cap > 0 { let _ = write!(s, " data[:{}]: [", cap); for i in 0..cap { - if i > 0 { let _ = write!(s, ", "); } + if i > 0 { + let _ = write!(s, ", "); + } // Convert flat index to multi-dim and read as f64 for display let idx = flat_to_indices(i, self.shape()); if let Ok(off) = self.layout.byte_offset(&idx) { - let v = read_as_f64(unsafe { self.raw.as_ptr().add(off) }, self.dtype, self.dtype.itemsize()); + let v = read_as_f64( + unsafe { self.raw.as_ptr().add(off) }, + self.dtype, + self.dtype.itemsize(), + ); let _ = write!(s, "{v:.4}"); } } @@ -1710,22 +1872,30 @@ impl std::fmt::Display for Buffer { } fn fmt_buffer_data( - buf: &Buffer, - f: &mut std::fmt::Formatter<'_>, - shape: &[usize], - idx: &mut Vec, - dim: usize, + buf: &Buffer, + f: &mut std::fmt::Formatter<'_>, + shape: &[usize], + idx: &mut Vec, + dim: usize, _indent: usize, ) -> std::fmt::Result { if shape.is_empty() { // Scalar let off = buf.layout.byte_offset(idx).map_err(|_| std::fmt::Error)?; - let v = read_as_f64(unsafe { buf.raw.as_ptr().add(off) }, buf.dtype, buf.dtype.itemsize()); + let v = read_as_f64( + unsafe { buf.raw.as_ptr().add(off) }, + buf.dtype, + buf.dtype.itemsize(), + ); return write!(f, "{v}"); } if dim == shape.len() { let off = buf.layout.byte_offset(idx).map_err(|_| std::fmt::Error)?; - let v = read_as_f64(unsafe { buf.raw.as_ptr().add(off) }, buf.dtype, buf.dtype.itemsize()); + let v = read_as_f64( + unsafe { buf.raw.as_ptr().add(off) }, + buf.dtype, + buf.dtype.itemsize(), + ); return write!(f, "{v}"); } @@ -1735,7 +1905,9 @@ fn fmt_buffer_data( write!(f, "[")?; let show = n.min(MAX_DISPLAY); for i in 0..show { - if i > 0 { write!(f, ", ")?; } + if i > 0 { + write!(f, ", ")?; + } idx[dim] = i; fmt_buffer_data(buf, f, shape, idx, dim + 1, _indent + 1)?; } @@ -1763,8 +1935,14 @@ impl PartialEq for Buffer { { return true; // literally the same view } - let a = match self.to_contiguous() { Ok(b) => b, Err(_) => return false }; - let b = match other.to_contiguous() { Ok(b) => b, Err(_) => return false }; + let a = match self.to_contiguous() { + Ok(b) => b, + Err(_) => return false, + }; + let b = match other.to_contiguous() { + Ok(b) => b, + Err(_) => return false, + }; let ab = unsafe { std::slice::from_raw_parts(a.as_ptr(), a.nbytes()) }; let bb = unsafe { std::slice::from_raw_parts(b.as_ptr(), b.nbytes()) }; ab == bb @@ -1774,19 +1952,29 @@ impl PartialEq for Buffer { // ─── From impls ─────────────────────────────────────────────────────────────── impl From> for Buffer { - fn from(v: Vec) -> Self { Buffer::from_vec(v).expect("from Vec") } + fn from(v: Vec) -> Self { + Buffer::from_vec(v).expect("from Vec") + } } impl From> for Buffer { - fn from(v: Vec) -> Self { Buffer::from_vec(v).expect("from Vec") } + fn from(v: Vec) -> Self { + Buffer::from_vec(v).expect("from Vec") + } } impl From> for Buffer { - fn from(v: Vec) -> Self { Buffer::from_vec(v).expect("from Vec") } + fn from(v: Vec) -> Self { + Buffer::from_vec(v).expect("from Vec") + } } impl From> for Buffer { - fn from(v: Vec) -> Self { Buffer::from_vec(v).expect("from Vec") } + fn from(v: Vec) -> Self { + Buffer::from_vec(v).expect("from Vec") + } } impl From> for Buffer { - fn from(v: Vec) -> Self { Buffer::from_vec(v).expect("from Vec") } + fn from(v: Vec) -> Self { + Buffer::from_vec(v).expect("from Vec") + } } // ─── Internal helpers ───────────────────────────────────────────────────────── @@ -1817,26 +2005,26 @@ fn read_as_f64(ptr: *const u8, dtype: DType, _itemsize: usize) -> f64 { use mohu_dtype::DType::*; match dtype { Bool => unsafe { *(ptr as *const u8) as f64 }, - I8 => unsafe { *(ptr as *const i8) as f64 }, - U8 => unsafe { *(ptr as *const u8) as f64 }, - I16 => unsafe { *(ptr as *const i16) as f64 }, - U16 => unsafe { *(ptr as *const u16) as f64 }, - I32 => unsafe { *(ptr as *const i32) as f64 }, - U32 => unsafe { *(ptr as *const u32) as f64 }, - I64 => unsafe { *(ptr as *const i64) as f64 }, - U64 => unsafe { *(ptr as *const u64) as f64 }, - F32 => unsafe { *(ptr as *const f32) as f64 }, - F64 => unsafe { *(ptr as *const f64) }, - F16 => { + I8 => unsafe { *(ptr as *const i8) as f64 }, + U8 => unsafe { *(ptr as *const u8) as f64 }, + I16 => unsafe { *(ptr as *const i16) as f64 }, + U16 => unsafe { *(ptr as *const u16) as f64 }, + I32 => unsafe { *(ptr as *const i32) as f64 }, + U32 => unsafe { *(ptr as *const u32) as f64 }, + I64 => unsafe { *(ptr as *const i64) as f64 }, + U64 => unsafe { *(ptr as *const u64) as f64 }, + F32 => unsafe { *(ptr as *const f32) as f64 }, + F64 => unsafe { *(ptr as *const f64) }, + F16 => { let bits = unsafe { (ptr as *const u16).read_unaligned() }; half::f16::from_bits(bits).to_f64() - } + }, BF16 => { let bits = unsafe { (ptr as *const u16).read_unaligned() }; half::bf16::from_bits(bits).to_f64() - } - C64 => unsafe { *(ptr as *const f32) as f64 }, // real part - C128 => unsafe { *(ptr as *const f64) }, // real part + }, + C64 => unsafe { *(ptr as *const f32) as f64 }, // real part + C128 => unsafe { *(ptr as *const f64) }, // real part } } @@ -1851,6 +2039,6 @@ fn flat_to_indices(mut flat: usize, shape: &[usize]) -> Vec { idx } -use num_traits; #[cfg(unix)] use libc; +use num_traits; diff --git a/crates/mohu-buffer/src/layout.rs b/crates/mohu-buffer/src/layout.rs index 1c6a5c3..b5d3bd9 100644 --- a/crates/mohu-buffer/src/layout.rs +++ b/crates/mohu-buffer/src/layout.rs @@ -16,8 +16,7 @@ use mohu_error::{MohuError, MohuResult}; use crate::strides::{ - self, ShapeVec, StrideVec, broadcast_strides, c_strides, f_strides, - validate_strides, + self, ShapeVec, StrideVec, broadcast_strides, c_strides, f_strides, validate_strides, }; // ─── Order ──────────────────────────────────────────────────────────────────── @@ -41,13 +40,17 @@ pub enum Order { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct SliceArg { pub start: Option, - pub stop: Option, - pub step: Option, + pub stop: Option, + pub step: Option, } impl SliceArg { /// A slice that selects the full axis (`..`). - pub const FULL: Self = Self { start: None, stop: None, step: None }; + pub const FULL: Self = Self { + start: None, + stop: None, + step: None, + }; /// Resolves this `SliceArg` against an axis of length `dim`, returning /// `(start_index, element_count, step)` in element units. @@ -79,18 +82,28 @@ impl SliceArg { let e = self .stop .map(|v| { - if v < 0 { ((v + idim).max(-1)) as usize } else { (v as usize).min(dim) } + if v < 0 { + ((v + idim).max(-1)) as usize + } else { + (v as usize).min(dim) + } }) .unwrap_or(usize::MAX); // sentinel for "before index 0" (s, e) }; let count = if step > 0 { - if stop <= start { 0 } else { (stop - start + (step as usize) - 1) / (step as usize) } + if stop <= start { + 0 + } else { + (stop - start).div_ceil(step as usize) + } } else { let abs_step = (-step) as usize; - if e_reversed_empty(start, stop) { 0 } else { - (start.saturating_sub(stop) + abs_step - 1) / abs_step + if e_reversed_empty(start, stop) { + 0 + } else { + start.saturating_sub(stop).div_ceil(abs_step) } }; @@ -111,10 +124,10 @@ fn e_reversed_empty(start: usize, stop: usize) -> bool { /// Does not own memory — it is always paired with a backing `Buffer`. #[derive(Clone, PartialEq, Eq)] pub struct Layout { - shape: ShapeVec, - strides: StrideVec, + shape: ShapeVec, + strides: StrideVec, /// Byte offset from buffer start to element `[0, 0, …, 0]`. - offset: usize, + offset: usize, /// Size in bytes of one scalar element. itemsize: usize, } @@ -140,15 +153,16 @@ impl Layout { /// - `shape.len() == strides.len()` /// - All non-broadcast strides are multiples of `itemsize` pub fn new_custom( - shape: &[usize], - strides: &[isize], - offset: usize, + shape: &[usize], + strides: &[isize], + offset: usize, itemsize: usize, ) -> MohuResult { if shape.len() != strides.len() { return Err(MohuError::bug(format!( "Layout::new_custom: shape.len()={} != strides.len()={}", - shape.len(), strides.len() + shape.len(), + strides.len() ))); } if itemsize == 0 { @@ -156,8 +170,8 @@ impl Layout { } validate_strides(shape, strides, itemsize, false)?; Ok(Self { - shape: ShapeVec::from_slice(shape), - strides: StrideVec::from_slice(strides), + shape: ShapeVec::from_slice(shape), + strides: StrideVec::from_slice(strides), offset, itemsize, }) @@ -166,9 +180,9 @@ impl Layout { /// Creates a 0-dimensional (scalar) layout holding exactly one element. pub fn scalar(itemsize: usize) -> Self { Self { - shape: ShapeVec::new(), - strides: StrideVec::new(), - offset: 0, + shape: ShapeVec::new(), + strides: StrideVec::new(), + offset: 0, itemsize, } } @@ -176,23 +190,40 @@ impl Layout { // ─── Properties ─────────────────────────────────────────────────────────── /// Number of dimensions (axes). - #[inline] pub fn ndim(&self) -> usize { self.shape.len() } + #[inline] + pub fn ndim(&self) -> usize { + self.shape.len() + } /// Total number of elements (product of shape). #[inline] - pub fn size(&self) -> usize { self.shape.iter().product() } + pub fn size(&self) -> usize { + self.shape.iter().product() + } /// The shape of the array, as a slice of dimension sizes. - #[inline] pub fn shape(&self) -> &[usize] { &self.shape } + #[inline] + pub fn shape(&self) -> &[usize] { + &self.shape + } /// The byte strides of the array. - #[inline] pub fn strides(&self) -> &[isize] { &self.strides } + #[inline] + pub fn strides(&self) -> &[isize] { + &self.strides + } /// Byte offset of element `[0, 0, …, 0]` from the buffer start. - #[inline] pub fn offset(&self) -> usize { self.offset } + #[inline] + pub fn offset(&self) -> usize { + self.offset + } /// Size in bytes of a single element. - #[inline] pub fn itemsize(&self) -> usize { self.itemsize } + #[inline] + pub fn itemsize(&self) -> usize { + self.itemsize + } /// Total bytes in a contiguous copy of this array (`size * itemsize`). /// @@ -214,10 +245,15 @@ impl Layout { } /// Returns `true` if this is a 0-dimensional (scalar) array. - #[inline] pub fn is_scalar(&self) -> bool { self.ndim() == 0 } + #[inline] + pub fn is_scalar(&self) -> bool { + self.ndim() == 0 + } /// Returns `true` if any dimension is 0 (zero-element array). - pub fn is_empty(&self) -> bool { self.shape.iter().any(|&d| d == 0) } + pub fn is_empty(&self) -> bool { + self.shape.contains(&0) + } // ─── Contiguity checks ──────────────────────────────────────────────────── @@ -254,14 +290,17 @@ impl Layout { pub fn permute(&self, axes: &[usize]) -> MohuResult { let ndim = self.ndim(); if axes.len() != ndim { - return Err(MohuError::DimensionMismatch { expected: ndim, got: axes.len() }); + return Err(MohuError::DimensionMismatch { + expected: ndim, + got: axes.len(), + }); } // Validate that `axes` is a proper permutation. let mut seen = vec![false; ndim]; for &ax in axes { if ax >= ndim { return Err(MohuError::AxisOutOfRange { - axis: ax as i64, + axis: ax as i64, ndim, valid: format!("0..{ndim}"), }); @@ -273,12 +312,12 @@ impl Layout { } seen[ax] = true; } - let new_shape: ShapeVec = axes.iter().map(|&a| self.shape[a]).collect(); + let new_shape: ShapeVec = axes.iter().map(|&a| self.shape[a]).collect(); let new_strides: StrideVec = axes.iter().map(|&a| self.strides[a]).collect(); Ok(Self { - shape: new_shape, - strides: new_strides, - offset: self.offset, + shape: new_shape, + strides: new_strides, + offset: self.offset, itemsize: self.itemsize, }) } @@ -303,12 +342,12 @@ impl Layout { let ndim = self.ndim(); if axis > ndim { return Err(MohuError::AxisOutOfRange { - axis: axis as i64, + axis: axis as i64, ndim, valid: format!("0..={ndim}"), }); } - let mut new_shape = ShapeVec::with_capacity(ndim + 1); + let mut new_shape = ShapeVec::with_capacity(ndim + 1); let mut new_strides = StrideVec::with_capacity(ndim + 1); for i in 0..=ndim { if i == axis { @@ -321,9 +360,9 @@ impl Layout { } } Ok(Self { - shape: new_shape, - strides: new_strides, - offset: self.offset, + shape: new_shape, + strides: new_strides, + offset: self.offset, itemsize: self.itemsize, }) } @@ -332,15 +371,17 @@ impl Layout { /// /// Equivalent to `np.squeeze(a)`. pub fn squeeze(&self) -> Self { - let new_shape: ShapeVec = self.shape.iter().copied().filter(|&d| d != 1).collect(); + let new_shape: ShapeVec = self.shape.iter().copied().filter(|&d| d != 1).collect(); let new_strides: StrideVec = self - .shape.iter().zip(self.strides.iter()) + .shape + .iter() + .zip(self.strides.iter()) .filter_map(|(&d, &s)| if d != 1 { Some(s) } else { None }) .collect(); Self { - shape: new_shape, - strides: new_strides, - offset: self.offset, + shape: new_shape, + strides: new_strides, + offset: self.offset, itemsize: self.itemsize, } } @@ -352,7 +393,7 @@ impl Layout { let ndim = self.ndim(); if axis >= ndim { return Err(MohuError::AxisOutOfRange { - axis: axis as i64, + axis: axis as i64, ndim, valid: format!("0..{ndim}"), }); @@ -363,14 +404,22 @@ impl Layout { self.shape[axis] ))); } - let new_shape: ShapeVec = self.shape.iter().enumerate() - .filter_map(|(i, &d)| if i != axis { Some(d) } else { None }).collect(); - let new_strides: StrideVec = self.strides.iter().enumerate() - .filter_map(|(i, &s)| if i != axis { Some(s) } else { None }).collect(); + let new_shape: ShapeVec = self + .shape + .iter() + .enumerate() + .filter_map(|(i, &d)| if i != axis { Some(d) } else { None }) + .collect(); + let new_strides: StrideVec = self + .strides + .iter() + .enumerate() + .filter_map(|(i, &s)| if i != axis { Some(s) } else { None }) + .collect(); Ok(Self { - shape: new_shape, - strides: new_strides, - offset: self.offset, + shape: new_shape, + strides: new_strides, + offset: self.offset, itemsize: self.itemsize, }) } @@ -394,9 +443,9 @@ impl Layout { } let new_strides = c_strides(new_shape, self.itemsize); Ok(Self { - shape: ShapeVec::from_slice(new_shape), - strides: new_strides, - offset: self.offset, + shape: ShapeVec::from_slice(new_shape), + strides: new_strides, + offset: self.offset, itemsize: self.itemsize, }) } @@ -409,7 +458,7 @@ impl Layout { let ndim = self.ndim(); if axis >= ndim { return Err(MohuError::AxisOutOfRange { - axis: axis as i64, + axis: axis as i64, ndim, valid: format!("0..{ndim}"), }); @@ -418,22 +467,20 @@ impl Layout { let (start, count, step) = arg.resolve(dim)?; // Advance the base offset by `start` steps along this axis. - let new_offset = (self.offset as isize - + start as isize * self.strides[axis]) - as usize; + let new_offset = (self.offset as isize + start as isize * self.strides[axis]) as usize; // Multiply the stride by the step size. let new_stride = self.strides[axis] * step; - let mut new_shape = self.shape.clone(); + let mut new_shape = self.shape.clone(); let mut new_strides = self.strides.clone(); - new_shape[axis] = count; + new_shape[axis] = count; new_strides[axis] = new_stride; Ok(Self { - shape: new_shape, - strides: new_strides, - offset: new_offset, + shape: new_shape, + strides: new_strides, + offset: new_offset, itemsize: self.itemsize, }) } @@ -445,9 +492,9 @@ impl Layout { pub fn broadcast_to(&self, new_shape: &[usize]) -> MohuResult { let new_strides = broadcast_strides(&self.shape, &self.strides, new_shape)?; Ok(Self { - shape: ShapeVec::from_slice(new_shape), - strides: new_strides, - offset: self.offset, + shape: ShapeVec::from_slice(new_shape), + strides: new_strides, + offset: self.offset, itemsize: self.itemsize, }) } @@ -461,7 +508,7 @@ impl Layout { if indices.len() != self.ndim() { return Err(MohuError::TooManyIndices { given: indices.len(), - ndim: self.ndim(), + ndim: self.ndim(), }); } for (axis, (&idx, &dim)) in indices.iter().zip(self.shape.iter()).enumerate() { @@ -469,7 +516,7 @@ impl Layout { return Err(MohuError::IndexOutOfBounds { index: idx as i64, axis, - size: dim, + size: dim, }); } } @@ -500,9 +547,15 @@ impl Layout { let mut lo = self.offset as isize; let mut hi = self.offset as isize; for (&dim, &stride) in self.shape.iter().zip(self.strides.iter()) { - if stride == 0 || dim == 0 { continue; } + if stride == 0 || dim == 0 { + continue; + } let span = stride * (dim as isize - 1); - if span > 0 { hi += span; } else { lo += span; } + if span > 0 { + hi += span; + } else { + lo += span; + } } (lo as usize, hi as usize) } @@ -530,9 +583,9 @@ impl Layout { /// suitable for a freshly allocated contiguous copy. pub fn to_c_contiguous(&self) -> Self { Self { - shape: self.shape.clone(), - strides: c_strides(&self.shape, self.itemsize), - offset: 0, + shape: self.shape.clone(), + strides: c_strides(&self.shape, self.itemsize), + offset: 0, itemsize: self.itemsize, } } @@ -557,9 +610,9 @@ impl Layout { impl std::fmt::Debug for Layout { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Layout") - .field("shape", &self.shape.as_slice()) - .field("strides", &self.strides.as_slice()) - .field("offset", &self.offset) + .field("shape", &self.shape.as_slice()) + .field("strides", &self.strides.as_slice()) + .field("offset", &self.offset) .field("itemsize", &self.itemsize) .finish() } @@ -567,9 +620,12 @@ impl std::fmt::Debug for Layout { impl std::fmt::Display for Layout { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "Layout(shape={:?}, strides={:?}, itemsize={})", - self.shape.as_slice(), - self.strides.as_slice(), - self.itemsize) + write!( + f, + "Layout(shape={:?}, strides={:?}, itemsize={})", + self.shape.as_slice(), + self.strides.as_slice(), + self.itemsize + ) } } diff --git a/crates/mohu-buffer/src/lib.rs b/crates/mohu-buffer/src/lib.rs index afe4c4a..23a7be7 100644 --- a/crates/mohu-buffer/src/lib.rs +++ b/crates/mohu-buffer/src/lib.rs @@ -23,30 +23,24 @@ pub mod view; // ─── Re-exports ─────────────────────────────────────────────────────────────── -pub use alloc::{ - AllocHandle, AllocStats, Strategy, - CACHE_LINE, MMAP_THRESHOLD, SIMD_ALIGN, -}; +pub use alloc::{AllocHandle, AllocStats, CACHE_LINE, MMAP_THRESHOLD, SIMD_ALIGN, Strategy}; pub use buffer::{ - Buffer, BufferFlags, RawBuffer, - DLManagedTensor, DLTensor, RawDLDataType, RawDLDevice, + Buffer, BufferFlags, DLManagedTensor, DLTensor, RawBuffer, RawDLDataType, RawDLDevice, }; pub use layout::{Layout, Order, SliceArg}; pub use ops::{ - cast_copy, copy_to_contiguous, fill, fill_one, fill_raw, fill_zero, - parallel_inplace, parallel_map, reduce, + cast_copy, copy_to_contiguous, fill, fill_one, fill_raw, fill_zero, parallel_inplace, + parallel_map, reduce, }; -pub use pool::{BufferPool, PoolStats, GLOBAL_POOL}; +pub use pool::{BufferPool, GLOBAL_POOL, PoolStats}; pub use strides::{ - NdIndexIter, ShapeVec, StrideVec, StridedByteIter, - broadcast_strides, c_strides, contiguous_nbytes, f_strides, - ravel_multi_index, shape_size, unravel_index, validate_strides, - byte_offset, + NdIndexIter, ShapeVec, StrideVec, StridedByteIter, broadcast_strides, byte_offset, c_strides, + contiguous_nbytes, f_strides, ravel_multi_index, shape_size, unravel_index, validate_strides, }; pub use view::{BufferView, BufferViewMut}; diff --git a/crates/mohu-buffer/src/ops.rs b/crates/mohu-buffer/src/ops.rs index 0024248..6a7dc01 100644 --- a/crates/mohu-buffer/src/ops.rs +++ b/crates/mohu-buffer/src/ops.rs @@ -14,15 +14,12 @@ use rayon::prelude::*; use mohu_dtype::{ cast::cast_scalar_unchecked, dispatch_dtype, - promote::{can_cast, CastMode}, + promote::{CastMode, can_cast}, scalar::Scalar, }; use mohu_error::{MohuError, MohuResult}; -use crate::{ - buffer::Buffer, - strides::StridedByteIter, -}; +use crate::{buffer::Buffer, strides::StridedByteIter}; // ─── fill_raw ──────────────────────────────────────────────────────────────── @@ -37,7 +34,8 @@ pub fn fill_raw(buf: &mut Buffer, fill_bytes: &[u8]) -> MohuResult<()> { if fill_bytes.len() != itemsize { return Err(MohuError::bug(format!( "fill_raw: fill_bytes.len()={} != itemsize={}", - fill_bytes.len(), itemsize + fill_bytes.len(), + itemsize ))); } if !buf.is_writeable() { @@ -48,9 +46,7 @@ pub fn fill_raw(buf: &mut Buffer, fill_bytes: &[u8]) -> MohuResult<()> { if buf.is_c_contiguous() { let total = buf.len() * itemsize; // SAFETY: buf is uniquely owned, C-contiguous, pointer is valid. - let slice = unsafe { - std::slice::from_raw_parts_mut(buf.as_mut_ptr(), total) - }; + let slice = unsafe { std::slice::from_raw_parts_mut(buf.as_mut_ptr(), total) }; // Parallel fill: split into 4KiB chunks and fill each on Rayon threads. let chunk_size = 4096.max(itemsize * 64); slice.par_chunks_mut(chunk_size).for_each(|chunk| { @@ -66,11 +62,7 @@ pub fn fill_raw(buf: &mut Buffer, fill_bytes: &[u8]) -> MohuResult<()> { for off in StridedByteIter::new(buf.shape(), buf.strides(), buf.offset()) { // SAFETY: stride iterator yields valid byte offsets within the buffer. unsafe { - std::ptr::copy_nonoverlapping( - fill_bytes.as_ptr(), - raw_ptr.add(off), - itemsize, - ); + std::ptr::copy_nonoverlapping(fill_bytes.as_ptr(), raw_ptr.add(off), itemsize); } } } @@ -89,7 +81,7 @@ where if T::DTYPE != buf.dtype() { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: buf.dtype().to_string(), + got: buf.dtype().to_string(), }); } if !buf.is_writeable() { @@ -140,9 +132,7 @@ pub fn fill_zero(buf: &mut Buffer) -> MohuResult<()> { /// Dispatches at runtime over all 15 dtypes. pub fn fill_one(buf: &mut Buffer) -> MohuResult<()> { macro_rules! do_fill_one { - ($T:ty) => {{ - fill(buf, <$T as Scalar>::ONE) - }}; + ($T:ty) => {{ fill(buf, <$T as Scalar>::ONE) }}; } dispatch_dtype!(buf.dtype(), do_fill_one) } @@ -157,13 +147,13 @@ pub fn copy_to_contiguous(src: &Buffer, dst: &mut Buffer) -> MohuResult<()> { if src.dtype() != dst.dtype() { return Err(MohuError::DTypeMismatch { expected: src.dtype().to_string(), - got: dst.dtype().to_string(), + got: dst.dtype().to_string(), }); } if src.len() != dst.len() { return Err(MohuError::ShapeMismatch { expected: src.shape().to_vec(), - got: dst.shape().to_vec(), + got: dst.shape().to_vec(), }); } if !dst.is_writeable() { @@ -176,7 +166,7 @@ pub fn copy_to_contiguous(src: &Buffer, dst: &mut Buffer) -> MohuResult<()> { if src.is_c_contiguous() && dst.is_c_contiguous() { // Parallel memcpy: split into cache-friendly 64 KiB chunks. let src_bytes = src.len() * itemsize; - let chunk = 65536_usize.max(itemsize); + let chunk = 65536_usize.max(itemsize); unsafe { let s = std::slice::from_raw_parts(src.as_ptr(), src_bytes); let d = std::slice::from_raw_parts_mut(dst.as_mut_ptr(), src_bytes); @@ -190,16 +180,12 @@ pub fn copy_to_contiguous(src: &Buffer, dst: &mut Buffer) -> MohuResult<()> { let dst_raw = unsafe { dst.as_mut_ptr() }; let dst_itemsize = dst.itemsize(); - for (elem_idx, src_off) in StridedByteIter::new( - src.shape(), src.strides(), src.offset(), - ).enumerate() { + for (elem_idx, src_off) in + StridedByteIter::new(src.shape(), src.strides(), src.offset()).enumerate() + { let dst_off = elem_idx * dst_itemsize; unsafe { - std::ptr::copy_nonoverlapping( - src_raw.add(src_off), - dst_raw.add(dst_off), - itemsize, - ); + std::ptr::copy_nonoverlapping(src_raw.add(src_off), dst_raw.add(dst_off), itemsize); } } } @@ -217,7 +203,7 @@ pub fn cast_copy(src: &Buffer, dst: &mut Buffer, mode: CastMode) -> MohuResult<( if src.len() != dst.len() { return Err(MohuError::ShapeMismatch { expected: src.shape().to_vec(), - got: dst.shape().to_vec(), + got: dst.shape().to_vec(), }); } if src.dtype() == dst.dtype() { @@ -225,8 +211,8 @@ pub fn cast_copy(src: &Buffer, dst: &mut Buffer, mode: CastMode) -> MohuResult<( } if !can_cast(src.dtype(), dst.dtype(), mode) { return Err(MohuError::InvalidCast { - from: src.dtype().to_string(), - to: dst.dtype().to_string(), + from: src.dtype().to_string(), + to: dst.dtype().to_string(), reason: format!("{mode:?} cast is not permitted"), }); } @@ -241,9 +227,7 @@ pub fn cast_copy(src: &Buffer, dst: &mut Buffer, mode: CastMode) -> MohuResult<( macro_rules! cast_src { ($S:ty) => {{ macro_rules! cast_dst { - ($D:ty) => {{ - cast_typed::<$S, $D>(src, dst) - }}; + ($D:ty) => {{ cast_typed::<$S, $D>(src, dst) }}; } dispatch_dtype!(dst_dtype, cast_dst) }}; @@ -278,9 +262,9 @@ fn cast_typed( let dst_raw = unsafe { dst.as_mut_ptr() }; let dst_itemsize = D::ITEMSIZE; - for (elem_idx, src_off) in StridedByteIter::new( - src.shape(), src.strides(), src.offset(), - ).enumerate() { + for (elem_idx, src_off) in + StridedByteIter::new(src.shape(), src.strides(), src.offset()).enumerate() + { let dst_off = elem_idx * dst_itemsize; unsafe { let s = (src_raw.add(src_off) as *const S).read_unaligned(); @@ -308,7 +292,7 @@ where if src.len() != dst.len() { return Err(MohuError::ShapeMismatch { expected: src.shape().to_vec(), - got: dst.shape().to_vec(), + got: dst.shape().to_vec(), }); } if !dst.is_writeable() { @@ -359,12 +343,7 @@ where /// /// Computes `init` combined with every element using `combine`. /// Uses Rayon's `map_reduce` to parallelise across chunks. -pub fn reduce( - buf: &Buffer, - init: R, - map_fn: F, - combine: G, -) -> MohuResult +pub fn reduce(buf: &Buffer, init: R, map_fn: F, combine: G) -> MohuResult where T: Scalar + Send + Sync, R: Clone + Send + Sync, @@ -374,7 +353,7 @@ where if T::DTYPE != buf.dtype() { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: buf.dtype().to_string(), + got: buf.dtype().to_string(), }); } if !buf.is_c_contiguous() { @@ -403,7 +382,7 @@ where if T::DTYPE != buf.dtype() { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: buf.dtype().to_string(), + got: buf.dtype().to_string(), }); } if !buf.is_writeable() { @@ -419,7 +398,9 @@ where slice.par_iter_mut().enumerate().for_each(|(i, v)| { // start + step * i — computed without shared mutable state let mut acc = start; - for _ in 0..i { acc = acc + step; } // naive but correct; compiler may vectorise + for _ in 0..i { + acc = acc + step; + } // naive but correct; compiler may vectorise *v = acc; }); } else { @@ -427,7 +408,9 @@ where let raw_ptr = unsafe { buf.as_mut_ptr() }; for (i, off) in StridedByteIter::new(buf.shape(), buf.strides(), buf.offset()).enumerate() { let mut acc = start; - for _ in 0..i { acc = acc + step; } + for _ in 0..i { + acc = acc + step; + } unsafe { (raw_ptr.add(off) as *mut T).write_unaligned(acc); } @@ -450,12 +433,7 @@ where /// 3. **Up-sweep**: each chunk adds its carry offset in parallel. /// /// This achieves O(n) work and O(log n) span, identical to serial complexity. -pub fn parallel_scan( - src: &Buffer, - dst: &mut Buffer, - identity: T, - f: F, -) -> MohuResult<()> +pub fn parallel_scan(src: &Buffer, dst: &mut Buffer, identity: T, f: F) -> MohuResult<()> where T: Scalar + Copy + Send + Sync, F: Fn(T, T) -> T + Send + Sync + Copy, @@ -463,13 +441,13 @@ where if T::DTYPE != src.dtype() || T::DTYPE != dst.dtype() { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: src.dtype().to_string(), + got: src.dtype().to_string(), }); } if src.len() != dst.len() { return Err(MohuError::ShapeMismatch { expected: src.shape().to_vec(), - got: dst.shape().to_vec(), + got: dst.shape().to_vec(), }); } if !src.is_c_contiguous() || !dst.is_c_contiguous() { @@ -480,14 +458,14 @@ where } dst.make_unique()?; - let len = src.len(); - let src_s = unsafe { std::slice::from_raw_parts(src.as_ptr() as *const T, len) }; - let dst_s = unsafe { std::slice::from_raw_parts_mut(dst.as_mut_ptr() as *mut T, len) }; + let len = src.len(); + let src_s = unsafe { std::slice::from_raw_parts(src.as_ptr() as *const T, len) }; + let dst_s = unsafe { std::slice::from_raw_parts_mut(dst.as_mut_ptr() as *mut T, len) }; // Chunk size: balance parallelism vs. carry overhead. - let n_threads = rayon::current_num_threads().max(1); + let n_threads = rayon::current_num_threads().max(1); let chunk_size = (len / n_threads).max(256).next_power_of_two(); - let n_chunks = len.div_ceil(chunk_size); + let n_chunks = len.div_ceil(chunk_size); // Step 1: local scans + collect chunk totals. let mut chunk_totals: Vec = vec![identity; n_chunks]; @@ -535,29 +513,29 @@ where /// All four buffers must be C-contiguous. pub fn where_select( mask: &Buffer, - a: &Buffer, - b: &Buffer, - dst: &mut Buffer, + a: &Buffer, + b: &Buffer, + dst: &mut Buffer, ) -> MohuResult<()> { use mohu_dtype::DType; if mask.dtype() != DType::U8 { return Err(MohuError::DTypeMismatch { expected: "U8".to_string(), - got: mask.dtype().to_string(), + got: mask.dtype().to_string(), }); } for (name, buf) in [("a", a), ("b", b)] { if T::DTYPE != buf.dtype() { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: buf.dtype().to_string(), + got: buf.dtype().to_string(), }); } if buf.shape() != mask.shape() { return Err(MohuError::ShapeMismatch { expected: mask.shape().to_vec(), - got: buf.shape().to_vec(), + got: buf.shape().to_vec(), }); } if !buf.is_c_contiguous() { @@ -565,20 +543,22 @@ pub fn where_select( return Err(MohuError::NonContiguous); } } - if !dst.is_writeable() { return Err(MohuError::ReadOnly); } + if !dst.is_writeable() { + return Err(MohuError::ReadOnly); + } if dst.shape() != mask.shape() { return Err(MohuError::ShapeMismatch { expected: mask.shape().to_vec(), - got: dst.shape().to_vec(), + got: dst.shape().to_vec(), }); } dst.make_unique()?; - let len = mask.len(); - let m_ptr = mask.as_ptr(); - let a_ptr = a.as_ptr() as *const T; - let b_ptr = b.as_ptr() as *const T; - let d_ptr = unsafe { dst.as_mut_ptr() } as *mut T; + let len = mask.len(); + let m_ptr = mask.as_ptr(); + let a_ptr = a.as_ptr() as *const T; + let b_ptr = b.as_ptr() as *const T; + let d_ptr = unsafe { dst.as_mut_ptr() } as *mut T; let m_s = unsafe { std::slice::from_raw_parts(m_ptr, len) }; let a_s = unsafe { std::slice::from_raw_parts(a_ptr, len) }; @@ -608,29 +588,37 @@ where if T::DTYPE != src.dtype() || T::DTYPE != dst.dtype() { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: src.dtype().to_string(), + got: src.dtype().to_string(), }); } if src.len() != dst.len() { return Err(MohuError::ShapeMismatch { expected: src.shape().to_vec(), - got: dst.shape().to_vec(), + got: dst.shape().to_vec(), }); } if !src.is_c_contiguous() || !dst.is_c_contiguous() { return Err(MohuError::NonContiguous); } - if !dst.is_writeable() { return Err(MohuError::ReadOnly); } + if !dst.is_writeable() { + return Err(MohuError::ReadOnly); + } dst.make_unique()?; - let len = src.len(); + let len = src.len(); let s_ptr = src.as_ptr() as *const T; let d_ptr = unsafe { dst.as_mut_ptr() } as *mut T; - let s_s = unsafe { std::slice::from_raw_parts(s_ptr, len) }; - let d_s = unsafe { std::slice::from_raw_parts_mut(d_ptr, len) }; + let s_s = unsafe { std::slice::from_raw_parts(s_ptr, len) }; + let d_s = unsafe { std::slice::from_raw_parts_mut(d_ptr, len) }; s_s.par_iter().zip(d_s.par_iter_mut()).for_each(|(s, d)| { - *d = if *s < lo { lo } else if *s > hi { hi } else { *s }; + *d = if *s < lo { + lo + } else if *s > hi { + hi + } else { + *s + }; }); Ok(()) @@ -648,13 +636,13 @@ pub fn gather(src: &Buffer, indices: &Buffer, dst: &mut Buffer) -> MohuResult<() if indices.dtype() != DType::I64 { return Err(MohuError::DTypeMismatch { expected: "I64".to_string(), - got: indices.dtype().to_string(), + got: indices.dtype().to_string(), }); } if src.dtype() != dst.dtype() { return Err(MohuError::DTypeMismatch { expected: src.dtype().to_string(), - got: dst.dtype().to_string(), + got: dst.dtype().to_string(), }); } if !src.is_c_contiguous() || !indices.is_c_contiguous() || !dst.is_c_contiguous() { @@ -663,39 +651,43 @@ pub fn gather(src: &Buffer, indices: &Buffer, dst: &mut Buffer) -> MohuResult<() if indices.len() != dst.len() { return Err(MohuError::ShapeMismatch { expected: indices.shape().to_vec(), - got: dst.shape().to_vec(), + got: dst.shape().to_vec(), }); } - if !dst.is_writeable() { return Err(MohuError::ReadOnly); } + if !dst.is_writeable() { + return Err(MohuError::ReadOnly); + } dst.make_unique()?; - let itemsize = src.dtype().itemsize(); - let src_len = src.len(); - let n_idx = indices.len(); + let itemsize = src.dtype().itemsize(); + let src_len = src.len(); + let n_idx = indices.len(); // Build typed slices from the raw pointers. All three buffers are // exclusively borrowed (src immutably, dst mutably) for this call. // Using &[u8] for src makes it Sync, enabling safe sharing across Rayon threads. - let src_bytes: &[u8] = unsafe { - std::slice::from_raw_parts(src.as_ptr(), src_len * itemsize) - }; - let idx_s: &[i64] = unsafe { - std::slice::from_raw_parts(indices.as_ptr() as *const i64, n_idx) - }; - let dst_b: &mut [u8] = unsafe { - std::slice::from_raw_parts_mut(dst.as_mut_ptr(), n_idx * itemsize) - }; - - idx_s.par_iter().zip(dst_b.par_chunks_mut(itemsize)).for_each(|(&idx, out)| { - let i = if idx < 0 { - (src_len as i64 + idx) as usize - } else { - idx as usize - }; - debug_assert!(i < src_len, "gather: index {i} out of bounds (len {src_len})"); - let i = i.min(src_len.saturating_sub(1)); // clamp in release - out.copy_from_slice(&src_bytes[i * itemsize..(i + 1) * itemsize]); - }); + let src_bytes: &[u8] = unsafe { std::slice::from_raw_parts(src.as_ptr(), src_len * itemsize) }; + let idx_s: &[i64] = + unsafe { std::slice::from_raw_parts(indices.as_ptr() as *const i64, n_idx) }; + let dst_b: &mut [u8] = + unsafe { std::slice::from_raw_parts_mut(dst.as_mut_ptr(), n_idx * itemsize) }; + + idx_s + .par_iter() + .zip(dst_b.par_chunks_mut(itemsize)) + .for_each(|(&idx, out)| { + let i = if idx < 0 { + (src_len as i64 + idx) as usize + } else { + idx as usize + }; + debug_assert!( + i < src_len, + "gather: index {i} out of bounds (len {src_len})" + ); + let i = i.min(src_len.saturating_sub(1)); // clamp in release + out.copy_from_slice(&src_bytes[i * itemsize..(i + 1) * itemsize]); + }); Ok(()) } @@ -715,13 +707,13 @@ pub fn scatter(dst: &mut Buffer, indices: &Buffer, src: &Buffer) -> MohuResult<( if indices.dtype() != DType::I64 { return Err(MohuError::DTypeMismatch { expected: "I64".to_string(), - got: indices.dtype().to_string(), + got: indices.dtype().to_string(), }); } if src.dtype() != dst.dtype() { return Err(MohuError::DTypeMismatch { expected: src.dtype().to_string(), - got: dst.dtype().to_string(), + got: dst.dtype().to_string(), }); } if !src.is_c_contiguous() || !indices.is_c_contiguous() || !dst.is_c_contiguous() { @@ -730,18 +722,20 @@ pub fn scatter(dst: &mut Buffer, indices: &Buffer, src: &Buffer) -> MohuResult<( if indices.len() != src.len() { return Err(MohuError::ShapeMismatch { expected: indices.shape().to_vec(), - got: src.shape().to_vec(), + got: src.shape().to_vec(), }); } - if !dst.is_writeable() { return Err(MohuError::ReadOnly); } + if !dst.is_writeable() { + return Err(MohuError::ReadOnly); + } dst.make_unique()?; - let itemsize = src.dtype().itemsize(); - let dst_len = dst.len(); - let n_src = src.len(); - let idx_ptr = indices.as_ptr() as *const i64; - let src_ptr = src.as_ptr(); - let dst_ptr = unsafe { dst.as_mut_ptr() }; + let itemsize = src.dtype().itemsize(); + let dst_len = dst.len(); + let n_src = src.len(); + let idx_ptr = indices.as_ptr() as *const i64; + let src_ptr = src.as_ptr(); + let dst_ptr = unsafe { dst.as_mut_ptr() }; let idx_s = unsafe { std::slice::from_raw_parts(idx_ptr, n_src) }; let src_b = unsafe { std::slice::from_raw_parts(src_ptr, n_src * itemsize) }; @@ -826,18 +820,20 @@ pub fn abs_copy(src: &Buffer, dst: &mut Buffer) -> MohuResult<()> { } match src.dtype() { // Unsigned and bool: abs is the identity. - DType::Bool | DType::U8 | DType::U16 | DType::U32 | DType::U64 => copy_to_contiguous(src, dst), + DType::Bool | DType::U8 | DType::U16 | DType::U32 | DType::U64 => { + copy_to_contiguous(src, dst) + }, // Signed integers and reals: round-trip through f64. - DType::I8 => abs_via_f64!(i8), - DType::I16 => abs_via_f64!(i16), - DType::I32 => abs_via_f64!(i32), - DType::I64 => abs_via_f64!(i64), - DType::F16 => abs_via_f64!(::half::f16), + DType::I8 => abs_via_f64!(i8), + DType::I16 => abs_via_f64!(i16), + DType::I32 => abs_via_f64!(i32), + DType::I64 => abs_via_f64!(i64), + DType::F16 => abs_via_f64!(::half::f16), DType::BF16 => abs_via_f64!(::half::bf16), - DType::F32 => parallel_map::(src, dst, |x| x.abs()), - DType::F64 => parallel_map::(src, dst, |x| x.abs()), + DType::F32 => parallel_map::(src, dst, |x| x.abs()), + DType::F64 => parallel_map::(src, dst, |x| x.abs()), // Complex: abs each component independently (preserves dtype). - DType::C64 => parallel_map::, Complex, _>(src, dst, |x| { + DType::C64 => parallel_map::, Complex, _>(src, dst, |x| { Complex::new(x.re.abs(), x.im.abs()) }), DType::C128 => parallel_map::, Complex, _>(src, dst, |x| { @@ -868,27 +864,23 @@ pub fn neg_copy(src: &Buffer, dst: &mut Buffer) -> MohuResult<()> { match src.dtype() { // Bool has no meaningful negation; produce a copy. DType::Bool => copy_to_contiguous(src, dst), - DType::I8 => neg_via_i128!(i8), - DType::I16 => neg_via_i128!(i16), - DType::I32 => neg_via_i128!(i32), - DType::I64 => neg_via_i128!(i64), - DType::U8 => parallel_map::(src, dst, |x| x.wrapping_neg()), - DType::U16 => parallel_map::(src, dst, |x| x.wrapping_neg()), - DType::U32 => parallel_map::(src, dst, |x| x.wrapping_neg()), - DType::U64 => parallel_map::(src, dst, |x| x.wrapping_neg()), - DType::F16 => { - parallel_map::<::half::f16, ::half::f16, _>(src, dst, |x| { - ::half::f16::from_f32(-x.to_f32()) - }) - } - DType::BF16 => { - parallel_map::<::half::bf16, ::half::bf16, _>(src, dst, |x| { - ::half::bf16::from_f32(-x.to_f32()) - }) - } - DType::F32 => parallel_map::(src, dst, |x| -x), - DType::F64 => parallel_map::(src, dst, |x| -x), - DType::C64 => parallel_map::, Complex, _>(src, dst, |x| -x), + DType::I8 => neg_via_i128!(i8), + DType::I16 => neg_via_i128!(i16), + DType::I32 => neg_via_i128!(i32), + DType::I64 => neg_via_i128!(i64), + DType::U8 => parallel_map::(src, dst, |x| x.wrapping_neg()), + DType::U16 => parallel_map::(src, dst, |x| x.wrapping_neg()), + DType::U32 => parallel_map::(src, dst, |x| x.wrapping_neg()), + DType::U64 => parallel_map::(src, dst, |x| x.wrapping_neg()), + DType::F16 => parallel_map::<::half::f16, ::half::f16, _>(src, dst, |x| { + ::half::f16::from_f32(-x.to_f32()) + }), + DType::BF16 => parallel_map::<::half::bf16, ::half::bf16, _>(src, dst, |x| { + ::half::bf16::from_f32(-x.to_f32()) + }), + DType::F32 => parallel_map::(src, dst, |x| -x), + DType::F64 => parallel_map::(src, dst, |x| -x), + DType::C64 => parallel_map::, Complex, _>(src, dst, |x| -x), DType::C128 => parallel_map::, Complex, _>(src, dst, |x| -x), } } @@ -903,7 +895,7 @@ pub fn sqrt_copy(src: &Buffer, dst: &mut Buffer) -> MohuResult<()> { DType::F64 => parallel_map::(src, dst, |x| x.sqrt()), other => Err(MohuError::DTypeMismatch { expected: "F32 or F64".to_string(), - got: other.to_string(), + got: other.to_string(), }), } } @@ -916,7 +908,7 @@ pub fn ln_copy(src: &Buffer, dst: &mut Buffer) -> MohuResult<()> { DType::F64 => parallel_map::(src, dst, |x| x.ln()), other => Err(MohuError::DTypeMismatch { expected: "F32 or F64".to_string(), - got: other.to_string(), + got: other.to_string(), }), } } @@ -929,7 +921,7 @@ pub fn exp_copy(src: &Buffer, dst: &mut Buffer) -> MohuResult<()> { DType::F64 => parallel_map::(src, dst, |x| x.exp()), other => Err(MohuError::DTypeMismatch { expected: "F32 or F64".to_string(), - got: other.to_string(), + got: other.to_string(), }), } } @@ -944,48 +936,49 @@ pub fn flip_axis_copy(src: &Buffer, dst: &mut Buffer, axis: usize) -> MohuResult if src.dtype() != dst.dtype() { return Err(MohuError::DTypeMismatch { expected: src.dtype().to_string(), - got: dst.dtype().to_string(), + got: dst.dtype().to_string(), }); } if src.shape() != dst.shape() { return Err(MohuError::ShapeMismatch { expected: src.shape().to_vec(), - got: dst.shape().to_vec(), + got: dst.shape().to_vec(), }); } if axis >= src.ndim() { return Err(MohuError::bug(format!( - "flip_axis_copy: axis {axis} out of bounds for ndim {}", src.ndim() + "flip_axis_copy: axis {axis} out of bounds for ndim {}", + src.ndim() ))); } if !dst.is_c_contiguous() { return Err(MohuError::NonContiguous); } - if !dst.is_writeable() { return Err(MohuError::ReadOnly); } + if !dst.is_writeable() { + return Err(MohuError::ReadOnly); + } dst.make_unique()?; let itemsize = src.dtype().itemsize(); - let src_raw = src.as_ptr(); - let dst_raw = unsafe { dst.as_mut_ptr() }; + let src_raw = src.as_ptr(); + let dst_raw = unsafe { dst.as_mut_ptr() }; use crate::strides::NdIndexIter; // Walk destination in C order; compute the flipped source index. for (dst_flat, mut idx) in NdIndexIter::new(dst.shape()).enumerate() { // Flip the requested axis. - let dim = src.shape()[axis]; + let dim = src.shape()[axis]; idx[axis] = dim - 1 - idx[axis]; - let src_off = src.layout().byte_offset(idx.as_slice()) + let src_off = src + .layout() + .byte_offset(idx.as_slice()) .expect("NdIndexIter always in bounds"); let dst_off = dst_flat * itemsize; unsafe { - std::ptr::copy_nonoverlapping( - src_raw.add(src_off), - dst_raw.add(dst_off), - itemsize, - ); + std::ptr::copy_nonoverlapping(src_raw.add(src_off), dst_raw.add(dst_off), itemsize); } } @@ -1021,38 +1014,36 @@ pub fn sum_all_f64(buf: &Buffer) -> MohuResult { macro_rules! do_sum { ($T:ty) => {{ if buf.is_c_contiguous() { - let s = unsafe { - std::slice::from_raw_parts(buf.as_ptr() as *const $T, buf.len()) - }; - let sum: f64 = s.par_iter().map(|&x| { - num_traits::cast::<$T, f64>(x).unwrap_or(0.0) - }).sum(); + let s = unsafe { std::slice::from_raw_parts(buf.as_ptr() as *const $T, buf.len()) }; + let sum: f64 = s + .par_iter() + .map(|&x| num_traits::cast::<$T, f64>(x).unwrap_or(0.0)) + .sum(); Ok(sum) } else { let c = buf.to_contiguous()?; - let s = unsafe { - std::slice::from_raw_parts(c.as_ptr() as *const $T, c.len()) - }; - let sum: f64 = s.par_iter().map(|&x| { - num_traits::cast::<$T, f64>(x).unwrap_or(0.0) - }).sum(); + let s = unsafe { std::slice::from_raw_parts(c.as_ptr() as *const $T, c.len()) }; + let sum: f64 = s + .par_iter() + .map(|&x| num_traits::cast::<$T, f64>(x).unwrap_or(0.0)) + .sum(); Ok(sum) } }}; } match buf.dtype() { - DType::I8 => do_sum!(i8), - DType::I16 => do_sum!(i16), - DType::I32 => do_sum!(i32), - DType::I64 => do_sum!(i64), - DType::U8 => do_sum!(u8), - DType::U16 => do_sum!(u16), - DType::U32 => do_sum!(u32), - DType::U64 => do_sum!(u64), - DType::F16 => do_sum!(::half::f16), + DType::I8 => do_sum!(i8), + DType::I16 => do_sum!(i16), + DType::I32 => do_sum!(i32), + DType::I64 => do_sum!(i64), + DType::U8 => do_sum!(u8), + DType::U16 => do_sum!(u16), + DType::U32 => do_sum!(u32), + DType::U64 => do_sum!(u64), + DType::F16 => do_sum!(::half::f16), DType::BF16 => do_sum!(::half::bf16), - DType::F32 => do_sum!(f32), - DType::F64 => do_sum!(f64), + DType::F32 => do_sum!(f32), + DType::F64 => do_sum!(f64), // Bool and C64/C128 handled above. _ => unreachable!(), } @@ -1068,7 +1059,9 @@ pub fn min_all_f64(buf: &Buffer) -> MohuResult { }); } if buf.dtype() == DType::Bool { - if buf.is_empty() { return Ok(f64::INFINITY); } + if buf.is_empty() { + return Ok(f64::INFINITY); + } let c; let s: &[u8] = if buf.is_c_contiguous() { unsafe { std::slice::from_raw_parts(buf.as_ptr(), buf.len()) } @@ -1090,28 +1083,30 @@ pub fn min_all_f64(buf: &Buffer) -> MohuResult { c = buf.to_contiguous()?; unsafe { std::slice::from_raw_parts(c.as_ptr() as *const $T, c.len()) } }; - let m = s.par_iter().cloned().reduce_with(|a, b| { - match a.partial_cmp(&b) { + let m = s + .par_iter() + .cloned() + .reduce_with(|a, b| match a.partial_cmp(&b) { Some(std::cmp::Ordering::Less) | Some(std::cmp::Ordering::Equal) => a, _ => b, - } - }).unwrap_or(s[0]); + }) + .unwrap_or(s[0]); Ok(num_traits::cast::<$T, f64>(m).unwrap_or(f64::NAN)) }}; } match buf.dtype() { - DType::I8 => do_min!(i8), - DType::I16 => do_min!(i16), - DType::I32 => do_min!(i32), - DType::I64 => do_min!(i64), - DType::U8 => do_min!(u8), - DType::U16 => do_min!(u16), - DType::U32 => do_min!(u32), - DType::U64 => do_min!(u64), - DType::F16 => do_min!(::half::f16), + DType::I8 => do_min!(i8), + DType::I16 => do_min!(i16), + DType::I32 => do_min!(i32), + DType::I64 => do_min!(i64), + DType::U8 => do_min!(u8), + DType::U16 => do_min!(u16), + DType::U32 => do_min!(u32), + DType::U64 => do_min!(u64), + DType::F16 => do_min!(::half::f16), DType::BF16 => do_min!(::half::bf16), - DType::F32 => do_min!(f32), - DType::F64 => do_min!(f64), + DType::F32 => do_min!(f32), + DType::F64 => do_min!(f64), _ => unreachable!(), } } @@ -1126,7 +1121,9 @@ pub fn max_all_f64(buf: &Buffer) -> MohuResult { }); } if buf.dtype() == DType::Bool { - if buf.is_empty() { return Ok(f64::NEG_INFINITY); } + if buf.is_empty() { + return Ok(f64::NEG_INFINITY); + } let c; let s: &[u8] = if buf.is_c_contiguous() { unsafe { std::slice::from_raw_parts(buf.as_ptr(), buf.len()) } @@ -1148,28 +1145,30 @@ pub fn max_all_f64(buf: &Buffer) -> MohuResult { c = buf.to_contiguous()?; unsafe { std::slice::from_raw_parts(c.as_ptr() as *const $T, c.len()) } }; - let m = s.par_iter().cloned().reduce_with(|a, b| { - match a.partial_cmp(&b) { + let m = s + .par_iter() + .cloned() + .reduce_with(|a, b| match a.partial_cmp(&b) { Some(std::cmp::Ordering::Greater) | Some(std::cmp::Ordering::Equal) => a, _ => b, - } - }).unwrap_or(s[0]); + }) + .unwrap_or(s[0]); Ok(num_traits::cast::<$T, f64>(m).unwrap_or(f64::NAN)) }}; } match buf.dtype() { - DType::I8 => do_max!(i8), - DType::I16 => do_max!(i16), - DType::I32 => do_max!(i32), - DType::I64 => do_max!(i64), - DType::U8 => do_max!(u8), - DType::U16 => do_max!(u16), - DType::U32 => do_max!(u32), - DType::U64 => do_max!(u64), - DType::F16 => do_max!(::half::f16), + DType::I8 => do_max!(i8), + DType::I16 => do_max!(i16), + DType::I32 => do_max!(i32), + DType::I64 => do_max!(i64), + DType::U8 => do_max!(u8), + DType::U16 => do_max!(u16), + DType::U32 => do_max!(u32), + DType::U64 => do_max!(u64), + DType::F16 => do_max!(::half::f16), DType::BF16 => do_max!(::half::bf16), - DType::F32 => do_max!(f32), - DType::F64 => do_max!(f64), + DType::F32 => do_max!(f32), + DType::F64 => do_max!(f64), _ => unreachable!(), } } @@ -1184,7 +1183,9 @@ pub fn argmin_flat(buf: &Buffer) -> MohuResult { }); } if buf.dtype() == DType::Bool { - if buf.is_empty() { return Ok(0); } + if buf.is_empty() { + return Ok(0); + } let c; let s: &[u8] = if buf.is_c_contiguous() { unsafe { std::slice::from_raw_parts(buf.as_ptr(), buf.len()) } @@ -1192,14 +1193,18 @@ pub fn argmin_flat(buf: &Buffer) -> MohuResult { c = buf.to_contiguous()?; unsafe { std::slice::from_raw_parts(c.as_ptr(), c.len()) } }; - return Ok(s.iter().enumerate() + return Ok(s + .iter() + .enumerate() .min_by_key(|&(_, &v)| v) .map(|(i, _)| i) .unwrap_or(0)); } macro_rules! do_argmin { ($T:ty) => {{ - if buf.is_empty() { return Ok(0); } + if buf.is_empty() { + return Ok(0); + } let c; let s: &[$T] = if buf.is_c_contiguous() { unsafe { std::slice::from_raw_parts(buf.as_ptr() as *const $T, buf.len()) } @@ -1207,27 +1212,30 @@ pub fn argmin_flat(buf: &Buffer) -> MohuResult { c = buf.to_contiguous()?; unsafe { std::slice::from_raw_parts(c.as_ptr() as *const $T, c.len()) } }; - Ok(s.par_iter().enumerate() - .reduce_with(|(ai, av), (bi, bv)| { - if av <= bv { (ai, av) } else { (bi, bv) } - }) + Ok(s.par_iter() + .enumerate() + .reduce_with( + |(ai, av), (bi, bv)| { + if av <= bv { (ai, av) } else { (bi, bv) } + }, + ) .map(|(i, _)| i) .unwrap_or(0)) }}; } match buf.dtype() { - DType::I8 => do_argmin!(i8), - DType::I16 => do_argmin!(i16), - DType::I32 => do_argmin!(i32), - DType::I64 => do_argmin!(i64), - DType::U8 => do_argmin!(u8), - DType::U16 => do_argmin!(u16), - DType::U32 => do_argmin!(u32), - DType::U64 => do_argmin!(u64), - DType::F16 => do_argmin!(::half::f16), + DType::I8 => do_argmin!(i8), + DType::I16 => do_argmin!(i16), + DType::I32 => do_argmin!(i32), + DType::I64 => do_argmin!(i64), + DType::U8 => do_argmin!(u8), + DType::U16 => do_argmin!(u16), + DType::U32 => do_argmin!(u32), + DType::U64 => do_argmin!(u64), + DType::F16 => do_argmin!(::half::f16), DType::BF16 => do_argmin!(::half::bf16), - DType::F32 => do_argmin!(f32), - DType::F64 => do_argmin!(f64), + DType::F32 => do_argmin!(f32), + DType::F64 => do_argmin!(f64), _ => unreachable!(), } } @@ -1242,7 +1250,9 @@ pub fn argmax_flat(buf: &Buffer) -> MohuResult { }); } if buf.dtype() == DType::Bool { - if buf.is_empty() { return Ok(0); } + if buf.is_empty() { + return Ok(0); + } let c; let s: &[u8] = if buf.is_c_contiguous() { unsafe { std::slice::from_raw_parts(buf.as_ptr(), buf.len()) } @@ -1250,14 +1260,18 @@ pub fn argmax_flat(buf: &Buffer) -> MohuResult { c = buf.to_contiguous()?; unsafe { std::slice::from_raw_parts(c.as_ptr(), c.len()) } }; - return Ok(s.iter().enumerate() + return Ok(s + .iter() + .enumerate() .max_by_key(|&(_, &v)| v) .map(|(i, _)| i) .unwrap_or(0)); } macro_rules! do_argmax { ($T:ty) => {{ - if buf.is_empty() { return Ok(0); } + if buf.is_empty() { + return Ok(0); + } let c; let s: &[$T] = if buf.is_c_contiguous() { unsafe { std::slice::from_raw_parts(buf.as_ptr() as *const $T, buf.len()) } @@ -1265,27 +1279,30 @@ pub fn argmax_flat(buf: &Buffer) -> MohuResult { c = buf.to_contiguous()?; unsafe { std::slice::from_raw_parts(c.as_ptr() as *const $T, c.len()) } }; - Ok(s.par_iter().enumerate() - .reduce_with(|(ai, av), (bi, bv)| { - if av >= bv { (ai, av) } else { (bi, bv) } - }) + Ok(s.par_iter() + .enumerate() + .reduce_with( + |(ai, av), (bi, bv)| { + if av >= bv { (ai, av) } else { (bi, bv) } + }, + ) .map(|(i, _)| i) .unwrap_or(0)) }}; } match buf.dtype() { - DType::I8 => do_argmax!(i8), - DType::I16 => do_argmax!(i16), - DType::I32 => do_argmax!(i32), - DType::I64 => do_argmax!(i64), - DType::U8 => do_argmax!(u8), - DType::U16 => do_argmax!(u16), - DType::U32 => do_argmax!(u32), - DType::U64 => do_argmax!(u64), - DType::F16 => do_argmax!(::half::f16), + DType::I8 => do_argmax!(i8), + DType::I16 => do_argmax!(i16), + DType::I32 => do_argmax!(i32), + DType::I64 => do_argmax!(i64), + DType::U8 => do_argmax!(u8), + DType::U16 => do_argmax!(u16), + DType::U32 => do_argmax!(u32), + DType::U64 => do_argmax!(u64), + DType::F16 => do_argmax!(::half::f16), DType::BF16 => do_argmax!(::half::bf16), - DType::F32 => do_argmax!(f32), - DType::F64 => do_argmax!(f64), + DType::F32 => do_argmax!(f32), + DType::F64 => do_argmax!(f64), _ => unreachable!(), } } @@ -1296,17 +1313,21 @@ pub fn argmax_flat(buf: &Buffer) -> MohuResult { /// /// Bypasses the CPU cache — achieves peak DRAM write bandwidth for buffers /// > a few MiB where the data will not be immediately re-read. -/// Falls back to the standard Rayon fill on non-x86_64 platforms. +/// > Falls back to the standard Rayon fill on non-x86_64 platforms. pub fn fill_nontemporal_f32_buf(buf: &mut Buffer, value: f32) -> MohuResult<()> { use mohu_dtype::DType; if buf.dtype() != DType::F32 { return Err(MohuError::DTypeMismatch { expected: "F32".to_string(), - got: buf.dtype().to_string(), + got: buf.dtype().to_string(), }); } - if !buf.is_c_contiguous() { return Err(MohuError::NonContiguous); } - if !buf.is_writeable() { return Err(MohuError::ReadOnly); } + if !buf.is_c_contiguous() { + return Err(MohuError::NonContiguous); + } + if !buf.is_writeable() { + return Err(MohuError::ReadOnly); + } buf.make_unique()?; let len = buf.len(); @@ -1316,15 +1337,17 @@ pub fn fill_nontemporal_f32_buf(buf: &mut Buffer, value: f32) -> MohuResult<()> { // Fast path: non-temporal stores for aligned 32-byte regions. // Align the pointer upward to 32 bytes; fill the prefix with regular stores. - let addr = ptr as usize; + let addr = ptr as usize; let align_off = (32 - addr % 32) % 32 / std::mem::size_of::(); - let prefix = align_off.min(len); + let prefix = align_off.min(len); for i in 0..prefix { - unsafe { ptr.add(i).write(value); } + unsafe { + ptr.add(i).write(value); + } } - let aligned_ptr = unsafe { ptr.add(prefix) }; - let aligned_len = len - prefix; - let nt_len = aligned_len / 8 * 8; // round down to multiple of 8 + let aligned_ptr = unsafe { ptr.add(prefix) }; + let aligned_len = len - prefix; + let nt_len = aligned_len / 8 * 8; // round down to multiple of 8 if nt_len > 0 { // SAFETY: aligned_ptr is 32-byte aligned, nt_len is multiple of 8. @@ -1339,9 +1362,11 @@ pub fn fill_nontemporal_f32_buf(buf: &mut Buffer, value: f32) -> MohuResult<()> } // Scalar tail for i in (prefix + nt_len)..len { - unsafe { ptr.add(i).write(value); } + unsafe { + ptr.add(i).write(value); + } } - return Ok(()); + Ok(()) } #[cfg(not(target_arch = "x86_64"))] @@ -1356,12 +1381,7 @@ pub fn fill_nontemporal_f32_buf(buf: &mut Buffer, value: f32) -> MohuResult<()> /// Applies `f(a[i], b[i]) -> T` over all elements of two same-shape C-contiguous /// buffers, writing results into `dst`. Generalises element-wise arithmetic. -pub fn parallel_zip( - a: &Buffer, - b: &Buffer, - dst: &mut Buffer, - f: F, -) -> MohuResult<()> +pub fn parallel_zip(a: &Buffer, b: &Buffer, dst: &mut Buffer, f: F) -> MohuResult<()> where S: Scalar + Copy + Send + Sync, D: Scalar + Copy + Send + Sync, @@ -1370,26 +1390,28 @@ where if S::DTYPE != a.dtype() || S::DTYPE != b.dtype() { return Err(MohuError::DTypeMismatch { expected: S::DTYPE.to_string(), - got: a.dtype().to_string(), + got: a.dtype().to_string(), }); } if a.len() != b.len() || a.len() != dst.len() { return Err(MohuError::ShapeMismatch { expected: a.shape().to_vec(), - got: b.shape().to_vec(), + got: b.shape().to_vec(), }); } if !a.is_c_contiguous() || !b.is_c_contiguous() || !dst.is_c_contiguous() { return Err(MohuError::NonContiguous); } - if !dst.is_writeable() { return Err(MohuError::ReadOnly); } + if !dst.is_writeable() { + return Err(MohuError::ReadOnly); + } dst.make_unique()?; - let len = a.len(); - let a_s = unsafe { std::slice::from_raw_parts(a.as_ptr() as *const S, len) }; - let b_s = unsafe { std::slice::from_raw_parts(b.as_ptr() as *const S, len) }; + let len = a.len(); + let a_s = unsafe { std::slice::from_raw_parts(a.as_ptr() as *const S, len) }; + let b_s = unsafe { std::slice::from_raw_parts(b.as_ptr() as *const S, len) }; let d_ptr = unsafe { dst.as_mut_ptr() } as *mut D; - let d_s = unsafe { std::slice::from_raw_parts_mut(d_ptr, len) }; + let d_s = unsafe { std::slice::from_raw_parts_mut(d_ptr, len) }; a_s.par_iter() .zip(b_s.par_iter()) diff --git a/crates/mohu-buffer/src/pool.rs b/crates/mohu-buffer/src/pool.rs index d9ce290..9240757 100644 --- a/crates/mohu-buffer/src/pool.rs +++ b/crates/mohu-buffer/src/pool.rs @@ -60,11 +60,11 @@ thread_local! { #[derive(Debug, Clone, Copy, Default)] pub struct TlStats { /// Number of TL cache hits. - pub hits: u64, + pub hits: u64, /// Number of TL cache misses (fell through to global pool). - pub misses: u64, + pub misses: u64, /// Number of handles returned to the TL cache. - pub returns: u64, + pub returns: u64, /// Total bytes currently cached in this thread's TL cache. pub cached_bytes: usize, } @@ -80,20 +80,25 @@ impl TlStats { #[inline] fn size_class(n: usize) -> usize { - if n == 0 { return 1; } + if n == 0 { + return 1; + } n.next_power_of_two() } // ─── PoolBucket ─────────────────────────────────────────────────────────────── struct PoolBucket { - handles: Vec, + handles: Vec, cached_bytes: usize, } impl PoolBucket { fn new() -> Self { - Self { handles: Vec::new(), cached_bytes: 0 } + Self { + handles: Vec::new(), + cached_bytes: 0, + } } fn push(&mut self, handle: AllocHandle) { @@ -113,7 +118,9 @@ impl PoolBucket { self.cached_bytes = 0; } - fn len(&self) -> usize { self.handles.len() } + fn len(&self) -> usize { + self.handles.len() + } } // ─── SizeClassStats ─────────────────────────────────────────────────────────── @@ -122,11 +129,11 @@ impl PoolBucket { #[derive(Debug, Clone, Copy)] pub struct SizeClassStats { /// Size class in bytes (always a power of two). - pub size_class: usize, + pub size_class: usize, /// Number of cached handles. pub cached_handles: usize, /// Total cached bytes in this class. - pub cached_bytes: usize, + pub cached_bytes: usize, } // ─── BufferPool ─────────────────────────────────────────────────────────────── @@ -138,25 +145,25 @@ pub struct SizeClassStats { /// Use [`acquire`](Self::acquire) / [`release`](Self::release) for the global /// pool directly (no TL cache involvement). pub struct BufferPool { - inner: Mutex, + inner: Mutex, max_cached_bytes: usize, } struct PoolInner { - buckets: BTreeMap, + buckets: BTreeMap, cached_bytes: usize, - hit_count: u64, - miss_count: u64, + hit_count: u64, + miss_count: u64, return_count: u64, } impl PoolInner { fn new() -> Self { Self { - buckets: BTreeMap::new(), + buckets: BTreeMap::new(), cached_bytes: 0, - hit_count: 0, - miss_count: 0, + hit_count: 0, + miss_count: 0, return_count: 0, } } @@ -166,19 +173,19 @@ impl PoolInner { #[derive(Debug, Clone, Copy)] pub struct PoolStats { /// Total bytes currently cached in the pool. - pub cached_bytes: usize, + pub cached_bytes: usize, /// Total number of cached allocation handles. pub cached_blocks: usize, /// Number of successful acquisitions from the cache. - pub hit_count: u64, + pub hit_count: u64, /// Number of acquisitions that required a new allocation. - pub miss_count: u64, + pub miss_count: u64, /// Number of handles returned to the pool. - pub return_count: u64, + pub return_count: u64, /// Cache hit rate as a fraction in `[0.0, 1.0]`. - pub hit_rate: f64, + pub hit_rate: f64, /// Number of distinct active size classes. - pub size_classes: usize, + pub size_classes: usize, } impl BufferPool { @@ -218,7 +225,9 @@ impl BufferPool { /// Returns an `AllocHandle` to the **global** pool. pub fn release(&self, handle: AllocHandle) { - if handle.is_empty() { return; } + if handle.is_empty() { + return; + } let class = size_class(handle.len()); let mut inner = self.lock(); if inner.cached_bytes + handle.len() > self.max_cached_bytes { @@ -229,7 +238,8 @@ impl BufferPool { let handle_len = handle.len(); inner.cached_bytes += handle_len; inner.return_count += 1; - inner.buckets + inner + .buckets .entry(class) .or_insert_with(PoolBucket::new) .push(handle); @@ -252,12 +262,10 @@ impl BufferPool { if class <= TL_MAX_BYTES { let from_tl = TL_CACHE.with(|cache| { let mut c = cache.borrow_mut(); - c.iter() - .position(|(k, _)| *k == class) - .map(|pos| { - let (_, handle) = c.swap_remove(pos); - handle - }) + c.iter().position(|(k, _)| *k == class).map(|pos| { + let (_, handle) = c.swap_remove(pos); + handle + }) }); if let Some(handle) = from_tl { TL_STATS.with(|s| { @@ -279,7 +287,9 @@ impl BufferPool { /// /// Optionally poisons the handle before caching (debug builds only). pub fn fast_release(&self, mut handle: AllocHandle) { - if handle.is_empty() { return; } + if handle.is_empty() { + return; + } handle.poison(); // no-op in release builds let class = size_class(handle.len()); @@ -290,11 +300,7 @@ impl BufferPool { let c = cache.borrow_mut(); // Compute current TL cached bytes let cur_bytes: usize = c.iter().map(|(_, h)| h.len()).sum(); - if c.len() < TL_SLOTS && cur_bytes + handle.len() <= TL_MAX_BYTES { - true - } else { - false - } + c.len() < TL_SLOTS && cur_bytes + handle.len() <= TL_MAX_BYTES }); if accepted { let hlen = handle.len(); @@ -343,7 +349,9 @@ impl BufferPool { /// ``` pub fn warm(&self, sizes: &[usize], count_per_size: usize) -> MohuResult<()> { for &size in sizes { - if size == 0 { continue; } + if size == 0 { + continue; + } let class = size_class(size); // Allocate outside the lock. let mut handles = Vec::with_capacity(count_per_size); @@ -359,7 +367,8 @@ impl BufferPool { let hlen = h.len(); inner.cached_bytes += hlen; inner.return_count += 1; - inner.buckets + inner + .buckets .entry(class) .or_insert_with(PoolBucket::new) .push(h); @@ -421,11 +430,11 @@ impl BufferPool { let inner = self.lock(); let total_calls = inner.hit_count + inner.miss_count; PoolStats { - cached_bytes: inner.cached_bytes, + cached_bytes: inner.cached_bytes, cached_blocks: inner.buckets.values().map(|b| b.len()).sum(), - hit_count: inner.hit_count, - miss_count: inner.miss_count, - return_count: inner.return_count, + hit_count: inner.hit_count, + miss_count: inner.miss_count, + return_count: inner.return_count, hit_rate: if total_calls == 0 { 0.0 } else { @@ -438,11 +447,15 @@ impl BufferPool { /// Returns per-size-class breakdown of cached blocks. pub fn size_class_stats(&self) -> Vec { let inner = self.lock(); - inner.buckets.iter().map(|(&class, b)| SizeClassStats { - size_class: class, - cached_handles: b.len(), - cached_bytes: b.cached_bytes, - }).collect() + inner + .buckets + .iter() + .map(|(&class, b)| SizeClassStats { + size_class: class, + cached_handles: b.len(), + cached_bytes: b.cached_bytes, + }) + .collect() } /// Current number of cached bytes (global pool only). @@ -471,11 +484,11 @@ impl std::fmt::Debug for BufferPool { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let stats = self.stats(); f.debug_struct("BufferPool") - .field("cached_bytes", &stats.cached_bytes) + .field("cached_bytes", &stats.cached_bytes) .field("cached_blocks", &stats.cached_blocks) - .field("size_classes", &stats.size_classes) - .field("max_bytes", &self.max_cached_bytes) - .field("hit_rate", &format_args!("{:.1}%", stats.hit_rate * 100.0)) + .field("size_classes", &stats.size_classes) + .field("max_bytes", &self.max_cached_bytes) + .field("hit_rate", &format_args!("{:.1}%", stats.hit_rate * 100.0)) .finish() } } diff --git a/crates/mohu-buffer/src/strides.rs b/crates/mohu-buffer/src/strides.rs index 3c51571..8c3cb7d 100644 --- a/crates/mohu-buffer/src/strides.rs +++ b/crates/mohu-buffer/src/strides.rs @@ -70,11 +70,13 @@ pub fn f_strides(shape: &[usize], itemsize: usize) -> StrideVec { pub fn shape_size(shape: &[usize]) -> MohuResult { let mut total: usize = 1; for &dim in shape { - total = total.checked_mul(dim).ok_or(MohuError::ShapeOverflow { - max: usize::MAX, - })?; + total = total + .checked_mul(dim) + .ok_or(MohuError::ShapeOverflow { max: usize::MAX })?; if total > isize::MAX as usize { - return Err(MohuError::ShapeOverflow { max: isize::MAX as usize }); + return Err(MohuError::ShapeOverflow { + max: isize::MAX as usize, + }); } } Ok(total) @@ -104,9 +106,9 @@ pub fn contiguous_nbytes(shape: &[usize], itemsize: usize) -> MohuResult /// /// Returns `Err(BroadcastError)` if broadcasting is impossible. pub fn broadcast_strides( - src_shape: &[usize], + src_shape: &[usize], src_strides: &[isize], - tgt_shape: &[usize], + tgt_shape: &[usize], ) -> MohuResult { let src_ndim = src_shape.len(); let tgt_ndim = tgt_shape.len(); @@ -127,9 +129,7 @@ pub fn broadcast_strides( } // Align and validate the trailing dimensions. - for (axis, (&s_dim, &s_stride)) in - src_shape.iter().zip(src_strides.iter()).enumerate() - { + for (axis, (&s_dim, &s_stride)) in src_shape.iter().zip(src_strides.iter()).enumerate() { let t_dim = tgt_shape[offset + axis]; if s_dim == t_dim { out.push(s_stride); @@ -156,7 +156,7 @@ pub fn unravel_index(flat: usize, shape: &[usize]) -> MohuResult { if flat >= size && size > 0 { return Err(MohuError::IndexOutOfBounds { index: flat as i64, - axis: 0, + axis: 0, size, }); } @@ -177,7 +177,7 @@ pub fn ravel_multi_index(indices: &[usize], shape: &[usize]) -> MohuResult MohuResult isize { +pub fn byte_offset(indices: &[usize], strides: &[isize], base_offset: usize) -> isize { let mut off: isize = base_offset as isize; for (&idx, &stride) in indices.iter().zip(strides.iter()) { off += idx as isize * stride; @@ -226,11 +222,11 @@ pub fn byte_offset( /// ``` #[derive(Debug, Clone)] pub struct NdIndexIter { - shape: ShapeVec, + shape: ShapeVec, current: ShapeVec, - done: bool, - count: usize, - total: usize, + done: bool, + count: usize, + total: usize, } impl NdIndexIter { @@ -251,7 +247,9 @@ impl NdIndexIter { } /// Returns the total number of indices this iterator will yield. - pub fn total(&self) -> usize { self.total } + pub fn total(&self) -> usize { + self.total + } fn advance(&mut self) { let ndim = self.shape.len(); @@ -309,8 +307,8 @@ impl ExactSizeIterator for NdIndexIter {} /// without an allocation. #[derive(Debug, Clone)] pub struct StridedByteIter { - nd_iter: NdIndexIter, - strides: StrideVec, + nd_iter: NdIndexIter, + strides: StrideVec, base_offset: usize, } @@ -350,10 +348,10 @@ impl ExactSizeIterator for StridedByteIter {} /// Broadcast strides (stride = 0) are excluded from the overlap check because /// they represent read-only virtual replication of a single element. pub fn validate_strides( - shape: &[usize], - strides: &[isize], - itemsize: usize, - mutable: bool, + shape: &[usize], + strides: &[isize], + itemsize: usize, + mutable: bool, ) -> MohuResult<()> { for (axis, (&stride, &dim)) in strides.iter().zip(shape.iter()).enumerate() { if dim <= 1 { @@ -380,11 +378,7 @@ pub fn validate_strides( /// Returns `Err(OverlappingStrides)` if the stride+shape combination would /// cause two distinct elements to share a byte address. -fn check_no_overlap( - shape: &[usize], - strides: &[isize], - itemsize: usize, -) -> MohuResult<()> { +fn check_no_overlap(shape: &[usize], strides: &[isize], itemsize: usize) -> MohuResult<()> { // Fast path: only one axis has stride != 0 — can't overlap. let non_broadcast: Vec<_> = strides .iter() @@ -412,8 +406,8 @@ fn check_no_overlap( // The span of the current axis must not exceed the spacing of the prev. if curr_stride.unsigned_abs() < prev_stride.unsigned_abs() * prev_dim { return Err(MohuError::OverlappingStrides { - shape: shape.to_vec(), - strides: strides.to_vec(), + shape: shape.to_vec(), + strides: strides.to_vec(), element_size: itemsize, }); } diff --git a/crates/mohu-buffer/src/view.rs b/crates/mohu-buffer/src/view.rs index 27d26ea..e28574e 100644 --- a/crates/mohu-buffer/src/view.rs +++ b/crates/mohu-buffer/src/view.rs @@ -16,10 +16,7 @@ use std::marker::PhantomData; use mohu_dtype::{dtype::DType, scalar::Scalar}; use mohu_error::{MohuError, MohuResult}; -use crate::{ - buffer::Buffer, - strides::StridedByteIter, -}; +use crate::{buffer::Buffer, strides::StridedByteIter}; // ─── BufferView<'buf, T> ────────────────────────────────────────────────────── @@ -30,7 +27,7 @@ use crate::{ /// The view borrows `'buf` from the `Buffer`. The `Buffer` (and its backing /// `Arc`) must outlive this view. pub struct BufferView<'buf, T: Scalar> { - buf: &'buf Buffer, + buf: &'buf Buffer, _marker: PhantomData<&'buf T>, } @@ -42,28 +39,52 @@ impl<'buf, T: Scalar> BufferView<'buf, T> { if T::DTYPE != buf.dtype() { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: buf.dtype().to_string(), + got: buf.dtype().to_string(), }); } - Ok(Self { buf, _marker: PhantomData }) + Ok(Self { + buf, + _marker: PhantomData, + }) } // ─── Properties ─────────────────────────────────────────────────────────── /// Returns the element data type. - #[inline] pub fn dtype(&self) -> DType { self.buf.dtype() } + #[inline] + pub fn dtype(&self) -> DType { + self.buf.dtype() + } /// Returns the number of dimensions. - #[inline] pub fn ndim(&self) -> usize { self.buf.ndim() } + #[inline] + pub fn ndim(&self) -> usize { + self.buf.ndim() + } /// Returns the shape as a slice of dimension sizes. - #[inline] pub fn shape(&self) -> &[usize] { self.buf.shape() } + #[inline] + pub fn shape(&self) -> &[usize] { + self.buf.shape() + } /// Returns the byte strides. - #[inline] pub fn strides(&self) -> &[isize] { self.buf.strides() } + #[inline] + pub fn strides(&self) -> &[isize] { + self.buf.strides() + } /// Returns the total number of elements. - #[inline] pub fn len(&self) -> usize { self.buf.len() } + #[inline] + pub fn len(&self) -> usize { + self.buf.len() + } /// Returns `true` if the buffer has zero elements. - #[inline] pub fn is_empty(&self) -> bool { self.buf.is_empty() } + #[inline] + pub fn is_empty(&self) -> bool { + self.buf.is_empty() + } /// Returns `true` if the buffer is C-contiguous. - #[inline] pub fn is_c_contiguous(&self) -> bool { self.buf.is_c_contiguous() } + #[inline] + pub fn is_c_contiguous(&self) -> bool { + self.buf.is_c_contiguous() + } // ─── Slice access (contiguous only) ─────────────────────────────────────── @@ -106,12 +127,8 @@ impl<'buf, T: Scalar> BufferView<'buf, T> { /// For non-contiguous arrays it uses stride-based byte-offset iteration. pub fn iter(&self) -> impl Iterator + '_ { ViewIter { - raw_ptr: self.buf.as_ptr(), - iter: StridedByteIter::new( - self.buf.shape(), - self.buf.strides(), - self.buf.offset(), - ), + raw_ptr: self.buf.as_ptr(), + iter: StridedByteIter::new(self.buf.shape(), self.buf.strides(), self.buf.offset()), _phantom: PhantomData, } } @@ -125,7 +142,7 @@ impl<'buf, T: Scalar> BufferView<'buf, T> { /// [`Buffer::make_unique`] — if the backing bytes are shared, they are /// copied first. pub struct BufferViewMut<'buf, T: Scalar> { - buf: &'buf mut Buffer, + buf: &'buf mut Buffer, _marker: PhantomData<&'buf mut T>, } @@ -138,28 +155,46 @@ impl<'buf, T: Scalar> BufferViewMut<'buf, T> { if T::DTYPE != buf.dtype() { return Err(MohuError::DTypeMismatch { expected: T::DTYPE.to_string(), - got: buf.dtype().to_string(), + got: buf.dtype().to_string(), }); } if !buf.is_writeable() { return Err(MohuError::ReadOnly); } buf.make_unique()?; - Ok(Self { buf, _marker: PhantomData }) + Ok(Self { + buf, + _marker: PhantomData, + }) } // ─── Properties ─────────────────────────────────────────────────────────── /// Returns the element data type. - #[inline] pub fn dtype(&self) -> DType { self.buf.dtype() } + #[inline] + pub fn dtype(&self) -> DType { + self.buf.dtype() + } /// Returns the number of dimensions. - #[inline] pub fn ndim(&self) -> usize { self.buf.ndim() } + #[inline] + pub fn ndim(&self) -> usize { + self.buf.ndim() + } /// Returns the shape as a slice of dimension sizes. - #[inline] pub fn shape(&self) -> &[usize] { self.buf.shape() } + #[inline] + pub fn shape(&self) -> &[usize] { + self.buf.shape() + } /// Returns the total number of elements. - #[inline] pub fn len(&self) -> usize { self.buf.len() } + #[inline] + pub fn len(&self) -> usize { + self.buf.len() + } /// Returns `true` if the buffer has zero elements. - #[inline] pub fn is_empty(&self) -> bool { self.buf.is_empty() } + #[inline] + pub fn is_empty(&self) -> bool { + self.buf.is_empty() + } // ─── Slice access ───────────────────────────────────────────────────────── @@ -200,9 +235,9 @@ impl<'buf, T: Scalar> BufferViewMut<'buf, T> { /// Returns an iterator over mutable references to all elements in C order. pub fn iter_mut(&mut self) -> impl Iterator + '_ { let raw_ptr = unsafe { self.buf.as_mut_ptr() }; - let shape = self.buf.shape().to_vec(); + let shape = self.buf.shape().to_vec(); let strides = self.buf.strides().to_vec(); - let offset = self.buf.offset(); + let offset = self.buf.offset(); ViewIterMut { raw_ptr, iter: StridedByteIter::new(&shape, &strides, offset), @@ -214,8 +249,8 @@ impl<'buf, T: Scalar> BufferViewMut<'buf, T> { // ─── ViewIter ──────────────────────────────────────────────────────────────── struct ViewIter<'a, T> { - raw_ptr: *const u8, - iter: StridedByteIter, + raw_ptr: *const u8, + iter: StridedByteIter, _phantom: PhantomData<&'a T>, } @@ -236,8 +271,8 @@ impl<'a, T: Scalar> Iterator for ViewIter<'a, T> { // ─── ViewIterMut ───────────────────────────────────────────────────────────── struct ViewIterMut<'a, T> { - raw_ptr: *mut u8, - iter: StridedByteIter, + raw_ptr: *mut u8, + iter: StridedByteIter, _phantom: PhantomData<&'a mut T>, } diff --git a/crates/mohu-buffer/tests/integration.rs b/crates/mohu-buffer/tests/integration.rs index 720c5a9..5cddb77 100644 --- a/crates/mohu-buffer/tests/integration.rs +++ b/crates/mohu-buffer/tests/integration.rs @@ -1,9 +1,8 @@ //! Integration tests for mohu-buffer — exercises every major subsystem. use mohu_buffer::{ - Buffer, Order, SliceArg, - ops, GLOBAL_POOL, - strides::{c_strides, f_strides, broadcast_strides, NdIndexIter}, + Buffer, GLOBAL_POOL, Order, SliceArg, ops, + strides::{NdIndexIter, broadcast_strides, c_strides, f_strides}, }; use mohu_dtype::{DType, promote::CastMode}; @@ -28,7 +27,12 @@ fn full_fills_bytes() { // full() takes raw bytes — pass f32 3.14 as bytes let fill: f32 = 3.14; let buf = Buffer::full(DType::F32, &[5, 5], &fill.to_le_bytes()).unwrap(); - assert!(buf.as_slice::().unwrap().iter().all(|&x| (x - 3.14_f32).abs() < 1e-6)); + assert!( + buf.as_slice::() + .unwrap() + .iter() + .all(|&x| (x - 3.14_f32).abs() < 1e-6) + ); } // ── 2. from_slice + reshape + get/set ──────────────────────────────────────── @@ -81,7 +85,10 @@ fn transpose_2d_shape_and_values() { #[test] fn permute_3d_shape() { let data: Vec = (0..24).map(|x| x as f32).collect(); - let buf = Buffer::from_slice(&data).unwrap().reshape(&[2, 3, 4]).unwrap(); + let buf = Buffer::from_slice(&data) + .unwrap() + .reshape(&[2, 3, 4]) + .unwrap(); let p = buf.permute(&[2, 0, 1]).unwrap(); assert_eq!(p.shape(), &[4, 2, 3]); } @@ -92,7 +99,16 @@ fn permute_3d_shape() { fn slice_axis_rows() { let data: Vec = (0..12).map(|x| x as f64).collect(); let buf = Buffer::from_slice(&data).unwrap().reshape(&[4, 3]).unwrap(); - let s = buf.slice_axis(0, SliceArg { start: Some(1), stop: Some(3), step: Some(1) }).unwrap(); + let s = buf + .slice_axis( + 0, + SliceArg { + start: Some(1), + stop: Some(3), + step: Some(1), + }, + ) + .unwrap(); assert_eq!(s.shape(), &[2, 3]); // row 1 of original starts at index 3 → [3, 4, 5] assert_eq!(s.get::(&[0, 0]).unwrap(), 3.0_f64); @@ -104,7 +120,16 @@ fn slice_axis_with_step() { let data: Vec = (0..10).collect(); let buf = Buffer::from_slice(&data).unwrap(); // every other element: 0, 2, 4, 6, 8 - let s = buf.slice_axis(0, SliceArg { start: Some(0), stop: Some(10), step: Some(2) }).unwrap(); + let s = buf + .slice_axis( + 0, + SliceArg { + start: Some(0), + stop: Some(10), + step: Some(2), + }, + ) + .unwrap(); assert_eq!(s.shape(), &[5]); assert_eq!(s.get::(&[2]).unwrap(), 4); } @@ -208,7 +233,11 @@ fn fill_zero_clears_values() { fn copy_to_contiguous_from_transposed() { // Transposed 3×3: source is non-contiguous let data: Vec = (0..9).map(|x| x as f64).collect(); - let src = Buffer::from_slice(&data).unwrap().reshape(&[3, 3]).unwrap().transpose(); + let src = Buffer::from_slice(&data) + .unwrap() + .reshape(&[3, 3]) + .unwrap() + .transpose(); assert!(!src.is_c_contiguous()); let mut dst = Buffer::alloc(DType::F64, &[3, 3], Order::C).unwrap(); ops::copy_to_contiguous(&src, &mut dst).unwrap(); @@ -261,11 +290,11 @@ fn f_strides_correct() { #[test] fn broadcast_strides_zero_for_size_one_axes() { - let src_shape = [1usize, 4]; + let src_shape = [1usize, 4]; let src_strides = c_strides(&src_shape, 4); - let tgt_shape = [3usize, 4]; + let tgt_shape = [3usize, 4]; let bs = broadcast_strides(&src_shape, &src_strides, &tgt_shape).unwrap(); - assert_eq!(bs[0], 0); // size-1 axis → 0-stride + assert_eq!(bs[0], 0); // size-1 axis → 0-stride assert_ne!(bs[1], 0); } diff --git a/crates/mohu-dtype/examples/dtype_basics.rs b/crates/mohu-dtype/examples/dtype_basics.rs index 3263aba..8415b1d 100644 --- a/crates/mohu-dtype/examples/dtype_basics.rs +++ b/crates/mohu-dtype/examples/dtype_basics.rs @@ -9,31 +9,37 @@ // np.finfo(np.float32) // np.iinfo(np.int32) -use mohu_dtype::{ - DType, FloatInfo, IntInfo, ALL_DTYPES, -}; +use mohu_dtype::{ALL_DTYPES, DType, FloatInfo, IntInfo}; fn main() { // ── DType construction and display ───────────────────────────────────── // NumPy: np.dtype('float32') let f32_dt = DType::F32; - println!("DType: {f32_dt}"); // "float32" + println!("DType: {f32_dt}"); // "float32" println!(" itemsize: {} bytes", f32_dt.itemsize()); // 4 - println!(" alignment: {} bytes", f32_dt.alignment());// 4 - println!(" bit_width: {} bits", f32_dt.bit_width()); // 32 + println!(" alignment: {} bytes", f32_dt.alignment()); // 4 + println!(" bit_width: {} bits", f32_dt.bit_width()); // 32 // ── Parsing from NumPy-compatible strings ────────────────────────────── // NumPy: np.dtype('int64'), np.dtype('f4'), np.dtype('complex128') let dt1 = DType::from_str("int64").unwrap(); - let dt2 = DType::from_str("f4").unwrap(); // shorthand for float32 + let dt2 = DType::from_str("f4").unwrap(); // shorthand for float32 let dt3 = DType::from_str("complex128").unwrap(); println!("\nParsed dtypes: {dt1}, {dt2}, {dt3}"); // ── Classification predicates ────────────────────────────────────────── // NumPy: np.issubdtype(np.float32, np.floating) println!("\n── Classification ──"); - for &dt in &[DType::Bool, DType::I32, DType::U8, DType::F64, DType::C128, DType::BF16] { - println!("{:>12}: integer={}, float={}, complex={}, numeric={}, ordered={}", + for &dt in &[ + DType::Bool, + DType::I32, + DType::U8, + DType::F64, + DType::C128, + DType::BF16, + ] { + println!( + "{:>12}: integer={}, float={}, complex={}, numeric={}, ordered={}", dt.numpy_str(), dt.is_integer(), dt.is_float(), @@ -46,20 +52,20 @@ fn main() { // ── Type conversion helpers ──────────────────────────────────────────── println!("\n── Type conversions ──"); println!("F32.complex_dtype() = {}", DType::F32.complex_dtype()); // C64 - println!("C128.real_dtype() = {}", DType::C128.real_dtype()); // F64 - println!("U16.as_signed() = {}", DType::U16.as_signed()); // I16 - println!("I32.widen() = {}", DType::I32.widen()); // I64 - println!("I64.to_float() = {}", DType::I64.to_float()); // F64 + println!("C128.real_dtype() = {}", DType::C128.real_dtype()); // F64 + println!("U16.as_signed() = {}", DType::U16.as_signed()); // I16 + println!("I32.widen() = {}", DType::I32.widen()); // I64 + println!("I64.to_float() = {}", DType::I64.to_float()); // F64 // ── NumPy string representations ─────────────────────────────────────── // NumPy: np.dtype('float32').str => '21} max={}", - dt.numpy_str(), ii.bits, ii.is_signed, ii.min, ii.max); + println!( + " {:<10} bits={:2} signed={:5} min={:>21} max={}", + dt.numpy_str(), + ii.bits, + ii.is_signed, + ii.min, + ii.max + ); } // ── Iterate all dtypes ───────────────────────────────────────────────── diff --git a/crates/mohu-dtype/examples/type_promotion.rs b/crates/mohu-dtype/examples/type_promotion.rs index ccbc54e..d3a778e 100644 --- a/crates/mohu-dtype/examples/type_promotion.rs +++ b/crates/mohu-dtype/examples/type_promotion.rs @@ -11,8 +11,10 @@ use mohu_dtype::{ DType, - promote::{can_cast, common_type, minimum_scalar_type, promote, result_type, weak_promote, CastMode}, cast::{cast_scalar, cast_slice}, + promote::{ + CastMode, can_cast, common_type, minimum_scalar_type, promote, result_type, weak_promote, + }, }; fn main() { @@ -20,13 +22,13 @@ fn main() { // NumPy: np.result_type(np.int32, np.float32) => np.float64 println!("── Type promotion ──"); let pairs = [ - (DType::I32, DType::F32), // integer + float → wider float - (DType::F16, DType::F32), // float + float → wider float - (DType::I64, DType::F32), // wide int + float → F64 - (DType::C64, DType::F64), // complex + float → wider complex - (DType::Bool, DType::I32), // bool + anything → anything - (DType::U8, DType::I8), // mixed sign → signed, one step wider - (DType::U32, DType::I32), // U32 + I32 → I64 + (DType::I32, DType::F32), // integer + float → wider float + (DType::F16, DType::F32), // float + float → wider float + (DType::I64, DType::F32), // wide int + float → F64 + (DType::C64, DType::F64), // complex + float → wider complex + (DType::Bool, DType::I32), // bool + anything → anything + (DType::U8, DType::I8), // mixed sign → signed, one step wider + (DType::U32, DType::I32), // U32 + I32 → I64 ]; for (a, b) in pairs { let result = promote(a, b); @@ -37,7 +39,10 @@ fn main() { println!("\n Symmetry check:"); println!(" promote(I32, F32) = {}", promote(DType::I32, DType::F32)); println!(" promote(F32, I32) = {}", promote(DType::F32, DType::I32)); - assert_eq!(promote(DType::I32, DType::F32), promote(DType::F32, DType::I32)); + assert_eq!( + promote(DType::I32, DType::F32), + promote(DType::F32, DType::I32) + ); // ── common_type (reduce over multiple dtypes) ────────────────────────── // NumPy: np.result_type(np.int8, np.float32, np.int16) @@ -58,44 +63,74 @@ fn main() { // Safe: no information loss println!(" Safe casts:"); - println!(" I8 → F32: {}", can_cast(DType::I8, DType::F32, CastMode::Safe)); // true - println!(" F64 → F32: {}", can_cast(DType::F64, DType::F32, CastMode::Safe)); // false - println!(" U8 → I16: {}", can_cast(DType::U8, DType::I16, CastMode::Safe)); // true - println!(" U8 → I8: {}", can_cast(DType::U8, DType::I8, CastMode::Safe)); // false + println!( + " I8 → F32: {}", + can_cast(DType::I8, DType::F32, CastMode::Safe) + ); // true + println!( + " F64 → F32: {}", + can_cast(DType::F64, DType::F32, CastMode::Safe) + ); // false + println!( + " U8 → I16: {}", + can_cast(DType::U8, DType::I16, CastMode::Safe) + ); // true + println!( + " U8 → I8: {}", + can_cast(DType::U8, DType::I8, CastMode::Safe) + ); // false // SameKind: within same kind, precision loss OK println!(" SameKind casts:"); - println!(" F64 → F32: {}", can_cast(DType::F64, DType::F32, CastMode::SameKind)); // true - println!(" I32 → I16: {}", can_cast(DType::I32, DType::I16, CastMode::SameKind)); // true - println!(" F32 → I32: {}", can_cast(DType::F32, DType::I32, CastMode::SameKind)); // false + println!( + " F64 → F32: {}", + can_cast(DType::F64, DType::F32, CastMode::SameKind) + ); // true + println!( + " I32 → I16: {}", + can_cast(DType::I32, DType::I16, CastMode::SameKind) + ); // true + println!( + " F32 → I32: {}", + can_cast(DType::F32, DType::I32, CastMode::SameKind) + ); // false // Unsafe: any cast allowed println!(" Unsafe casts:"); - println!(" F32 → I32: {}", can_cast(DType::F32, DType::I32, CastMode::Unsafe)); // true - println!(" C64 → F32: {}", can_cast(DType::C64, DType::F32, CastMode::Unsafe)); // true + println!( + " F32 → I32: {}", + can_cast(DType::F32, DType::I32, CastMode::Unsafe) + ); // true + println!( + " C64 → F32: {}", + can_cast(DType::C64, DType::F32, CastMode::Unsafe) + ); // true // ── Scalar casting (cast_scalar / cast_slice) ────────────────────────── // NumPy: int(np.float64(3.7)) => 3 (truncation) println!("\n── Scalar casting ──"); let v: i32 = cast_scalar::(3.7, CastMode::Unsafe).unwrap(); - println!(" cast f64(3.7) → i32 = {v}"); // 3 (truncated) + println!(" cast f64(3.7) → i32 = {v}"); // 3 (truncated) let v: f64 = cast_scalar::(42_i16, CastMode::Safe).unwrap(); - println!(" cast i16(42) → f64 = {v}"); // 42.0 + println!(" cast i16(42) → f64 = {v}"); // 42.0 // Slice casting let src: Vec = vec![1.1, 2.5, 3.9, -4.2]; let mut dst = vec![0i32; 4]; cast_slice::(&src, &mut dst, CastMode::Unsafe).unwrap(); - println!(" cast_slice f32 → i32: {dst:?}"); // [1, 2, 3, -4] + println!(" cast_slice f32 → i32: {dst:?}"); // [1, 2, 3, -4] // ── minimum_scalar_type (np.result_type for scalars) ─────────────────── // NumPy: np.result_type(42) => dtype('int8') println!("\n── Minimum scalar type ──"); let values = [0.0, 200.0, -1.5, 70000.0, 1e308, -130.0]; for v in values { - println!(" minimum_scalar_type({v:>10}) = {}", minimum_scalar_type(v)); + println!( + " minimum_scalar_type({v:>10}) = {}", + minimum_scalar_type(v) + ); } // ── Weak promotion (NumPy 2.0 semantics) ────────────────────────────── diff --git a/crates/mohu-dtype/src/cast.rs b/crates/mohu-dtype/src/cast.rs index 925a692..f4ea2b8 100644 --- a/crates/mohu-dtype/src/cast.rs +++ b/crates/mohu-dtype/src/cast.rs @@ -18,10 +18,7 @@ /// output slice. use mohu_error::{MohuError, MohuResult}; -use crate::{ - promote::CastMode, - scalar::Scalar, -}; +use crate::{promote::CastMode, scalar::Scalar}; // ─── cast_scalar ───────────────────────────────────────────────────────────── @@ -40,8 +37,8 @@ use crate::{ pub fn cast_scalar(value: S, mode: CastMode) -> MohuResult { if !crate::promote::can_cast(S::DTYPE, D::DTYPE, mode) { return Err(MohuError::InvalidCast { - from: S::DTYPE.to_string(), - to: D::DTYPE.to_string(), + from: S::DTYPE.to_string(), + to: D::DTYPE.to_string(), reason: format!("{mode:?} cast is not allowed"), }); } @@ -82,14 +79,14 @@ pub fn cast_slice( if src.len() != dst.len() { return Err(MohuError::ShapeMismatch { expected: vec![src.len()], - got: vec![dst.len()], + got: vec![dst.len()], }); } // Check mode once before the loop. if !crate::promote::can_cast(S::DTYPE, D::DTYPE, mode) { return Err(MohuError::InvalidCast { - from: S::DTYPE.to_string(), - to: D::DTYPE.to_string(), + from: S::DTYPE.to_string(), + to: D::DTYPE.to_string(), reason: format!("{mode:?} cast is not allowed"), }); } @@ -111,10 +108,7 @@ fn byte_copy_cast(value: S) -> D { // `D::ITEMSIZE == S::ITEMSIZE` (asserted in debug builds above), so the // ptr::read_unaligned reads exactly the right number of bytes. unsafe { - let bytes = std::slice::from_raw_parts( - &value as *const S as *const u8, - S::ITEMSIZE, - ); + let bytes = std::slice::from_raw_parts(&value as *const S as *const u8, S::ITEMSIZE); buf[..S::ITEMSIZE].copy_from_slice(bytes); std::ptr::read_unaligned(buf.as_ptr() as *const D) } @@ -151,8 +145,12 @@ where } let min_f = D::min_value().to_f64_lossy(); let max_f = D::max_value().to_f64_lossy(); - if v <= min_f { return D::min_value(); } - if v >= max_f { return D::max_value(); } + if v <= min_f { + return D::min_value(); + } + if v >= max_f { + return D::max_value(); + } D::from_f64_lossy(v) } @@ -161,9 +159,9 @@ where impl std::fmt::Display for CastMode { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - Self::Safe => write!(f, "safe"), + Self::Safe => write!(f, "safe"), Self::SameKind => write!(f, "same_kind"), - Self::Unsafe => write!(f, "unsafe"), + Self::Unsafe => write!(f, "unsafe"), } } } diff --git a/crates/mohu-dtype/src/compat.rs b/crates/mohu-dtype/src/compat.rs index 071924e..31c1fbc 100644 --- a/crates/mohu-dtype/src/compat.rs +++ b/crates/mohu-dtype/src/compat.rs @@ -42,9 +42,9 @@ impl ByteOrder { /// Returns the struct-module prefix character for this byte order. pub const fn struct_prefix(self) -> char { match self { - Self::Native => '=', - Self::Little => '<', - Self::Big => '>', + Self::Native => '=', + Self::Little => '<', + Self::Big => '>', Self::NotApplicable => '|', } } @@ -53,7 +53,11 @@ impl ByteOrder { /// item size. Single-byte types use `NotApplicable`; multi-byte types /// use `Native`. pub const fn for_itemsize(itemsize: usize) -> Self { - if itemsize == 1 { Self::NotApplicable } else { Self::Native } + if itemsize == 1 { + Self::NotApplicable + } else { + Self::Native + } } /// Returns `true` if the current host is little-endian. @@ -63,7 +67,11 @@ impl ByteOrder { /// Returns the concrete byte order of the host. pub fn host() -> Self { - if Self::host_is_little_endian() { Self::Little } else { Self::Big } + if Self::host_is_little_endian() { + Self::Little + } else { + Self::Big + } } } @@ -88,20 +96,20 @@ impl DType { pub const fn struct_format_char(self) -> Option { match self { Self::Bool => Some('?'), - Self::I8 => Some('b'), - Self::I16 => Some('h'), - Self::I32 => Some('i'), - Self::I64 => Some('q'), - Self::U8 => Some('B'), - Self::U16 => Some('H'), - Self::U32 => Some('I'), - Self::U64 => Some('Q'), - Self::F16 => Some('e'), - Self::BF16 => None, // no struct code - Self::F32 => Some('f'), - Self::F64 => Some('d'), + Self::I8 => Some('b'), + Self::I16 => Some('h'), + Self::I32 => Some('i'), + Self::I64 => Some('q'), + Self::U8 => Some('B'), + Self::U16 => Some('H'), + Self::U32 => Some('I'), + Self::U64 => Some('Q'), + Self::F16 => Some('e'), + Self::BF16 => None, // no struct code + Self::F32 => Some('f'), + Self::F64 => Some('d'), // Complex: struct encodes as two consecutive floats - Self::C64 => None, + Self::C64 => None, Self::C128 => None, } } @@ -115,12 +123,12 @@ impl DType { /// assert_eq!(DType::I64.to_struct_format().unwrap(), "=q"); /// ``` pub fn to_struct_format(self) -> MohuResult { - let ch = self.struct_format_char().ok_or_else(|| { - MohuError::UnsupportedDType { - op: "struct format", + let ch = self + .struct_format_char() + .ok_or_else(|| MohuError::UnsupportedDType { + op: "struct format", dtype: self.to_string(), - } - })?; + })?; Ok(format!("={ch}")) } @@ -146,19 +154,19 @@ impl DType { pub fn buffer_format(self) -> &'static str { match self { Self::Bool => "|?", - Self::I8 => "|b", - Self::U8 => "|B", - Self::I16 => " " " " " " " " " " " "|b", + Self::U8 => "|B", + Self::I16 => " " " " " " " " " " " " Option<&'static str> { match self { Self::Bool => Some("c_bool"), - Self::I8 => Some("c_int8"), - Self::I16 => Some("c_int16"), - Self::I32 => Some("c_int32"), - Self::I64 => Some("c_int64"), - Self::U8 => Some("c_uint8"), - Self::U16 => Some("c_uint16"), - Self::U32 => Some("c_uint32"), - Self::U64 => Some("c_uint64"), - Self::F32 => Some("c_float"), - Self::F64 => Some("c_double"), - _ => None, // F16, BF16, complex have no ctypes equivalent + Self::I8 => Some("c_int8"), + Self::I16 => Some("c_int16"), + Self::I32 => Some("c_int32"), + Self::I64 => Some("c_int64"), + Self::U8 => Some("c_uint8"), + Self::U16 => Some("c_uint16"), + Self::U32 => Some("c_uint32"), + Self::U64 => Some("c_uint64"), + Self::F32 => Some("c_float"), + Self::F64 => Some("c_double"), + _ => None, // F16, BF16, complex have no ctypes equivalent } } @@ -220,19 +228,19 @@ impl DType { pub const fn c_type_name(self) -> &'static str { match self { Self::Bool => "_Bool", - Self::I8 => "int8_t", - Self::I16 => "int16_t", - Self::I32 => "int32_t", - Self::I64 => "int64_t", - Self::U8 => "uint8_t", - Self::U16 => "uint16_t", - Self::U32 => "uint32_t", - Self::U64 => "uint64_t", - Self::F16 => "__fp16", + Self::I8 => "int8_t", + Self::I16 => "int16_t", + Self::I32 => "int32_t", + Self::I64 => "int64_t", + Self::U8 => "uint8_t", + Self::U16 => "uint16_t", + Self::U32 => "uint32_t", + Self::U64 => "uint64_t", + Self::F16 => "__fp16", Self::BF16 => "__bfloat16", - Self::F32 => "float", - Self::F64 => "double", - Self::C64 => "float _Complex", + Self::F32 => "float", + Self::F64 => "double", + Self::C64 => "float _Complex", Self::C128 => "double _Complex", } } @@ -243,19 +251,19 @@ impl DType { pub const fn rust_type_name(self) -> &'static str { match self { Self::Bool => "bool", - Self::I8 => "i8", - Self::I16 => "i16", - Self::I32 => "i32", - Self::I64 => "i64", - Self::U8 => "u8", - Self::U16 => "u16", - Self::U32 => "u32", - Self::U64 => "u64", - Self::F16 => "half::f16", + Self::I8 => "i8", + Self::I16 => "i16", + Self::I32 => "i32", + Self::I64 => "i64", + Self::U8 => "u8", + Self::U16 => "u16", + Self::U32 => "u32", + Self::U64 => "u64", + Self::F16 => "half::f16", Self::BF16 => "half::bf16", - Self::F32 => "f32", - Self::F64 => "f64", - Self::C64 => "num_complex::Complex", + Self::F32 => "f32", + Self::F64 => "f64", + Self::C64 => "num_complex::Complex", Self::C128 => "num_complex::Complex", } } @@ -278,17 +286,17 @@ pub mod arrow_compat { pub fn to_arrow(self) -> Option { match self { Self::Bool => Some(ArrowDataType::Boolean), - Self::I8 => Some(ArrowDataType::Int8), - Self::I16 => Some(ArrowDataType::Int16), - Self::I32 => Some(ArrowDataType::Int32), - Self::I64 => Some(ArrowDataType::Int64), - Self::U8 => Some(ArrowDataType::UInt8), - Self::U16 => Some(ArrowDataType::UInt16), - Self::U32 => Some(ArrowDataType::UInt32), - Self::U64 => Some(ArrowDataType::UInt64), - Self::F16 => Some(ArrowDataType::Float16), - Self::F32 => Some(ArrowDataType::Float32), - Self::F64 => Some(ArrowDataType::Float64), + Self::I8 => Some(ArrowDataType::Int8), + Self::I16 => Some(ArrowDataType::Int16), + Self::I32 => Some(ArrowDataType::Int32), + Self::I64 => Some(ArrowDataType::Int64), + Self::U8 => Some(ArrowDataType::UInt8), + Self::U16 => Some(ArrowDataType::UInt16), + Self::U32 => Some(ArrowDataType::UInt32), + Self::U64 => Some(ArrowDataType::UInt64), + Self::F16 => Some(ArrowDataType::Float16), + Self::F32 => Some(ArrowDataType::Float32), + Self::F64 => Some(ArrowDataType::Float64), // BF16 / complex have no standard Arrow equivalent Self::BF16 | Self::C64 | Self::C128 => None, } @@ -300,18 +308,18 @@ pub mod arrow_compat { /// no mohu equivalent (timestamps, dictionaries, lists, etc.). pub fn from_arrow(dt: &ArrowDataType) -> MohuResult { match dt { - ArrowDataType::Boolean => Ok(Self::Bool), - ArrowDataType::Int8 => Ok(Self::I8), - ArrowDataType::Int16 => Ok(Self::I16), - ArrowDataType::Int32 => Ok(Self::I32), - ArrowDataType::Int64 => Ok(Self::I64), - ArrowDataType::UInt8 => Ok(Self::U8), - ArrowDataType::UInt16 => Ok(Self::U16), - ArrowDataType::UInt32 => Ok(Self::U32), - ArrowDataType::UInt64 => Ok(Self::U64), - ArrowDataType::Float16 => Ok(Self::F16), - ArrowDataType::Float32 => Ok(Self::F32), - ArrowDataType::Float64 => Ok(Self::F64), + ArrowDataType::Boolean => Ok(Self::Bool), + ArrowDataType::Int8 => Ok(Self::I8), + ArrowDataType::Int16 => Ok(Self::I16), + ArrowDataType::Int32 => Ok(Self::I32), + ArrowDataType::Int64 => Ok(Self::I64), + ArrowDataType::UInt8 => Ok(Self::U8), + ArrowDataType::UInt16 => Ok(Self::U16), + ArrowDataType::UInt32 => Ok(Self::U32), + ArrowDataType::UInt64 => Ok(Self::U64), + ArrowDataType::Float16 => Ok(Self::F16), + ArrowDataType::Float32 => Ok(Self::F32), + ArrowDataType::Float64 => Ok(Self::F64), other => Err(MohuError::ArrowUnsupportedType { arrow_type: format!("{other:?}"), }), diff --git a/crates/mohu-dtype/src/dlpack.rs b/crates/mohu-dtype/src/dlpack.rs index ce6d3e6..f5325a4 100644 --- a/crates/mohu-dtype/src/dlpack.rs +++ b/crates/mohu-dtype/src/dlpack.rs @@ -33,10 +33,10 @@ use crate::dtype::DType; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] #[repr(u8)] pub enum DLDataTypeCode { - Int = 0, - UInt = 1, - Float = 2, - BFloat = 4, + Int = 0, + UInt = 1, + Float = 2, + BFloat = 4, Complex = 5, } @@ -62,9 +62,9 @@ impl DLDataTypeCode { /// A parsed representation of a DLPack `DLDataType` struct. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct DLDataType { - pub code: DLDataTypeCode, + pub code: DLDataTypeCode, /// Bits per scalar element (per lane). - pub bits: u8, + pub bits: u8, /// Number of SIMD lanes. Must be 1 for mohu arrays. pub lanes: u16, } @@ -72,7 +72,11 @@ pub struct DLDataType { impl DLDataType { /// Constructs a `DLDataType` with `lanes = 1`. pub const fn scalar(code: DLDataTypeCode, bits: u8) -> Self { - Self { code, bits, lanes: 1 } + Self { + code, + bits, + lanes: 1, + } } /// Parses from raw `(code, bits, lanes)` triple. @@ -104,21 +108,21 @@ impl DType { pub const fn to_dlpack(self) -> DLDataType { use DLDataTypeCode::*; match self { - Self::Bool => DLDataType::scalar(UInt, 8), - Self::I8 => DLDataType::scalar(Int, 8), - Self::I16 => DLDataType::scalar(Int, 16), - Self::I32 => DLDataType::scalar(Int, 32), - Self::I64 => DLDataType::scalar(Int, 64), - Self::U8 => DLDataType::scalar(UInt, 8), - Self::U16 => DLDataType::scalar(UInt, 16), - Self::U32 => DLDataType::scalar(UInt, 32), - Self::U64 => DLDataType::scalar(UInt, 64), - Self::F16 => DLDataType::scalar(Float, 16), + Self::Bool => DLDataType::scalar(UInt, 8), + Self::I8 => DLDataType::scalar(Int, 8), + Self::I16 => DLDataType::scalar(Int, 16), + Self::I32 => DLDataType::scalar(Int, 32), + Self::I64 => DLDataType::scalar(Int, 64), + Self::U8 => DLDataType::scalar(UInt, 8), + Self::U16 => DLDataType::scalar(UInt, 16), + Self::U32 => DLDataType::scalar(UInt, 32), + Self::U64 => DLDataType::scalar(UInt, 64), + Self::F16 => DLDataType::scalar(Float, 16), Self::BF16 => DLDataType::scalar(BFloat, 16), - Self::F32 => DLDataType::scalar(Float, 32), - Self::F64 => DLDataType::scalar(Float, 64), - Self::C64 => DLDataType::scalar(Complex, 64), - Self::C128 => DLDataType::scalar(Complex,128), + Self::F32 => DLDataType::scalar(Float, 32), + Self::F64 => DLDataType::scalar(Float, 64), + Self::C64 => DLDataType::scalar(Complex, 64), + Self::C128 => DLDataType::scalar(Complex, 128), } } @@ -141,20 +145,20 @@ impl DType { } let kind = DLDataTypeCode::from_u8(code)?; match (kind, bits) { - (DLDataTypeCode::Int, 8) => Ok(Self::I8), - (DLDataTypeCode::Int, 16) => Ok(Self::I16), - (DLDataTypeCode::Int, 32) => Ok(Self::I32), - (DLDataTypeCode::Int, 64) => Ok(Self::I64), - (DLDataTypeCode::UInt, 8) => Ok(Self::U8), // Bool also maps here - (DLDataTypeCode::UInt, 16) => Ok(Self::U16), - (DLDataTypeCode::UInt, 32) => Ok(Self::U32), - (DLDataTypeCode::UInt, 64) => Ok(Self::U64), - (DLDataTypeCode::Float, 16) => Ok(Self::F16), - (DLDataTypeCode::Float, 32) => Ok(Self::F32), - (DLDataTypeCode::Float, 64) => Ok(Self::F64), - (DLDataTypeCode::BFloat, 16) => Ok(Self::BF16), + (DLDataTypeCode::Int, 8) => Ok(Self::I8), + (DLDataTypeCode::Int, 16) => Ok(Self::I16), + (DLDataTypeCode::Int, 32) => Ok(Self::I32), + (DLDataTypeCode::Int, 64) => Ok(Self::I64), + (DLDataTypeCode::UInt, 8) => Ok(Self::U8), // Bool also maps here + (DLDataTypeCode::UInt, 16) => Ok(Self::U16), + (DLDataTypeCode::UInt, 32) => Ok(Self::U32), + (DLDataTypeCode::UInt, 64) => Ok(Self::U64), + (DLDataTypeCode::Float, 16) => Ok(Self::F16), + (DLDataTypeCode::Float, 32) => Ok(Self::F32), + (DLDataTypeCode::Float, 64) => Ok(Self::F64), + (DLDataTypeCode::BFloat, 16) => Ok(Self::BF16), (DLDataTypeCode::Complex, 64) => Ok(Self::C64), - (DLDataTypeCode::Complex,128) => Ok(Self::C128), + (DLDataTypeCode::Complex, 128) => Ok(Self::C128), (_, b) => Err(MohuError::DLPackUnsupportedDType { code, bits: b, @@ -173,39 +177,39 @@ impl DType { #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] #[repr(i32)] pub enum DLDeviceType { - Cpu = 1, - Cuda = 2, - CpuPinned = 3, - OpenCL = 4, - Vulkan = 7, - Metal = 8, - Vpi = 9, - Rocm = 10, - ExtDev = 12, - CudaManaged= 13, - OneApi = 14, - WebGpu = 15, - Hexagon = 16, + Cpu = 1, + Cuda = 2, + CpuPinned = 3, + OpenCL = 4, + Vulkan = 7, + Metal = 8, + Vpi = 9, + Rocm = 10, + ExtDev = 12, + CudaManaged = 13, + OneApi = 14, + WebGpu = 15, + Hexagon = 16, } impl DLDeviceType { /// Parses a device type from its raw `i32` code. pub fn from_i32(code: i32) -> Option { match code { - 1 => Some(Self::Cpu), - 2 => Some(Self::Cuda), - 3 => Some(Self::CpuPinned), - 4 => Some(Self::OpenCL), - 7 => Some(Self::Vulkan), - 8 => Some(Self::Metal), - 9 => Some(Self::Vpi), + 1 => Some(Self::Cpu), + 2 => Some(Self::Cuda), + 3 => Some(Self::CpuPinned), + 4 => Some(Self::OpenCL), + 7 => Some(Self::Vulkan), + 8 => Some(Self::Metal), + 9 => Some(Self::Vpi), 10 => Some(Self::Rocm), 12 => Some(Self::ExtDev), 13 => Some(Self::CudaManaged), 14 => Some(Self::OneApi), 15 => Some(Self::WebGpu), 16 => Some(Self::Hexagon), - _ => None, + _ => None, } } @@ -219,19 +223,19 @@ impl DLDeviceType { /// Human-readable device name. pub fn name(self) -> &'static str { match self { - Self::Cpu => "CPU", - Self::Cuda => "CUDA", - Self::CpuPinned => "CPU (pinned)", - Self::OpenCL => "OpenCL", - Self::Vulkan => "Vulkan", - Self::Metal => "Metal", - Self::Vpi => "VPI", - Self::Rocm => "ROCm", - Self::ExtDev => "ExtDev", + Self::Cpu => "CPU", + Self::Cuda => "CUDA", + Self::CpuPinned => "CPU (pinned)", + Self::OpenCL => "OpenCL", + Self::Vulkan => "Vulkan", + Self::Metal => "Metal", + Self::Vpi => "VPI", + Self::Rocm => "ROCm", + Self::ExtDev => "ExtDev", Self::CudaManaged => "CUDA (managed)", - Self::OneApi => "OneAPI", - Self::WebGpu => "WebGPU", - Self::Hexagon => "Hexagon", + Self::OneApi => "OneAPI", + Self::WebGpu => "WebGPU", + Self::Hexagon => "Hexagon", } } } diff --git a/crates/mohu-dtype/src/dtype.rs b/crates/mohu-dtype/src/dtype.rs index 81bf0ea..130a1b4 100644 --- a/crates/mohu-dtype/src/dtype.rs +++ b/crates/mohu-dtype/src/dtype.rs @@ -1,5 +1,5 @@ -use std::fmt; use mohu_error::{MohuError, MohuResult}; +use std::fmt; /// The runtime type tag for every element type mohu supports. /// @@ -37,19 +37,19 @@ use mohu_error::{MohuError, MohuResult}; #[repr(u8)] pub enum DType { Bool = 0, - I8 = 1, - I16 = 2, - I32 = 3, - I64 = 4, - U8 = 5, - U16 = 6, - U32 = 7, - U64 = 8, - F16 = 9, + I8 = 1, + I16 = 2, + I32 = 3, + I64 = 4, + U8 = 5, + U16 = 6, + U32 = 7, + U64 = 8, + F16 = 9, BF16 = 10, - F32 = 11, - F64 = 12, - C64 = 13, + F32 = 11, + F64 = 12, + C64 = 13, C128 = 14, } @@ -59,10 +59,20 @@ pub const DTYPE_COUNT: usize = 15; /// All `DType` variants in definition order. pub const ALL_DTYPES: [DType; DTYPE_COUNT] = [ DType::Bool, - DType::I8, DType::I16, DType::I32, DType::I64, - DType::U8, DType::U16, DType::U32, DType::U64, - DType::F16, DType::BF16, DType::F32, DType::F64, - DType::C64, DType::C128, + DType::I8, + DType::I16, + DType::I32, + DType::I64, + DType::U8, + DType::U16, + DType::U32, + DType::U64, + DType::F16, + DType::BF16, + DType::F32, + DType::F64, + DType::C64, + DType::C128, ]; impl DType { @@ -80,11 +90,11 @@ impl DType { /// ``` pub const fn itemsize(self) -> usize { match self { - Self::Bool | Self::I8 | Self::U8 => 1, - Self::I16 | Self::U16 | Self::F16 | Self::BF16 => 2, - Self::I32 | Self::U32 | Self::F32 => 4, - Self::I64 | Self::U64 | Self::F64 | Self::C64 => 8, - Self::C128 => 16, + Self::Bool | Self::I8 | Self::U8 => 1, + Self::I16 | Self::U16 | Self::F16 | Self::BF16 => 2, + Self::I32 | Self::U32 | Self::F32 => 4, + Self::I64 | Self::U64 | Self::F64 | Self::C64 => 8, + Self::C128 => 16, } } @@ -92,10 +102,10 @@ impl DType { /// corresponding Rust primitive. pub const fn alignment(self) -> usize { match self { - Self::Bool | Self::I8 | Self::U8 => 1, - Self::I16 | Self::U16 | Self::F16 | Self::BF16 => 2, - Self::I32 | Self::U32 | Self::F32 | Self::C64 => 4, - Self::I64 | Self::U64 | Self::F64 | Self::C128 => 8, + Self::Bool | Self::I8 | Self::U8 => 1, + Self::I16 | Self::U16 | Self::F16 | Self::BF16 => 2, + Self::I32 | Self::U32 | Self::F32 | Self::C64 => 4, + Self::I64 | Self::U64 | Self::F64 | Self::C128 => 8, } } @@ -109,61 +119,81 @@ impl DType { // ------------------------------------------------------------------------- /// Returns `true` if this is `Bool`. - #[inline] pub const fn is_bool(self) -> bool { + #[inline] + pub const fn is_bool(self) -> bool { matches!(self, Self::Bool) } /// Returns `true` if this is any integer type (signed or unsigned). - #[inline] pub const fn is_integer(self) -> bool { - matches!(self, Self::I8|Self::I16|Self::I32|Self::I64 - | Self::U8|Self::U16|Self::U32|Self::U64) + #[inline] + pub const fn is_integer(self) -> bool { + matches!( + self, + Self::I8 + | Self::I16 + | Self::I32 + | Self::I64 + | Self::U8 + | Self::U16 + | Self::U32 + | Self::U64 + ) } /// Returns `true` if this is a signed integer. - #[inline] pub const fn is_signed_integer(self) -> bool { - matches!(self, Self::I8|Self::I16|Self::I32|Self::I64) + #[inline] + pub const fn is_signed_integer(self) -> bool { + matches!(self, Self::I8 | Self::I16 | Self::I32 | Self::I64) } /// Returns `true` if this is an unsigned integer. - #[inline] pub const fn is_unsigned_integer(self) -> bool { - matches!(self, Self::U8|Self::U16|Self::U32|Self::U64) + #[inline] + pub const fn is_unsigned_integer(self) -> bool { + matches!(self, Self::U8 | Self::U16 | Self::U32 | Self::U64) } /// Returns `true` if this is a non-complex floating-point type /// (F16, BF16, F32, or F64). - #[inline] pub const fn is_float(self) -> bool { - matches!(self, Self::F16|Self::BF16|Self::F32|Self::F64) + #[inline] + pub const fn is_float(self) -> bool { + matches!(self, Self::F16 | Self::BF16 | Self::F32 | Self::F64) } /// Returns `true` if this is a complex type (C64 or C128). - #[inline] pub const fn is_complex(self) -> bool { - matches!(self, Self::C64|Self::C128) + #[inline] + pub const fn is_complex(self) -> bool { + matches!(self, Self::C64 | Self::C128) } /// Returns `true` for any floating-point type, including complex. - #[inline] pub const fn is_floating_point(self) -> bool { + #[inline] + pub const fn is_floating_point(self) -> bool { self.is_float() || self.is_complex() } /// Returns `true` for any type that supports arithmetic (everything /// except `Bool`). - #[inline] pub const fn is_numeric(self) -> bool { + #[inline] + pub const fn is_numeric(self) -> bool { !matches!(self, Self::Bool) } /// Returns `true` for types that have a total order (integers and /// real floats — not complex, not bool). - #[inline] pub const fn is_ordered(self) -> bool { + #[inline] + pub const fn is_ordered(self) -> bool { self.is_integer() || self.is_float() } /// Returns `true` for standard IEEE 754 precision (F32 or F64). - #[inline] pub const fn is_standard_float(self) -> bool { - matches!(self, Self::F32|Self::F64) + #[inline] + pub const fn is_standard_float(self) -> bool { + matches!(self, Self::F32 | Self::F64) } /// Returns `true` for `BF16` (brain float). - #[inline] pub const fn is_brain_float(self) -> bool { + #[inline] + pub const fn is_brain_float(self) -> bool { matches!(self, Self::BF16) } @@ -175,7 +205,7 @@ impl DType { /// C64 → F32, C128 → F64. For all other types, returns `self`. pub const fn real_dtype(self) -> DType { match self { - Self::C64 => Self::F32, + Self::C64 => Self::F32, Self::C128 => Self::F64, other => other, } @@ -195,7 +225,7 @@ impl DType { /// U8 → I8, U16 → I16, U32 → I32, U64 → I64. Others return `self`. pub const fn as_signed(self) -> DType { match self { - Self::U8 => Self::I8, + Self::U8 => Self::I8, Self::U16 => Self::I16, Self::U32 => Self::I32, Self::U64 => Self::I64, @@ -207,7 +237,7 @@ impl DType { /// I8 → U8, I16 → U16, I32 → U32, I64 → U64. Others return `self`. pub const fn as_unsigned(self) -> DType { match self { - Self::I8 => Self::U8, + Self::I8 => Self::U8, Self::I16 => Self::U16, Self::I32 => Self::U32, Self::I64 => Self::U64, @@ -221,19 +251,19 @@ impl DType { pub const fn widen(self) -> DType { match self { Self::Bool => Self::U8, - Self::I8 => Self::I16, - Self::I16 => Self::I32, - Self::I32 => Self::I64, - Self::I64 => Self::I64, - Self::U8 => Self::U16, - Self::U16 => Self::U32, - Self::U32 => Self::U64, - Self::U64 => Self::U64, - Self::F16 => Self::F32, + Self::I8 => Self::I16, + Self::I16 => Self::I32, + Self::I32 => Self::I64, + Self::I64 => Self::I64, + Self::U8 => Self::U16, + Self::U16 => Self::U32, + Self::U32 => Self::U64, + Self::U64 => Self::U64, + Self::F16 => Self::F32, Self::BF16 => Self::F32, - Self::F32 => Self::F64, - Self::F64 => Self::F64, - Self::C64 => Self::C128, + Self::F32 => Self::F64, + Self::F64 => Self::F64, + Self::C64 => Self::C128, Self::C128 => Self::C128, } } @@ -242,14 +272,14 @@ impl DType { /// types (Bool, I8, U8, F16, BF16). pub const fn narrow(self) -> Option { match self { - Self::I16 => Some(Self::I8), - Self::I32 => Some(Self::I16), - Self::I64 => Some(Self::I32), - Self::U16 => Some(Self::U8), - Self::U32 => Some(Self::U16), - Self::U64 => Some(Self::U32), - Self::F32 => Some(Self::F16), - Self::F64 => Some(Self::F32), + Self::I16 => Some(Self::I8), + Self::I32 => Some(Self::I16), + Self::I64 => Some(Self::I32), + Self::U16 => Some(Self::U8), + Self::U32 => Some(Self::U16), + Self::U64 => Some(Self::U32), + Self::F32 => Some(Self::F16), + Self::F64 => Some(Self::F32), Self::C128 => Some(Self::C64), _ => None, } @@ -262,10 +292,9 @@ impl DType { /// For float and complex, returns `self`. For Bool, returns F16. pub const fn to_float(self) -> DType { match self { - Self::Bool | Self::I8 | Self::U8 => Self::F16, - Self::I16 | Self::U16 => Self::F32, - Self::I32 | Self::U32 - | Self::I64 | Self::U64 => Self::F64, + Self::Bool | Self::I8 | Self::U8 => Self::F16, + Self::I16 | Self::U16 => Self::F32, + Self::I32 | Self::U32 | Self::I64 | Self::U64 => Self::F64, other => other, } } @@ -278,19 +307,19 @@ impl DType { pub const fn numpy_str(self) -> &'static str { match self { Self::Bool => "bool", - Self::I8 => "int8", - Self::I16 => "int16", - Self::I32 => "int32", - Self::I64 => "int64", - Self::U8 => "uint8", - Self::U16 => "uint16", - Self::U32 => "uint32", - Self::U64 => "uint64", - Self::F16 => "float16", + Self::I8 => "int8", + Self::I16 => "int16", + Self::I32 => "int32", + Self::I64 => "int64", + Self::U8 => "uint8", + Self::U16 => "uint16", + Self::U32 => "uint32", + Self::U64 => "uint64", + Self::F16 => "float16", Self::BF16 => "bfloat16", - Self::F32 => "float32", - Self::F64 => "float64", - Self::C64 => "complex64", + Self::F32 => "float32", + Self::F64 => "float64", + Self::C64 => "complex64", Self::C128 => "complex128", } } @@ -310,19 +339,19 @@ impl DType { pub const fn numpy_char(self) -> Option { match self { Self::Bool => Some('?'), - Self::I8 => Some('b'), - Self::I16 => Some('h'), - Self::I32 => Some('i'), - Self::I64 => Some('l'), - Self::U8 => Some('B'), - Self::U16 => Some('H'), - Self::U32 => Some('I'), - Self::U64 => Some('L'), - Self::F16 => Some('e'), + Self::I8 => Some('b'), + Self::I16 => Some('h'), + Self::I32 => Some('i'), + Self::I64 => Some('l'), + Self::U8 => Some('B'), + Self::U16 => Some('H'), + Self::U32 => Some('I'), + Self::U64 => Some('L'), + Self::F16 => Some('e'), Self::BF16 => None, - Self::F32 => Some('f'), - Self::F64 => Some('d'), - Self::C64 => Some('F'), + Self::F32 => Some('f'), + Self::F64 => Some('d'), + Self::C64 => Some('F'), Self::C128 => Some('D'), } } @@ -333,10 +362,10 @@ impl DType { pub const fn kind_char(self) -> char { match self { Self::Bool => 'b', - Self::I8|Self::I16|Self::I32|Self::I64 => 'i', - Self::U8|Self::U16|Self::U32|Self::U64 => 'u', - Self::F16|Self::BF16|Self::F32|Self::F64 => 'f', - Self::C64|Self::C128 => 'c', + Self::I8 | Self::I16 | Self::I32 | Self::I64 => 'i', + Self::U8 | Self::U16 | Self::U32 | Self::U64 => 'u', + Self::F16 | Self::BF16 | Self::F32 | Self::F64 => 'f', + Self::C64 | Self::C128 => 'c', } } @@ -355,15 +384,15 @@ impl DType { pub fn min_as_f64(self) -> Option { match self { Self::Bool => Some(0.0), - Self::I8 => Some(i8::MIN as f64), - Self::I16 => Some(i16::MIN as f64), - Self::I32 => Some(i32::MIN as f64), - Self::I64 => Some(i64::MIN as f64), + Self::I8 => Some(i8::MIN as f64), + Self::I16 => Some(i16::MIN as f64), + Self::I32 => Some(i32::MIN as f64), + Self::I64 => Some(i64::MIN as f64), Self::U8 | Self::U16 | Self::U32 | Self::U64 => Some(0.0), - Self::F16 => Some(half::f16::MIN.to_f64()), + Self::F16 => Some(half::f16::MIN.to_f64()), Self::BF16 => Some(half::bf16::MIN.to_f64()), - Self::F32 => Some(f32::MIN as f64), - Self::F64 => Some(f64::MIN), + Self::F32 => Some(f32::MIN as f64), + Self::F64 => Some(f64::MIN), Self::C64 | Self::C128 => None, } } @@ -372,18 +401,18 @@ impl DType { pub fn max_as_f64(self) -> Option { match self { Self::Bool => Some(1.0), - Self::I8 => Some(i8::MAX as f64), - Self::I16 => Some(i16::MAX as f64), - Self::I32 => Some(i32::MAX as f64), - Self::I64 => Some(i64::MAX as f64), - Self::U8 => Some(u8::MAX as f64), - Self::U16 => Some(u16::MAX as f64), - Self::U32 => Some(u32::MAX as f64), - Self::U64 => Some(u64::MAX as f64), - Self::F16 => Some(half::f16::MAX.to_f64()), + Self::I8 => Some(i8::MAX as f64), + Self::I16 => Some(i16::MAX as f64), + Self::I32 => Some(i32::MAX as f64), + Self::I64 => Some(i64::MAX as f64), + Self::U8 => Some(u8::MAX as f64), + Self::U16 => Some(u16::MAX as f64), + Self::U32 => Some(u32::MAX as f64), + Self::U64 => Some(u64::MAX as f64), + Self::F16 => Some(half::f16::MAX.to_f64()), Self::BF16 => Some(half::bf16::MAX.to_f64()), - Self::F32 => Some(f32::MAX as f64), - Self::F64 => Some(f64::MAX), + Self::F32 => Some(f32::MAX as f64), + Self::F64 => Some(f64::MAX), Self::C64 | Self::C128 => None, } } @@ -391,10 +420,10 @@ impl DType { /// Machine epsilon for floating-point types, or `None` for all others. pub fn epsilon_as_f64(self) -> Option { match self { - Self::F16 => Some(half::f16::EPSILON.to_f64()), + Self::F16 => Some(half::f16::EPSILON.to_f64()), Self::BF16 => Some(half::bf16::EPSILON.to_f64()), - Self::F32 => Some(f32::EPSILON as f64), - Self::F64 => Some(f64::EPSILON), + Self::F32 => Some(f32::EPSILON as f64), + Self::F64 => Some(f64::EPSILON), _ => None, } } @@ -402,10 +431,10 @@ impl DType { /// Smallest positive normalised value for float types, or `None`. pub fn min_positive_as_f64(self) -> Option { match self { - Self::F16 => Some(half::f16::MIN_POSITIVE.to_f64()), + Self::F16 => Some(half::f16::MIN_POSITIVE.to_f64()), Self::BF16 => Some(half::bf16::MIN_POSITIVE.to_f64()), - Self::F32 => Some(f32::MIN_POSITIVE as f64), - Self::F64 => Some(f64::MIN_POSITIVE), + Self::F32 => Some(f32::MIN_POSITIVE as f64), + Self::F64 => Some(f64::MIN_POSITIVE), _ => None, } } @@ -413,8 +442,8 @@ impl DType { /// Number of significant decimal digits for integer types, or `None`. pub const fn max_decimal_digits(self) -> Option { match self { - Self::Bool => Some(1), - Self::I8 | Self::U8 => Some(3), + Self::Bool => Some(1), + Self::I8 | Self::U8 => Some(3), Self::I16 | Self::U16 => Some(5), Self::I32 | Self::U32 => Some(10), Self::I64 | Self::U64 => Some(20), @@ -427,27 +456,29 @@ impl DType { // ------------------------------------------------------------------------- /// Returns the stable `u8` discriminant code for this dtype. - pub const fn as_u8(self) -> u8 { self as u8 } + pub const fn as_u8(self) -> u8 { + self as u8 + } /// Constructs a `DType` from its `u8` code. pub fn from_u8(code: u8) -> MohuResult { match code { - 0 => Ok(Self::Bool), - 1 => Ok(Self::I8), - 2 => Ok(Self::I16), - 3 => Ok(Self::I32), - 4 => Ok(Self::I64), - 5 => Ok(Self::U8), - 6 => Ok(Self::U16), - 7 => Ok(Self::U32), - 8 => Ok(Self::U64), - 9 => Ok(Self::F16), + 0 => Ok(Self::Bool), + 1 => Ok(Self::I8), + 2 => Ok(Self::I16), + 3 => Ok(Self::I32), + 4 => Ok(Self::I64), + 5 => Ok(Self::U8), + 6 => Ok(Self::U16), + 7 => Ok(Self::U32), + 8 => Ok(Self::U64), + 9 => Ok(Self::F16), 10 => Ok(Self::BF16), 11 => Ok(Self::F32), 12 => Ok(Self::F64), 13 => Ok(Self::C64), 14 => Ok(Self::C128), - n => Err(MohuError::UnknownDType(format!("dtype code {n}"))), + n => Err(MohuError::UnknownDType(format!("dtype code {n}"))), } } @@ -461,38 +492,27 @@ impl DType { /// and common aliases. Case-insensitive except for uppercase char codes /// (B, H, I, L, F, D). pub fn from_str(s: &str) -> MohuResult { + Self::parse_str(s) + } + + fn parse_str(s: &str) -> MohuResult { let lower = s.trim().to_ascii_lowercase(); match lower.as_str() { - "bool" | "bool_" | "?" | "b1" - => Ok(Self::Bool), - "int8" | "i1" | "b" | "byte" - => Ok(Self::I8), - "int16" | "i2" | "h" | "short" - => Ok(Self::I16), - "int32" | "i4" | "i" | "int" | "int_" - => Ok(Self::I32), - "int64" | "i8" | "l" | "long" | "longlong" - => Ok(Self::I64), - "uint8" | "u1" | "ubyte" - => Ok(Self::U8), - "uint16" | "u2" | "ushort" - => Ok(Self::U16), - "uint32" | "u4" | "uint" - => Ok(Self::U32), - "uint64" | "u8" | "ulong" - => Ok(Self::U64), - "float16" | "f2" | "half" | "f16" - => Ok(Self::F16), - "bfloat16" | "bf16" - => Ok(Self::BF16), - "float32" | "f4" | "float" | "single" | "f32" - => Ok(Self::F32), - "float64" | "f8" | "double" | "f64" - => Ok(Self::F64), - "complex64" | "c8" | "csingle" | "c64" - => Ok(Self::C64), - "complex128" | "c16" | "cdouble" | "c128" - => Ok(Self::C128), + "bool" | "bool_" | "?" | "b1" => Ok(Self::Bool), + "int8" | "i1" | "b" | "byte" => Ok(Self::I8), + "int16" | "i2" | "h" | "short" => Ok(Self::I16), + "int32" | "i4" | "i" | "int" | "int_" => Ok(Self::I32), + "int64" | "i8" | "l" | "long" | "longlong" => Ok(Self::I64), + "uint8" | "u1" | "ubyte" => Ok(Self::U8), + "uint16" | "u2" | "ushort" => Ok(Self::U16), + "uint32" | "u4" | "uint" => Ok(Self::U32), + "uint64" | "u8" | "ulong" => Ok(Self::U64), + "float16" | "f2" | "half" | "f16" => Ok(Self::F16), + "bfloat16" | "bf16" => Ok(Self::BF16), + "float32" | "f4" | "float" | "single" | "f32" => Ok(Self::F32), + "float64" | "f8" | "double" | "f64" => Ok(Self::F64), + "complex64" | "c8" | "csingle" | "c64" => Ok(Self::C64), + "complex128" | "c16" | "cdouble" | "c128" => Ok(Self::C128), _ => { // Uppercase single-char codes survive the lowercase pass // only if the original string was uppercase. @@ -505,7 +525,7 @@ impl DType { "D" => Ok(Self::C128), other => Err(MohuError::UnknownDType(other.to_string())), } - } + }, } } @@ -523,16 +543,28 @@ impl fmt::Display for DType { } } +impl std::str::FromStr for DType { + type Err = MohuError; + + fn from_str(s: &str) -> Result { + DType::parse_str(s) + } +} + // ─── TryFrom ───────────────────────────────────────────────────────────────── impl TryFrom<&str> for DType { type Error = MohuError; - fn try_from(s: &str) -> MohuResult { DType::from_str(s) } + fn try_from(s: &str) -> MohuResult { + DType::parse_str(s) + } } impl TryFrom for DType { type Error = MohuError; - fn try_from(s: String) -> MohuResult { DType::from_str(&s) } + fn try_from(s: String) -> MohuResult { + DType::parse_str(&s) + } } // ─── serde (feature-gated) ─────────────────────────────────────────────────── @@ -548,7 +580,7 @@ impl serde::Serialize for DType { impl<'de> serde::Deserialize<'de> for DType { fn deserialize>(d: D) -> Result { let s = String::deserialize(d)?; - DType::from_str(&s).map_err(serde::de::Error::custom) + DType::parse_str(&s).map_err(serde::de::Error::custom) } } diff --git a/crates/mohu-dtype/src/finfo.rs b/crates/mohu-dtype/src/finfo.rs index 44b6f1b..bcc3a4b 100644 --- a/crates/mohu-dtype/src/finfo.rs +++ b/crates/mohu-dtype/src/finfo.rs @@ -77,10 +77,10 @@ impl FloatInfo { /// not a floating-point type. pub fn of(dtype: DType) -> MohuResult { match dtype { - DType::F16 => Ok(Self::f16()), + DType::F16 => Ok(Self::f16()), DType::BF16 => Ok(Self::bf16()), - DType::F32 => Ok(Self::f32()), - DType::F64 => Ok(Self::f64()), + DType::F32 => Ok(Self::f32()), + DType::F64 => Ok(Self::f64()), other => Err(MohuError::UnsupportedDType { op: "finfo", dtype: other.to_string(), @@ -93,28 +93,28 @@ impl FloatInfo { /// Returns `FloatInfo` for IEEE 754 binary16 (half-precision). pub fn f16() -> Self { // IEEE 754 binary16: sign=1, exp=5, mantissa=10 - let eps = half::f16::EPSILON.to_f64(); - let max = half::f16::MAX.to_f64(); - let tiny = half::f16::MIN_POSITIVE.to_f64(); + let eps = half::f16::EPSILON.to_f64(); + let max = half::f16::MAX.to_f64(); + let tiny = half::f16::MIN_POSITIVE.to_f64(); // Smallest subnormal: 2^{-24} - let smallest_sub = 5.960_464_477_539_063e-8_f64; - let precision = (10_u32 as f64 * f64::log10(2.0)).floor() as u32; // 3 + let smallest_sub = 5.960_464_477_539_063e-8_f64; + let precision = (10_f64 * f64::log10(2.0)).floor() as u32; // 3 Self { - dtype: DType::F16, - bits: 16, - nmant: 10, - nexp: 5, - maxexp: 16, - minexp: -13, + dtype: DType::F16, + bits: 16, + nmant: 10, + nexp: 5, + maxexp: 16, + minexp: -13, eps, max, - min: -max, + min: -max, tiny, - smallest_normal: tiny, + smallest_normal: tiny, smallest_subnormal: smallest_sub, precision, - resolution: 10f64.powi(-(precision as i32)), - epsneg: eps, + resolution: 10f64.powi(-(precision as i32)), + epsneg: eps, } } @@ -123,27 +123,27 @@ impl FloatInfo { /// Returns `FloatInfo` for Google Brain Float16 (bfloat16). pub fn bf16() -> Self { // Google Brain Float16: sign=1, exp=8, mantissa=7 - let eps = half::bf16::EPSILON.to_f64(); - let max = half::bf16::MAX.to_f64(); - let tiny = half::bf16::MIN_POSITIVE.to_f64(); - let smallest_sub = 9.183_549_615_799_121e-41_f64; - let precision = (7_u32 as f64 * f64::log10(2.0)).floor() as u32; // 2 + let eps = half::bf16::EPSILON.to_f64(); + let max = half::bf16::MAX.to_f64(); + let tiny = half::bf16::MIN_POSITIVE.to_f64(); + let smallest_sub = 9.183_549_615_799_121e-41_f64; + let precision = (7_f64 * f64::log10(2.0)).floor() as u32; // 2 Self { - dtype: DType::BF16, - bits: 16, - nmant: 7, - nexp: 8, - maxexp: 128, - minexp: -125, + dtype: DType::BF16, + bits: 16, + nmant: 7, + nexp: 8, + maxexp: 128, + minexp: -125, eps, max, - min: -max, + min: -max, tiny, - smallest_normal: tiny, + smallest_normal: tiny, smallest_subnormal: smallest_sub, precision, - resolution: 10f64.powi(-(precision as i32)), - epsneg: eps, + resolution: 10f64.powi(-(precision as i32)), + epsneg: eps, } } @@ -151,27 +151,27 @@ impl FloatInfo { /// Returns `FloatInfo` for IEEE 754 binary32 (single-precision). pub fn f32() -> Self { - let eps = f32::EPSILON as f64; - let max = f32::MAX as f64; - let tiny = f32::MIN_POSITIVE as f64; - let smallest_sub = 1.401_298_464_324_817e-45_f64; - let precision = (23_u32 as f64 * f64::log10(2.0)).floor() as u32; // 6 + let eps = f32::EPSILON as f64; + let max = f32::MAX as f64; + let tiny = f32::MIN_POSITIVE as f64; + let smallest_sub = 1.401_298_464_324_817e-45_f64; + let precision = (23_f64 * f64::log10(2.0)).floor() as u32; // 6 Self { - dtype: DType::F32, - bits: 32, - nmant: 23, - nexp: 8, - maxexp: 128, - minexp: -125, + dtype: DType::F32, + bits: 32, + nmant: 23, + nexp: 8, + maxexp: 128, + minexp: -125, eps, max, - min: -max, + min: -max, tiny, - smallest_normal: tiny, + smallest_normal: tiny, smallest_subnormal: smallest_sub, precision, - resolution: 10f64.powi(-(precision as i32)), - epsneg: eps, + resolution: 10f64.powi(-(precision as i32)), + epsneg: eps, } } @@ -179,27 +179,27 @@ impl FloatInfo { /// Returns `FloatInfo` for IEEE 754 binary64 (double-precision). pub fn f64() -> Self { - let eps = f64::EPSILON; - let max = f64::MAX; - let tiny = f64::MIN_POSITIVE; - let smallest_sub = 5.0e-324_f64; - let precision = (52_u32 as f64 * f64::log10(2.0)).floor() as u32; // 15 + let eps = f64::EPSILON; + let max = f64::MAX; + let tiny = f64::MIN_POSITIVE; + let smallest_sub = 5.0e-324_f64; + let precision = (52_f64 * f64::log10(2.0)).floor() as u32; // 15 Self { - dtype: DType::F64, - bits: 64, - nmant: 52, - nexp: 11, - maxexp: 1024, - minexp: -1021, + dtype: DType::F64, + bits: 64, + nmant: 52, + nexp: 11, + maxexp: 1024, + minexp: -1021, eps, max, - min: -max, + min: -max, tiny, - smallest_normal: tiny, + smallest_normal: tiny, smallest_subnormal: smallest_sub, precision, - resolution: 10f64.powi(-(precision as i32)), - epsneg: eps, + resolution: 10f64.powi(-(precision as i32)), + epsneg: eps, } } @@ -222,7 +222,9 @@ impl FloatInfo { /// Returns `true` if `|a - b| <= n_ulps * ULP(max(|a|, |b|))`. pub fn within_ulps(&self, a: f64, b: f64, n_ulps: u64) -> bool { let scale = a.abs().max(b.abs()); - if scale == 0.0 { return a == b; } + if scale == 0.0 { + return a == b; + } let ulp = self.ulp_at(scale); (a - b).abs() <= n_ulps as f64 * ulp } @@ -234,12 +236,12 @@ impl std::fmt::Display for FloatInfo { f, "FloatInfo({dtype}: bits={bits}, nmant={nmant}, nexp={nexp}, eps={eps:.3e}, max={max:.3e}, tiny={tiny:.3e})", dtype = self.dtype, - bits = self.bits, + bits = self.bits, nmant = self.nmant, - nexp = self.nexp, - eps = self.eps, - max = self.max, - tiny = self.tiny, + nexp = self.nexp, + eps = self.eps, + max = self.max, + tiny = self.tiny, ) } } diff --git a/crates/mohu-dtype/src/iinfo.rs b/crates/mohu-dtype/src/iinfo.rs index 3aba6ba..90d7eae 100644 --- a/crates/mohu-dtype/src/iinfo.rs +++ b/crates/mohu-dtype/src/iinfo.rs @@ -41,14 +41,14 @@ impl IntInfo { /// not an integer type. pub fn of(dtype: DType) -> MohuResult { match dtype { - DType::I8 => Ok(Self::i8()), - DType::I16 => Ok(Self::i16()), - DType::I32 => Ok(Self::i32()), - DType::I64 => Ok(Self::i64()), - DType::U8 => Ok(Self::u8()), - DType::U16 => Ok(Self::u16()), - DType::U32 => Ok(Self::u32()), - DType::U64 => Ok(Self::u64()), + DType::I8 => Ok(Self::i8()), + DType::I16 => Ok(Self::i16()), + DType::I32 => Ok(Self::i32()), + DType::I64 => Ok(Self::i64()), + DType::U8 => Ok(Self::u8()), + DType::U16 => Ok(Self::u16()), + DType::U32 => Ok(Self::u32()), + DType::U64 => Ok(Self::u64()), other => Err(MohuError::UnsupportedDType { op: "iinfo", dtype: other.to_string(), @@ -60,46 +60,86 @@ impl IntInfo { /// Returns `IntInfo` for `i8` (signed 8-bit integer). pub const fn i8() -> Self { - Self { dtype: DType::I8, bits: 8, is_signed: true, - min: i8::MIN as i128, max: i8::MAX as u128 } + Self { + dtype: DType::I8, + bits: 8, + is_signed: true, + min: i8::MIN as i128, + max: i8::MAX as u128, + } } /// Returns `IntInfo` for `i16` (signed 16-bit integer). pub const fn i16() -> Self { - Self { dtype: DType::I16, bits: 16, is_signed: true, - min: i16::MIN as i128, max: i16::MAX as u128 } + Self { + dtype: DType::I16, + bits: 16, + is_signed: true, + min: i16::MIN as i128, + max: i16::MAX as u128, + } } /// Returns `IntInfo` for `i32` (signed 32-bit integer). pub const fn i32() -> Self { - Self { dtype: DType::I32, bits: 32, is_signed: true, - min: i32::MIN as i128, max: i32::MAX as u128 } + Self { + dtype: DType::I32, + bits: 32, + is_signed: true, + min: i32::MIN as i128, + max: i32::MAX as u128, + } } /// Returns `IntInfo` for `i64` (signed 64-bit integer). pub const fn i64() -> Self { - Self { dtype: DType::I64, bits: 64, is_signed: true, - min: i64::MIN as i128, max: i64::MAX as u128 } + Self { + dtype: DType::I64, + bits: 64, + is_signed: true, + min: i64::MIN as i128, + max: i64::MAX as u128, + } } // ─── unsigned ────────────────────────────────────────────────────────────── /// Returns `IntInfo` for `u8` (unsigned 8-bit integer). pub const fn u8() -> Self { - Self { dtype: DType::U8, bits: 8, is_signed: false, - min: 0, max: u8::MAX as u128 } + Self { + dtype: DType::U8, + bits: 8, + is_signed: false, + min: 0, + max: u8::MAX as u128, + } } /// Returns `IntInfo` for `u16` (unsigned 16-bit integer). pub const fn u16() -> Self { - Self { dtype: DType::U16, bits: 16, is_signed: false, - min: 0, max: u16::MAX as u128 } + Self { + dtype: DType::U16, + bits: 16, + is_signed: false, + min: 0, + max: u16::MAX as u128, + } } /// Returns `IntInfo` for `u32` (unsigned 32-bit integer). pub const fn u32() -> Self { - Self { dtype: DType::U32, bits: 32, is_signed: false, - min: 0, max: u32::MAX as u128 } + Self { + dtype: DType::U32, + bits: 32, + is_signed: false, + min: 0, + max: u32::MAX as u128, + } } /// Returns `IntInfo` for `u64` (unsigned 64-bit integer). pub const fn u64() -> Self { - Self { dtype: DType::U64, bits: 64, is_signed: false, - min: 0, max: u64::MAX as u128 } + Self { + dtype: DType::U64, + bits: 64, + is_signed: false, + min: 0, + max: u64::MAX as u128, + } } // ─── convenience ─────────────────────────────────────────────────────────── @@ -122,17 +162,29 @@ impl IntInfo { /// Returns the minimum scalar type (smallest integer dtype) that can /// represent the given signed value without overflow. pub fn minimum_signed_type_for(v: i64) -> DType { - if v >= i8::MIN as i64 && v <= i8::MAX as i64 { return DType::I8; } - if v >= i16::MIN as i64 && v <= i16::MAX as i64 { return DType::I16; } - if v >= i32::MIN as i64 && v <= i32::MAX as i64 { return DType::I32; } + if v >= i8::MIN as i64 && v <= i8::MAX as i64 { + return DType::I8; + } + if v >= i16::MIN as i64 && v <= i16::MAX as i64 { + return DType::I16; + } + if v >= i32::MIN as i64 && v <= i32::MAX as i64 { + return DType::I32; + } DType::I64 } /// Returns the minimum scalar type that can represent the given unsigned value. pub fn minimum_unsigned_type_for(v: u64) -> DType { - if v <= u8::MAX as u64 { return DType::U8; } - if v <= u16::MAX as u64 { return DType::U16; } - if v <= u32::MAX as u64 { return DType::U32; } + if v <= u8::MAX as u64 { + return DType::U8; + } + if v <= u16::MAX as u64 { + return DType::U16; + } + if v <= u32::MAX as u64 { + return DType::U32; + } DType::U64 } } @@ -142,10 +194,10 @@ impl std::fmt::Display for IntInfo { write!( f, "IntInfo({dtype}: bits={bits}, min={min}, max={max}, signed={signed})", - dtype = self.dtype, - bits = self.bits, - min = self.min, - max = self.max, + dtype = self.dtype, + bits = self.bits, + min = self.min, + max = self.max, signed = self.is_signed, ) } diff --git a/crates/mohu-dtype/src/lib.rs b/crates/mohu-dtype/src/lib.rs index 5be2f3f..04ca005 100644 --- a/crates/mohu-dtype/src/lib.rs +++ b/crates/mohu-dtype/src/lib.rs @@ -40,11 +40,11 @@ pub mod macros; pub mod promote; pub mod scalar; -pub use dtype::{DType, ALL_DTYPES, DTYPE_COUNT}; +pub use dtype::{ALL_DTYPES, DTYPE_COUNT, DType}; pub use finfo::FloatInfo; pub use iinfo::IntInfo; pub use promote::{ - can_cast, common_type, minimum_scalar_type, promote, result_type, weak_promote, CastMode, + CastMode, can_cast, common_type, minimum_scalar_type, promote, result_type, weak_promote, }; pub use scalar::{ ComplexScalar, FloatScalar, IntScalar, RealScalar, Scalar, SignedScalar, UnsignedScalar, diff --git a/crates/mohu-dtype/src/macros.rs b/crates/mohu-dtype/src/macros.rs index 9220144..e551ee1 100644 --- a/crates/mohu-dtype/src/macros.rs +++ b/crates/mohu-dtype/src/macros.rs @@ -1,24 +1,24 @@ -/// Compile-time and runtime dispatch macros for scalar types. -/// -/// These macros are the foundation of how mohu achieves zero-overhead -/// generic dispatch over runtime `DType` values. Every hot-path kernel -/// in `mohu-buffer`, `mohu-array`, and `mohu-ops` uses `dispatch_dtype!` -/// to monomorphise a single generic function over the correct scalar type -/// without a vtable or allocation. -/// -/// # Macro reference -/// -/// | Macro | Purpose | -/// |-------|---------| -/// | [`dtype_of!`] | `DType` constant for a Rust type literal | -/// | [`dispatch_dtype!`] | runtime DType → monomorphised call | -/// | [`dispatch_numeric!`] | same, excluding `Bool` | -/// | [`dispatch_integer!`] | integers only | -/// | [`dispatch_float!`] | floats only (F16/BF16/F32/F64) | -/// | [`dispatch_real!`] | integers + real floats (no complex, no bool) | -/// | [`dispatch_signed!`] | signed integers + floats | -/// | [`for_each_dtype!`] | invoke a macro for every dtype (codegen helper) | -/// | [`assert_dtype!`] | assert a DType at runtime or return an error | +//! Compile-time and runtime dispatch macros for scalar types. +//! +//! These macros are the foundation of how mohu achieves zero-overhead +//! generic dispatch over runtime `DType` values. Every hot-path kernel +//! in `mohu-buffer`, `mohu-array`, and `mohu-ops` uses `dispatch_dtype!` +//! to monomorphise a single generic function over the correct scalar type +//! without a vtable or allocation. +//! +//! # Macro reference +//! +//! | Macro | Purpose | +//! |-------|---------| +//! | [`dtype_of!`] | `DType` constant for a Rust type literal | +//! | [`dispatch_dtype!`] | runtime DType → monomorphised call | +//! | [`dispatch_numeric!`] | same, excluding `Bool` | +//! | [`dispatch_integer!`] | integers only | +//! | [`dispatch_float!`] | floats only (F16/BF16/F32/F64) | +//! | [`dispatch_real!`] | integers + real floats (no complex, no bool) | +//! | [`dispatch_signed!`] | signed integers + floats | +//! | [`for_each_dtype!`] | invoke a macro for every dtype (codegen helper) | +//! | [`assert_dtype!`] | assert a DType at runtime or return an error | // ─── dtype_of! ─────────────────────────────────────────────────────────────── @@ -36,25 +36,63 @@ /// ``` #[macro_export] macro_rules! dtype_of { - (bool) => { $crate::dtype::DType::Bool }; - (i8) => { $crate::dtype::DType::I8 }; - (i16) => { $crate::dtype::DType::I16 }; - (i32) => { $crate::dtype::DType::I32 }; - (i64) => { $crate::dtype::DType::I64 }; - (u8) => { $crate::dtype::DType::U8 }; - (u16) => { $crate::dtype::DType::U16 }; - (u32) => { $crate::dtype::DType::U32 }; - (u64) => { $crate::dtype::DType::U64 }; - (f16) => { $crate::dtype::DType::F16 }; - (::half::f16) => { $crate::dtype::DType::F16 }; - (bf16) => { $crate::dtype::DType::BF16 }; - (::half::bf16) => { $crate::dtype::DType::BF16 }; - (f32) => { $crate::dtype::DType::F32 }; - (f64) => { $crate::dtype::DType::F64 }; - (Complex) => { $crate::dtype::DType::C64 }; - (::num_complex::Complex) => { $crate::dtype::DType::C64 }; - (Complex) => { $crate::dtype::DType::C128 }; - (::num_complex::Complex) => { $crate::dtype::DType::C128 }; + (bool) => { + $crate::dtype::DType::Bool + }; + (i8) => { + $crate::dtype::DType::I8 + }; + (i16) => { + $crate::dtype::DType::I16 + }; + (i32) => { + $crate::dtype::DType::I32 + }; + (i64) => { + $crate::dtype::DType::I64 + }; + (u8) => { + $crate::dtype::DType::U8 + }; + (u16) => { + $crate::dtype::DType::U16 + }; + (u32) => { + $crate::dtype::DType::U32 + }; + (u64) => { + $crate::dtype::DType::U64 + }; + (f16) => { + $crate::dtype::DType::F16 + }; + (::half::f16) => { + $crate::dtype::DType::F16 + }; + (bf16) => { + $crate::dtype::DType::BF16 + }; + (::half::bf16) => { + $crate::dtype::DType::BF16 + }; + (f32) => { + $crate::dtype::DType::F32 + }; + (f64) => { + $crate::dtype::DType::F64 + }; + (Complex) => { + $crate::dtype::DType::C64 + }; + (::num_complex::Complex) => { + $crate::dtype::DType::C64 + }; + (Complex) => { + $crate::dtype::DType::C128 + }; + (::num_complex::Complex) => { + $crate::dtype::DType::C128 + }; } // ─── dispatch_dtype! ───────────────────────────────────────────────────────── @@ -158,19 +196,19 @@ macro_rules! dispatch_numeric { op: "numeric dispatch", dtype: "bool".to_string(), }), - $crate::dtype::DType::I8 => Ok($macro!(i8)), - $crate::dtype::DType::I16 => Ok($macro!(i16)), - $crate::dtype::DType::I32 => Ok($macro!(i32)), - $crate::dtype::DType::I64 => Ok($macro!(i64)), - $crate::dtype::DType::U8 => Ok($macro!(u8)), - $crate::dtype::DType::U16 => Ok($macro!(u16)), - $crate::dtype::DType::U32 => Ok($macro!(u32)), - $crate::dtype::DType::U64 => Ok($macro!(u64)), - $crate::dtype::DType::F16 => Ok($macro!(::half::f16)), + $crate::dtype::DType::I8 => Ok($macro!(i8)), + $crate::dtype::DType::I16 => Ok($macro!(i16)), + $crate::dtype::DType::I32 => Ok($macro!(i32)), + $crate::dtype::DType::I64 => Ok($macro!(i64)), + $crate::dtype::DType::U8 => Ok($macro!(u8)), + $crate::dtype::DType::U16 => Ok($macro!(u16)), + $crate::dtype::DType::U32 => Ok($macro!(u32)), + $crate::dtype::DType::U64 => Ok($macro!(u64)), + $crate::dtype::DType::F16 => Ok($macro!(::half::f16)), $crate::dtype::DType::BF16 => Ok($macro!(::half::bf16)), - $crate::dtype::DType::F32 => Ok($macro!(f32)), - $crate::dtype::DType::F64 => Ok($macro!(f64)), - $crate::dtype::DType::C64 => Ok($macro!(::num_complex::Complex)), + $crate::dtype::DType::F32 => Ok($macro!(f32)), + $crate::dtype::DType::F64 => Ok($macro!(f64)), + $crate::dtype::DType::C64 => Ok($macro!(::num_complex::Complex)), $crate::dtype::DType::C128 => Ok($macro!(::num_complex::Complex)), } }; @@ -185,14 +223,14 @@ macro_rules! dispatch_numeric { macro_rules! dispatch_integer { ($dtype:expr, $macro:ident) => { match $dtype { - $crate::dtype::DType::I8 => Ok($macro!(i8)), - $crate::dtype::DType::I16 => Ok($macro!(i16)), - $crate::dtype::DType::I32 => Ok($macro!(i32)), - $crate::dtype::DType::I64 => Ok($macro!(i64)), - $crate::dtype::DType::U8 => Ok($macro!(u8)), - $crate::dtype::DType::U16 => Ok($macro!(u16)), - $crate::dtype::DType::U32 => Ok($macro!(u32)), - $crate::dtype::DType::U64 => Ok($macro!(u64)), + $crate::dtype::DType::I8 => Ok($macro!(i8)), + $crate::dtype::DType::I16 => Ok($macro!(i16)), + $crate::dtype::DType::I32 => Ok($macro!(i32)), + $crate::dtype::DType::I64 => Ok($macro!(i64)), + $crate::dtype::DType::U8 => Ok($macro!(u8)), + $crate::dtype::DType::U16 => Ok($macro!(u16)), + $crate::dtype::DType::U32 => Ok($macro!(u32)), + $crate::dtype::DType::U64 => Ok($macro!(u64)), other => Err($crate::MohuError::UnsupportedDType { op: "integer dispatch", dtype: other.to_string(), @@ -210,10 +248,10 @@ macro_rules! dispatch_integer { macro_rules! dispatch_float { ($dtype:expr, $macro:ident) => { match $dtype { - $crate::dtype::DType::F16 => Ok($macro!(::half::f16)), + $crate::dtype::DType::F16 => Ok($macro!(::half::f16)), $crate::dtype::DType::BF16 => Ok($macro!(::half::bf16)), - $crate::dtype::DType::F32 => Ok($macro!(f32)), - $crate::dtype::DType::F64 => Ok($macro!(f64)), + $crate::dtype::DType::F32 => Ok($macro!(f32)), + $crate::dtype::DType::F64 => Ok($macro!(f64)), other => Err($crate::MohuError::UnsupportedDType { op: "float dispatch", dtype: other.to_string(), @@ -231,18 +269,18 @@ macro_rules! dispatch_float { macro_rules! dispatch_real { ($dtype:expr, $macro:ident) => { match $dtype { - $crate::dtype::DType::I8 => Ok($macro!(i8)), - $crate::dtype::DType::I16 => Ok($macro!(i16)), - $crate::dtype::DType::I32 => Ok($macro!(i32)), - $crate::dtype::DType::I64 => Ok($macro!(i64)), - $crate::dtype::DType::U8 => Ok($macro!(u8)), - $crate::dtype::DType::U16 => Ok($macro!(u16)), - $crate::dtype::DType::U32 => Ok($macro!(u32)), - $crate::dtype::DType::U64 => Ok($macro!(u64)), - $crate::dtype::DType::F16 => Ok($macro!(::half::f16)), + $crate::dtype::DType::I8 => Ok($macro!(i8)), + $crate::dtype::DType::I16 => Ok($macro!(i16)), + $crate::dtype::DType::I32 => Ok($macro!(i32)), + $crate::dtype::DType::I64 => Ok($macro!(i64)), + $crate::dtype::DType::U8 => Ok($macro!(u8)), + $crate::dtype::DType::U16 => Ok($macro!(u16)), + $crate::dtype::DType::U32 => Ok($macro!(u32)), + $crate::dtype::DType::U64 => Ok($macro!(u64)), + $crate::dtype::DType::F16 => Ok($macro!(::half::f16)), $crate::dtype::DType::BF16 => Ok($macro!(::half::bf16)), - $crate::dtype::DType::F32 => Ok($macro!(f32)), - $crate::dtype::DType::F64 => Ok($macro!(f64)), + $crate::dtype::DType::F32 => Ok($macro!(f32)), + $crate::dtype::DType::F64 => Ok($macro!(f64)), other => Err($crate::MohuError::UnsupportedDType { op: "real dispatch", dtype: other.to_string(), @@ -260,14 +298,14 @@ macro_rules! dispatch_real { macro_rules! dispatch_signed { ($dtype:expr, $macro:ident) => { match $dtype { - $crate::dtype::DType::I8 => Ok($macro!(i8)), - $crate::dtype::DType::I16 => Ok($macro!(i16)), - $crate::dtype::DType::I32 => Ok($macro!(i32)), - $crate::dtype::DType::I64 => Ok($macro!(i64)), - $crate::dtype::DType::F16 => Ok($macro!(::half::f16)), + $crate::dtype::DType::I8 => Ok($macro!(i8)), + $crate::dtype::DType::I16 => Ok($macro!(i16)), + $crate::dtype::DType::I32 => Ok($macro!(i32)), + $crate::dtype::DType::I64 => Ok($macro!(i64)), + $crate::dtype::DType::F16 => Ok($macro!(::half::f16)), $crate::dtype::DType::BF16 => Ok($macro!(::half::bf16)), - $crate::dtype::DType::F32 => Ok($macro!(f32)), - $crate::dtype::DType::F64 => Ok($macro!(f64)), + $crate::dtype::DType::F32 => Ok($macro!(f32)), + $crate::dtype::DType::F64 => Ok($macro!(f64)), other => Err($crate::MohuError::UnsupportedDType { op: "signed dispatch", dtype: other.to_string(), @@ -297,21 +335,21 @@ macro_rules! dispatch_signed { #[macro_export] macro_rules! for_each_dtype { ($macro:ident) => { - $macro!(bool, $crate::dtype::DType::Bool); - $macro!(i8, $crate::dtype::DType::I8); - $macro!(i16, $crate::dtype::DType::I16); - $macro!(i32, $crate::dtype::DType::I32); - $macro!(i64, $crate::dtype::DType::I64); - $macro!(u8, $crate::dtype::DType::U8); - $macro!(u16, $crate::dtype::DType::U16); - $macro!(u32, $crate::dtype::DType::U32); - $macro!(u64, $crate::dtype::DType::U64); - $macro!(::half::f16, $crate::dtype::DType::F16); - $macro!(::half::bf16, $crate::dtype::DType::BF16); - $macro!(f32, $crate::dtype::DType::F32); - $macro!(f64, $crate::dtype::DType::F64); - $macro!(::num_complex::Complex, $crate::dtype::DType::C64); - $macro!(::num_complex::Complex, $crate::dtype::DType::C128); + $macro!(bool, $crate::dtype::DType::Bool); + $macro!(i8, $crate::dtype::DType::I8); + $macro!(i16, $crate::dtype::DType::I16); + $macro!(i32, $crate::dtype::DType::I32); + $macro!(i64, $crate::dtype::DType::I64); + $macro!(u8, $crate::dtype::DType::U8); + $macro!(u16, $crate::dtype::DType::U16); + $macro!(u32, $crate::dtype::DType::U32); + $macro!(u64, $crate::dtype::DType::U64); + $macro!(::half::f16, $crate::dtype::DType::F16); + $macro!(::half::bf16, $crate::dtype::DType::BF16); + $macro!(f32, $crate::dtype::DType::F32); + $macro!(f64, $crate::dtype::DType::F64); + $macro!(::num_complex::Complex, $crate::dtype::DType::C64); + $macro!(::num_complex::Complex, $crate::dtype::DType::C128); }; } @@ -335,7 +373,7 @@ macro_rules! assert_dtype { if $actual != $expected { return Err($crate::MohuError::DTypeMismatch { expected: $expected.to_string(), - got: $actual.to_string(), + got: $actual.to_string(), }); } }; @@ -349,7 +387,7 @@ macro_rules! require_float { ($dtype:expr, $op:expr) => { if !$dtype.is_float() && !$dtype.is_complex() { return Err($crate::MohuError::UnsupportedDType { - op: $op, + op: $op, dtype: $dtype.to_string(), }); } @@ -364,7 +402,7 @@ macro_rules! require_numeric { ($dtype:expr, $op:expr) => { if $dtype.is_bool() { return Err($crate::MohuError::UnsupportedDType { - op: $op, + op: $op, dtype: "bool".to_string(), }); } @@ -379,7 +417,7 @@ macro_rules! require_real { ($dtype:expr, $op:expr) => { if $dtype.is_complex() || $dtype.is_bool() { return Err($crate::MohuError::UnsupportedDType { - op: $op, + op: $op, dtype: $dtype.to_string(), }); } diff --git a/crates/mohu-dtype/src/promote.rs b/crates/mohu-dtype/src/promote.rs index f4c58ef..176f0f0 100644 --- a/crates/mohu-dtype/src/promote.rs +++ b/crates/mohu-dtype/src/promote.rs @@ -20,7 +20,7 @@ /// | Safe | No information loss guaranteed (i8→i16, f32→f64) | /// | SameKind | Within same kind, precision loss OK (f64→f32) | /// | Unsafe | Any cast, including float→int, complex→real | -use crate::dtype::{DType, DTYPE_COUNT}; +use crate::dtype::{DTYPE_COUNT, DType}; // ─── CastMode ──────────────────────────────────────────────────────────────── @@ -69,7 +69,7 @@ pub const fn promote(a: DType, b: DType) -> DType { /// stored as its `u8` discriminant and cast back via `DType::from_u8` at /// lookup time. const PROMOTION_TABLE: [DType; DTYPE_COUNT * DTYPE_COUNT] = { - use DType::{Bool,I8,I16,I32,I64,U8,U16,U32,U64,F16,BF16,F32,F64,C64,C128}; + use DType::{BF16, Bool, C64, C128, F16, F32, F64, I8, I16, I32, I64, U8, U16, U32, U64}; // Row-major layout, symmetric. // Index mapping: Bool=0,I8=1,I16=2,I32=3,I64=4,U8=5,U16=6,U32=7,U64=8, @@ -78,21 +78,23 @@ const PROMOTION_TABLE: [DType; DTYPE_COUNT * DTYPE_COUNT] = { // Read as: promote(ROW, COL) = entry [ // Bool I8 I16 I32 I64 U8 U16 U32 U64 F16 BF16 F32 F64 C64 C128 - /* Bool */ Bool, I8, I16, I32, I64, U8, U16, U32, U64, F16, BF16, F32, F64, C64, C128, - /* I8 */ I8, I8, I16, I32, I64, I16, I32, I64, F64, F32, BF16, F32, F64, C64, C128, - /* I16 */ I16, I16, I16, I32, I64, I16, I32, I64, F64, F32, F32, F32, F64, C64, C128, - /* I32 */ I32, I32, I32, I32, I64, I32, I32, I64, F64, F64, F64, F64, F64, C128, C128, - /* I64 */ I64, I64, I64, I64, I64, I64, I64, I64, F64, F64, F64, F64, F64, C128, C128, - /* U8 */ U8, I16, I16, I32, I64, U8, U16, U32, U64, F16, BF16, F32, F64, C64, C128, - /* U16 */ U16, I32, I32, I32, I64, U16, U16, U32, U64, F32, F32, F32, F64, C64, C128, - /* U32 */ U32, I64, I64, I64, I64, U32, U32, U32, U64, F64, F64, F64, F64, C128, C128, - /* U64 */ U64, F64, F64, F64, F64, U64, U64, U64, U64, F64, F64, F64, F64, C128, C128, - /* F16 */ F16, F32, F32, F64, F64, F16, F32, F64, F64, F16, F32, F32, F64, C64, C128, - /* BF16 */ BF16, BF16, F32, F64, F64, BF16, F32, F64, F64, F32, BF16, F32, F64, C64, C128, - /* F32 */ F32, F32, F32, F64, F64, F32, F32, F64, F64, F32, F32, F32, F64, C64, C128, - /* F64 */ F64, F64, F64, F64, F64, F64, F64, F64, F64, F64, F64, F64, F64, C128, C128, - /* C64 */ C64, C64, C64, C128, C128, C64, C64, C128, C128, C64, C64, C64, C128, C64, C128, - /* C128 */ C128, C128, C128, C128, C128, C128, C128, C128, C128, C128, C128, C128, C128, C128, C128, + /* Bool */ + Bool, I8, I16, I32, I64, U8, U16, U32, U64, F16, BF16, F32, F64, C64, C128, + /* I8 */ I8, I8, I16, I32, I64, I16, I32, I64, F64, F32, BF16, F32, F64, C64, C128, + /* I16 */ I16, I16, I16, I32, I64, I16, I32, I64, F64, F32, F32, F32, F64, C64, C128, + /* I32 */ I32, I32, I32, I32, I64, I32, I32, I64, F64, F64, F64, F64, F64, C128, + C128, /* I64 */ I64, I64, I64, I64, I64, I64, I64, I64, F64, F64, F64, F64, F64, + C128, C128, /* U8 */ U8, I16, I16, I32, I64, U8, U16, U32, U64, F16, BF16, F32, F64, + C64, C128, /* U16 */ U16, I32, I32, I32, I64, U16, U16, U32, U64, F32, F32, F32, F64, + C64, C128, /* U32 */ U32, I64, I64, I64, I64, U32, U32, U32, U64, F64, F64, F64, F64, + C128, C128, /* U64 */ U64, F64, F64, F64, F64, U64, U64, U64, U64, F64, F64, F64, + F64, C128, C128, /* F16 */ F16, F32, F32, F64, F64, F16, F32, F64, F64, F16, F32, + F32, F64, C64, C128, /* BF16 */ BF16, BF16, F32, F64, F64, BF16, F32, F64, F64, F32, + BF16, F32, F64, C64, C128, /* F32 */ F32, F32, F32, F64, F64, F32, F32, F64, F64, + F32, F32, F32, F64, C64, C128, /* F64 */ F64, F64, F64, F64, F64, F64, F64, F64, F64, + F64, F64, F64, F64, C128, C128, /* C64 */ C64, C64, C64, C128, C128, C64, C64, C128, + C128, C64, C64, C64, C128, C64, C128, /* C128 */ C128, C128, C128, C128, C128, C128, + C128, C128, C128, C128, C128, C128, C128, C128, C128, ] }; @@ -109,9 +111,9 @@ const PROMOTION_TABLE: [DType; DTYPE_COUNT * DTYPE_COUNT] = { /// ``` pub const fn can_cast(from: DType, to: DType, mode: CastMode) -> bool { match mode { - CastMode::Safe => SAFE_CAST_TABLE[from as usize * DTYPE_COUNT + to as usize], + CastMode::Safe => SAFE_CAST_TABLE[from as usize * DTYPE_COUNT + to as usize], CastMode::SameKind => SAMEKIND_CAST_TABLE[from as usize * DTYPE_COUNT + to as usize], - CastMode::Unsafe => true, // all casts are valid in unsafe mode + CastMode::Unsafe => true, // all casts are valid in unsafe mode } } @@ -135,21 +137,16 @@ const SAFE_CAST_TABLE: [bool; DTYPE_COUNT * DTYPE_COUNT] = { // Bool I8 I16 I32 I64 U8 U16 U32 U64 F16 BF16 F32 F64 C64 C128 [ - /* Bool */ T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, - /* I8 */ F, T, T, T, T, F, F, F, F, T, T, T, T, T, T, - /* I16 */ F, F, T, T, T, F, F, F, F, F, F, T, T, T, T, - /* I32 */ F, F, F, T, T, F, F, F, F, F, F, F, T, F, T, - /* I64 */ F, F, F, F, T, F, F, F, F, F, F, F, T, F, T, - /* U8 */ F, F, T, T, T, T, T, T, T, T, T, T, T, T, T, - /* U16 */ F, F, F, T, T, F, T, T, T, F, F, T, T, T, T, - /* U32 */ F, F, F, F, T, F, F, T, T, F, F, F, T, F, T, - /* U64 */ F, F, F, F, F, F, F, F, T, F, F, F, T, F, T, - /* F16 */ F, F, F, F, F, F, F, F, F, T, F, T, T, T, T, - /* BF16 */ F, F, F, F, F, F, F, F, F, F, T, T, T, T, T, - /* F32 */ F, F, F, F, F, F, F, F, F, F, F, T, T, T, T, - /* F64 */ F, F, F, F, F, F, F, F, F, F, F, F, T, F, T, - /* C64 */ F, F, F, F, F, F, F, F, F, F, F, F, F, T, T, - /* C128 */ F, F, F, F, F, F, F, F, F, F, F, F, F, F, T, + /* Bool */ T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, /* I8 */ F, T, T, T, T, + F, F, F, F, T, T, T, T, T, T, /* I16 */ F, F, T, T, T, F, F, F, F, F, F, T, T, T, T, + /* I32 */ F, F, F, T, T, F, F, F, F, F, F, F, T, F, T, /* I64 */ F, F, F, F, T, + F, F, F, F, F, F, F, T, F, T, /* U8 */ F, F, T, T, T, T, T, T, T, T, T, T, T, T, T, + /* U16 */ F, F, F, T, T, F, T, T, T, F, F, T, T, T, T, /* U32 */ F, F, F, F, T, + F, F, T, T, F, F, F, T, F, T, /* U64 */ F, F, F, F, F, F, F, F, T, F, F, F, T, F, T, + /* F16 */ F, F, F, F, F, F, F, F, F, T, F, T, T, T, T, /* BF16 */ F, F, F, F, F, + F, F, F, F, F, T, T, T, T, T, /* F32 */ F, F, F, F, F, F, F, F, F, F, F, T, T, T, T, + /* F64 */ F, F, F, F, F, F, F, F, F, F, F, F, T, F, T, /* C64 */ F, F, F, F, F, + F, F, F, F, F, F, F, F, T, T, /* C128 */ F, F, F, F, F, F, F, F, F, F, F, F, F, F, T, ] }; @@ -165,21 +162,16 @@ const SAMEKIND_CAST_TABLE: [bool; DTYPE_COUNT * DTYPE_COUNT] = { // Bool I8 I16 I32 I64 U8 U16 U32 U64 F16 BF16 F32 F64 C64 C128 [ - /* Bool */ T, F, F, F, F, F, F, F, F, F, F, F, F, F, F, - /* I8 */ F, T, T, T, T, T, T, T, T, F, F, F, F, F, F, - /* I16 */ F, T, T, T, T, T, T, T, T, F, F, F, F, F, F, - /* I32 */ F, T, T, T, T, T, T, T, T, F, F, F, F, F, F, - /* I64 */ F, T, T, T, T, T, T, T, T, F, F, F, F, F, F, - /* U8 */ F, T, T, T, T, T, T, T, T, F, F, F, F, F, F, - /* U16 */ F, T, T, T, T, T, T, T, T, F, F, F, F, F, F, - /* U32 */ F, T, T, T, T, T, T, T, T, F, F, F, F, F, F, - /* U64 */ F, T, T, T, T, T, T, T, T, F, F, F, F, F, F, - /* F16 */ F, F, F, F, F, F, F, F, F, T, T, T, T, F, F, - /* BF16 */ F, F, F, F, F, F, F, F, F, T, T, T, T, F, F, - /* F32 */ F, F, F, F, F, F, F, F, F, T, T, T, T, F, F, - /* F64 */ F, F, F, F, F, F, F, F, F, T, T, T, T, F, F, - /* C64 */ F, F, F, F, F, F, F, F, F, F, F, F, F, T, T, - /* C128 */ F, F, F, F, F, F, F, F, F, F, F, F, F, T, T, + /* Bool */ T, F, F, F, F, F, F, F, F, F, F, F, F, F, F, /* I8 */ F, T, T, T, T, + T, T, T, T, F, F, F, F, F, F, /* I16 */ F, T, T, T, T, T, T, T, T, F, F, F, F, F, F, + /* I32 */ F, T, T, T, T, T, T, T, T, F, F, F, F, F, F, /* I64 */ F, T, T, T, T, + T, T, T, T, F, F, F, F, F, F, /* U8 */ F, T, T, T, T, T, T, T, T, F, F, F, F, F, F, + /* U16 */ F, T, T, T, T, T, T, T, T, F, F, F, F, F, F, /* U32 */ F, T, T, T, T, + T, T, T, T, F, F, F, F, F, F, /* U64 */ F, T, T, T, T, T, T, T, T, F, F, F, F, F, F, + /* F16 */ F, F, F, F, F, F, F, F, F, T, T, T, T, F, F, /* BF16 */ F, F, F, F, F, + F, F, F, F, T, T, T, T, F, F, /* F32 */ F, F, F, F, F, F, F, F, F, T, T, T, T, F, F, + /* F64 */ F, F, F, F, F, F, F, F, F, T, T, T, T, F, F, /* C64 */ F, F, F, F, F, + F, F, F, F, F, F, F, F, T, T, /* C128 */ F, F, F, F, F, F, F, F, F, F, F, F, F, T, T, ] }; @@ -242,21 +234,37 @@ pub fn minimum_scalar_type(v: f64) -> DType { if v.fract() == 0.0 && v.is_finite() { if v >= 0.0 && v <= u64::MAX as f64 { let u = v as u64; - if u <= u8::MAX as u64 { return DType::U8; } - if u <= u16::MAX as u64 { return DType::U16; } - if u <= u32::MAX as u64 { return DType::U32; } + if u <= u8::MAX as u64 { + return DType::U8; + } + if u <= u16::MAX as u64 { + return DType::U16; + } + if u <= u32::MAX as u64 { + return DType::U32; + } return DType::U64; } else if v < 0.0 && v >= i64::MIN as f64 { let i = v as i64; - if i >= i8::MIN as i64 { return DType::I8; } - if i >= i16::MIN as i64 { return DType::I16; } - if i >= i32::MIN as i64 { return DType::I32; } + if i >= i8::MIN as i64 { + return DType::I8; + } + if i >= i16::MIN as i64 { + return DType::I16; + } + if i >= i32::MIN as i64 { + return DType::I32; + } return DType::I64; } // Value is an integer but too large for i64/u64 — fall through to float. } // Non-integer: check if f32 can represent it faithfully. - if (v as f32) as f64 == v { DType::F32 } else { DType::F64 } + if (v as f32) as f64 == v { + DType::F32 + } else { + DType::F64 + } } // ─── weak_promote (NumPy 2.0 style) ───────────────────────────────────────── diff --git a/crates/mohu-dtype/src/scalar.rs b/crates/mohu-dtype/src/scalar.rs index f3eda38..e895b45 100644 --- a/crates/mohu-dtype/src/scalar.rs +++ b/crates/mohu-dtype/src/scalar.rs @@ -23,25 +23,25 @@ use crate::dtype::DType; // ─── sealing mechanism ─────────────────────────────────────────────────────── mod private { + use half::{bf16, f16}; use num_complex::Complex; - use half::{f16, bf16}; /// Private marker trait — only types in this module can implement `Scalar`. pub trait Sealed {} impl Sealed for bool {} - impl Sealed for i8 {} + impl Sealed for i8 {} impl Sealed for i16 {} impl Sealed for i32 {} impl Sealed for i64 {} - impl Sealed for u8 {} + impl Sealed for u8 {} impl Sealed for u16 {} impl Sealed for u32 {} impl Sealed for u64 {} - impl Sealed for f16 {} + impl Sealed for f16 {} impl Sealed for bf16 {} - impl Sealed for f32 {} - impl Sealed for f64 {} + impl Sealed for f32 {} + impl Sealed for f64 {} impl Sealed for Complex {} impl Sealed for Complex {} } @@ -338,36 +338,122 @@ macro_rules! impl_scalar_int { $ty:ty, $dtype:expr, $zero:expr, $one:expr ) => { impl Scalar for $ty { - const DTYPE: DType = $dtype; - const ZERO: $ty = $zero; - const ONE: $ty = $one; - const ITEMSIZE: usize = std::mem::size_of::<$ty>(); - - #[inline] fn to_f64_lossy(self) -> f64 { self as f64 } - #[inline] fn from_f64_lossy(v: f64) -> $ty { v as $ty } + const DTYPE: DType = $dtype; + const ZERO: $ty = $zero; + const ONE: $ty = $one; + const ITEMSIZE: usize = std::mem::size_of::<$ty>(); + + #[inline] + fn to_f64_lossy(self) -> f64 { + self as f64 + } + #[inline] + fn from_f64_lossy(v: f64) -> $ty { + v as $ty + } } }; } -impl_scalar_int!(i8, DType::I8, 0i8, 1i8); +impl_scalar_int!(i8, DType::I8, 0i8, 1i8); impl_scalar_int!(i16, DType::I16, 0i16, 1i16); impl_scalar_int!(i32, DType::I32, 0i32, 1i32); impl_scalar_int!(i64, DType::I64, 0i64, 1i64); -impl_scalar_int!(u8, DType::U8, 0u8, 1u8); +impl_scalar_int!(u8, DType::U8, 0u8, 1u8); impl_scalar_int!(u16, DType::U16, 0u16, 1u16); impl_scalar_int!(u32, DType::U32, 0u32, 1u32); impl_scalar_int!(u64, DType::U64, 0u64, 1u64); // ─── RealScalar for integers ───────────────────────────────────────────────── -impl RealScalar for i8 { fn min_value()->Self{i8::MIN} fn max_value()->Self{i8::MAX} fn abs(self)->Self{self.wrapping_abs()} } -impl RealScalar for i16 { fn min_value()->Self{i16::MIN} fn max_value()->Self{i16::MAX} fn abs(self)->Self{self.wrapping_abs()} } -impl RealScalar for i32 { fn min_value()->Self{i32::MIN} fn max_value()->Self{i32::MAX} fn abs(self)->Self{self.wrapping_abs()} } -impl RealScalar for i64 { fn min_value()->Self{i64::MIN} fn max_value()->Self{i64::MAX} fn abs(self)->Self{self.wrapping_abs()} } -impl RealScalar for u8 { fn min_value()->Self{0} fn max_value()->Self{u8::MAX} fn abs(self)->Self{self} } -impl RealScalar for u16 { fn min_value()->Self{0} fn max_value()->Self{u16::MAX} fn abs(self)->Self{self} } -impl RealScalar for u32 { fn min_value()->Self{0} fn max_value()->Self{u32::MAX} fn abs(self)->Self{self} } -impl RealScalar for u64 { fn min_value()->Self{0} fn max_value()->Self{u64::MAX} fn abs(self)->Self{self} } +impl RealScalar for i8 { + fn min_value() -> Self { + i8::MIN + } + fn max_value() -> Self { + i8::MAX + } + fn abs(self) -> Self { + self.wrapping_abs() + } +} +impl RealScalar for i16 { + fn min_value() -> Self { + i16::MIN + } + fn max_value() -> Self { + i16::MAX + } + fn abs(self) -> Self { + self.wrapping_abs() + } +} +impl RealScalar for i32 { + fn min_value() -> Self { + i32::MIN + } + fn max_value() -> Self { + i32::MAX + } + fn abs(self) -> Self { + self.wrapping_abs() + } +} +impl RealScalar for i64 { + fn min_value() -> Self { + i64::MIN + } + fn max_value() -> Self { + i64::MAX + } + fn abs(self) -> Self { + self.wrapping_abs() + } +} +impl RealScalar for u8 { + fn min_value() -> Self { + 0 + } + fn max_value() -> Self { + u8::MAX + } + fn abs(self) -> Self { + self + } +} +impl RealScalar for u16 { + fn min_value() -> Self { + 0 + } + fn max_value() -> Self { + u16::MAX + } + fn abs(self) -> Self { + self + } +} +impl RealScalar for u32 { + fn min_value() -> Self { + 0 + } + fn max_value() -> Self { + u32::MAX + } + fn abs(self) -> Self { + self + } +} +impl RealScalar for u64 { + fn min_value() -> Self { + 0 + } + fn max_value() -> Self { + u64::MAX + } + fn abs(self) -> Self { + self + } +} // ─── IntScalar for all integer types ──────────────────────────────────────── @@ -375,17 +461,50 @@ macro_rules! impl_int_scalar { ($ty:ty) => { impl IntScalar for $ty { const BITS: u32 = <$ty>::BITS; - #[inline] fn overflowing_add(self, r: Self) -> (Self, bool) { <$ty>::overflowing_add(self, r) } - #[inline] fn overflowing_sub(self, r: Self) -> (Self, bool) { <$ty>::overflowing_sub(self, r) } - #[inline] fn overflowing_mul(self, r: Self) -> (Self, bool) { <$ty>::overflowing_mul(self, r) } - #[inline] fn saturating_add(self, r: Self) -> Self { <$ty>::saturating_add(self, r) } - #[inline] fn saturating_sub(self, r: Self) -> Self { <$ty>::saturating_sub(self, r) } - #[inline] fn checked_add(self, r: Self) -> Option { <$ty>::checked_add(self, r) } - #[inline] fn checked_sub(self, r: Self) -> Option { <$ty>::checked_sub(self, r) } - #[inline] fn count_ones(self) -> u32 { <$ty>::count_ones(self) } - #[inline] fn leading_zeros(self) -> u32 { <$ty>::leading_zeros(self) } - #[inline] fn trailing_zeros(self) -> u32 { <$ty>::trailing_zeros(self) } - #[inline] fn to_u64_bits(self) -> u64 { self as u64 } + #[inline] + fn overflowing_add(self, r: Self) -> (Self, bool) { + <$ty>::overflowing_add(self, r) + } + #[inline] + fn overflowing_sub(self, r: Self) -> (Self, bool) { + <$ty>::overflowing_sub(self, r) + } + #[inline] + fn overflowing_mul(self, r: Self) -> (Self, bool) { + <$ty>::overflowing_mul(self, r) + } + #[inline] + fn saturating_add(self, r: Self) -> Self { + <$ty>::saturating_add(self, r) + } + #[inline] + fn saturating_sub(self, r: Self) -> Self { + <$ty>::saturating_sub(self, r) + } + #[inline] + fn checked_add(self, r: Self) -> Option { + <$ty>::checked_add(self, r) + } + #[inline] + fn checked_sub(self, r: Self) -> Option { + <$ty>::checked_sub(self, r) + } + #[inline] + fn count_ones(self) -> u32 { + <$ty>::count_ones(self) + } + #[inline] + fn leading_zeros(self) -> u32 { + <$ty>::leading_zeros(self) + } + #[inline] + fn trailing_zeros(self) -> u32 { + <$ty>::trailing_zeros(self) + } + #[inline] + fn to_u64_bits(self) -> u64 { + self as u64 + } } }; } @@ -404,8 +523,14 @@ impl_int_scalar!(u64); macro_rules! impl_signed_scalar { ($ty:ty) => { impl SignedScalar for $ty { - #[inline] fn saturating_abs(self) -> Self { <$ty>::saturating_abs(self) } - #[inline] fn signum(self) -> Self { <$ty>::signum(self) } + #[inline] + fn saturating_abs(self) -> Self { + <$ty>::saturating_abs(self) + } + #[inline] + fn signum(self) -> Self { + <$ty>::signum(self) + } } }; } @@ -417,7 +542,7 @@ impl_signed_scalar!(i64); // ─── UnsignedScalar ────────────────────────────────────────────────────────── -impl UnsignedScalar for u8 {} +impl UnsignedScalar for u8 {} impl UnsignedScalar for u16 {} impl UnsignedScalar for u32 {} impl UnsignedScalar for u64 {} @@ -425,83 +550,137 @@ impl UnsignedScalar for u64 {} // ─── Bool ──────────────────────────────────────────────────────────────────── impl Scalar for bool { - const DTYPE: DType = DType::Bool; - const ZERO: bool = false; - const ONE: bool = true; + const DTYPE: DType = DType::Bool; + const ZERO: bool = false; + const ONE: bool = true; const ITEMSIZE: usize = 1; - #[inline] fn to_f64_lossy(self) -> f64 { self as u8 as f64 } - #[inline] fn from_f64_lossy(v: f64) -> bool { v != 0.0 } + #[inline] + fn to_f64_lossy(self) -> f64 { + self as u8 as f64 + } + #[inline] + fn from_f64_lossy(v: f64) -> bool { + v != 0.0 + } } // ─── Scalar for native floats ───────────────────────────────────────────────── impl Scalar for f32 { - const DTYPE: DType = DType::F32; - const ZERO: f32 = 0.0_f32; - const ONE: f32 = 1.0_f32; + const DTYPE: DType = DType::F32; + const ZERO: f32 = 0.0_f32; + const ONE: f32 = 1.0_f32; const ITEMSIZE: usize = 4; - #[inline] fn to_f64_lossy(self) -> f64 { self as f64 } - #[inline] fn from_f64_lossy(v: f64) -> f32 { v as f32 } + #[inline] + fn to_f64_lossy(self) -> f64 { + self as f64 + } + #[inline] + fn from_f64_lossy(v: f64) -> f32 { + v as f32 + } } impl Scalar for f64 { - const DTYPE: DType = DType::F64; - const ZERO: f64 = 0.0_f64; - const ONE: f64 = 1.0_f64; + const DTYPE: DType = DType::F64; + const ZERO: f64 = 0.0_f64; + const ONE: f64 = 1.0_f64; const ITEMSIZE: usize = 8; - #[inline] fn to_f64_lossy(self) -> f64 { self } - #[inline] fn from_f64_lossy(v: f64) -> f64 { v } + #[inline] + fn to_f64_lossy(self) -> f64 { + self + } + #[inline] + fn from_f64_lossy(v: f64) -> f64 { + v + } } // ─── Scalar for half-precision types ───────────────────────────────────────── impl Scalar for half::f16 { - const DTYPE: DType = DType::F16; - const ZERO: half::f16 = half::f16::ZERO; - const ONE: half::f16 = half::f16::ONE; - const ITEMSIZE: usize = 2; - - #[inline] fn to_f64_lossy(self) -> f64 { self.to_f64() } - #[inline] fn from_f64_lossy(v: f64) -> half::f16 { half::f16::from_f64(v) } + const DTYPE: DType = DType::F16; + const ZERO: half::f16 = half::f16::ZERO; + const ONE: half::f16 = half::f16::ONE; + const ITEMSIZE: usize = 2; + + #[inline] + fn to_f64_lossy(self) -> f64 { + self.to_f64() + } + #[inline] + fn from_f64_lossy(v: f64) -> half::f16 { + half::f16::from_f64(v) + } } impl Scalar for half::bf16 { - const DTYPE: DType = DType::BF16; - const ZERO: half::bf16 = half::bf16::ZERO; - const ONE: half::bf16 = half::bf16::ONE; - const ITEMSIZE: usize = 2; - - #[inline] fn to_f64_lossy(self) -> f64 { self.to_f64() } - #[inline] fn from_f64_lossy(v: f64) -> half::bf16 { half::bf16::from_f64(v) } + const DTYPE: DType = DType::BF16; + const ZERO: half::bf16 = half::bf16::ZERO; + const ONE: half::bf16 = half::bf16::ONE; + const ITEMSIZE: usize = 2; + + #[inline] + fn to_f64_lossy(self) -> f64 { + self.to_f64() + } + #[inline] + fn from_f64_lossy(v: f64) -> half::bf16 { + half::bf16::from_f64(v) + } } // ─── RealScalar for floats ──────────────────────────────────────────────────── impl RealScalar for f32 { - fn min_value() -> f32 { f32::MIN } - fn max_value() -> f32 { f32::MAX } - fn abs(self) -> f32 { f32::abs(self) } + fn min_value() -> f32 { + f32::MIN + } + fn max_value() -> f32 { + f32::MAX + } + fn abs(self) -> f32 { + f32::abs(self) + } } impl RealScalar for f64 { - fn min_value() -> f64 { f64::MIN } - fn max_value() -> f64 { f64::MAX } - fn abs(self) -> f64 { f64::abs(self) } + fn min_value() -> f64 { + f64::MIN + } + fn max_value() -> f64 { + f64::MAX + } + fn abs(self) -> f64 { + f64::abs(self) + } } impl RealScalar for half::f16 { - fn min_value() -> Self { half::f16::MIN } - fn max_value() -> Self { half::f16::MAX } - fn abs(self) -> Self { half::f16::from_f32(self.to_f32().abs()) } + fn min_value() -> Self { + half::f16::MIN + } + fn max_value() -> Self { + half::f16::MAX + } + fn abs(self) -> Self { + half::f16::from_f32(self.to_f32().abs()) + } } impl RealScalar for half::bf16 { - fn min_value() -> Self { half::bf16::MIN } - fn max_value() -> Self { half::bf16::MAX } - fn abs(self) -> Self { half::bf16::from_f32(self.to_f32().abs()) } + fn min_value() -> Self { + half::bf16::MIN + } + fn max_value() -> Self { + half::bf16::MAX + } + fn abs(self) -> Self { + half::bf16::from_f32(self.to_f32().abs()) + } } // ─── FloatScalar for f32 / f64 ─────────────────────────────────────────────── @@ -509,36 +688,109 @@ impl RealScalar for half::bf16 { macro_rules! impl_float_scalar_native { ($ty:ty) => { impl FloatScalar for $ty { - fn nan() -> $ty { <$ty>::NAN } - fn infinity() -> $ty { <$ty>::INFINITY } - fn neg_infinity() -> $ty { <$ty>::NEG_INFINITY } - - #[inline] fn is_nan(self) -> bool { <$ty>::is_nan(self) } - #[inline] fn is_infinite(self) -> bool { <$ty>::is_infinite(self) } - #[inline] fn is_finite(self) -> bool { <$ty>::is_finite(self) } - #[inline] fn is_sign_positive(self) -> bool { <$ty>::is_sign_positive(self) } - #[inline] fn is_sign_negative(self) -> bool { <$ty>::is_sign_negative(self) } - - #[inline] fn sqrt(self) -> $ty { <$ty>::sqrt(self) } - #[inline] fn ln(self) -> $ty { <$ty>::ln(self) } - #[inline] fn log2(self) -> $ty { <$ty>::log2(self) } - #[inline] fn log10(self) -> $ty { <$ty>::log10(self) } - #[inline] fn exp(self) -> $ty { <$ty>::exp(self) } - #[inline] fn exp2(self) -> $ty { <$ty>::exp2(self) } - #[inline] fn powi(self, n: i32) -> $ty { <$ty>::powi(self, n) } - #[inline] fn powf(self, n: $ty) -> $ty { <$ty>::powf(self, n) } - #[inline] fn floor(self) -> $ty { <$ty>::floor(self) } - #[inline] fn ceil(self) -> $ty { <$ty>::ceil(self) } - #[inline] fn round(self) -> $ty { <$ty>::round(self) } - #[inline] fn trunc(self) -> $ty { <$ty>::trunc(self) } - #[inline] fn fract(self) -> $ty { <$ty>::fract(self) } - #[inline] fn mul_add(self, a: $ty, b: $ty) -> $ty { <$ty>::mul_add(self, a, b) } - - fn epsilon() -> $ty { <$ty>::EPSILON } - fn min_positive() -> $ty { <$ty>::MIN_POSITIVE } - - #[inline] fn to_f32(self) -> f32 { self as f32 } - #[inline] fn to_f64(self) -> f64 { self as f64 } + fn nan() -> $ty { + <$ty>::NAN + } + fn infinity() -> $ty { + <$ty>::INFINITY + } + fn neg_infinity() -> $ty { + <$ty>::NEG_INFINITY + } + + #[inline] + fn is_nan(self) -> bool { + <$ty>::is_nan(self) + } + #[inline] + fn is_infinite(self) -> bool { + <$ty>::is_infinite(self) + } + #[inline] + fn is_finite(self) -> bool { + <$ty>::is_finite(self) + } + #[inline] + fn is_sign_positive(self) -> bool { + <$ty>::is_sign_positive(self) + } + #[inline] + fn is_sign_negative(self) -> bool { + <$ty>::is_sign_negative(self) + } + + #[inline] + fn sqrt(self) -> $ty { + <$ty>::sqrt(self) + } + #[inline] + fn ln(self) -> $ty { + <$ty>::ln(self) + } + #[inline] + fn log2(self) -> $ty { + <$ty>::log2(self) + } + #[inline] + fn log10(self) -> $ty { + <$ty>::log10(self) + } + #[inline] + fn exp(self) -> $ty { + <$ty>::exp(self) + } + #[inline] + fn exp2(self) -> $ty { + <$ty>::exp2(self) + } + #[inline] + fn powi(self, n: i32) -> $ty { + <$ty>::powi(self, n) + } + #[inline] + fn powf(self, n: $ty) -> $ty { + <$ty>::powf(self, n) + } + #[inline] + fn floor(self) -> $ty { + <$ty>::floor(self) + } + #[inline] + fn ceil(self) -> $ty { + <$ty>::ceil(self) + } + #[inline] + fn round(self) -> $ty { + <$ty>::round(self) + } + #[inline] + fn trunc(self) -> $ty { + <$ty>::trunc(self) + } + #[inline] + fn fract(self) -> $ty { + <$ty>::fract(self) + } + #[inline] + fn mul_add(self, a: $ty, b: $ty) -> $ty { + <$ty>::mul_add(self, a, b) + } + + fn epsilon() -> $ty { + <$ty>::EPSILON + } + fn min_positive() -> $ty { + <$ty>::MIN_POSITIVE + } + + #[inline] + fn to_f32(self) -> f32 { + self as f32 + } + #[inline] + fn to_f64(self) -> f64 { + self as f64 + } } }; } @@ -551,65 +803,151 @@ impl_float_scalar_native!(f64); macro_rules! impl_float_scalar_half { ($ty:ty, $from_f32:expr, $from_f64:expr) => { impl FloatScalar for $ty { - fn nan() -> $ty { <$ty>::NAN } - fn infinity() -> $ty { <$ty>::INFINITY } - fn neg_infinity() -> $ty { <$ty>::NEG_INFINITY } + fn nan() -> $ty { + <$ty>::NAN + } + fn infinity() -> $ty { + <$ty>::INFINITY + } + fn neg_infinity() -> $ty { + <$ty>::NEG_INFINITY + } - #[inline] fn is_nan(self) -> bool { <$ty>::is_nan(self) } - #[inline] fn is_infinite(self) -> bool { <$ty>::is_infinite(self) } - #[inline] fn is_finite(self) -> bool { <$ty>::is_finite(self) } - #[inline] fn is_sign_positive(self) -> bool { self.to_f32() >= 0.0 } - #[inline] fn is_sign_negative(self) -> bool { self.to_f32() < 0.0 } + #[inline] + fn is_nan(self) -> bool { + <$ty>::is_nan(self) + } + #[inline] + fn is_infinite(self) -> bool { + <$ty>::is_infinite(self) + } + #[inline] + fn is_finite(self) -> bool { + <$ty>::is_finite(self) + } + #[inline] + fn is_sign_positive(self) -> bool { + self.to_f32() >= 0.0 + } + #[inline] + fn is_sign_negative(self) -> bool { + self.to_f32() < 0.0 + } // All ops go via f32 — sufficient for half-precision accuracy. - #[inline] fn sqrt(self) -> $ty { $from_f32(self.to_f32().sqrt()) } - #[inline] fn ln(self) -> $ty { $from_f32(self.to_f32().ln()) } - #[inline] fn log2(self) -> $ty { $from_f32(self.to_f32().log2()) } - #[inline] fn log10(self) -> $ty { $from_f32(self.to_f32().log10()) } - #[inline] fn exp(self) -> $ty { $from_f32(self.to_f32().exp()) } - #[inline] fn exp2(self) -> $ty { $from_f32(self.to_f32().exp2()) } - #[inline] fn powi(self, n: i32) -> $ty { $from_f32(self.to_f32().powi(n)) } - #[inline] fn powf(self, n: $ty) -> $ty { $from_f32(self.to_f32().powf(n.to_f32())) } - #[inline] fn floor(self) -> $ty { $from_f32(self.to_f32().floor()) } - #[inline] fn ceil(self) -> $ty { $from_f32(self.to_f32().ceil()) } - #[inline] fn round(self) -> $ty { $from_f32(self.to_f32().round()) } - #[inline] fn trunc(self) -> $ty { $from_f32(self.to_f32().trunc()) } - #[inline] fn fract(self) -> $ty { $from_f32(self.to_f32().fract()) } - #[inline] fn mul_add(self, a: $ty, b: $ty) -> $ty { + #[inline] + fn sqrt(self) -> $ty { + $from_f32(self.to_f32().sqrt()) + } + #[inline] + fn ln(self) -> $ty { + $from_f32(self.to_f32().ln()) + } + #[inline] + fn log2(self) -> $ty { + $from_f32(self.to_f32().log2()) + } + #[inline] + fn log10(self) -> $ty { + $from_f32(self.to_f32().log10()) + } + #[inline] + fn exp(self) -> $ty { + $from_f32(self.to_f32().exp()) + } + #[inline] + fn exp2(self) -> $ty { + $from_f32(self.to_f32().exp2()) + } + #[inline] + fn powi(self, n: i32) -> $ty { + $from_f32(self.to_f32().powi(n)) + } + #[inline] + fn powf(self, n: $ty) -> $ty { + $from_f32(self.to_f32().powf(n.to_f32())) + } + #[inline] + fn floor(self) -> $ty { + $from_f32(self.to_f32().floor()) + } + #[inline] + fn ceil(self) -> $ty { + $from_f32(self.to_f32().ceil()) + } + #[inline] + fn round(self) -> $ty { + $from_f32(self.to_f32().round()) + } + #[inline] + fn trunc(self) -> $ty { + $from_f32(self.to_f32().trunc()) + } + #[inline] + fn fract(self) -> $ty { + $from_f32(self.to_f32().fract()) + } + #[inline] + fn mul_add(self, a: $ty, b: $ty) -> $ty { $from_f32(self.to_f32().mul_add(a.to_f32(), b.to_f32())) } - fn epsilon() -> $ty { <$ty>::EPSILON } - fn min_positive() -> $ty { <$ty>::MIN_POSITIVE } + fn epsilon() -> $ty { + <$ty>::EPSILON + } + fn min_positive() -> $ty { + <$ty>::MIN_POSITIVE + } - #[inline] fn to_f32(self) -> f32 { <$ty>::to_f32(self) } - #[inline] fn to_f64(self) -> f64 { <$ty>::to_f64(self) } + #[inline] + fn to_f32(self) -> f32 { + <$ty>::to_f32(self) + } + #[inline] + fn to_f64(self) -> f64 { + <$ty>::to_f64(self) + } } }; } -impl_float_scalar_half!(half::f16, half::f16::from_f32, half::f16::from_f64); +impl_float_scalar_half!(half::f16, half::f16::from_f32, half::f16::from_f64); impl_float_scalar_half!(half::bf16, half::bf16::from_f32, half::bf16::from_f64); // ─── Scalar for Complex types ───────────────────────────────────────────────── impl Scalar for Complex { - const DTYPE: DType = DType::C64; - const ZERO: Complex = Complex { re: 0.0, im: 0.0 }; - const ONE: Complex = Complex { re: 1.0, im: 0.0 }; - const ITEMSIZE: usize = 8; + const DTYPE: DType = DType::C64; + const ZERO: Complex = Complex { re: 0.0, im: 0.0 }; + const ONE: Complex = Complex { re: 1.0, im: 0.0 }; + const ITEMSIZE: usize = 8; - #[inline] fn to_f64_lossy(self) -> f64 { self.norm() as f64 } - #[inline] fn from_f64_lossy(v: f64) -> Self { Complex { re: v as f32, im: 0.0 } } + #[inline] + fn to_f64_lossy(self) -> f64 { + self.norm() as f64 + } + #[inline] + fn from_f64_lossy(v: f64) -> Self { + Complex { + re: v as f32, + im: 0.0, + } + } } impl Scalar for Complex { - const DTYPE: DType = DType::C128; - const ZERO: Complex = Complex { re: 0.0, im: 0.0 }; - const ONE: Complex = Complex { re: 1.0, im: 0.0 }; - const ITEMSIZE: usize = 16; - - #[inline] fn to_f64_lossy(self) -> f64 { self.norm() } - #[inline] fn from_f64_lossy(v: f64) -> Self { Complex { re: v, im: 0.0 } } + const DTYPE: DType = DType::C128; + const ZERO: Complex = Complex { re: 0.0, im: 0.0 }; + const ONE: Complex = Complex { re: 1.0, im: 0.0 }; + const ITEMSIZE: usize = 16; + + #[inline] + fn to_f64_lossy(self) -> f64 { + self.norm() + } + #[inline] + fn from_f64_lossy(v: f64) -> Self { + Complex { re: v, im: 0.0 } + } } // ─── ComplexScalar impls ────────────────────────────────────────────────────── @@ -619,16 +957,46 @@ macro_rules! impl_complex_scalar { impl ComplexScalar for Complex<$real> { type Real = $real; - #[inline] fn from_re_im(re: $real, im: $real) -> Self { Complex { re, im } } - #[inline] fn re(self) -> $real { self.re } - #[inline] fn im(self) -> $real { self.im } - #[inline] fn conj(self) -> Self { Complex::conj(&self) } - #[inline] fn norm(self) -> $real { Complex::norm(self) } - #[inline] fn norm_sqr(self) -> $real { Complex::norm_sqr(&self) } - #[inline] fn arg(self) -> $real { Complex::arg(self) } - #[inline] fn is_nan(self) -> bool { self.re.is_nan() || self.im.is_nan() } - #[inline] fn is_infinite(self) -> bool { self.re.is_infinite() || self.im.is_infinite() } - #[inline] fn is_finite(self) -> bool { self.re.is_finite() && self.im.is_finite() } + #[inline] + fn from_re_im(re: $real, im: $real) -> Self { + Complex { re, im } + } + #[inline] + fn re(self) -> $real { + self.re + } + #[inline] + fn im(self) -> $real { + self.im + } + #[inline] + fn conj(self) -> Self { + Complex::conj(&self) + } + #[inline] + fn norm(self) -> $real { + Complex::norm(self) + } + #[inline] + fn norm_sqr(self) -> $real { + Complex::norm_sqr(&self) + } + #[inline] + fn arg(self) -> $real { + Complex::arg(self) + } + #[inline] + fn is_nan(self) -> bool { + self.re.is_nan() || self.im.is_nan() + } + #[inline] + fn is_infinite(self) -> bool { + self.re.is_infinite() || self.im.is_infinite() + } + #[inline] + fn is_finite(self) -> bool { + self.re.is_finite() && self.im.is_finite() + } } }; } @@ -642,9 +1010,18 @@ impl_complex_scalar!(Complex, f64, DType::C128); const _SCALAR_IMPLS: () = { const fn check() {} check::(); - check::(); check::(); check::(); check::(); - check::(); check::(); check::(); check::(); - check::(); check::(); - check::(); check::(); - check::>(); check::>(); + check::(); + check::(); + check::(); + check::(); + check::(); + check::(); + check::(); + check::(); + check::(); + check::(); + check::(); + check::(); + check::>(); + check::>(); }; diff --git a/crates/mohu-error/src/codes.rs b/crates/mohu-error/src/codes.rs index ad33228..e3d4671 100644 --- a/crates/mohu-error/src/codes.rs +++ b/crates/mohu-error/src/codes.rs @@ -36,92 +36,92 @@ #[non_exhaustive] pub enum ErrorCode { // Shape (1xxx) - ShapeMismatch = 1000, - BroadcastError = 1001, - DimensionMismatch = 1002, - AxisOutOfRange = 1003, - ScalarArray = 1004, - ZeroSizedDimension = 1005, - ShapeOverflow = 1006, - ReshapeIncompatible = 1007, - EmptyStackSequence = 1008, - ConcatShapeMismatch = 1009, + ShapeMismatch = 1000, + BroadcastError = 1001, + DimensionMismatch = 1002, + AxisOutOfRange = 1003, + ScalarArray = 1004, + ZeroSizedDimension = 1005, + ShapeOverflow = 1006, + ReshapeIncompatible = 1007, + EmptyStackSequence = 1008, + ConcatShapeMismatch = 1009, // DType (2xxx) - DTypeMismatch = 2000, - InvalidCast = 2001, - Overflow = 2002, - Underflow = 2003, - UnknownDType = 2004, - UnsupportedDType = 2005, - AmbiguousPromotion = 2006, + DTypeMismatch = 2000, + InvalidCast = 2001, + Overflow = 2002, + Underflow = 2003, + UnknownDType = 2004, + UnsupportedDType = 2005, + AmbiguousPromotion = 2006, // Index / slice (3xxx) - IndexOutOfBounds = 3000, - TooManyIndices = 3001, - ZeroSliceStep = 3002, - SliceOutOfBounds = 3003, + IndexOutOfBounds = 3000, + TooManyIndices = 3001, + ZeroSliceStep = 3002, + SliceOutOfBounds = 3003, BoolIndexShapeMismatch = 3004, - FancyIndexOutOfBounds = 3005, + FancyIndexOutOfBounds = 3005, // Buffer / memory (4xxx) - AllocationFailed = 4000, - AlignmentError = 4001, - BufferTooSmall = 4002, - InvalidStride = 4003, - OverlappingStrides = 4004, - NonContiguous = 4005, - ReadOnly = 4006, - CannotResizeShared = 4007, - OffsetOverflow = 4008, + AllocationFailed = 4000, + AlignmentError = 4001, + BufferTooSmall = 4002, + InvalidStride = 4003, + OverlappingStrides = 4004, + NonContiguous = 4005, + ReadOnly = 4006, + CannotResizeShared = 4007, + OffsetOverflow = 4008, // Compute / math (5xxx) - SingularMatrix = 5000, - NonConvergence = 5001, - DomainError = 5002, - DivisionByZero = 5003, - MatrixDimensionMismatch = 5004, - EigenDecompositionFailed = 5005, - NotPositiveDefinite = 5006, - QRRankDeficient = 5007, - SVDNonConvergence = 5008, - UnsupportedNormOrder = 5009, + SingularMatrix = 5000, + NonConvergence = 5001, + DomainError = 5002, + DivisionByZero = 5003, + MatrixDimensionMismatch = 5004, + EigenDecompositionFailed = 5005, + NotPositiveDefinite = 5006, + QRRankDeficient = 5007, + SVDNonConvergence = 5008, + UnsupportedNormOrder = 5009, // I/O (6xxx) - Io = 6000, - InvalidMagic = 6001, - UnsupportedVersion = 6002, - CorruptData = 6003, - UnexpectedEof = 6004, - UnsupportedCodec = 6005, - NpyHeaderError = 6006, - NpzEntryNotFound = 6007, - CsvParseError = 6008, + Io = 6000, + InvalidMagic = 6001, + UnsupportedVersion = 6002, + CorruptData = 6003, + UnexpectedEof = 6004, + UnsupportedCodec = 6005, + NpyHeaderError = 6006, + NpzEntryNotFound = 6007, + CsvParseError = 6008, // DLPack (7xxx) - DLPackUnsupportedDevice = 7000, - DLPackVersionMismatch = 7001, - DLPackNullPointer = 7002, - DLPackUnsupportedDType = 7003, - DLPackInvalid = 7004, + DLPackUnsupportedDevice = 7000, + DLPackVersionMismatch = 7001, + DLPackNullPointer = 7002, + DLPackUnsupportedDType = 7003, + DLPackInvalid = 7004, // Arrow (8xxx) - ArrowSchema = 8000, - ArrowIpc = 8001, - ArrowUnsupportedType = 8002, - ArrowValidityError = 8003, + ArrowSchema = 8000, + ArrowIpc = 8001, + ArrowUnsupportedType = 8002, + ArrowValidityError = 8003, // Python / PyO3 (9xxx) - PythonType = 9000, - PythonValue = 9001, - PythonBuffer = 9002, - PythonNoBuffer = 9003, - PythonUnsupportedBufferFormat = 9004, + PythonType = 9000, + PythonValue = 9001, + PythonBuffer = 9002, + PythonNoBuffer = 9003, + PythonUnsupportedBufferFormat = 9004, // General (10xxx) - Context = 10000, - NotImplemented = 10001, - Internal = 10002, + Context = 10000, + NotImplemented = 10001, + Internal = 10002, } impl ErrorCode { @@ -137,36 +137,54 @@ impl ErrorCode { 7000..=7999 => "dlpack", 8000..=8999 => "arrow", 9000..=9999 => "python", - _ => "general", + _ => "general", } } /// Returns `true` if this code falls in the shape domain (1000–1999). - pub fn is_shape(self) -> bool { matches!(self as u32, 1000..=1999) } + pub fn is_shape(self) -> bool { + matches!(self as u32, 1000..=1999) + } /// Returns `true` if this code falls in the dtype domain (2000–2999). - pub fn is_dtype(self) -> bool { matches!(self as u32, 2000..=2999) } + pub fn is_dtype(self) -> bool { + matches!(self as u32, 2000..=2999) + } /// Returns `true` if this code falls in the index domain (3000–3999). - pub fn is_index(self) -> bool { matches!(self as u32, 3000..=3999) } + pub fn is_index(self) -> bool { + matches!(self as u32, 3000..=3999) + } /// Returns `true` if this code falls in the buffer domain (4000–4999). - pub fn is_buffer(self) -> bool { matches!(self as u32, 4000..=4999) } + pub fn is_buffer(self) -> bool { + matches!(self as u32, 4000..=4999) + } /// Returns `true` if this code falls in the compute domain (5000–5999). - pub fn is_compute(self) -> bool { matches!(self as u32, 5000..=5999) } + pub fn is_compute(self) -> bool { + matches!(self as u32, 5000..=5999) + } /// Returns `true` if this code falls in the I/O domain (6000–6999). - pub fn is_io(self) -> bool { matches!(self as u32, 6000..=6999) } + pub fn is_io(self) -> bool { + matches!(self as u32, 6000..=6999) + } /// Returns `true` if this code falls in the DLPack domain (7000–7999). - pub fn is_dlpack(self) -> bool { matches!(self as u32, 7000..=7999) } + pub fn is_dlpack(self) -> bool { + matches!(self as u32, 7000..=7999) + } /// Returns `true` if this code falls in the Arrow domain (8000–8999). - pub fn is_arrow(self) -> bool { matches!(self as u32, 8000..=8999) } + pub fn is_arrow(self) -> bool { + matches!(self as u32, 8000..=8999) + } /// Returns `true` if this code falls in the Python domain (9000–9999). - pub fn is_python(self) -> bool { matches!(self as u32, 9000..=9999) } + pub fn is_python(self) -> bool { + matches!(self as u32, 9000..=9999) + } } impl std::fmt::Display for ErrorCode { diff --git a/crates/mohu-error/src/context.rs b/crates/mohu-error/src/context.rs index f0a7e63..0d457f7 100644 --- a/crates/mohu-error/src/context.rs +++ b/crates/mohu-error/src/context.rs @@ -142,4 +142,3 @@ impl ResultExt for Option { self.ok_or(err) } } - diff --git a/crates/mohu-error/src/error.rs b/crates/mohu-error/src/error.rs index 347a409..eb5666b 100644 --- a/crates/mohu-error/src/error.rs +++ b/crates/mohu-error/src/error.rs @@ -27,13 +27,15 @@ pub enum MohuError { // ------------------------------------------------------------------------- // Shape & dimension errors (1xxx) // ------------------------------------------------------------------------- - /// Two arrays have incompatible shapes for an element-wise operation. #[error( "shape mismatch: expected {expected:?}, got {got:?}\n\ hint: shapes must be identical for this operation, or broadcastable" )] - ShapeMismatch { expected: Vec, got: Vec }, + ShapeMismatch { + expected: Vec, + got: Vec, + }, /// Two shapes cannot be broadcast together under NumPy-style rules. #[error( @@ -54,7 +56,11 @@ pub enum MohuError { "axis {axis} is out of range for a {ndim}D array\n\ hint: valid axes are {valid}" )] - AxisOutOfRange { axis: i64, ndim: usize, valid: String }, + AxisOutOfRange { + axis: i64, + ndim: usize, + valid: String, + }, /// An operation that requires at least one dimension was called on a scalar. #[error( @@ -104,7 +110,6 @@ pub enum MohuError { // ------------------------------------------------------------------------- // DType errors (2xxx) // ------------------------------------------------------------------------- - /// An operation received arrays with incompatible data types. #[error( "dtype mismatch: expected {expected}, got {got}\n\ @@ -114,7 +119,11 @@ pub enum MohuError { /// A type cast between two dtypes is not valid or would lose data. #[error("cannot cast {from} to {to}: {reason}")] - InvalidCast { from: String, to: String, reason: String }, + InvalidCast { + from: String, + to: String, + reason: String, + }, /// A value exceeds the representable range of the target dtype. #[error( @@ -151,13 +160,16 @@ pub enum MohuError { // ------------------------------------------------------------------------- // Index & slice errors (3xxx) // ------------------------------------------------------------------------- - /// An integer index is outside the valid range for its axis. #[error( "index {index} is out of bounds for axis {axis} with size {size}\n\ hint: valid indices are -{size}..{size}" )] - IndexOutOfBounds { index: i64, axis: usize, size: usize }, + IndexOutOfBounds { + index: i64, + axis: usize, + size: usize, + }, /// More indices were provided than the array has dimensions. #[error( @@ -172,7 +184,12 @@ pub enum MohuError { /// A slice range is out of bounds for the axis it indexes. #[error("slice [{start}:{stop}:{step}] is invalid for axis with size {size}")] - SliceOutOfBounds { start: i64, stop: i64, step: i64, size: usize }, + SliceOutOfBounds { + start: i64, + stop: i64, + step: i64, + size: usize, + }, /// A boolean mask has a different shape than the array it indexes. #[error( @@ -189,12 +206,15 @@ pub enum MohuError { "fancy index on axis {axis} is out of bounds: \ index value {index} exceeds axis size {size}" )] - FancyIndexOutOfBounds { index: i64, axis: usize, size: usize }, + FancyIndexOutOfBounds { + index: i64, + axis: usize, + size: usize, + }, // ------------------------------------------------------------------------- // Buffer & memory errors (4xxx) // ------------------------------------------------------------------------- - /// A memory allocation request failed (likely OOM). #[error( "memory allocation failed: requested {bytes} bytes ({human})\n\ @@ -221,7 +241,11 @@ pub enum MohuError { "invalid stride on axis {axis}: stride {stride} is not a multiple \ of the element size {element_size}" )] - InvalidStride { axis: usize, stride: isize, element_size: usize }, + InvalidStride { + axis: usize, + stride: isize, + element_size: usize, + }, /// The combination of shape and strides would produce overlapping elements. #[error( @@ -269,7 +293,6 @@ pub enum MohuError { // ------------------------------------------------------------------------- // Compute / math errors (5xxx) // ------------------------------------------------------------------------- - /// A matrix is singular (rank-deficient) and cannot be inverted or solved. #[error( "singular matrix: rank-deficient and cannot be inverted or solved\n\ @@ -328,9 +351,7 @@ pub enum MohuError { NotPositiveDefinite, /// QR decomposition found a rank-deficient matrix when full rank was expected. - #[error( - "QR decomposition failed: matrix has rank {actual}, expected full rank {expected}" - )] + #[error("QR decomposition failed: matrix has rank {actual}, expected full rank {expected}")] QRRankDeficient { expected: usize, actual: usize }, /// SVD iteration did not converge. @@ -347,7 +368,6 @@ pub enum MohuError { // ------------------------------------------------------------------------- // I/O errors (6xxx) // ------------------------------------------------------------------------- - #[error("I/O error: {0}")] Io(#[from] std::io::Error), @@ -374,7 +394,10 @@ pub enum MohuError { }, #[error("corrupt or malformed {format} data: {detail}")] - CorruptData { format: &'static str, detail: String }, + CorruptData { + format: &'static str, + detail: String, + }, #[error("unexpected end of file while reading {format} at byte offset {offset}")] UnexpectedEof { format: &'static str, offset: u64 }, @@ -396,12 +419,15 @@ pub enum MohuError { NpzEntryNotFound { name: String }, #[error("CSV parse error at row {row}, column {col}: {detail}")] - CsvParseError { row: usize, col: usize, detail: String }, + CsvParseError { + row: usize, + col: usize, + detail: String, + }, // ------------------------------------------------------------------------- // DLPack errors (7xxx) // ------------------------------------------------------------------------- - #[error( "DLPack: unsupported device type {device_type}\n\ hint: mohu supports CPU (device_type=1) only; \ @@ -434,7 +460,6 @@ pub enum MohuError { // ------------------------------------------------------------------------- // Arrow errors (8xxx) // ------------------------------------------------------------------------- - #[error("Arrow schema mismatch: {0}")] ArrowSchema(String), @@ -447,15 +472,12 @@ pub enum MohuError { )] ArrowUnsupportedType { arrow_type: String }, - #[error( - "Arrow validity bitmap is inconsistent with array length {length}: {detail}" - )] + #[error("Arrow validity bitmap is inconsistent with array length {length}: {detail}")] ArrowValidityError { length: usize, detail: String }, // ------------------------------------------------------------------------- // Python / PyO3 errors (9xxx) // ------------------------------------------------------------------------- - #[error("Python type error: expected {expected}, got {got}")] PythonType { expected: &'static str, got: String }, @@ -480,7 +502,6 @@ pub enum MohuError { // ------------------------------------------------------------------------- // Contextual / structural errors (10xxx) // ------------------------------------------------------------------------- - /// Wraps a lower-level `MohuError` with a human-readable context string. /// Use the [`ResultExt`](crate::context::ResultExt) trait instead of /// constructing this variant directly. @@ -520,84 +541,84 @@ impl MohuError { pub fn code(&self) -> crate::codes::ErrorCode { use crate::codes::ErrorCode; match self { - Self::ShapeMismatch { .. } => ErrorCode::ShapeMismatch, - Self::BroadcastError { .. } => ErrorCode::BroadcastError, - Self::DimensionMismatch { .. } => ErrorCode::DimensionMismatch, - Self::AxisOutOfRange { .. } => ErrorCode::AxisOutOfRange, - Self::ScalarArray => ErrorCode::ScalarArray, - Self::ZeroSizedDimension { .. } => ErrorCode::ZeroSizedDimension, - Self::ShapeOverflow { .. } => ErrorCode::ShapeOverflow, - Self::ReshapeIncompatible { .. } => ErrorCode::ReshapeIncompatible, - Self::EmptyStackSequence => ErrorCode::EmptyStackSequence, - Self::ConcatShapeMismatch { .. } => ErrorCode::ConcatShapeMismatch, - - Self::DTypeMismatch { .. } => ErrorCode::DTypeMismatch, - Self::InvalidCast { .. } => ErrorCode::InvalidCast, - Self::Overflow { .. } => ErrorCode::Overflow, - Self::Underflow { .. } => ErrorCode::Underflow, - Self::UnknownDType(_) => ErrorCode::UnknownDType, - Self::UnsupportedDType { .. } => ErrorCode::UnsupportedDType, - Self::AmbiguousPromotion { .. } => ErrorCode::AmbiguousPromotion, - - Self::IndexOutOfBounds { .. } => ErrorCode::IndexOutOfBounds, - Self::TooManyIndices { .. } => ErrorCode::TooManyIndices, - Self::ZeroSliceStep => ErrorCode::ZeroSliceStep, - Self::SliceOutOfBounds { .. } => ErrorCode::SliceOutOfBounds, - Self::BoolIndexShapeMismatch { .. } => ErrorCode::BoolIndexShapeMismatch, - Self::FancyIndexOutOfBounds { .. } => ErrorCode::FancyIndexOutOfBounds, - - Self::AllocationFailed { .. } => ErrorCode::AllocationFailed, - Self::AlignmentError { .. } => ErrorCode::AlignmentError, - Self::BufferTooSmall { .. } => ErrorCode::BufferTooSmall, - Self::InvalidStride { .. } => ErrorCode::InvalidStride, - Self::OverlappingStrides { .. } => ErrorCode::OverlappingStrides, - Self::NonContiguous => ErrorCode::NonContiguous, - Self::ReadOnly => ErrorCode::ReadOnly, - Self::CannotResizeShared => ErrorCode::CannotResizeShared, - Self::OffsetOverflow { .. } => ErrorCode::OffsetOverflow, - - Self::SingularMatrix => ErrorCode::SingularMatrix, - Self::NonConvergence { .. } => ErrorCode::NonConvergence, - Self::DomainError { .. } => ErrorCode::DomainError, - Self::DivisionByZero => ErrorCode::DivisionByZero, - Self::MatrixDimensionMismatch { .. } => ErrorCode::MatrixDimensionMismatch, - Self::EigenDecompositionFailed { .. } => ErrorCode::EigenDecompositionFailed, - Self::NotPositiveDefinite => ErrorCode::NotPositiveDefinite, - Self::QRRankDeficient { .. } => ErrorCode::QRRankDeficient, - Self::SVDNonConvergence { .. } => ErrorCode::SVDNonConvergence, - Self::UnsupportedNormOrder { .. } => ErrorCode::UnsupportedNormOrder, - - Self::Io(_) => ErrorCode::Io, - Self::InvalidMagic { .. } => ErrorCode::InvalidMagic, - Self::UnsupportedVersion { .. } => ErrorCode::UnsupportedVersion, - Self::CorruptData { .. } => ErrorCode::CorruptData, - Self::UnexpectedEof { .. } => ErrorCode::UnexpectedEof, - Self::UnsupportedCodec { .. } => ErrorCode::UnsupportedCodec, - Self::NpyHeaderError { .. } => ErrorCode::NpyHeaderError, - Self::NpzEntryNotFound { .. } => ErrorCode::NpzEntryNotFound, - Self::CsvParseError { .. } => ErrorCode::CsvParseError, - - Self::DLPackUnsupportedDevice { .. } => ErrorCode::DLPackUnsupportedDevice, - Self::DLPackVersionMismatch { .. } => ErrorCode::DLPackVersionMismatch, - Self::DLPackNullPointer => ErrorCode::DLPackNullPointer, - Self::DLPackUnsupportedDType { .. } => ErrorCode::DLPackUnsupportedDType, - Self::DLPackInvalid(_) => ErrorCode::DLPackInvalid, - - Self::ArrowSchema(_) => ErrorCode::ArrowSchema, - Self::ArrowIpc(_) => ErrorCode::ArrowIpc, - Self::ArrowUnsupportedType { .. } => ErrorCode::ArrowUnsupportedType, - Self::ArrowValidityError { .. } => ErrorCode::ArrowValidityError, - - Self::PythonType { .. } => ErrorCode::PythonType, - Self::PythonValue(_) => ErrorCode::PythonValue, - Self::PythonBuffer(_) => ErrorCode::PythonBuffer, - Self::PythonNoBuffer => ErrorCode::PythonNoBuffer, - Self::PythonUnsupportedBufferFormat{..}=> ErrorCode::PythonUnsupportedBufferFormat, - - Self::Context { .. } => ErrorCode::Context, - Self::NotImplemented(_) => ErrorCode::NotImplemented, - Self::Internal(_) => ErrorCode::Internal, - Self::Multiple(_) => ErrorCode::Internal, + Self::ShapeMismatch { .. } => ErrorCode::ShapeMismatch, + Self::BroadcastError { .. } => ErrorCode::BroadcastError, + Self::DimensionMismatch { .. } => ErrorCode::DimensionMismatch, + Self::AxisOutOfRange { .. } => ErrorCode::AxisOutOfRange, + Self::ScalarArray => ErrorCode::ScalarArray, + Self::ZeroSizedDimension { .. } => ErrorCode::ZeroSizedDimension, + Self::ShapeOverflow { .. } => ErrorCode::ShapeOverflow, + Self::ReshapeIncompatible { .. } => ErrorCode::ReshapeIncompatible, + Self::EmptyStackSequence => ErrorCode::EmptyStackSequence, + Self::ConcatShapeMismatch { .. } => ErrorCode::ConcatShapeMismatch, + + Self::DTypeMismatch { .. } => ErrorCode::DTypeMismatch, + Self::InvalidCast { .. } => ErrorCode::InvalidCast, + Self::Overflow { .. } => ErrorCode::Overflow, + Self::Underflow { .. } => ErrorCode::Underflow, + Self::UnknownDType(_) => ErrorCode::UnknownDType, + Self::UnsupportedDType { .. } => ErrorCode::UnsupportedDType, + Self::AmbiguousPromotion { .. } => ErrorCode::AmbiguousPromotion, + + Self::IndexOutOfBounds { .. } => ErrorCode::IndexOutOfBounds, + Self::TooManyIndices { .. } => ErrorCode::TooManyIndices, + Self::ZeroSliceStep => ErrorCode::ZeroSliceStep, + Self::SliceOutOfBounds { .. } => ErrorCode::SliceOutOfBounds, + Self::BoolIndexShapeMismatch { .. } => ErrorCode::BoolIndexShapeMismatch, + Self::FancyIndexOutOfBounds { .. } => ErrorCode::FancyIndexOutOfBounds, + + Self::AllocationFailed { .. } => ErrorCode::AllocationFailed, + Self::AlignmentError { .. } => ErrorCode::AlignmentError, + Self::BufferTooSmall { .. } => ErrorCode::BufferTooSmall, + Self::InvalidStride { .. } => ErrorCode::InvalidStride, + Self::OverlappingStrides { .. } => ErrorCode::OverlappingStrides, + Self::NonContiguous => ErrorCode::NonContiguous, + Self::ReadOnly => ErrorCode::ReadOnly, + Self::CannotResizeShared => ErrorCode::CannotResizeShared, + Self::OffsetOverflow { .. } => ErrorCode::OffsetOverflow, + + Self::SingularMatrix => ErrorCode::SingularMatrix, + Self::NonConvergence { .. } => ErrorCode::NonConvergence, + Self::DomainError { .. } => ErrorCode::DomainError, + Self::DivisionByZero => ErrorCode::DivisionByZero, + Self::MatrixDimensionMismatch { .. } => ErrorCode::MatrixDimensionMismatch, + Self::EigenDecompositionFailed { .. } => ErrorCode::EigenDecompositionFailed, + Self::NotPositiveDefinite => ErrorCode::NotPositiveDefinite, + Self::QRRankDeficient { .. } => ErrorCode::QRRankDeficient, + Self::SVDNonConvergence { .. } => ErrorCode::SVDNonConvergence, + Self::UnsupportedNormOrder { .. } => ErrorCode::UnsupportedNormOrder, + + Self::Io(_) => ErrorCode::Io, + Self::InvalidMagic { .. } => ErrorCode::InvalidMagic, + Self::UnsupportedVersion { .. } => ErrorCode::UnsupportedVersion, + Self::CorruptData { .. } => ErrorCode::CorruptData, + Self::UnexpectedEof { .. } => ErrorCode::UnexpectedEof, + Self::UnsupportedCodec { .. } => ErrorCode::UnsupportedCodec, + Self::NpyHeaderError { .. } => ErrorCode::NpyHeaderError, + Self::NpzEntryNotFound { .. } => ErrorCode::NpzEntryNotFound, + Self::CsvParseError { .. } => ErrorCode::CsvParseError, + + Self::DLPackUnsupportedDevice { .. } => ErrorCode::DLPackUnsupportedDevice, + Self::DLPackVersionMismatch { .. } => ErrorCode::DLPackVersionMismatch, + Self::DLPackNullPointer => ErrorCode::DLPackNullPointer, + Self::DLPackUnsupportedDType { .. } => ErrorCode::DLPackUnsupportedDType, + Self::DLPackInvalid(_) => ErrorCode::DLPackInvalid, + + Self::ArrowSchema(_) => ErrorCode::ArrowSchema, + Self::ArrowIpc(_) => ErrorCode::ArrowIpc, + Self::ArrowUnsupportedType { .. } => ErrorCode::ArrowUnsupportedType, + Self::ArrowValidityError { .. } => ErrorCode::ArrowValidityError, + + Self::PythonType { .. } => ErrorCode::PythonType, + Self::PythonValue(_) => ErrorCode::PythonValue, + Self::PythonBuffer(_) => ErrorCode::PythonBuffer, + Self::PythonNoBuffer => ErrorCode::PythonNoBuffer, + Self::PythonUnsupportedBufferFormat { .. } => ErrorCode::PythonUnsupportedBufferFormat, + + Self::Context { .. } => ErrorCode::Context, + Self::NotImplemented(_) => ErrorCode::NotImplemented, + Self::Internal(_) => ErrorCode::Internal, + Self::Multiple(_) => ErrorCode::Internal, } } @@ -610,7 +631,8 @@ impl MohuError { // Internal errors are always Fatal-kind. Self::Internal(_) => ErrorKind::Internal, // Multiple: take the worst kind across all inner errors. - Self::Multiple(m) => m.iter() + Self::Multiple(m) => m + .iter() .map(|e| e.kind()) .max_by_key(|k| *k as u8) .unwrap_or(ErrorKind::Internal), @@ -654,7 +676,10 @@ impl MohuError { /// Builds an [`AllocationFailed`](Self::AllocationFailed) error with a /// human-readable size string computed automatically. pub fn alloc(bytes: usize) -> Self { - Self::AllocationFailed { bytes, human: fmt_bytes(bytes) } + Self::AllocationFailed { + bytes, + human: fmt_bytes(bytes), + } } /// Builds an [`Internal`](Self::Internal) error. Use for assertion-style @@ -665,7 +690,10 @@ impl MohuError { /// Builds a [`DomainError`](Self::DomainError). pub fn domain(op: &'static str, reason: impl Into) -> Self { - Self::DomainError { op, reason: reason.into() } + Self::DomainError { + op, + reason: reason.into(), + } } /// Builds a [`MatrixDimensionMismatch`](Self::MatrixDimensionMismatch). diff --git a/crates/mohu-error/src/kind.rs b/crates/mohu-error/src/kind.rs index 5153946..7f5e0a6 100644 --- a/crates/mohu-error/src/kind.rs +++ b/crates/mohu-error/src/kind.rs @@ -26,16 +26,16 @@ pub enum ErrorKind { /// The caller passed invalid arguments — wrong shapes, out-of-bounds /// indices, incompatible dtypes, etc. These errors indicate a /// programming mistake and should never be retried as-is. - Usage = 0, + Usage = 0, /// A well-formed operation failed at runtime due to the mathematical /// properties of the data — singular matrix, non-convergence, /// domain error, etc. - Runtime = 1, + Runtime = 1, /// A system-level failure outside mohu's control — I/O, memory /// allocation, DLPack version mismatch, Arrow IPC failure. - System = 2, + System = 2, /// An invariant inside mohu was violated. These should never appear /// in production and always indicate a bug in mohu itself. @@ -55,9 +55,9 @@ impl ErrorKind { /// Human-readable one-word label for this kind. pub fn label(self) -> &'static str { match self { - Self::Usage => "usage", - Self::Runtime => "runtime", - Self::System => "system", + Self::Usage => "usage", + Self::Runtime => "runtime", + Self::System => "system", Self::Internal => "internal", } } @@ -78,8 +78,8 @@ impl From for ErrorKind { // Buffer — mostly caller mistakes (bad strides, read-only), // but allocation failure is a system error. - 4000..=4002 => ErrorKind::System, // Alloc, Align, BufSmall - 4003..=4999 => ErrorKind::Usage, + 4000..=4002 => ErrorKind::System, // Alloc, Align, BufSmall + 4003..=4999 => ErrorKind::Usage, // Compute — runtime mathematical failures 5000..=5999 => ErrorKind::Runtime, @@ -89,7 +89,7 @@ impl From for ErrorKind { // DLPack — mostly usage (wrong device, bad version) // except null pointer which is an internal invariant violation - 7002 => ErrorKind::Internal, // DLPackNullPointer + 7002 => ErrorKind::Internal, // DLPackNullPointer 7000..=7999 => ErrorKind::Usage, // Arrow — system/IPC @@ -101,9 +101,9 @@ impl From for ErrorKind { // Context: delegate to inner error — handled in MohuError::kind() // NotImplemented: runtime // Internal: internal - 10000 => ErrorKind::Runtime, // Context (placeholder; overridden) - 10001 => ErrorKind::Runtime, // NotImplemented - 10002 => ErrorKind::Internal, // Internal + 10000 => ErrorKind::Runtime, // Context (placeholder; overridden) + 10001 => ErrorKind::Runtime, // NotImplemented + 10002 => ErrorKind::Internal, // Internal _ => ErrorKind::Internal, } diff --git a/crates/mohu-error/src/macros.rs b/crates/mohu-error/src/macros.rs index 322876d..341fe3a 100644 --- a/crates/mohu-error/src/macros.rs +++ b/crates/mohu-error/src/macros.rs @@ -95,7 +95,11 @@ macro_rules! assert_axis_valid { }, }); } - if ax < 0 { (nd_i + ax) as usize } else { ax as usize } + if ax < 0 { + (nd_i + ax) as usize + } else { + ax as usize + } }}; } @@ -117,6 +121,10 @@ macro_rules! assert_in_bounds { size: sz, }); } - if idx < 0 { (sz_i + idx) as usize } else { idx as usize } + if idx < 0 { + (sz_i + idx) as usize + } else { + idx as usize + } }}; } diff --git a/crates/mohu-error/src/multi.rs b/crates/mohu-error/src/multi.rs index fa67e9c..7796a87 100644 --- a/crates/mohu-error/src/multi.rs +++ b/crates/mohu-error/src/multi.rs @@ -38,7 +38,9 @@ impl MultiError { /// Creates a `MultiError` pre-allocated for `capacity` errors. pub fn with_capacity(capacity: usize) -> Self { - Self { errors: Vec::with_capacity(capacity) } + Self { + errors: Vec::with_capacity(capacity), + } } /// Adds an error to the collection. @@ -145,7 +147,9 @@ impl<'a> IntoIterator for &'a MultiError { impl FromIterator for MultiError { fn from_iter>(iter: I) -> Self { - Self { errors: iter.into_iter().collect() } + Self { + errors: iter.into_iter().collect(), + } } } diff --git a/crates/mohu-error/src/python.rs b/crates/mohu-error/src/python.rs index f31c3c5..0909dcd 100644 --- a/crates/mohu-error/src/python.rs +++ b/crates/mohu-error/src/python.rs @@ -17,8 +17,8 @@ /// | ReadOnly, NonContiguous | `BufferError` | /// | everything else | `RuntimeError` | use pyo3::exceptions::{ - PyBufferError, PyIndexError, PyMemoryError, PyNotImplementedError, PyOSError, - PyRuntimeError, PyTypeError, PyValueError, + PyBufferError, PyIndexError, PyMemoryError, PyNotImplementedError, PyOSError, PyRuntimeError, + PyTypeError, PyValueError, }; use pyo3::prelude::*; @@ -74,8 +74,9 @@ impl From for PyErr { | MohuError::ArrowValidityError { .. } => PyValueError::new_err(msg), // ── TypeError ───────────────────────────────────────────────── - MohuError::PythonType { .. } - | MohuError::PythonUnsupportedBufferFormat { .. } => PyTypeError::new_err(msg), + MohuError::PythonType { .. } | MohuError::PythonUnsupportedBufferFormat { .. } => { + PyTypeError::new_err(msg) + }, // ── MemoryError ─────────────────────────────────────────────── MohuError::AllocationFailed { .. } diff --git a/crates/mohu-error/src/reporter.rs b/crates/mohu-error/src/reporter.rs index 0b04560..f409d3e 100644 --- a/crates/mohu-error/src/reporter.rs +++ b/crates/mohu-error/src/reporter.rs @@ -39,15 +39,15 @@ /// ``` use std::fmt; -use crate::{chain::ErrorChain, MohuError}; +use crate::{MohuError, chain::ErrorChain}; // ─── ANSI escape sequences ──────────────────────────────────────────────────── mod ansi { - pub const RESET: &str = "\x1b[0m"; - pub const BOLD: &str = "\x1b[1m"; - pub const DIM: &str = "\x1b[2m"; - pub const BOLD_RED: &str = "\x1b[1;31m"; + pub const RESET: &str = "\x1b[0m"; + pub const BOLD: &str = "\x1b[1m"; + pub const DIM: &str = "\x1b[2m"; + pub const BOLD_RED: &str = "\x1b[1;31m"; pub const BOLD_CYAN: &str = "\x1b[1;36m"; } @@ -56,7 +56,7 @@ mod ansi { fn colour_enabled() -> bool { match std::env::var("MOHU_COLOR").as_deref() { Ok("always") => true, - Ok("never") => false, + Ok("never") => false, // "auto" or unset: check for NO_COLOR and TERM _ => { if std::env::var("NO_COLOR").is_ok() { @@ -68,7 +68,7 @@ fn colour_enabled() -> bool { Ok("dumb") | Err(_) => false, Ok(_) => true, } - } + }, } } @@ -93,14 +93,18 @@ pub enum ReportMode { /// and similar macros. pub struct ErrorReporter<'a> { error: &'a MohuError, - mode: ReportMode, + mode: ReportMode, color: bool, } impl<'a> ErrorReporter<'a> { /// Creates a new reporter with the given mode. pub fn new(error: &'a MohuError, mode: ReportMode) -> Self { - Self { error, mode, color: colour_enabled() } + Self { + error, + mode, + color: colour_enabled(), + } } /// Compact single-line reporter. @@ -115,7 +119,11 @@ impl<'a> ErrorReporter<'a> { /// JSON reporter (machine-readable, no ANSI). pub fn json(error: &'a MohuError) -> Self { - Self { error, mode: ReportMode::Json, color: false } + Self { + error, + mode: ReportMode::Json, + color: false, + } } /// Forces colour on or off regardless of environment detection. @@ -141,16 +149,16 @@ impl<'a> ErrorReporter<'a> { f, "{bold_red}error[{code}]{reset} {bold}{msg}{reset}", bold_red = self.c(ansi::BOLD_RED), - code = code, - reset = self.reset(), - bold = self.c(ansi::BOLD), - msg = self.error, + code = code, + reset = self.reset(), + bold = self.c(ansi::BOLD), + msg = self.error, ) } fn fmt_full(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let root = ErrorChain::root(self.error); - let code = root.code(); + let root = ErrorChain::root(self.error); + let code = root.code(); let depth = ErrorChain::depth(self.error); // ── header line ─────────────────────────────────────────────────── @@ -158,10 +166,10 @@ impl<'a> ErrorReporter<'a> { f, "{bold_red}error[{code}]{reset}{bold}: {msg}{reset}", bold_red = self.c(ansi::BOLD_RED), - code = code, - reset = self.reset(), - bold = self.c(ansi::BOLD), - msg = root, + code = code, + reset = self.reset(), + bold = self.c(ansi::BOLD), + msg = root, )?; // ── context chain (outermost → innermost) ───────────────────────── @@ -170,17 +178,17 @@ impl<'a> ErrorReporter<'a> { writeln!( f, "{dim} context chain:{reset}", - dim = self.c(ansi::DIM), + dim = self.c(ansi::DIM), reset = self.reset(), )?; for (i, ctx) in ctxs.iter().enumerate().rev() { writeln!( f, " {dim}{arrow}{reset} {ctx}", - dim = self.c(ansi::DIM), + dim = self.c(ansi::DIM), arrow = if i == 0 { "└─" } else { "├─" }, reset = self.reset(), - ctx = ctx, + ctx = ctx, )?; } } @@ -192,9 +200,9 @@ impl<'a> ErrorReporter<'a> { writeln!( f, " {cyan}hint{reset}: {hint}", - cyan = self.c(ansi::BOLD_CYAN), + cyan = self.c(ansi::BOLD_CYAN), reset = self.reset(), - hint = hint, + hint = hint, )?; } } @@ -203,24 +211,25 @@ impl<'a> ErrorReporter<'a> { writeln!( f, " {dim}[{code}] {domain} error{reset}", - dim = self.c(ansi::DIM), - code = code, + dim = self.c(ansi::DIM), + code = code, domain = code.domain(), - reset = self.reset(), + reset = self.reset(), )?; Ok(()) } fn fmt_json(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let root = ErrorChain::root(self.error); - let code = root.code() as u32; - let kind = crate::kind::ErrorKind::from(root.code()); + let root = ErrorChain::root(self.error); + let code = root.code() as u32; + let kind = crate::kind::ErrorKind::from(root.code()); let depth = ErrorChain::depth(self.error); // Build context array - let ctxs = ErrorChain::context_messages(self.error); - let ctx_json: Vec = ctxs.iter() + let ctxs = ErrorChain::context_messages(self.error); + let ctx_json: Vec = ctxs + .iter() .map(|s| format!("\"{}\"", s.replace('"', "\\\""))) .collect(); @@ -237,12 +246,12 @@ impl<'a> ErrorReporter<'a> { write!( f, r#"{{"code":{code},"kind":"{kind}","message":"{primary}","context":[{ctx}],"hints":[{hints}],"chain_depth":{depth}}}"#, - code = code, - kind = kind, + code = code, + kind = kind, primary = primary, - ctx = ctx_json.join(","), - hints = hints.join(","), - depth = depth, + ctx = ctx_json.join(","), + hints = hints.join(","), + depth = depth, ) } } @@ -251,8 +260,8 @@ impl fmt::Display for ErrorReporter<'_> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self.mode { ReportMode::Compact => self.fmt_compact(f), - ReportMode::Full => self.fmt_full(f), - ReportMode::Json => self.fmt_json(f), + ReportMode::Full => self.fmt_full(f), + ReportMode::Json => self.fmt_json(f), } } } @@ -295,8 +304,8 @@ impl Severity { pub fn label(self) -> &'static str { match self { Self::Warning => "warning", - Self::Error => "error", - Self::Fatal => "fatal", + Self::Error => "error", + Self::Fatal => "fatal", } } } @@ -316,7 +325,7 @@ impl MohuError { pub fn severity(&self) -> Severity { match ErrorChain::root(self) { MohuError::Internal(_) => Severity::Fatal, - _ => Severity::Error, + _ => Severity::Error, } } } diff --git a/crates/mohu-error/src/test_utils.rs b/crates/mohu-error/src/test_utils.rs index 97ebada..3e5f239 100644 --- a/crates/mohu-error/src/test_utils.rs +++ b/crates/mohu-error/src/test_utils.rs @@ -11,7 +11,7 @@ /// - [`assert_err_kind!`] — assert the broad error kind /// - [`assert_shape_err!`] — assert a specific shape mismatch /// - [`assert_err_chain!`] — assert the depth of the context chain -use crate::{codes::ErrorCode, kind::ErrorKind, MohuError, MohuResult}; +use crate::{MohuError, MohuResult, codes::ErrorCode, kind::ErrorKind}; // ─── assertion helpers (non-macro) ─────────────────────────────────────────── @@ -27,15 +27,10 @@ use crate::{codes::ErrorCode, kind::ErrorKind, MohuError, MohuResult}; /// let err = assert_err(r, "should fail on zero divisor"); /// assert!(matches!(err, MohuError::DivisionByZero)); /// ``` -pub fn assert_err( - result: MohuResult, - context: &str, -) -> MohuError { +pub fn assert_err(result: MohuResult, context: &str) -> MohuError { match result { Err(e) => e, - Ok(v) => panic!( - "assert_err failed ({context}): expected Err(_), got Ok({v:?})" - ), + Ok(v) => panic!("assert_err failed ({context}): expected Err(_), got Ok({v:?})"), } } @@ -45,9 +40,7 @@ pub fn assert_err( pub fn assert_ok(result: MohuResult, context: &str) -> T { match result { Ok(v) => v, - Err(e) => panic!( - "assert_ok failed ({context}): expected Ok(_), got Err({e})" - ), + Err(e) => panic!("assert_ok failed ({context}): expected Ok(_), got Err({e})"), } } @@ -108,10 +101,8 @@ pub fn assert_shape_err( got ShapeMismatch {{ expected: {expected:?}, got: {got:?} }}" ); } - } - other => panic!( - "assert_shape_err: expected ShapeMismatch, got {other:?}" - ), + }, + other => panic!("assert_shape_err: expected ShapeMismatch, got {other:?}"), } } diff --git a/crates/mohu-fft/Cargo.toml b/crates/mohu-fft/Cargo.toml index 627b4c1..3a87cfb 100644 --- a/crates/mohu-fft/Cargo.toml +++ b/crates/mohu-fft/Cargo.toml @@ -23,3 +23,9 @@ rayon.workspace = true rustfft.workspace = true num-complex.workspace = true num-traits.workspace = true + +[dev-dependencies] +mohu-testing.workspace = true + +[package.metadata.cargo-machete] +ignored = ["mohu-buffer", "mohu-dtype", "num-traits", "rayon"] diff --git "a/crates/mohu-fft/src/freq\nhelpers\nnd\nnorm\nplan\nreal\ntransform.rs" "b/crates/mohu-fft/src/freq\nhelpers\nnd\nnorm\nplan\nreal\ntransform.rs" deleted file mode 100644 index 2912da6..0000000 --- "a/crates/mohu-fft/src/freq\nhelpers\nnd\nnorm\nplan\nreal\ntransform.rs" +++ /dev/null @@ -1,7 +0,0 @@ -// freq -helpers -nd -norm -plan -real -transform — implementation pending diff --git a/crates/mohu-fft/src/freq.rs b/crates/mohu-fft/src/freq.rs index 08629c5..6d8dcff 100644 --- a/crates/mohu-fft/src/freq.rs +++ b/crates/mohu-fft/src/freq.rs @@ -1 +1,78 @@ -// freq — implementation pending +/// Frequency-axis helpers similar to NumPy's `fftfreq` and `fftshift`. + +/// Return the Discrete Fourier Transform sample frequencies for a window of +/// length `n` and sample spacing `d` (default 1.0). +pub fn fftfreq(n: usize, d: f64) -> Vec { + if n == 0 { + return Vec::new(); + } + let val = 1.0 / (n as f64 * d); + let mut freqs = Vec::with_capacity(n); + // For even `n` the Nyquist frequency (n/2) should be negative (-0.5/d). + // Use `pos_len = (n + 1) / 2` as the number of non-negative frequency bins. + let pos_len = n.div_ceil(2); + for i in 0..pos_len { + freqs.push(i as f64 * val); + } + for i in pos_len..n { + freqs.push(-((n - i) as f64) * val); + } + freqs +} + +/// Return the non-negative frequency bins for real-input transforms. +pub fn rfftfreq(n: usize, d: f64) -> Vec { + if n == 0 { + return Vec::new(); + } + let val = 1.0 / (n as f64 * d); + let count = n / 2 + 1; + (0..count).map(|i| i as f64 * val).collect() +} + +/// Shift the zero-frequency component to the center of the spectrum. +pub fn fftshift(v: &[T]) -> Vec { + let n = v.len(); + if n == 0 { + return Vec::new(); + } + let mid = n / 2; + let mut out = Vec::with_capacity(n); + out.extend_from_slice(&v[mid..]); + out.extend_from_slice(&v[..mid]); + out +} + +/// The inverse of `fftshift`. +pub fn ifftshift(v: &[T]) -> Vec { + let n = v.len(); + if n == 0 { + return Vec::new(); + } + let mid = n.div_ceil(2); + let mut out = Vec::with_capacity(n); + out.extend_from_slice(&v[mid..]); + out.extend_from_slice(&v[..mid]); + out +} + +#[cfg(test)] +mod tests { + use super::*; + use mohu_testing::assert_allclose; + + #[test] + fn test_fftfreq_len() { + let f = fftfreq(4, 1.0); + assert_allclose!(f, vec![0.0, 0.25, -0.5, -0.25], atol = 1e-12); + } + + #[test] + fn test_fftshift() { + let v = vec![0, 1, 2, 3]; + let s = fftshift(&v); + assert_eq!(s, vec![2, 3, 0, 1]); + let r = ifftshift(&s); + assert_eq!(r, v); + } +} diff --git a/crates/mohu-fft/src/lib.rs b/crates/mohu-fft/src/lib.rs index 516688d..d417531 100644 --- a/crates/mohu-fft/src/lib.rs +++ b/crates/mohu-fft/src/lib.rs @@ -17,7 +17,8 @@ /// | `fft2(a, s, axes)` | `np.fft.fft2` | /// | `fftn(a, s, axes)` | `np.fft.fftn` | /// | `fftfreq(n, d)` | `np.fft.fftfreq` | -/// | `fftshift(a, axes)` | `np.fft.fftshift` | +/// | `fftshift(v)` | `np.fft.fftshift` | +/// | `ifftshift(v)` | `np.fft.ifftshift` | /// /// # Normalization modes /// @@ -26,7 +27,6 @@ /// | `Backward` | 1 | 1/n (default) | /// | `Ortho` | 1/sqrt(n) | 1/sqrt(n) | /// | `Forward` | 1/n | 1 | - pub mod freq; pub mod helpers; pub mod nd; diff --git a/crates/mohu-fft/src/transform.rs b/crates/mohu-fft/src/transform.rs index 343a727..8563e7d 100644 --- a/crates/mohu-fft/src/transform.rs +++ b/crates/mohu-fft/src/transform.rs @@ -1 +1,128 @@ -// transform — implementation pending +use num_complex::Complex; +use rustfft::{FftPlanner, num_complex::Complex as RComplex}; + +use crate::Norm; + +/// Compute the 1-D FFT of `input` with optional length `n` and `norm` mode. +/// If `n` is larger than `input.len()` the input is zero-padded; if smaller, +/// it is truncated. +pub fn fft(input: &[Complex], n: Option, norm: Norm) -> Vec> { + let len = n.unwrap_or(input.len()); + if len == 0 { + return Vec::new(); + } + let mut buf: Vec> = vec![RComplex::new(0.0, 0.0); len]; + for (i, v) in input.iter().take(len).enumerate() { + buf[i] = RComplex::new(v.re, v.im); + } + + let mut planner = FftPlanner::new(); + let fft = planner.plan_fft_forward(len); + fft.process(&mut buf); + + // apply forward normalization + let scale = match norm { + Norm::Backward => 1.0, + Norm::Ortho => 1.0 / (len as f64).sqrt(), + Norm::Forward => 1.0 / (len as f64), + }; + + buf.into_iter() + .map(|c| Complex::new(c.re * scale, c.im * scale)) + .collect() +} + +/// Compute the 1-D inverse FFT (IFFT) of `input` with optional length `n` and `norm` mode. +pub fn ifft(input: &[Complex], n: Option, norm: Norm) -> Vec> { + let len = n.unwrap_or(input.len()); + if len == 0 { + return Vec::new(); + } + let mut buf: Vec> = vec![RComplex::new(0.0, 0.0); len]; + for (i, v) in input.iter().take(len).enumerate() { + buf[i] = RComplex::new(v.re, v.im); + } + + let mut planner = FftPlanner::new(); + let ifft = planner.plan_fft_inverse(len); + ifft.process(&mut buf); + + // apply backward normalization + let scale = match norm { + Norm::Backward => 1.0 / (len as f64), + Norm::Ortho => 1.0 / (len as f64).sqrt(), + Norm::Forward => 1.0, + }; + + buf.into_iter() + .map(|c| Complex::new(c.re * scale, c.im * scale)) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use mohu_testing::assert_allclose; + use num_complex::Complex; + + fn assert_complex_close(actual: &[Complex], expected: &[Complex]) { + let actual_re: Vec = actual.iter().map(|value| value.re).collect(); + let actual_im: Vec = actual.iter().map(|value| value.im).collect(); + let expected_re: Vec = expected.iter().map(|value| value.re).collect(); + let expected_im: Vec = expected.iter().map(|value| value.im).collect(); + + assert_allclose!(actual_re, expected_re, atol = 1e-9); + assert_allclose!(actual_im, expected_im, atol = 1e-9); + } + + #[test] + fn roundtrip_fft_ifft_backward() { + let input: Vec> = (0..8).map(|i| Complex::new(i as f64, 0.0)).collect(); + let out = fft(&input, None, Norm::Backward); + let back = ifft(&out, None, Norm::Backward); + assert_complex_close(&back, &input); + } + + #[test] + fn roundtrip_fft_ifft_ortho() { + let input: Vec> = (0..8) + .map(|i| Complex::new((i as f64) * 0.5, 0.0)) + .collect(); + let out = fft(&input, None, Norm::Ortho); + let back = ifft(&out, None, Norm::Ortho); + assert_complex_close(&back, &input); + } + + #[test] + fn roundtrip_fft_ifft_forward() { + let input: Vec> = (0..8) + .map(|i| Complex::new((i as f64) - 3.0, 0.0)) + .collect(); + let out = fft(&input, None, Norm::Forward); + let back = ifft(&out, None, Norm::Forward); + assert_complex_close(&back, &input); + } + + #[test] + fn fft_padding_roundtrip() { + let input = vec![Complex::new(1.0, 0.0), Complex::new(2.0, 0.0)]; + let out = fft(&input, Some(4), Norm::Backward); + let back = ifft(&out, Some(4), Norm::Backward); + let expected = vec![ + Complex::new(1.0, 0.0), + Complex::new(2.0, 0.0), + Complex::new(0.0, 0.0), + Complex::new(0.0, 0.0), + ]; + assert_complex_close(&back, &expected); + } + + #[test] + fn fft_truncation_roundtrip() { + let input: Vec> = (0..4).map(|i| Complex::new(i as f64, 0.0)).collect(); + let out = fft(&input, Some(2), Norm::Backward); + let back = ifft(&out, Some(2), Norm::Backward); + let expected = vec![Complex::new(0.0, 0.0), Complex::new(1.0, 0.0)]; + assert_complex_close(&back, &expected); + } +} diff --git "a/crates/mohu-index/src/boolean\nfancy\ngather\nslice\ntake\nwhere_op.rs" "b/crates/mohu-index/src/boolean\nfancy\ngather\nslice\ntake\nwhere_op.rs" deleted file mode 100644 index 903bd67..0000000 --- "a/crates/mohu-index/src/boolean\nfancy\ngather\nslice\ntake\nwhere_op.rs" +++ /dev/null @@ -1,6 +0,0 @@ -// boolean -fancy -gather -slice -take -where_op — implementation pending diff --git a/crates/mohu-index/src/lib.rs b/crates/mohu-index/src/lib.rs index 278cbaf..abbfbb1 100644 --- a/crates/mohu-index/src/lib.rs +++ b/crates/mohu-index/src/lib.rs @@ -22,7 +22,6 @@ /// /// Boolean indexing also returns a copy because the output length is not /// known until the mask is scanned. - pub mod boolean; pub mod fancy; pub mod gather; diff --git a/crates/mohu-io/Cargo.toml b/crates/mohu-io/Cargo.toml index fda7c36..4fb38be 100644 --- a/crates/mohu-io/Cargo.toml +++ b/crates/mohu-io/Cargo.toml @@ -16,3 +16,4 @@ arrow.workspace = true serde.workspace = true memmap2.workspace = true thiserror.workspace = true +csv.workspace = true diff --git a/crates/mohu-io/src/arrow.rs b/crates/mohu-io/src/arrow.rs index e69de29..8b13789 100644 --- a/crates/mohu-io/src/arrow.rs +++ b/crates/mohu-io/src/arrow.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-io/src/csv.rs b/crates/mohu-io/src/csv.rs index e69de29..199f51b 100644 --- a/crates/mohu-io/src/csv.rs +++ b/crates/mohu-io/src/csv.rs @@ -0,0 +1,416 @@ +//! CSV file I/O for mohu arrays. +//! +//! This module provides a small, typed CSV reader and writer with custom +//! delimiters, header handling, missing-value sentinels, and streaming parse +//! behavior suitable for large inputs. + +use std::{ + fs::File, + io::{BufReader, BufWriter, Read, Write}, + path::Path, +}; + +use thiserror::Error; + +/// Errors produced while reading or writing CSV data. +#[derive(Debug, Error)] +pub enum CsvError { + /// Underlying I/O failure. + #[error("I/O error: {0}")] + Io(#[from] std::io::Error), + + /// Failure reported by the `csv` crate. + #[error("CSV parse error: {0}")] + Parse(#[from] csv::Error), + + /// UTF-8 conversion failed while producing a `String` output. + #[error("UTF-8 conversion error: {0}")] + Utf8(#[from] std::string::FromUtf8Error), + + /// No header or data rows were found. + #[error("empty file: no records found")] + EmptyFile, + + /// A data row had a different number of columns than expected. + #[error("row {row} has {got} columns, expected {expected}")] + ColumnMismatch { + /// 1-based row index within the data section. + row: usize, + /// Expected number of columns. + expected: usize, + /// Actual number of columns. + got: usize, + }, +} + +/// Result type used by CSV reader and writer operations. +pub type CsvResult = Result; + +/// A single typed cell value inferred from a CSV field. +#[derive(Debug, Clone, PartialEq)] +pub enum CsvValue { + /// Signed integer value. + Int(i64), + /// Floating-point value. + Float(f64), + /// Boolean value. + Bool(bool), + /// String value. + Str(String), + /// Missing value placeholder. + Missing, +} + +impl CsvValue { + /// Infer a typed value from a CSV field. + fn infer(raw: &str, missing_values: &[&str]) -> Self { + let trimmed = raw.trim(); + + if trimmed.is_empty() || missing_values.contains(&trimmed) { + return Self::Missing; + } + + if let Ok(value) = trimmed.parse::() { + return Self::Int(value); + } + + if let Ok(value) = trimmed.parse::() { + return Self::Float(value); + } + + match trimmed.to_ascii_lowercase().as_str() { + "true" | "yes" => Self::Bool(true), + "false" | "no" => Self::Bool(false), + _ => Self::Str(trimmed.to_owned()), + } + } + + /// Convert the value into a CSV-ready string. + pub fn to_csv_string(&self) -> String { + match self { + Self::Int(value) => value.to_string(), + Self::Float(value) => value.to_string(), + Self::Bool(value) => value.to_string(), + Self::Str(value) => value.clone(), + Self::Missing => String::new(), + } + } +} + +/// Read configuration for [`CsvReader`]. +#[derive(Debug, Clone)] +pub struct ReadOptions { + /// Field delimiter. + pub delimiter: u8, + /// Whether the first non-comment row should be treated as a header. + pub has_header: bool, + /// Values treated as missing. + pub missing_values: Vec, + /// Optional comment prefix. + pub comment: Option, + /// Maximum number of data rows to read. + pub max_rows: Option, + /// Number of data rows to skip after the header. + pub skip_rows: usize, +} + +impl Default for ReadOptions { + fn default() -> Self { + Self { + delimiter: b',', + has_header: true, + missing_values: vec![ + String::new(), + "NA".to_owned(), + "N/A".to_owned(), + "nan".to_owned(), + "NaN".to_owned(), + "null".to_owned(), + "NULL".to_owned(), + ], + comment: None, + max_rows: None, + skip_rows: 0, + } + } +} + +/// Parsed CSV content stored in row-major form. +#[derive(Debug, Clone)] +pub struct CsvTable { + /// Column names. Empty when the input had no header. + pub headers: Vec, + /// Row-major cells. + pub data: Vec>, + /// Number of columns. + pub ncols: usize, +} + +impl CsvTable { + /// Number of data rows. + pub fn nrows(&self) -> usize { + self.data.len() + } + + /// Return a column by index. + pub fn column(&self, idx: usize) -> Option> { + if idx >= self.ncols { + return None; + } + + Some(self.data.iter().map(|row| &row[idx]).collect()) + } + + /// Return a column by header name. + pub fn column_by_name(&self, name: &str) -> Option> { + let idx = self.headers.iter().position(|header| header == name)?; + self.column(idx) + } +} + +/// Reads a CSV file into a [`CsvTable`]. +/// +/// # Example +/// ```rust,no_run +/// use mohu_io::csv::{CsvReader, ReadOptions}; +/// +/// let table = CsvReader::new(ReadOptions::default()) +/// .read_file("data.csv") +/// .unwrap(); +/// +/// assert!(table.nrows() > 0); +/// ``` +pub struct CsvReader { + opts: ReadOptions, +} + +impl CsvReader { + /// Create a reader with the given options. + pub fn new(opts: ReadOptions) -> Self { + Self { opts } + } + + /// Read CSV content from a file path. + pub fn read_file>(&self, path: P) -> CsvResult { + let file = File::open(path)?; + self.read_impl(BufReader::new(file)) + } + + /// Read CSV content from an in-memory string. + pub fn read_str(&self, src: &str) -> CsvResult { + self.read_impl(src.as_bytes()) + } + + fn read_impl(&self, reader: R) -> CsvResult { + let missing_values: Vec<&str> = self + .opts + .missing_values + .iter() + .map(String::as_str) + .collect(); + + let mut csv_reader = csv::ReaderBuilder::new() + .delimiter(self.opts.delimiter) + .has_headers(false) + .comment(self.opts.comment) + .trim(csv::Trim::All) + .flexible(true) + .from_reader(reader); + + let mut records = csv_reader.records(); + + let headers = if self.opts.has_header { + let header_record = match records.next() { + Some(result) => result?, + None => return Err(CsvError::EmptyFile), + }; + + header_record + .iter() + .map(|value| value.trim().to_owned()) + .collect::>() + } else { + Vec::new() + }; + + let mut data = Vec::new(); + let mut ncols = if headers.is_empty() { + None + } else { + Some(headers.len()) + }; + + for (record_index, result) in records.enumerate() { + if record_index < self.opts.skip_rows { + continue; + } + + if let Some(limit) = self.opts.max_rows { + if data.len() >= limit { + break; + } + } + + let record = result?; + let cells = record + .iter() + .map(|field| CsvValue::infer(field, &missing_values)) + .collect::>(); + + match ncols { + None => ncols = Some(cells.len()), + Some(expected) if cells.len() != expected => { + return Err(CsvError::ColumnMismatch { + row: record_index + 1, + expected, + got: cells.len(), + }); + }, + Some(_) => {}, + } + + data.push(cells); + } + + let ncols = ncols.unwrap_or(0); + + if data.is_empty() { + return Err(CsvError::EmptyFile); + } + + Ok(CsvTable { + headers, + data, + ncols, + }) + } +} + +/// Write configuration for [`CsvWriter`]. +#[derive(Debug, Clone)] +pub struct WriteOptions { + /// Field delimiter. + pub delimiter: u8, + /// Whether to write the header row. + pub write_header: bool, + /// Representation to use for missing values. + pub missing_repr: String, + /// Line terminator written after each row. + pub line_terminator: String, +} + +impl Default for WriteOptions { + fn default() -> Self { + Self { + delimiter: b',', + write_header: true, + missing_repr: String::new(), + line_terminator: "\n".to_owned(), + } + } +} + +/// Writes a [`CsvTable`] to a file or string. +/// +/// # Example +/// ```rust,no_run +/// use mohu_io::csv::{CsvWriter, WriteOptions}; +/// +/// # let table = mohu_io::csv::CsvTable { +/// # headers: vec!["a".to_owned()], +/// # data: vec![vec![mohu_io::csv::CsvValue::Int(1)]], +/// # ncols: 1, +/// # }; +/// let csv = CsvWriter::new(WriteOptions::default()).write_str(&table).unwrap(); +/// assert!(csv.contains('a')); +/// ``` +pub struct CsvWriter { + opts: WriteOptions, +} + +impl CsvWriter { + /// Create a writer with the given options. + pub fn new(opts: WriteOptions) -> Self { + Self { opts } + } + + /// Write a [`CsvTable`] to a file. + pub fn write_file>(&self, table: &CsvTable, path: P) -> CsvResult<()> { + let file = File::create(path)?; + self.write_impl(table, BufWriter::new(file)) + } + + /// Write a [`CsvTable`] to a string. + pub fn write_str(&self, table: &CsvTable) -> CsvResult { + let mut buffer = Vec::new(); + self.write_impl(table, &mut buffer)?; + Ok(String::from_utf8(buffer)?) + } + + fn write_impl(&self, table: &CsvTable, mut writer: W) -> CsvResult<()> { + if self.opts.write_header && !table.headers.is_empty() { + self.write_row(&mut writer, table.headers.iter().map(String::as_str))?; + } + + for row in &table.data { + let values = row + .iter() + .map(|value| { + if matches!(value, CsvValue::Missing) { + self.opts.missing_repr.as_str().to_owned() + } else { + value.to_csv_string() + } + }) + .collect::>(); + + self.write_row(&mut writer, values.iter().map(String::as_str))?; + } + + Ok(()) + } + + fn write_row(&self, writer: &mut W, fields: I) -> CsvResult<()> + where + W: Write, + I: IntoIterator, + S: AsRef<[u8]>, + { + let mut row_bytes = Vec::new(); + { + let mut csv_writer = csv::WriterBuilder::new() + .delimiter(self.opts.delimiter) + .has_headers(false) + .terminator(csv::Terminator::Any(b'\n')) + .from_writer(&mut row_bytes); + + csv_writer.write_record(fields)?; + csv_writer.flush()?; + } + + if self.opts.line_terminator == "\n" { + writer.write_all(&row_bytes)?; + return Ok(()); + } + + if row_bytes.ends_with(b"\r\n") { + row_bytes.truncate(row_bytes.len().saturating_sub(2)); + } else if row_bytes.ends_with(b"\n") { + row_bytes.pop(); + } + + writer.write_all(&row_bytes)?; + writer.write_all(self.opts.line_terminator.as_bytes())?; + Ok(()) + } +} + +/// Read a CSV file using default options. +pub fn read_csv>(path: P) -> CsvResult { + CsvReader::new(ReadOptions::default()).read_file(path) +} + +/// Write a [`CsvTable`] to a file using default options. +pub fn write_csv>(table: &CsvTable, path: P) -> CsvResult<()> { + CsvWriter::new(WriteOptions::default()).write_file(table, path) +} diff --git a/crates/mohu-io/src/lib.rs b/crates/mohu-io/src/lib.rs index b756580..6204111 100644 --- a/crates/mohu-io/src/lib.rs +++ b/crates/mohu-io/src/lib.rs @@ -2,3 +2,8 @@ pub mod arrow; pub mod csv; pub mod mmap; pub mod npy; + +pub use csv::{ + CsvError, CsvReader, CsvResult, CsvTable, CsvValue, CsvWriter, ReadOptions, WriteOptions, + read_csv, write_csv, +}; diff --git a/crates/mohu-io/src/mmap.rs b/crates/mohu-io/src/mmap.rs index e69de29..8b13789 100644 --- a/crates/mohu-io/src/mmap.rs +++ b/crates/mohu-io/src/mmap.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-io/src/npy.rs b/crates/mohu-io/src/npy.rs index e69de29..8b13789 100644 --- a/crates/mohu-io/src/npy.rs +++ b/crates/mohu-io/src/npy.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-io/tests/csv_tests.rs b/crates/mohu-io/tests/csv_tests.rs new file mode 100644 index 0000000..82c0f8d --- /dev/null +++ b/crates/mohu-io/tests/csv_tests.rs @@ -0,0 +1,201 @@ +use mohu_io::csv::{CsvError, CsvReader, CsvValue, CsvWriter, ReadOptions, WriteOptions}; + +const SAMPLE_CSV: &str = "\ +name,age,score,active +Alice,30,9.5,true +Bob,25,8.1,false +Charlie,,7.7,true +"; + +const TAB_CSV: &str = "x\ty\tz\n1\t2\t3\n4\t5\t6\n"; + +#[test] +fn test_read_basic_headers() { + let table = CsvReader::new(ReadOptions::default()) + .read_str(SAMPLE_CSV) + .unwrap(); + + assert_eq!(table.headers, vec!["name", "age", "score", "active"]); + assert_eq!(table.nrows(), 3); + assert_eq!(table.ncols, 4); +} + +#[test] +fn test_type_inference_int() { + let table = CsvReader::new(ReadOptions::default()) + .read_str(SAMPLE_CSV) + .unwrap(); + + assert_eq!(table.data[0][1], CsvValue::Int(30)); + assert_eq!(table.data[1][1], CsvValue::Int(25)); +} + +#[test] +fn test_type_inference_float() { + let table = CsvReader::new(ReadOptions::default()) + .read_str(SAMPLE_CSV) + .unwrap(); + + assert_eq!(table.data[0][2], CsvValue::Float(9.5)); +} + +#[test] +fn test_type_inference_bool() { + let table = CsvReader::new(ReadOptions::default()) + .read_str(SAMPLE_CSV) + .unwrap(); + + assert_eq!(table.data[0][3], CsvValue::Bool(true)); + assert_eq!(table.data[1][3], CsvValue::Bool(false)); +} + +#[test] +fn test_missing_value_detected() { + let table = CsvReader::new(ReadOptions::default()) + .read_str(SAMPLE_CSV) + .unwrap(); + + assert_eq!(table.data[2][1], CsvValue::Missing); +} + +#[test] +fn test_custom_missing_sentinel() { + let csv = "a,b\n1,N/A\n2,3\n"; + let opts = ReadOptions { + missing_values: vec!["N/A".to_owned()], + ..Default::default() + }; + + let table = CsvReader::new(opts).read_str(csv).unwrap(); + assert_eq!(table.data[0][1], CsvValue::Missing); + assert_eq!(table.data[1][1], CsvValue::Int(3)); +} + +#[test] +fn test_tab_delimiter() { + let opts = ReadOptions { + delimiter: b'\t', + ..Default::default() + }; + + let table = CsvReader::new(opts).read_str(TAB_CSV).unwrap(); + assert_eq!(table.headers, vec!["x", "y", "z"]); + assert_eq!(table.data[0][0], CsvValue::Int(1)); +} + +#[test] +fn test_no_header() { + let csv = "1,2,3\n4,5,6\n"; + let opts = ReadOptions { + has_header: false, + ..Default::default() + }; + + let table = CsvReader::new(opts).read_str(csv).unwrap(); + assert!(table.headers.is_empty()); + assert_eq!(table.nrows(), 2); +} + +#[test] +fn test_max_rows() { + let opts = ReadOptions { + max_rows: Some(1), + ..Default::default() + }; + + let table = CsvReader::new(opts).read_str(SAMPLE_CSV).unwrap(); + assert_eq!(table.nrows(), 1); +} + +#[test] +fn test_skip_rows() { + let opts = ReadOptions { + skip_rows: 1, + ..Default::default() + }; + + let table = CsvReader::new(opts).read_str(SAMPLE_CSV).unwrap(); + assert_eq!(table.data[0][0], CsvValue::Str("Bob".to_owned())); +} + +#[test] +fn test_column_by_name() { + let table = CsvReader::new(ReadOptions::default()) + .read_str(SAMPLE_CSV) + .unwrap(); + + let names: Vec<_> = table.column_by_name("name").unwrap(); + assert_eq!(names[0], &CsvValue::Str("Alice".to_owned())); +} + +#[test] +fn test_empty_file_error() { + let result = CsvReader::new(ReadOptions::default()).read_str("name,age\n"); + assert!(matches!(result, Err(CsvError::EmptyFile))); +} + +#[test] +fn test_round_trip() { + let original = CsvReader::new(ReadOptions::default()) + .read_str(SAMPLE_CSV) + .unwrap(); + + let written = CsvWriter::new(WriteOptions::default()) + .write_str(&original) + .unwrap(); + + let recovered = CsvReader::new(ReadOptions::default()) + .read_str(&written) + .unwrap(); + + assert_eq!(original.headers, recovered.headers); + assert_eq!(original.nrows(), recovered.nrows()); + assert_eq!(original.data[0][0], recovered.data[0][0]); +} + +#[test] +fn test_write_tab_delimiter() { + let table = CsvReader::new(ReadOptions::default()) + .read_str(SAMPLE_CSV) + .unwrap(); + + let opts = WriteOptions { + delimiter: b'\t', + ..Default::default() + }; + + let out = CsvWriter::new(opts).write_str(&table).unwrap(); + assert!(out.contains('\t')); + assert!(!out.contains(',')); +} + +#[test] +fn test_write_missing_repr() { + let csv = "a,b\n1,\n2,3\n"; + let table = CsvReader::new(ReadOptions::default()) + .read_str(csv) + .unwrap(); + + let opts = WriteOptions { + missing_repr: "NA".to_owned(), + ..Default::default() + }; + + let out = CsvWriter::new(opts).write_str(&table).unwrap(); + assert!(out.contains("NA")); +} + +#[test] +fn test_write_no_header() { + let table = CsvReader::new(ReadOptions::default()) + .read_str(SAMPLE_CSV) + .unwrap(); + + let opts = WriteOptions { + write_header: false, + ..Default::default() + }; + + let out = CsvWriter::new(opts).write_str(&table).unwrap(); + assert!(!out.starts_with("name")); +} diff --git "a/crates/mohu-masked/src/arith\narray\ncompress\nfill\nio\nmask_ops\nreduce.rs" "b/crates/mohu-masked/src/arith\narray\ncompress\nfill\nio\nmask_ops\nreduce.rs" deleted file mode 100644 index 81a59d0..0000000 --- "a/crates/mohu-masked/src/arith\narray\ncompress\nfill\nio\nmask_ops\nreduce.rs" +++ /dev/null @@ -1,7 +0,0 @@ -// arith -array -compress -fill -io -mask_ops -reduce — implementation pending diff --git a/crates/mohu-masked/src/lib.rs b/crates/mohu-masked/src/lib.rs index 8fc7571..d9c7182 100644 --- a/crates/mohu-masked/src/lib.rs +++ b/crates/mohu-masked/src/lib.rs @@ -23,7 +23,6 @@ /// | [`fill`] | `filled` — replace masked with fill_value | /// | [`mask_ops`] | `masked_where`, `masked_equal`, `getmask`, `getdata` | /// | [`io`] | serialise/deserialise masked arrays (NPY extension) | - pub mod arith; pub mod array; pub mod compress; diff --git a/crates/mohu-ops/src/arith.rs b/crates/mohu-ops/src/arith.rs index e69de29..8b13789 100644 --- a/crates/mohu-ops/src/arith.rs +++ b/crates/mohu-ops/src/arith.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-ops/src/broadcast.rs b/crates/mohu-ops/src/broadcast.rs index e69de29..8b13789 100644 --- a/crates/mohu-ops/src/broadcast.rs +++ b/crates/mohu-ops/src/broadcast.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-ops/src/cmp.rs b/crates/mohu-ops/src/cmp.rs index e69de29..8b13789 100644 --- a/crates/mohu-ops/src/cmp.rs +++ b/crates/mohu-ops/src/cmp.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-ops/src/logical.rs b/crates/mohu-ops/src/logical.rs index e69de29..8b13789 100644 --- a/crates/mohu-ops/src/logical.rs +++ b/crates/mohu-ops/src/logical.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-ops/src/reduce.rs b/crates/mohu-ops/src/reduce.rs index e69de29..8b13789 100644 --- a/crates/mohu-ops/src/reduce.rs +++ b/crates/mohu-ops/src/reduce.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-ops/src/unary.rs b/crates/mohu-ops/src/unary.rs index e69de29..8b13789 100644 --- a/crates/mohu-ops/src/unary.rs +++ b/crates/mohu-ops/src/unary.rs @@ -0,0 +1 @@ + diff --git "a/crates/mohu-random/src/continuous\ndiscrete\nentropy\ngenerator\nmultivariate\npermutation\nseeding.rs" "b/crates/mohu-random/src/continuous\ndiscrete\nentropy\ngenerator\nmultivariate\npermutation\nseeding.rs" deleted file mode 100644 index 391727c..0000000 --- "a/crates/mohu-random/src/continuous\ndiscrete\nentropy\ngenerator\nmultivariate\npermutation\nseeding.rs" +++ /dev/null @@ -1,7 +0,0 @@ -// continuous -discrete -entropy -generator -multivariate -permutation -seeding — implementation pending diff --git a/crates/mohu-random/src/generator.rs b/crates/mohu-random/src/generator.rs index 6f6d205..68d3cd7 100644 --- a/crates/mohu-random/src/generator.rs +++ b/crates/mohu-random/src/generator.rs @@ -3,7 +3,9 @@ /// Trait implemented by all mohu PRNG engines. pub trait Generator: Send + Sync { /// Seed the generator from a u64. - fn seed(seed: u64) -> Self where Self: Sized; + fn seed(seed: u64) -> Self + where + Self: Sized; /// Fill a byte slice with random bytes. fn fill_bytes(&mut self, dest: &mut [u8]); /// Return a random u64. @@ -20,7 +22,8 @@ impl Pcg64 { const MULTIPLIER: u128 = 0x2360_ED05_1FC6_5DA4_4385_DF64_9FCC_F645; fn step(&mut self) { - self.state = self.state + self.state = self + .state .wrapping_mul(Self::MULTIPLIER) .wrapping_add(self.inc); } @@ -82,12 +85,7 @@ impl Philox4x64 { const M1: u64 = 0xCA5A826395121157; let (hi0, lo0) = mul128(M0, ctr[0]); let (hi1, lo1) = mul128(M1, ctr[2]); - [ - hi1 ^ ctr[1] ^ key[0], - lo1, - hi0 ^ ctr[3] ^ key[1], - lo0, - ] + [hi1 ^ ctr[1] ^ key[0], lo1, hi0 ^ ctr[3] ^ key[1], lo0] } fn generate(&mut self) { diff --git a/crates/mohu-random/src/lib.rs b/crates/mohu-random/src/lib.rs index 517b00f..d5f3ecf 100644 --- a/crates/mohu-random/src/lib.rs +++ b/crates/mohu-random/src/lib.rs @@ -31,7 +31,6 @@ /// /// All generators implement `Seed` — the same seed always produces the /// same sequence regardless of CPU count or mohu version within a major. - pub mod continuous; pub mod discrete; pub mod entropy; diff --git "a/crates/mohu-simd/src/arith\nbitwise\ncast\ncmp\ncopy\nfill\nfma\nmath\nreduce\ndetect.rs" "b/crates/mohu-simd/src/arith\nbitwise\ncast\ncmp\ncopy\nfill\nfma\nmath\nreduce\ndetect.rs" deleted file mode 100644 index 1e0359c..0000000 --- "a/crates/mohu-simd/src/arith\nbitwise\ncast\ncmp\ncopy\nfill\nfma\nmath\nreduce\ndetect.rs" +++ /dev/null @@ -1,10 +0,0 @@ -// arith -bitwise -cast -cmp -copy -fill -fma -math -reduce -detect — implementation pending diff --git a/crates/mohu-simd/src/lib.rs b/crates/mohu-simd/src/lib.rs index b6d972e..bae8b04 100644 --- a/crates/mohu-simd/src/lib.rs +++ b/crates/mohu-simd/src/lib.rs @@ -30,7 +30,6 @@ /// | [`math`] | sqrt, rsqrt, exp, log, sin, cos (approx + exact)| /// | [`fma`] | fused multiply-add / multiply-subtract | /// | [`bitwise`] | and, or, xor, not, shl, shr for integer types | - pub mod arith; pub mod bitwise; pub mod cast; diff --git "a/crates/mohu-sparse/src/arith\nbsr\ncoo\ncsc\ncsr\nconvert\ndia\nlinalg\nslice\nspmm\nspmv.rs" "b/crates/mohu-sparse/src/arith\nbsr\ncoo\ncsc\ncsr\nconvert\ndia\nlinalg\nslice\nspmm\nspmv.rs" deleted file mode 100644 index 666cef3..0000000 --- "a/crates/mohu-sparse/src/arith\nbsr\ncoo\ncsc\ncsr\nconvert\ndia\nlinalg\nslice\nspmm\nspmv.rs" +++ /dev/null @@ -1,11 +0,0 @@ -// arith -bsr -coo -csc -csr -convert -dia -linalg -slice -spmm -spmv — implementation pending diff --git a/crates/mohu-sparse/src/lib.rs b/crates/mohu-sparse/src/lib.rs index f3b3c8d..614fd5f 100644 --- a/crates/mohu-sparse/src/lib.rs +++ b/crates/mohu-sparse/src/lib.rs @@ -33,13 +33,12 @@ /// coo.push(5, 7, 2.71); /// let csr = CsrMatrix::from(coo); /// ``` - pub mod arith; pub mod bsr; +pub mod convert; pub mod coo; pub mod csc; pub mod csr; -pub mod convert; pub mod dia; pub mod linalg; pub mod slice; diff --git "a/crates/mohu-special/src/beta\nbessel\nerf\nexpint\ngamma\nmisc\nstats_fn\ntrig.rs" "b/crates/mohu-special/src/beta\nbessel\nerf\nexpint\ngamma\nmisc\nstats_fn\ntrig.rs" deleted file mode 100644 index 6d83c44..0000000 --- "a/crates/mohu-special/src/beta\nbessel\nerf\nexpint\ngamma\nmisc\nstats_fn\ntrig.rs" +++ /dev/null @@ -1,8 +0,0 @@ -// beta -bessel -erf -expint -gamma -misc -stats_fn -trig — implementation pending diff --git a/crates/mohu-special/src/lib.rs b/crates/mohu-special/src/lib.rs index a7c468d..06bf5fb 100644 --- a/crates/mohu-special/src/lib.rs +++ b/crates/mohu-special/src/lib.rs @@ -1,3 +1,4 @@ +pub mod bessel; /// Special mathematical functions for mohu. /// /// Equivalent to `scipy.special` — pure-Rust implementations with @@ -28,9 +29,7 @@ /// Every scalar function is `#[inline(always)]` and designed to auto-vectorise /// under LLVM. `mohu-simd` provides hand-written AVX2 versions for the /// most common (erf, gamma, exp, log) on x86-64. - pub mod beta; -pub mod bessel; pub mod erf; pub mod expint; pub mod gamma; diff --git a/crates/mohu-stats/src/descriptive.rs b/crates/mohu-stats/src/descriptive.rs index e69de29..8b13789 100644 --- a/crates/mohu-stats/src/descriptive.rs +++ b/crates/mohu-stats/src/descriptive.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-stats/src/distributions.rs b/crates/mohu-stats/src/distributions.rs index e69de29..8b13789 100644 --- a/crates/mohu-stats/src/distributions.rs +++ b/crates/mohu-stats/src/distributions.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-stats/src/random.rs b/crates/mohu-stats/src/random.rs index e69de29..8b13789 100644 --- a/crates/mohu-stats/src/random.rs +++ b/crates/mohu-stats/src/random.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-stats/src/sampling.rs b/crates/mohu-stats/src/sampling.rs index e69de29..8b13789 100644 --- a/crates/mohu-stats/src/sampling.rs +++ b/crates/mohu-stats/src/sampling.rs @@ -0,0 +1 @@ + diff --git "a/crates/mohu-testing/src/approx\nassert\ndtype\nfixtures\ngen\nperf.rs" "b/crates/mohu-testing/src/approx\nassert\ndtype\nfixtures\ngen\nperf.rs" deleted file mode 100644 index 3917164..0000000 --- "a/crates/mohu-testing/src/approx\nassert\ndtype\nfixtures\ngen\nperf.rs" +++ /dev/null @@ -1,6 +0,0 @@ -// approx -assert -dtype -fixtures -gen -perf — implementation pending diff --git a/crates/mohu-testing/src/assert.rs b/crates/mohu-testing/src/assert.rs index 0a0ff45..aa59356 100644 --- a/crates/mohu-testing/src/assert.rs +++ b/crates/mohu-testing/src/assert.rs @@ -1 +1,42 @@ -// assert — implementation pending +/// Assert that two equal-length float slices are close within tolerances. +/// +/// # Example +/// +/// ```rust,ignore +/// use mohu_testing::assert_allclose; +/// +/// assert_allclose!(vec![1.0, 2.0], vec![1.0, 2.0000001], atol = 1e-6); +/// ``` +#[macro_export] +macro_rules! assert_allclose { + ($actual:expr, $expected:expr, atol = $atol:expr $(,)?) => { + $crate::assert_allclose!($actual, $expected, rtol = 0.0, atol = $atol) + }; + ($actual:expr, $expected:expr, rtol = $rtol:expr, atol = $atol:expr $(,)?) => {{ + let actual_value = $actual; + let expected_value = $expected; + let actual = ::core::convert::AsRef::<[_]>::as_ref(&actual_value); + let expected = ::core::convert::AsRef::<[_]>::as_ref(&expected_value); + + assert_eq!( + actual.len(), + expected.len(), + "length mismatch: left = {}, right = {}", + actual.len(), + expected.len() + ); + + for (index, (&lhs, &rhs)) in actual.iter().zip(expected.iter()).enumerate() { + let difference = (lhs - rhs).abs(); + let tolerance = $atol + $rtol * lhs.abs().max(rhs.abs()); + assert!( + difference <= tolerance, + "values differ at index {index}: left = {:?}, right = {:?}, diff = {:?}, tolerance = {:?}", + lhs, + rhs, + difference, + tolerance + ); + } + }}; +} diff --git a/crates/mohu-testing/src/lib.rs b/crates/mohu-testing/src/lib.rs index dd7901c..6c26d2b 100644 --- a/crates/mohu-testing/src/lib.rs +++ b/crates/mohu-testing/src/lib.rs @@ -32,13 +32,12 @@ /// } /// } /// ``` - pub mod approx; pub mod assert; pub mod dtype; pub mod fixtures; -pub mod strategies; pub mod perf; +pub mod strategies; pub use mohu_error::{MohuError, MohuResult}; diff --git "a/crates/mohu-ufunc/src/broadcast\ndispatch\nloop_impl\nmacros\nmethods\nreduce\nresolver\ntraits.rs" "b/crates/mohu-ufunc/src/broadcast\ndispatch\nloop_impl\nmacros\nmethods\nreduce\nresolver\ntraits.rs" deleted file mode 100644 index 65d3c32..0000000 --- "a/crates/mohu-ufunc/src/broadcast\ndispatch\nloop_impl\nmacros\nmethods\nreduce\nresolver\ntraits.rs" +++ /dev/null @@ -1,8 +0,0 @@ -// broadcast -dispatch -loop_impl -macros -methods -reduce -resolver -traits — implementation pending diff --git a/crates/mohu-ufunc/src/lib.rs b/crates/mohu-ufunc/src/lib.rs index 9d15f16..50d465f 100644 --- a/crates/mohu-ufunc/src/lib.rs +++ b/crates/mohu-ufunc/src/lib.rs @@ -32,7 +32,6 @@ /// /// Implement [`Ufunc`] and register it in the dispatch table. The macro /// [`define_ufunc!`] generates the boilerplate for common binary/unary cases. - pub mod broadcast; pub mod dispatch; pub mod loop_impl; diff --git a/scripts/README.md b/scripts/README.md new file mode 100644 index 0000000..ed97679 --- /dev/null +++ b/scripts/README.md @@ -0,0 +1,16 @@ +# Windows helper + +Use `use-llvm-mingw.ps1` to run `cargo` with the LLVM-MinGW toolchain that +works on this workspace’s Windows setup. + +Example: + +```powershell +.\scripts\use-llvm-mingw.ps1 +``` + +You can pass extra cargo args after the script name: + +```powershell +.\scripts\use-llvm-mingw.ps1 test -p mohu-io --target x86_64-pc-windows-gnu +``` \ No newline at end of file diff --git a/scripts/use-llvm-mingw.ps1 b/scripts/use-llvm-mingw.ps1 new file mode 100644 index 0000000..29c93ce --- /dev/null +++ b/scripts/use-llvm-mingw.ps1 @@ -0,0 +1,43 @@ +param( + [Parameter(ValueFromRemainingArguments = $true)] + [string[]]$CargoArgs = @('test', '-p', 'mohu-io', '--target', 'x86_64-pc-windows-gnu') +) + +$ErrorActionPreference = 'Stop' + +$repoRoot = Join-Path $PSScriptRoot '..' +Set-Location $repoRoot + +$packageRoot = Join-Path $env:LOCALAPPDATA 'Microsoft\WinGet\Packages' +$llvmPackage = Get-ChildItem $packageRoot -Directory | + Where-Object { $_.Name -like 'MartinStorsjo.LLVM-MinGW.UCRT_*' } | + Sort-Object Name -Descending | + Select-Object -First 1 + +if (-not $llvmPackage) { + throw 'LLVM-MinGW UCRT is not installed. Install MartinStorsjo.LLVM-MinGW.UCRT with winget first.' +} + +$llvmDlltool = Get-ChildItem $llvmPackage.FullName -Recurse -Filter dlltool.exe -File | + Select-Object -First 1 + +if (-not $llvmDlltool) { + throw "Could not find dlltool.exe under $($llvmPackage.FullName)" +} + +$llvmBin = Split-Path $llvmDlltool.FullName -Parent +$selfContained = Join-Path $env:USERPROFILE '.rustup\toolchains\stable-x86_64-pc-windows-gnu\lib\rustlib\x86_64-pc-windows-gnu\bin\self-contained' +$rustLld = Join-Path $env:USERPROFILE '.rustup\toolchains\stable-x86_64-pc-windows-gnu\lib\rustlib\x86_64-pc-windows-gnu\bin\rust-lld.exe' + +if (-not (Test-Path $rustLld)) { + throw "Could not find rust-lld.exe at $rustLld" +} + +$env:Path = "$llvmBin;$selfContained;$env:Path" +$env:CARGO_TARGET_X86_64_PC_WINDOWS_GNU_LINKER = $rustLld + +Write-Host "dlltool => $((Get-Command dlltool.exe).Source)" +Write-Host "gcc => $((Get-Command gcc.exe).Source)" + +& (Join-Path $env:USERPROFILE '.cargo\bin\rustup.exe') run stable-x86_64-pc-windows-gnu ` + (Join-Path $env:USERPROFILE '.cargo\bin\cargo.exe') @CargoArgs \ No newline at end of file