Skip to content

Commit 0bce325

Browse files
EricLBuehlerEricLBuehler
committed
Fix cpu
1 parent 3276c80 commit 0bce325

2 files changed

Lines changed: 124 additions & 30 deletions

File tree

‎candle-core/src/cpu/avx.rs‎

Lines changed: 81 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
1-
use super::{Cpu, CpuF16};
1+
use super::{Cpu, CpuBF16, CpuF16};
22
#[cfg(target_arch = "x86")]
33
use core::arch::x86::*;
44
#[cfg(target_arch = "x86_64")]
55
use core::arch::x86_64::*;
66

7-
use half::f16;
7+
use half::{bf16, f16};
88

99
pub struct CurrentCpu {}
1010

@@ -146,3 +146,82 @@ impl CpuF16<ARR> for CurrentCpuF16 {
146146
*y = _mm_cvtss_f32(_mm_hadd_ps(t1, t1));
147147
}
148148
}
149+
150+
pub struct CurrentCpuBF16 {}
151+
impl CpuBF16<ARR> for CurrentCpuBF16 {
152+
type Unit = __m256;
153+
type Array = [__m256; ARR];
154+
155+
const STEP: usize = STEP;
156+
const EPR: usize = EPR;
157+
158+
fn n() -> usize {
159+
ARR
160+
}
161+
162+
unsafe fn zero() -> Self::Unit {
163+
_mm256_setzero_ps()
164+
}
165+
166+
unsafe fn zero_array() -> Self::Array {
167+
[Self::zero(); ARR]
168+
}
169+
170+
unsafe fn from_f32(v: f32) -> Self::Unit {
171+
_mm256_set1_ps(v)
172+
}
173+
174+
#[cfg(target_feature = "f16c")]
175+
unsafe fn load(mem_addr: *const bf16) -> Self::Unit {
176+
_mm256_cvtph_ps(_mm_loadu_si128(mem_addr as *const __m128i))
177+
}
178+
179+
#[cfg(not(target_feature = "f16c"))]
180+
unsafe fn load(mem_addr: *const bf16) -> Self::Unit {
181+
let mut tmp = [0.0f32; 8];
182+
for i in 0..8 {
183+
tmp[i] = (*mem_addr.add(i)).to_f32();
184+
}
185+
_mm256_loadu_ps(tmp.as_ptr())
186+
}
187+
188+
unsafe fn vec_add(a: Self::Unit, b: Self::Unit) -> Self::Unit {
189+
_mm256_add_ps(a, b)
190+
}
191+
192+
unsafe fn vec_fma(a: Self::Unit, b: Self::Unit, c: Self::Unit) -> Self::Unit {
193+
_mm256_add_ps(_mm256_mul_ps(b, c), a)
194+
}
195+
196+
#[cfg(target_feature = "f16c")]
197+
unsafe fn vec_store(mem_addr: *mut bf16, a: Self::Unit) {
198+
_mm_storeu_si128(mem_addr as *mut __m128i, _mm256_cvtps_ph(a, 0))
199+
}
200+
201+
#[cfg(not(target_feature = "f16c"))]
202+
unsafe fn vec_store(mem_addr: *mut bf16, a: Self::Unit) {
203+
let mut tmp = [0.0f32; 8];
204+
_mm256_storeu_ps(tmp.as_mut_ptr(), a);
205+
for i in 0..8 {
206+
*mem_addr.add(i) = bf16::from_f32(tmp[i]);
207+
}
208+
}
209+
210+
unsafe fn vec_reduce(mut x: Self::Array, y: *mut f32) {
211+
let mut offset = ARR >> 1;
212+
for i in 0..offset {
213+
x[i] = _mm256_add_ps(x[i], x[offset + i]);
214+
}
215+
offset >>= 1;
216+
for i in 0..offset {
217+
x[i] = _mm256_add_ps(x[i], x[offset + i]);
218+
}
219+
offset >>= 1;
220+
for i in 0..offset {
221+
x[i] = _mm256_add_ps(x[i], x[offset + i]);
222+
}
223+
let t0 = _mm_add_ps(_mm256_castps256_ps128(x[0]), _mm256_extractf128_ps(x[0], 1));
224+
let t1 = _mm_hadd_ps(t0, t0);
225+
*y = _mm_cvtss_f32(_mm_hadd_ps(t1, t1));
226+
}
227+
}

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

Lines changed: 43 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
//! Traits and methods for CPU-backed Tensors
2-
31
pub mod erf;
42
pub mod kernels;
53

@@ -38,16 +36,33 @@ trait CpuF16<const ARR: usize> {
3836
unsafe fn from_f32(v: f32) -> Self::Unit;
3937
unsafe fn vec_store(mem_addr: *mut f16, a: Self::Unit);
4038
}
41-
#[cfg(not(target_feature = "avx"))]
42-
use half::bf16;
43-
use half::f16;
39+
40+
#[allow(unused)]
41+
trait CpuBF16<const ARR: usize> {
42+
type Unit;
43+
type Array;
44+
const STEP: usize;
45+
const EPR: usize;
46+
47+
fn n() -> usize;
48+
unsafe fn zero() -> Self::Unit;
49+
unsafe fn zero_array() -> Self::Array;
50+
unsafe fn load(mem_addr: *const bf16) -> Self::Unit;
51+
unsafe fn vec_add(a: Self::Unit, b: Self::Unit) -> Self::Unit;
52+
unsafe fn vec_fma(a: Self::Unit, b: Self::Unit, c: Self::Unit) -> Self::Unit;
53+
unsafe fn vec_reduce(x: Self::Array, y: *mut f32);
54+
unsafe fn from_f32(v: f32) -> Self::Unit;
55+
unsafe fn vec_store(mem_addr: *mut bf16, a: Self::Unit);
56+
}
57+
58+
use half::{bf16, f16};
4459

4560
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
4661
#[cfg(target_feature = "avx")]
4762
pub mod avx;
4863
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
4964
#[cfg(target_feature = "avx")]
50-
pub use avx::{CurrentCpu, CurrentCpuF16};
65+
pub use avx::{CurrentCpu, CurrentCpuBF16, CurrentCpuF16};
5166

5267
#[cfg(target_arch = "wasm32")]
5368
#[cfg(target_feature = "simd128")]
@@ -148,24 +163,24 @@ pub(crate) unsafe fn vec_sum(row: *const f32, b: *mut f32, k: usize) {
148163

149164
#[cfg(target_feature = "avx")]
150165
#[inline(always)]
151-
pub(crate) unsafe fn vec_dot_bf16(a_row: *const bf16, b_row: *const bf16, c: *mut f32, k: usize) {
166+
pub(crate) unsafe fn vec_dot_f16(a_row: *const f16, b_row: *const f16, c: *mut f32, k: usize) {
152167
let mut sumf = 0.0f32;
153-
let np = k & !(CurrentCpuBF16::STEP - 1);
168+
let np = k & !(CurrentCpuF16::STEP - 1);
154169

155-
let mut sum = CurrentCpuBF16::zero_array();
156-
let mut ax = CurrentCpuBF16::zero_array();
157-
let mut ay = CurrentCpuBF16::zero_array();
170+
let mut sum = CurrentCpuF16::zero_array();
171+
let mut ax = CurrentCpuF16::zero_array();
172+
let mut ay = CurrentCpuF16::zero_array();
158173

159-
for i in (0..np).step_by(CurrentCpuBF16::STEP) {
160-
for j in 0..CurrentCpuBF16::n() {
161-
ax[j] = CurrentCpuBF16::load(a_row.add(i + j * CurrentCpuBF16::EPR));
162-
ay[j] = CurrentCpuBF16::load(b_row.add(i + j * CurrentCpuBF16::EPR));
174+
for i in (0..np).step_by(CurrentCpuF16::STEP) {
175+
for j in 0..CurrentCpuF16::n() {
176+
ax[j] = CurrentCpuF16::load(a_row.add(i + j * CurrentCpuF16::EPR));
177+
ay[j] = CurrentCpuF16::load(b_row.add(i + j * CurrentCpuF16::EPR));
163178

164-
sum[j] = CurrentCpuBF16::vec_fma(sum[j], ax[j], ay[j]);
179+
sum[j] = CurrentCpuF16::vec_fma(sum[j], ax[j], ay[j]);
165180
}
166181
}
167182

168-
CurrentCpuBF16::vec_reduce(sum, &mut sumf);
183+
CurrentCpuF16::vec_reduce(sum, &mut sumf);
169184

170185
// leftovers
171186
for i in np..k {
@@ -176,24 +191,24 @@ pub(crate) unsafe fn vec_dot_bf16(a_row: *const bf16, b_row: *const bf16, c: *mu
176191

177192
#[cfg(target_feature = "avx")]
178193
#[inline(always)]
179-
pub(crate) unsafe fn vec_dot_f16(a_row: *const f16, b_row: *const f16, c: *mut f32, k: usize) {
194+
pub(crate) unsafe fn vec_dot_bf16(a_row: *const bf16, b_row: *const bf16, c: *mut f32, k: usize) {
180195
let mut sumf = 0.0f32;
181-
let np = k & !(CurrentCpuF16::STEP - 1);
196+
let np = k & !(CurrentCpuBF16::STEP - 1);
182197

183-
let mut sum = CurrentCpuF16::zero_array();
184-
let mut ax = CurrentCpuF16::zero_array();
185-
let mut ay = CurrentCpuF16::zero_array();
198+
let mut sum = CurrentCpuBF16::zero_array();
199+
let mut ax = CurrentCpuBF16::zero_array();
200+
let mut ay = CurrentCpuBF16::zero_array();
186201

187-
for i in (0..np).step_by(CurrentCpuF16::STEP) {
188-
for j in 0..CurrentCpuF16::n() {
189-
ax[j] = CurrentCpuF16::load(a_row.add(i + j * CurrentCpuF16::EPR));
190-
ay[j] = CurrentCpuF16::load(b_row.add(i + j * CurrentCpuF16::EPR));
202+
for i in (0..np).step_by(CurrentCpuBF16::STEP) {
203+
for j in 0..CurrentCpuBF16::n() {
204+
ax[j] = CurrentCpuBF16::load(a_row.add(i + j * CurrentCpuBF16::EPR));
205+
ay[j] = CurrentCpuBF16::load(b_row.add(i + j * CurrentCpuBF16::EPR));
191206

192-
sum[j] = CurrentCpuF16::vec_fma(sum[j], ax[j], ay[j]);
207+
sum[j] = CurrentCpuBF16::vec_fma(sum[j], ax[j], ay[j]);
193208
}
194209
}
195210

196-
CurrentCpuF16::vec_reduce(sum, &mut sumf);
211+
CurrentCpuBF16::vec_reduce(sum, &mut sumf);
197212

198213
// leftovers
199214
for i in np..k {

0 commit comments

Comments
 (0)