Warnings in `trace_link.py` when running Chakra on AMD GPUs

When using chakra_trace_link on AMD Instinct MI210 GPUs, a bunch of warnings crop up when linking the Kineto and ET JSON files:

[2024-07-11 19:37:52,768] trace_linker.py:721 [WARNING]: Missing parent CPU operator for GPU op 'CopyHostToDevice'. Orphaned GPU operator.
[2024-07-11 19:37:52,768] trace_linker.py:775 [WARNING]: No CUDA runtime operator found for correlation ID -1. This is not a common case, and there should be a corresponding CUDA runtime operator for a given GPU kernel operator. It can be a case where CUDA runtime operators are not properly identified and added to the map, kineto_correlation_cuda_runtime_map. Please manually check if the corresponding CUDA runtime operator with the correlation is dropped by mistake. It is likely that it is because of incomplete map, cuda_launch_operations, in is_cuda_launch_op. Please update the map properly to cover all CUDA runtime launch operators.
[2024-07-11 19:37:52,768] trace_linker.py:721 [WARNING]: Missing parent CPU operator for GPU op 'void at::native::modern::elementwise_kernel<at::native::mse_kernel_cuda(at::TensorIteratorBase&)::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda(float, float)#1}, at::detail::Array<char*, 3> >(int, at::native::mse_kernel_cuda(at::TensorIteratorBase&)::{lambda()#1}::operator()() const::{lambda()#2}::operator()() const::{lambda(float, float)#1}, at::detail::Array<char*, 3>)'. Orphaned GPU operator.
[2024-07-11 19:37:52,768] trace_linker.py:775 [WARNING]: No CUDA runtime operator found for correlation ID -1. This is not a common case, and there should be a corresponding CUDA runtime operator for a given GPU kernel operator. It can be a case where CUDA runtime operators are not properly identified and added to the map, kineto_correlation_cuda_runtime_map. Please manually check if the corresponding CUDA runtime operator with the correlation is dropped by mistake. It is likely that it is because of incomplete map, cuda_launch_operations, in is_cuda_launch_op. Please update the map properly to cover all CUDA runtime launch operators.

Steps to Reproduce

1. Copy the code from toy_model_train.py to your local AMD GPU setup.
2. Run python3 toy_model_train.py
3. Two files KINETO_demo.json and ET_demo.json are generated.
4. Attempt running chakra_trace_link --pytorch-et-file ET_demo.json --kineto-file KINETO_demo.json --output-file LINKED.json
5. You will see the warnings seen above.

Environment Details

• python 3.9.19
• PyTorch version 2.3.0a0+gitae01701
• Chakra installed from source using commit ID fc6bfa8c5adb18ea581498445aeb448d5af1bb60
• Param installed from source using commit ID 884a1f0154a16e2c170e456f8027f2646c9108ae

Possible Causes

1. Chakra assumes that the PyTorch Kineto traces contains the correlation field in the JSON objects. However, on AMD GPUs, the PyTorch Kineto traces do not contain the correlation field - (see this PyTorch issue for more information).
2. In the is_cuda_launch_op function (link), the cuda_launch_operations list does not contain operation names such as hipLaunchKernel.