[torch] Unify batch_box_cox implementations into perfkernels folder (pytorch#86569)

Summary:
1) Adding MKL/AVX2 based implementation into perfkernels. This implementation is similar to caffe2/operators/batch_box_cox_op.cc
2) Migrating batch_box_cox_op of caffe2 use this implementation

Test Plan: CI

Differential Revision: D40208074

Pull Request resolved: pytorch#86569
Approved by: https://github.com/hyuen

Author: efiks

caffe2/operators/batch_box_cox_op.cc

@@ -2,72 +2,34 @@
 
 #include "caffe2/core/operator.h"
 #include "caffe2/core/tensor.h"
-
-#ifdef CAFFE2_USE_MKL
-#include <mkl.h>
-#endif // CAFFE2_USE_MKL
+#include "caffe2/perfkernels/batch_box_cox.h"
 
 namespace caffe2 {
 
-#ifdef CAFFE2_USE_MKL
 namespace {
-
-// Helpers for copying parameters.
 template <typename T>
-void TileArrayIntoVector(const T* a, int D, int K, vector<T>* b) {
-  b->resize(K * D);
-  for (int k = 0; k < K; k++) {
-    std::copy(a, a + D, b->begin() + k * D);
-  }
-}
-
-void TileIndicesInPlace(vector<int>* v, int D, int K) {
-  int n = v->size();
-  v->resize(K * n);
-  for (int k = 1; k < K; k++) {
-    for (int j = 0; j < n; j++) {
-      (*v)[k * n + j] = (*v)[j] + k * D;
+void BoxCoxNaive(
+    int64_t N,
+    int64_t D,
+    const T* data_ptr,
+    const T* lambda1_ptr,
+    const T* lambda2_ptr,
+    T* output_ptr) {
+  constexpr T k_eps = static_cast<T>(1e-6);
+  for (int64_t i = 0; i < N; i++) {
+    for (int64_t j = 0; j < D; j++, data_ptr++, output_ptr++) {
+      T lambda1_v = lambda1_ptr[j];
+      T lambda2_v = lambda2_ptr[j];
+      T tmp = std::max(*data_ptr + lambda2_v, k_eps);
+      if (lambda1_v == 0) {
+        *output_ptr = std::log(tmp);
+      } else {
+        *output_ptr = (std::pow(tmp, lambda1_v) - 1) / lambda1_v;
+      }
     }
   }
 }
-
-// MKL VML function templates.
-template <typename T>
-void PackV(const int N, const T* a, const int* ia, T* y);
-template <typename T>
-void UnpackV(const int N, const T* a, T* y, const int* iy);
-template <typename T>
-void Pow(const int N, const T* a, const T* b, T* y);
-
-#define DELEGATE_PACKV_FUNCTION(T, OriginalFunc)                \
-  template <>                                                   \
-  void PackV<T>(const int N, const T* a, const int* ia, T* y) { \
-    OriginalFunc(N, a, ia, y);                                  \
-  }
-DELEGATE_PACKV_FUNCTION(float, vsPackV)
-DELEGATE_PACKV_FUNCTION(double, vdPackV)
-#undef DELEGATE_PACKV_FUNCTION
-
-#define DELEGATE_UNPACKV_FUNCTION(T, OriginalFunc)                \
-  template <>                                                     \
-  void UnpackV<T>(const int N, const T* a, T* y, const int* iy) { \
-    OriginalFunc(N, a, y, iy);                                    \
-  }
-DELEGATE_UNPACKV_FUNCTION(float, vsUnpackV)
-DELEGATE_UNPACKV_FUNCTION(double, vdUnpackV)
-#undef DELEGATE_UNPACKV_FUNCTION
-
-#define DELEGATE_SIMPLE_BINARY_FUNCTION(T, Funcname, OriginalFunc) \
-  template <>                                                      \
-  void Funcname<T>(const int N, const T* a, const T* b, T* y) {    \
-    OriginalFunc(N, a, b, y);                                      \
-  }
-DELEGATE_SIMPLE_BINARY_FUNCTION(float, Pow, vsPow)
-DELEGATE_SIMPLE_BINARY_FUNCTION(double, Pow, vdPow)
-#undef DELEGATE_SIMPLE_BINARY_FUNCTION
-
-} // namespace
-#endif // CAFFE2_USE_MKL
+}
 
 template <>
 template <typename T>
@@ -93,227 +55,19 @@ bool BatchBoxCoxOp<CPUContext>::DoRunWithType() {
   const auto* lambda1_ptr = lambda1.template data<T>();
   const auto* lambda2_ptr = lambda2.template data<T>();
 
-  const T k_eps = static_cast<T>(1e-6);
-
 #ifdef CAFFE2_USE_MKL
   if (min_block_size_ < 1) {
-    BoxCoxNaive(N, D, data_ptr, lambda1_ptr, lambda2_ptr, k_eps, output_ptr);
-  } else {
-    // Find zero-valued columns, since they get special treatment.
-    nonzeros_.clear();
-    zeros_.clear();
-    nonzeros_.reserve(D);
-    zeros_.reserve(D);
-    for (int64_t j = 0; j < D; j++) {
-      if (lambda1_ptr[j] == 0) {
-        zeros_.push_back(j);
-      } else {
-        nonzeros_.push_back(j);
-      }
-    }
-
-    // Process K rows at a time for effective vectorization with small rows.
-    const int K = std::min(N, (min_block_size_ + D - 1) / D);
-
-    // Avoid copying data if all lambda1 values are zero, or if all are nonzero.
-    // In each of the three cases here, when K > 1, first process batches of K
-    // rows by replicating the input parameters K times. Then finish row-by-row.
-    TypedCachedBuffers<T>& b = GetBuffers<T>();
-    if (nonzeros_.size() == D) {
-      int64_t i = 0;
-      if (K > 1) {
-        TileArrayIntoVector(lambda1_ptr, D, K, &b.lambda1_);
-        TileArrayIntoVector(lambda2_ptr, D, K, &b.lambda2_);
-        TORCH_DCHECK_EQ(K * D, b.lambda1_.size());
-        TORCH_DCHECK_EQ(K * D, b.lambda2_.size());
-        for (; i < N - K + 1; i += K, data_ptr += K * D, output_ptr += K * D) {
-          BoxCoxNonzeroLambda(
-              K * D,
-              data_ptr,
-              b.lambda1_.data(),
-              b.lambda2_.data(),
-              k_eps,
-              output_ptr);
-        }
-      }
-      for (; i < N; i++, data_ptr += D, output_ptr += D) {
-        BoxCoxNonzeroLambda(
-            D, data_ptr, lambda1_ptr, lambda2_ptr, k_eps, output_ptr);
-      }
-    } else if (zeros_.size() == D) {
-      int64_t i = 0;
-      if (K > 1) {
-        TileArrayIntoVector(lambda2_ptr, D, K, &b.lambda2_z_);
-        TORCH_DCHECK_EQ(K * D, b.lambda2_z_.size());
-        for (; i < N - K + 1; i += K, data_ptr += K * D, output_ptr += K * D) {
-          BoxCoxZeroLambda(
-              K * D, data_ptr, b.lambda2_z_.data(), k_eps, output_ptr);
-        }
-      }
-      for (; i < N; i++, data_ptr += D, output_ptr += D) {
-        BoxCoxZeroLambda(D, data_ptr, lambda2_ptr, k_eps, output_ptr);
-      }
-    } else { // General case of mixed zero and non-zero lambda1 values.
-      int n = nonzeros_.size();
-      if (K > 1) {
-        TileIndicesInPlace(&nonzeros_, 0, K);
-        TileIndicesInPlace(&zeros_, 0, K);
-      }
-
-      // Gather parameter values into contiguous memory.
-      b.lambda1_.resize(nonzeros_.size());
-      b.lambda2_.resize(nonzeros_.size());
-      b.lambda2_z_.resize(zeros_.size());
-      PackV(nonzeros_.size(), lambda1_ptr, nonzeros_.data(), b.lambda1_.data());
-      PackV(nonzeros_.size(), lambda2_ptr, nonzeros_.data(), b.lambda2_.data());
-      PackV(zeros_.size(), lambda2_ptr, zeros_.data(), b.lambda2_z_.data());
-
-      int64_t i = 0;
-      b.accumulator_.resize(std::max(nonzeros_.size(), zeros_.size()));
-      if (K > 1) {
-        // Truncate to original size, and re-tile with offsets this time.
-        nonzeros_.resize(n);
-        zeros_.resize(D - n);
-        TileIndicesInPlace(&nonzeros_, D, K);
-        TileIndicesInPlace(&zeros_, D, K);
-        TORCH_DCHECK_EQ(nonzeros_.size(), b.lambda1_.size());
-        TORCH_DCHECK_EQ(nonzeros_.size(), b.lambda2_.size());
-        TORCH_DCHECK_EQ(zeros_.size(), b.lambda2_z_.size());
-        for (; i < N - K + 1; i += K, data_ptr += K * D, output_ptr += K * D) {
-          BoxCoxMixedLambda(
-              data_ptr,
-              nonzeros_,
-              zeros_,
-              b.lambda1_.data(),
-              b.lambda2_.data(),
-              b.lambda2_z_.data(),
-              k_eps,
-              b.accumulator_.data(),
-              output_ptr);
-        }
-        // Truncate to original size.
-        nonzeros_.resize(n);
-        zeros_.resize(D - n);
-      }
-      for (; i < N; i++, data_ptr += D, output_ptr += D) {
-        BoxCoxMixedLambda(
-            data_ptr,
-            nonzeros_,
-            zeros_,
-            b.lambda1_.data(),
-            b.lambda2_.data(),
-            b.lambda2_z_.data(),
-            k_eps,
-            b.accumulator_.data(),
-            output_ptr);
-      }
-    }
+    BoxCoxNaive(N, D, data_ptr, lambda1_ptr, lambda2_ptr, output_ptr);
+    return true;
   }
-#else // CAFFE2_USE_MKL
-  BoxCoxNaive(N, D, data_ptr, lambda1_ptr, lambda2_ptr, k_eps, output_ptr);
-#endif // CAFFE2_USE_MKL
+  caffe2::compute_batch_box_cox(
+    N, D, min_block_size_, data_ptr, lambda1_ptr, lambda2_ptr, output_ptr);
+#else
+  BoxCoxNaive(N, D, data_ptr, lambda1_ptr, lambda2_ptr, output_ptr);
+#endif
   return true;
 }
 
-template <>
-template <typename T>
-void BatchBoxCoxOp<CPUContext>::BoxCoxNaive(
-    int64_t N,
-    int64_t D,
-    const T* data_ptr,
-    const T* lambda1_ptr,
-    const T* lambda2_ptr,
-    T k_eps,
-    T* output_ptr) {
-  for (int64_t i = 0; i < N; i++) {
-    for (int64_t j = 0; j < D; j++, data_ptr++, output_ptr++) {
-      T lambda1_v = lambda1_ptr[j];
-      T lambda2_v = lambda2_ptr[j];
-      T tmp = std::max(*data_ptr + lambda2_v, k_eps);
-      if (lambda1_v == 0) {
-        *output_ptr = std::log(tmp);
-      } else {
-        *output_ptr = (std::pow(tmp, lambda1_v) - 1) / lambda1_v;
-      }
-    }
-  }
-}
-
-#ifdef CAFFE2_USE_MKL
-
-template <>
-template <typename T>
-void BatchBoxCoxOp<CPUContext>::BoxCoxNonzeroLambda(
-    int64_t D,
-    const T* data_ptr,
-    const T* lambda1,
-    const T* lambda2,
-    T k_eps,
-    T* out) {
-  caffe2::math::Add(D, data_ptr, lambda2, out, &context_);
-  for (int64_t j = 0; j < D; j++) {
-    out[j] = std::max(out[j], k_eps);
-  }
-  Pow(D, out, lambda1, out);
-  for (int64_t j = 0; j < D; j++) {
-    out[j] -= 1.0;
-  }
-  caffe2::math::Div(D, out, lambda1, out, &context_);
-}
-
-template <>
-template <typename T>
-void BatchBoxCoxOp<CPUContext>::BoxCoxZeroLambda(
-    int64_t D,
-    const T* data_ptr,
-    const T* lambda2,
-    T k_eps,
-    T* output_ptr) {
-  caffe2::math::Add(D, data_ptr, lambda2, output_ptr, &context_);
-  for (int64_t j = 0; j < D; j++) {
-    output_ptr[j] = std::max(output_ptr[j], k_eps);
-  }
-  caffe2::math::Log(D, output_ptr, output_ptr, &context_);
-}
-
-template <>
-template <typename T>
-void BatchBoxCoxOp<CPUContext>::BoxCoxMixedLambda(
-    const T* data_ptr,
-    const vector<int>& nonzeros,
-    const vector<int>& zeros,
-    const T* lambda1,
-    const T* lambda2,
-    const T* lambda2_z,
-    T k_eps,
-    T* buffer,
-    T* output_ptr) {
-  PackV(nonzeros.size(), data_ptr, nonzeros.data(), buffer);
-  BoxCoxNonzeroLambda(nonzeros.size(), buffer, lambda1, lambda2, k_eps, buffer);
-  UnpackV(nonzeros.size(), buffer, output_ptr, nonzeros.data());
-
-  PackV(zeros.size(), data_ptr, zeros.data(), buffer);
-  BoxCoxZeroLambda(zeros.size(), buffer, lambda2_z, k_eps, buffer);
-  UnpackV(zeros.size(), buffer, output_ptr, zeros.data());
-}
-
-// Helpers to access cached buffers.
-#define DEFINE_CACHED_BUFFERS(T, tag)                                         \
-  template <>                                                                 \
-  template <>                                                                 \
-  BatchBoxCoxOp<CPUContext>::TypedCachedBuffers<T>&                           \
-  BatchBoxCoxOp<CPUContext>::GetBuffers<T>() {                                \
-    if (!buffers_ || buffers_->type_ != tag) {                                \
-      buffers_.reset(new BatchBoxCoxOp<CPUContext>::TypedCachedBuffers<T>()); \
-      buffers_->type_ = tag;                                                  \
-    }                                                                         \
-    return *static_cast<TypedCachedBuffers<T>*>(buffers_.get());              \
-  }
-DEFINE_CACHED_BUFFERS(float, 1);
-DEFINE_CACHED_BUFFERS(double, 2);
-#undef DEFINE_CACHED_BUFFERS
-
-#endif // CAFFE2_USE_MKL
 
 namespace {
 

caffe2/operators/batch_box_cox_op.h

@@ -29,65 +29,7 @@ class BatchBoxCoxOp final : public Operator<Context> {
   bool DoRunWithType();
 
  protected:
-  template <typename T>
-  void BoxCoxNaive(
-      int64_t N,
-      int64_t D,
-      const T* data_ptr,
-      const T* lambda1_ptr,
-      const T* lambda2_ptr,
-      T k_eps,
-      T* output_ptr);
-
-#ifdef CAFFE2_USE_MKL
-  template <typename T>
-  void BoxCoxNonzeroLambda(
-      int64_t D,
-      const T* data_ptr,
-      const T* lambda1,
-      const T* lambda2,
-      T k_eps,
-      T* output_ptr);
-
-  template <typename T>
-  void BoxCoxZeroLambda(
-      int64_t D,
-      const T* data_ptr,
-      const T* lambda2,
-      T k_eps,
-      T* output_ptr);
-
-  template <typename T>
-  void BoxCoxMixedLambda(
-      const T* data_ptr,
-      const vector<int>& nonzeros,
-      const vector<int>& zeros,
-      const T* lambda1,
-      const T* lambda2,
-      const T* lambda2_z,
-      T k_eps,
-      T* buffer,
-      T* output_ptr);
-
-  vector<int> nonzeros_, zeros_;
-
-  // Buffers used by the MKL version are cached across calls.
-  struct CachedBuffers {
-    virtual ~CachedBuffers() {}
-    int type_;
-  };
-  template <typename T>
-  struct TypedCachedBuffers : public CachedBuffers {
-    vector<T> lambda1_, lambda2_, lambda2_z_;
-    vector<T> accumulator_;
-  };
-  template <typename T>
-  TypedCachedBuffers<T>& GetBuffers();
-  unique_ptr<CachedBuffers> buffers_;
-
-#endif // CAFFE2_USE_MKL
-
-  int min_block_size_;
+  std::size_t min_block_size_;
 
   INPUT_TAGS(DATA, LAMBDA1, LAMBDA2);
 };