actually verified test coverage

[ghstack-poisoned]

Author: swolchok

kernels/portable/CMakeLists.txt

@@ -69,8 +69,15 @@ if(EXECUTORCH_BUILD_PTHREADPOOL AND EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
   target_compile_options(optimized_portable_kernels PUBLIC ${_common_compile_options})
   target_include_directories(optimized_portable_kernels PRIVATE ${TORCH_INCLUDE_DIRS})
   target_compile_definitions(optimized_portable_kernels PRIVATE ET_USE_PYTORCH_HEADERS)
+  gen_selected_ops(LIB_NAME "optimized_portable_ops_lib" OPS_SCHEMA_YAML "${_yaml}")
+  generate_bindings_for_kernels(
+    LIB_NAME "optimized_portable_ops_lib" FUNCTIONS_YAML "${_yaml}"
+  )
+  gen_operators_lib(
+    LIB_NAME "optimized_portable_ops_lib" KERNEL_LIBS optimized_portable_kernels DEPS executorch_core
+  )
   install(
-    TARGETS optimized_portable_kernels
+    TARGETS optimized_portable_kernels optimized_portable_ops_lib
     DESTINATION lib
   )
 endif()

kernels/portable/cpu/op_add.cpp

@@ -102,14 +102,18 @@ Tensor& add_scalar_out(
   static constexpr const char op_name[] = "add.Scalar_out";
 
   ET_SWITCH_REALB_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() {
+    CTYPE_COMPUTE val_b = utils::scalar_to<CTYPE_COMPUTE>(b);
+    CTYPE_COMPUTE val_alpha = utils::scalar_to<CTYPE_COMPUTE>(alpha);
+    auto val_alpha_times_b = val_alpha * val_b;
     utils::apply_unitensor_elementwise_fn<
         CTYPE_COMPUTE,
         op_name,
         utils::SupportedTensorDtypes::SAME_AS_COMMON>(
-        [b, alpha](const auto val_a) {
-          CTYPE_COMPUTE val_b = utils::scalar_to<CTYPE_COMPUTE>(b);
-          CTYPE_COMPUTE val_alpha = utils::scalar_to<CTYPE_COMPUTE>(alpha);
-          return val_a + val_alpha * val_b;
+        [val_alpha_times_b](const auto val_a) {
+          // Cast here supports vectorization; either it does nothing
+          // or it casts from CTYPE_COMPUTE to
+          // Vectorized<CTYPE_COMPUTE>.
+          return val_a + decltype(val_a)(val_alpha_times_b);
         },
         ctx,
         a,

kernels/portable/cpu/op_clamp.cpp

@@ -138,9 +138,8 @@ Tensor& clamp_out(
         CTYPE_COMPUTE,
         op_name,
         utils::SupportedTensorDtypes::SAME_AS_COMMON>(
-        [has_min, min_opt, has_max, max_opt](const CTYPE_COMPUTE val_in) {
-          // TODO: rewrite this to be vectorization-capable.
-          CTYPE_COMPUTE val_out = val_in;
+        [has_min, min_opt, has_max, max_opt](const auto val_in) {
+          auto val_out = val_in;
           if (has_min) {
             val_out = utils::max_override(
                 val_out, utils::scalar_to<CTYPE_COMPUTE>(min_opt.value()));

kernels/portable/cpu/op_native_dropout.cpp

@@ -57,8 +57,11 @@ std::tuple<Tensor&, Tensor&> native_dropout_out(
     }
     ET_SWITCH_FLOATHBF16_TYPES(
         input.scalar_type(), ctx, op_name, CTYPE_COMPUTE, [&]() {
-          utils::apply_bitensor_elementwise_fn<CTYPE_COMPUTE, op_name>(
-              [](const auto val, const auto mask_val) {
+          utils::apply_bitensor_elementwise_fn<
+              CTYPE_COMPUTE,
+              op_name,
+              utils::SupportedTensorDtypes::SAME_AS_COMMON>(
+              [](const CTYPE_COMPUTE val, const CTYPE_COMPUTE mask_val) {
                 if (!mask_val) {
                   return static_cast<decltype(val)>(0);
                 }
@@ -70,8 +73,7 @@ std::tuple<Tensor&, Tensor&> native_dropout_out(
               mask,
               // TODO: should really be just BOOL
               utils::SupportedTensorDtypes::BOOL_OR_BYTE,
-              out,
-              utils::SupportedTensorDtypes::SAME_AS_COMMON);
+              out);
         });
   } else if (input.numel() > 0) {
     std::memcpy(out.mutable_data_ptr(), input.data_ptr(), input.nbytes());

kernels/portable/cpu/op_pow.cpp

@@ -58,7 +58,6 @@ Tensor& pow_Tensor_Tensor_out(
         op_name,
         utils::SupportedTensorDtypes::REALHBF16>(
         [](const auto val_a, const auto val_b) {
-          // TODO: rewrite this to be vectorization-capable.
           return executorch::math::pow(val_a, val_b);
         },
         ctx,
@@ -111,8 +110,13 @@ Tensor& pow_Tensor_Scalar_out(
         CTYPE_COMPUTE,
         op_name,
         utils::SupportedTensorDtypes::REALHBF16>(
-        // TODO: rewrite this to be vectorization-capable.
-        [val_b](const CTYPE_COMPUTE val_a) { return std::pow(val_a, val_b); },
+        // Casting val_b here supports vectorization; it does
+        // nothing if we are not vectorizing (casts to
+        // CTYPE_COMPUTE) and casts to a vectorized type
+        // otherwise.
+        [val_b](const auto val_a) {
+          return executorch::math::pow(val_a, decltype(val_a)(val_b));
+        },
         ctx,
         a,
         utils::SupportedTensorDtypes::REALHBBF16,
@@ -161,8 +165,13 @@ Tensor& pow_Scalar_out(
         CTYPE_COMPUTE,
         op_name,
         utils::SupportedTensorDtypes::REALHBF16>(
-        // TODO: rewrite this to be vectorization-capable.
-        [val_a](const CTYPE_COMPUTE val_b) { return std::pow(val_a, val_b); },
+        // Casting val_a here supports vectorization; it does
+        // nothing if we are not vectorizing (casts to
+        // CTYPE_COMPUTE) and casts to a vectorized type
+        // otherwise.
+        [val_a](const auto val_b) {
+          return executorch::math::pow(decltype(val_b)(val_a), val_b);
+        },
         ctx,
         b,
         utils::SupportedTensorDtypes::REALHBBF16,

kernels/portable/cpu/util/elementwise_util.h

@@ -131,11 +131,22 @@ inline void dtype_specialized_elementwise_fn_impl(
             const auto vectorized_end = end - (end % Vec::size());
             // Scalar prologue.
             for (const auto idx : c10::irange(begin, vectorized_begin)) {
+              // In debug mode, always use Vectorized so that even
+              // small-sized tests will test whether using Vectorized broke our
+              // lambda.
+#ifndef NDEBUG
+              std::array<Vec, kNumInputs> loaded_inputs;
+#else // NDEBUG
               std::array<CTYPE_COMPUTE, kNumInputs> loaded_inputs;
+#endif // NDEBUG
               for (const auto input_idx : c10::irange(kNumInputs)) {
                 loaded_inputs[input_idx] = inputs_data_ptrs[input_idx][idx];
               }
+#ifndef NDEBUG
+              std::apply(compute_fun, loaded_inputs).store(&data_out[idx], 1);
+#else // NDEBUG
               data_out[idx] = std::apply(compute_fun, loaded_inputs);
+#endif // NDEBUG
             }
 
             // Main vectorized loop.
@@ -152,11 +163,19 @@ inline void dtype_specialized_elementwise_fn_impl(
 
             // Scalar epilogue.
             for (const auto idx : c10::irange(vectorized_end, end)) {
+#ifndef NDEBUG
+              std::array<Vec, kNumInputs> loaded_inputs;
+#else // NDEBUG
               std::array<CTYPE_COMPUTE, kNumInputs> loaded_inputs;
+#endif // NDEBUG
               for (const auto input_idx : c10::irange(kNumInputs)) {
                 loaded_inputs[input_idx] = inputs_data_ptrs[input_idx][idx];
               }
+#ifndef NDEBUG
+              std::apply(compute_fun, loaded_inputs).store(&data_out[idx], 1);
+#else // NDEBUG
               data_out[idx] = std::apply(compute_fun, loaded_inputs);
+#endif // NDEBUG
             }
           });
       return;

kernels/portable/cpu/util/math_util.h

@@ -150,12 +150,23 @@ at::vec::Vectorized<T> min_override(
   return at::vec::minimum(a, b);
 }
 
+template <typename T>
+at::vec::Vectorized<T> min_override(at::vec::Vectorized<T> a, T b) {
+  return min_override(a, at::vec::Vectorized<T>(b));
+}
+
 template <typename T>
 at::vec::Vectorized<T> max_override(
     at::vec::Vectorized<T> a,
     at::vec::Vectorized<T> b) {
   return at::vec::maximum(a, b);
 }
+
+template <typename T>
+at::vec::Vectorized<T> max_override(at::vec::Vectorized<T> a, T b) {
+  return max_override(a, at::vec::Vectorized<T>(b));
+}
+
 #endif
 /**
  * There is a slight difference in how std::fmod works compared to how ATen

kernels/test/CMakeLists.txt

@@ -17,7 +17,7 @@ set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..)
 
 include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake)
 
-set(_kernels portable optimized quantized)
+set(_kernels portable optimized_portable optimized quantized)
 foreach(kernel ${_kernels})
   set(_wrapper_dir
       "${CMAKE_CURRENT_BINARY_DIR}/include/${kernel}/executorch/kernels/test"
@@ -37,13 +37,17 @@ foreach(kernel ${_kernels})
     VERBATIM
   )
 
+  set(_supported_features_kernel ${kernel})
+  if(${kernel} STREQUAL "optimized_portable")
+    set(_supported_features_kernel "portable")
+  endif()
   add_custom_command(
     OUTPUT "${_wrapper_dir}/supported_features.cpp"
            "${_wrapper_dir}/supported_features.h"
     COMMAND mkdir -p ${_wrapper_dir}
     COMMAND
       ${PYTHON_EXECUTABLE} kernels/test/gen_supported_features.py
-      kernels/${kernel}/test/supported_features_def.yaml >
+      kernels/${_supported_features_kernel}/test/supported_features_def.yaml >
       ${_wrapper_dir}/supported_features.cpp
     COMMAND
       ${PYTHON_EXECUTABLE} kernels/test/gen_supported_features.py
@@ -57,6 +61,11 @@ foreach(kernel ${_kernels})
     set(_kernel_ops_lib "optimized_native_cpu_ops_lib")
     set(_kernel_ops_lib_path
         "${CMAKE_CURRENT_BINARY_DIR}/../../configurations/optimized_native_cpu_ops_lib"
+      )
+  elseif(${kernel} STREQUAL "optimized_portable")
+    set(_kernel_ops_lib "${kernel}_ops_lib")
+    set(_kernel_ops_lib_path
+        "${CMAKE_CURRENT_BINARY_DIR}/../../kernels/portable/${kernel}_ops_lib"
     )
   else()
     set(_kernel_ops_lib "${kernel}_ops_lib")
@@ -88,6 +97,9 @@ add_custom_target(
     "${CMAKE_CURRENT_BINARY_DIR}/include/optimized/executorch/kernels/test/FunctionHeaderWrapper.h"
     "${CMAKE_CURRENT_BINARY_DIR}/include/optimized/executorch/kernels/test/supported_features.h"
     "${CMAKE_CURRENT_BINARY_DIR}/include/optimized/executorch/kernels/test/supported_features.cpp"
+    "${CMAKE_CURRENT_BINARY_DIR}/include/optimized_portable/executorch/kernels/test/FunctionHeaderWrapper.h"
+    "${CMAKE_CURRENT_BINARY_DIR}/include/optimized_portable/executorch/kernels/test/supported_features.h"
+    "${CMAKE_CURRENT_BINARY_DIR}/include/optimized_portable/executorch/kernels/test/supported_features.cpp"
     "${CMAKE_CURRENT_BINARY_DIR}/include/quantized/executorch/kernels/test/FunctionHeaderWrapper.h"
     "${CMAKE_CURRENT_BINARY_DIR}/include/quantized/executorch/kernels/test/supported_features.h"
     "${CMAKE_CURRENT_BINARY_DIR}/include/quantized/executorch/kernels/test/supported_features.cpp"
@@ -297,6 +309,22 @@ set(_optimized_kernels_test_sources
 if(TARGET optimized_portable_kernels)
   list(APPEND _optimized_kernels_test_sources ${all_test_sources})
   list(REMOVE_DUPLICATES _optimized_kernels_test_sources)
+
+  # Make sure that we still test optimized versions of portable
+  # kernels even if they would currently be shadowed by specific
+  # optimized implementations.
+  et_cxx_test(
+    optimized_portable_kernels_test
+    SOURCES
+    ${all_test_sources}
+    ${CMAKE_CURRENT_BINARY_DIR}/include/optimized_portable/executorch/kernels/test/supported_features.cpp
+    EXTRA_LIBS
+    optimized_portable_kernels
+  )
+   add_dependencies(optimized_portable_kernels_test generate_wrapper)
+  target_include_directories(
+    optimized_portable_kernels_test PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/include/optimized_portable"
+  )
 endif()
 
 et_cxx_test(

kernels/test/op_atan2_test.cpp

@@ -47,17 +47,16 @@ TEST(OpAtan2OutTest, SmokeTest) {
   EXPECT_TENSOR_CLOSE(out, out_expected);
 }
 
-TEST(OpAtan2OutTest, SmokeTestNoBroadcasting) {
+TEST(OpAtan2OutTest, SmokeTestNoBroadcastingSameDtype) {
   TensorFactory<ScalarType::Double> tfDouble;
-  TensorFactory<ScalarType::Float> tfFloat;
 
   std::vector<double> a(18);
   std::iota(a.begin(), a.end(), -8);
   std::vector<double> b(18, 2.0);
   Tensor self = tfDouble.make({18}, a);
   Tensor other = tfDouble.make({18}, b);
-  Tensor out = tfFloat.zeros({18});
-  Tensor out_expected = tfFloat.make(
+  Tensor out = tfDouble.zeros({18});
+  Tensor out_expected = tfDouble.make(
       {18},
       {-1.3258176636680326,
        -1.2924966677897853,

kernels/test/op_clamp_test.cpp

@@ -31,6 +31,15 @@ using torch::executor::testing::TensorFactory;
 
 using OptScalar = executorch::aten::optional<Scalar>;
 
+namespace {
+template <typename T>
+std::vector<T> arange(T stop) {
+  std::vector<T> result(stop);
+  std::iota(result.begin(), result.end(), 0);
+  return result;
+}
+} // namespace
+
 class OpClampOutTest : public OperatorTest {
  protected:
   Tensor& op_clamp_out(
@@ -114,6 +123,31 @@ class OpClampOutTest : public OperatorTest {
             // Should set all elements to max.
             {6, 6, 6, 6}, // expected_data
         },
+        {
+            std::string(__func__) + ": Simple clamp larger data",
+            {18}, // sizes
+            arange<typename ClampTestCase<DTYPE>::ctype>(18), // input_data
+            OptScalar(1), // min
+            OptScalar(6), // max
+            {1,
+             1,
+             2,
+             3,
+             4,
+             5,
+             6,
+             6,
+             6,
+             6,
+             6,
+             6,
+             6,
+             6,
+             6,
+             6,
+             6,
+             6}, // expected_data
+        },
     };
 
     run_test_cases(test_cases);