[mypy] added type annotations to codegen_nodes methods (#119080)

Yang Chen · pytorchmergebot · commit b2e0f8d82d72 · 2024-02-05T18:33:52.000Z
added correct type annotations to scheduler and backends' codegen_nodes methods Pull Request resolved: #119080 Approved by: https://github.com/eellison
diff --git a/torch/_inductor/codegen/cpp.py b/torch/_inductor/codegen/cpp.py
@@ -3089,18 +3089,18 @@ def is_memory_copy_scheduler_node(node: SchedulerNode):
                 body: ir.LoopBody = node._body
                 _legalize_lowp_fp(body)
 
-    def codegen_nodes(self, nodes):
+    def codegen_nodes(self, nodes: List[SchedulerNode]):
         # Legalize BF16 node by adding to_dtype explicitly
         self.legalize_lowp_fp_dtype(nodes)
         self.data_type_propagation(nodes)
 
         assert len(nodes) >= 1
         first_node = nodes[0]
         vec_dtype = (
-            first_node._lowp_fp_type
+            first_node._lowp_fp_type  # type: ignore[attr-defined]
             if all(
                 hasattr(_node, "_lowp_fp_type")
-                and _node._lowp_fp_type == first_node._lowp_fp_type
+                and _node._lowp_fp_type == first_node._lowp_fp_type  # type: ignore[attr-defined]
                 for _node in nodes
             )
             else torch.float
@@ -3318,7 +3318,7 @@ def can_fuse_horizontal(self, node1, node2):
     def can_fuse_vertical(self, node1, node2):
         return self._can_fuse_horizontal_impl(node1, node2) and not node1.is_reduction()
 
-    def codegen_nodes(self, nodes):
+    def codegen_nodes(self, nodes: List[SchedulerNode]):
         """
         Turn an set of pre-fused nodes into a C++ kernel.
         """
diff --git a/torch/_inductor/codegen/cuda_combined_scheduling.py b/torch/_inductor/codegen/cuda_combined_scheduling.py
@@ -59,7 +59,7 @@ def codegen_template(
                 template_node, epilogue_nodes
             )
 
-    def codegen_nodes(self, nodes: List[BaseSchedulerNode]):
+    def codegen_nodes(self, nodes: List[SchedulerNode]):
         return self._triton_scheduling.codegen_nodes(nodes)
 
     def codegen_sync(self):
diff --git a/torch/_inductor/codegen/triton.py b/torch/_inductor/codegen/triton.py
@@ -3142,7 +3142,7 @@ def requires_closing_previous_reduction(node, node_schedule):
 
         return node_schedule
 
-    def codegen_nodes(self, nodes):
+    def codegen_nodes(self, nodes: List[scheduler.SchedulerNode]):
         """
         Given a set of pre-fused nodes, generate a Triton kernel.
         """
diff --git a/torch/_inductor/scheduler.py b/torch/_inductor/scheduler.py
@@ -2343,7 +2343,7 @@ def codegen_template(
         """
         raise NotImplementedError()
 
-    def codegen_nodes(self, nodes: List[BaseSchedulerNode]):
+    def codegen_nodes(self, nodes: List[SchedulerNode]):
         """
         Generate a kernel given a list of pre-fused nodes.
         """

Original file line number	Diff line number	Diff line change
`@@ -59,7 +59,7 @@ def codegen_template(`
`59`	`59`	`template_node, epilogue_nodes`
`60`	`60`	`)`
`61`	`61`
`62`		`- def codegen_nodes(self, nodes: List[BaseSchedulerNode]):`
	`62`	`+ def codegen_nodes(self, nodes: List[SchedulerNode]):`
`63`	`63`	`return self._triton_scheduling.codegen_nodes(nodes)`
`64`	`64`
`65`	`65`	`def codegen_sync(self):`