Update on "[ET-VK] Replacing the use of uvec3 with WorkgroupSize class to reduce memory usage and improve processing speed"

This diff replaces the use of `uvec3` with `WorkgroupSize` class to reduce memory usage and improve processing speed in the Vulkan backend of Executorch.

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

[ghstack-poisoned]

Author: trivedivivek

.buckconfig

@@ -8,14 +8,15 @@
   root = .
   prelude = third-party/prelude
   shim = shim
+  shim_et = shim_et
 
 [repository_aliases]
   config = prelude
   ovr_config = prelude
-  toolchains = shim
-  fbcode = shim
-  fbcode_macros = shim
-  fbsource = shim
+  toolchains = shim_et
+  fbcode = shim_et
+  fbcode_macros = shim_et
+  fbsource = shim_et
   buck = shim
 
 [cxx]

backends/arm/test/misc/test_partition_decomposed_quantized_ops.py

@@ -60,6 +60,9 @@ def test_softplus_tosa_BI(test_data: input_t1):
     pipeline.pop_stage("check_not.exir")
     # check that all ops in exir_op except add are rejected
     pipeline.add_stage_after(
-        "partition", pipeline.tester.check, exir_op[1:], suffix="exir_post_partition"
+        "to_edge_transform_and_lower",
+        pipeline.tester.check,
+        exir_op[1:],
+        suffix="exir_post_partition",
     )
     pipeline.run()

backends/arm/test/ops/test_conv2d.py

@@ -370,7 +370,7 @@ def test_conv2d_tosa_BI(test_module):
     pipeline = TosaPipelineBI[input_t](
         test_module, test_module.get_inputs(), aten_op, exir_op
     )
-    pipeline.change_args("run_method_and_compare_outputs.0", qtol=1)
+    pipeline.change_args("run_method_and_compare_outputs", qtol=1)
     pipeline.run()
 
 

backends/arm/test/passes/test_cast_int64_pass.py

@@ -8,7 +8,7 @@
 import torch
 from executorch.backends.arm._passes.cast_int64_pass import CastInt64ToInt32Pass
 
-from executorch.backends.arm.test.tester.test_pipeline import TestPassPipeline
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
 
 input_t = Tuple[torch.Tensor]  # Input x
 
@@ -28,7 +28,7 @@ def test_int64_model_tosa_BI():
         "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1,
         "executorch_exir_dialects_edge__ops_aten_add_Tensor": 1,
     }
-    pipeline = TestPassPipeline[input_t](
+    pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
         tosa_version="TOSA-0.80+BI",

backends/arm/test/passes/test_fold_qdq_pass.py

@@ -9,7 +9,7 @@
 from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import (
     FoldAndAnnotateQParamsPass,
 )
-from executorch.backends.arm.test.tester.test_pipeline import TestPassPipeline
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
 
 
 input_t = Tuple[torch.Tensor, torch.Tensor]  # Input x, y
@@ -32,7 +32,7 @@ def test_fold_qdq_pass_tosa_BI():
     is removed from the representation.
     """
     module = SimpleQuantizeModel()
-    pipeline = TestPassPipeline[input_t](
+    pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
         tosa_version="TOSA-0.80+BI",

backends/arm/test/passes/test_fuse_batchnorm_pass.py

@@ -8,7 +8,7 @@
 import torch
 from executorch.backends.arm._passes.fuse_batchnorm2d_pass import FuseBatchnorm2DPass
 from executorch.backends.arm.test import common
-from executorch.backends.arm.test.tester.test_pipeline import TestPassPipeline
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
 
 input_t = Tuple[torch.Tensor]  # Input x
 
@@ -138,7 +138,7 @@ def forward(self, x):
 @common.parametrize("module", modules)
 def test_fuse_batchnorm_tosa_MI(module):
     """Test various cases where the batchnorm should and shouldn't be fused."""
-    pipeline = TestPassPipeline[input_t](
+    pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
         tosa_version="TOSA-0.80+MI",

backends/arm/test/passes/test_insert_table_ops_pass.py

@@ -11,7 +11,7 @@
     FoldAndAnnotateQParamsPass,
 )
 from executorch.backends.arm._passes.insert_table_ops import InsertTableOpsPass
-from executorch.backends.arm.test.tester.test_pipeline import TestPassPipeline
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
 
 input_t = Tuple[torch.Tensor]  # Input x
 
@@ -27,7 +27,7 @@ def get_inputs(self) -> input_t:
 
 def test_insert_table_tosa_BI():
     module = Sigmoid()
-    pipeline = TestPassPipeline[input_t](
+    pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
         tosa_version="TOSA-0.80+BI",

backends/arm/test/passes/test_meandim_to_averagepool2d.py

@@ -11,7 +11,7 @@
     ConvertMeanDimToAveragePoolPass,
 )
 from executorch.backends.arm.test import common
-from executorch.backends.arm.test.tester.test_pipeline import TestPassPipeline
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
 
 
 input_t = Tuple[torch.Tensor, torch.Tensor]  # Input x
@@ -65,7 +65,7 @@ def test_meandim_to_avgpool_tosa_BI(module):
     Tests the MeanDimToAveragePool2dPass which converts mean.dim to average_pool2d
     for the special case where dim is [-1, -2] and keepdim is True.
     """
-    pipeline = TestPassPipeline[input_t](
+    pipeline = PassPipeline[input_t](
         module,
         module.get_inputs(),
         tosa_version="TOSA-0.80+BI",

backends/arm/test/passes/test_unsqueeze_before_repeat_pass.py

@@ -10,7 +10,7 @@
     UnsqueezeBeforeRepeatPass,
 )
 from executorch.backends.arm.test import common
-from executorch.backends.arm.test.tester.test_pipeline import TestPassPipeline
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
 
 input_t = Tuple[
     torch.Tensor, Dict[str, int], list[str]
@@ -47,7 +47,7 @@ def test_unsqueeze_before_repeat_tosa_MI(test_data):
     """
     module = Repeat()
     data, ops_after_pass, ops_not_after_pass = test_data
-    pipeline = TestPassPipeline(
+    pipeline = PassPipeline(
         module,
         data,
         tosa_version="TOSA-0.80+MI",