artiq/artiq/frontend/artiq_ddb_template.py

#!/usr/bin/env python3

import argparse
import sys
import json
import textwrap
from collections import defaultdict


def process_header(output, description):
    if description["target"] != "kasli":
        raise NotImplementedError

    print(textwrap.dedent("""
        # Autogenerated for the {variant} variant
        core_addr = "{core_addr}"

        device_db = {{
            "core": {{
                "type": "local",
                "module": "artiq.coredevice.core",
                "class": "Core",
                "arguments": {{"host": core_addr, "ref_period": {ref_period}}}
            }},
            "core_log": {{
                "type": "controller",
                "host": "::1",
                "port": 1068,
                "command": "aqctl_corelog -p {{port}} --bind {{bind}} " + core_addr
            }},
            "core_cache": {{
                "type": "local",
                "module": "artiq.coredevice.cache",
                "class": "CoreCache"
            }},
            "core_dma": {{
                "type": "local",
                "module": "artiq.coredevice.dma",
                "class": "CoreDMA"
            }},

            "i2c_switch0": {{
                "type": "local",
                "module": "artiq.coredevice.i2c",
                "class": "PCA9548",
                "arguments": {{"address": 0xe0}}
            }},
            "i2c_switch1": {{
                "type": "local",
                "module": "artiq.coredevice.i2c",
                "class": "PCA9548",
                "arguments": {{"address": 0xe2}}
            }},
        }}
        """).format(
            variant=description["variant"],
            core_addr=description.get("core_addr", "192.168.1.70"),
            ref_period=1/(8*description.get("rtio_frequency", 125e6))),
        file=output)


class PeripheralManager:
    def __init__(self, output, master_description):
        self.counts = defaultdict(int)
        self.output = output
        self.master_description = master_description

    def get_name(self, ty):
        count = self.counts[ty]
        self.counts[ty] = count + 1
        return "{}{}".format(ty, count)

    def process_dio(self, rtio_offset, peripheral):
        class_names = {
            "input": "TTLInOut",
            "output": "TTLOut"
        }
        classes = [
            class_names[peripheral["bank_direction_low"]],
            class_names[peripheral["bank_direction_high"]]
        ]
        for i in range(8):
            print(textwrap.dedent("""
                device_db["{name}"] = {{
                    "type": "local",
                    "module": "artiq.coredevice.ttl",
                    "class": "{class_name}",
                    "arguments": {{"channel": 0x{channel:06x}}},
                }}
                """.format(
                    name=self.get_name("ttl"),
                    class_name=classes[i//4],
                    channel=rtio_offset+i)),
                file=self.output)
        return 8

    def process_urukul(self, rtio_offset, peripheral):
        return 0


    def process_sampler(self, rtio_offset, peripheral):
        return 0


    def process_zotino(self, rtio_offset, peripheral):
        return 0


    def process_grabber(self, rtio_offset, peripheral):
        return 0

    def process(self, rtio_offset, peripheral):
        processor = getattr(self, "process_"+str(peripheral["type"]))
        return processor(rtio_offset, peripheral)


def process(output, master_description, satellites):
    base = master_description["base"]
    if base not in ("standalone", "master"):
        raise ValueError("Invalid master base")

    if base == "standalone" and satellites:
        raise ValueError("A standalone system cannot have satellites")

    process_header(output, master_description)

    pm = PeripheralManager(output, master_description)

    print("# {} peripherals".format(base), file=output)
    rtio_offset = 0
    for peripheral in master_description["peripherals"]:
        n_channels = pm.process(rtio_offset, peripheral)
        rtio_offset += n_channels

    for destination, description in satellites:
        if description["base"] != "satellite":
            raise ValueError("Invalid base for satellite at destination {}".format(destination))

        print("# DEST#{} peripherals".format(destination), file=output)
        rtio_offset = destination << 16
        for peripheral in description["peripherals"]:
            n_channels = pm.process(rtio_offset, peripheral)
            rtio_offset += n_channels


def main():
    parser = argparse.ArgumentParser(
        description="ARTIQ device database template builder")
    parser.add_argument("master_description", metavar="MASTER_DESCRIPTION",
                        help="JSON system description file for the standalone or master node")
    parser.add_argument("-o", "--output",
                        help="output file, defaults to standard output if omitted")
    parser.add_argument("-s", "--satellite", nargs=2, action="append",
                        default=[], metavar=("DESTINATION", "DESCRIPTION"), type=str,
                        help="add DRTIO satellite at the given destination number with "
                             "devices from the given JSON description")

    args = parser.parse_args()

    with open(args.master_description, "r") as f:
        master_description = json.load(f)

    satellites = []
    for destination, description in args.satellite:
        with open(description, "r") as f:
            satellites.append((destination, json.load(f)))

    if args.output is not None:
        with open(args.output, "w") as f:
            process(f, master_description, satellites)
    else:
        process(sys.stdout, master_description, satellites)


if __name__ == "__main__":
    main()