CUTLASS 2.6 (#298)

CUTLASS 2.6

tools/library/scripts/manifest.py

@@ -101,6 +101,70 @@ void initialize_all_${operation_name}_operations(Manifest &manifest) {
     self.top_level_file.write(self.epilogue_template)
     self.top_level_file.close()
 
+class EmitInterfaceLibrary:
+  def __init__(self, generated_path, operation_count, args):
+    self.generated_path = generated_path
+    self.args = args
+
+
+    self.prototypes = []
+    self.fn_calls = []
+    self.operation_count = str(operation_count)
+
+    self.top_level_hdr_template = '''
+/*
+ Generated by manifest.py - Do not edit.
+*/
+'''
+    self.top_level_prologue = '''
+
+#include "cutlass/library/library.h"
+#include "cutlass/library/manifest.h"
+
+namespace cutlass {
+\tnamespace library {
+
+${prototypes}
+
+\t\tvoid initialize_all(Manifest &manifest) {
+\t\t\tmanifest.reserve(${operation_count});\n\n
+${fn_calls}
+\t\t\t}
+
+\t} // namespace library
+} // namespace cutlass
+
+'''
+
+  #
+  def __enter__(self):
+    self.top_level_path = os.path.join(self.generated_path, 'initialize_all.cpp')
+
+    self.top_level_file = open(self.top_level_path, "w")
+    self.top_level_file.write(self.top_level_hdr_template)
+
+    self.source_files = [self.top_level_path,]
+
+    return self
+
+  #
+  def emit(self, operation_name):
+    self.prototypes.append(SubstituteTemplate(
+       "\t\tvoid initialize_all_${operation_kind}_operations(Manifest &manifest);",
+       {'operation_kind': operation_name}))
+    self.fn_calls.append(SubstituteTemplate(
+       "\t\t\tinitialize_all_${operation_kind}_operations(manifest);",
+       {'operation_kind': operation_name}))
+    
+
+
+  #
+  def __exit__(self, exception_type, exception_value, traceback):
+    self.top_level_file.write(SubstituteTemplate(self.top_level_prologue, {'prototypes':"\n".join(self.prototypes),
+                                                                           'fn_calls':"\n".join(self.fn_calls),
+                                                                           'operation_count': self.operation_count}))
+    self.top_level_file.close()
+
 ###################################################################################################
 ###################################################################################################
 
@@ -150,27 +214,6 @@ class Manifest:
 
     self.operation_count = 0
     self.operations_by_name = {}
-    self.top_level_prologue = '''
-
-#include "cutlass/library/library.h"
-#include "cutlass/library/manifest.h"
-
-namespace cutlass {
-namespace library {
-
-${prototypes}
-
-void initialize_all(Manifest &manifest) {
-
-'''
-    self.top_level_reserve = '  manifest.reserve(${operation_count});\n\n'
-    self.top_level_epilogue = '''
-}
-
-} // namespace library
-} // namespace cutlass
-
-'''
 
 
   def get_kernel_filters (self, kernelListFile):
@@ -288,6 +331,9 @@ void initialize_all(Manifest &manifest) {
     operation_emitters = {
       GeneratorTarget.Library: EmitOperationKindLibrary 
     }
+    interface_emitters = {
+      GeneratorTarget.Library: EmitInterfaceLibrary 
+    }
 
     generated_path = os.path.join(self.args.curr_build_dir, 'generated')
 
@@ -299,38 +345,20 @@ void initialize_all(Manifest &manifest) {
 
     source_files = []
 
-    top_level_path = os.path.join(generated_path, 'initialize_all.cpp')
-    with open(top_level_path, 'w') as top_level_file:
-
-      if target == GeneratorTarget.Library:
-        source_files.append(top_level_path)
-
-      prototypes = []
+    with interface_emitters[target](generated_path, self.operation_count, self.args) as iface_emitter:
       for operation_kind, configurations in self.operations.items():
-        prototypes.append(SubstituteTemplate(
-          "void initialize_all_${operation_kind}_operations(Manifest &manifest);",
-          {'operation_kind': OperationKindNames[operation_kind]}))
+          iface_emitter.emit(OperationKindNames[operation_kind])
 
-      top_level_file.write(SubstituteTemplate(self.top_level_prologue,
-        {'prototypes': "\n".join(prototypes)}))
+      source_files += iface_emitter.source_files
 
-      top_level_file.write(SubstituteTemplate(
-        self.top_level_reserve, {'operation_count': str(self.operation_count)}))
 
-      # for each operation kind, emit initializer for all configurations
-      for operation_kind, configurations in self.operations.items():
-        
-        with operation_emitters[target](generated_path, operation_kind, self.args) as operation_kind_emitter:
-          for configuration_name, operations in configurations.items():
-            operation_kind_emitter.emit(configuration_name, operations)
+    # for each operation kind, emit initializer for all configurations
+    for operation_kind, configurations in self.operations.items():
+      with operation_emitters[target](generated_path, operation_kind, self.args) as operation_kind_emitter:
+        for configuration_name, operations in configurations.items():
+          operation_kind_emitter.emit(configuration_name, operations)
 
-          source_files += operation_kind_emitter.source_files
-
-        top_level_file.write(SubstituteTemplate(
-          "  initialize_all_${operation_kind}_operations(manifest);\n",
-          {'operation_kind': OperationKindNames[operation_kind]}))
-
-      top_level_file.write(self.top_level_epilogue)
+        source_files += operation_kind_emitter.source_files
 
     # write the manifest.cmake file containing paths from all targets
     manifest_path = os.path.join(generated_path, "manifest.cmake")