Update on "[Executorch][LLM] Use caching allocator for runner"

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

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

[ghstack-poisoned]

Author: kimishpatel

backends/arm/operators/op_index_select.py

@@ -19,7 +19,7 @@
 )
 from executorch.backends.arm.tosa.mapping import TosaArg
 
-from executorch.backends.arm.tosa.utils import build_reshape_tosa_1_0
+from executorch.backends.arm.tosa.utils import build_reshape_tosa
 from torch.fx import Node
 
 
@@ -67,7 +67,7 @@ def define_node(
                 weights_new_shape,
                 weights.dtype,
             )
-            build_reshape_tosa_1_0(
+            build_reshape_tosa(
                 tosa_graph, weights.name, weights_new_shape, weights_reshaped.name
             )
 
@@ -89,7 +89,7 @@ def define_node(
             indices_new_shape,
             indices.dtype,
         )
-        build_reshape_tosa_1_0(
+        build_reshape_tosa(
             tosa_graph, indices.name, indices_new_shape, indices_reshaped.name
         )
 
@@ -106,6 +106,4 @@ def define_node(
 
         if len(weights.shape) == 2:
             output_real_shape = [output.shape[0], output.shape[1]]
-            build_reshape_tosa_1_0(
-                tosa_graph, output_name, output_real_shape, output.name
-            )
+            build_reshape_tosa(tosa_graph, output_name, output_real_shape, output.name)

backends/arm/operators/op_index_tensor.py

@@ -180,7 +180,7 @@ def define_node(
                 gather_idx_shape,
                 index_dtype,
             )
-            tutils.build_reshape_tosa_1_0(
+            tutils.build_reshape_tosa(
                 tosa_graph,
                 stride_shifted_indices.name,
                 gather_idx_shape,
@@ -212,7 +212,7 @@ def define_node(
         gather_vals_shape = [N, K, C]
         reshaped_input = tosa_graph.addIntermediate(gather_vals_shape, values.dtype)
 
-        tutils.build_reshape_tosa_1_0(
+        tutils.build_reshape_tosa(
             tosa_graph,
             values.name,
             gather_vals_shape,
@@ -238,7 +238,7 @@ def define_node(
 
         output_shape = tutils.tosa_shape(output.shape, output.dim_order)
 
-        tutils.build_reshape_tosa_1_0(
+        tutils.build_reshape_tosa(
             tosa_graph,
             gather_out.name,
             list(output_shape),

backends/arm/test/ops/test_tanh.py

@@ -8,21 +8,15 @@
 
 import pytest
 import torch
-from executorch.backends.arm.quantizer.arm_quantizer import (
-    get_symmetric_a16w8_quantization_config,
-    TOSAQuantizer,
-)
 
-from executorch.backends.arm.test import common, conftest
+from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineINT,
     EthosU85PipelineINT,
     TosaPipelineFP,
     TosaPipelineINT,
     VgfPipeline,
 )
-from executorch.backends.arm.tosa.specification import TosaSpecification
-from executorch.backends.xnnpack.test.tester import Quantize
 
 aten_op = "torch.ops.aten.tanh.default"
 input_t1 = Tuple[torch.Tensor]  # Input x
@@ -114,29 +108,6 @@ def test_tanh_vgf_INT(test_data: Tuple):
     pipeline.run()
 
 
-def get_symmetric_a16w8_tanh_quantizer(per_channel_quantization=False):
-    tosa_version = conftest.get_option("tosa_version")
-    tosa_profiles = {
-        "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT+int16"),
-    }
-
-    quantizer = TOSAQuantizer(tosa_profiles[tosa_version])
-
-    # Use a smaller episilon value to not greatly inflate [qmin, qmax]
-    quantizer.set_global(
-        get_symmetric_a16w8_quantization_config(
-            is_per_channel=per_channel_quantization, epsilon=2**-16
-        )
-    )
-
-    return Quantize(
-        quantizer,
-        get_symmetric_a16w8_quantization_config(
-            is_per_channel=per_channel_quantization, epsilon=2**-16
-        ),
-    )
-
-
 @common.parametrize("test_data", test_data_suite)
 def test_tanh_16a8w_tosa_INT(test_data: torch.Tensor):
     """Test tanh operation with 16A8W quantization (16-bit activations, 8-bit weights)"""
@@ -150,13 +121,8 @@ def test_tanh_16a8w_tosa_INT(test_data: torch.Tensor):
         per_channel_quantization=per_channel_quantization,
         use_to_edge_transform_and_lower=True,
         tosa_extensions=["int16"],
-    )
-
-    pipeline.change_args(
-        "quantize",
-        get_symmetric_a16w8_tanh_quantizer(
-            per_channel_quantization=per_channel_quantization
-        ),
+        epsilon=2**-16,
+        rtol=2e-03,
     )
     pipeline.run()
 
@@ -177,13 +143,9 @@ def test_tanh_16a8w_u55_INT16(test_data: torch.Tensor):
         exir_ops=[],
         per_channel_quantization=per_channel_quantization,
         use_to_edge_transform_and_lower=True,
-    )
-
-    pipeline.change_args(
-        "quantize",
-        get_symmetric_a16w8_tanh_quantizer(
-            per_channel_quantization=per_channel_quantization
-        ),
+        a16w8_quantization=True,
+        epsilon=2**-16,
+        rtol=2e-03,
     )
     pipeline.run()
 
@@ -201,12 +163,8 @@ def test_tanh_16a8w_u85_INT16(test_data: torch.Tensor):
         exir_ops=[],
         per_channel_quantization=per_channel_quantization,
         use_to_edge_transform_and_lower=True,
-    )
-
-    pipeline.change_args(
-        "quantize",
-        get_symmetric_a16w8_tanh_quantizer(
-            per_channel_quantization=per_channel_quantization
-        ),
+        a16w8_quantization=True,
+        epsilon=2**-16,
+        rtol=2e-03,
     )
     pipeline.run()

backends/arm/test/tester/test_pipeline.py

@@ -376,7 +376,7 @@ def __init__(
         qtol: int = 1,
         dynamic_shapes: Optional[Tuple[Any]] = None,
         tosa_extensions: Optional[List[str]] = None,
-        epsilon: float = 2**12,
+        epsilon: float = 2**-12,
     ):
         if tosa_extensions is None:
             tosa_extensions = []
@@ -570,7 +570,7 @@ def __init__(
         atol: float = 1e-03,
         rtol: float = 1e-03,
         qtol: int = 1,
-        epsilon: float = 2**12,
+        epsilon: float = 2**-12,
     ):
         compile_spec = common.get_u55_compile_spec(
             custom_path=custom_path,
@@ -671,7 +671,7 @@ def __init__(
         atol: float = 1e-03,
         rtol: float = 1e-03,
         qtol: int = 1,
-        epsilon: float = 2**12,
+        epsilon: float = 2**-12,
     ):
         compile_spec = common.get_u85_compile_spec(
             custom_path=custom_path,

backends/arm/tosa/utils.py

@@ -110,7 +110,7 @@ def broadcast_tensors(
             tens_dtype,
         )
 
-        build_reshape_tosa_1_0(tosa_fb, node.name, new_shape, reshaped.name)
+        build_reshape_tosa(tosa_fb, node.name, new_shape, reshaped.name)
 
         tiled = tosa_fb.addIntermediate(common_shape, tens_dtype)
         multipliers = [
@@ -137,7 +137,7 @@ def broadcast_tensors(
     return broadcast_tensors
 
 
-def build_reshape_tosa_1_0(
+def build_reshape_tosa(
     tosa_graph, input_name, new_shape, output_name, shape_name_override=""
 ):
     """Insert a TOSA reshape operator using the v1.0 semantics.

backends/vulkan/runtime/graph/ops/glsl/clone.glsl

@@ -16,7 +16,10 @@ layout(std430) buffer;
 
 ${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)}
 ${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)}
-${layout_declare_ubo(B, "ivec3", "out_limits")}
+
+layout(push_constant) uniform restrict Block {
+  ivec3 out_limits;
+};
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 

backends/vulkan/runtime/graph/ops/impl/Clone.cpp

@@ -48,9 +48,9 @@ void add_clone_node(
       // Inputs and Outputs
       {{out, vkapi::kWrite}, {in, vkapi::kRead}},
       // Parameter Buffers
-      {graph.logical_limits_ubo(out)},
-      // Push Constants
       {},
+      // Push Constants
+      {graph.logical_limits_pc_of(out)},
       // Specialization Constants
       {},
       // Resize Args

examples/arm/executor_runner/arm_executor_runner.cpp

@@ -87,6 +87,7 @@
 #include <errno.h>
 #include <executorch/extension/data_loader/buffer_data_loader.h>
 #include <executorch/extension/runner_util/inputs.h>
+#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
 #include <executorch/runtime/core/memory_allocator.h>
 #include <executorch/runtime/executor/program.h>
 #include <executorch/runtime/platform/log.h>
@@ -95,6 +96,7 @@
 #include <stdio.h>
 #include <unistd.h>
 #include <memory>
+#include <type_traits>
 #include <vector>
 
 #include "arm_memory_allocator.h"
@@ -183,6 +185,7 @@ using executorch::runtime::Result;
 using executorch::runtime::Span;
 using executorch::runtime::Tag;
 using executorch::runtime::TensorInfo;
+using executorch::runtime::toString;
 #if defined(ET_BUNDLE_IO)
 using executorch::bundled_program::compute_method_output_error_stats;
 using executorch::bundled_program::ErrorStats;
@@ -395,6 +398,19 @@ class Box {
   }
 };
 
+template <typename ValueType>
+void fill_tensor_with_default_value(Tensor& tensor) {
+  ValueType fill_value{};
+  if constexpr (std::is_same_v<ValueType, bool>) {
+    fill_value = true;
+  } else {
+    fill_value = ValueType(1);
+  }
+
+  ValueType* data_ptr = tensor.mutable_data_ptr<ValueType>();
+  std::fill(data_ptr, data_ptr + tensor.numel(), fill_value);
+}
+
 Error prepare_input_tensors(
     Method& method,
     MemoryAllocator& allocator,
@@ -452,32 +468,17 @@ Error prepare_input_tensors(
       if (input_evalues[i].isTensor()) {
         Tensor& tensor = input_evalues[i].toTensor();
         switch (tensor.scalar_type()) {
-          case ScalarType::Int:
-            std::fill(
-                tensor.mutable_data_ptr<int>(),
-                tensor.mutable_data_ptr<int>() + tensor.numel(),
-                1);
-            break;
-          case ScalarType::Float:
-            std::fill(
-                tensor.mutable_data_ptr<float>(),
-                tensor.mutable_data_ptr<float>() + tensor.numel(),
-                1.0);
-            break;
-          case ScalarType::Char:
-            std::fill(
-                tensor.mutable_data_ptr<int8_t>(),
-                tensor.mutable_data_ptr<int8_t>() + tensor.numel(),
-                1);
-            break;
-          case ScalarType::Bool:
-            std::fill(
-                tensor.mutable_data_ptr<int8_t>(),
-                tensor.mutable_data_ptr<int8_t>() + tensor.numel(),
-                1);
-            break;
+#define HANDLE_SCALAR_TYPE(cpp_type, scalar_name)     \
+  case ScalarType::scalar_name:                       \
+    fill_tensor_with_default_value<cpp_type>(tensor); \
+    break;
+          ET_FORALL_SCALAR_TYPES(HANDLE_SCALAR_TYPE)
+#undef HANDLE_SCALAR_TYPE
           default:
-            ET_LOG(Error, "Unhandled ScalarType");
+            ET_LOG(
+                Error,
+                "Unhandled ScalarType %s",
+                toString(tensor.scalar_type()));
             err = Error::InvalidArgument;
             break;
         }