Remove cuDNN/MIOpen activation broadcast overloads (fixes #504, #509)

.gitignore

@@ -11,7 +11,7 @@ deps/usr
 deps.jl
 *.log
 .vscode/
-/Manifest.toml
+Manifest.toml
 test/Manifest.toml
 benchmark/Manifest.toml
 benchmark/*.json

benchmark/cuda/Project.toml

@@ -0,0 +1,9 @@
+[deps]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
+
+[compat]
+BenchmarkTools = "1.3"
+julia = "1.10"

benchmark/cuda/README.md

@@ -0,0 +1,68 @@
+# CUDA activation broadcast benchmarks
+
+Benchmarks accompanying
+[PR #686](https://github.com/FluxML/NNlib.jl/pull/686), which removes the custom
+cuDNN-routed broadcast overloads for the activations `relu`, `σ`, `elu` and
+`tanh` on `CuArray`s.
+
+Before the PR, `relu.(x::CuArray)` (and `materialize!`/in-place forms) were
+pirated onto cuDNN's `cudnnActivationForward!`. The PR drops those overloads and
+relies on CUDA.jl's native broadcast, because:
+
+- the type piracy of `Base.materialize`/`materialize!` caused method
+  invalidations and load-time latency (#504), and
+- cuDNN does not propagate `NaN`s by default, so `relu.(cu([NaN]))` returned `0`
+  instead of `NaN` (#509).
+
+The throughput rationale was that these elementwise ops are memory-bandwidth
+bound, so the native broadcast should be "just as fast". These benchmarks check
+that claim.
+
+## Running
+
+```
+julia --project=benchmark/cuda -e 'using Pkg; Pkg.develop(path="."); Pkg.instantiate()'
+julia --project=benchmark/cuda benchmark/cuda/activations.jl
+```
+
+The script measures both paths in the same process, so the comparison is
+apples-to-apples on the same machine/driver:
+
+| column   | path                                   | corresponds to |
+| -------- | -------------------------------------- | -------------- |
+| `native` | `f.(x)` / `broadcast!(f, dst, x)`      | **post-PR**    |
+| `cudnn`  | `cudnnActivationForward!(dst, x; mode)`| **pre-PR**     |
+
+`ratio = native / cudnn`; a ratio `> 1` means the post-PR native path is slower.
+Times are the minimum of 1000 GPU-synced samples, in microseconds.
+
+## Results
+
+See [`results.txt`](results.txt) for a full run. Machine used:
+RTX 5090, CUDA 13.2, cuDNN 9.2, Julia 1.12.
+
+Summary of what the numbers show:
+
+- **Float32 / Float64, large (memory-bound) tensors** — native broadcast is
+  competitive. For `elu`/`relu` in Float64 the native path is on par or *faster*
+  than cuDNN (cuDNN's ELU/RELU kernels do extra work); for `tanh`/`σ` it is
+  within a few percent. This matches the PR's "just as fast" claim.
+- **Small tensors** — cuDNN has lower CPU-side launch overhead, so the native
+  path is a few µs slower in absolute terms (launch-overhead dominated regime).
+- **Float16** — native broadcast is markedly slower (up to ~5–6×). CUDA.jl's
+  native broadcast does not vectorize Float16 (`half2`), so the Float16 native
+  time is essentially identical to the Float32 native time (no bandwidth
+  benefit), whereas cuDNN's Float16 kernel is faster than its Float32 one as
+  expected for a bandwidth-bound op.
+- **fast variants** (`tanh_fast`, `sigmoid_fast`; relu/elu have none) — these
+  were never cuDNN-routed, so the PR doesn't change them. On GPU they are no
+  faster than the plain native versions for Float16/Float32 (GPUs have
+  hardware-fast transcendentals, so the approximations save nothing), so they do
+  *not* close the Float16 gap. The exception is Float64 `tanh_fast`, ~10–14%
+  faster than native `tanh` — enough to beat cuDNN.
+
+**Takeaway:** the correctness (NaN propagation) and latency (invalidation)
+arguments for the PR stand on their own. On pure throughput, native broadcast
+is competitive for Float32/Float64 but leaves Float16 performance on the table.
+If Float16 elementwise throughput becomes important, the fix belongs in CUDA.jl's
+broadcast (Float16 vectorization), not in re-introducing the cuDNN piracy.

benchmark/cuda/activations.jl

@@ -0,0 +1,105 @@
+# Benchmarks for the removal of the custom cuDNN activation broadcast overloads
+# (PR https://github.com/FluxML/NNlib.jl/pull/686).
+#
+# Before that PR, broadcasting `relu`, `σ`, `elu`, `tanh` over a `CuArray` was
+# routed through cuDNN's `cudnnActivationForward!` by pirating
+# `Base.materialize`/`materialize!`. The PR removes those overloads and relies on
+# CUDA.jl's native broadcast instead. The claim is that for these
+# memory-bandwidth-bound elementwise ops the native broadcast is just as fast
+# (while also propagating NaNs correctly and avoiding method invalidations).
+#
+# This script measures both paths side by side in the same process so we get an
+# apples-to-apples "pre vs post" comparison without checking out two commits:
+#
+#   * "native"  -> `f.(x)` / `f.(x)` into `dst`     (the post-PR behaviour)
+#   * "cudnn"   -> `cudnnActivationForward!(dst, x)` (the pre-PR behaviour)
+#
+# Run with:
+#   julia --project=benchmark/cuda benchmark/cuda/activations.jl
+
+using CUDA
+using cuDNN
+using NNlib
+using BenchmarkTools
+using Printf
+
+using cuDNN: cudnnActivationForward!,
+             CUDNN_ACTIVATION_TANH, CUDNN_ACTIVATION_SIGMOID,
+             CUDNN_ACTIVATION_ELU, CUDNN_ACTIVATION_RELU
+
+CUDA.allowscalar(false)
+
+# (name, native activation fn, fast variant or `nothing`, cuDNN mode) for the four
+# activations that used to have a cuDNN-routed broadcast overload. `tanh_fast` and
+# `sigmoid_fast` are NNlib's faster approximations (relu/elu have no fast variant);
+# they were never routed through cuDNN, so they are an alternative native path.
+const ACTS = [
+    ("relu",    relu,    nothing,             CUDNN_ACTIVATION_RELU),
+    ("sigmoid", NNlib.σ, NNlib.sigmoid_fast,  CUDNN_ACTIVATION_SIGMOID),
+    ("elu",     elu,     nothing,             CUDNN_ACTIVATION_ELU),
+    ("tanh",    tanh,    NNlib.tanh_fast,     CUDNN_ACTIVATION_TANH),
+]
+
+# cuDNN supports Float16/Float32/Float64 for activations.
+const ELTYPES = (Float16, Float32, Float64)
+
+# A few representative shapes: a square matrix (as in the CPU benchmarks) and a
+# couple of conv-like 4D tensors of growing size.
+const SIZES = [
+    (1024, 1024),
+    (224, 224, 3, 32),
+    (56, 56, 64, 64),
+]
+
+# Median time in seconds for a GPU-synced benchmark of `f`.
+function gpu_time(f)
+    b = @benchmark CUDA.@sync($f()) samples=1000 evals=1 seconds=3
+    return minimum(b).time / 1e9  # seconds, use the minimum to reduce noise
+end
+
+# helper: "-" when no fast variant exists / time is missing
+fmt(t) = t === nothing ? "       -" : @sprintf("%8.2f", t * 1e6)
+ratio(a, b) = (a === nothing || b === nothing) ? "      -" : @sprintf("%6.2fx", a / b)
+
+function run_suite()
+    @printf("%-9s %-8s %-16s %8s %8s %8s   %7s %7s\n",
+            "act", "eltype", "size", "native", "fast", "cudnn", "nat/cu", "fst/cu")
+    println("-"^80)
+    results = Tuple[]
+    for (name, act, fast, mode) in ACTS, et in ELTYPES, sz in SIZES
+        x   = CUDA.randn(et, sz...)
+        dst = similar(x)
+
+        # warm up / compile all paths
+        act.(x)
+        broadcast!(act, dst, x)
+        cudnnActivationForward!(dst, x; mode)
+        fast === nothing || fast.(x)
+
+        # out-of-place
+        t_native_oop = gpu_time(() -> act.(x))
+        t_fast_oop   = fast === nothing ? nothing : gpu_time(() -> fast.(x))
+        t_cudnn_oop  = gpu_time(() -> cudnnActivationForward!(similar(x), x; mode))
+        # in-place
+        t_native_ip  = gpu_time(() -> broadcast!(act, dst, x))
+        t_fast_ip    = fast === nothing ? nothing : gpu_time(() -> broadcast!(fast, dst, x))
+        t_cudnn_ip   = gpu_time(() -> cudnnActivationForward!(dst, x; mode))
+
+        szstr = join(sz, "x")
+        @printf("%-9s %-8s %-16s %s %s %s   %s %s   (out-of-place)\n",
+                name, et, szstr, fmt(t_native_oop), fmt(t_fast_oop), fmt(t_cudnn_oop),
+                ratio(t_native_oop, t_cudnn_oop), ratio(t_fast_oop, t_cudnn_oop))
+        @printf("%-9s %-8s %-16s %s %s %s   %s %s   (in-place)\n",
+                "", "", "", fmt(t_native_ip), fmt(t_fast_ip), fmt(t_cudnn_ip),
+                ratio(t_native_ip, t_cudnn_ip), ratio(t_fast_ip, t_cudnn_ip))
+        push!(results, (name, et, szstr,
+                        t_native_oop, t_fast_oop, t_cudnn_oop,
+                        t_native_ip, t_fast_ip, t_cudnn_ip))
+    end
+    return results
+end
+
+if abspath(PROGRAM_FILE) == @__FILE__
+    @info "CUDA / cuDNN environment" CUDA.device() cuDNN.version()
+    run_suite()
+end

benchmark/cuda/results.txt

@@ -0,0 +1,74 @@
+act       eltype   size               native     fast    cudnn    nat/cu  fst/cu
+--------------------------------------------------------------------------------
+relu      Float16  1024x1024           12.15        -     8.42     1.44x       -   (out-of-place)
+                                       10.36        -     6.72     1.54x       -   (in-place)
+relu      Float16  224x224x3x32        40.05        -    11.65     3.44x       -   (out-of-place)
+                                       39.46        -    11.04     3.58x       -   (in-place)
+relu      Float16  56x56x64x64        108.42        -    19.30     5.62x       -   (out-of-place)
+                                      107.14        -    17.41     6.16x       -   (in-place)
+relu      Float32  1024x1024           11.71        -     8.74     1.34x       -   (out-of-place)
+                                       10.29        -     7.81     1.32x       -   (in-place)
+relu      Float32  224x224x3x32        41.10        -    14.55     2.82x       -   (out-of-place)
+                                       39.03        -    12.18     3.20x       -   (in-place)
+relu      Float32  56x56x64x64        109.34        -    58.33     1.87x       -   (out-of-place)
+                                      109.26        -    31.00     3.53x       -   (in-place)
+relu      Float64  1024x1024           12.89        -    10.95     1.18x       -   (out-of-place)
+                                       11.07        -     9.72     1.14x       -   (in-place)
+relu      Float64  224x224x3x32        41.30        -    38.80     1.06x       -   (out-of-place)
+                                       39.65        -    18.91     2.10x       -   (in-place)
+relu      Float64  56x56x64x64        134.89        -   135.47     1.00x       -   (out-of-place)
+                                      139.45        -   140.21     0.99x       -   (in-place)
+sigmoid   Float16  1024x1024           13.96    13.03     7.72     1.81x   1.69x   (out-of-place)
+                                       12.25    12.25     7.09     1.73x   1.73x   (in-place)
+sigmoid   Float16  224x224x3x32        46.06    45.56    12.96     3.55x   3.52x   (out-of-place)
+                                       44.81    44.71    10.98     4.08x   4.07x   (in-place)
+sigmoid   Float16  56x56x64x64        122.23   122.60    22.40     5.46x   5.47x   (out-of-place)
+                                      121.04   120.31    20.85     5.80x   5.77x   (in-place)
+sigmoid   Float32  1024x1024           12.79    12.77     8.69     1.47x   1.47x   (out-of-place)
+                                       11.52    11.72     7.48     1.54x   1.57x   (in-place)
+sigmoid   Float32  224x224x3x32        43.50    43.13    14.04     3.10x   3.07x   (out-of-place)
+                                       42.12    42.05    12.58     3.35x   3.34x   (in-place)
+sigmoid   Float32  56x56x64x64        117.44   117.77    49.67     2.36x   2.37x   (out-of-place)
+                                      118.62   118.41    34.28     3.46x   3.45x   (in-place)
+sigmoid   Float64  1024x1024           44.34    41.41    39.89     1.11x   1.04x   (out-of-place)
+                                       42.32    39.88    38.61     1.10x   1.03x   (in-place)
+sigmoid   Float64  224x224x3x32       161.35   149.83   143.90     1.12x   1.04x   (out-of-place)
+                                      159.73   148.34   142.01     1.12x   1.04x   (in-place)
+sigmoid   Float64  56x56x64x64        425.62   417.06   367.80     1.16x   1.13x   (out-of-place)
+                                      425.69   415.81   366.70     1.16x   1.13x   (in-place)
+elu       Float16  1024x1024           13.96        -     8.28     1.69x       -   (out-of-place)
+                                       11.96        -     6.80     1.76x       -   (in-place)
+elu       Float16  224x224x3x32        44.49        -    12.28     3.62x       -   (out-of-place)
+                                       43.19        -    10.83     3.99x       -   (in-place)
+elu       Float16  56x56x64x64        117.62        -    20.86     5.64x       -   (out-of-place)
+                                      115.89        -    19.76     5.87x       -   (in-place)
+elu       Float32  1024x1024           12.47        -     8.04     1.55x       -   (out-of-place)
+                                       10.41        -     7.30     1.43x       -   (in-place)
+elu       Float32  224x224x3x32        40.88        -    14.06     2.91x       -   (out-of-place)
+                                       39.44        -    12.52     3.15x       -   (in-place)
+elu       Float32  56x56x64x64        108.74        -    49.04     2.22x       -   (out-of-place)
+                                      109.44        -    33.94     3.22x       -   (in-place)
+elu       Float64  1024x1024           27.30        -    35.10     0.78x       -   (out-of-place)
+                                       25.71        -    33.54     0.77x       -   (in-place)
+elu       Float64  224x224x3x32        87.19        -   121.99     0.71x       -   (out-of-place)
+                                       85.86        -   120.53     0.71x       -   (in-place)
+elu       Float64  56x56x64x64        237.12        -   310.82     0.76x       -   (out-of-place)
+                                      237.16        -   310.27     0.76x       -   (in-place)
+tanh      Float16  1024x1024           12.00    12.76     8.29     1.45x   1.54x   (out-of-place)
+                                       10.66    10.72     6.84     1.56x   1.57x   (in-place)
+tanh      Float16  224x224x3x32        42.04    41.84    12.43     3.38x   3.37x   (out-of-place)
+                                       40.35    40.41    10.77     3.75x   3.75x   (in-place)
+tanh      Float16  56x56x64x64        111.33   111.18    21.15     5.26x   5.26x   (out-of-place)
+                                      109.48   110.30    19.67     5.57x   5.61x   (in-place)
+tanh      Float32  1024x1024           12.30    12.73     8.17     1.51x   1.56x   (out-of-place)
+                                       10.82    11.00     7.38     1.47x   1.49x   (in-place)
+tanh      Float32  224x224x3x32        41.73    43.17    14.21     2.94x   3.04x   (out-of-place)
+                                       40.51    41.36    12.31     3.29x   3.36x   (in-place)
+tanh      Float32  56x56x64x64        110.87   115.30    48.40     2.29x   2.38x   (out-of-place)
+                                      112.22   116.47    33.78     3.32x   3.45x   (in-place)
+tanh      Float64  1024x1024           61.47    53.15    61.47     1.00x   0.86x   (out-of-place)
+                                       59.81    52.82    60.49     0.99x   0.87x   (in-place)
+tanh      Float64  224x224x3x32       237.99   202.50   230.99     1.03x   0.88x   (out-of-place)
+                                      236.80   201.35   228.99     1.03x   0.88x   (in-place)
+tanh      Float64  56x56x64x64        623.87   559.06   596.92     1.05x   0.94x   (out-of-place)
+                                      623.32   558.84   596.30     1.05x   0.94x   (in-place)

ext/NNlibAMDGPUExt/activations.jl

@@ -1,16 +1,6 @@
-for (f, op) in [
-        NNlib.relu => MIOpen.relu,
-        NNlib.relu6 => x -> MIOpen.clippedrelu(x, 6),
-        NNlib.softplus => MIOpen.softrelu,
-        NNlib.σ => MIOpen.sigmoid,
-        Base.tanh => MIOpen.tanh,
-        # TODO define for leakyrelu, elu, etc.?
-    ], N in 1:5
-    @eval function Base.materialize(
-        bc::Broadcast.Broadcasted{<:Any,<:Any,typeof($f),<:Tuple{ROCArray{<:MIOPENFloat,$N}}}
-    )
-        return $op(bc.args[1])
-    end
-end
+# We deliberately do NOT route activation broadcasts (relu, σ, tanh, ...) through
+# MIOpen. Those overloads used to pirate `Base.materialize` (hurting latency via
+# invalidations, #504) and, like cuDNN, MIOpen does not propagate NaNs (#509).
+# AMDGPU's native broadcast is correct and, for these elementwise ops, just as fast.
 
 Base.broadcasted(::typeof(identity), x::ROCArray{T}) where {T<:MIOPENFloat} = x

ext/NNlibCUDACUDNNExt/activations.jl

@@ -1,40 +1,13 @@
 
 # Activation
 
-using Base.Broadcast
-using cuDNN: cudnnActivationForward!, cudnnOpTensor!,
-             CUDNN_ACTIVATION_TANH, CUDNN_ACTIVATION_SIGMOID, CUDNN_ACTIVATION_ELU,
-             CUDNN_ACTIVATION_RELU, CUDNN_ACTIVATION_CLIPPED_RELU, CUDNN_OP_TENSOR_MAX,
-             CUDNN_ACTIVATION_IDENTITY
-
-for (f, op) in [
-    CUDA.tanh       => (src,dst)->cudnnActivationForward!(dst, src, mode=CUDNN_ACTIVATION_TANH),
-    NNlib.σ         => (src,dst)->cudnnActivationForward!(dst, src, mode=CUDNN_ACTIVATION_SIGMOID),
-    NNlib.elu       => (src,dst)->cudnnActivationForward!(dst, src, mode=CUDNN_ACTIVATION_ELU),
-    NNlib.relu      => (src,dst)->cudnnActivationForward!(dst, src, mode=CUDNN_ACTIVATION_RELU),
-    # NNlib.relu6     => (src,dst)->cudnnActivationForward!(dst, src, mode=CUDNN_ACTIVATION_CLIPPED_RELU, coef=6.0),
-    # NNlib.leakyrelu => (src,dst)->cudnnOpTensor!(dst, src, src; op=CUDNN_OP_TENSOR_MAX, alpha1=0.01),
-    ]
-
-    @eval begin
-        # in-place
-        function Base.materialize!(dst::DenseCuArray{<:CUDNNFloat},
-                                   bc::Broadcast.Broadcasted{<:Any,<:Any,typeof($f),<:Tuple{DenseCuArray{<:CUDNNFloat}}})
-            $op(bc.args[1], dst)
-            return dst
-        end
-
-        # out of place
-        function Base.materialize(bc::Broadcast.Broadcasted{<:Any,<:Any,typeof($f),<:Tuple{DenseCuArray{<:CUDNNFloat}}})
-            ElType = Broadcast.combine_eltypes(bc.f, bc.args)
-            dst = similar(bc, ElType)
-            $op(bc.args[1], dst)
-            return dst
-        end
-    end
-end
+# We deliberately do NOT route activation broadcasts (relu, σ, elu, tanh, ...)
+# through cuDNN's `cudnnActivationForward!`. Those overloads used to pirate
+# `Base.materialize`/`materialize!` (hurting latency via invalidations, #504) and,
+# worse, cuDNN does not propagate NaNs by default, so e.g. `relu.(cu([NaN]))`
+# returned `0` instead of `NaN` (#509). CUDA.jl's native broadcast is correct and,
+# for these memory-bandwidth-bound elementwise ops, just as fast.
 
 # CUDNN_ACTIVATION_IDENTITY does not work with cudnnActivationForward
 # FIXME: put this optimization in GPUArrays' `copyto!` (like Base.Broadcast's `copyto!`)
 Base.broadcasted(::typeof(identity), x::DenseCuArray{T}) where {T<:CUDNNFloat} = x
-

test/ext_amdgpu/activations.jl

@@ -9,3 +9,11 @@
         end
     end
 end
+
+@testset "NaN propagation" begin
+    # MIOpen's activation path used to swallow NaNs (returning e.g. 0 for relu),
+    # diverging from the CPU. Make sure the native broadcast propagates them (#509).
+    for f in (NNlib.relu, NNlib.relu6, NNlib.softplus, tanh, NNlib.σ)
+        @test all(isnan, f.(ROCArray([NaN32, NaN32])))
+    end
+end

test/ext_cuda/activations.jl

@@ -23,7 +23,15 @@ end
     @test f(cs) ≈ collect(f(CuArray(cs)))
 end
 
-@testset "softplus" begin 
+@testset "NaN propagation" begin
+    # cuDNN's activation path used to swallow NaNs (returning e.g. 0 for relu),
+    # diverging from the CPU. Make sure the native broadcast propagates them (#509).
+    for f in (relu, sigmoid, tanh, elu)
+        @test all(isnan, f.(CuArray([NaN32, NaN32])))
+    end
+end
+
+@testset "softplus" begin
   # softplus does not give `Inf` for large arguments
    x = CuArray([1000.])
    @test all(softplus.(x) .== x)