Move to the 2024 edition

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "fastnn"
 version = "0.3.2"
-edition = "2021"
+edition = "2024"
 rust-version = "1.87"
 description = "A robust deep learning library in Rust with CUDA kernels"
 license = "MIT"

benches/tensor_ops.rs

@@ -1,7 +1,7 @@
 //! Hot path throughput on the GPU if there is one, otherwise the CPU.
 //! GPU timings include the launch.
 
-use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
 use fastnn::prelude::*;
 
 fn device() -> (Device, &'static str) {

examples/gan.rs

@@ -56,7 +56,10 @@ fn main() {
             let var = sample.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / sample.len() as f32;
             println!(
                 "step {:4}  d_loss {:.3}  g_loss {:.3}  generated mean {mean:.2} std {:.2}  (target {target_mean} / {target_std})",
-                step + 1, d_loss.item(), g_loss.item(), var.sqrt()
+                step + 1,
+                d_loss.item(),
+                g_loss.item(),
+                var.sqrt()
             );
         }
     }

src/autograd/ops/activation.rs

@@ -3,8 +3,8 @@
 
 use crate::autograd::Backward;
 use crate::cuda::{ffi, kernels};
-use crate::tensor::storage::Storage;
 use crate::tensor::Tensor;
+use crate::tensor::storage::Storage;
 
 use super::unary::elementwise;
 
@@ -114,11 +114,7 @@ impl Backward for LeakyReluBackward {
     fn backward(&self, grad: &Tensor) -> Vec<Tensor> {
         let slope = self.slope;
         vec![elementwise(grad, &self.input, move |x| {
-            if x > 0.0 {
-                1.0
-            } else {
-                slope
-            }
+            if x > 0.0 { 1.0 } else { slope }
         })]
     }
     fn name(&self) -> &'static str {

src/autograd/ops/conv.rs

@@ -3,9 +3,9 @@
 
 use crate::autograd::Backward;
 use crate::cuda::kernels;
-use crate::tensor::ops::conv::{fold, unfold, Window};
-use crate::tensor::storage::Storage;
 use crate::tensor::Tensor;
+use crate::tensor::ops::conv::{Window, fold, unfold};
+use crate::tensor::storage::Storage;
 
 /// Gradient of `im2col`: fold the columns back, summing overlaps.
 pub struct Col2ImBackward {

src/autograd/ops/index.rs

@@ -2,8 +2,8 @@
 
 use crate::autograd::Backward;
 use crate::cuda::kernels;
-use crate::tensor::storage::Storage;
 use crate::tensor::Tensor;
+use crate::tensor::storage::Storage;
 
 /// Scatter-add each row's gradient into its source row, so repeats accumulate.
 pub struct IndexSelectBackward {

src/autograd/ops/norm.rs

@@ -2,8 +2,8 @@
 
 use crate::autograd::Backward;
 use crate::cuda::kernels;
-use crate::tensor::storage::Storage;
 use crate::tensor::Tensor;
+use crate::tensor::storage::Storage;
 
 /// LayerNorm on the host. Returns gradients for `(input, gamma, beta)`.
 pub struct LayerNormBackward {

src/autograd/ops/pool.rs

@@ -1,9 +1,9 @@
 //! Derivatives of pooling, which route the gradient the way the forward did.
 
 use crate::autograd::Backward;
+use crate::tensor::Tensor;
 use crate::tensor::ops::conv::Window;
 use crate::tensor::ops::pool::adaptive_range;
-use crate::tensor::Tensor;
 
 /// Only the winning input gets the gradient.
 pub struct MaxPool2dBackward {

src/autograd/ops/reduce.rs

@@ -1,8 +1,8 @@
 //! Derivatives of reductions, which broadcast the gradient back out.
 
 use crate::autograd::Backward;
-use crate::tensor::shape;
 use crate::tensor::Tensor;
+use crate::tensor::shape;
 
 /// Every element gets the whole gradient.
 pub struct SumBackward {

src/autograd/ops/view.rs

@@ -1,8 +1,8 @@
 //! Derivatives of the rearranging ops, each the inverse rearrangement.
 
 use crate::autograd::Backward;
-use crate::tensor::shape;
 use crate::tensor::Tensor;
+use crate::tensor::shape;
 
 use super::reduce_to;
 

src/cuda/buffer.rs

@@ -161,7 +161,7 @@ fn raw_alloc(len: usize) -> Result<*mut f32> {
 /// Free list of blocks keyed by element count. Safe to reuse right away because every kernel
 /// runs on the ordered default stream.
 mod cache {
-    use super::{ffi, HashMap, Mutex};
+    use super::{HashMap, Mutex, ffi};
 
     /// Spare blocks kept per size before memory goes back to the driver.
     const MAX_PER_SIZE: usize = 32;

src/cuda/ffi.rs

@@ -6,7 +6,7 @@
 
 use std::ffi::c_int;
 
-extern "C" {
+unsafe extern "C" {
     pub fn fastnn_cuda_init(device_id: c_int) -> c_int;
     pub fn fastnn_cuda_device_count() -> c_int;
     pub fn fastnn_cuda_synchronize() -> c_int;
@@ -81,7 +81,7 @@ extern "C" {
         cols: c_int,
     ) -> c_int;
     pub fn fastnn_cuda_log_softmax(x: *const f32, out: *mut f32, rows: c_int, cols: c_int)
-        -> c_int;
+    -> c_int;
 
     pub fn fastnn_cuda_layer_norm_forward(
         x: *const f32,

src/cuda/mod.rs

@@ -6,7 +6,7 @@ pub mod buffer;
 pub mod ffi;
 pub mod kernels;
 
-pub use buffer::{empty_cache, CudaBuffer};
+pub use buffer::{CudaBuffer, empty_cache};
 
 use std::sync::OnceLock;
 

src/data/dataset.rs

@@ -1,6 +1,6 @@
 //! The dataset trait and an in-memory implementation.
 
-use crate::tensor::{shape, Tensor};
+use crate::tensor::{Tensor, shape};
 
 /// A finite, indexable collection of `(input, target)` pairs.
 ///

src/data/loader.rs

@@ -1,7 +1,7 @@
 //! Batching and shuffling.
 
 use crate::rng;
-use crate::tensor::{ops::index::index_from_f32, Device, Tensor};
+use crate::tensor::{Device, Tensor, ops::index::index_from_f32};
 
 use super::dataset::Dataset;
 

src/lib.rs

@@ -70,18 +70,18 @@ pub mod prelude {
     pub use crate::autograd::{detect_anomaly, no_grad};
     pub use crate::data::{Batch, DataLoader, Dataset, TensorDataset};
     pub use crate::nn::{
-        bce, bce_with_logits, cross_entropy, huber, kl_divergence, mae, mse, nll,
         AdaptiveAvgPool2d, AvgPool2d, BatchNorm2d, BeamSearch, Buffer, Conv1d, Conv2d,
-        ConvTranspose1d, ConvTranspose2d, CrossEntropyLoss, Dropout, Embedding, Flatten, KvCache,
-        LayerNorm, LeakyReLU, Linear, MaxPool2d, Module, MultiHeadAttention, Param,
+        ConvTranspose1d, ConvTranspose2d, CrossEntropyLoss, Dropout, Embedding, Flatten, GELU, GRU,
+        KvCache, LSTM, LayerNorm, LeakyReLU, Linear, MaxPool2d, Module, MultiHeadAttention, Param,
         PositionalEncoding, RMSNorm, ReLU, Reduction, Reshape, Sampler, Sequential, SiLU, Sigmoid,
         Softmax, StackCache, Tanh, Transformer, TransformerBlock, TransformerDecoder,
-        TransformerDecoderBlock, TransformerStack, GELU, GRU, LSTM,
+        TransformerDecoderBlock, TransformerStack, bce, bce_with_logits, cross_entropy, huber,
+        kl_divergence, mae, mse, nll,
     };
     pub use crate::optim::{
-        clip_grad_norm, clip_grad_value, Adadelta, Adagrad, Adam, AdamW, Constant, CosineAnnealing,
-        Ema, Lion, Lookahead, LrSchedule, OneCycle, Optimizer, OptimizerState, RAdam, RMSprop,
-        StepDecay, Warmup, SGD,
+        Adadelta, Adagrad, Adam, AdamW, Constant, CosineAnnealing, Ema, Lion, Lookahead,
+        LrSchedule, OneCycle, Optimizer, OptimizerState, RAdam, RMSprop, SGD, StepDecay, Warmup,
+        clip_grad_norm, clip_grad_value,
     };
     pub use crate::rng::manual_seed;
     pub use crate::serialize::{

src/nn/attention.rs

@@ -8,7 +8,7 @@ use crate::tensor::{Device, Tensor};
 use super::cache::KvCache;
 use super::dropout::Dropout;
 use super::linear::Linear;
-use super::module::{scoped, Module};
+use super::module::{Module, scoped};
 use super::param::Param;
 
 /// Large negative added to masked scores, so softmax gives them zero weight.

src/nn/decoder.rs

@@ -6,7 +6,7 @@ use crate::tensor::Tensor;
 use super::attention::MultiHeadAttention;
 use super::dropout::Dropout;
 use super::linear::Linear;
-use super::module::{scoped, Module};
+use super::module::{Module, scoped};
 use super::norm::LayerNorm;
 use super::param::Param;
 use super::transformer::{Activation, TransformerStack};

src/nn/embedding.rs

@@ -1,6 +1,6 @@
 //! Embedding table and positional encoding.
 
-use crate::tensor::{ops::index::index_from_f32, Tensor};
+use crate::tensor::{Tensor, ops::index::index_from_f32};
 
 use super::module::Module;
 use super::param::Param;

src/nn/mod.rs

@@ -24,20 +24,20 @@ pub mod sequential;
 pub mod shape;
 pub mod transformer;
 
-pub use activation::{LeakyReLU, ReLU, SiLU, Sigmoid, Softmax, Tanh, GELU};
+pub use activation::{GELU, LeakyReLU, ReLU, SiLU, Sigmoid, Softmax, Tanh};
 pub use attention::MultiHeadAttention;
 pub use beam::BeamSearch;
 pub use cache::{KvCache, StackCache};
 pub use conv::Conv2d;
-pub use conv1d::Conv1d;
 pub use conv_transpose::{ConvTranspose1d, ConvTranspose2d};
+pub use conv1d::Conv1d;
 pub use decoder::{Transformer, TransformerDecoder, TransformerDecoderBlock};
 pub use dropout::Dropout;
 pub use embedding::{Embedding, PositionalEncoding};
 pub use linear::Linear;
 pub use loss::{
-    bce, bce_with_logits, cross_entropy, dice, focal_bce_with_logits, gaussian_nll, huber,
-    kl_divergence, mae, mse, nll, poisson_nll, smooth_l1, CrossEntropyLoss, Reduction,
+    CrossEntropyLoss, Reduction, bce, bce_with_logits, cross_entropy, dice, focal_bce_with_logits,
+    gaussian_nll, huber, kl_divergence, mae, mse, nll, poisson_nll, smooth_l1,
 };
 pub use metric_losses::{contrastive, cosine_embedding, info_nce, margin_ranking, triplet_margin};
 pub use module::Module;

src/nn/sequential.rs

@@ -2,7 +2,7 @@
 
 use crate::tensor::Tensor;
 
-use super::module::{scoped, Module};
+use super::module::{Module, scoped};
 use super::param::{Buffer, Param};
 
 /// Runs its layers in order.

src/nn/transformer.rs

@@ -6,7 +6,7 @@ use super::attention::MultiHeadAttention;
 use super::cache::{KvCache, StackCache};
 use super::dropout::Dropout;
 use super::linear::Linear;
-use super::module::{scoped, Module};
+use super::module::{Module, scoped};
 use super::norm::LayerNorm;
 use super::param::Param;
 

src/serialize/half.rs

@@ -60,11 +60,7 @@ pub fn f32_from_f16(bits: u16) -> f32 {
         (0, 0) => f32::from_bits(sign),
         (0, m) => {
             let magnitude = m as f32 * f32::from_bits(0x3380_0000); // 2^-24
-            if sign != 0 {
-                -magnitude
-            } else {
-                magnitude
-            }
+            if sign != 0 { -magnitude } else { magnitude }
         }
         (0x1F, 0) => f32::from_bits(sign | 0x7F80_0000),
         (0x1F, _) => f32::NAN,

src/serialize/mod.rs

@@ -6,6 +6,6 @@ pub mod half;
 pub mod safetensors;
 pub mod training;
 
-pub use checkpoint::{load, load_tensors, save, save_tensors, Checkpoint};
-pub use safetensors::{load_safetensors, load_safetensors_renamed, save_safetensors, LoadReport};
+pub use checkpoint::{Checkpoint, load, load_tensors, save, save_tensors};
+pub use safetensors::{LoadReport, load_safetensors, load_safetensors_renamed, save_safetensors};
 pub use training::{load_training, save_training};

src/serialize/safetensors.rs

@@ -278,8 +278,8 @@ impl<'a> Header<'a> {
                         "BF16" => Dtype::Bf16,
                         other => {
                             return Err(bad(format!(
-                        "'{name}' has dtype {other}, and only F32/F16/BF16 load into f32 storage"
-                    )))
+                                "'{name}' has dtype {other}, and only F32/F16/BF16 load into f32 storage"
+                            )));
                         }
                     })
                 }

src/serialize/training.rs

@@ -27,7 +27,7 @@ use crate::nn::Module;
 use crate::optim::{Optimizer, OptimizerState};
 use crate::tensor::Tensor;
 
-use super::checkpoint::{load_tensors, save_tensors, Checkpoint};
+use super::checkpoint::{Checkpoint, load_tensors, save_tensors};
 
 /// Name prefixes inside the file. `/` can't clash with the `.` in module paths.
 const PARAM: &str = "param/";

src/tensor/checked.rs

@@ -3,8 +3,8 @@
 
 use crate::error::{Error, Result};
 
-use super::shape;
 use super::Tensor;
+use super::shape;
 
 impl Tensor {
     /// [`matmul`](Tensor::matmul), reporting a mismatch instead of panicking.

src/tensor/core.rs

@@ -4,9 +4,9 @@ use std::sync::Arc;
 
 use crate::autograd::{self, Backward, GradSlot, Node};
 
+use super::Device;
 use super::shape;
 use super::storage::Storage;
-use super::Device;
 
 /// An n-dimensional `f32` array on the CPU or a GPU.
 ///

src/tensor/init.rs

@@ -2,8 +2,8 @@
 
 use crate::rng;
 
-use super::shape;
 use super::Tensor;
+use super::shape;
 
 impl Tensor {
     /// All zeros.

src/tensor/ops/activation.rs

@@ -6,8 +6,8 @@ use crate::autograd::ops::activation::{
     SiluBackward, SoftmaxBackward, TanhBackward,
 };
 use crate::cuda::{ffi, kernels};
-use crate::tensor::storage::Storage;
 use crate::tensor::Tensor;
+use crate::tensor::storage::Storage;
 
 use super::unary_op;
 

src/tensor/ops/conv.rs

@@ -3,8 +3,8 @@
 
 use crate::autograd::ops::conv::{Col2ImBackward, Im2ColBackward};
 use crate::cuda::kernels;
-use crate::tensor::storage::Storage;
 use crate::tensor::Tensor;
+use crate::tensor::storage::Storage;
 
 /// A 2-D sliding window. Dilation `d` spreads the taps to `(k-1)·d + 1` pixels.
 #[derive(Clone, Copy, Debug)]

src/tensor/ops/index.rs

@@ -2,8 +2,8 @@
 
 use crate::autograd::ops::index::IndexSelectBackward;
 use crate::cuda::kernels;
-use crate::tensor::storage::Storage;
 use crate::tensor::Tensor;
+use crate::tensor::storage::Storage;
 
 /// Convert a float index without truncating or saturating it.
 pub(crate) fn index_from_f32(value: f32) -> usize {

src/tensor/ops/matmul.rs

@@ -6,8 +6,8 @@ use rayon::prelude::*;
 
 use crate::autograd::ops::matmul::{MatmulBackward, MatmulNtBackward, MatmulTnBackward};
 use crate::cuda::kernels;
-use crate::tensor::storage::Storage;
 use crate::tensor::Tensor;
+use crate::tensor::storage::Storage;
 
 /// Below this many multiply-adds the thread pool is not worth waking.
 const SERIAL_FLOPS: usize = 16_384;

src/tensor/ops/mod.rs

@@ -16,9 +16,9 @@ use rayon::prelude::*;
 
 use crate::cuda::{ffi, kernels};
 
+use super::Tensor;
 use super::shape;
 use super::storage::Storage;
-use super::Tensor;
 
 /// Below this many elements rayon costs more than it saves.
 const PARALLEL_THRESHOLD: usize = 1 << 15;

src/tensor/ops/norm.rs

@@ -3,8 +3,8 @@
 
 use crate::autograd::ops::norm::{BatchNorm2dBackward, LayerNormBackward, LayerNormCudaBackward};
 use crate::cuda::kernels;
-use crate::tensor::storage::Storage;
 use crate::tensor::Tensor;
+use crate::tensor::storage::Storage;
 
 /// Per-channel statistics of one batch, for updating running estimates.
 pub struct BatchStats {