LU.hs

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

-- | Typed 'stablehlo.custom_call' builder for LU factorization with partial pivoting.
module EigenHHLO.IR.LU
    ( luBuilder
    ) 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 (luConfig)
import EigenHHLO.Core.Types (BackendType(..), EigenSession(..))

-- | LU factorization with partial pivoting via LAPACK @dgetrf@.
--
-- Returns @(LU, pivots)@ where @LU@ contains both L and U packed
-- together, and @pivots@ are the 0-based permutation indices.
luBuilder :: forall m n k. (KnownNat m, KnownNat n, KnownNat k)
          => EigenSession
          -> Tensor '[m, n] 'F64
          -> Builder (Tensor '[m, n] 'F64, Tensor '[k] 'I32)
luBuilder esess a =
    let m = fromIntegral (natVal (Proxy @m)) :: Integer
        n = fromIntegral (natVal (Proxy @n)) :: Integer
        k = fromIntegral (natVal (Proxy @k)) :: Integer
        cfg = luConfig m n
        apiVer = case eigenBackend esess of
            CPU -> 2
            GPU -> 3
        inType = TensorType [m, n] F64
        outType1 = TensorType [m, n] F64
        outType2 = TensorType [k] I32
        vids = [tensorValue a]
        inTypes = [inType]
    in do
        vidsRes <- customCallRaw "eigenhhlo_dgetrf" vids inTypes cfg False apiVer [outType1, outType2]
        case vidsRes of
            [v1, v2] -> return (Tensor v1, Tensor v2)
            _        -> error "luBuilder: expected exactly two results"