operator.h 24.8 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
Q
Qiao Longfei 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14 15 16

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. */

#pragma once

D
dongzhihong 已提交
17
#include <algorithm>
18
#include <atomic>
L
luotao1 已提交
19
#include <memory>
20
#include <mutex>  // NOLINT
Q
Qiao Longfei 已提交
21
#include <string>
D
dzhwinter 已提交
22
#include <tuple>
Q
Qiao Longfei 已提交
23
#include <unordered_map>
L
luotao1 已提交
24
#include <utility>
Q
Qiao Longfei 已提交
25 26
#include <vector>

Y
Yu Yang 已提交
27
#include "glog/logging.h"  // For VLOG
Y
Yi Wang 已提交
28 29
#include "paddle/fluid/framework/attribute.h"
#include "paddle/fluid/framework/block_desc.h"
30
#include "paddle/fluid/framework/convert_utils.h"
Y
Yi Wang 已提交
31 32 33
#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/framework/op_info.h"
#include "paddle/fluid/framework/op_kernel_type.h"
34
#include "paddle/fluid/framework/phi_utils.h"
Y
Yi Wang 已提交
35
#include "paddle/fluid/framework/scope.h"
36
#include "paddle/fluid/framework/selected_rows_utils.h"
Y
Yi Wang 已提交
37
#include "paddle/fluid/framework/tensor.h"
38
#include "paddle/fluid/framework/unused_var_check.h"
39
#include "paddle/fluid/memory/malloc.h"
Y
Yi Wang 已提交
40 41
#include "paddle/fluid/platform/device_context.h"
#include "paddle/fluid/platform/variant.h"
42 43 44
#include "paddle/phi/core/compat/arg_map_context.h"
#include "paddle/phi/core/compat/op_utils.h"
#include "paddle/phi/core/kernel_factory.h"
45
#include "paddle/utils/flat_hash_map.h"
46

W
wanghuancoder 已提交
47 48 49 50 51 52 53 54 55
namespace paddle {
namespace framework {
class InferShapeContext;
class OpInfo;
class Scope;
class Variable;
}  // namespace framework
}  // namespace paddle

56 57 58 59
namespace phi {
class KernelContext;
}

Q
Qiao Longfei 已提交
60 61
DECLARE_int32(inner_op_parallelism);

Q
Qiao Longfei 已提交
62 63 64
namespace paddle {
namespace framework {

65
/// If a variable is a empty variable, that name will be used.
66
constexpr char kEmptyVarName[] = "@EMPTY@";
67 68 69

/// If a variable is a temporary variable, that name will be set in Python,
/// but it will be convert to a unique name in scope after OpCreator.
70
constexpr char kTempVarName[] = "@TEMP@";
71 72

/// If a variable's name has a certain suffix, it means that the
T
tianshuo78520a 已提交
73 74
/// variable is the gradient of another variable.
/// e.g. Variable "x@GRAD" is the gradient of variable "x".
75
constexpr char kGradVarSuffix[] = "@GRAD";
76

M
minqiyang 已提交
77 78
constexpr size_t kGradVarSuffixSize = 5U;

79
/// Variables with this suffix are supposed to be filled up with zeros.
80
constexpr char kZeroVarSuffix[] = "@ZERO";
81

C
chengduo 已提交
82 83 84
/// Variables with this suffix are the new Gradient.
constexpr char kNewGradSuffix[] = "@NEWGRAD@";

L
luotao1 已提交
85 86 87 88 89 90 91 92
/// RuntimeContext is used to relate input/output names of Operator with
/// the corresponding variables in name scope.
/// If an Op has attribute kEnableCacheRuntimeContext, it means that in a same
/// name scope, since the input/output names of this Op do not change in the
/// execution, RuntimeContext could be created only at the first iteration of
/// this Op's execution to save the elapsed time.
constexpr char kEnableCacheRuntimeContext[] = "@ENABLE_CACHE_RUNTIME_CONTEXT@";

L
luotao1 已提交
93 94 95 96 97 98 99 100 101
/// If an Op has this attribute, all its kernels should calculate output
/// variable's shape in the corresponding Compute() function. And
/// OperatorWithKernel::RunImpl() would skip call this Op's InferShape()
/// function in its runtime for speedup.
/// TODO(luotao): Note that this temporal attribute would be deleted after all
/// ops contain it.
constexpr char kAllKernelsMustComputeRuntimeShape[] =
    "@ALL_KERNELS_MUST_COMPUTE_RUNTIME_SHAPE@";

D
dzhwinter 已提交
102
// define some kernel priority
103
/* Define multiple kernel type fallback order*/
D
dzhwinter 已提交
104 105
extern std::vector<std::tuple<platform::Place, LibraryType>> kKernelPriority;

106
inline std::string GradVarName(const std::string& var_name) {
M
minqiyang 已提交
107 108 109 110 111
  std::string result;
  result.reserve(var_name.size() + kGradVarSuffixSize);
  result += var_name;
  result += kGradVarSuffix;
  return result;
112 113
}

M
minqiyang 已提交
114
inline std::string GradOriginalVarName(const std::string& grad_var_name) {
M
minqiyang 已提交
115
  std::size_t pos = grad_var_name.rfind(kGradVarSuffix);
116 117 118 119 120
  if (pos == std::string::npos) {
    return grad_var_name;
  } else {
    return grad_var_name.substr(0, pos);
  }
121 122
}

123
inline bool VarIsTensor(const Variable& var) {
124
  return var.IsType<LoDTensor>() || var.IsType<phi::SelectedRows>();
125 126
}

C
chengduo 已提交
127 128
const Tensor* GetLoDTensorOrSelectedRowsValueFromVar(const Variable& var);
Tensor* GetMutableLoDTensorOrSelectedRowsValueFromVar(Variable* var);
Q
qiaolongfei 已提交
129

130
class ExecutionContext;
W
wanghuancoder 已提交
131
class OperatorBase;
132

X
Xin Pan 已提交
133 134
class RuntimeContext {
 public:
X
Xin Pan 已提交
135
  RuntimeContext(const VariableNameMap& innames,
136 137
                 const VariableNameMap& outnames,
                 const Scope& scope);
X
Xin Pan 已提交
138

X
Xin Pan 已提交
139 140 141 142
  RuntimeContext(const VariableValueMap& invars,
                 const VariableValueMap& outvars)
      : inputs(invars), outputs(outvars) {}

X
Xin Pan 已提交
143 144 145 146
  VariableValueMap inputs;
  VariableValueMap outputs;
};

Q
Qiao Longfei 已提交
147
/**
X
Xin Pan 已提交
148
 * OperatorBase has the basic elements that Net will call to do computation.
Q
Qiao Longfei 已提交
149 150 151 152 153 154
 * Only CreateOperator from OpRegistry will new Operator directly. User
 * should always construct a proto message OpDesc and call
 * OpRegistry::CreateOp(op_desc) to get an Operator instance.
 */
class OperatorBase {
 public:
155 156 157 158
  OperatorBase(const std::string& type,
               const VariableNameMap& inputs,
               const VariableNameMap& outputs,
               const AttributeMap& attrs);
159

Q
Qiao Longfei 已提交
160 161
  virtual ~OperatorBase() {}

162
  /// Executor will call this interface function to Run an op.
163 164
  //  The implementation should be written at RunImpl
  void Run(const Scope& scope, const platform::Place& place);
Y
Yu Yang 已提交
165

T
typhoonzero 已提交
166 167 168
  // FIXME(typhoonzero): this is only used for recv_op to stop event_loop.
  virtual void Stop() {}

169
  /// if scope is not null, also show dimensions of arguments
170
  virtual std::string DebugStringEx(const ScopeBase* scope) const;
171
  std::string DebugString() const { return DebugStringEx(nullptr); }
Y
Yu Yang 已提交
172

173
  virtual bool SupportGPU() const { return false; }
B
Baibaifan 已提交
174
  virtual bool SupportNPU() const { return false; }
F
fwenguang 已提交
175
  virtual bool SupportMLU() const { return false; }
176

177 178
  const std::string& Type() const { return type_; }

M
Michal Gallus 已提交
179
  bool HasAttr(const std::string& name) const { return attrs_.count(name); }
180 181
  template <typename T>
  inline const T& Attr(const std::string& name) const {
182
    PADDLE_ENFORCE_NE(
183 184
        attrs_.find(name),
        attrs_.end(),
185
        platform::errors::NotFound("(%s) is not found in AttributeMap.", name));
186
    return BOOST_GET_CONST(T, attrs_.at(name));
187
  }
188 189
  void SetAttr(const std::string& name, const Attribute& v) {
    PADDLE_ENFORCE_EQ(
190 191
        HasAttr(name),
        true,
192 193 194 195 196
        platform::errors::NotFound(
            "The attribute %s is not found in operator %s", name, Type()));

    attrs_[name] = v;
  }
197
  const AttributeMap& Attrs() const { return attrs_; }
D
dongzhihong 已提交
198

Y
Yu Yang 已提交
199 200
  const VariableNameMap& Inputs() const { return inputs_; }
  const VariableNameMap& Outputs() const { return outputs_; }
201 202
  VariableNameMap& Inputs() { return inputs_; }
  VariableNameMap& Outputs() { return outputs_; }
203

S
sneaxiy 已提交
204
  const OpInfo& Info() const {
205
    PADDLE_ENFORCE_NOT_NULL(
206 207 208
        info_,
        platform::errors::NotFound("OpInfo of operator (%s) is not found.",
                                   type_));
S
sneaxiy 已提交
209 210 211
    return *info_;
  }

212
  bool HasInputs(const std::string& name) const;
Y
Yu Yang 已提交
213
  //! Get a input with argument's name described in `op_proto`
214
  std::string Input(const std::string& name) const;
Y
Yu Yang 已提交
215
  //! Get a input which has multiple variables.
Y
Yu Yang 已提交
216
  const std::vector<std::string>& Inputs(const std::string& name) const;
217
  //! Get all inputs variable names
Q
qijun 已提交
218 219
  std::vector<std::string> InputVars() const;

220
  bool HasOutputs(const std::string& name) const;
Y
Yu Yang 已提交
221
  //! Get a output with argument's name described in `op_proto`
222
  std::string Output(const std::string& name) const;
Y
Yu Yang 已提交
223 224
  //! Get an output which has multiple variables.
  //! TODO add a vector_view to prevent memory copy.
Y
Yu Yang 已提交
225
  const std::vector<std::string>& Outputs(const std::string& name) const;
226
  //! Get all outputs variable names
Y
Yu Yang 已提交
227
  virtual std::vector<std::string> OutputVars(bool has_intermediate) const;
228

229
  void SetIsCalledByExecutor(bool x) { run_by_executor_ = x; }
230

B
baojun-nervana 已提交
231
  virtual void RuntimeInferShape(const Scope& scope,
X
Xin Pan 已提交
232 233
                                 const platform::Place& place,
                                 const RuntimeContext& ctx) const {}
234

Z
Zhang Ting 已提交
235 236 237 238 239
  virtual platform::Place GetExecutionPlace(
      const platform::Place& place) const {
    return place;
  }

Q
qiaolongfei 已提交
240
 protected:
Q
Qiao Longfei 已提交
241
  std::string type_;
D
dongzhihong 已提交
242
  // NOTE: in case of OpGrad, inputs_ contains:
243
  // I (Inputs)
D
dongzhihong 已提交
244 245
  // O (Outputs)
  // OG (Output Gradients)
Y
Yu Yang 已提交
246
  VariableNameMap inputs_;
Y
Yu Yang 已提交
247

D
dongzhihong 已提交
248 249
  // NOTE: in case of OpGrad, outputs_ contains
  // IG (Inputs Gradients)
Y
Yu Yang 已提交
250
  VariableNameMap outputs_;
Q
Qiao Longfei 已提交
251
  AttributeMap attrs_;
S
sneaxiy 已提交
252 253 254 255

  // OpInfo
  const OpInfo* info_;

256 257
  // Whether this operator executes in an Executor.
  bool run_by_executor_{true};
258 259 260 261

 private:
  void GenerateTemporaryNames();
  void CheckAllInputOutputSet() const;
262 263
  virtual void RunImpl(const Scope& scope,
                       const platform::Place& place) const = 0;
Y
Yan Chunwei 已提交
264 265
};

266
class ExecutionContext {
Y
Yan Chunwei 已提交
267
 public:
268 269
  ExecutionContext(const OperatorBase& op,
                   const Scope& scope,
X
Xin Pan 已提交
270
                   const platform::DeviceContext& device_context,
271 272
                   const RuntimeContext& ctx)
      : op_(op), scope_(scope), device_context_(device_context), ctx_(ctx) {}
H
hong 已提交
273
  virtual ~ExecutionContext() {}
274

H
hong 已提交
275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294
  virtual std::string InputName(const std::string& name) const {
    return op_.Input(name);
  }
  virtual std::vector<std::string> InputNames(const std::string& name) const {
    return op_.Inputs(name);
  }
  virtual std::string OutputName(const std::string& name) const {
    return op_.Output(name);
  }

  virtual std::vector<std::string> OutputNames(const std::string& name) const {
    return op_.Outputs(name);
  }

  virtual bool HasAttr(const std::string& name) const {
    return op_.HasAttr(name);
  }
  virtual const AttributeMap& Attrs() const { return op_.Attrs(); }

  const std::string& Type() const { return op_.Type(); }
Q
qiaolongfei 已提交
295 296 297

  const Scope& scope() const { return scope_; }

Q
qiaolongfei 已提交
298
  template <typename T>
Y
Yu Yang 已提交
299
  inline const T& Attr(const std::string& name) const {
300
    return BOOST_GET_CONST(T, GetAttr(name));
Q
qiaolongfei 已提交
301 302
  }

H
hong 已提交
303 304 305
  virtual const Attribute& GetAttr(const std::string& name) const {
    return op_.Attrs().at(name);
  }
306

H
hong 已提交
307
  virtual bool HasInput(const std::string& name) const;
308

309 310
  virtual bool HasInputs(const std::string& name) const;

H
hong 已提交
311
  virtual bool HasOutput(const std::string& name) const;
312

H
hong 已提交
313
  virtual size_t InputSize(const std::string& name) const {
Y
Yu Yang 已提交
314
    return op_.Inputs(name).size();
Y
Yan Chunwei 已提交
315 316
  }

H
hong 已提交
317
  virtual size_t OutputSize(const std::string& name) const {
Y
Yu Yang 已提交
318
    return op_.Outputs(name).size();
Y
Yan Chunwei 已提交
319 320
  }

H
hong 已提交
321
  virtual const Variable* InputVar(const std::string& name) const;
Y
Yan Chunwei 已提交
322

H
hong 已提交
323
  virtual Variable* OutputVar(const std::string& name) const;
Y
Yan Chunwei 已提交
324

H
hong 已提交
325
  virtual const std::vector<Variable*> MultiInputVar(
326
      const std::string& name) const {
327 328
    LogVarUsageIfUnusedVarCheckEnabled(name);

X
Xin Pan 已提交
329 330 331 332
    auto it = ctx_.inputs.find(name);
    if (it == ctx_.inputs.end()) {
      return {};
    }
G
Gabor Buella 已提交
333
    return {it->second.begin(), it->second.end()};
X
Xin Pan 已提交
334 335
  }

H
hong 已提交
336
  virtual std::vector<Variable*> MultiOutputVar(const std::string& name) const {
X
Xin Pan 已提交
337 338 339 340 341 342 343
    auto it = ctx_.outputs.find(name);
    if (it == ctx_.outputs.end()) {
      return {};
    }
    return it->second;
  }

C
Chen Weihang 已提交
344 345
  virtual paddle::small_vector<const std::string*> InNameList() const {
    paddle::small_vector<const std::string*> vec_temp;
H
hong 已提交
346 347 348
    vec_temp.reserve(ctx_.inputs.size());

    for (auto& input : ctx_.inputs) {
349
      vec_temp.push_back(&input.first);
H
hong 已提交
350 351 352 353 354
    }

    return vec_temp;
  }

355 356
  template <typename T>
  const T* Input(const std::string& name) const {
Y
Yu Yang 已提交
357
    auto* var = InputVar(name);
358
    return var == nullptr ? nullptr : &var->Get<T>();
359 360 361 362
  }

  template <typename T>
  T* Output(const std::string& name) const {
363
    auto var = OutputVar(name);
364
    return var == nullptr ? nullptr : var->GetMutable<T>();
365 366 367 368
  }

  template <typename T>
  const std::vector<const T*> MultiInput(const std::string& name) const {
369 370
    LogVarUsageIfUnusedVarCheckEnabled(name);

H
hong 已提交
371 372
    auto vars = MultiInputVar(name);
    if (vars.size() == 0) {
X
Xin Pan 已提交
373 374 375 376
      return {};
    }
    std::vector<const T*> res;
    res.reserve(vars.size());
377 378 379
    std::transform(vars.begin(),
                   vars.end(),
                   std::back_inserter(res),
H
hong 已提交
380
                   [&](const Variable* var) -> const T* {
X
Xin Pan 已提交
381 382 383 384 385 386 387
                     return var == nullptr ? nullptr : &var->Get<T>();
                   });
    return res;
  }

  template <typename T>
  std::vector<T*> MultiOutput(const std::string& name) const {
H
hong 已提交
388 389 390
    auto vars = MultiOutputVar(name);

    if (vars.size() == 0) {
X
Xin Pan 已提交
391 392
      return {};
    }
H
hong 已提交
393

X
Xin Pan 已提交
394 395
    std::vector<T*> res;
    res.reserve(vars.size());
396 397 398
    std::transform(vars.begin(),
                   vars.end(),
                   std::back_inserter(res),
X
Xin Pan 已提交
399 400 401
                   [&](Variable* var) -> T* {
                     return var == nullptr ? nullptr : var->GetMutable<T>();
                   });
H
hong 已提交
402

X
Xin Pan 已提交
403 404 405
    return res;
  }

406
  platform::Place GetPlace() const { return device_context_.GetPlace(); }
Q
qijun 已提交
407

Q
QI JUN 已提交
408 409 410 411 412
  template <typename DeviceContextType>
  const DeviceContextType& device_context() const {
    return *reinterpret_cast<const DeviceContextType*>(&device_context_);
  }

413
  const platform::DeviceContext& device_context() const {
Q
qijun 已提交
414
    return device_context_;
Q
qijun 已提交
415
  }
Q
qijun 已提交
416

417
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
Q
QI JUN 已提交
418
  const inline platform::CUDADeviceContext& cuda_device_context() const {
419 420
    PADDLE_ENFORCE_EQ(platform::is_gpu_place(device_context_.GetPlace()),
                      true,
421 422
                      platform::errors::PreconditionNotMet(
                          "Current device context place is not GPUPlace."));
Q
QI JUN 已提交
423 424 425 426 427
    return *reinterpret_cast<const platform::CUDADeviceContext*>(
        &device_context_);
  }
#endif

X
Xin Pan 已提交
428 429 430
  template <typename T, typename DevContext>
  Tensor AllocateTmpTensor(const framework::DDim& dim,
                           const DevContext& dev_ctx) const {
431
    auto tmp_allocation_ptr = memory::Alloc(dev_ctx, product(dim) * sizeof(T));
X
Xin Pan 已提交
432 433
    auto& deleter = tmp_allocation_ptr.get_deleter();
    auto* allocation_ptr = tmp_allocation_ptr.release();
434
    auto shared_allocation =
435
        std::shared_ptr<phi::Allocation>(allocation_ptr, deleter);
X
Xin Pan 已提交
436

437
    PADDLE_ENFORCE_GE(
438 439
        allocation_ptr->size(),
        phi::product(dim) * sizeof(T),
440 441 442
        platform::errors::PreconditionNotMet(
            "The data memory size(%d) is less than the tensor needed memory "
            "size(%d).",
443 444
            allocation_ptr->size(),
            phi::product(dim) * sizeof(T)));
X
Xin Pan 已提交
445

446
    paddle::framework::Tensor temp_tensor(framework::TransToPhiDataType(
447
        framework::ToDataType(std::type_index(typeid(T)))));
X
Xin Pan 已提交
448 449 450 451 452
    temp_tensor.Resize(dim);
    temp_tensor.ResetHolder(std::move(shared_allocation));
    return temp_tensor;
  }

H
hong 已提交
453 454 455
  const RuntimeContext Context() const { return ctx_; }

  std::string DebugString() const { return op_.DebugString(); }
456
  const OperatorBase& GetOp() const { return op_; }
H
hong 已提交
457

458
 private:
459 460
  const OperatorBase& op_;
  const Scope& scope_;
461
  const platform::DeviceContext& device_context_;
X
Xin Pan 已提交
462
  const RuntimeContext& ctx_;
Q
Qiao Longfei 已提交
463 464
};

465
// TODO(chenweihang): split impl based OpProto or Dygraph if needed
466
class ExecutionArgumentMappingContext : public phi::ArgumentMappingContext {
467 468 469 470 471
 public:
  explicit ExecutionArgumentMappingContext(const ExecutionContext& ctx)
      : ctx_(ctx) {}

  bool HasInput(const std::string& name) const override {
472
    return ctx_.HasInputs(name);
473 474 475 476 477 478
  }

  bool HasOutput(const std::string& name) const override {
    return ctx_.HasOutput(name);
  }

479 480 481 482
  bool HasAttr(const std::string& name) const override {
    return ctx_.HasAttr(name);
  }

483 484 485
  paddle::any Attr(const std::string& name) const override {
    auto& attr = ctx_.GetAttr(name);
    return GetAttrValue(attr);
486 487 488
  }

  size_t InputSize(const std::string& name) const override {
489
    return ctx_.MultiInputVar(name).size();
490 491 492
  }

  size_t OutputSize(const std::string& name) const override {
493
    return ctx_.MultiOutputVar(name).size();
494 495 496
  }

  bool IsDenseTensorInput(const std::string& name) const override {
497 498 499 500 501
    const auto* var = ctx_.InputVar(name);
    return var->IsType<phi::DenseTensor>();
  }

  bool IsDenseTensorInputs(const std::string& name) const override {
502 503 504 505
    auto vars = ctx_.MultiInputVar(name);
    return std::all_of(vars.begin(), vars.end(), [](const Variable* var) {
      return var->IsType<phi::DenseTensor>();
    });
506 507 508
  }

  bool IsSelectedRowsInput(const std::string& name) const override {
509 510
    const auto* var = ctx_.InputVar(name);
    return var->IsType<phi::SelectedRows>();
511 512
  }

513
  bool IsDenseTensorVectorInput(const std::string& name) const override {
514 515 516 517
    auto vars = ctx_.MultiInputVar(name);
    return std::all_of(vars.begin(), vars.end(), [](const Variable* var) {
      return var->IsType<framework::LoDTensorArray>();
    });
518 519
  }

520
  bool IsDenseTensorOutput(const std::string& name) const override {
521 522 523 524
    auto vars = ctx_.MultiOutputVar(name);
    return std::all_of(vars.begin(), vars.end(), [](const Variable* var) {
      return var->IsType<phi::DenseTensor>();
    });
525 526 527
  }

  bool IsSelectedRowsOutput(const std::string& name) const override {
528 529 530 531
    auto vars = ctx_.MultiOutputVar(name);
    return std::all_of(vars.begin(), vars.end(), [](const Variable* var) {
      return var->IsType<phi::SelectedRows>();
    });
532 533
  }

534 535
  bool IsForInferShape() const override { return false; }

536 537 538 539
 private:
  const ExecutionContext& ctx_;
};

540 541 542 543 544 545 546 547
template <>
const std::vector<const Tensor*> ExecutionContext::MultiInput<Tensor>(
    const std::string& name) const;

template <>
std::vector<Tensor*> ExecutionContext::MultiOutput<Tensor>(
    const std::string& name) const;

Y
Yu Yang 已提交
548
class OpKernelBase {
Q
qijun 已提交
549
 public:
Q
qijun 已提交
550
  /**
551
   * ExecutionContext is the only parameter of Kernel Run function.
Q
qijun 已提交
552 553
   * Run will get input/output variables, state such as momentum and
   * device resource such as CUDA stream, cublas handle, etc. from
554
   * ExecutionContext. User should construct it before run the Operator.
Q
qijun 已提交
555 556
   */

557
  virtual void Compute(const ExecutionContext& context) const = 0;
Y
Yu Yang 已提交
558

Y
Yu Yang 已提交
559 560 561 562 563 564 565
  virtual ~OpKernelBase() = default;
};

template <typename T>
class OpKernel : public OpKernelBase {
 public:
  using ELEMENT_TYPE = T;
Y
Yu Yang 已提交
566 567
};

Y
Yu Yang 已提交
568 569
class OperatorWithKernel : public OperatorBase {
 public:
Y
yuyang18 已提交
570
  using OpKernelFunc = std::function<void(const ExecutionContext&)>;
Y
Yu Yang 已提交
571
  using OpKernelMap =
Y
yuyang18 已提交
572
      std::unordered_map<OpKernelType, OpKernelFunc, OpKernelType::Hash>;
Q
Qiao Longfei 已提交
573

574 575 576 577
  OperatorWithKernel(const std::string& type,
                     const VariableNameMap& inputs,
                     const VariableNameMap& outputs,
                     const AttributeMap& attrs)
Y
Yu Yang 已提交
578 579
      : OperatorBase(type, inputs, outputs, attrs) {}

C
chentianyu03 已提交
580
  static paddle::flat_hash_map<std::string /* op_type */, OpKernelMap>&
Y
Yu Yang 已提交
581
  AllOpKernels() {
C
chentianyu03 已提交
582
    static paddle::flat_hash_map<std::string, OpKernelMap> g_all_op_kernels;
Y
Yu Yang 已提交
583
    return g_all_op_kernels;
Y
Yu Yang 已提交
584
  }
Y
Yan Chunwei 已提交
585

586 587 588
  bool SupportGPU() const override;

  bool SupportNPU() const override;
589

F
fwenguang 已提交
590
  bool SupportMLU() const override {
591
    // TODO(zhiqiu): support phi if needed?
F
fwenguang 已提交
592
    auto& op_kernels = OperatorWithKernel::AllOpKernels().at(type_);
593 594
    return std::any_of(op_kernels.begin(),
                       op_kernels.end(),
F
fwenguang 已提交
595 596 597 598
                       [](OpKernelMap::const_reference kern_pair) {
                         return platform::is_mlu_place(kern_pair.first.place_);
                       });
  }
599
  bool SupportsMKLDNN(proto::VarType::Type data_type) const;
600

601 602
  bool SupportsKernelType(const OpKernelType& kernel_type) const;

603 604
  bool CanMKLDNNBeUsed(const framework::ExecutionContext& ctx,
                       proto::VarType::Type data_type) const;
605

606
  virtual void InferShape(InferShapeContext* ctx) const;
Y
Yu Yang 已提交
607

608 609
  void RuntimeInferShape(const Scope& scope,
                         const platform::Place& place,
X
Xin Pan 已提交
610
                         const RuntimeContext& ctx) const override;
B
baojun-nervana 已提交
611

612 613 614
  proto::VarType::Type IndicateVarDataType(const ExecutionContext& ctx,
                                           const std::string& name) const;

615
  proto::VarType::Type IndicateOrPromoteVarDataTypes(
616 617
      const ExecutionContext& ctx,
      const std::string& name1,
618 619
      const std::string& name2) const;

620
  virtual OpKernelType GetExpectedKernelType(const ExecutionContext& ctx) const;
X
Xin Pan 已提交
621

622 623
  // change this to public so that in dygraph mode we can call it to check if we
  // need transform data
624
  virtual OpKernelType GetKernelTypeForVar(
625 626
      const std::string& var_name,
      const Tensor& tensor,
627
      const OpKernelType& expected_kernel_type) const;
Y
Yu Yang 已提交
628

629 630
  platform::Place GetExecutionPlace(
      const platform::Place& platform) const override {
Z
Zhang Ting 已提交
631 632 633
    return kernel_type_->place_;
  }

634
  /* member functions for adapting to phi lib */
635
  /** In the Tensor calculation library, the new Kernel adopts a clearer and
636 637 638 639 640 641
   * more streamlined design. The arguments of the Kernel and the input and
   * output arguments registered in the original OpMaker do not match in some
   * cases, so we use map to record the arguments required by the kernel.
   * When selecting Kernel during Op execution, select the arguments of the
   * original Op according to the GetExpectedPhiKernelArgs returned arguments.
   */
642
  phi::KernelSignature GetExpectedPhiKernelArgs(
643 644
      const ExecutionContext& ctx) const;

645 646
  /* member functions for adapting to phi lib */
  phi::KernelKey ChoosePhiKernel(const ExecutionContext& ctx) const;
647

648
  void ChooseKernel(const ExecutionContext& ctx) const;
649
  /**
650
   * Transfer data place for phi kernel
651 652
   * Is this really needed?
   */
653 654
  Scope* PreparePhiData(const Scope& scope,
                        const phi::Kernel& pt_kernel,
655 656
                        const phi::KernelSignature& pt_kernel_signature,
                        RuntimeContext* ctx) const;
657

658 659 660
  void BuildPhiKernelContext(const RuntimeContext& ctx,
                             platform::DeviceContext* dev_ctx,
                             phi::KernelContext* pt_kernel_context) const;
661

662
  phi::KernelSignature* PhiKernelSignature() const {
663
    return kernel_signature_.get();
664 665
  }

666
  phi::Kernel* PhiKernel() const { return pt_kernel_.get(); }
667

668
  void ResetPhiKernel(phi::Kernel* kernel) const {
669 670 671
    return pt_kernel_.reset(kernel);
  }

672
  const OpKernelType* kernel_type() const { return kernel_type_.get(); }
673
  const OpKernelFunc* kernel_func() const { return kernel_func_.get(); }
674

675 676 677 678
  void ResetKernelType(OpKernelType* kernel_type) {
    kernel_type_.reset(kernel_type);
  }

Y
Yu Yang 已提交
679
 private:
680
  void RunImpl(const Scope& scope, const platform::Place& place) const final;
681 682
  void RunImpl(const Scope& scope,
               const platform::Place& place,
L
luotao1 已提交
683
               RuntimeContext* runtime_ctx) const;
Y
yuyang18 已提交
684 685

  /**
T
tianshuo78520a 已提交
686 687
   * Transfer data from scope to a transferred scope. If there is no data need
   * to
Y
yuyang18 已提交
688 689 690 691
   * be tranfered, it returns nullptr.
   *
   * * transfered_inplace_vars is a output vector.
   */
X
Xin Pan 已提交
692 693 694 695
  Scope* PrepareData(const Scope& scope,
                     const OpKernelType& expected_kernel_key,
                     std::vector<std::string>* transfered_inplace_vars,
                     RuntimeContext* ctx) const;
Y
yuyang18 已提交
696 697 698 699

  void TransferInplaceVarsBack(const Scope& scope,
                               const std::vector<std::string>& inplace_vars,
                               const Scope& exec_scope) const;
700

701 702
  OpKernelType InnerGetExpectedKernelType(const ExecutionContext& ctx) const;

703 704 705
  void HandleComplexGradToRealGrad(const Scope& scope,
                                   RuntimeContext* ctx) const;

706 707 708 709 710
  /* Inner assist methods */
  // indicate kernel DataType by input data.
  // By default all input data must be same.
  proto::VarType::Type IndicateDataType(const ExecutionContext& ctx) const;
  // used for IndicateDataType
711 712
  void ParseInputDataType(const Variable* vars,
                          const std::string& name,
713
                          proto::VarType::Type* data_type) const;
714 715 716
  void ParseMultiInputDataType(const std::vector<Variable*>& vars,
                               const std::string& name,
                               proto::VarType::Type* data_type) const;
717 718 719 720
  // used for IndicateOrPromoteVarDataTypes
  Tensor* GetTensorFormInputSafely(const ExecutionContext& ctx,
                                   const std::string& name) const;

721
 protected:
L
Liu Yiqun 已提交
722 723
  mutable std::unique_ptr<OpKernelType> kernel_type_;
  mutable std::unique_ptr<OpKernelFunc> kernel_func_;
L
luotao1 已提交
724 725
  mutable std::unique_ptr<RuntimeContext> runtime_ctx_;
  mutable const Scope* pre_scope_ = nullptr;
726
  mutable bool need_prepare_data_ = true;
727
  mutable bool need_prepare_phi_data_ = false;
728 729
  mutable bool enable_cache_runtime_context_ = false;
  mutable bool all_kernels_must_compute_runtime_shape_ = false;
730
  mutable std::mutex cache_update_mutex_;
731
  mutable bool enable_cache_transfer_scope_ = false;
732
  // NOTE(chenweihang): Similar op members are used to adapt to
733
  // new phi kernel, if there is a better design in the future,
734
  // we may polish the implementation here
735
  mutable bool run_phi_kernel_ = false;
L
Liu-xiandong 已提交
736
  mutable bool run_kp_kernel = false;
737
  mutable std::unique_ptr<phi::KernelSignature> kernel_signature_;
738
  mutable std::unique_ptr<phi::Kernel> pt_kernel_;
739
  mutable std::unique_ptr<phi::ArgumentMappingFn> arg_map_fn_;
740 741 742

  struct CacheImpl;
  mutable CacheImpl* impl_{nullptr};
Q
Qiao Longfei 已提交
743 744
};

Y
Yu Yang 已提交
745 746
extern bool OpSupportGPU(const std::string& op_type);

Q
Qiao Longfei 已提交
747 748
}  // namespace framework
}  // namespace paddle