wehub-resource-sync
97 lines
2.6 KiB
Plaintext
97 lines
2.6 KiB
Plaintext
#include "ArgMaxExecution.hpp"
|
|
#include "core/MusaBackend.hpp"
|
|
|
|
namespace MNN {
|
|
namespace MUSA {
|
|
|
|
template<typename T>
|
|
__global__ void ArgMaxKernel(const T* input, int* output,
|
|
int outerSize, int axisSize, int innerSize) {
|
|
int outerIdx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (outerIdx < outerSize) {
|
|
T maxVal = input[outerIdx * axisSize * innerSize];
|
|
int maxIdx = 0;
|
|
|
|
for (int i = 0; i < axisSize; i++) {
|
|
for (int j = 0; j < innerSize; j++) {
|
|
int idx = (outerIdx * axisSize + i) * innerSize + j;
|
|
if (input[idx] > maxVal) {
|
|
maxVal = input[idx];
|
|
maxIdx = i;
|
|
}
|
|
}
|
|
}
|
|
|
|
output[outerIdx] = maxIdx;
|
|
}
|
|
}
|
|
|
|
ArgMaxExecution::ArgMaxExecution(const std::vector<Tensor*>& inputs, const MNN::Op* op, Backend* backend)
|
|
: Execution(inputs, {}, backend) {
|
|
mBackend = static_cast<MusaBackend*>(backend);
|
|
mOp = op->main_as_ArgMax();
|
|
}
|
|
|
|
ErrorCode ArgMaxExecution::onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
auto input = inputs[0];
|
|
auto output = outputs[0];
|
|
|
|
mAxis = mOp->axis();
|
|
if (mAxis < 0) {
|
|
mAxis += input->dimensions();
|
|
}
|
|
|
|
mOuterSize = 1;
|
|
for (int i = 0; i < mAxis; i++) {
|
|
mOuterSize *= input->length(i);
|
|
}
|
|
|
|
mAxisSize = input->length(mAxis);
|
|
|
|
mInnerSize = 1;
|
|
for (int i = mAxis + 1; i < input->dimensions(); i++) {
|
|
mInnerSize *= input->length(i);
|
|
}
|
|
|
|
int threads = 256;
|
|
int blocks = (mOuterSize + threads - 1) / threads;
|
|
|
|
mDim3Grid = {blocks, 1, 1};
|
|
mDim3Block = {threads, 1, 1};
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode ArgMaxExecution::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<int>();
|
|
|
|
ArgMaxKernel<<<mDim3Grid, mDim3Block>>>(
|
|
inputPtr, outputPtr,
|
|
mOuterSize, mAxisSize, mInnerSize
|
|
);
|
|
|
|
musaError_t err = musaGetLastError();
|
|
if (err != musaSuccess) {
|
|
return COMPUTE_NO_SUPPORT;
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
class ArgMaxCreator : public Creator {
|
|
public:
|
|
virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const MNN::Op* op, Backend* backend) const override {
|
|
return new ArgMaxExecution(inputs, op, backend);
|
|
}
|
|
};
|
|
|
|
MNNCreatorRegister<ArgMaxCreator> gArgMaxRegistration(OpType_ArgMax);
|
|
|
|
} // namespace MUSA
|
|
} // namespace MNN
|