wehub-resource-sync
133 lines
3.7 KiB
Plaintext
133 lines
3.7 KiB
Plaintext
//
|
|
// SoftmaxExecution.cu
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2026/02/25.
|
|
// Copyright © 2026, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "SoftmaxExecution.hpp"
|
|
#include "core/Macro.h"
|
|
#include "core/TensorUtils.hpp"
|
|
#include "backend/musa/core/MusaBackend.hpp"
|
|
#include <musa_runtime.h>
|
|
|
|
namespace MNN {
|
|
namespace MUSA {
|
|
|
|
// MUSA kernel for softmax operation
|
|
__global__ void SoftmaxKernel(const float* input, float* output, int outerCount, int depth, int innerCount) {
|
|
int outerIdx = blockIdx.x;
|
|
int innerIdx = threadIdx.x + blockIdx.y * blockDim.x;
|
|
|
|
if (outerIdx >= outerCount || innerIdx >= innerCount) return;
|
|
|
|
const float* inPtr = input + outerIdx * depth * innerCount;
|
|
float* outPtr = output + outerIdx * depth * innerCount;
|
|
|
|
// Find max value for numerical stability
|
|
float maxVal = -FLT_MAX;
|
|
for (int i = 0; i < depth; i++) {
|
|
float val = inPtr[i * innerCount + innerIdx];
|
|
if (val > maxVal) {
|
|
maxVal = val;
|
|
}
|
|
}
|
|
|
|
// Compute exp and sum
|
|
float sum = 0.0f;
|
|
for (int i = 0; i < depth; i++) {
|
|
float expVal = expf(inPtr[i * innerCount + innerIdx] - maxVal);
|
|
outPtr[i * innerCount + innerIdx] = expVal;
|
|
sum += expVal;
|
|
}
|
|
|
|
// Normalize
|
|
float invSum = 1.0f / sum;
|
|
for (int i = 0; i < depth; i++) {
|
|
outPtr[i * innerCount + innerIdx] *= invSum;
|
|
}
|
|
}
|
|
|
|
SoftmaxExecution::SoftmaxExecution(int axis, Backend* backend) : Execution(backend) {
|
|
auto musaBackend = static_cast<MusaBackend*>(backend);
|
|
mRuntime = musaBackend->getMusaRuntime();
|
|
mAxis = axis;
|
|
}
|
|
|
|
ErrorCode SoftmaxExecution::onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
auto input = inputs[0];
|
|
auto shape = input->shape();
|
|
int dims = shape.size();
|
|
|
|
if (mAxis < 0) {
|
|
mAxis = dims + mAxis;
|
|
}
|
|
|
|
mOuterCount = 1;
|
|
for (int i = 0; i < mAxis; i++) {
|
|
mOuterCount *= shape[i];
|
|
}
|
|
|
|
mDepth = shape[mAxis];
|
|
|
|
mInnerCount = 1;
|
|
for (int i = mAxis + 1; i < dims; i++) {
|
|
mInnerCount *= shape[i];
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode SoftmaxExecution::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
#ifdef LOG_VERBOSE
|
|
MNN_PRINT("start SoftmaxExecution onExecute...\n");
|
|
#endif
|
|
|
|
auto input = inputs[0]->deviceId();
|
|
auto output = outputs[0]->deviceId();
|
|
|
|
int threadsPerBlock = 256;
|
|
dim3 blockDim(threadsPerBlock);
|
|
dim3 gridDim(mOuterCount, (mInnerCount + threadsPerBlock - 1) / threadsPerBlock);
|
|
|
|
SoftmaxKernel<<<gridDim, blockDim>>>((const float*)input, (float*)output, mOuterCount, mDepth, mInnerCount);
|
|
|
|
// Check for kernel launch errors
|
|
musaError_t err = musaGetLastError();
|
|
if (err != musaSuccess) {
|
|
MNN_ERROR("MUSA Softmax kernel launch failed: %s\n", musaGetErrorString(err));
|
|
}
|
|
|
|
// Synchronize to ensure completion
|
|
mRuntime->device_sync();
|
|
|
|
#ifdef LOG_VERBOSE
|
|
MNN_PRINT("end SoftmaxExecution onExecute...\n");
|
|
#endif
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
// Creator for Softmax operations
|
|
class SoftmaxCreator : 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 {
|
|
int axis = 1;
|
|
if (op->type() == OpType_Softmax) {
|
|
auto softmax = op->main_as_Softmax();
|
|
if (softmax != nullptr && softmax->axis() != -1) {
|
|
axis = softmax->axis();
|
|
}
|
|
return new SoftmaxExecution(axis, backend);
|
|
}
|
|
return nullptr;
|
|
}
|
|
};
|
|
|
|
MusaCreatorRegister<SoftmaxCreator> __SoftmaxExecution(OpType_Softmax);
|
|
|
|
} // namespace MUSA
|
|
} // namespace MNN
|