Skip to content

Commit

Permalink
Search for 'ptxas' only for cuda backend in 'supports_tma' function
Browse files Browse the repository at this point in the history
Signed-off-by: Anatoly Myachev <[email protected]>
  • Loading branch information
anmyachev authored and peterbell10 committed Dec 4, 2024
1 parent d3a94e0 commit d25ee7f
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion python/triton/_internal_testing.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,11 +142,13 @@ def to_numpy(x):


def supports_tma(byval_only=False):
if not is_cuda():
return False
_, cuda_version = _path_to_binary("ptxas")
min_cuda_version = (12, 0) if byval_only else (12, 3)
cuda_version_tuple = tuple(map(int, cuda_version.split(".")))
assert len(cuda_version_tuple) == 2, cuda_version_tuple
return is_cuda() and torch.cuda.get_device_capability()[0] >= 9 and cuda_version_tuple >= min_cuda_version
return torch.cuda.get_device_capability()[0] >= 9 and cuda_version_tuple >= min_cuda_version


def tma_skip_msg(byval_only=False):
Expand Down

0 comments on commit d25ee7f

Please sign in to comment.