/* ******************************************************************************
 *
 *
 * 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 raver119@gmail.com
//  @author Yurii Shyrma (iuriish@yahoo.com)

#ifndef PAIRWISE_CU
#define PAIRWISE_CU
#include <loops/pairwise_instantiations.h>

#include "../pairwise_transform.h"
using namespace simdOps;

////////////////////////////////////////////////////////////////////////////////
template <typename X, typename Y, typename Z, typename OpType>
SD_KERNEL static void pairwiseSimpleShaped(void const* vx, sd::LongType const* xShapeInfo,
                                           void const* vy, sd::LongType const* yShapeInfo,
                                           void* vz, sd::LongType const* zShapeInfo,
                                           void* vextraParams) {
  auto x          = static_cast<X const*>(vx);
  auto y          = static_cast<Y const*>(vy);
  auto z          = static_cast<Z*>(vz);
  auto extraParams= static_cast<Z*>(vextraParams);

  const int tid = blockIdx.x * blockDim.x + threadIdx.x;

  __shared__ sd::LongType len;
  __shared__ sd::LongType xRank;
  __shared__ sd::LongType yRank;
  __shared__ sd::LongType zRank;

  __shared__ const sd::LongType* xShape;
  __shared__ const sd::LongType* xStride;
  __shared__ const sd::LongType* yShape;
  __shared__ const sd::LongType* yStride;
  __shared__ const sd::LongType* zShape;
  __shared__ const sd::LongType* zStride;

  if (threadIdx.x == 0) {
    len     = shape::length(xShapeInfo);

    xRank   = shape::rank(xShapeInfo);
    xShape  = shape::shapeOf(xShapeInfo);
    xStride = shape::stride(xShapeInfo);

    yRank   = shape::rank(yShapeInfo);
    yShape  = shape::shapeOf(yShapeInfo);
    yStride = shape::stride(yShapeInfo);

    zRank   = shape::rank(zShapeInfo);
    zShape  = shape::shapeOf(zShapeInfo);
    zStride = shape::stride(zShapeInfo);
  }
  __syncthreads();

  for (sd::LongType i = tid; i < len; i += (gridDim.x * blockDim.x)) {
    sd::LongType xCoords[SD_MAX_RANK];
    sd::LongType yCoords[SD_MAX_RANK];
    sd::LongType zCoords[SD_MAX_RANK];
    sd::LongType xOffset;
    sd::LongType yOffset;
    sd::LongType zOffset;

    INDEX2COORDS(i, xRank, xShape, xCoords);
    COORDS2INDEX(xRank, xStride, xCoords, xOffset);

    INDEX2COORDS(i, yRank, yShape, yCoords);
    COORDS2INDEX(yRank, yStride, yCoords, yOffset);

    INDEX2COORDS(i, zRank, zShape, zCoords);
    COORDS2INDEX(zRank, zStride, zCoords, zOffset);

    z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
  }
}

namespace functions {
namespace pairwise_transforms {

////////////////////////////////////////////////////////////////////////////////
template <typename X, typename Y, typename Z>
template <typename OpType>
void SD_HOST PairWiseTransform<X, Y, Z>::intermediateShaped(dim3& launchDims,
                                                            cudaStream_t* stream,
                                                            void const* vx,
                                                            sd::LongType const* xShapeInfo,
                                                            void const* vy,
                                                            sd::LongType const* yShapeInfo,
                                                            void* vz,
                                                            sd::LongType const* zShapeInfo,
                                                            void* vextraParams) {
  pairwiseSimpleShaped<X, Y, Z, OpType>
  <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
      vx, xShapeInfo,
      vy, yShapeInfo,
      vz, zShapeInfo,
      vextraParams);
  sd::DebugHelper::checkErrorCode(stream, "PairWiseTransform intermediateShaped(...) failed");
}

////////////////////////////////////////////////////////////////////////////////
template <typename X, typename Y, typename Z>
void SD_HOST PairWiseTransform<X, Y, Z>::executeCudaShaped(dim3& launchDims,
                                                           cudaStream_t* stream,
                                                           int opNum,
                                                           void const* vx,
                                                           sd::LongType const* xShapeInfo,
                                                           void const* vy,
                                                           sd::LongType const* yShapeInfo,
                                                           void* vz,
                                                           sd::LongType const* zShapeInfo,
                                                           void* vextraParams) {
  DISPATCH_BY_OPNUM_TTT(
      intermediateShaped,
      PARAMS(launchDims, stream, vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, vextraParams),
      PAIRWISE_TRANSFORM_OPS);
}

}  // namespace pairwise_transforms
}  // namespace functions

#endif  // PAIRWISE_CU