Control where TB writes output (stdout, file, or both) (#1249)

Author: vloncar

hls4ml/backends/vitis/vitis_backend.py

@@ -50,6 +50,7 @@ def create_initial_config(
         namespace=None,
         write_weights_txt=True,
         write_tar=False,
+        tb_output_stream='both',
         **_,
     ):
         """Create initial configuration of the Vitis backend.
@@ -64,6 +65,8 @@ def create_initial_config(
             write_weights_txt (bool, optional): If True, writes weights to .txt files which speeds up compilation.
                 Defaults to True.
             write_tar (bool, optional): If True, compresses the output directory into a .tar.gz file. Defaults to False.
+            tb_output_stream (str, optional): Controls where to write the output. Options are 'stdout', 'file' and 'both'.
+                Defaults to 'both'.
 
         Returns:
             dict: initial configuration.
@@ -79,6 +82,7 @@ def create_initial_config(
             'Namespace': namespace,
             'WriteWeightsTxt': write_weights_txt,
             'WriteTar': write_tar,
+            'TBOutputStream': tb_output_stream,
         }
 
         return config

hls4ml/backends/vivado/vivado_backend.py

@@ -202,6 +202,7 @@ def create_initial_config(
         namespace=None,
         write_weights_txt=True,
         write_tar=False,
+        tb_output_stream='both',
         **_,
     ):
         """Create initial configuration of the Vivado backend.
@@ -216,6 +217,8 @@ def create_initial_config(
             write_weights_txt (bool, optional): If True, writes weights to .txt files which speeds up compilation.
                 Defaults to True.
             write_tar (bool, optional): If True, compresses the output directory into a .tar.gz file. Defaults to False.
+            tb_output_stream (str, optional): Controls where to write the output. Options are 'stdout', 'file' and 'both'.
+                Defaults to 'both'.
 
         Returns:
             dict: initial configuration.
@@ -231,6 +234,7 @@ def create_initial_config(
             'Namespace': namespace,
             'WriteWeightsTxt': write_weights_txt,
             'WriteTar': write_tar,
+            'TBOutputStream': tb_output_stream,
         }
 
         return config

hls4ml/model/graph.py

@@ -60,6 +60,7 @@ def __init__(self, config):
                 'Namespace': None,
                 'WriteWeightsTxt': True,
                 'WriteTar': False,
+                'TBOutputStream': 'both',
             }
 
         self._parse_hls_config()

hls4ml/writer/vivado_writer.py

@@ -570,20 +570,25 @@ def write_test_bench(self, model):
 
             elif '// hls-fpga-machine-learning insert tb-output' in line:
                 newline = line
-                for out in model_outputs:
-                    newline += indent + 'nnet::print_result<{}, {}>({}, fout);\n'.format(
-                        out.type.name, out.size_cpp(), out.name
-                    )  # TODO enable this
+                tb_stream = model.config.get_writer_config().get('TBOutputStream', 'both')
+                if tb_stream != 'stdout':
+                    for out in model_outputs:
+                        newline += indent + 'nnet::print_result<{}, {}>({}, fout);\n'.format(
+                            out.type.name, out.size_cpp(), out.name
+                        )  # TODO enable this
 
             elif (
                 '// hls-fpga-machine-learning insert output' in line
                 or '// hls-fpga-machine-learning insert quantized' in line
             ):
                 newline = line
-                for out in model_outputs:
-                    newline += indent + 'nnet::print_result<{}, {}>({}, std::cout, true);\n'.format(
-                        out.type.name, out.size_cpp(), out.name
-                    )
+                tb_stream = model.config.get_writer_config().get('TBOutputStream', 'both')
+                keep_output = str(tb_stream != 'stdout').lower()  # We keep output if we need to write it to file too.
+                if tb_stream != 'file':
+                    for out in model_outputs:
+                        newline += indent + 'nnet::print_result<{}, {}>({}, std::cout, {});\n'.format(
+                            out.type.name, out.size_cpp(), out.name, keep_output
+                        )
 
             elif '// hls-fpga-machine-learning insert namespace' in line:
                 newline = ''

test/pytest/test_writer_config.py

@@ -14,7 +14,7 @@
 @pytest.fixture(scope='module')
 def keras_model():
     model = Sequential()
-    model.add(Dense(10, activation='softmax', input_shape=(15,)))
+    model.add(Dense(10, kernel_initializer='zeros', use_bias=False, input_shape=(15,)))
     model.compile()
     return model
 
@@ -69,3 +69,40 @@ def test_write_weights_txt(keras_model, write_weights_txt, backend):
 
     txt_written = os.path.exists(odir + '/firmware/weights/w2.txt')
     assert txt_written == write_weights_txt
+
+
+@pytest.mark.skip(reason='Skipping for now as it needs the installation of the compiler.')
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis'])
+@pytest.mark.parametrize('tb_output_stream', ['stdout', 'file', 'both'])
+def test_tb_output_stream(capfd, keras_model, tb_output_stream, backend):
+
+    config = hls4ml.utils.config_from_keras_model(keras_model, granularity='name')
+    odir = str(test_root_path / f'hls4mlprj_tb_output_stream_{tb_output_stream}_{backend}')
+    if os.path.exists(odir):
+        shutil.rmtree(odir)
+
+    hls_model = hls4ml.converters.convert_from_keras_model(
+        keras_model,
+        io_type='io_stream',
+        hls_config=config,
+        output_dir=odir,
+        backend=backend,
+        tb_output_stream=tb_output_stream,
+    )
+    hls_model.build(csim=True, synth=False)
+
+    # Check the output based on tb_output_stream
+    tb_file_path = os.path.join(odir, 'tb_data/csim_results.log')
+
+    with open(tb_file_path) as tb_file:
+        tb_content = tb_file.read()
+        if tb_output_stream in ['file', 'both']:
+            assert len(tb_content) > 0, 'Testbench output file expected to contain model outputs, but is empty'
+        else:
+            assert len(tb_content) == 0, 'Testbench output file expected to be empty, but contains data'
+
+    captured = capfd.readouterr()
+    if tb_output_stream in ['stdout', 'both']:
+        assert '0 0 0 0 0 0 0 0 0 0' in captured.out, 'Expected model output not found in stdout'
+    else:
+        assert '0 0 0 0 0 0 0 0 0 0' not in captured.out, 'Model output should not be printed to stdout'