tenferro_linalg/
extension.rs

1use std::any::Any;
2use std::hash::Hasher;
3use std::sync::Arc;
4
5use num_complex::{Complex32, Complex64};
6use tenferro_extension_macros::define_extension_runtime;
7use tenferro_ops::SymDim;
8use tenferro_runtime::extension::{ExtensionExecutionContext, ExtensionOp, HostReference};
9use tenferro_tensor::{
10    DType, DeviceKind, Error, GpuBackendKind, MemoryKind, Placement, Tensor, TensorRead,
11};
12
13use crate::backend::LinalgBackend;
14
15mod gauge;
16#[cfg(all(test, not(feature = "cuda")))]
17mod tests;
18
19pub(crate) use gauge::{apply_eigh_gauge, apply_qr_gauge};
20
21pub const LINALG_EXTENSION_FAMILY_ID: &str = "tenferro-linalg.linalg.v1";
22
23/// Default derivative regularization used by decomposition AD rules.
24///
25/// This epsilon is used only when differentiating decomposition formulas with
26/// repeated or nearly repeated spectral values. It is not a solver tolerance.
27///
28/// # Examples
29///
30/// ```rust
31/// use tenferro_linalg::{SvdOptions, DEFAULT_DECOMPOSITION_DERIVATIVE_EPS};
32///
33/// let options = SvdOptions::default();
34/// assert_eq!(options.derivative_eps, DEFAULT_DECOMPOSITION_DERIVATIVE_EPS);
35/// ```
36pub const DEFAULT_DECOMPOSITION_DERIVATIVE_EPS: f64 = 1e-12;
37
38/// Singular-vector gauge convention used by [`SvdOptions`].
39///
40/// # Examples
41///
42/// ```rust
43/// use tenferro_linalg::{SvdGauge, SvdOptions};
44///
45/// let options = SvdOptions::default().gauge(SvdGauge::CanonicalPivot);
46/// assert_eq!(options.gauge, SvdGauge::CanonicalPivot);
47/// ```
48#[derive(Clone, Copy, Debug, PartialEq, Eq)]
49pub enum SvdGauge {
50    /// Leave the backend's raw singular vector signs or phases unchanged.
51    Raw,
52    /// Make each left singular vector's max-absolute pivot entry positive-real
53    /// and adjust the matching `VT` row so reconstruction is preserved.
54    CanonicalPivot,
55}
56
57/// Eigenvector gauge convention used by [`EighOptions`].
58///
59/// # Examples
60///
61/// ```rust
62/// use tenferro_linalg::{EighGauge, EighOptions};
63///
64/// let options = EighOptions::default().gauge(EighGauge::CanonicalPivot);
65/// assert_eq!(options.gauge, EighGauge::CanonicalPivot);
66/// ```
67#[derive(Clone, Copy, Debug, PartialEq, Eq)]
68pub enum EighGauge {
69    /// Leave the backend's raw eigenvector signs or phases unchanged.
70    Raw,
71    /// Make each eigenvector's max-absolute pivot entry positive-real.
72    CanonicalPivot,
73}
74
75/// QR factor gauge convention used by [`QrOptions`].
76///
77/// # Examples
78///
79/// ```rust
80/// use tenferro_linalg::{QrGauge, QrOptions};
81///
82/// let options = QrOptions::default().gauge(QrGauge::PositiveDiagonal);
83/// assert_eq!(options.gauge, QrGauge::PositiveDiagonal);
84/// ```
85#[derive(Clone, Copy, Debug, PartialEq, Eq)]
86pub enum QrGauge {
87    /// Leave the backend's raw QR signs or phases unchanged.
88    Raw,
89    /// Make each `R` diagonal entry positive-real, compensating `Q`.
90    PositiveDiagonal,
91}
92
93/// Options for singular value decomposition.
94///
95/// # Examples
96///
97/// ```rust
98/// use tenferro_linalg::{SvdGauge, SvdOptions};
99///
100/// let options = SvdOptions::default()
101///     .gauge(SvdGauge::CanonicalPivot)
102///     .derivative_eps(1.0e-10);
103/// assert_eq!(options.gauge, SvdGauge::CanonicalPivot);
104/// assert_eq!(options.derivative_eps, 1.0e-10);
105/// ```
106#[derive(Clone, Copy, Debug, PartialEq)]
107pub struct SvdOptions {
108    /// Singular-vector gauge convention.
109    pub gauge: SvdGauge,
110    /// AD derivative regularization for repeated or nearly repeated singular values.
111    pub derivative_eps: f64,
112}
113
114impl Default for SvdOptions {
115    fn default() -> Self {
116        Self {
117            gauge: SvdGauge::Raw,
118            derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
119        }
120    }
121}
122
123impl SvdOptions {
124    /// Return options with the requested singular-vector gauge.
125    ///
126    /// # Examples
127    ///
128    /// ```rust
129    /// use tenferro_linalg::{SvdGauge, SvdOptions};
130    ///
131    /// let options = SvdOptions::default().gauge(SvdGauge::CanonicalPivot);
132    /// assert_eq!(options.gauge, SvdGauge::CanonicalPivot);
133    /// ```
134    pub fn gauge(mut self, gauge: SvdGauge) -> Self {
135        self.gauge = gauge;
136        self
137    }
138
139    /// Return options with an explicit derivative epsilon.
140    ///
141    /// # Examples
142    ///
143    /// ```rust
144    /// use tenferro_linalg::SvdOptions;
145    ///
146    /// let options = SvdOptions::default().derivative_eps(1.0e-9);
147    /// assert_eq!(options.derivative_eps, 1.0e-9);
148    /// ```
149    pub fn derivative_eps(mut self, derivative_eps: f64) -> Self {
150        self.derivative_eps = derivative_eps;
151        self
152    }
153}
154
155/// Options for Hermitian eigenvalue decomposition.
156///
157/// # Examples
158///
159/// ```rust
160/// use tenferro_linalg::EighOptions;
161///
162/// let options = EighOptions::default().derivative_eps(1.0e-10);
163/// assert_eq!(options.derivative_eps, 1.0e-10);
164/// ```
165#[derive(Clone, Copy, Debug, PartialEq)]
166pub struct EighOptions {
167    /// Eigenvector gauge convention.
168    pub gauge: EighGauge,
169    /// AD derivative regularization for repeated or nearly repeated eigenvalues.
170    pub derivative_eps: f64,
171}
172
173impl Default for EighOptions {
174    fn default() -> Self {
175        Self {
176            gauge: EighGauge::Raw,
177            derivative_eps: DEFAULT_DECOMPOSITION_DERIVATIVE_EPS,
178        }
179    }
180}
181
182impl EighOptions {
183    /// Return options with the requested eigenvector gauge.
184    ///
185    /// # Examples
186    ///
187    /// ```rust
188    /// use tenferro_linalg::{EighGauge, EighOptions};
189    ///
190    /// let options = EighOptions::default().gauge(EighGauge::CanonicalPivot);
191    /// assert_eq!(options.gauge, EighGauge::CanonicalPivot);
192    /// ```
193    pub fn gauge(mut self, gauge: EighGauge) -> Self {
194        self.gauge = gauge;
195        self
196    }
197
198    /// Return options with an explicit derivative epsilon.
199    ///
200    /// # Examples
201    ///
202    /// ```rust
203    /// use tenferro_linalg::EighOptions;
204    ///
205    /// let options = EighOptions::default().derivative_eps(1.0e-9);
206    /// assert_eq!(options.derivative_eps, 1.0e-9);
207    /// ```
208    pub fn derivative_eps(mut self, derivative_eps: f64) -> Self {
209        self.derivative_eps = derivative_eps;
210        self
211    }
212}
213
214/// Options for QR decomposition.
215///
216/// # Examples
217///
218/// ```rust
219/// use tenferro_linalg::{QrGauge, QrOptions};
220///
221/// let options = QrOptions::default().gauge(QrGauge::PositiveDiagonal);
222/// assert_eq!(options.gauge, QrGauge::PositiveDiagonal);
223/// ```
224#[derive(Clone, Copy, Debug, PartialEq, Eq)]
225pub struct QrOptions {
226    /// QR sign or phase convention.
227    pub gauge: QrGauge,
228}
229
230impl Default for QrOptions {
231    fn default() -> Self {
232        Self {
233            gauge: QrGauge::Raw,
234        }
235    }
236}
237
238impl QrOptions {
239    /// Return options with the requested QR gauge.
240    ///
241    /// # Examples
242    ///
243    /// ```rust
244    /// use tenferro_linalg::{QrGauge, QrOptions};
245    ///
246    /// let options = QrOptions::default().gauge(QrGauge::PositiveDiagonal);
247    /// assert_eq!(options.gauge, QrGauge::PositiveDiagonal);
248    /// ```
249    pub fn gauge(mut self, gauge: QrGauge) -> Self {
250        self.gauge = gauge;
251        self
252    }
253}
254
255pub(crate) fn validate_derivative_eps(
256    op: &'static str,
257    derivative_eps: f64,
258) -> tenferro_tensor::Result<()> {
259    if derivative_eps.is_finite() && derivative_eps > 0.0 {
260        Ok(())
261    } else {
262        Err(Error::InvalidConfig {
263            op,
264            message: format!("derivative_eps must be positive and finite, got {derivative_eps}"),
265        })
266    }
267}
268
269#[derive(Clone, Copy, Debug, PartialEq)]
270#[doc(hidden)]
271pub(crate) enum LinalgOp {
272    Cholesky,
273    Lu,
274    LuFactor,
275    LuSolvePrepared {
276        transpose_a: bool,
277        conjugate_a: bool,
278    },
279    FullPivLu,
280    FullPivLuSolve {
281        transpose_a: bool,
282    },
283    Svd {
284        derivative_eps: f64,
285        gauge: SvdGauge,
286    },
287    SvdVals {
288        derivative_eps: f64,
289    },
290    Qr {
291        gauge: QrGauge,
292    },
293    Eigh {
294        derivative_eps: f64,
295        gauge: EighGauge,
296    },
297    EighVals {
298        derivative_eps: f64,
299    },
300    Eig {
301        input_dtype: DType,
302    },
303    EigVals {
304        input_dtype: DType,
305    },
306    TriangularSolve {
307        left_side: bool,
308        lower: bool,
309        transpose_a: bool,
310        unit_diagonal: bool,
311    },
312}
313
314impl LinalgOp {
315    fn output_count(self) -> usize {
316        match self {
317            Self::Cholesky
318            | Self::EighVals { .. }
319            | Self::EigVals { .. }
320            | Self::FullPivLuSolve { .. }
321            | Self::LuSolvePrepared { .. }
322            | Self::SvdVals { .. }
323            | Self::TriangularSolve { .. } => 1,
324            Self::Svd { .. } => 3,
325            Self::Qr { .. } | Self::Eigh { .. } | Self::Eig { .. } => 2,
326            Self::LuFactor => 3,
327            Self::Lu => 4,
328            Self::FullPivLu => 5,
329        }
330    }
331
332    fn input_count(self) -> usize {
333        match self {
334            Self::FullPivLuSolve { .. } | Self::TriangularSolve { .. } => 2,
335            Self::LuSolvePrepared { .. } => 4,
336            _ => 1,
337        }
338    }
339
340    fn tag(self) -> u8 {
341        match self {
342            Self::Cholesky => 0,
343            Self::Lu => 1,
344            Self::FullPivLu => 2,
345            Self::FullPivLuSolve { .. } => 3,
346            Self::Svd { .. } => 4,
347            Self::Qr { .. } => 5,
348            Self::Eigh { .. } => 6,
349            Self::Eig { .. } => 7,
350            Self::TriangularSolve { .. } => 9,
351            Self::LuFactor => 10,
352            Self::LuSolvePrepared { .. } => 11,
353            Self::SvdVals { .. } => 12,
354            Self::EighVals { .. } => 13,
355            Self::EigVals { .. } => 14,
356        }
357    }
358}
359
360#[derive(Clone, Debug, PartialEq)]
361#[doc(hidden)]
362pub(crate) struct LinalgExtensionOp {
363    op: LinalgOp,
364}
365
366impl LinalgExtensionOp {
367    pub(crate) fn new(op: LinalgOp) -> Self {
368        Self { op }
369    }
370
371    pub(crate) fn op(&self) -> LinalgOp {
372        self.op
373    }
374}
375
376#[derive(Clone, Copy, Debug, PartialEq, Eq)]
377enum EagerLinalgDevice {
378    Cpu,
379    Cuda(usize),
380}
381
382fn tensor_placement(input: &Tensor) -> &Placement {
383    input.placement()
384}
385
386fn input_eager_device(input: &Tensor) -> tenferro_tensor::Result<EagerLinalgDevice> {
387    let placement = tensor_placement(input);
388    match (&placement.memory_kind, placement.device.as_ref()) {
389        (MemoryKind::Device, Some(device)) => match &device.kind {
390            DeviceKind::Gpu(GpuBackendKind::Cuda) => Ok(EagerLinalgDevice::Cuda(device.ordinal)),
391            DeviceKind::Gpu(kind) => Err(Error::backend_failure(
392                "linalg_host_reference",
393                format!("unsupported GPU backend {kind:?} for eager linalg"),
394            )),
395            kind => Err(Error::backend_failure(
396                "linalg_host_reference",
397                format!("unsupported device kind {kind:?} for eager linalg"),
398            )),
399        },
400        (MemoryKind::Device, None) => Err(Error::backend_failure(
401            "linalg_host_reference",
402            "device tensor is missing placement device metadata",
403        )),
404        _ => Ok(EagerLinalgDevice::Cpu),
405    }
406}
407
408fn eager_linalg_device(inputs: &[&Tensor]) -> tenferro_tensor::Result<EagerLinalgDevice> {
409    let mut selected = None;
410    for input in inputs {
411        let device = input_eager_device(input)?;
412        match (selected, device) {
413            (None, next) => selected = Some(next),
414            (Some(EagerLinalgDevice::Cpu), EagerLinalgDevice::Cpu) => {}
415            (Some(EagerLinalgDevice::Cuda(lhs)), EagerLinalgDevice::Cuda(rhs)) if lhs == rhs => {}
416            (Some(lhs), rhs) => {
417                return Err(Error::backend_failure(
418                    "linalg_host_reference",
419                    format!("all eager linalg inputs must be on the same device, got {lhs:?} and {rhs:?}"),
420                ));
421            }
422        }
423    }
424    Ok(selected.unwrap_or(EagerLinalgDevice::Cpu))
425}
426
427#[cfg(feature = "cuda")]
428fn execute_cuda_eager_linalg(
429    op: LinalgOp,
430    inputs: &[&Tensor],
431    device_ordinal: usize,
432) -> tenferro_tensor::Result<Vec<Tensor>> {
433    let mut backend = tenferro_gpu::CudaBackend::new(device_ordinal)?;
434    execute_linalg(op, inputs, &mut backend)
435}
436
437#[cfg(not(feature = "cuda"))]
438fn execute_cuda_eager_linalg(
439    _op: LinalgOp,
440    _inputs: &[&Tensor],
441    device_ordinal: usize,
442) -> tenferro_tensor::Result<Vec<Tensor>> {
443    Err(Error::backend_failure(
444        "linalg_host_reference",
445        format!(
446            "received CUDA tensor on cuda:{device_ordinal}, but tenferro-linalg was built \
447             without the cuda feature; enable the cuda feature or download the tensor to CPU \
448             before eager linalg"
449        ),
450    ))
451}
452
453impl ExtensionOp for LinalgExtensionOp {
454    fn family_id(&self) -> &'static str {
455        LINALG_EXTENSION_FAMILY_ID
456    }
457
458    fn payload_hash(&self, hasher: &mut dyn Hasher) {
459        hasher.write_u8(self.op.tag());
460        match self.op {
461            LinalgOp::Svd {
462                derivative_eps,
463                gauge,
464            } => {
465                hasher.write_u64(derivative_eps.to_bits());
466                hash_svd_gauge(hasher, gauge);
467            }
468            LinalgOp::SvdVals { derivative_eps } | LinalgOp::EighVals { derivative_eps } => {
469                hasher.write_u64(derivative_eps.to_bits());
470            }
471            LinalgOp::Qr { gauge } => {
472                hash_qr_gauge(hasher, gauge);
473            }
474            LinalgOp::Eigh {
475                derivative_eps,
476                gauge,
477            } => {
478                hasher.write_u64(derivative_eps.to_bits());
479                hash_eigh_gauge(hasher, gauge);
480            }
481            LinalgOp::Eig { input_dtype } | LinalgOp::EigVals { input_dtype } => {
482                hash_dtype(hasher, input_dtype);
483            }
484            LinalgOp::FullPivLuSolve { transpose_a } => {
485                hasher.write_u8(u8::from(transpose_a));
486            }
487            LinalgOp::LuSolvePrepared {
488                transpose_a,
489                conjugate_a,
490            } => {
491                hasher.write_u8(u8::from(transpose_a));
492                hasher.write_u8(u8::from(conjugate_a));
493            }
494            LinalgOp::TriangularSolve {
495                left_side,
496                lower,
497                transpose_a,
498                unit_diagonal,
499            } => {
500                hasher.write_u8(u8::from(left_side));
501                hasher.write_u8(u8::from(lower));
502                hasher.write_u8(u8::from(transpose_a));
503                hasher.write_u8(u8::from(unit_diagonal));
504            }
505            LinalgOp::Cholesky | LinalgOp::Lu | LinalgOp::LuFactor | LinalgOp::FullPivLu => {}
506        }
507    }
508
509    fn payload_eq(&self, other: &dyn ExtensionOp) -> bool {
510        other
511            .as_any()
512            .downcast_ref::<Self>()
513            .is_some_and(|that| self == that)
514    }
515
516    fn clone_arc(&self) -> Arc<dyn ExtensionOp> {
517        Arc::new(self.clone())
518    }
519
520    fn as_any(&self) -> &dyn Any {
521        self
522    }
523
524    fn input_count(&self) -> usize {
525        self.op.input_count()
526    }
527
528    fn output_count(&self) -> usize {
529        self.op.output_count()
530    }
531
532    fn prune_outputs(&self, live_outputs: &[bool]) -> Option<Arc<dyn ExtensionOp>> {
533        match self.op {
534            LinalgOp::Svd { derivative_eps, .. } if live_outputs == [false, true, false] => {
535                Some(Arc::new(Self::new(LinalgOp::SvdVals { derivative_eps })))
536            }
537            LinalgOp::Eigh { derivative_eps, .. } if live_outputs == [true, false] => {
538                Some(Arc::new(Self::new(LinalgOp::EighVals { derivative_eps })))
539            }
540            LinalgOp::Eig { input_dtype } if live_outputs == [true, false] => {
541                Some(Arc::new(Self::new(LinalgOp::EigVals { input_dtype })))
542            }
543            _ => None,
544        }
545    }
546
547    fn infer_output_meta(
548        &self,
549        input_dtypes: &[DType],
550        input_shapes: &[&[SymDim]],
551    ) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
552        if input_dtypes.len() != self.input_count() || input_shapes.len() != self.input_count() {
553            return Err(Error::InvalidConfig {
554                op: "tenferro-linalg",
555                message: format!(
556                    "expected {} input metadata entries, got dtypes={} shapes={}",
557                    self.input_count(),
558                    input_dtypes.len(),
559                    input_shapes.len()
560                ),
561            });
562        }
563        let metas = match self.op {
564            LinalgOp::Cholesky => {
565                require_matrix_meta("tenferro-linalg.cholesky", input_shapes[0])?;
566                vec![(promote_dtypes(input_dtypes), input_shapes[0].to_vec())]
567            }
568            LinalgOp::FullPivLuSolve { .. } => {
569                require_matrix_meta("tenferro-linalg.full_piv_lu_solve", input_shapes[0])?;
570                require_matrix_meta("tenferro-linalg.full_piv_lu_solve", input_shapes[1])?;
571                vec![(promote_dtypes(input_dtypes), input_shapes[1].to_vec())]
572            }
573            LinalgOp::TriangularSolve { .. } => {
574                require_matrix_meta("tenferro-linalg.triangular_solve", input_shapes[0])?;
575                require_matrix_meta("tenferro-linalg.triangular_solve", input_shapes[1])?;
576                vec![(promote_dtypes(input_dtypes), input_shapes[1].to_vec())]
577            }
578            LinalgOp::LuSolvePrepared { .. } => {
579                require_matrix_meta("tenferro-linalg.lu_solve_prepared_lu", input_shapes[0])?;
580                require_matrix_meta("tenferro-linalg.lu_solve_prepared_rhs", input_shapes[3])?;
581                vec![(
582                    promote_dtypes(&[input_dtypes[0], input_dtypes[3]]),
583                    input_shapes[3].to_vec(),
584                )]
585            }
586            LinalgOp::Lu => lu_meta(input_dtypes[0], input_shapes[0])?,
587            LinalgOp::LuFactor => lu_factor_meta(input_dtypes[0], input_shapes[0])?,
588            LinalgOp::FullPivLu => full_piv_lu_meta(input_dtypes[0], input_shapes[0])?,
589            LinalgOp::Svd { .. } => svd_meta(input_dtypes[0], input_shapes[0])?,
590            LinalgOp::SvdVals { .. } => {
591                vec![svd_values_meta(input_dtypes[0], input_shapes[0])?]
592            }
593            LinalgOp::Qr { .. } => qr_meta(input_dtypes[0], input_shapes[0])?,
594            LinalgOp::Eigh { .. } => eigh_meta(input_dtypes[0], input_shapes[0])?,
595            LinalgOp::EighVals { .. } => vec![eigh_values_meta(input_dtypes[0], input_shapes[0])?],
596            LinalgOp::Eig { input_dtype } => eig_meta(input_dtype, input_shapes[0])?,
597            LinalgOp::EigVals { input_dtype } => {
598                vec![eig_values_meta(input_dtype, input_shapes[0])?]
599            }
600        };
601        Ok(metas)
602    }
603
604    fn host_reference(&self) -> Option<&dyn HostReference> {
605        Some(self)
606    }
607}
608
609impl HostReference for LinalgExtensionOp {
610    fn execute(&self, inputs: &[&Tensor]) -> tenferro_tensor::Result<Vec<Tensor>> {
611        let expected = self.input_count();
612        if inputs.len() != expected {
613            return Err(Error::InvalidConfig {
614                op: "linalg_host_reference",
615                message: format!(
616                    "expected {expected} inputs for {:?}, got {}",
617                    self.op,
618                    inputs.len()
619                ),
620            });
621        }
622
623        match eager_linalg_device(inputs)? {
624            EagerLinalgDevice::Cpu => {
625                let mut backend = tenferro_cpu::CpuBackend::new();
626                execute_linalg(self.op, inputs, &mut backend)
627            }
628            EagerLinalgDevice::Cuda(device_ordinal) => {
629                execute_cuda_eager_linalg(self.op, inputs, device_ordinal)
630            }
631        }
632    }
633}
634
635fn execute_linalg_extension<B: LinalgBackend + 'static>(
636    op: &LinalgExtensionOp,
637    inputs: &[&Tensor],
638    ctx: &mut ExtensionExecutionContext<'_, B>,
639) -> tenferro_tensor::Result<Vec<Tensor>> {
640    execute_linalg(op.op(), inputs, ctx.backend_mut())
641}
642
643fn execute_linalg_extension_reads<B: LinalgBackend + 'static>(
644    op: &LinalgExtensionOp,
645    inputs: &[TensorRead<'_>],
646    ctx: &mut ExtensionExecutionContext<'_, B>,
647) -> tenferro_tensor::Result<Vec<Tensor>> {
648    // Linalg backends currently operate on compact tensors; materialization is
649    // explicit here so borrowed views cannot silently bypass backend errors.
650    let materialized_inputs: Vec<Tensor> = inputs
651        .iter()
652        .map(TensorRead::to_tensor)
653        .collect::<tenferro_tensor::Result<_>>()?;
654    let input_refs: Vec<&Tensor> = materialized_inputs.iter().collect();
655    execute_linalg_extension(op, &input_refs, ctx)
656}
657
658define_extension_runtime! {
659    runtime = LinalgRuntime,
660    family_id = LINALG_EXTENSION_FAMILY_ID,
661    op_type = LinalgExtensionOp,
662    execute = execute_linalg_extension,
663    execute_reads = execute_linalg_extension_reads,
664    register_fn = register_runtime,
665    backend_bound = LinalgBackend,
666}
667
668fn execute_linalg<B: LinalgBackend>(
669    op: LinalgOp,
670    inputs: &[&Tensor],
671    backend: &mut B,
672) -> tenferro_tensor::Result<Vec<Tensor>> {
673    match op {
674        LinalgOp::Cholesky => Ok(vec![backend.cholesky(inputs[0])?]),
675        LinalgOp::Lu => backend.lu(inputs[0]),
676        LinalgOp::LuFactor => backend.lu_factor(inputs[0]),
677        LinalgOp::LuSolvePrepared {
678            transpose_a,
679            conjugate_a,
680        } => Ok(vec![backend.lu_solve_prepared(
681            inputs[0],
682            inputs[1],
683            inputs[2],
684            inputs[3],
685            transpose_a,
686            conjugate_a,
687        )?]),
688        LinalgOp::FullPivLu => backend.full_piv_lu(inputs[0]),
689        LinalgOp::FullPivLuSolve { transpose_a } => Ok(vec![backend.full_piv_lu_solve(
690            inputs[0],
691            inputs[1],
692            transpose_a,
693        )?]),
694        LinalgOp::Svd {
695            derivative_eps,
696            gauge,
697        } => backend.svd_with_options(
698            inputs[0],
699            SvdOptions {
700                derivative_eps,
701                gauge,
702            },
703        ),
704        LinalgOp::SvdVals { .. } => Ok(vec![backend.svd_values(inputs[0])?]),
705        LinalgOp::Qr { gauge } => backend.qr_with_options(inputs[0], QrOptions { gauge }),
706        LinalgOp::Eigh {
707            derivative_eps,
708            gauge,
709        } => backend.eigh_with_options(
710            inputs[0],
711            EighOptions {
712                derivative_eps,
713                gauge,
714            },
715        ),
716        LinalgOp::EighVals { .. } => Ok(vec![backend.eigh_values(inputs[0])?]),
717        LinalgOp::Eig { .. } => backend.eig(inputs[0]),
718        LinalgOp::EigVals { .. } => Ok(vec![backend.eig_values(inputs[0])?]),
719        LinalgOp::TriangularSolve {
720            left_side,
721            lower,
722            transpose_a,
723            unit_diagonal,
724        } => Ok(vec![backend.triangular_solve(
725            inputs[0],
726            inputs[1],
727            left_side,
728            lower,
729            transpose_a,
730            unit_diagonal,
731        )?]),
732    }
733}
734
735pub(crate) fn apply_svd_gauge(
736    gauge: SvdGauge,
737    outputs: &mut [Tensor],
738) -> tenferro_tensor::Result<()> {
739    match gauge {
740        SvdGauge::Raw => Ok(()),
741        SvdGauge::CanonicalPivot => apply_canonical_pivot_svd_gauge(outputs),
742    }
743}
744
745fn apply_canonical_pivot_svd_gauge(outputs: &mut [Tensor]) -> tenferro_tensor::Result<()> {
746    if outputs.len() != 3 {
747        return Err(Error::InvalidConfig {
748            op: "tenferro-linalg.svd",
749            message: format!(
750                "canonical SVD gauge expected three outputs, got {}",
751                outputs.len()
752            ),
753        });
754    }
755
756    let (u_slice, rest) = outputs.split_at_mut(1);
757    let (singular_slice, vt_slice) = rest.split_at_mut(1);
758    let u = &mut u_slice[0];
759    let singular_values = &singular_slice[0];
760    let vt = &mut vt_slice[0];
761    let u_shape = u.shape().to_vec();
762    let s_shape = singular_values.shape().to_vec();
763    let vt_shape = vt.shape().to_vec();
764    if u_shape.len() < 2 || vt_shape.len() < 2 || s_shape.is_empty() {
765        return Err(Error::InvalidConfig {
766            op: "tenferro-linalg.svd",
767            message: format!(
768                "canonical SVD gauge expected U rank >= 2, S rank >= 1, VT rank >= 2; got U={u_shape:?}, S={s_shape:?}, VT={vt_shape:?}"
769            ),
770        });
771    }
772
773    let m = u_shape[0];
774    let k = u_shape[1];
775    let n = vt_shape[1];
776    if s_shape[0] != k
777        || vt_shape[0] != k
778        || u_shape[2..] != vt_shape[2..]
779        || s_shape[1..] != u_shape[2..]
780    {
781        return Err(Error::InvalidConfig {
782            op: "tenferro-linalg.svd",
783            message: format!(
784                "canonical SVD gauge expected compatible compact SVD shapes, got U={u_shape:?}, S={s_shape:?}, VT={vt_shape:?}"
785            ),
786        });
787    }
788    let batch_count = u_shape[2..].iter().product::<usize>();
789
790    match (u, vt) {
791        (Tensor::F64(u), Tensor::F64(vt)) => canonicalize_svd_gauge_f64(
792            u.host_data_mut()?,
793            vt.host_data_mut()?,
794            m,
795            k,
796            n,
797            batch_count,
798        ),
799        (Tensor::F32(u), Tensor::F32(vt)) => canonicalize_svd_gauge_f32(
800            u.host_data_mut()?,
801            vt.host_data_mut()?,
802            m,
803            k,
804            n,
805            batch_count,
806        ),
807        (Tensor::C64(u), Tensor::C64(vt)) => canonicalize_svd_gauge_c64(
808            u.host_data_mut()?,
809            vt.host_data_mut()?,
810            m,
811            k,
812            n,
813            batch_count,
814        ),
815        (Tensor::C32(u), Tensor::C32(vt)) => canonicalize_svd_gauge_c32(
816            u.host_data_mut()?,
817            vt.host_data_mut()?,
818            m,
819            k,
820            n,
821            batch_count,
822        ),
823        (u, vt) => Err(Error::DTypeMismatch {
824            op: "tenferro-linalg.svd",
825            lhs: u.dtype(),
826            rhs: vt.dtype(),
827        }),
828    }
829}
830
831fn canonicalize_svd_gauge_f64(
832    u: &mut [f64],
833    vt: &mut [f64],
834    m: usize,
835    k: usize,
836    n: usize,
837    batch_count: usize,
838) -> tenferro_tensor::Result<()> {
839    for batch in 0..batch_count {
840        let u_batch = batch * m * k;
841        let vt_batch = batch * k * n;
842        for col in 0..k {
843            let pivot = max_abs_pivot_f64(u, u_batch, m, col);
844            let pivot_value = u[u_batch + pivot + m * col];
845            if pivot_value < 0.0 {
846                for row in 0..m {
847                    let offset = u_batch + row + m * col;
848                    u[offset] = -u[offset];
849                }
850                for vt_col in 0..n {
851                    let offset = vt_batch + col + k * vt_col;
852                    vt[offset] = -vt[offset];
853                }
854            }
855        }
856    }
857    Ok(())
858}
859
860fn canonicalize_svd_gauge_f32(
861    u: &mut [f32],
862    vt: &mut [f32],
863    m: usize,
864    k: usize,
865    n: usize,
866    batch_count: usize,
867) -> tenferro_tensor::Result<()> {
868    for batch in 0..batch_count {
869        let u_batch = batch * m * k;
870        let vt_batch = batch * k * n;
871        for col in 0..k {
872            let pivot = max_abs_pivot_f32(u, u_batch, m, col);
873            let pivot_value = u[u_batch + pivot + m * col];
874            if pivot_value < 0.0 {
875                for row in 0..m {
876                    let offset = u_batch + row + m * col;
877                    u[offset] = -u[offset];
878                }
879                for vt_col in 0..n {
880                    let offset = vt_batch + col + k * vt_col;
881                    vt[offset] = -vt[offset];
882                }
883            }
884        }
885    }
886    Ok(())
887}
888
889fn canonicalize_svd_gauge_c64(
890    u: &mut [Complex64],
891    vt: &mut [Complex64],
892    m: usize,
893    k: usize,
894    n: usize,
895    batch_count: usize,
896) -> tenferro_tensor::Result<()> {
897    for batch in 0..batch_count {
898        let u_batch = batch * m * k;
899        let vt_batch = batch * k * n;
900        for col in 0..k {
901            let pivot = max_abs_pivot_c64(u, u_batch, m, col);
902            let pivot_value = u[u_batch + pivot + m * col];
903            let pivot_norm = pivot_value.norm();
904            if pivot_norm == 0.0 {
905                continue;
906            }
907            let phase = pivot_value.conj() / pivot_norm;
908            let vt_phase = phase.conj();
909            for row in 0..m {
910                let offset = u_batch + row + m * col;
911                u[offset] *= phase;
912            }
913            for vt_col in 0..n {
914                let offset = vt_batch + col + k * vt_col;
915                vt[offset] *= vt_phase;
916            }
917        }
918    }
919    Ok(())
920}
921
922fn canonicalize_svd_gauge_c32(
923    u: &mut [Complex32],
924    vt: &mut [Complex32],
925    m: usize,
926    k: usize,
927    n: usize,
928    batch_count: usize,
929) -> tenferro_tensor::Result<()> {
930    for batch in 0..batch_count {
931        let u_batch = batch * m * k;
932        let vt_batch = batch * k * n;
933        for col in 0..k {
934            let pivot = max_abs_pivot_c32(u, u_batch, m, col);
935            let pivot_value = u[u_batch + pivot + m * col];
936            let pivot_norm = pivot_value.norm();
937            if pivot_norm == 0.0 {
938                continue;
939            }
940            let phase = pivot_value.conj() / pivot_norm;
941            let vt_phase = phase.conj();
942            for row in 0..m {
943                let offset = u_batch + row + m * col;
944                u[offset] *= phase;
945            }
946            for vt_col in 0..n {
947                let offset = vt_batch + col + k * vt_col;
948                vt[offset] *= vt_phase;
949            }
950        }
951    }
952    Ok(())
953}
954
955fn max_abs_pivot_f64(u: &[f64], u_batch: usize, m: usize, col: usize) -> usize {
956    let mut pivot = 0;
957    let mut pivot_abs = u[u_batch + m * col].abs();
958    for row in 1..m {
959        let candidate_abs = u[u_batch + row + m * col].abs();
960        if candidate_abs > pivot_abs {
961            pivot = row;
962            pivot_abs = candidate_abs;
963        }
964    }
965    pivot
966}
967
968fn max_abs_pivot_f32(u: &[f32], u_batch: usize, m: usize, col: usize) -> usize {
969    let mut pivot = 0;
970    let mut pivot_abs = u[u_batch + m * col].abs();
971    for row in 1..m {
972        let candidate_abs = u[u_batch + row + m * col].abs();
973        if candidate_abs > pivot_abs {
974            pivot = row;
975            pivot_abs = candidate_abs;
976        }
977    }
978    pivot
979}
980
981fn max_abs_pivot_c64(u: &[Complex64], u_batch: usize, m: usize, col: usize) -> usize {
982    let mut pivot = 0;
983    let mut pivot_abs = u[u_batch + m * col].norm_sqr();
984    for row in 1..m {
985        let candidate_abs = u[u_batch + row + m * col].norm_sqr();
986        if candidate_abs > pivot_abs {
987            pivot = row;
988            pivot_abs = candidate_abs;
989        }
990    }
991    pivot
992}
993
994fn max_abs_pivot_c32(u: &[Complex32], u_batch: usize, m: usize, col: usize) -> usize {
995    let mut pivot = 0;
996    let mut pivot_abs = u[u_batch + m * col].norm_sqr();
997    for row in 1..m {
998        let candidate_abs = u[u_batch + row + m * col].norm_sqr();
999        if candidate_abs > pivot_abs {
1000            pivot = row;
1001            pivot_abs = candidate_abs;
1002        }
1003    }
1004    pivot
1005}
1006
1007fn require_matrix_meta(op: &'static str, shape: &[SymDim]) -> tenferro_tensor::Result<()> {
1008    if shape.len() < 2 {
1009        return Err(Error::RankMismatch {
1010            op,
1011            expected: 2,
1012            actual: shape.len(),
1013        });
1014    }
1015    Ok(())
1016}
1017
1018fn matrix_meta_parts<'a>(
1019    op: &'static str,
1020    shape: &'a [SymDim],
1021) -> tenferro_tensor::Result<(SymDim, SymDim, &'a [SymDim])> {
1022    require_matrix_meta(op, shape)?;
1023    Ok((shape[0].clone(), shape[1].clone(), &shape[2..]))
1024}
1025
1026fn lu_meta(dtype: DType, shape: &[SymDim]) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1027    let (m, n, batch) = matrix_meta_parts("tenferro-linalg.lu", shape)?;
1028    let k = m.clone().min(n.clone());
1029    Ok(vec![
1030        (dtype, matrix_shape(m.clone(), m, batch)),
1031        (dtype, matrix_shape(shape[0].clone(), k.clone(), batch)),
1032        (dtype, matrix_shape(k, n, batch)),
1033        (dtype, batch.to_vec()),
1034    ])
1035}
1036
1037fn lu_factor_meta(
1038    dtype: DType,
1039    shape: &[SymDim],
1040) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1041    let (m, n, batch) = matrix_meta_parts("tenferro-linalg.lu_factor", shape)?;
1042    let k = m.min(n);
1043    Ok(vec![
1044        (dtype, shape.to_vec()),
1045        (DType::I32, vector_shape(k, batch)),
1046        (dtype, batch.to_vec()),
1047    ])
1048}
1049
1050fn full_piv_lu_meta(
1051    dtype: DType,
1052    shape: &[SymDim],
1053) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1054    let (n, _, batch) = matrix_meta_parts("tenferro-linalg.full_piv_lu", shape)?;
1055    Ok(vec![
1056        (dtype, matrix_shape(n.clone(), n.clone(), batch)),
1057        (dtype, matrix_shape(n.clone(), n.clone(), batch)),
1058        (dtype, matrix_shape(n.clone(), n.clone(), batch)),
1059        (dtype, matrix_shape(n.clone(), n, batch)),
1060        (singular_values_dtype(dtype), batch.to_vec()),
1061    ])
1062}
1063
1064fn svd_meta(dtype: DType, shape: &[SymDim]) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1065    let (m, n, batch) = matrix_meta_parts("tenferro-linalg.svd", shape)?;
1066    let k = m.clone().min(n.clone());
1067    Ok(vec![
1068        (dtype, matrix_shape(m, k.clone(), batch)),
1069        (singular_values_dtype(dtype), vector_shape(k.clone(), batch)),
1070        (dtype, matrix_shape(k, n, batch)),
1071    ])
1072}
1073
1074fn svd_values_meta(
1075    dtype: DType,
1076    shape: &[SymDim],
1077) -> tenferro_tensor::Result<(DType, Vec<SymDim>)> {
1078    let (m, n, batch) = matrix_meta_parts("tenferro-linalg.svd_values", shape)?;
1079    let k = m.min(n);
1080    Ok((singular_values_dtype(dtype), vector_shape(k, batch)))
1081}
1082
1083fn qr_meta(dtype: DType, shape: &[SymDim]) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1084    let (m, n, batch) = matrix_meta_parts("tenferro-linalg.qr", shape)?;
1085    let k = m.clone().min(n.clone());
1086    Ok(vec![
1087        (dtype, matrix_shape(m, k.clone(), batch)),
1088        (dtype, matrix_shape(k, n, batch)),
1089    ])
1090}
1091
1092fn eigh_meta(dtype: DType, shape: &[SymDim]) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1093    let (n, _, batch) = matrix_meta_parts("tenferro-linalg.eigh", shape)?;
1094    Ok(vec![
1095        (singular_values_dtype(dtype), vector_shape(n.clone(), batch)),
1096        (dtype, matrix_shape(n.clone(), n, batch)),
1097    ])
1098}
1099
1100fn eigh_values_meta(
1101    dtype: DType,
1102    shape: &[SymDim],
1103) -> tenferro_tensor::Result<(DType, Vec<SymDim>)> {
1104    let (n, _, batch) = matrix_meta_parts("tenferro-linalg.eigh_values", shape)?;
1105    Ok((singular_values_dtype(dtype), vector_shape(n, batch)))
1106}
1107
1108fn eig_meta(
1109    input_dtype: DType,
1110    shape: &[SymDim],
1111) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
1112    let dtype = eig_output_dtype(input_dtype);
1113    let (n, _, batch) = matrix_meta_parts("tenferro-linalg.eig", shape)?;
1114    Ok(vec![
1115        (dtype, vector_shape(n.clone(), batch)),
1116        (dtype, matrix_shape(n.clone(), n, batch)),
1117    ])
1118}
1119
1120fn eig_values_meta(
1121    input_dtype: DType,
1122    shape: &[SymDim],
1123) -> tenferro_tensor::Result<(DType, Vec<SymDim>)> {
1124    let dtype = eig_output_dtype(input_dtype);
1125    let (n, _, batch) = matrix_meta_parts("tenferro-linalg.eig_values", shape)?;
1126    Ok((dtype, vector_shape(n, batch)))
1127}
1128
1129fn matrix_shape(rows: SymDim, cols: SymDim, batch: &[SymDim]) -> Vec<SymDim> {
1130    let mut shape = vec![rows, cols];
1131    shape.extend_from_slice(batch);
1132    shape
1133}
1134
1135fn vector_shape(len: SymDim, batch: &[SymDim]) -> Vec<SymDim> {
1136    let mut shape = vec![len];
1137    shape.extend_from_slice(batch);
1138    shape
1139}
1140
1141fn eig_output_dtype(dtype: DType) -> DType {
1142    match dtype {
1143        DType::F64 | DType::C64 => DType::C64,
1144        DType::F32 | DType::C32 => DType::C32,
1145        DType::I32 | DType::I64 | DType::Bool => DType::C64,
1146    }
1147}
1148
1149fn singular_values_dtype(dtype: DType) -> DType {
1150    match dtype {
1151        DType::C64 => DType::F64,
1152        DType::C32 => DType::F32,
1153        other => other,
1154    }
1155}
1156
1157fn promote_dtypes(dtypes: &[DType]) -> DType {
1158    dtypes
1159        .iter()
1160        .copied()
1161        .reduce(tenferro_tensor::validate::promote_dtype)
1162        .unwrap_or(DType::F64)
1163}
1164
1165fn hash_dtype(hasher: &mut dyn Hasher, dtype: DType) {
1166    let tag = match dtype {
1167        DType::F64 => 0,
1168        DType::F32 => 1,
1169        DType::I64 => 2,
1170        DType::C64 => 3,
1171        DType::C32 => 4,
1172        DType::I32 => 5,
1173        DType::Bool => 6,
1174    };
1175    hasher.write_u8(tag);
1176}
1177
1178fn hash_svd_gauge(hasher: &mut dyn Hasher, gauge: SvdGauge) {
1179    let tag = match gauge {
1180        SvdGauge::Raw => 0,
1181        SvdGauge::CanonicalPivot => 1,
1182    };
1183    hasher.write_u8(tag);
1184}
1185
1186fn hash_eigh_gauge(hasher: &mut dyn Hasher, gauge: EighGauge) {
1187    let tag = match gauge {
1188        EighGauge::Raw => 0,
1189        EighGauge::CanonicalPivot => 1,
1190    };
1191    hasher.write_u8(tag);
1192}
1193
1194fn hash_qr_gauge(hasher: &mut dyn Hasher, gauge: QrGauge) {
1195    let tag = match gauge {
1196        QrGauge::Raw => 0,
1197        QrGauge::PositiveDiagonal => 1,
1198    };
1199    hasher.write_u8(tag);
1200}
tenferro_linalg/extension.rs

tenferro_linalg/
extension.rs
