Cholesky.hs

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

-- | Typed 'stablehlo.custom_call' builder for Cholesky factorization.
module EigenHHLO.IR.Cholesky
    ( cholBuilder
    ) 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 (cholConfig)
import EigenHHLO.Core.Types (BackendType(..), EigenSession(..))

-- | Cholesky factorization via LAPACK @dpotrf@.
--
-- Factorizes a symmetric positive-definite matrix A into L · Lᵀ
-- (lower triangular).  The @uplo@ parameter is taken from the
-- 'BackendConfig' (currently hard-coded to @'L'@).
cholBuilder :: forall n. KnownNat n
            => EigenSession
            -> Tensor '[n, n] 'F64
            -> Builder (Tensor '[n, n] 'F64)
cholBuilder esess a =
    let n = fromIntegral (natVal (Proxy @n)) :: Integer
        cfg = cholConfig n 'L'
        apiVer = case eigenBackend esess of
            CPU -> 2
            GPU -> 3
        inType = TensorType [n, n] F64
        outType = TensorType [n, n] F64
        vids = [tensorValue a]
        inTypes = [inType]
    in do
        vidsRes <- customCallRaw "eigenhhlo_dpotrf" vids inTypes cfg False apiVer [outType]
        case vidsRes of
            [v1] -> return (Tensor v1)
            _    -> error "cholBuilder: expected exactly one result"