wehub-resource-sync
133 lines
5.7 KiB
Plaintext
133 lines
5.7 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 <helpers/ConstantTadHelper.h>
|
|
#include <helpers/MmulHelper.h>
|
|
#include <helpers/ShapeUtils.h>
|
|
#include <ops/declarable/helpers/lstsq.h>
|
|
#include <ops/declarable/helpers/lup.h>
|
|
#include <ops/declarable/helpers/qr.h>
|
|
#include <ops/declarable/helpers/triangular_solve.h>
|
|
#include <system/op_boilerplate.h>
|
|
|
|
#include "execution/cuda/LaunchDims.h"
|
|
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
template <typename T>
|
|
static SD_KERNEL void fillRegularizerKernel(T* ioMatrixData, const LongType* ioMatrixShape,
|
|
const LongType* ioMatrixTads, const LongType* ioMatrixOffsets,
|
|
LongType batchSize, LongType rows, T const value) {
|
|
for (auto x = blockIdx.x; x < batchSize; x += gridDim.x) {
|
|
auto z = ioMatrixData + ioMatrixOffsets[x];
|
|
for (auto r = threadIdx.x; r < rows; r += blockDim.x) {
|
|
LongType pos[] = {r, r};
|
|
LongType zIndex;
|
|
COORDS2INDEX(2, shape::stride(ioMatrixTads), pos, zIndex);
|
|
z[zIndex] = value;
|
|
}
|
|
}
|
|
}
|
|
template <typename T>
|
|
static void fillRegularizer(LaunchContext* context, NDArray* ioMatrix, double const value) {
|
|
std::vector<LongType> dims = {-2, -1};
|
|
auto lastDimsTads = ConstantTadHelper::getInstance().tadForDimensions(ioMatrix->shapeInfo(), &dims);
|
|
auto stream = context->getCudaStream();
|
|
auto rows = ioMatrix->sizeAt(-2);
|
|
dim3 launchDims = getLaunchDims("lstsq_reg");
|
|
fillRegularizerKernel<T><<<launchDims.y,launchDims.x,launchDims.z, *stream>>>(
|
|
ioMatrix->dataBuffer()->specialAsT<T>(), ioMatrix->specialShapeInfo(), lastDimsTads->specialShapeInfo(),
|
|
lastDimsTads->specialOffsets(), lastDimsTads->numberOfTads(), rows, (T)value);
|
|
|
|
}
|
|
|
|
template <typename T>
|
|
Status leastSquaresSolveFunctor_(LaunchContext* context, NDArray* leftInput, NDArray* rightInput,
|
|
double const l2Regularizer, bool const fast, NDArray* output) {
|
|
if (fast) { // Cholesky decomposition approach
|
|
// Equation for solve A^T * Ax = A^T * b, so
|
|
// 1. Computing A2:
|
|
auto tAtShape = ShapeUtils::evalShapeForMatmul(leftInput->shapeInfo(), leftInput->shapeInfo(), true, false);
|
|
// tAtShape[tAtShape.size() - 2] = output->sizeAt(-2);
|
|
NDArray leftOutput(leftInput->ordering(), tAtShape, output->dataType(), context);
|
|
MmulHelper::matmul(leftInput, leftInput, &leftOutput, true, false,1.0,0.0,&leftOutput); // Computing A2 = A^T * A
|
|
// 2. Computing B' = A^T * b
|
|
auto rightOutput = output->ulike();
|
|
|
|
MmulHelper::matmul(leftInput, rightInput, rightOutput, true, false,1.0,0.0,rightOutput); // Computing B' = A^T * b
|
|
// 3. Regularization ( indeed A' = A2 - l2Regularizer * I)
|
|
if (l2Regularizer != 0.0) {
|
|
auto regularizer = leftOutput.ulike();
|
|
regularizer->nullify();
|
|
fillRegularizer<T>(context, regularizer, (T)l2Regularizer);
|
|
leftOutput += *regularizer;
|
|
}
|
|
|
|
// 4. Cholesky decomposition -- output matrix is square and lower triangular
|
|
cholesky(context, &leftOutput, &leftOutput, true); // inplace decomposition
|
|
// 5. Solve two triangular systems:
|
|
auto rightB = rightOutput->ulike();
|
|
rightB->nullify();
|
|
|
|
triangularSolveFunctor(context, &leftOutput, rightOutput, true, false, rightB);
|
|
|
|
adjointMatrix(context, &leftOutput, true, &leftOutput);
|
|
triangularSolveFunctor(context, &leftOutput, rightB, false, false, output);
|
|
// All done
|
|
} else { // QR decomposition approach
|
|
// Equation for solve Rx = Q^T * b, where A = Q * R, where Q - orthogonal matrix, and R - upper triangular
|
|
// 1. QR decomposition
|
|
auto* qShapePtr = leftInput->getShapeAsVector();
|
|
std::vector<LongType> qShape = *qShapePtr;
|
|
delete qShapePtr;
|
|
auto* rShapePtr = leftInput->getShapeAsVector();
|
|
std::vector<LongType> rShape = *rShapePtr;
|
|
delete rShapePtr;
|
|
qShape[leftInput->rankOf() - 1] = leftInput->sizeAt(-2);
|
|
|
|
NDArray Q(leftInput->ordering(), qShape, leftInput->dataType(), context);
|
|
NDArray R(leftInput->ordering(), rShape, leftInput->dataType(), context);
|
|
qr(context, leftInput, &Q, &R, true);
|
|
// 2. b` = Q^t * b:
|
|
auto rightOutput = rightInput->ulike();
|
|
MmulHelper::matmul(&Q, rightInput, rightOutput, true, false,1.0,0.0,rightOutput);
|
|
// 3. Solve triangular system
|
|
triangularSolveFunctor(context, &R, rightOutput, false, false, output);
|
|
}
|
|
return Status::OK;
|
|
}
|
|
|
|
Status leastSquaresSolveFunctor(LaunchContext* context, NDArray* leftInput, NDArray* rightInput,
|
|
double const l2Regularizer, bool const fast, NDArray* output) {
|
|
BUILD_SINGLE_SELECTOR(leftInput->dataType(), return leastSquaresSolveFunctor_,
|
|
(context, leftInput, rightInput, l2Regularizer, fast, output), SD_FLOAT_TYPES);
|
|
}
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|