xref: /llvm-project/llvm/test/CodeGen/MLRegAlloc/Inputs/interactive_main.py (revision 085587e1a9cdf1625efca61153a7dfe30946a6ce)
1import log_reader
2import interactive_host
3import sys
4
5from typing import Sequence
6
7def main(args):
8    # this advisor just picks the first legal register to evict, which is
9    # identifiable by the "mask" feature
10    class Advisor:
11        to_return = False
12
13        def advice(self, tensor_values: Sequence[log_reader.TensorValue]):
14            for tv in tensor_values:
15                if tv.spec().name != "mask":
16                    continue
17                for i, v in enumerate(tv):
18                    if v == 1:
19                        return i
20            # i.e. invalid:
21            return -1
22
23    a = Advisor()
24    interactive_host.run_interactive(args[0], a.advice, args[1:])
25
26
27if __name__ == "__main__":
28    main(sys.argv[1:])
29