aten Converter Capabilities Scoping

[gs-olive]

aten Converter Capabilities Scoping

TL;DR

How can we differentiate and select converters on a more granular scale than simply supported/unsupported?

Goal(s)

Python-implemented aten converters are used across the array of frontends which torch_tensorrt provides. Specifying the capabilities of each converter requires more precision than simply categorizing into "present" vs "absent". For instance, certain converters could be typed to only work for dynamic-shape tensors or only static-shape tensors. This RFC details how we might implement this sort of converter differentiation scheme.

Usecases

When selecting the specific converter to translate a node, we can take into account global information about the runtime, including compilation settings and other selections. For example, consider the following node:

%sub_tensor : [#users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%arg0, %arg0), kwargs = {})

In the current implementation, we simply check if torch.ops.aten.sub.Tensor is in the converter registry dictionary. If so, we label the node supported; if not, we label it unsupported. In a future implementation, we could differentiate on more granular criteria, such as whether the input tensors have dynamic shapes. In this scenario, we might have a new converter decorator like:

@tensorrt_converter(torch.ops.aten.sub.Tensor, dynamic=True, static=False)

An even more advanced implementation might have arguments such as:

@tensorrt_converter(torch.ops.aten.sub.Tensor, check_args=arg_check_fn)

The above would have a corresponding checker function, arg_check_fn, which validates the converter can correctly convert its input node. One example of a case where this might be useful is for _to_copy, which is effectively an alias of torch.to. As discussed in #2058, this sort of operator can only be converted if its operands are of a certain type. For instance, we cannot convert tensor.to(torch.long), but we can convert tensor.to(torch.float).

The vast majority of converters would have check_args=lambda node: True, but for a few, this construct would be very helpful in avoiding unnecessary errors and strongly-typing and enforcing conversion constraints. Furthermore, this check_args function can be relaxed as the converter implementation is improved.

Proposed APIs / UX

There would be no major changes needed to the existing user-contribution system for aten converters. The design of this feature is non-breaking, meaning that all of the existing converters will still work, as-is, with no changes to the decorator. A future feature/task can make all necessary changes to update the existing decorators.

Example Workflow

First, a user would implement a converter as follows:

def aten_sub_arg_check_fn(node: torch.fx.Node):
    # Check some aspect of the node inputs
    return check_args(node.args) and check_kwargs(node.kwargs)
 
@tensorrt_converter(torch.ops.aten.sub.Tensor, dynamic=True, static=True, check_args=aten_sub_arg_check_fn)
def aten_sub_converter(...):
    ...

Then, during partitioning (here the dynamo_compile partitioner is shown), we call the check_args_fn and validate the node:

if (
    node.target in CONVERTERS.keys()
    and node_name not in self.torch_executed_ops
    and CONVERTERS[node.target].check_args(node)
):

Internal Implementation

Design

Firstly, the CONVERTERS global registry will need some modifications. Instead of mapping node operations to their corresponding converter functions, it will instead map these node operations to a dataclass struct containing an argument-checking function, a dynamic shape boolean flag, and any other necessary differentiation flags. For instance, see below:

CONVERTERS: Dict[torch.fx.node.Target, Sequence[ConverterSupport]] = {}

@dataclass(frozen=True)
class ConverterSupport:
    check_args: Callable[torch.fx.Node, bool] = field(default=lambda node: True)
    dynamic: bool = True
    static: bool = True
    converter_implementation: Callable

Extensions Required to Core API implementations

The partitioning system in both compile and export paths would need to be tweaked to filter converters accordingly with the above. The new CONVERTERS dictionary stores a target mapped to a sequence of potential converters. The partitioner's role is to determine which of those converter implementations is best-suited to handle the specific node. One sample of a converter priority list we might adhere to is strictest-first:

1. If any inputs are dynamic, first try a converter with dynamic=True, static=False, and vice-versa for static.
2. Verify check_args passes. If not, repeat step 1 with the next-strictest converter. If no more unchecked converters remain, return that the op is unsupported.
3. Select the passing converter's implementation as the one to use.

Data Structures

See above for a discussion of new data structures.

Implementation Phases

Prototype - Medium

• Implement the above dataclass and converter differentiation scheme, minus the dynamic and static flags
• Implement the decorator augmentations necessary to support the new registration scheme without breaking existing converters
• Augment the dynamo_compile partitioner to handle this use case and properly filter the converters, using the new DYNAMO_CONVERTERS registry built here: fix/feat: Add Dynamo-only converter registry #1944

MVP - Small

• Add the above functionality to Dynamo export
• Add support for dynamic and static (dependent on Dynamic shapes RFC Dynamic shape support in dynamo #2014)

Extension Phase 1 - Small

• Add any other necessary converter differentiation arguments to the data class

[apbose]

Thanks for the RFC George. This looks good. Couple of points to add

1. In the use cases, as discussed, we can add the condition in which the converter can specify if input is runtime ITensor or a torch.Tensor. Depending on use cases it can then be an evaluator or a converter. One example would be 1827, where the expand operator for shape tensor errors out when it gets ITensor as input
2. Also why do we need both dynamic and static arguments in the below
   @tensorrt_converter(torch.ops.aten.sub.Tensor, dynamic=True, static=False)
3. It is a good design to map the converter registry to the classes instead of the converter implementation. I guess it would be an implementation thing, but I am assuming in that case the decorator would look something like-

@tensorrt_converter(torch.ops.aten.sub.Tensor, dynamic=True, static=True, check_args=aten_sub_arg_check_fn, converter_impl=aten_sub_converter_impl)
def aten_sub_converter(...)
         assign converterSupport class

[gs-olive]

Thanks for the comments @apbose - to address the questions:

1. This sounds good - I think with this check_args handling, we will be able to have a mix of converters and evaluators depending on the specific use case.
2. I was thinking there could be converters which are dynamic-only, static-only, and static-dynamic. I am not sure whether there are any specific cases of this yet, as the Dynamo dynamic shape path is under development still.
3. I agree with the above, except we will not need converter_impl=aten_sub_converter_impl, since the decorator should automatically handle the caching of the converter implementation which it decorates.

[apbose]

So in the third point, the aten_sub_converter would not contain the class assignment but the implementation call as before? Where would the class assignment take place, in tensorrt_converter ?

[gs-olive]

Yes, my intent was for the mapping of Target to converter implementation would take place in the tensorrt_converter as it is done currently in the prototype, such as here:

TensorRT/py/torch_tensorrt/dynamo/converter_registry.py

Line 86 in 53378e9

DYNAMO_ATEN_CONVERTERS[key] = [converter_support]

The motivation was to ensure that existing converter implementations would all work as-is, without any additions needed.