llvm/test/CodeGen/MLRegAlloc/Inputs/interactive_main.py

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