TVM API

TVM API Explaination

Posted by Treaseven on December 25, 2024

TVM学习资源

https://www.zhihu.com/people/yang-cheng-68-6/posts?page=2
http://giantpandacv.com/project/%E9%83%A8%E7%BD%B2%E4%BC%98%E5%8C%96/%E6%B7%B1%E5%BA%A6%E5%AD%A6%E4%B9%A0%E7%BC%96%E8%AF%91%E5%99%A8/TVM%20%E5%AD%A6%E4%B9%A0%E6%8C%87%E5%8D%97/
https://discuss.tvm.apache.org/t/tvm/15159/16
https://www.cnblogs.com/wanger-sjtu
https://tvm.hyper.ai/docs/
https://tvm.d2l.ai/

tvm

  • tvm.error
  • tvm.ir
  • tvm.instrument
  • tvm.transform
    class tvm.transform.Sequential: 在TVM中实现一个序列化的编译优化通道管理器,将多个编译优化通道(Pass)按顺序组合起来,形成一个复合的优化流程 class Sequential(Pass)
    class tvm.transform.PassContext: 提供控制和配置编译优化流程的上下文 class PassContext(tvm.runtime.Object)
  • tvm.target
    class tvm.target.Target: 用于描述和管理目标设备的信息 tvm.target.Target(target, host=None)
  • tvm.driver

tvm.runtime

  • tvm.runtime
  • tvm.runtime.ndarray
  • tvm.runtime.relax_vm
  • tvm.runtime.disco
  • tvm.runtime.profiling

tvm.relax

  • tvm.relax
  • tvm.relax.analysis
  • tvm.relax.block_builder
  • tvm.relax.frontend
  • tvm.relax.op
  • tvm.relax.transform

tvm.tir

  • tvm.tir tvm.tir.IntImm: 表示整数常量的类 class tvm.tir.IntImm(dtype:str, value:int, span:Span|None = None)
  • tvm.tir.analysis
  • tvm.tir.schedule
  • tvm.tir.stmt_functor
  • tvm.tir.transform

tvm.te

  • tvm.te
    tvm.te.create_schedule: 创建计算调度,接收一个或多个操作作为输入,返回一个schedule对象 tvm.te.create_schedule(ops)
    tvm.te.placeholder: 创建一个占位符张量 tvm.te.placeholder(shape, dtype=None, name=’placeholder’)
    tvm.te.compute: 定义张量计算操作的核心函数 tvm.te.compute(shape, fcompute, name=’compute’, tag=’’, attrs=None, varargs_names=None)
    tvm.te.reduce_axis: 创建归约操作的迭代轴 tvm.te.reduce_axis(dom, name=’rv’, thread_tag=’’, span=None)
    tvm.te.sum: 执行归约求和操作的函数 tvm.te.sum(expr, axis, where=None, init=None, *args)
  • tvm.te.hybrid
  • tvm.topi from tvm.topi import tag tag.is_injective: 用来标记和判断张量运算是否属于injective类型的函数, injective类型的函数是指输出元素的计算只依赖于输入张量中固定位置的元素,且不存在数据依赖关系如add、multiply、sigmoid、reshape、transpose、broadcast

tvm.meta_schedule

  • tvm.meta_schedule

tvm.dlight

  • tvm.dlight

Misc

  • tvm.rpc
  • tvm.contrib

Legacy

  • tvm.relay
    tvm.relay.var: 定义输入变量, tvm.relay.var(name_hint, type_annotation=None, shape=None, dtype=’float32’, span=None)
    tvm.relay.greater: 判断左侧值是否严格大于右侧值,返回一个布尔变量,tvm.relay.greater(lhs, rhs)
    tvm.relay.greater_equal: 判断左侧值是否大于或等于右侧值,返回一个布尔变量,tvm.relay.greater_equal(lhs, rhs)
    class tvm.relay.Function: 创建函数表达式的类,定义一个可计算的函数 class tvm.relay.Function(params, body, ret_type=None, type_params=None, attrs=None, span=None) 其方法 astext,用于获取函数的文本格式表示
    tvm.relay.build: 将Relay IR模块编译成可执行的形式 tvm.relay.build(ir_mod, target=None, target_host=None, executor=graph{“link-params”:T.bool(False)}, runtime=cpp, workspace_memory_pools=None, constant_memory_pools=None, params=None, mod_name=’default’)
    tvm.relay.abs: 对输入的张量数据进行逐元素的绝对值运算 tvm.relay.abs(data)
    tvm.relay.sum: 计算数组元素求和的函数 tvm.relay.sum(data, axis=None, keepdims=False, exclude=False)
    tvm.relay.divide: 对两个张量执行逐元素除法 tvm.relay.divide(lhs, rhs)
    tvm.relay.sqrt: 计算张量元素平方根的函数 tvm.relay.sqrt(data)
    tvm.relay.power: 计算幂运算的函数 tvm.relay.power(lhs, rhs)
    tvm.relay.unique: 查找一维张量中唯一元素的函数 tvm.relay.unique(data, is_sorted=True, return_counts=False)
    tvm.relay.multiply: 执行乘法运算的函数,支持广播机制 tvm.relay.multiply(lhs, rhs)
    tvm.relay.substract: 执行减法运算的函数,支持广播机制 tvm.relay.subtract(lhs, rhs)
    tvm.relay.argmax: 用于沿指定的轴计算输入数据中的最大值的索引 tvm.relay.argmax(data, axis=None, keepdims=False, exclude=False, select_last_index=False)
    tvm.relay.argmin: 用于沿指定的轴计算输入数据中的最小值的索引 tvm.relay.argmin(data, axis=None, keepdims=False, exclude=False, select_last_index=False)
    tvm.relay.take: 从输入数组data中,根据指定的索引indices提取元素 tvm.relay.take(data, indices, axis=None, batch_dims=0, mode=’clip’)
  • tvm.relay.frontend tvm.relay.frontend.from_pytorch: 将Pytorch模型转换为TVM Relay IR tvm.relay.frontend.from_pytorch(script_module, input_infos, custom_convert_map=None, default_dtype=”float32”, use_parser_friendly_name=False, keep_quantized_weight=False)
  • tvm.relay.nn
    tvm.relay.nn.sparse_dense: 执行稀疏矩阵与密集矩阵相乘的操作,返回类型为tvm.relay.Expr的计算结果矩阵 tvm.relay.nn.sparse_dense(dense_mat, sparse_mat, sparse_lhs=False)
    tvm.relay.nn.dense: 矩阵乘法运算 tvm.relay.nn.dense(data, weight, units=None, out_dtype=’’)
    tvm.relay.nn.bias_add: 将一维偏置张量加到输入数据的指定轴上 tvm.relay.nn.bias_add(data, bias, axis=1)
  • tvm.relay.vision
  • tvm.relay.image
  • tvm.relay.transform
    tvm.relay.transform.InferType(): 用于推断一个relay.Expr表达式的类型 tvm.relay.transform.InferType()
    tvm.relay.transform.RemoveUnusedFunctions: 创建一个TVM优化通道(Pass),用于移除Relay模块中未被使用的全局函数 tvm.realy.transform.RemoveUnusedFunctions(entry_functions=None)
    tvm.relay.transform.ConvertLayout: 创建一个用于转换数据布局的优化通道,解决不同硬件对数据布局有不同偏好的问题 tvm.relay.transform.ConvertLayout(desired_layouts)
  • tvm.relay.analysis
    tvm.relay.analysis.free_type_vars: 获取表达式或类型中的自由类型变量(free type variables)、以后序深度优化搜索顺序返回
    tvm.relay.analysis.free_vars: 获取表达式中的自由变量(free variables)、以后序深度优先搜索顺序返回
  • tvm.relay.backend
  • tvm.relay.dataflow_pattern
  • tvm.relay.testing
  • tvm.autotvm
    class tvm.autotvm.measure.measure_methods.LocalBuilder: 在本地机器上执行TVM编译任务,用于自动调优过程中,负责将优化后的计算图编译成可执行代码 tvm.autotvm.measure.measure_methods.LocalBuilder(timeout=10, n_parallel=None, build_kwargs=None, build_func=’default’, do_fork=False, runtime=None)
    class tvm.autotvm.measure.measure_methods.RPCRunner: 用于远程执行和测量性能的关键组件 tvm.autotvm.measure.measure_methods.RPCRunner(key, host, port, priority=1, timeout=10, n_parallel=None, number=4, repeat=3, min_repeat_ms=0, coldown_interval=0.1, enable_cpu_cache_flush=False, module_loader=None)
    class tvm.autotvm.measure.measure_methods.LocalRunner: 用于在本地设备上执行和测量代码性能的组件 tvm.autotvm.measure.measure_methods.LocalRunner(timeout=10, number=4, repeat=3, min_repeat_ms=0, coldown_interval=0.1, enable_cpu_cache_flush=False, module_loader=None)
    class tvm.autotvm.measure.measure_option: 用于配置性能测量选项的函数,将构建(builder)和运行(runner)的配置组合在一起,形成完整的测量方案 tvm.autotvm.measure.measure_option(builder, runner)
    tvm.autotvm.measure.create_measure_batch: 创建性能测量函数的工具,为自动调优系统提供性能评估能力 tvm.autotvm.measure.create_measure_batch(task, option)
    tvm.autotvm.tuner.RandomTuner: 以随机顺序枚举空间
    tvm.autotvm.tuner.GridSearchTuner: 以网格顺序枚举空间
    tvm.autotvm.tuner.GATuner: 使用遗传算法搜索空间
    tvm.autotvm.tuner.XGBTuner: 用基于模型的方法训练一个XGBoost模型,来预测降级IR的速度,并根据预测值选择下一批配置
    tvm.autotvm.task.extract_from_program: 从一个Relay程序中提取可调优的任务 tvm.autotvm.task.extract_from_program(mod, params, target, target_host=None, ops=None)
    from tvm.autotvm.measure.measure import MeasureInput: from tvm.autotvm.measure.measure import create_measure_batch:
  • tvm.auto_scheduler
    from tvm.auto_scheduler import register_workload_tensors: 接受一组张量(输入、输出),自动提取这些张量对应的计算表达式,将这些张量的形状、数据类型和计算关系注册为唯一的工作负载标识符
    tvm.auto_scheduler.SearchTask: 定义和管理自动调度优化任务的核心类,作用是封装一个需要优化的计算任务,并为该任务自动搜索最优的硬件调度策略
    tvm.auto_scheduler.task_scheduler.make_search_policies: 为多个优化任务生成搜索策略,指导任务调度器如何高效分配资源、选择任务优化顺序,并动态调整每个任务的搜索过程 tvm.auto_scheduler.extract_tasks: 从给定的计算图中识别和提取可以优化的计算任务 tvm.auto_scheduler.extrac_tasks(mod, params, target, target_host=None, hardware_params=None, include_simple_tasks=False, dump_workload_to_dag_log=None, opt_level=3, other_targets=None)
    class tvm.auto_scheduler.TaskScheduler: 用于在同时调优多个任务时分配时间资源 tvm.auto_scheduler.TaskScheduler(tasks, task_weights=None, objective_func=None, strategy=’gradient’, load_model_file: str|None = None, load_log_file: str|None = None, alpha:float = 0.2, beta:float = 2, gamma:float = 0.5, backward_window_size: int = 3, callbacks=None) 其方法 tune(tune_option, search_policy=’default’, search_policy_params=None, adaptive_training=False, per_task_early_stopping=None)
    class tvm.auto_scheduler.TuningOptions: 用于控制性能调优选项的配置 tvm.auto_scheduler.TuningOptions(num_measure_trials=0, early_stopping=None, num_measures_per_round=64, verbose=1, builder=’local’, runner=’local’, measure_callbacks=None)
    class tvm.auto_scheduler.LocalRunner: 用于在本地CPU/GPU上测量程序运行时间 classtvm.auto_scheduler.LocalRunner(timeout=10, number=3, repeat=1, min_repeat_ms=100, cooldown_interval=0.0, enable_cpu_cache_flush=False, device=0)
    class tvm.auto_scheduler.RecordToFile: 用于将测量记录写入文件 tvm.auto_scheduler.RecordToFile(filename)
    class tvm.auto_scheduler.ApplyHistoryBest 用于应用历史最佳配置的类 tvm.auto_scheduler.ApplyHistoryBest(records, n_lines=None, include_compatible=False) 其方法 获取特定目标和工作负载的记录条目 static get_workload_entry(best_records, target_key, workload_key) 加载调优记录 load(records, n_lines=None) 更新工作负载的配置 update(target, workload_key, state)
  • tvm.contrib.graph_executor
    class tvm.contrib.graph_executor.GraphModule: 运行编译后的模型 tvm.contrib.graph_executor.GraphModule(module),其方法 set_input(key=None, value=None, **params) 设置输入数据 set_input_zero_copy(key=None, value=None, **params) 零拷贝方式设置输入 set_output_zero_copy(key, value) 设置零拷贝输出 run(**input_dict) 获取输出 get_num_outputs() 获取输出数量 get_num_inputs() 获取输入数量 get_input(index, out=None) 获取指定输入 get_input_index(name) 获取输入索引 get_input_info() 获取输入信息 get_output(index, out=None) 获取输出 debug_get_output(node, out) 获取中间节点输出 load_params(params_bytes) 加载参数 share_params(other, params_bytes) 共享参数 benchmark(device, func_name=’run’, repeat=5, min_repeat_ms=None, limit_zero_time_iterations=100, end_to_end=False, cooldown_interval_ms=0, repeats_cooldown=1, **kwargs) 性能基准测试

Other API

  • tvm.IRModule
    tvm.IRModule.from_expr: 将Relay表达式转换为IR Module(中间表示模块)的函数 tvm.IRModule.from_expr(expr)
  • tvm.transform.PassContext
    tvm.transform.PassContext(opt_level=3): 创建一个优化上下文,优化级别范围是0-3
  • tvm.build
    tvm.build: 将计算描述编译成可执行的代码 tvm.build(schedule, args, target=”llvm”, target_host=None, name=”default_functino”, binds=None) 返回值是tvm.runtime.Module对象,其属性和功能: 时间测量 func.time_evaluator(func.entry_name, tvm.cpu(), number=10) 获取入口函数名 func.entry_name 获取已导出的函数列表 func.get_function_list() 获取特定函数 func.get_function(“function_name”) 保存到文件 func.save(“compiler_func.o”) 加载模块 func.export_library(“compiled_lib.so”) 从文件加载 tvm.runtime.load_module(“compiled_lib.so”) 获取源代码 func.get_source() 获取IR代码 func.get_ir() 设置线程数 func.set_num_threads 配置GPU工作组大小 func.set_cuda_props(max_shared_memory_per_block=49152) 模块类型 func.type_key 代码格式 func.format 是否为二进制格式 func.is_binary
  • tvm.lower
    tvm.lower: 将高层次的计算图转换为更底层的具体实现代码 tvm.lower(schedule, inputs, simple_mode, name)

other important things

leaf_iter_vars: 叶子迭代变量
for v in stage.leaf_iter_vars:
    print("变量名", v.var.name)
    print("迭代范围", v.dom.min, v.dom.extent)
    print("迭代类型", v.iter_type)
    print("循环类型", v.for_type)
    print("线程标签", v.thread_tag)
stage: 表示计算阶段的重要对象
属性相关内容
- op相关
* stage.op: 当前阶段的操作
    * op.name: 操作的名称
    * op.axis: 空间维度的迭代轴
    * op.reduce_axis: 归约维度的迭代轴
    * op.input_tensors: 输入张量
    * op.output: 输出张量
- 迭代变量
* stage.leaf_iter_vars: 叶子迭代变量列表
* stage.all_iter_vars: 所有迭代变量列表
* stage.leaf_var_map: 叶子变量映射
- 调度信息
* stage.attach_type: 绑定类型
* stage.attach_ivar: 绑定的迭代变量
* stage.attach_stage: 绑定的计算阶段
* stage.store_pred: 存储谓词
* stage.compute_pred: 计算谓词
- 优化相关
* stage.schedule: 所属的调度对象
* stage.group: 计算组信息
* stage.scope: 存储域范围
* stage.is_output: 是否是输出阶段
* stage.is_specialized: 是否是特化的阶段
op.axis: 计算定义中的原始轴,反映计算的数学结构
leaf_iter_vars: 经过调度变换后的最终迭代变量,反映了实际代码生成时的循环结构
-----------------------------------------
常用方法
outer, inner = stage.split(axis, factor) # 分割迭代轴
stage.reorder(axis1, axis2, ...) # 重排序迭代轴
stage.compute_at(other_stage, axis) # 计算绑定
stage.store_at(other_stage, axis) # 存储绑定
stage.parallel(axis) # 并行化
stage.vectorize(axis) # 向量化
stage.unroll(axis) # 展开
stage.cache_read(tensor, scope, readers) # 缓存数据
stage.cache_write(tensor, scope) # 缓存数据
stage.fuse(axis1, axis2) # 融合操作
stage.bind(axis, thread_type) # 线程绑定

TVM学习帖子

为什么需要编译器?
- 性能:编译器或者推理引擎推理效率更高;一方面,pytorch执行推理的时候,计算图其实有点类似解析执行的,而编译器或者推理引擎要干的事情,是把计算图编译成目标硬件上可高效执行的代码;另一方面,pytorch训练框架自带的推理,是不做任何优化的,但是编译器加入了很多优化:1.面向图结构的:算符融合、死代码删除、常量折叠;2.面向体系结构的:prefetch、读写缓存、并行计算
- 异构硬件:编译器把图转换成一系列的分层抽象的IR,在不同的层次完成不同的优化工作,在最底层完成到目标硬件的代码生成

IR公共的一些核心概念
编程语法最基本的核心概念3个:类型、运算符、表达式,在IR这里分别对应Type、Op、Expr
- Type:PrimType(最原始的Type,可以直接映射到low-level IR的基本数据类型)、FuncType(函数类型)
- Expr: PrimExpr(原始Expr,主要在tir模块中定义,可以相对直接地映射到low-level code)、RelayExpr(所有的非PrimExpr,比如tensor,function,adt等其他一等公民)、GlobalVar(全局变量,只有函数才会引用GlobalVar,通常用来实现函数的递归调用)、IntImm(Int64常量)、FloatImm(double变量)、Bool(布尔变量)、Integer(继承自IntImm)、Range(表示范围)
- Op: 表示所有系统定义的原始算子/内联函数的通用类
TVM代码库结构概述
- src:用于算子编译和部署runtime的C++代码
- src/relay: Relay的实现,一种用于深度学习框架的新功能IR;管理计算图的组件,图结构中的节点使用src其余部分实现的基础架构进行编译和执行
- python:Python前端,用于包装src中实现的C++函数和对象;C++ API和执行编译的驱动代码,提供Python绑定,与节点对应的算子注册在src/relay/op中
- src/topi: 标准神经网络算子的计算定义和后端调度
向量加法示例:

1. 定义张量
Tensor
python/tvm/te/tensor.py   python Tensor 由C++支持,在include/tvm/te/tensor.h和src/te/tensor.c中实现。TVM中所有Python类型都可以视为具有相同名称的底层C++类型的句柄
Tensor有一个与之相关的Operation对象
python/tvm/te/tensor.py   include/tvm/te/operation.h和src/tvm/te/operation  每个Operation对象都有input_tensors()方法,该方法返回一个输入Tensor列表

2.调度
Schedule
C.op传递给python/tvm/te/schedule.py中的tvm.te.create_schedule()函数  include/tvm/schedule.h
Schedule由Stage和输出Operation的集合组成
Schedule/Stage  python/tvm/te/schedule.py include/tvm/te/schedule.h src/te/schedule/schedule_ops.cc

3.绑定
tvm.build()
python/tvm/driver/build_module.py 接收一个schedule,输入和输出Tensor以及一个target,然后返回一个tvm.runtime.Module
tvm.runtime.Module对象包含一个可以用函数调用语法来调用的已编译函数
tvm.build()的过程可以分为两个步骤
* 降级,高级的、初始的循环嵌套结构被转化为最终的、底层的IR
* 代码生成,由底层的IR来生成目标机器代码
降级:tvm.lower() python/tvm/build_module.py 先进行边界推断(推断所有循环边界和中间缓冲区大小的过程,如果目标式cuda后端,并且使用共享内存,所需的最小尺寸就会在这里自动推断 src/te/schedule/bound.cc src/te/schedule/graph.cc src/te/schedule/message_passing.cc),然后创建一个初始循环嵌套结构(ScheduleOps() src/te/schedule/schedule_ops.cc)
代码生成:build_module() python/tvm/target/codegen.py src/target/codegen/codegen.cc中的Build()返回runtime::Module include/tvm/runtime/module.h src/runtime/module.cc

4.执行
DeviceAPI类
include/tvm/runtime/device_api.h

PackedFunc:用于跨语言函数调用,可以包装任意函数签名,支持不同语言间的互操作性,通过TVMValue和DLDataType统一类型系统
相关文件:
* packed_func.h  用于C++ API
* c_runtime_api.cc 用于C API以及如何提供回调
实现方式
using PackedFunc = std::function<TVMRetValue(TVMArgs args)>;
- TVMArgs: 参数容器,包含参数值和类型信息
- TVMRetValue: 返回值类型,支持常见数据类型转换
注册机制
TVM_REGISTER_GLOBAL("my_packed_func")
.set_body([](TVMArgs args, TVMRetValue* rv){
    // function implementation
});
使用示例
// C++定义
void MyAdd(int a, int b){
    return a + b;
}
TVM_REGISTER_GLOBAL("my_add").set_body_typed(MyAdd);
// Python调用
import tvm
f = tvm.get_global_func("my_add")
result = f(1, 2)
TE和Relay的区别:
TE是低级张量计算抽象,直接定义计算和内存访问模式;用于单个算子的定义和优化,可以精细控制调度,如循环变换、内存分配,需要显式定义计算过程和优化策略
Relay是高级函数式计算图表示,更贴近深度学习框架的抽象层次;适合表示和优化完整的神经网络模型,提供端到端的编译流程,自动处理数据依赖关系和形状推断

使用TE定义和优化矩阵乘法
1.使用TE定义计算 2.创建计算调度 3.编译生成可执行模块 4.分配设备内存并执行计算 5. 执行编译后的函数
te.placeholder   te.create_schedule  tvm.build     tvm.device       func
使用Relay定义和优化矩阵乘法
1.定义Relay计算图 2.构建Relay函数 3.构建模块并设置优化级别 4.创建执行器并执行计算 5.设置输入数据、执行计算、获取结果
relay.var       relay.Function    tvm.IRModule.from_expr tvm.transform.PassContext relay.build tvm.device  tvm.contrib.graph_executor.GraphModule  model.set_input  model.run  model.get_output
向Relay中添加算子
1.添加一个属性节点,声明在编译时已知的固定参数
2.为算子编写一个类型关系,以整合到Relay的类型系统中
3.使用C++中的RELAY_REGISTER_OP宏,为编译器注册算子的数量、类型和其他提示
4.编写算子的计算方式
5.用Relay注册算子和schedule
6.定义一个为算子产生调用节点的C++函数,并为该函数注册一个python API hook
7.将上述Python API hook放在一个更简洁的接口中
8.为新的Relay算子编写测试
向Relay中添加Compiler Pass
编写pass包括两个关键组成部分:创建一个或多个遍历程序的C++类、将遍历实现及其在pass manager API中的元数据包装,从而方便与Pass Infrastructure轻松交互

AST遍历器(Traversers)
用于遍历Relay程序的基类使ExprFunctor.它提供的公共接口是一个VisitExpr方法,该方法接收一个表达式以及零个或多个参数,并返回某种类型的实例. 扩展此类时,可以通过覆盖每种表达式类型的VisitExpr_实现,来定义AST遍历模式

表达式访问器(Expression Visitors)
ExprVisitor不用于修改程序的pass,而是用于实施程序分析和收集信息的pass. 使用这个类,VisitExpr和私有counterparts不会返回任何内容,此类提供的VisitExpr_实现只是访问表达式的所有表达式字段

表达式修改器(Expression Mutators)
ExprMutator用于以某方式转换程序的pass.通过这个类,VisitExpr及其对应的私有部分返回Expr.此类提供的默认VisitExpr_实现访问表达式的所有表达式字段,并将字段设置为访问它们的结果

Pass是编译器中的一个基本概念,指的是对程序的源代码或中间表示(IR)进行单次遍历和转换的处理步骤
在TVM的Relay编译器中,Pass具体指:
1.分析和转换的单元:每个Pass负责一种特定类型的分析或转换,如常量折叠、死代码消除、内联函数等
2.AST处理工具:Pass通过遍历和处理抽象语法树来改变程序的结构或收集信息
3.模块化组件:通过将不同的优化和分析拆分成独立的Pass,编译器可以更模块化地构建,便于扩展和维护
实现常量折叠Pass的简化版
class ConstantFolder : public ExprMutator {
    private:
        // 尝试评估常量表达式
        Optional<Constant> TryEvalConstExpr(const Expr& expr){

        }

    public:
        Expr VisitExpr_(const AddNode* op) override {
            // 首先递归处理操作数
            Expr lhs = this->Mutate(op->lhs);
            Expr rhs = this->Mutate(op->rhs);

            //检查两个操作数是否都是常量
            if (lhs->IsInstance<ConstantNode>() && rhs->IsInstance<ConstantNode>()) {
                //如果都是常量,尝试计算结果
                Optional<Constant> result = TryEvalConstAdd(lhs, rhs);
                if (result) {
                    return result.value(); // 返回计算结果的常量
                }
            }

            // 如果无法折叠,创建新的加法表达式
            if (lhs.same_as(op->lhs) && rhs.same_as(op->rhs)) {
                return GetRef<Expr>(op); // 没有变化,返回原始表达式
            } else {
                return Add(lhs, rhs); // 创建新的加法表达式
            }
        }

        Expr Fold(const Expr& expr){
            return this->Mutate(expr);
        }
};

TVM_REGISTER_GLOBAL("relay._transform.FoldConstant")
.set_body_typed([](Function f, IRModule m){
    ConstantFolder folder;
    return folder.Fold(f);
});


上面涉及的语法知识:
Optional是一个模块类,用于表示可能有值或没有值的情况
Expr是TVM/Relay中表示表达式的类型
Constant是TVM/Relay中表示常量值的类型
override关键字表示此方法重写了基类中的虚函数
向TVM中添加Codegen
创建的codegen 必须位于 src/relay/backend/contrib/<your-codegenname>/
文件中实现的两个类,两个类的关系如下:
            subgraph                        subgraph
TVM backend ---------------------> CSourceCodegen ------------------> CodegenC
    ^                                   |   ^                           |
    |                                   |   |                           |
    -------------------------------------    ----------------------------
      generated C source runtime module          generated C code
当TVM后端发现Relay计算图中的函数,用注册的编译器标签进行了注释时,TVM后端就会调用CSourceCodegen并传递子图。CSourceCodegen的成员函数CreateCSourceModule将:
1.为子图生成C代码
2.将生成的C代码包装到C source runtime模块中,以便TVM后端进行编译和部署

实现 CodegenC
算子的代码生成
首先实现VisitExpr_(const CallNode* call),该函数在遍历子图时会访问所有调用节点,每个调用节点都包含一个我们想要卸载到硬件的算子,需要按照拓扑顺序生成具有正确算子的相应C代码
1. 生成函数声明:函数名、算子的类型、输出张量shape
2. 生成函数调用
3. 生成输出数组
4. 更新输出数组
输入变量的代码生成
Code Emitting: 实现CSourceCodegen、实现GenCFunc、实现CreateCSourceModule

注册CodegenC
为表征实现Codegen

实现ExampleJsonCodeGen
实现自定义runtime
自定义runtime应位于src/runtime/contrib/<your-runtime-name>/, 一个自定义runtime类必须由TVM ModuleNode派生,以保证与其他TVM runtime模块兼容
以下这些从ModuleNode派生的函数,必须在ExampleJsonModule中实现
- 构造函数:这个类的构造函数,应该接收一个表征中的子图,用户可以自行决定处理和存储的格式保存的子图可以被以下两个函数使用
- GetFunction: 当TVM runtime要使用编译器标签执行子图时,它会从自定义runtime模块中调用此函数,它提供函数名及runtime参数,GetFunctiion会返回一个打包的函数实现,以供TVM runtime要使用编译器标签执行子图时
- SaveToBinary和LoadFromBinary: SaveToBinary将runtime模块序列化为二进制格式异构后续部署,用户使用export_library API时,TVM会调用这个函数
- GetSource
用schedule模板和AutoTVM优化算子
auto-tuning包括两个步骤
1. 定义搜索空间
2. 运行搜索算法来探索这个空间
AVX是Intel和AMD处理器上的SIMD指令集扩展
|指令集|AVX2|AVX512|
|"---"|"---"|"---"|
|寄存器宽度|使用256位宽的YMM寄存器|使用512位宽的ZMM寄存器,一次可处理更多数据|
|并行处理能力|同时处理8个32位浮点数或整数|同时处理16个32位浮点数或整数|
|指令集丰富度|比AVX512差|添加了许多新指令,包括更灵活的数据排列、更多的归约操作和更复杂的数学函数|
|处理器支持|在第四代Core及以后的Intel处理器和较新的AMD处理器上支持|在Intel Xeon Phi|
|功耗和频率影响|频率影响较小|导致CPU降频以保持在功耗和温度限制内|

TVM源码学习内容

src/runtime/runtime_base.h: 用于为运行时API函数提供统一的异常处理和错误码返回机制,将C++异常转换为C风格的错误码,确保语言调用的兼容性



import tvm

在最开始的__init__.py文件

  1. 引入_ffi文件里面加载.so文件,初始化相应变量 RPC_SESS_MASK = 128 TVMArrayHandle = ctypes.POINTER(TVMArray) TVMPyCapsuleDestructor = ctypes.CFUNCTYPE(None, ctypes.c_void_p) _c_str_dltensor = c_str(“dltensor”) _c_str_used_dltensor = c_str(“used_dltensor”) _c_dlpack_deleter = TVMPyCapsuleDestructor(_dlpack_deleter) _TVM_COMPATS = {} _TVM_ND_CLS = {} _CLASS_NDARRAY = None ObjectHandle = ctypes.c_void_p init_by_constructor = None OBJECT_TYPE = {} _CLASS_OBJECT = None RETURN_SWITCH[ArgTypeCode.OBJECT_HANDLE] = _return_object C_TO_PY_ARG_SWITCH[ArgTypeCode.OBJECT_HANDLE] = _wrap_arg_func( _return_object, ArgTypeCode.OBJECT_HANDLE ) C_TO_PY_ARG_SWITCH[ArgTypeCode.OBJECT_RVALUE_REF_ARG] = _wrap_arg_func( _return_object, ArgTypeCode.OBJECT_RVALUE_REF_ARG ) PackedFuncHandle = ctypes.c_void_p ModuleHandle = ctypes.c_void_p ObjectHandle = ctypes.c_void_p TVMRetValueHandle = ctypes.c_void_p TVM_FREE_PYOBJ = TVMCFuncFinalizer(_ctypes_free_resource) ctypes.pythonapi.Py_IncRef(ctypes.py_object(TVM_FREE_PYOBJ)) _object.init_by_constructor = init_handle_by_constructor RETURN_SWITCH[ArgTypeCode.PACKED_FUNC_HANDLE] = _handle_return_func RETURN_SWITCH[ArgTypeCode.MODULE_HANDLE] = _return_module RETURN_SWITCH[ArgTypeCode.NDARRAY_HANDLE] = lambda x: _make_array(x.v_handle, False, True) C_TO_PY_ARG_SWITCH[ArgTypeCode.PACKED_FUNC_HANDLE] = _wrap_arg_func( _handle_return_func, ArgTypeCode.PACKED_FUNC_HANDLE ) C_TO_PY_ARG_SWITCH[ArgTypeCode.MODULE_HANDLE] = _wrap_arg_func( _return_module, ArgTypeCode.MODULE_HANDLE ) C_TO_PY_ARG_SWITCH[ArgTypeCode.DLTENSOR_HANDLE] = lambda x: _make_array(x.v_handle, True, False) C_TO_PY_ARG_SWITCH[ArgTypeCode.NDARRAY_HANDLE] = _wrap_arg_func( lambda x: _make_array(x.v_handle, False, True), ArgTypeCode.NDARRAY_HANDLE ) _CLASS_MODULE = None _CLASS_PACKED_FUNC = None _CLASS_OBJECT_GENERIC = None _FUNC_CONVERT_TO_OBJECT = None TVMPackedCFunc = ctypes.CFUNCTYPE( ctypes.c_int, ctypes.POINTER(TVMValue), ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ) TVMCFuncFinalizer = ctypes.CFUNCTYPE(None, ctypes.c_void_p) RETURN_SWITCH = { ArgTypeCode.INT: lambda x: x.v_int64, ArgTypeCode.FLOAT: lambda x: x.v_float64, ArgTypeCode.HANDLE: _return_handle, ArgTypeCode.NULL: lambda x: None, ArgTypeCode.STR: lambda x: py_str(x.v_str), ArgTypeCode.BYTES: _return_bytes, ArgTypeCode.DLDEVICE: _return_device, } C_TO_PY_ARG_SWITCH = { ArgTypeCode.INT: lambda x: x.v_int64, ArgTypeCode.FLOAT: lambda x: x.v_float64, ArgTypeCode.HANDLE: _return_handle, ArgTypeCode.NULL: lambda x: None, ArgTypeCode.STR: lambda x: py_str(x.v_str), ArgTypeCode.BYTES: _return_bytes, ArgTypeCode.DLDEVICE: _return_device, }
  2. 加载runtime文件,将关键变量赋值,并注册runtime、node、runtime.profiling 赋值操作 _CLASS_PACKED_FUNC = PackedFunc _ClASS_OBJECT = Object _CLASS_OBJECT_GENERIC = ObjectGeneric _FUNC_CONVERT_TO_OBJECT = convert_to_object _CLASS_NDARRY = NDArray _CLASS_MODULE = Module

runtime文件夹里面.cc文件和.h文件

TVM_DLL宏: 根据不同平台和编译场景,自动展开成正确的导出或导入标记 TVM_ATTRIBUTE_UNUSED宏: 用于消除编译器对未使用变量或函数的警告 TVM_STR_CONCAT_(_x, __y): 直接拼接,如果__x或__y是其他宏,它们可能不会先展开 TVM_STR_CONCAT(__x, __y): 间接展开后拼接,通过中间层调用TVM_STR_CONCAT,确保__x和__y在拼接前先被展开 __COUNTER__宏: 编译器扩展宏,会自动展开为一个从0开始递增的整数,且每次递增后值不可逆 TVM_ALWAYS_INLINE宏: 编译器特定的强制内联标记 TVM_CHECK_FUNC宏: 生成类型通用的条件检查函数,用于验证两个值是否满足特定关系(如等于、不等于、大于等),并在不满足条件时生成错误信息 TVM_DECLARE_BASE_OBJECT_INFO宏: 实现类型继承系统,分配唯一类型索引 TVM_DECLARE_FINAL_OBJECT_INFO宏: 声明一个不可被继承的“最终类” TVM_DEFINE_OBJECT_REF_METHODS宏: 为自定义的引用类型自动生成构造函数、拷贝/移动语义及对象访问接口 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS宏: 定义不可为空对象引用类 TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(TypeName, ParentType, ObjectName): 定义可变对象引用类型的宏,主要的目的是为某个对象类型(ObjectName)生成一个类型安全的智能指针包装类(TypeName),并提供便捷的访问方法 TVM_DEFINE_OBJECT_REF_COW_METHOD宏: 为对象引用类型添加写时复制功能的宏

c_runtime_api.cc涉及到的头文件 c_runtime_api.h、c_backend_api.h、device_api.h、module.h、packed_func.h、registry.h、object_internal.h、runtime_base.h、ndarray.h、memory.h、object.h、data_type.h、logging.h、serializer.h runtime/container/base.h、array.h、map.h、optional.h、shape_tuple.h、string.h registry.h(涉及全局函数注册相关的信息) object.h、logging.h(涉及一些宏的定义) c_runtime_api.cc里面定义的函数来自于c_runtime_api.h、c_backend_api.h、device_api.h、data_type.h

container.cc文件 注册相关全局函数 cpu_device_api.cc文件 注册device_api.cpu全局函数 dso_library.cc文件 DSOLibrary类继承自Library file_utils.cc文件 实现的函数来自于file_utils.h、meta_data.h 注册runtime.SaveParams、runtime.LoadParams library_module.cc文件 实现的函数来自于library_module.h 注册runtime.module.loadfile_so metadata_module.cc文件 实现的函数来自于meta_data.h 注册runtime.module.loadbinary_metadata module.cc文件 实现的函数来自于module.h 注册runtime.RuntimeEnabled、runtime.ModuleGetSource、runtime.ModuleImportSize、runtime.ModuleGetImport、runtime.ModuleGetTypeKey、runtime.ModuleGetFormat、runtime.ModuleLoadFromFile、runtime.ModuleSaveToFile ndarray.cc文件 实现的函数来自于ndarray.h、c_runtime_api.h 注册runtime.TVMArrayAllocWithScope object.cc文件 实现的函数来自于object.h、c_runtime_api.h 注册runtime.ObjectPtrHash、runtime.DumpTypeTable profiling.cc文件 实现的函数来自于profiling.h 注册runtime.profiling.AsTable、runtime.profiling.AsCSV、runtime.profiling.AsJSON、runtime.profiling.FromJSON、runtime.profiling.DeviceWrapper、runtime.profiling.ProfileFunction registry.cc文件 实现的函数来自于registry.h、c_runtime_api.h、c_backend_api.h source_utils.cc文件 实现的函数来自于source_utils.h system_library.cc文件 实现的函数来自于c_backend_api.h 注册runtime.SystemLib thread_pool.cc文件 实现的函数来自于c_backend_api.h、threading_backend.h 注册runtime.config_threadpool threading_backend.cc文件 实现的函数来自于threading_backend.h workspace_pool.cc文件 实现的函数来自于workspace.h

auto_scheduler注册的全局函数 auto_scheduler auto_scheduler.cc: TuningOptions、AutoSchedule compute_dag.cc: ComputeDAG、ComputeDAGApplyStepsFromState、ComputeDAGPrintPythonCodeFromState、ComputeDAGPrintDAG、ComputeDAGInferBoundFromState、ComputeDAGRewriteLayoutFromState、RewriteIndexForNewLayout、GetShapeFromRewrittenLayout cost_model.cc: RandomModel、PythonBasedModel、CostModelUpdate、CostModelPredict feature.cc: GetPerStoreFeaturesFromFile、GetPerStoreFeaturesFromMeasurePairs、GetPerStoreFeaturesFromStates、GetPerStoreFeatureNames loop_state.cc: StateBind、StateParallel、StateUnroll、StateVectorize、StateFuse、StatePragma、StateReorder、StateSplit、StateFollowSplit、StateFollowFusedSplit、StateStorageAlign、StateComputeAt、StateComputeInline、StateComputeRoot、StateCacheRead、StateCacheWrite、StateRfactor、StateEqual measure_record.cc: RecordToFile、RecordReader、RecordReaderReadLines、RecordReaderReadNext、ReadMeasureRecord、WriteMeasureRecords、SaveRecords、SerializeMeasureInput、DeserializeMeasureInput、SerializeSearchTask、DeserializeSearchTask measure.cc: MeasureInput、BuildResult、MeasureResult、PythonBasedMeasureCallback、ProgramMeasurer、ProgramBuilderBuild、ProgramRunnerRun、LocalBuilder、LocalRunner、RPCRunner、 search_task.cc: HardwareParams、GetDefaultHardwareParams、SearchTask empty_policy.cc: EmptyPolicy search_policy.cc: SearchPolicyRunCallbacks、SearchPolicyContinueSearchOneRound、SearchPolicySetVerbose、PreloadMeasuredStates utils.cc: SearchPolicyUtilsGetConsumers、SearchPolicyUtilsIsElementwiseMatch、SearchPolicyUtilsIsTiled、SearchPolicyUtilsHasCacheReadStage、SearchPolicyUtilsHasCacheWriteStage、SearchPolicyUtilsHasRfactorStage、SearchPolicyUtilsHasCrossThreadReduction sketch_policy.cc: SketchPolicy、SketchPolicyGenerateSketches、SketchPolicySampleInitialPopulation、SketchPolicyEvolutionarySearch、PrintTitle、PreloadCustomSketchRule

auto_scheduler注册的节点类型 auto_scheduler.cc: TuningOptionsNode compute_dag.cc: ComputeDAGNode cost_model.cc: CostModelNode、RandomModelNode、PythonBasedModelNode loop_state.cc: StepNode、StageNode、StateNode、IteratorNode measure_record.cc: RecordToFileNode、RecordReaderNode、 measure.cc: MeasureInputNode、BuildResultNode、MeasureResultNode、MeasureCallbackNode、PythonBasedMeasureCallbackNode、ProgramRunnerNode、ProgramBuilderNode、ProgramMeasurerNode、LocalBuilderNode、LocalRunnerNode、RPCRunnerNode search_task.cc: HardwareParamsNode、SearchTaskNode transform_step.cc: AnnotationStepNode、FuseStepNode、PragmaStepNode、ReorderStepNode、SplitStepNode、FollowSplitStepNode、FollowFusedSplitStepNode、StorageAlignStepNode、ComputeAtStepNode、ComputeInlineStepNode、ComputeRootStepNode、CacheReadStepNode、CacheWriteStepNode、RfactorStepNode empty_policy.cc: EmptyPolicyNode search_policy.cc: SearchCallbackNode、SearchPolicyNode、PreloadMeasuredStatesNode sketch_policy.cc: SketchPolicyNode、PreloadCustomSketchRuleNode

需要类型注册的场景: 当Python类需要与C++类双向交互(如构造实例、调用成员方法)时,必须用@register_object显示注册 无需类型注册的场景: 调用全局C++函数(如_ffi_api.AutoSchedule)时,FFI直接通过函数名调用,参数的类型检查由FFI系统自动完成 上面这个例子可以用@tvm._ffi.register_object(“auto_scheduler.SearchTask”) 和 _ffi_api.AutoSchedule(search_policy, tuning_options) TVM_REGISTER_OBJECT_TYPE注册对象类型

  1. 底层C++类的注册
  2. Python绑定的生成,TVM的构建系统会扫描C++代码中的注册信息,生成Python绑定的胶水代码,例如: 生成一个名为_ffi_api.TuningOptions的函数,对应C++的构造函数,生成__init_handle_by_constructor__方法,用于将Python参数传递给C++构造函数
  3. python对象的初始化 @tvm._ffi.register_func: 注册全局函数(将Python函数暴露给C++代码调用,或在python中调用C++注册函数) python测注册函数供C++调用,C++代码中可通过tvm::runtime::Registry::Get(“custom.add”)调用该函数 调用c++注册函数, tvm._ffi.get_global_func(xx) @tvm._ffi.register_object: 定义Python类与C++类的映射关系,使得python对象可以操作C++对象

tvm 版本迭代过程 0.4版本 autotvm 0.7版本 auto_scheduler 0.8版本 meta_scheduler 0.12版本 减去parser 0.16版本 增加dlight、relax 0.19版本 减去micro

auto_scheduler tvm/auto_scheduler/measure.py class LocalRPCMeasureContext : 在本地环境中创建一个轻量级的RPC测量环境,用于执行和测量任务 代码流程: 1. 检测CUDA环境并设置架构 2. 启动Tracker 3. 启动Server 4. 闯将RPCRunner 5.等待0.5秒确保服务启动 6. 析构时关闭Tracker和Server

Ansor调用全局函数过程

ContinueSearchOneRound过程 SearchOneRound -> PickStatesWithEpsGreedy -> measure->Measure -> program_cost_model->Update(inputs, results)

SearchOneRound过程 GenerateSketches() -> SampleInitPopulation -> GeneratePotentialSpaceStates -> RandomSampleStates -> EvolutionarySearch

SampleInitPopulation函数: 负责从草图中采样初始种群,为后续的进化搜索提供起点

  1. 并行地从草图中随机选择状态并应用初始化规则
  2. 过滤掉无效的和重复的状态
  3. 使用成本模型评估状态质量
  4. 动态调整目标种群大小,避免过长的采样时间
  5. 输出足够数量的有效状态,作为进化搜索的初始种群

EvolutionSearch函数: 一种基于遗传算法的搜索策略,用于在庞大的调度空间中高效寻找优质的调度方案

  1. 从初始种群开始,通过迭代进行搜索
  2. 在每次迭代中,使用成本模型评估当前种群
  3. 维护一个最优状态堆,保存找到的最佳状态
  4. 通过带权重的选择和变异生成下一代种群
  5. 支持常规搜索和潜力空间两种模式
  6. 最终返回按质量排序的最佳状态集合

PickStateWithEpsGreedy函数: 实现 ε-贪心策略, 用于在自动调度器中平衡探索与利用

  1. 根据配置的 ε-贪心参数,计算要选择的最佳状态和随机状态数量
  2. 首先选择足够数量的最佳状态
  3. 然后选择足够数量的随机状态
  4. 确保不会重复测量已经测量过的状态
  5. 将选择的状态转换为测量输入并返回

ctypes语法 结构体和联合 结构体和联合必须派生自Structure和Union基类,这两个基类是在ctypes模块中定义的,每个子类都必须定义_fields_属性。_fields_必须是一个二元组的列表,其中包含一个字段名称和一个字段类型

初始化过程 base.py string_types = (str,) integer_types = (int, np.int32) numeric_types = integer_types + (float, np.float32) py_str = lambda x: x.decode(“utf-8”) version = libinfo.version _LIB, _LIB_NAME = _load_lib() _RUNTIME_ONLY = “runtime” in _LIB_NAME _FFI_MODE = os.environ.get(“TVM_FFI”, “auto”) ERROR_TYPE = {}

runtime_ctypes.py tvm_shape_index_t = ctypes.c_int64 定义64位整数类型,用于表示张量的形状索引 RPC_SESS_MASK = 128 TVMArrayHandle = ctypes.POINTER(TVMArray) 指向TVMArray结构体的指针

object.py ObjectHandle = ctypes.c_void_p init_by_constructor = None OBJECT_TYPE = {} _CLASS_OBJECT = None RETURN_SWITCH[ArgTypeCode.OBJECT_HANDLE] = _return_object C_TO_PY_ARG_SWITCH[ArgTypeCode.OBJECT_HANDLE] = _wrap_arg_func( _return_object, ArgTypeCode.OBJECT_HANDLE ) C_TO_PY_ARG_SWITCH[ArgTypeCode.OBJECT_RVALUE_REF_ARG] = _wrap_arg_func( _return_object, ArgTypeCode.OBJECT_RVALUE_REF_ARG )

ndarray.py TVMPyCapsuleDestructor = ctypes.CFUNCTYPE(None, ctypes.c_void_p) _c_str_dltensor = c_str(“dltensor”) _c_str_used_dltensor = c_str(“used_dltensor”) _c_dlpack_deleter = TVMPyCapsuleDestructor(_dlpack_deleter) _TVM_COMPATS = () _TVM_ND_CLS = {} _CLASS_NDARRAY = None

packed_func.py PackedFuncHandle = ctypes.c_void_p ModuleHandle = ctypes.c_void_p ObjectHandle = ctypes.c_void_p TVMRetValueHandle = ctypes.c_void_p TVM_FREE_PYOBJ = TVMCFuncFinalizer(_ctypes_free_resource) ctypes.pythonapi.Py_IncRef(ctypes.py_object(TVM_FREE_PYOBJ)) _object.init_by_constructor = init_handle_by_constructor RETURN_SWITCH[ArgTypeCode.PACKED_FUNC_HANDLE] = _handle_return_func RETURN_SWITCH[ArgTypeCode.MODULE_HANDLE] = _return_module RETURN_SWITCH[ArgTypeCode.NDARRAY_HANDLE] = lambda x: _make_array(x.v_handle, False, True) C_TO_PY_ARG_SWITCH[ArgTypeCode.PACKED_FUNC_HANDLE] = _wrap_arg_func( _handle_return_func, ArgTypeCode.PACKED_FUNC_HANDLE ) C_TO_PY_ARG_SWITCH[ArgTypeCode.MODULE_HANDLE] = _wrap_arg_func( _return_module, ArgTypeCode.MODULE_HANDLE ) C_TO_PY_ARG_SWITCH[ArgTypeCode.DLTENSOR_HANDLE] = lambda x: _make_array(x.v_handle, True, False) C_TO_PY_ARG_SWITCH[ArgTypeCode.NDARRAY_HANDLE] = _wrap_arg_func( lambda x: _make_array(x.v_handle, False, True), ArgTypeCode.NDARRAY_HANDLE ) _CLASS_MODULE = None _CLASS_PACKED_FUNC = None _CLASS_OBJECT_GENERIC = None _FUNC_CONVERT_TO_OBJECT = None

types.py TVMPackedCFunc = ctypes.CFUNCTYPE( ctypes.c_int, ctypes.POINTER(TVMValue), ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ) TVMCFuncFinalizer = ctypes.CFUNCTYPE(None, ctypes.c_void_p) RETURN_SWITCH = { ArgTypeCode.INT: lambda x: x.v_int64, ArgTypeCode.FLOAT: lambda x: x.v_float64, ArgTypeCode.HANDLE: _return_handle, ArgTypeCode.NULL: lambda x: None, ArgTypeCode.STR: lambda x: py_str(x.v_str), ArgTypeCode.BYTES: _return_bytes, ArgTypeCode.TVM_CONTEXT: _return_context, } C_TO_PY_ARG_SWITCH = { ArgTypeCode.INT: lambda x: x.v_int64, ArgTypeCode.FLOAT: lambda x: x.v_float64, ArgTypeCode.HANDLE: _return_handle, ArgTypeCode.NULL: lambda x: None, ArgTypeCode.STR: lambda x: py_str(x.v_str), ArgTypeCode.BYTES: _return_bytes, ArgTypeCode.TVM_CONTEXT: _return_context, }

runtime/packed_func.py PackedFuncHandle = ctypes.c_void_p _set_class_packed_func(PackedFunc)

runtime/object.py _set_class_object(Object)

runtime/object_generic.py ObjectTypes = (ObjectBase, NDArrayBase, Module, ObjectRValueRef, PyNativeObject) _set_class_object_generic(ObjectGeneric, convert_to_object)

runtime/ndarray.py cl = opencl mtl = metal _set_class_ndarray(NDArray)

runtime/module.py ProfileResult = namedtuple(“ProfileResult”, [“mean”, “results”]) _set_class_module(Module)

_CLASS_PACKED_FUNC = PackedFunc _CLASS_OBJECT = Object _CLASS_OBJECT_GENERIC = ObjectGeneric _FUNC_CONVERT_TO_OBJECT = convert_to_object _CLASS_NDARRAY = NDArray _CLASS_MODULE = Module

_ffi_api.GetPerStoreFeaturesFromMeasurePairsPAM时能够直接调用PackedFuncBased.__call__方法的原因解释如下 python对象属性访问和方法调用机制: 在Python中使用点运算符访问对象的属性和方法时,python会:1.首先在对象的__dict__中查找该属性或方法 2. 如果没找到,则在对象的类及其父类中查找 3. 如果找到了方法,则调用该方法 当Python解释器遇到_ffi_api.GetPerStoreFeaturesFromMeasurePairsPAM(inputs, results, …)时:

  1. _ffi_api是一个模块对象 2. GetPerStoreFeaturesFromMeasurePairsPAM是这个模块的一个属性,其值是一个PackedFuncBase对象 3. 然后Python解释器尝试调用这个对象,即PackedFuncBase对象(inputs, results, …) 4. 这会触发Python的方法调用机制,调用PackedFuncBase.call(self, inputs, results, …)

关键属性 PackedFunc TypePackedFunc TVMArgs TVMPODValue_ TVMArgValue TVMMovableArgValue_ TVMMovableArgValueWithContext_ TVMRetValue TVMArgsSetter

结构体 ObjectTypeChecker、PackedFuncValueConverter

TVMArgValue(values[i], type_codes[i])