wehub-resource-sync
1377 lines
60 KiB
C++
1377 lines
60 KiB
C++
/* ******************************************************************************
|
|
*
|
|
*
|
|
* 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), created on 14.03.2019
|
|
//
|
|
|
|
#ifndef LIBND4J_LOOPS_H
|
|
#define LIBND4J_LOOPS_H
|
|
#include <array/DataTypeUtils.h>
|
|
#include <execution/Threads.h>
|
|
#include <helpers/ConstantTadHelper.h>
|
|
#include <helpers/LoopKind.h>
|
|
#include <helpers/OmpLaunchHelper.h>
|
|
#include <helpers/shape.h>
|
|
#include <loops/indexreduce.h>
|
|
#include <ops/ops.h>
|
|
|
|
|
|
#include <functional>
|
|
|
|
namespace sd {
|
|
|
|
template <typename X, typename Z, typename E>
|
|
class ReductionLoops {
|
|
protected:
|
|
public:
|
|
template <typename OpType>
|
|
static SD_INLINE void loopReduce(memory::Workspace* workspace, const X* x, const LongType* xShapeInfo, Z* z,
|
|
const LongType* zShapeInfo, const LongType* dims, E* extraParams);
|
|
};
|
|
|
|
template <typename X, typename Z>
|
|
class SD_LIB_HIDDEN ReductionFloatLoops : public ReductionLoops<X, Z, Z> {
|
|
public:
|
|
static void wrapper(int opNum, memory::Workspace* workspace, const X* x, const LongType* xShapeInfo, Z* z,
|
|
const LongType* zShapeInfo, const LongType* dims, Z* extraParams);
|
|
|
|
template <typename OpType>
|
|
static void innerloopReduce(memory::Workspace* workspace, const X* x, const LongType* xShapeInfo, Z* z,
|
|
const LongType* zShapeInfo, const LongType* dims, Z* extraParams);
|
|
};
|
|
|
|
template <typename X, typename Z>
|
|
class SD_LIB_HIDDEN ReductionBoolLoops : public ReductionLoops<X, Z, X> {
|
|
public:
|
|
static void wrapper(int opNum, memory::Workspace* workspace, const X* x, const LongType* xShapeInfo, Z* z,
|
|
const LongType* zShapeInfo, const LongType* dims, X* extraParams);
|
|
|
|
template <typename OpType>
|
|
static void innerloopReduce(memory::Workspace* workspace, const X* x, const LongType* xShapeInfo, Z* z,
|
|
const LongType* zShapeInfo, const LongType* dims, X* extraParams);
|
|
};
|
|
|
|
template <typename X, typename Z>
|
|
class SD_LIB_HIDDEN ReductionLongLoops : public ReductionLoops<X, Z, X> {
|
|
public:
|
|
static void wrapper(int opNum, memory::Workspace* workspace, const X* x, const LongType* xShapeInfo, Z* z,
|
|
const LongType* zShapeInfo, const LongType* dims, X* extraParams);
|
|
|
|
template <typename OpType>
|
|
static void innerloopReduce(memory::Workspace* workspace, const X* x, const LongType* xShapeInfo, Z* z,
|
|
const LongType* zShapeInfo, const LongType* dims, X* extraParams);
|
|
};
|
|
|
|
template <typename X>
|
|
class ReductionSameLoops : public ReductionLoops<X, X, X> {
|
|
public:
|
|
static void wrapper(int opNum, memory::Workspace* workspace, const X* x, const LongType* xShapeInfo, X* z,
|
|
const LongType* zShapeInfo, const LongType* dims, X* extraParams);
|
|
|
|
template <typename OpType>
|
|
static void innerloopReduce(memory::Workspace* workspace, const X* x, const LongType* xShapeInfo, X* z,
|
|
const LongType* zShapeInfo, const LongType* dims, X* extraParams);
|
|
};
|
|
|
|
template <typename X, typename Z>
|
|
class IndexReductionLoops {
|
|
public:
|
|
static void wrapIndexReduce(int opNum, const void* x, const LongType* xShapeInfo, void* z,
|
|
const LongType* zShapeInfo, const LongType* tadShapeInfo,
|
|
const LongType* tadOffsets, void* extraParams);
|
|
|
|
template <typename OpType>
|
|
static void loopIndexReduce(X* x, const LongType* xShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* tadShapeInfo, const LongType* tadOffsets, void* extraParams);
|
|
};
|
|
|
|
template <typename X, typename Z, typename E>
|
|
class SD_LIB_HIDDEN TransformLoops {
|
|
public:
|
|
template <typename OpType>
|
|
static SD_INLINE void loopTransform(const X* x, const LongType* xShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
E* extraParams, LongType threadId, LongType numThreads);
|
|
};
|
|
|
|
template <typename X, typename Z>
|
|
class Reduction3Loops {
|
|
public:
|
|
template <typename OpType>
|
|
static SD_INLINE void loopReduce3(const X* x, const LongType* xShapeInfo, const X* y,
|
|
const LongType* yShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
LongType* dims,
|
|
int dimsLen, Z* extraParams, int64_t start, int64_t stop);
|
|
|
|
template <typename OpType>
|
|
static SD_INLINE void loopReduce3All(const X* x, const LongType* xShapeInfo, const X* y,
|
|
const LongType* yShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* xTadShapeInfo, const LongType* xTadOffsets,
|
|
const LongType* yTadShapeInfo, const LongType* yTadOffsets,
|
|
Z* extraParams, int64_t start, int64_t stop);
|
|
|
|
static void wrapper(int opNum, const X* x, const LongType* xShapeInfo, const X* y, const LongType* yShapeInfo,
|
|
Z* z, const LongType* zShapeInfo, LongType* dims, int dimsLen, Z* extraParams, int64_t start,
|
|
int64_t stop);
|
|
|
|
static void wrapperAll(int opNum, const X* x, const LongType* xShapeInfo, const X* y,
|
|
const LongType* yShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* xTadShapeInfo, const LongType* xTadOffsets,
|
|
const LongType* yTadShapeInfo, const LongType* yTadOffsets, Z* extraParams,
|
|
int64_t start, int64_t stop);
|
|
|
|
template <typename OpType>
|
|
static void innerloopReduce3(const X* x, const LongType* xShapeInfo, const X* y, const LongType* yShapeInfo,
|
|
Z* z, const LongType* zShapeInfo, LongType* dims, int dimsLen, Z* extraParams,
|
|
int64_t start, int64_t stop);
|
|
|
|
template <typename OpType>
|
|
static void innerloopReduce3All(const X* x, const LongType* xShapeInfo, const X* y,
|
|
const LongType* yShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* xTadShapeInfo, const LongType* xTadOffsets,
|
|
const LongType* yTadShapeInfo, const LongType* yTadOffsets, Z* extraParams,
|
|
int64_t start, int64_t stop);
|
|
};
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z, typename E, typename OpType>
|
|
static void reduceExec21(const X* x, const LongType* xShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* dims, E* extraParams) {
|
|
const LongType xAxis0 = shape::sizeAt(xShapeInfo, dims[0]);
|
|
const LongType xStrd0 = shape::strideAt(xShapeInfo, dims[0]);
|
|
const LongType zStrd0 = shape::strideAt(zShapeInfo, static_cast<LongType>(0));
|
|
|
|
const LongType xAxis1 = shape::sizeAt(xShapeInfo, dims[1]);
|
|
const LongType xStrd1 = shape::strideAt(xShapeInfo, dims[1]);
|
|
|
|
// Convert parameters to Z* for OpType compatibility
|
|
Z* compatibleExtraParams = nullptr;
|
|
Z convertedParams[8];
|
|
|
|
if (extraParams != nullptr) {
|
|
if constexpr (std::is_same_v<E, Z>) {
|
|
compatibleExtraParams = reinterpret_cast<Z*>(extraParams);
|
|
} else {
|
|
// Convert parameters to Z type
|
|
for (int i = 0; i < 8; ++i) {
|
|
convertedParams[i] = static_cast<Z>(extraParams[i]);
|
|
}
|
|
compatibleExtraParams = convertedParams;
|
|
}
|
|
}
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto i0 = start; i0 < stop; ++i0) {
|
|
auto x0 = x + i0 * xStrd0;
|
|
auto z0 = z + i0 * zStrd0;
|
|
|
|
auto s = static_cast<typename OpType::InterType>(OpType::startingValue(x0));
|
|
|
|
if (xStrd1 == 1)
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1) {
|
|
s = OpType::update(s, OpType::op(x0[i1], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1) {
|
|
s = OpType::update(s, OpType::op(x0[i1 * xStrd1], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
|
|
*z0 = OpType::postProcess(s, static_cast<LongType>(xAxis1), compatibleExtraParams);
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, xAxis0);
|
|
}
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z, typename E, typename OpType>
|
|
static void reduceExec31(const X* x, const LongType* xShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* dims, E* extraParams) {
|
|
const LongType xAxis0 = shape::sizeAt(xShapeInfo, dims[0]);
|
|
const LongType xStrd0 = shape::strideAt(xShapeInfo, dims[0]);
|
|
const LongType zStrd0 = shape::strideAt(zShapeInfo, static_cast<LongType>(0));
|
|
|
|
const LongType xAxis1 = shape::sizeAt(xShapeInfo, dims[1]);
|
|
const LongType xStrd1 = shape::strideAt(xShapeInfo, dims[1]);
|
|
|
|
const LongType xAxis2 = shape::sizeAt(xShapeInfo, dims[2]);
|
|
const LongType xStrd2 = shape::strideAt(xShapeInfo, dims[2]);
|
|
|
|
const LongType tadLen = static_cast<LongType>(xAxis1 * xAxis2);
|
|
|
|
Z* compatibleExtraParams = nullptr;
|
|
Z convertedParams[8];
|
|
|
|
if (extraParams != nullptr) {
|
|
if constexpr (std::is_same_v<E, Z>) {
|
|
compatibleExtraParams = reinterpret_cast<Z*>(extraParams);
|
|
} else {
|
|
// Convert parameters to Z type
|
|
for (int i = 0; i < 8; ++i) {
|
|
convertedParams[i] = static_cast<Z>(extraParams[i]);
|
|
}
|
|
compatibleExtraParams = convertedParams;
|
|
}
|
|
}
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto i0 = start; i0 < stop; ++i0) {
|
|
auto x0 = x + i0 * xStrd0;
|
|
auto z0 = z + i0 * zStrd0;
|
|
|
|
auto s = static_cast<typename OpType::InterType>(OpType::startingValue(x0));
|
|
|
|
if (xStrd1 == 1)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2)
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2 is %lld,%lld,%lld reduceExec31\n", i0,i1,i2);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x0[i1 + i2 * xStrd2], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else if (xStrd2 == 1)
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2 is %lld,%lld,%lld offset is %lld reduceExec31\n", i0,i1,i2,i1 * xStrd1 + i2);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x0[i1 * xStrd1 + i2], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2 is %lld,%lld,%lld offset is %lld reduceExec31\n", i0,i1,i2,i1 * xStrd1 + i2 * xStrd2);
|
|
#endif
|
|
|
|
s = OpType::update(s, OpType::op(x0[i1 * xStrd1 + i2 * xStrd2], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
|
|
*z0 = OpType::postProcess(s, tadLen, compatibleExtraParams);
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, xAxis0);
|
|
}
|
|
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z, typename E, typename OpType>
|
|
SD_LIB_HIDDEN void reduceExec32(const X* x, const LongType* xShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* dims, E* extraParams) {
|
|
|
|
const LongType xAxis0 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[0] : dims[1]);
|
|
const LongType xStrd0 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[0] : dims[1]);
|
|
const LongType zStrd0 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(0) : static_cast<LongType>(1));
|
|
|
|
const LongType xAxis1 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[1] : dims[0]);
|
|
const LongType xStrd1 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[1] : dims[0]);
|
|
const LongType zStrd1 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(1) : static_cast<LongType>(0));
|
|
|
|
const LongType xAxis2 = shape::sizeAt(xShapeInfo, dims[2]);
|
|
const LongType xStrd2 = shape::strideAt(xShapeInfo, dims[2]);
|
|
|
|
|
|
Z convertedParams[8] = {0};
|
|
Z* compatibleExtraParams = nullptr;
|
|
|
|
if (extraParams != nullptr) {
|
|
if constexpr (std::is_same_v<E, Z>) {
|
|
compatibleExtraParams = reinterpret_cast<Z*>(extraParams);
|
|
} else {
|
|
for (int i = 0; i < 8; ++i) {
|
|
convertedParams[i] = static_cast<Z>(extraParams[i]);
|
|
}
|
|
compatibleExtraParams = convertedParams;
|
|
}
|
|
}
|
|
|
|
auto func = PRAGMA_THREADS_FOR_2D {
|
|
for (auto i0 = start_x; i0 < stop_x; ++i0) {
|
|
for (auto i1 = start_y; i1 < stop_y; ++i1) {
|
|
auto x1 = x + i0 * xStrd0 + i1 * xStrd1;
|
|
auto z1 = z + i0 * zStrd0 + i1 * zStrd1;
|
|
|
|
auto s = static_cast<typename OpType::InterType>(OpType::startingValue(x1));
|
|
|
|
if (xStrd2 == 1)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2 is %lld,%lld,%lld reduceExec32\n", i0,i1,i2);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x1[i2], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2 is %lld,%lld,%lld reduceExec32\n", i0,i1,i2);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x1[i2 * xStrd2], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
*z1 = OpType::postProcess(s, static_cast<LongType>(xAxis2), compatibleExtraParams);
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, xAxis0, 1, 0, xAxis1, 1);
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z, typename E, typename OpType>
|
|
SD_LIB_HIDDEN void reduceExec41(const X* x,
|
|
const LongType* xShapeInfo,
|
|
Z* z, const
|
|
LongType* zShapeInfo,
|
|
const LongType* dims,
|
|
E* extraParams) {
|
|
LongType xRank = shape::rank(xShapeInfo);
|
|
LongType zRank = shape::rank(zShapeInfo);
|
|
|
|
const LongType xAxis0 = shape::sizeAt(xShapeInfo, dims[0]);
|
|
const LongType xStrd0 = shape::strideAt(xShapeInfo, dims[0]);
|
|
const LongType zStrd0 = shape::strideAt(zShapeInfo, static_cast<LongType>(0));
|
|
|
|
const LongType xAxis1 = shape::sizeAt(xShapeInfo, dims[1]);
|
|
const LongType xStrd1 = shape::strideAt(xShapeInfo, dims[1]);
|
|
|
|
const LongType xAxis2 = shape::sizeAt(xShapeInfo, dims[2]);
|
|
const LongType xStrd2 = shape::strideAt(xShapeInfo, dims[2]);
|
|
|
|
const LongType xAxis3 = shape::sizeAt(xShapeInfo, dims[3]);
|
|
const LongType xStrd3 = shape::strideAt(xShapeInfo, dims[3]);
|
|
|
|
const LongType tadLen = static_cast<LongType>(xAxis1 * xAxis2 * xAxis3);
|
|
|
|
Z convertedParams[8] = {0};
|
|
Z* compatibleExtraParams = nullptr;
|
|
|
|
if (extraParams != nullptr) {
|
|
if constexpr (std::is_same_v<E, Z>) {
|
|
compatibleExtraParams = reinterpret_cast<Z*>(extraParams);
|
|
} else {
|
|
for (int i = 0; i < 8; ++i) {
|
|
convertedParams[i] = static_cast<Z>(extraParams[i]);
|
|
}
|
|
compatibleExtraParams = convertedParams;
|
|
}
|
|
}
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto i0 = start; i0 < stop; ++i0) {
|
|
auto x0 = x + i0 * xStrd0;
|
|
auto z0 = z + i0 * zStrd0;
|
|
|
|
auto s = static_cast<typename OpType::InterType>(OpType::startingValue(x0));
|
|
|
|
if (xStrd1 == 1)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2)
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3 is %lld,%lld,%lld,%lld offset is %lld reduceExec41\n", i0,i1,i2,i3,i1 + i2 * xStrd2 + i3 * xStrd3);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x0[i1 + i2 * xStrd2 + i3 * xStrd3], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else if (xStrd2 == 1)
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3 is %lld,%lld,%lld,%lld offset is %lld reduceExec41\n", i0,i1,i2,i3,i1 * xStrd1 + i2 + i3 * xStrd3);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x0[i1 * xStrd1 + i2 + i3 * xStrd3], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
|
|
else if (xStrd3 == 1)
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3 is %lld,%lld,%lld,%lld offset is %lld reduceExec41\n", i0,i1,i2,i3,i1 * xStrd1 + i2 * xStrd2 + i3);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x0[i1 * xStrd1 + i2 * xStrd2 + i3], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3 is %lld,%lld,%lld,%lld offset is %lld reduceExec41\n", i0,i1,i2,i3,i1 * xStrd1 + i2 * xStrd2 + i3 * xStrd3);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x0[i1 * xStrd1 + i2 * xStrd2 + i3 * xStrd3], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
*z0 = OpType::postProcess(s, tadLen, compatibleExtraParams);
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, xAxis0);
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z, typename E, typename OpType>
|
|
SD_LIB_HIDDEN void reduceExec42(const X* x, const LongType* xShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* dims, E* extraParams) {
|
|
const LongType xAxis0 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[0] : dims[1]);
|
|
const LongType xStrd0 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[0] : dims[1]);
|
|
const LongType zStrd0 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(0) : static_cast<LongType>(1));
|
|
|
|
const LongType xAxis1 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[1] : dims[0]);
|
|
const LongType xStrd1 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[1] : dims[0]);
|
|
const LongType zStrd1 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(1) : static_cast<LongType>(0));
|
|
|
|
const LongType xAxis2 = shape::sizeAt(xShapeInfo, dims[2]);
|
|
const LongType xStrd2 = shape::strideAt(xShapeInfo, dims[2]);
|
|
|
|
const LongType xAxis3 = shape::sizeAt(xShapeInfo, dims[3]);
|
|
const LongType xStrd3 = shape::strideAt(xShapeInfo, dims[3]);
|
|
|
|
const LongType tadLen = static_cast<LongType>(xAxis2 * xAxis3);
|
|
|
|
LongType xRank = shape::rank(xShapeInfo);
|
|
|
|
Z convertedParams[8] = {0};
|
|
Z* compatibleExtraParams = nullptr;
|
|
|
|
if (extraParams != nullptr) {
|
|
if constexpr (std::is_same_v<E, Z>) {
|
|
compatibleExtraParams = reinterpret_cast<Z*>(extraParams);
|
|
} else {
|
|
for (int i = 0; i < 8; ++i) {
|
|
convertedParams[i] = static_cast<Z>(extraParams[i]);
|
|
}
|
|
compatibleExtraParams = convertedParams;
|
|
}
|
|
}
|
|
|
|
auto func = PRAGMA_THREADS_FOR_2D {
|
|
for (auto i0 = start_x; i0 < stop_x; ++i0) {
|
|
for (auto i1 = start_y; i1 < stop_y; ++i1) {
|
|
auto x1 = x + i0 * xStrd0 + i1 * xStrd1;
|
|
auto z1 = z + i0 * zStrd0 + i1 * zStrd1;
|
|
|
|
auto s = static_cast<typename OpType::InterType>(OpType::startingValue(x1));
|
|
|
|
if (xStrd2 == 1)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3 is %lld,%lld,%lld,%lld reduceExec42\n", i0,i1,i2,i3);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x1[i2 + i3 * xStrd3], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else if (xStrd3 == 1)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3 is %lld,%lld,%lld,%lld offset %lld reduceExec42\n", i0,i1,i2,i3,i2 * xStrd2 + i3);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x1[i2 * xStrd2 + i3], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3 is %lld,%lld,%lld,%lld offset %lld reduceExec42\n", i0,i1,i2,i3,i2 * xStrd2 + i3 * xStrd3);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x1[i2 * xStrd2 + i3 * xStrd3], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
|
|
*z1 = OpType::postProcess(s, tadLen, compatibleExtraParams);
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, xAxis0, 1, 0, xAxis1, 1);
|
|
}
|
|
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z, typename E, typename OpType>
|
|
SD_LIB_HIDDEN void reduceExec43(const X* x, const LongType* xShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* dims, E* extraParams) {
|
|
const LongType xAxis0 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[0] : dims[2]);
|
|
const LongType xStrd0 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[0] : dims[2]);
|
|
const LongType zStrd0 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(0) : static_cast<LongType>(2));
|
|
|
|
const LongType xAxis1 = shape::sizeAt(xShapeInfo, dims[1]);
|
|
const LongType xStrd1 = shape::strideAt(xShapeInfo, dims[1]);
|
|
const LongType zStrd1 = shape::strideAt(zShapeInfo, static_cast<LongType>(1));
|
|
|
|
const LongType xAxis2 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[2] : dims[0]);
|
|
const LongType xStrd2 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[2] : dims[0]);
|
|
const LongType zStrd2 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(2) : static_cast<LongType>(0));
|
|
|
|
const LongType xAxis3 = shape::sizeAt(xShapeInfo, dims[3]);
|
|
const LongType xStrd3 = shape::strideAt(xShapeInfo, dims[3]);
|
|
LongType xRank = shape::rank(xShapeInfo);
|
|
|
|
Z convertedParams[8] = {0};
|
|
Z* compatibleExtraParams = nullptr;
|
|
|
|
if (extraParams != nullptr) {
|
|
if constexpr (std::is_same_v<E, Z>) {
|
|
compatibleExtraParams = reinterpret_cast<Z*>(extraParams);
|
|
} else {
|
|
for (int i = 0; i < 8; ++i) {
|
|
convertedParams[i] = static_cast<Z>(extraParams[i]);
|
|
}
|
|
compatibleExtraParams = convertedParams;
|
|
}
|
|
}
|
|
|
|
auto func = PRAGMA_THREADS_FOR_3D {
|
|
for (auto i0 = start_x; i0 < stop_x; ++i0) {
|
|
for (auto i1 = start_y; i1 < stop_y; ++i1) {
|
|
for (auto i2 = start_z; i2 < stop_z; ++i2) {
|
|
auto x2 = x + i0 * xStrd0 + i1 * xStrd1 + i2 * xStrd2;
|
|
auto z2 = z + i0 * zStrd0 + i1 * zStrd1 + i2 * zStrd2;
|
|
|
|
auto s = static_cast<typename OpType::InterType>(OpType::startingValue(x2));
|
|
|
|
if (xStrd3 == 1)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3 is %lld,%lld,%lld,%lld reduceExec43\n", i0,i1,i2,i3);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x2[i3], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3 is %lld,%lld,%lld,%lld reduceExec43\n", i0,i1,i2,i3);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x2[i3 * xStrd3], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
|
|
*z2 = OpType::postProcess(s, static_cast<LongType>(xAxis3), compatibleExtraParams);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, xAxis0, 1, 0, xAxis1, 1, 0, xAxis2, 1);
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z, typename E, typename OpType>
|
|
SD_LIB_HIDDEN void reduceExec51(const X* x, const LongType* xShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* dims, E* extraParams) {
|
|
const LongType xAxis0 = shape::sizeAt(xShapeInfo, dims[0]);
|
|
const LongType xStrd0 = shape::strideAt(xShapeInfo, dims[0]);
|
|
const LongType zStrd0 = shape::strideAt(zShapeInfo, static_cast<LongType>(0));
|
|
|
|
const LongType xAxis1 = shape::sizeAt(xShapeInfo, dims[1]);
|
|
const LongType xStrd1 = shape::strideAt(xShapeInfo, dims[1]);
|
|
|
|
const LongType xAxis2 = shape::sizeAt(xShapeInfo, dims[2]);
|
|
const LongType xStrd2 = shape::strideAt(xShapeInfo, dims[2]);
|
|
|
|
const LongType xAxis3 = shape::sizeAt(xShapeInfo, dims[3]);
|
|
const LongType xStrd3 = shape::strideAt(xShapeInfo, dims[3]);
|
|
|
|
const LongType xAxis4 = shape::sizeAt(xShapeInfo, dims[4]);
|
|
const LongType xStrd4 = shape::strideAt(xShapeInfo, dims[4]);
|
|
|
|
const LongType tadLen = static_cast<LongType>(xAxis1 * xAxis2 * xAxis3 * xAxis4);
|
|
|
|
LongType xRank = shape::rank(xShapeInfo);
|
|
|
|
Z convertedParams[8] = {0};
|
|
Z* compatibleExtraParams = nullptr;
|
|
|
|
if (extraParams != nullptr) {
|
|
if constexpr (std::is_same_v<E, Z>) {
|
|
compatibleExtraParams = reinterpret_cast<Z*>(extraParams);
|
|
} else {
|
|
for (int i = 0; i < 8; ++i) {
|
|
convertedParams[i] = static_cast<Z>(extraParams[i]);
|
|
}
|
|
compatibleExtraParams = convertedParams;
|
|
}
|
|
}
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto i0 = start; i0 < stop; ++i0) {
|
|
auto x0 = x + i0 * xStrd0;
|
|
auto z0 = z + i0 * zStrd0;
|
|
|
|
auto s = static_cast<typename OpType::InterType>(OpType::startingValue(x0));
|
|
|
|
if (xStrd1 == 1)
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2)
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3,i4 is %lld,%lld,%lld,%lld,%lld reduceExec51\n", i0,i1,i2,i3,i4);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x0[i1 + i2 * xStrd2 + i3 * xStrd3 + i4 * xStrd4], compatibleExtraParams),
|
|
compatibleExtraParams);
|
|
}
|
|
else if (xStrd2 == 1)
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3)
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3,i4 is %lld,%lld,%lld,%lld,%lld offset %lldreduceExec51\n", i0,i1,i2,i3,i4,i1 * xStrd1 + i2 + i3 * xStrd3 + i4 * xStrd4);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x0[i1 * xStrd1 + i2 + i3 * xStrd3 + i4 * xStrd4], compatibleExtraParams),
|
|
compatibleExtraParams);
|
|
}
|
|
else if (xStrd3 == 1)
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2)
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3) {
|
|
#if defined(PRINT_INDICES)
|
|
printf("Index i0,i1,i2,i3,i4 is %lld,%lld,%lld,%lld,%lld offset %lld reduceExec51\n", i0,i1,i2,i3,i4,i1 * xStrd1 + i2 * xStrd2 + i3 + i4 * xStrd4);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x0[i1 * xStrd1 + i2 * xStrd2 + i3 + i4 * xStrd4], compatibleExtraParams),
|
|
compatibleExtraParams);
|
|
}
|
|
else if (xStrd4 == 1)
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3)
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4) {
|
|
#if defined(PRINT_INDICES)
|
|
printf("Index i0,i1,i2,i3,i4 is %lld,%lld,%lld,%lld,%lld offset %lld reduceExec51\n", i0,i1,i2,i3,i4,i1 * xStrd1 + i2 * xStrd2 + i3 * xStrd3 + i4);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x0[i1 * xStrd1 + i2 * xStrd2 + i3 * xStrd3 + i4], compatibleExtraParams),
|
|
compatibleExtraParams);
|
|
}
|
|
else
|
|
for (LongType i1 = 0; i1 < xAxis1; ++i1)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3)
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4) {
|
|
#if defined(PRINT_INDICES)
|
|
printf("Index i0,i1,i2,i3,i4 is %lld,%lld,%lld,%lld,%lld offset %lld reduceExec51\n", i0,i1,i2,i3,i4,i1 * xStrd1 + i2 * xStrd2 + i3 * xStrd3 + i4 * xStrd4);
|
|
#endif
|
|
s = OpType::update(
|
|
s, OpType::op(x0[i1 * xStrd1 + i2 * xStrd2 + i3 * xStrd3 + i4 * xStrd4], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
*z0 = OpType::postProcess(s, tadLen, compatibleExtraParams);
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, xAxis0);
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z, typename E, typename OpType>
|
|
SD_LIB_HIDDEN void reduceExec52(const X* x, const LongType* xShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* dims, E* extraParams) {
|
|
const LongType xAxis0 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[0] : dims[1]);
|
|
const LongType xStrd0 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[0] : dims[1]);
|
|
const LongType zStrd0 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(0) : static_cast<LongType>(1));
|
|
|
|
const LongType xAxis1 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[1] : dims[0]);
|
|
const LongType xStrd1 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[1] : dims[0]);
|
|
const LongType zStrd1 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(1) : static_cast<LongType>(0));
|
|
|
|
const LongType xAxis2 = shape::sizeAt(xShapeInfo, dims[2]);
|
|
const LongType xStrd2 = shape::strideAt(xShapeInfo, dims[2]);
|
|
|
|
const LongType xAxis3 = shape::sizeAt(xShapeInfo, dims[3]);
|
|
const LongType xStrd3 = shape::strideAt(xShapeInfo, dims[3]);
|
|
|
|
const LongType xAxis4 = shape::sizeAt(xShapeInfo, dims[4]);
|
|
const LongType xStrd4 = shape::strideAt(xShapeInfo, dims[4]);
|
|
|
|
const LongType tadLen = static_cast<LongType>(xAxis2 * xAxis3 * xAxis4);
|
|
|
|
LongType xRank = shape::rank(xShapeInfo);
|
|
|
|
Z convertedParams[8] = {0};
|
|
Z* compatibleExtraParams = nullptr;
|
|
|
|
if (extraParams != nullptr) {
|
|
if constexpr (std::is_same_v<E, Z>) {
|
|
compatibleExtraParams = reinterpret_cast<Z*>(extraParams);
|
|
} else {
|
|
for (int i = 0; i < 8; ++i) {
|
|
convertedParams[i] = static_cast<Z>(extraParams[i]);
|
|
}
|
|
compatibleExtraParams = convertedParams;
|
|
}
|
|
}
|
|
|
|
auto func = PRAGMA_THREADS_FOR_2D {
|
|
for (auto i0 = start_x; i0 < stop_x; ++i0) {
|
|
for (auto i1 = start_y; i1 < stop_y; ++i1) {
|
|
auto x1 = x + i0 * xStrd0 + i1 * xStrd1;
|
|
auto z1 = z + i0 * zStrd0 + i1 * zStrd1;
|
|
|
|
auto s = static_cast<typename OpType::InterType>(OpType::startingValue(x1));
|
|
|
|
if (xStrd2 == 1)
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3,i4 is %lld,%lld,%lld,%lld,%lld offset %lld reduceExec52\n", i0,i1,i2,i3,i4,i2 + i3 * xStrd3 + i4 * xStrd4);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x1[i2 + i3 * xStrd3 + i4 * xStrd4], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else if (xStrd3 == 1)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2)
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3,i4 is %lld,%lld,%lld,%lld,%lld offset %lld reduceExec52\n", i0,i1,i2,i3,i4,i2 + i3 * xStrd3 + i4 * xStrd4);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x1[i2 * xStrd2 + i3 + i4 * xStrd4], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else if (xStrd4 == 1)
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3)
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3,i4 is %lld,%lld,%lld,%lld,%lld offset %lld reduceExec52\n", i0,i1,i2,i3,i4,i2 * xStrd2 + i3 + i4 * xStrd4);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x1[i2 * xStrd2 + i3 * xStrd3 + i4], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else
|
|
for (LongType i2 = 0; i2 < xAxis2; ++i2)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3)
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4) {
|
|
#if defined(PRINT_INDICES)
|
|
printf("Index i0,i1,i2,i3,i4 is %lld,%lld,%lld,%lld,%lld offset %lld reduceExec52\n", i0,i1,i2,i3,i4,i2 * xStrd2 + i3 * xStrd3 + i4 * xStrd4);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x1[i2 * xStrd2 + i3 * xStrd3 + i4 * xStrd4], compatibleExtraParams),
|
|
compatibleExtraParams);
|
|
}
|
|
|
|
*z1 = OpType::postProcess(s, tadLen, compatibleExtraParams);
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, xAxis0, 1, 0, xAxis1, 1);
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z, typename E, typename OpType>
|
|
SD_LIB_HIDDEN void reduceExec53(const X* x, const LongType* xShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* dims, E* extraParams) {
|
|
const LongType xAxis0 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[0] : dims[2]);
|
|
const LongType xStrd0 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[0] : dims[2]);
|
|
const LongType zStrd0 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(0) : static_cast<LongType>(2));
|
|
|
|
const LongType xAxis1 = shape::sizeAt(xShapeInfo, dims[1]);
|
|
const LongType xStrd1 = shape::strideAt(xShapeInfo, dims[1]);
|
|
const LongType zStrd1 = shape::strideAt(zShapeInfo, static_cast<LongType>(1));
|
|
|
|
const LongType xAxis2 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[2] : dims[0]);
|
|
const LongType xStrd2 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[2] : dims[0]);
|
|
const LongType zStrd2 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(2) : static_cast<LongType>(0));
|
|
|
|
const LongType xAxis3 = shape::sizeAt(xShapeInfo, dims[3]);
|
|
const LongType xStrd3 = shape::strideAt(xShapeInfo, dims[3]);
|
|
|
|
const LongType xAxis4 = shape::sizeAt(xShapeInfo, dims[4]);
|
|
const LongType xStrd4 = shape::strideAt(xShapeInfo, dims[4]);
|
|
|
|
const LongType tadLen = static_cast<LongType>(xAxis3 * xAxis4);
|
|
|
|
LongType xRank = shape::rank(xShapeInfo);
|
|
|
|
Z convertedParams[8] = {0};
|
|
Z* compatibleExtraParams = nullptr;
|
|
|
|
if (extraParams != nullptr) {
|
|
if constexpr (std::is_same_v<E, Z>) {
|
|
compatibleExtraParams = reinterpret_cast<Z*>(extraParams);
|
|
} else {
|
|
for (int i = 0; i < 8; ++i) {
|
|
convertedParams[i] = static_cast<Z>(extraParams[i]);
|
|
}
|
|
compatibleExtraParams = convertedParams;
|
|
}
|
|
}
|
|
|
|
auto func = PRAGMA_THREADS_FOR_3D {
|
|
for (auto i0 = start_x; i0 < stop_x; ++i0) {
|
|
for (auto i1 = start_y; i1 < stop_y; ++i1) {
|
|
for (auto i2 = start_z; i2 < stop_z; ++i2) {
|
|
auto x2 = x + i0 * xStrd0 + i1 * xStrd1 + i2 * xStrd2;
|
|
auto z2 = z + i0 * zStrd0 + i1 * zStrd1 + i2 * zStrd2;
|
|
|
|
auto s = static_cast<typename OpType::InterType>(OpType::startingValue(x2));
|
|
|
|
if (xStrd3 == 1)
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3,i4 is %lld,%lld,%lld,%lld,%lld offset %lld reduceExec53\n", i0,i1,i2,i3,i4,i3 + i4 * xStrd4);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x2[i3 + i4 * xStrd4], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else if (xStrd4 == 1)
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3)
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
printf("Index i0,i1,i2,i3,i4 is %lld,%lld,%lld,%lld,%lld offset %lld reduceExec53\n", i0,i1,i2,i3,i4,i3 * xStrd3 + i4);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x2[i3 * xStrd3 + i4], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else
|
|
for (LongType i3 = 0; i3 < xAxis3; ++i3)
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4) {
|
|
#if defined(PRINT_INDICES)
|
|
printf("Index i0,i1,i2,i3,i4 is %lld,%lld,%lld,%lld,%lld offset %lld reduceExec53\n", i0,i1,i2,i3,i4,i3 * xStrd3 + i4 * xStrd4);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x2[i3 * xStrd3 + i4 * xStrd4], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
*z2 = OpType::postProcess(s, tadLen, compatibleExtraParams);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, xAxis0, 1, 0, xAxis1, 1, 0, xAxis2, 1);
|
|
}
|
|
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z, typename E, typename OpType>
|
|
SD_LIB_HIDDEN void reduceExec54(const X* x, const LongType* xShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* dims, E* extraParams) {
|
|
const LongType xAxis0 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[0] : dims[3]);
|
|
const LongType xStrd0 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[0] : dims[3]);
|
|
const LongType zStrd0 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(0) : static_cast<LongType>(3));
|
|
|
|
const LongType xAxis1 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[1] : dims[2]);
|
|
const LongType xStrd1 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[1] : dims[2]);
|
|
const LongType zStrd1 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(1) : static_cast<LongType>(2));
|
|
|
|
const LongType xAxis2 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[2] : dims[1]);
|
|
const LongType xStrd2 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[2] : dims[1]);
|
|
const LongType zStrd2 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(2) : static_cast<LongType>(1));
|
|
|
|
const LongType xAxis3 = shape::sizeAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[3] : dims[0]);
|
|
const LongType xStrd3 = shape::strideAt(xShapeInfo, shape::order(zShapeInfo) == 'c' ? dims[3] : dims[0]);
|
|
const LongType zStrd3 = shape::strideAt(zShapeInfo, shape::order(zShapeInfo) == 'c' ? static_cast<LongType>(3) : static_cast<LongType>(0));
|
|
|
|
const LongType xAxis4 = shape::sizeAt(xShapeInfo, dims[4]);
|
|
const LongType xStrd4 = shape::strideAt(xShapeInfo, dims[4]);
|
|
|
|
LongType xRank = shape::rank(xShapeInfo);
|
|
|
|
Z convertedParams[8] = {0};
|
|
Z* compatibleExtraParams = nullptr;
|
|
|
|
if (extraParams != nullptr) {
|
|
if constexpr (std::is_same_v<E, Z>) {
|
|
compatibleExtraParams = reinterpret_cast<Z*>(extraParams);
|
|
} else {
|
|
for (int i = 0; i < 8; ++i) {
|
|
convertedParams[i] = static_cast<Z>(extraParams[i]);
|
|
}
|
|
compatibleExtraParams = convertedParams;
|
|
}
|
|
}
|
|
|
|
auto func = PRAGMA_THREADS_FOR_3D {
|
|
for (auto i0 = start_x; i0 < stop_x; ++i0) {
|
|
for (auto i1 = start_y; i1 < stop_y; ++i1) {
|
|
for (auto i2 = start_z; i2 < stop_z; ++i2) {
|
|
for (auto i3 = 0; i3 < xAxis3; ++i3) {
|
|
auto x3 = x + i0 * xStrd0 + i1 * xStrd1 + i2 * xStrd2 + i3 * xStrd3;
|
|
auto z3 = z + i0 * zStrd0 + i1 * zStrd1 + i2 * zStrd2 + i3 * zStrd3;
|
|
|
|
auto s = static_cast<typename OpType::InterType>(OpType::startingValue(x3));
|
|
|
|
if (xStrd4 == 1)
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x3[i4], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
else
|
|
for (LongType i4 = 0; i4 < xAxis4; ++i4) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xShapeInfo);
|
|
shape::printShapeInfo(zShapeInfo);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(x3[i4 * xStrd4], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
*z3 = OpType::postProcess(s, static_cast<LongType>(xAxis4), compatibleExtraParams);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, xAxis0, 1, 0, xAxis1, 1, 0, xAxis2, 1);
|
|
}
|
|
|
|
////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z, typename E, typename OpType>
|
|
SD_LIB_HIDDEN void reduceDefault(memory::Workspace* workspace, const X* x, const LongType* xShapeInfo, Z* z,
|
|
const LongType* zShapeInfo, const LongType* dims, E* extraParams) {
|
|
const int zRank = shape::rank(zShapeInfo);
|
|
const int tadRank = shape::rank(xShapeInfo) - zRank;
|
|
|
|
LongType* outerXTadShapeInfo = ShapeBuilders::createSubArrShapeInfo(xShapeInfo, dims, zRank);
|
|
LongType* innerXTadShapeInfo = ShapeBuilders::createSubArrShapeInfo(xShapeInfo, dims + zRank, tadRank);
|
|
|
|
const bool sameOffsets1 = shape::haveSameShapeAndStrides(zShapeInfo, outerXTadShapeInfo);
|
|
const bool sameOffsets2 = shape::haveSameShapeAndStrides(zShapeInfo, innerXTadShapeInfo);
|
|
|
|
const LongType zLen = shape::length(zShapeInfo);
|
|
const LongType tadLen = shape::length(innerXTadShapeInfo);
|
|
|
|
LongType* zOffsets = nullptr;
|
|
ALLOCATE(zOffsets, workspace, zLen, sd::LongType);
|
|
shape::calcOffsets(zShapeInfo, zOffsets);
|
|
|
|
LongType* outerXTadOffsets = zOffsets;
|
|
if (!sameOffsets1) {
|
|
ALLOCATE(outerXTadOffsets, workspace, zLen, sd::LongType);
|
|
shape::calcOffsets(outerXTadShapeInfo, outerXTadOffsets);
|
|
}
|
|
|
|
LongType* innerXTadOffsets = zOffsets;
|
|
if (!sameOffsets2) {
|
|
ALLOCATE(innerXTadOffsets, workspace, tadLen, sd::LongType);
|
|
shape::calcOffsets(innerXTadShapeInfo, innerXTadOffsets);
|
|
}
|
|
|
|
LongType xRank = shape::rank(xShapeInfo);
|
|
|
|
Z convertedParams[8] = {0};
|
|
Z* compatibleExtraParams = nullptr;
|
|
|
|
if (extraParams != nullptr) {
|
|
if constexpr (std::is_same_v<E, Z>) {
|
|
compatibleExtraParams = reinterpret_cast<Z*>(extraParams);
|
|
} else {
|
|
for (int i = 0; i < 8; ++i) {
|
|
convertedParams[i] = static_cast<Z>(extraParams[i]);
|
|
}
|
|
compatibleExtraParams = convertedParams;
|
|
}
|
|
}
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto i = start; i < stop; ++i) {
|
|
const auto tad = x + outerXTadOffsets[i];
|
|
auto s = static_cast<typename OpType::InterType>(OpType::startingValue(tad));
|
|
|
|
for (LongType j = 0; j < tadLen; j++) {
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(outerXTadShapeInfo);
|
|
shape::printShapeInfo(innerXTadShapeInfo);
|
|
printf("Index i,j is %lld,%lld offset %lld reduceDefault\n", i,j,innerXTadOffsets[j]);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(tad[innerXTadOffsets[j]], compatibleExtraParams), compatibleExtraParams);
|
|
}
|
|
z[zOffsets[i]] = OpType::postProcess(s, tadLen, compatibleExtraParams);
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, shape::length(zShapeInfo));
|
|
|
|
RELEASE(outerXTadShapeInfo, workspace);
|
|
RELEASE(innerXTadShapeInfo, workspace);
|
|
RELEASE(zOffsets, workspace);
|
|
if (!sameOffsets1) RELEASE(outerXTadOffsets, workspace);
|
|
if (!sameOffsets2) RELEASE(innerXTadOffsets, workspace);
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z, typename E>
|
|
template <typename OpType>
|
|
SD_LIB_HIDDEN void ReductionLoops<X, Z, E>::loopReduce(memory::Workspace* workspace, const X* x,
|
|
const LongType* xShapeInfo, Z* z,
|
|
const LongType* zShapeInfo, const LongType* dims,
|
|
E* extraParams) {
|
|
const LongType xRank = shape::rank(xShapeInfo);
|
|
const LongType zRank = shape::rank(zShapeInfo);
|
|
|
|
if (xRank == 2 && zRank == 1)
|
|
reduceExec21<X, Z, E, OpType>(x, xShapeInfo, z, zShapeInfo, dims, extraParams);
|
|
else if (xRank == 3 && zRank == 1)
|
|
reduceExec31<X, Z, E, OpType>(x, xShapeInfo, z, zShapeInfo, dims, extraParams);
|
|
else if (xRank == 3 && zRank == 2)
|
|
reduceExec32<X, Z, E, OpType>(x, xShapeInfo, z, zShapeInfo, dims, extraParams);
|
|
else if (xRank == 4 && zRank == 1)
|
|
reduceExec41<X, Z, E, OpType>(x, xShapeInfo, z, zShapeInfo, dims, extraParams);
|
|
else if (xRank == 4 && zRank == 2)
|
|
reduceExec42<X, Z, E, OpType>(x, xShapeInfo, z, zShapeInfo, dims, extraParams);
|
|
else if (xRank == 4 && zRank == 3)
|
|
reduceExec43<X, Z, E, OpType>(x, xShapeInfo, z, zShapeInfo, dims, extraParams);
|
|
else if (xRank == 5 && zRank == 1)
|
|
reduceExec51<X, Z, E, OpType>(x, xShapeInfo, z, zShapeInfo, dims, extraParams);
|
|
else if (xRank == 5 && zRank == 2)
|
|
reduceExec52<X, Z, E, OpType>(x, xShapeInfo, z, zShapeInfo, dims, extraParams);
|
|
else if (xRank == 5 && zRank == 3)
|
|
reduceExec53<X, Z, E, OpType>(x, xShapeInfo, z, zShapeInfo, dims, extraParams);
|
|
else if (xRank == 5 && zRank == 4)
|
|
reduceExec54<X, Z, E, OpType>(x, xShapeInfo, z, zShapeInfo, dims, extraParams);
|
|
else
|
|
reduceDefault<X, Z, E, OpType>(workspace, x, xShapeInfo, z, zShapeInfo, dims, extraParams);
|
|
}
|
|
|
|
|
|
template <typename X, typename Z, typename E>
|
|
template <typename OpType>
|
|
SD_LIB_HIDDEN void TransformLoops<X, Z, E>::loopTransform(const X* x,
|
|
const LongType* xShapeInfo, Z* z,
|
|
const LongType* zShapeInfo,
|
|
E* extraParams,
|
|
LongType threadId, LongType numThreads) {
|
|
// Basic validation
|
|
if(xShapeInfo == nullptr) {
|
|
THROW_EXCEPTION("TransformLoops::loopTransform: Input x shape info was null!");
|
|
}
|
|
if(zShapeInfo == nullptr) {
|
|
THROW_EXCEPTION("TransformLoops::loopTransform: Input z shape info was null!");
|
|
}
|
|
if(x == nullptr) {
|
|
THROW_EXCEPTION("TransformLoops::loopTransform: Input x buffer was null!");
|
|
}
|
|
if(z == nullptr) {
|
|
THROW_EXCEPTION("TransformLoops::loopTransform: Output z buffer was null!");
|
|
}
|
|
if(xShapeInfo[0] > SD_MAX_RANK || xShapeInfo[0] < 0) {
|
|
THROW_EXCEPTION("TransformLoops::loopTransform: x shape info appears to be corrupt (invalid rank).");
|
|
}
|
|
if(zShapeInfo[0] > SD_MAX_RANK || zShapeInfo[0] < 0) {
|
|
THROW_EXCEPTION("TransformLoops::loopTransform: z shape info appears to be corrupt (invalid rank).");
|
|
}
|
|
|
|
// Get basic shape information
|
|
const LongType zLen = shape::length(zShapeInfo);
|
|
const LongType xLen = shape::length(xShapeInfo);
|
|
const LongType xRank = shape::rank(xShapeInfo);
|
|
const LongType zRank = shape::rank(zShapeInfo);
|
|
|
|
// For element-wise transforms (like Assign), input and output lengths must match
|
|
// This validation catches mismatched TAD shapes that could cause out-of-bounds access
|
|
if (xLen != zLen) {
|
|
std::string error = "TransformLoops::loopTransform: input length ";
|
|
error += std::to_string(xLen);
|
|
error += " != output length ";
|
|
error += std::to_string(zLen);
|
|
error += ". This indicates a shape mismatch between input and output arrays.";
|
|
THROW_EXCEPTION(error.c_str());
|
|
}
|
|
|
|
const LongType* xShape = shape::shapeOf(const_cast<LongType*>(xShapeInfo));
|
|
const LongType* zShape = shape::shapeOf(const_cast<LongType*>(zShapeInfo));
|
|
const LongType* xStride = shape::stride(const_cast<LongType*>(xShapeInfo));
|
|
const LongType* zStride = shape::stride(const_cast<LongType*>(zShapeInfo));
|
|
|
|
// Determine thread work span
|
|
auto span = samediff::Span::build(threadId, numThreads, 0, zLen, 1);
|
|
|
|
// Process each element
|
|
for (auto i = span.startX(); i < span.stopX(); i++) {
|
|
// Calculate coordinates and offsets
|
|
LongType xCoords[SD_MAX_RANK];
|
|
LongType zCoords[SD_MAX_RANK];
|
|
|
|
INDEX2COORDS(i, xRank, xShape, xCoords);
|
|
INDEX2COORDS(i, zRank, zShape, zCoords);
|
|
|
|
LongType xOffset;
|
|
LongType zOffset;
|
|
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
|
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
|
|
|
// Bounds checking for calculated offsets
|
|
// These checks prevent segfaults when TAD shapes don't match or offsets are miscalculated
|
|
if (xOffset < 0 || xOffset >= xLen) {
|
|
// Skip invalid offsets rather than crash - this can happen with malformed TADs
|
|
continue;
|
|
}
|
|
if (zOffset < 0 || zOffset >= zLen) {
|
|
// Skip invalid offsets rather than crash
|
|
continue;
|
|
}
|
|
|
|
// Simple heuristic based on operation patterns observed
|
|
Z result;
|
|
|
|
// For X -> bool transformations (any input type to bool output)
|
|
if constexpr (std::is_same_v<Z, bool>) {
|
|
// Operations that expect X* (input type pointer)
|
|
constexpr bool expects_x_ptr =
|
|
std::is_same_v<OpType, simdOps::IsInf<X, Z>> ||
|
|
std::is_same_v<OpType, simdOps::IsNan<X, Z>> ||
|
|
std::is_same_v<OpType, simdOps::MatchConditionBool<X, Z>> ||
|
|
std::is_same_v<OpType, simdOps::IsPositive<X, Z>> ||
|
|
std::is_same_v<OpType, simdOps::Not<X, Z>> ||
|
|
std::is_same_v<OpType, simdOps::IsNegative<X, Z>> ||
|
|
std::is_same_v<OpType, simdOps::Assign<X, Z>>;
|
|
|
|
if constexpr (expects_x_ptr) {
|
|
result = OpType::op(x[xOffset], static_cast<X*>(nullptr));
|
|
} else {
|
|
// Operations that expect Z* (bool*): IsFinite, IsInfOrNan, etc.
|
|
result = OpType::op(x[xOffset], static_cast<Z*>(nullptr));
|
|
}
|
|
}
|
|
// For same-type transformations, use input type
|
|
else if constexpr (std::is_same_v<X, Z>) {
|
|
result = OpType::op(x[xOffset], static_cast<X*>(nullptr));
|
|
}
|
|
// Default: try with extraParams if types match, otherwise check for operations that expect Z*
|
|
else {
|
|
// Operations that typically expect Z* (output type pointer)
|
|
constexpr bool expects_z_ptr =
|
|
std::is_same_v<OpType, simdOps::Sqrt<X, Z>> ||
|
|
std::is_same_v<OpType, simdOps::RSqrt<X, Z>>; // Added RSqrt here
|
|
|
|
if constexpr (expects_z_ptr) {
|
|
result = OpType::op(x[xOffset], static_cast<Z*>(nullptr));
|
|
} else if constexpr (std::is_same_v<E, X>) {
|
|
result = OpType::op(x[xOffset], extraParams);
|
|
} else {
|
|
result = OpType::op(x[xOffset], static_cast<X*>(nullptr));
|
|
}
|
|
}
|
|
|
|
z[zOffset] = result;
|
|
}
|
|
}
|
|
|
|
|
|
//////////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z>
|
|
template <typename OpType>
|
|
void Reduction3Loops<X, Z>::loopReduce3(const X* x, const LongType* xShapeInfo, const X* y,
|
|
const LongType* yShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
LongType* dims, int dimsLen, Z* extraParameters, int64_t start, int64_t stop) {
|
|
// both tads have same shape, however strides and ews may differ
|
|
|
|
Z param0(OpType::startingValue(x)), param1(OpType::startingValue(x)),
|
|
param2(extraParameters ? extraParameters[0] : OpType::startingValue(x));
|
|
|
|
const LongType xLen = shape::length(xShapeInfo);
|
|
const LongType yLen = shape::length(yShapeInfo);
|
|
|
|
const LongType *xTadShapeInfo = nullptr, *yTadShapeInfo = nullptr, *xTadOffsets = nullptr, *yTadOffsets = nullptr;
|
|
std::shared_ptr<TadPack> tadPackX, tadPackY;
|
|
std::vector<LongType> zeroOffsets;
|
|
|
|
if (xLen == yLen) {
|
|
tadPackX = ConstantTadHelper::getInstance().tadForDimensions(const_cast<sd::LongType*>(xShapeInfo), dims, dimsLen);
|
|
tadPackY = ConstantTadHelper::getInstance().tadForDimensions(const_cast<sd::LongType*>(yShapeInfo), dims, dimsLen);
|
|
xTadShapeInfo = tadPackX->primaryShapeInfo();
|
|
yTadShapeInfo = tadPackY->primaryShapeInfo();
|
|
xTadOffsets = tadPackX->primaryOffsets();
|
|
yTadOffsets = tadPackY->primaryOffsets();
|
|
} else if (yLen > xLen) {
|
|
tadPackY = ConstantTadHelper::getInstance().tadForDimensions(const_cast<sd::LongType*>(yShapeInfo), dims, dimsLen);
|
|
xTadShapeInfo = xShapeInfo;
|
|
yTadShapeInfo = tadPackY->primaryShapeInfo();
|
|
yTadOffsets = tadPackY->primaryOffsets();
|
|
} else {
|
|
tadPackX = ConstantTadHelper::getInstance().tadForDimensions(const_cast<sd::LongType*>(xShapeInfo), dims, dimsLen);
|
|
yTadShapeInfo = yShapeInfo;
|
|
xTadShapeInfo = tadPackX->primaryShapeInfo();
|
|
xTadOffsets = tadPackX->primaryOffsets();
|
|
}
|
|
|
|
const auto zLen = shape::length(zShapeInfo);
|
|
const auto tadLen = shape::length(xTadShapeInfo);
|
|
|
|
const auto tadShape = shape::shapeOf(xTadShapeInfo);
|
|
const auto xTadStride = shape::stride(xTadShapeInfo);
|
|
const auto yTadStride = shape::stride(xTadShapeInfo);
|
|
|
|
int numThreads = OmpLaunchHelper::tadThreads(tadLen, zLen);
|
|
|
|
LongType castXTadShapeInfo[SD_MAX_RANK];
|
|
const bool canCastXTad = DataTypeUtils::castShapeInfo<LongType>(xTadShapeInfo, castXTadShapeInfo);
|
|
LongType castYTadShapeInfo[SD_MAX_RANK];
|
|
const bool canCastYTad = DataTypeUtils::castShapeInfo<LongType>(yTadShapeInfo, castYTadShapeInfo);
|
|
|
|
sd::LongType zRank = shape::rank(zShapeInfo);
|
|
sd::LongType *zShape = shape::shapeOf(zShapeInfo);
|
|
sd::LongType *zStride = shape::stride(zShapeInfo);
|
|
Z extraParams[3];
|
|
for (auto i = start; i < stop; i++) {
|
|
extraParams[0] = param0;
|
|
extraParams[1] = param1;
|
|
extraParams[2] = param2;
|
|
|
|
const auto xTad = xTadOffsets ? x + xTadOffsets[i] : x;
|
|
const auto yTad = yTadOffsets ? y + yTadOffsets[i] : y;
|
|
auto s = OpType::startingValue(xTad);
|
|
|
|
// Calculate z coordinates for this iteration
|
|
LongType zCoords[SD_MAX_RANK];
|
|
INDEX2COORDS(i,zRank, zShape, zCoords);
|
|
LongType zOffset;
|
|
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
|
|
|
sd::LongType xTadRank = shape::rank(xTadShapeInfo);
|
|
sd::LongType yTadRank = shape::rank(yTadShapeInfo);
|
|
sd::LongType *xTadShape = shape::shapeOf(xTadShapeInfo);
|
|
sd::LongType *yTadShape = shape::shapeOf(yTadShapeInfo);
|
|
|
|
for (LongType j = 0; j < tadLen; ++j) {
|
|
LongType coords[SD_MAX_RANK];
|
|
LongType yCoords[SD_MAX_RANK];
|
|
INDEX2COORDS(j, xTadRank, xTadShape, coords);
|
|
INDEX2COORDS(j, yTadRank,yTadShape, yCoords);
|
|
LongType xTadOffset, yTadOffset;
|
|
COORDS2INDEX(xTadRank, xTadStride, coords, xTadOffset);
|
|
COORDS2INDEX(yTadRank,yTadStride, yCoords, yTadOffset);
|
|
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xTadShapeInfo);
|
|
shape::printShapeInfo(yTadShapeInfo);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(xTad[xTadOffset], yTad[yTadOffset], extraParams), extraParams);
|
|
}
|
|
|
|
z[zOffset] = OpType::postProcess(s, tadLen, extraParams);
|
|
};
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Z>
|
|
template <typename OpType>
|
|
void Reduction3Loops<X, Z>::loopReduce3All(const X* x, const LongType* xShapeInfo, const X* y,
|
|
const LongType* yShapeInfo, Z* z, const LongType* zShapeInfo,
|
|
const LongType* xTadShapeInfo, const LongType* xTadOffsets,
|
|
const LongType* yTadShapeInfo, const LongType* yTadOffsets,
|
|
Z* extraParameters, int64_t start, int64_t stop) {
|
|
// both tads have same shape, however strides and ews may differ
|
|
|
|
Z param0(OpType::startingValue(x)), param1(OpType::startingValue(x)),
|
|
param2(extraParameters ? extraParameters[0] : OpType::startingValue(x));
|
|
|
|
|
|
|
|
const auto zLen = shape::length(zShapeInfo);
|
|
const auto tadLen = shape::length(xTadShapeInfo);
|
|
|
|
const auto numXTads = shape::length(xShapeInfo) / tadLen;
|
|
const auto numYTads = shape::length(yShapeInfo) / tadLen;
|
|
|
|
const auto tadShape = shape::shapeOf(xTadShapeInfo);
|
|
const auto xTadStride = shape::stride(xTadShapeInfo);
|
|
const auto yTadStride = shape::stride(yTadShapeInfo);
|
|
|
|
const auto startVal = OpType::startingValue(x);
|
|
|
|
int numThreads = OmpLaunchHelper::tadThreads(tadLen, numXTads * numYTads);
|
|
|
|
//*********************************************//
|
|
LongType castXTadShapeInfo[SD_MAX_RANK];
|
|
LongType castYTadShapeInfo[SD_MAX_RANK];
|
|
const bool canCastYTad = DataTypeUtils::castShapeInfo<LongType>(yTadShapeInfo, castYTadShapeInfo);
|
|
|
|
Z extraParams[3];
|
|
sd::LongType *xTadShape = shape::shapeOf(xTadShapeInfo);
|
|
sd::LongType *yTadShape = shape::shapeOf(yTadShapeInfo);
|
|
sd::LongType xTadRank = shape::rank(xTadShapeInfo);
|
|
sd::LongType zRank = shape::rank(zShapeInfo);
|
|
sd::LongType *zShape = shape::shapeOf(zShapeInfo);
|
|
sd::LongType *zStride = shape::stride(zShapeInfo);
|
|
sd::LongType yTadRank = shape::rank(yTadShapeInfo);
|
|
|
|
for (LongType ix = 0; ix < numXTads; ix++) {
|
|
for (LongType iy = 0; iy < numYTads; iy++) {
|
|
extraParams[0] = param0;
|
|
extraParams[1] = param1;
|
|
extraParams[2] = param2;
|
|
|
|
const auto xTad = x + xTadOffsets[ix];
|
|
const auto yTad = y + yTadOffsets[iy];
|
|
auto s = startVal;
|
|
|
|
for (LongType j = 0; j < tadLen; ++j) {
|
|
LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(j, xTadRank, xTadShape, coords);
|
|
LongType xTadOffset, yTadOffset;
|
|
COORDS2INDEX(xTadRank, xTadStride, coords, xTadOffset);
|
|
COORDS2INDEX(yTadRank, yTadStride, coords, yTadOffset);
|
|
#if defined(PRINT_INDICES)
|
|
shape::printShapeInfo(xTadShapeInfo);
|
|
shape::printShapeInfo(yTadShapeInfo);
|
|
printf("Index is %lld offset is %lld loop kind: default Reduction3Loops<X, Z>::loopReduce3All\n", ix * numYTads + iy, j);
|
|
#endif
|
|
s = OpType::update(s, OpType::op(xTad[xTadOffset], yTad[yTadOffset], extraParams), extraParams);
|
|
}
|
|
z[ix * numYTads + iy] = OpType::postProcess(s, tadLen, extraParams);
|
|
}
|
|
}
|
|
|
|
}
|
|
|
|
} // namespace sd
|
|
|
|
// Include implementation for ReductionLongLoops to make innerloopReduce visible for inlining
|
|
#ifdef SD_LOOPS_INLINED
|
|
#include <helpers/cpu/loops/ReductionLoops_long.cpp>
|
|
#endif
|
|
|
|
#endif // LIBND4J_LOOPS_H
|