wehub-resource-sync
143 lines
4.6 KiB
Plaintext
143 lines
4.6 KiB
Plaintext
//
|
|
// BatchNormExecution.cu
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2026/02/25.
|
|
// Copyright © 2026, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "core/MusaBackend.hpp"
|
|
#include "core/TensorUtils.hpp"
|
|
#include "MNN_generated.h"
|
|
#include <musa_runtime.h>
|
|
|
|
namespace MNN {
|
|
namespace MUSA {
|
|
|
|
// MUSA kernel for batch normalization
|
|
__global__ void BatchNormKernel(const float* input, float* output,
|
|
const float* scale, const float* bias,
|
|
const float* mean, const float* variance,
|
|
float epsilon, int batchSize, int channels, int spatialSize) {
|
|
int c = blockIdx.x * blockDim.x + threadIdx.x;
|
|
int s = blockIdx.y * blockDim.y + threadIdx.y;
|
|
|
|
if (c >= channels || s >= spatialSize) return;
|
|
|
|
float invStd = 1.0f / sqrtf(variance[c] + epsilon);
|
|
float m = mean[c];
|
|
float b = bias[c];
|
|
float s_val = scale[c];
|
|
|
|
for (int b = 0; b < batchSize; ++b) {
|
|
int idx = (b * channels + c) * spatialSize + s;
|
|
float x = input[idx];
|
|
float y = (x - m) * invStd * s_val + b;
|
|
output[idx] = y;
|
|
}
|
|
}
|
|
|
|
class BatchNormExecution : public Execution {
|
|
public:
|
|
BatchNormExecution(Backend* backend) : Execution(backend) {}
|
|
|
|
virtual ErrorCode onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) override {
|
|
auto input = inputs[0];
|
|
auto inputShape = input->shape();
|
|
|
|
mBatchSize = inputShape[0];
|
|
mChannels = inputShape[1];
|
|
mSpatialSize = 1;
|
|
for (size_t i = 2; i < inputShape.size(); ++i) {
|
|
mSpatialSize *= inputShape[i];
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
virtual ErrorCode onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) override {
|
|
#ifdef LOG_VERBOSE
|
|
MNN_PRINT("start BatchNormExecution onExecute...\n");
|
|
#endif
|
|
|
|
auto input = inputs[0];
|
|
auto output = outputs[0];
|
|
auto op = this->op();
|
|
|
|
auto batchNorm = op->main_as_BatchNorm();
|
|
|
|
void* inputPtr = (void*)input->deviceId();
|
|
void* outputPtr = (void*)output->deviceId();
|
|
|
|
// Get scale, bias, mean, variance from the op
|
|
auto scaleData = batchNorm->scaleData();
|
|
auto biasData = batchNorm->biasData();
|
|
auto meanData = batchNorm->meanData();
|
|
auto varianceData = batchNorm->varianceData();
|
|
float epsilon = batchNorm->eps();
|
|
|
|
// Copy parameters to device
|
|
float *dScale, *dBias, *dMean, *dVariance;
|
|
size_t dataSize = sizeof(float) * mChannels;
|
|
|
|
musaMalloc(&dScale, dataSize);
|
|
musaMalloc(&dBias, dataSize);
|
|
musaMalloc(&dMean, dataSize);
|
|
musaMalloc(&dVariance, dataSize);
|
|
|
|
musaMemcpy(dScale, scaleData->data(), dataSize, MNNMemcpyHostToDevice);
|
|
musaMemcpy(dBias, biasData->data(), dataSize, MNNMemcpyHostToDevice);
|
|
musaMemcpy(dMean, meanData->data(), dataSize, MNNMemcpyHostToDevice);
|
|
musaMemcpy(dVariance, varianceData->data(), dataSize, MNNMemcpyHostToDevice);
|
|
|
|
dim3 threadsPerBlock(16, 16);
|
|
dim3 blocksPerGrid((mChannels + 15) / 16, (mSpatialSize + 15) / 16);
|
|
|
|
BatchNormKernel<<<blocksPerGrid, threadsPerBlock>>>(
|
|
(const float*)inputPtr, (float*)outputPtr,
|
|
dScale, dBias, dMean, dVariance,
|
|
epsilon, mBatchSize, mChannels, mSpatialSize);
|
|
|
|
// Check for kernel launch errors
|
|
musaError_t err = musaGetLastError();
|
|
if (err != musaSuccess) {
|
|
MNN_ERROR("MUSA BatchNorm kernel launch failed: %s\n", musaGetErrorString(err));
|
|
}
|
|
|
|
// Synchronize to ensure completion
|
|
auto musaBackend = static_cast<MusaBackend*>(backend());
|
|
musaBackend->getMusaRuntime()->device_sync();
|
|
|
|
// Free temporary device memory
|
|
musaFree(dScale);
|
|
musaFree(dBias);
|
|
musaFree(dMean);
|
|
musaFree(dVariance);
|
|
|
|
#ifdef LOG_VERBOSE
|
|
MNN_PRINT("end BatchNormExecution onExecute...\n");
|
|
#endif
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
private:
|
|
int mBatchSize;
|
|
int mChannels;
|
|
int mSpatialSize;
|
|
};
|
|
|
|
// Creator for BatchNorm operations
|
|
class BatchNormCreator : public MusaBackend::Creator {
|
|
public:
|
|
virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
|
const MNN::Op* op, Backend* backend) const override {
|
|
return new BatchNormExecution(backend);
|
|
}
|
|
};
|
|
|
|
MusaCreatorRegister<BatchNormCreator> __BatchNormExecution(OpType_BatchNorm);
|
|
|
|
} // namespace MUSA
|
|
} // namespace MNN
|