llvm/utils/imisched.py: Accept input from stdin

This allows something like:
`llc test.ll --debug-only=machine-scheduler 2>&1 | imisched.py -`

Author: gbossu

llvm/test/CodeGen/AIE/aie2/end-to-end/Conv2D-red.ll

@@ -6,6 +6,9 @@
 ;
 ; (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
 ; RUN: llc -O2 -mtriple=aie2 --enable-pipeliner=0 %s -o - | FileCheck %s --check-prefix=ASM
+; RUN: llc -O2 -mtriple=aie2 --enable-pipeliner=0 %s -o - --debug-only=machine-scheduler  \
+; RUN:    2>&1 | %imisched -d - \
+; RUN:    | FileCheck %s --check-prefix=SCHED-DUMP
 
 
 ; This is a reduced version of the Conv2D_0 MLLib benchmark which only contains
@@ -23,6 +26,13 @@
 ; The test is meant as a quick way to spot QoR regressions, but if the ASM is
 ; too unstable, we can use different FileCheck lines.
 
+
+; SCHED-DUMP: Region: conv2d.loop.nest:bb.2
+; SCHED-DUMP: Region: conv2d.loop.nest:bb.4
+; SCHED-DUMP: Region: conv2d.loop.nest:bb.3
+; SCHED-DUMP: Region: conv2d.loop.nest:bb.1
+; SCHED-DUMP: Region: conv2d.loop.nest:bb.0
+
 define dso_local void @conv2d.loop.nest(ptr %add.ptr6.i51, ptr %add.ptr5, ptr %cond, ptr %cond.i50, <16 x i32> %0, i32 %cond67.i79, i20 %idx.ext.i.i81, i20 %idx.ext.i404.i, i20 %idx.ext.i410.i, i20 %idx.ext.i434.i85, i32 %1, i20 %2, i20 %3, i20 %4, i20 %5, i20 %6, i32 %7, i32 %8, i32 %or9.i.i.i.i.i96, i32 %9, i20 %idx.ext.i422.i82, i20 %10, i20 %11, i20 %12, i20 %13, i20 %14, i20 %15, i20 %16, i20 %17, i20 %18, i20 %19, i20 %20, i20 %21, i20 %22, i20 %23, i32 %conv192.i107, i20 %24, i20 %idx.ext.i428.i, i20 %25, i20 %26, i20 %27, i32 %28) #1 {
 ; ASM-LABEL: conv2d.loop.nest:
 ; ASM:         .p2align 4

llvm/test/CodeGen/AIE/aie2/lit.local.cfg

@@ -5,6 +5,11 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 #
 # (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
+import os
+
 config.substitutions.insert(0, ('%last-mi-pass', 'machinemoduleinfo'))
 config.substitutions.insert(0, ('%topdown-single', '--issue-limit=1 --aie-bottomup-cycles=0'))
 config.substitutions.insert(0, ('%topdown-multi', '--issue-limit=6 --aie-bottomup-cycles=0'))
+
+imisched_path = os.path.join(config.llvm_src_root, 'utils', 'imisched.py')
+config.substitutions.insert(0, ('%imisched', f"{config.python_executable} {imisched_path}"))

llvm/utils/imisched.py

@@ -5,18 +5,35 @@
 # (c) Copyright 2024 Advanced Micro Devices, Inc. or its affiliates
 
 import argparse
+import io
+import os
+import sys
+
 from mischedutils import schedlogparser
 
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("sched_log")
-    parser.add_argument("-d", "--dump", action="store_true")
+    parser.add_argument("sched_log", help="Log file to parse, or '-' to read stdin")
+    parser.add_argument(
+        "-d", "--dump", action="store_true", help="dump the sched regions"
+    )
     args = parser.parse_args()
 
     regions = []
-    with open(args.sched_log, "r") as file:
-        regions = schedlogparser.process_log(file.read())
+    if args.sched_log != "-":
+        if not sys.__stdin__.isatty():
+            print(
+                "Error: cannot get data from stdin and a file at the same time",
+                file=sys.stderr,
+            )
+            exit(1)
+        print(f"Parsing '{args.sched_log}'...")
+        with open(args.sched_log, "r") as file:
+            regions = schedlogparser.process_log(file.read())
+    else:
+        print("Parsing stdin...")
+        regions = schedlogparser.process_log(sys.__stdin__.read())
 
     if args.dump:
         for region in regions:
@@ -29,6 +46,10 @@ def main():
     else:
         from mischedutils import interactive_sched
 
+        # Make sure Textual gets a fresh stdin for user keyboard inputs.
+        sys.__stdin__.flush()
+        sys.__stdin__ = io.TextIOWrapper(io.FileIO(os.open(os.ctermid(), os.O_RDONLY)))
+
         app = interactive_sched.InteractiveMISchedApp(regions)
         app.run()