[1/2] Script to compare measured (trace) runtimes with estimated runt…

…imes (pytorch#109076)

Summary:
X-link: pytorch/benchmark#1891


Reviewed By: xmfan, xuzhao9, xw285cornell

Differential Revision: D48523883

torch/_inductor/codegen/triton.py

@@ -2543,7 +2543,7 @@ def codegen_node_schedule(self, node_schedule, numel, reduction_numel):
                     node.mark_run()
 
         kernel_name = self.define_kernel(src_code, node_schedule)
-
+        log.debug("Generating kernel code with kernel_name: %s", kernel_name)
         self.codegen_comment(node_schedule)
         kernel.call_kernel(kernel_name)
         V.graph.removed_buffers |= kernel.removed_buffers

torch/_inductor/scheduler.py

@@ -500,6 +500,8 @@ def get_estimated_runtime(self) -> float:
         dtype = None
         if not self.node:
             assert self.snodes
+            if not self.snodes[0].node:
+                return 0
             layout = self.snodes[0].node.get_layout()
             dtype = self.snodes[0].node.get_dtype()
         else:
@@ -530,19 +532,23 @@ def get_estimated_runtime(self) -> float:
                     from .ir import ir_node_to_tensor
 
                     fake_inputs = [
-                        ir_node_to_tensor(input) for input in self.node.inputs
+                        ir_node_to_tensor(input, guard_shape=False)
+                        for input in self.node.inputs
                     ]
                     cls = self.node.__class__
                     cls.process_kernel(op, *fake_inputs, **self.node.kwargs)
 
                     # TODO(xmfan): find a better heuristic to model FLOPS/latency relationship
-                    factor = 0.5
+                    factor = 1.0
                     counted_flops = flop_counter_mode.get_total_flops()
-                    return factor * counted_flops / gpu_flops
+
+                    # Return estimated runtime in nanoseconds
+                    return (factor * counted_flops / gpu_flops) * 1e9
 
         elif isinstance(self, FusedSchedulerNode) or isinstance(
             self.node, ComputedBuffer
         ):
+            # Return estimated runtime in nanoseconds (bytes / gbps)
             return self.get_read_write_buffers_sizes() / gpu_memory_bandwidth
 
         # TODO(xmfan): add support for CollectiveKernel
@@ -944,6 +950,7 @@ def __init__(
         else:
             self.scheduler = scheduler
             self.snodes = nodes
+            self.node = None
 
             self.node = None
             self.users = None
@@ -1822,6 +1829,18 @@ def get_order(n):
     @dynamo_timed
     def codegen(self):
         for node in self.nodes:
+            try:
+                log.debug(
+                    "Generating code for node %s with estimated runtime %f",
+                    node.get_name(),
+                    node.get_estimated_runtime(),
+                )
+            except Exception:
+                log.error(
+                    "Generating code for node %s with estimated runtime 0.0",
+                    node.get_name(),
+                )
+
             self.enter_context(node)
 
             if not isinstance(node, NopKernelSchedulerNode):