diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,20 @@ +# Revision history for eigen-hhlo + +## 0.1.0.0 -- 2025-04-27 + +* Initial release. +* CPU LAPACK custom-call wrappers: Cholesky (`dpotrf`), SVD (`dgesvd`), + QR (`dgeqrf` / `dorgqr`), symmetric eigenvalue (`dsyevd`), and + LU with partial pivoting (`dgetrf`). +* Typed EDSL builders for all five decompositions via + `stablehlo.custom_call` with `api_version = 2` (CPU) and + `api_version = 3` (GPU) ABIs. +* Session management (`withEigenCPU` / `withEigenGPU`) with backend + tagging and GPU custom-call registration infrastructure. +* Example suite demonstrating MLIR generation for each decomposition + plus a composed pipeline (Gram matrix → Cholesky). +* Known limitation: CPU end-to-end execution is blocked by the + PJRT CPU plugin, which does not expose an external custom-call + registration API. MLIR generation and compilation work correctly; + runtime execution requires either a custom PJRT CPU plugin build + or the GPU backend. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 overshiki + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,292 @@ +# eigen-hhlo + +Dense linear algebra decompositions on [HHLO](https://github.com/overshiki/hhlo) via XLA custom calls. + +**eigen-hhlo** provides singular value decomposition (SVD), QR decomposition, symmetric eigenvalue decomposition, Cholesky factorization, and LU decomposition with partial pivoting as first-class operations within the HHLO ecosystem. Computations run on XLA through PJRT using **GPU via cuSOLVER**. + +> **Note on CPU support:** A CPU backend via LAPACK is prepared but deferred pending upstream PJRT CPU plugin support. See [Appendix: CPU Backend](#appendix-cpu-backend) for details. + +--- + +## Background + +The HHLO ecosystem (Haskell-frontend for StableHLO/XLA) provides a typed EDSL for building MLIR modules and executing them through PJRT. While it covers neural network primitives, reductions, and control flow, it lacks dense linear algebra factorizations. StableHLO itself does not define decomposition ops; the standard way to add them is via `stablehlo.custom_call` backed by a runtime library. eigen-hhlo fills this gap by exposing cuSOLVER routines through XLA custom calls. + +--- + +## Design + +The Haskell EDSL emits `stablehlo.custom_call` operations whose `call_target_name` matches symbols in a CUDA shared library (e.g. `eigenhhlo_dpotrf`). At session startup, `withEigenGPU` registers these symbols with the PJRT CUDA plugin via `registerGpuCustomCall`. When XLA compiles the StableHLO module, it maps each custom call to the registered CUDA function; at runtime the wrapper handles device buffer management, cuSOLVER workspace allocation, and kernel launch. + +- **HHLO-only**: No dependency on `hmatrix` or other host linear algebra libraries. Everything lives inside the XLA graph. +- **GPU via cuSOLVER**: C++ wrappers around cuSOLVER (`cusolverDnDpotrf`, `cusolverDnDgesvd`, `cusolverDnDgeqrf`/`cusolverDnDorgqr`, `cusolverDnDsyevd`, `cusolverDnDgetrf`) conforming to the XLA GPU custom-call ABI (`api_version = 3`). +- **Type-safe shapes**: All ops are statically shaped using GHC `TypeLits` (e.g. `Tensor '[m, n] 'F64`). + +--- + +## Implementation + +### C GPU Library + +`cbits/gpu/eigenhhlo_cusolver.cu` wraps cuSOLVER routines for the XLA GPU custom-call ABI: + +| Symbol | cuSOLVER | Operation | +|--------|----------|-----------| +| `eigenhhlo_dpotrf` | `cusolverDnDpotrf` | Cholesky factorization | +| `eigenhhlo_dgesvd` | `cusolverDnDgesvd` | SVD | +| `eigenhhlo_dgeqrf` | `cusolverDnDgeqrf` | QR factorization (reflectors) | +| `eigenhhlo_dorgqr` | `cusolverDnDorgqr` | Generate Q from QR reflectors | +| `eigenhhlo_dsyevd` | `cusolverDnDsyevd` | Symmetric eigenvalue decomposition | +| `eigenhhlo_dgetrf` | `cusolverDnDgetrf` | LU with partial pivoting | + +The opaque `backend_config` string carries dimensions and options (e.g. `"n=2,uplo=L"`). + +### Haskell Modules + +| Module | Purpose | +|--------|---------| +| `EigenHHLO.IR.Cholesky` | `cholBuilder` — typed `customCallRaw` wrapper for Cholesky | +| `EigenHHLO.IR.SVD` | `svdBuilder` — 3-output custom call (U, S, Vt) | +| `EigenHHLO.IR.QR` | `qrBuilder` / `qBuilder` — QR + explicit Q generation | +| `EigenHHLO.IR.Eigenvalue` | `eigBuilder` — eigenvalues + eigenvectors | +| `EigenHHLO.IR.LU` | `luBuilder` — LU factors + pivot indices | +| `EigenHHLO.EDSL.Decomposition` | Re-exports with simpler names (`chol`, `svd`, `qr`, etc.) | +| `EigenHHLO.Runtime.Session` | `withEigenGPU` — session setup + library loading | +| `EigenHHLO.Core.BackendConfig` | Opaque config string builders | + +--- + +## Building + +### ⚠️ Required Manual Step: Compile the CUDA Custom-Call Library + +**eigen-hhlo does NOT compile its CUDA code automatically through Cabal.** +You must run the provided `build.sh` script manually before the package can execute any custom calls at runtime. + +```bash +cd cbits/gpu && bash build.sh && cd ../.. +# Produces: lib/libeigenhhlo_gpu.so +``` + +**Why this is necessary:** +The custom-call functions (`eigenhhlo_dpotrf`, `eigenhhlo_dgesvd`, etc.) live in a separate shared library that is loaded at runtime by `withEigenGPU` (via `registerGpuCustomCall`). Cabal's `extra-source-files` only ships the source code in the tarball; it does **not** invoke `nvcc` during the Haskell build. This is a deliberate design choice because CUDA compilation requires `nvcc`, which cannot be assumed to be present on every machine that installs the Haskell package. + +**Runtime path resolution:** +The session helper resolves the custom-call library path using the same three-tier strategy as HHLO's `getPluginPath`: + +1. **Environment variable** (highest priority): `EIGENHHLO_GPU_LIB` +2. **Default relative path**: `lib/libeigenhhlo_gpu.so` +3. **Clear runtime error** with build instructions if the file is not found + +This means you can run your executable from any working directory by setting an absolute path: + +```bash +export EIGENHHLO_GPU_LIB=/home/you/projects/myapp/lib/libeigenhhlo_gpu.so +./myapp +``` + +### Prerequisites + +- GHC 9.6+ +- Cabal 3.10+ +- CUDA toolkit 11.8+ with cuSOLVER (for GPU backend) + +### Step-by-step build + +```bash +# 1. Build the GPU custom-call library (REQUIRED) +cd cbits/gpu && bash build.sh && cd ../.. +# Produces: lib/libeigenhhlo_gpu.so + +# 2. Build the Haskell project +cabal build all + +# 3. Run tests +cabal test test:eigen-hhlo-test +``` + +### Build and run the examples + +```bash +# 1. Compile GPU custom-call library FIRST (required) +cd cbits/gpu && bash build.sh && cd ../.. + +# 2. Build all examples +cabal build eigen-hhlo:exes --flags=examples + +# 3. Run an individual example +cabal run example-cholesky --flags=examples +``` + +--- + +## Usage + +### Cholesky Example + +```haskell +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} + +import HHLO.Core.Types (DType(..)) +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder (arg, moduleFromBuilder) +import HHLO.IR.Pretty (render) + +import EigenHHLO.EDSL.Decomposition (chol) +import EigenHHLO.Runtime.Session (withEigenGPU) + +main :: IO () +main = withEigenGPU $ \esess -> do + let modu = moduleFromBuilder @'[2,2] @'F64 "main" + [ FuncArg "a" (TensorType [2,2] F64) ] + $ do a <- arg @'[2,2] @'F64 + chol esess a + + putStrLn $ T.unpack (render modu) +``` + +Output: +```mlir +module { + func.func @main(%arg0: tensor<2x2xf64>) -> tensor<2x2xf64> { + %0 = stablehlo.custom_call @eigenhhlo_dpotrf(%arg0) + {call_target_name = "eigenhhlo_dpotrf", has_side_effect = false, + backend_config = "n=2,uplo=L", api_version = 3 : i32} + : (tensor<2x2xf64>) -> tensor<2x2xf64> + return %0 : tensor<2x2xf64> + } +} +``` + +### Available Operations + +```haskell +-- Cholesky: A = L · Lᵀ +chol :: KnownNat n => EigenSession -> Tensor '[n,n] 'F64 -> Builder (Tensor '[n,n] 'F64) + +-- SVD: A = U · diag(S) · Vt +svd :: (KnownNat m, KnownNat n, KnownNat k) => EigenSession -> Tensor '[m,n] 'F64 + -> Builder (Tensor '[m,k] 'F64, Tensor '[k] 'F64, Tensor '[k,n] 'F64) + +-- QR factorization: A = Q · R +qr :: (KnownNat m, KnownNat n, KnownNat k) => EigenSession -> Tensor '[m,n] 'F64 + -> Builder (Tensor '[m,n] 'F64, Tensor '[k] 'F64) + +-- Generate explicit Q matrix from QR reflectors +q :: (KnownNat m, KnownNat n, KnownNat k) => EigenSession -> Tensor '[m,n] 'F64 -> Tensor '[k] 'F64 + -> Builder (Tensor '[m,m] 'F64) + +-- Symmetric eigenvalue: A = V · Λ · Vᵀ +eig :: KnownNat n => EigenSession -> Tensor '[n,n] 'F64 + -> Builder (Tensor '[n] 'F64, Tensor '[n,n] 'F64) + +-- LU with partial pivoting +lu :: (KnownNat m, KnownNat n, KnownNat k) => EigenSession -> Tensor '[m,n] 'F64 + -> Builder (Tensor '[m,n] 'F64, Tensor '[k] 'I32) +``` + +### Session Setup + +```haskell +import EigenHHLO.Runtime.Session (withEigenGPU) + +-- GPU: registers symbols with PJRT CUDA plugin via cuSOLVER +withEigenGPU $ \esess -> do + -- build, compile, run ... + pure () +``` + +--- + +## Current Status + +- **MLIR generation**: ✅ Complete for all five decompositions. +- **Test suite**: ✅ 8/8 tests pass (5 MLIR generation + 3 GPU numerical end-to-end). +- **GPU execution**: ✅ **Working** on CUDA hardware via cuSOLVER custom calls. + +--- + +## Project Structure + +``` +eigen-hhlo/ +├── app/Main.hs -- CLI demo / smoke test (GPU) +├── cbits/ +│ └── gpu/ +│ ├── eigenhhlo_cusolver.cu -- cuSOLVER wrappers (XLA GPU ABI) +│ └── build.sh -- Builds lib/libeigenhhlo_gpu.so +├── examples/ +│ ├── 01-cholesky.hs -- Cholesky factorization (GPU) +│ ├── 02-svd.hs -- Singular value decomposition (GPU) +│ ├── 03-qr.hs -- QR factorization -> explicit Q (GPU) +│ ├── 04-eigenvalue.hs -- Symmetric eigenvalue decomposition (GPU) +│ ├── 05-lu.hs -- LU with partial pivoting (GPU) +│ ├── 06-pipeline.hs -- Composed pipeline (Gram -> Cholesky) (GPU) +│ └── 07-gpu-setup.hs -- GPU custom-call registration demo +├── lib/ +│ └── libeigenhhlo_gpu.so -- GPU custom-call shared library +├── src/ +│ ├── EigenHHLO/Core/ +│ │ ├── BackendConfig.hs -- Opaque config string builders +│ │ └── Types.hs -- BackendType, EigenSession +│ ├── EigenHHLO/EDSL/ +│ │ └── Decomposition.hs -- User-facing API +│ ├── EigenHHLO/IR/ +│ │ ├── Cholesky.hs +│ │ ├── Eigenvalue.hs +│ │ ├── LU.hs +│ │ ├── QR.hs +│ │ └── SVD.hs -- Typed custom_call builders +│ └── EigenHHLO/Runtime/ +│ └── Session.hs -- withEigenGPU +├── test/ -- Tasty/HUnit MLIR generation + GPU numerical tests +├── cabal.project +├── eigen-hhlo.cabal +└── devlog/ -- Design docs & analysis +``` + +--- + +## Appendix: CPU Backend + +The CPU backend is **deferred** — the source code is present but end-to-end execution is not yet possible. + +### Why it is deferred + +The PJRT CPU plugin uses an internal C++ `CustomCallTargetRegistry` singleton to resolve custom calls. This registry can only be populated by code **linked directly into the plugin binary** itself. Our LAPACK wrapper library (`libeigenhhlo_cpu.so`) would be loaded separately via `dlopen(RTLD_GLOBAL)`, and its symbols are globally visible — yet the plugin never looks them up. We verified this experimentally (`LD_PRELOAD`, `dlsym` tests, and symbol-table inspection all confirm the plugin ignores externally-loaded symbols). + +**What this means in practice:** +- ✅ MLIR generation and compilation work correctly on the CPU code path. +- ❌ Runtime execution crashes with `PJRTException "No registered implementation for untyped custom call to eigenhhlo_* for Host"`. +- This is **not a bug in eigen-hhlo**; it is an architectural limitation of the current PJRT CPU plugin. + +**Paths forward for CPU:** +1. **Build a custom PJRT CPU plugin** from OpenXLA source with the LAPACK wrappers linked in. +2. **Wait for upstream** — OpenXLA is tracking CPU custom-call registration APIs (see [openxla/xla#26928](https://github.com/openxla/xla/issues/26928)). + +### CPU source code + +| File | Purpose | +|------|---------| +| `cbits/cpu/eigenhhlo_lapack.c` | LAPACK wrappers (`dpotrf`, `dgesvd`, `dgeqrf`/`dorgqr`, `dsyevd`, `dgetrf`) for XLA CPU ABI (`api_version = 2`) | +| `cbits/cpu/build.sh` | Build script for `lib/libeigenhhlo_cpu.so` | +| `lib/libeigenhhlo_cpu.so` | Output shared library (not usable until upstream support arrives) | + +### Building the CPU library (optional) + +If you want to experiment with the CPU backend: + +```bash +# Prerequisites: LAPACK / OpenBLAS (liblapack.so.3, libopenblas.so.0) +cd cbits/cpu && bash build.sh && cd ../.. +# Produces: lib/libeigenhhlo_cpu.so +``` + +The Haskell API also exposes `withEigenCPU`, which resolves the library path via `EIGENHHLO_CPU_LIB` (environment variable) or the default `lib/libeigenhhlo_cpu.so`, following the same three-tier strategy as the GPU backend. + +--- + +## License + +MIT diff --git a/app/Main.hs b/app/Main.hs new file mode 100644 --- /dev/null +++ b/app/Main.hs @@ -0,0 +1,34 @@ +{-# 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)." diff --git a/cbits/cpu/build.sh b/cbits/cpu/build.sh new file mode 100644 --- /dev/null +++ b/cbits/cpu/build.sh @@ -0,0 +1,20 @@ +#!/bin/bash +# Build script for eigen-hhlo CPU custom-call library (LAPACK wrappers) + +set -e +CC=${CC:-gcc} +OUTDIR="${1:-../../lib}" + +mkdir -p "$OUTDIR" + +echo "Building eigen-hhlo CPU custom-call library..." +echo " Output: $OUTDIR/libeigenhhlo_cpu.so" + +$CC -shared -o "$OUTDIR/libeigenhhlo_cpu.so" \ + -fPIC \ + eigenhhlo_lapack.c \ + -llapack -lblas -O2 + +echo "Build complete." +echo "Load at runtime with:" +echo " loadCustomCallLibrary \"$OUTDIR/libeigenhhlo_cpu.so\"" diff --git a/cbits/cpu/eigenhhlo_lapack.c b/cbits/cpu/eigenhhlo_lapack.c new file mode 100644 --- /dev/null +++ b/cbits/cpu/eigenhhlo_lapack.c @@ -0,0 +1,313 @@ +/* eigen-hhlo CPU custom-call library — LAPACK wrappers. + * + * Conforms to XLA CPU custom-call ABI (API version 2): + * void target(void* out, void** in, + * const char* opaque, size_t opaque_len, + * XlaCustomCallStatus* status); + * + * For multi-result custom calls, `out` is a void** array of output buffers. + * + * Compile: + * cd cbits/cpu && bash build.sh + */ + +#include +#include +#include +#include + +/* -------------------------------------------------------------------------- + * XlaCustomCallStatus placeholder (opaque struct pointer) + * -------------------------------------------------------------------------- */ +typedef struct XlaCustomCallStatus_ XlaCustomCallStatus; + +/* -------------------------------------------------------------------------- + * Opaque string parser: key=value,... + * -------------------------------------------------------------------------- */ +static int parse_int_kv(const char* opaque, const char* key, int default_val) { + if (!opaque || !key) return default_val; + size_t keylen = strlen(key); + const char* p = opaque; + while (*p) { + if (strncmp(p, key, keylen) == 0 && p[keylen] == '=') { + p += keylen + 1; + return (int)strtol(p, NULL, 10); + } + while (*p && *p != ',') p++; + if (*p == ',') p++; + } + return default_val; +} + +static char parse_char_kv(const char* opaque, const char* key, char default_val) { + if (!opaque || !key) return default_val; + size_t keylen = strlen(key); + const char* p = opaque; + while (*p) { + if (strncmp(p, key, keylen) == 0 && p[keylen] == '=') { + p += keylen + 1; + return *p; + } + while (*p && *p != ',') p++; + if (*p == ',') p++; + } + return default_val; +} + +/* -------------------------------------------------------------------------- + * LAPACK declarations (Fortran ABI: all args by pointer, trailing underscore) + * -------------------------------------------------------------------------- */ +extern void dpotrf_(char* uplo, int* n, double* a, int* lda, int* info); +extern void dgetrf_(int* m, int* n, double* a, int* lda, int* ipiv, int* info); +extern void dgeqrf_(int* m, int* n, double* a, int* lda, double* tau, + double* work, int* lwork, int* info); +extern void dorgqr_(int* m, int* n, int* k, double* a, int* lda, double* tau, + double* work, int* lwork, int* info); +extern void dsyevd_(char* jobz, char* uplo, int* n, double* a, int* lda, + double* w, double* work, int* lwork, int* iwork, + int* liwork, int* info); +extern void dgesvd_(char* jobu, char* jobvt, int* m, int* n, double* a, + int* lda, double* s, double* u, int* ldu, double* vt, + int* ldvt, double* work, int* lwork, int* info); + +/* -------------------------------------------------------------------------- + * Helper: query LAPACK workspace size + * -------------------------------------------------------------------------- */ +static int query_lwork_dgeqrf(int m, int n) { + double work_query; + int lwork = -1, info; + dgeqrf_(&m, &n, NULL, &m, NULL, &work_query, &lwork, &info); + return (int)work_query; +} + +static int query_lwork_dorgqr(int m, int n, int k) { + double work_query; + int lwork = -1, info; + dorgqr_(&m, &n, &k, NULL, &m, NULL, &work_query, &lwork, &info); + return (int)work_query; +} + +static int query_lwork_dsyevd(int n) { + double work_query; + int iwork_query; + int lwork = -1, liwork = -1, info; + char jobz = 'V', uplo = 'L'; + dsyevd_(&jobz, &uplo, &n, NULL, &n, NULL, &work_query, &lwork, + &iwork_query, &liwork, &info); + return (int)work_query; +} + +static int query_liwork_dsyevd(int n) { + double work_query; + int iwork_query; + int lwork = -1, liwork = -1, info; + char jobz = 'V', uplo = 'L'; + dsyevd_(&jobz, &uplo, &n, NULL, &n, NULL, &work_query, &lwork, + &iwork_query, &liwork, &info); + return iwork_query; +} + +static int query_lwork_dgesvd(int m, int n) { + double work_query; + int lwork = -1, info; + char jobu = 'A', jobvt = 'A'; + dgesvd_(&jobu, &jobvt, &m, &n, NULL, &m, NULL, NULL, &m, NULL, &n, + &work_query, &lwork, &info); + return (int)work_query; +} + +/* -------------------------------------------------------------------------- + * Cholesky: eigenhhlo_dpotrf + * in[0] = A (n×n, double) + * out[0] = L or U (n×n, double, overwrites triangle) + * -------------------------------------------------------------------------- */ +void eigenhhlo_dpotrf(void* out, void** in, + const char* opaque, size_t opaque_len, + XlaCustomCallStatus* status) +{ + int n = parse_int_kv(opaque, "n", 1); + char uplo = parse_char_kv(opaque, "uplo", 'L'); + + double* A = (double*)in[0]; + double* outA = (double*)out; /* single result */ + + /* Copy input to output (LAPACK overwrites in-place) */ + memcpy(outA, A, n * n * sizeof(double)); + + int lda = n; + int info; + dpotrf_(&uplo, &n, outA, &lda, &info); + (void)info; /* TODO: propagate error via status */ +} + +/* -------------------------------------------------------------------------- + * LU: eigenhhlo_dgetrf + * in[0] = A (m×n, double) + * out[0] = LU (m×n, double) + * out[1] = pivots (min(m,n), int32) + * -------------------------------------------------------------------------- */ +void eigenhhlo_dgetrf(void* out, void** in, + const char* opaque, size_t opaque_len, + XlaCustomCallStatus* status) +{ + int m = parse_int_kv(opaque, "m", 1); + int n = parse_int_kv(opaque, "n", 1); + + double* A = (double*)in[0]; + void** outputs = (void**)out; + double* outLU = (double*)outputs[0]; + int32_t* outPiv = (int32_t*)outputs[1]; + + memcpy(outLU, A, m * n * sizeof(double)); + + int lda = m; + int info; + int* piv = (int*)malloc((m < n ? m : n) * sizeof(int)); + dgetrf_(&m, &n, outLU, &lda, piv, &info); + + /* Copy pivots (1-based in LAPACK, convert to 0-based) */ + int min_mn = (m < n) ? m : n; + for (int i = 0; i < min_mn; i++) { + outPiv[i] = (int32_t)(piv[i] - 1); + } + free(piv); + (void)info; +} + +/* -------------------------------------------------------------------------- + * QR factorization: eigenhhlo_dgeqrf + * in[0] = A (m×n, double) + * out[0] = A overwritten with R+reflectors (m×n, double) + * out[1] = tau (min(m,n), double) + * -------------------------------------------------------------------------- */ +void eigenhhlo_dgeqrf(void* out, void** in, + const char* opaque, size_t opaque_len, + XlaCustomCallStatus* status) +{ + int m = parse_int_kv(opaque, "m", 1); + int n = parse_int_kv(opaque, "n", 1); + + double* A = (double*)in[0]; + void** outputs = (void**)out; + double* outA = (double*)outputs[0]; + double* outTau = (double*)outputs[1]; + + memcpy(outA, A, m * n * sizeof(double)); + + int lda = m; + int min_mn = (m < n) ? m : n; + int lwork = query_lwork_dgeqrf(m, n); + double* work = (double*)malloc(lwork * sizeof(double)); + int info; + + dgeqrf_(&m, &n, outA, &lda, outTau, work, &lwork, &info); + free(work); + (void)info; +} + +/* -------------------------------------------------------------------------- + * Generate Q from QR: eigenhhlo_dorgqr + * in[0] = A (m×n, double, contains reflectors from dgeqrf) + * in[1] = tau (k, double) + * out[0] = Q (m×m, double) [full Q] + * -------------------------------------------------------------------------- */ +void eigenhhlo_dorgqr(void* out, void** in, + const char* opaque, size_t opaque_len, + XlaCustomCallStatus* status) +{ + int m = parse_int_kv(opaque, "m", 1); + int n = parse_int_kv(opaque, "n", 1); + int k = parse_int_kv(opaque, "k", 1); + + double* A = (double*)in[0]; + double* tau = (double*)in[1]; + double* outQ = (double*)out; /* single result */ + + /* Copy reflectors to output */ + memcpy(outQ, A, m * n * sizeof(double)); + + int lda = m; + int lwork = query_lwork_dorgqr(m, n, k); + double* work = (double*)malloc(lwork * sizeof(double)); + int info; + + dorgqr_(&m, &n, &k, outQ, &lda, tau, work, &lwork, &info); + free(work); + (void)info; +} + +/* -------------------------------------------------------------------------- + * Symmetric eigenvalue: eigenhhlo_dsyevd + * in[0] = A (n×n, double, symmetric) + * out[0] = eigenvalues (n, double) + * out[1] = eigenvectors (n×n, double) + * -------------------------------------------------------------------------- */ +void eigenhhlo_dsyevd(void* out, void** in, + const char* opaque, size_t opaque_len, + XlaCustomCallStatus* status) +{ + int n = parse_int_kv(opaque, "n", 1); + char uplo = parse_char_kv(opaque, "uplo", 'L'); + + double* A = (double*)in[0]; + void** outputs = (void**)out; + double* outW = (double*)outputs[0]; + double* outV = (double*)outputs[1]; + + /* Copy A to eigenvectors (LAPACK overwrites in-place) */ + memcpy(outV, A, n * n * sizeof(double)); + + int lda = n; + int lwork = query_lwork_dsyevd(n); + int liwork = query_liwork_dsyevd(n); + double* work = (double*)malloc(lwork * sizeof(double)); + int* iwork = (int*)malloc(liwork * sizeof(int)); + int info; + char jobz = 'V'; + + dsyevd_(&jobz, &uplo, &n, outV, &lda, outW, work, &lwork, + iwork, &liwork, &info); + free(work); + free(iwork); + (void)info; +} + +/* -------------------------------------------------------------------------- + * SVD: eigenhhlo_dgesvd + * in[0] = A (m×n, double) + * out[0] = U (m×m or m×min, double) + * out[1] = S (min(m,n), double) + * out[2] = Vt (n×n or min×n, double) + * -------------------------------------------------------------------------- */ +void eigenhhlo_dgesvd(void* out, void** in, + const char* opaque, size_t opaque_len, + XlaCustomCallStatus* status) +{ + int m = parse_int_kv(opaque, "m", 1); + int n = parse_int_kv(opaque, "n", 1); + char jobu = parse_char_kv(opaque, "jobu", 'A'); + char jobvt = parse_char_kv(opaque, "jobvt", 'A'); + + double* A = (double*)in[0]; + void** outputs = (void**)out; + double* outU = (double*)outputs[0]; + double* outS = (double*)outputs[1]; + double* outVt = (double*)outputs[2]; + + /* Copy A to workspace (LAPACK overwrites) */ + double* awork = (double*)malloc(m * n * sizeof(double)); + memcpy(awork, A, m * n * sizeof(double)); + + int lda = m; + int ldu = m; + int ldvt = n; + int lwork = query_lwork_dgesvd(m, n); + double* work = (double*)malloc(lwork * sizeof(double)); + int info; + + dgesvd_(&jobu, &jobvt, &m, &n, awork, &lda, outS, outU, &ldu, + outVt, &ldvt, work, &lwork, &info); + free(work); + free(awork); + (void)info; +} diff --git a/cbits/gpu/build.sh b/cbits/gpu/build.sh new file mode 100644 --- /dev/null +++ b/cbits/gpu/build.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# Build script for eigen-hhlo GPU custom-call library (cuSOLVER wrappers) + +set -e +NVCC=${NVCC:-nvcc} +ARCH=${ARCH:-sm_75} +OUTDIR="${1:-../../lib}" + +mkdir -p "$OUTDIR" + +echo "Building eigen-hhlo GPU custom-call library..." +echo " NVCC: $NVCC" +echo " ARCH: $ARCH" +echo " Output: $OUTDIR/libeigenhhlo_gpu.so" + +$NVCC -shared -o "$OUTDIR/libeigenhhlo_gpu.so" \ + -Xcompiler -fPIC \ + -gencode "arch=compute_${ARCH#sm_},code=$ARCH" \ + eigenhhlo_cusolver.cu \ + -lcusolver -lcudart -O3 + +echo "Build complete." +echo "Register at runtime with:" +echo " registerGpuCustomCall api \"$OUTDIR/libeigenhhlo_gpu.so\" \"eigenhhlo_dpotrf\"" diff --git a/cbits/gpu/eigenhhlo_cusolver.cu b/cbits/gpu/eigenhhlo_cusolver.cu new file mode 100644 --- /dev/null +++ b/cbits/gpu/eigenhhlo_cusolver.cu @@ -0,0 +1,381 @@ +/* eigen-hhlo GPU custom-call library — cuSOLVER wrappers. + * + * Conforms to XLA GPU custom-call ABI (API version 3): + * void target(CUstream stream, void** buffers, + * const char* opaque, size_t opaque_len, + * XlaCustomCallStatus* status); + * + * Buffer layout: [in0, in1, ..., out0, out1, ...] + * + * Compile: + * cd cbits/gpu && bash build.sh + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +/* -------------------------------------------------------------------------- + * XlaCustomCallStatus placeholder (opaque struct pointer) + * -------------------------------------------------------------------------- */ +struct XlaCustomCallStatus; + +/* -------------------------------------------------------------------------- + * Opaque string parser: key=value,... + * -------------------------------------------------------------------------- */ +static int parse_int_kv(const char* opaque, const char* key, int default_val) { + if (!opaque || !key) return default_val; + size_t keylen = strlen(key); + const char* p = opaque; + while (*p) { + if (strncmp(p, key, keylen) == 0 && p[keylen] == '=') { + p += keylen + 1; + return (int)strtol(p, NULL, 10); + } + while (*p && *p != ',') p++; + if (*p == ',') p++; + } + return default_val; +} + +static char parse_char_kv(const char* opaque, const char* key, char default_val) { + if (!opaque || !key) return default_val; + size_t keylen = strlen(key); + const char* p = opaque; + while (*p) { + if (strncmp(p, key, keylen) == 0 && p[keylen] == '=') { + p += keylen + 1; + return *p; + } + while (*p && *p != ',') p++; + if (*p == ',') p++; + } + return default_val; +} + +/* -------------------------------------------------------------------------- + * cuSOLVER handle cache: one handle per CUDA stream + * -------------------------------------------------------------------------- */ +static thread_local std::unordered_map g_handleCache; + +static cusolverDnHandle_t getHandle(CUstream stream) { + auto it = g_handleCache.find(stream); + if (it != g_handleCache.end()) return it->second; + cusolverDnHandle_t h; + cusolverDnCreate(&h); + cusolverDnSetStream(h, (cudaStream_t)stream); + g_handleCache[stream] = h; + return h; +} + +/* -------------------------------------------------------------------------- + * Helper: map char to cublasFillMode_t + * -------------------------------------------------------------------------- */ +static cublasFillMode_t charToFillMode(char uplo) { + return (uplo == 'U' || uplo == 'u') ? CUBLAS_FILL_MODE_UPPER + : CUBLAS_FILL_MODE_LOWER; +} + +/* -------------------------------------------------------------------------- + * Helper: check cuSOLVER status and print errors + * -------------------------------------------------------------------------- */ +static void checkCusolver(cusolverStatus_t status, const char* func) { + if (status != CUSOLVER_STATUS_SUCCESS) { + fprintf(stderr, "eigen-hhlo GPU: %s failed with status %d\n", func, status); + } +} + +/* ========================================================================== + * Cholesky: eigenhhlo_dpotrf + * buffers[0] = A (n×n, double, input) + * buffers[1] = L or U (n×n, double, output) + * ========================================================================== */ +extern "C" __attribute__((visibility("default"))) +void eigenhhlo_dpotrf(CUstream stream, void** buffers, + const char* opaque, size_t opaque_len, + XlaCustomCallStatus* status) +{ + int n = parse_int_kv(opaque, "n", 1); + char uplo_char = parse_char_kv(opaque, "uplo", 'L'); + cublasFillMode_t uplo = charToFillMode(uplo_char); + + const double* A = (const double*)buffers[0]; + double* outA = (double*)buffers[1]; + + cudaMemcpyAsync(outA, A, n * n * sizeof(double), cudaMemcpyDeviceToDevice, stream); + + cusolverDnHandle_t handle = getHandle(stream); + + int lwork; + checkCusolver(cusolverDnDpotrf_bufferSize(handle, uplo, n, outA, n, &lwork), + "cusolverDnDpotrf_bufferSize"); + + double* work; + cudaMallocAsync(&work, lwork * sizeof(double), stream); + int* devInfo; + cudaMallocAsync(&devInfo, sizeof(int), stream); + + checkCusolver(cusolverDnDpotrf(handle, uplo, n, outA, n, work, lwork, devInfo), + "cusolverDnDpotrf"); + + cudaFreeAsync(work, stream); + cudaFreeAsync(devInfo, stream); +} + +/* ========================================================================== + * LU: eigenhhlo_dgetrf + * buffers[0] = A (m×n, double, input) + * buffers[1] = LU (m×n, double, output) + * buffers[2] = pivots (min(m,n), int32, output) + * ========================================================================== */ +extern "C" __attribute__((visibility("default"))) +void eigenhhlo_dgetrf(CUstream stream, void** buffers, + const char* opaque, size_t opaque_len, + XlaCustomCallStatus* status) +{ + int m = parse_int_kv(opaque, "m", 1); + int n = parse_int_kv(opaque, "n", 1); + + const double* A = (const double*)buffers[0]; + double* outLU = (double*)buffers[1]; + int32_t* outPiv = (int32_t*)buffers[2]; + + cudaMemcpyAsync(outLU, A, m * n * sizeof(double), cudaMemcpyDeviceToDevice, stream); + + cusolverDnHandle_t handle = getHandle(stream); + + int lwork; + checkCusolver(cusolverDnDgetrf_bufferSize(handle, m, n, outLU, m, &lwork), + "cusolverDnDgetrf_bufferSize"); + + double* work; + cudaMallocAsync(&work, lwork * sizeof(double), stream); + int* devInfo; + cudaMallocAsync(&devInfo, sizeof(int), stream); + + checkCusolver(cusolverDnDgetrf(handle, m, n, outLU, m, work, (int*)outPiv, devInfo), + "cusolverDnDgetrf"); + + cudaFreeAsync(work, stream); + cudaFreeAsync(devInfo, stream); +} + +/* ========================================================================== + * QR factorization: eigenhhlo_dgeqrf + * buffers[0] = A (m×n, double, input) + * buffers[1] = A overwritten with R+reflectors (m×n, double, output) + * buffers[2] = tau (min(m,n), double, output) + * ========================================================================== */ +extern "C" __attribute__((visibility("default"))) +void eigenhhlo_dgeqrf(CUstream stream, void** buffers, + const char* opaque, size_t opaque_len, + XlaCustomCallStatus* status) +{ + int m = parse_int_kv(opaque, "m", 1); + int n = parse_int_kv(opaque, "n", 1); + + const double* A = (const double*)buffers[0]; + double* outA = (double*)buffers[1]; + double* outTau = (double*)buffers[2]; + + cudaMemcpyAsync(outA, A, m * n * sizeof(double), cudaMemcpyDeviceToDevice, stream); + + cusolverDnHandle_t handle = getHandle(stream); + + int lwork; + checkCusolver(cusolverDnDgeqrf_bufferSize(handle, m, n, outA, m, &lwork), + "cusolverDnDgeqrf_bufferSize"); + + double* work; + cudaMallocAsync(&work, lwork * sizeof(double), stream); + int* devInfo; + cudaMallocAsync(&devInfo, sizeof(int), stream); + + checkCusolver(cusolverDnDgeqrf(handle, m, n, outA, m, outTau, work, lwork, devInfo), + "cusolverDnDgeqrf"); + + cudaFreeAsync(work, stream); + cudaFreeAsync(devInfo, stream); +} + +/* ========================================================================== + * Generate Q from QR: eigenhhlo_dorgqr + * buffers[0] = A (m×n, double, contains reflectors from dgeqrf) + * buffers[1] = tau (k, double) + * buffers[2] = Q (m×m, double, output) + * ========================================================================== */ +extern "C" __attribute__((visibility("default"))) +void eigenhhlo_dorgqr(CUstream stream, void** buffers, + const char* opaque, size_t opaque_len, + XlaCustomCallStatus* status) +{ + int m = parse_int_kv(opaque, "m", 1); + int n = parse_int_kv(opaque, "n", 1); + int k = parse_int_kv(opaque, "k", 1); + + const double* A = (const double*)buffers[0]; + const double* tau = (const double*)buffers[1]; + double* outQ = (double*)buffers[2]; + + /* Initialize Q to identity by zeroing and setting diagonal, then copy reflectors */ + cudaMemsetAsync(outQ, 0, m * m * sizeof(double), stream); + /* Copy reflectors to the left part of Q */ + cudaMemcpyAsync(outQ, A, m * n * sizeof(double), cudaMemcpyDeviceToDevice, stream); + + cusolverDnHandle_t handle = getHandle(stream); + + int lwork; + checkCusolver(cusolverDnDorgqr_bufferSize(handle, m, m, k, outQ, m, tau, &lwork), + "cusolverDnDorgqr_bufferSize"); + + double* work; + cudaMallocAsync(&work, lwork * sizeof(double), stream); + int* devInfo; + cudaMallocAsync(&devInfo, sizeof(int), stream); + + checkCusolver(cusolverDnDorgqr(handle, m, m, k, outQ, m, tau, work, lwork, devInfo), + "cusolverDnDorgqr"); + + cudaFreeAsync(work, stream); + cudaFreeAsync(devInfo, stream); +} + +/* ========================================================================== + * Symmetric eigenvalue: eigenhhlo_dsyevd + * buffers[0] = A (n×n, double, symmetric, input) + * buffers[1] = eigenvalues (n, double, output) + * buffers[2] = eigenvectors (n×n, double, output) + * ========================================================================== */ +extern "C" __attribute__((visibility("default"))) +void eigenhhlo_dsyevd(CUstream stream, void** buffers, + const char* opaque, size_t opaque_len, + XlaCustomCallStatus* status) +{ + int n = parse_int_kv(opaque, "n", 1); + char uplo_char = parse_char_kv(opaque, "uplo", 'L'); + cublasFillMode_t uplo = charToFillMode(uplo_char); + + const double* A = (const double*)buffers[0]; + double* outW = (double*)buffers[1]; + double* outV = (double*)buffers[2]; + + cudaMemcpyAsync(outV, A, n * n * sizeof(double), cudaMemcpyDeviceToDevice, stream); + + cusolverDnHandle_t handle = getHandle(stream); + + int lwork; + checkCusolver(cusolverDnDsyevd_bufferSize(handle, CUSOLVER_EIG_MODE_VECTOR, uplo, n, outV, n, outW, &lwork), + "cusolverDnDsyevd_bufferSize"); + + double* work; + cudaMallocAsync(&work, lwork * sizeof(double), stream); + int* devInfo; + cudaMallocAsync(&devInfo, sizeof(int), stream); + + checkCusolver(cusolverDnDsyevd(handle, CUSOLVER_EIG_MODE_VECTOR, uplo, n, outV, n, outW, work, lwork, devInfo), + "cusolverDnDsyevd"); + + cudaFreeAsync(work, stream); + cudaFreeAsync(devInfo, stream); +} + +/* -------------------------------------------------------------------------- + * Transpose an m×n column-major matrix to an n×m column-major matrix. + * -------------------------------------------------------------------------- */ +__global__ void transpose_colmajor(double* out, const double* in, int m, int n) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < m * n) { + int i = idx % m; // row in input + int j = idx / m; // col in input + out[j + i * n] = in[i + j * m]; + } +} + +/* ========================================================================== + * SVD: eigenhhlo_dgesvd + * buffers[0] = A (m×n, double, input) + * buffers[1] = U (m×m or m×min, double, output) + * buffers[2] = S (min(m,n), double, output) + * buffers[3] = Vt (n×n or min×n, double, output) + * + * Workaround: cusolverDnDgesvd returns CUSOLVER_STATUS_INVALID_VALUE + * when m < n on some driver/GPU combinations. We transpose A, call + * dgesvd on A^T (where n >= m), and swap the U / Vt output buffers. + * ========================================================================== */ +extern "C" __attribute__((visibility("default"))) +void eigenhhlo_dgesvd(CUstream stream, void** buffers, + const char* opaque, size_t opaque_len, + XlaCustomCallStatus* status) +{ + int m = parse_int_kv(opaque, "m", 1); + int n = parse_int_kv(opaque, "n", 1); + char jobu = parse_char_kv(opaque, "jobu", 'A'); + char jobvt = parse_char_kv(opaque, "jobvt", 'A'); + + const double* A = (const double*)buffers[0]; + double* outU = (double*)buffers[1]; + double* outS = (double*)buffers[2]; + double* outVt = (double*)buffers[3]; + + /* cuSOLVER destroys A; allocate a temporary workspace copy */ + double* awork; + cudaMallocAsync(&awork, m * n * sizeof(double), stream); + cudaMemcpyAsync(awork, A, m * n * sizeof(double), cudaMemcpyDeviceToDevice, stream); + + cusolverDnHandle_t handle = getHandle(stream); + + if (m < n) { + /* Transpose A to A^T (n×m) so that n >= m */ + double* temp; + cudaMallocAsync(&temp, m * n * sizeof(double), stream); + int threads = 256; + int blocks = (m * n + threads - 1) / threads; + transpose_colmajor<<>>(temp, awork, m, n); + + int lwork; + checkCusolver(cusolverDnDgesvd_bufferSize(handle, n, m, &lwork), + "cusolverDnDgesvd_bufferSize"); + double* work; + cudaMallocAsync(&work, lwork * sizeof(double), stream); + int* devInfo; + cudaMallocAsync(&devInfo, sizeof(int), stream); + + /* Swap U and VT outputs: dgesvd on A^T writes U' to what we call VT + * and VT' to what we call U. ldu=n (rows of U'), ldvt=m (rows of VT'). */ + checkCusolver(cusolverDnDgesvd(handle, jobvt, jobu, n, m, temp, n, + outS, outVt, n, outU, m, + work, lwork, NULL, devInfo), + "cusolverDnDgesvd"); + + cudaFreeAsync(work, stream); + cudaFreeAsync(devInfo, stream); + cudaFreeAsync(temp, stream); + } else { + int k = (m < n) ? m : n; + int ldvt = (jobvt == 'S' || jobvt == 's') ? k : n; + + int lwork; + checkCusolver(cusolverDnDgesvd_bufferSize(handle, m, n, &lwork), + "cusolverDnDgesvd_bufferSize"); + double* work; + cudaMallocAsync(&work, lwork * sizeof(double), stream); + int* devInfo; + cudaMallocAsync(&devInfo, sizeof(int), stream); + + checkCusolver(cusolverDnDgesvd(handle, jobu, jobvt, m, n, awork, m, + outS, outU, m, outVt, ldvt, + work, lwork, NULL, devInfo), + "cusolverDnDgesvd"); + + cudaFreeAsync(work, stream); + cudaFreeAsync(devInfo, stream); + } + + cudaFreeAsync(awork, stream); +} diff --git a/eigen-hhlo.cabal b/eigen-hhlo.cabal new file mode 100644 --- /dev/null +++ b/eigen-hhlo.cabal @@ -0,0 +1,193 @@ +cabal-version: 3.0 +name: eigen-hhlo +version: 0.1.0.0 +synopsis: Dense linear algebra on HHLO (SVD, QR, eigenvalue, Cholesky, LU) +description: + eigen-hhlo provides singular value decomposition, QR decomposition, + symmetric eigenvalue decomposition, Cholesky factorization, and LU + decomposition as first-class operations within the HHLO ecosystem. + . + The library targets both CPU (via LAPACK custom calls) and GPU + (via cuSOLVER custom calls, future work) through PJRT, with no + host-side numerical fallback. + . + See the README for important runtime notes regarding CPU custom-call + registration limitations in the current PJRT CPU plugin. +homepage: https://github.com/overshiki/eigen-hhlo +bug-reports: https://github.com/overshiki/eigen-hhlo/issues +license: MIT +license-file: LICENSE +author: overshiki +maintainer: le.niu@hotmail.com +category: Math, Numeric, Machine Learning +tested-with: GHC == 9.6.7 +build-type: Simple +extra-doc-files: CHANGELOG.md + README.md +extra-source-files: + cbits/**/*.c + cbits/**/*.cu + cbits/**/*.sh + +source-repository head + type: git + location: https://github.com/overshiki/eigen-hhlo.git + +common warnings + ghc-options: -Wall + +library + import: warnings + hs-source-dirs: src + exposed-modules: + EigenHHLO.Core.Types + EigenHHLO.Core.BackendConfig + EigenHHLO.IR.Cholesky + EigenHHLO.IR.Eigenvalue + EigenHHLO.IR.LU + EigenHHLO.IR.QR + EigenHHLO.IR.SVD + EigenHHLO.EDSL.Decomposition + EigenHHLO.Runtime.Session + build-depends: + base >= 4.18.2 && < 5, + directory >= 1.3 && < 1.4, + hhlo >= 0.10.0.0 && < 0.11, + text >= 2.0 && < 2.2 + default-language: Haskell2010 + +executable eigen-hhlo + import: warnings + main-is: Main.hs + hs-source-dirs: app + build-depends: + base >= 4.18.2 && < 5, + eigen-hhlo, + hhlo >= 0.10.0.0 && < 0.11, + text >= 2.0 && < 2.2, + vector >= 0.13 && < 0.14 + default-language: Haskell2010 + +test-suite eigen-hhlo-test + import: warnings + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: test + other-modules: + EigenHHLO.Test.Cholesky + EigenHHLO.Test.Eigenvalue + EigenHHLO.Test.LU + EigenHHLO.Test.QR + EigenHHLO.Test.SVD + build-depends: + base >= 4.18.2 && < 5, + eigen-hhlo, + hhlo >= 0.10.0.0 && < 0.11, + text >= 2.0 && < 2.2, + vector >= 0.13 && < 0.14, + tasty >= 1.4 && < 1.6, + tasty-hunit >= 0.10 && < 0.11 + default-language: Haskell2010 + +flag examples + description: Build the example executables + default: False + manual: True + +executable example-cholesky + import: warnings + main-is: 01-cholesky.hs + hs-source-dirs: examples + build-depends: + base >= 4.18.2 && < 5, + eigen-hhlo, + hhlo >= 0.10.0.0 && < 0.11, + text >= 2.0 && < 2.2, + vector >= 0.13 && < 0.14 + default-language: Haskell2010 + if !flag(examples) + buildable: False + +executable example-svd + import: warnings + main-is: 02-svd.hs + hs-source-dirs: examples + build-depends: + base >= 4.18.2 && < 5, + eigen-hhlo, + hhlo >= 0.10.0.0 && < 0.11, + text >= 2.0 && < 2.2, + vector >= 0.13 && < 0.14 + default-language: Haskell2010 + if !flag(examples) + buildable: False + +executable example-qr + import: warnings + main-is: 03-qr.hs + hs-source-dirs: examples + build-depends: + base >= 4.18.2 && < 5, + eigen-hhlo, + hhlo >= 0.10.0.0 && < 0.11, + text >= 2.0 && < 2.2, + vector >= 0.13 && < 0.14 + default-language: Haskell2010 + if !flag(examples) + buildable: False + +executable example-eigenvalue + import: warnings + main-is: 04-eigenvalue.hs + hs-source-dirs: examples + build-depends: + base >= 4.18.2 && < 5, + eigen-hhlo, + hhlo >= 0.10.0.0 && < 0.11, + text >= 2.0 && < 2.2, + vector >= 0.13 && < 0.14 + default-language: Haskell2010 + if !flag(examples) + buildable: False + +executable example-lu + import: warnings + main-is: 05-lu.hs + hs-source-dirs: examples + build-depends: + base >= 4.18.2 && < 5, + eigen-hhlo, + hhlo >= 0.10.0.0 && < 0.11, + text >= 2.0 && < 2.2, + vector >= 0.13 && < 0.14 + default-language: Haskell2010 + if !flag(examples) + buildable: False + +executable example-pipeline + import: warnings + main-is: 06-pipeline.hs + hs-source-dirs: examples + build-depends: + base >= 4.18.2 && < 5, + eigen-hhlo, + hhlo >= 0.10.0.0 && < 0.11, + text >= 2.0 && < 2.2, + vector >= 0.13 && < 0.14 + default-language: Haskell2010 + if !flag(examples) + buildable: False + +executable example-gpu-setup + import: warnings + main-is: 07-gpu-setup.hs + hs-source-dirs: examples + build-depends: + base >= 4.18.2 && < 5, + eigen-hhlo, + hhlo >= 0.10.0.0 && < 0.11, + text >= 2.0 && < 2.2, + vector >= 0.13 && < 0.14 + default-language: Haskell2010 + if !flag(examples) + buildable: False diff --git a/examples/01-cholesky.hs b/examples/01-cholesky.hs new file mode 100644 --- /dev/null +++ b/examples/01-cholesky.hs @@ -0,0 +1,51 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Example 1: Cholesky factorization of a 3x3 symmetric positive-definite matrix. +-- +-- Build and run with: +-- cabal run example-cholesky --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 (chol) +import EigenHHLO.Runtime.Session (withEigenGPU) + +main :: IO () +main = withEigenGPU $ \esess -> do + putStrLn "=== Example 1: Cholesky Factorization (GPU) ===" + + let modu = moduleFromBuilder @'[3,3] @'F64 "main" + [ FuncArg "a" (TensorType [3,3] F64) ] + $ do a <- arg @'[3,3] @'F64 + chol esess a + + putStrLn "\nGenerated MLIR:" + putStrLn (T.unpack $ render modu) + + compiled <- compile (eigenSession esess) modu + + -- A = [[25, 15, 5], + -- [15, 18, 0], + -- [ 5, 0, 11]] (column-major) + let input = hostFromList @'[3,3] @'F64 + [25, 15, 5, 15, 18, 0, 5, 0, 11] + + putStrLn "\nInput matrix (column-major):" + print (hostToList input) + + result <- run (eigenSession esess) compiled input + :: IO (HostTensor '[3,3] 'F64) + + putStrLn "\nCholesky factor L (column-major):" + print (hostToList result) diff --git a/examples/02-svd.hs b/examples/02-svd.hs new file mode 100644 --- /dev/null +++ b/examples/02-svd.hs @@ -0,0 +1,59 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Example 2: Singular Value Decomposition of a 3x2 matrix. +-- +-- Demonstrates a multi-output custom call (U, S, Vt). +-- +-- Build and run with: +-- cabal run example-svd --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 (returnTuple3) +import HHLO.IR.Builder (arg, moduleFromBuilder3) +import HHLO.IR.Pretty (render) +import HHLO.Session (HostTensor, compile, hostFromList, hostToList, run) + +import EigenHHLO.Core.Types (eigenSession) +import EigenHHLO.EDSL.Decomposition (svd) +import EigenHHLO.Runtime.Session (withEigenGPU) + +main :: IO () +main = withEigenGPU $ \esess -> do + putStrLn "=== Example 2: Singular Value Decomposition (GPU) ===" + + let modu = moduleFromBuilder3 @'[3,2] @'F64 @'[2] @'F64 @'[2,2] @'F64 "main" + [ FuncArg "a" (TensorType [3,2] F64) ] + $ do a <- arg @'[3,2] @'F64 + (u, s, vt) <- svd esess a + returnTuple3 u s vt + + putStrLn "\nGenerated MLIR:" + putStrLn (T.unpack $ render modu) + + compiled <- compile (eigenSession esess) modu + + -- A = [[3, 0], + -- [0, 2], + -- [0, 0]] (3x2, column-major) + let input = hostFromList @'[3,2] @'F64 + [3, 0, 0, 0, 2, 0] + + putStrLn "\nInput matrix (column-major):" + print (hostToList input) + + (uResult, sResult, vtResult) <- run (eigenSession esess) compiled input + :: IO (HostTensor '[3,3] 'F64, HostTensor '[2] 'F64, HostTensor '[2,2] 'F64) + + putStrLn "\nU (3x3, column-major):" + print (hostToList uResult) + putStrLn "S (2,):" + print (hostToList sResult) + putStrLn "Vt (2x2, column-major):" + print (hostToList vtResult) diff --git a/examples/03-qr.hs b/examples/03-qr.hs new file mode 100644 --- /dev/null +++ b/examples/03-qr.hs @@ -0,0 +1,57 @@ +{-# 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) diff --git a/examples/04-eigenvalue.hs b/examples/04-eigenvalue.hs new file mode 100644 --- /dev/null +++ b/examples/04-eigenvalue.hs @@ -0,0 +1,57 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Example 4: Symmetric eigenvalue decomposition. +-- +-- Returns eigenvalues and eigenvectors: A = V · Λ · Vᵀ. +-- +-- Build and run with: +-- cabal run example-eigenvalue --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 (eig) +import EigenHHLO.Runtime.Session (withEigenGPU) + +main :: IO () +main = withEigenGPU $ \esess -> do + putStrLn "=== Example 4: Symmetric Eigenvalue Decomposition (GPU) ===" + + let modu = moduleFromBuilder2 @'[3] @'F64 @'[3,3] @'F64 "main" + [ FuncArg "a" (TensorType [3,3] F64) ] + $ do a <- arg @'[3,3] @'F64 + (w, v) <- eig esess a + returnTuple2 w v + + putStrLn "\nGenerated MLIR:" + putStrLn (T.unpack $ render modu) + + compiled <- compile (eigenSession esess) modu + + -- A = [[4, 2, 1], + -- [2, 5, 3], + -- [1, 3, 6]] (column-major, symmetric) + let input = hostFromList @'[3,3] @'F64 + [4, 2, 1, 2, 5, 3, 1, 3, 6] + + putStrLn "\nInput matrix (column-major):" + print (hostToList input) + + (wResult, vResult) <- run (eigenSession esess) compiled input + :: IO (HostTensor '[3] 'F64, HostTensor '[3,3] 'F64) + + putStrLn "\nEigenvalues (3,):" + print (hostToList wResult) + putStrLn "\nEigenvectors (3x3, column-major):" + print (hostToList vResult) diff --git a/examples/05-lu.hs b/examples/05-lu.hs new file mode 100644 --- /dev/null +++ b/examples/05-lu.hs @@ -0,0 +1,57 @@ +{-# 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) diff --git a/examples/06-pipeline.hs b/examples/06-pipeline.hs new file mode 100644 --- /dev/null +++ b/examples/06-pipeline.hs @@ -0,0 +1,58 @@ +{-# 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) diff --git a/examples/07-gpu-setup.hs b/examples/07-gpu-setup.hs new file mode 100644 --- /dev/null +++ b/examples/07-gpu-setup.hs @@ -0,0 +1,46 @@ +{-# 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'." diff --git a/src/EigenHHLO/Core/BackendConfig.hs b/src/EigenHHLO/Core/BackendConfig.hs new file mode 100644 --- /dev/null +++ b/src/EigenHHLO/Core/BackendConfig.hs @@ -0,0 +1,50 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Build opaque 'backend_config' strings for XLA custom calls. +module EigenHHLO.Core.BackendConfig + ( svdConfig + , qrConfig + , eigConfig + , cholConfig + , luConfig + ) where + +import Data.Text (Text) +import qualified Data.Text as T + +-- | Build backend_config string for SVD. +svdConfig :: Integer -> Integer -> Char -> Char -> Text +svdConfig m n jobu jobvt = T.intercalate "," $ + [ "m=" <> T.pack (show m) + , "n=" <> T.pack (show n) + , "jobu=" <> T.pack [jobu] + , "jobvt=" <> T.pack [jobvt] + ] + +-- | Build backend_config string for QR factorization. +qrConfig :: Integer -> Integer -> Text +qrConfig m n = T.intercalate "," $ + [ "m=" <> T.pack (show m) + , "n=" <> T.pack (show n) + ] + +-- | Build backend_config string for symmetric eigenvalue. +eigConfig :: Integer -> Char -> Text +eigConfig n uplo = T.intercalate "," $ + [ "n=" <> T.pack (show n) + , "uplo=" <> T.pack [uplo] + ] + +-- | Build backend_config string for Cholesky. +cholConfig :: Integer -> Char -> Text +cholConfig n uplo = T.intercalate "," $ + [ "n=" <> T.pack (show n) + , "uplo=" <> T.pack [uplo] + ] + +-- | Build backend_config string for LU. +luConfig :: Integer -> Integer -> Text +luConfig m n = T.intercalate "," $ + [ "m=" <> T.pack (show m) + , "n=" <> T.pack (show n) + ] diff --git a/src/EigenHHLO/Core/Types.hs b/src/EigenHHLO/Core/Types.hs new file mode 100644 --- /dev/null +++ b/src/EigenHHLO/Core/Types.hs @@ -0,0 +1,23 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} + +-- | Core types for eigen-hhlo: backend tagging and session metadata. +module EigenHHLO.Core.Types + ( BackendType(..) + , EigenSession(..) + ) where + +import HHLO.Session (Session) + +-- | Which PJRT backend is in use. +data BackendType = CPU | GPU + deriving stock (Eq, Show) + +-- | A session extended with backend metadata. +-- +-- Use 'withEigenCPU' or 'withEigenGPU' (future) from +-- 'EigenHHLO.Runtime.Session' to construct. +data EigenSession = EigenSession + { eigenSession :: !Session + , eigenBackend :: !BackendType + } diff --git a/src/EigenHHLO/EDSL/Decomposition.hs b/src/EigenHHLO/EDSL/Decomposition.hs new file mode 100644 --- /dev/null +++ b/src/EigenHHLO/EDSL/Decomposition.hs @@ -0,0 +1,68 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE TypeApplications #-} + +-- | User-facing EDSL API for dense linear algebra decompositions. +module EigenHHLO.EDSL.Decomposition + ( svd + , qr + , q + , eig + , chol + , lu + ) where + +import GHC.TypeLits (KnownNat) +import HHLO.Core.Types (DType(..)) +import HHLO.IR.Builder (Builder, Tensor) + +import EigenHHLO.Core.Types (EigenSession) +import qualified EigenHHLO.IR.Cholesky as IR +import qualified EigenHHLO.IR.Eigenvalue as IR +import qualified EigenHHLO.IR.LU as IR +import qualified EigenHHLO.IR.QR as IR +import qualified EigenHHLO.IR.SVD as IR + +-- | Singular value decomposition. +svd :: forall m n k. (KnownNat m, KnownNat n, KnownNat k) + => EigenSession + -> Tensor '[m, n] 'F64 + -> Builder (Tensor '[m, k] 'F64, Tensor '[k] 'F64, Tensor '[k, n] 'F64) +svd = IR.svdBuilder + +-- | QR factorization: A = Q · R. +-- Returns (A overwritten with reflectors, tau). +qr :: forall m n k. (KnownNat m, KnownNat n, KnownNat k) + => EigenSession + -> Tensor '[m, n] 'F64 + -> Builder (Tensor '[m, n] 'F64, Tensor '[k] 'F64) +qr = IR.qrBuilder + +-- | Generate explicit Q from QR reflectors. +q :: forall m n k. (KnownNat m, KnownNat n, KnownNat k) + => EigenSession + -> Tensor '[m, n] 'F64 + -> Tensor '[k] 'F64 + -> Builder (Tensor '[m, m] 'F64) +q = IR.qBuilder + +-- | Symmetric eigenvalue decomposition. +eig :: forall n. (KnownNat n) + => EigenSession + -> Tensor '[n, n] 'F64 + -> Builder (Tensor '[n] 'F64, Tensor '[n, n] 'F64) +eig = IR.eigBuilder + +-- | Cholesky factorization. +chol :: forall n. (KnownNat n) + => EigenSession + -> Tensor '[n, n] 'F64 + -> Builder (Tensor '[n, n] 'F64) +chol = IR.cholBuilder + +-- | LU factorization with partial pivoting. +lu :: forall m n k. (KnownNat m, KnownNat n, KnownNat k) + => EigenSession + -> Tensor '[m, n] 'F64 + -> Builder (Tensor '[m, n] 'F64, Tensor '[k] 'I32) +lu = IR.luBuilder diff --git a/src/EigenHHLO/IR/Cholesky.hs b/src/EigenHHLO/IR/Cholesky.hs new file mode 100644 --- /dev/null +++ b/src/EigenHHLO/IR/Cholesky.hs @@ -0,0 +1,46 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Typed 'stablehlo.custom_call' builder for Cholesky factorization. +module EigenHHLO.IR.Cholesky + ( cholBuilder + ) where + +import Data.Proxy (Proxy(..)) +import GHC.TypeLits (KnownNat, natVal) + +import HHLO.Core.Types (DType(..)) +import HHLO.EDSL.Ops (customCallRaw) +import HHLO.IR.AST (TensorType(..)) +import HHLO.IR.Builder (Builder, Tensor(..), tensorValue) + +import EigenHHLO.Core.BackendConfig (cholConfig) +import EigenHHLO.Core.Types (BackendType(..), EigenSession(..)) + +-- | Cholesky factorization via LAPACK @dpotrf@. +-- +-- Factorizes a symmetric positive-definite matrix A into L · Lᵀ +-- (lower triangular). The @uplo@ parameter is taken from the +-- 'BackendConfig' (currently hard-coded to @'L'@). +cholBuilder :: forall n. KnownNat n + => EigenSession + -> Tensor '[n, n] 'F64 + -> Builder (Tensor '[n, n] 'F64) +cholBuilder esess a = + let n = fromIntegral (natVal (Proxy @n)) :: Integer + cfg = cholConfig n 'L' + apiVer = case eigenBackend esess of + CPU -> 2 + GPU -> 3 + inType = TensorType [n, n] F64 + outType = TensorType [n, n] F64 + vids = [tensorValue a] + inTypes = [inType] + in do + vidsRes <- customCallRaw "eigenhhlo_dpotrf" vids inTypes cfg False apiVer [outType] + case vidsRes of + [v1] -> return (Tensor v1) + _ -> error "cholBuilder: expected exactly one result" diff --git a/src/EigenHHLO/IR/Eigenvalue.hs b/src/EigenHHLO/IR/Eigenvalue.hs new file mode 100644 --- /dev/null +++ b/src/EigenHHLO/IR/Eigenvalue.hs @@ -0,0 +1,44 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Typed 'stablehlo.custom_call' builder for symmetric eigenvalue decomposition. +module EigenHHLO.IR.Eigenvalue + ( eigBuilder + ) where + +import Data.Proxy (Proxy(..)) +import GHC.TypeLits (KnownNat, natVal) + +import HHLO.Core.Types (DType(..)) +import HHLO.EDSL.Ops (customCallRaw) +import HHLO.IR.AST (TensorType(..)) +import HHLO.IR.Builder (Builder, Tensor(..), tensorValue) + +import EigenHHLO.Core.BackendConfig (eigConfig) +import EigenHHLO.Core.Types (BackendType(..), EigenSession(..)) + +-- | Symmetric eigenvalue decomposition: A = V · Λ · V^T. +-- Returns (eigenvalues, eigenvectors). +eigBuilder :: forall n. KnownNat n + => EigenSession + -> Tensor '[n, n] 'F64 + -> Builder (Tensor '[n] 'F64, Tensor '[n, n] 'F64) +eigBuilder esess a = + let n = fromIntegral (natVal (Proxy @n)) :: Integer + cfg = eigConfig n 'L' + apiVer = case eigenBackend esess of + CPU -> 2 + GPU -> 3 + inType = TensorType [n, n] F64 + outType1 = TensorType [n] F64 + outType2 = TensorType [n, n] F64 + vids = [tensorValue a] + inTypes = [inType] + in do + vidsRes <- customCallRaw "eigenhhlo_dsyevd" vids inTypes cfg False apiVer [outType1, outType2] + case vidsRes of + [v1, v2] -> return (Tensor v1, Tensor v2) + _ -> error "eigBuilder: expected exactly two results" diff --git a/src/EigenHHLO/IR/LU.hs b/src/EigenHHLO/IR/LU.hs new file mode 100644 --- /dev/null +++ b/src/EigenHHLO/IR/LU.hs @@ -0,0 +1,48 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Typed 'stablehlo.custom_call' builder for LU factorization with partial pivoting. +module EigenHHLO.IR.LU + ( luBuilder + ) where + +import Data.Proxy (Proxy(..)) +import GHC.TypeLits (KnownNat, natVal) + +import HHLO.Core.Types (DType(..)) +import HHLO.EDSL.Ops (customCallRaw) +import HHLO.IR.AST (TensorType(..)) +import HHLO.IR.Builder (Builder, Tensor(..), tensorValue) + +import EigenHHLO.Core.BackendConfig (luConfig) +import EigenHHLO.Core.Types (BackendType(..), EigenSession(..)) + +-- | LU factorization with partial pivoting via LAPACK @dgetrf@. +-- +-- Returns @(LU, pivots)@ where @LU@ contains both L and U packed +-- together, and @pivots@ are the 0-based permutation indices. +luBuilder :: forall m n k. (KnownNat m, KnownNat n, KnownNat k) + => EigenSession + -> Tensor '[m, n] 'F64 + -> Builder (Tensor '[m, n] 'F64, Tensor '[k] 'I32) +luBuilder esess a = + let m = fromIntegral (natVal (Proxy @m)) :: Integer + n = fromIntegral (natVal (Proxy @n)) :: Integer + k = fromIntegral (natVal (Proxy @k)) :: Integer + cfg = luConfig m n + apiVer = case eigenBackend esess of + CPU -> 2 + GPU -> 3 + inType = TensorType [m, n] F64 + outType1 = TensorType [m, n] F64 + outType2 = TensorType [k] I32 + vids = [tensorValue a] + inTypes = [inType] + in do + vidsRes <- customCallRaw "eigenhhlo_dgetrf" vids inTypes cfg False apiVer [outType1, outType2] + case vidsRes of + [v1, v2] -> return (Tensor v1, Tensor v2) + _ -> error "luBuilder: expected exactly two results" diff --git a/src/EigenHHLO/IR/QR.hs b/src/EigenHHLO/IR/QR.hs new file mode 100644 --- /dev/null +++ b/src/EigenHHLO/IR/QR.hs @@ -0,0 +1,73 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Typed 'stablehlo.custom_call' builders for QR factorization and explicit Q generation. +module EigenHHLO.IR.QR + ( qrBuilder + , qBuilder + ) where + +import Data.Proxy (Proxy(..)) +import GHC.TypeLits (KnownNat, natVal) + +import HHLO.Core.Types (DType(..)) +import HHLO.EDSL.Ops (customCallRaw) +import HHLO.IR.AST (TensorType(..)) +import HHLO.IR.Builder (Builder, Tensor(..), tensorValue) +import qualified Data.Text as T + +import EigenHHLO.Core.BackendConfig (qrConfig) +import EigenHHLO.Core.Types (BackendType(..), EigenSession(..)) + +-- | QR factorization: A = Q · R. +-- Returns (A overwritten with R+reflectors, tau). +qrBuilder :: forall m n k. (KnownNat m, KnownNat n, KnownNat k) + => EigenSession + -> Tensor '[m, n] 'F64 + -> Builder (Tensor '[m, n] 'F64, Tensor '[k] 'F64) +qrBuilder esess a = + let m = fromIntegral (natVal (Proxy @m)) :: Integer + n = fromIntegral (natVal (Proxy @n)) :: Integer + k = fromIntegral (natVal (Proxy @k)) :: Integer + cfg = qrConfig m n + apiVer = case eigenBackend esess of + CPU -> 2 + GPU -> 3 + inType = TensorType [m, n] F64 + outType1 = TensorType [m, n] F64 + outType2 = TensorType [k] F64 + vids = [tensorValue a] + inTypes = [inType] + in do + vidsRes <- customCallRaw "eigenhhlo_dgeqrf" vids inTypes cfg False apiVer [outType1, outType2] + case vidsRes of + [v1, v2] -> return (Tensor v1, Tensor v2) + _ -> error "qrBuilder: expected exactly two results" + +-- | Generate explicit Q matrix from QR reflectors. +qBuilder :: forall m n k. (KnownNat m, KnownNat n, KnownNat k) + => EigenSession + -> Tensor '[m, n] 'F64 -- ^ reflectors from qrBuilder + -> Tensor '[k] 'F64 -- ^ tau from qrBuilder + -> Builder (Tensor '[m, m] 'F64) +qBuilder esess a tau = + let m = fromIntegral (natVal (Proxy @m)) :: Integer + n = fromIntegral (natVal (Proxy @n)) :: Integer + k = fromIntegral (natVal (Proxy @k)) :: Integer + cfg = qrConfig m n <> ",k=" <> (T.pack $ show k) + apiVer = case eigenBackend esess of + CPU -> 2 + GPU -> 3 + inType1 = TensorType [m, n] F64 + inType2 = TensorType [k] F64 + outType = TensorType [m, m] F64 + vids = [tensorValue a, tensorValue tau] + inTypes = [inType1, inType2] + in do + vidsRes <- customCallRaw "eigenhhlo_dorgqr" vids inTypes cfg False apiVer [outType] + case vidsRes of + [v1] -> return (Tensor v1) + _ -> error "qBuilder: expected exactly one result" diff --git a/src/EigenHHLO/IR/SVD.hs b/src/EigenHHLO/IR/SVD.hs new file mode 100644 --- /dev/null +++ b/src/EigenHHLO/IR/SVD.hs @@ -0,0 +1,47 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +-- | Typed 'stablehlo.custom_call' builder for singular value decomposition. +module EigenHHLO.IR.SVD + ( svdBuilder + ) where + +import Data.Proxy (Proxy(..)) +import GHC.TypeLits (KnownNat, natVal) + +import HHLO.Core.Types (DType(..)) +import HHLO.EDSL.Ops (customCallRaw) +import HHLO.IR.AST (TensorType(..)) +import HHLO.IR.Builder (Builder, Tensor(..), tensorValue) + +import EigenHHLO.Core.BackendConfig (svdConfig) +import EigenHHLO.Core.Types (BackendType(..), EigenSession(..)) + +-- | Singular value decomposition: A = U · diag(S) · Vt. +-- Returns (U, S, Vt). +svdBuilder :: forall m n k. (KnownNat m, KnownNat n, KnownNat k) + => EigenSession + -> Tensor '[m, n] 'F64 + -> Builder (Tensor '[m, k] 'F64, Tensor '[k] 'F64, Tensor '[k, n] 'F64) +svdBuilder esess a = + let m = fromIntegral (natVal (Proxy @m)) :: Integer + n = fromIntegral (natVal (Proxy @n)) :: Integer + k = fromIntegral (natVal (Proxy @k)) :: Integer + cfg = svdConfig m n 'S' 'S' + apiVer = case eigenBackend esess of + CPU -> 2 + GPU -> 3 + inType = TensorType [m, n] F64 + outType1 = TensorType [m, k] F64 + outType2 = TensorType [k] F64 + outType3 = TensorType [k, n] F64 + vids = [tensorValue a] + inTypes = [inType] + in do + vidsRes <- customCallRaw "eigenhhlo_dgesvd" vids inTypes cfg False apiVer [outType1, outType2, outType3] + case vidsRes of + [v1, v2, v3] -> return (Tensor v1, Tensor v2, Tensor v3) + _ -> error "svdBuilder: expected exactly three results" diff --git a/src/EigenHHLO/Runtime/Session.hs b/src/EigenHHLO/Runtime/Session.hs new file mode 100644 --- /dev/null +++ b/src/EigenHHLO/Runtime/Session.hs @@ -0,0 +1,68 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Session lifecycle management for eigen-hhlo backends. +module EigenHHLO.Runtime.Session + ( withEigenCPU + , withEigenGPU + ) where + +import Data.Char (toLower, toUpper) +import System.Directory (doesFileExist) +import System.Environment (lookupEnv) + +import HHLO.Runtime.CustomCall (loadCustomCallLibrary, registerGpuCustomCall) +import HHLO.Session (Session(..), withCPU, withGPU) + +import EigenHHLO.Core.Types (BackendType(..), EigenSession(..)) + +-- | Resolve the path to a custom-call shared library. +-- +-- Priority: +-- 1. @EIGENHHLO__LIB@ environment variable +-- 2. @lib/libeigenhhlo_.so@ relative to CWD +-- 3. Runtime error with instructions +getCustomCallLibPath :: String -> FilePath -> IO FilePath +getCustomCallLibPath backend defaultName = do + mEnv <- lookupEnv ("EIGENHHLO_" ++ map toUpper backend ++ "_LIB") + case mEnv of + Just p -> return p + Nothing -> do + let defaultPath = "lib/" ++ defaultName + exists <- doesFileExist defaultPath + if exists + then return defaultPath + else error $ unlines + [ "eigen-hhlo " ++ backend ++ " custom-call library not found at: " ++ defaultPath + , "" + , "To fix this, either:" + , " 1. Build the library:" + , " cd cbits/" ++ map toLower backend ++ " && bash build.sh" + , " 2. Set the environment variable:" + , " export EIGENHHLO_" ++ map toUpper backend ++ "_LIB=/path/to/" ++ defaultName + ] + +-- | Create an eigen-hhlo session on the CPU backend. +-- Loads the CPU custom-call library via RTLD_GLOBAL. +-- +-- The library path is resolved via 'getCustomCallLibPath'. +withEigenCPU :: (EigenSession -> IO a) -> IO a +withEigenCPU action = withCPU $ \sess -> do + libPath <- getCustomCallLibPath "CPU" "libeigenhhlo_cpu.so" + loadCustomCallLibrary libPath + action (EigenSession sess CPU) + +-- | Create an eigen-hhlo session on the GPU backend. +-- Registers all decomposition symbols with the PJRT CUDA plugin. +-- +-- The library path is resolved via 'getCustomCallLibPath'. +withEigenGPU :: (EigenSession -> IO a) -> IO a +withEigenGPU action = withGPU $ \sess -> do + libPath <- getCustomCallLibPath "GPU" "libeigenhhlo_gpu.so" + -- Register all symbols with the PJRT GPU plugin + registerGpuCustomCall (sessionApi sess) libPath "eigenhhlo_dgesvd" + registerGpuCustomCall (sessionApi sess) libPath "eigenhhlo_dgeqrf" + registerGpuCustomCall (sessionApi sess) libPath "eigenhhlo_dorgqr" + registerGpuCustomCall (sessionApi sess) libPath "eigenhhlo_dsyevd" + registerGpuCustomCall (sessionApi sess) libPath "eigenhhlo_dpotrf" + registerGpuCustomCall (sessionApi sess) libPath "eigenhhlo_dgetrf" + action (EigenSession sess GPU) diff --git a/test/EigenHHLO/Test/Cholesky.hs b/test/EigenHHLO/Test/Cholesky.hs new file mode 100644 --- /dev/null +++ b/test/EigenHHLO/Test/Cholesky.hs @@ -0,0 +1,35 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module EigenHHLO.Test.Cholesky (tests) where + +import qualified Data.Text as T +import Foreign.Ptr (nullPtr) +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.Core.Types (DType(..)) +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder (arg, moduleFromBuilder) +import HHLO.IR.Pretty (render) +import HHLO.Runtime.PJRT.Types (PJRTApi(..), PJRTClient(..), PJRTDevice(..)) +import HHLO.Session (sessionFrom) + +import EigenHHLO.IR.Cholesky (cholBuilder) +import EigenHHLO.Core.Types (BackendType(..), EigenSession(..)) + +dummySess :: EigenSession +dummySess = EigenSession (sessionFrom (PJRTApi nullPtr) (PJRTClient nullPtr) (PJRTDevice nullPtr)) CPU + +tests :: TestTree +tests = testGroup "Cholesky" + [ testCase "MLIR generation" $ do + let modu = moduleFromBuilder @'[2,2] @'F64 "main" + [ FuncArg "a" (TensorType [2,2] F64) ] + $ do a <- arg @'[2,2] @'F64 + cholBuilder dummySess a + let mlir = render modu + assertBool "contains dpotrf" ("eigenhhlo_dpotrf" `T.isInfixOf` mlir) + assertBool "contains api_version = 2" ("api_version = 2" `T.isInfixOf` mlir) + ] diff --git a/test/EigenHHLO/Test/Eigenvalue.hs b/test/EigenHHLO/Test/Eigenvalue.hs new file mode 100644 --- /dev/null +++ b/test/EigenHHLO/Test/Eigenvalue.hs @@ -0,0 +1,35 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module EigenHHLO.Test.Eigenvalue (tests) where + +import qualified Data.Text as T +import Foreign.Ptr (nullPtr) +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.Core.Types (DType(..)) +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder (arg, moduleFromBuilder) +import HHLO.IR.Pretty (render) +import HHLO.Runtime.PJRT.Types (PJRTApi(..), PJRTClient(..), PJRTDevice(..)) +import HHLO.Session (sessionFrom) + +import EigenHHLO.IR.Eigenvalue (eigBuilder) +import EigenHHLO.Core.Types (BackendType(..), EigenSession(..)) + +dummySess :: EigenSession +dummySess = EigenSession (sessionFrom (PJRTApi nullPtr) (PJRTClient nullPtr) (PJRTDevice nullPtr)) CPU + +tests :: TestTree +tests = testGroup "Eigenvalue" + [ testCase "MLIR generation contains dsyevd" $ do + let modu = moduleFromBuilder @'[2] @'F64 "main" + [ FuncArg "a" (TensorType [2,2] F64) ] + $ do a <- arg @'[2,2] @'F64 + (w, _v) <- eigBuilder dummySess a + return w + let mlir = render modu + assertBool "contains dsyevd" ("eigenhhlo_dsyevd" `T.isInfixOf` mlir) + ] diff --git a/test/EigenHHLO/Test/LU.hs b/test/EigenHHLO/Test/LU.hs new file mode 100644 --- /dev/null +++ b/test/EigenHHLO/Test/LU.hs @@ -0,0 +1,35 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module EigenHHLO.Test.LU (tests) where + +import qualified Data.Text as T +import Foreign.Ptr (nullPtr) +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.Core.Types (DType(..)) +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder (arg, moduleFromBuilder) +import HHLO.IR.Pretty (render) +import HHLO.Runtime.PJRT.Types (PJRTApi(..), PJRTClient(..), PJRTDevice(..)) +import HHLO.Session (sessionFrom) + +import EigenHHLO.IR.LU (luBuilder) +import EigenHHLO.Core.Types (BackendType(..), EigenSession(..)) + +dummySess :: EigenSession +dummySess = EigenSession (sessionFrom (PJRTApi nullPtr) (PJRTClient nullPtr) (PJRTDevice nullPtr)) CPU + +tests :: TestTree +tests = testGroup "LU" + [ testCase "MLIR generation contains dgetrf" $ do + let modu = moduleFromBuilder @'[2,2] @'F64 "main" + [ FuncArg "a" (TensorType [2,2] F64) ] + $ do a <- arg @'[2,2] @'F64 + (lu, _piv) <- luBuilder @2 @2 @2 dummySess a + return lu + let mlir = render modu + assertBool "contains dgetrf" ("eigenhhlo_dgetrf" `T.isInfixOf` mlir) + ] diff --git a/test/EigenHHLO/Test/QR.hs b/test/EigenHHLO/Test/QR.hs new file mode 100644 --- /dev/null +++ b/test/EigenHHLO/Test/QR.hs @@ -0,0 +1,35 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module EigenHHLO.Test.QR (tests) where + +import qualified Data.Text as T +import Foreign.Ptr (nullPtr) +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.Core.Types (DType(..)) +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder (arg, moduleFromBuilder) +import HHLO.IR.Pretty (render) +import HHLO.Runtime.PJRT.Types (PJRTApi(..), PJRTClient(..), PJRTDevice(..)) +import HHLO.Session (sessionFrom) + +import EigenHHLO.IR.QR (qrBuilder) +import EigenHHLO.Core.Types (BackendType(..), EigenSession(..)) + +dummySess :: EigenSession +dummySess = EigenSession (sessionFrom (PJRTApi nullPtr) (PJRTClient nullPtr) (PJRTDevice nullPtr)) CPU + +tests :: TestTree +tests = testGroup "QR" + [ testCase "MLIR generation contains dgeqrf" $ do + let modu = moduleFromBuilder @'[2,2] @'F64 "main" + [ FuncArg "a" (TensorType [2,2] F64) ] + $ do a <- arg @'[2,2] @'F64 + (r, _tau) <- qrBuilder @2 @2 @2 dummySess a + return r + let mlir = render modu + assertBool "contains dgeqrf" ("eigenhhlo_dgeqrf" `T.isInfixOf` mlir) + ] diff --git a/test/EigenHHLO/Test/SVD.hs b/test/EigenHHLO/Test/SVD.hs new file mode 100644 --- /dev/null +++ b/test/EigenHHLO/Test/SVD.hs @@ -0,0 +1,47 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module EigenHHLO.Test.SVD (tests) where + +import qualified Data.Text as T +import Foreign.Ptr (nullPtr) +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.Core.Types (DType(..)) +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder (arg, moduleFromBuilder) +import HHLO.IR.Pretty (render) +import HHLO.Runtime.PJRT.Types (PJRTApi(..), PJRTClient(..), PJRTDevice(..)) +import HHLO.Session (sessionFrom) + +import EigenHHLO.IR.SVD (svdBuilder) +import EigenHHLO.Core.Types (BackendType(..), EigenSession(..)) + +dummySess :: EigenSession +dummySess = EigenSession (sessionFrom (PJRTApi nullPtr) (PJRTClient nullPtr) (PJRTDevice nullPtr)) CPU + +tests :: TestTree +tests = testGroup "SVD" + [ testCase "MLIR generation contains dgesvd" $ do + let modu = moduleFromBuilder @'[2,2] @'F64 "main" + [ FuncArg "a" (TensorType [2,3] F64) ] + $ do a <- arg @'[2,3] @'F64 + (u, _s, _vt) <- svdBuilder @2 @3 @2 dummySess a + return u + let mlir = render modu + assertBool "contains dgesvd" ("eigenhhlo_dgesvd" `T.isInfixOf` mlir) + + , testCase "MLIR output shapes are correct for thin SVD" $ do + let modu = moduleFromBuilder @'[3] @'F64 "main" + [ FuncArg "a" (TensorType [4,3] F64) ] + $ do a <- arg @'[4,3] @'F64 + (_u, s, _vt) <- svdBuilder @4 @3 @3 dummySess a + return s + let mlir = render modu + -- U should be 4x3, S should be 3, Vt should be 3x3 + assertBool "U shape 4x3" ("tensor<4x3xf64>" `T.isInfixOf` mlir) + assertBool "S shape 3" ("tensor<3xf64>" `T.isInfixOf` mlir) + assertBool "Vt shape 3x3" ("tensor<3x3xf64>" `T.isInfixOf` mlir) + ] diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,20 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Main (main) where + +import Test.Tasty + +import qualified EigenHHLO.Test.Cholesky +import qualified EigenHHLO.Test.Eigenvalue +import qualified EigenHHLO.Test.LU +import qualified EigenHHLO.Test.QR +import qualified EigenHHLO.Test.SVD + +main :: IO () +main = defaultMain $ testGroup "eigen-hhlo" + [ EigenHHLO.Test.Cholesky.tests + , EigenHHLO.Test.Eigenvalue.tests + , EigenHHLO.Test.LU.tests + , EigenHHLO.Test.QR.tests + , EigenHHLO.Test.SVD.tests + ]