operator.cc 83.2 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
      platform::RecordEvent op_type_record_event(
C
chenjian 已提交
266 267 268 269 270
          Type(), platform::TracerEventType::Operator, 1);
      // auto op_name = platform::OpName(outputs_, Type());
      // platform::RecordEvent op_name_record_event(
      //     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
              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.";
      }
    }
1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224
#ifdef PADDLE_WITH_XPU
    bool is_xpu_unsupport =
        paddle::platform::is_xpu_place(kernel_type_->place_) &&
            !paddle::platform::is_xpu_support_op(type_, *kernel_type_.get()) ||
        paddle::platform::is_in_xpu_black_list(type_);
#endif
    if (pt_kernel_->IsValid()
#ifdef PADDLE_WITH_XPU
        && !is_xpu_unsupport
#endif
        ) {
1225 1226 1227 1228 1229 1230 1231 1232
      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
1233
          || is_xpu_unsupport
1234
#endif
1235
          ) {
1236 1237 1238
        auto pt_cpu_kernel_key =
            FallBackToCpu(*kernel_type_.get(), pt_kernel_key, *this);
        pt_kernel_.reset(
1239
            new phi::Kernel(phi::KernelFactory::Instance().SelectKernel(
1240 1241 1242 1243 1244 1245 1246 1247 1248 1249
                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;
        }
      }
1250 1251 1252 1253 1254
    }
  }
  if (!run_pten_kernel_) {
    if (kernel_type_.get() == nullptr || kernel_func_.get() == nullptr) {
      ChooseKernel(exe_ctx);
1255
      dev_ctx = pool.Get(kernel_type_->place_);
1256
    }
1257 1258
  }

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

1275
  if (!all_kernels_must_compute_runtime_shape_) {
1276
    platform::RecordEvent record_event("infer_shape",
C
chenjian 已提交
1277 1278
                                       platform::TracerEventType::OperatorInner,
                                       1, platform::EventRole::kInnerOp);
1279
    RuntimeInferShapeContext infer_shape_ctx(*this, *runtime_ctx);
1280
    this->Info().infer_shape_(&infer_shape_ctx);
1281
  }
1282 1283 1284 1285 1286

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

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

Y
yuyang18 已提交
1307
  if (!transfered_inplace_vars.empty()) {
T
tianshuo78520a 已提交
1308
    // there is inplace variable has been transferred.
Y
yuyang18 已提交
1309
    TransferInplaceVarsBack(scope, transfered_inplace_vars, *transfer_scope);
1310
  }
1311 1312 1313 1314 1315 1316 1317

  // 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);
  }

1318 1319 1320 1321 1322 1323 1324 1325
  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);
    }
  }
1326

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

  if (FLAGS_check_nan_inf) {
W
WangXi 已提交
1337
    framework::details::CheckOpHasNanOrInf(*this, exec_scope, place);
C
chengduoZH 已提交
1338
  }
1339 1340 1341 1342 1343 1344 1345

  // 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 已提交
1346
}
X
Xin Pan 已提交
1347

1348 1349 1350
OpKernelType OperatorWithKernel::InnerGetExpectedKernelType(
    const ExecutionContext& ctx) const {
  auto expected_kernel_key = this->GetExpectedKernelType(ctx);
1351 1352 1353
  if (HasAttr("op_device")) {
    if (Attr<std::string>("op_device") == "cpu") {
      expected_kernel_key.place_ = platform::CPUPlace();
1354 1355 1356 1357 1358 1359 1360 1361 1362 1363
    } 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.";
      }
1364 1365
      // 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.
1366 1367
      expected_kernel_key.place_ = platform::CPUPlace();
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
1368
      if (SupportGPU()) {
1369
        auto& dev_ctx = ctx.device_context();
1370
        expected_kernel_key.place_ = dev_ctx.GetPlace();
1371 1372 1373 1374 1375
      }
#endif
#ifdef PADDLE_WITH_ASCEND_CL
      if (SupportNPU()) {
        auto& dev_ctx = ctx.device_context();
1376
        expected_kernel_key.place_ = dev_ctx.GetPlace();
1377 1378 1379
      }
#endif
      if (platform::is_cpu_place(expected_kernel_key.place_)) {
1380 1381 1382 1383 1384 1385
        LOG_FIRST_N(WARNING, 1)
            << "Op(" << type_
            << ") has no CUDA implementation. It will be assigned to CPUPlace.";
      }
    }
  }
C
cc 已提交
1386 1387
  VLOG(3) << "op type:" << type_
          << ", expected_kernel_key:" << expected_kernel_key;
1388 1389 1390
  return expected_kernel_key;
}

1391
phi::KernelKey OperatorWithKernel::ChoosePtenKernel(
1392
    const ExecutionContext& ctx) const {
1393
  pt_kernel_signature_.reset(
1394
      new KernelSignature(std::move(GetExpectedPtenKernelArgs(ctx))));
1395
  VLOG(6) << *pt_kernel_signature_.get();
1396 1397 1398 1399

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

Y
YuanRisheng 已提交
1400
  auto pt_kernel_name = pt_kernel_signature_->name;
1401
  auto pt_kernel_key = TransOpKernelTypeToPtenKernelKey(*kernel_type_.get());
1402 1403
  pt_kernel_.reset(new phi::Kernel(phi::KernelFactory::Instance().SelectKernel(
      pt_kernel_name, pt_kernel_key)));
1404 1405

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

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 已提交
1429 1430

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

L
Liu Yiqun 已提交
1432 1433 1434 1435 1436 1437 1438 1439 1440
#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);
  }
1441 1442
#endif
#ifdef PADDLE_WITH_XPU
1443
  if (platform::is_xpu_place(expected_kernel_key.place_) &&
Q
QingshuChen 已提交
1444 1445 1446
      (kernel_iter == kernels.end() ||
       !paddle::platform::is_xpu_support_op(type_, expected_kernel_key) ||
       paddle::platform::is_in_xpu_black_list(type_))) {
1447 1448 1449 1450 1451 1452
    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);
  }
1453
#endif
L
Liu-xiandong 已提交
1454 1455 1456 1457 1458 1459 1460 1461 1462 1463 1464 1465 1466 1467 1468 1469

#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 已提交
1470 1471 1472 1473 1474 1475 1476 1477 1478 1479
#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
1480 1481
#ifdef PADDLE_WITH_ASCEND_CL
  if (kernel_iter == kernels.end() &&
1482
      platform::is_npu_place(expected_kernel_key.place_)) {
1483 1484 1485 1486 1487 1488
    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 已提交
1489 1490 1491
#endif
#ifdef PADDLE_WITH_MLU
  if (kernel_iter == kernels.end() &&
1492
      platform::is_mlu_place(expected_kernel_key.place_)) {
F
fwenguang 已提交
1493 1494 1495 1496 1497 1498
    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 已提交
1499
#endif
1500 1501 1502 1503
  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 已提交
1504

1505 1506 1507 1508 1509
  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 已提交
1510 1511
}

Y
yuyang18 已提交
1512 1513 1514 1515
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 已提交
1516
    VLOG(3) << "share inplace var " + var_name + " back to it's original scope";
C
chengduo 已提交
1517
    auto* origin_var = scope.FindVar(var_name);
1518 1519 1520
    PADDLE_ENFORCE_NOT_NULL(origin_var,
                            platform::errors::InvalidArgument(
                                "The variable[%s] is nullptr.", var_name));
C
chengduo 已提交
1521
    auto* original_tensor =
C
chengduo 已提交
1522
        GetMutableLoDTensorOrSelectedRowsValueFromVar(origin_var);
C
chengduo 已提交
1523
    auto* var = transfer_scope.FindVar(var_name);
1524 1525
    PADDLE_ENFORCE_NOT_NULL(var, platform::errors::InvalidArgument(
                                     "The variable[%s] is nullptr.", var_name));
C
chengduo 已提交
1526
    auto* transformed_tensor = GetLoDTensorOrSelectedRowsValueFromVar(*var);
1527
    auto original_dims = original_tensor->dims();
Y
yuyang18 已提交
1528
    original_tensor->ShareDataWith(*transformed_tensor);
B
Baibaifan 已提交
1529 1530 1531 1532 1533
    // 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 已提交
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 1560 1561 1562 1563 1564 1565
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
1566
      auto src_type = framework::TransToProtoVarType(grad_tensor->dtype());
1567 1568 1569 1570 1571 1572 1573 1574 1575 1576 1577 1578 1579 1580 1581 1582 1583 1584 1585
      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
1586
      auto dst_type = framework::TransToProtoVarType(tensor->dtype());
1587 1588 1589 1590 1591 1592 1593 1594 1595 1596 1597 1598 1599 1600 1601 1602 1603
      // 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 已提交
1604
Scope* OperatorWithKernel::PrepareData(
Y
yuyang18 已提交
1605
    const Scope& scope, const OpKernelType& expected_kernel_key,
X
Xin Pan 已提交
1606 1607
    std::vector<std::string>* transfered_inplace_vars,
    RuntimeContext* ctx) const {
Y
yuyang18 已提交
1608
  Scope* new_scope = nullptr;
S
sneaxiy 已提交
1609

1610
  const std::unordered_set<std::string>* no_buffer_ins = nullptr;
S
sneaxiy 已提交
1611 1612 1613 1614
  if (info_) {
    auto& no_buffer_inferer = info_->NoNeedBufferVarsInferer();
    // Some op may not register NoNeedBufferVarsInferer
    if (no_buffer_inferer) {
1615 1616
      no_buffer_ins = &(no_buffer_inferer(Inputs(), Outputs(), Attrs()));
      if (no_buffer_ins->empty()) no_buffer_ins = nullptr;
S
sneaxiy 已提交
1617 1618 1619
    }
  }

Y
yuyang18 已提交
1620
  for (auto& var_name_item : Inputs()) {
1621 1622
    bool should_skip_input =
        no_buffer_ins && no_buffer_ins->count(var_name_item.first) > 0;
S
sneaxiy 已提交
1623

X
Xin Pan 已提交
1624 1625 1626 1627
    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 已提交
1628
      auto* var = input_vars[i];
X
Xin Pan 已提交
1629

Y
yuyang18 已提交
1630
      // Only tensor can be tranfer to another device.
C
chengduo 已提交
1631
      if (var == nullptr || !VarIsTensor(*var)) {
Y
yuyang18 已提交
1632 1633 1634
        continue;
      }

C
chengduo 已提交
1635
      auto* tensor_in = GetLoDTensorOrSelectedRowsValueFromVar(*var);
1636 1637 1638 1639 1640 1641 1642 1643 1644 1645 1646 1647 1648 1649 1650

      // 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) &&
1651 1652
            (paddle::platform::MKLDNNDeviceContext::tls()
                 .get_cur_paddle_data_layout() == DataLayout::kNHWC)) {
1653 1654 1655 1656 1657 1658 1659 1660 1661 1662 1663 1664 1665 1666 1667 1668 1669 1670 1671 1672 1673
          // 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 已提交
1674 1675 1676 1677 1678 1679 1680 1681 1682 1683 1684
      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 已提交
1685 1686
      VLOG(3) << "Transform Variable " << var_name << " from "
              << kernel_type_for_var << " to " << expected_kernel_key;
Y
yuyang18 已提交
1687

1688 1689 1690
      // 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.
1691
      // We use a thread_local cache to fix that issue, the key in the cache is
1692 1693 1694 1695 1696
      // 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.
1697 1698
      // If this op is not called by an Executor or ParallelExecutor, it should
      // called by a NaiveExecutor, the NaiveExecutor will cache the scopes and
1699
      // variables, that behavior a lot different.
1700 1701 1702 1703 1704 1705 1706 1707 1708
      //
      // 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_))) {
1709 1710
        new_scope = TryCreateTransferScope(kernel_type_for_var,
                                           expected_kernel_key, &scope);
1711
        enable_cache_transfer_scope_ = true;
1712
      }
1713
      if (!new_scope) {
Y
yuyang18 已提交
1714 1715
        new_scope = &scope.NewScope();
      }
1716 1717 1718 1719
      // 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.
1720
      // However, if enable_cache_runtime_context_, we get the cpu tensor each
1721 1722
      // time, not the gpu tensor. Thus, we set pre_scope_ = nullptr
      // to trigger `new RuntimeContext()` in RunImpl().
1723
      if (enable_cache_runtime_context_) {
1724 1725
        pre_scope_ = nullptr;
      }
L
Leo Chen 已提交
1726 1727

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

      // 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 已提交
1747
      Tensor out;
Y
yuyang18 已提交
1748
      TransformData(expected_kernel_key, kernel_type_for_var, *tensor_in, &out);
Y
yuyang18 已提交
1749 1750 1751
      SetTensorToVariable(*var, out, trans_var);
    }
  }
L
Leo Chen 已提交
1752

1753 1754 1755 1756 1757 1758
  // 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 已提交
1759 1760 1761 1762 1763 1764

  // 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) {
1765 1766
    need_prepare_data_ = false;
  }
Y
yuyang18 已提交
1767 1768 1769

  return new_scope;
}
Q
Qiao Longfei 已提交
1770

1771
void OperatorWithKernel::ParseInputDataType(
1772
    const std::vector<Variable*>& vars, const std::string& name,
1773
    proto::VarType::Type* data_type) const {
1774
  proto::VarType::Type default_data_type =
1775 1776 1777 1778 1779 1780 1781 1782 1783
      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>();
1784 1785
      } else if (var->IsType<phi::SelectedRows>()) {
        t = &(var->Get<phi::SelectedRows>().value());
1786
      } else if (var->IsType<LoDTensorArray>()) {
1787 1788 1789 1790
        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));
1791 1792
          }
        }
1793 1794
      }
      if (t != nullptr) {
1795 1796
        PADDLE_ENFORCE_EQ(
            t->IsInitialized(), true,
1797 1798 1799
            platform::errors::InvalidArgument("The %s Op's Input Variable `%s` "
                                              "contains uninitialized Tensor.",
                                              Type(), name));
1800 1801
        proto::VarType::Type tmp =
            paddle::framework::TransToProtoVarType(t->dtype());
1802 1803 1804 1805 1806 1807 1808 1809 1810
        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)));
1811 1812 1813 1814 1815 1816
        *data_type = tmp;
      }
    }
  }
}

1817
proto::VarType::Type OperatorWithKernel::IndicateDataType(
Y
Yu Yang 已提交
1818
    const ExecutionContext& ctx) const {
1819 1820 1821
  proto::VarType::Type dafault_data_type =
      static_cast<proto::VarType::Type>(-1);
  proto::VarType::Type data_type = dafault_data_type;
H
hong 已提交
1822
  for (auto& input : ctx.InNameList()) {
1823 1824
    const std::vector<Variable*> vars = ctx.MultiInputVar(input);
    ParseInputDataType(vars, input, &data_type);
Y
Yu Yang 已提交
1825
  }
1826 1827 1828 1829
  PADDLE_ENFORCE_NE(
      data_type, dafault_data_type,
      platform::errors::NotFound(
          "DataType should be indicated by input Variable at %s.", Type()));
1830 1831 1832 1833 1834 1835 1836 1837
  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;
1838
  ParseInputDataType(ctx.MultiInputVar(name), name, &data_type);
1839 1840
  PADDLE_ENFORCE_NE(
      data_type, dafault_data_type,
1841 1842 1843 1844 1845
      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()));
1846
  return data_type;
Y
Yu Yang 已提交
1847
}
1848

1849 1850 1851 1852 1853 1854 1855 1856 1857 1858 1859 1860 1861 1862 1863 1864 1865 1866
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>();
1867 1868
  } else if (var->IsType<phi::SelectedRows>()) {
    t = var->GetMutable<phi::SelectedRows>()->mutable_value();
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 1894 1895 1896 1897 1898 1899
  } 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
1900 1901
  auto type_a = framework::TransToProtoVarType(tensor_a->dtype());
  auto type_b = framework::TransToProtoVarType(tensor_b->dtype());
1902 1903 1904 1905 1906 1907 1908

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

  return target_type;
}

1909 1910 1911 1912 1913 1914 1915 1916
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 已提交
1917 1918
  return OpKernelType(expected_kernel_type.data_type_, tensor.place(),
                      tensor.layout());
1919 1920
}

1921 1922
KernelSignature OperatorWithKernel::GetExpectedPtenKernelArgs(
    const ExecutionContext& ctx) const {
1923 1924
  InitDefaultKernelSignatureMap();
  ExecutionArgumentMappingContext arg_mapping_ctx(ctx);
1925
  return phi::OpUtilsMap::Instance().GetArgumentMappingFn(Type())(
1926
      arg_mapping_ctx);
1927 1928
}

1929
Scope* OperatorWithKernel::PreparePtenData(
1930
    const Scope& scope, const phi::Kernel& pt_kernel,
1931 1932 1933 1934 1935 1936 1937 1938 1939
    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;
1940
  auto& name_map = Inputs();
Y
YuanRisheng 已提交
1941 1942 1943 1944 1945 1946 1947 1948 1949 1950
  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;
    }
  }

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

1961 1962 1963 1964 1965 1966 1967
    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 已提交
1968 1969 1970 1971 1972 1973 1974 1975 1976

      // 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;
      }

1977 1978 1979 1980
      if (!tensor_in->IsInitialized()) {
        continue;
      }

1981 1982 1983
      if (in_def.backend == phi::Backend::ALL_BACKEND) {
        continue;
      }
1984
      auto expected_place = phi::TransToPtenPlace(in_def.backend);
1985 1986 1987 1988
      if (platform::is_same_place(tensor_in->place(), expected_place)) {
        continue;
      }

1989 1990
      VLOG(3) << "PTen Transform Variable " << input_names[i] << " from "
              << tensor_in->place() << " to " << expected_place;
1991

1992 1993 1994
      if (!new_scope) {
        new_scope = &scope.NewScope();
      }
1995

1996
      // Create new var with the same name in transfer scopes
1997
      auto* trans_var = new_scope->Var(name_vec[offset]);
1998
      ins_vector[offset] = trans_var;
1999

2000 2001 2002 2003
      // Do transfer
      Tensor out;
      framework::TensorCopySync(*tensor_in, expected_place, &out);
      SetTensorToVariable(*var, out, trans_var);
2004 2005 2006 2007 2008 2009
    }
  }

  return new_scope;
}

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

  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 已提交
2042
    auto it = ctx.inputs.find(input_names[i]);
2043 2044 2045

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

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

2073
      pt_kernel_context->EmplaceBackInputWithoutSetRange(tensor_in);
2074
    }
2075
    pt_kernel_context->AssignInputRange(std::make_pair(start_idx, end_idx), i);
2076 2077 2078
  }

  for (size_t i = 0; i < output_names.size(); ++i) {
H
hong 已提交
2079
    auto it = ctx.outputs.find(output_names[i]);
2080
    size_t start_idx =
2081
        (i == 0 ? 0 : pt_kernel_context->OutputRangeAt(i - 1).second);
H
hong 已提交
2082 2083 2084 2085 2086 2087 2088 2089 2090 2091 2092 2093 2094 2095

    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;

2096
    size_t end_idx = start_idx + outs_vector.size();
2097 2098

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

2111 2112
      experimental::ResetTensorDtypeAndLayoutByArgDef(tensor_out,
                                                      output_defs.at(i));
2113
      SetAllocationForOutputTenosr(
2114
          tensor_out, phi::TransToPtenPlace(output_defs.at(i).backend));
2115 2116

      pt_kernel_context->EmplaceBackOutputWithoutSetRange(tensor_out);
2117
    }
2118

2119
    pt_kernel_context->AssignOutputRange(std::make_pair(start_idx, end_idx), i);
2120 2121 2122
  }

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

2185 2186
    } else {
      // TODO(chenweihang): support other attrs later
2187
      auto& attr = Attrs().at(attr_names[i]);
2188
      if (attr_defs[i].type_index == std::type_index(typeid(int))) {
2189
        pt_kernel_context->EmplaceBackAttr(BOOST_GET_CONST(int, attr));
2190
      } else if (attr_defs[i].type_index == std::type_index(typeid(float))) {
2191
        pt_kernel_context->EmplaceBackAttr(BOOST_GET_CONST(float, attr));
2192
      } else if (attr_defs[i].type_index == std::type_index(typeid(bool))) {
2193
        pt_kernel_context->EmplaceBackAttr(BOOST_GET_CONST(bool, attr));
H
hong 已提交
2194 2195
      } 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 已提交
2196 2197 2198
      } else if (attr_defs[i].type_index ==
                 std::type_index(typeid(std::string))) {
        pt_kernel_context->EmplaceBackAttr(BOOST_GET_CONST(std::string, attr));
2199
      } else if (attr_defs[i].type_index ==
2200
                 std::type_index(typeid(phi::DataType))) {
2201
        auto data_type = paddle::framework::TransToPtenDataType(
2202 2203
            static_cast<framework::proto::VarType::Type>(
                BOOST_GET_CONST(int, attr)));
2204
        pt_kernel_context->EmplaceBackAttr(data_type);
2205 2206 2207 2208 2209 2210 2211 2212
      } 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());
2213
          pt_kernel_context->EmplaceBackAttr(vector_int64_attr);
2214 2215 2216
        }
        // TODO(YuanRisheng) Need support vector<int64_t> attr

H
hong 已提交
2217 2218 2219 2220
      } 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);
2221 2222
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
2223
            "Unsupported cast op attribute `%s` when construct "
2224 2225 2226 2227 2228 2229 2230
            "KernelContext.",
            attr_names[i]));
      }
    }
  }
}

Q
Qiao Longfei 已提交
2231
}  // namespace framework
L
liaogang 已提交
2232
}  // namespace paddle