Skip to content

Commit

Permalink
Rename is_arrow_op to is_ac2g_op and add excerpt from PyTorch documen…
Browse files Browse the repository at this point in the history
…tation
  • Loading branch information
TaekyungHeo committed Jul 5, 2024
1 parent 86da209 commit 9dddc64
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 2 deletions.
11 changes: 10 additions & 1 deletion src/trace_link/kineto_operator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""
Expand Down
2 changes: 1 addition & 1 deletion src/trace_link/trace_linker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = (
Expand Down

0 comments on commit 9dddc64

Please sign in to comment.