tenferro_tensor/tensor/constructors/
rng.rs

1#[cfg(feature = "cuda")]
2use std::sync::Arc;
3
4use tenferro_algebra::Scalar;
5use tenferro_device::{with_default_generator, Generator, LogicalMemorySpace, Result};
6
7use super::super::Tensor;
8#[cfg(feature = "cuda")]
9use super::super::TensorParts;
10use crate::MemoryOrder;
11
12#[cfg(feature = "cuda")]
13use crate::layout::compute_contiguous_strides;
14#[cfg(feature = "cuda")]
15use crate::DataBuffer;
16#[cfg(feature = "cuda")]
17use tenferro_device::cuda::runtime as device_cuda;
18
19fn with_generator<R, F>(
20    memory_space: LogicalMemorySpace,
21    generator: Option<&mut Generator>,
22    f: F,
23) -> Result<R>
24where
25    F: FnOnce(&mut Generator) -> Result<R>,
26{
27    match generator {
28        Some(generator) => f(generator),
29        None => with_default_generator(memory_space, f),
30    }
31}
32
33fn finish_generated_allocation<T: Scalar>(
34    data: Vec<T>,
35    dims: &[usize],
36    memory_space: LogicalMemorySpace,
37    order: MemoryOrder,
38) -> Result<Tensor<T>> {
39    Tensor::finish_allocation(
40        Tensor::main_memory_contiguous(data, dims, order),
41        memory_space,
42    )
43}
44
45#[cfg(feature = "cuda")]
46fn gpu_generated_tensor<T: Scalar>(
47    dims: &[usize],
48    memory_space: LogicalMemorySpace,
49    order: MemoryOrder,
50) -> Result<Tensor<T>> {
51    let LogicalMemorySpace::GpuMemory { device_id } = memory_space else {
52        return Err(Error::DeviceError(format!(
53            "expected CUDA memory space, got {memory_space:?}"
54        )));
55    };
56    let runtime = device_cuda::get_or_init(device_id)?;
57    let allocation = runtime.alloc::<T>(dims.iter().product())?;
58    let tensor = Tensor::from_parts(TensorParts {
59        buffer: unsafe {
60            DataBuffer::from_gpu_parts(
61                allocation.device_ptr(),
62                allocation.len(),
63                memory_space,
64                move || drop(allocation),
65            )
66        },
67        dims: Arc::from(dims),
68        strides: Arc::from(compute_contiguous_strides(dims, order)),
69        offset: 0,
70        logical_memory_space: memory_space,
71        preferred_compute_device: None,
72        event: None,
73        conjugated: false,
74        fw_grad: None,
75    });
76    Ok(tensor)
77}
78
79impl Tensor<f64> {
80    /// Create a tensor filled with uniform samples on `[0, 1)`.
81    ///
82    /// # Examples
83    ///
84    /// ```ignore
85    /// use tenferro_tensor::{MemoryOrder, Tensor};
86    ///
87    /// let t = Tensor::<f64>::rand(
88    ///     &[2, 2],
89    ///     tenferro_device::LogicalMemorySpace::MainMemory,
90    ///     MemoryOrder::ColumnMajor,
91    ///     None,
92    /// ).unwrap();
93    /// assert_eq!(t.dims(), &[2, 2]);
94    /// ```
95    pub fn rand(
96        dims: &[usize],
97        memory_space: LogicalMemorySpace,
98        order: MemoryOrder,
99        generator: Option<&mut Generator>,
100    ) -> Result<Self> {
101        with_generator(memory_space, generator, |generator| {
102            #[cfg(feature = "cuda")]
103            if matches!(memory_space, LogicalMemorySpace::GpuMemory { .. }) {
104                let tensor = gpu_generated_tensor::<f64>(dims, memory_space, order)?;
105                let LogicalMemorySpace::GpuMemory { device_id } = memory_space else {
106                    unreachable!("gpu_generated_tensor only accepts GPU memory");
107                };
108                let runtime = device_cuda::get_or_init(device_id)?;
109                let dst = tensor.buffer().as_device_ptr().ok_or_else(|| {
110                    Error::DeviceError("CUDA RNG destination is not on GPU".into())
111                })? as *mut f64;
112                let dst_len = tensor.buffer().len();
113                unsafe {
114                    runtime.rng_fill_uniform_f64_raw(
115                        generator,
116                        dst,
117                        dst_len,
118                        tensor.dims(),
119                        tensor.strides(),
120                        tensor.offset(),
121                    )?;
122                }
123                return Ok(tensor);
124            }
125            let n_elements: usize = dims.iter().product();
126            let mut data = Vec::with_capacity(n_elements);
127            for _ in 0..n_elements {
128                data.push(generator.sample_uniform_f64());
129            }
130            finish_generated_allocation(data, dims, memory_space, order)
131        })
132    }
133
134    /// Create a tensor filled with standard-normal samples.
135    ///
136    /// # Examples
137    ///
138    /// ```ignore
139    /// use tenferro_tensor::{MemoryOrder, Tensor};
140    ///
141    /// let t = Tensor::<f64>::randn(
142    ///     &[4],
143    ///     tenferro_device::LogicalMemorySpace::MainMemory,
144    ///     MemoryOrder::ColumnMajor,
145    ///     None,
146    /// ).unwrap();
147    /// assert_eq!(t.dims(), &[4]);
148    /// ```
149    pub fn randn(
150        dims: &[usize],
151        memory_space: LogicalMemorySpace,
152        order: MemoryOrder,
153        generator: Option<&mut Generator>,
154    ) -> Result<Self> {
155        with_generator(memory_space, generator, |generator| {
156            #[cfg(feature = "cuda")]
157            if matches!(memory_space, LogicalMemorySpace::GpuMemory { .. }) {
158                let tensor = gpu_generated_tensor::<f64>(dims, memory_space, order)?;
159                let LogicalMemorySpace::GpuMemory { device_id } = memory_space else {
160                    unreachable!("gpu_generated_tensor only accepts GPU memory");
161                };
162                let runtime = device_cuda::get_or_init(device_id)?;
163                let dst = tensor.buffer().as_device_ptr().ok_or_else(|| {
164                    Error::DeviceError("CUDA RNG destination is not on GPU".into())
165                })? as *mut f64;
166                let dst_len = tensor.buffer().len();
167                unsafe {
168                    runtime.rng_fill_normal_f64_raw(
169                        generator,
170                        dst,
171                        dst_len,
172                        tensor.dims(),
173                        tensor.strides(),
174                        tensor.offset(),
175                    )?;
176                }
177                return Ok(tensor);
178            }
179            let n_elements: usize = dims.iter().product();
180            let mut data = Vec::with_capacity(n_elements);
181            for _ in 0..n_elements {
182                data.push(generator.sample_standard_normal_f64());
183            }
184            finish_generated_allocation(data, dims, memory_space, order)
185        })
186    }
187
188    /// Create a tensor with the same shape/layout convention as another tensor and fill it with
189    /// uniform samples on `[0, 1)`.
190    ///
191    /// # Examples
192    ///
193    /// ```ignore
194    /// use tenferro_tensor::{MemoryOrder, Tensor};
195    ///
196    /// let base = Tensor::<f64>::zeros(&[2, 3], tenferro_device::LogicalMemorySpace::MainMemory, MemoryOrder::RowMajor).unwrap();
197    /// let t = Tensor::<f64>::rand_like(&base, None).unwrap();
198    /// assert_eq!(t.dims(), base.dims());
199    /// ```
200    pub fn rand_like(reference: &Self, generator: Option<&mut Generator>) -> Result<Self> {
201        Self::rand(
202            reference.dims(),
203            reference.logical_memory_space(),
204            Self::like_order(reference),
205            generator,
206        )
207    }
208
209    /// Create a tensor with the same shape/layout convention as another tensor and fill it with
210    /// standard-normal samples.
211    ///
212    /// # Examples
213    ///
214    /// ```ignore
215    /// use tenferro_tensor::{MemoryOrder, Tensor};
216    ///
217    /// let base = Tensor::<f64>::zeros(&[2, 3], tenferro_device::LogicalMemorySpace::MainMemory, MemoryOrder::RowMajor).unwrap();
218    /// let t = Tensor::<f64>::randn_like(&base, None).unwrap();
219    /// assert_eq!(t.dims(), base.dims());
220    /// ```
221    pub fn randn_like(reference: &Self, generator: Option<&mut Generator>) -> Result<Self> {
222        Self::randn(
223            reference.dims(),
224            reference.logical_memory_space(),
225            Self::like_order(reference),
226            generator,
227        )
228    }
229}
230
231impl Tensor<i32> {
232    /// Create a tensor filled with integer samples in `[low, high)`.
233    ///
234    /// # Examples
235    ///
236    /// ```ignore
237    /// use tenferro_tensor::{MemoryOrder, Tensor};
238    ///
239    /// let t = Tensor::<i32>::randint(
240    ///     -2,
241    ///     5,
242    ///     &[2, 2],
243    ///     tenferro_device::LogicalMemorySpace::MainMemory,
244    ///     MemoryOrder::ColumnMajor,
245    ///     None,
246    /// ).unwrap();
247    /// assert_eq!(t.dims(), &[2, 2]);
248    /// ```
249    pub fn randint(
250        low: i32,
251        high: i32,
252        dims: &[usize],
253        memory_space: LogicalMemorySpace,
254        order: MemoryOrder,
255        generator: Option<&mut Generator>,
256    ) -> Result<Self> {
257        with_generator(memory_space, generator, |generator| {
258            #[cfg(feature = "cuda")]
259            if matches!(memory_space, LogicalMemorySpace::GpuMemory { .. }) {
260                let tensor = gpu_generated_tensor::<i32>(dims, memory_space, order)?;
261                let LogicalMemorySpace::GpuMemory { device_id } = memory_space else {
262                    unreachable!("gpu_generated_tensor only accepts GPU memory");
263                };
264                let runtime = device_cuda::get_or_init(device_id)?;
265                let dst = tensor.buffer().as_device_ptr().ok_or_else(|| {
266                    Error::DeviceError("CUDA RNG destination is not on GPU".into())
267                })? as *mut i32;
268                let dst_len = tensor.buffer().len();
269                unsafe {
270                    runtime.rng_fill_i32_raw(
271                        generator,
272                        low,
273                        high,
274                        dst,
275                        dst_len,
276                        tensor.dims(),
277                        tensor.strides(),
278                        tensor.offset(),
279                    )?;
280                }
281                return Ok(tensor);
282            }
283            let n_elements: usize = dims.iter().product();
284            let mut data = Vec::with_capacity(n_elements);
285            for _ in 0..n_elements {
286                data.push(generator.sample_integer_i32(low, high)?);
287            }
288            finish_generated_allocation(data, dims, memory_space, order)
289        })
290    }
291
292    /// Create a tensor with the same shape/layout convention as another tensor and fill it with
293    /// integer samples in `[low, high)`.
294    ///
295    /// # Examples
296    ///
297    /// ```ignore
298    /// use tenferro_tensor::{MemoryOrder, Tensor};
299    ///
300    /// let base = Tensor::<i32>::zeros(&[2, 3], tenferro_device::LogicalMemorySpace::MainMemory, MemoryOrder::ColumnMajor).unwrap();
301    /// let t = Tensor::<i32>::randint_like(&base, -2, 5, None).unwrap();
302    /// assert_eq!(t.dims(), base.dims());
303    /// ```
304    pub fn randint_like(
305        reference: &Self,
306        low: i32,
307        high: i32,
308        generator: Option<&mut Generator>,
309    ) -> Result<Self> {
310        Self::randint(
311            low,
312            high,
313            reference.dims(),
314            reference.logical_memory_space(),
315            Self::like_order(reference),
316            generator,
317        )
318    }
319}
tenferro_tensor/tensor/constructors/rng.rs

tenferro_tensor/tensor/constructors/
rng.rs
