Skip to content

feat: support aten.copy dynamo converter #2550

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jan 2, 2024
Merged

feat: support aten.copy dynamo converter #2550

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
aeb9ab9
feat: support aten.copy dynamo converter
zewenli98 Dec 18, 2023
05e02c0
fix bug
zewenli98 Dec 28, 2023
2686ee1
fix bug
zewenli98 Dec 29, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2547,3 +2547,28 @@ def aten_ops_trunc(
name,
args[0],
)


@dynamo_tensorrt_converter(torch.ops.aten.copy.default)
@enforce_tensor_types(
{
1: (TRTTensor,),
}
)
def aten_ops_copy(
ctx: ConversionContext,
target: Target,
args: Tuple[Argument, ...],
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
src = args[1]
return impl.cast.to_copy(
ctx,
target,
SourceIR.ATEN,
name,
src,
src.dtype,
force_layer=True,
)
Comment on lines +2566 to +2574
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The force_layer=True does seem useful for the purpose of the test case, but in a normal model we should prefer force_layer=False. This should be a future feature, but to_copy should have an intelligent force_layer mechanism which can select whether to insert the layer based on whether the node is an input node to the TRTEngine.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I agree with you, but for now TensorRT doesn't support the same input and output.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added #2561 to track this proposed feature (in Torch-TRT)

31 changes: 31 additions & 0 deletions tests/py/dynamo/conversion/test_copy_aten.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
import torch
import torch.nn as nn
from parameterized import parameterized
from torch.testing._internal.common_utils import run_tests

from .harness import DispatchTestCase


class TestCopyConverter(DispatchTestCase):
@parameterized.expand(
[
((3,), (3,), False),
((1, 10), (1, 10), False),
((2, 3, 4), (2, 3, 4), True),
((2, 3, 4, 5), (2, 3, 4, 5), True),
]
)
def test_copy_float(self, input_shape, src_shape, non_blocking):
class Copy(nn.Module):
def forward(self, input, src):
return torch.ops.aten.copy.default(input, src, non_blocking)

inputs = [torch.randn(input_shape), torch.randn(src_shape)]
self.run_test(
Copy(),
inputs,
)


if __name__ == "__main__":
run_tests()