timekiller
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

add broadcast feature for matmul op

This commit is contained in:
wangyanling 2021-10-22 11:33:42 +08:00
parent d3aed289bf
commit 3a0a1c6b72
19 changed files with 484 additions and 106 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/matmul_infer.c
View File

@ -25,7 +25,9 @@ int CheckMatmulInputShape(int *a_shape, size_t a_shape_size, int *b_shape, size_
return NNACL_PARAM_INVALID;
}
for (size_t i = 0; i < (a_shape_size - 2) && i < (b_shape_size - 2); ++i) {
if (a_shape[i] != b_shape[i]) {
int min_value = MSMIN(a_shape[i], b_shape[i]);
int max_value = MSMAX(a_shape[i], b_shape[i]);
if (max_value % min_value != 0) {
return NNACL_INPUT_TENSOR_ERROR;
}
}
@ -47,6 +49,17 @@ int CheckMatmulInputShape(int *a_shape, size_t a_shape_size, int *b_shape, size_
return NNACL_OK;
}
bool BroadcastInfer(int *a_shape, size_t a_shape_size, int *b_shape, size_t b_shape_size) {
for (size_t i = 0; i < (a_shape_size - 2) && i < (b_shape_size - 2); ++i) {
int min_value = MSMIN(a_shape[i], b_shape[i]);
int max_value = MSMAX(a_shape[i], b_shape[i]);
if (a_shape[i] != b_shape[i] && max_value % min_value == 0) {
return true;
}
}
return false;
}
int MatmulInferShape(const TensorC *const *inputs, size_t inputs_size, TensorC **outputs, size_t outputs_size,
OpParameter *parameter) {
int check_ret = CheckAugmentNullSizeInputTwo(inputs, inputs_size, outputs, outputs_size, parameter, 2, 3, 1);
@ -110,6 +123,12 @@ int MatmulInferShape(const TensorC *const *inputs, size_t inputs_size, TensorC *
if (del_end) {
c_shape_size--;
}
if (BroadcastInfer(a_shape, a_shape_size, b_shape, b_shape_size)) {
for (size_t i = 0; i < (a_shape_size - 2) && i < (b_shape_size - 2); ++i) {
c_shape[i] = MSMAX(a_shape[i], b_shape[i]);
}
}
SetShapeArray(output, c_shape, c_shape_size);
return NNACL_OK;
}

8
mindspore/lite/README.md
View File

@ -57,12 +57,12 @@ For more details please check out our [MindSpore Lite Architecture Guide](https:
## MindSpore Lite benchmark test result
We test a couple of networks on HUAWEI Mate40 (Hisilicon Kirin9000) mobile phone, and get the test results below for your reference.
We test a couple of networks on HUAWEI Mate40 (Hisilicon Kirin9000e) mobile phone, and get the test results below for your reference.
| NetWork | Thread Number | Average Run Time(ms) |
| ------------------- | ------------- | -------------------- |
| basic_squeezenet | 4 | 7.246 |
| basic_squeezenet | 4 | 6.415 |
| inception_v3 | 4 | 36.767 |
| mobilenet_v1_10_224 | 4 | 5.187 |
| mobilenet_v2_10_224 | 4 | 4.153 |
| mobilenet_v1_10_224 | 4 | 4.936 |
| mobilenet_v2_10_224 | 4 | 3.644 |
| resnet_v2_50 | 4 | 25.071 |

8
mindspore/lite/README_CN.md
View File

@ -65,12 +65,12 @@ MindSpore Lite是MindSpore推出的端云协同的、轻量化、高性能AI推
## MindSpore Lite性能参考数据
我们在HUAWEI Mate40（Hisilicon Kirin9000）手机上，测试了一组端侧常见网络的性能数据，供您参考：
我们在HUAWEI Mate40（Hisilicon Kirin9000e）手机上，测试了一组端侧常见网络的性能数据，供您参考：
| 网络 | 线程数 | 平均推理时间(毫秒) |
| ------------------- | ----- | --------------- |
| basic_squeezenet | 4 | 7.246 |
| basic_squeezenet | 4 | 6.415 |
| inception_v3 | 4 | 36.767 |
| mobilenet_v1_10_224 | 4 | 5.187 |
| mobilenet_v2_10_224 | 4 | 4.153 |
| mobilenet_v1_10_224 | 4 | 4.936 |
| mobilenet_v2_10_224 | 4 | 3.644 |
| resnet_v2_50 | 4 | 25.071 |

2
mindspore/lite/src/common/common.h
View File

@ -31,6 +31,8 @@ enum CHWK_SHAPE { CHWK_C = 0, CHWK_H = 1, CHWK_W = 2, CHWK_K = 3 };
enum KHWC_SHAPE { KHWC_K = 0, KHWC_H = 1, KHWC_W = 2, KHWC_C = 3 };
enum CHW_SHAPE { CHW_C = 0, CHW_H = 1, CHW_W = 2 };
enum HWC_SHAPE { HWC_H = 0, HWC_W = 1, HWC_C = 2 };
static constexpr int kHWDimNumber = 2;
static constexpr int kCHWDimNumber = 3;
static constexpr int kNCHWDimNumber = 4;
static constexpr int kNHWCDimNumber = 4;

2
mindspore/lite/src/runtime/kernel/arm/fp16/fullconnection_fp16.cc
View File

@ -56,6 +56,8 @@ int FullconnectionFP16CPUKernel::Prepare() {
row_tile_ = C12NUM;
#endif
params_->batch = 1;
a_batch_ = 1;
b_batch_ = 1;
params_->a_transpose_ = false;
params_->b_transpose_ = true;

103
mindspore/lite/src/runtime/kernel/arm/fp16/matmul_base_fp16.cc
View File

@ -15,10 +15,13 @@
*/
#include "src/runtime/kernel/arm/fp16/matmul_base_fp16.h"
#include <algorithm>
#include "nnacl/fp16/matmul_fp16.h"
#include "nnacl/fp16/cast_fp16.h"
#include "include/errorcode.h"
using mindspore::lite::kCHWDimNumber;
using mindspore::lite::kHWDimNumber;
using mindspore::lite::RET_ERROR;
using mindspore::lite::RET_INPUT_TENSOR_ERROR;
using mindspore::lite::RET_MEMORY_FAILED;
@ -112,6 +115,7 @@ int MatmulBaseFP16CPUKernel::ReSize() {
thread_count_ = MSMIN(op_parameter_->thread_num_, UP_DIV(params_->col_, C8NUM));
thread_stride_ = UP_DIV(UP_DIV(params_->col_, C8NUM), thread_count_) * C8NUM;
}
return RET_OK;
}
@ -137,12 +141,12 @@ void MatmulBaseFP16CPUKernel::ResizeParameter() {
int MatmulBaseFP16CPUKernel::InitBufferA() {
a_pack_ptr_ = reinterpret_cast<float16_t *>(
ms_context_->allocator->Malloc(params_->batch * params_->row_align_ * params_->deep_ * sizeof(float16_t)));
ms_context_->allocator->Malloc(a_batch_ * params_->row_align_ * params_->deep_ * sizeof(float16_t)));
if (a_pack_ptr_ == nullptr) {
return RET_MEMORY_FAILED;
}
memset(a_pack_ptr_, 0, params_->batch * params_->row_align_ * params_->deep_ * sizeof(float16_t));
memset(a_pack_ptr_, 0, a_batch_ * params_->row_align_ * params_->deep_ * sizeof(float16_t));
return RET_OK;
}
@ -152,12 +156,12 @@ int MatmulBaseFP16CPUKernel::InitBufferB() {
}
b_pack_ptr_ = reinterpret_cast<float16_t *>(
ms_context_->allocator->Malloc(params_->batch * params_->col_align_ * params_->deep_ * sizeof(float16_t)));
ms_context_->allocator->Malloc(b_batch_ * params_->col_align_ * params_->deep_ * sizeof(float16_t)));
if (b_pack_ptr_ == nullptr) {
return RET_MEMORY_FAILED;
}
memset(b_pack_ptr_, 0, params_->batch * params_->col_align_ * params_->deep_ * sizeof(float16_t));
memset(b_pack_ptr_, 0, b_batch_ * params_->col_align_ * params_->deep_ * sizeof(float16_t));
return RET_OK;
}
@ -167,15 +171,15 @@ void MatmulBaseFP16CPUKernel::InitMatrixA(const void *src_ptr) {
if (vec_matmul_) {
if (src_data_type == kNumberTypeFloat32) {
Float32ToFloat16(reinterpret_cast<const float *>(src_ptr), a_pack_ptr_, params_->batch * params_->deep_);
Float32ToFloat16(reinterpret_cast<const float *>(src_ptr), a_pack_ptr_, a_batch_ * params_->deep_);
} else {
memcpy(a_pack_ptr_, src_ptr, params_->batch * params_->deep_ * sizeof(float16_t));
memcpy(a_pack_ptr_, src_ptr, a_batch_ * params_->deep_ * sizeof(float16_t));
}
return;
}
const int8_t *int8_src = reinterpret_cast<const int8_t *>(src_ptr);
for (int i = 0; i < params_->batch; i++) {
for (int i = 0; i < a_batch_; i++) {
const int8_t *src = int8_src + i * params_->deep_ * params_->row_ * lite::DataTypeSize(src_data_type);
float16_t *dst = a_pack_ptr_ + i * params_->deep_ * params_->row_align_;
if (params_->a_transpose_) {
@ -203,20 +207,20 @@ void MatmulBaseFP16CPUKernel::InitMatrixB(const void *src_ptr, TypeId src_data_t
if (params_->b_transpose_) {
if (src_data_type == kNumberTypeFloat32) {
Float32ToFloat16(reinterpret_cast<const float *>(src_ptr), b_pack_ptr_,
params_->batch * params_->col_ * params_->deep_);
b_batch_ * params_->col_ * params_->deep_);
} else {
#ifdef ENABLE_ARM64
for (auto i = 0; i < params_->batch; ++i) {
for (auto i = 0; i < b_batch_; ++i) {
const auto *b_src = reinterpret_cast<const float16_t *>(src_ptr) + i * params_->col_ * params_->deep_;
auto *dst = b_pack_ptr_ + i * params_->col_align_ * params_->deep_;
RowMajor2Col16MajorFp16Opt(b_src, dst, params_->col_, params_->deep_);
}
#else
memcpy(b_pack_ptr_, src_ptr, params_->batch * params_->col_ * params_->deep_ * sizeof(float16_t));
memcpy(b_pack_ptr_, src_ptr, b_batch_ * params_->col_ * params_->deep_ * sizeof(float16_t));
#endif
}
} else {
for (int i = 0; i < params_->batch; i++) {
for (int i = 0; i < b_batch_; i++) {
#ifdef ENABLE_ARM64
const auto *b_src = reinterpret_cast<const float16_t *>(src_ptr) + i * params_->col_ * params_->deep_;
auto *dst = b_pack_ptr_ + i * params_->col_align_ * params_->deep_;
@ -231,7 +235,7 @@ void MatmulBaseFP16CPUKernel::InitMatrixB(const void *src_ptr, TypeId src_data_t
return;
}
for (int i = 0; i < params_->batch; i++) {
for (int i = 0; i < b_batch_; i++) {
const int8_t *src = int8_src + i * params_->deep_ * params_->col_ * lite::DataTypeSize(src_data_type);
float16_t *dst = b_pack_ptr_ + i * params_->deep_ * params_->col_align_;
if (params_->b_transpose_) {
@ -262,17 +266,16 @@ int MatmulBaseFP16CPUKernel::Prepare() {
auto b_tensor = in_tensors_[1];
MS_ASSERT(b_tensor != nullptr);
MS_ASSERT(b_tensor->data() != nullptr);
src_b_ = reinterpret_cast<float16_t *>(malloc(params_->batch * params_->col_ * params_->deep_ * sizeof(float16_t)));
src_b_ = reinterpret_cast<float16_t *>(malloc(b_batch_ * params_->col_ * params_->deep_ * sizeof(float16_t)));
if (src_b_ == nullptr) {
MS_LOG(ERROR) << "Matmul fp16 malloc src_b_ failed";
return RET_ERROR;
}
if (b_tensor->data_type() == kNumberTypeFloat32) {
Float32ToFloat16(reinterpret_cast<float *>(b_tensor->data()), src_b_,
params_->batch * params_->col_ * params_->deep_);
Float32ToFloat16(reinterpret_cast<float *>(b_tensor->data()), src_b_, b_batch_ * params_->col_ * params_->deep_);
} else {
memcpy(src_b_, b_tensor->data(), params_->batch * params_->col_ * params_->deep_ * sizeof(float16_t));
memcpy(src_b_, b_tensor->data(), b_batch_ * params_->col_ * params_->deep_ * sizeof(float16_t));
}
}
@ -313,25 +316,38 @@ int MatmulBaseFP16CPUKernel::RunImpl(int task_id) {
return RET_OK;
}
int MatmulBaseFP16CPUKernel::Run() {
int MatmulBaseFP16CPUKernel::BroadcastMatmulRun() {
auto c_ptr = reinterpret_cast<float16_t *>(out_tensors_[0]->data());
CHECK_NULL_RETURN(c_ptr);
if ((params_->a_const_ == false) || IsRepack()) {
if (RET_OK != InitBufferA()) {
for (int i = 0; i < params_->batch; ++i) {
if (vec_matmul_) {
batch_a_ptr_ = a_pack_ptr_ + a_offset_[i] * params_->deep_;
#ifdef ENABLE_ARM64
batch_b_ptr_ = b_pack_ptr_ + b_offset_[i] * params_->deep_ * params_->col_align_;
#else
batch_b_ptr_ = b_pack_ptr_ + b_offset_[i] * params_->deep_ * params_->col_;
#endif
batch_c_ptr_ = c_ptr + i * params_->row_ * params_->col_;
} else {
batch_a_ptr_ = a_pack_ptr_ + a_offset_[i] * params_->row_align_ * params_->deep_;
batch_b_ptr_ = b_pack_ptr_ + b_offset_[i] * params_->deep_ * params_->col_align_;
batch_c_ptr_ = c_ptr + i * params_->row_ * params_->col_;
}
auto ret = ParallelLaunch(this->ms_context_, MatmulBaseFP16Run, this, thread_count_);
if (ret != RET_OK) {
MS_LOG(ERROR) << "MatmulBaseFloatRun failed";
return RET_ERROR;
}
InitMatrixA(in_tensors_[0]->data());
}
if ((params_->b_const_ == false) || IsRepack()) {
if (RET_OK != InitBufferB()) {
FreeResizeBufA();
return RET_ERROR;
}
InitMatrixB(in_tensors_[1]->data(), in_tensors_[1]->data_type());
InitMatrixB(in_tensors_.at(1)->data(), in_tensors_.at(1)->data_type());
InitBias();
}
return RET_OK;
}
int MatmulBaseFP16CPUKernel::NormalMatmulRun() {
auto c_ptr = reinterpret_cast<float16_t *>(out_tensors_.at(0)->data());
CHECK_NULL_RETURN(c_ptr);
for (int i = 0; i < params_->batch; ++i) {
if (vec_matmul_) {
batch_a_ptr_ = a_pack_ptr_ + i * params_->deep_;
@ -353,6 +369,37 @@ int MatmulBaseFP16CPUKernel::Run() {
}
}
return RET_OK;
}
int MatmulBaseFP16CPUKernel::Run() {
if ((params_->a_const_ == false) || IsRepack()) {
if (RET_OK != InitBufferA()) {
return RET_ERROR;
}
InitMatrixA(in_tensors_[0]->data());
}
if ((params_->b_const_ == false) || IsRepack()) {
if (RET_OK != InitBufferB()) {
FreeResizeBufA();
return RET_ERROR;
}
InitMatrixB(in_tensors_[1]->data(), in_tensors_[1]->data_type());
InitBias();
}
if (!a_broadcast_ && !b_broadcast_) {
auto ret = NormalMatmulRun();
if (ret != RET_OK) {
MS_LOG(ERROR) << "NormalMatmulRun failed";
}
} else {
auto ret = BroadcastMatmulRun();
if (ret != RET_OK) {
MS_LOG(ERROR) << "BroadcastMatmulRun failed";
}
}
if (params_->a_const_ == false) {
FreeResizeBufA();
}

12
mindspore/lite/src/runtime/kernel/arm/fp16/matmul_base_fp16.h
View File

@ -22,6 +22,7 @@
#endif
#include <vector>
#include "src/inner_kernel.h"
#include "src/common/common.h"
#include "nnacl/matmul_parameter.h"
namespace mindspore::kernel {
@ -54,10 +55,21 @@ class MatmulBaseFP16CPUKernel : public InnerKernel {
void InitMatrixB(const void *src_ptr, TypeId data_type);
void FreeResizeBufA();
void FreeResizeBufB();
int NormalMatmulRun();
int BroadcastMatmulRun();
protected:
MatMulParameter *params_ = nullptr;
int row_tile_ = 0;
bool a_broadcast_ = false;
bool b_broadcast_ = false;
int a_batch_ = 1;
int b_batch_ = 1;
int batch_sizes_[MAX_SHAPE_SIZE] = {0};
int a_batch_sizes_[MAX_SHAPE_SIZE] = {0};
int b_batch_sizes_[MAX_SHAPE_SIZE] = {0};
std::vector<int> a_offset_;
std::vector<int> b_offset_;
private:
int thread_stride_ = 0;

80
mindspore/lite/src/runtime/kernel/arm/fp16/matmul_fp16.cc
View File

@ -15,10 +15,14 @@
*/
#include "src/runtime/kernel/arm/fp16/matmul_fp16.h"
#include <algorithm>
#include "include/errorcode.h"
#include "src/kernel_registry.h"
using mindspore::lite::kCHWDimNumber;
using mindspore::lite::KernelRegistrar;
using mindspore::lite::kHWDimNumber;
using mindspore::lite::kNCHWDimNumber;
using mindspore::lite::RET_ERROR;
using mindspore::lite::RET_OK;
using mindspore::schema::PrimitiveType_MatMul;
@ -34,7 +38,7 @@ void MatmulFP16CPUKernel::InitAShape() {
for (size_t i = 0; i < a_shape.size() - 2; ++i) {
batch *= a_shape[i];
}
params_->batch = batch;
a_batch_ = batch;
params_->row_ = params_->a_transpose_ ? a_shape[a_shape.size() - 1] : a_shape[a_shape.size() - 2];
params_->deep_ = params_->a_transpose_ ? a_shape[a_shape.size() - 2] : a_shape[a_shape.size() - 1];
params_->row_16_ = UP_ROUND(params_->row_, row_tile_);
@ -50,7 +54,7 @@ void MatmulFP16CPUKernel::InitBShape() {
for (size_t i = 0; i < b_shape.size() - 2; ++i) {
batch *= b_shape[i];
}
params_->batch = batch;
b_batch_ = batch;
params_->col_ = params_->b_transpose_ ? b_shape[b_shape.size() - 2] : b_shape[b_shape.size() - 1];
params_->col_8_ = UP_ROUND(params_->col_, 8);
params_->deep_ = params_->b_transpose_ ? b_shape[b_shape.size() - 1] : b_shape[b_shape.size() - 2];
@ -84,9 +88,81 @@ int MatmulFP16CPUKernel::Prepare() {
return ReSize();
}
int MatmulFP16CPUKernel::InitBroadcastParams() {
auto a_shape = in_tensors_[kInputIndex]->shape();
if (a_shape.size() < kNCHWDimNumber) {
int add_nums = kNCHWDimNumber - a_shape.size();
for (size_t i = 0; i < add_nums; ++i) {
a_shape.insert(a_shape.begin(), 1);
}
}
auto b_shape = in_tensors_[kWeightIndex]->shape();
if (b_shape.size() < kNCHWDimNumber) {
int add_nums = kNCHWDimNumber - b_shape.size();
for (size_t i = 0; i < add_nums; ++i) {
b_shape.insert(b_shape.begin(), 1);
}
}
for (int i = a_shape.size() - kCHWDimNumber; i >= 0; --i) {
if (static_cast<int>(a_shape.size() - kCHWDimNumber) == i) {
batch_sizes_[i] = std::max(a_shape[i], b_shape[i]);
a_batch_sizes_[i] = a_shape[i];
b_batch_sizes_[i] = b_shape[i];
} else {
batch_sizes_[i] = batch_sizes_[i + 1] * std::max(a_shape[i], b_shape[i]);
a_batch_sizes_[i] = a_batch_sizes_[i + 1] * a_shape[i];
b_batch_sizes_[i] = b_batch_sizes_[i + 1] * b_shape[i];
}
}
int out_batch = 1;
for (size_t i = 0; i < a_shape.size() - kHWDimNumber; ++i) {
out_batch *= MSMAX(a_shape[i], b_shape[i]);
if (a_shape[i] < b_shape[i] && b_shape[i] % a_shape[i] == 0) {
a_broadcast_ = true;
} else if (a_shape[i] > b_shape[i] && a_shape[i] % b_shape[i] == 0) {
b_broadcast_ = true;
} else if (a_shape[i] != b_shape[i]) {
MS_LOG(ERROR) << "matmul don't support broadcast for dimension " << a_shape << " and " << b_shape;
return RET_ERROR;
}
}
params_->batch = out_batch;
return RET_OK;
}
int MatmulFP16CPUKernel::ReSize() {
InitAShape();
InitBShape();
InitBroadcastParams();
a_offset_.resize(params_->batch, 0);
b_offset_.resize(params_->batch, 0);
auto a_shape = in_tensors_[kInputIndex]->shape();
auto b_shape = in_tensors_[kWeightIndex]->shape();
for (int i = 0; i < params_->batch; ++i) {
int delta = i;
int a_offset = 0;
int b_offset = 0;
for (size_t j = 0; j < a_shape.size() - kHWDimNumber; ++j) {
if (j > 0) {
delta = delta % batch_sizes_[j];
}
if (j < (a_shape.size() - kCHWDimNumber)) {
a_offset +=
(delta / batch_sizes_[j + 1] * a_shape[j] / std::max(a_shape[j], b_shape[j])) * a_batch_sizes_[j + 1];
b_offset +=
(delta / batch_sizes_[j + 1] * b_shape[j] / std::max(a_shape[j], b_shape[j])) * b_batch_sizes_[j + 1];
} else {
a_offset += (delta * a_shape[j] / std::max(a_shape[j], b_shape[j]));
b_offset += (delta * b_shape[j] / std::max(a_shape[j], b_shape[j]));
}
}
a_offset_[i] = a_offset;
b_offset_[i] = b_offset;
}
return MatmulBaseFP16CPUKernel::ReSize();
}

1
mindspore/lite/src/runtime/kernel/arm/fp16/matmul_fp16.h
View File

@ -35,6 +35,7 @@ class MatmulFP16CPUKernel : public MatmulBaseFP16CPUKernel {
private:
void InitAShape();
void InitBShape();
int InitBroadcastParams();
};
} // namespace mindspore::kernel

2
mindspore/lite/src/runtime/kernel/arm/fp32/fullconnection_fp32.cc
View File

@ -45,6 +45,8 @@ int FullconnectionCPUKernel::Prepare() {
}
params_->batch = 1;
a_batch_ = 1;
b_batch_ = 1;
params_->a_transpose_ = false;
params_->b_transpose_ = true;

79
mindspore/lite/src/runtime/kernel/arm/fp32/matmul_fp32.cc
View File

@ -15,11 +15,15 @@
*/
#include "src/runtime/kernel/arm/fp32/matmul_fp32.h"
#include <algorithm>
#include "include/errorcode.h"
#include "nnacl/fp32/matmul_fp32.h"
#include "src/kernel_registry.h"
using mindspore::lite::kCHWDimNumber;
using mindspore::lite::KernelRegistrar;
using mindspore::lite::kHWDimNumber;
using mindspore::lite::kNCHWDimNumber;
using mindspore::lite::RET_ERROR;
using mindspore::lite::RET_OK;
using mindspore::schema::PrimitiveType_MatMul;
@ -32,7 +36,7 @@ void MatmulCPUKernel::InitShapeA() {
for (size_t i = 0; i < a_shape.size() - 2; ++i) {
batch *= a_shape[i];
}
params_->batch = batch;
a_batch_ = batch;
params_->row_ = params_->a_transpose_ ? a_shape[a_shape.size() - 1] : a_shape[a_shape.size() - 2];
params_->deep_ = params_->a_transpose_ ? a_shape[a_shape.size() - 2] : a_shape[a_shape.size() - 1];
}
@ -44,7 +48,7 @@ void MatmulCPUKernel::InitShapeB() {
for (size_t i = 0; i < b_shape.size() - 2; ++i) {
batch *= b_shape[i];
}
params_->batch = batch;
b_batch_ = batch;
params_->col_ = params_->b_transpose_ ? b_shape[b_shape.size() - 2] : b_shape[b_shape.size() - 1];
params_->deep_ = params_->b_transpose_ ? b_shape[b_shape.size() - 1] : b_shape[b_shape.size() - 2];
}
@ -73,9 +77,80 @@ int MatmulCPUKernel::Prepare() {
return ReSize();
}
int MatmulCPUKernel::InitBroadcastParams() {
auto a_shape = in_tensors_[kInputIndex]->shape();
if (a_shape.size() < kNCHWDimNumber) {
size_t add_nums = kNCHWDimNumber - a_shape.size();
for (size_t i = 0; i < add_nums; ++i) {
a_shape.insert(a_shape.begin(), 1);
}
}
auto b_shape = in_tensors_[kWeightIndex]->shape();
if (b_shape.size() < kNCHWDimNumber) {
size_t add_nums = kNCHWDimNumber - b_shape.size();
for (size_t i = 0; i < add_nums; ++i) {
b_shape.insert(b_shape.begin(), 1);
}
}
for (int i = a_shape.size() - kCHWDimNumber; i >= 0; --i) {
if (static_cast<int>(a_shape.size() - kCHWDimNumber) == i) {
batch_sizes_[i] = std::max(a_shape[i], b_shape[i]);
a_batch_sizes_[i] = a_shape[i];
b_batch_sizes_[i] = b_shape[i];
} else {
batch_sizes_[i] = batch_sizes_[i + 1] * std::max(a_shape[i], b_shape[i]);
a_batch_sizes_[i] = a_batch_sizes_[i + 1] * a_shape[i];
b_batch_sizes_[i] = b_batch_sizes_[i + 1] * b_shape[i];
}
}
int out_batch = 1;
for (size_t i = 0; i < a_shape.size() - kHWDimNumber; ++i) {
out_batch *= MSMAX(a_shape[i], b_shape[i]);
if (a_shape[i] < b_shape[i] && b_shape[i] % a_shape[i] == 0) {
a_broadcast_ = true;
} else if (a_shape[i] > b_shape[i] && a_shape[i] % b_shape[i] == 0) {
b_broadcast_ = true;
} else if (a_shape[i] != b_shape[i]) {
MS_LOG(ERROR) << "matmul don't support broadcast for dimension " << a_shape << " and " << b_shape;
return RET_ERROR;
}
}
params_->batch = out_batch;
return RET_OK;
}
int MatmulCPUKernel::ReSize() {
InitShapeA();
InitShapeB();
InitBroadcastParams();
a_offset_.resize(params_->batch, 0);
b_offset_.resize(params_->batch, 0);
auto a_shape = in_tensors_[kInputIndex]->shape();
auto b_shape = in_tensors_[kWeightIndex]->shape();
for (int i = 0; i < params_->batch; ++i) {
int delta = i;
int a_offset = 0;
int b_offset = 0;
for (size_t j = 0; j < a_shape.size() - kHWDimNumber; ++j) {
if (j > 0) {
delta = delta % batch_sizes_[j];
}
if (j < (a_shape.size() - kCHWDimNumber)) {
a_offset +=
(delta / batch_sizes_[j + 1] * a_shape[j] / std::max(a_shape[j], b_shape[j])) * a_batch_sizes_[j + 1];
b_offset +=
(delta / batch_sizes_[j + 1] * b_shape[j] / std::max(a_shape[j], b_shape[j])) * b_batch_sizes_[j + 1];
} else {
a_offset += (delta * a_shape[j] / std::max(a_shape[j], b_shape[j]));
b_offset += (delta * b_shape[j] / std::max(a_shape[j], b_shape[j]));
}
}
a_offset_[i] = a_offset;
b_offset_[i] = b_offset;
}
return MatmulFp32BaseCPUKernel::ReSize();
}

1
mindspore/lite/src/runtime/kernel/arm/fp32/matmul_fp32.h
View File

@ -35,6 +35,7 @@ class MatmulCPUKernel : public MatmulFp32BaseCPUKernel {
private:
void InitShapeA();
void InitShapeB();
int InitBroadcastParams();
};
} // namespace mindspore::kernel
#endif // MINDSPORE_LITE_SRC_RUNTIME_KERNEL_ARM_FP32_MATMUL_H_

81
mindspore/lite/src/runtime/kernel/arm/fp32/matmul_fp32_base.cc
View File

@ -15,6 +15,7 @@
*/
#include "src/runtime/kernel/arm/fp32/matmul_fp32_base.h"
#include <algorithm>
#include "nnacl/fp32/matmul_fp32.h"
#include "nnacl/fp32/pack_fp32.h"
@ -159,11 +160,11 @@ int MatmulFp32BaseCPUKernel::InitMatrixA(const float *src_ptr) {
}
#else
if (vec_matmul_) {
memcpy(a_pack_ptr_, src_ptr, params_->batch * params_->deep_ * static_cast<int>(sizeof(float)));
memcpy(a_pack_ptr_, src_ptr, a_batch_ * params_->deep_ * static_cast<int>(sizeof(float)));
return RET_OK;
}
#endif
for (int i = 0; i < params_->batch; i++) {
for (int i = 0; i < a_batch_; i++) {
const float *src = src_ptr + i * params_->deep_ * params_->row_;
float *dst = a_pack_ptr_ + i * params_->deep_ * params_->row_align_;
if (params_->a_transpose_) {
@ -177,7 +178,7 @@ int MatmulFp32BaseCPUKernel::InitMatrixA(const float *src_ptr) {
int MatmulFp32BaseCPUKernel::InitMatrixB(const float *src_ptr) {
CHECK_NULL_RETURN(src_ptr);
for (int i = 0; i < params_->batch; i++) {
for (int i = 0; i < b_batch_; i++) {
const float *src = src_ptr + i * params_->deep_ * params_->col_;
float *dst = b_pack_ptr_ + i * params_->deep_ * params_->col_align_;
if (params_->b_transpose_) {
@ -297,7 +298,7 @@ int MatmulFp32BaseCPUKernel::Prepare() {
CHECK_LESS_RETURN(in_tensors_.size(), C2NUM);
CHECK_LESS_RETURN(out_tensors_.size(), 1);
init_global_variable();
matrix_a_pack_size_ = params_->batch * params_->row_align_ * params_->deep_;
matrix_a_pack_size_ = a_batch_ * params_->row_align_ * params_->deep_;
if (matrix_a_pack_size_ < 0) {
MS_LOG(ERROR) << "Matrix pack size is negative "
<< "matrix_a_pack_size=" << matrix_a_pack_size_;
@ -322,13 +323,13 @@ int MatmulFp32BaseCPUKernel::Prepare() {
// only copy weight data
// resize or run to pack
auto b_tensor = in_tensors_.at(1);
src_b_ = reinterpret_cast<float *>(
malloc(params_->batch * params_->deep_ * params_->col_ * static_cast<int>(sizeof(float))));
src_b_ =
reinterpret_cast<float *>(malloc(b_batch_ * params_->deep_ * params_->col_ * static_cast<int>(sizeof(float))));
if (src_b_ == nullptr) {
MS_LOG(ERROR) << "matmul fp16 src_b_ is failed!";
return RET_ERROR;
}
memcpy(src_b_, b_tensor->data(), params_->batch * params_->deep_ * params_->col_ * static_cast<int>(sizeof(float)));
memcpy(src_b_, b_tensor->data(), b_batch_ * params_->deep_ * params_->col_ * static_cast<int>(sizeof(float)));
}
return RET_OK;
}
@ -342,8 +343,8 @@ void MatmulFp32BaseCPUKernel::FreeBuffSrcB() {
int MatmulFp32BaseCPUKernel::ReSize() {
ResizeParameter();
matrix_a_pack_size_ = params_->batch * params_->row_align_ * params_->deep_;
matrix_b_pack_size_ = params_->batch * params_->col_align_ * params_->deep_;
matrix_a_pack_size_ = a_batch_ * params_->row_align_ * params_->deep_;
matrix_b_pack_size_ = b_batch_ * params_->col_align_ * params_->deep_;
if (matrix_a_pack_size_ < 0 || matrix_b_pack_size_ < 0) {
MS_LOG(ERROR) << "Matrix pack size is negative "
<< "matrix_a_pack_size=" << matrix_a_pack_size_ << "matrix_b_pack_size=" << matrix_b_pack_size_;
@ -371,6 +372,7 @@ int MatmulFp32BaseCPUKernel::ReSize() {
#else
thread_stride_ = UP_DIV(UP_DIV(params_->col_align_, col_tile_), thread_count_);
#endif
return RET_OK;
}
@ -397,6 +399,47 @@ int MatmulFp32BaseCPUKernel::InitTmpOutBuffer() {
return RET_OK;
}
int MatmulFp32BaseCPUKernel::NormalMatmulRun() {
for (int i = 0; i < params_->batch; ++i) {
batch_a_ptr_ = a_pack_ptr_ + i * params_->row_align_ * params_->deep_;
batch_b_ptr_ = b_pack_ptr_ + i * params_->deep_ * params_->col_align_;
#ifdef ENABLE_AVX
batch_c_ptr_ = output_data_ + i * params_->row_ * params_->col_align_;
#else
// need not aligned
batch_c_ptr_ = output_data_ + i * params_->row_ * params_->col_;
#endif
auto ret = ParallelLaunch(this->ms_context_, MatmulBaseFloatRun, this, thread_count_);
if (ret != RET_OK) {
MS_LOG(ERROR) << "MatmulBaseFloatRun failed";
}
}
return RET_OK;
}
int MatmulFp32BaseCPUKernel::BroadcastMatmulRun() {
for (int i = 0; i < params_->batch; ++i) {
batch_a_ptr_ = a_pack_ptr_ + a_offset_[i] * params_->row_align_ * params_->deep_;
batch_b_ptr_ = b_pack_ptr_ + b_offset_[i] * params_->deep_ * params_->col_align_;
#ifdef ENABLE_AVX
batch_c_ptr_ = output_data_ + i * params_->row_ * params_->col_align_;
#else
// need not aligned
batch_c_ptr_ = output_data_ + i * params_->row_ * params_->col_;
#endif
auto ret = ParallelLaunch(this->ms_context_, MatmulBaseFloatRun, this, thread_count_);
if (ret != RET_OK) {
MS_LOG(ERROR) << "MatmulBaseFloatRun failed";
return RET_ERROR;
}
}
return RET_OK;
}
int MatmulFp32BaseCPUKernel::Run() {
if (!params_->a_const_) {
auto a_ptr = reinterpret_cast<float *>(in_tensors_[0]->data());
@ -432,19 +475,15 @@ int MatmulFp32BaseCPUKernel::Run() {
return ret;
}
for (int i = 0; i < params_->batch; ++i) {
batch_a_ptr_ = a_pack_ptr_ + i * params_->row_align_ * params_->deep_;
batch_b_ptr_ = b_pack_ptr_ + i * params_->deep_ * params_->col_align_;
#ifdef ENABLE_AVX
batch_c_ptr_ = output_data_ + i * params_->row_ * params_->col_align_;
#else
// need not aligned
batch_c_ptr_ = output_data_ + i * params_->row_ * params_->col_;
#endif
ret = ParallelLaunch(this->ms_context_, MatmulBaseFloatRun, this, thread_count_);
if (!a_broadcast_ && !b_broadcast_) {
ret = NormalMatmulRun();
if (ret != RET_OK) {
MS_LOG(ERROR) << "MatmulBaseFloatRun failed";
MS_LOG(ERROR) << "NormalMatmulRun failed";
}
} else {
ret = BroadcastMatmulRun();
if (ret != RET_OK) {
MS_LOG(ERROR) << "BroadcastMatmulRun failed";
}
}

12
mindspore/lite/src/runtime/kernel/arm/fp32/matmul_fp32_base.h
View File

@ -21,6 +21,7 @@
#include "src/inner_kernel.h"
#include "nnacl/matmul_parameter.h"
#include "include/errorcode.h"
#include "src/common/common.h"
using mindspore::lite::RET_ERROR;
using mindspore::lite::RET_MEMORY_FAILED;
@ -61,11 +62,22 @@ class MatmulFp32BaseCPUKernel : public InnerKernel {
void FreeBuffSrcB();
int CalBroadCastBiasDataElements();
int InitTmpOutBuffer();
int NormalMatmulRun();
int BroadcastMatmulRun();
protected:
MatMulParameter *params_ = nullptr;
float *a_pack_ptr_ = nullptr;
float *b_pack_ptr_ = nullptr;
bool a_broadcast_ = false;
bool b_broadcast_ = false;
int a_batch_ = 1;
int b_batch_ = 1;
int batch_sizes_[MAX_SHAPE_SIZE] = {0};
int a_batch_sizes_[MAX_SHAPE_SIZE] = {0};
int b_batch_sizes_[MAX_SHAPE_SIZE] = {0};
std::vector<int> a_offset_;
std::vector<int> b_offset_;
private:
int col_tile_ = 0;

1
mindspore/lite/test/config/models_tf.cfg
View File

@ -108,3 +108,4 @@ fsr_270_mindspore.pb
fsr_360_mindspore.pb
fsr_720_mindspore.pb
hiai_asr_last_e1_cpu_fast_wavenet_batch1_frame1_one_cache.pb;2
tt_raw_h4800_mel80_ms_fe001_ex_20210506_joint_decoder.pb;14;4:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:1,640

1
mindspore/lite/test/config/models_tf_fp16.cfg
View File

@ -91,3 +91,4 @@ fsr_270_mindspore.pb 6.0
fsr_360_mindspore.pb 6.5
fsr_720_mindspore.pb 2.0
hiai_asr_last_e1_cpu_fast_wavenet_batch1_frame1_one_cache.pb;2
tt_raw_h4800_mel80_ms_fe001_ex_20210506_joint_decoder.pb;14;4:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:1,640 0.5

4
mindspore/lite/test/config/models_tflite.cfg
View File

@ -199,3 +199,7 @@ add_uint8.tflite;2:input0,input1
coco_ssd_mobilenet_v1_1.0.tflite
hiai_asr_last_e1_cpu_fast_wavenet_batch1_frame1_one_cache_fp32.tflite;2
hiai_asr_ctc.tflite;2
# weight quant model
tt_raw_h4800_mel80_ms_fe001_ex_20210506_encoder.tflite;25;1,15,80:1,15,80:1,15,80:1,15,80:1,15,80:1,15,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,31,80:1,640 8.5
# weight quant model
tt_raw_h4800_mel80_ms_fe001_ex_20210506_joint_decoder.tflite;14;4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:4,7,64:1,640:4 10.5

170
mindspore/lite/tools/optimizer/fusion/batchmatmul_fusion.cc
View File

@ -26,6 +26,8 @@
namespace mindspore::opt {
namespace {
constexpr int64_t kFcRightInputDims = 3;
constexpr float kFpPrecision = 1e-6;
void *GetInputAddr(const AnfNodePtr &node, size_t input_index) {
MS_ASSERT(node != nullptr);
if (!node->isa<CNode>()) {
@ -145,21 +147,7 @@ std::shared_ptr<ops::MatMul> BuildMatMulPrim(const CNodePtr &stack_cnode) {
matmul_cvalue->AddAttr("quant_params", quant_params_holder);
return matmul_cvalue;
}
} // namespace
const BaseRef BatchMatMulFusion::DefinePattern() const {
auto is_stack = std::make_shared<CondVar>(IsSpecifiedNode<&prim::kPrimStack>);
MS_CHECK_TRUE_RET(is_stack != nullptr, {});
auto is_fullconnect1 = std::make_shared<CondVar>(IsSpecifiedNode<&prim::kPrimFullConnection>);
MS_CHECK_TRUE_RET(is_fullconnect1 != nullptr, {});
auto is_fullconnect2 = std::make_shared<CondVar>(IsSpecifiedNode<&prim::kPrimFullConnection>);
MS_CHECK_TRUE_RET(is_fullconnect2 != nullptr, {});
auto is_seq_var = std::make_shared<SeqVar>();
MS_CHECK_TRUE_RET(is_seq_var != nullptr, {});
return VectorRef({is_stack, is_fullconnect1, is_fullconnect2, is_seq_var});
}
namespace {
constexpr float kFpPrecision = 1e-6;
bool IsTensorZero(const tensor::TensorPtr &tensor) {
MS_ASSERT(tensor != nullptr);
if (tensor->data_type() != TypeId::kNumberTypeFloat32) {
@ -207,68 +195,150 @@ bool IsFCNonBias(const CNodePtr &fc) {
}
return true;
}
bool ConnectTransposeConcat(const AnfNodePtr &node) {
auto cnode = node->cast<CNodePtr>();
if (cnode == nullptr) {
MS_LOG(ERROR) << "cnode is null";
return false;
}
auto right_transpose_node = cnode->input(1);
auto right_transpose_cnode = right_transpose_node->cast<CNodePtr>();
if (right_transpose_cnode == nullptr) {
MS_LOG(ERROR) << "cnode is null";
return false;
}
auto front_node = right_transpose_cnode->input(1);
auto front_cnode = front_node->cast<CNodePtr>();
if (front_cnode == nullptr) {
MS_LOG(ERROR) << "cnode is null";
return false;
}
if (CheckPrimitiveType(right_transpose_cnode, prim::kPrimTranspose) &&
(CheckPrimitiveType(front_cnode, prim::kPrimTranspose) || CheckPrimitiveType(front_cnode, prim::kPrimConcat))) {
return true;
}
return false;
}
int ResetReshapeParameters(const AnfNodePtr &reshape_node) {
auto reshape_cnode = reshape_node->cast<CNodePtr>();
MS_ASSERT(reshape_cnode != nullptr);
auto reshape_shape_param = reshape_cnode->input(kInputIndexTwo)->cast<ParameterPtr>();
MS_ASSERT(reshape_shape_param != nullptr);
auto shape_tensor = std::dynamic_pointer_cast<tensor::Tensor>(reshape_shape_param->default_param());
auto rmatmul_input_shape = shape_tensor->shape();
std::vector<int64_t> shape(1, 0);
if (rmatmul_input_shape.size() <= 0) {
MS_LOG(ERROR) << "Create tensor info failed";
return RET_ERROR;
} else if (shape[0] < kFcRightInputDims) {
shape[0] = rmatmul_input_shape[0] + 1;
}
auto tensor_info = std::make_shared<tensor::Tensor>(shape_tensor->data_type(), shape);
if (tensor_info == nullptr) {
MS_LOG(ERROR) << "Create tensor info failed";
return RET_ERROR;
}
int *tensor_data = reinterpret_cast<int *>(tensor_info->data_c());
tensor_data[0] = 1;
int *reshape_data = reinterpret_cast<int *>(shape_tensor->data_c());
for (int64_t i = 1; i < shape[0]; ++i) {
tensor_data[i] = reshape_data[i - 1];
}
lite::InitParameterFromTensorInfo(reshape_shape_param, tensor_info);
return RET_OK;
}
} // namespace
// slice +fullconnect ->batchmatmul
const AnfNodePtr BatchMatMulFusion::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
const EquivPtr &) const {
MS_ASSERT(func_graph != nullptr);
MS_ASSERT(node != nullptr);
auto stack_cnode = node->cast<CNodePtr>();
const BaseRef BatchMatMulFusion::DefinePattern() const {
auto is_stack = std::make_shared<CondVar>(IsSpecifiedNode<&prim::kPrimStack>);
MS_CHECK_TRUE_RET(is_stack != nullptr, {});
auto is_fullconnect1 = std::make_shared<CondVar>(IsSpecifiedNode<&prim::kPrimFullConnection>);
MS_CHECK_TRUE_RET(is_fullconnect1 != nullptr, {});
auto is_fullconnect2 = std::make_shared<CondVar>(IsSpecifiedNode<&prim::kPrimFullConnection>);
MS_CHECK_TRUE_RET(is_fullconnect2 != nullptr, {});
auto is_seq_var = std::make_shared<SeqVar>();
MS_CHECK_TRUE_RET(is_seq_var != nullptr, {});
return VectorRef({is_stack, is_fullconnect1, is_fullconnect2, is_seq_var});
}
bool BatchMatMulFusion::CheckCnodeProper(const CNodePtr &stack_cnode, const CNodePtr &fullconnect_cnode,
const CNodePtr &left_slice_cnode) const {
if (IsMarkedTrainOp(stack_cnode)) {
return nullptr;
return false;
}
// check stack node all inputs must fullconnect
for (size_t i = 1; i < stack_cnode->inputs().size(); i++) {
auto input_node = stack_cnode->input(i);
if (!CheckPrimitiveType(input_node, prim::kPrimFullConnection)) {
MS_LOG(WARNING) << "batchmatmulfusion stack node all inputs must fullconnect type";
return nullptr;
return false;
}
}
auto fullconnect_node = stack_cnode->input(1);
MS_ASSERT(fullconnect_node != nullptr);
auto fullconnect_cnode = fullconnect_node->cast<CNodePtr>();
MS_ASSERT(fullconnect_cnode != nullptr);
if (IsMarkedTrainOp(fullconnect_cnode)) {
return nullptr;
return false;
}
if (!IsFCNonBias(fullconnect_cnode)) {
return nullptr;
return false;
}
auto left_slice_node = fullconnect_cnode->input(1);
auto left_slice_cnode = left_slice_node->cast<CNodePtr>();
if (IsMarkedTrainOp(left_slice_cnode)) {
return nullptr;
return false;
}
MS_CHECK_TRUE_RET(left_slice_cnode != nullptr, nullptr);
if (!CheckPrimitiveType(left_slice_cnode, prim::kPrimSliceFusion)) {
if (!CheckPrimitiveType(left_slice_cnode, prim::kPrimReshape)) {
return nullptr;
return false;
}
auto &left_reshape_cnode = left_slice_cnode;
left_slice_cnode = left_reshape_cnode->input(1)->cast<CNodePtr>();
if (IsMarkedTrainOp(left_slice_cnode)) {
return nullptr;
auto up_slice_cnode = left_slice_cnode->input(1)->cast<CNodePtr>();
if (IsMarkedTrainOp(up_slice_cnode)) {
return false;
}
if (left_slice_cnode == nullptr || !CheckPrimitiveType(left_slice_cnode, prim::kPrimSliceFusion)) {
return nullptr;
if (up_slice_cnode == nullptr || !CheckPrimitiveType(up_slice_cnode, prim::kPrimSliceFusion)) {
return false;
}
}
return true;
}
const AnfNodePtr BatchMatMulFusion::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
const EquivPtr &) const {
MS_CHECK_TRUE_RET(func_graph != nullptr, nullptr);
MS_CHECK_TRUE_RET(node != nullptr, nullptr);
auto stack_cnode = node->cast<CNodePtr>();
auto fullconnect_node = stack_cnode->input(1);
auto fullconnect_cnode = fullconnect_node->cast<CNodePtr>();
MS_CHECK_TRUE_RET(fullconnect_cnode != nullptr, nullptr);
auto left_slice_node = fullconnect_cnode->input(1);
auto left_slice_cnode = left_slice_node->cast<CNodePtr>();
MS_CHECK_TRUE_RET(left_slice_cnode != nullptr, nullptr);
if (!CheckCnodeProper(stack_cnode, fullconnect_cnode, left_slice_cnode)) {
MS_LOG(WARNING) << stack_cnode->fullname_with_scope() << " can't fusion into matmul. Fusion failed";
return nullptr;
}
if (CheckPrimitiveType(left_slice_cnode, prim::kPrimReshape)) {
auto &left_reshape_cnode = left_slice_cnode;
left_slice_cnode = left_reshape_cnode->input(1)->cast<CNodePtr>();
}
// slice +fullconnect ->batchmatmul
auto left_matmul_input = left_slice_cnode->input(1);
auto right_reshape_node = fullconnect_cnode->input(kInputIndexTwo);
MS_ASSERT(right_reshape_node != nullptr);
auto matmul_cvalue = BuildMatMulPrim(stack_cnode);
if (matmul_cvalue == nullptr) {
MS_LOG(ERROR) << "new MatMul failed";
return nullptr;
}
MS_CHECK_TRUE_RET(matmul_cvalue != nullptr, nullptr);
auto matmul_value_node = NewValueNode(std::shared_ptr<ops::PrimitiveC>(matmul_cvalue));
MS_CHECK_TRUE_RET(matmul_value_node != nullptr, nullptr);
std::vector<AnfNodePtr> matmul_inputs = {matmul_value_node, left_matmul_input};
// batchmatmul right node may be const
bool right_transpose = false;
if (right_reshape_node->isa<Parameter>()) {
auto rmatmul_paramter = func_graph->add_parameter();
MS_CHECK_TRUE_RET(rmatmul_paramter != nullptr, nullptr);
@ -280,6 +350,14 @@ const AnfNodePtr BatchMatMulFusion::Process(const FuncGraphPtr &func_graph, cons
MS_ASSERT(prim_matmul != nullptr);
prim_matmul->set_transpose_b(true);
matmul_inputs.push_back(rmatmul_paramter);
} else if (ConnectTransposeConcat(right_reshape_node)) {
right_transpose = true;
auto ret = ResetReshapeParameters(right_reshape_node);
if (ret != RET_OK) {
MS_LOG(ERROR) << "reset reshape parameters failed";
return nullptr;
}
matmul_inputs.push_back(right_reshape_node);
} else {
auto right_reshape_cnode = right_reshape_node->cast<CNodePtr>();
MS_CHECK_TRUE_RET(right_reshape_cnode != nullptr, nullptr);
@ -303,6 +381,10 @@ const AnfNodePtr BatchMatMulFusion::Process(const FuncGraphPtr &func_graph, cons
matmul_cnode->set_fullname_with_scope("matmul_" + stack_cnode->fullname_with_scope());
MS_CHECK_TRUE_RET(stack_cnode->abstract() != nullptr, nullptr);
matmul_cnode->set_abstract(stack_cnode->abstract()->Clone());
if (right_transpose) {
auto matmul_primitive = GetValueNode<std::shared_ptr<ops::MatMul>>(matmul_cnode->input(0));
matmul_primitive->set_transpose_b(true);
}
MS_LOG(INFO) << "stack node:" << stack_cnode->fullname_with_scope() << " batchmatmul fusion success";
return matmul_cnode;
}

2
mindspore/lite/tools/optimizer/fusion/batchmatmul_fusion.h
View File

@ -30,6 +30,8 @@ class BatchMatMulFusion : public PatternProcessPass {
private:
const BaseRef DefinePattern() const override;
const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override;
bool CheckCnodeProper(const CNodePtr &stack_cnode, const CNodePtr &fullconnect_cnode,
const CNodePtr &left_slice_cnode) const;
};
} // namespace opt
} // namespace mindspore