wehub-resource-sync
231 lines
9.2 KiB
C++
231 lines
9.2 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
|
|
******************************************************************************/
|
|
|
|
//
|
|
// Created by george on 05.04.18.
|
|
//
|
|
#include <execution/Threads.h>
|
|
#include <ops/declarable/helpers/dynamic.h>
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
template <typename T>
|
|
static void _dynamicPartitionFunctor(NDArray * input, NDArray * indices, std::vector<NDArray*>& outputList) {
|
|
std::vector<std::pair<NDArray*, sd::LongType>> outputs(outputList.size());
|
|
int sourceDimsLen = input->rankOf() - indices->rankOf();
|
|
if (sourceDimsLen) {
|
|
std::vector<sd::LongType> sourceDims(sourceDimsLen);
|
|
|
|
for (sd::LongType i = sourceDimsLen; i > 0; i--) sourceDims[sourceDimsLen - i] = input->rankOf() - i;
|
|
|
|
ResultSet listOfTensors = input->allTensorsAlongDimension(sourceDims);
|
|
|
|
sd::LongType outSize = outputList.size();
|
|
|
|
for (sd::LongType i = 0; i < outSize; i++) {
|
|
outputs[i].first = outputList[i];
|
|
std::vector<sd::LongType > outDims(outputs[i].first->rankOf() - 1);
|
|
|
|
sd::LongType r = outputs[i].first->rankOf();
|
|
|
|
for (sd::LongType k = 1; k < r; k++) outDims[k - 1] = k;
|
|
|
|
ResultSet listOutForCurrent = outputs[i].first->allTensorsAlongDimension(outDims);
|
|
|
|
outputs[i].second = 0;
|
|
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e)
|
|
if ((*indices).e<sd::LongType>(e) == i) {
|
|
listOutForCurrent.at(outputs[i].second++)->assign(listOfTensors.at(e));
|
|
}
|
|
}
|
|
|
|
} else {
|
|
sd::LongType outSize = outputList.size();
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto i = start; i < stop; i++) {
|
|
outputs[i].first = outputList[i];
|
|
outputs[i].second = 0;
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e)
|
|
if (indices->e<sd::LongType>(e) == i) outputs[i].first->p(outputs[i].second++, input->e<T>(e));
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_tad(func, 0, outSize);
|
|
}
|
|
}
|
|
template <typename T>
|
|
static sd::Status _dynamicStitchFunctor(std::vector<NDArray*> const& inputs, std::vector<NDArray*> const& indices,
|
|
NDArray* output) {
|
|
sd::LongType numOfData = inputs.size();
|
|
|
|
if (output->isVector()) {
|
|
for (sd::LongType e = 0; e < numOfData; e++) {
|
|
auto data = inputs[e];
|
|
auto index = indices[e];
|
|
for (sd::LongType i = 0; i < index->lengthOf(); i++) {
|
|
sd::LongType pos = index->e<sd::LongType>(i);
|
|
if (pos < 0) {
|
|
sd_printf("dynamic_stitch: Index value should be non-negative. But %i was given", pos);
|
|
return sd::Status::VALIDATION;
|
|
}
|
|
if (pos >= output->lengthOf()) {
|
|
sd_printf("dynamic_stitch: Index should be less than %i. But %i was given", output->lengthOf(), pos);
|
|
return sd::Status::VALIDATION;
|
|
}
|
|
output->p<T>(pos, data->e<T>(i));
|
|
}
|
|
}
|
|
} else {
|
|
std::vector<sd::LongType > restDims(output->rankOf() - 1);
|
|
for (auto i = restDims.size(); i > 0; i--) restDims[restDims.size() - i] = output->rankOf() - i;
|
|
|
|
ResultSet listOfOutTensors = output->allTensorsAlongDimension(restDims);
|
|
for (int e = 0; e < numOfData; e++) {
|
|
auto data = inputs[e];
|
|
auto index = indices[e];
|
|
std::vector<sd::LongType > sourceDims(data->rankOf() - index->rankOf());
|
|
for (auto i = sourceDims.size(); i > 0; i--) sourceDims[sourceDims.size() - i] = data->rankOf() - i;
|
|
|
|
ResultSet listOfTensors = data->allTensorsAlongDimension(sourceDims);
|
|
|
|
for (sd::LongType i = 0; i < index->lengthOf(); i++) {
|
|
auto pos = index->e<sd::LongType>(i);
|
|
if (pos < 0) {
|
|
sd_printf("dynamic_stitch: Index value should be non-negative. But %i was given", pos);
|
|
return sd::Status::VALIDATION;
|
|
}
|
|
if (pos >= output->lengthOf()) {
|
|
sd_printf("dynamic_stitch: Index should be less than %i. But %i was given", output->lengthOf(), pos);
|
|
return sd::Status::VALIDATION;
|
|
}
|
|
|
|
listOfOutTensors.at(pos)->assign(listOfTensors.at(i));
|
|
}
|
|
}
|
|
}
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
template <typename T>
|
|
static void _dynamicPartitionFunctorBP(NDArray * input, NDArray * indices,
|
|
std::vector<NDArray*> const& inputGradientList,
|
|
std::vector<NDArray*>& outputList) {
|
|
std::vector<std::pair<NDArray*, sd::LongType>> outputs(inputGradientList.size());
|
|
|
|
int sourceDimsLen = input->rankOf() - indices->rankOf();
|
|
if (sourceDimsLen) { // multidimensional case
|
|
std::vector<sd::LongType > sourceDims(sourceDimsLen);
|
|
|
|
for (sd::LongType i = sourceDimsLen; i > 0; i--) sourceDims[sourceDimsLen - i] = input->rankOf() - i;
|
|
|
|
ResultSet listOfTensors = outputList[0]->allTensorsAlongDimension(sourceDims);
|
|
|
|
for (size_t i = 0; i < inputGradientList.size(); i++) {
|
|
outputs[i].first = inputGradientList[i];
|
|
if (outputs[i].first->rankOf() < 1) continue; // skip empty gradient outs
|
|
std::vector<sd::LongType > outDims(outputs[i].first->rankOf() - 1);
|
|
|
|
for (int k = 1; k < outputs[i].first->rankOf(); k++) outDims[k - 1] = k;
|
|
|
|
ResultSet listOutForCurrent = outputs[i].first->allTensorsAlongDimension(outDims);
|
|
|
|
outputs[i].second = 0;
|
|
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e)
|
|
if (indices->e<sd::LongType>(e) == static_cast<sd::LongType>(i)) listOfTensors.at(e)->assign(listOutForCurrent.at(outputs[i].second++));
|
|
}
|
|
} else { // one-dimensional case
|
|
auto output = outputList[0];
|
|
unsigned int gradsSize = inputGradientList.size();
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto i = start; i < stop; i++) {
|
|
outputs[i].first = inputGradientList[i];
|
|
outputs[i].second = 0;
|
|
for (sd::LongType e = 0; e < indices->lengthOf(); ++e)
|
|
if (indices->e<sd::LongType>(e) == i) output->p<T>(e, outputs[i].first->e<T>(outputs[i].second++));
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_tad(func, 0, gradsSize);
|
|
}
|
|
|
|
outputList[1]->assign(indices);
|
|
}
|
|
|
|
void dynamicPartitionFunctor(sd::LaunchContext* context, NDArray * input, NDArray * indices,
|
|
std::vector<NDArray*>& outputList) {
|
|
auto xType = input->dataType();
|
|
|
|
BUILD_SINGLE_SELECTOR(xType, _dynamicPartitionFunctor, (input, indices, outputList), SD_COMMON_TYPES);
|
|
}
|
|
|
|
template <typename T>
|
|
static sd::Status _dynamicStitchFunctorBP(std::vector<NDArray*> const& inputs, std::vector<NDArray*> const& indices,
|
|
NDArray * gradInput, std::vector<NDArray*>& outputList) {
|
|
THROW_EXCEPTION("Not implemented yet");
|
|
}
|
|
|
|
sd::Status dynamicStitchFunctor(sd::LaunchContext* context, std::vector<NDArray*> const& inputs,
|
|
std::vector<NDArray*> const& indices, NDArray* output) {
|
|
auto xType = inputs.at(0)->dataType();
|
|
|
|
BUILD_SINGLE_SELECTOR(xType, return _dynamicStitchFunctor, (inputs, indices, output), SD_COMMON_TYPES);
|
|
}
|
|
|
|
sd::Status dynamicStitchFunctorBP(sd::LaunchContext* context, std::vector<NDArray*> const& inputs,
|
|
std::vector<NDArray*> const& indices, NDArray * gradInput,
|
|
std::vector<NDArray*>& outputList) {
|
|
auto xType = inputs.at(0)->dataType();
|
|
|
|
BUILD_SINGLE_SELECTOR(xType, return _dynamicStitchFunctorBP, (inputs, indices, gradInput, outputList),
|
|
SD_COMMON_TYPES);
|
|
}
|
|
|
|
void dynamicPartitionFunctorBP(sd::LaunchContext* context, NDArray * input, NDArray * indices,
|
|
std::vector<NDArray*> const& inputGradientList, std::vector<NDArray*>& outputList) {
|
|
auto xType = input->dataType();
|
|
|
|
BUILD_SINGLE_SELECTOR(xType, _dynamicPartitionFunctorBP, (input, indices, inputGradientList, outputList),
|
|
SD_COMMON_TYPES);
|
|
}
|
|
|
|
BUILD_SINGLE_TEMPLATE( void _dynamicPartitionFunctorBP,
|
|
(NDArray * input, NDArray * indices, std::vector<NDArray*> const& inputGradientList,
|
|
std::vector<NDArray*>& outputList);
|
|
, SD_COMMON_TYPES);
|
|
BUILD_SINGLE_TEMPLATE( sd::Status _dynamicStitchFunctorBP,
|
|
(std::vector<NDArray*> const& inputs, std::vector<NDArray*> const& indices,
|
|
NDArray * gradInput, std::vector<NDArray*>& outputList);
|
|
, SD_COMMON_TYPES);
|
|
|
|
BUILD_SINGLE_TEMPLATE( void _dynamicPartitionFunctor,
|
|
(NDArray * input, NDArray * indices, std::vector<NDArray*>& outputList);
|
|
, SD_COMMON_TYPES);
|
|
BUILD_SINGLE_TEMPLATE( sd::Status _dynamicStitchFunctor,
|
|
(std::vector<NDArray*> const& inputs, std::vector<NDArray*> const& indices, NDArray* output);
|
|
, SD_COMMON_TYPES);
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|