Merge pull request #591 from LuxDL/ap/fused_ops

Used New Fused Ops from LuxLib

.github/workflows/Downstream.yml

@@ -50,7 +50,7 @@ jobs:
             # force it to use this PR's version of the package
             Pkg.develop(PackageSpec(path="."))  # resolver may fail with main deps
             Pkg.update()
-            Pkg.test()  # resolver may fail with test time deps
+            Pkg.test(; coverage=true)  # resolver may fail with test time deps
           catch err
             err isa Pkg.Resolve.ResolverError || rethrow()
             # If we can't resolve that means this is incompatible by SemVer and this is fine

Project.toml

@@ -65,7 +65,7 @@ ADTypes = "0.2, 1"
 Adapt = "4"
 Aqua = "0.8.4"
 ArrayInterface = "7.9"
-CUDA = "5.2"
+CUDA = "5.3.2"
 ChainRulesCore = "1.23"
 ComponentArrays = "0.15.11"
 ConcreteStructs = "0.2.3"
@@ -84,7 +84,7 @@ LuxAMDGPU = "0.2.2"
 LuxCUDA = "0.3.2"
 LuxCore = "0.1.14"
 LuxDeviceUtils = "0.1.19"
-LuxLib = "0.3.11"
+LuxLib = "0.3.18"
 LuxTestUtils = "0.1.15"
 MLUtils = "0.4.3"
 MPI = "0.20.19"

examples/NeuralODE/main.jl

@@ -171,8 +171,12 @@ function train(model_function; cpu::Bool=false, kwargs...)
 end
 
 train(NeuralODECompact)
+nothing #hide
+
+#-
 
 train(NeuralODE)
+nothing #hide
 
 # We can also change the sensealg and train the model! `GaussAdjoint` allows you to use
 # any arbitrary parameter structure and not just a flat vector (`ComponentArray`).

src/layers/basic.jl

@@ -197,7 +197,7 @@ end
 function initialparameters(rng::AbstractRNG, d::Dense{use_bias}) where {use_bias}
     if use_bias
         return (weight=d.init_weight(rng, d.out_dims, d.in_dims),
-            bias=d.init_bias(rng, d.out_dims, 1))
+            bias=d.init_bias(rng, d.out_dims, 1)) #TODO: In v0.6 make it a vector
     else
         return (weight=d.init_weight(rng, d.out_dims, d.in_dims),)
     end
@@ -210,32 +210,19 @@ statelength(d::Dense) = 0
 
 outputsize(d::Dense) = (d.out_dims,)
 
-@inline function (d::Dense{false})(x::AbstractVecOrMat, ps, st::NamedTuple)
-    return apply_activation(d.activation, ps.weight * x), st
+@inline function (d::Dense)(x::AbstractVector, ps, st::NamedTuple)
+    return vec(first(d(reshape(x, :, 1), ps, st))), st
 end
 
-@inline function (d::Dense{false})(x::AbstractArray, ps, st::NamedTuple)
-    x_reshaped = reshape(x, size(x, 1), :)
+@inline function (d::Dense)(x::AbstractMatrix, ps, st::NamedTuple)
     return (
-        reshape(apply_activation(d.activation, ps.weight * x_reshaped),
-            d.out_dims, size(x)[2:end]...),
+        fused_dense_bias_activation(
+            d.activation, ps.weight, x, _vec(_getproperty(ps, Val(:bias)))),
         st)
 end
 
-@inline function (d::Dense{true})(x::AbstractVector, ps, st::NamedTuple)
-    return apply_bias_activation(d.activation, ps.weight * x, vec(ps.bias)), st
-end
-
-@inline function (d::Dense{true})(x::AbstractMatrix, ps, st::NamedTuple)
-    return apply_bias_activation(d.activation, ps.weight * x, ps.bias), st
-end
-
-@inline function (d::Dense{true})(x::AbstractArray, ps, st::NamedTuple)
-    x_reshaped = reshape(x, size(x, 1), :)
-    return (
-        reshape(apply_bias_activation(d.activation, ps.weight * x_reshaped, ps.bias),
-            d.out_dims, size(x)[2:end]...),
-        st)
+@inline function (d::Dense)(x::AbstractArray, ps, st::NamedTuple)
+    return reshape(first(d(reshape(x, size(x, 1), :), ps, st)), :, size(x)[2:end]...), st
 end
 
 """

src/layers/conv.jl

@@ -114,16 +114,12 @@ function parameterlength(c::Conv{N, use_bias}) where {N, use_bias}
            (use_bias ? c.out_chs : 0)
 end
 
-@inline function (c::Conv{N, false})(x::AbstractArray, ps, st::NamedTuple) where {N}
-    cdims = DenseConvDims(
-        x, ps.weight; stride=c.stride, padding=c.pad, dilation=c.dilation, groups=c.groups)
-    return apply_activation(c.activation, _conv(x, ps.weight, cdims)), st
-end
-
-@inline function (c::Conv{N, true})(x::AbstractArray, ps, st::NamedTuple) where {N}
-    cdims = DenseConvDims(
-        x, ps.weight; stride=c.stride, padding=c.pad, dilation=c.dilation, groups=c.groups)
-    return apply_bias_activation(c.activation, _conv(x, ps.weight, cdims), ps.bias), st
+@inline function (c::Conv)(x::AbstractArray, ps, st::NamedTuple)
+    cdims = DenseConvDims(x, ps.weight; c.stride, padding=c.pad, c.dilation, c.groups)
+    return (
+        fused_conv_bias_activation(
+            c.activation, ps.weight, x, _getproperty(ps, Val(:bias)), cdims),
+        st)
 end
 
 function Base.show(io::IO, l::Conv)
@@ -617,16 +613,13 @@ function parameterlength(c::CrossCor{N, use_bias}) where {N, use_bias}
     return prod(c.kernel_size) * c.in_chs * c.out_chs + (use_bias ? c.out_chs : 0)
 end
 
-@inline function (c::CrossCor{N, false})(x::AbstractArray, ps, st::NamedTuple) where {N}
+@inline function (c::CrossCor)(x::AbstractArray, ps, st::NamedTuple)
     cdims = DenseConvDims(
         DenseConvDims(x, ps.weight; c.stride, padding=c.pad, c.dilation); F=true)
-    return apply_activation(c.activation, _conv(x, ps.weight, cdims)), st
-end
-
-@inline function (c::CrossCor{N, true})(x::AbstractArray, ps, st::NamedTuple) where {N}
-    cdims = DenseConvDims(
-        DenseConvDims(x, ps.weight; c.stride, padding=c.pad, c.dilation); F=true)
-    return apply_bias_activation(c.activation, _conv(x, ps.weight, cdims), ps.bias), st
+    return (
+        fused_conv_bias_activation(
+            c.activation, ps.weight, x, _getproperty(ps, Val(:bias)), cdims),
+        st)
 end
 
 function Base.show(io::IO, l::CrossCor)

src/layers/normalize.jl

@@ -129,15 +129,15 @@ statelength(l::BatchNorm) = (_track_stats(l) ? 2 * l.chs : 0) + 1
 
 function (BN::BatchNorm)(x::AbstractArray, ps, st::NamedTuple)
     y, stats = batchnorm(x, _getproperty(ps, Val(:scale)), _getproperty(ps, Val(:bias)),
-        _getproperty(st, Val(:running_mean)),
-        _getproperty(st, Val(:running_var)); BN.momentum, BN.epsilon, st.training)
+        _getproperty(st, Val(:running_mean)), _getproperty(st, Val(:running_var)),
+        BN.activation; BN.momentum, BN.epsilon, st.training)
 
     if _track_stats(BN)
         @set! st.running_mean = stats.running_mean
         @set! st.running_var = stats.running_var
     end
 
-    return apply_activation(BN.activation, y), st
+    return y, st
 end
 
 function Base.show(io::IO, l::BatchNorm)
@@ -241,9 +241,9 @@ end
 parameterlength(l::GroupNorm) = _affine(l) ? (l.chs * 2) : 0
 
 function (GN::GroupNorm)(x::AbstractArray, ps, st::NamedTuple)
-    y = groupnorm(x, _getproperty(ps, Val(:scale)),
-        _getproperty(ps, Val(:bias)); GN.groups, GN.epsilon)
-    return apply_activation(GN.activation, y), st
+    y = groupnorm(x, _getproperty(ps, Val(:scale)), _getproperty(ps, Val(:bias)),
+        GN.activation; GN.groups, GN.epsilon)
+    return y, st
 end
 
 function Base.show(io::IO, l::GroupNorm)
@@ -355,9 +355,9 @@ initialstates(rng::AbstractRNG, l::InstanceNorm) = (; training=Val(true))
 parameterlength(l::InstanceNorm) = _affine(l) ? (l.chs * 2) : 0
 
 function (IN::InstanceNorm)(x::AbstractArray, ps, st::NamedTuple)
-    y, stats = instancenorm(x, _getproperty(ps, Val(:scale)),
-        _getproperty(ps, Val(:bias)); IN.epsilon, st.training)
-    return apply_activation(IN.activation, y), st
+    y, stats = instancenorm(x, _getproperty(ps, Val(:scale)), _getproperty(ps, Val(:bias)),
+        IN.activation; IN.epsilon, st.training)
+    return y, st
 end
 
 function Base.show(io::IO, l::InstanceNorm)
@@ -574,9 +574,9 @@ function initialparameters(rng::AbstractRNG, ln::LayerNorm)
 end
 
 function (l::LayerNorm)(x::AbstractArray, ps, st::NamedTuple)
-    y = layernorm(
-        x, _getproperty(ps, Val(:scale)), _getproperty(ps, Val(:bias)); l.dims, l.epsilon)
-    return apply_activation(l.activation, y), st
+    y = layernorm(x, _getproperty(ps, Val(:scale)),
+        _getproperty(ps, Val(:bias)), l.activation; l.dims, l.epsilon)
+    return y, st
 end
 
 function Base.show(io::IO, l::LayerNorm)

src/utils.jl

@@ -29,7 +29,7 @@ _maybetuple_string(pad::Tuple) = all(==(pad[1]), pad) ? string(pad[1]) : string(
 # Padding
 struct SamePad end
 
-function _calc_padding(pad, k::NTuple{N}, dilation, stride) where {N}
+function _calc_padding(pad, ::NTuple{N}, dilation, stride) where {N}
     return _expand(Val(2 * N), pad)
 end
 
@@ -43,9 +43,6 @@ function _calc_padding(::SamePad, k::NTuple, dilation, stride)
     return Tuple(mapfoldl(i -> [cld(i, 2), fld(i, 2)], vcat, pad_amt))
 end
 
-# Getting typename
-get_typename(::T) where {T} = Base.typename(T).wrapper
-
 # RNN Utilities
 @inline _gate(h::Int, n::Int) = (1:h) .+ h * (n - 1)
 @inline _gate(x::AbstractVector, h::Int, n::Int) = view(x, _gate(h, n))
@@ -120,11 +117,6 @@ end
 
 # Backend Integration
 ## Convolution
-@inline _conv(x, weight, cdims) = conv(x, weight, cdims)
-@inline function _conv(x::SubArray{T, N, <:AbstractArray}, weight, cdims) where {T, N}
-    return _conv(copy(x), weight, cdims)
-end
-
 @inline _conv_transpose(x, weight, cdims) = ∇conv_data(x, weight, cdims)
 @inline function _conv_transpose(
         x::SubArray{T, N, <:GPUArraysCore.AnyGPUArray}, weight, cdims) where {T, N}