Add MIOpen (AMDGPU) fast path for batchnorm
- 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