File size: 1,707 Bytes
ad0aa3a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import os
from contextlib import contextmanager
import torch
try:
import torch_tensorrt
_TRT_VERSION = tuple(int(x) for x in torch_tensorrt.__version__.split(".")[:2])
except (ImportError, AttributeError, ValueError):
_TRT_VERSION = (999, 0) # Assume latest version if not installed
@contextmanager
def nullcontext_decorator(*args, **kwargs):
"""Works with torch_tensorrt >= 2.9.0"""
yield
class NullContextDecoratorCompat:
"""Class-based implementation for torch_tensorrt < 2.9.0 compatibility.
Unlike @contextmanager decorated functions which create _GeneratorContextManager,
this class-based implementation doesn't cause graph breaks in torch.dynamo.
Supports both context manager usage and decorator usage:
- with NullContextDecoratorCompat("name"): ...
- @NullContextDecoratorCompat("name")
"""
def __init__(self, *args, **kwargs):
pass
def __enter__(self):
return self
def __exit__(self, *args):
return False
def __call__(self, func):
"""Support usage as a decorator - just return the function unchanged."""
return func
# Use class-based implementation for torch_tensorrt < 2.9.0 to avoid graph break
_USE_COMPAT_NULLCONTEXT = _TRT_VERSION < (2, 9)
ENABLE_NVTX_DECORATOR = os.environ.get("ENABLE_NVTX_DECORATOR", "0") == "1"
if ENABLE_NVTX_DECORATOR:
NVTXRangeDecorator = torch.cuda.nvtx.range
elif _USE_COMPAT_NULLCONTEXT:
NVTXRangeDecorator = NullContextDecoratorCompat
else:
NVTXRangeDecorator = nullcontext_decorator
|