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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104

//! [`DevicePayload`], one partner's exchange blocks with their columns in device memory.

use std::sync::Arc;

use cudarc::driver::{CudaSlice, CudaStream};

use super::columns::{grow, grow_keep};
use super::error::GpuError;
use crate::engine::partitioned::exchange_block::BlockHeader;

/// One exchange block with its columns on a device: the header and CSR offsets on the host, `x`/`z`/`coeff` on the device the sender ran on.
/// The columns are grow-only and `header.rows` says how much of them is live.
pub(crate) struct DeviceBlock<const W: usize> {
    pub(crate) header: BlockHeader,
    pub(crate) offsets: Vec<u32>,
    pub(crate) x: CudaSlice<u64>,
    pub(crate) z: CudaSlice<u64>,
    pub(crate) coefficient: CudaSlice<f64>,
}

impl<const W: usize> DeviceBlock<W> {
    fn new(stream: &Arc<CudaStream>) -> Result<Self, GpuError> {
        Ok(Self {
            header: BlockHeader { num_buckets: 0, rows: 0, width: W as u32, entry: 0 },
            offsets: Vec::new(),
            x: stream.alloc_zeros(W)?,
            z: stream.alloc_zeros(W)?,
            coefficient: stream.alloc_zeros(2)?,
        })
    }

    /// The block's header and offsets for `rows` rows over `buckets` positions; `offsets` must already hold the `buckets + 1` CSR offsets.
    pub(crate) fn set_header(&mut self, entry: u32, buckets: usize) {
        debug_assert_eq!(self.offsets.len(), buckets + 1);
        self.header = BlockHeader {
            num_buckets: buckets as u32,
            rows: self.offsets[buckets],
            width: W as u32,
            entry,
        };
    }

    /// Room for `rows` rows in every column, discarding what is there.
    pub(crate) fn grow(
        &mut self,
        stream: &Arc<CudaStream>,
        rows: usize,
        ordinal: u32,
    ) -> Result<(), GpuError> {
        grow(stream, &mut self.x, rows * W, ordinal)?;
        grow(stream, &mut self.z, rows * W, ordinal)?;
        grow(stream, &mut self.coefficient, 2 * rows, ordinal)
    }

    /// Room for `need` rows, keeping the first `live` rows of every column; grows geometrically.
    pub(crate) fn reserve_keep(
        &mut self,
        stream: &Arc<CudaStream>,
        live: usize,
        need: usize,
        ordinal: u32,
    ) -> Result<(), GpuError> {
        let capacity = self.coefficient.len() / 2;
        if need <= capacity {
            return Ok(());
        }
        let rows = need.max(2 * capacity);
        let bytes = (rows * (2 * W + 2) * 8) as u64;
        grow_keep(stream, &mut self.x, rows * W, live * W, ordinal, bytes)?;
        grow_keep(stream, &mut self.z, rows * W, live * W, ordinal, bytes)?;
        grow_keep(stream, &mut self.coefficient, 2 * rows, 2 * live, ordinal, bytes)
    }

    pub(crate) fn rows(&self) -> usize {
        self.header.rows as usize
    }

    /// Bytes the block occupies on a wire: header, offsets and the three live columns.
    pub(crate) fn bytes(&self) -> usize {
        std::mem::size_of::<BlockHeader>()
            + self.offsets.len() * std::mem::size_of::<u32>()
            + self.rows() * (2 * W + 2) * std::mem::size_of::<u64>()
    }
}

/// One partner's exchange blocks in ascending remote-delta index, pooled by the partition that fills them.
#[derive(Default)]
pub(crate) struct DevicePayload<const W: usize> {
    pub(crate) blocks: Vec<DeviceBlock<W>>,
}

impl<const W: usize> DevicePayload<W> {
    /// Block `index`, allocating on `stream` up to it.
    pub(crate) fn block_mut(
        &mut self,
        index: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<&mut DeviceBlock<W>, GpuError> {
        while self.blocks.len() <= index {
            self.blocks.push(DeviceBlock::new(stream)?);
        }
        Ok(&mut self.blocks[index])
    }
}