[Code gen trafo] Small fixes in code gen trafo and finished test of function

src/finn/transformation/code_gen_transformation.py

@@ -13,8 +13,7 @@ def code_gen_transformation(node, context, model):
         # lookup op_type in registry of CustomOps
         inst = registry.custom_op[op_type](node)
 
-        # get the path of the code generation directory if already set
-        # check instance and check node attributes for value
+        # get the path of the code generation directory
         code_gen_dir = inst.code_gen_dir
         code_gen_dir = code_gen_dir.s.decode("UTF-8")
 
@@ -50,6 +49,7 @@ def code_gen_transformation(node, context, model):
                             model.set_attribute(node, "code_gen_dir", tmp_dir)
                     else:
                         raise Exception("Code was not generated!")
+                # else: attribute is correctly set
             else:
                 inst.code_gen_dir = tmp_dir
                 model.set_attribute(node, "code_gen_dir", tmp_dir)

tests/transformation/test_code_gen_trafo.py

+import os
+
 import numpy as np
 from onnx import TensorProto, helper
 
@@ -59,3 +61,18 @@ def test_code_gen_trafo():
     model.set_initializer("thresh", T)
 
     model = model.transform(CodeGen())
+    for node in model.graph.node:
+        code_gen_attribute = util.get_by_name(node.attribute, "code_gen_dir")
+        tmp_dir = code_gen_attribute.s.decode("UTF-8")
+        assert os.path.isdir(
+            tmp_dir
+        ), """Code generation directory of node with
+            op type {} does not exist!""".format(
+            node.op_type
+        )
+        assert (
+            len(os.listdir(tmp_dir)) != 0
+        ), """Code generation directory of node with
+            op type {} is empty!""".format(
+            node.op_type
+        )