🔗 Helpful links

• 🧪  See artifacts and rendered test results at hud.pytorch.org/pr/50247
• 📄  Preview docs built from this PR
• 📄  Preview C++ docs built from this PR
• 🔧  Opt-in to CIFlow to control what jobs run on your PRs

💊 CI failures summary and remediations

As of commit 0c7bd64 (more details on the Dr. CI page):

💚 💚 Looks good so far! There are no failures yet. 💚 💚

This patch adds native support of channels last memory format for nn.ChannelShuffle.
Currently nn.ChannelShuffle cannot propagate input memory format, channels last would be treated as non-contiguous NCHW tensor.

From performance perspective, ChannelShuffle favors NCHW over NHWC, since on NHWC it would end up with matrix transpose. I made fast paths with vectorized transpose for groups = 2, 4, 8, 16, ... so that NHWC has similar perf as NCHW.

Performance Benchmarking

I tested pytorch operator benchmark for ChannelShuffle, Intel(R) Xeon(R) Gold 6248 CPU @ 2.50GHz, single core.

Unit: (us)
Jemalloc is applied, the benchmark has quite a large memory footprint.

The benchmark naming is self explanatory.

short tag:

python -m pt.channel_shuffle_test --omp_num_threads 1 --mkl_num_threads 1

long tag:

python -m pt.channel_shuffle_test --omp_num_threads 1 --mkl_num_threads 1 --tag_filter long

This stack has been rebased, please help review :)

@VitalyFedyunin has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

Please rebase stack, landing tool requires it.

updated the original test case from test_nn.py:

python test_nn.py TestNN.test_channel_shuffle

The test case has validated all 3 memory formats: torch.contiguous, torch.channels_last, torch.channels_last_3d

The stack is newly rebased!

Sorry, can you please rebase again, I'm getting ghstack errors Internal error: couldn't understand base diff

This stack is newly newly based, please review! @VitalyFedyunin

Also #58348 has been updated: the backward OPs are removed to align with e6c3aa3

@VitalyFedyunin has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

@mingfeima why does this PR introduce native_channel_shuffle in native_functions.yaml when it doesnt really need to provide aten op by the same name and thus be usable from torch namespace?

@mingfeima why does this PR introduce native_channel_shuffle in native_functions.yaml when it doesnt really need to provide aten op by the same name and thus be usable from torch namespace?

It's because ChannelShuffle has a math implementation, while native_ indicates a directly kernel level implementation. Similar to other OPs with math implementation like LayerNorm or GroupNorm, they have native_ dispatch as well.

Also one more reason is that channel_shuffle could possibly run into XNNPack (decided at compile time)

#if defined(C10_MOBILE) && defined(USE_XNNPACK)
  if (self.is_contiguous(MemoryFormat::ChannelsLast) &&
      xnnpack::use_channel_shuffle(self, groups)) {
    return xnnpack::channel_shuffle(self, groups);
  }

If i merge native_channel_shuffle into channel_shuffle, i could not handle it.

Another tricky thing about it is the native_ implementation only applies for CPU and CUDA will fall into dispatch key of CompositeImplicitAutograd. It would be more decent if the cuda kernel is added as well but usually we (intel folks) only touched the CPU part.

Feel free to make a change if better options can be made :)
Actually we cares more about the performance part and also make sure that it won't break MemoryFormat propagation (so that it won't introduce reorder in following OPs).

@mingfeima could you please submit PR with python docstrings for this function.