Skip to content

Add Reactant optimization support - #738

Open
jpsamaroo wants to merge 2 commits into
masterfrom
jps/reactant
Open

Add Reactant optimization support#738
jpsamaroo wants to merge 2 commits into
masterfrom
jps/reactant

Conversation

@jpsamaroo

@jpsamaroo jpsamaroo commented Aug 18, 2026

Copy link
Copy Markdown
Member

Adds opt-in Reactant support via Dagger.@reactant to either individual tasks (:inner) or entire DAGs (:full).

Written by Claude Opus

jpsamaroo and others added 2 commits August 18, 2026 12:55
@jpsamaroo @cursoragent
DArray: Accept any array type from a partial reduction
635d88e 
A partial reduction which produced an array that is not a `Base.Array` was
wrapped in a 1x1 array, rather than being passed along as the array it already
is, which added a level of nesting to the result of the reduction.

Co-authored-by: Cursor <cursoragent@cursor.com>
@jpsamaroo @cursoragent
Add Reactant.jl integration
849cd5c 
`Dagger.@reactant expr` hands the work that `expr` launches to Reactant.jl,
which compiles it through MLIR/XLA and optimizes it in ways that Julia's own
compiler does not. Two modes are available:

- `mode=:inner` (the default) compiles each task's function on its own, on the
  processor that runs it, leaving Dagger's scheduling and data movement alone.
- `mode=:full` traces a whole `spawn_datadeps` region as one program, which lets
  Reactant fuse and reorder across task boundaries, at the cost of Datadeps
  doing no planning or scheduling for that region.

Whatever Reactant cannot do is done without it: a missing Reactant.jl, a task
function that cannot be traced, and a region that cannot be compiled or run are
all warned about and then run as they would have been otherwise, so the same
application code works with and without Reactant. `must_opt=true` and
`must_load=true` make those cases errors instead, for code which cannot afford
to quietly stop being accelerated.

Compiled programs are cached and reused whenever the values Reactant baked into
them are known to match. Reactant is a weak dependency, and all of the logic
that uses it lives in the ReactantExt extension.

Cholesky factorization needs two kernels that Reactant cannot use as written:
`LAPACK.potrf!` is a `ccall`, and `BLAS.syrk!` currently lowers to an op which
updates the opposite triangle of `C` from what BLAS does, so both are given
traceable implementations here.

Reactant tests run in their own environment and CI job, as it is far too heavy
a dependency for the main testsuite. `test/reactantenv/bench.jl` compares each
mode against plain Dagger; what it finds is written up in the docs.

Co-authored-by: Cursor <cursoragent@cursor.com>
@github-actions

github-actions Bot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Dagger benchmarks: dirty vs master

master dirty master / dirty
array/dagger/N=1024 (block 512)/add (X + X) 3.87 ± 0.59 ms 4.08 ± 0.41 ms 0.948 ± 0.17
array/dagger/N=1024 (block 512)/alloc (rand) 3.45 ± 0.53 ms 3.41 ± 0.19 ms 1.01 ± 0.17
array/dagger/N=1024 (block 512)/broadcast (X .+ 1) 2.71 ± 0.86 ms 2.59 ± 0.96 ms 1.05 ± 0.51
array/dagger/N=1024 (block 512)/map (sin.(X)) 7.1 ± 0.96 ms 7.06 ± 0.84 ms 1 ± 0.18
array/dagger/N=1024 (block 512)/norm 1.62 ± 0.4 ms 1.31 ± 0.032 ms 1.24 ± 0.3
array/dagger/N=1024 (block 512)/reduce (sum) 3.11 ± 2.2 ms 2.76 ± 0.42 ms 1.13 ± 0.81
array/dagger/N=1024 (block 512)/transpose (permutedims) 7.19 ± 0.54 ms 7.29 ± 0.64 ms 0.987 ± 0.11
array/dagger/N=256 (block 256)/add (X + X) 0.962 ± 0.059 ms 1.05 ± 0.085 ms 0.914 ± 0.093
array/dagger/N=256 (block 256)/alloc (rand) 0.85 ± 0.046 ms 1.03 ± 0.2 ms 0.826 ± 0.16
array/dagger/N=256 (block 256)/broadcast (X .+ 1) 0.747 ± 0.085 ms 0.856 ± 0.067 ms 0.872 ± 0.12
array/dagger/N=256 (block 256)/map (sin.(X)) 1.28 ± 0.41 ms 1.28 ± 0.04 ms 0.998 ± 0.32
array/dagger/N=256 (block 256)/norm 0.523 ± 1.3 ms 0.533 ± 0.22 ms 0.981 ± 2.5
array/dagger/N=256 (block 256)/reduce (sum) 4.73 ± 3.9 ms 4.14 ± 4 ms 1.14 ± 1.5
array/dagger/N=256 (block 256)/transpose (permutedims) 0.936 ± 0.086 ms 1.05 ± 0.27 ms 0.89 ± 0.24
linalg/dagger/N=1024 (block 512)/cholesky 24.7 ± 7.1 ms 23.7 ± 2.7 ms 1.04 ± 0.32
linalg/dagger/N=1024 (block 512)/lu 0.0438 ± 0.0011 s 0.0475 ± 0.0033 s 0.923 ± 0.069
linalg/dagger/N=1024 (block 512)/matmul (A*A) 0.0551 ± 0.013 s 0.057 ± 0.015 s 0.966 ± 0.34
linalg/dagger/N=1024 (block 512)/matvec (A*x) 2.96 ± 0.21 ms 3.32 ± 1.3 ms 0.893 ± 0.36
linalg/dagger/N=1024 (block 512)/qr 0.116 ± 0.0064 s 0.119 ± 0.0059 s 0.974 ± 0.072
linalg/dagger/N=1024 (block 512)/solve (A\b via lu) 0.0548 ± 0.0099 s 0.0527 ± 0.0043 s 1.04 ± 0.21
linalg/dagger/N=1024 (block 512)/svd 0.0368 h 0.0366 h 1
linalg/dagger/N=1024 (block 512)/syrk (A'*A) 0.0375 ± 0.0043 s 0.0388 ± 0.0021 s 0.965 ± 0.12
linalg/dagger/N=256 (block 256)/cholesky 5.22 ± 1.4 ms 3.62 ± 0.97 ms 1.44 ± 0.55
linalg/dagger/N=256 (block 256)/lu 4.55 ± 0.67 ms 4.47 ± 0.39 ms 1.02 ± 0.17
linalg/dagger/N=256 (block 256)/matmul (A*A) 2.34 ± 0.63 ms 2.98 ± 0.74 ms 0.782 ± 0.29
linalg/dagger/N=256 (block 256)/matvec (A*x) 1.34 ± 0.28 ms 1.37 ± 0.2 ms 0.973 ± 0.25
linalg/dagger/N=256 (block 256)/qr 5.87 ± 0.79 ms 5.35 ± 0.46 ms 1.1 ± 0.17
linalg/dagger/N=256 (block 256)/solve (A\b via lu) 10.2 ± 1.7 ms 11.3 ± 6.8 ms 0.902 ± 0.56
linalg/dagger/N=256 (block 256)/svd 0.536 ± 0.022 s 0.558 ± 0.025 s 0.96 ± 0.058
linalg/dagger/N=256 (block 256)/syrk (A'*A) 3.9 ± 1.5 ms 3.84 ± 0.7 ms 1.02 ± 0.42
stencil/dagger/N=1024 (block 512)/alloc (neighbors Wrap) 9.2 ± 1.2 ms 8.65 ± 1.3 ms 1.06 ± 0.21
stencil/dagger/N=1024 (block 512)/assign (const) 1.53 ± 0.6 ms 1.72 ± 2.2 ms 0.887 ± 1.2
stencil/dagger/N=1024 (block 512)/multi-expr 5.36 ± 0.33 ms 4.13 ± 0.57 ms 1.3 ± 0.2
stencil/dagger/N=1024 (block 512)/neighbors (Clamp) 7.57 ± 0.38 ms 8.14 ± 1.5 ms 0.93 ± 0.17
stencil/dagger/N=1024 (block 512)/neighbors (Pad) 7.96 ± 1.6 ms 6.01 ± 0.63 ms 1.32 ± 0.3
stencil/dagger/N=1024 (block 512)/neighbors (Reflect) 6.93 ± 1.3 ms 7.05 ± 0.2 ms 0.982 ± 0.19
stencil/dagger/N=1024 (block 512)/neighbors (Wrap) 7.63 ± 0.3 ms 8.09 ± 0.93 ms 0.943 ± 0.11
stencil/dagger/N=1024 (block 512)/update (+) 2.22 ± 0.92 ms 2.26 ± 2.5 ms 0.985 ± 1.2
stencil/dagger/N=256 (block 256)/alloc (neighbors Wrap) 2.12 ± 0.13 ms 1.88 ± 0.28 ms 1.13 ± 0.18
stencil/dagger/N=256 (block 256)/assign (const) 0.664 ± 0.12 ms 0.706 ± 0.013 ms 0.941 ± 0.18
stencil/dagger/N=256 (block 256)/multi-expr 1.29 ± 0.19 ms 1.72 ± 0.53 ms 0.754 ± 0.26
stencil/dagger/N=256 (block 256)/neighbors (Clamp) 2.04 ± 0.04 ms 1.77 ± 0.064 ms 1.16 ± 0.048
stencil/dagger/N=256 (block 256)/neighbors (Pad) 1.82 ± 0.27 ms 1.72 ± 0.12 ms 1.06 ± 0.17
stencil/dagger/N=256 (block 256)/neighbors (Reflect) 1.76 ± 0.058 ms 1.9 ± 0.13 ms 0.926 ± 0.073
stencil/dagger/N=256 (block 256)/neighbors (Wrap) 1.98 ± 0.18 ms 1.67 ± 0.16 ms 1.19 ± 0.16
stencil/dagger/N=256 (block 256)/update (+) 1.09 ± 0.047 ms 0.9 ± 0.081 ms 1.22 ± 0.12
time_to_load 1.07 ± 0.0064 s 1.07 ± 0.009 s 0.995 ± 0.01

Plots

⚠️ Regressions (> 25.0%)

  • stencil/dagger/N=256 (block 256)/multi-expr: +32.6%
  • linalg/dagger/N=256 (block 256)/matmul (A*A): +27.8%

Improvements (> 25.0% faster)

  • linalg/dagger/N=256 (block 256)/cholesky: -30.7%

Full results and plots (download the benchmark-results artifact).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

datadeps performance reactant

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@jpsamaroo