/* ******************************************************************************
 *
 *
 * 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 AbdelRauf
//

#include <array/NDArrayFactory.h>
#include <ops/declarable/OpRegistrator.h>

#include "cudnnUtils.h"

namespace sd {
namespace ops {
namespace platforms {

// our implementation designed for 1 physical layer
constexpr int numLayers = 1;

// we will copy without using cudnnGetRNNLinLayerMatrixParams : 1 pseudo layer , isBidirectional : 2 pseudo layer
void copyWeights(const cudaStream_t &stream, bool isBidirectional, uint8_t *weightsSpace, size_t weightsSize,
                 uint8_t *inputWeightsData, uint8_t *recurrentWeightsData, uint8_t *biasesData, LongType inputSize,
                 int hiddenSize, int dataTypeSize) {
  int pseudo_layer_count = isBidirectional ? 2 : 1;
  uint8_t *wptr = weightsSpace;
  auto wEnd = wptr + weightsSize;

  // copy size for 1 full pseudo layer
  // in bidirectional 1 layer consist of 2 pseduo layers
  auto input_pseudo_size = 4 * inputSize * hiddenSize * dataTypeSize;
  auto hidden_pseudo_size = 4 * hiddenSize * hiddenSize * dataTypeSize;
  for (LongType i = 0; i < pseudo_layer_count; i++) {
    if (wptr + input_pseudo_size + hidden_pseudo_size > wEnd) return;
    // copy input weights
    if (inputWeightsData) {
      cudaMemcpyAsync(wptr, inputWeightsData, input_pseudo_size, cudaMemcpyDeviceToDevice, stream);
      inputWeightsData += input_pseudo_size;
    }
    wptr += input_pseudo_size;
    // copy recurrent weights
    if (recurrentWeightsData) {
      cudaMemcpyAsync(wptr, recurrentWeightsData, hidden_pseudo_size, cudaMemcpyDeviceToDevice, stream);
      recurrentWeightsData += hidden_pseudo_size;
    }
    wptr += hidden_pseudo_size;
  }

  // copy bias first 4
  auto bias_size = 4 * hiddenSize * dataTypeSize;
  for (int i = 0; i < pseudo_layer_count; i++) {
    // refill first 4 biases
    if (biasesData && wptr + bias_size < wEnd) {
      cudaMemcpyAsync(wptr, biasesData, bias_size, cudaMemcpyDeviceToDevice, stream);
      biasesData += bias_size;
    }
    wptr += bias_size;
    // refill next 4 with zeros
    if (wptr + bias_size < wEnd) {
      cudaMemsetAsync(wptr, 0, bias_size, stream);
      wptr += bias_size;
    }
  }
  // memset the rest
  if (wEnd - wptr) cudaMemsetAsync(wptr, 0, wEnd - wptr, stream);
}

void cudnn_rnn_old(LaunchContext *contextPtr, int dataFormat, NDArray *input, NDArray *inputWeights,
                   NDArray *recurrentWeights, NDArray *biases, NDArray *prevAct, NDArray *prevMemCell,
                   NDArray *outputActivations, NDArray *finalTimeStepActivations, NDArray *finalMemCellState,
                   LongType maxSeqLength, LongType batchSize, LongType inputSize, LongType hiddenSize, double cellClip,
                   bool isBidirectional) {
  sd_debug("cudnn rnn api %s \n", "v6");

  bool training = false;
  cudnnHandle_t handle = *(reinterpret_cast<cudnnHandle_t *>(contextPtr->getCuDnnHandle()));

  auto stream = *(contextPtr->getCudaStream());
  CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(handle, stream));

  CudnnTensorList xDescList(maxSeqLength);
  CudnnTensorList yDescList(maxSeqLength);

  auto cudnnType = cudnnDataType(input->dataType());
  auto dataTypeSize = input->sizeOfT();

  CudnnTensor hxDesc, cxDesc, hyDesc, cyDesc;

  constexpr int rankOf = 3;
  const int numDirections = isBidirectional ? 2 : 1;

  const int dimsX[rankOf] = {static_cast<int>(batchSize), static_cast<int>(inputSize), 1};
  const int stridesX[rankOf] = {static_cast<int>(inputSize), 1, 1};

  const int dimsY[rankOf] = {static_cast<int>(batchSize), static_cast<int>(hiddenSize * numDirections), 1};
  const int stridesY[rankOf] = {static_cast<int>(hiddenSize * numDirections), 1, 1};

  const int dimC[rankOf] = {static_cast<int>(numLayers * numDirections), static_cast<int>(batchSize), static_cast<int>(hiddenSize)};
  const int strideC[rankOf] = {static_cast<int>(batchSize * hiddenSize), static_cast<int>(hiddenSize), 1};
  for (int i = 0; i < maxSeqLength; i++) {
    xDescList.set(i, cudnnType, rankOf, dimsX, stridesX);
    yDescList.set(i, cudnnType, rankOf, dimsY, stridesY);
  }

  auto xDesc0 = xDescList.get(0);

  hxDesc.set(cudnnType, rankOf, dimC, strideC);
  cxDesc.set(cudnnType, rankOf, dimC, strideC);
  hyDesc.set(cudnnType, rankOf, dimC, strideC);
  cyDesc.set(cudnnType, rankOf, dimC, strideC);

  PointersManager manager(contextPtr, __func__);
  // dropout section
  DropoutDesc dropoutDesc(nullptr);
  // dropout
  float dropout = 0;
  size_t sizeInBytes = 0;
  void *droupoutMem = nullptr;
  uint64_t seed = 1;  // seed
  if (dropout != 0) {
    dropoutDesc.create();
    CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnDropoutGetStatesSize), cudnnDropoutGetStatesSize(handle, &sizeInBytes));
    // allocate and set
    droupoutMem = manager.allocateDevMem(sizeInBytes);
    dropoutDesc.set(handle, dropout, droupoutMem, sizeInBytes, seed);
  }

  // RNN
  RnnDesc rnnDesc;
  cudnnRNNMode_t rnnCellMode = CUDNN_LSTM;
  cudnnRNNAlgo_t algo = CUDNN_RNN_ALGO_STANDARD;

  auto direction = isBidirectional ? CUDNN_BIDIRECTIONAL : CUDNN_UNIDIRECTIONAL;
  auto mathPrec = cudnnType;

  // Note: We will set some parameters manually
  constexpr auto inputMode = CUDNN_LINEAR_INPUT;
  rnnDesc.setUsingOldAPI(handle, inputMode, direction, rnnCellMode, algo, mathPrec, hiddenSize, numLayers, dropoutDesc);
#if CUDNN_VERSION >= CUDNN_CLIPPING_API_VER
  if (cellClip > 0 && cudnnGetVersion() >= CUDNN_CLIPPING_API_VER) {
    CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnRNNSetClip), cudnnRNNSetClip(handle, rnnDesc, CUDNN_RNN_CLIP_MINMAX,
                                                                        CUDNN_PROPAGATE_NAN, -cellClip, cellClip));
  }
#endif
  // set up parameters
  size_t weightsSize = 0;
  CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetRNNParamsSize),
                          cudnnGetRNNParamsSize(handle, rnnDesc, xDesc0, &weightsSize, cudnnType));

  FilterDesc wDesc;
  int dimW[] = {static_cast<int>(weightsSize / dataTypeSize), 1, 1};

  wDesc.set(cudnnType, CUDNN_TENSOR_NCHW, 3, dimW);
  // allocation
  void *weightsSpace = manager.allocateDevMem(weightsSize);

  size_t workSpaceSizeInBytes = 0;
  size_t reserveSpaceSizeInBytes = 0;

  CHECK_CUDNN_FAILURE_MSG(
      STRINGIZE(cudnnGetRNNWorkspaceSize),
      cudnnGetRNNWorkspaceSize(handle, rnnDesc, maxSeqLength, xDescList.getDescriptors(), &workSpaceSizeInBytes));

  void *workSpace = manager.allocateDevMem(workSpaceSizeInBytes);
  void *reserveSpace = nullptr;
  // training
  if (training) {
    CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetRNNTrainingReserveSize),
                            cudnnGetRNNTrainingReserveSize(handle, rnnDesc, maxSeqLength, xDescList.getDescriptors(),
                                                           &reserveSpaceSizeInBytes));
    reserveSpace = manager.allocateDevMem(reserveSpaceSizeInBytes);
  }

  NDArray::prepareSpecialUse({outputActivations, finalTimeStepActivations, finalMemCellState},
                             {input, inputWeights, recurrentWeights, biases, prevAct, prevMemCell});

  uint8_t *biasesData = biases ? (uint8_t *)biases->specialBuffer() : nullptr;
  auto prevActData = prevAct ? prevAct->specialBuffer() : nullptr;
  auto prevMemCellData = prevMemCell ? prevMemCell->specialBuffer() : nullptr;
  auto finalTimeStepActivationsData = finalTimeStepActivations ? finalTimeStepActivations->specialBuffer() : nullptr;
  auto finalMemCellStateData = finalMemCellState ? finalMemCellState->specialBuffer() : nullptr;

  // dimension 4*nOut implies order it, ft, c't, ot
  // input gate, forget gate, new gate, output gate, input gate, forget gate, new gate, output gate
  // Note: our weights should be transposed and duplicated with C order to match cudnn ones

  NDArray inputWeightsT, recurrentWeightsT;
  uint8_t *inputWeightsData = nullptr;
  uint8_t *recurrentWeightsData = nullptr;
  if (inputWeights) {
    inputWeightsT =
        inputWeights->rankOf() == 3 ? inputWeights->permute({0, 2, 1}, 0, false).dup('c') : inputWeights->transpose().dup('c');
    inputWeightsData = (uint8_t *)inputWeightsT.specialBuffer();
  }
  if (recurrentWeights) {
    recurrentWeightsT = recurrentWeights->rankOf() == 3 ? recurrentWeights->permute({0, 2, 1}, 0, false).dup('c')
                                                        : recurrentWeights->transpose().dup('c');
    recurrentWeightsData = (uint8_t *)recurrentWeightsT.specialBuffer();
  }

  // copy without cudnnGetRNNLinLayerMatrixParams
  copyWeights(stream, isBidirectional, (uint8_t *)weightsSpace, weightsSize, inputWeightsData, recurrentWeightsData,
              biasesData, inputSize, hiddenSize, dataTypeSize);

  // permute based on dataformat
  NDArray *argX = input;
  NDArray *argOutput = outputActivations;
  NDArray permutedX, outputH;

  if (outputActivations != nullptr && (dataFormat != 0 || outputActivations->ordering() != 'c')) {
    outputH = NDArray('c', std::vector<LongType>{maxSeqLength, batchSize, (numDirections * hiddenSize)},
                      outputActivations->dataType(), contextPtr);
    argOutput = &outputH;
  }

  if (dataFormat == 1) {
    permutedX = input->permute({1, 0, 2}, 0, false).dup('c');
    argX = &permutedX;
  }

  auto xData = argX->specialBuffer();
  auto yData = argOutput ? argOutput->specialBuffer() : nullptr;

  if (training) {
    CHECK_CUDNN_FAILURE_MSG(
        STRINGIZE(cudnnRNNForwardTraining),
        cudnnRNNForwardTraining(handle, rnnDesc, (int)maxSeqLength, xDescList.getDescriptors(), xData, hxDesc,
                                prevActData, cxDesc, prevMemCellData, wDesc, weightsSpace, yDescList.getDescriptors(),
                                yData, hyDesc, finalTimeStepActivationsData, cyDesc, finalMemCellStateData, workSpace,
                                workSpaceSizeInBytes, reserveSpace, reserveSpaceSizeInBytes));
  } else {
    CHECK_CUDNN_FAILURE_MSG(
        STRINGIZE(cudnnRNNForwardInference),
        cudnnRNNForwardInference(handle, rnnDesc, (int)maxSeqLength, xDescList.getDescriptors(), xData, hxDesc,
                                 prevActData, cxDesc, prevMemCellData, wDesc, weightsSpace, yDescList.getDescriptors(),
                                 yData, hyDesc, finalTimeStepActivationsData, cyDesc, finalMemCellStateData, workSpace,
                                 workSpaceSizeInBytes));
  }

  // remap output
  if (outputActivations != nullptr && argOutput != outputActivations) {
    // refill output
    if (dataFormat == 1) {
      std::vector<sd::LongType> permute = {1,0,2};
      NDArray assign = argOutput->permute(permute, 0, false);
      outputActivations->assign(&assign);
    }
  }
  NDArray::registerSpecialUse({outputActivations, finalTimeStepActivations, finalMemCellState},
                              {input, inputWeights, recurrentWeights, biases, prevAct, prevMemCell});

  return;
}

#if CUDNN_VERSION >= CUDNN_NEW_RNN_API_VER

void cudnn_rnn_v8(LaunchContext *contextPtr, int dataFormat, NDArray *input, NDArray *seqLengthArray,
                  NDArray *inputWeights, NDArray *recurrentWeights, NDArray *biases, NDArray *prevAct,
                  NDArray *prevMemCell, NDArray *outputActivations, NDArray *finalTimeStepActivations,
                  NDArray *finalMemCellState, int maxSeqLength, int batchSize, int inputSize, int hiddenSize,
                  double cellClip, bool isBidirectional) {
  sd_debug("cudnn rnn api %s \n", "v8");
  // seqLengthArray should be int
  NDArray *argSeqNdArray = nullptr;
  NDArray seqArrIntData;
  if (seqLengthArray) {
    if (seqLengthArray->ews() == 1 && seqLengthArray->dataType() == INT32) {
      argSeqNdArray = seqLengthArray;
    } else {
      if (seqLengthArray->dataType() != INT32) {
        seqArrIntData = seqLengthArray->cast(INT32);
        if (seqArrIntData.ews() != 1) seqArrIntData = seqArrIntData.dup('c');
      } else {
        seqArrIntData = seqLengthArray->dup('c');
      }
      argSeqNdArray = &seqArrIntData;
    }
  } else {
    seqArrIntData = NDArray('c', std::vector<LongType>{batchSize}, INT32, contextPtr);
    seqArrIntData.assign(maxSeqLength);
    argSeqNdArray = &seqArrIntData;
  }
  PointersManager manager(contextPtr, __func__);
  bool training = false;
  cudnnHandle_t handle = *(reinterpret_cast<cudnnHandle_t *>(contextPtr->getCuDnnHandle()));
  auto stream = *(contextPtr->getCudaStream());
  CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(handle, stream));

  auto cudnnType = cudnnDataType(input->dataType());
  auto dataTypeSize = input->sizeOfT();

  CudnnTensor hDesc, cDesc;

  constexpr int rankOf = 3;
  const int numDirections = isBidirectional ? 2 : 1;

  const int dimC[rankOf] = {numLayers * numDirections, batchSize, hiddenSize};
  const int strideC[rankOf] = {batchSize * hiddenSize, hiddenSize, 1};

  hDesc.set(cudnnType, rankOf, dimC, strideC);
  cDesc.set(cudnnType, rankOf, dimC, strideC);

  // dropout section
  DropoutDesc dropoutDesc(nullptr);
  // dropout
  float dropout = 0;
  size_t sizeInBytes = 0;
  void *droupoutMem = nullptr;
  uint64_t seed = 1;  // seed
  if (dropout != 0) {
    dropoutDesc.create();
    CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnDropoutGetStatesSize), cudnnDropoutGetStatesSize(handle, &sizeInBytes));
    // allocate and set
    droupoutMem = manager.allocateDevMem(sizeInBytes);
    dropoutDesc.set(handle, dropout, droupoutMem, sizeInBytes, seed);
  }

  // RNN
  RnnDesc rnnDesc;
  cudnnRNNMode_t rnnCellMode = CUDNN_LSTM;
  cudnnRNNAlgo_t algo = CUDNN_RNN_ALGO_STANDARD;
  auto direction = isBidirectional ? CUDNN_BIDIRECTIONAL : CUDNN_UNIDIRECTIONAL;
  auto mathPrec = cudnnType;

  // Note: We will set some parameters manually. Some of them could be parameter in future
  constexpr auto inputMode = CUDNN_LINEAR_INPUT;
  bool use_tensor_ops = false;  // could be parameter in future
#if CUDNN_VERSION >= CUDNN_NEW_RNN_API_VER
  cudnnMathType_t mathType = use_tensor_ops ? CUDNN_TENSOR_OP_MATH : CUDNN_FMA_MATH;
#else
  cudnnMathType_t mathType = use_tensor_ops ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH;
#endif
  // disable projection
  int projSize = hiddenSize;
  cudnnRNNBiasMode_t bias_mode = CUDNN_RNN_DOUBLE_BIAS;
  uint32_t aux_flags = CUDNN_RNN_PADDED_IO_ENABLED;

  rnnDesc.set(algo, rnnCellMode, bias_mode, direction, inputMode, cudnnType, mathPrec, mathType, inputSize, hiddenSize,
              projSize, numLayers, dropoutDesc, aux_flags);
  if (cellClip > 0) {
    CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnRNNSetClip), cudnnRNNSetClip(handle, rnnDesc, CUDNN_RNN_CLIP_MINMAX,
                                                                        CUDNN_PROPAGATE_NAN, -cellClip, cellClip));
  }
  // set Data desc
  RnnDataDesc xDataDesc, yDataDesc;
  bool time_major = false;
  float padding_fill = 0.0f;
  auto hostSeqArr = bufferInHost<int>(*argSeqNdArray);
  cudnnRNNDataLayout_t layout =
      dataFormat == 0 ? CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_UNPACKED : CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED;
  xDataDesc.set(cudnnType, layout, maxSeqLength, batchSize, inputSize, hostSeqArr, (void *)&padding_fill);
  yDataDesc.set(cudnnType, layout, maxSeqLength, batchSize, hiddenSize * numDirections, hostSeqArr,
                (void *)&padding_fill);
  // set up parameters
  size_t weightsSize = 0;
  CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetRNNWeightSpaceSize),
                          cudnnGetRNNWeightSpaceSize(handle, rnnDesc, &weightsSize));

  // allocation
  void *weightsSpace = manager.allocateDevMem(weightsSize);

  // Set up work space and reserved memory
  void *workSpace = nullptr;
  void *reserveSpace = nullptr;

  size_t workSpaceSizeInBytes = 0;
  size_t reserveSpaceSizeInBytes = 0;

  cudnnForwardMode_t fwdMode = training ? CUDNN_FWD_MODE_TRAINING : CUDNN_FWD_MODE_INFERENCE;
  CHECK_CUDNN_FAILURE_MSG(
      STRINGIZE(cudnnGetRNNTempSpaceSizes),
      cudnnGetRNNTempSpaceSizes(handle, rnnDesc, fwdMode, xDataDesc, &workSpaceSizeInBytes, &reserveSpaceSizeInBytes));
  workSpace = manager.allocateDevMem(workSpaceSizeInBytes);
  // training
  if (training) {
    reserveSpace = manager.allocateDevMem(reserveSpaceSizeInBytes);
  }

  NDArray::prepareSpecialUse({outputActivations, finalTimeStepActivations, finalMemCellState},
                             {input, inputWeights, recurrentWeights, biases, prevAct, prevMemCell, argSeqNdArray});

  auto xData = input->specialBuffer();
  uint8_t *biasesData = biases ? (uint8_t *)biases->specialBuffer() : nullptr;
  auto prevActData = prevAct ? prevAct->specialBuffer() : nullptr;
  auto prevMemCellData = prevMemCell ? prevMemCell->specialBuffer() : nullptr;
  auto yData = outputActivations ? outputActivations->specialBuffer() : nullptr;
  auto finalTimeStepActivationsData = finalTimeStepActivations ? finalTimeStepActivations->specialBuffer() : nullptr;
  auto finalMemCellStateData = finalMemCellState ? finalMemCellState->specialBuffer() : nullptr;

  // dimension 4*nOut implies order it, ft, c't, ot
  // input gate, forget gate, new gate, output gate, input gate, forget gate, new gate, output gate
  // Note: our weights should be transposed and duplicated with C order to match cudnn ones

  NDArray inputWeightsT, recurrentWeightsT;
  uint8_t *inputWeightsData = nullptr;
  uint8_t *recurrentWeightsData = nullptr;
  if (inputWeights) {
    inputWeightsT =
        inputWeights->rankOf() == 3 ? inputWeights->permute({0, 2, 1}).dup('c') : inputWeights->transpose().dup('c');
    inputWeightsData = (uint8_t *)inputWeightsT.specialBuffer();
  }
  if (recurrentWeights) {
    recurrentWeightsT = recurrentWeights->rankOf() == 3 ? recurrentWeights->permute({0, 2, 1}).dup('c')
                                                        : recurrentWeights->transpose().dup('c');
    recurrentWeightsData = (uint8_t *)recurrentWeightsT.specialBuffer();
  }

  // copy without cudnnGetRNNLinLayerMatrixParams
  copyWeights(stream, isBidirectional, (uint8_t *)weightsSpace, weightsSize, inputWeightsData, recurrentWeightsData,
              biasesData, inputSize, hiddenSize, dataTypeSize);

  CHECK_CUDNN_FAILURE_MSG(
      STRINGIZE(cudnnRNNForward),
      cudnnRNNForward(handle, rnnDesc, fwdMode, (const int32_t *)argSeqNdArray->specialBuffer(), xDataDesc, xData,
                      yDataDesc, yData, hDesc, prevActData, finalTimeStepActivationsData, cDesc, prevMemCellData,
                      finalMemCellStateData, weightsSize, weightsSpace, workSpaceSizeInBytes, workSpace,
                      reserveSpaceSizeInBytes, reserveSpace));

  NDArray::registerSpecialUse({outputActivations, finalTimeStepActivations, finalMemCellState},
                              {input, inputWeights, recurrentWeights, biases, prevAct, prevMemCell});

  return;
}

#endif

//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(lstmLayer, ENGINE_CUDA) {
  const auto dataFormat = INT_ARG(0);  // for unidirectional: 0 = [sL, bS, nIn], 1 = [bS, sL ,nIn], 2 = [bS, nIn, sL],
                                       // for bidirectional: 3 = [sL, 2, bS, nOut] (for ONNX)
  const LongType directionMode =
      INT_ARG(1);  // direction: 0 = fwd, 1 = bwd, 2 = bidirectional sum, 3 = bidirectional concat, 4 = bidirectional
                   // extra output dim (in conjunction with format dataFormat = 3)

  const auto hasBiases = B_ARG(0);       // indicates whether biases array is provided
  const auto hasSeqLenArray = B_ARG(1);  // indicates whether seqLen array is provided
  const auto hasInitH = B_ARG(2);        // indicates whether initial output is provided
  const auto hasInitC = B_ARG(3);        // indicates whether initial cell state is provided
  const auto hasPH = B_ARG(4);           // indicates whether peephole connections are present
  const auto retFullSeq = B_ARG(5);      // indicates whether to return whole time sequence h {h_0, h_1, ... , h_sL-1}
  const auto retLastH = B_ARG(6);  // indicates whether to return output at last time step only, in this case shape
                                   // would be [bS, nOut] (exact shape depends on dataFormat argument)
  const auto retLastC = B_ARG(7);  // indicates whether to return cells state at last time step only, in this case shape
                                   // would be [bS, nOut] (exact shape depends on dataFormat argument)

  const auto cellClip = T_ARG(0);  // cell clipping value, if it = 0 then do not apply clipping

  const auto x = INPUT_VARIABLE(0);   // input
  const auto Wx = INPUT_VARIABLE(1);  // input weights
  const auto Wr = INPUT_VARIABLE(2);  // recurrent weights

  int count = 3;
  const auto b = hasBiases ? INPUT_VARIABLE(count++) : nullptr;                    // biases
  const auto seqLengthArray = hasSeqLenArray ? INPUT_VARIABLE(count++) : nullptr;  // seqLen vector
  const auto hI = hasInitH ? INPUT_VARIABLE(count++) : nullptr;                    // initial output
  const auto cI = hasInitC ? INPUT_VARIABLE(count++) : nullptr;                    // initial cell state
  const auto Wp = hasPH ? INPUT_VARIABLE(count++) : nullptr;                       // peephole weights

  count = 0;
  auto h = retFullSeq ? OUTPUT_VARIABLE(count++) : nullptr;  // output
  auto hL = retLastH ? OUTPUT_VARIABLE(count++) : nullptr;   // output at last step
  auto cL = retLastC ? OUTPUT_VARIABLE(count++) : nullptr;   // cell state at last step

  REQUIRE_TRUE(cellClip >= 0, 0, "LSTM_LAYER operation: cell clipping value should be nonnegative (>=0) !");
  REQUIRE_TRUE(retFullSeq || retLastH || retLastC, 0,
               "LSTM_LAYER operation: please specify what output arrays to produce !");
  // evaluate dimensions
  const LongType seqLength = dataFormat == 3 ? x->sizeAt(0) : x->sizeAt(dataFormat);
  const LongType bS = dataFormat == 1 || dataFormat == 2 ? x->sizeAt(0) : x->sizeAt(1);
  const LongType nIn = dataFormat == 2 ? x->sizeAt(1) : x->sizeAt(2);
  const LongType nOut = Wx->sizeAt(-1) / 4;
  const LongType hiddenSize = nOut;

  auto contextPtr = block.launchContext();
  bool isBidirectional = directionMode >= 2;

  if (!isBidirectional) {  // no bidirectional
    // Wx validation
    if (Wx->rankOf() != 2 || Wx->sizeAt(0) != nIn)
      REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of input weights, expected is %s, but got %s instead !",
                   ShapeUtils::shapeAsString({nIn, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
    // Wr validation
    if (Wr->rankOf() != 2 || Wr->sizeAt(0) != nOut || Wr->sizeAt(1) != 4 * nOut)
      REQUIRE_TRUE(false, 0,
                   "LSTM_LAYER operation: wrong shape of recurrent weights, expected is %s, but got %s instead !",
                   ShapeUtils::shapeAsString({nOut, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wr).c_str());
    // biases validation
    if (b != nullptr && (b->rankOf() != 1 || b->sizeAt(0) != 4 * nOut))
      REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of biases, expected is %s, but got %s instead !",
                   ShapeUtils::shapeAsString({4 * nOut}).c_str(), ShapeUtils::shapeAsString(b).c_str());
    // initial output validation
    if (hI != nullptr && (hI->rankOf() != 2 || hI->sizeAt(0) != bS || hI->sizeAt(1) != nOut))
      REQUIRE_TRUE(false, 0,
                   "LSTM_LAYER operation: wrong shape of initial output, expected is %s, but got %s instead !",
                   ShapeUtils::shapeAsString({bS, nOut}).c_str(), ShapeUtils::shapeAsString(hI).c_str());
    // initial cell  validation
    if (cI != nullptr && (cI->rankOf() != 2 || cI->sizeAt(0) != bS || cI->sizeAt(1) != nOut))
      REQUIRE_TRUE(false, 0,
                   "LSTM_LAYER operation: wrong shape of initial cell state, expected is %s, but got %s instead !",
                   ShapeUtils::shapeAsString({bS, nOut}).c_str(), ShapeUtils::shapeAsString(cI).c_str());
  } else {  // bidirectional
            // Wx validation
    if (Wx->rankOf() != 3 || Wx->sizeAt(0) != 2 || Wx->sizeAt(1) != nIn)
      REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of input weights, expected is %s, but got %s instead !",
                   ShapeUtils::shapeAsString({2, nIn, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
    // Wr validation
    if (Wr->rankOf() != 3 || Wr->sizeAt(0) != 2 || Wr->sizeAt(1) != nOut || Wr->sizeAt(2) != 4 * nOut)
      REQUIRE_TRUE(false, 0,
                   "LSTM_LAYER operation: wrong shape of recurrent weights, expected is %s, but got %s instead !",
                   ShapeUtils::shapeAsString({2, nOut, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wr).c_str());
    // biases validation
    if (b != nullptr && (b->rankOf() != 2 || b->sizeAt(0) != 2 || b->sizeAt(1) != 4 * nOut))
      REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of biases, expected is %s, but got %s instead !",
                   ShapeUtils::shapeAsString({2, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(b).c_str());
    // initial output validation
    if (hI != nullptr && (hI->rankOf() != 3 || hI->sizeAt(0) != 2 || hI->sizeAt(1) != bS || hI->sizeAt(2) != nOut))
      REQUIRE_TRUE(false, 0,
                   "LSTM_LAYER operation: wrong shape of initial output, expected is %s, but got %s instead !",
                   ShapeUtils::shapeAsString({2, bS, nOut}).c_str(), ShapeUtils::shapeAsString(hI).c_str());
    // initial cell  validation
    if (cI != nullptr && (cI->rankOf() != 3 || cI->sizeAt(0) != 2 || cI->sizeAt(1) != bS || cI->sizeAt(2) != nOut))
      REQUIRE_TRUE(false, 0,
                   "LSTM_LAYER operation: wrong shape of initial cell state, expected is %s, but got %s instead !",
                   ShapeUtils::shapeAsString({2, bS, nOut}).c_str(), ShapeUtils::shapeAsString(cI).c_str());
  }

#if CUDNN_VERSION < CUDNN_NEW_RNN_API_VER
  cudnn_rnn_old(contextPtr, dataFormat, x, Wx, Wr, b, hI, cI, h, hL, cL, seqLength, bS, nIn, hiddenSize,
                (double)cellClip, isBidirectional);
#else
  if (cudnnGetVersion() >= CUDNN_NEW_RNN_API_VER) {
    cudnn_rnn_v8(contextPtr, dataFormat, x, seqLengthArray, Wx, Wr, b, hI, cI, h, hL, cL, seqLength, bS, nIn,
                 hiddenSize, (double)cellClip, isBidirectional);
  } else {
    cudnn_rnn_old(contextPtr, dataFormat, x, Wx, Wr, b, hI, cI, h, hL, cL, seqLength, bS, nIn, hiddenSize,
                  (double)cellClip, isBidirectional);
  }
#endif

  return Status::OK;
}

// Cudnn Lstm:
// Forward inference implemented using v6, and v8 (when version > 8.0.1) api calls.
// As our Cuda Lstm implementation has 1 layer. Cudnn implementation was implemented for 1 physical layer
// Cudnn helper restrictions:
//  - all NDArrays should be the same type
//  - dataFormat should be 0 or 1
//  - only unidirectional (directionMode == 0) and bidirectional concat (directionMode == 3)
//  - no peephole connection
//  - Clipping is allowed for cudnn version >= 7.2.1
//  - SeqLen array is allowed for cudnn version >= 8.0.1
//  - gateActivation: sigmoid, cellActivation and outputActivation: tanh
//  - NDArrays (excluding the weight arrays, as we have to transpose or permute it) should follow 'c' order and ews()==1
PLATFORM_CHECK(lstmLayer, ENGINE_CUDA) {
  const auto dataFormat = INT_ARG(0);  // for unidirectional: 0 = [sL, bS, nIn], 1 = [bS, sL ,nIn], 2 = [bS, nIn, sL],
                                       // for bidirectional: 3 = [sL, 2, bS, nOut] (for ONNX)
  const auto directionMode =
      INT_ARG(1);  // direction: 0 = fwd, 1 = bwd, 2 = bidirectional sum, 3 = bidirectional concat, 4 = bidirectional
                   // extra output dim (in conjunction with format dataFormat = 3)
  // integer numbers corresponding to activations: 0=tanh, 1=relu, 2=sigmoid, 3=affine, 4=leaky relu, 5= thresholded
  // relu, 6=scaled tanh, 7=hard sigmoid, 8=ELU, 9=softsign, 10=softplus
  const auto gateAct = INT_ARG(2);  // activation for input (i), forget (f) and output (o) gates
  const auto cellAct = INT_ARG(3);  // activation for cell state (c)
  const auto outAct = INT_ARG(4);   // activation for output (h)

  const auto hasBiases = B_ARG(0);       // indicates whether biases array is provided
  const auto hasSeqLenArray = B_ARG(1);  // indicates whether seqLen array is provided
  const auto hasInitH = B_ARG(2);        // indicates whether initial output is provided
  const auto hasInitC = B_ARG(3);        // indicates whether initial cell state is provided
  const auto hasPH = B_ARG(4);           // indicates whether peephole connections are present
  const auto retFullSeq = B_ARG(5);      // indicates whether to return whole time sequence h {h_0, h_1, ... , h_sL-1}
  const auto retLastH = B_ARG(6);  // indicates whether to return output at last time step only, in this case shape
                                   // would be [bS, nOut] (exact shape depends on dataFormat argument)
  const auto retLastC = B_ARG(7);  // indicates whether to return cells state at last time step only, in this case shape
                                   // would be [bS, nOut] (exact shape depends on dataFormat argument)

  const auto cellClip = T_ARG(0);  // cell clipping value, if it = 0 then do not apply clipping

  const auto x = INPUT_VARIABLE(0);   // input
  const auto Wx = INPUT_VARIABLE(1);  // input weights
  const auto Wr = INPUT_VARIABLE(2);  // recurrent weights

  int count = 3;
  const auto b = hasBiases ? INPUT_VARIABLE(count++) : nullptr;  // biases
  const auto hI = hasInitH ? INPUT_VARIABLE(count++) : nullptr;  // initial output
  const auto cI = hasInitC ? INPUT_VARIABLE(count++) : nullptr;  // initial cell state

  count = 0;
  auto h = retFullSeq ? OUTPUT_VARIABLE(count++) : nullptr;  // output
  auto hL = retLastH ? OUTPUT_VARIABLE(count++) : nullptr;   // output at last step
  auto cL = retLastC ? OUTPUT_VARIABLE(count++) : nullptr;   // cell state at last step

  DataType xType = x->dataType();
  DataType WxType = Wx->dataType();
  DataType WrType = Wr->dataType();

  Requirements req("CUDNN LSTMLAYER OP");
  // cudnn related restrictions    //gateAct: sigmoid, cellAct: tanh adn et cetera
  // integer numbers corresponding to activations: 0=tanh, 1=relu, 2=sigmoid, 3=affine,
  // 4=leaky relu, 5= thresholded relu, 6=scaled tanh, 7=hard sigmoid, 8=ELU, 9=softsign, 10=softplus
  req.expectEq(makeInfoVariable(gateAct, "gate Activation"), makeInfoVariable(2, "sigmoid")) &&
      req.expectEq(makeInfoVariable(cellAct, "cell Activation"), makeInfoVariable(2, "tanh")) &&
      req.expectEq(makeInfoVariable(outAct, "out Activation"), makeInfoVariable(2, "tanh")) &&
      req.expectFalse(makeInfoVariable(hasPH, HAVE_PEEPHOLE), EXPECTED_NOT_SUPPORTED) &&
      req.expectIn(makeInfoVariable(directionMode, "directionMode"), {0, 3}) &&
      req.expectIn(makeInfoVariable(dataFormat, "data Format"), {0, 1});

  if (req) {
    // cudnn api version related restrictions in our helpers
    size_t cudnn_version = cudnnGetVersion();
    // though seqlengthArray was added in earlier versions we do not handle it below 8.0.0.1
#if CUDNN_VERSION < CUDNN_NEW_RNN_API_VER
    // implRestrictions = implRestrictions && !hasSeqLenArray;
    req.expectFalse(makeInfoVariable(hasSeqLenArray, HAVE_SEQLENARR), EXPECTED_NOT_SUPPORTED);
#else
    // implRestrictions = implRestrictions && (cudnn_version >= CUDNN_NEW_RNN_API_VER || !hasSeqLenArray);
    if (cudnn_version < CUDNN_NEW_RNN_API_VER) {
      req.expectFalse(makeInfoVariable(hasSeqLenArray, HAVE_SEQLENARR), EXPECTED_NOT_SUPPORTED);
    }
#endif
    // implRestrictions = implRestrictions && (cudnn_version >= CUDNN_CLIPPING_API_VER || cellClip==0);
    if (cudnn_version < CUDNN_CLIPPING_API_VER) {
      req.expectEq(makeInfoVariable(cellClip, MSG_CELL_CLIPPING), 0);
    }
  }
  // restriction that comes either from not setting Descriptor or not handling manipulation:
  // restrict0: the same types
  req.expectEq(makeInfoVariable(x->ordering(), ORDERING_MSG_INPUT0), 'c') &&
      req.expectEq(makeInfoVariable(WxType, TYPE_MSG_INPUT1), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
      req.expectEq(makeInfoVariable(WrType, TYPE_MSG_INPUT2), makeInfoVariable(xType, TYPE_MSG_INPUT0));
  if (b)
    req.expectEq(makeInfoVariable(b->dataType(), TYPE_MSG_INPUT_ "#bias"), makeInfoVariable(xType, TYPE_MSG_INPUT0));
  if (hI) {
    req.expectEq(makeInfoVariable(hI->dataType(), TYPE_MSG_INPUT_ "#hI"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
        req.expectEq(makeInfoVariable(hI->ordering(), ORDERING_MSG_INPUT_ "#hI"), 'c') &&
  }
  if (cI) {
    req.expectEq(makeInfoVariable(cI->dataType(), TYPE_MSG_INPUT_ "#cI"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
        req.expectEq(makeInfoVariable(cI->ordering(), ORDERING_MSG_INPUT_ "#cI"), 'c') &&
  }
  if (h) {
    req.expectEq(makeInfoVariable(h->dataType(), TYPE_MSG_OUTPUT_ "#h"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
        req.expectEq(makeInfoVariable(h->ordering(), ORDERING_MSG_OUTPUT_ "#h"), 'c') &&
  }
  if (hL) {
    req.expectEq(makeInfoVariable(hL->dataType(), TYPE_MSG_OUTPUT_ "#hL"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
        req.expectEq(makeInfoVariable(hL->ordering(), ORDERING_MSG_OUTPUT_ "#hL"), 'c') &&
  }
  if (cL) {
    req.expectEq(makeInfoVariable(cL->dataType(), TYPE_MSG_OUTPUT_ "#cL"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
        req.expectEq(makeInfoVariable(cL->ordering(), ORDERING_MSG_OUTPUT_ "#cL"), 'c') &&
  }
  req.logTheSuccess();
  return req;
}

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