huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 18:48:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

UG metal integration. (#2580)

Browse Source
This commit is contained in:
Laurent Mazare
2024-10-27 15:20:37 +01:00
committed by GitHub
parent 594d984f9c
commit 0e2c8c17fb
8 changed files with 92 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
candle-core/tests/custom_op_tests.rs
View File

@ -144,7 +144,7 @@ fn inplace_op1() -> Result<()> {
Ok(())
}
#[cfg(feature = "cuda")]
#[cfg(any(feature = "cuda", feature = "metal"))]
#[allow(clippy::approx_constant)]
#[test]
fn ug_op() -> Result<()> {
@ -160,15 +160,21 @@ fn ug_op() -> Result<()> {
let opts: ug::lower_op::Opts = Default::default();
kernel.lower(&opts.with_global(0, 12))?
};
let device = Device::new_cuda(0)?;
let device = if candle_core::utils::cuda_is_available() {
Device::new_cuda(0)?
} else if candle_core::utils::metal_is_available() {
Device::new_metal(0)?
} else {
candle_core::bail!("metal/cuda is mandatory for this test")
};
let op = candle_core::UgIOp1::new("test", kernel, &device)?;
let t = Tensor::arange(0u32, 12u32, &device)?.to_dtype(DType::F32)?;
t.inplace_op1(&op)?;
assert_eq!(
to_vec1_round(&t, 4)?,
to_vec1_round(&t, 2)?,
&[
1.0, 2.7183, 7.3891, 20.0855, 54.5982, 148.4132, 403.4287, 1096.6334, 2980.9578,
8103.0806, 22026.469, 59874.133
1.0, 2.72, 7.39, 20.09, 54.6, 148.41, 403.43, 1096.63, 2980.96, 8103.08, 22026.47,
59874.13
]
);
Ok(())