paulistrings 0.1.2

Classical simulation of quantum circuits in the Pauli basis, parallel across threads, NUMA domains, MPI ranks and GPUs
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

//! Device-wide exclusive scan of a `u32` column, the three launches of `kernels/scan.cu`.

use std::sync::Arc;

use cudarc::driver::{CudaSlice, CudaStream, CudaView, CudaViewMut, LaunchConfig, PushKernelArg};

use super::columns::grow;
use super::error::GpuError;
use super::module::KernelSet;

/// Elements per scan block; must match `SCAN_BLOCK` in `kernels/scan.cu`.
const SCAN_BLOCK: usize = 4096;
const SCAN_THREADS: u32 = 1024;

/// Grow-only per-block buffers a scan needs, kept between layers so a scan allocates nothing.
pub(crate) struct ScanScratch {
    block_sum: CudaSlice<u32>,
    block_max: CudaSlice<u32>,
}

impl ScanScratch {
    pub(crate) fn new(stream: &Arc<CudaStream>) -> Result<Self, GpuError> {
        Ok(Self { block_sum: stream.alloc_zeros(1)?, block_max: stream.alloc_zeros(1)? })
    }
}

/// `out[i] = Σ_{j<i} input[j]` for `i ≤ n`, so `out[n]` is the total, with `total_max` (two elements) receiving `[total, max]` of `input[0..n]`; enqueued on `stream`, not synchronized.
pub(crate) fn exclusive_scan(
    stream: &Arc<CudaStream>,
    kernels: &KernelSet,
    input: &CudaView<'_, u32>,
    out: &mut CudaViewMut<'_, u32>,
    n: usize,
    scratch: &mut ScanScratch,
    total_max: &mut CudaSlice<u32>,
) -> Result<(), GpuError> {
    let nb = n.div_ceil(SCAN_BLOCK).max(1);
    if nb > SCAN_BLOCK {
        return Err(GpuError::Unsupported("scan of more than 2^24 elements"));
    }
    debug_assert!(input.len() >= n && out.len() > n && total_max.len() >= 2);
    let ordinal = stream.context().ordinal() as u32;
    // Every block writes its own slot before `scan_single` reads it, so no zero fill is needed.
    grow(stream, &mut scratch.block_sum, nb, ordinal)?;
    grow(stream, &mut scratch.block_max, nb, ordinal)?;
    let block_sum = &mut scratch.block_sum;
    let block_max = &mut scratch.block_max;
    let (n32, nb32) = (n as u32, nb as u32);
    let block = (SCAN_THREADS, 1, 1);
    // SAFETY: argument lists match the `extern "C"` signatures in scan.cu, and every buffer holds the elements indexed.
    unsafe {
        stream
            .launch_builder(&kernels.scan_block)
            .arg(input)
            .arg(&mut *out)
            .arg(&mut *block_sum)
            .arg(&mut *block_max)
            .arg(&n32)
            .launch(LaunchConfig {
                grid_dim: (nb32, 1, 1),
                block_dim: block,
                shared_mem_bytes: 0,
            })?;
        stream
            .launch_builder(&kernels.scan_single)
            .arg(&mut *block_sum)
            .arg(&*block_max)
            .arg(&nb32)
            .arg(&mut *out)
            .arg(&n32)
            .arg(&mut *total_max)
            .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: block, shared_mem_bytes: 0 })?;
        stream.launch_builder(&kernels.scan_add).arg(&mut *out).arg(&*block_sum).arg(&n32).launch(
            LaunchConfig {
                grid_dim: ((n as u32).div_ceil(SCAN_THREADS).max(1), 1, 1),
                block_dim: block,
                shared_mem_bytes: 0,
            },
        )?;
    }
    Ok(())
}