build(dnn/cuda): fix cmake compile dependency for cutlass kernels

GitOrigin-RevId: ebe71f5a12

dnn/scripts/cutlass_generator/conv2d_operation.py

@@ -10,7 +10,6 @@ import os.path
 import shutil
 from typing import Tuple, List
 
-from lazy_file import LazyFile
 from library import *
 
 ###################################################################################################
@@ -584,7 +583,7 @@ void initialize_${operation_name}(Manifest &manifest) {
   #
   def __enter__(self):
     self.kernel_path = os.path.join(self.kernel_path, "%s.cu" % self.operation.procedural_name()) 
-    self.kernel_file = LazyFile(self.kernel_path)
+    self.kernel_file = open(self.kernel_path, "w")
     self.kernel_file.write(self.header_template)
     return self
 

dnn/scripts/cutlass_generator/gemm_operation.py

@@ -10,7 +10,6 @@ import shutil
 import functools
 import operator
 
-from lazy_file import LazyFile
 from library import *
 
 
@@ -1045,7 +1044,7 @@ void initialize_${operation_name}(Manifest &manifest) {
   #
   def __enter__(self):
     self.kernel_path = os.path.join(self.kernel_path, "%s.cu" % self.operation.procedural_name()) 
-    self.kernel_file = LazyFile(self.kernel_path)
+    self.kernel_file = open(self.kernel_path, "w")
     self.kernel_file.write(self.header_template)
     return self
 
@@ -1109,7 +1108,7 @@ ${operation_instance}
   #
   def __enter__(self):
     self.kernel_path = os.path.join(self.kernel_path, "%s.cu" % self.operation.procedural_name())
-    self.kernel_file = LazyFile(self.kernel_path)
+    self.kernel_file = open(self.kernel_path, "w")
     self.kernel_file.write(SubstituteTemplate(self.header_template, {
       'wrapper_path': self.wrapper_path,
       }))

dnn/scripts/cutlass_generator/lazy_file.py

@@ -1,27 +0,0 @@
-#
-# \file lazy_file.py
-#
-# \brief LazyFile updates the target file only when the content is changed
-#        in order to avoid generating new cutlass kimpls each time cmake is called
-#
-
-import io
-import os
-
-class LazyFile:
-    def __init__(self, filename):
-        self.filename = filename
-        self.buffer = io.StringIO()
-
-    def write(self, data):
-        self.buffer.write(str(data))
-
-    def close(self):
-        if os.path.isfile(self.filename):
-            old_data = open(self.filename).read()
-        else:
-            old_data = ""
-        new_data = self.buffer.getvalue()
-        if old_data != new_data:
-            with open(self.filename, "w") as f:
-                f.write(new_data)

dnn/scripts/cutlass_generator/manifest.py

@@ -8,7 +8,6 @@ import enum
 import os.path
 import shutil
 
-from lazy_file import LazyFile
 from library import *
 from gemm_operation import *
 from conv2d_operation import *  
@@ -353,7 +352,7 @@ void initialize_all(Manifest &manifest) {
 
 def GenerateManifest(args, operations, output_dir):
   manifest_path = os.path.join(output_dir, "all_%s_%s_operations.cu" % (args.operations, args.type))
-  f = LazyFile(manifest_path)
+  f = open(manifest_path, "w")
   f.write("""
 /*
  Generated by generator.py - Do not edit.

dnn/src/CMakeLists.txt

@@ -116,11 +116,18 @@ if(MGE_WITH_CUDA)
 
     set(CUTLASS_GEN_SCRIPT ${CMAKE_CURRENT_SOURCE_DIR}/../scripts/cutlass_generator/generator.py)
     set(CUTLASS_GEN_DIR ${CMAKE_CURRENT_BINARY_DIR}/cuda/cutlass/generated)
-    function(gen_cutlass_kimpl op type)
+    set(CUTLASS_SOURCES "")
+    function(gen_cutlass_kimpl op type gen_files)
+        set(CURRENT_CUTLASS_STAGE_DIR ${CUTLASS_GEN_DIR}/${op}_${type}.stage)
         set(CURRENT_CUTLASS_GEN_DIR ${CUTLASS_GEN_DIR}/${op}_${type})
+        
+        set_directory_properties(PROPERTIES CMAKE_CONFIGURE_DEPENDS ${CUTLASS_GEN_SCRIPT})
+        
+        file(REMOVE_RECURSE ${CURRENT_CUTLASS_STAGE_DIR})
+        file(MAKE_DIRECTORY ${CURRENT_CUTLASS_STAGE_DIR})
         file(MAKE_DIRECTORY ${CURRENT_CUTLASS_GEN_DIR})
         execute_process(
-            COMMAND ${PYTHON3_EXECUTABLE_WITHOUT_VERSION} ${CUTLASS_GEN_SCRIPT} --operations ${op} --type ${type} ${CURRENT_CUTLASS_GEN_DIR}
+            COMMAND ${PYTHON3_EXECUTABLE_WITHOUT_VERSION} ${CUTLASS_GEN_SCRIPT} --operations ${op} --type ${type} ${CURRENT_CUTLASS_STAGE_DIR}
             RESULT_VARIABLE gen_cutlass_result
             OUTPUT_FILE ${CURRENT_CUTLASS_GEN_DIR}/gen_cutlass.log
             ERROR_FILE ${CURRENT_CUTLASS_GEN_DIR}/gen_cutlass.log
@@ -128,14 +135,27 @@ if(MGE_WITH_CUDA)
         if (NOT gen_cutlass_result EQUAL 0)
             message(FATAL_ERROR "Error generating library instances. See ${CURRENT_CUTLASS_GEN_DIR}/gen_cutlass.log")
         endif()
+        file(GLOB CUTLASS_GEN_FILES RELATIVE "${CURRENT_CUTLASS_GEN_DIR}/" "${CURRENT_CUTLASS_GEN_DIR}/*.cu")
+        foreach(FILE ${CUTLASS_GEN_FILES}) 
+            if (NOT EXISTS "${CURRENT_CUTLASS_STAGE_DIR}/${FILE}")
+                file(REMOVE "${CURRENT_CUTLASS_GEN_DIR}/${FILE}")
+            endif()
+        endforeach()
+        file(GLOB CUTLASS_GEN_FILES RELATIVE "${CURRENT_CUTLASS_STAGE_DIR}" "${CURRENT_CUTLASS_STAGE_DIR}/*.cu")
+        foreach(FILE ${CUTLASS_GEN_FILES})
+            execute_process(COMMAND ${CMAKE_COMMAND} -E copy_if_different "${CURRENT_CUTLASS_STAGE_DIR}/${FILE}" "${CURRENT_CUTLASS_GEN_DIR}")
+        endforeach()
+        file(REMOVE_RECURSE ${CURRENT_CUTLASS_STAGE_DIR})
+        file(GLOB_RECURSE CUTLASS_GEN_FILES "${CURRENT_CUTLASS_GEN_DIR}/*.cu")
+        list(APPEND ${gen_files} ${CUTLASS_GEN_FILES})
+        set(${gen_files} "${${gen_files}}" PARENT_SCOPE)
     endfunction()
-    gen_cutlass_kimpl(gemm simt)
-    gen_cutlass_kimpl(gemv simt)
-    gen_cutlass_kimpl(deconv simt)
-    gen_cutlass_kimpl(conv2d simt)
-    gen_cutlass_kimpl(conv2d tensorop8816)
-    gen_cutlass_kimpl(conv2d tensorop8832)
-    file(GLOB_RECURSE CUTLASS_SOURCES ${CUTLASS_GEN_DIR}/*.cu)
+    gen_cutlass_kimpl(gemm simt CUTLASS_SOURCES)
+    gen_cutlass_kimpl(gemv simt CUTLASS_SOURCES)
+    gen_cutlass_kimpl(deconv simt CUTLASS_SOURCES)
+    gen_cutlass_kimpl(conv2d simt CUTLASS_SOURCES)
+    gen_cutlass_kimpl(conv2d tensorop8816 CUTLASS_SOURCES)
+    gen_cutlass_kimpl(conv2d tensorop8832 CUTLASS_SOURCES)
     list(APPEND SOURCES ${CUTLASS_SOURCES})
     list(APPEND SOURCES ${CUSOURCES})
 endif()