文档
注册
文档首页
评分
提单
论坛
小AI

动态Shape算子（注册算子选择器）

前提条件

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

基本原理

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

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

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

        • 函数原型： 
           1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
          		 
          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专题>样例使用”。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
 
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”个结果

当前产品无相关内容

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