Add Reactant optimization support

.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

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"

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

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.

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
+```