Skip to content

Commit 1584fcf

Browse files
committed
Merge from main@77709d3dca0bba519358ecf7583d865176d0e891
* src/main: Enable `triton.testing.get_dram_gbps` for XPU (#5455)
2 parents ef9b173 + 77709d3 commit 1584fcf

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
@@ -434,11 +434,10 @@ def perf_report(benchmarks):
434434

435435
def get_dram_gbps(device=None):
436436
''' return DRAM bandwidth in GB/s '''
437-
import torch
438437

439438
from .runtime import driver
440-
if not device:
441-
device = torch.cuda.current_device()
439+
if device is None:
440+
device = driver.active.get_device_interface().current_device()
442441
mem_clock_khz = driver.active.utils.get_device_properties(device)["mem_clock_rate"] # in kHz
443442
bus_width = driver.active.utils.get_device_properties(device)["mem_bus_width"]
444443
bw_gbps = mem_clock_khz * bus_width * 2 / 1e6 / 8 # In GB/s

0 commit comments

Comments
 (0)