Skip to content

[build] fix computer capability arch flags, add PTX, handle PTX #591

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jeffra merged 4 commits into from
Dec 11, 2020
Merged

[build] fix computer capability arch flags, add PTX, handle PTX #591

Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
03111ce
fix arch flags, add PTX
stas00 Dec 9, 2020
2803061
bug fix
stas00 Dec 9, 2020
63af698
Merge branch 'master' into correct-archs
jeffra Dec 11, 2020
fb22dd1
Merge branch 'master' into correct-archs
jeffra Dec 11, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions op_builder/builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,7 +221,7 @@ def compute_capability_args(self, cross_compile_archs=None):

1. `TORCH_CUDA_ARCH_LIST` takes priority over `cross_compile_archs`.
2. If neither is set default compute capabilities will be used
3. Under `jit_mode` compute capabilities of all visible cards will be used.
3. Under `jit_mode` compute capabilities of all visible cards will be used plus PTX

Format:

Expand All @@ -243,6 +243,7 @@ def compute_capability_args(self, cross_compile_archs=None):
if cc not in ccs:
ccs.append(cc)
ccs = sorted(ccs)
ccs[-1] += '+PTX'
else:
# Cross-compile mode, compile for various architectures
# env override takes priority
Expand All @@ -260,8 +261,10 @@ def compute_capability_args(self, cross_compile_archs=None):

args = []
for cc in ccs:
cc = cc.replace('.', '')
args.append(f'-gencode=arch=compute_{cc},code=compute_{cc}')
num = cc[0] + cc[2]
args.append(f'-gencode=arch=compute_{num},code=sm_{num}')
if cc.endswith('+PTX'):
args.append(f'-gencode=arch=compute_{num},code=compute_{num}')

return args

Expand Down