[rtlsim] enable stitched-IP rtlsim with pyxsi

src/finn/core/rtlsim_exec.py

@@ -30,7 +30,7 @@
 from pyverilator.util.axi_utils import reset_rtlsim, rtlsim_multi_io
 from qonnx.custom_op.registry import getCustomOp
 
-from finn.util.basic import pyverilate_get_liveness_threshold_cycles
+from finn.util.basic import make_build_dir, pyverilate_get_liveness_threshold_cycles
 from finn.util.data_packing import npy_to_rtlsim_input, rtlsim_output_to_npy
 from finn.util.pyverilator import pyverilate_stitched_ip
 
@@ -39,35 +39,10 @@
 except ModuleNotFoundError:
     PyVerilator = None
 
+import finn.util.pyxsi_rpcclient as pyxsi_rpcclient
 
-def rtlsim_exec(model, execution_context, pre_hook=None, post_hook=None):
-    """Use PyVerilator to execute given model with stitched IP. The execution
-    context contains the input values. Hook functions can be optionally
-    specified to observe/alter the state of the circuit, receiving the
-    PyVerilator sim object as their first argument:
-    - pre_hook : hook function to be called before sim start (after reset)
-    - post_hook : hook function to be called after sim end
-    """
-    if PyVerilator is None:
-        raise ImportError("Installation of PyVerilator is required.")
-    # ensure stitched ip project already exists
-    assert os.path.isfile(
-        model.get_metadata_prop("wrapper_filename")
-    ), """The
-    file name from metadata property "wrapper_filename" doesn't exist."""
-    assert os.path.isdir(
-        model.get_metadata_prop("vivado_stitch_proj")
-    ), """The
-    directory from metadata property "vivado_stitch_proj" doesn't exist"""
-    trace_file = model.get_metadata_prop("rtlsim_trace")
-    if trace_file is None:
-        trace_file = ""
-    extra_verilator_args = model.get_metadata_prop("extra_verilator_args")
-    if extra_verilator_args is None:
-        extra_verilator_args = []
-    else:
-        extra_verilator_args = eval(extra_verilator_args)
 
+def prep_rtlsim_io_dict(model, execution_context):
     # extract i/o info to prepare io_dict
     io_dict = {"inputs": {}, "outputs": {}}
     if_dict = eval(model.get_metadata_prop("vivado_stitch_ifnames"))
@@ -123,6 +98,108 @@ def rtlsim_exec(model, execution_context, pre_hook=None, post_hook=None):
         o_stream_w = last_node.get_outstream_width()
         o_tensor_info.append((o_stream_w, o_dt, o_folded_shape, o_shape))
         num_out_values += batchsize * last_node.get_number_output_values()
+    return io_dict, if_dict, num_out_values, o_tensor_info
+
+
+def file_to_basename(x):
+    return os.path.basename(os.path.realpath(x))
+
+
+def rtlsim_exec_pyxsi(model, execution_context, pre_hook=None, post_hook=None):
+    """Use PyXSI to execute given model with stitched IP. The execution
+    context contains the input values. Hook functions can be optionally
+    specified to observe/alter the state of the circuit, receiving the
+    PyXSI RPC sim handle as their first argument:
+    - pre_hook : hook function to be called before sim start (after reset)
+    - post_hook : hook function to be called after sim end
+    """
+    # ensure stitched ip project already exists
+    assert os.path.isfile(
+        model.get_metadata_prop("wrapper_filename")
+    ), """The
+    file name from metadata property "wrapper_filename" doesn't exist."""
+    assert os.path.isdir(
+        model.get_metadata_prop("vivado_stitch_proj")
+    ), """The
+    directory from metadata property "vivado_stitch_proj" doesn't exist"""
+    trace_file = model.get_metadata_prop("rtlsim_trace")
+    io_dict, if_dict, num_out_values, o_tensor_info = prep_rtlsim_io_dict(model, execution_context)
+
+    # prepare rtlsim model
+    rtlsim_so = model.get_metadata_prop("rtlsim_so")
+    if (rtlsim_so is None) or (not os.path.isfile(rtlsim_so)):
+        vivado_stitch_proj_dir = model.get_metadata_prop("vivado_stitch_proj")
+        with open(vivado_stitch_proj_dir + "/all_verilog_srcs.txt", "r") as f:
+            all_verilog_srcs = f.read().split()
+        top_module_file_name = file_to_basename(model.get_metadata_prop("wrapper_filename"))
+        top_module_name = top_module_file_name.strip(".v")
+        single_src_dir = make_build_dir("rtlsim_" + top_module_name + "_")
+
+        rtlsim_so = pyxsi_rpcclient.compile_sim_obj(
+            top_module_name, all_verilog_srcs, single_src_dir
+        )
+        # save generated lib filename in attribute
+        model.set_metadata_prop("rtlsim_so", rtlsim_so[0] + "/" + rtlsim_so[1])
+        sim_base, sim_rel = rtlsim_so
+        # pass in correct tracefile from attribute
+        if trace_file == "default":
+            trace_file = top_module_file_name + ".wdb"
+        sim = pyxsi_rpcclient.load_sim_obj(sim_base, sim_rel, trace_file)
+    else:
+        sim = pyxsi_rpcclient.load_sim_obj(sim_base, sim_rel, trace_file)
+
+    # reset and call rtlsim, including any pre/post hooks
+    pyxsi_rpcclient.reset_rtlsim(sim)
+    if pre_hook is not None:
+        pre_hook(sim)
+    n_cycles = pyxsi_rpcclient.rtlsim_multi_io(
+        sim,
+        io_dict,
+        num_out_values,
+        sname="_",
+        liveness_threshold=pyverilate_get_liveness_threshold_cycles(),
+    )
+    if post_hook is not None:
+        post_hook(sim)
+    # important to call close_rtlsim for pyxsi to flush traces and stop
+    # the RPC server process
+    pyxsi_rpcclient.close_rtlsim(sim)
+
+    # unpack outputs and put back into execution context
+    for o, o_vi in enumerate(model.graph.output):
+        o_name = o_vi.name
+        if_name = if_dict["m_axis"][o][0]
+        o_stream_w, o_dt, o_folded_shape, o_shape = o_tensor_info[o]
+        packed_output = io_dict["outputs"][if_name]
+        o_folded_tensor = rtlsim_output_to_npy(
+            packed_output, None, o_dt, o_folded_shape, o_stream_w, o_dt.bitwidth()
+        )
+        execution_context[o_name] = o_folded_tensor.reshape(o_shape)
+
+    model.set_metadata_prop("cycles_rtlsim", str(n_cycles))
+
+
+def rtlsim_exec_pyverilator(model, execution_context, pre_hook=None, post_hook=None):
+    if PyVerilator is None:
+        raise ImportError("Installation of PyVerilator is required.")
+    # ensure stitched ip project already exists
+    assert os.path.isfile(
+        model.get_metadata_prop("wrapper_filename")
+    ), """The
+    file name from metadata property "wrapper_filename" doesn't exist."""
+    assert os.path.isdir(
+        model.get_metadata_prop("vivado_stitch_proj")
+    ), """The
+    directory from metadata property "vivado_stitch_proj" doesn't exist"""
+    trace_file = model.get_metadata_prop("rtlsim_trace")
+    if trace_file is None:
+        trace_file = ""
+    extra_verilator_args = model.get_metadata_prop("extra_verilator_args")
+    if extra_verilator_args is None:
+        extra_verilator_args = []
+    else:
+        extra_verilator_args = eval(extra_verilator_args)
+    io_dict, if_dict, num_out_values, o_tensor_info = prep_rtlsim_io_dict(model, execution_context)
 
     # prepare pyverilator model
     rtlsim_so = model.get_metadata_prop("rtlsim_so")
@@ -159,3 +236,21 @@ def rtlsim_exec(model, execution_context, pre_hook=None, post_hook=None):
         execution_context[o_name] = o_folded_tensor.reshape(o_shape)
 
     model.set_metadata_prop("cycles_rtlsim", str(n_cycles))
+
+
+def rtlsim_exec(model, execution_context, pre_hook=None, post_hook=None):
+    """Use PyVerilator or PyXSI to execute given model with stitched IP, depending
+    on the rtlsim_backend metadata_prop on the model. The execution
+    context contains the input values. Hook functions can be optionally
+    specified to observe/alter the state of the circuit, receiving the
+    PyVerilator sim object as their first argument:
+    - pre_hook : hook function to be called before sim start (after reset)
+    - post_hook : hook function to be called after sim end
+    """
+    backend = model.get_metadata_prop("rtlsim_backend")
+    if backend == "pyverilator":
+        rtlsim_exec_pyverilator(model, execution_context, pre_hook, post_hook)
+    elif backend == "pyxsi":
+        rtlsim_exec_pyxsi(model, execution_context, pre_hook, post_hook)
+    else:
+        assert False, f"Unrecognized rtlsim_backend value: {backend}"