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 1 commit
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
7 changes: 5 additions & 2 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 @@ -261,7 +262,9 @@ 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}')
args.append(f'-gencode=arch=compute_{cc},code=sm_{cc}')
if cc.endswith('+PTX'):
args.append(f'-gencode=arch=compute_{cc},code=compute_{cc}')

@RezaYazdaniAminabadi RezaYazdaniAminabadi Dec 9, 2020 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't we use the compute-capability (cc) number alone, rather than concatenated with '+PTX'?

@stas00 stas00 Dec 9, 2020 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

edit: I misunderstood what you said, you were pointing at a bug - fixed now. thank you!

But I will leave what I originally wrote below if it helps to explain things in general.

There are 2 unrelated things.

  1. The code you quoted above does the right thing (edit: after I fixed the bug below) if a user passes TORCH_CUDA_ARCH_LIST that contains +PTX, as my previous PR missed that part (I had no clue how +PTX was converted to compile flags at that time, but now I do)

  2. The question is whether to add +PTX for jit_mode or not here:
    https://github.com/microsoft/DeepSpeed/blob/03111cef09b33a25792246785ab0ded68be1733c/op_builder/builder.py#L246
    I added it to sync with pytorch CUDAExtension implementation, which is how it'll do it once pytorch-1.8 is released. But you can choose to not include it. I have no idea how and when jit_mode is used to tell whether it's needed or not.

So to conclude the code you quoted above does the right thing if any of the ccs contains '+PTX, which may come from user-defined TORCH_CUDA_ARCH_LIST`- so it's a must. But wrt to (2) you can choose not to include it.
If you do - please let me know and I will remove https://github.com/microsoft/DeepSpeed/blob/03111cef09b33a25792246785ab0ded68be1733c/op_builder/builder.py#L246

Let me know.

@stas00 stas00 Dec 9, 2020 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh, I see what you meant, there is a bug in my code - sorry - I forgot to check the result of test runs - thank you for flagging that, @RezaYazdaniAminabadi.


return args

Expand Down