wehub-resource-sync
108 lines
3.1 KiB
Plaintext
108 lines
3.1 KiB
Plaintext
#include "TopKV2Execution.hpp"
|
|
#include "core/MusaBackend.hpp"
|
|
|
|
namespace MNN {
|
|
namespace MUSA {
|
|
|
|
template<typename T>
|
|
__global__ void TopKKernel(const T* input, T* outValues, int* outIndices,
|
|
int outerSize, int k, int innerSize) {
|
|
int outerIdx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (outerIdx < outerSize) {
|
|
const T* inputPtr = input + outerIdx * k * innerSize;
|
|
T* outValPtr = outValues + outerIdx * k * innerSize;
|
|
int* outIdxPtr = outIndices + outerIdx * k * innerSize;
|
|
|
|
// Simple selection sort for top k
|
|
for (int i = 0; i < k; i++) {
|
|
T maxVal = inputPtr[i * innerSize];
|
|
int maxIdx = i;
|
|
|
|
for (int j = i + 1; j < innerSize; j++) {
|
|
if (inputPtr[j] > maxVal) {
|
|
maxVal = inputPtr[j];
|
|
maxIdx = j;
|
|
}
|
|
}
|
|
|
|
// Swap
|
|
if (maxIdx != i) {
|
|
T tempVal = inputPtr[i * innerSize];
|
|
inputPtr[i * innerSize] = maxVal;
|
|
inputPtr[maxIdx * innerSize] = tempVal;
|
|
}
|
|
|
|
outValPtr[i * innerSize] = maxVal;
|
|
outIdxPtr[i * innerSize] = maxIdx;
|
|
}
|
|
}
|
|
}
|
|
|
|
TopKV2Execution::TopKV2Execution(const std::vector<Tensor*>& inputs, const MNN::Op* op, Backend* backend)
|
|
: Execution(inputs, {}, backend) {
|
|
mBackend = static_cast<MusaBackend*>(backend);
|
|
mOp = op->main_as_TopKV2();
|
|
}
|
|
|
|
ErrorCode TopKV2Execution::onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
auto input = inputs[0];
|
|
auto kTensor = inputs[1];
|
|
|
|
mAxis = mOp->axis();
|
|
if (mAxis < 0) {
|
|
mAxis += input->dimensions();
|
|
}
|
|
|
|
mK = kTensor->host<int>()[0];
|
|
|
|
mOuterSize = 1;
|
|
for (int i = 0; i < mAxis; i++) {
|
|
mOuterSize *= input->length(i);
|
|
}
|
|
|
|
mInnerSize = input->length(mAxis);
|
|
|
|
int threads = 256;
|
|
int blocks = (mOuterSize + threads - 1) / threads;
|
|
|
|
mDim3Grid = {blocks, 1, 1};
|
|
mDim3Block = {threads, 1, 1};
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode TopKV2Execution::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
auto input = inputs[0];
|
|
auto outputValues = outputs[0];
|
|
auto outputIndices = outputs[1];
|
|
|
|
auto inputPtr = input->host<float>();
|
|
auto outputValuesPtr = outputValues->host<float>();
|
|
auto outputIndicesPtr = outputIndices->host<int>();
|
|
|
|
TopKKernel<<<mDim3Grid, mDim3Block>>>(
|
|
inputPtr, outputValuesPtr, outputIndicesPtr,
|
|
mOuterSize, mK, mInnerSize
|
|
);
|
|
|
|
musaError_t err = musaGetLastError();
|
|
if (err != musaSuccess) {
|
|
return COMPUTE_NO_SUPPORT;
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
class TopKV2Creator : public Creator {
|
|
public:
|
|
virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const MNN::Op* op, Backend* backend) const override {
|
|
return new TopKV2Execution(inputs, op, backend);
|
|
}
|
|
};
|
|
|
|
MNNCreatorRegister<TopKV2Creator> gTopKV2Registration(OpType_TopKV2);
|
|
|
|
} // namespace MUSA
|
|
} // namespace MNN
|