|
| 1 | +module LuxLibReactantExt |
| 2 | + |
| 3 | +using Reactant: Reactant, MLIR, Ops, TracedUtils, TracedRArray, AnyTracedRArray, |
| 4 | + AnyTracedRVector, TracedRNumber |
| 5 | +using Static: StaticBool, True, False |
| 6 | + |
| 7 | +using LuxLib: LuxLib, Impl, Optional, Utils |
| 8 | + |
| 9 | +# Most of the NN code gen happens in Reactant.jl via an extension on NNlib, however, |
| 10 | +# NNlib doesn't have certain ops implemented. In those cases we can emit more optimized |
| 11 | +# StableHLO |
| 12 | + |
| 13 | +function Impl.batchnorm( |
| 14 | + x::AnyTracedRArray{T}, |
| 15 | + γ::Optional{<:AnyTracedRVector}, β::Optional{<:AnyTracedRVector}, |
| 16 | + rμ::Optional{<:AnyTracedRVector}, rσ²::Optional{<:AnyTracedRVector}, |
| 17 | + training::StaticBool, act::F, momentum, ϵ |
| 18 | +) where {T, F} |
| 19 | + x = TracedUtils.materialize_traced_array(x) |
| 20 | + |
| 21 | + γ = if γ === nothing |
| 22 | + Ops.constant(fill(T(1), size(x, ndims(x) - 1))) |
| 23 | + else |
| 24 | + TracedUtils.materialize_traced_array(γ) |
| 25 | + end |
| 26 | + β = if β === nothing |
| 27 | + Ops.constant(fill(T(0), size(x, ndims(x) - 1))) |
| 28 | + else |
| 29 | + TracedUtils.materialize_traced_array(β) |
| 30 | + end |
| 31 | + |
| 32 | + if training isa True |
| 33 | + op = MLIR.Dialects.stablehlo.batch_norm_training( |
| 34 | + TracedUtils.get_mlir_data(x), |
| 35 | + TracedUtils.get_mlir_data(γ), |
| 36 | + TracedUtils.get_mlir_data(β); |
| 37 | + epsilon=Float32(ϵ), |
| 38 | + feature_index=Int64(ndims(x) - 2) |
| 39 | + ) |
| 40 | + |
| 41 | + res = act.(TracedRArray{T, ndims(x)}((), MLIR.IR.result(op, 1), size(x))) |
| 42 | + μ = TracedRArray{T, 1}((), MLIR.IR.result(op, 2), size(x, ndims(x) - 1)) |
| 43 | + σ² = TracedRArray{T, 1}((), MLIR.IR.result(op, 3), size(x, ndims(x) - 1)) |
| 44 | + |
| 45 | + if rμ === nothing && rσ² === nothing |
| 46 | + return res, nothing, nothing |
| 47 | + else |
| 48 | + @assert rμ !== nothing && rσ² !== nothing |
| 49 | + m = T(Impl.accum_size(x, Impl.batchnorm_reduce_dims(x))) |
| 50 | + rμ, rσ² = Impl.update_running_statistics( |
| 51 | + rμ, rσ², μ, σ², momentum, momentum * m / (m - one(m)) |
| 52 | + ) |
| 53 | + return res, rμ, rσ² |
| 54 | + end |
| 55 | + else |
| 56 | + if rμ === nothing && rσ² === nothing |
| 57 | + μ, σ² = Impl.mean_var( |
| 58 | + x; dims=Utils.unsafe_known(Impl.batchnorm_reduce_dims(x)), corrected=false |
| 59 | + ) |
| 60 | + μ = TracedUtils.materialize_traced_array(vec(μ)) |
| 61 | + σ² = TracedUtils.materialize_traced_array(vec(σ²)) |
| 62 | + else |
| 63 | + @assert rμ !== nothing && rσ² !== nothing |
| 64 | + μ = TracedUtils.materialize_traced_array(rμ) |
| 65 | + σ² = TracedUtils.materialize_traced_array(rσ²) |
| 66 | + end |
| 67 | + |
| 68 | + res = MLIR.IR.result( |
| 69 | + MLIR.Dialects.stablehlo.batch_norm_inference( |
| 70 | + TracedUtils.get_mlir_data(x), |
| 71 | + TracedUtils.get_mlir_data(γ), |
| 72 | + TracedUtils.get_mlir_data(β), |
| 73 | + TracedUtils.get_mlir_data(μ), |
| 74 | + TracedUtils.get_mlir_data(σ²); |
| 75 | + epsilon=Float32(ϵ), |
| 76 | + feature_index=Int64(ndims(x) - 2) |
| 77 | + ), |
| 78 | + 1 |
| 79 | + ) |
| 80 | + |
| 81 | + return act.(TracedRArray{T, ndims(x)}((), res, size(x))), rμ, rσ² |
| 82 | + end |
| 83 | +end |
| 84 | + |
| 85 | +end |
0 commit comments