Merge branch 'main' of https://github.com/EnzymeAD/Reactant.jl into probprog-trace-operand

Author: sbrantq

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.143"
+version = "0.2.144"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -43,8 +43,8 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 YaoBlocks = "418bc28f-b43b-5e0b-a6e7-61bbc1a2c1df"
 
-[sources]
-ReactantCore = {path = "lib/ReactantCore"}
+[sources.ReactantCore]
+path = "lib/ReactantCore"
 
 [extensions]
 ReactantAbstractFFTsExt = "AbstractFFTs"
@@ -90,13 +90,13 @@ PythonCall = "0.9.25"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.15"
-Reactant_jll = "0.0.214"
+Reactant_jll = "0.0.216"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"
 SpecialFunctions = "2.4"
 Statistics = "1.10"
-YaoBlocks = "0.13"
+YaoBlocks = "0.13, 0.14"
 julia = "1.10"
 
 [extras]

deps/ReactantExtra/WORKSPACE

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

src/mlir/Dialects/EnzymeXLA.jl

@@ -219,6 +219,31 @@ function gpu_wrapper(
     )
 end
 
+function ml_gelu(
+    input::Value;
+    result=nothing::Union{Nothing,IR.Type},
+    gelu_approximation,
+    location=Location(),
+)
+    op_ty_results = IR.Type[]
+    operands = Value[input,]
+    owned_regions = Region[]
+    successors = Block[]
+    attributes = NamedAttribute[namedattribute("gelu_approximation", gelu_approximation),]
+    !isnothing(result) && push!(op_ty_results, result)
+
+    return create_operation(
+        "enzymexla.ml.gelu",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=(length(op_ty_results) == 0 ? nothing : op_ty_results),
+        result_inference=(length(op_ty_results) == 0 ? true : false),
+    )
+end
+
 function get_stream(; result::IR.Type, location=Location())
     op_ty_results = IR.Type[result,]
     operands = Value[]
@@ -491,6 +516,26 @@ function linalg_qr(
     )
 end
 
+function ml_relu(input::Value; result=nothing::Union{Nothing,IR.Type}, location=Location())
+    op_ty_results = IR.Type[]
+    operands = Value[input,]
+    owned_regions = Region[]
+    successors = Block[]
+    attributes = NamedAttribute[]
+    !isnothing(result) && push!(op_ty_results, result)
+
+    return create_operation(
+        "enzymexla.ml.relu",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=(length(op_ty_results) == 0 ? nothing : op_ty_results),
+        result_inference=(length(op_ty_results) == 0 ? true : false),
+    )
+end
+
 function rotate(
     operand::Value;
     result=nothing::Union{Nothing,IR.Type},