Merge branch 'main' into pb/ll-inline-roots

Author: Pangoraw

.gitignore

@@ -253,6 +253,7 @@ docs/site/
 Manifest.toml
 Manifest-v*.toml
 .CondaPkg
+LocalPreferences.toml
 
 .vscode/*
 .vscode/settings.json

Project.toml

@@ -1,6 +1,6 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
-version = "0.2.253"
+version = "0.2.254"
 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>"]
 
 [workspace]
@@ -149,7 +149,7 @@ PythonCall = "0.9.25"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.18"
-Reactant_jll = "0.0.371"
+Reactant_jll = "0.0.375"
 ScopedValues = "1.3.0"
 Scratch = "1.3"
 Serialization = "1.10"

deps/ReactantExtra/BUILD

@@ -1672,6 +1672,24 @@ gentbl_cc_library(
     ],
 )
 
+gentbl_cc_library(
+    name = "ImpulseJLIncGen",
+    tbl_outs = [
+        (
+            [
+                "--generator=jl-op-defs",
+                "--disable-module-wrap=0",
+            ],
+            "Impulse.jl",
+        ),
+    ],
+    tblgen = "//:mlir-jl-tblgen",
+    td_file = "@enzyme//:Enzyme/MLIR/Dialect/Impulse/ImpulseOps.td",
+    deps = [
+        "@enzyme//:ImpulseDialectTdFiles",
+    ],
+)
+
 gentbl_cc_library(
     name = "EnzymeXLAJLIncGen",
     tbl_outs = [

deps/ReactantExtra/WORKSPACE

@@ -4,7 +4,7 @@ NSYNC_COMMIT = "82b118aa7ace3132e517e2c467f8732978cf4023"
 
 NSYNC_SHA256 = ""
 
-ENZYMEXLA_COMMIT = "10d4571a0d5fa08a2fbf285dc585cf16a091b573"
+ENZYMEXLA_COMMIT = "99d7b16a3ef9cd5cac1f616c94e5be617a16f0d2"
 
 ENZYMEXLA_SHA256 = ""
 

deps/ReactantExtra/make-bindings.jl

@@ -16,6 +16,7 @@ dialect_files = [
     "Func.jl",
     "Enzyme.jl",
     "EnzymeXLA.jl",
+    "Impulse.jl",
     "StableHLO.jl",
     "CHLO.jl",
     "VHLO.jl",

docs/src/.vitepress/config.mts

@@ -127,6 +127,7 @@ export default defineConfig({
               { text: "EnzymeXLA", link: "/api/dialects/enzymexla" },
               { text: "Func", link: "/api/dialects/func" },
               { text: "GPU", link: "/api/dialects/gpu" },
+              { text: "Impulse", link: "/api/dialects/impulse" },
               { text: "LLVM", link: "/api/dialects/llvm" },
               { text: "MPI", link: "/api/dialects/mpi" },
               { text: "MemRef", link: "/api/dialects/memref" },
@@ -225,6 +226,7 @@ export default defineConfig({
               { text: "EnzymeXLA", link: "/api/dialects/enzymexla" },
               { text: "Func", link: "/api/dialects/func" },
               { text: "GPU", link: "/api/dialects/gpu" },
+              { text: "Impulse", link: "/api/dialects/impulse" },
               { text: "LLVM", link: "/api/dialects/llvm" },
               { text: "MPI", link: "/api/dialects/mpi" },
               { text: "MemRef", link: "/api/dialects/memref" },

docs/src/api/dialects/impulse.md

@@ -0,0 +1,9 @@
+```@meta
+CollapsedDocStrings = true
+```
+
+# Impulse Dialect
+
+```@autodocs
+Modules = [Reactant.MLIR.Dialects.impulse]
+```

src/Compiler.jl

@@ -277,6 +277,37 @@ Base.@nospecializeinfer function create_result(
     return result_cache[tocopy]
 end
 
+Base.@nospecializeinfer function create_result(
+    @nospecialize(tocopy::Enum),
+    @nospecialize(path::Tuple),
+    result_stores,
+    path_to_shard_info,
+    to_unreshard_results,
+    _unresharded_code::Vector{Expr},
+    _unresharded_arrays_cache,
+    used_shardinfo,
+    result_cache,
+    var_idx,
+    resultgen_code,
+)
+    if !haskey(result_cache, tocopy)
+        sym = Symbol("result", var_idx[])
+        var_idx[] += 1
+
+        result = Meta.quot(tocopy)
+
+        push!(
+            resultgen_code,
+            quote
+                $sym = $result
+            end,
+        )
+        result_cache[tocopy] = sym
+    end
+
+    return result_cache[tocopy]
+end
+
 function create_result(
     tocopy::ConcretePJRTNumber{T,D},
     @nospecialize(path::Tuple),
@@ -823,14 +854,14 @@ end
 # However, this errs as we cannot attach the transform with to the funcop itself [as we run a functionpass].
 const enzyme_pass::String = "enzyme{postpasses=\"arith-raise{stablehlo=true},canonicalize,cse,canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math,canonicalize,cse,canonicalize,arith-raise{stablehlo=true}\"}"
 
-function probprog_pass(;
+function impulse_pass(;
     debug_dump::Bool=DEBUG_PROBPROG_DUMP_VALUE[],
     disable_optimizations::Bool=DEBUG_PROBPROG_DISABLE_OPT[],
 )
     if !disable_optimizations
-        # TODO(#2063): Add probprog optimization passes
+        # TODO(#2063): Add impulse optimization passes
     end
-    return "probprog{debug-dump=$debug_dump postpasses=\"arith-raise{stablehlo=true}\"}"
+    return "expand-impulse{debug-dump=$debug_dump postpasses=\"arith-raise{stablehlo=true}\"}"
 end
 
 function run_pass_pipeline!(mod, pass_pipeline, key=""; enable_verifier=true)
@@ -1575,8 +1606,8 @@ function compile_mlir!(
                         raise_passes,
                         "enzyme-batch",
                         opt_passes2,
-                        probprog_pass(),
-                        "lower-probprog-to-stablehlo{backend=$backend}",
+                        impulse_pass(),
+                        "lower-impulse-to-stablehlo{backend=$backend}",
                         "outline-enzyme-regions",
                         enzyme_pass,
                         opt_passes2,
@@ -1592,7 +1623,7 @@ function compile_mlir!(
                         )...,
                         opt_passes2,
                         lower_enzymexla_passes,
-                        "lower-probprog-trace-ops{backend=$backend}",
+                        "lower-impulse-trace-ops{backend=$backend}",
                         jit,
                     ]
                 else
@@ -1601,8 +1632,8 @@ function compile_mlir!(
                         opt_passes,
                         "enzyme-batch",
                         opt_passes2,
-                        probprog_pass(),
-                        "lower-probprog-to-stablehlo{backend=$backend}",
+                        impulse_pass(),
+                        "lower-impulse-to-stablehlo{backend=$backend}",
                         "outline-enzyme-regions",
                         enzyme_pass,
                         opt_passes2,
@@ -1620,13 +1651,13 @@ function compile_mlir!(
                         kern,
                         raise_passes,
                         lower_enzymexla_passes,
-                        "lower-probprog-trace-ops{backend=$backend}",
+                        "lower-impulse-trace-ops{backend=$backend}",
                         jit,
                     ]
                 end,
                 ",",
             ),
-            "probprog",
+            "impulse",
         )
     elseif compile_options.optimization_passes === :only_enzyme
         run_pass_pipeline!(

src/Enzyme.jl

@@ -1,9 +1,11 @@
-const enzyme_out = 0
-const enzyme_dup = 1
-const enzyme_const = 2
-const enzyme_dupnoneed = 3
-const enzyme_outnoneed = 4
-const enzyme_constnoneed = 5
+@enumx EnzymeActivity begin
+    OUT = 0
+    DUPLICATED = 1
+    CONST = 2
+    DUPLICATED_NO_NEED = 3
+    OUT_NO_NEED = 4
+    CONST_NO_NEED = 5
+end
 
 struct StackedBatchDuplicated{T,N,M,V<:AbstractArray{T,N},W<:AbstractArray{T,M}} <:
        Annotation{V}
@@ -201,17 +203,17 @@ end
 end
 
 @inline function act_from_type(::Type{<:Active}, reverse, needs_primal)
-    return needs_primal ? enzyme_out : enzyme_outnoneed
+    return needs_primal ? EnzymeActivity.OUT : EnzymeActivity.OUT_NO_NEED
 end
 @inline function act_from_type(::Type{<:Const}, reverse, needs_primal)
-    return needs_primal ? enzyme_const : enzyme_constnoneed
+    return needs_primal ? EnzymeActivity.CONST : EnzymeActivity.CONST_NO_NEED
 end
 
 @inline function act_from_type(::Type{<:Duplicated}, reverse, needs_primal)
     if reverse
-        return needs_primal ? enzyme_out : enzyme_outnoneed
+        return needs_primal ? EnzymeActivity.OUT : EnzymeActivity.OUT_NO_NEED
     else
-        return needs_primal ? enzyme_dup : enzyme_dupnoneed
+        return needs_primal ? EnzymeActivity.DUPLICATED : EnzymeActivity.DUPLICATED_NO_NEED
     end
 end
 @inline function act_from_type(
@@ -221,7 +223,7 @@ end
 end
 
 @inline function act_from_type(::Type{<:DuplicatedNoNeed}, reverse, needs_primal)
-    return reverse ? enzyme_out : enzyme_dupnoneed
+    return reverse ? EnzymeActivity.OUT : EnzymeActivity.DUPLICATED_NO_NEED
 end
 @inline function act_from_type(
     ::Type{<:Union{BatchDuplicatedNoNeed,StackedBatchDuplicatedNoNeed}},
@@ -289,7 +291,9 @@ function set_act!(inp, path, reverse, tostore; emptypath=false, width=1)
 end
 
 function act_attr(val)
-    return MLIR.IR.Attribute(MLIR.API.enzymeActivityAttrGet(MLIR.IR.current_context(), val))
+    return MLIR.IR.Attribute(
+        MLIR.API.enzymeActivityAttrGet(MLIR.IR.current_context(), Int32(val))
+    )
 end
 
 function infer_activity(
@@ -342,7 +346,7 @@ function overload_autodiff(
     (; result, linear_args, in_tys, linear_results) = mlir_fn_res
     fnwrap = mlir_fn_res.fnwrapped
 
-    activity = Int32[]
+    activity = EnzymeActivity.T[]
     ad_inputs = MLIR.IR.Value[]
 
     reverse_seeds = Dict{Tuple,MLIR.IR.Value}()
@@ -353,7 +357,7 @@ function overload_autodiff(
         push!(activity, act_from_type(arg, reverse))
         push_acts!(ad_inputs, arg, path[3:end], reverse)
 
-        if CMode <: ReverseMode && act_from_type(arg, false) == enzyme_dup
+        if CMode <: ReverseMode && act_from_type(arg, false) == EnzymeActivity.DUPLICATED
             x = if width == 1
                 arg.dval
             elseif arg.dval isa AbstractArray
@@ -370,7 +374,7 @@ function overload_autodiff(
     end
 
     outtys = MLIR.IR.Type[]
-    ret_activity = Int32[]
+    ret_activity = EnzymeActivity.T[]
 
     for a in linear_results
         if TracedUtils.has_idx(a, resprefix)
@@ -395,7 +399,7 @@ function overload_autodiff(
 
             act = act_from_type(A, reverse, EnzymeCore.needs_primal(CMode))
             cst = nothing
-            if act == enzyme_out || act == enzyme_outnoneed
+            if act == EnzymeActivity.OUT || act == EnzymeActivity.OUT_NO_NEED
                 if width == 1
                     cst = @opcall fill(one(unwrapped_eltype(a)), size(a))
                 else
@@ -407,26 +411,28 @@ function overload_autodiff(
             if CMode <: ReverseMode && TracedUtils.has_idx(a, argprefix)
                 idx, path = TracedUtils.get_argidx(a, argprefix)
                 arg = idx == 1 && fnwrap ? f : args[idx - fnwrap]
-                if act_from_type(arg, false) == enzyme_dup
+                if act_from_type(arg, false) == EnzymeActivity.DUPLICATED
                     seed = reverse_seeds[path]
-                    if cst == nothing
-                        if act == enzyme_const
-                            act = enzyme_out
-                        elseif act == enzyme_constnoneed
-                            act = enzyme_outnoneed
+                    if cst === nothing
+                        if act == EnzymeActivity.CONST
+                            act = EnzymeActivity.OUT
+                        elseif act == EnzymeActivity.CONST_NO_NEED
+                            act = EnzymeActivity.OUT_NO_NEED
                         else
                             @assert false
                         end
                         cst = seed
                     else
-                        @assert act == enzyme_out || act == enzyme_outnoneed
+                        @assert (
+                            act == EnzymeActivity.OUT || act == EnzymeActivity.OUT_NO_NEED
+                        )
                         cst = MLIR.IR.result(MLIR.Dialects.stablehlo.add(cst, seed), 1)
                     end
                 end
             end
 
             push!(ret_activity, act)
-            if cst != nothing
+            if cst !== nothing
                 push!(ad_inputs, cst)
             end
         else
@@ -437,7 +443,7 @@ function overload_autodiff(
                 act = act_from_type(arg, reverse, true)
                 push!(ret_activity, act)
 
-                if act == enzyme_out || act == enzyme_outnoneed
+                if act == EnzymeActivity.OUT || act == EnzymeActivity.OUT_NO_NEED
                     seed = reverse_seeds[path]
                     push!(ad_inputs, seed)
                 end
@@ -453,7 +459,11 @@ function overload_autodiff(
     end
 
     for (i, act) in enumerate(activity)
-        if act == enzyme_out || act == enzyme_dup || act == enzyme_dupnoneed
+        if (
+            act == EnzymeActivity.OUT ||
+            act == EnzymeActivity.DUPLICATED ||
+            act == EnzymeActivity.DUPLICATED_NO_NEED
+        )
             push!(outtys, TracedUtils.batch_ty(width, in_tys[i]))
         end
     end
@@ -528,7 +538,7 @@ function overload_autodiff(
         idx, path = TracedUtils.get_argidx(a, argprefix)
 
         arg = idx == 1 && fnwrap ? f : args[idx - fnwrap]
-        act_from_type(arg, reverse) != enzyme_out && continue
+        act_from_type(arg, reverse) != EnzymeActivity.OUT && continue
 
         if idx == 1 && fnwrap && arg isa Active
             @assert false

src/Ops.jl

@@ -650,6 +650,8 @@ end
     dims::Vector{Int};
     location=mlir_stacktrace("reshape", @__FILE__, @__LINE__),
 ) where {T,N}
+    @assert length(x) == prod(dims)
+
     # HLO reshape semantics collapse the opposite way
     res1 = transpose(x, Int64[N:-1:1...])
     restype = mlir_type(TracedRArray{T,length(dims)}, collect(Int64, Base.reverse(dims)))