trt_op_teller_pass.cc 2.3 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
// Copyright (c) 2021 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.

#include "paddle/infrt/dialect/tensorrt/trt_op_teller_pass.h"

W
Wilber 已提交
17
#include <llvm/Support/Casting.h>
18
#include <mlir/IR/Builders.h>
W
Wilber 已提交
19
#include "paddle/infrt/dialect/dense_tensor.h"
20 21
#include "paddle/infrt/dialect/infrt/ir/basic_kernels.h"
#include "paddle/infrt/dialect/infrt/ir/infrt_dialect.h"
22
#include "paddle/infrt/dialect/pd/ir/pd_ops.h"
23
#include "paddle/infrt/dialect/phi/ir/infrt_phi_tensor.h"
24 25 26 27

namespace infrt {
namespace trt {
// Implementation of the trtOpTellerPass。
28
void TRTOpTellerPass::runOnFunction() {
29 30
  mlir::Block &body = getFunction().front();
  std::vector<mlir::Operation *> worklist;
31 32 33 34 35
  worklist.reserve(body.getOperations().size());
  for (auto &op : body) {
    worklist.push_back(&op);
  }
  // Build GraphOp.
36
  mlir::OpBuilder builder(&body, body.begin());
37 38 39 40
  while (!worklist.empty()) {
    auto *op = worklist.back();
    worklist.pop_back();
    if (op == nullptr) continue;
W
Wilber 已提交
41
    if (op->getName().getStringRef().substr(0, 3) != "pd.") continue;
42 43
    builder.setInsertionPoint(op);
    auto loc = getFunction().getLoc();
44
    auto graph_op = builder.create<::infrt::GraphOp>(
S
Shang Zhizhou 已提交
45
        loc, op->getResultTypes(), op->getOperands());
46

47
    ::llvm::SmallVector<mlir::Value, 4> tblgen_repl_values;
48
    for (auto v :
49
         ::llvm::SmallVector<mlir::Value, 4>{graph_op.getODSResults(0)}) {
50 51 52 53
      tblgen_repl_values.push_back(v);
    }
    op->replaceAllUsesWith(tblgen_repl_values);
    // Build graph op.
54
    mlir::Block *block = new mlir::Block;
55 56 57
    graph_op.body().push_back(block);
    op->moveBefore(block, block->begin());
    builder.setInsertionPointToEnd(block);
S
Shang Zhizhou 已提交
58
    builder.create<::infrt::ReturnOp>(loc, op->getResults());
59 60 61 62
  }
}
}  // namespace trt
}  // namespace infrt