Update on "[ET-VK] New implementation of cat operator"

## Changes

* Introduce `concat_texture.glsl` and `concat_buffer.glsl` to implement the `torch.cat` operator
* Introduce `Concat.cpp` to replace `Cat.cpp`
* Fix a bug with channels-packed buffer tensors where input data would be copied incorrectly with multiple dims have a stride of 1

## Motivation

> * Introduce `concat_texture.glsl` and `concat_buffer.glsl` to implement the `torch.cat` operator
> * Introduce `Concat.cpp` to replace `Cat.cpp`

The existing implementation of `torch.cat` uses the copy_channel_offset` shaders. However, these shaders have a critical bug where the output tensor is passed in separately with difference access types, i.e.

```
graph.execute_nodes().emplace_back(new DispatchNode(
        graph,
        VK_KERNEL_FROM_STR(kernel_name),
        global_size,
        local_size,
        // Inputs and Outputs
        {
            {out, vkapi::kWrite},
            {out, vkapi::kRead},
            {in, vkapi::kRead},
        },
```

This creates many validation layer errors because the memory barriers for the resource cannot be formed properly. The shader essentially relies on undefined behaviour to work correctly. The result is that the `cat` operator produces incorrect result on many platforms.

Rather than fix the `copy_offset` shaders, I decided to just introduce new shaders to perform the concat operation. The new implementation handles both buffer and texture inputs and is agnostic to memory layout.

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

[ghstack-poisoned]

Author: SS-JIA

.ci/scripts/test_llama.sh

@@ -156,8 +156,7 @@ cmake_install_executorch_libraries() {
         -DCMAKE_INSTALL_PREFIX=cmake-out \
         -DCMAKE_BUILD_TYPE="$CMAKE_BUILD_TYPE" \
         -DEXECUTORCH_BUILD_QNN="$QNN" \
-        -DQNN_SDK_ROOT="$QNN_SDK_ROOT" \
-        -Bcmake-out .
+        -DQNN_SDK_ROOT="$QNN_SDK_ROOT"
     cmake --build cmake-out -j9 --target install --config "$CMAKE_BUILD_TYPE"
 }
 

.ci/scripts/unittest-buck2.sh

@@ -15,7 +15,7 @@ buck2 query "//backends/apple/... + //backends/example/... + \
 //kernels/optimized/... + //kernels/portable/... + //kernels/quantized/... + \
 //kernels/test/... + //runtime/... + //schema/... + //test/... + //util/..."
 
-UNBUILDABLE_OPTIMIZED_OPS_REGEX="gelu|fft_r2c|log_softmax"
+UNBUILDABLE_OPTIMIZED_OPS_REGEX="_elu|gelu|fft|log_softmax"
 BUILDABLE_OPTIMIZED_OPS=$(buck2 query //kernels/optimized/cpu/... | grep -E -v $UNBUILDABLE_OPTIMIZED_OPS_REGEX)
 
 # TODO: build prim_ops_test_cpp again once supported_features works in
@@ -24,6 +24,10 @@ BUILDABLE_KERNELS_PRIM_OPS_TARGETS=$(buck2 query //kernels/prim_ops/... | grep -
 # TODO: expand the covered scope of Buck targets.
 # //runtime/kernel/... is failing because //third-party:torchgen_files's shell script can't find python on PATH.
 # //runtime/test/... requires Python torch, which we don't have in our OSS buck setup.
-buck2 test $BUILDABLE_OPTIMIZED_OPS //kernels/portable/... \
-      $BUILDABLE_KERNELS_PRIM_OPS_TARGETS //runtime/backend/... //runtime/core/... \
-      //runtime/executor: //runtime/kernel/... //runtime/platform/...
+for op in "build" "test"; do
+    buck2 $op $BUILDABLE_OPTIMIZED_OPS \
+          //examples/selective_build:select_all_dtype_selective_lib_portable_lib \
+          //kernels/portable/... \
+          $BUILDABLE_KERNELS_PRIM_OPS_TARGETS //runtime/backend/... //runtime/core/... \
+          //runtime/executor: //runtime/kernel/... //runtime/platform/...
+done

.github/workflows/trunk.yml

@@ -262,7 +262,7 @@ jobs:
         output=$(ls -la ${elf})
         arr=($output)
         size=${arr[4]}
-        threshold="103068" # ~100KiB
+        threshold="103268" # ~100KiB
         echo "size: $size, threshold: $threshold"
         if [[ "$size" -le "$threshold" ]]; then
           echo "Success $size <= $threshold"

CMakeLists.txt

@@ -430,14 +430,6 @@ endif()
 
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/configurations)
 
-#
-# gflags: Commandline flag host library.
-#
-
-if(EXECUTORCH_BUILD_GFLAGS)
-  add_subdirectory(third-party/gflags)
-endif()
-
 # Install `executorch` library as well as `executorch-config.cmake` under
 # ${CMAKE_INSTALL_PREFIX}/
 install(
@@ -522,17 +514,6 @@ if(EXECUTORCH_BUILD_CORTEX_M)
 endif()
 
 if(EXECUTORCH_BUILD_DEVTOOLS)
-  if(NOT EXECUTORCH_BUILD_ARM_BAREMETAL)
-    set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER
-        ON
-        CACHE BOOL "EXECUTORCH_BUILD_EXTENSION_DATA_LOADER" FORCE
-    )
-  else()
-    set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER
-        OFF
-        CACHE BOOL "EXECUTORCH_BUILD_EXTENSION_DATA_LOADER" FORCE
-    )
-  endif()
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/devtools)
 endif()
 
@@ -573,6 +554,10 @@ if(EXECUTORCH_BUILD_PTHREADPOOL AND EXECUTORCH_BUILD_CPUINFO)
 endif()
 
 if(EXECUTORCH_BUILD_PYBIND)
+
+  # Add codegen tools subdirectory for selective_build pybind module
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/codegen/tools)
+
   if(NOT EXECUTORCH_BUILD_EXTENSION_DATA_LOADER)
     add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/data_loader)
   endif()

backends/apple/coreml/scripts/build_tests.sh

@@ -33,8 +33,7 @@ cmake "$EXECUTORCH_ROOT_PATH" -B"$CMAKE_EXECUTORCH_BUILD_DIR_PATH" \
 -DPLATFORM=MAC_UNIVERSAL \
 -DDEPLOYMENT_TARGET=13.0 \
 -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=OFF \
--DEXECUTORCH_BUILD_XNNPACK=OFF \
--DEXECUTORCH_BUILD_GFLAGS=OFF
+-DEXECUTORCH_BUILD_XNNPACK=OFF
 
 cmake --build "$CMAKE_EXECUTORCH_BUILD_DIR_PATH"  -j9 -t executorch
 

backends/arm/_passes/arm_pass_manager.py

@@ -62,7 +62,10 @@
     UnsqueezeScalarPlaceholdersPass,
 )
 
-from executorch.backends.arm.tosa_specification import TosaSpecification
+from executorch.backends.arm.tosa_specification import (
+    TosaLoweringContext,
+    TosaSpecification,
+)
 from executorch.backends.transforms.decompose_sdpa import (
     DecomposeScaledDotProductAttention,
 )
@@ -80,7 +83,8 @@ def __init__(self, tosa_spec: TosaSpecification) -> None:
         super().__init__()
 
     def _transform(self, graph_module: GraphModule):
-        return self(graph_module).graph_module
+        with TosaLoweringContext(self.tosa_spec):
+            return self(graph_module).graph_module
 
     def _tosa_080_BI_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(FuseQuantizedActivationPass())

backends/arm/_passes/scalars_to_attribute_pass.py

@@ -12,8 +12,8 @@
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 
 from executorch.exir.pass_base import ExportPass, PassResult
-from torch.ao.quantization.fx.utils import get_new_attr_name_with_prefix
 from torch.fx import GraphModule, Node
+from torchao.quantization.pt2e.utils import get_new_attr_name_with_prefix
 
 
 class ScalarsToAttributePass(ExportPass):

backends/arm/quantizer/arm_quantizer.py

@@ -247,9 +247,9 @@ def set_module_name(
         quantizer.set_module_name("blocks.sub"), it will quantize all supported operator/operator
         patterns in the submodule with this module name with the given `quantization_config`
         """
-        assert (
-            quantization_config is not None
-        ), " quantization_config == None is not supported yet"
+        # Validate that quantization_config is provided
+        if quantization_config is None:
+            raise ValueError("quantization_config == None is not supported yet")
         self.module_name_config[module_name] = quantization_config
         return self
 

backends/arm/quantizer/quantization_annotator.py

@@ -221,10 +221,12 @@ def _match_pattern(
     torch.ops.aten.squeeze_copy.dim,
     torch.ops.aten.squeeze.dim,
     torch.ops.aten.squeeze.dims,
+    torch.ops.aten.unbind.int,
     torch.ops.aten.unsqueeze.default,
     torch.ops.aten.unsqueeze_copy.default,
     torch.ops.aten.reshape.default,
     torch.ops.aten.repeat.default,
+    torch.ops.aten.repeat_interleave.self_int,
     torch.ops.aten.expand_copy.default,
     torch.ops.aten.expand.default,
     # Disabling these as there seems to be an issue with support for complex
@@ -256,6 +258,7 @@ def _match_pattern(
     torch.ops.aten.amin.default,
     torch.ops.aten.clamp.default,
     torch.ops.aten.clamp.Tensor,
+    torch.ops.aten.unflatten.int,
 ]
 
 _one_to_one_shared_input_or_input_act_qspec = [
@@ -271,6 +274,7 @@ def _match_pattern(
     torch.ops.aten.avg_pool2d.default,
     torch.ops.aten.max_pool2d.default,
     torch.ops.aten.full.default,
+    torch.ops.aten.full,
     torch.ops.aten.flatten.using_ints,
     torch.ops.aten.dropout.default,
     torch.ops.aten.dropout_.default,
@@ -539,6 +543,7 @@ def annotate_graph(  # type: ignore[return]
         if node.target in [
             torch.ops.aten.full_like.default,
             torch.ops.aten.full.default,
+            torch.ops.aten.full,
             torch.ops.aten.scalar_tensor.default,
         ]:
             node.kwargs = {}

backends/arm/quantizer/quantization_config.py

@@ -29,30 +29,40 @@ def get_input_act_qspec(self) -> QuantizationSpec | None:
         """Returns QuantizationSpec 'input_activation' after asserting that input_activation.qscheme is valid."""
         if self.input_activation is None:
             return None
-        assert self.input_activation.qscheme in [
+        # Validate that input_activation uses a supported qscheme
+        if self.input_activation.qscheme not in [
             torch.per_tensor_affine,
             torch.per_tensor_symmetric,
-        ], f"Unsupported quantization_spec {self.input_activation} for input_activation."
+        ]:
+            raise ValueError(
+                f"Unsupported quantization_spec {self.input_activation} for input_activation."
+            )
         return self.input_activation
 
     def get_output_act_qspec(self) -> QuantizationSpec | None:
         """Returns QuantizationSpec 'output_activation' after asserting that output_activation.qscheme is valid."""
         if self.output_activation is None:
             return None
-        assert self.output_activation.qscheme in [
+        # Validate that output_activation uses a supported qscheme
+        if self.output_activation.qscheme not in [
             torch.per_tensor_affine,
             torch.per_tensor_symmetric,
-        ], f"Unsupported quantization_spec {self.output_activation} for output_activation."
+        ]:
+            raise ValueError(
+                f"Unsupported quantization_spec {self.output_activation} for output_activation."
+            )
         return self.output_activation
 
     def get_weight_qspec(self) -> QuantizationSpec | None:
         """Returns QuantizationSpec 'weight' after asserting that weight.qscheme is valid."""
         if self.weight is None:
             return None
-        assert self.weight.qscheme in [
+        # Validate that weight uses a supported qscheme
+        if self.weight.qscheme not in [
             torch.per_tensor_symmetric,
             torch.per_channel_symmetric,
-        ], f"Unsupported quantization_spec {self.weight} for weight"
+        ]:
+            raise ValueError(f"Unsupported quantization_spec {self.weight} for weight")
         return self.weight
 
     def get_bias_qspec(self, node: torch.fx.Node) -> QuantizationSpec | None:
@@ -61,11 +71,11 @@ def get_bias_qspec(self, node: torch.fx.Node) -> QuantizationSpec | None:
         def _derive_qparams_fn(
             obs_or_fqs: list[ObserverOrFakeQuantize],
         ) -> tuple[torch.Tensor, torch.Tensor]:
-            assert (
-                len(obs_or_fqs) == 2
-            ), "Expecting two obs/fqs, one for activation and one for weight, got: {}".format(
-                len(obs_or_fqs)
-            )
+            # Validate expected number of observers/fake-quantizes
+            if len(obs_or_fqs) != 2:
+                raise ValueError(
+                    f"Expecting two obs/fqs, one for activation and one for weight, got: {len(obs_or_fqs)}"
+                )
             act_obs_or_fq = obs_or_fqs[0]
             weight_obs_or_fq = obs_or_fqs[1]
             act_scale, act_zp = act_obs_or_fq.calculate_qparams()
@@ -94,9 +104,11 @@ def _derive_qparams_fn(
 
         if self.bias is None:
             return None
-        assert (
-            self.bias.dtype == torch.float
-        ), "Only float dtype for bias is supported for bias right now"
+        # Validate that bias dtype is floating-point
+        if self.bias.dtype != torch.float:
+            raise ValueError(
+                "Only float dtype for bias is supported for bias right now"
+            )
         return self.bias
 
     def get_fixed_qspec(