prepared_operator.cc 6.9 KB
Newer Older
J
Jiabin Yang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
// Copyright (c) 2019 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/fluid/imperative/prepared_operator.h"
#include <sstream>
17 18 19
#include "paddle/fluid/imperative/execution_context.h"
#include "paddle/fluid/imperative/infer_shape_context.h"
#include "paddle/fluid/imperative/infer_var_type_context.h"
J
Jiabin Yang 已提交
20 21 22 23 24 25 26 27 28 29 30 31 32 33

namespace paddle {
namespace imperative {

const framework::Tensor* GetTensorFromVar(const framework::Variable& var) {
  if (var.IsType<framework::LoDTensor>()) {
    return &(var.Get<framework::LoDTensor>());
  } else if (var.IsType<framework::SelectedRows>()) {
    return &(var.Get<framework::SelectedRows>().value());
  } else {
    return nullptr;
  }
}

34 35 36
template <typename VarType>
static void PrepareDataImpl(
    const platform::Place& place, const NameVarMap<VarType>& ins,
37 38 39 40
    const framework::OperatorWithKernel& op,
    const framework::OpKernelType& expected_kernel_key) {
  for (const auto& name_pair : ins) {
    for (const auto& var_base : name_pair.second) {
J
Jiabin Yang 已提交
41 42 43
      const auto* tensor = GetTensorFromVar(var_base->Var());
      if (tensor && tensor->IsInitialized()) {
        auto tmp_place = tensor->place();
H
hong 已提交
44

45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
        // TODO(jiabin): Support transform data layout when we Verify it on more
        // tests
        if (!(tmp_place == place)) {
          auto kernel_type_for_var = op.GetKernelTypeForVar(
              name_pair.first, *tensor, expected_kernel_key);
          if (!NeedTransform(kernel_type_for_var, expected_kernel_key)) {
            continue;
          } else {
            VLOG(3) << "Transform Variable " << var_base->Name() << " from "
                    << kernel_type_for_var << " to " << expected_kernel_key;
            framework::Tensor out;
            TransformData(expected_kernel_key, kernel_type_for_var, *tensor,
                          &out);
            SetTensorToVariable(var_base->Var(), out, var_base->MutableVar());
          }
        }
J
Jiabin Yang 已提交
61 62 63 64 65
      }
    }
  }
}

66 67 68 69 70 71 72 73 74 75 76 77 78 79
void PreparedOp::PrepareData(
    const platform::Place& place, const NameVarMap<VarBase>& ins,
    const framework::OperatorWithKernel& op,
    const framework::OpKernelType& expected_kernel_key) {
  PrepareDataImpl<VarBase>(place, ins, op, expected_kernel_key);
}

void PreparedOp::PrepareData(
    const platform::Place& place, const NameVarMap<VariableWrapper>& ins,
    const framework::OperatorWithKernel& op,
    const framework::OpKernelType& expected_kernel_key) {
  PrepareDataImpl<VariableWrapper>(place, ins, op, expected_kernel_key);
}

J
Jiabin Yang 已提交
80 81
PreparedOp::PreparedOp(const framework::OperatorBase& op,
                       const framework::RuntimeContext& ctx,
82
                       const framework::OperatorWithKernel::OpKernelFunc& func,
83 84
                       platform::DeviceContext* dev_ctx)
    : op_(op), ctx_(ctx), func_(func), dev_ctx_(dev_ctx) {}
J
Jiabin Yang 已提交
85

86 87 88 89 90 91
template <typename VarType>
PreparedOp PrepareOpImpl(const NameVarMap<VarType>& ins,
                         const NameVarMap<VarType>& outs,
                         const framework::OperatorWithKernel& op,
                         platform::Place place,
                         const framework::AttributeMap& attrs) {
92 93
  platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
  auto* dev_ctx = pool.Get(place);
J
Jiabin Yang 已提交
94 95 96 97 98 99 100 101 102 103 104 105

  // check if op[type] has kernel registered.
  auto& all_op_kernels = op.AllOpKernels();
  auto kernels_iter = all_op_kernels.find(op.Type());
  if (kernels_iter == all_op_kernels.end()) {
    PADDLE_THROW(
        "There are no kernels which are registered in the %s operator.",
        op.Type());
  }

  auto& kernels = kernels_iter->second;

H
hong 已提交
106
  framework::RuntimeContext ctx({}, {});
107 108
  auto expected_kernel_key =
      op.GetExpectedKernelType(DygraphExecutionContext<VarType>(
109
          op, framework::Scope(), *dev_ctx, ctx, ins, outs, attrs));
J
Jiabin Yang 已提交
110 111 112 113 114 115 116 117
  VLOG(3) << "expected_kernel_key:" << expected_kernel_key;

  auto kernel_iter = kernels.find(expected_kernel_key);
  // TODO(jiabin): Add operator.cc's line 1000 part back when we need that case
  if (kernel_iter == kernels.end()) {
    PADDLE_THROW("op %s does not have kernel for %s", op.Type(),
                 KernelTypeToString(expected_kernel_key));
  }
118 119 120 121 122 123

  if (!(expected_kernel_key.place_ == place)) {
    dev_ctx = pool.Get(expected_kernel_key.place_);
    place = dev_ctx->GetPlace();
  }

124
  PrepareDataImpl<VarType>(place, ins, op, expected_kernel_key);
125
  return PreparedOp(op, ctx, kernel_iter->second, dev_ctx);
J
Jiabin Yang 已提交
126 127
}

128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147
PreparedOp PreparedOp::Prepare(const NameVarMap<VarBase>& ins,
                               const NameVarMap<VarBase>& outs,
                               const framework::OperatorWithKernel& op,
                               const platform::Place& place,
                               const framework::AttributeMap& attrs) {
  return PrepareOpImpl<VarBase>(ins, outs, op, place, attrs);
}

PreparedOp PreparedOp::Prepare(const NameVarMap<VariableWrapper>& ins,
                               const NameVarMap<VariableWrapper>& outs,
                               const framework::OperatorWithKernel& op,
                               const platform::Place& place,
                               const framework::AttributeMap& attrs) {
  return PrepareOpImpl<VariableWrapper>(ins, outs, op, place, attrs);
}

template <typename VarType>
static void PreparedOpRunImpl(
    const framework::OperatorBase& op, const framework::RuntimeContext& ctx,
    const framework::OperatorWithKernel::OpKernelFunc& func,
148 149
    platform::DeviceContext* dev_ctx, const NameVarMap<VarType>& ins,
    const NameVarMap<VarType>& outs, const framework::AttributeMap& attrs) {
J
Jiabin Yang 已提交
150 151
  // TODO(zjl): remove scope in dygraph
  framework::Scope scope;
H
hong 已提交
152

153 154 155
  DygraphInferShapeContext<VarType> infer_shape_ctx(&ins, &outs, &attrs);
  static_cast<const framework::OperatorWithKernel&>(op).InferShape(
      &infer_shape_ctx);
H
hong 已提交
156

157 158
  func(DygraphExecutionContext<VarType>(op, scope, *dev_ctx, ctx, ins, outs,
                                        attrs));
159
}
H
hong 已提交
160

161 162 163
void PreparedOp::Run(const NameVarMap<VarBase>& ins,
                     const NameVarMap<VarBase>& outs,
                     const framework::AttributeMap& attrs) {
164
  PreparedOpRunImpl<VarBase>(op_, ctx_, func_, dev_ctx_, ins, outs, attrs);
165
}
H
hong 已提交
166

167 168 169
void PreparedOp::Run(const NameVarMap<VariableWrapper>& ins,
                     const NameVarMap<VariableWrapper>& outs,
                     const framework::AttributeMap& attrs) {
170 171
  PreparedOpRunImpl<VariableWrapper>(op_, ctx_, func_, dev_ctx_, ins, outs,
                                     attrs);
J
Jiabin Yang 已提交
172 173 174 175
}

}  // namespace imperative
}  // namespace paddle