Update on "[Executorch][custom ops] Change lib loading logic to account for package dir"

Just looking at the location of the source file. In this case custom_ops.py,
can, and does, yield to wrong location depending on where you import custom_ops
from.

If you are importing custom_ops from another source file inside extension
folder, e.g. builder.py that is in extensions/llm/export, then, I think,
custom_ops gets resolved to the one installed in site-packages or pip package.

But if this is imported from say
examples/models/llama/source_transformations/quantized_kv_cache.py (Like in the
in next PR), then it seems to resolve to the source location. In one of the CI
this is /pytorch/executorch.

Now depending on which directory your filepath resolves to, you will search for
lib in that. This of course does not work when filepath resolves to source
location.

This PR changes that to resolve to package location.

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

[ghstack-poisoned]

Author: kimishpatel

.github/scripts/extract_benchmark_results.py

@@ -451,7 +451,7 @@ def main() -> None:
             continue
 
         output_dir = os.path.join(args.output_dir, schema)
-        os.mkdir(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
 
         output_file = os.path.basename(args.artifacts)
         with open(f"{output_dir}/{output_file}", "w") as f:

CMakeLists.txt

@@ -200,8 +200,6 @@ option(EXECUTORCH_BUILD_EXTENSION_TENSOR "Build the Tensor extension" OFF)
 
 option(EXECUTORCH_BUILD_EXTENSION_TRAINING "Build the training extension" OFF)
 
-option(EXECUTORCH_BUILD_GTESTS "Build googletest based test binaries" OFF)
-
 option(EXECUTORCH_BUILD_MPS "Build the MPS backend" OFF)
 
 option(EXECUTORCH_BUILD_NEURON "Build the backends/mediatek directory" OFF)
@@ -216,6 +214,8 @@ option(EXECUTORCH_BUILD_KERNELS_QUANTIZED "Build the quantized kernels" OFF)
 
 option(EXECUTORCH_BUILD_DEVTOOLS "Build the ExecuTorch Developer Tools")
 
+option(EXECUTORCH_BUILD_TESTS "Build CMake-based unit tests" OFF)
+
 option(EXECUTORCH_NNLIB_OPT "Build Cadence backend Hifi nnlib kernel" OFF)
 
 option(EXECUTORCH_CADENCE_CPU_RUNNER "Build Cadence backend CPU runner" OFF)
@@ -330,6 +330,10 @@ if(EXECUTORCH_BUILD_PTHREADPOOL)
   )
 endif()
 
+if(EXECUTORCH_BUILD_TESTS)
+  include(CTest)
+endif()
+
 if(NOT PYTHON_EXECUTABLE)
   resolve_python_executable()
 endif()
@@ -625,7 +629,7 @@ cmake_dependent_option(
 )
 
 # Add googletest if any test targets should be built
-if(EXECUTORCH_BUILD_GTESTS)
+if(BUILD_TESTING)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third-party/googletest)
 endif()
 
@@ -829,5 +833,7 @@ if(EXECUTORCH_BUILD_VULKAN)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/vulkan)
 endif()
 
+include(Test.cmake)
+
 # Print all summary
 executorch_print_configuration_summary()

Test.cmake

@@ -0,0 +1,29 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+#
+# A helper CMake file to trigger C++ unit tests.
+#
+
+if(BUILD_TESTING)
+  # This contains the list of tests which are always built
+  add_subdirectory(extension/evalue_util/test)
+  add_subdirectory(extension/kernel_util/test)
+  add_subdirectory(extension/memory_allocator/test)
+  add_subdirectory(extension/parallel/test)
+  add_subdirectory(extension/pytree/test)
+  add_subdirectory(kernels/portable/cpu/util/test)
+  add_subdirectory(kernels/prim_ops/test)
+  add_subdirectory(kernels/test)
+  add_subdirectory(runtime/core/exec_aten/testing_util/test)
+  add_subdirectory(runtime/core/exec_aten/util/test)
+  add_subdirectory(runtime/core/portable_type/test)
+  add_subdirectory(runtime/core/test)
+  add_subdirectory(runtime/executor/test)
+  add_subdirectory(runtime/kernel/test)
+  add_subdirectory(runtime/platform/test)
+  add_subdirectory(test/utils)
+endif()

backends/arm/_passes/arm_pass_utils.py

@@ -156,7 +156,7 @@ def get_node_arg(args: list | dict, key: int | str | type, default_value=None):
                 f"Out of bounds index {key} for getting value in args (of size {len(args)})"
             )
     elif isinstance(key, str):
-        return args.get(key, default_value)
+        return args.get(key, default_value)  # pyre-ignore[16]
     elif isclass(key):
         for arg in args:
             if isinstance(arg, key):

backends/arm/_passes/keep_dims_false_to_squeeze_pass.py

@@ -64,12 +64,17 @@ def call(self, graph_module: torch.fx.GraphModule):
                 continue
 
             sum_node = cast(torch.fx.Node, node)
-            keep_dim = get_node_arg(sum_node.args, keep_dim_index, False)
+            keep_dim = get_node_arg(
+                # pyre-ignore[6]
+                sum_node.args,
+                keep_dim_index,
+                False,
+            )
 
             if keep_dim:
                 continue
 
-            dim_list = get_node_arg(sum_node.args, 1, [0])
+            dim_list = get_node_arg(sum_node.args, 1, [0])  # pyre-ignore[6]
 
             # Add keep_dim = True arg to sum node.
             set_node_arg(sum_node, 2, True)

backends/arm/operator_support/to_copy_support.py

@@ -33,7 +33,9 @@ class ToCopySupported(SupportedTOSAOperatorCheck):
 
     @staticmethod
     def _merge_supported_types(
-        dtypes1: SupportedTypeDict, dtypes2: SupportedTypeDict
+        # pyre-ignore[11]
+        dtypes1: SupportedTypeDict,
+        dtypes2: SupportedTypeDict,
     ) -> SupportedTypeDict:
         merged_dtypes = dtypes1
         for k, v in dtypes2.items():
@@ -92,7 +94,7 @@ def is_node_supported(self, node: fx.Node, tosa_spec: TosaSpecification) -> bool
         if input_dtype not in supported_dtypes:
             logger.info(
                 f"Input dtype {input_val.dtype} is not supported in "
-                f"{node.target.name()}."
+                f"{node.target.name()}."  # pyre-ignore[16]
             )
             return False
 
@@ -102,7 +104,7 @@ def is_node_supported(self, node: fx.Node, tosa_spec: TosaSpecification) -> bool
         if output_val.dtype not in supported_dtypes[input_dtype]:
             logger.info(
                 f"Output dtype {output_val.dtype} is not supported in "
-                f"{node.target.name()} for input dtype {input_dtype}. "
+                f"{node.target.name()} for input dtype {input_dtype}. "  # pyre-ignore[16]
                 f"Supported output types: "
                 f"{''.join(str(t) for t in supported_dtypes[input_dtype])}"
             )
@@ -113,7 +115,7 @@ def is_node_supported(self, node: fx.Node, tosa_spec: TosaSpecification) -> bool
             if node.kwargs["memory_format"] in (torch.preserve_format,):
                 logger.info(
                     f"Argument 'memory_format' is not supported for "
-                    f"{node.target.name()} right now."
+                    f"{node.target.name()} right now."  # pyre-ignore[16]
                 )
                 return False
 

backends/cadence/CMakeLists.txt

@@ -23,7 +23,6 @@ include(${EXECUTORCH_ROOT}/build/Utils.cmake)
 
 # Let files say "include <executorch/path/to/header.h>".
 set(_common_include_directories ${EXECUTORCH_ROOT}/..)
-set(TARGET_DIR reference)
 
 if(EXECUTORCH_CADENCE_CPU_RUNNER)
   include(${EXECUTORCH_ROOT}/build/Codegen.cmake)
@@ -61,6 +60,9 @@ if(EXECUTORCH_CADENCE_CPU_RUNNER)
                                       ${_common_include_directories}
   )
 
+  set(TARGET_DIR reference)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/${TARGET_DIR}/kernels)
+
   target_link_libraries(
   cadence_runner
   executorch

backends/cadence/aot/functions.yaml

@@ -142,6 +142,41 @@
     - arg_meta: null
       kernel_name: torch::executor::where_out
 
+- op: transpose_copy.int_out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::transpose_copy_int_out
+
+- op: eq.Scalar_out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::eq_scalar_out
+
+- op: logical_not.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::logical_not_out
+
+- op: any.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::any_out
+
+- op: native_group_norm.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::native_group_norm_out
+
+- op: sum.IntList_out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::sum_dim_out
+
+- op: select_copy.int_out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::select_copy_int_out
+
 # custom ops
 - func: cadence::quantize_per_tensor.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
   variants: function
@@ -183,3 +218,18 @@
   kernels:
     - arg_meta: null
       kernel_name: impl::reference::quantized_matmul_out
+
+- func: cadence::quantized_linear.per_tensor_out(Tensor src, Tensor weight, Tensor bias, SymInt src_zero_point, SymInt weight_zero_point, SymInt out_multiplier, SymInt out_shift, SymInt out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: impl::reference::quantized_linear_per_tensor_out
+
+- func: cadence::im2row.out(Tensor input, int[2] kernel_size, int[2] dilation, int[2] padding, int[2] stride, Tensor in_zero_point, bool channel_last=False, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: impl::reference::im2row_out
+
+- func: cadence::quantized_conv.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, bool channel_last=False, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: impl::reference::quantized_conv_per_tensor_out

backends/cadence/reference/operators/CMakeLists.txt

@@ -55,6 +55,16 @@ set(_aten_ops__srcs
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_expand_copy.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_gelu.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_empty.cpp"
+    "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_transpose_copy.cpp"
+    "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_eq.cpp"
+    "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_logical_not.cpp"
+    "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_any.cpp"
+    "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_native_group_norm.cpp"
+    "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_sum.cpp"
+    "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_select_copy.cpp"
+    "${EXECUTORCH_ROOT}/kernels/portable/cpu/util/dtype_util.cpp"
+    "${EXECUTORCH_ROOT}/kernels/portable/cpu/util/normalization_ops_util.cpp"
+    "${EXECUTORCH_ROOT}/kernels/portable/cpu/util/select_copy_util.cpp"
 )
 add_library(aten_ops_cadence ${_aten_ops__srcs})
 target_link_libraries(aten_ops_cadence PUBLIC executorch)
@@ -78,6 +88,7 @@ add_library(
   "quantize_per_tensor.cpp"
   "dequantize_per_tensor.cpp"
   "quantized_matmul_out.cpp"
+  "im2row_out.cpp"
 )
 target_include_directories(
   custom_ops PUBLIC ${ROOT_DIR}/.. ${CMAKE_BINARY_DIR}