Repository: pytorch
Issue: Replace native NHWC BN kernel
Labels: triaged,module: memory format

## 🚀 Feature

In PR #59981 we are discussing removing BN nhwc kernel with TensorIterator implementations. This PR is used to track potential performance regression.

## Additional context

Comparing the result between NHWC kernel and non-NHWC (I disabled NHWC kernel by excluding them in the cpp side).
We found that in general cases where spatial dimensions are tiny, NHWC kernel actually gives perf regression. However, for larger size tensor, it runs a few times faster.

Looking at the profile, I think kernel `at::native::batch_norm_backward_kernel` is the bottleneck.

Here's the script:

```
import torch

from torch import nn
from typing import Callable, Any, Optional, Tuple, List

import time
import itertools

batches = [8, 16, 32]
channels = [8, 16, 32, 64, 128, 256, 512, 768]
sizes = [1, 7, 8, 17, 32, 64]

WARM_UP = 10
ITER = 500

torch.backends.cudnn.enabled=False
dtype=torch.float

for b, c, hw in itertools.product(batches, channels, sizes):
  input_nhwc = torch.randn(b, c, hw, hw, device="cuda").to(dtype=dtype, memory_format=torch.channels_last).requires_grad_(True)
  grad_nhwc = torch.randn(b, c, hw, hw, device="cuda").to(dtype=dtype, memory_format=torch.channels_last)
  my_module = nn.BatchNorm2d(c).cuda().to(dtype=dtype)

  def func_nhwc():
    o0 = my_module(input_nhwc)
    o0.backward(grad_nhwc)
    #torch.cuda.synchronize()

  for i in range(0, WARM_UP):
    func_nhwc()
  torch.cuda.synchronize()
  start = time.time()
  for i in range(0, ITER):
    func_nhwc()
  torch.cuda.synchronize()
  end = time.time()
  print("sizes, {}x{}x{}, time, {} ms/iter".format(b, c, hw, (end-start)*1000.0/ITER))
```

Here's data. measuring ms per iteration.

&nbsp; | nhwc_float | non_nhwc_float | speedup
-- | -- | -- | --
8x8x1 | 0.296369075775146 | 0.258503437042236 | 0.872234852324341
8x8x7 | 0.306492328643799 | 0.280899047851562 | 0.916496178206207
8x8x8 | 0.313601493835449 | 0.27117919921875 | 0.864725470220627
8x8x17 | 0.305299758911133 | 0.27459716796875 | 0.899434604691844
8x8x32 | 0.355751991271973 | 0.27347469329834 | 0.768722874383768
8x8x64 | 0.348571300506592 | 0.275537967681885 | 0.790478066557502
8x16x1 | 0.305832386016846 | 0.259871482849121 | 0.849718652212349
8x16x7 | 0.297619342803955 | 0.280623435974121 | 0.942893809690893
8x16x8 | 0.298600196838379 | 0.274928092956543 | 0.920723080116894
8x16x17 | 0.298842906951904 | 0.276485919952393 | 0.925188162477921
8x16x32 | 0.347675323486328 | 0.284414768218994 | 0.818047037008591
8x16x64 | 0.346132755279541 | 0.276742458343506 | 0.799526927522376
8x32x1 | 0.302460193634033 | 0.261518478393555 | 0.864637674304948
8x32x7 | 0.311790943145752 | 0.285078048706055 | 0.914324341271164
8x32x8 | 0.294690132141113 | 0.288757801055908 | 0.979869257779
8x32x17 | 0.34625244140625 | 0.286839008331299 | 0.828410067424643
8x32x32 | 0.344767093658447 | 0.296492099761963 | 0.859977953858006
8x32x64 | 0.346677780151367 | 0.29764986038208 | 0.8585778420876
8x64x1 | 0.301368713378906 | 0.273605823516846 | 0.907877332219437
8x64x7 | 0.297780990600586 | 0.291343212127686 | 0.978380827936947
8x64x8 | 0.298311710357666 | 0.285955905914307 | 0.958580893694903
8x64x17 | 0.3441481590271 | 0.285999298095703 | 0.831035385760069
8x64x32 | 0.354624271392822 | 0.292672634124756 | 0.825303448563334
8x64x64 | 0.351122856140137 | 0.2923264503479 | 0.832547483696789
8x128x1 | 0.294944286346436 | 0.274373531341553 | 0.930255455158332
8x128x7 | 0.302114009857178 | 0.290267944335937 | 0.960789420103886
8x128x8 | 0.308194637298584 | 0.281272888183594 | 0.912646925491737
8x128x17 | 0.354531288146973 | 0.270200729370117 | 0.762135073556905
8x128x32 | 0.341504573822021 | 0.260371208190918 | 0.762423780381381
8x128x64 | 0.330519199371338 | 0.488327980041504 | 1.47745722781105
8x256x1 | 0.288928031921387 | 0.221806049346924 | 0.767686153094602
8x256x7 | 0.292155742645264 | 0.261454582214355 | 0.894915088257616
8x256x8 | 0.288446426391602 | 0.2607741355896 | 0.904064365961555
8x256x17 | 0.326685905456543 | 0.263041496276855 | 0.805181649662098
8x256x32 | 0.327075958251953 | 0.275776386260986 | 0.84315700773259
8x256x64 | 0.465546131134033 | 1.46136283874512 | 3.13902906933273
8x512x1 | 0.247006893157959 | 0.22305965423584 | 0.903050321325223
8x512x7 | 0.289066791534424 | 0.284266471862793 | 0.98339373524662
8x512x8 | 0.289161682128906 | 0.289201736450195 | 1.00013851877259
8x512x17 | 0.329724311828613 | 0.295010089874268 | 0.894717432991748
8x512x32 | 0.339834690093994 | 0.47899866104126 | 1.40950490048198
8x512x64 | 0.883801937103272 | 4.02801322937012 | 4.55759719487857
8x768x1 | 0.239550590515137 | 0.225006103515625 | 0.939284278246884
8x768x7 | 0.304868221282959 | 0.29229736328125 | 0.958766256617998
8x768x8 | 0.304521083831787 | 0.292906761169434 | 0.961860366066577
8x768x17 | 0.342549324035645 | 0.304738998413086 | 0.889620784600865
8x768x32 | 0.352364540100098 | 0.862393379211426 | 2.4474465534087
8x768x64 | 1.30678033828735 | 7.66752052307129 | 5.86748996630928
16x8x1 | 0.234663486480713 | 0.226991653442383 | 0.967307086614173
16x8x7 | 0.291863918304443 | 0.289965152740479 | 0.993494346354988
16x8x8 | 0.290300846099854 | 0.306094646453857 | 1.0544049408267
16x8x17 | 0.292605876922607 | 0.293336391448975 | 1.00249658186654
16x8x32 | 0.32784366607666 | 0.308577060699463 | 0.941232339157981
16x8x64 | 0.329196929931641 | 0.299880981445312 | 0.910947078113955
16x16x1 | 0.293085098266602 | 0.277911186218262 | 0.948226941123642
16x16x7 | 0.308955669403076 | 0.295466899871826 | 0.956340760610377
16x16x8 | 0.306412696838379 | 0.295915126800537 | 0.965740420856715
16x16x17 | 0.347569942474365 | 0.294765949249268 | 0.84807664077859
16x16x32 | 0.346548080444336 | 0.295411586761475 | 0.852440407064742
16x16x64 | 0.350766181945801 | 0.297928333282471 | 0.84936447302239
16x32x1 | 0.306893348693848 | 0.27112865447998 | 0.883462139645309
16x32x7 | 0.306521892547607 | 0.293563842773437 | 0.957725532533839
16x32x8 | 0.305248260498047 | 0.294153690338135 | 0.963653944688137
16x32x17 | 0.344708919525146 | 0.295713424682617 | 0.857864151267037
16x32x32 | 0.351790904998779 | 0.294126987457275 | 0.836084683480649
16x32x64 | 0.348780632019043 | 0.290832996368408 | 0.833856497950635
16x64x1 | 0.307173728942871 | 0.279016017913818 | 0.908332945249073
16x64x7 | 0.306697845458984 | 0.292588710784912 | 0.95399662931131
16x64x8 | 0.298262596130371 | 0.29512882232666 | 0.989493239030411
16x64x17 | 0.347235679626465 | 0.29587459564209 | 0.852085810883184
16x64x32 | 0.350239753723145 | 0.316976547241211 | 0.905027324487479
16x64x64 | 0.351851940155029 | 0.445825576782227 | 1.26708290022727
16x128x1 | 0.31151008605957 | 0.218288421630859 | 0.700742709143343
16x128x7 | 0.309917449951172 | 0.275523662567139 | 0.88902274657509
16x128x8 | 0.303652763366699 | 0.284367084503174 | 0.936487721535287
16x128x17 | 0.347876071929932 | 0.297348976135254 | 0.854755472216397
16x128x32 | 0.350400924682617 | 0.288174152374268 | 0.822412648126677
16x128x64 | 0.469685554504395 | 1.01043033599854 | 2.1512910633684
16x256x1 | 0.271766662597656 | 0.226653575897217 | 0.834000659723197
16x256x7 | 0.314548015594482 | 0.253268718719482 | 0.805183012332204
16x256x8 | 0.309408187866211 | 0.268337726593018 | 0.867261233271072
16x256x17 | 0.353528022766113 | 0.246078968048096 | 0.696066371550117
16x256x32 | 0.352646827697754 | 0.530220031738281 | 1.50354402788708
16x256x64 | 0.888144493103027 | 4.18027448654175 | 4.70675044320387
16x512x1 | 0.260415077209473 | 0.243691444396973 | 0.935780858037464
16x512x7 | 0.310879707336426 | 0.304919242858887 | 0.980827103420138
16x512x8 | 0.310995101928711 | 0.304154872894287 | 0.978005348019945
16x512x17 | 0.353450298309326 | 0.297888278961182 | 0.842801039880524
16x512x32 | 0.458423137664795 | 1.41802453994751 | 3.09326563918854
16x512x64 | 1.72479009628296 | 10.1610007286072 | 5.89115205989693
16x768x1 | 0.230637073516846 | 0.22266149520874 | 0.965419356972881
16x768x7 | 0.286874294281006 | 0.29221773147583 | 1.01862640641336
16x768x8 | 0.29481840133667 | 0.304117202758789 | 1.03154077689846
16x768x17 | 0.334246635437012 | 0.390329837799072 | 1.16778987853904
16x768x32 | 0.672853469848633 | 2.8329815864563 | 4.21039901465265
16x768x64 | 2.5929741859436 | 17.8364109992981 | 6.87874607313427
32x8x1 | 0.199938774108887 | 0.228226184844971 | 1.14148036498753
32x8x7 | 0.280766487121582 | 0.289992809295654 | 1.03286119461286
32x8x8 | 0.288478374481201 | 0.288513660430908 | 1.00012231748661
32x8x17 | 0.328126430511475 | 0.286970615386963 | 0.874573300723263
32x8x32 | 0.327083110809326 | 0.298518180847168 | 0.912667670637356
32x8x64 | 0.330361366271973 | 0.327936172485352 | 0.992658966712759
32x16x1 | 0.288805961608887 | 0.252796173095703 | 0.875314940479139
32x16x7 | 0.29206371307373 | 0.294923305511475 | 1.00979098843759
32x16x8 | 0.290713787078857 | 0.288819313049316 | 0.993483370539192
32x16x17 | 0.327481269836426 | 0.291618824005127 | 0.890490085588065
32x16x32 | 0.329592704772949 | 0.291948795318604 | 0.8857865817137
32x16x64 | 0.331625461578369 | 0.362840175628662 | 1.09412640966024
32x32x1 | 0.290622234344482 | 0.246495723724365 | 0.848165400284505
32x32x7 | 0.292683124542236 | 0.289180278778076 | 0.988031951723767
32x32x8 | 0.329047203063965 | 0.289097785949707 | 0.87859061939362
32x32x17 | 0.330667972564697 | 0.293543815612793 | 0.88772980744411
32x32x32 | 0.331005573272705 | 0.293840885162354 | 0.887721866000931
32x32x64 | 0.339998722076416 | 0.474749565124512 | 1.39632749860104
32x64x1 | 0.306783199310303 | 0.221792221069336 | 0.722960780016507
32x64x7 | 0.306890487670898 | 0.299352169036865 | 0.975436453924512
32x64x8 | 0.347555637359619 | 0.308742523193359 | 0.88832546506475
32x64x17 | 0.349190235137939 | 0.300435066223145 | 0.860376482476564
32x64x32 | 0.35122013092041 | 0.295378684997559 | 0.841007274336716
32x64x64 | 0.467539310455322 | 0.98347806930542 | 2.10351952726208
32x128x1 | 0.268299102783203 | 0.220008850097656 | 0.820013364992251
32x128x7 | 0.310078620910645 | 0.306240558624268 | 0.987622293097456
32x128x8 | 0.347151756286621 | 0.299485683441162 | 0.862693845033859
32x128x17 | 0.351085662841797 | 0.298224925994873 | 0.849436355734232
32x128x32 | 0.351006984710693 | 0.490714550018311 | 1.39801933121863
32x128x64 | 0.899456024169922 | 2.10693120956421 | 2.3424504955744
32x256x1 | 0.248291492462158 | 0.230325222015381 | 0.927640410597171
32x256x7 | 0.307347297668457 | 0.306653022766113 | 0.997741073672648
32x256x8 | 0.352444171905518 | 0.299785614013672 | 0.85059035702834
32x256x17 | 0.356542110443115 | 0.304477691650391 | 0.853973998392453
32x256x32 | 0.460453987121582 | 1.39680433273315 | 3.033537273648
32x256x64 | 1.74108743667603 | 10.5460143089294 | 6.05714227027175
32x512x1 | 0.248486042022705 | 0.226709365844727 | 0.91236257779023
32x512x7 | 0.303464412689209 | 0.301892280578613 | 0.99481938558573
32x512x8 | 0.348631858825684 | 0.304160594940186 | 0.872440619640175
32x512x17 | 0.352962493896484 | 0.528259754180908 | 1.49664557372443
32x512x32 | 0.882181167602539 | 3.77037191390991 | 4.27392020185204
32x512x64 | 3.4178466796875 | 23.9915480613708 | 7.01949218610486
32x768x1 | 0.224312782287598 | 0.228887557983398 | 1.02039462775659
32x768x7 | 0.286787509918213 | 0.304922103881836 | 1.06323355563426
32x768x8 | 0.331912040710449 | 0.302271366119385 | 0.910697199994253
32x768x17 | 0.393331050872803 | 0.946650981903076 | 2.40675375056827
32x768x32 | 1.30361366271973 | 7.7111029624939 | 5.91517501159526
32x768x64 | 5.10943365097046 | 44.0090584754944 | 8.61329483496386



cc @VitalyFedyunin @jamesr66a @ppwwyyxx