(*args, **kwargs)
| 106 | def timed_op(func): |
| 107 | |
| 108 | def log_wrapper(*args, **kwargs): |
| 109 | # Add enabled flag so that overhead to each comm op is two if conditions at most |
| 110 | if comms_logger.enabled: |
| 111 | if ('prof' in kwargs |
| 112 | and kwargs['prof']) or comms_logger.prof_all or ('log_name' in kwargs |
| 113 | and kwargs['log_name'] in comms_logger.prof_ops): |
| 114 | # Need func args for their defaults |
| 115 | func_args = get_default_args(func) |
| 116 | func_args.update(kwargs) |
| 117 | msg_size = get_msg_size_from_args(func, *args, **kwargs) |
| 118 | log_name = get_debug_log_name(func_args, comms_logger.debug) |
| 119 | timers(log_name).start() |
| 120 | # Return the op, then stop the op's timer |
| 121 | try: |
| 122 | return func(*args, **kwargs) |
| 123 | finally: |
| 124 | if comms_logger.enabled: |
| 125 | # Need to make op blocking for accurate logging |
| 126 | get_accelerator().synchronize() |
| 127 | # If we're using MPI, we can't simply sync the stream |
| 128 | if cdb.using_mpi: |
| 129 | cdb.barrier() |
| 130 | if ('prof' in kwargs and kwargs['prof']) or comms_logger.prof_all or ( |
| 131 | 'log_name' in kwargs and kwargs['log_name'] in comms_logger.prof_ops): |
| 132 | log_name = get_debug_log_name(func_args, comms_logger.debug) |
| 133 | raw_name = func.__name__ |
| 134 | timers(log_name).stop() |
| 135 | # need temp var since 'elapsed' resets events |
| 136 | time_elapsed = timers(log_name).elapsed(reset=False) |
| 137 | comms_logger.append(raw_name, log_name, time_elapsed, msg_size) |
| 138 | |
| 139 | return log_wrapper |
| 140 |
nothing calls this directly
no test coverage detected