SuperAPI AI 工具集

contrib-op-shape-inference-memory-safety

通用与其它

查看项目

更新于 10-16 00:31

技能介绍

用途:审计并修复 ONNX Runtime 算子形状推断函数中的输出越界写入缺陷。典型模式是 TypeAndShapeInference 用 getNumOutputs() > N 作为守卫,却写入大于 N 的输出索引;当节点声明较少输出时,写入位置超出推断上下文输出向量末尾。

范围:onnxruntime/core/graph/contrib_ops/*.cc 与 shape_inference_functions.cc 中的 schema 级形状推断,在 Graph::Resolve(模型加载时)运行一次,与执行提供器(EP)无关,CPU/CUDA/ROCm 无对应内核副本。

主要能力:识别守卫阈值与最高写入索引不匹配的模式;说明为何较少输出数合法(尾部 Optional 输出降低 min_output,checker 只校验 min_output <= N <= max_output);定位 sink——InferenceContextImpl::getOutputType 使用未做边界检查的 operator[],而 getInputType 使用 .at()。

使用方式:用 git grep -n 'getNumOutputs' 扫描所有形状推断守卫,对每个 if (ctx.getNumOutputs() > N) 块找出内部传给 propagateElemTypeFromInputToOutput、updateOutputShape、getOutputType 的最大索引,规则是写入索引 k 则守卫须保证 getNumOutputs() > k。修复时提高守卫阈值,并可在 getOutputType 增加边界检查作为纵深防御。测试需经 Model + Graph::Resolve 路径,负向用例必须非空转,正向用例覆盖全部输出,并兼容 ORT_NO_EXCEPTIONS。

适用场景:审查或加固算子实现及其形状推断,尤其是注意力类 contrib 算子。

限制:资料未提供具体收费、平台或效果信息;部分细节(如各算子完整输出名)需查阅对应 OpSchema。