run_noc_pass.py

"""Frontend to run the NoC pass."""

__copyright__ = """
Copyright (c) 2024 RapidStream Design Automation, Inc. and contributors.
All rights reserved. The contributor(s) of this file has/have agreed to the
RapidStream Contributor License Agreement.
"""

import copy
import json
import os
import subprocess
import sys
from enum import Enum, auto
from typing import Any

import click

from device import Device
from gen_vivado_bd import gen_arm_bd_ddr, gen_arm_bd_hbm
from ir_helper import (
    FREQUENCY,
    parse_floorplan,
    parse_inter_slot,
    parse_mmap_noc,
    parse_top_mod,
    round_up_to_noc_tdata,
)
from noc_pass import greedy_selector, ilp_noc_selector, random_selector
from noc_rtl_wrapper import add_dont_touch, noc_rtl_wrapper
from tcl_helper import (
    dump_neg_paths_summary,
    dump_streams_loc_tcl,
    export_constraint,
    export_control_s_axi_constraint,
    export_noc_constraint,
    gen_vivado_prj_tcl,
    parse_neg_paths,
    print_mmap_noc_loc_tcl,
    print_stream_noc_loc_tcl,
)
from vh1582_nocgraph import vh1582_nocgraph
from vp1802_nocgraph import vp1802_nocgraph


class DeviceEnum(Enum):
    """Supported FPGA devices."""

    VH1582 = auto()
    VP1802 = auto()


class SelectorEnum(Enum):
    """Supported NoC selectors."""

    NONE = auto()
    EMPTY = auto()
    RANDOM = auto()
    GREEDY = auto()
    GRB = auto()


@click.command()
@click.option(
    "--rapidstream-json",
    help="add_pipeline.json generated by RapidStream's no-flatten pass.",
)
@click.option("--tapa-xo", help="The .xo file generated by TAPA for the NONE selector.")
@click.option(
    "--device-name",
    required=True,
    type=click.Choice([e.name for e in DeviceEnum]),
    help="Device name.",
)
@click.option(
    "--mmap-port-json", required=True, help="Design's memory port configurations."
)
@click.option(
    "--selector",
    required=True,
    type=click.Choice([e.name for e in SelectorEnum]),
    help="Selector for mapping FIFOs to NoC.",
)
@click.option(
    "--build-dir",
    default=None,
    help="<Optional> Absolute path of build directory.",
)
@click.option(
    "--tar",
    default=None,
    help="<Optional> File name for the build directory tar." "Exclude '.tar.gz'.",
)
@click.option(
    "--top-mod-name",
    default=None,
    help="<Optional> Top module name." "Required for --mmap-ilp or --rapidstream-json.",
)
@click.option("--mmap-ilp", is_flag=True, help="Runs MMAP port mapping ILP.")
def parse_arguments(**kwargs: dict[str, Any]) -> dict[str, Any]:
    """Parse and validate command-line arguments.

    Returns a dictionary containing the parsed command-line arguments.
    """

    # Print all options that were used
    print("Used options:")
    for option, value in kwargs.items():
        if value is not None:
            print(f"{option}: {value}")

    if kwargs["build_dir"] and kwargs["tar"]:
        raise click.BadParameter("Choose either --build-dir or --tar.")

    if not kwargs["rapidstream_json"] and not kwargs["tapa_xo"]:
        raise click.BadParameter(
            "Either --rapidstream-json or --tapa-xo inputs is required."
        )
    if kwargs["tapa_xo"] and (
        kwargs["top_mod_name"] is None
        or str(kwargs["selector"]) != SelectorEnum.NONE.name
    ):
        raise click.BadParameter(
            "--top-mod-name and --selector NONE is required when using --tapa-xo."
        )
    if kwargs["rapidstream_json"] and kwargs["top_mod_name"]:
        raise click.BadParameter(
            "--top-mod-name should not be provided when using --rapidstream-json."
        )

    return kwargs


if __name__ == "__main__":
    if not (args := parse_arguments(standalone_mode=False)):
        sys.exit(1)
    rapidstream_json = args["rapidstream_json"]
    tapa_xo = args["tapa_xo"]
    device_name = args["device_name"]
    mmap_port_json = args["mmap_port_json"]
    selector = args["selector"]
    build_dir = args["build_dir"]
    tar = args["tar"]
    top_mod_name = args["top_mod_name"]
    mmap_ilp = args["mmap_ilp"]

    # currently hard-coded parameters
    IMPL_FREQUENCY = "300.0"
    HBM_INIT_FILE = "/home/jakeke/rapidstream-noc/test/serpens32_nasa4704.mem"
    TB_FILE = "/home/jakeke/rapidstream-noc/test/jacobi2d_mem8_tb.sv"
    USE_M_AXI_FPD = False
    MULTI_SITE_NOC = False

    # intermediate dumps
    BD_NAME = "top_arm"
    GROUPED_MOD_NAME = "axis_noc_if"
    I_MMAP_PORT_JSON = "mmap_port.json"
    SELECTED_STREAMS_JSON = "noc_streams.json"
    NOC_PASS_JSON = "noc_pass.json"
    NOC_PASS_WRAPPER_JSON = "noc_pass_wrapper.json"
    RTL_FOLDER = "rtl/"
    NOC_CONSTRAINT_TCL = "noc_constraint.tcl"
    NOC_STREAMS_ATTR = "noc_streams_attr.json"
    CONSTRAINT_TCL = "constraint.tcl"
    VIVADO_BD_TCL = "arm_bd.tcl"
    VIVADO_PRJ_TCL = "run.tcl"
    DUMP_NEG_PATHS_TCL = "dump_neg_paths.tcl"

    with open(mmap_port_json, "r", encoding="utf-8") as file:
        mmap_port_ir = json.load(file)

    if device_name == DeviceEnum.VP1802.name:
        G = vp1802_nocgraph()
        PART_NUM = "xcvp1802-lsvc4072-2MP-e-S"
        BOARD_PART = "xilinx.com:vpk180:part0:1.2"
        cr_mapping = [
            [
                "CLOCKREGION_X0Y1:CLOCKREGION_X4Y4",
                "CLOCKREGION_X0Y5:CLOCKREGION_X4Y7",
                "CLOCKREGION_X0Y8:CLOCKREGION_X4Y10",
                "CLOCKREGION_X0Y11:CLOCKREGION_X4Y13",
            ],
            [
                "CLOCKREGION_X5Y1:CLOCKREGION_X9Y4",
                "CLOCKREGION_X5Y5:CLOCKREGION_X9Y7",
                "CLOCKREGION_X5Y8:CLOCKREGION_X9Y10",
                "CLOCKREGION_X5Y11:CLOCKREGION_X9Y13",
            ],
        ]

        D = Device(
            part_num=PART_NUM,
            board_part=BOARD_PART,
            slot_width=2,
            slot_height=4,
            noc_graph=G,
            nmu_per_slot=[],  # generated
            nsu_per_slot=[],  # generated
            cr_mapping=cr_mapping,
        )
    elif device_name == DeviceEnum.VH1582.name:
        G = vh1582_nocgraph()
        PART_NUM = "xcvh1582-vsva3697-2MP-e-S"
        BOARD_PART = "xilinx.com:vhk158:part0:1.1"
        cr_mapping = [
            [
                "CLOCKREGION_X0Y1:CLOCKREGION_X4Y4",
                "CLOCKREGION_X0Y5:CLOCKREGION_X4Y7",
            ],
            [
                "CLOCKREGION_X5Y1:CLOCKREGION_X9Y4",
                "CLOCKREGION_X5Y5:CLOCKREGION_X9Y7",
            ],
        ]

        D = Device(
            part_num=PART_NUM,
            board_part=BOARD_PART,
            slot_width=2,
            slot_height=2,
            noc_graph=G,
            nmu_per_slot=[],  # generated
            nsu_per_slot=[],  # generated
            cr_mapping=cr_mapping,
        )
    else:
        raise NotImplementedError

    if tar:
        build_dir = tar
        # build_dir = f"{os.getcwd()}/{tar}"

    if build_dir:
        if os.path.exists(build_dir):
            print(f"The folder '{build_dir}' already exists. Aborting.")
            sys.exit(1)

    zsh_cmds = f"""
mkdir {build_dir}
cp {mmap_port_json} {build_dir}/{I_MMAP_PORT_JSON}
"""
    if rapidstream_json:
        zsh_cmds += f"cp {rapidstream_json} {build_dir}/"
    elif tapa_xo:
        zsh_cmds += f"cp {tapa_xo} {build_dir}/"
    else:
        raise NotImplementedError
    print(zsh_cmds)
    subprocess.run(["zsh", "-c", zsh_cmds], check=True)

    # Select MMAP ports to put over NoC

    # Select FIFOs to put over NoC
    if selector == SelectorEnum.NONE.name:
        streams_slots: dict[str, dict[str, str]] = {}
        noc_streams: list[str] = []
    else:
        with open(rapidstream_json, "r", encoding="utf-8") as file:
            rapidstream_ir = json.load(file)

        top_mod_name = rapidstream_ir["modules"]["top_name"]
        streams_slots, streams_widths = parse_inter_slot(parse_top_mod(rapidstream_ir))

        streams_bw = {s: w * FREQUENCY / 8 for s, w in streams_widths.items()}

        for s, attr in streams_slots.items():
            print(s, attr, streams_widths[s], streams_bw[s])
        assert len(streams_bw) == len(streams_slots), "parse_inter_slot ERROR"

        if selector == SelectorEnum.EMPTY.name:
            noc_streams = []
        elif selector == SelectorEnum.RANDOM.name:
            noc_streams = random_selector(streams_slots, D)
        elif selector == SelectorEnum.GREEDY.name:
            noc_streams = greedy_selector(streams_slots, D)
        elif selector == SelectorEnum.GRB.name:
            mmap_noc, mmap_bw = parse_mmap_noc(mmap_port_ir)
            noc_streams, node_loc = ilp_noc_selector(
                streams_slots, streams_bw, mmap_noc, mmap_bw, D
            )
        else:
            raise NotImplementedError

    print("Top module name:", top_mod_name)
    print("Number of inter-slot streams:", len(streams_slots))
    print("Selected streams for NoC", noc_streams)
    for s in noc_streams:
        print(f"{s}\t {streams_slots[s]}\t {streams_widths[s]}")

    # Dump outputs or
    if not build_dir:
        print("No build directory to dump outputs.")
        sys.exit(1)

    # dumps the selected streams json
    noc_stream_json = {GROUPED_MOD_NAME: noc_streams}
    with open(f"{build_dir}/{SELECTED_STREAMS_JSON}", "w", encoding="utf-8") as file:
        json.dump(noc_stream_json, file, indent=4)

    # generate grouped ir with the selected streams
    cc_ret_noc_stream: dict[str, dict[str, str]] = {}
    if selector == SelectorEnum.NONE.name:
        # skip generating grouped ir and wrapper
        zsh_cmds = f"""
unzip {tapa_xo} -d {build_dir}/tmp
mv {build_dir}/tmp/ip_repo/*/src {build_dir}/rtl
"""
    else:
        if selector == SelectorEnum.EMPTY.name:
            # skip generating grouped ir and wrapper
            # noc_pass_wrapper_ir = rapidstream_ir

            # but add dont_touch to pipelining registers
            noc_pass_wrapper_ir = copy.deepcopy(rapidstream_ir)
            add_dont_touch(noc_pass_wrapper_ir)
        else:
            zsh_cmds = f"""
source ~/.zshrc && amd
rapidstream-optimizer -i {rapidstream_json} -o {build_dir}/{NOC_PASS_JSON} \
create-group-wrapper --group-name-to-insts-json={build_dir}/{SELECTED_STREAMS_JSON}
"""
            print(zsh_cmds)
            subprocess.run(["zsh", "-c", zsh_cmds], check=True)

            # generate new rtl wrapper
            with open(f"{build_dir}/{NOC_PASS_JSON}", "r", encoding="utf-8") as file:
                noc_pass_ir = json.load(file)

            noc_pass_wrapper_ir, cc_ret_noc_stream = noc_rtl_wrapper(
                noc_pass_ir, GROUPED_MOD_NAME
            )
            for s, attr in cc_ret_noc_stream.items():
                print(f'{s}\t {attr["width"]}\t {attr["bandwidth"]}')

        with open(
            f"{build_dir}/{NOC_PASS_WRAPPER_JSON}", "w", encoding="utf-8"
        ) as file:
            json.dump(noc_pass_wrapper_ir, file, indent=4)

        zsh_cmds = f"""
rapidstream-exporter -i {build_dir}/{NOC_PASS_WRAPPER_JSON} -f {build_dir}/rtl
"""
    # generate rtl folder
    print(zsh_cmds)
    subprocess.run(["zsh", "-c", zsh_cmds], check=True)

    # export noc IPI constraints
    tcl = []
    if mmap_ilp:
        # single site NoC constraint found by ILP
        tcl = print_mmap_noc_loc_tcl(
            [attr["noc"] for n, attr in mmap_port_ir.items() if attr.get("noc")]
        )

    if MULTI_SITE_NOC:
        # multi-site NoC constraints
        tcl += dump_streams_loc_tcl(
            streams_slots | cc_ret_noc_stream,
            noc_streams + list(cc_ret_noc_stream.keys()),
            D,
        )
    elif selector == SelectorEnum.GRB.name:
        # single site NoC constraint found by ILP
        tcl += print_stream_noc_loc_tcl(node_loc)

    with open(f"{build_dir}/{NOC_CONSTRAINT_TCL}", "w", encoding="utf-8") as file:
        file.write("\n".join(tcl))

    # generate vivado bd tcl
    bd_attr = {
        "bd_name": BD_NAME,
        "top_mod": top_mod_name,
        "hbm_init_file": HBM_INIT_FILE,
        "frequency": IMPL_FREQUENCY,
    }

    noc_stream_attr: dict[str, dict[str, str]] = {}
    for s in noc_streams:
        noc_stream_attr[f"m_axis_{s}"] = {
            "dest": f"s_axis_{s}",
            "bandwidth": str(streams_bw[s]),
            "width": round_up_to_noc_tdata(str(streams_widths[s]), False),
        }

    for s, attr in cc_ret_noc_stream.items():
        noc_stream_attr[f"m_axis_{s}"] = {
            "dest": f"s_axis_{s}",
            "bandwidth": attr["bandwidth"],
            "width": attr["width"],
        }
    with open(f"{build_dir}/{NOC_STREAMS_ATTR}", "w", encoding="utf-8") as file:
        json.dump(noc_stream_attr, file, indent=4)

    if device_name == DeviceEnum.VP1802.name:
        tcl = gen_arm_bd_ddr(
            bd_attr=bd_attr,
            mmap_ports=mmap_port_ir,
            stream_attr=noc_stream_attr,
            fpd=USE_M_AXI_FPD,
        )
    elif device_name == DeviceEnum.VH1582.name:
        tcl = gen_arm_bd_hbm(
            bd_attr=bd_attr,
            mmap_ports=mmap_port_ir,
            stream_attr=noc_stream_attr,
            fpd=USE_M_AXI_FPD,
        )
    else:
        raise NotImplementedError
    with open(f"{build_dir}/{VIVADO_BD_TCL}", "w", encoding="utf-8") as file:
        file.write("\n".join(tcl))

    # export placement constraints
    if selector == SelectorEnum.NONE.name:
        tcl = []
    else:
        final_ir = (
            rapidstream_json
            if selector == SelectorEnum.EMPTY.name
            else f"{build_dir}/{NOC_PASS_WRAPPER_JSON}"
        )
        with open(final_ir, "r", encoding="utf-8") as file:
            noc_pass_wrapper_ir = json.load(file)

        floorplan = parse_floorplan(noc_pass_wrapper_ir, GROUPED_MOD_NAME)
        print("Number of modules:", sum(len(v) for v in floorplan.values()))
        print("Used slots: ", floorplan.keys())

        tcl = export_constraint(floorplan, D)

        # needed for multi-site NoC constraints
        if MULTI_SITE_NOC:
            tcl += export_noc_constraint(
                streams_slots | cc_ret_noc_stream,
                noc_streams + list(cc_ret_noc_stream.keys()),
                D,
            )

        if not USE_M_AXI_FPD:
            tcl += export_control_s_axi_constraint(floorplan, D)

    with open(f"{build_dir}/{CONSTRAINT_TCL}", "w", encoding="utf-8") as file:
        file.write("\n".join(tcl))

    # generate vivado prj tcl
    tcl = gen_vivado_prj_tcl(
        {
            "build_dir": build_dir,
            "part_num": D.part_num,
            "board_part": D.board_part,
            "bd_name": BD_NAME,
            "rtl_dir": RTL_FOLDER,
            "tb_file": TB_FILE,
            "constraint": CONSTRAINT_TCL,
            "bd_tcl": VIVADO_BD_TCL,
            "noc_tcl": NOC_CONSTRAINT_TCL,
        }
    )
    with open(f"{build_dir}/{VIVADO_PRJ_TCL}", "w", encoding="utf-8") as file:
        file.write("\n".join(tcl))

    tcl = dump_neg_paths_summary(build_dir)
    with open(f"{build_dir}/{DUMP_NEG_PATHS_TCL}", "w", encoding="utf-8") as file:
        file.write("\n".join(tcl))

    if tar:
        zsh_cmds = f"tar -czf {build_dir}.tar.gz {build_dir}\n"
        # delete the temporary tar directory
        zsh_cmds += f"rm -rf {build_dir}\n"
        print(zsh_cmds)
        subprocess.run(["zsh", "-c", zsh_cmds], check=True)
        sys.exit(1)

    # launch vivado
    zsh_cmds = f"""
source ~/.zshrc && amd
cd {build_dir}
vivado -mode batch -source {VIVADO_PRJ_TCL}
vivado -mode batch -source {DUMP_NEG_PATHS_TCL}
"""
    print(zsh_cmds)
    subprocess.run(["zsh", "-c", zsh_cmds], check=True)

    parse_neg_paths(build_dir, list(streams_slots.keys()), noc_streams)