MKLDNN Sum and Batch Normalizaton kernels (#115)

* mkldnn activations.relu

* updates after PR review 1

* mkldnn Sum and Batch Normalization Kernels

* BN,Sum PR Review changes

* Sum PR review 2

* Sum PR review 3

* PR Review 4: const and const refs

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)>());

onnxruntime/core/providers/mkldnn/nn/batch_norm.cc

@@ -0,0 +1,287 @@
+// 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/nn/batch_norm.h"
+#include "core/providers/mkldnn/mkldnn_fwd.h"
+#include "core/providers/mkldnn/memcpy_s.h"
+#include "core/providers/cpu/nn/batch_norm_helper.h"
+
+namespace onnxruntime {
+namespace mkl_dnn {
+
+namespace {
+// Struct which encapsulates parameters for MKLDNN BatchNorm primitive.
+struct BatchNormParams {
+  const mkldnn::memory::dims& src_dims;
+  const mkldnn::memory::dims& scale_dims;
+  const mkldnn::memory::dims& b_dims;
+  const mkldnn::memory::dims& mean_dims;
+  const mkldnn::memory::dims& var_dims;
+  const mkldnn::memory::dims& dst_dims;
+  const float epsilon;
+  const int num_dimensions;
+
+  BatchNormParams(const mkldnn::memory::dims& src_dims_mkl,
+    const mkldnn::memory::dims& scale_dims_mkl,
+    const mkldnn::memory::dims& b_dims_mkl, const mkldnn::memory::dims& mean_dims_mkl,
+    const mkldnn::memory::dims& var_dims_mkl, const mkldnn::memory::dims& dst_dims_mkl,
+    const float eps, const int dimensions)
+    : src_dims(src_dims_mkl),
+    scale_dims(scale_dims_mkl),
+    b_dims(b_dims_mkl),
+    mean_dims(mean_dims_mkl),
+    var_dims(var_dims_mkl),
+    dst_dims(dst_dims_mkl),
+    epsilon(eps),
+    num_dimensions(dimensions) {}
+
+  // Used as the key for BatchNorm Primitive Reuse Pool.
+  std::string ToString() const {
+    std::string key;
+    key.reserve(128);
+    key.append("BatchNorm_");
+    AddDimsToKey(key, src_dims);
+    AddDimsToKey(key, scale_dims);
+    AddDimsToKey(key, b_dims);
+    AddDimsToKey(key, mean_dims);
+    AddDimsToKey(key, var_dims);
+    AddDimsToKey(key, dst_dims);
+    return key;
+  }
+};
+
+template <typename T>
+class BatchNormPrimitive final : public PrimitiveBase {
+ public:
+  explicit BatchNormPrimitive(const BatchNormParams& params)
+      : cpu_engine_(GetEngine()) {
+    context_.stream.reset(new mkldnn::stream(mkldnn::stream::kind::eager));
+    if (context_.batchnorm_fwd == nullptr) {
+      Initialize(params);
+    }
+  }
+
+  ~BatchNormPrimitive() = default;
+
+  void Compute(const T* src_data, const T* scale_data, const T* b_data, 
+    const T* mean_data, const T* var_data, const T* dst_data, 
+    int scale_dims_channels) {
+    context_.src_mem->set_data_handle(
+      static_cast<void*>(const_cast<T*>(src_data)));
+    context_.mean_mem->set_data_handle(
+      static_cast<void*>(const_cast<T*>(mean_data)));
+    context_.var_mem->set_data_handle(
+      static_cast<void*>(const_cast<T*>(var_data)));
+    context_.dst_mem->set_data_handle(
+      static_cast<void*>(const_cast<T*>(dst_data)));
+
+    T* scaleShift_buf = static_cast<T*>(context_.scale_shift_mem->get_data_handle());
+
+    size_t src_bytes = sizeof(T) * scale_dims_channels;
+    size_t dst_bytes = sizeof(T) * scale_dims_channels;
+
+    MEMCPY_S(scaleShift_buf, scale_data, src_bytes, dst_bytes);
+    MEMCPY_S(&scaleShift_buf[scale_dims_channels], b_data, src_bytes, dst_bytes);
+    context_.stream->submit(context_.net);
+    return;
+  }
+
+  std::unique_ptr<mkldnn::convolution_forward::primitive_desc>
+  GetPrimitiveDesc() const {
+    return context_.conv_fwd_pd;
+  }
+
+ private:
+  struct BatchNormContext {
+    std::unique_ptr<mkldnn::memory> src_mem;
+    std::unique_ptr<mkldnn::memory> scale_shift_mem;
+    std::unique_ptr<mkldnn::memory> mean_mem;
+    std::unique_ptr<mkldnn::memory> var_mem;
+    std::unique_ptr<mkldnn::memory> dst_mem;
+
+    std::unique_ptr<mkldnn::memory::desc> src_md;
+    std::unique_ptr<mkldnn::memory::desc> scale_shift_md;
+    std::unique_ptr<mkldnn::memory::desc> mean_md;
+    std::unique_ptr<mkldnn::memory::desc> var_md;
+    std::unique_ptr<mkldnn::memory::desc> dst_md;
+
+    std::unique_ptr<mkldnn::batch_normalization_forward::desc> batchnorm_fwd;
+    std::unique_ptr<mkldnn::batch_normalization_forward::primitive_desc> 
+      batchnorm_fwd_pd;
+
+    std::unique_ptr<mkldnn::stream> stream;
+    std::vector<mkldnn::primitive> net;
+  };
+
+  void Initialize(const BatchNormParams& params) {
+    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; }
+    }
+    context_.src_md.reset(new mkldnn::memory::desc(
+      { params.src_dims }, MklDnnType<T>(), fmt));
+
+    context_.scale_shift_md.reset(new mkldnn::memory::desc(
+      { 2, params.scale_dims[0] }, MklDnnType<T>(), mkldnn::memory::format::nc));
+
+    context_.mean_md.reset(new mkldnn::memory::desc(
+      { params.mean_dims }, MklDnnType<T>(), mkldnn::memory::format::x));
+    context_.var_md.reset(new mkldnn::memory::desc(
+      { params.var_dims }, MklDnnType<T>(), mkldnn::memory::format::x));
+    context_.dst_md.reset(new mkldnn::memory::desc(
+      { params.dst_dims }, MklDnnType<T>(), fmt));
+
+    context_.src_mem.reset(
+      new mkldnn::memory({ *context_.src_md, cpu_engine_ }, nullptr));
+   
+   // scale_shift_mem will allocate 2*C*sizeof(float) buffer
+   //
+    context_.scale_shift_mem.reset(
+      new mkldnn::memory({ *context_.scale_shift_md, cpu_engine_ }));
+
+    context_.mean_mem.reset(
+      new mkldnn::memory({ *context_.mean_md, cpu_engine_ }, nullptr));
+    context_.var_mem.reset(
+      new mkldnn::memory({ *context_.var_md, cpu_engine_ }, nullptr));
+
+    context_.batchnorm_fwd.reset(new mkldnn::batch_normalization_forward::desc(
+      mkldnn::prop_kind::forward_inference, *context_.src_md, params.epsilon,
+      mkldnn::batch_normalization_flag::use_scale_shift |
+      mkldnn::batch_normalization_flag::use_global_stats));
+
+   context_.batchnorm_fwd_pd.reset(
+     new mkldnn::batch_normalization_forward::primitive_desc(
+       *context_.batchnorm_fwd, cpu_engine_));
+
+   context_.dst_mem.reset(
+     new mkldnn::memory(
+       context_.batchnorm_fwd_pd->dst_primitive_desc(), nullptr));
+
+   auto bn = mkldnn::batch_normalization_forward(
+     *context_.batchnorm_fwd_pd,
+     (const mkldnn::primitive::at)*context_.src_mem,
+     (const mkldnn::primitive::at)*context_.mean_mem,
+     (const mkldnn::primitive::at)*context_.var_mem,
+     (const mkldnn::memory)*context_.scale_shift_mem,
+     (const mkldnn::memory) *context_.dst_mem);
+   
+   context_.net.push_back(bn);
+  }
+
+  BatchNormContext context_;
+  mkldnn::engine& cpu_engine_;
+};
+
+// Pool which allows for reuse of MKLDNN BatchNorm 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 BatchNormPrimitivePool : public PrimitivePool<T> {
+ public:
+  static BatchNormPrimitive<T>* Get(const BatchNormParams& params) {
+    BatchNormPrimitive<T>* primitive = 
+      dynamic_cast<BatchNormPrimitive<T>*>(
+        BatchNormPrimitivePool<T>::GetInstance().GetPrimitive(params.ToString()));
+
+    if (primitive == nullptr) {
+      auto BatchNorm_primitive = std::make_unique<BatchNormPrimitive<T>>(params);
+      primitive = BatchNorm_primitive.get();
+      BatchNormPrimitivePool<T>::GetInstance().SetPrimitive(
+        params.ToString(), std::move(BatchNorm_primitive));
+    }
+    return primitive;
+  }
+
+ private:
+  BatchNormPrimitivePool() = default;
+  ~BatchNormPrimitivePool() = default;
+
+  static BatchNormPrimitivePool& GetInstance() {
+    static BatchNormPrimitivePool pool;
+    return pool;
+  }
+};
+} // namespace
+
+template <typename T>
+Status BatchNorm<T>::Compute(OpKernelContext* context) const {
+  const Tensor* X = context->Input<Tensor>(0);
+  
+  int num_dimensions = static_cast<int>(X->Shape().NumDimensions());
+  if (num_dimensions == 3) {
+    // Fall back CPU implementation
+    return onnxruntime::BatchNorm<T>::Compute(context);
+  }
+
+  const T* src_data = X->template Data<T>();
+
+  const Tensor* scale = context->Input<Tensor>(1);
+  const T* scale_data = scale->template Data<T>();
+
+  const Tensor* B = context->Input<Tensor>(2);
+  const T* b_data = B->template Data<T>();
+
+  const Tensor* mean = context->Input<Tensor>(3);
+  const T* mean_data = mean->template Data<T>();
+
+  const Tensor* var = context->Input<Tensor>(4);
+  const T* var_data = var->template Data<T>();
+
+  Tensor* Y = context->Output(0, X->Shape());
+  T* dst_data = Y->template MutableData<T>();
+
+  ONNXRUNTIME_RETURN_IF_ERROR(
+    BatchNormHelper::ValidateInputs(X, scale, B, mean, var));
+
+  mkldnn::memory::dims src_dims_mkl(
+    X->Shape().GetDims().begin(), X->Shape().GetDims().end());
+  mkldnn::memory::dims scale_dims_mkl(
+    scale->Shape().GetDims().begin(), scale->Shape().GetDims().end());
+  mkldnn::memory::dims b_dims_mkl(
+    B->Shape().GetDims().begin(), B->Shape().GetDims().end());
+  mkldnn::memory::dims mean_dims_mkl(
+    mean->Shape().GetDims().begin(), mean->Shape().GetDims().end());
+  mkldnn::memory::dims var_dims_mkl(
+    var->Shape().GetDims().begin(), var->Shape().GetDims().end());
+
+  mkldnn::memory::dims dst_dims_mkl(
+    Y->Shape().GetDims().begin(), Y->Shape().GetDims().end());
+
+  try {
+    BatchNormParams batchNorm_params(src_dims_mkl, scale_dims_mkl, 
+      b_dims_mkl, mean_dims_mkl, var_dims_mkl, dst_dims_mkl, 
+      onnxruntime::BatchNorm<T>::epsilon_, num_dimensions);
+    BatchNormPrimitive<T>* batchNorm_primitive = 
+      BatchNormPrimitivePool<T>::Get(batchNorm_params);
+    ONNXRUNTIME_RETURN_IF_NOT(batchNorm_primitive != nullptr);
+    batchNorm_primitive->Compute(src_data, scale_data, b_data, 
+      mean_data, var_data, dst_data, scale_dims_mkl[0]);
+
+  } 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(
+    BatchNormalization,
+    kOnnxDomain,
+    7,
+    kMklDnnExecutionProvider,
+    KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
+    BatchNorm<float>);
+
+}  // namespace mkl_dnn
+}  // namespace onnxruntime

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

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