cuSolverMp

.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

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"

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:

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.

lib/cunumeric_jl_wrapper/VERSION

@@ -1 +1 @@
-26.6.1
+26.6.2

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",

lib/cunumeric_jl_wrapper/src/wrapper.cpp

@@ -21,6 +21,7 @@
 #include <cstdint>
 #include <initializer_list>
 #include <iostream>
+#include <memory>
 #include <string>  //needed for return type of toString methods
 #include <type_traits>
 #include <vector>
@@ -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<legate::LogicalStorePartition> part,
+                uint64_t row, uint64_t col) {
+               std::vector<uint64_t> color{row, col};
+               task.add_input(part->get_child_store(color));
+             });
+  mod.method("add_input_column",
+             [](legate::ManualTask& task,
+                std::shared_ptr<legate::LogicalStorePartition> part,
+                int32_t col) {
+               task.add_input(*part, legate::SymbolicPoint{
+                   std::vector<legate::SymbolicExpr>{legate::dimension(0),
+                                                     legate::constant(col)}});
+             });
 
   mod.method("add_input_proj",
              [](legate::ManualTask& task,

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.

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/<mode>-<timestamp>` |
+| `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.

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.")

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")