diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ed1a360..f36e9af 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -72,11 +72,12 @@ jobs: os: [ubuntu-latest] julia-version: ['1'] julia-arch: [x64] - # Both legs must pass: the AMDGPU extension (T23) runs the suite on - # ROCBackend with rocBLAS/rocSOLVER (test/backends.jl adds it when - # AMDGPU is functional). Making the amdgpu check required is a branch - # ruleset setting of the owner. - backend: ['cuda', 'amdgpu'] + # Every leg must pass: the AMDGPU extension (T23) runs the suite on + # ROCBackend with rocBLAS/rocSOLVER, the oneAPI extension (T31) on + # oneAPIBackend with oneMKL (the runner `headroom`); test/backends.jl adds + # each backend when its package is functional. Making a check required is + # a branch ruleset setting of the owner. + backend: ['cuda', 'amdgpu', 'oneapi'] runs-on: [self-hosted, linux, X64, gpu, '${{ matrix.backend }}'] steps: @@ -95,9 +96,48 @@ jobs: if: matrix.backend == 'amdgpu' run: | julia --project=test -e 'using Pkg; Pkg.add("AMDGPU")' + # oneAPI.jl main: the released oneAPI.jl does not support KernelAbstractions 0.10 yet, + # which the oneAPI extension needs; switch to Pkg.add("oneAPI") once a release does + - name: Add oneAPI.jl to the test environment + if: matrix.backend == 'oneapi' + run: | + julia --project=test -e 'using Pkg; Pkg.add(url = "https://github.com/JuliaGPU/oneAPI.jl", rev = "main")' + julia --project=test -e 'using oneAPI; oneAPI.versioninfo()' - uses: julia-actions/julia-runtest@latest - uses: julia-actions/julia-processcoverage@v1 - uses: codecov/codecov-action@v4 with: file: lcov.info token: ${{ secrets.CODECOV_TOKEN }} + + # Metal (T31) on GitHub's hosted Apple Silicon runner. Metal has no Float64 and the extension takes Float32 + # only (`to_backend` raises NotSupportedError otherwise), so this leg tests Float32 whatever the + # workflow-level element types are. + test-metal: + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + SDS_TEST_ELTYPES: 'Float32' + SDS_TEST_CPU: '0' + SDS_TEST_SKIP: 'test_aqua' + # the hosted macOS arm64 runner has 3 cores and 7 GB + PTR_NUM_JOBS: '3' + runs-on: macos-latest + strategy: + fail-fast: false + matrix: + julia-version: ['1'] + julia-arch: [aarch64] + + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@latest + with: + version: ${{ matrix.julia-version }} + arch: ${{ matrix.julia-arch }} + - name: Machine + run: sysctl -n hw.ncpu machdep.cpu.brand_string && system_profiler SPDisplaysDataType | grep -E "Chipset|Metal" || true + - name: Add Metal.jl to the test environment + run: | + julia --project=test -e 'using Pkg; Pkg.add("Metal")' + julia --project=test -e 'using Metal; Metal.versioninfo(); Metal.functional() || error("Metal is not functional on this runner")' + - uses: julia-actions/julia-runtest@latest diff --git a/AGENTS.md b/AGENTS.md index 95a05aa..afca89d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -96,7 +96,9 @@ reviewing agent check; the owner merges by hand. The moving parts: `ci.yml` on the PR, and from the `amdgpu` leg of `ci.yml` on the self-hosted AMD runner (the AMDGPU extension, T23; no longer continue-on-error, a required check once the owner adds it to the branch - ruleset). + ruleset), and from the `oneapi` leg on the self-hosted Intel GPU runner + `headroom` (the oneAPI extension, T31; oneAPI.jl from its `main` branch until a + release supports KernelAbstractions 0.10). ### The owner's machine @@ -108,8 +110,8 @@ reviewing agent check; the owner merges by hand. The moving parts: KernelAbstractions CPU backend and on CUDA only. GitHub CI (`.github/workflows/ci.yml`, modelled on `../ExaPF.jl`) runs the CPU suite and the GPU suite on the exanauts self-hosted runners (CPU suite: the `kkt` - machine; GPU suite: the runners labelled `cuda` and `amdgpu`). oneAPI and Metal extension files are written by analogy - (task T23) and only precompile-checked. + machine; GPU suite: the runners labelled `cuda`, `amdgpu` and `oneapi`; Metal: + GitHub's hosted `macos-latest` runner, `Float32` only, job `test-metal`). * CUDA.jl 6.x is split into packages: `CUDACore` (`CuArray`, `CUDABackend`), `cuSPARSE` (`CuSparseMatrixCSR`, `CuSparseMatrixCSC`), `cuBLAS`, `cuSOLVER`. Use these as weak dependencies; tests load the umbrella `CUDA`. diff --git a/METAL_OPTIMIZATION_PLAN.md b/METAL_OPTIMIZATION_PLAN.md new file mode 100644 index 0000000..db6f3eb --- /dev/null +++ b/METAL_OPTIMIZATION_PLAN.md @@ -0,0 +1,294 @@ +# Metal Optimization Plan + +## Goal + +Improve Float32 performance on Apple GPUs without changing the portable sparse +symbolic analysis, factor layout, or public API. The primary target is warmed +Cholesky refactorization and solve on fixed-pattern systems. Correctness and +phase ordering take precedence over individual dense-kernel timings. + +The reference workload is `lap2d_300` (`n = 90_000`, `nnz(A) = 448_800`) so +results can be compared with the existing RTX 4080 and cuDSS records. Smaller +Laplacians remain in the benchmark to locate the CPU/Metal crossover. + +## Current Baseline + +M4 Max, Metal.jl 1.11.1, Float32, warmed median: + +| workload | refactorization | solve | relative residual | +| --- | ---: | ---: | ---: | +| `lap2d_120`, `n = 14_400` | 16.8 ms | 1.59 ms | 1.56e-4 | +| `lap2d_300`, `n = 90_000` | 142 ms | 5.19 ms | 1.08e-3 | +| `lap2d_300`, two IR steps | 142 ms | 15.1 ms | 4.63e-4 | + +Published RTX 4080 Float64 results for `lap2d_300`: + +| implementation | refactorization | solve | relative residual | +| --- | ---: | ---: | ---: | +| SparseDirectSolver CUDA | 32.9 ms | 3.48 ms | 1.78e-12 | +| cuDSS | 3.96 ms | 0.675 ms | 1.84e-12 | + +The precision difference must be shown with every CUDA comparison. + +## Available Libraries + +### Metal.jl / Metal Performance Shaders + +Use the existing low-level Metal.jl wrappers for: + +- `MPSMatrixDecompositionCholesky` +- `MPSMatrixSolveTriangular` +- `MPSMatrixMultiplication` +- `MPSMatrix`, including explicit buffer offsets and descriptors + +Do not use the high-level `cholesky!` and triangular `ldiv!` methods in the +numeric phase. They allocate scratch matrices, synchronize for status, and do +not accept the panel views used by the solver. + +### AppleAccelerate.jl + +Use Accelerate/libSparse only as an Apple CPU baseline and possible future +Float64 fallback. It is not a Metal GPU sparse solver and is not part of the +initial implementation dependency set. + +### MLX + +Treat MLX as a source of optimized Metal kernel designs, not as a dependency. +Reconsider only if direct MPS integration is slower and zero-copy `MTLBuffer` +interop can be demonstrated in a small prototype. + +## Phase 1: Measure Regime-C Dense Operations + +Create a standalone benchmark that uses the actual solver panel layout and +compares, by front shape: + +- current KA POTRF +- low-level MPS Cholesky +- current KA TRSM +- low-level MPS triangular solve +- current Metal.jl panel GEMM/SYRK +- low-level MPS matrix multiplication + +Shapes must be taken from `lap2d_300` and at least one condensed KKT matrix. +Measure complete operation cost, including descriptor construction and command +submission, both cold and warmed. + +Record for every shape: + +- dimensions and leading dimension +- median and minimum runtime +- host allocations +- Metal allocations +- synchronization behavior +- maximum relative error +- factorization status behavior on an indefinite input + +Gate: no operation enters the solver unless it is correct on panel offsets, +asynchronous, and faster than KA at a stable size threshold. + +## Phase 2: Offset-Aware MPS Bindings + +Add Metal extension methods that construct `MPSMatrix` directly over the +solver's existing factor and workspace buffers: + +```text +factor MtlVector + + panel byte offset + + rows, columns + + rowBytes / matrixBytes + -> MPSMatrix +``` + +Precompute and retain during analysis/allocation: + +- matrix descriptors per regime-C front shape +- MPS kernel objects per operation and shape +- status buffers +- any required scratch descriptors and buffers + +No per-refactorization device allocation is allowed. Host object allocation +must be measured and either eliminated or shown negligible relative to the +operation. + +The bindings stay in `SparseDirectSolverMetalExt`; numeric core code continues +to call the dense interface. + +## Phase 3: Ordered KA and MPS Execution + +Prototype command ordering before enabling MPS in the solver: + +1. KA assembly writes a front. +2. MPS factors or updates that front. +3. KA packs and extend-adds the result. +4. A later KA kernel consumes it. + +Verify visibility without intermediate host synchronization over at least 100 +repetitions and with Metal private storage. Use one of: + +- one shared command buffer/queue accepted by both KA and MPS; +- explicit command-buffer commits on one queue, relying on queue order; +- Metal events/fences if separate encoders require them. + +Gate: bitwise-stable status and residuals across repeated runs. If the ordering +cannot be made reliable without one synchronization per front, MPS POTRF/TRSM +is rejected. + +## Phase 4: Selective Regime-C Dispatch + +Enable MPS independently per operation and only above measured thresholds. +Expected initial policy: + +| operation | small/medium front | large front | +| --- | --- | --- | +| POTRF | KA | MPS if ordered and faster | +| TRSM | KA | MPS if ordered and faster | +| GEMM/SYRK | Metal.jl generic/native | MPS if faster | +| assembly/extend-add | KA | KA | + +Thresholds belong to Metal backend capabilities, not global options. Preserve +explicit `impl = :ka` as a correctness/reference route. + +Uniform batches remain on the current KA/per-member path until MPS batch +correctness and status handling are demonstrated separately. + +## Phase 5: Reduce Launch Count + +Dense acceleration alone cannot remove the hundreds of sparse launch groups. +After MPS integration, profile launch and GPU time separately and pursue: + +- merge adjacent same-level Metal launch groups where dependencies permit; +- preserve fused level solves for all Metal fronts; +- increase subtree fusion subject to the 32 KiB threadgroup-memory limit; +- test Apple-specific workgroup sizes for regime A and B; +- batch assembly kernels for consecutive regime-C fronts; +- use Metal command batching without forcing early flushes between KA kernels. + +Every launch-count change must retain owner-pull determinism and cannot assume +inter-threadgroup forward progress. + +## Phase 6: Accuracy Strategy + +Same-precision refinement is insufficient for `lap2d_300`. Keep performance +and accuracy results separate until mixed precision exists. + +Near-term: + +- report raw Float32 residuals; +- report same-precision IR residual and total solve time; +- document matrices that fail the standard Float32 tolerance. + +Future work: + +- Float32 Metal factors and solves; +- Float64 residual on the CPU using unified memory or an explicit conversion; +- FGMRES refinement with Float32 factors; +- AppleAccelerate/libSparse Float64 fallback for systems that do not converge. + +## Verification Matrix + +Run after each enabled MPS operation: + +```bash +SDS_TEST_CPU=0 SDS_TEST_ELTYPES=Float32 \ +SDS_TEST_ONLY="test_dense,test_numeric_cholesky_a,test_numeric_cholesky_b,test_numeric_cholesky_c,test_solve,test_ubatch,test_schur" \ +PTR_NUM_JOBS=1 julia +1.13.1 --project=. -e 'using Pkg; Pkg.test()' +``` + +Before merging: + +- full Float32 Metal suite; +- full Float32 CPU suite; +- Aqua; +- CUDA CI to verify backend capability defaults do not alter CUDA dispatch; +- `git diff --check`; +- no committed `Manifest.toml` or local GPU dependency in `test/Project.toml`. + +Correctness gates: + +- no test failures; +- exact pivot sequences and inertia where specified; +- failed batch members report the correct original pivot; +- 100 repeated refactorization/solve cycles without races; +- explicit user schedule and `impl` choices remain authoritative. + +## Performance Gates + +Use warmed medians with at least seven samples and synchronize after each +phase. Every timing is accepted only with a recorded residual. + +Minimum useful improvement for merging an optimization: + +- at least 10% lower `lap2d_300` refactorization time, or +- at least 10% lower solve time without slowing refactorization by more than + 3%, or +- at least 15% lower combined refactorization-plus-solve time. + +Report: + +- M4 Max before/after; +- M4 CPU before/after; +- published RTX 4080 SDS and cuDSS context; +- precision and residual differences; +- schedule composition and launch count. + +## Stop Conditions + +Do not merge an approach when it: + +- synchronizes once per front; +- allocates device storage per refactorization or solve; +- requires scalar indexing; +- weakens failed-pivot reporting; +- depends on unsupported inter-threadgroup progress; +- regresses CPU, CUDA, or ROCm defaults; +- adds MLX or another dependency without a measured advantage over direct MPS. + +## Planned Order of Work + +1. Build panel-shape benchmark and capture `lap2d_300` front distribution. +2. Prototype offset-aware MPS GEMM, then POTRF, then TRSM. +3. Prove KA/MPS command ordering over repeated runs. +4. Enable the fastest correct operations behind Metal capabilities. +5. Re-profile launch count and implement one launch-reduction experiment. +6. Run the verification matrix and publish the before/after table. + +## Execution Record + +Implemented and measured on the M4 Max: + +- Captured `lap2d_300`: 113 regime-C fronts, widths 7–300; 18 fronts have + width at least 64. +- Added an offset-aware low-level MPS Cholesky prototype over padded solver + panels. MPS crosses the KA POTRF near width 64 and is 5–14× faster at widths + 128–256 in the isolated benchmark. +- Rejected MPS TRSM: it was correct but made end-to-end refactorization slower. +- Enabled MPS POTRF for single-matrix Float32 fronts of width at least 64. +- Cached MPS descriptors, decomposition kernels and the status buffer binding + during numeric allocation; no panel or status device allocation occurs in a + refactorization. +- Proved queue ordering with 100 repeated refactorization/solve cycles using + private Metal storage. +- Preserved KA POTRF for narrow fronts and every uniform batch. +- Preserved KA diagonal validation after MPS so failed pivots retain the + package's LAPACK-style first-column semantics. +- Fused every Metal front in the level solve and selected 128 work items per + Metal solve workgroup, without changing other backends. + +Final `lap2d_300` result before the full regression gate: + +| phase | initial Metal | final Metal | speedup | +| --- | ---: | ---: | ---: | +| refactorization | 142 ms | 62 ms | 2.29× | +| solve | 5.19 ms | 5.18 ms | unchanged | + +Published RTX 4080 context for the same matrix (Float64, versus Metal +Float32): SparseDirectSolver CUDA records 32.9 ms refactorization and 3.48 ms +solve, so final Metal is 1.89× and 1.49× slower respectively. cuDSS records +3.96 ms and 0.675 ms, so final Metal is 15.7× and 7.7× slower. Metal's raw +Float32 residual is about 9.3e-4; these timings are not precision parity. + +The final MPS route creates command-buffer wrapper objects and measured about +2.1 MB of host allocation per refactorization. The MPS matrix descriptors, +kernels, panel storage and status storage are preallocated. Removing the +remaining wrapper allocation requires an upstream-style reusable command-buffer +encoding API and is follow-up work rather than a correctness blocker. diff --git a/Project.toml b/Project.toml index fdbc029..95ca428 100644 --- a/Project.toml +++ b/Project.toml @@ -20,6 +20,8 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2" Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" +Metal = "dde4c033-4e86-420c-a63e-0dd931031962" +oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" cuBLAS = "182d3088-87b7-4494-8cad-fc6afaa545bc" cuSOLVER = "887afef0-6a32-4de5-add4-7827692ba8fc" cuSPARSE = "b26da814-b3bc-49ef-b0ee-c816305aa060" @@ -29,6 +31,8 @@ SparseDirectSolverAMDGPUExt = "AMDGPU" SparseDirectSolverCUDAExt = ["CUDACore", "cuBLAS", "cuSOLVER", "cuSPARSE"] SparseDirectSolverKrylovExt = "Krylov" SparseDirectSolverMadNLPExt = "MadNLP" +SparseDirectSolverMetalExt = "Metal" +SparseDirectSolverOneAPIExt = "oneAPI" [compat] AMD = "0.5" @@ -43,9 +47,11 @@ KernelAbstractions = "0.9, 0.10" Krylov = "0.10" LinearAlgebra = "1.13" MadNLP = "0.8, 0.9, 0.10" +Metal = "1.11" Metis = "1" SparseArrays = "1.13" cuBLAS = "6" cuSOLVER = "6" cuSPARSE = "6" julia = "1.13" +oneAPI = "2.9.2" diff --git a/README.md b/README.md index e630a55..20e4264 100644 --- a/README.md +++ b/README.md @@ -9,8 +9,9 @@ and phases of [CUDSS.jl](https://github.com/exanauts/CUDSS.jl) so that MadNLP and other cuDSS users can switch to it mechanically, and targets CUDA, AMDGPU, oneAPI, Metal and the KernelAbstractions CPU backend. -**Status: v0.1 under construction.** Working today, on the CPU backend, on -CUDA and on AMDGPU (ROCm): +**Status: v0.1 under construction.** Working today on the CPU backend, CUDA, +AMDGPU (ROCm), Metal (`Float32`, tested on an M4 Max) and oneAPI (tested on an +Intel Data Center GPU Max 1550): * host symbolic analysis: AMD or nested dissection (METIS) ordering, elimination tree, supernodes with GPU-tuned @@ -60,8 +61,7 @@ CUDA and on AMDGPU (ROCm): through `SDS_LOG_LEVEL`. Not there yet: non-uniform batches, ND partition-tree export, mixed precision, -hybrid host memory, delayed pivots, and the oneAPI and Metal -extensions. Unsupported structures, phases and parameters raise +hybrid host memory, delayed pivots, and the oneAPI extension. Unsupported structures, phases and parameters raise `NotSupportedError` rather than falling back silently. The remaining gap to cuDSS is performance, not features (see below and [`PERFORMANCE.md`](PERFORMANCE.md)); a MadNLP end-to-end run on CUDA and AMD @@ -108,7 +108,9 @@ Pkg.add(url = "https://github.com/exanauts/SparseDirectSolver.jl") Loading CUDA.jl enables the CUDA extension (`CuSparseMatrixCSR`/`CuSparseMatrixCSC` constructors, cuBLAS/cuSOLVER dense kernels); loading AMDGPU.jl enables the AMDGPU extension (`ROCSparseMatrixCSR`/`ROCSparseMatrixCSC` constructors, -rocBLAS/rocSOLVER dense kernels; ROCm must be found by AMDGPU.jl). +rocBLAS/rocSOLVER dense kernels; ROCm must be found by AMDGPU.jl); loading +oneAPI.jl enables the oneAPI extension (`oneSparseMatrixCSR`/`oneSparseMatrixCSC` +constructors, oneMKL BLAS kernels; needs oneAPI.jl with KernelAbstractions 0.10). ## Usage diff --git a/TASKS.md b/TASKS.md index d7bfc7b..caa3d1f 100644 --- a/TASKS.md +++ b/TASKS.md @@ -3299,7 +3299,7 @@ minutes of work (`src/symbolic/ndtree.jl` with the tree export/import and a canonical supernode postorder, `test/test_ndtree.jl`), unreviewed and without CI. Start from it or from scratch; it was not checked against this text. -### T31 — oneAPI and Metal extensions, allocation-free KA dense fallbacks (issues #112, #53) `[ ]` +### T31 — oneAPI and Metal extensions, allocation-free KA dense fallbacks (issues #112, #53) `[~]` The remainder of T23, split off by owner decision when PR #113 delivered the AMDGPU extension and the per-backend local-memory cap alone (#112). Three @@ -3329,6 +3329,90 @@ both extensions (no hardware needed to load them on Linux; Metal by review only). Mark both as untested in the Report; the `:ka` allocation budget and the `select_impl` order are tested on the KA CPU backend. Closes #112 and #53. +### Report +- Status: [~] — parts (1) oneAPI and (2) Metal are done (PR #136); part (3), the allocation-free `:ka` + fallbacks (#53), is not. `:auto` prefers `:ka` on oneAPI for LAPACK (below), not for every op (#112). +- What was built (oneAPI, part 1): `ext/SparseDirectSolverOneAPIExt.jl` (weak dependency `oneAPI`, compat + `2.9.2`, i.e. oneAPI.jl `main`, the only version with KernelAbstractions 0.10 support; the released 2.9.2 is + KA 0.9 and does not co-resolve): `CSR(::oneSparseMatrixCSR)`/`CSR(::oneSparseMatrixCSC)` sharing the + arrays, `oneSparseMatrixCSR(::CSR)`, `to_backend` (Float64/ComplexF64 rejected on devices without FP64), + the T13 forwarders (`DirectSolver`, `update!`, `cholesky`/`ldlt`/`lu` and the wrappers), + `max_local_bytes(::oneAPIBackend)` from `maxSharedLocalMemory` (128 KiB on the Max 1550), oneMKL bindings + for `gemm`, `syrk`, `herk`, `trsm`, `potrf`/`potrf_info!`/`getrf` (through the C shims, which keep the + LAPACK info that oneAPI.jl's wrappers drop), strided-batched `gemm` and `trsm`. Core: the + `_auto_vendor_dense(op, backend, T)` hook (`src/dense/interface.jl`): every oneMKL call waits for its queue + and LAPACK allocates a scratchpad, so on oneAPI `:auto` keeps `potrf`/`getrf` on the KA kernels and the + generic LinearAlgebra path is off; `impl = :vendor` still reaches oneMKL. `dense_impls` lists only the + `:auto`-eligible vendor bindings, so `:auto` stays its first entry and `:ka` its last. Tests: oneAPI in + `test/backends.jl`, the oneMKL column and the `:auto` policy in `test_dense.jl`, the extension check, the + local-memory cap. CI: a `oneapi` leg of `test-gpu` on the runner `headroom` (labels `gpu`, `oneapi`; + oneAPI.jl `main`) and a `test-metal` job on GitHub's hosted `macos-latest`, Float32 only. Docs: + `lib/oneapi.md`, backend table, README. Base commit (also in this PR): `backend isa + KernelAbstractions.GPU` was true for the CPU backend under KA 0.10 (`GPU` is an alias of `Backend` there), + so the CPU took the device symbolic-map route; now `!(backend isa CPU)`. KA compat stays `"0.9, 0.10"`. +- Tests (oneAPI): Aurora compute node, Intel Data Center GPU Max 1550 (one tile, `ZE_FLAT_DEVICE_HIERARCHY=FLAT`), + oneAPI.jl `main` (94aa6f2) with its support library built by the oneAPI.jl CI build step, KA 0.10.0, all + four element types, `SDS_TEST_CPU=0 SDS_TEST_SKIP=test_aqua`, 16 workers: 64,637 pass / 0 fail / 4 broken + (the `@test_skip` of `test_api.jl`: no device allocation counter on oneAPI), 6.5 min. The CPU backend was + not run on Aurora by owner decision (CI's CPU job covers it); a 48-worker CPU run there aborted in POCL's + thread-pool start-up (`pthread_scheduler_init`, every worker starts a pool of all 208 hardware threads). +- What was built (Metal, part 2): `ext/SparseDirectSolverMetalExt.jl` (weak dependency `Metal`, compat `1.11`: the + offset-aware `MPSMatrix(buffer, desc, offset)` constructor and the `MPSMatrixDecompositionCholesky` + wrapper used by the extension): Float32 CSR upload (`to_backend`), `max_local_bytes(::MetalBackend) = 32768`, + backend hooks in `src/dense/capabilities.jl` (`_backend_schedule_options`, `_allow_batched_dense`, + `_allow_generic_dense`, `_serial_cholesky_stats`, `_host_reduce_cholesky_info`, `_dense_solve_front`, + `_solve_workgroup`, `_offset_potrf_status`/`_offset_potrf_plan`/`_offset_potrf_info!`), MPS Cholesky for + single-matrix regime-C fronts at least 64 wide (plans and the status buffer preallocated in + `allocate_numeric`, pivots checked after MPS by `ka_chol_check_info!`, #127 semantics), every Metal front in + the fused level solve, 128-item solve workgroups; Metal in `test/backends.jl`, docs, `bench/metal_*.jl`. + `_tri_decode` (`src/symbolic/schedule.jl`) uses an integer binary search (`_kernel_isqrt`) instead of + `Base.isqrt`, whose Julia 1.13 implementation goes through Float64 and does not compile for Metal, on + every backend. +- Tests (Metal, M4 Max, Float32, implementer's run): full suite 23,399 pass / 0 fail / 1 broken (existing). + Owner's machine after the review fixes, `SDS_TEST_ONLY="test_helpers,test_numeric_cholesky_a, + test_numeric_cholesky_b,test_numeric_cholesky_c,test_solve,test_api"` (+ `test_aqua` on CPU), all four + element types: CPU 5,817 pass / 0 fail; CUDA (`SDS_TEST_CPU=0`, RTX 4080) 5,584 pass / 0 fail. + Full suite and AMDGPU: CI on the PR. oneAPI: above. +- Measurements (M4 Max, Metal.jl 1.11.1, Float32, warmed median, `lap2d_300`, n = 90,000): refactorization + 142 → 62 ms (2.29×), solve 5.19 → 5.18 ms; 113 regime-C fronts, 18 at least 64 wide; MPS potrf 5–14× + faster than KA potrf at widths 128–256, crossover near 64; MPS trsm was correct but slower end to end + (not used). RTX 4080 Float64 context: SDS 32.9 / 3.48 ms, cuDSS 3.96 / 0.675 ms (not precision parity). +- Deviations from PLAN.md / this task: + * `:auto` on Metal resolves `gemm`/`syrk` to `:generic` (MPS through `mul!` on views, allocating MPS + descriptors per call), not `:ka` as the task asks. The extension states it is faster; the comparison + with `:ka` is not recorded and has to be measured. POTRF/TRSM stay `:ka` (POTRF: MPS for wide fronts, + above). + * The MPS potrf route creates an `MPSCommandBuffer` wrapper per wide front per refactorization, about + 2.1 MB of host allocation per refactorization on `lap2d_300` — an exception to "no allocation in the + numeric phase"; needs a reusable command-buffer encoding (follow-up). + * Float32 only: `ComplexF32` is not supported on Metal yet (`to_backend` raises `NotSupportedError`). + * LDLᵀ/LDLᴴ and LU on Metal run in regime C only (vendor factorization, no subtree budgets): their fused + kernels bind 32 buffers, Apple GPUs allow 31. Cholesky with default options uses + `regime_c_width = 32`, `regime_c_rows = 128` on Metal. + * Uniform batches on Metal: KA dense kernels only, and the Cholesky statistics/`info` reduction runs in + a serial kernel plus a host reduction (`_serial_cholesky_stats`, `_host_reduce_cholesky_info`). + * `test/test_numeric_cholesky_a.jl`: `NUMERIC_ABC_OPTS` gains `subtree_max_fronts = 32` on every backend. +- Open issues / follow-ups (oneAPI): upstream oneAPI.jl: (a) the `*_scratchpad_size` C shims do not catch + oneMKL exceptions, so `sytrf` (unsupported on the Max 1550) aborts the process — no `vendor_sytrf!` binding + here; (b) the generated `onemklZtrsm_batch_strided` binding declares `Ptr{ComplexF32}` arguments (and the + strided-batch trsm bindings host `Ptr`s, not `ZePtr`s) — the extension passes typed raw pointers; (c) every + shim ends in `wait_and_throw`, so a vendor call is a host synchronization; (d) no cumulative device + allocation counter for `device_allocated`. Batched `potrf`/`getrf` report one status per batch, so uniform + batches use the KA kernels. Width-dependent `:auto` (oneMKL only for wide fronts) is not measured. The POCL + thread oversubscription of the KA 0.10 CPU backend with many test workers (`JULIA_KA_CPU_THREADS`). +- Open issues / follow-ups (Metal): reusable MPS command-buffer encoding (allocation above); why the 256-item + `_cholesky_stats_kernel!` reduction needs the serial variant for Metal batches (kernel bug hidden by + CUDA/AMD warps, or a Metal compiler limit) — not yet diagnosed; confirm on CUDA that the `_tri_decode` + change does not slow regime A/B (`bench/`); `ComplexF32` on Metal; oneAPI and part (3) of this task. + Review fixes (PR #136): the MPS status output had been written into `N.aux` (clobbering the `max |aᵢⱼ|` + statistic); it now has its own `MPSMatrixDecompositionStatus` buffer. The test-helper round trip keeps + `ComplexF64` on CUDA/AMDGPU. The budget-class test in `test_numeric_cholesky_a.jl` expects the + rejection of an explicit budget above `max_local_bytes` (failed on AMDGPU once the test helpers passed + the backend's `max_local` to the analysis). +- Suggested plan changes: part (3) (#53) as its own task; PLAN §2.6 oneAPI column: gemm/syrk/herk/trsm, + potrf/getrf and strided-batched gemm/trsm yes, sytrf and per-member batched potrf/getrf no; record + Metal's 31-buffer argument limit next to the 32 KiB local-memory cap in PLAN §2.7. + ### External — MadNLPGPU integration (in the MadNLP repository) `[ ]` Add a `SparseDirectSolver`-backed `AbstractLinearSolver` next to diff --git a/bench/metal_float32.jl b/bench/metal_float32.jl new file mode 100644 index 0000000..d557ed9 --- /dev/null +++ b/bench/metal_float32.jl @@ -0,0 +1,59 @@ +# Float32 CPU/Metal comparison on generated SPD matrices. +# +# Run from an environment containing SparseDirectSolver and Metal: +# julia +1.13 --project= bench/metal_float32.jl + +using LinearAlgebra +using KernelAbstractions +using Metal +using Random +using SparseArrays +using SparseDirectSolver +using Statistics + +function laplacian2d(n::Int) + N = n * n + I = Int[] + J = Int[] + V = Float32[] + for y in 1:n, x in 1:n + k = x + (y - 1) * n + push!(I, k); push!(J, k); push!(V, 4.0f0) + for (xx, yy) in ((x - 1, y), (x + 1, y), (x, y - 1), (x, y + 1)) + if 1 <= xx <= n && 1 <= yy <= n + push!(I, k); push!(J, xx + (yy - 1) * n); push!(V, -1.0f0) + end + end + end + return sparse(I, J, V, N, N) +end + +function run_case(label, backend, A; samples = 10) + C = to_backend(tril(A), backend) + solver = DirectSolver(C, "SPD", 'L') + b = rand(Float32, size(A, 1)) + bd = KernelAbstractions.allocate(backend, Float32, length(b)) + xd = similar(bd) + copyto!(bd, b) + phase(p) = execute!(p, solver, xd, bd; asynchronous = false) + + ta = @elapsed phase("analysis") + phase("factorization") + phase("solve") + tf = median([@elapsed phase("refactorization") for _ in 1:samples]) + ts = median([@elapsed phase("solve") for _ in 1:samples]) + x = Array(xd) + residual = norm(b - A * x) / norm(b) + residual <= sqrt(eps(Float32)) || error("$label residual $residual") + println(join((label, size(A, 1), nnz(A), ta, tf, ts, residual), ',')) +end + +if abspath(PROGRAM_FILE) == @__FILE__ + Random.seed!(666) + println("backend,n,nnz,analysis_s,refactorization_s,solve_s,relative_residual") + for grid in (40, 80, 120) + A = laplacian2d(grid) + run_case("CPU-$grid", KernelAbstractions.CPU(), A) + run_case("Metal-$grid", MetalBackend(), A) + end +end diff --git a/bench/metal_mps_proto.jl b/bench/metal_mps_proto.jl new file mode 100644 index 0000000..d25b8b1 --- /dev/null +++ b/bench/metal_mps_proto.jl @@ -0,0 +1,85 @@ +# Low-level MPS prototype on offset/strided panel storage. This intentionally +# does not modify the solver; it gates direct MPS integration. + +using LinearAlgebra +using Metal +using Metal.MPS +using Metal.MTL +using SparseDirectSolver +using Statistics + +function mps_panel(A::MtlVector{T}, p0, rows, cols, ld) where {T} + desc = MPSMatrixDescriptor(cols, rows, sizeof(T) * ld, T) + return MPSMatrix(A, desc, (p0 - 1) * sizeof(T)) +end + +function mps_potrf!(A::MtlVector{Float32}, p0, n, ld) + dev = Metal.device() + status = MtlArray{MPS.MPSMatrixDecompositionStatus, 0, SharedStorage}(undef) + mat = mps_panel(A, p0, n, n, ld) + cmdbuf = MPSCommandBuffer(global_queue(dev)) + encode!(cmdbuf, MPSMatrixDecompositionCholesky(dev, false, n), mat, mat, status) + commit!(cmdbuf) + return status +end + +function mps_potrf_plan(A::MtlVector{Float32}, p0, n, ld) + dev = Metal.device() + status = MtlArray{MPS.MPSMatrixDecompositionStatus, 0, SharedStorage}(undef) + mat = mps_panel(A, p0, n, n, ld) + kernel = MPSMatrixDecompositionCholesky(dev, false, n) + return (; dev, status, mat, kernel) +end + +function mps_potrf!(plan) + cmdbuf = MPSCommandBuffer(global_queue(plan.dev)) + encode!(cmdbuf, plan.kernel, plan.mat, plan.mat, plan.status) + commit!(cmdbuf) + return plan.status +end + +function panel_buffer(A, f, w; lead = 7) + buf = zeros(Float32, lead + f * w + 5) + buf[(lead + 1):(lead + f * w)] .= vec(A) + return MtlArray(buf), lead + 1 +end + +function run_potrf(n; f = n + 13, samples = 20) + G = randn(Float32, n, n) + S = G * G' + n * I + P = zeros(Float32, f, n) + P[1:n, :] .= S + buf, p0 = panel_buffer(P, f, n) + ref = cholesky(Symmetric(S)).L + + mps_potrf!(buf, p0, n, f) + Metal.synchronize() + H = reshape(Array(buf)[p0:(p0 + f * n - 1)], f, n) + err = norm(LowerTriangular(H[1:n, :]) - ref) / norm(ref) + + plan = mps_potrf_plan(buf, p0, n, f) + times = Float64[] + for _ in 1:samples + copyto!(view(buf, p0:(p0 + f * n - 1)), MtlArray(vec(P))) + push!(times, @elapsed begin + mps_potrf!(plan) + Metal.synchronize() + end) + end + info = MtlVector{Int32}(undef, 1) + F11 = view(reshape(view(buf, p0:(p0 + f * n - 1)), f, n), 1:n, 1:n) + ka = Float64[] + for _ in 1:samples + copyto!(view(buf, p0:(p0 + f * n - 1)), MtlArray(vec(P))) + push!(ka, @elapsed begin + SparseDirectSolver.ka_potrf!('L', F11, info) + Metal.synchronize() + end) + end + return median(times), minimum(times), median(ka), minimum(ka), err +end + +println("n,mps_median_s,mps_min_s,ka_median_s,ka_min_s,error") +for n in (16, 32, 64, 96, 128, 192, 256) + println(join((n, run_potrf(n)...), ',')) +end diff --git a/bench/metal_tune.jl b/bench/metal_tune.jl new file mode 100644 index 0000000..8e09ddd --- /dev/null +++ b/bench/metal_tune.jl @@ -0,0 +1,33 @@ +# Metal scheduling sweep for Float32 Cholesky. Correctness gates every result. + +include("metal_float32.jl") + +function tune_case(A, label, opts; samples = 7) + solver = DirectSolver(to_backend(tril(A), MetalBackend()), "SPD", 'L') + solver.options = opts + b = rand(Float32, size(A, 1)) + bd, xd = MtlArray(b), MtlArray(similar(b)) + execute!("analysis", solver, nothing, nothing) + execute!("factorization", solver, nothing, nothing; asynchronous = false) + execute!("solve", solver, xd, bd; asynchronous = false) + residual = norm(b - A * Array(xd)) / norm(b) + residual <= sqrt(eps(Float32)) || error("$label residual $residual") + tf = median([@elapsed execute!("refactorization", solver, nothing, nothing; asynchronous = false) + for _ in 1:samples]) + ts = median([@elapsed execute!("solve", solver, xd, bd; asynchronous = false) + for _ in 1:samples]) + sc = solver.host_symbolic.schedule + regimes = count.((==(SparseDirectSolver.REGIME_A), ==(SparseDirectSolver.REGIME_B), + ==(SparseDirectSolver.REGIME_C)), Ref(sc.regime)) + println(join((label, regimes..., SparseDirectSolver.nsubtrees(sc), SparseDirectSolver.nlaunches(sc), + tf, ts, residual), ',')) +end + +println("label,A,B,C,subtrees,launches,refactorization_s,solve_s,residual") +A = laplacian2d(120) +for p in (512, 1024, 2048, 4096, 8192, 16384), mf in (0, 16, 32, 64) + tune_case(A, "p$(p)-f$(mf)", Options(subtree_parallelism = p, subtree_max_fronts = mf)) +end +for (cw, cr) in ((32, 128), (32, 256), (64, 256), (64, 512), (64, 1024), (128, 1024)) + tune_case(A, "c$(cw)x$(cr)", Options(regime_c_width = cw, regime_c_rows = cr)) +end diff --git a/docs/Project.toml b/docs/Project.toml index 49ad890..891afd6 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -6,6 +6,7 @@ KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseDirectSolver = "dc8e0fd9-41f1-440b-9624-22817132f700" +oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [sources] SparseDirectSolver = {path = ".."} @@ -14,3 +15,4 @@ SparseDirectSolver = {path = ".."} AMDGPU = "2" CUDA = "6" Documenter = "1" +oneAPI = "2.9.2" diff --git a/docs/make.jl b/docs/make.jl index 3efe95b..b32f3ae 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -2,8 +2,8 @@ # julia --project=docs docs/make.jl # The examples run on the KernelAbstractions CPU backend and on CUDA (the docs # are built on the self-hosted `cuda` runner, see .github/workflows/Documentation.yml). -# AMDGPU.jl is only loaded, for the docstrings of the AMDGPU extension: it loads -# without AMD hardware. +# AMDGPU.jl and oneAPI.jl are only loaded, for the docstrings of their extensions: +# they load without AMD or Intel hardware. using Documenter using LinearAlgebra @@ -11,6 +11,7 @@ using SparseArrays using SparseDirectSolver using CUDA using AMDGPU +using oneAPI CUDA.functional() || error("the documentation examples need a functional CUDA.jl") CUDA.versioninfo() @@ -18,6 +19,8 @@ CUDA.versioninfo() const CUDAExt = Base.get_extension(SparseDirectSolver, :SparseDirectSolverCUDAExt) const AMDGPUExt = Base.get_extension(SparseDirectSolver, :SparseDirectSolverAMDGPUExt) AMDGPUExt === nothing && error("the AMDGPU extension did not load") +const OneAPIExt = Base.get_extension(SparseDirectSolver, :SparseDirectSolverOneAPIExt) +OneAPIExt === nothing && error("the oneAPI extension did not load") # the cuDSS comparison plot is regenerated by hand in bench/ (bench/README.md) mkpath(joinpath(@__DIR__, "src", "assets")) @@ -37,7 +40,7 @@ makedocs( size_threshold = 1024 * 1024, size_threshold_warn = 512 * 1024, ), - modules = [SparseDirectSolver, CUDAExt, AMDGPUExt], + modules = [SparseDirectSolver, CUDAExt, AMDGPUExt, OneAPIExt], repo = Remotes.GitHub("exanauts", "SparseDirectSolver.jl"), checkdocs = :exports, pages = [ @@ -57,6 +60,7 @@ makedocs( "LinearAlgebra interface" => "lib/linearalgebra.md", "CUDA extension" => "lib/cuda.md", "AMDGPU extension" => "lib/amdgpu.md", + "oneAPI extension" => "lib/oneapi.md", "Internals" => "lib/internals.md", ], ], diff --git a/docs/src/index.md b/docs/src/index.md index 94978b4..d31ba91 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -18,9 +18,9 @@ cuDSS it is open source and not tied to one vendor: it runs on CUDA, AMDGPU, oneAPI, Metal and the KernelAbstractions CPU backend. !!! warning "Status" - Version 0.1 is under construction. The CPU backend, CUDA and AMDGPU are - tested; the oneAPI and Metal extensions, non-uniform batches and mixed - precision are not there yet. Unsupported structures, phases and parameters raise + Version 0.1 is under construction. The CPU backend, CUDA, AMDGPU and Metal + (`Float32`, tested on an M4 Max) are tested; the oneAPI extension, + non-uniform batches and mixed precision are not there yet. Unsupported structures, phases and parameters raise [`NotSupportedError`](@ref) rather than falling back silently. ## Features diff --git a/docs/src/lib/oneapi.md b/docs/src/lib/oneapi.md new file mode 100644 index 0000000..b5312f3 --- /dev/null +++ b/docs/src/lib/oneapi.md @@ -0,0 +1,13 @@ +# oneAPI extension + +Loaded with `using oneAPI`. It accepts the sparse matrices of oneMKL, routes the +level-3 BLAS operations of large fronts (`gemm`, `syrk`/`herk`, `trsm` and +their strided batches) to oneMKL, and lets the regime-A subtree kernels use the +64 KiB local-memory class on Intel GPUs. Every oneMKL call waits for its queue +and the LAPACK routines allocate a scratchpad, so the Cholesky and LU panel +factorizations stay on the KernelAbstractions kernels unless `impl = :vendor` +is requested. + +```@autodocs +Modules = [Base.get_extension(SparseDirectSolver, :SparseDirectSolverOneAPIExt)] +``` diff --git a/docs/src/man/backends.md b/docs/src/man/backends.md index 6db479d..577b699 100644 --- a/docs/src/man/backends.md +++ b/docs/src/man/backends.md @@ -9,7 +9,8 @@ copy and every numeric phase runs where they live. | CPU (`KernelAbstractions.CPU()`) | `SparseMatrixCSC`, [`CSR`](@ref) of `Vector`s | KernelAbstractions kernels | tested | | CUDA (`using CUDA`) | `CuSparseMatrixCSR`, `CuSparseMatrixCSC`, [`CSR`](@ref) of `CuArray`s | cuBLAS, cuSOLVER | tested | | AMDGPU (`using AMDGPU`) | `ROCSparseMatrixCSR`, `ROCSparseMatrixCSC`, [`CSR`](@ref) of `ROCArray`s | rocBLAS, rocSOLVER | tested (MI300X) | -| oneAPI, Metal | [`CSR`](@ref) of device arrays | KernelAbstractions kernels | planned extensions | +| Metal (`using Metal`) | [`CSR`](@ref) of `MtlArray`s | Metal.jl / KernelAbstractions kernels | Float32, tested (M4 Max) | +| oneAPI (`using oneAPI`) | `oneSparseMatrixCSR`, `oneSparseMatrixCSC`, [`CSR`](@ref) of `oneArray`s | oneMKL BLAS, KernelAbstractions kernels for the panel factorizations | tested (Intel Data Center GPU Max 1550) | [`to_backend`](@ref) moves a `SparseMatrixCSC` to a backend: @@ -25,4 +26,15 @@ assembly (an atomic-free variant exists for the one atomic kernel, selected by numeric and solve phases. The symbolic analysis runs on the host. Element types are `Float32`, `Float64`, `ComplexF32` and `ComplexF64`; index -types are `Int32` and `Int64`. +types are `Int32` and `Int64`. Metal currently supports `Float32` factors; +Apple GPUs do not expose Float64 arithmetic. + +On Metal, the default SPD schedule uses narrower regime-C thresholds than the +CUDA/ROCm schedule (`32` columns and `128` rows), Metal.jl's panel GEMM/SYRK, +and fused level solves for every front. Regime-C Cholesky fronts at least 64 +columns wide use offset-aware MPS decomposition objects cached at analysis. On +an M4 Max, the 90,000-row `lap2d_300` warmed refactorization fell from 142 ms to +62 ms while its 5.2 ms solve time was preserved. Explicit schedule options are +never replaced by this backend tuning. LDLT and LU currently use the regime-C +factorization path because their fused kernels exceed Metal's 31-buffer argument +limit. diff --git a/ext/SparseDirectSolverMetalExt.jl b/ext/SparseDirectSolverMetalExt.jl new file mode 100644 index 0000000..6c7fb55 --- /dev/null +++ b/ext/SparseDirectSolverMetalExt.jl @@ -0,0 +1,115 @@ +module SparseDirectSolverMetalExt + +using SparseDirectSolver +using SparseDirectSolver: CSR, INDEX_ONE +using SparseArrays +using Metal +using Metal.MPS +using Metal.MTL + +""" + to_backend(A::SparseMatrixCSC, ::MetalBackend; index = 'O') -> CSR + +Build the CSR arrays of `A` on the host and upload them to Metal. Metal has no +vendor sparse matrix container, so the returned [`CSR`](@ref) directly owns +`MtlVector` index and value arrays. +""" +function SparseDirectSolver.to_backend(A::SparseMatrixCSC, ::MetalBackend; index = INDEX_ONE) + eltype(A) === Float32 || throw(SparseDirectSolver.NotSupportedError( + "Metal currently supports Float32 matrices; got $(eltype(A))")) + B = CSR(A; index) + return CSR(MtlVector(B.rowptr), MtlVector(B.colval), MtlArray(B.nzval), B.nrows, B.ncols; + index = B.index, transposed = B.transposed) +end + +""" + max_local_bytes(::MetalBackend) -> Int + +Metal's maximum static threadgroup-memory allocation per workgroup (32 KiB on +Apple GPUs). The symbolic scheduler uses this to exclude the 48 and 64 KiB +regime-A kernel classes. +""" +SparseDirectSolver.max_local_bytes(::MetalBackend) = 32768 +SparseDirectSolver._allow_batched_dense(::MetalBackend) = false +SparseDirectSolver._serial_cholesky_stats(::MetalBackend, nbatch::Integer) = nbatch > 1 +SparseDirectSolver._host_reduce_cholesky_info(::MetalBackend, nbatch::Integer) = nbatch > 1 +SparseDirectSolver._dense_solve_front(::MetalBackend, schedule::SparseDirectSolver.Schedule, s::Integer) = false +SparseDirectSolver._solve_workgroup(::MetalBackend) = 128 + +const MPS_POTRF_MIN_WIDTH = 64 +SparseDirectSolver._offset_potrf_enabled(::MetalBackend) = true + +struct MetalPotrfPlan{M, K, S} + panel::M + kernel::K + status::S # keeps the MTLBuffer of the MPS status output alive +end + +# MPS writes an `MPSMatrixDecompositionStatus` per encode; one buffer of that type is +# shared by the wide fronts of a factorization. It is never read: the pivots are +# checked afterwards by `ka_chol_check_info!`. +SparseDirectSolver._offset_potrf_status(::MetalBackend) = + MtlArray{MPS.MPSMatrixDecompositionStatus, 0, Metal.SharedStorage}(undef) + +function SparseDirectSolver._offset_potrf_plan(::MetalBackend, factor::MtlVector{Float32}, p0, f, w, + status::MtlArray{MPS.MPSMatrixDecompositionStatus, 0}) + w >= MPS_POTRF_MIN_WIDTH || return nothing + desc = MPSMatrixDescriptor(w, w, sizeof(Float32) * f, Float32) + panel = MPSMatrix(factor, desc, (p0 - 1) * sizeof(Float32)) + kernel = MPSMatrixDecompositionCholesky(Metal.device(), false, w) + return MetalPotrfPlan(panel, kernel, status) +end + +function SparseDirectSolver._offset_potrf_info!(::MetalBackend, plan::MetalPotrfPlan, info, idx) + fill!(view(info, idx:idx), Int32(0)) + # the KA kernels before (assembly) and after (`ka_chol_check_info!`, trsm) run on the + # same `global_queue`, which orders them with this command buffer + cmdbuf = MPSCommandBuffer(global_queue(Metal.device())) + encode!(cmdbuf, plan.kernel, plan.panel, plan.panel, Base.unsafe_convert(MTLBuffer, plan.status)) + commit!(cmdbuf) + return true +end + +# Metal.jl's MPS Cholesky and triangular solves work on full MtlMatrix objects, +# but panel views dispatch to host LAPACK/BLAS. Panel GEMM/SYRK are supported +# and faster (`:generic` allocates MPS descriptors per call on Metal); TRSM stays +# on KA, POTRF too except for single-matrix fronts at least `MPS_POTRF_MIN_WIDTH` +# wide (`_offset_potrf_plan`). Uniform batches use only KA dense kernels. +function SparseDirectSolver._allow_generic_dense(op::Symbol, ::MetalBackend, ::Type{T}) where {T} + return op === :gemm || op === :syrk +end + +function SparseDirectSolver._backend_schedule_options(backend::MetalBackend, structure::SparseDirectSolver.Structure, + opts::SparseDirectSolver.Options, nbatch::Integer) + out = SparseDirectSolver._backend_schedule_options(backend, structure, opts) + nbatch == 1 && return out + out === opts && (out = copy(opts)) + out.factorization_alg = SparseDirectSolver.FACTORIZATION_VENDOR + empty!(out.subtree_budgets) + return out +end + +# The fused LDLT/LU kernels bind 32 device buffers; Apple GPUs expose 31 +# indirect argument-buffer resources. Route those structures through the +# existing regime-C kernels. Cholesky's smaller fused kernels remain enabled. +function SparseDirectSolver._backend_schedule_options(::MetalBackend, structure::SparseDirectSolver.Structure, + opts::SparseDirectSolver.Options) + if structure in (SparseDirectSolver.STRUCTURE_SPD, SparseDirectSolver.STRUCTURE_HPD) + opts.regime_c_width == 64 && opts.regime_c_rows == 512 && + opts.subtree_budgets == SparseDirectSolver.DEFAULT_SUBTREE_BUDGETS && + opts.subtree_parallelism == SparseDirectSolver.SUBTREE_PARALLELISM && + opts.subtree_max_fronts == 0 && opts.memory_budget == -1 || return opts + out = copy(opts) + out.regime_c_width = 32 + out.regime_c_rows = 128 + return out + end + structure in (SparseDirectSolver.STRUCTURE_SYMMETRIC, SparseDirectSolver.STRUCTURE_HERMITIAN, + SparseDirectSolver.STRUCTURE_GENERAL) || return opts + out = copy(opts) + out.factorization_alg = SparseDirectSolver.FACTORIZATION_VENDOR + empty!(out.subtree_budgets) + return out +end + +end # module SparseDirectSolverMetalExt diff --git a/ext/SparseDirectSolverOneAPIExt.jl b/ext/SparseDirectSolverOneAPIExt.jl new file mode 100644 index 0000000..e1292bf --- /dev/null +++ b/ext/SparseDirectSolverOneAPIExt.jl @@ -0,0 +1,274 @@ +module SparseDirectSolverOneAPIExt + +# oneAPI (Intel GPU) support (T31), mirroring `SparseDirectSolverAMDGPUExt`: CSR +# adapters for oneMKL sparse matrices, the public API on them (T13), the vendor dense +# bindings of the dense layer (PLAN §2.6) on oneMKL, and the backend's local memory +# per workgroup for the regime-A kernels (issue #60). +# +# Every oneMKL entry point of oneAPI.jl's support library waits for its queue before +# it returns (`wait_and_throw`), and the LAPACK routines need a scratchpad: a vendor +# call is a host synchronization, LAPACK also a device allocation. `impl = :auto` +# therefore keeps LAPACK (`potrf`, `getrf`) off the vendor path +# (`_auto_vendor_dense`); the level-3 BLAS calls (`gemm`, `syrk`/`herk`, `trsm` and +# their strided batches) stay on it. The generic LinearAlgebra path is disabled +# (`_allow_generic_dense`): it reaches the same oneMKL calls, plus the allocating +# LAPACK wrappers. +# +# LAPACK status: the C shims return oneMKL's `computation_error` info on the host; +# the Julia wrappers of oneAPI.jl drop it, so `potrf`/`getrf` call the C shims +# directly. Batched `potrf`/`getrf` report one status per batch, not per member: no +# binding (the KA kernels serve uniform batches). No `sytrf`: oneMKL rejects it on +# the Intel Data Center GPU Max with `unsupported_device`, thrown from the +# scratchpad-size shim, which does not catch it, so the process aborts. + +using SparseDirectSolver +using SparseDirectSolver: CSR, DirectSolver, INDEX_ONE, INDEX_ZERO, InvalidValueError, NotSupportedError +using SparseArrays +using LinearAlgebra +using KernelAbstractions +using oneAPI +using oneAPI: oneL0, oneMKL, Support, oneStridedMatrix, oneStridedArray, sycl_queue, global_queue +using oneAPI.oneL0: ZePtr +using oneAPI.oneMKL: oneSparseMatrixCSR, oneSparseMatrixCSC + +const SDS = SparseDirectSolver + +""" + CSR(A::oneSparseMatrixCSR) + +Wrap `A` without copying: the `CSR` shares `rowPtr`, `colVal` and `nzVal` +(one-based, `transposed = false`). +""" +function SparseDirectSolver.CSR(A::oneSparseMatrixCSR) + m, n = size(A) + return CSR(A.rowPtr, A.colVal, A.nzVal, m, n; index = INDEX_ONE, transposed = false) +end + +""" + CSR(A::oneSparseMatrixCSC) + +Wrap `A` (`m × n`) without copying as the CSR of `transpose(A)` (`n × m`, +`transposed = true`), sharing `colPtr`, `rowVal` and `nzVal`. +""" +function SparseDirectSolver.CSR(A::oneSparseMatrixCSC) + m, n = size(A) + return CSR(A.colPtr, A.rowVal, A.nzVal, n, m; index = INDEX_ONE, transposed = true) +end + +""" + oneSparseMatrixCSR(A::CSR) + +The stored CSR matrix of `A` (the `transposed` flag is not applied) as a +`oneSparseMatrixCSR`. One-based arrays are shared; zero-based index arrays are +copied and rebased. Batches (`nbatch(A) > 1`) are rejected. +""" +function oneMKL.oneSparseMatrixCSR(A::CSR{T, INT, <:oneVector{INT}, <:oneArray{T}}) where {T, INT} + nbatch(A) == 1 || throw(InvalidValueError("cannot convert a batch of $(nbatch(A)) matrices to oneSparseMatrixCSR")) + nzval = A.nzval isa oneVector ? A.nzval : vec(A.nzval) + rowptr, colval = A.rowptr, A.colval + if A.index == INDEX_ZERO + rowptr = rowptr .+ one(INT) + colval = colval .+ one(INT) + end + return oneSparseMatrixCSR{T, INT}(rowptr, colval, nzval, (A.nrows, A.ncols)) +end + +_check_eltype(::Type{T}) where {T} = + real(T) === Float64 && !KernelAbstractions.supports_float64(oneAPIBackend()) && + throw(NotSupportedError("the oneAPI device $(oneAPI.device()) has no Float64 support; got $T")) + +""" + to_backend(A::SparseMatrixCSC, ::oneAPIBackend; index = 'O') -> CSR + +CSR arrays of `A` built on the host, then uploaded as `oneVector`s. Float64 and +ComplexF64 raise `NotSupportedError` on a device without Float64 support. +""" +function SparseDirectSolver.to_backend(A::SparseMatrixCSC{T}, ::oneAPIBackend; index = INDEX_ONE) where {T} + _check_eltype(T) + B = CSR(A; index) + return CSR(oneVector(B.rowptr), oneVector(B.colval), oneVector(B.nzval), B.nrows, B.ncols; + index = B.index, transposed = B.transposed) +end + +""" + max_local_bytes(::oneAPIBackend) -> Int + +The shared local memory of one workgroup on the current oneAPI device +(`maxSharedLocalMemory` of its compute properties; 128 KiB on the Intel Data +Center GPU Max): the regime-A kernels may use their largest class, 64 KiB. +""" +SparseDirectSolver.max_local_bytes(::oneAPIBackend) = + Int(oneL0.compute_properties(oneAPI.device()).maxSharedLocalMemory) + +# --------------------------------------------------------------------------- +# public API (T13, PLAN §3.1): solver constructors, update!, the LinearAlgebra layer + +""" + DirectSolver(A::oneSparseMatrixCSR, structure::String, view::Char; index = 'O') + DirectSolver(A::oneSparseMatrixCSC, structure::String, view::Char; index = 'O') + +Solver on `A`'s arrays, without copies (the oneAPI twin of +`DirectSolver(::CuSparseMatrixCSR, …)`). A `oneSparseMatrixCSC` is read as the +CSR of its transpose with the view flipped, see [`DirectSolver`](@ref). +""" +SparseDirectSolver.DirectSolver(A::oneSparseMatrixCSR, structure, view; index = INDEX_ONE) = + DirectSolver(CSR(A), structure, view; index) +SparseDirectSolver.DirectSolver(A::oneSparseMatrixCSC, structure, view; index = INDEX_ONE) = + DirectSolver(CSR(A), structure, view; index) + +""" + update!(solver::DirectSolver, A::oneSparseMatrixCSR) + update!(solver::DirectSolver, A::oneSparseMatrixCSC) + +Point `solver` at the arrays of `A`, see [`update!`](@ref). +""" +SparseDirectSolver.update!(solver::DirectSolver, A::Union{oneSparseMatrixCSR, oneSparseMatrixCSC}) = + SparseDirectSolver.update!(solver, CSR(A)) + +""" + cholesky(A::oneSparseMatrixCSR, NoPivot(); view = 'F', check = false) -> DirectSolver + cholesky(Symmetric(A::oneSparseMatrixCSR)) / cholesky(Hermitian(A::oneSparseMatrixCSR)) + +LLᵀ/LLᴴ factorization of `A` on the Intel GPU; the wrappers pass their `uplo` as +the view. See `cholesky(::CSR)`. +""" +LinearAlgebra.cholesky(A::oneSparseMatrixCSR, p::NoPivot = NoPivot(); view::Char = 'F', check::Bool = false) = + cholesky(CSR(A), p; view, check) +LinearAlgebra.cholesky(A::Symmetric{T, <:oneSparseMatrixCSR{T}}, p::NoPivot = NoPivot(); + check::Bool = false) where {T <: Union{Float32, Float64}} = + cholesky(CSR(A.data), p; view = A.uplo, check) +LinearAlgebra.cholesky(A::Hermitian{T, <:oneSparseMatrixCSR{T}}, p::NoPivot = NoPivot(); check::Bool = false) where {T} = + cholesky(CSR(A.data), p; view = A.uplo, check) + +""" + ldlt(A::oneSparseMatrixCSR; view = 'F', check = false) -> DirectSolver + ldlt(Symmetric(A::oneSparseMatrixCSR)) / ldlt(Hermitian(A::oneSparseMatrixCSR)) + +LDLᵀ/LDLᴴ factorization of `A` on the Intel GPU; the wrappers pass their `uplo` +as the view. See `ldlt(::CSR)`. +""" +LinearAlgebra.ldlt(A::oneSparseMatrixCSR; view::Char = 'F', check::Bool = false) = ldlt(CSR(A); view, check) +LinearAlgebra.ldlt(A::Symmetric{T, <:oneSparseMatrixCSR{T}}; check::Bool = false) where {T <: Real} = + ldlt(CSR(A.data); view = A.uplo, check) +LinearAlgebra.ldlt(A::Hermitian{T, <:oneSparseMatrixCSR{T}}; check::Bool = false) where {T} = + ldlt(CSR(A.data); view = A.uplo, check) + +""" + lu(A::oneSparseMatrixCSR; check = false) -> DirectSolver + +`L D U` factorization (structure `"G"`) of `A` on the Intel GPU. See `lu(::CSR)`. +""" +LinearAlgebra.lu(A::oneSparseMatrixCSR; check::Bool = false) = lu(CSR(A); check) + +# --------------------------------------------------------------------------- +# dense implementation policy (see the header) + +SparseDirectSolver._allow_generic_dense(op::Symbol, ::oneAPIBackend, ::Type{T}) where {T} = false +SparseDirectSolver._auto_vendor_dense(op::Symbol, ::oneAPIBackend, ::Type{T}) where {T} = + op in (:gemm, :syrk, :herk, :trsm, :gemm_strided_batched, :trsm_strided_batched) + +# --------------------------------------------------------------------------- +# vendor dense bindings (see `src/dense/vendor.jl` for the contracts) + +const MklT = Union{Float32, Float64, ComplexF32, ComplexF64} +const MklC = Union{ComplexF32, ComplexF64} + +_lda(A) = max(1, stride(A, 2)) +_queue(A) = global_queue(oneAPI.context(A), oneAPI.device(A)) +# device address of `A` for the support-library calls declared with host `Ptr` arguments +_rawptr(A::oneStridedArray{T}) where {T} = reinterpret(Ptr{T}, Base.unsafe_convert(ZePtr{T}, A)) + +SDS.vendor_gemm!(tA::Char, tB::Char, α, A::oneStridedMatrix{T}, B::oneStridedMatrix{T}, β, + C::oneStridedMatrix{T}) where {T <: MklT} = + oneMKL.gemm!(tA, tB, T(α), A, B, T(β), C) +SDS.vendor_syrk!(uplo::Char, α, A::oneStridedMatrix{T}, β, C::oneStridedMatrix{T}) where {T <: MklT} = + oneMKL.syrk!(uplo, 'N', T(α), A, T(β), C) +SDS.vendor_herk!(uplo::Char, α, A::oneStridedMatrix{T}, β, C::oneStridedMatrix{T}) where {T <: MklC} = + oneMKL.herk!(uplo, 'N', real(T)(α), A, real(T)(β), C) +SDS.vendor_trsm!(side::Char, uplo::Char, trans::Char, diag::Char, α, A::oneStridedMatrix{T}, + B::oneStridedMatrix{T}) where {T <: MklT} = + oneMKL.trsm!(side, uplo, trans, diag, T(α), A, B) + +for (elty, potrf, potrf_size, getrf, getrf_size) in + ((:Float32, :onemklSpotrf, :onemklSpotrf_scratchpad_size, :onemklSgetrf, :onemklSgetrf_scratchpad_size), + (:Float64, :onemklDpotrf, :onemklDpotrf_scratchpad_size, :onemklDgetrf, :onemklDgetrf_scratchpad_size), + (:ComplexF32, :onemklCpotrf, :onemklCpotrf_scratchpad_size, :onemklCgetrf, :onemklCgetrf_scratchpad_size), + (:ComplexF64, :onemklZpotrf, :onemklZpotrf_scratchpad_size, :onemklZgetrf, :onemklZgetrf_scratchpad_size)) + @eval begin + # LAPACK info of oneMKL's potrf (0, the failing column, or -1 for a SYCL error) + function _potrf_info(uplo::Char, A::oneStridedMatrix{$elty}) + n = LinearAlgebra.checksquare(A) + q = _queue(A) + lwork = Support.$potrf_size(sycl_queue(q), uplo, n, _lda(A)) + work = oneVector{$elty}(undef, lwork) + info = Support.$potrf(sycl_queue(q), uplo, n, A, _lda(A), work, lwork) + oneAPI.unsafe_free!(work) + return Int(info) + end + + function _getrf_info(A::oneStridedMatrix{$elty}, ipiv::oneVector{Int64}) + m, n = size(A) + q = _queue(A) + lwork = Support.$getrf_size(sycl_queue(q), m, n, _lda(A)) + work = oneVector{$elty}(undef, lwork) + info = Support.$getrf(sycl_queue(q), m, n, A, _lda(A), ipiv, work, lwork) + oneAPI.unsafe_free!(work) + return Int(info) + end + + end +end + +SDS.vendor_potrf!(uplo::Char, A::oneStridedMatrix{<:MklT}) = _potrf_info(uplo, A) + +# the status is a host value (see the header): written into `info[idx]` on the device +function SDS.vendor_potrf_info!(uplo::Char, A::oneStridedMatrix{<:MklT}, info::oneVector{Int32}, idx::Integer) + 1 <= idx <= length(info) || throw(DimensionMismatch("info index $idx outside 1:$(length(info))")) + fill!(view(info, idx:idx), Int32(_potrf_info(uplo, A))) + return info +end + +# oneMKL pivots are `Int64`; other index vectors go through a temporary +function _with_int64(f, ipiv::oneVector, k::Integer) + ipiv isa oneVector{Int64} && length(ipiv) == k && return f(ipiv) + p = oneVector{Int64}(undef, k) + info = f(p) + view(ipiv, 1:k) .= p + oneAPI.unsafe_free!(p) + return info +end + +SDS.vendor_getrf!(A::oneStridedMatrix{<:MklT}, ipiv::oneVector{<:Integer}) = + _with_int64(p -> _getrf_info(A, p), ipiv, min(size(A)...)) + +function SDS.vendor_gemm_strided_batched!(tA::Char, tB::Char, α, A::oneStridedArray{T, 3}, B::oneStridedArray{T, 3}, β, + C::oneStridedArray{T, 3}) where {T <: MklT} + size(A, 3) == size(B, 3) == size(C, 3) || throw(DimensionMismatch("gemm_strided_batched: batch counts differ")) + size(C, 3) == 0 && return C + return oneMKL.gemm_strided_batched!(tA, tB, T(α), A, B, T(β), C) +end + +for (fname, elty) in ((:onemklStrsm_batch_strided, :Float32), (:onemklDtrsm_batch_strided, :Float64), + (:onemklCtrsm_batch_strided, :ComplexF32), (:onemklZtrsm_batch_strided, :ComplexF64)) + @eval function SDS.vendor_trsm_batched!(side::Char, uplo::Char, trans::Char, diag::Char, α, + A::oneStridedArray{$elty, 3}, B::oneStridedArray{$elty, 3}) + m, n, nb = size(B) + size(A, 1) == size(A, 2) == (side == 'L' ? m : n) && size(A, 3) == nb || + throw(DimensionMismatch("trsm_batched: A is $(size(A)), B is $(size(B)) (side $side)")) + nb == 0 && return B + alpha = Ref($elty(α)) + GC.@preserve A B alpha begin + # the declared host `Ptr`s carry device addresses, and every pointer is passed + # typed: oneAPI.jl's ComplexF64 binding declares them `Ptr{ComplexF32}`. The + # `sycl_queue` call comes after the argument preparation (oneAPI.jl's Julia → + # oneMKL ordering boundary). + a, b = _rawptr(A), _rawptr(B) + pα = Base.unsafe_convert(Ptr{$elty}, alpha) + Support.$fname(sycl_queue(_queue(B)), side, uplo, trans, diag, m, n, pα, a, _lda(A), + stride(A, 3), b, _lda(B), stride(B, 3), nb) + end + return B + end +end + +end # module SparseDirectSolverOneAPIExt diff --git a/src/dense/capabilities.jl b/src/dense/capabilities.jl index ea5f5b5..375dbd9 100644 --- a/src/dense/capabilities.jl +++ b/src/dense/capabilities.jl @@ -109,6 +109,20 @@ per-architecture tuning (`Fused.default_tuning`) keys on it. """ max_threads_per_multiprocessor(backend) = 0 +# Backend extensions may adjust scheduling options for hard kernel ABI limits. +_backend_schedule_options(backend, structure::Structure, opts::Options) = opts +_backend_schedule_options(backend, structure::Structure, opts::Options, nbatch::Integer) = + _backend_schedule_options(backend, structure, opts) +_allow_batched_dense(backend) = true +_serial_cholesky_stats(backend, nbatch::Integer) = false +_host_reduce_cholesky_info(backend, nbatch::Integer) = false +_dense_solve_front(backend, schedule, s::Integer) = takes_c_path(schedule, s) +_solve_workgroup(backend) = 64 +_offset_potrf_status(backend) = nothing +_offset_potrf_plan(backend, factor, p0, f, w, status) = nothing +_offset_potrf_info!(backend, plan, info, idx) = false +_offset_potrf_enabled(backend) = false + # --------------------------------------------------------------------------- # probes diff --git a/src/dense/interface.jl b/src/dense/interface.jl index 21ad7c3..6f4c2e2 100644 --- a/src/dense/interface.jl +++ b/src/dense/interface.jl @@ -32,6 +32,16 @@ const DENSE_OPS = ( const _DENSE_OP_FIELDS = Dict{Symbol, Tuple{Symbol, Symbol}}( op => (something(spec.generic, :none), something(spec.vendor, :none)) for (op, spec) in pairs(DENSE_OPS)) +# Backend extensions may reject a probed generic operation when the public +# entry point does not satisfy the solver's stronger panel/status contract. +_allow_generic_dense(op::Symbol, backend, ::Type{T}) where {T} = true + +# Backend extensions may keep `impl = :auto` off a probed vendor operation whose +# per-call cost does not fit the numeric phase (oneAPI: every oneMKL call waits for +# the queue, LAPACK allocates its scratchpad); an explicit `impl = :vendor` still +# reaches it. +_auto_vendor_dense(op::Symbol, backend, ::Type{T}) where {T} = true + function _flag_char(x, allowed::String, what::String) c = x isa Symbol ? (length(String(x)) == 1 ? only(String(x)) : '?') : x isa AbstractChar ? Char(x) : '?' c = uppercase(c) @@ -79,15 +89,16 @@ The implementations of dense op `op` (`:gemm`, `:syrk`, `:herk`, `:trsm`, `:potrf_batched`) available for element type `T` on `backend` according to [`capabilities`](@ref), in `:auto` preference order (`:vendor`, `:generic`, -`:ka`). `:ka` is always present. +`:ka`). `:ka` is always present. A vendor binding the backend keeps out of +`:auto` (oneAPI's LAPACK) is not listed; `impl = :vendor` still selects it. """ function dense_impls(op::Symbol, backend, ::Type{T}) where {T} haskey(DENSE_OPS, op) || throw(InvalidValueError("unknown dense op :$op; expected one of $(keys(DENSE_OPS))")) spec = DENSE_OPS[op] caps = capabilities(backend, T) impls = Symbol[] - spec.vendor === nothing || getfield(caps, spec.vendor) && push!(impls, :vendor) - spec.generic === nothing || getfield(caps, spec.generic) && push!(impls, :generic) + spec.vendor === nothing || getfield(caps, spec.vendor) && _auto_vendor_dense(op, backend, T) && push!(impls, :vendor) + spec.generic === nothing || getfield(caps, spec.generic) && _allow_generic_dense(op, backend, T) && push!(impls, :generic) push!(impls, :ka) return impls end @@ -108,8 +119,11 @@ function select_impl(op::Symbol, X::AbstractArray, impl::Symbol) gfield, vfield = _DENSE_OP_FIELDS[op] caps = capabilities(KernelAbstractions.get_backend(X), eltype(X)) vendor = vfield !== :none && getfield(caps, vfield)::Bool - generic = gfield !== :none && getfield(caps, gfield)::Bool - impl === :auto && return vendor ? :vendor : generic ? :generic : :ka + generic = gfield !== :none && getfield(caps, gfield)::Bool && + _allow_generic_dense(op, KernelAbstractions.get_backend(X), eltype(X)) + impl === :auto && + return vendor && _auto_vendor_dense(op, KernelAbstractions.get_backend(X), eltype(X)) ? :vendor : + generic ? :generic : :ka (impl === :vendor ? vendor : generic) || throw(NotSupportedError("impl = :$impl is not available for $op with $(eltype(X)) on $(KernelAbstractions.get_backend(X))")) return impl diff --git a/src/numeric/factorize.jl b/src/numeric/factorize.jl index 52b3e3c..ab43711 100644 --- a/src/numeric/factorize.jl +++ b/src/numeric/factorize.jl @@ -47,6 +47,28 @@ const STATS_WORKGROUP = 256 end end +@kernel function _cholesky_stats_serial_kernel!(stats_all, info_all, super_ptr, front_ncols, ns, schur, bm, + ::Val{NF}) where {NF} + j = @index(Global) + @inbounds if j <= bm.nact + k = _bm_member(bm, j) + stats = _mview(stats_all, k, bm.nbatch) + info = _iview(info_all, k, bm.nbatch) + best = typemax(Int64) + for s in 1:ns + fi = Int64(info[s]) + base = (s - 1) * NF + stats[base + 1] = s == schur ? Int64(0) : fi == 0 ? Int64(front_ncols[s]) : fi - 1 + for q in 2:(NF - 1) + stats[base + q] = 0 + end + stats[base + NF] = fi + fi > 0 && (best = min(best, Int64(super_ptr[s]) + fi - 1)) + end + info[ns + 1] = best == typemax(Int64) ? Int32(0) : Int32(best) + end +end + """ cholesky_stats!(numeric, symbolic) -> numeric @@ -58,9 +80,16 @@ factor column into `numeric.info[ns + 1]` (0 = none), for every active batch member. One workgroup per member, no atomics. Asynchronous. """ function cholesky_stats!(N::Numeric, S::Symbolic) - WG = STATS_WORKGROUP bm = batch_map(N) - kernel! = _cholesky_stats_kernel!(KernelAbstractions.get_backend(N.factor), WG) + backend = KernelAbstractions.get_backend(N.factor) + if _serial_cholesky_stats(backend, N.nbatch) + _cholesky_stats_serial_kernel!(backend, 1)(N.stats, N.info, S.super_ptr, S.front_ncols, + nsupernodes(S), S.schedule.schur, bm, + Val(FRONT_STATS_FIELDS); ndrange = bm.nact) + return N + end + WG = STATS_WORKGROUP + kernel! = _cholesky_stats_kernel!(backend, WG) kernel!(N.stats, N.info, S.super_ptr, S.front_ncols, nsupernodes(S), S.schedule.schur, bm, Val(WG), Val(FRONT_STATS_FIELDS), Val(_ilog2(WG)); ndrange = WG * bm.nact) @@ -100,7 +129,7 @@ end # batched path: `impl = :generic`, then every member takes the single-matrix path) function _front_impls(N::Numeric{T}, S::Symbolic, impl::Symbol) where {T} impl === :auto && !S.schedule.vendor_c && (impl = :ka) - batched = N.nbatch > 1 && impl !== :generic + batched = N.nbatch > 1 && impl !== :generic && _allow_batched_dense(KernelAbstractions.get_backend(N.factor)) bimpl(op) = batched ? (impl === :ka ? :ka : select_impl(op, N.factor, :auto)) : :none return (potrf = select_impl(:potrf, N.factor, impl), trsm = select_impl(:trsm, N.factor, impl), herk = select_impl(T <: Real ? :syrk : :herk, N.factor, impl), bpotrf = bimpl(:potrf_batched), @@ -112,8 +141,13 @@ end function _factor_front_c!(N::Numeric{T}, S::Symbolic, s::Int, p::NamedTuple, epsc::Real = 0.0) where {T} L, sc = S.layout, S.schedule if N.nbatch == 1 - _factor_panel_c!(N.factor, N.stack, N.work, N.info, s, L.panel_ptr[s], sc.rows[s], sc.width[s], L.cb_ptr[s], - p, epsc) + if _offset_potrf_enabled(KernelAbstractions.get_backend(N.factor)) + _factor_panel_c!(N.factor, N.stack, N.work, N.info, s, L.panel_ptr[s], sc.rows[s], sc.width[s], + L.cb_ptr[s], p, epsc, N.plan.vendor_ptrs[s]) + else + _factor_panel_c!(N.factor, N.stack, N.work, N.info, s, L.panel_ptr[s], sc.rows[s], sc.width[s], + L.cb_ptr[s], p, epsc) + end return nothing end plan = N.plan @@ -201,7 +235,7 @@ end function _factor_panel_c!(factor::AbstractVector{T}, stack::AbstractVector{T}, work::AbstractVector{T}, info::AbstractVector{Int32}, s::Int, p0::Int, f::Int, w::Int, c0::Int, - p::NamedTuple, epsc::Real = 0.0) where {T} + p::NamedTuple, epsc::Real = 0.0, vendor_plan = nothing) where {T} m = f - w P = reshape(view(factor, p0:(p0 + f * w - 1)), f, w) F11 = view(P, 1:w, 1:w) @@ -215,7 +249,14 @@ function _factor_panel_c!(factor::AbstractVector{T}, stack::AbstractVector{T}, w # elimination is caught after the fact by the sqrt(epsilon) factor-diagonal floor # inside _potrf_info_impl!, reported through info (the potrf route DETECTS where the # kernel routes perturb -- either way MadNLP gets its regularization signal, #127) - _potrf_info_impl!(p.potrf, 'L', F11, info, s, sqrt(real(T)(epsc))) + dmin = sqrt(real(T)(epsc)) + mps = vendor_plan !== nothing && p.potrf === :ka && + _offset_potrf_info!(KernelAbstractions.get_backend(factor), vendor_plan, info, s) + if mps + ka_chol_check_info!(info, s, F11, dmin) + else + _potrf_info_impl!(p.potrf, 'L', F11, info, s, dmin) + end m > 0 || return nothing F21 = view(P, (w + 1):f, 1:w) _trsm_impl!(p.trsm, 'R', 'L', 'C', 'N', one(T), F11, F21) @@ -324,7 +365,22 @@ function _numeric_info!(N::Numeric, S::Symbolic, read::Bool) plan = N.plan nb = N.nbatch if read - copyto!(plan.info_host, 1, N.info, nsupernodes(S) * nb + 1, nb) + ns = nsupernodes(S) + if !_is_ldlt_structure(S.structure) && S.structure != STRUCTURE_GENERAL && + _host_reduce_cholesky_info(KernelAbstractions.get_backend(N.factor), nb) + copyto!(plan.info_host, N.info) + for k in 1:nb + best = typemax(Int32) + for s in 1:ns + fi = plan.info_host[(s - 1) * nb + k] + fi > 0 && (best = min(best, Int32(S.partition.super_ptr[s]) + fi - 1)) + end + plan.info_host[ns * nb + k] = best == typemax(Int32) ? Int32(0) : best + end + copyto!(plan.info_host, 1, plan.info_host, ns * nb + 1, nb) + else + copyto!(plan.info_host, 1, N.info, ns * nb + 1, nb) + end else for j in 1:plan.nact[] plan.info_host[plan.members_host[j]] = 0 diff --git a/src/numeric/storage.jl b/src/numeric/storage.jl index 1a3d2c8..2d5a1ca 100644 --- a/src/numeric/storage.jl +++ b/src/numeric/storage.jl @@ -97,7 +97,8 @@ function NumericPlan(S, nb::Integer = 1) maxchild[t] = max(maxchild[t], nchild[sc.group_nodes[q]]) end end - return NumericPlan(first, last, maxchild, gfirst, glast, gmaxchild, gwidth, sfirst, slast, slocal, zeros(Int32, nb), + return NumericPlan(first, last, maxchild, gfirst, glast, gmaxchild, gwidth, sfirst, slast, slocal, + zeros(Int32, (ns + 1) * nb), Int32.(1:nb), Ref(Int(nb)), [1, nb + 1], Any[nothing for _ in 1:ns], schur_first, sc.schur > 0 ? nchild[sc.schur] : 0) end @@ -238,6 +239,14 @@ function allocate_numeric(S::Symbolic{INT}, ::Type{T}, backend::KernelAbstractio ufactor = KernelAbstractions.zeros(backend, T, lu ? nb * L.factor_len : 0) ustack = KernelAbstractions.zeros(backend, T, lu ? nb * L.stack_len : 0) plan = NumericPlan(S, nb) + if nb == 1 && !_is_ldlt_structure(S.structure) && S.structure != STRUCTURE_GENERAL + vstatus = _offset_potrf_status(backend) # vendor potrf status output, shared by the fronts (never read) + for s in 1:nsupernodes(S) + takes_c_path(S.schedule, s) || continue + plan.vendor_ptrs[s] = _offset_potrf_plan(backend, factor, L.panel_ptr[s], S.schedule.rows[s], + S.schedule.width[s], vstatus) + end + end nb > 1 && _vendor_batch_pointers!(plan.vendor_ptrs, factor, S, nb) return Numeric{T, typeof(factor), typeof(stats), typeof(info), typeof(pivot_kind)}(factor, d, stack, work, stats, info, piv, pivot_kind, totals, diff --git a/src/solve/sweeps.jl b/src/solve/sweeps.jl index 26716bc..60aa8a4 100644 --- a/src/solve/sweeps.jl +++ b/src/solve/sweeps.jl @@ -87,6 +87,7 @@ end function SolvePlan(S::Symbolic) sc = S.schedule + backend = KernelAbstractions.get_backend(S.rowval) kind, first, last, step = Int[], Int[], Int[], Int[] maxm = 0 for g in sc.groups @@ -95,7 +96,7 @@ function SolvePlan(S::Symbolic) elseif sc.schur > 0 && sc.group_nodes[g.first] == sc.schur k = SOLVE_SCHUR else - k = takes_c_path(sc, sc.group_nodes[g.first]) ? SOLVE_DENSE : SOLVE_FRONTS + k = _dense_solve_front(backend, sc, sc.group_nodes[g.first]) ? SOLVE_DENSE : SOLVE_FRONTS if k == SOLVE_DENSE for q in g.first:g.last s = sc.group_nodes[q] @@ -645,7 +646,7 @@ end # batched `b*` for a uniform batch, `:none` with `impl = :generic`) function _solve_impls(N::Numeric, S::Symbolic, impl::Symbol) impl === :auto && !S.schedule.vendor_c && (impl = :ka) - batched = N.nbatch > 1 && impl !== :generic + batched = N.nbatch > 1 && impl !== :generic && _allow_batched_dense(KernelAbstractions.get_backend(N.factor)) bimpl(op) = batched ? (impl === :ka ? :ka : select_impl(op, N.factor, :auto)) : :none return (trsm = select_impl(:trsm, N.factor, impl), gemm = select_impl(:gemm, N.factor, impl), btrsm = bimpl(:trsm_strided_batched), bgemm = bimpl(:gemm_strided_batched)) @@ -653,7 +654,7 @@ end function _launch_fwd!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, count::Int, nrhs::Int, sub::Val, det::Val, ldl::Val, factor::AbstractVector = N.factor) - WG = SOLVE_WORKGROUP + WG = _solve_workgroup(KernelAbstractions.get_backend(ws.Y)) solve_fwd_kernel!(KernelAbstractions.get_backend(ws.Y), WG)( ws.Y, ws.U, factor, N.piv, S.group_nodes, batch_map(N; first, nrhs), count, S.subtree_ptr, S.subtree_nodes, S.super_ptr, S.rowptr, S.rowval, S.front_ptr, S.front_nrows, S.front_ncols, S.child_ptr, S.child_list, @@ -663,7 +664,7 @@ end function _launch_bwd!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, count::Int, nrhs::Int, sub::Val, ldl::Val, cj::Val, factor::AbstractVector = N.factor) - WG = SOLVE_WORKGROUP + WG = _solve_workgroup(KernelAbstractions.get_backend(ws.Y)) solve_bwd_kernel!(KernelAbstractions.get_backend(ws.Y), WG)( ws.Y, ws.U, factor, N.piv, S.group_nodes, batch_map(N; first, nrhs), count, S.subtree_ptr, S.subtree_nodes, S.super_ptr, S.rowptr, S.rowval, S.front_ptr, S.front_nrows, S.front_ncols, S.child_ptr, sub, ldl, cj, Val(WG); @@ -724,7 +725,7 @@ end # LDLᵀ regime-C path: local pivot order of front s on the columns of Y (each with its member's order) function _piv_dense!(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, forward::Bool) - WG = SOLVE_WORKGROUP + WG = _solve_workgroup(KernelAbstractions.get_backend(ws.Y)) kernel! = _sv_piv_kernel!(KernelAbstractions.get_backend(ws.Y), WG) args = (ws.Y, ws.U, N.piv, s, S.super_ptr, S.rowptr, S.front_ptr, S.front_nrows, S.front_ncols, S.child_ptr, batch_map(N; nrhs)) @@ -901,13 +902,14 @@ function forward_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integ a, b = plan.first[k], plan.last[k] if plan.kind[k] == SOLVE_SCHUR # the Schur block stays as the children left it: Y₂ − L₂₁ Y₁ (the condensed right-hand side) - det && _sv_pull_kernel!(KernelAbstractions.get_backend(ws.Y), SOLVE_WORKGROUP)( + WG = _solve_workgroup(KernelAbstractions.get_backend(ws.Y)) + det && _sv_pull_kernel!(KernelAbstractions.get_backend(ws.Y), WG)( ws.Y, ws.U, S.group_nodes, a, b - a + 1, S.rowptr, S.super_ptr, S.front_ptr, S.front_nrows, - S.front_ncols, S.child_ptr, S.child_list, S.relind_ptr, S.relind, Val(SOLVE_WORKGROUP); - ndrange = SOLVE_WORKGROUP * (b - a + 1) * _ncols(N, nrhs)) + S.front_ncols, S.child_ptr, S.child_list, S.relind_ptr, S.relind, Val(WG); + ndrange = WG * (b - a + 1) * _ncols(N, nrhs)) elseif plan.kind[k] == SOLVE_DENSE if det - WG = SOLVE_WORKGROUP + WG = _solve_workgroup(KernelAbstractions.get_backend(ws.Y)) _sv_pull_kernel!(KernelAbstractions.get_backend(ws.Y), WG)( ws.Y, ws.U, S.group_nodes, a, b - a + 1, S.rowptr, S.super_ptr, S.front_ptr, S.front_nrows, S.front_ncols, S.child_ptr, S.child_list, S.relind_ptr, S.relind, Val(WG); diff --git a/src/solver.jl b/src/solver.jl index 703206a..2029383 100644 --- a/src/solver.jl +++ b/src/solver.jl @@ -457,7 +457,7 @@ end # the backend when the symbolic maps may be built on it (device.jl), else nothing _device_backend(solver::DirectSolver) = - solver.backend isa KernelAbstractions.GPU ? solver.backend : nothing + solver.backend isa KernelAbstractions.CPU ? nothing : solver.backend function _reorder!(solver::DirectSolver{T}) where {T} _check_analysis_supported(solver) @@ -551,9 +551,10 @@ function _symbolic!(solver::DirectSolver{T, INT}) where {T, INT} sp = schur_supernode_partition(factor_pattern(P, ord), ord.perm, length(solver.schur.indices), opts) schur = nsupernodes(sp) end - sc = build_schedule(sp, opts, T; reserve = subtree_local_reserve(solver.structure), - elsize = schedule_elsize(solver.structure, T), schur, - max_local = max_local_bytes(solver.backend)) + schedule_opts = _backend_schedule_options(solver.backend, solver.structure, opts, solver.nbatch) + sc = build_schedule(sp, schedule_opts, T; reserve = subtree_local_reserve(solver.structure), + elsize = schedule_elsize(solver.structure, T), schur, + max_local = max_local_bytes(solver.backend)) layout = build_layout(sp, sc; ldlt = _is_ldlt_structure(solver.structure)) Sh = Symbolic(sp, sc, layout, solver.host_rowptr, solver.analysis_colval, A.nrows, solver.structure; view = _stored_view(solver), index = A.index, device = _device_backend(solver)) diff --git a/src/symbolic/device.jl b/src/symbolic/device.jl index f3a36b3..6ac3ec5 100644 --- a/src/symbolic/device.jl +++ b/src/symbolic/device.jl @@ -24,7 +24,7 @@ sizes that fit the packed 64-bit sort keys. """ function device_maps_supported(backend, structure, view, n::Integer, nnz::Integer, factor_len::Integer) - backend isa KernelAbstractions.GPU || return false + backend isa KernelAbstractions.CPU && return false s = _structure(structure) (s == STRUCTURE_SPD || s == STRUCTURE_HPD || s == STRUCTURE_SYMMETRIC || s == STRUCTURE_HERMITIAN) || return false diff --git a/src/symbolic/schedule.jl b/src/symbolic/schedule.jl index 558ce9e..88b47a0 100644 --- a/src/symbolic/schedule.jl +++ b/src/symbolic/schedule.jl @@ -116,12 +116,28 @@ packed_length(m::Integer) = Int(m) * (Int(m) + 1) ÷ 2 # inverse of the packed lower-triangle enumeration: 0-based q in an r x r lower # triangle (col-major, n = r(r+1)/2 entries) -> 1-based (a, b), a >= b. Enumerating # only the triangle halves the work-item count of the guarded full-square loops. -# Integer isqrt (exact, and no Float64 in kernels: Metal has no fp64); the two -# boundary fixups are kept as a belt against any isqrt implementation quirk. +# Integer square root without Base.isqrt, whose Julia 1.13 implementation uses +# Float64 internally and therefore cannot compile for Metal kernels. +@inline function _kernel_isqrt(x::T) where {T <: Integer} + x <= one(T) && return x + lo, hi = zero(T), min(x, T(1) << ((8 * sizeof(T) + 1) >> 1)) + while lo + one(T) < hi + mid = lo + ((hi - lo) >> 1) + if mid <= x ÷ mid + lo = mid + else + hi = mid + end + end + return lo +end + +# The two boundary fixups are kept as a belt against integer overflow or a +# compiler quirk in the inverse triangular-number calculation. @inline function _tri_decode(q, r) n = r * (r + 1) ÷ 2 p = n - 1 - q - t = (isqrt(8 * p + 1) - 1) ÷ 2 + t = (_kernel_isqrt(8 * p + 1) - 1) ÷ 2 t * (t + 1) ÷ 2 > p && (t -= 1) (t + 1) * (t + 2) ÷ 2 <= p && (t += 1) b = r - t diff --git a/test/backends.jl b/test/backends.jl index cca6d4d..dfa1590 100644 --- a/test/backends.jl +++ b/test/backends.jl @@ -62,6 +62,21 @@ if AMDGPU_LOADED end const AMDGPU_FUNCTIONAL = AMDGPU_LOADED && AMDGPU.functional() +const METAL_LOADED = TEST_GPU && is_package_installed("Metal") +if METAL_LOADED + using Metal + Metal.allowscalar(false) +end +const METAL_FUNCTIONAL = METAL_LOADED && Metal.functional() + +const ONEAPI_LOADED = TEST_GPU && is_package_installed("oneAPI") +if ONEAPI_LOADED + using oneAPI + using oneAPI.oneMKL: oneSparseMatrixCSR, oneSparseMatrixCSC + oneAPI.GPUArrays.allowscalar(false) +end +const ONEAPI_FUNCTIONAL = ONEAPI_LOADED && oneAPI.functional() + """ BACKENDS @@ -72,14 +87,17 @@ const BACKENDS = Any[] TEST_CPU && push!(BACKENDS, CPU()) CUDA_FUNCTIONAL && push!(BACKENDS, CUDABackend()) AMDGPU_FUNCTIONAL && push!(BACKENDS, ROCBackend()) +METAL_FUNCTIONAL && push!(BACKENDS, MetalBackend()) +ONEAPI_FUNCTIONAL && push!(BACKENDS, oneAPIBackend()) isempty(BACKENDS) && error("no backend to test: SDS_TEST_CPU=0 and no functional GPU backend (SDS_TEST_GPU=$(get(ENV, "SDS_TEST_GPU", "1")), ", - "CUDA loaded: $CUDA_LOADED, AMDGPU loaded: $AMDGPU_LOADED)") + "CUDA loaded: $CUDA_LOADED, AMDGPU loaded: $AMDGPU_LOADED, Metal loaded: $METAL_LOADED, ", + "oneAPI loaded: $ONEAPI_LOADED)") """ backend_name(backend) -> String -Short name for testset titles: `"CPU"`, `"CUDA"`, `"ROCm"`. +Short name for testset titles: `"CPU"`, `"CUDA"`, `"ROCm"`, `"Metal"`, `"oneAPI"`. """ backend_name(::CPU) = "CPU" @@ -212,6 +230,33 @@ if AMDGPU_LOADED (n, n)) end +if METAL_LOADED + backend_name(::MetalBackend) = "Metal" + to_device(::MetalBackend, x::Array) = MtlArray(x) + to_device(backend::MetalBackend, A::SparseMatrixCSC{T, INT}) where {T, INT} = to_device(backend, A, INT) + to_device(::MetalBackend, A::SparseMatrixCSC{T}, ::Type{INT}) where {T, INT} = + SparseDirectSolver.to_backend(SparseMatrixCSC{T, INT}(A), MetalBackend()) + to_host(A::CSR{T, INT, <:MtlVector{INT}, <:MtlArray{T}}) where {T, INT} = SparseMatrixCSC(A) + index_eltype(::CSR{<:Any, INT, <:MtlVector{INT}}) where {INT} = INT + api_csr(::MetalBackend, rowptr::Vector{INT}, colval::Vector{INT}, nzval, n) where {INT} = + CSR(MtlVector{INT}(rowptr), MtlVector{INT}(colval), MtlArray(vec(nzval)), n, n) +end + +if ONEAPI_LOADED + backend_name(::oneAPIBackend) = "oneAPI" + to_device(::oneAPIBackend, x::Array) = oneArray(x) + to_device(backend::oneAPIBackend, A::SparseMatrixCSC{T, INT}) where {T, INT} = to_device(backend, A, INT) + function to_device(::oneAPIBackend, A::SparseMatrixCSC{T}, ::Type{INT}) where {T, INT} + rowptr, colval, nzval = host_csr_arrays(A, INT) + return oneSparseMatrixCSR{T, INT}(oneVector{INT}(rowptr), oneVector{INT}(colval), oneVector{T}(nzval), size(A)) + end + to_host(A::oneSparseMatrixCSR) = host_csc(A.rowPtr, A.colVal, A.nzVal, size(A)...) + index_eltype(A::oneSparseMatrixCSR) = eltype(A.rowPtr) + api_csr(::oneAPIBackend, rowptr::Vector{INT}, colval::Vector{INT}, nzval, n) where {INT} = + oneSparseMatrixCSR{eltype(nzval), INT}(oneVector{INT}(rowptr), oneVector{INT}(colval), oneVector(vec(nzval)), + (n, n)) +end + """ vendor_sparse_types(backend) -> Union{Nothing, Tuple} @@ -223,6 +268,7 @@ backend extensions loop over these. vendor_sparse_types(backend) = nothing CUDA_LOADED && (vendor_sparse_types(::CUDABackend) = (CuSparseMatrixCSR, CuSparseMatrixCSC, CuVector)) AMDGPU_LOADED && (vendor_sparse_types(::ROCBackend) = (ROCSparseMatrixCSR, ROCSparseMatrixCSC, ROCVector)) +ONEAPI_LOADED && (vendor_sparse_types(::oneAPIBackend) = (oneSparseMatrixCSR, oneSparseMatrixCSC, oneVector)) """ print_backends() @@ -234,8 +280,12 @@ function print_backends() gpus = String[] CUDA_FUNCTIONAL && push!(gpus, "CUDA ($(CUDA.name(CUDA.device())))") AMDGPU_FUNCTIONAL && push!(gpus, "ROCm ($(AMDGPU.device()))") + METAL_FUNCTIONAL && push!(gpus, "Metal ($(Metal.device().name))") + ONEAPI_FUNCTIONAL && push!(gpus, "oneAPI ($(oneAPI.oneL0.properties(oneAPI.device()).name))") CUDA_LOADED && !CUDA_FUNCTIONAL && push!(gpus, "CUDA installed but not functional: skipped") AMDGPU_LOADED && !AMDGPU_FUNCTIONAL && push!(gpus, "AMDGPU installed but not functional: skipped") + METAL_LOADED && !METAL_FUNCTIONAL && push!(gpus, "Metal installed but not functional: skipped") + ONEAPI_LOADED && !ONEAPI_FUNCTIONAL && push!(gpus, "oneAPI installed but not functional: skipped") TEST_GPU || push!(gpus, "GPU backends disabled by SDS_TEST_GPU=0") TEST_CPU || push!(gpus, "CPU backend disabled by SDS_TEST_CPU=0") println("Backends under test: ", join(backend_name.(BACKENDS), ", "), diff --git a/test/test_api.jl b/test/test_api.jl index 794d226..d7bfe47 100644 --- a/test/test_api.jl +++ b/test/test_api.jl @@ -186,7 +186,8 @@ end @test thrown(() -> execute!("solve", solver, x, to_device(backend, rand(T, 40, 2)))) isa DimensionMismatch # an n × nrhs × nbatch array is a uniform-batch layout (T17): 4 columns for X's 1 @test thrown(() -> execute!("solve", solver, x, to_device(backend, rand(T, 40, 2, 2)))) isa DimensionMismatch - other = T <: Complex ? ComplexF64 === T ? ComplexF32 : ComplexF64 : Float64 === T ? Float32 : Float64 + other = T === Float32 && backend_name(backend) == "Metal" ? Float16 : + T <: Complex ? ComplexF64 === T ? ComplexF32 : ComplexF64 : Float64 === T ? Float32 : Float64 bo = to_device(backend, rand(other, 40)) @test thrown(() -> execute!("solve", solver, bo, bo)) isa InvalidValueError backend isa CPU || @@ -482,7 +483,7 @@ end @test getparam(Fb, "info") > 0 @test thrown(() -> cholesky!(Fb, api_matrix(backend, tril(B), INT); check = true)) isa FactorizationError # the Symmetric/Hermitian wrappers of the vendor matrix types (GPU extensions) - if !(backend isa CPU) + if vendor_sparse_types(backend) !== nothing M = api_matrix(backend, tril(A), INT) Fh = cholesky(Hermitian(M, :L)) @test relres(A, to_host(Fh \ bd), b) <= tol(T) diff --git a/test/test_dense.jl b/test/test_dense.jl index d205a7b..5a38b51 100644 --- a/test/test_dense.jl +++ b/test/test_dense.jl @@ -33,6 +33,25 @@ dense_part(C, uplo::Char) = uplo == 'L' ? tril(C) : triu(C) end @test caps.vendor_herk == (T <: Complex) end + if backend_name(backend) == "oneAPI" # the oneMKL column of PLAN §2.6 (T31) + for f in (:vendor_gemm, :vendor_syrk, :vendor_trsm, :vendor_potrf, :vendor_getrf, + :vendor_gemm_strided_batched, :vendor_trsm_batched, :reshape_view_mul) + @test getfield(caps, f) + end + @test caps.vendor_herk == (T <: Complex) + # one status per batch in oneMKL: no per-member info, the KA kernels serve uniform batches; + # no sytrf on the GPU (oneMKL aborts with unsupported_device) + @test !caps.vendor_potrf_batched && !caps.vendor_getrf_batched && !caps.vendor_sytrf + # LAPACK waits for the queue and allocates its scratchpad: `:auto` takes the KA kernels, + # `impl = :vendor` still reaches oneMKL; the generic LinearAlgebra path is not used + @test SDS.dense_impls(:potrf, backend, T) == [:ka] + @test SDS.dense_impls(:getrf, backend, T) == [:ka] + @test SDS.dense_impls(:gemm, backend, T) == [:vendor, :ka] + @test SDS.dense_impls(:trsm, backend, T) == [:vendor, :ka] + X = to_device(backend, zeros(T, 2, 2)) + @test SDS.select_impl(:potrf, X, :vendor) === :vendor + @test_throws SDS.NotSupportedError SDS.select_impl(:potrf, X, :generic) + end for op in keys(SDS.DENSE_OPS) impls = SDS.dense_impls(op, backend, T) @test last(impls) === :ka diff --git a/test/test_fgmres.jl b/test/test_fgmres.jl index 6b68a37..f274721 100644 --- a/test/test_fgmres.jl +++ b/test/test_fgmres.jl @@ -13,6 +13,7 @@ fgmres_stall_matrix() = kkt_matrix(Float64, 150, 5, 0.0) # extension itself is INT-agnostic (it sees only the solve callback) (test-suite diet). RUN_SHARED && @testset "FGMRES-IR where plain IR stalls ($(backend_name(backend)), $INT)" for backend in BACKENDS, INT in (Int32,) + backend_name(backend) == "Metal" && continue Random.seed!(666) K = fgmres_stall_matrix() b = K * rand(size(K, 1)) @@ -103,6 +104,7 @@ end end RUN_SHARED && @testset "FGMRES-IR: interrupt and missing Krylov ($(backend_name(backend)))" for backend in BACKENDS + backend_name(backend) == "Metal" && continue Random.seed!(666) K = fgmres_stall_matrix() b = K * rand(size(K, 1)) diff --git a/test/test_helpers.jl b/test/test_helpers.jl index 8f7a05a..0253966 100644 --- a/test/test_helpers.jl +++ b/test/test_helpers.jl @@ -8,6 +8,10 @@ @test ("CUDA" in names) == cuda_expected rocm_expected = TEST_GPU && AMDGPU_LOADED && AMDGPU.functional() @test ("ROCm" in names) == rocm_expected + metal_expected = TEST_GPU && METAL_LOADED && Metal.functional() + @test ("Metal" in names) == metal_expected + oneapi_expected = TEST_GPU && ONEAPI_LOADED && oneAPI.functional() + @test ("oneAPI" in names) == oneapi_expected @info "Active backends: $(join(names, ", "))" if CUDA_LOADED # loading CUDA triggers the package extension @test Base.get_extension(SparseDirectSolver, :SparseDirectSolverCUDAExt) !== nothing @@ -15,17 +19,24 @@ if AMDGPU_LOADED # likewise AMDGPU (T23) @test Base.get_extension(SparseDirectSolver, :SparseDirectSolverAMDGPUExt) !== nothing end + if METAL_LOADED + @test Base.get_extension(SparseDirectSolver, :SparseDirectSolverMetalExt) !== nothing + end + if ONEAPI_LOADED # likewise oneAPI (T31) + @test Base.get_extension(SparseDirectSolver, :SparseDirectSolverOneAPIExt) !== nothing + end for backend in BACKENDS x = rand(Float32, 7) dx = to_device(backend, x) @test KernelAbstractions.get_backend(dx) == backend @test to_host(dx) == x - A = sprand(ComplexF64, 9, 9, 0.3) + I + TC = backend_name(backend) == "Metal" ? Float32 : ComplexF64 # Metal's to_backend takes Float32 only + A = sprand(TC, 9, 9, 0.3) + I @test to_host(to_device(backend, A)) == A for INT in INTTYPES dA = to_device(backend, A, INT) @test to_host(dA) == A - @test eltype(dA) == ComplexF64 + @test eltype(dA) == TC @test index_eltype(dA) == INT # issue #30: vendor constructors silently use Int32 end end diff --git a/test/test_matrix.jl b/test/test_matrix.jl index 5286260..cb30d63 100644 --- a/test/test_matrix.jl +++ b/test/test_matrix.jl @@ -153,7 +153,7 @@ end @test get_backend(d) == backend @test thrown(() -> update!(d, to_device(backend, rand(T, n * nb + 1)))) isa InvalidValueError @test thrown(() -> update!(d, to_device(backend, rand(T, n, nb + 1)))) isa InvalidValueError - @test thrown(() -> update!(d, to_device(backend, rand(T === Float32 ? Float64 : Float32, n * nb)))) isa + @test thrown(() -> update!(d, to_device(backend, rand(T === Float32 ? Float16 : Float32, n * nb)))) isa InvalidValueError @test d.data === x # failed updates leave the descriptor alone diff --git a/test/test_numeric_cholesky_a.jl b/test/test_numeric_cholesky_a.jl index 2413243..c581cd5 100644 --- a/test/test_numeric_cholesky_a.jl +++ b/test/test_numeric_cholesky_a.jl @@ -8,7 +8,7 @@ const NUMERIC_A_OPTS = Options(subtree_parallelism = 0) # small budgets and regime-C thresholds: all three regimes on a 2-D Laplacian const NUMERIC_ABC_OPTS = Options(subtree_budgets = [8192, 16384], regime_c_width = 16, regime_c_rows = 128, - subtree_parallelism = 0) + subtree_parallelism = 0, subtree_max_fronts = 32) numeric_a_matrices(::Type{T}) where {T} = (("laplacian2d(40,40) A+B+C", laplacian2d(T, 40, 40), NUMERIC_ABC_OPTS), @@ -112,9 +112,10 @@ end Options(subtree_budgets = [8192, 20000, 49152], regime_c_width = 32, subtree_parallelism = 0), Options(subtree_budgets = [8192, 20000, 1 << 20], regime_c_width = 32, subtree_parallelism = 0), Options(subtree_budgets = [32768], factorization_alg = "algo2", subtree_parallelism = 0)) - # the 64 KiB class (issue #60) runs where the backend has the local memory (CPU, ROCm); elsewhere (CUDA) - # the allocation of an analysis that uses it is rejected - if SDS.subtree_local_bytes(maximum(opts.subtree_budgets)) > SDS.max_local_bytes(backend) + # the 64 KiB class (issue #60) runs where the backend has the local memory (CPU, ROCm); elsewhere (CUDA, + # Metal) the analysis rejects it, and an explicit budget above `max_local_bytes` (the 1 MiB entry) is + # rejected on every GPU backend (`resolve_subtree_budgets`) + if maximum(opts.subtree_budgets) > SDS.max_local_bytes(backend) @test thrown(() -> numeric_setup(backend, A; opts)) isa InvalidValueError continue end diff --git a/test/test_numeric_ldlt.jl b/test/test_numeric_ldlt.jl index 2fa0cc9..b4cfbd7 100644 --- a/test/test_numeric_ldlt.jl +++ b/test/test_numeric_ldlt.jl @@ -94,7 +94,7 @@ end @test d_error(Nh, Nr) <= growth_tol(T, Nr) && panel_error(Nh, Nr) <= growth_tol(T, Nr) end S, Nr, Sd, Nd, nz = ldlt_setup(backend, A, Int64; opts) - @test !isempty(Nd.plan.sub_first) && SDS.pivot_stats(Nr).n2x2 > 0 + @test (backend_name(backend) == "Metal" || !isempty(Nd.plan.sub_first)) && SDS.pivot_stats(Nr).n2x2 > 0 SDS.factorize!(Nd, Sd, nz; opts) Nh = SDS.host_numeric(Nd) @test Nh.piv == Nr.piv && Nh.pivot_kind == Nr.pivot_kind @@ -260,9 +260,10 @@ end sc, sp = S.schedule, S.partition blocked_b = [s for s in 1:SDS.nsupernodes(S) if sc.regime[s] == SDS.REGIME_B && SDS.ldlt_blocked_path(sc, s) && !SDS.takes_c_path(sc, s)] - @test !isempty(blocked_b) + @test backend_name(backend) == "Metal" || !isempty(blocked_b) # 2×2 pivots inside such a front - @test any(s -> any(k -> Nr.pivot_kind[k] == SDS.PIVOT_KIND_2X2_FIRST, SDS.sncols(sp, s)), blocked_b) + @test backend_name(backend) == "Metal" || + any(s -> any(k -> Nr.pivot_kind[k] == SDS.PIVOT_KIND_2X2_FIRST, SDS.sncols(sp, s)), blocked_b) @test SDS.factorize!(Nd, Sd, nz; opts) == 0 Nh = SDS.host_numeric(Nd) @test Nh.piv == Nr.piv diff --git a/test/test_numeric_lu.jl b/test/test_numeric_lu.jl index f719c52..bf5919c 100644 --- a/test/test_numeric_lu.jl +++ b/test/test_numeric_lu.jl @@ -100,7 +100,7 @@ end @test d_error(Nh, Nr) <= gt && panel_error(Nh, Nr) <= gt && upanel_error(Nh, Nr) <= gt end S, Nr, Sd, Nd, nz = lu_setup(backend, A, Int64; opts) - @test !isempty(Nd.plan.sub_first) && Nr.piv != 1:400 + @test (backend_name(backend) == "Metal" || !isempty(Nd.plan.sub_first)) && Nr.piv != 1:400 SDS.factorize!(Nd, Sd, nz; opts) Nh = SDS.host_numeric(Nd) @test Nh.piv == Nr.piv && Nh.pivot_kind == Nr.pivot_kind diff --git a/test/test_refinement.jl b/test/test_refinement.jl index a93a8dd..f77ee33 100644 --- a/test/test_refinement.jl +++ b/test/test_refinement.jl @@ -16,6 +16,7 @@ function Base.CoreLogging.handle_message(L::InterruptAtStepLogger, level, messag end RUN_SHARED && @testset "refinement on a badly scaled SPD matrix ($(backend_name(backend)))" for backend in BACKENDS + backend_name(backend) == "Metal" && continue Random.seed!(666) T = Float64 A = badly_scaled_spd(T, 300, 0.02) # rows scaled by up to 10^(±8) @@ -43,6 +44,7 @@ end RUN_SHARED && @testset "refinement with static pivot perturbation ($(backend_name(backend)), $INT)" for backend in BACKENDS, INT in INTTYPES + backend_name(backend) == "Metal" && continue Random.seed!(666) # KKT matrix without 2×2 pairs: the zero (2,2) block is perturbed (pivot_epsilon), and refinement # removes the perturbation error (the case of the MadNLP K2 systems, issue #71) @@ -271,8 +273,9 @@ end RUN_SHARED && @testset "logging" begin Random.seed!(666) backend = first(BACKENDS) - A = random_spd(Float64, 80, 0.05) - b = rand(80) + T = backend_name(backend) == "Metal" ? Float32 : Float64 + A = random_spd(T, 80, 0.05) + b = rand(T, 80) old = SDS.set_log_level!("info") try solver = DirectSolver(api_matrix(backend, tril(A), Int32), "SPD", 'L') diff --git a/test/test_schur.jl b/test/test_schur.jl index 99f505e..918a7e3 100644 --- a/test/test_schur.jl +++ b/test/test_schur.jl @@ -44,11 +44,13 @@ schur_ldlt(structure) = structure in ("S", "H") sc = solver.host_symbolic.schedule @test sc.schur == SDS.nsupernodes(solver.host_symbolic) && sc.regime[sc.schur] == SDS.REGIME_C @test solver.host_symbolic.partition.perm[(n - ns + 1):n] == s - rname == "A" && @test count(==(SDS.REGIME_A), sc.regime) > length(sc.regime) ÷ 2 + rname == "A" && (backend_name(backend) == "Metal" && structure != "SPD" || + @test count(==(SDS.REGIME_A), sc.regime) > length(sc.regime) ÷ 2) # Laplacian: a regime-A subtree root hands its contribution block to the Schur root rname == "A" && structure in ("SPD", "HPD") && @test any(t -> solver.host_symbolic.partition.snparent[t] == sc.schur, sc.subtree_root) - rname == "B/C" && @test any(==(SDS.REGIME_B), sc.regime) && count(==(SDS.REGIME_C), sc.regime) > 1 + rname == "B/C" && (backend_name(backend) == "Metal" && structure != "SPD" || + @test any(==(SDS.REGIME_B), sc.regime) && count(==(SDS.REGIME_C), sc.regime) > 1) Sref = schur_reference(A, flags) P = schur_pattern_reference(A, flags) sym = structure != "G" @@ -183,7 +185,7 @@ end RUN_SHARED && @testset "Schur complement: orderings, edge cases, errors ($(backend_name(backend)))" for backend in BACKENDS Random.seed!(666) - T = Float64 + T = backend_name(backend) == "Metal" ? Float32 : Float64 # two disconnected blocks, Schur indices in the second one only: the first block is a separate tree A = blockdiag(laplacian2d(T, 6, 6), laplacian2d(T, 7, 5)) n = size(A, 1) @@ -238,7 +240,8 @@ RUN_SHARED && @testset "Schur complement: orderings, edge cases, errors ($(backe execute!("analysis", solver, nothing, nothing) @test getparam(solver, "schur_shape") == (4, 4, 10) @test thrown(() -> getparam(solver, "schur_matrix")) isa FactorizationError - for bad in (to_device(backend, zeros(T, 3, 3)), to_device(backend, zeros(Float32, 4, 4)), sparse(1.0I, 4, 4), + wrongT = T === Float32 ? Float16 : Float32 + for bad in (to_device(backend, zeros(T, 3, 3)), to_device(backend, zeros(wrongT, 4, 4)), sparse(1.0I, 4, 4), (to_device(backend, zeros(T, 4, 4)), 'L'), "S", CSR(to_device(backend, ones(Int32, 5)), to_device(backend, ones(Int32, 9)), to_device(backend, zeros(T, 9)), 4, 4)) diff --git a/test/test_symbolic_schedule.jl b/test/test_symbolic_schedule.jl index dd5a348..94fbb01 100644 --- a/test/test_symbolic_schedule.jl +++ b/test/test_symbolic_schedule.jl @@ -211,12 +211,14 @@ end @test SDS.SUBTREE_LOCAL_SIZES[end] == 65536 @test SDS.subtree_local_bytes(1 << 20) == 65536 @test SDS.subtree_local_bytes(65535) == 49152 - # the caps: no limit on the CPU backend, at least the CUDA-safe 48 KiB on every GPU backend + # the caps: no limit on CPU, 32 KiB on Metal, at least 48 KiB on the other GPU backends @test SDS.max_local_bytes(CPU()) >= 65536 for backend in BACKENDS - @test SDS.max_local_bytes(backend) >= 49152 + @test SDS.max_local_bytes(backend) >= (backend_name(backend) == "Metal" ? 32768 : 49152) CUDA_LOADED && backend isa CUDABackend && @test SDS.max_local_bytes(backend) == 49152 AMDGPU_LOADED && backend isa ROCBackend && @test SDS.max_local_bytes(backend) >= 65536 + METAL_LOADED && backend isa MetalBackend && @test SDS.max_local_bytes(backend) == 32768 + ONEAPI_LOADED && backend isa oneAPIBackend && @test SDS.max_local_bytes(backend) >= 65536 end # the defaults are clamped to the cap, explicit budgets above it are rejected @test SDS.resolve_subtree_budgets(SDS.DEFAULT_SUBTREE_BUDGETS, typemax(Int)) == [16384, 32768, 49152] @@ -245,7 +247,8 @@ end @test SDS.allocate_numeric(S64, Float64, CPU()) isa SDS.Numeric for backend in BACKENDS SDS.max_local_bytes(backend) < 65536 || continue - @test thrown(() -> SDS.allocate_numeric(SDS.adapt(backend, S64, Int32), Float64, backend)) isa InvalidValueError + T = backend_name(backend) == "Metal" ? Float32 : Float64 + @test thrown(() -> SDS.allocate_numeric(SDS.adapt(backend, S64, Int32), T, backend)) isa InvalidValueError end end diff --git a/test/utils.jl b/test/utils.jl index baf94e8..98babee 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -307,7 +307,10 @@ storage `Nd` and the values `nz` on `backend` (T09, T10). function numeric_setup(backend, A::SparseMatrixCSC{T}, ::Type{INT} = Int32; view = 'L', index = 'O', opts = Options(subtree_budgets = Int[])) where {T, INT} C = SparseDirectSolver.CSR(triangle_view(A, view); index) - S = SparseDirectSolver.symbolic_analysis(C, spd_structure(T), view; opts) + structure = SparseDirectSolver._structure(spd_structure(T)) + schedule_opts = SparseDirectSolver._backend_schedule_options(backend, structure, opts) + S = SparseDirectSolver.symbolic_analysis(C, structure, view; opts = schedule_opts, + max_local = SparseDirectSolver.max_local_bytes(backend)) Nr = SparseDirectSolver.allocate_numeric(S, T) info_ref = SparseDirectSolver.ref_factorize!(Nr, S, C.nzval) Sd = SparseDirectSolver.adapt(backend, S, INT) @@ -427,9 +430,10 @@ end # --- LDLᵀ/LDLᴴ reference helpers (T14; shared with the 2×2 pivot pair tests of #64) --- # analysis + numeric storage + LDLᵀ/LDLᴴ factorization of A given through `view` -function reference_ldlt(A::SparseMatrixCSC{T}; view = 'L', opts = Options(), structure = sym_structure(T)) where {T} +function reference_ldlt(A::SparseMatrixCSC{T}; view = 'L', opts = Options(), structure = sym_structure(T), + max_local = typemax(Int)) where {T} C = SparseDirectSolver.CSR(triangle_view(A, view)) - S = SparseDirectSolver.symbolic_analysis(C, structure, view; opts) + S = SparseDirectSolver.symbolic_analysis(C, structure, view; opts, max_local) N = SparseDirectSolver.allocate_numeric(S, T, CPU()) info = SparseDirectSolver.ref_ldlt!(N, S, C.nzval; opts) return S, N, info, C @@ -495,7 +499,10 @@ values `nz` on `backend` (T15). """ function ldlt_setup(backend, A::SparseMatrixCSC{T}, ::Type{INT} = Int32; view = 'L', opts = Options(), structure = sym_structure(T)) where {T, INT} - S, Nr, info, C = reference_ldlt(A; view, opts, structure) + st = SparseDirectSolver._structure(structure) + schedule_opts = SparseDirectSolver._backend_schedule_options(backend, st, opts) + S, Nr, info, C = reference_ldlt(A; view, opts = schedule_opts, structure = st, + max_local = SparseDirectSolver.max_local_bytes(backend)) info == 0 || error("ldlt_setup: the reference factorization failed") Sd = SparseDirectSolver.adapt(backend, S, INT) Nd = SparseDirectSolver.allocate_numeric(Sd, T, backend) @@ -666,7 +673,9 @@ with `INT` maps, device storage `Nd` and the values `nz` on `backend` (T19). """ function lu_setup(backend, A::SparseMatrixCSC{T}, ::Type{INT} = Int32; opts = Options()) where {T, INT} C = SparseDirectSolver.CSR(A) - S = SparseDirectSolver.symbolic_analysis(C, "G", 'F'; opts) + schedule_opts = SparseDirectSolver._backend_schedule_options(backend, SparseDirectSolver.STRUCTURE_GENERAL, opts) + S = SparseDirectSolver.symbolic_analysis(C, "G", 'F'; opts = schedule_opts, + max_local = SparseDirectSolver.max_local_bytes(backend)) Nr = SparseDirectSolver.allocate_numeric(S, T) SparseDirectSolver.ref_lu!(Nr, S, C.nzval; opts) == 0 || error("lu_setup: the reference factorization failed") Sd = SparseDirectSolver.adapt(backend, S, INT)