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