Skip to content

Commit f0f8991

Browse files
guoqingbaoguoqingbao
committed
Fix build
1 parent 0c9ea12 commit f0f8991

7 files changed

Lines changed: 701 additions & 884 deletions

File tree

‎candle-core/Cargo.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ ug-metal = { workspace = true, optional = true }
3232
yoke = { workspace = true }
3333
zip = { workspace = true }
3434

35-
ubridge = { git = "https://github.com/EnflameTechnology/Ubridge.git", version = "0.1.0", optional = true, rev = "36b9d09" }
35+
ubridge = { git = "https://github.com/EnflameTechnology/Ubridge.git", version = "0.1.0", optional = true, rev = "0ca4364" }
3636
uhal = { git = "https://github.com/EnflameTechnology/UHHI.git", version = "0.1.0", optional = true, rev = "0980208" }
3737
cust_core = { git = "https://github.com/EnflameTechnology/UHHI.git", version = "0.1.0", optional = true, rev = "0980208" }
3838

@@ -51,7 +51,7 @@ gcu = ["ubridge", "uhal", "cust_core", "ubridge/graph"]
5151
scorpio = ["ubridge", "uhal", "cust_core", "ubridge/scorpio"]
5252
cuda = ["cudarc", "dep:candle-kernels"]
5353
cudnn = ["cuda", "cudarc/cudnn"]
54-
graph = ["cuda", "cudarc/graph"]
54+
graph = ["ubridge/graph"]
5555
mkl = ["dep:libc", "dep:intel-mkl-src"]
5656
accelerate = ["dep:libc", "dep:accelerate-src"]
5757
metal = ["dep:metal", "dep:candle-metal-kernels", "dep:ug-metal"]

‎candle-core/src/cpu_backend/mod.rs‎

Lines changed: 3 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2167,9 +2167,7 @@ impl BackendStorage for CpuStorage {
21672167
) -> Result<()> {
21682168
match (self, dst) {
21692169
(Self::U8(src), Self::U8(dst)) => copy2d_(src, dst, d1, d2, src_s, dst_s, src_o, dst_o),
2170-
(Self::I8(src), Self::I8(dst)) => {
2171-
copy2d_(src, dst, d1, d2, src_s, dst_s, src_o, dst_o)
2172-
}
2170+
(Self::I8(src), Self::I8(dst)) => copy2d_(src, dst, d1, d2, src_s, dst_s, src_o, dst_o),
21732171
(Self::U32(src), Self::U32(dst)) => {
21742172
copy2d_(src, dst, d1, d2, src_s, dst_s, src_o, dst_o)
21752173
}
@@ -2555,11 +2553,7 @@ impl BackendDevice for CpuDevice {
25552553
let elem_count = shape.elem_count();
25562554
let mut rng = rand::rng();
25572555
match dtype {
2558-
DType::U8
2559-
| DType::I8
2560-
| DType::U32
2561-
| DType::I32
2562-
| DType::I64 => {
2556+
DType::U8 | DType::I8 | DType::U32 | DType::I32 | DType::I64 => {
25632557
Err(Error::UnsupportedDTypeForOp(dtype, "rand_uniform").bt())
25642558
}
25652559
DType::BF16 => {
@@ -2606,11 +2600,7 @@ impl BackendDevice for CpuDevice {
26062600
let elem_count = shape.elem_count();
26072601
let mut rng = rand::rng();
26082602
match dtype {
2609-
DType::U8
2610-
| DType::I8
2611-
| DType::U32
2612-
| DType::I32
2613-
| DType::I64 => {
2603+
DType::U8 | DType::I8 | DType::U32 | DType::I32 | DType::I64 => {
26142604
Err(Error::UnsupportedDTypeForOp(dtype, "rand_normal").bt())
26152605
}
26162606
DType::BF16 => {

‎candle-core/src/gcu_backend.rs‎

Lines changed: 35 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2477,15 +2477,42 @@ impl BackendStorage for GcuStorage {
24772477

24782478
fn copy2d(
24792479
&self,
2480-
_dst: &mut Self,
2481-
_d1: usize,
2482-
_d2: usize,
2483-
_src_s: usize,
2484-
_dst_s: usize,
2485-
_src_o: usize,
2486-
_dst_o: usize,
2480+
dst: &mut Self,
2481+
d1: usize,
2482+
d2: usize,
2483+
src_s: usize,
2484+
dst_s: usize,
2485+
src_o: usize,
2486+
dst_o: usize,
24872487
) -> Result<()> {
2488-
todo!()
2488+
if d1 == 0 || d2 == 0 {
2489+
return Ok(());
2490+
}
2491+
let dev = &self.device;
2492+
macro_rules! copy2d_impl {
2493+
($src_slice:expr, $dst_slice:expr) => {{
2494+
for i in 0..d1 {
2495+
let s_off = src_o + i * src_s;
2496+
let d_off = dst_o + i * dst_s;
2497+
let src_row = $src_slice.slice(s_off..s_off + d2);
2498+
let mut dst_row = $dst_slice.slice(d_off..d_off + d2);
2499+
dev.dtod_copy(&src_row, &mut dst_row).w()?;
2500+
}
2501+
}};
2502+
}
2503+
match (&self.slice, &mut dst.slice) {
2504+
(S::U8(s), S::U8(d)) => copy2d_impl!(s, d),
2505+
(S::I8(s), S::I8(d)) => copy2d_impl!(s, d),
2506+
(S::U32(s), S::U32(d)) => copy2d_impl!(s, d),
2507+
(S::I32(s), S::I32(d)) => copy2d_impl!(s, d),
2508+
(S::I64(s), S::I64(d)) => copy2d_impl!(s, d),
2509+
(S::BF16(s), S::BF16(d)) => copy2d_impl!(s, d),
2510+
(S::F16(s), S::F16(d)) => copy2d_impl!(s, d),
2511+
(S::F32(s), S::F32(d)) => copy2d_impl!(s, d),
2512+
(S::F64(s), S::F64(d)) => copy2d_impl!(s, d),
2513+
_ => Err(GcuError::InternalError("dtype mismatch in copy2d"))?,
2514+
}
2515+
Ok(())
24892516
}
24902517

24912518
fn copy_strided_src(&self, dst: &mut Self, dst_offset: usize, src_l: &Layout) -> Result<()> {

‎candle-core/src/offloadable.rs‎

Lines changed: 54 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -110,12 +110,10 @@ impl OffloadBuffer {
110110
}
111111
}
112112
#[cfg(feature = "cuda")]
113-
Device::Cuda(_) => {
114-
unsafe {
115-
ptr_host = result::malloc_host(size, 1).unwrap();
116-
std::ptr::copy(src.as_ptr() as *mut core::ffi::c_void, ptr_host, size);
117-
}
118-
}
113+
Device::Cuda(_) => unsafe {
114+
ptr_host = result::malloc_host(size, 1).unwrap();
115+
std::ptr::copy(src.as_ptr() as *mut core::ffi::c_void, ptr_host, size);
116+
},
119117
_ => {
120118
crate::bail!("offload buffer only for cuda or gcu device tensors")
121119
}
@@ -134,72 +132,77 @@ impl OffloadBuffer {
134132
#[cfg(feature = "gcu")]
135133
Device::Gcu(dev) => {
136134
let storage = match self.dtype {
137-
DType::BF16 => crate::Storage::Gcu(dev.storage_from_buffer(
135+
DType::BF16 => crate::Storage::Gcu(
136+
dev.storage_from_buffer(self.ptr_host as *mut bf16, self.len)?,
137+
),
138+
DType::F16 => crate::Storage::Gcu(
139+
dev.storage_from_buffer(self.ptr_host as *mut f16, self.len)?,
140+
),
141+
DType::F32 => crate::Storage::Gcu(
142+
dev.storage_from_buffer(self.ptr_host as *mut f32, self.len)?,
143+
),
144+
DType::U8 => crate::Storage::Gcu(
145+
dev.storage_from_buffer(self.ptr_host as *mut u8, self.len)?,
146+
),
147+
DType::U32 => crate::Storage::Gcu(
148+
dev.storage_from_buffer(self.ptr_host as *mut u32, self.len)?,
149+
),
150+
DType::I8 => crate::Storage::Gcu(
151+
dev.storage_from_buffer(self.ptr_host as *mut i8, self.len)?,
152+
),
153+
DType::I32 => crate::Storage::Gcu(
154+
dev.storage_from_buffer(self.ptr_host as *mut i32, self.len)?,
155+
),
156+
DType::I64 => crate::Storage::Gcu(
157+
dev.storage_from_buffer(self.ptr_host as *mut i64, self.len)?,
158+
),
159+
DType::F64 => crate::Storage::Gcu(
160+
dev.storage_from_buffer(self.ptr_host as *mut f64, self.len)?,
161+
),
162+
};
163+
Ok(storage)
164+
}
165+
#[cfg(feature = "cuda")]
166+
Device::Cuda(dev) => {
167+
let storage = match self.dtype {
168+
DType::BF16 => Storage::Cuda(storage_from_buffer(
138169
self.ptr_host as *mut bf16,
139170
self.len,
171+
dev,
140172
)?),
141-
DType::F16 => crate::Storage::Gcu(dev.storage_from_buffer(
173+
DType::F16 => Storage::Cuda(storage_from_buffer(
142174
self.ptr_host as *mut f16,
143175
self.len,
176+
dev,
144177
)?),
145-
DType::F32 => crate::Storage::Gcu(dev.storage_from_buffer(
178+
DType::F32 => Storage::Cuda(storage_from_buffer(
146179
self.ptr_host as *mut f32,
147180
self.len,
181+
dev,
148182
)?),
149-
DType::U8 => crate::Storage::Gcu(dev.storage_from_buffer(
183+
DType::U8 => Storage::Cuda(storage_from_buffer(
150184
self.ptr_host as *mut u8,
151185
self.len,
186+
dev,
152187
)?),
153-
DType::U32 => crate::Storage::Gcu(dev.storage_from_buffer(
188+
DType::U32 => Storage::Cuda(storage_from_buffer(
154189
self.ptr_host as *mut u32,
155190
self.len,
191+
dev,
156192
)?),
157-
DType::I8 => crate::Storage::Gcu(dev.storage_from_buffer(
158-
self.ptr_host as *mut i8,
159-
self.len,
160-
)?),
161-
DType::I32 => crate::Storage::Gcu(dev.storage_from_buffer(
162-
self.ptr_host as *mut i32,
163-
self.len,
164-
)?),
165-
DType::I64 => crate::Storage::Gcu(dev.storage_from_buffer(
193+
DType::I64 => Storage::Cuda(storage_from_buffer(
166194
self.ptr_host as *mut i64,
167195
self.len,
196+
dev,
168197
)?),
169-
DType::F64 => crate::Storage::Gcu(dev.storage_from_buffer(
198+
DType::F64 => Storage::Cuda(storage_from_buffer(
170199
self.ptr_host as *mut f64,
171200
self.len,
201+
dev,
172202
)?),
173203
};
174204
Ok(storage)
175205
}
176-
#[cfg(feature = "cuda")]
177-
Device::Cuda(dev) => {
178-
let storage = match self.dtype {
179-
DType::BF16 => {
180-
Storage::Cuda(storage_from_buffer(self.ptr_host as *mut bf16, self.len, dev)?)
181-
}
182-
DType::F16 => {
183-
Storage::Cuda(storage_from_buffer(self.ptr_host as *mut f16, self.len, dev)?)
184-
}
185-
DType::F32 => {
186-
Storage::Cuda(storage_from_buffer(self.ptr_host as *mut f32, self.len, dev)?)
187-
}
188-
DType::U8 => {
189-
Storage::Cuda(storage_from_buffer(self.ptr_host as *mut u8, self.len, dev)?)
190-
}
191-
DType::U32 => {
192-
Storage::Cuda(storage_from_buffer(self.ptr_host as *mut u32, self.len, dev)?)
193-
}
194-
DType::I64 => {
195-
Storage::Cuda(storage_from_buffer(self.ptr_host as *mut i64, self.len, dev)?)
196-
}
197-
DType::F64 => {
198-
Storage::Cuda(storage_from_buffer(self.ptr_host as *mut f64, self.len, dev)?)
199-
}
200-
};
201-
Ok(storage)
202-
}
203206
_ => crate::bail!("not supported device for cpu offloading"),
204207
}
205208
}
@@ -219,11 +222,9 @@ impl Drop for OffloadBuffer {
219222
}
220223
}
221224
#[cfg(feature = "cuda")]
222-
Device::Cuda(_) => {
223-
unsafe {
224-
let _ = result::free_host(self.ptr_host);
225-
}
226-
}
225+
Device::Cuda(_) => unsafe {
226+
let _ = result::free_host(self.ptr_host);
227+
},
227228
_ => {}
228229
}
229230
}

‎candle-nn/src/layer_norm.rs‎

Lines changed: 55 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -107,30 +107,51 @@ impl LayerNorm {
107107

108108
impl Module for LayerNorm {
109109
fn forward(&self, x: &Tensor) -> Result<Tensor> {
110-
if x.is_contiguous() && self.remove_mean {
111-
if let Some(bias) = self.bias.as_ref() {
112-
return crate::ops::layer_norm(x, &self.weight, bias, self.eps as f32);
113-
}
110+
#[cfg(feature = "gcu")]
111+
{
112+
let x = if x.is_contiguous() {
113+
x.clone()
114+
} else {
115+
x.contiguous()?
116+
};
117+
let bias_tensor = match &self.bias {
118+
Some(b) => b.clone(),
119+
None => Tensor::zeros_like(&self.weight)?,
120+
};
121+
let op = candle::gcu_backend::LayerNorm {
122+
eps: self.eps as f32,
123+
remove_mean: self.remove_mean,
124+
affine: self.bias.is_some(),
125+
};
126+
return x.apply_op3_no_bwd(&self.weight, &bias_tensor, &op);
114127
}
115-
let x_dtype = x.dtype();
116-
let internal_dtype = match x_dtype {
117-
DType::F16 | DType::BF16 => DType::F32,
118-
d => d,
119-
};
120-
let hidden_size = x.dim(D::Minus1)?;
121-
let x = x.to_dtype(internal_dtype)?;
122-
let x = if self.remove_mean {
123-
let mean_x = (x.sum_keepdim(D::Minus1)? / hidden_size as f64)?;
124-
x.broadcast_sub(&mean_x)?
125-
} else {
126-
x
127-
};
128-
let norm_x = (x.sqr()?.sum_keepdim(D::Minus1)? / hidden_size as f64)?;
129-
let x_normed = x.broadcast_div(&(norm_x + self.eps)?.sqrt()?)?;
130-
let x = x_normed.to_dtype(x_dtype)?.broadcast_mul(&self.weight)?;
131-
match &self.bias {
132-
None => Ok(x),
133-
Some(bias) => x.broadcast_add(bias),
128+
#[cfg(not(feature = "gcu"))]
129+
{
130+
if x.is_contiguous() && self.remove_mean {
131+
if let Some(bias) = self.bias.as_ref() {
132+
return crate::ops::layer_norm(x, &self.weight, bias, self.eps as f32);
133+
}
134+
}
135+
let x_dtype = x.dtype();
136+
let internal_dtype = match x_dtype {
137+
DType::F16 | DType::BF16 => DType::F32,
138+
d => d,
139+
};
140+
let hidden_size = x.dim(D::Minus1)?;
141+
let x = x.to_dtype(internal_dtype)?;
142+
let x = if self.remove_mean {
143+
let mean_x = (x.sum_keepdim(D::Minus1)? / hidden_size as f64)?;
144+
x.broadcast_sub(&mean_x)?
145+
} else {
146+
x
147+
};
148+
let norm_x = (x.sqr()?.sum_keepdim(D::Minus1)? / hidden_size as f64)?;
149+
let x_normed = x.broadcast_div(&(norm_x + self.eps)?.sqrt()?)?;
150+
let x = x_normed.to_dtype(x_dtype)?.broadcast_mul(&self.weight)?;
151+
match &self.bias {
152+
None => Ok(x),
153+
Some(bias) => x.broadcast_add(bias),
154+
}
134155
}
135156
}
136157
}
@@ -185,10 +206,17 @@ impl RmsNorm {
185206

186207
impl Module for RmsNorm {
187208
fn forward(&self, xs: &Tensor) -> Result<Tensor> {
188-
if xs.is_contiguous() {
189-
crate::ops::rms_norm(xs, &self.0.weight, self.0.eps as f32)
190-
} else {
191-
self.0.forward(xs)
209+
#[cfg(feature = "gcu")]
210+
{
211+
return self.0.forward(xs);
212+
}
213+
#[cfg(not(feature = "gcu"))]
214+
{
215+
if xs.is_contiguous() {
216+
crate::ops::rms_norm(xs, &self.0.weight, self.0.eps as f32)
217+
} else {
218+
self.0.forward(xs)
219+
}
192220
}
193221
}
194222
}

0 commit comments

Comments
 (0)