convert.h 13.0 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// 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.
#pragma once

16
#include <glog/logging.h>
W
Wilber 已提交
17
#include <llvm/Support/ErrorHandling.h>
18
#include <mlir/IR/Attributes.h>
19
#include <mlir/IR/Builders.h>
W
Wilber 已提交
20 21
#include <mlir/IR/BuiltinAttributes.h>
#include <mlir/IR/PatternMatch.h>
22
#include <mlir/Transforms/DialectConversion.h>
23
#include "paddle/infrt/dialect/infrt/common/types.h"
24
#include "paddle/infrt/dialect/infrt/ir/infrt_dialect.h"
25 26
#include "paddle/infrt/dialect/pd/ir/pd_ops.h"
#include "paddle/infrt/dialect/phi/ir/infrt_phi_tensor.h"
27
#include "paddle/infrt/dialect/tensorrt/trt_ops.h"
W
Wilber 已提交
28
#include "paddle/infrt/kernel/tensorrt/trt_helper.h"
29 30 31

namespace infrt {
namespace trt {
W
Wilber 已提交
32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59

#ifdef INFRT_WITH_TRT

#define STRING_TO_ENUM_TYPE(enum_type) enum_type
#define STRING_TO_ENUM_VALUE(enum_value) enum_value
#include <NvInfer.h>

#else  // INFRT_WITH_TRT

#define STRING_TO_ENUM_TYPE(enum_type) std::string
#define STRING_TO_ENUM_VALUE(enum_value) #enum_value

#endif  // INFRT_WITH_TRT

template <typename T>
::mlir::IntegerAttr createNvinferEnumAttr(
    ::mlir::PatternRewriter &rewriter,  // NOLINT
    T enum_value) {
  return rewriter.getSI32IntegerAttr((int32_t)enum_value);
}

template <>
::mlir::IntegerAttr createNvinferEnumAttr<std::string>(
    ::mlir::PatternRewriter &rewriter, std::string enum_value) {  // NOLINT
  (void)enum_value;
  return rewriter.getSI32IntegerAttr(-1);
}

60
static mlir::Value createTRTConv2dOp(mlir::PatternRewriter &rewriter,  // NOLINT
61 62 63
                                     mlir::Operation *op,
                                     mlir::Value input,
                                     mlir::Value filter) {
64
  auto conv_op = ::llvm::dyn_cast<infrt::pd::Conv2dOp>(op);
65
  ::mlir::SmallVector<::mlir::Value, 4> operands;
66 67
  operands.push_back(input);
  operands.push_back(filter);
68 69

  ::mlir::SmallVector<::mlir::Type, 4> resultTypes;
70
  for (auto v : conv_op.getODSResults(0)) {
71 72
    resultTypes.push_back(v.getType());
  }
73

74
  ::mlir::SmallVector<::mlir::NamedAttribute, 8> attributes;
75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93

  auto *filter_producer = filter.getDefiningOp();
  auto create_inited_tensor_op =
      llvm::dyn_cast<::infrt::phi::CreateHostInitedDenseTensorOp>(
          filter_producer);

  CHECK_NOTNULL(create_inited_tensor_op);
  mlir::ArrayAttr dims = create_inited_tensor_op.dims();
  CHECK_EQ(dims.size(), 4U);
  CHECK(dims[0].getType().isIntOrIndex());

  const int32_t n_output = dims[0].cast<mlir::IntegerAttr>().getInt();
  const int32_t filter_h = dims[2].cast<mlir::IntegerAttr>().getInt();
  const int32_t filter_w = dims[3].cast<mlir::IntegerAttr>().getInt();

  auto padding_attr = conv_op->getAttrOfType<::mlir::ArrayAttr>("paddings");
  llvm::SmallVector<int32_t, 4> paddings(padding_attr.size());
  for (size_t i = 0; i < padding_attr.size(); i++) {
    paddings[i] = padding_attr[i].cast<mlir::IntegerAttr>().getInt();
94
  }
95 96 97 98 99

  auto dilations_attr = conv_op->getAttrOfType<::mlir::ArrayAttr>("dilations");
  llvm::SmallVector<int32_t> dilations(dilations_attr.size());
  for (size_t i = 0; i < dilations_attr.size(); i++) {
    dilations[i] = dilations_attr[i].cast<mlir::IntegerAttr>().getInt();
100
  }
101 102 103 104 105 106 107 108 109 110 111 112

  llvm::SmallVector<int32_t, 2> nv_paddings(2);
  llvm::SmallVector<int32_t, 4> nv_pre_paddings(2);
  llvm::SmallVector<int32_t, 4> nv_post_paddings(2);
  llvm::SmallVector<int32_t, 2> nv_dilations({dilations[0], dilations[1]});
  int32_t nv_padding_mode = 0;  // nvinfer1::PaddingMode::kEXPLICIT_ROUND_DOWN
  auto padding_algorithm_attr =
      conv_op->getAttrOfType<::mlir::StringAttr>("padding_algorithm");
  if (padding_algorithm_attr.strref() == "VALID") {
    for (size_t i = 0; i < paddings.size(); i++) {
      paddings[i] = 0;
    }
113
  }
114 115 116 117
  if (padding_algorithm_attr.strref() == "SAME") {
    nv_padding_mode = 2;  // nvinfer1::PaddingMode::kSAME_UPPER
    nv_dilations[0] = 1;
    nv_dilations[1] = 1;
118
  }
119 120 121 122 123 124 125 126 127 128

  if (paddings.size() == 2) {
    nv_paddings[0] = paddings[0];
    nv_paddings[1] = paddings[1];
  } else {
    CHECK_EQ(paddings.size(), 4U);
    nv_pre_paddings[0] = paddings[0];
    nv_pre_paddings[1] = paddings[2];
    nv_post_paddings[0] = paddings[1];
    nv_post_paddings[1] = paddings[3];
129
  }
130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154

  attributes.emplace_back(rewriter.getStringAttr("out_channel_num"),
                          rewriter.getSI32IntegerAttr(n_output));

  attributes.emplace_back(rewriter.getStringAttr("kernel_size"),
                          rewriter.getI32ArrayAttr({filter_h, filter_w}));

  attributes.emplace_back(
      rewriter.getStringAttr("dilations"),
      rewriter.getI32ArrayAttr({nv_dilations[0], nv_dilations[1]}));

  attributes.emplace_back(rewriter.getStringAttr("padding_mode"),
                          rewriter.getSI32IntegerAttr(nv_padding_mode));

  attributes.emplace_back(rewriter.getStringAttr("paddings"),
                          rewriter.getI32ArrayAttr({paddings[0], paddings[1]}));

  attributes.emplace_back(
      rewriter.getStringAttr("pre_paddings"),
      rewriter.getI32ArrayAttr({nv_pre_paddings[0], nv_pre_paddings[1]}));

  attributes.emplace_back(
      rewriter.getStringAttr("post_paddings"),
      rewriter.getI32ArrayAttr({nv_post_paddings[0], nv_post_paddings[1]}));

155
  {
156
    auto tblgen_attr = conv_op->getAttrOfType<::mlir::IntegerAttr>("groups");
157 158 159
    attributes.emplace_back(rewriter.getStringAttr("groups"), tblgen_attr);
  }
  {
160 161
    auto tblgen_attr = conv_op->getAttrOfType<::mlir::ArrayAttr>("strides");
    attributes.emplace_back(rewriter.getStringAttr("strides"), tblgen_attr);
162 163
  }
  return rewriter.create<trt::ConvolutionOp>(
164
      conv_op->getLoc(), resultTypes, operands, attributes);
165
}
166

167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247
static inline mlir::ArrayAttr TransposeWeight(
    mlir::PatternRewriter &builder,  // NOLINT
    const mlir::ArrayAttr &weight,
    const mlir::ArrayAttr &dims) {
  CHECK_EQ(dims.size(), 2U);
  CHECK(!dims.empty());
  CHECK(dims[0].getType().isInteger(64));
  CHECK(!weight.empty());
  CHECK(weight[0].getType().isF32());

  int row = dims[0].cast<mlir::IntegerAttr>().getInt();
  int col = dims[1].cast<mlir::IntegerAttr>().getInt();
  std::vector<float> trans_weight(weight.size());
  for (int i = 0; i < row; ++i) {
    for (int j = 0; j < col; ++j) {
      trans_weight[j * row + i] =
          weight[i * col + j].cast<mlir::FloatAttr>().getValueAsDouble();
    }
  }
  return builder.getF32ArrayAttr(trans_weight);
}

// matmul_y and elt_y is weights.
inline ::llvm::SmallVector<::mlir::Value, 4> createTrtFcOp(
    mlir::PatternRewriter &builder,  // NOLINT
    mlir::Value matmul_x,
    mlir::Value matmul_y,
    mlir::Value elt_y,
    mlir::Value elt_out) {
  ::llvm::SmallVector<::mlir::Operation *, 4> tblgen_ops;

  auto *y_producer = matmul_y.getDefiningOp();
  auto create_inited_tensor_op =
      llvm::dyn_cast<::infrt::phi::CreateHostInitedDenseTensorOp>(y_producer);
  CHECK_NOTNULL(create_inited_tensor_op);

  mlir::ArrayAttr dims = create_inited_tensor_op.dims();
  CHECK_EQ(dims.size(), 2U);

  std::vector<int64_t> new_dims(dims.size());
  CHECK(!dims.empty());
  CHECK(dims[0].getType().isIntOrIndex());
  for (size_t i = 0; i < new_dims.size(); ++i) {
    new_dims[i] = dims[dims.size() - 1 - i].cast<mlir::IntegerAttr>().getInt();
  }
  auto insert_point = builder.saveInsertionPoint();
  builder.setInsertionPoint(create_inited_tensor_op);
  auto new_inited_op =
      builder.create<::infrt::phi::CreateHostInitedDenseTensorOp>(
          create_inited_tensor_op->getLoc(),
          create_inited_tensor_op.output().getType(),
          create_inited_tensor_op.context(),
          builder.getI64ArrayAttr(new_dims),
          ::infrt::LayoutAttr::get(builder.getContext(),
                                   ::infrt::LayoutType::NCHW),
          create_inited_tensor_op.lod(),
          TransposeWeight(builder, create_inited_tensor_op.values(), dims));
  builder.replaceOp(create_inited_tensor_op, new_inited_op->getResults());
  builder.restoreInsertionPoint(insert_point);

  auto ods_loc = builder.getFusedLoc({y_producer->getLoc()});
  ::infrt::trt::FullyConnectedOp fc_op;
  {
    ::mlir::SmallVector<::mlir::Type, 4> tblgen_types;

    fc_op = builder.create<::infrt::trt::FullyConnectedOp>(
        ods_loc,
        elt_out.getType(),
        matmul_x,
        new_inited_op.output(),
        elt_y,
        builder.getSI32IntegerAttr(new_dims[0]));
  }

  ::llvm::SmallVector<::mlir::Value, 4> tblgen_repl_values;
  for (auto v : ::llvm::SmallVector<::mlir::Value, 4>{fc_op.getODSResults(0)}) {
    tblgen_repl_values.push_back(v);
  }
  return tblgen_repl_values;
}

W
Wilber 已提交
248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322
inline mlir::IntegerAttr CreatePoolingType(
    mlir::PatternRewriter &builder,  // NOLINT
    mlir::StringAttr pool_type) {
  // pool_type.
  auto ptype = pool_type.str();
  if (ptype == "max") {
    return createNvinferEnumAttr(builder, nvinfer1::PoolingType::kMAX);
  } else if (ptype == "avg") {
    return createNvinferEnumAttr(builder, nvinfer1::PoolingType::kAVERAGE);
  } else {
    llvm_unreachable("unknown pool_type.");
    return {};
  }
}

inline mlir::IntegerAttr CreatePaddingMode(
    mlir::PatternRewriter &builder,  // NOLINT
    mlir::StringAttr padding_algorithm,
    mlir::BoolAttr ceil_mode) {
  // TODO(Inference): Phi pool kernel seems not process ceil_mode.
  auto padding_algo = padding_algorithm.str();
  if (padding_algo == "SAME") {
    return createNvinferEnumAttr(builder, nvinfer1::PaddingMode::kSAME_UPPER);
  }
  if (ceil_mode.getValue() && padding_algo != "SAME") {
    return createNvinferEnumAttr(builder,
                                 nvinfer1::PaddingMode::kEXPLICIT_ROUND_UP);
  } else {
    return createNvinferEnumAttr(builder,
                                 nvinfer1::PaddingMode::kEXPLICIT_ROUND_DOWN);
  }
}

inline ::llvm::SmallVector<::mlir::Value, 4> CreatePaddleTrtPoolingOp(
    mlir::PatternRewriter &builder,  // NOLINT
    mlir::Value input,
    mlir::StringAttr pool_type,
    mlir::ArrayAttr ksize,
    mlir::BoolAttr global_pooling,
    mlir::ArrayAttr strides,
    mlir::ArrayAttr paddings,
    mlir::BoolAttr exclusive,
    mlir::BoolAttr adaptive,
    mlir::BoolAttr ceil_mode,
    mlir::StringAttr data_format,
    mlir::StringAttr padding_algorithm) {
  ::llvm::SmallVector<::mlir::Value, 4> tblgen_repl_values;

  // TODO(inference): Support NHWC.
  if (data_format.str() != "NCHW") {
    CHECK(false) << "The pool2d converter now only support NCHW.";
  }

  // TODO(Wilber): How to support dynamic shape?

  auto *input_producer = input.getDefiningOp();

  // Process pool_type.
  auto pool_type_attr = CreatePoolingType(builder, pool_type);

  // Update padding.
  auto padding_algorithm_str = padding_algorithm.str();
  auto paddings_attr = paddings;
  if (padding_algorithm_str == "EXPLICIT") {
    // Do nothing on paddings.
  } else if (padding_algorithm_str == "SAME") {
    // We should process this case in trt network build phase.
  } else if (padding_algorithm_str == "VALID") {
    // Set padding to zero.
    paddings_attr = builder.getI32ArrayAttr({0, 0});
  } else {
    CHECK(false) << "Unknown padding_algotithm.";
  }

  // if global_pooling == true or adaptive == true, padding will be ignored
W
Wilber 已提交
323 324 325
  // if (global_pooling.getValue() || adaptive.getValue()) {
  //   paddings_attr = builder.getI32ArrayAttr({0, 0});
  // }
W
Wilber 已提交
326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348

  // if global_pooling == true, then we should update kernel size to input dims.
  if (global_pooling.getValue() == true) {
    // Update ksize to input dims.
  }

  // The adaptive logic should be processed when we get the context of
  // INetworkDefinition, so we place the logic in infrt runtime(trt compile
  // time).

  // The `exclusive` may be a naive attr, which can be forward to trt.

  auto padding_mode_attr =
      CreatePaddingMode(builder, padding_algorithm, ceil_mode);

  if (global_pooling.getValue() == true) {
    CHECK(false) << "Temporarily not support global_pool";
    return tblgen_repl_values;
  }

  PoolingOp pool_op;
  {
    auto ods_loc = builder.getFusedLoc({input_producer->getLoc()});
349 350 351 352 353 354 355 356 357 358 359
    pool_op = builder.create<PoolingOp>(ods_loc,
                                        input.getType(),
                                        input,
                                        pool_type_attr,
                                        ksize,
                                        strides,
                                        paddings_attr,
                                        padding_mode_attr,
                                        exclusive,
                                        adaptive,
                                        padding_algorithm);
W
Wilber 已提交
360 361 362 363 364 365 366 367 368
  }

  for (auto v :
       ::llvm::SmallVector<::mlir::Value, 4>{pool_op.getODSResults(0)}) {
    tblgen_repl_values.push_back(v);
  }
  return tblgen_repl_values;
}

369 370
}  // namespace trt
}  // namespace infrt