eigen-hhlo-0.1.0.0: examples/07-gpu-setup.hs

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

-- | Example 7: GPU custom-call registration and MLIR generation.
--
-- Demonstrates the GPU session setup path: loading the PJRT CUDA plugin,
-- registering all decomposition symbols with registerGpuCustomCall,
-- building a module, and compiling.
--
-- NOTE: This requires a CUDA-capable GPU and the PJRT CUDA plugin.
-- Without GPU hardware it will fail at session creation.
--
-- Build and run with:
--   cabal run example-gpu-setup --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 (compile)

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

main :: IO ()
main = withEigenGPU $ \esess -> do
    putStrLn "=== Example 7: GPU Custom-Call Registration ==="

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

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

    _compiled <- compile (eigenSession esess) modu

    putStrLn "\nCompilation succeeded."
    putStrLn "On a GPU machine, you could now upload buffers and execute with 'run'."