operator.h 24.4 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 136
  RuntimeContext(const VariableNameMap& innames,
                 const VariableNameMap& outnames, const Scope& scope);
X
Xin Pan 已提交
137

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  // OpInfo
  const OpInfo* info_;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

C
Chen Weihang 已提交
337 338
  virtual paddle::small_vector<const std::string*> InNameList() const {
    paddle::small_vector<const std::string*> vec_temp;
H
hong 已提交
339 340 341
    vec_temp.reserve(ctx_.inputs.size());

    for (auto& input : ctx_.inputs) {
342
      vec_temp.push_back(&input.first);
H
hong 已提交
343 344 345 346 347
    }

    return vec_temp;
  }

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

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

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

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

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

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

X
Xin Pan 已提交
385 386 387 388 389 390
    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 已提交
391

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  bool IsDenseTensorInput(const std::string& name) const override {
483 484 485 486 487
    const auto* var = ctx_.InputVar(name);
    return var->IsType<phi::DenseTensor>();
  }

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

  bool IsSelectedRowsInput(const std::string& name) const override {
495 496
    const auto* var = ctx_.InputVar(name);
    return var->IsType<phi::SelectedRows>();
497 498
  }

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

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

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

520 521
  bool IsForInferShape() const override { return false; }

522 523 524 525
 private:
  const ExecutionContext& ctx_;
};

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

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

Y
Yu Yang 已提交
545 546 547 548 549 550 551
  virtual ~OpKernelBase() = default;
};

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

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

Y
Yu Yang 已提交
560 561
  OperatorWithKernel(const std::string& type, const VariableNameMap& inputs,
                     const VariableNameMap& outputs, const AttributeMap& attrs)
Y
Yu Yang 已提交
562 563
      : OperatorBase(type, inputs, outputs, attrs) {}

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

570 571 572
  bool SupportGPU() const override;

  bool SupportNPU() const override;
573

F
fwenguang 已提交
574
  bool SupportMLU() const override {
575
    // TODO(zhiqiu): support phi if needed?
F
fwenguang 已提交
576 577 578 579 580 581
    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_);
                       });
  }
582
  bool SupportsMKLDNN(proto::VarType::Type data_type) const;
583

584 585
  bool SupportsKernelType(const OpKernelType& kernel_type) const;

586 587
  bool CanMKLDNNBeUsed(const framework::ExecutionContext& ctx,
                       proto::VarType::Type data_type) const;
588

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

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

594 595 596
  proto::VarType::Type IndicateVarDataType(const ExecutionContext& ctx,
                                           const std::string& name) const;

597 598 599 600
  proto::VarType::Type IndicateOrPromoteVarDataTypes(
      const ExecutionContext& ctx, const std::string& name1,
      const std::string& name2) const;

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

603 604
  // change this to public so that in dygraph mode we can call it to check if we
  // need transform data
605 606 607
  virtual OpKernelType GetKernelTypeForVar(
      const std::string& var_name, const Tensor& tensor,
      const OpKernelType& expected_kernel_type) const;
Y
Yu Yang 已提交
608

609 610
  platform::Place GetExecutionPlace(
      const platform::Place& platform) const override {
Z
Zhang Ting 已提交
611 612 613
    return kernel_type_->place_;
  }

614
  /* member functions for adapting to phi lib */
615
  /** In the Tensor calculation library, the new Kernel adopts a clearer and
616 617 618 619 620 621
   * 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.
   */
622
  phi::KernelSignature GetExpectedPhiKernelArgs(
623 624
      const ExecutionContext& ctx) const;

625 626
  /* member functions for adapting to phi lib */
  phi::KernelKey ChoosePhiKernel(const ExecutionContext& ctx) const;
627

628
  void ChooseKernel(const ExecutionContext& ctx) const;
629
  /**
630
   * Transfer data place for phi kernel
631 632
   * Is this really needed?
   */
633 634 635
  Scope* PreparePhiData(const Scope& scope, const phi::Kernel& pt_kernel,
                        const phi::KernelSignature& pt_kernel_signature,
                        RuntimeContext* ctx) const;
636

637 638 639
  void BuildPhiKernelContext(const RuntimeContext& ctx,
                             platform::DeviceContext* dev_ctx,
                             phi::KernelContext* pt_kernel_context) const;
640

641
  phi::KernelSignature* PhiKernelSignature() const {
642
    return kernel_signature_.get();
643 644
  }

645
  phi::Kernel* PhiKernel() const { return pt_kernel_.get(); }
646

647
  void ResetPhiKernel(phi::Kernel* kernel) const {
648 649 650
    return pt_kernel_.reset(kernel);
  }

651
  const OpKernelType* kernel_type() const { return kernel_type_.get(); }
652
  const OpKernelFunc* kernel_func() const { return kernel_func_.get(); }
653

654 655 656 657
  void ResetKernelType(OpKernelType* kernel_type) {
    kernel_type_.reset(kernel_type);
  }

Y
Yu Yang 已提交
658
 private:
659
  void RunImpl(const Scope& scope, const platform::Place& place) const final;
L
luotao1 已提交
660 661
  void RunImpl(const Scope& scope, const platform::Place& place,
               RuntimeContext* runtime_ctx) const;
Y
yuyang18 已提交
662 663

  /**
T
tianshuo78520a 已提交
664 665
   * Transfer data from scope to a transferred scope. If there is no data need
   * to
Y
yuyang18 已提交
666 667 668 669
   * be tranfered, it returns nullptr.
   *
   * * transfered_inplace_vars is a output vector.
   */
X
Xin Pan 已提交
670 671 672 673
  Scope* PrepareData(const Scope& scope,
                     const OpKernelType& expected_kernel_key,
                     std::vector<std::string>* transfered_inplace_vars,
                     RuntimeContext* ctx) const;
Y
yuyang18 已提交
674 675 676 677

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

679 680
  OpKernelType InnerGetExpectedKernelType(const ExecutionContext& ctx) const;

681 682 683
  void HandleComplexGradToRealGrad(const Scope& scope,
                                   RuntimeContext* ctx) const;

684 685 686 687 688
  /* 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
689
  void ParseInputDataType(const Variable* vars, const std::string& name,
690
                          proto::VarType::Type* data_type) const;
691 692 693
  void ParseMultiInputDataType(const std::vector<Variable*>& vars,
                               const std::string& name,
                               proto::VarType::Type* data_type) const;
694 695 696 697
  // used for IndicateOrPromoteVarDataTypes
  Tensor* GetTensorFormInputSafely(const ExecutionContext& ctx,
                                   const std::string& name) const;

698
 protected:
L
Liu Yiqun 已提交
699 700
  mutable std::unique_ptr<OpKernelType> kernel_type_;
  mutable std::unique_ptr<OpKernelFunc> kernel_func_;
L
luotao1 已提交
701 702
  mutable std::unique_ptr<RuntimeContext> runtime_ctx_;
  mutable const Scope* pre_scope_ = nullptr;
703
  mutable bool need_prepare_data_ = true;
704
  mutable bool need_prepare_phi_data_ = false;
705 706
  mutable bool enable_cache_runtime_context_ = false;
  mutable bool all_kernels_must_compute_runtime_shape_ = false;
707
  mutable std::mutex cache_update_mutex_;
708
  mutable bool enable_cache_transfer_scope_ = false;
709
  // NOTE(chenweihang): Similar op members are used to adapt to
710
  // new phi kernel, if there is a better design in the future,
711
  // we may polish the implementation here
712
  mutable bool run_phi_kernel_ = false;
L
Liu-xiandong 已提交
713
  mutable bool run_kp_kernel = false;
714
  mutable std::unique_ptr<phi::KernelSignature> kernel_signature_;
715
  mutable std::unique_ptr<phi::Kernel> pt_kernel_;
716
  mutable std::unique_ptr<phi::ArgumentMappingFn> arg_map_fn_;
717 718 719

  struct CacheImpl;
  mutable CacheImpl* impl_{nullptr};
Q
Qiao Longfei 已提交
720 721
};

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

Q
Qiao Longfei 已提交
724 725
}  // namespace framework
}  // namespace paddle