Eigenvalue.hs

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

-- | Typed 'stablehlo.custom_call' builder for symmetric eigenvalue decomposition.
module EigenHHLO.IR.Eigenvalue
    ( eigBuilder
    ) 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 (eigConfig)
import EigenHHLO.Core.Types (BackendType(..), EigenSession(..))

-- | Symmetric eigenvalue decomposition: A = V · Λ · V^T.
-- Returns (eigenvalues, eigenvectors).
eigBuilder :: forall n. KnownNat n
           => EigenSession
           -> Tensor '[n, n] 'F64
           -> Builder (Tensor '[n] 'F64, Tensor '[n, n] 'F64)
eigBuilder esess a =
    let n = fromIntegral (natVal (Proxy @n)) :: Integer
        cfg = eigConfig n 'L'
        apiVer = case eigenBackend esess of
            CPU -> 2
            GPU -> 3
        inType = TensorType [n, n] F64
        outType1 = TensorType [n] F64
        outType2 = TensorType [n, n] F64
        vids = [tensorValue a]
        inTypes = [inType]
    in do
        vidsRes <- customCallRaw "eigenhhlo_dsyevd" vids inTypes cfg False apiVer [outType1, outType2]
        case vidsRes of
            [v1, v2] -> return (Tensor v1, Tensor v2)
            _        -> error "eigBuilder: expected exactly two results"