numr 0.5.1

High-performance numerical computing with multi-backend GPU acceleration (CPU/CUDA/WebGPU)
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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

//! CUDA implementation of multivariate random distribution operations.
//!
//! This module provides implementations that delegate to generic algorithms
//! while using CUDA kernels for multinomial sampling.

use crate::error::Result;
use crate::ops::impl_generic::{
    DTypeSupport, MultinomialSamplingOps, dirichlet_impl, multinomial_samples_impl,
    multivariate_normal_impl, wishart_impl,
};
use crate::ops::traits::multivariate::MultivariateRandomOps;
use crate::ops::{BinaryOps, CumulativeOps, RandomOps, ReduceOps};
use crate::runtime::cuda::kernels::launch_multinomial_count;
use crate::runtime::cuda::{CudaClient, CudaRuntime};
use crate::tensor::Tensor;

impl MultivariateRandomOps<CudaRuntime> for CudaClient {
    fn multivariate_normal(
        &self,
        mean: &Tensor<CudaRuntime>,
        cov: &Tensor<CudaRuntime>,
        n_samples: usize,
    ) -> Result<Tensor<CudaRuntime>> {
        multivariate_normal_impl(self, mean, cov, n_samples, DTypeSupport::FULL)
    }

    fn wishart(
        &self,
        scale: &Tensor<CudaRuntime>,
        df: usize,
        n_samples: usize,
    ) -> Result<Tensor<CudaRuntime>> {
        wishart_impl(self, scale, df, n_samples, DTypeSupport::FULL)
    }

    fn dirichlet(
        &self,
        alpha: &Tensor<CudaRuntime>,
        n_samples: usize,
    ) -> Result<Tensor<CudaRuntime>> {
        dirichlet_impl(self, alpha, n_samples)
    }

    fn multinomial_samples(
        &self,
        probs: &Tensor<CudaRuntime>,
        n_trials: usize,
        n_samples: usize,
    ) -> Result<Tensor<CudaRuntime>> {
        multinomial_samples_impl(self, probs, n_trials, n_samples)
    }
}

/// CUDA implementation of multinomial sampling kernel.
///
/// Uses CUDA kernels for CDF computation and categorical sampling.
impl MultinomialSamplingOps<CudaRuntime> for CudaClient {
    fn multinomial_sample_kernel(
        &self,
        probs: &Tensor<CudaRuntime>,
        n_trials: usize,
        n_samples: usize,
    ) -> Result<Tensor<CudaRuntime>> {
        let dtype = probs.dtype();
        let k = probs.shape()[0];

        // Step 1: Normalize probabilities (GPU tensor ops)
        let sum_probs = self.sum(probs, &[0], false)?;
        let normalized = self.div(probs, &sum_probs)?;

        // Step 2: Compute CDF using cumsum (GPU)
        let cdf = self.cumsum(&normalized, 0)?;

        // Step 3: Generate uniform samples [n_samples, n_trials] (GPU)
        let uniforms = self.rand(&[n_samples, n_trials], dtype)?;

        // Step 4: Allocate output and launch CUDA kernel
        let output = Tensor::<CudaRuntime>::zeros(&[n_samples, k], dtype, &self.device);

        // Get device pointers
        let cdf_ptr = cdf.ptr();
        let uniforms_ptr = uniforms.ptr();
        let output_ptr = output.ptr();

        // Launch kernel
        unsafe {
            launch_multinomial_count(
                &self.context,
                &self.stream,
                self.device.index,
                dtype,
                cdf_ptr,
                uniforms_ptr,
                output_ptr,
                k,
                n_trials,
                n_samples,
            )?;
        }

        Ok(output)
    }
}