change to dialect generator

docs/src/api/affine.md

@@ -10,3 +10,7 @@ details.
 ```@autodocs
 Modules = [Reactant.MLIR.Dialects.affine]
 ```
+
+```@docs
+Reactant.MLIR.Dialects.affine.AtomicRMWKind
+```

docs/src/api/arith.md

@@ -10,3 +10,11 @@ more details.
 ```@autodocs
 Modules = [Reactant.MLIR.Dialects.arith]
 ```
+
+```@docs
+Reactant.MLIR.Dialects.arith.CmpFPredicate
+Reactant.MLIR.Dialects.arith.CmpIPredicate
+Reactant.MLIR.Dialects.arith.FastMathFlags
+Reactant.MLIR.Dialects.arith.IntegerOverflowFlags
+Reactant.MLIR.Dialects.arith.RoundingMode
+```

docs/src/api/chlo.md

@@ -10,3 +10,9 @@ for more details.
 ```@autodocs
 Modules = [Reactant.MLIR.Dialects.chlo]
 ```
+
+```@docs
+Reactant.MLIR.Dialects.chlo.ComparisonDirection
+Reactant.MLIR.Dialects.chlo.ComparisonType
+Reactant.MLIR.Dialects.chlo.Precision
+```

docs/src/api/enzyme.md

@@ -7,3 +7,7 @@ CollapsedDocStrings = true
 ```@autodocs
 Modules = [Reactant.MLIR.Dialects.enzyme]
 ```
+
+```@docs
+Reactant.MLIR.Dialects.enzyme.Activity
+```

docs/src/api/gpu.md

@@ -10,3 +10,13 @@ more details.
 ```@autodocs
 Modules = [Reactant.MLIR.Dialects.gpu]
 ```
+
+```@docs
+Reactant.MLIR.Dialects.gpu.AllReduceOperation
+Reactant.MLIR.Dialects.gpu.Dimension
+Reactant.MLIR.Dialects.gpu.MMAElementwiseOp
+Reactant.MLIR.Dialects.gpu.Prune2To4SpMatFlag
+Reactant.MLIR.Dialects.gpu.ShuffleMode
+Reactant.MLIR.Dialects.gpu.SpGEMMWorkEstimationOrComputeKind
+Reactant.MLIR.Dialects.gpu.TransposeMode
+```

docs/src/api/llvm.md

@@ -10,3 +10,15 @@ more details.
 ```@autodocs
 Modules = [Reactant.MLIR.Dialects.llvm]
 ```
+
+```@docs
+Reactant.MLIR.Dialects.llvm.AsmDialect
+Reactant.MLIR.Dialects.llvm.AtomicBinOp
+Reactant.MLIR.Dialects.llvm.AtomicOrdering
+Reactant.MLIR.Dialects.llvm.Comdat
+Reactant.MLIR.Dialects.llvm.FCmpPredicate
+Reactant.MLIR.Dialects.llvm.FastmathFlags
+Reactant.MLIR.Dialects.llvm.ICmpPredicate
+Reactant.MLIR.Dialects.llvm.UnnamedAddr
+Reactant.MLIR.Dialects.llvm.Visibility
+```

docs/src/api/mpi.md

@@ -10,3 +10,8 @@ more details.
 ```@autodocs
 Modules = [Reactant.MLIR.Dialects.mpi]
 ```
+
+```@docs
+Reactant.MLIR.Dialects.mpi.MPI_ErrorClassEnum
+Reactant.MLIR.Dialects.mpi.MPI_OpClassEnum
+```

docs/src/api/nvvm.md

@@ -10,3 +10,26 @@ more details.
 ```@autodocs
 Modules = [Reactant.MLIR.Dialects.nvvm]
 ```
+
+```@docs
+Reactant.MLIR.Dialects.nvvm.FPRoundingMode
+Reactant.MLIR.Dialects.nvvm.LoadCacheModifierKind
+Reactant.MLIR.Dialects.nvvm.MMAB1Op
+Reactant.MLIR.Dialects.nvvm.MMAFrag
+Reactant.MLIR.Dialects.nvvm.MMAIntOverflow
+Reactant.MLIR.Dialects.nvvm.MMALayout
+Reactant.MLIR.Dialects.nvvm.MMATypes
+Reactant.MLIR.Dialects.nvvm.MemScopeKind
+Reactant.MLIR.Dialects.nvvm.ProxyKind
+Reactant.MLIR.Dialects.nvvm.ReduxKind
+Reactant.MLIR.Dialects.nvvm.SaturationMode
+Reactant.MLIR.Dialects.nvvm.SetMaxRegisterAction
+Reactant.MLIR.Dialects.nvvm.SharedSpace
+Reactant.MLIR.Dialects.nvvm.ShflKind
+Reactant.MLIR.Dialects.nvvm.TMAReduxKind
+Reactant.MLIR.Dialects.nvvm.TMAStoreMode
+Reactant.MLIR.Dialects.nvvm.Tcgen05GroupKind
+Reactant.MLIR.Dialects.nvvm.WGMMAScaleIn
+Reactant.MLIR.Dialects.nvvm.WGMMAScaleOut
+Reactant.MLIR.Dialects.nvvm.WGMMATypes
+```

docs/src/api/shardy.md

@@ -9,3 +9,7 @@ Refer to the [official documentation](https://openxla.org/shardy) for more detai
 ```@autodocs
 Modules = [Reactant.MLIR.Dialects.sdy]
 ```
+
+```@docs
+Reactant.MLIR.Dialects.sdy.PropagationDirection
+```

docs/src/api/stablehlo.md

@@ -9,3 +9,14 @@ Refer to the [official documentation](https://openxla.org/stablehlo) for more de
 ```@autodocs
 Modules = [Reactant.MLIR.Dialects.stablehlo]
 ```
+
+```@docs
+Reactant.MLIR.Dialects.stablehlo.ComparisonDirection
+Reactant.MLIR.Dialects.stablehlo.ComparisonType
+Reactant.MLIR.Dialects.stablehlo.CustomCallApiVersion
+Reactant.MLIR.Dialects.stablehlo.FftType
+Reactant.MLIR.Dialects.stablehlo.Precision
+Reactant.MLIR.Dialects.stablehlo.RngAlgorithm
+Reactant.MLIR.Dialects.stablehlo.RngDistribution
+Reactant.MLIR.Dialects.stablehlo.Transpose
+```

docs/src/api/tpu.md

@@ -10,3 +10,11 @@ more details.
 ```@autodocs
 Modules = [Reactant.MLIR.Dialects.tpu]
 ```
+
+```@docs
+Reactant.MLIR.Dialects.tpu.ContractPrecision
+Reactant.MLIR.Dialects.tpu.CoreType
+Reactant.MLIR.Dialects.tpu.PackFormat
+Reactant.MLIR.Dialects.tpu.ReductionKind
+Reactant.MLIR.Dialects.tpu.RoundingMode
+```

docs/src/api/triton.md

@@ -10,3 +10,17 @@ more details.
 ```@autodocs
 Modules = [Reactant.MLIR.Dialects.tt]
 ```
+
+```@docs
+Reactant.MLIR.Dialects.tt.CacheModifier
+Reactant.MLIR.Dialects.tt.EvictionPolicy
+Reactant.MLIR.Dialects.tt.InputPrecision
+Reactant.MLIR.Dialects.tt.MemSemantic
+Reactant.MLIR.Dialects.tt.MemSyncScope
+Reactant.MLIR.Dialects.tt.PaddingOption
+Reactant.MLIR.Dialects.tt.ProgramIDDim
+Reactant.MLIR.Dialects.tt.PropagateNan
+Reactant.MLIR.Dialects.tt.RMWOp
+Reactant.MLIR.Dialects.tt.RoundingMode
+Reactant.MLIR.Dialects.tt.ScaleDotElemType
+```

ext/ReactantAbstractFFTsExt.jl

@@ -1,5 +1,5 @@
 module ReactantAbstractFFTsExt
-
+using Reactant.MLIR.Dialects: stablehlo
 using AbstractFFTs: AbstractFFTs
 using Reactant: Reactant, MLIR, Ops, TracedRArray
 
@@ -31,58 +31,62 @@ function compute_correct_pdims(x::AbstractArray, dims)
     end
 end
 
-for op in (:rfft, :fft, :ifft)
-    mode = uppercase(string(op))
-    @eval function AbstractFFTs.$(op)(x::TracedRArray, dims)
+for op in (stablehlo.FftType.RFFT, stablehlo.FftType.FFT, stablehlo.FftType.IFFT)
+    name = Symbol(lowercase(string(op)))
+    @eval function AbstractFFTs.$(name)(x::TracedRArray, dims)
         @assert maximum(dims) ≤ ndims(x) "dims out of range"
         if dims isa Integer
             if dims != 1
                 pdims = compute_correct_pdims(x, dims)
                 return permutedims(
-                    AbstractFFTs.$(op)(permutedims(x, pdims), 1), invperm(pdims)
+                    AbstractFFTs.$(name)(permutedims(x, pdims), 1), invperm(pdims)
                 )
             end
-            return generalized_fft(x, $(mode), nothing, length(dims))
+            return generalized_fft(x, $(op), nothing, length(dims))
         end
         if !check_contiguous_innermost_dims(dims, ndims(x))
             pdims = compute_correct_pdims(x, dims)
             return permutedims(
-                AbstractFFTs.$(op)(permutedims(x, pdims), 1:length(dims)), invperm(pdims)
+                AbstractFFTs.$(name)(permutedims(x, pdims), 1:length(dims)), invperm(pdims)
             )
         end
-        return generalized_fft(x, $(mode), nothing, length(dims))
+        return generalized_fft(x, $(op), nothing, length(dims))
     end
 end
 
-for op in (:irfft,)
-    mode = uppercase(string(op))
-    @eval function AbstractFFTs.$(op)(x::TracedRArray, d::Int, dims)
+for op in (stablehlo.FftType.IRFFT,)
+    name = Symbol(lowercase(string(op)))
+    @eval function AbstractFFTs.$(name)(x::TracedRArray, d::Int, dims)
         @assert maximum(dims) ≤ ndims(x) "dims out of range"
         if dims isa Integer
             if dims != 1
                 pdims = compute_correct_pdims(x, dims)
                 return permutedims(
-                    AbstractFFTs.$(op)(permutedims(x, pdims), d, 1), invperm(pdims)
+                    AbstractFFTs.$(name)(permutedims(x, pdims), d, 1), invperm(pdims)
                 )
             end
-            return generalized_fft(x, $(mode), d, length(dims))
+            return generalized_fft(x, $(op), d, length(dims))
         end
         if !check_contiguous_innermost_dims(dims, ndims(x))
             pdims = compute_correct_pdims(x, dims)
             return permutedims(
-                AbstractFFTs.$(op)(permutedims(x, pdims), d, 1:length(dims)), invperm(pdims)
+                AbstractFFTs.$(name)(permutedims(x, pdims), d, 1:length(dims)),
+                invperm(pdims),
             )
         end
-        return generalized_fft(x, $(mode), d, length(dims))
+        return generalized_fft(x, $(op), d, length(dims))
     end
 end
 
-function generalized_fft(x::TracedRArray{T,N}, mode::String, d, first_n::Int) where {T,N}
+function generalized_fft(
+    x::TracedRArray{T,N}, mode::stablehlo.FftType.T, d, first_n::Int
+) where {T,N}
     if d === nothing
-        @assert mode ∈ ("RFFT", "FFT", "IFFT")
+        @assert mode ∈
+            (stablehlo.FftType.RFFT, stablehlo.FftType.FFT, stablehlo.FftType.IFFT)
         fft_length = [size(x, i) for i in 1:first_n]
     else
-        @assert mode == "IRFFT"
+        @assert mode == stablehlo.FftType.IRFFT
         fft_length = [i == 1 ? d : size(x, i) for i in 1:first_n]
     end
 

ext/ReactantCUDAExt.jl

@@ -833,7 +833,7 @@ Reactant.@reactant_overlay @noinline function (func::LLVMFunc{F,tt})(
     wrapfunc = MLIR.IR.block!(MLIR.IR.body(mod)) do
         return MLIR.Dialects.llvm.func(;
             sym_name,
-            sym_visibility=MLIR.IR.Attribute("private"),
+            sym_visibility="private",
             function_type=wrapftype,
             body=MLIR.IR.Region(),
             CConv,
@@ -889,10 +889,7 @@ Reactant.@reactant_overlay @noinline function (func::LLVMFunc{F,tt})(
                 1,
             )
             alloc = MLIR.IR.result(
-                MLIR.Dialects.llvm.alloca(
-                    c1; elem_type=MLIR.IR.Attribute(argty), res=llvmptr
-                ),
-                1,
+                MLIR.Dialects.llvm.alloca(c1; elem_type=argty, res=llvmptr), 1
             )
             push!(allocs, (alloc, argty))
 
@@ -953,7 +950,7 @@ Reactant.@reactant_overlay @noinline function (func::LLVMFunc{F,tt})(
                         MLIR.IR.Value[];
                         res=llvmptr,
                         elem_type=i8,
-                        rawConstantIndices=MLIR.IR.Attribute([Int32(offset)]),
+                        rawConstantIndices=[Int32(offset)],
                     ),
                     1,
                 )
@@ -976,13 +973,11 @@ Reactant.@reactant_overlay @noinline function (func::LLVMFunc{F,tt})(
             wrapargs,
             MLIR.IR.Value[];
             callee=MLIR.IR.FlatSymbolRefAttribute(Base.String(fname)),
-            op_bundle_sizes=MLIR.IR.Attribute(Int32[]),
+            op_bundle_sizes=Int32[],
         )
         MLIR.Dialects.llvm.return_(nothing)
     end
 
-    output_operand_aliases = MLIR.IR.Attribute(aliases)
-
     blk_operands = MLIR.IR.Value[]
     for idx in
         (blockdim.x, blockdim.y, blockdim.z, threaddim.x, threaddim.y, threaddim.z, shmem)
@@ -997,9 +992,9 @@ Reactant.@reactant_overlay @noinline function (func::LLVMFunc{F,tt})(
     call = MLIR.Dialects.enzymexla.kernel_call(
         blk_operands...,
         mlir_args;
-        result_0=restys,
+        result=restys,
         fn=MLIR.IR.FlatSymbolRefAttribute(sym_name),
-        output_operand_aliases=MLIR.IR.Attribute(output_operand_aliases),
+        output_operand_aliases=aliases,
     )
 
     argidx = 1

ext/ReactantNNlibExt.jl

@@ -3,7 +3,7 @@ module ReactantNNlibExt
 using NNlib
 using GPUArraysCore: @allowscalar
 using Reactant: Reactant, Ops, TracedRArray, AnyTracedRArray, MLIR, TracedRNumber
-
+using Reactant.MLIR.Dialects: stablehlo
 using Reactant.TracedUtils:
     TracedUtils, materialize_traced_array, get_mlir_data, set_mlir_data!
 
@@ -94,7 +94,7 @@ function NNlib.conv!(
         Int64(output_batch_dim - 1),
         Int64(output_feature_dim - 1),
         length(output_spatial_dims), Int64[i - 1 for i in output_spatial_dims],
-    )
+    )#TODO:deal with this using a custom parser in julia code generation
     #! format: on
 
     padding = Reactant.MLIR.IR.DenseElementsAttribute(
@@ -110,11 +110,11 @@ function NNlib.conv!(
     conv = Reactant.MLIR.Dialects.stablehlo.convolution(
         get_mlir_data(x),
         get_mlir_data(weight);
-        result_0=result_type,
+        result=result_type,
         window_strides=collect(stride),
         padding,
         dimension_numbers,
-        lhs_dilation=1,
+        lhs_dilation=[1 for _ in dilation],
         rhs_dilation=collect(dilation),
         feature_group_count,
         batch_group_count=1,
@@ -176,13 +176,14 @@ function reduce_window(f, x::AnyTracedRArray{T,N}, pdims; init) where {T,N}
         end
 
     attr = fill(Reactant.MLIR.IR.Attribute(init), unranked)
+
     init_value = Reactant.MLIR.IR.result(
         Reactant.MLIR.Dialects.stablehlo.constant(; value=attr)
     )
     reduction = Reactant.MLIR.Dialects.stablehlo.reduce_window(
         [get_mlir_data(x)],
         [init_value];
-        result_0=[result_type],
+        result=[result_type],
         window_dimensions,
         window_strides,
         window_dilations,
@@ -415,7 +416,7 @@ function NNlib.∇conv_filter!(
     conv = MLIR.Dialects.stablehlo.convolution(
         get_mlir_data(x),
         get_mlir_data(dy);
-        result_0=result_type,
+        result=result_type,
         window_strides=collect(dilation),
         padding,
         dimension_numbers,
@@ -532,8 +533,8 @@ function NNlib.∇conv_data!(
     conv = MLIR.Dialects.stablehlo.convolution(
         get_mlir_data(dy),
         get_mlir_data(w);
-        result_0=result_type,
-        window_strides=1,
+        result=result_type,
+        window_strides=[1 for _ in dilation],
         padding,
         lhs_dilation=collect(stride),
         rhs_dilation=collect(dilation),

ext/ReactantRandom123Ext.jl

@@ -2,10 +2,11 @@ module ReactantRandom123Ext
 
 using Random123: Threefry4x, Threefry2x, Philox4x, Philox2x
 using Reactant: TracedRandom
+using Reactant.MLIR.Dialects: stablehlo
 
-TracedRandom.rng_algorithm(::Threefry4x) = "THREE_FRY"
-TracedRandom.rng_algorithm(::Threefry2x) = "THREE_FRY"
-TracedRandom.rng_algorithm(::Philox4x) = "PHILOX"
-TracedRandom.rng_algorithm(::Philox2x) = "PHILOX"
+TracedRandom.rng_algorithm(::Threefry4x) = stablehlo.RngAlgorithm.THREE_FRY
+TracedRandom.rng_algorithm(::Threefry2x) = stablehlo.RngAlgorithm.THREE_FRY
+TracedRandom.rng_algorithm(::Philox4x) = stablehlo.RngAlgorithm.PHILOX
+TracedRandom.rng_algorithm(::Philox2x) = stablehlo.RngAlgorithm.PHILOX
 
 end