From 635d88e1c75c6aae15849375a643c0d8a5576c9d Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 18 Aug 2026 12:55:52 -0700 Subject: [PATCH 1/2] DArray: Accept any array type from a partial reduction A partial reduction which produced an array that is not a `Base.Array` was wrapped in a 1x1 array, rather than being passed along as the array it already is, which added a level of nesting to the result of the reduction. Co-authored-by: Cursor --- src/array/map-reduce.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/array/map-reduce.jl b/src/array/map-reduce.jl index 6e16a912b..a8dc4ba97 100644 --- a/src/array/map-reduce.jl +++ b/src/array/map-reduce.jl @@ -83,7 +83,10 @@ function stage(ctx::Context, r::MapReduce{T,N}) where {T,N} A[1] = x return A end - to_array(x::Array, N) = x + # Any array is already the array that this needs to produce, whether or not it + # is a `Base.Array`: a partial reduction computed by Reactant, for instance, + # comes back as one of its own array types + to_array(x::AbstractArray, N) = x function treered_f(op, x, y, N) value = op.(x, y) return to_array(value, N) From 849cd5c173e26ebdc31814fd4f18d7bf2ef5e653 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 18 Aug 2026 12:56:09 -0700 Subject: [PATCH 2/2] Add Reactant.jl integration `Dagger.@reactant expr` hands the work that `expr` launches to Reactant.jl, which compiles it through MLIR/XLA and optimizes it in ways that Julia's own compiler does not. Two modes are available: - `mode=:inner` (the default) compiles each task's function on its own, on the processor that runs it, leaving Dagger's scheduling and data movement alone. - `mode=:full` traces a whole `spawn_datadeps` region as one program, which lets Reactant fuse and reorder across task boundaries, at the cost of Datadeps doing no planning or scheduling for that region. Whatever Reactant cannot do is done without it: a missing Reactant.jl, a task function that cannot be traced, and a region that cannot be compiled or run are all warned about and then run as they would have been otherwise, so the same application code works with and without Reactant. `must_opt=true` and `must_load=true` make those cases errors instead, for code which cannot afford to quietly stop being accelerated. Compiled programs are cached and reused whenever the values Reactant baked into them are known to match. Reactant is a weak dependency, and all of the logic that uses it lives in the ReactantExt extension. Cholesky factorization needs two kernels that Reactant cannot use as written: `LAPACK.potrf!` is a `ccall`, and `BLAS.syrk!` currently lowers to an op which updates the opposite triangle of `C` from what BLAS does, so both are given traceable implementations here. Reactant tests run in their own environment and CI job, as it is far too heavy a dependency for the main testsuite. `test/reactantenv/bench.jl` compares each mode against plain Dagger; what it finds is written up in the docs. Co-authored-by: Cursor --- .github/workflows/CI.yml | 22 + Project.toml | 3 + docs/make.jl | 1 + docs/src/datadeps.md | 19 + docs/src/reactant.md | 246 +++++++++++ ext/ReactantExt.jl | 787 ++++++++++++++++++++++++++++++++++ src/Dagger.jl | 1 + src/datadeps/queue.jl | 38 +- src/dtask.jl | 39 ++ src/options.jl | 3 + src/reactant.jl | 284 ++++++++++++ src/sch/Sch.jl | 7 +- src/utils/dagdebug.jl | 2 +- test/reactant.jl | 83 ++++ test/reactantenv/Project.toml | 14 + test/reactantenv/bench.jl | 196 +++++++++ test/reactantenv/full.jl | 126 ++++++ test/reactantenv/inner.jl | 171 ++++++++ test/reactantenv/runtests.jl | 22 + test/runtests.jl | 1 + 20 files changed, 2061 insertions(+), 4 deletions(-) create mode 100644 docs/src/reactant.md create mode 100644 ext/ReactantExt.jl create mode 100644 src/reactant.jl create mode 100644 test/reactant.jl create mode 100644 test/reactantenv/Project.toml create mode 100644 test/reactantenv/bench.jl create mode 100644 test/reactantenv/full.jl create mode 100644 test/reactantenv/inner.jl create mode 100644 test/reactantenv/runtests.jl diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 88d6d5309..d6127bb5a 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -90,6 +90,28 @@ jobs: token: ${{ secrets.CODECOV_TOKEN }} fail_ci_if_error: false + reactant: + name: Julia 1.12 (Reactant) + runs-on: ubuntu-latest + timeout-minutes: 90 + if: ${{ !contains(github.event.head_commit.message, '[skip tests]') }} + env: + JULIA_NUM_THREADS: '1' + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: '1.12' + arch: x64 + - uses: julia-actions/cache@v3 + # test/reactantenv pulls Reactant (which brings MLIR and XLA with it) plus + # the checked-out Dagger. It has its own environment, and thus its own job, + # because that is far too heavy to ask of the main testsuite. + - name: Instantiate Reactant environment + run: julia --project=test/reactantenv -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' + - name: Run Reactant tests + run: julia --project=test/reactantenv test/reactantenv/runtests.jl + opencl: name: Julia 1.11 (OpenCL) runs-on: ubuntu-latest diff --git a/Project.toml b/Project.toml index a9b5512de..e5078d286 100644 --- a/Project.toml +++ b/Project.toml @@ -45,6 +45,7 @@ Metal = "dde4c033-4e86-420c-a63e-0dd931031962" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] @@ -62,6 +63,7 @@ OpenCLExt = "OpenCL" PlotsExt = ["DataFrames", "Plots"] PythonExt = "PythonCall" ROCExt = "AMDGPU" +ReactantExt = "Reactant" [compat] AMDGPU = "1, 2" @@ -91,6 +93,7 @@ Plots = "1" PrecompileTools = "1.2" Preferences = "1.4.3" PythonCall = "0.9" +Reactant = "0.2.279" Requires = "1" ScopedValues = "1.1" Statistics = "1" diff --git a/docs/make.jl b/docs/make.jl index 611f7301d..e6c777080 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -33,6 +33,7 @@ makedocs(; ], "GPUs" => "gpu.md", "MPI" => "mpi.md", + "Reactant" => "reactant.md", "Option Propagation" => "propagation.md", "Logging and Visualization" => [ "Logging: Basics" => "logging.md", diff --git a/docs/src/datadeps.md b/docs/src/datadeps.md index 2d9e27c4f..1c90024fc 100644 --- a/docs/src/datadeps.md +++ b/docs/src/datadeps.md @@ -322,6 +322,25 @@ Hierarchical scheduling does not yet parallelize everything it could. These are - **Phases 2 and 3 are single-threaded.** Building the DAG and computing aliasing overlaps are incremental, order-dependent algorithms. They are cheap relative to phases 1 and 4 today, but will become the bottleneck as those scale. +## Executing a Region with Reactant + +A region's tasks are normally planned and scheduled by datadeps, as described +above. Alternatively, the whole region can be handed to Reactant.jl, which +compiles it as a single program and is therefore able to optimize across task +boundaries: + +```julia +Dagger.@reactant mode=:full begin + Dagger.spawn_datadeps() do + Dagger.@spawn my_task!(InOut(A)) + Dagger.@spawn another_task!(In(A), Out(B)) + end +end +``` + +See [Reactant](reactant.md) for what this involves, and for the alternative of +compiling each task on its own. + ## Chunk and DTask slicing with `view` The `view` function allows you to efficiently create a "view" of a `Chunk` or `DTask` that contains an array. This enables operations on specific parts of your distributed data using standard Julia array slicing, without needing to materialize the entire array. diff --git a/docs/src/reactant.md b/docs/src/reactant.md new file mode 100644 index 000000000..4857a049c --- /dev/null +++ b/docs/src/reactant.md @@ -0,0 +1,246 @@ +# Reactant + +Dagger normally runs each task's function as Julia compiled it, which means that +every task is optimized on its own, by Julia's compiler alone. +[Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) compiles Julia code +through MLIR and XLA instead, and applies optimizations that Julia does not: +operation fusion, layout selection, buffer reuse, and code generation for +accelerators. + +[`Dagger.@reactant`](@ref) hands work to Reactant, and is meant to be usable by +putting it in front of code you already have: + +```julia +using Dagger, Reactant, LinearAlgebra + +A = rand(Blocks(256, 256), 1024, 1024) +A = A * A' + 1024I + +chol = Dagger.@reactant cholesky(A) +``` + +Reactant is an optional dependency: if it is not loaded, `Dagger.@reactant` +warns once and runs the code as it normally would. The same application code +therefore works with and without Reactant, and there is no need to maintain two +copies of it. + +## Modes + +`Dagger.@reactant` accepts a `mode`, which decides how much of the work Reactant +is given. + +### `mode=:inner` (the default) + +Each task's function is compiled by Reactant individually, on the processor that +runs it: + +```julia +Dagger.@reactant mode=:inner cholesky(A) +``` + +Dagger's scheduling, data movement, and dependency handling are unchanged; the +only difference is that a task's function is executed by an XLA program rather +than by Julia directly. A task's array arguments are converted to Reactant arrays +before the call, and anything the program writes to them is copied back +afterwards, so tasks which mutate their arguments (as tasks in a +[Datadeps](datadeps.md) region normally do) behave as they always have. + +This is the most reliable mode, since each task is optimized in isolation and +nothing about the DAG changes. It also limits what Reactant can do: it never sees +more than one task at a time. + +### `mode=:full` + +Each [`spawn_datadeps`](@ref) region within the expression is handed to Reactant +as a single program: + +```julia +Dagger.@reactant mode=:full cholesky(A) +``` + +Dagger performs no planning or scheduling for such a region. The region's tasks +are recorded in submission order, all of their arguments are pulled to the +calling worker, and the whole algorithm is traced as one Julia function, which +Reactant then compiles. This lets Reactant optimize *across* task boundaries - +fusing operations from different tasks, keeping intermediates in registers or +scratch buffers, and reordering work - which is where the largest gains are. + +The trade-off is that the region no longer runs in parallel across workers, and +that a region which Reactant cannot trace is run by Dagger instead (see +[Falling back](@ref) below). Full mode is the more experimental of the two. + +Tasks launched outside of a Datadeps region are unaffected by `mode=:full`; use +`mode=:inner` for those. + +## Falling back + +Reactant cannot compile everything, and Dagger treats that as a fact of life +rather than an error: whatever Reactant cannot handle is run the way it would +have been without Reactant, and a warning is issued. + +- In `:inner` mode, a task function which Reactant cannot compile is executed + directly. Because whether a function can be traced is a property of its code, + it is only attempted once per function. +- In `:full` mode, a region which Reactant cannot compile or run is handed back + to Datadeps, which plans and schedules it as usual. Nothing in the region has + run at that point (Reactant works on copies of the region's data), so the + result is the same as if `Dagger.@reactant` had not been used. + +The most common reason for falling back is code which accesses arrays one element +at a time. XLA works on whole arrays, so a scalar loop such as + +```julia +for idx in eachindex(A) + A[idx] = idx +end +``` + +cannot be traced. This is what the kernels generated by +[`@stencil`](stencils.md) do, so stencils currently run without Reactant even +under `Dagger.@reactant`, although they still produce the same results. Code written in +terms of whole-array operations - matrix multiplication, factorizations, +broadcasts, reductions - is what Reactant is able to accelerate. + +To see which functions and regions fell back, and why, enable Dagger's +`:reactant` debug category: + +```julia +push!(Dagger.DAGDEBUG_CATEGORIES, :reactant) +ENV["JULIA_DEBUG"] = "Dagger" +``` + +## Requiring Reactant + +Falling back is what lets one program run with and without Reactant, but it also +means that a workload can stop being accelerated - because a kernel was rewritten, +or because a worker is missing `using Reactant` - and only say so in a warning. +Code which cannot afford that can require Reactant instead: + +```julia +Dagger.@reactant must_opt=true must_load=true cholesky(A) +``` + +- `must_opt=true` turns a task or region that Reactant cannot compile or run into + a [`Dagger.ReactantOptimizationError`](@ref), which reports the failure Reactant + hit as its cause. A task's error surfaces when the task is `fetch`ed, as any + other task failure does; a region's error is thrown by the region itself, before + any of it has run. +- `must_load=true` turns Reactant.jl not being loaded into a + [`Dagger.ReactantUnavailableError`](@ref). Since the requirement travels with + the tasks, this covers workers which do not have Reactant loaded, and names the + worker which was missing it. + +Both requirements are useful in tests and benchmarks, where the whole point is +that Reactant is doing the work. + +## Compiled programs are cached + +Compiling an XLA program is expensive, so Dagger caches the programs it compiles +and reuses them whenever it is safe to: the same function, called with arguments +of the same types and sizes, and with the same non-array arguments, reuses the +program compiled for it earlier. Sizes participate because XLA compiles a program +for fixed shapes, which is also why a heterogeneously-blocked `DArray` will +compile a program per distinct block shape. + +Anything that Reactant bakes into a program as a constant has to match exactly for +a cached program to be reused, which is why a task function that *captures a +mutable object* is compiled afresh every time it is called: whatever was baked in +could have been written to since. This is worth knowing about, because it is the +difference between a task that compiles once and one that compiles on every call +(see [Performance](@ref) below). + +The cache is per-process, and can be inspected and emptied: + +```@docs +Dagger.reactant_cache_size +Dagger.reactant_cache_clear! +``` + +## Performance + +Reactant is not a uniform win, and which mode helps depends on what the +algorithm does. The three cases below are measured by + +``` +julia --project=test/reactantenv -t auto test/reactantenv/bench.jl medium +``` + +and these numbers are from a 12-thread x86 CPU with `1024x1024` matrices in +`256x256` blocks, as multiples of what Dagger does without Reactant (higher is +better): + +| Algorithm | `:inner` | `:full` | +|:--|--:|--:| +| `cholesky(::DMatrix)` | 0.17x | 0.19x | +| Blocked matmul, as a Datadeps region | 1.41x | 2.01x | +| `sum(sqrt.(abs.(A) .+ 1) .* 2)` over a `DArray` | 0.01x | 2.58x | + +The pattern behind them: + +- **Full mode is where the gains are.** Seeing a whole region lets Reactant fuse + across task boundaries: the blocked matmul becomes one large matrix multiply, + and the elementwise pipeline becomes a single kernel rather than one task per + block per operation. +- **Inner mode is limited by what it adds.** Each task copies its arguments to + Reactant and its results back, and each task's program is optimized alone, so + it only pays off when a task's own computation is large enough to absorb that. +- **XLA's CPU backend is not competitive with LAPACK for dense factorizations.** + Compiling `cholesky` of a 1024x1024 matrix with Reactant alone - no Dagger + involved - takes 0.18s against LAPACK's 0.023s, which is the whole of the + slowdown above. `A * B` for the same size is 0.024s against BLAS's 0.022s, so + matrix multiplication is at parity and the gains above come from fusion. This + is a property of the backend rather than of this integration, and a GPU backend + changes the picture entirely. +- **A task whose program cannot be cached is very slow.** The `0.01x` above is + that: the stage closures which `DArray` broadcasting and reductions build + capture the `DArray` itself, and a captured mutable object makes the program + ineligible for caching (see [Compiled programs are cached](@ref)), so every task + of every call compiles from scratch. Full mode compiles the region once, and so + is unaffected. + +## Distributed use + +The cache, and Reactant itself, live in the process that runs a task, so Reactant +must be loaded on every worker which will run tasks under `Dagger.@reactant`: + +```julia +using Distributed +addprocs(4) +@everywhere using Dagger, Reactant +``` + +Workers which do not have Reactant loaded run their tasks without it. + +## Limitations + +- Reactant only holds strided arrays of numbers (and `Bool`) as device buffers. + Other arguments are baked into the compiled program as constants, so in `:full` + mode a task which writes to a mutable non-array argument (a `Ref` used as a + scalar output, say) makes the region fall back to Datadeps. +- In `:inner` mode, each array argument becomes a buffer of its own, `view`s + included, and is copied back after the call. Tasks which write to overlapping + views of one array within a single call therefore see the last write win, rather + than both. +- In `:full` mode, a task's result may only be `fetch`ed after its region has + finished, not from within the region: the region's tasks have not run yet while + it is being recorded. Regions which do fetch their own tasks' results (such as + the symmetry check that `cholesky` performs) still work, as the check happens + after its own region. +- A task from a `:full` mode region may not be passed as an argument to a task + outside of any region, as it was completed outside of the scheduler. +- Non-positive-definiteness is not detected by `cholesky` under Reactant: the + check is a host-side branch on a value that only exists on the device. + +## API + +```@docs +Dagger.@reactant +Dagger.with_reactant +Dagger.reactant_available +Dagger.reactant_mode +Dagger.ReactantMode +Dagger.ReactantInner +Dagger.ReactantFull +Dagger.ReactantOptimizationError +Dagger.ReactantUnavailableError +``` diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl new file mode 100644 index 000000000..c1d6001e0 --- /dev/null +++ b/ext/ReactantExt.jl @@ -0,0 +1,787 @@ +module ReactantExt + +import Dagger +import Dagger: ReactantMode, ReactantInner, ReactantFull +import Dagger: ReactantOptimizationError +import Dagger: Chunk, DTask, DTaskPair, Processor, In, Out, InOut, Deps +import Dagger: ScopedValue, with +import Dagger: REACTANT_COMPILE_LOCK, REACTANT_COMPILE_CACHE + +import Adapt + +import MemPool + +import Reactant + +import LinearAlgebra + +function __init__() + Dagger.REACTANT_LOADED[] = true + return +end + +############################################################################# +# Conversion between plain Julia values and Reactant values +############################################################################# + +"Element types which Reactant can represent as device buffers." +const RElType = Union{Bool, + Int8, Int16, Int32, Int64, + UInt8, UInt16, UInt32, UInt64, + Float16, Float32, Float64, + ComplexF32, ComplexF64} + +""" + is_traceable(x) -> Bool + +Whether `x` is an array which should be handed to Reactant as a device buffer, +and thus traced. Everything else is left alone, which means Reactant treats it as +a constant of the program it compiles. +""" +function is_traceable(@nospecialize(x)) + x isa Reactant.RArray && return false + x isa AbstractArray || return false + eltype(x) <: RElType || return false + return x isa StridedArray +end + +to_reactant_array(x::StridedArray) = Reactant.to_rarray(x isa Array ? x : Array(x)) + +""" + ToReactant() + ToReactant(writebacks) + +An Adapt.jl adaptor which replaces the arrays within a value by Reactant arrays, +leaving everything Reactant cannot hold in a buffer alone. Adapt is what walks the +value, so arrays nested in tuples, named tuples, and other containers are +converted too; arrays themselves are converted whole, including `view`s and other +wrappers, rather than by adapting what they wrap. + +Given a `writebacks` vector, each converted array is recorded there alongside the +array it came from, which is what [`write_back!`](@ref) later uses to make writes +to those buffers visible again. +""" +struct ToReactant + writebacks::Union{Vector{Pair{Any,Any}},Nothing} + # An array passed twice must stay one buffer, or a task which writes to it + # through one argument would not see the write through the other + converted::IdDict{Any,Any} + + ToReactant(writebacks=nothing) = new(writebacks, IdDict{Any,Any}()) +end + +function convert_array(to::ToReactant, @nospecialize(x::AbstractArray)) + is_traceable(x) || return x + return get!(to.converted, x) do + converted = to_reactant_array(x) + if to.writebacks !== nothing + push!(to.writebacks, x => converted) + end + return converted + end +end + +Adapt.adapt_storage(to::ToReactant, @nospecialize(x::AbstractArray)) = convert_array(to, x) + +# Adapt would convert an array wrapper by rebuilding it around its converted +# parent, which is not what a task's argument should become: Datadeps runs tasks +# which write to disjoint `view`s of one array concurrently, and writing the whole +# parent back on behalf of each of them would lose all but one of their updates. +# Each array argument therefore becomes a buffer of its own, whatever it wraps. +for Wrapper in (SubArray, PermutedDimsArray, Base.ReshapedArray, Base.LogicalIndex, + Base.NonReshapedReinterpretArray, Base.ReshapedReinterpretArray, + LinearAlgebra.Adjoint, LinearAlgebra.Transpose, + LinearAlgebra.LowerTriangular, LinearAlgebra.UnitLowerTriangular, + LinearAlgebra.UpperTriangular, LinearAlgebra.UnitUpperTriangular, + LinearAlgebra.Diagonal, LinearAlgebra.Tridiagonal, LinearAlgebra.Symmetric) + @eval Adapt.adapt_structure(to::ToReactant, @nospecialize(x::$Wrapper)) = + convert_array(to, x) +end + +"An Adapt.jl adaptor which replaces Reactant values by plain Julia ones." +struct FromReactant end + +Adapt.adapt_storage(::FromReactant, x::Reactant.AbstractConcreteArray) = Array(x) +Adapt.adapt_storage(::FromReactant, x::Reactant.AbstractConcreteNumber{T}) where T = + convert(T, x) + +# Adapt treats an `Array` as a leaf, but a compiled program can return one whose +# elements are Reactant values: Dagger's reductions, for instance, wrap each +# partial result in a 1x1 array. Those elements have to be converted too, or a +# Reactant value ends up in what `sum(::DArray)` hands back to the caller. +function Adapt.adapt_structure(to::FromReactant, x::Array) + isbitstype(eltype(x)) && return x + return map(element -> Adapt.adapt(to, element), x) +end + +to_reactant(@nospecialize(x), adaptor::ToReactant=ToReactant()) = Adapt.adapt(adaptor, x) +from_reactant(@nospecialize(x)) = Adapt.adapt(FromReactant(), x) + +""" + write_back!(writebacks) + +Copies the contents of each Reactant array recorded in `writebacks` back into the +array it was converted from. Reactant writes the results of a mutating computation +into the buffers it was given, so this is what makes in-place task functions +(`BLAS.gemm!` and friends) observable to Dagger. +""" +function write_back!(writebacks::Vector{Pair{Any,Any}}) + for (orig, converted) in writebacks + copyto!(orig, Array(converted)) + end + return +end + +############################################################################# +# Compilation cache +############################################################################# + +""" + Uncacheable + +Marks a value which Reactant bakes into a compiled program in a way we cannot +summarize cheaply, making the enclosing call ineligible for caching. +""" +struct Uncacheable end + +""" + cache_key(x) -> key + +A summary of `x` which distinguishes any two values that Reactant would compile +differently. + +Whatever Reactant bakes into a compiled executable - constants, and the closure +that the resulting program carries - must be *identical* for a cached executable +to be valid, so those values participate in the key by value. Reactant arrays are +passed in at call time, so only their type and size matter; the size does matter, +because an executable is compiled for fixed shapes. + +A mutable value is summarized as [`Uncacheable`](@ref), since whatever was baked +into the program could since have been written to. That is what makes a task +function which captures a `DArray` - as the stages of a `DArray` broadcast do - +compile afresh on every call. +""" +cache_key(x::Reactant.AbstractConcreteArray) = (typeof(x), size(x)) +cache_key(x::Reactant.AbstractConcreteNumber) = typeof(x) +cache_key(x::Union{Number,Char,Symbol,AbstractString,Type,Nothing,Missing}) = x +cache_key(x::Tuple) = map(cache_key, x) +cache_key(x::NamedTuple) = (typeof(x), map(cache_key, values(x))) +function cache_key(x::T) where T + # `Array` and friends land here too, and are `Uncacheable` for being mutable: + # what Reactant baked into the program is their contents + (isstructtype(T) && !ismutabletype(T)) || return Uncacheable() + nfields = fieldcount(T) + nfields == 0 && return T + return (T, ntuple(idx -> isdefined(x, idx) ? cache_key(getfield(x, idx)) : Uncacheable(), + nfields)) +end + +is_cacheable(::Uncacheable) = false +is_cacheable(x::Tuple) = all(is_cacheable, x) +is_cacheable(@nospecialize(x)) = true + +function compile_program(f, args::Tuple, kwargs::NamedTuple) + Dagger.@dagdebug nothing :reactant "Compiling $(typeof(f)) for Reactant" + return Reactant.compile(f, args; fn_kwargs=kwargs, sync=true) +end + +""" + compiled_program(key, f, args, kwargs) -> program + +Returns a Reactant-compiled version of `f(args...; kwargs...)`, reusing a +previously compiled one when `key` says that it is safe to do so. A `key` which +is not [`is_cacheable`](@ref) forces a fresh compilation. +""" +function compiled_program(key, f, args::Tuple, kwargs::NamedTuple) + # Compilation happens under the lock even when its result cannot be cached, + # since what the lock is for is keeping Reactant from compiling on several + # threads at once + return Base.@lock REACTANT_COMPILE_LOCK begin + is_cacheable(key) || return compile_program(f, args, kwargs) + get!(REACTANT_COMPILE_CACHE, key) do + compile_program(f, args, kwargs) + end + end +end + +# Functions which Reactant was unable to compile, and which are therefore run +# directly from now on. Keyed by function type, as traceability is a property of +# the code a function runs rather than of the arguments it is called with. +const UNTRACEABLE = Set{Type}() + +is_untraceable(@nospecialize(F::Type)) = + Base.@lock REACTANT_COMPILE_LOCK (F in UNTRACEABLE) + +"How a task's function is named in the warnings and errors that mention it." +function describe_callable(@nospecialize(F::Type)) + if F <: Function && isdefined(F, :instance) + return "the task function `$(nameof(F.instance))`" + end + # A closure, or a callable object: there is no name worth printing, and the + # type at least says which captures it was compiled for + return "the task function of type `$F`" +end + +function mark_untraceable!(@nospecialize(F::Type), err, bt) + fresh = Base.@lock REACTANT_COMPILE_LOCK begin + F in UNTRACEABLE ? false : (push!(UNTRACEABLE, F); true) + end + fresh || return + # Only the error itself, as tracing failures tend to come with backtraces + # thousands of frames deep; the full report is available with the `:reactant` + # debug category enabled + @warn """Reactant could not compile $(describe_callable(F)), so it will be run without Reactant. + This is expected for task functions which access their arrays elementwise, such as the kernels of `Dagger.@stencil`. + Pass `must_opt=true` to make this an error instead.""" exception=err + Dagger.@dagdebug nothing :reactant "Tracing $F failed:\n$(sprint(showerror, err, bt))" + return +end + +############################################################################# +# Inner mode: compile each task's function individually +############################################################################# + +""" + ReactantCall(f, must_opt) + +Wraps a task's function so that it is compiled and executed by Reactant. Kept as +a callable, rather than calling Reactant directly from +`Dagger.reactant_execute!`, so that the call still goes through +`Dagger.execute!` and thus retains all of the processor-specific setup Dagger +normally performs around a task's function. +""" +struct ReactantCall{F} + f::F + must_opt::Bool +end + +function (rc::ReactantCall)(args...; kwargs...) + # A function known to be untraceable is not attempted again, unless the caller + # asked to be told about it + if !rc.must_opt && is_untraceable(typeof(rc.f)) + return rc.f(args...; kwargs...) + end + writebacks = Pair{Any,Any}[] + adaptor = ToReactant(writebacks) + rargs = to_reactant(args, adaptor) + rkwargs = to_reactant((; kwargs...), adaptor) + key = (cache_key(rc.f), cache_key(rargs), cache_key(rkwargs)) + # Tracing failures are recoverable: the arguments Reactant was given are + # copies, so nothing the task was passed has been touched yet + program, failure = try + compiled_program(key, rc.f, rargs, rkwargs), nothing + catch err + nothing, (err, catch_backtrace()) + end + if failure !== nothing + err, bt = failure + # Thrown from out here, rather than from the handler above, so that the + # exception Dagger reports is this one rather than its cause + rc.must_opt && throw(ReactantOptimizationError(describe_callable(typeof(rc.f)), err)) + mark_untraceable!(typeof(rc.f), err, bt) + return rc.f(args...; kwargs...) + end + result = program(rargs...) + # Reflect any writes the compiled program made to its arguments + write_back!(writebacks) + converted = from_reactant(result) + Dagger.@dagdebug nothing :reactant "Ran $(typeof(rc.f)), returning $(typeof(result)) as $(typeof(converted))" + return converted +end + +function Dagger.reactant_execute!(mode::ReactantInner, to_proc::Processor, f, args...; kwargs...) + @nospecialize f args kwargs + return Dagger.execute!(to_proc, ReactantCall(f, mode.must_opt), args...; kwargs...) +end + +############################################################################# +# Full mode: hand a whole Datadeps region to Reactant +############################################################################# + +# Set while a Datadeps region is being captured or traced, so that regions nested +# within it become part of the same program instead of starting their own +const FULL_ACTIVE = ScopedValue{Bool}(false) + +""" + CaptureQueue() + +A task queue which records tasks without launching them, used to observe the raw +algorithm of a Datadeps region: no planning, no scheduling, no data movement. +""" +struct CaptureQueue <: Dagger.AbstractTaskQueue + pairs::Vector{DTaskPair} + + CaptureQueue() = new(DTaskPair[]) +end +Dagger.enqueue!(queue::CaptureQueue, pair::DTaskPair) = push!(queue.pairs, pair) +Dagger.enqueue!(queue::CaptureQueue, pairs::Vector{DTaskPair}) = append!(queue.pairs, pairs) + +""" + RegionUnsupported(reason) + +Thrown when a Datadeps region cannot be expressed as a Reactant program, which +makes Dagger run the region itself instead (see [`fall_back!`](@ref)). +""" +struct RegionUnsupported <: Exception + reason::String +end +Base.showerror(io::IO, err::RegionUnsupported) = print(io, err.reason) + +# Where each argument of a captured task comes from when the region is traced +struct FromInput + idx::Int +end +struct FromTask + idx::Int +end +struct FromConstant + value::Any +end + +struct RegionCall + f::Any + args::Vector{Any} + kwargs::Vector{Pair{Symbol,Any}} +end + +""" + RegionProgram + +The raw algorithm of a Datadeps region: its tasks in submission order, with every +argument resolved to one of the region's inputs, the result of an earlier task, or +a constant. + +`sources` records where each input came from (a `Chunk`, or an array passed +directly to a task) so that results can be published back to it, and `data` holds +the input's data as pulled to the calling worker. +""" +struct RegionProgram + calls::Vector{RegionCall} + sources::Vector{Any} + data::Vector{Any} +end +RegionProgram() = RegionProgram(RegionCall[], Any[], Any[]) + +"Identity of a region input which isn't backed by a `Chunk`." +struct ObjectKey + id::UInt +end + +# Two `Chunk`s referring to the same `DRef` are the same input +input_key(chunk::Chunk) = chunk.handle +input_key(@nospecialize(x)) = ObjectKey(objectid(x)) + +pull_local(chunk::Chunk) = Dagger.move(Dagger.OSProc(), chunk) +pull_local(@nospecialize(x)) = x + +unwrap_dep(dep::In) = dep.x +unwrap_dep(dep::Out) = dep.x +unwrap_dep(dep::InOut) = dep.x +unwrap_dep(dep::Deps) = dep.x +unwrap_dep(@nospecialize(x)) = x + +is_function_argument(arg) = Dagger.ispositional(arg) && Dagger.raw_position(arg) == 0 + +""" + build_program(pairs) -> RegionProgram + +Turns the tasks captured from a Datadeps region into a program that can be traced: +their functions, in submission order, with each argument resolved to one of the +region's inputs, the result of an earlier task, or a constant. + +The `In`/`Out`/`InOut` annotations play no part in this. Datadeps guarantees that +a region behaves as if its tasks ran sequentially in submission order, which is +exactly what tracing them in that order produces; Reactant is then free to +recover the parallelism from the data flow it can see. +""" +function build_program(pairs::Vector{DTaskPair}) + program = RegionProgram() + input_indices = Dict{Any,Int}() + task_indices = IdDict{DTask,Int}() + + for pair in pairs + f = nothing + args = Any[] + kwargs = Pair{Symbol,Any}[] + for arg in pair.spec.fargs + value = Dagger.value(arg) + if is_function_argument(arg) + if value isa DTask + throw(RegionUnsupported("a task's function is itself the result of a task")) + end + f = value isa Chunk ? pull_local(value) : value + continue + end + source = describe_argument!(program, input_indices, task_indices, value) + if Dagger.ispositional(arg) + push!(args, source) + else + push!(kwargs, Dagger.pos_kw(arg) => source) + end + end + push!(program.calls, RegionCall(f, args, kwargs)) + task_indices[pair.task] = length(program.calls) + end + + return program +end + +function describe_argument!(program::RegionProgram, input_indices, task_indices, value) + value = unwrap_dep(value) + + if value isa DTask + idx = get(task_indices, value, nothing) + idx === nothing || return FromTask(idx) + # A task from outside this region, such as a `DArray` chunk which is + # still the task that produced it; its result is an input to the region + if !Base.istaskstarted(value) + throw(RegionUnsupported("a task argument was neither created within this region nor launched")) + end + value = fetch(value; raw=true) + end + + if value isa Chunk || is_traceable(value) + key = input_key(value) + idx = get(input_indices, key, nothing) + if idx === nothing + data = pull_local(value) + is_traceable(data) || return as_constant(data) + push!(program.sources, value) + push!(program.data, data) + idx = length(program.data) + input_indices[key] = idx + end + return FromInput(idx) + end + + return as_constant(value) +end + +# Anything Reactant cannot hold in a buffer is baked into the program instead, +# which is only sound if it cannot change: a mutable argument (a `Ref` used as a +# scalar output, say) would be written to while tracing and never again +function as_constant(@nospecialize(value)) + if ismutable(value) + throw(RegionUnsupported("a task takes a mutable $(typeof(value)), which Reactant cannot write to")) + end + return FromConstant(value) +end + +program_cache_key(program::RegionProgram) = + (:reactant_full, + Tuple(map(call_cache_key, program.calls)), + Tuple(map(data -> (typeof(data), size(data)), program.data))) +call_cache_key(call::RegionCall) = + (cache_key(call.f), + Tuple(map(argument_cache_key, call.args)), + Tuple(map(kwarg -> (first(kwarg), argument_cache_key(last(kwarg))), call.kwargs))) +argument_cache_key(arg::FromInput) = arg +argument_cache_key(arg::FromTask) = arg +argument_cache_key(arg::FromConstant) = cache_key(arg.value) + +""" + TracedTask(value) + +Stands in for a `DTask` while code is being traced by Reactant. Since tasks are +executed inline into the trace, the task's result is already available. +""" +struct TracedTask{T} + value::T +end +Base.fetch(task::TracedTask; kwargs...) = task.value +Base.wait(::TracedTask) = nothing +Base.isready(::TracedTask) = true + +# Where a task's result is to be found once the compiled program has run +struct ResultOutput + idx::Int # position in the program's return value +end +struct ResultInput + idx::Int # one of the region's inputs, which the program wrote in place +end +struct ResultConstant + value::Any # what the result evaluated to while tracing +end + +""" + RegionTrace(program) + +The bookkeeping that [`trace_region`](@ref) fills in while a region is traced: +where each task's result is to be found afterwards, and the traced values that +the compiled program must return for that to be possible. +""" +struct RegionTrace + program::RegionProgram + results::Vector{Any} + outputs::Vector{Any} +end +RegionTrace(program::RegionProgram) = + RegionTrace(program, Vector{Any}(undef, length(program.calls)), Any[]) + +"A Reactant-compiled Datadeps region, together with how to read its results." +struct CompiledRegion + program::Any + results::Vector{Any} +end +CompiledRegion(compiled, trace::RegionTrace) = CompiledRegion(compiled, trace.results) + +is_traced(@nospecialize(x)) = x isa Reactant.TracedRArray || x isa Reactant.TracedRNumber +contains_traced(@nospecialize(x)) = is_traced(x) +contains_traced(x::Union{Tuple,NamedTuple}) = any(contains_traced, x) + +# Values which a compiled program produces identically on every execution, and +# which can therefore be recorded once, while tracing +is_trace_constant(@nospecialize(x)) = + x isa Union{Nothing,Missing,Number,Char,Symbol,AbstractString,Type} +is_trace_constant(x::Union{Tuple,NamedTuple}) = all(is_trace_constant, x) + +""" + classify_result!(trace, result, input_ids) -> ResultOutput | ResultInput | ResultConstant + +Decides how `result`, produced by a task while tracing, will be recovered once the +compiled program has run. + +A result which is one of the program's own buffers - which is what the in-place +tasks of a Datadeps region typically return - is read back from that buffer. This +matters for more than tidiness: returning every task's result separately would +cost a buffer per task, which for an algorithm like a blocked Cholesky is far more +memory than the arrays being factored. + +A result which is not traced at all was baked into the program by Reactant, so it +is the same on every execution and can be recorded here. Anything else - a plain +array, say, which may well have been freshly computed - would not be, so it makes +the region unsupported. +""" +function classify_result!(trace::RegionTrace, @nospecialize(result), input_ids::IdDict{Any,Int}) + if contains_traced(result) + idx = is_traced(result) ? get(input_ids, result, 0) : 0 + idx == 0 || return ResultInput(idx) + push!(trace.outputs, result) + return ResultOutput(length(trace.outputs)) + elseif is_trace_constant(result) + return ResultConstant(result) + end + throw(RegionUnsupported("a task returned a $(typeof(result)), which Reactant cannot return from a compiled program")) +end + +# The trace currently being built. Read at trace time only, which is what lets +# `trace_region` below be a plain function: were the region carried in a closure +# instead, Reactant would trace through it (and through every `Chunk` and `DArray` +# that the region references) on every compilation. +const CURRENT_TRACE = ScopedValue{Union{RegionTrace,Nothing}}(nothing) + +resolve_argument(arg::FromInput, inputs, results) = inputs[arg.idx] +resolve_argument(arg::FromTask, inputs, results) = results[arg.idx] +resolve_argument(arg::FromConstant, inputs, results) = arg.value + +function trace_region(inputs::Vararg{Any,N}) where N + trace = CURRENT_TRACE[]::RegionTrace + program = trace.program + + # Reactant may trace a program more than once, so start from a clean slate + empty!(trace.outputs) + input_ids = IdDict{Any,Int}() + for (idx, input) in enumerate(inputs) + input_ids[input] = idx + end + + results = Vector{Any}(undef, length(program.calls)) + for (idx, call) in enumerate(program.calls) + args = Any[resolve_argument(arg, inputs, results) for arg in call.args] + kwargs = NamedTuple(key => resolve_argument(arg, inputs, results) + for (key, arg) in call.kwargs) + result = call.f(args...; kwargs...) + results[idx] = result + # Note where each task's result will be found, so that the task it came + # from can be completed with it and thus be `fetch`ed as usual + trace.results[idx] = classify_result!(trace, result, input_ids) + end + return Tuple(trace.outputs) +end + +""" + compiled_region(program, inputs) -> CompiledRegion + +Compiles `program` for `inputs`, reusing a previously compiled region when it is +safe to do so (see [`cache_key`](@ref)). +""" +function compiled_region(program::RegionProgram, inputs::Tuple) + trace = RegionTrace(program) + compile() = with(CURRENT_TRACE => trace) do + CompiledRegion(compile_program(trace_region, inputs, NamedTuple()), trace) + end + key = program_cache_key(program) + return Base.@lock REACTANT_COMPILE_LOCK begin + is_cacheable(key) || return compile() + get!(compile, REACTANT_COMPILE_CACHE, key)::CompiledRegion + end +end + +function Dagger.reactant_spawn_datadeps(mode::ReactantFull, f) + # A nested region is already part of the enclosing region's program + FULL_ACTIVE[] && return f() + return with(FULL_ACTIVE => true) do + run_full_region(f, mode.must_opt) + end +end + +function run_full_region(f, must_opt::Bool) + # Capture the region's raw algorithm, without planning or scheduling it + queue = CaptureQueue() + result = Dagger.with_options(f; task_queue=queue) + pairs = queue.pairs + isempty(pairs) && return result + + # Nothing below touches the region's own data until the program has run to + # completion (Reactant works on copies of it), so any failure along the way + # can still be answered by handing the region back to Dagger + local region, outputs, updated + failure = nothing + try + program = build_program(pairs) + Dagger.@dagdebug nothing :reactant "Tracing $(length(program.calls)) task(s) over $(length(program.data)) input(s)" + inputs = ntuple(idx -> to_reactant_array(program.data[idx]), length(program.data)) + region = compiled_region(program, inputs) + outputs = region.program(inputs...) + # Publish the results back to where the region's arguments live + updated = map(Array, inputs) + for idx in 1:length(program.sources) + writeback_input!(program.sources[idx], updated[idx]) + end + catch err + failure = (err, catch_backtrace()) + end + if failure !== nothing + err, bt = failure + # Thrown from out here, rather than from the handler above, so that the + # exception Dagger reports is this one rather than its cause + must_opt && throw(ReactantOptimizationError("a Datadeps region of $(length(pairs)) task(s)", err)) + fall_back!(pairs, err, bt) + return result + end + + complete_tasks!(pairs, region, outputs, updated) + + return result +end + +# Make each captured task's result available, as the scheduler would have +function complete_tasks!(pairs::Vector{DTaskPair}, region::CompiledRegion, + outputs::Tuple, updated::Tuple) + for (idx, pair) in enumerate(pairs) + value = result_value(region.results[idx], outputs, updated) + Dagger.complete_unlaunched!(pair.task, Dagger.tochunk(value)) + end + return +end + +result_value(result::ResultOutput, outputs, updated) = from_reactant(outputs[result.idx]) +result_value(result::ResultInput, outputs, updated) = updated[result.idx] +result_value(result::ResultConstant, outputs, updated) = result.value + +""" + fall_back!(pairs, err, bt) + +Runs the region made up of `pairs` through Datadeps, because Reactant could not +run it (`err`). Regions differ widely in what they ask of Reactant, and a region +it cannot handle is not a reason to fail: the same code should keep working, just +without Reactant. +""" +function fall_back!(pairs::Vector{DTaskPair}, err, bt) + @warn """Reactant could not run a Datadeps region, so it will be run by Dagger instead. + Enable the `:reactant` debug category for the regions and errors involved, or pass `must_opt=true` to make this an error.""" exception=err maxlog=1 + Dagger.@dagdebug nothing :reactant "Region of $(length(pairs)) task(s) failed:\n$(sprint(showerror, err, bt))" + Dagger.launch_datadeps_tasks!(pairs) + return +end + +function writeback_input!(source::Chunk, updated) + MemPool.access_ref(source.handle, updated) do stored, updated + stored === updated || copyto!(stored, updated) + return + end + return +end +writeback_input!(source, updated) = (copyto!(source, updated); nothing) + +############################################################################# +# Overlays: how Dagger's API behaves within Reactant-traced code +############################################################################# + +# Traced code has no scheduler to submit to and no worker to run on, so a spawned +# task is simply executed inline, becoming part of the program being traced. +Reactant.@reactant_overlay function Dagger.spawn(f, args...; kwargs...) + return traced_spawn(f, args, kwargs) +end +Reactant.@reactant_overlay function Dagger.typed_spawn(f, args...; kwargs...) + return traced_spawn(f, args, kwargs) +end + +function traced_spawn(f, args, kwargs) + if length(args) >= 1 && first(args) isa Dagger.Options + args = args[2:end] + end + new_args = map(traced_argument, args) + new_kwargs = NamedTuple(key => traced_argument(value) for (key, value) in kwargs) + return TracedTask(f(new_args...; new_kwargs...)) +end + +function traced_argument(arg) + arg = unwrap_dep(arg) + arg isa TracedTask && return arg.value + arg isa Chunk && return pull_local(arg) + return arg +end + +# `task_processor` is defined in terms of a running `DTask`, which traced code +# need not be; report the processor that the trace is being built on instead. +Reactant.@reactant_overlay Dagger.task_processor() = + Dagger.in_task() ? Dagger.get_tls().processor : + Dagger.ThreadProc(Dagger.myid(), Threads.threadid()) + +############################################################################# +# Kernels which Reactant cannot trace as written +############################################################################# + +# `LAPACK.potrf!` is a `ccall` into LAPACK, which Reactant cannot trace, so +# Dagger's Cholesky panel factorization gets a traced implementation here. +# +# `info` is returned as a plain `0`, rather than a traced value, so that the +# positive-definiteness check in `potrf_checked!` remains a trace-time branch; +# the consequence is that a non-positive-definite matrix is not detected under +# Reactant. XLA's Cholesky also zeroes the triangle that `potrf!` would have left +# untouched, which is unobservable through the `Cholesky` factorization object +# that Dagger returns. +function Dagger.potrf_checked!(uplo, A::Reactant.AnyTracedRArray{T,2}, info_arr) where T + lower = is_lower(uplo) + factors = Reactant.Ops.cholesky(Reactant.TracedUtils.materialize_traced_array(A); lower) + copyto!(A, factors) + return A, 0 +end + +is_lower(uplo::AbstractChar) = uplo == 'L' || uplo == 'l' + +# Reactant lowers `BLAS.syrk!` to an `enzymexla.blas_syrk` op which, as of +# Reactant v0.2.279, interprets `uplo` in row-major terms and so updates the +# opposite triangle of `C` from what BLAS does. Dagger's Cholesky relies on +# `syrk!` for its trailing updates, so compute it here from operations whose +# behavior is unambiguous - the same way Reactant itself implements `syr2k!` and +# `herk!`. +Reactant.@reactant_overlay function LinearAlgebra.BLAS.syrk!(uplo::AbstractChar, + trans::AbstractChar, + alpha::Number, + A::Reactant.AnyTracedRMatrix, + beta::Number, + C::Reactant.AnyTracedRMatrix) + A = Reactant.TracedUtils.materialize_traced_array(A) + product = trans == 'N' || trans == 'n' ? A * transpose(A) : transpose(A) * A + updated = alpha .* product .+ beta .* C + if is_lower(uplo) + LinearAlgebra.LowerTriangular(C) .= LinearAlgebra.LowerTriangular(updated) + else + LinearAlgebra.UpperTriangular(C) .= LinearAlgebra.UpperTriangular(updated) + end + return C +end + +end # module diff --git a/src/Dagger.jl b/src/Dagger.jl index 3ef2c59d5..8f483d412 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -75,6 +75,7 @@ include("chunks.jl") include("utils/signature.jl") include("thunkid.jl") include("utils/lfucache.jl") +include("reactant.jl") include("options.jl") include("dtask.jl") include("cancellation.jl") diff --git a/src/datadeps/queue.jl b/src/datadeps/queue.jl index 86f11acbf..2214bd99d 100644 --- a/src/datadeps/queue.jl +++ b/src/datadeps/queue.jl @@ -76,6 +76,40 @@ function spawn_datadeps(f::Base.Callable; static::Bool=true, if !aliasing throw(ArgumentError("Aliasing analysis is no longer optional")) end + + # Under `Dagger.@reactant mode=:full`, the region's raw algorithm is handed + # to Reactant instead of being planned and scheduled by Datadeps + reactant_mode = get_options(:reactant, nothing) + if reactant_mode isa ReactantFull + return reactant_spawn_datadeps(reactant_mode, f) + end + + return _spawn_datadeps(f; scheduler, launch_wait, hierarchical) +end +_spawn_datadeps(f::Base.Callable; kwargs...) = + _datadeps_region(queue -> with_options(f; task_queue=queue); kwargs...) + +""" + launch_datadeps_tasks!(pairs::Vector{DTaskPair}) -> nothing + +Plans and launches `pairs`, waiting for them to finish, as if they had been +submitted to a Datadeps region directly. This is for tasks which were captured +from a region without being launched: [`Dagger.@reactant`](@ref)'s `:full` mode +does that in order to hand the region to Reactant, and needs to fall back to +Datadeps for the regions that Reactant turns out to be unable to run. +""" +function launch_datadeps_tasks!(pairs::Vector{DTaskPair}; kwargs...) + _datadeps_region(queue -> enqueue!(queue, pairs); kwargs...) + return +end + +# Runs one Datadeps region: `fill_queue!` submits the region's tasks to the +# queue it is given (and its return value becomes the region's result), then the +# whole set is planned, launched, and waited on. +function _datadeps_region(fill_queue!; + scheduler::Union{DataDepsScheduler,Nothing}=nothing, + launch_wait::Union{Bool,Nothing}=nothing, + hierarchical::Union{Bool,Nothing}=nothing) wait_all(; check_errors=true) do scheduler = something(scheduler, DATADEPS_SCHEDULER[], RoundRobinScheduler()) launch_wait = something(launch_wait, DATADEPS_LAUNCH_WAIT[], false)::Bool @@ -93,12 +127,12 @@ function spawn_datadeps(f::Base.Callable; static::Bool=true, if launch_wait result = spawn_bulk() do queue = DataDepsTaskQueue(get_options(:task_queue); scheduler) - with_options(f; task_queue=queue) + fill_queue!(queue) run_distribute(queue) end else queue = DataDepsTaskQueue(get_options(:task_queue); scheduler) - result = with_options(f; task_queue=queue) + result = fill_queue!(queue) run_distribute(queue) end return result diff --git a/src/dtask.jl b/src/dtask.jl index 3bdefb543..c0a207a82 100644 --- a/src/dtask.jl +++ b/src/dtask.jl @@ -127,6 +127,45 @@ function Base.show(io::IO, t::DTask) print(io, "DTask ($status)") end istask(t::DTask) = true + +# Stands in for the `Thunk` of a task completed by `complete_unlaunched!`, which +# has none. Allocated once, on first use, and shared by all such tasks. +const UNLAUNCHED_THUNK_REF = Ref{Union{DRef,Nothing}}(nothing) +const UNLAUNCHED_THUNK_REF_LOCK = Threads.ReentrantLock() +function unlaunched_thunk_ref() + Base.@lock UNLAUNCHED_THUNK_REF_LOCK begin + ref = UNLAUNCHED_THUNK_REF[] + ref === nothing || return ref + ref = poolset(nothing; size=0, device=MemPool.CPURAMDevice()) + UNLAUNCHED_THUNK_REF[] = ref + return ref + end +end + +""" + Dagger.complete_unlaunched!(task::DTask, value; error::Bool=false) -> DTask + +Completes `task` with `value`, making it behave like a finished task: it may be +`fetch`ed and `wait`ed on, and reports itself as done. + +This exists for `task`s which were taken from a task queue and executed outside +of Dagger's scheduler, as [`Dagger.@reactant`](@ref)'s `:full` mode does with the +tasks of a Datadeps region. Because such a task has no `Thunk` behind it, it may +not be passed as an argument to a task which *is* run by the scheduler. + +`value` should be a `Chunk` (as produced by [`Dagger.tochunk`](@ref)), matching +what the scheduler would have stored; `error=true` marks `value` as the +exception that the task failed with. +""" +function complete_unlaunched!(task::DTask, value; error::Bool=false) + if istaskstarted(task) + throw(ConcurrencyViolationError("Cannot complete a launched `DTask`")) + end + task.thunk_ref = unlaunched_thunk_ref() + put!(task.future, value; error) + return task +end + function Base.convert(::Type{ThunkSyncdep}, task::Dagger.DTask) return ThunkSyncdep(ThunkID(task.uid, isdefined(task, :thunk_ref) ? task.thunk_ref : nothing)) end diff --git a/src/options.jl b/src/options.jl index aabd20002..a74f67d3b 100644 --- a/src/options.jl +++ b/src/options.jl @@ -33,6 +33,7 @@ Stores per-task options to be passed to the scheduler. - `stream_max_evals::Union{Int,Nothing}=nothing`: (Streaming only) Specifies the maximum number of times the task will be evaluated before returning a result. Defaults to infinite evaluations. - `acceleration::Union{Acceleration,Nothing}=nothing`: The acceleration backend used to plan and execute this task (e.g. `DistributedAcceleration`, `MPIAcceleration`). When `nothing`, the current acceleration (`Dagger.current_acceleration()`) is used. - `return_type::Union{Type,Nothing}=nothing`: The expected return type of the task's function. When set to a concrete type, it is used as the task's `chunktype` before the task has run (e.g. so downstream metadata and, under MPI, cross-rank type uniformity are known ahead of execution). When `nothing`, the type is left unknown until the result is available. +- `reactant::Union{ReactantMode,Nothing}=nothing`: If not `nothing`, requests that this task (and, as this option propagates, the tasks it spawns) be executed through Reactant.jl in the given mode. Usually set via [`Dagger.@reactant`](@ref) rather than directly. """ Base.@kwdef mutable struct Options propagates::Union{Vector{Symbol},Nothing} = nothing @@ -74,6 +75,8 @@ Base.@kwdef mutable struct Options acceleration::Union{Acceleration,Nothing} = nothing return_type::Union{Type,Nothing} = nothing + + reactant::Union{ReactantMode,Nothing} = nothing end Options(::Nothing) = Options() function Options(old_options::NamedTuple) diff --git a/src/reactant.jl b/src/reactant.jl new file mode 100644 index 000000000..795c2f212 --- /dev/null +++ b/src/reactant.jl @@ -0,0 +1,284 @@ +# Reactant.jl integration +# +# Dagger can hand the code within a task, or the algorithm of a whole Datadeps +# region, to Reactant.jl, which compiles it through MLIR/XLA and applies +# optimizations (fusion, layout selection, heterogeneous code generation) that +# Julia's own compiler does not perform. +# +# The user-facing entrypoint is `Dagger.@reactant`, which selects a mode and +# makes it visible to the rest of Dagger through the scoped `reactant` option. +# All of the Reactant-specific logic lives in the ReactantExt extension; the +# functions below are the hooks that it specializes, and their fallbacks run the +# code as usual (after warning once) when Reactant is not loaded. + +""" + Dagger.ReactantMode + +Selects how Dagger hands work to Reactant.jl. See [`Dagger.ReactantInner`](@ref) +and [`Dagger.ReactantFull`](@ref). + +Every mode carries the two requirements that [`Dagger.@reactant`](@ref) can +impose: `must_opt`, which turns a failure to compile into an error rather than +running the code without Reactant, and `must_load`, which does the same for +Reactant not being loaded. They travel with the mode, so they apply on every +worker that the mode reaches. +""" +abstract type ReactantMode end + +""" + Dagger.ReactantInner(; must_opt=false, must_load=false) + +Compile each task's function with Reactant, separately, on the processor that +runs it. The task's arguments are converted to Reactant arrays before the call +and converted back (including writes to mutated arguments) afterwards, so +Dagger's scheduling, data movement, and dependency handling are unchanged. + +This is the default mode of [`Dagger.@reactant`](@ref), and the most reliable, +as each task is optimized in isolation. A task whose function Reactant cannot +compile is run without Reactant, unless `must_opt` is set. +""" +Base.@kwdef struct ReactantInner <: ReactantMode + must_opt::Bool = false + must_load::Bool = false +end + +""" + Dagger.ReactantFull(; must_opt=false, must_load=false) + +Hand the entire algorithm of a [`spawn_datadeps`](@ref) region to Reactant as a +single traced program. Dagger performs no planning or scheduling for the region: +all arguments are pulled to the calling worker, the region's tasks are traced +in submission order, and Reactant is responsible for optimizing and executing +the resulting DAG (which allows it to fuse and reorder across task boundaries). + +Results are written back into the region's arguments once execution finishes, and +each of the region's tasks is completed with its result, so the region remains +observationally equivalent to running it under Datadeps. A region which Reactant +cannot compile or run is handed back to Datadeps, which runs it as usual, unless +`must_opt` is set. + +Tasks launched outside of a Datadeps region are unaffected by this mode. +""" +Base.@kwdef struct ReactantFull <: ReactantMode + must_opt::Bool = false + must_load::Bool = false +end + +""" + Dagger.reactant_mode(mode; must_opt=false, must_load=false) -> ReactantMode + +The [`Dagger.ReactantMode`](@ref) that `mode` names: `:inner`, `:full`, or a mode +itself. Requirements given here are added to those the mode already carries. +""" +reactant_mode(mode::ReactantMode; must_opt::Bool=false, must_load::Bool=false) = + typeof(mode)(must_opt || mode.must_opt, must_load || mode.must_load) +function reactant_mode(mode::Symbol; must_opt::Bool=false, must_load::Bool=false) + if mode === :inner + return ReactantInner(; must_opt, must_load) + elseif mode === :full + return ReactantFull(; must_opt, must_load) + end + throw(ArgumentError("Invalid Reactant mode: $(repr(mode))\nValid modes are :inner and :full")) +end +reactant_mode(mode; must_opt::Bool=false, must_load::Bool=false) = + throw(ArgumentError("Invalid Reactant mode: $(repr(mode))\nExpected a Symbol (:inner or :full) or a Dagger.ReactantMode")) + +""" + Dagger.ReactantUnavailableError(worker) + +Thrown when Reactant.jl is not loaded in the process that needs it, and +`Dagger.@reactant`'s `must_load=true` was used to ask that this be an error rather +than a warning. +""" +struct ReactantUnavailableError <: Exception + worker::Int +end +ReactantUnavailableError() = ReactantUnavailableError(myid()) +Base.showerror(io::IO, err::ReactantUnavailableError) = + print(io, """ReactantUnavailableError: Reactant.jl is not loaded on worker $(err.worker), and `must_load=true` was requested. + Add `using Reactant` (on every worker which will run tasks) to enable Reactant-accelerated execution.""") + +""" + Dagger.ReactantOptimizationError(what, cause) + +Thrown when Reactant could not compile or run `what`, and `Dagger.@reactant`'s +`must_opt=true` was used to ask that this be an error rather than a fall back to +running the code without Reactant. `cause` is the failure that Reactant reported. +""" +struct ReactantOptimizationError <: Exception + what::String + cause::Any +end +function Base.showerror(io::IO, err::ReactantOptimizationError) + print(io, "ReactantOptimizationError: Reactant could not optimize $(err.what), and `must_opt=true` was requested.\nCaused by: ") + showerror(io, err.cause) + return +end + +"Set by ReactantExt when Reactant.jl is loaded into this process." +const REACTANT_LOADED = Ref(false) + +""" + Dagger.reactant_available() -> Bool + +Returns `true` if Reactant.jl is loaded in this process, and thus if +[`Dagger.@reactant`](@ref) will actually use Reactant. +""" +reactant_available() = REACTANT_LOADED[] + +function warn_reactant_unavailable() + @warn """Dagger.@reactant was used, but Reactant.jl is not loaded in this process; running without Reactant. + Add `using Reactant` (on every worker which will run tasks) to enable Reactant-accelerated execution, or `must_load=true` to make this an error.""" maxlog=1 + return +end + +""" + Dagger.@reactant expr + Dagger.@reactant mode=:inner expr + Dagger.@reactant mode=:full expr + Dagger.@reactant must_opt=true expr + Dagger.@reactant must_load=true expr + +Executes `expr` with Reactant.jl integration enabled, so that Dagger tasks +launched by `expr` (including those launched by library code, such as the tasks +of `cholesky(::DArray)`) are compiled and executed by Reactant. + +Two modes are available: + +- `mode=:inner` (the default): each task's function is compiled by Reactant + individually, on the processor that runs it. Dagger's scheduling and data + movement are unchanged. See [`Dagger.ReactantInner`](@ref). +- `mode=:full`: each [`spawn_datadeps`](@ref) region within `expr` is traced as + a whole and handed to Reactant as a single program, bypassing Dagger's + planning and scheduling for that region. See [`Dagger.ReactantFull`](@ref). + +By default, whatever Reactant cannot do is done without it: if Reactant.jl is not +loaded, a warning is emitted once per session and `expr` runs as it normally +would, and the same goes for individual tasks and regions which Reactant turns +out to be unable to compile. This is what allows the same application code to be +used with and without Reactant, but it also means that a workload can quietly +stop being accelerated. Two options make those cases loud instead: + +- `must_opt=true`: a task or region which Reactant cannot compile or run throws a + [`Dagger.ReactantOptimizationError`](@ref) instead of running without Reactant. +- `must_load=true`: Reactant.jl not being loaded (here, or on a worker that runs + one of the tasks) throws a [`Dagger.ReactantUnavailableError`](@ref) instead of + warning. + +# Examples + +```julia +using Dagger, Reactant + +A = rand(Blocks(32, 32), 128, 128) +A = A * A' + 128I +chol = Dagger.@reactant cholesky(A) + +# Fail rather than silently run the factorization without Reactant +chol = Dagger.@reactant must_opt=true must_load=true cholesky(A) +``` +""" +macro reactant(exs...) + if isempty(exs) + throw(ArgumentError("@reactant requires an expression to execute")) + end + inner_ex = last(exs) + mode_ex = QuoteNode(:inner) + must_opt_ex = false + must_load_ex = false + for opt in exs[1:end-1] + if !(Meta.isexpr(opt, :(=)) && length(opt.args) == 2) + throw(ArgumentError("@reactant: invalid option `$opt` (expected `name=value`)")) + end + name, value_ex = opt.args + if name === :mode + mode_ex = value_ex + elseif name === :must_opt + must_opt_ex = value_ex + elseif name === :must_load + must_load_ex = value_ex + else + throw(ArgumentError("@reactant: unknown option `$name` (valid options are `mode`, `must_opt`, and `must_load`)")) + end + end + return quote + $with_reactant($reactant_mode($(esc(mode_ex)); + must_opt=$(esc(must_opt_ex)), + must_load=$(esc(must_load_ex)))) do + $(esc(inner_ex)) + end + end +end + +""" + Dagger.with_reactant(f, mode::ReactantMode) -> Any + Dagger.with_reactant(f, mode; must_opt=false, must_load=false) -> Any + +Calls `f()` with Reactant integration enabled in `mode`. This is the function +form of [`Dagger.@reactant`](@ref). +""" +function with_reactant(f, mode::ReactantMode) + if !reactant_available() + mode.must_load && throw(ReactantUnavailableError()) + warn_reactant_unavailable() + return f() + end + return with_options(f; reactant=mode) +end +with_reactant(f, mode; kwargs...) = with_reactant(f, reactant_mode(mode; kwargs...)) + +""" + Dagger.reactant_execute!(mode::ReactantMode, to_proc, f, args...; kwargs...) + +Executes `f(args...; kwargs...)` on `to_proc` through Reactant, as requested by +`mode`. Specialized by the ReactantExt extension; the fallback here runs the +call without Reactant, as this process does not have Reactant loaded. +""" +function reactant_execute!(mode::ReactantMode, to_proc::Processor, f, args...; kwargs...) + @nospecialize f args kwargs + mode.must_load && throw(ReactantUnavailableError()) + warn_reactant_unavailable() + return execute!(to_proc, f, args...; kwargs...) +end + +""" + Dagger.reactant_spawn_datadeps(mode::ReactantMode, f) -> Any + +Executes the Datadeps region `f` through Reactant, as requested by `mode`. +Specialized by the ReactantExt extension; the fallback here runs the region +through Datadeps as usual, as this process does not have Reactant loaded. +""" +function reactant_spawn_datadeps(mode::ReactantMode, f) + mode.must_load && throw(ReactantUnavailableError()) + warn_reactant_unavailable() + return _spawn_datadeps(f) +end + +# Cache of Reactant-compiled programs, filled in by ReactantExt. +# +# N.B. Compilation is serialized under this lock: it is expensive, not +# guaranteed to be thread-safe, and Dagger will happily run many tasks +# concurrently. Executing an already-compiled program happens outside the lock. +const REACTANT_COMPILE_LOCK = Threads.ReentrantLock() +const REACTANT_COMPILE_CACHE = Dict{Any,Any}() + +""" + Dagger.reactant_cache_clear!() + +Discards every Reactant program that Dagger has compiled and cached in this +process. Compiled programs are reused across tasks and calls whenever it is safe +to do so, so this is mostly useful for benchmarking compilation, or to release +the memory that they hold. +""" +function reactant_cache_clear!() + Base.@lock REACTANT_COMPILE_LOCK empty!(REACTANT_COMPILE_CACHE) + return +end + +""" + Dagger.reactant_cache_size() -> Int + +The number of Reactant programs that Dagger currently has cached in this process. +""" +reactant_cache_size() = + Base.@lock REACTANT_COMPILE_LOCK length(REACTANT_COMPILE_CACHE) diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index 545f34497..9c0cefd5b 100644 --- a/src/sch/Sch.jl +++ b/src/sch/Sch.jl @@ -2001,9 +2001,14 @@ Executes a single task specified by `task` on `to_proc`. acceleration=Dagger.current_acceleration(), )) + reactant_mode = options.reactant result = Dagger.with_options(propagated) do # Execute - execute!(to_proc, f, fetched_args...; fetched_kwargs...) + if reactant_mode isa Dagger.ReactantInner + Dagger.reactant_execute!(reactant_mode, to_proc, f, fetched_args...; fetched_kwargs...) + else + execute!(to_proc, f, fetched_args...; fetched_kwargs...) + end end # Check if result is safe to store diff --git a/src/utils/dagdebug.jl b/src/utils/dagdebug.jl index ed4749d82..d3b084227 100644 --- a/src/utils/dagdebug.jl +++ b/src/utils/dagdebug.jl @@ -14,7 +14,7 @@ function task_id end # a user has explicitly asked for tracing. const DAGDEBUG_VALID_CATEGORIES = (:all, :global, :submit, :schedule, :scope, :take, :execute, :move, :processor, :finish, - :cancel, :stream, :validate) + :cancel, :stream, :validate, :reactant) const DAGDEBUG_CATEGORIES = Set{Symbol}() # Out-of-line emission keeps call-site IR minimal: just one `in` check + one diff --git a/test/reactant.jl b/test/reactant.jl new file mode 100644 index 000000000..bdc03dfb0 --- /dev/null +++ b/test/reactant.jl @@ -0,0 +1,83 @@ +# Reactant integration, as seen by a session which has *not* loaded Reactant.jl. +# The tests which actually exercise Reactant live in `test/reactantenv`, which has +# its own project, since Reactant is far too heavy to depend on here. + +import Logging + +@testset "Mode selection" begin + @test Dagger.reactant_mode(:inner) === Dagger.ReactantInner() + @test Dagger.reactant_mode(:full) === Dagger.ReactantFull() + @test Dagger.reactant_mode(Dagger.ReactantInner()) === Dagger.ReactantInner() + @test_throws ArgumentError Dagger.reactant_mode(:bogus) + @test_throws ArgumentError Dagger.reactant_mode(42) + + # Requirements are carried by the mode, and are additive + @test Dagger.reactant_mode(:inner; must_opt=true) === Dagger.ReactantInner(; must_opt=true) + @test Dagger.reactant_mode(:full; must_load=true) === Dagger.ReactantFull(; must_load=true) + @test Dagger.reactant_mode(Dagger.ReactantFull(; must_opt=true); must_load=true) === + Dagger.ReactantFull(; must_opt=true, must_load=true) + + @test_throws ArgumentError macroexpand(@__MODULE__, :(Dagger.@reactant bogus=true nothing)) + @test_throws ArgumentError macroexpand(@__MODULE__, :(Dagger.@reactant)) +end + +@testset "Running without Reactant" begin + @test !Dagger.reactant_available() + + A = rand(8, 8) + B = rand(8, 8) + + # The same application code must run with and without Reactant, so a missing + # Reactant is a warning rather than an error, and is only reported once + logs, C = Test.collect_test_logs(min_level=Logging.Warn) do + Dagger.@reactant fetch(Dagger.@spawn A * B) + Dagger.@reactant mode=:inner fetch(Dagger.@spawn A * B) + end + @test C ≈ A * B + @test length(logs) == 1 + @test occursin(r"Reactant.jl is not loaded", logs[1].message) + + # `:full` falls back to Datadeps executing the region as usual + C = zeros(8, 8) + Dagger.@reactant mode=:full begin + Dagger.spawn_datadeps() do + Dagger.@spawn mul!(Out(C), In(A), In(B)) + end + end + @test C ≈ A * B + + @test Dagger.with_reactant(() -> Dagger.get_options(:reactant, nothing), :inner) === nothing +end + +@testset "must_load" begin + A = rand(8, 8) + + # Code which cannot afford to silently run without Reactant asks to be told + @test_throws Dagger.ReactantUnavailableError begin + Dagger.@reactant must_load=true fetch(Dagger.@spawn sum(A)) + end + @test_throws Dagger.ReactantUnavailableError begin + Dagger.@reactant mode=:full must_load=true begin + Dagger.spawn_datadeps() do + Dagger.@spawn sum(In(A)) + end + end + end + + err = Dagger.ReactantUnavailableError(2) + @test occursin("worker 2", sprint(showerror, err)) + + # `must_opt` has nothing to act on without Reactant, so it is `must_load`'s job + # alone to notice that Reactant is missing + @test (@test_logs (:warn, r"Reactant.jl is not loaded") match_mode=:any begin + Dagger.@reactant must_opt=true fetch(Dagger.@spawn sum(A)) + end) ≈ sum(A) +end + +@testset "Compilation cache" begin + # Nothing can have been compiled without Reactant, but the cache is still + # queryable, as application code may report on it + @test Dagger.reactant_cache_size() == 0 + Dagger.reactant_cache_clear!() + @test Dagger.reactant_cache_size() == 0 +end diff --git a/test/reactantenv/Project.toml b/test/reactantenv/Project.toml new file mode 100644 index 000000000..646651b25 --- /dev/null +++ b/test/reactantenv/Project.toml @@ -0,0 +1,14 @@ +[deps] +Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +Logging = "56ddb016-857b-54e1-b83d-db4d58db5568" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +TimespanLogging = "a526e669-04d3-4846-9525-c66122c55f63" + +[sources] +Dagger = {path = "../.."} + +[compat] +Reactant = "0.2.279" diff --git a/test/reactantenv/bench.jl b/test/reactantenv/bench.jl new file mode 100644 index 000000000..ca7656f77 --- /dev/null +++ b/test/reactantenv/bench.jl @@ -0,0 +1,196 @@ +# Does Reactant actually make Dagger faster? This compares each algorithm as +# Dagger runs it, as `Dagger.@reactant mode=:inner` runs it, and as +# `Dagger.@reactant mode=:full` runs it. +# +# Run with: +# julia --project=test/reactantenv -t auto test/reactantenv/bench.jl [size] +# +# where `size` is `small` (the default), `medium`, or `large`. + +using Dagger +using Reactant +using LinearAlgebra +using Random +using Printf + +Reactant.set_default_backend("cpu") + +const SIZES = Dict("small" => (n=512, bs=128, reps=3), + "medium" => (n=1024, bs=256, reps=3), + "large" => (n=2048, bs=512, reps=2)) +const SIZE = get(SIZES, isempty(ARGS) ? "small" : first(ARGS)) do + error("Unknown size $(repr(first(ARGS))); expected one of $(join(sort(collect(keys(SIZES))), ", "))") +end + +""" + measure(prepare, run!; reps) -> (first, best) + +Times `run!(prepare())`, returning the time of the first call - which is where +Reactant compiles, and which is therefore what a one-shot script would see - and +the best of `reps` further calls, which is the steady-state cost once every +program involved has been compiled and cached. + +`prepare` runs untimed before each call, as these algorithms overwrite their +inputs. +""" +function measure(prepare, run!; reps::Int) + first_time = @elapsed run!(prepare()) + best = Inf + for _ in 1:reps + state = prepare() + best = min(best, @elapsed run!(state)) + end + return first_time, best +end + +const RESULTS = Tuple{String,String,Float64,Float64,Bool}[] + +function run_variant!(algorithm::String, variant::String, prepare, run!, check; reps::Int) + Dagger.reactant_cache_clear!() + correct = check(run!(prepare())) + first_time, best = measure(prepare, run!; reps) + push!(RESULTS, (algorithm, variant, first_time, best, correct)) + @printf(" %-26s %8.3fs first %8.3fs best%s\n", + variant, first_time, best, correct ? "" : " [WRONG RESULT]") + return +end + +############################################################################# +# Cholesky factorization of a DMatrix +############################################################################# + +chol_dagger(DX) = cholesky(DX) +chol_inner(DX) = Dagger.@reactant cholesky(DX) +chol_full(DX) = Dagger.@reactant mode=:full cholesky(DX) + +function bench_cholesky(n, bs, reps) + println("cholesky($n x $n, $bs x $bs blocks)") + Random.seed!(1234) + X = rand(n, n) + X = X * X' + n * I + reference = cholesky(copy(X)).U + + prepare() = distribute(copy(X), Blocks(bs, bs)) + check(chol) = isapprox(collect(chol.U), reference; rtol=1e-6) + + run_variant!("cholesky", "Dagger", prepare, chol_dagger, check; reps) + run_variant!("cholesky", "Reactant (inner)", prepare, chol_inner, check; reps) + run_variant!("cholesky", "Reactant (full)", prepare, chol_full, check; reps) + return +end + +############################################################################# +# A blocked matrix multiply, written as a Datadeps region +############################################################################# + +function blocked_matmul!(state) + C, A, B, blocks = state + Dagger.spawn_datadeps() do + for I in blocks, J in blocks, K in blocks + Dagger.@spawn mul!(InOut(view(C, I, J)), In(view(A, I, K)), In(view(B, K, J)), + 1.0, 1.0) + end + end + return C +end +matmul_dagger(state) = blocked_matmul!(state) +matmul_inner(state) = Dagger.@reactant blocked_matmul!(state) +matmul_full(state) = Dagger.@reactant mode=:full blocked_matmul!(state) + +function bench_matmul(n, bs, reps) + nblocks = cld(n, bs) + println("blocked matmul($n x $n, $(nblocks)x$(nblocks) blocks, $(nblocks^3) tasks)") + Random.seed!(1234) + A = rand(n, n) + B = rand(n, n) + reference = A * B + blocks = [idx:min(idx + bs - 1, n) for idx in 1:bs:n] + + prepare() = (zeros(n, n), A, B, blocks) + check(C) = isapprox(C, reference; rtol=1e-6) + + run_variant!("matmul", "Dagger", prepare, matmul_dagger, check; reps) + run_variant!("matmul", "Reactant (inner)", prepare, matmul_inner, check; reps) + run_variant!("matmul", "Reactant (full)", prepare, matmul_full, check; reps) + return +end + +############################################################################# +# An elementwise pipeline over a DArray, which is not a Datadeps region and so +# is only affected by inner mode +############################################################################# + +pipeline(DA) = sum(sqrt.(abs.(DA) .+ 1) .* 2) +pipeline_dagger(DA) = pipeline(DA) +pipeline_inner(DA) = Dagger.@reactant pipeline(DA) +pipeline_full(DA) = Dagger.@reactant mode=:full pipeline(DA) + +function bench_pipeline(n, bs, reps) + println("broadcast pipeline($n x $n, $bs x $bs blocks)") + Random.seed!(1234) + X = rand(n, n) + reference = pipeline(X) + + DA = distribute(X, Blocks(bs, bs)) + prepare() = DA + check(total) = isapprox(total, reference; rtol=1e-6) + + run_variant!("pipeline", "Dagger", prepare, pipeline_dagger, check; reps) + run_variant!("pipeline", "Reactant (inner)", prepare, pipeline_inner, check; reps) + run_variant!("pipeline", "Reactant (full)", prepare, pipeline_full, check; reps) + return +end + +############################################################################# +# What XLA itself costs, with no Dagger involved: the ceiling that either mode +# is working towards, and the explanation for most of what the numbers above show +############################################################################# + +function bench_reference(n, reps) + println("reference, without Dagger ($n x $n)") + Random.seed!(1234) + X = rand(n, n) + X = X * X' + n * I + A = rand(n, n) + B = rand(n, n) + + rX = Reactant.to_rarray(X) + rA = Reactant.to_rarray(A) + rB = Reactant.to_rarray(B) + chol_program = Reactant.compile(cholesky, (rX,); sync=true) + mul_program = Reactant.compile(*, (rA, rB); sync=true) + + _, lapack = measure(() -> copy(X), cholesky; reps) + _, xla = measure(() -> rX, chol_program; reps) + @printf(" %-26s %8.3fs LAPACK %8.3fs XLA\n", "cholesky", lapack, xla) + + _, blas = measure(() -> (A, B), state -> state[1] * state[2]; reps) + _, xla = measure(() -> (rA, rB), state -> mul_program(state...); reps) + @printf(" %-26s %8.3fs BLAS %8.3fs XLA\n", "matmul", blas, xla) + return +end + +############################################################################# + +function summarize() + println() + println("Summary (best-of-N, relative to Dagger without Reactant)") + for algorithm in unique(first.(RESULTS)) + rows = filter(row -> row[1] == algorithm, RESULTS) + baseline = rows[findfirst(row -> row[2] == "Dagger", rows)][4] + println(" $algorithm:") + for (_, variant, first_time, best, correct) in rows + @printf(" %-26s %8.3fs %6.2fx%s\n", + variant, best, baseline / best, correct ? "" : " [WRONG RESULT]") + end + end + return +end + +@info "Environment" julia=VERSION threads=Threads.nthreads() blas_threads=BLAS.get_num_threads() xla_platform=Reactant.XLA.platform_name(Reactant.XLA.default_backend()) + +bench_cholesky(SIZE.n, SIZE.bs, SIZE.reps) +bench_matmul(SIZE.n, SIZE.bs, SIZE.reps) +bench_pipeline(SIZE.n, SIZE.bs, SIZE.reps) +bench_reference(SIZE.n, SIZE.reps) +summarize() diff --git a/test/reactantenv/full.jl b/test/reactantenv/full.jl new file mode 100644 index 000000000..4436ea9e3 --- /dev/null +++ b/test/reactantenv/full.jl @@ -0,0 +1,126 @@ +@testset "Full mode" begin + @testset "A region becomes one compiled program" begin + Dagger.reactant_cache_clear!() + + A = rand(8, 8) + B = rand(8, 8) + C = zeros(8, 8) + Dagger.@reactant mode=:full begin + Dagger.spawn_datadeps() do + Dagger.@spawn mul!(Out(C), In(A), In(B)) + end + end + @test C ≈ A * B + @test Dagger.reactant_cache_size() == 1 + + # Running the same region again reuses the program compiled for it + for _ in 1:2 + Dagger.@reactant mode=:full begin + Dagger.spawn_datadeps() do + Dagger.@spawn mul!(Out(C), In(A), In(B)) + end + end + end + @test C ≈ A * B + @test Dagger.reactant_cache_size() == 1 + end + + @testset "Tasks are chained through their dependencies" begin + A = rand(8, 8) + B = zeros(8, 8) + C = zeros(8, 8) + increment!(X) = X .+= 1 + Dagger.@reactant mode=:full begin + Dagger.spawn_datadeps() do + Dagger.@spawn copyto!(Out(B), In(A)) + Dagger.@spawn increment!(InOut(B)) + Dagger.@spawn mul!(Out(C), In(B), In(A)) + end + end + @test B ≈ A .+ 1 + @test C ≈ (A .+ 1) * A + end + + @testset "Task results are available once the region has run" begin + A = rand(8, 8) + task = Ref{Any}(nothing) + Dagger.@reactant mode=:full begin + Dagger.spawn_datadeps() do + task[] = Dagger.@spawn sum(In(A)) + end + end + @test fetch(task[]) ≈ sum(A) + end + + @testset "Regions which Reactant cannot run are run by Dagger" begin + function scalar_fill!(A) + for idx in eachindex(A) + A[idx] = idx + end + return nothing + end + A = zeros(4, 4) + Dagger.@reactant mode=:full begin + Dagger.spawn_datadeps() do + Dagger.@spawn scalar_fill!(InOut(A)) + end + end + @test A == reshape(collect(1.0:16.0), 4, 4) + + # A `Ref` used as a scalar output cannot be written to by a compiled + # program, so the region must be handed back to Dagger + total = Ref(0.0) + B = rand(4, 4) + accumulate!(total, X) = (total[] = sum(X); nothing) + Dagger.@reactant mode=:full begin + Dagger.spawn_datadeps() do + Dagger.@spawn accumulate!(Out(total), In(B)) + end + end + @test total[] ≈ sum(B) + end + + @testset "must_opt reports a region Reactant cannot run" begin + function scalar_fill!(A) + for idx in eachindex(A) + A[idx] = idx + end + return nothing + end + A = zeros(4, 4) + @test_throws Dagger.ReactantOptimizationError begin + Dagger.@reactant mode=:full must_opt=true begin + Dagger.spawn_datadeps() do + Dagger.@spawn scalar_fill!(InOut(A)) + end + end + end + # The region is reported as un-runnable before any of it has run + @test all(A .== 0) + end + + @testset "cholesky($(bs)x$(bs) blocks)" for bs in (8, 16) + Random.seed!(1234) + X = rand(16, 16) + X = X * X' + 16I + DX = distribute(copy(X), Blocks(bs, bs)) + + chol = Dagger.@reactant mode=:full cholesky(DX) + ref = cholesky(X) + @test collect(chol.U) ≈ collect(ref.U) rtol=1e-8 + @test collect(chol.L) * collect(chol.U) ≈ X rtol=1e-8 + end + + @testset "@stencil" begin + A = zeros(Blocks(2, 2), Int, 4, 4) + A[2, 2] = 1 + source = collect(A) + B = zeros(Blocks(2, 2), Int, 4, 4) + Dagger.@reactant mode=:full begin + @stencil B[idx] = sum(@neighbors(A[idx], 1, Wrap())) + end + expected = [sum(source[mod1(i + di, 4), mod1(j + dj, 4)] for di in -1:1, dj in -1:1) + for i in 1:4, j in 1:4] + @test collect(B) == expected + end +end diff --git a/test/reactantenv/inner.jl b/test/reactantenv/inner.jl new file mode 100644 index 000000000..5618be7e8 --- /dev/null +++ b/test/reactantenv/inner.jl @@ -0,0 +1,171 @@ +@testset "Inner mode" begin + @testset "Task functions are compiled and run by Reactant" begin + Dagger.reactant_cache_clear!() + + A = rand(8, 8) + B = rand(8, 8) + C = Dagger.@reactant fetch(Dagger.@spawn A * B) + @test C isa Matrix{Float64} + @test C ≈ A * B + @test Dagger.reactant_cache_size() == 1 + + # The same call with the same shapes reuses the compiled program + C = Dagger.@reactant fetch(Dagger.@spawn A * B) + @test C ≈ A * B + @test Dagger.reactant_cache_size() == 1 + + # Different shapes are a different program, as XLA compiles for fixed shapes + A2 = rand(4, 4) + @test Dagger.@reactant(fetch(Dagger.@spawn A2 * A2)) ≈ A2 * A2 + @test Dagger.reactant_cache_size() == 2 + + Dagger.reactant_cache_clear!() + @test Dagger.reactant_cache_size() == 0 + end + + @testset "Writes to arguments are visible to Dagger" begin + A = rand(8, 8) + B = rand(8, 8) + C = zeros(8, 8) + Dagger.@reactant begin + Dagger.spawn_datadeps() do + Dagger.@spawn mul!(Out(C), In(A), In(B)) + end + end + @test C ≈ A * B + end + + @testset "Keyword arguments" begin + f(A; scale) = A .* scale + A = rand(8, 8) + @test Dagger.@reactant(fetch(Dagger.@spawn f(A; scale=3))) ≈ A .* 3 + end + + @testset "Nested and wrapped array arguments" begin + # Arguments are converted by Adapt, so arrays nested in containers and + # behind array wrappers are converted (and written back) as well + function scale_all!(pair, scale) + pair.first .*= scale + pair.second .*= scale + return nothing + end + A = ones(4, 4) + B = ones(4, 4) + Dagger.@reactant fetch(Dagger.@spawn scale_all!((first=A, second=B), 3)) + @test all(A .== 3) && all(B .== 3) + + # A write through a `view` lands in the array it is a view of + fill_with!(X, value) = (X .= value; nothing) + C = zeros(4, 4) + Dagger.@reactant fetch(Dagger.@spawn fill_with!(view(C, 1:2, :), 5)) + @test all(C[1:2, :] .== 5) && all(C[3:4, :] .== 0) + + # An array passed twice is one buffer, so both writes are kept + increment_both!(X, Y) = (X .+= 1; Y .+= 1; nothing) + D = zeros(4, 4) + Dagger.@reactant fetch(Dagger.@spawn increment_both!(D, D)) + @test all(D .== 2) + end + + @testset "Concurrent writes to views of one array" begin + # Datadeps runs tasks which write to disjoint views of an array + # concurrently, so each task's writes must be published back to its own + # view rather than to the whole array + fill_with!(X, value) = (X .= value; nothing) + blocks = [idx:(idx + 15) for idx in 1:16:64] + A = zeros(64, 64) + Dagger.@reactant begin + Dagger.spawn_datadeps() do + for (bi, I) in enumerate(blocks), (bj, J) in enumerate(blocks) + Dagger.@spawn fill_with!(Out(view(A, I, J)), bi * 10 + bj) + end + end + end + @test A == [fld1(i, 16) * 10 + fld1(j, 16) for i in 1:64, j in 1:64] + end + + @testset "Results are plain Julia values" begin + # Reactant values must not escape into the caller's hands, including from + # within the arrays that Dagger's reductions pass between their tasks + DA = distribute(rand(16, 16), Blocks(8, 8)) + total = Dagger.@reactant sum(DA) + @test total isa Float64 + @test total ≈ sum(collect(DA)) + + column_sums = Dagger.@reactant collect(sum(DA; dims=1)) + @test column_sums isa Matrix{Float64} + @test column_sums ≈ sum(collect(DA); dims=1) + end + + @testset "Tasks which Reactant cannot compile are run without it" begin + # A scalar-indexed kernel cannot be traced, and must not be a failure: + # the task simply runs as it normally would + function scalar_fill!(A) + for idx in eachindex(A) + A[idx] = idx + end + return A + end + A = zeros(4, 4) + result = Dagger.@reactant fetch(Dagger.@spawn scalar_fill!(A)) + @test result == reshape(collect(1.0:16.0), 4, 4) + + # ... and is not attempted again + A2 = zeros(4, 4) + result = @test_logs min_level=Logging.Warn begin + Dagger.@reactant fetch(Dagger.@spawn scalar_fill!(A2)) + end + @test result == reshape(collect(1.0:16.0), 4, 4) + + # Unless the caller cannot afford to lose Reactant, in which case the + # failure is reported rather than worked around + A3 = zeros(4, 4) + err = try + Dagger.@reactant must_opt=true fetch(Dagger.@spawn scalar_fill!(A3)) + nothing + catch caught + caught + end + @test Dagger.Sch.unwrap_nested_exception(err) isa Dagger.ReactantOptimizationError + @test all(A3 .== 0) + end + + @testset "cholesky($T, $(bs)x$(bs) blocks)" for T in (Float32, Float64), bs in (4, 8) + Random.seed!(1234) + X = rand(T, 16, 16) + X = X * X' + 16I + DX = distribute(copy(X), Blocks(bs, bs)) + + chol = Dagger.@reactant cholesky(DX) + ref = cholesky(X) + rtol = T == Float32 ? 1e-3 : 1e-8 + @test collect(chol.U) ≈ collect(ref.U) rtol=rtol + @test collect(chol.L) * collect(chol.U) ≈ X rtol=rtol + end + + @testset "@stencil" begin + A = zeros(Blocks(2, 2), Int, 4, 4) + B = zeros(Blocks(2, 2), Int, 4, 4) + Dagger.@reactant begin + @stencil begin + A[idx] = 1 + B[idx] = A[idx] + 2 + end + end + @test all(collect(A) .== 1) + @test all(collect(B) .== 3) + + # Neighborhood reads, which go through a halo and so are the least + # Reactant-friendly thing `@stencil` generates + C = zeros(Blocks(2, 2), Int, 4, 4) + C[2, 2] = 1 + source = collect(C) + D = zeros(Blocks(2, 2), Int, 4, 4) + Dagger.@reactant begin + @stencil D[idx] = sum(@neighbors(C[idx], 1, Wrap())) + end + expected = [sum(source[mod1(i + di, 4), mod1(j + dj, 4)] for di in -1:1, dj in -1:1) + for i in 1:4, j in 1:4] + @test collect(D) == expected + end +end diff --git a/test/reactantenv/runtests.jl b/test/reactantenv/runtests.jl new file mode 100644 index 000000000..832437fa2 --- /dev/null +++ b/test/reactantenv/runtests.jl @@ -0,0 +1,22 @@ +# Reactant integration tests. These live in their own project (and their own CI +# job), because Reactant is a heavy dependency to build and load. +# +# Run with: +# julia --project=test/reactantenv test/reactantenv/runtests.jl + +using Test +using Dagger +using Reactant +using LinearAlgebra +using Random +import Logging + +import Dagger: @stencil, Wrap + +# XLA's GPU backend is not what is under test here, and CI has no GPU +Reactant.set_default_backend("cpu") + +@test Dagger.reactant_available() + +include("inner.jl") +include("full.jl") diff --git a/test/runtests.jl b/test/runtests.jl index bfc347ce0..82ae4568d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -21,6 +21,7 @@ tests = [ ("Task Queues", "task-queues.jl"), ("Task Affinity", "task-affinity.jl"), ("Datadeps", "datadeps.jl"), + ("Reactant", "reactant.jl"), ("Streaming", "streaming.jl"), ("Domain Utilities", "domain.jl"), ("Array - Allocation", "array/allocation.jl"),