Skip to content

feat: support Output Allocator runtime mode for plugin autogen #3442

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 1 commit into from
Mar 17, 2025
Merged

feat: support Output Allocator runtime mode for plugin autogen #3442

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
7 changes: 5 additions & 2 deletions examples/dynamo/auto_generate_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,14 +108,17 @@ def _(x: torch.Tensor, y: torch.Tensor, b: float = 0.2, a: int = 2) -> torch.Ten
# # -------------------------------------------------------------------
# # Given that we have defined the custom operator in PyTorch and TensorRT, we can now generate the converter for the operation.
# # As long as the namespace and names match, the following function will automatically generate the converter for the operation.
# # If plugins require an output allocator to dynamically allocate output buffers, like data dependent operators, please set requires_output_allocator to True.
torch_tensorrt.dynamo.conversion.plugins.generate_plugin_converter(
"torchtrt_ex::elementwise_scale_mul", supports_dynamic_shapes=True
"torchtrt_ex::elementwise_scale_mul",
supports_dynamic_shapes=True,
requires_output_allocator=False,
)


# # %%
# # Above two commands can be replaced with the following single one line:
# torch_tensorrt.dynamo.conversion.plugins.custom_op("torchtrt_ex::elementwise_scale_mul", supports_dynamic_shapes=True)
# torch_tensorrt.dynamo.conversion.plugins.custom_op("torchtrt_ex::elementwise_scale_mul", supports_dynamic_shapes=True, requires_output_allocator=False)


# %%
Expand Down
8 changes: 7 additions & 1 deletion py/torch_tensorrt/dynamo/conversion/plugins/_custom_op.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ def custom_op(
capability_validator: Optional[Callable[[Node, CompilationSettings], bool]] = None,
priority: ConverterPriority = ConverterPriority.STANDARD,
supports_dynamic_shapes: bool = False,
requires_output_allocator: bool = False,
) -> None:
"""
Generate the Plugin and corresponding Plugin Converter using external kernels and TensorRT Quick Deployable Plugin APIs.
Expand All @@ -26,8 +27,13 @@ def custom_op(
partitioner will make sure this Node is run in PyTorch in the compiled graph.
priority: Allows developers to override existing converters in the converter registry
supports_dynamic_shapes: if dynamic shape is supported
requires_output_allocator: if the converter creates operators which require an Output Allocator to run (e.g. data dependent operators)
"""
generate_plugin(op_name)
generate_plugin_converter(
op_name, capability_validator, priority, supports_dynamic_shapes
op_name,
capability_validator,
priority,
supports_dynamic_shapes,
requires_output_allocator,
)
4 changes: 4 additions & 0 deletions py/torch_tensorrt/dynamo/conversion/plugins/_generate_plugin_converter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ def _generate_plugin_converter(
capability_validator: Optional[Callable[[Node, CompilationSettings], bool]] = None,
priority: ConverterPriority = ConverterPriority.STANDARD,
supports_dynamic_shapes: bool = False,
requires_output_allocator: bool = False,
) -> DynamoConverterImplSignature:
torch_target = getattr(getattr(torch.ops, namespace), op_name)
overload_str = overload if overload else ""
Expand Down Expand Up @@ -87,6 +88,7 @@ def custom_kernel_converter(
capability_validator=capability_validator,
priority=priority,
supports_dynamic_shapes=supports_dynamic_shapes,
requires_output_allocator=requires_output_allocator,
)(custom_kernel_converter)
assert (
torch_overload in DYNAMO_CONVERTERS
Expand All @@ -99,6 +101,7 @@ def generate_plugin_converter(
capability_validator: Optional[Callable[[Node, CompilationSettings], bool]] = None,
priority: ConverterPriority = ConverterPriority.STANDARD,
supports_dynamic_shapes: bool = False,
requires_output_allocator: bool = False,
) -> DynamoConverterImplSignature:
plugin_ns, plugin_name = plugin_id.split("::")
return _generate_plugin_converter(
Expand All @@ -107,4 +110,5 @@ def generate_plugin_converter(
capability_validator=capability_validator,
priority=priority,
supports_dynamic_shapes=supports_dynamic_shapes,
requires_output_allocator=requires_output_allocator,
)
Loading