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.

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

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
Open

fix: prelu perf gap on Unet #3717

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
2f19666
decompose prelu
zewenli98 Jul 24, 2025
79b3153
revert decomp and disable forcing to insert ICast Layer
zewenli98 Jul 28, 2025
cdc31d7
update prelu
zewenli98 Jul 28, 2025
d1d18b9
revert prelu impl
zewenli98 Jul 29, 2025
cf9e7bd
fix identity issue
zewenli98 Aug 8, 2025
11550df
set force_layer=True for some aten ops
zewenli98 Aug 8, 2025
dbb5c78
fix typo
zewenli98 Aug 12, 2025
9745a53
add comments
zewenli98 Aug 12, 2025
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
View file
Open in desktop
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_oppy,
# 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
View file
Open in desktop
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,
Copy link
Collaborator

Choose a reason for hiding this comment

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

I think there might be cases where we might need to actually force_layer=True. Do you know when that would be useful?. Also consider adding a comment here conveying that force_layer=False results in better performance.

Copy link
Collaborator Author

@zewenli98 zewenli98 Aug 12, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

Do you know when we need force_layer=True? My understanding is that 1) Since cast_trt_tensor would be called even if the force_layer=False, it only adds Cast Layer if necessary. 2) If an operator satisfies not is_only_operator_on_placeholder, that means there's other operator on the placeholder or the lowering pass repair_input_as_output to promise "input is not output", so we don't have to explicitly add Cast Layer.

)


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