LU.hs

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

module EigenHHLO.Test.LU (tests) where

import qualified Data.Text as T
import Foreign.Ptr (nullPtr)
import Test.Tasty
import Test.Tasty.HUnit

import HHLO.Core.Types (DType(..))
import HHLO.IR.AST (FuncArg(..), TensorType(..))
import HHLO.IR.Builder (arg, moduleFromBuilder)
import HHLO.IR.Pretty (render)
import HHLO.Runtime.PJRT.Types (PJRTApi(..), PJRTClient(..), PJRTDevice(..))
import HHLO.Session (sessionFrom)

import EigenHHLO.IR.LU (luBuilder)
import EigenHHLO.Core.Types (BackendType(..), EigenSession(..))

dummySess :: EigenSession
dummySess = EigenSession (sessionFrom (PJRTApi nullPtr) (PJRTClient nullPtr) (PJRTDevice nullPtr)) CPU

tests :: TestTree
tests = testGroup "LU"
    [ testCase "MLIR generation contains dgetrf" $ do
        let modu = moduleFromBuilder @'[2,2] @'F64 "main"
                [ FuncArg "a" (TensorType [2,2] F64) ]
                $ do a <- arg @'[2,2] @'F64
                     (lu, _piv) <- luBuilder @2 @2 @2 dummySess a
                     return lu
        let mlir = render modu
        assertBool "contains dgetrf" ("eigenhhlo_dgetrf" `T.isInfixOf` mlir)
    ]