chore: import upstream snapshot with attribution

libnd4j/include/loops/cuda/reduce/reduce_float.chpp

@@ -0,0 +1,502 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//
+#include <execution/LaunchContext.h>
+#include <exceptions/cuda_exception.h>
+#include <system/op_boilerplate.h>
+#include <loops/reduce_float.h>
+#include <loops/scalar.h>
+#include <loops/legacy_ops.h>
+#include <helpers/DebugHelper.h>
+#include <types/types.h>
+#include <ops/specials_cuda.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+    using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+SD_KERNEL void simpleReduce(
+    const void* x,
+    const sd::LongType* outerXTadShapeInfo,
+    const sd::LongType* innerXTadShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    const sd::LongType* zShapeInfo) {
+
+  functions::reduce::ReduceFloatFunction<X,Z>::template transformCuda<OpType>(
+      x,
+      outerXTadShapeInfo,
+      innerXTadShapeInfo,
+      extraParams,
+      vreductionBuffer,
+      z,
+      zShapeInfo);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z, typename OpType>
+SD_KERNEL void simpleScalar(
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    sd::LongType* dimension,
+    long long int dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  functions::reduce::ReduceFloatFunction<X, Z>::template execScalarCuda<OpType>(
+      x,
+      xShapeInfo,
+      extraParams,
+      z,
+      zShapeInfo,
+      reductionBuffer,
+      tadOnlyShapeInfo);
+}
+
+namespace functions {
+namespace reduce {
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void ReduceFloatFunction<X,Z>::aggregatePartials(
+    void* vsPartials,
+    sd::LongType tid,
+    sd::LongType numItems,
+    void* vextraParams) {
+
+  using Y             = typename OpType::InterType;
+  auto sPartials      = reinterpret_cast<Y*>(vsPartials);
+  auto extraParams    = reinterpret_cast<Z*>(vextraParams);
+
+  sd::LongType floorPow2 = numItems;
+
+  if (floorPow2 & (floorPow2 - 1)) {
+    while (floorPow2 & (floorPow2 - 1)) {
+      floorPow2 &= (floorPow2 - 1);
+    }
+    if (tid >= floorPow2) {
+      sPartials[tid - floorPow2] =
+          OpType::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
+    }
+    __syncthreads();
+  }
+
+  for (sd::LongType activeThreads = (floorPow2 >> 1); activeThreads; activeThreads >>= 1) {
+    if (tid < activeThreads && (tid + activeThreads) < numItems) {
+      sPartials[tid] =
+          OpType::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
+    }
+    __syncthreads();
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void ReduceFloatFunction<X,Z>::transformCuda(
+    const void* vx,
+    const sd::LongType* outerXTadShapeInfo,
+    const sd::LongType* innerXTadShapeInfo,
+    void* vextraParams,
+    void* vreductionBuffer,
+    void* vz,
+    const sd::LongType* zShapeInfo) {
+
+  auto x          = reinterpret_cast<const X*>(vx);
+  auto z          = reinterpret_cast<Z*>(vz);
+  auto extraParams= reinterpret_cast<Z*>(vextraParams);
+
+  __shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
+  __shared__ sd::LongType tadLen;
+  __shared__ sd::LongType numTads;
+  __shared__ bool sameOffsets;
+
+  // Cache ranks/shape/stride for outer and z
+  __shared__ sd::LongType outerRank;
+  __shared__ const sd::LongType* outerStridePtr;
+  __shared__ const sd::LongType* outerShapePtr;
+
+  __shared__ sd::LongType zRank;
+  __shared__ const sd::LongType* zStridePtr;
+  __shared__ const sd::LongType* zShapePtr;
+
+  // Cache ranks/shape/stride for inner as well if needed
+  __shared__ sd::LongType innerRank;
+  __shared__ const sd::LongType* innerStridePtr;
+  __shared__ const sd::LongType* innerShapePtr;
+
+  if (threadIdx.x == 0) {
+    outerRank      = shape::rank(outerXTadShapeInfo);
+    outerShapePtr  = shape::shapeOf(outerXTadShapeInfo);
+    outerStridePtr = shape::stride(outerXTadShapeInfo);
+
+    zRank          = shape::rank(zShapeInfo);
+    zShapePtr      = shape::shapeOf(zShapeInfo);
+    zStridePtr     = shape::stride(zShapeInfo);
+
+    innerRank      = shape::rank(innerXTadShapeInfo);
+    innerShapePtr  = shape::shapeOf(innerXTadShapeInfo);
+    innerStridePtr = shape::stride(innerXTadShapeInfo);
+
+    sameOffsets    = shape::haveSameShapeAndStrides(zShapeInfo, outerXTadShapeInfo);
+    tadLen         = shape::length(innerXTadShapeInfo);
+    numTads        = shape::length(outerXTadShapeInfo);
+  }
+  __syncthreads();
+
+  sd::LongType coords[SD_MAX_RANK];
+  sd::LongType zCoords[SD_MAX_RANK];
+  for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
+    INDEX2COORDS(r, outerRank, outerShapePtr, coords);
+
+    sd::LongType outerOffset;
+    COORDS2INDEX(outerRank, outerStridePtr, coords, outerOffset);
+
+    sd::LongType zOffset;
+    if (sameOffsets) {
+      zOffset = outerOffset;
+    } else {
+      INDEX2COORDS(r, zRank, zShapePtr, zCoords);
+      COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
+    }
+
+    const X* xTad = x + outerOffset;
+    sPartials[threadIdx.x] = OpType::startingValue(xTad);
+
+    // For the inner dimension
+    for (sd::LongType i = threadIdx.x; i < tadLen; i += blockDim.x) {
+      sd::LongType iCoords[SD_MAX_RANK];
+      sd::LongType innerOffset;
+
+      INDEX2COORDS(i, innerRank, innerShapePtr, iCoords);
+      COORDS2INDEX(innerRank, innerStridePtr, iCoords, innerOffset);
+
+      sPartials[threadIdx.x] = OpType::update(
+          sPartials[threadIdx.x],
+          OpType::op(xTad[innerOffset], extraParams),
+          extraParams);
+    }
+    __syncthreads();
+
+    aggregatePartials<OpType>(
+        sPartials,
+        threadIdx.x,
+        sd::math::sd_min<int>(blockDim.x, tadLen),
+        extraParams);
+    __syncthreads();
+
+    if (threadIdx.x == 0) {
+      z[zOffset] =
+          OpType::postProcess(sPartials[0], tadLen, extraParams);
+    }
+    __syncthreads();
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void ReduceFloatFunction<X, Z>::execScalarCuda(
+    const void* vx,
+    const sd::LongType* xShapeInfo,
+    void* vextraParams,
+    void* vz,
+    const sd::LongType* zShapeInfo,
+    void* vreductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  auto x           = reinterpret_cast<const X*>(vx);
+  auto z           = reinterpret_cast<Z*>(vz);
+  auto extraParams = reinterpret_cast<Z*>(vextraParams);
+  auto reductionBuffer = reinterpret_cast<Z*>(vreductionBuffer);
+
+  using Y = typename OpType::InterType;
+
+  __shared__ Y sPartials[SD_CUDA_BLOCK_SIZE];
+  __shared__ sd::LongType length;
+
+  // Cache rank/shape/stride
+  __shared__ sd::LongType xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  int tid = blockDim.x * blockIdx.x + threadIdx.x;
+  if (threadIdx.x == 0) {
+    length     = shape::length(xShapeInfo);
+    xRank      = shape::rank(xShapeInfo);
+    xShapePtr  = shape::shapeOf(xShapeInfo);
+    xStridePtr = shape::stride(xShapeInfo);
+  }
+  __syncthreads();
+
+  sPartials[threadIdx.x] = OpType::startingValue(x);
+
+  sd::LongType gridSize = gridDim.x * blockDim.x;
+  for (sd::LongType i = tid; i < length; i += gridSize) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+
+    INDEX2COORDS(i, xRank, xShapePtr, xCoords);
+    COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
+
+    if(xOffset < length) {
+      sPartials[threadIdx.x] = OpType::update(
+          sPartials[threadIdx.x],
+          OpType::op(x[xOffset], extraParams),
+          extraParams);
+    }
+
+  }
+  __syncthreads();
+
+  aggregatePartials<OpType>(
+      sPartials,
+      threadIdx.x,
+      sd::math::sd_min<int>(blockDim.x, length),
+      extraParams);
+  __syncthreads();
+
+  if (gridDim.x > 1) {
+    auto tc = reinterpret_cast<unsigned int*>(reductionBuffer);
+    __shared__ bool amLast;
+
+    if (threadIdx.x == 0) {
+      reductionBuffer[blockIdx.x] = sPartials[0];
+    }
+    __threadfence();
+    __syncthreads();
+
+    if (threadIdx.x == 0) {
+      unsigned int ticket = atomicInc(&tc[16384], gridDim.x);
+      amLast              = (ticket == (gridDim.x - 1));
+    }
+    __syncthreads();
+
+    if (amLast) {
+      tc[16384]           = 0;
+      sPartials[threadIdx.x] = OpType::startingValue(x);
+
+      for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
+        sPartials[threadIdx.x] =
+            OpType::update(sPartials[threadIdx.x],
+                           reductionBuffer[i],
+                           extraParams);
+      }
+      __syncthreads();
+
+      aggregatePartials<OpType>(
+          sPartials,
+          threadIdx.x,
+          sd::math::sd_min<int>(gridDim.x, blockDim.x),
+          extraParams);
+      __syncthreads();
+
+      if (threadIdx.x == 0) {
+        z[0] = OpType::postProcess(sPartials[0], length, extraParams);
+      }
+    }
+  } else {
+    if (threadIdx.x == 0) {
+      auto tc   = reinterpret_cast<unsigned int*>(reductionBuffer);
+      tc[16384] = 0;
+      z[0]      = OpType::postProcess(sPartials[0], length, extraParams);
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template<typename OpType>
+SD_HOST void ReduceFloatFunction<X,Z>::intermediate(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const void* x,
+    const sd::LongType* dXShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    const sd::LongType* dZShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    const sd::LongType* dims) {
+
+  if (shape::isEmptyConst(hXShapeInfo)) {
+    const auto startingVal = std::is_same<OpType, simdOps::Mean<X,Z>>::value
+                                 ? sd::DataTypeUtils::nanOrZero<Z>()
+                                 : static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));
+
+    auto res = cudaMemcpyAsync(
+        sd::LaunchContext::defaultContext()->getScalarPointer(),
+        &startingVal,
+        sizeof(Z),
+        cudaMemcpyHostToDevice,
+        *stream);
+    if (res != 0) {
+      throw sd::cuda_exception::build(
+          "ReduceFloatFunction<X,Z>::intermediate: failed to copy temporary scalar", res);
+    }
+
+    auto ptr = sd::LaunchContext::defaultContext()->getScalarPointer();
+    // scalar assign
+    functions::scalar::ScalarTransform<Z, Z, Z>::executeCudaShaped(
+        launchDims,
+        stream,
+        14,
+        z,
+        dZShapeInfo,
+        hZShapeInfo,
+        z,
+        dZShapeInfo,
+        hZShapeInfo,
+        ptr,
+        nullptr);
+  } else {
+    const int zRank   = shape::rank(hZShapeInfo);
+    const int tadRank = shape::rank(hXShapeInfo) - zRank;
+
+    auto outerPack = sd::ConstantShapeHelper::getInstance().createSubArrShapeInfo(
+        const_cast<sd::LongType*>(hXShapeInfo), const_cast<sd::LongType*>(dims), zRank);
+    auto innerPack = sd::ConstantShapeHelper::getInstance().createSubArrShapeInfo(
+        const_cast<sd::LongType*>(hXShapeInfo), const_cast<sd::LongType*>(dims + zRank), tadRank);
+
+    simpleReduce<X, Z, OpType>
+        <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+            x,
+            reinterpret_cast<const sd::LongType*>(outerPack->special()),
+            reinterpret_cast<const sd::LongType*>(innerPack->special()),
+            extraParams,
+            vreductionBuffer,
+            z,
+            dZShapeInfo);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template<typename OpType>
+SD_HOST void ReduceFloatFunction<X,Z>::intermediateScalar(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* dZShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    sd::LongType* dimension,
+    sd::LongType dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  if (shape::isEmptyConst(hXShapeInfo)) {
+    const auto startingVal = std::is_same<OpType, simdOps::Mean<X,Z>>::value
+                                 ? sd::DataTypeUtils::nanOrZero<Z>()
+                                 : static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));
+
+    auto res = cudaMemcpyAsync(z, &startingVal, sizeof(Z), cudaMemcpyHostToDevice, *stream);
+    if (res != 0) {
+      throw sd::cuda_exception::build(
+          "ReduceFloatFunction<X,Z>::intermediateScalar: failed to copy resulting scalar", res);
+    }
+  } else {
+    simpleScalar<X, Z, OpType><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+        x,
+        xShapeInfo,
+        extraParams,
+        z,
+        dZShapeInfo,
+        dimension,
+        dimensionLength,
+        reductionBuffer,
+        tadOnlyShapeInfo);
+  }
+  sd::DebugHelper::checkErrorCode(stream, "ReduceFloatFunction intermediateScalar(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_HOST void ReduceFloatFunction<X,Y>::execReduceScalar(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const int opNum,
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* z,
+    const sd::LongType* dZShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    sd::LongType* dimension,
+    long long int dimensionLength,
+    void* reductionBuffer,
+    const sd::LongType* tadOnlyShapeInfo) {
+
+  DISPATCH_BY_OPNUM_TT(
+      intermediateScalar,
+      PARAMS(
+          launchDims, stream, x, xShapeInfo, hXShapeInfo, extraParams, z, dZShapeInfo, hZShapeInfo, dimension,
+          dimensionLength, reductionBuffer, tadOnlyShapeInfo),
+      OPS_A(REDUCE_FLOAT_OPS));
+  sd::DebugHelper::checkErrorCode(stream, "execReduceScalarFloat(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y>
+SD_HOST void ReduceFloatFunction<X,Y>::execReduce(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const int opNum,
+    const void* x,
+    const sd::LongType* dXShapeInfo,
+    const sd::LongType* hXShapeInfo,
+    void* extraParams,
+    void* vreductionBuffer,
+    void* z,
+    const sd::LongType* dZShapeInfo,
+    const sd::LongType* hZShapeInfo,
+    const sd::LongType* dims) {
+
+  if (shape::length(hZShapeInfo) == 1) {
+    ReduceFloatFunction<X,Y>::execReduceScalar(
+        launchDims, stream, opNum, x, dXShapeInfo, hXShapeInfo,
+        extraParams, z, dZShapeInfo, hZShapeInfo,
+        nullptr, 0, vreductionBuffer, nullptr);
+  } else {
+    DISPATCH_BY_OPNUM_TT(
+        intermediate,
+        PARAMS(
+            launchDims, stream, x, dXShapeInfo, hXShapeInfo, extraParams, vreductionBuffer, z, dZShapeInfo,
+            hZShapeInfo, dims),
+        OPS_A(REDUCE_FLOAT_OPS));
+  }
+  sd::DebugHelper::checkErrorCode(stream, "ReduceFloatFunction execReduce(...) failed");
+}
+
+}  // namespace reduce
+}  // namespace functions