From a759e85f5f93752433a019cbd17d7beadfb0e218 Mon Sep 17 00:00:00 2001
From: Christian Sigg <chsigg@users.noreply.github.com>
Date: Thu, 30 Nov 2023 06:25:40 +0100
Subject: [PATCH] Add subclass declarations to generated files. (#1193)

---
 python/cutlass_library/manifest.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/python/cutlass_library/manifest.py b/python/cutlass_library/manifest.py
index 3e3b477f..20a2181a 100644
--- a/python/cutlass_library/manifest.py
+++ b/python/cutlass_library/manifest.py
@@ -181,6 +181,7 @@ void initialize_all_sm${min_cc}_${subclass_name}_${operation_name}_operations(Ma
 """
     self.configuration_prototype_template = "void initialize_${configuration_name}(Manifest &manifest);\n"
     self.configuration_template = "  initialize_${configuration_name}(manifest);\n"
+    self.subclass_prototype_template = "void initialize_all_sm${min_cc}_${subclass_name}_${operation_name}_operations(Manifest &manifest);\n"
     self.subclass_call_template = "  initialize_all_sm${min_cc}_${subclass_name}_${operation_name}_operations(manifest);\n"
 
     self.epilogue_template ="""}
@@ -251,6 +252,14 @@ void initialize_all_sm${min_cc}_${subclass_name}_${operation_name}_operations(Ma
   #
   def __exit__(self, exception_type, exception_value, traceback):
 
+    for subclass_name, subclass_file in sorted(self.subclass_files.items()):
+      subclass_cfg = {
+        'min_cc': str(self.min_cc),
+        'subclass_name': subclass_name,
+        'operation_name': OperationKindNames[self.kind]
+      }
+      self.top_level_file.write(SubstituteTemplate(self.subclass_prototype_template, subclass_cfg))
+
     self.top_level_file.write(
       SubstituteTemplate(self.entry_template, {
         'min_cc': str(self.min_cc),