diff --git a/src/finn/core/modelwrapper.py b/src/finn/core/modelwrapper.py
index 1d5d3a978948998dce95d74ee780dfb998376465..473a33c47885a9e21e0c331233b85fb265933456 100644
--- a/src/finn/core/modelwrapper.py
+++ b/src/finn/core/modelwrapper.py
@@ -115,6 +115,19 @@ class ModelWrapper:
qa.quant_parameter_tensor_names.append(dt)
qnt_annotations.append(qa)
+ def get_tensor_valueinfo(self, tensor_name):
+ """Returns ValueInfoProto of tensor with given name, if it has one."""
+ graph = self._model_proto.graph
+ vi_names = [(x.name, x) for x in graph.input]
+ vi_names += [(x.name, x) for x in graph.output]
+ vi_names += [(x.name, x) for x in graph.value_info]
+ try:
+ vi_ind = [x[0] for x in vi_names].index(tensor_name)
+ vi = vi_names[vi_ind][1]
+ return vi
+ except ValueError:
+ return None
+
def get_tensor_shape(self, tensor_name):
"""Returns the shape of tensor with given name, if it has ValueInfoProto."""
graph = self._model_proto.graph
diff --git a/src/finn/custom_op/registry.py b/src/finn/custom_op/registry.py
index bd1605c96d8bb7a0687686b48912bd3b61a8f6cc..c38a28dae46f6e39ed45e58647385b79fe271bcc 100644
--- a/src/finn/custom_op/registry.py
+++ b/src/finn/custom_op/registry.py
@@ -7,6 +7,7 @@ from finn.custom_op.fpgadataflow.streamingfclayer_batch import StreamingFCLayer_
from finn.custom_op.fpgadataflow.streamingmaxpool_batch import StreamingMaxPool_Batch
from finn.custom_op.fpgadataflow.tlastmarker import TLastMarker
from finn.custom_op.multithreshold import MultiThreshold
+from finn.custom_op.streamingdataflowpartition import StreamingDataflowPartition
from finn.custom_op.xnorpopcount import XnorPopcountMatMul
# create a mapping of all known CustomOp names and classes
@@ -18,6 +19,7 @@ custom_op["StreamingMaxPool_Batch"] = StreamingMaxPool_Batch
custom_op["StreamingFCLayer_Batch"] = StreamingFCLayer_Batch
custom_op["ConvolutionInputGenerator"] = ConvolutionInputGenerator
custom_op["TLastMarker"] = TLastMarker
+custom_op["StreamingDataflowPartition"] = StreamingDataflowPartition
def getCustomOp(node):
diff --git a/src/finn/custom_op/streamingdataflowpartition.py b/src/finn/custom_op/streamingdataflowpartition.py
new file mode 100644
index 0000000000000000000000000000000000000000..856ba6d5604c37440985648bb91b18fffeaf79fd
--- /dev/null
+++ b/src/finn/custom_op/streamingdataflowpartition.py
@@ -0,0 +1,65 @@
+from finn.custom_op import CustomOp
+
+# note that the StreamingDataflowPartition node is only a meta/container node,
+# it does not produce any HLS or bitfile by itself. it's a placeholder for
+# a group of fpgadataflow nodes that have been separated out into a FINN-ONNX
+# model of its own.
+
+
+class StreamingDataflowPartition(CustomOp):
+ def get_nodeattr_types(self):
+ return {
+ "model": ("s", True, ""),
+ }
+
+ def make_shape_compatible_op(self):
+ pass
+
+ def infer_node_datatype(self, model):
+ pass
+
+ def execute_node(self, context, graph):
+ # TODO add RPC execution with synthesized bitfile?
+ # whole-design rtlsim with PyVerilator may also be an alternative
+ pass
+
+ def verify_node(self):
+ info_messages = []
+
+ # verify number of attributes
+ num_of_attr = 1
+ if len(self.onnx_node.attribute) == num_of_attr:
+ info_messages.append("The number of attributes is correct")
+ else:
+ info_messages.append(
+ """The number of attributes is incorrect,
+ {} should have {} attributes""".format(
+ self.onnx_node.op_type, num_of_attr
+ )
+ )
+
+ # verify that "domain" is set to "finn"
+ domain_value = self.onnx_node.domain
+ if domain_value == "finn":
+ info_messages.append("Attribute domain is set correctly")
+ else:
+ info_messages.append('Attribute domain should be set to "finn"')
+
+ # verify that all necessary attributes exist
+ try:
+ self.get_nodeattr("model")
+ info_messages.append("All necessary attributes exist")
+ except Exception:
+ info_messages.append(
+ """The necessary attributes do not exist.
+ StreamingDataflowPartition needs the following attribute(s):
+ model"""
+ )
+
+ # verify the number of inputs
+ if len(self.onnx_node.input) == 1:
+ info_messages.append("The number of inputs is correct")
+ else:
+ info_messages.append("StreamingDataflowPartition needs 1 data input")
+
+ return info_messages
diff --git a/src/finn/data/onnx/finn-hls-model/tfc_w1_a1_after_conv_to_hls.onnx b/src/finn/data/onnx/finn-hls-model/tfc_w1_a1_after_conv_to_hls.onnx
new file mode 100644
index 0000000000000000000000000000000000000000..aada6f07e9d3910122d2eb357d8a8c1224e9fbab
Binary files /dev/null and b/src/finn/data/onnx/finn-hls-model/tfc_w1_a1_after_conv_to_hls.onnx differ
diff --git a/src/finn/transformation/fpgadataflow/create_dataflow_partition.py b/src/finn/transformation/fpgadataflow/create_dataflow_partition.py
new file mode 100644
index 0000000000000000000000000000000000000000..af43b7ea55c30155ff0ff326d2a26942ab69bad2
--- /dev/null
+++ b/src/finn/transformation/fpgadataflow/create_dataflow_partition.py
@@ -0,0 +1,74 @@
+import copy
+
+from onnx import helper
+
+from finn.transformation import Transformation
+from finn.util.basic import get_by_name, make_build_dir
+
+
+class CreateDataflowPartition(Transformation):
+ """Split a graph into two graphs; one which contains non-FINN-dataflow nodes
+ and a StreamingDataflowPartition node, and another which only contains
+ FINN dataflow nodes. The StreamingDataflowPartition has a model attribute
+ that indicates the filename for the second graph that only contains
+ dataflow nodes. No action is taken if there are no dataflow nodes."""
+
+ def __init__(self):
+ super().__init__()
+
+ def apply(self, model):
+ # TODO we currently assume that all dataflow nodes are connected to
+ # each other, forming a single partition. check the assumption and/or
+ # improve this.
+ all_nodes = list(model.graph.node)
+ df_nodes = filter(
+ lambda x: get_by_name(x.attribute, "backend") is not None, all_nodes
+ )
+ df_nodes = filter(
+ lambda x: get_by_name(x.attribute, "backend").s.decode("UTF-8")
+ == "fpgadataflow",
+ df_nodes,
+ )
+ df_nodes = list(df_nodes)
+ non_df_nodes = filter(lambda x: x not in df_nodes, all_nodes)
+ non_df_nodes = list(non_df_nodes)
+
+ if len(df_nodes) == 0:
+ # no changes if no dataflow nodes are present
+ return (model, False)
+ else:
+ # partition the model into two models
+ df_model = copy.deepcopy(model)
+ non_df_model = model
+ # remove all non-dataflow nodes from the dataflow model
+ for node_to_remove in non_df_nodes:
+ df_model.graph.node.remove(node_to_remove)
+ # identify the entry and exit points for the dataflow part
+ df_in = df_model.graph.node[0].input[0]
+ df_out = df_model.graph.node[-1].output[0]
+ df_in_vi = df_model.get_tensor_valueinfo(df_in)
+ df_out_vi = df_model.get_tensor_valueinfo(df_out)
+ # set df graph in/out to be df_in/df_out
+ df_model.graph.input.remove(df_model.graph.input[0])
+ df_model.graph.input.insert(0, df_in_vi)
+ df_model.graph.output.remove(df_model.graph.output[0])
+ df_model.graph.output.insert(0, df_out_vi)
+ df_model_dir = make_build_dir("dataflow_partition_")
+ df_model_filename = df_model_dir + "/df_model.onnx"
+ df_model.save(df_model_filename)
+ # remove all dataflow nodes from the non-dataflow model
+ # keep track of where the dataflow part starts
+ df_start_ind = all_nodes.index(df_nodes[0])
+ for node_to_remove in df_nodes:
+ non_df_model.graph.node.remove(node_to_remove)
+ # create StreamingDataflow node with df_in/df_out io
+ df_node = helper.make_node(
+ "StreamingDataflowPartition",
+ [df_in],
+ [df_out],
+ # use the model attribute to mark the df model
+ model=df_model_filename,
+ )
+ non_df_model.graph.node.insert(df_start_ind, df_node)
+
+ return (non_df_model, False)
diff --git a/tests/fpgadataflow/test_create_dataflow_partition.py b/tests/fpgadataflow/test_create_dataflow_partition.py
new file mode 100644
index 0000000000000000000000000000000000000000..e0b8b491b0ed926a56331a8e125f3de4ecd91615
--- /dev/null
+++ b/tests/fpgadataflow/test_create_dataflow_partition.py
@@ -0,0 +1,21 @@
+import os.path
+from pkgutil import get_data
+
+from finn.core.modelwrapper import ModelWrapper
+from finn.custom_op.registry import getCustomOp
+from finn.transformation.fpgadataflow.create_dataflow_partition import (
+ CreateDataflowPartition,
+)
+
+
+def test_create_dataflow_partition():
+ # load the onnx model
+ raw_m = get_data(
+ "finn", "data/onnx/finn-hls-model/tfc_w1_a1_after_conv_to_hls.onnx"
+ )
+ model = ModelWrapper(raw_m)
+ model = model.transform(CreateDataflowPartition())
+ assert model.graph.node[2].op_type == "StreamingDataflowPartition"
+ sdp_node = getCustomOp(model.graph.node[2])
+ assert sdp_node.__class__.__name__ == "StreamingDataflowPartition"
+ assert os.path.isfile(sdp_node.get_nodeattr("model"))