Open kevinstephano opened 1 year ago
I looked at the NHWC repro on A100. Looks like it's a heuristics issue. We are using the block-parallel outer normalization scheduler, and I don't think we've done any specific tuning for a long time.
This case is probably one of the cases we would find with Liqiang's benchmarking effort. I'd expect we would find multiple performance issues and need to think about priority.
Christian's suggestions.
[12:12] Christian Sarofeen Even just put in a really dirty hack of "if exactly these sizes for this normalization on this device do this". [12:13] Christian Sarofeen Even more on this, we could even skip the segmentation and heuristics directly and only do this for instance norm 3d [12:13] Christian Sarofeen We don't need generic fusion support, we just need an explicit nvFuser instance norm, and we can accumulate some small technical debt here.
Sizes at types that are cared about:
Here are the cudnn layer files describing the shapes we are using https://gitlab-master.nvidia.com/dl/JoC/model_profiling/-/tree/main/cuDNN/mlperf/mxnet/image_segmentation. in the wild, only non-spatial cases (1x8x7, 1x8x1) are used. If that's hard to read/parse, three largest layers are (DxHxWxC):
- 128x128x128x32
- 64x64x64x64
- 32x32x32x128
dtypes are the standard set: fp16, amp, tf32, fp32, ordered according to priority
The dtypes are fp16 and fp32; amp and tf32 doesn't mean anything to us.
@jacobhinkle I also think we want to pull the logic you wrote in apex directly into an nvFuser contrib or something of the sort. We don't need the flexibility you had and can basically hard code the operator just for instance norm 3D at this time.
It looks like the NHWC, with an outer reduction, takes double the time of the NCHW layout. Two things that look suspect about the NHWC kernel is that it uses 128 register compared to 36 registers for the NCHW layout suggesting perhaps a non-optimal heuristic choice. Additionally, the NHWC quantizes its blocks with 400 threads that leads to 12.5 warps.
NCHW Layout Repro:
Perf Result: (119 us)
NHWC Layout Repro:
Scheduler Params:
Perf result: (225 us)