/* ******************************************************************************
 *
 *
 * This program and the accompanying materials are made available under the
 * terms of the Apache License, Version 2.0 which is available at
 * https://www.apache.org/licenses/LICENSE-2.0.
 *
 *  See the NOTICE file distributed with this work for additional
 *  information regarding copyright ownership.
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations
 * under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 ******************************************************************************/

//
// @author Yurii Shyrma (iuriish@yahoo.com)
//

#include <ops/declarable/helpers/convolutions.h>

#include "cudnnUtils.h"

namespace sd {
namespace ops {
namespace platforms {

//////////////////////////////////////////////////////////////////////////
static void batchnormCUDNN(const LaunchContext* context, NDArray* input, NDArray* mean,
                           NDArray* variance, NDArray* gamma, NDArray* beta, NDArray* output,
                           const double epsilon, const bool isSpatialMode) {
  // input, output -> 4D:nchw, 5D:ncdhw
  // mean, variance, gamma, beta -> 1xCx1x1 for 4D and 1xCx1x1x1 for 5D for BATCHNORM_MODE_SPATIAL mode
  //                             -> 1xCxHxW for 4D and 1xCxDxHxW for 5D for BATCHNORM_MODE_PER_ACTIVATION mode

  const cudnnDataType_t dataType = cudnnDataType(input->dataType());

  const LongType xRank = input->rankOf();

  auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
  CHECK_CUDNN_FAILURE(cudnnSetStream(*handle, *context->getCudaStream()));

  const std::vector<int> xShape = input->getShapeAsVectorInt();  // input and output have same shapes

  std::vector<int> paramsShape, paramsStrides;  // mean, variance, gamma and beta have same shapes
  if (isSpatialMode) {                          // 1xCx1x1
    const int iC = static_cast<int>(mean->lengthOf());
    const int stride0 = static_cast<int>(mean->strideAt(0));
    paramsShape = xRank == 4 ? std::vector<int>({1, iC, 1, 1}) : std::vector<int>({1, iC, 1, 1, 1});
    paramsStrides = xRank == 4 ? std::vector<int>({iC * stride0, stride0, 1, 1})
                               : std::vector<int>({iC * stride0, stride0, 1, 1, 1});
  } else {
    paramsShape = std::vector<int>(mean->getShapeAsVector().begin(), mean->getShapeAsVector().end());
    paramsStrides = xRank == 4
                    ? std::vector<int>({static_cast<int>(mean->strideAt(0)), static_cast<int>(mean->strideAt(1)), static_cast<int>(mean->strideAt(2)),
                                        static_cast<int>(mean->strideAt(3))})
                    : std::vector<int>({static_cast<int>(mean->strideAt(0)), static_cast<int>(mean->strideAt(1)), static_cast<int>(mean->strideAt(2)),
                                        static_cast<int>(mean->strideAt(3)), static_cast<int>(mean->strideAt(4))});
  }

  std::vector<int> xStrides = {static_cast<int>(input->strideAt(0)), static_cast<int>(input->strideAt(1)), static_cast<int>(input->strideAt(2)),
                               static_cast<int>(input->strideAt(3))};
  std::vector<int> zStrides = {static_cast<int>(output->strideAt(0)), static_cast<int>(output->strideAt(1)), static_cast<int>(output->strideAt(2)),
                               static_cast<int>(output->strideAt(3))};
  if (xRank > 4) {  // 5D
    xStrides.push_back((LongType)input->strideAt(4));
    zStrides.push_back((LongType)output->strideAt(4));
  }

  cudnnTensorFormat_t format = CUDNN_TENSOR_NCHW;

  // input descriptor
    x.set(dataType, xRank, xShape.data(), xStrides.data());

  // output descriptor
  CudnnTensor z;
    z.set(dataType, xRank, xShape.data(), zStrides.data());

  // mean, variance, gamma and beta descriptor, the same descriptor for all of them
  CudnnTensor params;
    params.set(dataType, xRank, paramsShape.data(), paramsStrides.data());

  // provide scaling parameters
  const float alpha32(1), beta32(0);
  const double alpha64(1), beta64(0);
  const void* ptrAlpha =
      output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
  const void* ptrBeta =
      output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);

  NDArray::prepareSpecialUse({output}, {input, mean, variance, gamma, beta});

  // calculations
  CHECK_CUDNN_FAILURE_MSG(
      STRINGIZE(cudnnBatchNormalizationForwardInference),
      cudnnBatchNormalizationForwardInference(
          *handle, isSpatialMode ? CUDNN_BATCHNORM_SPATIAL : CUDNN_BATCHNORM_PER_ACTIVATION, ptrAlpha, ptrBeta, x,
          input->specialBuffer(), z, output->specialBuffer(), params, gamma->specialBuffer(), beta->specialBuffer(),
          mean->specialBuffer(), variance->specialBuffer(), epsilon));

  auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
  if (cudaErr != 0) throw cuda_exception::build("batchnormCUDNN: cudaStreamSynchronize failed !", cudaErr);

  NDArray::registerSpecialUse({output}, {input, mean, variance, gamma, beta});
}

//////////////////////////////////////////////////////////////////////////
static void batchnormBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* mean,
                             NDArray* variance, NDArray* gamma, NDArray* gradO, NDArray* gradI,
                             NDArray* gradG, NDArray* gradB, const double epsilon, const bool isSpatialMode) {
  // input, gradO, gradI -> 4D:nchw, 5D:ncdhw
  // mean, variance, gamma, beta, gradM, gradV, gradG, gradB -> 1xCx1x1 for 4D and 1xCx1x1x1 for 5D for
  // BATCHNORM_MODE_SPATIAL mode
  //                                                         -> 1xCxHxW for 4D and 1xCxDxHxW for 5D for
  //                                                         BATCHNORM_MODE_PER_ACTIVATION mode

  const cudnnDataType_t dataType = cudnnDataType(input->dataType());

  const int xRank = input->rankOf();

  auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
  cudnnStatus_t err = cudnnSetStream(*handle, *context->getCudaStream());

  const std::vector<int> xShape = input->getShapeAsVectorInt();  // input and output have same shapes

  std::vector<int> paramsShape, paramsStrides;  // mean, variance, gamma and beta have same shapes
  if (isSpatialMode) {                          // 1xCx1x1
    const int iC = static_cast<int>(mean->lengthOf());
    const int stride0 = static_cast<int>(mean->strideAt(0));
    paramsShape = xRank == 4 ? std::vector<int>({1, iC, 1, 1}) : std::vector<int>({1, iC, 1, 1, 1});
    paramsStrides = xRank == 4 ? std::vector<int>({iC * stride0, stride0, 1, 1})
                               : std::vector<int>({iC * stride0, stride0, 1, 1, 1});
  } else {
    paramsShape = std::vector<int>(mean->getShapeAsVector().begin(), mean->getShapeAsVector().end());
    paramsStrides = xRank == 4
                    ? std::vector<int>({static_cast<int>(mean->strideAt(0)), static_cast<int>(mean->strideAt(1)), static_cast<int>(mean->strideAt(2)),
                                        static_cast<int>(mean->strideAt(3))})
                    : std::vector<int>({static_cast<int>(mean->strideAt(0)), static_cast<int>(mean->strideAt(1)), static_cast<int>(mean->strideAt(2)),
                                        static_cast<int>(mean->strideAt(3)), static_cast<int>(mean->strideAt(4))});
  }

  std::vector<int> xStrides = {static_cast<int>(input->strideAt(0)), static_cast<int>(input->strideAt(1)), static_cast<int>(input->strideAt(2)),
                               static_cast<int>(input->strideAt(3))};
  std::vector<int> dxStrides = {static_cast<int>(gradI->strideAt(0)), static_cast<int>(gradI->strideAt(1)), static_cast<int>(gradI->strideAt(2)),
                                static_cast<int>(gradI->strideAt(3))};
  std::vector<int> dzStrides = {static_cast<int>(gradO->strideAt(0)), static_cast<int>(gradO->strideAt(1)), static_cast<int>(gradO->strideAt(2)),
                                static_cast<int>(gradO->strideAt(3))};

  if (xRank > 4) {  // 5D
    xStrides.push_back(static_cast<int>(input->strideAt(4)));
    dxStrides.push_back(static_cast<int>(gradI->strideAt(4)));
    dzStrides.push_back(static_cast<int>(gradO->strideAt(4)));
  }
  cudnnTensorFormat_t format = CUDNN_TENSOR_NCHW;

  // input descriptor
  CudnnTensor x;

    x.set(dataType, xRank, xShape.data(), xStrides.data());

  // gradO descriptor
  CudnnTensor dz;
    dz.set(dataType, xRank, xShape.data(), dzStrides.data());

  // gradI descriptor
  CudnnTensor dx;
    dx.set(dataType, xRank, xShape.data(), dxStrides.data());

  // mean, variance, gamma, gradG and gradB descriptor, the same descriptor for all of them
  CudnnTensor params;
    params.set(dataType, xRank, paramsShape.data(), paramsStrides.data());

  // provide scaling parameters
  const float alpha32(1), beta32(0);
  double alpha64(1), beta64(0);
  const void* ptrAlpha =
      input->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
  const void* ptrBeta =
      input->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);

  NDArray::prepareSpecialUse({gradI, gradG, gradB}, {input, mean, variance, gamma, gradO});

  // calculations
  // TODO: we can use cache here
  CHECK_CUDNN_FAILURE_MSG(
      STRINGIZE(cudnnBatchNormalizationBackward),
      cudnnBatchNormalizationBackward(*handle, isSpatialMode ? CUDNN_BATCHNORM_SPATIAL : CUDNN_BATCHNORM_PER_ACTIVATION,
                                      ptrAlpha, ptrBeta, ptrAlpha, ptrBeta, x, input->specialBuffer(), dz,
                                      gradO->specialBuffer(), dx, gradI->specialBuffer(), params,
                                      gamma->specialBuffer(), gradG->specialBuffer(), gradB->specialBuffer(), epsilon,
                                      nullptr /*mean->specialBuffer()*/, nullptr /*variance->specialBuffer()*/));

  auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
  if (cudaErr != 0) throw cuda_exception::build("batchnormBpCUDNN: cudaStreamSynchronize failed !", cudaErr);

  NDArray::registerSpecialUse({gradI, gradG, gradB}, {input, mean, variance, gamma, gradO});
}

//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(batchnorm, ENGINE_CUDA) {
  auto input = INPUT_VARIABLE(0);
  auto mean = INPUT_VARIABLE(1);
  auto variance = INPUT_VARIABLE(2);

  NDArray* gamma = nullptr;
  NDArray* beta = nullptr;

  auto output = OUTPUT_VARIABLE(0);

  const bool applyScale = (bool)INT_ARG(0);
  const bool applyOffset = (bool)INT_ARG(1);
  const double epsilon = T_ARG(0);

  if (applyScale) gamma = INPUT_VARIABLE(3);
  if (applyOffset) beta = INPUT_VARIABLE(3 + (int)applyScale);

  const int numOfIntArgs = block.getIArguments()->size();
  const int inRank = input->rankOf();

  // get axes args to normalize input array over
  std::vector<int> axes;
  if (numOfIntArgs > 2)
    for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
  else
    axes.push_back(inRank - 1);  // default dimension to reduce along is last dimension

  const int numOfAxes = axes.size();
  REQUIRE_TRUE(numOfAxes <= inRank, 0,
               "BATCHNORM CUDNN op: too big number of input axes to normalize over, expected number should be less or "
               "equal to rank of input array, but got %i and %i correspondingly !",
               numOfAxes, inRank);

  // evaluate expected shape for mean, variance and gamma. These 3 arrays should have identical shapes
  // for example if input shape is {2,3,4,5,6} and axes = {1,3}, then expected shape would be {1,3,1,5,1}, and if axes =
  // {3}, then expected shape would be {5}
  std::vector<LongType> expShape;
  if (numOfAxes == 1)
    expShape.push_back(input->sizeAt(axes[0]));
  else {  // get, for example, something like {1, inputDim1, 1, inputDim3, 1} if axes = {1, 3}
    expShape = std::vector<LongType>(inRank, 1);
    for (LongType i = 0; i < numOfAxes; ++i) expShape[axes[i]] = input->sizeAt(axes[i]);
  }

  REQUIRE_TRUE(mean->isSameShape(expShape), 0,
               "BATCHNORM CUDNN op: wrong shape of mean array, expected is %s, but got %s instead !",
               ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(mean).c_str());
  REQUIRE_TRUE(variance->isSameShape(expShape), 0,
               "BATCHNORM CUDNN op: wrong shape of variance array, expected is %s, but got %s instead !",
               ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(variance).c_str());
  if (gamma)
  REQUIRE_TRUE(gamma->isSameShape(expShape), 0,
               "BATCHNORM CUDNN op: wrong shape of gamma array, expected is %s, but got %s instead !",
               ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(gamma).c_str());
  if (beta)
  REQUIRE_TRUE(beta->isSameShape(expShape), 0,
               "BATCHNORM CUDNN op: wrong shape of beta array, expected is %s, but got %s instead !",
               ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(beta).c_str());

  // types of all input arrays should be the same
  for (int i = 1; i < block.width(); ++i)
  REQUIRE_TRUE(INPUT_VARIABLE(0)->dataType() == INPUT_VARIABLE(i)->dataType(), 0,
               "BATCHNORM CUDNN op: types of all input arrays should be the same !");

  // cudnn supports NCHW format only
  const bool needPermut = axes.size() == 1 && mean->lengthOf() == input->sizeAt(-1);

  std::unique_ptr<NDArray> tmpGamma = {}, tmpBeta = {}, tmpInput = {}, tmpOutput = {};
  if (needPermut) {  // if NHWC
    std::vector<LongType> perm =
        inRank == 4 ? std::vector<LongType>({0, 3, 1, 2}) : std::vector<LongType>({0, 4, 1, 2, 3});  // NHWC -> NCHW
    tmpInput.reset(new NDArray(input->permute(perm)));
    tmpOutput.reset(new NDArray(output->permute(perm)));
    input = tmpInput.get();
    output = tmpOutput.get();
  }

  // cudnn requires gamma and beta to be non-nullptr
  if (!applyScale) {
    tmpGamma.reset(new NDArray(mean));
    gamma = tmpGamma.get();
    *gamma = 1;
  }
  if (!applyOffset) {
    tmpBeta.reset(new NDArray(mean));
    beta = tmpBeta.get();
    *beta = 0;
  }

  // calculations
  batchnormCUDNN(block.launchContext(), input, mean, variance, gamma, beta, output, epsilon, axes.size() == 1);

  return Status::OK;
}

//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(batchnorm, ENGINE_CUDA) {
  const bool applyScale = (bool)INT_ARG(0);
  const bool applyOffset = (bool)INT_ARG(1);

  NDArray* input = INPUT_VARIABLE(0);
  NDArray* mean = INPUT_VARIABLE(1);
  NDArray* variance = INPUT_VARIABLE(2);
  NDArray* gamma = applyScale ? INPUT_VARIABLE(3) : nullptr;
  NDArray* beta = applyOffset ? INPUT_VARIABLE(3 + (int)applyScale) : nullptr;

  const int numOfIntArgs = block.getIArguments()->size();
  const int xRank = input->rankOf();

  // *********************************** //
  // get axes args to normalize input array over
  std::vector<int> axes;
  if (numOfIntArgs > 2)
    for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
  else
    axes.push_back(xRank - 1);  // default dimension to reduce along is last dimension

  Requirements req("CUDNN BATCHNORM OP");
  req.expectIn(makeInfoVariable(xRank, RANK_MSG_INPUT0), {4, 5}) &&
  req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
               {HALF, FLOAT32, DOUBLE}) &&
  req.expectIn(makeInfoVariable(axes.size(), "axes.size()"), {1, 3, 4}) &&
  req.expect(
      makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1), makeShapeInfoVariable(variance, SHAPE_MSG_INPUT2),
      [](const decltype(mean)& l, const decltype(variance)& r) {
        return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
      },
      EXPECTED_EQ_MSG);
  if (gamma) {
    req.expect(
        makeShapeInfoVariable(gamma, SHAPE_MSG_INPUT_ "#gamma"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
        [](const decltype(gamma)& l, const decltype(mean)& r) {
          return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
        },
        EXPECTED_EQ_MSG);
  }
  if (beta) {
    req.expect(
        makeShapeInfoVariable(beta, SHAPE_MSG_INPUT_ "#beta"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
        [](const decltype(beta)& l, const decltype(mean)& r) {
          return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
        },
        EXPECTED_EQ_MSG);
  }
  if (axes.size() == 1) {
    req.expectIn(makeInfoVariable(mean->lengthOf(), LENGTH_MSG_INPUT1), {-1, 1});
  } else {
    auto inputShapeModif = input->getShapeAsVector();  // [dim0,dim1,dim2,dim3] 4D or [dim0,dim1,dim2,dim3,dim4]
    inputShapeModif[0] = 1;
    // mean [1,dim1,dim2,dim3] 4D or [1,dim1,dim2,dim3,dim4]
    req.expect(
        makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
        makeShapeInfoVariable(inputShapeModif, SHAPE_MSG_INPUT_ "#expect"),
        [](const decltype(mean)& l, const decltype(inputShapeModif)& r) { return l->isSameShape(r); }, EXPECTED_EQ_MSG);
  }
  req.logTheSuccess();
  return req;
}

//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(batchnorm_bp, ENGINE_CUDA) {
  NDArray* input = INPUT_VARIABLE(0);
  NDArray* mean = INPUT_VARIABLE(1);
  NDArray* variance = INPUT_VARIABLE(2);
  NDArray* gamma = nullptr;
  NDArray* beta = nullptr;
  NDArray* gradO = INPUT_VARIABLE(block.width() - 1);  // next epsilon

  NDArray* gradI = OUTPUT_VARIABLE(0);
  NDArray* gradM = OUTPUT_VARIABLE(1);
  NDArray* gradV = OUTPUT_VARIABLE(2);
  NDArray* gradG = nullptr;
  NDArray* gradB = nullptr;

  const bool applyScale = (bool)INT_ARG(0);
  const bool applyOffset = (bool)INT_ARG(1);
  const float epsilon = T_ARG(0);

  if (applyScale) {
    gamma = INPUT_VARIABLE(3);
    gradG = OUTPUT_VARIABLE(3);
  }
  if (applyOffset) {
    beta = INPUT_VARIABLE(3 + (int)applyScale);
    gradB = OUTPUT_VARIABLE(3 + (int)applyScale);
  }

  const int numOfIntArgs = block.getIArguments()->size();
  const int inRank = input->rankOf();

  // get axes args to normalize input array over
  std::vector<int> axes;
  if (numOfIntArgs > 2)
    for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
  else
    axes.push_back(inRank - 1);  // default dimension to reduce along is last dimension

  const int numOfAxes = axes.size();
  REQUIRE_TRUE(numOfAxes <= inRank, 0,
               "BATCHNORM_BP CUDNN op: too big number of input axes to normalize over, expected number should be less "
               "or equal to rank of input array, but got %i and %i correspondingly !",
               numOfAxes, inRank);

  // evaluate expected shape for mean, variance and gamma. These 3 arrays should have identical shapes
  // for example if input shape is {2,3,4,5,6} and axes = {1,3}, then expected shape would be {1,3,1,5,1}, and if axes =
  // {3}, then expected shape would be {5}
  std::vector<LongType> expShape;
  if (numOfAxes == 1)
    expShape.push_back(input->sizeAt(axes[0]));
  else {  // get, for example, something like {1, inputDim1, 1, inputDim3, 1} if axes = {1, 3}
    expShape = std::vector<LongType>(inRank, 1);
    for (LongType i = 0; i < numOfAxes; ++i) expShape[axes[i]] = input->sizeAt(axes[i]);
  }

  REQUIRE_TRUE(mean->isSameShape(expShape), 0,
               "BATCHNORM_BP CUDNN op: wrong shape of mean array, expected is %s, but got %s instead !",
               ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(mean).c_str());
  REQUIRE_TRUE(variance->isSameShape(expShape), 0,
               "BATCHNORM_BP CUDNN op: wrong shape of variance array, expected is %s, but got %s instead !",
               ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(variance).c_str());
  if (gamma)
  REQUIRE_TRUE(gamma->isSameShape(expShape), 0,
               "BATCHNORM_BP CUDNN op: wrong shape of gamma array, expected is %s, but got %s instead !",
               ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(gamma).c_str());
  if (beta)
  REQUIRE_TRUE(beta->isSameShape(expShape), 0,
               "BATCHNORM_BP CUDNN op: wrong shape of beta array, expected is %s, but got %s instead !",
               ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(beta).c_str());

  REQUIRE_TRUE(input->isSameShape(gradO), 0,
               "BATCHNORM_BP CUDNN op: wrong shape of output gradients array, expected is %s, but got %s instead !",
               ShapeUtils::shapeAsString(input).c_str(), ShapeUtils::shapeAsString(gradO).c_str());

  // types of all input arrays should be the same (except gradO)
  for (int i = 1; i < block.width() - 2; ++i)
  REQUIRE_TRUE(INPUT_VARIABLE(0)->dataType() == INPUT_VARIABLE(i)->dataType(), 0,
               "BATCHNORM_BP CUDNN op: types of arrays (input, mean, variance, gamma, beta) should be the same !");

  // cudnn supports NCHW format only
  const bool needPermut = axes.size() == 1 && mean->lengthOf() != input->sizeAt(1);
  std::unique_ptr<NDArray> tmpGamma = {}, tmpGradG = {}, tmpGradB = {}, tmpInput = {}, tmpGradI = {}, tmpGradO = {};
  if (needPermut) {  // if NHWC
    std::vector<LongType> perm =
        inRank == 4 ? std::vector<LongType>({0, 3, 1, 2}) : std::vector<LongType>({0, 4, 1, 2, 3});  // NHWC -> NCHW
    tmpInput.reset(new NDArray(input->permute(perm)));
    tmpGradO.reset(new NDArray(gradO->permute(perm)));
    tmpGradI.reset(new NDArray(gradI->permute(perm)));
    input = tmpInput.get();
    gradO = tmpGradO.get();
    gradI = tmpGradI.get();
  }

  // cudnn requires gamma, gradG, gradB to be non-nullptr
  if (!applyScale) {
    tmpGamma.reset(new NDArray(mean));
    tmpGradG.reset(new NDArray(mean));
    gamma = tmpGamma.get();
    gradG = tmpGradG.get();
    *gamma = 1;
  }
  if (!applyOffset) {
    tmpGradB.reset(new NDArray(mean));
    gradB = tmpGradB.get();
  }

  // calculations
  batchnormBpCUDNN(block.launchContext(), input, mean, variance, gamma, gradO, gradI, gradG, gradB, epsilon,
                   axes.size() == 1);

  *gradM = 0;  // put zeros so far
  *gradV = 0;  // put zeros so far

  return Status::OK;
}

PLATFORM_CHECK(batchnorm_bp, ENGINE_CUDA) {
  NDArray* input = INPUT_VARIABLE(0);
  NDArray* mean = INPUT_VARIABLE(1);
  NDArray* variance = INPUT_VARIABLE(2);
  NDArray* gamma = nullptr;
  NDArray* beta = nullptr;
  NDArray* gradO = INPUT_VARIABLE(block.width() - 1);  // next epsilon

  NDArray* gradI = OUTPUT_VARIABLE(0);
  NDArray* gradM = OUTPUT_VARIABLE(1);
  NDArray* gradV = OUTPUT_VARIABLE(2);
  NDArray* gradG = nullptr;
  NDArray* gradB = nullptr;

  const int numOfIntArgs = block.getIArguments()->size();
  const int xRank = input->rankOf();

  // *********************************** //
  // get axes args to normalize input array over
  std::vector<int> axes;
  if (numOfIntArgs > 2)
    for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
  else
    axes.push_back(xRank - 1);  // default dimension to reduce along is last dimension

  Requirements req("CUDNN BATCHNORM_BP OP");
  req.expectIn(makeInfoVariable(xRank, RANK_MSG_INPUT0), {4, 5}) &&
  req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
               {HALF, FLOAT32, DOUBLE}) &&
  req.expectIn(makeInfoVariable(axes.size(), "axes.size()"), {1, 3, 4}) &&
  req.expect(
      makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1), makeShapeInfoVariable(variance, SHAPE_MSG_INPUT2),
      [](const decltype(mean)& l, const decltype(variance)& r) {
        return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
      },
      EXPECTED_EQ_MSG);
  if (gamma) {
    req.expect(
        makeShapeInfoVariable(gamma, SHAPE_MSG_INPUT_ "#gamma"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
        [](const decltype(gamma)& l, const decltype(mean)& r) {
          return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
        },
        EXPECTED_EQ_MSG);
  }
  if (gradG) {
    req.expect(
        makeShapeInfoVariable(gradG, SHAPE_MSG_INPUT_ "#gradG"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
        [](const decltype(gradG)& l, const decltype(mean)& r) {
          return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
        },
        EXPECTED_EQ_MSG);
  }
  if (gradB) {
    req.expect(
        makeShapeInfoVariable(gradB, SHAPE_MSG_INPUT_ "#gradB"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
        [](const decltype(gradB)& l, const decltype(mean)& r) {
          return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
        },
        EXPECTED_EQ_MSG);
  }
  if (axes.size() == 1) {
    //     isFormatGood = mean->lengthOf() == input->sizeAt(1) || mean->lengthOf() == input->sizeAt(-1);   // mean [C]
    req.expectIn(makeInfoVariable(mean->lengthOf(), LENGTH_MSG_INPUT1), {-1, 1});
  } else {
    auto inputShapeModif = input->getShapeAsVector();  // [dim0,dim1,dim2,dim3] 4D or [dim0,dim1,dim2,dim3,dim4]
    inputShapeModif[0] = 1;
    //     isFormatGood = mean->isSameShape(inputShapeModif);    // mean [1,dim1,dim2,dim3] 4D or
    //     [1,dim1,dim2,dim3,dim4]
    req.expect(
        makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
        makeShapeInfoVariable(inputShapeModif, SHAPE_MSG_INPUT_ "#expect"),
        [](const decltype(mean)& l, const decltype(inputShapeModif)& r) { return l->isSameShape(r); }, EXPECTED_EQ_MSG);
  }
  req.logTheSuccess();
  return req;
}

}  // namespace platforms
}  // namespace ops
}  // namespace sd