chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,355 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 07.03.2019
|
||||
//
|
||||
#include <execution/Threads.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
#include <ops/declarable/helpers/gather.h>
|
||||
#include <legacy/NativeOpExecutioner.h>
|
||||
|
||||
#include <numeric>
|
||||
#if NOT_EXCLUDED(OP_gather)
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
void gather(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output,
|
||||
const std::vector<LongType>& intArgs) {
|
||||
sd::LongType axis = intArgs.size() > 0 ? intArgs[0] : 0;
|
||||
const sd::LongType inputRank = input->rankOf();
|
||||
if (axis < 0) axis += inputRank;
|
||||
|
||||
const sd::LongType numOfIntArgs = intArgs.size();
|
||||
|
||||
// Special handling for 1D input with axis=0
|
||||
// This handles cases like gathering from shape arrays where we want flat indexing
|
||||
bool is1DFlatGather = (inputRank == 1 && axis == 0);
|
||||
|
||||
if (indices != nullptr) {
|
||||
// Validate indices
|
||||
for (sd::LongType i = 0; i < indices->lengthOf(); ++i) {
|
||||
auto idx = indices->e<sd::LongType>(i);
|
||||
|
||||
if (is1DFlatGather) {
|
||||
// For 1D arrays with axis=0, treat as flat array access
|
||||
if (idx >= input->lengthOf() || idx < 0) {
|
||||
std::string error = "helpers::gather function: invalid flat index ";
|
||||
error += std::to_string(idx);
|
||||
error += " at position ";
|
||||
error += std::to_string(i);
|
||||
error += ". Input is 1D with length ";
|
||||
error += std::to_string(input->lengthOf());
|
||||
error += ", valid range is [0, ";
|
||||
error += std::to_string(input->lengthOf() - 1);
|
||||
error += "]";
|
||||
THROW_EXCEPTION(error.c_str());
|
||||
}
|
||||
} else {
|
||||
// Standard axis-based validation
|
||||
if (idx >= input->sizeAt(axis) || idx < 0) {
|
||||
std::string error = "helpers::gather function: invalid index ";
|
||||
error += std::to_string(idx);
|
||||
error += " at position ";
|
||||
error += std::to_string(i);
|
||||
error += ". Input shape ";
|
||||
error += ShapeUtils::shapeAsString(input->shapeInfo());
|
||||
error += ", axis ";
|
||||
error += std::to_string(axis);
|
||||
error += ", valid range is [0, ";
|
||||
error += std::to_string(input->sizeAt(axis) - 1);
|
||||
error += "]";
|
||||
THROW_EXCEPTION(error.c_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (is1DFlatGather) {
|
||||
// Special case: 1D input with axis=0 - treat as flat array gather
|
||||
// This handles gathering from shape arrays like [1, 512] -> gather index 1 -> get 512
|
||||
auto func = PRAGMA_THREADS_FOR {
|
||||
for (auto i = start; i < stop; i++) {
|
||||
auto idx = indices->e<sd::LongType>(i);
|
||||
auto value = input->e<double>(idx); // Get value at flat index
|
||||
output->p(i, value); // Put in output at position i
|
||||
}
|
||||
};
|
||||
samediff::Threads::parallel_for(func, 0, indices->lengthOf());
|
||||
|
||||
} else {
|
||||
// Standard gather implementation
|
||||
//
|
||||
// For gather with axis=A on input shape [..., dimA, ...] and indices shape [I1, I2, ...]:
|
||||
// - Output shape is: input[0:A] + indices_shape + input[A+1:]
|
||||
// - Input TADs: iterate along axis A, each TAD has shape input[A+1:]
|
||||
// - Output TADs: iterate along indices dimensions, each TAD has same shape as input TAD
|
||||
//
|
||||
// tadForDimensions takes dimensions to KEEP in each TAD (not to exclude)
|
||||
// It then internally calls evalDimsToExclude to find which dims to iterate over
|
||||
|
||||
std::vector<sd::LongType> axesVec = {axis};
|
||||
auto dimensions = ShapeUtils::evalDimsToExclude(input->rankOf(), 1, axesVec.data());
|
||||
|
||||
// For output TADs, we want the same shape as input TADs
|
||||
// Input TAD shape = all dims except axis
|
||||
// Output shape = input[0:axis] + indices_shape + input[axis+1:]
|
||||
// Output TAD dims should be: dims 0 to axis-1, then dims axis+indicesRank to end
|
||||
// This gives TAD shape matching input's TAD shape
|
||||
std::vector<sd::LongType> outputTadDims;
|
||||
sd::LongType indicesRank = indices->rankOf();
|
||||
|
||||
// Add dimensions before the indices dimensions (0 to axis-1)
|
||||
for (sd::LongType d = 0; d < axis; d++) {
|
||||
outputTadDims.push_back(d);
|
||||
}
|
||||
// Add dimensions after the indices dimensions (axis+indicesRank to outputRank-1)
|
||||
for (sd::LongType d = axis + indicesRank; d < output->rankOf(); d++) {
|
||||
outputTadDims.push_back(d);
|
||||
}
|
||||
|
||||
// If outputTadDims is empty, it means each TAD is a scalar - handle this case
|
||||
// by using the same approach as input (which would also have empty TAD dims)
|
||||
|
||||
// Get TAD packs - these are cached and should not be deleted
|
||||
auto tadPack = sd::ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), dimensions);
|
||||
auto tadPackOut = sd::ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), &outputTadDims);
|
||||
|
||||
// Validate TAD packs before use
|
||||
if (tadPack == nullptr || tadPackOut == nullptr) {
|
||||
if (dimensions) delete dimensions;
|
||||
THROW_EXCEPTION("gather: Failed to create TAD packs");
|
||||
}
|
||||
|
||||
// Now safe to delete dimensions as TAD helper has made internal copy
|
||||
delete dimensions;
|
||||
|
||||
auto tadShapeInfo = tadPack->primaryShapeInfo();
|
||||
auto tadOffsets = tadPack->primaryOffsets();
|
||||
auto tadShapeInfoOut = tadPackOut->primaryShapeInfo();
|
||||
auto tadOffsetsOut = tadPackOut->primaryOffsets();
|
||||
|
||||
// Validate that input and output TAD shapes match
|
||||
auto inputTadLength = shape::length(tadShapeInfo);
|
||||
auto outputTadLength = shape::length(tadShapeInfoOut);
|
||||
if (inputTadLength != outputTadLength) {
|
||||
std::string error = "gather: TAD shape mismatch - input TAD length ";
|
||||
error += std::to_string(inputTadLength);
|
||||
error += " != output TAD length ";
|
||||
error += std::to_string(outputTadLength);
|
||||
error += ". Input shape: ";
|
||||
error += ShapeUtils::shapeAsString(input->shapeInfo());
|
||||
error += ", Output shape: ";
|
||||
error += ShapeUtils::shapeAsString(output->shapeInfo());
|
||||
error += ", Indices shape: ";
|
||||
error += ShapeUtils::shapeAsString(indices->shapeInfo());
|
||||
error += ", axis: ";
|
||||
error += std::to_string(axis);
|
||||
THROW_EXCEPTION(error.c_str());
|
||||
}
|
||||
|
||||
auto tadShapeInfoCast = const_cast<sd::LongType *>(tadShapeInfo);
|
||||
auto tadShapeInfoOutCast = const_cast<sd::LongType *>(tadShapeInfoOut);
|
||||
|
||||
// Calculate the number of gather operations (equal to indices length)
|
||||
const sd::LongType numGatherOps = indices->lengthOf();
|
||||
|
||||
// Validate bounds before parallel execution
|
||||
if (numGatherOps > tadPackOut->numberOfTads()) {
|
||||
std::string error = "gather: indices length ";
|
||||
error += std::to_string(numGatherOps);
|
||||
error += " exceeds output TAD count ";
|
||||
error += std::to_string(tadPackOut->numberOfTads());
|
||||
THROW_EXCEPTION(error.c_str());
|
||||
}
|
||||
|
||||
auto func = PRAGMA_THREADS_FOR {
|
||||
for (auto i = start; i < stop; i++) {
|
||||
auto idx = indices->e<sd::LongType>(i);
|
||||
|
||||
// Bounds check for input TAD access
|
||||
if (idx >= tadPack->numberOfTads() || idx < 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Bounds check for output TAD access
|
||||
if (i >= tadPackOut->numberOfTads()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
auto offsetIn = tadOffsets[idx];
|
||||
auto offsetOut = tadOffsetsOut[i];
|
||||
|
||||
NativeOpExecutioner::execTransformAny(input->getContext(),
|
||||
transform::Assign,
|
||||
input->bufferWithOffset(offsetIn), tadShapeInfoCast,
|
||||
nullptr, nullptr,
|
||||
output->bufferWithOffset(offsetOut), tadShapeInfoOutCast,
|
||||
nullptr, nullptr,
|
||||
nullptr, false);
|
||||
}
|
||||
};
|
||||
samediff::Threads::parallel_tad(func, 0, numGatherOps);
|
||||
}
|
||||
|
||||
} else {
|
||||
// Integer arguments case
|
||||
for (int i = 1; i < numOfIntArgs; ++i) {
|
||||
if (is1DFlatGather) {
|
||||
// For 1D arrays with axis=0, validate against total length
|
||||
if (intArgs[i] >= input->lengthOf() || intArgs[i] < 0) {
|
||||
std::string error = "helpers::gather function: invalid flat index ";
|
||||
error += std::to_string(intArgs[i]);
|
||||
error += " at position ";
|
||||
error += std::to_string(i-1);
|
||||
error += ". Input is 1D with length ";
|
||||
error += std::to_string(input->lengthOf());
|
||||
error += ", valid range is [0, ";
|
||||
error += std::to_string(input->lengthOf() - 1);
|
||||
error += "]";
|
||||
THROW_EXCEPTION(error.c_str());
|
||||
}
|
||||
} else {
|
||||
// Standard validation
|
||||
if (intArgs[i] >= input->sizeAt(axis) || intArgs[i] < 0) {
|
||||
std::string error = "helpers::gather function: invalid index ";
|
||||
error += std::to_string(intArgs[i]);
|
||||
error += " at position ";
|
||||
error += std::to_string(i-1);
|
||||
error += ". Input shape ";
|
||||
error += ShapeUtils::shapeAsString(input->shapeInfo());
|
||||
error += ", axis ";
|
||||
error += std::to_string(axis);
|
||||
error += ", valid range is [0, ";
|
||||
error += std::to_string(input->sizeAt(axis) - 1);
|
||||
error += "]";
|
||||
THROW_EXCEPTION(error.c_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (numOfIntArgs == 2) {
|
||||
if (is1DFlatGather) {
|
||||
// For 1D flat gather with single index
|
||||
auto value = input->e<double>(intArgs[1]);
|
||||
output->assign(value);
|
||||
} else {
|
||||
// Standard single index gather
|
||||
NDArray *copy = (*input)(intArgs[1], {axis});
|
||||
output->assign(copy);
|
||||
delete copy;
|
||||
}
|
||||
} else {
|
||||
if (is1DFlatGather) {
|
||||
// Multiple indices for 1D flat gather
|
||||
for (int i = 1; i < numOfIntArgs; ++i) {
|
||||
auto idx = intArgs[i];
|
||||
auto value = input->e<double>(idx);
|
||||
output->p(i - 1, value);
|
||||
}
|
||||
} else {
|
||||
// Standard multiple indices gather
|
||||
// Use the same dimension calculation for input and output TADs
|
||||
std::vector<sd::LongType> axesVec = {axis};
|
||||
auto dimensions = ShapeUtils::evalDimsToExclude(input->rankOf(), 1, axesVec.data());
|
||||
|
||||
// Get TAD packs - these are cached and should not be deleted
|
||||
auto tadPack = sd::ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), dimensions);
|
||||
auto tadPackOut = sd::ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), dimensions);
|
||||
|
||||
// Validate TAD packs before use
|
||||
if (tadPack == nullptr || tadPackOut == nullptr) {
|
||||
if (dimensions) delete dimensions;
|
||||
THROW_EXCEPTION("gather: Failed to create TAD packs");
|
||||
}
|
||||
|
||||
// Now safe to delete dimensions as TAD helper has made internal copy
|
||||
delete dimensions;
|
||||
|
||||
auto tadShapeInfo = tadPack->primaryShapeInfo();
|
||||
auto tadOffsets = tadPack->primaryOffsets();
|
||||
auto tadShapeInfoOut = tadPackOut->primaryShapeInfo();
|
||||
auto tadOffsetsOut = tadPackOut->primaryOffsets();
|
||||
|
||||
// Validate that input and output TAD shapes match
|
||||
auto inputTadLength = shape::length(tadShapeInfo);
|
||||
auto outputTadLength = shape::length(tadShapeInfoOut);
|
||||
if (inputTadLength != outputTadLength) {
|
||||
std::string error = "gather: TAD shape mismatch - input TAD length ";
|
||||
error += std::to_string(inputTadLength);
|
||||
error += " != output TAD length ";
|
||||
error += std::to_string(outputTadLength);
|
||||
error += ". Input shape: ";
|
||||
error += ShapeUtils::shapeAsString(input->shapeInfo());
|
||||
error += ", Output shape: ";
|
||||
error += ShapeUtils::shapeAsString(output->shapeInfo());
|
||||
error += ", axis: ";
|
||||
error += std::to_string(axis);
|
||||
THROW_EXCEPTION(error.c_str());
|
||||
}
|
||||
|
||||
// Number of gather operations (number of indices provided as int args)
|
||||
const sd::LongType numGatherOps = numOfIntArgs - 1;
|
||||
|
||||
// Validate bounds before parallel execution
|
||||
if (numGatherOps > tadPackOut->numberOfTads()) {
|
||||
std::string error = "gather: number of indices ";
|
||||
error += std::to_string(numGatherOps);
|
||||
error += " exceeds output TAD count ";
|
||||
error += std::to_string(tadPackOut->numberOfTads());
|
||||
THROW_EXCEPTION(error.c_str());
|
||||
}
|
||||
|
||||
auto func = PRAGMA_THREADS_FOR {
|
||||
for (auto i = start; i < stop; i++) {
|
||||
auto idx = intArgs[i + 1];
|
||||
|
||||
// Bounds check for input TAD access
|
||||
if (idx >= tadPack->numberOfTads() || idx < 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Bounds check for output TAD access
|
||||
if (i >= tadPackOut->numberOfTads()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
auto offsetIn = tadOffsets[idx];
|
||||
auto offsetOut = tadOffsetsOut[i];
|
||||
|
||||
NativeOpExecutioner::execTransformAny(input->getContext(),
|
||||
transform::Assign,
|
||||
input->bufferWithOffset(offsetIn), const_cast<sd::LongType*>(tadShapeInfo),
|
||||
nullptr, nullptr,
|
||||
output->bufferWithOffset(offsetOut), const_cast<sd::LongType*>(tadShapeInfoOut),
|
||||
nullptr, nullptr,
|
||||
nullptr, false);
|
||||
}
|
||||
};
|
||||
samediff::Threads::parallel_tad(func, 0, numGatherOps);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
#endif
|
||||
Reference in New Issue
Block a user