infershape_utils.cc 27.6 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
  }

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

  bool IsDenseTensorInput(const std::string& name) const override {
73 74 75 76 77
    auto var_type = ctx_.GetInputVarType(name);
    return var_type == proto::VarType::LOD_TENSOR;
  }

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

  bool IsSelectedRowsInput(const std::string& name) const override {
86 87
    auto var_type = ctx_.GetInputVarType(name);
    return var_type == proto::VarType::SELECTED_ROWS;
C
Chen Weihang 已提交
88 89
  }

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

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

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

114 115
  bool IsForInferShape() const override { return true; }

116 117
  bool IsRuntime() const override { return ctx_.IsRuntime(); }

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

122 123 124 125 126 127 128
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 已提交
129
  }
130
}
C
Chen Weihang 已提交
131

132 133 134 135 136 137 138 139 140 141 142
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 已提交
143
    } else {
144 145 146
      PADDLE_THROW(platform::errors::Unimplemented(
          "Currently, only can get dims from DenseTensor or SelectedRows or "
          "DenseTensorArray."));
C
Chen Weihang 已提交
147
    }
148 149 150 151 152
  } else {
    auto* var = BOOST_GET_CONST(VarDesc*, var_);

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

156 157 158 159 160 161 162 163 164 165 166
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 已提交
167
    } else {
168 169
      PADDLE_THROW(platform::errors::Unimplemented(
          "Currently, only can get dtype from DenseTensor or SelectedRows."));
C
Chen Weihang 已提交
170
    }
171 172 173
  } else {
    auto* var = BOOST_GET_CONST(VarDesc*, var_);
    return paddle::framework::TransToPhiDataType(var->GetDataType());
C
Chen Weihang 已提交
174
  }
175
}
C
Chen Weihang 已提交
176

177 178 179 180 181 182 183 184
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>()) {
185
      // NOTE(chenweihang): do nothing
186 187 188 189 190 191
      // 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 已提交
192
    }
193 194 195 196
  } else {
    // NOTE(chenweihang): do nothing
    // Unsupported get layout for VarDesc now
    return DataLayout::UNDEFINED;
C
Chen Weihang 已提交
197
  }
198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218
}

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 已提交
219
    } else {
220 221
      PADDLE_THROW(platform::errors::Unimplemented(
          "Currently, only can set dims from DenseTensor or SelectedRows."));
C
Chen Weihang 已提交
222
    }
223 224 225
  } else {
    auto* var = BOOST_GET(VarDesc*, var_);
    var->SetShape(vectorize(dims));
C
Chen Weihang 已提交
226
  }
227 228 229 230 231 232 233 234 235 236 237 238 239 240
}

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 已提交
241
    } else {
242 243
      PADDLE_THROW(platform::errors::Unimplemented(
          "Currently, only can set dtype from DenseTensor or SelectedRows."));
C
Chen Weihang 已提交
244
    }
245 246 247
  } else {
    auto* var = BOOST_GET(VarDesc*, var_);
    var->SetDataType(paddle::framework::TransToProtoVarType(dtype));
C
Chen Weihang 已提交
248
  }
249 250 251 252 253 254 255 256 257 258 259 260
}

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>()) {
261
      // NOTE(chenweihang): do nothing
262 263 264 265 266
      // Unsupported set dtype for LoDTensorArray now
    } else {
      PADDLE_THROW(platform::errors::Unimplemented(
          "Currently, only can set layout from DenseTensor or "
          "SelectedRows."));
C
Chen Weihang 已提交
267
    }
268 269 270
  } else {
    // NOTE(chenweihang): do nothing
    // Unsupported set layout for VarDesc now
C
Chen Weihang 已提交
271
  }
272 273 274 275 276 277 278 279 280
}

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 已提交
281
    } else {
282 283
      // NOTE(chenweihang): do nothing
      // only LoDTensor need to share lod
C
Chen Weihang 已提交
284
    }
285 286 287 288
  } else {
    auto* var = BOOST_GET(VarDesc*, var_);
    var->SetLoDLevel(
        static_cast<const CompatMetaTensor&>(meta_tensor).GetCompileTimeLoD());
C
Chen Weihang 已提交
289
  }
290 291 292 293 294 295 296 297 298 299 300 301
}

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());
302
    }
303
  }
304 305 306 307 308 309 310 311 312
}

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 已提交
313

314 315 316 317 318 319 320 321 322 323 324 325
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(
C
Chen Weihang 已提交
326
    paddle::small_vector<CompatMetaTensor, phi::kInputSmallVectorSize> inputs) {
327 328 329 330 331 332 333 334
  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(
C
Chen Weihang 已提交
335
    paddle::small_vector<CompatMetaTensor, phi::kOutputSmallVectorSize>
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 405 406
        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) {
407
  // 1. get kernel args
408
  auto* arg_map_fn = ctx->GetPhiArgumentMappingFn();
409
  InferShapeArgumentMappingContext arg_map_context(*ctx);
410 411 412
  phi::KernelSignature signature = arg_map_fn
                                       ? (*arg_map_fn)(arg_map_context)
                                       : *ctx->GetPhiDefaultKernelSignature();
413 414 415
  VLOG(3) << "BuildInferMetaContext: op kernel signature - " << signature;

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

419 420 421
  const auto& input_names = signature.input_names;
  const auto& attr_names = signature.attr_names;
  const auto& output_names = signature.output_names;
422

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

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

448 449
  VLOG(6) << "BuildInferMetaContext: Done inputs";

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

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

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

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

666 667
  VLOG(6) << "BuildInferMetaContext: Done attrs";

668
  for (auto& out_name : output_names) {
669
    if (ctx->HasOutputs(out_name, true)) {
670
      auto output_var = std::move(ctx->GetOutputVarPtrs(out_name));
671
      if (output_var.size() == 1) {
672 673
        infer_meta_context.EmplaceBackOutput(
            std::move(CompatMetaTensor(output_var[0], ctx->IsRuntime())));
674
      } else {
C
Chen Weihang 已提交
675
        paddle::small_vector<CompatMetaTensor, phi::kOutputSmallVectorSize>
676
            outputs;
677
        for (const auto& out : output_var) {
678 679 680
          if (ctx->IsRuntime()) {
            if (BOOST_GET_CONST(Variable*, out)) {
              outputs.emplace_back(
681
                  std::move(CompatMetaTensor(out, ctx->IsRuntime())));
682 683 684 685
              continue;
            }
          } else if (BOOST_GET_CONST(VarDesc*, out)) {
            outputs.emplace_back(
686
                std::move(CompatMetaTensor(out, ctx->IsRuntime())));
687 688
            continue;
          }
689
          outputs.emplace_back(std::move(CompatMetaTensor(ctx->IsRuntime())));
690 691 692 693
        }
        infer_meta_context.EmplaceBackOutputs(std::move(outputs));
      }
    } else {
694 695
      infer_meta_context.EmplaceBackOutput(
          std::move(CompatMetaTensor(ctx->IsRuntime())));
696
    }
697 698
  }

699 700
  VLOG(6) << "BuildInferMetaContext: Done outputs";

701 702 703
  return infer_meta_context;
}

C
Chen Weihang 已提交
704 705
}  // namespace framework
}  // namespace paddle