文档
注册

动态Shape算子(注册算子选择器)

前提条件

在加载与执行动态Shape算子前,您需要参见TBE&AI CPU算子开发指南中的“专题 > TIK自定义算子动态Shape专题”中的说明开发自定义算子以及生成对应的二进制文件。

基本原理

动态Shape场景下,算子加载与执行的流程如下:

  1. 资源初始化,包括pyACL初始化、设置单算子模型文件的加载目录、指定用于运算的Device等。
    • 调用acl.init接口实现pyACL初始化。
    • 调用pyACL接口注册要编译的自定义算子:
      • 调用acl.op.register_compile_func接口注册算子选择器(即选择Tiling策略的函数),用于在算子执行时,能针对不同Shape,选择相应的Tiling策略。

        算子选择器需由用户提前定义并实现:

        • 函数原型:
          op_selector(in_num, in_desc, out_num, out_desc, op_attr, op_kernel_desc)
              """
              算子选择器:函数和参数名称可自定义,参数个数和类型必须匹配。
              :param in_num: 输入tensor描述数。
              :param in_desc: 输入tensor描述list。
              :param out_num: 输出tensor描述数。
              :param out_desc: 输出tensor描述list。
              :param op_attr: 算子属性的地址对象,用于设置算子属性信息。
              :param op_kernel_desc:算子内核描述地址对象,用于动态Shape场景下,设置算子Workspace参数。
              :return:
              """
        • 函数实现:

          用户自行编写代码逻辑实现Tiling策略选择、Tiling参数生成,并将调用acl.op.set_kernel_args接口,设置算子Tiling参数、执行并发数等。

      • 调用acl.op.create_kernel接口将算子注册到系统内部,用于在算子执行时,查找到算子实现代码。
    • 调用acl.rt.set_device接口指定运算的Device。
    • 调用acl.rt.create_context接口显式创建一个Context,调用acl.rt.create_stream接口显式创建一个Stream。

      若没有显式创建Stream,则使用默认Stream,默认Stream是在调用acl.rt.set_device接口时隐式创建的,默认Stream作为接口入参时,直接传0。

  2. 用户自行构造算子描述信息(输入输出Tensor描述、算子属性等)、申请存放算子输入输出数据的内存。
  3. 将算子输入数据从Host复制到Device上。
    • 调用acl.rt.memcpy接口实现同步内存复制,内存使用结束后需及时释放。
    • 调用acl.rt.memcpy_async接口实现异步内存复制,内存使用结束后需及时释放。
  4. 编译单算子。

    调用acl.op.update_params接口编译指定算子,触发算子选择器的调用逻辑。

  5. 执行单算子。

    调用acl.op.execute_v2接口加载并执行算子。

  6. 将算子运算的输出数据从Device上复制到Host上(提前申请Host上的内存)。
    • 调用acl.rt.memcpy接口实现同步内存复制,内存使用结束后需及时释放。
    • 调用acl.rt.memcpy_async接口实现异步内存复制,内存使用结束后需及时释放。
  7. 按顺序先释放Stream资源,再释放Context资源,最后释放Device资源。
  8. 调用acl.finalize接口实现pyACL去初始化。

示例代码

在样例代码中,调用接口后,需增加异常处理的分支,并记录报错日志、提示日志,此处不一一列举。以下是关键步骤的代码示例,不可以直接拷贝编译运行,仅供参考。

动态Shape算子(注册算子选择器)样例的获取、编译运行请参见TBE&AI CPU算子开发指南中的“专题 > TIK自定义算子动态Shape专题>样例使用”

import acl
import numpy as np
# ......

ACL_ENGINE_AICORE = 1
ACL_FLOAT16 = 1
ACL_FORMAT_ND = 2
ACL_MEM_MALLOC_NORMAL_ONLY = 2
ACL_MEMCPY_HOST_TO_DEVICE = 1
ACL_MEMCPY_DEVICE_TO_HOST= 2

device_id = 0

# 1.资源初始化。
ret = acl.init()
ret = acl.rt.set_device(device_id)
stream, ret = acl.rt.create_stream()
ret = acl.op.register_compile_func("add", op_select)
# 将算子Kernel的*.o文件需要用户提前编译好,并调用numpy加载.o文件并转成地址对象,op_data_size_0表示第一个.o文件占用的内存大小。
# 如果有多个算子Kernel的*.o文件,需要多次调用该接口。
ret = acl.op.create_kernel("add", "cce_add_11_33_float16_11_33_float16__kernel0", "cce_add_11_33_float16_11_33_float16__kernel0", np_op_0_ptr, op_data_size_0, ACL_ENGINE_AICORE, 0)

# 2.构造add算子的输入输出Tensor、输入输出Tensor描述,并申请存放算子输入数据、输出数据的内存。
# 输入参数。
a = np.random.rand(2, 1).astype(np.float16)
b = np.random.rand(2, 1).astype(np.float16)
bytes_data = a.tobytes()
a_ptr = acl.util.bytes_to_ptr(bytes_data)
bytes_data = b.tobytes()
b_ptr = acl.util.bytes_to_ptr(bytes_data)
input_desc_list = [acl.create_tensor_desc(ACL_FLOAT16, [2, 1], ACL_FORMAT_ND), acl.create_tensor_desc(ACL_FLOAT16, [2, 1], ACL_FORMAT_ND)]
output_desc_list = [acl.create_tensor_desc(ACL_FLOAT16, [2, 1], ACL_FORMAT_ND)]
# 申请device侧内存。
size_a = acl.get_tensor_desc_size(input_desc_list[0])
size_b = acl.get_tensor_desc_size(input_desc_list[1])
size_c = acl.get_tensor_desc_size(output_desc_list[0])
dev_a, ret = acl.rt.malloc(size_a, ACL_MEM_MALLOC_NORMAL_ONLY)
dev_b, ret = acl.rt.malloc(size_b, ACL_MEM_MALLOC_NORMAL_ONLY)
dev_c, ret = acl.rt.malloc(size_c, ACL_MEM_MALLOC_NORMAL_ONLY)

# 3.将算子输入数据从Host复制到Device上。
ret = acl.rt.memcpy(dev_a, size_a, a_ptr, size_a, ACL_MEMCPY_HOST_TO_DEVICE)
ret = acl.rt.memcpy(dev_b, size_b, b_ptr, size_b, ACL_MEMCPY_HOST_TO_DEVICE)

# 4.调用acl.op.update_params接口编译算子。
op_attr = acl.op.create_attr()
ret = acl.op.update_params("add", input_desc_list, output_desc_list, op_attr )

# 5.调用acl.op.execute接口加载并执行算子。
in_data_list = [acl.create_data_buffer(dev_a, size_a), acl.create_data_buffer(dev_b, size_b)]
out_data_list = [acl.create_data_buffer(dev_c, size_c)]
ret = acl.op.execute_v2("add", input_desc_list, in_data_list, output_desc_list, out_data_list, op_attr, stream)

# 6.将算子运算的输出数据从Device上复制到Host上(提前申请Host上的内存)。
host_ptr, ret = acl.rt.malloc_host(size_c)
ret = acl.rt.memcpy(host_ptr, size_c, dev_c, size_c, ACL_MEMCPY_DEVICE_TO_HOST)
bytes_out = acl.util.ptr_to_bytes(host_ptr, size_c)
out_np = np.frombuffer(bytes_out, dtype=np.byte).reshape((size_c,))

# 7.按顺序释放资源。
# 7.1 释放算子的输入输出Tensor描述。
# 7.2 释放Host上的内存。
# 7.3 释放Device上的内存。
# 7.4 释放设备管理资源。
ret = acl.rt.destroy_stream(stream)
ret = acl.rt.reset_device(device_id)
ret = acl.finalize()
# ......
搜索结果
找到“0”个结果

当前产品无相关内容

未找到相关内容,请尝试其他搜索词