Repository navigation
Commit bc40309
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
- portable/cpu
- util
- test
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
9 | 9 | | |
10 | 10 | | |
11 | 11 | | |
| 12 | + | |
12 | 13 | | |
13 | 14 | | |
14 | 15 | | |
| |||
29 | 30 | | |
30 | 31 | | |
31 | 32 | | |
32 | | - | |
| 33 | + | |
33 | 34 | | |
34 | 35 | | |
35 | 36 | | |
| |||
72 | 73 | | |
73 | 74 | | |
74 | 75 | | |
75 | | - | |
76 | | - | |
77 | | - | |
| 76 | + | |
| 77 | + | |
78 | 78 | | |
79 | | - | |
| 79 | + | |
80 | 80 | | |
81 | 81 | | |
82 | 82 | | |
| |||
97 | 97 | | |
98 | 98 | | |
99 | 99 | | |
100 | | - | |
| 100 | + | |
| 101 | + | |
101 | 102 | | |
102 | 103 | | |
103 | 104 | | |
| |||
0 commit comments