Matmul算子(TIK方式)
功能描述
本样例使用TIK方式进行了matmul算子的实现,给定两个输入的张量A和B,进行矩阵相乘,输出结果张量。
说明:此样例未实现算子适配插件,若开发者想在原始框架为caffe的网络模型中使用此自定义算子,需要自定义实现对应的算子适配插件。
算子分析
使用TIK方式开发matmul算子前,我们需要确定算子功能、输入、输出,算子开发方式、算子类型以及算子实现函数名称等。
- 明确算子的功能以及数学表达式。
matmul算子的数学表达式为:
y=x1*x2
计算过程是:将两个输入矩阵相乘,得到最终结果矩阵z并将其返回。
- 明确输入和输出。
- matmul算子有两个输入:x1与x2,输出为y。
- 算子输入的数据类型为float16,输出的数据类型为float32。
- 算子输入支持固定shape,输入shape分别为(16,64)和(64,1024),输出shape为(16,1024)。
- 算子输入支持的format为:ND。
- 确定算子开发方式及使用的计算接口。
- 计算过程中主要涉及矩阵乘法操作,初分析可使用matmul()接口实现矩阵x1*x2的操作。
- 由于在整个matmul的计算过程中会涉及到数据搬运操作,可使用data_move()接口把数据从Global Memory搬迁到L1 Buffer中。
- 矩阵计算完成后,对结果进行处理,可使用fixpipe()把数据从L1OUT Buffer搬迁到Global Memory中。
- 明确算子实现文件名称、算子实现函数名称以及算子的类型(OpType)。
- 算子类型需要采用大驼峰的命名方式,即采用大写字符区分不同的语义。
- 算子文件名称和算子函数名称,可选用以下任意一种命名规则:
- 用户自定义,此时需要在算子信息定义中配置opFile.value与opInterface.value。
- 不配置算子信息定义中的opFile.value与opInterface.value,FE会将OpType按照如下方式进行转换后进行算子文件名和算子函数名的匹配。转换规则如下:
- 首字符的大写字符转换为小写字符。
例如:Abc -> abc
- 小写字符后的大写字符转换为下划线+小写字符。
例如:AbcDef -> abc_def
- 紧跟数字以及大写字符后的大写字符,作为同一语义字符串,查找此字符串后的第一个小写字符,并将此小写字符的前一个大写字符转换为下划线+小写字符,其余大写字符转换为小写字符。若此字符串后不存在小写字符,则直接将此字符串中的大写字符转换为小写字符。
例如:ABCDef -> abc_def;Abc2DEf -> abc2d_ef;Abc2DEF -> abc2def;ABC2dEF -> abc2d_ef。
- 首字符的大写字符转换为小写字符。
本例中,算子类型定义为MatmulTik;算子的实现文件名称及实现函数名称定义为matmul_tik。通过以上分析,得到MatmulTik算子的设计规格如下:
表14-6 MatmulTik算子设计规格算子类型(OpType)
MatmulTik
算子输入
name:x1
shape:(16,64)
data type:float16
format:ND
name:x2
shape:(64,1024)
data type:float16
format:ND
算子输出
name:y
shape:(16,1024)
data type:float32
format:ND
算子实现使用主要TIK接口
matmul()
data_move()
fixpipe()
算子实现文件/实现函数名称
matmul_tik
算子实现
本章节介绍算子实现中的关键功能点。
算子代码实现
- 首先取出算子输入的shape以及dtype,并配置好对应的tiling参数和循环次数,其中tiling_size表示方向切分大小;cycle_size表示方向的循环次数(例:m = m_tiling_size*m_cycle_num),thread_num表示方向是否开启double buffer功能,将参数统一配置后,调用算子计算函数。
def matmul_tik(input_x1, input_x2, output_y={}, kernel_name="simple_matmul"): shape_a = input_x1.get("ori_shape") shape_b = input_x2.get("ori_shape") m = shape_a[0] k = shape_a[1] n = shape_b[1] data_type = input_x1.get("dtype").lower() params = { 'M': m, 'K': k, 'N': n, 'data_type': data_type, 'm_tiling_size': 16, 'm_cycle_times': 1, 'm_thread_num': 1, 'n_tiling_size': 64, 'n_cycle_times': 16, 'n_thread_num': 1, 'k_tiling_size': 32, 'k_cycle_times': 2, 'k_thread_num': 2 } matmul_tik_compute(params, kernel_name)
- 算子计算函数的实现逻辑如下所示。
- 对输入tensor进行占位。
- 根据设置的参数通过for_range( )循环实现tiling。
- 调用matmul API实现矩阵相乘计算。
- 调用data_move、fixpipe API实现计算结果数据搬运。
def matmul_tik_compute(params, kernel_name, new_ws=None, cnt=0): te_set_l2_mode(1) tik_instance = tik.Tik() if not isinstance(params, dict): params = params.__dict__ m, k, n = params['M'], params['K'], params['N'] data_type = params["data_type"] m_tiling_size = int(params["m_tiling_size"]) n_tiling_size = int(params["n_tiling_size"]) k_tiling_size = int(params['k_tiling_size']) m_cycle_times = params["m_cycle_times"] n_cycle_times = params["n_cycle_times"] k_cycle_times = params["k_cycle_times"] if data_type == "float16": C_loc_out_type = "float32" K0 = 16 else: C_loc_out_type = "int32" K0 = 32 block_size = 16 n_thread_num = params['n_thread_num'] m_thread_num = params['m_thread_num'] k_thread_num = params['k_thread_num'] # 对输入输出tensor进行占位 C_gm = tik_instance.Tensor(C_loc_out_type, (n // block_size, m, block_size), name="C_gm", scope=tik.scope_gm) A_gm = tik_instance.Tensor(params["data_type"], (k//K0, m, K0), name="A_gm", scope=tik.scope_gm) B_gm = tik_instance.Tensor(params["data_type"], (k//K0, n, K0), name="B_gm", scope=tik.scope_gm) # 执行循环,开启double buffer和多核 with tik_instance.for_range(0, 2, block_num=2) as core_id: with tik_instance.for_range(0, n_cycle_times//2, thread_num=n_thread_num) as n_idx: with tik_instance.for_range(0, m_cycle_times, thread_num=m_thread_num) as m_idx: # 根据切分大小定义L1OUT Buffer中的Tensor dst_l0c = tik_instance.Tensor(C_loc_out_type, [n_tiling_size//16, m_tiling_size, 16], name='dst_l0c', scope=tik.scope_cbuf_out) # k方向进行循环,分别对A、B进行数据搬运 with tik_instance.for_range(0, k_cycle_times, thread_num=k_thread_num) as k_idx: A_l1 = tik_instance.Tensor(params['data_type'], [k_tiling_size//K0, m_tiling_size, K0], name="A_tiling_l1", scope=tik.scope_cbuf) tik_instance.data_move(A_l1, A_gm[k_idx * k_tiling_size // K0, m_idx*m_tiling_size, :],0, k_tiling_size//K0, m_tiling_size, m - m_tiling_size, 0) B_l1 = tik_instance.Tensor(params["data_type"], [k_tiling_size//K0, n_tiling_size, K0], name="B_tiling_l1", scope=tik.scope_cbuf) if n-n_tiling_size>65535: with tik_instance.for_range(0, k_tiling_size//K0) as dma_k_idx: tik_instance.data_move(B_l1[dma_k_idx, :, :], B_gm[k_idx*k_tiling_size//K0 + dma_k_idx, (core_id*n_cycle_times//2+n_idx)*n_tiling_size, :], 0, 1, n_tiling_size, 0, 0) else: tik_instance.data_move(B_l1, B_gm[k_idx*k_tiling_size//K0, (core_id*n_cycle_times//2+n_idx)*n_tiling_size, :], 0, k_tiling_size//K0, n_tiling_size, n-n_tiling_size, 0) # 调用matmul API进行计算 with tik_instance.if_scope(k_idx == 0): tik_instance.matmul(dst_l0c, A_l1, B_l1, m_tiling_size, k_tiling_size, n_tiling_size, init_l1out=True) with tik_instance.else_scope(): tik_instance.matmul(dst_l0c, A_l1, B_l1, m_tiling_size, k_tiling_size, n_tiling_size, init_l1out=False) # 矩阵计算完成后,对结果进行处理,把数据搬迁到Global Memory中 tik_instance.fixpipe(C_gm[n_tiling_size//16*(core_id*n_cycle_times//2+n_idx), m_idx*m_tiling_size, :], dst_l0c, n_tiling_size//16, m_tiling_size*16*DTYPE_SIZE[C_loc_out_type]//32, (m-m_tiling_size)*16*DTYPE_SIZE[C_loc_out_type]//32, 0)
- 调用BuildCCE()进行编译。
tik_instance.BuildCCE(kernel_name=kernel_name, inputs=[A_gm, B_gm], outputs=[C_gm])
算子适配插件实现
此样例未实现算子适配插件,若开发者想在第三方网络模型中使用此自定义算子,需要自定义实现对应的算子适配插件。
算子原型定义
MatumulTik算子的原型定义详细代码可参见“op_proto/matmul_tik.h”与“op_proto/matmul_tik.cpp”文件。
matmul_tik.h对MatmulTik算子进行原型定义。
namespace ge { REG_OP(MatmulTik) .INPUT(x1, TensorType({DT_FLOAT, DT_FLOAT16, DT_INT32})) .INPUT(x2, TensorType({DT_FLOAT, DT_FLOAT16, DT_INT32})) .OUTPUT(y, TensorType({DT_FLOAT, DT_FLOAT16, DT_INT32})) .OP_END_FACTORY_REG(MatmulTik) }
matmul_tik.cpp根据算子的输入张量描述、算子逻辑及算子属性,推理出算子的输出张量描述,包括张量的形状、数据类型及数据排布格式等信息。这样算子构图准备阶段就可以为所有的张量静态分配内存,避免动态内存分配带来的开销。
namespace ge { IMPLEMT_VERIFIER(MatmulTik, MatmulTikVerify) { std::vector<DataType> support_list; support_list.push_back(DT_FLOAT16); support_list.push_back(DT_FLOAT); support_list.push_back(DT_INT32); support_list.push_back(DT_INT8); support_list.push_back(DT_UINT8); return GRAPH_SUCCESS; } // Obtains the processing function of the output tensor description. IMPLEMT_COMMON_INFERFUNC(MatmulTikInferShape) { TensorDesc tensordesc_output = op.GetOutputDesc("y"); ge::TensorDesc inputTensorDescX = op.GetInputDesc("x1"); ge::TensorDesc inputTensorDescY = op.GetInputDesc("x2"); ge::Shape shapeX = inputTensorDescX.GetShape(); ge::Shape shapeY = inputTensorDescY.GetShape(); DataType dtype = inputTensorDescX.GetDataType(); bool transposeA = false; bool transposeB = false; std::vector<int64_t> dimVector; dimVector.push_back(shapeX.GetDim(0)); dimVector.push_back(shapeY.GetDim(1)); ge::Shape outputShape(dimVector); tensordesc_output.SetShape(outputShape); tensordesc_output.SetDataType(op.GetInputDesc("x1").GetDataType()); (void)op.UpdateOutputDesc("y", tensordesc_output); return GRAPH_SUCCESS; } //Registered inferfunction COMMON_INFER_FUNC_REG(MatmulTik, MatmulTikInferShape); //Registered verify function VERIFY_FUNC_REG(MatmulTik, MatmulTikVerify); }
算子信息定义
MatmulTik算子的信息定义文件请参见“tbe/op_info_cfg/ai_core/<soc_version>/matmul_tik.ini”。
说明:当前MatmulTik算子仅提供了昇腾310 AI处理器上的信息定义文件,若开发者想将此算子运行于其他昇腾AI处理器,请参考实现对应的算子信息定义文件并将其存放到对应的<soc_version>目录下。