From 201322d8ea940340b61c43f192c52eb71b6cbbde Mon Sep 17 00:00:00 2001 From: jyz Date: Fri, 26 Jun 2026 01:13:36 +0000 Subject: [PATCH] Fix #1864: get_device_tflops device-aware --- torch/_inductor/analysis/device_info.py | 112 +++++++++++++++++++----- torch/utils/_runtime_estimation.py | 36 ++++++-- 2 files changed, 117 insertions(+), 31 deletions(-) diff --git a/torch/_inductor/analysis/device_info.py b/torch/_inductor/analysis/device_info.py index bcd7cec5e320e..73d0787527720 100644 --- a/torch/_inductor/analysis/device_info.py +++ b/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 diff --git a/torch/utils/_runtime_estimation.py b/torch/utils/_runtime_estimation.py index 17c11ee561456..990a499ce4fe2 100644 --- a/torch/utils/_runtime_estimation.py +++ b/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,7 +166,11 @@ 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) ) @@ -158,6 +178,6 @@ def get_transfer_time(flat_args_kwargs, flat_outs) -> float: # type: ignore[no- 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