(func, *args, **kwargs)
| 100 | |
| 101 | |
| 102 | def get_msg_size_from_args(func, *args, **kwargs): |
| 103 | # 3 cases: |
| 104 | # - tensor arg is in args |
| 105 | # - tensor arg is in kwargs |
| 106 | # - tensor arg is not present (e.g. barrier) |
| 107 | tensor_arg_position = -1 |
| 108 | tensor_arg = None |
| 109 | # check if tensor arg is in args |
| 110 | if len(args) > 0: |
| 111 | tensor_arg_position = get_tensor_position(func) |
| 112 | if tensor_arg_position > -1: |
| 113 | tensor_arg = args[get_tensor_position(func)] |
| 114 | # check if tensor arg is in kwargs |
| 115 | if tensor_arg is None and len(kwargs) > 0: |
| 116 | tensor_arg = get_tensor_kwarg(func, kwargs) |
| 117 | # if tensor arg is not present, no data is being transmitted |
| 118 | if tensor_arg is None: |
| 119 | return 0 |
| 120 | else: |
| 121 | # Sum of tensor sizes for list colls such as torch's all_to_all |
| 122 | # NOTE: msg_size for list colls will not be the actual size transmitted by a given MPI/NCCL call within the coll op. Instead, it's the total amount of data transmitted. |
| 123 | if type(tensor_arg) is list: |
| 124 | return sum(x.element_size() * x.nelement() for x in tensor_arg) |
| 125 | else: |
| 126 | return tensor_arg.element_size() * tensor_arg.nelement() |
| 127 | |
| 128 | |
| 129 | def get_debug_log_name(func_args, debug): |
no test coverage detected