| |
| |
|
|
| 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) |
|
|
|
|
| @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_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 |
|
|
