57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
|
#
|
|
# Use of this software is governed by the terms and conditions of the
|
|
# NVIDIA End User License Agreement (EULA), available at:
|
|
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
|
|
#
|
|
# Any use, reproduction, disclosure, or distribution of this software
|
|
# and related documentation outside the scope permitted by the EULA
|
|
# is strictly prohibited.
|
|
|
|
"""
|
|
This module provides a timing helper functions
|
|
"""
|
|
from functools import wraps
|
|
|
|
from .logger import log
|
|
|
|
|
|
# TODO: revisit this part when mlir timing manager is ready for pybind.
|
|
def timer(*dargs, **kwargs):
|
|
enable = kwargs.get("enable", True)
|
|
|
|
def decorator(func):
|
|
@wraps(func)
|
|
def func_wrapper(*args, **kwargs):
|
|
if not enable:
|
|
return func(*args, **kwargs)
|
|
from time import time
|
|
|
|
start = time()
|
|
result = func(*args, **kwargs)
|
|
end = time()
|
|
|
|
# Convert time from seconds to us
|
|
spend_us = (end - start) * 1e6
|
|
|
|
# Determine the function type and format the log message
|
|
if hasattr(func, "__name__"):
|
|
func_name = func.__name__
|
|
log_message = f"[JIT-TIMER] Function: {func_name} | Execution Time: {spend_us:.2f} µs"
|
|
elif "CFunctionType" in str(type(func)):
|
|
log_message = f"[JIT-TIMER] C API Function: {str(func)} | Execution Time: {spend_us:.2f} µs"
|
|
else:
|
|
log_message = f"[JIT-TIMER] Anonymous Function | Execution Time: {spend_us:.2f} µs"
|
|
|
|
log().info(log_message)
|
|
|
|
return result
|
|
|
|
return func_wrapper
|
|
|
|
if len(dargs) == 1 and callable(dargs[0]):
|
|
return decorator(dargs[0])
|
|
else:
|
|
return decorator
|