[build] fix computer capability arch flags, add PTX, handle PTX (micr…

…osoft#591) * fix arch flags, add PTX * bug fix Co-authored-by: Jeff Rasley <[email protected]>
bobisai · Dec 11, 2020 · 8a184b6 · 8a184b6
1 parent 0518252
commit 8a184b6
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/op_builder/builder.py b/op_builder/builder.py
@@ -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:
 
@@ -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
@@ -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