support dtk2304

openfold/utils/kernel/csrc/softmax_cuda_kernel.cu

@@ -32,14 +32,14 @@
 
 __inline__ __device__ float WarpAllReduceMax(float val) {
     for (int mask = 1; mask < 32; mask *= 2) {
-        val = max(val, __shfl_xor_sync(0xffffffff, val, mask));
+        val = max(val, __shfl_xorc(val, mask));
     }
     return val;
 }
 
 __inline__ __device__ float WarpAllReduceSum(float val) {
     for (int mask = 1; mask < 32; mask *= 2) {
-        val += __shfl_xor_sync(0xffffffff, val, mask);
+        val += __shfl_xor( val, mask);
     }
     return val;
 }

setup.py

@@ -32,30 +32,28 @@ version_dependent_macros = [
 extra_cuda_flags = [
     '-std=c++14',
     #'-maxrregcount=50',
-    #'-U__CUDA_NO_HALF_OPERATORS__',
-    #'-U__CUDA_NO_HALF_CONVERSIONS__',
-    '--expt-relaxed-constexpr',
-    '--expt-extended-lambda'
+    '-U__CUDA_NO_HALF_OPERATORS__',
+    '-U__CUDA_NO_HALF_CONVERSIONS__'
 ]
 
-cc_flag = ['-DAMDGPU_TARGETS', 'gfx900;gfx906']
+cc_flag = ['-DAMDGPU_TARGETS', 'gfx900;gfx906;gfx926']
 
 extra_cuda_flags += cc_flag
 
 
-def get_sha(pytorch_root: Union[str, Path]) -> str:
+def get_sha(root: Union[str, Path]) -> str:
     try:
-        return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=pytorch_root).decode('ascii').strip()
+        return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=root).decode('ascii').strip()
     except Exception:
         return 'Unknown'
 
 
 def get_version_add(sha: Optional[str] = None) -> str:
+    openfold_root = os.path.dirname(os.path.abspath(__file__))
     add_version_path = "version.py"
     if sha != 'Unknown':
         if sha is None:
-            sha_path = os.getenv('OPENFOLD_DOWNLOAD_PATH', "")
-            sha = get_sha(sha_path)
+            sha = get_sha(openfold_root)
         version = 'git' + sha[:7]
 
     if os.getenv('OPENFOLD_BUILD_VERSION'):