hpt 0.1.3

High Performance Tensor (HPT) - A fast, efficient, and user-friendly tensor computation library for Rust
Documentation
use crate::backends::cuda::{
    cuda_utils::get_module_name_1, utils::unary::unary::uary_fn_with_out_simd,
};
use crate::{backend::Cuda, tensor_base::_Tensor};
use cudarc::driver::DeviceRepr;
use hpt_allocator::traits::{Allocator, AllocatorOutputRetrive};
use hpt_common::error::base::TensorError;
use hpt_traits::ops::creation::TensorCreator;
use hpt_traits::ops::windows::WindowOps;
use hpt_traits::tensor::CommonBounds;
use hpt_types::cuda_types::scalar::Scalar;
use hpt_types::dtype::CudaType;
use hpt_types::{
    dtype::FloatConst,
    into_scalar::Cast,
    type_promote::{FloatOutBinary, FloatOutUnary},
};

impl<T, const DEVICE_ID: usize, Al> WindowOps for _Tensor<T, Cuda, DEVICE_ID, Al>
where
    f64: Cast<T>,
    T: CommonBounds
        + FloatOutBinary<Output = T>
        + FloatOutUnary<Output = T>
        + FloatConst
        + DeviceRepr
        + CudaType,
    usize: Cast<T>,
    i64: Cast<T>,
    Al: Allocator,
    Al::Output: AllocatorOutputRetrive,
{
    type Output = _Tensor<T, Cuda, DEVICE_ID, Al>;
    type Meta = T;
    #[track_caller]
    fn hamming_window(
        window_length: i64,
        periodic: bool,
    ) -> std::result::Result<Self::Output, TensorError> {
        __hamming_window::<T, DEVICE_ID, Al>(window_length, (0.54).cast(), (0.46).cast(), periodic)
    }

    #[track_caller]
    fn hann_window(
        window_length: i64,
        periodic: bool,
    ) -> std::result::Result<Self::Output, TensorError> {
        __hamming_window::<T, DEVICE_ID, Al>(window_length, (0.5).cast(), (0.5).cast(), periodic)
    }

    #[track_caller]
    fn blackman_window(
        window_length: i64,
        periodic: bool,
    ) -> std::result::Result<Self::Output, TensorError> {
        let length_usize = if periodic {
            window_length
        } else {
            window_length - 1
        };
        let length: T = length_usize.cast();
        let ret = _Tensor::<T, Cuda, DEVICE_ID, Al>::empty(&[length_usize])?;
        uary_fn_with_out_simd(
            &ret,
            &get_module_name_1("blackman_window", &ret),
            |out, idx| {
                let res = match T::STR {
                    "f32" => {
                        format!(
                            "
                            float n = (float){idx};
                            float w1 = 2.0f * M_PI * n / {length}f;
                            float w2 = 2.0f * w1;  // 4π * n / (N-1)
                            {out} = 0.42f - 0.5f * cosf(w1) + 0.08f * cosf(w2);"
                        )
                    }
                    "f64" => {
                        format!(
                            "
                            double n = (double){idx};
                            double w1 = 2.0 * M_PI * n / {length};
                            double w2 = 2.0 * w1;
                            {out} = 0.42 - 0.5 * cos(w1) + 0.08 * cos(w2);"
                        )
                    }
                    "f16" => {
                        format!(
                            "
                            float n = (float){idx};
                            float w1 = 2.0f * M_PI * n / {length}f;
                            float w2 = 2.0f * w1;
                            {out} = __float2half(0.42f - 0.5f * cosf(w1) + 0.08f * cosf(w2));"
                        )
                    }
                    _ => unreachable!(),
                };
                Scalar::<T>::new(res)
            },
            None::<Self::Output>,
        )
    }
}

#[track_caller]
fn __hamming_window<T, const DEVICE_ID: usize, Al>(
    window_length: i64,
    alpha: T,
    beta: T,
    periodic: bool,
) -> std::result::Result<_Tensor<T, Cuda, DEVICE_ID, Al>, TensorError>
where
    T: CommonBounds + DeviceRepr + CudaType,
    usize: Cast<T>,
    Al: Allocator,
    Al::Output: AllocatorOutputRetrive,
{
    let length_usize = (if periodic {
        window_length
    } else {
        window_length - 1
    }) as usize;
    let length: T = length_usize.cast();
    let ret = _Tensor::<T, Cuda, DEVICE_ID, Al>::empty(&[length_usize as i64])?;
    uary_fn_with_out_simd(
        &ret,
        &get_module_name_1("hamming_window", &ret),
        |out, idx| {
            let res = match T::STR {
                "f32" => {
                    format!(
                        "
                    float n = (float){idx};
                    {out} = {alpha} - {beta} * cosf(2.0f * M_PI * n / {length});"
                    )
                }
                "f64" => {
                    format!(
                        "
                    double n = (double){idx};
                    {out} = {alpha} - {beta} * cos(2.0 * M_PI * n / {length});"
                    )
                }
                "f16" => {
                    format!(
                        "
                    float n = (float){idx};
                    {out} = __float2half({alpha}f - {beta}f * cosf(2.0f * M_PI * n / {length}));"
                    )
                }
                _ => unreachable!(),
            };
            Scalar::<T>::new(res)
        },
        None::<_Tensor<T, Cuda, DEVICE_ID, Al>>,
    )
}