| 442 | } |
| 443 | |
| 444 | void RuntimeInferShapeContext::ShareLoD(const std::string& in, |
| 445 | const std::string& out, |
| 446 | size_t i, |
| 447 | size_t j) const { |
| 448 | if (can_skip_lod_) { |
| 449 | return; |
| 450 | } |
| 451 | auto in_it = ctx_.inputs.find(in); |
| 452 | auto out_it = ctx_.outputs.find(out); |
| 453 | PADDLE_ENFORCE_NE(in_it, |
| 454 | ctx_.inputs.end(), |
| 455 | common::errors::NotFound("Input %s does not exist.", in)); |
| 456 | PADDLE_ENFORCE_NE(out_it, |
| 457 | ctx_.outputs.end(), |
| 458 | common::errors::NotFound("Output %s does not exist.", out)); |
| 459 | PADDLE_ENFORCE_LT(i, |
| 460 | in_it->second.size(), |
| 461 | common::errors::InvalidArgument( |
| 462 | "The index of input dimension is out of range, " |
| 463 | "excepted index less than %zu, but received %zu.", |
| 464 | in_it->second.size(), |
| 465 | i)); |
| 466 | PADDLE_ENFORCE_LT(j, |
| 467 | out_it->second.size(), |
| 468 | common::errors::InvalidArgument( |
| 469 | "The index of output dimension is out of range, " |
| 470 | "excepted index less than %zu, but received %zu.", |
| 471 | out_it->second.size(), |
| 472 | j)); |
| 473 | |
| 474 | Variable* in_var = in_it->second.at(i); |
| 475 | if (!in_var->IsType<DenseTensor>()) return; |
| 476 | Variable* out_var = out_it->second.at(j); |
| 477 | PADDLE_ENFORCE_EQ( |
| 478 | out_var->IsType<DenseTensor>(), |
| 479 | true, |
| 480 | common::errors::InvalidArgument( |
| 481 | "The %zu-th output of Output(%s) must be phi::DenseTensor.", j, out)); |
| 482 | auto& in_tensor = in_var->Get<DenseTensor>(); |
| 483 | auto* out_tensor = out_var->GetMutable<DenseTensor>(); |
| 484 | out_tensor->set_lod(in_tensor.lod()); |
| 485 | |
| 486 | // TODO(dzhwinter) : reuse ShareLoD in most operators. |
| 487 | // Need to call ShareLayout explicitly in sequence related ops. |
| 488 | // Shall we have a better method to shared info between in/out phi::DenseTensor? |
| 489 | #ifdef PADDLE_WITH_DNNL |
| 490 | // Fix me: ugly workaround below |
| 491 | // Correct solution: |
| 492 | // set_layout() should NOT be called here (i.e. ShareLoD). Instead, |
| 493 | // layout of output tensor should be set "manually" in Compute() |
| 494 | // of each OPKernel. The reason layout should NOT be shared between |
| 495 | // input and output "automatically" (now by InferShape()->ShareLoD()) |
| 496 | // is that layout transform may occur after InferShape(). |
| 497 | // Workaround: |
| 498 | // Skip set_layout() when input layout is kMKLDNN |
| 499 | // This is to avoid kMKLDNN is populated wrongly into a non-MKLDNN |
| 500 | // OPKernel. In all OneDNN OPkernel, set_layout(kMKLDNN) should be called |
| 501 | // in Compute() |
no test coverage detected