wehub-resource-sync
169 lines
8.2 KiB
C++
169 lines
8.2 KiB
C++
#include "KleidiAIConvolutionDepthwise.hpp"
|
|
|
|
#ifdef MNN_KLEIDIAI_ENABLED
|
|
|
|
#include <string.h>
|
|
#include "core/Concurrency.h"
|
|
#include "backend/cpu/compute/Int8FunctionsOpt.h"
|
|
#include "core/Macro.h"
|
|
#include "core/TensorUtils.hpp"
|
|
#include "backend/cpu/compute/CommonOptFunction.h"
|
|
#include "backend/cpu/compute/ConvOpt.h"
|
|
|
|
namespace MNN {
|
|
template<typename T>
|
|
void nchw_to_nhwc_optimized(const T* src, T* dst,
|
|
int batch, int channel, int height, int width) {
|
|
const int hw = height * width;
|
|
const int chw = channel * hw;
|
|
const int wc = width * channel;
|
|
|
|
for (int n = 0; n < batch; ++n) {
|
|
const T* src_batch = src + n * chw;
|
|
T* dst_batch = dst + n * chw;
|
|
|
|
for (int c = 0; c < channel; ++c) {
|
|
const T* src_channel = src_batch + c * hw;
|
|
|
|
for (int h = 0; h < height; ++h) {
|
|
const T* src_row = src_channel + h * width;
|
|
T* dst_row = dst_batch + h * wc + c;
|
|
|
|
for (int w = 0; w < width; ++w) {
|
|
dst_row[w * channel] = src_row[w];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
KleidiAIConvolutionDepthwise::KleidiAIDepthwiseExecution::KleidiAIDepthwiseExecution(const Convolution2DCommon* common, Backend* b,
|
|
const float* originWeight, size_t originWeightSize,
|
|
const float* bias, size_t biasSize)
|
|
: MNN::CPUConvolution(common, b) {
|
|
int kernel_height = common->kernelY();
|
|
int kernel_width = common->kernelX();
|
|
int channels = common->outputCount();
|
|
int packedRhsSize = kai_rhs_get_dst_size_dwconv_pack_x32p1vlx1b_x32_x32_sme(kernel_height, kernel_width, channels);
|
|
mPackedRhs.reset(Tensor::createDevice<uint8_t>(std::vector<int>{packedRhsSize}));
|
|
bool success = b->onAcquireBuffer(mPackedRhs.get(), Backend::STATIC);
|
|
if (!success) {
|
|
MNN_ERROR("Error for alloc memory for CPUConvolutionDepthwise\n");
|
|
mValid = false;
|
|
return;
|
|
}
|
|
mNumber = ((CPUBackend*)b)->threadNumber();
|
|
mWeightTemp.reset(Tensor::createDevice<uint8_t>(std::vector<int>{channels * kernel_height * kernel_width * (int)sizeof(float)}));
|
|
success = b->onAcquireBuffer(mWeightTemp.get(), Backend::STATIC);
|
|
if (!success) {
|
|
MNN_ERROR("Error for alloc memory for CPUConvolutionDepthwise\n");
|
|
mValid = false;
|
|
return;
|
|
}
|
|
auto weightTempPtr = mWeightTemp->host<float>();
|
|
nchw_to_nhwc_optimized(originWeight, weightTempPtr, 1, channels, kernel_height, kernel_width);
|
|
kai_run_rhs_dwconv_pack_x32p1vlx1b_x32_x32_sme(kernel_height, kernel_width, kernel_height, kernel_width, channels, weightTempPtr, bias, mPackedRhs.get()->host<void>());
|
|
b->onReleaseBuffer(mWeightTemp.get(), Backend::STATIC);
|
|
}
|
|
|
|
ErrorCode KleidiAIConvolutionDepthwise::KleidiAIDepthwiseExecution::onResize(const std::vector<Tensor*>& inputs,
|
|
const std::vector<Tensor*>& outputs) {
|
|
CPUConvolution::onResize(inputs, outputs);
|
|
auto input = inputs[0];
|
|
auto output = outputs[0];
|
|
TensorUtils::getDescribe(&mOutputNHWC)->dimensionFormat = MNN_DATA_FORMAT_NHWC;
|
|
mOutputNHWC.buffer().dimensions = 4;
|
|
mOutputNHWC.buffer().dim[0].extent = output->batch();
|
|
mOutputNHWC.buffer().dim[1].extent = output->height();
|
|
mOutputNHWC.buffer().dim[2].extent = output->width();
|
|
mOutputNHWC.buffer().dim[3].extent = output->channel();
|
|
mOutputNHWC.buffer().type = output->getType();
|
|
auto success = backend()->onAcquireBuffer(&mOutputNHWC, Backend::DYNAMIC);
|
|
if (!success) {
|
|
return OUT_OF_MEMORY;
|
|
}
|
|
|
|
TensorUtils::getDescribe(&mInputNHWC)->dimensionFormat = MNN_DATA_FORMAT_NHWC;
|
|
mInputNHWC.buffer().dimensions = 4;
|
|
mInputNHWC.buffer().dim[0].extent = input->batch();
|
|
mInputNHWC.buffer().dim[1].extent = input->height();
|
|
mInputNHWC.buffer().dim[2].extent = input->width();
|
|
mInputNHWC.buffer().dim[3].extent = input->channel();
|
|
mInputNHWC.buffer().type = input->getType();
|
|
success = backend()->onAcquireBuffer(&mInputNHWC, Backend::DYNAMIC);
|
|
if (!success) {
|
|
return OUT_OF_MEMORY;
|
|
}
|
|
|
|
backend()->onReleaseBuffer(&mOutputNHWC, Backend::DYNAMIC);
|
|
backend()->onReleaseBuffer(&mInputNHWC, Backend::DYNAMIC);
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode KleidiAIConvolutionDepthwise::KleidiAIDepthwiseExecution::onExecute(const std::vector<Tensor*>& inputs,
|
|
const std::vector<Tensor*>& outputs) {
|
|
auto inputTensor = inputs[0];
|
|
auto outputTensor = outputs[0];
|
|
const auto srcOrigin = mInputNHWC.host<uint8_t>();
|
|
auto dstOrigin = mOutputNHWC.host<uint8_t>();
|
|
auto postData = getPostParameters();
|
|
auto output_height = outputTensor->height();
|
|
auto core = static_cast<CPUBackend*>(backend())->functions();
|
|
auto batch = inputTensor->batch();
|
|
|
|
MNN_CONCURRENCY_BEGIN(tId, mNumber) {
|
|
CPUTensorConverter::convert(inputTensor, &mInputNHWC, core, tId, mNumber);
|
|
}
|
|
MNN_CONCURRENCY_END();
|
|
|
|
//CPUTensorConverter::convert(inputTensor, &mInputNHWC, core);
|
|
|
|
constexpr size_t rows_handled = 4; // no of rows kernel handles each time.
|
|
for(size_t b = 0; b < batch; b++) {
|
|
const auto srcOriginBatch = srcOrigin + b * inputTensor->height() * inputTensor->width() * inputTensor->channel() * sizeof(float);
|
|
auto dstOriginBatch = dstOrigin + b * outputTensor->height() * outputTensor->width() * outputTensor->channel() * sizeof(float);
|
|
for (size_t out_row = 0; out_row < output_height; out_row += rows_handled) {
|
|
// Variables below used to calculate start of input pointer.
|
|
const int start_in_row = out_row - mPadY;
|
|
const size_t pad_top = (start_in_row < 0) ? (-start_in_row) : 0;
|
|
const size_t in_row = (start_in_row < 0) ? 0 : start_in_row;
|
|
|
|
// Calculate row strides for pointer.
|
|
const size_t in_row_stride_bytes = (inputTensor->width() * inputTensor->channel() * sizeof(float));
|
|
const size_t out_row_stride_bytes = (outputTensor->width() * outputTensor->channel() * sizeof(float));
|
|
|
|
// Number of input rows that can be read, number of output rows to calculate.
|
|
const size_t valid_input_rows = (in_row < inputTensor->height()) ? (inputTensor->height() - in_row) : 0;
|
|
const size_t valid_out_rows = (outputTensor->height() - out_row);
|
|
|
|
// Increment output/input pointers according to tile being calculated.
|
|
auto out_offset = kai_get_dst_offset_dwconv_clamp_f32_f32_f32p1vlx1b_3x3_s1_4xc_sme2_mla(
|
|
out_row, out_row_stride_bytes);
|
|
auto in_offset = kai_get_src_offset_dwconv_clamp_f32_f32_f32p1vlx1b_3x3_s1_4xc_sme2_mla(
|
|
in_row, in_row_stride_bytes);
|
|
const auto inptr = (uint8_t*)srcOriginBatch + in_offset;
|
|
auto outptr = (uint8_t*)dstOriginBatch + out_offset;
|
|
|
|
// NOTE: Kernel expects strides to be passed as bytes.
|
|
// f32_f32_f32p1vlx1b -> f32 output, f32 LHS, packed F32 rhs (with bias) as 1VL blocks.
|
|
// 3x3_s : 3x3 filter with stride 1
|
|
// 4xc : 4 rows across all output channels (plane c) is produced.
|
|
kai_run_dwconv_clamp_f32_f32_f32p1vlx1b_3x3_s1_4xc_sme2_mla(
|
|
inptr, mPackedRhs.get()->host<void>(), outptr, in_row_stride_bytes, inputTensor->channel() * sizeof(float),
|
|
out_row_stride_bytes, outputTensor->channel() * sizeof(float), valid_input_rows, valid_out_rows,
|
|
mPadX, pad_top, 0.0f, postData[2], postData[3]);
|
|
}
|
|
}
|
|
|
|
MNN_CONCURRENCY_BEGIN(tId, mNumber) {
|
|
CPUTensorConverter::convert(&mOutputNHWC, outputTensor, core, tId, mNumber);
|
|
}
|
|
MNN_CONCURRENCY_END();
|
|
//CPUTensorConverter::convert(&mOutputNHWC, outputTensor, core);
|
|
return NO_ERROR;
|
|
}
|
|
|
|
} // namespace MNN
|
|
|
|
#endif // defined(MNN_KLEIDIAI_ENABLED)
|