eigen-hhlo-0.1.0.0: examples/05-lu.hs

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

-- | Example 5: LU factorization with partial pivoting.
--
-- Demonstrates mixed-dtype output: LU factors (F64) + pivot indices (I32).
--
-- Build and run with:
--   cabal run example-lu --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 (lu)
import EigenHHLO.Runtime.Session (withEigenGPU)

main :: IO ()
main = withEigenGPU $ \esess -> do
    putStrLn "=== Example 5: LU Factorization with Partial Pivoting (GPU) ==="

    let modu = moduleFromBuilder2 @'[3,3] @'F64 @'[3] @'I32 "main"
            [ FuncArg "a" (TensorType [3,3] F64) ]
            $ do a <- arg @'[3,3] @'F64
                 (luFact, piv) <- lu @3 @3 @3 esess a
                 returnTuple2 luFact piv

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

    compiled <- compile (eigenSession esess) modu

    -- A = [[2, 1, 1],
    --      [4, 3, 3],
    --      [8, 7, 9]]  (column-major)
    let input = hostFromList @'[3,3] @'F64
            [2, 4, 8, 1, 3, 7, 1, 3, 9]

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

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

    putStrLn "\nLU factors (3x3, column-major):"
    print (hostToList luResult)
    putStrLn "\nPivot indices (3,):"
    print (hostToList pivResult)