Skip to content

Commit

Permalink
infer shape pass and avg shard pass (PaddlePaddle#165)
Browse files Browse the repository at this point in the history
* infer shape pass and avg shard pass

* rm useless test
  • Loading branch information
yaozhixin committed Sep 17, 2021
1 parent 142c9e8 commit 11708d4
Show file tree
Hide file tree
Showing 27 changed files with 297 additions and 76 deletions.
2 changes: 2 additions & 0 deletions paddle/fluid/framework/ipu/ipu_backend.cc
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,8 @@ void IpuBackend::SetIpuStrategy(const IpuStrategy& strategy) {
executor_->SetIpuStrategy(strategy);
}

const IpuStrategy* IpuBackend::GetIpuStrategy() { return ipu_strategy_; }

size_t IpuBackend::GetNumDevices() {
// IpuModel
bool ipu_model = GetBoolEnv("POPLAR_IPUMODEL");
Expand Down
3 changes: 2 additions & 1 deletion paddle/fluid/framework/ipu/ipu_backend.h
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,8 @@ limitations under the License. */
#include <popart/names.hpp>

#include "paddle/fluid/framework/feed_fetch_type.h"
#include "paddle/fluid/framework/ipu/ipu_compiler.h"
#include "paddle/fluid/framework/ipu/device.h"
#include "paddle/fluid/framework/ipu/ipu_compiler.h"
#include "paddle/fluid/framework/ipu/ipu_executor.h"
#include "paddle/fluid/framework/ipu/ipu_strategy.h"
#include "paddle/fluid/framework/scope.h"
Expand Down Expand Up @@ -64,6 +64,7 @@ class IpuBackend {

void SetScope(Scope &scope);
void SetIpuStrategy(const IpuStrategy &strategy);
const IpuStrategy *GetIpuStrategy();

// Device
size_t GetNumDevices();
Expand Down
27 changes: 0 additions & 27 deletions paddle/fluid/framework/ipu/ipu_compiler.cc
Original file line number Diff line number Diff line change
Expand Up @@ -281,19 +281,6 @@ void Compiler::SetIpuIndexStage(const std::vector<std::string>& tensor_ids,
auto tensor_ids_set =
std::set<std::string>(tensor_ids.begin(), tensor_ids.end());

// TODO(yaozhixin) : workaround for Paddle inference
// auto debug_id =
// builder_->getInt64NodeAttribute("__debug_info_id", tensor_ids_set);
// if (debug_id < 300) {
// builder_->virtualGraph(tensor_ids_set, 0);
// } else if (debug_id < 600) {
// builder_->virtualGraph(tensor_ids_set, 1);
// } else if (debug_id < 900) {
// builder_->virtualGraph(tensor_ids_set, 2);
// } else {
// builder_->virtualGraph(tensor_ids_set, 3);
// }

if (op_desc->HasAttr(sIpuIndexAttr)) {
auto ipu_index = BOOST_GET_CONST(int, op_desc->GetAttr(sIpuIndexAttr));
builder_->virtualGraph(tensor_ids_set, ipu_index);
Expand All @@ -313,20 +300,6 @@ void Compiler::SetIpuIndexStage(const std::string& tensor_id,
const OpDesc* op_desc) {
VLOG(10) << "enter Compiler::SetIpuIndexStage";

// TODO(yaozhixin) : workaround for Paddle inference
// auto tensor_id_set = std::set<std::string>{tensor_id};
// auto debug_id =
// builder_->getInt64NodeAttribute("__debug_info_id", tensor_id_set);
// if (debug_id < 300) {
// builder_->virtualGraph(tensor_id, 0);
// } else if (debug_id < 600) {
// builder_->virtualGraph(tensor_id, 1);
// } else if (debug_id < 900) {
// builder_->virtualGraph(tensor_id, 2);
// } else {
// builder_->virtualGraph(tensor_id, 3);
// }

if (op_desc->HasAttr(sIpuIndexAttr)) {
auto ipu_index = BOOST_GET_CONST(int, op_desc->GetAttr(sIpuIndexAttr));
builder_->virtualGraph(tensor_id, ipu_index);
Expand Down
2 changes: 2 additions & 0 deletions paddle/fluid/framework/ipu/ipu_strategy.h
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ struct IpuStrategy {
bool save_init_onnx = false;
bool save_last_onnx = true;
popart::SessionOptions popart_options_;
bool need_infer_shape = false;
bool need_avg_shard = false;
};

} // namespace ipu
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,9 @@ namespace {
Node *elementwise_op_handler(Graph *graph, Node *node,
const std::string &type) {
auto *op = node->Op();
auto x_shape = op->Block()->FindVar(op->Input("X").front())->GetShape();
auto x_shape = GetInputNode("X", node)->Var()->GetShape();
int64_t x_rank = x_shape.size();
auto y_shape = op->Block()->FindVar(op->Input("Y").front())->GetShape();
auto y_shape = GetInputNode("Y", node)->Var()->GetShape();
int64_t y_rank = y_shape.size();

auto axis = BOOST_GET_CONST(int, op->GetAttr("axis"));
Expand Down
12 changes: 4 additions & 8 deletions paddle/fluid/framework/ipu/popart_canonicalization/math_ops.cc
Original file line number Diff line number Diff line change
Expand Up @@ -68,8 +68,8 @@ Node *mul_handler(Graph *graph, Node *node) {
auto *op = node->Op();
auto x_num_col_dims = BOOST_GET_CONST(int, op->GetAttr("x_num_col_dims"));
auto y_num_col_dims = BOOST_GET_CONST(int, op->GetAttr("y_num_col_dims"));
auto x_shape_ = op->Block()->FindVar(op->Input("X")[0])->GetShape();
auto y_shape_ = op->Block()->FindVar(op->Input("Y")[0])->GetShape();
auto x_shape_ = GetInputNode("X", node)->Var()->GetShape();
auto y_shape_ = GetInputNode("Y", node)->Var()->GetShape();

// build the shape for reshape
std::vector<int64_t> reshape_shape_{};
Expand All @@ -89,10 +89,6 @@ Node *mul_handler(Graph *graph, Node *node) {
CreateBaseOp(graph, node, "popart_matmul",
{x_flatten->outputs[0], y_flatten->outputs[0]}, {}, {});

// TODO(yaozhixin): workaround for Paddle inference
// reshape_shape_[0] = 1;
// reshape_shape_[1] = -1;

auto reshape_const = CreateConst(
graph, node, {}, {},
{{"value", reshape_shape_},
Expand Down Expand Up @@ -170,7 +166,7 @@ Node *scale_handler(Graph *graph, Node *node) {
auto bias_ = BOOST_GET_CONST(float, op->GetAttr("bias"));
auto bias_after_scale_ =
BOOST_GET_CONST(bool, op->GetAttr("bias_after_scale"));
auto data_type_ = op->Block()->FindVar(op->Input("X")[0])->GetDataType();
auto data_type_ = GetInputNode("X", node)->Var()->GetDataType();

auto new_node_bias_var =
CreateConst(graph, node, {}, {}, {{"value", std::vector<float>{bias_}},
Expand Down Expand Up @@ -219,7 +215,7 @@ Node *cross_entropy2_handler(Graph *graph, Node *node) {
auto ignoreIndex = BOOST_GET_CONST(int, op->GetAttr("ignore_index"));
auto new_cast = CreateCast(graph, node, {GetInputNode("Label", node)}, {},
proto::VarType::INT32);
auto label_shape_ = op->Block()->FindVar(op->Input("Label")[0])->GetShape();
auto label_shape_ = GetInputNode("Label", node)->Var()->GetShape();
if (label_shape_.size() == 1) {
return CreateBaseOp(graph, node, "popart_nllloss",
{GetInputNode("X", node), new_cast->outputs[0]},
Expand Down
9 changes: 1 addition & 8 deletions paddle/fluid/framework/ipu/popart_canonicalization/nn_ops.cc
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ Node *instance_norm_handler(Graph *graph, Node *node) {
Node *layer_norm_handler(Graph *graph, Node *node) {
auto *op = node->Op();
auto begin_norm_axis_ = BOOST_GET_CONST(int, op->GetAttr("begin_norm_axis"));
auto input_shape_ = op->Block()->FindVar(op->Input("X")[0])->GetShape();
auto input_shape_ = GetInputNode("X", node)->Var()->GetShape();

std::vector<int64_t> norm_shape_{1, 1};
for (int i = 0; i < input_shape_.size(); i++) {
Expand All @@ -191,13 +191,6 @@ Node *layer_norm_handler(Graph *graph, Node *node) {
}
}

// TODO(yaozhixin): workaround for Paddle inference
// if (norm_shape_[0] < -1) {
// norm_shape_[0] = -1;
// }
// input_shape_[0] = 1;
// input_shape_[1] = -1;

auto attrs1 = AttributeMap{
{"value", norm_shape_},
{"dims", std::vector<int64_t>{static_cast<int64_t>(norm_shape_.size())}},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,7 @@ Node *gather_handler(Graph *graph, Node *node) {
Node *squeeze_handler(Graph *graph, Node *node) {
auto *op = node->Op();
auto axes_ = BOOST_GET_CONST(std::vector<int>, op->GetAttr("axes"));
auto input_shape_ = op->Block()->FindVar(op->Input("X")[0])->GetShape();
auto input_shape_ = GetInputNode("X", node)->Var()->GetShape();

std::vector<int64_t> axes{axes_.begin(), axes_.end()};
if (axes_.empty()) {
Expand Down Expand Up @@ -170,7 +170,7 @@ Node *cast_handler(Graph *graph, Node *node) {
Node *lookup_table_handler(Graph *graph, Node *node) {
auto *op = node->Op();
auto padding_idx_ = BOOST_GET_CONST(int64_t, op->GetAttr("padding_idx"));
auto w_shape_ = op->Block()->FindVar(op->Input("W")[0])->GetShape();
auto w_shape_ = GetInputNode("W", node)->Var()->GetShape();
auto table_size_ = w_shape_[0];
auto emb_size_ = w_shape_[1];

Expand Down
2 changes: 2 additions & 0 deletions paddle/fluid/framework/ir/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,8 @@ if(WITH_IPU)
pass_library(inference_attr_extract_pass base DIR ipu DEPS ipu_pass_base)
pass_library(popart_canonicalization_pass base DIR ipu DEPS ipu_pass_base)
pass_library(ipu_inplace_pass base DIR ipu DEPS ipu_pass_base)
pass_library(infer_shape_pass base DIR ipu DEPS ipu_pass_base)
pass_library(avg_shard_pass base DIR ipu DEPS ipu_pass_base)
target_link_libraries(popart_canonicalization_pass -Wl,--whole-archive popart_canonicalization_utils -Wl,--no-whole-archive)
endif()

Expand Down
55 changes: 55 additions & 0 deletions paddle/fluid/framework/ir/ipu/avg_shard_pass.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
// 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/avg_shard_pass.h"

#include "paddle/fluid/framework/ipu/ipu_backend.h"

#include "paddle/fluid/framework/ir/graph_helper.h"
#include "paddle/fluid/framework/ir/pass_tester_helper.h"

namespace paddle {
namespace framework {
namespace ir {

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

std::shared_ptr<ipu::IpuBackend> ipu_backend = ipu::IpuBackend::GetInstance();

if (ipu_backend->GetIpuStrategy()->need_avg_shard) {
VLOG(10) << "start AvgShardPass";
auto nodes = ir::TopologySortOperations(*graph);
auto num_ipus = ipu_backend->GetIpuStrategy()->num_ipus;

int shard_position = nodes.size() / num_ipus;
int index_and_stage = -1;
for (int i = 0; i < nodes.size(); i++) {
if ((i % shard_position) == 0 && index_and_stage < num_ipus - 1) {
index_and_stage++;
}
nodes[i]->Op()->SetAttr("ipu_index", index_and_stage);
nodes[i]->Op()->SetAttr("ipu_stage", index_and_stage);
}
VLOG(10) << "end AvgShardPass";
}

VLOG(10) << "leave AvgShardPass::ApplyImpl";
}

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

REGISTER_PASS(avg_shard_pass, paddle::framework::ir::AvgShardPass);
30 changes: 30 additions & 0 deletions paddle/fluid/framework/ir/ipu/avg_shard_pass.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
// 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.

#pragma once

#include "paddle/fluid/framework/ir/ipu/ipu_pass_base.h"

namespace paddle {
namespace framework {
namespace ir {

class AvgShardPass : public IPUPassBase {
protected:
void ApplyImpl(ir::Graph* graph) const override;
};

} // namespace ir
} // namespace framework
} // namespace paddle
102 changes: 102 additions & 0 deletions paddle/fluid/framework/ir/ipu/infer_shape_pass.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
// 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/ddim.h"
#include "paddle/fluid/framework/ipu/ipu_backend.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"

namespace paddle {
namespace framework {
namespace ir {

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

// graph_viz_pass
// auto graph_viz_pass = PassRegistry::Instance().Get("graph_viz_pass");
// graph_viz_pass->Set("graph_viz_path",
// new std::string("before_pass.dot"));
// graph_viz_pass->Apply(graph);

std::shared_ptr<ipu::IpuBackend> ipu_backend = ipu::IpuBackend::GetInstance();

if (ipu_backend->GetIpuStrategy()->need_infer_shape) {
VLOG(10) << "start InferShapePass";
// temp scope for shape inference
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(paddle::framework::make_ddim(var_desc->GetShape()));

// skip lowerWeights in IpuBackend if persistable node has inputs
if (var_desc->Persistable() && node->inputs.size()) {
var_desc->SetPersistable(false);
}
}

// infer shape
auto nodes = ir::TopologySortOperations(*graph);
for (auto node : nodes) {
auto op_desc = node->Op();
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 = paddle::framework::vectorize(dim);
for (auto output_node : node->outputs) {
if (output_node->Name() == output_name) {
output_node->Var()->SetShape(new_shape);
}
}
}
}
}
// release the temp scope
scope.reset();
VLOG(10) << "end InferShapePass";
}

// graph_viz_pass
// graph_viz_pass->Erase("graph_viz_path");
// graph_viz_pass->Set("graph_viz_path",
// new std::string("after_pass.dot"));
// graph_viz_pass->Apply(graph);

VLOG(10) << "leave InferShapePass::ApplyImpl";
}

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

REGISTER_PASS(infer_shape_pass, paddle::framework::ir::InferShapePass);
Loading

0 comments on commit 11708d4

Please sign in to comment.