Use optimized library if op available (#324)

Summary:
Pull Request resolved: #324

Some models (transformer models) has ops like bmm and optimized library providers a better performance than portable.

With optimized: https://www.internalfb.com/intern/aibench/details/125435623160736 (750 ms per iteration)
With portable: https://www.internalfb.com/intern/aibench/details/521699054450309 (1200 ms per iteration)

Reviewed By: mergennachin

Differential Revision: D49240517

fbshipit-source-id: 7cbc9dfe9ddf7eeff28daa41a0965a74f913497d

Author: kirklandsign

kernels/optimized/targets.bzl

@@ -34,6 +34,9 @@ def define_common_targets():
     et_operator_library(
         name = "optimized_oplist",
         ops_schema_yaml_target = ":optimized.yaml",
+        visibility = [
+            "@EXECUTORCH_CLIENTS",
+        ],
     )
 
     # Used mainly for operator testing. In practice, a generated lib specific

sdk/runners/targets.bzl

@@ -4,6 +4,7 @@ load(
     "CXX",
 )
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+load("@fbsource//xplat/executorch/codegen:codegen.bzl", "executorch_generated_lib")
 load("@fbsource//xplat/executorch/extension/pybindings:pybindings.bzl", "MODELS_ATEN_OPS_ATEN_MODE_GENERATED_LIB", "MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB")
 
 def define_common_targets():
@@ -13,11 +14,34 @@ def define_common_targets():
     TARGETS and BUCK files that call this function.
     """
 
-    # Test driver for models, uses all portable kernels.
-    for aten_mode in (True, False):
+    executorch_generated_lib(
+        name = "generated_op_lib_for_runner",
+        deps = [
+            "//executorch/kernels/optimized:optimized_operators",
+            "//executorch/kernels/optimized:optimized_oplist",
+            "//executorch/kernels/portable:executorch_aten_ops",
+            "//executorch/kernels/portable:executorch_custom_ops",
+            "//executorch/kernels/portable:operators",
+        ],
+        custom_ops_aten_kernel_deps = [
+            "//executorch/kernels/portable:operators_aten",
+        ],
+        functions_yaml_target = "//executorch/kernels/optimized:optimized.yaml",
+        custom_ops_yaml_target = "//executorch/kernels/portable:custom_ops.yaml",
+        fallback_yaml_target = "//executorch/kernels/portable:functions.yaml",
+        define_static_targets = True,
+    )
+
+    # Test driver for models, uses one of the following kernels:
+    # (1) all portable kernels
+    # (2) aten kernels
+    # (3) optimized kernels, with portable kernels as fallback
+    for kernel_mode in ["portable", "aten", "optimized"]:
+        aten_mode = kernel_mode == "aten"
         aten_suffix = ("_aten" if aten_mode else "")
+        binary_name = "executor_runner" if kernel_mode == "portable" else ("executor_runner_" + kernel_mode)
         runtime.cxx_binary(
-            name = "executor_runner" + aten_suffix,
+            name = binary_name,
             srcs = ["executor_runner.cpp"],
             deps = [
                 "//executorch/runtime/executor/test:test_backend_compiler_lib" + aten_suffix,
@@ -31,7 +55,9 @@ def define_common_targets():
             ] + (MODELS_ATEN_OPS_ATEN_MODE_GENERATED_LIB if aten_mode else [
                 "//executorch/configurations:executor_cpu_optimized",
                 "//executorch/kernels/quantized:generated_lib",
-            ] + MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB),
+            ] + (MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB if kernel_mode == "portable" else [
+                ":generated_op_lib_for_runner",
+            ])),
             preprocessor_flags = ["-DUSE_ATEN_LIB"] if aten_mode else [],
             external_deps = [
                 "gflags",