wehub-resource-sync
145 lines
5.5 KiB
Plaintext
145 lines
5.5 KiB
Plaintext
/*
|
|
* ******************************************************************************
|
|
* *
|
|
* *
|
|
* * 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 <array/NDArray.h>
|
|
#include <array/NDArrayFactory.h>
|
|
#include <execution/Threads.h>
|
|
#include <helpers/ConstantTadHelper.h>
|
|
#include <helpers/MmulHelper.h>
|
|
#include <system/op_boilerplate.h>
|
|
|
|
#include "../lup.h"
|
|
#include "../solve.h"
|
|
#include "../triangular_solve.h"
|
|
#include "execution/cuda/LaunchDims.h"
|
|
#include "helpers/DebugHelper.h"
|
|
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
|
|
|
|
|
|
template <typename T>
|
|
static Status solveFunctor_(LaunchContext* context, NDArray* leftInput, NDArray* rightInput, bool adjoint,
|
|
NDArray* output) {
|
|
// TODO: note: this is the cpu implementation.
|
|
// it's not preferred but cuda has enough edge cases
|
|
// that I would prefer to have a working solution for now.
|
|
NDArray::preparePrimaryUse({output}, {leftInput, rightInput});
|
|
|
|
// stage 1: LU decomposition batched
|
|
auto leftOutput = leftInput->ulike();
|
|
|
|
auto permuShape = rightInput->getShapeAsVector();
|
|
permuShape.pop_back();
|
|
auto permutations = NDArrayFactory::create<LongType>('c', permuShape, context);
|
|
lu(context, leftInput, leftOutput, &permutations);
|
|
auto leftLower = leftOutput->dup();
|
|
|
|
auto rightOutput = rightInput->ulike();
|
|
|
|
const std::vector<LongType> dims1 = {-2, -1};
|
|
|
|
auto P = leftInput->ulike();
|
|
P->nullify();
|
|
auto PPart = P->allTensorsAlongDimension({-2, -1});
|
|
auto permutationsPart = permutations.allTensorsAlongDimension({-1});
|
|
for (auto batch = 0; batch < permutationsPart.size(); batch++) {
|
|
for (LongType row = 0; row < PPart[batch]->rows(); row++) {
|
|
std::vector<LongType> vec = {row, permutationsPart[batch]->t<LongType>(row)};
|
|
PPart[batch]->r<T>(row, permutationsPart[batch]->t<LongType>(row)) = T(1.f);
|
|
}
|
|
}
|
|
|
|
P->tickWriteHost();
|
|
|
|
auto rightPart = rightInput->ulike();
|
|
|
|
MmulHelper::matmul(P, rightInput, rightPart,false,false, 0.0, 0.0,rightPart);
|
|
ResultSet leftLowerPart = leftLower.allTensorsAlongDimension({-2, -1});
|
|
for (auto i = 0; i < leftLowerPart.size(); i++) {
|
|
for (LongType r = 0; r < leftLowerPart[i]->rows(); r++) leftLowerPart[i]->r<T>(r, r) = (T)1.f;
|
|
}
|
|
triangularSolveFunctor(context, &leftLower, rightPart, true, false, rightOutput);
|
|
triangularSolveFunctor(context, leftOutput, rightOutput, false, false, output);
|
|
NDArray::registerPrimaryUse({output}, {leftInput, rightInput});
|
|
|
|
return Status::OK;
|
|
}
|
|
|
|
Status solveFunctor(LaunchContext* context, NDArray* leftInput, NDArray* rightInput, bool adjoint,
|
|
NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(leftInput->dataType(), return solveFunctor_, (context, leftInput, rightInput, adjoint, output),
|
|
SD_FLOAT_TYPES);
|
|
}
|
|
|
|
template <typename T>
|
|
static SD_KERNEL void adjointKernel(T* output, LongType batchSize, LongType rows, LongType columns,
|
|
LongType const* outputTads, LongType const* outputOffsets) {
|
|
for (auto b = blockIdx.x; b < batchSize; b += gridDim.x) {
|
|
auto outputPart = output + outputOffsets[b];
|
|
for (auto r = threadIdx.x; r < rows; r += blockDim.x) {
|
|
for (auto c = threadIdx.y; c < r; c += blockDim.y) {
|
|
LongType zPos[] = {r, c};
|
|
LongType xPos[] = {c, r};
|
|
LongType zIndex, xIndex;
|
|
COORDS2INDEX(shape::rank(outputTads), shape::stride(outputTads), zPos, zIndex);
|
|
COORDS2INDEX(shape::rank(outputTads), shape::stride(outputTads), xPos, xIndex);
|
|
math::sd_swap(outputPart[zIndex], outputPart[xIndex]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
static void adjointMatrix_(LaunchContext* context, NDArray * input, NDArray* output) {
|
|
NDArray::prepareSpecialUse({output}, {input});
|
|
const std::vector<LongType> dims1 = {-2, -1};
|
|
auto outputTads = ConstantTadHelper::getInstance().tadForDimensions(
|
|
output->shapeInfo(), const_cast<LongType*>(dims1.data()), dims1.size());
|
|
auto stream = context->getCudaStream();
|
|
auto outputBuf = reinterpret_cast<T*>(output->specialBuffer());
|
|
auto rows = input->sizeAt(-2);
|
|
auto columns = input->sizeAt(-1);
|
|
output->assign(input);
|
|
dim3 solveDims = getLaunchDims("solve");
|
|
|
|
adjointKernel<T><<<solveDims.x,solveDims.y, solveDims.z, *stream>>>(outputBuf, outputTads->numberOfTads(), rows, columns,
|
|
outputTads->specialShapeInfo(), outputTads->specialOffsets());
|
|
|
|
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "adjointKernel failed");
|
|
|
|
NDArray::registerSpecialUse({output}, {input});
|
|
}
|
|
|
|
void adjointMatrix(LaunchContext* context, NDArray * input, NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(input->dataType(), adjointMatrix_, (context, input, output), SD_FLOAT_TYPES);
|
|
}
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|