xref: /llvm-project/llvm/test/CodeGen/MLRegAlloc/Inputs/interactive_main.py (revision 085587e1a9cdf1625efca61153a7dfe30946a6ce)
165b40f27SMatt Arsenaultimport log_reader
265b40f27SMatt Arsenaultimport interactive_host
365b40f27SMatt Arsenaultimport sys
465b40f27SMatt Arsenault
5*085587e1SAiden Grossmanfrom typing import Sequence
665b40f27SMatt Arsenault
765b40f27SMatt Arsenaultdef main(args):
865b40f27SMatt Arsenault    # this advisor just picks the first legal register to evict, which is
965b40f27SMatt Arsenault    # identifiable by the "mask" feature
1065b40f27SMatt Arsenault    class Advisor:
1165b40f27SMatt Arsenault        to_return = False
1265b40f27SMatt Arsenault
13*085587e1SAiden Grossman        def advice(self, tensor_values: Sequence[log_reader.TensorValue]):
1465b40f27SMatt Arsenault            for tv in tensor_values:
1565b40f27SMatt Arsenault                if tv.spec().name != "mask":
1665b40f27SMatt Arsenault                    continue
1765b40f27SMatt Arsenault                for i, v in enumerate(tv):
1865b40f27SMatt Arsenault                    if v == 1:
1965b40f27SMatt Arsenault                        return i
2065b40f27SMatt Arsenault            # i.e. invalid:
2165b40f27SMatt Arsenault            return -1
2265b40f27SMatt Arsenault
2365b40f27SMatt Arsenault    a = Advisor()
2465b40f27SMatt Arsenault    interactive_host.run_interactive(args[0], a.advice, args[1:])
2565b40f27SMatt Arsenault
2665b40f27SMatt Arsenault
2765b40f27SMatt Arsenaultif __name__ == "__main__":
2865b40f27SMatt Arsenault    main(sys.argv[1:])
29