wehub-resource-sync
98 lines
2.8 KiB
Plaintext
98 lines
2.8 KiB
Plaintext
#include "RasterExecution.hpp"
|
|
#include "core/MusaBackend.hpp"
|
|
|
|
namespace MNN {
|
|
namespace MUSA {
|
|
|
|
template<typename T>
|
|
__global__ void RasterKernel(const T** inputs, T* output, const int* regionInfos,
|
|
int totalRegions, int totalSize) {
|
|
int index = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (index < totalSize) {
|
|
int regionIdx = 0;
|
|
int offset = index;
|
|
|
|
// Find which region this index belongs to
|
|
for (int i = 0; i < totalRegions; i++) {
|
|
int regionSize = regionInfos[i * 4 + 3];
|
|
if (offset < regionSize) {
|
|
regionIdx = i;
|
|
break;
|
|
}
|
|
offset -= regionSize;
|
|
}
|
|
|
|
int srcIdx = regionInfos[regionIdx * 4 + 0];
|
|
int srcOffset = regionInfos[regionIdx * 4 + 1];
|
|
int dstOffset = regionInfos[regionIdx * 4 + 2];
|
|
|
|
output[dstOffset + offset] = inputs[srcIdx][srcOffset + offset];
|
|
}
|
|
}
|
|
|
|
RasterExecution::RasterExecution(const std::vector<Tensor*>& inputs, const MNN::Op* op, Backend* backend)
|
|
: Execution(inputs, {}, backend) {
|
|
mBackend = static_cast<MusaBackend*>(backend);
|
|
}
|
|
|
|
ErrorCode RasterExecution::onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
auto output = outputs[0];
|
|
|
|
mTotalSize = 1;
|
|
for (int i = 0; i < output->dimensions(); i++) {
|
|
mTotalSize *= output->length(i);
|
|
}
|
|
|
|
mTotalRegions = inputs.size();
|
|
|
|
int threads = 256;
|
|
int blocks = (mTotalSize + threads - 1) / threads;
|
|
|
|
mDim3Grid = {blocks, 1, 1};
|
|
mDim3Block = {threads, 1, 1};
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode RasterExecution::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
auto output = outputs[0];
|
|
|
|
// Prepare input pointers
|
|
std::vector<float*> inputPtrs(inputs.size());
|
|
for (size_t i = 0; i < inputs.size(); i++) {
|
|
inputPtrs[i] = inputs[i]->host<float>();
|
|
}
|
|
|
|
auto outputPtr = output->host<float>();
|
|
|
|
// Simple copy for single input
|
|
if (inputs.size() == 1) {
|
|
auto inputPtr = inputs[0]->host<float>();
|
|
for (int i = 0; i < mTotalSize; i++) {
|
|
outputPtr[i] = inputPtr[i];
|
|
}
|
|
} else {
|
|
// Multiple inputs - need region info
|
|
// For now, just copy from first input
|
|
auto inputPtr = inputs[0]->host<float>();
|
|
for (int i = 0; i < mTotalSize; i++) {
|
|
outputPtr[i] = inputPtr[i];
|
|
}
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
class RasterCreator : public Creator {
|
|
public:
|
|
virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const MNN::Op* op, Backend* backend) const override {
|
|
return new RasterExecution(inputs, op, backend);
|
|
}
|
|
};
|
|
|
|
MNNCreatorRegister<RasterCreator> gRasterRegistration(OpType_Raster);
|
|
|
|
} // namespace MUSA
|
|
} // namespace MNN
|