FluxML / FluxML/NNlib.jl

Add MIOpen (AMDGPU) fast path for batchnorm

Open
#752 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
258
Forks
138
Avg merge
2d 1h
Merged PRs (30d)
3

Description

Following the addition of the functional normalization operators (`batchnorm`, `instancenorm`, `groupnorm`, `layernorm`) in NNlib v0.9.41, the cuDNN-accelerated `batchnorm` lives in `NNlibCUDACUDNNExt` and is selected automatically for `CuArray`s. There is no equivalent MIOpen fast path for AMDGPU `ROCArray`s — they fall back to the generic implementation.

Flux is migrating its normalization layers to wrap `NNlib.batchnorm` (FluxML/Flux.jl#2700) and, as part of that, is **removing** its own AMDGPU MIOpen `BatchNorm` specialization (previously in `ext/FluxAMDGPUExt/batchnorm.jl`). To keep the accelerated AMD path available, it would be good to move that logic into NNlib as an AMDGPU/MIOpen extension, mirroring `NNlibCUDACUDNNExt`.

The code removed from Flux, as a starting point:

```julia
function _amdgpu_batchnorm(x, γ, β; μ, σ², ϵ, within_grad::Bool)
if within_grad
return AMDGPU.MIOpen.batchnorm_training(x, γ, β, μ, σ²; ϵ=Float64(ϵ), iteration=0) # TODO iteration
else
return AMDGPU.MIOpen.batchnorm_inference(x, γ, β, μ, σ²; ϵ=Float64(ϵ))
end
end

function ChainRulesCore.rrule(::typeof(_amdgpu_batchnorm), x, γ, β; μ, σ², ϵ, within_grad::Bool)
y, μ_saved, ν_saved = _amdgpu_batchnorm(x, γ, β; μ, σ², ϵ, within_grad)
function _batchnorm_pullback(Δ)
dx, dγ, dβ = AMDGPU.MIOpen.∇batchnorm(unthunk(Δ), x, γ, β, μ_saved, ν_saved)
(NoTangent(), dx, dγ, dβ)
end
y, _batchnorm_pullback
end
```

Ideally the NNlib method would implement the full `batchnorm(g, b, x, running_mean, running_var, momentum; eps, training, track_stats)` signature (matching the generic and cuDNN methods) and its `∇batchnorm`, honoring `training` / `track_stats` / running-statistic updates rather than relying solely on `within_gradient` as the old Flux code did. It should also cover the 2D/4D/5D input shapes and address the `iteration` TODO.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reading the cuDNN implementation in NNlibCUDACUDNNExt and the removed Flux ext/FluxAMDGPUExt/batchnorm.jl logic. Trace the generic batchnorm and ∇batchnorm signatures, then compare the AMDGPU/MIOpen APIs for 2D, 4D, and 5D inputs. Done means an NNlib AMDGPU extension supports the full signature, training and running-statistic behavior, gradients, and a resolved iteration value.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.