Update on "[ET-VK] Adding push constant and ubo verison of select and slice ops to improve memory and performance."

Adding push constant and ubo verison of select and slice ops to improve memory and performance.
*   Updated `transfer_buffer.yaml` and `transfer_texture.yaml` to include `UBO_PARAMS` parameter and generate variants for `select` and `slice` ops with UBO parameters.
*   Updated `transfer.glsl` to generate ubo and push constant versions of `select` and `slice` ops with UBO parameters.

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

[ghstack-poisoned]

Author: trivedivivek

backends/cadence/aot/compiler.py

@@ -8,7 +8,7 @@
 
 import logging
 from pathlib import Path
-from typing import Optional
+from typing import Callable, cast, Optional
 
 import executorch.backends.cadence.aot.ops_registrations  # noqa
 import torch
@@ -32,6 +32,7 @@
     ExecutorchBackendConfig,
     ExecutorchProgramManager,
 )
+from executorch.exir.pass_base import PassResult
 from executorch.exir.passes import ToOutVarPass
 from executorch.exir.passes.sym_shape_eval_pass import HintBasedSymShapeEvalPass
 from executorch.exir.program._program import to_edge_with_preserved_ops
@@ -40,7 +41,7 @@
 from torch.export.exported_program import ExportedProgram
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
-from .passes import apply_exir_ops_passes, apply_torch_ops_passes
+from .passes import get_cadence_passes
 
 from .utils import print_ops_info
 
@@ -261,20 +262,14 @@ def export_to_edge(
     inputs: tuple[object, ...],
     dump_graphs: bool = False,
     constant_methods: Optional[dict[str, object]] = None,
-    core_aten_exceptions: Optional[list[torch._ops.OpOverload]] = None,
 ) -> EdgeProgramManager:
     assert isinstance(model, torch.nn.Module), "model should be an nn.Module"
 
     # Export the model into an ExportedProgram.
     expo_program = trace(model, inputs)
 
-    # Apply passes which transform the ExportedProgram before it gets lowered to edge.
-    expo_program = apply_torch_ops_passes(expo_program)
-
     # Lower the model to edge IR.
-    edge_prog_manager = _lower_ep_to_edge(
-        expo_program, dump_graphs, constant_methods, core_aten_exceptions
-    )
+    edge_prog_manager = _lower_ep_to_edge(expo_program, dump_graphs, constant_methods)
 
     return edge_prog_manager
 
@@ -316,7 +311,14 @@ def _lower_ep_to_cadence(
     Lower an existing ExportedProgram to edge IR and apply frontend optimization passes.
     """
     edge_prog_manager = _lower_ep_to_edge(program, dump_graphs=dump_graphs)
-    cadence_prog_manager = apply_exir_ops_passes(opt_level, edge_prog_manager)
+    cadence_passes = get_cadence_passes(opt_level)
+
+    # Run a couple required passes for quant/dequant ops
+    cadence_prog_manager = edge_prog_manager.transform(
+        cast(
+            list[Callable[[torch.fx.GraphModule], Optional[PassResult]]], cadence_passes
+        )
+    )
     return cadence_prog_manager
 
 
@@ -327,7 +329,14 @@ def export_to_cadence(
     opt_level: int = 1,
 ) -> EdgeProgramManager:
     edge_prog_manager = export_to_edge(model, inputs, dump_graphs=dump_graphs)
-    cadence_prog_manager = apply_exir_ops_passes(opt_level, edge_prog_manager)
+    cadence_passes = get_cadence_passes(opt_level)
+
+    # Run a couple required passes for quant/dequant ops
+    cadence_prog_manager = edge_prog_manager.transform(
+        cast(
+            list[Callable[[torch.fx.GraphModule], Optional[PassResult]]], cadence_passes
+        )
+    )
     return cadence_prog_manager
 
 
@@ -364,8 +373,15 @@ def export_to_executorch_gen_etrecord(
     memory_config: Optional[MemoryConfig] = None,
     dump_graphs: bool = False,
 ) -> ExecutorchProgramManager:
+    cadence_passes = get_cadence_passes(opt_level)
     edge_prog_manager = export_to_edge(model, inputs, dump_graphs)
-    cadence_prog_manager = apply_exir_ops_passes(opt_level, edge_prog_manager)
+
+    # Run a couple required passes for quant/dequant ops
+    cadence_prog_manager = edge_prog_manager.transform(
+        cast(
+            list[Callable[[torch.fx.GraphModule], Optional[PassResult]]], cadence_passes
+        )
+    )
 
     # Print some information to terminal
     print_ops_info(

backends/cadence/aot/fuse_ops.py

@@ -1127,7 +1127,6 @@ class CadenceFuseOpsInGraph:
         FuseCascadedTransposeOrPermuteOps,
         FuseCascadedViewOps,
         FuseQuantDequantToRequantizePass,
-        FuseMulTensorIntoQuantPass,
         FuseMulTensorIntoDequantPass,
         FuseMulScalarIntoDequantPass,
         FuseFullThenReshapePass,

backends/cadence/aot/passes.py

@@ -6,7 +6,7 @@
 
 # pyre-strict
 
-from typing import Any, Callable, cast, List, Optional
+from typing import Any, List, Optional
 
 import torch
 import torch.fx
@@ -28,18 +28,13 @@
     RemoveRedundantOps,
 )
 from executorch.backends.cadence.aot.reorder_ops import CadenceReorderOpsInGraph
-from executorch.backends.cadence.aot.replace_ops import (
-    CadenceReplaceOpsInGraph,
-    ReplaceMulTensorWithMulAndFullOpsPass,
-)
+from executorch.backends.cadence.aot.replace_ops import CadenceReplaceOpsInGraph
 from executorch.backends.cadence.aot.simplify_ops import CadenceSimplifyOpsInGraph
-from executorch.exir import EdgeProgramManager
 from executorch.exir.pass_base import ExportPass, PassResult
 from executorch.exir.pass_manager import PassManager, PassType
 from executorch.exir.passes import dead_code_elimination_pass
 from executorch.exir.passes.scalar_to_tensor_pass import ScalarToTensorPass
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
-from torch.export.exported_program import ExportedProgram
 
 
 @register_cadence_pass(CadencePassAttribute(opt_level=0))
@@ -94,37 +89,14 @@ def get_passes_in_default_order() -> List[ExportPass]:
     return pytree.tree_flatten(passes)[0]
 
 
-def apply_exir_ops_passes(
+def get_cadence_passes(
     opt_level: int,
-    edge_prog_manager: EdgeProgramManager,
-) -> EdgeProgramManager:
+) -> List[Optional[PassResult]]:
     passes = get_passes_in_default_order()
     pass_filter = create_cadence_pass_filter(opt_level)
-    cadence_passes = [
-        (
-            lambda graph_module, filtered_pass=filtered_pass: filtered_pass()(
-                graph_module
-            )
-        )
+    filtered_passes = [
+        # pyre-ignore[20]: Expect argument graph_module
+        filtered_pass()
         for filtered_pass in list(filter(pass_filter, passes))
     ]
-    cadence_prog_manager = edge_prog_manager.transform(
-        cast(
-            list[Callable[[torch.fx.GraphModule], Optional[PassResult]]], cadence_passes
-        )
-    )
-    return cadence_prog_manager
-
-
-def apply_torch_ops_passes(expo_program: ExportedProgram) -> ExportedProgram:
-    """
-    Applies compiler passes on torch.ops IR, including torch.ops.aten, torch.ops.cadence, etc.
-    expo_program is expected to be the output of the torch.export.export().
-    """
-
-    aten_passes: List[Callable[[torch.fx.GraphModule], Optional[PassResult]]] = [
-        ReplaceMulTensorWithMulAndFullOpsPass()
-    ]
-    # TODO(T230417247): Use PassResult which is currently ignored.
-    PassManager(aten_passes)(expo_program.graph_module)
-    return expo_program
+    return filtered_passes

examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/project.pbxproj

@@ -485,6 +485,8 @@
 				03CF43A52CEC5CEC00C7113B /* kernels_custom_debug */,
 				03CF43A72CEC5CEC00C7113B /* kernels_optimized */,
 				03CF43A92CEC5CEC00C7113B /* kernels_optimized_debug */,
+				03CF43AB2CEC5CEC00C7113B /* kernels_portable */,
+				03CF43AD2CEC5CEC00C7113B /* kernels_portable_debug */,
 				03CF43AF2CEC5CEC00C7113B /* kernels_quantized */,
 				03CF43B12CEC5CEC00C7113B /* kernels_quantized_debug */,
 			);
@@ -1010,6 +1012,16 @@
 			package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */;
 			productName = kernels_optimized_debug;
 		};
+		03CF43AB2CEC5CEC00C7113B /* kernels_portable */ = {
+			isa = XCSwiftPackageProductDependency;
+			package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */;
+			productName = kernels_portable;
+		};
+		03CF43AD2CEC5CEC00C7113B /* kernels_portable_debug */ = {
+			isa = XCSwiftPackageProductDependency;
+			package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */;
+			productName = kernels_portable_debug;
+		};
 		03CF43AF2CEC5CEC00C7113B /* kernels_quantized */ = {
 			isa = XCSwiftPackageProductDependency;
 			package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */;

examples/demo-apps/react-native/rnllama/ios/LlamaBridge.h

@@ -1,9 +1,9 @@
 #ifndef LLaMABridge_h
 #define LLaMABridge_h
 
-#import <LLaMARunner/LLaMARunner.h>
 #import <React/RCTBridgeModule.h>
 #import <React/RCTEventEmitter.h>
+#import "LLaMARunner.h"
 
 NS_ASSUME_NONNULL_BEGIN
 

examples/demo-apps/react-native/rnllama/ios/Release.xcconfig

@@ -0,0 +1,24 @@
+ET_PLATFORM[sdk=iphonesimulator*] = simulator
+ET_PLATFORM[sdk=iphoneos*] = ios
+ET_PLATFORM[sdk=macos*] = macos
+
+// Link the Debug version of ExecuTorch runtime to keep the logs.
+// Switch to Release for better performance if logs are not needed.
+OTHER_LDFLAGS = $(inherited) \
+    -force_load $(BUILT_PRODUCTS_DIR)/libexecutorch_debug_$(ET_PLATFORM).a \
+    -force_load $(BUILT_PRODUCTS_DIR)/libbackend_coreml_$(ET_PLATFORM).a \
+    -force_load $(BUILT_PRODUCTS_DIR)/libbackend_mps_$(ET_PLATFORM).a \
+    -force_load $(BUILT_PRODUCTS_DIR)/libbackend_xnnpack_$(ET_PLATFORM).a \
+    -force_load $(BUILT_PRODUCTS_DIR)/libkernels_custom_$(ET_PLATFORM).a \
+    -force_load $(BUILT_PRODUCTS_DIR)/libkernels_optimized_$(ET_PLATFORM).a \
+    -force_load $(BUILT_PRODUCTS_DIR)/libkernels_quantized_$(ET_PLATFORM).a \
+    @$(TEMP_DIR)/cmake/linker_flags
+
+// LLaMARunner requires additional dependencies built with CMake in a custom run script phase.
+// Include headers and libraries from $(TEMP_DIR)/cmake for it.
+HEADER_SEARCH_PATHS = $(inherited) \
+    $(SRCROOT)/../../../../../.. \
+    $(TEMP_DIR)/cmake/include
+
+LIBRARY_SEARCH_PATHS = $(inherited) \
+    $(TEMP_DIR)/cmake/lib