Skip to main content

singe_cusparse/
spsv.rs

1use singe_cuda::{data_type::DataTypeLike, types::DevicePtr};
2
3use singe_cusparse_sys as sys;
4
5use crate::{
6    context::Context,
7    error::Result,
8    matrix::SparseMatrixDescriptor,
9    operation::SpSvDescriptor,
10    scalar::Scalar,
11    try_ffi,
12    types::{Operation, SpSvAlgorithm, SpSvUpdate},
13    utility::to_usize,
14    vector::DenseVectorDescriptor,
15};
16
17pub fn spsv_buffer_size<Compute: DataTypeLike>(
18    ctx: &Context,
19    operation: Operation,
20    alpha: Scalar<'_, Compute>,
21    matrix: &SparseMatrixDescriptor,
22    x: &DenseVectorDescriptor,
23    y: &mut DenseVectorDescriptor,
24    algorithm: SpSvAlgorithm,
25    descriptor: &SpSvDescriptor,
26) -> Result<usize> {
27    descriptor.ensure_context(ctx)?;
28    matrix.ensure_context(ctx)?;
29    x.ensure_context(ctx)?;
30    y.ensure_context(ctx)?;
31    ctx.bind()?;
32    if ctx.scalar_pointer_mode()? != alpha.pointer_mode() {
33        ctx.set_scalar_pointer_mode(alpha.pointer_mode())?;
34    }
35
36    let mut size = 0;
37    unsafe {
38        try_ffi!(sys::cusparseSpSV_bufferSize(
39            ctx.as_raw(),
40            operation.into(),
41            alpha.ptr().cast(),
42            matrix.as_raw_const(),
43            x.as_raw_const(),
44            y.as_raw(),
45            Compute::data_type().into(),
46            algorithm.into(),
47            descriptor.as_raw(),
48            &raw mut size,
49        ))?;
50    }
51    to_usize(size, "spsv buffer size")
52}
53
54pub fn spsv_analysis<Compute: DataTypeLike>(
55    ctx: &Context,
56    operation: Operation,
57    alpha: Scalar<'_, Compute>,
58    matrix: &SparseMatrixDescriptor,
59    x: &DenseVectorDescriptor,
60    y: &mut DenseVectorDescriptor,
61    algorithm: SpSvAlgorithm,
62    descriptor: &SpSvDescriptor,
63    external_buffer: Option<DevicePtr>,
64) -> Result<()> {
65    descriptor.ensure_context(ctx)?;
66    matrix.ensure_context(ctx)?;
67    x.ensure_context(ctx)?;
68    y.ensure_context(ctx)?;
69    ctx.bind()?;
70    if ctx.scalar_pointer_mode()? != alpha.pointer_mode() {
71        ctx.set_scalar_pointer_mode(alpha.pointer_mode())?;
72    }
73
74    unsafe {
75        try_ffi!(sys::cusparseSpSV_analysis(
76            ctx.as_raw(),
77            operation.into(),
78            alpha.ptr().cast(),
79            matrix.as_raw_const(),
80            x.as_raw_const(),
81            y.as_raw(),
82            Compute::data_type().into(),
83            algorithm.into(),
84            descriptor.as_raw(),
85            external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
86        ))?;
87    }
88    Ok(())
89}
90
91pub fn spsv_solve<Compute: DataTypeLike>(
92    ctx: &Context,
93    operation: Operation,
94    alpha: Scalar<'_, Compute>,
95    matrix: &SparseMatrixDescriptor,
96    x: &DenseVectorDescriptor,
97    y: &mut DenseVectorDescriptor,
98    algorithm: SpSvAlgorithm,
99    descriptor: &SpSvDescriptor,
100) -> Result<()> {
101    descriptor.ensure_context(ctx)?;
102    matrix.ensure_context(ctx)?;
103    x.ensure_context(ctx)?;
104    y.ensure_context(ctx)?;
105    ctx.bind()?;
106    if ctx.scalar_pointer_mode()? != alpha.pointer_mode() {
107        ctx.set_scalar_pointer_mode(alpha.pointer_mode())?;
108    }
109
110    unsafe {
111        try_ffi!(sys::cusparseSpSV_solve(
112            ctx.as_raw(),
113            operation.into(),
114            alpha.ptr().cast(),
115            matrix.as_raw_const(),
116            x.as_raw_const(),
117            y.as_raw(),
118            Compute::data_type().into(),
119            algorithm.into(),
120            descriptor.as_raw(),
121        ))?;
122    }
123    Ok(())
124}
125
126pub fn spsv_update_matrix(
127    ctx: &Context,
128    descriptor: &SpSvDescriptor,
129    new_values: DevicePtr,
130    update: SpSvUpdate,
131) -> Result<()> {
132    descriptor.ensure_context(ctx)?;
133    ctx.bind()?;
134    unsafe {
135        try_ffi!(sys::cusparseSpSV_updateMatrix(
136            ctx.as_raw(),
137            descriptor.as_raw(),
138            new_values.as_ptr() as _,
139            update.into(),
140        ))?;
141    }
142    Ok(())
143}