Add BatchNorm and Concat to ACL EP (#4190)

* Fix acl padding

* Add BatchNormalization operator to ACL Execution Provider

* Add Concat operator to ACL Execution Provider

Co-authored-by: Andrei-Alexandru <andrei-alexandru.avram@nxp.com>

onnxruntime/core/providers/acl/acl_common.cc

@@ -1,5 +1,5 @@
 // Copyright(C) 2018 Intel Corporation
-// Copyright (c) 2019, NXP Semiconductor, Inc. All rights reserved.
+// Copyright (c) 2019-2020, NXP Semiconductor, Inc. All rights reserved.
 // Licensed under the MIT License
 
 #ifdef _WIN32
@@ -59,5 +59,24 @@ arm_compute::Status ACLImportMemory(arm_compute::TensorAllocator* allocator, voi
 #endif
 }
 
+template <typename T>
+void importDataFromTensor(arm_compute::Tensor* tensor, T* data){
+
+  arm_compute::Window aclInpuWindow;
+  aclInpuWindow.use_tensor_dimensions(tensor->info()->tensor_shape());
+
+  arm_compute::Iterator aclInputIt(tensor, aclInpuWindow);
+  int index = 0;
+  // copy input tensor into the larger buffer
+  arm_compute::execute_window_loop(
+      aclInpuWindow,
+      [&](const arm_compute::Coordinates& co) {
+        data[index] = *reinterpret_cast<float*>(aclInputIt.ptr());
+        index++;
+      },
+      aclInputIt);
+}
+template void importDataFromTensor<float>(arm_compute::Tensor*, float*);
+
 }  // namespace acl
 }  // namespace onnxruntime

onnxruntime/core/providers/acl/acl_common.h

@@ -1,5 +1,5 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
-// Copyright (c) 2019, NXP Semiconductor, Inc. All rights reserved.
+// Copyright (c) 2019-2020, NXP Semiconductor, Inc. All rights reserved.
 // Licensed under the MIT License.
 
 #pragma once
@@ -18,6 +18,8 @@ arm_compute::TensorShape ACLTensorShape(const TensorShape& tensorShape, unsigned
 void ACLPrintTensorShape(const char*, arm_compute::Tensor& t);
 std::shared_ptr<arm_compute::MemoryManagerOnDemand> ACLCreateMemoryManager();
 arm_compute::Status ACLImportMemory(arm_compute::TensorAllocator* allocator, void* memory, size_t size);
+template <typename T>
+void importDataFromTensor(arm_compute::Tensor* tensor, T* data);
 
 }  // namespace acl
 }  // namespace onnxruntime

onnxruntime/core/providers/acl/acl_execution_provider.cc

@@ -37,6 +37,9 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kAclExecutionProvider, kOn
 // Opset 10
 class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kAclExecutionProvider, kOnnxDomain, 10, 10, float, AveragePool);
 
+class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kAclExecutionProvider, kOnnxDomain, 7, 9, BatchNormalization);
+class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kAclExecutionProvider, kOnnxDomain, 4, 10, Concat);
+
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kAclExecutionProvider, kMSDomain, 1, float, FusedConv);
 
 static void RegisterACLKernels(KernelRegistry& kernel_registry) {
@@ -61,6 +64,9 @@ static void RegisterACLKernels(KernelRegistry& kernel_registry) {
   // Opset 10
   kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kAclExecutionProvider, kOnnxDomain, 10, 10, float, AveragePool)>());
 
+  kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kAclExecutionProvider, kOnnxDomain, 7, 9, BatchNormalization)>());
+  kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kAclExecutionProvider, kOnnxDomain, 4, 10, Concat)>());
+
   kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kAclExecutionProvider, kMSDomain, 1, float, FusedConv)>());
 }
 

onnxruntime/core/providers/acl/math/gemm.h

@@ -21,25 +21,6 @@
 //NEON
 #include "arm_compute/runtime/NEON/functions/NEFullyConnectedLayer.h"
 
-template <typename T>
-void importDataFromTensor(arm_compute::Tensor* tensor, T* data){
-
-  arm_compute::Window aclInpuWindow;
-  aclInpuWindow.use_tensor_dimensions(tensor->info()->tensor_shape());
-
-  arm_compute::Iterator aclInputIt(tensor, aclInpuWindow);
-  const unsigned int aclWidth = tensor->info()->dimension(0);
-  const unsigned int aclHeight = tensor->info()->dimension(1);
-
-  // copy input tensor into the larger buffer
-  arm_compute::execute_window_loop(
-      aclInpuWindow,
-      [&](const arm_compute::Coordinates& co) {
-        data[co.z() * (aclWidth * aclHeight) + co.y() * aclWidth + co.x()] = *reinterpret_cast<float*>(aclInputIt.ptr());
-      },
-      aclInputIt);
-}
-
 namespace onnxruntime {
 namespace acl {
 

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

@@ -0,0 +1,155 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) 2020, NXP Semiconductor, Inc. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/common/common.h"
+#include "core/common/exceptions.h"
+#include "core/framework/op_kernel.h"
+#include "core/providers/cpu/nn/autopad_type.h"
+#include "core/framework/tensor.h"
+#include "core/util/math_cpuonly.h"
+
+#include <thread>
+#include <mutex>
+
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+#include "core/util/math.h"
+#include "core/util/math_cpuonly.h"
+
+#include "core/providers/acl/nn/batch_norm.h"
+#include "core/providers/acl/acl_common.h"
+#include "core/providers/acl/acl_fwd.h"
+#include "core/providers/cpu/nn/batch_norm_helper.h"
+
+// ACL
+#include "arm_compute/runtime/Tensor.h"
+#include "arm_compute/core/TensorInfo.h"
+#include "arm_compute/runtime/TensorAllocator.h"
+
+// NEON
+#include "arm_compute/runtime/NEON/functions/NEBatchNormalizationLayer.h"
+
+namespace onnxruntime {
+namespace acl {
+
+template <typename T>
+thread_local std::map<OpKernel*, ACLNEBatchNorm> BatchNorm<T>::batchNormLayers;
+
+template <typename T>
+Status BatchNorm<T>::Compute(OpKernelContext* context) const {
+
+  const Tensor* X = context->Input<Tensor>(0);
+  const Tensor* S = context->Input<Tensor>(1);//scale
+  const Tensor* B = context->Input<Tensor>(2);
+  const Tensor* M = context->Input<Tensor>(3);//mean
+  const Tensor* V = context->Input<Tensor>(4);//var
+
+  ORT_RETURN_IF_ERROR(BatchNormHelper::ValidateInputs(X, S, B, M, V));
+
+  const T* x_data = X->template Data<T>();
+
+  Tensor* Y = context->Output(0, X->Shape());
+
+  T* y_data = Y->template MutableData<T>();
+
+  ACLNEBatchNorm* pBatchNorm;
+  BatchNormLayersIterator it = batchNormLayers.find((OpKernel*)this);
+  if (it == batchNormLayers.end()) {
+
+    ACLNEBatchNorm tbatch_norm;
+    tbatch_norm.in = std::make_shared<arm_compute::Tensor>();
+    tbatch_norm.mean = std::make_shared<arm_compute::Tensor>();
+    tbatch_norm.var = std::make_shared<arm_compute::Tensor>();
+    tbatch_norm.b = std::make_shared<arm_compute::Tensor>();
+    tbatch_norm.scale = std::make_shared<arm_compute::Tensor>();
+    tbatch_norm.out = std::make_shared<arm_compute::Tensor>();
+
+    auto layer = std::make_shared<arm_compute::NEBatchNormalizationLayer>();
+
+    tbatch_norm.in->allocator()->init(arm_compute::TensorInfo(ACLTensorShape(X->Shape()), arm_compute::Format::F32));
+    tbatch_norm.out->allocator()->init(arm_compute::TensorInfo(tbatch_norm.in->info()->tensor_shape(), arm_compute::Format::F32));
+
+    tbatch_norm.scale->allocator()->init(arm_compute::TensorInfo(ACLTensorShape(S->Shape()), arm_compute::Format::F32));
+    tbatch_norm.b->allocator()->init(arm_compute::TensorInfo(ACLTensorShape(B->Shape()), arm_compute::Format::F32));
+    tbatch_norm.mean->allocator()->init(arm_compute::TensorInfo(ACLTensorShape(M->Shape()), arm_compute::Format::F32));
+    tbatch_norm.var->allocator()->init(arm_compute::TensorInfo(ACLTensorShape(V->Shape()), arm_compute::Format::F32));
+
+    layer->configure(tbatch_norm.in.get(), tbatch_norm.out.get(),
+      tbatch_norm.mean.get(), tbatch_norm.var.get(), B != nullptr ? tbatch_norm.b.get() : nullptr, S != nullptr ? tbatch_norm.scale.get() : nullptr,
+      epsilon_);//no activation in onnx
+
+    const T* scale_data = S->template Data<T>();
+    const T* b_data = B->template Data<T>();
+    const T* mean_data = M->template Data<T>();
+    const T* var_data = V->template Data<T>();
+
+    ACLImportMemory(tbatch_norm.mean->allocator(), (void*)mean_data, M->Shape().Size() * 4);
+    ACLImportMemory(tbatch_norm.var->allocator(), (void*)var_data, V->Shape().Size() * 4);
+    ACLImportMemory(tbatch_norm.b->allocator(), (void*)b_data, B->Shape().Size() * 4);
+    ACLImportMemory(tbatch_norm.scale->allocator(), (void*)scale_data, S->Shape().Size() * 4);
+
+    // allocate space for input tensor to accomodate paddings and strides
+    tbatch_norm.in->allocator()->allocate();
+
+    tbatch_norm.layer = std::move(layer);
+
+    std::pair<BatchNormLayersIterator, bool> ret;
+    ret = batchNormLayers.insert(std::pair<OpKernel*, ACLNEBatchNorm>((OpKernel*)this, tbatch_norm));
+    pBatchNorm = &tbatch_norm;
+
+  } else {
+    pBatchNorm = &it->second;
+  }
+
+
+  if(X->Shape().Size() != 0 && pBatchNorm->in->info()->has_padding() ){
+    arm_compute::Window aclInpuWindow;
+    aclInpuWindow.use_tensor_dimensions(pBatchNorm->in->info()->tensor_shape());
+
+    arm_compute::Iterator aclInputIt(pBatchNorm->in.get(), aclInpuWindow);
+    const unsigned int aclWidth = pBatchNorm->in->info()->dimension(0);
+    const unsigned int aclHeight = pBatchNorm->in->info()->dimension(1);
+
+    // copy input tensor into the larger buffer
+    arm_compute::execute_window_loop(
+      aclInpuWindow,
+      [&](const arm_compute::Coordinates& co) {
+        *reinterpret_cast<float*>(aclInputIt.ptr()) = x_data[co.z() * (aclWidth * aclHeight) + co.y() * aclHeight + co.x()];
+      },
+      aclInputIt);
+  }else{
+    ACLImportMemory(pBatchNorm->in->allocator(), (void*)x_data, X->Shape().Size() * 4);
+  }
+
+
+  if(Y->Shape().Size() != 0 && pBatchNorm->out->info()->has_padding() ){
+    pBatchNorm->out->allocator()->allocate();
+  } else {
+    ACLImportMemory(pBatchNorm->out->allocator(), (void*)y_data, Y->Shape().Size() * 4);
+  }
+
+  pBatchNorm->layer->run();
+
+  if(Y->Shape().Size() != 0 && pBatchNorm->out->info()->has_padding() ){
+      importDataFromTensor<T>(pBatchNorm->out.get(), y_data);
+  }
+
+  return Status::OK();
+}
+
+ONNX_OPERATOR_VERSIONED_KERNEL_EX(
+    BatchNormalization,
+    kOnnxDomain,
+    7, 9,
+    kAclExecutionProvider,
+    KernelDefBuilder()
+      .TypeConstraint("X", DataTypeImpl::GetTensorType<float>())
+      .TypeConstraint("scale", DataTypeImpl::GetTensorType<float>())
+      .TypeConstraint("B", DataTypeImpl::GetTensorType<float>())
+      .TypeConstraint("mean", DataTypeImpl::GetTensorType<float>())
+      .TypeConstraint("var", DataTypeImpl::GetTensorType<float>()),
+    BatchNorm<float>);
+
+}  // namespace acl
+}  // namespace onnxruntime

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

@@ -0,0 +1,61 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) 2020, NXP Semiconductor, Inc. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+#include "core/framework/op_kernel.h"
+#include "core/providers/cpu/nn/batch_norm.h"
+#include "core/providers/acl/acl_execution_provider.h"
+
+// ACL
+#include "arm_compute/runtime/Tensor.h"
+#include "arm_compute/core/TensorInfo.h"
+#include "arm_compute/runtime/TensorAllocator.h"
+#include "arm_compute/runtime/Allocator.h"
+#include "arm_compute/runtime/PoolManager.h"
+#include "arm_compute/runtime/BlobLifetimeManager.h"
+#include "arm_compute/runtime/MemoryManagerOnDemand.h"
+
+// NEON
+#include "arm_compute/runtime/NEON/functions/NEBatchNormalizationLayer.h"
+
+namespace onnxruntime {
+namespace acl {
+
+typedef struct {
+  std::shared_ptr<arm_compute::NEBatchNormalizationLayer> layer;
+  std::shared_ptr<arm_compute::Tensor> in, out;
+  std::shared_ptr<arm_compute::Tensor> scale, b, mean, var;
+} ACLNEBatchNorm;
+
+typedef std::map<OpKernel*, ACLNEBatchNorm>::iterator BatchNormLayersIterator;
+
+template <typename T>
+class BatchNorm final : public OpKernel {
+ public:
+  explicit BatchNorm(const OpKernelInfo& info) : OpKernel(info) {
+    auto st = info.GetAttr<float>("epsilon", &epsilon_);
+    ORT_ENFORCE(st.IsOK(), st.ErrorMessage());
+
+    provider_ = (const_cast<ACLExecutionProvider*>(
+        dynamic_cast<const ACLExecutionProvider*>(info.GetExecutionProvider())));
+  }
+
+  ~BatchNorm() {
+	batchNormLayers.erase(this);
+  }
+
+  Status Compute(OpKernelContext* context) const override;
+
+ protected:
+   float epsilon_;
+
+ private:
+  ACLExecutionProvider* provider_;
+  static thread_local std::map<OpKernel*, ACLNEBatchNorm> batchNormLayers;
+};
+
+
+
+}  // namespace acl
+}  // namespace onnxruntime

onnxruntime/core/providers/acl/nn/conv.cc

@@ -156,6 +156,7 @@ Status Conv<T>::Compute(OpKernelContext* context) const {
     aclStrides[1] = strides[0];
 
     std::vector<int64_t> aclPads(4);
+    // The pad order in acl is: pad_left, pad_right, pad_top, pad_bottom
     if (pads.size() == 2) {
       if (strides.size() == 1) {
         aclPads[0] = 0;
@@ -170,8 +171,8 @@ Status Conv<T>::Compute(OpKernelContext* context) const {
       }
     } else {
       aclPads[0] = pads[1];
-      aclPads[1] = pads[0];
-      aclPads[2] = pads[3];
+      aclPads[1] = pads[3];
+      aclPads[2] = pads[0];
       aclPads[3] = pads[2];
     }
 

onnxruntime/core/providers/acl/nn/pool.cc

@@ -68,6 +68,7 @@ ACLNEPool PoolOperation(onnxruntime::OpKernelContext* context,
       aclStrides[1] = strides[0];
 
       std::vector<int64_t> aclPads(4);
+    // The pad order in acl is: pad_left, pad_right, pad_top, pad_bottom
       if (pads.size() == 2) {
         if (strides.size() == 1) {
           aclPads[0] = 0;
@@ -82,8 +83,8 @@ ACLNEPool PoolOperation(onnxruntime::OpKernelContext* context,
         }
       } else {
         aclPads[0] = pads[1];
-        aclPads[1] = pads[0];
-        aclPads[2] = pads[3];
+        aclPads[1] = pads[3];
+        aclPads[2] = pads[0];
         aclPads[3] = pads[2];
       }
 
@@ -96,7 +97,9 @@ ACLNEPool PoolOperation(onnxruntime::OpKernelContext* context,
 
       arm_compute::Size2D aclSize(aclKernelShape[0], aclKernelShape[1]);
 
-      arm_compute::PoolingLayerInfo pool_info(pool_type, aclSize, aclPadStride);
+      bool excludePadding = (pool_type == arm_compute::PoolingType::AVG && pool_attrs.count_include_pad) ? false : true;
+
+      arm_compute::PoolingLayerInfo pool_info(pool_type, aclSize, aclPadStride, excludePadding);
       layer->configure(tpool.in.get(), tpool.out.get(), pool_info);
     }
 

onnxruntime/core/providers/acl/tensor/concat.cc

@@ -0,0 +1,123 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) 2020, NXP Semiconductor, Inc. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/acl/tensor/concat.h"
+#include "core/providers/common.h"
+#include "core/framework/TensorSeq.h"
+#include "arm_compute/core/utils/misc/ShapeCalculator.h"
+
+#include "core/providers/acl/acl_common.h"
+#include "core/providers/acl/acl_fwd.h"
+
+#define PREF_DIM 4
+
+namespace onnxruntime {
+namespace acl {
+
+template <typename T>
+Status Concat<T>::Compute(OpKernelContext* ctx) const {
+
+  // Number of input tensors to concatenate
+  auto input_count = Node().InputArgCount().front();
+
+  // Hold pointers to the input tensors to be used in the PrepareForCompute() step
+  std::vector<const Tensor*> input_tensors;
+  input_tensors.reserve(input_count);
+  for (int i = 0; i < input_count; ++i) {
+    input_tensors.push_back(ctx->Input<Tensor>(i));
+  }
+
+  std::vector<int64_t> output_dims = input_tensors[0]->Shape().GetDims();
+
+  // 'Concat' mode
+  if (!is_stack_) {
+    // While concating, the rank of the output is the same as the input rank(s)
+
+    // Calculate the size of the concatenated axis
+    size_t concat_axis_size = 0;
+    for (int64_t index = 0; index < input_count; index++) {
+      concat_axis_size += input_tensors[index]->Shape()[static_cast<int>(axis_)];
+    }
+
+    output_dims[axis_] = concat_axis_size;
+  } else { // 'Stack' mode
+    // While stacking, the rank of the output is one more than the input rank(s).
+    // Stacking may be thought of as adding an unit dimension (of value 1) in the input tensors,
+    // and concatenating them on thie new axis.
+    // The value in the corresponding axis of the output will be the number of inputs that are being stacked.
+    output_dims.insert(output_dims.begin() + axis_, static_cast<int64_t>(input_count));
+  }
+
+  if(output_dims.size() > 4 || axis_ > 3)
+    return onnxruntime::Concat::Compute(ctx);
+
+  TensorShape output_shape(output_dims);
+  Tensor* Y = ctx->Output(0, output_shape);
+
+  arm_compute::Tensor output;
+  std::vector<arm_compute::ITensor*> inputs_vector;
+  for (int i = 0; i < input_count; i++) {
+    arm_compute::Tensor* input = new arm_compute::Tensor();
+    auto X = input_tensors[i];
+    input->allocator()->init(arm_compute::TensorInfo(ACLTensorShape(X->Shape(), PREF_DIM), arm_compute::Format::F32));
+
+    inputs_vector.push_back(input);
+  }
+
+  arm_compute::NEConcatenateLayer layer;
+  layer.configure(inputs_vector, &output, 3 - axis_);
+
+  for (int i = 0; i < input_count; i++) {
+    auto X = input_tensors[i];
+    const T* x_data = X->template Data<T>();
+    arm_compute::Tensor* in = static_cast<arm_compute::Tensor*>(inputs_vector[i]);
+
+    if (X->Shape().Size() != 0 && in->info()->has_padding() ){
+      in->allocator()->allocate();
+      arm_compute::Window aclInpuWindow;
+      aclInpuWindow.use_tensor_dimensions(in->info()->tensor_shape());
+
+      arm_compute::Iterator aclInputIt(in, aclInpuWindow);
+      int index = 0;
+
+      // copy input tensor into the larger buffer
+      arm_compute::execute_window_loop(
+        aclInpuWindow,
+        [&](const arm_compute::Coordinates& co) {
+          *reinterpret_cast<float*>(aclInputIt.ptr()) = x_data[index];
+          index++;
+        },
+        aclInputIt);
+    } else {
+      ACLImportMemory(in->allocator(), (void*)x_data, X->Shape().Size() * 4);
+    }
+  }
+
+  T* y_data = Y->template MutableData<T>();
+
+  if(Y->Shape().Size() != 0 && output.info()->has_padding() ){
+    output.allocator()->allocate();
+  } else {
+    ACLImportMemory(output.allocator(), (void*)y_data, Y->Shape().Size() * 4);
+  }
+
+  layer.run();
+
+  if (Y->Shape().Size() != 0 && output.info()->has_padding() ){
+    importDataFromTensor<T>(&output, y_data);
+  }
+
+  return Status::OK();
+}
+
+ONNX_OPERATOR_VERSIONED_KERNEL_EX(
+    Concat,
+    kOnnxDomain,
+    4, 10,
+    kAclExecutionProvider,
+    KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
+    Concat<float>);
+
+}  // namespace acl
+}  // namespace onnxruntime

onnxruntime/core/providers/acl/tensor/concat.h

@@ -0,0 +1,45 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) 2020, NXP Semiconductor, Inc. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+#include "core/providers/cpu/tensor/concat.h"
+#include "core/providers/acl/acl_execution_provider.h"
+
+// ACL
+#include "arm_compute/runtime/Tensor.h"
+#include "arm_compute/core/TensorInfo.h"
+#include "arm_compute/runtime/TensorAllocator.h"
+#include "arm_compute/runtime/Allocator.h"
+#include "arm_compute/runtime/PoolManager.h"
+#include "arm_compute/runtime/BlobLifetimeManager.h"
+#include "arm_compute/runtime/MemoryManagerOnDemand.h"
+
+// NEON
+#include "arm_compute/runtime/NEON/functions/NEConcatenateLayer.h"
+
+namespace onnxruntime {
+namespace acl {
+
+template <typename T>
+class Concat : public onnxruntime::Concat {
+ public:
+  explicit Concat(const OpKernelInfo& info) : onnxruntime::Concat(info) {
+
+    provider_ = (const_cast<ACLExecutionProvider*>(
+        dynamic_cast<const ACLExecutionProvider*>(info.GetExecutionProvider())));
+  }
+
+  ~Concat() {}
+
+  Status Compute(OpKernelContext* context) const override;
+
+ private:
+  ACLExecutionProvider* provider_;
+};
+
+
+}  // namespace acl
+}  // namespace onnxruntime

onnxruntime/core/providers/cpu/tensor/concat.h

@@ -51,7 +51,7 @@ class ConcatBase {
   bool is_sequence_op_;
 };
 
-class Concat final : public OpKernel, public ConcatBase {
+class Concat : public OpKernel, public ConcatBase {
  public:
   Concat(const OpKernelInfo& info) : OpKernel(info), ConcatBase(info) {}