Merge branch 'main' into nounroll

Author: wsmoses

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <wmoses@mit.edu>", "Valentin Churavy <vchuravy@mit.edu>", "Sergio Sánchez Ramírez <sergio.sanchez.ramirez@bsc.es>", "Paul Berg <paul@plutojl.org>", "Avik Pal <avikpal@mit.edu>", "Mosè Giordano <mose@gnu.org>"]
-version = "0.2.44"
+version = "0.2.46"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -85,8 +85,8 @@ Preferences = "1.4"
 PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
-ReactantCore = "0.1.5"
-Reactant_jll = "0.0.92"
+ReactantCore = "0.1.6"
+Reactant_jll = "0.0.93"
 Scratch = "1.2"
 Sockets = "1.10"
 SpecialFunctions = "2.4"

deps/ReactantExtra/WORKSPACE

@@ -9,7 +9,7 @@ http_archive(
     urls = ["https://github.com/wsmoses/nsync/archive/{commit}.tar.gz".format(commit = NSYNC_COMMIT)],
 )
 
-ENZYMEXLA_COMMIT = "d39a1fa814d329293645ff9771e68c9e9be8ceae"
+ENZYMEXLA_COMMIT = "9fe4c5e3f2ec044db7860ead3d42838589a08104"
 ENZYMEXLA_SHA256 = ""
 
 http_archive(

docs/make.jl

@@ -36,6 +36,7 @@ pages = [
         "Overview" => "tutorials/index.md",
         "Profiling" => "tutorials/profiling.md",
         "Distributed" => "tutorials/multihost.md",
+        "Local build" => "tutorials/local-build.md",
     ],
     "API Reference" => [
         "Reactant API" => "api/api.md",

docs/src/.vitepress/config.mts

@@ -64,6 +64,7 @@ export default defineConfig({
           {text: "Overview", link: "/tutorials/"},
           {text: "Profiling", link: "/tutorials/profiling"},
           {text: "Distributed", link: "/tutorials/multihost"},
+          {text: "Local build", link: "/tutorials/local-build"},
         ],
       },
       {
@@ -124,6 +125,7 @@ export default defineConfig({
           { text: "Overview", link: "/tutorials/" },
           { text: "Profiling", link: "/tutorials/profiling" },
           { text: "Distributed", link: "/tutorials/multihost" },
+          { text: "Local build", link: "/tutorials/local-build" },
         ],
       },
       "/api/": {

docs/src/tutorials/index.md

@@ -2,5 +2,6 @@
 
  - [Profiling](@ref profiling).
  - [Multi-Host Environments](@ref distributed).
+ - [Local build of ReactantExtra](@ref local-build).
 
 We are currently working on adding more tutorials to Reactant!! Please check back soon!

docs/src/tutorials/local-build.md

@@ -0,0 +1,74 @@
+# [Local build of ReactantExtra](@ref local-build)
+
+In the `deps/` subdirectory of the Reactant repository there is a script to do local builds of ReactantExtra, including debug builds.
+
+## Requirements
+
+* Julia.  If you don't have it already, you can obtain it from the [official Julia website](https://julialang.org/downloads/)
+* A reasonably recent C/C++ compiler, ideally GCC 12+.
+  Older compilers may not work.
+* Bazel. If you don't have it already, you can download a build for your platform from [the latest `bazelbuild/bazelisk` release](https://github.com/bazelbuild/bazelisk/releases/latest) and put the `bazel` executable in `PATH`
+* not necessary in general, but for debug builds with CUDA support, you'll need a fast linker, like `lld` or `mold`
+  Binutils `ld` won't work, don't even try using it.
+  You can obtain `mold` for your platform from the [latest `rui314/mold` release](https://github.com/rui314/mold/releases/latest) and put the `mold` executable in `PATH`
+
+## Building
+
+At a high-level, after you `cd` to the `deps/` directory you can run the commands
+
+```bash
+julia --project -e 'using Pkg; Pkg.instantiate()' # needed only the first time to install dependencies for this script
+julia -O0 --color=yes --project build_local.jl
+```
+
+There are a few of options you may want to use to tweak the build.
+For more information run the command (what's show below may not be up to date, run the command locally to see the options available to you):
+
+```console
+% julia -O0 --project build_local.jl --help
+usage: build_local.jl [--debug] [--backend BACKEND]
+                      [--gcc_host_compiler_path GCC_HOST_COMPILER_PATH]
+                      [--cc CC]
+                      [--hermetic_python_version HERMETIC_PYTHON_VERSION]
+                      [--jobs JOBS] [--copt COPT] [--cxxopt CXXOPT]
+                      [--extraopt EXTRAOPT] [--color COLOR] [-h]
+
+optional arguments:
+  --debug               Build with debug mode (-c dbg).
+  --backend BACKEND     Build with the specified backend (auto, cpu,
+                        cuda). (default: "auto")
+  --gcc_host_compiler_path GCC_HOST_COMPILER_PATH
+                        Path to the gcc host compiler. (default:
+                        "/usr/bin/gcc")
+  --cc CC                (default: "/usr/bin/cc")
+  --hermetic_python_version HERMETIC_PYTHON_VERSION
+                        Hermetic Python version. (default: "3.10")
+  --jobs JOBS           Number of parallel jobs. (type: Int64,
+                        default: <MAXIMUM NUMBER OF CPUs>)
+  --copt COPT           Options to be passed to the C compiler.  Can
+                        be used multiple times.
+  --cxxopt CXXOPT       Options to be passed to the C++ compiler.  Can
+                        be used multiple times.
+  --extraopt EXTRAOPT   Extra options to be passed to Bazel.  Can be
+                        used multiple times.
+  --color COLOR         Set to `yes` to enable color output, or `no`
+                        to disable it. Defaults to same color setting
+                        as the Julia process. (default: "no")
+  -h, --help            show this help message and exit
+```
+
+### Doing a build on a system with memoryor number of processes restrictions
+
+If you try to do the build on certain systems where there are in place restrictions on the number of processes or memory that your user can use (for example login node of clusters), you may have to limit the number of parallel jobs used by Bazel.
+By default Bazel would try to use the maximum number of CPUs available on the system, if you need reduce that pass the `--jobs JOBS` flag option.
+The Bazel server may be terminated abruptly if using too much memory (e.g. if concurrent compiler processes are cumulatively using a large amount of memory), also in this case reducing the number of parallel jobs may be beneficial.
+
+### CUDA debug build
+
+A CUDA debug build (`--debug --backend=cuda`) requires a recent GCC compiler (at least v12) and also a fast linker (see requirements above).
+You can tell GCC to use either `lld` or `mold` with `--extraopt '--linkopt=-fuse-ld=lld'` or `--extraopt '--linkopt=-fuse-ld=mold'` respectively.
+NOTE: the option `-fuse-ld=mold` was added in GCC 12, if you're trying to use an older version you can have some luck by making a symlink named `ld` pointing to `mold` in `PATH`, with higher precendce than Binutils `ld`.
+
+### Using ccache
+
+If you want to use `ccache` as your compiler, you may have to add the flag `--extraopt "--sandbox_writable_path=/path/to/ccache/directory"` to let `ccache` write to its own directory.

lib/ReactantCore/Project.toml

@@ -1,7 +1,7 @@
 name = "ReactantCore"
 uuid = "a3311ec8-5e00-46d5-b541-4f83e724a433"
 authors = ["William Moses <wmoses@mit.edu>", "Valentin Churavy <vchuravy@mit.edu>", "Sergio Sánchez Ramírez <sergio.sanchez.ramirez@bsc.es>", "Paul Berg <paul@plutojl.org>", "Avik Pal <avikpal@mit.edu>"]
-version = "0.1.5"
+version = "0.1.6"
 
 [deps]
 ExpressionExplorer = "21656369-7473-754a-2065-74616d696c43"

lib/ReactantCore/src/ReactantCore.jl

@@ -6,7 +6,18 @@ using MacroTools: MacroTools
 export @trace, within_compile, MissingTracedValue
 
 # Traits
-is_traced(x) = false
+function is_traced((@nospecialize x::T), seen=Base.IdSet()) where {T}
+    if !isprimitivetype(x)
+        for fn in fieldnames(T)
+            f = getfield(x, fn)
+            if !(f in seen)
+                push!(seen, f)
+                is_traced(f, seen) && return true
+            end
+        end
+    end
+    return false
+end
 
 # New Type signifying that a value is missing
 mutable struct MissingTracedValue

src/Compiler.jl

@@ -1055,6 +1055,17 @@ function compile_mlir!(
     MLIR.API.mlirOperationDestroy(compiled_f.operation)
     compiled_f.operation = MLIR.API.MlirOperation(C_NULL)
 
+    # Add a `donated` attr to the function arguments. This doesn't affect XLA, but lets us
+    # check which arguments were donated.
+    preserved_args_idx = last.(preserved_args)
+    for (i, arg) in enumerate(linear_args)
+        if i ∉ preserved_args_idx
+            MLIR.API.mlirFuncSetArgAttr(
+                func3, i - 1, "reactant.donated", MLIR.IR.UnitAttribute()
+            )
+        end
+    end
+
     return Reactant.TracedUtils.CompiledMlirFnResult(
         fnwrapped,
         func3,

src/ConcreteRArray.jl

@@ -341,6 +341,7 @@ end
 
 # TODO replace this copy for `setindex!` maybe? how to copy data to already existing buffer? (i.e. `copyto!`)
 function Base.copy(bc::Base.Broadcast.Broadcasted{Broadcast.ArrayStyle{ConcretePJRTArray}})
+    bc = Broadcast.flatten(bc)
     for x in bc.args
         x isa ConcretePJRTArray && wait(x)
     end
@@ -370,6 +371,7 @@ function Base.copy(bc::Base.Broadcast.Broadcasted{Broadcast.ArrayStyle{ConcreteP
 end
 
 function Base.copy(bc::Base.Broadcast.Broadcasted{Broadcast.ArrayStyle{ConcreteIFRTArray}})
+    bc = Broadcast.flatten(bc)
     fn = compile(Broadcast.BroadcastFunction(bc.f), (bc.args...,))
     return fn(bc.args...)
 end

src/Ops.jl

@@ -36,6 +36,7 @@ end
 
 const DEBUG_MODE::Ref{Bool} = Ref(false)
 const LARGE_CONSTANT_THRESHOLD = Ref(100 << 20) # 100 MiB
+const LARGE_CONSTANT_RAISE_ERROR = Ref(true)
 
 function with_debug(f)
     old = DEBUG_MODE[]
@@ -89,8 +90,18 @@ end
 @noinline function constant(
     x::DenseArray{T,N}; location=mlir_stacktrace("constant", @__FILE__, @__LINE__)
 ) where {T,N}
-    sizeof(x) > LARGE_CONSTANT_THRESHOLD[] &&
-        error("Generating a constant larger than $(LARGE_CONSTANT_THRESHOLD[]) bytes.")
+    if sizeof(x) > LARGE_CONSTANT_THRESHOLD[]
+        if LARGE_CONSTANT_RAISE_ERROR[]
+            error(
+                "Generating a constant of $(sizeof(x)) bytes, which larger than the $(LARGE_CONSTANT_THRESHOLD[]) bytes threshold",
+            )
+        else
+            location = with_debug() do
+                mlir_stacktrace("constant", @__FILE__, @__LINE__)
+            end
+        end
+    end
+
     value = MLIR.IR.DenseElementsAttribute(x)
     constants = constant_context()[2]
     if haskey(constants, value)
@@ -1713,7 +1724,10 @@ end
     traced_args = Vector{Any}(undef, N)
     for i in 1:N
         @inbounds traced_args[i] = Reactant.make_tracer(
-            seen_args, args[i], (), Reactant.NoStopTracedTrack; track_numbers=Number
+            seen_args,
+            args[i],
+            (),
+            Reactant.NoStopTracedTrack, #; track_numbers=Number
         )
     end
 

src/TracedRArray.jl

@@ -22,6 +22,7 @@ using ..TracedUtils: TracedUtils, get_mlir_data, set_mlir_data!, materialize_tra
 using ReactantCore: ReactantCore
 using GPUArraysCore: GPUArraysCore, @allowscalar
 
+ReactantCore.is_traced(::TracedRArray, seen) = true
 ReactantCore.is_traced(::TracedRArray) = true
 
 Base.strides(x::TracedRArray) = Base.size_to_strides(1, size(x)...)

src/TracedRNumber.jl

@@ -4,6 +4,7 @@ using ..Reactant:
     Reactant, TracedRNumber, TracedRArray, TracedUtils, Ops, MLIR, unwrapped_eltype
 using ReactantCore
 
+ReactantCore.is_traced(::TracedRNumber, seen) = true
 ReactantCore.is_traced(::TracedRNumber) = true
 
 Base.getindex(a::TracedRNumber{T}) where {T} = a

src/TracedUtils.jl

@@ -629,6 +629,7 @@ function elem_apply(f, args::Vararg{Any,Nargs}) where {Nargs}
 end
 
 function broadcast_to_size(arg::AbstractArray{<:TracedRNumber}, rsize)
+    collect(size(arg)) == collect(rsize) && return arg
     if Reactant.ancestor(arg) isa TracedRArray
         return broadcast_to_size(materialize_traced_array(arg), rsize)
     end
@@ -658,7 +659,7 @@ function broadcast_to_size(arg::AbstractIrrational, rsize)
 end
 
 function broadcast_to_size(arg::ReactantPrimitive, rsize)
-    return Ops.constant(Base.fill(arg, Tuple(rsize)))
+    return Ops.fill(arg, rsize)
 end
 
 function broadcast_to_size(arg::TracedRNumber{T}, rsize) where {T}

src/Tracing.jl

@@ -551,6 +551,20 @@ Base.@nospecializeinfer function traced_type_inner(
     end
 end
 
+Base.@nospecializeinfer function Reactant.traced_type_inner(
+    @nospecialize(OA::Type{SubArray{T,N,P,I,L}}),
+    seen,
+    mode::Reactant.TraceMode,
+    @nospecialize(track_numbers::Type),
+    @nospecialize(sharding),
+    @nospecialize(runtime)
+) where {T,N,P,I,L}
+    P2 = Reactant.traced_type_inner(P, seen, mode, track_numbers, sharding, runtime)
+    I2 = Reactant.traced_type_inner(I, seen, mode, track_numbers, sharding, runtime)
+    T2 = eltype(P2)
+    return SubArray{T2,N,P2,I2,L}
+end
+
 for P in (Ptr, Core.LLVMPtr, Base.RefValue)
     @eval Base.@nospecializeinfer function traced_type_inner(
         @nospecialize(PT::Type{$P}),
@@ -918,6 +932,81 @@ function make_tracer(
 end
 append_path(@nospecialize(path), i) = (path..., i)
 
+function make_tracer_via_immutable_constructor(
+    seen,
+    @nospecialize(prev),
+    @nospecialize(path),
+    mode;
+    @nospecialize(track_numbers::Type = Union{}),
+    @nospecialize(sharding = Sharding.NoSharding()),
+    @nospecialize(runtime = nothing),
+    kwargs...,
+)
+    RT = Core.Typeof(prev)
+    if haskey(seen, prev)
+        if mode == TracedToTypes
+            id = seen[prev]
+            push!(path, id)
+            return nothing
+        elseif mode != NoStopTracedTrack && haskey(seen, prev)
+            return seen[prev]
+        end
+    elseif mode == TracedToTypes
+        push!(path, RT)
+        seen[prev] = VisitedObject(length(seen) + 1)
+    end
+    TT = traced_type(RT, Val(mode), track_numbers, sharding, runtime)
+    @assert !Base.isabstracttype(RT)
+    @assert Base.isconcretetype(RT)
+    nf = fieldcount(RT)
+
+    @assert !ismutabletype(TT)
+
+    if nf == 0
+        if mode == TracedToTypes
+            push!(path, prev)
+            return nothing
+        end
+        return prev
+    end
+
+    flds = Vector{Any}(undef, nf)
+    changed = false
+    for i in 1:nf
+        if isdefined(prev, i)
+            newpath = mode == TracedToTypes ? path : append_path(path, i)
+            xi = Base.getfield(prev, i)
+            xi2 = make_tracer(
+                seen,
+                xi,
+                newpath,
+                mode;
+                track_numbers,
+                sharding=Base.getproperty(sharding, i),
+                runtime,
+                kwargs...,
+            )
+            if xi !== xi2
+                changed = true
+            end
+            flds[i] = xi2
+        else
+            nf = i - 1 # rest of tail must be undefined values
+            break
+        end
+    end
+    if mode == TracedToTypes
+        return nothing
+    end
+    if !changed
+        seen[prev] = prev
+        return prev
+    end
+    y = TT(flds...)
+    seen[prev] = y
+    return y
+end
+
 function make_tracer(
     seen,
     @nospecialize(prev),

test/basic.jl

@@ -1014,3 +1014,9 @@ end
     @test Array(x_ra) ==
         [0.0 0.0 1.0 1.0; 0.0 0.0 1.0 1.0; 1.0 1.0 1.0 1.0; 1.0 1.0 1.0 1.0]
 end
+
+@testset "copy(::Broadcast.Broadcasted{ArrayStyle{ConcreteRArray}})" begin
+    x_ra = Reactant.to_rarray(ones(4, 4))
+    res = copy(Broadcast.broadcasted(-, Broadcast.broadcasted(+, x_ra, 1)))
+    @test res ≈ -(Array(x_ra) .+ 1)
+end