tpu.py (1529B)
1 2class TPUDict(dict): 3 def __setattr__(self, __name, __value): 4 if __name in self.__dict__: 5 raise ValueError("Duplicate tpu location") 6 return super().__setattr__(__name, __value) 7 8def parse(asm): 9 part = { "tpus": TPUDict({}) } 10 current = None 11 lines = asm.strip().split("\n") 12 getxy = lambda l: [int(v[1:]) for v in l.split()[1:3]] 13 for l in lines: 14 if l.startswith("stdin"): 15 part["stdin"] = [*getxy(l.rsplit(" ", 1)[0]), l.split()[-1]] 16 elif l.startswith("stdout"): 17 part["stdout"] = [*getxy(l.rsplit(" ", 1)[0]), l.split()[-1]] 18 elif l.startswith("tpu"): 19 x,y = getxy(l) 20 part["tpus"][(x, y)] = [] 21 current = part["tpus"][(x,y)] 22 elif l.startswith("end"): 23 if "addr" not in part: 24 part["addr"] = int(current[0].strip().split()[-1], 0) 25 continue 26 elif l != "": 27 assert(current is not None) 28 current.append(l) 29 return part 30 31def write(part, filepath): 32 with open(filepath, "w+") as f: 33 if "stdin" in part: 34 x,y,d = part["stdin"] 35 f.write(f"stdin X{x} Y{y} {d}\n") 36 if "stdout" in part: 37 x,y,d = part["stdout"] 38 f.write(f"stdout X{x} Y{y} {d}\n") 39 f.write("\n") 40 for (x,y),ls in part["tpus"].items(): 41 f.write(f"tpu X{x} Y{y}\n") 42 for l in ls: 43 f.write(f"{l}\n") 44 f.write("end\n") 45 f.write("\n") 46 47