Arm backend: Add FP8 VGF operator tests (pytorch#24148)

Add E4M3 and E5M2 operator coverage with numerical comparisons where
PyTorch provides the required FP8 CPU kernels.

Centralize missing CPU reference handling in VgfPipelineFP8. Check VGF
execution and the expected missing CPU kernel error in one test instead
of two separate tests that each compile the model. Keeping this logic in
the pipeline makes future changes to the testing policy easier.

For cases marked expect_missing_reference=True, run VGF first without
output comparison, then check for the expected missing FP8 CPU kernel
error. Keep normal comparisons for cases with an available reference.

Handle expected errors only during reference execution so compiler and
runtime failures cannot be hidden. Fail when an expected missing CPU
reference starts working, prompting the test to enable comparison.

Authored with assistance from OpenAI Codex.

Change-Id: I762a13f5dd05c1682a2dd1b3a9631e51cd041eef
Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>

Author: YufengShi-dudu

backends/arm/test/ops/test_adaptive_avg_pool2d.py

@@ -18,6 +18,7 @@
     TosaPipelineFP,
     TosaPipelineINT,
     VgfPipeline,
+    VgfPipelineFP8,
 )
 
 exir_op = "executorch_exir_dialects_edge__ops_aten_avg_pool2d_default"
@@ -112,7 +113,7 @@ def forward(self, *args, **kwargs):
     ),
 }
 
-test_modules_fp8 = {
+test_modules_fp8_no_comparing = {
     "output_2x2_fp8e4m3": lambda: (
         AdaptiveAvgPool2d((2, 2)),
         (torch.rand(1, 4, 10, 10).to(torch.float8_e4m3fn),),
@@ -139,7 +140,7 @@ def test_adaptive_avg_pool2d_tosa_FP(test_module):
     pipeline.run()
 
 
-@common.parametrize("test_module", test_modules_fp8)
+@common.parametrize("test_module", test_modules_fp8_no_comparing)
 def test_adaptive_avg_pool2d_tosa_FP_fp8(test_module):
     model, input_tensor, tosa_extension = test_module()
 
@@ -250,6 +251,20 @@ def test_adaptive_avg_pool2d_u85_INT_a16w8(test_module):
     pipeline.run()
 
 
+@common.parametrize("test_module", test_modules_fp8_no_comparing)
+@common.SkipIfNoModelConverter
+def test_adaptive_avg_pool2d_vgf_no_quant_fp8(test_module):
+    model, input_tensor, _ = test_module()
+    pipeline = VgfPipelineFP8[input_t](
+        model,
+        input_tensor,
+        aten_op=[],
+        exir_op=exir_op,
+        expect_missing_reference=True,
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_module", test_modules)
 @common.SkipIfNoModelConverter
 def test_adaptive_avg_pool2d_vgf_no_quant(test_module):

backends/arm/test/ops/test_avg_pool2d.py

@@ -20,6 +20,7 @@
     TosaPipelineFP,
     TosaPipelineINT,
     VgfPipeline,
+    VgfPipelineFP8,
 )
 
 aten_op = "avg_pool2d.default"
@@ -152,7 +153,7 @@ def forward(self, x: torch.Tensor):
         (torch.rand(1, 4, 12, 12, dtype=torch.float16),),
     ),
 }
-test_modules_fp8 = {
+test_modules_fp8_no_comparing = {
     "rand_fp8e4m3": lambda: (
         AvgPool2d(4, 2, 0, False),
         (torch.rand(1, 16, 50, 32).to(torch.float8_e4m3fn),),
@@ -180,7 +181,7 @@ def test_avg_pool2d_tosa_FP(test_module):
     pipeline.run()
 
 
-@common.parametrize("test_module", test_modules_fp8)
+@common.parametrize("test_module", test_modules_fp8_no_comparing)
 def test_avg_pool2d_tosa_FP_fp8(test_module):
     model, input_tensor, tosa_extension = test_module()
     pipeline = TosaPipelineFP[input_t](
@@ -288,6 +289,20 @@ def test_avg_pool2d_16a8w_u85_INT(test_module):
     pipeline.run()
 
 
+@common.parametrize("test_module", test_modules_fp8_no_comparing)
+@common.SkipIfNoModelConverter
+def test_avg_pool2d_vgf_no_quant_fp8(test_module):
+    model, input_tensor, _ = test_module()
+    pipeline = VgfPipelineFP8[input_t](
+        model,
+        input_tensor,
+        aten_op,
+        exir_op,
+        expect_missing_reference=True,
+    )
+    pipeline.run()
+
+
 @common.parametrize(
     "test_module",
     test_modules | test_modules_bf16 | test_modules_fp16,

backends/arm/test/ops/test_cat.py

@@ -201,6 +201,23 @@ def test_cat_u85_INT(test_data: Tuple):
     pipeline.run()
 
 
+@common.parametrize("test_data", Cat.test_parameters_fp8)
+@common.SkipIfNoModelConverter
+def test_cat_vgf_no_quant_fp8(test_data: Tuple):
+    tensors, dim, _ = test_data()
+    pipeline = VgfPipeline[input_t1](
+        Cat(),
+        (tensors, dim),
+        aten_op,
+        exir_op,
+        tosa_extensions=["fp8e4m3", "fp8e5m2"],
+        quantize=False,
+        atol=0,
+        rtol=0,
+    )
+    pipeline.run()
+
+
 @common.parametrize(
     "test_data",
     Cat.test_parameters | Cat.test_parameters_bf16 | Cat.test_parameters_fp16,

backends/arm/test/ops/test_constant_pad_nd.py

@@ -404,6 +404,23 @@ def test_constant_pad_nd_u55_INT():
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_suite_fp8)
+@common.SkipIfNoModelConverter
+def test_constant_pad_nd_vgf_no_quant_fp8(test_data: Tuple):
+    test_data, padding, value, mode, _ = test_data()
+    pipeline = VgfPipeline[input_t1](
+        ConstantPadND(padding, value, mode),
+        (test_data,),
+        aten_op,
+        exir_op,
+        tosa_extensions=["fp8e4m3", "fp8e5m2"],
+        quantize=False,
+        atol=0,
+        rtol=0,
+    )
+    pipeline.run()
+
+
 @common.parametrize(
     "test_data", test_data_suite | test_data_suite_bf16 | test_data_suite_fp16
 )

backends/arm/test/ops/test_conv2d.py

@@ -19,6 +19,7 @@
     TosaPipelineFP,
     TosaPipelineINT,
     VgfPipeline,
+    VgfPipelineFP8,
 )
 
 aten_op = "torch.ops.aten.conv2d.default"
@@ -522,7 +523,7 @@ def conv2d_fp16_1x1():
     "fp16_3x3": conv2d_fp16_3x3,
     "fp16_1x1": conv2d_fp16_1x1,
 }
-test_data_FP_fp8 = {
+test_data_FP_fp8_no_comparing = {
     "fp8e4m3": lambda: (
         Conv2d(
             height=8,
@@ -607,7 +608,7 @@ def test_convolution_2d_tosa_FP(test_data):
     pipeline.run()
 
 
-@common.parametrize("test_data", test_data_FP_fp8)
+@common.parametrize("test_data", test_data_FP_fp8_no_comparing)
 def test_convolution_2d_tosa_FP_fp8(test_data):
     model, tosa_extension = test_data()
     pipeline = TosaPipelineFP[input_t](
@@ -718,6 +719,20 @@ def test_convolution_2d_u85_INT_a8w4(test_data):
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_FP_fp8_no_comparing)
+@common.SkipIfNoModelConverter
+def test_convolution_2d_vgf_no_quant_fp8(test_data):
+    model, _ = test_data()
+    pipeline = VgfPipelineFP8[input_t](
+        model,
+        model.get_inputs(),
+        aten_op,
+        exir_op,
+        expect_missing_reference=True,
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", test_data_FP | test_data_FP_bf16 | test_data_FP_fp16)
 @common.SkipIfNoModelConverter
 def test_convolution_2d_vgf_no_quant(test_data):

backends/arm/test/ops/test_conv3d.py

@@ -20,6 +20,7 @@
     TosaPipelineFP,
     TosaPipelineINT,
     VgfPipeline,
+    VgfPipelineFP8,
 )
 
 aten_op = "torch.ops.aten.conv3d.default"
@@ -483,7 +484,7 @@ def forward(self, x):
     "5x5_3x2x24x24_st1": lambda: conv3d_5x5_3x2x24x24_st1,
     "3x3_1x3x28x28_st2_pd1": lambda: conv3d_3x3_1x3x28x28_st2_pd1,
 }
-test_data_FP_fp8 = {
+test_data_FP_fp8_no_comparing = {
     "basic_fp8e4m3": lambda: (
         Conv3d(
             height=6,
@@ -608,7 +609,7 @@ def test_convolution_3d_tosa_FP(test_data):
     pipeline.run()
 
 
-@common.parametrize("test_data", test_data_FP_fp8)
+@common.parametrize("test_data", test_data_FP_fp8_no_comparing)
 def test_convolution_3d_tosa_FP_fp8(test_data):
     model, tosa_extension = test_data()
     pipeline = TosaPipelineFP[input_t](
@@ -773,6 +774,20 @@ def test_convolution_3d_u85_INT_a8w4(test_data):
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_FP_fp8_no_comparing)
+@common.SkipIfNoModelConverter
+def test_convolution_3d_vgf_no_quant_fp8(test_data):
+    model, _ = test_data()
+    pipeline = VgfPipelineFP8[input_t](
+        model,
+        model.get_inputs(),
+        aten_op,
+        exir_op,
+        expect_missing_reference=True,
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", test_data_FP | test_data_FP_bf16 | test_data_FP_fp16)
 @common.SkipIfNoModelConverter
 def test_convolution_3d_vgf_no_quant(test_data):

backends/arm/test/ops/test_depthwise_conv.py

@@ -20,6 +20,7 @@
     TosaPipelineFP,
     TosaPipelineINT,
     VgfPipeline,
+    VgfPipelineFP8,
 )
 
 input_t = Tuple[torch.Tensor]  # Input x
@@ -195,7 +196,7 @@
         dtype=torch.float16,
     ),
 }
-test_data_conv2d_FP_fp8 = {
+test_data_conv2d_FP_fp8_no_comparing = {
     "fp8e4m3_3x3_gp3": lambda: (
         Conv2d(
             in_channels=3,
@@ -289,7 +290,7 @@ def test_convolution_2d_tosa_FP_depthwise(test_data: torch.nn.Module):
     pipeline.run()
 
 
-@common.parametrize("test_data", test_data_conv2d_FP_fp8)
+@common.parametrize("test_data", test_data_conv2d_FP_fp8_no_comparing)
 def test_convolution_2d_tosa_FP_fp8_depthwise(test_data):
     model, tosa_extension = test_data()
     pipeline = TosaPipelineFP[input_t](
@@ -340,6 +341,20 @@ def test_convolution_2d_tosa_INT_a8w4_depthwise(test_data):
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_conv2d_FP_fp8_no_comparing)
+@common.SkipIfNoModelConverter
+def test_convolution_2d_vgf_no_quant_fp8_depthwise(test_data):
+    model, _ = test_data()
+    pipeline = VgfPipelineFP8[input_t](
+        model,
+        model.get_inputs(),
+        aten_op=[],
+        exir_op=exir_op,
+        expect_missing_reference=True,
+    )
+    pipeline.run()
+
+
 @common.parametrize(
     "test_data",
     test_data_conv1d_FP

backends/arm/test/ops/test_embedding.py

@@ -141,6 +141,25 @@ def test_embedding_tosa_INT_expand():
     pipeline.run()
 
 
+@common.parametrize("test_input", test_input_fp8)
+@common.SkipIfNoModelConverter
+def test_embedding_vgf_no_quant_fp8(test_input):
+    inputs, _ = test_input
+    op = Embedding()
+    pipeline = VgfPipeline[input_params](
+        op,
+        inputs,
+        op.aten_op,
+        op.exir_op,
+        transform_passes=[InsertInt32CastsAfterInt64PlaceholdersPass()],
+        tosa_extensions=["fp8e4m3", "fp8e5m2"],
+        quantize=False,
+        atol=0,
+        rtol=0,
+    )
+    pipeline.run()
+
+
 @pytest.mark.skip("reason=MLETORCH-1274 Improve data type checks during partitioning")
 @common.parametrize("test_input", test_input)
 @common.SkipIfNoModelConverter

backends/arm/test/ops/test_flip.py

@@ -225,6 +225,23 @@ def test_flip_u85_INT(test_data):
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_suite_fp8)
+@common.SkipIfNoModelConverter
+def test_flip_vgf_no_quant_fp8(test_data):
+    data, dims, _ = test_data()
+    pipeline = VgfPipeline[input_t1](
+        Flip(dims),
+        (data,),
+        aten_op,
+        exir_op,
+        tosa_extensions=["fp8e4m3", "fp8e5m2"],
+        quantize=False,
+        atol=0,
+        rtol=0,
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", test_data_suite)
 @common.SkipIfNoModelConverter
 def test_flip_vgf_no_quant(test_data):

backends/arm/test/ops/test_gather.py

@@ -16,6 +16,7 @@
     TosaPipelineFP,
     TosaPipelineINT,
     VgfPipeline,
+    VgfPipelineFP8,
 )
 
 
@@ -87,7 +88,7 @@ def forward(self, input_: torch.Tensor, dim_, index_: torch.Tensor):
         ),  # Shape: [N=2, W=2, C=2]
     ),
 }
-test_data_fp_fp8: dict[str, tuple[input_params, str]] = {
+test_data_fp_fp8_no_comparing: dict[str, tuple[input_params, str]] = {
     "test_fp8e4m3_2d": (
         (
             torch.tensor(
@@ -175,7 +176,7 @@ def test_gather_tosa_FP(test_data: input_params):
     pipeline.run()
 
 
-@common.parametrize("test_data", test_data_fp_fp8)
+@common.parametrize("test_data", test_data_fp_fp8_no_comparing)
 def test_gather_tosa_FP_fp8(test_data: tuple[input_params, str]):
     input_data, tosa_extension = test_data
     pipeline = TosaPipelineFP[input_params](
@@ -231,6 +232,23 @@ def test_gather_u85_INT(test_data: input_params):
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_fp_fp8_no_comparing)
+@common.SkipIfNoModelConverter
+def test_gather_vgf_no_quant_fp8(test_data: tuple[input_params, str]):
+    input_data, _ = test_data
+    pipeline = VgfPipelineFP8[input_params](
+        Gather(),
+        input_data,
+        aten_op=Gather.aten_op,
+        exir_op=Gather.exir_op,
+        transform_passes=[InsertInt32CastsAfterInt64PlaceholdersPass()],
+        expect_missing_reference=True,
+        atol=0,
+        rtol=0,
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", test_data_fp | test_data_fp_bf16 | test_data_int)
 @common.SkipIfNoModelConverter
 def test_gather_vgf_no_quant(test_data: input_params):