Update on "[ET-VK] 5/n Split dispatches between multiple command buffers. Track previously submitted command buffers in context and add function to execute all previous command buffers."

The diff adds changes to store command buffers submitted with final_use set to false.

Storing these buffers is necessary for `execute()` function. Since, `encode_execute()` function is typically called once but `execute()` can be called multiple times, `submit_all_non_final_cmds` function is added so all recorded command buffers with `final_use = False` can be called multiple times in `execute()`.

#### Key Changes
*   Added a flag `execute_pending_first_submission` to the `ComputeGraph` class to track whether execute nodes have been freshly encoded and need to be submitted first.
*   Added a new function `submit_all_non_final_cmds` to the `Context` class, which submits all non-final command buffers to the GPU.
*   Modified the `submit_cmd_to_gpu` function to add the submitted command buffer to the `non_final_cmds_` list if it's not marked as final use.
*   Updated the `execute` function in `ComputeGraph` to submit all non-final command buffers before executing the graph.

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

[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,