Skip to content

Commit 77709d3

Browse files
authored
Enable triton.testing.get_dram_gbps for XPU (#5455)
Signed-off-by: Anatoly Myachev <anatoly.myachev@intel.com>
1 parent abe34a7 commit 77709d3

File tree

1 file changed

+2
-3
lines changed

1 file changed

+2
-3
lines changed

python/triton/testing.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -420,11 +420,10 @@ def perf_report(benchmarks):
420420

421421
def get_dram_gbps(device=None):
422422
''' return DRAM bandwidth in GB/s '''
423-
import torch
424423

425424
from .runtime import driver
426-
if not device:
427-
device = torch.cuda.current_device()
425+
if device is None:
426+
device = driver.active.get_device_interface().current_device()
428427
mem_clock_khz = driver.active.utils.get_device_properties(device)["mem_clock_rate"] # in kHz
429428
bus_width = driver.active.utils.get_device_properties(device)["mem_bus_width"]
430429
bw_gbps = mem_clock_khz * bus_width * 2 / 1e6 / 8 # In GB/s

0 commit comments

Comments
 (0)