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