diff --git a/src/trace_link/kineto_operator.py b/src/trace_link/kineto_operator.py index 46a45770..7bbf2d05 100644 --- a/src/trace_link/kineto_operator.py +++ b/src/trace_link/kineto_operator.py @@ -138,10 +138,19 @@ def is_gpu_op(self) -> bool: gpu_categories = {"kernel", "gpu_memcpy"} return self.category in gpu_categories - def is_arrow_op(self) -> bool: + def is_ac2g_op(self) -> bool: """ Check if the operator is categorized as 'ac2g', which stands for arrows from CPU to GPU. + Excerpt from https://pytorch.org/docs/stable/torch.compiler_profiling_torch_compile.html + ``` + Every kernel on the GPU occurs after being launched by code running on the CPU. The profiler can draw + connections (i.e. "flows") between the GPU and CPU events to show which CPU event launched a GPU kernel. + This is particularly helpful because, with a few exceptions, GPU kernels are launched asynchronously. + + To view a flow connection, click on a GPU kernel and click "ac2g". + ```` + Returns bool: True if the operator is an 'ac2g' type, otherwise False. """ diff --git a/src/trace_link/trace_linker.py b/src/trace_link/trace_linker.py index 0e2372bd..d9dc6c1c 100644 --- a/src/trace_link/trace_linker.py +++ b/src/trace_link/trace_linker.py @@ -244,7 +244,7 @@ def construct_kineto_data_structures(self, kineto_ops: List[KinetoOperator]) -> kineto_gpu_ops.append(op) logging.debug(f"Added GPU op: {op.name}") - elif op.is_arrow_op(): + elif op.is_ac2g_op(): # arrow from CPU to GPU assert (op.phase == "s") or (op.phase == "f") if op.id is None: error_msg = (