Skip to content

↔ [Converter] Add support for Granular Op-Support Checks in Dynamo #2058

Closed
@gs-olive

Description

@gs-olive

Add support for operator.get and torch.ops.aten._to_copy and torch.ops.aten.clone

For operator.get we need an evaluator-style implementation which simply dispatches the call out to the Python implementation. For torch.ops.aten._to_copy, the implementation will be a bit more challenging, since we can only support a subset of these operations. Currently, a simple way to perform the differentiation of which _to_copy ops we can handle is to modify the support logic to check if the to converter has a cast which can be performed in TRT. An example of an allowed cast might be torch.float -> torch.int, whereas torch.int ->torch.long is not allowed:

def is_node_supported(
self, submodules: Dict[str, torch.nn.Module], node: torch.fx.Node
) -> bool:
node_name = (
_get_qualified_name(node.target)
if not isinstance(node.target, str)
else node.target
)
if (
node.target in CONVERTERS.keys()
and node_name not in self.torch_executed_ops
):
# If node is a proper, supported computational node, store the operator
if not node.is_impure():
self.supported_operators.add(node_name)
return True
else:
if not node.is_impure():
self.unsupported_operators.add(node_name)
return False

Long Term Solution

A long term solution to this issue of selective conversion is to alter the CONVERSION dictionary to also store an auxiliary function in addition to its converter. This function will take an input node of the converter's type and return whether that node can be converted. For most converters, this function will be: lambda node: True. For _to_copy, however, it might be something like:

def converts_to_copy(node: torch.fx.Node):
    return "dtype" in kwargs and kwargs["dtype"] in APPROVED_CAST_DTYPES

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions