diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index b3f72a3c9..b61fbc9b4 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -111,5 +111,17 @@ steps: julia --project=test/metalenv -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' julia --project=test/metalenv test/run_mpi.jl 2 2 test/mpi_metal.jl + - label: Julia 1.11 (Finch) + timeout_in_minutes: 90 + if: build.message !~ /\[skip tests\]/ + plugins: + - JuliaCI/julia#v1: + version: "1.11" + - JuliaCI/julia-test#v1: ~ + - JuliaCI/julia-coverage#v1: + codecov: true + env: + CI_TEST_FINCH: "1" + # env: # SECRET_CODECOV_TOKEN: "" diff --git a/Project.toml b/Project.toml index 21eb4a00d..3f64360ae 100644 --- a/Project.toml +++ b/Project.toml @@ -23,7 +23,6 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Requires = "ae029012-a4dd-5104-9daa-d747884805df" ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63" Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b" -SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" TaskLocalValues = "ed4db957-447d-4319-bfb6-7fa9ae7ecf34" @@ -33,40 +32,62 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" [weakdeps] AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" +AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +Finch = "9177782c-1635-4eb9-9bfb-d9dfa25e6bce" +IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895" GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinuxPerf = "b4c46c6c-4fb0-484d-a11a-41bc3392d094" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" +Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +PureKLU = "0c0d3e7f-3a8b-4f7e-b6f1-9a4d2e7c1f01" +PureUMFPACK = "b7e1f0a2-3c4d-4e5f-9a0b-1c2d3e4f5a6b" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] AbstractFFTsExt = "AbstractFFTs" +AlgebraicMultigridExt = ["AlgebraicMultigrid", "SparseArrays"] CUDAExt = "CUDA" +CUDASparseArraysExt = ["CUDA", "SparseArrays"] DistributionsExt = "Distributions" +FinchExt = "Finch" GraphVizExt = "GraphViz" GraphVizSimpleExt = "Colors" +IncompleteLUExt = ["IncompleteLU", "SparseArrays"] IntelExt = "oneAPI" +IntelSparseArraysExt = ["oneAPI", "SparseArrays"] JSON3Ext = "JSON3" +KrylovExt = "Krylov" LinuxPerfExt = "LinuxPerf" -MPIExt = "MPI" +MPIExt = ["MPI", "SparseArrays"] MetalExt = "Metal" +MetalSparseArraysExt = ["Metal", "SparseArrays"] +MetisExt = ["Metis", "SparseArrays"] OpenCLExt = "OpenCL" +OpenCLSparseArraysExt = ["OpenCL", "SparseArrays"] PlotsExt = ["DataFrames", "Plots"] +PureKLUExt = ["PureKLU", "SparseArrays"] +PureUMFPACKExt = ["PureUMFPACK", "SparseArrays"] PythonExt = "PythonCall" ROCExt = "AMDGPU" +ROCSparseArraysExt = ["AMDGPU", "SparseArrays"] +SparseArraysExt = "SparseArrays" [compat] AMDGPU = "1, 2" AbstractFFTs = "1.5.0" Adapt = "4" +AlgebraicMultigrid = "1" CUDA = "3, 4, 5" Colors = "0.12, 0.13" DataFrames = "1" @@ -74,22 +95,28 @@ DataStructures = "0.18, 0.19" DistributedNext = "1.0.0" Distributions = "0.25" FillArrays = "1.13.0" +Finch = "1.2.9" GPUArraysCore = "0.2.0" GraphViz = "0.2" Graphs = "1" +IncompleteLU = "0.2" JSON3 = "1" KernelAbstractions = "0.9" +Krylov = "0.10" LinuxPerf = "0.4.2" MPI = "0.20.22" MacroTools = "0.5" MemPool = "0.4.17" Metal = "1.1" +Metis = "1" NextLA = "0.2.2" OnlineStats = "1" OpenCL = "0.10" Plots = "1" PrecompileTools = "1.2" Preferences = "1.4.3" +PureKLU = "1" +PureUMFPACK = "0.1" PythonCall = "0.9" Requires = "1" ScopedValues = "1.1" diff --git a/docs/make.jl b/docs/make.jl index 611f7301d..d7f6f0168 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -23,6 +23,8 @@ makedocs(; "Task Affinity" => "task-affinity.md", "Data Management" => "data-management.md", "Distributed Arrays" => "darray.md", + "Sparse Arrays" => "sparse-arrays.md", + "Iterative Solvers" => "iterative-solving.md", "Streaming Tasks" => "streaming.md", "Scopes" => "scopes.md", "Processors" => "processors.md", diff --git a/docs/src/index.md b/docs/src/index.md index 4b79aaa72..3e4b8ca59 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -363,6 +363,76 @@ collect(DA) # returns a `Matrix{Float64}` ----- +## Quickstart: Sparse Arrays + +A `DArray` can hold sparse tiles, giving a distributed, tiled sparse matrix. +Load a sparse backend (`SparseArrays`) to enable it. + +For more details: [Sparse Distributed Arrays](@ref) + +### Distribute or allocate a sparse `DArray` + +```julia +using SparseArrays + +# Partition an existing sparse matrix into sparse tiles +DA = distribute(sprand(1000, 1000, 0.01), Blocks(250, 250)) + +# Or allocate sparse DArrays directly +Z = spzeros(Blocks(250, 250), Float64, 1000, 1000) +R = sprand(Blocks(250, 250), Float64, (1000, 1000), 0.01) +``` + +### Multiply sparse arrays + +```julia +using LinearAlgebra +x = distribute(rand(1000), Blocks(250)) +y = DA * x # distributed sparse matrix-vector multiply (dense result) +C = DA * DA # distributed sparse-sparse matmul (sparse result) +``` + +`collect(DA)` returns a dense `Array`; operate on the `DArray` to stay sparse. + +----- + +## Quickstart: Iterative Solvers + +Dagger provides distributed, matrix-free Krylov solvers (`cg`, `minres`, +`gmres`, `bicgstab`) for `A x = b`. Load `Krylov` to enable them. + +For more details: [Iterative Solvers](@ref) + +### Solve a sparse system + +```julia +using SparseArrays, Krylov, LinearAlgebra + +n = 1000 +A = spdiagm(-1 => fill(-1.0, n-1), 0 => fill(2.0, n), 1 => fill(-1.0, n-1)) +DA = distribute(A, Blocks(250, 250)) # square tiles +b = distribute(rand(n), Blocks(250)) + +x, stats = Dagger.cg(DA, b) +@show stats.solved, stats.niter +``` + +### Add a preconditioner + +```julia +using AlgebraicMultigrid # enables Dagger.AMGPreconditioner + +P = Dagger.AMGPreconditioner(DA) # build once +x, stats = Dagger.cg(DA, b; M = P) # pass as `M` +``` + +Other preconditioners: `Dagger.JacobiPreconditioner`, +`Dagger.BlockJacobiPreconditioner` (core), and `Dagger.BlockILUPreconditioner` +(load `IncompleteLU`). Any object implementing `mul!(y, A, x)` over `DVector`s +can be used as a matrix-free operator `A`. + +----- + ## Quickstart: Stencil Operations Dagger's `@stencil` macro allows for easy specification of stencil operations on `DArray`s, often used in simulations and image processing. These operations typically involve updating an element based on the values of its neighbors. diff --git a/docs/src/iterative-solving.md b/docs/src/iterative-solving.md new file mode 100644 index 000000000..31b62d755 --- /dev/null +++ b/docs/src/iterative-solving.md @@ -0,0 +1,226 @@ +# Iterative Solvers + +Dagger provides distributed, **matrix-free** Krylov solvers for linear systems +`A x = b`, built on top of [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl). +They run entirely over `DArray`s: the operator `A` may be a dense or +[sparse](@ref "Sparse Distributed Arrays") `DMatrix`, or any object that +implements distributed `mul!(y, A, x)` over `DVector`s. This makes them a natural +fit for large, sparse systems arising from ODE/PDE discretizations. + +Solvers are provided through a package extension; load `Krylov` to enable them: + +```julia +using Distributed +addprocs(4) +using Dagger, SparseArrays, Krylov +``` + +If you call a solver without `Krylov` loaded, you'll get a clear error telling +you to `using Krylov`. + +## Solving a system + +```julia +using SparseArrays, Krylov + +# A symmetric positive-definite operator (1-D Laplacian) and a right-hand side +n = 1000 +A = spdiagm(-1 => fill(-1.0, n-1), 0 => fill(2.0, n), 1 => fill(-1.0, n-1)) +DA = distribute(A, Blocks(250, 250)) # square tiles (see below) +b = distribute(rand(n), Blocks(250)) + +x, stats = Dagger.cg(DA, b) +@show stats.solved, stats.niter +``` + +Each solver returns a tuple `(x, stats)`, where `x` is a `DVector` and `stats` +is Krylov's statistics object (`stats.solved`, `stats.niter`, residual history, +etc.). Keyword arguments such as `atol`, `rtol`, and `itmax` are forwarded +to Krylov. + +### Available methods + +| Function | Operator class | Notes | +|---------------------|-------------------------------|------------------------------------| +| [`Dagger.cg`](@ref) | symmetric positive-definite | cheapest; the PDE workhorse | +| [`Dagger.minres`](@ref) | symmetric (indefinite ok) | saddle-point / indefinite systems | +| [`Dagger.gmres`](@ref) | general nonsymmetric | robust; `restart`/`memory` to bound memory | +| [`Dagger.bicgstab`](@ref) | general nonsymmetric | short recurrence, low memory | + +[`Dagger.krylov_solve`](@ref) is a generic entry point taking the method as a +symbol (`:cg`, `:minres`, `:gmres`, `:bicgstab`): + +```julia +x, stats = Dagger.krylov_solve(:gmres, DA, b; memory=50) +``` + +## Matrix-free operators + +The solvers never form `A⁻¹`, and they never require `A` to be a materialized +matrix. They only need `mul!(y, A, x)` (and `mul!(y, A', x)` for two-sided +methods like GMRES/BiCGStab) to work over `DVector`s. So you can pass any custom +operator type that implements distributed `mul!`: + +```julia +struct MyStencilOperator + # ...your distributed state... +end +function LinearAlgebra.mul!(y::Dagger.DVector, A::MyStencilOperator, x::Dagger.DVector) + # fill y with A*x using Dagger tasks / datadeps + return y +end + +x, stats = Dagger.cg(MyStencilOperator(...), b) +``` + +Workspace vectors are allocated via `similar(b)`, so every internal vector +inherits `b`'s element type **and** partitioning. The distributed BLAS-1 building +blocks the solvers rely on — `dot`, `norm`, `axpy!`, `axpby!`, `rmul!`, +`copyto!`, `fill!`, broadcasting — are all implemented for `DVector` and align +mismatched partitionings automatically. + +!!! note "Square tiles for `DMatrix` operators" + When the operator is a `DMatrix`, use **square tiles** (`Blocks(k, k)`). The + solver's workspace vectors are all allocated as `similar(b)` (one + partitioning), and each must serve as both the length-`n` input and the + length-`n` output of `mul!(y, A, x)`. The distributed SpMV requires the input + to match `A`'s column blocks and the output to match `A`'s row blocks, which + is only simultaneously possible when those block sizes are equal. Any uniform + square tile size works (a ragged final block is fine). This constraint is on + how *you* partition the operator — generic code calling `mul!`/`dot` never + needs to know about it. + +## Preconditioners + +A preconditioner accelerates convergence by approximating `A⁻¹`. Dagger's +preconditioners follow Krylov's `ldiv=false` convention: the object **represents +the inverse operator `M⁻¹`** and is applied via `mul!(y, P, x)` (computing +`y = M⁻¹ x`). Pass one as the `M` keyword: + +```julia +P = Dagger.JacobiPreconditioner(DA) +x, stats = Dagger.cg(DA, b; M = P) +``` + +The built-in preconditioners, from cheapest to strongest: + +| Preconditioner | Needs | Idea | +|--------------------------------------|------------------------|--------------------------------------------| +| [`Dagger.JacobiPreconditioner`](@ref) | (core) | scale by `1 ./ diag(A)` | +| [`Dagger.BlockJacobiPreconditioner`](@ref) | (core) | exact `lu` solve per diagonal tile | +| [`Dagger.BlockILUPreconditioner`](@ref) | `IncompleteLU` | incomplete-LU (drop tol `τ`) per tile | +| [`Dagger.AMGPreconditioner`](@ref) | `AlgebraicMultigrid` | AMG V-cycle per tile | + +```julia +using AlgebraicMultigrid, IncompleteLU + +x, _ = Dagger.cg(DA, b; M = Dagger.BlockJacobiPreconditioner(DA)) +x, _ = Dagger.cg(DA, b; M = Dagger.BlockILUPreconditioner(DA; τ = 0.01)) +x, _ = Dagger.gmres(DA, b; M = Dagger.AMGPreconditioner(DA; method = :ruge_stuben)) +``` + +### How block preconditioners are distributed + +`BlockJacobiPreconditioner`, `BlockILUPreconditioner`, and `AMGPreconditioner` +are all **block-diagonal** preconditioners: they build one operator per diagonal +tile of `A` and apply them independently per block. They share a common +mechanism ([`Dagger.AbstractBlockPreconditioner`](@ref)): + +- The per-tile operator (an `lu`/ILU factorization, or an AMG hierarchy) is built + **once**, at construction. +- A factorization/hierarchy generally cannot be moved between workers (sparse + `lu` factors and AMG hierarchies hold process-bound resources). So each + operator is **pinned** to the worker owning its tile, and every apply for that + block is scheduled there — only the (small, movable) vector chunks are + transferred. + +A useful consequence of the per-tile design: with a **single tile** +(`Blocks(n, n)`), any of these becomes a *global* preconditioner over the whole +matrix (e.g. a global AMG, or an exact direct solve for block-Jacobi). With many +tiles, it becomes a scalable block-Jacobi / additive-Schwarz variant that trades +some convergence for parallelism. As with the operator, these require square +diagonal tiles. + +### Choosing a preconditioner + +- **SPD elliptic (Poisson-like) problems:** `AMGPreconditioner` gives near + mesh-independent convergence and is usually the best choice; `cg` as the + solver. +- **General sparse systems:** `BlockILUPreconditioner` is a solid, cheap-setup + general-purpose option; pair with `gmres` or `bicgstab`. +- **Quick baseline / very well-conditioned systems:** `JacobiPreconditioner` (or + none) may suffice. +- **Strong per-subdomain coupling:** `BlockJacobiPreconditioner` (exact tile + solves) is stronger than diagonal Jacobi. + +## Sparse direct solvers + +For systems that fit on a single worker, Dagger also offers **direct** sparse +solves via pure-Julia factorization backends. Unlike the C-bound `UmfpackLU`, +these factorizations are plain Julia data, so Dagger can move and schedule them +freely. + +Load `PureKLU` (KLU; good for unsymmetric/circuit systems) or `PureUMFPACK` +(UMFPACK-style multifrontal LU): + +```julia +using SparseArrays, PureKLU, PureUMFPACK + +A = distribute(sprand(2000, 2000, 0.005) + 10I, Blocks(500, 500)) +b = distribute(rand(2000), Blocks(500)) + +F = Dagger.klu(A) # or Dagger.splu(A) +x = F \ b # returns a DVector partitioned like b +``` + +`Dagger.klu`/`Dagger.splu` gather the sparse `DMatrix` into one +`SparseMatrixCSC` (without densifying), factor it once, and return a +[`Dagger.DaggerSparseLU`](@ref) supporting `F \ b` and `ldiv!(x, F, b)`. Factor +once, solve many right-hand sides cheaply. + +There are also **block direct preconditioners** that factor each diagonal tile +exactly (`Dagger.BlockKLUPreconditioner`, `Dagger.BlockUMFPACKPreconditioner`), +usable like the other block preconditioners. With a single tile they are exact +whole-matrix solves; with many tiles they are exact-block-Jacobi preconditioners +for the iterative solvers. + +## A worked example: implicit time stepping + +Implicit ODE/PDE integrators repeatedly solve systems with the same operator +`(I - Δt·L)` and changing right-hand sides. Build the preconditioner once and +reuse it across steps: + +```julia +using SparseArrays, Krylov, AlgebraicMultigrid + +L = distribute(laplacian, Blocks(k, k)) # discretized operator (square tiles) +A = I - Δt * L # or a custom matrix-free operator +P = Dagger.AMGPreconditioner(L) # build hierarchy once + +u = distribute(u0, Blocks(k)) +for step in 1:nsteps + rhs = ... # depends on current state + u, stats = Dagger.cg(A, rhs; M = P, rtol = 1e-8) +end +``` + +## API + +```@docs +Dagger.cg +Dagger.minres +Dagger.gmres +Dagger.bicgstab +Dagger.krylov_solve +Dagger.AbstractDaggerPreconditioner +Dagger.JacobiPreconditioner +Dagger.AbstractBlockPreconditioner +Dagger.BlockJacobiPreconditioner +Dagger.BlockILUPreconditioner +Dagger.AMGPreconditioner +Dagger.BlockKLUPreconditioner +Dagger.BlockUMFPACKPreconditioner +Dagger.klu +Dagger.splu +Dagger.DaggerSparseLU +``` diff --git a/docs/src/sparse-arrays.md b/docs/src/sparse-arrays.md new file mode 100644 index 000000000..9935b033c --- /dev/null +++ b/docs/src/sparse-arrays.md @@ -0,0 +1,177 @@ +# Sparse Distributed Arrays + +Dagger's [`DArray`](@ref) can hold **sparse** tiles, giving you a distributed, +tiled sparse matrix (or vector) that participates in the same scheduling, +[Datadeps](@ref), and linear-algebra machinery as dense `DArray`s. This is the +foundation for distributed sparse matrix multiplication and the +[matrix-free iterative solvers](@ref "Iterative Solvers"). + +Sparse support is provided through package extensions, so you opt in by loading a +sparse backend: + +- **`SparseArrays`** (the standard library) — tiles are `SparseMatrixCSC` / + `SparseVector` on the CPU. This is the default, well-supported backend. +- **GPU + `SparseArrays`** — under a GPU compute scope (`cuda_gpu`, `rocm_gpu`, + `cl_device`, `metal_gpu`, `intel_gpu`), tiles use the vendor sparse type when + available (CUDA cuSPARSE / AMDGPU rocSPARSE) or else + `Dagger.DeviceSparseMatrixCSC` (OpenCL / Metal / oneAPI) with host SpGEMM/SpMV + fallbacks. Load the GPU package together with `SparseArrays`. +- **`Finch`** — tiles are `Finch.Tensor`s, enabling a wider range of sparse + formats. This backend is more experimental (CPU only for now). + +```julia +using Distributed +addprocs(4) +using Dagger, SparseArrays +``` + +!!! note "Load order with workers" + As with all Dagger usage, add your workers *before* `using Dagger` and the + backend package, so the packages load on every worker. See the note at the + top of the [home page](@ref "Dagger: A framework for out-of-core and parallel execution"). + +## Creating a sparse `DArray` + +### From an existing sparse array + +`distribute` accepts a sparse matrix or vector and partitions it into sparse +tiles according to a `Blocks` specification: + +```julia +using SparseArrays +A = sprand(1000, 1000, 0.01) # a SparseMatrixCSC +DA = distribute(A, Blocks(250, 250)) # a 4×4 grid of sparse tiles +``` + +Each tile is a sparse matrix in its own right, stored on one of the workers. + +### Allocating directly + +You can also allocate a sparse `DArray` without first building a local sparse +array, using the `Blocks`-aware methods of `spzeros` and `sprand`: + +```julia +using SparseArrays + +# All-zeros sparse DArray, Float64, 1000×1000 in 250×250 tiles +Z = spzeros(Blocks(250, 250), Float64, 1000, 1000) + +# Random sparse DArray with ~1% nonzeros per tile +R = sprand(Blocks(250, 250), Float64, (1000, 1000), 0.01) +``` + +These run the per-tile allocation on the owning worker, so no large sparse array +is ever materialized on a single process. + +### Converting back to a dense array + +`collect` gathers the tiles and returns a **dense** `Array`: + +```julia +M = collect(DA) # dense Matrix{Float64} +``` + +To keep data sparse and distributed, operate on the `DArray` directly rather +than collecting. + +## How it works + +### The `DSparseArray` wrapper + +Internally, each sparse tile is wrapped in a [`Dagger.DSparseArray`](@ref) — a +small mutable container holding the actual sparse storage (`mat`): + +```julia +mutable struct DSparseArray{T,N} <: AbstractArray{T,N} + mat # e.g. a SparseMatrixCSC, SparseVector, or Finch.Tensor +end +``` + +`DSparseVector{T}` and `DSparseMatrix{T}` are the 1- and 2-dimensional aliases. + +The wrapper exists because **sparse storage is reallocated on writes**. Many +sparse operations (e.g. `A*B`, or anything that changes the sparsity pattern) +cannot update their result in place — they produce a brand-new sparse array of a +different size. Datadeps, however, tracks data dependencies by the *identity* of +the objects it manages, and it does not support objects that grow or shrink. The +`DSparseArray` wrapper solves this: its identity is stable, and a write simply +swaps the inner `mat` for the new storage: + +```julia +# Conceptually, how an in-place sparse update is modeled: +tile.mat = tile.mat * other # identity of `tile` is unchanged +``` + +### Aliasing as a whole + +Because the inner storage may move, it is never safe to alias *part* of a sparse +tile (e.g. via a `view` or a strided sub-region). Dagger therefore treats a +`DSparseArray` as an **indivisible aliasing unit**: any access — including +through `view`, `transpose`, `adjoint`, or `reshape` — resolves to the +container's stable whole-object aliasing. This is what keeps Datadeps correct +when sparse writes reallocate storage. (For the curious: the type opts in via +`Dagger.aliases_as_whole`, and Datadeps' `aliasing_root` unwraps any wrapper of a +`DSparseArray` before computing aliasing. Calling `pointer` on a `DSparseArray` +intentionally errors, to catch any code path that tries to treat it as raw +strided memory.) + +The practical upshot: you can pass sparse tiles, or views of sparse `DArray`s, +into `Dagger.spawn_datadeps` regions and trust that read/write ordering is +tracked correctly. + +## Operations + +Sparse `DArray`s support the array operations that have distributed +implementations, including: + +- **Matrix–matrix multiply** (`A * B`, `mul!`), sparse × sparse, producing a + sparse result. +- **Sparse matrix–vector multiply** (SpMV: `A * x`, `mul!(y, A, x)`) with a + sparse matrix and dense vectors — the workhorse of iterative solvers. +- **Transpose/adjoint**, **`collect`**, and elementwise/`norm` operations. + +```julia +using SparseArrays, LinearAlgebra +A = distribute(sprand(1000, 1000, 0.01), Blocks(250, 250)) +x = distribute(rand(1000), Blocks(250)) + +y = A * x # distributed SpMV -> dense DVector +C = A * A # distributed sparse-sparse matmul -> sparse DArray +``` + +### Partitioning guidance + +- Choose tile sizes so each tile comfortably fits on a worker, and so the number + of tiles is at least the number of workers (for parallelism). +- For **square operators used with the iterative solvers**, use **square tiles** + (`Blocks(k, k)`); see [Iterative Solvers](@ref) for why. +- Operands with mismatched partitionings are aligned automatically (by buffered + copy) where needed, but matching partitionings avoid that overhead. + +## Backends + +### `SparseArrays` (recommended) + +Tiles are `SparseMatrixCSC` (matrices) or `SparseVector` (vectors). This backend +provides efficient SpMV (including transposed/adjoint operands) and uses +`SparseArrays`' own `*` for sparse–sparse products. + +### `Finch` (experimental) + +Loading `Finch` makes tiles `Finch.Tensor`s, supporting a broader set of sparse +and structured formats. Finch support is newer and exercised by a dedicated test +suite; prefer `SparseArrays` unless you specifically need a Finch format. + +## Limitations + +- `collect` densifies; there is no sparse-preserving global gather. +- A sparse tile is aliased as a whole — Datadeps cannot track independent writes + to disjoint sub-regions of a single sparse tile (use finer tiling instead). +- Not every dense `DArray` operation has a sparse counterpart yet; sparse support + focuses on multiplication and the building blocks needed for iterative solving. + +## API + +```@docs +Dagger.DSparseArray +``` diff --git a/docs/src/stencils.md b/docs/src/stencils.md index bfb4838df..27da654e0 100644 --- a/docs/src/stencils.md +++ b/docs/src/stencils.md @@ -36,6 +36,32 @@ B = zeros(Blocks(2, 2), Int, 4, 4) In these examples, the stencil operation is executed for each chunk of the target array. The `idx` variable corresponds to the indices within each chunk. +## Choosing a Block Size + +`@stencil` spawns one task per chunk per expression, so the array's block size sets +both the available parallelism and the number of tasks. Scheduling a task is not +free -- it costs on the order of a hundred microseconds -- so blocks want to be +large enough that the work in a chunk dwarfs that, and numerous enough to keep the +processors busy. Roughly one chunk per processor is a good starting point, which is +what `AutoBlocks` picks: + +```julia +A = rand(AutoBlocks(), Float64, 4096, 4096) +``` + +Partitioning far more finely than the processor count is the common mistake, and it +is expensive: a 4096x4096 Jacobi sweep over 12 threads runs about 1.7x slower with +`Blocks(512, 512)` (64 chunks) than with `AutoBlocks()` (12 chunks), because the +per-task cost is paid 64 times per sweep instead of 12 while the arithmetic is +unchanged. Coarsening past one chunk per processor is bounded by the other side: +chunks are the unit of parallelism, so fewer chunks than processors leaves +processors idle. + +Note that this argues only against *needlessly* fine partitioning. Smaller chunks +still pay for themselves when they buy something else: load balance across uneven +processors, overlap of communication with computation, or fitting a working set +into cache or device memory. + ## Allocation and Return Syntax The `@stencil` macro also supports a functional syntax that automatically allocates an output `DArray` and returns it. This is triggered when the stencil expression (or the last expression in a block) is a "naked" expression (not an assignment). diff --git a/ext/AlgebraicMultigridExt.jl b/ext/AlgebraicMultigridExt.jl new file mode 100644 index 000000000..dd68bbcfa --- /dev/null +++ b/ext/AlgebraicMultigridExt.jl @@ -0,0 +1,45 @@ +module AlgebraicMultigridExt + +import AlgebraicMultigrid +import SparseArrays +import Dagger +import Dagger: DMatrix +import LinearAlgebra + +# Algebraic-multigrid preconditioner, built per diagonal tile on top of Dagger's +# block-preconditioner machinery (see `src/array/iterativesolvers.jl`). The AMG +# hierarchy's expensive setup runs once per tile at construction; each apply runs +# one V-cycle (a sequence of SpMVs, smoother sweeps, and a coarse solve). + +_as_sparse(A::SparseArrays.SparseMatrixCSC) = A +_as_sparse(A::AbstractMatrix) = SparseArrays.sparse(A) + +function _amg_operator(tile, method::Symbol; kwargs...) + S = _as_sparse(Dagger._tile_matrix(tile)) + ml = if method === :ruge_stuben + AlgebraicMultigrid.ruge_stuben(S; kwargs...) + elseif method === :smoothed_aggregation + AlgebraicMultigrid.smoothed_aggregation(S; kwargs...) + else + throw(ArgumentError("AMGPreconditioner: unknown method $(method); use \ + :ruge_stuben or :smoothed_aggregation")) + end + return AlgebraicMultigrid.aspreconditioner(ml) +end + +function Dagger.AMGPreconditioner(A::DMatrix; method::Symbol=:ruge_stuben, kwargs...) + build = tile -> _amg_operator(tile, method; kwargs...) + return Dagger._build_block_preconditioner(Dagger.AMGPreconditioner, A, + "AMGPreconditioner", build) +end + +# An AMG preconditioner applies a V-cycle via `ldiv!` (`\` is not defined). +Dagger._block_apply!(y, p::AlgebraicMultigrid.Preconditioner, x) = + (LinearAlgebra.ldiv!(y, p, x); return nothing) + +# The hierarchy is read-only and pinned to its worker; it is never written and +# its internals (a `Vector{Level}`) aren't introspectable by datadeps, so treat +# it as non-aliasing rather than recursing into it. +Dagger.aliasing(::AlgebraicMultigrid.Preconditioner) = Dagger.NoAliasing() + +end # module AlgebraicMultigridExt diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index 813763ce1..02cff791f 100644 --- a/ext/CUDAExt.jl +++ b/ext/CUDAExt.jl @@ -390,13 +390,9 @@ Adapt.adapt_structure(to::CUDA.KernelAdaptor, H::Dagger.HaloArray) = H.halo_width; own_center=H.own_center) function Dagger.inner_stencil_proc!(::CuArrayDeviceProc, f, output, read_vars) - Dagger.Kernel(_inner_stencil!)(f, output, read_vars; ndrange=size(output)) + Dagger.gpu_stencil_sweep!(f, output, read_vars) return end -@kernel function _inner_stencil!(f, output, read_vars) - idx = @index(Global, Cartesian) - f(idx, output, read_vars) -end Dagger.gpu_processor(::Val{:CUDA}) = CuArrayDeviceProc Dagger.gpu_can_compute(::Val{:CUDA}) = CUDA.has_cuda() diff --git a/ext/CUDASparseArraysExt.jl b/ext/CUDASparseArraysExt.jl new file mode 100644 index 000000000..072577ff3 --- /dev/null +++ b/ext/CUDASparseArraysExt.jl @@ -0,0 +1,177 @@ +module CUDASparseArraysExt + +# Sparse DArray support for CUDA / cuSPARSE. Loaded when both CUDA and +# SparseArrays are available (see Project.toml combo extension). + +import Dagger +import CUDA +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import CUDA: CuArray +import CUDA.CUSPARSE: CuSparseMatrixCSC, CuSparseMatrixCSR, CuSparseVector + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +const CUDAExt = Base.get_extension(Dagger, :CUDAExt)::Module +using .CUDAExt: CuArrayDeviceProc +_with_context(f, proc) = CUDAExt.with_context(f, proc) + +#----- Memory / aliasing ------------------------------------------------------- + +Dagger.value_memory_space(x::CuSparseMatrixCSC) = Dagger.memory_space(x.nzVal) +Dagger.value_memory_space(x::CuSparseMatrixCSR) = Dagger.memory_space(x.nzVal) +Dagger.value_memory_space(x::CuSparseVector) = Dagger.memory_space(x.nzVal) +Dagger.memory_space(x::CuSparseMatrixCSC) = Dagger.value_memory_space(x) +Dagger.memory_space(x::CuSparseMatrixCSR) = Dagger.value_memory_space(x) +Dagger.memory_space(x::CuSparseVector) = Dagger.value_memory_space(x) + +function Dagger.aliasing(x::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}, _=identity) + space = Dagger.value_memory_space(x) + ptr = Dagger.RemotePtr{Cvoid}(UInt(pointer_from_objref(x)), space) + return Dagger.ObjectAliasing(ptr, sizeof(typeof(x))) +end +Dagger.aliases_as_whole(::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}) = true +Dagger.maybe_wrap_tile(x::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}) = + Dagger.DSparseArray(x) + +#----- Allocation -------------------------------------------------------------- + +Dagger.allocate_sparse_zeros(::CuArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + CuSparseMatrixCSC(SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_zeros(::CuArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + CuSparseVector(SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_rand(::CuArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + CuSparseMatrixCSC(SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::CuArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + CuSparseVector(SparseArrays.sprand(T, dims..., sparsity)) + +#----- Similar / collect ------------------------------------------------------- + +Dagger._sparse_similar(::CuSparseMatrixCSC, ::Type{T}, dims::Dims{2}) where T = + CuSparseMatrixCSC(SparseArrays.spzeros(T, dims...)) +Dagger._sparse_similar(::CuSparseVector, ::Type{T}, dims::Dims{1}) where T = + CuSparseVector(SparseArrays.spzeros(T, dims...)) +Dagger._sparse_collect(A::CuSparseMatrixCSC) = SparseMatrixCSC(A) +Dagger._sparse_collect(A::CuSparseMatrixCSR) = SparseMatrixCSC(CuSparseMatrixCSC(A)) +Dagger._sparse_collect(A::CuSparseVector) = SparseVector(A) +Dagger._sparse_copy(A::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}) = copy(A) + +#----- Host ↔ device helpers --------------------------------------------------- + +_to_cu_sparse(x::SparseMatrixCSC) = CuSparseMatrixCSC(x) +_to_cu_sparse(x::SparseVector) = CuSparseVector(x) +_to_cu_sparse(x::CuSparseMatrixCSC) = x +_to_cu_sparse(x::CuSparseMatrixCSR) = CuSparseMatrixCSC(x) +_to_cu_sparse(x::CuSparseVector) = x +_to_cu_sparse(x::Dagger.DeviceSparseMatrixCSC) = CuSparseMatrixCSC(SparseMatrixCSC(x)) +_to_cu_dsparse(x::Dagger.DSparseArray) = Dagger.DSparseArray(_to_cu_sparse(x.mat)) + +_to_host_sparse(x::CuSparseMatrixCSC) = SparseMatrixCSC(x) +_to_host_sparse(x::CuSparseMatrixCSR) = SparseMatrixCSC(CuSparseMatrixCSC(x)) +_to_host_sparse(x::CuSparseVector) = SparseVector(x) +_to_host_sparse(x::SparseMatrixCSC) = x +_to_host_sparse(x::SparseVector) = x +_to_host_dsparse(x::Dagger.DSparseArray) = Dagger.DSparseArray(_to_host_sparse(x.mat)) + +_to_csc(A::CuSparseMatrixCSC) = A +_to_csc(A::CuSparseMatrixCSR) = CuSparseMatrixCSC(A) +_to_csc(A) = CuSparseMatrixCSC(A) + +function _op_csc(X, t::Char) + Xc = _to_csc(X) + t == 'N' && return Xc + Sh = SparseMatrixCSC(Xc) + t == 'T' && return CuSparseMatrixCSC(SparseArrays.sparse(transpose(Sh))) + t == 'C' && return CuSparseMatrixCSC(SparseArrays.sparse(adjoint(Sh))) + throw(ArgumentError("Invalid trans char: $t")) +end + +_apply_trans(X, t::Char) = + t == 'N' ? X : t == 'T' ? transpose(X) : t == 'C' ? adjoint(X) : + throw(ArgumentError("Invalid trans char: $t")) + +#----- Move -------------------------------------------------------------------- + +function Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(CuSparseMatrixCSC(x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::SparseVector) + _with_context(to_proc) do + return Dagger.DSparseArray(CuSparseVector(x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + return _to_cu_dsparse(x) + end +end +function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + CUDA.synchronize() + return _to_host_dsparse(x) + end +end +function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CPUProc, + x::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}) + _with_context(from_proc) do + CUDA.synchronize() + return Dagger.DSparseArray(_to_host_sparse(x)) + end +end +function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CuArrayDeviceProc, x::Dagger.DSparseArray) + if from_proc == to_proc + _with_context(CUDA.synchronize, from_proc) + return x + end + _with_context(to_proc) do + return _to_cu_dsparse(_to_host_dsparse(x)) + end +end + +#----- Tile kernels ------------------------------------------------------------ + +function Dagger.matmatmul!( + C::Dagger.DSparseMatrix, + transA::Char, transB::Char, + A::Union{CuSparseMatrixCSC,CuSparseMatrixCSR}, + B::Union{CuSparseMatrixCSC,CuSparseMatrixCSR}, + alpha, beta +) + AB = _to_csc(_op_csc(A, transA) * _op_csc(B, transB)) + prod = _to_csc(isone(alpha) ? AB : alpha * AB) + if iszero(beta) + C.mat = prod + elseif isone(beta) + C.mat = _to_csc(prod + _to_csc(C.mat)) + else + C.mat = _to_csc(prod + beta * _to_csc(C.mat)) + end + return C +end + +function Dagger.matvecmul!(C::CuArray, transA::Char, + A::Union{CuSparseMatrixCSC,CuSparseMatrixCSR}, + B::CuArray, alpha, beta) + LinearAlgebra.mul!(C, _apply_trans(A, transA), B, alpha, beta) + return C +end + +function Dagger.transpose_tile(B::CuSparseMatrixCSC) + return CuSparseMatrixCSC(SparseArrays.sparse(SparseMatrixCSC(B)')) +end +function Dagger.transpose_tile(B::CuSparseMatrixCSC, uplo::Char) + Bh = SparseMatrixCSC(B) + Bt = uplo == 'U' ? SparseArrays.triu(Bh) : + uplo == 'L' ? SparseArrays.tril(Bh) : + throw(ArgumentError("uplo must be 'U' or 'L', got $uplo")) + Ct = Bt + Bt' + for i in 1:LinearAlgebra.checksquare(Bh) + Ct[i, i] = Bh[i, i] + end + return CuSparseMatrixCSC(Ct) +end + +end # module CUDASparseArraysExt diff --git a/ext/FinchExt.jl b/ext/FinchExt.jl new file mode 100644 index 000000000..68bf8f3d8 --- /dev/null +++ b/ext/FinchExt.jl @@ -0,0 +1,186 @@ +module FinchExt + +import Finch +import LinearAlgebra +import Dagger +import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseArray, DSparseMatrix + +Dagger._sparse_collect(A::Finch.Tensor) = Array(A) +# Finch's generic `similar` produces tensor formats that destabilize later +# `@finch`/`@einsum` kernels (observed as segfaults during tile moves), so +# allocate an empty COO-backed tile explicitly instead. +Dagger._sparse_similar(::Finch.Tensor, ::Type{T}, dims::Dims) where {T} = + Finch.fspzeros(T, dims...) +Dagger.maybe_wrap_tile(x::Finch.Tensor) = DSparseArray(x) + +function Finch.fspzeros(p::BlocksOrAuto, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Dagger.Tapes.resolve_partitioning(T, dims, p, assignment) + d = Dagger.ArrayDomain(map(x->1:x, dims)) + N = length(dims) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Finch.fspzeros(T, _dims...)), false, d, Dagger.partition(part, d), part, assign; + return_type=DSparseArray{T,N}) + return Dagger.Tapes.track!(Dagger._to_darray(a), plan) +end +Finch.fspzeros(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = + Finch.fspzeros(p, T, dims; assignment) +Finch.fspzeros(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbitrary) = + Finch.fspzeros(p, Float64, dims; assignment) +Finch.fspzeros(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = + Finch.fspzeros(p, Float64, dims; assignment) + +function Finch.fsprand(p::BlocksOrAuto, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) + part, assign, plan = Dagger.Tapes.resolve_partitioning(T, dims, p, assignment) + d = Dagger.ArrayDomain(map(x->1:x, dims)) + N = length(dims) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Finch.fsprand(T, _dims..., sparsity)), false, d, Dagger.partition(part, d), part, assign; + return_type=DSparseArray{T,N}) + return Dagger.Tapes.track!(Dagger._to_darray(a), plan) +end +Finch.fsprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = + Finch.fsprand(p, T, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) +Finch.fsprand(p::BlocksOrAuto, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = + Finch.fsprand(p, Float64, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) +Finch.fsprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = + Finch.fsprand(p, Float64, dims, sparsity; assignment) + +# Materialize a (possibly lazy/`SwizzleArray`) result into a concrete sparse +# `Tensor`. `@einsum` returns a lazy `SwizzleArray` for transposed-output +# patterns, which is not itself a `Finch.Tensor` and corrupts subsequent +# accumulation steps; copying through `@finch` produces a clean `SparseCOO`. +function _finch_materialize(X, ::Type{T}, dims) where {T} + out = Finch.fspzeros(T, dims...) + Finch.@finch begin + out .= 0 + for j = _, i = _ + if X[i, j] != 0 + out[i, j] = X[i, j] + end + end + end + return out +end + +# Finch tensors do not define `Base.copy`; build a fresh equivalent tensor. +Dagger._sparse_copy(mat::Finch.Tensor) = _finch_materialize(mat, eltype(mat), size(mat)) + +# Finch tensors do not support `setindex!`, so copy-buffering writeback into a +# Finch tile densifies, applies the (partial-range) copy, then re-sparsifies. +function Dagger._sparse_copyto_view!(mat::Finch.Tensor, Brange, src) + dense = Array(mat) + copyto!(view(dense, Brange), src) + return _finch_materialize(dense, eltype(mat), size(mat)) +end + +# Compute `C_out += alpha * op(A) * op(B)` via Finch's `@einsum`, choosing the +# right index pattern for each transpose option. We deliberately use *explicit* +# index patterns (`A[k,i]`) rather than `Finch.swizzle`, because feeding a +# swizzled tensor into `@einsum` makes it return a lazy `SwizzleArray` (rather +# than a materialized `Tensor`), which breaks subsequent accumulation steps. +# Adjoint ('C') additionally applies `conj` inside the einsum. +function _finch_einsum_matmul!(C_out, A, B, transA::Char, transB::Char, alpha) + aT = transA == 'T' || transA == 'C' + aC = transA == 'C' + bT = transB == 'T' || transB == 'C' + bC = transB == 'C' + if !aT && !bT + bC ? Finch.@einsum(C_out[i,j] += alpha * (A[i,k] * conj(B[k,j]))) : + Finch.@einsum(C_out[i,j] += alpha * (A[i,k] * B[k,j])) + elseif !aT && bT + bC ? Finch.@einsum(C_out[i,j] += alpha * (A[i,k] * conj(B[j,k]))) : + Finch.@einsum(C_out[i,j] += alpha * (A[i,k] * B[j,k])) + elseif aT && !bT + aC ? Finch.@einsum(C_out[i,j] += alpha * (conj(A[k,i]) * B[k,j])) : + Finch.@einsum(C_out[i,j] += alpha * (A[k,i] * B[k,j])) + else # aT && bT + if aC && bC + Finch.@einsum C_out[i,j] += alpha * (conj(A[k,i]) * conj(B[j,k])) + elseif aC + Finch.@einsum C_out[i,j] += alpha * (conj(A[k,i]) * B[j,k]) + elseif bC + Finch.@einsum C_out[i,j] += alpha * (A[k,i] * conj(B[j,k])) + else + Finch.@einsum C_out[i,j] += alpha * (A[k,i] * B[j,k]) + end + end + return C_out +end + +function Dagger.matmatmul!( + C::DSparseMatrix, + transA::Char, + transB::Char, + A::Finch.Tensor, + B::Finch.Tensor, + alpha, + beta +) + if !isa(C.mat, Finch.Tensor) + # Not supported here, forward to generic matmatmul! + return Dagger.matmatmul!(C.mat, transA, transB, A, B, alpha, beta) + end + transA in ('N', 'T', 'C') || throw(ArgumentError("Invalid transA: $transA")) + transB in ('N', 'T', 'C') || throw(ArgumentError("Invalid transB: $transB")) + + # `@einsum` cannot write into an existing sparse tensor in place, so we build + # a fresh output tensor and (re)assign it to the wrapper. `DSparseMatrix` + # hides this reallocation from Datadeps. + # + # N.B. BEWARE: `@einsum Cm[i,j] = ...` (assignment into an existing tensor) + # does not work; use a fresh `C_out` with `+=`. + Cm = C.mat + C_out = Finch.fspzeros(eltype(Cm), size(Cm)...) + C_out = _finch_einsum_matmul!(C_out, A, B, transA, transB, alpha) + # `@einsum` may return a lazy `SwizzleArray`; materialize to a concrete + # sparse `Tensor` so accumulation and later iterations stay well-typed. + C_out = _finch_materialize(C_out, eltype(Cm), size(Cm)) + if beta != 0 + C_out = C_out + beta * Cm + end + C.mat = C_out + + return C +end + +# Sparse matrix-vector multiply tile kernel: `C = alpha*op(A)*B + beta*C` with a +# Finch tensor `A` and dense vectors `B`/`C`. Finch's `@einsum` does not reliably +# accumulate into a dense output, so we densify the single sparse tile and use a +# dense BLAS-backed `mul!`; this is a per-tile densification, not the whole matrix. +function Dagger.matvecmul!(C::AbstractVector, transA::Char, A::Finch.Tensor, B::AbstractVector, alpha, beta) + Ad = Array(A) + opA = transA == 'N' ? Ad : + transA == 'C' ? adjoint(Ad) : + transA == 'T' ? transpose(Ad) : + throw(ArgumentError("Invalid transA: $transA")) + LinearAlgebra.mul!(C, opA, B, alpha, beta) + return C +end + +# Tile transpose/symmetrization used by `copytri!`. Finch tensors don't support +# in-place `setindex!`, so we densify, build the result, and re-sparsify via +# `_finch_materialize`. +# - `uplo === nothing`: off-diagonal tile, return the conjugate transpose. +# - `uplo == 'U'`/`'L'`: diagonal tile, build the full Hermitian tile from the +# given triangle (matching the dense `copydiagtile!` semantics). +function Dagger.transpose_tile(B::Finch.Tensor, uplo=nothing) + _B = Array(B) + if uplo === nothing + C = adjoint(_B) + return _finch_materialize(C, eltype(_B), size(C)) + end + n = LinearAlgebra.checksquare(_B) + C = zeros(eltype(_B), n, n) + if uplo == 'U' + for j in 1:n, i in 1:n + C[i, j] = i <= j ? _B[i, j] : conj(_B[j, i]) + end + elseif uplo == 'L' + for j in 1:n, i in 1:n + C[i, j] = i >= j ? _B[i, j] : conj(_B[j, i]) + end + else + throw(ArgumentError("uplo must be 'U' or 'L', got $uplo")) + end + return _finch_materialize(C, eltype(_B), size(C)) +end + +end # module FinchExt \ No newline at end of file diff --git a/ext/IncompleteLUExt.jl b/ext/IncompleteLUExt.jl new file mode 100644 index 000000000..06ecab522 --- /dev/null +++ b/ext/IncompleteLUExt.jl @@ -0,0 +1,36 @@ +module IncompleteLUExt + +import IncompleteLU +import SparseArrays +import Dagger +import Dagger: DMatrix +import LinearAlgebra + +# Block incomplete-LU preconditioner, built per diagonal tile on top of Dagger's +# block-preconditioner machinery (see `src/array/iterativesolvers.jl`). Each tile +# gets an ILU factorization (drop tolerance `τ`) once at construction; the apply +# is a forward/backward substitution per block. + +_as_sparse(A::SparseArrays.SparseMatrixCSC) = A +_as_sparse(A::AbstractMatrix) = SparseArrays.sparse(A) + +function _ilu_operator(tile; τ=0.001, kwargs...) + S = _as_sparse(Dagger._tile_matrix(tile)) + return IncompleteLU.ilu(S; τ=τ, kwargs...) +end + +function Dagger.BlockILUPreconditioner(A::DMatrix; kwargs...) + build = tile -> _ilu_operator(tile; kwargs...) + return Dagger._build_block_preconditioner(Dagger.BlockILUPreconditioner, A, + "BlockILUPreconditioner", build) +end + +# An ILU factorization applies via `ldiv!` (`\` is not defined). +Dagger._block_apply!(y, F::IncompleteLU.ILUFactorization, x) = + (LinearAlgebra.ldiv!(y, F, x); return nothing) + +# The factorization is read-only and pinned to its worker; treat it as +# non-aliasing rather than recursing into its internals. +Dagger.aliasing(::IncompleteLU.ILUFactorization) = Dagger.NoAliasing() + +end # module IncompleteLUExt diff --git a/ext/IntelExt.jl b/ext/IntelExt.jl index 87bb7138f..2f25f7e3d 100644 --- a/ext/IntelExt.jl +++ b/ext/IntelExt.jl @@ -366,13 +366,9 @@ Adapt.adapt_structure(to::oneAPI.KernelAdaptor, H::Dagger.HaloArray) = H.halo_width; own_center=H.own_center) function Dagger.inner_stencil_proc!(::oneArrayDeviceProc, f, output, read_vars) - Dagger.Kernel(_inner_stencil!)(f, output, read_vars; ndrange=size(output)) + Dagger.gpu_stencil_sweep!(f, output, read_vars) return end -@kernel function _inner_stencil!(f, output, read_vars) - idx = @index(Global, Cartesian) - f(idx, output, read_vars) -end Dagger.gpu_processor(::Val{:oneAPI}) = oneArrayDeviceProc Dagger.gpu_can_compute(::Val{:oneAPI}) = oneAPI.functional() diff --git a/ext/IntelSparseArraysExt.jl b/ext/IntelSparseArraysExt.jl new file mode 100644 index 000000000..9cf2f2c01 --- /dev/null +++ b/ext/IntelSparseArraysExt.jl @@ -0,0 +1,68 @@ +module IntelSparseArraysExt + +# Sparse DArray support for oneAPI via DeviceSparseMatrixCSC (no vendor sparse +# library assumed). SpGEMM/SpMV fall back to host SparseArrays. + +import Dagger +import oneAPI +import Adapt +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import oneAPI: oneArray + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +const IntelExt = Base.get_extension(Dagger, :IntelExt)::Module +using .IntelExt: oneArrayDeviceProc +_with_context(f, proc) = IntelExt.with_context(f, proc) + +Adapt.adapt_structure(::Type{<:oneArray}, S::SparseMatrixCSC) = + Dagger.device_sparse_from_host(oneArray, S) + +Dagger.allocate_sparse_zeros(::oneArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + Dagger.device_sparse_from_host(oneArray, SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_zeros(::oneArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + SparseArrays.spzeros(T, dims...) +Dagger.allocate_sparse_rand(::oneArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + Dagger.device_sparse_from_host(oneArray, SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::oneArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) + +function Dagger.move(from_proc::CPUProc, to_proc::oneArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(Dagger.device_sparse_from_host(oneArray, x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::oneArrayDeviceProc, x::SparseVector) + return Dagger.DSparseArray(copy(x)) +end +function Dagger.move(from_proc::CPUProc, to_proc::oneArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + mat = x.mat + S = mat isa SparseMatrixCSC ? mat : SparseMatrixCSC(mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(oneArray, S)) + end +end +function Dagger.move(from_proc::oneArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + oneAPI.synchronize() + mat = x.mat + if mat isa Dagger.DeviceSparseMatrixCSC + return Dagger.DSparseArray(SparseMatrixCSC(mat)) + else + return Dagger.DSparseArray(copy(mat)) + end + end +end +function Dagger.move(from_proc::oneArrayDeviceProc, to_proc::oneArrayDeviceProc, x::Dagger.DSparseArray) + if from_proc == to_proc + return x + end + _with_context(to_proc) do + S = x.mat isa Dagger.DeviceSparseMatrixCSC ? SparseMatrixCSC(x.mat) : SparseMatrixCSC(x.mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(oneArray, S)) + end +end + +end # module IntelSparseArraysExt diff --git a/ext/KrylovExt.jl b/ext/KrylovExt.jl new file mode 100644 index 000000000..5d36ba2fa --- /dev/null +++ b/ext/KrylovExt.jl @@ -0,0 +1,36 @@ +module KrylovExt + +import Krylov +import Dagger +import Dagger: DVector, DMatrix +import LinearAlgebra + +# Distributed iterative solvers, implemented on top of Krylov.jl. +# +# We never form `A⁻¹`: Krylov only needs `mul!(y, A, x)` (and `mul!(y, A', x)` +# for two-sided methods), which Dagger provides over `DVector`s via its +# distributed SpMV/`gemv_dagger!` path. The remaining requirements -- `dot`, +# `norm`, `axpy!`, `axpby!`, `rmul!`, `copyto!`, `fill!`, broadcasting -- are the +# distributed BLAS-1 ops in `Dagger`. +# +# Workspace vectors are allocated through a `KrylovConstructor` built from +# `similar(b)`, so every internal vector inherits `b`'s element type *and* +# partitioning (`similar` preserves both). This keeps all `mul!`/`dot`/`axpy!` +# operands on a compatible chunk layout, which the distributed kernels require. + +function _dagger_krylov(method::Symbol, A, b::DVector; kwargs...) + kc = Krylov.KrylovConstructor(similar(b)) + workspace = Krylov.krylov_workspace(Val(method), kc) + Krylov.krylov_solve!(workspace, A, b; kwargs...) + return Krylov.solution(workspace), Krylov.statistics(workspace) +end + +Dagger.krylov_solve(method::Symbol, A, b::DVector; kwargs...) = + _dagger_krylov(method, A, b; kwargs...) + +Dagger.cg(A, b::DVector; kwargs...) = _dagger_krylov(:cg, A, b; kwargs...) +Dagger.minres(A, b::DVector; kwargs...) = _dagger_krylov(:minres, A, b; kwargs...) +Dagger.gmres(A, b::DVector; kwargs...) = _dagger_krylov(:gmres, A, b; kwargs...) +Dagger.bicgstab(A, b::DVector; kwargs...) = _dagger_krylov(:bicgstab, A, b; kwargs...) + +end # module KrylovExt diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 9274496b4..586a3ccd9 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -8,7 +8,8 @@ import Dagger: @dagdebug, @opcounter # extends with new methods (for MPI-specific types) or calls directly. import Dagger: AbstractAliasing, accelerate!, accel_matches_proc, aliased_object!, - AliasedObjectCache, AliasedObjectCacheStore, aliasing, bind_moved_argument, + AliasedObjectCache, AliasedObjectCacheStore, aliasing, aliasing_unwrapped, + bind_moved_argument, chunktype, ChunkView, check_uniform, check_uniformity!, CHECK_UNIFORMITY, cleanup_tasks_accel!, constrain, CPURAMMemorySpace, current_acceleration, CyclicProcGrid, datasize, default_enabled, @@ -24,7 +25,7 @@ import Dagger: move_rewrap_child_types, move_rewrap_header_mode, move_rewrap_parts, move_rewrap_result_type, move_type, multi_span_gather!, multi_span_scatter!, next_id, NoAliasing, Options, OSProc, - post_stage_array_chunks!, processors, ProcessScope, proc_in_scope, + partition_affinity_id, post_stage_array_chunks!, processors, ProcessScope, proc_in_scope, remotecall_endpoint_toplevel, RemotePtr, root_worker_id, same_node, schedule_argument_move, sch_handle, select_processors_uniform!, set_key_stored!, set_stored!, short_name, span_len, stage_acquire!, stage_release!, @@ -141,7 +142,15 @@ end MPIAcceleration() = MPIAcceleration(MPI.COMM_WORLD) function aliasing(accel::MPIAcceleration, x::Chunk, T) - handle = x.handle::MPIRef + handle = x.handle + # Chunks created under a temporary DistributedAcceleration (or a worker + # thread that did not inherit MPI TLS) carry a DRef; fall back to the + # Distributed unwrap path rather than hard-failing the typeassert. + if !(handle isa MPIRef) + return _with_default_acceleration() do + aliasing(x, T) + end + end @assert accel.comm == handle.comm "MPIAcceleration comm mismatch" tag = to_tag() check_uniform(tag) @@ -472,6 +481,12 @@ function memory_spaces(proc::MPIProcessor) end root_worker_id(mem_space::MPIMemorySpace) = myid() +# Hierarchical Datadeps partitions by this id (analogous to Distributed worker +# id). MPI always reports `root_worker_id == myid()`, so affinity must key on +# the owning rank instead. +partition_affinity_id(space::MPIMemorySpace) = space.rank +partition_affinity_id(proc::MPIProcessor) = proc.rank +partition_affinity_id(proc::MPIOSProc) = proc.rank function processors(memSpace::MPIMemorySpace) rawProc = Processor[] @@ -1757,7 +1772,9 @@ function aliasing(accel::MPIAcceleration, x::ChunkView, dep_mod) if handle.rank == rank ainfo = _with_default_acceleration() do v = view(unwrap(x.chunk), x.slices...) - aliasing(v, dep_mod) + # Resolve whole-object containers (e.g. `DSparseArray`) where `v` + # lives; see `aliasing_unwrapped`. + aliasing_unwrapped(v, dep_mod) end ainfo = mpi_remap_ainfo(ainfo, handle.rank) @opcounter :aliasing_bcast_send_yield @@ -2104,7 +2121,11 @@ function mpi_propagate_chunk_types!(tasks, accel::MPIAcceleration, expected_type for t in tasks if t isa Thunk if t.options !== nothing - t.options.return_type = expected_type + # Respect a more specific tile type already set by the spawner + # (e.g. `DSparseArray` from sparse allocators). + if t.options.return_type === nothing + t.options.return_type = expected_type + end else t.options = Options(return_type=expected_type) end diff --git a/ext/MetalExt.jl b/ext/MetalExt.jl index 7e175993c..f0b6efda0 100644 --- a/ext/MetalExt.jl +++ b/ext/MetalExt.jl @@ -432,13 +432,9 @@ Adapt.adapt_structure(to::Metal.Adaptor, H::Dagger.HaloArray) = H.halo_width; own_center=H.own_center) function Dagger.inner_stencil_proc!(::MtlArrayDeviceProc, f, output, read_vars) - Dagger.Kernel(_inner_stencil!)(f, output, read_vars; ndrange=size(output)) + Dagger.gpu_stencil_sweep!(f, output, read_vars) return end -@kernel function _inner_stencil!(f, output, read_vars) - idx = @index(Global, Cartesian) - f(idx, output, read_vars) -end function Base.show(io::IO, proc::MtlArrayDeviceProc) print(io, "MtlArrayDeviceProc(worker $(proc.owner), device $(something(_get_metal_device(proc)).name))") diff --git a/ext/MetalSparseArraysExt.jl b/ext/MetalSparseArraysExt.jl new file mode 100644 index 000000000..0d47f97b2 --- /dev/null +++ b/ext/MetalSparseArraysExt.jl @@ -0,0 +1,68 @@ +module MetalSparseArraysExt + +# Sparse DArray support for Metal via DeviceSparseMatrixCSC (no vendor sparse +# library). SpGEMM/SpMV fall back to host SparseArrays. + +import Dagger +import Metal +import Adapt +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import Metal: MtlArray + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +const MetalExt = Base.get_extension(Dagger, :MetalExt)::Module +using .MetalExt: MtlArrayDeviceProc +_with_context(f, proc) = MetalExt.with_context(f, proc) + +Adapt.adapt_structure(::Type{<:MtlArray}, S::SparseMatrixCSC) = + Dagger.device_sparse_from_host(MtlArray, S) + +Dagger.allocate_sparse_zeros(::MtlArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + Dagger.device_sparse_from_host(MtlArray, SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_zeros(::MtlArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + SparseArrays.spzeros(T, dims...) +Dagger.allocate_sparse_rand(::MtlArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + Dagger.device_sparse_from_host(MtlArray, SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::MtlArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) + +function Dagger.move(from_proc::CPUProc, to_proc::MtlArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(Dagger.device_sparse_from_host(MtlArray, x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::MtlArrayDeviceProc, x::SparseVector) + return Dagger.DSparseArray(copy(x)) +end +function Dagger.move(from_proc::CPUProc, to_proc::MtlArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + mat = x.mat + S = mat isa SparseMatrixCSC ? mat : SparseMatrixCSC(mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(MtlArray, S)) + end +end +function Dagger.move(from_proc::MtlArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + Metal.synchronize() + mat = x.mat + if mat isa Dagger.DeviceSparseMatrixCSC + return Dagger.DSparseArray(SparseMatrixCSC(mat)) + else + return Dagger.DSparseArray(copy(mat)) + end + end +end +function Dagger.move(from_proc::MtlArrayDeviceProc, to_proc::MtlArrayDeviceProc, x::Dagger.DSparseArray) + if from_proc == to_proc + return x + end + _with_context(to_proc) do + S = x.mat isa Dagger.DeviceSparseMatrixCSC ? SparseMatrixCSC(x.mat) : SparseMatrixCSC(x.mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(MtlArray, S)) + end +end + +end # module MetalSparseArraysExt diff --git a/ext/MetisExt.jl b/ext/MetisExt.jl new file mode 100644 index 000000000..9f4a2623b --- /dev/null +++ b/ext/MetisExt.jl @@ -0,0 +1,75 @@ +module MetisExt + +import Metis +import SparseArrays +import SparseArrays: SparseMatrixCSC, nnz, rowvals, nzrange +import LinearAlgebra +import Dagger + +# Single-level (non-recursive) METIS k-way partition + vertex-separator +# construction for Stage 4c Schur-complement domain decomposition. +# Recursive / nested dissection is future work. + +# Symmetric adjacency pattern of `A + Aᵀ` (structure only, unit weights). +function _symmetrize_pattern(A::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti} + n = size(A, 1) + S = A + transpose(A) + return SparseMatrixCSC{Tv,Ti}(n, n, S.colptr, S.rowval, ones(Tv, nnz(S))) +end + +""" + Dagger._nested_dissection_partition(A, nparts) -> (; perm, interiors, separator) + +K-way METIS partition of the undirected adjacency of `A + Aᵀ`, then a vertex +separator `Γ` = every vertex with a neighbor in a different part. Interiors are +the remaining vertices grouped by part (empty parts dropped). The permutation is +`[I₁; …; Iₖ; Γ]`. + +For tiny graphs (or `nparts == 1`), returns a trivial partition: one interior +covering `1:n` and an empty separator (caller should fall back to a serial +factorization). +""" +function Dagger._nested_dissection_partition(A::SparseMatrixCSC, nparts::Integer) + n = LinearAlgebra.checksquare(A) + nparts = Int(nparts) + nparts >= 1 || throw(ArgumentError("nparts must be ≥ 1, got $nparts")) + + # Too small to partition usefully → trivial (serial) partition. + if nparts == 1 || n < 4 || n <= nparts + return (; perm=collect(1:n), interiors=[collect(1:n)], separator=Int[]) + end + + pattern = _symmetrize_pattern(A) + g = Metis.graph(pattern) + part = Int.(Metis.partition(g, nparts)) + + is_sep = falses(n) + rows = rowvals(pattern) + @inbounds for v in 1:n + for k in nzrange(pattern, v) + u = rows[k] + u == v && continue + if part[u] != part[v] + is_sep[v] = true + break + end + end + end + + separator = findall(is_sep) + interiors = Vector{Vector{Int}}() + for p in 1:nparts + Ii = findall(i -> !is_sep[i] && part[i] == p, 1:n) + isempty(Ii) || push!(interiors, Ii) + end + + # Everything landed in the separator (or no interiors) → serial fallback. + if isempty(interiors) + return (; perm=collect(1:n), interiors=[collect(1:n)], separator=Int[]) + end + + perm = vcat(interiors..., separator) + return (; perm, interiors, separator) +end + +end # module MetisExt diff --git a/ext/OpenCLExt.jl b/ext/OpenCLExt.jl index eceb9cf6f..a61c0b310 100644 --- a/ext/OpenCLExt.jl +++ b/ext/OpenCLExt.jl @@ -361,13 +361,9 @@ Adapt.adapt_structure(to::OpenCL.KernelAdaptor, H::Dagger.HaloArray) = H.halo_width; own_center=H.own_center) function Dagger.inner_stencil_proc!(::CLArrayDeviceProc, f, output, read_vars) - Dagger.Kernel(_inner_stencil!)(f, output, read_vars; ndrange=size(output)) + Dagger.gpu_stencil_sweep!(f, output, read_vars) return end -@kernel function _inner_stencil!(f, output, read_vars) - idx = @index(Global, Cartesian) - f(idx, output, read_vars) -end Dagger.gpu_processor(::Val{:OpenCL}) = CLArrayDeviceProc Dagger.gpu_can_compute(::Val{:OpenCL}) = length(cl.platforms()) > 0 diff --git a/ext/OpenCLSparseArraysExt.jl b/ext/OpenCLSparseArraysExt.jl new file mode 100644 index 000000000..2a2f52772 --- /dev/null +++ b/ext/OpenCLSparseArraysExt.jl @@ -0,0 +1,82 @@ +module OpenCLSparseArraysExt + +# Sparse DArray support for OpenCL via DeviceSparseMatrixCSC (no vendor sparse +# library). SpGEMM/SpMV fall back to host SparseArrays. + +import Dagger +import OpenCL +import Adapt +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import OpenCL: CLArray + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +const OpenCLExt = Base.get_extension(Dagger, :OpenCLExt)::Module +using .OpenCLExt: CLArrayDeviceProc +_with_context(f, proc) = OpenCLExt.with_context(f, proc) + +# Do not densify via Adapt (generic `adapt(CLArray, SparseMatrixCSC)` would). +Adapt.adapt_structure(::Type{<:CLArray}, S::SparseMatrixCSC) = + Dagger.device_sparse_from_host(CLArray, S) + +_empty_device_sparse(::Type{T}, dims::Dims{2}) where T = + Dagger.device_sparse_from_host(CLArray, SparseArrays.spzeros(T, dims...)) +_empty_device_sparse(::Type{T}, dims::Dims{1}) where T = + # Store 1-D sparse as a 1-column DeviceSparseMatrixCSC + Dagger.device_sparse_from_host(CLArray, SparseMatrixCSC(SparseArrays.spzeros(T, dims...))) + +Dagger.allocate_sparse_zeros(::CLArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + _empty_device_sparse(T, dims) +Dagger.allocate_sparse_zeros(::CLArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + SparseArrays.spzeros(T, dims...) # vectors stay host CSC; wrap still applies +Dagger.allocate_sparse_rand(::CLArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + Dagger.device_sparse_from_host(CLArray, SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::CLArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) + +function Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x::SparseVector) + # Keep SparseVector on host (wrapped); OpenCL SpMV host-falls-back anyway. + return Dagger.DSparseArray(copy(x)) +end +function Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + mat = x.mat + if mat isa SparseMatrixCSC + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, mat)) + elseif mat isa Dagger.DeviceSparseMatrixCSC + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, SparseMatrixCSC(mat))) + else + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, SparseMatrixCSC(mat))) + end + end +end +function Dagger.move(from_proc::CLArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + OpenCL.cl.finish(OpenCL.cl.queue()) + mat = x.mat + if mat isa Dagger.DeviceSparseMatrixCSC + return Dagger.DSparseArray(SparseMatrixCSC(mat)) + else + return Dagger.DSparseArray(copy(mat)) + end + end +end +function Dagger.move(from_proc::CLArrayDeviceProc, to_proc::CLArrayDeviceProc, x::Dagger.DSparseArray) + if from_proc == to_proc + return x + end + _with_context(to_proc) do + mat = x.mat + S = mat isa Dagger.DeviceSparseMatrixCSC ? SparseMatrixCSC(mat) : SparseMatrixCSC(mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, S)) + end +end + +end # module OpenCLSparseArraysExt diff --git a/ext/PureKLUExt.jl b/ext/PureKLUExt.jl new file mode 100644 index 000000000..c2bbf22f5 --- /dev/null +++ b/ext/PureKLUExt.jl @@ -0,0 +1,38 @@ +module PureKLUExt + +import PureKLU +import SparseArrays +import SparseArrays: SparseMatrixCSC +import Dagger +import Dagger: DMatrix +import LinearAlgebra + +# Sparse direct (KLU) integration. KLU factorizations are pure-Julia data +# (serializable/movable), so they back both a whole-matrix direct solve +# (`Dagger.klu`) and per-tile block direct solves (`BlockKLUPreconditioner`), +# the latter via Dagger's block-preconditioner machinery. + +_as_sparse(A::SparseMatrixCSC) = A +_as_sparse(A::AbstractMatrix) = SparseArrays.sparse(A) + +# Whole-matrix direct solve: gather+factor on one worker, pin the factor there. +function Dagger.klu(A::DMatrix; kwargs...) + return Dagger._spawn_direct_factorization(A, S -> PureKLU.klu(S; kwargs...)) +end + +# Per-tile block direct preconditioner. +function Dagger.BlockKLUPreconditioner(A::DMatrix; kwargs...) + build = tile -> PureKLU.klu(_as_sparse(Dagger._tile_matrix(tile)); kwargs...) + return Dagger._build_block_preconditioner(Dagger.BlockKLUPreconditioner, A, + "BlockKLUPreconditioner", build) +end + +# KLU supports in-place `ldiv!`, so apply each block solve without allocating. +Dagger._block_apply!(y, F::PureKLU.KLUFactorization, x) = + (LinearAlgebra.ldiv!(y, F, x); return nothing) + +# The factorization is read-only and pinned to its worker; treat it as +# non-aliasing rather than recursing into its internals. +Dagger.aliasing(::PureKLU.KLUFactorization) = Dagger.NoAliasing() + +end # module PureKLUExt diff --git a/ext/PureUMFPACKExt.jl b/ext/PureUMFPACKExt.jl new file mode 100644 index 000000000..1817a5010 --- /dev/null +++ b/ext/PureUMFPACKExt.jl @@ -0,0 +1,63 @@ +module PureUMFPACKExt + +import PureUMFPACK +import SparseArrays +import SparseArrays: SparseMatrixCSC +import Dagger +import Dagger: DMatrix +import LinearAlgebra + +# Sparse direct (UMFPACK-style, pure-Julia) integration. Like KLU, these +# factorizations are plain Julia data (movable), backing both a whole-matrix +# direct solve (`Dagger.splu`) and per-tile block direct solves +# (`BlockUMFPACKPreconditioner`). With `distributed=true`: +# - `method=:trsv` (default) — Stage 4b tiled triangular solves +# - `method=:schur` — Stage 4c Schur-complement domain decomposition + +_as_sparse(A::SparseMatrixCSC) = A +_as_sparse(A::AbstractMatrix) = SparseArrays.sparse(A) + +function _umfpack_factorize(; kwargs...) + return S -> PureUMFPACK.splu(_as_sparse(S); kwargs...) +end + +# Whole-matrix direct solve: gather+factor on one worker, pin the factor there. +# Opt-in distributed paths select Stage 4b (`:trsv`) or Stage 4c (`:schur`). +function Dagger.splu(A::DMatrix; distributed::Bool=false, method::Symbol=:trsv, + blocksize=nothing, nparts=nothing, kwargs...) + factorize = _umfpack_factorize(; kwargs...) + if !distributed + return Dagger._spawn_direct_factorization(A, factorize) + end + if method === :trsv + F = Dagger._spawn_direct_factorization(A, factorize) + return Dagger._make_distributed_sparse_lu(F, A; blocksize) + elseif method === :schur + return Dagger._make_schur_sparse_lu(A, factorize; nparts) + else + throw(ArgumentError( + "Dagger.splu distributed method must be :trsv or :schur, got $(repr(method))")) + end +end + +# `(Rs .* A)[p, q] == L * U` with unit-lower `L` (explicit stored ones on the +# diagonal; solve treats it as unit diagonal via `UnitLowerTriangular`). +function Dagger._extract_lu_factors(F::PureUMFPACK.PureLU) + return (F.L, F.U, copy(F.p), copy(F.q), copy(F.Rs)) +end + +# Per-tile block direct preconditioner. +function Dagger.BlockUMFPACKPreconditioner(A::DMatrix; kwargs...) + build = tile -> PureUMFPACK.splu(_as_sparse(Dagger._tile_matrix(tile)); kwargs...) + return Dagger._build_block_preconditioner(Dagger.BlockUMFPACKPreconditioner, A, + "BlockUMFPACKPreconditioner", build) +end + +# `PureLU` provides `\` (no `ldiv!`); the default `_block_apply!` (`y .= op \ x`) +# already uses it, so no apply override is needed. + +# The factorization is read-only and pinned to its worker; treat it as +# non-aliasing rather than recursing into its internals. +Dagger.aliasing(::PureUMFPACK.PureLU) = Dagger.NoAliasing() + +end # module PureUMFPACKExt diff --git a/ext/ROCExt.jl b/ext/ROCExt.jl index 70d836369..ddd85d203 100644 --- a/ext/ROCExt.jl +++ b/ext/ROCExt.jl @@ -105,10 +105,42 @@ function with_context!(space::ROCVRAMMemorySpace) end Dagger.with_context!(proc::ROCArrayDeviceProc) = with_context!(proc) Dagger.with_context!(space::ROCVRAMMemorySpace) = with_context!(space) + +""" + task_stream_slot() + restore_stream_slot!(old_stream) + +Save and restore the running task's default HIP stream. + +`AMDGPU.stream()` creates a stream when the current task doesn't have one yet, +so using it to record what to restore would allocate a HIP stream for every +Dagger task that touches a ROC processor. Those streams live until they are +finalized, and a run goes through enough tasks that stream creation itself +eventually stalls inside the driver. Reading the task-local slot directly +reports "no stream" as `nothing` instead of manufacturing one. +""" +function task_stream_slot() + state = AMDGPU.task_local_state() + state === nothing && return nothing + return state.streams[AMDGPU.device_id(state.device)] +end +function restore_stream_slot!(old_stream) + if old_stream !== nothing + stream!(old_stream) + return + end + # The task had no stream of its own; put the slot back the way we found it + # rather than leaving it pointing at Dagger's per-device stream. + state = AMDGPU.task_local_state() + state === nothing && return + state.streams[AMDGPU.device_id(state.device)] = nothing + return +end + function with_context(f, x) old_ctx = context() old_device = AMDGPU.device() - old_stream = stream() + old_stream = task_stream_slot() with_context!(x) try @@ -116,7 +148,7 @@ function with_context(f, x) finally context!(old_ctx) AMDGPU.device!(old_device) - stream!(old_stream) + restore_stream_slot!(old_stream) end end @@ -387,13 +419,9 @@ Adapt.adapt_structure(to::AMDGPU.Runtime.Adaptor, H::Dagger.HaloArray) = H.halo_width; own_center=H.own_center) function Dagger.inner_stencil_proc!(::ROCArrayDeviceProc, f, output, read_vars) - Dagger.Kernel(_inner_stencil!)(f, output, read_vars; ndrange=size(output)) + Dagger.gpu_stencil_sweep!(f, output, read_vars) return end -@kernel function _inner_stencil!(f, output, read_vars) - idx = @index(Global, Cartesian) - f(idx, output, read_vars) -end Dagger.gpu_processor(::Val{:ROC}) = ROCArrayDeviceProc Dagger.gpu_can_compute(::Val{:ROC}) = AMDGPU.functional() diff --git a/ext/ROCSparseArraysExt.jl b/ext/ROCSparseArraysExt.jl new file mode 100644 index 000000000..7340109d5 --- /dev/null +++ b/ext/ROCSparseArraysExt.jl @@ -0,0 +1,185 @@ +module ROCSparseArraysExt + +# Sparse DArray support for AMDGPU / rocSPARSE. Loaded when both AMDGPU and +# SparseArrays are available (see Project.toml combo extension). + +import Dagger +import AMDGPU +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import AMDGPU: ROCArray +import AMDGPU.rocSPARSE: ROCSparseMatrixCSC, ROCSparseMatrixCSR, ROCSparseVector + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +# ROCExt is triggered by AMDGPU and provides the processor + context helpers. +const ROCExt = Base.get_extension(Dagger, :ROCExt)::Module +using .ROCExt: ROCArrayDeviceProc +_with_context(f, proc) = ROCExt.with_context(f, proc) + +#----- Memory / aliasing ------------------------------------------------------- + +Dagger.value_memory_space(x::ROCSparseMatrixCSC) = Dagger.memory_space(x.nzVal) +Dagger.value_memory_space(x::ROCSparseMatrixCSR) = Dagger.memory_space(x.nzVal) +Dagger.value_memory_space(x::ROCSparseVector) = Dagger.memory_space(x.nzVal) +# Datadeps `aliased_object!` keys slots by `memory_space(x)` (see DSparseArray). +Dagger.memory_space(x::ROCSparseMatrixCSC) = Dagger.value_memory_space(x) +Dagger.memory_space(x::ROCSparseMatrixCSR) = Dagger.value_memory_space(x) +Dagger.memory_space(x::ROCSparseVector) = Dagger.value_memory_space(x) + +function Dagger.aliasing(x::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}, _=identity) + space = Dagger.value_memory_space(x) + ptr = Dagger.RemotePtr{Cvoid}(UInt(pointer_from_objref(x)), space) + return Dagger.ObjectAliasing(ptr, sizeof(typeof(x))) +end +Dagger.aliases_as_whole(::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}) = true +Dagger.maybe_wrap_tile(x::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}) = + Dagger.DSparseArray(x) + +#----- Allocation -------------------------------------------------------------- + +Dagger.allocate_sparse_zeros(::ROCArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + ROCSparseMatrixCSC(SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_zeros(::ROCArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + ROCSparseVector(SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_rand(::ROCArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + ROCSparseMatrixCSC(SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::ROCArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + ROCSparseVector(SparseArrays.sprand(T, dims..., sparsity)) + +#----- Similar / collect ------------------------------------------------------- + +Dagger._sparse_similar(::ROCSparseMatrixCSC, ::Type{T}, dims::Dims{2}) where T = + ROCSparseMatrixCSC(SparseArrays.spzeros(T, dims...)) +Dagger._sparse_similar(::ROCSparseVector, ::Type{T}, dims::Dims{1}) where T = + ROCSparseVector(SparseArrays.spzeros(T, dims...)) +Dagger._sparse_collect(A::ROCSparseMatrixCSC) = SparseMatrixCSC(A) +Dagger._sparse_collect(A::ROCSparseMatrixCSR) = SparseMatrixCSC(ROCSparseMatrixCSC(A)) +Dagger._sparse_collect(A::ROCSparseVector) = SparseVector(A) +Dagger._sparse_copy(A::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}) = copy(A) + +#----- Host ↔ device helpers --------------------------------------------------- + +_to_roc_sparse(x::SparseMatrixCSC) = ROCSparseMatrixCSC(x) +_to_roc_sparse(x::SparseVector) = ROCSparseVector(x) +_to_roc_sparse(x::ROCSparseMatrixCSC) = x +_to_roc_sparse(x::ROCSparseMatrixCSR) = ROCSparseMatrixCSC(x) +_to_roc_sparse(x::ROCSparseVector) = x +_to_roc_sparse(x::Dagger.DeviceSparseMatrixCSC) = ROCSparseMatrixCSC(SparseMatrixCSC(x)) +_to_roc_dsparse(x::Dagger.DSparseArray) = Dagger.DSparseArray(_to_roc_sparse(x.mat)) + +_to_host_sparse(x::ROCSparseMatrixCSC) = SparseMatrixCSC(x) +_to_host_sparse(x::ROCSparseMatrixCSR) = SparseMatrixCSC(ROCSparseMatrixCSC(x)) +_to_host_sparse(x::ROCSparseVector) = SparseVector(x) +_to_host_sparse(x::SparseMatrixCSC) = x +_to_host_sparse(x::SparseVector) = x +_to_host_dsparse(x::Dagger.DSparseArray) = Dagger.DSparseArray(_to_host_sparse(x.mat)) + +_to_csc(A::ROCSparseMatrixCSC) = A +_to_csc(A::ROCSparseMatrixCSR) = ROCSparseMatrixCSC(A) +_to_csc(A) = ROCSparseMatrixCSC(A) + +# Materialize transpose/adjoint as CSC so SpGEMM stays on rocSPARSE / host CSC +# paths; `A * transpose(B)` would otherwise fall into scalar-indexing generic mul. +function _op_csc(X, t::Char) + Xc = _to_csc(X) + t == 'N' && return Xc + Sh = SparseMatrixCSC(Xc) + t == 'T' && return ROCSparseMatrixCSC(SparseArrays.sparse(transpose(Sh))) + t == 'C' && return ROCSparseMatrixCSC(SparseArrays.sparse(adjoint(Sh))) + throw(ArgumentError("Invalid trans char: $t")) +end + +#----- Move (preserve sparsity; wrap in DSparseArray) -------------------------- + +function Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(ROCSparseMatrixCSC(x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::SparseVector) + _with_context(to_proc) do + return Dagger.DSparseArray(ROCSparseVector(x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + return _to_roc_dsparse(x) + end +end +function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + AMDGPU.synchronize() + return _to_host_dsparse(x) + end +end +function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::CPUProc, + x::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}) + _with_context(from_proc) do + AMDGPU.synchronize() + return Dagger.DSparseArray(_to_host_sparse(x)) + end +end +function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::ROCArrayDeviceProc, x::Dagger.DSparseArray) + # Same device: identity (like dense ROCArray). A copy here would discard + # in-place SpGEMM writes under Datadeps/Sch argument moves. + if from_proc == to_proc + _with_context(AMDGPU.synchronize, from_proc) + return x + end + # Distinct devices: stage through host (no peer-copy helper yet). + _with_context(to_proc) do + return _to_roc_dsparse(_to_host_dsparse(x)) + end +end + +#----- Tile kernels ------------------------------------------------------------ + +function Dagger.matmatmul!( + C::Dagger.DSparseMatrix, + transA::Char, transB::Char, + A::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR}, + B::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR}, + alpha, beta +) + AB = _to_csc(_op_csc(A, transA) * _op_csc(B, transB)) + prod = _to_csc(isone(alpha) ? AB : alpha * AB) + if iszero(beta) + C.mat = prod + elseif isone(beta) + C.mat = _to_csc(prod + _to_csc(C.mat)) + else + C.mat = _to_csc(prod + beta * _to_csc(C.mat)) + end + return C +end + +_apply_trans(X, t::Char) = + t == 'N' ? X : t == 'T' ? transpose(X) : t == 'C' ? adjoint(X) : + throw(ArgumentError("Invalid trans char: $t")) + +function Dagger.matvecmul!(C::ROCArray, transA::Char, + A::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR}, + B::ROCArray, alpha, beta) + # rocSPARSE SpMV supports transpose/adjoint wrappers. + LinearAlgebra.mul!(C, _apply_trans(A, transA), B, alpha, beta) + return C +end + +function Dagger.transpose_tile(B::ROCSparseMatrixCSC) + return ROCSparseMatrixCSC(SparseArrays.sparse(SparseMatrixCSC(B)')) +end +function Dagger.transpose_tile(B::ROCSparseMatrixCSC, uplo::Char) + Bh = SparseMatrixCSC(B) + Bt = uplo == 'U' ? SparseArrays.triu(Bh) : + uplo == 'L' ? SparseArrays.tril(Bh) : + throw(ArgumentError("uplo must be 'U' or 'L', got $uplo")) + Ct = Bt + Bt' + for i in 1:LinearAlgebra.checksquare(Bh) + Ct[i, i] = Bh[i, i] + end + return ROCSparseMatrixCSC(Ct) +end + +end # module ROCSparseArraysExt diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl new file mode 100644 index 000000000..757660d1c --- /dev/null +++ b/ext/SparseArraysExt.jl @@ -0,0 +1,236 @@ +module SparseArraysExt + +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import Dagger +import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseArray, DSparseMatrix + +# Keep tiles sparse through `collect`/`cat`; the outer `collect` densifies. +Dagger._sparse_collect(M::SparseMatrixCSC) = copy(M) + +# Assemble already-local tiles into one global `SparseMatrixCSC` without +# densifying: unwrap each tile, offset its (i,j) by the precomputed subdomain +# offsets, and build from triplets. Intended to run inside a worker-scoped task +# (the scheduler moves the tile chunks there); used by `Dagger.klu` / `Dagger.splu`. +function Dagger._gather_sparse(::Type{T}, tiles, row_offsets, col_offsets, m, n) where T + Is = Int[]; Js = Int[]; Vs = T[] + for k in 1:length(tiles) + tile = SparseMatrixCSC(Dagger._tile_matrix(tiles[k])) + ti, tj, tv = SparseArrays.findnz(tile) + append!(Is, ti .+ row_offsets[k]) + append!(Js, tj .+ col_offsets[k]) + append!(Vs, tv) + end + return SparseArrays.sparse(Is, Js, Vs, m, n) +end + +# Dense → sparse for Stage-4c Schur complements (fill-in is expected). +Dagger._sparse_copy_of(S::AbstractMatrix) = SparseArrays.sparse(S) +Dagger._sparse_copy_of(S::SparseMatrixCSC) = S + +# Wrap bare sparse tiles (e.g. from `distribute`) so Datadeps sees a stable container. +Dagger.maybe_wrap_tile(x::SparseMatrixCSC) = DSparseArray(x) +Dagger.maybe_wrap_tile(x::SparseVector) = DSparseArray(x) +Dagger.maybe_wrap_tile(x::Dagger.DeviceSparseMatrixCSC) = DSparseArray(x) + +# Host defaults for scoped sparse allocation (GPU Exts override per-processor). +Dagger.allocate_sparse_zeros_default(::Dagger.Processor, ::Type{T}, dims::Dims{2}) where T = + SparseArrays.spzeros(T, dims...) +Dagger.allocate_sparse_zeros_default(::Dagger.Processor, ::Type{T}, dims::Dims{1}) where T = + SparseArrays.spzeros(T, dims...) +Dagger.allocate_sparse_rand_default(::Dagger.Processor, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) +Dagger.allocate_sparse_rand_default(::Dagger.Processor, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) + +# DeviceSparseMatrixCSC ↔ SparseMatrixCSC +function SparseArrays.SparseMatrixCSC(A::Dagger.DeviceSparseMatrixCSC{Tv,Ti}) where {Tv,Ti} + return SparseMatrixCSC{Tv,Ti}(A.m, A.n, Array(A.colptr), Array(A.rowval), Array(A.nzval)) +end +""" + device_sparse_from_host(Arr, S::SparseMatrixCSC) -> DeviceSparseMatrixCSC + +Upload a host CSC onto device vectors of type `Arr` (e.g. `CLArray`, `MtlArray`, +`oneArray`). Indices are converted to `Int32` for device friendliness. +""" +function Dagger.device_sparse_from_host(::Type{Arr}, S::SparseMatrixCSC{Tv}) where {Tv,Arr} + colptr = Arr(Int32.(S.colptr)) + rowval = Arr(Int32.(S.rowval)) + nzval = Arr(Array(S.nzval)) + return Dagger.DeviceSparseMatrixCSC(S.m, S.n, colptr, rowval, nzval) +end +Base.copy(A::Dagger.DeviceSparseMatrixCSC) = + Dagger.DeviceSparseMatrixCSC(A.m, A.n, copy(A.colptr), copy(A.rowval), copy(A.nzval)) +Dagger._sparse_copy(A::Dagger.DeviceSparseMatrixCSC) = copy(A) +Dagger._sparse_collect(A::Dagger.DeviceSparseMatrixCSC) = SparseMatrixCSC(A) +function Dagger._sparse_similar(A::Dagger.DeviceSparseMatrixCSC{Tv,Ti}, ::Type{T}, dims::Dims{2}) where {Tv,Ti,T} + n = dims[2] + colptr = similar(A.colptr, Ti, n + 1) + fill!(colptr, one(Ti)) + rowval = similar(A.rowval, Ti, 0) + nzval = similar(A.nzval, T, 0) + return Dagger.DeviceSparseMatrixCSC(dims[1], n, colptr, rowval, nzval) +end + +# Rebuild a DeviceSparseMatrixCSC on the same device array type as `like`. +function _to_device_sparse(like::Dagger.DeviceSparseMatrixCSC, S::SparseMatrixCSC) + colptr = similar(like.colptr, eltype(like.colptr), length(S.colptr)) + rowval = similar(like.rowval, eltype(like.rowval), length(S.rowval)) + nzval = similar(like.nzval, eltype(S), length(S.nzval)) + copyto!(colptr, eltype(colptr).(S.colptr)) + copyto!(rowval, eltype(rowval).(S.rowval)) + copyto!(nzval, S.nzval) + return Dagger.DeviceSparseMatrixCSC(S.m, S.n, colptr, rowval, nzval) +end + +function SparseArrays.spzeros(p::BlocksOrAuto, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Dagger.Tapes.resolve_partitioning(T, dims, p, assignment) + d = Dagger.ArrayDomain(map(x->1:x, dims)) + N = length(dims) + # Route through `allocate_sparse_zeros` so a GPU compute scope yields + # device-resident sparse tiles (vendor sparse or DeviceSparseMatrixCSC). + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Dagger.allocate_sparse_zeros(Dagger.task_processor(), T, _dims)), false, d, Dagger.partition(part, d), part, assign; + return_type=DSparseArray{T,N}) + return Dagger.Tapes.track!(Dagger._to_darray(a), plan) +end +SparseArrays.spzeros(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = + SparseArrays.spzeros(p, T, dims; assignment) +SparseArrays.spzeros(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbitrary) = + SparseArrays.spzeros(p, Float64, dims; assignment) +SparseArrays.spzeros(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = + SparseArrays.spzeros(p, Float64, dims; assignment) + +function SparseArrays.sprand(p::BlocksOrAuto, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) + part, assign, plan = Dagger.Tapes.resolve_partitioning(T, dims, p, assignment) + d = Dagger.ArrayDomain(map(x->1:x, dims)) + N = length(dims) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Dagger.allocate_sparse_rand(Dagger.task_processor(), T, _dims, sparsity)), false, d, Dagger.partition(part, d), part, assign; + return_type=DSparseArray{T,N}) + return Dagger.Tapes.track!(Dagger._to_darray(a), plan) +end +SparseArrays.sprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = + SparseArrays.sprand(p, T, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) +SparseArrays.sprand(p::BlocksOrAuto, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = + SparseArrays.sprand(p, Float64, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) +SparseArrays.sprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = + SparseArrays.sprand(p, Float64, dims, sparsity; assignment) + +_apply_trans(X, t::Char) = + t == 'N' ? X : + t == 'T' ? transpose(X) : + t == 'C' ? adjoint(X) : + throw(ArgumentError("Invalid trans char: $t")) + +function _sparse_gemm_assign!(C::DSparseMatrix, prod, beta) + if iszero(beta) + C.mat = prod + elseif isone(beta) + C.mat = prod + C.mat + else + C.mat = prod + beta * C.mat + end + return C +end + +function Dagger.matmatmul!( + C::DSparseMatrix, + transA::Char, + transB::Char, + A::SparseMatrixCSC, + B::SparseMatrixCSC, + alpha, + beta +) + opA = _apply_trans(A, transA) + opB = _apply_trans(B, transB) + # Sparse*sparse yields a freshly-allocated sparse matrix, which we reassign + # into the wrapper (`DSparseMatrix` hides this reallocation from Datadeps). + # `SparseArrays` provides no efficient 5-arg `mul!` into a sparse `C` -- the + # output sparsity pattern is determined by the product -- so we form the + # product out-of-place and apply only the alpha/beta scaling that is actually + # needed. The transposed-operand products dispatch to specialized SparseArrays + # methods, so `opA`/`opB` are not materialized. + AB = opA * opB + prod = isone(alpha) ? AB : alpha * AB + return _sparse_gemm_assign!(C, prod, beta) +end + +# DeviceSparseMatrixCSC SpGEMM: gather to host, multiply, scatter back. +function Dagger.matmatmul!( + C::DSparseMatrix, + transA::Char, + transB::Char, + A::Dagger.DeviceSparseMatrixCSC, + B::Dagger.DeviceSparseMatrixCSC, + alpha, + beta +) + Ah = SparseMatrixCSC(A) + Bh = SparseMatrixCSC(B) + Ch = C.mat isa Dagger.DeviceSparseMatrixCSC ? SparseMatrixCSC(C.mat) : + C.mat isa SparseMatrixCSC ? C.mat : SparseMatrixCSC(C.mat) + opA = _apply_trans(Ah, transA) + opB = _apply_trans(Bh, transB) + AB = opA * opB + prod = isone(alpha) ? AB : alpha * AB + if iszero(beta) + result = prod + elseif isone(beta) + result = prod + Ch + else + result = prod + beta * Ch + end + C.mat = _to_device_sparse(A, SparseMatrixCSC(result)) + return C +end + +# Sparse matrix-vector multiply tile kernel: `C = alpha*op(A)*B + beta*C` with a +# `SparseMatrixCSC` `A` and dense vectors `B`/`C`. SparseArrays provides an +# efficient 5-arg `mul!` (SpMV) into a dense output, including for transposed and +# adjoint operands, so this updates `C` in place with no allocation. +function Dagger.matvecmul!(C::AbstractVector, transA::Char, A::SparseMatrixCSC, B::AbstractVector, alpha, beta) + LinearAlgebra.mul!(C, _apply_trans(A, transA), B, alpha, beta) + return C +end + +# DeviceSparseMatrixCSC SpMV: host fallback (works for any dense vector type +# that supports Array(::)/copyto!). +function Dagger.matvecmul!(C::AbstractVector, transA::Char, A::Dagger.DeviceSparseMatrixCSC, B::AbstractVector, alpha, beta) + Ah = SparseMatrixCSC(A) + Bh = Array(B) + Ch = Array(C) + LinearAlgebra.mul!(Ch, _apply_trans(Ah, transA), Bh, alpha, beta) + copyto!(C, Ch) + return C +end + +# Off-diagonal tile copy in `copytri!`: produce the (conjugate) transpose tile. +function Dagger.transpose_tile(B::SparseMatrixCSC) + return SparseArrays.sparse(B') +end +function Dagger.transpose_tile(B::Dagger.DeviceSparseMatrixCSC) + return _to_device_sparse(B, SparseArrays.sparse(SparseMatrixCSC(B)')) +end +# Diagonal tile symmetrization in `copytri!`: build the full Hermitian tile from +# its `uplo` triangle (matching the dense `copydiagtile!` semantics). +function Dagger.transpose_tile(B::SparseMatrixCSC, uplo::Char) + if uplo == 'U' + Bt = SparseArrays.triu(B) + elseif uplo == 'L' + Bt = SparseArrays.tril(B) + else + throw(ArgumentError("uplo must be 'U' or 'L', got $uplo")) + end + C = Bt + Bt' + # The shared diagonal was added twice; restore the original tile's diagonal. + for i in 1:LinearAlgebra.checksquare(B) + C[i, i] = B[i, i] + end + return C +end +function Dagger.transpose_tile(B::Dagger.DeviceSparseMatrixCSC, uplo::Char) + return _to_device_sparse(B, Dagger.transpose_tile(SparseMatrixCSC(B), uplo)) +end + +end # module SparseArraysExt diff --git a/src/Dagger.jl b/src/Dagger.jl index 81881716e..059f0b465 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -2,7 +2,6 @@ module Dagger import Serialization import Serialization: AbstractSerializer, serialize, deserialize -import SparseArrays: sprand, SparseMatrixCSC import MemPool import MemPool: DRef, FileRef, poolget, poolset @@ -103,11 +102,13 @@ include("sch/Sch.jl"); using .Sch include("tochunk.jl") # Data dependency task queue +include("datadeps/types.jl") include("datadeps/aliasing.jl") include("datadeps/chunkview.jl") include("datadeps/remainders.jl") include("datadeps/scheduling.jl") include("datadeps/queue.jl") +include("datadeps/hierarchical.jl") # Stencils include("utils/haloarray.jl") @@ -124,6 +125,7 @@ include("file-io.jl") # Array computations include("procgrid.jl") include("array/darray.jl") +include("tapes/Tapes.jl"); using .Tapes include("array/alloc.jl") include("array/map-reduce.jl") include("array/copy.jl") @@ -132,7 +134,7 @@ include("array/operators.jl") include("array/indexing.jl") include("array/setindex.jl") include("array/matrix.jl") -include("array/sparse_partition.jl") +include("array/sparse.jl") include("array/sort.jl") include("array/permute.jl") include("array/linalg.jl") @@ -142,6 +144,8 @@ include("array/trsm.jl") include("array/lu.jl") include("array/qr.jl") include("array/svd.jl") +include("array/iterativesolvers.jl") +include("array/sparsedirect.jl") # GPU include("gpu.jl") @@ -239,6 +243,10 @@ function __init__() catch err @warn "Error parsing JULIA_DAGGER_DEBUG" exception=err end + + # Enable tapes + Tapes.clear!() + Tapes.enable!() end end # module diff --git a/src/acceleration.jl b/src/acceleration.jl index 99a71e1e2..ea411bb0a 100644 --- a/src/acceleration.jl +++ b/src/acceleration.jl @@ -114,6 +114,18 @@ processor_owner_rank(accel::Acceleration, proc::Processor) = first(fire_order_ke processor_order_key(accel::Acceleration, proc::Processor) = (fire_order_key(proc)..., short_name(proc)) +""" + partition_affinity_id(x) -> Int + +Logical owner identity used by hierarchical Datadeps partitioning. Defaults to +`root_worker_id` (Distributed worker id). Accelerations with a different notion +of ownership (e.g. MPI ranks) override for their spaces/processors so affinity +partitioning spreads work across ranks the same way Distributed spreads across +workers. +""" +partition_affinity_id(space::MemorySpace) = root_worker_id(space) +partition_affinity_id(proc::Processor) = partition_affinity_id(only(memory_spaces(proc))) + "Preferred execution rank for `task`, or `nothing` if no preference." task_processor_preference(accel::Acceleration, task, state) = nothing diff --git a/src/array/alloc.jl b/src/array/alloc.jl index 93862d545..ef926816e 100644 --- a/src/array/alloc.jl +++ b/src/array/alloc.jl @@ -10,13 +10,20 @@ mutable struct AllocateArray{T,N} <: ArrayOp{T,N} domainchunks partitioning::AbstractBlocks{N} procgrid::Union{AbstractProcGrid{N}, Nothing} - - function AllocateArray(eltype::Type{T}, f, want_index::Bool, d::ArrayDomain{N}, domainchunks, p::AbstractBlocks{N}, assignment::Union{AssignmentType{N},Nothing} = nothing) where {T,N} + # Optional concrete tile type for staged `@spawn`s. When `nothing`, leave + # `return_type` unset so inference (and MPI `post_stage_array_chunks!`, + # which defaults dense tiles to `Array{T,N}`) decide. Sparse allocators + # pass e.g. `DSparseArray{T,N}` so MPI does not overwrite with `Array`. + return_type::Union{Type,Nothing} + + function AllocateArray(eltype::Type{T}, f, want_index::Bool, d::ArrayDomain{N}, + domainchunks, p::AbstractBlocks{N}, + assignment::Union{AssignmentType{N},Nothing} = nothing; + return_type::Union{Type,Nothing} = nothing) where {T,N} sizeA = map(length, d.indexes) procgrid = build_procgrid(something(assignment, :arbitrary), Tuple(sizeA), p.blocksize, current_acceleration()) - return new{T,N}(eltype, f, want_index, d, domainchunks, p, procgrid) + return new{T,N}(eltype, f, want_index, d, domainchunks, p, procgrid, return_type) end - end size(a::AllocateArray) = size(a.domain) @@ -45,12 +52,22 @@ function stage(ctx, A::AllocateArray) tasks = emit_chunk_tasks!(A.domainchunks, A.procgrid, A.eltype, (scope, I, i) -> begin x = A.domainchunks[I] - N = ndims(A.domainchunks) - ret_type = Array{A.eltype, N} - if A.want_index - Dagger.@spawn compute_scope=scope return_type=ret_type allocate_array(A.f, A.eltype, i, size(x)) + if A.f isa DArray + # `similar(::DArray)` carries the source tile type (e.g. sparse) forward. + chunk = A.f.chunks[I] + Dagger.@spawn compute_scope=scope similar(chunk, A.eltype, size(x)) + elseif A.want_index + if A.return_type !== nothing + Dagger.@spawn compute_scope=scope return_type=A.return_type allocate_array(A.f, A.eltype, i, size(x)) + else + Dagger.@spawn compute_scope=scope allocate_array(A.f, A.eltype, i, size(x)) + end else - Dagger.@spawn compute_scope=scope return_type=ret_type allocate_array(A.f, A.eltype, size(x)) + if A.return_type !== nothing + Dagger.@spawn compute_scope=scope return_type=A.return_type allocate_array(A.f, A.eltype, size(x)) + else + Dagger.@spawn compute_scope=scope allocate_array(A.f, A.eltype, size(x)) + end end end) return DArray(A.eltype, A.domain, A.domainchunks, tasks, A.partitioning) @@ -58,10 +75,11 @@ end const BlocksOrAuto = Union{Blocks{N} where N, AutoBlocks} -function Base.rand(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) +function Base.rand(p::BlocksOrAuto, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Tapes.resolve_partitioning(T, dims, p, assignment) d = ArrayDomain(map(x->1:x, dims)) - a = AllocateArray(T, rand, false, d, partition(p, d), p, assignment) - return _to_darray(a) + a = AllocateArray(T, rand, false, d, partition(part, d), part, assign) + return Tapes.track!(_to_darray(a), plan) end Base.rand(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = rand(p, T, dims; assignment) @@ -69,13 +87,12 @@ Base.rand(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbit rand(p, Float64, dims; assignment) Base.rand(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = rand(p, Float64, dims; assignment) -Base.rand(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = - rand(auto_blocks(dims), T, dims; assignment) -function Base.randn(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) +function Base.randn(p::BlocksOrAuto, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Tapes.resolve_partitioning(T, dims, p, assignment) d = ArrayDomain(map(x->1:x, dims)) - a = AllocateArray(T, randn, false, d, partition(p, d), p, assignment) - return _to_darray(a) + a = AllocateArray(T, randn, false, d, partition(part, d), part, assign) + return Tapes.track!(_to_darray(a), plan) end Base.randn(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = randn(p, T, dims; assignment) @@ -83,27 +100,12 @@ Base.randn(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbi randn(p, Float64, dims; assignment) Base.randn(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = randn(p, Float64, dims; assignment) -Base.randn(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = - randn(auto_blocks(dims), T, dims; assignment) -function sprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) +function Base.ones(p::BlocksOrAuto, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Tapes.resolve_partitioning(T, dims, p, assignment) d = ArrayDomain(map(x->1:x, dims)) - a = AllocateArray(T, (T, _dims) -> sprand(T, _dims..., sparsity), false, d, partition(p, d), p, assignment) - return _to_darray(a) -end -sprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = - sprand(p, T, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) -sprand(p::BlocksOrAuto, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = - sprand(p, Float64, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) -sprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = - sprand(p, Float64, dims, sparsity; assignment) -sprand(::AutoBlocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = - sprand(auto_blocks(dims), T, dims, sparsity; assignment) - -function Base.ones(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) - d = ArrayDomain(map(x->1:x, dims)) - a = AllocateArray(T, ones, false, d, partition(p, d), p, assignment) - return _to_darray(a) + a = AllocateArray(T, ones, false, d, partition(part, d), part, assign) + return Tapes.track!(_to_darray(a), plan) end Base.ones(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = ones(p, T, dims; assignment) @@ -111,13 +113,12 @@ Base.ones(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbit ones(p, Float64, dims; assignment) Base.ones(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = ones(p, Float64, dims; assignment) -Base.ones(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = - ones(auto_blocks(dims), T, dims; assignment) -function Base.zeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) +function Base.zeros(p::BlocksOrAuto, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Tapes.resolve_partitioning(T, dims, p, assignment) d = ArrayDomain(map(x->1:x, dims)) - a = AllocateArray(T, zeros, false, d, partition(p, d), p, assignment) - return _to_darray(a) + a = AllocateArray(T, zeros, false, d, partition(part, d), part, assign) + return Tapes.track!(_to_darray(a), plan) end Base.zeros(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = zeros(p, T, dims; assignment) @@ -125,8 +126,6 @@ Base.zeros(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbi zeros(p, Float64, dims; assignment) Base.zeros(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = zeros(p, Float64, dims; assignment) -Base.zeros(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = - zeros(auto_blocks(dims), T, dims; assignment) function Base.zero(x::DArray{T,N}) where {T,N} dims = ntuple(i->x.domain.indexes[i].stop, N) @@ -140,12 +139,12 @@ end function LinearAlgebra._zeros(::Type{T}, B::DVector, n::Integer) where T m = max(size(B, 1), n) sz = (m,) - return zeros(auto_blocks(sz), T, sz) + return zeros(AutoBlocks(), T, sz) end function LinearAlgebra._zeros(::Type{T}, B::DMatrix, n::Integer) where T m = max(size(B, 1), n) sz = (m, size(B, 2)) - return zeros(auto_blocks(sz), T, sz) + return zeros(AutoBlocks(), T, sz) end function Base.view(A::AbstractArray{T,N}, p::Blocks{N}) where {T,N} @@ -167,3 +166,14 @@ function unsafe_free!(A::DArray) end end end + +# Initializers + +function Base.fill!(A::DArray, x) + spawn_datadeps() do + for chunk in A.chunks + Dagger.@spawn fill!(chunk, x) + end + end + return A +end diff --git a/src/array/cholesky.jl b/src/array/cholesky.jl index 45ba39185..da780714e 100644 --- a/src/array/cholesky.jl +++ b/src/array/cholesky.jl @@ -17,6 +17,7 @@ function potrf_checked!(uplo, A, info_arr) return _A, info end function LinearAlgebra._chol!(A::DArray{T,2}, ::Type{UpperTriangular}) where T + Dagger.@record_op :cholesky! A LinearAlgebra.checksquare(A) zone = one(T) @@ -62,6 +63,7 @@ function LinearAlgebra._chol!(A::DArray{T,2}, ::Type{UpperTriangular}) where T return UpperTriangular(A), info[1] end function LinearAlgebra._chol!(A::DArray{T,2}, ::Type{LowerTriangular}) where T + Dagger.@record_op :cholesky! A LinearAlgebra.checksquare(A) zone = one(T) diff --git a/src/array/copy.jl b/src/array/copy.jl index 7d92566e6..28d68b835 100644 --- a/src/array/copy.jl +++ b/src/array/copy.jl @@ -36,6 +36,13 @@ function allocate_copy_buffer(part::Blocks{N}, A::DArray{T,N}) where {T,N} return DArray{T}(undef, part, size(A)) end +to_range(x::UnitRange) = x +to_range(x::Integer) = x:x +to_range(x::Base.OneTo{Int}) = UnitRange(x) +to_range(x::Base.Slice{Base.OneTo{Int}}) = UnitRange(x) +to_range(::StepRange) = throw(ArgumentError("Cannot convert StepRange to UnitRange")) +to_range(x) = throw(ArgumentError("Cannot convert $(typeof(x)) to UnitRange")) + function darray_copyto!(B::DArray{TB,NB}, A::DArray{TA,NA}, Binds=parentindices(B), Ainds=parentindices(A)) where {TB,NB,TA,NA} Nmax = max(NA, NB) @@ -45,13 +52,6 @@ function darray_copyto!(B::DArray{TB,NB}, A::DArray{TA,NA}, Binds=parentindices( padNmax(x) = ntuple(i->pad1range(x, i), Nmax) padNmax(x::ArrayDomain) = padNmax(x.indexes) - to_range(x::UnitRange) = x - to_range(x::Integer) = x:x - to_range(x::Base.OneTo{Int}) = UnitRange(x) - to_range(x::Base.Slice{Base.OneTo{Int}}) = UnitRange(x) - to_range(::StepRange) = throw(ArgumentError("Non-continuous ranges are not yet supported for DArray copy")) - to_range(x) = throw(ArgumentError("Unsupported range type for DArray copy: $(typeof(x))")) - if any(x->x isa Vector, Binds) || any(x->x isa Vector, Ainds) # Split the copy into multiple copies dims_with_vector = findall(x->x[1] isa Vector || x[2] isa Vector, collect(zip(Binds, Ainds))) @@ -143,14 +143,30 @@ function copyto_view!(Bpart, Brange, Apart, Arange) return end -Base.copyto!(B::DArray{T,N}, A::DArray{T,N}) where {T,N} = - darray_copyto!(B, A) -Base.copyto!(B::DArray{T,N}, A::Array{T,N}) where {T,N} = - darray_copyto!(B, view(A, B.partitioning)) +function Base.copyto!(B::DArray{T,N}, A::DArray{T,N}) where {T,N} + Dagger.@record_op :copyto! B + return darray_copyto!(B, A) +end +function Base.copyto!(B::DArray{T,N}, A::Array{T,N}) where {T,N} + Dagger.@record_op :copyto! B + return darray_copyto!(B, view(A, B.partitioning)) +end Base.copyto!(B::Array{T,N}, A::DArray{T,N}) where {T,N} = darray_copyto!(view(B, A.partitioning), A) StridedDArray{T,N} = Union{<:DArray{T,N}, SubArray{T,N,<:DArray{T,NP}} where NP} -Base.copyto!(B::StridedDArray, A::StridedDArray) = - darray_copyto!(parent(B), parent(A), parentindices(B), parentindices(A)) +function Base.copyto!(B::StridedDArray, A::StridedDArray) + Dagger.@record_op :copyto! parent(B) + return darray_copyto!(parent(B), parent(A), parentindices(B), parentindices(A)) +end +function Base.copyto!(B::Array, A::StridedDArray) + DB = view(B, AutoBlocks()) + darray_copyto!(DB, parent(A), parentindices(DB), parentindices(A)) + return B +end +function Base.copyto!(B::SubArray, A::StridedDArray) + DB = view(parent(B), AutoBlocks()) + darray_copyto!(DB, parent(A), parentindices(B), parentindices(A)) + return B +end diff --git a/src/array/darray.jl b/src/array/darray.jl index bd0e5efc5..6467a2ea5 100644 --- a/src/array/darray.jl +++ b/src/array/darray.jl @@ -180,6 +180,12 @@ stage(ctx, c::DArray) = c # closures get non-uniform ArgumentWrapper hashes). _collect_cat(concat, dims::Int, xs...) = concat(xs...; dims) +# Finalize a gathered DArray to a dense `Array{T,N}`. Sparse tile cats may yield +# a `DSparseArray` / `SparseMatrixCSC`, and `Base.collect` does not densify those. +_collect_dense(::Type{T}, ::Val{N}, x::Array{T,N}) where {T,N} = x +_collect_dense(::Type{T}, ::Val{N}, x::AbstractArray) where {T,N} = Array{T,N}(x) +_collect_dense(::Type, ::Val, x) = x + function Base.collect(d::DArray{T,N}; tree=true, copyto=false) where {T,N} a = fetch(d) if isempty(d.chunks) @@ -208,13 +214,21 @@ function Base.collect(d::DArray{T,N}; tree=true, copyto=false) where {T,N} result = Dagger.spawn_datadeps() do treereduce_nd(spawn_catfuncs, a.chunks) end - return collect(fetch(result)) + return _collect_dense(T, Val(N), fetch(result)) end # Distributed: fetch chunks directly and concat in-process. This avoids # routing chunk data through datadeps aliasing, which requires an # aliasing-resolvable (e.g. isbits) element type. + # + # Sparse tiles (`DSparseArray`) are unwrapped to host storage before `cat`: + # GPU sparse types (CuSparse/ROCSparse/DeviceSparseMatrixCSC) disallow the + # scalar indexing that generic `cat` would use. dimcatfuncs = [(x...) -> concat(x..., dims=i) for i in 1:N] - return collect(treereduce_nd(dimcatfuncs, asyncmap(fetch, a.chunks))) + tiles = asyncmap(a.chunks) do c + x = fetch(c) + x isa DSparseArray ? _sparse_collect(x.mat) : x + end + return _collect_dense(T, Val(N), treereduce_nd(dimcatfuncs, tiles)) end Array{T,N}(A::DArray{S,N}) where {T,N,S} = convert(Array{T,N}, collect(A)) @@ -342,8 +356,16 @@ aliasing(x::DArray) = memory_space(x::DArray) = throw(ConcurrencyViolationError("DArray memory spaces may be mixed and unstable")) -Base.similar(D::DArray{T,N} where T, ::Type{S}, dims::Dims{N}) where {S,N} = - DArray{S,N}(undef, D.partitioning, dims) +function Base.similar(D::DArray{T,N} where T, ::Type{S}, dims::Dims{N}) where {S,N} + if dims == size(D) + domain = ArrayDomain(map(x->1:x, dims)) + subdomains = partition(D.partitioning, domain) + a = AllocateArray(S, D, false, domain, subdomains, D.partitioning, nothing) + return _to_darray(a) + else + return DArray{S,N}(undef, D.partitioning, dims) + end +end Base.similar(D::DArray{T,N1} where T, ::Type{S}, dims::Dims{N2}) where {S,N1,N2} = DArray{S,N2}(undef, auto_blocks(dims), dims) @@ -413,31 +435,21 @@ end """ Base.fetch(c::DArray) -If a `DArray` tree has a `Thunk` in it, make the whole thing a big thunk. +If a `DArray` tree has a `DTask` in it, `fetch` each one (as a `Chunk`, +without moving its data) and return a new `DArray` with all partitions +resolved to `Chunk`s. """ function Base.fetch(c::DArray{T}) where T - if any(istask, chunks(c)) - if uniform_execution() - # SPMD (MPI): every rank holds every task locally, and chunk - # records are rank-local views (placeholders for non-owned data). - # Resolve them locally instead of assembling on one rank and - # broadcasting (which would clobber the local records with the - # assembling rank's placeholders). - thunks = chunks(c) - new_chunks = Any[istask(t) ? fetch(t; raw=true) : t for t in thunks] - return DArray(T, domain(c), domainchunks(c), - reshape(new_chunks, size(thunks)), - c.partitioning, c.concat) - end - thunks = chunks(c) - sz = size(thunks) + thunks = chunks(c) + if any(istask, thunks) dmn = domain(c) dmnchunks = domainchunks(c) - return fetch(Dagger.spawn(Options(meta=true), thunks...) do results... - t = eltype(fetch(results[1])) - DArray(t, dmn, dmnchunks, reshape(Any[results...], sz), - c.partitioning, c.concat) - end) + new_chunks = similar(thunks, Any) + for idx in eachindex(thunks) + part = thunks[idx] + new_chunks[idx] = istask(part) ? fetch(part; raw=true) : part + end + return DArray(T, dmn, dmnchunks, new_chunks, c.partitioning, c.concat) else return c end @@ -503,7 +515,7 @@ function stage(ctx::Context, d::Distribute) cs = emit_chunk_tasks!(d.domainchunks, d.procgrid, T, (scope, I, i) -> begin c = d.domainchunks[I] - Dagger.@spawn compute_scope=scope identity(d.data[c]) + Dagger.@spawn compute_scope=scope maybe_wrap_tile(d.data[c]) end) end return DArray(eltype(d.data), @@ -534,7 +546,11 @@ function distribute(A::AbstractArray{T,N}, dist::Blocks{N}, assignment::Assignme return _to_darray(Distribute(dist, A, procgrid)) end -distribute(A::AbstractArray, ::AutoBlocks, assignment::AssignmentType = :arbitrary) = distribute(A, auto_blocks(A), assignment) +function distribute(A::AbstractArray{T,N}, ::AutoBlocks, + assignment::AssignmentType = :arbitrary) where {T,N} + part, assign, plan = Tapes.resolve_partitioning(T, size(A), AutoBlocks(), assignment) + return Tapes.track!(distribute(A, part, assign), plan) +end function distribute(x::AbstractArray{T,N}, n::NTuple{N}, assignment::AssignmentType{N} = :arbitrary) where {T,N} p = map((d, dn)->ceil(Int, d / dn), size(x), n) distribute(x, Blocks(p), assignment) @@ -550,33 +566,39 @@ DVector(A::AbstractVector{T}, assignment::AssignmentType{1} = :arbitrary) where DMatrix(A::AbstractMatrix{T}, assignment::AssignmentType{2} = :arbitrary) where T = DMatrix(A, AutoBlocks(), assignment) DArray(A::AbstractArray, assignment::AssignmentType = :arbitrary) = DArray(A, AutoBlocks(), assignment) -DVector(A::AbstractVector{T}, ::AutoBlocks, assignment::AssignmentType{1} = :arbitrary) where T = DVector(A, auto_blocks(A), assignment) -DMatrix(A::AbstractMatrix{T}, ::AutoBlocks, assignment::AssignmentType{2} = :arbitrary) where T = DMatrix(A, auto_blocks(A), assignment) -DArray(A::AbstractArray, ::AutoBlocks, assignment::AssignmentType = :arbitrary) = DArray(A, auto_blocks(A), assignment) +DVector(A::AbstractVector{T}, ::AutoBlocks, assignment::AssignmentType{1} = :arbitrary) where T = + distribute(A, AutoBlocks(), assignment) +DMatrix(A::AbstractMatrix{T}, ::AutoBlocks, assignment::AssignmentType{2} = :arbitrary) where T = + distribute(A, AutoBlocks(), assignment) +DArray(A::AbstractArray, ::AutoBlocks, assignment::AssignmentType = :arbitrary) = + distribute(A, AutoBlocks(), assignment) struct AllocateUndef{S} end (::AllocateUndef{S})(T, dims::Dims{N}) where {S,N} = Array{S,N}(undef, dims) -function DArray{T,N}(::UndefInitializer, dist::Blocks{N}, dims::NTuple{N,Int}; assignment::AssignmentType{N} = :arbitrary) where {T,N} +function DArray{T,N}(::UndefInitializer, dist::Union{Blocks{N},AutoBlocks}, dims::NTuple{N,Int}; + assignment::AssignmentType = :arbitrary) where {T,N} + part, assign, plan = Tapes.resolve_partitioning(T, dims, dist, assignment) domain = ArrayDomain(map(x->1:x, dims)) - subdomains = partition(dist, domain) - a = AllocateArray(T, AllocateUndef{T}(), false, domain, subdomains, dist, assignment) - return _to_darray(a) + subdomains = partition(part, domain) + a = AllocateArray(T, AllocateUndef{T}(), false, domain, subdomains, part, assign) + return Tapes.track!(_to_darray(a), plan) end DArray{T,N}(::UndefInitializer, dist::Blocks{N}, dims::Vararg{Int,N}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = DArray{T,N}(undef, dist, (dims...,); assignment) DArray{T,N}(::UndefInitializer, dims::NTuple{N,Int}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = - DArray{T,N}(undef, auto_blocks(dims), dims; assignment) + DArray{T,N}(undef, AutoBlocks(), dims; assignment) DArray{T,N}(::UndefInitializer, dims::Vararg{Int,N}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = - DArray{T,N}(undef, auto_blocks((dims...,)), (dims...,); assignment) + DArray{T,N}(undef, AutoBlocks(), (dims...,); assignment) -DArray{T}(::UndefInitializer, dist::Blocks{N}, dims::NTuple{N,Int}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = +DArray{T}(::UndefInitializer, dist::Union{Blocks{N},AutoBlocks}, dims::NTuple{N,Int}; + assignment::AssignmentType = :arbitrary) where {T,N} = DArray{T,N}(undef, dist, dims; assignment) DArray{T}(::UndefInitializer, dist::Blocks{N}, dims::Vararg{Int,N}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = DArray{T,N}(undef, dist, (dims...,); assignment) DArray{T}(::UndefInitializer, dims::NTuple{N,Int}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = - DArray{T,N}(undef, auto_blocks(dims), dims; assignment) + DArray{T,N}(undef, AutoBlocks(), dims; assignment) DArray{T}(::UndefInitializer, dims::Vararg{Int,N}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = - DArray{T,N}(undef, auto_blocks((dims...,)), (dims...,); assignment) + DArray{T,N}(undef, AutoBlocks(), (dims...,); assignment) function DArray(A::WrappedDArray{T}; assignment::AssignmentType = :arbitrary) where T B = DArray{T}(undef, size(A); assignment) diff --git a/src/array/iterativesolvers.jl b/src/array/iterativesolvers.jl new file mode 100644 index 000000000..f3d2f116c --- /dev/null +++ b/src/array/iterativesolvers.jl @@ -0,0 +1,358 @@ +# Distributed iterative (Krylov) linear solvers. +# +# The user-facing entry points (`Dagger.cg`, `Dagger.minres`, `Dagger.gmres`, +# `Dagger.bicgstab`, and the generic `Dagger.krylov_solve`) live here so the +# solver backend can evolve without changing user code. The actual solves are +# implemented in `ext/KrylovExt.jl`, which is loaded when `Krylov` is available +# (`using Krylov`). The matrix-free building blocks they rely on -- distributed +# `mul!`/SpMV and the BLAS-1 vector ops -- live in `mul.jl`/`linalg.jl`. + +""" + cg(A, b::DVector; M=I, atol, rtol, itmax, ...) -> (x::DVector, stats) + +Solve the symmetric positive-definite system `A x = b` with the conjugate +gradient method, distributed over `A`'s and `b`'s chunks. `A` may be a +`DMatrix` (dense or sparse-backed) or any object supporting `mul!(y, A, x)` over +`DVector`s (matrix-free). Requires `Krylov.jl` to be loaded. + +See also [`minres`](@ref), [`gmres`](@ref), [`bicgstab`](@ref), and +[`krylov_solve`](@ref). Keyword arguments are forwarded to `Krylov.cg!`. +""" +function cg end + +""" + minres(A, b::DVector; M=I, atol, rtol, itmax, ...) -> (x::DVector, stats) + +Solve the symmetric (possibly indefinite) system `A x = b` with MINRES. +Requires `Krylov.jl` to be loaded. See [`cg`](@ref). +""" +function minres end + +""" + gmres(A, b::DVector; M=I, N=I, restart=false, memory=20, ...) -> (x::DVector, stats) + +Solve the general (nonsymmetric) system `A x = b` with restarted GMRES. +`M`/`N` are left/right preconditioners. Requires `Krylov.jl` to be loaded. +See [`cg`](@ref). +""" +function gmres end + +""" + bicgstab(A, b::DVector; M=I, N=I, ...) -> (x::DVector, stats) + +Solve the general (nonsymmetric) system `A x = b` with BiCGStab (short +recurrence, low memory). Requires `Krylov.jl` to be loaded. See [`cg`](@ref). +""" +function bicgstab end + +""" + krylov_solve(method::Symbol, A, b::DVector; kwargs...) -> (x::DVector, stats) + +Generic entry point dispatching to the iterative `method` (`:cg`, `:minres`, +`:gmres`, `:bicgstab`). Requires `Krylov.jl` to be loaded. +""" +function krylov_solve end + +# Friendly fallbacks: these generic methods are shadowed by the more specific +# `(A, b::DVector)` methods added in `ext/KrylovExt.jl` once Krylov is loaded. +_krylov_required(name) = throw(ArgumentError( + "Dagger.$name requires Krylov.jl. Run `using Krylov` (or `import Krylov`) \ + to enable distributed iterative solvers.")) +cg(A, b; kwargs...) = _krylov_required(:cg) +minres(A, b; kwargs...) = _krylov_required(:minres) +gmres(A, b; kwargs...) = _krylov_required(:gmres) +bicgstab(A, b; kwargs...) = _krylov_required(:bicgstab) +krylov_solve(method::Symbol, A, b; kwargs...) = _krylov_required(:krylov_solve) + +# --- Preconditioners ------------------------------------------------------ +# A preconditioner object `P` represents the (approximate) *inverse* operator +# `M⁻¹`: applying it (`mul!(y, P, x)`) computes `y = M⁻¹ x`, which is exactly an +# ordinary matrix-vector product by the operator `P` stands for (no inversion +# happens at apply time -- any reciprocals/factorizations are precomputed once). +# This matches Krylov's `ldiv=false` convention, where the object passed as `M` +# is applied via `mul!(y, M, x)` to compute `y ← M⁻¹ x`, so these are passed +# straight through as `M=P`. +# +# These preconditioners are backend-agnostic: they operate on `DVector`s (and a +# `DMatrix`'s diagonal tiles) and need no solver backend, so they live in core. +# +# A note on the square-tile requirement below: the iterative solvers allocate +# *all* workspace vectors as `similar(b)` (one partitioning), and each must +# serve as both the length-`n` input and length-`n` output of `mul!(y, A, x)`. +# `gemv_dagger!` requires the input to match `A`'s column blocks and the output +# to match `A`'s row blocks, which is only simultaneously possible when those +# block sizes are equal. So the operator must use square tiles regardless of +# preconditioning; the preconditioners simply share that constraint. (Any +# uniform square tile size is fine, including a ragged final block.) + +""" + AbstractDaggerPreconditioner + +Supertype for Dagger's distributed preconditioners. A preconditioner `P` +represents an (approximate) inverse operator `M⁻¹` and applies it via +`mul!(y, P, x)` (`y = M⁻¹ x`) over `DVector`s, so it can be passed to the +solvers as `M=P` (with `ldiv=false`, the default). +""" +abstract type AbstractDaggerPreconditioner end + +""" + JacobiPreconditioner(A::DMatrix) + +Diagonal (Jacobi) preconditioner. The object represents the inverse-diagonal +operator `M⁻¹ = inv(Diagonal(A))`; it precomputes and stores the reciprocal +diagonal `dinv = 1 ./ diag(A)`, so applying it (`mul!(y, P, x)`) is a single +elementwise multiply `y = dinv .* x` per chunk -- an ordinary product by the +stored operator, not an inversion. Cheap to build (one diagonal extraction per +diagonal tile) and to apply. + +Requires square diagonal tiles (equal row/column block sizes); see the note +above on why the solver itself imposes this. Repartition `A` with equal block +sizes if needed. +""" +struct JacobiPreconditioner{V<:DVector} <: AbstractDaggerPreconditioner + dinv::V +end + +JacobiPreconditioner(A::DMatrix) = JacobiPreconditioner(_jacobi_dinv(A)) + +# Per-tile kernel: write the reciprocal diagonal of `tile` into `out`. +_set_inv_diag!(out, tile) = (out .= inv.(LinearAlgebra.diag(tile)); return nothing) + +# Validate that `A` has square diagonal tiles and return (n, Ac, mt, blocksize). +function _check_square_tiles(A::DMatrix, who) + n = LinearAlgebra.checksquare(A) + Ac = A.chunks + mt, nt = size(Ac) + mb, nb = A.partitioning.blocksize + # For a square matrix with square tiles, mt == nt follows automatically; we + # only need to reject non-square tiles (mb != nb). + mb == nb || throw(ArgumentError( + "$who requires square diagonal tiles (got block size $(mb)x$(nb)); \ + repartition A with equal block sizes (this matches the iterative \ + solver's own requirement on the operator)")) + return n, Ac, mt, mb +end + +function _jacobi_dinv(A::DMatrix{T}) where T + n, Ac, mt, mb = _check_square_tiles(A, "JacobiPreconditioner") + dinv = DVector{T}(undef, Blocks(mb), n) + dc = dinv.chunks + Dagger.spawn_datadeps() do + for i in 1:mt + Dagger.@spawn _set_inv_diag!(Out(dc[i]), In(Ac[i, i])) + end + end + return dinv +end + +function LinearAlgebra.mul!(y::DVector, P::JacobiPreconditioner, x::DVector) + part = P.dinv.partitioning + maybe_copy_buffered(P.dinv => part, x => part, y => part) do dinv, x, y + dc, xc, yc = dinv.chunks, x.chunks, y.chunks + Dagger.spawn_datadeps() do + for i in eachindex(yc) + Dagger.@spawn _jacobi_apply!(Out(yc[i]), In(dc[i]), In(xc[i])) + end + end + end + return y +end +_jacobi_apply!(y, dinv, x) = (y .= dinv .* x; return nothing) + +# --- Block-diagonal preconditioners --------------------------------------- +# A family of preconditioners of the form `M⁻¹ = blockdiag(op₁, …, op_k)`, where +# each `opⱼ` approximates the inverse of `A`'s `j`-th diagonal tile `Aⱼⱼ`. They +# share all machinery and differ only in how a per-tile operator is *built*: +# +# - BlockJacobiPreconditioner: exact tile solve via `lu` (dense or sparse). +# - BlockILUPreconditioner: incomplete LU per tile (needs IncompleteLU.jl). +# - AMGPreconditioner: an AMG hierarchy per tile (needs AlgebraicMultigrid.jl). +# +# Building per tile is embarrassingly parallel and a natural fit for the tiled +# layout. A single tile (`Blocks(n, n)`) makes any of these a *global* +# preconditioner over the whole matrix; many tiles make it a scalable +# block-Jacobi / additive-Schwarz variant that trades some convergence for +# parallelism. All require square diagonal tiles (see `JacobiPreconditioner`). +# +# A per-tile operator (`lu`/ILU factor, AMG hierarchy) generally cannot be moved +# between workers (dense `LU` has no `move!`; `UmfpackLU`/AMG hold process-bound +# resources). So each operator is built *once* and pinned (via +# `tochunk(..., ProcessScope)`) to the worker that owns its tile, and every apply +# for that block is scheduled there (`compute_scope`); datadeps then moves only +# the (movable) vector chunks to the operator, never the operator itself. + +""" + AbstractBlockPreconditioner <: AbstractDaggerPreconditioner + +Block-diagonal preconditioner: holds one per-diagonal-tile operator `opⱼ` +(`y ← opⱼ⁻¹ x`), pinned to its tile's worker, and applies them independently per +block. Concrete subtypes (`BlockJacobiPreconditioner`, `BlockILUPreconditioner`, +`AMGPreconditioner`) share these fields: `ops`, `scopes`, `part`, `n`. +""" +abstract type AbstractBlockPreconditioner <: AbstractDaggerPreconditioner end + +_tile_scope(c::Chunk) = ProcessScope(root_worker_id(c)) +_tile_scope(t::DTask) = ProcessScope(root_worker_id(fetch(t; raw=true))) + +# Build the per-tile operator on the current worker and pin the result there: the +# returned `Chunk` is process-scoped, so it can never be moved off this worker. +function _build_tile_pinned(build, tile) + op = build(tile) + proc = Dagger.task_processor() + return tochunk(op, proc, ProcessScope(root_worker_id(proc))) +end + +# Build a block preconditioner of type `Ctor` by applying `build` to each +# diagonal tile (pinned to the tile's worker). `build(tile) -> op` returns a +# per-tile operator supporting the block apply below. +function _build_block_preconditioner(Ctor, A::DMatrix, who, build) + n, Ac, mt, mb = _check_square_tiles(A, who) + ops = Vector{Any}(undef, mt) + scopes = Vector{ProcessScope}(undef, mt) + for i in 1:mt + tile = Ac[i, i] + scope = _tile_scope(tile) + scopes[i] = scope + ops[i] = Dagger.@spawn compute_scope=scope _build_tile_pinned(build, tile) + end + return Ctor(ops, scopes, Blocks(mb), n) +end + +function LinearAlgebra.mul!(y::DVector, P::AbstractBlockPreconditioner, x::DVector) + part = P.part + maybe_copy_buffered(x => part, y => part) do x, y + xc, yc = x.chunks, y.chunks + length(yc) == length(P.ops) || throw(DimensionMismatch( + "$(nameof(typeof(P))) has $(length(P.ops)) blocks but the vector has \ + $(length(yc)) chunks")) + Dagger.spawn_datadeps() do + for i in eachindex(yc) + # Pin the apply to the operator's worker; the operator is passed as + # an untracked arg (read-only, already-pinned) so datadeps never + # moves it -- only the vector chunks are moved to this worker. + Dagger.@spawn compute_scope=P.scopes[i] _block_apply!(Out(yc[i]), P.ops[i], In(xc[i])) + end + end + end + return y +end + +# Apply one block operator: `y = op⁻¹ x`. Default uses `\`; backends override for +# operator types where `\` is unavailable (e.g. ILU/AMG use `ldiv!`). +# Host factorizations (UMFPACK/`lu` of gathered CSC) need a host RHS; GPU vector +# chunks are gathered, solved, and written back. CPU `Array` short-circuits. +function _block_apply!(y, op, x) + if x isa Array + y .= op \ x + else + copyto!(y, op \ Adapt.adapt(Array, x)) + end + return nothing +end + +# Underlying matrix of a tile (overridden for `DSparseArray` in `sparse.jl`). +_tile_matrix(A) = A + +""" + BlockJacobiPreconditioner(A::DMatrix) + +Block-Jacobi preconditioner: `M⁻¹ = blockdiag(A₁₁, …, A_kk)⁻¹`. Each diagonal +tile is factorized *once* with `lu` (sparse or dense) and the apply solves +`Aᵢᵢ yᵢ = xᵢ`. Stronger than `JacobiPreconditioner` (captures intra-block +coupling); a single tile recovers an exact solve. See +[`AbstractBlockPreconditioner`](@ref). Requires square diagonal tiles. +""" +struct BlockJacobiPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F # cached per-tile factorizations (pinned to their workers) + scopes::S # the `ProcessScope` each operator/apply is pinned to + part::Blocks{1} # partitioning of the vectors it applies to + n::Int +end + +# Factorize a diagonal tile. Backends override for their storage (e.g. sparse +# tiles factorize the inner `SparseMatrixCSC`); the default is a dense LU factor. +_factorize_tile(A) = LinearAlgebra.lu(A) + +BlockJacobiPreconditioner(A::DMatrix) = + _build_block_preconditioner(BlockJacobiPreconditioner, A, "BlockJacobiPreconditioner", _factorize_tile) + +""" + BlockILUPreconditioner(A::DMatrix; τ=0.001, kwargs...) + +Block incomplete-LU preconditioner: an ILU factorization (with drop tolerance +`τ`) of each diagonal tile, applied per block. Cheaper setup than a full block +solve, good as a general-purpose preconditioner. Requires `IncompleteLU.jl` to +be loaded and sparse-backed tiles. See [`AbstractBlockPreconditioner`](@ref). +""" +struct BlockILUPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F + scopes::S + part::Blocks{1} + n::Int +end +# Friendly fallback (shadowed by the `::DMatrix` method added in `ext/IncompleteLUExt.jl`). +BlockILUPreconditioner(A; kwargs...) = throw(ArgumentError( + "Dagger.BlockILUPreconditioner requires IncompleteLU.jl. Run `using IncompleteLU` \ + to enable block incomplete-LU preconditioning.")) + +""" + AMGPreconditioner(A::DMatrix; method=:ruge_stuben, kwargs...) + +Algebraic-multigrid preconditioner: builds an AMG hierarchy (`method` is +`:ruge_stuben` or `:smoothed_aggregation`) for each diagonal tile and applies a +V-cycle per block. Near mesh-independent convergence for elliptic (Poisson-like) +operators. With one tile this is global AMG; with many tiles it is a scalable +block/additive-Schwarz AMG. Requires `AlgebraicMultigrid.jl` to be loaded and +sparse-backed tiles. See [`AbstractBlockPreconditioner`](@ref). +""" +struct AMGPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F + scopes::S + part::Blocks{1} + n::Int +end +# Friendly fallback (shadowed by the `::DMatrix` method added in `ext/AlgebraicMultigridExt.jl`). +AMGPreconditioner(A; kwargs...) = throw(ArgumentError( + "Dagger.AMGPreconditioner requires AlgebraicMultigrid.jl. Run \ + `using AlgebraicMultigrid` to enable algebraic-multigrid preconditioning.")) + +""" + BlockKLUPreconditioner(A::DMatrix; kwargs...) + +Block direct preconditioner using KLU: an exact (pure-Julia) KLU factorization of +each diagonal tile, applied per block. With a single tile this is an exact direct +solve; with many tiles it is an exact-block-Jacobi preconditioner. Requires +`PureKLU.jl` to be loaded and sparse-backed tiles. See also +[`AbstractBlockPreconditioner`](@ref) and the whole-matrix solver +[`Dagger.klu`](@ref). +""" +struct BlockKLUPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F + scopes::S + part::Blocks{1} + n::Int +end +# Friendly fallback (shadowed by the `::DMatrix` method added in `ext/PureKLUExt.jl`). +BlockKLUPreconditioner(A; kwargs...) = throw(ArgumentError( + "Dagger.BlockKLUPreconditioner requires PureKLU.jl. Run `using PureKLU` to \ + enable block KLU preconditioning.")) + +""" + BlockUMFPACKPreconditioner(A::DMatrix; kwargs...) + +Block direct preconditioner using a pure-Julia UMFPACK-style LU of each diagonal +tile, applied per block. With a single tile this is an exact direct solve; with +many tiles it is an exact-block-Jacobi preconditioner. Requires `PureUMFPACK.jl` +to be loaded and sparse-backed tiles. See also [`AbstractBlockPreconditioner`](@ref) +and the whole-matrix solver [`Dagger.splu`](@ref). +""" +struct BlockUMFPACKPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F + scopes::S + part::Blocks{1} + n::Int +end +# Friendly fallback (shadowed by the `::DMatrix` method added in `ext/PureUMFPACKExt.jl`). +BlockUMFPACKPreconditioner(A; kwargs...) = throw(ArgumentError( + "Dagger.BlockUMFPACKPreconditioner requires PureUMFPACK.jl. Run \ + `using PureUMFPACK` to enable block UMFPACK preconditioning.")) diff --git a/src/array/linalg.jl b/src/array/linalg.jl index 2fcaa42fa..89d0575e1 100644 --- a/src/array/linalg.jl +++ b/src/array/linalg.jl @@ -4,6 +4,71 @@ function LinearAlgebra.norm2(A::DArray{T,N}) where {T,N} zeroRT = zero(real(T)) return sqrt(sum(map(norm->fetch(norm)::real(T), norms); init=zeroRT)) end + +# --- BLAS-1 vector operations --------------------------------------------- +# Efficient, distributed level-1 operations needed by iterative (Krylov) +# solvers. These run one BLAS-1 kernel per chunk (avoiding the generic +# scalar-indexing fallbacks, which are slow and trigger scalar access on GPUs). +# Operands may be partitioned differently: we align them with +# `maybe_copy_buffered` (a no-op when they already share a layout, and which +# copies results back to the originals for in-place operations), so users keep +# full control over their own (possibly sub-optimal) partitioning. + +function LinearAlgebra.dot(x::DArray{Tx,N}, y::DArray{Ty,N}) where {Tx,Ty,N} + size(x) == size(y) || throw(DimensionMismatch("dot: x has size $(size(x)), y has size $(size(y))")) + R = typeof(LinearAlgebra.dot(zero(Tx), zero(Ty))) + return maybe_copy_buffered(x => x.partitioning, y => x.partitioning) do x, y + xc, yc = x.chunks, y.chunks + parts = [Dagger.@spawn LinearAlgebra.dot(xc[i], yc[i]) for i in eachindex(xc)] + sum(fetch, parts; init=zero(R)) + end +end + +function LinearAlgebra.axpy!(a::Number, x::DArray{Tx,N}, y::DArray{Ty,N}) where {Tx,Ty,N} + size(x) == size(y) || throw(DimensionMismatch("axpy!: x has size $(size(x)), y has size $(size(y))")) + # Align the read-only `x` to `y`'s layout so only `x` may be buffered; `y` is + # updated in place (and copied back by `maybe_copy_buffered` if it was buffered). + maybe_copy_buffered(x => y.partitioning, y => y.partitioning) do x, y + xc, yc = x.chunks, y.chunks + Dagger.spawn_datadeps() do + for i in eachindex(xc) + Dagger.@spawn LinearAlgebra.axpy!(a, In(xc[i]), InOut(yc[i])) + end + end + end + return y +end + +function LinearAlgebra.axpby!(a::Number, x::DArray{Tx,N}, b::Number, y::DArray{Ty,N}) where {Tx,Ty,N} + size(x) == size(y) || throw(DimensionMismatch("axpby!: x has size $(size(x)), y has size $(size(y))")) + maybe_copy_buffered(x => y.partitioning, y => y.partitioning) do x, y + xc, yc = x.chunks, y.chunks + Dagger.spawn_datadeps() do + for i in eachindex(xc) + Dagger.@spawn LinearAlgebra.axpby!(a, In(xc[i]), b, InOut(yc[i])) + end + end + end + return y +end + +function LinearAlgebra.rmul!(x::DArray, a::Number) + Dagger.spawn_datadeps() do + for c in x.chunks + Dagger.@spawn LinearAlgebra.rmul!(InOut(c), a) + end + end + return x +end + +function LinearAlgebra.lmul!(a::Number, x::DArray) + Dagger.spawn_datadeps() do + for c in x.chunks + Dagger.@spawn LinearAlgebra.lmul!(a, InOut(c)) + end + end + return x +end function LinearAlgebra.norm2(A::UpperTriangular{T,<:DArray{T,2}}) where T Ac = parent(A).chunks Ac_upper = [] diff --git a/src/array/lu.jl b/src/array/lu.jl index 45d0c2c38..8ede04e6f 100644 --- a/src/array/lu.jl +++ b/src/array/lu.jl @@ -3,10 +3,15 @@ LinearAlgebra.lu(A::DMatrix{T}, pivot::Union{LinearAlgebra.RowMaximum,LinearAlge LinearAlgebra.lu!(A::DMatrix{T}, pivot::Union{LinearAlgebra.RowMaximum,LinearAlgebra.NoPivot} = LinearAlgebra.RowMaximum(); check::Bool=true, allowsingular::Bool=false) where {T<:LinearAlgebra.BlasFloat} = LinearAlgebra.lu(A, pivot; check=check, allowsingular=allowsingular) function LinearAlgebra.lu(A::DMatrix{T}, ::LinearAlgebra.NoPivot; check::Bool = true, allowsingular::Bool = false) where {T<:LinearAlgebra.BlasFloat} + # Record against the caller's array, not the internal copy: otherwise the + # allocation site that produced `A` never sees `:lu!` and the tape cannot + # learn from the usual `F = lu(A)` pattern. + Dagger.@record_op :lu! A A_copy = LinearAlgebra._lucopy(A, LinearAlgebra.lutype(T)) return LinearAlgebra.lu!(A_copy, LinearAlgebra.NoPivot(); check) end function LinearAlgebra.lu!(A::DMatrix{T}, ::LinearAlgebra.NoPivot; check::Bool = true, allowsingular::Bool = false) where {T<:LinearAlgebra.BlasFloat} + Dagger.@record_op :lu! A check && LinearAlgebra.LAPACK.chkfinite(A) zone = one(T) @@ -260,10 +265,13 @@ end # Implementation of https://inria.hal.science/hal-04984070v1/file/ipdps_paper.pdf function LinearAlgebra.lu(A::DMatrix{T}, ::LinearAlgebra.RowMaximum; check::Bool = true, allowsingular::Bool = false) where {T<:LinearAlgebra.BlasFloat} + # See NoPivot method: attribute the op to the source array, not the copy. + Dagger.@record_op :lu! A A_copy = LinearAlgebra._lucopy(A, LinearAlgebra.lutype(T)) return LinearAlgebra.lu!(A_copy, LinearAlgebra.RowMaximum(); check, allowsingular) end function LinearAlgebra.lu!(A::DMatrix{T}, ::LinearAlgebra.RowMaximum; check::Bool = true, allowsingular::Bool = false) where {T<:LinearAlgebra.BlasFloat} + Dagger.@record_op :lu! A check && LinearAlgebra.LAPACK.chkfinite(A) zone = one(T) diff --git a/src/array/map-reduce.jl b/src/array/map-reduce.jl index 6e16a912b..b0d373943 100644 --- a/src/array/map-reduce.jl +++ b/src/array/map-reduce.jl @@ -41,7 +41,10 @@ function stage(ctx::Context, node::Map) return DArray(RT, domain(primary), domainchunks(primary), thunks, partitioning, concat) end -Base.map(f, x::DArray, xs::DArray...) = _to_darray(Map(f, (x, xs...))) +function Base.map(f, x::DArray, xs::DArray...) + Dagger.@record_op :map x + return _to_darray(Map(f, (x, xs...))) +end #### MapReduce @@ -139,8 +142,15 @@ end import Statistics: mean, var, std import OnlineStats -Base.mapreduce(f::Function, op::Function, x::DArray; dims=nothing, init=Base._InitialValue()) = - _mapreduce_maybesync(f, nothing, op, x, dims, init) +function Base.mapreduce(f::Function, op::Function, x::DArray; dims=nothing, init=Base._InitialValue()) + Dagger.@record_op :mapreduce x + return _mapreduce_maybesync(f, nothing, op, x, dims, init) +end + +function Base.reduce(op::Function, x::DArray; dims=nothing, init=Base._InitialValue()) + Dagger.@record_op :reduce x + return _mapreduce_maybesync(identity, nothing, op, x, dims, init) +end Base.sum(x::DArray; dims=nothing, init=Base._InitialValue()) = sum(identity, x; dims, init) diff --git a/src/array/matrix.jl b/src/array/matrix.jl index b9935f662..9b7e77c1d 100644 --- a/src/array/matrix.jl +++ b/src/array/matrix.jl @@ -1,6 +1,7 @@ # Transpose/Adjoint function copydiag(f, A::DArray{T, 2}) where T + Dagger.@record_op (f === Transpose ? :transpose : :adjoint) A Ac = A.chunks Ac_copy = Matrix{Any}(undef, size(Ac, 2), size(Ac, 1)) _copytile(f, Ac) = copy(f(Ac)) diff --git a/src/array/mul.jl b/src/array/mul.jl index 14ed9f0ba..4d07c9f0e 100644 --- a/src/array/mul.jl +++ b/src/array/mul.jl @@ -1,3 +1,74 @@ +""" + matmatmul!(C, transA::Char, transB::Char, A, B, alpha, beta) + +A general-purpose matrix-matrix multiply, like `LinearAlgebra.generic_matmatmul!`, +but with extra functionality. May internally convert `A` and `B` to a type that +better matches `C` and provides optimal portability and, when possible, +better performance. The actual matrix multiply operation should happen in +`LinearAlgebra.generic_matmatmul!` or an equivalent call. + +The following automatic conversions are performed: +- If no `LinearAlgebra.generic_matmatmul!` method is available, convert `A` and `B` to dense Array-like +- If `C` is a `DSparseMatrix`, perform the operation out-of-place and then update `C` in-place +""" +function matmatmul!( + C, + transA::Char, + transB::Char, + A, + B, + alpha, + beta +) + EC = eltype(C) + EA = eltype(A) + EB = eltype(B) + + TC = typeof(C) + TA = typeof(A) + TB = typeof(B) + + mam = LinearAlgebra.MulAddMul(alpha, beta) + + # Check if C doesn't support in-place operations (e.g. DSparseMatrix) + # We'll get here if A and B don't have equivalent types + if isa(C, DSparseMatrix) + C.mat = alpha * A * B + beta * C.mat + return C + end + + # Check if the call will fail due to MethodError + sig = Tuple{TC, Char, Char, TA, TB, typeof(mam)} + if !hasmethod(LinearAlgebra.generic_matmatmul!, sig) + # Convert to Array-like + # FIXME: GPU support + C_new = C + A_new = collect(A) + B_new = collect(B) + alpha_new = alpha + beta_new = beta + # FIXME: Re-check hasmethod, and if no method, then convert and bounce C + @goto ready + end + + C_new = C + A_new = A + B_new = B + alpha_new = alpha + beta_new = beta + + @label ready + mam_new = LinearAlgebra.MulAddMul(alpha_new, beta_new) + return LinearAlgebra.generic_matmatmul!( + C_new, + transA, + transB, + A_new, + B_new, + mam_new + ) +end + function LinearAlgebra.generic_matmatmul!( C::DMatrix{T}, transA::Char, @@ -17,6 +88,7 @@ function LinearAlgebra.generic_matmatmul!( alpha::Number, beta::Number, ) where {T} + Dagger.@record_op :mul! C A B partC, partA, partB = _repartition_matmatmul(C, A, B, transA, transB) if all(in(('N', 'T', 'C')), (transA, transB)) @@ -41,7 +113,8 @@ function LinearAlgebra.generic_matmatmul!( return gemm_dagger!(C, transA, transB, A, B, alpha, beta) end end -function _repartition_matmatmul(C, A, B, transA::Char, transB::Char)::Tuple{Blocks{2}, Blocks{2}, Blocks{2}} +# FIXME: Mixed-precision methods +function _repartition_matmatmul(C, A, B, transA::Char, transB::Char) partA = A.partitioning.blocksize partB = B.partitioning.blocksize istransA = transA == 'T' || transA == 'C' @@ -141,28 +214,28 @@ function gemm_dagger!( # A: NoTrans / B: NoTrans for k in range(1, Ant) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemm!( + Dagger.@spawn matmatmul!( + InOut(Cc[m, n]), transA, transB, - alpha, In(Ac[m, k]), In(Bc[k, n]), + alpha, mzone, - InOut(Cc[m, n]), ) end else # A: NoTrans / B: [Conj]Trans for k in range(1, Ant) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemm!( + Dagger.@spawn matmatmul!( + InOut(Cc[m, n]), transA, transB, - alpha, In(Ac[m, k]), In(Bc[n, k]), + alpha, mzone, - InOut(Cc[m, n]), ) end end @@ -171,28 +244,28 @@ function gemm_dagger!( # A: [Conj]Trans / B: NoTrans for k in range(1, Amt) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemm!( + Dagger.@spawn matmatmul!( + InOut(Cc[m, n]), transA, transB, - alpha, In(Ac[k, m]), In(Bc[k, n]), + alpha, mzone, - InOut(Cc[m, n]), ) end else # A: [Conj]Trans / B: [Conj]Trans for k in range(1, Amt) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemm!( + Dagger.@spawn matmatmul!( + InOut(Cc[m, n]), transA, transB, - alpha, In(Ac[k, m]), In(Bc[n, k]), + alpha, mzone, - InOut(Cc[m, n]), ) end end @@ -224,6 +297,7 @@ function syrk_dagger!( _alpha, _beta, ) where {T} + Dagger.@record_op :syrk! C A Ac = A.chunks Cc = C.chunks @@ -240,6 +314,7 @@ function syrk_dagger!( iscomplex = T <: Complex transs = iscomplex ? 'C' : 'T' + anti_transs = trans == 'N' ? transs : 'N' Dagger.spawn_datadeps() do for n in range(1, Cnt) @@ -247,116 +322,63 @@ function syrk_dagger!( # NoTrans for k in range(1, Ant) mzone = k == 1 ? real(beta) : one(real(T)) - if iscomplex - Dagger.@spawn BLAS.herk!( - uplo, + _alpha = iscomplex ? real(alpha) : alpha + Dagger.@spawn matmatmul!( + InOut(Cc[n, n]), + trans, + anti_transs, + In(Ac[n, k]), + In(Ac[n, k]), + _alpha, + mzone, + ) + end + # NoTrans / Upper + for m in range(n + 1, Cmt) + for k in range(1, Ant) + mzone = k == 1 ? beta : one(T) + Dagger.@spawn matmatmul!( + InOut(Cc[n, m]), trans, - real(alpha), + transs, In(Ac[n, k]), - mzone, - InOut(Cc[n, n]), - ) - else - Dagger.@spawn BLAS.syrk!( - uplo, - trans, + In(Ac[m, k]), alpha, - In(Ac[n, k]), mzone, - InOut(Cc[n, n]), ) end end - if uplo == 'L' - # NoTrans / Lower - for m in range(n + 1, Cmt) - for k in range(1, Ant) - mzone = k == 1 ? beta : one(T) - Dagger.@spawn BLAS.gemm!( - trans, - transs, - alpha, - In(Ac[m, k]), - In(Ac[n, k]), - mzone, - InOut(Cc[m, n]), - ) - end - end - else - # NoTrans / Upper - for m in range(n + 1, Cmt) - for k in range(1, Ant) - mzone = k == 1 ? beta : one(T) - Dagger.@spawn BLAS.gemm!( - trans, - transs, - alpha, - In(Ac[n, k]), - In(Ac[m, k]), - mzone, - InOut(Cc[n, m]), - ) - end - end - end else # [Conj]Trans for k in range(1, Amt) mzone = k == 1 ? real(beta) : one(real(T)) - if iscomplex - Dagger.@spawn BLAS.herk!( - uplo, + _alpha = iscomplex ? real(alpha) : alpha + _trans = iscomplex ? transs : trans + Dagger.@spawn matmatmul!( + InOut(Cc[n, n]), + _trans, + anti_transs, + In(Ac[k, n]), + In(Ac[k, n]), + _alpha, + mzone, + ) + end + # [Conj]Trans / Upper + for m in range(n + 1, Cmt) + for k in range(1, Amt) + mzone = k == 1 ? beta : one(T) + Dagger.@spawn matmatmul!( + InOut(Cc[n, m]), transs, - real(alpha), + 'N', In(Ac[k, n]), - mzone, - InOut(Cc[n, n]), - ) - else - Dagger.@spawn BLAS.syrk!( - uplo, - trans, + In(Ac[k, m]), alpha, - In(Ac[k, n]), mzone, - InOut(Cc[n, n]), ) end end - if uplo == 'L' - # [Conj]Trans / Lower - for m in range(n + 1, Cmt) - for k in range(1, Amt) - mzone = k == 1 ? beta : one(T) - Dagger.@spawn BLAS.gemm!( - transs, - 'N', - alpha, - In(Ac[k, m]), - In(Ac[k, n]), - mzone, - InOut(Cc[m, n]), - ) - end - end - else - # [Conj]Trans / Upper - for m in range(n + 1, Cmt) - for k in range(1, Amt) - mzone = k == 1 ? beta : one(T) - Dagger.@spawn BLAS.gemm!( - transs, - 'N', - alpha, - In(Ac[k, n]), - In(Ac[k, m]), - mzone, - InOut(Cc[n, m]), - ) - end - end - end end end end @@ -398,7 +420,7 @@ end return A end -@inline function copytile!(A::AbstractMatrix{T}, B::AbstractMatrix{T})::Nothing where {T} +function copytile!(A, B) m, n = size(A) C = B' @@ -408,15 +430,14 @@ end return nothing end -@inline function copydiagtile!(A::AbstractMatrix{T}, uplo::AbstractChar)::Nothing where {T} +function copydiagtile!(A, uplo) m, n = size(A) - Acpy = copy(A) if uplo == 'U' - C = UpperTriangular(Acpy)' + UpperTriangular(Acpy) + C = UpperTriangular(A)' + UpperTriangular(A) C[diagind(C)] .= A[diagind(A)] elseif uplo == 'L' - C = LowerTriangular(Acpy)' + Acpy - UpperTriangular(Acpy) + C = LowerTriangular(A)' + A - UpperTriangular(A) C[diagind(C)] .= A[diagind(A)] end @@ -425,6 +446,7 @@ end end return nothing end + function LinearAlgebra.generic_matvecmul!( C::DVector{T}, transA::Char, @@ -442,6 +464,7 @@ function LinearAlgebra.generic_matvecmul!( _alpha::Number, _beta::Number, ) where {T} + Dagger.@record_op :mul! C A B partC, partA, partB = _repartition_matvecmul(C, A, B, transA) return maybe_copy_buffered(C=>partC, A=>partA, B=>partB) do C, A, B return gemv_dagger!(C, transA, A, B, _alpha, _beta) @@ -471,6 +494,21 @@ function _repartition_matvecmul(C, A, B, transA::Char)::Tuple{Blocks{1}, Blocks{ partC = (dimA,) return Blocks(partC...), Blocks(partA...), Blocks(partB...) end +""" + matvecmul!(C, transA::Char, A, B, alpha, beta) + +Tile-level matrix-vector multiply computing `C = alpha*op(A)*B + beta*C` in +place on the (dense) output vector `C`, where `op` is determined by `transA` +(`'N'`, `'T'`, `'C'`). Dispatches on the tile types: dense tiles use BLAS, while +sparse tiles (e.g. `DSparseArray`) provide their own method (in a package +extension) using a sparse matrix-vector product. This is the matvec analogue of +[`matmatmul!`](@ref). +""" +function matvecmul!(C, transA::Char, A, B, alpha, beta) + BLAS.gemv!(transA, alpha, A, B, beta, C) + return C +end + function gemv_dagger!( C::DVector{T}, transA::Char, @@ -513,26 +551,26 @@ function gemv_dagger!( # A: NoTrans for k in range(1, Ant) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemv!( + Dagger.@spawn matvecmul!( + InOut(Cc[m]), transA, - alpha, In(Ac[m, k]), In(Bc[k]), + alpha, mzone, - InOut(Cc[m]), ) end else # A: [Conj]Trans — C's blocks index A's column-blocks for k in range(1, Amt) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemv!( + Dagger.@spawn matvecmul!( + InOut(Cc[m]), transA, - alpha, In(Ac[k, m]), In(Bc[k]), + alpha, mzone, - InOut(Cc[m]), ) end end @@ -541,3 +579,53 @@ function gemv_dagger!( return C end + +wrap_as_darray(A::DArray) = A +wrap_as_darray(A::Array) = view(A, AutoBlocks()) +function wrap_as_darray(A::SubArray{T,Nv,Array{T,Na}}) where {T,Nv,Na} + Ap = parent(A) + part = auto_blocks(map(last, parentindices(Ap))) + partsize = part.blocksize + inds = parentindices(A) + inds_ranges_parent = ntuple(i->to_range(inds[i]), Val(Na)) + inds_ranges_view = ntuple(i->to_range(inds[i]), Val(Nv)) + subdomains = partition(part, ArrayDomain(inds_ranges_parent)) + nparts = size(subdomains) + chunks = Array{Any,Na}(undef, nparts...) + for idx in CartesianIndices(nparts) + subdomain_view = subdomains[idx] + subdomain_parent = ArrayDomain(ntuple(i->Nv >= i ? subdomain_view.indexes[i] : inds_ranges_parent[i], Val(Na))) + subinds = ntuple(i->subdomain_parent.indexes[i], Val(Na)) + subA = view(Ap, subinds...) + chunks[idx] = tochunk(subA) + end + return DArray(T, ArrayDomain(inds_ranges_parent), subdomains, chunks, part) +end + +# Generate generic_matvecmul! methods for all combinations of DArray, Array, and SubArray +for CT in (DVector, Vector, SubArray{<:Any,1,<:Array}), + AT in (DMatrix, Matrix, SubArray{<:Any,2,<:Array}), + BT in (DVector, Vector, SubArray{<:Any,1,<:Array}) + + # Don't commit type piracy + CT isa DArray || AT isa DArray || BT isa DArray || continue + + @eval function LinearAlgebra.generic_matvecmul!( + C::$(CT), + transA::Char, + A::$(AT), + B::$(BT), + _add::LinearAlgebra.MulAddMul, + ) + new_C = wrap_as_darray(C) + new_A = wrap_as_darray(A) + new_B = wrap_as_darray(B) + return LinearAlgebra.generic_matvecmul!( + new_C, + transA, + new_A, + new_B, + _add, + ) + end +end \ No newline at end of file diff --git a/src/array/operators.jl b/src/array/operators.jl index f93f09a75..837f89fe0 100644 --- a/src/array/operators.jl +++ b/src/array/operators.jl @@ -177,6 +177,7 @@ function imap!(f, A) end function Base.map!(f, a::DArray{T}) where T + Dagger.@record_op :map! a Dagger.spawn_datadeps() do for ca in chunks(a) Dagger.@spawn imap!(f, InOut(ca)) @@ -186,6 +187,7 @@ function Base.map!(f, a::DArray{T}) where T end function Base.map!(f, a::DArray{T}, b::AbstractArray{U}) where {T, U} + Dagger.@record_op :map! a b2 = view(b, a.partitioning) Dagger.spawn_datadeps() do for (c_a, c_b2) in zip(chunks(a), chunks(b2)) diff --git a/src/array/permute.jl b/src/array/permute.jl index 9da51e70b..08b204ca8 100644 --- a/src/array/permute.jl +++ b/src/array/permute.jl @@ -1,4 +1,5 @@ function Base.permutedims(A::DArray{T,N}, perm) where {T,N} + Dagger.@record_op :permutedims A dc = domainchunks(A) new_dc = DomainBlocks( ntuple(i -> dc.start[perm[i]], N), diff --git a/src/array/qr.jl b/src/array/qr.jl index 4409cc1c5..e35a7fcd7 100644 --- a/src/array/qr.jl +++ b/src/array/qr.jl @@ -833,6 +833,7 @@ In-place tiled Compact-WY QR factorization of a distributed matrix. than `p=1`. """ function LinearAlgebra.qr!(A::DMatrix{T}; ib::Union{Int,Nothing}=nothing, p::Int=1) where {T<:Number} + Dagger.@record_op :qr! A p >= 1 || throw(ArgumentError("p must be >= 1, got $p")) ib === nothing || ib >= 1 || throw(ArgumentError("ib must be >= 1, got $ib")) diff --git a/src/array/sort.jl b/src/array/sort.jl index 7abbe373f..109647c46 100644 --- a/src/array/sort.jl +++ b/src/array/sort.jl @@ -301,6 +301,7 @@ function Base.sort(v::ArrayOp; nsamples=2000, order::Ordering=default_ord) v1 = fetch(v) + Dagger.@record_op :sort! v1 ord = Base.Sort.ord(lt,by,rev,order) nchunks = nchunks === nothing ? length(v1.chunks) : nchunks cs = dsort_chunks(v1.chunks, nchunks, nsamples, diff --git a/src/array/sparse.jl b/src/array/sparse.jl new file mode 100644 index 000000000..c5ab69b91 --- /dev/null +++ b/src/array/sparse.jl @@ -0,0 +1,259 @@ +""" + DSparseArray{T,N} <: AbstractArray{T,N} + +A sparse array container, for which the contained array may be replaced with a +new one to support in-place operations. Designed to work well with Datadeps +algorithms: writes that reallocate (and grow/shrink) the inner sparse storage +are hidden behind the wrapper's stable identity, so Datadeps aliasing tracking +remains valid (see `aliases_as_whole`). + +`DSparseVector{T}` and `DSparseMatrix{T}` are aliases for the 1- and 2-dimensional +cases. The wrapper is general over `N` so it can hold sparse vectors, matrices, +and (eventually) higher-order sparse tensors (e.g. Finch tensors). +""" +mutable struct DSparseArray{T,N} <: AbstractArray{T,N} + mat +end +const DSparseVector{T} = DSparseArray{T,1} +const DSparseMatrix{T} = DSparseArray{T,2} + +# Convenience constructor inferring element type and dimensionality from `x`. +DSparseArray(x) = DSparseArray{eltype(x),ndims(x)}(x) + +_sparse_collect(M) = collect(M) +# Allocate a destination tile matching the inner storage's sparse backend. The +# default delegates to the storage type's own `similar`, which carries the +# concrete backend (e.g. `SparseMatrixCSC`) forward to the destination -- this is +# how `similar(::DArray)` propagates a tile's sparse type to a newly-allocated +# result DArray. Backends whose `similar` is unsuitable (e.g. Finch, whose generic +# `similar` yields tensor formats that destabilize later kernels) override this. +_sparse_similar(mat, ::Type{T}, dims::Dims) where {T} = similar(mat, T, dims) + +Base.eltype(M::DSparseArray{T}) where T = T +Base.size(M::DSparseArray) = size(M.mat) +Base.iterate(M::DSparseArray) = iterate(M.mat) +Base.iterate(M::DSparseArray, state) = iterate(M.mat, state) +Base.similar(M::DSparseArray, ::Type{T}, dims::Dims) where T = + DSparseArray{T,length(dims)}(_sparse_similar(M.mat, T, dims)) +Base.collect(M::DSparseArray) = _sparse_collect(M.mat) + +# N.B. hash and aliasing shouldn't change even if M.mat changes. +# This is the key property that makes `DSparseArray` safe for Datadeps: +# the aliasing identity is tied to the (stable) wrapper object, not to the +# inner storage, which may be reallocated (and grow/shrink) on writes. +Base.hash(M::DSparseArray, h::UInt) = hash(objectid(M), hash(DSparseArray, h)) + +# Sparse containers must alias as a single, indivisible unit. Their storage is +# reallocated (and may grow/shrink) on writes, so it is unsafe to alias any +# *part* of the container independently. We therefore force *all* access to a +# `DSparseArray` to resolve to the container's stable whole-object +# `ObjectAliasing` -- including access via views, transposes, adjoints, and +# reshapes -- rather than e.g. a `StridedAliasing` over storage that may have +# since moved. `aliases_as_whole` opts the type into this behavior, and Datadeps' +# `aliasing_root` resolves any wrapper of a `DSparseArray` back to the container +# before computing aliasing. Note: this must return a *bare* `ObjectAliasing`, as +# `aliases_as_whole(::ObjectAliasing)` relies on it to drive whole-object copies. +# +# The ObjectAliasing is stamped with the inner storage's memory space so GPU- +# resident sparse tiles (e.g. CuSparseMatrixCSC / ROCSparseMatrixCSC, or +# DeviceSparseMatrixCSC) are tracked in VRAM rather than host RAM. The pointer +# itself is still the host wrapper's `pointer_from_objref` (identity only); +# Datadeps never span-copies these objects (see `aliases_as_whole` / `FullCopy`). +function aliasing(M::DSparseArray, _=identity) + space = value_memory_space(M.mat) + ptr = RemotePtr{Cvoid}(UInt(pointer_from_objref(M)), space) + return ObjectAliasing(ptr, sizeof(typeof(M))) +end +aliases_as_whole(::DSparseArray) = true +# Chunk space follows the inner sparse storage (host CSC → CPURAM, GPU sparse → VRAM). +# `memory_space` (not only `value_memory_space`) must resolve correctly: Datadeps' +# `aliased_object!` uses `memory_space(x)` when deciding whether a same-space +# slot can share the original object. The generic fallback labels unknown values +# as CPURAM, which would force a discarded GPU copy and drop SpGEMM writes. +value_memory_space(M::DSparseArray) = value_memory_space(M.mat) +memory_space(M::DSparseArray) = value_memory_space(M) + +# A `DSparseArray` has no meaningful raw data pointer (its storage may be +# reallocated/resized). This trap ensures that if aliasing ever tries to treat a +# wrapper of a `DSparseArray` as strided memory -- e.g. a new array-wrapper type +# that `aliasing_root` failed to resolve via `Base.parent` -- it fails loudly +# instead of silently corrupting Datadeps' aliasing tracking. +function Base.pointer(::DSparseArray) + throw(ArgumentError("`pointer(::DSparseArray)` is intentionally unsupported: \ + a DSparseArray may reallocate its storage, so it must be aliased as a \ + whole object via `Dagger.aliasing`. If you reached here through Datadeps \ + aliasing of an array wrapper, ensure that wrapper implements `Base.parent` \ + so that `Dagger.aliasing_root` can resolve it to the DSparseArray.")) +end + +# Forward indexing to the inner array. Ranges/colons are supported so that +# views (used by copy-buffering between mismatched partitionings) work. +Base.getindex(M::DSparseArray, I...) = getindex(M.mat, I...) +Base.setindex!(M::DSparseArray, v, I...) = setindex!(M.mat, v, I...) + +# Whole-tile copy. This is how Datadeps moves a tile between workers via +# `move!`: we *reallocate* the inner storage rather than mutating in place, +# since sparse storage cannot generally be updated in place. The wrapper's +# identity (and thus its aliasing) is preserved. +function Base.copyto!(dst::DSparseArray, src::DSparseArray) + dst.mat = _sparse_copy(src.mat) + return dst +end +# Backend-overridable deep copy of the inner storage (Finch tensors don't define +# `Base.copy`). +_sparse_copy(mat) = copy(mat) + +# Adapt preserves the `DSparseArray` wrapper while adapting the inner storage +# (e.g. SparseMatrixCSC → CuSparseMatrixCSC under `adapt(CuArray, …)`). GPU +# package extensions override `move` for sparse types so OpenCL/Metal/oneAPI do +# not densify via the generic `adapt(DeviceArray, SparseMatrixCSC)` path. +Adapt.adapt_structure(to, M::DSparseArray) = DSparseArray(Adapt.adapt(to, M.mat)) + +# Whole-object in-place move for Datadeps `FullCopy`. Reassigns `to.mat` from a +# space-converted copy of `from.mat` rather than span-copying storage. +function move!(to_space::MemorySpace, from_space::MemorySpace, to::DSparseArray, from::DSparseArray) + if to_space == from_space + to.mat = _sparse_copy(from.mat) + else + to_proc = first(processors(to_space)) + from_proc = first(processors(from_space)) + moved = move(from_proc, to_proc, from) + to.mat = moved.mat + end + return +end + +# Wrapping hook used when materializing tiles (e.g. in `distribute`). Backends +# overload this (e.g. in package extensions) to wrap freshly-created tiles in a +# container that Datadeps can track (such as `DSparseArray` for sparse tiles); +# everything else is left untouched. +maybe_wrap_tile(x) = x + +# Partial-range tile copy used by copy-buffering (e.g. when matmul operands need +# to be repartitioned). Some backends (Finch) forbid `setindex!`, so we route +# through a backend hook that returns the (possibly reallocated) inner storage. +# `DSparseArray` hides the reallocation from Datadeps. +function copyto_view!(Bpart::DSparseArray, Brange, Apart, Arange) + Bpart.mat = _sparse_copyto_view!(Bpart.mat, Brange, view(Apart, Arange)) + return +end +# Default: in-place element-wise copy (works for `setindex!`-capable backends +# such as `SparseArrays`). May reallocate and return new storage for backends +# that can't update in place, so callers must use the returned value. +function _sparse_copyto_view!(mat, Brange, src) + copyto!(view(mat, Brange), src) + return mat +end + +#Base.BroadcastStyle(::Type{<:DSparseArray}) = FinchStyle() + +LinearAlgebra.norm2(M::DSparseArray) = LinearAlgebra.norm2(M.mat) + +# Matrix-specific operations dispatch on the 2-D alias `DSparseMatrix`. +function matmatmul!( + C::DSparseMatrix, + transA::Char, + transB::Char, + A::DSparseMatrix, + B::DSparseMatrix, + alpha, + beta +) + # Forward to sparse matmatmul!, but allow in-place update of C + return matmatmul!(C, transA, transB, A.mat, B.mat, alpha, beta) +end + +# Sparse matrix-vector multiply: `C = alpha*op(A)*B + beta*C` with a sparse +# matrix tile `A` and dense vectors `B`/`C`. Unwrap to the inner storage so each +# backend can provide an efficient (in-place on dense `C`) method. +function matvecmul!(C, transA::Char, A::DSparseMatrix, B, alpha, beta) + return matvecmul!(C, transA, A.mat, B, alpha, beta) +end + +# Factorize via host CSC so GPU sparse tiles (rocSPARSE / DeviceSparseMatrixCSC) +# can still feed block-Jacobi; CPU SparseMatrixCSC is a no-op collect. +_factorize_tile(M::DSparseArray) = LinearAlgebra.lu(_sparse_collect(M.mat)) + +# Expose the inner sparse storage so preconditioner backends (ILU, AMG) build on +# the actual `SparseMatrixCSC` rather than the `DSparseArray` wrapper. GPU mats +# are gathered to host CSC (those backends are host-only). +_tile_matrix(M::DSparseArray) = _sparse_collect(M.mat) + +# Jacobi diagonal extract without GPU scalar indexing. `diag(::SparseMatrixCSC)` +# returns a `SparseVector`; densify before `copyto!` into a device vector. +function _set_inv_diag!(out, tile::DSparseArray) + d = Vector(LinearAlgebra.diag(_sparse_collect(tile.mat))) + copyto!(out, inv.(d)) + return nothing +end + +function transpose_tile end +function copytile!(A::DSparseMatrix, B::DSparseMatrix) + A.mat = transpose_tile(B.mat) + return +end +function copydiagtile!(A::DSparseMatrix, uplo) + A.mat = transpose_tile(A.mat, uplo) + return +end + +#============================================================================== + Device-resident CSC for GPU backends without a vendor sparse library + (OpenCL / Metal / oneAPI). Storage vectors live in device memory; SpGEMM/SpMV + kernels fall back to host SparseArrays (see SparseArraysExt). CUDA/ROCm use + their native CuSparse/ROCSparse types instead. +==============================================================================# + +""" + DeviceSparseMatrixCSC{Tv,Ti,V,I} + +A CSC sparse matrix whose `colptr`/`rowval`/`nzval` live in device arrays +(`CLArray`, `MtlArray`, `oneArray`, …). Used as the inner storage of a +[`DSparseArray`](@ref) on GPU backends that lack a vendor sparse library. +""" +struct DeviceSparseMatrixCSC{Tv,Ti,V<:AbstractVector{Tv},I<:AbstractVector{Ti}} <: AbstractMatrix{Tv} + m::Int + n::Int + colptr::I + rowval::I + nzval::V +end +Base.size(A::DeviceSparseMatrixCSC) = (A.m, A.n) +Base.eltype(::DeviceSparseMatrixCSC{Tv}) where Tv = Tv +Base.IndexStyle(::Type{<:DeviceSparseMatrixCSC}) = IndexCartesian() +function Base.getindex(A::DeviceSparseMatrixCSC{Tv}, i::Integer, j::Integer) where Tv + @boundscheck checkbounds(A, i, j) + # Scalar device indexing; acceptable for rare host-side inspection. + col_start = Int(A.colptr[j]) + col_end = Int(A.colptr[j + 1]) - 1 + for p in col_start:col_end + Int(A.rowval[p]) == i && return A.nzval[p] + end + return zero(Tv) +end +# Memory space follows the nonzero values buffer. +value_memory_space(A::DeviceSparseMatrixCSC) = value_memory_space(A.nzval) +memory_space(A::DeviceSparseMatrixCSC) = value_memory_space(A) +# Whole-object aliasing (storage may be replaced on write); stamp with device space. +function aliasing(A::DeviceSparseMatrixCSC, _=identity) + space = value_memory_space(A) + ptr = RemotePtr{Cvoid}(UInt(pointer_from_objref(A)), space) + return ObjectAliasing(ptr, sizeof(typeof(A))) +end +aliases_as_whole(::DeviceSparseMatrixCSC) = true +function Base.pointer(::DeviceSparseMatrixCSC) + throw(ArgumentError("`pointer(::DeviceSparseMatrixCSC)` is unsupported; \ + alias as a whole object via `Dagger.aliasing`.")) +end + +# Allocation hooks used by `spzeros`/`sprand` under a GPU compute scope. +# GPU Exts override these for their processor type. +allocate_sparse_zeros(proc::Processor, ::Type{T}, dims::Dims) where T = + allocate_sparse_zeros_default(proc, T, dims) +allocate_sparse_rand(proc::Processor, ::Type{T}, dims::Dims, sparsity::AbstractFloat) where T = + allocate_sparse_rand_default(proc, T, dims, sparsity) +# Defaults are filled in by SparseArraysExt (host SparseMatrixCSC / SparseVector). +function allocate_sparse_zeros_default end +function allocate_sparse_rand_default end +# Upload host CSC onto device vectors of type `Arr` (filled in by SparseArraysExt). +function device_sparse_from_host end \ No newline at end of file diff --git a/src/array/sparse_partition.jl b/src/array/sparse_partition.jl deleted file mode 100644 index 26efe3e6a..000000000 --- a/src/array/sparse_partition.jl +++ /dev/null @@ -1,37 +0,0 @@ -export partition_sparse - -function partition_sparse(colptr, nnz, sz, nparts) - nnz_per_chunk = nnz/nparts - dom = ArrayDomain(map(x->1:x, sz)) - - nxt = nnz_per_chunk - colstart = 1 - splits = UnitRange[] - - for i=1:nparts-1 - idxs = searchsorted(colptr, nxt) - if length(idxs) == 0 - # the exact index was not found. latch on to the nearest column boundary - idx = first(idxs) - if abs(colptr[idx]-nxt) < abs(colptr[idx-1]-nxt) - nextidx = idx - else - nextidx = idx-1 - end - else - nextidx = last(idxs) - end - push!(splits, colstart:nextidx) - colstart=nextidx+1 - nxt=nxt + nnz_per_chunk - end - push!(splits, colstart:(length(colptr)-1)) - - - cumlength = cumsum(map(length, splits)) - subdomains = DomainBlocks((1,1), ([size(dom, 1)], cumlength)) - DomainSplit(dom, subdomains) -end - -partition_sparse(S::SparseMatrixCSC, nparts) = - partition_sparse(S.colptr, length(S.nzval), size(S), nparts) diff --git a/src/array/sparsedirect.jl b/src/array/sparsedirect.jl new file mode 100644 index 000000000..508b24ac4 --- /dev/null +++ b/src/array/sparsedirect.jl @@ -0,0 +1,518 @@ +# Sparse direct solvers. +# +# Pure-Julia sparse LU backends (KLU via `PureKLU.jl`, UMFPACK-style via +# `PureUMFPACK.jl`) provide *direct* solves for sparse systems. Unlike the +# C-bound `SparseArrays.UmfpackLU` (whose factorization holds process-local +# UMFPACK resources and can't be moved between workers), these factorizations are +# plain Julia data: serializable and movable, so Dagger can schedule them freely. +# +# Two entry points are exposed (implemented in `ext/PureKLUExt.jl` / +# `ext/PureUMFPACKExt.jl`): +# +# - A whole-matrix direct solve: `Dagger.klu(A)` / `Dagger.splu(A)` gather the +# (sparse) `DMatrix` onto a single worker (the one owning the most tiles), +# assemble one `SparseMatrixCSC`, factor it once, and return a +# `DaggerSparseLU` whose factor is pinned there. Solves move only the RHS / +# solution vectors. This is a correctness-first, single-worker factorization +# for systems that fit on one worker (e.g. coarse-grid solves). +# - An opt-in distributed solve path: `Dagger.splu(A; distributed=true)` +# (PureUMFPACK only) reuses that factorization, extracts the triangular +# factors into sparse `DMatrix`es, and solves with a tiled datadeps +# forward/backward substitution (`DistributedSparseLU`, `method=:trsv`). +# - A Schur-complement domain-decomposition path: +# `Dagger.splu(A; distributed=true, method=:schur)` factors interior blocks +# in parallel across workers and solves via a separator Schur system +# (`DistributedSchurLU`; requires Metis.jl). +# - Per-tile block direct solves (`BlockKLUPreconditioner` / +# `BlockUMFPACKPreconditioner`, in `iterativesolvers.jl`), which reuse the +# block-preconditioner machinery. + +""" + klu(A::DMatrix; kwargs...) -> DaggerSparseLU + +Direct sparse LU factorization of a sparse `DMatrix` using KLU. Tiles are +gathered onto one worker, assembled into a single `SparseMatrixCSC`, and factored +once; the factor stays pinned there. Solve with `F \\ b` or `ldiv!(x, F, b)` over +`DVector`s. Requires `PureKLU.jl` to be loaded. + +KLU is well suited to unsymmetric systems with near-triangular structure (e.g. +circuit simulation). See also [`splu`](@ref). +""" +function klu end + +""" + splu(A::DMatrix; distributed=false, method=:trsv, blocksize=nothing, nparts=nothing, kwargs...) + +Direct sparse LU factorization of a sparse `DMatrix` using a pure-Julia +UMFPACK-style multifrontal solver. + +- `distributed=false` → [`DaggerSparseLU`](@ref): gather+factor on one worker. +- `distributed=true, method=:trsv` (default) → [`DistributedSparseLU`](@ref): + Stage-4a factor plus tiled datadeps triangular solves (`blocksize` controls + the solve tiling). +- `distributed=true, method=:schur` → [`DistributedSchurLU`](@ref): single-level + Schur-complement domain decomposition (METIS separator; `nparts` defaults to + `max(2, nworkers())`). Requires `Metis.jl`. + +Requires `PureUMFPACK.jl`. See also [`klu`](@ref). +""" +function splu end + +_klu_required() = throw(ArgumentError( + "Dagger.klu requires PureKLU.jl. Run `using PureKLU` to enable sparse direct \ + (KLU) solves.")) +_splu_required() = throw(ArgumentError( + "Dagger.splu requires PureUMFPACK.jl. Run `using PureUMFPACK` to enable sparse \ + direct (UMFPACK) solves.")) +klu(A; kwargs...) = _klu_required() +splu(A; kwargs...) = _splu_required() + +""" + DaggerSparseLU + +A direct sparse factorization of a sparse `DMatrix`, produced by +[`Dagger.klu`](@ref) or [`Dagger.splu`](@ref). The underlying factorization is +pure-Julia and pinned to a single worker (`fact` is a `Chunk`/`DTask` with +`scope`). Solve `A x = b` with `F \\ b` (returns a `DVector` partitioned like +`b`) or `ldiv!(x, F, b)`. +""" +struct DaggerSparseLU{F,S,P} + fact::F # pinned factorization (`Chunk` / `DTask`) + scope::S # `ProcessScope` of the factor's worker + n::Int # system size (square) + part::P # default vector partitioning (column blocks of `A`) +end + +Base.size(F::DaggerSparseLU) = (F.n, F.n) +Base.size(F::DaggerSparseLU, i::Integer) = i <= 2 ? F.n : 1 + +""" + DistributedSparseLU + +A sparse LU factorization with triangular factors stored as sparse `DMatrix`es, +produced by `Dagger.splu(A; distributed=true)`. Solves use a tiled datadeps +forward/backward substitution over `L` and `U` (PureUMFPACK only). +""" +struct DistributedSparseLU{L,U,P,Q,R,Part} + L::L # unit-lower triangular factor (`DMatrix` of sparse tiles) + U::U # upper triangular factor (`DMatrix` of sparse tiles) + p::P # row permutation + q::Q # column permutation + Rs::R # row scaling + n::Int # system size (square) + part::Part # vector partitioning matching `L`/`U` tile rows +end + +Base.size(F::DistributedSparseLU) = (F.n, F.n) +Base.size(F::DistributedSparseLU, i::Integer) = i <= 2 ? F.n : 1 + +# Assemble local tiles into one `SparseMatrixCSC` (no densification). The real +# method is added by `SparseArraysExt`; without `SparseArrays` there are no +# sparse tiles to gather. +function _gather_sparse end +_gather_sparse(args...) = throw(ArgumentError( + "gathering a sparse DMatrix requires SparseArrays to be loaded")) + +# Backend-specific extraction of `(L, U, p, q, Rs)` from a factored object. +# Implemented for `PureUMFPACK.PureLU` in `PureUMFPACKExt`. +function _extract_lu_factors end +_extract_lu_factors(F) = throw(ArgumentError( + "distributed sparse LU requires a PureUMFPACK.PureLU factor (got $(typeof(F))). \ + Use `Dagger.splu(A; distributed=true)` with PureUMFPACK.jl loaded.")) + +# Choose the worker that already owns the most tiles of `A` (minimizes gather +# traffic). Ties break toward the lowest worker id. +function _select_factor_scope(A::DMatrix) + counts = Dict{Int,Int}() + for c in A.chunks + wid = _tile_scope(c).wid + counts[wid] = get(counts, wid, 0) + 1 + end + best_wid, best_count = 0, -1 + for (wid, cnt) in counts + if cnt > best_count || (cnt == best_count && wid < best_wid) + best_wid, best_count = wid, cnt + end + end + return ProcessScope(best_wid) +end + +# Runs on the chosen worker: assemble from tiles (already moved here by the +# scheduler), factorize, and pin the factor to this process. +function _assemble_and_factor_pinned(factorize, ::Type{T}, row_offsets, col_offsets, + m, n, tiles...) where T + S = _gather_sparse(T, tiles, row_offsets, col_offsets, m, n) + F = factorize(S) + proc = task_processor() + return tochunk(F, proc, ProcessScope(root_worker_id(proc))) +end + +# Shared path for `klu` / `splu`: pick a worker, spawn assemble+factor there. +function _spawn_direct_factorization(A::DMatrix{T}, factorize) where T + n = LinearAlgebra.checksquare(A) + scope = _select_factor_scope(A) + Ac = A.chunks + mt, nt = size(Ac) + ntiles = mt * nt + row_offsets = Vector{Int}(undef, ntiles) + col_offsets = Vector{Int}(undef, ntiles) + tiles = Vector{Any}(undef, ntiles) + idx = 1 + for i in 1:mt, j in 1:nt + dom = A.subdomains[i, j] + row_offsets[idx] = first(dom.indexes[1]) - 1 + col_offsets[idx] = first(dom.indexes[2]) - 1 + tiles[idx] = Ac[i, j] + idx += 1 + end + mA, nA = size(A) + part = Blocks(A.partitioning.blocksize[2]) + fact = spawn(_assemble_and_factor_pinned, Options(; compute_scope=scope), + factorize, T, row_offsets, col_offsets, mA, nA, tiles...) + return DaggerSparseLU(fact, scope, n, part) +end + +# Gather RHS chunks, solve against the pinned factor, return the dense solution. +function _direct_solve(fact, bparts...) + b = length(bparts) == 1 ? bparts[1] : reduce(vcat, bparts) + return fact \ b +end + +function Base.:\(F::DaggerSparseLU, b::DVector) + length(b) == F.n || throw(DimensionMismatch( + "factorization is $(F.n)×$(F.n) but b has length $(length(b))")) + # Solve on the factor's worker (factor stays pinned); only the O(n) solution + # returns here to be redistributed like `b`. + x = fetch(spawn(_direct_solve, Options(; compute_scope=F.scope), + F.fact, b.chunks...)) + return distribute(x, b.partitioning) +end + +function LinearAlgebra.ldiv!(x::DVector, F::DaggerSparseLU, b::DVector) + return copyto!(x, F \ b) +end + +# ---- Stage 4b: distributed triangular solves -------------------------------- + +_fetch_lu_factors(fact) = _extract_lu_factors(fact) + +function _resolve_solve_blocksize(A::DMatrix, blocksize) + if blocksize === nothing + return Int(min(A.partitioning.blocksize...)) + end + bs = Int(blocksize) + bs > 0 || throw(ArgumentError("blocksize must be positive, got $bs")) + return bs +end + +# Re-tile PureLU factors into sparse `DMatrix`es after a Stage-4a factorization. +function _make_distributed_sparse_lu(F::DaggerSparseLU, A::DMatrix; blocksize=nothing) + bs = _resolve_solve_blocksize(A, blocksize) + L, U, p, q, Rs = fetch(spawn(_fetch_lu_factors, Options(; compute_scope=F.scope), F.fact)) + part2 = Blocks(bs, bs) + part1 = Blocks(bs) + return DistributedSparseLU(distribute(L, part2), distribute(U, part2), + p, q, Rs, F.n, part1) +end + +# Tile kernels: unwrap `DSparseArray` → `SparseMatrixCSC`, never densify. +function _sparse_unit_ltrsv!(Ltile, x::AbstractVector) + LinearAlgebra.ldiv!(UnitLowerTriangular(_tile_matrix(Ltile)), x) + return +end + +function _sparse_utrsv!(Utile, x::AbstractVector) + LinearAlgebra.ldiv!(UpperTriangular(_tile_matrix(Utile)), x) + return +end + +function _sparse_gemv_sub!(Atile, x::AbstractVector, y::AbstractVector) + matvecmul!(y, 'N', Atile, x, -one(eltype(y)), one(eltype(y))) + return +end + +# Blocked forward substitution for unit-lower `L` (mirrors `trsv!` lower/'N'). +function _sparse_trsv_forward!(L::DMatrix, c::DVector) + Lc, Cc = L.chunks, c.chunks + nt = length(Cc) + size(Lc, 1) == nt && size(Lc, 2) == nt || throw(DimensionMismatch( + "L tile grid $(size(Lc)) incompatible with vector tiles $nt")) + Dagger.spawn_datadeps() do + for k in 1:nt + Dagger.@spawn _sparse_unit_ltrsv!(In(Lc[k, k]), InOut(Cc[k])) + for i in (k + 1):nt + Dagger.@spawn _sparse_gemv_sub!(In(Lc[i, k]), In(Cc[k]), InOut(Cc[i])) + end + end + end + return c +end + +# Blocked backward substitution for upper `U` (mirrors `trsv!` upper/'N'). +function _sparse_trsv_backward!(U::DMatrix, c::DVector) + Uc, Cc = U.chunks, c.chunks + nt = length(Cc) + size(Uc, 1) == nt && size(Uc, 2) == nt || throw(DimensionMismatch( + "U tile grid $(size(Uc)) incompatible with vector tiles $nt")) + Dagger.spawn_datadeps() do + for k in reverse(1:nt) + Dagger.@spawn _sparse_utrsv!(In(Uc[k, k]), InOut(Cc[k])) + for i in 1:(k - 1) + Dagger.@spawn _sparse_gemv_sub!(In(Uc[i, k]), In(Cc[k]), InOut(Cc[i])) + end + end + end + return c +end + +function Base.:\(F::DistributedSparseLU, b::DVector) + length(b) == F.n || throw(DimensionMismatch( + "factorization is $(F.n)×$(F.n) but b has length $(length(b))")) + # Permute + row-scale (driver O(n)); matches PureUMFPACK `_solve_factor`. + b_local = collect(b) + T = promote_type(eltype(b_local), eltype(F.Rs), eltype(F.U)) + c_local = Vector{T}(undef, F.n) + p, q, Rs = F.p, F.q, F.Rs + @inbounds for k in eachindex(p) + c_local[k] = Rs[p[k]] * b_local[p[k]] + end + c = distribute(c_local, F.part) + _sparse_trsv_forward!(F.L, c) + _sparse_trsv_backward!(F.U, c) + c_local = collect(c) + x_local = similar(c_local) + @inbounds for k in eachindex(q) + x_local[q[k]] = c_local[k] + end + return distribute(x_local, b.partitioning) +end + +function LinearAlgebra.ldiv!(x::DVector, F::DistributedSparseLU, b::DVector) + return copyto!(x, F \ b) +end + +# ---- Stage 4c: Schur-complement domain decomposition ------------------------- +# +# Single-level (non-recursive) k-way Schur factorization. Recursive / nested +# dissection of the separator is future work. + +""" + DistributedSchurLU + +Sparse LU via single-level Schur-complement domain decomposition, produced by +`Dagger.splu(A; distributed=true, method=:schur)`. Interior blocks are factored +in parallel (pinned round-robin across workers); the separator Schur system is +factored on one worker. Solve with `F \\ b` / `ldiv!(x, F, b)`. +""" +struct DistributedSchurLU{F,W,G,SF,S,P} + interior_facts::Vector{F} # pinned `Aᵢᵢ` factors + W::Vector{W} # pinned `Wᵢ = Aᵢᵢ⁻¹ AᵢΓ` (dense columns) + AΓi::Vector{G} # pinned `A_Γi` blocks + interior_scopes::Vector{S} + Sfact::SF # pinned Schur factor (`nothing` if `|Γ|==0`) + Sscope::S + interiors::Vector{Vector{Int}} + separator::Vector{Int} + perm::Vector{Int} + n::Int + part::P +end + +Base.size(F::DistributedSchurLU) = (F.n, F.n) +Base.size(F::DistributedSchurLU, i::Integer) = i <= 2 ? F.n : 1 + +# METIS partition + separator (implemented in `MetisExt`). +function _nested_dissection_partition end +_nested_dissection_partition(A, nparts) = throw(ArgumentError( + "Dagger.splu(...; method=:schur) requires Metis.jl. Run `using Metis` to \ + enable Schur-complement distributed sparse LU.")) + +# Gather a sparse `DMatrix` into one `SparseMatrixCSC` on the caller (pattern +# assembly for the partitioner; no densification). +function _collect_sparse_dmatrix(A::DMatrix{T}) where T + Ac = A.chunks + mt, nt = size(Ac) + ntiles = mt * nt + row_offsets = Vector{Int}(undef, ntiles) + col_offsets = Vector{Int}(undef, ntiles) + tiles = Vector{Any}(undef, ntiles) + idx = 1 + for i in 1:mt, j in 1:nt + dom = A.subdomains[i, j] + row_offsets[idx] = first(dom.indexes[1]) - 1 + col_offsets[idx] = first(dom.indexes[2]) - 1 + tiles[idx] = fetch(Ac[i, j]) + idx += 1 + end + mA, nA = size(A) + return _gather_sparse(T, tiles, row_offsets, col_offsets, mA, nA) +end + +_schur_worker_ids() = (w = workers(); isempty(w) ? Int[myid()] : Int[w...]) + +# Factor `Aᵢᵢ`, form dense `Wᵢ = Aᵢᵢ⁻¹ AᵢΓ` (column solves; densifying the thin +# RHS is simpler than sparse multi-RHS), and local Schur contribution `A_Γi Wᵢ`. +# Pins `(fact, W, A_Γi)` to the current worker. +function _schur_factor_interior_pinned(factorize, Aii, AiΓ, AΓi) + F = factorize(Aii) + ni, ng = size(AiΓ) + T = eltype(AiΓ) + W = zeros(T, ni, ng) + for j in 1:ng + W[:, j] = F \ Vector{T}(AiΓ[:, j]) + end + Si = AΓi * W + proc = task_processor() + scope = ProcessScope(root_worker_id(proc)) + return (tochunk(F, proc, scope), tochunk(W, proc, scope), + tochunk(AΓi, proc, scope), Si) +end + +function _schur_factor_interface_pinned(factorize, S) + F = factorize(S) + proc = task_processor() + scope = ProcessScope(root_worker_id(proc)) + return tochunk(F, proc, scope) +end + +# Forward: `y = Aᵢᵢ⁻¹ bᵢ`, `contrib = A_Γi y`. Pin `y` for the back-substitution. +function _schur_forward_pinned(fact, AΓi, bI) + y = fact \ bI + contrib = AΓi * y + proc = task_processor() + scope = ProcessScope(root_worker_id(proc)) + return tochunk(y, proc, scope), contrib +end + +_schur_back_pinned(y, W, xΓ) = y - W * xΓ + +_schur_interface_solve(Sfact, g) = Sfact \ g + +function _make_schur_sparse_lu(A::DMatrix, factorize; nparts=nothing) + n = LinearAlgebra.checksquare(A) + np = nparts === nothing ? max(2, nworkers()) : Int(nparts) + np >= 1 || throw(ArgumentError("nparts must be ≥ 1, got $np")) + + Asp = _collect_sparse_dmatrix(A) + nd = _nested_dissection_partition(Asp, np) + interiors = nd.interiors + separator = nd.separator + perm = nd.perm + + # Trivial partition → Stage-4a serial factorization. + if length(interiors) <= 1 && isempty(separator) + return _spawn_direct_factorization(A, factorize) + end + + wids = _schur_worker_ids() + k = length(interiors) + interior_facts = Vector{Any}(undef, k) + Ws = Vector{Any}(undef, k) + AΓis = Vector{Any}(undef, k) + scopes = Vector{ProcessScope}(undef, k) + Si_tasks = Vector{Any}(undef, k) + + Γ = separator + ng = length(Γ) + for i in 1:k + Ii = interiors[i] + scope = ProcessScope(wids[mod1(i, length(wids))]) + scopes[i] = scope + Aii = Asp[Ii, Ii] + AiΓ = Asp[Ii, Γ] + AΓi = Asp[Γ, Ii] + task = spawn(_schur_factor_interior_pinned, Options(; compute_scope=scope), + factorize, Aii, AiΓ, AΓi) + Si_tasks[i] = task + end + + # Fetch pinned chunks + local Schur contributions; reduce S on the driver. + T = eltype(Asp) + S = isempty(Γ) ? zeros(T, 0, 0) : Matrix{T}(Asp[Γ, Γ]) + for i in 1:k + Fi, Wi, AΓi_c, Si = fetch(Si_tasks[i]) + interior_facts[i] = Fi + Ws[i] = Wi + AΓis[i] = AΓi_c + isempty(Γ) || (S .-= Si) + end + + if isempty(Γ) + Sfact = nothing + Sscope = scopes[1] + else + Sscope = ProcessScope(wids[1]) + # Schur fill-in is expected; hand a sparse view to the backend factorizer. + Ssp = _sparse_copy_of(S) + Sfact = fetch(spawn(_schur_factor_interface_pinned, + Options(; compute_scope=Sscope), factorize, Ssp)) + end + + part = Blocks(A.partitioning.blocksize[2]) + return DistributedSchurLU(interior_facts, Ws, AΓis, scopes, Sfact, Sscope, + interiors, separator, perm, n, part) +end + +# Convert a dense Schur matrix to the sparse type the backend expects. Defined +# as a hook so the core package need not depend on SparseArrays; the real method +# lives in `SparseArraysExt`. +function _sparse_copy_of end +_sparse_copy_of(S) = throw(ArgumentError( + "Schur-complement sparse LU requires SparseArrays to be loaded")) + +function Base.:\(F::DistributedSchurLU, b::DVector) + length(b) == F.n || throw(DimensionMismatch( + "factorization is $(F.n)×$(F.n) but b has length $(length(b))")) + b_local = collect(b) + bp = b_local[F.perm] + k = length(F.interiors) + ng = length(F.separator) + + # Split permuted RHS into interior segments + separator. + ys = Vector{Any}(undef, k) + contrib_tasks = Vector{Any}(undef, k) + off = 0 + for i in 1:k + ni = length(F.interiors[i]) + bI = bp[off+1:off+ni] + off += ni + contrib_tasks[i] = spawn(_schur_forward_pinned, + Options(; compute_scope=F.interior_scopes[i]), + F.interior_facts[i], F.AΓi[i], bI) + end + + g = ng == 0 ? similar(bp, 0) : copy(bp[off+1:end]) + for i in 1:k + yi, contrib = fetch(contrib_tasks[i]) + ys[i] = yi + ng == 0 || (g .-= contrib) + end + + xΓ = if ng == 0 + similar(g) + else + fetch(spawn(_schur_interface_solve, Options(; compute_scope=F.Sscope), + F.Sfact, g)) + end + + x_parts = Vector{Any}(undef, k) + back_tasks = Vector{Any}(undef, k) + for i in 1:k + back_tasks[i] = spawn(_schur_back_pinned, + Options(; compute_scope=F.interior_scopes[i]), + ys[i], F.W[i], xΓ) + end + for i in 1:k + x_parts[i] = fetch(back_tasks[i]) + end + + xp = vcat(x_parts..., xΓ) + x_local = similar(b_local) + x_local[F.perm] = xp + return distribute(x_local, b.partitioning) +end + +function LinearAlgebra.ldiv!(x::DVector, F::DistributedSchurLU, b::DVector) + return copyto!(x, F \ b) +end + diff --git a/src/array/stencil.jl b/src/array/stencil.jl index 0f3af7f22..c0c9e272b 100644 --- a/src/array/stencil.jl +++ b/src/array/stencil.jl @@ -57,18 +57,6 @@ function load_neighbor_region(arr, region_code::NTuple{N,Int}, neigh_dist) where return move(task_processor(), copy(@view arr[start_idx:stop_idx])) end -# In-place variant: load region directly into a pre-allocated destination buffer. -function load_neighbor_region_into!(dest, arr, region_code::NTuple{N,Int}, neigh_dist) where N - validate_neigh_dist(neigh_dist, size(arr)) - start_idx = CartesianIndex(ntuple(N) do i - region_code[i] == -1 ? lastindex(arr, i) - get_neigh_dist(neigh_dist, i) + 1 : firstindex(arr, i) - end) - stop_idx = CartesianIndex(ntuple(N) do i - region_code[i] == +1 ? firstindex(arr, i) + get_neigh_dist(neigh_dist, i) - 1 : lastindex(arr, i) - end) - copyto!(dest, @view arr[start_idx:stop_idx]) -end - is_past_boundary(size, idx) = any(ntuple(i -> idx[i] < 1 || idx[i] > size[i], length(size))) ############################################################################# @@ -142,9 +130,6 @@ boundary_transition(::Wrap, idx, size) = load_boundary_region(::Wrap, arr, region_code, neigh_dist, boundary_dims) = load_neighbor_region(arr, region_code, neigh_dist) -load_boundary_region_into!(dest, ::Wrap, arr, region_code, neigh_dist, boundary_dims) = - load_neighbor_region_into!(dest, arr, region_code, neigh_dist) - function boundary_source_index(::Wrap, arr, rc, nd, idx_d, d) if rc == -1 return lastindex(arr, d) - nd + idx_d @@ -180,9 +165,6 @@ function load_boundary_region(pad::Pad, arr, region_code::NTuple{N,Int}, neigh_d return move(task_processor(), result) end -load_boundary_region_into!(dest, pad::Pad, arr, region_code, neigh_dist, boundary_dims) = - fill!(dest, pad.padval) - # Use edge as source index (value will be overridden by apply_boundary_value) boundary_source_index(::Pad, arr, rc, nd, idx_d, d) = rc == -1 ? firstindex(arr, d) : (rc == +1 ? lastindex(arr, d) : idx_d) @@ -247,10 +229,6 @@ function load_boundary_region(::Clamp, arr, region_code::NTuple{N,Int}, neigh_di return move(task_processor(), result) end -function load_boundary_region_into!(dest, ::Clamp, arr, region_code::NTuple{N,Int}, neigh_dist, boundary_dims::NTuple{N,Bool}) where N - Kernel(load_boundary_region_kernel)(Clamp(), dest, arr, region_code, neigh_dist, boundary_dims; ndrange=length(dest)) -end - function boundary_source_index(::Clamp, arr, rc, nd, idx_d, d) if rc == -1 return firstindex(arr, d) @@ -362,18 +340,6 @@ function load_boundary_region(::LinearExtrapolate, arr::AbstractArray{T}, region return move(task_processor(), result) end -function load_boundary_region_into!(dest, ::LinearExtrapolate, arr::AbstractArray{T}, region_code::NTuple{N,Int}, neigh_dist, boundary_dims::NTuple{N,Bool}) where {T<:Real,N} - extrap_dim = 0 - for d in 1:N - if boundary_dims[d] && region_code[d] != 0 - extrap_dim = d - break - end - end - nd = get_neigh_dist(neigh_dist, extrap_dim) - Kernel(load_boundary_region_kernel)(LinearExtrapolate(), dest, arr, region_code, neigh_dist, boundary_dims, Val(extrap_dim), Val(nd); ndrange=length(dest)) -end - # Use edge as source index (value will be computed by apply_boundary_value) boundary_source_index(::LinearExtrapolate, arr, rc, nd, idx_d, d) = rc == -1 ? firstindex(arr, d) : (rc == +1 ? lastindex(arr, d) : idx_d) @@ -476,41 +442,6 @@ function load_boundary_region(::Reflect{Symm}, arr, region_code::NTuple{N,Int}, return region end -function load_boundary_region_into!(dest, ::Reflect{Symm}, arr, region_code::NTuple{N,Int}, neigh_dist, boundary_dims::NTuple{N,Bool}) where {N, Symm} - flipped_code = ntuple(N) do i - (region_code[i] != 0 && boundary_dims[i]) ? -region_code[i] : region_code[i] - end - skip = Symm ? 0 : 1 - start_idx = CartesianIndex(ntuple(N) do i - needs_skip = boundary_dims[i] && region_code[i] != 0 - actual_skip = needs_skip ? skip : 0 - if flipped_code[i] == -1 - lastindex(arr, i) - get_neigh_dist(neigh_dist, i) + 1 - actual_skip - elseif flipped_code[i] == +1 - firstindex(arr, i) + actual_skip - else - firstindex(arr, i) - end - end) - stop_idx = CartesianIndex(ntuple(N) do i - needs_skip = boundary_dims[i] && region_code[i] != 0 - actual_skip = needs_skip ? skip : 0 - if flipped_code[i] == +1 - firstindex(arr, i) + get_neigh_dist(neigh_dist, i) - 1 + actual_skip - elseif flipped_code[i] == -1 - lastindex(arr, i) - actual_skip - else - lastindex(arr, i) - end - end) - copyto!(dest, @view arr[start_idx:stop_idx]) - for i in 1:N - GPUArraysCore.@allowscalar if region_code[i] != 0 && boundary_dims[i] - reverse!(dest, dims=i) - end - end -end - function boundary_source_index(::Reflect{Symm}, arr, rc, nd, idx_d, d) where Symm skip = Symm ? 0 : 1 if rc == -1 @@ -641,10 +572,6 @@ function load_boundary_region(boundary::Tuple, arr, region_code::NTuple{N,Int}, return move(task_processor(), result) end -function load_boundary_region_into!(dest, boundary::Tuple, arr, region_code::NTuple{N,Int}, neigh_dist, boundary_dims::NTuple{N,Bool}) where N - Kernel(load_boundary_region_kernel)(boundary, dest, arr, region_code, neigh_dist, boundary_dims; ndrange=length(dest)) -end - ############################################################################# # Chunk Selection and Halo Building ############################################################################# @@ -854,97 +781,107 @@ function select_neighborhood_info(chunks, idx, neigh_dist, boundary) return region_metadata, neighbor_chunks end -# Per-thread cache: WeakKeyDict{DArray, Dict{(chunk_idx, halo_width), HaloArray}}. -# WeakKeyDict is used for the outer level so that the cache does not hold a strong reference -# to the source DArray — allowing its GC finalizer to fire when user code drops its last -# reference (see below). Using chunk_idx as part of the inner key ensures that within one -# DArray, every chunk has its own dedicated buffer — so if a single worker thread processes -# multiple same-shaped chunks in the same iteration sequentially, each gets a distinct -# HaloArray and there is no aliasing with a concurrently running inner-stencil task. -# Filling a cached buffer in-place is safe because spawn_datadeps blocks until all inner -# tasks complete before the next iteration's build_halo_consolidated calls run. -const HALO_ARRAY_CACHE = TaskLocalValue{WeakKeyDict{Any,Dict{Any,Any}}}(()->WeakKeyDict{Any,Dict{Any,Any}}()) - -# Called on the main task (outside any Dagger.@spawn) to get or create the per-DArray inner -# cache dict. Keeping all WeakKeyDict operations here — rather than inside spawned tasks — -# avoids every pitfall of passing a DArray to @spawn (Dagger resolving it to a ROCArray on -# GPU, serialization producing a copy with a different objectid under Distributed, etc.). -# A finalizer registered on first encounter frees all cached HaloArrays when the DArray is -# collected; the WeakKeyDict ensures the cache itself does not prevent that collection. -function get_halo_inner_cache(read_darray) - outer_cache = HALO_ARRAY_CACHE[] - if !haskey(outer_cache, read_darray) - inner_cache = Dict{Any,Any}() - outer_cache[read_darray] = inner_cache - finalizer(read_darray) do _ - # GC finalizers cannot yield (no task switches allowed), but unsafe_free! on - # GPU/distributed Chunks calls remotecall_fetch which acquires locks and yields. - # Defer cleanup to a scheduled task so it runs outside the finalizer context. - errormonitor(Threads.@spawn begin - for halo in values(inner_cache) - try - unsafe_free!(halo) - catch e - # fill_halo_inplace! runs inside spawn_datadeps and may be placed - # on a Distributed worker; if that worker has already exited by - # cleanup time, skip the free. Use nameof to work with both - # Distributed and DistributedNext backends. - string(nameof(typeof(e))) == "ProcessExitedException" || rethrow(e) - end - end - end) - end - end - return outer_cache[read_darray] +function build_halo(neigh_dist, boundary, center, all_halos...; own_center::Bool=false) + N = ndims(center) + expected_halos = 3^N - 1 + @assert length(all_halos) == expected_halos "Halo mismatch: N=$N expected $expected_halos halos, got $(length(all_halos))" + center_data = own_center ? copy(center) : center + return HaloArray(center_data, (all_halos...,), ntuple(i->get_neigh_dist(neigh_dist, i), N); own_center) end -# Cache-miss path: allocate an empty HaloArray using similar() so the buffers are created -# on the same device as the center chunk (GPU or CPU). No filling — that happens inside -# spawn_datadeps via fill_halo_inplace!. -function alloc_halo(neigh_dist, center) - N = ndims(center) - center_size = size(center) - halo_width = ntuple(i -> get_neigh_dist(neigh_dist, i), N) - codes = all_region_codes(Val(N)) - halos = ntuple(length(codes)) do i - region_size = halo_region_size(center_size, halo_width, codes[i]) - similar(center, region_size...) +############################################################################# +# Fused halo construction +############################################################################# +# +# The `@stencil` fast path builds a chunk's `HaloArray` inside the same task that +# sweeps it, from the neighboring chunks passed as `In` dependencies. Compared to +# filling a cached `HaloArray` in a separate task this avoids (a) a second task per +# chunk per expression and (b) copying the whole center chunk, which for a +# double-buffered stencil is as much memory traffic as the sweep itself. + +struct ViewHalos end +struct CopyHalos end + +# Whether a boundary condition's halo regions are plain slices of a neighboring +# chunk, in which case they can be `view`s rather than materialized copies. +@inline halo_build_style(@nospecialize(boundary)) = CopyHalos() +@inline halo_build_style(::Wrap) = ViewHalos() +@inline halo_build_style(boundary::Tuple) = _combine_halo_style(map(halo_build_style, boundary)) +@inline _combine_halo_style(::Tuple{}) = ViewHalos() +@inline _combine_halo_style(styles::Tuple) = + first(styles) isa ViewHalos ? _combine_halo_style(Base.tail(styles)) : CopyHalos() + +@inline function neighbor_region_view(arr, region_code::NTuple{N,Int}, neigh_dist) where N + ranges = ntuple(Val(N)) do i + nd = get_neigh_dist(neigh_dist, i) + start_i = region_code[i] == -1 ? lastindex(arr, i) - nd + 1 : firstindex(arr, i) + stop_i = region_code[i] == +1 ? firstindex(arr, i) + nd - 1 : lastindex(arr, i) + start_i:stop_i end - return HaloArray(similar(center), halos, halo_width; own_center=true) + return view(arr, ranges...) end -# Cache-hit path: fill an existing HaloArray in-place and return it. No cache operations — -# the HaloArray was looked up and passed by the main task before spawning. -function fill_halo_inplace!(halo::HaloArray, neigh_dist, boundary, center, region_metadata, neighbor_chunks...) - expected_halos = length(region_metadata) - @assert length(neighbor_chunks) == expected_halos - copyto!(halo.center, center) - for i in 1:expected_halos +@inline function _build_fused_halos(::ViewHalos, neigh_dist, boundary, region_metadata, + neighbor_chunks::NTuple{NH,Any}) where NH + return ntuple(Val(NH)) do i + region_code, _, _ = region_metadata[i] + neighbor_region_view(neighbor_chunks[i], region_code, neigh_dist) + end +end +@inline function _build_fused_halos(::CopyHalos, neigh_dist, boundary, region_metadata, + neighbor_chunks::NTuple{NH,Any}) where NH + return ntuple(Val(NH)) do i region_code, is_boundary, boundary_dims = region_metadata[i] chunk = neighbor_chunks[i] if is_boundary - load_boundary_region_into!(halo.halos[i], boundary, chunk, region_code, neigh_dist, boundary_dims) + load_boundary_region(boundary, chunk, region_code, neigh_dist, boundary_dims) else - load_neighbor_region_into!(halo.halos[i], chunk, region_code, neigh_dist) + load_neighbor_region(chunk, region_code, neigh_dist) end end - return halo end -function build_halo(neigh_dist, boundary, center, all_halos...; own_center::Bool=false) - N = ndims(center) - expected_halos = 3^N - 1 - @assert length(all_halos) == expected_halos "Halo mismatch: N=$N expected $expected_halos halos, got $(length(all_halos))" - center_data = own_center ? copy(center) : center - return HaloArray(center_data, (all_halos...,), ntuple(i->get_neigh_dist(neigh_dist, i), N); own_center) +""" + build_fused_halo(neigh_dist, boundary, region_metadata, center, neighbor_chunks...) + +Wraps `center` (used in place, not copied) plus halo regions taken from +`neighbor_chunks` into a `HaloArray`. `region_metadata` is the per-region +`(region_code, is_boundary, boundary_dims)` triple precomputed on the submitting +task by `select_neighborhood_info`. +""" +function build_fused_halo(neigh_dist, boundary, region_metadata, + center::AbstractArray{T,N}, neighbor_chunks::Vararg{Any,NH}) where {T,N,NH} + validate_neigh_dist(neigh_dist, size(center)) + halo_width = ntuple(i -> get_neigh_dist(neigh_dist, i), Val(N)) + halos = _build_fused_halos(halo_build_style(boundary), neigh_dist, boundary, + region_metadata, neighbor_chunks) + return HaloArray(center, halos, halo_width; own_center=false) end -function load_neighborhood(arr::HaloArray{T,N}, idx) where {T,N} - start_idx = idx - CartesianIndex(ntuple(i->arr.halo_width[i], ndims(arr))) - stop_idx = idx + CartesianIndex(ntuple(i->arr.halo_width[i], ndims(arr))) - return @view arr[start_idx:stop_idx] +""" + stencil_source_chunks(read_chunks, write_chunks) -> chunks + +The chunk array a neighborhood read should be taken from. Normally that is +`read_chunks` itself, but when the expression writes back into the same chunks it +reads (`A[idx] = f(@neighbors(A[idx]))`), neighbors must come from a snapshot +taken before any chunk is overwritten. +""" +function stencil_source_chunks(read_chunks, write_chunks) + write_set = IdDict{Any,Nothing}() + for chunk in write_chunks + write_set[chunk] = nothing + end + overlaps = any(chunk -> haskey(write_set, chunk), read_chunks) + overlaps || return read_chunks + snapshot_tasks = map(chunk -> Dagger.@spawn(name="stencil_snapshot", copy(chunk)), read_chunks) + return map(task -> fetch(task; raw=true), snapshot_tasks) end +@inline load_neighborhood(arr::HaloArray{T,N}, idx) where {T,N} = + StencilNeighborhood(arr, idx, arr.halo_width) +@inline load_neighborhood(arr::HaloInterior{T,N}, idx) where {T,N} = + StencilNeighborhood(arr.parent, idx, arr.halo_width) + function inner_stencil!(f, output, read_vars) processor = task_processor() inner_stencil_proc!(processor, f, output, read_vars) @@ -954,8 +891,140 @@ end # Non-KA (for CPUs) function inner_stencil_proc!(::ThreadProc, f, output, read_vars) - for idx in CartesianIndices(output) - f(idx, output, read_vars) + cpu_stencil_sweep!(f, output, read_vars) + return +end + +# Widest halo required by any neighborhood variable, per dimension. Zero in every +# dimension means no variable is accessed through a halo, so the whole chunk can +# take the fast path. +@inline _widen_halo(w::NTuple{N,Int}, ::Any) where N = w +@inline _widen_halo(w::NTuple{N,Int}, A::HaloArray{T,N}) where {T,N} = + ntuple(i -> max(w[i], A.halo_width[i]), Val(N)) +@inline _max_halo_width(vars::Tuple{}, w::NTuple{N,Int}) where N = w +@inline _max_halo_width(vars::Tuple, w::NTuple{N,Int}) where N = + _max_halo_width(Base.tail(vars), _widen_halo(w, first(vars))) + +# Interior stand-ins: HaloArrays become HaloInterior (direct center indexing), +# everything else is passed through untouched. +@inline _interior_var(A::HaloArray) = HaloInterior(A.center, A.halo_width) +@inline _interior_var(A) = A + +""" + cpu_stencil_sweep!(f, output, read_vars) + +Applies `f` at every index of `output`, splitting the chunk into an interior +where all neighborhood accesses land in the center array, and a thin boundary +shell where they may reach into halos. + +The split is what makes stencils fast: the interior (essentially the whole chunk) +runs against plain arrays under `@inbounds @simd` with `f` force-inlined, while +the generic, branch-heavy `HaloArray` path is confined to the shell. +""" +function cpu_stencil_sweep!(f::F, output::AbstractArray{T,N}, read_vars::NamedTuple) where {F,T,N} + w = _max_halo_width(values(read_vars), ntuple(_ -> 0, Val(N))) + ax = axes(output) + + interior_ax = ntuple(Val(N)) do i + (first(ax[i]) + w[i]):(last(ax[i]) - w[i]) + end + if any(isempty, interior_ax) + # Halos are as wide as the chunk itself; there is no interior to split off. + for idx in CartesianIndices(output) + @inline f(idx, output, read_vars) + end + return + end + + interior_vars = map(_interior_var, read_vars) + @inbounds @simd for idx in CartesianIndices(interior_ax) + @inline f(idx, output, interior_vars) + end + + all(iszero, w) && return + + # Boundary shell, decomposed into 2N disjoint slabs: for dimension `d`, the + # low and high slabs span the full extent in dimensions after `d` and only the + # interior extent in dimensions before `d`. + for d in 1:N + for low_side in (true, false) + slab_ax = ntuple(Val(N)) do i + if i < d + interior_ax[i] + elseif i == d + low_side ? (first(ax[i]):(first(ax[i]) + w[i] - 1)) : + ((last(ax[i]) - w[i] + 1):last(ax[i])) + else + first(ax[i]):last(ax[i]) + end + end + for idx in CartesianIndices(slab_ax) + @inline f(idx, output, read_vars) + end + end + end + return +end + +@kernel function _stencil_box_kernel!(f, output, read_vars, offset) + idx = @index(Global, Cartesian) + @inline f(idx + offset, output, read_vars) +end + +# Sweep `f` over one box of `output`. KernelAbstractions indexes an `ndrange` +# from 1, so a sub-box is expressed as a shifted full launch. +@inline function _stencil_box!(f, output, read_vars, ranges::NTuple{N,AbstractUnitRange}) where N + nd = map(length, ranges) + any(iszero, nd) && return + offset = CartesianIndex(ntuple(i -> first(ranges[i]) - 1, Val(N))) + Kernel(_stencil_box_kernel!)(f, output, read_vars, offset; ndrange=nd) + return +end + +""" + gpu_stencil_sweep!(f, output, read_vars) + +Applies `f` at every index of `output` using KernelAbstractions, with the same +interior/boundary split as [`cpu_stencil_sweep!`](@ref). + +The split matters more here than on the CPU. Every `@neighbors` access on the +general path recomputes a region code, branches on it, and remaps the index; a +GPU runs that scalar bookkeeping for all `(2w+1)^N` accesses of every element. +Sweeping the interior against `HaloInterior` stand-ins reduces it to a direct +load from the center array, leaving the general path to the `2N` boundary slabs. +""" +function gpu_stencil_sweep!(f::F, output::AbstractArray{T,N}, read_vars::NamedTuple) where {F,T,N} + w = _max_halo_width(values(read_vars), ntuple(_ -> 0, Val(N))) + ax = axes(output) + full_ax = ntuple(i -> first(ax[i]):last(ax[i]), Val(N)) + + interior_ax = ntuple(Val(N)) do i + (first(ax[i]) + w[i]):(last(ax[i]) - w[i]) + end + if any(isempty, interior_ax) + # Halos are as wide as the chunk itself; there is no interior to split off. + _stencil_box!(f, output, read_vars, full_ax) + return + end + + _stencil_box!(f, output, map(_interior_var, read_vars), interior_ax) + + all(iszero, w) && return + + for d in 1:N + for low_side in (true, false) + slab_ax = ntuple(Val(N)) do i + if i < d + interior_ax[i] + elseif i == d + low_side ? (first(ax[i]):(first(ax[i]) + w[i] - 1)) : + ((last(ax[i]) - w[i] + 1):last(ax[i])) + else + full_ax[i] + end + end + _stencil_box!(f, output, read_vars, slab_ax) + end end return end @@ -1136,130 +1205,111 @@ macro stencil(orig_ex) end new_inner_f = :(($inner_index_var, $inner_write_var, $inner_vars)->$new_inner_ex_body) actual_read_vars = filter(v -> (v != write_var) || (v in keys(neighborhoods)), collect(read_vars)) - new_inner_ex = quote - $inner_vars = (;$(actual_read_vars...)) - $inner_stencil!($new_inner_f, $inner_write_var, $inner_vars) - end - inner_fn = Expr(:->, Expr(:tuple, Expr(:parameters, inner_write_var, actual_read_vars...)), new_inner_ex) - # 2a. For each neighborhood read_var: pre-compute region metadata on the main task - # (to avoid passing DArrays into @spawn), and for cache misses spawn alloc_halo - # outside spawn_datadeps so that HaloArray buffers are allocated on the correct - # device (GPU or CPU) before filling. + # 2a. For each neighborhood read_var, pre-compute on the main task (so no DArray + # is ever passed into @spawn) the chunk array to read neighbors from and, per + # chunk, the region metadata plus the neighboring chunks themselves. # - # fill_halo_inplace! is spawned *inside* spawn_datadeps (step 2c) with explicit - # Read deps on the DArray chunks, so Datadeps automatically enforces the ordering - # between fill tasks and inner stencil tasks — including the cross-chunk case where - # fill_task[C] reads chunk[N] as a neighbor while inner_stencil[N] writes chunk[N]. - # This replaces the explicit wait-barrier that was needed when fills ran outside. - cache_sym_map = Dict{Symbol, NamedTuple}() + # `stencil_source_chunks` substitutes a snapshot when the expression writes back + # into the chunks it reads (`A[idx] = f(@neighbors(A[idx]))`); everywhere else it + # is the identity, and the neighbor chunks are read directly. + neigh_sym_map = Dict{Symbol, NamedTuple}() for read_var in read_vars if read_var in keys(neighborhoods) neigh_dist, boundary = neighborhoods[read_var] - @gensym region_info_table fill_tasks region_meta neighbor_cks halo_width_var inner_cache_var cache_key_var - cache_sym_map[read_var] = (; fill_tasks, halo_width_var, inner_cache_var, cache_key_var, region_info_table) - # Validate and compute halo_width on the main task so errors are immediate. + @gensym region_info_table src_chunks region_meta neighbor_cks + neigh_sym_map[read_var] = (; region_info_table, src_chunks) push!(final_ex.args, :($validate_neigh_dist($neigh_dist, ndims($read_var)))) - push!(final_ex.args, :($halo_width_var = ntuple(i -> $get_neigh_dist($neigh_dist, i), ndims($read_var)))) - push!(final_ex.args, :($inner_cache_var = $get_halo_inner_cache($read_var))) - # Pre-compute region metadata for every chunk on the main task. - push!(final_ex.args, :($region_info_table = Array{Any}(undef, size($chunks($read_var))))) - push!(final_ex.args, :($fill_tasks = Array{$DTask}(undef, size($chunks($read_var))))) + push!(final_ex.args, :($src_chunks = $stencil_source_chunks($chunks($read_var), $chunks($write_var)))) + push!(final_ex.args, :($region_info_table = Array{Any}(undef, size($src_chunks)))) push!(final_ex.args, quote - for $chunk_idx in $CartesianIndices($chunks($read_var)) - ($region_meta, $neighbor_cks) = $select_neighborhood_info($chunks($read_var), $chunk_idx, $neigh_dist, $boundary) + for $chunk_idx in $CartesianIndices($src_chunks) + ($region_meta, $neighbor_cks) = $select_neighborhood_info($src_chunks, $chunk_idx, $neigh_dist, $boundary) $region_info_table[$chunk_idx] = (tuple($region_meta...), $neighbor_cks) - $cache_key_var = ($chunk_idx, $halo_width_var) - if !haskey($inner_cache_var, $cache_key_var) - # Cache miss: allocate empty buffers on the correct device. - $inner_cache_var[$cache_key_var] = Dagger.@spawn name="stencil_alloc_halo" $alloc_halo($neigh_dist, $chunks($read_var)[$chunk_idx]) - end end end) end end - # 2b. Build the inner-stencil @spawn expression. Neighborhood deps reference - # fill_tasks[chunk_idx] populated by the fill loop (see 2c below). - deps_ex = Any[] - if write_var in read_vars - push!(deps_ex, Expr(:kw, inner_write_var, :($ReadWrite($chunks($write_var)[$chunk_idx])))) - else - push!(deps_ex, Expr(:kw, inner_write_var, :($Write($chunks($write_var)[$chunk_idx])))) - end + # 2b. Build the per-chunk task. Neighborhood variables are passed as their own + # chunk followed by their `3^N - 1` neighboring chunks (a runtime-length group, + # hence the positional splat); the task assembles the HaloArray itself. + prologue_exs = Expr[] # runs per chunk on the submitting task + arg_exs = Any[] # positional args of the spawned task + unpack_exs = Expr[] # runs inside the task, binds each read var + local_vars = Any[] # task-local binding holding each read var's chunk data + @gensym task_args arg_offset for read_var in actual_read_vars - if read_var in keys(neighborhoods) - syms = cache_sym_map[read_var] - # Reference fill_tasks[chunk_idx] — populated by the fill loop before the - # stencil loop runs, so it is always assigned when this expression executes. - push!(deps_ex, Expr(:kw, read_var, :($Read($(syms.fill_tasks)[$chunk_idx])))) - else - if read_var != write_var - push!(deps_ex, Expr(:kw, read_var, :($Read($chunks($read_var)[$chunk_idx])))) - end - end - end - inner_spawn_ex = :(Dagger.@spawn name="stencil_inner_fn" $inner_fn(;$(deps_ex...))) - - # Build the fill @spawn expression — one per neighborhood read_var. - fill_spawn_exs = Expr[] - for read_var in read_vars + # N.B. Bind into a gensym rather than `read_var` itself: the task closure is + # nested inside the user's scope, so assigning `read_var` there would rebind + # the user's DArray variable instead of creating a task-local binding. + local_var = gensym(read_var) if read_var in keys(neighborhoods) neigh_dist, boundary = neighborhoods[read_var] - syms = cache_sym_map[read_var] - @gensym _rmt _nck _rn - push!(fill_spawn_exs, quote - ($_rmt, $_nck) = $(syms.region_info_table)[$chunk_idx] - $_rn = map($Read, $_nck) - $(syms.fill_tasks)[$chunk_idx] = Dagger.@spawn name="stencil_fill_halo" $fill_halo_inplace!( - $ReadWrite($(syms.inner_cache_var)[($chunk_idx, $(syms.halo_width_var))]), - $neigh_dist, $boundary, - $Read($chunks($read_var)[$chunk_idx]), - $_rmt, - $_rn... - ) + syms = neigh_sym_map[read_var] + @gensym region_meta neighbor_cks nneighbors + push!(prologue_exs, quote + ($region_meta, $neighbor_cks) = $(syms.region_info_table)[$chunk_idx] + $nneighbors = length($neighbor_cks) + end) + push!(arg_exs, :($Read($(syms.src_chunks)[$chunk_idx]))) + push!(arg_exs, Expr(:..., :(map($Read, $neighbor_cks)))) + push!(unpack_exs, quote + $local_var = $build_fused_halo($neigh_dist, $boundary, $region_meta, + $task_args[$arg_offset + 1], + $task_args[($arg_offset + 2):($arg_offset + 1 + $nneighbors)]...) + $arg_offset += 1 + $nneighbors + end) + elseif read_var != write_var + push!(arg_exs, :($Read($chunks($read_var)[$chunk_idx]))) + push!(unpack_exs, quote + $local_var = $task_args[$arg_offset + 1] + $arg_offset += 1 end) end + push!(local_vars, local_var) + end + write_dep_ex = if write_var in read_vars + :($ReadWrite($chunks($write_var)[$chunk_idx])) + else + :($Write($chunks($write_var)[$chunk_idx])) end - # 2c. Each expression gets its own spawn_datadeps region with TWO separate loops: - # first all fills, then all stencils. This ordering is critical when write_var is - # also a neighborhood read_var: submitting all fill tasks before any stencil tasks - # ensures Datadeps sees fill[C]'s Read(chunk[N]) BEFORE stencil[N]'s - # ReadWrite(chunk[N]), so stencil[N] is correctly ordered after fill[C]. - # (An interleaved single loop would register stencil[C]'s write to chunk[C] before - # fill[C+1] reads chunk[C], inverting the dependency.) - fill_loop_body = Expr(:block, fill_spawn_exs...) + # The kernel body may mention an accessed variable bare (e.g. `length(A)`), which + # would otherwise capture the user's DArray into the spawned closure -- illegal, + # since Datadeps analyzes the closure's captures for aliasing. Shadow each + # accessed name with a parameter bound to that chunk's local data, so the body + # sees the chunk (with halos, where applicable) instead. + shadow_params = copy(actual_read_vars) + shadow_args = copy(local_vars) + if !(write_var in actual_read_vars) + push!(shadow_params, write_var) + push!(shadow_args, inner_write_var) + end + shadow_body = quote + $inner_vars = (;$([Expr(:kw, v, v) for v in actual_read_vars]...)) + $inner_stencil!($new_inner_f, $inner_write_var, $inner_vars) + end + shadow_fn = Expr(:->, Expr(:tuple, shadow_params...), shadow_body) + inner_fn_body = quote + $arg_offset = 0 + $(unpack_exs...) + $shadow_fn($(shadow_args...)) + end + inner_fn = Expr(:->, Expr(:tuple, inner_write_var, Expr(:..., task_args)), inner_fn_body) + inner_spawn_ex = Expr(:block, prologue_exs..., + :(Dagger.@spawn name="stencil_inner_fn" $inner_fn($write_dep_ex, $(arg_exs...)))) + + # 2c. One spawn_datadeps region per expression, one task per chunk. Because the + # region blocks until all of its tasks finish, the next expression's tasks are + # only submitted once this expression has been applied everywhere, which is what + # gives `@stencil` its "all at once" semantics. + push!(final_ex.args, :(Dagger.@record_op :stencil $write_var)) push!(final_ex.args, :(Dagger.spawn_datadeps() do - for $chunk_idx in $CartesianIndices($chunks($write_var)) - $fill_loop_body - end for $chunk_idx in $CartesianIndices($chunks($write_var)) $inner_spawn_ex end end)) - - # 2d. After spawn_datadeps completes, resolve any alloc DTasks in the cache to - # their origin Chunks. Do *not* store `fetch(fill_task)`: Datadeps may have run - # fill on a libc-backed remote slot copy and then `unsafe_free!`'d that copy at - # region end, so the fill task's return value can dangle. The origin Chunk (from - # alloc_halo / a prior iteration) receives the updated data via remainder sync - # and is what must stay in the cache for the next iteration. - for read_var in read_vars - if read_var in keys(neighborhoods) - syms = cache_sym_map[read_var] - @gensym cache_entry - push!(final_ex.args, quote - for $chunk_idx in $CartesianIndices($chunks($read_var)) - $(syms.cache_key_var) = ($chunk_idx, $(syms.halo_width_var)) - $cache_entry = $(syms.inner_cache_var)[$(syms.cache_key_var)] - if $cache_entry isa $DTask - $(syms.inner_cache_var)[$(syms.cache_key_var)] = fetch($cache_entry; raw=true) - end - end - end) - end - end end # 3. Return last allocated var if applicable diff --git a/src/array/trsm.jl b/src/array/trsm.jl index c0c025468..81e098e7a 100644 --- a/src/array/trsm.jl +++ b/src/array/trsm.jl @@ -1,4 +1,5 @@ function trsv!(uplo::Char, trans::Char, diag::Char, alpha::T, A::DMatrix{T}, B::DVector{T}) where T + Dagger.@record_op :trsv! A B zone = one(T) mzone = -one(T) @@ -57,6 +58,7 @@ function trsv!(uplo::Char, trans::Char, diag::Char, alpha::T, A::DMatrix{T}, B:: end function trsm!(side::Char, uplo::Char, trans::Char, diag::Char, alpha::T, A::DMatrix{T}, B::DVecOrMat{T}) where T + Dagger.@record_op :trsm! A B zone = one(T) mzone = -one(T) diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index dc6668486..42760020a 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -1,4 +1,5 @@ -import Graphs: SimpleDiGraph, add_edge!, add_vertex!, inneighbors, outneighbors, nv +import Graphs: SimpleDiGraph, add_edge!, add_vertex!, inneighbors, outneighbors, nv, + weakly_connected_components, topological_sort_by_dfs, vertices export In, Out, InOut, Deps, spawn_datadeps @@ -240,6 +241,11 @@ struct HistoryEntry ainfo::AliasingWrapper space::MemorySpace write_num::Int + # Producer of this write/copy. Remainder syncdeps wait on this task + # directly instead of re-resolving through live `ainfos_owner`, which may + # later name a different task for the same ainfo or miss the producer when + # only an overlapping ainfo is consulted. + task::DTask end struct AliasedObjectCacheStore @@ -388,6 +394,9 @@ function aliased_object!(f, cache::AliasedObjectCache, x; ainfo=aliasing(cache.a end struct DataDepsState + # The DAG spec for the task + dag_spec::DAGSpec + # The mapping of original raw argument to its Chunk # N.B. Values are Chunks, or raw remote handles (e.g. ChunkView under MPI) raw_arg_to_chunk::IdDict{Any,Any} @@ -457,7 +466,7 @@ struct DataDepsState ainfos_owner::Dict{AliasingWrapper,Union{Pair{DTask,Int},Nothing}} ainfos_readers::Dict{AliasingWrapper,Vector{Pair{DTask,Int}}} - function DataDepsState() + function DataDepsState(dag_spec::DAGSpec) arg_to_chunk = IdDict{Any,Any}() arg_origin = IdDict{Any,MemorySpace}() remote_args = Dict{MemorySpace,IdDict{Any,Any}}() @@ -482,7 +491,7 @@ struct DataDepsState ainfos_owner = Dict{AliasingWrapper,Union{Pair{DTask,Int},Nothing}}() ainfos_readers = Dict{AliasingWrapper,Vector{Pair{DTask,Int}}}() - return new(arg_to_chunk, arg_origin, remote_args, remote_arg_to_original, remote_arg_w, ainfo_arg, arg_history, arg_owner, arg_current, arg_overlaps, ainfo_backing_chunk, + return new(dag_spec, arg_to_chunk, arg_origin, remote_args, remote_arg_to_original, remote_arg_w, ainfo_arg, arg_history, arg_owner, arg_current, arg_overlaps, ainfo_backing_chunk, supports_inplace_cache, ainfo_cache, ainfos_lookup, ainfos_overlaps, ainfos_owner, ainfos_readers) end end @@ -610,7 +619,12 @@ function aliasing!(state::DataDepsState, target_space::MemorySpace, arg_w::Argum return state.ainfo_cache[remote_arg_w] end - # Calculate the ainfo + # Calculate the ainfo via the current acceleration so SPMD backends (MPI) + # can broadcast owner-computed aliasing and keep ranks uniform. The generic + # `aliasing(::Acceleration, ...)` fallback uses `aliasing_unwrapped`, which + # resolves wrappers of whole-object containers (e.g. a view of a + # `DSparseArray`) to the container itself. For `Chunk`s, unwrap+aliasing + # happen where the data lives (remotely / on the owning rank). ainfo = AliasingWrapper(aliasing(current_acceleration(), remote_arg, arg_w.dep_mod)) # Cache the result @@ -661,10 +675,8 @@ function merge_history!(state::DataDepsState, arg_w::ArgumentWrapper, other_arg_ history = state.arg_history[arg_w] @opcounter :merge_history @opcounter :merge_history_complexity length(history) - origin_space = state.arg_origin[other_arg_w.arg] for other_entry in state.arg_history[other_arg_w] - write_num_tuple = HistoryEntry(AliasingWrapper(NoAliasing()), origin_space, other_entry.write_num) - range = searchsorted(history, write_num_tuple; by=x->x.write_num) + range = searchsorted(history, other_entry; by=x->x.write_num) if !isempty(range) # Find and skip duplicates match = false @@ -672,7 +684,8 @@ function merge_history!(state::DataDepsState, arg_w::ArgumentWrapper, other_arg_ source_entry = history[source_idx] if source_entry.ainfo == other_entry.ainfo && source_entry.space == other_entry.space && - source_entry.write_num == other_entry.write_num + source_entry.write_num == other_entry.write_num && + source_entry.task == other_entry.task match = true break end @@ -839,12 +852,12 @@ function add_writer!(state::DataDepsState, arg_w::ArgumentWrapper, dest_space::M empty!(state.arg_history[arg_w]) # Add our own history - push!(state.arg_history[arg_w], HistoryEntry(ainfo, dest_space, write_num)) + push!(state.arg_history[arg_w], HistoryEntry(ainfo, dest_space, write_num, task)) # Find overlapping arguments and update their history for other_arg_w in state.arg_overlaps[arg_w] other_arg_w == arg_w && continue - push!(state.arg_history[other_arg_w], HistoryEntry(ainfo, dest_space, write_num)) + push!(state.arg_history[other_arg_w], HistoryEntry(ainfo, dest_space, write_num, task)) end # Track which spaces hold a fully-current replica of this region @@ -889,6 +902,29 @@ end isremotehandle(x) = false isremotehandle(x::DTask) = true isremotehandle(x::Chunk) = true +""" + slot_is_already_in_place(data, orig_space, dest_space) -> Bool + +Whether `data` can serve as its own Datadeps slot in `dest_space`. + +Only a `Chunk` that already lives in `dest_space` and wraps a `move_rewrap` +*leaf* qualifies. `move_rewrap` does more than move bytes: it rebuilds wrappers +and resolves handles (a `ChunkView` becomes a `Chunk` over a real `SubArray`, a +nested `Chunk` is flattened) so that what reaches the task is a plain +destination-space value. Those results differ from the input even when nothing +moves, so only leaves may be passed through untouched. + +Restricted to locally-owned chunks under non-uniform execution because deciding +this requires unwrapping `data`, and because MPI chunk handles are rank-relative. +""" +function slot_is_already_in_place(data, orig_space, dest_space) + data isa Chunk || return false + orig_space == dest_space || return false + uniform_execution() && return false + root_worker_id(data.processor) == myid() || return false + return move_rewrap_parts(unwrap(data)) === nothing +end + function generate_slot!(state::DataDepsState, dest_space, data) # N.B. We do not perform any sync/copy with the current owner of the data, # because all we want here is to make a copy of some version of the data, @@ -906,8 +942,18 @@ function generate_slot!(state::DataDepsState, dest_space, data) id = rand(Int) @maybelog ctx timespan_start(ctx, :move, (;thunk_id=0, id, position=ArgPosition(), processor=to_proc), (;f=nothing, data)) tid = something(DATADEPS_CURRENT_TASK[], (;uid=0)).uid - data_chunk = with(DATADEPS_THUNK_ID=>tid) do - remotecall_endpoint_toplevel(move_rewrap, current_acceleration(), aliased_object_cache, from_proc, to_proc, orig_space, dest_space, data) + data_chunk = if slot_is_already_in_place(data, orig_space, dest_space) + # Nothing to move: the slot for data already in `dest_space` is the data + # itself. Going through `move_rewrap` here would allocate a second Chunk + # (and DRef) over the very same memory, which costs a MemPool round-trip + # per argument per region and buys nothing. Still route through + # `aliased_object!` so the cache records this object as the (never-freed) + # original for its aliasing key, exactly as the general path does. + aliased_object!(Returns(data), aliased_object_cache, data)::Chunk + else + with(DATADEPS_THUNK_ID=>tid) do + remotecall_endpoint_toplevel(move_rewrap, current_acceleration(), aliased_object_cache, from_proc, to_proc, orig_space, dest_space, data) + end end @maybelog ctx timespan_finish(ctx, :move, (;thunk_id=0, id, position=ArgPosition(), processor=to_proc), (;f=nothing, data=data_chunk)) @assert memory_space(data_chunk) == dest_space "space mismatch! $dest_space (dest) != $(memory_space(data_chunk)) (actual) ($(typeof(data)) (data) vs. $(typeof(data_chunk)) (chunk)), spaces ($orig_space -> $dest_space)" diff --git a/src/datadeps/chunkview.jl b/src/datadeps/chunkview.jl index 903785f93..86964e6a9 100644 --- a/src/datadeps/chunkview.jl +++ b/src/datadeps/chunkview.jl @@ -69,7 +69,10 @@ function aliasing(accel::Acceleration, x::ChunkView{N}, dep_mod) where N return remotecall_fetch(root_worker_id(x.chunk.processor), x.chunk, x.slices) do x, slices x = unwrap(x) v = view(x, slices...) - return aliasing(accel, v, dep_mod) + # A view of a whole-object container (e.g. `DSparseArray`) must alias the + # entire container; `aliasing_unwrapped` resolves that (otherwise it just + # aliases the view), and crucially does so here where `v` lives. + return aliasing_unwrapped(v) end end aliasing(x::ChunkView) = aliasing(current_acceleration(), x, identity) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl new file mode 100644 index 000000000..8c00d8b51 --- /dev/null +++ b/src/datadeps/hierarchical.jl @@ -0,0 +1,1358 @@ +# Hierarchical scheduling for datadeps +# Spreads scheduling work across multiple threads/workers/MPI ranks via a +# 4-phase pipeline: +# Phase 1: Parallel aliasing info construction +# Phase 2: Sequential DAG construction from aliasing overlaps +# Phase 3: Data-affinity DAG partitioning (by Distributed worker or MPI rank +# via `partition_affinity_id`) +# Phase 4: Per-partition scheduling -- under Distributed each partition runs on +# its own task concurrently; under MPI (uniform execution) partitions +# are scheduled sequentially in global topological order for SPMD +# safety. Each partition computes its *own* (partition-local) AOT +# schedule over just its processors. +# +# N.B. Concurrent task submission from many partition tasks used to intermittently +# hang in the core eager scheduler. The root cause was `Distributed.Future`, +# which is unsafe under concurrent same-process access; Dagger now backs task +# futures with `MemPool.DFuture` (a thread-safe, `DEvent`-based future), so +# parallel Phase 4 is safe. The actual task submission (`enqueue!`) is still +# serialized via `LockedEnqueueQueue`, but the (expensive) per-task +# `distribute_task!` preparation runs concurrently across partitions. + +struct HierarchicalTaskInfo + arg_w::ArgumentWrapper + readdep::Bool + writedep::Bool +end + +""" + BatchedEnqueueQueue(inner, lock; limit=DATADEPS_BATCH_LIMIT[]) + +Buffers tasks locally and submits them to `inner` in batches of at most `limit`. + +Submitting one task at a time costs a scheduler round-trip each; batching pays +that once per group. It also collapses the lock traffic, which matters because +every partition of the hierarchical path shares a single lock over the region's +queue -- with one acquisition per task the partitions spend most of their time +convoyed behind each other rather than preparing tasks in parallel. + +`limit` trades that amortization against latency: buffered tasks cannot start +executing, so an unbounded batch would keep a partition's whole workload idle +until it finished planning, serializing planning against execution. A small +limit keeps the pipeline fed while still amortizing most of the per-submission +cost. + +`flush_batch!` must additionally be called before anything outside this +partition can observe the buffered tasks (see `schedule_partition_full!`), since +a task's syncdeps must already be submitted when a dependent task is prepared. +""" +struct BatchedEnqueueQueue <: AbstractTaskQueue + inner::AbstractTaskQueue + lock::ReentrantLock + pending::Vector{DTaskPair} + limit::Int +end +BatchedEnqueueQueue(inner::AbstractTaskQueue, lock::ReentrantLock; + limit::Int=DATADEPS_BATCH_LIMIT[]) = + BatchedEnqueueQueue(inner, lock, DTaskPair[], limit) +function enqueue!(beq::BatchedEnqueueQueue, pair::DTaskPair) + push!(beq.pending, pair) + length(beq.pending) >= beq.limit && flush_batch!(beq) + return +end +function enqueue!(beq::BatchedEnqueueQueue, pairs::Vector{DTaskPair}) + append!(beq.pending, pairs) + length(beq.pending) >= beq.limit && flush_batch!(beq) + return +end +function flush_batch!(beq::BatchedEnqueueQueue) + isempty(beq.pending) && return + @lock beq.lock enqueue!(beq.inner, beq.pending) + empty!(beq.pending) + return +end + +"Maximum number of tasks a hierarchical partition buffers before submitting." +const DATADEPS_BATCH_LIMIT = Ref(4) + +struct HierarchicalTaskMeta + pair::DTaskPair + # Wrapped argument identity for the task's first aliasing arg. Under MPI + # this may be a raw `ChunkView` (kept unwrapped by `datadeps_arg_wrap`). + arg_chunk::Any + may_alias::Bool + inplace_move::Bool + deps::Vector{HierarchicalTaskInfo} + # Indices of same-region producer tasks whose results are passed as + # arguments (e.g. `In(t1)`). These cannot be `fetch`'d during the + # pre-scan because they are not launched yet; they become hard DAG edges. + value_deps::Vector{Int} +end + +### Cross-partition chunk ownership ### +# +# Each partition schedules with its own `DataDepsState`, so `arg_owner` / +# `arg_history` / physical slots are per-partition. When a backing chunk is +# written by tasks in >=2 partitions that live in *different* memory spaces, +# each partition would otherwise copy that chunk from its origin, write its own +# sub-range, and record itself as owner -- the physical copies then diverge and +# the final copy-back keeps only one, silently losing the others' writes. +# +# The registry below carries the single authoritative version ("ownership") of +# each such shared chunk across partition boundaries. It is deadlock-free by +# construction: a single lock (never per-argument locks acquired in different +# orders) plus the global DAG ordering -- a consumer partition always `wait`s on +# its cross-partition predecessor's `task_submitted` event before scheduling, so +# the producer's ownership commit is always visible before the consumer reads it. +# The lock only provides memory-safety for concurrent access to *different* +# chunks; per-chunk correctness comes from the DAG order. +"Authoritative, cross-partition ownership state for a single shared backing chunk." +mutable struct ChunkOwnership + owner_space::MemorySpace # space holding the current version + owner_slot::Any # physical slot chunk in `owner_space` + owner_task::Union{DTask,Nothing} # producer of the current version (nothing => origin data) + owner_state::Union{DataDepsState,Nothing} # owning partition's state (for copy-back) + const origin_space::MemorySpace # the chunk's home space +end + +struct SharedChunkRegistry + entries::IdDict{Any,ChunkOwnership} # backing chunk (identity) => ownership + lock::ReentrantLock +end +SharedChunkRegistry() = SharedChunkRegistry(IdDict{Any,ChunkOwnership}(), ReentrantLock()) + +is_shared_chunk(::Nothing, @nospecialize(chunk)) = false +is_shared_chunk(reg::SharedChunkRegistry, @nospecialize(chunk)) = haskey(reg.entries, chunk) + +""" + build_shared_chunk_registry(task_metas, vertex_to_partition, partition_space) -> SharedChunkRegistry or nothing + +Detects backing chunks accessed by partitions spanning >=2 distinct memory +spaces (the only case that can split-brain) and returns a registry seeded with +origin ownership for each. Returns `nothing` when all partitions share a single +space (e.g. single-worker), so the fast path is entirely unchanged there. +""" +function build_shared_chunk_registry(task_metas::Vector{HierarchicalTaskMeta}, + vertex_to_partition::Vector{Int}, + partition_space::Vector{<:MemorySpace}) + length(unique(partition_space)) <= 1 && return nothing + + chunk_spaces = IdDict{Any,Set{MemorySpace}}() + chunk_origin = IdDict{Any,MemorySpace}() + for v in eachindex(task_metas) + pspace = partition_space[vertex_to_partition[v]] + for dep in task_metas[v].deps + chunk = dep.arg_w.arg + push!(get!(()->Set{MemorySpace}(), chunk_spaces, chunk), pspace) + haskey(chunk_origin, chunk) || (chunk_origin[chunk] = memory_space(chunk)) + end + end + + reg = SharedChunkRegistry() + for (chunk, spaces) in chunk_spaces + length(spaces) >= 2 || continue + origin = chunk_origin[chunk] + reg.entries[chunk] = ChunkOwnership(origin, chunk, nothing, nothing, origin) + end + isempty(reg.entries) && return nothing + return reg +end + +""" + _sync_incoming_ownership!(state, registry, our_space, task_arg_ws, write_num) + +Before a task's copy-to phase, for each shared backing-chunk argument whose +globally-current owner (per `registry`) lives in a space other than `our_space`, +seed this partition's `state` so the existing copy-to machinery pulls a fresh +whole-chunk copy from the true owner (with a syncdep on the producing task). A +no-op for private chunks or when we already hold the authoritative version. +""" +function _sync_incoming_ownership!(state::DataDepsState, registry::SharedChunkRegistry, + our_space::MemorySpace, task_arg_ws, write_num::Int) + map_or_ntuple(task_arg_ws) do idx + arg_ws = task_arg_ws[idx] + (arg_ws.may_alias && arg_ws.inplace_move) || return + chunk = arg_ws.arg + entry = get(registry.entries, chunk, nothing) + entry === nothing && return + + owner_space, owner_slot, owner_task = @lock registry.lock begin + (entry.owner_space, entry.owner_slot, entry.owner_task) + end + owner_space == our_space && return # we already hold the authoritative copy + + # Register the owner's physical slot so slot/aliasing lookups in this + # partition reuse it (rather than generating a fresh, stale copy). + dest_args = get!(IdDict{Any,Any}, state.remote_args, owner_space) + if !haskey(dest_args, chunk) + dest_args[chunk] = owner_slot + state.remote_arg_to_original[owner_slot] = chunk + end + + map_or_ntuple(arg_ws.deps) do dep_idx + dep = arg_ws.deps[dep_idx] + arg_w = dep.arg_w + # Point ownership at the owner space and clear any locally-merged + # history, so `compute_remainder_for_arg!` takes the `FullCopy` + # (whole-chunk) path from the owner rather than a partial remainder. + state.arg_owner[arg_w] = owner_space + haskey(state.arg_history, arg_w) && empty!(state.arg_history[arg_w]) + src_ainfo = aliasing!(state, owner_space, arg_w) + if owner_task !== nothing + # Make the ensuing copy-to (via `get_read_deps!`) wait on the + # producer, so we never copy the owner slot before it is written. + state.ainfos_owner[src_ainfo] = owner_task => (write_num - 1) + end + return + end + return + end + return +end + +""" + _commit_ownership!(state, registry, our_space, task, task_arg_ws) + +After a task is recorded as a writer, publish it as the new authoritative owner +of each shared backing chunk it writes, so subsequent cross-partition consumers +pull from here. Ordering (and thus visibility) is guaranteed by the DAG's +`task_submitted` handshake; the lock only guards concurrent commits for other +chunks. +""" +function _commit_ownership!(state::DataDepsState, registry::SharedChunkRegistry, + our_space::MemorySpace, task::DTask, task_arg_ws) + map_or_ntuple(task_arg_ws) do idx + arg_ws = task_arg_ws[idx] + (arg_ws.may_alias && arg_ws.inplace_move) || return + chunk = arg_ws.arg + entry = get(registry.entries, chunk, nothing) + entry === nothing && return + + wrote = false + map_or_ntuple(arg_ws.deps) do dep_idx + arg_ws.deps[dep_idx].writedep && (wrote = true) + return + end + wrote || return + + dest_args = get(state.remote_args, our_space, nothing) + slot = dest_args === nothing ? nothing : get(dest_args, chunk, nothing) + @lock registry.lock begin + entry.owner_space = our_space + slot !== nothing && (entry.owner_slot = slot) + entry.owner_task = task + entry.owner_state = state + end + return + end + return +end + +# Below this many tasks, the fixed costs of spawning threads and merging +# per-thread results outweigh the benefit of parallelizing the pre-scan. +const COLLECT_ALIASED_ARGS_MIN_CHUNK = 256 + +# Thread-safe "get or compute" against a shared `IdDict` cache. The +# expensive computation (`f`) is performed outside of the lock so that +# multiple threads can make progress on distinct keys concurrently; if two +# threads race on the same key, the loser's result is simply discarded (the +# corresponding `Chunk`/Bool is cheap to let the GC reclaim) so that every +# thread agrees on a single canonical value (e.g. `Chunk`) per raw argument. +@inline function _cached_get!(f, cache::IdDict{Any,V}, cache_lock::Union{ReentrantLock,Nothing}, key) where V + if cache_lock === nothing + return get!(f, cache, key) + end + @lock cache_lock begin + haskey(cache, key) && return cache[key]::V + end + result = f()::V + @lock cache_lock begin + return get!(cache, key, result)::V + end +end + +""" + collect_aliased_args(seen_tasks) -> (task_metas, unique_arg_ws) + +Pre-scans all tasks to collect per-task dependency metadata and the set of +unique `ArgumentWrapper`s that need aliasing analysis. This mirrors the logic +in `populate_task_info!` but only inspects arguments without modifying any +scheduling state. + +For large batches of tasks (the common case for e.g. panel-factorization +algorithms which submit many small tasks per `spawn_datadeps` region), the +pre-scan itself (not just the aliasing computation in +`build_aliasing_parallel`) can dominate scheduling time, since it touches +every argument of every task. This is parallelized across threads: each +thread scans a contiguous range of `seen_tasks` into its own disjoint slice +of `task_metas` and its own local `unique_arg_ws` map (merged at the end), +while sharing (lock-protected) caches for `Chunk`-wrapping and +`supports_inplace_move`, ensuring a single canonical `Chunk` identity per +raw argument regardless of which thread first observes it. +""" +function collect_aliased_args(seen_tasks::Vector{DTaskPair}) + n = length(seen_tasks) + task_metas = Vector{HierarchicalTaskMeta}(undef, n) + n == 0 && return task_metas, Dict{ArgumentWrapper,ArgumentWrapper}() + + # Map in-region tasks to vertex indices so we can record value deps + # without fetching unlaunched DTasks during the pre-scan. + task_to_idx = IdDict{DTask,Int}() + for (i, pair) in enumerate(seen_tasks) + task_to_idx[pair.task] = i + end + + supports_cache = IdDict{Any,Bool}() + # Values are Chunks, or raw remote handles kept by `datadeps_arg_wrap` + # (e.g. ChunkView under MPI). + raw_arg_cache = IdDict{Any,Any}() + + # Under uniform execution (MPI), `tochunk` / `MPIRefID` allocation must stay + # sequential on the root task, and spawned threads would not inherit the + # TaskLocalValue acceleration (falling back to Distributed → DRef chunks). + nchunks = if uniform_execution() || Threads.nthreads() <= 1 + 1 + else + min(Threads.nthreads(), cld(n, COLLECT_ALIASED_ARGS_MIN_CHUNK)) + end + + if nchunks <= 1 + unique_arg_ws = Dict{ArgumentWrapper, ArgumentWrapper}() + _collect_aliased_args_range!(task_metas, unique_arg_ws, seen_tasks, 1:n, + supports_cache, raw_arg_cache, nothing, task_to_idx) + return task_metas, unique_arg_ws + end + + cache_lock = ReentrantLock() + chunk_size = cld(n, nchunks) + starts = collect(1:chunk_size:n) + per_chunk_arg_ws = Vector{Dict{ArgumentWrapper,ArgumentWrapper}}(undef, length(starts)) + # Propagate the caller's acceleration into each worker thread; TaskLocalValue + # does not inherit across `Threads.@spawn`. + parent_accel = current_acceleration() + + @sync for (ci, start) in enumerate(starts) + range = start:min(start+chunk_size-1, n) + Threads.@spawn begin + set_task_acceleration!(parent_accel) + local_arg_ws = Dict{ArgumentWrapper,ArgumentWrapper}() + _collect_aliased_args_range!(task_metas, local_arg_ws, seen_tasks, range, + supports_cache, raw_arg_cache, cache_lock, task_to_idx) + per_chunk_arg_ws[ci] = local_arg_ws + end + end + + unique_arg_ws = per_chunk_arg_ws[1] + for ci in 2:length(per_chunk_arg_ws) + merge!(unique_arg_ws, per_chunk_arg_ws[ci]) + end + + return task_metas, unique_arg_ws +end + +function _collect_aliased_args_range!(task_metas::Vector{HierarchicalTaskMeta}, + unique_arg_ws::Dict{ArgumentWrapper,ArgumentWrapper}, + seen_tasks::Vector{DTaskPair}, + range::UnitRange{Int}, + supports_cache::IdDict{Any,Bool}, + raw_arg_cache::IdDict{Any,Any}, + cache_lock::Union{ReentrantLock,Nothing}, + task_to_idx::IdDict{DTask,Int}) + for task_idx in range + pair = seen_tasks[task_idx] + spec = pair.spec + task = pair.task + fargs = spec.fargs + + all_deps = HierarchicalTaskInfo[] + value_deps = Int[] + first_chunk = nothing + task_may_alias = false + task_inplace = false + + for arg_idx in (is_typed(spec) ? (1:length(fargs)) : eachindex(fargs)) + _arg = fargs[arg_idx] + _arg_with_deps = value(_arg) + + arg_pre_unwrap, deps = unwrap_inout(_arg_with_deps) + + # Same-region DTask arguments are not launched yet, so we cannot + # `fetch` them here (that is what the sequential `distribute_task!` + # path does *after* launching the producer). Record a value + # dependency instead; aliasing of the result is handled later in + # `distribute_task!` once the producer has been submitted. + if arg_pre_unwrap isa DTask && !istaskstarted(arg_pre_unwrap) + pred_idx = get(task_to_idx, arg_pre_unwrap, 0) + if pred_idx != 0 && pred_idx != task_idx + push!(value_deps, pred_idx) + end + continue + end + + arg = arg_pre_unwrap isa DTask ? fetch(arg_pre_unwrap; raw=true) : arg_pre_unwrap + + may_alias = type_may_alias(typeof(arg)) + inplace_move = may_alias && _cached_get!(supports_cache, cache_lock, arg) do + supports_inplace_move(arg) + end + + if !may_alias || !inplace_move + continue + end + + arg_chunk = _cached_get!(raw_arg_cache, cache_lock, arg) do + # Match `populate_task_info!`: acceleration-aware wrap (MPIRef + # under MPI) rather than a bare `tochunk` that can produce a + # rank-local DRef when TLS acceleration is unset. + arg isa Chunk ? arg : datadeps_arg_wrap(arg) + end + + if first_chunk === nothing + first_chunk = arg_chunk + task_may_alias = true + task_inplace = true + end + + for (dep_mod, readdep, writedep) in deps + arg_w = ArgumentWrapper(arg_chunk, dep_mod) + unique_arg_ws[arg_w] = arg_w + push!(all_deps, HierarchicalTaskInfo(arg_w, readdep, writedep)) + end + end + + task_metas[task_idx] = HierarchicalTaskMeta( + pair, first_chunk, task_may_alias, task_inplace, all_deps, value_deps + ) + end +end + +""" + build_aliasing_parallel(unique_arg_ws) -> (lookup, ainfos_overlaps, arg_to_ainfo) + +Phase 1: Computes `AliasingWrapper` for every unique `ArgumentWrapper` in +parallel. On each worker, threads are used to compute aliasing info for local +data. Results are gathered and reduced into a single `AliasingLookup` with +overlap information. +""" +function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWrapper}) + arg_ws_vec = collect(values(unique_arg_ws)) + + by_worker = Dict{Int, Vector{ArgumentWrapper}}() + for arg_w in arg_ws_vec + wid = root_worker_id(memory_space(arg_w.arg)) + worker_args = get!(Vector{ArgumentWrapper}, by_worker, wid) + push!(worker_args, arg_w) + end + + arg_to_ainfo = Dict{ArgumentWrapper, AliasingWrapper}() + + if length(by_worker) == 1 + # Common single-worker case: avoid the `@sync`/`Threads.@spawn`/lock + # overhead entirely, since there's nothing to run concurrently with. + # `_compute_aliasing_batch` still uses threads internally when there + # are enough args to make it worthwhile. + wid, worker_args = only(by_worker) + results = wid == myid() ? _compute_aliasing_batch(worker_args) : + remotecall_fetch(_compute_aliasing_batch, wid, worker_args) + # Key by the *local* `arg_w`, not the pair's: for a remote worker the + # returned `ArgumentWrapper` is a deserialized copy that need not be + # identity/hash-equal to the entry in `arg_ws_vec` we later look up + # (which would raise a `KeyError`). `_compute_aliasing_batch` preserves + # input order, so pair by index. + for i in eachindex(worker_args) + arg_to_ainfo[worker_args[i]] = results[i].second + end + else + all_results_lock = ReentrantLock() + @sync for (wid, worker_args) in by_worker + Threads.@spawn begin + results = if wid == myid() + _compute_aliasing_batch(worker_args) + else + remotecall_fetch(_compute_aliasing_batch, wid, worker_args) + end + # Key by the *local* `arg_w` (see single-worker note above): a + # remote worker returns deserialized `ArgumentWrapper` copies + # that may not compare equal to our `arg_ws_vec` lookup keys. + @lock all_results_lock begin + for i in eachindex(worker_args) + arg_to_ainfo[worker_args[i]] = results[i].second + end + end + end + end + end + + lookup = AliasingLookup() + ainfos_overlaps = Dict{AliasingWrapper, Set{AliasingWrapper}}() + + for arg_w in arg_ws_vec + ainfo = arg_to_ainfo[arg_w] + if haskey(ainfos_overlaps, ainfo) + continue + end + + ainfo_idx = push!(lookup, ainfo) + overlaps = Set{AliasingWrapper}() + push!(overlaps, ainfo) + for other_ainfo in intersect(lookup, ainfo; ainfo_idx) + ainfo == other_ainfo && continue + push!(overlaps, other_ainfo) + push!(ainfos_overlaps[other_ainfo], ainfo) + end + ainfos_overlaps[ainfo] = overlaps + end + + return lookup, ainfos_overlaps, arg_to_ainfo +end + +# Below this many args, the fixed cost of forking/joining `Threads.@threads` +# outweighs the benefit of parallelizing the (typically cheap) `aliasing()` calls. +const COMPUTE_ALIASING_BATCH_MIN_PARALLEL = 8 + +function _compute_aliasing_batch(arg_ws::Vector{ArgumentWrapper}) + n = length(arg_ws) + results = Vector{Pair{ArgumentWrapper, AliasingWrapper}}(undef, n) + accel = current_acceleration() + # Under uniform execution (MPI), aliasing may perform collectives that must + # run in the same sequential order on every rank -- never Threads.@threads. + # Also dispatch through `aliasing(accel, ...)` so MPI Chunks go through the + # owner-broadcast path rather than a local unwrap. + can_parallel = !uniform_execution(accel) && + n >= COMPUTE_ALIASING_BATCH_MIN_PARALLEL && + Threads.nthreads() > 1 + if can_parallel + Threads.@threads for i in 1:n + arg_w = arg_ws[i] + ainfo = AliasingWrapper(aliasing(accel, arg_w.arg, arg_w.dep_mod)) + results[i] = arg_w => ainfo + end + else + for i in 1:n + arg_w = arg_ws[i] + ainfo = AliasingWrapper(aliasing(accel, arg_w.arg, arg_w.dep_mod)) + results[i] = arg_w => ainfo + end + end + return results +end + +""" + build_dependency_dag(task_metas, arg_to_ainfo, ainfos_overlaps) + -> SimpleDiGraph + +Phase 2: Walks tasks in submission order and builds a `SimpleDiGraph` encoding +data dependencies based on the pre-computed aliasing overlaps. Uses the same +WAW / RAW / WAR rules as `get_write_deps!` / `get_read_deps!`. +""" +function build_dependency_dag(task_metas::Vector{HierarchicalTaskMeta}, + arg_to_ainfo::Dict{ArgumentWrapper, AliasingWrapper}, + ainfos_overlaps::Dict{AliasingWrapper, Set{AliasingWrapper}}) + n = length(task_metas) + dag = SimpleDiGraph(n) + + ainfos_owner = Dict{AliasingWrapper, Union{Nothing, Pair{Int,Int}}}() + ainfos_readers = Dict{AliasingWrapper, Vector{Pair{Int,Int}}}() + + write_num = 1 + for v in 1:n + meta = task_metas[v] + + # Hard edges from same-region DTask value arguments (producer must + # be launched before we can fetch its result in distribute_task!). + for pred_v in meta.value_deps + if pred_v != v + add_edge!(dag, pred_v, v) + end + end + + # Add dependency edges + for dep in meta.deps + ainfo = get(arg_to_ainfo, dep.arg_w, nothing) + ainfo === nothing && continue + + if !haskey(ainfos_owner, ainfo) + ainfos_owner[ainfo] = nothing + ainfos_readers[ainfo] = Pair{Int,Int}[] + end + + overlaps = get(ainfos_overlaps, ainfo, Set{AliasingWrapper}()) + + if dep.writedep + for other_ainfo in overlaps + owner = get(ainfos_owner, other_ainfo, nothing) + if owner !== nothing + pred_v, pred_wn = owner + if pred_wn != write_num && pred_v != v + add_edge!(dag, pred_v, v) + end + end + for (reader_v, reader_wn) in get(ainfos_readers, other_ainfo, Pair{Int,Int}[]) + if reader_wn != write_num && reader_v != v + add_edge!(dag, reader_v, v) + end + end + end + else + for other_ainfo in overlaps + owner = get(ainfos_owner, other_ainfo, nothing) + if owner !== nothing + pred_v, pred_wn = owner + if pred_wn != write_num && pred_v != v + add_edge!(dag, pred_v, v) + end + end + end + end + end + + # Update ownership tracking + for dep in meta.deps + ainfo = get(arg_to_ainfo, dep.arg_w, nothing) + ainfo === nothing && continue + + if !haskey(ainfos_owner, ainfo) + ainfos_owner[ainfo] = nothing + ainfos_readers[ainfo] = Pair{Int,Int}[] + end + + if dep.writedep + ainfos_owner[ainfo] = v => write_num + empty!(ainfos_readers[ainfo]) + else + push!(ainfos_readers[ainfo], v => write_num) + end + end + + write_num += 1 + end + + return dag +end + +""" +Target number of tasks per partition when partitioning across the local +processors of a single owner. + +Partitioning a single owner parallelizes *planning* only -- every partition +still places work on the same processors -- so it has to earn back what it +duplicates. Each partition builds its own `DataDepsState`, so an argument shared +between partitions has its aliasing and slot bookkeeping redone once per +partition; with neighbour-sharing workloads (stencils) that duplication grows +quickly with the partition count. +""" +const HIER_TASKS_PER_PARTITION = 16 + +""" + single_owner_partition_count(ntasks, nprocs) -> Int + +How many partitions to split `ntasks` into when all processors share an owner. + +Capped at half of `nprocs` because partition planning tasks and the processor +runners execute on the same threads: one partition per processor leaves nothing +to run the work being planned, and measures slower than not partitioning at all. +""" +single_owner_partition_count(ntasks::Int, nprocs::Int) = + clamp(ntasks ÷ HIER_TASKS_PER_PARTITION, 1, max(1, nprocs ÷ 2)) + +""" + partition_dag(dag, task_metas, all_procs) -> (vertex_to_partition, n_partitions, partition_procs) + +Phase 3: Assigns each task vertex to a partition using data-affinity. For +multi-owner setups (Distributed workers, or MPI ranks via +`partition_affinity_id`), tasks are assigned to the owner holding the most +argument data. For single-owner multi-threaded setups, tasks are balanced +across available processors in topological order. +""" +function partition_dag(dag::SimpleDiGraph, task_metas::Vector{HierarchicalTaskMeta}, + all_procs::Vector{<:Processor}) + n = length(task_metas) + # Stable order so SPMD ranks (and Distributed workers) agree on partition + # indexing when affinity ids are collected from an unordered processor set. + owners = sort!(unique(partition_affinity_id.(all_procs))) + n_owners = length(owners) + + procs_by_owner = Dict{Int, Vector{Processor}}() + for proc in all_procs + oid = partition_affinity_id(proc) + push!(get!(Vector{Processor}, procs_by_owner, oid), proc) + end + + multi_owner = n_owners > 1 + if multi_owner + n_partitions = n_owners + partition_owner = owners + else + n_partitions = single_owner_partition_count(n, length(all_procs)) + partition_owner = fill(first(owners), n_partitions) + end + + vertex_to_partition = Vector{Int}(undef, n) + + if multi_owner + owner_to_partition = Dict(o => i for (i, o) in enumerate(owners)) + default_scope = DefaultScope() + for v in 1:n + meta = task_metas[v] + task_scope = @something(meta.pair.spec.options.compute_scope, meta.pair.spec.options.scope, default_scope) + + # Owners whose processors are eligible under this task's scope. + # A non-default scope that spans multiple owners must still spread + # work across those owners (via affinity / round-robin) -- picking + # only the first match pins everything to owner 1 and breaks + # multi-owner execution. + if task_scope == default_scope + matching = collect(1:n_partitions) + else + matching = Int[] + for (pid, oid) in enumerate(owners) + oprocs = procs_by_owner[oid] + if any(proc -> proc_in_scope(proc, task_scope), oprocs) + push!(matching, pid) + end + end + if isempty(matching) + matching = [1] + end + end + + if length(matching) == 1 + vertex_to_partition[v] = only(matching) + continue + end + + affinity = zeros(Int, n_owners) + for dep in meta.deps + arg_space = memory_space(dep.arg_w.arg) + arg_oid = partition_affinity_id(arg_space) + idx = get(owner_to_partition, arg_oid, 0) + if idx > 0 && idx in matching + affinity[idx] += 1 + end + end + best_pid = matching[1] + best_aff = -1 + for pid in matching + if affinity[pid] > best_aff + best_aff = affinity[pid] + best_pid = pid + end + end + if best_aff <= 0 + vertex_to_partition[v] = matching[mod1(v, length(matching))] + else + vertex_to_partition[v] = best_pid + end + end + else + topo = try + topological_sort_by_dfs(dag) + catch + collect(1:n) + end + + default_scope = DefaultScope() + partition_load = zeros(Int, n_partitions) + for v in topo + meta = task_metas[v] + task_scope = @something(meta.pair.spec.options.compute_scope, meta.pair.spec.options.scope, default_scope) + + if task_scope != default_scope + assigned = false + for pid in 1:n_partitions + pidx = mod1(pid, length(all_procs)) + if proc_in_scope(all_procs[pidx], task_scope) + vertex_to_partition[v] = pid + partition_load[pid] += 1 + assigned = true + break + end + end + if !assigned + best = argmin(partition_load) + vertex_to_partition[v] = best + partition_load[best] += 1 + end + else + best = argmin(partition_load) + vertex_to_partition[v] = best + partition_load[best] += 1 + end + end + end + + partition_procs = Vector{Vector{Processor}}(undef, n_partitions) + if multi_owner + for pid in 1:n_partitions + oid = partition_owner[pid] + partition_procs[pid] = procs_by_owner[oid] + end + else + for pid in 1:n_partitions + partition_procs[pid] = copy(all_procs) + end + end + + return vertex_to_partition, n_partitions, partition_procs, multi_owner +end + + +""" + _schedule_vertex!(v, partition_id, temp_queue, state, local_procs, + local_scope, dag, seen_tasks, vertex_to_partition, + schedule, proc_to_scope_lfu, write_num) -> write_num + +Schedules a single (already topologically-ordered) task vertex `v` into its +partition's `state` via `distribute_task!`, returning the updated `write_num`. + +Records cross-partition predecessors as explicit `ThunkSyncdep`s (same-partition +deps are derived from `state` by `distribute_task!`), and passes the +partition-local AOT-assigned processor when one is available and usable for this +partition (else falls back to JIT scheduling). + +Callers must guarantee that every predecessor of `v` (in any partition) has +already been submitted before calling this, so the `ThunkSyncdep`s are valid. +""" +function _schedule_vertex!(v::Int, partition_id::Int, + temp_queue::DataDepsTaskQueue, + state::DataDepsState, + local_procs::Vector{<:Processor}, + local_scope, + dag::SimpleDiGraph, + seen_tasks::Vector{DTaskPair}, + vertex_to_partition::Vector{Int}, + schedule::Dict{DTask,Processor}, + proc_to_scope_lfu, + write_num::Int, + registry::Union{SharedChunkRegistry,Nothing}) + pair = seen_tasks[v] + spec = pair.spec + task = pair.task + + if spec.options.syncdeps === nothing + spec.options.syncdeps = Set{ThunkSyncdep}() + end + for pred_v in inneighbors(dag, v) + if vertex_to_partition[pred_v] != partition_id + pred_task = seen_tasks[pred_v].task + push!(spec.options.syncdeps, ThunkSyncdep(pred_task)) + end + end + + # Use the AOT-assigned processor when one is available and it lies within + # this partition's processors; otherwise fall back to JIT scheduling + # (`proc === nothing`) via `temp_queue`'s scheduler. (An AOT proc can lie + # outside `local_procs` only in the multi-worker case, where partitioning + # restricts each partition to one worker's processors; single-worker + # partitions always span all procs, so the AOT assignment is always usable + # there.) + proc = get(schedule, task, nothing) + if proc !== nothing && !(proc in local_procs) + proc = nothing + end + + return distribute_task!(temp_queue, state, local_procs, local_scope, + spec, task, spec.fargs, + proc_to_scope_lfu, write_num; proc, ownership=registry) +end + +""" + schedule_partition_full!(queue, batch_queue, partition_id, partition_verts, + dag, seen_tasks, local_procs, + vertex_to_partition, task_submitted, + region_uids, value_dep_verts, registry) -> DataDepsState + +Per-partition scheduling for both single-worker and multi-worker hierarchical +paths. Uses existing `distribute_task!` logic with per-partition +`DataDepsState`, `all_procs` limited to this partition's processors, and +cross-partition syncdeps from the precomputed DAG. + +AOT scheduling is run *locally per partition*: a partition-local `DAGSpec` is +built from just this partition's tasks and scheduled over just `local_procs` via +`datadeps_build_schedule!`. `region_uids` (all in-region task uids) lets that +builder bail cleanly when a task depends on a same-region producer in another +partition (which it must not `fetch`). Tasks without an AOT assignment fall back +to JIT scheduling in `distribute_task!`. +""" +function schedule_partition_full!(queue::DataDepsTaskQueue, + batch_queue::BatchedEnqueueQueue, + partition_id::Int, + partition_verts::Vector{Int}, + dag::SimpleDiGraph, + seen_tasks::Vector{DTaskPair}, + local_procs::Vector{<:Processor}, + vertex_to_partition::Vector{Int}, + task_submitted::Vector{Base.Event}, + region_uids::Set{UInt}, + value_dep_verts::Set{Int}, + registry::Union{SharedChunkRegistry,Nothing}) + if isempty(partition_verts) || isempty(local_procs) + return DataDepsState(DAGSpec()) + end + + local_scope = UnionScope(map(ExactScope, local_procs)) + + # N.B. Hierarchical scheduling doesn't build a global `DAGSpec` (that's + # only used by the AOT DAG-caching path in `distribute_tasks!`), so each + # partition just gets an empty one; `distribute_task!` never reads + # `state.dag_spec` directly. + state = DataDepsState(DAGSpec()) + write_num = 1 + proc_to_scope_lfu = BasicLFUCache{Processor,AbstractScope}(1024) + + vert_set = Set{Int}(partition_verts) + topo = try + topological_sort_by_dfs(dag) + catch + collect(vertices(dag)) + end + ordered_verts = filter(v -> v in vert_set, topo) + + # N.B. Each partition gets its own fresh scheduler shard via `similar` + # rather than sharing `queue.scheduler` across all partitions. Two + # independent problems would arise from sharing a single scheduler + # instance here: + # 1) Data race: e.g. `RoundRobinScheduler.proc_idx` would be + # concurrently read/written by every partition's `Threads.@spawn` + # task with no synchronization. + # 2) Semantic bug (worse than the race, and *not* fixed by adding a + # lock): each partition schedules only onto its own worker's + # `local_procs`, which generally has a *different length* than + # other partitions' (or the global) processor list. A `proc_idx` + # counter advanced by one partition's `local_procs` is meaningless + # -- and can be out-of-bounds -- when applied to another + # partition's differently-sized `local_procs`. This reliably + # crashes with a `BoundsError` under multi-worker hierarchical + # scheduling. Giving each partition its own scheduler instance, + # scoped to its own `local_procs`, fixes both issues at once. + temp_queue = DataDepsTaskQueue(batch_queue; scheduler=similar(queue.scheduler)) + + # Per-partition (local) AOT scheduling over just this partition's procs. + # `partition_verts` is in ascending vertex (= submission) order. + partition_pairs = DTaskPair[seen_tasks[v] for v in partition_verts] + _pdag, schedule = datadeps_build_schedule!(temp_queue.scheduler, partition_pairs, + local_procs, local_scope; region_uids) + + # N.B. If this partition throws partway through (e.g. from + # `distribute_task!`), any of our vertices that haven't yet been + # `notify`'d will never be, which would leave other partitions blocked + # forever in `wait(task_submitted[pred_v])` below -- turning a normal, + # reportable exception into a silent, permanent hang (since the + # enclosing `@sync` in `distribute_tasks_hierarchical!` can't finish, + # and thus can't propagate our exception, until *every* spawned + # partition task completes, including the ones stuck waiting on us). + # The `finally` ensures every one of our events gets notified no matter + # how we exit, so that sibling partitions can unblock (and themselves + # fail/finish) and our real exception can actually surface. + # A vertex whose successors all live in this partition is only ever observed + # by us, so its submission can stay buffered; one with a successor elsewhere + # must be submitted before we notify, because that successor will record it + # as a syncdep. + has_external_successor = Set{Int}() + for v in ordered_verts + for succ_v in outneighbors(dag, v) + if vertex_to_partition[succ_v] != partition_id + push!(has_external_successor, v) + break + end + end + end + + try + for v in ordered_verts + for pred_v in inneighbors(dag, v) + if vertex_to_partition[pred_v] != partition_id + # Anything we have buffered may be a dependency of tasks the + # producing partition is about to prepare; get it submitted + # before we block, or two partitions can deadlock waiting on + # each other's unflushed tasks. + flush_batch!(batch_queue) + wait(task_submitted[pred_v]) + end + end + + # A task taking another in-region task's *value* as an argument has + # `distribute_task!` `fetch` that producer, which requires it to have + # actually been launched. A same-partition producer is topologically + # earlier, so it is prepared but possibly still sitting in our batch. + v in value_dep_verts && flush_batch!(batch_queue) + + # N.B. The per-task `distribute_task!` preparation runs concurrently + # across partitions; only the final task submission is serialized + # (via `BatchedEnqueueQueue`, `temp_queue`'s upper queue). This is + # safe now that task futures are backed by `MemPool.DFuture` rather + # than the concurrency-unsafe `Distributed.Future`. The + # cross-partition `wait`s above happen before scheduling `v`, so the + # `ThunkSyncdep`s recorded for `v` are valid. + write_num = _schedule_vertex!( + v, partition_id, temp_queue, state, local_procs, local_scope, + dag, seen_tasks, vertex_to_partition, schedule, + proc_to_scope_lfu, write_num, registry) + + v in has_external_successor && flush_batch!(batch_queue) + notify(task_submitted[v]) + end + flush_batch!(batch_queue) + finally + flush_batch!(batch_queue) + for v in ordered_verts + notify(task_submitted[v]) + end + end + + return state +end + +""" + schedule_partitions_sequential!(...) -> Vector{DataDepsState} + +SPMD-safe Phase 4: schedule every partition's tasks on the root task in global +topological order. Processor assignment remains partition-local (MPI-rank / +worker affinity), but a *single* shared `DataDepsState` is used so argument +history, remainders, and final copy-back stay coherent across ranks -- the same +model as flat `distribute_tasks!`. Per-partition states would otherwise +split-brain overlapping writes (e.g. whole-array + view + triangular dep_mods). + +Returns a one-element vector containing the shared state (for the hierarchical +copy-from/free epilogue). +""" +function schedule_partitions_sequential!(queue::DataDepsTaskQueue, + queue_lock::ReentrantLock, + partitions::Vector{Vector{Int}}, + dag::SimpleDiGraph, + seen_tasks::Vector{DTaskPair}, + partition_procs::Vector{<:Vector{<:Processor}}, + vertex_to_partition::Vector{Int}, + region_uids::Set{UInt}, + registry::Union{SharedChunkRegistry,Nothing}, + wait_all_queue) + n_partitions = length(partitions) + temp_queues = Vector{DataDepsTaskQueue}(undef, n_partitions) + local_scopes = Vector{AbstractScope}(undef, n_partitions) + schedules = Vector{Dict{DTask,Processor}}(undef, n_partitions) + proc_to_scope_lfus = [BasicLFUCache{Processor,AbstractScope}(1024) for _ in 1:n_partitions] + shared_state = DataDepsState(DAGSpec()) + write_num = 1 + # Shared state already tracks global ownership/history like flat + # `distribute_tasks!`. Do not pass `registry` as `ownership`: sync/commit + # would fight the single-state history, and the epilogue must then ignore + # the registry too (see `distribute_tasks_hierarchical!`). + ownership = nothing + + for pid in 1:n_partitions + local_procs = partition_procs[pid] + if isempty(partitions[pid]) || isempty(local_procs) + temp_queues[pid] = DataDepsTaskQueue(wait_all_queue; scheduler=similar(queue.scheduler)) + local_scopes[pid] = DefaultScope() + schedules[pid] = Dict{DTask,Processor}() + continue + end + local_scope = UnionScope(map(ExactScope, local_procs)) + local_scopes[pid] = local_scope + locked_queue = LockedEnqueueQueue(wait_all_queue, queue_lock) + temp_queues[pid] = DataDepsTaskQueue(locked_queue; scheduler=similar(queue.scheduler)) + partition_pairs = DTaskPair[seen_tasks[v] for v in partitions[pid]] + _pdag, schedule = datadeps_build_schedule!(temp_queues[pid].scheduler, partition_pairs, + local_procs, local_scope; region_uids) + schedules[pid] = schedule + end + + topo = try + topological_sort_by_dfs(dag) + catch + collect(vertices(dag)) + end + + with_options(; task_queue=LockedEnqueueQueue(wait_all_queue, queue_lock)) do + for v in topo + pid = vertex_to_partition[v] + local_procs = partition_procs[pid] + isempty(local_procs) && continue + write_num = _schedule_vertex!( + v, pid, temp_queues[pid], shared_state, local_procs, + local_scopes[pid], dag, seen_tasks, vertex_to_partition, + schedules[pid], proc_to_scope_lfus[pid], write_num, ownership) + end + end + + return DataDepsState[shared_state] +end + +struct LockedEnqueueQueue <: AbstractTaskQueue + inner::AbstractTaskQueue + lock::ReentrantLock +end +function enqueue!(leq::LockedEnqueueQueue, pair::DTaskPair) + @lock leq.lock enqueue!(leq.inner, pair) +end +function enqueue!(leq::LockedEnqueueQueue, pairs::Vector{DTaskPair}) + @lock leq.lock enqueue!(leq.inner, pairs) +end + +# Parallel partition scheduling wraps errors in TaskFailedException / +# CompositeException via `@sync`/`Threads.@spawn`. Unwrap so callers and tests +# see the root `SchedulingException` / scheduler error. +function _unwrap_partition_exception(e) + while true + if e isa CompositeException && !isempty(e.exceptions) + e = e.exceptions[1] + elseif e isa TaskFailedException + e = something(e.task.exception, e) + else + return e + end + end +end + +""" + distribute_tasks_hierarchical!(queue) + +Main entry point for hierarchical scheduling. Runs the 4-phase pipeline: +1. Parallel aliasing construction +2. DAG construction +3. Partitioning (by Distributed-worker / MPI-rank affinity, or across local + procs on one owner) +4. Per-partition scheduling via `distribute_task!` + +Each partition computes its *own* partition-local AOT schedule over just its +processors (see `schedule_partition_full!`); there is no global AOT pass or +global synchronization, which is what allows this to scale under Distributed. +AOT scheduling here therefore need not match the flat `distribute_tasks!` path +exactly. Both drivers use `distribute_task!` for argument preparation and +`DataDepsScheduler` dispatch; the old single-worker "batch enqueue with DAG +syncdeps only" path is intentionally not used (it skipped `distribute_task!` and +broke `ChunkView` / custom schedulers). + +Under uniform execution (MPI), Phase 4 runs sequentially on the root task so +SPMD tag / `MPIRefID` allocation stays deterministic across ranks. +""" +function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) + seen_tasks = queue.seen_tasks + if isempty(seen_tasks) + return + end + + # Match flat `distribute_tasks!`: acceleration-aware processor enumeration + # (Distributed `OSProc`s, or MPI `MPIOSProc`s → `MPIProcessor`s). + accel = current_acceleration() + accel_procs = filter(procs(Sch.eager_context())) do proc + accel_matches_proc(accel, proc) + end + all_procs = unique(vcat([collect(get_processors(gp)) for gp in accel_procs]...)) + select_processors_uniform!(all_procs, accel) + scope = get_compute_scope() + filter!(proc->proc_in_scope(proc, scope), all_procs) + if isempty(all_procs) + throw(Sch.SchedulingException("No processors available, try widening scope")) + end + if uniform_execution(accel) + for proc in all_procs + check_uniform(proc) + end + end + + # With a single owner, partitioning buys parallel planning and nothing else + # (see `HIER_TASKS_PER_PARTITION`). A region too small to fill more than one + # partition would pay for the extra prescan/DAG/partition phases and the + # per-partition state without ever running two planners at once, so plan it + # flat instead. Multi-owner regions always partition: there the split is by + # data ownership, not a throughput heuristic. + if !uniform_execution(accel) && + allequal(partition_affinity_id(proc) for proc in all_procs) && + single_owner_partition_count(length(seen_tasks), length(all_procs)) == 1 + return distribute_tasks!(queue) + end + + # All in-region task uids, so each partition's local AOT-schedule builder + # can tell same-region producers (which it must not `fetch`) apart from + # already-materialized external values. + region_uids = Set{UInt}(pair.task.uid for pair in seen_tasks) + + # Phase 1: Collect arguments and compute aliasing in parallel + task_metas, unique_arg_ws = collect_aliased_args(seen_tasks) + _lookup, ainfos_overlaps, arg_to_ainfo = build_aliasing_parallel(unique_arg_ws) + + # Phase 2: Build dependency DAG + dag = build_dependency_dag(task_metas, arg_to_ainfo, ainfos_overlaps) + + # Phase 3: Partition the DAG + vertex_to_partition, n_partitions, partition_procs, _multi_owner = + partition_dag(dag, task_metas, all_procs) + + # Detect backing chunks shared across partitions in different memory spaces. + # These need runtime ownership transfer to avoid split-brain concurrent + # writes; `nothing` when all partitions share one space (single-owner). + partition_space = MemorySpace[only(memory_spaces(first(pp))) for pp in partition_procs] + registry = build_shared_chunk_registry(task_metas, vertex_to_partition, partition_space) + + # Vertices whose `distribute_task!` will `fetch` an in-region producer's + # value, and so cannot run until that producer has actually been submitted. + value_dep_verts = Set{Int}(v for v in 1:length(task_metas) + if !isempty(task_metas[v].value_deps)) + + # Group vertices by partition + partitions = [Int[] for _ in 1:n_partitions] + for v in 1:length(seen_tasks) + pid = vertex_to_partition[v] + push!(partitions[pid], v) + end + + queue_lock = ReentrantLock() + task_submitted = [Base.Event() for _ in 1:length(seen_tasks)] + wait_all_queue = get_options(:task_queue) + + # Phase 4: per-partition scheduling. Under Distributed, each partition runs + # on its own task concurrently (see the header note). Under uniform + # execution (MPI), planning stays sequential on the root task in global + # topological order so `to_tag` / generic `MPIRefID` allocation is + # SPMD-deterministic and cross-partition deps cannot deadlock. + mpi_shared_state = uniform_execution(accel) + partition_states = try + if mpi_shared_state + schedule_partitions_sequential!( + queue, queue_lock, partitions, dag, seen_tasks, + partition_procs, vertex_to_partition, region_uids, registry, + wait_all_queue) + else + states = Vector{DataDepsState}(undef, n_partitions) + @sync for pid in 1:n_partitions + Threads.@spawn begin + batch_queue = BatchedEnqueueQueue(wait_all_queue, queue_lock) + # Copy tasks spawned from within `distribute_task!` pick the + # queue up from options, so they batch alongside their task. + with_options(; task_queue=batch_queue) do + try + states[pid] = schedule_partition_full!( + queue, batch_queue, pid, partitions[pid], + dag, seen_tasks, + partition_procs[pid], vertex_to_partition, + task_submitted, region_uids, value_dep_verts, + registry + ) + finally + flush_batch!(batch_queue) + end + end + end + end + states + end + catch e + rethrow(_unwrap_partition_exception(e)) + end + + # Sequential MPI path returns a one-element shared-state vector and does not + # commit into `registry` (`ownership=nothing` there); copy-back must therefore + # ignore the registry and use the shared state's full history (same as flat). + # Passing the registry would skip every multi-partition chunk: the registry + # path needs `owner_state` (never set), and the private path skips shared keys. + # Distributed keeps per-partition states and needs the registry for coherent + # cross-partition write-back. + epilogue_registry = mpi_shared_state ? nothing : registry + _hierarchical_copy_from_and_free!(partition_states, length(partition_states), epilogue_registry) +end + +function _hierarchical_max_write_num(state::DataDepsState, arg_w::ArgumentWrapper) + wn = 0 + if haskey(state.arg_history, arg_w) + for entry in state.arg_history[arg_w] + wn = max(wn, entry.write_num) + end + end + return wn +end + +function _hierarchical_copy_from!(state::DataDepsState, arg_w::ArgumentWrapper, write_num::Int) + haskey(state.arg_origin, arg_w.arg) || return + origin_space = state.arg_origin[arg_w.arg] + remainder, _ = compute_remainder_for_arg!(state, origin_space, arg_w, write_num) + if remainder isa MultiRemainderAliasing + origin_scope = UnionScope(map(ExactScope, collect(processors(origin_space)))...) + enqueue_remainder_copy_from!(state, origin_space, arg_w, remainder, origin_scope, write_num) + elseif remainder isa FullCopy + origin_scope = UnionScope(map(ExactScope, collect(processors(origin_space)))...) + enqueue_copy_from!(state, origin_space, arg_w, origin_scope, write_num) + end + return +end + +# Deterministic sort key for shared-chunk registry / free-list iteration so +# SPMD ranks enqueue copy-back and free tasks in the same order. +_hierarchical_chunk_sort_key(chunk::Chunk) = (short_name(chunk.space), hash(chunk.handle)) +_hierarchical_chunk_sort_key(chunk) = ("", _identity_hash(chunk)) + +function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsState}, n_partitions::Int, + registry::Union{SharedChunkRegistry,Nothing}) + # 1. Shared chunks: write back from the registry's authoritative owner state + # (whose per-partition history is coherent), not the cross-partition + # max-`write_num` heuristic below (per-partition write_nums are not + # comparable across partitions). + if registry !== nothing + shared_chunks = sort!(collect(keys(registry.entries)); by=_hierarchical_chunk_sort_key) + for chunk in shared_chunks + entry = registry.entries[chunk] + state = entry.owner_state + state === nothing && continue # never written + entry.owner_space == entry.origin_space && continue # already in-place at origin + arg_ws = sort!(ArgumentWrapper[arg_w for (arg_w, _) in state.arg_owner if arg_w.arg === chunk]; + by=arg_w->arg_w.hash) + for arg_w in arg_ws + _hierarchical_copy_from!(state, arg_w, _hierarchical_max_write_num(state, arg_w) + 1) + end + end + end + + # 2. Private chunks: the last writer across partitions is the authoritative + # owner (shared chunks are skipped here, handled above). + merged_arg_owner = Dict{ArgumentWrapper, Tuple{MemorySpace, Int, DataDepsState}}() + for pid in 1:n_partitions + state = partition_states[pid] + for (arg_w, space) in state.arg_owner + is_shared_chunk(registry, arg_w.arg) && continue + wn = _hierarchical_max_write_num(state, arg_w) + if !haskey(merged_arg_owner, arg_w) || wn > merged_arg_owner[arg_w][2] + merged_arg_owner[arg_w] = (space, wn, state) + end + end + end + + for arg_w in sort(collect(keys(merged_arg_owner)); by=arg_w->arg_w.hash) + _space, wn, state = merged_arg_owner[arg_w] + _hierarchical_copy_from!(state, arg_w, wn + 1) + end + + # 3. Free Datadeps-allocated slots. For shared chunks, also sync on the final + # global writer: an intermediate owner's slot may still be read by a + # cross-partition boundary copy recorded in *another* partition's state, + # and the final writer transitively depends on all such copies. + # Iteration is sorted for SPMD-uniform free-task tagging/enqueue order. + for pid in 1:n_partitions + state = partition_states[pid] + obj_cache = unwrap(state.ainfo_backing_chunk) + write_num = typemax(Int) - 1 + remote_spaces = sort!(collect(keys(obj_cache.values)); by=short_name) + for remote_space in remote_spaces + space_entries = sort!(collect(obj_cache.values[remote_space]); by=p->hash(p.first)) + for (ainfo, remote_arg) in space_entries + if !is_original(obj_cache, remote_space, ainfo) + remote_proc = first(processors(remote_space)) + free_scope = ExactScope(remote_proc) + free_syncdeps = Set{ThunkSyncdep}() + if haskey(state.ainfo_arg, ainfo) + get_write_deps!(state, remote_space, ainfo, write_num, free_syncdeps) + end + if registry !== nothing + orig = get(state.remote_arg_to_original, remote_arg, nothing) + if orig !== nothing + entry = get(registry.entries, orig, nothing) + if entry !== nothing && entry.owner_task !== nothing + push!(free_syncdeps, ThunkSyncdep(entry.owner_task)) + end + end + end + Dagger.@spawn scope=free_scope syncdeps=free_syncdeps tag=datadeps_task_tag() Dagger.unsafe_free!(remote_arg) + end + end + end + end +end diff --git a/src/datadeps/queue.jl b/src/datadeps/queue.jl index 9441f3bbe..e7dfceafe 100644 --- a/src/datadeps/queue.jl +++ b/src/datadeps/queue.jl @@ -65,7 +65,8 @@ function spawn_datadeps(f::Base.Callable; static::Bool=true, traversal::Symbol=:inorder, scheduler::Union{DataDepsScheduler,Nothing}=nothing, aliasing::Bool=true, - launch_wait::Union{Bool,Nothing}=nothing) + launch_wait::Union{Bool,Nothing}=nothing, + hierarchical::Union{Bool,Nothing}=nothing) if !static throw(ArgumentError("Dynamic scheduling is no longer available")) end @@ -78,22 +79,127 @@ function spawn_datadeps(f::Base.Callable; static::Bool=true, wait_all(; check_errors=true) do scheduler = something(scheduler, DATADEPS_SCHEDULER[], RoundRobinScheduler()) launch_wait = something(launch_wait, DATADEPS_LAUNCH_WAIT[], false)::Bool + hierarchical = something(hierarchical, DATADEPS_HIERARCHICAL[], true)::Bool + run_distribute = queue -> begin + if hierarchical + distribute_tasks_hierarchical!(queue) + else + distribute_tasks!(queue) + end + end if launch_wait result = spawn_bulk() do queue = DataDepsTaskQueue(get_options(:task_queue); scheduler) with_options(f; task_queue=queue) - distribute_tasks!(queue) + run_distribute(queue) end else queue = DataDepsTaskQueue(get_options(:task_queue); scheduler) result = with_options(f; task_queue=queue) - distribute_tasks!(queue) + run_distribute(queue) end return result end end const DATADEPS_SCHEDULER = ScopedValue{Union{DataDepsScheduler,Nothing}}(nothing) const DATADEPS_LAUNCH_WAIT = ScopedValue{Union{Bool,Nothing}}(nothing) +const DATADEPS_HIERARCHICAL = ScopedValue{Union{Bool,Nothing}}(nothing) + +"Returns `true` if any argument of `spec` is a same-region (`region_uids`) `DTask`." +function _spec_has_region_dtask_arg(spec::DTaskSpec, region_uids::Set{UInt}) + for _arg in spec.fargs + arg, _ = unwrap_inout(value(_arg)) + if arg isa DTask && arg.uid in region_uids + return true + end + end + return false +end + +""" + datadeps_build_schedule!(scheduler, pairs, all_procs, all_scope; + region_uids=nothing) -> (dag_spec, schedule) + +Builds a DAG spec from `pairs` (in submission order) and computes an AOT +processor assignment (`schedule::Dict{DTask,Processor}`) over `all_procs` / +`all_scope`, reusing a cached schedule when an equivalent DAG has been seen +before (and otherwise caching the freshly-computed one for future reuse). + +The schedule cache is task-local per scheduler type, so callers running on +distinct tasks (e.g. the per-partition scheduling tasks of the hierarchical +path) each maintain an independent cache. Schedulers that don't implement +`datadeps_schedule_dag_aot!` leave the schedule empty, in which case tasks fall +back to JIT scheduling in `distribute_task!`. + +`region_uids`, when provided, is the set of *all* in-region task uids. It lets +this function bail (fall back to JIT for the remaining tasks) before +`dag_add_task!` would try to `fetch` an unlaunched same-region producer that is +absent from this (sub)set of `pairs` -- which is essential when `pairs` covers +only one partition of the region, since a producer in another partition would +otherwise be treated as an external value and block forever. For the flat path +(where `pairs` is the whole region), `dag_add_task!`'s own in-DAG check already +handles this and `region_uids` can be omitted. + +Used by both the flat (`distribute_tasks!`, whole region) and hierarchical +(`distribute_tasks_hierarchical!`, per partition) paths. +""" +function datadeps_build_schedule!(scheduler::DataDepsScheduler, + pairs::Vector{DTaskPair}, + all_procs, all_scope; + region_uids::Union{Set{UInt},Nothing}=nothing) + # Compute the DAG spec. `dag_add_task!` ignores the state argument (it only + # inspects the task spec), so a throwaway state is fine here. + dag_spec = DAGSpec() + dummy_state = DataDepsState(dag_spec) + for (spec, task) in pairs + if region_uids !== nothing && _spec_has_region_dtask_arg(spec, region_uids) + # Depends on an in-region producer that may live outside `pairs`; + # defer the rest to JIT scheduling. + break + end + if !dag_add_task!(dag_spec, dummy_state, spec, task) + # This task depends on an in-region task's result; defer the rest + # to JIT scheduling (they won't appear in `schedule`). + break + end + end + + # Attempt to find any matching DAG specs and reuse their schedule + schedule = Dict{DTask, Processor}() + schedule_cache = datadeps_schedule_cache(scheduler) + cache_hit = false + for (other_spec, spec_schedule) in schedule_cache + if datadeps_dag_equivalent(scheduler, dag_spec, other_spec) + @dagdebug nothing :spawn_datadeps "Found matching DAG spec!" + for (id, proc) in spec_schedule.id_to_proc + uid = dag_spec.id_to_uid[id] + task_idx = findfirst(spec_task -> spec_task.task.uid == uid, pairs) + task = pairs[task_idx].task + schedule[task] = proc + end + cache_hit = true + break + end + end + + if !cache_hit && !isempty(dag_spec) + # Compute a fresh AOT schedule (no-op for schedulers that fall back + # to JIT in distribute_task!) + datadeps_schedule_dag_aot!(scheduler, schedule, dag_spec, all_procs, all_scope) + + # Persist the schedule for reuse by future equivalent DAGs + if !isempty(schedule) + spec_schedule = DAGSpecSchedule() + for (task, proc) in schedule + id = dag_spec.uid_to_id[task.uid] + spec_schedule.id_to_proc[id] = proc + end + push!(schedule_cache, dag_spec => spec_schedule) + end + end + + return dag_spec, schedule +end # Current task uid, propagated into `tochunk` so uniform-execution backends # (MPIExt) can derive deterministic, rank-agreed handle IDs. Core datadeps sets @@ -133,17 +239,21 @@ function distribute_tasks!(queue::DataDepsTaskQueue) all_scope = UnionScope(map(ExactScope, all_procs)) exec_spaces = unique(vcat(map(proc->collect(memory_spaces(proc)), all_procs)...)) - # Round-robin assign tasks to processors upper_queue = get_options(:task_queue) + # Compute the DAG spec and an AOT processor assignment (reused from cache + # when an equivalent DAG has been seen before). + dag_spec, schedule = datadeps_build_schedule!(queue.scheduler, queue.seen_tasks, all_procs, all_scope) + # Start launching tasks and necessary copies - state = DataDepsState() + state = DataDepsState(dag_spec) write_num = 1 proc_to_scope_lfu = BasicLFUCache{Processor,AbstractScope}(1024) for pair in queue.seen_tasks spec = pair.spec task = pair.task - write_num = distribute_task!(queue, state, all_procs, all_scope, spec, task, spec.fargs, proc_to_scope_lfu, write_num) + proc = get(schedule, task, nothing) + write_num = distribute_task!(queue, state, all_procs, all_scope, spec, task, spec.fargs, proc_to_scope_lfu, write_num; proc) end # Copy args from remote to local @@ -153,16 +263,13 @@ function distribute_tasks!(queue::DataDepsTaskQueue) check_uniform(arg_w) arg = arg_w.arg origin_space = state.arg_origin[arg] - # When the origin still holds a fully-current replica (the argument was - # only read, or copies merely propagated it), the write-back is elided. - # This is only safe here at region end: mid-region, the copy tasks also - # serialize readers against later writers, so they must not be skipped. - current = get(state.arg_current, arg_w, nothing) - if current !== nothing && origin_space in current - remainder = NoAliasing() - else - remainder, _ = compute_remainder_for_arg!(state, origin_space, arg_w, write_num) - end + # Always compute write-back from history. Eliding via `arg_current` + # (`origin in arg_current` ⇒ skip) is unsound under multi-worker + # Distributed: Krylov `similar` workspaces can leave origin marked + # current while a remote true write holds the only initialized bytes. + # Read-only copy-ins may therefore schedule a redundant write-back + # (copy history advances `arg_owner`); that is intentional and harmless. + remainder, _ = compute_remainder_for_arg!(state, origin_space, arg_w, write_num) if remainder isa MultiRemainderAliasing origin_scope = UnionScope(map(ExactScope, collect(processors(origin_space)))...) enqueue_remainder_copy_from!(state, origin_space, arg_w, remainder, origin_scope, write_num) @@ -245,7 +352,7 @@ map_or_ntuple(f, xs::Vector) = map(f, 1:length(xs)) # N.B. Accept any `Tuple` (typed specs produce heterogeneous tuples of # `TypedArgument{T}`, not a homogeneous `NTuple{N,T}`). @inline map_or_ntuple(@specialize(f), xs::Tuple) = ntuple(f, Val(length(xs))) -function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_procs, all_scope, spec::DTaskSpec{typed}, task::DTask, fargs, proc_to_scope_lfu, write_num::Int) where typed +function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_procs, all_scope, spec::DTaskSpec{typed}, task::DTask, fargs, proc_to_scope_lfu, write_num::Int; proc::Union{Processor,Nothing}=nothing, ownership=nothing) where typed @specialize spec fargs if typed @@ -257,8 +364,14 @@ function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_pr DATADEPS_CURRENT_TASK[] = task task_scope = @something(spec.options.compute_scope, spec.options.scope, DefaultScope()) - scheduler = queue.scheduler - our_proc = datadeps_schedule_task(scheduler, state, all_procs, all_scope, task_scope, spec, task) + + if proc === nothing + # Schedule and JIT assign tasks to processors + our_proc = datadeps_schedule_task_jit!(queue.scheduler, all_procs, all_scope, task_scope, spec, task) + else + # Use the provided processor from AOT scheduling + our_proc = proc + end @assert our_proc in all_procs our_space = only(memory_spaces(our_proc)) check_uniform(our_proc) @@ -305,6 +418,14 @@ function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_pr return end + # Hierarchical scheduling only: for shared backing chunks whose current + # version was produced by another partition, seed this partition's state so + # the copy-to below pulls a fresh whole-chunk copy from the true owner (and + # syncs on its producer). No-op on the flat path (`ownership === nothing`). + if ownership !== nothing + _sync_incoming_ownership!(state, ownership, our_space, task_arg_ws, write_num) + end + # Copy args from local to remote remote_args = map_or_ntuple(task_arg_ws) do idx arg_ws = task_arg_ws[idx] @@ -433,6 +554,13 @@ function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_pr return end + # Hierarchical scheduling only: publish this task as the new authoritative + # owner of each shared backing chunk it writes, so later cross-partition + # consumers pull the up-to-date version from here. + if ownership !== nothing + _commit_ownership!(state, ownership, our_space, task, task_arg_ws) + end + write_num += 1 DATADEPS_CURRENT_TASK[] = nothing diff --git a/src/datadeps/remainders.jl b/src/datadeps/remainders.jl index 1000471b2..cbd494fef 100644 --- a/src/datadeps/remainders.jl +++ b/src/datadeps/remainders.jl @@ -43,6 +43,7 @@ memory_spans(mra::MultiRemainderAliasing) = vcat(memory_spans.(mra.remainders).. Base.hash(mra::MultiRemainderAliasing, h::UInt) = hash(mra.remainders, hash(MultiRemainderAliasing, h)) Base.:(==)(mra1::MultiRemainderAliasing, mra2::MultiRemainderAliasing) = mra1.remainders == mra2.remainders +# A whole-object copy from the argument's recorded owner space. struct FullCopy end """ @@ -87,9 +88,34 @@ function compute_remainder_for_arg!(state::DataDepsState, target_space::MemorySpace, arg_w::ArgumentWrapper, write_num::Int; compute_syncdeps::Bool=true) + owner_space = state.arg_owner[arg_w] + + # Whole-object containers (e.g. `DSparseArray`) reallocate their entire + # storage on every write, so they can never be *partially* current in a + # memory space -- the whole object is current in exactly one space: the most + # recent writer (or the origin owner if never written). There is thus never a + # meaningful span "remainder" to compute, so we short-circuit to a whole-object + # `FullCopy` from the owner space (or `NoAliasing` if the target already holds + # it). This avoids the span-diff computation and the construction + transfer of + # a `RemainderAliasing` entirely. We detect this from the aliasing info, which + # is a bare `ObjectAliasing` for such containers (see `aliases_as_whole`). + # + # `arg_owner` is the most-recent *direct* writer of `arg_w`, while the merged + # `arg_history` records the latest write through *any* overlapping argument. + # For every path Datadeps currently exercises a whole-object container is only + # ever reached through a single argument, so these agree -- the assertion below + # guards that invariant. They could only diverge under cross-argument aliasing + # of a whole-object container (the same container written through one argument + # and read through a different, overlapping one); supporting that would require + # sourcing the copy from `history[end].space` instead of `arg_owner`. + if aliases_as_whole(aliasing!(state, target_space, arg_w)) + history = state.arg_history[arg_w] + @assert isempty(history) || history[end].space == owner_space "whole-object container reached through a different overlapping argument than its owner; cross-argument aliasing of whole-object containers is unsupported" + return owner_space == target_space ? (NoAliasing(), 0) : (FullCopy(), 0) + end + spaces_set = Set{MemorySpace}() push!(spaces_set, target_space) - owner_space = state.arg_owner[arg_w] push!(spaces_set, owner_space) @label restart @@ -152,6 +178,7 @@ function compute_remainder_for_arg!(state::DataDepsState, break end + other_entry = nothing if idx > 0 other_entry = state.arg_history[arg_w][idx] other_ainfo = other_entry.ainfo @@ -199,7 +226,17 @@ function compute_remainder_for_arg!(state::DataDepsState, has_overlap = schedule_remainder!(tracker_other_space[1], other_space_idx, target_space_idx, remainder, other_many_spans) if compute_syncdeps && has_overlap @assert haskey(state.ainfos_owner, other_ainfo) "[idx $idx] ainfo $(typeof(other_ainfo)) has no owner" - get_read_deps!(state, other_space, other_ainfo, write_num, tracker_other_space[3]) + if other_entry !== nothing + # Wait on the producer recorded in history. Re-resolving via + # live `ainfos_owner`/`get_read_deps!` is not equivalent: + # ownership of `other_ainfo` may have moved to a later copy, or + # the producer may only be registered on an overlapping ainfo + # missing from `ainfos_overlaps[other_ainfo]`. + push!(tracker_other_space[3], ThunkSyncdep(other_entry.task)) + else + # idx==0 owner fallback: no HistoryEntry, use live writer set. + get_read_deps!(state, other_space, other_ainfo, write_num, tracker_other_space[3]) + end push!(tracker_other_space[2], other_ainfo) end end @@ -437,6 +474,9 @@ function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space: multi_span_copy!(to_s, from_s, dep_mod.spans) return end + # N.B. Whole-object containers (e.g. `DSparseArray`) never reach here: their + # storage reallocates on write, so they can't be span-copied. They are routed + # to a whole-object `FullCopy` in `compute_remainder_for_arg!` instead. # Gather source spans into a host buffer (device arrays pack via KA first). # Pin when the source is device memory so DtoH hits page-locked host RAM. @@ -548,17 +588,3 @@ function write_remainder!(copies::Vector{UInt8}, copies_offset::UInt64, to::Base write_remainder!(copies, copies_offset, to[], to_ptr, n) end end - -function find_object_holding_ptr(A::SparseMatrixCSC, ptr::UInt64) - span = LocalMemorySpan(pointer(A.nzval), length(A.nzval)*sizeof(eltype(A.nzval))) - if span_start(span) <= ptr <= span_end(span) - return A.nzval - end - span = LocalMemorySpan(pointer(A.colptr), length(A.colptr)*sizeof(eltype(A.colptr))) - if span_start(span) <= ptr <= span_end(span) - return A.colptr - end - span = LocalMemorySpan(pointer(A.rowval), length(A.rowval)*sizeof(eltype(A.rowval))) - @assert span_start(span) <= ptr <= span_end(span) "Pointer $ptr not found in SparseMatrixCSC" - return A.rowval -end diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 62a350c72..af72869df 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -1,10 +1,259 @@ -abstract type DataDepsScheduler end +### DAG Analysis ### + +Base.length(spec::DAGSpec) = nv(spec.g) +Base.isempty(spec::DAGSpec) = length(spec) == 0 + +function dag_add_task!(dspec::DAGSpec, state, tspec::DTaskSpec, task::DTask) + # Check if this task depends on any other tasks within the DAG, + # which we are not yet ready to handle + for (idx, _arg) in enumerate(tspec.fargs) + arg, deps = unwrap_inout(value(_arg)) + for (dep_mod, readdep, writedep) in deps + if arg isa DTask + if arg.uid in keys(dspec.uid_to_id) + # Within-DAG dependency, bail out + return false + end + end + end + end + + add_vertex!(dspec.g) + id = nv(dspec.g) + + # Record function signature + dspec.id_to_functype[id] = chunktype(tspec.fargs[1]) + argtypes = DatadepsArgSpec[] + for (idx, _arg) in enumerate(tspec.fargs) + arg, deps = unwrap_inout(value(_arg)) + pos = raw_position(_arg) + for (dep_mod, readdep, writedep) in deps + if arg isa DTask + #= TODO: Re-enable this when we can handle within-DAG dependencies + if arg.uid in keys(dspec.uid_to_id) + # Within-DAG dependency + arg_id = dspec.uid_to_id[arg.uid] + push!(dspec.id_to_argtypes[arg_id], DatadepsArgSpec(pos, DTaskDAGID{arg_id}, dep_mod, UnknownAliasing())) + add_edge!(dspec.g, arg_id, id) + continue + end + =# + + # External DTask, so fetch this and track it as a raw value + arg = fetch(arg; raw=true) + end + ainfo = aliasing(arg, dep_mod) + # FIXME: Generate syncdeps and add edges + push!(argtypes, DatadepsArgSpec(pos, typeof(arg), dep_mod, ainfo)) + end + end + dspec.id_to_argtypes[id] = argtypes + dspec.id_to_scope[id] = @something(tspec.options.compute_scope, + tspec.options.scope, + DefaultScope()) + + # FIXME: Record syncdeps + dspec.id_to_uid[id] = task.uid + dspec.uid_to_id[task.uid] = id + dspec.id_to_spec[id] = tspec + dspec.id_to_task[id] = task + + return true +end +function dag_build_edges!(dag_spec::DAGSpec) + # Naively build edges based on exact argument comparisons (no aliasing) + arg_writes = Dict{Any,Int}() + for idx in 1:nv(dag_spec.g) + task_spec = dag_spec.id_to_spec[idx] + + # Get the raw arguments for this task + task_raw_args = Vector{Any}() + for arg in task_spec.fargs + arg, deps = unwrap_inout(arg) + for (dep_mod, readdep, writedep) in deps + # Get the raw argument + raw_arg = arg isa DTask ? fetch(arg; raw=true) : arg + + # Did any previous task write to this argument? + if haskey(arg_writes, raw_arg) && arg_writes[raw_arg] != idx + prev_task_id = arg_writes[raw_arg] + add_edge!(dag_spec.g, prev_task_id, idx) + end + + if writedep + # Record this write + arg_writes[raw_arg] = idx + end + end + end + end +end +function dag_has_task(dspec::DAGSpec, task::DTask) + return task.uid in keys(dspec.uid_to_id) +end + +### DAGSpec Equivalence (scheduler-dispatched) ### + +""" + datadeps_dag_equivalent(scheduler, dspec1::DAGSpec, dspec2::DAGSpec) -> Bool + +Returns `true` if a schedule cached for `dspec2` may safely be reused for +`dspec1`, as judged by `scheduler`. This is the top-level entry point used by +`distribute_tasks!` when looking up a cached schedule. + +The default implementation requires: +- Same number of vertices and edges in the dependency graph +- Identical edge sets (per-vertex `outneighbors`) +- Per-vertex agreement on function type and on each task's compute scope +- Per-vertex agreement on the multiset of argspecs, compared via + `datadeps_argspec_equivalent` + +Schedulers that want completely custom matching (or want to opt out of caching +entirely) can override this directly. Schedulers that only want to tweak how +individual arguments are compared should instead override +`datadeps_argspec_equivalent` or `datadeps_ainfo_equivalent`. +""" +function datadeps_dag_equivalent(scheduler::DataDepsScheduler, + dspec1::DAGSpec, dspec2::DAGSpec) + # Graph shape + nv(dspec1.g) == nv(dspec2.g) || return false + ne(dspec1.g) == ne(dspec2.g) || return false + + @inbounds for id in 1:nv(dspec1.g) + # outneighbors covers all edges; inneighbors would be redundant + outneighbors(dspec1.g, id) == outneighbors(dspec2.g, id) || return false + + # Function type must match exactly + dspec1.id_to_functype[id] === dspec2.id_to_functype[id] || return false + + # Per-task compute scope must match (different scopes can produce + # different schedules and must not be aliased) + dspec1.id_to_scope[id] == dspec2.id_to_scope[id] || return false + + # Argspecs must match as a multiset (Deps can put multiple argspecs at + # the same position) + _argspecs_equivalent(scheduler, + dspec1.id_to_argtypes[id], + dspec2.id_to_argtypes[id]) || return false + end + + return true +end + +# Backwards-compatible default `==`: use the no-scheduler default (i.e. as if +# all schedulers behaved like the base `DataDepsScheduler`). The runtime path +# in `distribute_tasks!` calls `datadeps_dag_equivalent` directly. +Base.:(==)(dspec1::DAGSpec, dspec2::DAGSpec) = + datadeps_dag_equivalent(_DefaultEquivalenceScheduler(), dspec1, dspec2) + +# A private marker scheduler used to provide a default for `Base.:(==)` on +# `DAGSpec`. Not exported and not intended for direct use. +struct _DefaultEquivalenceScheduler <: DataDepsScheduler end + +""" + datadeps_argspec_equivalent(scheduler, + a1::DatadepsArgSpec, + a2::DatadepsArgSpec) -> Bool + +Returns `true` if argspecs `a1` and `a2` are interchangeable for the purposes +of `scheduler`'s cached-schedule lookup. The default requires equal positions, +equal value types, equal dep_mods, and structurally-equivalent ainfos (per +`datadeps_ainfo_equivalent`). +""" +function datadeps_argspec_equivalent(scheduler::DataDepsScheduler, + a1::DatadepsArgSpec, a2::DatadepsArgSpec) + a1.pos == a2.pos || return false + a1.value_type === a2.value_type || return false + a1.dep_mod === a2.dep_mod || return false + return datadeps_ainfo_equivalent(scheduler, a1.ainfo, a2.ainfo) +end + +""" + datadeps_ainfo_equivalent(scheduler, + a1::AbstractAliasing, + a2::AbstractAliasing) -> Bool + +Returns `true` if aliasings `a1` and `a2` are interchangeable for the purposes +of `scheduler`'s cached-schedule lookup. The default uses +`equivalent_structure`, which compares shape/strides/lengths while ignoring +absolute pointer addresses, enabling reuse across re-allocations. + +Schedulers can override this to choose a different equivalence strategy, e.g. +pointer-identical (strictest), locality-only (memory-space only), or fully +permissive. +""" +datadeps_ainfo_equivalent(::DataDepsScheduler, + a1::AbstractAliasing, a2::AbstractAliasing) = + equivalent_structure(a1, a2) + +# Compare two argspec vectors as multisets, since `Deps` can place multiple +# argspecs at the same position. We pair each argspec in `as1` with a not-yet- +# matched argspec in `as2`; both lists must be exhausted simultaneously. +function _argspecs_equivalent(scheduler::DataDepsScheduler, + as1::Vector{DatadepsArgSpec}, + as2::Vector{DatadepsArgSpec}) + length(as1) == length(as2) || return false + n = length(as1) + n == 0 && return true + matched = falses(n) + @inbounds for a1 in as1 + found = false + for j in 1:n + matched[j] && continue + if datadeps_argspec_equivalent(scheduler, a1, as2[j]) + matched[j] = true + found = true + break + end + end + found || return false + end + return true +end + +### Schedule Cache (scheduler-owned) ### + +struct DAGSpecSchedule + id_to_proc::Dict{Int, Processor} + DAGSpecSchedule() = new(Dict{Int, Processor}()) +end + +# Per-scheduler-type cache. Each entry in the inner Vector is a (DAGSpec => +# DAGSpecSchedule) pair recorded by a prior call. The outer Dict partitions +# the cache by `typeof(scheduler)` so schedulers don't contaminate each other. +const DATADEPS_DAG_SPECS = + TaskLocalValue{Dict{Type, Vector{Pair{DAGSpec, DAGSpecSchedule}}}}( + ()->Dict{Type, Vector{Pair{DAGSpec, DAGSpecSchedule}}}()) + +""" + datadeps_schedule_cache(scheduler) -> Vector{Pair{DAGSpec, DAGSpecSchedule}} + +Returns the schedule cache that `scheduler` should consult for prior schedules +and append newly-computed schedules to. The default implementation returns a +task-local, per-scheduler-type cache. + +Override this to implement custom caching strategies (e.g., bounded LRU, no +cache at all, cross-task shared cache). +""" +function datadeps_schedule_cache(scheduler::DataDepsScheduler) + cache_by_type = DATADEPS_DAG_SPECS[] + return get!(Vector{Pair{DAGSpec, DAGSpecSchedule}}, + cache_by_type, typeof(scheduler)) +end + +### JIT Schedulers ### + +# Default for user-defined schedulers with a zero-arg constructor. Schedulers +# that carry mutable state should specialize `similar` to return a fresh shard +# (used when hierarchical scheduling clones a scheduler per partition). +Base.similar(s::DataDepsScheduler) = typeof(s)() mutable struct RoundRobinScheduler <: DataDepsScheduler proc_idx::Int RoundRobinScheduler() = new(1) end -function datadeps_schedule_task(sched::RoundRobinScheduler, state::DataDepsState, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) +Base.similar(::RoundRobinScheduler) = RoundRobinScheduler() +function datadeps_schedule_task_jit!(sched::RoundRobinScheduler, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) proc_idx = sched.proc_idx our_proc = all_procs[proc_idx] if task_scope == all_scope @@ -24,7 +273,8 @@ function datadeps_schedule_task(sched::RoundRobinScheduler, state::DataDepsState end struct NaiveScheduler <: DataDepsScheduler end -function datadeps_schedule_task(sched::NaiveScheduler, state::DataDepsState, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) +Base.similar(::NaiveScheduler) = NaiveScheduler() +function datadeps_schedule_task_jit!(sched::NaiveScheduler, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) raw_args = map(arg->tochunk(value(arg)), spec.fargs) our_proc = remotecall_fetch(1, all_procs, raw_args) do all_procs, raw_args Sch.init_eager() @@ -59,7 +309,8 @@ struct UltraScheduler <: DataDepsScheduler Dict{MemorySpace,Int}()) end end -function datadeps_schedule_task(sched::UltraScheduler, state::DataDepsState, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) +Base.similar(::UltraScheduler) = UltraScheduler() +function datadeps_schedule_task_jit!(sched::UltraScheduler, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) args = Base.mapany(spec.fargs) do arg pos, data = arg data, _ = unwrap_inout(data) @@ -121,3 +372,61 @@ function datadeps_schedule_task(sched::UltraScheduler, state::DataDepsState, all return our_proc end + +### AOT Schedulers ### + +function datadeps_schedule_dag_aot!(scheduler, schedule, dag_spec, all_procs, all_scope) + # Fallback to JIT scheduling (done in distribute_task!) + return +end + +struct LayeredScheduler <: DataDepsScheduler end +function datadeps_schedule_dag_aot!(scheduler::LayeredScheduler, schedule, dag_spec, all_procs, all_scope) + layer = 1 + layer_data = Vector{Any}() + layers = Vector{Vector{Int}}() + push!(layers, Vector{Int}()) + for idx in 1:nv(dag_spec.g) + spec = dag_spec.id_to_spec[idx] + + # Get the raw arguments for this task + task_raw_args = Vector{Any}() + for arg in spec.fargs + arg, deps = unwrap_inout(arg) + for (dep_mod, readdep, writedep) in deps + # We only care about write dependencies + writedep || continue + + # Get the raw argument + raw_arg = arg isa DTask ? fetch(arg; raw=true) : arg + + push!(task_raw_args, raw_arg) + end + end + + # Determine if this task stays in this layer + if any(raw_arg -> raw_arg in layer_data, task_raw_args) + # This argument is already written by a previous task in this layer + # Generate new layer + layer += 1 + empty!(layer_data) + push!(layers, Vector{Int}()) + end + + # Add our data and this task to the current layer + append!(layer_data, task_raw_args) + push!(layers[layer], idx) + end + + # Perform round-robin scheduling within each layer + for layer in layers + sub_scheduler = RoundRobinScheduler() + for idx in layer + spec = dag_spec.id_to_spec[idx] + task = dag_spec.id_to_task[idx] + task_scope = @something(spec.options.compute_scope, spec.options.scope, DefaultScope()) + our_proc = datadeps_schedule_task_jit!(sub_scheduler, all_procs, all_scope, task_scope, spec, task) + schedule[task] = our_proc + end + end +end diff --git a/src/datadeps/types.jl b/src/datadeps/types.jl new file mode 100644 index 000000000..70212e10e --- /dev/null +++ b/src/datadeps/types.jl @@ -0,0 +1,30 @@ +import Graphs: SimpleDiGraph, add_edge!, add_vertex!, inneighbors, outneighbors, nv, ne + +### Scheduling ### + +struct DatadepsArgSpec + pos::Union{Int, Symbol} + value_type::Type + dep_mod::Any + ainfo::AbstractAliasing +end +struct DTaskDAGID{id} end +struct DAGSpec + g::SimpleDiGraph{Int} + id_to_uid::Dict{Int, UInt} + uid_to_id::Dict{UInt, Int} + id_to_functype::Dict{Int, Type} # FIXME: DatadepsArgSpec + id_to_argtypes::Dict{Int, Vector{DatadepsArgSpec}} + id_to_scope::Dict{Int, AbstractScope} + id_to_spec::Dict{Int, DTaskSpec} + id_to_task::Dict{Int, DTask} + DAGSpec() = new(SimpleDiGraph{Int}(), + Dict{Int, UInt}(), Dict{UInt, Int}(), + Dict{Int, Type}(), + Dict{Int, Vector{DatadepsArgSpec}}(), + Dict{Int, AbstractScope}(), + Dict{Int, DTaskSpec}(), + Dict{Int, DTask}()) +end + +abstract type DataDepsScheduler end \ No newline at end of file diff --git a/src/file-io.jl b/src/file-io.jl index 1407fbced..15405d049 100644 --- a/src/file-io.jl +++ b/src/file-io.jl @@ -1,5 +1,4 @@ export save, load -using SparseArrays import MemPool: GenericFileDevice, FileRef struct File diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index 617f7ea8b..0074916ca 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -125,8 +125,27 @@ mutable struct AliasingWrapper <: AbstractAliasing AliasingWrapper(inner::AbstractAliasing) = new(inner, hash(inner)) end memory_spans(x::AliasingWrapper) = memory_spans(x.inner) +""" + equivalent_structure(x::AbstractAliasing, y::AbstractAliasing) -> Bool + +Returns `true` if `x` and `y` describe aliasing regions of the same *structure*, +ignoring absolute memory addresses. Two aliasings that originate from different +allocations of identically-shaped data should compare equal. + +Used as the default `datadeps_ainfo_equivalent` for `DAGSpec` cache lookups, +which lets schedulers reuse a previously-computed schedule across re-runs of +the same algorithm with freshly-allocated inputs. + +Default fallback returns `false` when the two arguments are of differing +concrete subtypes; specialized methods exist for each concrete subtype. +""" +equivalent_structure(x::AbstractAliasing, y::AbstractAliasing) = false equivalent_structure(x::AliasingWrapper, y::AliasingWrapper) = x.hash == y.hash || equivalent_structure(x.inner, y.inner) +equivalent_structure(x::AliasingWrapper, y::AbstractAliasing) = + equivalent_structure(x.inner, y) +equivalent_structure(x::AbstractAliasing, y::AliasingWrapper) = + equivalent_structure(x, y.inner) Base.hash(x::AliasingWrapper, h::UInt64) = hash(x.hash, h) Base.isequal(x::AliasingWrapper, y::AliasingWrapper) = x.hash == y.hash Base.:(==)(x::AliasingWrapper, y::AliasingWrapper) = x.hash == y.hash @@ -331,8 +350,13 @@ end struct NoAliasing <: AbstractAliasing end memory_spans(::NoAliasing) = MemorySpan{CPURAMMemorySpace}[] +equivalent_structure(::NoAliasing, ::NoAliasing) = true struct UnknownAliasing <: AbstractAliasing end memory_spans(::UnknownAliasing) = [MemorySpan{CPURAMMemorySpace}(C_NULL, typemax(UInt))] +# Two UnknownAliasings have no known structure to compare; conservatively +# treat them as equivalent so cache lookups don't permanently miss when +# aliasing analysis falls back to UnknownAliasing. +equivalent_structure(::UnknownAliasing, ::UnknownAliasing) = true error_unknown_aliasing(T) = throw(ConcurrencyViolationError("Cannot resolve aliasing for object of type $T, execution may become sequential")) @@ -357,6 +381,13 @@ Base.:(==)(ca1::CombinedAliasing, ca2::CombinedAliasing) = ca1.sub_ainfos == ca2.sub_ainfos Base.hash(ca1::CombinedAliasing, h::UInt) = hash(ca1.sub_ainfos, hash(CombinedAliasing, h)) +function equivalent_structure(ca1::CombinedAliasing, ca2::CombinedAliasing) + length(ca1.sub_ainfos) == length(ca2.sub_ainfos) || return false + @inbounds for i in eachindex(ca1.sub_ainfos) + equivalent_structure(ca1.sub_ainfos[i], ca2.sub_ainfos[i]) || return false + end + return true +end struct ObjectAliasing{S<:MemorySpace} <: AbstractAliasing ptr::RemotePtr{Cvoid,S} @@ -374,8 +405,14 @@ function memory_spans(oa::ObjectAliasing{S}) where S span = MemorySpan{S}(oa.ptr, oa.sz) return [span] end +equivalent_structure(x::ObjectAliasing{S}, y::ObjectAliasing{S}) where S = + x.sz == y.sz -aliasing(accel::Acceleration, x, T) = aliasing(x, T) +# Acceleration entry point used by Datadeps. Unwrap whole-object containers +# (e.g. views of `DSparseArray`) before computing aliasing. SPMD backends +# (MPI) overload this for `Chunk`/`ChunkView` to broadcast owner-computed +# aliasing; those methods call `aliasing_unwrapped` on the owning rank. +aliasing(accel::Acceleration, x, T) = aliasing_unwrapped(x, T) function aliasing(x, dep_mod) if dep_mod isa Symbol return aliasing(getfield(x, dep_mod)) @@ -411,22 +448,106 @@ end aliasing(::String) = NoAliasing() # FIXME: Not necessarily true aliasing(::Symbol) = NoAliasing() aliasing(::Type) = NoAliasing() + +""" + aliases_as_whole(x) -> Bool + +Whether `x` -- either an object, or an already-computed aliasing -- must be +treated as a single, indivisible unit, i.e. it is never safe to alias (or to +consider current) only *part* of it. + +For an object: containers whose backing storage may be reallocated or resized on +write (such as `DSparseArray`) return `true`. This causes [`aliasing_root`](@ref) +to resolve any view/wrapper of them to the whole container, so all access funnels +through the container's own `aliasing`. By contract, such a container's `aliasing` +must be a bare [`ObjectAliasing`](@ref). + +For an aliasing: a bare `ObjectAliasing` (the aliasing of a whole-object +container) reports `true`. Datadeps uses this to copy such arguments as a whole +(a `FullCopy`) rather than computing per-span remainders, since they can never be +*partially* current in a memory space. +""" +aliases_as_whole(@nospecialize x) = false +aliases_as_whole(ainfo::AliasingWrapper) = aliases_as_whole(ainfo.inner) +aliases_as_whole(::ObjectAliasing) = true + +""" + aliasing_root(x) + +Resolve `x` down to the whole-object container it wraps: peel array wrappers +(views, transposes, adjoints, reshapes, permutations, ...) off of `x` until +reaching a container that must alias as a whole (see [`aliases_as_whole`](@ref)), +and return that container. If no such container is wrapped, return `x` unchanged. + +This relies solely on the standard `Base.parent` interface, so it transparently +handles *any* array wrapper without per-wrapper `aliasing` methods: a new wrapper +type needs no special-casing here, and a new whole-object container only needs to +define [`aliases_as_whole`](@ref) (plus its own `aliasing` method, which must +return a bare `ObjectAliasing`). A trapping `Base.pointer` on such containers +guards against a wrapper slipping through and being misinterpreted as strided +memory. + +!!! warning + The resolved object's *identity* defines its aliasing, so this is only + meaningful in the memory space where `x` physically resides. Never return its + result across a worker boundary and *then* compute `aliasing` -- the transfer + copies the object and changes its aliasing. Use [`aliasing_unwrapped`](@ref), + which fuses the two operations so they always run together. +""" +aliasing_root(@nospecialize x) = x +function aliasing_root(x::AbstractArray) + aliases_as_whole(x) && return x + p = parent(x) + # `Base.parent` returns the argument itself for non-wrapper arrays, which + # terminates the recursion. + p === x && return x + root = aliasing_root(p) + return aliases_as_whole(root) ? root : x +end + +""" + aliasing_unwrapped(x[, dep_mod]) + +Compute the `aliasing` of `x`, first resolving (via [`aliasing_root`](@ref)) any +whole-object container that `x` wraps. This is the entry point Datadeps uses to +compute the aliasing of a (possibly wrapped) argument. + +Unwrapping and `aliasing` are intentionally fused into a single call so they +always execute in the same place. It MUST be evaluated where `x` physically lives +-- e.g. *inside* a `Chunk`'s `remotecall_fetch` block -- because the unwrapped +object's identity determines its aliasing; returning the unwrapped object across +a worker boundary first would copy it and silently change the result. +""" +aliasing_unwrapped(x) = aliasing(aliasing_root(x)) +aliasing_unwrapped(x, dep_mod) = aliasing(aliasing_root(x), dep_mod) + function aliasing(x::Chunk, T) + # Under uniform execution (MPI), `root_worker_id` is always `myid()` and is + # not a valid owner key -- defer to the acceleration so non-owning ranks + # take the owner-broadcast path instead of a local unwrap. + accel = current_acceleration() + if uniform_execution(accel) + return aliasing(accel, x, T) + end if root_worker_id(x.processor) == myid() - return aliasing(unwrap(x), T) + return aliasing_unwrapped(unwrap(x), T) end @assert x.handle isa DRef return remotecall_fetch(root_worker_id(x.processor), x, T) do x, T - aliasing(unwrap(x), T) + aliasing_unwrapped(unwrap(x), T) end end function aliasing(x::Chunk) + accel = current_acceleration() + if uniform_execution(accel) + return aliasing(accel, x, identity) + end if root_worker_id(x.processor) == myid() - return aliasing(unwrap(x)) + return aliasing_unwrapped(unwrap(x)) end @assert x.handle isa DRef return remotecall_fetch(root_worker_id(x.processor), x) do x - aliasing(unwrap(x)) + aliasing_unwrapped(unwrap(x)) end end aliasing(x::DTask, T) = aliasing(fetch(x; move_value=false, unwrap=false), T) @@ -449,6 +570,8 @@ end memory_spans(a::ContiguousAliasing{S}) where S = MemorySpan{S}[a.span] will_alias(x::ContiguousAliasing{S}, y::ContiguousAliasing{S}) where S = will_alias(x.span, y.span) +equivalent_structure(x::ContiguousAliasing{S}, y::ContiguousAliasing{S}) where S = + x.span.len == y.span.len struct IteratedAliasing{T} <: AbstractAliasing x::T end @@ -517,6 +640,14 @@ function aliasing(x::SubArray{T,N}) where {T,N} return UnknownAliasing() end end +function equivalent_structure(x::StridedAliasing{T,N,S}, y::StridedAliasing{T,N,S}) where {T,N,S} + x.base_inds == y.base_inds || return false + x.lengths == y.lengths || return false + x.strides == y.strides || return false + # Compare the offset of the view into its parent, not the absolute pointers, + # so views into different (but identically-shaped) parents match. + return (x.ptr.addr - x.base_ptr.addr) == (y.ptr.addr - y.base_ptr.addr) +end function will_alias(x::StridedAliasing{T1,N1,S1}, y::StridedAliasing{T2,N2,S2}) where {T1,T2,N1,N2,S1,S2} # Check if the base pointers are the same # FIXME: Conservatively incorrect via `unsafe_wrap` and friends @@ -566,6 +697,8 @@ function memory_spans(a::TriangularAliasing{T,S}) where {T,S} end return spans end +equivalent_structure(x::TriangularAliasing{T,S}, y::TriangularAliasing{T,S}) where {T,S} = + x.stride == y.stride && x.isupper == y.isupper && x.diagonal == y.diagonal function aliasing(x::UpperTriangular{T}) where T p = parent(x) space = memory_space(p) @@ -606,6 +739,8 @@ function aliasing(x::AbstractMatrix{T}, ::Type{Diagonal}) where T rptr = RemotePtr{Cvoid}(ptr, S) return DiagonalAliasing{T,typeof(S)}(rptr, size(parent(x), 1)) end +equivalent_structure(x::DiagonalAliasing{T,S}, y::DiagonalAliasing{T,S}) where {T,S} = + x.stride == y.stride # FIXME: Bidiagonal # FIXME: Tridiagonal diff --git a/src/precompile.jl b/src/precompile.jl index 34b14c503..639474459 100644 --- a/src/precompile.jl +++ b/src/precompile.jl @@ -9,6 +9,28 @@ t2 = spawn(+, 1, t1) fetch(t2) + # Exercise the tape planner and default cost models while enabled, so the + # first user-facing enabled allocation does not pay a cold-compile spike. + # Use `:lexical` here: `:backtrace` site IDs are session-local and the + # precompile process's stack is not a useful key to bake into the image. + # Avoid allocating real DArrays here — that would spawn tasks the cleanup + # below is not sized to drain. + let old_site_id = Tapes.CONFIG.site_id + # `calibrate=false`: do not bake build-host FLOP/bandwidth rates into + # the image; runtime `enable!` will measure on first use. + Tapes.enable!(site_id=:lexical, calibrate=false) + try + Tapes.backtrace_hash() + Tapes.explain(devnull, Float64, (128, 128)) + Tapes.plan_allocation(Float64, (64, 64); requested = AutoBlocks()) + Tapes.resolve_partitioning(Float64, (64, 64), AutoBlocks(), :arbitrary) + finally + Tapes.disable!() + Tapes.clear!() + Tapes.CONFIG.site_id = old_site_id + end + end + # Clean up refs t1 = nothing; t2 = nothing state = Sch.EAGER_STATE[] diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index 6a353e28c..15edd343c 100644 --- a/src/sch/Sch.jl +++ b/src/sch/Sch.jl @@ -1198,15 +1198,21 @@ Base.hash(task::TaskSpec, h::UInt) = hash(task.thunk_id, hash(TaskSpec, h)) sch_handle, state.uid)) end + # N.B. These must be migratable. Otherwise they pile up on whichever thread + # is submitting work: a submitter that runs a long stretch of + # `schedule_one!` without yielding (e.g. Datadeps planning a whole region) + # holds off every fire task it just created, so nothing starts executing + # until planning finishes. Letting them migrate is what allows execution to + # overlap submission. if !isempty(to_send) if root_worker_id(gproc) == myid() - @reusable_tasks :fire_tasks!_task_cache 32 _->nothing "fire_tasks!" FireTaskSpec(proc, state.chan, to_send) + @reusable_tasks :fire_tasks!_task_cache 32 Dagger.set_task_migratable! "fire_tasks!" FireTaskSpec(proc, state.chan, to_send) else # N.B. We don't batch these because we might get a deserialization # error due to something not being defined on the worker, and then we don't # know which task failed. for task_spec in to_send - @reusable_tasks :fire_tasks!_task_cache 32 _->nothing "fire_tasks!" FireTaskSpec(proc, state.chan, task_spec) + @reusable_tasks :fire_tasks!_task_cache 32 Dagger.set_task_migratable! "fire_tasks!" FireTaskSpec(proc, state.chan, task_spec) end end end @@ -1565,7 +1571,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re if tid !== nothing Dagger.set_task_tid!(t, tid) else - t.sticky = false + Dagger.set_task_migratable!(t) end end "thunk $thunk_id" DoTaskSpec(to_proc, return_queue, task, cancel_token) @@ -1579,7 +1585,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re if tid !== nothing Dagger.set_task_tid!(proc_run_task, tid) else - proc_run_task.sticky = false + Dagger.set_task_migratable!(proc_run_task) end return errormonitor_tracked("processor $to_proc", schedule(proc_run_task)) end diff --git a/src/sch/util.jl b/src/sch/util.jl index 9057f5740..9242c7bbb 100644 --- a/src/sch/util.jl +++ b/src/sch/util.jl @@ -637,6 +637,7 @@ function estimate_task_costs(state, procs, task; sig=nothing) return sorted_procs, costs end const DEFAULT_TRANSFER_RATE = UInt64(1_000_000) +const EMPTY_TRANSFER_RATES = Dict{Processor,UInt64}() @reuse_scope function estimate_task_costs!(sorted_procs, costs, state, procs, task; sig=nothing) # Find all Chunks @@ -654,28 +655,47 @@ const DEFAULT_TRANSFER_RATE = UInt64(1_000_000) end est_time_util = lock(state.signature_time_cost) do stc; get(stc, sig, 1000^3); end - # Estimate total cost for executing this task on each candidate processor + # Estimate network transfer cost per *parent* processor. Chunks are located + # per worker, so this depends only on `get_parent(proc)`, and `procs` is + # commonly every thread of a single worker -- computing it per processor + # would redo identical work (a scan of every chunk) once per thread. + # N.B. We treat same-worker transfers as having zero transfer cost + # TODO: For non-Chunk, model cost from scheduler to worker + # TODO: Measure and model processor move overhead + tx_costs = @reusable_dict :estimate_task_costs_tx_costs Processor Float64 OSProc() 0.0 8 + tx_costs_cleanup = @reuse_defer_cleanup empty!(tx_costs) for proc in procs gproc = get_parent(proc) + haskey(tx_costs, gproc) && continue chunks_filt = Iterators.filter(c->get_parent(processor(c)) != gproc, chunks) + tx_costs[gproc] = impute_sum(datasize(chunk) for chunk in chunks_filt) + end + chunks_cleanup() - # Estimate network transfer costs based on data size - # N.B. We treat same-worker transfers as having zero transfer cost - # TODO: For non-Chunk, model cost from scheduler to worker - # TODO: Measure and model processor move overhead - tx_cost = impute_sum(datasize(chunk) for chunk in chunks_filt) - - # Add fixed cost for cross-worker task transfer (esimated at 1ms) - # TODO: Actually estimate/benchmark this - task_xfer_cost = root_worker_id(gproc) != myid() ? 1_000_000 : 0 # 1ms - pid = Dagger.root_worker_id(gproc) - - tx_rate = lock(state.worker_transfer_rate) do wtr - get(get(wtr, pid, Dict{Processor,UInt64}()), proc, DEFAULT_TRANSFER_RATE) + # Estimate total cost for executing this task on each candidate processor. + # The transfer-rate table is taken once rather than once per processor. + all_equal = true + lock(state.worker_transfer_rate) do wtr + local first_cost = 0.0 + for (idx, proc) in enumerate(procs) + gproc = get_parent(proc) + pid = Dagger.root_worker_id(gproc) + + # Add fixed cost for cross-worker task transfer (esimated at 1ms) + # TODO: Actually estimate/benchmark this + task_xfer_cost = pid != myid() ? 1_000_000 : 0 # 1ms + + tx_rate = get(get(wtr, pid, EMPTY_TRANSFER_RATES), proc, DEFAULT_TRANSFER_RATE) + cost = est_time_util + (tx_costs[gproc]/tx_rate) + task_xfer_cost + costs[proc] = cost + if idx == 1 + first_cost = cost + elseif cost != first_cost + all_equal = false + end end - costs[proc] = est_time_util + (tx_cost/tx_rate) + task_xfer_cost end - chunks_cleanup() + tx_costs_cleanup() # Shuffle procs around, so equally-costly procs are equally considered np = length(procs) @@ -688,8 +708,10 @@ const DEFAULT_TRANSFER_RATE = UInt64(1_000_000) end end - # Sort by lowest cost first - sort!(sorted_procs, by=p->costs[p]) + # Sort by lowest cost first. Skipped when every processor costs the same + # (the usual single-worker case), where the shuffle above already gives the + # arbitrary-but-fair order that sorting would leave us with anyway. + all_equal || sort!(sorted_procs, by=p->costs[p]) end """ diff --git a/src/tapes/README.md b/src/tapes/README.md new file mode 100644 index 000000000..e17e2e512 --- /dev/null +++ b/src/tapes/README.md @@ -0,0 +1,140 @@ +# Dagger.Tapes — history-driven partitioning selection + +Speculative selection of `DArray` partitionings from recorded operation +history. Opt-in, disabled by default, and inert when disabled. + +## Files + +``` +tapes/Tapes.jl module entry, TapeConfig, enable!/disable! +tapes/tape.jl keys, site identification, the trie, recording, prediction +tapes/cost.jl MachineModel, ArgView, @cost_model, default models +tapes/plan.jl candidate generation, DP planner, regret gating +tapes/api.jl plan_allocation/track!, @record_op, @expect_ops, introspection +tapes/integration.jl resolve_partitioning, @tracked_alloc, patch points +tapes_tests.jl test suite (rename to test/tapes.jl) +``` + +Copy `tapes/` to `src/tapes/`, add to `src/Dagger.jl` immediately after the +`include("array/darray.jl")` line: + +```julia +include("tapes/Tapes.jl"); using .Tapes +``` + +That alone is inert. The allocation and instrumentation patch points are +written out in the `PATCH POINTS` comment block at the bottom of +`tapes/integration.jl`. + +## How it works + +1. **Identify.** Each allocation gets a `SiteKey`: calling context, element + type, rank, and a log-bucketed size. Three context strategies — + `:backtrace` (unwind and hash raw IPs), `:context` (incremental + probabilistic calling context, O(1)), `:lexical` (macro-expansion site + only, free). + +2. **Record.** `@record_op :cholesky! A` appends `(op, argument position, + arity)` to the array's tape, committed into a per-site prefix trie + immediately. Argument position is part of the identity: being the `A` of + `mul!(C, A, B)` implies a different layout preference from being the `C`. + +3. **Predict.** On a later allocation from the same site, walk the trie + greedily and emit `PredictedOp`s carrying the cumulative probability of + reaching each step. + +4. **Plan.** Cost each candidate layout against each predicted step, then + dynamic-program the chain including redistribution cost on the edges. + Linear chain, so exact in O(n·|L|²). + +5. **Gate.** Adopt the plan only if it beats the fallback by + `gate_margin`, and — below `regret_confidence_threshold` — only if no + single predicted operation is more than `max_regret_ratio` worse than its + own best layout. + +Only `steps[1]` is committed; the planner re-runs at each observed operation +against the branch actually taken. + +## Usage + +```julia +Dagger.Tapes.enable!() + +for i in 1:10 + A = rand(AutoBlocks(), Float64, 4096, 4096) + B = A * A' + 4096I + cholesky!(B) +end + +Dagger.Tapes.report() # what was learned +Dagger.Tapes.explain(Float64, (4096, 4096)) # why a layout was chosen +``` + +Ahead-of-time, no warm-up required: + +```julia +Dagger.@expect_ops [:mul!, :cholesky!, :trsm!] begin + A = rand(AutoBlocks(), Float64, n, n) + ... +end +``` + +Without an operation list, `@expect_ops` just establishes a stable lexical +site key — which makes `site_id = :lexical` as precise as `:backtrace` for the +code you care about, at zero cost. + +Declaring a cost model: + +```julia +Dagger.Tapes.@cost_model my_solve(A, B) = begin + nt = nblocks(A, 1) + serial_time(nt * blocksize(A, 1)^2) + + flops_time(size(A, 1)^2 * size(B, 2)) * imbalance(nt) + + task_time(nt^2 / 2) +end +``` + +Argument order must match the `@record_op` call for that operation. Pass only +`DArray`s. + +## Benchmarking + +Performance becomes warm-up-dependent and bimodal once tapes are active, so +`dagger_bench.py` and anything else reporting a distribution across +repetitions needs to control cold vs warm explicitly: + +- `Tapes.clear!()` between repetitions for cold measurements +- `Tapes.pin!(T, dims, Blocks(...), :cyclicrow)` to make a site deterministic + and history-independent +- `Tapes.disable!()` for a true baseline +- `Tapes.stats().hit_rate` to confirm predictions are actually landing + +## Verify in this order + +1. **`@cost_model` expansion.** The one construct I could not test — no Julia + in the build container. It names the method via + `Expr(:., TAPES_MODULE, QuoteNode(:op_cost))`. If it misbehaves, the two + alternatives are a bare unescaped `op_cost` (relies on hygiene resolving a + *definition*) or `GlobalRef(TAPES_MODULE, :op_cost)`. +2. `Tapes.enable!(); Tapes.explain(Float64, (4096,4096))` — exercises site + identification, candidate generation, the planner and all default models in + one call. +3. The test suite, which runs without a cluster except for the last two + groups. +4. `@record_op` overhead with `site_id = :backtrace` against your 100µs task + floor. + +## TODOs in the code + +Every one carries its rationale inline. The load-bearing ones: + +| where | what | +|:--|:--| +| `record_op!` | **Joint planning.** Tapes are per-array but layout decisions are joint — `mul!(C,A,B)` needs mutually compatible layouts. Currently undervalues layouts that only pay off when all operands agree. Wants union-find over co-participating arrays using the Datadeps aliasing analysis. | +| `integration.jl` | **Residency tracker.** Precondition for read-only replication, precise `redistribution_cost`, heterogeneous cost models and OOC. Highest-leverage missing piece. | +| `integration.jl` | **Datadeps lookahead.** Inside `spawn_datadeps` the full dependency structure is already known — exact lookahead, no prediction. Strictly better where it applies; build before broadening `@record_op` coverage. | +| `maybe_repartition!` | **Repartition mechanism.** Planner emits per-step layouts; nothing acts on them. Hard part is arrays captured by already-submitted-but-unscheduled thunks. | +| `current_machine` | **Calibration.** FLOP rate and bandwidth are guesses; should be measured at `enable!` and refined from `TimespanLogging` task durations. Natural autotuner seam. | +| `backtrace_hash` | **Persistence.** Raw IPs are session-local. Symbolicating `(file, line)` with a per-IP cache would make tapes survive restarts — which is what makes the *first* run fast rather than the second. | +| `ArgSpec` | **Assignment provenance.** `DArray` does not store the assignment it was built with, so recorded specs assume `:arbitrary`. | +| `OP_AFFINITY` | **Learn it.** Hand-maintained table that will rot; the tape already produces the triples needed to fit it. | diff --git a/src/tapes/Tapes.jl b/src/tapes/Tapes.jl new file mode 100644 index 000000000..8b15e9d72 --- /dev/null +++ b/src/tapes/Tapes.jl @@ -0,0 +1,244 @@ +""" + Dagger.Tapes + +Speculative, history-driven selection of `DArray` partitionings. + +# Rationale + +Choosing a partitioning at allocation time is a decision made with no knowledge +of what the array will be *used for*. Any fixed default (currently +`auto_blocks`, which splits the last dimension by processor count) is +necessarily wrong for large classes of operations: a Cholesky factorization +wants square tiles laid out 2D-cyclically, a row-wise `mapreduce` wants fat row +blocks, a triangular solve against a vector wants something else again. + +This module exploits the *principle of persistence*: in real applications, +allocations happen at a small number of program points, and the sequence of +operations applied to the resulting array is highly repeatable across +executions of that program point. So: + +1. At each allocation, identify the *site* (call-stack context + element type + + coarse size bucket). +2. Record the ordered sequence of operations subsequently applied to the array + from that site, into a per-site prefix trie (`TapeRoot`). +3. On a later allocation from the same site, walk the trie to predict the + operation sequence that is about to follow, with a probability attached to + each step. +4. Feed that prediction, plus per-operation cost models, into a dynamic program + that picks the layout minimising expected total cost across the whole + predicted chain (including the cost of any mid-chain repartitioning). + +A misprediction costs a suboptimal layout, never a wrong answer. + +# Status / safety + +The whole subsystem is **opt-in** and disabled by default; when disabled every +hook is a single predictable branch on a `const` struct field, which the +compiler hoists out of hot paths. When enabled, it never overrides an +*explicit* `Blocks(...)` supplied by the user unless +`CONFIG.override_explicit_blocks` is set — it only fills in the choice that +`AutoBlocks()` would otherwise have made blindly. + +# Quick start + +```julia +Dagger.Tapes.enable!() # :backtrace site identification + +# Run your workload once to populate tapes, then again to benefit: +for i in 1:10 + A = rand(AutoBlocks(), Float64, 4096, 4096) + A = A * A' + 4096I + cholesky!(A) +end + +Dagger.Tapes.report() # what was learned +Dagger.Tapes.explain(Float64, (4096, 4096)) # why a layout was chosen +``` + +Cheaper site identification, and ahead-of-time declaration, are available via +[`@expect_ops`](@ref): + +```julia +Dagger.@expect_ops [:mul!, :cholesky!, :trsm!] begin + A = rand(AutoBlocks(), Float64, 4096, 4096) + ... +end +``` + +# Public API + +- [`enable!`](@ref), [`disable!`](@ref), [`CONFIG`](@ref) +- [`@record_op`](@ref) / [`record_op!`](@ref) — instrument an operation +- [`@cost_model`](@ref) — declare a cost model for an operation +- [`@expect_ops`](@ref) — declare an operation chain ahead of time +- [`plan_allocation`](@ref) / [`track!`](@ref) — allocation-site hooks +- [`report`](@ref), [`explain`](@ref), [`dump_tapes`](@ref), [`clear!`](@ref), + [`pin!`](@ref), [`unpin!`](@ref) — introspection and manual override +""" +module Tapes + +using ..Dagger +import ..Dagger: DArray, Blocks, AutoBlocks, AbstractBlocks, AssignmentType +import ..Dagger: auto_blocks, num_processors + +using ScopedValues +using MacroTools: @capture + +export @record_op, @cost_model, @expect_ops + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + +""" + TapeConfig + +Tunables for the tape subsystem. The live instance is [`CONFIG`](@ref); mutate +its fields directly, or use [`enable!`](@ref) which takes the same names as +keyword arguments. +""" +Base.@kwdef mutable struct TapeConfig + "Master switch. When `false` every hook is a single branch and returns immediately." + enabled::Bool = false + + """ + How allocation sites are identified. One of: + - `:backtrace` — hash the raw instruction pointers of `backtrace()`. Most + precise, costs roughly 100-500us per allocation depending on stack depth. + Keys are only valid within a session (see TODO on symbolication). + - `:context` — probabilistic calling context (Bond & McKinley, OOPSLA'07): + an incrementally-maintained hash mixed at `@record_op` / `@expect_ops` + boundaries. O(1), but only as precise as the instrumentation density. + - `:lexical` — the macro-expansion site of the allocation call only. + Effectively free, but cannot distinguish two callers of the same wrapper. + """ + site_id::Symbol = :backtrace + + "Number of stack frames hashed in `:backtrace` mode (nearest-first, after `backtrace_skip`)." + backtrace_depth::Int = 24 + "Frames to discard from the top of the stack (Dagger's own plumbing)." + backtrace_skip::Int = 2 + + "Powers-of-two subdivisions per octave when bucketing sizes into a site key. `0` = exact sizes." + size_buckets_per_octave::Int = 2 + + "Maximum number of operations recorded per array (trie depth cap)." + max_tape_length::Int = 24 + "How far ahead the planner looks when predicting." + horizon::Int = 12 + "A trie branch is only followed during prediction if it holds at least this share of its parent's count." + min_branch_prob::Float64 = 0.15 + "Minimum cumulative probability for the first predicted op before we will speculate at all." + min_confidence::Float64 = 0.5 + "Minimum times a site must have been observed before its predictions are used." + min_observations::Int = 2 + + """ + The planned layout must beat the fallback layout by at least this factor to + be adopted (`0.9` = must be at least 10% cheaper). Guards against churn from + cost-model noise. + """ + gate_margin::Float64 = 0.9 + + """ + Reject a plan if, for any *individual* predicted operation, it is more than + this many times worse than the best layout for that operation alone. This is + the minimax-regret safety valve: early on, prefer a layout that is decent for + everything over one that is optimal for the modal chain and catastrophic + otherwise. + """ + max_regret_ratio::Float64 = 3.0 + "Below this cumulative confidence, `max_regret_ratio` is enforced strictly." + regret_confidence_threshold::Float64 = 0.85 + + "Upper bound on candidate layouts considered per allocation (planner is O(n*L^2))." + max_candidates::Int = 16 + + "Override a partitioning the user gave explicitly. Off by default: explicit user intent wins." + override_explicit_blocks::Bool = false + + """ + Permit mid-chain repartitioning when the plan calls for it. Currently a + no-op stub; see `maybe_repartition!`. + """ + allow_repartition::Bool = false + + "LRU cap on the number of distinct sites retained." + max_sites::Int = 4096 + "LRU cap on trie nodes across all sites." + max_nodes::Int = 200_000 + + "Emit `@debug`-style commentary about every decision." + verbose::Bool = false +end + +""" + CONFIG::TapeConfig + +The live configuration. See [`TapeConfig`](@ref) for field documentation. +""" +const CONFIG = TapeConfig() + +""" + is_enabled() -> Bool + +Master predicate guarding every hook. Deliberately trivial so that the branch +is cheap and predictable in the disabled case. +""" +@inline is_enabled() = CONFIG.enabled + +""" + enable!(; calibrate=true, kwargs...) + +Turn the tape subsystem on. Keyword arguments set the matching +[`TapeConfig`](@ref) fields. + +By default this also runs [`calibrate_machine!`](@ref) once, so layout ranking +uses measured FLOP/s and memory bandwidth rather than the pessimistic +placeholders. Pass `calibrate=false` to skip (e.g. after an explicit +[`set_machine!`](@ref), or in tests that need a fixed model). + +```julia +Dagger.Tapes.enable!(site_id=:lexical, horizon=8, verbose=true) +``` +""" +function enable!(; calibrate::Bool = true, kwargs...) + for (k, v) in kwargs + hasfield(TapeConfig, k) || throw(ArgumentError("unknown Tapes config option: $k")) + setfield!(CONFIG, k, convert(fieldtype(TapeConfig, k), v)) + end + CONFIG.site_id in (:backtrace, :context, :lexical) || + throw(ArgumentError("site_id must be :backtrace, :context or :lexical, got $(CONFIG.site_id)")) + CONFIG.enabled = true + # Calibrate once per session unless the user already installed rates via + # `set_machine!` (or passed `calibrate=false` to skip, e.g. in tests). + if calibrate + m = _ensure_machine() + m.calibrated || calibrate_machine!(m) + end + return CONFIG +end + +""" + disable!() + +Turn the tape subsystem off. Learned tapes are retained; use [`clear!`](@ref) +to discard them. +""" +function disable!() + CONFIG.enabled = false + return CONFIG +end + +@inline function vlog(args...) + CONFIG.verbose && @info string("[Tapes] ", args...) + nothing +end + +include("tape.jl") +include("cost.jl") +include("plan.jl") +include("api.jl") +include("integration.jl") + +end # module Tapes diff --git a/src/tapes/api.jl b/src/tapes/api.jl new file mode 100644 index 000000000..512ee29ca --- /dev/null +++ b/src/tapes/api.jl @@ -0,0 +1,533 @@ +# =========================================================================== +# Public API: allocation hooks, instrumentation macros, and introspection. +# =========================================================================== + +# --------------------------------------------------------------------------- +# Ahead-of-time declaration +# --------------------------------------------------------------------------- + +""" + Declaration + +An active [`@expect_ops`](@ref) region. `ops === nothing` means the region only +establishes a lexical scope (so allocations inside it are keyed cheaply and +consistently); a non-`nothing` list is a user assertion about what will happen, +adopted with `prob = declared_prob`. +""" +struct Declaration + token::UInt64 + ops::Union{Nothing,Vector{Symbol}} + declared_prob::Float64 +end + +const ACTIVE_DECLARATION = ScopedValue{Union{Nothing,Declaration}}(nothing) + +""" + @expect_ops begin ... end + @expect_ops [:mul!, :cholesky!, :trsm!] begin ... end + @expect_ops [:cholesky!] prob=0.9 begin ... end + +Declare an operation chain ahead of time for allocations made inside the +region. + +Two distinct uses, both feeding the same planner: + +1. **With an operation list.** You are asserting what will happen to arrays + allocated in this region. The list is adopted as a forecast with confidence + `prob` (default `1.0`) and the planner optimises against it *immediately* — + no warm-up run required, which is the main reason to reach for this. + +2. **Without a list.** The region establishes a stable lexical site key and + mixes into the calling-context hash. Allocations inside are keyed by the + region rather than by an unwound stack, which makes + `CONFIG.site_id = :lexical` (free) or `:context` (nearly free) as precise as + `:backtrace` for the code you care about. Recording and prediction then + proceed normally from observation. + +The declared operations are matched positionally against argument position 1 +with arity 1. That is a simplification — it cannot express "this array is the +`B` of a `trsm!`" — but it covers the common case where the declared array is +the primary operand. + +TODO(richer-declarations): accept `:trsm! => 2` or a small tuple syntax to +declare argument position and arity, and accept per-array declarations rather +than region-wide ones (currently *every* allocation in the region adopts the +same chain, which is wrong when a region allocates both a matrix and its +right-hand side). + +# Example + +```julia +Dagger.@expect_ops [:mul!, :cholesky!, :trsm!] begin + A = rand(AutoBlocks(), Float64, n, n) + B = A * A' + n*I + cholesky!(B) + ldiv!(B, rhs) +end +``` +""" +macro expect_ops(args...) + isempty(args) && error("@expect_ops requires a body") + body = args[end] + rest = args[1:end-1] + + ops = nothing + prob = 1.0 + for a in rest + if @capture(a, prob = p_) + prob = p + elseif a isa Expr && (a.head === :vect || a.head === :tuple) + ops = a + else + error("@expect_ops: unexpected argument `$a`; expected an op list and/or `prob=...`") + end + end + + token = lexical_token(__module__, __source__) + opsexpr = ops === nothing ? :(nothing) : :(Symbol[$(map(esc, ops.args)...)]) + + quote + if $is_enabled() + local __parent__ = $(ACTIVE_DECLARATION)[] + local __tok__ = $mix_context(__parent__ === nothing ? $(CONTEXT_HASH)[] : + __parent__.token, $token) + local __decl__ = $Declaration(__tok__, $opsexpr, Float64($(esc(prob)))) + $(ScopedValues.with)($(ACTIVE_DECLARATION) => __decl__, + $(CONTEXT_HASH) => __tok__) do + $(esc(body)) + end + else + $(esc(body)) + end + end +end + +"Build a synthetic forecast from a user declaration." +function declared_prediction(d::Declaration, self::ArgSpec) + d.ops === nothing && return PredictedOp[] + p = clamp(d.declared_prob, 0.0, 1.0) + return [PredictedOp(OpKey(op, 1, 1), p, ArgSpec[self]) for op in d.ops] +end + +""" + site_for(token::UInt64) -> UInt64 + +The site identifier for the current context, accounting for an enclosing +[`@expect_ops`](@ref) region. + +**Every entry point that keys into the tape store must go through this.** +`plan_allocation`, `pin!` and `explain` each independently derive a site key, +and if they disagree then `pin!` silently pins a site nobody allocates from and +`explain` reports on a site that does not exist — both failures being invisible +rather than loud. +""" +@inline function site_for(token::UInt64) + decl = ACTIVE_DECLARATION[] + return decl === nothing ? current_site(token) : mix_context(decl.token, token) +end + +"Build the full store key for an allocation of `T` and shape `dims` at `token`." +@inline site_key(::Type{T}, dims::Tuple, token::UInt64) where {T} = + SiteKey(site_for(token), T, Int32(length(dims)), size_bucket(dims)) + +# --------------------------------------------------------------------------- +# Operation instrumentation +# --------------------------------------------------------------------------- + +""" + @record_op :opname A B C + +Record that operation `:opname` is about to be applied to the given arrays. +Place this immediately before the operation's task submission. + +Expands to a single `is_enabled()` branch plus a call, so it costs essentially +nothing when the subsystem is off. The macro-expansion site is baked in as a +literal and used both as the lexical site token and as the mixing constant for +the calling-context hash. + +**Pass only `DArray` arguments**, in a stable order, and keep that order in +sync with the matching [`@cost_model`](@ref): argument position is part of an +operation's recorded identity, because layout preference depends on role. + +```julia +function LinearAlgebra.mul!(C::DMatrix, A::DMatrix, B::DMatrix) + Dagger.@record_op :mul! C A B + ... +end +``` +""" +macro record_op(op, args...) + token = lexical_token(__module__, __source__) + quote + if $is_enabled() + $_record_op!($(esc(op)), $token, ($(map(esc, args)...),)) + end + nothing + end +end + +# --------------------------------------------------------------------------- +# Allocation hooks +# --------------------------------------------------------------------------- + +""" + AllocationPlan + +What [`plan_allocation`](@ref) decided. `partitioning` and `assignment` are +what the caller should actually allocate with; the rest is bookkeeping handed +back to [`track!`](@ref). +""" +struct AllocationPlan + partitioning::Union{Blocks,AutoBlocks} + assignment::Any + root::Union{Nothing,TapeRoot} + self::Union{Nothing,ArgSpec} + layout::Union{Nothing,LayoutChoice} + predicted::Vector{PredictedOp} + plan::Union{Nothing,LayoutPlan} +end + +"A plan that changes nothing — returned whenever the subsystem declines to act." +passthrough(part, assignment) = + AllocationPlan(part, assignment, nothing, nothing, nothing, PredictedOp[], nothing) + +""" + plan_allocation(::Type{T}, dims; requested=AutoBlocks(), assignment=:arbitrary, + token=UInt64(0)) -> AllocationPlan + +Decide the partitioning for an array about to be allocated. + +Declines to act — returning `requested` unchanged — when: + +- the subsystem is disabled; +- `requested` is a concrete `Blocks` and `CONFIG.override_explicit_blocks` is + false (explicit user intent wins; we only fill in what `AutoBlocks` would + otherwise have guessed); +- the site has been seen fewer than `CONFIG.min_observations` times; +- the forecast's confidence is below `CONFIG.min_confidence`; +- the planner's gate rejects the plan (insufficient margin over the fallback, + or too much regret at low confidence). + +Always registers the site, so a declined allocation still *learns*. + +Pair with [`track!`](@ref): + +```julia +p = Tapes.plan_allocation(T, dims; requested=dist, assignment=assignment) +A = _allocate(T, dims, p.partitioning, p.assignment) +return Tapes.track!(A, p) +``` +""" +function plan_allocation(::Type{T}, dims::Tuple; + requested = AutoBlocks(), + assignment = :arbitrary, + token::UInt64 = UInt64(0)) where {T} + is_enabled() || return passthrough(requested, assignment) + (requested isa Blocks && !CONFIG.override_explicit_blocks) && + return passthrough(requested, assignment) + isempty(dims) && return passthrough(requested, assignment) + + decl = ACTIVE_DECLARATION[] + key = site_key(T, dims, token) + root = get_root!(key) + + fb = fallback_layout(T, dims) + self = ArgSpec(T, dims, blocksize(fb), :arbitrary) + + # Manual override short-circuits everything. + pinned = root.pinned + if pinned !== nothing + return AllocationPlan(to_blocks(pinned), pinned.assignment, + root, ArgSpec(T, dims, blocksize(pinned), pinned.assignment), + pinned, PredictedOp[], nothing) + end + + # Forecast: a user declaration if present, otherwise the learned trie. + pred = PredictedOp[] + if decl !== nothing && decl.ops !== nothing + pred = declared_prediction(decl, self) + elseif root.nobservations >= CONFIG.min_observations + pred = predict(root.root) + end + + if isempty(pred) || confidence(pred) < CONFIG.min_confidence + vlog("no usable forecast for $(key.eltype)$(dims); using fallback $fb") + return AllocationPlan(requested, assignment, root, self, fb, pred, nothing) + end + + m = current_machine() + cands = candidate_layouts(T, dims, m) + lp = plan_chain(pred, cands, self, m) + + if !lp.accepted + vlog("plan rejected ($(lp.reason)); using fallback $fb") + return AllocationPlan(requested, assignment, root, self, fb, pred, lp) + end + + chosen = lp.steps[1] + vlog("chose $chosen for $(key.eltype)$(dims) ", + "(cost $(round(lp.cost; sigdigits=3))s vs fallback ", + "$(round(lp.fallback_cost; sigdigits=3))s, regret $(round(lp.max_regret; digits=2)))") + return AllocationPlan(to_blocks(chosen), chosen.assignment, root, + ArgSpec(T, dims, blocksize(chosen), chosen.assignment), + chosen, pred, lp) +end + +""" + track!(A::DArray, p::AllocationPlan) -> A + +Attach the recording state produced by [`plan_allocation`](@ref) to the +freshly-allocated array, and return the array so this can be used in tail +position. + +The trace is held in a `WeakKeyDict`, so it never extends `A`'s lifetime and +disappears with it. +""" +function track!(A::DArray, p::AllocationPlan) + root = p.root + (root === nothing || !is_enabled()) && return A + self = p.self + layout = p.layout + self === nothing && return A + layout === nothing && return A + + steps = p.plan === nothing ? LayoutChoice[layout] : p.plan.steps + trace = LiveTrace(root, self, layout, p.predicted, steps) + TRACES[A] = trace + lock(STORE_LOCK) do + root.nobservations += 1 + root.root.count += 1 + end + return A +end + +""" + suggest_partitioning(::Type{T}, dims; kwargs...) -> (partitioning, assignment) + +One-shot convenience wrapper for call sites that cannot conveniently thread an +[`AllocationPlan`](@ref) through to a [`track!`](@ref) call. The array is not +tracked, so it contributes nothing to future predictions — prefer +`plan_allocation` + `track!` wherever the allocation site can be edited +properly. +""" +function suggest_partitioning(::Type{T}, dims::Tuple; kwargs...) where {T} + p = plan_allocation(T, dims; kwargs...) + return (p.partitioning, p.assignment) +end + +# --------------------------------------------------------------------------- +# Manual override +# --------------------------------------------------------------------------- + +""" + pin!(::Type{T}, dims, blocks::Blocks, assignment=:arbitrary; token=UInt64(0)) + +Force a layout for a site, bypassing prediction entirely. Intended for +debugging and for benchmark reproducibility: with a pin in place, the +subsystem's behaviour at that site is deterministic and independent of warm-up +history. + +Note that pinning still requires the *site* to be identified, so in +`:backtrace` mode you must call `pin!` from the same context the allocation +happens in. In `:lexical` mode, wrap the allocation in [`@expect_ops`](@ref) +and pin against that region's token. +""" +function pin!(::Type{T}, dims::Tuple, blocks::Blocks, assignment::Symbol = :arbitrary; + token::UInt64 = UInt64(0)) where {T} + root = get_root!(site_key(T, dims, token)) + lock(STORE_LOCK) do + root.pinned = LayoutChoice(blocks.blocksize, assignment, :pinned) + end + return root +end + +"Remove all pins." +function unpin!() + lock(STORE_LOCK) do + for (_, r) in STORE + r.pinned = nothing + end + end + return nothing +end + +""" + clear!() + +Discard all learned tapes. Does not change [`CONFIG`](@ref). + +Essential for benchmarking: with tapes active, performance becomes +warm-up-dependent and bimodal, so any harness that reports a distribution +across repetitions needs to control explicitly whether each repetition starts +cold or warm. +""" +function clear!() + lock(STORE_LOCK) do + empty!(STORE) + TOTAL_NODES[] = 0 + end + empty!(TRACES) + reset_machine!() + return nothing +end + +# --------------------------------------------------------------------------- +# Introspection +# --------------------------------------------------------------------------- + +""" + stats() -> NamedTuple + +Aggregate counters: sites tracked, trie nodes, total observations, and +forecast hit/miss counts. +""" +function stats() + lock(STORE_LOCK) do + nobs = 0; hits = 0; misses = 0; nodes = 0 + for (_, r) in STORE + nobs += r.nobservations; hits += r.hits; misses += r.misses; nodes += r.nnodes + end + return (sites = length(STORE), nodes = nodes, observations = nobs, + hits = hits, misses = misses, + hit_rate = (hits + misses) > 0 ? hits / (hits + misses) : NaN, + live_traces = length(TRACES)) + end +end + +""" + report([io]; limit=20) + +Human-readable overview of what has been learned: one line per site with its +observation count and modal predicted chain. + +(Named `report` rather than `summary` to avoid shadowing `Base.summary` inside +this module.) +""" +function report(io::IO = stdout; limit::Int = 20) + s = stats() + println(io, "Dagger.Tapes: ", CONFIG.enabled ? "enabled" : "disabled", + " (site_id=:", CONFIG.site_id, ")") + println(io, " sites=", s.sites, " nodes=", s.nodes, + " observations=", s.observations, " live=", s.live_traces) + roots = lock(STORE_LOCK) do + sort!(collect(values(STORE)); by = r -> -r.nobservations) + end + for (i, r) in enumerate(roots) + i > limit && (println(io, " ... ", length(roots) - limit, " more"); break) + pred = predict(r.root) + chain = isempty(pred) ? "(no prediction)" : + join([string(p.key.op, "@", round(p.prob; digits = 2)) for p in pred], " -> ") + println(io, " [", r.nobservations, "x] ", r.key.eltype, "^", r.key.ndims, + " site=", string(r.key.site; base = 16)[1:min(end, 8)], + r.pinned === nothing ? "" : " PINNED=$(r.pinned)") + println(io, " ", chain) + end + return nothing +end + +""" + dump_tapes([io]) + +Print the full trie for every site, with per-branch counts. Verbose; intended +for debugging prediction quality when `report` shows something surprising. +""" +function dump_tapes(io::IO = stdout) + roots = lock(STORE_LOCK) do + sort!(collect(values(STORE)); by = r -> -r.nobservations) + end + for r in roots + println(io, "site ", string(r.key.site; base = 16), " ", r.key.eltype, + " ndims=", r.key.ndims, " obs=", r.nobservations) + _dump_node(io, r.root, 1) + end + return nothing +end + +function _dump_node(io::IO, n::TapeNode, indent::Int) + kids = n.children + kids === nothing && return nothing + for (k, c) in sort!(collect(kids); by = kv -> -kv[2].count) + stops = max(0, c.count - child_total(c)) + println(io, " "^(2 * indent), k, " x", c.count, + stops > 0 ? " (chain ended here x$stops)" : "") + _dump_node(io, c, indent + 1) + end + return nothing +end + +""" + explain([io], ::Type{T}, dims; requested=AutoBlocks(), assignment=:arbitrary) + +Show, for a hypothetical allocation of `T` and shape `dims` at the *current* +call site, what the subsystem predicts, which candidates it considered, what +each costs, and why it accepted or rejected the resulting plan. + +This is the first thing to reach for when a layout choice looks wrong. Call it +from the same context as the real allocation so the site key matches. + +```julia +Dagger.Tapes.explain(Float64, (4096, 4096)) +``` +""" +function explain(io::IO, ::Type{T}, dims::Tuple; + requested = AutoBlocks(), assignment = :arbitrary) where {T} + if !is_enabled() + println(io, "Tapes is disabled; every allocation uses the caller's request.") + return nothing + end + decl = ACTIVE_DECLARATION[] + key = site_key(T, dims, UInt64(0)) + site = key.site + root = get_root(key) + + println(io, "site = 0x", string(site; base = 16)) + println(io, "key = ", T, " ndims=", length(dims), " bucket=", + unpad_dims(size_bucket(dims), length(dims))) + if root === nothing + println(io, "status = unseen; would use fallback ", fallback_layout(T, dims)) + return nothing + end + println(io, "obs = ", root.nobservations, + " (min required ", CONFIG.min_observations, ")") + + pred = decl !== nothing && decl.ops !== nothing ? + declared_prediction(decl, ArgSpec(T, dims, blocksize(fallback_layout(T, dims)), :arbitrary)) : + predict(root.root) + if isempty(pred) + println(io, "forecast = none; would use fallback ", fallback_layout(T, dims)) + return nothing + end + println(io, "forecast =") + for (i, p) in enumerate(pred) + println(io, " ", i, ". ", p.key, " p=", round(p.prob; digits = 3), + has_cost_model(p.key.op) ? " [modeled]" : " [generic]") + end + + m = current_machine() + self = ArgSpec(T, dims, blocksize(fallback_layout(T, dims)), :arbitrary) + cands = candidate_layouts(T, dims, m) + println(io, "candidates (per-step expected cost, seconds):") + for c in cands + costs = [step_cost(p, self, c, m) for p in pred] + println(io, " ", rpad(string(c), 44), " total=", + round(sum(costs); sigdigits = 4), " ", + join([string(round(x; sigdigits = 3)) for x in costs], " ")) + end + + lp = plan_chain(pred, cands, self, m) + println(io, "plan = ", isempty(lp.steps) ? "(none)" : string(lp.steps[1]), + " (committed; receding horizon)") + length(lp.steps) > 1 && println(io, " full: ", join(string.(lp.steps), " | ")) + println(io, "cost = ", round(lp.cost; sigdigits = 4), + " fallback = ", round(lp.fallback_cost; sigdigits = 4), + " ratio = ", round(lp.cost / lp.fallback_cost; digits = 3), + " (must be <= ", CONFIG.gate_margin, ")") + println(io, "regret = ", round(lp.max_regret; digits = 3), + " (limit ", CONFIG.max_regret_ratio, " below confidence ", + CONFIG.regret_confidence_threshold, ")") + println(io, "verdict= ", lp.accepted ? "ACCEPTED" : "REJECTED ($(lp.reason))") + return nothing +end + +explain(::Type{T}, dims::Tuple; kwargs...) where {T} = explain(stdout, T, dims; kwargs...) +explain(::Type{T}, dims::Integer...; kwargs...) where {T} = explain(stdout, T, dims; kwargs...) diff --git a/src/tapes/cost.jl b/src/tapes/cost.jl new file mode 100644 index 000000000..d526b6fca --- /dev/null +++ b/src/tapes/cost.jl @@ -0,0 +1,798 @@ +# =========================================================================== +# Cost models. +# +# The tape tells us *which* operations are coming. To turn that into a layout +# decision we need `cost(op, layout, sizes) -> seconds` for combinations we may +# never have executed. A pure measurement table cannot answer counterfactuals +# ("what would Cholesky cost with 512-square tiles, given we have only ever run +# it with row blocks?") and the (op x layout x size) space is far too large to +# explore online. So the primary model is analytic and parametric; measurements +# are intended to *correct* it, not replace it. +# =========================================================================== + +import LinearAlgebra + +"The module these macros expand references into. Captured so `@cost_model` can be used from anywhere." +const TAPES_MODULE = @__MODULE__ + +""" + MachineModel + +Coarse performance parameters of the current machine. Deliberately crude: the +planner only ever compares layouts against each other, so systematic error in +these constants largely cancels. What matters is that the *ratios* between +compute, communication and per-task overhead are roughly right. +""" +Base.@kwdef mutable struct MachineModel + "Number of compute processors available." + nprocs::Int = 1 + "Aggregate achievable FLOP/s across all processors, for Float64." + flops_per_sec::Float64 = 5.0e10 + "Effective inter-processor bandwidth in bytes/s." + bandwidth::Float64 = 5.0e9 + "Per-transfer latency in seconds." + latency::Float64 = 1.0e-5 + """ + Per-task scheduler cost in seconds, **as serialized wall-clock time**, not as + work to be divided among processors. + + Most of what a Dagger task costs before it runs — Datadeps aliasing + analysis, dependency-edge construction, queue insertion — happens on the + submitting thread, so adding processors barely reduces it. Measured on a + 6-thread node in a `spawn_datadeps` region: ~230us/task for the + three-argument, tile-sharing pattern that blocked linear algebra submits, + and ~150us/task for trivial tasks holding one private chunk each. Dropping + to a single thread only raises the latter to ~360us, so six times the + processors buys a 2.4x reduction: this is far closer to a serial cost than a + parallel one, and cost models must not divide it by `nprocs`. Doing so was + why the planner preferred tile counts an order of magnitude too high. + + The default suits a handful of threads on one node; use + [`calibrate_task_overhead!`](@ref) to measure it. + """ + task_overhead::Float64 = 2.5e-4 + "Usable memory per processor in bytes; used to reject layouts that will not fit." + mem_per_proc::Float64 = 8.0e9 + """ + Whether [`calibrate_machine!`](@ref) (or the user, via [`set_machine!`](@ref)) + has supplied non-default `flops_per_sec` / `bandwidth`. Untuned defaults + skew candidate generation and layout ranking badly on modern hardware. + """ + calibrated::Bool = false +end + +const MACHINE = Ref{Union{Nothing,MachineModel}}(nothing) + +""" + _ensure_machine() -> MachineModel + +Return the cached machine model, constructing an uncalibrated default from +`Dagger.num_processors()` if needed. Does not calibrate (see +[`current_machine`](@ref)). +""" +function _ensure_machine() + m = MACHINE[] + m === nothing || return m + np = try + max(1, num_processors()) + catch + 1 + end + m = MachineModel(; nprocs = np) + MACHINE[] = m + return m +end + +""" + current_machine() -> MachineModel + +The machine model in effect. Constructed lazily from `Dagger.num_processors()` +on first use and cached. When the tape subsystem is enabled and the model has +not yet been calibrated, the first call runs [`calibrate_machine!`](@ref). + +TODO(online-refinement): refine these rates from the scheduler's existing task +timings — Dagger already records per-task durations through `TimespanLogging`, +which is exactly the signal needed. This is also the natural place to hook the +autotuner: it can supply measured `(op, layout, size) -> seconds` points that +override the analytic model where they exist and leave it in place where they +do not. + +TODO(heterogeneity): a single scalar FLOP rate is wrong on a CPU+GPU node, +which is the configuration Dagger most cares about. The model should carry a +per-processor-kind rate and the planner should weight by the fraction of tiles +each kind is expected to receive under the candidate assignment. Note that +`build_procgrid` currently filters to `ThreadProc` for symbolic assignments, so +today's assignment vocabulary cannot express a heterogeneous mapping anyway; +these two limitations should be lifted together. +""" +function current_machine() + m = _ensure_machine() + if is_enabled() && !m.calibrated + calibrate_machine!(m) + end + return m +end + +""" + calibrate_machine!(m = _ensure_machine(); force=false) -> MachineModel + +Measure aggregate GEMM throughput and streaming memory bandwidth with small +local kernels, and write the results into `m`. Takes a few tens of milliseconds +and is amortised over the session. + +Called automatically on the first [`current_machine`](@ref) after +[`enable!`](@ref), and from `enable!` itself. Pass `force=true` to overwrite a +previous calibration (including one the user installed via [`set_machine!`](@ref)). + +On a multi-threaded BLAS this measures *aggregate* FLOP/s (do not multiply by +`nprocs`). Bandwidth is a single-node DRAM-copy proxy for the "inter-processor" +term the cost models use; on multi-socket / distributed setups it is an +optimistic upper bound. +""" +function calibrate_machine!(m::MachineModel = _ensure_machine(); force::Bool = false) + (!force && m.calibrated) && return m + + try + m.nprocs = max(1, num_processors()) + catch + end + + # --- Aggregate GEMM throughput ---------------------------------------- + # BLAS may already be multi-threaded, so the result is aggregate across + # cores, not per-core. Keep this small: `enable!` should stay snappy. + n = 768 + A = rand(Float64, n, n) + B = rand(Float64, n, n) + C = Matrix{Float64}(undef, n, n) + LinearAlgebra.mul!(C, A, B) # warmup / JIT + best = Inf + for _ in 1:3 + best = min(best, @elapsed LinearAlgebra.mul!(C, A, B)) + end + if isfinite(best) && best > 0 + m.flops_per_sec = max(1.0, 2.0 * Float64(n)^3 / best) + end + + # --- Streaming memory bandwidth --------------------------------------- + # Buffer sized to exceed typical LLC so we measure DRAM, not L3. Proxy for + # shared-memory tile movement; optimistic for multi-node. + # + # The destination is freshly allocated on every repetition rather than + # reused. That is deliberate: `bandwidth` is only ever used to price moving + # data *between Dagger chunks* (redistribution, re-blocking, halo + # exchange), and every such move allocates its destination chunks, so it + # always pays first-touch page faults. Reusing a warm destination measures + # something real but irrelevant — on a 6-thread machine it reads 16.7 GB/s + # against 7.8 GB/s cold, and a measured DArray re-block achieves 4.2 GB/s. + # Overstating bandwidth makes every layout change look nearly free, which + # is exactly the error that lets a bad plan clear `CONFIG.gate_margin`. + bytes = 64 * 1024 * 1024 + nelem = bytes ÷ sizeof(Float64) + x = rand(Float64, nelem) + copyto!(Vector{Float64}(undef, nelem), x) # warmup / JIT + bestb = Inf + for _ in 1:3 + bestb = min(bestb, @elapsed copyto!(Vector{Float64}(undef, nelem), x)) + end + if isfinite(bestb) && bestb > 0 + # copy reads x and writes a fresh y. + m.bandwidth = max(1.0, 2.0 * Float64(bytes) / bestb) + end + + # `task_overhead` is left at its default. Probing it needs live Dagger + # spawns, which is unsafe from inside `enable!`/`current_machine`: it can + # deadlock against an in-flight scheduler or BLAS threadpool. Call + # `calibrate_task_overhead!` explicitly from a quiescent point instead, or + # refine it from TimespanLogging task durations (see TODO on + # `current_machine`). + + m.calibrated = true + vlog("calibrated machine: nprocs=", m.nprocs, + " flops_per_sec=", m.flops_per_sec, + " bandwidth=", m.bandwidth) + return m +end + +"An empty task body, used only by `calibrate_task_overhead!`." +_overhead_probe!(c, a, b) = nothing + +""" + calibrate_task_overhead!(m = _ensure_machine(); nt=12) -> MachineModel + +Measure `m.task_overhead` by submitting a GEMM-shaped dependency pattern of +empty tasks over tiny tiles, so the elapsed time is essentially all scheduler +cost. Uses the three-argument, tile-sharing shape that blocked linear algebra +submits, because Datadeps' cost is dominated by aliasing analysis over shared +tiles and a pattern of independent tasks measures ~1.5x too low. + +**Not** called from [`enable!`](@ref) or [`calibrate_machine!`](@ref): it spawns +real Dagger tasks, which can deadlock if a scheduler or BLAS threadpool is +already in flight. Call it once from a quiescent point: + +```julia +Dagger.Tapes.enable!() +Dagger.Tapes.calibrate_task_overhead!() +``` +""" +function calibrate_task_overhead!(m::MachineModel = _ensure_machine(); nt::Int = 12) + tiles = [Dagger.tochunk(zeros(4, 4)) for _ in 1:nt, _ in 1:nt] + ntasks = 0 + submit = function () + n = 0 + Dagger.spawn_datadeps() do + for k in 1:nt, j in (k+1):nt, i in (k+1):nt + Dagger.@spawn _overhead_probe!(Dagger.InOut(tiles[i, j]), + Dagger.In(tiles[i, k]), + Dagger.In(tiles[k, j])) + n += 1 + end + end + return n + end + try + ntasks = submit() # warmup / JIT + ntasks <= 0 && return m + best = Inf + for _ in 1:3 + best = min(best, @elapsed submit()) + end + isfinite(best) && best > 0 && + (m.task_overhead = max(1.0e-6, best / ntasks)) + catch err + vlog("task-overhead calibration failed ($err); keeping ", m.task_overhead) + return m + end + vlog("calibrated task_overhead=", m.task_overhead, " over ", ntasks, " tasks") + return m +end + +""" + set_machine!(m::MachineModel) + set_machine!(; kwargs...) + +Install or update the machine model; keyword form updates fields in place. +Supplying `flops_per_sec` or `bandwidth` marks the model as calibrated so +[`enable!`](@ref) will not overwrite it. +""" +set_machine!(m::MachineModel) = (MACHINE[] = m) +function set_machine!(; kwargs...) + m = _ensure_machine() + for (k, v) in kwargs + hasfield(MachineModel, k) || throw(ArgumentError("unknown MachineModel field: $k")) + setfield!(m, k, convert(fieldtype(MachineModel, k), v)) + end + if any(k -> k === :flops_per_sec || k === :bandwidth || k === :calibrated, keys(kwargs)) + # User-supplied rates (or an explicit flag) count as calibrated. + haskey(kwargs, :calibrated) || (m.calibrated = true) + end + return m +end + +"Reset the cached machine model so it is rebuilt (and recalibrated) on next use." +reset_machine!() = (MACHINE[] = nothing) + +# --------------------------------------------------------------------------- +# ArgView: what a cost model actually sees +# --------------------------------------------------------------------------- + +""" + ArgView + +An operation argument as presented to a cost model: shape and element type from +the recorded [`ArgSpec`](@ref), paired with the [`LayoutChoice`](@ref) being +evaluated for it. +""" +struct ArgView + spec::ArgSpec + layout::LayoutChoice +end + +Base.eltype(v::ArgView) = v.spec.eltype +Base.ndims(v::ArgView) = ndims(v.spec) +Base.size(v::ArgView) = size(v.spec) +Base.size(v::ArgView, d::Integer) = size(v.spec, d) +Base.length(v::ArgView) = prod(size(v); init = 1) + +_elsize(::Type{Nothing}) = 0 +_elsize(::Type{T}) where {T} = isbitstype(T) ? sizeof(T) : 8 + +"Bytes per element." +elsize(v::ArgView) = _elsize(v.spec.eltype) + +"Total bytes in the (logical) array." +nbytes(v::ArgView) = Float64(length(v)) * elsize(v) + +"Block size along dimension `d`, or the full tuple." +blocksize(v::ArgView) = blocksize(v.layout) +blocksize(v::ArgView, d::Integer) = blocksize(v.layout, d) + +"Number of blocks along dimension `d`, or the full tuple." +nblocks(v::ArgView) = nblocks(v.layout, size(v)) +nblocks(v::ArgView, d::Integer) = + d <= ndims(v) ? max(1, cld(size(v, d), max(1, blocksize(v, d)))) : 1 + +"Total number of blocks." +ntiles(v::ArgView) = prod(nblocks(v); init = 1) + +"Bytes in one block." +tilebytes(v::ArgView) = Float64(prod(blocksize(v); init = 1)) * elsize(v) + +"Processor assignment strategy of the candidate layout." +assignment(v::ArgView) = v.layout.assignment + +""" + square_tile(v::ArgView) -> (edge, reblocked) + +The block edge a blocked dense kernel will *actually* run with, and whether +getting there costs a copy. + +Dagger's factorizations and triangular solves do not honour a non-square +blocking: `lu!`, `cholesky!`, `qr!` and the `ldiv!` family all funnel through +`maybe_copy_buffered(A => Blocks(b, b))` with `b = min(blocksize...)`, which +copies the whole array into square tiles, runs, and copies back. A cost model +that reads `blocksize(A, 1)` directly is therefore costing a layout the runtime +never executes — for `Blocks(4096, 683)` it charges a 4096-tall panel +factorization when what runs is a 683-square one. + +Getting this wrong does not merely misprice one candidate: `AutoBlocks` is +exactly the non-square case, so it inflates the *baseline* every plan is gated +against, which turns `CONFIG.gate_margin` from a safety check into a rubber +stamp. +""" +function square_tile(v::ArgView) + n = ndims(v) + n == 0 && return (1, false) + bs = blocksize(v) + edges = ntuple(d -> max(1, min(bs[d], size(v, d))), n) + edge = minimum(edges) + return (edge, any(!=(edge), edges)) +end + +""" + reblock_cost(v::ArgView, m::MachineModel) -> Float64 + +Cost of the copy-in/copy-out `maybe_copy_buffered` performs when a layout has to +be squared off before a dense kernel will accept it: the array's bytes move +twice, over one task per tile in each of the two blockings. +""" +function reblock_cost(v::ArgView, m::MachineModel) + edge, reblocked = square_tile(v) + reblocked || return 0.0 + nt_sq = prod(ntuple(d -> max(1, cld(size(v, d), edge)), ndims(v)); init = 1) + return 2 * (m.latency + nbytes(v) / m.bandwidth) + + (ntiles(v) + nt_sq) * m.task_overhead +end + +"Longest/shortest block edge ratio; `1.0` for a perfect cube." +function aspect(v::ArgView) + n = ndims(v) + n <= 1 && return 1.0 + bs = blocksize(v) + lo, hi = typemax(Int), 0 + for d in 1:n + b = min(bs[d], size(v, d)) + b <= 0 && continue + lo = min(lo, b); hi = max(hi, b) + end + (lo == typemax(Int) || lo <= 0) && return 1.0 + return hi / lo +end + +# --------------------------------------------------------------------------- +# The cost model registry +# --------------------------------------------------------------------------- + +"Operations with a hand-written [`@cost_model`](@ref), as opposed to the generic fallback." +const MODELED_OPS = Set{Symbol}() + +""" + op_cost(::Val{op}, args::Vector{ArgView}, m::MachineModel) -> Float64 + +Estimated wall-clock seconds for one execution of `op` with the given argument +layouts. Define methods with [`@cost_model`](@ref). + +The fallback delegates to [`generic_op_cost`](@ref), so unmodelled operations +still exert sensible (if blunt) pull on the planner rather than being ignored. +""" +op_cost(::Val{Op}, args::Vector{ArgView}, m::MachineModel) where {Op} = + generic_op_cost(args, m, Op) + +"Whether `op` has a hand-written cost model." +has_cost_model(op::Symbol) = op in MODELED_OPS + +"Guard against NaN / negative / infinite model output poisoning the planner." +@inline function clamp_cost(x::Float64) + (isnan(x) || x < 0) && return Inf + return x +end + +""" + @cost_model f(a, b, ...) = expr + @cost_model function f(a, b, ...) ... end + +Declare a cost model for the operation named `f`, returning estimated seconds. +Usable inside Dagger or from user code as `Dagger.Tapes.@cost_model`. + +The declared argument names are bound to [`ArgView`](@ref)s, **in the order the +operation was recorded with** — that is, the order of arguments passed to +[`@record_op`](@ref), not the order in the underlying function signature. By +convention only `DArray`s are recorded, so scalar arguments such as `alpha` or +`uplo` do not appear. + +An `ArgView` supports `size`, `ndims`, `eltype`, `length`, plus these, all of +which are bound as locals inside the body so no imports are needed: +`blocksize`, `nblocks`, `ntiles`, `tilebytes`, `nbytes`, `elsize`, `aspect`, +`assignment`. + +These helpers are also in scope: + +| helper | meaning | +|:-----------------|:------------------------------------------------------------| +| `flops_time(f)` | seconds for `f` FLOPs spread perfectly over all processors | +| `serial_time(f)` | seconds for `f` FLOPs on one processor (i.e. critical path) | +| `bytes_time(b)` | seconds to move `b` bytes, including one latency | +| `task_time(n)` | scheduler cost for `n` tasks; serial, see `task_overhead` | +| `reblock_time(v)` | cost of squaring off `v`'s blocking, `0.0` if already square | +| `imbalance(n)` | load-imbalance multiplier for `n` tiles over `nprocs` | +| `machine` | the [`MachineModel`](@ref) | +| `nprocs_avail()` | processor count | + +`square_tile(v)` is also in scope, returning the `(edge, reblocked)` pair a +blocked dense kernel will really run with; prefer it over `blocksize(v, 1)` in +any model for an operation that squares its input. + +If the operation is recorded with fewer arguments than the model declares, the +generic fallback is used rather than erroring. + +# Example + +```julia +Dagger.Tapes.@cost_model my_solve(A, B) = begin + nt = nblocks(A, 1) + serial_time(nt * blocksize(A, 1)^2) + + flops_time(size(A, 1)^2 * size(B, 2)) * imbalance(nt) + + task_time(nt^2 / 2) +end +``` +""" +macro cost_model(ex) + fname = nothing; fargs = nothing; body = nothing + if @capture(ex, function fnm_(fa__) fb__ end) + fname = fnm; fargs = fa; body = Expr(:block, fb...) + elseif @capture(ex, fnm_(fa__) = fb_) + fname = fnm; fargs = fa; body = fb + else + error("@cost_model expects `f(args...) = expr` or `function f(args...) ... end`") + end + fname isa Symbol || error("@cost_model: operation name must be a plain symbol, got `$fname`") + all(a -> a isa Symbol, fargs) || + error("@cost_model: arguments must be plain names (no types, defaults or slurps)") + + nargs = length(fargs) + A = esc(gensym(:args)) + M = esc(gensym(:machine)) + opq = QuoteNode(fname) + + # Bind the argument names, and re-export the ArgView accessor vocabulary as + # locals so the body works regardless of what the caller has imported. + accessors = (:blocksize, :nblocks, :ntiles, :tilebytes, :nbytes, :elsize, + :aspect, :assignment, :square_tile) + accbinds = [:(local $(esc(s)) = $(getfield(TAPES_MODULE, s))) for s in accessors] + argbinds = [:(local $(esc(a)) = $A[$i]) for (i, a) in enumerate(fargs)] + + # Name the method as `.op_cost` with the module interpolated as a + # value. This is the form that works regardless of where the macro is + # expanded from: an unescaped bare `op_cost` relies on hygiene resolving to + # the macro's defining module (fragile for definitions), and an escaped + # `Dagger.Tapes.op_cost` needs `Dagger` to be bound in the caller's scope, + # which it is not while `Tapes` is itself still being defined. + fnexpr = Expr(:., TAPES_MODULE, QuoteNode(:op_cost)) + + quote + function $fnexpr(::$(Val{fname}), + $A::Vector{$ArgView}, + $M::$MachineModel) + if length($A) < $nargs + return $generic_op_cost($A, $M, $opq) + end + local $(esc(:machine)) = $M + local $(esc(:nprocs_avail)) = () -> $M.nprocs + local $(esc(:flops_time)) = f -> f / $M.flops_per_sec + local $(esc(:serial_time)) = f -> f / ($M.flops_per_sec / max(1, $M.nprocs)) + local $(esc(:bytes_time)) = b -> ($M.latency + b / $M.bandwidth) + local $(esc(:task_time)) = n -> n * $M.task_overhead + local $(esc(:reblock_time)) = v -> $reblock_cost(v, $M) + local $(esc(:imbalance)) = n -> (n <= 0 ? 1.0 : + let np = max(1, $M.nprocs); ceil(n / np) / (n / np) end) + $(accbinds...) + $(argbinds...) + return $clamp_cost(Float64($(esc(body)))) + end + push!($MODELED_OPS, $opq) + $opq + end +end + +# --------------------------------------------------------------------------- +# Generic fallback +# --------------------------------------------------------------------------- + +""" + OP_AFFINITY + +Coarse block-shape preference per operation, consulted by +[`generic_op_cost`](@ref) when no analytic model exists. This is how an +unmodelled operation still exerts *some* pull on the planner. + +- `:square` — wants near-cubic tiles (dense factorizations, GEMM) +- `:rowwise` — wants whole rows resident per block +- `:colwise` — wants whole columns resident per block +- `:large` — indifferent to shape, wants few large blocks (elementwise, reductions) +- `:any` — genuinely indifferent + +TODO(learning): this table is hand-maintained and will rot as operations are +added. It should be *learned*: track measured cost against the aspect ratio and +tile count actually used per op, and fit the affinity rather than asserting it. +The tape already produces the `(op, layout, outcome)` triples needed, and this +is a much smaller fitting problem than learning a full cost model. +""" +const OP_AFFINITY = Dict{Symbol,Symbol}( + :cholesky => :square, :cholesky! => :square, :potrf! => :square, + :lu => :square, :lu! => :square, :getrf! => :square, + :qr => :square, :qr! => :square, :geqrf! => :square, + :svd => :square, :svd! => :square, + :mul! => :square, :gemm! => :square, :syrk! => :square, + :trsm! => :square, :ldiv! => :square, :rdiv! => :square, + :trsv! => :rowwise, + :transpose => :square, :adjoint => :square, :permutedims => :square, + :map => :large, :map! => :large, :broadcast => :large, + :materialize! => :large, :copyto! => :large, :fill! => :large, + :reduce => :large, :mapreduce => :large, :sum => :large, :norm => :large, + :stencil => :square, + :sort => :rowwise, :sort! => :rowwise, +) + +""" + shape_penalty(v::ArgView, op::Symbol, m::MachineModel) -> Float64 + +Multiplicative penalty in `[1, Inf)` for a layout that is a poor fit for `op`'s +known affinity, or that is unreasonable on its own terms: tiles too small to +amortise task overhead, tiles too large to fit a processor's memory, or too few +tiles to keep every processor busy. + +These three self-evident penalties matter more than the affinity term. They are +what make layout errors *asymmetric* in the model the way they are in reality — +too-small blocks collapse the scheduler, too-large blocks run out of memory — +so the planner will not trade a small predicted win for a catastrophic risk. +""" +function shape_penalty(v::ArgView, op::Symbol, m::MachineModel) + p = 1.0 + tb = tilebytes(v) + nt = ntiles(v) + + # A task must do enough work to be worth launching. + if tb > 0 + min_useful = m.task_overhead * m.bandwidth + tb < min_useful && (p *= min_useful / tb) + end + + # A tile (and the handful a blocked algorithm holds live) must fit. + budget = m.mem_per_proc / 4 + tb > budget && (p *= 1.0 + 8.0 * (tb / budget - 1.0)) + + # There must be enough tiles to occupy every processor. + nt < m.nprocs && (p *= m.nprocs / max(1, nt)) + + aff = get(OP_AFFINITY, op, :any) + if aff === :square + p *= 1.0 + 0.25 * log2(max(1.0, aspect(v))) + elseif aff === :rowwise + nd = ndims(v) + nd >= 2 && nblocks(v, nd) > 1 && (p *= 1.0 + 0.5 * log2(nblocks(v, nd))) + elseif aff === :colwise + ndims(v) >= 1 && nblocks(v, 1) > 1 && (p *= 1.0 + 0.5 * log2(nblocks(v, 1))) + elseif aff === :large + p *= 1.0 + 0.05 * log2(max(1, nt)) + end + + return p +end + +""" + generic_op_cost(args, m, op=:_unknown) -> Float64 + +Shape-agnostic cost estimate used when no [`@cost_model`](@ref) matches. + +Charges for data volume touched once, per-tile scheduler overhead, load +imbalance, per-tile latency, and the [`shape_penalty`](@ref). + +This also doubles as the planner's "unknown future operation" stand-in: the +residual probability mass of a prediction is charged at this rate, so layouts +that are pathological in general are penalised even when they are perfect for +the modal chain. That is what keeps a misprediction merely suboptimal instead +of catastrophic. +""" +function generic_op_cost(args::Vector{ArgView}, m::MachineModel, op::Symbol = :_unknown) + isempty(args) && return 0.0 + np = max(1, m.nprocs) + total = 0.0 + for v in args + eltype(v) === Nothing && continue + nt = ntiles(v) + c = nbytes(v) / m.bandwidth # touch the data once + c += nt * m.task_overhead # per-tile scheduling (serial) + c += nt * m.latency # per-tile boundary crossing + c *= nt <= 0 ? 1.0 : ceil(nt / np) / (nt / np) # load imbalance + c *= shape_penalty(v, op, m) + total += c + end + return clamp_cost(total) +end + +""" + cost_of(op, args, m=current_machine()) -> Float64 + +Dispatch to the registered model for `op`, falling back to the affinity-aware +generic model. +""" +cost_of(op::Symbol, args::Vector{ArgView}, m::MachineModel = current_machine()) = + clamp_cost(op_cost(Val(op), args, m)) + +""" + redistribution_cost(from, to, spec, m) -> Float64 + +Cost of changing an array's layout mid-chain. Charged as a full all-to-all of +the array's bytes plus per-tile latency and task overhead on both sides, +because in general every element lands on a different processor. + +TODO(precision): this is a deliberate overestimate in the cases where it is +wrong — `Blocks(1024,1024) -> Blocks(512,512)` under the same assignment is a +purely local subdivision that moves nothing. Computing the true volume needs +both proc grids and an intersection of the two `DomainBlocks`; `build_procgrid` +plus `src/datadeps/aliasing.jl` has the pieces. Overestimating is the safe +direction (it suppresses speculative repartitioning), so this is not urgent, +but it does leave real wins on the table for refinement-style changes. +""" +function redistribution_cost(from::LayoutChoice, to::LayoutChoice, spec::ArgSpec, + m::MachineModel = current_machine()) + from == to && return 0.0 + sz = size(spec) + bytes = Float64(prod(sz; init = 1)) * _elsize(spec.eltype) + nt_to = prod(nblocks(to, sz); init = 1) + nt_from = prod(nblocks(from, sz); init = 1) + return bytes / m.bandwidth + + (nt_to + nt_from) * (m.latency + m.task_overhead) +end + +# --------------------------------------------------------------------------- +# Default models for Dagger's built-in operations. +# +# These are first-order: they capture the terms that actually *differ* between +# layouts (tile count, critical-path length, communication volume) and ignore +# constants that do not, since the planner only ever compares layouts against +# each other. They are meant to be replaced by measurement-corrected versions +# once the autotuner is wired in. +# +# The argument order below matches the order these operations are instrumented +# with in `integration.jl` — keep the two in sync. +# --------------------------------------------------------------------------- + +@cost_model cholesky!(A) = begin + n = size(A, 1) + # `cholesky!` squares off its input before factorizing; cost what runs. + b, _ = square_tile(A) + nt = max(1, cld(n, b)) + # nt POTRF panels sit on the critical path; the rest is parallel SYRK/GEMM/TRSM. + serial_time(nt * b^3 / 3) + + flops_time(Float64(n)^3 / 3) * imbalance(nt * (nt + 1) / 2) + + task_time(nt^3 / 3 + nt^2) + + bytes_time(nt^2 * b^2 * elsize(A)) + + reblock_time(A) +end + +@cost_model lu!(A) = begin + n = size(A, 1) + b, _ = square_tile(A) + nt = max(1, cld(n, b)) + serial_time(nt * b^3) + + flops_time(2 * Float64(n)^3 / 3) * imbalance(nt * nt) + + # Row swaps against the trailing submatrix dominate the task count. + task_time(nt^3 / 2 + nt^3 / 3 + nt^2) + + bytes_time(nt * nt * b^2 * elsize(A)) + + reblock_time(A) +end + +@cost_model qr!(A) = begin + mm = size(A, 1); nn = size(A, 2) + b, _ = square_tile(A) + ntr = max(1, cld(mm, b)); ntc = max(1, cld(nn, b)) + serial_time(ntc * b^3) + + flops_time(2 * Float64(mm) * nn^2 - 2 * Float64(nn)^3 / 3) * imbalance(ntr * ntc) + + task_time(ntr * ntc^2) + + bytes_time(ntr * ntc * b^2 * elsize(A)) + + reblock_time(A) +end + +@cost_model mul!(C, A, B) = begin + mm = size(A, 1); kk = size(A, 2); nn = size(B, 2) + ntm = nblocks(C, 1); ntn = nblocks(C, 2); ntk = nblocks(A, 2) + flops_time(2.0 * mm * kk * nn) * imbalance(ntm * ntn) + + task_time(ntm * ntn * ntk) + + # Each C tile pulls a block row of A and a block column of B. + bytes_time(ntm * ntn * ntk * + (prod(blocksize(A); init = 1) + prod(blocksize(B); init = 1)) * elsize(A)) + + # Mismatched inner blocking forces a repack of one operand. + (blocksize(A, 2) == blocksize(B, 1) ? 0.0 : bytes_time(nbytes(B))) +end + +@cost_model syrk!(C, A) = begin + nn = size(C, 1); kk = size(A, 2) + ntn = nblocks(C, 1); ntk = nblocks(A, 2) + flops_time(Float64(nn)^2 * kk) * imbalance(ntn * (ntn + 1) / 2) + + task_time(ntn * (ntn + 1) * ntk / 2) + + bytes_time(ntn * ntn * ntk * prod(blocksize(A); init = 1) * elsize(A)) +end + +@cost_model trsm!(A, B) = begin + nn = size(A, 1) + nd = ndims(B) + nrhs = nd >= 2 ? size(B, 2) : 1 + # `ldiv!` squares `A` off (and re-blocks `B` to match) before calling trsm!. + b, _ = square_tile(A) + nta = max(1, cld(nn, b)); ntb = max(1, cld(nrhs, b)) + # Triangular solve is inherently sequential along the block diagonal. + serial_time(nta * Float64(b)^2 * b) + + flops_time(Float64(nn)^2 * nrhs) * imbalance(nta * ntb) + + task_time(nta * nta * ntb / 2) + + bytes_time(nta * ntb * Float64(b)^2 * elsize(B)) + + reblock_time(A) +end + +@cost_model trsv!(A, B) = begin + nn = size(A, 1) + b, _ = square_tile(A) + nta = max(1, cld(nn, b)) + serial_time(nta * Float64(b)^2) + + flops_time(Float64(nn)^2) * imbalance(nta) + + task_time(nta * (nta + 1) / 2) + + bytes_time(nta * nta * Float64(b) * elsize(A)) + + reblock_time(A) +end + +@cost_model map(A) = flops_time(Float64(length(A))) * imbalance(ntiles(A)) + task_time(ntiles(A)) +@cost_model map!(A) = flops_time(Float64(length(A))) * imbalance(ntiles(A)) + task_time(ntiles(A)) +@cost_model copyto!(A) = bytes_time(nbytes(A)) / max(1, nprocs_avail()) + task_time(ntiles(A)) + +@cost_model reduce(A) = begin + nt = ntiles(A) + flops_time(Float64(length(A))) * imbalance(nt) + + task_time(nt) + + (nt > 1 ? bytes_time(nt * elsize(A)) * log2(nt) : 0.0) # tree reduction over tiles +end + +@cost_model mapreduce(A) = begin + nt = ntiles(A) + flops_time(Float64(length(A))) * imbalance(nt) + + task_time(nt) + + (nt > 1 ? bytes_time(nt * elsize(A)) * log2(nt) : 0.0) +end + +@cost_model transpose(A) = begin + # A full transpose is an all-to-all unless the blocking is symmetric. + sym = ndims(A) >= 2 && blocksize(A, 1) == blocksize(A, 2) + bytes_time(nbytes(A)) * (sym ? 0.5 : 1.0) + task_time(ntiles(A)) +end + +@cost_model permutedims(A) = bytes_time(nbytes(A)) + task_time(ntiles(A)) + +@cost_model stencil(A) = begin + nt = ntiles(A) + tile = Float64(prod(blocksize(A); init = 1)) + halo = 0.0 + for d in 1:ndims(A) + halo += tile / max(1, blocksize(A, d)) # one face per dimension + end + flops_time(Float64(length(A))) * imbalance(nt) + + task_time(nt) + + bytes_time(2 * nt * halo * elsize(A)) +end diff --git a/src/tapes/integration.jl b/src/tapes/integration.jl new file mode 100644 index 000000000..1f12cc5bc --- /dev/null +++ b/src/tapes/integration.jl @@ -0,0 +1,263 @@ +# =========================================================================== +# Integration helpers. +# +# This file holds the glue that touches Dagger's existing code paths. It is +# kept separate from the rest of the subsystem so that the blast radius of the +# feature is obvious: everything else in `tapes/` is self-contained and has no +# effect until something here calls into it. +# +# It is included from `Tapes.jl`, so everything here lives in `Dagger.Tapes` +# and is referenced from Dagger as `Tapes.resolve_partitioning(...)` etc. +# =========================================================================== + +""" + resolve_partitioning(::Type{T}, dims, requested, assignment) -> (part, assignment, plan) + +The single entry point for allocators. Returns the partitioning and assignment +to allocate with, plus the [`AllocationPlan`](@ref) to hand to +[`track!`](@ref). + +When the subsystem is disabled this resolves `AutoBlocks()` exactly as before +(`auto_blocks(dims)`), so behaviour is bit-identical to the status quo. +""" +function resolve_partitioning(::Type{T}, dims::Tuple, requested, assignment) where {T} + if !is_enabled() + part = requested isa AutoBlocks ? auto_blocks(map(Int, dims)::Dims) : requested + return (part, assignment, passthrough(part, assignment)) + end + p = plan_allocation(T, dims; requested = requested, assignment = assignment) + part = p.partitioning isa AutoBlocks ? auto_blocks(map(Int, dims)::Dims) : p.partitioning + return (part, p.assignment, p) +end + +""" + @tracked_alloc T dims requested assignment expr + +Convenience wrapper for an allocation site. Binds `part` and `assign` for use +inside `expr` (which must evaluate to the new `DArray`), and tracks the result. + +```julia +function Base.zeros(p::BlocksOrAuto, T::Type, dims::Dims; assignment = :arbitrary) + Dagger.Tapes.@tracked_alloc T dims p assignment begin + _zeros_impl(part, T, dims, assign) + end +end +``` +""" +macro tracked_alloc(T, dims, requested, assignment, expr) + quote + local __T__ = $(esc(T)) + local __d__ = $(esc(dims)) + local __part__, __assign__, __plan__ = + $resolve_partitioning(__T__, __d__, $(esc(requested)), $(esc(assignment))) + local $(esc(:part)) = __part__ + local $(esc(:assign)) = __assign__ + local __A__ = $(esc(expr)) + __A__ isa $DArray ? $track!(__A__, __plan__) : __A__ + end +end + +# =========================================================================== +# PATCH POINTS +# +# The changes below are written out rather than applied, because they touch +# files this subsystem does not own. Each is small and independently +# revertible. Nothing here changes behaviour while `CONFIG.enabled == false`. +# =========================================================================== + +#= + +## 1. `src/Dagger.jl` + +Insert after the `array/darray.jl` include (currently line 126), before +`array/alloc.jl`: + +```julia +include("tapes/Tapes.jl"); using .Tapes +include("tapes/integration.jl") +``` + +`using .Tapes` brings `@record_op`, `@cost_model` and `@expect_ops` into +`Dagger`, so users write `Dagger.@record_op` and `Dagger.@expect_ops`. + +This must come *after* `array/darray.jl` (which defines `DArray`, `Blocks`, +`AutoBlocks` and `auto_blocks`) and *before* `array/alloc.jl` and the +linear-algebra files, which call into it. + +## 2. `src/array/alloc.jl` — the allocation sites + +Every constructor that accepts `BlocksOrAuto` is a decision point. The pattern +is the same for `rand`, `randn`, `ones`, `zeros`, `sprand` and the +`DArray{T,N}(undef, ...)` family. Taking `zeros` as the example: + +```julia +function Base.zeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Tapes.resolve_partitioning(T, dims, p, assignment) + d = ArrayDomain(map(x -> 1:x, dims)) + s = reduce(vcat, partition(part, d)) + procgrid = build_procgrid(assign, dims, part.blocksize, current_acceleration()) + a = AllocateArray(T, AllocateZeros{T}(), false, d, s, part, procgrid) + return Tapes.track!(_to_darray(a), plan) +end + +# The AutoBlocks method must stop short-circuiting to `auto_blocks` so the +# planner gets a say: +Base.zeros(p::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = + _zeros_planned(T, dims, p, assignment) +``` + +The important structural point: **`AutoBlocks` must reach +`resolve_partitioning` un-lowered**. At present `AutoBlocks` methods call +`auto_blocks(A)` immediately and forward to the `Blocks` method, which erases +the distinction between "the user asked for automatic" and "the user asked for +this specific blocking". The tape system needs that distinction — it declines +to override explicit `Blocks` by default. Threading `requested` through as +shown preserves it. + +## 3. `src/array/darray.jl` — `distribute` + +`distribute(A, AutoBlocks(), assignment)` is the other main entry point: + +```julia +function distribute(A::AbstractArray{T,N}, ::AutoBlocks, + assignment::AssignmentType = :arbitrary) where {T,N} + part, assign, plan = Tapes.resolve_partitioning(T, size(A), AutoBlocks(), assignment) + return Tapes.track!(distribute(A, part, assign), plan) +end +``` + +## 4. Operation instrumentation + +One `@record_op` per operation, immediately before task submission. Argument +order must match the corresponding `@cost_model` in `tapes/cost.jl` — pass only +`DArray`s. + +- `src/array/cholesky.jl`, in `LinearAlgebra._chol!(A::DArray{T,2}, ...)`: + ```julia + Dagger.@record_op :cholesky! A + ``` + +- `src/array/lu.jl`, in the `lu!` entry point: + ```julia + Dagger.@record_op :lu! A + ``` + +- `src/array/qr.jl`, in the `qr!` entry point: + ```julia + Dagger.@record_op :qr! A + ``` + +- `src/array/mul.jl`, in `LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, alpha, beta)`: + ```julia + Dagger.@record_op :mul! C A B + ``` + and in `syrk_dagger!(C, transA, A, alpha, beta)`: + ```julia + Dagger.@record_op :syrk! C A + ``` + Note `generic_matvecmul!` should record as `:mul!` too, with `B` a `DVector`; + the cost model handles `ndims(B) == 1` via its `nd` branch. + +- `src/array/trsm.jl`: + ```julia + Dagger.@record_op :trsm! A B # in trsm! + Dagger.@record_op :trsv! A B # in trsv! + ``` + +- `src/array/map-reduce.jl`: `:map`, `:map!`, `:reduce`, `:mapreduce`. +- `src/array/copy.jl`: `:copyto!`. +- `src/array/permute.jl`: `:transpose`, `:permutedims`. +- `src/array/stencil.jl`: `:stencil`. +- `src/array/sort.jl`: `:sort!`. + +TODO(coverage): `src/array/operators.jl` broadcasting is the highest-value +missing instrumentation — broadcast fusion is where most real user code spends +its operations, and a fused broadcast's layout preference (`:large`, shape +indifferent) is a genuinely useful signal that the current chain misses. It +needs care because a single `materialize!` covers an arbitrary expression tree; +recording it as one `:broadcast` op loses which arrays participated. + +TODO(datadeps): operations submitted inside `spawn_datadeps` are the case where +this system should shine, because Datadeps already knows the full dependency +structure of the region. Rather than instrumenting each operation, the Datadeps +queue itself could record the op sequence per aliased region at +`distribute_tasks!` time — exact lookahead for the region, no prediction +required, which subsumes the tape for that scope. That is strictly better where +it applies and should be built before broadening `@record_op` coverage much +further. + +## 5. Precompilation + +`src/precompile.jl` should run at least one workload with the subsystem enabled +so the planner and macro-generated cost models are precompiled; otherwise the +first *enabled* allocation pays a large latency spike that will be mistaken for +tape overhead. + +=# + +# =========================================================================== +# DESIGN TODOs deferred by explicit decision +# =========================================================================== + +#= + +## Deferred materialization + +The cleanest solution to the whole problem is not to predict at all: keep an +array as a *logical description* until its first consumer is known, then +materialise it with a layout chosen from actual knowledge rather than a +forecast. That gets the first operation's layout exactly right with zero +prediction risk, and reduces the tape's job to optimising operations 2..n. + +Deferred: this requires restructuring how Datadeps tracks and materialises +arrays. Specifically, `DArray` would need an unmaterialised state that +`chunks` accessors can trigger materialisation from, every internal +`A.chunks` access would need to route through that, and the Datadeps aliasing +analysis would need to reason about a domain decomposition that does not exist +yet. `AllocateArray`/`stage` is the natural seam but the change reaches much +further than that. + +Note the interaction: once deferred materialization exists, the tape's role +shifts from "choose the layout" to "choose the layout for the operations after +the first", which makes the confidence gating *more* important, not less — +the easy win is gone and only the speculative part remains. + +## Multiple materialised layouts for read-only data + +For data that is read but not written, maintaining two physical layouts costs +memory but requires no coherence protocol at all. This is exactly what column +stores do with multiple sort-order projections (C-Store / Vertica). For a chain +where operation A wants square tiles and operations B..F want row blocks, +replicating can easily dominate repartitioning. + +Deferred: doing this without introducing OOM failures requires a memory +accounting system that Dagger does not have — something that tracks per-space +residency, knows the high-water mark of the current task graph, and can evict +or refuse a replica under pressure. MemPool has pieces of this +(`MemPool.approx_size`, the storage device abstraction) but there is no +global admission control. Building replication on top of the current +best-effort memory handling would turn a performance feature into a +reliability regression. + +The precondition is a data-residency tracker: (array, layout) -> space, bytes, +last-touch, read-only flag. That is the same structure needed for a proper +`redistribution_cost`, for heterogeneous cost modelling, and for out-of-core +support, so it is probably the highest-leverage missing piece overall. + +## Verified lookahead vs prediction + +A submission-window scheme — buffer N submitted tasks or T microseconds before +scheduling — gives *exact* lookahead with zero prediction risk, and beats this +system on any workload whose operation chain fits in one scope. The tape earns +its keep only where a window cannot reach: operations scattered across user +code, deep call stacks, library boundaries, and loop iterations that a bounded +buffer cannot span. + +These are complements, not alternatives, and the strongest configuration is +both: exact layouts for the verified near-term operations, predicted layouts +for the tail. `plan_chain` already accepts an arbitrary `Vector{PredictedOp}`, +so a window would simply contribute entries with `prob = 1.0` ahead of the +predicted ones. Wiring that up is cheap once a window exists; `@expect_ops` +with an explicit list is the manual version of the same thing. + +=# diff --git a/src/tapes/plan.jl b/src/tapes/plan.jl new file mode 100644 index 000000000..676983f4b --- /dev/null +++ b/src/tapes/plan.jl @@ -0,0 +1,395 @@ +# =========================================================================== +# Candidate generation and layout planning. +# +# Given a predicted operation chain and a set of candidate layouts, choosing +# the layout sequence that minimises total cost (including repartitioning on +# the edges) is exactly the classical automatic-data-layout problem of Kennedy +# & Kremer (TOPLAS 1998). They needed 0-1 integer programming because their +# phase graph was general. Ours is a *linear chain*, so it is a shortest path +# through a layered DAG and an exact O(n * |L|^2) dynamic program suffices. +# =========================================================================== + +const ASSIGNMENTS_ALL = (:arbitrary, :blockrow, :blockcol, :cyclicrow, :cycliccol) + +""" + fallback_layout(::Type{T}, dims) -> LayoutChoice + +The blind default: whatever `Dagger.auto_blocks` would pick, with `:arbitrary` +assignment (which spreads by compute capacity). Not good for much in +particular, but acceptable for most things, which is exactly what a fallback +should be. Every speculative plan is measured against this. +""" +function fallback_layout(::Type{T}, dims::Tuple) where {T} + ab = try + auto_blocks(map(Int, dims)::Dims) + catch + Blocks(ntuple(i -> max(1, Int(dims[i])), length(dims))) + end + return LayoutChoice(ab.blocksize, :arbitrary, :auto) +end + +""" + candidate_layouts(::Type{T}, dims, m) -> Vector{LayoutChoice} + +Enumerate plausible layouts for an array of element type `T` and shape `dims`. + +The set is deliberately small and structured rather than a dense sweep: the +planner is quadratic in candidate count, and the interesting decisions are +between *families* (square tiles vs row blocks vs column blocks) and between a +few block sizes within a family, not between 511 and 512. + +Block sizes are drawn from powers of two bounded below by "big enough that a +task amortises `task_overhead`" and above by "small enough that there are at +least a few tiles per processor". + +TODO(rank>2): the N-d cases here are a reasonable guess (split leading dim, +split trailing dim, near-cubic) but untested against real workloads. Tensor +contractions in particular want layouts derived from the contraction indices, +which the tape does not currently record. +""" +function candidate_layouts(::Type{T}, dims::Tuple, m::MachineModel = current_machine()) where {T} + N = length(dims) + N == 0 && return LayoutChoice[] + esz = max(1, _elsize(T)) + np = max(1, m.nprocs) + out = LayoutChoice[] + + # Smallest tile worth scheduling, and largest that leaves >= 2 tiles/proc. + # `task_overhead` is serialized wall-clock, so it is *not* divided by `np` + # here: a tile has to be big enough to hide the whole submission cost, not + # `1/np` of it. + min_tile_bytes = m.task_overhead * m.bandwidth + min_edge = max(32, ceil(Int, (min_tile_bytes / esz)^(1 / N))) + total_bytes = Float64(prod(dims; init = 1)) * esz + max_tile_bytes = max(min_tile_bytes, total_bytes / (2 * np)) + max_edge = max(min_edge, floor(Int, (max_tile_bytes / esz)^(1 / N))) + + function push_unique!(l::LayoutChoice) + l in out || push!(out, l) + return nothing + end + + # --- Family 1: near-cubic tiles at a few sizes ------------------------- + edges = Int[] + e = 64 + while e <= 8192 + (e >= min_edge ÷ 2 && e <= max_edge * 2) && push!(edges, e) + e *= 2 + end + # Also the edge that yields exactly one tile per processor per dimension. + push!(edges, max(1, cld(maximum(dims), max(1, round(Int, np^(1 / N)))))) + # And the square form of the blind default's block size. Powers of two miss + # it (it is `cld(dims[end], np)`, e.g. 683 for 4096 over 6 procs) and it is + # the layout Dagger's own `maybe_copy_buffered` re-blocks `AutoBlocks` to, + # so it is both a strong candidate and the honest baseline for comparison. + push!(edges, max(1, minimum(blocksize(fallback_layout(T, dims))))) + for edge in unique!(sort!(edges)) + bs = ntuple(i -> clamp(edge, 1, max(1, Int(dims[i]))), N) + for a in (N >= 2 ? (:cyclicrow, :cycliccol, :arbitrary) : (:arbitrary,)) + push_unique!(LayoutChoice(bs, a, :square)) + end + end + + # --- Family 2: block rows (whole trailing dims resident) --------------- + for div in (np, 2np, 4np) + bs = ntuple(i -> i == 1 ? max(1, cld(Int(dims[1]), div)) : Int(dims[i]), N) + push_unique!(LayoutChoice(bs, :blockrow, :rowblock)) + push_unique!(LayoutChoice(bs, :arbitrary, :rowblock)) + end + + # --- Family 3: block columns (whole leading dims resident) ------------- + if N >= 2 + for div in (np, 2np, 4np) + bs = ntuple(i -> i == N ? max(1, cld(Int(dims[N]), div)) : Int(dims[i]), N) + push_unique!(LayoutChoice(bs, :blockcol, :colblock)) + push_unique!(LayoutChoice(bs, :arbitrary, :colblock)) + end + end + + # --- Family 4: whatever the blind default would have been -------------- + fb = fallback_layout(T, dims) + push_unique!(fb) + + # Trim to budget, spreading across families. The fallback is seeded first + # so it always survives: the planner needs it present as the baseline. + if length(out) > CONFIG.max_candidates + keep = LayoutChoice[fb] + per_family = max(1, cld(CONFIG.max_candidates - 1, 3)) + for fam in (:square, :rowblock, :colblock) + fams = filter(l -> l.label === fam, out) + isempty(fams) && continue + stride = max(1, cld(length(fams), per_family)) + for l in fams[1:stride:end] + l in keep || push!(keep, l) + end + end + out = length(keep) > CONFIG.max_candidates ? keep[1:CONFIG.max_candidates] : keep + end + return out +end + +# --------------------------------------------------------------------------- +# Turning a predicted operation into cost-model input +# --------------------------------------------------------------------------- + +""" + views_for(pred::PredictedOp, self::ArgSpec, layout) -> Vector{ArgView} + +Build the [`ArgView`](@ref) list a cost model expects for one predicted +operation: the tracked array at position `pred.key.pos` takes the candidate +`layout`; the other arguments take whatever layout they were last observed +with. + +TODO(co-argument layouts): using the *last observed* layout for co-arguments is +the weakest link in the whole chain. It means the planner evaluates "what if I +change only this array" rather than solving for a jointly consistent +assignment, so it systematically undervalues layouts that are only good when +all operands agree — precisely the layouts that matter for GEMM and TRSM. Fix +this together with the joint-planning TODO in `record_op!`: once arrays are +grouped into connected components, `views_for` should take a component-wide +assignment rather than a single layout. + +TODO(distributional specs): co-argument specs are the last observation, not a +distribution. If the same site alternates between a 1-column and 512-column +right-hand side, we silently plan for whichever ran last. +""" +function views_for(pred::PredictedOp, self::ArgSpec, layout::LayoutChoice) + pos = Int(pred.key.pos) + arity = Int(pred.key.arity) + specs = pred.argspecs + n = max(arity, length(specs), pos) + views = Vector{ArgView}(undef, n) + @inbounds for i in 1:n + spec = i <= length(specs) ? specs[i] : self + if i == pos + views[i] = ArgView(self, layout) + else + lay = LayoutChoice(unpad_dims(spec.blocksize, spec.ndims), + spec.assignment, :observed) + views[i] = ArgView(spec, lay) + end + end + return views +end + +""" + step_cost(pred, self, layout, m) -> Float64 + +Expected cost contribution of one predicted operation under `layout`. + +Weighted two ways: +- with probability `pred.prob` the predicted operation happens, and costs what + its model says; +- with the residual probability `1 - pred.prob` *something else* happens, which + we charge at the generic rate. That residual term is what makes the objective + risk-aware: a layout that is superb for the predicted chain but absurd in + general (one giant tile, or a million tiny ones) pays for that here. +""" +function step_cost(pred::PredictedOp, self::ArgSpec, layout::LayoutChoice, m::MachineModel) + views = views_for(pred, self, layout) + c_hit = cost_of(pred.key.op, views, m) + p = clamp(pred.prob, 0.0, 1.0) + p >= 1.0 && return c_hit + c_miss = generic_op_cost(ArgView[ArgView(self, layout)], m) + return p * c_hit + (1 - p) * c_miss +end + +# --------------------------------------------------------------------------- +# The planner +# --------------------------------------------------------------------------- + +""" + LayoutPlan + +Result of planning. `steps[i]` is the layout the planner wants for predicted +operation `i`; `steps[1]` is the only one actually committed under the +receding-horizon policy (see [`plan_chain`](@ref)). +""" +struct LayoutPlan + steps::Vector{LayoutChoice} + cost::Float64 + fallback_cost::Float64 + "Worst per-operation ratio of chosen layout to best-possible layout." + max_regret::Float64 + accepted::Bool + reason::Symbol +end + +""" + plan_chain(pred, cands, self, m; start=nothing) -> LayoutPlan + +Exact dynamic program over the predicted chain. + +`D[i, l]` is the minimum expected cost of executing predicted operations +`1..i` with operation `i` running under candidate `l`, including any +repartitioning charged on the edges. Backtracking the argmin over the final +layer gives the optimal layout sequence. + +If `start` is given it is the array's *current* layout, and the first step is +charged the cost of moving from it (used when re-planning mid-chain). + +# Policy note + +The plan covers the whole horizon, but callers commit only `steps[1]` and +re-plan at each subsequent observed operation. That is deliberate: a +speculative repartition is O(data) network traffic, so plan-and-commit exposes +the whole chain's worth of data movement to a single misprediction, whereas +receding-horizon control keeps the multi-operation lookahead benefit while +bounding exposure to one decision at a time. +""" +function plan_chain(pred::Vector{PredictedOp}, cands::Vector{LayoutChoice}, + self::ArgSpec, m::MachineModel; + start::Union{Nothing,LayoutChoice} = nothing) + n = length(pred) + L = length(cands) + fb = fallback_layout(self.eltype, size(self)) + + (n == 0 || L == 0) && + return LayoutPlan([fb], Inf, Inf, 1.0, false, :no_prediction) + + D = fill(Inf, n, L) + P = zeros(Int, n, L) + + # Per-operation cost of every candidate, reused for the regret bound. + C = Matrix{Float64}(undef, n, L) + @inbounds for i in 1:n, l in 1:L + C[i, l] = step_cost(pred[i], self, cands[l], m) + end + + @inbounds for l in 1:L + entry = start === nothing ? 0.0 : redistribution_cost(start, cands[l], self, m) + D[1, l] = C[1, l] + entry + end + @inbounds for i in 2:n + for l in 1:L + best = Inf; bi = 0 + for lp in 1:L + c = D[i-1, lp] + redistribution_cost(cands[lp], cands[l], self, m) + if c < best + best = c; bi = lp + end + end + D[i, l] = best + C[i, l] + P[i, l] = bi + end + end + + # Backtrack. + bestl = 1; bestc = Inf + @inbounds for l in 1:L + D[n, l] < bestc && (bestc = D[n, l]; bestl = l) + end + steps = Vector{LayoutChoice}(undef, n) + l = bestl + for i in n:-1:1 + steps[i] = cands[l] + i > 1 && (l = P[i, l]) + end + + # Same DP, pinned to the fallback layout throughout: the baseline to beat. + fb_idx = findfirst(==(fb), cands) + fb_cost = 0.0 + if fb_idx === nothing + fbviews_cost = 0.0 + for i in 1:n + fbviews_cost += step_cost(pred[i], self, fb, m) + end + fb_cost = fbviews_cost + (start === nothing ? 0.0 : + redistribution_cost(start, fb, self, m)) + else + for i in 1:n + fb_cost += C[i, fb_idx] + end + fb_cost += start === nothing ? 0.0 : redistribution_cost(start, cands[fb_idx], self, m) + end + + # Regret bound: how much worse is the committed layout than the best + # possible layout, for the single operation where it does worst? + committed = steps[1] + ci = findfirst(==(committed), cands) + max_regret = 1.0 + if ci !== nothing + @inbounds for i in 1:n + best_i = minimum(@view C[i, :]) + best_i > 0 && isfinite(best_i) && + (max_regret = max(max_regret, C[i, ci] / best_i)) + end + end + + conf = confidence(pred) + accepted = true + reason = :ok + if !isfinite(bestc) + accepted = false; reason = :nonfinite_cost + elseif bestc > CONFIG.gate_margin * fb_cost + # Not enough predicted win to justify departing from the known-adequate default. + accepted = false; reason = :insufficient_margin + elseif conf < CONFIG.regret_confidence_threshold && max_regret > CONFIG.max_regret_ratio + # Low confidence: refuse layouts that are catastrophic for any single + # predicted operation, even if they win on the weighted sum. + accepted = false; reason = :regret_too_high + end + + return LayoutPlan(steps, bestc, fb_cost, max_regret, accepted, reason) +end + +# --------------------------------------------------------------------------- +# Receding-horizon replanning +# --------------------------------------------------------------------------- + +""" + maybe_repartition!(A::DArray, trace::LiveTrace) + +Hook called after each recorded operation when `CONFIG.allow_repartition` is +set: re-plan from the array's current position in the trie and, if the plan +calls for a different layout and the predicted win exceeds the redistribution +cost, physically repartition `A`. + +**Currently a no-op stub.** The planning half is real (`replan` below); the +acting half is not, and deliberately so — see the TODO. + +TODO(repartition-mechanism): performing this safely requires a real +redistribution primitive that (a) is expressed as Datadeps tasks so it +interleaves with in-flight work rather than serialising behind it, (b) mutates +the `DArray` in place — `domain`, `subdomains`, `chunks` and `partitioning` all +have to change atomically from the perspective of concurrent readers — and +(c) is safe against the array being captured by tasks already submitted but not +yet scheduled. (c) is the hard part: Dagger's eager submission means there may +be queued thunks holding `Chunk` references into the old block structure. The +plausible approach is to route repartitioning through the Datadeps queue as a +normal `InOut` task on the whole array, which makes the dependency explicit and +lets the existing scheduler serialise it correctly, at the cost of a +synchronisation point. + +Until that exists, the planner's multi-step output is informational: it shows +up in `explain` and can guide manual `pin!` decisions, but only `steps[1]` is +ever acted on. +""" +function maybe_repartition!(A::DArray, trace::LiveTrace) + CONFIG.allow_repartition || return nothing + plan = replan(trace) + plan === nothing && return nothing + isempty(plan.steps) && return nothing + want = plan.steps[1] + want == trace.layout && return nothing + plan.accepted || return nothing + vlog("would repartition from $(trace.layout) to $want (not implemented)") + return nothing +end + +""" + replan(trace::LiveTrace) -> Union{Nothing,LayoutPlan} + +Re-run the planner from the array's current position in the trie, taking its +present layout as the starting state. This is the receding-horizon step: the +prediction is now conditioned on the prefix actually observed, so a chain that +diverged from the forecast is re-forecast against the branch it really took. +""" +function replan(trace::LiveTrace) + pred = predict(trace.node) + isempty(pred) && return nothing + self = trace.self + m = current_machine() + cands = candidate_layouts(self.eltype, size(self), m) + return plan_chain(pred, cands, self, m; start = trace.layout) +end diff --git a/src/tapes/tape.jl b/src/tapes/tape.jl new file mode 100644 index 000000000..0a9ced602 --- /dev/null +++ b/src/tapes/tape.jl @@ -0,0 +1,612 @@ +# =========================================================================== +# Core types, allocation-site identification, and the operation-sequence trie. +# =========================================================================== + +""" +Maximum array rank tracked in fixed-width tuples. Ranks above this are still +handled, but their trailing dimensions collapse into the last slot, which only +costs precision in the site key (never correctness). +""" +const MAX_TRACKED_DIMS = 8 + +const DimTuple = NTuple{MAX_TRACKED_DIMS,Int} + +@inline function pad_dims(t::Tuple)::DimTuple + n = length(t) + ntuple(i -> i <= n ? Int(t[i]) : 0, MAX_TRACKED_DIMS) +end + +@inline unpad_dims(t::DimTuple, n::Integer) = ntuple(i -> t[i], Int(n)) + +# --------------------------------------------------------------------------- +# Keys +# --------------------------------------------------------------------------- + +""" + SiteKey + +Identifies an allocation *site* together with the coarse shape of what it +allocates. Two allocations share a key only if they came from the same program +context, have the same element type and rank, and fall into the same size +bucket — the last of these matters because the best block size for a 1024^2 +matrix is not the best block size for a 100000^2 matrix even at an identical +call site. +""" +struct SiteKey + site::UInt64 + eltype::DataType + ndims::Int32 + szbucket::DimTuple +end + +Base.hash(k::SiteKey, h::UInt) = + hash(k.szbucket, hash(k.ndims, hash(k.eltype, hash(k.site, h)))) +Base.:(==)(a::SiteKey, b::SiteKey) = + a.site == b.site && a.eltype === b.eltype && a.ndims == b.ndims && a.szbucket == b.szbucket + +""" + size_bucket(dims) -> DimTuple + +Quantise `dims` logarithmically so that nearby problem sizes share a tape. +`CONFIG.size_buckets_per_octave == 0` disables bucketing (exact sizes). +""" +function size_bucket(dims::Tuple)::DimTuple + bpo = CONFIG.size_buckets_per_octave + if bpo <= 0 + return pad_dims(dims) + end + n = length(dims) + ntuple(MAX_TRACKED_DIMS) do i + i > n && return 0 + d = Int(dims[i]) + d <= 0 ? 0 : round(Int, log2(d) * bpo) + end +end + +""" + OpKey + +An operation as seen *from the perspective of one of its arguments*. The +argument position is part of the identity because layout preference depends on +role: being the `A` of `mul!(C, A, B)` implies a different preference from +being the `C`. +""" +struct OpKey + op::Symbol + pos::Int16 + arity::Int16 +end +OpKey(op::Symbol, pos::Integer, arity::Integer) = OpKey(op, Int16(pos), Int16(arity)) + +Base.hash(k::OpKey, h::UInt) = hash(k.arity, hash(k.pos, hash(k.op, h))) +Base.:(==)(a::OpKey, b::OpKey) = a.op === b.op && a.pos == b.pos && a.arity == b.arity +Base.show(io::IO, k::OpKey) = print(io, k.op, "[", k.pos, "/", k.arity, "]") + +const ROOT_OPKEY = OpKey(Symbol("#root"), 0, 0) + +# --------------------------------------------------------------------------- +# Argument and layout descriptions +# --------------------------------------------------------------------------- + +""" + ArgSpec + +Metadata-only snapshot of a `DArray` argument. Deliberately holds no reference +to the array itself so that recorded tapes never keep data alive. +""" +struct ArgSpec + eltype::DataType + ndims::Int32 + size::DimTuple + blocksize::DimTuple + assignment::Symbol +end + +function ArgSpec(A::DArray{T,N}) where {T,N} + bs = A.partitioning isa Blocks ? A.partitioning.blocksize : size(A) + # N.B. `DArray` does not persist the assignment used to build it, so we + # cannot recover it here. + # TODO(assignment-provenance): add an `assignment` (or `procgrid`) field to + # `DArray` so recorded specs and cost models can reason about the actual + # processor mapping rather than assuming `:arbitrary`. Without it the + # planner can pick a good *block shape* but has to guess at the mapping + # quality of the arrays it did not itself allocate. + ArgSpec(T, Int32(N), pad_dims(size(A)), pad_dims(bs), :arbitrary) +end + +ArgSpec(::Type{T}, dims::Tuple, bs::Tuple, assignment::Symbol) where {T} = + ArgSpec(T, Int32(length(dims)), pad_dims(dims), pad_dims(bs), assignment) + +Base.eltype(s::ArgSpec) = s.eltype +Base.ndims(s::ArgSpec) = Int(s.ndims) +Base.size(s::ArgSpec) = unpad_dims(s.size, s.ndims) +Base.size(s::ArgSpec, d::Integer) = d <= s.ndims ? s.size[d] : 1 + +""" + LayoutChoice + +A candidate partitioning: a block shape plus a processor assignment strategy. +`label` records the family it came from, purely for explanation output. +""" +struct LayoutChoice + blocksize::DimTuple + ndims::Int32 + assignment::Symbol + label::Symbol +end + +LayoutChoice(bs::Tuple, assignment::Symbol, label::Symbol=:custom) = + LayoutChoice(pad_dims(bs), Int32(length(bs)), assignment, label) + +blocksize(l::LayoutChoice) = unpad_dims(l.blocksize, l.ndims) +blocksize(l::LayoutChoice, d::Integer) = d <= l.ndims ? l.blocksize[d] : 1 + +"Convert to a Dagger `Blocks{N}` suitable for passing to an allocator." +to_blocks(l::LayoutChoice) = Blocks(blocksize(l)) + +Base.:(==)(a::LayoutChoice, b::LayoutChoice) = + a.blocksize == b.blocksize && a.ndims == b.ndims && a.assignment === b.assignment +Base.hash(l::LayoutChoice, h::UInt) = hash(l.assignment, hash(l.ndims, hash(l.blocksize, h))) + +function Base.show(io::IO, l::LayoutChoice) + print(io, "Blocks", blocksize(l), " / :", l.assignment) + l.label === :custom || print(io, " (", l.label, ")") +end + +"Number of blocks along each dimension for an array of `dims` under `l`." +nblocks(l::LayoutChoice, dims::Tuple) = + ntuple(i -> max(1, cld(Int(dims[i]), max(1, blocksize(l, i)))), length(dims)) + +# --------------------------------------------------------------------------- +# Allocation-site identification +# --------------------------------------------------------------------------- + +""" + lexical_token(mod, source) -> UInt64 + +A compile-time-stable identifier for a macro expansion site. Computed at +macro-expansion time and baked into the expansion as a literal, so it costs +nothing at runtime. +""" +function lexical_token(mod::Module, source::LineNumberNode) + h = hash(nameof(mod), UInt64(0x9e3779b97f4a7c15)) + h = hash(something(source.file, :unknown), h) + h = hash(source.line, h) + return h % UInt64 +end + +""" + CONTEXT_HASH + +Probabilistic calling context, in the sense of Bond & McKinley (OOPSLA 2007). +Maintained incrementally as `V := 3V + site` at instrumentation boundaries +rather than by unwinding, giving context sensitivity for the price of a +multiply-add. Propagates into child tasks automatically because +`ScopedValue`s do. +""" +const CONTEXT_HASH = ScopedValue{UInt64}(UInt64(0)) + +@inline mix_context(v::UInt64, token::UInt64) = 3 * v + token + +""" + backtrace_hash() -> UInt64 + +Hash the raw instruction pointers of the current stack, skipping +`CONFIG.backtrace_skip` frames and taking at most `CONFIG.backtrace_depth`. + +Raw pointers are hashed rather than symbolicated frames because symbolication +dominates the cost by orders of magnitude. The consequence is that keys are +only meaningful within a single session and can shift if a frame is +recompiled at a different specialisation. + +TODO(persistence): to persist tapes across sessions — which is what would make +the very first run of a program fast rather than the second — we need a stable +key. Options, in increasing order of cost: (a) symbolicate lazily only when +serialising and re-resolve on load; (b) hash `(file, line)` pairs obtained from +`StackTraces.lookup`, cached per instruction pointer in an `IdDict` so each +unique frame is symbolicated once; (c) require `:lexical`/`:context` mode for +persistable tapes. (b) is probably the right default — the cache makes steady- +state cost comparable to raw pointer hashing. + +TODO(cost): even unsymbolicated, `backtrace()` is the single most expensive +thing this subsystem does. Consider (i) caching the hash in a task-local slot +keyed by a cheap discriminator so repeated allocations in a loop body unwind +once, or (ii) `jl_unw_stepn`-style bounded unwinding via a small C shim rather +than the full `backtrace()` allocation of a `Vector{Ptr}`. +""" +function backtrace_hash() + bt = backtrace() + n = length(bt) + lo = min(CONFIG.backtrace_skip + 1, n + 1) + hi = min(lo + CONFIG.backtrace_depth - 1, n) + h = UInt64(0xcbf29ce484222325) + @inbounds for i in lo:hi + # Julia 1.12+ backtraces can mix `Ptr` frames with non-isbits + # `Base.InterpreterIP` entries; only pointers can be reinterpreted. + frame = bt[i] + ip = frame isa Ptr ? (UInt(frame) % UInt64) : (hash(frame) % UInt64) + h = (h ⊻ ip) * UInt64(0x100000001b3) + end + return h +end + +""" + current_site(token::UInt64) -> UInt64 + +Resolve the site identifier for the active `CONFIG.site_id` strategy. `token` +is the lexical token of the calling macro expansion. +""" +@inline function current_site(token::UInt64) + mode = CONFIG.site_id + if mode === :lexical + return token + elseif mode === :context + return mix_context(CONTEXT_HASH[], token) + else # :backtrace + # Fold in the ambient context too: a backtrace cannot see through a + # `Threads.@spawn` or `Dagger.@spawn` boundary, but the scoped context + # hash can, and `@expect_ops` regions set it deliberately. + return mix_context(CONTEXT_HASH[], backtrace_hash()) + end +end + +# --------------------------------------------------------------------------- +# The tape: a prefix trie over operation sequences +# --------------------------------------------------------------------------- + +""" + TapeNode + +One node of a per-site prefix trie. The path from the root to a node is an +observed operation sequence; `count` is how many times that exact prefix was +observed. + +A trie rather than a flat list of recorded sequences because it gives, for +free: incremental commit (no need to wait for an array to die before learning +from it), prefix-conditional prediction (re-planning mid-chain naturally +restricts to the branch actually taken), and per-step branch probabilities +that feed straight into the expected-cost objective. +""" +mutable struct TapeNode + const op::OpKey + const depth::Int + """ + Number of times this exact prefix was observed. For the root this is the + number of arrays allocated at the site (bumped by `track!`), which makes it + the correct denominator for `P(next op = k) = child.count / node.count`. + The residual `node.count - sum(children counts)` is the probability that + the chain *stopped* here — an array that is allocated and never used must + reduce our confidence, or a site that is 90% dead weight would look + perfectly predictable. + """ + count::Int + children::Union{Nothing,Dict{OpKey,TapeNode}} + "Most recently observed specs of *all* arguments of this operation." + argspecs::Vector{ArgSpec} +end + +TapeNode(op::OpKey, depth::Int) = TapeNode(op, depth, 0, nothing, ArgSpec[]) + +"Sum of child counts; `count - child_total` is the number of chains that stopped here." +function child_total(n::TapeNode) + kids = n.children + kids === nothing && return 0 + t = 0 + for (_, c) in kids + t += c.count + end + return t +end + +@inline function children!(n::TapeNode) + c = n.children + c === nothing || return c + d = Dict{OpKey,TapeNode}() + n.children = d + return d +end + +nchildren(n::TapeNode) = n.children === nothing ? 0 : length(n.children) + +""" + TapeRoot + +All history for one [`SiteKey`](@ref). +""" +mutable struct TapeRoot + const key::SiteKey + const root::TapeNode + nobservations::Int + nnodes::Int + hits::Int + misses::Int + last_used::Float64 + "Manual override installed by [`pin!`](@ref); bypasses prediction entirely." + pinned::Union{Nothing,LayoutChoice} +end + +TapeRoot(key::SiteKey) = + TapeRoot(key, TapeNode(ROOT_OPKEY, 0), 0, 1, 0, 0, time(), nothing) + +const STORE = Dict{SiteKey,TapeRoot}() +const STORE_LOCK = ReentrantLock() +const TOTAL_NODES = Ref(0) + +function get_root!(key::SiteKey) + lock(STORE_LOCK) do + r = get(STORE, key, nothing) + if r === nothing + maybe_evict!() + r = TapeRoot(key) + STORE[key] = r + TOTAL_NODES[] += 1 + end + r.last_used = time() + return r + end +end + +get_root(key::SiteKey) = lock(STORE_LOCK) do + get(STORE, key, nothing) +end + +""" +Evict least-recently-used sites when over budget. Called with `STORE_LOCK` held. + +TODO(eviction): LRU on sites is crude — a site observed 10000 times and not +touched for a minute is more valuable than one observed twice a second ago. +Weight by `nobservations` and by realised benefit (`hits`), i.e. evict by +`last_used - w*log(1+hits)`. +""" +function maybe_evict!() + (length(STORE) < CONFIG.max_sites && TOTAL_NODES[] < CONFIG.max_nodes) && return nothing + victims = sort!(collect(STORE); by = kv -> kv[2].last_used) + ndrop = max(1, length(victims) ÷ 4) + for i in 1:ndrop + k, r = victims[i] + TOTAL_NODES[] -= r.nnodes + delete!(STORE, k) + end + vlog("evicted $ndrop site(s); $(length(STORE)) remain") + return nothing +end + +# --------------------------------------------------------------------------- +# Live traces: the association between an allocated array and its tape +# --------------------------------------------------------------------------- + +""" + PredictedOp + +One step of a forecast: the operation, the cumulative probability of reaching +it, and a representative snapshot of the arguments it was last seen with. +""" +struct PredictedOp + key::OpKey + prob::Float64 + argspecs::Vector{ArgSpec} +end + +""" + LiveTrace + +Per-array recording state. Attached to a `DArray` in a `WeakKeyDict`, so it +disappears with the array and never extends its lifetime. +""" +mutable struct LiveTrace + const root::TapeRoot + node::TapeNode + nops::Int + truncated::Bool + """ + Spec of the array itself. Held separately from `root.key` because the key + stores a *bucketed* size, which is right for tape lookup but wrong for + costing — a plan computed against `2^round(log2(n))` would be off by up to + 41% in each dimension. + """ + const self::ArgSpec + "Layout actually chosen at allocation time." + const layout::LayoutChoice + "Forecast made at allocation time (for hit/miss accounting and re-planning)." + predicted::Vector{PredictedOp} + "Per-step layouts the planner would like; `plan[1]` is what was committed." + plan::Vector{LayoutChoice} + "How many recorded operations matched the forecast." + matched::Int + diverged::Bool +end + +LiveTrace(root::TapeRoot, self::ArgSpec, layout::LayoutChoice, + predicted::Vector{PredictedOp}, plan::Vector{LayoutChoice}) = + LiveTrace(root, root.root, 0, false, self, layout, predicted, plan, 0, false) + +const TRACES = WeakKeyDict{DArray,LiveTrace}() + +get_trace(A::DArray) = get(TRACES, A, nothing) +get_trace(@nospecialize(_)) = nothing + +# --------------------------------------------------------------------------- +# Recording +# --------------------------------------------------------------------------- + +""" + advance!(trace, opkey, specs) + +Extend `trace` by one observed operation, committing it into the trie +immediately. +""" +function advance!(trace::LiveTrace, opkey::OpKey, specs::Vector{ArgSpec}) + if trace.nops >= CONFIG.max_tape_length + trace.truncated = true + return nothing + end + + # Divergence accounting against the forecast made at allocation time. + idx = trace.nops + 1 + if idx <= length(trace.predicted) + # Unlocked on purpose: these are diagnostic counters on the hot path, + # and a torn increment costs us an inaccurate hit-rate readout, not + # correctness. Everything the planner actually reads is either + # trace-local or guarded by STORE_LOCK. + if trace.predicted[idx].key == opkey + trace.matched += 1 + trace.root.hits += 1 + else + trace.diverged = true + trace.root.misses += 1 + end + end + + node = trace.node + kids = children!(node) + child = get(kids, opkey, nothing) + if child === nothing + child = TapeNode(opkey, node.depth + 1) + kids[opkey] = child + lock(STORE_LOCK) do + trace.root.nnodes += 1 + TOTAL_NODES[] += 1 + end + end + child.count += 1 + child.argspecs = specs + trace.node = child + trace.nops += 1 + return nothing +end + +""" + record_op!(op::Symbol, token::UInt64, args::Tuple) + +Record that operation `op` is about to be applied to `args`. Called by +[`@record_op`](@ref); the macro supplies `token`. + +Every `DArray` in `args` that carries a live trace has the operation appended, +tagged with its own argument position. Arrays with no trace are ignored. + +TODO(joint-planning): this treats each argument's tape independently, but +layout decisions are *joint* — `mul!(C, A, B)` needs the three layouts to be +mutually compatible, not merely individually good. That is the alignment half +of the classical automatic-data-layout problem (Kennedy & Kremer, TOPLAS 1998) +and it is the harder half. The right structure is to record the co-participating +arrays' identities here, union-find them into connected components, and solve +layout over the component rather than per array. Dagger's Datadeps aliasing +machinery (`src/datadeps/aliasing.jl`) already computes most of the required +overlap information. + +TODO(adoption): when a traced array meets an untraced one (e.g. a user-supplied +`DArray` from `distribute`), we currently learn nothing about the untraced one. +It should be adopted into the traced one's component so the planner at least +knows a redistribution may be needed. +""" +function _record_op!(op::Symbol, token::UInt64, args::Tuple) + is_enabled() || return nothing + arity = length(args) + + # Fast path: if nothing here is traced, we only pay for the scan. + any_traced = false + @inbounds for a in args + if a isa DArray && haskey(TRACES, a) + any_traced = true + break + end + end + any_traced || return nothing + + specs = build_specs(args) + @inbounds for i in 1:arity + a = args[i] + a isa DArray || continue + trace = get_trace(a) + trace === nothing && continue + advance!(trace, OpKey(op, i, arity), specs) + CONFIG.allow_repartition && maybe_repartition!(a, trace) + end + return nothing +end + +""" + record_op!(op::Symbol, args...) + +Function form of [`@record_op`](@ref), for callers that build their argument +list dynamically. Carries no lexical token, so in `:context` site-identification +mode it contributes nothing to the context hash; prefer the macro where the +call site is statically known. +""" +record_op!(op::Symbol, args...) = _record_op!(op, UInt64(0), args) + +""" + build_specs(args) -> Vector{ArgSpec} + +Snapshot every `DArray` argument. Non-`DArray` arguments get a degenerate spec +so that positions line up with the caller's argument list. +""" +function build_specs(args::Tuple) + specs = Vector{ArgSpec}(undef, length(args)) + @inbounds for i in eachindex(args) + a = args[i] + specs[i] = a isa DArray ? ArgSpec(a) : ArgSpec(Nothing, (), (), :none) + end + return specs +end + +# --------------------------------------------------------------------------- +# Prediction +# --------------------------------------------------------------------------- + +""" + predict(node::TapeNode; horizon, min_prob) -> Vector{PredictedOp} + +Walk forward from `node`, greedily following the most-taken branch, and stop +when the branch probability or the cumulative probability falls below +threshold, when the chain is more likely than not to have ended, or when +`horizon` steps have been produced. + +This is a first-order Markov / prediction-by-partial-match walk over the trie. +Each returned step carries the *cumulative* probability of reaching it, which +the planner uses directly as the weight on that operation's cost — an operation +we are 40% sure about should contribute 40% of its cost to the objective. +""" +function predict(node::TapeNode; + horizon::Int = CONFIG.horizon, + min_prob::Float64 = CONFIG.min_branch_prob) + out = PredictedOp[] + cur = node + cum = 1.0 + for _ in 1:horizon + kids = cur.children + (kids === nothing || isempty(kids)) && break + + # Denominator is how many times we were *here*, not how many times we + # left: chains that stopped at this node count against the forecast. + total = max(cur.count, child_total(cur)) + total <= 0 && break + + best = nothing + bestcount = 0 + for (_, c) in kids + if c.count > bestcount + bestcount = c.count + best = c + end + end + best === nothing && break + + p = bestcount / total + p < min_prob && break + cum *= p + cum < min_prob && break + + push!(out, PredictedOp(best.op, cum, best.argspecs)) + cur = best + end + return out +end + +""" + confidence(pred) -> Float64 + +Cumulative probability of the first predicted step, i.e. how much we believe +anything at all is about to happen. +""" +confidence(pred::Vector{PredictedOp}) = isempty(pred) ? 0.0 : pred[1].prob diff --git a/src/utils/haloarray.jl b/src/utils/haloarray.jl index a9bb4a0d4..48e4f39c2 100644 --- a/src/utils/haloarray.jl +++ b/src/utils/haloarray.jl @@ -118,6 +118,54 @@ end end end +""" + halo_load(halos, idx, I) + halo_store!(halos, idx, I, value) + +Read or write element `I` of the `idx`-th halo region. + +`halos` is not necessarily homogeneous. GPU backends return a plain device array +for a `view` that happens to be contiguous and a `SubArray` for one that is not, +so a halo tuple built from views of neighboring chunks can hold both. Indexing +such a tuple with a runtime `idx` infers to a union wide enough that the load +that follows becomes a dynamic dispatch, which GPU backends cannot compile at +all. Selecting the region with an unrolled comparison chain instead gives every +branch one concrete array type. + +Homogeneous tuples -- what the CPU path always produces -- keep the plain +indexed load, so they pay nothing for this. +""" +@generated function halo_load(halos::Tuple, idx::Int, I::NTuple{N,Int}) where N + Ts = fieldtypes(halos) + body = if allequal(Ts) + :(@inbounds halos[idx][I...]) + else + # Last region is the fallback: `region_index` never returns out of range, + # and a trailing `throw` would only add unreachable code the GPU backend + # still has to compile. + ex = :(@inbounds halos[$(length(Ts))][I...]) + for i in (length(Ts) - 1):-1:1 + ex = :(idx === $i ? @inbounds(halos[$i][I...]) : $ex) + end + ex + end + return Expr(:block, Expr(:meta, :inline), body) +end + +@generated function halo_store!(halos::Tuple, idx::Int, I::NTuple{N,Int}, value) where N + Ts = fieldtypes(halos) + body = if allequal(Ts) + :(@inbounds halos[idx][I...] = value) + else + ex = :(@inbounds halos[$(length(Ts))][I...] = value) + for i in (length(Ts) - 1):-1:1 + ex = :(idx === $i ? @inbounds(halos[$i][I...] = value) : $ex) + end + ex + end + return Expr(:block, Expr(:meta, :inline), body) +end + # Define getindex for HaloArray @inline function Base.getindex(tile::HaloArray{T,N}, I::Vararg{Int,N}) where {T,N} Base.@boundscheck checkbounds(tile, I...) @@ -130,7 +178,7 @@ end # Halo region idx = region_index(code) local_idx = compute_local_index(tile, I, code) - return @inbounds tile.halos[idx][local_idx...] + return halo_load(tile.halos, idx, local_idx) end end @@ -146,18 +194,83 @@ end # Halo region idx = region_index(code) local_idx = compute_local_index(tile, I, code) - return @inbounds tile.halos[idx][local_idx...] = value + return halo_store!(tile.halos, idx, local_idx, value) end end Base.IndexStyle(::Type{<:HaloArray}) = IndexCartesian() -# GPU-friendly reductions for SubArray views of HaloArray. +############################################################################# +# HaloInterior +############################################################################# + +""" + HaloInterior(parent, halo_width) + +A stand-in for a `HaloArray` used while sweeping the region of a chunk whose +whole neighborhood lies inside the center. Indexing goes straight to `parent`, +skipping the region-code dispatch that `HaloArray` must perform, which is what +lets the interior sweep vectorize. +""" +struct HaloInterior{T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N} + parent::A + halo_width::NTuple{N,Int} +end + +Base.size(A::HaloInterior) = size(A.parent) +Base.axes(A::HaloInterior) = axes(A.parent) +Base.IndexStyle(::Type{<:HaloInterior}) = IndexCartesian() +@inline function Base.getindex(A::HaloInterior{T,N}, I::Vararg{Int,N}) where {T,N} + Base.@boundscheck checkbounds(A.parent, I...) + return @inbounds A.parent[I...] +end +@inline function Base.setindex!(A::HaloInterior{T,N}, value, I::Vararg{Int,N}) where {T,N} + Base.@boundscheck checkbounds(A.parent, I...) + return @inbounds A.parent[I...] = value +end + +############################################################################# +# StencilNeighborhood +############################################################################# + +""" + StencilNeighborhood(parent, center, halo_width) + +The value produced by `@neighbors`: a `(2w+1)^N` window of `parent` centered on +`center`, with 1-based indices. This is a plain offset-index wrapper rather than +a `SubArray` because the latter's index translation blocks vectorization of the +sweep loop (roughly a 5x difference on the interior). +""" +struct StencilNeighborhood{T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N} + parent::A + center::CartesianIndex{N} + halo_width::NTuple{N,Int} +end + +Base.size(nb::StencilNeighborhood{T,N}) where {T,N} = + ntuple(i -> 2 * nb.halo_width[i] + 1, Val(N)) +Base.axes(nb::StencilNeighborhood{T,N}) where {T,N} = + ntuple(i -> Base.OneTo(2 * nb.halo_width[i] + 1), Val(N)) +Base.IndexStyle(::Type{<:StencilNeighborhood}) = IndexCartesian() + +@inline function _neighborhood_index(nb::StencilNeighborhood{T,N}, I::NTuple{N,Int}) where {T,N} + return nb.center + CartesianIndex(ntuple(i -> I[i] - nb.halo_width[i] - 1, Val(N))) +end +@inline function Base.getindex(nb::StencilNeighborhood{T,N}, I::Vararg{Int,N}) where {T,N} + Base.@boundscheck checkbounds(nb, I...) + return @inbounds nb.parent[_neighborhood_index(nb, I)] +end +@inline function Base.setindex!(nb::StencilNeighborhood{T,N}, value, I::Vararg{Int,N}) where {T,N} + Base.@boundscheck checkbounds(nb, I...) + return @inbounds nb.parent[_neighborhood_index(nb, I)] = value +end + +# GPU-friendly reductions over neighborhoods. # The standard reduction path uses _foldl_impl/iterate which produces Union # types that SPIR-V and other GPU compilers can't handle, and passes through # keyword-argument forwarding that triggers dynamic dispatch on GPU. # These overrides use CartesianIndices iteration which compiles cleanly. -@inline function Base.mapreduce(f::F, op::OP, A::SubArray{T,N,<:HaloArray}) where {F,OP,T,N} +@inline function Base.mapreduce(f::F, op::OP, A::StencilNeighborhood{T,N}) where {F,OP,T,N} first_idx = CartesianIndex(ntuple(d -> firstindex(A, d), Val(N))) result = f(@inbounds A[first_idx]) @inbounds for idx in CartesianIndices(A) @@ -166,7 +279,7 @@ Base.IndexStyle(::Type{<:HaloArray}) = IndexCartesian() end return result end -@inline function Base.sum(A::SubArray{T,N,<:HaloArray}) where {T,N} +@inline function Base.sum(A::StencilNeighborhood{T,N}) where {T,N} first_idx = CartesianIndex(ntuple(d -> firstindex(A, d), Val(N))) result = @inbounds A[first_idx] @inbounds for idx in CartesianIndices(A) @@ -175,7 +288,7 @@ end end return result end -@inline function Base.sum(f::F, A::SubArray{T,N,<:HaloArray}) where {F,T,N} +@inline function Base.sum(f::F, A::StencilNeighborhood{T,N}) where {F,T,N} first_idx = CartesianIndex(ntuple(d -> firstindex(A, d), Val(N))) result = f(@inbounds A[first_idx]) @inbounds for idx in CartesianIndices(A) @@ -191,6 +304,9 @@ Adapt.adapt_structure(to, H::Dagger.HaloArray) = H.halo_width; own_center=H.own_center) +Adapt.adapt_structure(to, A::Dagger.HaloInterior) = + HaloInterior(Adapt.adapt(to, A.parent), A.halo_width) + function aliasing(A::HaloArray) return CombinedAliasing([aliasing(A.center), map(aliasing, A.halos)...]) end diff --git a/src/utils/tasks.jl b/src/utils/tasks.jl index ddd8da2ee..beab71cd7 100644 --- a/src/utils/tasks.jl +++ b/src/utils/tasks.jl @@ -1,3 +1,23 @@ +""" + set_task_migratable!(task) -> task + +Allow `task` to run on any thread of the default pool. + +Clearing `sticky` is not sufficient on its own. A hand-built `@task` has no +thread pool assigned, and `Base.enq_work` routes such a task onto the *current* +thread's work queue rather than the shared multi-queue, so it cannot be picked +up until the thread that scheduled it yields. `Threads.@spawn` sets both; task +pools built with `@task` must do the same or every task they spawn is pinned, +in effect, to whichever thread happened to schedule it. +""" +function set_task_migratable!(task::Task) + task.sticky = false + @static if isdefined(Base.Threads, :_spawn_set_thrpool) + Base.Threads._spawn_set_thrpool(task, :default) + end + return task +end + function set_task_tid!(task::Task, tid::Integer) task.sticky = true ctr = 0 diff --git a/test/Project.toml b/test/Project.toml index edc738961..47894c0bf 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c" ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63" Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" @@ -9,15 +10,20 @@ FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" +IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinuxPerf = "b4c46c6c-4fb0-484d-a11a-41bc3392d094" MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94" +Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b" OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" Preferences = "21216c6a-2e73-6563-6e65-726566657250" +PureKLU = "0c0d3e7f-3a8b-4f7e-b6f1-9a4d2e7c1f01" +PureUMFPACK = "b7e1f0a2-3c4d-4e5f-9a0b-1c2d3e4f5a6b" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b" diff --git a/test/array/allocation.jl b/test/array/allocation.jl index 83dba36d2..4548f3ed1 100644 --- a/test/array/allocation.jl +++ b/test/array/allocation.jl @@ -99,7 +99,8 @@ end @test size(Xsp) == dims AT = length(dims) == 2 ? SparseMatrixCSC : SparseVector Ach = fetch(Xsp.chunks[1]) - @test Ach isa AT{T} + @test Ach isa Dagger.DSparseArray{T,length(dims)} + @test Ach.mat isa AT{T} AXsp = collect(Xsp) @test AXsp isa Array{T,length(dims)} @test AXsp == collect(Xsp) diff --git a/test/array/linalg/core.jl b/test/array/linalg/core.jl index 4c66ac34a..5de6fefd1 100644 --- a/test/array/linalg/core.jl +++ b/test/array/linalg/core.jl @@ -23,3 +23,57 @@ end DA = DArray(A) @test isapprox(norm(A), norm(DA)) end + +@testset "BLAS-1 vector ops" begin + @testset "T=$T part=$(part.blocksize)" for T in (Float64, ComplexF64), part in (Blocks(16), Blocks(4)) + n = 16 + x = rand(T, n) + y = rand(T, n) + a = T(2) + b = T(3) + Dx = distribute(x, part) + + # dot (conjugates first arg for complex) + @test dot(Dx, distribute(y, part)) ≈ dot(x, y) + + # axpy!: y += a*x + Dy = distribute(copy(y), part) + axpy!(a, Dx, Dy) + @test collect(Dy) ≈ a .* x .+ y + + # axpby!: y = a*x + b*y + Dy = distribute(copy(y), part) + axpby!(a, Dx, b, Dy) + @test collect(Dy) ≈ a .* x .+ b .* y + + # rmul!/lmul!: scale in place + Dz = distribute(copy(x), part) + rmul!(Dz, a) + @test collect(Dz) ≈ a .* x + Dz = distribute(copy(x), part) + lmul!(a, Dz) + @test collect(Dz) ≈ a .* x + end + + # Operands with *different* partitionings must still work (aligned internally + # via maybe_copy_buffered), and the in-place result keeps its own layout. + @testset "mismatched partitionings T=$T" for T in (Float64, ComplexF64) + n = 16 + x = rand(T, n) + y = rand(T, n) + a, b = T(2), T(3) + Dx = distribute(x, Blocks(4)) + Dy = distribute(y, Blocks(8)) + + @test dot(Dx, Dy) ≈ dot(x, y) + + Dy2 = distribute(copy(y), Blocks(8)) + axpy!(a, Dx, Dy2) + @test collect(Dy2) ≈ a .* x .+ y + @test Dy2.partitioning == Blocks(8) # output layout preserved + + Dy3 = distribute(copy(y), Blocks(8)) + axpby!(a, Dx, b, Dy3) + @test collect(Dy3) ≈ a .* x .+ b .* y + end +end diff --git a/test/array/linalg/iterativesolvers.jl b/test/array/linalg/iterativesolvers.jl new file mode 100644 index 000000000..ca7865f29 --- /dev/null +++ b/test/array/linalg/iterativesolvers.jl @@ -0,0 +1,206 @@ +# Distributed iterative (Krylov) linear-solver tests. +# +# Exercises the matrix-free Krylov integration (`Dagger.cg`/`minres`/`gmres`/ +# `bicgstab` + the generic `krylov_solve`) over both dense and sparse-backed +# `DMatrix` operators, plus the Jacobi preconditioner. Reference solutions come +# from a dense direct solve. +# +# julia test/runtests.jl --test array/linalg/iterativesolvers + +using Krylov +using AlgebraicMultigrid +using IncompleteLU + +# Strongly diagonally-dominant tridiagonal SPD matrix. The large diagonal keeps +# the condition number small so the Krylov methods converge in a handful of +# iterations (keeping the distributed test fast). `inv(diag) == 1/4`. +const SPD_DIAG = 4.0 + +laplacian_1d(T, n) = SparseArrays.spdiagm( + -1 => fill(-one(T), n - 1), + 0 => fill(T(SPD_DIAG), n), + 1 => fill(-one(T), n - 1), +) + +# Add a first-order advection term -> nonsymmetric, still well-conditioned. +function advection_diffusion_1d(T, n) + return laplacian_1d(T, n) + SparseArrays.spdiagm( + -1 => fill(T(-3) / 10, n - 1), + 1 => fill(T(3) / 10, n - 1), + ) +end + +@testset "Iterative solvers (Krylov)" begin + n = 64 + k = 16 + Db_part = Blocks(k) + A_part = Blocks(k, k) + + @testset "SPD operator ($(backend))" for backend in (:dense, :sparse) + Asp = laplacian_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + @testset "$(nameof(solver))" for solver in (Dagger.cg, Dagger.minres, Dagger.gmres, Dagger.bicgstab) + x, stats = solver(DA, Db; atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test x isa Dagger.DVector + @test collect(x) ≈ xref rtol = 1e-6 + end + + # Generic entry point. + x, stats = Dagger.krylov_solve(:cg, DA, Db; atol = 1e-12, rtol = 1e-10) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + @testset "nonsymmetric operator ($(backend))" for backend in (:dense, :sparse) + Asp = advection_diffusion_1d(Float64, n) + b = rand(n) + xref = Matrix(Asp) \ b + + DA = backend === :dense ? distribute(Matrix(Asp), A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + @testset "$(nameof(solver))" for solver in (Dagger.gmres, Dagger.bicgstab) + x, stats = solver(DA, Db; atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + end + + @testset "complex SPD (Hermitian) operator" begin + # Real SPD tridiagonal is Hermitian as a complex matrix. + Asp = SparseArrays.spdiagm( + -1 => fill(ComplexF64(-1), n - 1), + 0 => fill(ComplexF64(SPD_DIAG), n), + 1 => fill(ComplexF64(-1), n - 1), + ) + b = rand(ComplexF64, n) + xref = Matrix(Asp) \ b + DA = distribute(Asp, A_part) + Db = distribute(b, Db_part) + x, stats = Dagger.cg(DA, Db; atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + @testset "Jacobi preconditioner" begin + Asp = laplacian_1d(Float64, n) + b = rand(n) + xref = Matrix(Asp) \ b + + @testset "build + apply ($(backend))" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Matrix(Asp), A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + P = Dagger.JacobiPreconditioner(DA) + @test collect(P.dinv) ≈ fill(1 / SPD_DIAG, n) # 1/diag + + # Apply: y = M⁻¹ x = dinv .* x. + y = similar(Db) + mul!(y, P, Db) + @test collect(y) ≈ (1 / SPD_DIAG) .* b + + x, stats = Dagger.cg(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + # Non-square block grid must be rejected with a helpful error. + DA_ragged = distribute(Matrix(Asp), Blocks(k, k ÷ 2)) + @test_throws ArgumentError Dagger.JacobiPreconditioner(DA_ragged) + end + + @testset "block-Jacobi preconditioner" begin + Asp = laplacian_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + # Reference: apply the exact block-diagonal inverse. + yref = similar(b) + for s in 1:k:n + r = s:min(s + k - 1, n) + yref[r] = Adense[r, r] \ b[r] + end + + @testset "build + apply ($(backend))" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + P = Dagger.BlockJacobiPreconditioner(DA) + y = similar(Db) + mul!(y, P, Db) + @test collect(y) ≈ yref + + x, stats = Dagger.cg(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + # A single tile makes block-Jacobi an *exact* solve, so PCG converges + # essentially immediately. + DA1 = distribute(Adense, Blocks(n, n)) + Db1 = distribute(b, Blocks(n)) + P1 = Dagger.BlockJacobiPreconditioner(DA1) + x1, s1 = Dagger.cg(DA1, Db1; M = P1, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test s1.solved + @test s1.niter <= 2 + @test collect(x1) ≈ xref rtol = 1e-8 + + @test_throws ArgumentError Dagger.BlockJacobiPreconditioner(distribute(Adense, Blocks(k, k ÷ 2))) + end + + @testset "AMG preconditioner ($(method))" for method in (:ruge_stuben, :smoothed_aggregation) + Asp = laplacian_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + @testset "$(backend)" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + P = Dagger.AMGPreconditioner(DA; method = method) + # Apply is a V-cycle approximating M⁻¹ x; just check it runs + is finite + # (and repeatable, exercising the cached, pinned hierarchy). + y1 = similar(Db); mul!(y1, P, Db) + y2 = similar(Db); mul!(y2, P, Db) + @test all(isfinite, collect(y1)) + @test collect(y1) ≈ collect(y2) + + x, stats = Dagger.cg(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + end + + @testset "block-ILU preconditioner" begin + Asp = laplacian_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + @testset "build + apply ($(backend))" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + P = Dagger.BlockILUPreconditioner(DA; τ = 0.01) + y1 = similar(Db); mul!(y1, P, Db) + y2 = similar(Db); mul!(y2, P, Db) + @test all(isfinite, collect(y1)) + @test collect(y1) ≈ collect(y2) + + x, stats = Dagger.cg(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + @test_throws ArgumentError Dagger.BlockILUPreconditioner(distribute(Adense, Blocks(k, k ÷ 2)); τ = 0.01) + end +end diff --git a/test/array/linalg/matmul.jl b/test/array/linalg/matmul.jl index 26fd0afe3..f0ec73287 100644 --- a/test/array/linalg/matmul.jl +++ b/test/array/linalg/matmul.jl @@ -29,36 +29,71 @@ function test_gemm!(T, szA, szB, partA, partB) DA = distribute(A, partA) DB = distribute(B, partB) + SA = sprand(T, szA..., 0.1) + SB = sprand(T, szA..., 0.1) + + DSA = distribute(SA, partA) + DSB = distribute(SB, partB) + ## Out-of-place gemm # No transA, No transB + # Dense DC = DA * DB C = A * B @test collect(DC) ≈ C + # Sparse + DSC = DSA * DSB + SC = SA * SB + @test collect(DSC) ≈ SC if szA == szB # No transA, transB + # Dense DC = DA * DB' C = A * B' @test collect(DC) ≈ C + # Sparse + DSC = DSA * DSB' + SC = SA * SB' + @test collect(DSC) ≈ SC # transA, No transB + # Dense DC = DA' * DB C = A' * B @test collect(DC) ≈ C + # Sparse + DSC = DSA' * DSB + SC = SA' * SB + @test collect(DSC) ≈ SC end # transA, transB + # Dense DC = DA' * DB' C = A' * B' @test collect(DC) ≈ C + #= Sparse + DSC = DSA' * DSB' + SC = SA' * SB' + @test collect(DSC) ≈ SC + =# ## In-place gemm # No transA, No transB + # Dense C = zeros(T, szC...) DC = distribute(C, partC) mul!(C, A, B) mul!(DC, DA, DB) @test collect(DC) ≈ C + #= Sparse + SC = zeros(T, szC...) + DSC = distribute(SC, partC) + mul!(SC, SA, SB) + mul!(DSC, DSA, DSB) + @test collect(DSC) ≈ SC + =# if szA == szB # No transA, transB diff --git a/test/array/linalg/matmul_finch.jl b/test/array/linalg/matmul_finch.jl new file mode 100644 index 000000000..061874f1e --- /dev/null +++ b/test/array/linalg/matmul_finch.jl @@ -0,0 +1,105 @@ +# Fast, focused sparse matrix-multiply tests for the Finch backend. +# +# Mirrors `matmul_sparse.jl` (SparseArrays) but exercises Finch tensor tiles via +# `DSparseMatrix`. Finch is an optional (weak) dependency, so this test loads it +# on all workers and skips gracefully if it cannot be loaded. +# +# Run with: +# +# julia test/runtests.jl --test array/linalg/matmul_finch +# +# N.B. Finch precompilation is heavy (several minutes); this test is intended for +# local validation rather than fast CI. + +# Make the test-project environment (which provides Finch) resolvable on the +# workers, which are launched with Dagger's project. +const _FINCH_TEST_ENV = abspath(joinpath(@__DIR__, "..", "..")) +@everywhere pushfirst!(LOAD_PATH, $_FINCH_TEST_ENV) + +const FINCH_AVAILABLE = try + @eval using Finch + @everywhere using Finch + true +catch err + @warn "Finch unavailable; skipping Finch matmul tests" exception=err + false +end + +if !FINCH_AVAILABLE + @testset "Finch GEMM (quick)" begin + @test_skip false + end +else + const FINCH_DENSITY = 0.3 + + # Self-consistent: build distributed Finch operands, take the dense + # `collect` as the reference, and compare against the distributed result. + function test_finch_gemm!(T, sz, partA, partB) + partC = Blocks(partA.blocksize[1], partB.blocksize[2]) + + DSA = Finch.fsprand(partA, T, sz, FINCH_DENSITY) + DSB = Finch.fsprand(partB, T, sz, FINCH_DENSITY) + A = collect(DSA) + B = collect(DSB) + + ## Out-of-place + @test collect(DSA * DSB) ≈ A * B # N / N + @test collect(DSA * DSB') ≈ A * B' # N / C + @test collect(DSA' * DSB) ≈ A' * B # C / N + @test collect(DSA' * DSB') ≈ A' * B' # C / C + + ## Symmetric rank-k (syrk path: A === B with one operand transposed) + @test collect(DSA' * DSA) ≈ A' * A + @test collect(DSA * DSA') ≈ A * A' + + ## In-place + DSC = Finch.fspzeros(partC, T, sz) + mul!(DSC, DSA, DSB) + @test collect(DSC) ≈ A * B + end + + const FINCH_QUICK_CASES = [ + ((8, 8), Blocks(4, 4), Blocks(4, 4)), + ((8, 8), Blocks(2, 4), Blocks(4, 2)), + ] + + @testset "Finch GEMM (quick)" begin + @testset "size=$sz part=$(partA.blocksize)/$(partB.blocksize)" for (sz, partA, partB) in FINCH_QUICK_CASES + @testset "T=$T" for T in (Float64, ComplexF64) + test_finch_gemm!(T, sz, partA, partB) + end + end + end + + # Sparse matrix-vector multiply (SpMV): Finch sparse matrix tiles, dense vectors. + function test_finch_spmv!(T, n, part) + bs = part.blocksize[1] + DSA = Finch.fsprand(part, T, (n, n), FINCH_DENSITY) + A = collect(DSA) + x = rand(T, n) + Dx = distribute(x, Blocks(bs)) + + @test collect(DSA * Dx) ≈ A * x + @test collect(transpose(DSA) * Dx) ≈ transpose(A) * x + @test collect(DSA' * Dx) ≈ A' * x + + y = rand(T, n) + Dy = distribute(copy(y), Blocks(bs)) + alpha, beta = T(2), T(3) + mul!(Dy, DSA, Dx, alpha, beta) + @test collect(Dy) ≈ alpha * (A * x) + beta * y + end + + const FINCH_SPMV_CASES = [ + (8, Blocks(4, 4)), + (8, Blocks(2, 2)), + ] + + @testset "Finch SpMV (quick)" begin + @testset "n=$n part=$(part.blocksize)" for (n, part) in FINCH_SPMV_CASES + @testset "T=$T" for T in (Float64, ComplexF64) + test_finch_spmv!(T, n, part) + end + end + end +end diff --git a/test/array/linalg/matmul_sparse.jl b/test/array/linalg/matmul_sparse.jl new file mode 100644 index 000000000..cc2e8274e --- /dev/null +++ b/test/array/linalg/matmul_sparse.jl @@ -0,0 +1,139 @@ +# Fast, focused sparse matrix-multiply tests. +# +# This is intentionally a *small* subset of the full GEMM suite in `matmul.jl`, +# meant for quick iteration while developing sparse-array support (the full +# suite takes ~1hr). Run with: +# +# julia test/runtests.jl --test array/linalg/matmul_sparse +# +# Keep this fast: only a couple of sizes/partitionings, just enough to exercise +# the out-of-place and in-place paths (incl. transposes) for both real and +# complex element types. + +const SPARSE_DENSITY = 0.3 + +function test_sparse_gemm!(T, sz, partA, partB) + rows, cols = sz + @assert rows == cols "sparse quick-test uses square matrices so transposes line up" + partC = Blocks(partA.blocksize[1], partB.blocksize[2]) + + SA = sprand(T, sz..., SPARSE_DENSITY) + SB = sprand(T, sz..., SPARSE_DENSITY) + + DSA = distribute(SA, partA) + DSB = distribute(SB, partB) + + ## Out-of-place + @test collect(DSA * DSB) ≈ SA * SB # N / N + @test collect(DSA * DSB') ≈ SA * SB' # N / T + @test collect(DSA' * DSB) ≈ SA' * SB # T / N + @test collect(DSA' * DSB') ≈ SA' * SB' # T / T + + ## Symmetric rank-k (syrk path: A === B with one operand transposed) + @test collect(DSA' * DSA) ≈ Array(SA)' * Array(SA) + @test collect(DSA * DSA') ≈ Array(SA) * Array(SA)' + + ## In-place + SC = SA * SB + DSC = distribute(sparse(zeros(T, sz...)), partC) + mul!(DSC, DSA, DSB) + @test collect(DSC) ≈ SC +end + +const SPARSE_QUICK_CASES = [ + ((8, 8), Blocks(4, 4), Blocks(4, 4)), + ((8, 8), Blocks(2, 4), Blocks(4, 2)), + ((8, 8), Blocks(4, 2), Blocks(2, 4)), +] + +@testset "Sparse GEMM (quick)" begin + @testset "size=$sz part=$(partA.blocksize)/$(partB.blocksize)" for (sz, partA, partB) in SPARSE_QUICK_CASES + @testset "T=$T" for T in (Float64, ComplexF64) + test_sparse_gemm!(T, sz, partA, partB) + end + end +end + +# Sparse matrix-vector multiply (SpMV): sparse matrix tiles, dense vectors. This +# is the core kernel for distributed iterative (Krylov) solvers. +function test_sparse_spmv!(T, n, part) + bs = part.blocksize[1] + SA = sprand(T, n, n, SPARSE_DENSITY) + x = rand(T, n) + + DSA = distribute(SA, part) + Dx = distribute(x, Blocks(bs)) + + ## Out-of-place (N / T / C) + @test collect(DSA * Dx) ≈ SA * x + @test collect(transpose(DSA) * Dx) ≈ transpose(SA) * x + @test collect(DSA' * Dx) ≈ SA' * x + + ## In-place 5-arg mul!: y = alpha*A*x + beta*y + y = rand(T, n) + Dy = distribute(copy(y), Blocks(bs)) + alpha, beta = T(2), T(3) + mul!(Dy, DSA, Dx, alpha, beta) + @test collect(Dy) ≈ alpha * (SA * x) + beta * y +end + +const SPARSE_SPMV_CASES = [ + (8, Blocks(4, 4)), + (8, Blocks(2, 2)), +] + +@testset "Sparse SpMV (quick)" begin + @testset "n=$n part=$(part.blocksize)" for (n, part) in SPARSE_SPMV_CASES + @testset "T=$T" for T in (Float64, ComplexF64) + test_sparse_spmv!(T, n, part) + end + end +end + +# Any *partial* or *reinterpreted* access to a sparse container (views, +# transposes, adjoints, reshapes, and combinations thereof) must alias the +# *entire* container, so that Datadeps never tracks stale sub-spans of storage +# that may have been reallocated on write. This is enforced by `aliasing_root`, +# which resolves any such wrapper back to the container itself (via `Base.parent`) +# -- exposed to Datadeps through the fused `aliasing_unwrapped` -- plus a trapping +# `Base.pointer` that fires if a wrapper ever slips through to the strided path. +@testset "Sparse whole-container aliasing" begin + M = Dagger.DSparseMatrix{Float64}(sprand(Float64, 6, 6, 0.4)) + + @test Dagger.aliases_as_whole(M) + @test !Dagger.aliases_as_whole(rand(Float64, 6, 6)) + # The container's aliasing is a bare `ObjectAliasing`, which also reports + # whole-object (drives the Datadeps whole-object copy short-circuit). + @test Dagger.aliasing(M) isa Dagger.ObjectAliasing + @test Dagger.aliases_as_whole(Dagger.aliasing(M)) + @test !Dagger.aliases_as_whole(Dagger.aliasing(rand(Float64, 6, 6))) + + # `aliasing_root` peels any array wrapper back to the sparse container. + @test Dagger.aliasing_root(M) === M + @test Dagger.aliasing_root(view(M, 1:3, :)) === M + @test Dagger.aliasing_root(view(M, 4:6, 2:4)) === M + @test Dagger.aliasing_root(transpose(M)) === M + @test Dagger.aliasing_root(M') === M + @test Dagger.aliasing_root(reshape(M, 36)) === M + @test Dagger.aliasing_root(view(transpose(M), 1:2, :)) === M + + # After unwrapping, all wrappers alias the same whole container. + a_full = Dagger.aliasing(M) + @test Dagger.will_alias(a_full, Dagger.aliasing_unwrapped(view(M, 1:3, :))) + @test Dagger.will_alias(Dagger.aliasing_unwrapped(view(M, 1:3, :)), + Dagger.aliasing_unwrapped(view(M, 4:6, :))) + + # Distinct containers must not alias. + M2 = Dagger.DSparseMatrix{Float64}(sprand(Float64, 6, 6, 0.4)) + @test !Dagger.will_alias(a_full, Dagger.aliasing(M2)) + + # Dense arrays are untouched: not whole-object, views stay strided, and + # `aliasing_root` returns them unchanged. + A = rand(Float64, 6, 6) + Av = view(A, 1:3, :) + @test Dagger.aliasing_root(Av) === Av + @test Dagger.aliasing_unwrapped(Av) isa Dagger.StridedAliasing + + # The `pointer` trap fires if a sparse container is treated as raw memory. + @test_throws ArgumentError pointer(M) +end diff --git a/test/array/linalg/sparsedirect.jl b/test/array/linalg/sparsedirect.jl new file mode 100644 index 000000000..e2f94ad50 --- /dev/null +++ b/test/array/linalg/sparsedirect.jl @@ -0,0 +1,249 @@ +# Sparse direct-solver tests (KLU via PureKLU, UMFPACK-style via PureUMFPACK). +# +# Exercises the whole-matrix direct solves (`Dagger.klu` / `Dagger.splu`) and the +# per-tile block direct preconditioners (`BlockKLUPreconditioner` / +# `BlockUMFPACKPreconditioner`). Reference solutions come from a dense direct +# solve. +# +# julia test/runtests.jl --test array/linalg/sparsedirect + +using PureKLU +using PureUMFPACK +using Metis +using Krylov + +# A nonsymmetric, well-conditioned sparse system (advection-diffusion-like). +function nonsym_1d(T, n) + return SparseArrays.spdiagm( + -1 => fill(-one(T), n - 1), + 0 => fill(T(3), n), + 1 => fill(T(-7) / 10, n - 1), + ) +end + +# SPD, diagonally dominant tridiagonal (stable under pivoting). +function spd_1d(T, n) + return SparseArrays.spdiagm( + -1 => fill(-one(T), n - 1), + 0 => fill(T(4), n), + 1 => fill(-one(T), n - 1), + ) +end + +# 2D 5-point Laplacian on an N×N grid (SPD; clean METIS separators). +function lap2d(T, N) + n = N * N + Is = Int[]; Js = Int[]; Vs = T[] + idx(i, j) = (j - 1) * N + i + for j in 1:N, i in 1:N + k = idx(i, j) + push!(Is, k); push!(Js, k); push!(Vs, T(4)) + if i > 1; push!(Is, k); push!(Js, idx(i - 1, j)); push!(Vs, -one(T)); end + if i < N; push!(Is, k); push!(Js, idx(i + 1, j)); push!(Vs, -one(T)); end + if j > 1; push!(Is, k); push!(Js, idx(i, j - 1)); push!(Vs, -one(T)); end + if j < N; push!(Is, k); push!(Js, idx(i, j + 1)); push!(Vs, -one(T)); end + end + return SparseArrays.sparse(Is, Js, Vs, n, n) +end + +# Nonsymmetric, diagonally dominant 2D advection-diffusion stencil. +function advdiff2d(T, N) + n = N * N + Is = Int[]; Js = Int[]; Vs = T[] + idx(i, j) = (j - 1) * N + i + for j in 1:N, i in 1:N + k = idx(i, j) + push!(Is, k); push!(Js, k); push!(Vs, T(5)) + if i > 1; push!(Is, k); push!(Js, idx(i - 1, j)); push!(Vs, T(-12) / 10); end + if i < N; push!(Is, k); push!(Js, idx(i + 1, j)); push!(Vs, T(-8) / 10); end + if j > 1; push!(Is, k); push!(Js, idx(i, j - 1)); push!(Vs, T(-11) / 10); end + if j < N; push!(Is, k); push!(Js, idx(i, j + 1)); push!(Vs, T(-9) / 10); end + end + return SparseArrays.sparse(Is, Js, Vs, n, n) +end + +@testset "Sparse direct solvers" begin + n = 64 + k = 16 + Db_part = Blocks(k) + A_part = Blocks(k, k) + + Asp = nonsym_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + @testset "whole-matrix direct ($(solver)) — $(backend)" for solver in (:klu, :splu), + backend in (:dense, :sparse, :singletile) + DA = if backend === :dense + distribute(Adense, A_part) + elseif backend === :sparse + distribute(Asp, A_part) + else + distribute(Asp, Blocks(n, n)) # single tile + end + Db = distribute(b, Db_part) + + F = solver === :klu ? Dagger.klu(DA) : Dagger.splu(DA) + @test size(F) == (n, n) + + # `F \ b` returns a DVector partitioned like `b`. + x = F \ Db + @test x isa Dagger.DVector + @test collect(x) ≈ xref + + # In-place `ldiv!` writes into a preallocated DVector. + y = similar(Db) + LinearAlgebra.ldiv!(y, F, Db) + @test collect(y) ≈ xref + + # Dimension mismatch is caught. + @test_throws DimensionMismatch F \ distribute(rand(n ÷ 2), Blocks(k)) + end + + @testset "block direct preconditioner ($(name))" for (name, build) in ( + ("KLU", Dagger.BlockKLUPreconditioner), + ("UMFPACK", Dagger.BlockUMFPACKPreconditioner), + ) + @testset "build + apply ($(backend))" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + # Reference: apply the exact block-diagonal inverse. + yref = similar(b) + for s in 1:k:n + r = s:min(s + k - 1, n) + yref[r] = Adense[r, r] \ b[r] + end + + P = build(DA) + y = similar(Db) + mul!(y, P, Db) + @test collect(y) ≈ yref + # Repeatable (exercises the cached, pinned factorization). + y2 = similar(Db) + mul!(y2, P, Db) + @test collect(y2) ≈ yref + + x, stats = Dagger.bicgstab(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + # A single tile makes the block solve an *exact* whole-matrix solve. + DA1 = distribute(Asp, Blocks(n, n)) + Db1 = distribute(b, Blocks(n)) + P1 = build(DA1) + y1 = similar(Db1) + mul!(y1, P1, Db1) + @test collect(y1) ≈ xref + + # Non-square block grid must be rejected with a helpful error. + @test_throws ArgumentError build(distribute(Adense, Blocks(k, k ÷ 2))) + end + + @testset "distributed triangular solve (splu)" begin + systems = ( + ("nonsym", nonsym_1d(Float64, n)), + ("spd", spd_1d(Float64, n)), + ) + # Include a case where A's tiling differs from the solve blocksize. + configs = ( + (Blocks(k, k), nothing), # default: match A's blocks + (Blocks(k, k), 8), # finer solve tiles than A + (Blocks(32, 16), 16), # non-square A tiling → square solve + (Blocks(n, n), 16), # single-tile A, multi-tile solve + ) + for (name, Asp_sys) in systems, (A_part_cfg, solve_bs) in configs + @testset "$name part=$(A_part_cfg.blocksize) bs=$(solve_bs)" begin + Adense_sys = Matrix(Asp_sys) + DA = distribute(Asp_sys, A_part_cfg) + Db = distribute(b, Db_part) + + F4a = Dagger.splu(DA) + F4b = if solve_bs === nothing + Dagger.splu(DA; distributed=true) + else + Dagger.splu(DA; distributed=true, blocksize=solve_bs) + end + @test F4b isa Dagger.DistributedSparseLU + @test size(F4b) == (n, n) + @test F4b.L isa Dagger.DMatrix + @test F4b.U isa Dagger.DMatrix + # Factors stay sparse (no densification). + @test fetch(F4b.L.chunks[1, 1]) isa Dagger.DSparseArray + @test fetch(F4b.U.chunks[1, 1]) isa Dagger.DSparseArray + + x4a = collect(F4a \ Db) + x = F4b \ Db + @test x isa Dagger.DVector + @test x.partitioning == Db.partitioning + xc = collect(x) + @test xc ≈ x4a rtol=1e-12 atol=1e-14 + @test norm(Adense_sys * xc - b) / norm(b) < 1e-10 + + # Multi-RHS / repeated solves against one factor. + for _ in 1:3 + bi = rand(n) + Dbi = distribute(bi, Db_part) + xi = collect(F4b \ Dbi) + @test xi ≈ collect(F4a \ Dbi) rtol=1e-12 atol=1e-14 + @test norm(Adense_sys * xi - bi) / norm(bi) < 1e-10 + end + + y = similar(Db) + LinearAlgebra.ldiv!(y, F4b, Db) + @test collect(y) ≈ x4a rtol=1e-12 atol=1e-14 + + @test_throws DimensionMismatch F4b \ distribute(rand(n ÷ 2), Blocks(k)) + end + end + end + + @testset "Schur distributed LU (4c)" begin + systems = ( + ("lap2d", N -> lap2d(Float64, N)), + ("advdiff2d", N -> advdiff2d(Float64, N)), + ) + for (name, buildA) in systems, N in (10, 20), nparts in (2, 4) + @testset "$name N=$N nparts=$nparts" begin + Asp_sys = buildA(N) + nsys = size(Asp_sys, 1) + bs = max(8, nsys ÷ 4) + DA = distribute(Asp_sys, Blocks(bs, bs)) + bsys = rand(nsys) + Db = distribute(bsys, Blocks(bs)) + + F4a = Dagger.splu(DA) + F4c = Dagger.splu(DA; distributed=true, method=:schur, nparts=nparts) + @test F4c isa Dagger.DistributedSchurLU + @test size(F4c) == (nsys, nsys) + + x = F4c \ Db + @test x isa Dagger.DVector + @test x.partitioning == Db.partitioning + xc = collect(x) + x4a = collect(F4a \ Db) + @test xc ≈ x4a rtol=1e-8 + @test norm(Asp_sys * xc - bsys) / norm(bsys) < 1e-9 + + y = similar(Db) + LinearAlgebra.ldiv!(y, F4c, Db) + @test collect(y) ≈ x4a rtol=1e-8 + @test norm(Asp_sys * collect(y) - bsys) / norm(bsys) < 1e-9 + + # Factor-once / solve-many. + for _ in 1:3 + bi = rand(nsys) + Dbi = distribute(bi, Blocks(bs)) + xi = collect(F4c \ Dbi) + @test xi ≈ collect(F4a \ Dbi) rtol=1e-8 + @test norm(Asp_sys * xi - bi) / norm(bi) < 1e-9 + end + + @test_throws DimensionMismatch F4c \ distribute(rand(nsys ÷ 2), Blocks(bs)) + end + end + end +end + diff --git a/test/cudaenv/Project.toml b/test/cudaenv/Project.toml index 6e245a12a..38ea2a004 100644 --- a/test/cudaenv/Project.toml +++ b/test/cudaenv/Project.toml @@ -4,6 +4,7 @@ Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources] diff --git a/test/datadeps.jl b/test/datadeps.jl index 1ad135cac..d0f7b4e67 100644 --- a/test/datadeps.jl +++ b/test/datadeps.jl @@ -998,7 +998,7 @@ end struct DummyErrorScheduler <: Dagger.DataDepsScheduler end struct DummySchedulerError <: Exception end -function Dagger.datadeps_schedule_task(::DummyErrorScheduler, state, all_procs, all_scope, task_scope, spec, task) +function Dagger.datadeps_schedule_task_jit!(::DummyErrorScheduler, all_procs, all_scope, task_scope, spec, task) throw(DummySchedulerError()) end diff --git a/test/datadeps/scheduling.jl b/test/datadeps/scheduling.jl new file mode 100644 index 000000000..2c2003640 --- /dev/null +++ b/test/datadeps/scheduling.jl @@ -0,0 +1,584 @@ +import Dagger: DAGSpec, DatadepsArgSpec, dag_add_task!, + equivalent_structure, + datadeps_dag_equivalent, datadeps_argspec_equivalent, + datadeps_ainfo_equivalent, datadeps_schedule_cache, + LayeredScheduler, RoundRobinScheduler, DataDepsScheduler, + NoAliasing, UnknownAliasing, ContiguousAliasing, + StridedAliasing, TriangularAliasing, DiagonalAliasing, + ObjectAliasing, CombinedAliasing, AliasingWrapper, + DATADEPS_SCHEDULER, In, Out, InOut, Deps +import Dagger +using LinearAlgebra +using Test + +# ---------- equivalent_structure unit tests ---------- + +@testset "equivalent_structure" begin + @testset "NoAliasing/UnknownAliasing" begin + @test equivalent_structure(NoAliasing(), NoAliasing()) + @test equivalent_structure(UnknownAliasing(), UnknownAliasing()) + @test !equivalent_structure(NoAliasing(), UnknownAliasing()) + @test !equivalent_structure(UnknownAliasing(), NoAliasing()) + end + + @testset "ContiguousAliasing" begin + # Same shape, different allocation → equivalent + @test equivalent_structure(Dagger.aliasing(rand(10, 10)), + Dagger.aliasing(rand(10, 10))) + @test equivalent_structure(Dagger.aliasing(zeros(Float64, 32)), + Dagger.aliasing(rand(Float64, 32))) + # Different element counts → not equivalent + @test !equivalent_structure(Dagger.aliasing(rand(10, 10)), + Dagger.aliasing(rand(5, 5))) + # Different element types → not equivalent (S/T params differ) + @test !equivalent_structure(Dagger.aliasing(rand(Float64, 10)), + Dagger.aliasing(rand(Float32, 10))) + end + + @testset "StridedAliasing" begin + A1 = rand(10, 10) + A2 = rand(10, 10) # same parent shape, different allocation + # Same view shape & relative offset → equivalent + @test equivalent_structure(Dagger.aliasing(view(A1, 2:5, 3:6)), + Dagger.aliasing(view(A2, 2:5, 3:6))) + # Same view shape, different relative offset → not equivalent + @test !equivalent_structure(Dagger.aliasing(view(A1, 2:5, 3:6)), + Dagger.aliasing(view(A1, 3:6, 4:7))) + # Different parent shapes → strides differ → not equivalent + B = rand(20, 20) + @test !equivalent_structure(Dagger.aliasing(view(A1, 2:5, 3:6)), + Dagger.aliasing(view(B, 2:5, 3:6))) + # Different view dim lengths → not equivalent + @test !equivalent_structure(Dagger.aliasing(view(A1, 2:5, 3:6)), + Dagger.aliasing(view(A1, 2:6, 3:6))) + end + + @testset "TriangularAliasing" begin + A1 = rand(10, 10); A2 = rand(10, 10) + @test equivalent_structure(Dagger.aliasing(UpperTriangular(A1)), + Dagger.aliasing(UpperTriangular(A2))) + @test equivalent_structure(Dagger.aliasing(LowerTriangular(A1)), + Dagger.aliasing(LowerTriangular(A2))) + # Upper vs Lower → not equivalent + @test !equivalent_structure(Dagger.aliasing(UpperTriangular(A1)), + Dagger.aliasing(LowerTriangular(A1))) + # Unit vs non-unit → not equivalent + @test !equivalent_structure(Dagger.aliasing(UpperTriangular(A1)), + Dagger.aliasing(UnitUpperTriangular(A1))) + # Different size → not equivalent + @test !equivalent_structure(Dagger.aliasing(UpperTriangular(A1)), + Dagger.aliasing(UpperTriangular(rand(5, 5)))) + end + + @testset "DiagonalAliasing" begin + A1 = rand(10, 10); A2 = rand(10, 10) + @test equivalent_structure(Dagger.aliasing(A1, Diagonal), + Dagger.aliasing(A2, Diagonal)) + @test !equivalent_structure(Dagger.aliasing(A1, Diagonal), + Dagger.aliasing(rand(5, 5), Diagonal)) + end + + @testset "ObjectAliasing" begin + r1 = Ref(rand(10)) + r2 = Ref(rand(10)) # same Ref{Vector{Float64}} shape + a1 = first(Dagger.aliasing(r1).sub_ainfos)::ObjectAliasing + a2 = first(Dagger.aliasing(r2).sub_ainfos)::ObjectAliasing + @test equivalent_structure(a1, a2) + end + + @testset "CombinedAliasing (Ref)" begin + @test equivalent_structure(Dagger.aliasing(Ref(rand(8))), + Dagger.aliasing(Ref(rand(8)))) + @test !equivalent_structure(Dagger.aliasing(Ref(rand(8))), + Dagger.aliasing(Ref(rand(4)))) + end + + @testset "AliasingWrapper delegation" begin + A1 = rand(10, 10); A2 = rand(10, 10) + w1 = AliasingWrapper(Dagger.aliasing(A1)) + w2 = AliasingWrapper(Dagger.aliasing(A2)) + @test equivalent_structure(w1, w2) + # Wrapper vs bare also works + @test equivalent_structure(w1, Dagger.aliasing(A2)) + @test equivalent_structure(Dagger.aliasing(A1), w2) + end + + @testset "Cross-type returns false" begin + @test !equivalent_structure(Dagger.aliasing(rand(10)), + Dagger.aliasing(view(rand(10, 10), 1:5, 1:5))) + @test !equivalent_structure(Dagger.aliasing(rand(10, 10)), + NoAliasing()) + @test !equivalent_structure(NoAliasing(), + Dagger.aliasing(rand(10, 10))) + end +end + +# ---------- Helpers for DAG capture via the schedule cache ---------- + +# N.B. These tests exercise the *flat* AOT DAG-scheduling/caching path +# (`distribute_tasks!` + `datadeps_build_schedule!`), so every `spawn_datadeps` +# call here runs with `hierarchical=false`. Hierarchical scheduling intentionally +# computes a separate, partition-local AOT schedule on each partition's own task +# (no global schedule cache on the calling task), so the caller-visible cache +# these tests inspect is only populated by the flat path. + +""" +Run `f()` inside `spawn_datadeps` with `scheduler`, returning the resulting +schedule cache snapshot (a Vector). The cache is cleared before the run so +each call starts from a known state. +""" +function run_with_fresh_cache(f, scheduler::DataDepsScheduler = LayeredScheduler()) + cache = datadeps_schedule_cache(scheduler) + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => scheduler) do + Dagger.spawn_datadeps(f) + end + return cache +end + +""" +Run a sequence of `n` invocations of `f()` under `scheduler` and return the +cache. Useful for asserting cache size after repeated runs. +""" +function run_n_times(f, n::Int; scheduler::DataDepsScheduler = LayeredScheduler()) + cache = datadeps_schedule_cache(scheduler) + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => scheduler) do + for _ in 1:n + Dagger.spawn_datadeps(f) + end + end + return cache +end + +# Simple test kernels +add!(X, Y) = (X .+= Y; X) +scale!(X, a) = (X .*= a; X) + +# ---------- DAGSpec construction ---------- + +@testset "DAGSpec construction" begin + @testset "Vertices and function types" begin + A = rand(64) + B = rand(64) + cache = run_with_fresh_cache() do + Dagger.@spawn add!(InOut(A), In(B)) + Dagger.@spawn scale!(InOut(A), 2.0) + end + @test length(cache) == 1 + spec = first(cache).first + @test length(spec) == 2 + @test spec.id_to_functype[1] === typeof(add!) + @test spec.id_to_functype[2] === typeof(scale!) + # `tspec.fargs` includes the function as fargs[1], so argtypes also + # includes the function argspec → 3 entries per vertex. + @test length(spec.id_to_argtypes[1]) == 3 + @test length(spec.id_to_argtypes[2]) == 3 + end + + @testset "Argspec positions, types, and dep_mods" begin + A = rand(32, 32) + cache = run_with_fresh_cache() do + Dagger.@spawn add!(InOut(A), In(A)) + end + spec = first(cache).first + argspecs = spec.id_to_argtypes[1] + # 3 argspecs: function (pos=0), A (InOut, pos=1), A (In, pos=2) + @test length(argspecs) == 3 + @test all(a -> a.pos isa Int, argspecs) + # The two data argspecs (positions 1 and 2) refer to a Matrix{Float64} + data_argspecs = filter(a -> a.pos != 0, argspecs) + @test length(data_argspecs) == 2 + @test all(a -> a.value_type === Matrix{Float64}, data_argspecs) + @test all(a -> a.dep_mod === identity, data_argspecs) + @test all(a -> a.ainfo isa ContiguousAliasing, data_argspecs) + # Function argspec is at position 0 with value_type === typeof(add!) + f_argspec = only(filter(a -> a.pos == 0, argspecs)) + @test f_argspec.value_type === typeof(add!) + end + + @testset "Deps with multiple dep_mods at one position" begin + X = Ref(rand(64)) + do_nothing(R) = nothing + cache = run_with_fresh_cache() do + Dagger.@spawn do_nothing(Deps(X, InOut(:x), In(:x))) + end + spec = first(cache).first + argspecs = spec.id_to_argtypes[1] + # Both Deps entries for :x should be present, even though they share pos. + x_deps = filter(a -> a.dep_mod === :x, argspecs) + @test length(x_deps) == 2 + @test all(a -> a.pos == x_deps[1].pos, x_deps) + end + + @testset "Per-task scope recorded" begin + A = rand(64) + my_scope = Dagger.scope(worker=1) + cache = run_with_fresh_cache() do + Dagger.@spawn scope=my_scope add!(InOut(A), In(A)) + end + spec = first(cache).first + @test spec.id_to_scope[1] == my_scope + end +end + +# ---------- DAGSpec equivalence: cache reuse ---------- + +@testset "DAGSpec equivalence (end-to-end cache reuse)" begin + @testset "Identical algorithm with fresh allocations → cache hit" begin + cache = run_n_times(3) do + A = rand(128) + B = rand(128) + Dagger.@spawn add!(InOut(A), In(B)) + Dagger.@spawn add!(InOut(A), In(B)) + end + @test length(cache) == 1 + end + + @testset "Different array sizes → cache miss" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do + for sz in (64, 128, 256) + A = rand(sz); B = rand(sz) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + end + @test length(cache) == 3 + end + + @testset "Different element types → cache miss" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do + for T in (Float64, Float32, Int) + A = T <: AbstractFloat ? rand(T, 64) : T.(rand(1:100, 64)) + B = T <: AbstractFloat ? rand(T, 64) : T.(rand(1:100, 64)) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + end + @test length(cache) == 3 + end + + @testset "Different task counts → cache miss" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do + for ntasks in (1, 2, 4) + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + for _ in 1:ntasks + Dagger.@spawn add!(InOut(A), In(B)) + end + end + end + end + @test length(cache) == 3 + end + + @testset "Different per-task scope → cache miss" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + scopes = [Dagger.DefaultScope(), + Dagger.ExactScope(Dagger.ThreadProc(1, 1))] + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do + for scp in scopes + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn scope=scp add!(InOut(A), In(B)) + end + end + end + @test length(cache) == length(scopes) + end + + @testset "Different functions → cache miss" begin + cache = run_with_fresh_cache() do + A = rand(64); B = rand(64) + Dagger.@spawn add!(InOut(A), In(B)) + end + @test length(cache) == 1 + # Repeat with the other function + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do + A = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn scale!(InOut(A), 2.0) + end + end + @test length(cache) == 2 + end +end + +# ---------- Multi-scheduler cache partitioning ---------- + +@testset "Cache partitioning between schedulers" begin + ls_cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(ls_cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + @test length(ls_cache) == 1 + + # RoundRobinScheduler doesn't AOT-schedule, so it shouldn't add to its + # cache, and it must NOT see LayeredScheduler's entry. + rr_cache = datadeps_schedule_cache(RoundRobinScheduler()) + @test rr_cache !== ls_cache + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => RoundRobinScheduler()) do + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + @test isempty(rr_cache) + @test length(ls_cache) == 1 +end + +# ---------- Benchmark algorithms ---------- + +# A small mock-Cholesky-like algorithm that doesn't depend on DArray, to keep +# the test fast and self-contained. +function mock_cholesky!(M::Matrix{<:Matrix}) + mt = size(M, 1) + for k in 1:mt + Dagger.@spawn LinearAlgebra.LAPACK.potrf!('L', InOut(M[k, k])) + for m in (k+1):mt + Dagger.@spawn LinearAlgebra.BLAS.trsm!('R', 'L', 'T', 'N', + 1.0, In(M[k, k]), + InOut(M[m, k])) + end + for n in (k+1):mt + Dagger.@spawn LinearAlgebra.BLAS.syrk!('L', 'N', -1.0, + In(M[n, k]), 1.0, + InOut(M[n, n])) + for m in (n+1):mt + Dagger.@spawn LinearAlgebra.BLAS.gemm!('N', 'T', -1.0, + In(M[m, k]), + In(M[n, k]), 1.0, + InOut(M[m, n])) + end + end + end +end + +make_spd_blocks(sz::Int, nb::Int) = begin + M_dense = rand(sz, sz); M_dense = M_dense * M_dense' + M_dense[diagind(M_dense)] .+= sz + return [M_dense[i:(i+nb-1), j:(j+nb-1)] for i in 1:nb:sz, j in 1:nb:sz] +end + +@testset "Benchmark: mock Cholesky" begin + @testset "Repeated runs, fresh allocations → 1 cache entry" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do + for _ in 1:3 + M = make_spd_blocks(64, 16) + Dagger.spawn_datadeps() do + mock_cholesky!(M) + end + end + end + @test length(cache) == 1 + spec = first(cache).first + # 4x4 block grid → potrf + trsms + syrks + gemms. + # Count expected: sum_{k=1}^{4} (1 + (4-k) + (4-k) + (4-k)(4-k-1)/2 * 2) + # Or just sanity check >0. + @test length(spec) > 0 + end + + @testset "Different block-grid sizes → distinct cache entries" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do + # 2x2 grid (sz=32, nb=16) vs 4x4 grid (sz=64, nb=16) → different task counts + for (sz, nb) in ((32, 16), (64, 16)) + M = make_spd_blocks(sz, nb) + Dagger.spawn_datadeps() do + mock_cholesky!(M) + end + end + end + @test length(cache) == 2 + end + + @testset "Same block-grid shape, different block size → distinct entries" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do + # Both 2x2 grids; block size 16 vs 32 → same task count, different ContiguousAliasing length + for (sz, nb) in ((32, 16), (64, 32)) + M = make_spd_blocks(sz, nb) + Dagger.spawn_datadeps() do + mock_cholesky!(M) + end + end + end + @test length(cache) == 2 + end +end + +@testset "Benchmark: tree reduce" begin + function tree_reduce!(As) + to_reduce = Vector[] + push!(to_reduce, As) + while !isempty(to_reduce) + xs = pop!(to_reduce) + n = length(xs) + if n == 2 + Dagger.@spawn Base.mapreducedim!(identity, +, InOut(xs[1]), In(xs[2])) + elseif n > 2 + push!(to_reduce, [xs[1], xs[div(n, 2)+1]]) + push!(to_reduce, xs[1:div(n, 2)]) + push!(to_reduce, xs[div(n, 2)+1:end]) + end + end + end + + @testset "Repeated runs → 1 cache entry" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do + for _ in 1:3 + As = [rand(64) for _ in 1:8] + Dagger.spawn_datadeps() do + tree_reduce!(As) + end + end + end + @test length(cache) == 1 + end + + @testset "Different reduction widths → distinct entries" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do + for n in (4, 8, 16) + As = [rand(64) for _ in 1:n] + Dagger.spawn_datadeps() do + tree_reduce!(As) + end + end + end + @test length(cache) == 3 + end +end + +# ---------- Direct datadeps_dag_equivalent semantics ---------- + +@testset "datadeps_dag_equivalent (direct)" begin + LS = LayeredScheduler() + + # Build two DAGSpecs from equivalent-but-freshly-allocated workloads. + cache = datadeps_schedule_cache(LS) + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LS) do + for _ in 1:2 + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + Dagger.@spawn scale!(InOut(A), 2.0) + end + end + end + # Cache hit on second run → still 1 entry + @test length(cache) == 1 + spec_a = first(cache).first + + # Now build a structurally-different spec. + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LS) do + A = rand(128); B = rand(128) # different size + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + Dagger.@spawn scale!(InOut(A), 2.0) + end + end + spec_b = first(cache).first + + @test datadeps_dag_equivalent(LS, spec_a, spec_a) + @test datadeps_dag_equivalent(LS, spec_b, spec_b) + @test !datadeps_dag_equivalent(LS, spec_a, spec_b) +end + +# ---------- Scheduler overrides ---------- + +# A scheduler that opts out of all caching by always returning false. +struct NoCacheScheduler <: DataDepsScheduler end +Dagger.datadeps_dag_equivalent(::NoCacheScheduler, ::DAGSpec, ::DAGSpec) = false +function Dagger.datadeps_schedule_dag_aot!(::NoCacheScheduler, schedule, dag_spec, all_procs, all_scope) + for idx in 1:Dagger.nv(dag_spec.g) + task = dag_spec.id_to_task[idx] + schedule[task] = first(all_procs) + end +end +Dagger.datadeps_schedule_task_jit!(::NoCacheScheduler, all_procs, all_scope, task_scope, spec, task) = + first(all_procs) + +@testset "Scheduler-overridable equivalence" begin + @testset "Always-false override never hits cache" begin + nc = NoCacheScheduler() + cache = datadeps_schedule_cache(nc) + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => nc) do + for _ in 1:3 + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + end + @test length(cache) == 3 + end +end + +# A scheduler that uses a strict pointer-identical comparison (the original +# intent of `equivalent_structure`). +struct PtrStrictScheduler <: DataDepsScheduler end +Dagger.datadeps_ainfo_equivalent(::PtrStrictScheduler, + a1::Dagger.AbstractAliasing, + a2::Dagger.AbstractAliasing) = + hash(a1) == hash(a2) +function Dagger.datadeps_schedule_dag_aot!(::PtrStrictScheduler, schedule, dag_spec, all_procs, all_scope) + for idx in 1:Dagger.nv(dag_spec.g) + task = dag_spec.id_to_task[idx] + schedule[task] = first(all_procs) + end +end +Dagger.datadeps_schedule_task_jit!(::PtrStrictScheduler, all_procs, all_scope, task_scope, spec, task) = + first(all_procs) + +@testset "PtrStrictScheduler: structural matches miss, pointer-identical hits" begin + ps = PtrStrictScheduler() + cache = datadeps_schedule_cache(ps) + + # Fresh allocations each run → pointer hash differs → all miss. + empty!(cache) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => ps) do + for _ in 1:3 + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + end + @test length(cache) == 3 + + # Same allocation reused each run → pointer hash matches → cache hits. + empty!(cache) + A = rand(64); B = rand(64) + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => ps) do + for _ in 1:3 + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + end + @test length(cache) == 1 +end diff --git a/test/gpu.jl b/test/gpu.jl index 474a0bd08..28b908eb6 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -1,11 +1,113 @@ using Random using LinearAlgebra +using SparseArrays +using Krylov @everywhere begin using Distributed, Dagger import Dagger: Kernel using KernelAbstractions end + +# Sparse DArray SpGEMM / SpMV under a GPU compute scope. `check_tile(chunk)` +# verifies a distributed tile is device-resident (backend-specific). +function test_gpu_sparse_darray(scope; check_tile) + SA = sprand(Float32, 8, 8, 0.35) + SB = sprand(Float32, 8, 8, 0.35) + x = rand(Float32, 8) + Dagger.with_options(;scope) do + DSA = distribute(SA, Blocks(4, 4)) + DSB = distribute(SB, Blocks(4, 4)) + for chunk in DSA.chunks + @test check_tile(fetch(chunk; raw=true)) + end + + @test collect(DSA * DSB) ≈ SA * SB + @test collect(DSA * DSB') ≈ SA * SB' + @test collect(DSA' * DSB) ≈ SA' * SB + + DSC = similar(DSA) + mul!(DSC, DSA, DSB) + @test collect(DSC) ≈ SA * SB + + Dx = distribute(x, Blocks(4)) + @test collect(DSA * Dx) ≈ SA * x + + Z = SparseArrays.spzeros(Blocks(4, 4), Float32, 8, 8) + # `collect` densifies; check emptiness on device tiles / dense gather. + @test iszero(sum(abs, collect(Z))) + for chunk in Z.chunks + @test check_tile(fetch(chunk; raw=true)) + end + end +end + +# Sparse Krylov solvers under a GPU scope (mirrors the sparse backend cases in +# `array/linalg/iterativesolvers.jl`, trimmed for GPU CI time). +function _gpu_sparse_laplacian(T, n) + return SparseArrays.spdiagm( + -1 => fill(-one(T), n - 1), + 0 => fill(T(4), n), + 1 => fill(-one(T), n - 1), + ) +end +function _gpu_sparse_advection(T, n) + return _gpu_sparse_laplacian(T, n) + SparseArrays.spdiagm( + -1 => fill(T(-3) / 10, n - 1), + 1 => fill(T(3) / 10, n - 1), + ) +end +function test_gpu_sparse_solvers(scope; check_tile=nothing) + n, k = 32, 8 + A_part, b_part = Blocks(k, k), Blocks(k) + T = Float32 + + Dagger.with_options(;scope) do + Asp = _gpu_sparse_laplacian(T, n) + b = rand(T, n) + xref = Matrix(Asp) \ b + DA = distribute(Asp, A_part) + Db = distribute(b, b_part) + if check_tile !== nothing + @test check_tile(fetch(DA.chunks[1]; raw=true)) + end + + @testset "$(nameof(solver))" for solver in (Dagger.cg, Dagger.minres, Dagger.gmres, Dagger.bicgstab) + x, stats = solver(DA, Db; atol = T(1e-6), rtol = T(1e-5), itmax = 500) + @test stats.solved + @test x isa Dagger.DVector + @test collect(x) ≈ xref rtol = 1e-3 + end + + x, stats = Dagger.krylov_solve(:cg, DA, Db; atol = T(1e-6), rtol = T(1e-5)) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-3 + + # Diagonal Jacobi is device-friendly. Block-Jacobi uses host UMFPACK + # factors that Datadeps cannot place under a GPU-only compute scope, so + # it is covered on CPU in `array/linalg/iterativesolvers.jl` instead. + P = Dagger.JacobiPreconditioner(DA) + @test collect(P.dinv) ≈ fill(T(1) / T(4), n) + y = similar(Db) + mul!(y, P, Db) + @test collect(y) ≈ (T(1) / T(4)) .* b + x, stats = Dagger.cg(DA, Db; M = P, atol = T(1e-6), rtol = T(1e-5), itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-3 + + Anonsym = _gpu_sparse_advection(T, n) + bn = rand(T, n) + xrefn = Matrix(Anonsym) \ bn + DAn = distribute(Anonsym, A_part) + Dbn = distribute(bn, b_part) + @testset "$(nameof(solver)) nonsym" for solver in (Dagger.gmres, Dagger.bicgstab) + x, stats = solver(DAn, Dbn; atol = T(1e-6), rtol = T(1e-5), itmax = 500) + @test stats.solved + @test collect(x) ≈ xrefn rtol = 1e-3 + end + end +end + @everywhere begin function isongpu(X) return !(X isa Array) @@ -227,6 +329,19 @@ end end @test collect(Db) ≈ b_ref rtol=1e-5 end + + @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs + scope = Dagger.scope(worker=1, cuda_gpu=gpu) + CUDAExt = Base.get_extension(Dagger, :CUDAExt) + check_tile = chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa CUDA.CUSPARSE.CuSparseMatrixCSC && + chunk.space isa CUDAExt.CUDAVRAMMemorySpace + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) + end end end @@ -398,6 +513,19 @@ end end @test collect(Db) ≈ b_ref rtol=1e-5 end + + @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs + scope = Dagger.scope(worker=1, rocm_gpu=gpu) + ROCExt = Base.get_extension(Dagger, :ROCExt) + check_tile = chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa AMDGPU.rocSPARSE.ROCSparseMatrixCSC && + chunk.space isa ROCExt.ROCVRAMMemorySpace + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) + end end end @@ -569,6 +697,19 @@ end end @test collect(Db) ≈ b_ref rtol=1e-5 end + + @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs + scope = Dagger.scope(worker=1, intel_gpu=gpu) + IntelExt = Base.get_extension(Dagger, :IntelExt) + check_tile = chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa Dagger.DeviceSparseMatrixCSC && + chunk.space isa IntelExt.IntelVRAMMemorySpace + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) + end end end @@ -714,6 +855,19 @@ end @test_broken array[2, 1] == 4.0f0 @test_broken array[2, 2] == 5.0f0 end + + @testset "Sparse DArray" begin + scope = Dagger.scope(worker=1, metal_gpu=1) + MetalExt = Base.get_extension(Dagger, :MetalExt) + check_tile = chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa Dagger.DeviceSparseMatrixCSC && + chunk.space isa MetalExt.MetalVRAMMemorySpace + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) + end end end @@ -820,5 +974,18 @@ end end @test A ≈ ref .+ 1 end + + @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs + scope = Dagger.scope(worker=1, cl_device=gpu) + OpenCLExt = Base.get_extension(Dagger, :OpenCLExt) + check_tile = chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa Dagger.DeviceSparseMatrixCSC && + chunk.space isa OpenCLExt.CLMemorySpace + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) + end end end diff --git a/test/metalenv/Project.toml b/test/metalenv/Project.toml index b3d5e9632..903bfcb7e 100644 --- a/test/metalenv/Project.toml +++ b/test/metalenv/Project.toml @@ -4,6 +4,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources] diff --git a/test/mpi.jl b/test/mpi.jl index f5e7ffd6f..120974761 100644 --- a/test/mpi.jl +++ b/test/mpi.jl @@ -543,14 +543,16 @@ read_only(X) = (sum(X); nothing) r1 = min(1, nranks-1) cross = r1 != 0 # single-rank runs degenerate to zero copies - # Read-only cross-rank arg: one copy-in, no writeback + # Read-only cross-rank arg: copy-in plus region-end write-back. Write-back + # is not elided via `arg_current`; for a pure `In` the bytes are unchanged, + # but planning still inserts the copy-from. A = ones(4, 4) logs = with_logs() do Dagger.spawn_datadeps() do Dagger.@spawn scope=rank_scope(r1) read_only(In(A)) end end - @test count_move_tasks(logs) == (cross ? 1 : 0) + @test count_move_tasks(logs) == (cross ? 2 : 0) @test Dagger.check_uniform(count_move_tasks(logs)) # Same-rank InOut: the origin slot aliases the original, zero copies diff --git a/test/mpienv/Project.toml b/test/mpienv/Project.toml index 71e617954..581efcb0e 100644 --- a/test/mpienv/Project.toml +++ b/test/mpienv/Project.toml @@ -3,6 +3,7 @@ Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources] diff --git a/test/oneapienv/Project.toml b/test/oneapienv/Project.toml index 30e7631f3..1c5664279 100644 --- a/test/oneapienv/Project.toml +++ b/test/oneapienv/Project.toml @@ -3,6 +3,7 @@ Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" diff --git a/test/openclenv/Project.toml b/test/openclenv/Project.toml index b7718eca7..8c25145c4 100644 --- a/test/openclenv/Project.toml +++ b/test/openclenv/Project.toml @@ -1,10 +1,13 @@ [deps] Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd" [sources] Dagger = {path = "../.."} diff --git a/test/rocmenv/Project.toml b/test/rocmenv/Project.toml index d350b4c4a..7b8ac67f3 100644 --- a/test/rocmenv/Project.toml +++ b/test/rocmenv/Project.toml @@ -1,9 +1,11 @@ [deps] AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources] diff --git a/test/runtests.jl b/test/runtests.jl index bfc347ce0..2ac87b2cd 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -7,6 +7,11 @@ USE_METAL = parse(Bool, get(ENV, "CI_USE_METAL", "0")) USE_OPENCL = parse(Bool, get(ENV, "CI_USE_OPENCL", "0")) USE_GPU = USE_CUDA || USE_ROCM || USE_ONEAPI || USE_METAL || USE_OPENCL +# Finch is a heavy optional dependency (precompilation takes minutes), so it is +# not a default test dependency. The dedicated Finch CI job signals via this +# variable to `Pkg.add` Finch and run only the Finch-backed tests. +USE_FINCH = parse(Bool, get(ENV, "CI_TEST_FINCH", "0")) + tests = [ ("Thunk", "thunk.jl"), ("Scheduler", "scheduler.jl"), @@ -21,8 +26,10 @@ tests = [ ("Task Queues", "task-queues.jl"), ("Task Affinity", "task-affinity.jl"), ("Datadeps", "datadeps.jl"), + ("Datadeps - Scheduling", "datadeps/scheduling.jl"), ("Streaming", "streaming.jl"), ("Domain Utilities", "domain.jl"), + ("Tapes", "tapes.jl"), ("Array - Allocation", "array/allocation.jl"), ("Array - Indexing", "array/indexing.jl"), ("Array - Core", "array/core.jl"), @@ -32,11 +39,15 @@ tests = [ ("Array - LinearAlgebra - Core", "array/linalg/core.jl"), ("Array - LinearAlgebra - Arithmetic", "array/linalg/arithmetic.jl"), ("Array - LinearAlgebra - Matmul", "array/linalg/matmul.jl"), + ("Array - LinearAlgebra - Matmul (Sparse, quick)", "array/linalg/matmul_sparse.jl"), + ("Array - LinearAlgebra - Matmul (Finch, quick)", "array/linalg/matmul_finch.jl"), ("Array - LinearAlgebra - Cholesky", "array/linalg/cholesky.jl"), ("Array - LinearAlgebra - LU", "array/linalg/lu.jl"), ("Array - LinearAlgebra - Solve", "array/linalg/solve.jl"), ("Array - LinearAlgebra - QR", "array/linalg/qr.jl"), ("Array - LinearAlgebra - SVD", "array/linalg/svd.jl"), + ("Array - LinearAlgebra - Iterative Solvers", "array/linalg/iterativesolvers.jl"), + ("Array - LinearAlgebra - Sparse Direct", "array/linalg/sparsedirect.jl"), ("Array - Permute", "array/permute.jl"), ("Array - Random", "array/random.jl"), ("Array - Stencils", "array/stencil.jl"), @@ -57,8 +68,22 @@ if USE_GPU ("Array - Stencils", "array/stencil.jl"), ] end +if USE_FINCH + # Only run the Finch-backed tests in the dedicated Finch CI job. + tests = [ + ("Array - LinearAlgebra - Matmul (Finch)", "array/linalg/matmul_finch.jl"), + ] +end all_test_names = map(test -> replace(last(test), ".jl"=>""), tests) +# Tests excluded from default runs; they only run when explicitly requested via +# `--test` or enabled by an environment signal. Finch's precompilation is very +# heavy, so its tests are opt-in (except in the dedicated Finch job, where we +# explicitly want them to run). +optin_test_names = USE_FINCH ? String[] : String[ + "array/linalg/matmul_finch", +] + additional_workers::Int = 3 worker_threads::Int = 1 @@ -71,6 +96,10 @@ if PROGRAM_FILE != "" && realpath(PROGRAM_FILE) == @__FILE__ Pkg.instantiate() catch end + if USE_FINCH + # Finch is not a default test dependency; add it on demand when signaled. + Pkg.add("Finch") + end using ArgParse s = ArgParseSettings(description = "Dagger Testsuite") @@ -109,7 +138,7 @@ if PROGRAM_FILE != "" && realpath(PROGRAM_FILE) == @__FILE__ parsed_args = parse_args(s) to_test = String[] if isempty(parsed_args["test"]) - to_test = copy(all_test_names) + to_test = filter(!in(optin_test_names), copy(all_test_names)) else for _test in parsed_args["test"] test = only(_test) @@ -166,7 +195,7 @@ if PROGRAM_FILE != "" && realpath(PROGRAM_FILE) == @__FILE__ Pkg.offline(true) end else - to_test = all_test_names + to_test = filter(!in(optin_test_names), all_test_names) @info "Running all tests" end diff --git a/test/tapes.jl b/test/tapes.jl new file mode 100644 index 000000000..63f8561a3 --- /dev/null +++ b/test/tapes.jl @@ -0,0 +1,328 @@ +# Tests for the operation-tape layout predictor. +# +# Drop into `test/` and add `include("tapes.jl")` to `test/runtests.jl`. +# +# Most of these exercise the pure logic (site keys, trie, prediction, cost +# models, planner) with synthetic specs, so they run without a cluster and in +# milliseconds. The end-to-end group at the bottom needs real DArrays. + +using Dagger: Tapes +using Dagger.Tapes: SiteKey, OpKey, ArgSpec, ArgView, LayoutChoice, PredictedOp, + TapeNode, TapeRoot, MachineModel, + size_bucket, predict, advance!, plan_chain, candidate_layouts, + fallback_layout, cost_of, generic_op_cost, step_cost, + op_cost, has_cost_model, blocksize, nblocks, ntiles, + to_blocks, confidence, child_total + +# A deterministic machine so cost comparisons are reproducible. +const TESTMACHINE = MachineModel(nprocs = 16, + flops_per_sec = 1.0e12, + bandwidth = 1.0e10, + latency = 1.0e-5, + task_overhead = 1.0e-4, + mem_per_proc = 8.0e9, + calibrated = true) + +spec(T, dims, bs; assign = :arbitrary) = ArgSpec(T, dims, bs, assign) + +function fresh!() + Tapes.clear!() + Tapes.set_machine!(TESTMACHINE) + return nothing +end + +was_enabled = Tapes.is_enabled() + +@testset "site keys" begin + fresh!() + Tapes.CONFIG.size_buckets_per_octave = 2 + + # Nearby sizes share a bucket; distant ones do not. + @test size_bucket((1024, 1024)) == size_bucket((1100, 1100)) + @test size_bucket((1024, 1024)) != size_bucket((8192, 8192)) + + k1 = SiteKey(UInt64(7), Float64, Int32(2), size_bucket((1024, 1024))) + k2 = SiteKey(UInt64(7), Float64, Int32(2), size_bucket((1024, 1024))) + k3 = SiteKey(UInt64(7), Float32, Int32(2), size_bucket((1024, 1024))) + @test k1 == k2 + @test hash(k1) == hash(k2) + @test k1 != k3 # element type is part of identity + + Tapes.CONFIG.size_buckets_per_octave = 0 + @test size_bucket((1024, 1024)) != size_bucket((1100, 1100)) + Tapes.CONFIG.size_buckets_per_octave = 2 +end + +@testset "trie recording and prediction" begin + fresh!() + key = SiteKey(UInt64(1), Float64, Int32(2), size_bucket((1024, 1024))) + root = Tapes.get_root!(key) + self = spec(Float64, (1024, 1024), (256, 256)) + layout = LayoutChoice((256, 256), :arbitrary, :square) + + chain = [OpKey(:mul!, 1, 3), OpKey(:cholesky!, 1, 1), OpKey(:trsm!, 1, 2)] + + # Ten identical executions of the same chain. + for _ in 1:10 + tr = Tapes.LiveTrace(root, self, layout, PredictedOp[], LayoutChoice[layout]) + root.nobservations += 1 + root.root.count += 1 + for op in chain + advance!(tr, op, ArgSpec[self]) + end + end + + pred = predict(root.root) + @test length(pred) == 3 + @test [p.key.op for p in pred] == [:mul!, :cholesky!, :trsm!] + @test confidence(pred) ≈ 1.0 + + # Divergence: five runs take a different second step. + for _ in 1:5 + tr = Tapes.LiveTrace(root, self, layout, PredictedOp[], LayoutChoice[layout]) + root.nobservations += 1 + root.root.count += 1 + advance!(tr, OpKey(:mul!, 1, 3), ArgSpec[self]) + advance!(tr, OpKey(:lu!, 1, 1), ArgSpec[self]) + end + + pred = predict(root.root) + @test pred[1].key.op == :mul! + @test pred[1].prob ≈ 1.0 # first step is still certain + @test pred[2].key.op == :cholesky! # modal branch wins + @test 0.6 < pred[2].prob < 0.7 # 10/15 + + # Prefix conditioning: from the :mul! node the answer is the same, but + # from a node we have never reached there is nothing to say. + mul_node = root.root.children[OpKey(:mul!, 1, 3)] + @test predict(mul_node)[1].key.op == :cholesky! + @test child_total(mul_node) == 15 +end + +@testset "allocated-but-unused reduces confidence" begin + fresh!() + key = SiteKey(UInt64(2), Float64, Int32(2), size_bucket((512, 512))) + root = Tapes.get_root!(key) + self = spec(Float64, (512, 512), (128, 128)) + layout = LayoutChoice((128, 128), :arbitrary, :square) + + # One array gets used, nine are allocated and abandoned. + tr = Tapes.LiveTrace(root, self, layout, PredictedOp[], LayoutChoice[layout]) + root.nobservations += 1; root.root.count += 1 + advance!(tr, OpKey(:cholesky!, 1, 1), ArgSpec[self]) + for _ in 1:9 + root.nobservations += 1; root.root.count += 1 + end + + pred = predict(root.root) + # 1/10 is below min_branch_prob, so we should predict nothing at all + # rather than confidently predicting :cholesky!. + @test isempty(pred) +end + +@testset "cost models" begin + fresh!() + @test has_cost_model(:cholesky!) + @test has_cost_model(:mul!) + @test !has_cost_model(:definitely_not_an_op) + + s = spec(Float64, (8192, 8192), (512, 512)) + + square = LayoutChoice((512, 512), :cyclicrow, :square) + rowblk = LayoutChoice((512, 8192), :blockrow, :rowblock) + views(l) = ArgView[ArgView(s, l)] + + # Cholesky must prefer square tiles to fat row blocks. + @test cost_of(:cholesky!, views(square), TESTMACHINE) < + cost_of(:cholesky!, views(rowblk), TESTMACHINE) + + # An unmodelled op falls through to the generic model without erroring. + c = cost_of(:some_unknown_op, views(square), TESTMACHINE) + @test isfinite(c) && c > 0 + + # Degenerate layouts must be punished, not merely disfavoured. + tiny = LayoutChoice((1, 1), :arbitrary, :square) + huge = LayoutChoice((8192, 8192), :arbitrary, :square) + base = generic_op_cost(views(square), TESTMACHINE) + @test generic_op_cost(views(tiny), TESTMACHINE) > 10 * base + @test generic_op_cost(views(huge), TESTMACHINE) > base +end + +@testset "@cost_model defines a usable model" begin + fresh!() + Tapes.@cost_model __test_op(A, B) = flops_time(Float64(length(A)) * size(B, 1)) + + task_time(ntiles(A)) + @test has_cost_model(:__test_op) + sA = spec(Float64, (100, 100), (50, 50)) + sB = spec(Float64, (100, 10), (50, 10)) + lA = LayoutChoice((50, 50), :arbitrary, :square) + lB = LayoutChoice((50, 10), :arbitrary, :colblock) + v = ArgView[ArgView(sA, lA), ArgView(sB, lB)] + @test cost_of(:__test_op, v, TESTMACHINE) > 0 + + # Too few recorded arguments must degrade, not throw. + @test cost_of(:__test_op, ArgView[ArgView(sA, lA)], TESTMACHINE) > 0 +end + +@testset "candidate generation" begin + fresh!() + cands = candidate_layouts(Float64, (8192, 8192), TESTMACHINE) + @test !isempty(cands) + @test length(cands) <= Tapes.CONFIG.max_candidates + @test fallback_layout(Float64, (8192, 8192)) in cands # baseline always present + @test any(c -> c.label === :square, cands) + @test any(c -> c.label === :rowblock, cands) + @test all(c -> all(>(0), blocksize(c)), cands) + + # Vectors must not produce degenerate 2D candidates. + cv = candidate_layouts(Float64, (100_000,), TESTMACHINE) + @test !isempty(cv) + @test all(c -> c.ndims == 1, cv) +end + +@testset "planner" begin + fresh!() + s = spec(Float64, (8192, 8192), (512, 512)) + cands = candidate_layouts(Float64, (8192, 8192), TESTMACHINE) + + # A confident, homogeneous Cholesky chain should land on square tiles. + pred = [PredictedOp(OpKey(:cholesky!, 1, 1), 1.0, ArgSpec[s]) for _ in 1:3] + lp = plan_chain(pred, cands, s, TESTMACHINE) + @test lp.accepted + @test lp.cost <= lp.fallback_cost + @test length(lp.steps) == 3 + @test Dagger.Tapes.aspect(ArgView(s, lp.steps[1])) < 4.0 + + # An empty forecast must decline rather than guess. + lp0 = plan_chain(PredictedOp[], cands, s, TESTMACHINE) + @test !lp0.accepted + @test lp0.reason === :no_prediction + + # A low-confidence forecast is subject to the regret bound. + predlow = [PredictedOp(OpKey(:cholesky!, 1, 1), 0.3, ArgSpec[s])] + lplow = plan_chain(predlow, cands, s, TESTMACHINE) + @test lplow.accepted == false || lplow.max_regret <= Tapes.CONFIG.max_regret_ratio + + # The gate must reject a plan that does not beat the fallback by the margin. + old = Tapes.CONFIG.gate_margin + Tapes.CONFIG.gate_margin = 0.0 # nothing can beat this + @test !plan_chain(pred, cands, s, TESTMACHINE).accepted + Tapes.CONFIG.gate_margin = old +end + +@testset "planner accounts for repartitioning" begin + fresh!() + s = spec(Float64, (8192, 8192), (512, 512)) + cands = candidate_layouts(Float64, (8192, 8192), TESTMACHINE) + + # A chain that alternates between two operations with opposed preferences. + # Whatever the planner decides, it must not thrash: the number of distinct + # layouts in the plan should be small relative to the chain length, because + # each change is charged a full redistribution. + pred = PredictedOp[] + for i in 1:8 + op = isodd(i) ? :cholesky! : :reduce + push!(pred, PredictedOp(OpKey(op, 1, 1), 1.0, ArgSpec[s])) + end + lp = plan_chain(pred, cands, s, TESTMACHINE) + @test length(unique(lp.steps)) <= 3 +end + +@testset "enable/disable is inert by default" begin + fresh!() + Tapes.disable!() + @test !Tapes.is_enabled() + p = Tapes.plan_allocation(Float64, (1024, 1024)) + @test p.partitioning isa AutoBlocks # request passed through untouched + @test p.root === nothing + + Tapes.enable!(site_id = :lexical) + @test Tapes.is_enabled() + # An explicit Blocks must survive untouched unless explicitly overridden. + p = Tapes.plan_allocation(Float64, (1024, 1024); requested = Blocks(64, 64)) + @test p.partitioning == Blocks(64, 64) + Tapes.disable!() +end + +@testset "pinning overrides prediction" begin + fresh!() + Tapes.enable!(site_id = :lexical) + Tapes.pin!(Float64, (2048, 2048), Blocks(256, 256), :cyclicrow) + p = Tapes.plan_allocation(Float64, (2048, 2048)) + @test p.partitioning == Blocks(256, 256) + @test p.assignment === :cyclicrow + Tapes.unpin!() + Tapes.disable!() +end + +@testset "introspection does not throw" begin + fresh!() + Tapes.enable!(site_id = :lexical) + Tapes.plan_allocation(Float64, (1024, 1024)) + io = IOBuffer() + @test Tapes.report(io) === nothing + @test Tapes.dump_tapes(io) === nothing + @test Tapes.explain(io, Float64, (1024, 1024)) === nothing + @test Tapes.stats().sites >= 1 + Tapes.disable!() +end + +# --------------------------------------------------------------------------- +# End-to-end. Requires the `@record_op` / `plan_allocation` patch points from +# `tapes/integration.jl` to actually be applied to Dagger's allocators and +# linear-algebra routines; skipped otherwise. +# --------------------------------------------------------------------------- + +@testset "end-to-end learning" begin + fresh!() + Tapes.enable!(site_id = :backtrace, min_observations = 1, verbose = false) + try + n = 512 + function workload() + A = rand(AutoBlocks(), Float64, n, n) + B = A * A' + n * I + return cholesky!(B) + end + + workload() # cold: learns nothing yet + s1 = Tapes.stats() + workload() # warm: should now predict + s2 = Tapes.stats() + + @test s2.observations >= s1.observations + if s2.observations == 0 + @info "Tapes end-to-end: no allocations tracked; integration patches not applied" + else + @test s2.sites >= 1 + end + finally + Tapes.disable!() + Tapes.clear!() + Tapes.reset_machine!() + end +end + +@testset "@expect_ops declares ahead of time" begin + fresh!() + Tapes.enable!(site_id = :lexical, min_observations = 1) + try + n = 512 + Dagger.@expect_ops [:cholesky!, :trsm!] begin + # With a declaration in scope, the very first allocation should be + # planned rather than falling back, since confidence is asserted. + p = Tapes.plan_allocation(Float64, (n, n)) + @test p.root !== nothing + @test !isempty(p.predicted) + @test p.predicted[1].key.op === :cholesky! + @test p.predicted[1].prob ≈ 1.0 + end + finally + Tapes.disable!() + Tapes.clear!() + Tapes.reset_machine!() + end +end + +if was_enabled + Tapes.enable!() +end \ No newline at end of file