prepared_operator.cc 24.0 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 22
#include "paddle/pten/common/scalar.h"
#include "paddle/utils/small_vector.h"
Q
QingshuChen 已提交
23
#ifdef PADDLE_WITH_XPU
24
#include "paddle/fluid/platform/device/xpu/xpu_op_list.h"
Q
QingshuChen 已提交
25
#endif
26
DECLARE_bool(check_nan_inf);
27
DECLARE_bool(run_pten_kernel);
28
DECLARE_bool(benchmark);
29

J
Jiabin Yang 已提交
30 31 32
namespace paddle {
namespace imperative {

33 34 35 36 37 38 39 40 41 42
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 已提交
43 44 45 46 47 48 49 50 51 52
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;
  }
}

53 54 55 56 57 58 59 60 61 62 63 64 65 66 67
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;
}

68
template <typename VarType>
69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85
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 已提交
86
      if (tensor && tensor->IsInitialized()) {
87 88 89 90 91 92 93 94
        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 已提交
95 96 97 98 99 100 101
      }
    }
  }
}

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

111 112 113 114 115
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,
116
                       pten::KernelContext* pt_kernel_context,
117 118 119 120 121 122 123 124
                       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),
125 126
      pt_kernel_(pt_kernel),
      pt_kernel_context_(pt_kernel_context) {}
127

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

139 140 141 142 143 144 145 146
  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());
147 148 149 150
    mutable_op_attrs = default_attrs;
    for (auto& attr : attrs) {
      mutable_op_attrs[attr.first] = attr.second;
    }
151 152
  }
#endif
J
Jiabin Yang 已提交
153

154
  // 1. get expected kernel key
155 156 157
  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);
158 159
  VLOG(3) << "expected_kernel_key:" << expected_kernel_key;

160 161 162
  if (FLAGS_run_pten_kernel &&
      pten::KernelFactory::Instance().HasCompatiblePtenKernel(op.Type())) {
    auto pt_kernel_signature = op.GetExpectedPtenKernelArgs(dygraph_exe_ctx);
C
Chen Weihang 已提交
163
    VLOG(6) << framework::KernelSignatureToString(pt_kernel_signature);
164 165 166 167 168 169 170

    auto pt_kernel_name = pten::KernelName(pt_kernel_signature.name);
    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 已提交
171
      VLOG(6) << "Dynamic mode PrepareImpl - kernel name: " << pt_kernel_name
172 173 174 175 176
              << " | 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,
177
                        pt_kernel, pt_kernel_context, dev_ctx);
178
    } else {
C
Chen Weihang 已提交
179
      VLOG(6) << "Dynamic mode ChoosePtenKernel - kernel `" << pt_kernel_name
180 181 182 183
              << "` not found.";
    }
  }

184
  // 2. check if op[type] has kernel registered.
J
Jiabin Yang 已提交
185 186
  auto& all_op_kernels = op.AllOpKernels();
  auto kernels_iter = all_op_kernels.find(op.Type());
187 188 189 190 191
  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 已提交
192 193 194

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

224 225 226 227
  if (!(expected_kernel_key.place_ == place)) {
    dev_ctx = pool.Get(expected_kernel_key.place_);
  }

228
  return PreparedOp(op, ctx, expected_kernel_key, kernel_iter->second, dev_ctx);
229 230
}

231 232 233 234
PreparedOp PreparedOp::Prepare(const NameVarMap<VarBase>& ins,
                               const NameVarMap<VarBase>& outs,
                               const framework::OperatorWithKernel& op,
                               const platform::Place& place,
235
                               const framework::AttributeMap& attrs,
236 237 238 239
                               const framework::AttributeMap& default_attrs,
                               pten::KernelContext* pt_kernel_context) {
  return PrepareImpl<VarBase>(ins, outs, op, place, attrs, default_attrs,
                              pt_kernel_context);
240 241 242 243 244 245
}

PreparedOp PreparedOp::Prepare(const NameVarMap<VariableWrapper>& ins,
                               const NameVarMap<VariableWrapper>& outs,
                               const framework::OperatorWithKernel& op,
                               const platform::Place& place,
246
                               const framework::AttributeMap& attrs,
247 248
                               const framework::AttributeMap& default_attrs,
                               pten::KernelContext* pt_kernel_context) {
249
  return PrepareImpl<VariableWrapper>(ins, outs, op, place, attrs,
250
                                      default_attrs, pt_kernel_context);
251 252
}

253
template <typename VarType>
254
static void BuildDygraphPtenKernelContext(
255 256 257 258
    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,
259
    platform::DeviceContext* dev_ctx, pten::KernelContext* kernel_ctx) {
260 261 262 263 264 265 266
  // TODO(chenweihang): now only work for very simple case,
  // many cases need to be deal with later:
  // 1. the input and output are not tensor
  // 2. the dispensbale, duplicable input and output
  // 3. needless attributes remove
  // 4. use pt Tensor directly
  // 5. kernel input is not DenseTensor
267
  kernel_ctx->SetDeviceContext(dev_ctx);
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

  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]);
298 299 300 301 302 303 304 305 306 307

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

    // The current size of input/output in pt_kernel_context_ is at least equal
    // the start_idx. For the reason of reusing the allocted of inputs or
    // outputs in pt_kernel_context_, the current size of input/output can be
    // greater then the index of which the tensort wanted to set to, so it will
    // use ReMakePtenDenseTensorFromVar to make pten tensor.
    if (kernel_ctx->InputsSize() == start_idx) {
308 309 310 311 312 313 314
      paddle::SmallVector<std::shared_ptr<pten::TensorBase>> tmp_inputs;
      for (const auto& var : ins_vector) {
        const auto& variable = var->Var();
        tmp_inputs.emplace_back(
            experimental::MakePtenTensorBaseFromVar(variable, in_def));
      }
      kernel_ctx->EmplaceBackInputs(std::move(tmp_inputs));
315
    } else if (kernel_ctx->InputsSize() > start_idx) {
316 317
      size_t input_size = kernel_ctx->InputsSize();
      for (size_t j = 0; j < ins_vector.size(); ++j) {
318
        if (input_size > start_idx + j) {
319 320
          experimental::ReMakePtenDenseTensorFromVar(
              ins_vector[j]->Var(), in_def,
321 322 323 324 325 326 327 328
              kernel_ctx->MutableInputAt<pten::DenseTensor>(start_idx + j));
          // TODO(chentianyu03): When multi input kernel, open this code
          /*
          } else {
            kernel_ctx->EmplaceBackInputWithoutSetRange(
                experimental::MakePtenTensorBaseFromVar(ins_vector[j]->Var(),
                                                        in_def));
          */
329 330
        }
      }
331 332 333 334 335 336 337
      kernel_ctx->MutableInputRangeAt(i) = std::make_pair(start_idx, end_idx);
    } else {
      PADDLE_THROW(platform::errors::PreconditionNotMet(
          "Error start index when trying to set new tensor to inputs, start "
          "index is `%d`, but current pt_kernel_context_.inputs.size() is "
          "`%d`.",
          start_idx, kernel_ctx->InputsSize()));
338 339 340 341 342 343
    }
  }

  for (size_t i = 0; i < output_names.size(); ++i) {
    auto& out_def = output_defs.at(i);
    auto& outs_vector = outs.at(output_names[i]);
344 345 346 347 348 349 350 351 352 353

    size_t start_idx = (i == 0 ? 0 : kernel_ctx->OutputRangeAt(i - 1).second);
    size_t end_idx = start_idx + outs_vector.size();

    // The current size of input/output in pt_kernel_context_ is at least equal
    // the start_idx. For the reason of reusing the allocted of inputs or
    // outputs in pt_kernel_context_, the current size of input/output can be
    // greater then the index of which the tensort wanted to set to, so it will
    // use ReMakePtenDenseTensorFromVar to make pten tensor.
    if (kernel_ctx->OutputsSize() == start_idx) {
354 355 356 357 358 359 360
      paddle::SmallVector<std::shared_ptr<pten::TensorBase>> tmp_outputs;
      for (auto& var : outs_vector) {
        auto* variable = var->MutableVar();
        tmp_outputs.emplace_back(
            experimental::MakePtenTensorBaseFromVar(variable, out_def));
      }
      kernel_ctx->EmplaceBackOutputs(std::move(tmp_outputs));
361
    } else if (kernel_ctx->OutputsSize() > start_idx) {
362 363 364 365 366 367
      size_t output_size = kernel_ctx->OutputsSize();
      for (size_t j = 0; j < outs_vector.size(); ++j) {
        if (output_size > i + j) {
          experimental::ReMakePtenDenseTensorFromVar(
              outs_vector[j]->MutableVar(), out_def,
              kernel_ctx->MutableOutputAt<pten::DenseTensor>(i + j));
368 369 370 371 372 373 374
          // TODO(chentianyu03): When multi output kernel, open this code
          /*
          } else {
            kernel_ctx->EmplaceBackOutputWithoutSetRange(
                experimental::MakePtenTensorBaseFromVar(
                    outs_vector[j]->MutableVar(), out_def));
          */
375 376
        }
      }
377 378 379 380 381 382 383
      kernel_ctx->MutableOutputRangeAt(i) = std::make_pair(start_idx, end_idx);
    } else {
      PADDLE_THROW(platform::errors::PreconditionNotMet(
          "Error start index when trying to set new tensor to inputs, start "
          "index is `%d`, but current pt_kernel_context_.outputs.size() is "
          "`%d`.",
          start_idx, kernel_ctx->OutputsSize()));
384 385 386 387 388 389 390 391 392 393
    }
  }

  for (size_t i = 0; i < attr_names.size(); ++i) {
    auto& attr = GetAttr(attrs, default_attrs, attr_names[i]);
    if (attr_defs[i].type_index == std::type_index(typeid(pten::Scalar))) {
      // TODO(chenweihang): support other attrs later
      // TODO(zhangyunfei): Scalar should hold scaler type, and we should check
      // attribtue type by attr_defs
      if (std::type_index(attr.type()) == std::type_index(typeid(float))) {
394
        kernel_ctx->EmplaceBackAttr(
395
            std::move(pten::Scalar(BOOST_GET_CONST(float, attr))));
396 397
      } else if (std::type_index(attr.type()) ==
                 std::type_index(typeid(std::string))) {
398
        kernel_ctx->EmplaceBackAttr(
399
            std::move(pten::Scalar(BOOST_GET_CONST(std::string, attr))));
400 401 402 403 404 405 406 407 408
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
            "unsupported cast op attribute `%s` to Scalar when construct "
            "KernelContext in dygraph.",
            attr_names[i]));
      }
    } else {
      // TODO(chenweihang): support other attrs later
      if (attr_defs[i].type_index == std::type_index(typeid(int))) {
409
        kernel_ctx->EmplaceBackAttr(BOOST_GET_CONST(int, attr));
410
      } else if (attr_defs[i].type_index == std::type_index(typeid(float))) {
411
        kernel_ctx->EmplaceBackAttr(BOOST_GET_CONST(float, attr));
412
      } else if (attr_defs[i].type_index == std::type_index(typeid(bool))) {
413
        kernel_ctx->EmplaceBackAttr(BOOST_GET_CONST(bool, attr));
414
      } else if (attr_defs[i].type_index ==
415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430
                 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
431 432 433 434 435 436 437 438 439 440
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
            "unsupported cast op attribute `%s` when construct "
            "KernelContext in dygraph.",
            attr_names[i]));
      }
    }
  }
}

441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460
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) {
    auto& outs_vector = outs.at(output_names[i]);

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

    for (size_t j = 0; j < pten_outs.size(); ++j) {
      experimental::MakeVariableFromPtenTensor(pten_outs[j],
                                               outs_vector[j]->MutableVar());
    }
  }
}

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

472
  DygraphInferShapeContext<VarType> infer_shape_ctx(&ins, &outs, &attrs,
473
                                                    &default_attrs, op.Type());
474 475
  static_cast<const framework::OperatorWithKernel&>(op).InferShape(
      &infer_shape_ctx);
H
hong 已提交
476

477
  func(DygraphExecutionContext<VarType>(op, scope, *dev_ctx, ctx, ins, outs,
478
                                        attrs, default_attrs));
479

480 481 482 483 484
  if (FLAGS_check_nan_inf) {
    framework::details::CheckOpHasNanOrInfInDygraph<VarType>(
        op.Type(), outs, dev_ctx->GetPlace());
  }

485 486 487 488 489 490 491 492 493 494 495 496 497 498 499
  /**
   * [ 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);
  }
500
}
H
hong 已提交
501

502 503 504 505
template <typename VarType>
static void PreparedOpRunPtImpl(
    const framework::OperatorBase& op,
    const framework::KernelSignature& pt_kernel_signature,
506 507 508
    const pten::Kernel& pt_kernel, pten::KernelContext* pt_kernel_context,
    platform::DeviceContext* dev_ctx, const NameVarMap<VarType>& ins,
    const NameVarMap<VarType>& outs, const framework::AttributeMap& attrs,
509 510 511 512 513 514
    const framework::AttributeMap& default_attrs) {
  DygraphInferShapeContext<VarType> infer_shape_ctx(&ins, &outs, &attrs,
                                                    &default_attrs, op.Type());
  static_cast<const framework::OperatorWithKernel&>(op).InferShape(
      &infer_shape_ctx);

515 516 517 518 519
  BuildDygraphPtenKernelContext<VarType>(pt_kernel_signature, pt_kernel, ins,
                                         outs, attrs, default_attrs, dev_ctx,
                                         pt_kernel_context);

  pt_kernel(pt_kernel_context);
520

521 522 523 524 525 526 527 528 529 530 531 532 533 534
  if (FLAGS_benchmark) {
    dev_ctx->Wait();
#if defined(PADDLE_WITH_CUDA)
    PADDLE_ENFORCE_CUDA_SUCCESS(cudaGetLastError());
    VLOG(4) << "Operator(" << op.Type() << "): context wait and get last error";
#endif
#if defined(PADDLE_WITH_HIP)
    PADDLE_ENFORCE_CUDA_SUCCESS(hipGetLastError());
    VLOG(4) << "Operator(" << op.Type() << "): context wait and get last error";
#endif
  }

  WriteBackToOutputs<VarType>(pt_kernel_signature, outs, pt_kernel_context);

535 536
  // Ensure that it does not affect the VarBase life cycle management
  pt_kernel_context->ClearData();
537 538 539 540 541

  // TODO(chenweihang): add debug flags later
  // TODO(chenweihang): deal with complex cases later
}

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

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

}  // namespace imperative
}  // namespace paddle