feat: use new enum based API (#2766)

* Regenerate MLIR Bindings

* fix: update to use enums

* fix: workaround raising issue

* fix: more

* test: syrk tests are busted

---------

Co-authored-by: enzyme-ci-bot[bot] <78882869+enzyme-ci-bot[bot]@users.noreply.github.com>

Author: avik-pal

ext/ReactantLogExpFunctionsExt.jl

@@ -96,11 +96,16 @@ end
 function LogExpFunctions.softmax!(
     r::AnyTracedRArray{<:Real}, x::AnyTracedRArray{<:Real}=r; dims=:
 )
-    return LogExpFunctions._softmax!(r, x, dims)
+    dims isa Colon && (dims = 1:ndims(x))
+    res = @opcall softmax(x; dims=vec(collect(Int64, dims)))
+    copyto!(r, res)
+    return r
 end
 
 function LogExpFunctions.softmax(x::AnyTracedRArray{<:Real}; dims=:)
-    return LogExpFunctions._softmax!(similar(x, float(eltype(x))), x, dims)
+    dims isa Colon && (dims = 1:ndims(x))
+    res = @opcall softmax(x; dims=vec(collect(Int64, dims)))
+    return res
 end
 
 for (T1, T2) in [(TracedRNumber, Number), (Number, TracedRNumber)]

ext/ReactantMPIExt/Ops.jl

@@ -184,36 +184,20 @@ end
     return recvbuf
 end
 
-@enum MPIOpEnum begin
-    MPI_OP_NULL_ENUM = 0
-    MPI_BAND_ENUM = 1
-    MPI_BOR_ENUM = 2
-    MPI_BXOR_ENUM = 3
-    MPI_LAND_ENUM = 4
-    MPI_LOR_ENUM = 5
-    MPI_LXOR_ENUM = 6
-    MPI_MAX_ENUM = 7
-    MPI_MIN_ENUM = 8
-    MPI_PROD_ENUM = 9
-    MPI_REPLACE_ENUM = 10
-    MPI_SUM_ENUM = 11
-    MPI_NO_OP_ENUM = 12
-end
-
 const MPI_OP_MAP = Dict(
-    MPI.OP_NULL.val => MPI_OP_NULL_ENUM,
-    MPI.BAND.val => MPI_BAND_ENUM,
-    MPI.BOR.val => MPI_BOR_ENUM,
-    MPI.BXOR.val => MPI_BXOR_ENUM,
-    MPI.LAND.val => MPI_LAND_ENUM,
-    MPI.LOR.val => MPI_LOR_ENUM,
-    MPI.LXOR.val => MPI_LXOR_ENUM,
-    MPI.MAX.val => MPI_MAX_ENUM,
-    MPI.MIN.val => MPI_MIN_ENUM,
-    MPI.PROD.val => MPI_PROD_ENUM,
-    MPI.REPLACE.val => MPI_REPLACE_ENUM,
-    MPI.SUM.val => MPI_SUM_ENUM,
-    MPI.NO_OP.val => MPI_NO_OP_ENUM,
+    MPI.OP_NULL.val => MLIR.API.ENZYMEXLA_MPI_OP_NULL,
+    MPI.BAND.val => MLIR.API.ENZYMEXLA_MPI_BAND,
+    MPI.BOR.val => MLIR.API.ENZYMEXLA_MPI_BOR,
+    MPI.BXOR.val => MLIR.API.ENZYMEXLA_MPI_BXOR,
+    MPI.LAND.val => MLIR.API.ENZYMEXLA_MPI_LAND,
+    MPI.LOR.val => MLIR.API.ENZYMEXLA_MPI_LOR,
+    MPI.LXOR.val => MLIR.API.ENZYMEXLA_MPI_LXOR,
+    MPI.MAX.val => MLIR.API.ENZYMEXLA_MPI_MAX,
+    MPI.MIN.val => MLIR.API.ENZYMEXLA_MPI_MIN,
+    MPI.PROD.val => MLIR.API.ENZYMEXLA_MPI_PROD,
+    MPI.REPLACE.val => MLIR.API.ENZYMEXLA_MPI_REPLACE,
+    MPI.SUM.val => MLIR.API.ENZYMEXLA_MPI_SUM,
+    MPI.NO_OP.val => MLIR.API.ENZYMEXLA_MPI_NO_OP,
 )
 
 function get_mpi_op_enum(op)
@@ -222,64 +206,34 @@ function get_mpi_op_enum(op)
     end
 end
 
-@enum MPIDataTypeEnum begin
-    MPI_DATATYPE_NULL_ENUM = 0
-    MPI_INT8_T_ENUM = 1
-    MPI_UINT8_T_ENUM = 2
-    MPI_INT16_T_ENUM = 3
-    MPI_UINT16_T_ENUM = 4
-    MPI_INT32_T_ENUM = 5
-    MPI_UINT32_T_ENUM = 6
-    MPI_INT64_T_ENUM = 7
-    MPI_UINT64_T_ENUM = 8
-    MPI_BYTE_ENUM = 9
-    MPI_SHORT_ENUM = 10
-    MPI_UNSIGNED_SHORT_ENUM = 11
-    MPI_INT_ENUM = 12
-    MPI_UNSIGNED_ENUM = 13
-    MPI_LONG_ENUM = 14
-    MPI_UNSIGNED_LONG_ENUM = 15
-    MPI_LONG_LONG_INT_ENUM = 16
-    MPI_UNSIGNED_LONG_LONG_ENUM = 17
-    MPI_CHAR_ENUM = 18
-    MPI_SIGNED_CHAR_ENUM = 19
-    MPI_UNSIGNED_CHAR_ENUM = 20
-    MPI_WCHAR_ENUM = 21
-    MPI_FLOAT_ENUM = 22
-    MPI_DOUBLE_ENUM = 23
-    MPI_C_FLOAT_COMPLEX_ENUM = 24
-    MPI_C_DOUBLE_COMPLEX_ENUM = 25
-    MPI_C_BOOL_ENUM = 26
-end
-
 const MPI_DATATYPE_MAP = Dict(
-    MPI.DATATYPE_NULL.val => MPI_DATATYPE_NULL_ENUM,
-    MPI.INT8_T.val => MPI_INT8_T_ENUM,
-    MPI.UINT8_T.val => MPI_UINT8_T_ENUM,
-    MPI.INT16_T.val => MPI_INT16_T_ENUM,
-    MPI.UINT16_T.val => MPI_UINT16_T_ENUM,
-    MPI.INT32_T.val => MPI_INT32_T_ENUM,
-    MPI.UINT32_T.val => MPI_UINT32_T_ENUM,
-    MPI.INT64_T.val => MPI_INT64_T_ENUM,
-    MPI.UINT64_T.val => MPI_UINT64_T_ENUM,
-    MPI.BYTE.val => MPI_BYTE_ENUM,
-    MPI.SHORT.val => MPI_SHORT_ENUM,
-    MPI.UNSIGNED_SHORT.val => MPI_UNSIGNED_SHORT_ENUM,
-    MPI.INT.val => MPI_INT_ENUM,
-    MPI.UNSIGNED.val => MPI_UNSIGNED_ENUM,
-    MPI.LONG.val => MPI_LONG_ENUM,
-    MPI.UNSIGNED_LONG.val => MPI_UNSIGNED_LONG_ENUM,
-    MPI.LONG_LONG_INT.val => MPI_LONG_LONG_INT_ENUM,
-    MPI.UNSIGNED_LONG_LONG.val => MPI_UNSIGNED_LONG_LONG_ENUM,
-    MPI.CHAR.val => MPI_CHAR_ENUM,
-    MPI.SIGNED_CHAR.val => MPI_SIGNED_CHAR_ENUM,
-    MPI.UNSIGNED_CHAR.val => MPI_UNSIGNED_CHAR_ENUM,
-    MPI.WCHAR.val => MPI_WCHAR_ENUM,
-    MPI.FLOAT.val => MPI_FLOAT_ENUM,
-    MPI.DOUBLE.val => MPI_DOUBLE_ENUM,
-    MPI.C_FLOAT_COMPLEX.val => MPI_C_FLOAT_COMPLEX_ENUM,
-    MPI.C_DOUBLE_COMPLEX.val => MPI_C_DOUBLE_COMPLEX_ENUM,
-    MPI.C_BOOL.val => MPI_C_BOOL_ENUM,
+    MPI.DATATYPE_NULL.val => MLIR.API.ENZYMEXLA_MPI_DATATYPE_NULL,
+    MPI.INT8_T.val => MLIR.API.ENZYMEXLA_MPI_INT8_T,
+    MPI.UINT8_T.val => MLIR.API.ENZYMEXLA_MPI_UINT8_T,
+    MPI.INT16_T.val => MLIR.API.ENZYMEXLA_MPI_INT16_T,
+    MPI.UINT16_T.val => MLIR.API.ENZYMEXLA_MPI_UINT16_T,
+    MPI.INT32_T.val => MLIR.API.ENZYMEXLA_MPI_INT32_T,
+    MPI.UINT32_T.val => MLIR.API.ENZYMEXLA_MPI_UINT32_T,
+    MPI.INT64_T.val => MLIR.API.ENZYMEXLA_MPI_INT64_T,
+    MPI.UINT64_T.val => MLIR.API.ENZYMEXLA_MPI_UINT64_T,
+    MPI.BYTE.val => MLIR.API.ENZYMEXLA_MPI_BYTE,
+    MPI.SHORT.val => MLIR.API.ENZYMEXLA_MPI_SHORT,
+    MPI.UNSIGNED_SHORT.val => MLIR.API.ENZYMEXLA_MPI_UNSIGNED_SHORT,
+    MPI.INT.val => MLIR.API.ENZYMEXLA_MPI_INT,
+    MPI.UNSIGNED.val => MLIR.API.ENZYMEXLA_MPI_UNSIGNED,
+    MPI.LONG.val => MLIR.API.ENZYMEXLA_MPI_LONG,
+    MPI.UNSIGNED_LONG.val => MLIR.API.ENZYMEXLA_MPI_UNSIGNED_LONG,
+    MPI.LONG_LONG_INT.val => MLIR.API.ENZYMEXLA_MPI_LONG_LONG_INT,
+    MPI.UNSIGNED_LONG_LONG.val => MLIR.API.ENZYMEXLA_MPI_UNSIGNED_LONG_LONG,
+    MPI.CHAR.val => MLIR.API.ENZYMEXLA_MPI_CHAR,
+    MPI.SIGNED_CHAR.val => MLIR.API.ENZYMEXLA_MPI_SIGNED_CHAR,
+    MPI.UNSIGNED_CHAR.val => MLIR.API.ENZYMEXLA_MPI_UNSIGNED_CHAR,
+    MPI.WCHAR.val => MLIR.API.ENZYMEXLA_MPI_WCHAR,
+    MPI.FLOAT.val => MLIR.API.ENZYMEXLA_MPI_FLOAT,
+    MPI.DOUBLE.val => MLIR.API.ENZYMEXLA_MPI_DOUBLE,
+    MPI.C_FLOAT_COMPLEX.val => MLIR.API.ENZYMEXLA_MPI_C_FLOAT_COMPLEX,
+    MPI.C_DOUBLE_COMPLEX.val => MLIR.API.ENZYMEXLA_MPI_C_DOUBLE_COMPLEX,
+    MPI.C_BOOL.val => MLIR.API.ENZYMEXLA_MPI_C_BOOL,
 )
 
 function get_mpi_datatype_enum(datatype)

ext/ReactantNNlibExt/Implementations.jl

@@ -27,31 +27,15 @@ end
 
 function NNlib.softmax!(out::AnyTracedRArray{T,N}, x::AbstractArray; dims=1) where {T,N}
     x = T.(materialize_traced_array(x))
-    max_ = maximum(x; dims)
-    diff = exp.(x .- max_)
-    # TOOD: re-enable conditional once https://github.com/EnzymeAD/Reactant.jl/issues/1581
-    # fixed
-    # @trace if all(isfinite, max_)
-    @. out = diff
-    # else
-    #     @. out = ifelse(isinf(max_), ifelse(isinf(x), T(1), T(0)), diff)
-    # end
-    out ./= sum(out; dims)
+    res = @opcall softmax(x; dims=vec(collect(Int64, dims)))
+    copyto!(out, res)
     return out
 end
 
 function NNlib.logsoftmax!(out::AnyTracedRArray{T}, x::AbstractArray; dims=1) where {T}
     x = T.(materialize_traced_array(x))
-    max_ = maximum(x; dims)
-    diff = x .- max_
-    # TOOD: re-enable conditional once https://github.com/EnzymeAD/Reactant.jl/issues/1581
-    # fixed
-    # @trace if all(isfinite, max_)
-    @. out = diff
-    # else
-    #     @. out = ifelse(isinf(max_), ifelse(isinf(x), T(0), -T(Inf)), diff)
-    # end
-    out .-= log.(sum(exp, out; dims))
+    res = @opcall logsoftmax(x; dims=vec(collect(Int64, dims)))
+    copyto!(out, res)
     return out
 end
 

src/Ops.jl

@@ -15,6 +15,31 @@ using ..Reactant:
 using ReactantCore: ReactantCore
 using GPUArraysCore: GPUArraysCore
 
+const GELU_APPROXIMATION_MAP = Dict(
+    "NONE" => MLIR.API.ENZYMEXLA_GELU_APPROXIMATION_NONE,
+    "TANH" => MLIR.API.ENZYMEXLA_GELU_APPROXIMATION_TANH,
+    "SIGMOID" => MLIR.API.ENZYMEXLA_GELU_APPROXIMATION_SIGMOID,
+)
+
+const LAPACK_TRANSPOSE_MAP = Dict(
+    'N' => MLIR.API.ENZYMEXLA_LAPACK_TRANSPOSE_NONE,
+    'T' => MLIR.API.ENZYMEXLA_LAPACK_TRANSPOSE_TRANSPOSE,
+    'C' => MLIR.API.ENZYMEXLA_LAPACK_TRANSPOSE_CONJUGATE_TRANSPOSE,
+)
+
+const LAPACK_UPLO_MAP = Dict(
+    'U' => MLIR.API.ENZYMEXLA_LAPACK_UPLO_UPPER,
+    'L' => MLIR.API.ENZYMEXLA_LAPACK_UPLO_LOWER,
+    'F' => MLIR.API.ENZYMEXLA_LAPACK_UPLO_FULL,
+)
+
+const SVD_ALGORITHM_MAP = Dict(
+    "DEFAULT" => MLIR.API.ENZYMEXLA_SVD_ALGORITHM_NONE,
+    "QRIteration" => MLIR.API.ENZYMEXLA_SVD_ALGORITHM_QRITERATION,
+    "Jacobi" => MLIR.API.ENZYMEXLA_SVD_ALGORITHM_JACOBI,
+    "DivideAndConquer" => MLIR.API.ENZYMEXLA_SVD_ALGORITHM_DIVIDEANDCONQUER,
+)
+
 function _function_macro_error()
     throw(ArgumentError("`caller_function` is not available in this context"))
 end
@@ -3876,26 +3901,16 @@ end
     Vt_size = (batch_sizes..., full ? n : r, n)
     info_size = batch_sizes
 
-    if algorithm == "DEFAULT"
-        algint = 0
-    elseif algorithm == "QRIteration"
-        algint = 1
-    elseif algorithm == "DivideAndConquer"
-        algint = 2
-    elseif algorithm == "Jacobi"
-        algint = 3
-    else
-        error("Unsupported SVD algorithm: $algorithm")
-    end
-
     svd_op = enzymexla.linalg_svd(
         x.mlir_data;
         U=mlir_type(TracedRArray{T,N}, U_size),
         S=mlir_type(TracedRArray{Base.real(T),N - 1}, S_size),
         Vt=mlir_type(TracedRArray{T,N}, Vt_size),
         info=mlir_type(TracedRArray{iT,N - 2}, info_size),
         full=full,
-        algorithm=MLIR.API.enzymexlaSVDAlgorithmAttrGet(MLIR.IR.current_context(), algint),
+        algorithm=MLIR.API.enzymexlaSVDAlgorithmAttrGet(
+            MLIR.IR.current_context(), SVD_ALGORITHM_MAP[algorithm]
+        ),
         location,
     )
 
@@ -4103,21 +4118,15 @@ end
     approximation::String;
     location=mlir_stacktrace("ml.gelu", @__FILE__, @__LINE__),
 )
-    approx = if approximation == "NONE"
-        0
-    elseif approximation == "TANH"
-        1
-    elseif approximation == "SIGMOID"
-        2
-    else
-        error("Invalid gelu approximation: $approximation")
-    end
-    approx = MLIR.API.enzymexlaGeluApproximationAttrGet(
-        MLIR.IR.current_context(), Int32(approx)
-    )
-
     res = MLIR.IR.result(
-        enzymexla.ml_gelu(x.mlir_data; gelu_approximation=approx, location), 1
+        enzymexla.ml_gelu(
+            x.mlir_data;
+            gelu_approximation=MLIR.API.enzymexlaGeluApproximationAttrGet(
+                MLIR.IR.current_context(), GELU_APPROXIMATION_MAP[approximation]
+            ),
+            location,
+        ),
+        1,
     )
 
     if x isa TracedRArray
@@ -4210,41 +4219,19 @@ end
     location=mlir_stacktrace("syrk", @__FILE__, @__LINE__),
 ) where {T,N}
     ctx = MLIR.IR.current_context()
-    uplo_attr = MLIR.API.enzymexlaLapackUploAttrGet(
-        ctx,
-        if uplo == 'U'
-            Int32(1)
-        elseif uplo == 'L'
-            Int32(0)
-        else
-            Int32(2)
-        end,
-    )
-    transpose_attr = MLIR.API.enzymexlaLapackTransposeAttrGet(
-        ctx,
-        if transpose_a == 'N'
-            Int32(0)
-        elseif transpose_a == 'T'
-            Int32(1)
-        elseif transpose_a == 'C'
-            Int32(2)
-        else
-            error("Unknown transpose mode: $transpose_a")
-        end,
-    )
-
-    alpha_ = constant(alpha; location)
-    beta_ = constant(beta; location)
+    uplo_attr = MLIR.API.enzymexlaLapackUploAttrGet(ctx, LAPACK_UPLO_MAP[uplo])
 
     res = MLIR.IR.result(
         enzymexla.blas_syrk(
             A.mlir_data,
             C.mlir_data,
-            alpha_.mlir_data,
-            beta_.mlir_data;
+            constant(alpha; location).mlir_data,
+            constant(beta; location).mlir_data;
             uplo=uplo_attr,
             output_uplo=uplo_attr,
-            transpose=transpose_attr,
+            transpose=MLIR.API.enzymexlaLapackTransposeAttrGet(
+                ctx, LAPACK_TRANSPOSE_MAP[transpose_a]
+            ),
             output=mlir_type(TracedRArray{T,N}, size(C)),
             location,
         ),
@@ -4510,4 +4497,29 @@ function julia_callback(
     return Tuple(results)
 end
 
+@noinline function softmax(
+    x::TracedRArray{T,N};
+    dims::Vector{Int64},
+    location=mlir_stacktrace("softmax", @__FILE__, @__LINE__),
+) where {T,N}
+    max_val = Reactant.call_with_reactant(Core.kwcall, (; dims,), Base.maximum, x)
+    exp_diff = exponential(x .- max_val; location)
+    denom = Reactant.call_with_reactant(Core.kwcall, (; dims,), Base.sum, exp_diff)
+    return exp_diff ./ denom
+end
+
+@noinline function logsoftmax(
+    x::TracedRArray{T,N};
+    dims::Vector{Int64},
+    location=mlir_stacktrace("logsoftmax", @__FILE__, @__LINE__),
+) where {T,N}
+    max_val = Reactant.call_with_reactant(Core.kwcall, (; dims,), Base.maximum, x)
+    diff = x .- max_val
+    exp_diff = exponential(diff; location)
+    reduced_exp_diff = Reactant.call_with_reactant(
+        Core.kwcall, (; dims,), Base.sum, exp_diff
+    )
+    return diff .- log(reduced_exp_diff; location)
+end
+
 end # module Ops