run.py

import subprocess
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--futhark", type=str, default="./futhark-nightly/bin/futhark",
                    help="Path to Futhark binary.")
parser.add_argument("--backend", type=str, default="cuda",
                    help="The backend used when compiling programs.")
parser.add_argument("-d", "--device", type=int, default=0,
                    help="Device number.")
parser.add_argument("-r", "--runs", type=int, default=10)
parser.add_argument("command", type=str, choices=["bench", "autotune"])
parser.add_argument("filename", type=str)

args = parser.parse_args()

# parser.add_argument("-R", "--reg", type=int, default=None,
#                     help="Default register tile size.")
# parser.add_argument("-T", "--tile", type=int, default=None,
#                     help="Default tile size.")

reg_tile = [
  # {"reg": None, "tile": None},
  {"reg": 5, "tile": 16},
  # {"reg": 4, "tile": 32},
]

def run(reg=None, tile=None, compile=True):
  cmd = ([
      args.futhark,
      args.command,
      f"--backend={args.backend}",
      f"--pass-option=--device=#{args.device}",
      f"--tuning=reg{reg}-tile{tile}.tuning",
      f"--runs={args.runs}",
    ] + ([
      f"--json={args.filename}.reg{reg}-tile{tile}.json",
    ] if args.command == "bench" else [] ) + ([
      f"--pass-option=--default-reg-tile-size={reg}",
      f"--pass-option=--default-tile-size={tile}",
    ] if reg is not None and tile is not None else []) + ([
      "--skip-compilation",
    ] if not compile else [])
    + [ args.filename ]
  )
  print()
  print(" ".join(cmd))
  subprocess.run(cmd)

if __name__ == "__main__":
  compile = True
  for kwargs in reg_tile:
    run(**kwargs, compile=compile)
    compile = False