operator.cc 83.1 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

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. */
D
dzhwinter 已提交
14

15 16
#include "paddle/fluid/framework/operator.h"

17
#include <glog/logging.h>
P
peizhilin 已提交
18 19
#include <sstream>
#include <string>
20

21
#include "gflags/gflags.h"
22
#include "paddle/fluid/framework/convert_utils.h"
Y
Yi Wang 已提交
23
#include "paddle/fluid/framework/data_transform.h"
24
#include "paddle/fluid/framework/data_type_transform.h"
W
WangXi 已提交
25
#include "paddle/fluid/framework/details/nan_inf_utils.h"
26
#include "paddle/fluid/framework/op_call_stack.h"
27
#include "paddle/fluid/framework/pten_utils.h"
Y
Yi Wang 已提交
28
#include "paddle/fluid/framework/shape_inference.h"
29
#include "paddle/fluid/framework/transfer_scope_cache.h"
30
#include "paddle/fluid/framework/unused_var_check.h"
Y
Yi Wang 已提交
31
#include "paddle/fluid/framework/var_type.h"
32
#include "paddle/fluid/platform/device/device_wrapper.h"
L
Leo Chen 已提交
33
#include "paddle/fluid/platform/enforce.h"
34
#include "paddle/fluid/platform/profiler.h"
C
chenjian 已提交
35
#include "paddle/fluid/platform/profiler/event_tracing.h"
36 37 38 39
#include "paddle/phi/common/scalar.h"
#include "paddle/phi/common/scalar_array.h"
#include "paddle/phi/core/kernel_factory.h"
#include "paddle/phi/ops/compat/signatures.h"
40

41
namespace phi {
42
class DenseTensor;
43
}  // namespace phi
44

45
#ifdef PADDLE_WITH_XPU
46 47
#include "paddle/fluid/platform/device/xpu/xpu_info.h"
#include "paddle/fluid/platform/device/xpu/xpu_op_list.h"
48
#endif
Q
Qiao Longfei 已提交
49

50 51 52 53
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/fluid/platform/mkldnn_helper.h"
#endif

F
fwenguang 已提交
54 55 56 57
#ifdef PADDLE_WITH_MLU
#include "paddle/fluid/platform/device/mlu/mlu_info.h"
#endif

D
dzhwinter 已提交
58
DECLARE_bool(benchmark);
59
DECLARE_bool(check_nan_inf);
60
DECLARE_bool(enable_unused_var_check);
61 62
PADDLE_DEFINE_EXPORTED_int32(inner_op_parallelism, 0,
                             "number of threads for inner op");
F
Feng Xing 已提交
63
DECLARE_bool(run_kp_kernel);
D
dzhwinter 已提交
64

Q
Qiao Longfei 已提交
65 66 67
namespace paddle {
namespace framework {

68 69 70 71 72 73
std::vector<std::tuple<platform::Place, LibraryType>> kKernelPriority = {
    std::make_tuple(platform::CUDAPlace(0), LibraryType::kCUDNN),
    std::make_tuple(platform::CUDAPlace(0), LibraryType::kPlain),
    std::make_tuple(platform::CPUPlace(), LibraryType::kMKLDNN),
    std::make_tuple(platform::CPUPlace(), LibraryType::kPlain),
};
D
dzhwinter 已提交
74

75
static DDim GetDimsDebug(const ScopeBase& scope, const std::string& name,
76
                         bool get_actual_dim = false) {
77
  Variable* var = scope.FindVar(name);
Q
qiaolongfei 已提交
78 79
  if (var == nullptr) {
    return DDim({-1});
Q
Qiao Longfei 已提交
80 81
  }

M
minqiyang 已提交
82 83 84
  if (var->IsType<LoDTensor>()) {
    const LoDTensor& tensor = var->Get<LoDTensor>();
    return tensor.dims();
85
  } else if (var->IsType<phi::SelectedRows>()) {
M
minqiyang 已提交
86
    if (get_actual_dim) {
87
      return var->Get<phi::SelectedRows>().value().dims();
M
minqiyang 已提交
88
    } else {
89
      return var->Get<phi::SelectedRows>().GetCompleteDims();
M
minqiyang 已提交
90
    }
S
Steffy-zxf 已提交
91 92
  } else if (var->IsType<Strings>()) {
    return DDim({static_cast<int64_t>(var->Get<Strings>().size())});
93 94 95 96 97
  } else {
    return DDim({-1});
  }
}

98
static bool VarInited(const ScopeBase& scope, const std::string& name) {
Q
Qiao Longfei 已提交
99 100 101 102 103
  Variable* var = scope.FindVar(name);
  if (var == nullptr) return false;
  return var->IsInitialized();
}

104
static std::string GetDtype(const ScopeBase& scope, const std::string& name) {
D
dzhwinter 已提交
105 106 107 108
  Variable* var = scope.FindVar(name);
  if (var == nullptr) {
    return "";
  }
109

M
minqiyang 已提交
110 111 112
  if (var->IsType<LoDTensor>()) {
    const LoDTensor& tensor = var->Get<LoDTensor>();
    if (UNLIKELY(!tensor.IsInitialized())) {
113 114
      return "";
    }
115
    return DataTypeToString(framework::TransToProtoVarType(tensor.dtype()));
116 117
  } else if (var->IsType<phi::SelectedRows>()) {
    auto tensor = var->Get<phi::SelectedRows>().value();
Q
Qiao Longfei 已提交
118 119 120
    if (UNLIKELY(!tensor.IsInitialized())) {
      return "uninited";
    } else {
121
      return DataTypeToString(framework::TransToProtoVarType(tensor.dtype()));
Q
Qiao Longfei 已提交
122
    }
S
Steffy-zxf 已提交
123 124
  } else if (var->IsType<Strings>()) {
    return "strings";
D
dzhwinter 已提交
125 126 127 128 129
  } else {
    return "";
  }
}

130
static std::string GetPlace(const ScopeBase& scope, const std::string& name) {
L
Leo Chen 已提交
131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146
  Variable* var = scope.FindVar(name);
  if (var == nullptr) {
    return "";
  }
  auto to_string = [](const platform::Place& p) {
    std::stringstream sstream;
    sstream << p;
    return sstream.str();
  };

  if (var->IsType<LoDTensor>()) {
    const LoDTensor& tensor = var->Get<LoDTensor>();
    if (UNLIKELY(!tensor.IsInitialized())) {
      return "";
    }
    return to_string(tensor.place());
147 148
  } else if (var->IsType<phi::SelectedRows>()) {
    auto tensor = var->Get<phi::SelectedRows>().value();
L
Leo Chen 已提交
149 150 151 152 153 154 155 156 157 158
    if (UNLIKELY(!tensor.IsInitialized())) {
      return "uninited";
    } else {
      return to_string(tensor.place());
    }
  } else {
    return "";
  }
}

159
static int GetRowSize(const ScopeBase& scope, const std::string& name) {
160 161 162 163 164
  Variable* var = scope.FindVar(name);
  if (var == nullptr) {
    return -1;
  }

165 166
  if (var->IsType<phi::SelectedRows>()) {
    return var->Get<phi::SelectedRows>().rows().size();
167 168 169 170 171
  }

  return -1;
}

172
static LoD GetLoDDebug(const ScopeBase& scope, const std::string& name) {
Q
Qiao Longfei 已提交
173 174 175 176 177 178 179
  Variable* var = scope.FindVar(name);
  auto default_lod = LoD({{}});

  if (var == nullptr) {
    return default_lod;
  }

M
minqiyang 已提交
180 181 182
  if (var->IsType<LoDTensor>()) {
    const LoDTensor& tensor = var->Get<LoDTensor>();
    return tensor.lod();
Q
Qiao Longfei 已提交
183 184 185 186 187
  } else {
    return default_lod;
  }
}

X
Xin Pan 已提交
188 189 190 191 192
RuntimeContext::RuntimeContext(const VariableNameMap& innames,
                               const VariableNameMap& outnames,
                               const Scope& scope) {
  for (auto& var_name_item : innames) {
    std::vector<Variable*>& input_vars = inputs[var_name_item.first];
X
Xin Pan 已提交
193
    input_vars.reserve(var_name_item.second.size());
X
Xin Pan 已提交
194 195 196 197 198 199
    for (auto& var_name : var_name_item.second) {
      input_vars.push_back(scope.FindVar(var_name));
    }
  }
  for (auto& var_name_item : outnames) {
    std::vector<Variable*>& output_vars = outputs[var_name_item.first];
X
Xin Pan 已提交
200
    output_vars.reserve(var_name_item.second.size());
X
Xin Pan 已提交
201 202 203 204 205 206
    for (auto& var_name : var_name_item.second) {
      output_vars.push_back(scope.FindVar(var_name));
    }
  }
}

207
void OperatorBase::Run(const Scope& scope, const platform::Place& place) {
P
peizhilin 已提交
208 209 210
  try {
    VLOG(4) << place << " " << DebugStringEx(&scope);
    if (platform::is_gpu_place(place)) {
211
#if !defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP)
212 213 214 215
      PADDLE_THROW(platform::errors::Unavailable(
          "Cannot run operator on place %s, please recompile paddle or "
          "reinstall Paddle with CUDA support.",
          place));
216
#else
217
      auto dev_id = place.device;
P
peizhilin 已提交
218
      platform::SetDeviceId(dev_id);
219 220 221
#endif
    } else if (platform::is_xpu_place(place)) {
#ifndef PADDLE_WITH_XPU
222 223 224 225
      PADDLE_THROW(platform::errors::Unavailable(
          "Cannot run operator on place %s, please recompile paddle or "
          "reinstall Paddle with XPU support.",
          place));
226
#else
227
      auto dev_id = place.device;
228
      platform::SetXPUDeviceId(dev_id);
229 230 231 232 233 234 235 236
#endif
    } else if (platform::is_npu_place(place)) {
#ifndef PADDLE_WITH_ASCEND_CL
      PADDLE_THROW(platform::errors::Unavailable(
          "Cannot run operator on place %s, please recompile paddle or "
          "reinstall Paddle with NPU support.",
          place));
#else
237
      auto dev_id = place.device;
238
      platform::SetNPUDeviceId(dev_id);
F
fwenguang 已提交
239 240 241 242 243 244 245 246
#endif
    } else if (platform::is_mlu_place(place)) {
#ifndef PADDLE_WITH_MLU
      PADDLE_THROW(platform::errors::Unavailable(
          "Cannot run operator on place %s, please recompile paddle or "
          "reinstall Paddle with MLU support.",
          place));
#else
247
      auto dev_id = place.device;
F
fwenguang 已提交
248
      platform::SetMLUDeviceId(dev_id);
249 250 251 252 253 254 255 256 257
#endif
    } else if (platform::is_custom_place(place)) {
#ifndef PADDLE_WITH_CUSTOM_DEVICE
      PADDLE_THROW(platform::errors::Unavailable(
          "Cannot run operator on place %s, please recompile paddle or "
          "reinstall Paddle with CustomDevice support.",
          place));
#else
      platform::DeviceManager::SetDevice(place);
258
#endif
P
peizhilin 已提交
259
    }
P
peizhilin 已提交
260

261
    {
262 263 264
      // TODO(wangchaochaohu) : refine code to use only one RecordEvent)
      // in order to record different op type cost time
      // and different op name cost time,we set two event.
C
chenjian 已提交
265 266
      platform::RecordEvent op_type_record_event(
          Type().c_str(), platform::TracerEventType::Operator, 1);
267 268
      auto op_name = platform::OpName(outputs_, Type());
      platform::RecordEvent op_name_record_event(
C
chenjian 已提交
269 270
          op_name, platform::TracerEventType::Operator, 1,
          platform::EventRole::kUniqueOp);
P
peizhilin 已提交
271 272
      RunImpl(scope, place);
    }
273

Z
Zhang Ting 已提交
274
    VLOG(3) << GetExecutionPlace(place) << " " << DebugStringEx(&scope);
275
  } catch (platform::EnforceNotMet& exception) {
276
    framework::InsertCallStackInfo(Type(), Attrs(), &exception);
277
    throw std::move(exception);
278 279 280 281 282 283
  } catch (platform::EOFException&) {
    std::rethrow_exception(std::current_exception());
  } catch (std::exception& ex) {
    LOG(WARNING) << Type() << " raises an exception "
                 << platform::demangle(typeid(ex).name()) << ", " << ex.what();
    std::rethrow_exception(std::current_exception());
P
peizhilin 已提交
284
  } catch (...) {
285
    LOG(WARNING) << Type() << " raises an unknown exception";
P
peizhilin 已提交
286
    std::rethrow_exception(std::current_exception());
287
  }
288 289
}

290
bool OperatorBase::HasInputs(const std::string& name) const {
M
minqiyang 已提交
291
  return inputs_.find(name) != inputs_.end();
292 293
}

294
std::string OperatorBase::Input(const std::string& name) const {
Y
Yu Yang 已提交
295
  auto& ins = Inputs(name);
296 297
  PADDLE_ENFORCE_LE(
      ins.size(), 1UL,
298
      platform::errors::InvalidArgument(
299 300
          "Operator %s's input %s should contain only one variable.", type_,
          name));
Y
Yu Yang 已提交
301
  return ins.empty() ? kEmptyVarName : ins[0];
Y
Yan Chunwei 已提交
302 303
}

Y
Yu Yang 已提交
304 305
const std::vector<std::string>& OperatorBase::Inputs(
    const std::string& name) const {
Y
Yu Yang 已提交
306
  auto it = inputs_.find(name);
307 308 309 310
  PADDLE_ENFORCE_NE(
      it, inputs_.end(),
      platform::errors::NotFound("Operator %s does not have the input %s.",
                                 type_, name));
Y
Yu Yang 已提交
311
  return it->second;
Y
Yan Chunwei 已提交
312 313
}

314
bool OperatorBase::HasOutputs(const std::string& name) const {
315
  if (outputs_.find(name) != outputs_.end()) {
316 317 318 319 320 321
    return true;
  } else {
    return false;
  }
}

322
std::string OperatorBase::Output(const std::string& name) const {
Y
Yu Yang 已提交
323
  auto& outs = Outputs(name);
324 325 326 327 328
  PADDLE_ENFORCE_LE(
      outs.size(), 1UL,
      platform::errors::InvalidArgument(
          "Operator %s's output %s should contain only one variable.", type_,
          name));
Y
Yu Yang 已提交
329
  return outs.empty() ? kEmptyVarName : outs[0];
Y
Yan Chunwei 已提交
330 331
}

Y
Yu Yang 已提交
332 333
const std::vector<std::string>& OperatorBase::Outputs(
    const std::string& name) const {
Y
Yu Yang 已提交
334
  auto it = outputs_.find(name);
335 336 337 338
  PADDLE_ENFORCE_NE(
      it, outputs_.end(),
      platform::errors::NotFound(
          "Operator %s does not have an output called %s.", type_, name));
Y
Yu Yang 已提交
339
  return it->second;
Y
Yan Chunwei 已提交
340 341
}

342
std::string OperatorBase::DebugStringEx(const ScopeBase* scope) const {
Q
Qiao Longfei 已提交
343
  std::stringstream ss;
Y
Yu Yang 已提交
344
  ss << "Op(" << type_ << "), inputs:{";
345

346
  const std::unordered_set<std::string>* no_need_buffer_vars = nullptr;
347 348
  if (info_ && info_->NoNeedBufferVarsInferer()) {
    no_need_buffer_vars =
349 350
        &(Info().NoNeedBufferVarsInferer()(Inputs(), Outputs(), Attrs()));
    if (no_need_buffer_vars->empty()) no_need_buffer_vars = nullptr;
351 352
  }

Y
Yu Yang 已提交
353 354
  for (auto it = inputs_.begin(); it != inputs_.end();) {
    auto& input = *it;
355 356
    bool is_no_need_buffer_var =
        (no_need_buffer_vars && no_need_buffer_vars->count(input.first) > 0);
Y
Yu Yang 已提交
357 358
    ss << input.first << "[";
    for (size_t i = 0; i < input.second.size(); ++i) {
Q
Qiao Longfei 已提交
359 360
      auto var_name = input.second[i];
      ss << var_name;
361
      if (scope) {
Q
Qiao Longfei 已提交
362 363 364 365 366 367 368
        if (!VarInited(*scope, var_name)) {
          ss << "[uninited]";
        } else {
          int row_size = GetRowSize(*scope, var_name);
          if (row_size >= 0) {
            ss << "[row_size=" << row_size << "]";
          }
369 370 371
          std::string dtype = is_no_need_buffer_var
                                  ? "unknown_dtype"
                                  : GetDtype(*scope, var_name);
Q
Qiao Longfei 已提交
372
          ss << ":" << dtype;
373 374
          ss << "[" << GetDimsDebug(*scope, var_name, true) << "]";
          ss << "(" << GetLoDDebug(*scope, var_name) << ")";
L
Leo Chen 已提交
375
          ss << "(" << GetPlace(*scope, var_name) << ")";
376
        }
377
      }
Y
Yu Yang 已提交
378 379 380
      if (i != input.second.size() - 1) {
        ss << ", ";
      }
381
    }
Y
Yu Yang 已提交
382
    ss << "]";
Y
Yu Yang 已提交
383 384
    ++it;
    if (it != inputs_.end()) {
385 386
      ss << ", ";
    }
Q
Qiao Longfei 已提交
387
  }
Y
Yu Yang 已提交
388
  ss << "}, outputs:{";
Y
Yu Yang 已提交
389 390
  for (auto it = outputs_.begin(); it != outputs_.end();) {
    auto& output = *it;
Y
Yu Yang 已提交
391 392
    ss << output.first << "[";
    for (size_t i = 0; i < output.second.size(); ++i) {
Q
Qiao Longfei 已提交
393 394
      auto var_name = output.second[i];
      ss << var_name;
395
      if (scope) {
Q
Qiao Longfei 已提交
396 397 398 399 400 401 402
        if (!VarInited(*scope, var_name)) {
          ss << "[uninited]";
        } else {
          int row_size = GetRowSize(*scope, output.second[i]);
          if (row_size >= 0) {
            ss << "[row_size=" << row_size << "]";
          }
C
chengduo 已提交
403 404
          std::string dtype = GetDtype(*scope, output.second[i]);
          ss << ":" << dtype;
405 406
          ss << "[" << GetDimsDebug(*scope, var_name, true) << "]";
          ss << "(" << GetLoDDebug(*scope, var_name) << ")";
L
Leo Chen 已提交
407
          ss << "(" << GetPlace(*scope, var_name) << ")";
408
        }
409
      }
Y
Yu Yang 已提交
410 411 412
      if (i != output.second.size() - 1) {
        ss << ", ";
      }
413
    }
Y
Yu Yang 已提交
414
    ss << "]";
Y
Yu Yang 已提交
415 416
    ++it;
    if (it != outputs_.end()) {
417 418
      ss << ", ";
    }
Q
Qiao Longfei 已提交
419
  }
Y
Yu Yang 已提交
420
  ss << "}.";
Q
Qiao Longfei 已提交
421 422 423
  return ss.str();
}

Y
Yu Yang 已提交
424
OperatorBase::OperatorBase(const std::string& type,
Y
Yu Yang 已提交
425 426
                           const VariableNameMap& inputs,
                           const VariableNameMap& outputs,
Y
Yu Yang 已提交
427
                           const AttributeMap& attrs)
S
sneaxiy 已提交
428 429 430 431 432 433
    : type_(type),
      inputs_(inputs),
      outputs_(outputs),
      attrs_(attrs),
      // NOTE(zjl): why op_info may be nullptr?
      info_(OpInfoMap::Instance().GetNullable(type)) {
H
hong 已提交
434 435 436 437 438 439 440 441
  // In dygraph mode, all the OperatorBase will be constructed by function:
  // framework::OpRegistry::CreateOp(type, {}, {}, {}, false).
  // Inputs, outputs and attrs will be set to empty map
  // to improve the execution efficiency of dygraph.
  if (inputs_.size() > 0 || outputs_.size() > 0) {
    GenerateTemporaryNames();
    CheckAllInputOutputSet();
  }
Y
Yu Yang 已提交
442
}
443

Q
qijun 已提交
444 445
std::vector<std::string> OperatorBase::InputVars() const {
  std::vector<std::string> ret_val;
Y
Yu Yang 已提交
446
  for (auto& o : inputs_) {
Q
qijun 已提交
447 448 449 450 451 452
    ret_val.reserve(ret_val.size() + o.second.size());
    ret_val.insert(ret_val.end(), o.second.begin(), o.second.end());
  }
  return ret_val;
}

Y
Yu Yang 已提交
453 454 455 456 457 458 459 460 461 462
std::vector<std::string> OperatorBase::OutputVars(bool has_intermediate) const {
  std::vector<std::string> ret_val;
  if (has_intermediate) {
    // push all outputs into ret_val
    for (auto& o : outputs_) {
      ret_val.reserve(ret_val.size() + o.second.size());
      ret_val.insert(ret_val.end(), o.second.begin(), o.second.end());
    }
    return ret_val;
  }
S
sneaxiy 已提交
463
  auto& info = Info();
Y
Yu Yang 已提交
464 465

  // get all OpProto::Var for outputs
Y
Yu Yang 已提交
466
  for (auto& o : info.Proto().outputs()) {
Y
Yu Yang 已提交
467 468 469 470 471 472 473 474 475
    // ignore all intermediate output
    if (o.intermediate()) continue;
    auto out = outputs_.find(o.name());
    if (out != outputs_.end()) {
      ret_val.reserve(ret_val.size() + out->second.size());
      ret_val.insert(ret_val.end(), out->second.begin(), out->second.end());
    }
  }
  return ret_val;
D
dongzhihong 已提交
476 477
}

478
void OperatorBase::CheckAllInputOutputSet() const {
S
sneaxiy 已提交
479
  if (info_ == nullptr || info_->proto_ == nullptr) return;
480

S
sneaxiy 已提交
481
  for (auto& in : info_->Proto().inputs()) {
482
    if (!in.dispensable() && !in.extra()) {
483 484 485 486
      PADDLE_ENFORCE_NE(
          inputs_.find(in.name()), inputs_.end(),
          platform::errors::NotFound("Operator %s's input (%s) is not set.",
                                     Type(), in.name()));
487
    }
488 489
  }

S
sneaxiy 已提交
490
  for (auto& out : info_->Proto().outputs()) {
491
    if (!out.dispensable() && !out.extra()) {
492 493 494 495
      PADDLE_ENFORCE_NE(
          outputs_.find(out.name()), outputs_.end(),
          platform::errors::NotFound("Operator %s's output (%s) is not set.",
                                     Type(), out.name()));
496
    }
497 498 499 500 501 502 503 504 505 506 507 508 509 510 511
  }
}

void OperatorBase::GenerateTemporaryNames() {
  static std::atomic<size_t> gUniqId(0UL);
  for (auto& output : outputs_) {
    for (auto& output_name : output.second) {
      if (output_name == kTempVarName) {
        output_name += type_;
        output_name += "@";
        output_name += std::to_string(gUniqId.fetch_add(1));
      }
    }
  }
}
512

C
chengduo 已提交
513
const Tensor* GetLoDTensorOrSelectedRowsValueFromVar(const Variable& var) {
C
chengduo 已提交
514 515
  if (var.IsType<LoDTensor>()) {
    return static_cast<const Tensor*>(&(var.Get<LoDTensor>()));
516 517
  } else if (var.IsType<phi::SelectedRows>()) {
    return &(var.Get<phi::SelectedRows>().value());
Q
QI JUN 已提交
518
  } else {
519 520 521
    PADDLE_THROW(platform::errors::InvalidArgument(
        "Variable type is %s, expect LoDTensor or SelectedRows.",
        ToTypeName(var.Type())));
Q
QI JUN 已提交
522 523 524
  }
}

C
chengduo 已提交
525
Tensor* GetMutableLoDTensorOrSelectedRowsValueFromVar(Variable* var) {
Q
QI JUN 已提交
526
  if (var->IsType<LoDTensor>()) {
527
    return var->GetMutable<LoDTensor>();
528 529
  } else if (var->IsType<phi::SelectedRows>()) {
    return var->GetMutable<phi::SelectedRows>()->mutable_value();
Q
QI JUN 已提交
530
  } else {
531 532 533
    PADDLE_THROW(platform::errors::InvalidArgument(
        "Variable type is %s, expect LoDTensor or SelectedRows.",
        ToTypeName(var->Type())));
Q
QI JUN 已提交
534 535 536
  }
}

537
bool ExecutionContext::HasInput(const std::string& name) const {
538
  auto* var = InputVar(name);
539 540 541 542
  return var != nullptr;
}

bool ExecutionContext::HasOutput(const std::string& name) const {
543
  auto* var = OutputVar(name);
544 545 546
  return var != nullptr;
}

X
Xin Pan 已提交
547
const Variable* ExecutionContext::InputVar(const std::string& name) const {
548 549
  LogVarUsageIfUnusedVarCheckEnabled(name);

X
Xin Pan 已提交
550 551 552
  auto it = ctx_.inputs.find(name);
  if (it == ctx_.inputs.end()) return nullptr;

553 554
  PADDLE_ENFORCE_LE(
      it->second.size(), 1UL,
555
      platform::errors::InvalidArgument(
556 557
          "Operator %s's input %s should contain only one variable.",
          op_.Type(), name));
X
Xin Pan 已提交
558 559 560
  return it->second.empty() ? nullptr : it->second[0];
}

X
clean  
Xin Pan 已提交
561
Variable* ExecutionContext::OutputVar(const std::string& name) const {
X
Xin Pan 已提交
562 563 564
  auto it = ctx_.outputs.find(name);
  if (it == ctx_.outputs.end()) return nullptr;

565 566 567 568 569
  PADDLE_ENFORCE_LE(
      it->second.size(), 1UL,
      platform::errors::InvalidArgument(
          "Operator %s's output %s should contain only one variable.",
          op_.Type(), name));
X
Xin Pan 已提交
570 571 572
  return it->second.empty() ? nullptr : it->second[0];
}

573
template <>
574
const std::vector<const Tensor*> ExecutionContext::MultiInput<Tensor>(
575
    const std::string& name) const {
576 577
  LogVarUsageIfUnusedVarCheckEnabled(name);

H
hong 已提交
578 579
  auto vars = MultiInputVar(name);
  if (vars.size() == 0) {
X
Xin Pan 已提交
580 581 582 583 584
    return {};
  }
  std::vector<const Tensor*> res;
  res.reserve(vars.size());
  std::transform(vars.begin(), vars.end(), std::back_inserter(res),
H
hong 已提交
585
                 [&](const Variable* var) -> const Tensor* {
X
Xin Pan 已提交
586
                   if (var == nullptr) return nullptr;
587 588 589 590 591
                   PADDLE_ENFORCE_EQ(var->IsType<LoDTensor>(), true,
                                     platform::errors::InvalidArgument(
                                         "Input variable should be LoDTensor, "
                                         "but the received type is %s.",
                                         ToTypeName(var->Type())));
X
Xin Pan 已提交
592 593 594 595 596
                   return &(var->Get<LoDTensor>());
                 });
  return res;
}

597
template <>
598
std::vector<Tensor*> ExecutionContext::MultiOutput<Tensor>(
599
    const std::string& name) const {
H
hong 已提交
600 601 602
  auto vars = MultiOutputVar(name);

  if (vars.size() == 0) {
603 604
    return {};
  }
605
  std::vector<Tensor*> res;
606 607 608 609 610
  res.reserve(vars.size());
  std::transform(vars.begin(), vars.end(), std::back_inserter(res),
                 [&](Variable* var) -> Tensor* {
                   return var == nullptr ? nullptr
                                         : var->GetMutable<LoDTensor>();
611
                 });
612 613 614
  return res;
}

Y
Yu Yang 已提交
615
bool OpSupportGPU(const std::string& op_type) {
H
hong 已提交
616
  // check in new Function kernel first
617
  auto& kernel_factory = phi::KernelFactory::Instance();
H
hong 已提交
618
  auto kernel_key_map =
619
      kernel_factory.SelectKernelMap(phi::TransToPtenKernelName(op_type));
H
hong 已提交
620
  for (auto& kernel : kernel_key_map) {
621
    if (platform::is_gpu_place(phi::TransToPtenPlace(kernel.first.backend()))) {
H
hong 已提交
622 623 624 625
      return true;
    }
  }

Y
Yu Yang 已提交
626 627 628 629 630 631 632 633 634 635 636
  auto& all_kernels = OperatorWithKernel::AllOpKernels();
  auto it = all_kernels.find(op_type);
  if (it == all_kernels.end()) {
    // All control operator must support GPU
    return true;
  }
  for (auto& kern_pair : it->second) {
    if (platform::is_gpu_place(kern_pair.first.place_)) {
      return true;
    }
  }
H
hong 已提交
637

Y
Yu Yang 已提交
638 639 640
  return false;
}

641 642
class RuntimeInferShapeContext : public InferShapeContext {
 public:
643
  RuntimeInferShapeContext(const OperatorBase& op, const RuntimeContext& ctx)
G
Gabor Buella 已提交
644
      : op_(op), ctx_(ctx) {}
645 646

  bool HasInput(const std::string& name) const override {
647
    // has only one input
X
Xin Pan 已提交
648
    const auto& ins = ctx_.inputs;
649 650
    auto it = ins.find(name);
    if (it == ins.end()) {
651 652
      return false;
    }
653
    const auto& in = it->second;
X
Xin Pan 已提交
654
    if (in.size() == 0) return false;
655 656 657 658
    PADDLE_ENFORCE_EQ(
        in.size(), 1UL,
        platform::errors::InvalidArgument(
            "Input %s should not contain more than one inputs.", name));
X
Xin Pan 已提交
659
    return in[0] != nullptr;
660 661 662
  }

  bool HasOutput(const std::string& name) const override {
663
    // has only one output
X
Xin Pan 已提交
664
    const auto& outs = ctx_.outputs;
665 666
    auto it = outs.find(name);
    if (it == outs.end()) {
667 668
      return false;
    }
669
    const auto& out = it->second;
X
Xin Pan 已提交
670
    if (out.size() == 0) {
671 672
      return false;
    }
673 674 675 676
    PADDLE_ENFORCE_EQ(
        out.size(), 1UL,
        platform::errors::InvalidArgument(
            "Output %s should not contain more than one outputs.", name));
X
Xin Pan 已提交
677
    return out[0] != nullptr;
678 679
  }

680 681 682 683
  bool HasAttr(const std::string& name) const override {
    return op_.HasAttr(name);
  }

684
  bool HasInputs(const std::string& name) const override {
X
Xin Pan 已提交
685 686
    const auto& ins = ctx_.inputs;
    auto it = ins.find(name);
X
fix  
Xin Pan 已提交
687
    if (it == ins.end() || it->second.empty()) {
688 689
      return false;
    }
X
Xin Pan 已提交
690 691
    for (auto& input : it->second) {
      if (input == nullptr) {
692 693 694 695 696 697 698
        return false;
      }
    }
    return true;
  }

  bool HasOutputs(const std::string& name) const override {
X
Xin Pan 已提交
699 700
    const auto& outs = ctx_.outputs;
    auto it = outs.find(name);
X
fix  
Xin Pan 已提交
701
    if (it == outs.end() || it->second.empty()) {
702 703
      return false;
    }
X
Xin Pan 已提交
704 705
    for (auto& output : it->second) {
      if (output == nullptr) {
706 707 708 709 710 711 712 713
        return false;
      }
    }
    return true;
  }

  AttrReader Attrs() const override { return AttrReader(op_.Attrs()); }

H
hong 已提交
714
  std::vector<std::string> Inputs(const std::string& name) const override {
715 716 717
    return op_.Inputs(name);
  }

H
hong 已提交
718
  std::vector<std::string> Outputs(const std::string& name) const override {
719 720 721
    return op_.Outputs(name);
  }

722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744
  std::string GetInputNameByIdx(size_t idx) const override {
    auto& op_proto =
        paddle::framework::OpInfoMap::Instance().Get(op_.Type()).proto_;
    PADDLE_ENFORCE_LT(idx, op_proto->inputs().size(),
                      platform::errors::OutOfRange(
                          "The index should be less than the size of inputs of "
                          "operator %s, but got index is %d and size is %d",
                          op_.Type(), idx, op_proto->inputs().size()));
    return op_proto->inputs()[idx].name();
  }

  std::string GetOutputNameByIdx(size_t idx) const override {
    auto& op_proto =
        paddle::framework::OpInfoMap::Instance().Get(op_.Type()).proto_;
    PADDLE_ENFORCE_LT(
        idx, op_proto->outputs().size(),
        platform::errors::OutOfRange(
            "The index should be less than the size of outputs of "
            "operator %s, but got index is %d and size is %d",
            op_.Type(), idx, op_proto->outputs().size()));
    return op_proto->outputs()[idx].name();
  }

745 746
  void ShareDim(const std::string& in, const std::string& out, size_t i = 0,
                size_t j = 0) override {
X
Xin Pan 已提交
747 748
    auto in_it = ctx_.inputs.find(in);
    auto out_it = ctx_.outputs.find(out);
749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764
    PADDLE_ENFORCE_NE(
        in_it, ctx_.inputs.end(),
        platform::errors::NotFound("Input %s does not exist.", in));
    PADDLE_ENFORCE_NE(
        out_it, ctx_.outputs.end(),
        platform::errors::NotFound("Output %s does not exist.", out));
    PADDLE_ENFORCE_LT(i, in_it->second.size(),
                      platform::errors::InvalidArgument(
                          "The index of input dimension is out of range, "
                          "excepted index less than %zu, but received %zu.",
                          in_it->second.size(), i));
    PADDLE_ENFORCE_LT(j, out_it->second.size(),
                      platform::errors::InvalidArgument(
                          "The index of output dimension is out of range, "
                          "excepted index less than %zu, but received %zu.",
                          out_it->second.size(), j));
X
Xin Pan 已提交
765 766 767

    Variable* in_var = in_it->second[i];
    Variable* out_var = out_it->second[j];
768

769 770 771 772 773
    PADDLE_ENFORCE_EQ(
        in_var->Type(), out_var->Type(),
        platform::errors::InvalidArgument(
            "The type of input (%s) and output (%s) are inconsistent.", in,
            out));
774

775 776 777
    if (in_var->IsType<phi::SelectedRows>()) {
      auto& in_sele_rows = in_var->Get<phi::SelectedRows>();
      auto out_sele_rows = out_var->GetMutable<phi::SelectedRows>();
778 779 780 781 782 783 784 785
      out_sele_rows->mutable_value()->Resize(in_sele_rows.value().dims());
      out_sele_rows->set_rows(in_sele_rows.rows());
      out_sele_rows->set_height(in_sele_rows.height());
    } else if (in_var->IsType<framework::LoDTensor>()) {
      auto& in_lod_tensor = in_var->Get<framework::LoDTensor>();
      auto* out_lod_tensor = out_var->GetMutable<framework::LoDTensor>();
      out_lod_tensor->Resize(in_lod_tensor.dims());
    } else {
786
      PADDLE_THROW(platform::errors::Unimplemented(
787
          "Currently, the input type of ShareDim only can be LoDTensor "
788
          "or SelectedRows."));
789 790 791
    }
  }

H
hong 已提交
792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809
  void ShareAllLoD(const std::string& in,
                   const std::string& out) const override {
    auto in_it = ctx_.inputs.find(in);
    auto out_it = ctx_.outputs.find(out);
    PADDLE_ENFORCE_NE(in_it, ctx_.inputs.end(),
                      platform::errors::NotFound(
                          "Input [%s] found error in Op [%s]", in, op_.Type()));
    PADDLE_ENFORCE_NE(
        out_it, ctx_.outputs.end(),
        platform::errors::NotFound("Output [%s] found error in Op [%s]", out,
                                   op_.Type()));

    auto& in_var_list = in_it->second;
    auto& out_var_list = out_it->second;

    PADDLE_ENFORCE_EQ(
        in_var_list.size(), out_var_list.size(),
        platform::errors::PreconditionNotMet(
T
tianshuo78520a 已提交
810
            "Op [%s]: Input var size should be equal with output var size",
H
hong 已提交
811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836
            op_.Type()));

    auto& out_var_names = op_.Outputs(out);

    for (size_t i = 0; i < in_var_list.size(); ++i) {
      if (out_var_names[i] == framework::kEmptyVarName) {
        continue;
      }

      Variable* in_var = in_var_list[i];
      if (!in_var->IsType<LoDTensor>()) return;
      Variable* out_var = out_var_list[i];
      PADDLE_ENFORCE_EQ(out_var->IsType<LoDTensor>(), true,
                        platform::errors::PreconditionNotMet(
                            "The %d-th output of Output(%s) must be LoDTensor.",
                            i, out_var_names[i]));
      auto& in_tensor = in_var->Get<LoDTensor>();
      auto* out_tensor = out_var->GetMutable<LoDTensor>();
      out_tensor->set_lod(in_tensor.lod());
#ifdef PADDLE_WITH_MKLDNN
      if (in_tensor.layout() != DataLayout::kMKLDNN)
#endif
        out_tensor->set_layout(in_tensor.layout());
    }
  }

Q
Qiao Longfei 已提交
837 838
  void ShareLoD(const std::string& in, const std::string& out, size_t i = 0,
                size_t j = 0) const override {
X
Xin Pan 已提交
839 840
    auto in_it = ctx_.inputs.find(in);
    auto out_it = ctx_.outputs.find(out);
841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856
    PADDLE_ENFORCE_NE(
        in_it, ctx_.inputs.end(),
        platform::errors::NotFound("Input %s does not exist.", in));
    PADDLE_ENFORCE_NE(
        out_it, ctx_.outputs.end(),
        platform::errors::NotFound("Output %s does not exist.", out));
    PADDLE_ENFORCE_LT(i, in_it->second.size(),
                      platform::errors::InvalidArgument(
                          "The index of input dimension is out of range, "
                          "excepted index less than %zu, but received %zu.",
                          in_it->second.size(), i));
    PADDLE_ENFORCE_LT(j, out_it->second.size(),
                      platform::errors::InvalidArgument(
                          "The index of output dimension is out of range, "
                          "excepted index less than %zu, but received %zu.",
                          out_it->second.size(), j));
X
Xin Pan 已提交
857 858

    Variable* in_var = in_it->second.at(i);
Q
Qiao Longfei 已提交
859
    if (!in_var->IsType<LoDTensor>()) return;
X
Xin Pan 已提交
860
    Variable* out_var = out_it->second.at(j);
861 862 863 864
    PADDLE_ENFORCE_EQ(
        out_var->IsType<LoDTensor>(), true,
        platform::errors::InvalidArgument(
            "The %zu-th output of Output(%s) must be LoDTensor.", j, out));
865
    auto& in_tensor = in_var->Get<LoDTensor>();
Q
Qiao Longfei 已提交
866 867
    auto* out_tensor = out_var->GetMutable<LoDTensor>();
    out_tensor->set_lod(in_tensor.lod());
D
dzhwinter 已提交
868

M
mozga-intel 已提交
869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887
// TODO(dzhwinter) : reuse ShareLoD in most operators.
// Need to call ShareLayout explicitly in sequence related ops.
// Shall we have a better method to shared info between in/out Tensor?
#ifdef PADDLE_WITH_MKLDNN
    // Fix me: ugly workaround below
    // Correct solution:
    //    set_layout() should NOT be called here (i.e. ShareLoD). Instead,
    //    layout of output tensor should be set "manually" in Compute()
    //    of each OPKernel. The reason layout should NOT be shared between
    //    input and output "automatically" (now by InferShape()->ShareLoD())
    //    is that layout transform may occur after InferShape().
    // Workaround:
    //    Skip set_layout() when input layout is kMKLDNN
    //    This is to avoid kMKLDNN is populated wrongly into a non-MKLDNN
    //    OPKernel. In all MKLDNN OPkernel, set_layout(kMKLDNN) should be called
    //    in Compute()
    if (in_tensor.layout() != DataLayout::kMKLDNN)
#endif
      out_tensor->set_layout(in_tensor.layout());
D
dzhwinter 已提交
888 889
  }

890
  int32_t GetLoDLevel(const std::string& in, size_t i = 0) const override {
891
    PADDLE_THROW(platform::errors::PreconditionNotMet(
892
        "GetLoDLevel is only used in compile time. The calculation of "
893
        "output's actual lod is different among operators so that should be "
894
        "set in the runtime kernel."));
895 896
  }

897 898
  void SetLoDLevel(const std::string& out, int32_t lod_level,
                   size_t j = 0) const override {
899
    PADDLE_THROW(platform::errors::PreconditionNotMet(
900
        "SetLoDLevel is only used in compile time. The calculation of "
901
        "output's actual lod is different among operators so that should be "
902
        "set in the runtime kernel."));
C
chengduo 已提交
903 904
  }

905 906
  bool IsRuntime() const override { return true; }

907 908 909 910 911 912 913 914 915 916 917
  bool IsRunMKLDNNKernel() const override {
    try {
      auto& op_with_kernel = dynamic_cast<const OperatorWithKernel&>(op_);
      return ((op_with_kernel.kernel_type()) &&
              (op_with_kernel.kernel_type()->data_layout_ ==
               framework::DataLayout::kMKLDNN));
    } catch (std::bad_cast exp) {
      return false;
    }
  }

918 919
  // TODO(paddle-dev): Can this be template?
  std::vector<InferShapeVarPtr> GetInputVarPtrs(
920
      const std::string& name) const override {
921 922 923 924 925 926 927 928
    const std::vector<Variable*>& vars = InputVars(name);
    std::vector<InferShapeVarPtr> res;
    res.reserve(vars.size());
    res.insert(res.begin(), vars.begin(), vars.end());
    return res;
  }

  std::vector<InferShapeVarPtr> GetOutputVarPtrs(
929
      const std::string& name) const override {
930 931 932 933 934 935 936
    const std::vector<Variable*>& vars = OutputVars(name);
    std::vector<InferShapeVarPtr> res;
    res.reserve(vars.size());
    res.insert(res.begin(), vars.begin(), vars.end());
    return res;
  }

X
Xin Pan 已提交
937 938
  DDim GetInputDim(const std::string& name) const override {
    const std::vector<Variable*>& vars = InputVars(name);
939 940 941 942 943
    PADDLE_ENFORCE_EQ(
        vars.size(), 1UL,
        platform::errors::InvalidArgument(
            "Input(%s) should hold one element, but now it holds %zu elements.",
            name, vars.size()));
X
Xin Pan 已提交
944 945 946 947 948 949 950 951
    return this->GetDim(vars[0]);
  }

  std::vector<DDim> GetInputsDim(const std::string& name) const override {
    const std::vector<Variable*>& vars = InputVars(name);
    return GetDims(vars);
  }

X
Xin Pan 已提交
952 953 954 955 956 957 958 959 960 961
  std::vector<proto::VarType::Type> GetInputsVarType(
      const std::string& name) const override {
    return GetVarTypes(InputVars(name));
  }

  std::vector<proto::VarType::Type> GetOutputsVarType(
      const std::string& name) const override {
    return GetVarTypes(OutputVars(name));
  }

X
Xin Pan 已提交
962 963
  void SetOutputDim(const std::string& name, const DDim& dim) override {
    auto& vars = OutputVars(name);
964 965 966 967 968
    PADDLE_ENFORCE_EQ(
        vars.size(), 1UL,
        platform::errors::InvalidArgument("Output(%s) should hold one element, "
                                          "but now it holds %zu elements.",
                                          name, vars.size()));
X
Xin Pan 已提交
969 970 971 972 973 974 975 976 977
    SetDim(vars[0], dim);
  }

  void SetOutputsDim(const std::string& name,
                     const std::vector<DDim>& dims) override {
    auto& vars = OutputVars(name);
    SetDims(vars, dims);
  }

978
 protected:
X
Xin Pan 已提交
979
  DDim GetDim(Variable* var) const {
980 981
    PADDLE_ENFORCE_NOT_NULL(
        var, platform::errors::InvalidArgument("Input variable is nullptr."));
982 983
    if (var->IsType<LoDTensor>()) {
      return var->Get<LoDTensor>().dims();
984 985
    } else if (var->IsType<phi::SelectedRows>()) {
      return var->Get<phi::SelectedRows>().GetCompleteDims();
986
    } else {
987 988 989 990
      PADDLE_THROW(platform::errors::InvalidArgument(
          "Only LoDTensor or SelectedRows support 'GetDim', but input "
          "Variable's type is %s.",
          ToTypeName(var->Type())));
F
fengjiayi 已提交
991 992 993
    }
  }

X
Xin Pan 已提交
994 995 996 997 998 999 1000 1001
  std::vector<DDim> GetDims(const std::vector<Variable*>& vars) const {
    std::vector<DDim> ret;
    ret.reserve(vars.size());
    std::transform(vars.begin(), vars.end(), std::back_inserter(ret),
                   [this](Variable* var) { return this->GetDim(var); });
    return ret;
  }

F
fengjiayi 已提交
1002
  std::vector<DDim> GetRepeatedDims(const std::string& name) const override {
1003 1004
    PADDLE_THROW(platform::errors::PreconditionNotMet(
        "GetRepeatedDims method only ban be used in compile time."));
1005 1006
  }

X
Xin Pan 已提交
1007
  void SetDim(Variable* var, const DDim& dim) {
1008 1009
    if (var->IsType<LoDTensor>()) {
      var->GetMutable<LoDTensor>()->Resize(dim);
1010 1011
    } else if (var->IsType<phi::SelectedRows>()) {
      var->GetMutable<phi::SelectedRows>()->set_height(dim[0]);
1012
    } else {
1013 1014 1015 1016
      PADDLE_THROW(platform::errors::Unimplemented(
          "Variable type error, expect LoDTensor or SelectedRows, but received "
          "(%s).",
          ToTypeName(var->Type())));
X
Xin Pan 已提交
1017 1018 1019 1020 1021 1022
    }
  }

  void SetDims(const std::vector<Variable*>& vars,
               const std::vector<DDim>& dims) {
    size_t length = vars.size();
1023 1024 1025 1026 1027 1028
    PADDLE_ENFORCE_EQ(length, dims.size(),
                      platform::errors::InvalidArgument(
                          "The number of input variables do not match the "
                          "number of input dimensions, the number of variables "
                          "is %zu, the number of dimensions is %zu.",
                          length, dims.size()));
X
Xin Pan 已提交
1029 1030 1031 1032 1033
    for (size_t i = 0; i < length; ++i) {
      if (vars[i] == nullptr) {
        continue;
      }
      SetDim(vars[i], dims[i]);
1034 1035 1036
    }
  }

F
fengjiayi 已提交
1037 1038
  void SetRepeatedDims(const std::string& name,
                       const std::vector<DDim>& dims) override {
1039 1040
    PADDLE_THROW(platform::errors::PreconditionNotMet(
        "SetRepeatedDims method only can be used in compile time."));
F
fengjiayi 已提交
1041 1042
  }

X
Xin Pan 已提交
1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053
  std::vector<proto::VarType::Type> GetVarTypes(
      const std::vector<Variable*>& vars) const {
    std::vector<proto::VarType::Type> retv;
    retv.resize(vars.size());
    std::transform(vars.begin(), vars.end(), retv.begin(),
                   std::bind(std::mem_fn(&RuntimeInferShapeContext::GetVarType),
                             this, std::placeholders::_1));
    return retv;
  }

  proto::VarType::Type GetVarType(Variable* var) const {
1054 1055 1056
    return ToVarType(var->Type());
  }

1057 1058 1059
 private:
  const std::vector<Variable*>& InputVars(const std::string& name) const {
    auto it = ctx_.inputs.find(name);
1060 1061 1062 1063
    PADDLE_ENFORCE_NE(
        it, ctx_.inputs.end(),
        platform::errors::NotFound(
            "Operator (%s) does not have the input (%s).", op_.Type(), name));
1064 1065 1066 1067 1068
    return it->second;
  }

  const std::vector<Variable*>& OutputVars(const std::string& name) const {
    auto it = ctx_.outputs.find(name);
1069 1070 1071 1072
    PADDLE_ENFORCE_NE(
        it, ctx_.outputs.end(),
        platform::errors::NotFound(
            "Operator (%s) does not have the outputs (%s).", op_.Type(), name));
1073
    return it->second;
F
fengjiayi 已提交
1074 1075
  }

1076
  const OperatorBase& op_;
X
Xin Pan 已提交
1077
  const RuntimeContext& ctx_;
1078 1079
};

1080 1081
static void CheckTensorNANOrInf(const std::string& op_type,
                                const std::string& name,
C
chengduoZH 已提交
1082 1083 1084 1085
                                const framework::Tensor& tensor) {
  if (tensor.memory_size() == 0) {
    return;
  }
1086 1087
  if (framework::TransToProtoVarType(tensor.dtype()) != proto::VarType::FP32 &&
      framework::TransToProtoVarType(tensor.dtype()) != proto::VarType::FP64) {
C
chengduoZH 已提交
1088 1089
    return;
  }
1090 1091 1092 1093 1094 1095 1096 1097
  PADDLE_ENFORCE_NE(
      framework::TensorContainsInf(tensor), true,
      platform::errors::Fatal("Operator %s output Tensor %s contains Inf.",
                              op_type, name));
  PADDLE_ENFORCE_NE(
      framework::TensorContainsNAN(tensor), true,
      platform::errors::Fatal("Operator %s output Tensor %s contains NAN.",
                              op_type, name));
C
chengduoZH 已提交
1098 1099
}

1100 1101
bool OperatorWithKernel::SupportsMKLDNN(
    const proto::VarType::Type data_type) const {
1102 1103
  auto& op_kernels = OperatorWithKernel::AllOpKernels().at(type_);
  return std::any_of(op_kernels.begin(), op_kernels.end(),
1104
                     [data_type](OpKernelMap::const_reference kern_pair) {
1105 1106
                       return platform::is_cpu_place(kern_pair.first.place_) &&
                              kern_pair.first.library_type_ ==
1107 1108
                                  LibraryType::kMKLDNN &&
                              kern_pair.first.data_type_ == data_type;
1109 1110 1111
                     });
}

1112 1113
bool OperatorWithKernel::CanMKLDNNBeUsed(const framework::ExecutionContext& ctx,
                                         proto::VarType::Type data_type) const {
1114 1115 1116
  bool use_mkldnn_ctx = ctx.HasAttr("use_mkldnn") &&
                        ctx.Attr<bool>("use_mkldnn") &&
                        platform::is_cpu_place(ctx.GetPlace());
1117
  return use_mkldnn_ctx && this->SupportsMKLDNN(data_type);
1118 1119
}

1120 1121 1122 1123 1124 1125 1126
void OperatorWithKernel::InferShape(InferShapeContext* ctx) const {
  PADDLE_THROW(platform::errors::PermissionDenied(
      "The default InferShape function of OperatorWithKernel is not allowed to "
      "be called, please override corresponding InferShape function in the "
      "specific operator."));
}

B
baojun-nervana 已提交
1127
void OperatorWithKernel::RuntimeInferShape(const Scope& scope,
X
Xin Pan 已提交
1128 1129
                                           const platform::Place& place,
                                           const RuntimeContext& ctx) const {
1130
  RuntimeInferShapeContext infer_shape_ctx(*this, ctx);
1131
  this->Info().infer_shape_(&infer_shape_ctx);
B
baojun-nervana 已提交
1132 1133
}

L
luotao1 已提交
1134 1135
void OperatorWithKernel::RunImpl(const Scope& scope,
                                 const platform::Place& place) const {
L
luotao1 已提交
1136 1137
  // To reduce the elapsed time of HasAttr, we use bool variable to record the
  // result of HasAttr.
1138 1139 1140
  if (!enable_cache_runtime_context_ && HasAttr(kEnableCacheRuntimeContext))
    enable_cache_runtime_context_ = true;
  if (!all_kernels_must_compute_runtime_shape_ &&
L
luotao1 已提交
1141
      HasAttr(kAllKernelsMustComputeRuntimeShape))
1142
    all_kernels_must_compute_runtime_shape_ = true;
1143
  const Scope* cur_scope = &scope;
1144
  if (!enable_cache_runtime_context_) {
L
luotao1 已提交
1145 1146
    RuntimeContext ctx(Inputs(), Outputs(), scope);
    RunImpl(scope, place, &ctx);
1147
    pre_scope_ = cur_scope;
L
luotao1 已提交
1148
  } else {
1149 1150 1151 1152 1153 1154
    if (runtime_ctx_.get() == nullptr || pre_scope_ != cur_scope) {
      std::lock_guard<std::mutex> lock(cache_update_mutex_);
      if (runtime_ctx_.get() == nullptr || pre_scope_ != cur_scope) {
        runtime_ctx_.reset(new RuntimeContext(Inputs(), Outputs(), scope));
        pre_scope_ = cur_scope;
      }
L
luotao1 已提交
1155 1156 1157 1158 1159 1160 1161 1162
    }
    RunImpl(scope, place, runtime_ctx_.get());
  }
}

void OperatorWithKernel::RunImpl(const Scope& scope,
                                 const platform::Place& place,
                                 RuntimeContext* runtime_ctx) const {
Y
Yu Yang 已提交
1163
  platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
1164
  auto* dev_ctx = pool.Get(place);
1165

1166 1167 1168 1169 1170 1171 1172 1173 1174 1175
#ifdef PADDLE_WITH_ASCEND_CL
  // NOTE(wangxi): nan/inf cannot be detected on NPU by checking the variable
  // values, but only through special `float_status` to checks whether
  // the operation is overflow. More about `float_status`, see:
  // https://gitee.com/ascend/modelzoo/issues/I3NF8V?from=project-issue
  if (FLAGS_check_nan_inf) {
    framework::details::NPUAllocAndClearFloatStatus(*this, scope, place);
  }
#endif

1176
  auto exe_ctx = ExecutionContext(*this, scope, *dev_ctx, *runtime_ctx);
1177 1178 1179 1180
  // using cache
  if (kernel_type_.get()) {
    dev_ctx = pool.Get(kernel_type_->place_);
  }
1181 1182 1183 1184 1185 1186

  // TODO(chenweihang): Now we are still reusing a lot of the original fluid
  // implementation, this is a gradual replacement process
  // TODO(chenweihang): in the first phase of project, we only support CPU, CUDA
  // and RCOM backend, the XPU, NPU and MKLDNN will be supported in the second
  // phase
1187
  phi::KernelKey pt_kernel_key;
1188
  std::string pt_kernel_name;
1189
  if (phi::KernelFactory::Instance().HasCompatiblePtenKernel(type_)) {
1190
    if (pt_kernel_signature_ == nullptr || pt_kernel_ == nullptr) {
1191 1192
      pt_kernel_signature_.reset(
          new KernelSignature(std::move(GetExpectedPtenKernelArgs(exe_ctx))));
1193 1194 1195 1196 1197 1198 1199 1200 1201
      VLOG(6) << *pt_kernel_signature_.get();

      kernel_type_.reset(
          new OpKernelType(std::move(InnerGetExpectedKernelType(exe_ctx))));
      dev_ctx = pool.Get(kernel_type_->place_);

      pt_kernel_name = pt_kernel_signature_->name;
      pt_kernel_key = TransOpKernelTypeToPtenKernelKey(*kernel_type_.get());
      pt_kernel_.reset(
1202
          new phi::Kernel(phi::KernelFactory::Instance().SelectKernel(
1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232
              pt_kernel_name, pt_kernel_key)));

      if (pt_kernel_->IsValid()) {
        VLOG(6) << "Static mode ChoosePtenKernel - kernel name: "
                << pt_kernel_name << " | kernel key: " << pt_kernel_key
                << " | kernel: " << *pt_kernel_;
      } else {
        VLOG(6) << "Static mode ChoosePtenKernel - kernel `" << pt_kernel_name
                << "` not found.";
      }
    }
    if (pt_kernel_->IsValid()) {
      run_pten_kernel_ = true;
    } else {
      auto& all_op_kernels = AllOpKernels();
      auto kernels_iter = all_op_kernels.find(type_);
      if (kernels_iter == all_op_kernels.end() ||
          kernels_iter->second.find(*kernel_type_.get()) ==
              kernels_iter->second.end()
#ifdef PADDLE_WITH_XPU
          ||
          paddle::platform::is_xpu_place(kernel_type_->place_) &&  // NOLINT
              !paddle::platform::is_xpu_support_op(
                  type_, *kernel_type_.get())  // NOLINT
          || paddle::platform::is_in_xpu_black_list(type_)
#endif
              ) {
        auto pt_cpu_kernel_key =
            FallBackToCpu(*kernel_type_.get(), pt_kernel_key, *this);
        pt_kernel_.reset(
1233
            new phi::Kernel(phi::KernelFactory::Instance().SelectKernel(
1234 1235 1236 1237 1238 1239 1240 1241 1242 1243
                pt_kernel_name, pt_cpu_kernel_key)));

        dev_ctx = pool.Get(platform::CPUPlace());
        if (pt_kernel_->IsValid()) {
          VLOG(6) << "Static mode PrepareImpl - kernel name: " << pt_kernel_name
                  << " | kernel key: " << pt_cpu_kernel_key
                  << " | kernel: " << *pt_kernel_;
          run_pten_kernel_ = true;
        }
      }
1244 1245 1246 1247 1248
    }
  }
  if (!run_pten_kernel_) {
    if (kernel_type_.get() == nullptr || kernel_func_.get() == nullptr) {
      ChooseKernel(exe_ctx);
1249
      dev_ctx = pool.Get(kernel_type_->place_);
1250
    }
1251 1252
  }

Y
yuyang18 已提交
1253 1254
  // do data transformScope &transfer_scope;
  std::vector<std::string> transfered_inplace_vars;
1255 1256
  Scope* transfer_scope = nullptr;
  {
1257
    platform::RecordEvent record_event("prepare_data",
C
chenjian 已提交
1258 1259
                                       platform::TracerEventType::OperatorInner,
                                       1, platform::EventRole::kInnerOp);
1260 1261 1262 1263
    if (need_prepare_data_) {
      transfer_scope = PrepareData(scope, *kernel_type_,
                                   &transfered_inplace_vars, runtime_ctx);
    }
1264
  }
Y
yuyang18 已提交
1265 1266 1267 1268
  // exec scope is the scope that kernel actually executed on.
  const Scope& exec_scope =
      (transfer_scope == nullptr ? scope : *transfer_scope);

1269
  if (!all_kernels_must_compute_runtime_shape_) {
1270
    platform::RecordEvent record_event("infer_shape",
C
chenjian 已提交
1271 1272
                                       platform::TracerEventType::OperatorInner,
                                       1, platform::EventRole::kInnerOp);
1273
    RuntimeInferShapeContext infer_shape_ctx(*this, *runtime_ctx);
1274
    this->Info().infer_shape_(&infer_shape_ctx);
1275
  }
1276 1277 1278 1279 1280

  if (FLAGS_enable_unused_var_check) {
    GetThreadLocalUsedVarNameSet()->clear();
  }

X
clean  
Xin Pan 已提交
1281 1282
  // TODO(panyx0718): ExecutionContext should only depend on RuntimeContext
  // not Scope. Imperative mode only pass inputs and get outputs.
1283
  {
1284
    platform::RecordEvent record_event("compute",
C
chenjian 已提交
1285 1286
                                       platform::TracerEventType::OperatorInner,
                                       1, platform::EventRole::kInnerOp);
1287
    if (run_pten_kernel_) {
1288
      phi::KernelContext pt_kernel_context;
1289
      // Do data transform before building KernelContext
1290
      // TODO(zhiqiu): support TransferInplaceVarsBack
1291 1292
      PreparePtenData(exec_scope, *pt_kernel_, *pt_kernel_signature_,
                      runtime_ctx);
1293 1294
      BuildPtenKernelContext(*runtime_ctx, dev_ctx, &pt_kernel_context);
      (*pt_kernel_)(&pt_kernel_context);
1295 1296 1297 1298
    } else {
      (*kernel_func_)(
          ExecutionContext(*this, exec_scope, *dev_ctx, *runtime_ctx));
    }
1299
  }
D
dzhwinter 已提交
1300

Y
yuyang18 已提交
1301
  if (!transfered_inplace_vars.empty()) {
T
tianshuo78520a 已提交
1302
    // there is inplace variable has been transferred.
Y
yuyang18 已提交
1303
    TransferInplaceVarsBack(scope, transfered_inplace_vars, *transfer_scope);
1304
  }
1305 1306 1307 1308 1309 1310 1311

  // See [ Why need handle complex gradient to real gradient? ]
  // Only handle the case where the current kernel data type is complex
  if (framework::IsComplexType(kernel_type_->data_type_)) {
    HandleComplexGradToRealGrad(scope, runtime_ctx);
  }

1312 1313 1314 1315 1316 1317 1318 1319
  if (FLAGS_enable_unused_var_check) {
    // skip op that uses mkldnn because it has different memory reuse strategy.
    // use attr here because some GradMakers (like ActivationGradOpMaker) add
    // input when use_mkldnn=true;
    if (!(HasAttr("use_mkldnn") && Attr<bool>("use_mkldnn"))) {
      CheckUnusedVar(*this, scope);
    }
  }
1320

D
dzhwinter 已提交
1321
  /*For profiling/benchmark only*/
D
dzhwinter 已提交
1322
  if (FLAGS_benchmark) {
Y
yuyang18 已提交
1323
    dev_ctx->Wait();
1324 1325
#if defined(PADDLE_WITH_CUDA) || defined(PADLDE_WITH_ROCM)
    PADDLE_ENFORCE_GPU_SUCCESS(platform::GpuGetLastError());
1326 1327
#endif
    VLOG(4) << "Operator(" << Type() << "): context wait and get last error";
D
dzhwinter 已提交
1328
  }
C
chengduoZH 已提交
1329 1330

  if (FLAGS_check_nan_inf) {
W
WangXi 已提交
1331
    framework::details::CheckOpHasNanOrInf(*this, exec_scope, place);
C
chengduoZH 已提交
1332
  }
1333 1334 1335 1336 1337 1338 1339

  // To solve issue #15032, have a discussion with @Luotao for cpu inference,
  // do not cache transfer scope, hence in this case delete transfer scope
  // after run to avoid memory leak
  if (transfer_scope && !run_by_executor_ && !enable_cache_transfer_scope_) {
    scope.DeleteScope(transfer_scope);
  }
Q
Qiao Longfei 已提交
1340
}
X
Xin Pan 已提交
1341

1342 1343 1344
OpKernelType OperatorWithKernel::InnerGetExpectedKernelType(
    const ExecutionContext& ctx) const {
  auto expected_kernel_key = this->GetExpectedKernelType(ctx);
1345 1346 1347
  if (HasAttr("op_device")) {
    if (Attr<std::string>("op_device") == "cpu") {
      expected_kernel_key.place_ = platform::CPUPlace();
1348 1349 1350 1351 1352 1353 1354 1355 1356 1357
    } else if (Attr<std::string>("op_device").find("gpu") !=
               std::string::npos) {
      auto device = Attr<std::string>("op_device");
      size_t pos = device.find(':');
      if (pos != std::string::npos) {
        device = device.substr(0, pos);
        LOG_FIRST_N(WARNING, 1)
            << "Device index is only supported under pipeline parallelism, "
            << "so it will be ignored.";
      }
1358 1359
      // when the Op that only has CPUKernel is assigned to GPU, the CPUKernel
      // will be executed and a warning will be given at the same time.
1360 1361
      expected_kernel_key.place_ = platform::CPUPlace();
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
1362
      if (SupportGPU()) {
1363
        auto& dev_ctx = ctx.device_context();
1364
        expected_kernel_key.place_ = dev_ctx.GetPlace();
1365 1366 1367 1368 1369
      }
#endif
#ifdef PADDLE_WITH_ASCEND_CL
      if (SupportNPU()) {
        auto& dev_ctx = ctx.device_context();
1370
        expected_kernel_key.place_ = dev_ctx.GetPlace();
1371 1372 1373
      }
#endif
      if (platform::is_cpu_place(expected_kernel_key.place_)) {
1374 1375 1376 1377 1378 1379
        LOG_FIRST_N(WARNING, 1)
            << "Op(" << type_
            << ") has no CUDA implementation. It will be assigned to CPUPlace.";
      }
    }
  }
C
cc 已提交
1380 1381
  VLOG(3) << "op type:" << type_
          << ", expected_kernel_key:" << expected_kernel_key;
1382 1383 1384
  return expected_kernel_key;
}

1385
phi::KernelKey OperatorWithKernel::ChoosePtenKernel(
1386
    const ExecutionContext& ctx) const {
1387
  pt_kernel_signature_.reset(
1388
      new KernelSignature(std::move(GetExpectedPtenKernelArgs(ctx))));
1389
  VLOG(6) << *pt_kernel_signature_.get();
1390 1391 1392 1393

  kernel_type_.reset(
      new OpKernelType(std::move(InnerGetExpectedKernelType(ctx))));

Y
YuanRisheng 已提交
1394
  auto pt_kernel_name = pt_kernel_signature_->name;
1395
  auto pt_kernel_key = TransOpKernelTypeToPtenKernelKey(*kernel_type_.get());
1396 1397
  pt_kernel_.reset(new phi::Kernel(phi::KernelFactory::Instance().SelectKernel(
      pt_kernel_name, pt_kernel_key)));
1398 1399

  if (pt_kernel_->IsValid()) {
C
Chen Weihang 已提交
1400
    VLOG(6) << "Static mode ChoosePtenKernel - kernel name: " << pt_kernel_name
1401 1402 1403
            << " | kernel key: " << pt_kernel_key
            << " | kernel: " << *pt_kernel_;
  } else {
C
Chen Weihang 已提交
1404
    VLOG(6) << "Static mode ChoosePtenKernel - kernel `" << pt_kernel_name
1405 1406
            << "` not found.";
  }
1407
  return pt_kernel_key;
1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419 1420 1421 1422
}

void OperatorWithKernel::ChooseKernel(const ExecutionContext& ctx) const {
  // check if op[type] has kernel registered.
  auto& all_op_kernels = AllOpKernels();
  auto kernels_iter = all_op_kernels.find(type_);
  PADDLE_ENFORCE_NE(
      kernels_iter, all_op_kernels.end(),
      platform::errors::Unavailable(
          "There are no kernels which are registered in the %s operator.",
          type_));

  OpKernelMap& kernels = kernels_iter->second;

  auto expected_kernel_key = InnerGetExpectedKernelType(ctx);
L
Liu Yiqun 已提交
1423 1424

  auto kernel_iter = kernels.find(expected_kernel_key);
L
Liu-xiandong 已提交
1425

L
Liu Yiqun 已提交
1426 1427 1428 1429 1430 1431 1432 1433 1434
#ifdef PADDLE_WITH_MKLDNN
  // workaround for missing MKLDNN kernel when FLAGS_use_mkldnn env var is set
  if (kernel_iter == kernels.end() &&
      expected_kernel_key.library_type_ == LibraryType::kMKLDNN) {
    VLOG(3) << "missing MKLDNN kernel: fallbacking to PLAIN one";
    expected_kernel_key.library_type_ = LibraryType::kPlain;
    expected_kernel_key.data_layout_ = DataLayout::kAnyLayout;
    kernel_iter = kernels.find(expected_kernel_key);
  }
1435 1436
#endif
#ifdef PADDLE_WITH_XPU
1437
  if (platform::is_xpu_place(expected_kernel_key.place_) &&
Q
QingshuChen 已提交
1438 1439 1440
      (kernel_iter == kernels.end() ||
       !paddle::platform::is_xpu_support_op(type_, expected_kernel_key) ||
       paddle::platform::is_in_xpu_black_list(type_))) {
1441 1442 1443 1444 1445 1446
    VLOG(3) << "missing XPU kernel: " << type_
            << ", expected_kernel_key:" << expected_kernel_key
            << ", fallbacking to CPU one!";
    expected_kernel_key.place_ = platform::CPUPlace();
    kernel_iter = kernels.find(expected_kernel_key);
  }
1447
#endif
L
Liu-xiandong 已提交
1448 1449 1450 1451 1452 1453 1454 1455 1456 1457 1458 1459 1460 1461 1462 1463

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

A
Allen Guo 已提交
1464 1465 1466 1467 1468 1469 1470 1471 1472 1473
#ifdef PADDLE_WITH_IPU
  if (kernel_iter == kernels.end() &&
      platform::is_ipu_place(expected_kernel_key.place_)) {
    VLOG(3) << "missing IPU kernel: " << type_
            << ", expected_kernel_key:" << expected_kernel_key
            << ", fallbacking to CPU one!";
    expected_kernel_key.place_ = platform::CPUPlace();
    kernel_iter = kernels.find(expected_kernel_key);
  }
#endif
1474 1475
#ifdef PADDLE_WITH_ASCEND_CL
  if (kernel_iter == kernels.end() &&
1476
      platform::is_npu_place(expected_kernel_key.place_)) {
1477 1478 1479 1480 1481 1482
    VLOG(3) << "missing NPU kernel: " << type_
            << ", expected_kernel_key:" << expected_kernel_key
            << ", fallbacking to CPU one!";
    expected_kernel_key.place_ = platform::CPUPlace();
    kernel_iter = kernels.find(expected_kernel_key);
  }
F
fwenguang 已提交
1483 1484 1485
#endif
#ifdef PADDLE_WITH_MLU
  if (kernel_iter == kernels.end() &&
1486
      platform::is_mlu_place(expected_kernel_key.place_)) {
F
fwenguang 已提交
1487 1488 1489 1490 1491 1492
    VLOG(3) << "missing MLU kernel: " << type_
            << ", expected_kernel_key:" << expected_kernel_key
            << ", fallbacking to CPU one!";
    expected_kernel_key.place_ = platform::CPUPlace();
    kernel_iter = kernels.find(expected_kernel_key);
  }
L
Liu Yiqun 已提交
1493
#endif
1494 1495 1496 1497
  PADDLE_ENFORCE_NE(kernel_iter, kernels.end(),
                    platform::errors::NotFound(
                        "Operator (%s) does not have kernel for %s.", type_,
                        KernelTypeToString(expected_kernel_key)));
L
Liu Yiqun 已提交
1498

1499 1500 1501 1502 1503
  std::lock_guard<std::mutex> lock(cache_update_mutex_);
  if (kernel_type_.get() == nullptr || kernel_func_.get() == nullptr) {
    kernel_type_.reset(new OpKernelType(expected_kernel_key));
    kernel_func_.reset(new OpKernelFunc(kernel_iter->second));
  }
L
Liu Yiqun 已提交
1504 1505
}

Y
yuyang18 已提交
1506 1507 1508 1509
void OperatorWithKernel::TransferInplaceVarsBack(
    const Scope& scope, const std::vector<std::string>& inplace_vars,
    const Scope& transfer_scope) const {
  for (auto& var_name : inplace_vars) {
M
minqiyang 已提交
1510
    VLOG(3) << "share inplace var " + var_name + " back to it's original scope";
C
chengduo 已提交
1511
    auto* origin_var = scope.FindVar(var_name);
1512 1513 1514
    PADDLE_ENFORCE_NOT_NULL(origin_var,
                            platform::errors::InvalidArgument(
                                "The variable[%s] is nullptr.", var_name));
C
chengduo 已提交
1515
    auto* original_tensor =
C
chengduo 已提交
1516
        GetMutableLoDTensorOrSelectedRowsValueFromVar(origin_var);
C
chengduo 已提交
1517
    auto* var = transfer_scope.FindVar(var_name);
1518 1519
    PADDLE_ENFORCE_NOT_NULL(var, platform::errors::InvalidArgument(
                                     "The variable[%s] is nullptr.", var_name));
C
chengduo 已提交
1520
    auto* transformed_tensor = GetLoDTensorOrSelectedRowsValueFromVar(*var);
1521
    auto original_dims = original_tensor->dims();
Y
yuyang18 已提交
1522
    original_tensor->ShareDataWith(*transformed_tensor);
B
Baibaifan 已提交
1523 1524 1525 1526 1527
    // In order to solve the problem that the output latitude of NPU reshape
    // operator is not changed when inplace.
    if (type_ != "reshape2" && type_ != "reshape2_grad") {
      original_tensor->Resize(original_dims);
    }
Y
yuyang18 已提交
1528 1529 1530
  }
}

1531 1532 1533 1534 1535 1536 1537 1538 1539 1540 1541 1542 1543 1544 1545 1546 1547 1548 1549 1550 1551 1552 1553 1554 1555 1556 1557 1558 1559
void OperatorWithKernel::HandleComplexGradToRealGrad(
    const Scope& scope, RuntimeContext* ctx) const {
  for (auto& var_name_item : Outputs()) {
    std::vector<Variable*>& output_vars = ctx->outputs[var_name_item.first];
    for (size_t i = 0; i < var_name_item.second.size(); ++i) {
      // 1. find grad_var & check whether is complex tensor
      auto var_name = var_name_item.second[i];
      auto orig_var_name = GradOriginalVarName(var_name);
      // only focus on gradient var
      if (var_name == orig_var_name) {
        continue;
      }
      auto* grad_var = output_vars[i];
      // skip nullptr var
      if (grad_var == nullptr) {
        continue;
      }
      // don't process LoDTensorArray temporarily,
      // add support if necessary for complex number calculations in the future
      if (!VarIsTensor(*grad_var)) {
        continue;
      }
      auto* grad_tensor =
          GetMutableLoDTensorOrSelectedRowsValueFromVar(grad_var);
      // skip nullptr tensor
      if (grad_tensor == nullptr || !grad_tensor->IsInitialized()) {
        continue;
      }
      // only focus on complex dtype now
1560
      auto src_type = framework::TransToProtoVarType(grad_tensor->dtype());
1561 1562 1563 1564 1565 1566 1567 1568 1569 1570 1571 1572 1573 1574 1575 1576 1577 1578 1579
      if (!IsComplexType(src_type)) {
        continue;
      }

      // 2. find forward var & check whether need to cast
      auto* var = scope.FindVar(orig_var_name);
      // if forward var not exists, do nothing
      if (var == nullptr) {
        continue;
      }
      if (!VarIsTensor(*var)) {
        continue;
      }
      const auto* tensor = GetLoDTensorOrSelectedRowsValueFromVar(*var);
      PADDLE_ENFORCE_NOT_NULL(
          tensor,
          platform::errors::Unavailable(
              "Forward tensor is nullptr when handle complex data to real."));
      // only need record type, the allocation may have been released
1580
      auto dst_type = framework::TransToProtoVarType(tensor->dtype());
1581 1582 1583 1584 1585 1586 1587 1588 1589 1590 1591 1592 1593 1594 1595 1596 1597
      // only focus on real dtype and need casting
      if (IsComplexType(dst_type)) {
        continue;
      }

      // 3. cast complex grad to real grad
      VLOG(6) << "Transform " << framework::DataTypeToString(src_type)
              << " var `" << var_name << "` to "
              << framework::DataTypeToString(dst_type)
              << " real var in static graph.";
      Tensor out;
      TransComplexToReal(dst_type, src_type, *grad_tensor, &out);
      SetTensorToVariable(*grad_var, out, grad_var);
    }
  }
}

X
Xin Pan 已提交
1598
Scope* OperatorWithKernel::PrepareData(
Y
yuyang18 已提交
1599
    const Scope& scope, const OpKernelType& expected_kernel_key,
X
Xin Pan 已提交
1600 1601
    std::vector<std::string>* transfered_inplace_vars,
    RuntimeContext* ctx) const {
Y
yuyang18 已提交
1602
  Scope* new_scope = nullptr;
S
sneaxiy 已提交
1603

1604
  const std::unordered_set<std::string>* no_buffer_ins = nullptr;
S
sneaxiy 已提交
1605 1606 1607 1608
  if (info_) {
    auto& no_buffer_inferer = info_->NoNeedBufferVarsInferer();
    // Some op may not register NoNeedBufferVarsInferer
    if (no_buffer_inferer) {
1609 1610
      no_buffer_ins = &(no_buffer_inferer(Inputs(), Outputs(), Attrs()));
      if (no_buffer_ins->empty()) no_buffer_ins = nullptr;
S
sneaxiy 已提交
1611 1612 1613
    }
  }

Y
yuyang18 已提交
1614
  for (auto& var_name_item : Inputs()) {
1615 1616
    bool should_skip_input =
        no_buffer_ins && no_buffer_ins->count(var_name_item.first) > 0;
S
sneaxiy 已提交
1617

X
Xin Pan 已提交
1618 1619 1620 1621
    std::vector<Variable*>& input_vars = ctx->inputs[var_name_item.first];

    for (size_t i = 0; i < var_name_item.second.size(); ++i) {
      auto& var_name = var_name_item.second[i];
X
Xin Pan 已提交
1622
      auto* var = input_vars[i];
X
Xin Pan 已提交
1623

Y
yuyang18 已提交
1624
      // Only tensor can be tranfer to another device.
C
chengduo 已提交
1625
      if (var == nullptr || !VarIsTensor(*var)) {
Y
yuyang18 已提交
1626 1627 1628
        continue;
      }

C
chengduo 已提交
1629
      auto* tensor_in = GetLoDTensorOrSelectedRowsValueFromVar(*var);
1630 1631 1632 1633 1634 1635 1636 1637 1638 1639 1640 1641 1642 1643 1644

      // When no_buffer_ins then checking of Tensor::holder_ is
      // not a thread safe. And for infershape scenario checks
      // to be omitted are not really needed
      if (should_skip_input == true) {
#ifdef PADDLE_WITH_MKLDNN
        // Var without buffer may be needed
        // for some situation like InferShape().
        // In this situation We cannot skip Var analysis, as
        // MKL-DNN shape of Var may differ from kNHWC Var
        // In such situation corressponding resized Var
        // has to be created and registered
        if ((tensor_in->layout() == DataLayout::kMKLDNN) &&
            (var->IsType<LoDTensor>() == true) &&
            (expected_kernel_key.data_layout_ != DataLayout::kMKLDNN) &&
1645 1646
            (paddle::platform::MKLDNNDeviceContext::tls()
                 .get_cur_paddle_data_layout() == DataLayout::kNHWC)) {
1647 1648 1649 1650 1651 1652 1653 1654 1655 1656 1657 1658 1659 1660 1661 1662 1663 1664 1665 1666 1667
          // Mixed execution : MKL-DNN and GPU is not supported!
          if (!new_scope) {
            new_scope = &scope.NewScope();
          }
          auto* trans_var = new_scope->Var(var_name);
          input_vars[i] = trans_var;
          auto out = trans_var->GetMutable<LoDTensor>();
          out->Resize(tensor_in->dims());
          platform::MatchShapeToLayout(out, tensor_in->layout(),
                                       DataLayout::kNHWC);
          VLOG(7) << "Created reshaped dummy input based on MKL-DNN Tensor , "
                     "but kNHWC layout"
                  << var_name_item.first << " in Operator " << type_;
        } else {
          VLOG(7) << "Skip scanning input " << var_name_item.first
                  << " in Operator " << type_;
        }
#endif
        continue;
      }

Y
yuyang18 已提交
1668 1669 1670 1671 1672 1673 1674 1675 1676 1677 1678
      if (!tensor_in->IsInitialized()) {
        continue;
      }

      auto kernel_type_for_var = GetKernelTypeForVar(
          var_name_item.first, *tensor_in, expected_kernel_key);

      if (!NeedTransform(kernel_type_for_var, expected_kernel_key)) {
        continue;
      }

M
minqiyang 已提交
1679 1680
      VLOG(3) << "Transform Variable " << var_name << " from "
              << kernel_type_for_var << " to " << expected_kernel_key;
Y
yuyang18 已提交
1681

1682 1683 1684
      // In the inference scenerio, the scopes will be reused across the
      // batches, so the `new_scope` here will result in GPU memroy explosion
      // over the  running of operators.
1685
      // We use a thread_local cache to fix that issue, the key in the cache is
1686 1687 1688 1689 1690
      // the combination of the `scope` argument, from_kernel_type,
      // target_kernel_type.
      // Have a discussion with @Superjomn or the inference developers if some
      // changes on this logic for this macro might not tested on the other
      // scenerios.
1691 1692
      // If this op is not called by an Executor or ParallelExecutor, it should
      // called by a NaiveExecutor, the NaiveExecutor will cache the scopes and
1693
      // variables, that behavior a lot different.
1694 1695 1696 1697 1698 1699 1700 1701 1702
      //
      // To solve issue #15032, have a discussion with @Luotao for cpu
      // inference, for all cpu kernels cases without GPU participation, here
      // not do transfer scope caching, and cpu inference performance is not
      // impacted by test.
      enable_cache_transfer_scope_ = false;
      if (!run_by_executor_ &&
          (platform::is_gpu_place(kernel_type_for_var.place_) ||
           platform::is_gpu_place(expected_kernel_key.place_))) {
1703 1704
        new_scope = TryCreateTransferScope(kernel_type_for_var,
                                           expected_kernel_key, &scope);
1705
        enable_cache_transfer_scope_ = true;
1706
      }
1707
      if (!new_scope) {
Y
yuyang18 已提交
1708 1709
        new_scope = &scope.NewScope();
      }
1710 1711 1712 1713
      // For inference, if a gpu model has an op which could only run on CPU,
      // each result of different input will be the same with the first one.
      // The reason is that if a gpu tensor is the input of a cpu kernel,
      // we will create a new cpu tensor in new scope.
1714
      // However, if enable_cache_runtime_context_, we get the cpu tensor each
1715 1716
      // time, not the gpu tensor. Thus, we set pre_scope_ = nullptr
      // to trigger `new RuntimeContext()` in RunImpl().
1717
      if (enable_cache_runtime_context_) {
1718 1719
        pre_scope_ = nullptr;
      }
L
Leo Chen 已提交
1720 1721

      // Create new var with the same name in transfer scopes
Y
yuyang18 已提交
1722
      auto* trans_var = new_scope->Var(var_name);
X
fix  
Xin Pan 已提交
1723
      input_vars[i] = trans_var;
L
Leo Chen 已提交
1724 1725 1726 1727 1728 1729 1730 1731 1732 1733 1734 1735 1736 1737 1738 1739 1740

      // Find if inplace exists between input and output
      // If inplace exists, set the new created var to inplaced output, and
      // record its name in transfered_inplace_vars.
      for (auto& pair : Outputs()) {
        for (size_t j = 0; j < pair.second.size(); ++j) {
          if (pair.second[j] == var_name) {
            VLOG(4) << "Found inplace between input(" << var_name_item.first
                    << ") and output(" << pair.first
                    << "), the variable name is " << var_name;
            ctx->outputs[pair.first][j] = trans_var;
            transfered_inplace_vars->emplace_back(var_name);
          }
        }
      }

      // Do transfer
Y
yuyang18 已提交
1741
      Tensor out;
Y
yuyang18 已提交
1742
      TransformData(expected_kernel_key, kernel_type_for_var, *tensor_in, &out);
Y
yuyang18 已提交
1743 1744 1745
      SetTensorToVariable(*var, out, trans_var);
    }
  }
L
Leo Chen 已提交
1746

1747 1748 1749 1750 1751 1752
  // If pre_scope = &scope, it means that scope is cached and the op is not in
  // while block. If new_scope = nullptr, it means that for each input of this
  // Op, there is no need to do PrepareData. So PrepareData could be skipped at
  // the rest iterations to save the elapsed time.
  // We do not support skipping PrepareData in while block, because the Op's
  // input may be changed by subsequent Ops, which may cause an error.
W
wenbin 已提交
1753 1754 1755 1756 1757 1758

  // For inference, ops that behind conditional branch aren't supported well,
  // so disable prepare optimization conservatively.
  bool force_prepare_data = HasAttr("inference_force_prepare_data") &&
                            Attr<bool>("inference_force_prepare_data");
  if (pre_scope_ == &scope && new_scope == nullptr && !force_prepare_data) {
1759 1760
    need_prepare_data_ = false;
  }
Y
yuyang18 已提交
1761 1762 1763

  return new_scope;
}
Q
Qiao Longfei 已提交
1764

1765
void OperatorWithKernel::ParseInputDataType(
1766
    const std::vector<Variable*>& vars, const std::string& name,
1767
    proto::VarType::Type* data_type) const {
1768
  proto::VarType::Type default_data_type =
1769 1770 1771 1772 1773 1774 1775 1776 1777
      static_cast<proto::VarType::Type>(-1);
  for (size_t i = 0; i < vars.size(); ++i) {
    const Variable* var = vars[i];
    if (var != nullptr) {
      const Tensor* t = nullptr;
      if (var->IsType<Tensor>()) {
        t = &var->Get<Tensor>();
      } else if (var->IsType<LoDTensor>()) {
        t = &var->Get<LoDTensor>();
1778 1779
      } else if (var->IsType<phi::SelectedRows>()) {
        t = &(var->Get<phi::SelectedRows>().value());
1780
      } else if (var->IsType<LoDTensorArray>()) {
1781 1782 1783 1784
        auto t_arr = &var->Get<LoDTensorArray>();
        for (size_t j = 0; j < t_arr->size(); j++) {
          if (t_arr->at(j).IsInitialized()) {
            t = &(t_arr->at(j));
1785 1786
          }
        }
1787 1788
      }
      if (t != nullptr) {
1789 1790
        PADDLE_ENFORCE_EQ(
            t->IsInitialized(), true,
1791 1792 1793
            platform::errors::InvalidArgument("The %s Op's Input Variable `%s` "
                                              "contains uninitialized Tensor.",
                                              Type(), name));
1794 1795
        proto::VarType::Type tmp =
            paddle::framework::TransToProtoVarType(t->dtype());
1796 1797 1798 1799 1800 1801 1802 1803 1804
        PADDLE_ENFORCE(tmp == *data_type || *data_type == default_data_type,
                       platform::errors::InvalidArgument(
                           "The DataType of %s Op's duplicable or different "
                           "slot Variable %s must be "
                           "consistent or reigster GetExpectedKernelType. The "
                           "current variable type is (%s), but the "
                           "previous variable type is (%s).",
                           Type(), name, DataTypeToString(tmp),
                           DataTypeToString(*data_type)));
1805 1806 1807 1808 1809 1810
        *data_type = tmp;
      }
    }
  }
}

1811
proto::VarType::Type OperatorWithKernel::IndicateDataType(
Y
Yu Yang 已提交
1812
    const ExecutionContext& ctx) const {
1813 1814 1815
  proto::VarType::Type dafault_data_type =
      static_cast<proto::VarType::Type>(-1);
  proto::VarType::Type data_type = dafault_data_type;
H
hong 已提交
1816
  for (auto& input : ctx.InNameList()) {
1817 1818
    const std::vector<Variable*> vars = ctx.MultiInputVar(input);
    ParseInputDataType(vars, input, &data_type);
Y
Yu Yang 已提交
1819
  }
1820 1821 1822 1823
  PADDLE_ENFORCE_NE(
      data_type, dafault_data_type,
      platform::errors::NotFound(
          "DataType should be indicated by input Variable at %s.", Type()));
1824 1825 1826 1827 1828 1829 1830 1831
  return data_type;
}

proto::VarType::Type OperatorWithKernel::IndicateVarDataType(
    const ExecutionContext& ctx, const std::string& name) const {
  proto::VarType::Type dafault_data_type =
      static_cast<proto::VarType::Type>(-1);
  proto::VarType::Type data_type = dafault_data_type;
1832
  ParseInputDataType(ctx.MultiInputVar(name), name, &data_type);
1833 1834
  PADDLE_ENFORCE_NE(
      data_type, dafault_data_type,
1835 1836 1837 1838 1839
      platform::errors::InvalidArgument(
          "The Input Variable(%s) of (%s) Operator used to determine kernel "
          "data type is empty or not LoDTensor or SelectedRows or "
          "LoDTensorArray.",
          name, Type()));
1840
  return data_type;
Y
Yu Yang 已提交
1841
}
1842

1843 1844 1845 1846 1847 1848 1849 1850 1851 1852 1853 1854 1855 1856 1857 1858 1859 1860
Tensor* OperatorWithKernel::GetTensorFormInputSafely(
    const ExecutionContext& ctx, const std::string& name) const {
  // 1. get variable and check
  // NOTE: only supports signal input var now
  // NOTE: using const_cast is because we don't have method
  // can get single mutable var, and here will not change
  // the var's data, only use some attribute
  Variable* var = const_cast<Variable*>(ctx.InputVar(name));
  PADDLE_ENFORCE_NOT_NULL(
      var,
      platform::errors::NotFound(
          "The variable %s is not found when promote complex types.", name));
  // 2. get tensor and check
  Tensor* t = nullptr;
  if (var->IsType<Tensor>()) {
    t = var->GetMutable<Tensor>();
  } else if (var->IsType<LoDTensor>()) {
    t = var->GetMutable<LoDTensor>();
1861 1862
  } else if (var->IsType<phi::SelectedRows>()) {
    t = var->GetMutable<phi::SelectedRows>()->mutable_value();
1863 1864 1865 1866 1867 1868 1869 1870 1871 1872 1873 1874 1875 1876 1877 1878 1879 1880 1881 1882 1883 1884 1885 1886 1887 1888 1889 1890 1891 1892 1893
  } else {
    PADDLE_THROW(platform::errors::Unimplemented(
        "Unsupported input variable type in complex type promotion."));
  }
  PADDLE_ENFORCE_NOT_NULL(
      t,
      platform::errors::InvalidArgument(
          "The Tensor of variable %s is nullptr when promote complex types."));
  PADDLE_ENFORCE_EQ(t->IsInitialized(), true,
                    platform::errors::InvalidArgument(
                        "The Tensor in the %s Op's Input Variable %s(%s) is "
                        "not initialized.",
                        Type(), name, ctx.InputName(name)));
  return t;
}

/** NOTE(chenweihang): For safety reasons, we now only
 * perform type promotes for binary operations with
 * complex type inputs, which is used to support the
 * paddle quantum function.
 * In other cases, the first input data type is used as
 * the kernel data type.
 */
proto::VarType::Type OperatorWithKernel::IndicateOrPromoteVarDataTypes(
    const ExecutionContext& ctx, const std::string& name1,
    const std::string& name2) const {
  // 1. Get tensor
  auto* tensor_a = GetTensorFormInputSafely(ctx, name1);
  auto* tensor_b = GetTensorFormInputSafely(ctx, name2);

  // 2. Get two input types
1894 1895
  auto type_a = framework::TransToProtoVarType(tensor_a->dtype());
  auto type_b = framework::TransToProtoVarType(tensor_b->dtype());
1896 1897 1898 1899 1900 1901 1902

  // 3. Get first input type or promote complex types
  auto target_type = PromoteTypesIfComplexExists(type_a, type_b);

  return target_type;
}

1903 1904 1905 1906 1907 1908 1909 1910
OpKernelType OperatorWithKernel::GetExpectedKernelType(
    const ExecutionContext& ctx) const {
  return OpKernelType(IndicateDataType(ctx), ctx.GetPlace());
}

OpKernelType OperatorWithKernel::GetKernelTypeForVar(
    const std::string& var_name, const Tensor& tensor,
    const OpKernelType& expected_kernel_type) const {
M
mozga-intel 已提交
1911 1912
  return OpKernelType(expected_kernel_type.data_type_, tensor.place(),
                      tensor.layout());
1913 1914
}

1915 1916
KernelSignature OperatorWithKernel::GetExpectedPtenKernelArgs(
    const ExecutionContext& ctx) const {
1917 1918
  InitDefaultKernelSignatureMap();
  ExecutionArgumentMappingContext arg_mapping_ctx(ctx);
1919
  return phi::OpUtilsMap::Instance().GetArgumentMappingFn(Type())(
1920
      arg_mapping_ctx);
1921 1922
}

1923
Scope* OperatorWithKernel::PreparePtenData(
1924
    const Scope& scope, const phi::Kernel& pt_kernel,
1925 1926 1927 1928 1929 1930 1931 1932 1933
    const KernelSignature& pt_kernel_signature, RuntimeContext* ctx) const {
  auto& input_names = std::get<0>(pt_kernel_signature.args);
  auto input_defs = pt_kernel.args_def().input_defs();
  PADDLE_ENFORCE_EQ(input_names.size(), input_defs.size(),
                    platform::errors::InvalidArgument(
                        "The size of inputs_args names (%d) must be equal to "
                        "the size of kernel input_defs (%d).",
                        input_names.size(), input_defs.size()));
  Scope* new_scope = nullptr;
1934
  auto& name_map = Inputs();
Y
YuanRisheng 已提交
1935 1936 1937 1938 1939 1940 1941 1942 1943 1944
  const std::unordered_set<std::string>* no_buffer_ins = nullptr;
  if (info_) {
    auto& no_buffer_inferer = info_->NoNeedBufferVarsInferer();
    // Some op may not register NoNeedBufferVarsInferer
    if (no_buffer_inferer) {
      no_buffer_ins = &(no_buffer_inferer(Inputs(), Outputs(), Attrs()));
      if (no_buffer_ins->empty()) no_buffer_ins = nullptr;
    }
  }

1945 1946
  for (size_t i = 0; i < input_defs.size(); ++i) {
    auto& in_def = input_defs.at(i);
1947
    if (ctx->inputs.find(input_names[i]) == ctx->inputs.end()) {
H
hong 已提交
1948 1949
      continue;
    }
1950
    auto& ins_vector = ctx->inputs.at(input_names[i]);
1951
    auto& name_vec = name_map.at(input_names[i]);
Y
YuanRisheng 已提交
1952 1953 1954
    bool should_skip_input =
        no_buffer_ins && no_buffer_ins->count(input_names[i]) > 0;

1955 1956 1957 1958 1959 1960 1961
    for (size_t offset = 0; offset < ins_vector.size(); ++offset) {
      // Only tensor can be tranfer to another device.
      auto* var = ins_vector[offset];
      if (var == nullptr || !VarIsTensor(*var)) {
        continue;
      }
      auto* tensor_in = GetLoDTensorOrSelectedRowsValueFromVar(*var);
Y
YuanRisheng 已提交
1962 1963 1964 1965 1966 1967 1968 1969 1970

      // When no_buffer_ins then checking of Tensor::holder_ is
      // not a thread safe. And for infershape scenario checks
      // to be omitted are not really needed
      if (should_skip_input == true) {
        // TODO(YuanRisheng) : There need to supplement MKLDNN code later
        continue;
      }

1971 1972 1973 1974
      if (!tensor_in->IsInitialized()) {
        continue;
      }

1975 1976 1977
      if (in_def.backend == phi::Backend::ALL_BACKEND) {
        continue;
      }
1978
      auto expected_place = phi::TransToPtenPlace(in_def.backend);
1979 1980 1981 1982
      if (platform::is_same_place(tensor_in->place(), expected_place)) {
        continue;
      }

1983 1984
      VLOG(3) << "PTen Transform Variable " << input_names[i] << " from "
              << tensor_in->place() << " to " << expected_place;
1985

1986 1987 1988
      if (!new_scope) {
        new_scope = &scope.NewScope();
      }
1989

1990
      // Create new var with the same name in transfer scopes
1991
      auto* trans_var = new_scope->Var(name_vec[offset]);
1992
      ins_vector[offset] = trans_var;
1993

1994 1995 1996 1997
      // Do transfer
      Tensor out;
      framework::TensorCopySync(*tensor_in, expected_place, &out);
      SetTensorToVariable(*var, out, trans_var);
1998 1999 2000 2001 2002 2003
    }
  }

  return new_scope;
}

2004
void OperatorWithKernel::BuildPtenKernelContext(
2005
    const RuntimeContext& ctx, platform::DeviceContext* dev_ctx,
2006
    phi::KernelContext* pt_kernel_context) const {
2007
  pt_kernel_context->SetDeviceContext(dev_ctx);
2008 2009 2010 2011 2012 2013 2014 2015 2016 2017 2018 2019 2020 2021 2022 2023 2024 2025 2026 2027 2028 2029 2030 2031 2032 2033 2034 2035

  auto& input_names = std::get<0>(pt_kernel_signature_->args);
  auto& attr_names = std::get<1>(pt_kernel_signature_->args);
  auto& output_names = std::get<2>(pt_kernel_signature_->args);

  auto input_defs = pt_kernel_->args_def().input_defs();
  auto attr_defs = pt_kernel_->args_def().attribute_defs();
  auto output_defs = pt_kernel_->args_def().output_defs();

  PADDLE_ENFORCE_EQ(input_names.size(), input_defs.size(),
                    platform::errors::InvalidArgument(
                        "The size of inputs_args names (%d) must be equal to "
                        "the size of kernel input_defs (%d).",
                        input_names.size(), input_defs.size()));

  PADDLE_ENFORCE_EQ(output_names.size(), output_defs.size(),
                    platform::errors::InvalidArgument(
                        "The size of outputs_args names (%d) must be equal to "
                        "the size of kernel output_defs (%d).",
                        output_names.size(), output_defs.size()));

  PADDLE_ENFORCE_EQ(attr_names.size(), attr_defs.size(),
                    platform::errors::InvalidArgument(
                        "The size of attribute_args names (%d) must be equal "
                        "to the size of kernel attribute_defs (%d).",
                        attr_names.size(), attr_defs.size()));

  for (size_t i = 0; i < input_names.size(); ++i) {
H
hong 已提交
2036
    auto it = ctx.inputs.find(input_names[i]);
2037 2038 2039

    // calcute the start and end index of the input tensors
    size_t start_idx =
2040
        (i == 0 ? 0 : pt_kernel_context->InputRangeAt(i - 1).second);
2041

H
hong 已提交
2042
    // deal with optional here
2043
    if ((it == ctx.inputs.end() || it->second.size() == 0) &&
H
hong 已提交
2044
        (input_defs[i].type_index ==
2045
         std::type_index(typeid(paddle::optional<const phi::DenseTensor&>)))) {
H
hong 已提交
2046 2047 2048 2049 2050 2051 2052 2053
      pt_kernel_context->EmplaceBackInputWithoutSetRange(nullptr);
      auto end_idx = start_idx + 1;
      pt_kernel_context->AssignInputRange(std::make_pair(start_idx, end_idx),
                                          i);
      continue;
    }
    auto ins_vector = it->second;
    size_t end_idx = start_idx + ins_vector.size();
2054
    for (size_t offset = 0; offset < ins_vector.size(); ++offset) {
2055
      const phi::TensorBase* tensor_in = nullptr;
2056
      auto* var = ins_vector[offset];
H
hong 已提交
2057 2058
      if (var->IsType<framework::LoDTensor>()) {
        tensor_in = &(var->Get<framework::LoDTensor>());
2059 2060
      } else if (var->IsType<phi::SelectedRows>()) {
        tensor_in = &(var->Get<phi::SelectedRows>());
2061 2062 2063 2064
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
            "Unsupported input `%s` type when call pt kernel.",
            framework::ToTypeName(var->Type())));
2065
      }
H
hong 已提交
2066

2067
      pt_kernel_context->EmplaceBackInputWithoutSetRange(tensor_in);
2068
    }
2069
    pt_kernel_context->AssignInputRange(std::make_pair(start_idx, end_idx), i);
2070 2071 2072
  }

  for (size_t i = 0; i < output_names.size(); ++i) {
H
hong 已提交
2073
    auto it = ctx.outputs.find(output_names[i]);
2074
    size_t start_idx =
2075
        (i == 0 ? 0 : pt_kernel_context->OutputRangeAt(i - 1).second);
H
hong 已提交
2076 2077 2078 2079 2080 2081 2082 2083 2084 2085 2086 2087 2088 2089

    if (it == ctx.outputs.end() || it->second.empty()) {
      // Deal with the case that some outputs are not found or be NULL when run
      // the kernel.
      // For example : the outputs of matmul_grad are dx and dy,
      // sometimes dx or dy may be NULL.
      pt_kernel_context->EmplaceBackOutputWithoutSetRange(nullptr);
      auto end_idx = start_idx + 1;
      pt_kernel_context->AssignOutputRange(std::make_pair(start_idx, end_idx),
                                           i);
      continue;
    }
    auto& outs_vector = it->second;

2090
    size_t end_idx = start_idx + outs_vector.size();
2091 2092

    for (size_t offset = 0; offset < outs_vector.size(); ++offset) {
2093
      phi::TensorBase* tensor_out = nullptr;
2094
      auto* var = outs_vector[offset];
H
hong 已提交
2095 2096
      if (var->template IsType<framework::LoDTensor>()) {
        tensor_out = var->template GetMutable<framework::LoDTensor>();
2097 2098
      } else if (var->template IsType<phi::SelectedRows>()) {
        tensor_out = var->template GetMutable<phi::SelectedRows>();
2099 2100 2101 2102
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
            "Unsupported output `%s` type when call pt kernel.",
            framework::ToTypeName(var->Type())));
2103
      }
2104

2105 2106
      experimental::ResetTensorDtypeAndLayoutByArgDef(tensor_out,
                                                      output_defs.at(i));
2107
      SetAllocationForOutputTenosr(
2108
          tensor_out, phi::TransToPtenPlace(output_defs.at(i).backend));
2109 2110

      pt_kernel_context->EmplaceBackOutputWithoutSetRange(tensor_out);
2111
    }
2112

2113
    pt_kernel_context->AssignOutputRange(std::make_pair(start_idx, end_idx), i);
2114 2115 2116
  }

  for (size_t i = 0; i < attr_names.size(); ++i) {
2117
    if (attr_defs[i].type_index == std::type_index(typeid(phi::ScalarArray))) {
2118 2119 2120 2121
      auto attr_iter = Attrs().find(attr_names[i]);
      if (attr_iter != Attrs().end()) {  // shape is in the attribute
        if (std::type_index(attr_iter->second.type()) ==
            std::type_index(typeid(std::vector<int64_t>))) {
2122
          pt_kernel_context->EmplaceBackAttr(std::move(phi::ScalarArray(
2123
              BOOST_GET_CONST(std::vector<int64_t>, attr_iter->second))));
2124 2125
        } else if (std::type_index(attr_iter->second.type()) ==
                   std::type_index(typeid(std::vector<int32_t>))) {
2126
          pt_kernel_context->EmplaceBackAttr(std::move(phi::ScalarArray(
2127
              BOOST_GET_CONST(std::vector<int32_t>, attr_iter->second))));
C
chentianyu03 已提交
2128 2129
        } else if (std::type_index(attr_iter->second.type()) ==
                   std::type_index(typeid(int32_t))) {
2130
          pt_kernel_context->EmplaceBackAttr(std::move(phi::ScalarArray(
C
chentianyu03 已提交
2131
              &BOOST_GET_CONST(int32_t, attr_iter->second), 1)));
2132 2133 2134 2135 2136 2137 2138 2139 2140
        } else {
          PADDLE_THROW(platform::errors::Unimplemented(
              "Unsupported cast op attribute `%s` to ScalarArray when "
              "construct KernelContext.",
              attr_names[i]));
        }
      } else {  // shape is in the input
        auto& ins_vector = ctx.inputs.at(attr_names[i]);
        if (ins_vector.size() == 1) {  // ShapeTensor
2141
          pt_kernel_context->EmplaceBackAttr(std::move(
2142 2143
              experimental::MakePtenScalarArrayFromVar(*ins_vector.front())));
        } else {  // ShapeTensorList
2144
          pt_kernel_context->EmplaceBackAttr(std::move(
2145 2146 2147 2148
              experimental::MakePtenScalarArrayFromVarList(ins_vector)));
        }
      }
    } else if (attr_defs[i].type_index ==
2149
               std::type_index(typeid(phi::Scalar))) {
2150 2151 2152
      // TODO(chenweihang): support other attrs later
      // TODO(zhangyunfei): Scalar should hold scaler type, and we should check
      // attribtue type by attr_defs
2153 2154 2155 2156
      auto attr_iter = Attrs().find(attr_names[i]);
      if (attr_iter != Attrs().end()) {  // scalar is in the attribute
        auto& attr = Attrs().at(attr_names[i]);
        if (std::type_index(attr.type()) == std::type_index(typeid(float))) {
2157
          pt_kernel_context->EmplaceBackAttr(
2158
              std::move(phi::Scalar(BOOST_GET_CONST(float, attr))));
2159 2160
        } else if (std::type_index(attr.type()) ==
                   std::type_index(typeid(std::string))) {
2161
          pt_kernel_context->EmplaceBackAttr(
2162
              std::move(phi::Scalar(BOOST_GET_CONST(std::string, attr))));
2163 2164 2165
        } else if (std::type_index(attr.type()) ==
                   std::type_index(typeid(int))) {
          pt_kernel_context->EmplaceBackAttr(
2166
              std::move(phi::Scalar(BOOST_GET_CONST(int, attr))));
2167 2168 2169 2170 2171 2172
        } else {
          PADDLE_THROW(platform::errors::Unimplemented(
              "Unsupported cast op attribute `%s` to Scalar when construct "
              "KernelContext.",
              attr_names[i]));
        }
2173
      } else {
2174
        auto& ins_vector = ctx.inputs.at(attr_names[i]);
2175
        pt_kernel_context->EmplaceBackAttr(std::move(
2176
            experimental::MakePtenScalarFromVar(*ins_vector.front())));
2177
      }
2178

2179 2180
    } else {
      // TODO(chenweihang): support other attrs later
2181
      auto& attr = Attrs().at(attr_names[i]);
2182
      if (attr_defs[i].type_index == std::type_index(typeid(int))) {
2183
        pt_kernel_context->EmplaceBackAttr(BOOST_GET_CONST(int, attr));
2184
      } else if (attr_defs[i].type_index == std::type_index(typeid(float))) {
2185
        pt_kernel_context->EmplaceBackAttr(BOOST_GET_CONST(float, attr));
2186
      } else if (attr_defs[i].type_index == std::type_index(typeid(bool))) {
2187
        pt_kernel_context->EmplaceBackAttr(BOOST_GET_CONST(bool, attr));
H
hong 已提交
2188 2189
      } else if (attr_defs[i].type_index == std::type_index(typeid(int64_t))) {
        pt_kernel_context->EmplaceBackAttr(BOOST_GET_CONST(int64_t, attr));
H
hong 已提交
2190 2191 2192
      } else if (attr_defs[i].type_index ==
                 std::type_index(typeid(std::string))) {
        pt_kernel_context->EmplaceBackAttr(BOOST_GET_CONST(std::string, attr));
2193
      } else if (attr_defs[i].type_index ==
2194
                 std::type_index(typeid(phi::DataType))) {
2195
        auto data_type = paddle::framework::TransToPtenDataType(
2196 2197
            static_cast<framework::proto::VarType::Type>(
                BOOST_GET_CONST(int, attr)));
2198
        pt_kernel_context->EmplaceBackAttr(data_type);
2199 2200 2201 2202 2203 2204 2205 2206
      } else if (attr_defs[i].type_index ==
                 std::type_index(typeid(std::vector<int64_t>))) {
        if (std::type_index(attr.type()) ==
            std::type_index(typeid(std::vector<int>))) {
          // Emplace Back Attr according to the type of Pten_Kernel args.
          const auto& vector_int_attr = BOOST_GET_CONST(std::vector<int>, attr);
          const std::vector<int64_t> vector_int64_attr(vector_int_attr.begin(),
                                                       vector_int_attr.end());
2207
          pt_kernel_context->EmplaceBackAttr(vector_int64_attr);
2208 2209 2210
        }
        // TODO(YuanRisheng) Need support vector<int64_t> attr

H
hong 已提交
2211 2212 2213 2214
      } else if (attr_defs[i].type_index ==
                 std::type_index(typeid(std::vector<int32_t>))) {
        const auto& vector_int_attr = BOOST_GET_CONST(std::vector<int>, attr);
        pt_kernel_context->EmplaceBackAttr(vector_int_attr);
2215 2216
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
2217
            "Unsupported cast op attribute `%s` when construct "
2218 2219 2220 2221 2222 2223 2224
            "KernelContext.",
            attr_names[i]));
      }
    }
  }
}

Q
Qiao Longfei 已提交
2225
}  // namespace framework
L
liaogang 已提交
2226
}  // namespace paddle