/* * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

/*!
 * \file runtime/contrib/tensorrt/tensorrt_ops.cc
 * \brief Converters from Relay ops into TensorRT layers. Converters should
 * inherit from TensorRTOpConverter and implement the Convert() method.
 */

#include "tensorrt_ops.h"

#include <algorithm>
#include <cmath>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>

#include "../json/json_node.h"
#include "NvInfer.h"
#include "tensorrt_utils.h"

namespace tvm {
namespace runtime {
namespace contrib {

TensorRTOpConverter::TensorRTOpConverter(const std::vector<TensorRTInputType>& input_types,
                                         bool variable_input_count)
    : input_types(input_types), variable_input_count(variable_input_count) {}

nvinfer1::ITensor* TensorRTOpConverter::Reshape(TensorRTOpConverterParams* params,
                                                nvinfer1::ITensor* input,
                                                const std::vector<int>& new_shape) const {
  auto layer = params->network->addShuffle(*input);
  ICHECK(layer != nullptr);
  layer->setReshapeDimensions(VectorToTrtDims(new_shape));
  return layer->getOutput(0);
}

nvinfer1::ITensor* TensorRTOpConverter::Transpose(TensorRTOpConverterParams* params,
                                                  nvinfer1::ITensor* input,
                                                  const std::vector<int>& order) const {
  auto layer = params->network->addShuffle(*input);
  ICHECK(layer != nullptr);
  nvinfer1::Permutation perm;
  if (TRT_HAS_IMPLICIT_BATCH(params)) {
    // Batch dimension cannot be modified.
    ICHECK_EQ(input->getDimensions().nbDims, order.size() - 1);
    ICHECK_EQ(order[0], 0);
    for (size_t i = 0; i < order.size(); ++i) {
      perm.order[i] = order[i + 1] - 1;
    }
  } else {
    ICHECK_EQ(input->getDimensions().nbDims, order.size());
    for (size_t i = 0; i < order.size(); ++i) {
      perm.order[i] = order[i];
    }
  }
  layer->setFirstTranspose(perm);
  return layer->getOutput(0);
}

int TensorRTOpConverter::ConvertAxis(TensorRTOpConverterParams* params, int axis,
                                     int input_rank) const {
  // Add 1 for missing batch dim.
  if (TRT_HAS_IMPLICIT_BATCH(params)) {
    input_rank += 1;
  }
  ICHECK(axis >= -input_rank && axis < input_rank);
  if (axis < 0) axis += input_rank;
  if (TRT_HAS_IMPLICIT_BATCH(params)) {
    // Can't modify batch dimenson.
    ICHECK_NE(axis, 0);
    // Subtract 1 for implicit batch dim.
    axis -= 1;
  }
  return axis;
}

nvinfer1::ITensor* TensorRTOpConverter::CreateScalar(
    TensorRTOpConverterParams* params, float value, const nvinfer1::Dims& broadcast_to_dims) const {
  nvinfer1::Dims dims;
  dims.nbDims = broadcast_to_dims.nbDims;
  std::fill_n(dims.d, dims.nbDims, 1);
  float* values = new float[1];
  values[0] = value;
  nvinfer1::Weights weights{nvinfer1::DataType::kFLOAT, static_cast<void*>(values), 1};
  params->trt_weights->push_back(weights);
  return params->network->addConstant(dims, weights)->getOutput(0);
}

void TensorRTOpConverter::GetPadding(const std::vector<std::string>& padding,
                                     bool* use_asymmetric_padding, nvinfer1::DimsHW* prepadding,
                                     nvinfer1::DimsHW* postpadding) const {
  ICHECK(padding.size() == 1 || padding.size() == 2 || padding.size() == 4);
  if (padding.size() == 4) {
    // four int : padding width in the order of (top, left, bottom, right).
    *prepadding = nvinfer1::DimsHW(std::stoi(padding[0]), std::stoi(padding[1]));
    *postpadding = nvinfer1::DimsHW(std::stoi(padding[2]), std::stoi(padding[3]));
    *use_asymmetric_padding = true;
  } else if (padding.size() == 2) {
    // two int : bottom, right will use same padding as top, left
    *prepadding = nvinfer1::DimsHW(std::stoi(padding[0]), std::stoi(padding[1]));
    *postpadding = *prepadding;
    *use_asymmetric_padding = false;
  } else {
    // one int : same padding used on all sides
    *prepadding = nvinfer1::DimsHW(std::stoi(padding[0]), std::stoi(padding[0]));
    *postpadding = *prepadding;
    *use_asymmetric_padding = false;
  }
}

void TensorRTOpConverter::GetPadding3D(const std::vector<std::string>& padding,
                                       bool* use_asymmetric_padding, nvinfer1::Dims* prepadding,
                                       nvinfer1::Dims* postpadding) const {
  ICHECK(padding.size() == 1 || padding.size() == 3 || padding.size() == 6);
  if (padding.size() == 6) {
    // six int : padding width in the order of (front, top, left, back, bottom, right)
    *prepadding =
        nvinfer1::Dims3(std::stoi(padding[0]), std::stoi(padding[1]), std::stoi(padding[2]));
    *postpadding =
        nvinfer1::Dims3(std::stoi(padding[3]), std::stoi(padding[4]), std::stoi(padding[5]));
    *use_asymmetric_padding = true;
  } else if (padding.size() == 3) {
    // three int : back, bottom, right will use same padding as front, top, left
    *prepadding =
        nvinfer1::Dims3(std::stoi(padding[0]), std::stoi(padding[1]), std::stoi(padding[2]));
    *postpadding = *prepadding;
    *use_asymmetric_padding = false;
  } else {
    // one int : same padding used on all sides
    *prepadding =
        nvinfer1::Dims3(std::stoi(padding[0]), std::stoi(padding[0]), std::stoi(padding[0]));
    *postpadding = *prepadding;
    *use_asymmetric_padding = false;
  }
}

class ActivationOpConverter : public TensorRTOpConverter {
 public:
  ActivationOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    static const std::unordered_map<std::string, nvinfer1::ActivationType> op_map = {
      {"nn.relu", nvinfer1::ActivationType::kRELU},
      {"sigmoid", nvinfer1::ActivationType::kSIGMOID},
      {"tanh", nvinfer1::ActivationType::kTANH},
#if TRT_VERSION_GE(5, 1, 5)
      {"clip", nvinfer1::ActivationType::kCLIP},
      {"nn.leaky_relu", nvinfer1::ActivationType::kLEAKY_RELU},
#endif
    };
    auto it = op_map.find(params->op_name);
    ICHECK(it != op_map.end()) << "Unsupported activation type " << params->op_name;
    nvinfer1::IActivationLayer* act_layer =
        params->network->addActivation(*params->inputs.at(0).tensor, it->second);
#if TRT_VERSION_GE(5, 1, 5)
    if (params->op_name == "clip") {
      float a_min = std::stof(params->node.GetAttr<std::vector<std::string>>("a_min")[0]);
      float a_max = std::stof(params->node.GetAttr<std::vector<std::string>>("a_max")[0]);
      act_layer->setAlpha(a_min);
      act_layer->setBeta(a_max);
    } else if (params->op_name == "nn.leaky_relu") {
      float alpha = std::stof(params->node.GetAttr<std::vector<std::string>>("alpha")[0]);
      act_layer->setAlpha(alpha);
    }
#endif
    ICHECK(act_layer != nullptr);
    params->outputs.push_back(act_layer->getOutput(0));
  }
};

class ElementWiseBinaryOpConverter : public TensorRTOpConverter {
 public:
  ElementWiseBinaryOpConverter() : TensorRTOpConverter({kTensor, kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    static const std::unordered_map<std::string, nvinfer1::ElementWiseOperation> op_map = {
        {"add", nvinfer1::ElementWiseOperation::kSUM},
        {"subtract", nvinfer1::ElementWiseOperation::kSUB},
        {"multiply", nvinfer1::ElementWiseOperation::kPROD},
        {"divide", nvinfer1::ElementWiseOperation::kDIV},
        {"power", nvinfer1::ElementWiseOperation::kPOW},
        {"maximum", nvinfer1::ElementWiseOperation::kMAX},
        {"minimum", nvinfer1::ElementWiseOperation::kMIN}};
    auto it = op_map.find(params->op_name);
    ICHECK(it != op_map.end()) << "Unsupported elementwise type " << params->op_name;
    // Broadcast
    auto input0 = params->inputs.at(0).tensor;
    auto input0_dims = TrtDimsToVector(input0->getDimensions());
    auto input1 = params->inputs.at(1).tensor;
    auto input1_dims = TrtDimsToVector(input1->getDimensions());
    const bool need_broadcast = input0_dims.size() != input1_dims.size();
    if (need_broadcast) {
      if (input0_dims.size() < input1_dims.size()) {
        std::vector<int> new_shape(input0_dims);
        while (new_shape.size() < input1_dims.size()) new_shape.insert(new_shape.begin(), 1);
        input0 = Reshape(params, input0, new_shape);
      } else if (input1_dims.size() < input0_dims.size()) {
        std::vector<int> new_shape(input1_dims);
        while (new_shape.size() < input0_dims.size()) new_shape.insert(new_shape.begin(), 1);
        input1 = Reshape(params, input1, new_shape);
      }
    }

    nvinfer1::IElementWiseLayer* elemwise_layer =
        params->network->addElementWise(*input0, *input1, it->second);
    ICHECK(elemwise_layer != nullptr);
    params->outputs.push_back(elemwise_layer->getOutput(0));
  }
};

class Conv2DOpConverter : public TensorRTOpConverter {
 public:
  Conv2DOpConverter() : TensorRTOpConverter({kTensor, kWeight}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input_tensor = params->inputs.at(0).tensor;
    auto input_dims = TrtDimsToVector(input_tensor->getDimensions());
    auto weight_shape = params->inputs.at(1).weight_shape;
    ICHECK_EQ(params->node.GetAttr<std::vector<std::string>>("data_layout")[0], "NCHW");
    ICHECK(params->node.GetAttr<std::vector<std::string>>("out_layout")[0] == "" ||
           params->node.GetAttr<std::vector<std::string>>("out_layout")[0] == "NCHW");
    ICHECK_EQ(params->node.GetAttr<std::vector<std::string>>("kernel_layout")[0], "OIHW");
    auto str_strides = params->node.GetAttr<std::vector<std::string>>("strides");
    auto str_dilation = params->node.GetAttr<std::vector<std::string>>("dilation");
    auto str_padding = params->node.GetAttr<std::vector<std::string>>("padding");
    int groups = std::stoi(params->node.GetAttr<std::vector<std::string>>("groups")[0]);
    int channels = std::stoi(params->node.GetAttr<std::vector<std::string>>("channels")[0]);
    // TRT conv2d op doesn't support asymmetric padding before 5.1, so we
    // workaround by adding a padding layer before the pooling op.
    nvinfer1::DimsHW prepadding, postpadding;
    bool use_asymmetric_padding;
    GetPadding(str_padding, &use_asymmetric_padding, &prepadding, &postpadding);
#if !TRT_VERSION_GE(5, 1, 5)
    if (use_asymmetric_padding) {
      auto pad_layer = params->network->addPadding(*input_tensor, prepadding, postpadding);
      ICHECK(pad_layer != nullptr);
      input_tensor = pad_layer->getOutput(0);
      // No need for conv op to do any padding.
      use_asymmetric_padding = false;
      prepadding = nvinfer1::DimsHW(0, 0);
    }
#endif

    const auto kernel_size = nvinfer1::DimsHW(weight_shape[2], weight_shape[3]);
    nvinfer1::Weights bias{nvinfer1::DataType::kFLOAT, nullptr, 0};
    auto conv_layer = params->network->addConvolution(*input_tensor, channels, kernel_size,
                                                      params->inputs.at(1).weight, bias);
    ICHECK(conv_layer != nullptr);
    if (use_asymmetric_padding) {
#if TRT_VERSION_GE(5, 1, 5)
      conv_layer->setPrePadding(prepadding);
      conv_layer->setPostPadding(postpadding);
#endif
    } else {
      conv_layer->setPadding(prepadding);
    }
    ICHECK_EQ(str_strides.size(), 2);
    const auto strides = nvinfer1::DimsHW(std::stoi(str_strides[0]), std::stoi(str_strides[1]));
    conv_layer->setStride(strides);
    ICHECK_EQ(str_dilation.size(), 2);
    const auto dilation = nvinfer1::DimsHW(std::stoi(str_dilation[0]), std::stoi(str_dilation[1]));
    conv_layer->setDilation(dilation);
    conv_layer->setNbGroups(groups);
    params->outputs.push_back(conv_layer->getOutput(0));
  }
};

#if TRT_VERSION_GE(6, 0, 1)
class Conv3DOpConverter : public TensorRTOpConverter {
 public:
  Conv3DOpConverter() : TensorRTOpConverter({kTensor, kWeight}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input_tensor = params->inputs.at(0).tensor;
    auto input_dims = TrtDimsToVector(input_tensor->getDimensions());
    auto weight_shape = params->inputs.at(1).weight_shape;
    ICHECK_EQ(params->node.GetAttr<std::vector<std::string>>("data_layout")[0], "NCDHW");
    ICHECK(params->node.GetAttr<std::vector<std::string>>("out_layout")[0] == "" ||
           params->node.GetAttr<std::vector<std::string>>("out_layout")[0] == "NCDHW");
    ICHECK_EQ(params->node.GetAttr<std::vector<std::string>>("kernel_layout")[0], "OIDHW");
    auto str_strides = params->node.GetAttr<std::vector<std::string>>("strides");
    auto str_dilation = params->node.GetAttr<std::vector<std::string>>("dilation");
    auto str_padding = params->node.GetAttr<std::vector<std::string>>("padding");
    int groups = std::stoi(params->node.GetAttr<std::vector<std::string>>("groups")[0]);

    nvinfer1::Dims prepadding, postpadding;
    bool use_asymmetric_padding;
    GetPadding3D(str_padding, &use_asymmetric_padding, &prepadding, &postpadding);

    // Could use attrs->channels.as<IntImmNode>()->value
    const int num_outputs = weight_shape[0];
    const auto kernel_size = nvinfer1::Dims3(weight_shape[2], weight_shape[3], weight_shape[4]);
    nvinfer1::Weights bias{nvinfer1::DataType::kFLOAT, nullptr, 0};
    auto conv_layer = params->network->addConvolutionNd(*input_tensor, num_outputs, kernel_size,
                                                        params->inputs.at(1).weight, bias);
    ICHECK(conv_layer != nullptr);
    if (use_asymmetric_padding) {
      conv_layer->setPrePadding(prepadding);
      conv_layer->setPostPadding(postpadding);
    } else {
      conv_layer->setPaddingNd(prepadding);
    }
    ICHECK_EQ(str_strides.size(), 3);
    const auto strides = nvinfer1::Dims3(std::stoi(str_strides[0]), std::stoi(str_strides[1]),
                                         std::stoi(str_strides[2]));
    conv_layer->setStrideNd(strides);
    ICHECK_EQ(str_dilation.size(), 3);
    const auto dilation = nvinfer1::Dims3(std::stoi(str_dilation[0]), std::stoi(str_dilation[1]),
                                          std::stoi(str_dilation[2]));
    conv_layer->setDilationNd(dilation);
    conv_layer->setNbGroups(groups);
    params->outputs.push_back(conv_layer->getOutput(0));
  }
};
#endif  // TRT_VERSION_GE(6, 0, 1)

class DenseOpConverter : public TensorRTOpConverter {
 public:
  DenseOpConverter() : TensorRTOpConverter({kTensor, kWeight}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input_tensor = params->inputs.at(0).tensor;
    auto input_dims = TrtDimsToVector(input_tensor->getDimensions());
    ICHECK(input_dims.size() > 0 && input_dims.size() <= 3);
    const size_t required_rank = TRT_HAS_IMPLICIT_BATCH(params) ? 3 : 4;
    const bool need_reshape_on_input = input_dims.size() != required_rank;
    if (need_reshape_on_input) {
      // Add dims of size 1 until rank is required_rank.
      std::vector<int> new_shape(input_dims);
      while (new_shape.size() < required_rank) new_shape.insert(new_shape.end(), 1);
      input_tensor = Reshape(params, input_tensor, new_shape);
    }
    // Weights are in KC format.
    ICHECK_EQ(params->inputs.at(1).weight_shape.size(), 2);
    const int num_units = params->inputs.at(1).weight_shape[0];
    nvinfer1::Weights bias{nvinfer1::DataType::kFLOAT, nullptr, 0};
    nvinfer1::IFullyConnectedLayer* fc_layer = params->network->addFullyConnected(
        *input_tensor, num_units, params->inputs.at(1).weight, bias);
    ICHECK(fc_layer != nullptr);
    auto output_tensor = fc_layer->getOutput(0);
    if (need_reshape_on_input) {
      // Remove added dims.
      input_dims[input_dims.size() - 1] = num_units;
      output_tensor = Reshape(params, output_tensor, input_dims);
    }
    params->outputs.push_back(output_tensor);
  }
};

class BatchNormOpConverter : public TensorRTOpConverter {
 public:
  BatchNormOpConverter() : TensorRTOpConverter({kTensor, kWeight, kWeight, kWeight, kWeight}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input = params->inputs.at(0).tensor;
    auto gamma = params->inputs.at(1).weight;
    auto beta = params->inputs.at(2).weight;
    auto mean = params->inputs.at(3).weight;
    auto var = params->inputs.at(4).weight;
    ICHECK_EQ(gamma.count, beta.count);
    ICHECK_EQ(gamma.count, mean.count);
    ICHECK_EQ(gamma.count, var.count);
    const float epsilon = std::stof(params->node.GetAttr<std::vector<std::string>>("epsilon")[0]);
    const int axis = std::stoi(params->node.GetAttr<std::vector<std::string>>("axis")[0]);
    const bool scale = std::stoi(params->node.GetAttr<std::vector<std::string>>("scale")[0]);
    const bool center = std::stoi(params->node.GetAttr<std::vector<std::string>>("center")[0]);
    ICHECK(axis == 1 || axis == 3);
    const bool need_transpose = axis == 3;

    void* weight_scale_ptr = new float[gamma.count];
    nvinfer1::Weights weight_scale{nvinfer1::DataType::kFLOAT, weight_scale_ptr, gamma.count};
    params->trt_weights->push_back(weight_scale);
    void* weight_shift_ptr = new float[gamma.count];
    nvinfer1::Weights weight_shift{nvinfer1::DataType::kFLOAT, weight_shift_ptr, gamma.count};
    params->trt_weights->push_back(weight_shift);
    nvinfer1::Weights power{nvinfer1::DataType::kFLOAT, nullptr, 0};

    // fill in the content of weights for the Scale layer
    const float* gamma_ptr = reinterpret_cast<const float*>(gamma.values);
    const float* beta_ptr = reinterpret_cast<const float*>(beta.values);
    const float* mean_ptr = reinterpret_cast<const float*>(mean.values);
    const float* var_ptr = reinterpret_cast<const float*>(var.values);
    float* scale_ptr = reinterpret_cast<float*>(weight_scale_ptr);
    float* shift_ptr = reinterpret_cast<float*>(weight_shift_ptr);
    for (int i = 0; i < gamma.count; ++i) {
      scale_ptr[i] = 1.0 / std::sqrt(var_ptr[i] + epsilon);
      if (scale) {
        scale_ptr[i] *= gamma_ptr[i];
      }
      shift_ptr[i] = -mean_ptr[i] * scale_ptr[i];
      if (center) {
        shift_ptr[i] += beta_ptr[i];
      }
    }
    if (need_transpose) {
      input = Transpose(params, input, {0, 3, 1, 2});
    }
    nvinfer1::IScaleLayer* scale_layer = params->network->addScale(
        *input, nvinfer1::ScaleMode::kCHANNEL, weight_shift, weight_scale, power);
    ICHECK(scale_layer != nullptr);
    auto output = scale_layer->getOutput(0);
    if (need_transpose) {
      output = Transpose(params, output, {0, 2, 3, 1});
    }
    params->outputs.push_back(output);
  }
};

class BatchFlattenOpConverter : public TensorRTOpConverter {
 public:
  BatchFlattenOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    std::vector<int> new_shape{-1};
    if (!TRT_HAS_IMPLICIT_BATCH(params)) {
      new_shape.insert(new_shape.begin(), params->inputs.at(0).tensor->getDimensions().d[0]);
    }
    params->outputs.push_back(Reshape(params, params->inputs.at(0).tensor, new_shape));
  }
};

class SoftmaxOpConverter : public TensorRTOpConverter {
 public:
  SoftmaxOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input = params->inputs.at(0).tensor;
    const int input_rank = input->getDimensions().nbDims;
    const int original_axis = std::stoi(params->node.GetAttr<std::vector<std::string>>("axis")[0]);
    const int axis = ConvertAxis(params, original_axis, input_rank);
    nvinfer1::ISoftMaxLayer* softmax_layer = params->network->addSoftMax(*input);
    softmax_layer->setAxes(1 << axis);
    ICHECK(softmax_layer != nullptr);
    params->outputs.push_back(softmax_layer->getOutput(0));
  }
};

class PoolingOpConverter : public TensorRTOpConverter {
 public:
  PoolingOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input = params->inputs.at(0).tensor;
    static const std::unordered_map<std::string, nvinfer1::PoolingType> op_map = {
        {"nn.max_pool2d", nvinfer1::PoolingType::kMAX},
        {"nn.avg_pool2d", nvinfer1::PoolingType::kAVERAGE}};
    auto it = op_map.find(params->op_name);
    ICHECK(it != op_map.end()) << "Unsupported pooling type " << params->op_name << " in TensorRT";
    ICHECK_EQ(params->node.GetAttr<std::vector<std::string>>("layout")[0], "NCHW");
    auto str_pool_size = params->node.GetAttr<std::vector<std::string>>("pool_size");
    auto str_padding = params->node.GetAttr<std::vector<std::string>>("padding");
    auto str_strides = params->node.GetAttr<std::vector<std::string>>("strides");
    nvinfer1::DimsHW prepadding, postpadding;
    bool use_asymmetric_padding;
    GetPadding(str_padding, &use_asymmetric_padding, &prepadding, &postpadding);
    bool ceil_mode = std::stoi(params->node.GetAttr<std::vector<std::string>>("ceil_mode")[0]);

// TRT pooling op doesn't support asymmetric padding before 5.1, so we
// workaround by adding a padding layer before the pooling op.
#if !TRT_VERSION_GE(5, 1, 5)
    if (use_asymmetric_padding) {
      auto pad_layer = params->network->addPadding(*input, prepadding, postpadding);
      ICHECK(pad_layer != nullptr);
      input = pad_layer->getOutput(0);
      // No need for pooling op to do any padding.
      use_asymmetric_padding = false;
      prepadding = nvinfer1::DimsHW(0, 0);
    }
#endif

    nvinfer1::DimsHW window_size =
        nvinfer1::DimsHW(std::stoi(str_pool_size[0]), std::stoi(str_pool_size[1]));
    auto pool_layer = params->network->addPooling(*input, it->second, window_size);
    ICHECK(pool_layer != nullptr);
    nvinfer1::DimsHW strides =
        nvinfer1::DimsHW(std::stoi(str_strides[0]), std::stoi(str_strides[1]));
    pool_layer->setStride(strides);
    if (use_asymmetric_padding) {
#if TRT_VERSION_GE(5, 1, 5)
      pool_layer->setPrePadding(prepadding);
      pool_layer->setPostPadding(postpadding);
#endif
    } else {
      pool_layer->setPadding(prepadding);
    }
    if (params->op_name == "nn.avg_pool2d") {
      bool count_include_pad =
          std::stoi(params->node.GetAttr<std::vector<std::string>>("count_include_pad")[0]);
      // count_include_pad=True is useless if there is no padding. TRT doesn't
      // like count_include_pad in combination with strides even when there is
      // no padding or assymetric padding even, so turn off inclusive to avoid
      // error message. Note: Padding will always be symmetric with
      // count_include_pad since partitioner will prevent unsupported case.
      if (prepadding.h() == 0 && prepadding.w() == 0) {
        count_include_pad = false;
      }
      pool_layer->setAverageCountExcludesPadding(!count_include_pad);
    }
#if TRT_VERSION_GE(5, 1, 5)
    if (ceil_mode) {
      pool_layer->setPaddingMode(nvinfer1::PaddingMode::kEXPLICIT_ROUND_UP);
    }
#else
    ICHECK(!ceil_mode);
#endif
    params->outputs.push_back(pool_layer->getOutput(0));
  }
};

#if TRT_VERSION_GE(6, 0, 1)
class Pooling3DOpConverter : public TensorRTOpConverter {
 public:
  Pooling3DOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input = params->inputs.at(0).tensor;
    static const std::unordered_map<std::string, nvinfer1::PoolingType> op_map = {
        {"nn.max_pool3d", nvinfer1::PoolingType::kMAX},
        {"nn.avg_pool3d", nvinfer1::PoolingType::kAVERAGE}};
    auto it = op_map.find(params->op_name);
    ICHECK(it != op_map.end()) << "Unsupported pooling type " << params->op_name << " in TensorRT";
    ICHECK_EQ(params->node.GetAttr<std::vector<std::string>>("layout")[0], "NCDHW");
    auto str_pool_size = params->node.GetAttr<std::vector<std::string>>("pool_size");
    auto str_padding = params->node.GetAttr<std::vector<std::string>>("padding");
    auto str_strides = params->node.GetAttr<std::vector<std::string>>("strides");
    nvinfer1::DimsHW prepadding, postpadding;
    bool use_asymmetric_padding;
    GetPadding3D(str_padding, &use_asymmetric_padding, &prepadding, &postpadding);
    bool ceil_mode = std::stoi(params->node.GetAttr<std::vector<std::string>>("ceil_mode")[0]);
    nvinfer1::Dims window_size = nvinfer1::Dims3(
        std::stoi(str_pool_size[0]), std::stoi(str_pool_size[1]), std::stoi(str_pool_size[2]));
    auto pool_layer = params->network->addPoolingNd(*input, it->second, window_size);
    ICHECK(pool_layer != nullptr);
    nvinfer1::Dims strides = nvinfer1::Dims3(std::stoi(str_strides[0]), std::stoi(str_strides[1]),
                                             std::stoi(str_strides[2]));
    pool_layer->setStrideNd(strides);
    if (use_asymmetric_padding) {
      pool_layer->setPrePadding(prepadding);
      pool_layer->setPostPadding(postpadding);
    } else {
      pool_layer->setPaddingNd(prepadding);
    }
    if (params->op_name == "nn.avg_pool3d") {
      bool count_include_pad =
          std::stoi(params->node.GetAttr<std::vector<std::string>>("count_include_pad")[0]);
      pool_layer->setAverageCountExcludesPadding(!count_include_pad);
    }
    if (ceil_mode) {
      pool_layer->setPaddingMode(nvinfer1::PaddingMode::kEXPLICIT_ROUND_UP);
    }
    params->outputs.push_back(pool_layer->getOutput(0));
  }
};
#endif  // TRT_VERSION_GE(6, 0, 1)

class GlobalPoolingOpConverter : public TensorRTOpConverter {
 public:
  GlobalPoolingOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input_tensor = params->inputs.at(0).tensor;
    auto input_dims = TrtDimsToVector(input_tensor->getDimensions());
    static const std::unordered_map<std::string, nvinfer1::PoolingType> op_map = {
        {"nn.global_max_pool2d", nvinfer1::PoolingType::kMAX},
        {"nn.global_avg_pool2d", nvinfer1::PoolingType::kAVERAGE}};
    auto it = op_map.find(params->op_name);
    ICHECK(it != op_map.end()) << "Unsupported pooling type " << params->op_name << " in TensorRT";
    ICHECK_EQ(params->node.GetAttr<std::vector<std::string>>("layout")[0], "NCHW");
    const int h = TRT_HAS_IMPLICIT_BATCH(params) ? input_dims[1] : input_dims[2];
    const int w = TRT_HAS_IMPLICIT_BATCH(params) ? input_dims[2] : input_dims[3];
    auto pool_layer =
        params->network->addPooling(*input_tensor, it->second, nvinfer1::DimsHW(h, w));
    ICHECK(pool_layer != nullptr);
    params->outputs.push_back(pool_layer->getOutput(0));
  }
};

class ExpandDimsOpConverter : public TensorRTOpConverter {
 public:
  ExpandDimsOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input_tensor = params->inputs.at(0).tensor;
    auto input_dims = TrtDimsToVector(input_tensor->getDimensions());
    const int original_axis = std::stoi(params->node.GetAttr<std::vector<std::string>>("axis")[0]);
    const int num_newaxis =
        std::stoi(params->node.GetAttr<std::vector<std::string>>("num_newaxis")[0]);
    const int axis = ConvertAxis(params, original_axis, input_dims.size() + 1);
    for (int i = 0; i < num_newaxis; ++i) {
      input_dims.insert(input_dims.begin() + axis, 1);
    }
    params->outputs.push_back(Reshape(params, params->inputs.at(0).tensor, input_dims));
  }
};

class SqueezeOpConverter : public TensorRTOpConverter {
 public:
  SqueezeOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input_tensor = params->inputs.at(0).tensor;
    auto input_dims = TrtDimsToVector(input_tensor->getDimensions());
    auto str_axis = params->node.GetAttr<std::vector<std::string>>("axis");
    for (size_t i = 0; i < str_axis.size(); ++i) {
      const int axis = ConvertAxis(params, std::stoi(str_axis[i]), input_dims.size());
      input_dims[axis] = 0;
    }
    input_dims.erase(std::remove(input_dims.begin(), input_dims.end(), 0), input_dims.end());
    params->outputs.push_back(Reshape(params, params->inputs.at(0).tensor, input_dims));
  }
};

class UnaryOpConverter : public TensorRTOpConverter {
 public:
  UnaryOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    // The following ops are supported by TRT but don't exist in relay yet:
    // recip, tan, sinh, cosh, asin, acos, asinh, acosh, atanh
    static const std::unordered_map<std::string, nvinfer1::UnaryOperation> op_map = {
      {"exp", nvinfer1::UnaryOperation::kEXP},
      {"log", nvinfer1::UnaryOperation::kLOG},
      {"sqrt", nvinfer1::UnaryOperation::kSQRT},
      {"abs", nvinfer1::UnaryOperation::kABS},
      {"negative", nvinfer1::UnaryOperation::kNEG},
#if TRT_VERSION_GE(5, 1, 5)
      {"sin", nvinfer1::UnaryOperation::kSIN},
      {"cos", nvinfer1::UnaryOperation::kCOS},
      {"atan", nvinfer1::UnaryOperation::kATAN},
      {"ceil", nvinfer1::UnaryOperation::kCEIL},
      {"floor", nvinfer1::UnaryOperation::kFLOOR},
#endif
    };
    auto it = op_map.find(params->op_name);
    ICHECK(it != op_map.end()) << "Unsupported unary type " << params->op_name;
    nvinfer1::IUnaryLayer* unary_layer =
        params->network->addUnary(*params->inputs.at(0).tensor, it->second);
    ICHECK(unary_layer != nullptr);
    params->outputs.push_back(unary_layer->getOutput(0));
  }
};

class ConcatOpConverter : public TensorRTOpConverter {
 public:
  ConcatOpConverter() : TensorRTOpConverter({}, /*variable_input_count=*/true) {}

  void Convert(TensorRTOpConverterParams* params) const {
    const int num_inputs = params->inputs.size();
    ICHECK_GT(num_inputs, 0);
    const int input_rank = params->inputs[0].tensor->getDimensions().nbDims;
    std::vector<nvinfer1::ITensor*> input_tensors;
    for (auto input : params->inputs) {
      ICHECK(input.type == kTensor);
      ICHECK_EQ(input_rank, input.tensor->getDimensions().nbDims);
      input_tensors.push_back(input.tensor);
    }

    const int original_axis = std::stoi(params->node.GetAttr<std::vector<std::string>>("axis")[0]);
    const int axis = ConvertAxis(params, original_axis, input_rank);

    nvinfer1::IConcatenationLayer* concat_layer =
        params->network->addConcatenation(input_tensors.data(), input_tensors.size());
    ICHECK(concat_layer != nullptr);
    concat_layer->setAxis(axis);
    params->outputs.push_back(concat_layer->getOutput(0));
  }
};

class BiasAddOpConverter : public TensorRTOpConverter {
 public:
  BiasAddOpConverter() : TensorRTOpConverter({kTensor, kWeight}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input_tensor = params->inputs.at(0).tensor;
    auto input_dims = TrtDimsToVector(input_tensor->getDimensions());
    const size_t required_rank = TRT_HAS_IMPLICIT_BATCH(params) ? 3 : 4;
    ICHECK(input_dims.size() > 0 && input_dims.size() <= required_rank);
    const bool need_reshape_on_input = input_dims.size() != required_rank;
    if (need_reshape_on_input) {
      // Add dims of size 1 until rank is required_rank.
      std::vector<int> new_shape(input_dims);
      while (new_shape.size() < required_rank) new_shape.insert(new_shape.end(), 1);
      input_tensor = Reshape(params, input_tensor, new_shape);
    }

    nvinfer1::Weights shift{nvinfer1::DataType::kFLOAT, nullptr, 0};
    nvinfer1::Weights power{nvinfer1::DataType::kFLOAT, nullptr, 0};
    nvinfer1::IScaleLayer* scale_layer = params->network->addScale(
        *input_tensor, nvinfer1::ScaleMode::kCHANNEL, params->inputs.at(1).weight, shift, power);
    ICHECK(scale_layer != nullptr);
    auto output_tensor = scale_layer->getOutput(0);
    if (need_reshape_on_input) {
      // Remove added dims.
      output_tensor = Reshape(params, output_tensor, input_dims);
    }
    params->outputs.push_back(output_tensor);
  }
};

class Conv2DTransposeOpConverter : public TensorRTOpConverter {
 public:
  Conv2DTransposeOpConverter() : TensorRTOpConverter({kTensor, kWeight}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input_tensor = params->inputs.at(0).tensor;
    auto weight_shape = params->inputs.at(1).weight_shape;
    ICHECK_EQ(params->node.GetAttr<std::vector<std::string>>("data_layout")[0], "NCHW");
    ICHECK(params->node.GetAttr<std::vector<std::string>>("out_layout")[0] == "" ||
           params->node.GetAttr<std::vector<std::string>>("out_layout")[0] == "NCHW");
    ICHECK_EQ(params->node.GetAttr<std::vector<std::string>>("kernel_layout")[0], "OIHW");
    auto str_dilation = params->node.GetAttr<std::vector<std::string>>("dilation");
    ICHECK(std::stoi(str_dilation[0]) == 1 && std::stoi(str_dilation[1]) == 1);
    auto str_strides = params->node.GetAttr<std::vector<std::string>>("strides");
    auto str_padding = params->node.GetAttr<std::vector<std::string>>("padding");
    auto str_output_padding = params->node.GetAttr<std::vector<std::string>>("output_padding");
    int groups = std::stoi(params->node.GetAttr<std::vector<std::string>>("groups")[0]);

    // TRT deconv op doesn't support asymmetric padding before 5.1, so we
    // workaround by adding a padding layer before the pooling op.
    nvinfer1::DimsHW prepadding, postpadding;
    bool use_asymmetric_padding;
    GetPadding(str_padding, &use_asymmetric_padding, &prepadding, &postpadding);
#if !TRT_VERSION_GE(5, 1, 5)
    if (use_asymmetric_padding) {
      auto pad_layer = params->network->addPadding(*input_tensor, prepadding, postpadding);
      ICHECK(pad_layer != nullptr);
      input_tensor = pad_layer->getOutput(0);
      // No need for conv op to do any padding.
      use_asymmetric_padding = false;
      prepadding = nvinfer1::DimsHW(0, 0);
    }
#endif

    // Could use conv2d_attr->channels.as<IntImmNode>()->value
    const int num_outputs = weight_shape[1];
    const auto kernel_size = nvinfer1::DimsHW(weight_shape[2], weight_shape[3]);
    nvinfer1::Weights bias{nvinfer1::DataType::kFLOAT, nullptr, 0};
    auto deconv_layer = params->network->addDeconvolution(*input_tensor, num_outputs, kernel_size,
                                                          params->inputs.at(1).weight, bias);
    ICHECK(deconv_layer != nullptr);
    if (use_asymmetric_padding) {
#if TRT_VERSION_GE(5, 1, 5)
      deconv_layer->setPrePadding(prepadding);
      deconv_layer->setPostPadding(postpadding);
#endif
    } else {
      deconv_layer->setPadding(prepadding);
    }
    const auto strides = nvinfer1::DimsHW(std::stoi(str_strides[0]), std::stoi(str_strides[1]));
    deconv_layer->setStride(strides);
    deconv_layer->setNbGroups(groups);
    nvinfer1::ITensor* output = deconv_layer->getOutput(0);
    // Output padding.
    if (str_output_padding.size()) {
      GetPadding(str_output_padding, &use_asymmetric_padding, &prepadding, &postpadding);
      if (prepadding.h() != 0 || prepadding.w() != 0 || postpadding.h() != 0 ||
          postpadding.w() != 0) {
        // Output padding for Conv2D transpose is always asymmetric and applied to post only.
        prepadding = nvinfer1::DimsHW(0, 0);
        auto pad_layer = params->network->addPadding(*output, prepadding, postpadding);
        output = pad_layer->getOutput(0);
      }
    }
    params->outputs.push_back(output);
  }
};

#if TRT_VERSION_GE(6, 0, 1)
class Conv3DTransposeOpConverter : public TensorRTOpConverter {
 public:
  Conv3DTransposeOpConverter() : TensorRTOpConverter({kTensor, kWeight}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input_tensor = params->inputs.at(0).tensor;
    auto weight_shape = params->inputs.at(1).weight_shape;
    ICHECK_EQ(params->node.GetAttr<std::vector<std::string>>("data_layout")[0], "NCDHW");
    ICHECK(params->node.GetAttr<std::vector<std::string>>("out_layout")[0] == "" ||
           params->node.GetAttr<std::vector<std::string>>("out_layout")[0] == "NCDHW");
    ICHECK_EQ(params->node.GetAttr<std::vector<std::string>>("kernel_layout")[0], "OIDHW");
    auto str_dilation = params->node.GetAttr<std::vector<std::string>>("dilation");
    ICHECK_EQ(str_dilation.size(), 3);
    ICHECK(std::stoi(str_dilation[0]) == 1 && std::stoi(str_dilation[1]) == 1 &&
           std::stoi(str_dilation[2]) == 1);
    auto str_strides = params->node.GetAttr<std::vector<std::string>>("strides");
    auto str_padding = params->node.GetAttr<std::vector<std::string>>("padding");
    auto str_output_padding = params->node.GetAttr<std::vector<std::string>>("output_padding");
    int groups = std::stoi(params->node.GetAttr<std::vector<std::string>>("groups")[0]);
    nvinfer1::Dims prepadding, postpadding;
    bool use_asymmetric_padding;
    GetPadding3D(str_padding, &use_asymmetric_padding, &prepadding, &postpadding);

    // Could use attrs->channels.as<IntImmNode>()->value
    const int num_outputs = weight_shape[1];
    const auto kernel_size = nvinfer1::Dims3(weight_shape[2], weight_shape[3], weight_shape[4]);
    nvinfer1::Weights bias{nvinfer1::DataType::kFLOAT, nullptr, 0};
    auto deconv_layer = params->network->addDeconvolutionNd(*input_tensor, num_outputs, kernel_size,
                                                            params->inputs.at(1).weight, bias);
    ICHECK(deconv_layer != nullptr);
    if (use_asymmetric_padding) {
      deconv_layer->setPrePadding(prepadding);
      deconv_layer->setPostPadding(postpadding);
    } else {
      deconv_layer->setPaddingNd(prepadding);
    }
    ICHECK_EQ(str_strides.size(), 3);
    const auto strides = nvinfer1::Dims3(std::stoi(str_strides[0]), std::stoi(str_strides[1]),
                                         std::stoi(str_strides[2]));
    deconv_layer->setStrideNd(strides);
    deconv_layer->setNbGroups(groups);
    nvinfer1::ITensor* output = deconv_layer->getOutput(0);
    // Output padding.
    if (str_output_padding.size()) {
      GetPadding3D(str_output_padding, &use_asymmetric_padding, &prepadding, &postpadding);
      // Are any post-padding values non-zero?
      ICHECK(!std::any_of(postpadding.d, postpadding.d + postpadding.nbDims, [](int x) {
        return x != 0;
      })) << "TRT does not support padding on 3 dimensions.";
    }
    params->outputs.push_back(output);
  }
};
#endif  // TRT_VERSION_GE(6, 0, 1)

class TransposeOpConverter : public TensorRTOpConverter {
 public:
  TransposeOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input = params->inputs.at(0).tensor;
    auto str_axes = params->node.GetAttr<std::vector<std::string>>("axes");
    std::vector<int> order;
    for (size_t i = 0; i < str_axes.size(); ++i) {
      order.push_back(std::stoi(str_axes[i]));
    }
    params->outputs.push_back(Transpose(params, input, order));
  }
};

class LayoutTransformOpConverter : public TensorRTOpConverter {
 public:
  LayoutTransformOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input = params->inputs.at(0).tensor;
    auto src = params->node.GetAttr<std::vector<std::string>>("src_layout")[0];
    auto dst = params->node.GetAttr<std::vector<std::string>>("dst_layout")[0];
    std::vector<int> order;
    if (src == "NCHW" && dst == "NHWC") {
      order = {0, 2, 3, 1};
    } else if (src == "NHWC" && dst == "NCHW") {
      order = {0, 3, 1, 2};
    } else if (src == "NDHWC" && dst == "NCDHW") {
      order = {0, 4, 1, 2, 3};
    } else if (src == "NCDHW" && dst == "NDHWC") {
      order = {0, 2, 3, 4, 1};
    }
    params->outputs.push_back(Transpose(params, input, order));
  }
};

class ReshapeOpConverter : public TensorRTOpConverter {
 public:
  ReshapeOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input = params->inputs.at(0).tensor;
    ICHECK_EQ(std::stoi(params->node.GetAttr<std::vector<std::string>>("reverse")[0]), false);
    auto str_newshape = params->node.GetAttr<std::vector<std::string>>("newshape");
    std::vector<int> new_shape;
    const int start_index = TRT_HAS_IMPLICIT_BATCH(params) ? 1 : 0;
    for (size_t i = start_index; i < str_newshape.size(); ++i) {
      const int value = std::stoi(str_newshape[i]);
      ICHECK_GE(value, -1);
      new_shape.push_back(value);
    }
    params->outputs.push_back(Reshape(params, input, new_shape));
  }
};

class PadOpConverter : public TensorRTOpConverter {
 public:
  PadOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input = params->inputs.at(0).tensor;
    auto str_paddding = params->node.GetAttr<std::vector<std::string>>("padding");
    nvinfer1::DimsHW prepadding =
        nvinfer1::DimsHW(std::stoi(str_paddding[0]), std::stoi(str_paddding[1]));
    nvinfer1::DimsHW postpadding =
        nvinfer1::DimsHW(std::stoi(str_paddding[2]), std::stoi(str_paddding[3]));
    auto pad_layer = params->network->addPadding(*input, prepadding, postpadding);
    params->outputs.push_back(pad_layer->getOutput(0));
  }
};

class ReduceOpConverter : public TensorRTOpConverter {
 public:
  ReduceOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    static const std::unordered_map<std::string, nvinfer1::ReduceOperation> op_map = {
        {"sum", nvinfer1::ReduceOperation::kSUM},
        {"prod", nvinfer1::ReduceOperation::kPROD},
        {"max", nvinfer1::ReduceOperation::kMAX},
        {"min", nvinfer1::ReduceOperation::kMIN},
        {"mean", nvinfer1::ReduceOperation::kAVG}};
    auto it = op_map.find(params->op_name);
    ICHECK(it != op_map.end()) << "Unsupported reduce type " << params->op_name;

    auto input = params->inputs.at(0).tensor;
    ICHECK_EQ(std::stoi(params->node.GetAttr<std::vector<std::string>>("exclude")[0]), false);
    bool keepdims = std::stoi(params->node.GetAttr<std::vector<std::string>>("keepdims")[0]);
    auto str_axis = params->node.GetAttr<std::vector<std::string>>("axis");
    // TODO(trevmorr): Support reduce to scalar.
    ICHECK_GT(str_axis.size(), 0);
    uint32_t reduce_axes = 0;
    for (size_t i = 0; i < str_axis.size(); ++i) {
      const int axis = ConvertAxis(params, std::stoi(str_axis[i]), input->getDimensions().nbDims);
      reduce_axes |= 1 << axis;
    }
    auto reduce_layer = params->network->addReduce(*input, it->second, reduce_axes, keepdims);
    params->outputs.push_back(reduce_layer->getOutput(0));
  }
};

#if TRT_VERSION_GE(5, 1, 5)
class StridedSliceOpConverter : public TensorRTOpConverter {
 public:
  StridedSliceOpConverter() : TensorRTOpConverter({kTensor, kWeight, kWeight, kWeight}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input = params->inputs.at(0).tensor;
    auto input_dims = TrtDimsToVector(input->getDimensions());
    auto str_start = params->node.GetAttr<std::vector<std::string>>("start");
    auto str_size = params->node.GetAttr<std::vector<std::string>>("size");
    auto str_strides = params->node.GetAttr<std::vector<std::string>>("strides");
    std::vector<int> start, size, strides;
    std::transform(str_start.begin(), str_start.end(), std::back_inserter(start),
                   [](const std::string& s) { return std::stoi(s); });
    std::transform(str_size.begin(), str_size.end(), std::back_inserter(size),
                   [](const std::string& s) { return std::stoi(s); });
    std::transform(str_strides.begin(), str_strides.end(), std::back_inserter(strides),
                   [](const std::string& s) { return std::stoi(s); });
    if (TRT_HAS_IMPLICIT_BATCH(params)) {
      start.erase(start.begin());
      size.erase(size.begin());
      strides.erase(strides.begin());
    }
    auto slice_layer = params->network->addSlice(*input, VectorToTrtDims(start),
                                                 VectorToTrtDims(size), VectorToTrtDims(strides));
    params->outputs.push_back(slice_layer->getOutput(0));
  }
};
#endif

class AdaptivePoolingOpConverter : public TensorRTOpConverter {
 public:
  AdaptivePoolingOpConverter() : TensorRTOpConverter({kTensor}) {}

  void Convert(TensorRTOpConverterParams* params) const {
    auto input_tensor = params->inputs.at(0).tensor;
    auto input_dims = TrtDimsToVector(input_tensor->getDimensions());
    static const std::unordered_map<std::string, nvinfer1::PoolingType> op_map = {
        {"nn.adaptive_max_pool2d", nvinfer1::PoolingType::kMAX},
        {"nn.adaptive_avg_pool2d", nvinfer1::PoolingType::kAVERAGE}};
    auto it = op_map.find(params->op_name);
    ICHECK(it != op_map.end()) << "Unsupported pooling type " << params->op_name << " in TensorRT";
    ICHECK_EQ(params->node.GetAttr<std::vector<std::string>>("layout")[0], "NCHW");

    // This is an approximation of adaptive pooling. Results will not be
    // mathematically exact except when output_size is (1, 1).
    // Annotation rules will only allow output size of (1, 1).
    auto output_size = nvinfer1::DimsHW(1, 1);
    const int h = TRT_HAS_IMPLICIT_BATCH(params) ? input_dims[1] : input_dims[2];
    const int w = TRT_HAS_IMPLICIT_BATCH(params) ? input_dims[2] : input_dims[3];
    const auto stride = nvinfer1::DimsHW(h / output_size.h(), w / output_size.w());
    const auto window_size = nvinfer1::DimsHW(h - (output_size.h() - 1) * stride.h(),
                                              w - (output_size.w() - 1) * stride.w());
    auto pool_layer = params->network->addPooling(*input_tensor, it->second, window_size);
    ICHECK(pool_layer != nullptr);
    pool_layer->setStride(stride);
    params->outputs.push_back(pool_layer->getOutput(0));
  }
};

const std::shared_ptr<std::unordered_map<std::string, std::shared_ptr<TensorRTOpConverter>>>
GetOpConverters() {
  static auto map =
      std::make_shared<std::unordered_map<std::string, std::shared_ptr<TensorRTOpConverter>>>();
  if (!map->empty()) return map;
  map->emplace("nn.relu", std::make_shared<ActivationOpConverter>());
  map->emplace("sigmoid", std::make_shared<ActivationOpConverter>());
  map->emplace("tanh", std::make_shared<ActivationOpConverter>());
  map->emplace("nn.batch_norm", std::make_shared<BatchNormOpConverter>());
  map->emplace("nn.softmax", std::make_shared<SoftmaxOpConverter>());
  map->emplace("nn.conv2d", std::make_shared<Conv2DOpConverter>());
  map->emplace("nn.dense", std::make_shared<DenseOpConverter>());
  map->emplace("nn.bias_add", std::make_shared<BiasAddOpConverter>());
  map->emplace("add", std::make_shared<ElementWiseBinaryOpConverter>());
  map->emplace("subtract", std::make_shared<ElementWiseBinaryOpConverter>());
  map->emplace("multiply", std::make_shared<ElementWiseBinaryOpConverter>());
  map->emplace("divide", std::make_shared<ElementWiseBinaryOpConverter>());
  map->emplace("power", std::make_shared<ElementWiseBinaryOpConverter>());
  map->emplace("maximum", std::make_shared<ElementWiseBinaryOpConverter>());
  map->emplace("minimum", std::make_shared<ElementWiseBinaryOpConverter>());
  map->emplace("nn.max_pool2d", std::make_shared<PoolingOpConverter>());
  map->emplace("nn.avg_pool2d", std::make_shared<PoolingOpConverter>());
  map->emplace("nn.global_max_pool2d", std::make_shared<GlobalPoolingOpConverter>());
  map->emplace("nn.global_avg_pool2d", std::make_shared<GlobalPoolingOpConverter>());
  map->emplace("exp", std::make_shared<UnaryOpConverter>());
  map->emplace("log", std::make_shared<UnaryOpConverter>());
  map->emplace("sqrt", std::make_shared<UnaryOpConverter>());
  map->emplace("abs", std::make_shared<UnaryOpConverter>());
  map->emplace("negative", std::make_shared<UnaryOpConverter>());
  map->emplace("nn.batch_flatten", std::make_shared<BatchFlattenOpConverter>());
  map->emplace("expand_dims", std::make_shared<ExpandDimsOpConverter>());
  map->emplace("squeeze", std::make_shared<SqueezeOpConverter>());
  map->emplace("concatenate", std::make_shared<ConcatOpConverter>());
  map->emplace("nn.conv2d_transpose", std::make_shared<Conv2DTransposeOpConverter>());
  map->emplace("transpose", std::make_shared<TransposeOpConverter>());
  map->emplace("layout_transform", std::make_shared<LayoutTransformOpConverter>());
  map->emplace("reshape", std::make_shared<ReshapeOpConverter>());
  map->emplace("nn.pad", std::make_shared<PadOpConverter>());
  map->emplace("sum", std::make_shared<ReduceOpConverter>());
  map->emplace("prod", std::make_shared<ReduceOpConverter>());
  map->emplace("max", std::make_shared<ReduceOpConverter>());
  map->emplace("min", std::make_shared<ReduceOpConverter>());
  map->emplace("mean", std::make_shared<ReduceOpConverter>());
  map->emplace("nn.adaptive_max_pool2d", std::make_shared<AdaptivePoolingOpConverter>());
  map->emplace("nn.adaptive_avg_pool2d", std::make_shared<AdaptivePoolingOpConverter>());
#if TRT_VERSION_GE(5, 1, 5)
  map->emplace("clip", std::make_shared<ActivationOpConverter>());
  map->emplace("nn.leaky_relu", std::make_shared<ActivationOpConverter>());
  map->emplace("sin", std::make_shared<UnaryOpConverter>());
  map->emplace("cos", std::make_shared<UnaryOpConverter>());
  map->emplace("atan", std::make_shared<UnaryOpConverter>());
  map->emplace("ceil", std::make_shared<UnaryOpConverter>());
  map->emplace("floor", std::make_shared<UnaryOpConverter>());
  map->emplace("strided_slice", std::make_shared<StridedSliceOpConverter>());
#endif  // TRT_VERSION_GE(5, 1, 5)
#if TRT_VERSION_GE(6, 0, 1)
  map->emplace("nn.conv3d", std::make_shared<Conv3DOpConverter>());
  map->emplace("nn.max_pool3d", std::make_shared<Pooling3DOpConverter>());
  map->emplace("nn.avg_pool3d", std::make_shared<Pooling3DOpConverter>());
  map->emplace("nn.conv3d_transpose", std::make_shared<Conv3DTransposeOpConverter>());
#endif  // TRT_VERSION_GE(6, 0, 1)
  return map;
}

}  // namespace contrib
}  // namespace runtime
}  // namespace tvm
