SVD.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-- | Typed 'stablehlo.custom_call' builder for singular value decomposition.
module EigenHHLO.IR.SVD
    ( svdBuilder
    ) where

import Data.Proxy (Proxy(..))
import GHC.TypeLits (KnownNat, natVal)

import HHLO.Core.Types (DType(..))
import HHLO.EDSL.Ops (customCallRaw)
import HHLO.IR.AST (TensorType(..))
import HHLO.IR.Builder (Builder, Tensor(..), tensorValue)

import EigenHHLO.Core.BackendConfig (svdConfig)
import EigenHHLO.Core.Types (BackendType(..), EigenSession(..))

-- | Singular value decomposition: A = U · diag(S) · Vt.
-- Returns (U, S, Vt).
svdBuilder :: forall m n k. (KnownNat m, KnownNat n, KnownNat k)
           => EigenSession
           -> Tensor '[m, n] 'F64
           -> Builder (Tensor '[m, k] 'F64, Tensor '[k] 'F64, Tensor '[k, n] 'F64)
svdBuilder esess a =
    let m = fromIntegral (natVal (Proxy @m)) :: Integer
        n = fromIntegral (natVal (Proxy @n)) :: Integer
        k = fromIntegral (natVal (Proxy @k)) :: Integer
        cfg = svdConfig m n 'S' 'S'
        apiVer = case eigenBackend esess of
            CPU -> 2
            GPU -> 3
        inType = TensorType [m, n] F64
        outType1 = TensorType [m, k] F64
        outType2 = TensorType [k] F64
        outType3 = TensorType [k, n] F64
        vids = [tensorValue a]
        inTypes = [inType]
    in do
        vidsRes <- customCallRaw "eigenhhlo_dgesvd" vids inTypes cfg False apiVer [outType1, outType2, outType3]
        case vidsRes of
            [v1, v2, v3] -> return (Tensor v1, Tensor v2, Tensor v3)
            _            -> error "svdBuilder: expected exactly three results"