operator.h 23.9 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 30 31 32
#include "paddle/fluid/framework/attribute.h"
#include "paddle/fluid/framework/block_desc.h"
#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/framework/op_info.h"
#include "paddle/fluid/framework/op_kernel_type.h"
33
#include "paddle/fluid/framework/phi_utils.h"
Y
Yi Wang 已提交
34
#include "paddle/fluid/framework/scope.h"
35
#include "paddle/fluid/framework/selected_rows_utils.h"
Y
Yi Wang 已提交
36
#include "paddle/fluid/framework/tensor.h"
37
#include "paddle/fluid/framework/unused_var_check.h"
38
#include "paddle/fluid/memory/malloc.h"
Y
Yi Wang 已提交
39 40
#include "paddle/fluid/platform/device_context.h"
#include "paddle/fluid/platform/variant.h"
C
chentianyu03 已提交
41
#include "paddle/utils/flat_hash_map.h"
Q
Qiao Longfei 已提交
42

43
#include "paddle/fluid/framework/convert_utils.h"
44 45 46
#include "paddle/phi/core/compat/arg_map_context.h"
#include "paddle/phi/core/compat/op_utils.h"
#include "paddle/phi/core/kernel_factory.h"
47

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

57 58 59 60
namespace phi {
class KernelContext;
}

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

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

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

/// 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.
71
constexpr char kTempVarName[] = "@TEMP@";
72 73

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

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

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

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

L
luotao1 已提交
86 87 88 89 90 91 92 93
/// 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 已提交
94 95 96 97 98 99 100 101 102
/// 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 已提交
103
// define some kernel priority
104
/* Define multiple kernel type fallback order*/
D
dzhwinter 已提交
105 106
extern std::vector<std::tuple<platform::Place, LibraryType>> kKernelPriority;

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

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

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

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

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

X
Xin Pan 已提交
134 135
class RuntimeContext {
 public:
X
Xin Pan 已提交
136 137
  RuntimeContext(const VariableNameMap& innames,
                 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:
Y
Yu Yang 已提交
155 156
  OperatorBase(const std::string& type, const VariableNameMap& inputs,
               const VariableNameMap& outputs, const AttributeMap& attrs);
157

Q
Qiao Longfei 已提交
158 159
  virtual ~OperatorBase() {}

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

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

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

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

175 176
  const std::string& Type() const { return type_; }

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

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

Y
Yu Yang 已提交
195 196
  const VariableNameMap& Inputs() const { return inputs_; }
  const VariableNameMap& Outputs() const { return outputs_; }
197

S
sneaxiy 已提交
198
  const OpInfo& Info() const {
199 200 201
    PADDLE_ENFORCE_NOT_NULL(
        info_, platform::errors::NotFound(
                   "OpInfo of operator (%s) is not found.", type_));
S
sneaxiy 已提交
202 203 204
    return *info_;
  }

205
  bool HasInputs(const std::string& name) const;
Y
Yu Yang 已提交
206
  //! Get a input with argument's name described in `op_proto`
207
  std::string Input(const std::string& name) const;
Y
Yu Yang 已提交
208
  //! Get a input which has multiple variables.
Y
Yu Yang 已提交
209
  const std::vector<std::string>& Inputs(const std::string& name) const;
210
  //! Get all inputs variable names
Q
qijun 已提交
211 212
  std::vector<std::string> InputVars() const;

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

222
  void SetIsCalledByExecutor(bool x) { run_by_executor_ = x; }
223

B
baojun-nervana 已提交
224
  virtual void RuntimeInferShape(const Scope& scope,
X
Xin Pan 已提交
225 226
                                 const platform::Place& place,
                                 const RuntimeContext& ctx) const {}
227

Z
Zhang Ting 已提交
228 229 230 231 232
  virtual platform::Place GetExecutionPlace(
      const platform::Place& place) const {
    return place;
  }

Q
qiaolongfei 已提交
233
 protected:
Q
Qiao Longfei 已提交
234
  std::string type_;
D
dongzhihong 已提交
235
  // NOTE: in case of OpGrad, inputs_ contains:
236
  // I (Inputs)
D
dongzhihong 已提交
237 238
  // O (Outputs)
  // OG (Output Gradients)
Y
Yu Yang 已提交
239
  VariableNameMap inputs_;
Y
Yu Yang 已提交
240

D
dongzhihong 已提交
241 242
  // NOTE: in case of OpGrad, outputs_ contains
  // IG (Inputs Gradients)
Y
Yu Yang 已提交
243
  VariableNameMap outputs_;
Q
Qiao Longfei 已提交
244
  AttributeMap attrs_;
S
sneaxiy 已提交
245 246 247 248

  // OpInfo
  const OpInfo* info_;

249 250
  // Whether this operator executes in an Executor.
  bool run_by_executor_{true};
251 252 253 254

 private:
  void GenerateTemporaryNames();
  void CheckAllInputOutputSet() const;
255 256
  virtual void RunImpl(const Scope& scope,
                       const platform::Place& place) const = 0;
Y
Yan Chunwei 已提交
257 258
};

259
class ExecutionContext {
Y
Yan Chunwei 已提交
260
 public:
261
  ExecutionContext(const OperatorBase& op, const Scope& scope,
X
Xin Pan 已提交
262
                   const platform::DeviceContext& device_context,
263 264
                   const RuntimeContext& ctx)
      : op_(op), scope_(scope), device_context_(device_context), ctx_(ctx) {}
H
hong 已提交
265
  virtual ~ExecutionContext() {}
266

H
hong 已提交
267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286
  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 已提交
287 288 289

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

Q
qiaolongfei 已提交
290
  template <typename T>
Y
Yu Yang 已提交
291
  inline const T& Attr(const std::string& name) const {
292
    return BOOST_GET_CONST(T, GetAttr(name));
Q
qiaolongfei 已提交
293 294
  }

H
hong 已提交
295 296 297
  virtual const Attribute& GetAttr(const std::string& name) const {
    return op_.Attrs().at(name);
  }
298

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

301 302
  virtual bool HasInputs(const std::string& name) const;

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

H
hong 已提交
305
  virtual size_t InputSize(const std::string& name) const {
Y
Yu Yang 已提交
306
    return op_.Inputs(name).size();
Y
Yan Chunwei 已提交
307 308
  }

H
hong 已提交
309
  virtual size_t OutputSize(const std::string& name) const {
Y
Yu Yang 已提交
310
    return op_.Outputs(name).size();
Y
Yan Chunwei 已提交
311 312
  }

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

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

H
hong 已提交
317
  virtual const std::vector<Variable*> MultiInputVar(
318
      const std::string& name) const {
319 320
    LogVarUsageIfUnusedVarCheckEnabled(name);

X
Xin Pan 已提交
321 322 323 324
    auto it = ctx_.inputs.find(name);
    if (it == ctx_.inputs.end()) {
      return {};
    }
G
Gabor Buella 已提交
325
    return {it->second.begin(), it->second.end()};
X
Xin Pan 已提交
326 327
  }

H
hong 已提交
328
  virtual std::vector<Variable*> MultiOutputVar(const std::string& name) const {
X
Xin Pan 已提交
329 330 331 332 333 334 335
    auto it = ctx_.outputs.find(name);
    if (it == ctx_.outputs.end()) {
      return {};
    }
    return it->second;
  }

H
hong 已提交
336 337 338 339 340 341 342 343 344 345 346
  virtual std::vector<std::string> InNameList() const {
    std::vector<std::string> vec_temp;
    vec_temp.reserve(ctx_.inputs.size());

    for (auto& input : ctx_.inputs) {
      vec_temp.push_back(input.first);
    }

    return vec_temp;
  }

347 348
  template <typename T>
  const T* Input(const std::string& name) const {
Y
Yu Yang 已提交
349
    auto* var = InputVar(name);
350
    return var == nullptr ? nullptr : &var->Get<T>();
351 352 353 354
  }

  template <typename T>
  T* Output(const std::string& name) const {
355
    auto var = OutputVar(name);
356
    return var == nullptr ? nullptr : var->GetMutable<T>();
357 358 359 360
  }

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

H
hong 已提交
363 364
    auto vars = MultiInputVar(name);
    if (vars.size() == 0) {
X
Xin Pan 已提交
365 366 367 368 369
      return {};
    }
    std::vector<const T*> res;
    res.reserve(vars.size());
    std::transform(vars.begin(), vars.end(), std::back_inserter(res),
H
hong 已提交
370
                   [&](const Variable* var) -> const T* {
X
Xin Pan 已提交
371 372 373 374 375 376 377
                     return var == nullptr ? nullptr : &var->Get<T>();
                   });
    return res;
  }

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

    if (vars.size() == 0) {
X
Xin Pan 已提交
381 382
      return {};
    }
H
hong 已提交
383

X
Xin Pan 已提交
384 385 386 387 388 389
    std::vector<T*> res;
    res.reserve(vars.size());
    std::transform(vars.begin(), vars.end(), std::back_inserter(res),
                   [&](Variable* var) -> T* {
                     return var == nullptr ? nullptr : var->GetMutable<T>();
                   });
H
hong 已提交
390

X
Xin Pan 已提交
391 392 393
    return res;
  }

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

Q
QI JUN 已提交
396 397 398 399 400
  template <typename DeviceContextType>
  const DeviceContextType& device_context() const {
    return *reinterpret_cast<const DeviceContextType*>(&device_context_);
  }

401
  const platform::DeviceContext& device_context() const {
Q
qijun 已提交
402
    return device_context_;
Q
qijun 已提交
403
  }
Q
qijun 已提交
404

405
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
Q
QI JUN 已提交
406
  const inline platform::CUDADeviceContext& cuda_device_context() const {
407 408 409
    PADDLE_ENFORCE_EQ(platform::is_gpu_place(device_context_.GetPlace()), true,
                      platform::errors::PreconditionNotMet(
                          "Current device context place is not GPUPlace."));
Q
QI JUN 已提交
410 411 412 413 414
    return *reinterpret_cast<const platform::CUDADeviceContext*>(
        &device_context_);
  }
#endif

X
Xin Pan 已提交
415 416 417
  template <typename T, typename DevContext>
  Tensor AllocateTmpTensor(const framework::DDim& dim,
                           const DevContext& dev_ctx) const {
418
    auto tmp_allocation_ptr = memory::Alloc(dev_ctx, product(dim) * sizeof(T));
X
Xin Pan 已提交
419 420
    auto& deleter = tmp_allocation_ptr.get_deleter();
    auto* allocation_ptr = tmp_allocation_ptr.release();
421
    auto shared_allocation =
422
        std::shared_ptr<phi::Allocation>(allocation_ptr, deleter);
X
Xin Pan 已提交
423

424
    PADDLE_ENFORCE_GE(
425
        allocation_ptr->size(), phi::product(dim) * sizeof(T),
426 427 428
        platform::errors::PreconditionNotMet(
            "The data memory size(%d) is less than the tensor needed memory "
            "size(%d).",
429
            allocation_ptr->size(), phi::product(dim) * sizeof(T)));
X
Xin Pan 已提交
430

431
    paddle::framework::Tensor temp_tensor(framework::TransToPhiDataType(
432
        framework::ToDataType(std::type_index(typeid(T)))));
X
Xin Pan 已提交
433 434 435 436 437
    temp_tensor.Resize(dim);
    temp_tensor.ResetHolder(std::move(shared_allocation));
    return temp_tensor;
  }

H
hong 已提交
438 439 440
  const RuntimeContext Context() const { return ctx_; }

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

443
 private:
444 445
  const OperatorBase& op_;
  const Scope& scope_;
446
  const platform::DeviceContext& device_context_;
X
Xin Pan 已提交
447
  const RuntimeContext& ctx_;
Q
Qiao Longfei 已提交
448 449
};

450
// TODO(chenweihang): split impl based OpProto or Dygraph if needed
451
class ExecutionArgumentMappingContext : public phi::ArgumentMappingContext {
452 453 454 455 456
 public:
  explicit ExecutionArgumentMappingContext(const ExecutionContext& ctx)
      : ctx_(ctx) {}

  bool HasInput(const std::string& name) const override {
457
    return ctx_.HasInputs(name);
458 459 460 461 462 463
  }

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

464 465 466 467
  bool HasAttr(const std::string& name) const override {
    return ctx_.HasAttr(name);
  }

468 469 470
  paddle::any Attr(const std::string& name) const override {
    auto& attr = ctx_.GetAttr(name);
    return GetAttrValue(attr);
471 472 473
  }

  size_t InputSize(const std::string& name) const override {
474
    return ctx_.MultiInputVar(name).size();
475 476 477
  }

  size_t OutputSize(const std::string& name) const override {
478
    return ctx_.MultiOutputVar(name).size();
479 480 481
  }

  bool IsDenseTensorInput(const std::string& name) const override {
482 483 484 485
    auto vars = ctx_.MultiInputVar(name);
    return std::all_of(vars.begin(), vars.end(), [](const Variable* var) {
      return var->IsType<phi::DenseTensor>();
    });
486 487 488
  }

  bool IsSelectedRowsInput(const std::string& name) const override {
489 490 491 492
    auto vars = ctx_.MultiInputVar(name);
    return std::all_of(vars.begin(), vars.end(), [](const Variable* var) {
      return var->IsType<phi::SelectedRows>();
    });
493 494
  }

495
  bool IsDenseTensorVectorInput(const std::string& name) const override {
496 497 498 499
    auto vars = ctx_.MultiInputVar(name);
    return std::all_of(vars.begin(), vars.end(), [](const Variable* var) {
      return var->IsType<framework::LoDTensorArray>();
    });
500 501
  }

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

  bool IsSelectedRowsOutput(const std::string& name) const override {
510 511 512 513
    auto vars = ctx_.MultiOutputVar(name);
    return std::all_of(vars.begin(), vars.end(), [](const Variable* var) {
      return var->IsType<phi::SelectedRows>();
    });
514 515
  }

516 517
  bool IsForInferShape() const override { return false; }

518 519 520 521
 private:
  const ExecutionContext& ctx_;
};

522 523 524 525 526 527 528 529
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 已提交
530
class OpKernelBase {
Q
qijun 已提交
531
 public:
Q
qijun 已提交
532
  /**
533
   * ExecutionContext is the only parameter of Kernel Run function.
Q
qijun 已提交
534 535
   * Run will get input/output variables, state such as momentum and
   * device resource such as CUDA stream, cublas handle, etc. from
536
   * ExecutionContext. User should construct it before run the Operator.
Q
qijun 已提交
537 538
   */

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

Y
Yu Yang 已提交
541 542 543 544 545 546 547
  virtual ~OpKernelBase() = default;
};

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

Y
Yu Yang 已提交
550 551
class OperatorWithKernel : public OperatorBase {
 public:
Y
yuyang18 已提交
552
  using OpKernelFunc = std::function<void(const ExecutionContext&)>;
Y
Yu Yang 已提交
553
  using OpKernelMap =
Y
yuyang18 已提交
554
      std::unordered_map<OpKernelType, OpKernelFunc, OpKernelType::Hash>;
Q
Qiao Longfei 已提交
555

Y
Yu Yang 已提交
556 557
  OperatorWithKernel(const std::string& type, const VariableNameMap& inputs,
                     const VariableNameMap& outputs, const AttributeMap& attrs)
Y
Yu Yang 已提交
558 559
      : OperatorBase(type, inputs, outputs, attrs) {}

C
chentianyu03 已提交
560
  static paddle::flat_hash_map<std::string /* op_type */, OpKernelMap>&
Y
Yu Yang 已提交
561
  AllOpKernels() {
C
chentianyu03 已提交
562
    static paddle::flat_hash_map<std::string, OpKernelMap> g_all_op_kernels;
Y
Yu Yang 已提交
563
    return g_all_op_kernels;
Y
Yu Yang 已提交
564
  }
Y
Yan Chunwei 已提交
565

566 567 568
  bool SupportGPU() const override;

  bool SupportNPU() const override;
569

F
fwenguang 已提交
570
  bool SupportMLU() const override {
571
    // TODO(zhiqiu): support phi if needed?
F
fwenguang 已提交
572 573 574 575 576 577
    auto& op_kernels = OperatorWithKernel::AllOpKernels().at(type_);
    return std::any_of(op_kernels.begin(), op_kernels.end(),
                       [](OpKernelMap::const_reference kern_pair) {
                         return platform::is_mlu_place(kern_pair.first.place_);
                       });
  }
578
  bool SupportsMKLDNN(proto::VarType::Type data_type) const;
579

580 581
  bool CanMKLDNNBeUsed(const framework::ExecutionContext& ctx,
                       proto::VarType::Type data_type) const;
582

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

X
Xin Pan 已提交
585 586
  void RuntimeInferShape(const Scope& scope, const platform::Place& place,
                         const RuntimeContext& ctx) const override;
B
baojun-nervana 已提交
587

588 589 590
  proto::VarType::Type IndicateVarDataType(const ExecutionContext& ctx,
                                           const std::string& name) const;

591 592 593 594
  proto::VarType::Type IndicateOrPromoteVarDataTypes(
      const ExecutionContext& ctx, const std::string& name1,
      const std::string& name2) const;

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

597 598
  // change this to public so that in dygraph mode we can call it to check if we
  // need transform data
599 600 601
  virtual OpKernelType GetKernelTypeForVar(
      const std::string& var_name, const Tensor& tensor,
      const OpKernelType& expected_kernel_type) const;
Y
Yu Yang 已提交
602

603 604
  platform::Place GetExecutionPlace(
      const platform::Place& platform) const override {
Z
Zhang Ting 已提交
605 606 607
    return kernel_type_->place_;
  }

608
  /* member functions for adapting to phi lib */
609 610 611 612 613
  /** In the Tensor calculation library, the new Kernel adopts a clearer and
    * 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
614
    * original Op according to the GetExpectedPhiKernelArgs returned arguments.
615
    */
616
  phi::KernelSignature GetExpectedPhiKernelArgs(
617 618
      const ExecutionContext& ctx) const;

619 620
  /* member functions for adapting to phi lib */
  phi::KernelKey ChoosePhiKernel(const ExecutionContext& ctx) const;
621

622
  /**
623
   * Transfer data place for phi kernel
624 625
   * Is this really needed?
   */
626 627 628
  Scope* PreparePhiData(const Scope& scope, const phi::Kernel& pt_kernel,
                        const phi::KernelSignature& pt_kernel_signature,
                        RuntimeContext* ctx) const;
629

630 631 632
  void BuildPhiKernelContext(const RuntimeContext& ctx,
                             platform::DeviceContext* dev_ctx,
                             phi::KernelContext* pt_kernel_context) const;
633

634
  phi::KernelSignature* PhiKernelSignature() const {
635 636 637
    return pt_kernel_signature_.get();
  }

638
  phi::Kernel* PhiKernel() const { return pt_kernel_.get(); }
639

640
  void ResetPhiKernel(phi::Kernel* kernel) const {
641 642 643
    return pt_kernel_.reset(kernel);
  }

644 645
  const OpKernelType* kernel_type() const { return kernel_type_.get(); }

646 647 648 649
  void ResetKernelType(OpKernelType* kernel_type) {
    kernel_type_.reset(kernel_type);
  }

Y
Yu Yang 已提交
650
 private:
651
  void RunImpl(const Scope& scope, const platform::Place& place) const final;
L
luotao1 已提交
652 653
  void RunImpl(const Scope& scope, const platform::Place& place,
               RuntimeContext* runtime_ctx) const;
Y
yuyang18 已提交
654 655

  /**
T
tianshuo78520a 已提交
656 657
   * Transfer data from scope to a transferred scope. If there is no data need
   * to
Y
yuyang18 已提交
658 659 660 661
   * be tranfered, it returns nullptr.
   *
   * * transfered_inplace_vars is a output vector.
   */
X
Xin Pan 已提交
662 663 664 665
  Scope* PrepareData(const Scope& scope,
                     const OpKernelType& expected_kernel_key,
                     std::vector<std::string>* transfered_inplace_vars,
                     RuntimeContext* ctx) const;
Y
yuyang18 已提交
666 667 668 669

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

671 672 673
  OpKernelType InnerGetExpectedKernelType(const ExecutionContext& ctx) const;

  void ChooseKernel(const ExecutionContext& ctx) const;
L
Liu Yiqun 已提交
674

675 676 677
  void HandleComplexGradToRealGrad(const Scope& scope,
                                   RuntimeContext* ctx) const;

678 679 680 681 682
  /* 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
683 684 685
  void ParseInputDataType(const std::vector<Variable*>& vars,
                          const std::string& name,
                          proto::VarType::Type* data_type) const;
686 687 688 689
  // used for IndicateOrPromoteVarDataTypes
  Tensor* GetTensorFormInputSafely(const ExecutionContext& ctx,
                                   const std::string& name) const;

690
 protected:
L
Liu Yiqun 已提交
691 692
  mutable std::unique_ptr<OpKernelType> kernel_type_;
  mutable std::unique_ptr<OpKernelFunc> kernel_func_;
L
luotao1 已提交
693 694
  mutable std::unique_ptr<RuntimeContext> runtime_ctx_;
  mutable const Scope* pre_scope_ = nullptr;
695
  mutable bool need_prepare_data_ = true;
696 697
  mutable bool enable_cache_runtime_context_ = false;
  mutable bool all_kernels_must_compute_runtime_shape_ = false;
698
  mutable std::mutex cache_update_mutex_;
699
  mutable bool enable_cache_transfer_scope_ = false;
700
  // NOTE(chenweihang): Similar op members are used to adapt to
701
  // new phi kernel, if there is a better design in the future,
702
  // we may polish the implementation here
703
  mutable bool run_phi_kernel_ = false;
L
Liu-xiandong 已提交
704
  mutable bool run_kp_kernel = false;
705 706
  mutable std::unique_ptr<phi::KernelSignature> pt_kernel_signature_;
  mutable std::unique_ptr<phi::Kernel> pt_kernel_;
707
  mutable std::unique_ptr<phi::ArgumentMappingFn> arg_map_fn_;
Q
Qiao Longfei 已提交
708 709
};

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

Q
Qiao Longfei 已提交
712 713
}  // namespace framework
}  // namespace paddle