wehub-resource-sync
124 lines
5.1 KiB
C++
124 lines
5.1 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 GS <sgazeos@gmail.com>
|
|
//
|
|
#include "../solve.h"
|
|
|
|
#include <array/NDArray.h>
|
|
#include <array/NDArrayFactory.h>
|
|
#include <execution/Threads.h>
|
|
#include <helpers/MmulHelper.h>
|
|
#include <system/op_boilerplate.h>
|
|
|
|
#include "../lup.h"
|
|
#include "../triangular_solve.h"
|
|
#if NOT_EXCLUDED(OP_solve)
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
// ---------------------------------------------------------------------------------------------------------------------------------------
|
|
// //
|
|
template <typename T>
|
|
static void adjointMatrix_(sd::LaunchContext* context, NDArray * input, NDArray* output) {
|
|
auto inputPart = input->allTensorsAlongDimension({-2, -1});
|
|
auto outputPart = output->allTensorsAlongDimension({-2, -1});
|
|
auto rows = input->sizeAt(-2);
|
|
output->assign(input);
|
|
|
|
auto batchLoop = PRAGMA_THREADS_FOR {
|
|
for (auto batch = start; batch < stop; batch++) {
|
|
for (sd::LongType r = 0; r < rows; r++) {
|
|
for (sd::LongType c = 0; c < r; c++) {
|
|
math::sd_swap(outputPart[batch]->r<T>(r, c), outputPart[batch]->r<T>(c, r));
|
|
}
|
|
}
|
|
}
|
|
};
|
|
samediff::Threads::parallel_tad(batchLoop, 0, inputPart.size(), 1);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------------------------------------------------------------------
|
|
// //
|
|
template <typename T>
|
|
static sd::Status solveFunctor_(sd::LaunchContext* context, NDArray* leftInput, NDArray* rightInput, bool const adjoint,
|
|
NDArray* output) {
|
|
// stage 1: LU decomposition batched
|
|
auto leftOutput = leftInput->ulike();
|
|
auto permuShape = rightInput->getShapeAsVector();
|
|
permuShape->pop_back();
|
|
std::vector<sd::LongType> &shapeDeRef = *permuShape;
|
|
auto permutations = NDArrayFactory::create<sd::LongType>('c', shapeDeRef, context);
|
|
helpers::lu(context, leftInput, leftOutput, permutations);
|
|
|
|
|
|
auto P = leftInput->ulike(); // permutations batched matrix
|
|
P->nullify(); // to fill up matrices with zeros
|
|
auto PPart = P->allTensorsAlongDimension({-2, -1});
|
|
auto permutationsPart = permutations->allTensorsAlongDimension({-1});
|
|
for (auto batch = 0; batch < permutationsPart.size(); batch++) {
|
|
for (sd::LongType row = 0; row < PPart[batch]->rows(); row++) {
|
|
std::vector<sd::LongType> vec = {row,permutationsPart[batch]->t<sd::LongType>(row)};
|
|
PPart[batch]->r<T>(row, permutationsPart[batch]->t<sd::LongType>(row)) = T(1.f);
|
|
}
|
|
}
|
|
|
|
|
|
|
|
|
|
auto leftLower = leftOutput->dup(leftOutput->ordering());
|
|
auto rightOutput = rightInput->ulike();
|
|
auto rightPart = rightInput->ulike();
|
|
MmulHelper::matmul(P, rightInput, rightPart, 0.0, 0, 0, 0, rightPart);
|
|
ResultSet leftLowerPart = leftLower->allTensorsAlongDimension({-2, -1});
|
|
for (auto i = 0; i < leftLowerPart.size(); i++) {
|
|
for (sd::LongType r = 0; r < leftLowerPart[i]->rows(); r++) leftLowerPart[i]->r<T>(r, r) = (T)1.f;
|
|
}
|
|
|
|
// stage 2: triangularSolveFunctor for Lower with given b
|
|
helpers::triangularSolveFunctor(context, leftLower, rightPart, true, false, rightOutput);
|
|
// stage 3: triangularSolveFunctor for Upper with output of previous stage
|
|
helpers::triangularSolveFunctor(context, leftOutput, rightOutput, false, false, output);
|
|
|
|
delete permutations;
|
|
delete permuShape;
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------------------------------------------------------------------
|
|
// //
|
|
sd::Status solveFunctor(sd::LaunchContext* context, NDArray* leftInput, NDArray* rightInput, bool const adjoint,
|
|
NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(leftInput->dataType(), return solveFunctor_, (context, leftInput, rightInput, adjoint, output),
|
|
SD_FLOAT_TYPES);
|
|
}
|
|
// ---------------------------------------------------------------------------------------------------------------------------------------
|
|
// //
|
|
void adjointMatrix(sd::LaunchContext* context, NDArray * input, NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(input->dataType(), adjointMatrix_, (context, input, output), SD_FLOAT_TYPES);
|
|
}
|
|
// ---------------------------------------------------------------------------------------------------------------------------------------
|
|
// //
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|
|
#endif |