add channels last for MaxPool2d (#48917)

Summary: Pull Request resolved: #48917 max_pool2d channels last support forward path max_pool2d channels last support backward path vectorize channels last forward path rename the header file fix windows build combine PoolingKernel.h into Pool.h add data type check loosen test_max_pool2d_nhwc to cover device CPU Test Plan: Imported from OSS Reviewed By: glaringlee Differential Revision: D25399470 Pulled By: VitalyFedyunin fbshipit-source-id: b49b9581f1329a8c2b9c75bb10f12e2650e4c65a
pytorch · Apr 2, 2021 · f43eb59 · f43eb59
1 parent bdb225e
commit f43eb59
Show file tree

Hide file tree

Showing 10 changed files with 467 additions and 291 deletions.
diff --git a/aten/src/ATen/cpu/vec256/vec256_base.h b/aten/src/ATen/cpu/vec256/vec256_base.h
@@ -199,6 +199,17 @@ struct Vec256 {
     }
     return mask;
   }
+  Vec256<T> isnan() const {
+    Vec256<T> vec;
+    for (int64_t i = 0; i != size(); i++) {
+      if (_isnan(values[i])) {
+        std::memset(static_cast<void*>(vec.values + i), 0xFF, sizeof(T));
+      } else {
+        std::memset(static_cast<void*>(vec.values + i), 0, sizeof(T));
+      }
+    }
+    return vec;
+  }
   Vec256<T> map(T (*f)(T)) const {
     Vec256<T> ret;
     for (int64_t i = 0; i != size(); i++) {

diff --git a/aten/src/ATen/cpu/vec256/vec256_double.h b/aten/src/ATen/cpu/vec256/vec256_double.h
@@ -95,6 +95,9 @@ template <> class Vec256<double> {
     __m256d cmp = _mm256_cmp_pd(values, _mm256_set1_pd(0.0), _CMP_EQ_OQ);
     return _mm256_movemask_pd(cmp);
   }
+  Vec256<double> isnan() const {
+    return _mm256_cmp_pd(values, _mm256_set1_pd(0.0), _CMP_UNORD_Q);
+  }
   Vec256<double> map(double (*f)(double)) const {
     __at_align32__ double tmp[size()];
     store(tmp);

diff --git a/aten/src/ATen/cpu/vec256/vec256_float.h b/aten/src/ATen/cpu/vec256/vec256_float.h
@@ -102,6 +102,9 @@ template <> class Vec256<float> {
     __m256 cmp = _mm256_cmp_ps(values, _mm256_set1_ps(0.0f), _CMP_EQ_OQ);
     return _mm256_movemask_ps(cmp);
   }
+  Vec256<float> isnan() const {
+    return _mm256_cmp_ps(values, _mm256_set1_ps(0.0f), _CMP_UNORD_Q);
+  }
   Vec256<float> map(float (*f)(float)) const {
     __at_align32__ float tmp[size()];
     store(tmp);

diff --git a/aten/src/ATen/cpu/vec256/vec256_float_neon.h b/aten/src/ATen/cpu/vec256/vec256_float_neon.h
@@ -282,6 +282,19 @@ template <> class Vec256<float> {
     }
     return mask;
   }
+  Vec256<float> isnan() const {
+    __at_align32__ float tmp[size()];
+    __at_align32__ float res[size()];
+    store(tmp);
+    for (int64_t i = 0; i < size(); i++) {
+      if (_isnan(tmp[i])) {
+        std::memset(static_cast<void*>(&res[i]), 0xFF, sizeof(float));
+      } else {
+        std::memset(static_cast<void*>(&res[i]), 0, sizeof(float));
+      }
+    }
+    return loadu(res);
+  };
   Vec256<float> map(float (*f)(float)) const {
     __at_align32__ float tmp[size()];
     store(tmp);