[inductor][easy] add buffer layout to SchedulerNode.debug_str (pytorc…

…h#125090)

Pull Request resolved: pytorch#125090
Approved by: https://github.com/jansel

test/inductor/test_debug_trace.py

@@ -59,6 +59,8 @@ def fn(a, b):
 buf0.group.device = cpu
 buf0.group.iteration = ((256,), ())
 buf0.sizes = ([256], [])
+arg0_1_layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
+buf0_layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
 class buf0_loop_body:
     var_ranges = {z0: 256}
     index0 = z0
@@ -80,6 +82,8 @@ def body(self, ops):
 buf1.group.device = cpu
 buf1.group.iteration = ((256,), ())
 buf1.sizes = ([256], [])
+buf0_layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
+buf1_layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
 class buf1_loop_body:
     var_ranges = {z0: 256}
     index0 = z0
@@ -117,6 +121,8 @@ def body(self, ops):
     buf0.group.device = cpu
     buf0.group.iteration = ((256,), ())
     buf0.sizes = ([256], [])
+    arg0_1_layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
+    buf0_layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
     class buf0_loop_body:
         var_ranges = {z0: 256}
         index0 = z0
@@ -137,6 +143,8 @@ def body(self, ops):
     buf1.group.device = cpu
     buf1.group.iteration = ((256,), ())
     buf1.sizes = ([256], [])
+    buf0_layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
+    buf1_layout = FixedLayout('cpu', torch.float32, size=[16, 16], stride=[16, 1])
     class buf1_loop_body:
         var_ranges = {z0: 256}
         index0 = z0

torch/_inductor/scheduler.py

@@ -732,6 +732,10 @@ def debug_str_extra(self) -> str:
             f"{name}.group.iteration = {self.group[1]}",
             f"{name}.sizes = {self._sizes}",
         ]
+        for dep in self.read_writes.reads_and_writes():
+            buf_name = dep.name
+            buf = V.graph.get_buffer(buf_name)
+            lines.append(f"{buf_name}_layout = {pformat(buf.layout)}")
         if self.get_aliases():
             lines.append(f"{name}.aliases = {pformat(self.get_aliases())}")
         if self.get_mutations():