ronantakizawa/codereview-bench huggingface dataset:
Repository: huggingface/candle
File: candle-examples/examples/onnx/main.rs (Rust)

Code before:

            // Sort the predictions and take the top 5
            let mut top: Vec<_> = prs.iter().enumerate().collect();
            top.sort_by(|a, b| b.1.partial_cmp(a.1).unwrap());
            let top = top.into_iter().take(5).collect::<Vec<_>>();

            // Print the top predictions
            for &(i, p) in &top {
                println!(
                    "{:50}: {:.2}%",
                    candle_examples::imagenet::CLASSES[i],
                    p * 100.0
                );
            }
        }
        Which::EsrGan => {
            let max_pixel_val = candle::Tensor::try_from(255.0f32)?
                .to_device(prs.device())?
                .broadcast_as(prs.shape())?;
            let out = (prs * max_pixel_val)?.i(0)?.to_dtype(candle::DType::U8)?;

            let pb = std::path::PathBuf::from(args.image);
            let input_file_name = pb.file_name().unwrap();
            let mut output_file_name = std::ffi::OsString::from("super_");
            output_file_name.push(input_file_name);
            let mut output_path = std::path::PathBuf::from("./");
            output_path.push(output_file_name);

            save_image(&out, output_path).unwrap();
        }
    }

    Ok(())
}


Review comment (suggestion):
Why create a path of the form `./foo`, I would have thought that passing `output_filename` directly to `save_image` would do the same thing.

Code after:

            // Sort the predictions and take the top 5
            let mut top: Vec<_> = prs.iter().enumerate().collect();
            top.sort_by(|a, b| b.1.partial_cmp(a.1).unwrap());
            let top = top.into_iter().take(5).collect::<Vec<_>>();

            // Print the top predictions
            for &(i, p) in &top {
                println!(
                    "{:50}: {:.2}%",
                    candle_examples::imagenet::CLASSES[i],
                    p * 100.0
                );
            }
        }
        Which::EsrGan => {
            let max_pixel_val = candle::Tensor::try_from(255.0f32)?
                .to_device(prs.device())?
                .broadcast_as(prs.shape())?;
            let out = (prs * max_pixel_val)?.i(0)?.to_dtype(candle::DType::U8)?;

            let pb = std::path::PathBuf::from(args.image);
            let input_file_name = pb.file_name().unwrap();
            let mut output_file_name = std::ffi::OsString::from("super_");
            output_file_name.push(input_file_name);

            save_image(&out, output_file_name)?;
        }
    }

    Ok(())
}
Repository: huggingface/candle
File: candle-metal-kernels/src/lib.rs (Rust)

Code before:
pub mod err;
pub mod kernel;
pub mod kernels;
pub mod metal;
pub mod source;
pub mod utils;

pub use err::MetalKernelError;
pub use kernel::Kernels;
pub use kernels::{
    affine::*, call_binary_contiguous, call_binary_strided, call_mlx_gemm, cast::*, convolution::*,
    fill::*, indexing::*, quantized::*, random::*, reduce::*, sdpa::*, sort::*, ternary::*, unary,
    unary::*, GemmDType, GgmlDType,
};
use metal::{
    BlitCommandEncoder, Buffer, CommandQueue, ComputeCommandEncoder, ComputePipeline,
    ConstantValues, Device, Function, Library, MTLResourceOptions, Value,
};
use objc2_metal::{MTLCompileOptions, MTLMathMode, MTLSize};
use source::Source;
pub use utils::BufferOffset;
use utils::{get_block_dims, linear_split, EncoderParam, EncoderProvider};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DType {
    BF16,
    F16,
    F32,
    I64,
    U32,
    U8,
}

impl DType {
    fn size_in_bytes(&self) -> usize {
        match self {
            Self::U8 => 1,
            Self::U32 => 4,
            Self::I64 => 8,
            Self::BF16 => 2,
            Self::F16 => 2,
            Self::F32 => 4,
        }
    }
}

#[cfg(test)]
mod tests;


Review comment (refactor):
I sort of understand why it's here, but at the same time it feels like an abstraction leak. But I don't have the broader picture of how candle-core and this crate interact so just ignore me 😄 

Code after:
pub mod err;
pub mod kernel;
pub mod kernels;
pub mod metal;
pub mod source;
pub mod utils;

pub use err::MetalKernelError;
pub use kernel::Kernels;
pub use kernels::{
    affine::*, call_binary_contiguous, call_binary_strided, call_mlx_gemm, cast::*, convolution::*,
    fill::*, indexing::*, quantized::*, random::*, reduce::*, sdpa::*, sort::*, ternary::*, unary,
    unary::*, GemmDType, GgmlDType,
};
use metal::{
    BlitCommandEncoder, Buffer, CommandQueue, ComputeCommandEncoder, ComputePipeline,
    ConstantValues, Device, Function, Library, MTLResourceOptions, Value,
};
use objc2_metal::{MTLCompileOptions, MTLMathMode, MTLSize};
use source::Source;
pub use utils::BufferOffset;
use utils::{get_block_dims, linear_split, EncoderParam, EncoderProvider};

pub const RESOURCE_OPTIONS: MTLResourceOptions =
    objc2_metal::MTLResourceOptions(MTLResourceOptions::StorageModeShared.bits());
//| MTLResourceOptions::HazardTrackingModeUntracked.bits(),
//);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DType {
    BF16,
    F16,
    F32,
    I64,
    U32,
    U8,
}

impl DType {
    fn size_in_bytes(&self) -> usize {
        match self {
            Self::U8 => 1,
            Self::U32 => 4,
            Self::I64 => 8,
            Self::BF16 => 2,
            Self::F16 => 2,
            Self::F32 => 4,
        }
    }
}
Repository: huggingface/candle
File: candle-transformers/src/models/stable_diffusion/mod.rs (Rust)

Code before:
            in_channels,
            4,
            use_flash_attn,
            self.unet.clone(),
        )?;
        Ok(unet)
    }

    pub fn build_unet_sharded<P: AsRef<std::path::Path>>(
        &self,
        unet_weight_files: &[P],
        device: &Device,
        in_channels: usize,
        use_flash_attn: bool,
        dtype: DType,
    ) -> Result<unet_2d::UNet2DConditionModel> {
        let vs_unet =
            unsafe { nn::VarBuilder::from_mmaped_safetensors(unet_weight_files, dtype, device)? };
        let unet = unet_2d::UNet2DConditionModel::new(
            vs_unet,
            in_channels,
            4,
            use_flash_attn,
            self.unet.clone(),
        )?;
        Ok(unet)
    }

    pub fn build_scheduler(&self, n_steps: usize) -> Result<Box<dyn Scheduler>> {
        self.scheduler.build(n_steps)
    }
}

pub fn build_clip_transformer<P: AsRef<std::path::Path>>(
    clip: &clip::Config,
    clip_weights: P,
    device: &Device,
    dtype: DType,
) -> Result<clip::ClipTextTransformer> {
    let vs = unsafe { nn::VarBuilder::from_mmaped_safetensors(&[clip_weights], dtype, device)? };
    let text_model = clip::ClipTextTransformer::new(vs, clip)?;
    Ok(text_model)
}


Review comment (suggestion):
Nit, feel free to ignore 😉 
```suggestion
        unet_2d::UNet2DConditionModel::new(
            vs_unet,
            in_channels,
            4,
            use_flash_attn,
            self.unet.clone(),
        )
```

Code after:
            in_channels,
            4,
            use_flash_attn,
            self.unet.clone(),
        )?;
        Ok(unet)
    }

    pub fn build_unet_sharded<P: AsRef<std::path::Path>>(
        &self,
        unet_weight_files: &[P],
        device: &Device,
        in_channels: usize,
        use_flash_attn: bool,
        dtype: DType,
    ) -> Result<unet_2d::UNet2DConditionModel> {
        let vs_unet =
            unsafe { nn::VarBuilder::from_mmaped_safetensors(unet_weight_files, dtype, device)? };
        unet_2d::UNet2DConditionModel::new(
            vs_unet,
            in_channels,
            4,
            use_flash_attn,
            self.unet.clone(),
        )
    }

    pub fn build_scheduler(&self, n_steps: usize) -> Result<Box<dyn Scheduler>> {
        self.scheduler.build(n_steps)
    }
}

pub fn build_clip_transformer<P: AsRef<std::path::Path>>(
    clip: &clip::Config,
    clip_weights: P,
    device: &Device,
    dtype: DType,
) -> Result<clip::ClipTextTransformer> {
    let vs = unsafe { nn::VarBuilder::from_mmaped_safetensors(&[clip_weights], dtype, device)? };
    let text_model = clip::ClipTextTransformer::new(vs, clip)?;
    Ok(text_model)
}
Repository: huggingface/candle
File: candle-nn/benches/benchmarks/conv.rs (Rust)

Code before:

const B: usize = 1;
const C: usize = 1;

fn run(input: Tensor, weight: Tensor, bias: Option<Tensor>, config: Conv2dConfig) {
    Conv2d::new(weight, bias, config).forward(&input).unwrap();
}

fn run_conv2d_benchmark(
    c: &mut Criterion,
    device: &Device,
    dtype: DType,
    k_size: usize,
    m: usize,
    bias: bool,
    name: &str,
) {
    let weight = Tensor::ones((1, C, k_size, k_size), dtype, device)
        .unwrap()
        .to_dtype(dtype)
        .unwrap();
    let bias = if bias {
        Some(Tensor::zeros(m, dtype, device).unwrap())
    } else {
        None
    };
    let input = Tensor::ones((B, C, m, m), dtype, device).unwrap();

    let mut group = c.benchmark_group(device.bench_name(name));
    group.bench_function("iter", move |b| {
        b.iter_custom(|iters| {
            let start = Instant::now();
            for _i in 0..iters {
                run(
                    black_box(input.clone()),
                    black_box(weight.clone()),
                    black_box(bias.clone()),
                    Default::default(),
                );
            }
            device.sync().unwrap();
            start.elapsed()
        })
    });
    group.finish();
}

fn criterion_benchmark(c: &mut Criterion) {
    let device = BenchDeviceHandler::new().unwrap();
    for d in device.devices {
        run_conv2d_benchmark(c, &d, DType::F32, 3, 128, true, "conv2d_f32_i128_k3x3_b");

Review comment (suggestion):
```suggestion
    let input = Tensor::ones((B, C, m, m), dtype, device).unwrap();
    let name = format!("conv2d_{dtype:?}_i{m}_k{k}x{k}_{}", if bias { "b" } else { "nb" });
```

Code after:

const B: usize = 1;
const C: usize = 1;

fn run(input: Tensor, weight: Tensor, bias: Option<Tensor>, config: Conv2dConfig) {
    Conv2d::new(weight, bias, config).forward(&input).unwrap();
}

fn run_conv2d_benchmark(
    c: &mut Criterion,
    device: &Device,
    dtype: DType,
    k_size: usize,
    m: usize,
    bias: bool,
) {
    let weight = Tensor::ones((1, C, k_size, k_size), dtype, device)
        .unwrap()
        .to_dtype(dtype)
        .unwrap();
    let bias_t = if bias {
        Some(Tensor::zeros(m, dtype, device).unwrap())
    } else {
        None
    };
    let input = Tensor::ones((B, C, m, m), dtype, device).unwrap();
    let name = format!(
        "conv2d_{dtype:?}_i{m}_k{k_size}x{k_size}_{}",
        if bias { "b" } else { "nb" }
    );

    let mut group = c.benchmark_group(device.bench_name(name));
    group.bench_function("iter", move |b| {
        b.iter_custom(|iters| {
            let start = Instant::now();
            for _i in 0..iters {
                run(
                    black_box(input.clone()),
                    black_box(weight.clone()),
                    black_box(bias_t.clone()),
                    Default::default(),
                );
            }
            device.sync().unwrap();
            start.elapsed()
        })
    });
    group.finish();
}

fn criterion_benchmark(c: &mut Criterion) {
Repository: huggingface/candle
File: candle-nn/benches/benchmarks/conv.rs (Rust)

Code before:
            for _i in 0..iters {
                run(
                    black_box(input.clone()),
                    black_box(weight.clone()),
                    black_box(bias.clone()),
                    Default::default(),
                );
            }
            device.sync().unwrap();
            start.elapsed()
        })
    });
    group.finish();
}

fn criterion_benchmark(c: &mut Criterion) {
    let device = BenchDeviceHandler::new().unwrap();
    for d in device.devices {
        run_conv2d_benchmark(c, &d, DType::F32, 3, 128, true, "conv2d_f32_i128_k3x3_b");
        run_conv2d_benchmark(c, &d, DType::F32, 1, 128, false, "conv2d_f32_i128_k1x1_nb");
        run_conv2d_benchmark(c, &d, DType::F32, 5, 128, false, "conv2d_f32_i128_k5x5_nb");
        run_conv2d_benchmark(c, &d, DType::F32, 3, 512, false, "conv2d_f32_i512_k3x3_nb");
        run_conv2d_benchmark(c, &d, DType::F16, 3, 128, true, "conv2d_f16_i128_k3x3_b");
        run_conv2d_benchmark(c, &d, DType::F16, 1, 128, false, "conv2d_f16_i128_k1x1_nb");
        run_conv2d_benchmark(c, &d, DType::F16, 5, 128, false, "conv2d_f16_i128_k5x5_nb");
        run_conv2d_benchmark(c, &d, DType::F16, 5, 512, false, "conv2d_f16_i512_k3x3_nb");
    }
}

criterion_group!(benches, criterion_benchmark);


Review comment (suggestion):
```suggestion
        run_conv2d_benchmark(c, &d, DType::F32, 3, 128, true);
        run_conv2d_benchmark(c, &d, DType::F32, 1, 128, false);
        run_conv2d_benchmark(c, &d, DType::F32, 5, 128, false);
        run_conv2d_benchmark(c, &d, DType::F32, 3, 512, false);
        run_conv2d_benchmark(c, &d, DType::F16, 3, 128, true);
        run_conv2d_benchmark(c, &d, DType::F16, 1, 128, false);
        run_conv2d_benchmark(c, &d, DType::F16, 5, 128, false);
        run_conv2d_benchmark(c, &d, DType::F16, 5, 512, false);
```

Code after:
    group.bench_function("iter", move |b| {
        b.iter_custom(|iters| {
            let start = Instant::now();
            for _i in 0..iters {
                run(
                    black_box(input.clone()),
                    black_box(weight.clone()),
                    black_box(bias_t.clone()),
                    Default::default(),
                );
            }
            device.sync().unwrap();
            start.elapsed()
        })
    });
    group.finish();
}

fn criterion_benchmark(c: &mut Criterion) {
    let device = BenchDeviceHandler::new().unwrap();
    for d in device.devices {
        run_conv2d_benchmark(c, &d, DType::F32, 3, 128, true);
        run_conv2d_benchmark(c, &d, DType::F32, 1, 128, false);
        run_conv2d_benchmark(c, &d, DType::F32, 5, 128, false);
        run_conv2d_benchmark(c, &d, DType::F32, 3, 512, false);
        run_conv2d_benchmark(c, &d, DType::F16, 3, 128, true);
        run_conv2d_benchmark(c, &d, DType::F16, 1, 128, false);
        run_conv2d_benchmark(c, &d, DType::F16, 5, 128, false);
        run_conv2d_benchmark(c, &d, DType::F16, 5, 512, false);
    }
}

criterion_group!(benches, criterion_benchmark);
Repository: huggingface/candle
File: candle-metal-kernels/src/metal/commands.rs (Rust)

Code before:
        state: &mut EntryState,
        reset_to: usize,
    ) -> Result<(), MetalKernelError> {
        state.current.commit();
        let new_cb = create_command_buffer(&self.command_queue, Arc::clone(&entry.semaphore))?;
        let old_cb = std::mem::replace(&mut state.current, new_cb);
        state.in_flight.push(old_cb);
        entry.compute_count.store(reset_to, Ordering::Release);

        Ok(())
    }

    fn ensure_completed(cb: &CommandBuffer) -> Result<(), MetalKernelError> {
        match cb.status() {
            MTLCommandBufferStatus::NotEnqueued | MTLCommandBufferStatus::Enqueued => {
                cb.commit();
                cb.wait_until_completed();
            }
            MTLCommandBufferStatus::Committed | MTLCommandBufferStatus::Scheduled => {
                cb.wait_until_completed();
            }
            MTLCommandBufferStatus::Completed => {}
            MTLCommandBufferStatus::Error => {
                let msg = cb
                    .error()
                    .map(|e| e.into_owned())
                    .unwrap_or_else(|| "unknown error".to_string());
                return Err(MetalKernelError::CommandBufferError(msg));
            }
            _ => unreachable!(),
        }

        Ok(())
    }
}

impl Drop for Commands {
    fn drop(&mut self) {
        // TODO: Avoid redundant allocation before drop
        let _ = self.flush();
    }
}


Review comment (suggestion):
```suggestion
                    .map(|e| e.to_string())
```


Code after:
        state: &mut EntryState,
        reset_to: usize,
    ) -> Result<(), MetalKernelError> {
        state.current.commit();
        let new_cb = create_command_buffer(&self.command_queue, Arc::clone(&entry.semaphore))?;
        let old_cb = std::mem::replace(&mut state.current, new_cb);
        state.in_flight.push(old_cb);
        entry.compute_count.store(reset_to, Ordering::Release);

        Ok(())
    }

    fn ensure_completed(cb: &CommandBuffer) -> Result<(), MetalKernelError> {
        match cb.status() {
            MTLCommandBufferStatus::NotEnqueued | MTLCommandBufferStatus::Enqueued => {
                cb.commit();
                cb.wait_until_completed();
            }
            MTLCommandBufferStatus::Committed | MTLCommandBufferStatus::Scheduled => {
                cb.wait_until_completed();
            }
            MTLCommandBufferStatus::Completed => {}
            MTLCommandBufferStatus::Error => {
                let msg = cb
                    .error()
                    .map(|e| e.to_string())
                    .unwrap_or_else(|| "unknown error".to_string());
                return Err(MetalKernelError::CommandBufferError(msg));
            }
            _ => unreachable!(),
        }

        Ok(())
    }
}

impl Drop for Commands {
    fn drop(&mut self) {
        // TODO: Avoid redundant allocation before drop
        let _ = self.flush();
    }
}
Repository: huggingface/candle
File: candle-metal-kernels/src/tests.rs (Rust)

Code before:
    std::env::set_var("CANDLE_METAL_COMPUTE_PER_BUFFER", "2");

    let device = Device::system_default().unwrap();
    let queue = device.new_command_queue().unwrap();
    let commands = Commands::new(queue).unwrap();

    for _ in 0..10 {
        let (_flush, encoder) = commands.command_encoder().unwrap();
        drop(encoder);
    }

    commands.wait_until_completed().unwrap();
}

#[test]
fn commands_concurrent_acquisition() {
    std::env::set_var("CANDLE_METAL_COMPUTE_PER_BUFFER", "2");
    std::env::set_var("CANDLE_METAL_COMMAND_POOL_SIZE", "4");

    let device = Device::system_default().unwrap();
    let queue = device.new_command_queue().unwrap();
    let commands = Arc::new(Commands::new(queue).unwrap());

    let mut handles = vec![];

    for _ in 0..16 {
        let c = Arc::clone(&commands);
        handles.push(thread::spawn(move || {
            let (_flush, encoder) = c.command_encoder().unwrap();
            drop(encoder);
        }));
    }

    for h in handles {
        h.join().unwrap();
    }

    commands.wait_until_completed().unwrap();
}


Review comment (suggestion):
I like the test. Not sure I understand the significance of this line.
Is it to ensure there are no errors reported via the command buffer status?
If so worth noting that I've never actually gotten an error through that channel hehe

Code after:
}

#[test]
fn commands_concurrent_acquisition() {
    std::env::set_var("CANDLE_METAL_COMPUTE_PER_BUFFER", "2");
    std::env::set_var("CANDLE_METAL_COMMAND_POOL_SIZE", "4");

    let device = Device::system_default().unwrap();
    let queue = device.new_command_queue().unwrap();
    let commands = Arc::new(Commands::new(queue).unwrap());

    let mut handles = vec![];

    for _ in 0..16 {
        let c = Arc::clone(&commands);
        handles.push(thread::spawn(move || {
            let (_flush, encoder) = c.command_encoder().unwrap();
            drop(encoder);
        }));
    }

    for h in handles {
        h.join().unwrap();
    }

    commands.wait_until_completed().unwrap();
}
Repository: huggingface/candle
File: candle-nn/src/loss.rs (Rust)

Code before:
//! Loss Calculations
//!
use candle::{Result, Tensor};

pub trait Loss: Sized {
    type Config: Sized;
    fn new(config: Self::Config) -> Result<Self>;
    fn forward(&mut self, inp: &Tensor, target: &Tensor) -> Result<Tensor>;
}

/// HuberLoss
///
/// A robust loss function that combines `MAE` and `MSE` losses:
/// - When the absolute element-wise error is less than `delta`, it uses a squared term (MSE loss).
/// - When the absolute element-wise error is greater than or equal to `delta`, it uses a linear term (MAE loss scaled by `delta`).
/// # Formula
///
/// HuberLoss =
/// 0.5(x_n - y_n)^2, & |x_n - y_n| < delta
/// delta(|x_n - y_n| - 0.5delta), & |x_n - y_n| >= delta
///
pub struct HuberLoss {
    delta: f64,
}
impl Default for HuberLoss {
    fn default() -> Self {
        Self { delta: 1.0 }
    }
}
impl Loss for HuberLoss {
    type Config = f64;
    fn new(delta: f64) -> Result<Self> {
        Ok(Self { delta: delta })
    }

Review comment (suggestion):
Traits are at their most useful when you use them to describe generic behaviour, and in rust you typically either use generics ala `fn calculate<L: Loss>(t1: Tensor, t2: Tensor, loss: L)` or as a dyn trait `fn calculate(t1: Tensor, t2: Tensor, loss: dyn Loss)`.
However this trait is not [dyn-compatible](https://doc.rust-lang.org/reference/items/traits.html#dyn-compatibility), so you're left with generics.

I like generics a lot, but I don't think `Loss` is a particulary good fit as `<L: Loss>` would permeate all code that in any way relates to the usage of the loss function. I think it is better if we stick to the same approach as the other loss functions in this file.

Code after:
//! Loss Calculations
//!
use candle::{Result, Tensor};

/// The negative log likelihood loss.
///
/// Arguments
///
/// * [inp]: The input tensor of dimensions `N, C` where `N` is the batch size and `C` the number
///   of categories. This is expected to contain log probabilities.
/// * [target]: The ground truth labels as a tensor of u32 of dimension `N`.
///
/// The resulting tensor is a scalar containing the average value over the batch.
pub fn nll(inp: &Tensor, target: &Tensor) -> Result<Tensor> {
    let b_sz = match target.dims() {
        &[b_sz] => b_sz,
        dims => candle::bail!("the target tensor should have a single dimension ({dims:?})"),
    };
    match inp.dims() {
        &[inp_b_sz, _] => {
            if inp_b_sz != b_sz {
                candle::bail!("batch size mismatch between inp ({inp_b_sz}) and target ({b_sz})")
            }
        }
        dims => candle::bail!("the target tensor should have two dimensions ({dims:?})"),
    }
    inp.gather(&target.unsqueeze(1)?, 1)?
        .sum_all()?
        .affine(-1f64 / b_sz as f64, 0.)
}

/// The cross-entropy loss.
///
/// Arguments