prepared_operator.cc 23.2 KB
Newer Older
J
Jiabin Yang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
// 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"
16

17
#include "paddle/fluid/framework/data_type_transform.h"
18
#include "paddle/fluid/framework/details/nan_inf_utils.h"
19
#include "paddle/fluid/imperative/infer_shape_context.h"
20
#include "paddle/fluid/imperative/tracer.h"
21
#include "paddle/pten/common/scalar.h"
22
#include "paddle/pten/common/scalar_array.h"
23
#include "paddle/utils/small_vector.h"
Q
QingshuChen 已提交
24
#ifdef PADDLE_WITH_XPU
25
#include "paddle/fluid/platform/device/xpu/xpu_op_list.h"
Q
QingshuChen 已提交
26
#endif
27 28
#include "paddle/fluid/platform/device/gpu/gpu_info.h"

29
DECLARE_bool(check_nan_inf);
30
DECLARE_bool(run_pten_kernel);
31
DECLARE_bool(benchmark);
F
Feng Xing 已提交
32
DECLARE_bool(run_kp_kernel);
33

J
Jiabin Yang 已提交
34 35 36
namespace paddle {
namespace imperative {

37 38 39 40 41 42 43 44 45 46
const std::shared_ptr<VariableWrapper>& GetVariableWrapper(
    const std::shared_ptr<paddle::imperative::VarBase>& var) {
  return var->SharedVar();
}

const std::shared_ptr<VariableWrapper>& GetVariableWrapper(
    const std::shared_ptr<VariableWrapper>& var) {
  return var;
}

J
Jiabin Yang 已提交
47 48 49 50 51 52 53 54 55 56
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;
  }
}

57 58 59 60 61 62 63 64 65 66 67 68 69 70 71
static const framework::Attribute& GetAttr(
    const framework::AttributeMap& attrs,
    const framework::AttributeMap& default_attrs, const std::string& name) {
  auto it = attrs.find(name);
  bool found = it != attrs.end();
  if (!found) {
    it = default_attrs.find(name);
    found = it != default_attrs.end();
  }
  PADDLE_ENFORCE_EQ(
      found, true,
      platform::errors::NotFound("(%s) is not found in AttributeMap.", name));
  return it->second;
}

72
template <typename VarType>
73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89
static void HandleComplexGradToRealGrad(const NameVarMap<VarType>& outs) {
  for (auto& pair : outs) {
    for (auto& var : pair.second) {
      if (var == nullptr) {
        continue;
      }
      if (var->ForwardDataType() ==
          static_cast<framework::proto::VarType::Type>(-1)) {
        VLOG(6) << "Var (" << var->Name()
                << ")'s forward data type is not set.";
        continue;
      }
      if (!framework::IsComplexType(var->DataType()) ||
          framework::IsComplexType(var->ForwardDataType())) {
        continue;
      }
      const auto* tensor = GetTensorFromVar(var->Var());
J
Jiabin Yang 已提交
90
      if (tensor && tensor->IsInitialized()) {
91 92 93 94 95 96 97 98
        VLOG(6) << "Transform " << framework::DataTypeToString(var->DataType())
                << " var `" << var->Name() << "` to "
                << framework::DataTypeToString(var->ForwardDataType())
                << " real var in dynamic graph.";
        framework::Tensor out;
        framework::TransComplexToReal(var->ForwardDataType(), var->DataType(),
                                      *tensor, &out);
        SetTensorToVariable(var->Var(), out, var->MutableVar());
J
Jiabin Yang 已提交
99 100 101 102 103 104 105
      }
    }
  }
}

PreparedOp::PreparedOp(const framework::OperatorBase& op,
                       const framework::RuntimeContext& ctx,
106
                       const framework::OpKernelType& kernel_type,
107
                       const framework::OperatorWithKernel::OpKernelFunc& func,
108
                       platform::DeviceContext* dev_ctx)
109 110 111 112 113 114
    : op_(op),
      ctx_(ctx),
      kernel_type_(kernel_type),
      func_(func),
      dev_ctx_(dev_ctx) {}

115 116 117 118 119 120 121 122 123 124 125 126 127
PreparedOp::PreparedOp(const framework::OperatorBase& op,
                       const framework::RuntimeContext& ctx,
                       const framework::OpKernelType& kernel_type,
                       const framework::KernelSignature& kernel_signature,
                       const pten::Kernel& pt_kernel,
                       platform::DeviceContext* dev_ctx)
    : op_(op),
      ctx_(ctx),
      kernel_type_(kernel_type),
      func_(nullptr),
      dev_ctx_(dev_ctx),
      run_pten_kernel_(true),
      pt_kernel_signature_(kernel_signature),
128
      pt_kernel_(pt_kernel) {}
129

130 131 132 133 134
template <typename VarType>
PreparedOp PrepareImpl(const NameVarMap<VarType>& ins,
                       const NameVarMap<VarType>& outs,
                       const framework::OperatorWithKernel& op,
                       const platform::Place& place,
135
                       const framework::AttributeMap& attrs,
136
                       const framework::AttributeMap& default_attrs) {
137
  platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
138
  auto* dev_ctx = pool.Get(place);
139

140 141 142 143 144 145 146 147
  framework::RuntimeContext ctx({}, {});

#ifdef PADDLE_WITH_MKLDNN
  // MKLDNN variant of code reads attributes in some of GetKernelTypeForVar and
  // GetKernelType functions, so we need to copy the attributes there.
  // Const qualifier of Attrs had to be discarded to overwrite it.
  if (FLAGS_use_mkldnn) {
    auto& mutable_op_attrs = const_cast<framework::AttributeMap&>(op.Attrs());
148 149 150 151
    mutable_op_attrs = default_attrs;
    for (auto& attr : attrs) {
      mutable_op_attrs[attr.first] = attr.second;
    }
152 153
  }
#endif
J
Jiabin Yang 已提交
154

155
  // 1. get expected kernel key
156 157 158
  auto dygraph_exe_ctx = DygraphExecutionContext<VarType>(
      op, framework::Scope(), *dev_ctx, ctx, ins, outs, attrs, default_attrs);
  auto expected_kernel_key = op.GetExpectedKernelType(dygraph_exe_ctx);
159 160
  VLOG(3) << "expected_kernel_key:" << expected_kernel_key;

161 162 163
  if (FLAGS_run_pten_kernel &&
      pten::KernelFactory::Instance().HasCompatiblePtenKernel(op.Type())) {
    auto pt_kernel_signature = op.GetExpectedPtenKernelArgs(dygraph_exe_ctx);
164
    VLOG(6) << pt_kernel_signature;
165

Y
YuanRisheng 已提交
166
    auto pt_kernel_name = pt_kernel_signature.name;
167 168 169 170 171
    auto pt_kernel_key = TransOpKernelTypeToPtenKernelKey(expected_kernel_key);
    auto pt_kernel = pten::KernelFactory::Instance().SelectKernel(
        pt_kernel_name, pt_kernel_key);

    if (pt_kernel.IsValid()) {
C
Chen Weihang 已提交
172
      VLOG(6) << "Dynamic mode PrepareImpl - kernel name: " << pt_kernel_name
173 174 175 176 177
              << " | kernel key: " << pt_kernel_key
              << " | kernel: " << pt_kernel;

      // TODO(chenweihang): using CPUKernel when miss device kernel case
      return PreparedOp(op, ctx, expected_kernel_key, pt_kernel_signature,
178
                        pt_kernel, dev_ctx);
179
    } else {
C
Chen Weihang 已提交
180
      VLOG(6) << "Dynamic mode ChoosePtenKernel - kernel `" << pt_kernel_name
181 182 183 184
              << "` not found.";
    }
  }

185
  // 2. check if op[type] has kernel registered.
J
Jiabin Yang 已提交
186 187
  auto& all_op_kernels = op.AllOpKernels();
  auto kernels_iter = all_op_kernels.find(op.Type());
188 189 190 191 192
  PADDLE_ENFORCE_NE(
      kernels_iter, all_op_kernels.end(),
      platform::errors::NotFound(
          "There are no kernels which are registered in the %s operator.",
          op.Type()));
J
Jiabin Yang 已提交
193 194 195

  auto& kernels = kernels_iter->second;
  auto kernel_iter = kernels.find(expected_kernel_key);
196
#ifdef PADDLE_WITH_XPU
197
  if (paddle::platform::is_xpu_place(expected_kernel_key.place_) &&
Q
QingshuChen 已提交
198 199 200
      (kernel_iter == kernels.end() ||
       !paddle::platform::is_xpu_support_op(op.Type(), expected_kernel_key) ||
       paddle::platform::is_in_xpu_black_list(op.Type()))) {
201 202 203
    VLOG(3) << "missing XPU kernel: " << op.Type()
            << ", expected_kernel_key:" << expected_kernel_key
            << ", fallbacking to CPU one!";
204 205 206
    expected_kernel_key.place_ = platform::CPUPlace();
    kernel_iter = kernels.find(expected_kernel_key);
  }
207 208 209
#endif
#ifdef PADDLE_WITH_ASCEND_CL
  if (kernel_iter == kernels.end() &&
210
      paddle::platform::is_npu_place(expected_kernel_key.place_)) {
211 212 213
    VLOG(3) << "missing NPU kernel: " << op.Type()
            << ", expected_kernel_key:" << expected_kernel_key
            << ", fallbacking to CPU one!";
214 215 216
    expected_kernel_key.place_ = platform::CPUPlace();
    kernel_iter = kernels.find(expected_kernel_key);
  }
217 218 219
#endif
#ifdef PADDLE_WITH_MLU
  if (kernel_iter == kernels.end() &&
220
      paddle::platform::is_mlu_place(expected_kernel_key.place_)) {
221 222 223 224 225 226
    VLOG(3) << "missing MLU kernel: " << op.Type()
            << ", expected_kernel_key:" << expected_kernel_key
            << ", fallbacking to CPU one!";
    expected_kernel_key.place_ = platform::CPUPlace();
    kernel_iter = kernels.find(expected_kernel_key);
  }
227
#endif
228 229
  // TODO(jiabin): Add operator.cc's line 1000 part back when we need that
  // case
230 231 232 233
  PADDLE_ENFORCE_NE(kernel_iter, kernels.end(),
                    platform::errors::NotFound(
                        "Operator %s does not have kernel for %s.", op.Type(),
                        KernelTypeToString(expected_kernel_key)));
234

235 236 237 238
  if (!(expected_kernel_key.place_ == place)) {
    dev_ctx = pool.Get(expected_kernel_key.place_);
  }

239
  return PreparedOp(op, ctx, expected_kernel_key, kernel_iter->second, dev_ctx);
240 241
}

242 243 244 245
PreparedOp PreparedOp::Prepare(const NameVarMap<VarBase>& ins,
                               const NameVarMap<VarBase>& outs,
                               const framework::OperatorWithKernel& op,
                               const platform::Place& place,
246
                               const framework::AttributeMap& attrs,
247 248
                               const framework::AttributeMap& default_attrs) {
  return PrepareImpl<VarBase>(ins, outs, op, place, attrs, default_attrs);
249 250 251 252 253 254
}

PreparedOp PreparedOp::Prepare(const NameVarMap<VariableWrapper>& ins,
                               const NameVarMap<VariableWrapper>& outs,
                               const framework::OperatorWithKernel& op,
                               const platform::Place& place,
255
                               const framework::AttributeMap& attrs,
256
                               const framework::AttributeMap& default_attrs) {
257
  return PrepareImpl<VariableWrapper>(ins, outs, op, place, attrs,
258
                                      default_attrs);
259 260
}

261
template <typename VarType>
262
static void BuildDygraphPtenKernelContext(
263 264 265 266
    const framework::KernelSignature& pt_kernel_signature,
    const pten::Kernel& pt_kernel, const NameVarMap<VarType>& ins,
    const NameVarMap<VarType>& outs, const framework::AttributeMap& attrs,
    const framework::AttributeMap& default_attrs,
267 268
    platform::DeviceContext* dev_ctx, pten::KernelContext* kernel_ctx) {
  kernel_ctx->SetDeviceContext(dev_ctx);
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

  auto& input_names = std::get<0>(pt_kernel_signature.args);
  auto& attr_names = std::get<1>(pt_kernel_signature.args);
  auto& output_names = std::get<2>(pt_kernel_signature.args);

  auto& input_defs = pt_kernel.args_def().input_defs();
  auto& output_defs = pt_kernel.args_def().output_defs();
  auto& attr_defs = pt_kernel.args_def().attribute_defs();

  PADDLE_ENFORCE_EQ(input_names.size(), input_defs.size(),
                    platform::errors::InvalidArgument(
                        "the size of inputs_args names (%d) must be equal to "
                        "the size of kernel input_defs (%d).",
                        input_names.size(), input_defs.size()));

  PADDLE_ENFORCE_EQ(output_names.size(), output_defs.size(),
                    platform::errors::InvalidArgument(
                        "the size of outputs_args names (%d) must be equal to "
                        "the size of kernel output_defs (%d).",
                        output_names.size(), output_defs.size()));

  PADDLE_ENFORCE_EQ(attr_names.size(), attr_defs.size(),
                    platform::errors::InvalidArgument(
                        "the size of attribute_args names (%d) must be equal "
                        "to the size of kernel attribute_defs (%d).",
                        attr_names.size(), attr_defs.size()));

  for (size_t i = 0; i < input_names.size(); ++i) {
    auto& in_def = input_defs.at(i);
    auto& ins_vector = ins.at(input_names[i]);
299 300 301

    size_t start_idx = (i == 0 ? 0 : kernel_ctx->InputRangeAt(i - 1).second);
    size_t end_idx = start_idx + ins_vector.size();
302 303 304

    for (size_t offset = 0; offset < ins_vector.size(); ++offset) {
      const auto& variable = ins_vector[offset]->Var();
305 306
      kernel_ctx->EmplaceBackInputWithoutSetRange(
          paddle::experimental::MakePtenTensorBaseFromVar(variable, in_def));
307
    }
308
    kernel_ctx->AssignInputRange(std::make_pair(start_idx, end_idx), i);
309 310 311 312
  }

  for (size_t i = 0; i < output_names.size(); ++i) {
    auto& out_def = output_defs.at(i);
313 314

    size_t start_idx = (i == 0 ? 0 : kernel_ctx->OutputRangeAt(i - 1).second);
315 316 317

    auto iter = outs.find(output_names[i]);
    if (iter == outs.end()) {
318
      kernel_ctx->EmplaceBackOutputWithoutSetRange({nullptr});
319 320 321 322 323 324 325 326
      kernel_ctx->AssignOutputRange(std::make_pair(start_idx, start_idx + 1),
                                    i);
      continue;
    }

    auto& outs_vector = iter->second;
    size_t end_idx = start_idx + outs_vector.size();

327
    for (size_t offset = 0; offset < outs_vector.size(); ++offset) {
328 329 330
      kernel_ctx->EmplaceBackOutputWithoutSetRange(
          paddle::experimental::MakePtenTensorBaseFromVar(
              outs_vector[offset]->MutableVar(), out_def));
331
    }
332
    kernel_ctx->AssignOutputRange(std::make_pair(start_idx, end_idx), i);
333 334 335
  }

  for (size_t i = 0; i < attr_names.size(); ++i) {
336 337 338 339 340 341 342 343
    if (attr_defs[i].type_index == std::type_index(typeid(pten::ScalarArray))) {
      if (attrs.find(attr_names[i]) !=
          attrs.end()) {  // shape is in the attribute
        auto& attr = GetAttr(attrs, default_attrs, attr_names[i]);
        if (std::type_index(attr.type()) ==
            std::type_index(typeid(std::vector<int64_t>))) {
          kernel_ctx->EmplaceBackAttr(std::move(
              pten::ScalarArray(BOOST_GET_CONST(std::vector<int64_t>, attr))));
344 345 346 347
        } else if (std::type_index(attr.type()) ==
                   std::type_index(typeid(std::vector<int32_t>))) {
          kernel_ctx->EmplaceBackAttr(std::move(
              pten::ScalarArray(BOOST_GET_CONST(std::vector<int32_t>, attr))));
348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370
        } else {
          PADDLE_THROW(platform::errors::Unimplemented(
              "Unsupported cast op attribute `%s` to VectorTensor when "
              "construct KernelContext.",
              attr_names[i]));
        }
      } else {  // shape is in the input
        auto& ins_vector = ins.at(attr_names[i]);
        if (ins_vector.size() == 1) {  // ShapeTensor
          kernel_ctx->EmplaceBackAttr(std::move(
              experimental::MakePtenScalarArrayFromVar(ins_vector[0]->Var())));
        } else {  // ShapeTensorList
          std::vector<framework::Variable*> variables;
          variables.reserve(ins_vector.size());
          for (const auto& var_base : ins_vector) {
            variables.push_back(var_base->MutableVar());
          }
          kernel_ctx->EmplaceBackAttr(std::move(
              experimental::MakePtenScalarArrayFromVarList(variables)));
        }
      }
    } else if (attr_defs[i].type_index ==
               std::type_index(typeid(pten::Scalar))) {
371 372 373
      // TODO(chenweihang): support other attrs later
      // TODO(zhangyunfei): Scalar should hold scaler type, and we should check
      // attribtue type by attr_defs
374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394
      if (attrs.find(attr_names[i]) != attrs.end() ||
          default_attrs.find(attr_names[i]) !=
              default_attrs.end()) {  // scalar is in the attribute
        auto& attr = GetAttr(attrs, default_attrs, attr_names[i]);
        if (std::type_index(attr.type()) == std::type_index(typeid(float))) {
          kernel_ctx->EmplaceBackAttr(
              std::move(pten::Scalar(BOOST_GET_CONST(float, attr))));
        } else if (std::type_index(attr.type()) ==
                   std::type_index(typeid(std::string))) {
          kernel_ctx->EmplaceBackAttr(
              std::move(pten::Scalar(BOOST_GET_CONST(std::string, attr))));
        } else {
          PADDLE_THROW(platform::errors::Unimplemented(
              "Unsupported cast op attribute `%s` to Scalar when construct "
              "KernelContext in dygraph.",
              attr_names[i]));
        }
      } else {  // scalar is in the input
        auto& ins_vector = ins.at(attr_names[i]);
        kernel_ctx->EmplaceBackAttr(std::move(
            experimental::MakePtenScalarFromVar(ins_vector[0]->Var())));
395
      }
396

397 398
    } else {
      // TODO(chenweihang): support other attrs later
399
      auto& attr = GetAttr(attrs, default_attrs, attr_names[i]);
400
      if (attr_defs[i].type_index == std::type_index(typeid(int))) {
401
        kernel_ctx->EmplaceBackAttr(BOOST_GET_CONST(int, attr));
402
      } else if (attr_defs[i].type_index == std::type_index(typeid(float))) {
403
        kernel_ctx->EmplaceBackAttr(BOOST_GET_CONST(float, attr));
404
      } else if (attr_defs[i].type_index == std::type_index(typeid(bool))) {
405
        kernel_ctx->EmplaceBackAttr(BOOST_GET_CONST(bool, attr));
406
      } else if (attr_defs[i].type_index ==
407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422
                 std::type_index(typeid(pten::DataType))) {
        auto data_type = pten::TransToPtenDataType(
            static_cast<framework::proto::VarType::Type>(
                BOOST_GET_CONST(int, attr)));
        kernel_ctx->EmplaceBackAttr(data_type);
      } else if (attr_defs[i].type_index ==
                 std::type_index(typeid(std::vector<int64_t>))) {
        if (std::type_index(attr.type()) ==
            std::type_index(typeid(std::vector<int>))) {
          // Emplace Back Attr according to the type of Pten_Kernel args.
          const auto& vector_int_attr = BOOST_GET_CONST(std::vector<int>, attr);
          const std::vector<int64_t> vector_int64_attr(vector_int_attr.begin(),
                                                       vector_int_attr.end());
          kernel_ctx->EmplaceBackAttr(vector_int64_attr);
        }
        // TODO(YuanRisheng) Need support vector<int64_t> attr
423 424
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
425
            "Unsupported cast op attribute `%s` when construct "
426 427 428 429 430 431 432
            "KernelContext in dygraph.",
            attr_names[i]));
      }
    }
  }
}

433 434 435 436 437 438 439
template <typename VarType>
static void WriteBackToOutputs(
    const framework::KernelSignature& pt_kernel_signature,
    const NameVarMap<VarType>& outs, pten::KernelContext* kernel_ctx) {
  auto& output_names = std::get<2>(pt_kernel_signature.args);

  for (size_t i = 0; i < output_names.size(); ++i) {
440 441 442
    auto iter = outs.find(output_names[i]);
    if (iter != outs.end()) {
      auto& outs_vector = iter->second;
443

444 445 446
      auto& range_pair = kernel_ctx->OutputRangeAt(i);
      auto pten_outs = kernel_ctx->MutableOutputBetween<pten::DenseTensor>(
          range_pair.first, range_pair.second);
447

448 449 450 451
      for (size_t j = 0; j < pten_outs.size(); ++j) {
        experimental::MakeVariableFromPtenTensor(pten_outs[j],
                                                 outs_vector[j]->MutableVar());
      }
452 453 454 455
    }
  }
}

456 457 458
template <typename VarType>
static void PreparedOpRunImpl(
    const framework::OperatorBase& op, const framework::RuntimeContext& ctx,
459
    const framework::OpKernelType& kernel_type,
460
    const framework::OperatorWithKernel::OpKernelFunc& func,
461
    platform::DeviceContext* dev_ctx, const NameVarMap<VarType>& ins,
462 463
    const NameVarMap<VarType>& outs, const framework::AttributeMap& attrs,
    const framework::AttributeMap& default_attrs) {
J
Jiabin Yang 已提交
464 465
  // TODO(zjl): remove scope in dygraph
  framework::Scope scope;
H
hong 已提交
466

467 468
  DygraphInferShapeContext<VarType> infer_shape_ctx(
      &ins, &outs, &attrs, &default_attrs, op.Type(), &kernel_type);
469
  op.Info().infer_shape_(&infer_shape_ctx);
H
hong 已提交
470

471
  func(DygraphExecutionContext<VarType>(op, scope, *dev_ctx, ctx, ins, outs,
472
                                        attrs, default_attrs));
473

474 475 476 477 478
  if (FLAGS_check_nan_inf) {
    framework::details::CheckOpHasNanOrInfInDygraph<VarType>(
        op.Type(), outs, dev_ctx->GetPlace());
  }

L
Leo Chen 已提交
479 480 481 482 483 484 485 486
  if (FLAGS_benchmark) {
    dev_ctx->Wait();
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
    PADDLE_ENFORCE_GPU_SUCCESS(platform::GpuGetLastError());
    VLOG(4) << "Operator(" << op.Type() << "): context wait and get last error";
#endif
  }

487 488 489 490 491 492 493 494 495 496 497 498 499 500 501
  /**
   * [ Why need handle complex gradient to real gradient? ]
   *
   * After the introduction of complex number calculations, Ops that support
   * complex number calculations generally support type promotion, such as
   * x(float32) + y(complex64) = out(complex64), then the type of the grad
   * tensor should be dout(complex64), dx(float32), dy (complex64).
   *
   * But because the dout is complex64, the dx is also complex64 after
   * grad op kernel executed, we need to recognize this situation and
   * convert dx to float32 type. HandleComplexGradToRealGrad does this thing.
   */
  if (framework::IsComplexType(kernel_type.data_type_)) {
    HandleComplexGradToRealGrad<VarType>(outs);
  }
502
}
H
hong 已提交
503

504 505 506
template <typename VarType>
static void PreparedOpRunPtImpl(
    const framework::OperatorBase& op,
507
    const framework::OpKernelType& kernel_type,
508
    const framework::KernelSignature& pt_kernel_signature,
509 510 511
    const pten::Kernel& pt_kernel, platform::DeviceContext* dev_ctx,
    const NameVarMap<VarType>& ins, const NameVarMap<VarType>& outs,
    const framework::AttributeMap& attrs,
512
    const framework::AttributeMap& default_attrs) {
513 514
  DygraphInferShapeContext<VarType> infer_shape_ctx(
      &ins, &outs, &attrs, &default_attrs, op.Type(), &kernel_type);
515
  op.Info().infer_shape_(&infer_shape_ctx);
516

517
  pten::KernelContext pt_kernel_context;
518 519
  BuildDygraphPtenKernelContext<VarType>(pt_kernel_signature, pt_kernel, ins,
                                         outs, attrs, default_attrs, dev_ctx,
520
                                         &pt_kernel_context);
521

522
  pt_kernel(&pt_kernel_context);
523

524 525
  if (FLAGS_benchmark) {
    dev_ctx->Wait();
526 527
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
    PADDLE_ENFORCE_GPU_SUCCESS(platform::GpuGetLastError());
528 529 530 531
    VLOG(4) << "Operator(" << op.Type() << "): context wait and get last error";
#endif
  }

532
  WriteBackToOutputs<VarType>(pt_kernel_signature, outs, &pt_kernel_context);
533 534

  // TODO(chenweihang): add debug flags later
535 536 537
  if (framework::IsComplexType(kernel_type.data_type_)) {
    HandleComplexGradToRealGrad<VarType>(outs);
  }
538 539
}

540 541
void PreparedOp::Run(const NameVarMap<VarBase>& ins,
                     const NameVarMap<VarBase>& outs,
542 543
                     const framework::AttributeMap& attrs,
                     const framework::AttributeMap& default_attrs) {
544
  if (run_pten_kernel_) {
545
    PreparedOpRunPtImpl<VarBase>(op_, kernel_type_, pt_kernel_signature_,
546 547
                                 pt_kernel_, dev_ctx_, ins, outs, attrs,
                                 default_attrs);
548 549 550 551
  } else {
    PreparedOpRunImpl<VarBase>(op_, ctx_, kernel_type_, func_, dev_ctx_, ins,
                               outs, attrs, default_attrs);
  }
552
}
H
hong 已提交
553

554 555
void PreparedOp::Run(const NameVarMap<VariableWrapper>& ins,
                     const NameVarMap<VariableWrapper>& outs,
556 557
                     const framework::AttributeMap& attrs,
                     const framework::AttributeMap& default_attrs) {
558
  if (run_pten_kernel_) {
559
    PreparedOpRunPtImpl<VariableWrapper>(
560 561
        op_, kernel_type_, pt_kernel_signature_, pt_kernel_, dev_ctx_, ins,
        outs, attrs, default_attrs);
562 563 564 565
  } else {
    PreparedOpRunImpl<VariableWrapper>(op_, ctx_, kernel_type_, func_, dev_ctx_,
                                       ins, outs, attrs, default_attrs);
  }
J
Jiabin Yang 已提交
566 567 568 569
}

}  // namespace imperative
}  // namespace paddle