eigen-hhlo-0.1.0.0: examples/03-qr.hs

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

-- | Example 3: QR factorization followed by explicit Q generation.
--
-- Demonstrates chaining two dependent custom calls inside one module:
--   1. qr: A -> (reflectors, tau)
--   2. q:  (reflectors, tau) -> Q
--
-- Build and run with:
--   cabal run example-qr --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 (q, qr)
import EigenHHLO.Runtime.Session (withEigenGPU)

main :: IO ()
main = withEigenGPU $ \esess -> do
    putStrLn "=== Example 3: QR Factorization -> Explicit Q (GPU) ==="

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

    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 '[4,4] 'F64)

    putStrLn "\nExplicit Q (4x4, column-major):"
    print (hostToList result)