QR.hs

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

-- | Typed 'stablehlo.custom_call' builders for QR factorization and explicit Q generation.
module EigenHHLO.IR.QR
    ( qrBuilder
    , qBuilder
    ) 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 qualified Data.Text as T

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

-- | QR factorization: A = Q · R.
-- Returns (A overwritten with R+reflectors, tau).
qrBuilder :: forall m n k. (KnownNat m, KnownNat n, KnownNat k)
          => EigenSession
          -> Tensor '[m, n] 'F64
          -> Builder (Tensor '[m, n] 'F64, Tensor '[k] 'F64)
qrBuilder esess a =
    let m = fromIntegral (natVal (Proxy @m)) :: Integer
        n = fromIntegral (natVal (Proxy @n)) :: Integer
        k = fromIntegral (natVal (Proxy @k)) :: Integer
        cfg = qrConfig m n
        apiVer = case eigenBackend esess of
            CPU -> 2
            GPU -> 3
        inType = TensorType [m, n] F64
        outType1 = TensorType [m, n] F64
        outType2 = TensorType [k] F64
        vids = [tensorValue a]
        inTypes = [inType]
    in do
        vidsRes <- customCallRaw "eigenhhlo_dgeqrf" vids inTypes cfg False apiVer [outType1, outType2]
        case vidsRes of
            [v1, v2] -> return (Tensor v1, Tensor v2)
            _        -> error "qrBuilder: expected exactly two results"

-- | Generate explicit Q matrix from QR reflectors.
qBuilder :: forall m n k. (KnownNat m, KnownNat n, KnownNat k)
         => EigenSession
         -> Tensor '[m, n] 'F64   -- ^ reflectors from qrBuilder
         -> Tensor '[k] 'F64      -- ^ tau from qrBuilder
         -> Builder (Tensor '[m, m] 'F64)
qBuilder esess a tau =
    let m = fromIntegral (natVal (Proxy @m)) :: Integer
        n = fromIntegral (natVal (Proxy @n)) :: Integer
        k = fromIntegral (natVal (Proxy @k)) :: Integer
        cfg = qrConfig m n <> ",k=" <> (T.pack $ show k)
        apiVer = case eigenBackend esess of
            CPU -> 2
            GPU -> 3
        inType1 = TensorType [m, n] F64
        inType2 = TensorType [k] F64
        outType = TensorType [m, m] F64
        vids = [tensorValue a, tensorValue tau]
        inTypes = [inType1, inType2]
    in do
        vidsRes <- customCallRaw "eigenhhlo_dorgqr" vids inTypes cfg False apiVer [outType]
        case vidsRes of
            [v1] -> return (Tensor v1)
            _    -> error "qBuilder: expected exactly one result"