infershape_utils.cc 28.0 KB
Newer Older
C
Chen Weihang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/* Copyright (c) 2022 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/infershape_utils.h"

17
#include <algorithm>
18 19
#include <string>

20
#include "paddle/fluid/framework/convert_utils.h"
C
Chen Weihang 已提交
21
#include "paddle/fluid/framework/framework.pb.h"
22
#include "paddle/fluid/framework/phi_utils.h"
C
Chen Weihang 已提交
23
#include "paddle/fluid/platform/enforce.h"
24
#include "paddle/phi/common/int_array.h"
25
#include "paddle/phi/common/scalar.h"
26 27 28 29 30
#include "paddle/phi/core/compat/arg_map_context.h"
#include "paddle/phi/core/compat/convert_utils.h"
#include "paddle/phi/core/compat/op_utils.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/infermeta_utils.h"
31
#include "paddle/phi/core/kernel_factory.h"
32
#include "paddle/phi/core/tensor_utils.h"
C
Chen Weihang 已提交
33 34 35 36

namespace paddle {
namespace framework {

37
class InferShapeArgumentMappingContext : public phi::ArgumentMappingContext {
C
Chen Weihang 已提交
38 39 40 41 42 43 44 45 46 47 48 49
 public:
  explicit InferShapeArgumentMappingContext(const InferShapeContext& ctx)
      : ctx_(ctx) {}

  bool HasInput(const std::string& name) const override {
    return ctx_.HasInput(name);
  }

  bool HasOutput(const std::string& name) const override {
    return ctx_.HasOutput(name);
  }

50 51 52 53
  bool HasAttr(const std::string& name) const override {
    return ctx_.HasAttr(name);
  }

C
Chen Weihang 已提交
54 55 56 57 58 59
  paddle::any Attr(const std::string& name) const override {
    auto& attr = ctx_.Attrs().GetAttr(name);
    return GetAttrValue(attr);
  }

  size_t InputSize(const std::string& name) const override {
60 61 62 63 64 65
    if (ctx_.HasInputs(name)) {
      return ctx_.Inputs(name).size();
    } else if (ctx_.HasInput(name)) {
      return 1;
    }
    return 0;
C
Chen Weihang 已提交
66 67 68 69 70 71 72 73
  }

  size_t OutputSize(const std::string& name) const override {
    return ctx_.Outputs(name).size();
  }

  bool IsDenseTensorInput(const std::string& name) const override {
    auto var_types = ctx_.GetInputsVarType(name);
74 75 76 77
    return std::all_of(var_types.begin(), var_types.end(),
                       [](const proto::VarType::Type& type) {
                         return type == proto::VarType::LOD_TENSOR;
                       });
C
Chen Weihang 已提交
78 79 80 81
  }

  bool IsSelectedRowsInput(const std::string& name) const override {
    auto var_types = ctx_.GetInputsVarType(name);
82 83 84 85
    return std::all_of(var_types.begin(), var_types.end(),
                       [](const proto::VarType::Type& type) {
                         return type == proto::VarType::SELECTED_ROWS;
                       });
C
Chen Weihang 已提交
86 87
  }

88 89
  bool IsDenseTensorVectorInput(const std::string& name) const override {
    auto var_types = ctx_.GetInputsVarType(name);
90 91 92 93
    return std::all_of(var_types.begin(), var_types.end(),
                       [](const proto::VarType::Type& type) {
                         return type == proto::VarType::LOD_TENSOR_ARRAY;
                       });
94 95
  }

96 97
  bool IsDenseTensorOutput(const std::string& name) const override {
    auto var_types = ctx_.GetOutputsVarType(name);
98 99 100 101
    return std::all_of(var_types.begin(), var_types.end(),
                       [](const proto::VarType::Type& type) {
                         return type == proto::VarType::LOD_TENSOR;
                       });
102 103 104 105
  }

  bool IsSelectedRowsOutput(const std::string& name) const override {
    auto var_types = ctx_.GetOutputsVarType(name);
106 107 108 109
    return std::all_of(var_types.begin(), var_types.end(),
                       [](const proto::VarType::Type& type) {
                         return type == proto::VarType::SELECTED_ROWS;
                       });
110 111
  }

112 113
  bool IsForInferShape() const override { return true; }

114 115
  bool IsRuntime() const override { return ctx_.IsRuntime(); }

C
Chen Weihang 已提交
116 117 118 119
 private:
  const InferShapeContext& ctx_;
};

120 121 122 123 124 125 126
int64_t CompatMetaTensor::numel() const {
  if (is_runtime_) {
    auto* var = BOOST_GET_CONST(Variable*, var_);
    return var->Get<Tensor>().numel();
  } else {
    auto* var = BOOST_GET_CONST(VarDesc*, var_);
    return var->ElementSize();
C
Chen Weihang 已提交
127
  }
128
}
C
Chen Weihang 已提交
129

130 131 132 133 134 135 136 137 138 139 140
DDim CompatMetaTensor::dims() const {
  if (is_runtime_) {
    auto* var = BOOST_GET_CONST(Variable*, var_);
    if (var->IsType<phi::DenseTensor>()) {
      return var->Get<phi::DenseTensor>().dims();
    } else if (var->IsType<phi::SelectedRows>()) {
      return var->Get<phi::SelectedRows>().dims();
    } else if (var->IsType<framework::LoDTensorArray>()) {
      // use tensor array size as dims
      auto& tensor_array = var->Get<framework::LoDTensorArray>();
      return phi::make_ddim({static_cast<int64_t>(tensor_array.size())});
C
Chen Weihang 已提交
141
    } else {
142 143 144
      PADDLE_THROW(platform::errors::Unimplemented(
          "Currently, only can get dims from DenseTensor or SelectedRows or "
          "DenseTensorArray."));
C
Chen Weihang 已提交
145
    }
146 147 148 149 150
  } else {
    auto* var = BOOST_GET_CONST(VarDesc*, var_);

    return var->GetShape().empty() ? phi::make_ddim({0UL})
                                   : phi::make_ddim(var->GetShape());
C
Chen Weihang 已提交
151
  }
152
}
C
Chen Weihang 已提交
153

154 155 156 157 158 159 160 161 162 163 164
phi::DataType CompatMetaTensor::dtype() const {
  if (is_runtime_) {
    auto* var = BOOST_GET_CONST(Variable*, var_);
    if (var->IsType<phi::DenseTensor>()) {
      return var->Get<phi::DenseTensor>().dtype();
    } else if (var->IsType<phi::SelectedRows>()) {
      return var->Get<phi::SelectedRows>().dtype();
    } else if (var->IsType<framework::LoDTensorArray>()) {
      // NOTE(chenweihang): do nothing
      // Unsupported get dtype from LoDTensorArray now
      return phi::DataType::UNDEFINED;
C
Chen Weihang 已提交
165
    } else {
166 167
      PADDLE_THROW(platform::errors::Unimplemented(
          "Currently, only can get dtype from DenseTensor or SelectedRows."));
C
Chen Weihang 已提交
168
    }
169 170 171
  } else {
    auto* var = BOOST_GET_CONST(VarDesc*, var_);
    return paddle::framework::TransToPhiDataType(var->GetDataType());
C
Chen Weihang 已提交
172
  }
173
}
C
Chen Weihang 已提交
174

175 176 177 178 179 180 181 182
DataLayout CompatMetaTensor::layout() const {
  if (is_runtime_) {
    auto* var = BOOST_GET_CONST(Variable*, var_);
    if (var->IsType<phi::DenseTensor>()) {
      return var->Get<phi::DenseTensor>().layout();
    } else if (var->IsType<phi::SelectedRows>()) {
      return var->Get<phi::SelectedRows>().layout();
    } else if (var->IsType<framework::LoDTensorArray>()) {
183
      // NOTE(chenweihang): do nothing
184 185 186 187 188 189
      // Unsupported get layout from LoDTensorArray now
      return phi::DataLayout::UNDEFINED;
    } else {
      PADDLE_THROW(platform::errors::Unimplemented(
          "Currently, only can get layout from DenseTensor or "
          "SelectedRows."));
C
Chen Weihang 已提交
190
    }
191 192 193 194
  } else {
    // NOTE(chenweihang): do nothing
    // Unsupported get layout for VarDesc now
    return DataLayout::UNDEFINED;
C
Chen Weihang 已提交
195
  }
196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216
}

void CompatMetaTensor::set_dims(const DDim& dims) {
  if (is_runtime_) {
    auto* var = BOOST_GET(Variable*, var_);
    if (var->IsType<phi::DenseTensor>()) {
      auto* tensor = var->GetMutable<phi::DenseTensor>();
      phi::DenseTensorUtils::GetMutableMeta(tensor)->dims = dims;
    } else if (var->IsType<phi::SelectedRows>()) {
      auto* tensor = var->GetMutable<phi::SelectedRows>()->mutable_value();
      phi::DenseTensorUtils::GetMutableMeta(tensor)->dims = dims;
    } else if (var->IsType<framework::LoDTensorArray>()) {
      auto* tensor_array = var->GetMutable<framework::LoDTensorArray>();
      // Note: Here I want enforce `tensor_array->size() == 0UL`, because
      // inplace using on LoDTensorArray is dangerous, but the unittest
      // `test_list` contains this behavior
      PADDLE_ENFORCE_EQ(dims.size(), 1UL,
                        platform::errors::InvalidArgument(
                            "LoDTensorArray can only have one dimension."));
      // only set the array size for LoDTensorArray input
      tensor_array->resize(dims[0]);
C
Chen Weihang 已提交
217
    } else {
218 219
      PADDLE_THROW(platform::errors::Unimplemented(
          "Currently, only can set dims from DenseTensor or SelectedRows."));
C
Chen Weihang 已提交
220
    }
221 222 223
  } else {
    auto* var = BOOST_GET(VarDesc*, var_);
    var->SetShape(vectorize(dims));
C
Chen Weihang 已提交
224
  }
225 226 227 228 229 230 231 232 233 234 235 236 237 238
}

void CompatMetaTensor::set_dtype(phi::DataType dtype) {
  if (is_runtime_) {
    auto* var = BOOST_GET(Variable*, var_);
    if (var->IsType<phi::DenseTensor>()) {
      auto* tensor = var->GetMutable<phi::DenseTensor>();
      phi::DenseTensorUtils::GetMutableMeta(tensor)->dtype = dtype;
    } else if (var->IsType<phi::SelectedRows>()) {
      auto* tensor = var->GetMutable<phi::SelectedRows>()->mutable_value();
      phi::DenseTensorUtils::GetMutableMeta(tensor)->dtype = dtype;
    } else if (var->IsType<framework::LoDTensorArray>()) {
      // NOTE(chenweihang): do nothing
      // Unsupported set dtype for LoDTensorArray now
C
Chen Weihang 已提交
239
    } else {
240 241
      PADDLE_THROW(platform::errors::Unimplemented(
          "Currently, only can set dtype from DenseTensor or SelectedRows."));
C
Chen Weihang 已提交
242
    }
243 244 245
  } else {
    auto* var = BOOST_GET(VarDesc*, var_);
    var->SetDataType(paddle::framework::TransToProtoVarType(dtype));
C
Chen Weihang 已提交
246
  }
247 248 249 250 251 252 253 254 255 256 257 258
}

void CompatMetaTensor::set_layout(DataLayout layout) {
  if (is_runtime_) {
    auto* var = BOOST_GET(Variable*, var_);
    if (var->IsType<phi::DenseTensor>()) {
      auto* tensor = var->GetMutable<phi::DenseTensor>();
      phi::DenseTensorUtils::GetMutableMeta(tensor)->layout = layout;
    } else if (var->IsType<phi::SelectedRows>()) {
      auto* tensor = var->GetMutable<phi::SelectedRows>()->mutable_value();
      phi::DenseTensorUtils::GetMutableMeta(tensor)->layout = layout;
    } else if (var->IsType<framework::LoDTensorArray>()) {
259
      // NOTE(chenweihang): do nothing
260 261 262 263 264
      // Unsupported set dtype for LoDTensorArray now
    } else {
      PADDLE_THROW(platform::errors::Unimplemented(
          "Currently, only can set layout from DenseTensor or "
          "SelectedRows."));
C
Chen Weihang 已提交
265
    }
266 267 268
  } else {
    // NOTE(chenweihang): do nothing
    // Unsupported set layout for VarDesc now
C
Chen Weihang 已提交
269
  }
270 271 272 273 274 275 276 277 278
}

void CompatMetaTensor::share_lod(const MetaTensor& meta_tensor) {
  if (is_runtime_) {
    auto* var = BOOST_GET(Variable*, var_);
    if (var->IsType<phi::DenseTensor>()) {
      auto* tensor = var->GetMutable<phi::DenseTensor>();
      phi::DenseTensorUtils::GetMutableMeta(tensor)->lod =
          static_cast<const CompatMetaTensor&>(meta_tensor).GetRuntimeLoD();
C
Chen Weihang 已提交
279
    } else {
280 281
      // NOTE(chenweihang): do nothing
      // only LoDTensor need to share lod
C
Chen Weihang 已提交
282
    }
283 284 285 286
  } else {
    auto* var = BOOST_GET(VarDesc*, var_);
    var->SetLoDLevel(
        static_cast<const CompatMetaTensor&>(meta_tensor).GetCompileTimeLoD());
C
Chen Weihang 已提交
287
  }
288 289 290 291 292 293 294 295 296 297 298 299
}

void CompatMetaTensor::share_dims(const MetaTensor& meta_tensor) {
  set_dims(meta_tensor.dims());
  if (is_runtime_) {
    auto* var = BOOST_GET(Variable*, var_);
    if (var->IsType<phi::SelectedRows>()) {
      auto* selected_rows = var->GetMutable<phi::SelectedRows>();
      auto& input_selected_rows =
          static_cast<const CompatMetaTensor&>(meta_tensor).GetSelectedRows();
      selected_rows->set_rows(input_selected_rows.rows());
      selected_rows->set_height(input_selected_rows.height());
300
    }
301
  }
302 303 304 305 306 307 308 309 310
}

void CompatMetaTensor::share_meta(const MetaTensor& meta_tensor) {
  share_dims(meta_tensor);
  set_dtype(meta_tensor.dtype());
  set_layout(meta_tensor.layout());
  // special case: share lod of LoDTensor
  share_lod(meta_tensor);
}
C
Chen Weihang 已提交
311

312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404
void CompatInferMetaContext::EmplaceBackInput(CompatMetaTensor input) {
  int index = compat_inputs_.size();
  compat_inputs_.emplace_back(std::move(input));
  input_range_.emplace_back(std::pair<int, int>(index, index + 1));
}
void CompatInferMetaContext::EmplaceBackOutput(CompatMetaTensor output) {
  int index = compat_outputs_.size();
  compat_outputs_.emplace_back(std::move(output));
  output_range_.emplace_back(std::pair<int, int>(index, index + 1));
}

void CompatInferMetaContext::EmplaceBackInputs(
    paddle::SmallVector<CompatMetaTensor, phi::kInputSmallVectorSize> inputs) {
  int index = compat_inputs_.size();
  input_range_.emplace_back(std::pair<int, int>(index, index + inputs.size()));
  compat_inputs_.insert(compat_inputs_.end(),
                        std::make_move_iterator(inputs.begin()),
                        std::make_move_iterator(inputs.end()));
}

void CompatInferMetaContext::EmplaceBackOutputs(
    paddle::SmallVector<CompatMetaTensor, phi::kOutputSmallVectorSize>
        outputs) {
  int index = compat_outputs_.size();
  output_range_.emplace_back(
      std::pair<int, int>(index, index + outputs.size()));
  compat_outputs_.insert(compat_outputs_.end(),
                         std::make_move_iterator(outputs.begin()),
                         std::make_move_iterator(outputs.end()));
}

const phi::MetaTensor& CompatInferMetaContext::InputAt(size_t idx) const {
  return compat_inputs_.at(idx);
}

paddle::optional<const phi::MetaTensor&>
CompatInferMetaContext::OptionalInputAt(size_t idx) const {
  const auto& input = compat_inputs_.at(idx);
  return input.initialized()
             ? paddle::optional<const phi::MetaTensor&>{input}
             : paddle::optional<const phi::MetaTensor&>{paddle::none};
}

std::vector<const phi::MetaTensor*> CompatInferMetaContext::InputsBetween(
    size_t start, size_t end) const {
  std::vector<const phi::MetaTensor*> result;
  result.reserve(end - start);

  for (size_t i = start; i < end; ++i) {
    auto& in = compat_inputs_.at(i);
    result.emplace_back(in.initialized() ? &in : nullptr);
  }

  return result;
}

paddle::optional<const std::vector<const phi::MetaTensor*>>
CompatInferMetaContext::OptionalInputsBetween(size_t start, size_t end) const {
  const auto& first = compat_inputs_.at(start);

  if (first.initialized()) {
    std::vector<const phi::MetaTensor*> result;
    result.reserve(end - start);

    for (size_t i = start; i < end; ++i) {
      auto& in = compat_inputs_.at(i);
      result.emplace_back(in.initialized() ? &in : nullptr);
    }

    return paddle::optional<const std::vector<const phi::MetaTensor*>>(result);
  }
  return paddle::optional<const std::vector<const phi::MetaTensor*>>(
      paddle::none);
}

phi::MetaTensor* CompatInferMetaContext::MutableOutputAt(size_t idx) {
  auto& out = compat_outputs_.at(idx);
  return out.initialized() ? &out : nullptr;
}

std::vector<phi::MetaTensor*> CompatInferMetaContext::MutableOutputBetween(
    size_t start, size_t end) {
  std::vector<phi::MetaTensor*> result;
  result.reserve(end - start);
  for (size_t i = start; i < end; ++i) {
    auto& out = compat_outputs_.at(i);
    result.emplace_back(out.initialized() ? &out : nullptr);
  }
  return result;
}

CompatInferMetaContext BuildInferMetaContext(InferShapeContext* ctx,
                                             const std::string& op_type) {
405
  // 1. get kernel args
406
  auto* arg_map_fn = ctx->GetPhiArgumentMappingFn();
407
  InferShapeArgumentMappingContext arg_map_context(*ctx);
408 409 410
  phi::KernelSignature signature = arg_map_fn
                                       ? (*arg_map_fn)(arg_map_context)
                                       : *ctx->GetPhiDefaultKernelSignature();
411 412 413
  VLOG(3) << "BuildInferMetaContext: op kernel signature - " << signature;

  // 2. build infermeta context
414
  CompatInferMetaContext infer_meta_context(
F
From00 已提交
415
      {ctx->IsRuntime(), ctx->IsRunMKLDNNKernel()});
416

417 418 419
  const auto& input_names = signature.input_names;
  const auto& attr_names = signature.attr_names;
  const auto& output_names = signature.output_names;
420

421 422 423
  const auto& args_def =
      phi::KernelFactory::Instance().GetFirstKernelArgsDef(signature.name);
  const auto& attr_defs = args_def.attribute_defs();
424

425
  for (auto& in_name : input_names) {
426
    if (ctx->HasInputs(in_name)) {
427
      auto input_var = std::move(ctx->GetInputVarPtrs(in_name));
428 429
      if (input_var.size() == 1) {
        infer_meta_context.EmplaceBackInput(
430
            std::move(CompatMetaTensor(input_var[0], ctx->IsRuntime())));
431
      } else {
432 433
        paddle::SmallVector<CompatMetaTensor, phi::kInputSmallVectorSize>
            inputs;
434
        for (const auto& in : input_var) {
435 436
          inputs.emplace_back(
              std::move(CompatMetaTensor(in, ctx->IsRuntime())));
437 438 439
        }
        infer_meta_context.EmplaceBackInputs(std::move(inputs));
      }
440
    } else {
441 442
      infer_meta_context.EmplaceBackInput(
          std::move(CompatMetaTensor(ctx->IsRuntime())));
443
    }
444
  }
445

446 447
  VLOG(6) << "BuildInferMetaContext: Done inputs";

448
  auto attr_reader = ctx->Attrs();
449
  for (size_t i = 0; i < attr_names.size(); ++i) {
450
    auto& attr_name = attr_names[i];
451
    if (attr_defs[i].type_index == phi::AttributeType::INT_ARRAY) {
452
      // When attr is a vector_tensor or tensor, transform it to IntArray
453
      if (ctx->HasInputs(attr_name) || ctx->HasInput(attr_name)) {
454
        auto infershape_inputs = std::move(ctx->GetInputVarPtrs(attr_name));
455
        if (ctx->IsRuntime()) {
456
          // If is in runtime, we will get tensor's value for IntArray
457 458 459 460 461 462 463 464
          // and push it into attrs
          std::vector<Variable*> vars;
          vars.reserve(infershape_inputs.size());
          for (size_t i = 0; i < infershape_inputs.size(); i++) {
            vars.push_back(BOOST_GET_CONST(Variable*, infershape_inputs[i]));
          }
          if (infershape_inputs.size() != 1) {
            infer_meta_context.EmplaceBackAttr(
465
                std::move(experimental::MakePhiIntArrayFromVarList(vars)));
466 467
          } else {
            infer_meta_context.EmplaceBackAttr(
468
                std::move(experimental::MakePhiIntArrayFromVar(*vars[0])));
469 470
          }
        } else {
471
          // If is not in runtime, we will set default value(-1) for IntArray
472 473
          std::vector<VarDesc*> vars;
          vars.reserve(infershape_inputs.size());
474
          for (size_t i = 0; i < infershape_inputs.size(); ++i) {
475 476
            vars.push_back(BOOST_GET_CONST(VarDesc*, infershape_inputs[i]));
          }
477

478
          int64_t num_ele = 0;
479 480 481
          if (vars.size() == 1) {
            num_ele = 1;
            const auto& tensor_dims = vars[0]->GetShape();
482 483 484
            for (size_t i = 0; i < tensor_dims.size(); ++i) {
              num_ele *= tensor_dims[i];
            }
485 486 487

            if (num_ele <= 0) {
              PADDLE_THROW(platform::errors::Unimplemented(
488
                  "Invalid number for construct phi::IntArray, expected "
489 490 491 492
                  "number > 0, but actually is %d. ",
                  num_ele));
            }

493
          } else {
494
            num_ele = vars.size();
495
          }
496
          phi::IntArray tensor_attr(std::vector<int32_t>(num_ele, -1));
497 498 499 500 501 502 503 504
          tensor_attr.SetFromTensor(true);
          infer_meta_context.EmplaceBackAttr(std::move(tensor_attr));
        }
      } else if (ctx->HasAttr(attr_name)) {
        auto& attr = attr_reader.GetAttr(attr_name);
        if (std::type_index(attr.type()) ==
            std::type_index(typeid(std::vector<int32_t>))) {
          infer_meta_context.EmplaceBackAttr(std::move(
505
              phi::IntArray(BOOST_GET_CONST(std::vector<int32_t>, attr))));
506 507 508
        } else if (std::type_index(attr.type()) ==
                   std::type_index(typeid(std::vector<int64_t>))) {
          infer_meta_context.EmplaceBackAttr(std::move(
509
              phi::IntArray(BOOST_GET_CONST(std::vector<int64_t>, attr))));
510 511 512
        } else if (std::type_index(attr.type()) ==
                   std::type_index(typeid(int))) {
          infer_meta_context.EmplaceBackAttr(
513
              phi::IntArray({BOOST_GET_CONST(int, attr)}));
514 515
        } else {
          PADDLE_THROW(platform::errors::Unimplemented(
516
              "Unsupported cast op attribute `%s` to IntArray when "
517
              "construct InferMetaContext.",
518 519 520
              attr_name));
        }
      }
521
    } else if (attr_defs[i].type_index == phi::AttributeType::SCALAR) {
522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542
      if (ctx->HasAttr(attr_name)) {
        // TODO(chentianyu03): support other attrs later
        auto& attr = attr_reader.GetAttr(attr_name);
        if (std::type_index(attr.type()) == std::type_index(typeid(float))) {
          infer_meta_context.EmplaceBackAttr(
              phi::Scalar(BOOST_GET_CONST(float, attr)));
        } else if (std::type_index(attr.type()) ==
                   std::type_index(typeid(std::string))) {
          infer_meta_context.EmplaceBackAttr(
              phi::Scalar(BOOST_GET_CONST(std::string, attr)));
        } else if (std::type_index(attr.type()) ==
                   std::type_index(typeid(int))) {
          infer_meta_context.EmplaceBackAttr(
              phi::Scalar(BOOST_GET_CONST(int, attr)));
        } else {
          PADDLE_THROW(platform::errors::Unimplemented(
              "Unsupported cast op attribute `%s` to Scalar when construct "
              "InferMetaContext.",
              attr_name));
        }
      } else if (ctx->HasInput(attr_name)) {
543
        auto infershape_input = std::move(ctx->GetInputVarPtrs(attr_name));
544 545 546 547
        if (infershape_input.size() == 1) {
          if (ctx->IsRuntime()) {
            Variable* var = BOOST_GET_CONST(Variable*, infershape_input[0]);
            infer_meta_context.EmplaceBackAttr(
548
                std::move(experimental::MakePhiScalarFromVar(*var)));
549 550 551 552 553 554 555 556 557 558 559 560
          } else {
            phi::Scalar tensor_scalar(-1);
            tensor_scalar.SetFromTensor(true);
            infer_meta_context.EmplaceBackAttr(std::move(tensor_scalar));
          }
        } else {
          PADDLE_THROW(platform::errors::InvalidArgument(
              "Invalid input.size() when cast op attribute `%s` to Scalar, "
              "expected 1, but actually is %d .",
              attr_name, infershape_input.size()));
        }
      }
561
    } else if (attr_defs[i].type_index == phi::AttributeType::SCALARS) {
562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604
      auto& attr = attr_reader.GetAttr(attr_name);
      if (std::type_index(attr.type()) ==
          std::type_index(typeid(std::vector<int32_t>))) {
        const auto& vec = BOOST_GET_CONST(std::vector<int32_t>, attr);
        std::vector<phi::Scalar> scalar_list;
        scalar_list.reserve(vec.size());
        for (const auto& val : vec) {
          scalar_list.emplace_back(val);
        }
        infer_meta_context.EmplaceBackAttr(std::move(scalar_list));
      } else if (std::type_index(attr.type()) ==
                 std::type_index(typeid(std::vector<int64_t>))) {
        const auto& vec = BOOST_GET_CONST(std::vector<int64_t>, attr);
        std::vector<phi::Scalar> scalar_list;
        scalar_list.reserve(vec.size());
        for (const auto& val : vec) {
          scalar_list.emplace_back(val);
        }
        infer_meta_context.EmplaceBackAttr(std::move(scalar_list));
      } else if (std::type_index(attr.type()) ==
                 std::type_index(typeid(std::vector<float>))) {
        const auto& vec = BOOST_GET_CONST(std::vector<float>, attr);
        std::vector<phi::Scalar> scalar_list;
        scalar_list.reserve(vec.size());
        for (const auto& val : vec) {
          scalar_list.emplace_back(val);
        }
        infer_meta_context.EmplaceBackAttr(std::move(scalar_list));
      } else if (std::type_index(attr.type()) ==
                 std::type_index(typeid(std::vector<double>))) {
        const auto& vec = BOOST_GET_CONST(std::vector<double>, attr);
        std::vector<phi::Scalar> scalar_list;
        scalar_list.reserve(vec.size());
        for (const auto& val : vec) {
          scalar_list.emplace_back(val);
        }
        infer_meta_context.EmplaceBackAttr(std::move(scalar_list));
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
            "Unsupported cast op attribute `%s` to vector<Scalar> when "
            "construct InferMetaContext.",
            attr_names[i]));
      }
605 606
    } else if (ctx->HasAttr(attr_name)) {
      // Emplace Back Attr according to the type of attr.
607
      auto& attr = attr_reader.GetAttr(attr_name);
608
      if (attr_defs[i].type_index == phi::AttributeType::BOOL) {
609
        infer_meta_context.EmplaceBackAttr(BOOST_GET_CONST(bool, attr));
610
      } else if (attr_defs[i].type_index == phi::AttributeType::INT32) {
611
        infer_meta_context.EmplaceBackAttr(BOOST_GET_CONST(int, attr));
612
      } else if (attr_defs[i].type_index == phi::AttributeType::INT64) {
613
        infer_meta_context.EmplaceBackAttr(BOOST_GET_CONST(int64_t, attr));
614
      } else if (attr_defs[i].type_index == phi::AttributeType::FLOAT32) {
615
        infer_meta_context.EmplaceBackAttr(BOOST_GET_CONST(float, attr));
616
      } else if (attr_defs[i].type_index == phi::AttributeType::STRING) {
617
        infer_meta_context.EmplaceBackAttr(BOOST_GET_CONST(std::string, attr));
618
      } else if (attr_defs[i].type_index == phi::AttributeType::BOOLS) {
619 620
        infer_meta_context.EmplaceBackAttr(
            BOOST_GET_CONST(std::vector<bool>, attr));
621
      } else if (attr_defs[i].type_index == phi::AttributeType::INT32S) {
622 623
        infer_meta_context.EmplaceBackAttr(
            BOOST_GET_CONST(std::vector<int>, attr));
624
      } else if (attr_defs[i].type_index == phi::AttributeType::INT64S) {
625 626 627 628 629 630 631 632 633 634 635
        if (std::type_index(attr.type()) ==
            std::type_index(typeid(std::vector<int>))) {
          // Emplace Back Attr according to the type of Phi_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());
          infer_meta_context.EmplaceBackAttr(vector_int64_attr);
        } else {
          infer_meta_context.EmplaceBackAttr(
              BOOST_GET_CONST(std::vector<int64_t>, attr));
        }
636
      } else if (attr_defs[i].type_index == phi::AttributeType::FLOAT32S) {
637 638
        infer_meta_context.EmplaceBackAttr(
            BOOST_GET_CONST(std::vector<float>, attr));
639
      } else if (attr_defs[i].type_index == phi::AttributeType::FLOAT64S) {
640 641
        infer_meta_context.EmplaceBackAttr(
            BOOST_GET_CONST(std::vector<double>, attr));
642
      } else if (attr_defs[i].type_index == phi::AttributeType::STRINGS) {
643 644
        infer_meta_context.EmplaceBackAttr(
            BOOST_GET_CONST(std::vector<std::string>, attr));
645
      } else if (attr_defs[i].type_index == phi::AttributeType::DATA_TYPE) {
646
        auto data_type = paddle::framework::TransToPhiDataType(
647 648 649
            static_cast<framework::proto::VarType::Type>(
                BOOST_GET_CONST(int, attr)));
        infer_meta_context.EmplaceBackAttr(data_type);
650
      } else {
651 652 653
        PADDLE_THROW(platform::errors::Unimplemented(
            "Unsupported attribute type is received when call "
            "InferShapeFunctor."));
654
      }
H
hong 已提交
655 656
    } else if (ctx->HasInput(attr_name)) {
      // convert from data
657
      if (attr_defs[i].type_index == phi::AttributeType::INT32) {
H
hong 已提交
658
        if (ctx->IsRuntime()) {
659
          auto infershape_inputs = std::move(ctx->GetInputVarPtrs(attr_name));
H
hong 已提交
660 661 662 663 664 665 666 667 668 669 670
          auto var_temp = BOOST_GET_CONST(Variable*, infershape_inputs[i]);
          auto val = experimental::MakePhiScalarFromVar(*var_temp);
          int32_t val_int = val.template to<int32_t>();
          infer_meta_context.EmplaceBackAttr(val_int);
        } else {
          infer_meta_context.EmplaceBackAttr(-1);
        }
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
            "Get value from variable only support int yet"));
      }
671 672 673
    }
  }

674 675
  VLOG(6) << "BuildInferMetaContext: Done attrs";

676
  for (auto& out_name : output_names) {
677
    if (ctx->HasOutputs(out_name, true)) {
678
      auto output_var = std::move(ctx->GetOutputVarPtrs(out_name));
679
      if (output_var.size() == 1) {
680 681
        infer_meta_context.EmplaceBackOutput(
            std::move(CompatMetaTensor(output_var[0], ctx->IsRuntime())));
682
      } else {
683 684
        paddle::SmallVector<CompatMetaTensor, phi::kOutputSmallVectorSize>
            outputs;
685
        for (const auto& out : output_var) {
686 687 688
          if (ctx->IsRuntime()) {
            if (BOOST_GET_CONST(Variable*, out)) {
              outputs.emplace_back(
689
                  std::move(CompatMetaTensor(out, ctx->IsRuntime())));
690 691 692 693
              continue;
            }
          } else if (BOOST_GET_CONST(VarDesc*, out)) {
            outputs.emplace_back(
694
                std::move(CompatMetaTensor(out, ctx->IsRuntime())));
695 696
            continue;
          }
697
          outputs.emplace_back(std::move(CompatMetaTensor(ctx->IsRuntime())));
698 699 700 701
        }
        infer_meta_context.EmplaceBackOutputs(std::move(outputs));
      }
    } else {
702 703
      infer_meta_context.EmplaceBackOutput(
          std::move(CompatMetaTensor(ctx->IsRuntime())));
704
    }
705 706
  }

707 708
  VLOG(6) << "BuildInferMetaContext: Done outputs";

709 710 711
  return infer_meta_context;
}

C
Chen Weihang 已提交
712 713
}  // namespace framework
}  // namespace paddle