add MTIA as valid device type for prof averages (#130340)
Summary: Add MTIA as valid device option for getting profile averages
Test Plan: Tested with auto-trace on MTIA
Differential Revision: D59486392
Pull Request resolved: https://github.com/pytorch/pytorch/pull/130340
Approved by: https://github.com/aaronenyeshi
diff --git a/torch/autograd/profiler_util.py b/torch/autograd/profiler_util.py
index a5cff1e..dc2c909 100644
--- a/torch/autograd/profiler_util.py
+++ b/torch/autograd/profiler_util.py
@@ -580,7 +580,11 @@
# each legacy cpu events has a single (fake) kernel
return sum(kinfo.duration for kinfo in self.kernels)
else:
- assert self.device_type in [DeviceType.CUDA, DeviceType.PrivateUse1]
+ assert self.device_type in [
+ DeviceType.CUDA,
+ DeviceType.PrivateUse1,
+ DeviceType.MTIA,
+ ]
return self.time_range.elapsed_us()
@property
@@ -600,7 +604,11 @@
[child.device_time_total for child in self.cpu_children]
)
else:
- assert self.device_type in [DeviceType.CUDA, DeviceType.PrivateUse1]
+ assert self.device_type in [
+ DeviceType.CUDA,
+ DeviceType.PrivateUse1,
+ DeviceType.MTIA,
+ ]
return self.device_time_total
@property
@@ -961,7 +969,11 @@
if evt.device_type == DeviceType.CPU and evt.is_legacy:
# in legacy profiler, kernel info is stored in cpu events
sum_self_device_time_total += evt.self_device_time_total
- elif evt.device_type in [DeviceType.CUDA, DeviceType.PrivateUse1]:
+ elif evt.device_type in [
+ DeviceType.CUDA,
+ DeviceType.PrivateUse1,
+ DeviceType.MTIA,
+ ]:
# in kineto profiler, there're events with the correct device type (e.g. CUDA)
sum_self_device_time_total += evt.self_device_time_total
