Update on "[ET-VK] 6/n Split dispatches between multiple command buffers. Repurpose submit_current_cmd_and_wait to wait based on input flag."

This diff makes changes to the `submit_current_cmd_and_wait` function in the `ComputeGraph` class to repurpose it to wait for command buffer completion based on an input flag. The function is renamed to `submit_current_cmd` and now takes an additional `wait` parameter. If `wait` is `true`, the function submits the command buffer to the GPU and waits for its completion. Otherwise, it only submits the command buffer without waiting.

Differential Revision: [D78360042](https://our.internmc.facebook.com/intern/diff/D78360042/)

[ghstack-poisoned]

Author: trivedivivek

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-7cda4017ddda554752e89069ae205be5e8388f59
+90f1e7bed15ca5e48c61c5b6dc5ad4810524f82f

.github/workflows/trunk.yml

@@ -197,7 +197,7 @@ jobs:
       docker-image: executorch-ubuntu-22.04-arm-sdk
       submodules: 'recursive'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-      timeout: 90
+      timeout: 120
       script: |
         # The generic Linux job chooses to use base env, not the one setup by the image
         CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")

backends/arm/quantizer/arm_quantizer.py

@@ -60,7 +60,7 @@
 
 @functools.lru_cache
 def get_symmetric_quantization_config(
-    is_per_channel: bool = False,
+    is_per_channel: bool = True,
     is_qat: bool = False,
     is_dynamic: bool = False,
     act_qmin: int = -128,

backends/arm/test/common.py

@@ -18,6 +18,7 @@
     arm_executor_runner_exists,
     corstone300_installed,
     corstone320_installed,
+    model_converter_installed,
 )
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from executorch.exir.backend.compile_spec_schema import CompileSpec
@@ -245,6 +246,13 @@ def get_u85_compile_spec_unbuilt(
 )
 """Xfails a test if Corsone320 FVP is not installed, or if the executor runner is not built"""
 
+SkipIfNoModelConverter = pytest.mark.skipif(
+    condition=not (model_converter_installed()),
+    raises=FileNotFoundError,
+    reason="Did not find model-converter on path",
+)
+"""Xfails a test if model-converter is not installed"""
+
 xfail_type = str | tuple[str, type[Exception]]
 
 

backends/arm/test/misc/test_bn_relu_folding_qat.py

@@ -59,7 +59,9 @@ def test_qat_tosa_BI(model: torch.nn.Module):
         "quantize",
         Quantize(
             quantizer=quantizer,
-            quantization_config=get_symmetric_quantization_config(is_qat=True),
+            quantization_config=get_symmetric_quantization_config(
+                is_qat=True, is_per_channel=False
+            ),
             is_qat=True,
         ),
     )

backends/arm/test/ops/test_add.py

@@ -7,8 +7,6 @@
 
 from typing import Tuple
 
-import pytest
-
 import torch
 from executorch.backends.arm.arm_backend import get_tosa_spec
 from executorch.backends.arm.quantizer import arm_quantizer
@@ -190,7 +188,7 @@ def test_add_tensor_u85_BI_2(test_data: input_t2):
 
 
 @common.parametrize("test_data", Add.test_data)
-@pytest.mark.skip(reason="Model converter not yet made available")
+@common.SkipIfNoModelConverter
 def test_add_tensor_vgf_fp(test_data: input_t1):
     pipeline = VgfPipeline[input_t1](
         Add(), test_data(), aten_op, exir_op, tosa_version="TOSA-1.0+FP"
@@ -199,7 +197,7 @@ def test_add_tensor_vgf_fp(test_data: input_t1):
 
 
 @common.parametrize("test_data", Add.test_data)
-@pytest.mark.skip(reason="Model converter not yet made available")
+@common.SkipIfNoModelConverter
 def test_add_tensor_vgf_int(test_data: input_t1):
     pipeline = VgfPipeline[input_t1](
         Add(),

backends/arm/test/ops/test_multihead_attention.py

@@ -53,7 +53,14 @@ def test_multihead_attention_tosa_MI(test_data: input_t1):
 )
 def test_multihead_attention_tosa_BI(test_data):
     test_data, module = test_data()
-    pipeline = TosaPipelineBI(module, (*test_data, *test_data, *test_data), [], [])
+    pipeline = TosaPipelineBI(
+        module,
+        (*test_data, *test_data, *test_data),
+        [],
+        [],
+        # TODO: Per-channel quantization is broken (MLETORCH-1144)
+        per_channel_quantization=False,
+    )
     pipeline.run()
 
 
@@ -72,6 +79,8 @@ def test_multihead_attention_u55_BI(test_data: input_t1):
         [],
         use_to_edge_transform_and_lower=True,
         run_on_fvp=True,
+        # TODO: Per-channel quantization is broken (MLETORCH-1144)
+        per_channel_quantization=False,
     )
     pipeline.pop_stage("check_count.exir")
     pipeline.run()
@@ -92,5 +101,7 @@ def test_multihead_attention_u85_BI(test_data: input_t1):
         [],
         use_to_edge_transform_and_lower=True,
         run_on_fvp=True,
+        # TODO: Per-channel quantization is broken (MLETORCH-1144)
+        per_channel_quantization=False,
     )
     pipeline.run()

backends/arm/test/runner_utils.py

@@ -549,6 +549,15 @@ def corstone320_installed() -> bool:
     return True
 
 
+def model_converter_installed() -> bool:
+    cmd = ["model-converter", "--version"]
+    try:
+        _run_cmd(cmd, check=True)
+    except:
+        return False
+    return True
+
+
 def get_elf_path(target_board):
     elf_path = os.path.join(
         "arm_test",

backends/arm/test/test_arm_baremetal.sh

@@ -210,7 +210,9 @@ test_models_ethos-u55() { # End to End model tests using model_test.py
     python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u55-64  --model=mv3  --extra_flags="-DET_ATOL=5.00 -DET_RTOL=5.00"
     python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u55-256 --model=lstm --extra_flags="-DET_ATOL=0.03 -DET_RTOL=0.03"
     python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u55-128 --model=resnet18 --extra_flags="-DET_ATOL=0.2 -DET_RTOL=0.2"
-    python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u55-128 --model=resnet50 --extra_flags="-DET_ATOL=0.2 -DET_RTOL=0.2"
+    # TODO: Output performance for resnet50 is bad with per-channel quantization (MLETORCH-1149).
+    # Also we get OOM when running this model. Disable it for now.
+    #python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u55-128 --model=resnet50 --extra_flags="-DET_ATOL=6.2 -DET_RTOL=6.2"
 
     echo "${TEST_SUITE_NAME}: PASS"
     }

backends/arm/test/tester/test_pipeline.py

@@ -300,7 +300,7 @@ def __init__(
         run_on_tosa_ref_model: bool = True,
         tosa_version: str = "TOSA-0.80+BI",
         symmetric_io_quantization: bool = False,
-        per_channel_quantization: bool = False,
+        per_channel_quantization: bool = True,
         use_to_edge_transform_and_lower: bool = True,
         custom_path: str = None,
         atol: float = 1e-03,
@@ -317,16 +317,14 @@ def __init__(
         compile_spec = common.get_tosa_compile_spec(
             tosa_profiles[tosa_version], custom_path=custom_path
         )
-        if symmetric_io_quantization or per_channel_quantization:
-            quantizer = TOSAQuantizer(tosa_profiles[tosa_version])
-            quantization_config = get_symmetric_quantization_config(
-                is_per_channel=per_channel_quantization
-            )
-            if symmetric_io_quantization:
-                quantizer.set_io(quantization_config)
-            quant_stage = Quantize(quantizer, quantization_config)
-        else:
-            quant_stage = None
+
+        quantizer = TOSAQuantizer(tosa_profiles[tosa_version])
+        quantization_config = get_symmetric_quantization_config(
+            is_per_channel=per_channel_quantization
+        )
+        if symmetric_io_quantization:
+            quantizer.set_io(quantization_config)
+        quant_stage = Quantize(quantizer, quantization_config)
 
         super().__init__(
             module,
@@ -475,24 +473,21 @@ def __init__(
         exir_ops: Optional[str | List[str]] = None,
         run_on_fvp: bool = True,
         symmetric_io_quantization: bool = False,
-        per_channel_quantization: bool = False,
+        per_channel_quantization: bool = True,
         use_to_edge_transform_and_lower: bool = True,
         custom_path: str = None,
         atol: float = 1e-03,
         rtol: float = 1e-03,
         qtol: int = 1,
     ):
         compile_spec = common.get_u55_compile_spec(custom_path=custom_path)
-        if symmetric_io_quantization or per_channel_quantization:
-            quantizer = EthosUQuantizer(compile_spec)
-            quantization_config = get_symmetric_quantization_config(
-                is_per_channel=per_channel_quantization
-            )
-            if symmetric_io_quantization:
-                quantizer.set_io(quantization_config)
-            quant_stage = Quantize(quantizer, quantization_config)
-        else:
-            quant_stage = None
+        quantizer = EthosUQuantizer(compile_spec)
+        quantization_config = get_symmetric_quantization_config(
+            is_per_channel=per_channel_quantization
+        )
+        if symmetric_io_quantization:
+            quantizer.set_io(quantization_config)
+        quant_stage = Quantize(quantizer, quantization_config)
 
         super().__init__(
             module,
@@ -565,24 +560,21 @@ def __init__(
         exir_ops: str | List[str] = None,
         run_on_fvp: bool = True,
         symmetric_io_quantization: bool = False,
-        per_channel_quantization: bool = False,
+        per_channel_quantization: bool = True,
         use_to_edge_transform_and_lower: bool = True,
         custom_path: str = None,
         atol: float = 1e-03,
         rtol: float = 1e-03,
         qtol: int = 1,
     ):
         compile_spec = common.get_u85_compile_spec(custom_path=custom_path)
-        if symmetric_io_quantization or per_channel_quantization:
-            quantizer = EthosUQuantizer(compile_spec)
-            quantization_config = get_symmetric_quantization_config(
-                is_per_channel=per_channel_quantization
-            )
-            if symmetric_io_quantization:
-                quantizer.set_io(quantization_config)
-            quant_stage = Quantize(quantizer, quantization_config)
-        else:
-            quant_stage = None
+        quantizer = EthosUQuantizer(compile_spec)
+        quantization_config = get_symmetric_quantization_config(
+            is_per_channel=per_channel_quantization
+        )
+        if symmetric_io_quantization:
+            quantizer.set_io(quantization_config)
+        quant_stage = Quantize(quantizer, quantization_config)
 
         super().__init__(
             module,