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}