MKLDNN Sum and Batch Normalizaton kernels

onnxruntime/core/providers/cpu/math/element_wise_ops.h

@@ -139,7 +139,7 @@ class Log final : public OpKernel {
 };
 
 template <typename T>
-class Sum_6 final : public OpKernel {
+class Sum_6 : public OpKernel {
  public:
   Sum_6(const OpKernelInfo& info) : OpKernel(info) {
   }

onnxruntime/core/providers/cpu/nn/batch_norm.h

@@ -27,7 +27,7 @@
 namespace onnxruntime {
 
 template <typename T>
-class BatchNorm final : public OpKernel {
+class BatchNorm : public OpKernel {
  public:
   BatchNorm(const OpKernelInfo& op_kernel_info) : OpKernel(op_kernel_info) {
     float tmp_eplison;
@@ -38,7 +38,7 @@ class BatchNorm final : public OpKernel {
 
   Status Compute(OpKernelContext* p_op_kernel_context) const override;
 
- private:
+ protected:
   float epsilon_ = 1e-5f;
   int64_t is_test_;  // ignored in this implementation since we're doing inferencing only.
 };

onnxruntime/core/providers/mkldnn/math/sum.cc

@@ -0,0 +1,225 @@
+// Copyright(C) 2018 Intel Corporation
+// Licensed under the MIT License
+
+#ifdef _WIN32
+#pragma warning(disable : 4244)
+#endif
+
+#include "core/providers/mkldnn/mkldnn_common.h"
+#include "core/providers/mkldnn/math/sum.h"
+#include "core/providers/mkldnn/mkldnn_fwd.h"
+
+namespace onnxruntime {
+namespace mkl_dnn {
+
+namespace {
+// Struct which encapsulates parameters for MKLDNN Sum primitives.
+struct SumParams {
+  const std::vector<mkldnn::memory::dims>& src_dims;
+  const mkldnn::memory::dims& dst_dim;
+  const int num_inputs;
+  const int num_dimensions;
+
+  SumParams(const std::vector<mkldnn::memory::dims>& dims,
+    const mkldnn::memory::dims& dst_dims, const int numinputs, 
+    const int dimensions)
+    : src_dims(dims),
+    dst_dim(dst_dims),
+    num_inputs(numinputs),
+    num_dimensions(dimensions) {}
+
+  // Used as the key for Sum Primitive Reuse Sum.
+  std::string ToString() const {
+    std::string key;
+    key.reserve(64);
+    key.append("sum_");
+    AddDimsToKey(key, src_dims[0]);
+    AddDimsToKey(key, dst_dim);
+    return key;
+  }
+};
+
+template <typename T>
+class SumPrimitive final : public PrimitiveBase {
+ public:
+  explicit SumPrimitive(const SumParams& params)
+      : cpu_engine_(GetEngine()) {
+    context_.stream.reset(new mkldnn::stream(mkldnn::stream::kind::eager));
+    if (context_.sum_pd == nullptr) {
+      Initialize(params);
+    }
+  }
+
+  ~SumPrimitive() = default;
+
+  void Compute(OpKernelContext* context, int numinputs) {
+    const Tensor* X1 = context->Input<Tensor>(0);
+    Tensor* Y = context->Output(0, X1->Shape());
+    T* dst_data = Y->template MutableData<T>();
+
+    context_.dst_mem->set_data_handle(
+      static_cast<void*>(static_cast<T*>(dst_data)));
+
+    for (int i = 0; i < numinputs; i++) {
+      const Tensor* X = context->Input<Tensor>(i);
+      const T* src_data = X->template Data<T>();
+      context_.srcs_memory[i].set_data_handle(
+        static_cast<void*>(const_cast<T*>(src_data)));
+    }
+
+    context_.stream->submit(context_.net);
+  }
+
+  std::unique_ptr<mkldnn::memory::desc> GetDstMemoryDesc() const { 
+    return context_.dst_md; 
+  }
+
+  std::unique_ptr<mkldnn::sum::primitive_desc>
+  GetPrimitiveDesc() const {
+    return context_.sum_pd;
+  }
+
+ private:
+  struct SumContext {
+    std::unique_ptr<mkldnn::memory::desc> src_md;
+    std::unique_ptr<mkldnn::memory::desc> dst_md;
+
+    std::vector<mkldnn::memory> srcs_memory;
+    std::unique_ptr<mkldnn::memory> dst_mem;
+
+    std::vector<mkldnn::memory::primitive_desc> srcs_pd;
+    std::unique_ptr<mkldnn::memory::primitive_desc> src_mpd;
+    std::unique_ptr<mkldnn::memory::primitive_desc> dst_pd;
+    std::unique_ptr<mkldnn::sum::primitive_desc> sum_pd;
+
+    std::unique_ptr<mkldnn::stream> stream;
+    std::vector<mkldnn::primitive> net;
+  };
+
+  void Initialize(const SumParams& params) {
+    std::vector<float> coeff;
+
+    mkldnn::memory::format fmt = mkldnn::memory::format::any;
+    switch (params.num_dimensions) {
+    case 1: { fmt = mkldnn::memory::format::x; break; }
+    case 2: { fmt = mkldnn::memory::format::nc; break; }
+    case 3: { fmt = mkldnn::memory::format::ntc; break; }
+    case 4: { fmt = mkldnn::memory::format::nchw; break; }
+    case 5: { fmt = mkldnn::memory::format::ncdhw; break; }
+    default: {  fmt = mkldnn::memory::format::any; break; }
+    }
+
+    for (int i = 0; i < params.num_inputs; i++) {
+      context_.src_md.reset(
+        new mkldnn::memory::desc({params.src_dims[i]}, MklDnnType<T>(), fmt));
+      auto mpd = mkldnn::memory::primitive_desc(*context_.src_md, cpu_engine_);
+      auto src_memory = mkldnn::memory(mpd, nullptr);
+
+      context_.srcs_pd.push_back(mpd);
+      context_.srcs_memory.push_back(src_memory);
+      coeff.push_back(1.0);
+    }
+
+    std::unique_ptr<mkldnn::memory> dst;
+    context_.dst_md.reset(new mkldnn::memory::desc(
+      {params.dst_dim}, MklDnnType<T>(), mkldnn::memory::format::any));
+    context_.sum_pd.reset(new mkldnn::sum::primitive_desc(
+      *context_.dst_md, coeff, context_.srcs_pd));
+    context_.dst_mem.reset(new mkldnn::memory(
+      context_.sum_pd->dst_primitive_desc(), nullptr));
+
+    std::vector<mkldnn::primitive::at> inputs;
+    for (int i = 0; i < params.num_inputs; i++) {
+      inputs.push_back(context_.srcs_memory[i]);
+    }
+    auto c = mkldnn::sum(*context_.sum_pd, inputs, *context_.dst_mem);
+    context_.net.push_back(c);
+  }
+
+  SumContext context_;
+  mkldnn::engine& cpu_engine_;
+};
+
+// Pool which allows for reuse of MKLDNN Sum primitives which are 
+// expensive to instantiate. To address thread safety, the primitives
+// are stored in a map on thread local storage.
+
+template <typename T>
+class SumPrimitivePool : public PrimitivePool<T> {
+ public:
+  static SumPrimitive<T>* Get(const SumParams& params) {
+    SumPrimitive<T>* primitive = dynamic_cast<SumPrimitive<T>*>(
+        SumPrimitivePool<T>::GetInstance().GetPrimitive(params.ToString()));
+
+    if (primitive == nullptr) {
+      auto sum_primitive = std::make_unique<SumPrimitive<T>>(params);
+      primitive = sum_primitive.get();
+      SumPrimitivePool<T>::GetInstance().SetPrimitive(
+        params.ToString(), std::move(sum_primitive));
+    }
+    return primitive;
+  }
+
+ private:
+  SumPrimitivePool() = default;
+  ~SumPrimitivePool() = default;
+
+  static SumPrimitivePool& GetInstance() {
+    static SumPrimitivePool pool;
+    return pool;
+  }
+};
+} // namespace_
+
+template <typename T>
+Status Sum<T>::Compute(OpKernelContext* context) const {
+  int num_inputs = static_cast<int>(OpKernel::Node().InputDefs().size());
+
+  ONNXRUNTIME_ENFORCE(num_inputs > 0, "MKLDNN Sum kernel: Must have at least one input");
+
+  if (num_inputs == 1) {
+    return onnxruntime::Sum_6<T>::Compute(context);
+  }
+
+  std::vector<mkldnn::memory::dims> src_dims;
+
+  const Tensor* X1 = context->Input<Tensor>(0);
+  Tensor* Y = context->Output(0, X1->Shape());
+  int dimensions = static_cast<int>(X1->Shape().NumDimensions());
+
+  const TensorShape& x_shape = X1->Shape();
+  const auto& x_dims = x_shape.GetDims();
+  mkldnn::memory::dims src_dim(x_dims.begin(), x_dims.end());
+
+  mkldnn::memory::dims dst_dims_mkl(
+    Y->Shape().GetDims().begin(), Y->Shape().GetDims().end());
+
+  for (int i = 0; i < num_inputs; i++) {
+    const Tensor* X = context->Input<Tensor>(i);
+    mkldnn::memory::dims src_dims_mkl(
+      X->Shape().GetDims().begin(), X->Shape().GetDims().end());
+    src_dims.push_back(src_dims_mkl);
+  }
+  try {
+    SumParams parameters(src_dims, dst_dims_mkl, num_inputs, dimensions);
+    SumPrimitive<T>* sum_primitive = SumPrimitivePool<T>::Get(parameters);
+    ONNXRUNTIME_RETURN_IF_NOT(sum_primitive != nullptr);
+    sum_primitive->Compute(context, num_inputs);
+  } catch (const mkldnn::error& e) {
+    return ONNXRUNTIME_MAKE_STATUS(ONNXRUNTIME, FAIL, "Status: ", e.status, 
+      ", message: ", e.message.c_str());
+  }
+
+  return Status::OK();
+}
+
+ONNX_OPERATOR_KERNEL_EX(
+    Sum,
+    kOnnxDomain,
+    6,
+    kMklDnnExecutionProvider,
+    KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
+    Sum<float>);
+
+} // namespace mkl_dnn
+} // namespace onnxruntime

onnxruntime/core/providers/mkldnn/math/sum.h

@@ -0,0 +1,22 @@
+// Copyright(C) 2018 Intel Corporation
+// Licensed under the MIT License
+
+#pragma once
+#include "core/framework/op_kernel.h"
+#include "core/providers/cpu/math/element_wise_ops.h"
+#include "../mkldnn_execution_provider.h"
+
+namespace onnxruntime {
+namespace mkl_dnn {
+
+template <typename T>
+class Sum final : public onnxruntime::Sum_6<T> {
+ public:
+   Sum(const OpKernelInfo& info) : onnxruntime::Sum_6<T>(info) {}
+
+  Status Compute(OpKernelContext* context) const override;
+
+private:
+};
+}  // namespace mkl_dnn
+}  // namespace onnxruntime

onnxruntime/core/providers/mkldnn/memcpy_s.h

@@ -0,0 +1,13 @@
+// Copyright(C) 2018 Intel Corporation
+// Licensed under the MIT License
+
+#pragma once
+
+#ifdef _WIN32
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+#endif
+
+// memcpy is deprecated. Replacing it with more secure equivalent memcpy_s
+//
+#define MEMCPY_S(dest, src, destsz, srcsz) memcpy(dest, src, MIN(destsz, srcsz))
+

onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.cc

@@ -66,6 +66,8 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 7,
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 1, MemcpyFromHost);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 1, MemcpyToHost);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 6, Relu);
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 6, Sum);
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 7, BatchNormalization);
 class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 7, 8, float, AveragePool);
 class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 1, 8, float, GlobalAveragePool);
 class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 1, 7, float, MaxPool);
@@ -78,7 +80,9 @@ void RegisterMKLDNNKernels(std::function<void(KernelCreateInfo&&)> fn) {
   fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 7, Gemm)>());
   fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 1, MemcpyFromHost)>());
   fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 1, MemcpyToHost)>());
-  fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 6, Relu)>()); 
+  fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 6, Relu)>());
+  fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 6, Sum)>());
+  fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 7, BatchNormalization)>()); 
   fn(BuildKernel<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 7, 8, float, AveragePool)>());
   fn(BuildKernel<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 1, 8, float, GlobalAveragePool)>());
   fn(BuildKernel<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kMklDnnExecutionProvider, kOnnxDomain, 1, 7, float, MaxPool)>());