编写 Dynamo ATen Lowering Passes¶
Lowering Pass 的基础知识¶
ATen Lowering Pass 是 Python 函数,它接受 ATen 算子图作为输入,应用一些所需的修改,例如算子合并/融合、算子替换、子图重写、自定义算子插入,或对 torch.fx.GraphModule 进行其他操作,然后将修改后的图返回给调用者。这些 Lowering Pass 通常会就地修改图并返回相同的输入对象。
Lowering Pass 要求¶
Torch-TRT 中的 ATen Lowering Pass 函数必须满足两个要求: - 函数必须接受 torch.fx.GraphModule 和 torch Tensor 序列 Sequence[torch.Tensor] 作为输入,并返回 Lowering 后的 torch.fx.GraphModule - 函数必须使图处于有效且可调用的状态,包括执行任何必要的 linting 和重新编译
有关 FX 中 图操作 (Graph Manipulations) 的信息,请参阅此链接。下面是一个 Lowering Pass 示例,它修复输入同时也是输出的图,这种情况是 TRT Engines 不允许的配置。
Lowering Pass 示例¶
def repair_input_as_output(gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]) -> torch.fx.GraphModule:
"""Repair scenarios where inputs are also outputs of the graph
TRT does not allow such cases, so we insert a clone (identity) layer
"""
modified_graph = False
# Extract graph placeholder Tensors
placeholders = [
node
for node in gm.graph.nodes
if (
node.op == "placeholder"
and isinstance(node.type, type)
and issubclass(node.type, torch.Tensor)
)
]
for placeholder in placeholders:
# If any placeholder has any users which are direct graph outputs
if len(placeholder.users) >= 1 and any(
user.op == "output" for user in placeholder.users
):
modified_graph = True
# Get direct graph outputs which are direct uses of placeholders
direct_outputs = [user for user in placeholder.users if user.op == "output"]
# Insert clone node for placeholder to ensure
# placeholder is not a direct output
with gm.graph.inserting_after(placeholder):
cloned_placeholder = gm.graph.call_function(
torch.ops.aten.clone.default,
args=(placeholder,),
)
# Replace placeholder as output with cloned version
for output in direct_outputs:
output.replace_input_with(placeholder, cloned_placeholder)
# If the graph was modified, clean up the graph and ensure it is up-to-date
if modified_graph:
gm.graph.eliminate_dead_code()
gm.graph.lint()
gm.recompile()
logger.debug(f"Graph after repair_input_as_output:\n{gm.graph}")
return gm
注册 Lowering Pass¶
Lowering Pass 目前注册在 py/torch_tensorrt/dynamo/lowering/passes/__init__.py 中,使用 torch.fx.passes.pass_manager.PassManager 工具按期望顺序组装 pass 列表。直接添加到该列表的新 pass 将应用于 Torch-TensorRT torch.compile 后端中的图。目前,我们提供一个 ATen Lowering Pass 注册装饰器以方便使用,可以直接调用,也可以使用可选的 index 关键字参数来控制 Lowering Pass 在 pass 列表中的插入位置。
例如,要在默认位置(列表末尾)插入 pass,可以使用以下代码
@_aten_lowering_pass
def my_custom_pass(gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]) -> torch.fx.GraphModule:
...
或者,要在 pass 列表中的自定义索引(例如列表开头)插入 pass,可以使用以下代码
@_aten_lowering_pass(index=0)
def my_custom_pass(gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]) -> torch.fx.GraphModule:
...
在 torch_tensorrt.dynamo.lowering.passes 中还提供了实用工具,用于显示当前可用的 Lowering Pass 列表,将这些 pass 应用于任意 torch.fx.GraphModule,以及删除特定索引处的 Lowering Pass。
# Print all lowering passes in the list
print(dump_lowering_passes())
# Apply lowering passes to a GraphModule
apply_lowering_passes(graph_module, sample_inputs)
# Remove the lowering pass at index 1
_remove_lowering_pass(index=1)
注意: 上述 API 可能会发生变化,因为 Lowering Pass 系统正在发展中。