Merge pull request #1040 from JanFSchulte/updatepytests

Add Vitis to pytorch API tests

Author: jmitrevs

test/pytest/test_pytorch_api.py

@@ -22,7 +22,7 @@ def forward(self, x):
         return self.linear(x)
 
 
-@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus'])
 @pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
 def test_linear(backend, io_type):
     model = LinearModel()
@@ -73,7 +73,7 @@ def test_linear(backend, io_type):
         nn.Threshold(threshold=1.0, value=0.0),
     ],
 )
-@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus'])
 @pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
 def test_activations(activation_function, backend, io_type):
     model = torch.nn.Sequential(nn.Linear(1, 1), activation_function).to()
@@ -164,7 +164,7 @@ def forward(self, x):
         ThresholdModel(),
     ],
 )
-@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus'])
 @pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
 def test_activation_functionals(activation_function, backend, io_type):
     model = activation_function
@@ -201,7 +201,7 @@ def test_activation_functionals(activation_function, backend, io_type):
 
 
 @pytest.mark.parametrize('padds', padds_options)
-@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus'])
 @pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
 def test_conv1d(padds, backend, io_type):
     n_in = 2
@@ -242,7 +242,7 @@ def test_conv1d(padds, backend, io_type):
 
     if io_type == 'io_stream':
         # Vivado inserts and additional layer for 'same' padding in io_stream
-        if backend == "Vivado" and padds == 1:
+        if (backend == "Vivado" or backend == "Vitis") and padds == 1:
             assert nNodes == len(hls_model.get_layers())
         else:
             assert nNodes - 1 == len(hls_model.get_layers())
@@ -269,13 +269,13 @@ def test_conv1d(padds, backend, io_type):
     # if not (backend == 'Vivado' and io_type == 'io_stream' and padds == 1):
     conv_index = 2
     act_index = 3
-    if io_type == "io_stream" and not (backend == "Vivado" and padds == 1):
+    if io_type == "io_stream" and not ((backend == "Vivado" or backend == "Vitis") and padds == 1):
         conv_index = 1
         act_index = 2
     assert list(hls_model.get_layers())[conv_index].attributes['name'] == convNode.name
     assert list(hls_model.get_layers())[conv_index].attributes['class_name'] == 'Conv1D'
     assert list(hls_model.get_layers())[act_index].attributes['activation'] == class_object_relu.__class__.__name__
-    if io_type == "io_stream" and backend == "Vivado" and padds == 1:
+    if io_type == "io_stream" and (backend == "Vivado" or backend == "Vitis") and padds == 1:
         assert list(hls_model.get_layers())[conv_index].attributes["in_width"] == size_in + 2
     else:
         assert list(hls_model.get_layers())[conv_index].attributes["in_width"] == size_in
@@ -287,7 +287,7 @@ def test_conv1d(padds, backend, io_type):
         padding = 0
     else:
         padding = 1
-    if io_type == "io_stream" and backend == "Vivado" and padds == 1:
+    if io_type == "io_stream" and (backend == "Vivado" or backend == "Vitis") and padds == 1:
         padding = 1
         padds = 0
 
@@ -311,7 +311,7 @@ def test_conv1d(padds, backend, io_type):
 
 
 @pytest.mark.parametrize('padds', padds_options)
-@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus'])
 @pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
 def test_conv2d(padds, backend, io_type):
     n_in = 2
@@ -409,7 +409,7 @@ def test_conv2d(padds, backend, io_type):
     # results are not very good at the moment
     np.testing.assert_allclose(hls_prediction, pytorch_prediction, rtol=0, atol=5e-2)
 
-    if not (backend == 'Vivado' and io_type == 'io_stream' and padds == 1):
+    if not ((backend == 'Vivado' or backend == 'Vitis') and io_type == 'io_stream' and padds == 1):
         # Vivado inserts and additional layer for 'same' padding in io_stream
         conv_index = 2
         act_index = 3
@@ -464,7 +464,7 @@ def test_conv2d(padds, backend, io_type):
 
 @pytest.mark.parametrize('pooling', pooling_layers)
 @pytest.mark.parametrize('padds', padds_options)
-@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus'])
 def test_pooling(pooling, padds, backend):
     assert '1d' in pooling.__name__ or '2d' in pooling.__name__
 
@@ -588,7 +588,7 @@ def forward(self, x):
         return self.bn(x)
 
 
-@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus'])
 @pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
 def test_bn(backend, io_type):
     model = BatchNormModel()
@@ -631,7 +631,7 @@ def forward(self, x):
         return x
 
 
-@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus'])
 @pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
 def test_squeeze(backend, io_type):
     model = SqueezeModel()
@@ -667,7 +667,7 @@ def test_squeeze(backend, io_type):
         assert list(hls_model.get_layers())[3].attributes['target_shape'] == [3]
 
 
-@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus'])
 def test_flatten(backend):
     input = torch.randn(1, 1, 5, 5)
     model = nn.Sequential(nn.Conv2d(1, 32, 5, 1, 1), nn.Flatten(), nn.ReLU())
@@ -711,7 +711,7 @@ def forward(self, x):
         return x
 
 
-@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus'])
 @pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
 def test_skipped_layers(backend, io_type):
     model = ModelSkippedLayers()
@@ -742,7 +742,7 @@ def test_skipped_layers(backend, io_type):
     np.testing.assert_allclose(hls_prediction, pytorch_prediction, rtol=0, atol=5e-2)
 
 
-@pytest.mark.parametrize('backend', ['Vivado', 'Quartus'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus'])
 @pytest.mark.parametrize('io_type', ['io_parallel'])  # Only io_parallel for now
 @pytest.mark.parametrize('tensor_rank', [2, 3])
 def test_remove_transpose(backend, io_type, tensor_rank):