Kernels
KernelMC commited on
Commit
eba5a41
·
verified ·
1 Parent(s): fe4ef5e

Update build/torch-universal/triton_kernels/target_info.py

Browse files

If `num_sms` is called on a system running HIP, it currently returns `None`. But, the expression in the `is_cuda()` branch of this function (`torch.cuda.get_device_properties(0).multi_processor_count`) can also be used on a HIP system. This can be verified by evaluating this expression in a docker container running `rocm/pytorch:rocm7.2_ubuntu24.04_py3.12_pytorch_release_2.9.1` on a system with a supported AMD GPU or APU. Thus, I propose this branch should be taken if `is_cuda() or is_hip()`.

build/torch-universal/triton_kernels/target_info.py CHANGED
@@ -92,7 +92,7 @@ def has_native_mxfp():
92
 
93
 
94
  def num_sms():
95
- if is_cuda():
96
  return torch.cuda.get_device_properties(0).multi_processor_count
97
  if is_xpu():
98
  return torch.xpu.get_device_properties(0).max_compute_units
 
92
 
93
 
94
  def num_sms():
95
+ if is_cuda() or is_hip():
96
  return torch.cuda.get_device_properties(0).multi_processor_count
97
  if is_xpu():
98
  return torch.xpu.get_device_properties(0).max_compute_units