QR.hs

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

module EigenHHLO.Test.QR (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.QR (qrBuilder)
import EigenHHLO.Core.Types (BackendType(..), EigenSession(..))

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

tests :: TestTree
tests = testGroup "QR"
    [ testCase "MLIR generation contains dgeqrf" $ do
        let modu = moduleFromBuilder @'[2,2] @'F64 "main"
                [ FuncArg "a" (TensorType [2,2] F64) ]
                $ do a <- arg @'[2,2] @'F64
                     (r, _tau) <- qrBuilder @2 @2 @2 dummySess a
                     return r
        let mlir = render modu
        assertBool "contains dgeqrf" ("eigenhhlo_dgeqrf" `T.isInfixOf` mlir)
    ]