Session.hs

{-# LANGUAGE OverloadedStrings #-}

-- | Session lifecycle management for eigen-hhlo backends.
module EigenHHLO.Runtime.Session
    ( withEigenCPU
    , withEigenGPU
    ) where

import Data.Char (toLower, toUpper)
import System.Directory (doesFileExist)
import System.Environment (lookupEnv)

import HHLO.Runtime.CustomCall (loadCustomCallLibrary, registerGpuCustomCall)
import HHLO.Session (Session(..), withCPU, withGPU)

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

-- | Resolve the path to a custom-call shared library.
--
-- Priority:
--   1. @EIGENHHLO_<BACKEND>_LIB@ environment variable
--   2. @lib/libeigenhhlo_<backend>.so@ relative to CWD
--   3. Runtime error with instructions
getCustomCallLibPath :: String -> FilePath -> IO FilePath
getCustomCallLibPath backend defaultName = do
    mEnv <- lookupEnv ("EIGENHHLO_" ++ map toUpper backend ++ "_LIB")
    case mEnv of
        Just p  -> return p
        Nothing -> do
            let defaultPath = "lib/" ++ defaultName
            exists <- doesFileExist defaultPath
            if exists
                then return defaultPath
                else error $ unlines
                    [ "eigen-hhlo " ++ backend ++ " custom-call library not found at: " ++ defaultPath
                    , ""
                    , "To fix this, either:"
                    , "  1. Build the library:"
                    , "       cd cbits/" ++ map toLower backend ++ " && bash build.sh"
                    , "  2. Set the environment variable:"
                    , "       export EIGENHHLO_" ++ map toUpper backend ++ "_LIB=/path/to/" ++ defaultName
                    ]

-- | Create an eigen-hhlo session on the CPU backend.
-- Loads the CPU custom-call library via RTLD_GLOBAL.
--
-- The library path is resolved via 'getCustomCallLibPath'.
withEigenCPU :: (EigenSession -> IO a) -> IO a
withEigenCPU action = withCPU $ \sess -> do
    libPath <- getCustomCallLibPath "CPU" "libeigenhhlo_cpu.so"
    loadCustomCallLibrary libPath
    action (EigenSession sess CPU)

-- | Create an eigen-hhlo session on the GPU backend.
-- Registers all decomposition symbols with the PJRT CUDA plugin.
--
-- The library path is resolved via 'getCustomCallLibPath'.
withEigenGPU :: (EigenSession -> IO a) -> IO a
withEigenGPU action = withGPU $ \sess -> do
    libPath <- getCustomCallLibPath "GPU" "libeigenhhlo_gpu.so"
    -- Register all symbols with the PJRT GPU plugin
    registerGpuCustomCall (sessionApi sess) libPath "eigenhhlo_dgesvd"
    registerGpuCustomCall (sessionApi sess) libPath "eigenhhlo_dgeqrf"
    registerGpuCustomCall (sessionApi sess) libPath "eigenhhlo_dorgqr"
    registerGpuCustomCall (sessionApi sess) libPath "eigenhhlo_dsyevd"
    registerGpuCustomCall (sessionApi sess) libPath "eigenhhlo_dpotrf"
    registerGpuCustomCall (sessionApi sess) libPath "eigenhhlo_dgetrf"
    action (EigenSession sess GPU)