wehub-resource-sync
79 lines
2.3 KiB
Plaintext
79 lines
2.3 KiB
Plaintext
#include "RangeExecution.hpp"
|
|
#include "core/MusaBackend.hpp"
|
|
|
|
namespace MNN {
|
|
namespace MUSA {
|
|
|
|
template<typename T>
|
|
__global__ void RangeKernel(T* output, T start, T delta, int size) {
|
|
int index = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (index < size) {
|
|
output[index] = start + static_cast<T>(index) * delta;
|
|
}
|
|
}
|
|
|
|
RangeExecution::RangeExecution(const std::vector<Tensor*>& inputs, const MNN::Op* op, Backend* backend)
|
|
: Execution(inputs, {}, backend) {
|
|
mBackend = static_cast<MusaBackend*>(backend);
|
|
}
|
|
|
|
ErrorCode RangeExecution::onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
auto output = outputs[0];
|
|
|
|
mSize = 1;
|
|
for (int i = 0; i < output->dimensions(); i++) {
|
|
mSize *= output->length(i);
|
|
}
|
|
|
|
int threads = 256;
|
|
int blocks = (mSize + threads - 1) / threads;
|
|
|
|
mDim3Grid = {blocks, 1, 1};
|
|
mDim3Block = {threads, 1, 1};
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode RangeExecution::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
auto output = outputs[0];
|
|
auto op = mOp->main_as_Range();
|
|
|
|
auto start = op->start();
|
|
auto limit = op->limit();
|
|
auto delta = op->delta();
|
|
|
|
// Compute size from start, limit, delta
|
|
mSize = static_cast<int>((limit - start) / delta);
|
|
|
|
// Launch kernel based on data type
|
|
if (op->type() == DataType_DT_FLOAT) {
|
|
auto outputPtr = output->host<float>();
|
|
RangeKernel<<<mDim3Grid, mDim3Block>>>(outputPtr, static_cast<float>(start), static_cast<float>(delta), mSize);
|
|
} else if (op->type() == DataType_DT_INT32) {
|
|
auto outputPtr = output->host<int>();
|
|
RangeKernel<<<mDim3Grid, mDim3Block>>>(outputPtr, static_cast<int>(start), static_cast<int>(delta), mSize);
|
|
} else {
|
|
return COMPUTE_NO_SUPPORT;
|
|
}
|
|
|
|
musaError_t err = musaGetLastError();
|
|
if (err != musaSuccess) {
|
|
return COMPUTE_NO_SUPPORT;
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
class RangeCreator : public Creator {
|
|
public:
|
|
virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const MNN::Op* op, Backend* backend) const override {
|
|
return new RangeExecution(inputs, op, backend);
|
|
}
|
|
};
|
|
|
|
MNNCreatorRegister<RangeCreator> gRangeRegistration(OpType_Range);
|
|
|
|
} // namespace MUSA
|
|
} // namespace MNN
|