wehub-resource-sync
924 lines
37 KiB
C++
924 lines
37 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 raver119@gmail.com
|
|
//
|
|
#include <execution/Threads.h>
|
|
#include <helpers/ConstantTadHelper.h>
|
|
#include <helpers/LoopKind.h>
|
|
#include <helpers/ShapeUtils.h>
|
|
#include <loops/broadcasting.h>
|
|
#include <loops/legacy_ops.h>
|
|
#include <system/op_boilerplate.h>
|
|
#include <types/types.h>
|
|
#include <cstdio>
|
|
|
|
using namespace simdOps;
|
|
|
|
namespace functions {
|
|
namespace broadcast {
|
|
|
|
template <typename X, typename Y, typename Z>
|
|
void Broadcast<X, Y, Z>::execInverse(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
|
|
const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
|
|
sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo,
|
|
const sd::LongType *xTadOffset, const sd::LongType *zTadShapeInfo,
|
|
const sd::LongType *zTadOffset, sd::LongType start, sd::LongType stop) {
|
|
DISPATCH_BY_OPNUM_TTT(execInverse,
|
|
PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, dimension, dimensionLength, xTadShapeInfo,
|
|
xTadOffset, zTadShapeInfo, zTadOffset, start, stop),
|
|
BROADCAST_OPS);
|
|
}
|
|
|
|
template <typename X, typename Y, typename Z>
|
|
void Broadcast<X, Y, Z>::exec(int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
|
|
const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo,
|
|
sd::LongType *dimension, sd::LongType dimensionLength, const sd::LongType *xTadShapeInfo,
|
|
const sd::LongType *xTadOffset,
|
|
const sd::LongType *zTadShapeInfo, const sd::LongType *zTadOffset,
|
|
sd::LoopKind::Kind loopKind, sd::LongType start, sd::LongType stop) {
|
|
DISPATCH_BY_OPNUM_TTT(exec,
|
|
PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, dimension, dimensionLength, xTadShapeInfo,
|
|
xTadOffset, zTadShapeInfo, zTadOffset, loopKind, start, stop),
|
|
BROADCAST_OPS);
|
|
}
|
|
template <typename X, typename Y, typename Z>
|
|
template <typename OpType>
|
|
void Broadcast<X, Y, Z>::exec(const void* vx, const sd::LongType* xShapeInfo,
|
|
const void* vy, const sd::LongType* yShapeInfo,
|
|
void* vz, const sd::LongType* zShapeInfo,
|
|
sd::LongType* dimension, sd::LongType dimensionLength,
|
|
const sd::LongType* xTadShapeInfo, const sd::LongType* xTadOffset,
|
|
const sd::LongType* zTadShapeInfo, const sd::LongType* zTadOffset,
|
|
sd::LoopKind::Kind loopKind, sd::LongType start, sd::LongType stop) {
|
|
|
|
auto x = reinterpret_cast<const X*>(vx);
|
|
auto y = reinterpret_cast<const Y*>(vy);
|
|
auto z = reinterpret_cast<Z*>(vz);
|
|
|
|
// Get rank information
|
|
const int xRank = shape::rank(xShapeInfo);
|
|
const int yRank = shape::rank(yShapeInfo);
|
|
const int zRank = shape::rank(zShapeInfo);
|
|
const int xTadRank = xTadShapeInfo ? shape::rank(xTadShapeInfo) : xRank;
|
|
const int zTadRank = zTadShapeInfo ? shape::rank(zTadShapeInfo) : zRank;
|
|
|
|
// Get shape information
|
|
const sd::LongType* xShape = shape::shapeOf(xShapeInfo);
|
|
const sd::LongType* yShape = shape::shapeOf(yShapeInfo);
|
|
const sd::LongType* zShape = shape::shapeOf(zShapeInfo);
|
|
const sd::LongType* xTadShape = shape::shapeOf(xTadShapeInfo);
|
|
const sd::LongType* zTadShape = shape::shapeOf(zTadShapeInfo);
|
|
|
|
// Get stride information
|
|
const sd::LongType* xStrides = shape::stride(xShapeInfo);
|
|
const sd::LongType* yStrides = shape::stride(yShapeInfo);
|
|
const sd::LongType* zStrides = shape::stride(zShapeInfo);
|
|
const sd::LongType* xTadStrides = shape::stride(xTadShapeInfo);
|
|
const sd::LongType* zTadStrides = shape::stride(zTadShapeInfo);
|
|
|
|
// Classify array types
|
|
// For X array or X TAD
|
|
bool isXScalar = xTadRank == 0 || (xTadRank == 1 && xTadShape[0] == 1);
|
|
bool isXVector = (xTadRank == 1 && xTadShape[0] > 1) ||
|
|
(xTadRank == 2 && (xTadShape[0] == 1 || xTadShape[1] == 1));
|
|
bool isXRowVector = (xTadRank == 1 && xTadShape[0] > 1) ||
|
|
(xTadRank == 2 && xTadShape[0] == 1 && xTadShape[1] > 1);
|
|
bool isXColumnVector = (xTadRank == 2 && xTadShape[0] > 1 && xTadShape[1] == 1);
|
|
|
|
// For Y array
|
|
bool isYScalar = yRank == 0 || (yRank == 1 && yShape[0] == 1);
|
|
bool isYVector = (yRank == 1 && yShape[0] > 1) ||
|
|
(yRank == 2 && (yShape[0] == 1 || yShape[1] == 1));
|
|
bool isYRowVector = (yRank == 1 && yShape[0] > 1) ||
|
|
(yRank == 2 && yShape[0] == 1 && yShape[1] > 1);
|
|
bool isYColumnVector = (yRank == 2 && yShape[0] > 1 && yShape[1] == 1);
|
|
|
|
// For Z array or Z TAD
|
|
bool isZScalar = zTadRank == 0 || (zTadRank == 1 && zTadShape[0] == 1);
|
|
bool isZVector = (zTadRank == 1 && zTadShape[0] > 1) ||
|
|
(zTadRank == 2 && (zTadShape[0] == 1 || zTadShape[1] == 1));
|
|
bool isZRowVector = (zTadRank == 1 && zTadShape[0] > 1) ||
|
|
(zTadRank == 2 && zTadShape[0] == 1 && zTadShape[1] > 1);
|
|
bool isZColumnVector = (zTadRank == 2 && zTadShape[0] > 1 && zTadShape[1] == 1);
|
|
|
|
// Handle scalar broadcasting as a special case first
|
|
if (isYScalar) {
|
|
// Scalar broadcast - apply same value to every element
|
|
const Y scalarY = y[0];
|
|
sd::LongType length = shape::length(xTadShapeInfo ? xTadShapeInfo : xShapeInfo);
|
|
|
|
if (xTadShapeInfo && zTadShapeInfo) {
|
|
// TAD case
|
|
for (auto i = start; i < stop; i++) {
|
|
auto oX = x + xTadOffset[i];
|
|
auto oZ = z + zTadOffset[i];
|
|
|
|
// Handle different X and Z shapes
|
|
if (isXVector && isZVector) {
|
|
sd::LongType len = shape::length(xTadShapeInfo);
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType f = 0; f < len; f++) {
|
|
sd::LongType xOffset = f * xTadStrides[xTadRank-1];
|
|
sd::LongType zOffset = f * zTadStrides[zTadRank-1];
|
|
oZ[zOffset] = OpType::op(oX[xOffset], scalarY);
|
|
}
|
|
} else {
|
|
// General case
|
|
for (sd::LongType f = 0; f < length; f++) {
|
|
// Calculate proper offsets for current position
|
|
sd::LongType xCoord[SD_MAX_RANK], zCoord[SD_MAX_RANK];
|
|
sd::LongType xOffset, zOffset;
|
|
|
|
INDEX2COORDS(f, xTadRank, xTadShape, xCoord);
|
|
INDEX2COORDS(f, zTadRank, zTadShape, zCoord);
|
|
|
|
COORDS2INDEX(xTadRank, xTadStrides, xCoord, xOffset);
|
|
COORDS2INDEX(zTadRank, zTadStrides, zCoord, zOffset);
|
|
|
|
oZ[zOffset] = OpType::op(oX[xOffset], scalarY);
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
// Non-TAD case
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType f = 0; f < length; f++)
|
|
z[f] = OpType::op(x[f], scalarY);
|
|
}
|
|
}
|
|
// Handle 2D broadcasting
|
|
else if (loopKind == sd::LoopKind::BROADCAST_2D) {
|
|
// Determine shapes for broadcasting
|
|
sd::LongType nRows = zTadShape[0];
|
|
sd::LongType nCols = zTadRank > 1 ? zTadShape[1] : shape::length(zTadShapeInfo);
|
|
|
|
// Special vector broadcasting cases
|
|
if (isYVector && (isXRowVector || isXColumnVector || isXVector)) {
|
|
// Vector to vector broadcasting
|
|
if (isYRowVector && (isXRowVector || isXVector)) {
|
|
// Row vector to row vector
|
|
for (auto i = start; i < stop; i++) {
|
|
auto baseX = x + xTadOffset[i];
|
|
auto baseZ = z + zTadOffset[i];
|
|
|
|
sd::LongType xStride = xTadRank > 1 ? xTadStrides[xTadRank - 1] : xTadStrides[0];
|
|
sd::LongType yStride = yRank == 1 ? yStrides[0] : yStrides[1];
|
|
sd::LongType zStride = zTadRank ? zTadStrides[zTadRank - 1] : zTadStrides[0];
|
|
|
|
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType i1 = 0; i1 < nCols; i1++) {
|
|
auto rX = baseX + i1 * xStride;
|
|
auto rY = y + i1 * yStride;
|
|
auto rZ = baseZ + i1 * zStride;
|
|
|
|
*rZ = OpType::op(*rX, *rY);
|
|
}
|
|
}
|
|
}
|
|
else if (isYColumnVector && (isXColumnVector || isXVector)) {
|
|
// Column vector to column vector
|
|
// Row vector to row vector
|
|
for (auto i = start; i < stop; i++) {
|
|
auto baseX = x + (xTadOffset ? xTadOffset[i] : 0);
|
|
auto baseZ = z + (zTadOffset ? zTadOffset[i] : 0);
|
|
|
|
// Calculate correct strides based on shape and rank
|
|
sd::LongType xStride;
|
|
if (xTadRank == 1) {
|
|
xStride = xTadStrides[0];
|
|
} else { // xTadRank == 2
|
|
xStride = xTadStrides[0]; // For 2D column vector, use row stride
|
|
}
|
|
|
|
sd::LongType yStride;
|
|
if (yRank == 1) {
|
|
yStride = yStrides[0];
|
|
} else { // yRank == 2
|
|
yStride = yStrides[0]; // For 2D column vector, use row stride
|
|
}
|
|
|
|
sd::LongType zStride;
|
|
if (zTadRank == 1) {
|
|
zStride = zTadStrides[0];
|
|
} else { // zTadRank == 2
|
|
zStride = zTadStrides[0]; // For 2D column vector, use row stride
|
|
}
|
|
|
|
// Verify dimensions match
|
|
sd::LongType xLen = isXColumnVector ? (xTadRank == 2 ? xTadShape[0] : xTadShape[0]) : xTadShape[0];
|
|
sd::LongType yLen = yRank == 2 ? yShape[0] : yShape[0];
|
|
printf("xLen: %lld; yLen: %lld nRows %lld,xStride %lld,yStride %lld, zStride %lld\n", xLen, yLen,nRows,xStride,yStride,zStride);
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType i1 = 0; i1 < nRows; i1++) {
|
|
auto rX = baseX + i1 * xStride;
|
|
auto rY = y + i1 * yStride;
|
|
auto rZ = baseZ + i1 * zStride;
|
|
|
|
*rZ = OpType::op(*rX, *rY);
|
|
}
|
|
}
|
|
}
|
|
else if (isYColumnVector && isXRowVector) {
|
|
// Column vector to row vector (outer product)
|
|
for (auto i = start; i < stop; i++) {
|
|
auto baseX = x + (xTadOffset ? xTadOffset[i] : 0);
|
|
auto baseZ = z + (zTadOffset ? zTadOffset[i] : 0);
|
|
|
|
for (sd::LongType i0 = 0; i0 < nRows; i0++) {
|
|
auto colValue = y[i0 * yStrides[0]];
|
|
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType i1 = 0; i1 < nCols; i1++) {
|
|
auto rX = baseX + i1 * xTadStrides[xTadRank-1];
|
|
auto rZ = baseZ + i0 * zTadStrides[0] + i1 * zTadStrides[1];
|
|
|
|
*rZ = OpType::op(*rX, colValue);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else if (isYRowVector && isXColumnVector) {
|
|
// Row vector to column vector (outer product)
|
|
for (auto i = start; i < stop; i++) {
|
|
printf("4 2d tad: %lld\n", i);
|
|
fflush(stdout);
|
|
auto baseX = x + (xTadOffset ? xTadOffset[i] : 0);
|
|
auto baseZ = z + (zTadOffset ? zTadOffset[i] : 0);
|
|
|
|
for (sd::LongType i0 = 0; i0 < nRows; i0++) {
|
|
auto xValue = baseX[i0 * xTadStrides[0]];
|
|
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType i1 = 0; i1 < nCols; i1++) {
|
|
auto rY = y + i1 * (yRank == 1 ? yStrides[0] : yStrides[1]);
|
|
auto rZ = baseZ + i0 * zTadStrides[0] + i1 * zTadStrides[1];
|
|
|
|
*rZ = OpType::op(xValue, *rY);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// Matrix with vector broadcasting
|
|
else if ((isXRowVector && isYRowVector) || (isXColumnVector && isYColumnVector)) {
|
|
// Matching vectors - element-wise operation
|
|
for (auto i = start; i < stop; i++) {
|
|
auto baseX = x + (xTadOffset ? xTadOffset[i] : 0);
|
|
auto baseZ = z + (zTadOffset ? zTadOffset[i] : 0);
|
|
|
|
sd::LongType vecLength = isXRowVector ? nCols : nRows;
|
|
sd::LongType xStride = isXRowVector ? xTadStrides[xTadRank-1] : xTadStrides[0];
|
|
sd::LongType yStride = isYRowVector ? (yRank == 1 ? yStrides[0] : yStrides[1]) : yStrides[0];
|
|
sd::LongType zStride = isZRowVector ? zTadStrides[zTadRank-1] : zTadStrides[0];
|
|
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType i1 = 0; i1 < vecLength; i1++) {
|
|
auto rX = baseX + i1 * xStride;
|
|
auto rY = y + i1 * yStride;
|
|
auto rZ = baseZ + i1 * zStride;
|
|
|
|
*rZ = OpType::op(*rX, *rY);
|
|
}
|
|
}
|
|
}
|
|
// Matrix with row vector broadcasting
|
|
else if (isYRowVector && xTadRank == 2 && zTadRank == 2 &&
|
|
xTadShape[1] == (yRank == 1 ? yShape[0] : yShape[1])) {
|
|
// Broadcasting row vector (each element applied to a column)
|
|
for (auto i0 = start; i0 < stop; i0++) {
|
|
auto baseX = x + (xTadOffset ? xTadOffset[i0] : 0);
|
|
auto baseZ = z + (zTadOffset ? zTadOffset[i0] : 0);
|
|
|
|
for (sd::LongType i1 = 0; i1 < nRows; i1++) {
|
|
for (sd::LongType i2 = 0; i2 < nCols; i2++) {
|
|
// Get element from X at current position
|
|
auto xOffset = i1 * xTadStrides[0] + i2 * xTadStrides[1];
|
|
// Get element from Y row vector based on column index only
|
|
auto yOffset = i2 * (yRank == 1 ? yStrides[0] : yStrides[1]);
|
|
// Get destination element in Z at current position
|
|
auto zOffset = i1 * zTadStrides[0] + i2 * zTadStrides[1];
|
|
|
|
// Apply operation
|
|
baseZ[zOffset] = OpType::op(baseX[xOffset], y[yOffset]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// Matrix with column vector broadcasting
|
|
else if (isYColumnVector) {
|
|
// Broadcasting column vector (each element applied to a row)
|
|
for (auto i0 = start; i0 < stop; i0++) {
|
|
auto baseX = x + (xTadOffset ? xTadOffset[i0] : 0);
|
|
auto baseZ = z + (zTadOffset ? zTadOffset[i0] : 0);
|
|
|
|
for (sd::LongType i1 = 0; i1 < nRows; i1++) {
|
|
// Get element from column vector based on row index
|
|
auto rY = y + i1 * yStrides[0];
|
|
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType i2 = 0; i2 < nCols; i2++) {
|
|
auto rX = baseX + i1 * xTadStrides[0] + i2 * xTadStrides[1];
|
|
auto rZ = baseZ + i1 * zTadStrides[0] + i2 * zTadStrides[1];
|
|
|
|
*rZ = OpType::op(*rX, *rY);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// Standard 2D broadcasting
|
|
else {
|
|
for (auto i0 = start; i0 < stop; i0++) {
|
|
auto baseX = x + (xTadOffset ? xTadOffset[i0] : 0);
|
|
auto baseZ = z + (zTadOffset ? zTadOffset[i0] : 0);
|
|
|
|
for (sd::LongType i1 = 0; i1 < nRows; i1++) {
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType i2 = 0; i2 < nCols; i2++) {
|
|
auto rX = baseX + i1 * xTadStrides[0] + i2 * xTadStrides[1];
|
|
auto rY = y + i1 * yStrides[0] + i2 * yStrides[1];
|
|
auto rZ = baseZ + i1 * zTadStrides[0] + i2 * zTadStrides[1];
|
|
|
|
*rZ = OpType::op(*rX, *rY);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// Handle remaining loop kinds
|
|
else if (loopKind == sd::LoopKind::BROADCAST_SCALAR_X) {
|
|
sd::LongType tadLength = shape::length(xTadShapeInfo);
|
|
for (auto i = start; i < stop; i++) {
|
|
auto oY = y + (i * tadLength);
|
|
auto oZ = z + (i * tadLength);
|
|
const auto oX = x[i];
|
|
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType f = 0; f < tadLength; f++)
|
|
oZ[f] = OpType::op(oX, oY[f]);
|
|
}
|
|
}
|
|
else if (loopKind == sd::LoopKind::BROADCAST_SCALAR_Y) {
|
|
sd::LongType tadLength = shape::length(xTadShapeInfo);
|
|
for (auto i = start; i < stop; i++) {
|
|
auto oX = x + (i * tadLength);
|
|
auto oZ = z + (i * tadLength);
|
|
const auto oY = y[i];
|
|
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType f = 0; f < tadLength; f++)
|
|
oZ[f] = OpType::op(oX[f], oY);
|
|
}
|
|
}
|
|
// Handle 3D broadcasting (generalized like 2D case)
|
|
else if (loopKind == sd::LoopKind::BROADCAST_3D) {
|
|
// Get TAD info
|
|
const sd::LongType tadRank = xTadShapeInfo ? shape::rank(xTadShapeInfo) : 3;
|
|
const sd::LongType* tadShape = xTadShapeInfo ? shape::shapeOf(xTadShapeInfo) : xShape;
|
|
const sd::LongType* tadStride = xTadShapeInfo ? shape::stride(xTadShapeInfo) : xStrides;
|
|
const sd::LongType tadLength = xTadShapeInfo ? shape::length(xTadShapeInfo) : shape::length(xShapeInfo);
|
|
if (isYVector) {
|
|
// Vector broadcasting
|
|
const sd::LongType yLength = yRank == 1 ? yShape[0] : (yShape[0] == 1 ? yShape[1] : yShape[0]);
|
|
const sd::LongType yStride = yRank == 1 ? yStrides[0] : (yShape[0] == 1 ? yStrides[1] : yStrides[0]);
|
|
|
|
// Determine which dimension this vector should be broadcast along
|
|
// For a 3D TAD, check if vector length matches any dimension
|
|
if (tadRank == 3) {
|
|
if (yLength == tadShape[2]) {
|
|
// Broadcast along last dimension
|
|
for (auto i = start; i < stop; i++) {
|
|
auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
|
|
auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
|
|
|
|
for (sd::LongType j = 0; j < tadLength; j++) {
|
|
// Calculate TAD coords
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(j, tadRank, tadShape, coords);
|
|
|
|
// Get offsets
|
|
sd::LongType xOffset, zOffset;
|
|
COORDS2INDEX(tadRank, tadStride, coords, xOffset);
|
|
COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
|
|
|
|
// Get Y index - use the last dimension (coords[2])
|
|
sd::LongType yOffset = coords[2] * yStride;
|
|
|
|
// Apply operation
|
|
oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
|
|
}
|
|
}
|
|
}
|
|
else if (yLength == tadShape[1]) {
|
|
// Broadcast along middle dimension
|
|
PRAGMA_OMP_SIMD
|
|
for (auto i = start; i < stop; i++) {
|
|
auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
|
|
auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
|
|
|
|
for (sd::LongType j = 0; j < tadLength; j++) {
|
|
// Calculate TAD coords
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(j, tadRank, tadShape, coords);
|
|
|
|
// Get offsets
|
|
sd::LongType xOffset, zOffset;
|
|
COORDS2INDEX(tadRank, tadStride, coords, xOffset);
|
|
COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
|
|
|
|
// Get Y index - use the middle dimension (coords[1])
|
|
sd::LongType yOffset = coords[1] * yStride;
|
|
|
|
// Apply operation
|
|
oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
|
|
}
|
|
}
|
|
}
|
|
else if (yLength == tadShape[0]) {
|
|
// Broadcast along first dimension
|
|
PRAGMA_OMP_SIMD
|
|
for (auto i = start; i < stop; i++) {
|
|
auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
|
|
auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
|
|
|
|
for (sd::LongType j = 0; j < tadLength; j++) {
|
|
// Calculate TAD coords
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(j, tadRank, tadShape, coords);
|
|
|
|
// Get offsets
|
|
sd::LongType xOffset, zOffset;
|
|
COORDS2INDEX(tadRank, tadStride, coords, xOffset);
|
|
COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
|
|
|
|
// Get Y index - use the first dimension (coords[0])
|
|
sd::LongType yOffset = coords[0] * yStride;
|
|
|
|
// Apply operation
|
|
oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
|
|
}
|
|
}
|
|
}
|
|
else {
|
|
// Default broadcasting behavior - broadcast along the last dimension
|
|
PRAGMA_OMP_SIMD
|
|
for (auto i = start; i < stop; i++) {
|
|
auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
|
|
auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
|
|
|
|
for (sd::LongType j = 0; j < tadLength; j++) {
|
|
// Calculate TAD coords
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(j, tadRank, tadShape, coords);
|
|
|
|
// Get offsets
|
|
sd::LongType xOffset, zOffset;
|
|
COORDS2INDEX(tadRank, tadStride, coords, xOffset);
|
|
COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
|
|
|
|
// Get Y index with wrapping/broadcasting
|
|
sd::LongType yOffset = (coords[2] % yLength) * yStride;
|
|
|
|
// Apply operation
|
|
oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else {
|
|
// Handle lower rank TADs (1D or 2D)
|
|
PRAGMA_OMP_SIMD
|
|
for (auto i = start; i < stop; i++) {
|
|
auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
|
|
auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
|
|
|
|
for (sd::LongType j = 0; j < tadLength; j++) {
|
|
// Calculate TAD coords
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(j, tadRank, tadShape, coords);
|
|
|
|
// Get offsets
|
|
sd::LongType xOffset, zOffset;
|
|
COORDS2INDEX(tadRank, tadStride, coords, xOffset);
|
|
COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
|
|
|
|
// Get Y index - for lower ranks, broadcast along the last available dimension
|
|
sd::LongType lastCoord = tadRank > 0 ? coords[tadRank - 1] : 0;
|
|
sd::LongType yOffset = (lastCoord % yLength) * yStride;
|
|
|
|
// Apply operation
|
|
oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else if (yRank == 2) {
|
|
// Y is a 2D matrix - determine which dimensions it aligns with
|
|
for (auto i = start; i < stop; i++) {
|
|
auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
|
|
auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType j = 0; j < tadLength; j++) {
|
|
// Calculate TAD coords
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(j, tadRank, tadShape, coords);
|
|
|
|
// Get offsets
|
|
sd::LongType xOffset, zOffset;
|
|
COORDS2INDEX(tadRank, tadStride, coords, xOffset);
|
|
COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
|
|
|
|
// Calculate Y offset based on dimension matching
|
|
sd::LongType yOffset;
|
|
|
|
// Default behavior for different 2D matrix broadcasting patterns
|
|
if (tadRank == 3) {
|
|
if (yShape[0] == tadShape[0] && yShape[1] == tadShape[2]) {
|
|
// Y is aligned with dimensions 0 and 2
|
|
yOffset = coords[0] * yStrides[0] + coords[2] * yStrides[1];
|
|
}
|
|
else if (yShape[0] == tadShape[0] && yShape[1] == tadShape[1]) {
|
|
// Y is aligned with dimensions 0 and 1
|
|
yOffset = coords[0] * yStrides[0] + coords[1] * yStrides[1];
|
|
}
|
|
else if (yShape[0] == tadShape[1] && yShape[1] == tadShape[2]) {
|
|
// Y is aligned with dimensions 1 and 2
|
|
yOffset = coords[1] * yStrides[0] + coords[2] * yStrides[1];
|
|
}
|
|
else {
|
|
// Default: broadcast Y to match the last two dimensions with modulo
|
|
yOffset = (coords[1] % yShape[0]) * yStrides[0] + (coords[2] % yShape[1]) * yStrides[1];
|
|
}
|
|
}
|
|
else if (tadRank == 2) {
|
|
// Direct mapping for 2D TAD and 2D Y
|
|
yOffset = (coords[0] % yShape[0]) * yStrides[0] + (coords[1] % yShape[1]) * yStrides[1];
|
|
}
|
|
else {
|
|
// For 1D TAD, map to the first dimension of Y
|
|
yOffset = (coords[0] % yShape[0]) * yStrides[0];
|
|
}
|
|
|
|
// Apply operation
|
|
oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
|
|
}
|
|
}
|
|
}
|
|
else if (yRank == 3) {
|
|
// Y is a 3D tensor
|
|
PRAGMA_OMP_SIMD
|
|
for (auto i = start; i < stop; i++) {
|
|
auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
|
|
auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
|
|
|
|
for (sd::LongType j = 0; j < tadLength; j++) {
|
|
// Calculate TAD coords
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(j, tadRank, tadShape, coords);
|
|
|
|
// Get offsets
|
|
sd::LongType xOffset, zOffset;
|
|
COORDS2INDEX(tadRank, tadStride, coords, xOffset);
|
|
COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
|
|
|
|
// Calculate Y offset with modulo for broadcasting if needed
|
|
sd::LongType yCoords[3] = {0, 0, 0};
|
|
|
|
// Map coordinates appropriately based on ranks
|
|
if (tadRank == 3) {
|
|
yCoords[0] = coords[0] % yShape[0];
|
|
yCoords[1] = coords[1] % yShape[1];
|
|
yCoords[2] = coords[2] % yShape[2];
|
|
}
|
|
else if (tadRank == 2) {
|
|
// Map 2D to last 2 dimensions of 3D
|
|
yCoords[0] = 0; // First dimension is broadcasted
|
|
yCoords[1] = coords[0] % yShape[1];
|
|
yCoords[2] = coords[1] % yShape[2];
|
|
}
|
|
else {
|
|
// Map 1D to last dimension of 3D
|
|
yCoords[0] = 0;
|
|
yCoords[1] = 0;
|
|
yCoords[2] = coords[0] % yShape[2];
|
|
}
|
|
|
|
sd::LongType yOffset = yCoords[0] * yStrides[0] + yCoords[1] * yStrides[1] + yCoords[2] * yStrides[2];
|
|
|
|
// Apply operation
|
|
oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
|
|
}
|
|
}
|
|
}
|
|
else {
|
|
// General case for other ranks of Y
|
|
for (auto i = start; i < stop; i++) {
|
|
auto oX = x + (xTadOffset ? xTadOffset[i] : 0);
|
|
auto oZ = z + (zTadOffset ? zTadOffset[i] : 0);
|
|
|
|
for (sd::LongType j = 0; j < tadLength; j++) {
|
|
// Calculate TAD coords
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(j, tadRank, tadShape, coords);
|
|
|
|
// Get offsets
|
|
sd::LongType xOffset, zOffset;
|
|
COORDS2INDEX(tadRank, tadStride, coords, xOffset);
|
|
COORDS2INDEX(tadRank, zTadStrides, coords, zOffset);
|
|
|
|
// Calculate Y offset based on rank
|
|
sd::LongType yOffset = 0;
|
|
|
|
// Map coordinates to Y based on rank
|
|
for (int d = 0; d < tadRank && d < yRank; d++) {
|
|
yOffset += (coords[d] % yShape[d]) * yStrides[d];
|
|
}
|
|
|
|
// Apply operation
|
|
oZ[zOffset] = OpType::op(oX[xOffset], y[yOffset]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else if (loopKind == sd::LoopKind::BROADCAST_4D) {
|
|
const sd::LongType nSize1 = shape::sizeAt(zShapeInfo, 1);
|
|
const sd::LongType nSize2 = shape::sizeAt(zShapeInfo, 2);
|
|
const sd::LongType nSize3 = shape::sizeAt(zShapeInfo, 3);
|
|
|
|
for (auto i = start; i < stop; i++) {
|
|
uint64_t i0 = i / nSize1;
|
|
uint64_t i1 = i % nSize1;
|
|
|
|
for (sd::LongType i2 = 0; i2 < nSize2; i2++) {
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType i3 = 0; i3 < nSize3; i3++) {
|
|
auto rX = x + (xStrides[0] * i0 + xStrides[1] * i1 + xStrides[2] * i2 + xStrides[3] * i3);
|
|
auto rY = y + (yStrides[0] * i0 + yStrides[1] * i1 + yStrides[2] * i2 + yStrides[3] * i3);
|
|
auto rZ = z + (zStrides[0] * i0 + zStrides[1] * i1 + zStrides[2] * i2 + zStrides[3] * i3);
|
|
|
|
*rZ = OpType::op(*rX, *rY);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else if (loopKind == sd::LoopKind::BROADCAST_5D) {
|
|
const sd::LongType nSize1 = shape::sizeAt(zShapeInfo, 1);
|
|
const sd::LongType nSize2 = shape::sizeAt(zShapeInfo, 2);
|
|
const sd::LongType nSize3 = shape::sizeAt(zShapeInfo, 3);
|
|
const sd::LongType nSize4 = shape::sizeAt(zShapeInfo, 4);
|
|
|
|
for (auto i = start; i < stop; i++) {
|
|
uint32_t i0 = i / nSize1;
|
|
uint32_t i1 = i % nSize1;
|
|
|
|
for (sd::LongType i2 = 0; i2 < nSize2; i2++) {
|
|
for (sd::LongType i3 = 0; i3 < nSize3; i3++) {
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType i4 = 0; i4 < nSize4; i4++) {
|
|
auto rX = x + (xStrides[0] * i0 + xStrides[1] * i1 + xStrides[2] * i2 + xStrides[3] * i3 + xStrides[4] * i4);
|
|
auto rY = y + (yStrides[0] * i0 + yStrides[1] * i1 + yStrides[2] * i2 + yStrides[3] * i3 + yStrides[4] * i4);
|
|
auto rZ = z + (zStrides[0] * i0 + zStrides[1] * i1 + zStrides[2] * i2 + zStrides[3] * i3 + zStrides[4] * i4);
|
|
|
|
*rZ = OpType::op(*rX, *rY);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else {
|
|
// Default case for other ranks - general purpose implementation
|
|
sd::LongType xCoords[SD_MAX_RANK];
|
|
sd::LongType yCoords[SD_MAX_RANK];
|
|
sd::LongType zCoords[SD_MAX_RANK];
|
|
|
|
for (auto i = start; i < stop; i++) {
|
|
// Calculate independent coordinates for each array
|
|
INDEX2COORDS(i, xRank, xShape, xCoords);
|
|
INDEX2COORDS(i, yRank, yShape, yCoords);
|
|
INDEX2COORDS(i, zRank, zShape, zCoords);
|
|
|
|
// Calculate offsets based on each array's coordinates and strides
|
|
sd::LongType xOffset, yOffset, zOffset;
|
|
COORDS2INDEX(xRank, xStrides, xCoords, xOffset);
|
|
COORDS2INDEX(yRank, yStrides, yCoords, yOffset);
|
|
COORDS2INDEX(zRank, zStrides, zCoords, zOffset);
|
|
|
|
z[zOffset] = OpType::op(x[xOffset], y[yOffset]);
|
|
}
|
|
}
|
|
}
|
|
template <typename X, typename Y, typename Z>
|
|
template <typename OpType>
|
|
void Broadcast<X, Y, Z>::execInverse(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
|
|
const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo,
|
|
sd::LongType *dimension, sd::LongType dimensionLength,
|
|
const sd::LongType *yTadShapeInfo, const sd::LongType *yTadOffset,
|
|
const sd::LongType *zTadShapeInfo, const sd::LongType *zTadOffset,
|
|
sd::LongType start, sd::LongType stop) {
|
|
auto x = reinterpret_cast<const X *>(vx);
|
|
auto y = reinterpret_cast<const Y *>(vy);
|
|
auto z = reinterpret_cast<Z *>(vz);
|
|
|
|
// Handle TAD setup
|
|
auto yTadShapeShapeInfo = yTadShapeInfo;
|
|
auto tadOffsets = yTadOffset;
|
|
|
|
// When shared_ptr goes out of scope, it deletes the TadPack and invalidates pointers!
|
|
std::shared_ptr<sd::TadPack> tadPack = nullptr;
|
|
|
|
if (yTadShapeInfo == nullptr || tadOffsets == nullptr) {
|
|
tadPack = sd::ConstantTadHelper::getInstance().tadForDimensions(const_cast<sd::LongType*>(yShapeInfo), dimension,
|
|
dimensionLength);
|
|
yTadShapeShapeInfo = tadPack->primaryShapeInfo();
|
|
tadOffsets = tadPack->primaryOffsets();
|
|
}
|
|
|
|
if (zTadShapeInfo == nullptr) {
|
|
zTadShapeInfo = yTadShapeShapeInfo;
|
|
zTadOffset = tadOffsets;
|
|
}
|
|
|
|
// Get shape information
|
|
const auto xRank = shape::rank(xShapeInfo);
|
|
const auto yTadRank = shape::rank(yTadShapeShapeInfo);
|
|
const auto zTadRank = shape::rank(zTadShapeInfo);
|
|
|
|
const auto xStrides = shape::stride(xShapeInfo);
|
|
const auto yTadStrides = shape::stride(yTadShapeShapeInfo);
|
|
const auto zTadStrides = shape::stride(zTadShapeInfo);
|
|
|
|
const auto xShape = shape::shapeOf(xShapeInfo);
|
|
const auto yTadShape = shape::shapeOf(yTadShapeShapeInfo);
|
|
const auto zTadShape = shape::shapeOf(zTadShapeInfo);
|
|
|
|
const sd::LongType tadLength = shape::length(yTadShapeShapeInfo);
|
|
|
|
if (yTadRank <= 3) {
|
|
// Optimized path for lower ranks
|
|
for (auto i = start; i < stop; i++) {
|
|
auto oZ = z + zTadOffset[i];
|
|
auto oY = y + tadOffsets[i];
|
|
|
|
if (yTadRank == 1) {
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType j = 0; j < tadLength; j++) {
|
|
oZ[j * zTadStrides[0]] = OpType::op(x[j * xStrides[0]], oY[j * yTadStrides[0]]);
|
|
}
|
|
}
|
|
else if (yTadRank == 2) {
|
|
const sd::LongType dim0 = yTadShape[0];
|
|
const sd::LongType dim1 = yTadShape[1];
|
|
|
|
for (sd::LongType j0 = 0; j0 < dim0; j0++) {
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType j1 = 0; j1 < dim1; j1++) {
|
|
const auto xOffset = j0 * xStrides[0] + j1 * xStrides[1];
|
|
const auto yOffset = j0 * yTadStrides[0] + j1 * yTadStrides[1];
|
|
const auto zOffset = j0 * zTadStrides[0] + j1 * zTadStrides[1];
|
|
oZ[zOffset] = OpType::op(x[xOffset], oY[yOffset]);
|
|
}
|
|
}
|
|
}
|
|
else { // rank 3
|
|
const sd::LongType dim0 = yTadShape[0];
|
|
const sd::LongType dim1 = yTadShape[1];
|
|
const sd::LongType dim2 = yTadShape[2];
|
|
|
|
for (sd::LongType j0 = 0; j0 < dim0; j0++) {
|
|
for (sd::LongType j1 = 0; j1 < dim1; j1++) {
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType j2 = 0; j2 < dim2; j2++) {
|
|
const auto xOffset = j0 * xStrides[0] + j1 * xStrides[1] + j2 * xStrides[2];
|
|
const auto yOffset = j0 * yTadStrides[0] + j1 * yTadStrides[1] + j2 * yTadStrides[2];
|
|
const auto zOffset = j0 * zTadStrides[0] + j1 * zTadStrides[1] + j2 * zTadStrides[2];
|
|
oZ[zOffset] = OpType::op(x[xOffset], oY[yOffset]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else {
|
|
// Use macros for higher ranks
|
|
for (auto i = start; i < stop; i++) {
|
|
auto oZ = z + zTadOffset[i];
|
|
auto oY = y + tadOffsets[i];
|
|
|
|
PRAGMA_OMP_SIMD
|
|
for (sd::LongType f = 0; f < tadLength; f++) {
|
|
sd::LongType coords[SD_MAX_RANK];
|
|
INDEX2COORDS(f, yTadRank, yTadShape, coords);
|
|
|
|
sd::LongType xOffset, yOffset, zOffset;
|
|
COORDS2INDEX(xRank, xStrides, coords, xOffset);
|
|
COORDS2INDEX(yTadRank, yTadStrides, coords, yOffset);
|
|
COORDS2INDEX(zTadRank, zTadStrides, coords, zOffset);
|
|
|
|
oZ[zOffset] = OpType::op(x[xOffset], oY[yOffset]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename X, typename Y, typename Z>
|
|
void Broadcast<X, Y, Z>::exec(const int opNum, const void *x, const sd::LongType *xShapeInfo, const void *y,
|
|
const sd::LongType *yShapeInfo, void *z, const sd::LongType *zShapeInfo) {
|
|
DISPATCH_BY_OPNUM_TTT(exec, PARAMS(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo), BROADCAST_OPS);
|
|
}
|
|
|
|
template <typename X, typename Y, typename Z, typename OpType>
|
|
static void execDefault(const X *x, const sd::LongType *xShapeInfo, const Y *y, const sd::LongType *yShapeInfo, Z *z,
|
|
const sd::LongType *zShapeInfo) {
|
|
// Cache shape-related values
|
|
sd::LongType xRank = shape::rank(xShapeInfo);
|
|
sd::LongType yRank = shape::rank(yShapeInfo);
|
|
sd::LongType zRank = shape::rank(zShapeInfo);
|
|
|
|
// C-style arrays CANNOT be captured by value in lambdas - they decay to pointers
|
|
// that point to stack memory. std::array CAN be captured by value, ensuring each
|
|
// parallel thread gets its own copy of the data with guaranteed lifetime.
|
|
std::array<sd::LongType, SD_MAX_RANK> xShapeLocal;
|
|
std::array<sd::LongType, SD_MAX_RANK> yShapeLocal;
|
|
std::array<sd::LongType, SD_MAX_RANK> zShapeLocal;
|
|
std::array<sd::LongType, SD_MAX_RANK> xStrideLocal;
|
|
std::array<sd::LongType, SD_MAX_RANK> yStrideLocal;
|
|
std::array<sd::LongType, SD_MAX_RANK> zStrideLocal;
|
|
|
|
// Copy actual data from shapeInfo into std::arrays
|
|
std::memcpy(xShapeLocal.data(), shape::shapeOf(xShapeInfo), xRank * sizeof(sd::LongType));
|
|
std::memcpy(yShapeLocal.data(), shape::shapeOf(yShapeInfo), yRank * sizeof(sd::LongType));
|
|
std::memcpy(zShapeLocal.data(), shape::shapeOf(zShapeInfo), zRank * sizeof(sd::LongType));
|
|
std::memcpy(xStrideLocal.data(), shape::stride(xShapeInfo), xRank * sizeof(sd::LongType));
|
|
std::memcpy(yStrideLocal.data(), shape::stride(yShapeInfo), yRank * sizeof(sd::LongType));
|
|
std::memcpy(zStrideLocal.data(), shape::stride(zShapeInfo), zRank * sizeof(sd::LongType));
|
|
|
|
// Capture std::arrays by value - C++ will copy the entire array contents into the lambda's closure.
|
|
// This ensures each parallel thread has its own copy of the data with no dangling pointers.
|
|
auto func = [x, y, z, xRank, yRank, zRank, xShapeLocal, yShapeLocal, zShapeLocal, xStrideLocal, yStrideLocal, zStrideLocal](
|
|
sd::LongType thread_id, sd::LongType start, sd::LongType stop, sd::LongType increment) -> void {
|
|
for (auto i = start; i < stop; ++i) {
|
|
sd::LongType zCoords[SD_MAX_RANK];
|
|
sd::LongType xCoords[SD_MAX_RANK];
|
|
sd::LongType yCoords[SD_MAX_RANK];
|
|
|
|
// Convert linear index to coordinates based on Z (output) shape
|
|
INDEX2COORDS(i, zRank, zShapeLocal.data(), zCoords);
|
|
|
|
// Broadcast Z coordinates to X and Y shapes
|
|
// For broadcasting, we map Z coords to X and Y coords using modulo for smaller dimensions
|
|
// When a dimension is 1 in X or Y but larger in Z, we use index 0 (broadcast)
|
|
for (sd::LongType d = 0; d < xRank; d++) {
|
|
xCoords[d] = xShapeLocal[d] == 1 ? 0 : (zCoords[d] % xShapeLocal[d]);
|
|
}
|
|
for (sd::LongType d = 0; d < yRank; d++) {
|
|
yCoords[d] = yShapeLocal[d] == 1 ? 0 : (zCoords[d] % yShapeLocal[d]);
|
|
}
|
|
|
|
sd::LongType xOffset, yOffset, zOffset;
|
|
COORDS2INDEX(xRank, xStrideLocal.data(), xCoords, xOffset);
|
|
COORDS2INDEX(yRank, yStrideLocal.data(), yCoords, yOffset);
|
|
COORDS2INDEX(zRank, zStrideLocal.data(), zCoords, zOffset);
|
|
|
|
z[zOffset] = OpType::op(x[xOffset], y[yOffset]);
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, static_cast<sd::LongType>(0), shape::length(zShapeInfo));
|
|
}
|
|
|
|
template <typename X, typename Y, typename Z>
|
|
template <typename OpType>
|
|
void Broadcast<X, Y, Z>::exec(const void *vx, const sd::LongType *xShapeInfo, const void *vy,
|
|
const sd::LongType *yShapeInfo, void *vz, const sd::LongType *zShapeInfo) {
|
|
const X *x = reinterpret_cast<const X *>(vx);
|
|
const Y *y = reinterpret_cast<const Y *>(vy);
|
|
Z *z = reinterpret_cast<Z *>(vz);
|
|
|
|
const int rank = shape::rank(zShapeInfo); // xRank = yRank = zRank
|
|
|
|
switch (rank) {
|
|
default:
|
|
execDefault<X, Y, Z, OpType>(x, xShapeInfo, y, yShapeInfo, z, zShapeInfo);
|
|
}
|
|
}
|
|
|
|
|
|
} // namespace broadcast
|
|
} // namespace functions
|