operator.h 23.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 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/pten_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/pten/core/compat/arg_map_context.h"
44
#include "paddle/pten/core/compat/op_utils.h"
45 46
#include "paddle/pten/core/kernel_context.h"
#include "paddle/pten/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

Q
Qiao Longfei 已提交
57 58
DECLARE_int32(inner_op_parallelism);

Q
Qiao Longfei 已提交
59 60 61
namespace paddle {
namespace framework {

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

/// 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.
67
constexpr char kTempVarName[] = "@TEMP@";
68 69

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

M
minqiyang 已提交
74 75
constexpr size_t kGradVarSuffixSize = 5U;

76
/// Variables with this suffix are supposed to be filled up with zeros.
77
constexpr char kZeroVarSuffix[] = "@ZERO";
78

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

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

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

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

120
inline bool VarIsTensor(const Variable& var) {
121
  return var.IsType<LoDTensor>() || var.IsType<pten::SelectedRows>();
122 123
}

C
chengduo 已提交
124 125
const Tensor* GetLoDTensorOrSelectedRowsValueFromVar(const Variable& var);
Tensor* GetMutableLoDTensorOrSelectedRowsValueFromVar(Variable* var);
Q
qiaolongfei 已提交
126

127
class ExecutionContext;
W
wanghuancoder 已提交
128
class OperatorBase;
129

X
Xin Pan 已提交
130 131
class RuntimeContext {
 public:
X
Xin Pan 已提交
132 133
  RuntimeContext(const VariableNameMap& innames,
                 const VariableNameMap& outnames, const Scope& scope);
X
Xin Pan 已提交
134

X
Xin Pan 已提交
135 136 137 138
  RuntimeContext(const VariableValueMap& invars,
                 const VariableValueMap& outvars)
      : inputs(invars), outputs(outvars) {}

X
Xin Pan 已提交
139 140 141 142
  VariableValueMap inputs;
  VariableValueMap outputs;
};

Q
Qiao Longfei 已提交
143
/**
X
Xin Pan 已提交
144
 * OperatorBase has the basic elements that Net will call to do computation.
Q
Qiao Longfei 已提交
145 146 147 148 149 150
 * 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 已提交
151 152
  OperatorBase(const std::string& type, const VariableNameMap& inputs,
               const VariableNameMap& outputs, const AttributeMap& attrs);
153

Q
Qiao Longfei 已提交
154 155
  virtual ~OperatorBase() {}

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

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

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

167
  virtual bool SupportGPU() const { return false; }
B
Baibaifan 已提交
168
  virtual bool SupportNPU() const { return false; }
F
fwenguang 已提交
169
  virtual bool SupportMLU() const { return false; }
170

171 172
  const std::string& Type() const { return type_; }

M
Michal Gallus 已提交
173
  bool HasAttr(const std::string& name) const { return attrs_.count(name); }
174 175
  template <typename T>
  inline const T& Attr(const std::string& name) const {
176 177 178
    PADDLE_ENFORCE_NE(
        attrs_.find(name), attrs_.end(),
        platform::errors::NotFound("(%s) is not found in AttributeMap.", name));
179
    return BOOST_GET_CONST(T, attrs_.at(name));
180
  }
181 182 183 184 185 186 187 188
  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;
  }
189
  const AttributeMap& Attrs() const { return attrs_; }
D
dongzhihong 已提交
190

Y
Yu Yang 已提交
191 192
  const VariableNameMap& Inputs() const { return inputs_; }
  const VariableNameMap& Outputs() const { return outputs_; }
193

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

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

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

218
  void SetIsCalledByExecutor(bool x) { run_by_executor_ = x; }
219

B
baojun-nervana 已提交
220
  virtual void RuntimeInferShape(const Scope& scope,
X
Xin Pan 已提交
221 222
                                 const platform::Place& place,
                                 const RuntimeContext& ctx) const {}
223

Z
Zhang Ting 已提交
224 225 226 227 228
  virtual platform::Place GetExecutionPlace(
      const platform::Place& place) const {
    return place;
  }

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

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

  // OpInfo
  const OpInfo* info_;

245 246
  // Whether this operator executes in an Executor.
  bool run_by_executor_{true};
247 248 249 250

 private:
  void GenerateTemporaryNames();
  void CheckAllInputOutputSet() const;
251 252
  virtual void RunImpl(const Scope& scope,
                       const platform::Place& place) const = 0;
Y
Yan Chunwei 已提交
253 254
};

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

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

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

Q
qiaolongfei 已提交
286
  template <typename T>
Y
Yu Yang 已提交
287
  inline const T& Attr(const std::string& name) const {
288
    return BOOST_GET_CONST(T, GetAttr(name));
Q
qiaolongfei 已提交
289 290
  }

H
hong 已提交
291 292 293
  virtual const Attribute& GetAttr(const std::string& name) const {
    return op_.Attrs().at(name);
  }
294

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

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

H
hong 已提交
299
  virtual size_t InputSize(const std::string& name) const {
Y
Yu Yang 已提交
300
    return op_.Inputs(name).size();
Y
Yan Chunwei 已提交
301 302
  }

H
hong 已提交
303
  virtual size_t OutputSize(const std::string& name) const {
Y
Yu Yang 已提交
304
    return op_.Outputs(name).size();
Y
Yan Chunwei 已提交
305 306
  }

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

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

H
hong 已提交
311
  virtual const std::vector<Variable*> MultiInputVar(
312
      const std::string& name) const {
313 314
    LogVarUsageIfUnusedVarCheckEnabled(name);

X
Xin Pan 已提交
315 316 317 318
    auto it = ctx_.inputs.find(name);
    if (it == ctx_.inputs.end()) {
      return {};
    }
G
Gabor Buella 已提交
319
    return {it->second.begin(), it->second.end()};
X
Xin Pan 已提交
320 321
  }

H
hong 已提交
322
  virtual std::vector<Variable*> MultiOutputVar(const std::string& name) const {
X
Xin Pan 已提交
323 324 325 326 327 328 329
    auto it = ctx_.outputs.find(name);
    if (it == ctx_.outputs.end()) {
      return {};
    }
    return it->second;
  }

H
hong 已提交
330 331 332 333 334 335 336 337 338 339 340
  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;
  }

341 342
  template <typename T>
  const T* Input(const std::string& name) const {
Y
Yu Yang 已提交
343
    auto* var = InputVar(name);
344
    return var == nullptr ? nullptr : &var->Get<T>();
345 346 347 348
  }

  template <typename T>
  T* Output(const std::string& name) const {
349
    auto var = OutputVar(name);
350
    return var == nullptr ? nullptr : var->GetMutable<T>();
351 352 353 354
  }

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

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

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

    if (vars.size() == 0) {
X
Xin Pan 已提交
375 376
      return {};
    }
H
hong 已提交
377

X
Xin Pan 已提交
378 379 380 381 382 383
    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 已提交
384

X
Xin Pan 已提交
385 386 387
    return res;
  }

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

Q
QI JUN 已提交
390 391 392 393 394
  template <typename DeviceContextType>
  const DeviceContextType& device_context() const {
    return *reinterpret_cast<const DeviceContextType*>(&device_context_);
  }

395
  const platform::DeviceContext& device_context() const {
Q
qijun 已提交
396
    return device_context_;
Q
qijun 已提交
397
  }
Q
qijun 已提交
398

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

X
Xin Pan 已提交
409 410 411
  template <typename T, typename DevContext>
  Tensor AllocateTmpTensor(const framework::DDim& dim,
                           const DevContext& dev_ctx) const {
412
    auto tmp_allocation_ptr = memory::Alloc(dev_ctx, product(dim) * sizeof(T));
X
Xin Pan 已提交
413 414
    auto& deleter = tmp_allocation_ptr.get_deleter();
    auto* allocation_ptr = tmp_allocation_ptr.release();
415 416
    auto shared_allocation =
        std::shared_ptr<pten::Allocation>(allocation_ptr, deleter);
X
Xin Pan 已提交
417

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

    paddle::framework::Tensor temp_tensor(
        framework::ToDataType(std::type_index(typeid(T))));
    temp_tensor.Resize(dim);
    temp_tensor.ResetHolder(std::move(shared_allocation));
    return temp_tensor;
  }

H
hong 已提交
432 433 434
  const RuntimeContext Context() const { return ctx_; }

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

437
 private:
438 439
  const OperatorBase& op_;
  const Scope& scope_;
440
  const platform::DeviceContext& device_context_;
X
Xin Pan 已提交
441
  const RuntimeContext& ctx_;
Q
Qiao Longfei 已提交
442 443
};

444 445 446 447 448 449 450 451 452 453 454 455 456 457
// TODO(chenweihang): split impl based OpProto or Dygraph if needed
class ExecutionArgumentMappingContext : public pten::ArgumentMappingContext {
 public:
  explicit ExecutionArgumentMappingContext(const ExecutionContext& ctx)
      : ctx_(ctx) {}

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

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

458 459 460
  paddle::any Attr(const std::string& name) const override {
    auto& attr = ctx_.GetAttr(name);
    return GetAttrValue(attr);
461 462 463
  }

  size_t InputSize(const std::string& name) const override {
464
    return ctx_.MultiInputVar(name).size();
465 466 467
  }

  size_t OutputSize(const std::string& name) const override {
468
    return ctx_.MultiOutputVar(name).size();
469 470 471
  }

  bool IsDenseTensorInput(const std::string& name) const override {
472
    return ctx_.InputVar(name)->IsType<framework::LoDTensor>();
473 474 475
  }

  bool IsSelectedRowsInput(const std::string& name) const override {
476
    return ctx_.InputVar(name)->IsType<pten::SelectedRows>();
477 478
  }

479 480 481 482 483 484 485 486
  bool IsDenseTensorOutput(const std::string& name) const override {
    return ctx_.OutputVar(name)->IsType<framework::LoDTensor>();
  }

  bool IsSelectedRowsOutput(const std::string& name) const override {
    return ctx_.OutputVar(name)->IsType<pten::SelectedRows>();
  }

487 488 489 490
 private:
  const ExecutionContext& ctx_;
};

491 492 493 494 495 496 497 498
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 已提交
499
class OpKernelBase {
Q
qijun 已提交
500
 public:
Q
qijun 已提交
501
  /**
502
   * ExecutionContext is the only parameter of Kernel Run function.
Q
qijun 已提交
503 504
   * Run will get input/output variables, state such as momentum and
   * device resource such as CUDA stream, cublas handle, etc. from
505
   * ExecutionContext. User should construct it before run the Operator.
Q
qijun 已提交
506 507
   */

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

Y
Yu Yang 已提交
510 511 512 513 514 515 516
  virtual ~OpKernelBase() = default;
};

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

Y
Yu Yang 已提交
519 520
class OperatorWithKernel : public OperatorBase {
 public:
Y
yuyang18 已提交
521
  using OpKernelFunc = std::function<void(const ExecutionContext&)>;
Y
Yu Yang 已提交
522
  using OpKernelMap =
Y
yuyang18 已提交
523
      std::unordered_map<OpKernelType, OpKernelFunc, OpKernelType::Hash>;
Q
Qiao Longfei 已提交
524

Y
Yu Yang 已提交
525 526
  OperatorWithKernel(const std::string& type, const VariableNameMap& inputs,
                     const VariableNameMap& outputs, const AttributeMap& attrs)
Y
Yu Yang 已提交
527 528
      : OperatorBase(type, inputs, outputs, attrs) {}

C
chentianyu03 已提交
529
  static paddle::flat_hash_map<std::string /* op_type */, OpKernelMap>&
Y
Yu Yang 已提交
530
  AllOpKernels() {
C
chentianyu03 已提交
531
    static paddle::flat_hash_map<std::string, OpKernelMap> g_all_op_kernels;
Y
Yu Yang 已提交
532
    return g_all_op_kernels;
Y
Yu Yang 已提交
533
  }
Y
Yan Chunwei 已提交
534

535
  bool SupportGPU() const override {
536 537
    auto pten_kernels = pten::KernelFactory::Instance().SelectKernelMap(
        pten::TransToPtenKernelName(type_));
538 539 540 541 542
    auto has_pten_kernel =
        std::any_of(pten_kernels.begin(), pten_kernels.end(),
                    [](pten::KernelKeyMap::const_reference kern_pair) {
                      return kern_pair.first.backend() == pten::Backend::GPU;
                    });
543 544 545 546 547 548 549 550 551 552
    if (has_pten_kernel) {
      return true;
    } else {
      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_gpu_place(kern_pair.first.place_);
          });
    }
553
  }
554

B
Baibaifan 已提交
555
  bool SupportNPU() const override {
556
    // TODO(zhiqiu): support pten if needed?
B
Baibaifan 已提交
557 558 559 560 561 562
    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_npu_place(kern_pair.first.place_);
                       });
  }
F
fwenguang 已提交
563
  bool SupportMLU() const override {
564
    // TODO(zhiqiu): support pten if needed?
F
fwenguang 已提交
565 566 567 568 569 570
    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_);
                       });
  }
571
  bool SupportsMKLDNN(proto::VarType::Type data_type) const;
572

573 574
  bool CanMKLDNNBeUsed(const framework::ExecutionContext& ctx,
                       proto::VarType::Type data_type) const;
575

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

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

581 582 583
  proto::VarType::Type IndicateVarDataType(const ExecutionContext& ctx,
                                           const std::string& name) const;

584 585 586 587
  proto::VarType::Type IndicateOrPromoteVarDataTypes(
      const ExecutionContext& ctx, const std::string& name1,
      const std::string& name2) const;

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

590 591
  // change this to public so that in dygraph mode we can call it to check if we
  // need transform data
592 593 594
  virtual OpKernelType GetKernelTypeForVar(
      const std::string& var_name, const Tensor& tensor,
      const OpKernelType& expected_kernel_type) const;
Y
Yu Yang 已提交
595

596 597
  platform::Place GetExecutionPlace(
      const platform::Place& platform) const override {
Z
Zhang Ting 已提交
598 599 600
    return kernel_type_->place_;
  }

601 602 603 604 605 606 607 608
  /* member functions for adapting to pten lib */
  /** 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
    * original Op according to the GetExpectedPtenKernelArgs returned arguments.
    */
609
  pten::KernelSignature GetExpectedPtenKernelArgs(
610 611
      const ExecutionContext& ctx) const;

612
  /* member functions for adapting to pten lib */
613
  pten::KernelKey ChoosePtenKernel(const ExecutionContext& ctx) const;
614

615 616 617 618 619
  /**
   * Transfer data place for pten kernel
   * Is this really needed?
   */
  Scope* PreparePtenData(const Scope& scope, const pten::Kernel& pt_kernel,
620
                         const pten::KernelSignature& pt_kernel_signature,
621 622
                         RuntimeContext* ctx) const;

623
  void BuildPtenKernelContext(const RuntimeContext& ctx,
624 625
                              platform::DeviceContext* dev_ctx,
                              pten::KernelContext* pt_kernel_context) const;
626

627 628 629 630
  pten::KernelSignature* PtenKernelSignature() const {
    return pt_kernel_signature_.get();
  }

631 632
  pten::Kernel* PtenKernel() const { return pt_kernel_.get(); }

633 634 635 636
  void ResetPtenKernel(pten::Kernel* kernel) const {
    return pt_kernel_.reset(kernel);
  }

637 638
  const OpKernelType* kernel_type() const { return kernel_type_.get(); }

Y
Yu Yang 已提交
639
 private:
640
  void RunImpl(const Scope& scope, const platform::Place& place) const final;
L
luotao1 已提交
641 642
  void RunImpl(const Scope& scope, const platform::Place& place,
               RuntimeContext* runtime_ctx) const;
Y
yuyang18 已提交
643 644

  /**
T
tianshuo78520a 已提交
645 646
   * Transfer data from scope to a transferred scope. If there is no data need
   * to
Y
yuyang18 已提交
647 648 649 650
   * be tranfered, it returns nullptr.
   *
   * * transfered_inplace_vars is a output vector.
   */
X
Xin Pan 已提交
651 652 653 654
  Scope* PrepareData(const Scope& scope,
                     const OpKernelType& expected_kernel_key,
                     std::vector<std::string>* transfered_inplace_vars,
                     RuntimeContext* ctx) const;
Y
yuyang18 已提交
655 656 657 658

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

660 661 662
  OpKernelType InnerGetExpectedKernelType(const ExecutionContext& ctx) const;

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

664 665 666
  void HandleComplexGradToRealGrad(const Scope& scope,
                                   RuntimeContext* ctx) const;

667 668 669 670 671
  /* 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
672 673 674
  void ParseInputDataType(const std::vector<Variable*>& vars,
                          const std::string& name,
                          proto::VarType::Type* data_type) const;
675 676 677 678
  // used for IndicateOrPromoteVarDataTypes
  Tensor* GetTensorFormInputSafely(const ExecutionContext& ctx,
                                   const std::string& name) const;

679
 protected:
L
Liu Yiqun 已提交
680 681
  mutable std::unique_ptr<OpKernelType> kernel_type_;
  mutable std::unique_ptr<OpKernelFunc> kernel_func_;
L
luotao1 已提交
682 683
  mutable std::unique_ptr<RuntimeContext> runtime_ctx_;
  mutable const Scope* pre_scope_ = nullptr;
684
  mutable bool need_prepare_data_ = true;
685 686
  mutable bool enable_cache_runtime_context_ = false;
  mutable bool all_kernels_must_compute_runtime_shape_ = false;
687
  mutable std::mutex cache_update_mutex_;
688
  mutable bool enable_cache_transfer_scope_ = false;
689 690 691 692
  // NOTE(chenweihang): Similar op members are used to adapt to
  // new pten kernel, if there is a better design in the future,
  // we may polish the implementation here
  mutable bool run_pten_kernel_ = false;
L
Liu-xiandong 已提交
693
  mutable bool run_kp_kernel = false;
694
  mutable std::unique_ptr<pten::KernelSignature> pt_kernel_signature_;
695
  mutable std::unique_ptr<pten::Kernel> pt_kernel_;
Q
Qiao Longfei 已提交
696 697
};

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

Q
Qiao Longfei 已提交
700 701
}  // namespace framework
}  // namespace paddle