prepared_operator.cc 19.4 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

J
Jiabin Yang 已提交
17
#include "paddle/fluid/eager/eager_tensor.h"
18
#include "paddle/fluid/framework/data_type_transform.h"
19
#include "paddle/fluid/framework/details/nan_inf_utils.h"
20
#include "paddle/fluid/imperative/infer_shape_context.h"
21
#include "paddle/fluid/imperative/tracer.h"
22
#include "paddle/pten/common/scalar.h"
23
#include "paddle/pten/common/scalar_array.h"
24
#include "paddle/utils/small_vector.h"
Q
QingshuChen 已提交
25
#ifdef PADDLE_WITH_XPU
26
#include "paddle/fluid/platform/device/xpu/xpu_op_list.h"
Q
QingshuChen 已提交
27
#endif
L
Liu-xiandong 已提交
28
#include "paddle/fluid/framework/library_type.h"
29
#include "paddle/fluid/platform/device/gpu/gpu_info.h"
C
chenjian 已提交
30
#include "paddle/fluid/platform/profiler/event_tracing.h"
31

32
DECLARE_bool(check_nan_inf);
33
DECLARE_bool(benchmark);
F
Feng Xing 已提交
34
DECLARE_bool(run_kp_kernel);
35

J
Jiabin Yang 已提交
36 37 38
namespace paddle {
namespace imperative {

39 40 41 42 43 44 45 46 47 48
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 已提交
49 50 51
const framework::Tensor* GetTensorFromVar(const framework::Variable& var) {
  if (var.IsType<framework::LoDTensor>()) {
    return &(var.Get<framework::LoDTensor>());
52 53
  } else if (var.IsType<pten::SelectedRows>()) {
    return &(var.Get<pten::SelectedRows>().value());
J
Jiabin Yang 已提交
54 55 56 57 58
  } else {
    return nullptr;
  }
}

59
template <typename VarType>
J
Jiabin Yang 已提交
60
void HandleComplexGradToRealGrad(const NameVarMap<VarType>& outs) {
61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76
  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 已提交
77
      if (tensor && tensor->IsInitialized()) {
78 79 80 81 82 83 84 85
        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 已提交
86 87 88 89 90
      }
    }
  }
}

J
Jiabin Yang 已提交
91
template <>
92 93
void HandleComplexGradToRealGrad<egr::EagerVariable>(
    const NameVarMap<egr::EagerVariable>& outs) {
J
Jiabin Yang 已提交
94 95 96
  // TODO(jiabin): Support Complex here.
}

97 98 99 100 101
void TestHandleComplexGradToRealGradEager(
    const NameVarMap<egr::EagerVariable>& outs) {
  HandleComplexGradToRealGrad<egr::EagerVariable>(outs);
}

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

113 114 115 116 117 118 119 120 121 122 123 124 125
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),
126
      pt_kernel_(pt_kernel) {}
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
                       const framework::AttributeMap& default_attrs) {
135
  platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
136
  auto* dev_ctx = pool.Get(place);
137

138 139 140 141 142 143 144 145
  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());
146 147 148 149
    mutable_op_attrs = default_attrs;
    for (auto& attr : attrs) {
      mutable_op_attrs[attr.first] = attr.second;
    }
150 151
  }
#endif
152 153 154
  // NOTE(zhiqiu): for kernels on given device, for example NPU, the order to
  // choose is:
  // pten npu kernel > fluid npu kernel > pten cpu kernel > fluid cpu kernel
J
Jiabin Yang 已提交
155

156
  // 1. get expected kernel key
157 158 159
  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);
160

161 162 163 164 165
  framework::KernelSignature pt_kernel_signature;
  pten::KernelKey pt_kernel_key;
  std::string pt_kernel_name;
  if (pten::KernelFactory::Instance().HasCompatiblePtenKernel(op.Type())) {
    pt_kernel_signature = op.GetExpectedPtenKernelArgs(dygraph_exe_ctx);
166
    VLOG(6) << pt_kernel_signature;
167

168 169
    pt_kernel_name = pt_kernel_signature.name;
    pt_kernel_key = TransOpKernelTypeToPtenKernelKey(expected_kernel_key);
170 171 172 173
    auto pt_kernel = pten::KernelFactory::Instance().SelectKernel(
        pt_kernel_name, pt_kernel_key);

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

W
Wilber 已提交
178 179 180 181 182
      if (platform::is_cpu_place(expected_kernel_key.place_)) {
        auto* cpu_ctx = pool.Get(paddle::platform::CPUPlace());
        return PreparedOp(op, ctx, expected_kernel_key, pt_kernel_signature,
                          pt_kernel, cpu_ctx);
      }
183 184
      // TODO(chenweihang): using CPUKernel when miss device kernel case
      return PreparedOp(op, ctx, expected_kernel_key, pt_kernel_signature,
185
                        pt_kernel, dev_ctx);
186
    } else {
C
Chen Weihang 已提交
187
      VLOG(6) << "Dynamic mode ChoosePtenKernel - kernel `" << pt_kernel_name
188 189 190 191
              << "` not found.";
    }
  }

192
  // 2. check if op[type] has kernel registered.
J
Jiabin Yang 已提交
193 194
  auto& all_op_kernels = op.AllOpKernels();
  auto kernels_iter = all_op_kernels.find(op.Type());
195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222

  if ((kernels_iter == all_op_kernels.end() ||
       kernels_iter->second.find(expected_kernel_key) ==
           kernels_iter->second.end())
#ifdef PADDLE_WITH_XPU
      ||
      paddle::platform::is_xpu_place(expected_kernel_key.place_) &&
          !paddle::platform::is_xpu_support_op(op.Type(),
                                               expected_kernel_key) ||
      paddle::platform::is_in_xpu_black_list(op.Type())
#endif
          ) {
    if (pten::KernelFactory::Instance().HasCompatiblePtenKernel(op.Type())) {
      auto pt_cpu_kernel_key =
          FallBackToCpu(expected_kernel_key, pt_kernel_key, op);
      auto pt_cpu_kernel = pten::KernelFactory::Instance().SelectKernel(
          pt_kernel_name, pt_cpu_kernel_key);
      if (pt_cpu_kernel.IsValid()) {
        VLOG(6) << "Dynamic mode PrepareImpl - kernel name: " << pt_kernel_name
                << " | kernel key: " << pt_cpu_kernel_key
                << " | kernel: " << pt_cpu_kernel;
        auto* cpu_ctx = pool.Get(paddle::platform::CPUPlace());
        return PreparedOp(op, ctx, expected_kernel_key, pt_kernel_signature,
                          pt_cpu_kernel, cpu_ctx);
      }
    }
  }

223 224 225 226 227
  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 已提交
228 229
  auto& kernels = kernels_iter->second;
  auto kernel_iter = kernels.find(expected_kernel_key);
230

231
#ifdef PADDLE_WITH_XPU
232
  if (paddle::platform::is_xpu_place(expected_kernel_key.place_) &&
Q
QingshuChen 已提交
233 234 235
      (kernel_iter == kernels.end() ||
       !paddle::platform::is_xpu_support_op(op.Type(), expected_kernel_key) ||
       paddle::platform::is_in_xpu_black_list(op.Type()))) {
236 237 238
    VLOG(3) << "missing XPU kernel: " << op.Type()
            << ", expected_kernel_key:" << expected_kernel_key
            << ", fallbacking to CPU one!";
239 240 241
    expected_kernel_key.place_ = platform::CPUPlace();
    kernel_iter = kernels.find(expected_kernel_key);
  }
L
Liu-xiandong 已提交
242

243
#endif
L
Liu-xiandong 已提交
244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266

#ifdef PADDLE_WITH_XPU_KP
  bool use_xpu_kp_kernel_rt =
      FLAGS_run_kp_kernel &&
      paddle::platform::is_xpu_kp_support_op(op.Type(), expected_kernel_key);
  bool use_xpu_kp_kernel_debug =
      paddle::platform::is_in_xpu_kpwhite_list(op.Type());
  if (use_xpu_kp_kernel_rt) {
    VLOG(3) << "xpu_kp using rt mode ";
  }
  if (use_xpu_kp_kernel_debug) {
    VLOG(3) << "xpu_kp using debug mode ";
  }
  if (paddle::platform::is_xpu_place(expected_kernel_key.place_) &&
      (use_xpu_kp_kernel_rt || use_xpu_kp_kernel_debug)) {
    expected_kernel_key.place_ = platform::XPUPlace();
    expected_kernel_key.library_type_ = paddle::framework::LibraryType::kKP;
    kernel_iter = kernels.find(expected_kernel_key);
    VLOG(3) << "using XPU KP kernel: " << op.Type()
            << ", using_kernel_key:" << expected_kernel_key;
  }
#endif

267 268
#ifdef PADDLE_WITH_ASCEND_CL
  if (kernel_iter == kernels.end() &&
269
      paddle::platform::is_npu_place(expected_kernel_key.place_)) {
270 271 272
    VLOG(3) << "missing NPU kernel: " << op.Type()
            << ", expected_kernel_key:" << expected_kernel_key
            << ", fallbacking to CPU one!";
273 274 275
    expected_kernel_key.place_ = platform::CPUPlace();
    kernel_iter = kernels.find(expected_kernel_key);
  }
276 277 278
#endif
#ifdef PADDLE_WITH_MLU
  if (kernel_iter == kernels.end() &&
279
      paddle::platform::is_mlu_place(expected_kernel_key.place_)) {
280 281 282 283 284 285
    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);
  }
286 287 288 289 290 291 292 293 294 295
#endif
#ifdef PADDLE_WITH_CUSTOM_DEVICE
  if (kernel_iter == kernels.end() &&
      paddle::platform::is_custom_place(expected_kernel_key.place_)) {
    VLOG(3) << "missing " << place.GetDeviceType() << " 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);
  }
296
#endif
297 298
  // TODO(jiabin): Add operator.cc's line 1000 part back when we need that
  // case
299 300 301 302
  PADDLE_ENFORCE_NE(kernel_iter, kernels.end(),
                    platform::errors::NotFound(
                        "Operator %s does not have kernel for %s.", op.Type(),
                        KernelTypeToString(expected_kernel_key)));
303

304 305 306 307
  if (!(expected_kernel_key.place_ == place)) {
    dev_ctx = pool.Get(expected_kernel_key.place_);
  }

308
  return PreparedOp(op, ctx, expected_kernel_key, kernel_iter->second, dev_ctx);
309 310
}

311 312 313 314
PreparedOp PreparedOp::Prepare(const NameVarMap<VarBase>& ins,
                               const NameVarMap<VarBase>& outs,
                               const framework::OperatorWithKernel& op,
                               const platform::Place& place,
315
                               const framework::AttributeMap& attrs,
316 317
                               const framework::AttributeMap& default_attrs) {
  return PrepareImpl<VarBase>(ins, outs, op, place, attrs, default_attrs);
318 319 320 321 322 323
}

PreparedOp PreparedOp::Prepare(const NameVarMap<VariableWrapper>& ins,
                               const NameVarMap<VariableWrapper>& outs,
                               const framework::OperatorWithKernel& op,
                               const platform::Place& place,
324
                               const framework::AttributeMap& attrs,
325
                               const framework::AttributeMap& default_attrs) {
326
  return PrepareImpl<VariableWrapper>(ins, outs, op, place, attrs,
327
                                      default_attrs);
328 329
}

330 331
PreparedOp PreparedOp::Prepare(const NameVarMap<egr::EagerVariable>& ins,
                               const NameVarMap<egr::EagerVariable>& outs,
J
Jiabin Yang 已提交
332 333 334 335
                               const framework::OperatorWithKernel& op,
                               const platform::Place& place,
                               const framework::AttributeMap& attrs,
                               const framework::AttributeMap& default_attrs) {
336 337
  return PrepareImpl<egr::EagerVariable>(ins, outs, op, place, attrs,
                                         default_attrs);
J
Jiabin Yang 已提交
338
}
339 340 341
template <typename VarType>
static void PreparedOpRunImpl(
    const framework::OperatorBase& op, const framework::RuntimeContext& ctx,
342
    const framework::OpKernelType& kernel_type,
343
    const framework::OperatorWithKernel::OpKernelFunc& func,
344
    platform::DeviceContext* dev_ctx, const NameVarMap<VarType>& ins,
345 346
    const NameVarMap<VarType>& outs, const framework::AttributeMap& attrs,
    const framework::AttributeMap& default_attrs) {
J
Jiabin Yang 已提交
347 348
  // TODO(zjl): remove scope in dygraph
  framework::Scope scope;
H
hong 已提交
349

350
  {
C
chenjian 已提交
351 352 353
    platform::RecordEvent record_event(op.Type() + "::infer_shape",
                                       platform::TracerEventType::OperatorInner,
                                       1, platform::EventRole::kInnerOp);
354 355 356 357 358 359
    DygraphInferShapeContext<VarType> infer_shape_ctx(
        &ins, &outs, &attrs, &default_attrs, op.Type(), &kernel_type);
    op.Info().infer_shape_(&infer_shape_ctx);
  }

  {
C
chenjian 已提交
360 361 362
    platform::RecordEvent record_event(op.Type() + "::compute",
                                       platform::TracerEventType::OperatorInner,
                                       1, platform::EventRole::kInnerOp);
H
hong 已提交
363

364 365 366
    func(DygraphExecutionContext<VarType>(op, scope, *dev_ctx, ctx, ins, outs,
                                          attrs, default_attrs));
  }
367

368 369 370 371 372
  if (FLAGS_check_nan_inf) {
    framework::details::CheckOpHasNanOrInfInDygraph<VarType>(
        op.Type(), outs, dev_ctx->GetPlace());
  }

L
Leo Chen 已提交
373 374 375 376 377 378 379 380
  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
  }

381 382 383 384 385 386 387 388 389 390 391 392 393 394 395
  /**
   * [ 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);
  }
396
}
H
hong 已提交
397

398 399 400
template <typename VarType>
static void PreparedOpRunPtImpl(
    const framework::OperatorBase& op,
401
    const framework::OpKernelType& kernel_type,
402
    const framework::KernelSignature& pt_kernel_signature,
403 404 405
    const pten::Kernel& pt_kernel, platform::DeviceContext* dev_ctx,
    const NameVarMap<VarType>& ins, const NameVarMap<VarType>& outs,
    const framework::AttributeMap& attrs,
406
    const framework::AttributeMap& default_attrs) {
407
  {
C
chenjian 已提交
408 409 410
    platform::RecordEvent record_event(op.Type() + "::infer_shape",
                                       platform::TracerEventType::OperatorInner,
                                       1, platform::EventRole::kInnerOp);
411 412 413 414 415 416
    DygraphInferShapeContext<VarType> infer_shape_ctx(
        &ins, &outs, &attrs, &default_attrs, op.Type(), &kernel_type);
    op.Info().infer_shape_(&infer_shape_ctx);
  }

  {
C
chenjian 已提交
417 418 419
    platform::RecordEvent record_event(op.Type() + "::compute",
                                       platform::TracerEventType::OperatorInner,
                                       1, platform::EventRole::kInnerOp);
420

421
    PreparePtenData<VarType>(pt_kernel, pt_kernel_signature, ins);
422

423 424 425 426
    pten::KernelContext pt_kernel_context;
    BuildDygraphPtenKernelContext<VarType>(pt_kernel_signature, pt_kernel, ins,
                                           outs, attrs, default_attrs, dev_ctx,
                                           &pt_kernel_context);
427

428 429
    pt_kernel(&pt_kernel_context);
  }
430

431 432
  if (FLAGS_benchmark) {
    dev_ctx->Wait();
433 434
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
    PADDLE_ENFORCE_GPU_SUCCESS(platform::GpuGetLastError());
435 436 437 438
    VLOG(4) << "Operator(" << op.Type() << "): context wait and get last error";
#endif
  }

439
  // TODO(chenweihang): add debug flags later
440 441 442
  if (framework::IsComplexType(kernel_type.data_type_)) {
    HandleComplexGradToRealGrad<VarType>(outs);
  }
443 444
}

445 446
void PreparedOp::Run(const NameVarMap<VarBase>& ins,
                     const NameVarMap<VarBase>& outs,
447 448
                     const framework::AttributeMap& attrs,
                     const framework::AttributeMap& default_attrs) {
449
  if (run_pten_kernel_) {
450
    PreparedOpRunPtImpl<VarBase>(op_, kernel_type_, pt_kernel_signature_,
451 452
                                 pt_kernel_, dev_ctx_, ins, outs, attrs,
                                 default_attrs);
453 454 455 456
  } else {
    PreparedOpRunImpl<VarBase>(op_, ctx_, kernel_type_, func_, dev_ctx_, ins,
                               outs, attrs, default_attrs);
  }
457
}
H
hong 已提交
458

459 460
void PreparedOp::Run(const NameVarMap<VariableWrapper>& ins,
                     const NameVarMap<VariableWrapper>& outs,
461 462
                     const framework::AttributeMap& attrs,
                     const framework::AttributeMap& default_attrs) {
463
  if (run_pten_kernel_) {
464
    PreparedOpRunPtImpl<VariableWrapper>(
465 466
        op_, kernel_type_, pt_kernel_signature_, pt_kernel_, dev_ctx_, ins,
        outs, attrs, default_attrs);
467 468 469 470
  } else {
    PreparedOpRunImpl<VariableWrapper>(op_, ctx_, kernel_type_, func_, dev_ctx_,
                                       ins, outs, attrs, default_attrs);
  }
J
Jiabin Yang 已提交
471 472
}

473 474
void PreparedOp::Run(const NameVarMap<egr::EagerVariable>& ins,
                     const NameVarMap<egr::EagerVariable>& outs,
J
Jiabin Yang 已提交
475 476 477
                     const framework::AttributeMap& attrs,
                     const framework::AttributeMap& default_attrs) {
  if (run_pten_kernel_) {
478
    PreparedOpRunPtImpl<egr::EagerVariable>(
J
Jiabin Yang 已提交
479 480 481
        op_, kernel_type_, pt_kernel_signature_, pt_kernel_, dev_ctx_, ins,
        outs, attrs, default_attrs);
  } else {
482 483 484
    PreparedOpRunImpl<egr::EagerVariable>(op_, ctx_, kernel_type_, func_,
                                          dev_ctx_, ins, outs, attrs,
                                          default_attrs);
J
Jiabin Yang 已提交
485 486 487
  }
}

J
Jiabin Yang 已提交
488 489
}  // namespace imperative
}  // namespace paddle