eigen-hhlo-0.1.0.0: examples/04-eigenvalue.hs

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

-- | Example 4: Symmetric eigenvalue decomposition.
--
-- Returns eigenvalues and eigenvectors: A = V · Λ · Vᵀ.
--
-- Build and run with:
--   cabal run example-eigenvalue --flag=examples

module Main where

import qualified Data.Text as T

import HHLO.Core.Types (DType(..))
import HHLO.IR.AST (FuncArg(..), TensorType(..))
import HHLO.EDSL.Ops (returnTuple2)
import HHLO.IR.Builder (arg, moduleFromBuilder2)
import HHLO.IR.Pretty (render)
import HHLO.Session (HostTensor, compile, hostFromList, hostToList, run)

import EigenHHLO.Core.Types (eigenSession)
import EigenHHLO.EDSL.Decomposition (eig)
import EigenHHLO.Runtime.Session (withEigenGPU)

main :: IO ()
main = withEigenGPU $ \esess -> do
    putStrLn "=== Example 4: Symmetric Eigenvalue Decomposition (GPU) ==="

    let modu = moduleFromBuilder2 @'[3] @'F64 @'[3,3] @'F64 "main"
            [ FuncArg "a" (TensorType [3,3] F64) ]
            $ do a <- arg @'[3,3] @'F64
                 (w, v) <- eig esess a
                 returnTuple2 w v

    putStrLn "\nGenerated MLIR:"
    putStrLn (T.unpack $ render modu)

    compiled <- compile (eigenSession esess) modu

    -- A = [[4, 2, 1],
    --      [2, 5, 3],
    --      [1, 3, 6]]  (column-major, symmetric)
    let input = hostFromList @'[3,3] @'F64
            [4, 2, 1, 2, 5, 3, 1, 3, 6]

    putStrLn "\nInput matrix (column-major):"
    print (hostToList input)

    (wResult, vResult) <- run (eigenSession esess) compiled input
        :: IO (HostTensor '[3] 'F64, HostTensor '[3,3] 'F64)

    putStrLn "\nEigenvalues (3,):"
    print (hostToList wResult)
    putStrLn "\nEigenvectors (3x3, column-major):"
    print (hostToList vResult)