operator.cc 81.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"
Y
Yi Wang 已提交
22
#include "paddle/fluid/framework/data_transform.h"
23
#include "paddle/fluid/framework/data_type_transform.h"
W
WangXi 已提交
24
#include "paddle/fluid/framework/details/nan_inf_utils.h"
25
#include "paddle/fluid/framework/op_call_stack.h"
26
#include "paddle/fluid/framework/pten_utils.h"
Y
Yi Wang 已提交
27
#include "paddle/fluid/framework/shape_inference.h"
28
#include "paddle/fluid/framework/transfer_scope_cache.h"
29
#include "paddle/fluid/framework/unused_var_check.h"
Y
Yi Wang 已提交
30
#include "paddle/fluid/framework/var_type.h"
L
Leo Chen 已提交
31
#include "paddle/fluid/platform/enforce.h"
32
#include "paddle/fluid/platform/profiler.h"
33
#include "paddle/pten/common/scalar.h"
34
#include "paddle/pten/common/scalar_array.h"
H
hong 已提交
35
#include "paddle/pten/core/kernel_factory.h"
36
#include "paddle/pten/ops/compat/signatures.h"
37

38 39 40 41
namespace pten {
class DenseTensor;
}  // namespace pten

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

47 48 49 50
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/fluid/platform/mkldnn_helper.h"
#endif

F
fwenguang 已提交
51 52 53 54
#ifdef PADDLE_WITH_MLU
#include "paddle/fluid/platform/device/mlu/mlu_info.h"
#endif

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

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

65 66 67 68 69 70
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 已提交
71

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

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

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

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

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

127
static std::string GetPlace(const ScopeBase& scope, const std::string& name) {
L
Leo Chen 已提交
128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143
  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());
144 145
  } else if (var->IsType<pten::SelectedRows>()) {
    auto tensor = var->Get<pten::SelectedRows>().value();
L
Leo Chen 已提交
146 147 148 149 150 151 152 153 154 155
    if (UNLIKELY(!tensor.IsInitialized())) {
      return "uninited";
    } else {
      return to_string(tensor.place());
    }
  } else {
    return "";
  }
}

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

162 163
  if (var->IsType<pten::SelectedRows>()) {
    return var->Get<pten::SelectedRows>().rows().size();
164 165 166 167 168
  }

  return -1;
}

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

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

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

X
Xin Pan 已提交
185 186 187 188 189
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 已提交
190
    input_vars.reserve(var_name_item.second.size());
X
Xin Pan 已提交
191 192 193 194 195 196
    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 已提交
197
    output_vars.reserve(var_name_item.second.size());
X
Xin Pan 已提交
198 199 200 201 202 203
    for (auto& var_name : var_name_item.second) {
      output_vars.push_back(scope.FindVar(var_name));
    }
  }
}

204
void OperatorBase::Run(const Scope& scope, const platform::Place& place) {
P
peizhilin 已提交
205 206 207
  try {
    VLOG(4) << place << " " << DebugStringEx(&scope);
    if (platform::is_gpu_place(place)) {
208
#if !defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP)
209 210 211 212
      PADDLE_THROW(platform::errors::Unavailable(
          "Cannot run operator on place %s, please recompile paddle or "
          "reinstall Paddle with CUDA support.",
          place));
213
#else
214
      auto dev_id = place.device;
P
peizhilin 已提交
215
      platform::SetDeviceId(dev_id);
216 217 218
#endif
    } else if (platform::is_xpu_place(place)) {
#ifndef PADDLE_WITH_XPU
219 220 221 222
      PADDLE_THROW(platform::errors::Unavailable(
          "Cannot run operator on place %s, please recompile paddle or "
          "reinstall Paddle with XPU support.",
          place));
223
#else
224
      auto dev_id = place.device;
225
      platform::SetXPUDeviceId(dev_id);
226 227 228 229 230 231 232 233
#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
234
      auto dev_id = place.device;
235
      platform::SetNPUDeviceId(dev_id);
F
fwenguang 已提交
236 237 238 239 240 241 242 243
#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
244
      auto dev_id = place.device;
F
fwenguang 已提交
245
      platform::SetMLUDeviceId(dev_id);
246
#endif
P
peizhilin 已提交
247
    }
P
peizhilin 已提交
248

249
    {
250 251 252 253 254 255
      // 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.
      platform::RecordEvent op_type_record_event(Type());
      auto op_name = platform::OpName(outputs_, Type());
      platform::RecordEvent op_name_record_event(
256
          op_name, platform::EventRole::kUniqueOp);
P
peizhilin 已提交
257 258
      RunImpl(scope, place);
    }
259

Z
Zhang Ting 已提交
260
    VLOG(3) << GetExecutionPlace(place) << " " << DebugStringEx(&scope);
261
  } catch (platform::EnforceNotMet& exception) {
262
    framework::InsertCallStackInfo(Type(), Attrs(), &exception);
263
    throw std::move(exception);
264 265 266 267 268 269
  } 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 已提交
270
  } catch (...) {
271
    LOG(WARNING) << Type() << " raises an unknown exception";
P
peizhilin 已提交
272
    std::rethrow_exception(std::current_exception());
273
  }
274 275
}

276
bool OperatorBase::HasInputs(const std::string& name) const {
M
minqiyang 已提交
277
  return inputs_.find(name) != inputs_.end();
278 279
}

280
std::string OperatorBase::Input(const std::string& name) const {
Y
Yu Yang 已提交
281
  auto& ins = Inputs(name);
282 283
  PADDLE_ENFORCE_LE(
      ins.size(), 1UL,
284
      platform::errors::InvalidArgument(
285 286
          "Operator %s's input %s should contain only one variable.", type_,
          name));
Y
Yu Yang 已提交
287
  return ins.empty() ? kEmptyVarName : ins[0];
Y
Yan Chunwei 已提交
288 289
}

Y
Yu Yang 已提交
290 291
const std::vector<std::string>& OperatorBase::Inputs(
    const std::string& name) const {
Y
Yu Yang 已提交
292
  auto it = inputs_.find(name);
293 294 295 296
  PADDLE_ENFORCE_NE(
      it, inputs_.end(),
      platform::errors::NotFound("Operator %s does not have the input %s.",
                                 type_, name));
Y
Yu Yang 已提交
297
  return it->second;
Y
Yan Chunwei 已提交
298 299
}

300
bool OperatorBase::HasOutputs(const std::string& name) const {
301
  if (outputs_.find(name) != outputs_.end()) {
302 303 304 305 306 307
    return true;
  } else {
    return false;
  }
}

308
std::string OperatorBase::Output(const std::string& name) const {
Y
Yu Yang 已提交
309
  auto& outs = Outputs(name);
310 311 312 313 314
  PADDLE_ENFORCE_LE(
      outs.size(), 1UL,
      platform::errors::InvalidArgument(
          "Operator %s's output %s should contain only one variable.", type_,
          name));
Y
Yu Yang 已提交
315
  return outs.empty() ? kEmptyVarName : outs[0];
Y
Yan Chunwei 已提交
316 317
}

Y
Yu Yang 已提交
318 319
const std::vector<std::string>& OperatorBase::Outputs(
    const std::string& name) const {
Y
Yu Yang 已提交
320
  auto it = outputs_.find(name);
321 322 323 324
  PADDLE_ENFORCE_NE(
      it, outputs_.end(),
      platform::errors::NotFound(
          "Operator %s does not have an output called %s.", type_, name));
Y
Yu Yang 已提交
325
  return it->second;
Y
Yan Chunwei 已提交
326 327
}

328
std::string OperatorBase::DebugStringEx(const ScopeBase* scope) const {
Q
Qiao Longfei 已提交
329
  std::stringstream ss;
Y
Yu Yang 已提交
330
  ss << "Op(" << type_ << "), inputs:{";
331

332
  const std::unordered_set<std::string>* no_need_buffer_vars = nullptr;
333 334
  if (info_ && info_->NoNeedBufferVarsInferer()) {
    no_need_buffer_vars =
335 336
        &(Info().NoNeedBufferVarsInferer()(Inputs(), Outputs(), Attrs()));
    if (no_need_buffer_vars->empty()) no_need_buffer_vars = nullptr;
337 338
  }

Y
Yu Yang 已提交
339 340
  for (auto it = inputs_.begin(); it != inputs_.end();) {
    auto& input = *it;
341 342
    bool is_no_need_buffer_var =
        (no_need_buffer_vars && no_need_buffer_vars->count(input.first) > 0);
Y
Yu Yang 已提交
343 344
    ss << input.first << "[";
    for (size_t i = 0; i < input.second.size(); ++i) {
Q
Qiao Longfei 已提交
345 346
      auto var_name = input.second[i];
      ss << var_name;
347
      if (scope) {
Q
Qiao Longfei 已提交
348 349 350 351 352 353 354
        if (!VarInited(*scope, var_name)) {
          ss << "[uninited]";
        } else {
          int row_size = GetRowSize(*scope, var_name);
          if (row_size >= 0) {
            ss << "[row_size=" << row_size << "]";
          }
355 356 357
          std::string dtype = is_no_need_buffer_var
                                  ? "unknown_dtype"
                                  : GetDtype(*scope, var_name);
Q
Qiao Longfei 已提交
358
          ss << ":" << dtype;
359 360
          ss << "[" << GetDimsDebug(*scope, var_name, true) << "]";
          ss << "(" << GetLoDDebug(*scope, var_name) << ")";
L
Leo Chen 已提交
361
          ss << "(" << GetPlace(*scope, var_name) << ")";
362
        }
363
      }
Y
Yu Yang 已提交
364 365 366
      if (i != input.second.size() - 1) {
        ss << ", ";
      }
367
    }
Y
Yu Yang 已提交
368
    ss << "]";
Y
Yu Yang 已提交
369 370
    ++it;
    if (it != inputs_.end()) {
371 372
      ss << ", ";
    }
Q
Qiao Longfei 已提交
373
  }
Y
Yu Yang 已提交
374
  ss << "}, outputs:{";
Y
Yu Yang 已提交
375 376
  for (auto it = outputs_.begin(); it != outputs_.end();) {
    auto& output = *it;
Y
Yu Yang 已提交
377 378
    ss << output.first << "[";
    for (size_t i = 0; i < output.second.size(); ++i) {
Q
Qiao Longfei 已提交
379 380
      auto var_name = output.second[i];
      ss << var_name;
381
      if (scope) {
Q
Qiao Longfei 已提交
382 383 384 385 386 387 388
        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 已提交
389 390
          std::string dtype = GetDtype(*scope, output.second[i]);
          ss << ":" << dtype;
391 392
          ss << "[" << GetDimsDebug(*scope, var_name, true) << "]";
          ss << "(" << GetLoDDebug(*scope, var_name) << ")";
L
Leo Chen 已提交
393
          ss << "(" << GetPlace(*scope, var_name) << ")";
394
        }
395
      }
Y
Yu Yang 已提交
396 397 398
      if (i != output.second.size() - 1) {
        ss << ", ";
      }
399
    }
Y
Yu Yang 已提交
400
    ss << "]";
Y
Yu Yang 已提交
401 402
    ++it;
    if (it != outputs_.end()) {
403 404
      ss << ", ";
    }
Q
Qiao Longfei 已提交
405
  }
Y
Yu Yang 已提交
406
  ss << "}.";
Q
Qiao Longfei 已提交
407 408 409
  return ss.str();
}

Y
Yu Yang 已提交
410
OperatorBase::OperatorBase(const std::string& type,
Y
Yu Yang 已提交
411 412
                           const VariableNameMap& inputs,
                           const VariableNameMap& outputs,
Y
Yu Yang 已提交
413
                           const AttributeMap& attrs)
S
sneaxiy 已提交
414 415 416 417 418 419
    : type_(type),
      inputs_(inputs),
      outputs_(outputs),
      attrs_(attrs),
      // NOTE(zjl): why op_info may be nullptr?
      info_(OpInfoMap::Instance().GetNullable(type)) {
H
hong 已提交
420 421 422 423 424 425 426 427
  // 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 已提交
428
}
429

Q
qijun 已提交
430 431
std::vector<std::string> OperatorBase::InputVars() const {
  std::vector<std::string> ret_val;
Y
Yu Yang 已提交
432
  for (auto& o : inputs_) {
Q
qijun 已提交
433 434 435 436 437 438
    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 已提交
439 440 441 442 443 444 445 446 447 448
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 已提交
449
  auto& info = Info();
Y
Yu Yang 已提交
450 451

  // get all OpProto::Var for outputs
Y
Yu Yang 已提交
452
  for (auto& o : info.Proto().outputs()) {
Y
Yu Yang 已提交
453 454 455 456 457 458 459 460 461
    // 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 已提交
462 463
}

464
void OperatorBase::CheckAllInputOutputSet() const {
S
sneaxiy 已提交
465
  if (info_ == nullptr || info_->proto_ == nullptr) return;
466

S
sneaxiy 已提交
467
  for (auto& in : info_->Proto().inputs()) {
468
    if (!in.dispensable() && !in.extra()) {
469 470 471 472
      PADDLE_ENFORCE_NE(
          inputs_.find(in.name()), inputs_.end(),
          platform::errors::NotFound("Operator %s's input (%s) is not set.",
                                     Type(), in.name()));
473
    }
474 475
  }

S
sneaxiy 已提交
476
  for (auto& out : info_->Proto().outputs()) {
477
    if (!out.dispensable() && !out.extra()) {
478 479 480 481
      PADDLE_ENFORCE_NE(
          outputs_.find(out.name()), outputs_.end(),
          platform::errors::NotFound("Operator %s's output (%s) is not set.",
                                     Type(), out.name()));
482
    }
483 484 485 486 487 488 489 490 491 492 493 494 495 496 497
  }
}

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

C
chengduo 已提交
499
const Tensor* GetLoDTensorOrSelectedRowsValueFromVar(const Variable& var) {
C
chengduo 已提交
500 501
  if (var.IsType<LoDTensor>()) {
    return static_cast<const Tensor*>(&(var.Get<LoDTensor>()));
502 503
  } else if (var.IsType<pten::SelectedRows>()) {
    return &(var.Get<pten::SelectedRows>().value());
Q
QI JUN 已提交
504
  } else {
505 506 507
    PADDLE_THROW(platform::errors::InvalidArgument(
        "Variable type is %s, expect LoDTensor or SelectedRows.",
        ToTypeName(var.Type())));
Q
QI JUN 已提交
508 509 510
  }
}

C
chengduo 已提交
511
Tensor* GetMutableLoDTensorOrSelectedRowsValueFromVar(Variable* var) {
Q
QI JUN 已提交
512
  if (var->IsType<LoDTensor>()) {
513
    return var->GetMutable<LoDTensor>();
514 515
  } else if (var->IsType<pten::SelectedRows>()) {
    return var->GetMutable<pten::SelectedRows>()->mutable_value();
Q
QI JUN 已提交
516
  } else {
517 518 519
    PADDLE_THROW(platform::errors::InvalidArgument(
        "Variable type is %s, expect LoDTensor or SelectedRows.",
        ToTypeName(var->Type())));
Q
QI JUN 已提交
520 521 522
  }
}

523
bool ExecutionContext::HasInput(const std::string& name) const {
524
  auto* var = InputVar(name);
525 526 527 528
  return var != nullptr;
}

bool ExecutionContext::HasOutput(const std::string& name) const {
529
  auto* var = OutputVar(name);
530 531 532
  return var != nullptr;
}

X
Xin Pan 已提交
533
const Variable* ExecutionContext::InputVar(const std::string& name) const {
534 535
  LogVarUsageIfUnusedVarCheckEnabled(name);

X
Xin Pan 已提交
536 537 538
  auto it = ctx_.inputs.find(name);
  if (it == ctx_.inputs.end()) return nullptr;

539 540
  PADDLE_ENFORCE_LE(
      it->second.size(), 1UL,
541
      platform::errors::InvalidArgument(
542 543
          "Operator %s's input %s should contain only one variable.",
          op_.Type(), name));
X
Xin Pan 已提交
544 545 546
  return it->second.empty() ? nullptr : it->second[0];
}

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

551 552 553 554 555
  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 已提交
556 557 558
  return it->second.empty() ? nullptr : it->second[0];
}

559
template <>
560
const std::vector<const Tensor*> ExecutionContext::MultiInput<Tensor>(
561
    const std::string& name) const {
562 563
  LogVarUsageIfUnusedVarCheckEnabled(name);

H
hong 已提交
564 565
  auto vars = MultiInputVar(name);
  if (vars.size() == 0) {
X
Xin Pan 已提交
566 567 568 569 570
    return {};
  }
  std::vector<const Tensor*> res;
  res.reserve(vars.size());
  std::transform(vars.begin(), vars.end(), std::back_inserter(res),
H
hong 已提交
571
                 [&](const Variable* var) -> const Tensor* {
X
Xin Pan 已提交
572
                   if (var == nullptr) return nullptr;
573 574 575 576 577
                   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 已提交
578 579 580 581 582
                   return &(var->Get<LoDTensor>());
                 });
  return res;
}

583
template <>
584
std::vector<Tensor*> ExecutionContext::MultiOutput<Tensor>(
585
    const std::string& name) const {
H
hong 已提交
586 587 588
  auto vars = MultiOutputVar(name);

  if (vars.size() == 0) {
589 590
    return {};
  }
591
  std::vector<Tensor*> res;
592 593 594 595 596
  res.reserve(vars.size());
  std::transform(vars.begin(), vars.end(), std::back_inserter(res),
                 [&](Variable* var) -> Tensor* {
                   return var == nullptr ? nullptr
                                         : var->GetMutable<LoDTensor>();
597
                 });
598 599 600
  return res;
}

Y
Yu Yang 已提交
601
bool OpSupportGPU(const std::string& op_type) {
H
hong 已提交
602 603 604 605 606 607
  // check in new Function kernel first
  auto& kernel_factory = pten::KernelFactory::Instance();
  auto kernel_key_map =
      kernel_factory.SelectKernelMap(pten::TransToPtenKernelName(op_type));
  for (auto& kernel : kernel_key_map) {
    if (platform::is_gpu_place(
608
            pten::TransToPtenPlace(kernel.first.backend()))) {
H
hong 已提交
609 610 611 612
      return true;
    }
  }

Y
Yu Yang 已提交
613 614 615 616 617 618 619 620 621 622 623
  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 已提交
624

Y
Yu Yang 已提交
625 626 627
  return false;
}

628 629
class RuntimeInferShapeContext : public InferShapeContext {
 public:
630
  RuntimeInferShapeContext(const OperatorBase& op, const RuntimeContext& ctx)
G
Gabor Buella 已提交
631
      : op_(op), ctx_(ctx) {}
632 633

  bool HasInput(const std::string& name) const override {
634
    // has only one input
X
Xin Pan 已提交
635
    const auto& ins = ctx_.inputs;
636 637
    auto it = ins.find(name);
    if (it == ins.end()) {
638 639
      return false;
    }
640
    const auto& in = it->second;
X
Xin Pan 已提交
641
    if (in.size() == 0) return false;
642 643 644 645
    PADDLE_ENFORCE_EQ(
        in.size(), 1UL,
        platform::errors::InvalidArgument(
            "Input %s should not contain more than one inputs.", name));
X
Xin Pan 已提交
646
    return in[0] != nullptr;
647 648 649
  }

  bool HasOutput(const std::string& name) const override {
650
    // has only one output
X
Xin Pan 已提交
651
    const auto& outs = ctx_.outputs;
652 653
    auto it = outs.find(name);
    if (it == outs.end()) {
654 655
      return false;
    }
656
    const auto& out = it->second;
X
Xin Pan 已提交
657
    if (out.size() == 0) {
658 659
      return false;
    }
660 661 662 663
    PADDLE_ENFORCE_EQ(
        out.size(), 1UL,
        platform::errors::InvalidArgument(
            "Output %s should not contain more than one outputs.", name));
X
Xin Pan 已提交
664
    return out[0] != nullptr;
665 666 667
  }

  bool HasInputs(const std::string& name) const override {
X
Xin Pan 已提交
668 669
    const auto& ins = ctx_.inputs;
    auto it = ins.find(name);
X
fix  
Xin Pan 已提交
670
    if (it == ins.end() || it->second.empty()) {
671 672
      return false;
    }
X
Xin Pan 已提交
673 674
    for (auto& input : it->second) {
      if (input == nullptr) {
675 676 677 678 679 680 681
        return false;
      }
    }
    return true;
  }

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

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

H
hong 已提交
697
  std::vector<std::string> Inputs(const std::string& name) const override {
698 699 700
    return op_.Inputs(name);
  }

H
hong 已提交
701
  std::vector<std::string> Outputs(const std::string& name) const override {
702 703 704
    return op_.Outputs(name);
  }

705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727
  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();
  }

728 729
  void ShareDim(const std::string& in, const std::string& out, size_t i = 0,
                size_t j = 0) override {
X
Xin Pan 已提交
730 731
    auto in_it = ctx_.inputs.find(in);
    auto out_it = ctx_.outputs.find(out);
732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747
    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 已提交
748 749 750

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

752 753 754 755 756
    PADDLE_ENFORCE_EQ(
        in_var->Type(), out_var->Type(),
        platform::errors::InvalidArgument(
            "The type of input (%s) and output (%s) are inconsistent.", in,
            out));
757

758 759 760
    if (in_var->IsType<pten::SelectedRows>()) {
      auto& in_sele_rows = in_var->Get<pten::SelectedRows>();
      auto out_sele_rows = out_var->GetMutable<pten::SelectedRows>();
761 762 763 764 765 766 767 768
      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 {
769
      PADDLE_THROW(platform::errors::Unimplemented(
770
          "Currently, the input type of ShareDim only can be LoDTensor "
771
          "or SelectedRows."));
772 773 774
    }
  }

H
hong 已提交
775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792
  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 已提交
793
            "Op [%s]: Input var size should be equal with output var size",
H
hong 已提交
794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819
            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 已提交
820 821
  void ShareLoD(const std::string& in, const std::string& out, size_t i = 0,
                size_t j = 0) const override {
X
Xin Pan 已提交
822 823
    auto in_it = ctx_.inputs.find(in);
    auto out_it = ctx_.outputs.find(out);
824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839
    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 已提交
840 841

    Variable* in_var = in_it->second.at(i);
Q
Qiao Longfei 已提交
842
    if (!in_var->IsType<LoDTensor>()) return;
X
Xin Pan 已提交
843
    Variable* out_var = out_it->second.at(j);
844 845 846 847
    PADDLE_ENFORCE_EQ(
        out_var->IsType<LoDTensor>(), true,
        platform::errors::InvalidArgument(
            "The %zu-th output of Output(%s) must be LoDTensor.", j, out));
848
    auto& in_tensor = in_var->Get<LoDTensor>();
Q
Qiao Longfei 已提交
849 850
    auto* out_tensor = out_var->GetMutable<LoDTensor>();
    out_tensor->set_lod(in_tensor.lod());
D
dzhwinter 已提交
851

M
mozga-intel 已提交
852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870
// 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 已提交
871 872
  }

873
  int32_t GetLoDLevel(const std::string& in, size_t i = 0) const override {
874
    PADDLE_THROW(platform::errors::PreconditionNotMet(
875
        "GetLoDLevel is only used in compile time. The calculation of "
876
        "output's actual lod is different among operators so that should be "
877
        "set in the runtime kernel."));
878 879
  }

880 881
  void SetLoDLevel(const std::string& out, int32_t lod_level,
                   size_t j = 0) const override {
882
    PADDLE_THROW(platform::errors::PreconditionNotMet(
883
        "SetLoDLevel is only used in compile time. The calculation of "
884
        "output's actual lod is different among operators so that should be "
885
        "set in the runtime kernel."));
C
chengduo 已提交
886 887
  }

888 889
  bool IsRuntime() const override { return true; }

890 891 892 893 894 895 896 897 898 899 900
  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;
    }
  }

901 902
  // TODO(paddle-dev): Can this be template?
  std::vector<InferShapeVarPtr> GetInputVarPtrs(
903
      const std::string& name) const override {
904 905 906 907 908 909 910 911
    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(
912
      const std::string& name) const override {
913 914 915 916 917 918 919
    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 已提交
920 921
  DDim GetInputDim(const std::string& name) const override {
    const std::vector<Variable*>& vars = InputVars(name);
922 923 924 925 926
    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 已提交
927 928 929 930 931 932 933 934
    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 已提交
935 936 937 938 939 940 941 942 943 944
  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 已提交
945 946
  void SetOutputDim(const std::string& name, const DDim& dim) override {
    auto& vars = OutputVars(name);
947 948 949 950 951
    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 已提交
952 953 954 955 956 957 958 959 960
    SetDim(vars[0], dim);
  }

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

961
 protected:
X
Xin Pan 已提交
962
  DDim GetDim(Variable* var) const {
963 964
    PADDLE_ENFORCE_NOT_NULL(
        var, platform::errors::InvalidArgument("Input variable is nullptr."));
965 966
    if (var->IsType<LoDTensor>()) {
      return var->Get<LoDTensor>().dims();
967 968
    } else if (var->IsType<pten::SelectedRows>()) {
      return var->Get<pten::SelectedRows>().GetCompleteDims();
969
    } else {
970 971 972 973
      PADDLE_THROW(platform::errors::InvalidArgument(
          "Only LoDTensor or SelectedRows support 'GetDim', but input "
          "Variable's type is %s.",
          ToTypeName(var->Type())));
F
fengjiayi 已提交
974 975 976
    }
  }

X
Xin Pan 已提交
977 978 979 980 981 982 983 984
  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 已提交
985
  std::vector<DDim> GetRepeatedDims(const std::string& name) const override {
986 987
    PADDLE_THROW(platform::errors::PreconditionNotMet(
        "GetRepeatedDims method only ban be used in compile time."));
988 989
  }

X
Xin Pan 已提交
990
  void SetDim(Variable* var, const DDim& dim) {
991 992
    if (var->IsType<LoDTensor>()) {
      var->GetMutable<LoDTensor>()->Resize(dim);
993 994
    } else if (var->IsType<pten::SelectedRows>()) {
      var->GetMutable<pten::SelectedRows>()->set_height(dim[0]);
995
    } else {
996 997 998 999
      PADDLE_THROW(platform::errors::Unimplemented(
          "Variable type error, expect LoDTensor or SelectedRows, but received "
          "(%s).",
          ToTypeName(var->Type())));
X
Xin Pan 已提交
1000 1001 1002 1003 1004 1005
    }
  }

  void SetDims(const std::vector<Variable*>& vars,
               const std::vector<DDim>& dims) {
    size_t length = vars.size();
1006 1007 1008 1009 1010 1011
    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 已提交
1012 1013 1014 1015 1016
    for (size_t i = 0; i < length; ++i) {
      if (vars[i] == nullptr) {
        continue;
      }
      SetDim(vars[i], dims[i]);
1017 1018 1019
    }
  }

F
fengjiayi 已提交
1020 1021
  void SetRepeatedDims(const std::string& name,
                       const std::vector<DDim>& dims) override {
1022 1023
    PADDLE_THROW(platform::errors::PreconditionNotMet(
        "SetRepeatedDims method only can be used in compile time."));
F
fengjiayi 已提交
1024 1025
  }

X
Xin Pan 已提交
1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036
  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 {
1037 1038 1039
    return ToVarType(var->Type());
  }

1040 1041 1042
 private:
  const std::vector<Variable*>& InputVars(const std::string& name) const {
    auto it = ctx_.inputs.find(name);
1043 1044 1045 1046
    PADDLE_ENFORCE_NE(
        it, ctx_.inputs.end(),
        platform::errors::NotFound(
            "Operator (%s) does not have the input (%s).", op_.Type(), name));
1047 1048 1049 1050 1051
    return it->second;
  }

  const std::vector<Variable*>& OutputVars(const std::string& name) const {
    auto it = ctx_.outputs.find(name);
1052 1053 1054 1055
    PADDLE_ENFORCE_NE(
        it, ctx_.outputs.end(),
        platform::errors::NotFound(
            "Operator (%s) does not have the outputs (%s).", op_.Type(), name));
1056
    return it->second;
F
fengjiayi 已提交
1057 1058
  }

1059
  const OperatorBase& op_;
X
Xin Pan 已提交
1060
  const RuntimeContext& ctx_;
1061 1062
};

1063 1064
static void CheckTensorNANOrInf(const std::string& op_type,
                                const std::string& name,
C
chengduoZH 已提交
1065 1066 1067 1068
                                const framework::Tensor& tensor) {
  if (tensor.memory_size() == 0) {
    return;
  }
Y
Yu Yang 已提交
1069 1070
  if (tensor.type() != proto::VarType::FP32 &&
      tensor.type() != proto::VarType::FP64) {
C
chengduoZH 已提交
1071 1072
    return;
  }
1073 1074 1075 1076 1077 1078 1079 1080
  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 已提交
1081 1082
}

1083 1084
bool OperatorWithKernel::SupportsMKLDNN(
    const proto::VarType::Type data_type) const {
1085 1086
  auto& op_kernels = OperatorWithKernel::AllOpKernels().at(type_);
  return std::any_of(op_kernels.begin(), op_kernels.end(),
1087
                     [data_type](OpKernelMap::const_reference kern_pair) {
1088 1089
                       return platform::is_cpu_place(kern_pair.first.place_) &&
                              kern_pair.first.library_type_ ==
1090 1091
                                  LibraryType::kMKLDNN &&
                              kern_pair.first.data_type_ == data_type;
1092 1093 1094
                     });
}

1095 1096
bool OperatorWithKernel::CanMKLDNNBeUsed(const framework::ExecutionContext& ctx,
                                         proto::VarType::Type data_type) const {
1097 1098 1099
  bool use_mkldnn_ctx = ctx.HasAttr("use_mkldnn") &&
                        ctx.Attr<bool>("use_mkldnn") &&
                        platform::is_cpu_place(ctx.GetPlace());
1100
  return use_mkldnn_ctx && this->SupportsMKLDNN(data_type);
1101 1102
}

1103 1104 1105 1106 1107 1108 1109
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 已提交
1110
void OperatorWithKernel::RuntimeInferShape(const Scope& scope,
X
Xin Pan 已提交
1111 1112
                                           const platform::Place& place,
                                           const RuntimeContext& ctx) const {
1113
  RuntimeInferShapeContext infer_shape_ctx(*this, ctx);
1114
  this->Info().infer_shape_(&infer_shape_ctx);
B
baojun-nervana 已提交
1115 1116
}

L
luotao1 已提交
1117 1118
void OperatorWithKernel::RunImpl(const Scope& scope,
                                 const platform::Place& place) const {
L
luotao1 已提交
1119 1120
  // To reduce the elapsed time of HasAttr, we use bool variable to record the
  // result of HasAttr.
1121 1122 1123
  if (!enable_cache_runtime_context_ && HasAttr(kEnableCacheRuntimeContext))
    enable_cache_runtime_context_ = true;
  if (!all_kernels_must_compute_runtime_shape_ &&
L
luotao1 已提交
1124
      HasAttr(kAllKernelsMustComputeRuntimeShape))
1125
    all_kernels_must_compute_runtime_shape_ = true;
1126
  const Scope* cur_scope = &scope;
1127
  if (!enable_cache_runtime_context_) {
L
luotao1 已提交
1128 1129
    RuntimeContext ctx(Inputs(), Outputs(), scope);
    RunImpl(scope, place, &ctx);
1130
    pre_scope_ = cur_scope;
L
luotao1 已提交
1131
  } else {
1132 1133 1134 1135 1136 1137
    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 已提交
1138 1139 1140 1141 1142 1143 1144 1145
    }
    RunImpl(scope, place, runtime_ctx_.get());
  }
}

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

1149 1150 1151 1152 1153 1154 1155 1156 1157 1158
#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

1159
  auto exe_ctx = ExecutionContext(*this, scope, *dev_ctx, *runtime_ctx);
1160 1161 1162 1163
  // using cache
  if (kernel_type_.get()) {
    dev_ctx = pool.Get(kernel_type_->place_);
  }
1164 1165 1166 1167 1168 1169

  // 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
1170 1171 1172
  pten::KernelKey pt_kernel_key;
  std::string pt_kernel_name;
  if (pten::KernelFactory::Instance().HasCompatiblePtenKernel(type_)) {
1173
    if (pt_kernel_signature_ == nullptr || pt_kernel_ == nullptr) {
1174 1175
      pt_kernel_signature_.reset(
          new KernelSignature(std::move(GetExpectedPtenKernelArgs(exe_ctx))));
1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226
      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(
          new pten::Kernel(pten::KernelFactory::Instance().SelectKernel(
              pt_kernel_name, pt_kernel_key)));

      if (pt_kernel_->IsValid()) {
        VLOG(6) << "Static mode ChoosePtenKernel - kernel name: "
                << pt_kernel_name << " | kernel key: " << pt_kernel_key
                << " | kernel: " << *pt_kernel_;
      } else {
        VLOG(6) << "Static mode ChoosePtenKernel - kernel `" << pt_kernel_name
                << "` not found.";
      }
    }
    if (pt_kernel_->IsValid()) {
      run_pten_kernel_ = true;
    } else {
      auto& all_op_kernels = AllOpKernels();
      auto kernels_iter = all_op_kernels.find(type_);
      if (kernels_iter == all_op_kernels.end() ||
          kernels_iter->second.find(*kernel_type_.get()) ==
              kernels_iter->second.end()
#ifdef PADDLE_WITH_XPU
          ||
          paddle::platform::is_xpu_place(kernel_type_->place_) &&  // NOLINT
              !paddle::platform::is_xpu_support_op(
                  type_, *kernel_type_.get())  // NOLINT
          || paddle::platform::is_in_xpu_black_list(type_)
#endif
              ) {
        auto pt_cpu_kernel_key =
            FallBackToCpu(*kernel_type_.get(), pt_kernel_key, *this);
        pt_kernel_.reset(
            new pten::Kernel(pten::KernelFactory::Instance().SelectKernel(
                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;
        }
      }
1227 1228 1229 1230 1231
    }
  }
  if (!run_pten_kernel_) {
    if (kernel_type_.get() == nullptr || kernel_func_.get() == nullptr) {
      ChooseKernel(exe_ctx);
1232
      dev_ctx = pool.Get(kernel_type_->place_);
1233
    }
1234 1235
  }

Y
yuyang18 已提交
1236 1237
  // do data transformScope &transfer_scope;
  std::vector<std::string> transfered_inplace_vars;
1238 1239
  Scope* transfer_scope = nullptr;
  {
1240
    platform::RecordEvent record_event("prepare_data",
1241
                                       platform::EventRole::kInnerOp);
1242 1243 1244 1245
    if (need_prepare_data_) {
      transfer_scope = PrepareData(scope, *kernel_type_,
                                   &transfered_inplace_vars, runtime_ctx);
    }
1246
  }
Y
yuyang18 已提交
1247 1248 1249 1250
  // exec scope is the scope that kernel actually executed on.
  const Scope& exec_scope =
      (transfer_scope == nullptr ? scope : *transfer_scope);

1251
  if (!all_kernels_must_compute_runtime_shape_) {
1252
    platform::RecordEvent record_event("infer_shape",
1253
                                       platform::EventRole::kInnerOp);
1254
    RuntimeInferShapeContext infer_shape_ctx(*this, *runtime_ctx);
1255
    this->Info().infer_shape_(&infer_shape_ctx);
1256
  }
1257 1258 1259 1260 1261

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

X
clean  
Xin Pan 已提交
1262 1263
  // TODO(panyx0718): ExecutionContext should only depend on RuntimeContext
  // not Scope. Imperative mode only pass inputs and get outputs.
1264
  {
1265
    platform::RecordEvent record_event("compute",
1266
                                       platform::EventRole::kInnerOp);
1267
    if (run_pten_kernel_) {
1268
      pten::KernelContext pt_kernel_context;
1269
      // Do data transform before building KernelContext
1270
      // TODO(zhiqiu): support TransferInplaceVarsBack
1271 1272
      PreparePtenData(exec_scope, *pt_kernel_, *pt_kernel_signature_,
                      runtime_ctx);
1273 1274
      BuildPtenKernelContext(*runtime_ctx, dev_ctx, &pt_kernel_context);
      (*pt_kernel_)(&pt_kernel_context);
1275 1276 1277 1278
    } else {
      (*kernel_func_)(
          ExecutionContext(*this, exec_scope, *dev_ctx, *runtime_ctx));
    }
1279
  }
D
dzhwinter 已提交
1280

Y
yuyang18 已提交
1281
  if (!transfered_inplace_vars.empty()) {
T
tianshuo78520a 已提交
1282
    // there is inplace variable has been transferred.
Y
yuyang18 已提交
1283
    TransferInplaceVarsBack(scope, transfered_inplace_vars, *transfer_scope);
1284
  }
1285 1286 1287 1288 1289 1290 1291

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

1292 1293 1294 1295 1296 1297 1298 1299
  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);
    }
  }
1300

D
dzhwinter 已提交
1301
  /*For profiling/benchmark only*/
D
dzhwinter 已提交
1302
  if (FLAGS_benchmark) {
Y
yuyang18 已提交
1303
    dev_ctx->Wait();
1304 1305
#if defined(PADDLE_WITH_CUDA) || defined(PADLDE_WITH_ROCM)
    PADDLE_ENFORCE_GPU_SUCCESS(platform::GpuGetLastError());
1306 1307
#endif
    VLOG(4) << "Operator(" << Type() << "): context wait and get last error";
D
dzhwinter 已提交
1308
  }
C
chengduoZH 已提交
1309 1310

  if (FLAGS_check_nan_inf) {
W
WangXi 已提交
1311
    framework::details::CheckOpHasNanOrInf(*this, exec_scope, place);
C
chengduoZH 已提交
1312
  }
1313 1314 1315 1316 1317 1318 1319

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

1322 1323 1324
OpKernelType OperatorWithKernel::InnerGetExpectedKernelType(
    const ExecutionContext& ctx) const {
  auto& dev_ctx = ctx.device_context();
L
Liu Yiqun 已提交
1325

1326
  auto expected_kernel_key = this->GetExpectedKernelType(ctx);
1327 1328 1329
  if (HasAttr("op_device")) {
    if (Attr<std::string>("op_device") == "cpu") {
      expected_kernel_key.place_ = platform::CPUPlace();
1330 1331 1332 1333 1334 1335 1336 1337 1338 1339
    } 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.";
      }
1340 1341 1342
      // 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.
      if (SupportGPU()) {
1343
        expected_kernel_key.place_ = dev_ctx.GetPlace();
B
Baibaifan 已提交
1344
      } else if (SupportNPU()) {
1345
        expected_kernel_key.place_ = dev_ctx.GetPlace();
1346 1347 1348 1349 1350 1351 1352 1353
      } else {
        expected_kernel_key.place_ = platform::CPUPlace();
        LOG_FIRST_N(WARNING, 1)
            << "Op(" << type_
            << ") has no CUDA implementation. It will be assigned to CPUPlace.";
      }
    }
  }
C
cc 已提交
1354 1355
  VLOG(3) << "op type:" << type_
          << ", expected_kernel_key:" << expected_kernel_key;
1356 1357 1358
  return expected_kernel_key;
}

1359 1360
pten::KernelKey OperatorWithKernel::ChoosePtenKernel(
    const ExecutionContext& ctx) const {
1361
  pt_kernel_signature_.reset(
1362
      new KernelSignature(std::move(GetExpectedPtenKernelArgs(ctx))));
1363
  VLOG(6) << *pt_kernel_signature_.get();
1364 1365 1366 1367

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

Y
YuanRisheng 已提交
1368
  auto pt_kernel_name = pt_kernel_signature_->name;
1369 1370 1371 1372 1373 1374
  auto pt_kernel_key = TransOpKernelTypeToPtenKernelKey(*kernel_type_.get());
  pt_kernel_.reset(
      new pten::Kernel(pten::KernelFactory::Instance().SelectKernel(
          pt_kernel_name, pt_kernel_key)));

  if (pt_kernel_->IsValid()) {
C
Chen Weihang 已提交
1375
    VLOG(6) << "Static mode ChoosePtenKernel - kernel name: " << pt_kernel_name
1376 1377 1378
            << " | kernel key: " << pt_kernel_key
            << " | kernel: " << *pt_kernel_;
  } else {
C
Chen Weihang 已提交
1379
    VLOG(6) << "Static mode ChoosePtenKernel - kernel `" << pt_kernel_name
1380 1381
            << "` not found.";
  }
1382
  return pt_kernel_key;
1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396 1397
}

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 已提交
1398 1399

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

L
Liu Yiqun 已提交
1401 1402 1403 1404 1405 1406 1407 1408 1409
#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);
  }
1410 1411
#endif
#ifdef PADDLE_WITH_XPU
1412
  if (platform::is_xpu_place(expected_kernel_key.place_) &&
Q
QingshuChen 已提交
1413 1414 1415
      (kernel_iter == kernels.end() ||
       !paddle::platform::is_xpu_support_op(type_, expected_kernel_key) ||
       paddle::platform::is_in_xpu_black_list(type_))) {
1416 1417 1418 1419 1420 1421
    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);
  }
1422
#endif
L
Liu-xiandong 已提交
1423 1424 1425 1426 1427 1428 1429 1430 1431 1432 1433 1434 1435 1436 1437 1438

#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 已提交
1439 1440 1441 1442 1443 1444 1445 1446 1447 1448
#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
1449 1450
#ifdef PADDLE_WITH_ASCEND_CL
  if (kernel_iter == kernels.end() &&
1451
      platform::is_npu_place(expected_kernel_key.place_)) {
1452 1453 1454 1455 1456 1457
    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 已提交
1458 1459 1460
#endif
#ifdef PADDLE_WITH_MLU
  if (kernel_iter == kernels.end() &&
1461
      platform::is_mlu_place(expected_kernel_key.place_)) {
F
fwenguang 已提交
1462 1463 1464 1465 1466 1467
    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 已提交
1468
#endif
1469 1470 1471 1472
  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 已提交
1473

1474 1475 1476 1477 1478
  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 已提交
1479 1480
}

Y
yuyang18 已提交
1481 1482 1483 1484
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 已提交
1485
    VLOG(3) << "share inplace var " + var_name + " back to it's original scope";
C
chengduo 已提交
1486
    auto* origin_var = scope.FindVar(var_name);
1487 1488 1489
    PADDLE_ENFORCE_NOT_NULL(origin_var,
                            platform::errors::InvalidArgument(
                                "The variable[%s] is nullptr.", var_name));
C
chengduo 已提交
1490
    auto* original_tensor =
C
chengduo 已提交
1491
        GetMutableLoDTensorOrSelectedRowsValueFromVar(origin_var);
C
chengduo 已提交
1492
    auto* var = transfer_scope.FindVar(var_name);
1493 1494
    PADDLE_ENFORCE_NOT_NULL(var, platform::errors::InvalidArgument(
                                     "The variable[%s] is nullptr.", var_name));
C
chengduo 已提交
1495
    auto* transformed_tensor = GetLoDTensorOrSelectedRowsValueFromVar(*var);
1496
    auto original_dims = original_tensor->dims();
Y
yuyang18 已提交
1497
    original_tensor->ShareDataWith(*transformed_tensor);
B
Baibaifan 已提交
1498 1499 1500 1501 1502
    // 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 已提交
1503 1504 1505
  }
}

1506 1507 1508 1509 1510 1511 1512 1513 1514 1515 1516 1517 1518 1519 1520 1521 1522 1523 1524 1525 1526 1527 1528 1529 1530 1531 1532 1533 1534 1535 1536 1537 1538 1539 1540 1541 1542 1543 1544 1545 1546 1547 1548 1549 1550 1551 1552 1553 1554 1555 1556 1557 1558 1559 1560 1561 1562 1563 1564 1565 1566 1567 1568 1569 1570 1571 1572
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
      auto src_type = grad_tensor->type();
      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
      auto dst_type = tensor->saved_type();
      // 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 已提交
1573
Scope* OperatorWithKernel::PrepareData(
Y
yuyang18 已提交
1574
    const Scope& scope, const OpKernelType& expected_kernel_key,
X
Xin Pan 已提交
1575 1576
    std::vector<std::string>* transfered_inplace_vars,
    RuntimeContext* ctx) const {
Y
yuyang18 已提交
1577
  Scope* new_scope = nullptr;
S
sneaxiy 已提交
1578

1579
  const std::unordered_set<std::string>* no_buffer_ins = nullptr;
S
sneaxiy 已提交
1580 1581 1582 1583
  if (info_) {
    auto& no_buffer_inferer = info_->NoNeedBufferVarsInferer();
    // Some op may not register NoNeedBufferVarsInferer
    if (no_buffer_inferer) {
1584 1585
      no_buffer_ins = &(no_buffer_inferer(Inputs(), Outputs(), Attrs()));
      if (no_buffer_ins->empty()) no_buffer_ins = nullptr;
S
sneaxiy 已提交
1586 1587 1588
    }
  }

Y
yuyang18 已提交
1589
  for (auto& var_name_item : Inputs()) {
1590 1591
    bool should_skip_input =
        no_buffer_ins && no_buffer_ins->count(var_name_item.first) > 0;
S
sneaxiy 已提交
1592

X
Xin Pan 已提交
1593 1594 1595 1596
    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 已提交
1597
      auto* var = input_vars[i];
X
Xin Pan 已提交
1598

Y
yuyang18 已提交
1599
      // Only tensor can be tranfer to another device.
C
chengduo 已提交
1600
      if (var == nullptr || !VarIsTensor(*var)) {
Y
yuyang18 已提交
1601 1602 1603
        continue;
      }

C
chengduo 已提交
1604
      auto* tensor_in = GetLoDTensorOrSelectedRowsValueFromVar(*var);
1605 1606 1607 1608 1609 1610 1611 1612 1613 1614 1615 1616 1617 1618 1619

      // 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) &&
1620 1621
            (paddle::platform::MKLDNNDeviceContext::tls()
                 .get_cur_paddle_data_layout() == DataLayout::kNHWC)) {
1622 1623 1624 1625 1626 1627 1628 1629 1630 1631 1632 1633 1634 1635 1636 1637 1638 1639 1640 1641 1642
          // 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 已提交
1643 1644 1645 1646 1647 1648 1649 1650 1651 1652 1653
      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 已提交
1654 1655
      VLOG(3) << "Transform Variable " << var_name << " from "
              << kernel_type_for_var << " to " << expected_kernel_key;
Y
yuyang18 已提交
1656

1657 1658 1659
      // 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.
1660
      // We use a thread_local cache to fix that issue, the key in the cache is
1661 1662 1663 1664 1665
      // 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.
1666 1667
      // If this op is not called by an Executor or ParallelExecutor, it should
      // called by a NaiveExecutor, the NaiveExecutor will cache the scopes and
1668
      // variables, that behavior a lot different.
1669 1670 1671 1672 1673 1674 1675 1676 1677
      //
      // 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_))) {
1678 1679
        new_scope = TryCreateTransferScope(kernel_type_for_var,
                                           expected_kernel_key, &scope);
1680
        enable_cache_transfer_scope_ = true;
1681
      }
1682
      if (!new_scope) {
Y
yuyang18 已提交
1683 1684
        new_scope = &scope.NewScope();
      }
1685 1686 1687 1688
      // 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.
1689
      // However, if enable_cache_runtime_context_, we get the cpu tensor each
1690 1691
      // time, not the gpu tensor. Thus, we set pre_scope_ = nullptr
      // to trigger `new RuntimeContext()` in RunImpl().
1692
      if (enable_cache_runtime_context_) {
1693 1694
        pre_scope_ = nullptr;
      }
L
Leo Chen 已提交
1695 1696

      // Create new var with the same name in transfer scopes
Y
yuyang18 已提交
1697
      auto* trans_var = new_scope->Var(var_name);
X
fix  
Xin Pan 已提交
1698
      input_vars[i] = trans_var;
L
Leo Chen 已提交
1699 1700 1701 1702 1703 1704 1705 1706 1707 1708 1709 1710 1711 1712 1713 1714 1715

      // 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 已提交
1716
      Tensor out;
Y
yuyang18 已提交
1717
      TransformData(expected_kernel_key, kernel_type_for_var, *tensor_in, &out);
Y
yuyang18 已提交
1718 1719 1720
      SetTensorToVariable(*var, out, trans_var);
    }
  }
L
Leo Chen 已提交
1721

1722 1723 1724 1725 1726 1727
  // 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 已提交
1728 1729 1730 1731 1732 1733

  // 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) {
1734 1735
    need_prepare_data_ = false;
  }
Y
yuyang18 已提交
1736 1737 1738

  return new_scope;
}
Q
Qiao Longfei 已提交
1739

1740
void OperatorWithKernel::ParseInputDataType(
1741
    const std::vector<Variable*>& vars, const std::string& name,
1742
    proto::VarType::Type* data_type) const {
1743
  proto::VarType::Type default_data_type =
1744 1745 1746 1747 1748 1749 1750 1751 1752
      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>();
1753 1754
      } else if (var->IsType<pten::SelectedRows>()) {
        t = &(var->Get<pten::SelectedRows>().value());
1755
      } else if (var->IsType<LoDTensorArray>()) {
1756 1757 1758 1759
        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));
1760 1761
          }
        }
1762 1763
      }
      if (t != nullptr) {
1764 1765
        PADDLE_ENFORCE_EQ(
            t->IsInitialized(), true,
1766 1767 1768
            platform::errors::InvalidArgument("The %s Op's Input Variable `%s` "
                                              "contains uninitialized Tensor.",
                                              Type(), name));
1769
        proto::VarType::Type tmp = t->type();
1770 1771 1772 1773 1774 1775 1776 1777 1778
        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)));
1779 1780 1781 1782 1783 1784
        *data_type = tmp;
      }
    }
  }
}

1785
proto::VarType::Type OperatorWithKernel::IndicateDataType(
Y
Yu Yang 已提交
1786
    const ExecutionContext& ctx) const {
1787 1788 1789
  proto::VarType::Type dafault_data_type =
      static_cast<proto::VarType::Type>(-1);
  proto::VarType::Type data_type = dafault_data_type;
H
hong 已提交
1790
  for (auto& input : ctx.InNameList()) {
1791 1792
    const std::vector<Variable*> vars = ctx.MultiInputVar(input);
    ParseInputDataType(vars, input, &data_type);
Y
Yu Yang 已提交
1793
  }
1794 1795 1796 1797
  PADDLE_ENFORCE_NE(
      data_type, dafault_data_type,
      platform::errors::NotFound(
          "DataType should be indicated by input Variable at %s.", Type()));
1798 1799 1800 1801 1802 1803 1804 1805
  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;
1806
  ParseInputDataType(ctx.MultiInputVar(name), name, &data_type);
1807 1808
  PADDLE_ENFORCE_NE(
      data_type, dafault_data_type,
1809 1810 1811 1812 1813
      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()));
1814
  return data_type;
Y
Yu Yang 已提交
1815
}
1816

1817 1818 1819 1820 1821 1822 1823 1824 1825 1826 1827 1828 1829 1830 1831 1832 1833 1834
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>();
1835 1836
  } else if (var->IsType<pten::SelectedRows>()) {
    t = var->GetMutable<pten::SelectedRows>()->mutable_value();
1837 1838 1839 1840 1841 1842 1843 1844 1845 1846 1847 1848 1849 1850 1851 1852 1853 1854 1855 1856 1857 1858 1859 1860 1861 1862 1863 1864 1865 1866 1867 1868 1869 1870 1871 1872 1873 1874 1875 1876
  } 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
  auto type_a = tensor_a->type();
  auto type_b = tensor_b->type();

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

  return target_type;
}

1877 1878 1879 1880 1881 1882 1883 1884
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 已提交
1885 1886
  return OpKernelType(expected_kernel_type.data_type_, tensor.place(),
                      tensor.layout());
1887 1888
}

1889 1890
KernelSignature OperatorWithKernel::GetExpectedPtenKernelArgs(
    const ExecutionContext& ctx) const {
1891 1892 1893 1894
  InitDefaultKernelSignatureMap();
  ExecutionArgumentMappingContext arg_mapping_ctx(ctx);
  return pten::OpUtilsMap::Instance().GetArgumentMappingFn(Type())(
      arg_mapping_ctx);
1895 1896
}

1897 1898 1899 1900 1901 1902 1903 1904 1905 1906 1907
Scope* OperatorWithKernel::PreparePtenData(
    const Scope& scope, const pten::Kernel& pt_kernel,
    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;
1908
  auto& name_map = Inputs();
Y
YuanRisheng 已提交
1909 1910 1911 1912 1913 1914 1915 1916 1917 1918
  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;
    }
  }

1919 1920
  for (size_t i = 0; i < input_defs.size(); ++i) {
    auto& in_def = input_defs.at(i);
H
hong 已提交
1921 1922 1923 1924 1925 1926
    auto it = ctx->inputs.find(input_names[i]);
    if (it == ctx->inputs.end()) {
      continue;
    }

    auto& ins_vector = it->second;
1927
    auto& name_vec = name_map.at(input_names[i]);
Y
YuanRisheng 已提交
1928 1929 1930
    bool should_skip_input =
        no_buffer_ins && no_buffer_ins->count(input_names[i]) > 0;

1931 1932 1933 1934 1935 1936 1937 1938
    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 已提交
1939 1940 1941 1942 1943 1944 1945 1946 1947

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

1948 1949 1950 1951
      if (!tensor_in->IsInitialized()) {
        continue;
      }

1952
      auto expected_place = pten::TransToPtenPlace(in_def.backend);
1953 1954 1955 1956
      if (platform::is_same_place(tensor_in->place(), expected_place)) {
        continue;
      }

1957 1958
      VLOG(3) << "PTen Transform Variable " << input_names[i] << " from "
              << tensor_in->place() << " to " << expected_place;
1959

1960 1961 1962
      if (!new_scope) {
        new_scope = &scope.NewScope();
      }
1963

1964
      // Create new var with the same name in transfer scopes
1965
      auto* trans_var = new_scope->Var(name_vec[offset]);
1966
      ins_vector[offset] = trans_var;
1967

1968 1969 1970 1971
      // Do transfer
      Tensor out;
      framework::TensorCopySync(*tensor_in, expected_place, &out);
      SetTensorToVariable(*var, out, trans_var);
1972 1973 1974 1975 1976 1977
    }
  }

  return new_scope;
}

1978
void OperatorWithKernel::BuildPtenKernelContext(
1979 1980 1981
    const RuntimeContext& ctx, platform::DeviceContext* dev_ctx,
    pten::KernelContext* pt_kernel_context) const {
  pt_kernel_context->SetDeviceContext(dev_ctx);
1982 1983 1984 1985 1986 1987 1988 1989 1990 1991 1992 1993 1994 1995 1996 1997 1998 1999 2000 2001 2002 2003 2004 2005 2006 2007 2008 2009

  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 已提交
2010
    auto it = ctx.inputs.find(input_names[i]);
2011 2012 2013

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

H
hong 已提交
2016 2017 2018 2019 2020 2021 2022 2023 2024 2025 2026 2027
    // deal with optional here
    if ((it == ctx.inputs.end()) &&
        (input_defs[i].type_index ==
         std::type_index(typeid(paddle::optional<const pten::DenseTensor&>)))) {
      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();
2028
    for (size_t offset = 0; offset < ins_vector.size(); ++offset) {
2029
      const pten::TensorBase* tensor_in = nullptr;
2030
      auto* var = ins_vector[offset];
H
hong 已提交
2031 2032
      if (var->IsType<framework::LoDTensor>()) {
        tensor_in = &(var->Get<framework::LoDTensor>());
2033 2034
      } else if (var->IsType<pten::SelectedRows>()) {
        tensor_in = &(var->Get<pten::SelectedRows>());
2035 2036 2037 2038
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
            "Unsupported input `%s` type when call pt kernel.",
            framework::ToTypeName(var->Type())));
2039
      }
H
hong 已提交
2040

2041
      pt_kernel_context->EmplaceBackInputWithoutSetRange(tensor_in);
2042
    }
2043
    pt_kernel_context->AssignInputRange(std::make_pair(start_idx, end_idx), i);
2044 2045 2046
  }

  for (size_t i = 0; i < output_names.size(); ++i) {
H
hong 已提交
2047
    auto it = ctx.outputs.find(output_names[i]);
2048
    size_t start_idx =
2049
        (i == 0 ? 0 : pt_kernel_context->OutputRangeAt(i - 1).second);
H
hong 已提交
2050 2051 2052 2053 2054 2055 2056 2057 2058 2059 2060 2061 2062 2063

    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;

2064
    size_t end_idx = start_idx + outs_vector.size();
2065 2066

    for (size_t offset = 0; offset < outs_vector.size(); ++offset) {
2067
      pten::TensorBase* tensor_out = nullptr;
2068
      auto* var = outs_vector[offset];
H
hong 已提交
2069 2070
      if (var->template IsType<framework::LoDTensor>()) {
        tensor_out = var->template GetMutable<framework::LoDTensor>();
2071 2072
      } else if (var->template IsType<pten::SelectedRows>()) {
        tensor_out = var->template GetMutable<pten::SelectedRows>();
2073 2074 2075 2076
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
            "Unsupported output `%s` type when call pt kernel.",
            framework::ToTypeName(var->Type())));
2077
      }
2078

2079 2080
      experimental::ResetTensorDtypeAndLayoutByArgDef(tensor_out,
                                                      output_defs.at(i));
2081
      SetAllocationForOutputTenosr(
2082
          tensor_out, pten::TransToPtenPlace(output_defs.at(i).backend));
2083 2084

      pt_kernel_context->EmplaceBackOutputWithoutSetRange(tensor_out);
2085
    }
2086

2087
    pt_kernel_context->AssignOutputRange(std::make_pair(start_idx, end_idx), i);
2088 2089 2090
  }

  for (size_t i = 0; i < attr_names.size(); ++i) {
2091 2092 2093 2094 2095
    if (attr_defs[i].type_index == std::type_index(typeid(pten::ScalarArray))) {
      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>))) {
2096
          pt_kernel_context->EmplaceBackAttr(std::move(pten::ScalarArray(
2097
              BOOST_GET_CONST(std::vector<int64_t>, attr_iter->second))));
2098 2099
        } else if (std::type_index(attr_iter->second.type()) ==
                   std::type_index(typeid(std::vector<int32_t>))) {
2100
          pt_kernel_context->EmplaceBackAttr(std::move(pten::ScalarArray(
2101
              BOOST_GET_CONST(std::vector<int32_t>, attr_iter->second))));
2102 2103 2104 2105 2106 2107 2108 2109 2110
        } 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
2111
          pt_kernel_context->EmplaceBackAttr(std::move(
2112 2113
              experimental::MakePtenScalarArrayFromVar(*ins_vector.front())));
        } else {  // ShapeTensorList
2114
          pt_kernel_context->EmplaceBackAttr(std::move(
2115 2116 2117 2118 2119
              experimental::MakePtenScalarArrayFromVarList(ins_vector)));
        }
      }
    } else if (attr_defs[i].type_index ==
               std::type_index(typeid(pten::Scalar))) {
2120 2121 2122
      // TODO(chenweihang): support other attrs later
      // TODO(zhangyunfei): Scalar should hold scaler type, and we should check
      // attribtue type by attr_defs
2123 2124 2125 2126
      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))) {
2127
          pt_kernel_context->EmplaceBackAttr(
2128 2129 2130
              std::move(pten::Scalar(BOOST_GET_CONST(float, attr))));
        } else if (std::type_index(attr.type()) ==
                   std::type_index(typeid(std::string))) {
2131
          pt_kernel_context->EmplaceBackAttr(
2132
              std::move(pten::Scalar(BOOST_GET_CONST(std::string, attr))));
2133 2134 2135 2136
        } else if (std::type_index(attr.type()) ==
                   std::type_index(typeid(int))) {
          pt_kernel_context->EmplaceBackAttr(
              std::move(pten::Scalar(BOOST_GET_CONST(int, attr))));
2137 2138 2139 2140 2141 2142
        } else {
          PADDLE_THROW(platform::errors::Unimplemented(
              "Unsupported cast op attribute `%s` to Scalar when construct "
              "KernelContext.",
              attr_names[i]));
        }
2143
      } else {
2144
        auto& ins_vector = ctx.inputs.at(attr_names[i]);
2145
        pt_kernel_context->EmplaceBackAttr(std::move(
2146
            experimental::MakePtenScalarFromVar(*ins_vector.front())));
2147
      }
2148

2149 2150
    } else {
      // TODO(chenweihang): support other attrs later
2151
      auto& attr = Attrs().at(attr_names[i]);
2152
      if (attr_defs[i].type_index == std::type_index(typeid(int))) {
2153
        pt_kernel_context->EmplaceBackAttr(BOOST_GET_CONST(int, attr));
2154
      } else if (attr_defs[i].type_index == std::type_index(typeid(float))) {
2155
        pt_kernel_context->EmplaceBackAttr(BOOST_GET_CONST(float, attr));
2156
      } else if (attr_defs[i].type_index == std::type_index(typeid(bool))) {
2157
        pt_kernel_context->EmplaceBackAttr(BOOST_GET_CONST(bool, attr));
H
hong 已提交
2158 2159 2160
      } else if (attr_defs[i].type_index ==
                 std::type_index(typeid(std::string))) {
        pt_kernel_context->EmplaceBackAttr(BOOST_GET_CONST(std::string, attr));
2161
      } else if (attr_defs[i].type_index ==
2162 2163 2164 2165
                 std::type_index(typeid(pten::DataType))) {
        auto data_type = pten::TransToPtenDataType(
            static_cast<framework::proto::VarType::Type>(
                BOOST_GET_CONST(int, attr)));
2166
        pt_kernel_context->EmplaceBackAttr(data_type);
2167 2168 2169 2170 2171 2172 2173 2174
      } 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());
2175
          pt_kernel_context->EmplaceBackAttr(vector_int64_attr);
2176 2177 2178
        }
        // TODO(YuanRisheng) Need support vector<int64_t> attr

H
hong 已提交
2179 2180 2181 2182
      } 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);
2183 2184
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
2185
            "Unsupported cast op attribute `%s` when construct "
2186 2187 2188 2189 2190 2191 2192
            "KernelContext.",
            attr_names[i]));
      }
    }
  }
}

Q
Qiao Longfei 已提交
2193
}  // namespace framework
L
liaogang 已提交
2194
}  // namespace paddle