mlas/arm64: add BF16 fast-math conv kernels for NCHW/NCHWc paths
Description
Add Arm64 BF16 fast-math convolution support in MLAS:
- direct NCHW conv
- depthwise 3x3 NCHWc conv
- pointwise 1x1 NCHWc conv
This change adds new AArch64 BF16 asm kernels, wires them into MLAS platform dispatch, keeps accumulated pointwise batches on the custom BF16 path instead of falling back to generic SBGEMM, and adds the required BF16 build flags.
The new paths are only used when Arm64 BF16 fast-math is enabled via the existing session option. Baseline FP32 behavior is unchanged.
Performance
Individual convolution improvements when running on c8g AWS instance where in columns base is FP32 execution, fast-math when enabled without this PR and PR is fast-math with this change:
| Type | Shape | fast-math vs base | PR w/ fast-math vs base | PR w/ fast-math vs fast-math |
|---|---|---|---|---|
| depthwise | N1 IC32 OC32 H112xW112->112x112 K3x3 S1x1 D1x1 P1/1/1/1 G32 | 0.991x | 1.047x | 1.057x |
| depthwise | N1 IC96 OC96 H112xW112->56x56 K3x3 S2x2 D1x1 P1/1/1/1 G96 | 1.015x | 1.015x | 1.000x |
| depthwise | N1 IC144 OC144 H56xW56->28x28 K3x3 S2x2 D1x1 P1/1/1/1 G144 | 1.020x | 1.004x | 0.984x |
| depthwise | N1 IC144 OC144 H56xW56->56x56 K3x3 S1x1 D1x1 P1/1/1/1 G144 | 1.034x | 1.138x | 1.101x |
| depthwise | N1 IC192 OC192 H28xW28->28x28 K3x3 S1x1 D1x1 P1/1/1/1 G192 | 0.997x | 1.033x | 1.037x |
| depthwise | N1 IC384 OC384 H28xW28->14x14 K3x3 S2x2 D1x1 P1/1/1/1 G384 | 1.016x | 1.021x | 1.005x |
| depthwise | N1 IC384 OC384 H28xW28->28x28 K3x3 S1x1 D1x1 P1/1/1/1 G384 | 1.011x | 1.090x | 1.077x |
| depthwise | N1 IC576 OC576 H14xW14->7x7 K3x3 S2x2 D1x1 P1/1/1/1 G576 | 1.029x | 0.995x | 0.967x |
| depthwise | N1 IC576 OC576 H14xW14->14x14 K3x3 S1x1 D1x1 P1/1/1/1 G576 | 1.025x | 1.006x | 0.982x |
| depthwise | N1 IC960 OC960 H7xW7->7x7 K3x3 S1x1 D1x1 P1/1/1/1 G960 | 1.002x | 0.941x | 0.939x |
| nchw | N1 IC3 OC32 H224xW224->112x112 K3x3 S2x2 D1x1 P1/1/1/1 G1 | 1.001x | 1.058x | 1.058x |
| pointwise | N1 IC16 OC96 H112xW112->112x112 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.213x | 1.328x | 1.095x |
| pointwise | N1 IC32 OC16 H112xW112->112x112 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.020x | 1.019x | 0.998x |
| pointwise | N1 IC32 OC32 H112xW112->112x112 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.118x | 1.196x | 1.069x |
| pointwise | N1 IC32 OC144 H56xW56->56x56 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.220x | 1.528x | 1.252x |
| pointwise | N1 IC32 OC192 H28xW28->28x28 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.199x | 1.418x | 1.183x |
| pointwise | N1 IC64 OC384 H28xW28->28x28 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.294x | 1.938x | 1.497x |
| pointwise | N1 IC96 OC32 H56xW56->56x56 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.080x | 1.426x | 1.320x |
| pointwise | N1 IC96 OC576 H14xW14->14x14 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.280x | 1.961x | 1.532x |
| pointwise | N1 IC144 OC32 H28xW28->28x28 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.132x | 1.351x | 1.193x |
| pointwise | N1 IC144 OC32 H56xW56->56x56 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.073x | 1.374x | 1.281x |
| pointwise | N1 IC160 OC960 H7xW7->7x7 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.133x | 1.744x | 1.539x |
| pointwise | N1 IC192 OC32 H28xW28->28x28 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.166x | 1.411x | 1.210x |
| pointwise | N1 IC192 OC64 H28xW28->28x28 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.212x | 1.763x | 1.454x |
| pointwise | N1 IC320 OC1280 H7xW7->7x7 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.136x | 2.059x | 1.812x |
| pointwise | N1 IC384 OC64 H28xW28->28x28 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.256x | 1.904x | 1.516x |
| pointwise | N1 IC384 OC96 H14xW14->14x14 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.206x | 1.929x | 1.600x |
| pointwise | N1 IC576 OC96 H14xW14->14x14 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.250x | 2.055x | 1.644x |
| pointwise | N1 IC576 OC160 H7xW7->7x7 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 0.902x | 1.423x | 1.577x |
| pointwise | N1 IC960 OC160 H7xW7->7x7 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 0.915x | 1.527x | 1.668x |
| pointwise | N1 IC960 OC320 H7xW7->7x7 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 1.020x | 1.756x | 1.723x |
| pointwise | N1 IC1280 OC1008 H1xW1->1x1 K1x1 S1x1 D1x1 P0/0/0/0 G1 | 0.747x | 1.149x | 1.538x |
When running the full models the performance improvements are on c8g (AWS Graviton 4) and Standard_D32plds_v6 (Azure Cobalt-100) when running MobileNet v2.7 with 8 threads are:
| Instance | PR w/ fast-math vs base | PR w/ fast-math vs fast-mat |
|---|---|---|
c8g | 1.892x | 1.647x |
Standard_D32plds_v6 | 2.884x | 1.692x |
(cc: @Rohanjames1997 @snadampal)