ROIAlign: add cuda kernel implementation (#1823)

onnxruntime/core/providers/cpu/object_detection/roialign.cc

@@ -168,7 +168,7 @@ void RoiAlignForward(
     const T* bottom_rois,
     int64_t num_roi_cols,
     T* top_data,
-    const std::string& mode,
+    RoiAlignMode mode,
     const int64_t* batch_indices_ptr,
     const ThreadPool* ttp) {
   int64_t n_rois = nthreads / channels / pooled_width / pooled_height;
@@ -231,7 +231,7 @@ void RoiAlignForward(
           int64_t index = index_n_c + ph * pooled_width + pw;
 
           T output_val = 0.;
-          if (mode == "avg") {  // avg pooling
+          if (mode == RoiAlignMode::avg) {  // avg pooling
             for (int64_t iy = 0; iy < roi_bin_grid_h; iy++) {
               for (int64_t ix = 0; ix < roi_bin_grid_w; ix++) {
                 PreCalc<T> pc = pre_calc[pre_calc_index];
@@ -334,22 +334,22 @@ Status RoiAlign<T>::Compute(OpKernelContext* context) const {
     return status;
   }
 
-  auto& Y = *context->Output(0, {num_rois, x_dims[1], output_height_, output_width_});
+  auto& Y = *context->Output(0, {num_rois, x_dims[1], this->output_height_, this->output_width_});
   int64_t output_size = Y.Shape().Size();
   RoiAlignForward<T>(
       output_size,  // num threads
       X_ptr->Data<T>(),
-      spatial_scale_,
+      this->spatial_scale_,
       x_dims[1],  // num channels
       x_dims[2],  // height
       x_dims[3],  // width
-      output_height_,
-      output_width_,
-      sampling_ratio_,
+      this->output_height_,
+      this->output_width_,
+      this->sampling_ratio_,
       rois_ptr->Data<T>(),
       num_roi_cols,
       Y.template MutableData<T>(),
-      mode_,
+      this->mode_,
       batch_indices_ptr->Data<int64_t>(),
       static_cast<OpKernelContextInternal*>(context)->GetOperatorThreadPool());
 

onnxruntime/core/providers/cpu/object_detection/roialign.h

@@ -10,18 +10,23 @@ namespace onnxruntime {
 
 Status CheckROIAlignValidInput(const Tensor* X_ptr, const Tensor* rois_ptr, const Tensor* batch_indices_ptr);
 
+enum struct RoiAlignMode {
+  avg = 0,
+  max
+};
+
 template <typename T>
-class RoiAlign final : public OpKernel {
+class RoiAlignBase {
  public:
-  explicit RoiAlign(const OpKernelInfo& info) : OpKernel(info) {
+  explicit RoiAlignBase(const OpKernelInfo& info) {
     // mode
-    std::string mode_tmp;
-    if (info.GetAttr<std::string>("mode", &mode_tmp).IsOK()) {
-      mode_ = mode_tmp;
-      std::transform(mode_.begin(), mode_.end(), mode_.begin(), [](auto& i) { return static_cast<char>(::tolower(i)); });
-      if (mode_ != "avg" && mode_ != "max") {
-        ORT_THROW("Invalid mode of value ", mode_, " specified. It should be either avg or max");
+    std::string mode;
+    if (info.GetAttr<std::string>("mode", &mode).IsOK()) {
+      std::transform(mode.begin(), mode.end(), mode.begin(), [](auto& i) { return static_cast<char>(::tolower(i)); });
+      if (mode != "avg" && mode != "max") {
+        ORT_THROW("Invalid mode of value ", mode, " specified. It should be either avg or max");
       }
+      mode_ = mode == "avg" ? RoiAlignMode::avg : RoiAlignMode::max;
     }
 
     // output_height
@@ -50,15 +55,25 @@ class RoiAlign final : public OpKernel {
     }
   }
 
-  Status Compute(OpKernelContext* context) const override;
-
- private:
-  std::string mode_{"avg"};
+ protected:
+  RoiAlignMode mode_{RoiAlignMode::avg};
   int64_t output_height_{1};
   int64_t output_width_{1};
   int64_t sampling_ratio_{0};
   float spatial_scale_{1.0f};
 
+ private:
+  ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(RoiAlignBase);
+};
+
+template <typename T>
+class RoiAlign final : public OpKernel, public RoiAlignBase<T> {
+ public:
+  explicit RoiAlign(const OpKernelInfo& info) : OpKernel(info), RoiAlignBase<T>(info) {}
+
+  Status Compute(OpKernelContext* context) const override;
+
+ private:
   ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(RoiAlign);
 };
 }  // namespace onnxruntime

onnxruntime/core/providers/cuda/cuda_execution_provider.cc

@@ -537,6 +537,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain,
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 9, MLFloat16, Less);
 
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, Dropout);
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, float, RoiAlign);
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, double, RoiAlign);
 
 static void RegisterCudaKernels(KernelRegistry& kernel_registry) {
   static const BuildKernelCreateInfoFn function_table[] = {
@@ -874,6 +876,8 @@ static void RegisterCudaKernels(KernelRegistry& kernel_registry) {
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 9, float, Less)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 9, double, Less)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 9, MLFloat16, Less)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, float, RoiAlign)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, double, RoiAlign)>,
   };
 
   for (auto& function_table_entry : function_table) {

onnxruntime/core/providers/cuda/object_detection/roialign.cc

@@ -0,0 +1,75 @@
+// Copyright (c) Microsoft Corporation. All rights reserved. 
+// Licensed under the MIT License. 
+
+#include "roialign.h"
+#include "roialign_impl.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+#define REGISTER_KERNEL_TYPED(T)                                         \
+  ONNX_OPERATOR_TYPED_KERNEL_EX(                                         \
+      RoiAlign,                                                          \
+      kOnnxDomain,                                                       \
+      10,                                                                \
+      T,                                                                 \
+      kCudaExecutionProvider,                                            \
+      KernelDefBuilder()                                                 \
+          .TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
+          .TypeConstraint("T2", DataTypeImpl::GetTensorType<int64_t>()), \
+      RoiAlign<T>);
+
+template <typename T>
+Status RoiAlign<T>::ComputeInternal(OpKernelContext* context) const {
+  // X
+  const auto* X_ptr = context->Input<Tensor>(0);
+  // rois
+  const auto* rois_ptr = context->Input<Tensor>(1);
+  // batch indices
+  const auto* batch_indices_ptr = context->Input<Tensor>(2);
+
+  const auto& x_dims = X_ptr->Shape();
+  const auto& rois_dims = rois_ptr->Shape();
+  const auto& batch_indices_dims = batch_indices_ptr->Shape();
+
+  auto num_rois = batch_indices_dims[0];
+  auto num_roi_cols = rois_dims[1];
+
+  auto status = CheckROIAlignValidInput(X_ptr, rois_ptr, batch_indices_ptr);
+  if (status != Status::OK()) {
+    return status;
+  }
+
+  auto& Y = *context->Output(0, {num_rois, x_dims[1], this->output_height_, this->output_width_});
+  int64_t output_size = Y.Shape().Size();
+
+  RoiAlignImpl(
+      output_size,  // num threads
+      reinterpret_cast<const typename ToCudaType<T>::MappedType*>(X_ptr->template Data<T>()),
+      ToCudaType<T>::FromFloat(this->spatial_scale_),
+      x_dims[1],  // num channels
+      x_dims[2],  // height
+      x_dims[3],  // width
+      this->output_height_,
+      this->output_width_,
+      this->sampling_ratio_,
+      reinterpret_cast<const typename ToCudaType<T>::MappedType*>(rois_ptr->template Data<T>()),
+      num_roi_cols,
+      reinterpret_cast<typename ToCudaType<T>::MappedType*>(Y.template MutableData<T>()),
+      this->mode_ == RoiAlignMode::avg,
+      batch_indices_ptr->template Data<int64_t>()
+  );
+
+  return Status::OK();
+}
+
+#define SPECIALIZED_COMPUTE(T) \
+  REGISTER_KERNEL_TYPED(T)     \
+  template Status RoiAlign<T>::ComputeInternal(OpKernelContext* ctx) const;
+
+SPECIALIZED_COMPUTE(float)
+SPECIALIZED_COMPUTE(double)
+//SPECIALIZED_COMPUTE(MLFloat16)
+
+}  // namespace cuda
+};  // namespace onnxruntime

onnxruntime/core/providers/cuda/object_detection/roialign.h

@@ -0,0 +1,24 @@
+// Copyright (c) Microsoft Corporation. All rights reserved. 
+// Licensed under the MIT License. 
+
+#pragma once
+
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+#include "core/providers/cuda/cuda_common.h"
+#include "core/providers/cpu/object_detection/roialign.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+template <typename T>
+struct RoiAlign final : CudaKernel, RoiAlignBase<T> {
+  RoiAlign(const OpKernelInfo& info) : CudaKernel(info), RoiAlignBase<T>(info) {}
+
+  Status ComputeInternal(OpKernelContext* context) const override;
+
+ private:
+  ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(RoiAlign);
+};
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/core/providers/cuda/object_detection/roialign_impl.cu

@@ -0,0 +1,230 @@
+/**
+* Copyright (c) 2016-present, Facebook, Inc.
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*     http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+/* Modifications Copyright (c) Microsoft. */
+
+#include "roialign_impl.h"
+#include "core/providers/cuda/cu_inc/common.cuh"
+
+namespace onnxruntime {
+namespace cuda {
+  
+template <typename T>
+__device__ T bilinear_interpolate(
+    const T* bottom_data,
+    const int height,
+    const int width,
+    T y,
+    T x,
+    const bool is_mode_avg,
+    const int index /* index for debug only*/) {
+  // deal with cases that inverse elements are out of feature map boundary
+  if (y < -1.0 || y > height || x < -1.0 || x > width) {
+    // empty
+    return 0;
+  }
+
+  if (y <= 0) {
+    y = 0;
+  }
+  if (x <= 0) {
+    x = 0;
+  }
+
+  int y_low = (int)y;
+  int x_low = (int)x;
+  int y_high;
+  int x_high;
+
+  if (y_low >= height - 1) {
+    y_high = y_low = height - 1;
+    y = (T)y_low;
+  } else {
+    y_high = y_low + 1;
+  }
+
+  if (x_low >= width - 1) {
+    x_high = x_low = width - 1;
+    x = (T)x_low;
+  } else {
+    x_high = x_low + 1;
+  }
+
+  T ly = y - y_low;
+  T lx = x - x_low;
+  T hy = 1. - ly, hx = 1. - lx;
+  // do bilinear interpolation
+  T v1 = bottom_data[y_low * width + x_low];
+  T v2 = bottom_data[y_low * width + x_high];
+  T v3 = bottom_data[y_high * width + x_low];
+  T v4 = bottom_data[y_high * width + x_high];
+  T w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
+
+  T val = is_mode_avg
+            ? (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4)  // mode Avg
+            : max(max(max(w1 * v1, w2 * v2), w3 * v3), w4 * v4);  // mode Max
+
+  return val;
+}
+
+template <typename T>
+__global__ void RoIAlignForward(
+    const int64_t nthreads,
+    const T* bottom_data,
+    const T spatial_scale,
+    const int64_t channels,
+    const int64_t height,
+    const int64_t width,
+    const int64_t pooled_height,
+    const int64_t pooled_width,
+    const int64_t sampling_ratio,
+    const T* bottom_rois,
+    int64_t roi_cols,
+    T* top_data,
+    const bool is_mode_avg,
+    const int64_t* batch_indices_ptr) {
+  for (size_t index = blockIdx.x * blockDim.x + threadIdx.x; index < nthreads; index += blockDim.x * gridDim.x) {
+    // (n, c, ph, pw) is an element in the pooled output
+    int pw = index % pooled_width;
+    int ph = (index / pooled_width) % pooled_height;
+    int c = (index / pooled_width / pooled_height) % channels;
+    int n = index / pooled_width / pooled_height / channels;
+
+    // RoI could have 4 or 5 columns
+    const T* offset_bottom_rois = bottom_rois + n * roi_cols;
+    const auto roi_batch_ind = batch_indices_ptr[n];
+
+    bool continuous_coordinate = false;
+    // Do not using rounding; this implementation detail is critical
+    T roi_offset = continuous_coordinate ? T(0.5) : T(0);
+    T roi_start_w = offset_bottom_rois[0] * spatial_scale - roi_offset;
+    T roi_start_h = offset_bottom_rois[1] * spatial_scale - roi_offset;
+    T roi_end_w = offset_bottom_rois[2] * spatial_scale - roi_offset;
+    T roi_end_h = offset_bottom_rois[3] * spatial_scale - roi_offset;
+
+    T roi_width = roi_end_w - roi_start_w;
+    T roi_height = roi_end_h - roi_start_h;
+    if (!continuous_coordinate) { // backward compatiblity
+      // Force malformed ROIs to be 1x1
+      roi_width = max(roi_width, (T)1.);
+      roi_height = max(roi_height, (T)1.);
+    }
+    T bin_size_h = static_cast<T>(roi_height) / static_cast<T>(pooled_height);
+    T bin_size_w = static_cast<T>(roi_width) / static_cast<T>(pooled_width);
+
+    const T* offset_bottom_data =
+        bottom_data + static_cast<int64_t>((roi_batch_ind * channels + c) * height * width);
+
+    // We use roi_bin_grid to sample the grid and mimic integral
+    int roi_bin_grid_h = (sampling_ratio > 0)
+        ? sampling_ratio
+        : ceil(roi_height / pooled_height); // e.g., = 2
+    int roi_bin_grid_w =
+        (sampling_ratio > 0) ? sampling_ratio : ceil(roi_width / pooled_width);
+
+    // We do average (integral) pooling inside a bin
+    const T count = roi_bin_grid_h * roi_bin_grid_w; // e.g. = 4
+
+    T output_val = 0.;
+    bool max_flag = false;
+    for (int iy = 0; iy < roi_bin_grid_h; iy++) // e.g., iy = 0, 1
+    {
+      const T y = roi_start_h + ph * bin_size_h +
+          static_cast<T>(iy + .5f) * bin_size_h /
+              static_cast<T>(roi_bin_grid_h); // e.g., 0.5, 1.5
+      for (int ix = 0; ix < roi_bin_grid_w; ix++) {
+        const T x = roi_start_w + pw * bin_size_w +
+            static_cast<T>(ix + .5f) * bin_size_w /
+                static_cast<T>(roi_bin_grid_w);
+
+        T val = bilinear_interpolate(
+            offset_bottom_data, height, width, y, x, is_mode_avg, index);
+        
+        if (is_mode_avg) {
+          output_val += val;
+        } else {
+          if (!max_flag) {
+            output_val = val;
+            max_flag = true;
+          } else {
+            output_val = max(output_val, val);
+          }
+        }
+      }
+    }
+    if (is_mode_avg) {
+      output_val /= count;
+    }
+
+    top_data[index] = output_val;
+  }
+}
+
+template <typename T>
+void RoiAlignImpl(
+  const int64_t nthreads,
+  const T* bottom_data,
+  const T spatial_scale,
+  const int64_t channels,
+  const int64_t height,
+  const int64_t width,
+  const int64_t pooled_height,
+  const int64_t pooled_width,
+  const int64_t sampling_ratio,
+  const T* bottom_rois,
+  int64_t roi_cols,
+  T* top_data,
+  const bool is_mode_avg,
+  const int64_t* batch_indices_ptr) {
+    int blocksPerGrid = (int)(ceil(static_cast<float>(nthreads) / GridDim::maxThreadsPerBlock)); 
+    RoIAlignForward<T><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+      nthreads,
+      bottom_data,
+      spatial_scale,
+      channels,
+      height,
+      width,
+      pooled_height,
+      pooled_width,
+      sampling_ratio,
+      bottom_rois,
+      roi_cols,
+      top_data,
+      is_mode_avg,
+      batch_indices_ptr);    
+}
+
+#define SPECIALIZED_IMPL(T)                     \
+  template void RoiAlignImpl<T>(                \
+        const int64_t nthreads,                 \
+        const T* bottom_data,                   \
+        const T spatial_scale,                  \
+        const int64_t channels,                 \
+        const int64_t height,                   \
+        const int64_t width,                    \
+        const int64_t pooled_height,            \
+        const int64_t pooled_width,             \
+        const int64_t sampling_ratio,           \
+        const T* bottom_rois,                   \
+        int64_t roi_cols,                       \
+        T* top_data,                            \
+        const bool is_mode_avg,                 \
+        const int64_t* batch_indices_ptr);
+
+SPECIALIZED_IMPL(float)
+SPECIALIZED_IMPL(double)
+  
+} // namespace cuda
+} // namespace onnxruntime

onnxruntime/core/providers/cuda/object_detection/roialign_impl.h

@@ -0,0 +1,31 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+#include <stdint.h>
+#include "core/providers/cuda/shared_inc/cuda_utils.h"
+#include "core/framework/data_types.h"
+#include "core/common/common.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+template <typename T>
+void RoiAlignImpl(
+    const int64_t nthreads,
+    const T* bottom_data,
+    const T spatial_scale,
+    const int64_t channels,
+    const int64_t height,
+    const int64_t width,
+    const int64_t pooled_height,
+    const int64_t pooled_width,
+    const int64_t sampling_ratio,
+    const T* bottom_rois,
+    int64_t roi_cols,
+    T* top_data,
+    const bool is_mode_avg,
+    const int64_t* batch_indices_ptr);
+
+}  // namespace cuda
+}  // namespace onnxruntime