Skip to content

fix: prelu perf gap on Unet #3717

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.

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

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
13 changes: 10 additions & 3 deletions py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py
Original file line number Diff line number Diff line change
Expand Up @@ -440,7 +440,7 @@ def check_weight_equal(
except Exception:
return torch.all(sd_weight == network_weight)

@needs_refit
@needs_refit # type: ignore[misc]
def _save_weight_mapping(self) -> None:
"""
Construct the weight name mapping from engine weight name to state_dict weight name.
Expand Down Expand Up @@ -577,7 +577,7 @@ def _save_weight_mapping(self) -> None:
gc.collect()
torch.cuda.empty_cache()

@needs_refit
@needs_refit # type: ignore[misc]
def _insert_engine_to_cache(self, hash_val: str, serialized_engine: bytes) -> None:
# TODO: @Evan is waiting for TRT's feature to cache the weight-stripped engine
# if not self.compilation_settings.strip_engine_weights:
Expand Down Expand Up @@ -605,7 +605,7 @@ def _insert_engine_to_cache(self, hash_val: str, serialized_engine: bytes) -> No
),
)

@needs_refit
@needs_refit # type: ignore[misc]
def _pull_cached_engine(self, hash_val: str) -> Optional[TRTInterpreterResult]:
# query the cached TRT engine
cached_data = self.engine_cache.check(hash_val) # type: ignore[union-attr]
Expand Down Expand Up @@ -941,7 +941,14 @@ def output(self, target: str, args: Any, kwargs: Any) -> List[Any]:
f"Specified output dtypes ({len(self.output_dtypes)}) differ from number of outputs ({len(outputs)})"
)

marked_outputs_ids = []
for i, output in enumerate(outputs):
# In some cases, the same output tensor may be marked multiple times, such as _to_copy,
# so we skip marking if the output is already marked
if id(output) in marked_outputs_ids:
Copy link
Collaborator

Choose a reason for hiding this comment

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

Where does this id function come from?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

id() is a built-in function that returns the "identity" of an object.

continue
marked_outputs_ids.append(id(output))

name = f"output{i}"

output_dtype = dtype.unknown
Expand Down
4 changes: 2 additions & 2 deletions py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
Original file line number Diff line number Diff line change
Expand Up @@ -1094,7 +1094,7 @@ def aten_ops_clone_copy_dtype(
name,
args[0],
kwargs.get("dtype", args[0].dtype),
force_layer=True,
force_layer=False, # force_layer=False results in better performance
)


Expand Down Expand Up @@ -1226,7 +1226,7 @@ def aten_ops_sum(
name,
sum_,
kwargs["output_dtype"],
force_layer=True,
force_layer=False, # force_layer=False results in better performance
)
else:
return sum_
Expand Down
Loading