linter

bytedance · Mar 14, 2024 · cb4ea24 · cb4ea24
1 parent c47baf0
commit cb4ea24
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/gpu4pyscf/lib/cupy_helper.py b/gpu4pyscf/lib/cupy_helper.py
@@ -632,7 +632,16 @@ def grouped_dot(As, Bs, Cs=None):
     Ms = np.array(Ms)
     Ns = np.array(Ns)
     Ks = np.array(Ks)
-    total_size = 68 * groups # 68 is the result of sizeof(cutlass::gemm::GemmCoord) + sizeof(typename DeviceKernel::ElementA*) + sizeof(typename DeviceKernel::ElementB*) + sizeof(typename DeviceKernel::ElementC*) + sizeof(typename DeviceKernel::ElementC*) + sizeof(int64_t) + sizeof(int64_t) + sizeof(int64_t)
+    total_size = 68 * groups
+    '''
+    68 is the result of
+    sizeof(cutlass::gemm::GemmCoord) +
+    sizeof(typename DeviceKernel::ElementA*) +
+    sizeof(typename DeviceKernel::ElementB*) +
+    sizeof(typename DeviceKernel::ElementC*) +
+    sizeof(typename DeviceKernel::ElementC*) +
+    sizeof(int64_t) + sizeof(int64_t) + sizeof(int64_t)
+    '''
     padding = 8 - (total_size % 8)
     total_size += padding
     cutlass_space = cupy.empty(total_size, dtype=cupy.uint8)
@@ -715,4 +724,4 @@ def pinv(a, lindep=1e-10):
     v1 = v[:,mask]
     j2c = cupy.dot(v1/w[mask], v1.conj().T)
     return j2c
-  
+
diff --git a/gpu4pyscf/lib/cupy_helper/CMakeLists.txt b/gpu4pyscf/lib/cupy_helper/CMakeLists.txt
@@ -17,8 +17,6 @@
 
 #set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -arch=sm_80")
 
-set(ARCH "70;80")
-
 add_library(cupy_helper SHARED
   transpose.cu
   block_diag.cu
@@ -34,4 +32,4 @@ add_library(cupy_helper SHARED
 target_include_directories(cupy_helper PRIVATE ${cutlass_SOURCE_DIR}/tools/util/include ${cutlass_SOURCE_DIR}/include)
 
 set_target_properties(cupy_helper PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR})
-set_target_properties(cupy_helper PROPERTIES CUDA_ARCHITECTURES "${ARCH}")
+set_target_properties(cupy_helper PROPERTIES CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}")