Merge branch 'main' into ap/save_sharded_array

Author: glwagner

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.64"
+version = "0.2.66"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -87,7 +87,7 @@ PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.9"
-Reactant_jll = "0.0.123"
+Reactant_jll = "0.0.128"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

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 = "fbb729606494eaa4c8f8f140817d26375e3c79f6"
+ENZYMEXLA_COMMIT = "3f1054ecec69f60575370f873fe88bb1296a18ba"
 ENZYMEXLA_SHA256 = ""
 
 http_archive(

src/Compiler.jl

@@ -459,6 +459,7 @@ end
 
 const WHILE_CONCAT = Ref(false)
 const DUS_TO_CONCAT = Ref(false)
+const SUM_TO_CONV = Ref(false)
 
 # Optimization passes via transform dialect
 function optimization_passes(;
@@ -654,8 +655,19 @@ function optimization_passes(;
         "elementwise_licm(0)",
         "concatenate_licm(0)",
         "slice_broadcast",
+        "while_pad_induction_reduction",
+        "while_licm<1>(1)",
+        "associative_common_mul_op_reordering",
+        "slice_select_to_select_slice",
+        "pad_concat_to_concat_pad",
+        "slice_if",
+        "dus_to_i32",
     ]
 
+    if SUM_TO_CONV[]
+        push!(transform_passes_list, "sum_to_conv")
+    end
+
     if WHILE_CONCAT[]
         push!(transform_passes_list, "while_concat")
     end
@@ -701,7 +713,7 @@ function optimization_passes(;
             [
                 "transpose_while",
                 "transpose_slice",
-                "transpose_elementwise",
+                "transpose_elementwise(0)",
                 "transpose_concat",
                 "transpose_iota",
                 "transpose_reduce",

src/ConcreteRArray.jl

@@ -33,9 +33,7 @@ end
 function Base.deepcopy_internal(
     x::Union{AbstractConcreteArray,AbstractConcreteNumber}, stackdict::IdDict
 )
-    if haskey(stackdict, x)
-        return stackdict[x]::typeof(x)
-    end
+    haskey(stackdict, x) && return stackdict[x]
     return deepcopy(x)
 end
 

src/Ops.jl

@@ -1572,16 +1572,16 @@ julia> Reactant.@jit(
     ftype_attr = MLIR.IR.attr(fn, "function_type")
     ftype = MLIR.IR.Type(ftype_attr)
 
-    @assert all(Base.Fix2(isa, Reactant.AnyTracedRArray), args) "hlo_call: all inputs to hlo_call should be reactant arrays"
+    @assert all(Base.Fix2(isa, Union{TracedRArray,TracedRNumber}), args) "hlo_call: all inputs to hlo_call should be reactant arrays or numbers"
     @assert MLIR.IR.ninputs(ftype) == length(args) "hlo_call: invalid number of arguments for function $func_name"
 
     for (i, arg) in enumerate(args)
         expected_type = MLIR.IR.input(ftype, i)
-        arg_type = MLIR.IR.type(arg.mlir_data)
+        arg_type = MLIR.IR.type(Reactant.TracedUtils.get_mlir_data(arg))
         @assert expected_type == arg_type "hlo_call: argument #$i has the wrong type (expected $expected_type, got $arg_type)"
     end
 
-    operands = [a.mlir_data for a in args]
+    operands = MLIR.IR.Value[Reactant.TracedUtils.get_mlir_data(a) for a in args]
     call = MLIR.Dialects.func.call(
         operands;
         result_0=[MLIR.IR.result(ftype, i) for i in 1:MLIR.IR.nresults(ftype)],

src/mlir/IR/Pass.jl

@@ -145,9 +145,8 @@ end
 Run the provided `passManager` on the given `module`.
 """
 function run!(pm::PassManager, mod::Module, key::String="")
-    # Dump MLIR before running the pass manager.  We set `pm` to nothing because
-    # the pass manager isn't called yet here.
-    DUMP_MLIR_ALWAYS[] && dump_mlir(mod, nothing, isempty(key) ? "pre_pm" : "pre_$(key)_pm")
+    # Dump MLIR before running the pass manager, but also print the list of passes that will be called later.
+    DUMP_MLIR_ALWAYS[] && dump_mlir(mod, pm, isempty(key) ? "pre_pm" : "pre_$(key)_pm")
     status = LogicalResult(@static if isdefined(API, :mlirPassManagerRunOnOp)
         API.mlirPassManagerRunOnOp(pm, Operation(mod))
     else

test/ops.jl

@@ -981,6 +981,25 @@ end
             y_reactant,
         )
     )[1] ≈ x .+ y
+
+    @test Float32(
+        only(
+            Reactant.@jit(
+                Ops.hlo_call(
+                    """
+                    module {
+                      func.func @main(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
+                        %0 = stablehlo.add %arg0, %arg1 : tensor<f32>
+                        return %0 : tensor<f32>
+                      }
+                    }
+                    """,
+                    Reactant.ConcreteRNumber(2.0f0),
+                    Reactant.ConcreteRNumber(2.0f0),
+                )
+            )
+        ),
+    ) == 4.0f0
 end
 
 function f_repeat(x, y)