feat(hipdnn): add GPU batchnorm backward reference

dnn-providers/integration-tests/gpu-ref/include/hipdnn-gpu-ref/GpuFpReferenceBatchnorm.hpp

@@ -40,6 +40,26 @@ inline std::vector<std::string> buildBatchnormFwdDefines()
     return defines;
 }
 
+template <typename DyDataType,
+          typename XDataType,
+          typename ScaleBiasDataType,
+          typename MeanVarianceDataType,
+          typename DxDataType,
+          typename ComputeDataType>
+inline std::vector<std::string> buildBatchnormBwdDefines()
+{
+    std::vector<std::string> defines;
+    defines.emplace_back(std::string("-DGRAD_OUTPUT_TYPE=") + HipRtcTypeName<DyDataType>::VALUE);
+    defines.emplace_back(std::string("-DINPUT_TYPE=") + HipRtcTypeName<XDataType>::VALUE);
+    defines.emplace_back(std::string("-DSCALE_BIAS_TYPE=")
+                         + HipRtcTypeName<ScaleBiasDataType>::VALUE);
+    defines.emplace_back(std::string("-DMEAN_VAR_TYPE=")
+                         + HipRtcTypeName<MeanVarianceDataType>::VALUE);
+    defines.emplace_back(std::string("-DGRAD_INPUT_TYPE=") + HipRtcTypeName<DxDataType>::VALUE);
+    defines.emplace_back(std::string("-DCOMPUTE_TYPE=") + HipRtcTypeName<ComputeDataType>::VALUE);
+    return defines;
+}
+
 } // namespace detail
 
 class GpuFpReferenceBatchnorm
@@ -207,6 +227,55 @@ class GpuFpReferenceBatchnorm
         }
     }
 
+    template <class DyDataType,
+              class XDataType,
+              class ScaleBiasDataType,
+              class MeanVarianceDataType = ScaleBiasDataType,
+              class DxDataType = XDataType,
+              class ComputeDataType = MeanVarianceDataType>
+    static void backward(hipdnn_data_sdk::utilities::TensorBase<DyDataType>& dy,
+                         hipdnn_data_sdk::utilities::TensorBase<XDataType>& x,
+                         hipdnn_data_sdk::utilities::TensorBase<ScaleBiasDataType>& scale,
+                         hipdnn_data_sdk::utilities::TensorBase<DxDataType>& dx,
+                         hipdnn_data_sdk::utilities::TensorBase<ScaleBiasDataType>& dscale,
+                         hipdnn_data_sdk::utilities::TensorBase<ScaleBiasDataType>& dbias,
+                         hipdnn_data_sdk::utilities::TensorBase<MeanVarianceDataType>* mean
+                         = nullptr,
+                         hipdnn_data_sdk::utilities::TensorBase<MeanVarianceDataType>* invVariance
+                         = nullptr,
+                         double epsilon = hipdnn_data_sdk::utilities::BATCHNORM_DEFAULT_EPSILON)
+    {
+        validateBwdInput<DyDataType,
+                         XDataType,
+                         ScaleBiasDataType,
+                         MeanVarianceDataType,
+                         DxDataType,
+                         ComputeDataType>(dy, x, scale, dx, dscale, dbias, mean, invVariance);
+
+        auto defines = detail::buildBatchnormBwdDefines<DyDataType,
+                                                        XDataType,
+                                                        ScaleBiasDataType,
+                                                        MeanVarianceDataType,
+                                                        DxDataType,
+                                                        ComputeDataType>();
+        launchBackward(dy.memory().deviceData(),
+                       x.memory().deviceData(),
+                       x.dims(),
+                       x.strides(),
+                       scale.memory().deviceData(),
+                       dx.memory().deviceData(),
+                       dscale.memory().deviceData(),
+                       dbias.memory().deviceData(),
+                       mean ? mean->memory().deviceData() : nullptr,
+                       invVariance ? invVariance->memory().deviceData() : nullptr,
+                       epsilon,
+                       defines);
+
+        dx.memory().markDeviceModified();
+        dscale.memory().markDeviceModified();
+        dbias.memory().markDeviceModified();
+    }
+
 private:
     // --- Validators ---
 
@@ -486,6 +555,72 @@ class GpuFpReferenceBatchnorm
                       "compute data types.");
     }
 
+    template <typename DyDataType,
+              typename XDataType,
+              typename ScaleBiasDataType,
+              typename MeanVarianceDataType,
+              typename DxDataType,
+              typename ComputeDataType>
+    static void validateBwdInput(
+        const hipdnn_data_sdk::utilities::TensorBase<DyDataType>& dy,
+        const hipdnn_data_sdk::utilities::TensorBase<XDataType>& x,
+        const hipdnn_data_sdk::utilities::TensorBase<ScaleBiasDataType>& scale,
+        const hipdnn_data_sdk::utilities::TensorBase<DxDataType>& dx,
+        const hipdnn_data_sdk::utilities::TensorBase<ScaleBiasDataType>& dscale,
+        const hipdnn_data_sdk::utilities::TensorBase<ScaleBiasDataType>& dbias,
+        const hipdnn_data_sdk::utilities::TensorBase<MeanVarianceDataType>* mean,
+        const hipdnn_data_sdk::utilities::TensorBase<MeanVarianceDataType>* invVariance)
+    {
+        if((mean == nullptr) != (invVariance == nullptr))
+        {
+            throw std::invalid_argument(
+                "Batchnorm backward requires mean and invVariance to be provided together.");
+        }
+
+        const TensorProps inputTensorProps("input", x.dims(), x.strides());
+        TensorProps outputTensorProps("dx", dx.dims(), dx.strides());
+        std::vector<TensorProps> affineTensorProps;
+        affineTensorProps.emplace_back("scale", scale.dims(), scale.strides());
+        affineTensorProps.emplace_back("dscale", dscale.dims(), dscale.strides());
+        affineTensorProps.emplace_back("dbias", dbias.dims(), dbias.strides());
+        if(mean != nullptr)
+        {
+            affineTensorProps.emplace_back("mean", mean->dims(), mean->strides());
+            affineTensorProps.emplace_back(
+                "invVariance", invVariance->dims(), invVariance->strides());
+        }
+
+        validateConsistentDimensions(inputTensorProps, outputTensorProps, affineTensorProps);
+        if(dy.dims() != x.dims())
+        {
+            throw std::invalid_argument(
+                "Batchnorm backward requires dy and input tensors to have the same shape.");
+        }
+
+        affineTensorProps.emplace_back("dy", dy.dims(), dy.strides());
+        affineTensorProps.push_back(std::move(outputTensorProps));
+        validateConsistentLayouts(inputTensorProps, affineTensorProps);
+
+        static_assert(IS_SUPPORTED_DATA_TYPE<DyDataType>,
+                      "Batchnorm backward supports only double, float, half, and bfloat16 dy data "
+                      "types.");
+        static_assert(IS_SUPPORTED_DATA_TYPE<XDataType>,
+                      "Batchnorm backward supports only double, float, half, and bfloat16 input "
+                      "data types.");
+        static_assert(IS_SUPPORTED_DATA_TYPE<ScaleBiasDataType>,
+                      "Batchnorm backward supports only double, float, half, and bfloat16 "
+                      "scale/bias data types.");
+        static_assert(IS_SUPPORTED_DATA_TYPE<MeanVarianceDataType>,
+                      "Batchnorm backward supports only double, float, half, and bfloat16 "
+                      "mean/invVariance data types.");
+        static_assert(IS_SUPPORTED_DATA_TYPE<DxDataType>,
+                      "Batchnorm backward supports only double, float, half, and bfloat16 dx data "
+                      "types.");
+        static_assert(std::is_same_v<ComputeDataType, float>
+                          || std::is_same_v<ComputeDataType, double>,
+                      "Batchnorm backward requires float or double compute data types.");
+    }
+
     // --- Kernel launchers (defined in GpuFpReferenceBatchnorm.cpp) ---
     static void launchFwdInf(const void* inputPtr,
                              const std::vector<int64_t>& inputDims,
@@ -523,6 +658,19 @@ class GpuFpReferenceBatchnorm
                                void* nextRunningMeanPtr,
                                void* nextRunningVariancePtr,
                                std::vector<std::string>& defines);
+
+    static void launchBackward(const void* dyPtr,
+                               const void* inputPtr,
+                               const std::vector<int64_t>& inputDims,
+                               const std::vector<int64_t>& inputStrides,
+                               const void* scalePtr,
+                               void* dxPtr,
+                               void* dscalePtr,
+                               void* dbiasPtr,
+                               const void* meanPtr,
+                               const void* invVariancePtr,
+                               double epsilon,
+                               std::vector<std::string>& defines);
 };
 
 } // namespace hipdnn_gpu_ref

dnn-providers/integration-tests/gpu-ref/kernels/CMakeLists.txt

@@ -15,6 +15,7 @@ set(GPU_REF_KERNEL_FILES
     ${CMAKE_CURRENT_SOURCE_DIR}/sdpa/GpuRefSdpaRaggedFwd.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/rmsnorm/GpuRefRMSNormFwd.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/rmsnorm/GpuRefRMSNormBwd.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/batchnorm/GpuRefBatchnormBwd.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/batchnorm/GpuRefBatchnormFwdInf.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/batchnorm/GpuRefBatchnormFwdTrain.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/pointwise/GpuRefPointwiseUnary.cpp

dnn-providers/integration-tests/gpu-ref/kernels/batchnorm/GpuRefBatchnormBwd.cpp

@@ -0,0 +1,153 @@
+// Copyright © Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+// GPU reference Batchnorm backward kernel.
+// Compiled via HipRTC with all tensor and compute types supplied as preprocessor defines.
+// Each block handles one channel and reduces over its N*spatial elements. It obtains the channel
+// statistics, reduces dscale and dbias, then computes dx from those reduced gradients.
+
+#include "GpuRefTypes.h"
+
+using namespace gpu_ref;
+
+extern "C" __global__ void BatchnormBwdRef(BatchnormBwdArgs args)
+{
+    auto* dy = static_cast<const GRAD_OUTPUT_TYPE*>(args.dy);
+    auto* input = static_cast<const INPUT_TYPE*>(args.input);
+    auto* scale = static_cast<const SCALE_BIAS_TYPE*>(args.scale);
+    auto* dx = static_cast<GRAD_INPUT_TYPE*>(args.dx);
+    auto* dscale = static_cast<SCALE_BIAS_TYPE*>(args.dscale);
+    auto* dbias = static_cast<SCALE_BIAS_TYPE*>(args.dbias);
+    auto* savedMean = static_cast<const MEAN_VAR_TYPE*>(args.mean);
+    auto* savedInvVariance = static_cast<const MEAN_VAR_TYPE*>(args.invVariance);
+
+    constexpr long long localSize = static_cast<long long>(LOCAL_SIZE);
+    constexpr bool isChannelLastLayout = static_cast<bool>(IS_CHANNEL_LAST_LAYOUT);
+    const long long channel = static_cast<long long>(blockIdx.x);
+    const long long lid = static_cast<long long>(threadIdx.x);
+    const long long nhw = args.n * args.hw;
+    const long long chw = args.c * args.hw;
+    const COMPUTE_TYPE invNhw = static_cast<COMPUTE_TYPE>(1.0) / static_cast<COMPUTE_TYPE>(nhw);
+
+    __shared__ COMPUTE_TYPE reduceA[localSize];
+    __shared__ COMPUTE_TYPE reduceB[localSize];
+
+    COMPUTE_TYPE channelMean;
+    COMPUTE_TYPE channelInvVariance;
+    if(savedMean == nullptr || savedInvVariance == nullptr)
+    {
+        // Compute the channel mean
+        COMPUTE_TYPE sum = static_cast<COMPUTE_TYPE>(0);
+        for(long long i = lid; i < nhw; i += localSize)
+        {
+            const long long nidx = i / args.hw;
+            const long long hwidx = i - nidx * args.hw;
+            const long long index = isChannelLastLayout ? nidx * chw + hwidx * args.c + channel
+                                                        : nidx * chw + channel * args.hw + hwidx;
+            const COMPUTE_TYPE xValue = toAccum(input[index]);
+            sum = sum + xValue;
+        }
+
+        reduceA[lid] = sum;
+        __syncthreads();
+        for(long long offset = localSize >> 1; offset > 0; offset >>= 1)
+        {
+            if(lid < offset)
+            {
+                reduceA[lid] = reduceA[lid] + reduceA[lid + offset];
+            }
+            __syncthreads();
+        }
+
+        channelMean = reduceA[0] * invNhw;
+
+        // Compute variance from deviations from the mean to avoid cancellation
+        COMPUTE_TYPE varianceSum = static_cast<COMPUTE_TYPE>(0);
+        for(long long i = lid; i < nhw; i += localSize)
+        {
+            const long long nidx = i / args.hw;
+            const long long hwidx = i - nidx * args.hw;
+            const long long index = isChannelLastLayout ? nidx * chw + hwidx * args.c + channel
+                                                        : nidx * chw + channel * args.hw + hwidx;
+            const COMPUTE_TYPE deviation = toAccum(input[index]) - channelMean;
+            varianceSum = varianceSum + deviation * deviation;
+        }
+
+        reduceB[lid] = varianceSum;
+        __syncthreads();
+        for(long long offset = localSize >> 1; offset > 0; offset >>= 1)
+        {
+            if(lid < offset)
+            {
+                reduceB[lid] = reduceB[lid] + reduceB[lid + offset];
+            }
+            __syncthreads();
+        }
+
+        const COMPUTE_TYPE variance = reduceB[0] * invNhw;
+        channelInvVariance = rsqrt(variance + toAccum(args.epsilon));
+    }
+    else
+    {
+        channelMean = toAccum(savedMean[channel]);
+        channelInvVariance = toAccum(savedInvVariance[channel]);
+    }
+    __syncthreads();
+
+    // Accumulate the per-channel scale and bias gradients
+    COMPUTE_TYPE dotProduct = static_cast<COMPUTE_TYPE>(0);
+    COMPUTE_TYPE sumDy = static_cast<COMPUTE_TYPE>(0);
+    for(long long i = lid; i < nhw; i += localSize)
+    {
+        const long long nidx = i / args.hw;
+        const long long hwidx = i - nidx * args.hw;
+        const long long index = isChannelLastLayout ? nidx * chw + hwidx * args.c + channel
+                                                    : nidx * chw + channel * args.hw + hwidx;
+        const COMPUTE_TYPE xValue = toAccum(input[index]);
+        const COMPUTE_TYPE dyValue = toAccum(dy[index]);
+        const COMPUTE_TYPE xHat = (xValue - channelMean) * channelInvVariance;
+        dotProduct = dotProduct + xHat * dyValue;
+        sumDy = sumDy + dyValue;
+    }
+
+    reduceA[lid] = dotProduct;
+    reduceB[lid] = sumDy;
+    __syncthreads();
+    for(long long offset = localSize >> 1; offset > 0; offset >>= 1)
+    {
+        if(lid < offset)
+        {
+            reduceA[lid] = reduceA[lid] + reduceA[lid + offset];
+            reduceB[lid] = reduceB[lid] + reduceB[lid + offset];
+        }
+        __syncthreads();
+    }
+
+    const COMPUTE_TYPE channelScale = toAccum(scale[channel]);
+    const COMPUTE_TYPE meanDyXhat = reduceA[0] * invNhw;
+    const COMPUTE_TYPE meanDy = reduceB[0] * invNhw;
+    const COMPUTE_TYPE scalarCoefficient = channelScale * channelInvVariance;
+
+    // Store the reduced affine gradients once per channel
+    if(lid == 0)
+    {
+        SCALE_BIAS_TYPE* tag = nullptr;
+        dscale[channel] = fromAccum(reduceA[0], tag);
+        dbias[channel] = fromAccum(reduceB[0], tag);
+    }
+
+    // Compute the input gradient for each element from the channel reductions
+    GRAD_INPUT_TYPE* tag = nullptr;
+    for(long long i = lid; i < nhw; i += localSize)
+    {
+        const long long nidx = i / args.hw;
+        const long long hwidx = i - nidx * args.hw;
+        const long long index = isChannelLastLayout ? nidx * chw + hwidx * args.c + channel
+                                                    : nidx * chw + channel * args.hw + hwidx;
+        const COMPUTE_TYPE xValue = toAccum(input[index]);
+        const COMPUTE_TYPE dyValue = toAccum(dy[index]);
+        const COMPUTE_TYPE xHat = (xValue - channelMean) * channelInvVariance;
+        const COMPUTE_TYPE dxValue = (dyValue - meanDy - xHat * meanDyXhat) * scalarCoefficient;
+        dx[index] = fromAccum(dxValue, tag);
+    }
+}

dnn-providers/integration-tests/gpu-ref/kernels/types/GpuRefBatchnormArgs.h

@@ -52,3 +52,19 @@ struct BatchnormFwdTrainArgs
     long long c;
     long long hw;
 };
+
+struct BatchnormBwdArgs
+{
+    const void* dy;
+    const void* input;
+    const void* scale;
+    void* dx;
+    void* dscale;
+    void* dbias;
+    const void* mean;
+    const void* invVariance;
+    double epsilon;
+    long long n;
+    long long c;
+    long long hw;
+};