[GraphOptimizer] Fix wrong layout assumption in OptimizeReduceMean?

[shajrawi]

I see the following code in the optimization:

      // In Glow, AvgPool expects NHWC.
      auto *TR1 = F->createTranspose(
          RM->getName().str() + ".transposeNCHW2NHWC", in, NCHW2NHWC);
      auto *AP = F->createAvgPool(RM->getName().str() + ".avgPool", TR1,
                                  kernels, strides, pads);
      auto *TR2 = F->createTranspose(
          RM->getName().str() + ".transposeNHWC2NCHW", AP, NHWC2NCHW);

This looks like a bug to me, the canonical representation in Glow is NHWC. Why do we expect the input to be in NCHW format and doing the two transposes?

[shajrawi]

cc @vuzelac-cadence

[dspmihai]

Will this code be ever executed? I see in lowerBatchedReduceMeanNode that ReduceMean can only have one axis. @jfix71, @mciprian13 , any idea how it's better to approach this again? I am thinking of lowering ReduceMean to AvgPool in Lower.cpp, if possible. Or maybe lower to multiple ReduceMean, then find the pattern in OptimizeReduceMean (possible error-prone after future optimizations)?

[vuzelac-cadence]

@dspmihai, lowering is backend controlled. For example, we do use this optimization.