hpt 0.1.3

High Performance Tensor (HPT) - A fast, efficient, and user-friendly tensor computation library for Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74

use crate::backends::common::reduce::rearrange_array;
use crate::backends::cuda::cuda_utils::load_ptx_and_get_data;
use crate::{backend::Cuda, tensor_base::_Tensor};
use cudarc::driver::DeviceRepr;
use cudarc::driver::LaunchAsync;
use hpt_allocator::traits::{Allocator, AllocatorOutputRetrive};
use hpt_common::error::base::TensorError;
use hpt_common::error::shape::ShapeError;
use hpt_cudakernels::SET_VAL;
use hpt_traits::ops::creation::TensorCreator;
use hpt_traits::ops::shape_manipulate::ShapeManipulate;
use hpt_traits::tensor::CommonBounds;
use hpt_traits::tensor::TensorInfo;
use hpt_types::dtype::CudaType;

use crate::backends::cuda::cuda_utils::compute_kernel_launch_config;

pub(crate) fn reduce_prepare<
    T: CommonBounds + DeviceRepr + CudaType,
    O: CommonBounds + DeviceRepr + CudaType,
    const DEVICE_ID: usize,
    Al,
>(
    a: &_Tensor<T, Cuda, DEVICE_ID, Al>,
    axes: &[usize],
    init_val: O,
    init_out: bool,
    c: Option<_Tensor<O, Cuda, DEVICE_ID, Al>>,
) -> std::result::Result<
    (
        _Tensor<T, Cuda, DEVICE_ID, Al>,
        _Tensor<O, Cuda, DEVICE_ID, Al>,
    ),
    TensorError,
>
where
    Al: Allocator,
    Al::Output: AllocatorOutputRetrive,
{
    // get permute order, we move to_reduce axes to the end
    let mut transposed_axis = rearrange_array(a.layout.ndim(), axes);
    // sort the transposed axis based on the stride, ordering the axis can increase the cpu cache hitting rate when we do iteration
    transposed_axis[..a.layout.ndim() - axes.len()]
        .sort_by(|x, y| a.layout.strides()[*y].cmp(&a.layout.strides()[*x]));
    transposed_axis[a.layout.ndim() - axes.len()..]
        .sort_by(|x, y| a.layout.strides()[*y].cmp(&a.layout.strides()[*x]));
    let res_layout = a.layout.reduce(axes, false)?;
    let res = if let Some(out) = c {
        ShapeError::check_inplace_out_layout_valid(res_layout.shape(), &out.layout())?;
        if init_out {
            let (kernel, reg_info) = load_ptx_and_get_data(
                "set_val",
                &format!("set_val_{}", O::STR),
                a.device(),
                a.device_cap(),
                &SET_VAL,
            )
            .unwrap();
            let cfg =
                compute_kernel_launch_config(out.device(), &reg_info, res_layout.size() as usize);
            unsafe {
                kernel.launch(
                    cfg,
                    (out.cuda_slice(), init_val, res_layout.size() as usize),
                )?
            };
        }
        out.reshape(res_layout.shape())?
    } else {
        let res = _Tensor::<O, Cuda, DEVICE_ID, Al>::empty(res_layout.shape())?;
        res
    };
    Ok((a.permute(transposed_axis)?, res))
}