diff --git a/tests/ttnn/unit_tests/operations/test_conv2d.py b/tests/ttnn/unit_tests/operations/test_conv2d.py index 9b57fb5b790..e6f02dccdf7 100644 --- a/tests/ttnn/unit_tests/operations/test_conv2d.py +++ b/tests/ttnn/unit_tests/operations/test_conv2d.py @@ -1632,3 +1632,109 @@ def test_yolov4_conv_groups_1_low_resolution( padded_input_channels=16 if input_channels == 3 else None, output_layout=output_layout, ) + + +@skip_for_wormhole_b0() +@pytest.mark.parametrize( + "batch_size, output_channels, input_channels, input_height, input_width, filter_height, filter_width, stride_h, stride_w, pad_h, pad_w, use_1d_systolic_array, config_override, use_shallow_conv_variant", + ( + (1, 64, 64, 80, 80, 3, 3, 1, 1, 1, 1, True, None, False), + (1, 255, 512, 20, 20, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 64, 128, 80, 80, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 32, 3, 320, 320, 3, 3, 1, 1, 1, 1, True, None, False), + (1, 128, 128, 80, 80, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 128, 256, 20, 20, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 512, 256, 20, 20, 3, 3, 1, 1, 1, 1, True, None, False), + (1, 64, 32, 320, 320, 3, 3, 2, 2, 1, 1, True, None, False), + (1, 256, 512, 20, 20, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 64, 64, 80, 80, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 256, 128, 40, 40, 3, 3, 1, 1, 1, 1, True, None, False), + (1, 512, 1024, 10, 10, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 512, 256, 40, 40, 3, 3, 2, 2, 1, 1, True, None, False), + (1, 64, 128, 160, 160, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 256, 512, 10, 10, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 256, 128, 40, 40, 3, 3, 2, 2, 1, 1, True, None, False), + (1, 64, 64, 160, 160, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 512, 512, 20, 20, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 512, 256, 20, 20, 3, 3, 2, 2, 1, 1, True, None, False), + (1, 128, 256, 40, 40, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 255, 256, 40, 40, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 512, 2048, 10, 10, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 512, 512, 10, 10, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 64, 32, 160, 160, 3, 3, 1, 1, 1, 1, True, None, False), + (1, 128, 128, 40, 40, 3, 3, 1, 1, 1, 1, True, None, False), + (1, 32, 64, 160, 160, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 256, 128, 80, 80, 3, 3, 2, 2, 1, 1, True, None, False), + (1, 256, 256, 20, 20, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 256, 256, 20, 20, 3, 3, 1, 1, 1, 1, True, None, False), + (1, 512, 512, 10, 10, 3, 3, 1, 1, 1, 1, True, None, False), + (1, 1024, 512, 10, 10, 3, 3, 1, 1, 1, 1, True, None, False), + (1, 128, 64, 160, 160, 3, 3, 2, 2, 1, 1, True, None, False), + (1, 255, 1024, 10, 10, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 1024, 1024, 10, 10, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 1024, 512, 20, 20, 3, 3, 2, 2, 1, 1, True, None, False), + (1, 256, 256, 40, 40, 1, 1, 1, 1, 0, 0, True, None, False), + (1, 128, 128, 40, 40, 1, 1, 1, 1, 0, 0, True, None, False), + ), +) +@pytest.mark.parametrize( + "weights_dtype", + [ttnn.bfloat8_b], +) +@pytest.mark.parametrize( + "activations_dtype", + # [ttnn.bfloat8_b, ttnn.bfloat16], + [ttnn.bfloat8_b], +) +@pytest.mark.parametrize("math_fidelity", [ttnn.MathFidelity.LoFi]) +# @pytest.mark.parametrize("output_layout", [ttnn.ROW_MAJOR_LAYOUT, ttnn.TILE_LAYOUT]) +@pytest.mark.parametrize("output_layout", [ttnn.TILE_LAYOUT]) +def test_yolov4_REPO_320_320( + device, + use_program_cache, + math_fidelity, + activations_dtype, + weights_dtype, + batch_size, + output_channels, + input_channels, + input_height, + input_width, + filter_height, + filter_width, + stride_h, + stride_w, + pad_h, + pad_w, + use_1d_systolic_array, + config_override, + use_shallow_conv_variant, + output_layout, +): + if output_layout == ttnn.ROW_MAJOR_LAYOUT and activations_dtype == ttnn.bfloat8_b: + pytest.skip("Row major layout not compatible with bfloat8_b") + if output_layout == ttnn.ROW_MAJOR_LAYOUT and input_height >= 1056: + pytest.skip("OOM") + run_conv( + device, + math_fidelity, + activations_dtype, + weights_dtype, + batch_size, + output_channels, + input_channels, + input_height, + input_width, + filter_height, + filter_width, + stride_h, + stride_w, + pad_h, + pad_w, + use_1d_systolic_array, + config_override, + use_shallow_conv_variant=use_shallow_conv_variant, + # groups=groups, + padded_input_channels=16 if input_channels == 3 else None, + output_layout=output_layout, + )