Fix #1864: get_device_tflops device-aware

torch/_inductor/analysis/device_info.py

@@ -30,7 +30,8 @@ class DeviceInfo:
     tops_sparsity_factor: int = 1
 
 
-# Indexing is based on `torch.cuda.get_device_name()`, normalized to upper-case.
+# Indexing is based on device names returned by torch.cuda.get_device_name()
+# or torch.xpu.get_device_name(), normalized to upper-case.
 # TODO investigate profiler support for tf32 and allow device to report correct number when it's turned on.
 _device_mapping: dict[str, DeviceInfo] = {
     # Source:
@@ -156,6 +157,37 @@ class DeviceInfo:
         dram_bw_gbs=5300.0,
         dram_gb=192.0,
     ),
+    # Intel Data Center GPU Max 1550 / PVC
+    # Sources:
+    # - Intel ARK lists Max 1550 with 128 Xe-cores, 1600 MHz max dynamic
+    #   clock, 128 GB HBM2e, 3276.8 GB/s memory bandwidth.
+    # - Intel Xe-HPC architecture guide lists Data Center GPU Max 1550
+    #   device-level throughput: FP32 MAD: 16384 * 2 FLOPs/clk.
+    "INTEL DATA CENTER GPU MAX 1550": DeviceInfo(
+        tops={
+            torch.float64: 52.43,
+            torch.float32: 52.43,
+            "torch.tf32": 52.43,
+            torch.bfloat16: 838.86,
+            torch.float16: 838.86,
+            torch.int8: 1677.72,
+        },
+        dram_bw_gbs=3276.8,
+        dram_gb=128.0,
+    ),
+    # Intel Data Center GPU Max 1100 / PVC (same generation as Max 1550)
+    "INTEL DATA CENTER GPU MAX 1100": DeviceInfo(
+        tops={
+            torch.float64: 52.43,
+            torch.float32: 52.43,
+            "torch.tf32": 52.43,
+            torch.bfloat16: 838.86,
+            torch.float16: 838.86,
+            torch.int8: 1677.72,
+        },
+        dram_bw_gbs=2457.6,
+        dram_gb=48.0,
+    ),
     # Source:
     # @lint-ignore https://www.amd.com/content/dam/amd/\
     # en/documents/instinct-business-docs/product-briefs/instinct-mi210-brochure.pdf
@@ -185,6 +217,13 @@ class DeviceInfo:
 _device_mapping["AMD INSTINCT MI300X"] = _device_mapping["AMD MI300X"]
 _device_mapping["AMD INSTINCT MI210X"] = _device_mapping["AMD MI210X"]
 
+_device_mapping["Intel(R) Data Center GPU Max 1550"] = _device_mapping[
+    "INTEL DATA CENTER GPU MAX 1550"
+]
+_device_mapping["Intel(R) Data Center GPU Max 1100"] = _device_mapping[
+    "INTEL DATA CENTER GPU MAX 1100"
+]
+
 # Enforce the upper-case-key invariant so entries cannot silently miss
 # `lookup_device_info` (which upper-cases the query before lookup).
 _device_mapping = {k.upper(): v for k, v in _device_mapping.items()}
@@ -202,7 +241,45 @@ def lookup_device_info(name: str) -> DeviceInfo | None:
     return _device_mapping.get(name.upper())
 
 
-def datasheet_tops(dtype: torch.dtype, is_tf32: bool = False) -> float | None:
+def _get_device_name(device: torch.device | None = None) -> str | None:
+    if device is None:
+        if torch.cuda.is_available():
+            device = torch.device("cuda")
+        elif hasattr(torch, "xpu") and torch.xpu.is_available():
+            device = torch.device("xpu")
+        else:
+            return None
+
+    if device.type == "cuda" and torch.cuda.is_available():
+        return torch.cuda.get_device_name(device)
+
+    if device.type == "xpu" and hasattr(torch, "xpu") and torch.xpu.is_available():
+        get_name = getattr(torch.xpu, "get_device_name", None)
+        if get_name is not None:
+            if device.index is None:
+                return get_name()
+            return get_name(device.index)
+
+    return None
+
+
+def datasheet_dram_bw(device: torch.device | None = None) -> float | None:
+    name = _get_device_name(device)
+    if name is None:
+        log.info("No datasheet device name found for %s, returning None", device)
+        return None
+
+    device_info = lookup_device_info(name)
+    if device_info is None:
+        log.info("Device %s not in datasheet, returning None", name)
+        return None
+
+    return device_info.dram_bw_gbs
+
+
+def datasheet_tops(
+    dtype: torch.dtype, is_tf32: bool = False, device: torch.device | None = None
+) -> float | None:
     """
     Get the theoretical *dense* TFLOPS of the device for a given dtype.
 
@@ -211,36 +288,25 @@ def datasheet_tops(dtype: torch.dtype, is_tf32: bool = False) -> float | None:
     callers always receive the throughput achievable by cuBLAS/cuDNN on
     non-sparse data.
     """
-    if torch.cuda.is_available():
-        name: str | None = torch.cuda.get_device_name()
-    elif torch.xpu.is_available():
-        name: str | None = torch.xpu.get_device_name()
-        log.info(
-            "No XPU devices are currently supported in the datasheet lookup. "
-            "To get accurate compute estimates, add your device to "
-            "torch/_inductor/analysis/device_info.py",
-        )
-    else:
-        log.info("No supported device available, skipping datasheet lookup")
-        return None
-
+    name = _get_device_name(device)
     if name is None:
-        log.info("No device found, returning None")
+        log.info("No datasheet device name found for %s, returning None", device)
         return None
+
     device_info = lookup_device_info(name)
     if device_info is None:
-        log_str = f"Device {name} not in datasheet, returning None"
-        log.info(log_str)
+        log.info("Device %s not in datasheet, returning None", name)
         return None
-    if dtype not in device_info.tops:
+
+    dtype_key = "torch.tf32" if dtype == torch.float32 and is_tf32 else dtype
+
+    if dtype_key not in device_info.tops:
         log.info(
             "Device %s does not have a datasheet entry for %s, returning None",
             name,
-            dtype,
+            dtype_key,
         )
         return None
 
-    tops = device_info.tops[
-        "torch.tf32" if dtype == torch.float32 and is_tf32 else dtype
-    ]
+    tops = device_info.tops[dtype_key]
     return tops / device_info.tops_sparsity_factor

torch/utils/_runtime_estimation.py

@@ -1,4 +1,5 @@
 import torch
+import torch.utils._pytree as pytree
 from torch._inductor.utils import get_device_tflops, get_gpu_dram_gbps
 from torch.fx.experimental.symbolic_shapes import (
     optimization_hint,
@@ -74,14 +75,27 @@
 _IGNORE_OPS = _VIEW_OPS | _CREATE_OPS
 
 
-def flops_to_ns(flops: float | int, dtype: "torch.dtype") -> float:
+def _get_device_from_value(value) -> torch.device | None:  # type: ignore[no-untyped-def]
+    flat_values, _ = pytree.tree_flatten(value)
+    for item in flat_values:
+        if isinstance(item, torch.Tensor):
+            return item.device
+    return None
+
+
+def flops_to_ns(
+    flops: float | int, dtype: "torch.dtype", device: torch.device | None = None
+) -> float:
     """Convert a FLOPs count to estimated nanoseconds on the current GPU.
 
     Uses 75% of theoretical peak and converts FLOPs to MACs (divide by 2).
     """
-    peak_gpu_flops = get_device_tflops(dtype) * 1e12
-    if peak_gpu_flops == 0:
+    device_tflops = get_device_tflops(dtype, device=device)
+
+    if device_tflops == 0:
         return 0.0
+
+    peak_gpu_flops = device_tflops * 1e12
     macs = flops / 2
     return (macs / (0.75 * peak_gpu_flops)) * 1e9
 
@@ -116,7 +130,9 @@ def get_compute_time(
         if node_meta is not None:
             extra_kwargs["_node_meta"] = node_meta
         flop_count = flop_count_func(*args, **kwargs, out_val=out, **extra_kwargs)
-        return flops_to_ns(flop_count, dtype)
+
+        device = _get_device_from_value(out)
+        return flops_to_ns(flop_count, dtype, device=device)
     return 0.0
 
 
@@ -150,14 +166,18 @@ def get_transfer_time(flat_args_kwargs, flat_outs) -> float:  # type: ignore[no-
     Returns:
         float: The estimated memory transfer time in nanoseconds.
     """
-    gpu_memory_bandwidth = get_gpu_dram_gbps()
+    device = _get_device_from_value((flat_args_kwargs, flat_outs))
+    gpu_memory_bandwidth = get_gpu_dram_gbps(device=device)
+    if gpu_memory_bandwidth <= 0:
+        return 0.0
+
     read_bytes = sum(
         get_num_bytes(t) for t in flat_args_kwargs if isinstance(t, torch.Tensor)
     )
     write_bytes = sum(
         get_num_bytes(t) for t in flat_outs if isinstance(t, torch.Tensor)
     )
     counted_bytes = read_bytes + write_bytes
-    # The GPU memory bandwidth is in GB/s so the transfer time is in nanoseconds
-    transfer_time = counted_bytes / gpu_memory_bandwidth
-    return transfer_time
+
+    # gpu_memory_bandwidth is GB/s. bytes / GB/s gives nanoseconds.
+    return counted_bytes / gpu_memory_bandwidth