Skip to content

Commit 51be269

Browse files
farhan-syahfarhan-syah
authored
Merge pull request #2 from ml-rust/0.3.0
numr 0.3.0
2 parents 3568f63 + a6405e2 commit 51be269

51 files changed

Lines changed: 5245 additions & 1733 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎src/autograd/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@ mod forward;
121121
pub mod ops;
122122

123123
// Reverse-mode exports
124+
pub use crate::tensor::id::TensorId;
124125
pub use backward::{backward, backward_with_graph};
125126
pub use grad_fn::GradFn;
126127
pub use grad_store::GradStore;

‎src/lib.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ pub mod prelude {
9696
// Core types
9797
pub use crate::dtype::DType;
9898
pub use crate::error::{Error, Result};
99-
pub use crate::tensor::{Layout, Tensor};
99+
pub use crate::tensor::{Layout, Shape, Strides, Tensor};
100100

101101
// Runtime traits
102102
pub use crate::runtime::{Device, Runtime, RuntimeClient};
@@ -116,7 +116,7 @@ pub mod prelude {
116116

117117
// Backend runtimes
118118
#[cfg(feature = "cpu")]
119-
pub use crate::runtime::cpu::{CpuClient, CpuDevice, CpuRuntime};
119+
pub use crate::runtime::cpu::{CpuClient, CpuDevice, CpuRuntime, ParallelismConfig};
120120

121121
#[cfg(feature = "cuda")]
122122
pub use crate::runtime::cuda::{CudaClient, CudaDevice, CudaRuntime};

‎src/ops/cpu/complex.rs‎

Lines changed: 149 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -29,17 +29,30 @@ impl ComplexOps<CpuRuntime> for CpuClient {
2929

3030
let a_ptr = a_contig.storage().ptr();
3131
let out_ptr = out.storage().ptr();
32+
let chunk_size = self.chunk_size_hint();
3233

33-
unsafe {
34-
match dtype {
35-
DType::Complex64 => {
36-
kernels::conj_complex64(a_ptr as *const _, out_ptr as *mut _, numel);
37-
}
38-
DType::Complex128 => {
39-
kernels::conj_complex128(a_ptr as *const _, out_ptr as *mut _, numel);
40-
}
41-
_ => unreachable!("conj called on non-complex dtype"),
34+
match dtype {
35+
DType::Complex64 => {
36+
self.install_parallelism(|| unsafe {
37+
kernels::conj_complex64(
38+
a_ptr as *const _,
39+
out_ptr as *mut _,
40+
numel,
41+
chunk_size,
42+
);
43+
});
44+
}
45+
DType::Complex128 => {
46+
self.install_parallelism(|| unsafe {
47+
kernels::conj_complex128(
48+
a_ptr as *const _,
49+
out_ptr as *mut _,
50+
numel,
51+
chunk_size,
52+
);
53+
});
4254
}
55+
_ => unreachable!("conj called on non-complex dtype"),
4356
}
4457

4558
Ok(out)
@@ -71,17 +84,30 @@ impl ComplexOps<CpuRuntime> for CpuClient {
7184

7285
let a_ptr = a_contig.storage().ptr();
7386
let out_ptr = out.storage().ptr();
87+
let chunk_size = self.chunk_size_hint();
7488

75-
unsafe {
76-
match dtype {
77-
DType::Complex64 => {
78-
kernels::real_complex64(a_ptr as *const _, out_ptr as *mut _, numel);
79-
}
80-
DType::Complex128 => {
81-
kernels::real_complex128(a_ptr as *const _, out_ptr as *mut _, numel);
82-
}
83-
_ => unreachable!("real called on non-complex dtype"),
89+
match dtype {
90+
DType::Complex64 => {
91+
self.install_parallelism(|| unsafe {
92+
kernels::real_complex64(
93+
a_ptr as *const _,
94+
out_ptr as *mut _,
95+
numel,
96+
chunk_size,
97+
);
98+
});
8499
}
100+
DType::Complex128 => {
101+
self.install_parallelism(|| unsafe {
102+
kernels::real_complex128(
103+
a_ptr as *const _,
104+
out_ptr as *mut _,
105+
numel,
106+
chunk_size,
107+
);
108+
});
109+
}
110+
_ => unreachable!("real called on non-complex dtype"),
85111
}
86112

87113
Ok(out)
@@ -113,17 +139,30 @@ impl ComplexOps<CpuRuntime> for CpuClient {
113139

114140
let a_ptr = a_contig.storage().ptr();
115141
let out_ptr = out.storage().ptr();
142+
let chunk_size = self.chunk_size_hint();
116143

117-
unsafe {
118-
match dtype {
119-
DType::Complex64 => {
120-
kernels::imag_complex64(a_ptr as *const _, out_ptr as *mut _, numel);
121-
}
122-
DType::Complex128 => {
123-
kernels::imag_complex128(a_ptr as *const _, out_ptr as *mut _, numel);
124-
}
125-
_ => unreachable!("imag called on non-complex dtype"),
144+
match dtype {
145+
DType::Complex64 => {
146+
self.install_parallelism(|| unsafe {
147+
kernels::imag_complex64(
148+
a_ptr as *const _,
149+
out_ptr as *mut _,
150+
numel,
151+
chunk_size,
152+
);
153+
});
126154
}
155+
DType::Complex128 => {
156+
self.install_parallelism(|| unsafe {
157+
kernels::imag_complex128(
158+
a_ptr as *const _,
159+
out_ptr as *mut _,
160+
numel,
161+
chunk_size,
162+
);
163+
});
164+
}
165+
_ => unreachable!("imag called on non-complex dtype"),
127166
}
128167

129168
Ok(out)
@@ -133,6 +172,7 @@ impl ComplexOps<CpuRuntime> for CpuClient {
133172
let dtype = a.dtype();
134173
let shape = a.shape();
135174
let numel = a.numel();
175+
let chunk_size = self.chunk_size_hint();
136176

137177
let a_contig = ensure_contiguous(a);
138178

@@ -148,19 +188,31 @@ impl ComplexOps<CpuRuntime> for CpuClient {
148188
let a_ptr = a_contig.storage().ptr();
149189
let out_ptr = out.storage().ptr();
150190

151-
unsafe {
152-
match dtype {
153-
DType::F32 => {
154-
kernels::angle_real_f32(a_ptr as *const _, out_ptr as *mut _, numel);
155-
}
156-
DType::F64 => {
157-
kernels::angle_real_f64(a_ptr as *const _, out_ptr as *mut _, numel);
158-
}
159-
_ => {
160-
// For integer types, angle doesn't make mathematical sense
161-
// Return zeros
162-
return Ok(Tensor::<CpuRuntime>::zeros(shape, dtype, &self.device));
163-
}
191+
match dtype {
192+
DType::F32 => {
193+
self.install_parallelism(|| unsafe {
194+
kernels::angle_real_f32(
195+
a_ptr as *const _,
196+
out_ptr as *mut _,
197+
numel,
198+
chunk_size,
199+
);
200+
});
201+
}
202+
DType::F64 => {
203+
self.install_parallelism(|| unsafe {
204+
kernels::angle_real_f64(
205+
a_ptr as *const _,
206+
out_ptr as *mut _,
207+
numel,
208+
chunk_size,
209+
);
210+
});
211+
}
212+
_ => {
213+
// For integer types, angle doesn't make mathematical sense
214+
// Return zeros
215+
return Ok(Tensor::<CpuRuntime>::zeros(shape, dtype, &self.device));
164216
}
165217
}
166218
return Ok(out);
@@ -181,16 +233,28 @@ impl ComplexOps<CpuRuntime> for CpuClient {
181233
let a_ptr = a_contig.storage().ptr();
182234
let out_ptr = out.storage().ptr();
183235

184-
unsafe {
185-
match dtype {
186-
DType::Complex64 => {
187-
kernels::angle_complex64(a_ptr as *const _, out_ptr as *mut _, numel);
188-
}
189-
DType::Complex128 => {
190-
kernels::angle_complex128(a_ptr as *const _, out_ptr as *mut _, numel);
191-
}
192-
_ => unreachable!("angle called on non-complex dtype"),
236+
match dtype {
237+
DType::Complex64 => {
238+
self.install_parallelism(|| unsafe {
239+
kernels::angle_complex64(
240+
a_ptr as *const _,
241+
out_ptr as *mut _,
242+
numel,
243+
chunk_size,
244+
);
245+
});
193246
}
247+
DType::Complex128 => {
248+
self.install_parallelism(|| unsafe {
249+
kernels::angle_complex128(
250+
a_ptr as *const _,
251+
out_ptr as *mut _,
252+
numel,
253+
chunk_size,
254+
);
255+
});
256+
}
257+
_ => unreachable!("angle called on non-complex dtype"),
194258
}
195259

196260
Ok(out)
@@ -226,27 +290,32 @@ impl ComplexOps<CpuRuntime> for CpuClient {
226290
let real_ptr = real_contig.storage().ptr();
227291
let imag_ptr = imag_contig.storage().ptr();
228292
let out_ptr = out.storage().ptr();
293+
let chunk_size = self.chunk_size_hint();
229294

230-
unsafe {
231-
match input_dtype {
232-
DType::F32 => {
295+
match input_dtype {
296+
DType::F32 => {
297+
self.install_parallelism(|| unsafe {
233298
kernels::from_real_imag_f32(
234299
real_ptr as *const _,
235300
imag_ptr as *const _,
236301
out_ptr as *mut _,
237302
numel,
303+
chunk_size,
238304
);
239-
}
240-
DType::F64 => {
305+
});
306+
}
307+
DType::F64 => {
308+
self.install_parallelism(|| unsafe {
241309
kernels::from_real_imag_f64(
242310
real_ptr as *const _,
243311
imag_ptr as *const _,
244312
out_ptr as *mut _,
245313
numel,
314+
chunk_size,
246315
);
247-
}
248-
_ => unreachable!("validated above"),
316+
});
249317
}
318+
_ => unreachable!("validated above"),
250319
}
251320

252321
Ok(out)
@@ -275,27 +344,32 @@ impl ComplexOps<CpuRuntime> for CpuClient {
275344
let complex_ptr = complex_contig.storage().ptr();
276345
let real_ptr = real_contig.storage().ptr();
277346
let out_ptr = out.storage().ptr();
347+
let chunk_size = self.chunk_size_hint();
278348

279-
unsafe {
280-
match dtype {
281-
DType::Complex64 => {
349+
match dtype {
350+
DType::Complex64 => {
351+
self.install_parallelism(|| unsafe {
282352
kernels::complex64_mul_real(
283353
complex_ptr as *const _,
284354
real_ptr as *const _,
285355
out_ptr as *mut _,
286356
numel,
357+
chunk_size,
287358
);
288-
}
289-
DType::Complex128 => {
359+
});
360+
}
361+
DType::Complex128 => {
362+
self.install_parallelism(|| unsafe {
290363
kernels::complex128_mul_real(
291364
complex_ptr as *const _,
292365
real_ptr as *const _,
293366
out_ptr as *mut _,
294367
numel,
368+
chunk_size,
295369
);
296-
}
297-
_ => unreachable!("validated above"),
370+
});
298371
}
372+
_ => unreachable!("validated above"),
299373
}
300374

301375
Ok(out)
@@ -324,27 +398,32 @@ impl ComplexOps<CpuRuntime> for CpuClient {
324398
let complex_ptr = complex_contig.storage().ptr();
325399
let real_ptr = real_contig.storage().ptr();
326400
let out_ptr = out.storage().ptr();
401+
let chunk_size = self.chunk_size_hint();
327402

328-
unsafe {
329-
match dtype {
330-
DType::Complex64 => {
403+
match dtype {
404+
DType::Complex64 => {
405+
self.install_parallelism(|| unsafe {
331406
kernels::complex64_div_real(
332407
complex_ptr as *const _,
333408
real_ptr as *const _,
334409
out_ptr as *mut _,
335410
numel,
411+
chunk_size,
336412
);
337-
}
338-
DType::Complex128 => {
413+
});
414+
}
415+
DType::Complex128 => {
416+
self.install_parallelism(|| unsafe {
339417
kernels::complex128_div_real(
340418
complex_ptr as *const _,
341419
real_ptr as *const _,
342420
out_ptr as *mut _,
343421
numel,
422+
chunk_size,
344423
);
345-
}
346-
_ => unreachable!("validated above"),
424+
});
347425
}
426+
_ => unreachable!("validated above"),
348427
}
349428

350429
Ok(out)

0 commit comments

Comments
 (0)