eigen-hhlo-0.1.0.0: examples/01-cholesky.hs

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

-- | Example 1: Cholesky factorization of a 3x3 symmetric positive-definite matrix.
--
-- Build and run with:
--   cabal run example-cholesky --flag=examples

module Main where

import qualified Data.Text as T

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

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

main :: IO ()
main = withEigenGPU $ \esess -> do
    putStrLn "=== Example 1: Cholesky Factorization (GPU) ==="

    let modu = moduleFromBuilder @'[3,3] @'F64 "main"
            [ FuncArg "a" (TensorType [3,3] F64) ]
            $ do a <- arg @'[3,3] @'F64
                 chol esess a

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

    compiled <- compile (eigenSession esess) modu

    -- A = [[25, 15,  5],
    --      [15, 18,  0],
    --      [ 5,  0, 11]]  (column-major)
    let input = hostFromList @'[3,3] @'F64
            [25, 15, 5, 15, 18, 0, 5, 0, 11]

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

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

    putStrLn "\nCholesky factor L (column-major):"
    print (hostToList result)