eigen-hhlo-0.1.0.0: app/Main.hs

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

module Main (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 EigenHHLO.IR.Cholesky (cholBuilder)
import EigenHHLO.Runtime.Session (withEigenCPU)

main :: IO ()
main = withEigenCPU $ \esess -> do
    -- Build a StableHLO module:  f(a) = chol(a)
    let modu = moduleFromBuilder @'[2,2] @'F64 "main"
            [ FuncArg "a" (TensorType [2,2] F64) ]
            $ do a <- arg @'[2,2] @'F64
                 cholBuilder esess a

    putStrLn "=== Generated StableHLO (MLIR) ==="
    putStrLn $ T.unpack (render modu)

    putStrLn "\n=== Note ==="
    putStrLn "CPU custom-call execution is currently limited by the PJRT CPU plugin:"
    putStrLn "the plugin does not expose an API to register external custom-call targets"
    putStrLn "(unlike the GPU plugin which has PJRT_Gpu_Custom_Call)."
    putStrLn "MLIR generation and compilation work; runtime execution requires either:"
    putStrLn "  1. A custom PJRT CPU plugin built with the LAPACK wrappers linked in, or"
    putStrLn "  2. GPU backend via withEigenGPU + registerGpuCustomCall (future work)."