SVD.hs

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

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

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

tests :: TestTree
tests = testGroup "SVD"
    [ testCase "MLIR generation contains dgesvd" $ do
        let modu = moduleFromBuilder @'[2,2] @'F64 "main"
                [ FuncArg "a" (TensorType [2,3] F64) ]
                $ do a <- arg @'[2,3] @'F64
                     (u, _s, _vt) <- svdBuilder @2 @3 @2 dummySess a
                     return u
        let mlir = render modu
        assertBool "contains dgesvd" ("eigenhhlo_dgesvd" `T.isInfixOf` mlir)

    , testCase "MLIR output shapes are correct for thin SVD" $ do
        let modu = moduleFromBuilder @'[3] @'F64 "main"
                [ FuncArg "a" (TensorType [4,3] F64) ]
                $ do a <- arg @'[4,3] @'F64
                     (_u, s, _vt) <- svdBuilder @4 @3 @3 dummySess a
                     return s
        let mlir = render modu
        -- U should be 4x3, S should be 3, Vt should be 3x3
        assertBool "U shape 4x3"   ("tensor<4x3xf64>" `T.isInfixOf` mlir)
        assertBool "S shape 3"     ("tensor<3xf64>"   `T.isInfixOf` mlir)
        assertBool "Vt shape 3x3"  ("tensor<3x3xf64>" `T.isInfixOf` mlir)
    ]