Update loop rolling to custom op registration changes

Author: auphelia

src/finn/transformation/fpgadataflow/loop_rolling.py

@@ -8,11 +8,11 @@
 from qonnx.core.modelwrapper import ModelWrapper
 from qonnx.transformation.base import Transformation
 from qonnx.transformation.fold_constants import FoldConstants
+from qonnx.custom_op.registry import is_custom_op
 from typing import List, Tuple
 
 from finn.util import onnxscript_helpers as osh
 
-
 def get_constant_from_value(value):
     """
     Get the constant value of a tensor.
@@ -567,10 +567,12 @@ def apply(self, model: ModelWrapper) -> Tuple[ModelWrapper, bool]:
         from finn.util.basic import getHWCustomOp
 
         for loop_node in model_wrapper.get_nodes_by_op_type("FINNLoop"):
-            loop_body_graph = get_by_name(loop_node.attribute, "body").g
-            for node in loop_body_graph.node:
+            loop_body = getHWCustomOp(loop_node).get_nodeattr("body")
+            for node in loop_body.graph.node:
+                if not is_custom_op(node.domain):
+                    continue
                 try:
-                    inst = getHWCustomOp(node, model_wrapper)
+                    inst = getHWCustomOp(node)
                     inst.adapt_for_loop_body(LoopBody.signature)
                 except (KeyError, AttributeError):
                     # Operator doesn't need adaptation or doesn't support it

src/finn/util/onnxscript_helpers.py

@@ -17,7 +17,7 @@
     RewriterContext,
     pattern_builder,
 )
-from qonnx.util.basic import is_finn_op
+from qonnx.custom_op.registry import is_custom_op
 from typing import List, Optional
 
 
@@ -323,10 +323,10 @@ def is_fpgadataflow_onnxir_node(node):
     """Returns True if given node is fpgadataflow node. Otherwise False."""
     is_node = False
     if node is not None:
-        if is_finn_op(node.domain):
+        if is_custom_op(node.domain):
             if "backend" in node.attributes:
                 backend_value = node.attributes["backend"].as_string()
-                if backend_value == "fpgadataflow":
+                if backend_value in ["fpgadataflow", "hls", "rtl"]:
                     is_node = True
 
     return is_node