arch can be equal to max CC
@@ -150,6 +150,7 @@ def _get_arch():
return min(arch, nvrtc_max_compute_capability)
+@_util.memoize(for_each_device=True)
def _get_arch_for_options_for_nvrtc(arch=None):
# NVRTC in CUDA 11.3+ generates PTX that cannot be run an earlier driver
# version than the one included in the used CUDA version, as
@@ -165,7 +166,7 @@ def _get_arch_for_options_for_nvrtc(arch=None):
version = _cuda_hip_version
if (
not _use_ptx and version >= 11010
- and arch < _get_max_compute_capability()
+ and arch <= _get_max_compute_capability()
):
return f'-arch=sm_{arch}', 'cubin'
return f'-arch=compute_{arch}', 'ptx'
