huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-19 19:58:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Wrapping code to call the custom op. (#225)

Browse Source 
* Wrapping code to call the custom op.

* Get the rms example to work.

* Get around rustfmt failing in the CI.

* Fix the rms computation.
This commit is contained in:
Laurent Mazare
2023-07-23 12:31:17 +02:00
committed by GitHub
parent b8a10425ad
commit e449ce53a2
5 changed files with 35 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
candle-core/src/cuda_backend.rs
View File

@ -2,6 +2,7 @@ use crate::backend::{BackendDevice, BackendStorage};
use crate::op::{BinaryOpT, CmpOp, ReduceOp, UnaryOpT};
use crate::{CpuStorage, DType, Layout, Result, Shape, WithDType};
use candle_kernels as kernels;
pub use cudarc;
use cudarc::cublas::{Gemm, GemmConfig, StridedBatchedConfig};
use cudarc::driver::{
CudaFunction, CudaSlice, DeviceRepr, DeviceSlice, LaunchAsync, LaunchConfig, ValidAsZeroBits,
@ -101,7 +102,7 @@ impl std::ops::Deref for CudaDevice {
}
}
trait WrapErr<O> {
pub trait WrapErr<O> {
fn w(self) -> std::result::Result<O, crate::Error>;
}
@ -171,7 +172,7 @@ impl CudaDevice {
})
}
fn get_or_load_func(&self, module_name: &str, ptx: &'static str) -> Result<CudaFunction> {
pub fn get_or_load_func(&self, module_name: &str, ptx: &'static str) -> Result<CudaFunction> {
if !self.has_func(module_name, module_name) {
// Leaking the string here is a bit sad but we need a &'static str and this is only
// done once per kernel name.