singe_cusparse/
spsm.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::{DenseMatrixDescriptor, SparseMatrixDescriptor},
9    operation::SpSmDescriptor,
10    scalar::Scalar,
11    try_ffi,
12    types::{Operation, SpSMAlgorithm, SpSmUpdate},
13    utility::to_usize,
14};
15
16pub fn spsm_buffer_size<Compute: DataTypeLike>(
17    ctx: &Context,
18    op_a: Operation,
19    op_b: Operation,
20    alpha: Scalar<'_, Compute>,
21    matrix_a: &SparseMatrixDescriptor,
22    matrix_b: &DenseMatrixDescriptor,
23    matrix_c: &mut DenseMatrixDescriptor,
24    algorithm: SpSMAlgorithm,
25    descriptor: &SpSmDescriptor,
26) -> Result<usize> {
27    descriptor.ensure_context(ctx)?;
28    matrix_a.ensure_context(ctx)?;
29    matrix_b.ensure_context(ctx)?;
30    matrix_c.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::cusparseSpSM_bufferSize(
39            ctx.as_raw(),
40            op_a.into(),
41            op_b.into(),
42            alpha.ptr().cast(),
43            matrix_a.as_raw_const(),
44            matrix_b.as_raw_const(),
45            matrix_c.as_raw(),
46            Compute::data_type().into(),
47            algorithm.into(),
48            descriptor.as_raw(),
49            &raw mut size,
50        ))?;
51    }
52    to_usize(size, "spsm buffer size")
53}
54
55pub fn spsm_analysis<Compute: DataTypeLike>(
56    ctx: &Context,
57    op_a: Operation,
58    op_b: Operation,
59    alpha: Scalar<'_, Compute>,
60    matrix_a: &SparseMatrixDescriptor,
61    matrix_b: &DenseMatrixDescriptor,
62    matrix_c: &mut DenseMatrixDescriptor,
63    algorithm: SpSMAlgorithm,
64    descriptor: &SpSmDescriptor,
65    external_buffer: Option<DevicePtr>,
66) -> Result<()> {
67    descriptor.ensure_context(ctx)?;
68    matrix_a.ensure_context(ctx)?;
69    matrix_b.ensure_context(ctx)?;
70    matrix_c.ensure_context(ctx)?;
71    ctx.bind()?;
72    if ctx.scalar_pointer_mode()? != alpha.pointer_mode() {
73        ctx.set_scalar_pointer_mode(alpha.pointer_mode())?;
74    }
75
76    unsafe {
77        try_ffi!(sys::cusparseSpSM_analysis(
78            ctx.as_raw(),
79            op_a.into(),
80            op_b.into(),
81            alpha.ptr().cast(),
82            matrix_a.as_raw_const(),
83            matrix_b.as_raw_const(),
84            matrix_c.as_raw(),
85            Compute::data_type().into(),
86            algorithm.into(),
87            descriptor.as_raw(),
88            external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
89        ))?;
90    }
91    Ok(())
92}
93
94pub fn spsm_solve<Compute: DataTypeLike>(
95    ctx: &Context,
96    op_a: Operation,
97    op_b: Operation,
98    alpha: Scalar<'_, Compute>,
99    matrix_a: &SparseMatrixDescriptor,
100    matrix_b: &DenseMatrixDescriptor,
101    matrix_c: &mut DenseMatrixDescriptor,
102    algorithm: SpSMAlgorithm,
103    descriptor: &SpSmDescriptor,
104) -> Result<()> {
105    descriptor.ensure_context(ctx)?;
106    matrix_a.ensure_context(ctx)?;
107    matrix_b.ensure_context(ctx)?;
108    matrix_c.ensure_context(ctx)?;
109    ctx.bind()?;
110    if ctx.scalar_pointer_mode()? != alpha.pointer_mode() {
111        ctx.set_scalar_pointer_mode(alpha.pointer_mode())?;
112    }
113
114    unsafe {
115        try_ffi!(sys::cusparseSpSM_solve(
116            ctx.as_raw(),
117            op_a.into(),
118            op_b.into(),
119            alpha.ptr().cast(),
120            matrix_a.as_raw_const(),
121            matrix_b.as_raw_const(),
122            matrix_c.as_raw(),
123            Compute::data_type().into(),
124            algorithm.into(),
125            descriptor.as_raw(),
126        ))?;
127    }
128    Ok(())
129}
130
131pub fn spsm_update_matrix(
132    ctx: &Context,
133    descriptor: &SpSmDescriptor,
134    new_values: DevicePtr,
135    update: SpSmUpdate,
136) -> Result<()> {
137    descriptor.ensure_context(ctx)?;
138    ctx.bind()?;
139    unsafe {
140        try_ffi!(sys::cusparseSpSM_updateMatrix(
141            ctx.as_raw(),
142            descriptor.as_raw(),
143            new_values.as_ptr() as _,
144            update.into(),
145        ))?;
146    }
147    Ok(())
148}
singe_cusparse/spsm.rs

singe_cusparse/
spsm.rs
