Add the intra-op parallelism for equal operator (#28810)

jianyuh · facebook-github-bot · commit 6a4b51aec103 · 2019-11-06T10:30:44.000-08:00
Summary: Pull Request resolved: #28810 Similar to #28464 and #28477, we would like to enable the intra-op parallelism for layer norm. This will be mapped to the parallel performance win for the BERT/RoBERTa model. Test Plan: CI Differential Revision: D18165752 fbshipit-source-id: 354cede4c36893acbd69711f49aa6a51dc94397f
diff --git a/aten/src/TH/generic/THTensorMoreMath.cpp b/aten/src/TH/generic/THTensorMoreMath.cpp
@@ -19,7 +19,7 @@ ptrdiff_t THTensor_(numel)(THTensor *t)
 
 int THTensor_(equal)(THTensor *ta, THTensor* tb)
 {
-  int equal = 1;
+  std::atomic<int> equal{1};
   if(!THTensor_(isSameSizeAs)(ta, tb))
     return 0;
 
@@ -28,9 +28,21 @@ int THTensor_(equal)(THTensor *ta, THTensor* tb)
     scalar_t *tbp = tb->data<scalar_t>();
     ptrdiff_t sz = THTensor_(nElement)(ta);
     ptrdiff_t i;
-    for (i=0; i<sz; ++i){
-      if(tap[i] != tbp[i]) return 0;
-    }
+    at::parallel_for(
+        0,
+        sz,
+        TH_OMP_OVERHEAD_THRESHOLD,
+        [&](int64_t begin, int64_t end) {
+          for (auto iter = begin; iter < end; iter++) {
+            if (!equal) {
+              break;
+            }
+            if (tap[iter] != tbp[iter]) {
+              equal = 0;
+              break;
+            }
+          }
+        });
   } else {
     // Short-circuit the apply function on inequality
     TH_TENSOR_APPLY2(scalar_t, ta, scalar_t, tb,
@@ -39,7 +51,7 @@ int THTensor_(equal)(THTensor *ta, THTensor* tb)
                         TH_TENSOR_APPLY_hasFinished = 1; break;
                      })
   }
-  return equal;
+  return equal.load();
 }
 
 // Helper function to be used in a reduction operation.
