Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Hold for pipeline test] Support control-flow #353

Draft
wants to merge 19 commits into
base: master
Choose a base branch
from
Draft
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
draft: if and loop op_define
  • Loading branch information
xysmlx committed Jul 15, 2021

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
commit 6cff7e72ad90dfc861c23d6d831b0c1a864032b0
2 changes: 2 additions & 0 deletions src/nnfusion/core/operators/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -42,9 +42,11 @@ set(SRC
op_define/gelu.cpp
op_define/greater_eq.cpp
op_define/greater.cpp
op_define/if.cpp
op_define/less_eq.cpp
op_define/less.cpp
op_define/log.cpp
op_define/loop.cpp
op_define/lrn.cpp
op_define/max_pool.cpp
op_define/max.cpp
58 changes: 58 additions & 0 deletions src/nnfusion/core/operators/op_define/if.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
//*****************************************************************************
// Copyright 2017-2020 Intel Corporation
//
// 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.
//*****************************************************************************

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

#include "if.hpp"

using namespace std;
using namespace nnfusion::op;

If::If(std::shared_ptr<nnfusion::graph::Graph>& then_branch_graph,
std::shared_ptr<nnfusion::graph::Graph>& else_branch_graph)
: Op("If")
, m_then_branch_graph(then_branch_graph)
, m_else_branch_graph(else_branch_graph)
{
}

void If::validate_and_infer_types(std::shared_ptr<graph::GNode> gnode)
{
nnfusion::Shape cond_shape = gnode->get_input_shape(0);
nnfusion::element::Type cond_et = gnode->get_input_element_type(0);
NNFUSION_CHECK(cond_shape.size() == 0)
<< "The condition tensor of the If operation mush be scalar.";
NNFUSION_CHECK(cond_et == nnfusion::element::boolean)
<< "The condition tensor of the If operation mush be boolean.";

auto then_branch_outputs = m_then_branch_graph->get_outputs();
auto else_branch_outputs = m_else_branch_graph->get_outputs();
NNFUSION_CHECK(then_branch_outputs.size() == else_branch_outputs.size())
<< "The outputs in the then_branch and else_branch must have the same shape and "
"same data type.";
for (size_t i = 0; i < then_branch_outputs.size(); i++)
{
NNFUSION_CHECK(then_branch_outputs[i]->get_shape() == else_branch_outputs[i]->get_shape() &&
then_branch_outputs[i]->get_element_type() ==
else_branch_outputs[i]->get_element_type())
<< "The outputs in the then_branch and else_branch must have the same shape and "
"same data type.";

gnode->set_output_type_and_shape(
i, then_branch_outputs[i]->get_element_type(), then_branch_outputs[i]->get_shape());
}
}
49 changes: 49 additions & 0 deletions src/nnfusion/core/operators/op_define/if.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
//*****************************************************************************
// Copyright 2017-2020 Intel Corporation
//
// 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.
//*****************************************************************************

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

#pragma once

#include "../op.hpp"
#include "nnfusion/core/graph/graph.hpp"

namespace nnfusion
{
namespace op
{
/// \brief If control-flow operation, with same definition as https://github.com/onnx/onnx/blob/master/docs/Changelog.md#If-1.
class If : public Op
{
public:
/// \brief Constructs an if operation
///
/// \param then_branch_graph The then_branch graph.<br>
/// `[f]`
/// \param else_branch_graph The else_branch graph.<br>
/// `[f]`
If(std::shared_ptr<nnfusion::graph::Graph>& then_branch_graph,
std::shared_ptr<nnfusion::graph::Graph>& else_branch_graph);

void validate_and_infer_types(std::shared_ptr<graph::GNode> gnode) override;

protected:
std::shared_ptr<nnfusion::graph::Graph> m_then_branch_graph;
std::shared_ptr<nnfusion::graph::Graph> m_else_branch_graph;
};
}
}
55 changes: 55 additions & 0 deletions src/nnfusion/core/operators/op_define/loop.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
//*****************************************************************************
// Copyright 2017-2020 Intel Corporation
//
// 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.
//*****************************************************************************

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

#include "loop.hpp"

using namespace std;
using namespace nnfusion::op;

Loop::Loop(std::shared_ptr<nnfusion::graph::Graph>& loop_body_graph,
const std::vector<nnfusion::PartialShape>& output_shapes,
const std::vector<nnfusion::element::Type>& output_types)
: Op("Loop")
, m_loop_body_graph(loop_body_graph)
, m_output_shapes(output_shapes)
, m_output_types(output_types)
{
}

void Loop::validate_and_infer_types(std::shared_ptr<graph::GNode> gnode)
{
nnfusion::Shape trip_shape = gnode->get_input_shape(0);
nnfusion::element::Type trip_et = gnode->get_input_element_type(0);
NNFUSION_CHECK(trip_shape.size() == 0)
<< "The trip-count tensor of the Loop operation mush be scalar.";
NNFUSION_CHECK(trip_et == nnfusion::element::i64)
<< "The trip-count tensor of the Loop operation mush be boolean.";

nnfusion::Shape cond_shape = gnode->get_input_shape(1);
nnfusion::element::Type cond_et = gnode->get_input_element_type(1);
NNFUSION_CHECK(cond_shape.size() == 0)
<< "The condition tensor of the Loop operation mush be scalar.";
NNFUSION_CHECK(cond_et == nnfusion::element::boolean)
<< "The condition tensor of the Loop operation mush be boolean.";

for (size_t i = 0; i < m_output_types.size(); i++)
{
gnode->set_output_type_and_shape(i, m_output_types[i], m_output_shapes[i]);
}
}
49 changes: 49 additions & 0 deletions src/nnfusion/core/operators/op_define/loop.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
//*****************************************************************************
// Copyright 2017-2020 Intel Corporation
//
// 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.
//*****************************************************************************

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

#pragma once

#include "../op.hpp"
#include "nnfusion/core/graph/graph.hpp"

namespace nnfusion
{
namespace op
{
/// \brief Loop control-flow operation, with same definition as https://github.com/onnx/onnx/blob/master/docs/Changelog.md#Loop-11.
class Loop : public Op
{
public:
/// \brief Constructs an if operation
///
/// \param loop_body_graph The loop body graph.<br>
/// `[f]`
Loop(std::shared_ptr<nnfusion::graph::Graph>& loop_body_graph,
const std::vector<nnfusion::PartialShape>& output_shapes,
const std::vector<nnfusion::element::Type>& output_types);

void validate_and_infer_types(std::shared_ptr<graph::GNode> gnode) override;

protected:
std::shared_ptr<nnfusion::graph::Graph> m_loop_body_graph;
std::vector<nnfusion::PartialShape> m_output_shapes;
std::vector<nnfusion::element::Type> m_output_types;
};
} // namespace op
} // namespace nnfusion