Skip to content

Commit bc40309

Browse files
authored
[CPU] Fix FP16/BF16 convolution and layer normalization (pytorch#23153)
### Summary Portable convolution accumulates FP16/BF16 products in the output dtype, so long reductions lose contributions: summing 3072 ones can return 2048 in FP16 or 256 in BF16. Both ordinary and transposed convolution now accumulate products and bias in FP32 and narrow once per output. Transposed convolution visits bounded input/kernel ranges for 8x32 output tiles and reuses weights across contributing inputs. Each tile uses a fixed stack buffer (1 KiB for FP32 accumulators), with one cast per output after bias. The tiled helper and dtype wrapper stay separate from scalar-type dispatch to preserve ordinary/depthwise convolution performance. Optimized layer normalization routes FP16/BF16 through the shared scalar helper because RowwiseMoments does not implement their vector loads. The helper now uses corrected two-pass moments to prevent cancellation and NaNs on large-mean rows, accumulating in FP32 for Half/BFloat16/Float and FP64 for Double. Both kernel wrappers preserve epsilon in compute precision. This also fixes the existing portable and small-row FP32/FP64 cases. The large-row FP32/FP64 SIMD path is unchanged; reduced-precision SIMD support remains a follow-up in pytorch#23159. This PR contains only CPU kernel fixes and unit tests. It is a prerequisite for the BF16 CPU path in pytorch#23152. ### Test plan Added FP16/BF16 tests for long ordinary and transposed convolution reductions, mixed FP32 bias, and grouped transposed convolution with stride, dilation, output padding, and both layouts. The grouped tests use 1024 input channels per group and match both the PyTorch FP32 reference and native FP16/BF16 results. The existing FP32 geometry case remains a regression guard. Built the same test sources against the original kernels at 00e5030 and the fixed kernels. Before: eight convolution tests and both reduced-precision optimized layer-normalization tests fail. After: the convolution, portable layer-normalization, and optimized layer-normalization suites pass. Also built the layer-normalization suite with USE_ATEN_LIB: BFloat16LargeRows fails with the original exact-mean assertion and passes with the tolerant assertion. The current convolution suite passes 29 tests with one expected skip. Four added geometry tests cover coprime and shared-factor stride/dilation, full and partial tiles, strides larger than a tile, and both layouts; all four also pass in ATen mode. Compared 1,711 convolution cases against PyTorch, including 512 generated from the FACTO convolution specification and 768 additional stride/dilation cases. All numerical checks and UBSan checks pass. Reduced-precision references run PyTorch convolution in FP32 on already-rounded inputs and cast the result once. All outputs are bitwise identical to the previous PR head, 115d394. The earlier Nemotron BF16 production regression passes across all four presets. Added six layer-normalization regressions covering large-mean and constant rows in all four floating dtypes, odd tails, affine normalization, and tiny epsilon values. Each fails against both portable and optimized kernels at the previous PR head, 21a4ad2. With the fix, each native suite passes 16 tests with one expected skip; the ATen suite passes 14 with three expected skips. Both native suites also pass under UBSan. An 800-case numerical grid covers random, large-mean, and constant inputs, widths from 1 to 4096, all combinations of optional weight/bias, and tiny epsilon. Scalar results match a PyTorch FP64 reference within dtype-appropriate tolerances, with no unexpected NaNs or infinities. Existing optimized FP32/FP64 SIMD outputs remain bitwise identical to 21a4ad2. The grid also passes under UBSan. These checks ran on macOS arm64. `lintrunner` and `git diff --check` pass. Supplemental test harnesses and results remain local. Authored with AI assistance from OpenAI Codex.
1 parent 3535c92 commit bc40309

6 files changed

Lines changed: 688 additions & 151 deletions

File tree

‎kernels/optimized/cpu/op_native_layer_norm.cpp‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#include <executorch/runtime/kernel/kernel_includes.h>
1010
#include <cmath>
1111
#include <tuple>
12+
#include <type_traits>
1213

1314
#include <ATen/cpu/vec/functional.h>
1415
#include <ATen/cpu/vec/vec.h>
@@ -29,7 +30,7 @@ void layer_norm(
2930
IntArrayRef normalized_shape,
3031
const optional<Tensor>& weight,
3132
const optional<Tensor>& bias,
32-
CTYPE eps,
33+
double eps,
3334
Tensor& out,
3435
Tensor& mean,
3536
Tensor& rstd) {
@@ -72,11 +73,10 @@ void layer_norm(
7273
const bool gamma_null = gamma_data == nullptr;
7374
const bool beta_null = beta_data == nullptr;
7475

75-
// For small normalized dimensions, fall back to the portable scalar
76-
// implementation since SIMD vectorization setup/tail-handling overhead
77-
// exceeds the benefit for small N.
76+
// RowwiseMoments lacks reduced-precision vector loads. The scalar path also
77+
// avoids SIMD setup/tail-handling overhead for small normalized dimensions.
7878
constexpr size_t kSmallNThreshold = 256;
79-
if (N < kSmallNThreshold) {
79+
if (N < kSmallNThreshold || !std::is_same_v<CTYPE, acc_t<CTYPE>>) {
8080
layer_norm_scalar<CTYPE>(
8181
input_data,
8282
gamma_data,
@@ -97,7 +97,8 @@ void layer_norm(
9797
acc_t<CTYPE> mean_val;
9898
acc_t<CTYPE> rstd_val;
9999
std::tie(mean_val, rstd_val) = RowwiseMoments(src_ptr, N);
100-
rstd_val = CTYPE(1) / std::sqrt(rstd_val + eps);
100+
rstd_val =
101+
acc_t<CTYPE>(1) / std::sqrt(rstd_val + static_cast<acc_t<CTYPE>>(eps));
101102

102103
const acc_t<CTYPE> scale = rstd_val;
103104
const acc_t<CTYPE> offset = -rstd_val * mean_val;

0 commit comments

Comments
 (0)