wehub-resource-sync
150 lines
4.8 KiB
Plaintext
150 lines
4.8 KiB
Plaintext
//
|
|
// SplitExecution.cu
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2026/02/25.
|
|
// Copyright © 2026, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "core/MusaBackend.hpp"
|
|
#include "core/TensorUtils.hpp"
|
|
#include "MNN_generated.h"
|
|
#include <musa_runtime.h>
|
|
#include <vector>
|
|
|
|
namespace MNN {
|
|
namespace MUSA {
|
|
|
|
// MUSA kernel for split operation
|
|
__global__ void SplitKernel(const float* input, float** outputs,
|
|
const int* outputOffsets, int numOutputs,
|
|
int splitSize, int outerSize, int innerSize) {
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
int totalSize = outerSize * splitSize * innerSize;
|
|
|
|
if (idx >= totalSize) return;
|
|
|
|
int innerIdx = idx % innerSize;
|
|
int tempIdx = idx / innerSize;
|
|
int splitIdx = tempIdx % splitSize;
|
|
int outerIdx = tempIdx / splitSize;
|
|
|
|
// Find which output tensor this element belongs to
|
|
int outputIdx = 0;
|
|
int localSplitIdx = splitIdx;
|
|
for (int i = 0; i < numOutputs - 1; ++i) {
|
|
int outputSize = outputOffsets[i + 1] - outputOffsets[i];
|
|
if (localSplitIdx < outputSize) {
|
|
break;
|
|
}
|
|
localSplitIdx -= outputSize;
|
|
outputIdx++;
|
|
}
|
|
|
|
int srcIdx = idx;
|
|
int dstIdx = (outerIdx * (outputOffsets[outputIdx + 1] - outputOffsets[outputIdx]) + localSplitIdx) * innerSize + innerIdx;
|
|
|
|
outputs[outputIdx][dstIdx] = input[srcIdx];
|
|
}
|
|
|
|
class SplitExecution : public Execution {
|
|
public:
|
|
SplitExecution(int axis, Backend* backend) : Execution(backend), mAxis(axis) {}
|
|
|
|
virtual ErrorCode onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) override {
|
|
mOutputs.resize(outputs.size());
|
|
mOutputSizes.resize(outputs.size());
|
|
|
|
for (size_t i = 0; i < outputs.size(); ++i) {
|
|
mOutputs[i] = outputs[i];
|
|
mOutputSizes[i] = outputs[i]->length(mAxis);
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
virtual ErrorCode onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) override {
|
|
#ifdef LOG_VERBOSE
|
|
MNN_PRINT("start SplitExecution onExecute...\n");
|
|
#endif
|
|
|
|
auto input = inputs[0];
|
|
|
|
// Collect output device pointers
|
|
std::vector<void*> outputPtrs(outputs.size());
|
|
for (size_t i = 0; i < outputs.size(); ++i) {
|
|
outputPtrs[i] = (void*)outputs[i]->deviceId();
|
|
}
|
|
void* inputPtr = (void*)input->deviceId();
|
|
|
|
// Copy device pointers to device memory
|
|
float** dOutputs = nullptr;
|
|
int* dOutputOffsets = nullptr;
|
|
size_t ptrSize = sizeof(float*) * outputs.size();
|
|
size_t sizeSize = sizeof(int) * (outputs.size() + 1);
|
|
|
|
musaMalloc(&dOutputs, ptrSize);
|
|
musaMalloc(&dOutputOffsets, sizeSize);
|
|
|
|
musaMemcpy(dOutputs, outputPtrs.data(), ptrSize, MNNMemcpyHostToDevice);
|
|
|
|
// Calculate output offsets
|
|
std::vector<int> outputOffsets(outputs.size() + 1, 0);
|
|
for (size_t i = 0; i < outputs.size(); ++i) {
|
|
outputOffsets[i + 1] = outputOffsets[i] + mOutputSizes[i];
|
|
}
|
|
musaMemcpy(dOutputOffsets, outputOffsets.data(), sizeSize, MNNMemcpyHostToDevice);
|
|
|
|
int totalSize = input->elementSize();
|
|
dim3 threadsPerBlock(256);
|
|
dim3 blocksPerGrid((totalSize + 255) / 256);
|
|
|
|
SplitKernel<<<blocksPerGrid, threadsPerBlock>>>(
|
|
(const float*)inputPtr, dOutputs, dOutputOffsets, outputs.size(),
|
|
input->length(mAxis), 1, totalSize / input->length(mAxis));
|
|
|
|
// Check for kernel launch errors
|
|
musaError_t err = musaGetLastError();
|
|
if (err != musaSuccess) {
|
|
MNN_ERROR("MUSA Split kernel launch failed: %s\n", musaGetErrorString(err));
|
|
}
|
|
|
|
// Synchronize to ensure completion
|
|
auto musaBackend = static_cast<MusaBackend*>(backend());
|
|
musaBackend->getMusaRuntime()->device_sync();
|
|
|
|
// Free temporary device memory
|
|
musaFree(dOutputs);
|
|
musaFree(dOutputOffsets);
|
|
|
|
#ifdef LOG_VERBOSE
|
|
MNN_PRINT("end SplitExecution onExecute...\n");
|
|
#endif
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
private:
|
|
int mAxis;
|
|
std::vector<Tensor*> mOutputs;
|
|
std::vector<int> mOutputSizes;
|
|
};
|
|
|
|
// Creator for Split operations
|
|
class SplitCreator : 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 = 0;
|
|
if (op->type() == OpType_Split) {
|
|
axis = op->main_as_Split()->axis();
|
|
}
|
|
return new SplitExecution(axis, backend);
|
|
}
|
|
};
|
|
|
|
MusaCreatorRegister<SplitCreator> __SplitExecution(OpType_Split);
|
|
|
|
} // namespace MUSA
|
|
} // namespace MNN
|