Update on "[ET-VK] Using push constants for unary op."

This diff transitions the unary op to utilize push constants, replacing the previous ubo implementation.

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

[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