wehub-resource-sync
125 lines
3.6 KiB
Plaintext
125 lines
3.6 KiB
Plaintext
#include "LayerNormExecution.hpp"
|
|
#include "core/MusaBackend.hpp"
|
|
|
|
namespace MNN {
|
|
namespace MUSA {
|
|
|
|
template<typename T>
|
|
__global__ void LayerNormKernel(const T* input, const T* gamma, const T* beta, T* output,
|
|
int outerSize, int innerSize,
|
|
T epsilon, int gammaSize, int betaSize) {
|
|
int outerIdx = blockIdx.x;
|
|
|
|
if (outerIdx < outerSize) {
|
|
// Compute mean
|
|
T sum = 0;
|
|
for (int i = 0; i < innerSize; i++) {
|
|
int idx = outerIdx * innerSize + i;
|
|
sum += input[idx];
|
|
}
|
|
T mean = sum / innerSize;
|
|
|
|
// Compute variance
|
|
T var = 0;
|
|
for (int i = 0; i < innerSize; i++) {
|
|
int idx = outerIdx * innerSize + i;
|
|
T diff = input[idx] - mean;
|
|
var += diff * diff;
|
|
}
|
|
var = var / innerSize;
|
|
|
|
// Normalize
|
|
T invStd = 1.0 / sqrt(var + epsilon);
|
|
|
|
for (int i = 0; i < innerSize; i++) {
|
|
int idx = outerIdx * innerSize + i;
|
|
T normalized = (input[idx] - mean) * invStd;
|
|
|
|
T g = (gamma != nullptr && gammaSize > 0) ? gamma[i % gammaSize] : 1.0;
|
|
T b = (beta != nullptr && betaSize > 0) ? beta[i % betaSize] : 0.0;
|
|
|
|
output[idx] = normalized * g + b;
|
|
}
|
|
}
|
|
}
|
|
|
|
LayerNormExecution::LayerNormExecution(const std::vector<Tensor*>& inputs, const MNN::Op* op, Backend* backend)
|
|
: Execution(inputs, {}, backend) {
|
|
mBackend = static_cast<MusaBackend*>(backend);
|
|
mOp = op->main_as_LayerNorm();
|
|
}
|
|
|
|
ErrorCode LayerNormExecution::onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
auto input = inputs[0];
|
|
auto output = outputs[0];
|
|
|
|
mEpsilon = mOp->eps();
|
|
|
|
mOuterSize = 1;
|
|
for (int i = 0; i < input->dimensions() - 1; i++) {
|
|
mOuterSize *= input->length(i);
|
|
}
|
|
mInnerSize = input->length(input->dimensions() - 1);
|
|
|
|
mGammaSize = 0;
|
|
mBetaSize = 0;
|
|
if (mOp->gamma() != nullptr) {
|
|
mGammaSize = mOp->gamma()->size();
|
|
}
|
|
if (mOp->beta() != nullptr) {
|
|
mBetaSize = mOp->beta()->size();
|
|
}
|
|
|
|
int threads = 256;
|
|
dim3 grid(mOuterSize, 1, 1);
|
|
dim3 block(threads, 1, 1);
|
|
|
|
mDim3Grid = grid;
|
|
mDim3Block = block;
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode LayerNormExecution::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
auto input = inputs[0];
|
|
auto output = outputs[0];
|
|
|
|
auto inputPtr = input->host<float>();
|
|
auto outputPtr = output->host<float>();
|
|
|
|
const float* gammaPtr = nullptr;
|
|
const float* betaPtr = nullptr;
|
|
|
|
if (mGammaSize > 0 && mOp->gamma() != nullptr) {
|
|
gammaPtr = mOp->gamma()->data();
|
|
}
|
|
if (mBetaSize > 0 && mOp->beta() != nullptr) {
|
|
betaPtr = mOp->beta()->data();
|
|
}
|
|
|
|
LayerNormKernel<<<mDim3Grid, mDim3Block>>>(
|
|
inputPtr, gammaPtr, betaPtr, outputPtr,
|
|
mOuterSize, mInnerSize,
|
|
static_cast<float>(mEpsilon), mGammaSize, mBetaSize
|
|
);
|
|
|
|
musaError_t err = musaGetLastError();
|
|
if (err != musaSuccess) {
|
|
return COMPUTE_NO_SUPPORT;
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
class LayerNormCreator : public Creator {
|
|
public:
|
|
virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const MNN::Op* op, Backend* backend) const override {
|
|
return new LayerNormExecution(inputs, op, backend);
|
|
}
|
|
};
|
|
|
|
MNNCreatorRegister<LayerNormCreator> gLayerNormRegistration(OpType_LayerNorm);
|
|
|
|
} // namespace MUSA
|
|
} // namespace MNN
|