Huawei_Technology
/
mindspore
Not watched
1
0
Fork 0
Code
Releases 13 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
Browse Source

[MSLITE][Develop] rewrite arm cpu fp16 op: matmul init

tags/v1.1.0
yangruoqi713 5 years ago
parent
f957e4f588
commit
a769b7d592
1 changed files with 84 additions and 81 deletions
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. +84
    -81
      mindspore/lite/src/runtime/kernel/arm/fp16/matmul_fp16.cc

+ 84
- 81
mindspore/lite/src/runtime/kernel/arm/fp16/matmul_fp16.cc View File

@@ -30,9 +30,7 @@ using mindspore::lite::RET_OK;
using mindspore::schema::PrimitiveType_MatMul;

namespace mindspore::kernel {
MatmulFP16CPUKernel::~MatmulFP16CPUKernel() { FreeTmpBuffer(); }

void MatmulFP16CPUKernel::FreeTmpBuffer() {
MatmulFP16CPUKernel::~MatmulFP16CPUKernel() {
if (a_pack_ptr_ != nullptr) {
free(a_pack_ptr_);
a_pack_ptr_ = nullptr;
@@ -47,6 +45,17 @@ void MatmulFP16CPUKernel::FreeTmpBuffer() {
}
}

void MatmulFP16CPUKernel::FreeTmpBuffer() {
if (a_pack_ptr_ != nullptr) {
params_->a_const_ ? free(a_pack_ptr_) : context_->allocator->Free(a_pack_ptr_);
a_pack_ptr_ = nullptr;
}
if (b_pack_ptr_ != nullptr) {
params_->b_const_ ? free(b_pack_ptr_) : context_->allocator->Free(b_pack_ptr_);
b_pack_ptr_ = nullptr;
}
}

int MatmulFP16CPUKernel::MallocMatrixABuffer() {
auto a_shape = in_tensors_[0]->shape();
int batch = 1;
@@ -57,9 +66,13 @@ int MatmulFP16CPUKernel::MallocMatrixABuffer() {
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_, C16NUM);

a_pack_ptr_ =
reinterpret_cast<float16_t *>(malloc(params_->batch * params_->row_16_ * params_->deep_ * sizeof(float16_t)));
if (params_->a_const_) {
a_pack_ptr_ =
reinterpret_cast<float16_t *>(malloc(params_->batch * params_->row_16_ * params_->deep_ * sizeof(float16_t)));
} else {
a_pack_ptr_ = reinterpret_cast<float16_t *>(
context_->allocator->Malloc(params_->batch * params_->row_16_ * params_->deep_ * sizeof(float16_t)));
}
if (a_pack_ptr_ == nullptr) {
FreeTmpBuffer();
return RET_MEMORY_FAILED;
@@ -82,8 +95,13 @@ int MatmulFP16CPUKernel::MallocMatrixBBuffer() {
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];

b_pack_ptr_ =
reinterpret_cast<float16_t *>(malloc(params_->batch * params_->col_8_ * params_->deep_ * sizeof(float16_t)));
if (params_->b_const_) {
b_pack_ptr_ =
reinterpret_cast<float16_t *>(malloc(params_->batch * params_->col_8_ * params_->deep_ * sizeof(float16_t)));
} else {
b_pack_ptr_ = reinterpret_cast<float16_t *>(
context_->allocator->Malloc(params_->batch * params_->col_8_ * params_->deep_ * sizeof(float16_t)));
}
if (b_pack_ptr_ == nullptr) {
FreeTmpBuffer();
return RET_MEMORY_FAILED;
@@ -95,59 +113,32 @@ int MatmulFP16CPUKernel::MallocMatrixBBuffer() {
}

int MatmulFP16CPUKernel::InitBias() {
auto b_shape = in_tensors_[1]->shape();
auto c_shape = out_tensors_[0]->shape();
params_->col_ = params_->b_const_
? (params_->b_transpose_ ? b_shape[b_shape.size() - 2] : b_shape[b_shape.size() - 1])
: (c_shape[c_shape.size() - 1]);
params_->col_8_ = UP_ROUND(params_->col_, 8);
bias_ptr_ = reinterpret_cast<float16_t *>(malloc(params_->col_8_ * sizeof(float16_t)));
if (bias_ptr_ == nullptr) {
FreeTmpBuffer();
return RET_MEMORY_FAILED;
}
memset(bias_ptr_, 0, params_->col_8_ * sizeof(float16_t));
if (in_tensors_.size() == 3) {
auto c_shape = out_tensors_[0]->shape();
auto bias_shape = in_tensors_[1]->shape();
if (bias_shape[bias_shape.size() - 1] != c_shape[c_shape.size() - 1]) {
MS_LOG(ERROR) << "The bias'dimension is not equal with colum";
FreeTmpBuffer();
return RET_INPUT_TENSOR_ERROR;
}
auto col = c_shape[c_shape.size() - 1];
auto col_8 = UP_ROUND(col, 8);
bias_ptr_ = reinterpret_cast<float16_t *>(malloc(col_8 * sizeof(float16_t)));
if (bias_ptr_ == nullptr) {
FreeTmpBuffer();
return RET_MEMORY_FAILED;
}
memset(bias_ptr_, 0, col_8 * sizeof(float16_t));
Float32ToFloat16(reinterpret_cast<float *>(in_tensors_[2]->data_c()), bias_ptr_, col);
Float32ToFloat16(reinterpret_cast<float *>(in_tensors_[2]->data_c()), bias_ptr_, params_->col_);
}
return RET_OK;
}

int MatmulFP16CPUKernel::ReSize() {
if (params_->a_const_ == false || params_->a_init_shape_ == false) {
if (a_pack_ptr_ != nullptr) {
free(a_pack_ptr_);
a_pack_ptr_ = nullptr;
}
auto ret = MallocMatrixABuffer();
if (ret != RET_OK) {
MS_LOG(ERROR) << "Matmul fp16 malloc matrix a buffer failed";
return RET_ERROR;
}
}
if (params_->b_const_ == false || params_->b_init_shape_ == false) {
if (b_pack_ptr_ != nullptr) {
free(b_pack_ptr_);
b_pack_ptr_ = nullptr;
}
auto ret = MallocMatrixBBuffer();
if (!params_->b_const_) {
auto ret = InitBias();
if (ret != RET_OK) {
MS_LOG(ERROR) << "Matmul fp16 malloc matrix b buffer failed";
MS_LOG(ERROR) << "Matmul fp16 init bias failed";
return RET_ERROR;
}
}
if (bias_ptr_ != nullptr) {
free(bias_ptr_);
bias_ptr_ = nullptr;
}
auto ret = InitBias();
if (ret != RET_OK) {
MS_LOG(ERROR) << "Matmul fp16 init bias failed";
return RET_ERROR;
}
return RET_OK;
}

@@ -200,47 +191,36 @@ void MatmulFP16CPUKernel::InitMatrixB(float16_t *b_ptr, float16_t *b_pack_ptr) {
}

int MatmulFP16CPUKernel::Init() {
params_->a_init_shape_ = (in_tensors_[0]->shape().size() != 0);
params_->b_init_shape_ = (in_tensors_[1]->shape().size() != 0);
if (params_->a_init_shape_ == true) {
params_->a_const_ = (in_tensors_[0]->data_c() != nullptr);
params_->b_const_ = (in_tensors_[1]->data_c() != nullptr);
if (params_->a_const_) {
auto ret = MallocMatrixABuffer();
if (ret != RET_OK) {
MS_LOG(ERROR) << "Matmul fp16 malloc matrix a buffer failed";
return RET_ERROR;
}
}
if (params_->b_init_shape_ == true) {
auto ret = MallocMatrixBBuffer();
if (ret != RET_OK) {
MS_LOG(ERROR) << "Matmul fp16 malloc matrix b buffer failed";
MS_LOG(ERROR) << "Matmul fp16 malloc matrix A buffer failed";
return RET_ERROR;
}
}

params_->a_const_ = (in_tensors_[0]->data_c() != nullptr);
params_->b_const_ = (in_tensors_[1]->data_c() != nullptr);
if (params_->a_const_ == true) {
if (in_tensors_[0]->data_type() == kNumberTypeFloat32) {
InitMatrixA(reinterpret_cast<float *>(in_tensors_[0]->data_c()), a_pack_ptr_);
} else {
InitMatrixA(reinterpret_cast<float16_t *>(in_tensors_[0]->data_c()), a_pack_ptr_);
}
}
if (params_->b_const_ == true) {
if (params_->b_const_) {
auto ret = MallocMatrixBBuffer();
if (ret != RET_OK) {
MS_LOG(ERROR) << "Matmul fp16 malloc matrix B buffer failed";
return RET_ERROR;
}
if (in_tensors_[1]->data_type() == kNumberTypeFloat32) {
InitMatrixB(reinterpret_cast<float *>(in_tensors_[1]->data_c()), b_pack_ptr_);
} else {
InitMatrixB(reinterpret_cast<float16_t *>(in_tensors_[1]->data_c()), b_pack_ptr_);
}
}

if (!InferShapeDone()) {
return RET_OK;
}
auto ret = InitBias();
if (ret != RET_OK) {
MS_LOG(ERROR) << "Matmul fp16 init bias failed";
return RET_ERROR;
ret = InitBias();
if (ret != RET_OK) {
MS_LOG(ERROR) << "Matmul fp16 init bias failed";
return RET_ERROR;
}
}
return RET_OK;
}
@@ -275,7 +255,7 @@ int MatmulFP16Run(void *cdata, int task_id) {
auto op = reinterpret_cast<MatmulFP16CPUKernel *>(cdata);
auto error_code = op->RunImpl(task_id);
if (error_code != RET_OK) {
MS_LOG(ERROR) << "MatmulFp32Run error task_id[" << task_id << "] error_code[" << error_code << "]";
MS_LOG(ERROR) << "MatmulFp16Run error task_id[" << task_id << "] error_code[" << error_code << "]";
return RET_ERROR;
}
return RET_OK;
@@ -294,14 +274,24 @@ int MatmulFP16CPUKernel::Run() {
} else {
c_ptr = reinterpret_cast<float16_t *>(out_tensor->data_c());
}
if (params_->a_const_ == false) {
if (!params_->a_const_) {
ret = MallocMatrixABuffer();
if (ret != RET_OK) {
MS_LOG(ERROR) << "Matmul fp16 malloc matrix A buffer failed";
return RET_ERROR;
}
if (in_tensors_[0]->data_type() == kNumberTypeFloat32) {
InitMatrixA(reinterpret_cast<float *>(in_tensors_[0]->data_c()), a_pack_ptr_);
} else {
InitMatrixA(reinterpret_cast<float16_t *>(in_tensors_[0]->data_c()), a_pack_ptr_);
}
}
if (params_->b_const_ == false) {
if (!params_->b_const_) {
ret = MallocMatrixBBuffer();
if (ret != RET_OK) {
MS_LOG(ERROR) << "Matmul fp16 malloc matrix B buffer failed";
return RET_ERROR;
}
if (in_tensors_[1]->data_type() == kNumberTypeFloat32) {
InitMatrixB(reinterpret_cast<float *>(in_tensors_[1]->data_c()), b_pack_ptr_);
} else {
@@ -312,7 +302,12 @@ int MatmulFP16CPUKernel::Run() {
current_a_ = a_pack_ptr_ + i * params_->row_16_ * params_->deep_;
current_b_ = b_pack_ptr_ + i * params_->deep_ * params_->col_8_;
current_c_ = c_ptr + i * params_->row_ * params_->col_;
ParallelLaunch(this->context_->thread_pool_, MatmulFP16Run, this, thread_count_);
ret = ParallelLaunch(this->context_->thread_pool_, MatmulFP16Run, this, thread_count_);
if (ret != RET_OK) {
MS_LOG(ERROR) << "Matmul fp16 run function MatmulFP16Run failed";
FreeTmpBuffer();
return RET_ERROR;
}
}
if (out_tensor->data_type() == kNumberTypeFloat32) {
auto size = out_tensor->ElementsNum();
@@ -320,6 +315,14 @@ int MatmulFP16CPUKernel::Run() {
Float16ToFloat32(output_ptr_, out_tensor_data, size);
ctx_->allocator->Free(output_ptr_);
}
if (!params_->a_const_) {
context_->allocator->Free(a_pack_ptr_);
a_pack_ptr_ = nullptr;
}
if (!params_->b_const_) {
context_->allocator->Free(b_pack_ptr_);
b_pack_ptr_ = nullptr;
}
return RET_OK;
}



Write Preview
Loading…
Cancel
Save