infer_shape_pass.cc

// 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/ir/ipu/infer_shape_pass.h"
#include "paddle/fluid/framework/ir/graph_helper.h"
#include "paddle/fluid/framework/ir/pass_tester_helper.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/variable_helper.h"
#include "paddle/fluid/platform/device/ipu/ipu_backend.h"
#include "paddle/pten/core/ddim.h"

namespace paddle {
namespace framework {
namespace ir {

void InferShapePass::ApplyImpl(ir::Graph* graph) const {
  VLOG(10) << "enter InferShapePass::ApplyImpl";
  VLOG(10) << "Raw Graph: ";
  VLOG(10) << DebugString(graph);

  // Make batch_size fixed
  bool need_infer_shape = false;
  auto ipu_backend = platform::ipu::IpuBackend::GetInstance();
  auto micro_batch_size = ipu_backend->GetIpuStrategy()->micro_batch_size;
  auto feed_list = Get<std::vector<std::string>>("feed_list");
  for (auto node : graph->Nodes()) {
    if (!node->IsVar()) {
      continue;
    }
    bool is_feed = std::find(feed_list.begin(), feed_list.end(),
                             node->Name()) != feed_list.end();
    if (is_feed) {
      auto input_shape = node->Var()->GetShape();
      if (input_shape[0] <= -1) {
        input_shape[0] = micro_batch_size;
        node->Var()->SetShape(input_shape);
        need_infer_shape = true;
      }
      // int64->int32
      if (node->Var()->GetDataType() == proto::VarType::INT64) {
        node->Var()->SetDataType(proto::VarType::INT32);
      }
    }
  }

  // temp scope for shape inference
  if (need_infer_shape) {
    std::shared_ptr<paddle::framework::Scope> scope(
        new paddle::framework::Scope());
    for (auto node : graph->Nodes()) {
      if (!node->IsVar()) {
        continue;
      }
      auto var_desc = node->Var();
      auto* ptr = scope->Var(var_desc->Name());
      paddle::framework::InitializeVariable(ptr, var_desc->GetType());

      auto tensor = ptr->GetMutable<paddle::framework::LoDTensor>();
      tensor->Resize(pten::make_ddim(var_desc->GetShape()));
    }

    // infer shape
    auto nodes = ir::TopologySortOperations(*graph);
    for (auto node : nodes) {
      VLOG(10) << "InferShapePass: Infer shape for Op (" << node->Name() << ")";
      auto op_desc = node->Op();
      if (op_desc->Type() == "popart_optimizer") {
        continue;
      }
      auto op = paddle::framework::OpRegistry::CreateOp(*op_desc);
      paddle::framework::RuntimeContext ctx(op->Inputs(), op->Outputs(),
                                            *scope);
      op->RuntimeInferShape(*scope, paddle::platform::CPUPlace(), ctx);

      for (auto it = ctx.outputs.begin(); it != ctx.outputs.end(); it++) {
        for (int i = 0; i < it->second.size(); i++) {
          auto output_name = op_desc->Output(it->first)[i];
          auto dim =
              it->second[i]->GetMutable<paddle::framework::LoDTensor>()->dims();
          auto new_shape = pten::vectorize(dim);
          for (auto output_node : node->outputs) {
            if (output_node->Name() == output_name) {
              output_node->Var()->SetShape(new_shape);
              if (VLOG_IS_ON(10)) {
                std::ostringstream sout;
                sout << "InferShapePass: output[" << output_node->Name()
                     << "], infer shape:[";
                for (auto s : new_shape) {
                  sout << std::to_string(s) << ", ";
                }
                sout << "]";
                VLOG(10) << sout.str();
              }
            }
          }
        }
      }
      VLOG(10) << "InferShapePass: Infer shape for Op (" << node->Name()
               << ") finished";
    }
    // release the temp scope
    scope.reset();
  }

  VLOG(10) << "Post Graph: ";
  VLOG(10) << DebugString(graph);
  VLOG(10) << "leave InferShapePass::ApplyImpl";
}

}  // namespace ir
}  // namespace framework
}  // namespace paddle

REGISTER_PASS(infer_shape_pass, paddle::framework::ir::InferShapePass)
    .RequirePassAttr("feed_list");