diff --git a/src/finn/custom_op/fpgadataflow/__init__.py b/src/finn/custom_op/fpgadataflow/__init__.py
index 14ba1b813c8d61145f7d221deee9c184aeb9bddc..71c731f96ca45519c443a5f932ead050770e17de 100644
--- a/src/finn/custom_op/fpgadataflow/__init__.py
+++ b/src/finn/custom_op/fpgadataflow/__init__.py
@@ -88,6 +88,8 @@ class HLSCustomOp(CustomOp):
"res_hls": ("s", False, ""),
"res_synth": ("s", False, ""),
"rtlsim_so": ("s", False, ""),
+ # partitioning info
+ "partition_id": ("i", False, 0),
# input and output FIFO depths
"inFIFODepth": ("i", False, 2),
"outFIFODepth": ("i", False, 2),
diff --git a/src/finn/custom_op/registry.py b/src/finn/custom_op/registry.py
index d1f8f02a00810804163918d5fd4336ab6523bde0..8b607abeba52f5f944eeaee6c6ee6eadefadff10 100644
--- a/src/finn/custom_op/registry.py
+++ b/src/finn/custom_op/registry.py
@@ -52,6 +52,7 @@ from finn.custom_op.fpgadataflow.labelselect_batch import LabelSelect_Batch
from finn.custom_op.quantavgpool2d import QuantAvgPool2d
from finn.custom_op.fpgadataflow.duplicatestreams_batch import DuplicateStreams_Batch
from finn.custom_op.fpgadataflow.channelwise_op_batch import ChannelwiseOp_Batch
+from finn.custom_op.fpgadataflow.iodma import IODMA
# create a mapping of all known CustomOp names and classes
custom_op = {}
@@ -76,6 +77,7 @@ custom_op["LabelSelect_Batch"] = LabelSelect_Batch
custom_op["QuantAvgPool2d"] = QuantAvgPool2d
custom_op["DuplicateStreams_Batch"] = DuplicateStreams_Batch
custom_op["ChannelwiseOp_Batch"] = ChannelwiseOp_Batch
+custom_op["IODMA"] = IODMA
def getCustomOp(node):
diff --git a/src/finn/transformation/fpgadataflow/floorplan.py b/src/finn/transformation/fpgadataflow/floorplan.py
new file mode 100644
index 0000000000000000000000000000000000000000..1d9a51875499d77f384c03f54009a9dd1001dea0
--- /dev/null
+++ b/src/finn/transformation/fpgadataflow/floorplan.py
@@ -0,0 +1,80 @@
+# Copyright (c) 2020, Xilinx
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# * Neither the name of FINN nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+from finn.custom_op.registry import getCustomOp
+from finn.transformation import Transformation
+from finn.util.basic import get_by_name
+
+
+class Floorplan(Transformation):
+ """Perform Floorplanning of the dataflow design. Separate DMAs into their own
+ partitions IDs, and TODO: split the design into sections of defined size"""
+
+ def __init__(self, limits=None):
+ super().__init__()
+ self.resource_limits = limits
+
+ def apply(self, model):
+ target_partition_id = 0
+ # we currently assume that all dataflow nodes belonging to the same partition
+ # are connected to each other and there is a single input/output to/from each.
+ all_nodes = list(model.graph.node)
+ df_nodes = list(
+ filter(lambda x: get_by_name(x.attribute, "backend") is not None, all_nodes)
+ )
+ dma_nodes = list(filter(lambda x: x.op_type == "IODMA", df_nodes))
+
+ non_dma_nodes = list(filter(lambda x: x not in dma_nodes, df_nodes))
+ dyn_tlastmarker_nodes = list(
+ filter(
+ lambda x: x.op_type == "TLastMarker"
+ and getCustomOp(x).get_nodeattr("DynIters") == "true",
+ non_dma_nodes,
+ )
+ )
+
+ non_dma_nodes = list(
+ filter(lambda x: x not in dyn_tlastmarker_nodes, non_dma_nodes)
+ )
+
+ for node in dma_nodes:
+ node_inst = getCustomOp(node)
+ node_inst.set_nodeattr("partition_id", target_partition_id)
+ target_partition_id += 1
+
+ for node in dyn_tlastmarker_nodes:
+ node_inst = getCustomOp(node)
+ node_inst.set_nodeattr("partition_id", target_partition_id)
+ target_partition_id += 1
+
+ for node in non_dma_nodes:
+ # TODO: implement proper floorplanning; for now just a single partition
+ node_inst = getCustomOp(node)
+ node_inst.set_nodeattr("partition_id", target_partition_id)
+
+ return (model, False)
diff --git a/tests/fpgadataflow/test_fpgadataflow_ipstitch.py b/tests/fpgadataflow/test_fpgadataflow_ipstitch.py
index b830693c32afe629dd6fc70868d0bddacac4c887..a9f5bf5ffa1f816b82ef701800e92249056b7c74 100644
--- a/tests/fpgadataflow/test_fpgadataflow_ipstitch.py
+++ b/tests/fpgadataflow/test_fpgadataflow_ipstitch.py
@@ -54,6 +54,10 @@ from finn.util.basic import gen_finn_dt_tensor, pynq_part_map
from finn.util.fpgadataflow import pyverilate_stitched_ip
from finn.util.test import load_test_checkpoint_or_skip
from finn.transformation.fpgadataflow.synth_ooc import SynthOutOfContext
+from finn.transformation.infer_data_layouts import InferDataLayouts
+from finn.transformation.fpgadataflow.insert_iodma import InsertIODMA
+from finn.transformation.fpgadataflow.floorplan import Floorplan
+
test_pynq_board = os.getenv("PYNQ_BOARD", default="Pynq-Z1")
test_fpga_part = pynq_part_map[test_pynq_board]
@@ -390,3 +394,19 @@ def test_fpgadataflow_ipstitch_remote_execution():
assert np.isclose(outp["outp"], x).all()
except KeyError:
pytest.skip("PYNQ board IP address not specified")
+
+
+def test_fpgadataflow_ipstitch_iodma_floorplan():
+ model = create_one_fc_model()
+ if model.graph.node[0].op_type == "StreamingDataflowPartition":
+ sdp_node = getCustomOp(model.graph.node[0])
+ assert sdp_node.__class__.__name__ == "StreamingDataflowPartition"
+ assert os.path.isfile(sdp_node.get_nodeattr("model"))
+ model = load_test_checkpoint_or_skip(sdp_node.get_nodeattr("model"))
+ model = model.transform(InferDataLayouts())
+ model = model.transform(InsertIODMA())
+ model = model.transform(Floorplan())
+ assert getCustomOp(model.graph.node[0]).get_nodeattr("partition_id") == 0
+ assert getCustomOp(model.graph.node[1]).get_nodeattr("partition_id") == 2
+ assert getCustomOp(model.graph.node[2]).get_nodeattr("partition_id") == 1
+ model.save(ip_stitch_model_dir + "/test_fpgadataflow_ipstitch_iodma_floorplan.onnx")