I believe we actually do compile with -mavx in fbcode. I'm going to patch this and verify the performance in OSS

@zheng-xq - If you want to make use of the dispatch we use in OSS you need to move 'native/Normalization.cpp' into 'native/cpu/Normalization.cpp'. We can also do that as a separate step afterwards.

FWIW for BatchNorm1D this falls back to scalar code because only the inner loop at line 101 is vectorized, and image_size=1

This specialization gives me a 1.6x speedup on BatchNorm1D:

--- a/aten/src/ATen/native/Normalization.cpp
+++ b/aten/src/ATen/native/Normalization.cpp
@@ -96,13 +96,24 @@ void batch_norm_cpu_inference_contiguous(Tensor* output, const Tensor& input,
   // No need to use parallel_for as this function is supposed to be
   // memory-limited.
   // Keep the loop struture simple to make sure compiler vetorization kicks in.
-  for (int64_t n = 0; n < n_batch; ++n) {
-    for (int64_t c = 0; c < n_channel; ++c) {
-      for (int64_t i = 0; i < image_size; ++i) {
-        // Keep all the offset calculation within the inner loop for simplicity.
-        // Compilers are very good at hoisting the common part outside.
-        int64_t offset = n * n_channel * image_size + c * image_size + i;
-        output_data[offset] = input_data[offset] * alpha_data[c] + beta_data[c];
+  if (image_size != 1) {
+    for (int64_t n = 0; n < n_batch; ++n) {
+      for (int64_t c = 0; c < n_channel; ++c) {
+        for (int64_t i = 0; i < image_size; ++i) {
+          // Keep all the offset calculation within the inner loop for simplicity.
+          // Compilers are very good at hoisting the common part outside.
+          int64_t offset = n * n_channel * image_size + c * image_size + i;
+          output_data[offset] = input_data[offset] * alpha_data[c] + beta_data[c];
+        }
+      }
+    }
+  } else {
+    for (int64_t n = 0; n < n_batch; ++n) {
+      for (int64_t c = 0; c < n_channel; ++c) {
+          // Keep all the offset calculation within the inner loop for simplicity.
+          // Compilers are very good at hoisting the common part outside.
+          int64_t offset = n * n_channel + c;
+          output_data[offset] = input_data[offset] * alpha_data[c] + beta_data[c];

@jamesr66a

This specialization gives me a 1.6x speedup on BatchNorm1D:

--- a/aten/src/ATen/native/Normalization.cpp
+++ b/aten/src/ATen/native/Normalization.cpp
@@ -96,13 +96,24 @@ void batch_norm_cpu_inference_contiguous(Tensor* output, const Tensor& input,
   // No need to use parallel_for as this function is supposed to be
   // memory-limited.
   // Keep the loop struture simple to make sure compiler vetorization kicks in.
-  for (int64_t n = 0; n < n_batch; ++n) {
-    for (int64_t c = 0; c < n_channel; ++c) {
-      for (int64_t i = 0; i < image_size; ++i) {
-        // Keep all the offset calculation within the inner loop for simplicity.
-        // Compilers are very good at hoisting the common part outside.
-        int64_t offset = n * n_channel * image_size + c * image_size + i;
-        output_data[offset] = input_data[offset] * alpha_data[c] + beta_data[c];
+  if (image_size != 1) {
+    for (int64_t n = 0; n < n_batch; ++n) {
+      for (int64_t c = 0; c < n_channel; ++c) {
+        for (int64_t i = 0; i < image_size; ++i) {
+          // Keep all the offset calculation within the inner loop for simplicity.
+          // Compilers are very good at hoisting the common part outside.
+          int64_t offset = n * n_channel * image_size + c * image_size + i;
+          output_data[offset] = input_data[offset] * alpha_data[c] + beta_data[c];
+        }
+      }
+    }
+  } else {
+    for (int64_t n = 0; n < n_batch; ++n) {
+      for (int64_t c = 0; c < n_channel; ++c) {
+          // Keep all the offset calculation within the inner loop for simplicity.
+          // Compilers are very good at hoisting the common part outside.
+          int64_t offset = n * n_channel + c;
+          output_data[offset] = input_data[offset] * alpha_data[c] + beta_data[c];

Wouldn't flattening the loop nest be a simpler fix?

Thanks for the comments! PTAL

This pull request has been merged in 5627940.

Do we actually want to compare PT vs Caffe2 implementations as well?

I see that in Caffe2 we use Eigen, could have it worked here?
Btw, do we want to start sharing ops rather than implementing new ones from scratch and keeping old ones? (cc @dzhulgakov )

about eigen in particular, we put some effort to develop our own packet abstraction (vec256.h) and a switchable threadpool. So Eigen doesn't add a lot of value, and maybe not necessary.

@salexspb - Eigen is an entirely separate topic in itself. So far we haven't seen any advantages in using it. Ping me directly if you want to talk more and see what we've done already. In general, a rigorous comparison to evaluate whether adding that dependency is worthwhile and a whole project in and of itself.

I am confused on the dependency part. Don't we already use it in Caffe2? How is unified build is working then?

@salexspb - yes, that's true we already have it as part of our build chain. In terms of dependency, I mean includes. It's adding complexity to the code.

In general, we should have comparison benchmarks for operators like this which are already supported by Caffe2 and then see if we can recycle them. But it's often also worthwhile to see whether we can write faster code by using ATen's abstractions.

For this particular operator, I agree that it could be good to write a side-by-side comparison in a single script to compare Caffe2's and PyTorch's implementation.