commit
7c06d292c8
@ -0,0 +1,120 @@
|
||||
/**
|
||||
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* 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 "pre_activate/ascend/ir_fusion/parameter_and_transop_fusion.h"
|
||||
#include <memory>
|
||||
#include "session/anf_runtime_algorithm.h"
|
||||
#include "utils/utils.h"
|
||||
#include "operator/ops.h"
|
||||
#include "device/kernel_info.h"
|
||||
#include "pre_activate/common/helper.h"
|
||||
#include "pre_activate/common/optimizer.h"
|
||||
#include "pre_activate/ascend/ascend_helper.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace opt {
|
||||
const AnfNodePtr ParamTransRoad(const FuncGraphPtr &func_graph, const AnfNodePtr &node, bool first_flag,
|
||||
std::vector<CNodePtr> *trans_road) {
|
||||
if (node == nullptr) {
|
||||
MS_LOG(ERROR) << "nullptr";
|
||||
return nullptr;
|
||||
}
|
||||
if (node->isa<CNode>()) {
|
||||
auto cnode = node->cast<CNodePtr>();
|
||||
auto op_name = AnfAlgo::GetCNodeName(cnode);
|
||||
auto manager = func_graph->manager();
|
||||
if (manager == nullptr) {
|
||||
return nullptr;
|
||||
}
|
||||
if (op_name == prim::kPrimCast->name() || op_name == prim::kPrimTranspose->name() ||
|
||||
op_name == prim::kPrimReshape->name() || op_name == kTransDataOpName) {
|
||||
auto users = manager->node_users()[node];
|
||||
if (users.size() > 1 && !first_flag) {
|
||||
return nullptr;
|
||||
}
|
||||
trans_road->push_back(cnode);
|
||||
first_flag = false;
|
||||
auto next_node = AnfAlgo::GetInputNode(cnode, 0);
|
||||
if (next_node->isa<Parameter>() || next_node->isa<ValueNode>()) {
|
||||
return next_node;
|
||||
}
|
||||
return ParamTransRoad(func_graph, next_node, first_flag, trans_road);
|
||||
}
|
||||
} else if (node->isa<Parameter>() || node->isa<ValueNode>()) {
|
||||
return node;
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
bool ParameterTransOpFusion::Run(const FuncGraphPtr &func_graph) {
|
||||
if (func_graph == nullptr) {
|
||||
MS_LOG(ERROR) << "Func graph is nullptr";
|
||||
return false;
|
||||
}
|
||||
auto manager = func_graph->manager();
|
||||
if (manager == nullptr) {
|
||||
return false;
|
||||
}
|
||||
std::vector<AnfNodePtr> node_list = TopoSort(func_graph->get_return());
|
||||
bool changed = false;
|
||||
for (auto node : node_list) {
|
||||
if (node == nullptr || !node->isa<CNode>()) {
|
||||
continue;
|
||||
}
|
||||
auto cnode = node->cast<CNodePtr>();
|
||||
auto node_name = AnfAlgo::GetCNodeName(cnode);
|
||||
if (node_name == prim::kPrimCast->name() || node_name == prim::kPrimTranspose->name() ||
|
||||
node_name == prim::kPrimReshape->name() || node_name == kTransDataOpName) {
|
||||
MS_LOG(DEBUG) << "Skip trans op";
|
||||
continue;
|
||||
}
|
||||
for (size_t input_index = 0; input_index < AnfAlgo::GetInputTensorNum(cnode); input_index++) {
|
||||
std::vector<CNodePtr> trans_road;
|
||||
bool first_flag = true;
|
||||
auto final_node = ParamTransRoad(func_graph, AnfAlgo::GetInputNode(cnode, input_index), first_flag, &trans_road);
|
||||
if (final_node != nullptr && trans_road.size() == 3 && AnfAlgo::GetCNodeName(trans_road[0]) == kTransDataOpName &&
|
||||
AnfAlgo::GetCNodeName(trans_road[1]) == prim::kPrimCast->name() &&
|
||||
AnfAlgo::GetCNodeName(trans_road[2]) == kTransDataOpName) {
|
||||
auto cur_transop = trans_road[0];
|
||||
auto format = AnfAlgo::GetOutputFormat(cur_transop, 0);
|
||||
auto dtype = AnfAlgo::GetOutputDeviceDataType(cur_transop, 0);
|
||||
auto param_format = AnfAlgo::GetOutputFormat(final_node, 0);
|
||||
auto param_dtype = AnfAlgo::GetOutputDeviceDataType(final_node, 0);
|
||||
|
||||
auto cast = trans_road[1];
|
||||
auto cast_format = AnfAlgo::GetOutputFormat(cast, 0);
|
||||
auto cast_build_info = cast->kernel_info()->select_kernel_build_info();
|
||||
kernel::KernelBuildInfo::KernelBuildInfoBuilder builder;
|
||||
builder.SetOutputsFormat({format});
|
||||
builder.SetInputsFormat({format});
|
||||
builder.SetInputsDeviceType({param_dtype});
|
||||
builder.SetOutputsDeviceType({dtype});
|
||||
builder.SetKernelType(cast_build_info->kernel_type());
|
||||
builder.SetFusionType(cast_build_info->fusion_type());
|
||||
builder.SetProcessor(cast_build_info->processor());
|
||||
AnfAlgo::SetSelectKernelBuildInfo(builder.Build(), cast.get());
|
||||
if (param_format == format && param_dtype != dtype) {
|
||||
manager->Replace(trans_road[2], final_node);
|
||||
manager->Replace(cur_transop, cast);
|
||||
}
|
||||
changed = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
return changed;
|
||||
}
|
||||
} // namespace opt
|
||||
} // namespace mindspore
|
@ -0,0 +1,41 @@
|
||||
/**
|
||||
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
#ifndef MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_IR_FUSION_PARAMETER_AND_TRANSOP_FUSION_H_
|
||||
#define MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_IR_FUSION_PARAMETER_AND_TRANSOP_FUSION_H_
|
||||
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <memory>
|
||||
#include "ir/anf.h"
|
||||
#include "pre_activate/common/pass.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace opt {
|
||||
class ParameterTransOpFusion : public Pass {
|
||||
public:
|
||||
explicit ParameterTransOpFusion(size_t groups = 1) : Pass("Parameter_and_transop_fusion"), groups_(groups) {}
|
||||
~ParameterTransOpFusion() override = default;
|
||||
bool Run(const FuncGraphPtr &graph) override;
|
||||
|
||||
private:
|
||||
size_t groups_ = 1;
|
||||
};
|
||||
} // namespace opt
|
||||
} // namespace mindspore
|
||||
|
||||
#endif
|
Loading…
Reference in new issue