eigen-hhlo-0.1.0.0: examples/06-pipeline.hs

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

-- | Example 6: Composed linear algebra pipeline.
--
-- Computes the Gram matrix (Aᵀ · A) then performs Cholesky factorization.
-- Demonstrates integration between standard HHLO ops and eigen-hhlo decompositions.
--
-- Build and run with:
--   cabal run example-pipeline --flag=examples

module Main where

import qualified Data.Text as T

import HHLO.Core.Types (DType(..))
import HHLO.EDSL.Ops (matmul, transpose, v2)
import HHLO.IR.AST (FuncArg(..), TensorType(..))
import HHLO.IR.Builder (Builder, Tensor, 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 6: Pipeline (Gram Matrix -> Cholesky) (GPU) ==="

    let modu = moduleFromBuilder @'[3,3] @'F64 "main"
            [ FuncArg "a" (TensorType [4,3] F64) ]
            $ do a <- arg @'[4,3] @'F64
                 at <- transpose (v2 1 0) a :: Builder (Tensor '[3,4] 'F64)
                 gram <- matmul at a
                 chol esess gram

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

    compiled <- compile (eigenSession esess) modu

    -- A = [[1, 0, 0],
    --      [1, 1, 0],
    --      [1, 1, 1],
    --      [1, 1, 1]]  (4x3, column-major)
    let input = hostFromList @'[4,3] @'F64
            [1, 1, 1, 1, 0, 1, 1, 1, 0, 0, 1, 1]

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

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

    putStrLn "\nCholesky of Gram matrix (column-major):"
    print (hostToList result)