phi_op_cvt_pass.cc

// 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/phi/pass/phi_op_cvt_pass.h"

#include <glog/logging.h>
#include <llvm/ADT/SetVector.h>
#include <mlir/Analysis/SliceAnalysis.h>
#include <mlir/IR/Builders.h>
#include <mlir/IR/Operation.h>
#include <mlir/IR/OperationSupport.h>
#include <list>
#include <unordered_set>
#include <vector>

#include "paddle/infrt/dialect/infrt/ir/infrt_dialect.h"
#include "paddle/infrt/dialect/phi/ir/infrt_phi_tensor.h"
#include "paddle/infrt/dialect/phi/pass/kernel_op_desc.h"
#include "paddle/infrt/dialect/phi/pass/proto_arg_map_context.h"
#include "paddle/phi/core/compat/op_utils.h"
#include "paddle/phi/ops/compat/signatures.h"

namespace {
class phiOpCvtPass
    : public mlir::PassWrapper<phiOpCvtPass, mlir::FunctionPass> {
 public:
  ::llvm::StringRef getName() const override { return "phiOpCvtPass"; }
  void runOnFunction() override;
  explicit phiOpCvtPass(
      std::vector<infrt::Place> valid_places = std::vector<infrt::Place>())
      : valid_places_(valid_places) {}

 private:
  void convertStage();
  void diapatchStage();
  std::vector<infrt::Place> valid_places_;
};

// Implementation of the phiOpCvtPass.
void phiOpCvtPass::runOnFunction() {
  convertStage();
  diapatchStage();
}
void phiOpCvtPass::convertStage() {
  mlir::Block &body = getFunction().front();
  std::vector<mlir::Operation *> worklist;
  for (auto &op : body.without_terminator()) {
    worklist.push_back(&op);
  }
  mlir::OpBuilder builder(&body, body.begin());
  while (!worklist.empty()) {
    auto *op = worklist.back();
    worklist.pop_back();
    if (op == nullptr) continue;

    std::string op_name = op->getName().getIdentifier().str();

    // only convert op in pd dialect.
    if (op_name.substr(0, 3) != "pd.") continue;
    op_name = op_name.substr(3);
    if (pd_dialect_inputs_info_map_.find(op_name) ==
            pd_dialect_inputs_info_map_.end() ||
        pd_dialect_outputs_info_map_.find(op_name) ==
            pd_dialect_outputs_info_map_.end()) {
      // Todo: print log
      continue;
    }

    ::phi::KernelSignature kernel_sign =
        ::phi::OpUtilsMap::Instance().GetArgumentMappingFn(op_name)(
            infrt::ProtoArgumentMappingContext(op));
    // resort input&output according to kernel_sign
    ::llvm::SmallVector<mlir::Value, 4> inputs, ori_output;
    ::llvm::SmallVector<mlir::Type, 4> output_types;
    for (const std::string &str : std::get<0>(kernel_sign.args)) {
      if (pd_dialect_inputs_info_map_.at(op_name).count(str) == 0) {
        // Todo: print error log
        return;
      }
      uint8_t index = pd_dialect_inputs_info_map_.at(op_name).at(str);
      inputs.push_back(op->getOperands()[index]);
    }

    for (const std::string &str : std::get<2>(kernel_sign.args)) {
      if (pd_dialect_outputs_info_map_.at(op_name).count(str) == 0) {
        // Todo: print error log
        return;
      }
      uint8_t index = pd_dialect_outputs_info_map_.at(op_name).at(str);
      output_types.push_back(op->getResultTypes()[index]);
      ori_output.push_back(op->getResult(index));
    }

    auto loc = getFunction().getLoc();
    builder.setInsertionPoint(op);
    auto kernel_op = builder.create<infrt::KernelOp>(
        loc, output_types, inputs, kernel_sign.name, op->getAttrDictionary());
    for (size_t index = 0; index < ori_output.size(); ++index) {
      ori_output[index].replaceAllUsesWith(kernel_op.getResult(index));
    }
    if (!op->use_empty()) {
      // Todo: print error log
      return;
    }
    op->erase();
  }
}
void phiOpCvtPass::diapatchStage() {
  std::vector<infrt::KernelOp> worklist;
  mlir::Block &block = getFunction().front();
  for (auto &op : block) {
    infrt::KernelOp kernel_op = ::llvm::dyn_cast_or_null<infrt::KernelOp>(&op);
    if (nullptr != kernel_op) worklist.push_back(kernel_op);
  }

  mlir::OpBuilder builder(&block, block.begin());
  std::map<infrt::TargetType, mlir::Value> phi_context;
  for (infrt::KernelOp kernel_op : worklist) {
    std::string kernel_name = kernel_op.name().str();
    std::vector<infrt::PhiKernelDesc> candidates =
        getCandidateKernels(kernel_name, valid_places_);
    if (candidates.empty()) {
      LOG(FATAL) << "No candidate kernels for op:" << kernel_name;
      continue;
    }
    builder.setInsertionPoint(kernel_op);

    // Todo: Implimentation the concrete pass pick strategy
    const infrt::PhiKernelDesc &phi_kernel_desc = candidates.front();

    kernel_name =
        infrt::getPhiTargetPrefix(phi_kernel_desc.kernelType.target) +
        kernel_name +
        infrt::getPhiPrecisionSuffix(phi_kernel_desc.kernelType.precision) +
        infrt::getPhiLayoutSuffix(phi_kernel_desc.kernelType.layout);

    mlir::OperationName operation_name(kernel_name, kernel_op.getContext());
    mlir::OperationState operation_state(kernel_op.getLoc(), operation_name);

    if (phi_context.find(phi_kernel_desc.kernelType.target) ==
        phi_context.end()) {
      switch (phi_kernel_desc.kernelType.target) {
        case infrt::TargetType::CPU: {
          auto context_value =
              builder
                  .create<infrt::phi::CreateCPUContextOp>(
                      kernel_op.getLoc(),
                      infrt::phi::ContextType::get(kernel_op.getContext(),
                                                   infrt::TargetType::CPU))
                  .output();
          phi_context[infrt::TargetType::CPU] = context_value;
        } break;
        case infrt::TargetType::GPU:
        case infrt::TargetType::UNK:
        default:
          LOG(FATAL) << "Unsupported TargetType";
          break;
      }
    }
    operation_state.addOperands(
        phi_context.at(phi_kernel_desc.kernelType.target));
    for (size_t index = 0; index < phi_kernel_desc.inputsType.size(); ++index) {
      mlir::Value input = kernel_op.getOperand(index);
      auto cvt_tensor_type_op = builder.create<infrt::CvtTensorOp>(
          kernel_op.getLoc(),
          infrt::DenseTensorType::get(
              kernel_op.getContext(),
              phi_kernel_desc.inputsType[index].target,
              phi_kernel_desc.inputsType[index].precision,
              phi_kernel_desc.inputsType[index].layout),
          input);
      operation_state.addOperands(cvt_tensor_type_op.output());
    }
    for (size_t index = 0; index < phi_kernel_desc.outputsType.size();
         ++index) {
      operation_state.addTypes(infrt::DenseTensorType::get(
          kernel_op.getContext(),
          phi_kernel_desc.outputsType[index].target,
          phi_kernel_desc.outputsType[index].precision,
          phi_kernel_desc.outputsType[index].layout));
    }
    operation_state.addAttributes(kernel_op.attrsAttr().getValue());
    mlir::Operation *phi_operation = builder.createOperation(operation_state);
    for (size_t index = 0; index < phi_kernel_desc.outputsType.size();
         ++index) {
      mlir::Value input = phi_operation->getResult(index);
      auto cvt_tensor_type_op = builder.create<infrt::CvtTensorOp>(
          kernel_op.getLoc(), kernel_op.getResultTypes()[index], input);
      kernel_op.getResult(index).replaceAllUsesWith(
          cvt_tensor_type_op.output());
    }
    kernel_op.erase();
  }
}

}  // namespace

std::unique_ptr<mlir::Pass> infrt::createPhiOpCvtPass(
    std::vector<Place> valid_places) {
  return std::make_unique<phiOpCvtPass>(valid_places);
}