data_transfer.cc 27.1 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
//
// 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.

#include "paddle/fluid/framework/new_executor/data_transfer.h"
16

17
#include "paddle/fluid/framework/convert_utils.h"
18 19
#include "paddle/phi/core/kernel_context.h"
#include "paddle/phi/core/kernel_factory.h"
20

21 22 23 24
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/phi/backends/onednn/onednn_context.h"
#endif

25 26 27 28 29 30 31 32 33
namespace paddle {
namespace framework {
namespace interpreter {

bool DataTranferHelper::apply(const OpKernelType& kernel_type_for_var,
                              const OpKernelType& expected_kernel_key,
                              const std::string& var_name,
                              std::string* new_var_name,
                              std::vector<OpFuncNode>* op_func_nodes,
34 35
                              bool use_local_scope,
                              bool is_fetch_v2) {
36 37 38 39 40
  bool is_transferred = false;
  auto* src_var_name = &var_name;

  // 1. layout transform
  if (need_layout_transform(kernel_type_for_var, expected_kernel_key)) {
41 42 43 44 45
    auto op = TransferLayout(*src_var_name,
                             new_var_name,
                             kernel_type_for_var.data_layout_,
                             expected_kernel_key.data_layout_,
                             var_scope_,
46
                             scope_,
47
                             is_fetch_v2);
L
Leo Chen 已提交
48
    if (op) {
49 50
      RunAndConstructOpFuncNode(
          op, *src_var_name, *new_var_name, op_func_nodes);
L
Leo Chen 已提交
51
    }
52 53 54 55 56 57
    // update src_var_name
    src_var_name = new_var_name;
    is_transferred = true;
  }
  // 2. dype transform
  if (need_dtype_transform(kernel_type_for_var, expected_kernel_key)) {
58 59 60 61 62
    auto op = TransferDtype(*src_var_name,
                            new_var_name,
                            kernel_type_for_var.data_type_,
                            expected_kernel_key.data_type_,
                            var_scope_,
63
                            scope_);
L
Leo Chen 已提交
64
    if (op) {
65 66
      RunAndConstructOpFuncNode(
          op, *src_var_name, *new_var_name, op_func_nodes);
L
Leo Chen 已提交
67
    }
68 69 70 71 72 73 74 75
    // update src_var_name
    src_var_name = new_var_name;
    is_transferred = true;
  }
  // 3. device transform
  if (need_device_transform(kernel_type_for_var, expected_kernel_key)) {
    auto src_place = kernel_type_for_var.place_;
    auto dst_place = expected_kernel_key.place_;
L
Leo Chen 已提交
76

77 78
    auto op = TransferDevice(
        *src_var_name, new_var_name, src_place, dst_place, var_scope_, scope_);
L
Leo Chen 已提交
79
    if (op) {
80 81
      RunAndConstructOpFuncNode(
          op, *src_var_name, *new_var_name, op_func_nodes);
L
Leo Chen 已提交
82
    }
83 84 85 86 87
    is_transferred = true;
  }
  return is_transferred;
}

88
void DataTranferHelper::RunAndConstructShareNode(
89 90
    const std::string& src_var_name,
    const std::string& dst_var_name,
91 92 93 94 95 96 97 98 99 100
    std::vector<OpFuncNode>* op_func_nodes) {
  VariableNameMap in_name_map = {{"X", {src_var_name}}};
  VariableNameMap out_name_map = {{"Out", {dst_var_name}}};
  AttributeMap attr_map;

  std::string op_type("share_data");
  auto& op_info = OpInfoMap::Instance().Get(op_type);
  auto op = std::shared_ptr<OperatorBase>(
      op_info.Creator()(op_type, in_name_map, out_name_map, attr_map));

101 102
  VLOG(3) << string::Sprintf(
      "Insert %s with %s -> %s.", op_type, src_var_name, dst_var_name);
103 104 105 106

  RunAndConstructOpFuncNode(op, src_var_name, dst_var_name, op_func_nodes);
}

107
void DataTranferHelper::RunAndConstructOpFuncNode(
108 109
    const std::shared_ptr<OperatorBase>& op,
    const std::string& var_name,
110 111 112 113 114 115
    const std::string& new_var_name,
    std::vector<OpFuncNode>* new_op_func_nodes) {
  auto& op_type = op->Type();

  // 1. Construct RuntimeContext
  RuntimeContext runtime_context({}, {});
116 117
  runtime_context.inputs["X"] = {scope_->FindVar(var_name)};
  runtime_context.outputs["Out"] = {scope_->Var(new_var_name)};
118
  InterpretercoreInferShapeContext infer_shape_ctx(*op, runtime_context);
119
  op.get()->Info().infer_shape_(&infer_shape_ctx);
120 121 122 123 124 125 126 127 128 129 130

  // 2. choose kernel

  // prepare a ptr to OperatorWithKernel
  OperatorBase* op_ptr = op.get();
  if (dynamic_cast<framework::OperatorWithKernel*>(op_ptr) == nullptr) {
    PADDLE_THROW(platform::errors::PreconditionNotMet(
        "%s should be OperatorWithKernel type.", op_ptr->Type()));
  }
  auto op_with_kernel = static_cast<framework::OperatorWithKernel*>(op_ptr);

131 132
  platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
  auto* dev_ctx = pool.Get(place_);
133 134 135 136 137 138 139 140 141
  auto exec_ctx = ExecutionContext(*op, Scope(), *dev_ctx, runtime_context);
  auto expected_kernel_key = op_with_kernel->GetExpectedKernelType(exec_ctx);

  VLOG(6) << "expected_kernel_key " << expected_kernel_key << "\n";
  VLOG(6) << "op_with_kernel Type() " << op_with_kernel->Type() << "\n";

  bool run_phi_kernel = false;

  // check if phi kernel exists
142 143
  if (phi::KernelFactory::Instance().HasCompatiblePhiKernel(
          op_with_kernel->Type())) {
144
    auto phi_kernel_key = op_with_kernel->ChoosePhiKernel(exec_ctx);
145
    auto phi_kernel_name = op_with_kernel->PhiKernelSignature()->name;
146
    VLOG(6) << "phi_kernel_key " << phi_kernel_key << "\n";
147
    VLOG(6) << "phi_kernel_name " << phi_kernel_name << "\n";
148

149 150 151
    if (op_with_kernel->PhiKernel()->IsValid()) {
      run_phi_kernel = true;
    }
152 153 154 155 156 157 158 159 160

    // For data transfer ops, they should not fallback to cpu.
    // Though they're device-independent operations,
    // their implementations are device-related.
    // For example, consider changing the layout of a gpu tensor
    // while the gpu kernel of transfer_layout op does not exist.
    // To use the cpu kernel, you must insert memcpy_d2h/mepcpy_h2d op
    // in addition. But such operation should not be done here.
    // Maybe in future we will support this.
161
  }
162 163 164 165 166

  // 3. Execute transfer op and construct OpFuncNode
  OpFuncNode new_op_func_node;
  new_op_func_node.input_index["X"] = {var_scope_->VarId(var_name)};
  new_op_func_node.output_index["Out"] = {var_scope_->VarId(new_var_name)};
167 168 169 170 171 172 173 174 175 176 177 178 179

  if (!run_phi_kernel) {
    op_with_kernel->ChooseKernel(exec_ctx);
    new_op_func_node.kernel_func_ = *op_with_kernel->kernel_func();
    new_op_func_node.kernel_func_(exec_ctx);
  } else {
    new_op_func_node.phi_kernel_ = op_with_kernel->PhiKernel();
    phi::KernelContext phi_kernel_context;
    op_with_kernel->BuildPhiKernelContext(
        runtime_context, dev_ctx, &phi_kernel_context);
    (*new_op_func_node.phi_kernel_)(&phi_kernel_context);
  }

180 181 182
  // NOTE(winter-wang): in npu device, D2H kernel is asynchronous. need to
  // explicit synchronization.
#ifdef PADDLE_WITH_ASCEND_CL
183 184 185 186 187 188
  if (op_type == kMemcpyD2H && platform::is_npu_place(dev_ctx->GetPlace())) {
    dev_ctx->Wait();
  }
#endif
#ifdef PADDLE_WITH_CUSTOM_DEVICE
  if (op_type == kMemcpyD2H && platform::is_custom_place(dev_ctx->GetPlace())) {
189 190 191
    dev_ctx->Wait();
  }
#endif
192 193 194 195 196 197 198 199 200 201
  // NOTE(Aurelius84): data_transform_op is expensive operation, so we tag them
  // as kQueueSync and execute them in thread pool.
  new_op_func_node.type_ = OpFuncType::kQueueSync;
  new_op_func_node.dev_ctx_ = dev_ctx;
  new_op_func_node.operator_base_ = op;
  VLOG(3) << "Run " << op_type << " done.";

  new_op_func_nodes->emplace_back(std::move(new_op_func_node));
}

L
Leo Chen 已提交
202 203 204 205 206 207
// Var is initialized && var contains tensor && tensor is initialized
bool IsTensorOfVarInitialized(Variable* var) {
  if (var->IsInitialized()) {
    if (var->IsType<LoDTensor>() || var->IsType<phi::SelectedRows>()) {
      return GetLoDTensorOrSelectedRowsValueFromVar(*var)->IsInitialized();
    } else if (var->IsType<LoDTensorArray>()) {
208 209
      return static_cast<const phi::DenseTensor*>(
                 &(var->Get<LoDTensorArray>()[0]))
L
Leo Chen 已提交
210 211 212 213 214 215
          ->IsInitialized();
    }
  }
  return false;
}

216 217 218 219 220 221 222
std::shared_ptr<OperatorBase> TransferLayout(const std::string& var_name,
                                             std::string* new_var_name,
                                             DataLayout in_layout,
                                             DataLayout out_layout,
                                             VariableScope* var_scope,
                                             framework::Scope* local_scope,
                                             bool is_fetch_v2) {
L
Leo Chen 已提交
223
#ifdef PADDLE_WITH_MKLDNN
224

L
Leo Chen 已提交
225 226 227
  // NOTE(zhiqiu): hot fix, follow the same logic in DataCopy() in fetch_op.cc
  if (in_layout == framework::DataLayout::kMKLDNN &&
      var_name == framework::GradVarName("Filter") && is_fetch_v2) {
228
    VLOG(4) << "Match special case(Filter && fetch_v2) " << var_name;
L
Leo Chen 已提交
229 230
    out_layout = framework::DataLayout::kNCHW;
  }
231

232 233
  if (in_layout == framework::DataLayout::ONEDNN &&
      out_layout != framework::DataLayout::ONEDNN) {
234
    auto target_layout = phi::OneDNNContext::tls().get_cur_paddle_data_layout();
235
    VLOG(4) << "TransDataLayoutFromOneDNN: " << in_layout << "->"
236 237 238 239 240 241 242 243 244
            << target_layout;

    if (out_layout == DataLayout::kNCHW &&
        var_name == framework::GradVarName("Filter")) {
      VLOG(4) << "Match special case(Filter) " << var_name;
      target_layout = out_layout;
    }
    out_layout = target_layout;
  }
L
Leo Chen 已提交
245 246
#endif

247
  // 1. Generate new_var_name and Initialize it
L
Leo Chen 已提交
248 249 250 251 252
  *new_var_name = var_name + "_layout_" +
                  std::to_string(static_cast<int>(in_layout)) + "_" +
                  std::to_string(static_cast<int>(out_layout));

  if (var_scope->HasVar(*new_var_name) &&
253
      IsTensorOfVarInitialized(local_scope->FindVar(*new_var_name))) {
L
Leo Chen 已提交
254 255 256 257
    // already has same var
    VLOG(4) << "Use cached variable: " << *new_var_name;
    return nullptr;
  }
258

L
Leo Chen 已提交
259
  auto* ptr = local_scope->Var(*new_var_name);
260
  auto var_type = local_scope->FindVar(var_name)->Type();
261
  InitializeVariable(ptr, static_cast<proto::VarType::Type>(var_type));
262 263 264
  VLOG(3) << "Create Variable " << *new_var_name
          << " locally, which pointer is " << ptr << "Variable Type "
          << var_type;
265 266
  var_scope->MutableDataTransferAddedVars().push_back(
      std::make_pair(*new_var_name, var_type));
267
  var_scope->AddVar(*new_var_name, nullptr);
268 269 270 271

  // 2. Construct VariableNameMap
  VariableNameMap in_name_map = {{"X", {var_name}}};
  VariableNameMap out_name_map = {{"Out", {*new_var_name}}};
272 273
  AttributeMap attr_map = {{"src_layout", static_cast<int>(in_layout)},
                           {"dst_layout", static_cast<int>(out_layout)}};
274

275
  // 3. Create transfer_layout_op
276 277 278 279 280
  std::string op_type("transfer_layout");
  auto& op_info = OpInfoMap::Instance().Get(op_type);
  auto op = std::shared_ptr<OperatorBase>(
      op_info.Creator()(op_type, in_name_map, out_name_map, attr_map));

281
  VLOG(3) << string::Sprintf("Insert %s for variable %s(%s) -> %s(%s).",
282 283 284 285
                             op_type,
                             var_name,
                             in_layout,
                             *new_var_name,
286
                             out_layout);
287 288 289 290 291 292 293
  return op;
}

std::shared_ptr<OperatorBase> TransferDtype(const std::string& var_name,
                                            std::string* new_var_name,
                                            proto::VarType::Type in_dtype,
                                            proto::VarType::Type out_dtype,
294
                                            framework::VariableScope* var_scope,
295 296
                                            framework::Scope* local_scope) {
  // 1. Generate new_var_name and Initialize it
L
Leo Chen 已提交
297 298 299 300
  *new_var_name = var_name + "_dtype_" +
                  std::to_string(static_cast<int>(in_dtype)) + "_" +
                  std::to_string(static_cast<int>(out_dtype));
  if (var_scope->HasVar(*new_var_name) &&
301
      IsTensorOfVarInitialized(local_scope->FindVar(*new_var_name))) {
L
Leo Chen 已提交
302 303 304 305
    // already has same var
    VLOG(4) << "Use cached variable: " << *new_var_name;
    return nullptr;
  }
306

L
Leo Chen 已提交
307
  auto* ptr = local_scope->Var(*new_var_name);
308
  auto var_type = local_scope->FindVar(var_name)->Type();
309
  InitializeVariable(ptr, static_cast<proto::VarType::Type>(var_type));
310 311 312
  VLOG(3) << "Create Variable " << *new_var_name
          << " locally, which pointer is " << ptr << "Variable Type "
          << var_type;
313 314
  var_scope->MutableDataTransferAddedVars().push_back(
      std::make_pair(*new_var_name, var_type));
315
  var_scope->AddVar(*new_var_name, nullptr);
316 317 318 319 320 321 322 323 324 325

  // 2. Construct VariableNameMap
  VariableNameMap in_name_map = {{"X", {var_name}}};
  VariableNameMap out_name_map = {{"Out", {*new_var_name}}};
  AttributeMap attr_map;
  attr_map["in_dtype"] = static_cast<int>(in_dtype);
  attr_map["out_dtype"] = static_cast<int>(out_dtype);
  // NOTE(Aurelius84): In whice case use_mkldnn = true?
  attr_map["use_mkldnn"] = false;

326
  // 3. Create transfer_dtype_op
327 328 329 330 331
  std::string op_type("transfer_dtype");
  auto& op_info = OpInfoMap::Instance().Get(op_type);
  auto op = std::shared_ptr<OperatorBase>(
      op_info.Creator()(op_type, in_name_map, out_name_map, attr_map));

332 333 334 335 336 337
  VLOG(3) << string::Sprintf("Insert %s with %s(%s) -> %s(%s).",
                             op_type,
                             var_name,
                             DataTypeToString(in_dtype),
                             *new_var_name,
                             DataTypeToString(out_dtype));
338 339 340 341 342 343 344 345 346 347
  return op;
}

std::shared_ptr<OperatorBase> TransferDevice(const std::string& var_name,
                                             std::string* new_var_name,
                                             const platform::Place& src_place,
                                             const platform::Place& dst_place,
                                             VariableScope* var_scope,
                                             framework::Scope* local_scope) {
  // 1. Generate new_var_name and Initialize it
L
Leo Chen 已提交
348 349 350
  *new_var_name = var_name + "_device_" + src_place.DebugString() + "_" +
                  dst_place.DebugString();

351
  if (var_scope->HasVar(*new_var_name) &&
352
      IsTensorOfVarInitialized(local_scope->FindVar(*new_var_name))) {
L
Leo Chen 已提交
353 354 355 356
    // already has same var
    VLOG(4) << "Use cached variable: " << *new_var_name;
    return nullptr;
  }
357

L
Leo Chen 已提交
358
  auto* ptr = local_scope->Var(*new_var_name);
359
  auto var_type = local_scope->FindVar(var_name)->Type();
360
  InitializeVariable(ptr, static_cast<proto::VarType::Type>(var_type));
361 362 363
  VLOG(3) << "Create Variable " << *new_var_name
          << " locally, which pointer is " << ptr << "Variable Type "
          << var_type;
364 365
  var_scope->MutableDataTransferAddedVars().push_back(
      std::make_pair(*new_var_name, var_type));
366
  var_scope->AddVar(*new_var_name, nullptr);
367 368 369 370 371

  // 2. Construct VariableNameMap
  VariableNameMap in_name_map = {{"X", {var_name}}};
  VariableNameMap out_name_map = {{"Out", {*new_var_name}}};

372
  // 3. Create memcpy_d2h_op or memcpy_h2d_op
373 374
  std::string op_type;
  AttributeMap attr_map;
375 376
  PADDLE_ENFORCE_EQ(platform::is_same_place(src_place, dst_place),
                    false,
377 378 379 380
                    platform::errors::PreconditionNotMet(
                        "Required src_place shall be different with dst_place, "
                        "but received same place: %s",
                        src_place));
L
Leo Chen 已提交
381
  if (IsSupportedHeterPlace(dst_place)) {
382
    op_type = kMemcpyH2D;
383 384 385 386 387 388
    int dst_place_type = platform::is_gpu_place(dst_place)      ? 0
                         : platform::is_npu_place(dst_place)    ? 1
                         : platform::is_ipu_place(dst_place)    ? 3
                         : platform::is_xpu_place(dst_place)    ? 2
                         : platform::is_custom_place(dst_place) ? 6
                                                                : -1;
389
    attr_map = {{"dst_place_type", dst_place_type}};
L
Leo Chen 已提交
390
  } else if (IsSupportedHeterPlace(src_place)) {
391 392 393 394 395 396 397 398 399 400
    op_type = kMemcpyD2H;
    int dst_place_type = platform::is_cpu_place(dst_place)           ? 0
                         : platform::is_cuda_pinned_place(dst_place) ? 1
                                                                     : -1;
    attr_map = {{"dst_place_type", dst_place_type}};
  } else {
    PADDLE_THROW(platform::errors::PreconditionNotMet(
        "Not support Memcpy typ : %s -> %s", src_place, dst_place));
  }

401 402 403 404
  auto& op_info = OpInfoMap::Instance().Get(op_type);
  auto op = std::shared_ptr<OperatorBase>(
      op_info.Creator()(op_type, in_name_map, out_name_map, attr_map));

405 406 407 408 409 410
  VLOG(3) << string::Sprintf("Insert %s with %s(%s) -> %s(%s).",
                             op_type,
                             var_name,
                             src_place,
                             *new_var_name,
                             dst_place);
411 412 413 414 415 416
  return op;
}

void ApplyDataTransform(const OpKernelType& expected_kernel_key,
                        const platform::Place& place,
                        VariableValueMap* ins_map_temp,
417
                        VariableValueMap* outs_map_temp,
418 419
                        VariableScope* var_scope,
                        OpFuncNode* op_func_node,
420 421
                        std::vector<OpFuncNode>* new_op_func_nodes,
                        bool use_local_scope) {
422 423 424
  Scope* local_scope = use_local_scope ? var_scope->GetMutableLocalScope()
                                       : var_scope->GetMutableScope();

425
  auto op_base = op_func_node->operator_base_.get();
426 427 428 429
  PADDLE_ENFORCE_NOT_NULL(op_base,
                          platform::errors::PreconditionNotMet(
                              "op_base is null, please pass a valid "
                              "op_base in apply_data_transform."));
430 431

  VariableNameMap new_ins(op_base->Inputs());
432
  VariableNameMap new_outs(op_base->Outputs());
433 434 435
  // record the no need transform variable index.
  std::unordered_set<int> no_data_transform_index;

L
Leo Chen 已提交
436 437 438
  const std::unordered_set<std::string>* no_buffer_ins = nullptr;
  auto& no_buffer_inferer = op_base->Info().NoNeedBufferVarsInferer();
  if (no_buffer_inferer) {
439 440
    no_buffer_ins = &(no_buffer_inferer(
        op_base->Inputs(), op_base->Outputs(), op_base->Attrs()));
L
Leo Chen 已提交
441 442 443 444 445
    if (no_buffer_ins->empty()) {
      no_buffer_ins = nullptr;
    }
  }

446
  bool transfered = false;
447
  DataTranferHelper data_transfer_helper(place, var_scope, local_scope);
448
  for (auto& var_name_item : *ins_map_temp) {
L
Leo Chen 已提交
449 450 451
    bool should_skip_input =
        no_buffer_ins && no_buffer_ins->count(var_name_item.first) > 0;

452 453
    for (size_t i = 0; i < var_name_item.second.size(); ++i) {
      auto var = var_name_item.second[i];
454
      auto var_name = new_ins[var_name_item.first].at(i);
455
      const phi::DenseTensor* tensor_in;
L
Leo Chen 已提交
456 457 458
      std::string new_var_name;
      bool is_transferred = false;

459
      if (var->IsType<LoDTensor>() || var->IsType<phi::SelectedRows>()) {
460 461
        tensor_in = GetLoDTensorOrSelectedRowsValueFromVar(*var);
      } else if (var->IsType<LoDTensorArray>()) {
462 463 464
        if (var->Get<LoDTensorArray>().size() == 0) {
          continue;
        }
465 466
        tensor_in = static_cast<const phi::DenseTensor*>(
            &(var->Get<LoDTensorArray>()[0]));
467
      } else {
468
        continue;
469
      }
L
Leo Chen 已提交
470
      // special case
471
      if (!tensor_in->IsInitialized()) {
L
Leo Chen 已提交
472 473 474 475 476 477 478 479 480 481 482 483 484
        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) &&
              (paddle::platform::MKLDNNDeviceContext::tls()
                   .get_cur_paddle_data_layout() == DataLayout::kNHWC)) {
485 486
            VLOG(7) << "Created reshaped dummy input based on MKL-DNN "
                       "phi::DenseTensor , "
L
Leo Chen 已提交
487 488 489
                       "but kNHWC layout"
                    << var_name_item.first << " in Operator "
                    << op_base->Type();
490 491 492 493 494 495 496
            auto op = TransferLayout(var_name,
                                     &new_var_name,
                                     tensor_in->layout(),
                                     DataLayout::kNHWC,
                                     var_scope,
                                     local_scope,
                                     op_base->Type() == "fetch_v2");
L
Leo Chen 已提交
497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512
            if (op) {
              data_transfer_helper.RunAndConstructOpFuncNode(
                  op, var_name, new_var_name, new_op_func_nodes);
            }
            is_transferred = true;
          } else {
            VLOG(7) << "Skip scanning input " << var_name_item.first
                    << " in Operator " << op_base->Type();
          }
#endif
        } else {
          continue;
        }
      } else {
        auto kernel_type_for_var =
            static_cast<const framework::OperatorWithKernel*>(op_base)
513 514
                ->GetKernelTypeForVar(
                    var_name_item.first, *tensor_in, expected_kernel_key);
L
Leo Chen 已提交
515
        // apply data transform
516 517 518 519 520 521 522 523
        is_transferred =
            data_transfer_helper.apply(kernel_type_for_var,
                                       expected_kernel_key,
                                       var_name,
                                       &new_var_name,
                                       new_op_func_nodes,
                                       use_local_scope,
                                       op_base->Type() == "fetch_v2");
524 525 526
      }

      if (is_transferred) {
527
        transfered = true;
528 529 530
        // update RuntimeContext.inputs and original op_func_node inputs
        op_func_node->input_index[var_name_item.first][i] =
            var_scope->VarId(new_var_name);
531
        var_name_item.second[i] = local_scope->FindVar(new_var_name);
532
        new_ins[var_name_item.first][i] = new_var_name;
533 534 535 536 537 538 539
        for (auto& pair : new_outs) {
          for (size_t j = 0; j < pair.second.size(); ++j) {
            VLOG(4) << pair.second[j] << " " << var_name;
            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;
540 541
              (*outs_map_temp)[pair.first][j] =
                  local_scope->FindVar(new_var_name);
542 543 544 545 546 547 548 549 550
              new_outs[pair.first][j] = new_var_name;
              op_func_node
                  ->inplace_back_map[var_scope->GetIdByName(new_var_name)] =
                  var_scope->GetIdByName(var_name);
              op_func_node->output_index[pair.first][j] =
                  var_scope->VarId(new_var_name);
            }
          }
        }
551 552 553 554 555 556 557 558 559 560 561 562 563 564 565
        // NOTE(Aurelius84): avoid deepcopy twice if we already insert data
        // transfer op.
        if (op_base->Type() == "fetch_v2") {
          op_base->SetAttr("deepcopy", false);
        }
      } else {
        // record no need data transformer input var_id
        VLOG(3) << op_base->Type()
                << " found no data_transform var: " << var_name
                << " with id: " << var_scope->VarId(var_name);
        no_data_transform_index.emplace(var_scope->VarId(var_name));
      }
    }
  }

566 567
  if (transfered) {
    // NOTE(zhiqiu): UPDATE the corresponding OeratorBase to make it consistent
568 569 570
    // with instruction.
    op_base->Inputs() = new_ins;
    op_base->Outputs() = new_outs;
571
  }
572 573 574
  op_func_node->no_data_transform_index = std::move(no_data_transform_index);
}

575 576 577 578 579 580 581
void HandleComplexGradToRealGrad(const OpFuncNode& op_func_node,
                                 const platform::Place& place,
                                 const VariableNameMap& out_names,
                                 VariableValueMap* out_vars,
                                 VariableScope* var_scope,
                                 std::vector<OpFuncNode>* op_func_nodes,
                                 framework::Scope* local_scope) {
582
  DataTranferHelper data_transfer_helper(place, var_scope, local_scope);
583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614
  for (auto& var_name_item : out_names) {
    std::vector<Variable*>& vars = out_vars->at(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 = framework::GradOriginalVarName(var_name);
      // only focus on gradient var
      if (var_name == orig_var_name) {
        VLOG(3) << "skip " << var_name << " with same name as "
                << orig_var_name;
        continue;
      }
      auto* grad_var = vars[i];
      // skip nullptr var
      if (grad_var == nullptr) {
        VLOG(3) << "skip grad_var with nullptr";
        continue;
      }
      // don't process LoDTensorArray temporarily,
      // add support if necessary for complex number calculations in the future
      if (!framework::VarIsTensor(*grad_var)) {
        VLOG(3) << "skip grad_var with LoDTensorArray type";
        continue;
      }
      auto* grad_tensor =
          framework::GetMutableLoDTensorOrSelectedRowsValueFromVar(grad_var);
      // skip nullptr tensor
      if (grad_tensor == nullptr || !grad_tensor->IsInitialized()) {
        VLOG(3) << "skip with grad_tensor not IsInitialized";
        continue;
      }
      // only focus on complex dtype now
615
      auto src_type = framework::TransToProtoVarType(grad_tensor->dtype());
616 617 618 619 620 621
      if (!framework::IsComplexType(src_type)) {
        VLOG(3) << "skip grad_tensor with not complexType";
        continue;
      }

      // 2. find forward var & check whether need to cast
622
      auto* var = local_scope->FindVar(orig_var_name);
623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638
      // if forward var not exists, do nothing
      if (var == nullptr) {
        VLOG(3) << "skip " << orig_var_name << " with not found in var_scope";
        continue;
      }
      if (!framework::VarIsTensor(*var)) {
        VLOG(3) << "skip " << orig_var_name << " with LoDTensorArray.";
        continue;
      }
      const auto* tensor =
          framework::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
639
      auto dst_type = framework::TransToProtoVarType(tensor->dtype());
640 641 642 643 644 645 646 647 648 649 650 651 652 653
      // only focus on real dtype and need casting
      if (framework::IsComplexType(dst_type)) {
        continue;
      }

      // 3. cast complex grad to real grad inplacely
      VLOG(3) << "Transform " << framework::DataTypeToString(src_type)
              << " var `" << var_name << "` to "
              << framework::DataTypeToString(dst_type)
              << " real var in static graph.";

      // NOTE(Aurelius84): Consider to define a complex2real op to deal this
      // case.
      std::string new_var_name;
654 655 656 657 658 659
      auto op = TransferDtype(
          var_name, &new_var_name, src_type, dst_type, var_scope, local_scope);
      data_transfer_helper.RunAndConstructOpFuncNode(
          op, var_name, new_var_name, op_func_nodes);
      data_transfer_helper.RunAndConstructShareNode(
          new_var_name, var_name, op_func_nodes);
660 661 662 663
    }
  }
}

664 665 666
}  // namespace interpreter
}  // namespace framework
}  // namespace paddle