Skip to content

feat: support many unary dynamo converters#2246

Merged
narendasan merged 1 commit intopytorch:mainfrom
zewenli98:unary_dynamo_converters
Aug 23, 2023
Merged

feat: support many unary dynamo converters#2246
narendasan merged 1 commit intopytorch:mainfrom
zewenli98:unary_dynamo_converters

Conversation

@zewenli98
Copy link
Copy Markdown
Collaborator

@zewenli98 zewenli98 commented Aug 19, 2023

Description

Implemented all 22 unary dynamo converters, including EXP, LOG, SQRT, RECIP, ABS, SIN, COS, TAN, SINH, COSH, ASIN, ACOS, ATAN, ASINH, ACOSH, ATANH, CEIL, FLOOR, NOT, SIGN, ROUND, ISINF.

Fixes #2199

Type of change

  • New feature (non-breaking change which adds functionality)
  • This change requires a documentation update

Checklist:

  • My code follows the style guidelines of this project (You can use the linters)
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas and hacks
  • I have made corresponding changes to the documentation
  • I have added tests to verify my fix or my feature
  • New and existing unit tests pass locally with my changes
  • I have added the relevant labels to my PR in so that relevant reviewers are notified

@github-actions github-actions bot added component: api [Python] Issues re: Python API component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: tests Issues re: Tests labels Aug 19, 2023
@github-actions github-actions bot requested a review from apbose August 19, 2023 00:30
narendasan
narendasan reviewed Aug 21, 2023
View reviewed changes
Copy link
Copy Markdown
Collaborator

@narendasan narendasan left a comment

Choose a reason for hiding this comment

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

@zewenli98 We need a layer of indirection between aten and impl which is why functions like sign had dedicated implementations.

For each op we need:

  1. aten registration for the actual converter in aten_converters.py
  2. a dedicated impl function

So for aten::cos
We need:

  1. aten_ops_cos (converison.aten_ops_converters)
  2. impl.cos (conversion.impl.unary.ops)

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py


@dynamo_tensorrt_converter(torch.ops.aten.clone.default)
@dynamo_tensorrt_converter(torch.ops.aten.clone.default) # type: ignore[misc]
Copy link
Copy Markdown
Collaborator

@narendasan narendasan Aug 21, 2023

Choose a reason for hiding this comment

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

Before this is committed check to see if theres a way to fix this mypy error @gs-olive can you look at the decorator as well?

Copy link
Copy Markdown
Contributor

@gs-olive gs-olive Aug 22, 2023

Choose a reason for hiding this comment

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

I recently noticed that mypy does not accurately show errors for these decorators. I will look more into this.

The decorator is already as strongly-typed as it can be though, as here:

TensorRT/py/torch_tensorrt/dynamo/conversion/converter_registry.py

Lines 64 to 69 in 585b2b6

def dynamo_tensorrt_converter(
key: Target,
enabled: bool = True,
capability_validator: Optional[Callable[[Node], bool]] = None,
priority: ConverterPriority = ConverterPriority.STANDARD,
) -> Callable[[Any], Union[TRTTensor, Sequence[TRTTensor]]]:

@zewenli98
Copy link
Copy Markdown
Collaborator Author

zewenli98 commented Aug 21, 2023

@zewenli98 We need a layer of indirection between aten and impl which is why functions like sign had dedicated implementations.

For each op we need:

  1. aten registration for the actual converter in aten_converters.py
  2. a dedicated impl function

So for aten::cos We need:

  1. aten_ops_cos (converison.aten_ops_converters)
  2. impl.cos (conversion.impl.unary.ops)

I thought there's almost same code for these converters supported by TensorRT, so I created unary_op in unary/ops.py, like elementwise_op in elementwise/ops.py. I was wondering why dedicated implementations are better?

@narendasan
Copy link
Copy Markdown
Collaborator

narendasan commented Aug 21, 2023 •
edited
Loading

Dedicated implementations are better because they use a standardized simpler interface for all ops. So if I want to make a new converter for say LogSumExp, instead of my code looking like this:

layer = impl.unary_op(in, ..., trt.IUnaryOperation.LOG)
layer = impl.reduce_op(layer.output, ..., trt.IReduceOperation.SUM)
return impl.unary_op(layer.output, ..., trt.IUnaryOperation.EXP)

Peoples code would look closer to this:

x = impl.log(in,...)
x = impl.sum(x,...)
return impl.exp(x,...)

So a converter writer does not need to have in depth knowledge about tensorrt or its APIs, they can write code more similar to pytorch instead.

@zewenli98
Copy link
Copy Markdown
Collaborator Author

zewenli98 commented Aug 21, 2023

Oh I see. It makes sense. Actually at first I implemented dedicated implementations but after seeing elementwise_op, I changed to the current code. I'll change back and also for elementwise.

@narendasan
Copy link
Copy Markdown
Collaborator

narendasan commented Aug 21, 2023 •
edited
Loading

So there can be a core function which is centralized. So unary_op is good, but effectively we just need an alias for specific ops so users dont need to do that themselves

So we would do something like aten_ops_cos -> impl.cos -> impl.unary.unary_op(..., trt.IUnaryOperation.COS)

@zewenli98
Copy link
Copy Markdown
Collaborator Author

zewenli98 commented Aug 21, 2023

OK, I'll change to something like aten_ops_cos -> impl.cos -> impl.unary.unary_op(..., trt.IUnaryOperation.COS)

@zewenli98 zewenli98 requested a review from narendasan August 22, 2023 00:06
narendasan
narendasan approved these changes Aug 22, 2023
View reviewed changes
Copy link
Copy Markdown
Collaborator

@narendasan narendasan left a comment

Choose a reason for hiding this comment

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

LGTM @zewenli98 rebase and should be good to merge

@zewenli98
Copy link
Copy Markdown
Collaborator Author

zewenli98 commented Aug 22, 2023

Do we still need # type: ignore[misc]?

@narendasan
Copy link
Copy Markdown
Collaborator

narendasan commented Aug 22, 2023

Its fine for now we can resolve it later

@zewenli98
feat: support many unary dynamo converters
36d221a 
fix bugs

change to dedicated implementations

move input validation into ops
@zewenli98 zewenli98 force-pushed the unary_dynamo_converters branch from 1a16b09 to 36d221a Compare August 22, 2023 22:25
@zewenli98
Copy link
Copy Markdown
Collaborator Author

zewenli98 commented Aug 22, 2023

rebased, but I don't have access to merge.

@narendasan narendasan merged commit fabfc55 into pytorch:main Aug 23, 2023
@zewenli98 zewenli98 mentioned this pull request Aug 31, 2023
Closed
7 tasks
@gs-olive gs-olive mentioned this pull request Aug 31, 2023
Merged
7 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@narendasan narendasan narendasan approved these changes

@apbose apbose Awaiting requested review from apbose

+1 more reviewer

@gs-olive gs-olive gs-olive left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

cla signed component: api [Python] Issues re: Python API component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: tests Issues re: Tests

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Get IUnaryLayer Ops exposed in dynamo.conversion.impl

4 participants

@zewenli98 @narendasan @gs-olive @facebook-github-bot