diff --git a/.gitignore b/.gitignore index 045e0e303..773ba73df 100644 --- a/.gitignore +++ b/.gitignore @@ -32,6 +32,9 @@ benchmark/results** benchmark/results/* benchmark/plots** +# Distributed linear algebra acceptance logs and profiles +/linalg-results/ + compile_wrapper.sh *.tar.gz diff --git a/Project.toml b/Project.toml index 0358c7d0b..59b83feec 100644 --- a/Project.toml +++ b/Project.toml @@ -51,6 +51,8 @@ Random = "1" StaticArrays = "1" StatsBase = "0.34" TensorOperations = "5.8" +# Developer builds must resolve before the new wrapper JLL is published. +# Raise to 26.6.2 on publication; __init__ checks the required wrapper bindings. cunumeric_jl_wrapper_jll = "26.6.1" cupynumeric_jll = "26.6.0" julia = "1.10" diff --git a/docs/src/developer_mode.md b/docs/src/developer_mode.md index 913977942..95905d520 100644 --- a/docs/src/developer_mode.md +++ b/docs/src/developer_mode.md @@ -53,6 +53,29 @@ If the build fails, check CMake / g++ (C++20) / CUDA toolkit availability as des 4. using cuNumeric ``` +## cuSolverMp wrapper update + +Distributed solve, Cholesky, and QR require wrapper source version **26.6.2**. +The cupynumeric backend remains on the 26.06 line. Until the new wrapper JLL +is published, rebuild this checkout in developer mode against `cupynumeric_jll`: + +```julia +using Pkg +Pkg.develop(PackageSpec(path="lib/CNPreferences")) +using CNPreferences +CNPreferences.use_developer_mode(; use_jll=true) +``` + +Restart Julia, run `Pkg.build("cuNumeric")` with this project active, then +restart again before loading cuNumeric. An old wrapper produces an explicit +upgrade/rebuild error instead of reaching a missing task binding. + +The project temporarily permits the published 26.6.1 wrapper JLL so package +resolution and developer builds can bootstrap. Its binary is insufficient for +this source tree. When 26.6.2 is published, raise the `cunumeric_jl_wrapper_jll` +compat lower bound to `26.6.2` before releasing these Julia changes. No Legate.jl +update is needed: version 0.2.1 already supplies explicit partition color shapes. + ## Switch back to JLLs When you no longer need a local wrapper: diff --git a/docs/src/linalg.md b/docs/src/linalg.md index b3a0f337a..49a9daf56 100644 --- a/docs/src/linalg.md +++ b/docs/src/linalg.md @@ -25,6 +25,39 @@ equivalent, so they get their own `batched_*` names. | `LinearAlgebra.qr(A)` | `NDArrayQR` | not supported | | `cuNumeric.solve(A, b)`, `A \ b` | `NDArray` | `cuNumeric.batched_solve(A, B)` | +## Distributed solves and factorizations + +The existing `A \ b`, `cuNumeric.solve`, `cholesky`, and `qr` APIs select +distributed tasks automatically. Selection follows cuPyNumeric 26.06: + +| Operation | cuSolverMp cutoff | Internal block size | +| --- | --- | --- | +| Square solve | dimension ≥ 2048 | 512 | +| Lower Cholesky | dimension ≥ 8192 | 2048 | +| Reduced QR | matrix contains ≥ 1048576 elements | 128 | + +These cutoffs and block sizes are named module constants. The loaded library +must support cuSolverMp and Legate must have more than one active GPU. The +library selects the algorithm; Legate handles placement, communication, and +redistribution. Configure resources before starting Julia, as described in +[Hardware Configuration](./configuration/hardware.md). + +Ordinary solve and QR tasks remain the fallback. Cholesky also uses the tiled +POTRF/TRSM/SYRK/GEMM algorithm for eligible multi-processor configurations; +below its partitioning cutoff this is a single tile. Batched operations still +distribute independent matrices, each of which must fit on one processor. +SVD and general eigen do not acquire distributed factorization paths. + +Results retain their existing Julia types, element promotion, and shapes. +Inputs are preserved. Kernel failures propagate when the runtime reports them; +failed collectives are not retried using another algorithm. No new factor-reuse, +triangular-solve, or CG API is introduced by this change. + +This requires the new C++ wrapper; see [Developer Mode](./developer_mode.md#cusolvermp-wrapper-update). +The acceptance runner and multi-node validation procedure are in +[`scripts/linalg/README.md`](https://github.com/JuliaLegate/cuNumeric.jl/blob/develop/scripts/linalg/README.md). +Multi-node support remains pending validation on the target cluster. + ## Matrix multiply For two 2D arrays, `*` performs matrix multiplication; use `.*` for an @@ -303,5 +336,4 @@ There is no public dense-matrix `lu`, matrix `inv`, or `ldiv!` yet (beyond the operations, not matrix inverse. Also missing: `eigh` / Hermitian eigen (needs `Hermitian` and `Symmetric` -support on `NDArray`), batched SVD and QR, and the multi-GPU cuSolverMp paths -for Cholesky and solve. +support on `NDArray`), batched SVD and QR, and conjugate gradient. diff --git a/lib/cunumeric_jl_wrapper/VERSION b/lib/cunumeric_jl_wrapper/VERSION index e8c5e343a..892cb42a2 100644 --- a/lib/cunumeric_jl_wrapper/VERSION +++ b/lib/cunumeric_jl_wrapper/VERSION @@ -1 +1 @@ -26.6.1 +26.6.2 diff --git a/lib/cunumeric_jl_wrapper/src/types.cpp b/lib/cunumeric_jl_wrapper/src/types.cpp index df218dcc8..77ab7bc6a 100644 --- a/lib/cunumeric_jl_wrapper/src/types.cpp +++ b/lib/cunumeric_jl_wrapper/src/types.cpp @@ -169,6 +169,17 @@ void wrap_linalg_ops(jlcxx::Module& mod) { legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_SOLVE}); mod.set_const("MP_SOLVE", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_MP_SOLVE}); + mod.set_const("MP_POTRF", + legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_MP_POTRF}); + mod.set_const("MP_QR", + legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_MP_QR}); + mod.set_const("POTRS", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_POTRS}); + mod.set_const("TRSM", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_TRSM}); + mod.set_const("SYRK", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_SYRK}); + mod.set_const("GEMM", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_GEMM}); + mod.set_const("TRANSPOSE_COPY_2D", + legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_TRANSPOSE_COPY_2D}); + mod.set_const("TRILU", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_TRILU}); mod.set_const("SVD", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_SVD}); mod.set_const("CQR", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_QR}); mod.set_const("POTRF", diff --git a/lib/cunumeric_jl_wrapper/src/wrapper.cpp b/lib/cunumeric_jl_wrapper/src/wrapper.cpp index 22e9540ea..68f280d87 100644 --- a/lib/cunumeric_jl_wrapper/src/wrapper.cpp +++ b/lib/cunumeric_jl_wrapper/src/wrapper.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include //needed for return type of toString methods #include #include @@ -116,6 +117,42 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod) { // True when the loaded cuSolver provides cusolverDnXgeev. Without it // cupynumeric has no GPU eigenvalue kernel for general matrices. mod.method("cusolver_has_geev", &cupynumeric_cusolver_has_geev); + mod.method("cusolvermp_available", &cupynumeric_has_cusolvermp); + + // Match cuPyNumeric 26.06: the MP kernels use a single NCCL communicator. + mod.method("add_nccl_communicator", [](legate::ManualTask& task) { + task.add_communicator("nccl"); + }); + mod.method("add_nccl_communicator", [](legate::AutoTask& task) { + task.add_communicator("nccl"); + }); + + // Tiled Cholesky launches subrectangles in the original partition's color + // space. Bounds are inclusive and zero-based, as in Legion::Rect. + mod.method("create_linalg_task", + [](legate::LocalTaskID id, int64_t row_lo, int64_t col_lo, + int64_t row_hi, int64_t col_hi) { + auto domain = Legion::Domain{Legion::Rect<2>{ + Legion::Point<2>{row_lo, col_lo}, + Legion::Point<2>{row_hi, col_hi}}}; + return legate::Runtime::get_runtime()->create_task( + get_lib(), id, domain); + }); + mod.method("add_input_tile", + [](legate::ManualTask& task, + std::shared_ptr part, + uint64_t row, uint64_t col) { + std::vector color{row, col}; + task.add_input(part->get_child_store(color)); + }); + mod.method("add_input_column", + [](legate::ManualTask& task, + std::shared_ptr part, + int32_t col) { + task.add_input(*part, legate::SymbolicPoint{ + std::vector{legate::dimension(0), + legate::constant(col)}}); + }); mod.method("add_input_proj", [](legate::ManualTask& task, diff --git a/scripts/README.md b/scripts/README.md index 4fa1d0c47..715411949 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -75,3 +75,8 @@ This error was shown for each reduction operator in `legion_redop.inl` for compl #endif #endif ``` + +## Distributed linear algebra + +See [linalg/README.md](linalg/README.md) for CPU, multi-GPU, and multi-node +acceptance runs, including isolated numerical-failure tests. diff --git a/scripts/linalg/README.md b/scripts/linalg/README.md new file mode 100644 index 000000000..1266146fe --- /dev/null +++ b/scripts/linalg/README.md @@ -0,0 +1,86 @@ +# Distributed linear algebra acceptance + +Build the in-tree wrapper in developer mode before running these scripts; see +[Developer Mode](../../docs/src/developer_mode.md). These scripts require Linux, +Julia, Bash, GNU `timeout`, and the usual Legate GPU/network configuration. +They do not install packages, change preferences, or start a scheduler allocation. + +From the repository root, in an allocation with the requested resources: + +```bash +bash scripts/linalg/run.sh cpu +bash scripts/linalg/run.sh one-gpu +CUNUMERIC_LINALG_GPUS_PER_NODE=4 bash scripts/linalg/run.sh multi-gpu +``` + +For multiple nodes, supply your cluster's working Legate launcher after `--`. +For example, in a Slurm allocation with two nodes and four GPUs per node: + +```bash +CUNUMERIC_LINALG_GPUS_PER_NODE=4 CUNUMERIC_LINALG_EXPECT_GPUS=8 \ + bash scripts/linalg/run.sh multi-node -- \ + srun --nodes=2 --ntasks=2 --ntasks-per-node=1 +``` + +Adapt the launcher to the cluster's supported Realm network and GPU binding. +The example runs one Julia/Legate process per node, with four GPUs visible to +each process. A working single-node run does not establish multi-node support. +The script checks the global GPU count reported by Legate against the expected +count; profile inspection must additionally establish participation from each node. + +## What runs + +- Production selection-policy tests around every cutoff. +- Small real/complex solve, Cholesky, and square/tall/wide QR checks, including + uneven partitions, empty reduced-output partitions, and input preservation. +- Small forced MP launches whenever multiple GPUs are active, using private + tile arguments rather than changing module constants. +- Tiled Cholesky, including its complex conjugate-transpose updates. +- Singular solve and non-positive-definite Cholesky, each in a separate process + under a timeout. Only the expected numerical error counts as success; an + unrelated exception, crash, or timeout fails the run. + +The ordinary package test suite continues to cover promotion and batched +regressions. Run it as well after rebuilding the wrapper: + +```bash +LEGATE_CONFIG='--cpus 2 --gpus 0' julia --project -e 'using Pkg; Pkg.test()' +``` + +To additionally exercise the public APIs at the production MP cutoffs: + +```bash +CUNUMERIC_LINALG_PRODUCTION=1 CUNUMERIC_LINALG_GPUS_PER_NODE=4 \ + bash scripts/linalg/run.sh multi-gpu +``` + +This adds Float64 identity smoke tests, including vector RHS through `A \ b` +and an 8192×8192 Cholesky. Provision enough system and GPU memory for inputs, +outputs, redistribution buffers, and solver workspaces. Cholesky samples columns +at that size; the small tests perform complete nontrivial reconstructions. + +## Configuration and evidence + +The defaults use two CPU processors and `--profile`. Set `LEGATE_CONFIG` to +provide memory, networking, or other resource flags; include `--profile` to keep +timeline evidence. The runner appends a distinct `--logdir` for each process +launch, so omit that flag from the supplied configuration. + +Additional environment variables: + +| Variable | Purpose | +| --- | --- | +| `CUNUMERIC_LINALG_JULIA` | Julia executable; defaults to `julia` | +| `CUNUMERIC_LINALG_TIMEOUT` | Timeout per launch in seconds; defaults to 1800 | +| `CUNUMERIC_LINALG_LOGDIR` | Output root; defaults to `linalg-results/-` | +| `CUNUMERIC_LINALG_PRODUCTION` | `1` enables production-cutoff smoke tests | + +Keep the logs containing Julia/library versions, resource configuration, GPU +count, residuals, and test summaries. Process the acceptance profiles with +`legate_prof`. Verify that `MP_SOLVE`, `MP_POTRF`, and `MP_QR` tasks run on the +requested GPUs across every node, rather than inferring distribution from a +passing residual alone. Failure-run profiles are stored separately. + +Hardware numerical tests and multi-node acceptance have not been run as part of +this implementation. Report a failure log or timeout as a validation failure; +do not interpret it as a successful fallback or reuse that runtime afterward. diff --git a/scripts/linalg/acceptance.jl b/scripts/linalg/acceptance.jl new file mode 100644 index 000000000..7221b8a7a --- /dev/null +++ b/scripts/linalg/acceptance.jl @@ -0,0 +1,62 @@ +# Run through run.sh so that runtime configuration precedes package loading. +using cuNumeric, LinearAlgebra, Test + +expected = parse(Int, ENV["CUNUMERIC_LINALG_EXPECT_GPUS"]) +actual = Int(cuNumeric.Legate.num_gpus()) +println("Julia: ", VERSION) +cuNumeric.versioninfo() +println("LEGATE_CONFIG: ", get(ENV, "LEGATE_CONFIG", "")) +println("Active GPUs: ", actual, "; expected: ", expected) +println("Active processors: ", cuNumeric.Legate.num_procs()) +println("cuSolverMp available: ", cuNumeric.cusolvermp_available()) +actual == expected || error("Runtime GPU count does not match the requested topology") +expected > 1 && !cuNumeric.cusolvermp_available() && error("cuSolverMp is unavailable") + +# Includes all four element types, public APIs, forced small MP launches, and +# tiled Cholesky. No production policy constants are changed by these tests. +include(joinpath(@__DIR__, "../../test/array/distributed_linalg.jl")) + +if "--production" in ARGS + @testset "public APIs at production cutoffs" begin + cn = cuNumeric + n = cn.MIN_SOLVE_MATRIX_SIZE + a = cn.NDArray{Float64}(I, n, n) + println("Production solve backend: ", typeof(cn._linalg_backend(Val(:solve), a))) + b = cn.ones(Float64, n) + x = a \ b + residual = norm(dl_host(x) .- 1) / sqrt(n) + println("Production vector solve relative error: ", residual) + @test residual <= dl_tol(Float64) + + n = cn.MIN_CHOLESKY_MATRIX_SIZE + a = cn.NDArray{Float64}(I, n, n) + println("Production Cholesky backend: ", typeof(cn._linalg_backend(Val(:cholesky), a))) + f = cholesky(a) + # Smoke-test sampled columns at this large cutoff; the small tests above + # check complete reconstructions on nontrivial real/complex matrices. + for cols in (1:4, (n - 3):n) + got = dl_host(copy(f.factors[:, cols])) + expected_columns = zeros(Float64, n, 4) + for (j, i) in enumerate(cols) + expected_columns[i, j] = 1 + end + residual = norm(got - expected_columns) + println("Production Cholesky sampled-column error: ", residual) + @test residual <= dl_tol(Float64) + end + + n = isqrt(cn.MIN_QR_MATRIX_SIZE) + m = cld(cn.MIN_QR_MATRIX_SIZE, n) + a = cn.NDArray{Float64}(I, m, n) + println("Production QR backend: ", typeof(cn._linalg_backend(Val(:qr), a))) + f = qr(a) + q, r = dl_host(f.Q), dl_host(f.R) + residual = norm(q * r - Matrix{Float64}(I, m, n)) / sqrt(n) + println("Production QR relative reconstruction error: ", residual) + @test residual <= dl_tol(Float64) + end +end + +cuNumeric.Legate.issue_execution_fence(true) +println("LINALG_ACCEPTANCE_PASSED") +println("Inspect the profile for MP_SOLVE, MP_POTRF, and MP_QR across the requested GPUs/nodes.") diff --git a/scripts/linalg/failure.jl b/scripts/linalg/failure.jl new file mode 100644 index 000000000..941fc5160 --- /dev/null +++ b/scripts/linalg/failure.jl @@ -0,0 +1,44 @@ +# Keep each failure in its own process; a failed collective may poison a runtime. +using cuNumeric, LinearAlgebra + +length(ARGS) == 2 || error("Usage: failure.jl solve|cholesky single|mp|tiled") +op, backend = ARGS +op in ("solve", "cholesky") || error("Unknown operation") +backend in ("single", "mp", "tiled") || error("Unknown backend") +op == "solve" && backend == "tiled" && error("No tiled solve backend") +cn = cuNumeric +backend == "mp" && cn._check_mp_launch(4) +expected_gpus = parse(Int, ENV["CUNUMERIC_LINALG_EXPECT_GPUS"]) +Int(cn.Legate.num_gpus()) == expected_gpus || error("Unexpected runtime GPU count") +n = 33 +a = cn.NDArray(zeros(Float64, n, n)) +b = cn.ones(Float64, n, 1) +out = cn.zeros(Float64, n, op == "solve" ? 1 : n) +cn.Legate.issue_execution_fence(true) +expected_message = op == "solve" ? "singular" : "positive definite" + +try + if op == "solve" + if backend == "mp" + cn._solve!(cn._CuSolverMpLinalg(), out, a, b; tile=4) + else + cn._solve!(cn._SingleProcLinalg(), out, a, b) + end + elseif backend == "mp" + cn._cholesky!(cn._CuSolverMpLinalg(), out, a; tile=4) + elseif backend == "tiled" + cn._cholesky!(cn._TiledCholesky(), out, a; min_matrix=0, min_tile=4) + else + cn._cholesky!(cn._SingleProcLinalg(), out, a) + end + cn.allowscalar() do + Array(out) # Demand the failed result; errors may be deferred. + end + cn.Legate.issue_execution_fence(true) +catch err + message = sprint(showerror, err) + occursin(expected_message, lowercase(message)) || rethrow() + println("EXPECTED_LINALG_FAILURE: ", message) + exit(0) +end +error("The invalid input unexpectedly succeeded") diff --git a/scripts/linalg/run.sh b/scripts/linalg/run.sh new file mode 100644 index 000000000..86fa970c4 --- /dev/null +++ b/scripts/linalg/run.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Optional command after -- is an external cluster launcher, passed as argv. +# All resource/memory/network options remain configurable via LEGATE_CONFIG. +mode=${1:?Usage: run.sh cpu|one-gpu|multi-gpu|multi-node [-- launcher arguments...]} +shift +case "$mode" in + cpu) gpus=0; expected=0 ;; + one-gpu) gpus=1; expected=1 ;; + multi-gpu) gpus=${CUNUMERIC_LINALG_GPUS_PER_NODE:-4}; expected=$gpus ;; + multi-node) + gpus=${CUNUMERIC_LINALG_GPUS_PER_NODE:-4} + expected=${CUNUMERIC_LINALG_EXPECT_GPUS:?Set the total GPU count across nodes} + ;; + *) echo "Unknown mode: $mode" >&2; exit 2 ;; +esac +if [[ ! $gpus =~ ^[0-9]+$ || ! $expected =~ ^[0-9]+$ ]]; then + echo "GPU counts must be nonnegative integers" >&2 + exit 2 +fi +if [[ $mode == multi-gpu && $gpus -lt 2 ]]; then + echo "multi-gpu requires at least two GPUs" >&2 + exit 2 +fi +if [[ $mode == multi-node && ( $gpus -lt 1 || $expected -le $gpus ) ]]; then + echo "multi-node expects GPUs on more than one node" >&2 + exit 2 +fi +launcher=() +if [[ ${1:-} == -- ]]; then + shift + launcher=("$@") +elif [[ $# -gt 0 ]]; then + echo "Launcher arguments must follow --" >&2 + exit 2 +fi +if [[ $mode == multi-node && ${#launcher[@]} -eq 0 ]]; then + echo "multi-node needs an external launcher after --" >&2 + exit 2 +fi + +repo=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) +logdir=${CUNUMERIC_LINALG_LOGDIR:-"$repo/linalg-results/$mode-$(date +%Y%m%d-%H%M%S)"} +mkdir -p "$logdir" +logdir=$(cd "$logdir" && pwd) +export LEGATE_AUTO_CONFIG=0 +export LEGATE_SKIP_RUNTIME=false +export LEGATE_SHOW_CONFIG=1 +export LEGATE_CONFIG="${LEGATE_CONFIG:---gpus $gpus --cpus 2 --profile}" +export CUNUMERIC_LINALG_EXPECT_GPUS=$expected +julia=${CUNUMERIC_LINALG_JULIA:-julia} +seconds=${CUNUMERIC_LINALG_TIMEOUT:-1800} +extra=() +[[ ${CUNUMERIC_LINALG_PRODUCTION:-0} == 1 ]] && extra+=(--production) + +printf 'Mode: %s\nLEGATE_CONFIG: %s\nLogs: %s\n' "$mode" "$LEGATE_CONFIG" "$logdir" +mkdir -p "$logdir/acceptance" +export CUNUMERIC_LINALG_VERBOSE=1 +LEGATE_CONFIG="$LEGATE_CONFIG --logdir $logdir/acceptance" \ +timeout --kill-after=30s "${seconds}s" "${launcher[@]}" "$julia" --project="$repo" \ + "$repo/scripts/linalg/acceptance.jl" "${extra[@]}" 2>&1 | tee "$logdir/acceptance.log" + +# Separate launches also isolate communicators and ensure failures cannot hang +# the acceptance driver forever. Nonzero exits/timeouts remain test failures. +backends=(single tiled) +[[ $expected -gt 1 ]] && backends+=(mp) +for backend in "${backends[@]}"; do + for op in solve cholesky; do + [[ $backend == tiled && $op == solve ]] && continue + logfile="$logdir/failure-$backend-$op.log" + mkdir -p "$logdir/failure-$backend-$op" + LEGATE_CONFIG="$LEGATE_CONFIG --logdir $logdir/failure-$backend-$op" \ + timeout --kill-after=30s "${seconds}s" "${launcher[@]}" "$julia" --project="$repo" \ + "$repo/scripts/linalg/failure.jl" "$op" "$backend" 2>&1 | tee "$logfile" + grep -q 'EXPECTED_LINALG_FAILURE:' "$logfile" + done +done +echo "Numerical and failure checks passed. Review the task profiles before accepting the topology." diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl index 3e636a7f7..60d63cfaa 100644 --- a/src/cuNumeric.jl +++ b/src/cuNumeric.jl @@ -161,6 +161,7 @@ const TASK_SCOPE_NAMES = CNPreferences.TASK_SCOPE_NAMES # NDArray internal include("ndarray/detail/ndarray.jl") +include("ndarray/detail/distributed_linalg.jl") include("ndarray/detail/linalg.jl") include("ndarray/detail/fft.jl") @@ -278,6 +279,12 @@ function __init__() # skip runtime here as well get(ENV, "LEGATE_SKIP_RUNTIME", false) == "true" && return nothing + isdefined(@__MODULE__, :cusolvermp_available) || error( + "cuNumeric requires C++ wrapper 26.6.2 or newer. " * + "Until that JLL is published, enable CNPreferences.use_developer_mode(; use_jll=true), " * + "restart Julia, run Pkg.build(\"cuNumeric\"), and restart Julia again.", + ) + # Start runtime, but only if not pre-compiling ensure_runtime!() diff --git a/src/ndarray/detail/distributed_linalg.jl b/src/ndarray/detail/distributed_linalg.jl new file mode 100644 index 000000000..1d9dd1bf3 --- /dev/null +++ b/src/ndarray/detail/distributed_linalg.jl @@ -0,0 +1,249 @@ +# Copyright 2024 NVIDIA Corporation +# SPDX-License-Identifier: Apache-2.0 +# Task construction follows cuPyNumeric 26.06's linalg/_solve.py, _qr.py, +# and _cholesky.py. +# Keep algorithm selection here; Legate owns placement and redistribution. +const MIN_SOLVE_MATRIX_SIZE = 2048 +const MIN_SOLVE_TILE_SIZE = 512 +const MIN_CHOLESKY_MATRIX_SIZE = 8192 +const MIN_CHOLESKY_TILE_SIZE = 2048 +const MIN_QR_MATRIX_SIZE = 1048576 +const QR_TILE_SIZE = 128 +const MAX_CHOLESKY_TILES_PER_PROC = 4 + +struct _SingleProcLinalg end +struct _CuSolverMpLinalg end +struct _TiledCholesky end + +# Neither the loaded library's capability nor the active machine is queried at +# module/precompile time. Explicit arguments also let policy tests run on CPUs. +function _linalg_backend(op, a::NDArray; kwargs...) + return _linalg_backend( + op, size(a), cusolvermp_available(), Int(Legate.num_gpus()), Int(Legate.num_procs()); + kwargs..., + ) +end + +_mp_eligible(available::Bool, gpus::Integer) = available && gpus > 1 + +# Tuple length carries dimensionality in its type. Stacked systems never enter +# the MP selector; only their leading batch axes may be distributed. +_linalg_backend(::Val{:solve}, ::Tuple, available::Bool, gpus, procs) = _SingleProcLinalg() + +function _linalg_backend(::Val{:solve}, shape::NTuple{2,Int}, available::Bool, gpus, procs) + use_mp = shape[1] >= MIN_SOLVE_MATRIX_SIZE && _mp_eligible(available, gpus) + return use_mp ? _CuSolverMpLinalg() : _SingleProcLinalg() +end + +function _linalg_backend(::Val{:qr}, shape::Tuple{Int,Int}, available::Bool, gpus, procs) + use_mp = ( + !iszero(min(shape...)) && prod(shape) >= MIN_QR_MATRIX_SIZE && + _mp_eligible(available, gpus) + ) + return use_mp ? _CuSolverMpLinalg() : _SingleProcLinalg() +end + +function _linalg_backend( + ::Val{:cholesky}, ::Tuple, available::Bool, gpus, procs; + lower::Bool=true, inplace::Bool=false, +) + return _SingleProcLinalg() +end + +function _linalg_backend( + ::Val{:cholesky}, shape::NTuple{2,Int}, available::Bool, gpus, procs; + lower::Bool=true, inplace::Bool=false, +) + (!lower || inplace || procs == 1) && return _SingleProcLinalg() + use_mp = shape[1] >= MIN_CHOLESKY_MATRIX_SIZE && _mp_eligible(available, gpus) + return use_mp ? _CuSolverMpLinalg() : _TiledCholesky() +end + +function _linalg_scalars!(task, args...) + for arg in args + Legate.add_scalar(task, Legate.Scalar(arg)) + end + return nothing +end + +function _linalg_manual_task(id, lo::Tuple{Int,Int}, hi::Tuple{Int,Int}; throws=false) + task = create_linalg_task(id, Int64(lo[1]), Int64(lo[2]), Int64(hi[1]), Int64(hi[2])) + task_throws_exception(task, throws) + return task +end + +_submit_linalg_task(task) = Legate.submit_manual_task(Legate.get_runtime(), task) + +# Partition along rows, just as Python does. Every rank uses identical color +# spaces even when a reduced QR output has fewer rows than the input. +function _mp_row_partition(n::Int, gpus::Integer) + n > 0 && gpus > 1 || throw(ArgumentError("MP tasks need nonempty inputs and multiple GPUs")) + rows = cld(n, gpus) + return rows, (cld(n, rows), 1) +end + +function _check_mp_launch(tile::Integer) + tile > 0 || throw(ArgumentError("cuSolverMp tile size must be positive")) + cusolvermp_available() && Legate.num_gpus() > 1 || + throw(ArgumentError("cuSolverMp requires a supporting library and multiple active GPUs")) + return nothing +end + +_solve!(::_SingleProcLinalg, x, a, b) = solve_batched(a, b, x) + +function _solve!(::_CuSolverMpLinalg, x, a, b; tile::Int=MIN_SOLVE_TILE_SIZE) + _check_mp_launch(tile) + n, nrhs = size(a, 1), size(b, 2) + rows, colors = _mp_row_partition(n, Int(Legate.num_gpus())) + pa = Legate.partition_by_tiling(nda_to_logical_store(a), (rows, n)) + pb = Legate.partition_by_tiling(nda_to_logical_store(b), (rows, nrhs)) + px = Legate.partition_by_tiling(nda_to_logical_store(x), (rows, nrhs)) + @task_scope "mp_solve" begin + task = _linalg_manual_task(MP_SOLVE, (0, 0), (colors[1] - 1, 0); throws=true) + Legate.add_input(task, pa) + Legate.add_input(task, pb) + Legate.add_output(task, px) + _linalg_scalars!(task, Int64(n), Int64(nrhs), Int64(tile)) + add_nccl_communicator(task) + _submit_linalg_task(task) + end + return x +end + +function _qr(::_CuSolverMpLinalg, a::NDArray{T,2}; tile::Int=QR_TILE_SIZE) where {T} + _check_mp_launch(tile) + m, n = size(a) + k = min(m, n) + q, r = cuNumeric.zeros(T, m, k), cuNumeric.zeros(T, k, n) + rows, colors = _mp_row_partition(m, Int(Legate.num_gpus())) + tiles = (rows, n) + pa = Legate.partition_by_tiling(nda_to_logical_store(a), tiles) + pq = Legate.partition_by_tiling(nda_to_logical_store(q), tiles, colors) + pr = Legate.partition_by_tiling(nda_to_logical_store(r), tiles, colors) + @task_scope "mp_qr" begin + task = _linalg_manual_task(MP_QR, (0, 0), (colors[1] - 1, 0); throws=true) + Legate.add_input(task, pa) + Legate.add_output(task, pq) + Legate.add_output(task, pr) + _linalg_scalars!(task, Int64(m), Int64(n), Int64(tile), Int64(tile)) + add_nccl_communicator(task) + _submit_linalg_task(task) + end + return q, r +end + +_cholesky!(::_SingleProcLinalg, out, a) = potrf!(out, a; lower=true, zeroout=true) + +function _cholesky!(::_CuSolverMpLinalg, out, a; tile::Int=MIN_CHOLESKY_TILE_SIZE) + _check_mp_launch(tile) + @task_scope "mp_potrf" begin + rt = Legate.get_runtime() + task = Legate.create_auto_task(rt, get_lib(), MP_POTRF) + task_throws_exception(task, true) + ai = Legate.add_input(task, nda_to_logical_array(a)) + oi = Legate.add_output(task, nda_to_logical_array(out)) + Legate.add_constraint(task, Legate.align(oi, ai)) + _linalg_scalars!(task, Int64(size(a, 1)), Int64(tile)) + add_nccl_communicator(task) + Legate.submit_auto_task(rt, task) + _cholesky_tril!(out) + end + return out +end + +function _cholesky_tril!(out::NDArray) + rt = Legate.get_runtime() + task = Legate.create_auto_task(rt, get_lib(), TRILU) + store = nda_to_logical_array(out) + Legate.add_output(task, store) + Legate.add_input(task, store) + # The third argument identifies Cholesky to the backend/mapper. + _linalg_scalars!(task, true, Int32(0), true) + Legate.submit_auto_task(rt, task) + return nothing +end + +function _cholesky_color_shape( + n::Int, procs::Integer; + min_matrix::Int=MIN_CHOLESKY_MATRIX_SIZE, min_tile::Int=MIN_CHOLESKY_TILE_SIZE, +) + n > 0 && procs > 0 && min_matrix >= 0 && min_tile > 0 || + throw(ArgumentError("invalid tiled Cholesky dimensions or tile policy")) + (procs == 1 || n <= min_matrix) && return (1, 1) + tiles = Int(procs) + while cld(n, tiles) > min_tile && 2 * tiles <= procs * MAX_CHOLESKY_TILES_PER_PROC + tiles *= 2 + end + return (tiles, tiles) +end + +# Each task reads only tiles ready at this stage; Legate records the DAG from +# these inputs/outputs. No execution fence or host copy is needed between steps. +function _cholesky!( + ::_TiledCholesky, out, a; + min_matrix::Int=MIN_CHOLESKY_MATRIX_SIZE, min_tile::Int=MIN_CHOLESKY_TILE_SIZE, +) + n = size(a, 1) + initial = _cholesky_color_shape(n, Int(Legate.num_procs()); min_matrix, min_tile) + tile = cld(n, initial[1]) + colors = cld(n, tile) + pa = Legate.partition_by_tiling(nda_to_logical_store(a), (tile, tile)) + po = Legate.partition_by_tiling(nda_to_logical_store(out), (tile, tile)) + @task_scope "tiled_cholesky" begin + task = _linalg_manual_task(TRANSPOSE_COPY_2D, (0, 0), (colors - 1, colors - 1)) + Legate.add_output(task, po) + Legate.add_input(task, pa) + _submit_linalg_task(task) + for i in 0:(colors - 1) + _cholesky_potrf!(po, i) + _cholesky_trsm!(po, i, colors) + for k in (i + 1):(colors - 1) + _cholesky_syrk!(po, k, i) + _cholesky_gemm!(po, k, i, colors) + end + end + task = _linalg_manual_task(TRILU, (0, 0), (colors - 1, colors - 1)) + Legate.add_output(task, po) + Legate.add_input(task, po) + _linalg_scalars!(task, true, Int32(0), true) + _submit_linalg_task(task) + end + return out +end + +function _cholesky_potrf!(p, i) + task = _linalg_manual_task(POTRF, (i, i), (i, i); throws=true) + Legate.add_output(task, p) + Legate.add_input(task, p) + _linalg_scalars!(task, true, false) + return _submit_linalg_task(task) +end + +function _cholesky_trsm!(p, i, colors) + i + 1 >= colors && return nothing + task = _linalg_manual_task(TRSM, (i + 1, i), (colors - 1, i); throws=true) + Legate.add_output(task, p) + add_input_tile(task, p.handle, UInt64(i), UInt64(i)) + Legate.add_input(task, p) + # Right-side solve with the conjugate transpose of the lower factor. + _linalg_scalars!(task, false, true, Int32(2), false) + return _submit_linalg_task(task) +end + +function _cholesky_syrk!(p, k, i) + task = _linalg_manual_task(SYRK, (k, k), (k, k)) + Legate.add_output(task, p) + add_input_tile(task, p.handle, UInt64(k), UInt64(i)) + Legate.add_input(task, p) + return _submit_linalg_task(task) +end + +function _cholesky_gemm!(p, k, i, colors) + k + 1 >= colors && return nothing + task = _linalg_manual_task(GEMM, (k + 1, k), (colors - 1, k)) + Legate.add_output(task, p) + add_input_column(task, p.handle, Int32(i)) + add_input_tile(task, p.handle, UInt64(k), UInt64(i)) + Legate.add_input(task, p) + return _submit_linalg_task(task) +end diff --git a/src/ndarray/detail/linalg.jl b/src/ndarray/detail/linalg.jl index 6722495c1..8957c24b1 100644 --- a/src/ndarray/detail/linalg.jl +++ b/src/ndarray/detail/linalg.jl @@ -131,9 +131,9 @@ function _solve(a::NDArray{T,N}, b::NDArray{S,N}) where {T,S,N} " (size $(size(b)[end-1]) is different from $(size(a)[end]))", ), ) - prod(size(a)) == 0 || prod(size(b)) == 0 && return cuNumeric.zeros(T, size(b)...) + (prod(size(a)) == 0 || prod(size(b)) == 0) && return cuNumeric.zeros(T, size(b)...) x = cuNumeric.zeros(T, size(b)...) - solve_batched(a, b, x) + _solve!(_linalg_backend(Val(:solve), a), x, a, b) return x end @@ -283,6 +283,7 @@ function qr_single(a::NDArray{T,N}, q::NDArray, r::NDArray) where {T,N} rt = Legate.get_runtime() lib = cuNumeric.get_lib() task = Legate.create_auto_task(rt, lib, cuNumeric.CQR) + cuNumeric.task_throws_exception(task, true) l_a = nda_to_logical_array(a) l_q = nda_to_logical_array(q) @@ -302,13 +303,18 @@ end function _qr(a::NDArray{T,2}) where {T} m, n = size(a) k = min(m, n) - # cuSolver requires full square buffers regardless of output shape - q_buf = cuNumeric.zeros(T, m, m) - r_buf = cuNumeric.zeros(T, n, n) - qr_single(a, q_buf, r_buf) - # Host conversion assumes contiguous storage, so materialize the economy slices. - q = copy(q_buf[:, 1:k]) - r = copy(r_buf[1:k, :]) + k == 0 && return cuNumeric.zeros(T, m, k), cuNumeric.zeros(T, k, n) + return _qr(_linalg_backend(Val(:qr), a), a) +end + +function _qr(::_SingleProcLinalg, a::NDArray{T,2}) where {T} + m, n = size(a) + k = min(m, n) + # CQR writes dense column-major economy factors with leading dimensions + # m for Q and k for R. Square buffers give R the wrong stride when m < n. + q = cuNumeric.zeros(T, m, k) + r = cuNumeric.zeros(T, k, n) + qr_single(a, q, r) return q, r end @@ -359,7 +365,7 @@ assumed Hermitian without being checked, matching cupynumeric. function _cholesky(a::NDArray{T,N}) where {T,N} _check_square_matrices(:cholesky, a) out = cuNumeric.zeros(T, size(a)...) - potrf!(out, a; lower=true, zeroout=true) + _cholesky!(_linalg_backend(Val(:cholesky), a), out, a) return out end diff --git a/src/ndarray/linalg.jl b/src/ndarray/linalg.jl index 879edafdd..5f41eaba9 100644 --- a/src/ndarray/linalg.jl +++ b/src/ndarray/linalg.jl @@ -9,6 +9,9 @@ Solve the linear system `A * x = b`. `A` must be a square `(m, m)` matrix. `b` must have shape `(m,)` or `(m, n)`. The result has the same shape as `b`. +Large systems automatically use cuSolverMp when it is available and Legate has +multiple active GPUs. Algorithm selection follows cuPyNumeric 26.06. + Accepted element types are `Float32`, `Float64`, `ComplexF32`, and `ComplexF64`. Integer and `Bool` inputs are converted to `Float64`. As everywhere else in the package, that conversion needs `@allowpromotion` only when it widens the @@ -51,6 +54,9 @@ being Hermitian. A non-positive-definite input raises an `ErrorException` from the task rather than `LinearAlgebra.PosDefException`, so the `check` keyword is not supported. +Large matrices use cuSolverMp when available with multiple active GPUs. +Other multi-processor configurations use cuPyNumeric's tiled Cholesky algorithm. + Accepted element types are `Float32`, `Float64`, `ComplexF32`, and `ComplexF64`. Integer and `Bool` inputs are converted to `Float64`. As everywhere else in the package, that conversion needs `@allowpromotion` only when it widens the @@ -183,6 +189,9 @@ end Reduced QR factorization of `A`, with `A ≈ F.Q * F.R`. For an `m × n` input and `k = min(m, n)`, `F.Q` is `m × k` and `F.R` is `k × n`. +Large matrices automatically use cuSolverMp when available with multiple active +GPUs, including tall and wide inputs. + See [`NDArrayQR`](@ref) for why this is not a `LinearAlgebra.QRCompactWY`. Accepted element types are `Float32`, `Float64`, `ComplexF32`, and `ComplexF64`. diff --git a/src/utilities/version.jl b/src/utilities/version.jl index 6f05603c6..71d3a581e 100644 --- a/src/utilities/version.jl +++ b/src/utilities/version.jl @@ -28,7 +28,9 @@ end versioninfo() Prints the cuNumeric build configuration summary, including package -metadata, Julia and compiler version, and paths to core dependencies. +metadata, Julia and compiler version, paths to core dependencies, and +cuSolverMp availability and linear algebra tuning constants. Runtime GPU +eligibility is reported separately from the per-operation size/shape policy. """ function versioninfo(io::IO=stdout) name = string(Base.nameof(@__MODULE__)) @@ -54,6 +56,14 @@ function versioninfo(io::IO=stdout) is_auto_config = legate_auto_config != "0" ? true : false legate_config = is_auto_config ? "auto" : get(ENV, "LEGATE_CONFIG", "not set") + # versioninfo is also called by the test driver with LEGATE_SKIP_RUNTIME. + # Do not start the runtime or query its machine just to print diagnostics. + active = runtime_started() + not_queried = "not queried (runtime inactive)" + mp_available = active ? cusolvermp_available() : not_queried + active_gpus = active ? Int(Legate.num_gpus()) : not_queried + mp_eligible = active ? _mp_eligible(mp_available, active_gpus) : not_queried + str = """ ─────────────────────────────────────────────── cuNumeric Build Configuration @@ -68,6 +78,18 @@ function versioninfo(io::IO=stdout) Brodcast Fusion: $(FUSE_BROADCAST_EXPRS) Brodcast Min Ops: $(FUSE_BROADCAST_MIN_OPS) + cuSolverMp / Linear Algebra: + Library support: $mp_available + Active GPUs: $active_gpus + MP eligible before size/shape checks: $mp_eligible + MIN_SOLVE_MATRIX_SIZE: $MIN_SOLVE_MATRIX_SIZE (dimension) + MIN_SOLVE_TILE_SIZE: $MIN_SOLVE_TILE_SIZE + MIN_CHOLESKY_MATRIX_SIZE: $MIN_CHOLESKY_MATRIX_SIZE (dimension) + MIN_CHOLESKY_TILE_SIZE: $MIN_CHOLESKY_TILE_SIZE + MIN_QR_MATRIX_SIZE: $MIN_QR_MATRIX_SIZE (elements) + QR_TILE_SIZE: $QR_TILE_SIZE + MAX_CHOLESKY_TILES_PER_PROC: $MAX_CHOLESKY_TILES_PER_PROC + Hostname: $hostname Julia Version: $(VERSION) C++ Compiler: $compiler diff --git a/test/array/distributed_linalg.jl b/test/array/distributed_linalg.jl new file mode 100644 index 000000000..e2f574db8 --- /dev/null +++ b/test/array/distributed_linalg.jl @@ -0,0 +1,168 @@ +using Test, LinearAlgebra, Random +import cuNumeric + +dl_host(a) = cuNumeric.allowscalar() do + Array(a) +end +dl_tol(::Type{T}) where {T} = 200 * eps(real(T)) +function dl_record(op, T, residual) + get(ENV, "CUNUMERIC_LINALG_VERBOSE", "0") == "1" && + println(op, " ", T, " residual: ", residual) + return nothing +end + +@testset "linear algebra task selection" begin + cn = cuNumeric + for (op, limit) in ( + (:solve, cn.MIN_SOLVE_MATRIX_SIZE), (:cholesky, cn.MIN_CHOLESKY_MATRIX_SIZE) + ) + for available in (false, true), gpus in (0, 1, 2, 4), n in (limit - 1, limit, limit + 1) + backend = cn._linalg_backend(Val(op), (n, n), available, gpus, max(2, gpus)) + @test (backend isa cn._CuSolverMpLinalg) == (available && gpus > 1 && n >= limit) + end + end + qr_volumes = ( + cn.MIN_QR_MATRIX_SIZE - 1, cn.MIN_QR_MATRIX_SIZE, cn.MIN_QR_MATRIX_SIZE + 1 + ) + for available in (false, true), gpus in (0, 1, 2, 4), volume in qr_volumes + backend = cn._linalg_backend(Val(:qr), (volume, 1), available, gpus, max(2, gpus)) + @test (backend isa cn._CuSolverMpLinalg) == ( + available && gpus > 1 && volume >= cn.MIN_QR_MATRIX_SIZE + ) + end + for op in (:solve, :cholesky) + @test cn._linalg_backend(Val(op), (4, 9000, 9000), true, 4, 4) isa cn._SingleProcLinalg + end + @test cn._linalg_backend( + Val(:cholesky), (9000, 9000), true, 4, 4; lower=false + ) isa cn._SingleProcLinalg + @test cn._linalg_backend( + Val(:cholesky), (9000, 9000), true, 4, 4; inplace=true + ) isa cn._SingleProcLinalg + @test cn._linalg_backend(Val(:cholesky), (9000, 9000), true, 1, 1) isa cn._SingleProcLinalg + @test cn._linalg_backend(Val(:cholesky), (9000, 9000), false, 4, 4) isa cn._TiledCholesky + @test cn._linalg_backend(Val(:qr), (0, 10), true, 4, 4) isa cn._SingleProcLinalg + @test cn._mp_row_partition(33, 4) == (9, (4, 1)) + @test cn._mp_row_partition(2, 4) == (1, (2, 1)) + @test cn._cholesky_color_shape(33, 4; min_matrix=0, min_tile=4) == (16, 16) + @test cn._cholesky_color_shape(cn.MIN_CHOLESKY_MATRIX_SIZE, 4) == (1, 1) +end + +function dl_check_solve(T; distributed=false) + cn = cuNumeric + rng = MersenneTwister(71) + n = 33 + a = randn(rng, T, n, n) + T(n) * I + da = cn.NDArray(a) + for nrhs in (1, 3) + b = randn(rng, T, n, nrhs) + db = cn.NDArray(b) + x = if distributed + out = cn.zeros(T, n, nrhs) + cn._solve!(cn._CuSolverMpLinalg(), out, da, db; tile=4) + else + da \ db + end + hx = dl_host(x) + residual = norm(a * hx - b) / (norm(a) * norm(hx) + norm(b)) + dl_record(distributed ? "MP_SOLVE" : "solve", T, residual) + @test residual <= dl_tol(T) + @test dl_host(da) == a + @test dl_host(db) == b + end + # Public vector-RHS reshape path, independent of the low-level MP tests. + b = randn(rng, T, n) + db = cn.NDArray(b) + x = da \ db + @test size(x) == (n,) + @test isapprox(dl_host(x), a \ b; rtol=dl_tol(T)) +end + +function dl_check_cholesky(T; backend=nothing) + cn = cuNumeric + rng = MersenneTwister(72) + n = 33 + z = randn(rng, T, n, n) + a = z * z' + T(n) * I + # Only the lower triangle is meaningful, including for complex input. + input = copy(a) + for j in 1:n, i in 1:(j - 1) + input[i, j] = T(123) + end + da = cn.NDArray(input) + factors = if backend === nothing + f = cholesky(da) + @test f isa Cholesky + f.factors + elseif backend isa cn._CuSolverMpLinalg + cn._cholesky!(backend, cn.zeros(T, n, n), da; tile=4) + else + cn._cholesky!(backend, cn.zeros(T, n, n), da; min_matrix=0, min_tile=4) + end + l = dl_host(factors) + residual = norm(a - l * l') / norm(a) + dl_record("cholesky ($(typeof(backend)))", T, residual) + @test residual <= dl_tol(T) + @test istril(l) + @test dl_host(da) == input +end + +function dl_check_qr(T; distributed=false) + cn = cuNumeric + rng = MersenneTwister(73) + @testset "QR shape ($m, $n)" for (m, n) in ( + (33, 33), (65, 17), (17, 65), (7, 1), (1, 7) + ) + a = randn(rng, T, m, n) + da = cn.NDArray(a) + q, r = if distributed + cn._qr(cn._CuSolverMpLinalg(), da; tile=4) + else + f = qr(da) + @test f isa cn.NDArrayQR + f.Q, f.R + end + k = min(m, n) + @test size(q) == (m, k) + @test size(r) == (k, n) + hq, hr = dl_host(q), dl_host(r) + residual = norm(a - hq * hr) / norm(a) + dl_record(distributed ? "MP_QR ($m, $n)" : "qr ($m, $n)", T, residual) + @test residual <= dl_tol(T) + @test norm(hq' * hq - I) / sqrt(k) <= dl_tol(T) + @test istriu(hr) + @test dl_host(da) == a + end +end + +@testset "distributed linear algebra numerics" begin + for T in (Float32, Float64, ComplexF32, ComplexF64) + @testset "$T" begin + dl_check_solve(T) + dl_check_cholesky(T) + dl_check_qr(T) + dl_check_cholesky(T; backend=cuNumeric._TiledCholesky()) + if cuNumeric.cusolvermp_available() && cuNumeric.Legate.num_gpus() > 1 + dl_check_solve(T; distributed=true) + dl_check_cholesky(T; backend=cuNumeric._CuSolverMpLinalg()) + dl_check_qr(T; distributed=true) + else + @test_skip "MP numerical tests require multiple active GPUs and cuSolverMp" + end + end + end +end + +@testset "empty and invalid solves" begin + cn = cuNumeric + @test size(cn.zeros(Float64, 0, 0) \ cn.zeros(Float64, 0)) == (0,) + @test size(cn.zeros(Float64, 0, 0) \ cn.zeros(Float64, 0, 3)) == (0, 3) + @test size(cn.zeros(Float64, 3, 3) \ cn.zeros(Float64, 3, 0)) == (3, 0) + @test_throws ArgumentError cn.zeros(Float64, 2, 3) \ cn.zeros(Float64, 2) + @test_throws ArgumentError cn.zeros(Float64, 3, 3) \ cn.zeros(Float64, 2) + for (m, n) in ((0, 0), (0, 3), (3, 0)) + f = qr(cn.zeros(Float64, m, n)) + @test size(f.Q) == (m, min(m, n)) + @test size(f.R) == (min(m, n), n) + end +end diff --git a/test/array/linalg_edge_cases.jl b/test/array/linalg_edge_cases.jl new file mode 100644 index 000000000..30b7ca97a --- /dev/null +++ b/test/array/linalg_edge_cases.jl @@ -0,0 +1,150 @@ +using Test, LinearAlgebra, Random +using cuNumeric: cuNumeric + +le_host(a) = cuNumeric.allowscalar() do + return Array(a) +end +le_tol(::Type{T}) where {T} = 200 * eps(real(T)) +le_residual(a, b) = norm(a - b) / max(norm(a), one(real(eltype(a)))) + +# Keep the padded parent so we can check that the operation preserves both +# its input and the elements outside the view. Do not copy the view before +# passing it to the operation: the backend must handle its layout. +function le_input(a, layout) + if layout == :slice + m, n = size(a) + parent = fill(eltype(a)(19), m + 2, n + 2) + parent[2:(m + 1), 2:(n + 1)] = a + dp = cuNumeric.NDArray(parent) + return view(dp, 2:(m + 1), 2:(n + 1)), dp, parent + elseif layout == :transpose + parent = copy(transpose(a)) + dp = cuNumeric.NDArray(parent) + return cuNumeric.transpose(dp), dp, parent + end + dp = cuNumeric.NDArray(a) + return dp, dp, a +end + +function le_qr(a, da) + m, n = size(a) + k = min(m, n) + f = qr(da) + q, r = le_host(f.Q), le_host(f.R) + @test size(q) == (m, k) + @test size(r) == (k, n) + @test le_residual(a, q * r) <= le_tol(eltype(a)) + @test norm(q' * q - I) <= le_tol(eltype(a)) * k + @test istriu(r) +end + +function le_svd(a, da) + m, n = size(a) + for full in (false, true) + f = svd(da; full) + u, s, vt = le_host(f.U), le_host(f.S), le_host(f.Vt) + @test size(u) == (m, full ? m : n) + @test size(s) == (n,) + @test size(vt) == (n, n) + @test le_residual(a, u[:, 1:n] * Diagonal(s) * vt) <= le_tol(eltype(a)) + @test norm(u' * u - I) <= le_tol(eltype(a)) * size(u, 2) + @test norm(vt * vt' - I) <= le_tol(eltype(a)) * n + @test all(s .>= 0) + @test issorted(s; rev=true) + @test isapprox(s, svdvals(a); atol=le_tol(eltype(a)), rtol=le_tol(eltype(a))) + end +end + +function le_eigen(a, da) + f = eigen(da) + w, v = le_host(f.values), le_host(f.vectors) + n = size(a, 1) + @test size(w) == (n,) + @test size(v) == (n, n) + @test norm(a * v - v * Diagonal(w)) / max(norm(a) * norm(v), 1) <= le_tol(eltype(a)) + @test all(j -> isapprox(norm(v[:, j]), 1; atol=le_tol(eltype(a))), 1:n) + # The fixtures are Hermitian, so their spectra are real, including repeated + # zero eigenvalues. Sorting by real part avoids arbitrary eigenvector order. + expected = eigvals(Hermitian(a)) + for values in (w, le_host(eigvals(da))) + @test maximum(abs, imag.(values)) <= le_tol(eltype(a)) * max(norm(a), 1) + @test isapprox( + sort(real.(values)), expected; atol=le_tol(eltype(a)), rtol=le_tol(eltype(a)) + ) + end +end + +@testset "linear algebra degenerate inputs" begin + @testset "$T" for T in (Float32, Float64, ComplexF32, ComplexF64) + @testset "QR/SVD $kind ($m, $n)" for kind in (:zero, :rank_one), + (m, n) in ((4, 4), (6, 4), (4, 6)) + + u = T.(1:m) + T <: Complex && (u .+= im .* reverse(u)) + a = kind == :zero ? zeros(T, m, n) : u * transpose(T.(1:n)) + da = cuNumeric.NDArray(a) + le_qr(a, da) + m >= n && le_svd(a, da) + @test le_host(da) == a + end + @testset "square $kind" for kind in (:zero, :rank_one) + a = zeros(T, 4, 4) + kind == :rank_one && (a[1, 1] = 3) + da = cuNumeric.NDArray(a) + le_eigen(a, da) + @test le_host(da) == a + # A zero pivot is exact here. Materialize inside @test_throws so + # asynchronous task errors are observed by the assertion. + for b in (ones(T, 4), ones(T, 4, 2)) + db = cuNumeric.NDArray(b) + @test_throws "Singular matrix" le_host(da \ db) + @test le_host(db) == b + @test le_host(da) == a + end + @test_throws "Matrix is not positive definite" le_host(cholesky(da).factors) + @test le_host(da) == a + end + end +end + +@testset "linear algebra input layouts" begin + @testset "$T $layout" for T in (Float32, Float64, ComplexF32, ComplexF64), + layout in (:slice, :transpose) + + rng = MersenneTwister(81) + @testset "QR/SVD ($m, $n)" for (m, n) in ((4, 4), (6, 4), (4, 6)) + a = randn(rng, T, m, n) + da, dp, parent = le_input(a, layout) + le_qr(a, da) + m >= n && le_svd(a, da) + @test le_host(dp) == parent + end + z = randn(rng, T, 4, 4) + a = z * z' + T(4) * I + da, dp, parent = le_input(a, layout) + l = le_host(cholesky(da).factors) + @test le_residual(a, l * l') <= le_tol(T) + @test istril(l) + @test le_host(dp) == parent + le_eigen(a, da) + @test le_host(dp) == parent + # Also exercise a non-Hermitian solve so transpose/conjugation mistakes + # cannot be hidden by the Cholesky/eigen fixture's symmetry. + a = z + T(8) * I + da, dp, parent = le_input(a, layout) + @testset "solve $nrhs RHS" for nrhs in (1, 2) + b = randn(rng, T, 4, nrhs) + db, bp, bparent = le_input(b, layout) + x = le_host(da \ db) + @test size(x) == size(b) + @test le_residual(b, a * x) <= le_tol(T) + @test le_host(bp) == bparent + @test le_host(dp) == parent + end + # A zero RHS is valid even though a zero coefficient matrix is not. + for b in (zeros(T, 4), zeros(T, 4, 2)) + @test le_host(da \ cuNumeric.NDArray(b)) == b + end + @test le_host(dp) == parent + end +end