From db3d6a0cd36691319fe7fdbe7ff68be52fa5efdf Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Wed, 9 Sep 2026 13:43:19 -0700 Subject: [PATCH] Add optimized Arm64 BF16 `_to_copy` (#22493) Summary: Add an optimized `_to_copy.out` kernel. On Arm64, contiguous FP32-to-BF16 and BF16-to-FP32 conversions process eight elements per NEON iteration and use the ExecuTorch threadpool above its grain size. Other dtype pairs, layouts, and platforms fall back to the portable implementation, which remains unchanged. Reviewed By: digantdesai Differential Revision: D118502073 --- kernels/optimized/cpu/op_to_copy.cpp | 184 ++++++++++++++ kernels/optimized/optimized.yaml | 5 + kernels/test/op_to_copy_test.cpp | 228 ++++++++++++++++++ kernels/test/targets.bzl | 2 +- .../executorch/build/build_variables.bzl | 2 + .../optimized/op_registration_util.bzl | 8 + 6 files changed, 428 insertions(+), 1 deletion(-) create mode 100644 kernels/optimized/cpu/op_to_copy.cpp diff --git a/kernels/optimized/cpu/op_to_copy.cpp b/kernels/optimized/cpu/op_to_copy.cpp new file mode 100644 index 00000000000..2720f46d042 --- /dev/null +++ b/kernels/optimized/cpu/op_to_copy.cpp @@ -0,0 +1,184 @@ +/* + * 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. + */ + +#include +#include +#include + +#if defined(__aarch64__) +#include + +#include +#endif + +#include +#include + +namespace torch { +namespace executor { +namespace native { + +using BFloat16 = executorch::aten::BFloat16; +using MemoryFormat = executorch::aten::MemoryFormat; +using ScalarType = executorch::aten::ScalarType; +using Tensor = executorch::aten::Tensor; + +Tensor& to_copy_out( + KernelRuntimeContext& ctx, + const Tensor& self, + bool non_blocking, + std::optional memory_format, + Tensor& out); + +namespace { + +#if defined(__aarch64__) +static_assert(sizeof(BFloat16) == sizeof(uint16_t)); +static_assert(std::is_trivially_copyable_v); + +void float_to_bfloat16_range( + const float* const input, + BFloat16* const output, + const int64_t begin, + const int64_t end) { + constexpr int64_t kVectorWidth = 8; + // Integer rounding preserves subnormals regardless of FPCR, like BFloat16. + const uint32x4_t magnitude_mask = vdupq_n_u32(0x7FFFFFFF); + const uint32x4_t infinity = vdupq_n_u32(0x7F800000); + const uint32x4_t mantissa_lsb_mask = vdupq_n_u32(1); + const uint32x4_t rounding_bias = vdupq_n_u32(0x7FFF); + const uint32x4_t canonical_nan = vdupq_n_u32(0x7FC00000); + + int64_t i = begin; +#if defined(__clang__) +#pragma unroll 4 +#elif defined(__GNUC__) +#pragma GCC unroll 4 +#endif + for (; i + kVectorWidth <= end; i += kVectorWidth) { + const uint32x4_t low_bits = vreinterpretq_u32_f32(vld1q_f32(input + i)); + const uint32x4_t high_bits = + vreinterpretq_u32_f32(vld1q_f32(input + i + 4)); + + const auto round_and_canonicalize = [&](const uint32x4_t bits) { + const uint32x4_t mantissa_lsb = + vandq_u32(vshrq_n_u32(bits, 16), mantissa_lsb_mask); + const uint32x4_t rounded = + vaddq_u32(bits, vaddq_u32(rounding_bias, mantissa_lsb)); + const uint32x4_t is_nan = + vcgtq_u32(vandq_u32(bits, magnitude_mask), infinity); + return vbslq_u32(is_nan, canonical_nan, rounded); + }; + + const uint16x4_t low = vshrn_n_u32(round_and_canonicalize(low_bits), 16); + const uint16x8_t result = + vshrn_high_n_u32(low, round_and_canonicalize(high_bits), 16); + std::memcpy(output + i, &result, sizeof(result)); + } + + for (; i < end; ++i) { + output[i] = static_cast(input[i]); + } +} + +void bfloat16_to_float_range( + const BFloat16* const input, + float* const output, + const int64_t begin, + const int64_t end) { + constexpr int64_t kVectorWidth = 8; + + int64_t i = begin; +#if defined(__clang__) +#pragma unroll 4 +#elif defined(__GNUC__) +#pragma GCC unroll 4 +#endif + for (; i + kVectorWidth <= end; i += kVectorWidth) { + uint16x8_t input_bits; + // Avoid aliasing BFloat16 storage as a NEON vector. + std::memcpy(&input_bits, input + i, sizeof(input_bits)); + const uint32x4_t low_bits = vshll_n_u16(vget_low_u16(input_bits), 16); + const uint32x4_t high_bits = vshll_high_n_u16(input_bits, 16); + vst1q_f32(output + i, vreinterpretq_f32_u32(low_bits)); + vst1q_f32(output + i + 4, vreinterpretq_f32_u32(high_bits)); + } + + for (; i < end; ++i) { + output[i] = static_cast(input[i]); + } +} + +template +bool convert_contiguous(const Tensor& self, Tensor& out) { + const auto numel = self.numel(); + if (numel == 0) { + return true; + } + + const auto* const input = self.const_data_ptr(); + auto* const output = out.mutable_data_ptr(); + const auto convert_range = [&](const auto begin, const auto end) { + if constexpr (std::is_same_v) { + float_to_bfloat16_range(input, output, begin, end); + } else { + bfloat16_to_float_range(input, output, begin, end); + } + }; + + if (numel > ::executorch::extension::internal::GRAIN_SIZE) { + return ::executorch::extension::parallel_for( + 0, numel, ::executorch::extension::internal::GRAIN_SIZE, convert_range); + } + convert_range(0, numel); + return true; +} +#endif + +} // namespace + +Tensor& opt_to_copy_out( + KernelRuntimeContext& ctx, + const Tensor& self, + bool non_blocking, + std::optional memory_format, + Tensor& out) { +#if defined(__aarch64__) + const bool float_to_bfloat16 = self.scalar_type() == ScalarType::Float && + out.scalar_type() == ScalarType::BFloat16; + const bool bfloat16_to_float = self.scalar_type() == ScalarType::BFloat16 && + out.scalar_type() == ScalarType::Float; + const bool supported_memory_format = !memory_format.has_value() || + memory_format.value() == MemoryFormat::Contiguous; + const bool can_use_optimized_kernel = + (float_to_bfloat16 || bfloat16_to_float) && !non_blocking && + supported_memory_format && tensor_is_default_dim_order(self) && + tensor_is_default_dim_order(out); + if (can_use_optimized_kernel) { + ET_KERNEL_CHECK( + ctx, + resize_tensor(out, self.sizes()) == Error::Ok, + InvalidArgument, + out); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dim_order(self, out), InvalidArgument, out); + + const bool success = float_to_bfloat16 + ? convert_contiguous(self, out) + : convert_contiguous(self, out); + ET_KERNEL_CHECK_MSG(ctx, success, Internal, out, "parallel_for failed"); + return out; + } +#endif + + return to_copy_out(ctx, self, non_blocking, memory_format, out); +} + +} // namespace native +} // namespace executor +} // namespace torch diff --git a/kernels/optimized/optimized.yaml b/kernels/optimized/optimized.yaml index 5a001afc7a0..1827411ff57 100644 --- a/kernels/optimized/optimized.yaml +++ b/kernels/optimized/optimized.yaml @@ -17,6 +17,11 @@ - arg_meta: null kernel_name: torch::executor::opt_log_softmax_out +- op: _to_copy.out + kernels: + - arg_meta: null + kernel_name: torch::executor::opt_to_copy_out + - op: add.out kernels: - arg_meta: null diff --git a/kernels/test/op_to_copy_test.cpp b/kernels/test/op_to_copy_test.cpp index 45b2b2f6020..a6fb6239390 100644 --- a/kernels/test/op_to_copy_test.cpp +++ b/kernels/test/op_to_copy_test.cpp @@ -7,7 +7,9 @@ */ #include +#include #include +#include #include #include @@ -22,6 +24,7 @@ #include using namespace ::testing; +using executorch::aten::BFloat16; using executorch::aten::MemoryFormat; using executorch::aten::ScalarType; using executorch::aten::Tensor; @@ -81,6 +84,22 @@ class OpToTest : public OperatorTest { const std::vector data_out; }; + static float float_from_bits(uint32_t bits) { + float value; + std::memcpy(&value, &bits, sizeof(value)); + return value; + } + + static uint32_t float_bits(float value) { + uint32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; + } + + static bool is_bfloat16_nan(BFloat16 value) { + return (value.x & 0x7FFF) > 0x7F80; + } + // Each test has different combination of input and output types. Therefore it // is a little bit mess if create template test case and custom data types for // both input data and output data. @@ -121,6 +140,81 @@ class OpToTest : public OperatorTest { } } + template < + typename INPUT_CTYPE, + ScalarType INPUT_DTYPE, + typename OUTPUT_CTYPE, + ScalarType OUTPUT_DTYPE> + void test_conversion_at_sizes( + const std::vector& sizes, + const std::vector& input_pattern, + const std::vector& expected_pattern) { + static_assert( + (std::is_same_v && + std::is_same_v) || + (std::is_same_v && + std::is_same_v), + "Only float/BFloat16 conversion pairs are supported"); + ASSERT_EQ(input_pattern.size(), expected_pattern.size()); + + TensorFactory tf_in; + TensorFactory tf_out; + for (const int32_t numel : sizes) { + SCOPED_TRACE(::testing::Message() << "numel=" << numel); + std::vector input_data; + std::vector expected_data; + input_data.reserve(numel); + expected_data.reserve(numel); + for (int32_t i = 0; i < numel; ++i) { + input_data.push_back(input_pattern[i % input_pattern.size()]); + expected_data.push_back(expected_pattern[i % expected_pattern.size()]); + } + + Tensor input = tf_in.make({numel}, input_data); + Tensor output = tf_out.zeros({numel}); + + Tensor& ret = op_to_copy_out( + input, + /*non_blocking=*/false, + executorch::aten::MemoryFormat::Contiguous, + output); + + EXPECT_EQ(&ret, &output); + const auto* const actual_data = ret.const_data_ptr(); + if (actual_data == nullptr) { + ADD_FAILURE() << "conversion returned a null data pointer"; + continue; + } + if constexpr (std::is_same_v) { + const bool is_aten = + torch::executor::testing::SupportedFeatures::get()->is_aten; + std::vector actual_bits; + std::vector expected_bits; + actual_bits.reserve(numel); + expected_bits.reserve(numel); + for (int32_t i = 0; i < numel; ++i) { + if (is_aten && is_bfloat16_nan(expected_data[i])) { + EXPECT_TRUE(is_bfloat16_nan(actual_data[i])) << "index=" << i; + continue; + } + actual_bits.push_back(actual_data[i].x); + expected_bits.push_back(expected_data[i].x); + } + EXPECT_EQ(actual_bits, expected_bits); + } else { + std::vector actual_bits; + std::vector expected_bits; + actual_bits.reserve(numel); + expected_bits.reserve(numel); + for (int32_t i = 0; i < numel; ++i) { + actual_bits.push_back(float_bits(actual_data[i])); + expected_bits.push_back(float_bits(expected_data[i])); + } + EXPECT_EQ(actual_bits, expected_bits); + } + } + } + template void test_runner_to_bool( std::vector test_case, @@ -360,6 +454,140 @@ TEST_F(OpToTest, NanInfSupported) { #undef TEST_KERNEL } +TEST_F(OpToTest, FloatToBFloat16RawBitsAtVectorAndGrainBoundaries) { + std::vector sizes; + for (int32_t size = 1; size <= 17; ++size) { + sizes.push_back(size); + } + sizes.insert(sizes.end(), {32767, 32768, 32769}); + + const std::vector input_pattern = { + float_from_bits(0x00000000), float_from_bits(0x80000000), + float_from_bits(0x00007FFF), float_from_bits(0x00008000), + float_from_bits(0x00008001), float_from_bits(0x00017FFF), + float_from_bits(0x00018000), float_from_bits(0x00018001), + float_from_bits(0x3F807FFF), float_from_bits(0x3F808000), + float_from_bits(0x3F808001), float_from_bits(0x3F817FFF), + float_from_bits(0x3F818000), float_from_bits(0x3F818001), + float_from_bits(0xBF807FFF), float_from_bits(0xBF808000), + float_from_bits(0xBF808001), float_from_bits(0xBF817FFF), + float_from_bits(0xBF818000), float_from_bits(0xBF818001), + float_from_bits(0x7F800000), float_from_bits(0xFF800000), + float_from_bits(0x7FC12345), float_from_bits(0xFFC12345), + float_from_bits(0x7FA12345), float_from_bits(0xFFA12345), + }; + const std::vector expected_pattern = { + BFloat16(0x0000, BFloat16::from_bits()), + BFloat16(0x8000, BFloat16::from_bits()), + BFloat16(0x0000, BFloat16::from_bits()), + BFloat16(0x0000, BFloat16::from_bits()), + BFloat16(0x0001, BFloat16::from_bits()), + BFloat16(0x0001, BFloat16::from_bits()), + BFloat16(0x0002, BFloat16::from_bits()), + BFloat16(0x0002, BFloat16::from_bits()), + BFloat16(0x3F80, BFloat16::from_bits()), + BFloat16(0x3F80, BFloat16::from_bits()), + BFloat16(0x3F81, BFloat16::from_bits()), + BFloat16(0x3F81, BFloat16::from_bits()), + BFloat16(0x3F82, BFloat16::from_bits()), + BFloat16(0x3F82, BFloat16::from_bits()), + BFloat16(0xBF80, BFloat16::from_bits()), + BFloat16(0xBF80, BFloat16::from_bits()), + BFloat16(0xBF81, BFloat16::from_bits()), + BFloat16(0xBF81, BFloat16::from_bits()), + BFloat16(0xBF82, BFloat16::from_bits()), + BFloat16(0xBF82, BFloat16::from_bits()), + BFloat16(0x7F80, BFloat16::from_bits()), + BFloat16(0xFF80, BFloat16::from_bits()), + BFloat16(0x7FC0, BFloat16::from_bits()), + BFloat16(0x7FC0, BFloat16::from_bits()), + BFloat16(0x7FC0, BFloat16::from_bits()), + BFloat16(0x7FC0, BFloat16::from_bits()), + }; + + test_conversion_at_sizes< + float, + ScalarType::Float, + BFloat16, + ScalarType::BFloat16>(sizes, input_pattern, expected_pattern); +} + +#if defined(__aarch64__) +TEST_F(OpToTest, FloatToBFloat16SubnormalsIgnoreFlushToZero) { + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen conversion may flush subnormals to zero"); + + struct RestoreFpcr { + uint64_t value{}; + RestoreFpcr() { + asm volatile("mrs %0, fpcr" : "=r"(value)); + } + ~RestoreFpcr() { + asm volatile("msr fpcr, %0" : : "r"(value) : "memory"); + } + } original_fpcr; + + constexpr uint64_t kFlushToZero = uint64_t{1} << 24; + for (const bool flush_to_zero : {false, true}) { + SCOPED_TRACE(::testing::Message() << "flush_to_zero=" << flush_to_zero); + const uint64_t fpcr = flush_to_zero ? original_fpcr.value | kFlushToZero + : original_fpcr.value & ~kFlushToZero; + asm volatile("msr fpcr, %0" : : "r"(fpcr) : "memory"); + test_conversion_at_sizes< + float, + ScalarType::Float, + BFloat16, + ScalarType::BFloat16>( + {1, 7, 8, 9, 15, 16, 17}, + {float_from_bits(0x00018000), float_from_bits(0x80018000)}, + {BFloat16(0x0002, BFloat16::from_bits()), + BFloat16(0x8002, BFloat16::from_bits())}); + } +} +#endif + +TEST_F(OpToTest, BFloat16ToFloatRawBitsAtVectorAndGrainBoundaries) { + std::vector sizes; + for (int32_t size = 1; size <= 17; ++size) { + sizes.push_back(size); + } + sizes.insert(sizes.end(), {32767, 32768, 32769}); + + const std::vector input_pattern = { + BFloat16(0x0000, BFloat16::from_bits()), + BFloat16(0x8000, BFloat16::from_bits()), + BFloat16(0x3F80, BFloat16::from_bits()), + BFloat16(0x3F81, BFloat16::from_bits()), + BFloat16(0x3F82, BFloat16::from_bits()), + BFloat16(0x7F80, BFloat16::from_bits()), + BFloat16(0xFF80, BFloat16::from_bits()), + BFloat16(0x7FC1, BFloat16::from_bits()), + BFloat16(0xFFC1, BFloat16::from_bits()), + BFloat16(0x7FA1, BFloat16::from_bits()), + BFloat16(0xFFA1, BFloat16::from_bits()), + }; + const std::vector expected_pattern = { + float_from_bits(0x00000000), + float_from_bits(0x80000000), + float_from_bits(0x3F800000), + float_from_bits(0x3F810000), + float_from_bits(0x3F820000), + float_from_bits(0x7F800000), + float_from_bits(0xFF800000), + float_from_bits(0x7FC10000), + float_from_bits(0xFFC10000), + float_from_bits(0x7FA10000), + float_from_bits(0xFFA10000), + }; + + test_conversion_at_sizes< + BFloat16, + ScalarType::BFloat16, + float, + ScalarType::Float>(sizes, input_pattern, expected_pattern); +} + TEST_F(OpToTest, HardcodeFloatConvertInt) { // Hardcode input and output generated from core PyTorch // clang-format off diff --git a/kernels/test/targets.bzl b/kernels/test/targets.bzl index 837c7327c4f..9084dd2b16d 100644 --- a/kernels/test/targets.bzl +++ b/kernels/test/targets.bzl @@ -352,7 +352,7 @@ def define_common_targets(): _common_op_test("op_t_copy_test", ["aten", "portable"]) _common_op_test("op_tan_test", ["aten", "portable"]) _common_op_test("op_tanh_test", ["aten", "portable"]) - _common_op_test("op_to_copy_test", ["aten", "portable"]) + _common_op_test("op_to_copy_test", ["aten", "portable", "optimized"]) _common_op_test("op_topk_test", ["aten", "portable"]) _common_op_test("op_transpose_copy_test", ["aten", "portable"]) _common_op_test("op_tril_test", ["aten", "portable"]) diff --git a/shim_et/xplat/executorch/build/build_variables.bzl b/shim_et/xplat/executorch/build/build_variables.bzl index c75436d3dc7..8dc8944b052 100644 --- a/shim_et/xplat/executorch/build/build_variables.bzl +++ b/shim_et/xplat/executorch/build/build_variables.bzl @@ -278,6 +278,7 @@ OPTIMIZED_KERNELS_SRCS = [ "kernels/optimized/cpu/op_native_layer_norm.cpp", "kernels/optimized/cpu/op_sub.cpp", "kernels/optimized/cpu/op_sum.cpp", + "kernels/optimized/cpu/op_to_copy.cpp", "kernels/optimized/cpu/op_where.cpp", ] @@ -320,6 +321,7 @@ OPTIMIZED_NATIVE_CPU_OPS_SRCS = [ "kernels/optimized/cpu/op_mul.cpp", "kernels/optimized/cpu/op_native_layer_norm.cpp", "kernels/optimized/cpu/op_sub.cpp", + "kernels/optimized/cpu/op_to_copy.cpp", "kernels/optimized/cpu/op_where.cpp", ] diff --git a/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl index fba89adde64..25041897059 100644 --- a/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl @@ -327,6 +327,14 @@ OPTIMIZED_ATEN_OPS = ( "//executorch/kernels/portable/cpu/util:reduce_util", ], ), + op_target( + name = "op_to_copy", + deps = [ + "//executorch/extension/threadpool:threadpool", + "//executorch/kernels/portable/cpu:op_to_copy", + "//executorch/kernels/portable/cpu/util:copy_ops_util", + ], + ), op_target( name = "op_where", deps = [