Update build/torch-universal/triton_kernels/target_info.py
Browse filesIf `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
|