1#![warn(missing_debug_implementations)]
16
17use core::ffi::{c_int, c_void};
18use std::marker::PhantomData;
19
20use baracuda_cusolver_sys::{
21 cuComplex, cuDoubleComplex, cublasFillMode_t, cublasOperation_t, cusolver, cusolverDnHandle_t,
22 cusolverEigMode_t, cusolverStatus_t,
23};
24use baracuda_driver::{DeviceBuffer, Stream};
25use baracuda_types::{Complex32, Complex64, DeviceRepr};
26
27pub use baracuda_cusolver_sys::{cublasFillMode_t as Fill, cusolverEigMode_t as EigMode};
28
29pub type Error = baracuda_core::Error<cusolverStatus_t>;
31pub type Result<T, E = Error> = core::result::Result<T, E>;
33
34#[inline]
35fn check(status: cusolverStatus_t) -> Result<()> {
36 Error::check(status)
37}
38
39fn alloc_fail<E>(_e: E) -> Error {
41 Error::Status {
42 status: cusolverStatus_t::ALLOC_FAILED,
43 }
44}
45
46pub struct DnHandle {
50 handle: cusolverDnHandle_t,
51}
52
53unsafe impl Send for DnHandle {}
54
55impl core::fmt::Debug for DnHandle {
56 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
57 f.debug_struct("cusolver::DnHandle")
58 .field("handle", &self.handle)
59 .finish()
60 }
61}
62
63impl DnHandle {
64 pub fn new() -> Result<Self> {
66 let c = cusolver()?;
67 let cu = c.cusolver_dn_create()?;
68 let mut h: cusolverDnHandle_t = core::ptr::null_mut();
69 check(unsafe { cu(&mut h) })?;
70 Ok(Self { handle: h })
71 }
72
73 pub fn set_stream(&self, stream: &Stream) -> Result<()> {
76 let c = cusolver()?;
77 let cu = c.cusolver_dn_set_stream()?;
78 check(unsafe { cu(self.handle, stream.as_raw() as _) })
79 }
80
81 pub fn stream(&self) -> Result<*mut c_void> {
87 let c = cusolver()?;
88 let cu = c.cusolver_dn_get_stream()?;
89 let mut s: *mut c_void = core::ptr::null_mut();
90 check(unsafe { cu(self.handle, &mut s as *mut *mut c_void as *mut _) })?;
91 Ok(s)
92 }
93
94 pub fn version() -> Result<i32> {
97 let c = cusolver()?;
98 let cu = c.cusolver_get_version()?;
99 let mut v: c_int = 0;
100 check(unsafe { cu(&mut v) })?;
101 Ok(v)
102 }
103
104 #[inline]
106 pub fn as_raw(&self) -> cusolverDnHandle_t {
107 self.handle
108 }
109}
110
111impl Drop for DnHandle {
112 fn drop(&mut self) {
113 if let Ok(c) = cusolver()
114 && let Ok(cu) = c.cusolver_dn_destroy()
115 {
116 let _ = unsafe { cu(self.handle) };
117 }
118 }
119}
120
121#[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
123pub enum Op {
124 #[default]
126 N,
127 T,
129 C,
131}
132
133impl Op {
134 fn raw(self) -> cublasOperation_t {
135 match self {
136 Op::N => cublasOperation_t::N,
137 Op::T => cublasOperation_t::T,
138 Op::C => cublasOperation_t::C,
139 }
140 }
141}
142
143pub trait SolverScalar: DeviceRepr + Copy + 'static + sealed::Sealed {
147 type Real: DeviceRepr + Copy + 'static;
150
151 #[doc(hidden)]
153 unsafe fn getrf_buf(
154 h: cusolverDnHandle_t,
155 m: c_int,
156 n: c_int,
157 a: *mut Self,
158 lda: c_int,
159 lwork: *mut c_int,
160 ) -> cusolverStatus_t;
161
162 #[doc(hidden)]
164 unsafe fn getrf(
165 h: cusolverDnHandle_t,
166 m: c_int,
167 n: c_int,
168 a: *mut Self,
169 lda: c_int,
170 workspace: *mut Self,
171 ipiv: *mut c_int,
172 info: *mut c_int,
173 ) -> cusolverStatus_t;
174
175 #[doc(hidden)]
177 unsafe fn getrs(
178 h: cusolverDnHandle_t,
179 trans: cublasOperation_t,
180 n: c_int,
181 nrhs: c_int,
182 a: *const Self,
183 lda: c_int,
184 ipiv: *const c_int,
185 b: *mut Self,
186 ldb: c_int,
187 info: *mut c_int,
188 ) -> cusolverStatus_t;
189
190 #[doc(hidden)]
192 unsafe fn geqrf_buf(
193 h: cusolverDnHandle_t,
194 m: c_int,
195 n: c_int,
196 a: *mut Self,
197 lda: c_int,
198 lwork: *mut c_int,
199 ) -> cusolverStatus_t;
200
201 #[doc(hidden)]
203 unsafe fn geqrf(
204 h: cusolverDnHandle_t,
205 m: c_int,
206 n: c_int,
207 a: *mut Self,
208 lda: c_int,
209 tau: *mut Self,
210 workspace: *mut Self,
211 lwork: c_int,
212 info: *mut c_int,
213 ) -> cusolverStatus_t;
214
215 #[doc(hidden)]
217 unsafe fn potrf_buf(
218 h: cusolverDnHandle_t,
219 uplo: cublasFillMode_t,
220 n: c_int,
221 a: *mut Self,
222 lda: c_int,
223 lwork: *mut c_int,
224 ) -> cusolverStatus_t;
225
226 #[doc(hidden)]
228 unsafe fn potrf(
229 h: cusolverDnHandle_t,
230 uplo: cublasFillMode_t,
231 n: c_int,
232 a: *mut Self,
233 lda: c_int,
234 workspace: *mut Self,
235 lwork: c_int,
236 info: *mut c_int,
237 ) -> cusolverStatus_t;
238
239 #[doc(hidden)]
241 unsafe fn potrs(
242 h: cusolverDnHandle_t,
243 uplo: cublasFillMode_t,
244 n: c_int,
245 nrhs: c_int,
246 a: *const Self,
247 lda: c_int,
248 b: *mut Self,
249 ldb: c_int,
250 info: *mut c_int,
251 ) -> cusolverStatus_t;
252
253 #[doc(hidden)]
255 unsafe fn gesvd_buf(
256 h: cusolverDnHandle_t,
257 m: c_int,
258 n: c_int,
259 lwork: *mut c_int,
260 ) -> cusolverStatus_t;
261
262 #[doc(hidden)]
264 unsafe fn gesvd(
265 h: cusolverDnHandle_t,
266 jobu: u8,
267 jobvt: u8,
268 m: c_int,
269 n: c_int,
270 a: *mut Self,
271 lda: c_int,
272 s: *mut Self::Real,
273 u: *mut Self,
274 ldu: c_int,
275 vt: *mut Self,
276 ldvt: c_int,
277 work: *mut Self,
278 lwork: c_int,
279 rwork: *mut Self::Real,
280 info: *mut c_int,
281 ) -> cusolverStatus_t;
282
283 #[doc(hidden)]
285 unsafe fn syevd_buf(
286 h: cusolverDnHandle_t,
287 jobz: cusolverEigMode_t,
288 uplo: cublasFillMode_t,
289 n: c_int,
290 a: *const Self,
291 lda: c_int,
292 w: *const Self::Real,
293 lwork: *mut c_int,
294 ) -> cusolverStatus_t;
295
296 #[doc(hidden)]
298 unsafe fn syevd(
299 h: cusolverDnHandle_t,
300 jobz: cusolverEigMode_t,
301 uplo: cublasFillMode_t,
302 n: c_int,
303 a: *mut Self,
304 lda: c_int,
305 w: *mut Self::Real,
306 work: *mut Self,
307 lwork: c_int,
308 info: *mut c_int,
309 ) -> cusolverStatus_t;
310}
311
312mod sealed {
313 use baracuda_types::{Complex32, Complex64};
314 pub trait Sealed {}
315 impl Sealed for f32 {}
316 impl Sealed for f64 {}
317 impl Sealed for Complex32 {}
318 impl Sealed for Complex64 {}
319}
320
321macro_rules! real_impl {
322 ($t:ty, $getrf_buf:ident, $getrf:ident, $getrs:ident,
323 $geqrf_buf:ident, $geqrf:ident,
324 $potrf_buf:ident, $potrf:ident, $potrs:ident,
325 $gesvd_buf:ident, $gesvd:ident,
326 $syevd_buf:ident, $syevd:ident) => {
327 impl SolverScalar for $t {
328 type Real = $t;
329
330 unsafe fn getrf_buf(
331 h: cusolverDnHandle_t,
332 m: c_int,
333 n: c_int,
334 a: *mut $t,
335 lda: c_int,
336 lwork: *mut c_int,
337 ) -> cusolverStatus_t {
338 unsafe {
339 match cusolver().and_then(|c| c.$getrf_buf()) {
340 Ok(f) => f(h, m, n, a, lda, lwork),
341 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
342 }
343 }
344 }
345 unsafe fn getrf(
346 h: cusolverDnHandle_t,
347 m: c_int,
348 n: c_int,
349 a: *mut $t,
350 lda: c_int,
351 work: *mut $t,
352 ipiv: *mut c_int,
353 info: *mut c_int,
354 ) -> cusolverStatus_t {
355 unsafe {
356 match cusolver().and_then(|c| c.$getrf()) {
357 Ok(f) => f(h, m, n, a, lda, work, ipiv, info),
358 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
359 }
360 }
361 }
362 unsafe fn getrs(
363 h: cusolverDnHandle_t,
364 trans: cublasOperation_t,
365 n: c_int,
366 nrhs: c_int,
367 a: *const $t,
368 lda: c_int,
369 ipiv: *const c_int,
370 b: *mut $t,
371 ldb: c_int,
372 info: *mut c_int,
373 ) -> cusolverStatus_t {
374 unsafe {
375 match cusolver().and_then(|c| c.$getrs()) {
376 Ok(f) => f(h, trans, n, nrhs, a, lda, ipiv, b, ldb, info),
377 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
378 }
379 }
380 }
381 unsafe fn geqrf_buf(
382 h: cusolverDnHandle_t,
383 m: c_int,
384 n: c_int,
385 a: *mut $t,
386 lda: c_int,
387 lwork: *mut c_int,
388 ) -> cusolverStatus_t {
389 unsafe {
390 match cusolver().and_then(|c| c.$geqrf_buf()) {
391 Ok(f) => f(h, m, n, a, lda, lwork),
392 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
393 }
394 }
395 }
396 unsafe fn geqrf(
397 h: cusolverDnHandle_t,
398 m: c_int,
399 n: c_int,
400 a: *mut $t,
401 lda: c_int,
402 tau: *mut $t,
403 work: *mut $t,
404 lwork: c_int,
405 info: *mut c_int,
406 ) -> cusolverStatus_t {
407 unsafe {
408 match cusolver().and_then(|c| c.$geqrf()) {
409 Ok(f) => f(h, m, n, a, lda, tau, work, lwork, info),
410 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
411 }
412 }
413 }
414 unsafe fn potrf_buf(
415 h: cusolverDnHandle_t,
416 uplo: cublasFillMode_t,
417 n: c_int,
418 a: *mut $t,
419 lda: c_int,
420 lwork: *mut c_int,
421 ) -> cusolverStatus_t {
422 unsafe {
423 match cusolver().and_then(|c| c.$potrf_buf()) {
424 Ok(f) => f(h, uplo, n, a, lda, lwork),
425 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
426 }
427 }
428 }
429 unsafe fn potrf(
430 h: cusolverDnHandle_t,
431 uplo: cublasFillMode_t,
432 n: c_int,
433 a: *mut $t,
434 lda: c_int,
435 work: *mut $t,
436 lwork: c_int,
437 info: *mut c_int,
438 ) -> cusolverStatus_t {
439 unsafe {
440 match cusolver().and_then(|c| c.$potrf()) {
441 Ok(f) => f(h, uplo, n, a, lda, work, lwork, info),
442 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
443 }
444 }
445 }
446 unsafe fn potrs(
447 h: cusolverDnHandle_t,
448 uplo: cublasFillMode_t,
449 n: c_int,
450 nrhs: c_int,
451 a: *const $t,
452 lda: c_int,
453 b: *mut $t,
454 ldb: c_int,
455 info: *mut c_int,
456 ) -> cusolverStatus_t {
457 unsafe {
458 match cusolver().and_then(|c| c.$potrs()) {
459 Ok(f) => f(h, uplo, n, nrhs, a, lda, b, ldb, info),
460 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
461 }
462 }
463 }
464 unsafe fn gesvd_buf(
465 h: cusolverDnHandle_t,
466 m: c_int,
467 n: c_int,
468 lwork: *mut c_int,
469 ) -> cusolverStatus_t {
470 unsafe {
471 match cusolver().and_then(|c| c.$gesvd_buf()) {
472 Ok(f) => f(h, m, n, lwork),
473 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
474 }
475 }
476 }
477 unsafe fn gesvd(
478 h: cusolverDnHandle_t,
479 jobu: u8,
480 jobvt: u8,
481 m: c_int,
482 n: c_int,
483 a: *mut $t,
484 lda: c_int,
485 s: *mut $t,
486 u: *mut $t,
487 ldu: c_int,
488 vt: *mut $t,
489 ldvt: c_int,
490 work: *mut $t,
491 lwork: c_int,
492 rwork: *mut $t,
493 info: *mut c_int,
494 ) -> cusolverStatus_t {
495 unsafe {
496 match cusolver().and_then(|c| c.$gesvd()) {
497 Ok(f) => f(
498 h, jobu, jobvt, m, n, a, lda, s, u, ldu, vt, ldvt, work, lwork, rwork,
499 info,
500 ),
501 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
502 }
503 }
504 }
505 unsafe fn syevd_buf(
506 h: cusolverDnHandle_t,
507 jobz: cusolverEigMode_t,
508 uplo: cublasFillMode_t,
509 n: c_int,
510 a: *const $t,
511 lda: c_int,
512 w: *const $t,
513 lwork: *mut c_int,
514 ) -> cusolverStatus_t {
515 unsafe {
516 match cusolver().and_then(|c| c.$syevd_buf()) {
517 Ok(f) => f(h, jobz, uplo, n, a, lda, w, lwork),
518 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
519 }
520 }
521 }
522 unsafe fn syevd(
523 h: cusolverDnHandle_t,
524 jobz: cusolverEigMode_t,
525 uplo: cublasFillMode_t,
526 n: c_int,
527 a: *mut $t,
528 lda: c_int,
529 w: *mut $t,
530 work: *mut $t,
531 lwork: c_int,
532 info: *mut c_int,
533 ) -> cusolverStatus_t {
534 unsafe {
535 match cusolver().and_then(|c| c.$syevd()) {
536 Ok(f) => f(h, jobz, uplo, n, a, lda, w, work, lwork, info),
537 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
538 }
539 }
540 }
541 }
542 };
543}
544
545macro_rules! complex_impl {
546 ($t:ty, $real:ty, $raw:ty,
547 $getrf_buf:ident, $getrf:ident, $getrs:ident,
548 $geqrf_buf:ident, $geqrf:ident,
549 $potrf_buf:ident, $potrf:ident, $potrs:ident,
550 $gesvd_buf:ident, $gesvd:ident,
551 $heevd_buf:ident, $heevd:ident) => {
552 impl SolverScalar for $t {
553 type Real = $real;
554
555 unsafe fn getrf_buf(
556 h: cusolverDnHandle_t,
557 m: c_int,
558 n: c_int,
559 a: *mut $t,
560 lda: c_int,
561 lwork: *mut c_int,
562 ) -> cusolverStatus_t {
563 unsafe {
564 match cusolver().and_then(|c| c.$getrf_buf()) {
565 Ok(f) => f(h, m, n, a as *mut $raw, lda, lwork),
566 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
567 }
568 }
569 }
570 unsafe fn getrf(
571 h: cusolverDnHandle_t,
572 m: c_int,
573 n: c_int,
574 a: *mut $t,
575 lda: c_int,
576 work: *mut $t,
577 ipiv: *mut c_int,
578 info: *mut c_int,
579 ) -> cusolverStatus_t {
580 unsafe {
581 match cusolver().and_then(|c| c.$getrf()) {
582 Ok(f) => f(h, m, n, a as *mut $raw, lda, work as *mut $raw, ipiv, info),
583 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
584 }
585 }
586 }
587 unsafe fn getrs(
588 h: cusolverDnHandle_t,
589 trans: cublasOperation_t,
590 n: c_int,
591 nrhs: c_int,
592 a: *const $t,
593 lda: c_int,
594 ipiv: *const c_int,
595 b: *mut $t,
596 ldb: c_int,
597 info: *mut c_int,
598 ) -> cusolverStatus_t {
599 unsafe {
600 match cusolver().and_then(|c| c.$getrs()) {
601 Ok(f) => f(
602 h,
603 trans,
604 n,
605 nrhs,
606 a as *const $raw,
607 lda,
608 ipiv,
609 b as *mut $raw,
610 ldb,
611 info,
612 ),
613 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
614 }
615 }
616 }
617 unsafe fn geqrf_buf(
618 h: cusolverDnHandle_t,
619 m: c_int,
620 n: c_int,
621 a: *mut $t,
622 lda: c_int,
623 lwork: *mut c_int,
624 ) -> cusolverStatus_t {
625 unsafe {
626 match cusolver().and_then(|c| c.$geqrf_buf()) {
627 Ok(f) => f(h, m, n, a as *mut $raw, lda, lwork),
628 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
629 }
630 }
631 }
632 unsafe fn geqrf(
633 h: cusolverDnHandle_t,
634 m: c_int,
635 n: c_int,
636 a: *mut $t,
637 lda: c_int,
638 tau: *mut $t,
639 work: *mut $t,
640 lwork: c_int,
641 info: *mut c_int,
642 ) -> cusolverStatus_t {
643 unsafe {
644 match cusolver().and_then(|c| c.$geqrf()) {
645 Ok(f) => f(
646 h,
647 m,
648 n,
649 a as *mut $raw,
650 lda,
651 tau as *mut $raw,
652 work as *mut $raw,
653 lwork,
654 info,
655 ),
656 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
657 }
658 }
659 }
660 unsafe fn potrf_buf(
661 h: cusolverDnHandle_t,
662 uplo: cublasFillMode_t,
663 n: c_int,
664 a: *mut $t,
665 lda: c_int,
666 lwork: *mut c_int,
667 ) -> cusolverStatus_t {
668 unsafe {
669 match cusolver().and_then(|c| c.$potrf_buf()) {
670 Ok(f) => f(h, uplo, n, a as *mut $raw, lda, lwork),
671 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
672 }
673 }
674 }
675 unsafe fn potrf(
676 h: cusolverDnHandle_t,
677 uplo: cublasFillMode_t,
678 n: c_int,
679 a: *mut $t,
680 lda: c_int,
681 work: *mut $t,
682 lwork: c_int,
683 info: *mut c_int,
684 ) -> cusolverStatus_t {
685 unsafe {
686 match cusolver().and_then(|c| c.$potrf()) {
687 Ok(f) => f(
688 h,
689 uplo,
690 n,
691 a as *mut $raw,
692 lda,
693 work as *mut $raw,
694 lwork,
695 info,
696 ),
697 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
698 }
699 }
700 }
701 unsafe fn potrs(
702 h: cusolverDnHandle_t,
703 uplo: cublasFillMode_t,
704 n: c_int,
705 nrhs: c_int,
706 a: *const $t,
707 lda: c_int,
708 b: *mut $t,
709 ldb: c_int,
710 info: *mut c_int,
711 ) -> cusolverStatus_t {
712 unsafe {
713 match cusolver().and_then(|c| c.$potrs()) {
714 Ok(f) => f(
715 h,
716 uplo,
717 n,
718 nrhs,
719 a as *const $raw,
720 lda,
721 b as *mut $raw,
722 ldb,
723 info,
724 ),
725 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
726 }
727 }
728 }
729 unsafe fn gesvd_buf(
730 h: cusolverDnHandle_t,
731 m: c_int,
732 n: c_int,
733 lwork: *mut c_int,
734 ) -> cusolverStatus_t {
735 unsafe {
736 match cusolver().and_then(|c| c.$gesvd_buf()) {
737 Ok(f) => f(h, m, n, lwork),
738 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
739 }
740 }
741 }
742 unsafe fn gesvd(
743 h: cusolverDnHandle_t,
744 jobu: u8,
745 jobvt: u8,
746 m: c_int,
747 n: c_int,
748 a: *mut $t,
749 lda: c_int,
750 s: *mut $real,
751 u: *mut $t,
752 ldu: c_int,
753 vt: *mut $t,
754 ldvt: c_int,
755 work: *mut $t,
756 lwork: c_int,
757 rwork: *mut $real,
758 info: *mut c_int,
759 ) -> cusolverStatus_t {
760 unsafe {
761 match cusolver().and_then(|c| c.$gesvd()) {
762 Ok(f) => f(
763 h,
764 jobu,
765 jobvt,
766 m,
767 n,
768 a as *mut $raw,
769 lda,
770 s,
771 u as *mut $raw,
772 ldu,
773 vt as *mut $raw,
774 ldvt,
775 work as *mut $raw,
776 lwork,
777 rwork,
778 info,
779 ),
780 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
781 }
782 }
783 }
784 unsafe fn syevd_buf(
785 h: cusolverDnHandle_t,
786 jobz: cusolverEigMode_t,
787 uplo: cublasFillMode_t,
788 n: c_int,
789 a: *const $t,
790 lda: c_int,
791 w: *const $real,
792 lwork: *mut c_int,
793 ) -> cusolverStatus_t {
794 unsafe {
795 match cusolver().and_then(|c| c.$heevd_buf()) {
796 Ok(f) => f(h, jobz, uplo, n, a as *const $raw, lda, w, lwork),
797 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
798 }
799 }
800 }
801 unsafe fn syevd(
802 h: cusolverDnHandle_t,
803 jobz: cusolverEigMode_t,
804 uplo: cublasFillMode_t,
805 n: c_int,
806 a: *mut $t,
807 lda: c_int,
808 w: *mut $real,
809 work: *mut $t,
810 lwork: c_int,
811 info: *mut c_int,
812 ) -> cusolverStatus_t {
813 unsafe {
814 match cusolver().and_then(|c| c.$heevd()) {
815 Ok(f) => f(
816 h,
817 jobz,
818 uplo,
819 n,
820 a as *mut $raw,
821 lda,
822 w,
823 work as *mut $raw,
824 lwork,
825 info,
826 ),
827 Err(_) => cusolverStatus_t::NOT_INITIALIZED,
828 }
829 }
830 }
831 }
832 };
833}
834
835real_impl!(
836 f32,
837 cusolver_dn_sgetrf_buffer_size,
838 cusolver_dn_sgetrf,
839 cusolver_dn_sgetrs,
840 cusolver_dn_sgeqrf_buffer_size,
841 cusolver_dn_sgeqrf,
842 cusolver_dn_spotrf_buffer_size,
843 cusolver_dn_spotrf,
844 cusolver_dn_spotrs,
845 cusolver_dn_sgesvd_buffer_size,
846 cusolver_dn_sgesvd,
847 cusolver_dn_ssyevd_buffer_size,
848 cusolver_dn_ssyevd
849);
850
851real_impl!(
852 f64,
853 cusolver_dn_dgetrf_buffer_size,
854 cusolver_dn_dgetrf,
855 cusolver_dn_dgetrs,
856 cusolver_dn_dgeqrf_buffer_size,
857 cusolver_dn_dgeqrf,
858 cusolver_dn_dpotrf_buffer_size,
859 cusolver_dn_dpotrf,
860 cusolver_dn_dpotrs,
861 cusolver_dn_dgesvd_buffer_size,
862 cusolver_dn_dgesvd,
863 cusolver_dn_dsyevd_buffer_size,
864 cusolver_dn_dsyevd
865);
866
867complex_impl!(
868 Complex32,
869 f32,
870 cuComplex,
871 cusolver_dn_cgetrf_buffer_size,
872 cusolver_dn_cgetrf,
873 cusolver_dn_cgetrs,
874 cusolver_dn_cgeqrf_buffer_size,
875 cusolver_dn_cgeqrf,
876 cusolver_dn_cpotrf_buffer_size,
877 cusolver_dn_cpotrf,
878 cusolver_dn_cpotrs,
879 cusolver_dn_cgesvd_buffer_size,
880 cusolver_dn_cgesvd,
881 cusolver_dn_cheevd_buffer_size,
882 cusolver_dn_cheevd
883);
884
885complex_impl!(
886 Complex64,
887 f64,
888 cuDoubleComplex,
889 cusolver_dn_zgetrf_buffer_size,
890 cusolver_dn_zgetrf,
891 cusolver_dn_zgetrs,
892 cusolver_dn_zgeqrf_buffer_size,
893 cusolver_dn_zgeqrf,
894 cusolver_dn_zpotrf_buffer_size,
895 cusolver_dn_zpotrf,
896 cusolver_dn_zpotrs,
897 cusolver_dn_zgesvd_buffer_size,
898 cusolver_dn_zgesvd,
899 cusolver_dn_zheevd_buffer_size,
900 cusolver_dn_zheevd
901);
902
903pub fn getrf<T: SolverScalar>(
935 handle: &DnHandle,
936 m: i32,
937 n: i32,
938 a: &mut DeviceBuffer<T>,
939 lda: i32,
940 ipiv: &mut DeviceBuffer<i32>,
941 info: &mut DeviceBuffer<i32>,
942) -> Result<()> {
943 let mut lwork: c_int = 0;
944 check(unsafe { T::getrf_buf(handle.handle, m, n, a.as_raw().0 as *mut T, lda, &mut lwork) })?;
945 let workspace = DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
946 check(unsafe {
947 T::getrf(
948 handle.handle,
949 m,
950 n,
951 a.as_raw().0 as *mut T,
952 lda,
953 workspace.as_raw().0 as *mut T,
954 ipiv.as_raw().0 as *mut c_int,
955 info.as_raw().0 as *mut c_int,
956 )
957 })
958}
959
960pub fn getrs<T: SolverScalar>(
962 handle: &DnHandle,
963 trans: Op,
964 n: i32,
965 nrhs: i32,
966 a: &DeviceBuffer<T>,
967 lda: i32,
968 ipiv: &DeviceBuffer<i32>,
969 b: &mut DeviceBuffer<T>,
970 ldb: i32,
971 info: &mut DeviceBuffer<i32>,
972) -> Result<()> {
973 check(unsafe {
974 T::getrs(
975 handle.handle,
976 trans.raw(),
977 n,
978 nrhs,
979 a.as_raw().0 as *const T,
980 lda,
981 ipiv.as_raw().0 as *const c_int,
982 b.as_raw().0 as *mut T,
983 ldb,
984 info.as_raw().0 as *mut c_int,
985 )
986 })
987}
988
989pub fn geqrf<T: SolverScalar>(
1017 handle: &DnHandle,
1018 m: i32,
1019 n: i32,
1020 a: &mut DeviceBuffer<T>,
1021 lda: i32,
1022 tau: &mut DeviceBuffer<T>,
1023 info: &mut DeviceBuffer<i32>,
1024) -> Result<()> {
1025 let mut lwork: c_int = 0;
1026 check(unsafe { T::geqrf_buf(handle.handle, m, n, a.as_raw().0 as *mut T, lda, &mut lwork) })?;
1027 let workspace = DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
1028 check(unsafe {
1029 T::geqrf(
1030 handle.handle,
1031 m,
1032 n,
1033 a.as_raw().0 as *mut T,
1034 lda,
1035 tau.as_raw().0 as *mut T,
1036 workspace.as_raw().0 as *mut T,
1037 lwork,
1038 info.as_raw().0 as *mut c_int,
1039 )
1040 })
1041}
1042
1043pub fn potrf<T: SolverScalar>(
1070 handle: &DnHandle,
1071 uplo: Fill,
1072 n: i32,
1073 a: &mut DeviceBuffer<T>,
1074 lda: i32,
1075 info: &mut DeviceBuffer<i32>,
1076) -> Result<()> {
1077 let mut lwork: c_int = 0;
1078 check(unsafe {
1079 T::potrf_buf(
1080 handle.handle,
1081 uplo,
1082 n,
1083 a.as_raw().0 as *mut T,
1084 lda,
1085 &mut lwork,
1086 )
1087 })?;
1088 let workspace = DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
1089 check(unsafe {
1090 T::potrf(
1091 handle.handle,
1092 uplo,
1093 n,
1094 a.as_raw().0 as *mut T,
1095 lda,
1096 workspace.as_raw().0 as *mut T,
1097 lwork,
1098 info.as_raw().0 as *mut c_int,
1099 )
1100 })
1101}
1102
1103pub fn potrs<T: SolverScalar>(
1105 handle: &DnHandle,
1106 uplo: Fill,
1107 n: i32,
1108 nrhs: i32,
1109 a: &DeviceBuffer<T>,
1110 lda: i32,
1111 b: &mut DeviceBuffer<T>,
1112 ldb: i32,
1113 info: &mut DeviceBuffer<i32>,
1114) -> Result<()> {
1115 check(unsafe {
1116 T::potrs(
1117 handle.handle,
1118 uplo,
1119 n,
1120 nrhs,
1121 a.as_raw().0 as *const T,
1122 lda,
1123 b.as_raw().0 as *mut T,
1124 ldb,
1125 info.as_raw().0 as *mut c_int,
1126 )
1127 })
1128}
1129
1130pub fn gesvd<T: SolverScalar>(
1169 handle: &DnHandle,
1170 jobu: u8,
1171 jobvt: u8,
1172 m: i32,
1173 n: i32,
1174 a: &mut DeviceBuffer<T>,
1175 lda: i32,
1176 s: &mut DeviceBuffer<T::Real>,
1177 u: &mut DeviceBuffer<T>,
1178 ldu: i32,
1179 vt: &mut DeviceBuffer<T>,
1180 ldvt: i32,
1181 rwork: &mut DeviceBuffer<T::Real>,
1182 info: &mut DeviceBuffer<i32>,
1183) -> Result<()> {
1184 let mut lwork: c_int = 0;
1185 check(unsafe { T::gesvd_buf(handle.handle, m, n, &mut lwork) })?;
1186 let workspace = DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
1187 check(unsafe {
1188 T::gesvd(
1189 handle.handle,
1190 jobu,
1191 jobvt,
1192 m,
1193 n,
1194 a.as_raw().0 as *mut T,
1195 lda,
1196 s.as_raw().0 as *mut T::Real,
1197 u.as_raw().0 as *mut T,
1198 ldu,
1199 vt.as_raw().0 as *mut T,
1200 ldvt,
1201 workspace.as_raw().0 as *mut T,
1202 lwork,
1203 rwork.as_raw().0 as *mut T::Real,
1204 info.as_raw().0 as *mut c_int,
1205 )
1206 })
1207}
1208
1209pub fn syevd<T: SolverScalar>(
1211 handle: &DnHandle,
1212 jobz: EigMode,
1213 uplo: Fill,
1214 n: i32,
1215 a: &mut DeviceBuffer<T>,
1216 lda: i32,
1217 w: &mut DeviceBuffer<T::Real>,
1218 info: &mut DeviceBuffer<i32>,
1219) -> Result<()> {
1220 let mut lwork: c_int = 0;
1221 check(unsafe {
1222 T::syevd_buf(
1223 handle.handle,
1224 jobz,
1225 uplo,
1226 n,
1227 a.as_raw().0 as *const T,
1228 lda,
1229 w.as_raw().0 as *const T::Real,
1230 &mut lwork,
1231 )
1232 })?;
1233 let workspace = DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
1234 check(unsafe {
1235 T::syevd(
1236 handle.handle,
1237 jobz,
1238 uplo,
1239 n,
1240 a.as_raw().0 as *mut T,
1241 lda,
1242 w.as_raw().0 as *mut T::Real,
1243 workspace.as_raw().0 as *mut T,
1244 lwork,
1245 info.as_raw().0 as *mut c_int,
1246 )
1247 })
1248}
1249
1250pub use baracuda_cusolver_sys::{gesvdjInfo_t as GesvdjInfoRaw, syevjInfo_t as SyevjInfoRaw};
1253
1254#[derive(Debug)]
1256pub struct SyevjInfo {
1257 raw: SyevjInfoRaw,
1258}
1259
1260impl SyevjInfo {
1261 pub fn new() -> Result<Self> {
1263 let c = cusolver()?;
1264 let cu = c.cusolver_dn_create_syevj_info()?;
1265 let mut raw: SyevjInfoRaw = core::ptr::null_mut();
1266 check(unsafe { cu(&mut raw) })?;
1267 Ok(Self { raw })
1268 }
1269
1270 pub fn set_tolerance(&self, tol: f64) -> Result<()> {
1273 let c = cusolver()?;
1274 let cu = c.cusolver_dn_xsyevj_set_tolerance()?;
1275 check(unsafe { cu(self.raw, tol) })
1276 }
1277
1278 pub fn set_max_sweeps(&self, n: i32) -> Result<()> {
1281 let c = cusolver()?;
1282 let cu = c.cusolver_dn_xsyevj_set_max_sweeps()?;
1283 check(unsafe { cu(self.raw, n) })
1284 }
1285
1286 pub fn as_raw(&self) -> SyevjInfoRaw {
1288 self.raw
1289 }
1290}
1291
1292impl Drop for SyevjInfo {
1293 fn drop(&mut self) {
1294 if let Ok(c) = cusolver()
1295 && let Ok(cu) = c.cusolver_dn_destroy_syevj_info()
1296 {
1297 let _ = unsafe { cu(self.raw) };
1298 }
1299 }
1300}
1301
1302#[derive(Debug)]
1304pub struct GesvdjInfo {
1305 raw: GesvdjInfoRaw,
1306}
1307
1308impl GesvdjInfo {
1309 pub fn new() -> Result<Self> {
1311 let c = cusolver()?;
1312 let cu = c.cusolver_dn_create_gesvdj_info()?;
1313 let mut raw: GesvdjInfoRaw = core::ptr::null_mut();
1314 check(unsafe { cu(&mut raw) })?;
1315 Ok(Self { raw })
1316 }
1317
1318 pub fn as_raw(&self) -> GesvdjInfoRaw {
1320 self.raw
1321 }
1322}
1323
1324impl Drop for GesvdjInfo {
1325 fn drop(&mut self) {
1326 if let Ok(c) = cusolver()
1327 && let Ok(cu) = c.cusolver_dn_destroy_gesvdj_info()
1328 {
1329 let _ = unsafe { cu(self.raw) };
1330 }
1331 }
1332}
1333
1334pub fn syevj<T: SolverScalar>(
1337 handle: &DnHandle,
1338 jobz: EigMode,
1339 uplo: Fill,
1340 n: i32,
1341 a: &mut DeviceBuffer<T>,
1342 lda: i32,
1343 w: &mut DeviceBuffer<T::Real>,
1344 info: &mut DeviceBuffer<i32>,
1345 params: &SyevjInfo,
1346) -> Result<()> {
1347 use baracuda_cusolver_sys::{cuComplex, cuDoubleComplex};
1348 use core::mem;
1349
1350 let mut lwork: c_int = 0;
1351
1352 macro_rules! dispatch_real {
1355 ($t:ty, $bufsize:ident, $solve:ident) => {{
1356 let c = cusolver()?;
1357 check(unsafe {
1358 (c.$bufsize()?)(
1359 handle.as_raw(),
1360 jobz,
1361 uplo,
1362 n,
1363 a.as_raw().0 as *const $t,
1364 lda,
1365 w.as_raw().0 as *const $t,
1366 &mut lwork,
1367 params.raw,
1368 )
1369 })?;
1370 let workspace =
1371 DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
1372 check(unsafe {
1373 (c.$solve()?)(
1374 handle.as_raw(),
1375 jobz,
1376 uplo,
1377 n,
1378 a.as_raw().0 as *mut $t,
1379 lda,
1380 w.as_raw().0 as *mut $t,
1381 workspace.as_raw().0 as *mut $t,
1382 lwork,
1383 info.as_raw().0 as *mut c_int,
1384 params.raw,
1385 )
1386 })
1387 }};
1388 }
1389 macro_rules! dispatch_complex {
1390 ($t:ty, $real:ty, $raw:ty, $bufsize:ident, $solve:ident) => {{
1391 let c = cusolver()?;
1392 check(unsafe {
1393 (c.$bufsize()?)(
1394 handle.as_raw(),
1395 jobz,
1396 uplo,
1397 n,
1398 a.as_raw().0 as *const $raw,
1399 lda,
1400 w.as_raw().0 as *const $real,
1401 &mut lwork,
1402 params.raw,
1403 )
1404 })?;
1405 let workspace =
1406 DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
1407 check(unsafe {
1408 (c.$solve()?)(
1409 handle.as_raw(),
1410 jobz,
1411 uplo,
1412 n,
1413 a.as_raw().0 as *mut $raw,
1414 lda,
1415 w.as_raw().0 as *mut $real,
1416 workspace.as_raw().0 as *mut $raw,
1417 lwork,
1418 info.as_raw().0 as *mut c_int,
1419 params.raw,
1420 )
1421 })
1422 }};
1423 }
1424
1425 if mem::size_of::<T>() == mem::size_of::<f32>() && mem::size_of::<T::Real>() == 4 {
1426 dispatch_real!(f32, cusolver_dn_ssyevj_buffer_size, cusolver_dn_ssyevj)
1427 } else if mem::size_of::<T>() == mem::size_of::<f64>() && mem::size_of::<T::Real>() == 8 {
1428 dispatch_real!(f64, cusolver_dn_dsyevj_buffer_size, cusolver_dn_dsyevj)
1429 } else if mem::size_of::<T>() == mem::size_of::<Complex32>() {
1430 dispatch_complex!(
1431 Complex32,
1432 f32,
1433 cuComplex,
1434 cusolver_dn_cheevj_buffer_size,
1435 cusolver_dn_cheevj
1436 )
1437 } else {
1438 dispatch_complex!(
1439 Complex64,
1440 f64,
1441 cuDoubleComplex,
1442 cusolver_dn_zheevj_buffer_size,
1443 cusolver_dn_zheevj
1444 )
1445 }
1446}
1447
1448pub fn gesvdj<T: SolverScalar>(
1450 handle: &DnHandle,
1451 jobz: EigMode,
1452 econ: bool,
1453 m: i32,
1454 n: i32,
1455 a: &mut DeviceBuffer<T>,
1456 lda: i32,
1457 s: &mut DeviceBuffer<T::Real>,
1458 u: &mut DeviceBuffer<T>,
1459 ldu: i32,
1460 v: &mut DeviceBuffer<T>,
1461 ldv: i32,
1462 info: &mut DeviceBuffer<i32>,
1463 params: &GesvdjInfo,
1464) -> Result<()> {
1465 use baracuda_cusolver_sys::{cuComplex, cuDoubleComplex};
1466 use core::mem;
1467
1468 let mut lwork: c_int = 0;
1469 let econ_i = if econ { 1 } else { 0 };
1470
1471 macro_rules! dispatch_real {
1472 ($t:ty, $bufsize:ident, $solve:ident) => {{
1473 let c = cusolver()?;
1474 check(unsafe {
1475 (c.$bufsize()?)(
1476 handle.as_raw(),
1477 jobz,
1478 econ_i,
1479 m,
1480 n,
1481 a.as_raw().0 as *const $t,
1482 lda,
1483 s.as_raw().0 as *const $t,
1484 u.as_raw().0 as *const $t,
1485 ldu,
1486 v.as_raw().0 as *const $t,
1487 ldv,
1488 &mut lwork,
1489 params.raw,
1490 )
1491 })?;
1492 let workspace =
1493 DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
1494 check(unsafe {
1495 (c.$solve()?)(
1496 handle.as_raw(),
1497 jobz,
1498 econ_i,
1499 m,
1500 n,
1501 a.as_raw().0 as *mut $t,
1502 lda,
1503 s.as_raw().0 as *mut $t,
1504 u.as_raw().0 as *mut $t,
1505 ldu,
1506 v.as_raw().0 as *mut $t,
1507 ldv,
1508 workspace.as_raw().0 as *mut $t,
1509 lwork,
1510 info.as_raw().0 as *mut c_int,
1511 params.raw,
1512 )
1513 })
1514 }};
1515 }
1516 macro_rules! dispatch_complex {
1517 ($t:ty, $real:ty, $raw:ty, $bufsize:ident, $solve:ident) => {{
1518 let c = cusolver()?;
1519 check(unsafe {
1520 (c.$bufsize()?)(
1521 handle.as_raw(),
1522 jobz,
1523 econ_i,
1524 m,
1525 n,
1526 a.as_raw().0 as *const $raw,
1527 lda,
1528 s.as_raw().0 as *const $real,
1529 u.as_raw().0 as *const $raw,
1530 ldu,
1531 v.as_raw().0 as *const $raw,
1532 ldv,
1533 &mut lwork,
1534 params.raw,
1535 )
1536 })?;
1537 let workspace =
1538 DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
1539 check(unsafe {
1540 (c.$solve()?)(
1541 handle.as_raw(),
1542 jobz,
1543 econ_i,
1544 m,
1545 n,
1546 a.as_raw().0 as *mut $raw,
1547 lda,
1548 s.as_raw().0 as *mut $real,
1549 u.as_raw().0 as *mut $raw,
1550 ldu,
1551 v.as_raw().0 as *mut $raw,
1552 ldv,
1553 workspace.as_raw().0 as *mut $raw,
1554 lwork,
1555 info.as_raw().0 as *mut c_int,
1556 params.raw,
1557 )
1558 })
1559 }};
1560 }
1561
1562 if mem::size_of::<T>() == mem::size_of::<f32>() && mem::size_of::<T::Real>() == 4 {
1563 dispatch_real!(f32, cusolver_dn_sgesvdj_buffer_size, cusolver_dn_sgesvdj)
1564 } else if mem::size_of::<T>() == mem::size_of::<f64>() && mem::size_of::<T::Real>() == 8 {
1565 dispatch_real!(f64, cusolver_dn_dgesvdj_buffer_size, cusolver_dn_dgesvdj)
1566 } else if mem::size_of::<T>() == mem::size_of::<Complex32>() {
1567 dispatch_complex!(
1568 Complex32,
1569 f32,
1570 cuComplex,
1571 cusolver_dn_cgesvdj_buffer_size,
1572 cusolver_dn_cgesvdj
1573 )
1574 } else {
1575 dispatch_complex!(
1576 Complex64,
1577 f64,
1578 cuDoubleComplex,
1579 cusolver_dn_zgesvdj_buffer_size,
1580 cusolver_dn_zgesvdj
1581 )
1582 }
1583}
1584
1585pub fn orgqr<T: SolverScalar>(
1590 handle: &DnHandle,
1591 m: i32,
1592 n: i32,
1593 k: i32,
1594 a: &mut DeviceBuffer<T>,
1595 lda: i32,
1596 tau: &DeviceBuffer<T>,
1597 info: &mut DeviceBuffer<i32>,
1598) -> Result<()> {
1599 use baracuda_cusolver_sys::{cuComplex, cuDoubleComplex};
1600 use core::mem;
1601
1602 let mut lwork: c_int = 0;
1603 macro_rules! dispatch {
1604 ($t:ty, $raw:ty, $bufsize:ident, $solve:ident) => {{
1605 let c = cusolver()?;
1606 check(unsafe {
1607 (c.$bufsize()?)(
1608 handle.as_raw(),
1609 m,
1610 n,
1611 k,
1612 a.as_raw().0 as *const $raw,
1613 lda,
1614 tau.as_raw().0 as *const $raw,
1615 &mut lwork,
1616 )
1617 })?;
1618 let workspace =
1619 DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
1620 check(unsafe {
1621 (c.$solve()?)(
1622 handle.as_raw(),
1623 m,
1624 n,
1625 k,
1626 a.as_raw().0 as *mut $raw,
1627 lda,
1628 tau.as_raw().0 as *const $raw,
1629 workspace.as_raw().0 as *mut $raw,
1630 lwork,
1631 info.as_raw().0 as *mut c_int,
1632 )
1633 })
1634 }};
1635 }
1636
1637 if mem::size_of::<T>() == mem::size_of::<f32>() && mem::size_of::<T::Real>() == 4 {
1638 dispatch!(f32, f32, cusolver_dn_sorgqr_buffer_size, cusolver_dn_sorgqr)
1639 } else if mem::size_of::<T>() == mem::size_of::<f64>() && mem::size_of::<T::Real>() == 8 {
1640 dispatch!(f64, f64, cusolver_dn_dorgqr_buffer_size, cusolver_dn_dorgqr)
1641 } else if mem::size_of::<T>() == mem::size_of::<Complex32>() {
1642 dispatch!(
1643 Complex32,
1644 cuComplex,
1645 cusolver_dn_cungqr_buffer_size,
1646 cusolver_dn_cungqr
1647 )
1648 } else {
1649 dispatch!(
1650 Complex64,
1651 cuDoubleComplex,
1652 cusolver_dn_zungqr_buffer_size,
1653 cusolver_dn_zungqr
1654 )
1655 }
1656}
1657
1658#[derive(Copy, Clone, Debug, Eq, PartialEq)]
1660pub enum Side {
1661 Left,
1663 Right,
1665}
1666
1667impl Side {
1668 fn raw(self) -> core::ffi::c_int {
1669 match self {
1670 Side::Left => 0,
1671 Side::Right => 1,
1672 }
1673 }
1674}
1675
1676pub fn ormqr<T: SolverScalar>(
1679 handle: &DnHandle,
1680 side: Side,
1681 trans: Op,
1682 m: i32,
1683 n: i32,
1684 k: i32,
1685 a: &DeviceBuffer<T>,
1686 lda: i32,
1687 tau: &DeviceBuffer<T>,
1688 c_mat: &mut DeviceBuffer<T>,
1689 ldc: i32,
1690 info: &mut DeviceBuffer<i32>,
1691) -> Result<()> {
1692 use baracuda_cusolver_sys::{cuComplex, cuDoubleComplex};
1693 use core::mem;
1694
1695 let mut lwork: c_int = 0;
1696 let side_i = side.raw();
1697 macro_rules! dispatch {
1698 ($t:ty, $raw:ty, $bufsize:ident, $solve:ident) => {{
1699 let ca = cusolver()?;
1700 check(unsafe {
1701 (ca.$bufsize()?)(
1702 handle.as_raw(),
1703 side_i,
1704 trans.raw(),
1705 m,
1706 n,
1707 k,
1708 a.as_raw().0 as *const $raw,
1709 lda,
1710 tau.as_raw().0 as *const $raw,
1711 c_mat.as_raw().0 as *const $raw,
1712 ldc,
1713 &mut lwork,
1714 )
1715 })?;
1716 let workspace =
1717 DeviceBuffer::<T>::new(c_mat.context(), lwork as usize).map_err(alloc_fail)?;
1718 check(unsafe {
1719 (ca.$solve()?)(
1720 handle.as_raw(),
1721 side_i,
1722 trans.raw(),
1723 m,
1724 n,
1725 k,
1726 a.as_raw().0 as *const $raw,
1727 lda,
1728 tau.as_raw().0 as *const $raw,
1729 c_mat.as_raw().0 as *mut $raw,
1730 ldc,
1731 workspace.as_raw().0 as *mut $raw,
1732 lwork,
1733 info.as_raw().0 as *mut c_int,
1734 )
1735 })
1736 }};
1737 }
1738
1739 if mem::size_of::<T>() == mem::size_of::<f32>() && mem::size_of::<T::Real>() == 4 {
1740 dispatch!(f32, f32, cusolver_dn_sormqr_buffer_size, cusolver_dn_sormqr)
1741 } else if mem::size_of::<T>() == mem::size_of::<f64>() && mem::size_of::<T::Real>() == 8 {
1742 dispatch!(f64, f64, cusolver_dn_dormqr_buffer_size, cusolver_dn_dormqr)
1743 } else if mem::size_of::<T>() == mem::size_of::<Complex32>() {
1744 dispatch!(
1745 Complex32,
1746 cuComplex,
1747 cusolver_dn_cunmqr_buffer_size,
1748 cusolver_dn_cunmqr
1749 )
1750 } else {
1751 dispatch!(
1752 Complex64,
1753 cuDoubleComplex,
1754 cusolver_dn_zunmqr_buffer_size,
1755 cusolver_dn_zunmqr
1756 )
1757 }
1758}
1759
1760pub fn gels<T: SolverScalar>(
1767 handle: &DnHandle,
1768 m: i32,
1769 n: i32,
1770 nrhs: i32,
1771 a: &mut DeviceBuffer<T>,
1772 lda: i32,
1773 b: &mut DeviceBuffer<T>,
1774 ldb: i32,
1775 x: &mut DeviceBuffer<T>,
1776 ldx: i32,
1777 info: &mut DeviceBuffer<i32>,
1778) -> Result<i32> {
1779 use baracuda_cusolver_sys::{cuComplex, cuDoubleComplex};
1780 use core::mem;
1781
1782 let mut bytes: usize = 0;
1783
1784 macro_rules! dispatch {
1785 ($t:ty, $raw:ty, $bufsize:ident, $solve:ident) => {{
1786 let cs = cusolver()?;
1787 check(unsafe {
1788 (cs.$bufsize()?)(
1789 handle.as_raw(),
1790 m,
1791 n,
1792 nrhs,
1793 a.as_raw().0 as *mut $raw,
1794 lda,
1795 b.as_raw().0 as *mut $raw,
1796 ldb,
1797 x.as_raw().0 as *mut $raw,
1798 ldx,
1799 core::ptr::null_mut(),
1800 &mut bytes,
1801 )
1802 })?;
1803 let units = bytes.div_ceil(mem::size_of::<T>());
1805 let workspace = DeviceBuffer::<T>::new(a.context(), units).map_err(alloc_fail)?;
1806 let mut iter: c_int = 0;
1807 check(unsafe {
1808 (cs.$solve()?)(
1809 handle.as_raw(),
1810 m,
1811 n,
1812 nrhs,
1813 a.as_raw().0 as *mut $raw,
1814 lda,
1815 b.as_raw().0 as *mut $raw,
1816 ldb,
1817 x.as_raw().0 as *mut $raw,
1818 ldx,
1819 workspace.as_raw().0 as *mut c_void,
1820 bytes,
1821 &mut iter,
1822 info.as_raw().0 as *mut c_int,
1823 )
1824 })?;
1825 Ok(iter)
1826 }};
1827 }
1828
1829 if mem::size_of::<T>() == mem::size_of::<f32>() && mem::size_of::<T::Real>() == 4 {
1830 dispatch!(f32, f32, cusolver_dn_ssgels_buffer_size, cusolver_dn_ssgels)
1831 } else if mem::size_of::<T>() == mem::size_of::<f64>() && mem::size_of::<T::Real>() == 8 {
1832 dispatch!(f64, f64, cusolver_dn_ddgels_buffer_size, cusolver_dn_ddgels)
1833 } else if mem::size_of::<T>() == mem::size_of::<Complex32>() {
1834 dispatch!(
1835 Complex32,
1836 cuComplex,
1837 cusolver_dn_ccgels_buffer_size,
1838 cusolver_dn_ccgels
1839 )
1840 } else {
1841 dispatch!(
1842 Complex64,
1843 cuDoubleComplex,
1844 cusolver_dn_zzgels_buffer_size,
1845 cusolver_dn_zzgels
1846 )
1847 }
1848}
1849
1850pub fn potri<T: SolverScalar>(
1856 handle: &DnHandle,
1857 uplo: Fill,
1858 n: i32,
1859 a: &mut DeviceBuffer<T>,
1860 lda: i32,
1861 info: &mut DeviceBuffer<i32>,
1862) -> Result<()> {
1863 use baracuda_cusolver_sys::{cuComplex, cuDoubleComplex};
1864 use core::mem;
1865
1866 let mut lwork: c_int = 0;
1867 macro_rules! dispatch {
1868 ($t:ty, $raw:ty, $bufsize:ident, $solve:ident) => {{
1869 let cs = cusolver()?;
1870 check(unsafe {
1871 (cs.$bufsize()?)(
1872 handle.as_raw(),
1873 uplo,
1874 n,
1875 a.as_raw().0 as *mut $raw,
1876 lda,
1877 &mut lwork,
1878 )
1879 })?;
1880 let workspace =
1881 DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
1882 check(unsafe {
1883 (cs.$solve()?)(
1884 handle.as_raw(),
1885 uplo,
1886 n,
1887 a.as_raw().0 as *mut $raw,
1888 lda,
1889 workspace.as_raw().0 as *mut $raw,
1890 lwork,
1891 info.as_raw().0 as *mut c_int,
1892 )
1893 })
1894 }};
1895 }
1896
1897 if mem::size_of::<T>() == mem::size_of::<f32>() && mem::size_of::<T::Real>() == 4 {
1898 dispatch!(f32, f32, cusolver_dn_spotri_buffer_size, cusolver_dn_spotri)
1899 } else if mem::size_of::<T>() == mem::size_of::<f64>() && mem::size_of::<T::Real>() == 8 {
1900 dispatch!(f64, f64, cusolver_dn_dpotri_buffer_size, cusolver_dn_dpotri)
1901 } else if mem::size_of::<T>() == mem::size_of::<Complex32>() {
1902 dispatch!(
1903 Complex32,
1904 cuComplex,
1905 cusolver_dn_cpotri_buffer_size,
1906 cusolver_dn_cpotri
1907 )
1908 } else {
1909 dispatch!(
1910 Complex64,
1911 cuDoubleComplex,
1912 cusolver_dn_zpotri_buffer_size,
1913 cusolver_dn_zpotri
1914 )
1915 }
1916}
1917
1918pub fn syevj_batched<T: SolverScalar>(
1924 handle: &DnHandle,
1925 jobz: EigMode,
1926 uplo: Fill,
1927 n: i32,
1928 a: &mut DeviceBuffer<T>,
1929 lda: i32,
1930 w: &mut DeviceBuffer<T::Real>,
1931 info: &mut DeviceBuffer<i32>,
1932 params: &SyevjInfo,
1933 batch_size: i32,
1934) -> Result<()> {
1935 use baracuda_cusolver_sys::{cuComplex, cuDoubleComplex};
1936 use core::mem;
1937
1938 let mut lwork: c_int = 0;
1939 macro_rules! dispatch_real {
1940 ($t:ty, $bufsize:ident, $solve:ident) => {{
1941 let c = cusolver()?;
1942 check(unsafe {
1943 (c.$bufsize()?)(
1944 handle.as_raw(),
1945 jobz,
1946 uplo,
1947 n,
1948 a.as_raw().0 as *const $t,
1949 lda,
1950 w.as_raw().0 as *const $t,
1951 &mut lwork,
1952 params.as_raw(),
1953 batch_size,
1954 )
1955 })?;
1956 let workspace =
1957 DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
1958 check(unsafe {
1959 (c.$solve()?)(
1960 handle.as_raw(),
1961 jobz,
1962 uplo,
1963 n,
1964 a.as_raw().0 as *mut $t,
1965 lda,
1966 w.as_raw().0 as *mut $t,
1967 workspace.as_raw().0 as *mut $t,
1968 lwork,
1969 info.as_raw().0 as *mut c_int,
1970 params.as_raw(),
1971 batch_size,
1972 )
1973 })
1974 }};
1975 }
1976 macro_rules! dispatch_complex {
1977 ($t:ty, $real:ty, $raw:ty, $bufsize:ident, $solve:ident) => {{
1978 let c = cusolver()?;
1979 check(unsafe {
1980 (c.$bufsize()?)(
1981 handle.as_raw(),
1982 jobz,
1983 uplo,
1984 n,
1985 a.as_raw().0 as *const $raw,
1986 lda,
1987 w.as_raw().0 as *const $real,
1988 &mut lwork,
1989 params.as_raw(),
1990 batch_size,
1991 )
1992 })?;
1993 let workspace =
1994 DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
1995 check(unsafe {
1996 (c.$solve()?)(
1997 handle.as_raw(),
1998 jobz,
1999 uplo,
2000 n,
2001 a.as_raw().0 as *mut $raw,
2002 lda,
2003 w.as_raw().0 as *mut $real,
2004 workspace.as_raw().0 as *mut $raw,
2005 lwork,
2006 info.as_raw().0 as *mut c_int,
2007 params.as_raw(),
2008 batch_size,
2009 )
2010 })
2011 }};
2012 }
2013
2014 if mem::size_of::<T>() == mem::size_of::<f32>() && mem::size_of::<T::Real>() == 4 {
2015 dispatch_real!(
2016 f32,
2017 cusolver_dn_ssyevj_batched_buffer_size,
2018 cusolver_dn_ssyevj_batched
2019 )
2020 } else if mem::size_of::<T>() == mem::size_of::<f64>() && mem::size_of::<T::Real>() == 8 {
2021 dispatch_real!(
2022 f64,
2023 cusolver_dn_dsyevj_batched_buffer_size,
2024 cusolver_dn_dsyevj_batched
2025 )
2026 } else if mem::size_of::<T>() == mem::size_of::<Complex32>() {
2027 dispatch_complex!(
2028 Complex32,
2029 f32,
2030 cuComplex,
2031 cusolver_dn_cheevj_batched_buffer_size,
2032 cusolver_dn_cheevj_batched
2033 )
2034 } else {
2035 dispatch_complex!(
2036 Complex64,
2037 f64,
2038 cuDoubleComplex,
2039 cusolver_dn_zheevj_batched_buffer_size,
2040 cusolver_dn_zheevj_batched
2041 )
2042 }
2043}
2044
2045pub fn gesvdj_batched<T: SolverScalar>(
2047 handle: &DnHandle,
2048 jobz: EigMode,
2049 m: i32,
2050 n: i32,
2051 a: &mut DeviceBuffer<T>,
2052 lda: i32,
2053 s: &mut DeviceBuffer<T::Real>,
2054 u: &mut DeviceBuffer<T>,
2055 ldu: i32,
2056 v: &mut DeviceBuffer<T>,
2057 ldv: i32,
2058 info: &mut DeviceBuffer<i32>,
2059 params: &GesvdjInfo,
2060 batch_size: i32,
2061) -> Result<()> {
2062 use baracuda_cusolver_sys::{cuComplex, cuDoubleComplex};
2063 use core::mem;
2064
2065 let mut lwork: c_int = 0;
2066 macro_rules! dispatch_real {
2067 ($t:ty, $bufsize:ident, $solve:ident) => {{
2068 let c = cusolver()?;
2069 check(unsafe {
2070 (c.$bufsize()?)(
2071 handle.as_raw(),
2072 jobz,
2073 m,
2074 n,
2075 a.as_raw().0 as *const $t,
2076 lda,
2077 s.as_raw().0 as *const $t,
2078 u.as_raw().0 as *const $t,
2079 ldu,
2080 v.as_raw().0 as *const $t,
2081 ldv,
2082 &mut lwork,
2083 params.as_raw(),
2084 batch_size,
2085 )
2086 })?;
2087 let workspace =
2088 DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
2089 check(unsafe {
2090 (c.$solve()?)(
2091 handle.as_raw(),
2092 jobz,
2093 m,
2094 n,
2095 a.as_raw().0 as *mut $t,
2096 lda,
2097 s.as_raw().0 as *mut $t,
2098 u.as_raw().0 as *mut $t,
2099 ldu,
2100 v.as_raw().0 as *mut $t,
2101 ldv,
2102 workspace.as_raw().0 as *mut $t,
2103 lwork,
2104 info.as_raw().0 as *mut c_int,
2105 params.as_raw(),
2106 batch_size,
2107 )
2108 })
2109 }};
2110 }
2111 macro_rules! dispatch_complex {
2112 ($t:ty, $real:ty, $raw:ty, $bufsize:ident, $solve:ident) => {{
2113 let c = cusolver()?;
2114 check(unsafe {
2115 (c.$bufsize()?)(
2116 handle.as_raw(),
2117 jobz,
2118 m,
2119 n,
2120 a.as_raw().0 as *const $raw,
2121 lda,
2122 s.as_raw().0 as *const $real,
2123 u.as_raw().0 as *const $raw,
2124 ldu,
2125 v.as_raw().0 as *const $raw,
2126 ldv,
2127 &mut lwork,
2128 params.as_raw(),
2129 batch_size,
2130 )
2131 })?;
2132 let workspace =
2133 DeviceBuffer::<T>::new(a.context(), lwork as usize).map_err(alloc_fail)?;
2134 check(unsafe {
2135 (c.$solve()?)(
2136 handle.as_raw(),
2137 jobz,
2138 m,
2139 n,
2140 a.as_raw().0 as *mut $raw,
2141 lda,
2142 s.as_raw().0 as *mut $real,
2143 u.as_raw().0 as *mut $raw,
2144 ldu,
2145 v.as_raw().0 as *mut $raw,
2146 ldv,
2147 workspace.as_raw().0 as *mut $raw,
2148 lwork,
2149 info.as_raw().0 as *mut c_int,
2150 params.as_raw(),
2151 batch_size,
2152 )
2153 })
2154 }};
2155 }
2156
2157 if mem::size_of::<T>() == mem::size_of::<f32>() && mem::size_of::<T::Real>() == 4 {
2158 dispatch_real!(
2159 f32,
2160 cusolver_dn_sgesvdj_batched_buffer_size,
2161 cusolver_dn_sgesvdj_batched
2162 )
2163 } else if mem::size_of::<T>() == mem::size_of::<f64>() && mem::size_of::<T::Real>() == 8 {
2164 dispatch_real!(
2165 f64,
2166 cusolver_dn_dgesvdj_batched_buffer_size,
2167 cusolver_dn_dgesvdj_batched
2168 )
2169 } else if mem::size_of::<T>() == mem::size_of::<Complex32>() {
2170 dispatch_complex!(
2171 Complex32,
2172 f32,
2173 cuComplex,
2174 cusolver_dn_cgesvdj_batched_buffer_size,
2175 cusolver_dn_cgesvdj_batched
2176 )
2177 } else {
2178 dispatch_complex!(
2179 Complex64,
2180 f64,
2181 cuDoubleComplex,
2182 cusolver_dn_zgesvdj_batched_buffer_size,
2183 cusolver_dn_zgesvdj_batched
2184 )
2185 }
2186}
2187
2188pub mod mg {
2191 use core::ffi::{c_int, c_void};
2196
2197 use baracuda_cusolver_sys::{
2198 cudaDataType, cudaLibMgGrid_t, cudaLibMgMatrixDesc_t, cusolver_mg, cusolverMgHandle_t,
2199 };
2200
2201 use super::{EigMode, Fill, Result, alloc_fail, check};
2202
2203 #[derive(Debug)]
2205 pub struct Handle {
2206 raw: cusolverMgHandle_t,
2207 }
2208
2209 impl Handle {
2210 pub fn new() -> Result<Self> {
2213 let mg = cusolver_mg()?;
2214 let cu = mg.cusolver_mg_create()?;
2215 let mut h: cusolverMgHandle_t = core::ptr::null_mut();
2216 check(unsafe { cu(&mut h) })?;
2217 Ok(Self { raw: h })
2218 }
2219
2220 pub fn device_select(&self, devices: &[i32]) -> Result<()> {
2223 let mg = cusolver_mg()?;
2224 let cu = mg.cusolver_mg_device_select()?;
2225 check(unsafe { cu(self.raw, devices.len() as c_int, devices.as_ptr()) })
2226 }
2227
2228 pub fn as_raw(&self) -> cusolverMgHandle_t {
2230 self.raw
2231 }
2232 }
2233
2234 impl Drop for Handle {
2235 fn drop(&mut self) {
2236 if let Ok(mg) = cusolver_mg()
2237 && let Ok(cu) = mg.cusolver_mg_destroy()
2238 {
2239 let _ = unsafe { cu(self.raw) };
2240 }
2241 }
2242 }
2243
2244 #[derive(Debug)]
2246 pub struct DeviceGrid {
2247 raw: cudaLibMgGrid_t,
2248 }
2249
2250 impl DeviceGrid {
2251 pub fn new(
2255 num_row_devices: i32,
2256 num_col_devices: i32,
2257 devices: &[i32],
2258 mapping: i32,
2259 ) -> Result<Self> {
2260 let mg = cusolver_mg()?;
2261 let cu = mg.cusolver_mg_create_device_grid()?;
2262 let mut raw: cudaLibMgGrid_t = core::ptr::null_mut();
2263 check(unsafe {
2264 cu(
2265 &mut raw,
2266 num_row_devices,
2267 num_col_devices,
2268 devices.as_ptr(),
2269 mapping,
2270 )
2271 })?;
2272 Ok(Self { raw })
2273 }
2274
2275 pub fn as_raw(&self) -> cudaLibMgGrid_t {
2277 self.raw
2278 }
2279 }
2280
2281 impl Drop for DeviceGrid {
2282 fn drop(&mut self) {
2283 if let Ok(mg) = cusolver_mg()
2284 && let Ok(cu) = mg.cusolver_mg_destroy_grid()
2285 {
2286 let _ = unsafe { cu(self.raw) };
2287 }
2288 }
2289 }
2290
2291 #[derive(Debug)]
2293 pub struct MatrixDesc {
2294 raw: cudaLibMgMatrixDesc_t,
2295 }
2296
2297 impl MatrixDesc {
2298 pub fn new(
2301 num_rows: i64,
2302 num_cols: i64,
2303 row_block_size: i64,
2304 col_block_size: i64,
2305 data_type: cudaDataType,
2306 grid: &DeviceGrid,
2307 ) -> Result<Self> {
2308 let mg = cusolver_mg()?;
2309 let cu = mg.cusolver_mg_create_matrix_desc()?;
2310 let mut raw: cudaLibMgMatrixDesc_t = core::ptr::null_mut();
2311 check(unsafe {
2312 cu(
2313 &mut raw,
2314 num_rows,
2315 num_cols,
2316 row_block_size,
2317 col_block_size,
2318 data_type,
2319 grid.as_raw(),
2320 )
2321 })?;
2322 Ok(Self { raw })
2323 }
2324
2325 pub fn as_raw(&self) -> cudaLibMgMatrixDesc_t {
2327 self.raw
2328 }
2329 }
2330
2331 impl Drop for MatrixDesc {
2332 fn drop(&mut self) {
2333 if let Ok(mg) = cusolver_mg()
2334 && let Ok(cu) = mg.cusolver_mg_destroy_matrix_desc()
2335 {
2336 let _ = unsafe { cu(self.raw) };
2337 }
2338 }
2339 }
2340
2341 pub unsafe fn getrf_buffer_size(
2347 handle: &Handle,
2348 m: i32,
2349 n: i32,
2350 array_d_a: *mut *mut c_void,
2351 ia: i32,
2352 ja: i32,
2353 desc_a: &MatrixDesc,
2354 array_d_ipiv: *mut *mut c_int,
2355 compute_type: cudaDataType,
2356 ) -> Result<i64> {
2357 unsafe {
2358 let mg = cusolver_mg()?;
2359 let cu = mg.cusolver_mg_getrf_buffer_size()?;
2360 let mut lwork: i64 = 0;
2361 check(cu(
2362 handle.as_raw(),
2363 m,
2364 n,
2365 array_d_a,
2366 ia,
2367 ja,
2368 desc_a.as_raw(),
2369 array_d_ipiv,
2370 compute_type,
2371 &mut lwork,
2372 ))?;
2373 Ok(lwork)
2374 }
2375 }
2376
2377 pub unsafe fn getrf(
2380 handle: &Handle,
2381 m: i32,
2382 n: i32,
2383 array_d_a: *mut *mut c_void,
2384 ia: i32,
2385 ja: i32,
2386 desc_a: &MatrixDesc,
2387 array_d_ipiv: *mut *mut c_int,
2388 compute_type: cudaDataType,
2389 array_d_work: *mut *mut c_void,
2390 lwork: i64,
2391 info: &mut [c_int],
2392 ) -> Result<()> {
2393 unsafe {
2394 let mg = cusolver_mg()?;
2395 let cu = mg.cusolver_mg_getrf()?;
2396 let _ = alloc_fail::<()>; check(cu(
2398 handle.as_raw(),
2399 m,
2400 n,
2401 array_d_a,
2402 ia,
2403 ja,
2404 desc_a.as_raw(),
2405 array_d_ipiv,
2406 compute_type,
2407 array_d_work,
2408 lwork,
2409 info.as_mut_ptr(),
2410 ))
2411 }
2412 }
2413
2414 pub unsafe fn potrf_buffer_size(
2419 handle: &Handle,
2420 uplo: Fill,
2421 n: i32,
2422 array_d_a: *mut *mut c_void,
2423 ia: i32,
2424 ja: i32,
2425 desc_a: &MatrixDesc,
2426 compute_type: cudaDataType,
2427 ) -> Result<i64> {
2428 unsafe {
2429 let mg = cusolver_mg()?;
2430 let cu = mg.cusolver_mg_potrf_buffer_size()?;
2431 let mut lwork: i64 = 0;
2432 check(cu(
2433 handle.as_raw(),
2434 uplo,
2435 n,
2436 array_d_a,
2437 ia,
2438 ja,
2439 desc_a.as_raw(),
2440 compute_type,
2441 &mut lwork,
2442 ))?;
2443 Ok(lwork)
2444 }
2445 }
2446
2447 pub unsafe fn potrf(
2450 handle: &Handle,
2451 uplo: Fill,
2452 n: i32,
2453 array_d_a: *mut *mut c_void,
2454 ia: i32,
2455 ja: i32,
2456 desc_a: &MatrixDesc,
2457 compute_type: cudaDataType,
2458 array_d_work: *mut *mut c_void,
2459 lwork: i64,
2460 info: &mut [c_int],
2461 ) -> Result<()> {
2462 unsafe {
2463 let mg = cusolver_mg()?;
2464 let cu = mg.cusolver_mg_potrf()?;
2465 check(cu(
2466 handle.as_raw(),
2467 uplo,
2468 n,
2469 array_d_a,
2470 ia,
2471 ja,
2472 desc_a.as_raw(),
2473 compute_type,
2474 array_d_work,
2475 lwork,
2476 info.as_mut_ptr(),
2477 ))
2478 }
2479 }
2480
2481 pub unsafe fn syevd_buffer_size(
2486 handle: &Handle,
2487 jobz: EigMode,
2488 uplo: Fill,
2489 n: i32,
2490 array_d_a: *mut *mut c_void,
2491 ia: i32,
2492 ja: i32,
2493 desc_a: &MatrixDesc,
2494 w: *mut c_void,
2495 data_type_w: cudaDataType,
2496 compute_type: cudaDataType,
2497 ) -> Result<i64> {
2498 unsafe {
2499 let mg = cusolver_mg()?;
2500 let cu = mg.cusolver_mg_syevd_buffer_size()?;
2501 let mut lwork: i64 = 0;
2502 check(cu(
2503 handle.as_raw(),
2504 jobz,
2505 uplo,
2506 n,
2507 array_d_a,
2508 ia,
2509 ja,
2510 desc_a.as_raw(),
2511 w,
2512 data_type_w,
2513 compute_type,
2514 &mut lwork,
2515 ))?;
2516 Ok(lwork)
2517 }
2518 }
2519
2520 pub unsafe fn syevd(
2523 handle: &Handle,
2524 jobz: EigMode,
2525 uplo: Fill,
2526 n: i32,
2527 array_d_a: *mut *mut c_void,
2528 ia: i32,
2529 ja: i32,
2530 desc_a: &MatrixDesc,
2531 w: *mut c_void,
2532 data_type_w: cudaDataType,
2533 compute_type: cudaDataType,
2534 array_d_work: *mut *mut c_void,
2535 lwork: i64,
2536 info: &mut [c_int],
2537 ) -> Result<()> {
2538 unsafe {
2539 let mg = cusolver_mg()?;
2540 let cu = mg.cusolver_mg_syevd()?;
2541 check(cu(
2542 handle.as_raw(),
2543 jobz,
2544 uplo,
2545 n,
2546 array_d_a,
2547 ia,
2548 ja,
2549 desc_a.as_raw(),
2550 w,
2551 data_type_w,
2552 compute_type,
2553 array_d_work,
2554 lwork,
2555 info.as_mut_ptr(),
2556 ))
2557 }
2558 }
2559}
2560
2561pub fn sgetrf(
2565 handle: &DnHandle,
2566 m: i32,
2567 n: i32,
2568 a: &mut DeviceBuffer<f32>,
2569 lda: i32,
2570 ipiv: &mut DeviceBuffer<i32>,
2571 info: &mut DeviceBuffer<i32>,
2572) -> Result<()> {
2573 getrf::<f32>(handle, m, n, a, lda, ipiv, info)
2574}
2575
2576pub fn sgetrs(
2578 handle: &DnHandle,
2579 trans: Op,
2580 n: i32,
2581 nrhs: i32,
2582 a: &DeviceBuffer<f32>,
2583 lda: i32,
2584 ipiv: &DeviceBuffer<i32>,
2585 b: &mut DeviceBuffer<f32>,
2586 ldb: i32,
2587 info: &mut DeviceBuffer<i32>,
2588) -> Result<()> {
2589 getrs::<f32>(handle, trans, n, nrhs, a, lda, ipiv, b, ldb, info)
2590}
2591
2592pub mod xapi {
2595 use super::*;
2601 use baracuda_cusolver_sys::{cudaDataType, cusolverDnParams_t};
2602
2603 #[derive(Debug)]
2606 pub struct Params {
2607 raw: cusolverDnParams_t,
2608 }
2609
2610 impl Params {
2611 pub fn new() -> Result<Self> {
2614 let c = cusolver()?;
2615 let cu = c.cusolver_dn_create_params()?;
2616 let mut p: cusolverDnParams_t = core::ptr::null_mut();
2617 check(unsafe { cu(&mut p) })?;
2618 Ok(Self { raw: p })
2619 }
2620
2621 pub fn as_raw(&self) -> cusolverDnParams_t {
2623 self.raw
2624 }
2625 }
2626
2627 impl Drop for Params {
2628 fn drop(&mut self) {
2629 if let Ok(c) = cusolver()
2630 && let Ok(cu) = c.cusolver_dn_destroy_params()
2631 {
2632 let _ = unsafe { cu(self.raw) };
2633 }
2634 }
2635 }
2636
2637 pub unsafe fn xgetrf_buffer_size(
2649 handle: &DnHandle,
2650 params: &Params,
2651 m: i64,
2652 n: i64,
2653 data_type_a: cudaDataType,
2654 a: *const c_void,
2655 lda: i64,
2656 compute_type: cudaDataType,
2657 ) -> Result<(usize, usize)> {
2658 let c = cusolver()?;
2659 let cu = c.cusolver_dn_xgetrf_buffer_size()?;
2660 let (mut dev, mut host) = (0usize, 0usize);
2661 check(unsafe {
2662 cu(
2663 handle.as_raw(),
2664 params.raw,
2665 m,
2666 n,
2667 data_type_a,
2668 a,
2669 lda,
2670 compute_type,
2671 &mut dev,
2672 &mut host,
2673 )
2674 })?;
2675 Ok((dev, host))
2676 }
2677
2678 pub unsafe fn xgetrf(
2687 handle: &DnHandle,
2688 params: &Params,
2689 m: i64,
2690 n: i64,
2691 data_type_a: cudaDataType,
2692 a: *mut c_void,
2693 lda: i64,
2694 ipiv: *mut i64,
2695 compute_type: cudaDataType,
2696 device_buf: *mut c_void,
2697 device_bytes: usize,
2698 host_buf: *mut c_void,
2699 host_bytes: usize,
2700 info: *mut c_int,
2701 ) -> Result<()> {
2702 unsafe {
2703 let c = cusolver()?;
2704 let cu = c.cusolver_dn_xgetrf()?;
2705 check(cu(
2706 handle.as_raw(),
2707 params.raw,
2708 m,
2709 n,
2710 data_type_a,
2711 a,
2712 lda,
2713 ipiv,
2714 compute_type,
2715 device_buf,
2716 device_bytes,
2717 host_buf,
2718 host_bytes,
2719 info,
2720 ))
2721 }
2722 }
2723
2724 pub unsafe fn xgetrs(
2733 handle: &DnHandle,
2734 params: &Params,
2735 trans: Op,
2736 n: i64,
2737 nrhs: i64,
2738 data_type_a: cudaDataType,
2739 a: *const c_void,
2740 lda: i64,
2741 ipiv: *const i64,
2742 data_type_b: cudaDataType,
2743 b: *mut c_void,
2744 ldb: i64,
2745 info: *mut c_int,
2746 ) -> Result<()> {
2747 unsafe {
2748 let c = cusolver()?;
2749 let cu = c.cusolver_dn_xgetrs()?;
2750 check(cu(
2751 handle.as_raw(),
2752 params.raw,
2753 trans.raw(),
2754 n,
2755 nrhs,
2756 data_type_a,
2757 a,
2758 lda,
2759 ipiv,
2760 data_type_b,
2761 b,
2762 ldb,
2763 info,
2764 ))
2765 }
2766 }
2767
2768 pub unsafe fn xpotrs(
2776 handle: &DnHandle,
2777 params: &Params,
2778 uplo: Fill,
2779 n: i64,
2780 nrhs: i64,
2781 data_type_a: cudaDataType,
2782 a: *const c_void,
2783 lda: i64,
2784 data_type_b: cudaDataType,
2785 b: *mut c_void,
2786 ldb: i64,
2787 info: *mut c_int,
2788 ) -> Result<()> {
2789 unsafe {
2790 let c = cusolver()?;
2791 let cu = c.cusolver_dn_xpotrs()?;
2792 check(cu(
2793 handle.as_raw(),
2794 params.raw,
2795 uplo,
2796 n,
2797 nrhs,
2798 data_type_a,
2799 a,
2800 lda,
2801 data_type_b,
2802 b,
2803 ldb,
2804 info,
2805 ))
2806 }
2807 }
2808}
2809
2810pub mod sparse {
2813 use super::*;
2816 use baracuda_cusolver_sys::cusolverSpHandle_t;
2817 use core::ffi::c_int;
2818
2819 #[derive(Debug)]
2821 pub struct SpHandle {
2822 raw: cusolverSpHandle_t,
2823 _not_send: PhantomData<*mut ()>,
2824 }
2825
2826 impl SpHandle {
2827 pub fn new() -> Result<Self> {
2830 let c = cusolver()?;
2831 let cu = c.cusolver_sp_create()?;
2832 let mut h: cusolverSpHandle_t = core::ptr::null_mut();
2833 check(unsafe { cu(&mut h) })?;
2834 Ok(Self {
2835 raw: h,
2836 _not_send: PhantomData,
2837 })
2838 }
2839
2840 pub fn set_stream(&self, stream: &Stream) -> Result<()> {
2843 let c = cusolver()?;
2844 let cu = c.cusolver_sp_set_stream()?;
2845 check(unsafe { cu(self.raw, stream.as_raw() as _) })
2846 }
2847
2848 pub fn as_raw(&self) -> cusolverSpHandle_t {
2850 self.raw
2851 }
2852 }
2853
2854 impl Drop for SpHandle {
2855 fn drop(&mut self) {
2856 if let Ok(c) = cusolver()
2857 && let Ok(cu) = c.cusolver_sp_destroy()
2858 {
2859 let _ = unsafe { cu(self.raw) };
2860 }
2861 }
2862 }
2863
2864 pub unsafe fn scsrlsvchol(
2871 handle: &SpHandle,
2872 m: i32,
2873 nnz: i32,
2874 descr_a: *mut c_void,
2875 csr_val: *const f32,
2876 csr_row_ptr: *const c_int,
2877 csr_col_ind: *const c_int,
2878 b: *const f32,
2879 tol: f32,
2880 reorder: i32,
2881 x: *mut f32,
2882 singularity: *mut c_int,
2883 ) -> Result<()> {
2884 unsafe {
2885 let c = cusolver()?;
2886 let cu = c.cusolver_sp_scsrlsvchol()?;
2887 check(cu(
2888 handle.raw,
2889 m,
2890 nnz,
2891 descr_a,
2892 csr_val,
2893 csr_row_ptr,
2894 csr_col_ind,
2895 b,
2896 tol,
2897 reorder,
2898 x,
2899 singularity,
2900 ))
2901 }
2902 }
2903
2904 pub unsafe fn scsrlsvqr(
2909 handle: &SpHandle,
2910 m: i32,
2911 nnz: i32,
2912 descr_a: *mut c_void,
2913 csr_val: *const f32,
2914 csr_row_ptr: *const c_int,
2915 csr_col_ind: *const c_int,
2916 b: *const f32,
2917 tol: f32,
2918 reorder: i32,
2919 x: *mut f32,
2920 singularity: *mut c_int,
2921 ) -> Result<()> {
2922 unsafe {
2923 let c = cusolver()?;
2924 let cu = c.cusolver_sp_scsrlsvqr()?;
2925 check(cu(
2926 handle.raw,
2927 m,
2928 nnz,
2929 descr_a,
2930 csr_val,
2931 csr_row_ptr,
2932 csr_col_ind,
2933 b,
2934 tol,
2935 reorder,
2936 x,
2937 singularity,
2938 ))
2939 }
2940 }
2941}
2942
2943pub mod refactor {
2946 use super::*;
2950 use baracuda_cusolver_sys::cusolverRfHandle_t;
2951
2952 #[derive(Debug)]
2954 pub struct RfHandle {
2955 raw: cusolverRfHandle_t,
2956 _not_send: PhantomData<*mut ()>,
2957 }
2958
2959 impl RfHandle {
2960 pub fn new() -> Result<Self> {
2963 let c = cusolver()?;
2964 let cu = c.cusolver_rf_create()?;
2965 let mut h: cusolverRfHandle_t = core::ptr::null_mut();
2966 check(unsafe { cu(&mut h) })?;
2967 Ok(Self {
2968 raw: h,
2969 _not_send: PhantomData,
2970 })
2971 }
2972
2973 pub fn as_raw(&self) -> cusolverRfHandle_t {
2975 self.raw
2976 }
2977
2978 pub fn analyze(&self) -> Result<()> {
2981 let c = cusolver()?;
2982 let cu = c.cusolver_rf_analyze()?;
2983 check(unsafe { cu(self.raw) })
2984 }
2985
2986 pub fn refactor(&self) -> Result<()> {
2989 let c = cusolver()?;
2990 let cu = c.cusolver_rf_refactor()?;
2991 check(unsafe { cu(self.raw) })
2992 }
2993 }
2994
2995 impl Drop for RfHandle {
2996 fn drop(&mut self) {
2997 if let Ok(c) = cusolver()
2998 && let Ok(cu) = c.cusolver_rf_destroy()
2999 {
3000 let _ = unsafe { cu(self.raw) };
3001 }
3002 }
3003 }
3004}