* Copyright (c) 2025 Huawei Technologies Co., Ltd.
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
* CANN Open Software License Agreement Version 2.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
*/
#include "linear_aclnn_runner.h"
#include <aclnn/opdev/op_errno.h>
#include "atb/utils/aclnn_util.h"
#include "atb/utils/log.h"
#include "atb/utils/utils_internal.h"
#include "atbops/params/params.h"
#include "atb/utils/operation_register.h"
namespace {
static const int MATMUL_SELF_ACLNN_TENSOR_IDX = 0;
static const int MATMUL_MAT2_ACLNN_TENSOR_IDX = 1;
static const int MATMUL_OUT_ACLNN_TENSOR_IDX = 0;
static const int ADDMM_SELF_ACLNN_TENSOR_IDX = 0;
static const int ADDMM_MAT1_ACLNN_TENSOR_IDX = 1;
static const int ADDMM_MAT2_ACLNN_TENSOR_IDX = 2;
static const int ADDMM_OUT_ACLNN_TENSOR_IDX = 0;
static const int DEFAULT_ALIGN = 16;
}
namespace atb {
AclnnMatmulGetWorkspaceSizeFunc LinearAclnnRunner::aclnnMatmulGetWorkspaceSizeFunc_ = nullptr;
AclnnMatmulExecuteFunc LinearAclnnRunner::aclnnMatmulExecuteFunc_ = nullptr;
AclnnAddmmGetWorkspaceSizeFunc LinearAclnnRunner::aclnnAddmmGetWorkspaceSizeFunc_ = nullptr;
AclnnAddmmExecuteFunc LinearAclnnRunner::aclnnAddmmExecuteFunc_ = nullptr;
AclnnMatmulWeightNzGetWorkspaceSizeFunc LinearAclnnRunner::aclnnMatmulWeightNzGetWorkspaceSizeFunc_ = nullptr;
AclnnMatmulWeightNzExecuteFunc LinearAclnnRunner::aclnnMatmulWeightNzExecuteFunc_ = nullptr;
AclnnAddmmWeightNzGetWorkspaceSizeFunc LinearAclnnRunner::aclnnAddmmWeightNzGetWorkspaceSizeFunc_ = nullptr;
AclnnAddmmWeightNzExecuteFunc LinearAclnnRunner::aclnnAddmmWeightNzExecuteFunc_ = nullptr;
AclnnBatchMatMulWeightNzGetWorkspaceSizeFunc LinearAclnnRunner::aclnnBatchMatMulWeightNzGetWorkspaceSizeFunc_ = nullptr;
AclnnBatchMatMulWeightNzExecuteFunc LinearAclnnRunner::aclnnBatchMatMulWeightNzExecuteFunc_ = nullptr;
LinearAclnnRunner::LinearAclnnRunner(const infer::LinearParam ¶m) : AclnnRunner("LinearAclnnRunner"), param_(param)
{
ATB_LOG(INFO) << GetLogPrefix() << "LinearAclnnRunner::LinearAclnnRunner";
GetTensorNum();
}
LinearAclnnRunner::~LinearAclnnRunner()
{
if (alpha_) {
if (aclDestroyScalar(alpha_) != ACL_SUCCESS) {
ATB_LOG(ERROR) << GetLogPrefix() << "alpha aclDestroyScalar error";
}
alpha_ = nullptr;
}
if (beta_) {
if (aclDestroyScalar(beta_) != ACL_SUCCESS) {
ATB_LOG(ERROR) << GetLogPrefix() << "beta aclDestroyScalar error";
}
beta_ = nullptr;
}
}
Status LinearAclnnRunner::LoadAclnnFuncs()
{
ATB_LOG(INFO) << "LinearAclnnRunner::LoadAclnnFuncs";
if ((aclnnMatmulGetWorkspaceSizeFunc_ && aclnnMatmulExecuteFunc_) &&
(aclnnAddmmGetWorkspaceSizeFunc_ && aclnnAddmmExecuteFunc_) &&
(aclnnMatmulWeightNzGetWorkspaceSizeFunc_ && aclnnMatmulWeightNzExecuteFunc_) &&
(aclnnAddmmWeightNzGetWorkspaceSizeFunc_ && aclnnAddmmWeightNzExecuteFunc_) &&
(aclnnBatchMatMulWeightNzGetWorkspaceSizeFunc_ && aclnnBatchMatMulWeightNzExecuteFunc_)) {
return NO_ERROR;
}
Status st = NO_ERROR;
if (!aclnnMatmulGetWorkspaceSizeFunc_ || !aclnnMatmulExecuteFunc_) {
st = LoadFromSharedObjectFile("aclnnMatmulGetWorkspaceSize", "aclnnMatmul", aclnnMatmulGetWorkspaceSizeFunc_,
aclnnMatmulExecuteFunc_);
if (st != NO_ERROR) {
return st;
}
}
if (!aclnnAddmmGetWorkspaceSizeFunc_ || !aclnnAddmmExecuteFunc_) {
st = LoadFromSharedObjectFile("aclnnAddmmGetWorkspaceSize", "aclnnAddmm", aclnnAddmmGetWorkspaceSizeFunc_,
aclnnAddmmExecuteFunc_);
if (st != NO_ERROR) {
return st;
}
}
if (!aclnnMatmulWeightNzGetWorkspaceSizeFunc_ || !aclnnMatmulWeightNzExecuteFunc_) {
st = LoadFromSharedObjectFile("aclnnMatmulWeightNzGetWorkspaceSize", "aclnnMatmulWeightNz",
aclnnMatmulWeightNzGetWorkspaceSizeFunc_, aclnnMatmulWeightNzExecuteFunc_);
if (st != NO_ERROR) {
return st;
}
}
if (!aclnnAddmmWeightNzGetWorkspaceSizeFunc_ || !aclnnAddmmWeightNzExecuteFunc_) {
st = LoadFromSharedObjectFile("aclnnAddmmWeightNzGetWorkspaceSize", "aclnnAddmmWeightNz",
aclnnAddmmWeightNzGetWorkspaceSizeFunc_, aclnnAddmmWeightNzExecuteFunc_);
if (st != NO_ERROR) {
return st;
}
}
if (!aclnnBatchMatMulWeightNzGetWorkspaceSizeFunc_ || !aclnnBatchMatMulWeightNzExecuteFunc_) {
st = LoadFromSharedObjectFile("aclnnBatchMatMulWeightNzGetWorkspaceSize", "aclnnBatchMatMulWeightNz",
aclnnBatchMatMulWeightNzGetWorkspaceSizeFunc_,
aclnnBatchMatMulWeightNzExecuteFunc_);
if (st != NO_ERROR) {
return st;
}
}
return NO_ERROR;
}
Status LinearAclnnRunner::BuildAclnnVariantPack(const RunnerVariantPack &runnerVariantPack)
{
ATB_LOG(INFO) << GetLogPrefix()
<< "LinearAclnnRunner::BuildAclnnVariantPack, runnerVariantPack: " << runnerVariantPack.ToString();
atbVariantPack_ = runnerVariantPack;
isWeightNz_ = runnerVariantPack.inTensors[1].desc.format == ACL_FORMAT_FRACTAL_NZ;
isBatch_ = runnerVariantPack.inTensors[1].desc.shape.dimNum == 3 ||
(runnerVariantPack.inTensors[1].desc.shape.dimNum == 4 &&
runnerVariantPack.inTensors[1].desc.shape.dims[0] != 1);
InitTensorIndex();
aclnnVariantPack_.aclInTensors.reserve(aclInTensorNum_);
aclnnVariantPack_.aclInTensors.resize(aclInTensorNum_);
aclnnVariantPack_.aclOutTensors.reserve(aclOutTensorNum_);
aclnnVariantPack_.aclOutTensors.resize(aclOutTensorNum_);
Status st = NO_ERROR;
if (param_.hasBias) {
st = CreateAddmmMat1AclnnTensor();
if (st != NO_ERROR) {
return st;
}
st = CreateAddmmMat2AclnnTensor();
if (st != NO_ERROR) {
return st;
}
st = CreateAddmmSelfAclnnTensor();
if (st != NO_ERROR) {
return st;
}
st = CreateAddmmOutAclnnTensor();
if (st != NO_ERROR) {
return st;
}
} else {
st = CreateMatmulSelfAclnnTensor();
if (st != NO_ERROR) {
return st;
}
st = CreateMatmulMat2AclnnTensor();
if (st != NO_ERROR) {
return st;
}
st = CreateMatmulOutAclnnTensor();
if (st != NO_ERROR) {
return st;
}
}
return NO_ERROR;
}
aclnnStatus LinearAclnnRunner::SetAclNNWorkspaceExecutor()
{
ATB_LOG(INFO) << GetLogPrefix() << "LinearAclnnRunner::SetAclNNWorkspaceExecutor";
if (isWeightNz_) {
if (param_.hasBias) {
return SetAclnnAddmmWeightNzWorkspaceExecutor();
} else {
if (isBatch_) {
return SetAclnnBatchMatMulWeightNzWorkspaceExecutor();
} else {
return SetAclnnMatmulWeightNzWorkspaceExecutor();
}
}
} else {
if (param_.hasBias) {
return SetAclnnAddmmWorkspaceExecutor();
} else {
return SetAclnnMatmulWorkspaceExecutor();
}
}
}
Status LinearAclnnRunner::LaunchAclnnKernel()
{
ATB_LOG(INFO) << GetLogPrefix() << "LinearAclnnRunner::LaunchAclnnKernel";
aclrtStream executeStream = GetExecuteStream(atbVariantPack_.context);
aclnnStatus ret = ACL_SUCCESS;
if (isWeightNz_) {
if (param_.hasBias) {
ret = aclnnAddmmWeightNzExecuteFunc_(atbVariantPack_.workspaceBuffer, atbVariantPack_.workspaceBufferSize,
aclnnExecutor_.get(), executeStream);
} else {
if (isBatch_) {
ret = aclnnBatchMatMulWeightNzExecuteFunc_(atbVariantPack_.workspaceBuffer,
atbVariantPack_.workspaceBufferSize, aclnnExecutor_.get(),
executeStream);
} else {
ret = aclnnMatmulWeightNzExecuteFunc_(atbVariantPack_.workspaceBuffer,
atbVariantPack_.workspaceBufferSize, aclnnExecutor_.get(),
executeStream);
}
}
} else {
if (param_.hasBias) {
ret = aclnnAddmmExecuteFunc_(atbVariantPack_.workspaceBuffer, atbVariantPack_.workspaceBufferSize,
aclnnExecutor_.get(), executeStream);
} else {
ret = aclnnMatmulExecuteFunc_(atbVariantPack_.workspaceBuffer, atbVariantPack_.workspaceBufferSize,
aclnnExecutor_.get(), executeStream);
}
}
if (ret != ACL_SUCCESS) {
ATB_LOG(ERROR) << GetLogPrefix() << "Atb aclnn op kernel launch failed with return value: " << ret;
return ERROR_CANN_ERROR;
}
ATB_LOG(INFO) << GetLogPrefix() << "LaunchAclnnKernel execute success.";
return NO_ERROR;
}
void LinearAclnnRunner::GetTensorNum()
{
if (param_.hasBias) {
aclInTensorNum_ = 3;
} else {
aclInTensorNum_ = 2;
}
aclOutTensorNum_ = 1;
}
void LinearAclnnRunner::InitTensorIndex()
{
atbInTensorIndex_ = 0;
aclInTensorIndex_ = 0;
atbOutTensorIndex_ = 0;
aclOutTensorIndex_ = 0;
matmulSelfAclTensorIndex_ = 0;
matmulMat2AclTensorIndex_ = 0;
matmulOutAclTensorIndex_ = 0;
addmmSelfAclTensorIndex_ = 0;
addmmMat1AclTensorIndex_ = 0;
addmmMat2AclTensorIndex_ = 0;
addmmOutAclTensorIndex_ = 0;
}
Status LinearAclnnRunner::CreateMatmulSelfAclnnTensor()
{
ATB_LOG(INFO) << "LinearAclnnRunner::CreateMatmulSelfAclnnTensor";
std::shared_ptr<AclNNTensor> aclnnTensorPtr = CreateXAclnnTensor(MATMUL_SELF_ACLNN_TENSOR_IDX);
if (!aclnnTensorPtr->tensor) {
ATB_LOG(ERROR) << GetLogPrefix() << "matmul self aclCreateTensor failed";
return ERROR_INTERNAL_ERROR;
}
aclnnVariantPack_.aclInTensors.at(aclInTensorIndex_) = aclnnTensorPtr;
matmulSelfAclTensorIndex_ = aclInTensorIndex_++;
return NO_ERROR;
}
Status LinearAclnnRunner::CreateMatmulMat2AclnnTensor()
{
ATB_LOG(INFO) << "LinearAclnnRunner::CreateMatmulMat2AclnnTensor";
std::shared_ptr<AclNNTensor> aclnnTensorPtr = isWeightNz_ ?
CreateWeightNzAclnnTensor(MATMUL_MAT2_ACLNN_TENSOR_IDX) :
CreateWeightAclnnTensor(MATMUL_MAT2_ACLNN_TENSOR_IDX);
if (!aclnnTensorPtr->tensor) {
ATB_LOG(ERROR) << GetLogPrefix() << "matmul mat2 aclCreateTensor failed";
return ERROR_INTERNAL_ERROR;
}
aclnnVariantPack_.aclInTensors.at(aclInTensorIndex_) = aclnnTensorPtr;
matmulMat2AclTensorIndex_ = aclInTensorIndex_++;
return NO_ERROR;
}
Status LinearAclnnRunner::CreateMatmulOutAclnnTensor()
{
ATB_LOG(INFO) << "LinearAclnnRunner::CreateMatmulOutAclnnTensor";
std::shared_ptr<AclNNTensor> aclnnTensorPtr = CreateOutputAclnnTensor(MATMUL_OUT_ACLNN_TENSOR_IDX);
if (!aclnnTensorPtr->tensor) {
ATB_LOG(ERROR) << GetLogPrefix() << "matmul out aclCreateTensor failed";
return ERROR_INTERNAL_ERROR;
}
aclnnVariantPack_.aclOutTensors.at(aclOutTensorIndex_) = aclnnTensorPtr;
matmulOutAclTensorIndex_ = aclOutTensorIndex_++;
return NO_ERROR;
}
Status LinearAclnnRunner::CreateAddmmSelfAclnnTensor()
{
ATB_LOG(INFO) << "LinearAclnnRunner::CreateAddmmSelfAclnnTensor";
std::shared_ptr<AclNNTensor> aclnnTensorPtr = CreateBiasAclnnTensor(ADDMM_SELF_ACLNN_TENSOR_IDX);
if (!aclnnTensorPtr->tensor) {
ATB_LOG(ERROR) << GetLogPrefix() << "addmm self aclCreateTensor failed";
return ERROR_INTERNAL_ERROR;
}
aclnnVariantPack_.aclInTensors.at(aclInTensorIndex_) = aclnnTensorPtr;
addmmSelfAclTensorIndex_ = aclInTensorIndex_++;
return NO_ERROR;
}
Status LinearAclnnRunner::CreateAddmmMat1AclnnTensor()
{
ATB_LOG(INFO) << "LinearAclnnRunner::CreateAddmmMat1AclnnTensor";
std::shared_ptr<AclNNTensor> aclnnTensorPtr = CreateXAclnnTensor(ADDMM_MAT1_ACLNN_TENSOR_IDX);
if (!aclnnTensorPtr->tensor) {
ATB_LOG(ERROR) << GetLogPrefix() << "addmm mat1 aclCreateTensor failed";
return ERROR_INTERNAL_ERROR;
}
aclnnVariantPack_.aclInTensors.at(aclInTensorIndex_) = aclnnTensorPtr;
addmmMat2AclTensorIndex_ = aclInTensorIndex_++;
return NO_ERROR;
}
Status LinearAclnnRunner::CreateAddmmMat2AclnnTensor()
{
ATB_LOG(INFO) << "LinearAclnnRunner::CreateAddmmMat2AclnnTensor";
std::shared_ptr<AclNNTensor> aclnnTensorPtr = isWeightNz_ ? CreateWeightNzAclnnTensor(ADDMM_MAT2_ACLNN_TENSOR_IDX) :
CreateWeightAclnnTensor(ADDMM_MAT2_ACLNN_TENSOR_IDX);
if (!aclnnTensorPtr->tensor) {
ATB_LOG(ERROR) << GetLogPrefix() << "addmm mat2 aclCreateTensor failed";
return ERROR_INTERNAL_ERROR;
}
aclnnVariantPack_.aclInTensors.at(aclInTensorIndex_) = aclnnTensorPtr;
addmmMat2AclTensorIndex_ = aclInTensorIndex_++;
return NO_ERROR;
}
Status LinearAclnnRunner::CreateAddmmOutAclnnTensor()
{
ATB_LOG(INFO) << "LinearAclnnRunner::CreateAddmmOutAclnnTensor";
std::shared_ptr<AclNNTensor> aclnnTensorPtr = CreateOutputAclnnTensor(ADDMM_OUT_ACLNN_TENSOR_IDX);
if (!aclnnTensorPtr->tensor) {
ATB_LOG(ERROR) << GetLogPrefix() << "addmm out aclCreateTensor failed";
return ERROR_INTERNAL_ERROR;
}
aclnnVariantPack_.aclOutTensors.at(aclOutTensorIndex_) = aclnnTensorPtr;
addmmOutAclTensorIndex_ = aclOutTensorIndex_++;
return NO_ERROR;
}
aclnnStatus LinearAclnnRunner::SetAclnnMatmulWorkspaceExecutor()
{
ATB_LOG(INFO) << "LinearAclnnRunner::SetAclnnMatmulWorkspaceExecutor";
aclTensor *self = aclnnVariantPack_.aclInTensors.at(matmulSelfAclTensorIndex_)->tensor;
aclTensor *mat2 = aclnnVariantPack_.aclInTensors.at(matmulMat2AclTensorIndex_)->tensor;
aclTensor *out = aclnnVariantPack_.aclOutTensors.at(matmulOutAclTensorIndex_)->tensor;
int8_t cubeMathType = 1;
aclOpExecutor *rawExecutePtr = aclnnExecutor_.get();
aclnnStatus ret = aclnnMatmulGetWorkspaceSizeFunc_(self, mat2, out, cubeMathType,
&(atbVariantPack_.workspaceBufferSize), &rawExecutePtr);
aclnnExecutor_ = std::shared_ptr<aclOpExecutor>(rawExecutePtr, [this](aclOpExecutor *ptr) {
if (ptr && executorRepeatable_) {
aclDestroyAclOpExecutor(ptr);
}
});
return ret;
}
aclnnStatus LinearAclnnRunner::SetAclnnAddmmWorkspaceExecutor()
{
ATB_LOG(INFO) << "LinearAclnnRunner::SetAclnnAddmmWorkspaceExecutor";
aclTensor *self = aclnnVariantPack_.aclInTensors.at(addmmSelfAclTensorIndex_)->tensor;
aclTensor *mat1 = aclnnVariantPack_.aclInTensors.at(addmmMat1AclTensorIndex_)->tensor;
aclTensor *mat2 = aclnnVariantPack_.aclInTensors.at(addmmMat2AclTensorIndex_)->tensor;
aclTensor *out = aclnnVariantPack_.aclOutTensors.at(addmmOutAclTensorIndex_)->tensor;
if (alpha_) {
if (aclDestroyScalar(alpha_) != ACL_SUCCESS) {
ATB_LOG(ERROR) << GetLogPrefix() << "alpha aclDestroyScalar error";
return ERROR_INTERNAL_ERROR;
}
alpha_ = nullptr;
}
if (beta_) {
if (aclDestroyScalar(beta_) != ACL_SUCCESS) {
ATB_LOG(ERROR) << GetLogPrefix() << "beta aclDestroyScalar error";
return ERROR_INTERNAL_ERROR;
}
beta_ = nullptr;
}
float alphaValue = 1.0f;
alpha_ = aclCreateScalar(&alphaValue, aclDataType::ACL_FLOAT);
float betaValue = 1.0f;
beta_ = aclCreateScalar(&betaValue, aclDataType::ACL_FLOAT);
int8_t cubeMathType = 1;
aclOpExecutor *rawExecutePtr = aclnnExecutor_.get();
aclnnStatus ret = aclnnAddmmGetWorkspaceSizeFunc_(self, mat1, mat2, beta_, alpha_, out, cubeMathType,
&(atbVariantPack_.workspaceBufferSize), &rawExecutePtr);
aclnnExecutor_ = std::shared_ptr<aclOpExecutor>(rawExecutePtr, [this](aclOpExecutor *ptr) {
if (ptr && executorRepeatable_) {
aclDestroyAclOpExecutor(ptr);
}
});
return ret;
}
aclnnStatus LinearAclnnRunner::SetAclnnMatmulWeightNzWorkspaceExecutor()
{
ATB_LOG(INFO) << "LinearAclnnRunner::SetAclnnMatmulWeightNzWorkspaceExecutor";
aclTensor *self = aclnnVariantPack_.aclInTensors.at(matmulSelfAclTensorIndex_)->tensor;
aclTensor *mat2 = aclnnVariantPack_.aclInTensors.at(matmulMat2AclTensorIndex_)->tensor;
aclTensor *out = aclnnVariantPack_.aclOutTensors.at(matmulOutAclTensorIndex_)->tensor;
int8_t cubeMathType = 1;
aclOpExecutor *rawExecutePtr = aclnnExecutor_.get();
aclnnStatus ret = aclnnMatmulWeightNzGetWorkspaceSizeFunc_(self, mat2, out, cubeMathType,
&(atbVariantPack_.workspaceBufferSize), &rawExecutePtr);
aclnnExecutor_ = std::shared_ptr<aclOpExecutor>(rawExecutePtr, [this](aclOpExecutor *ptr) {
if (ptr && executorRepeatable_) {
aclDestroyAclOpExecutor(ptr);
}
});
return ret;
}
aclnnStatus LinearAclnnRunner::SetAclnnAddmmWeightNzWorkspaceExecutor()
{
ATB_LOG(INFO) << "LinearAclnnRunner::SetAclnnAddmmWeightNzWorkspaceExecutor";
aclTensor *self = aclnnVariantPack_.aclInTensors.at(addmmSelfAclTensorIndex_)->tensor;
aclTensor *mat1 = aclnnVariantPack_.aclInTensors.at(addmmMat1AclTensorIndex_)->tensor;
aclTensor *mat2 = aclnnVariantPack_.aclInTensors.at(addmmMat2AclTensorIndex_)->tensor;
aclTensor *out = aclnnVariantPack_.aclOutTensors.at(addmmOutAclTensorIndex_)->tensor;
if (alpha_) {
if (aclDestroyScalar(alpha_) != ACL_SUCCESS) {
ATB_LOG(ERROR) << GetLogPrefix() << "alpha aclDestroyScalar error";
return ERROR_INTERNAL_ERROR;
}
alpha_ = nullptr;
}
if (beta_) {
if (aclDestroyScalar(beta_) != ACL_SUCCESS) {
ATB_LOG(ERROR) << GetLogPrefix() << "beta aclDestroyScalar error";
return ERROR_INTERNAL_ERROR;
}
beta_ = nullptr;
}
float alphaValue = 1.0f;
alpha_ = aclCreateScalar(&alphaValue, aclDataType::ACL_FLOAT);
float betaValue = 1.0f;
beta_ = aclCreateScalar(&betaValue, aclDataType::ACL_FLOAT);
int8_t cubeMathType = 1;
aclOpExecutor *rawExecutePtr = aclnnExecutor_.get();
aclnnStatus ret = aclnnAddmmWeightNzGetWorkspaceSizeFunc_(self, mat1, mat2, beta_, alpha_, out, cubeMathType,
&(atbVariantPack_.workspaceBufferSize), &rawExecutePtr);
aclnnExecutor_ = std::shared_ptr<aclOpExecutor>(rawExecutePtr, [this](aclOpExecutor *ptr) {
if (ptr && executorRepeatable_) {
aclDestroyAclOpExecutor(ptr);
}
});
return ret;
}
aclnnStatus LinearAclnnRunner::SetAclnnBatchMatMulWeightNzWorkspaceExecutor()
{
ATB_LOG(INFO) << "LinearAclnnRunner::SetAclnnBatchMatMulWeightNzWorkspaceExecutor";
aclTensor *self = aclnnVariantPack_.aclInTensors.at(matmulSelfAclTensorIndex_)->tensor;
aclTensor *mat2 = aclnnVariantPack_.aclInTensors.at(matmulMat2AclTensorIndex_)->tensor;
aclTensor *out = aclnnVariantPack_.aclOutTensors.at(matmulOutAclTensorIndex_)->tensor;
int8_t cubeMathType = 1;
aclOpExecutor *rawExecutePtr = aclnnExecutor_.get();
aclnnStatus ret = aclnnBatchMatMulWeightNzGetWorkspaceSizeFunc_(
self, mat2, out, cubeMathType, &(atbVariantPack_.workspaceBufferSize), &rawExecutePtr);
aclnnExecutor_ = std::shared_ptr<aclOpExecutor>(rawExecutePtr, [this](aclOpExecutor *ptr) {
if (ptr && executorRepeatable_) {
aclDestroyAclOpExecutor(ptr);
}
});
return ret;
}
std::shared_ptr<AclNNTensor> LinearAclnnRunner::CreateXAclnnTensor(int aclnnTensorIndex)
{
Tensor atbTensor = atbVariantPack_.inTensors.at(atbInTensorIndex_++);
std::shared_ptr<AclNNTensor> aclnnTensorPtr = InitAclnnTensor(atbTensor, aclnnTensorIndex);
Dims viewShape = atbTensor.desc.shape;
if (viewShape.dimNum == 3 && !isBatch_) {
viewShape.dims[0] = viewShape.dims[0] * viewShape.dims[1];
viewShape.dims[1] = viewShape.dims[2];
viewShape.dimNum = 2;
}
aclnnTensorPtr->strides = GetCopyTensorStride(viewShape);
if (param_.transposeA) {
int64_t m = viewShape.dims[viewShape.dimNum - 1];
int64_t k = viewShape.dims[viewShape.dimNum - 2];
viewShape.dims[viewShape.dimNum - 2] = m;
viewShape.dims[viewShape.dimNum - 1] = k;
int64_t lastStride = aclnnTensorPtr->strides[viewShape.dimNum - 1];
aclnnTensorPtr->strides[viewShape.dimNum - 1] = aclnnTensorPtr->strides[viewShape.dimNum - 2];
aclnnTensorPtr->strides[viewShape.dimNum - 2] = lastStride;
}
aclnnTensorPtr->tensor = aclCreateTensor(
viewShape.dims, viewShape.dimNum, atbTensor.desc.dtype, aclnnTensorPtr->strides.data(), 0,
atbTensor.desc.format, atbTensor.desc.shape.dims, atbTensor.desc.shape.dimNum, atbTensor.deviceData);
return aclnnTensorPtr;
}
std::shared_ptr<AclNNTensor> LinearAclnnRunner::CreateWeightAclnnTensor(int aclnnTensorIndex)
{
Tensor atbTensor = atbVariantPack_.inTensors.at(atbInTensorIndex_++);
std::shared_ptr<AclNNTensor> aclnnTensorPtr = InitAclnnTensor(atbTensor, aclnnTensorIndex);
Dims viewShape = atbTensor.desc.shape;
aclnnTensorPtr->strides = GetCopyTensorStride(viewShape);
if (param_.transposeB) {
int64_t k = viewShape.dims[viewShape.dimNum - 1];
int64_t n = viewShape.dims[viewShape.dimNum - 2];
viewShape.dims[viewShape.dimNum - 2] = k;
viewShape.dims[viewShape.dimNum - 1] = n;
int64_t lastStride = aclnnTensorPtr->strides[viewShape.dimNum - 1];
aclnnTensorPtr->strides[viewShape.dimNum - 1] = aclnnTensorPtr->strides[viewShape.dimNum - 2];
aclnnTensorPtr->strides[viewShape.dimNum - 2] = lastStride;
}
aclnnTensorPtr->tensor = aclCreateTensor(
viewShape.dims, viewShape.dimNum, atbTensor.desc.dtype, aclnnTensorPtr->strides.data(), 0,
atbTensor.desc.format, atbTensor.desc.shape.dims, atbTensor.desc.shape.dimNum, atbTensor.deviceData);
return aclnnTensorPtr;
}
std::shared_ptr<AclNNTensor> LinearAclnnRunner::CreateWeightNzAclnnTensor(int aclnnTensorIndex)
{
Tensor atbTensor = atbVariantPack_.inTensors.at(atbInTensorIndex_++);
std::shared_ptr<AclNNTensor> aclnnTensorPtr = InitAclnnTensor(atbTensor, aclnnTensorIndex);
Dims viewShape = atbTensor.desc.shape;
Dims storageShape = atbTensor.desc.shape;
if (viewShape.dimNum == 4) {
Dims oldShape = viewShape;
if (isBatch_) {
viewShape.dims[0] = oldShape.dims[0];
viewShape.dims[1] = oldShape.dims[2];
viewShape.dims[2] = oldShape.dims[1] * oldShape.dims[3];
viewShape.dimNum = 3;
ATB_LOG(INFO) << GetLogPrefix() << "viewShape: [" << viewShape.dims[0] << ", " << viewShape.dims[1] << ", "
<< viewShape.dims[2] << "]";
} else {
viewShape.dims[0] = oldShape.dims[2];
viewShape.dims[1] = oldShape.dims[1] * oldShape.dims[3];
viewShape.dimNum = 2;
ATB_LOG(INFO) << GetLogPrefix() << "viewShape: [" << viewShape.dims[0] << ", " << viewShape.dims[1] << "]";
}
}
aclnnTensorPtr->strides = GetCopyTensorStride(viewShape);
if (param_.transposeB) {
int64_t k = viewShape.dims[viewShape.dimNum - 1];
int64_t n = viewShape.dims[viewShape.dimNum - 2];
viewShape.dims[viewShape.dimNum - 2] = k;
viewShape.dims[viewShape.dimNum - 1] = n;
int64_t lastStride = aclnnTensorPtr->strides[viewShape.dimNum - 1];
aclnnTensorPtr->strides[viewShape.dimNum - 1] = aclnnTensorPtr->strides[viewShape.dimNum - 2];
aclnnTensorPtr->strides[viewShape.dimNum - 2] = lastStride;
}
if (storageShape.dimNum == 2) {
Dims oldShape = storageShape;
storageShape.dims[0] = UtilsInternal::AlignUp(oldShape.dims[1], DEFAULT_ALIGN) / DEFAULT_ALIGN;
storageShape.dims[1] = UtilsInternal::AlignUp(oldShape.dims[0], DEFAULT_ALIGN) / DEFAULT_ALIGN;
storageShape.dims[2] = DEFAULT_ALIGN;
storageShape.dims[3] = DEFAULT_ALIGN;
storageShape.dimNum = 4;
ATB_LOG(INFO) << GetLogPrefix() << "storageShape: [" << storageShape.dims[0] << ", " << storageShape.dims[1]
<< ", " << storageShape.dims[2] << ", " << storageShape.dims[3] << "]";
} else if (storageShape.dimNum == 3) {
Dims oldShape = storageShape;
storageShape.dims[0] = oldShape.dims[0];
storageShape.dims[1] = UtilsInternal::AlignUp(oldShape.dims[2], DEFAULT_ALIGN) / DEFAULT_ALIGN;
storageShape.dims[2] = UtilsInternal::AlignUp(oldShape.dims[1], DEFAULT_ALIGN) / DEFAULT_ALIGN;
storageShape.dims[3] = DEFAULT_ALIGN;
storageShape.dims[4] = DEFAULT_ALIGN;
storageShape.dimNum = 5;
ATB_LOG(INFO) << GetLogPrefix() << "storageShape: [" << storageShape.dims[0] << ", " << storageShape.dims[1]
<< ", " << storageShape.dims[2] << ", " << storageShape.dims[3] << ", " << storageShape.dims[4]
<< "]";
} else if (storageShape.dimNum == 4) {
if (isBatch_) {
Dims oldShape = storageShape;
storageShape.dims[0] = oldShape.dims[0];
storageShape.dims[1] = oldShape.dims[1];
storageShape.dims[2] = UtilsInternal::AlignUp(oldShape.dims[2], DEFAULT_ALIGN) / DEFAULT_ALIGN;
storageShape.dims[3] = DEFAULT_ALIGN;
storageShape.dims[4] = DEFAULT_ALIGN;
storageShape.dimNum = 5;
ATB_LOG(INFO) << GetLogPrefix() << "storageShape: [" << storageShape.dims[0] << ", " << storageShape.dims[1]
<< ", " << storageShape.dims[2] << ", " << storageShape.dims[3] << ", "
<< storageShape.dims[4] << "]";
} else {
Dims oldShape = storageShape;
storageShape.dims[0] = oldShape.dims[1];
storageShape.dims[1] = UtilsInternal::AlignUp(oldShape.dims[2], DEFAULT_ALIGN) / DEFAULT_ALIGN;
storageShape.dims[2] = DEFAULT_ALIGN;
storageShape.dims[3] = DEFAULT_ALIGN;
storageShape.dimNum = 4;
ATB_LOG(INFO) << GetLogPrefix() << "storageShape: [" << storageShape.dims[0] << ", " << storageShape.dims[1]
<< ", " << storageShape.dims[2] << ", " << storageShape.dims[3] << "]";
}
}
aclnnTensorPtr->tensor =
aclCreateTensor(viewShape.dims, viewShape.dimNum, atbTensor.desc.dtype, aclnnTensorPtr->strides.data(), 0,
atbTensor.desc.format, storageShape.dims, storageShape.dimNum, atbTensor.deviceData);
return aclnnTensorPtr;
}
std::shared_ptr<AclNNTensor> LinearAclnnRunner::CreateBiasAclnnTensor(int aclnnTensorIndex)
{
Tensor atbTensor = atbVariantPack_.inTensors.at(atbInTensorIndex_++);
std::shared_ptr<AclNNTensor> aclnnTensorPtr = InitAclnnTensor(atbTensor, aclnnTensorIndex);
Dims viewShape = atbTensor.desc.shape;
aclnnTensorPtr->strides = GetCopyTensorStride(viewShape);
aclnnTensorPtr->tensor = aclCreateTensor(
viewShape.dims, viewShape.dimNum, atbTensor.desc.dtype, aclnnTensorPtr->strides.data(), 0,
atbTensor.desc.format, atbTensor.desc.shape.dims, atbTensor.desc.shape.dimNum, atbTensor.deviceData);
return aclnnTensorPtr;
}
std::shared_ptr<AclNNTensor> LinearAclnnRunner::CreateOutputAclnnTensor(int aclnnTensorIndex)
{
Tensor atbTensor = atbVariantPack_.outTensors.at(atbOutTensorIndex_++);
std::shared_ptr<AclNNTensor> aclnnTensorPtr = InitAclnnTensor(atbTensor, aclnnTensorIndex);
Dims viewShape = atbTensor.desc.shape;
if (viewShape.dimNum == 3 && !isBatch_) {
viewShape.dims[0] = viewShape.dims[0] * viewShape.dims[1];
viewShape.dims[1] = viewShape.dims[2];
viewShape.dimNum = 2;
}
aclnnTensorPtr->strides = GetCopyTensorStride(viewShape);
aclnnTensorPtr->tensor = aclCreateTensor(
viewShape.dims, viewShape.dimNum, atbTensor.desc.dtype, aclnnTensorPtr->strides.data(), 0,
atbTensor.desc.format, atbTensor.desc.shape.dims, atbTensor.desc.shape.dimNum, atbTensor.deviceData);
return aclnnTensorPtr;
}
std::shared_ptr<AclNNTensor> LinearAclnnRunner::InitAclnnTensor(Tensor atbTensor, int aclnnTensorIndex)
{
std::shared_ptr<AclNNTensor> aclnnTensorPtr = std::make_shared<AclNNTensor>();
aclnnTensorPtr->atbTensor = atbTensor;
aclnnTensorPtr->tensorIdx = aclnnTensorIndex;
aclnnTensorPtr->needUpdateTensorDataPtr = true;
return aclnnTensorPtr;
}
REG_RUNNER_TYPE(LinearAclnnRunner);
}