set min/max shape

Signed-off-by: zhupuxu <zhupuxu@huawei.com>
This commit is contained in:
zhupuxu 2022-03-09 15:36:57 +08:00
parent 06a76f69bc
commit d0d7c26b47
55 changed files with 536 additions and 225 deletions

View File

@ -112,7 +112,16 @@ CNodePtr InsertCastForGraphKernel(const FuncGraphPtr &func_graph, const CNodePtr
MS_EXCEPTION_IF_NULL(cast);
cast->set_scope(cnode->scope());
ShapeVector out_shape = GetShape(cur_input);
auto abs_shape_ptr = std::make_shared<abstract::Shape>(abstract::Shape(out_shape));
BaseShapePtr abs_shape_ptr = nullptr;
auto is_dynamic = std::any_of(out_shape.begin(), out_shape.end(), [](int64_t s) { return s < 0; });
if (is_dynamic) {
auto max_shape = common::AnfAlgo::GetOutputMaxShape(in_node, in_index);
auto min_shape = common::AnfAlgo::GetOutputMinShape(in_node, in_index);
abs_shape_ptr = std::make_shared<abstract::Shape>(out_shape, min_shape, max_shape);
} else {
abs_shape_ptr = std::make_shared<abstract::Shape>(out_shape);
}
auto abstract =
std::make_shared<abstract::AbstractTensor>(TypeIdToType(TypeId::kNumberTypeFloat16), abs_shape_ptr);
cast->set_abstract(abstract);

View File

@ -160,9 +160,9 @@ AnfNodePtr CreateTupleGetItem(const AnfNodePtr &buffer_fusion_kernel, session::K
tuple_getitem_inputs_list.push_back(idx);
auto tuple_item = kernel_graph->NewCNode(tuple_getitem_inputs_list);
MS_EXCEPTION_IF_NULL(tuple_item);
common::AnfAlgo::SetOutputInferTypeAndShape(
common::AnfAlgo::SetOutputTypeAndDetailShape(
{common::AnfAlgo::GetOutputInferDataType(buffer_fusion_kernel, output_index)},
{common::AnfAlgo::GetOutputInferShape(buffer_fusion_kernel, output_index)}, tuple_item.get());
{common::AnfAlgo::GetOutputDetailShape(buffer_fusion_kernel, output_index)}, tuple_item.get());
return tuple_item;
}
@ -517,19 +517,19 @@ bool UbPatternFusion::ReplaceFusionOp(mindspore::HashMap<int64_t, BufferFusionIn
AnfAlgo::SetSelectKernelBuildInfo(buffer_fusion_info.kernel_build_info, buffer_fusion.get());
// Set abstract of fusion_op node
std::vector<TypeId> types;
std::vector<std::vector<size_t>> shapes;
std::vector<BaseShapePtr> shapes;
for (const auto &out_node : buffer_fusion_info.outputs_list) {
size_t out_num = common::AnfAlgo::GetOutputTensorNum(out_node);
for (size_t idx = 0; idx < out_num; ++idx) {
(void)types.emplace_back(common::AnfAlgo::GetOutputInferDataType(out_node, idx));
(void)shapes.emplace_back(common::AnfAlgo::GetOutputInferShape(out_node, idx));
(void)shapes.emplace_back(common::AnfAlgo::GetOutputDetailShape(out_node, idx));
}
}
if (types.empty() || shapes.empty()) {
MS_LOG(WARNING) << "The outputs_list of buffer_fusion_info is empty.";
return false;
}
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, buffer_fusion.get());
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, buffer_fusion.get());
common::AnfAlgo::SetNodeAttr(kAttrIsUBFusionOp, MakeValue(true), buffer_fusion);
SetFusionOpRefInfos(kernel_graph, buffer_fusion_info.outputs_list, buffer_fusion);
ReplaceOldNode(buffer_fusion_infos, fusion_id, buffer_fusion, kernel_graph);

View File

@ -16,6 +16,7 @@
#include "plugin/device/ascend/optimizer/enhancer/concat_outputs_for_all_gather.h"
#include <utility>
#include <algorithm>
#include "backend/common/session/anf_runtime_algorithm.h"
#include "include/common/utils/anfalgo.h"
@ -25,6 +26,8 @@ OutputInfo GetNodeOutputInfo(const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(node);
std::vector<TypeId> output_infer_dtype;
std::vector<std::vector<size_t>> output_infer_shape;
std::vector<std::vector<int64_t>> output_max_shape;
std::vector<std::vector<int64_t>> output_min_shape;
std::vector<std::string> output_format;
std::vector<TypeId> output_device_dtype;
auto type_ptr = node->Type();
@ -37,11 +40,14 @@ OutputInfo GetNodeOutputInfo(const AnfNodePtr &node) {
for (size_t i = 0; i < output_num; i++) {
(void)output_infer_dtype.emplace_back(common::AnfAlgo::GetOutputInferDataType(type_ptr, i));
(void)output_infer_shape.emplace_back(common::AnfAlgo::GetOutputInferShape(node, shape_ptr, i));
(void)output_min_shape.emplace_back(common::AnfAlgo::GetOutputMinShape(node, i));
(void)output_max_shape.emplace_back(common::AnfAlgo::GetOutputMaxShape(node, i));
(void)output_format.emplace_back(build_info->GetOutputFormat(i));
(void)output_device_dtype.emplace_back(build_info->GetOutputDeviceType(i));
}
return {output_infer_dtype, output_infer_shape, output_format, output_device_dtype};
return {output_infer_dtype, output_infer_shape, output_min_shape,
output_max_shape, output_format, output_device_dtype};
}
kernel::KernelBuildInfoPtr GenerateKernelBuildInfo(const AnfNodePtr &concat, const OutputInfo &allgather_output_info,
@ -55,8 +61,8 @@ kernel::KernelBuildInfoPtr GenerateKernelBuildInfo(const AnfNodePtr &concat, con
size_t concat_input_num = common::AnfAlgo::GetInputTensorNum(concat);
for (size_t i = 0; i < concat_input_num; ++i) {
size_t input_index = allgather_input_idx + i * allgather_input_num;
(void)inputs_device_format.emplace_back(std::get<kIndex2>(allgather_output_info)[input_index]);
(void)inputs_device_type.emplace_back(std::get<kIndex3>(allgather_output_info)[input_index]);
(void)inputs_device_format.emplace_back(std::get<kIndex4>(allgather_output_info)[input_index]);
(void)inputs_device_type.emplace_back(std::get<kIndex5>(allgather_output_info)[input_index]);
}
// Current only support default format & float16
auto cmp_format = inputs_device_format.begin();
@ -102,10 +108,20 @@ AnfNodePtr ConcatOutputsForAllGather::InsertConcatForOutput(const FuncGraphPtr &
MS_EXCEPTION_IF_NULL(concat);
MS_EXCEPTION_IF_NULL(new_tuple_getitems[i]);
const std::vector<TypeId> &dtypes = {std::get<0>(output_info)[i]};
const auto &shape = std::get<1>(output_info)[i];
std::vector<std::vector<size_t>> shapes = {shape};
shapes[0][0] *= LongToSize(rank_size);
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, concat.get());
auto shape = std::get<1>(output_info)[i];
shape[0] *= LongToSize(rank_size);
if (AnfUtils::IsShapeDynamic(shape)) {
ShapeVector tensor_shape;
auto min_shape = std::get<kIndex2>(output_info)[i];
auto max_shape = std::get<kIndex3>(output_info)[i];
max_shape[0] *= rank_size;
min_shape[0] *= rank_size;
std::transform(shape.begin(), shape.end(), std::back_inserter(tensor_shape), SizeToLong);
BaseShapePtr base_shape = std::make_shared<abstract::Shape>(tensor_shape, min_shape, max_shape);
common::AnfAlgo::SetOutputTypeAndDetailShape(dtypes, {base_shape}, concat.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, {shape}, concat.get());
}
common::AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue(static_cast<int64_t>(0)), concat);
common::AnfAlgo::SetNodeAttr(kAttrInputNums, MakeValue(rank_size), concat);
std::vector<int64_t> dyn_input_size{rank_size};
@ -159,8 +175,8 @@ const AnfNodePtr ConcatOutputsForAllGather::Process(const FuncGraphPtr &func_gra
idx->set_abstract(abstract_scalar);
auto tuple_getitem = func_graph->NewCNode({NewValueNode(prim::kPrimTupleGetItem), node, idx});
MS_EXCEPTION_IF_NULL(tuple_getitem);
common::AnfAlgo::SetOutputInferTypeAndShape({std::get<0>(output_info)[i]}, {std::get<1>(output_info)[i]},
tuple_getitem.get());
common::AnfAlgo::SetOutputTypeAndDetailShape({std::get<0>(output_info)[i]},
{common::AnfAlgo::GetOutputDetailShape(node, i)}, tuple_getitem.get());
(void)new_outputs.emplace_back(std::move(tuple_getitem));
}
return InsertConcatForOutput(func_graph, node, output_info, new_outputs, rank_size);

View File

@ -25,8 +25,8 @@
namespace mindspore {
namespace opt {
using OutputInfo =
std::tuple<std::vector<TypeId>, std::vector<std::vector<size_t>>, std::vector<std::string>, std::vector<TypeId>>;
using OutputInfo = std::tuple<std::vector<TypeId>, std::vector<std::vector<size_t>>, std::vector<std::vector<int64_t>>,
std::vector<std::vector<int64_t>>, std::vector<std::string>, std::vector<TypeId>>;
class ConcatOutputsForAllGather : public PatternProcessPass {
public:

View File

@ -27,21 +27,36 @@ std::vector<AnfNodePtr> SplitInputsForReduceScatter::InsertSplitForInput(const F
size_t inputs_size = common::AnfAlgo::GetInputTensorNum(node);
std::vector<AnfNodePtr> split_outputs;
size_t rank_size_t = LongToSize(rank_size);
if (rank_size_t == 0) {
MS_LOG(EXCEPTION) << "The rank size can not be zero.";
}
for (size_t i = 0; i < inputs_size; i++) {
std::vector<AnfNodePtr> split_inputs{NewValueNode(std::make_shared<Primitive>(prim::kPrimSplitV->name()))};
split_inputs.push_back(common::AnfAlgo::GetInputNode(node, i));
auto split = NewCNode(split_inputs, func_graph);
MS_EXCEPTION_IF_NULL(split);
std::vector<TypeId> dtypes(rank_size, common::AnfAlgo::GetPrevNodeOutputInferDataType(node, i));
std::vector<std::vector<size_t>> shapes;
std::vector<int> size_splits;
std::vector<size_t> output_node_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(node, i);
output_node_shape[0] /= rank_size_t;
if (AnfUtils::IsShapeDynamic(output_node_shape)) {
auto min_shape = common::AnfAlgo::GetInputMinShape(node, i);
auto max_shape = common::AnfAlgo::GetInputMaxShape(node, i);
min_shape[0] /= rank_size_t;
max_shape[0] /= rank_size_t;
ShapeVector shape_tmp;
std::transform(output_node_shape.begin(), output_node_shape.end(), std::back_inserter(shape_tmp), SizeToLong);
std::vector<BaseShapePtr> shapes(rank_size_t, std::make_shared<abstract::Shape>(shape_tmp, min_shape, max_shape));
common::AnfAlgo::SetOutputTypeAndDetailShape(dtypes, shapes, split.get());
} else {
std::vector<std::vector<size_t>> shapes(rank_size_t, output_node_shape);
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, split.get());
}
for (size_t j = 0; j < rank_size_t; j++) {
std::vector<size_t> output_node_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(node, i);
output_node_shape[0] /= rank_size_t;
shapes.push_back(output_node_shape);
size_splits.push_back(output_node_shape[0]);
}
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, split.get());
common::AnfAlgo::SetNodeAttr("split_dim", MakeValue(0L), split);
common::AnfAlgo::SetNodeAttr("num_split", MakeValue(rank_size), split);
common::AnfAlgo::SetNodeAttr("size_splits", MakeValue(size_splits), split);

View File

@ -47,10 +47,22 @@ CNodePtr Insert(const FuncGraphPtr &func_graph, const CNodePtr &cnode, const std
auto origin_type = common::AnfAlgo::GetPrevNodeOutputInferDataType(cnode, 1);
auto origin_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(cnode, 1);
auto dst_shape = {origin_shape[1], origin_shape[0]};
auto is_dynamic = AnfUtils::IsShapeDynamic(dst_shape);
transpose_inputs.push_back(common::AnfAlgo::GetInputNode(cnode, 1));
CNodePtr transpose = func_graph->NewCNode(transpose_inputs);
MS_EXCEPTION_IF_NULL(transpose);
common::AnfAlgo::SetOutputInferTypeAndShape({origin_type}, {dst_shape}, transpose.get());
if (is_dynamic) {
auto shape = {SizeToLong(origin_shape[1]), SizeToLong(origin_shape[0])};
auto max_shape = common::AnfAlgo::GetInputMaxShape(cnode, 1);
auto min_shape = common::AnfAlgo::GetInputMinShape(cnode, 1);
auto shape_tmp1 = {min_shape[1], min_shape[0]};
auto shape_tmp2 = {max_shape[1], max_shape[0]};
BaseShapePtr base_shape = std::make_shared<abstract::Shape>(shape, shape_tmp1, shape_tmp2);
common::AnfAlgo::SetOutputTypeAndDetailShape({origin_type}, {base_shape}, transpose.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape({origin_type}, {dst_shape}, transpose.get());
}
common::AnfAlgo::SetNodeAttr(kAttrPerm, MakeValue(std::vector<int64_t>{1, 0}), transpose);
common::AnfAlgo::SetNodeInput(cnode, transpose, 1);
if (kernel_graph == nullptr) {
@ -66,11 +78,21 @@ CNodePtr Insert(const FuncGraphPtr &func_graph, const CNodePtr &cnode, const std
auto origin_shape = common::AnfAlgo::GetOutputInferShape(cnode, output_idx);
if (origin_shape.size() > 1 && output_idx == 0) {
auto dtype = common::AnfAlgo::GetOutputInferDataType(cnode, output_idx);
auto dst_shape = {origin_shape[0], origin_shape[1]};
transpose_inputs.push_back(tuple_getitem);
CNodePtr transpose = func_graph->NewCNode(transpose_inputs);
MS_EXCEPTION_IF_NULL(transpose);
common::AnfAlgo::SetOutputInferTypeAndShape({dtype}, {dst_shape}, transpose.get());
if (AnfUtils::IsShapeDynamic(origin_shape)) {
auto dst_shape = {SizeToLong(origin_shape[0]), SizeToLong(origin_shape[1])};
auto min_shape = common::AnfAlgo::GetOutputMinShape(cnode, output_idx);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(cnode, output_idx);
auto shape_tmp1 = {min_shape[0], min_shape[1]};
auto shape_tmp2 = {max_shape[0], max_shape[1]};
BaseShapePtr base_shape = std::make_shared<abstract::Shape>(dst_shape, shape_tmp1, shape_tmp2);
common::AnfAlgo::SetOutputTypeAndDetailShape({dtype}, {base_shape}, transpose.get());
} else {
auto dst_shape = {origin_shape[0], origin_shape[1]};
common::AnfAlgo::SetOutputInferTypeAndShape({dtype}, {dst_shape}, transpose.get());
}
common::AnfAlgo::SetNodeAttr(kAttrPerm, MakeValue(std::vector<int64_t>{1, 0}), transpose);
make_tuple_inputs.push_back(transpose);
} else {

View File

@ -44,9 +44,9 @@ void BatchNormGradSplit::CreateOutputsOfUpdateGrad(const FuncGraphPtr &graph, co
auto types = {common::AnfAlgo::GetOutputInferDataType(bn_grad_node, 1),
common::AnfAlgo::GetOutputInferDataType(bn_grad_node, 2)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(bn_grad_node, 1),
common::AnfAlgo::GetOutputInferShape(bn_grad_node, 2)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, bn_update_grad.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(bn_grad_node, 1),
common::AnfAlgo::GetOutputDetailShape(bn_grad_node, 2)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, bn_update_grad.get());
common::AnfAlgo::CopyNodeAttr(kAttrEpsilon, bn_grad_node, bn_update_grad);
CreateMultipleOutputsOfAnfNode(graph, bn_update_grad, kBNTrainingUpdateGradOutputNum, bn_update_grad_outputs);
@ -79,8 +79,8 @@ void BatchNormGradSplit::CreateOutputsOfReduceGrad(const FuncGraphPtr &graph, co
bn_reduce_grad->set_scope(bn_grad_node->scope());
auto types = {common::AnfAlgo::GetOutputInferDataType(bn_grad_node, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(bn_grad_node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, bn_reduce_grad.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(bn_grad_node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, bn_reduce_grad.get());
common::AnfAlgo::CopyNodeAttr(kAttrEpsilon, bn_grad_node, bn_reduce_grad);
(*bn_reduce_grad_outputs).push_back(bn_reduce_grad);

View File

@ -40,8 +40,8 @@ AnfNodePtr BCEWithLogitsLossFission::AddReduceNode(const FuncGraphPtr &func_grap
MS_EXCEPTION_IF_NULL(new_cnode);
auto predict_input = cnode->inputs()[kIndex1];
auto new_node_dtype = {common::AnfAlgo::GetOutputInferDataType(predict_input, 0)};
auto new_node_shape = {common::AnfAlgo::GetOutputInferShape(predict_input, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(new_node_dtype, new_node_shape, new_cnode.get());
auto new_node_shape = {common::AnfAlgo::GetOutputDetailShape(predict_input, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(new_node_dtype, new_node_shape, new_cnode.get());
// Add reduce node
string reduction = common::AnfAlgo::GetNodeAttr<std::string>(node, kAttrReduction);
@ -61,8 +61,8 @@ AnfNodePtr BCEWithLogitsLossFission::AddReduceNode(const FuncGraphPtr &func_grap
if (type == kNumberTypeFloat16) {
type = kNumberTypeFloat32;
}
auto shape = {common::AnfAlgo::GetOutputInferShape(node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape({type}, shape, reduce_node.get());
auto shape = {common::AnfAlgo::GetOutputDetailShape(node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape({type}, shape, reduce_node.get());
common::AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue(std::vector<int64_t>{}), reduce_node);
common::AnfAlgo::SetNodeAttr("keep_dims", MakeValue(false), reduce_node);
common::AnfAlgo::SetNodeAttr("is_backend_insert", MakeValue(true), reduce_node);

View File

@ -204,8 +204,8 @@ AnfNodePtr InsertCast(const FuncGraphPtr &graph, const AnfNodePtr &input, const
MS_EXCEPTION_IF_NULL(input);
if (common::AnfAlgo::GetOutputInferDataType(input, 0) != dst_type) {
AnfNodePtr cast = graph->NewCNode({NewValueNode(std::make_shared<Primitive>(kCastOpName)), input});
common::AnfAlgo::SetOutputInferTypeAndShape({dst_type}, {common::AnfAlgo::GetOutputInferShape(input, 0)},
cast.get());
common::AnfAlgo::SetOutputTypeAndDetailShape({dst_type}, {common::AnfAlgo::GetOutputDetailShape(input, 0)},
cast.get());
common::AnfAlgo::SetNodeAttr(kIsBackendCast, MakeValue(true), cast);
cast->set_scope(input->scope());
return cast;

View File

@ -27,19 +27,19 @@ namespace {
CNodePtr AddCastNode(const FuncGraphPtr &func_graph, const TypeId dst_type, const CNodePtr &input_node,
const bool fir_flag) {
std::vector<AnfNodePtr> new_cast_inputs = {NewValueNode(std::make_shared<Primitive>(prim::kPrimCast->name()))};
std::vector<size_t> shape;
BaseShapePtr shape;
if (fir_flag) {
new_cast_inputs.emplace_back(input_node->inputs()[kIndex1]);
shape = common::AnfAlgo::GetOutputInferShape(input_node->inputs()[kIndex1], 0);
shape = common::AnfAlgo::GetOutputDetailShape(input_node->inputs()[kIndex1], 0);
} else {
new_cast_inputs.emplace_back(input_node);
shape = common::AnfAlgo::GetOutputInferShape(input_node, 0);
shape = common::AnfAlgo::GetOutputDetailShape(input_node, 0);
}
CNodePtr new_cast = NewCNode(new_cast_inputs, func_graph);
new_cast->set_scope(input_node->scope());
new_cast->set_abstract(input_node->abstract());
common::AnfAlgo::SetNodeAttr(kAttrDstType, MakeValue(static_cast<size_t>(dst_type)), new_cast);
common::AnfAlgo::SetOutputInferTypeAndShape({dst_type}, {shape}, new_cast.get());
common::AnfAlgo::SetOutputTypeAndDetailShape({dst_type}, {shape}, new_cast.get());
return new_cast;
}
} // namespace

View File

@ -129,8 +129,21 @@ CNodePtr GatherV2DsFission::CreateGatherV2Ds(const FuncGraphPtr &graph, const CN
auto shape = common::AnfAlgo::GetOutputInferShape(origin_node, 0);
shape[shape.size() - 1] = pad_dim_size;
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(origin_node, 0)}, {shape},
gather_v2.get());
if (AnfUtils::IsShapeDynamic(shape)) {
ShapeVector shape_tmp;
auto min_shape = common::AnfAlgo::GetOutputMinShape(origin_node, 0);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(origin_node, 0);
min_shape[min_shape.size() - 1] = pad_dim_size;
max_shape[max_shape.size() - 1] = pad_dim_size;
std::transform(shape.begin(), shape.end(), std::back_inserter(shape_tmp), SizeToLong);
std::vector<BaseShapePtr> shapes = {std::make_shared<abstract::Shape>(shape_tmp, min_shape, max_shape)};
common::AnfAlgo::SetOutputTypeAndDetailShape({common::AnfAlgo::GetOutputInferDataType(origin_node, 0)}, shapes,
gather_v2.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(origin_node, 0)}, {shape},
gather_v2.get());
}
common::AnfAlgo::SetNodeAttr(kAttrInputIsDynamicShape, MakeValue(true), gather_v2);
auto input_names = common::AnfAlgo::GetNodeAttr<std::vector<std::string>>(origin_node, kAttrInputNames);
common::AnfAlgo::SetNodeAttr(kAttrInputNames, MakeValue(input_names), gather_v2);

View File

@ -52,7 +52,7 @@ AnfNodePtr PackFission::CreateNewPack(const FuncGraphPtr &func_graph, const CNod
std::vector<size_t> new_shape = output_shape;
auto axis_l = LongToSize(axis);
if (axis_l < new_shape.size()) {
new_shape[axis_l] = offset;
new_shape[axis_l] = static_cast<int64_t>(offset);
}
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(origin_pack_cnode, 0)},
{new_shape}, new_pack.get());

View File

@ -27,19 +27,19 @@ namespace {
CNodePtr AddCastNode(const FuncGraphPtr &func_graph, const TypeId dst_type, const CNodePtr &input_node,
const bool fir_flag) {
std::vector<AnfNodePtr> new_cast_inputs = {NewValueNode(std::make_shared<Primitive>(prim::kPrimCast->name()))};
std::vector<size_t> shape;
BaseShapePtr shape;
if (fir_flag) {
new_cast_inputs.emplace_back(input_node->inputs()[kIndex1]);
shape = common::AnfAlgo::GetOutputInferShape(input_node->inputs()[kIndex1], 0);
shape = common::AnfAlgo::GetOutputDetailShape(input_node->inputs()[kIndex1], 0);
} else {
new_cast_inputs.emplace_back(input_node);
shape = common::AnfAlgo::GetOutputInferShape(input_node, 0);
shape = common::AnfAlgo::GetOutputDetailShape(input_node, 0);
}
CNodePtr new_cast = NewCNode(new_cast_inputs, func_graph);
new_cast->set_scope(input_node->scope());
new_cast->set_abstract(input_node->abstract());
common::AnfAlgo::SetNodeAttr(kAttrDstType, MakeValue(static_cast<size_t>(dst_type)), new_cast);
common::AnfAlgo::SetOutputInferTypeAndShape({dst_type}, {shape}, new_cast.get());
common::AnfAlgo::SetOutputTypeAndDetailShape({dst_type}, {shape}, new_cast.get());
return new_cast;
}
} // namespace
@ -51,7 +51,7 @@ const BaseRef ReduceSumFission::DefinePattern() const {
}
CNodePtr AddReduceSumNode(const FuncGraphPtr &func_graph, const CNodePtr &input_node, const bool &keep_dims,
const std::vector<int64_t> &axis, const std::vector<size_t> &out_shape) {
const std::vector<int64_t> &axis, const BaseShapePtr &out_shape) {
MS_EXCEPTION_IF_NULL(func_graph);
MS_EXCEPTION_IF_NULL(input_node);
auto input_type = common::AnfAlgo::GetOutputInferDataType(input_node, 0);
@ -62,7 +62,7 @@ CNodePtr AddReduceSumNode(const FuncGraphPtr &func_graph, const CNodePtr &input_
reduce_sum->set_scope(input_node->scope());
common::AnfAlgo::SetNodeAttr(kAttrKeepDims, MakeValue(keep_dims), reduce_sum);
common::AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue(axis), reduce_sum);
common::AnfAlgo::SetOutputInferTypeAndShape({input_type}, {out_shape}, reduce_sum.get());
common::AnfAlgo::SetOutputTypeAndDetailShape({input_type}, {out_shape}, reduce_sum.get());
return reduce_sum;
}
@ -74,7 +74,7 @@ const AnfNodePtr ReduceSumFission::Process(const FuncGraphPtr &graph, const AnfN
auto cnode = node->cast<CNodePtr>();
auto prim = common::AnfAlgo::GetCNodePrimitive(cnode);
auto keep_dims = common::AnfAlgo::GetNodeAttr<bool>(cnode, kAttrKeepDims);
auto out_shape = common::AnfAlgo::GetOutputInferShape(cnode, 0);
auto out_shape = common::AnfAlgo::GetOutputDetailShape(cnode, 0);
std::vector<int64_t> inp_axis;
auto axis_value = prim->GetAttr(kAttrAxis);
MS_EXCEPTION_IF_NULL(axis_value);

View File

@ -56,8 +56,20 @@ CNodePtr UnsortSegmentSumFission::CreatePadding(const FuncGraphPtr &graph, const
padding->set_scope(origin_node->scope());
auto shape = common::AnfAlgo::GetPrevNodeOutputInferShape(origin_node, 0);
shape[shape.size() - 1] = pad_dim_size;
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetPrevNodeOutputInferDataType(origin_node, 0)},
{shape}, padding.get());
if (AnfUtils::IsShapeDynamic(shape)) {
auto min_shape = common::AnfAlgo::GetInputMinShape(origin_node, 0);
auto max_shape = common::AnfAlgo::GetInputMaxShape(origin_node, 0);
min_shape[shape.size() - 1] = pad_dim_size;
max_shape[shape.size() - 1] = pad_dim_size;
ShapeVector shape_tmp;
std::transform(shape.begin(), shape.end(), std::back_inserter(shape_tmp), SizeToLong);
BaseShapePtr base_shape = std::make_shared<abstract::Shape>(shape_tmp, min_shape, max_shape);
common::AnfAlgo::SetOutputTypeAndDetailShape({common::AnfAlgo::GetPrevNodeOutputInferDataType(origin_node, 0)},
{base_shape}, padding.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetPrevNodeOutputInferDataType(origin_node, 0)},
{shape}, padding.get());
}
common::AnfAlgo::SetNodeAttr(kAttrPadDimSize, MakeValue(SizeToLong(pad_dim_size)), padding);
return padding;
}
@ -75,8 +87,21 @@ CNodePtr UnsortSegmentSumFission::CreateUnsortedSegmentSum(const FuncGraphPtr &g
unsorted_segment_sum->set_scope(origin_node->scope());
auto shape = common::AnfAlgo::GetOutputInferShape(origin_node, 0);
shape[shape.size() - 1] = pad_dim_size;
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(origin_node, 0)}, {shape},
unsorted_segment_sum.get());
if (AnfUtils::IsShapeDynamic(shape)) {
auto min_shape = common::AnfAlgo::GetOutputMinShape(origin_node, 0);
auto max_shape = common::AnfAlgo::GetInputMaxShape(origin_node, 0);
min_shape[shape.size() - 1] = pad_dim_size;
max_shape[shape.size() - 1] = pad_dim_size;
ShapeVector shape_tmp;
std::transform(shape.begin(), shape.end(), std::back_inserter(shape_tmp), SizeToLong);
BaseShapePtr base_shape = std::make_shared<abstract::Shape>(shape_tmp, min_shape, max_shape);
common::AnfAlgo::SetOutputTypeAndDetailShape({common::AnfAlgo::GetOutputInferDataType(origin_node, 0)},
{base_shape}, unsorted_segment_sum.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(origin_node, 0)}, {shape},
unsorted_segment_sum.get());
}
common::AnfAlgo::SetNodeAttr(kAttrNumSegments, MakeValue(SizeToLong(shape[0])), unsorted_segment_sum);
return unsorted_segment_sum;
}

View File

@ -309,9 +309,9 @@ const AnfNodePtr AdamApplyOneWithDecayRule::Process(const FuncGraphPtr &graph, c
MS_EXCEPTION_IF_NULL(add1);
auto types = {common::AnfAlgo::GetOutputInferDataType(add1, 0), common::AnfAlgo::GetOutputInferDataType(add0, 0),
common::AnfAlgo::GetOutputInferDataType(sub0, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(add1, 0), common::AnfAlgo::GetOutputInferShape(add0, 0),
common::AnfAlgo::GetOutputInferShape(sub0, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, fusion_node.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(add1, 0), common::AnfAlgo::GetOutputDetailShape(add0, 0),
common::AnfAlgo::GetOutputDetailShape(sub0, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, fusion_node.get());
std::vector<AnfNodePtr> fusion_node_outputs;
CreateMultipleOutputsOfAnfNode(graph, fusion_node, kAdamApplyOneWithDecayOutputNum, &fusion_node_outputs);

View File

@ -93,9 +93,9 @@ const AnfNodePtr BNReduceGradConv2dBackpropFilterFusion::Process(const FuncGraph
MS_EXCEPTION_IF_NULL(fused_dbn_dw);
auto types = {common::AnfAlgo::GetOutputInferDataType(bnreduce_grad, 0),
common::AnfAlgo::GetOutputInferDataType(conv_back_filter, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(bnreduce_grad, 0),
common::AnfAlgo::GetOutputInferShape(conv_back_filter, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, fused_dbn_dw.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(bnreduce_grad, 0),
common::AnfAlgo::GetOutputDetailShape(conv_back_filter, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, fused_dbn_dw.get());
fused_dbn_dw->set_scope(bnreduce_grad->scope());
common::AnfAlgo::CopyNodeAttr(kAttrFilterSizes, conv_back_filter, fused_dbn_dw);
common::AnfAlgo::CopyNodeAttr(kAttrStride, conv_back_filter, fused_dbn_dw);

View File

@ -65,8 +65,8 @@ const AnfNodePtr ClipByNormNoDivSquareSumFusion::Process(const FuncGraphPtr &gra
auto fusion_node = NewCNode(inputs, graph);
MS_EXCEPTION_IF_NULL(fusion_node);
auto types = {common::AnfAlgo::GetOutputInferDataType(node, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, fusion_node.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, fusion_node.get());
fusion_node->set_scope(node->scope());
return fusion_node;
}

View File

@ -91,8 +91,8 @@ const AnfNodePtr ClipByValueFusion::Process(const FuncGraphPtr &graph, const Anf
auto clip_by_value = NewCNode(inputs, graph);
MS_EXCEPTION_IF_NULL(clip_by_value);
auto types = {common::AnfAlgo::GetOutputInferDataType(node, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, clip_by_value.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, clip_by_value.get());
clip_by_value->set_scope(node->scope());
return clip_by_value;
}

View File

@ -113,8 +113,8 @@ CNodePtr ConfusionMulGradFusion::CreateFusionNode(const FuncGraphPtr &graph, con
common::AnfAlgo::CopyNodeAttr(kAttrKeepDims, reduce_sum, fusion_node);
auto types = {common::AnfAlgo::GetOutputInferDataType(mul0, 0),
common::AnfAlgo::GetOutputInferDataType(reduce_sum, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(mul0, 0), common::AnfAlgo::GetOutputInferShape(reduce_sum, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, fusion_node.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(mul0, 0), common::AnfAlgo::GetOutputDetailShape(reduce_sum, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, fusion_node.get());
return fusion_node;
}

View File

@ -183,9 +183,9 @@ const AnfNodePtr LambNextMVWithDecayV1Rule::Process(const FuncGraphPtr &func_gra
std::tie(add0, add1) = GetAdd0Add1Nodes(real_div0, real_div1);
auto types = {common::AnfAlgo::GetOutputInferDataType(node, 0), common::AnfAlgo::GetOutputInferDataType(add0, 0),
common::AnfAlgo::GetOutputInferDataType(add1, 0), common::AnfAlgo::GetOutputInferDataType(add5, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(node, 0), common::AnfAlgo::GetOutputInferShape(add0, 0),
common::AnfAlgo::GetOutputInferShape(add1, 0), common::AnfAlgo::GetOutputInferShape(add5, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, fusion_node.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(node, 0), common::AnfAlgo::GetOutputDetailShape(add0, 0),
common::AnfAlgo::GetOutputDetailShape(add1, 0), common::AnfAlgo::GetOutputDetailShape(add5, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, fusion_node.get());
std::vector<AnfNodePtr> fusion_node_outputs;
CreateMultipleOutputsOfAnfNode(func_graph, fusion_node, kLambNextMVWithDecayV1OutputNum, &fusion_node_outputs);

View File

@ -70,8 +70,8 @@ const AnfNodePtr LambUpdateWithLRRuleFusion::Process(const FuncGraphPtr &graph,
MS_EXCEPTION_IF_NULL(lamb_update_with_lr);
auto types = {common::AnfAlgo::GetOutputInferDataType(node, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, lamb_update_with_lr.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, lamb_update_with_lr.get());
lamb_update_with_lr->set_scope(node->scope());
return lamb_update_with_lr;
}

View File

@ -55,8 +55,8 @@ const AnfNodePtr SoftmaxDropoutDoMaskV3Fusion::Process(const FuncGraphPtr &graph
MS_EXCEPTION_IF_NULL(softmax_dropout);
auto types = {common::AnfAlgo::GetOutputInferDataType(softmax, 0),
common::AnfAlgo::GetOutputInferDataType(dropout, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(softmax, 0), common::AnfAlgo::GetOutputInferShape(dropout, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, softmax_dropout.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(softmax, 0), common::AnfAlgo::GetOutputDetailShape(dropout, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, softmax_dropout.get());
softmax_dropout->set_scope(softmax->scope());
common::AnfAlgo::CopyNodeAttr(kAttrAxis, softmax, softmax_dropout);
common::AnfAlgo::CopyNodeAttr(kAttrKeepProb, dropout, softmax_dropout);

View File

@ -60,8 +60,8 @@ CNodePtr SquareSumFusion::GenerateSquareSumV1(const FuncGraphPtr &graph, const C
MS_EXCEPTION_IF_NULL(kernel_info);
square_sumv1->set_kernel_info(kernel_info);
auto types = {common::AnfAlgo::GetOutputInferDataType(sum, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(sum, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, square_sumv1.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(sum, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, square_sumv1.get());
square_sumv1->set_scope(sum->scope());
common::AnfAlgo::CopyNodeAttr(kAttrAxis, sum, square_sumv1);
common::AnfAlgo::CopyNodeAttr(kAttrKeepDims, sum, square_sumv1);
@ -82,8 +82,8 @@ CNodePtr SquareSumFusion::GenerateSquareSumV2(const FuncGraphPtr &graph, const C
auto square_sumv2 = NewCNode(square_sumv2_inputs, graph);
MS_EXCEPTION_IF_NULL(square_sumv2);
auto types = {common::AnfAlgo::GetOutputInferDataType(sum, 0), common::AnfAlgo::GetOutputInferDataType(square, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(sum, 0), common::AnfAlgo::GetOutputInferShape(square, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, square_sumv2.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(sum, 0), common::AnfAlgo::GetOutputDetailShape(square, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, square_sumv2.get());
square_sumv2->set_scope(sum->scope());
common::AnfAlgo::CopyNodeAttr(kAttrAxis, sum, square_sumv2);
common::AnfAlgo::CopyNodeAttr(kAttrKeepDims, sum, square_sumv2);

View File

@ -86,8 +86,21 @@ CNodePtr AllToAllUnifyMindIR::CreateSplitNode(const FuncGraphPtr &graph, const C
}
shape[LongToSize(split_dim)] /= static_cast<size_t>(split_count);
std::vector<TypeId> dtypes(split_count, dtype);
std::vector<std::vector<size_t>> shapes(split_count, shape);
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, split_v.get());
if (AnfUtils::IsShapeDynamic(shape)) {
auto min_shape = common::AnfAlgo::GetOutputMinShape(all_to_all_input, 0);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(all_to_all_input, 0);
max_shape[LongToSize(split_dim)] /= split_count;
min_shape[LongToSize(split_dim)] /= split_count;
ShapeVector new_shape;
std::transform(shape.begin(), shape.end(), std::back_inserter(new_shape), SizeToLong);
std::vector<BaseShapePtr> shapes(split_count, std::make_shared<abstract::Shape>(new_shape, min_shape, max_shape));
common::AnfAlgo::SetOutputTypeAndDetailShape(dtypes, shapes, split_v.get());
} else {
std::vector<std::vector<size_t>> shapes(split_count, shape);
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, split_v.get());
}
common::AnfAlgo::SetNodeAttr(kAttrSplitDim, MakeValue<int64_t>(split_dim), split_v);
common::AnfAlgo::SetNodeAttr(kAttrNumSplit, MakeValue<int64_t>(split_count), split_v);
common::AnfAlgo::SetNodeAttr(kAttrSizeSplits,
@ -117,11 +130,11 @@ CNodePtr AllToAllUnifyMindIR::CreateAllToAllvNode(const FuncGraphPtr &graph, con
(void)all_to_all_v_input.insert(all_to_all_v_input.end(), split_outputs.begin(), split_outputs.end());
auto all_to_all_v = NewCNode(all_to_all_v_input, graph);
MS_EXCEPTION_IF_NULL(all_to_all_v);
auto single_shape = common::AnfAlgo::GetOutputInferShape(split_outputs[0], 0);
auto single_shape = common::AnfAlgo::GetOutputDetailShape(split_outputs[0], 0);
auto single_type = common::AnfAlgo::GetOutputInferDataType(split_outputs[0], 0);
std::vector<TypeId> dtypes(split_count, single_type);
std::vector<std::vector<size_t>> shapes(split_count, single_shape);
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, all_to_all_v.get());
std::vector<BaseShapePtr> shapes(split_count, single_shape);
common::AnfAlgo::SetOutputTypeAndDetailShape(dtypes, shapes, all_to_all_v.get());
uint32_t rank_size = GetRankSize(group);
std::vector<int64_t> rank_ids(rank_size, 0);
for (uint32_t i = 0; i < rank_size; ++i) {
@ -160,8 +173,21 @@ CNodePtr AllToAllUnifyMindIR::CreateConcatNode(const FuncGraphPtr &graph, const
<< trace::DumpSourceLines(all_to_all);
}
single_shape[LongToSize(concat_dim)] *= static_cast<size_t>(split_count);
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(all_to_all_v_outputs[0], 0)},
{single_shape}, concat.get());
if (AnfUtils::IsShapeDynamic(single_shape)) {
auto min_shape = common::AnfAlgo::GetOutputMinShape(all_to_all_v_outputs[0], 0);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(all_to_all_v_outputs[0], 0);
max_shape[LongToSize(concat_dim)] *= split_count;
min_shape[LongToSize(concat_dim)] *= split_count;
ShapeVector new_shape;
std::transform(single_shape.begin(), single_shape.end(), std::back_inserter(new_shape), SizeToLong);
common::AnfAlgo::SetOutputTypeAndDetailShape({common::AnfAlgo::GetOutputInferDataType(all_to_all_v_outputs[0], 0)},
{std::make_shared<abstract::Shape>(new_shape, min_shape, max_shape)},
concat.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(all_to_all_v_outputs[0], 0)},
{single_shape}, concat.get());
}
common::AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue<int64_t>(concat_dim), concat);
common::AnfAlgo::SetNodeAttr(kAttrInputNums, MakeValue(split_count), concat);
std::vector<int64_t> dyn_input_size{split_count};

View File

@ -123,8 +123,19 @@ CNodePtr CreateTranspose(const FuncGraphPtr &graph, const CNodePtr &conv2d, cons
<< out_shape.size() << trace::DumpSourceLines(conv2d);
}
std::swap(out_shape[kDim0], out_shape[kDim1]);
auto shapes = {out_shape};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, transpose.get());
if (AnfUtils::IsShapeDynamic(out_shape)) {
ShapeVector new_shape;
auto min_shape = common::AnfAlgo::GetOutputMinShape(input_node, 0);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(input_node, 0);
std::swap(min_shape[kDim0], min_shape[kDim1]);
std::swap(max_shape[kDim0], max_shape[kDim1]);
std::transform(out_shape.begin(), out_shape.end(), std::back_inserter(new_shape), SizeToLong);
common::AnfAlgo::SetOutputTypeAndDetailShape(
types, {std::make_shared<abstract::Shape>(new_shape, min_shape, max_shape)}, transpose.get());
} else {
auto shapes = {out_shape};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, transpose.get());
}
} else {
transpose->set_abstract(conv2d->abstract());
}
@ -317,8 +328,19 @@ CNodePtr Conv2DBackpropFilterUnifyMindIR::CreateDepthwiseConv2DBackpropFilter(co
<< out_shape.size() << trace::DumpSourceLines(conv2d_backfil);
}
std::swap(out_shape[0], out_shape[1]);
auto shapes = {out_shape};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, depth_conv_backfil.get());
if (AnfUtils::IsShapeDynamic(out_shape)) {
ShapeVector new_shape;
auto min_shape = common::AnfAlgo::GetOutputMinShape(conv2d_backfil, 0);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(conv2d_backfil, 0);
std::swap(min_shape[0], min_shape[1]);
std::swap(max_shape[0], max_shape[1]);
std::transform(out_shape.begin(), out_shape.end(), std::back_inserter(new_shape), SizeToLong);
common::AnfAlgo::SetOutputTypeAndDetailShape(
types, {std::make_shared<abstract::Shape>(new_shape, min_shape, max_shape)}, depth_conv_backfil.get());
} else {
auto shapes = {out_shape};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, depth_conv_backfil.get());
}
return depth_conv_backfil;
}

View File

@ -49,9 +49,9 @@ void FakeLearnedScaleQuantPerLayerGradUnifyMindIR::CreateOutputsOfLSQPerLayerGra
auto types = {common::AnfAlgo::GetOutputInferDataType(lsq_perlayer_grad_node, 0),
common::AnfAlgo::GetOutputInferDataType(lsq_perlayer_grad_node, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(lsq_perlayer_grad_node, 0),
common::AnfAlgo::GetOutputInferShape(lsq_perlayer_grad_node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, lsq_perlayer_grad_d.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(lsq_perlayer_grad_node, 0),
common::AnfAlgo::GetOutputDetailShape(lsq_perlayer_grad_node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, lsq_perlayer_grad_d.get());
common::AnfAlgo::CopyNodeAttr(kAttrNeg_trunc, lsq_perlayer_grad_node, lsq_perlayer_grad_d);
CreateMultipleOutputsOfAnfNode(graph, lsq_perlayer_grad_d, kFakeLearnedScaleQuantGradDOutputNum,
@ -84,8 +84,8 @@ void FakeLearnedScaleQuantPerLayerGradUnifyMindIR::CreateOutputsOfLSQPerLayerRed
lsq_perlayer_reduce_grad->set_scope(lsq_perlayer_grad_node->scope());
auto types = {common::AnfAlgo::GetOutputInferDataType(lsq_perlayer_grad_node, 1)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(lsq_perlayer_grad_node, 1)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, lsq_perlayer_reduce_grad.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(lsq_perlayer_grad_node, 1)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, lsq_perlayer_reduce_grad.get());
(*lsq_perlayer_reduce_grad_outputs).push_back(lsq_perlayer_reduce_grad);
}
@ -111,9 +111,9 @@ void FakeLearnedScaleQuantPerChannelGradUnifyMindIR::CreateOutputsOfLSQPerChanne
auto types = {common::AnfAlgo::GetOutputInferDataType(lsq_perchannel_grad_node, 0),
common::AnfAlgo::GetOutputInferDataType(lsq_perchannel_grad_node, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(lsq_perchannel_grad_node, 0),
common::AnfAlgo::GetOutputInferShape(lsq_perchannel_grad_node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, lsq_perchannel_grad_d.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(lsq_perchannel_grad_node, 0),
common::AnfAlgo::GetOutputDetailShape(lsq_perchannel_grad_node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, lsq_perchannel_grad_d.get());
common::AnfAlgo::CopyNodeAttr(kAttrNeg_trunc, lsq_perchannel_grad_node, lsq_perchannel_grad_d);
common::AnfAlgo::CopyNodeAttr(kAttrChannelAxis, lsq_perchannel_grad_node, lsq_perchannel_grad_d);
@ -147,8 +147,8 @@ void FakeLearnedScaleQuantPerChannelGradUnifyMindIR::CreateOutputsOfLSQPerChanne
lsq_perchannel_reduce_grad->set_scope(lsq_perchannel_grad_node->scope());
auto types = {common::AnfAlgo::GetOutputInferDataType(lsq_perchannel_grad_node, 1)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(lsq_perchannel_grad_node, 1)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, lsq_perchannel_reduce_grad.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(lsq_perchannel_grad_node, 1)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, lsq_perchannel_reduce_grad.get());
common::AnfAlgo::CopyNodeAttr(kAttrChannelAxis, lsq_perchannel_grad_node, lsq_perchannel_reduce_grad);
(*lsq_perchannel_reduce_grad_outputs).push_back(lsq_perchannel_reduce_grad);
}

View File

@ -63,9 +63,9 @@ CNodePtr MaxPool2MaxPoolWithArgmax::CreateMaxPoolWithArgmax(const FuncGraphPtr &
// MaxPoolWithArgmax's second output is argmax, whose datatype is uint16 and with same shape as first output
TypeId argmax_dtype = kNumberTypeUInt16;
auto types = {common::AnfAlgo::GetOutputInferDataType(maxpool, 0), argmax_dtype};
auto out_shape = common::AnfAlgo::GetOutputInferShape(maxpool, 0);
auto shapes = {out_shape, out_shape};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, maxpool_argmax.get());
auto out_shape = common::AnfAlgo::GetOutputDetailShape(maxpool, 0);
std::vector<BaseShapePtr> shapes = {out_shape, out_shape};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, maxpool_argmax.get());
return maxpool_argmax;
}

View File

@ -18,6 +18,7 @@
#include <algorithm>
#include <vector>
#include <string>
#include <utility>
#include "backend/common/session/anf_runtime_algorithm.h"
#include "include/common/utils/anfalgo.h"
#include "plugin/device/ascend/hal/hccl_adapter/hccl_adapter.h"
@ -61,9 +62,12 @@ bool IsBottom(const std::vector<int64_t> &send_rank_ids) {
// cal split attrs size_splits, shapes and num_split
int64_t CalSplitAttrs(const std::vector<size_t> &base_shape, const bool is_first, const bool is_last,
const size_t split_dim, const std::vector<int64_t> &send_lens, std::vector<int64_t> *size_splits,
std::vector<std::vector<size_t>> *shapes) {
std::vector<std::vector<size_t>> *shapes, std::vector<int64_t> *min_shape,
std::vector<int64_t> *max_shape, bool is_dynamic) {
MS_EXCEPTION_IF_NULL(size_splits);
MS_EXCEPTION_IF_NULL(shapes);
MS_EXCEPTION_IF_NULL(max_shape);
MS_EXCEPTION_IF_NULL(min_shape);
if (SizeToLong(base_shape.size()) != kShapeSize) {
MS_LOG(EXCEPTION) << "Wrong base_shape size: " << base_shape.size() << ", it should be equal to 4.";
}
@ -84,6 +88,8 @@ int64_t CalSplitAttrs(const std::vector<size_t> &base_shape, const bool is_first
split_middle_size -= first_size;
shape_tmp[split_dim] = static_cast<size_t>(first_size);
shapes->push_back(shape_tmp);
(*min_shape)[split_dim] = (is_dynamic) ? first_size : (*min_shape)[split_dim];
(*max_shape)[split_dim] = (is_dynamic) ? first_size : (*max_shape)[split_dim];
}
if (is_last) {
// middle
@ -92,6 +98,8 @@ int64_t CalSplitAttrs(const std::vector<size_t> &base_shape, const bool is_first
++num_split;
size_splits->push_back(split_middle_size);
shape_tmp[split_dim] = static_cast<size_t>(split_middle_size);
(*min_shape)[split_dim] = (is_dynamic) ? split_middle_size : (*min_shape)[split_dim];
(*max_shape)[split_dim] = (is_dynamic) ? split_middle_size : (*max_shape)[split_dim];
shapes->push_back(shape_tmp);
}
// last
@ -103,6 +111,8 @@ int64_t CalSplitAttrs(const std::vector<size_t> &base_shape, const bool is_first
++num_split;
size_splits->push_back(split_middle_size);
shape_tmp[split_dim] = static_cast<size_t>(split_middle_size);
(*min_shape)[split_dim] = (is_dynamic) ? split_middle_size : (*min_shape)[split_dim];
(*max_shape)[split_dim] = (is_dynamic) ? split_middle_size : (*max_shape)[split_dim];
shapes->push_back(shape_tmp);
}
return num_split;
@ -110,7 +120,8 @@ int64_t CalSplitAttrs(const std::vector<size_t> &base_shape, const bool is_first
CNodePtr CreateSplitNode(const FuncGraphPtr &graph, const std::vector<AnfNodePtr> &split_input,
const std::vector<size_t> &base_shape, bool is_first, bool is_last, size_t split_dim,
const std::vector<int64_t> &send_lens, TypeId input_dtype, int64_t *num_split,
const std::vector<int64_t> &send_lens, TypeId input_dtype,
std::pair<ShapeVector, ShapeVector> *shape_pair, int64_t *num_split,
const PatternProcessPass &pass) {
MS_EXCEPTION_IF_NULL(graph);
MS_EXCEPTION_IF_NULL(num_split);
@ -122,10 +133,23 @@ CNodePtr CreateSplitNode(const FuncGraphPtr &graph, const std::vector<AnfNodePtr
MS_EXCEPTION_IF_NULL(split_v);
std::vector<int64_t> size_splits = {};
std::vector<std::vector<size_t>> shapes = {};
*num_split = CalSplitAttrs(base_shape, is_first, is_last, split_dim, send_lens, &size_splits, &shapes);
auto is_dynamic = AnfUtils::IsShapeDynamic(base_shape);
*num_split = CalSplitAttrs(base_shape, is_first, is_last, split_dim, send_lens, &size_splits, &shapes,
&shape_pair->first, &shape_pair->first, is_dynamic);
std::vector<TypeId> dtypes(*num_split, input_dtype);
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, split_v.get());
if (is_dynamic) {
std::vector<BaseShapePtr> shapes_ptr;
for (const auto &shape : shapes) {
ShapeVector shape_tmp;
std::transform(shape.begin(), shape.end(), std::back_inserter(shape_tmp), SizeToLong);
BaseShapePtr shape_ptr = std::make_shared<abstract::Shape>(shape_tmp, shape_pair->first, shape_pair->second);
shapes_ptr.push_back(shape_ptr);
}
common::AnfAlgo::SetOutputTypeAndDetailShape(dtypes, shapes_ptr, split_v.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, split_v.get());
}
common::AnfAlgo::SetNodeAttr(kAttrSplitDim, MakeValue<int64_t>(split_dim), split_v);
common::AnfAlgo::SetNodeAttr(kAttrNumSplit, MakeValue<int64_t>(*num_split), split_v);
common::AnfAlgo::SetNodeAttr(kAttrSizeSplits, MakeValue<std::vector<int64_t>>(size_splits), split_v);
@ -404,6 +428,10 @@ std::vector<CNodePtr> NeighborExchangeV2UnifyMindIR::CreateSplitNodes(const Func
auto dtype = common::AnfAlgo::GetOutputInferDataType(neighbor_exchange_v2_input, 0);
auto shape = common::AnfAlgo::GetOutputInferShape(neighbor_exchange_v2_input, 0);
auto is_dynamic = AnfUtils::IsShapeDynamic(shape);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(neighbor_exchange_v2_input, 0);
auto min_shape = common::AnfAlgo::GetOutputMinShape(neighbor_exchange_v2_input, 0);
auto shape_pair = std::make_pair(min_shape, max_shape);
if (SizeToLong(shape.size()) != kShapeSize) { // only support NCHW now
MS_LOG(EXCEPTION) << "Invalid shape size " << shape.size() << ", only support NCHW input now!"
<< trace::DumpSourceLines(neighbor_exchange_v2);
@ -425,7 +453,7 @@ std::vector<CNodePtr> NeighborExchangeV2UnifyMindIR::CreateSplitNodes(const Func
neighbor_exchange_v2_input};
split_v = CreateSplitNode(graph, split_input, shape, splitvs_is_first[i], !splitvs_is_first[i], splitvs_dim[i],
send_lens, dtype, &num_split, *this);
send_lens, dtype, &shape_pair, &num_split, *this);
}
(void)split_nodes.emplace_back(split_v);
split_num->push_back(num_split);
@ -461,13 +489,17 @@ std::vector<CNodePtr> NeighborExchangeV2UnifyMindIR::CreateSplitNodes(const Func
if (corner_splitvs_is_input_top[i]) {
(void)split_input.insert(split_input.end(), split_outputs_top.begin(), split_outputs_top.begin() + 1);
shape_tmp[kHDim] = send_lens[0];
min_shape[kHDim] = (is_dynamic) ? send_lens[0] : min_shape[kHDim];
max_shape[kHDim] = (is_dynamic) ? send_lens[0] : max_shape[kHDim];
} else {
(void)split_input.insert(split_input.end(), split_outputs_bottom.end() - 1, split_outputs_bottom.end());
shape_tmp[kHDim] = send_lens[1];
min_shape[kHDim] = (is_dynamic) ? send_lens[1] : min_shape[kHDim];
max_shape[kHDim] = (is_dynamic) ? send_lens[1] : max_shape[kHDim];
}
auto pair_tmp = std::make_pair(min_shape, max_shape);
split_v = CreateSplitNode(graph, split_input, shape_tmp, corner_splitvs_is_first[i], !corner_splitvs_is_first[i],
kWDim, send_lens, dtype, &num_split, *this);
kWDim, send_lens, dtype, &pair_tmp, &num_split, *this);
}
(void)split_nodes.emplace_back(split_v);
split_num->push_back(num_split);
@ -477,14 +509,11 @@ std::vector<CNodePtr> NeighborExchangeV2UnifyMindIR::CreateSplitNodes(const Func
}
CNodePtr NeighborExchangeV2UnifyMindIR::CreateConcatNode(const FuncGraphPtr &graph,
const std::vector<AnfNodePtr> &concat_input,
const std::vector<std::vector<size_t>> &output_shape,
const std::vector<TypeId> &output_dtype, int64_t axis,
const std::vector<AnfNodePtr> &concat_input, int64_t axis,
int64_t input_nums) const {
MS_EXCEPTION_IF_NULL(graph);
auto concat = NewCNode(concat_input, graph);
MS_EXCEPTION_IF_NULL(concat);
common::AnfAlgo::SetOutputInferTypeAndShape(output_dtype, output_shape, concat.get());
common::AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue<int64_t>(axis), concat);
common::AnfAlgo::SetNodeAttr(kAttrInputNums, MakeValue(input_nums), concat);
std::vector<int64_t> dyn_input_size_empty{input_nums};
@ -507,14 +536,23 @@ CNodePtr NeighborExchangeV2UnifyMindIR::CreateLeftRightConcat(const FuncGraphPtr
auto single_shape =
common::AnfAlgo::GetOutputInferShape(all_to_all_v_outputs[AllToAllRealIds(middle_ids, recv_rank_ids)], 0);
auto max_shape =
common::AnfAlgo::GetOutputMaxShape(all_to_all_v_outputs[AllToAllRealIds(middle_ids, recv_rank_ids)], 0);
auto min_shape =
common::AnfAlgo::GetOutputMinShape(all_to_all_v_outputs[AllToAllRealIds(middle_ids, recv_rank_ids)], 0);
auto is_dynamic = AnfUtils::IsShapeDynamic(single_shape);
if (recv_rank_ids[first_ids] != kInvalidId) {
++input_num;
single_shape[kDim2] += static_cast<size_t>(recv_lens[0]); // H in NCHW
max_shape[kDim2] += (is_dynamic) ? recv_lens[0] : 0;
min_shape[kDim2] += (is_dynamic) ? recv_lens[0] : 0;
}
if (recv_rank_ids[last_ids] != kInvalidId) {
++input_num;
single_shape[kDim2] += static_cast<size_t>(recv_lens[1]); // H in NCHW
max_shape[kDim2] += (is_dynamic) ? recv_lens[1] : 0;
min_shape[kDim2] += (is_dynamic) ? recv_lens[1] : 0;
}
if (is_left) {
(void)concat_input.insert(concat_input.end(), all_to_all_v_outputs.rbegin(),
@ -526,8 +564,15 @@ CNodePtr NeighborExchangeV2UnifyMindIR::CreateLeftRightConcat(const FuncGraphPtr
std::vector<TypeId> concat_output_dtype = {
common::AnfAlgo::GetOutputInferDataType(all_to_all_v_outputs[AllToAllRealIds(middle_ids, recv_rank_ids)], 0)};
auto concat = CreateConcatNode(graph, concat_input, {single_shape}, concat_output_dtype, kHDim, input_num);
auto concat = CreateConcatNode(graph, concat_input, kHDim, input_num);
if (is_dynamic) {
ShapeVector shape;
std::transform(single_shape.begin(), single_shape.end(), std::back_inserter(shape), SizeToLong);
BaseShapePtr base_shape = std::make_shared<abstract::Shape>(shape, min_shape, max_shape);
common::AnfAlgo::SetOutputTypeAndDetailShape(concat_output_dtype, {base_shape}, concat.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape(concat_output_dtype, {single_shape}, concat.get());
}
return concat;
}
@ -538,6 +583,9 @@ CNodePtr NeighborExchangeV2UnifyMindIR::CreateMiddleConcat(
int64_t input_num_all = 0;
auto neighbor_exchange_v2_input = neighbor_exchange_v2->input(kNeighborExchangeV2InputIdx);
auto single_shape = common::AnfAlgo::GetOutputInferShape(neighbor_exchange_v2_input, 0);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(neighbor_exchange_v2_input, 0);
auto min_shape = common::AnfAlgo::GetOutputMinShape(neighbor_exchange_v2_input, 0);
auto is_dynamic = AnfUtils::IsShapeDynamic(single_shape);
size_t first_idx = concat_dim == kWDim ? 6 : 0;
size_t last_idx = concat_dim == kWDim ? 2 : 4;
size_t first_len = concat_dim == kWDim ? static_cast<size_t>(recv_lens[kDim2]) : static_cast<size_t>(recv_lens[0]);
@ -554,6 +602,8 @@ CNodePtr NeighborExchangeV2UnifyMindIR::CreateMiddleConcat(
++input_num_all;
single_shape[concat_dim] += first_len;
max_shape[concat_dim] += (is_dynamic) ? first_len : 0;
min_shape[concat_dim] += (is_dynamic) ? first_len : 0;
}
concat_input_all.push_back(neighbor_exchange_v2_input);
@ -571,11 +621,20 @@ CNodePtr NeighborExchangeV2UnifyMindIR::CreateMiddleConcat(
++input_num_all;
single_shape[concat_dim] += last_len;
max_shape[concat_dim] += (is_dynamic) ? last_len : 0;
min_shape[concat_dim] += (is_dynamic) ? last_len : 0;
}
std::vector<TypeId> concat_output_dtype = {common::AnfAlgo::GetOutputInferDataType(all_to_all_v_outputs[0], 0)};
auto concat_all =
CreateConcatNode(graph, concat_input_all, {single_shape}, concat_output_dtype, concat_dim, input_num_all);
auto concat_all = CreateConcatNode(graph, concat_input_all, concat_dim, input_num_all);
if (is_dynamic) {
ShapeVector shape;
std::transform(single_shape.begin(), single_shape.end(), std::back_inserter(shape), SizeToLong);
BaseShapePtr base_shape = std::make_shared<abstract::Shape>(shape, min_shape, max_shape);
common::AnfAlgo::SetOutputTypeAndDetailShape(concat_output_dtype, {base_shape}, concat_all.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape(concat_output_dtype, {single_shape}, concat_all.get());
}
return concat_all;
}
@ -687,8 +746,8 @@ CNodePtr NeighborExchangeV2UnifyMindIR::CreateConcatNodes(const FuncGraphPtr &gr
}
std::vector<TypeId> concat_right_output_dtype = {common::AnfAlgo::GetOutputInferDataType(concat_input_all[1], 0)};
auto concat_all =
CreateConcatNode(graph, concat_input_all, {shape_all}, concat_right_output_dtype, kWDim, input_nums_all);
auto concat_all = CreateConcatNode(graph, concat_input_all, kWDim, input_nums_all);
common::AnfAlgo::SetOutputInferTypeAndShape(concat_right_output_dtype, {shape_all}, concat_all.get());
return concat_all;
}
@ -713,6 +772,10 @@ std::vector<CNodePtr> NeighborExchangeV2GradUnifyMindIR::CreateSplitNodesForGrad
auto neighbor_exchange_v2_grad_input = neighbor_exchange_v2_grad->input(kNeighborExchangeV2InputIdx);
auto dtype = common::AnfAlgo::GetOutputInferDataType(neighbor_exchange_v2_grad_input, 0);
auto shape = common::AnfAlgo::GetOutputInferShape(neighbor_exchange_v2_grad_input, 0);
auto is_dynamic = AnfUtils::IsShapeDynamic(shape);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(neighbor_exchange_v2_grad_input, 0);
auto min_shape = common::AnfAlgo::GetOutputMinShape(neighbor_exchange_v2_grad_input, 0);
if (SizeToLong(shape.size()) != kShapeSize) {
MS_LOG(EXCEPTION) << "Invalid shape size " << shape.size() << ", only support NCHW input now!"
<< trace::DumpSourceLines(neighbor_exchange_v2_grad);
@ -727,8 +790,9 @@ std::vector<CNodePtr> NeighborExchangeV2GradUnifyMindIR::CreateSplitNodesForGrad
if (is_top || is_bottom) {
std::vector<AnfNodePtr> split_input = {NewValueNode(std::make_shared<Primitive>(prim::kPrimSplitV->name())),
neighbor_exchange_v2_grad_input};
split_v_top_bottom =
CreateSplitNode(graph, split_input, shape, is_top, is_bottom, kHDim, send_lens, dtype, &num_split_h, *this);
auto pair_tmp = std::make_pair(max_shape, min_shape);
split_v_top_bottom = CreateSplitNode(graph, split_input, shape, is_top, is_bottom, kHDim, send_lens, dtype,
&pair_tmp, &num_split_h, *this);
}
(void)split_nodes.emplace_back(split_v_top_bottom);
split_num->push_back(num_split_h);
@ -767,8 +831,11 @@ std::vector<CNodePtr> NeighborExchangeV2GradUnifyMindIR::CreateSplitNodesForGrad
int64_t num_split_w = 0;
std::vector<size_t> base_shape(shape);
base_shape[kHDim] = static_cast<size_t>(size_split_h[i]);
min_shape[kHDim] = (is_dynamic) ? size_split_h[i] : min_shape[kHDim];
max_shape[kHDim] = (is_dynamic) ? size_split_h[i] : max_shape[kHDim];
auto pair_tmp = std::make_pair(min_shape, max_shape);
auto split_v_left_right = CreateSplitNode(graph, split_input, base_shape, is_left, is_right, kWDim, send_lens,
dtype, &num_split_w, *this);
dtype, &pair_tmp, &num_split_w, *this);
(void)split_nodes.emplace_back(split_v_left_right);
split_num->push_back(num_split_w);
}
@ -788,19 +855,22 @@ std::vector<CNodePtr> NeighborExchangeV2GradUnifyMindIR::CreateSplitNodesForGrad
return split_nodes;
}
CNodePtr NeighborExchangeV2GradUnifyMindIR::CreatePadNode(const FuncGraphPtr &graph, const AnfNodePtr &input,
const std::vector<int64_t> &begin,
const std::vector<int64_t> &size,
const std::vector<size_t> &shape, TypeId dtype) const {
CNodePtr NeighborExchangeV2GradUnifyMindIR::CreatePadNode(
const FuncGraphPtr &graph, const AnfNodePtr &input, const std::vector<int64_t> &begin,
const std::vector<int64_t> &size, const std::pair<std::vector<size_t>, BaseShapePtr> &shape_info,
TypeId dtype) const {
MS_EXCEPTION_IF_NULL(graph);
MS_EXCEPTION_IF_NULL(input);
auto shape = shape_info.first;
auto shape_base = shape_info.second;
MS_EXCEPTION_IF_NULL(shape_base);
std::vector<AnfNodePtr> pad_inputs = {NewValueNode(std::make_shared<Primitive>(kPadOpName)), input};
auto pad = NewCNode(pad_inputs, graph);
std::vector<std::vector<int64_t>> paddings;
for (size_t i = 0; i < shape.size(); ++i) {
(void)paddings.emplace_back(std::vector<int64_t>{begin[i], static_cast<int64_t>(shape[i]) - begin[i] - size[i]});
}
common::AnfAlgo::SetOutputInferTypeAndShape({dtype}, {shape}, pad.get());
common::AnfAlgo::SetOutputTypeAndDetailShape({dtype}, {shape_base}, pad.get());
common::AnfAlgo::SetNodeAttr(kAttrPaddings, MakeValue(paddings), pad);
common::AnfAlgo::SetNodeAttr(kAttrInputNames, MakeValue(std::vector<std::string>{"x"}), pad);
return pad;
@ -824,6 +894,7 @@ CNodePtr NeighborExchangeV2GradUnifyMindIR::CreateSplitGradNodes(const FuncGraph
auto centerx = GetCenter(graph, neighbor_exchange_v2_grad, split_nodes, split_num, send_rank_ids);
auto centerx_dtype = common::AnfAlgo::GetOutputInferDataType(centerx, 0);
auto centerx_shape = common::AnfAlgo::GetOutputInferShape(centerx, 0);
auto base_shape = common::AnfAlgo::GetOutputDetailShape(centerx, 0);
// empty
int64_t all_to_all_output_num =
std::count_if(recv_rank_ids.begin(), recv_rank_ids.end(), [](int64_t ids) { return ids != kInvalidId; });
@ -872,8 +943,9 @@ CNodePtr NeighborExchangeV2GradUnifyMindIR::CreateSplitGradNodes(const FuncGraph
size_t output_index = 0;
for (size_t i = 0; i < recv_rank_ids.size(); ++i) {
if (recv_rank_ids[i] != kInvalidId) {
auto shape_info = std::make_pair(centerx_shape, base_shape);
auto pad =
CreatePadNode(graph, all_to_all_v_outputs[output_index], begins[i], sizes[i], centerx_shape, centerx_dtype);
CreatePadNode(graph, all_to_all_v_outputs[output_index], begins[i], sizes[i], shape_info, centerx_dtype);
++output_index;
(void)pad_nodes.emplace_back(pad);
}
@ -894,7 +966,7 @@ CNodePtr NeighborExchangeV2GradUnifyMindIR::CreateSplitGradNodes(const FuncGraph
}
auto addn = NewCNode(addn_inputs, graph);
MS_EXCEPTION_IF_NULL(addn);
common::AnfAlgo::SetOutputInferTypeAndShape({centerx_dtype}, {centerx_shape}, addn.get());
common::AnfAlgo::SetOutputTypeAndDetailShape({centerx_dtype}, {base_shape}, addn.get());
common::AnfAlgo::SetNodeAttr(kAttrDynInputSizes, MakeValue<std::vector<int64_t>>({pad_num}), addn);
common::AnfAlgo::SetNodeAttr(kAttrN, MakeValue(pad_num), addn);
MS_LOG(DEBUG) << "Create splitvs grad nodes success.";

View File

@ -18,6 +18,7 @@
#include <memory>
#include <vector>
#include <utility>
#include "backend/common/optimizer/optimizer.h"
#include "backend/common/session/anf_runtime_algorithm.h"
#include "include/common/utils/anfalgo.h"
@ -35,9 +36,8 @@ class NeighborExchangeV2UnifyMindIR : public PatternProcessPass {
private:
std::vector<CNodePtr> CreateSplitNodes(const FuncGraphPtr &graph, const CNodePtr &neighbor_exchange_v2,
std::vector<int64_t> *split_num) const;
CNodePtr CreateConcatNode(const FuncGraphPtr &graph, const std::vector<AnfNodePtr> &concat_input,
const std::vector<std::vector<size_t>> &output_shape,
const std::vector<TypeId> &output_dtype, int64_t axis, int64_t input_nums) const;
CNodePtr CreateConcatNode(const FuncGraphPtr &graph, const std::vector<AnfNodePtr> &concat_input, int64_t axis,
int64_t input_nums) const;
CNodePtr CreateLeftRightConcat(const FuncGraphPtr &graph, const std::vector<AnfNodePtr> &all_to_all_v_outputs,
const std::vector<int64_t> &recv_rank_ids, const std::vector<int64_t> &recv_lens,
bool is_left) const;
@ -63,7 +63,8 @@ class NeighborExchangeV2GradUnifyMindIR : public PatternProcessPass {
std::vector<CNodePtr> CreateSplitNodesForGrad(const FuncGraphPtr &graph, const CNodePtr &neighbor_exchange_v2_grad,
std::vector<int64_t> *split_num) const;
CNodePtr CreatePadNode(const FuncGraphPtr &graph, const AnfNodePtr &input, const std::vector<int64_t> &begin,
const std::vector<int64_t> &size, const std::vector<size_t> &shape, TypeId dtype) const;
const std::vector<int64_t> &size,
const std::pair<std::vector<size_t>, BaseShapePtr> &shape_info, TypeId dtype) const;
CNodePtr CreateSplitGradNodes(const FuncGraphPtr &graph, const CNodePtr &neighbor_exchange_v2_grad,
const CNodePtr &all_to_all_v, const std::vector<CNodePtr> &split_nodes,
const std::vector<int64_t> &split_num) const;

View File

@ -104,7 +104,19 @@ CNodePtr CreateOneHot(const FuncGraphPtr &graph, const CNodePtr &sparse_softmax_
one_hot_node->set_scope(sparse_softmax_node->scope());
std::vector<size_t> labels_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(sparse_softmax_node, 1);
labels_shape.emplace_back(depth);
common::AnfAlgo::SetOutputInferTypeAndShape({kNumberTypeFloat32}, {labels_shape}, one_hot_node.get());
if (AnfUtils::IsShapeDynamic(labels_shape)) {
auto kernel_info = common::AnfAlgo::GetPrevNodeOutput(sparse_softmax_node, 1);
auto min_shape = common::AnfAlgo::GetOutputMinShape(kernel_info.first, kernel_info.second);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(kernel_info.first, kernel_info.second);
std::vector<int64_t> shape_tmp;
std::transform(labels_shape.begin(), labels_shape.end(), std::back_inserter(shape_tmp), SizeToLong);
min_shape.emplace_back(depth);
max_shape.emplace_back(depth);
common::AnfAlgo::SetOutputTypeAndDetailShape(
{kNumberTypeFloat32}, {std::make_shared<abstract::Shape>(shape_tmp, min_shape, max_shape)}, one_hot_node.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape({kNumberTypeFloat32}, {labels_shape}, one_hot_node.get());
}
if (is_convert_const_to_attr) {
common::AnfAlgo::SetNodeAttr(kAttrDepth, MakeValue(depth), one_hot_node);
}
@ -131,10 +143,20 @@ CNodePtr CreateSoftmaxCrossEntropyWithLogits(const FuncGraphPtr &graph, const CN
MS_LOG(EXCEPTION) << "One_hot output's shape is empty." << trace::DumpSourceLines(one_hot_node);
}
auto shapes = {loss_shape, common::AnfAlgo::GetOutputInferShape(one_hot_node, 0)};
auto data_types = common::AnfAlgo::GetOutputInferDataType(one_hot_node, 0);
auto types = {data_types, data_types};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, softmax_node.get());
if (AnfUtils::IsShapeDynamic(labels_shape)) {
ShapeVector shape_tmp = {static_cast<int64_t>(labels_shape[0])};
auto min_shape = common::AnfAlgo::GetOutputMinShape(one_hot_node, 0);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(one_hot_node, 0);
std::vector<BaseShapePtr> shapes = {
std::make_shared<abstract::Shape>(shape_tmp, ShapeVector(min_shape[0]), ShapeVector(max_shape[0])),
common::AnfAlgo::GetOutputDetailShape(one_hot_node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, softmax_node.get());
} else {
auto shapes = {loss_shape, labels_shape};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, softmax_node.get());
}
return softmax_node;
}
@ -223,8 +245,21 @@ CNodePtr CreateExpandDims(const FuncGraphPtr &graph, const CNodePtr &real_div_no
expand_dims_node->set_scope(real_div_node->scope());
std::vector<size_t> y_shape = common::AnfAlgo::GetOutputInferShape(real_div_node, 0);
y_shape.emplace_back(1);
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(real_div_node, 0)}, {y_shape},
expand_dims_node.get());
if (AnfUtils::IsShapeDynamic(y_shape)) {
auto min_shape = common::AnfAlgo::GetOutputMinShape(real_div_node, 0);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(real_div_node, 0);
min_shape.emplace_back(1);
max_shape.emplace_back(1);
std::vector<int64_t> shape_tmp;
std::transform(y_shape.begin(), y_shape.end(), std::back_inserter(shape_tmp), SizeToLong);
common::AnfAlgo::SetOutputTypeAndDetailShape({common::AnfAlgo::GetOutputInferDataType(real_div_node, 0)},
{std::make_shared<abstract::Shape>(shape_tmp, min_shape, max_shape)},
expand_dims_node.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(real_div_node, 0)}, {y_shape},
expand_dims_node.get());
}
return expand_dims_node;
}
@ -247,8 +282,20 @@ CNodePtr CreateExpandDimsPynative(const FuncGraphPtr &graph, const CNodePtr &rea
expand_dims_node->set_scope(real_div_node->scope());
std::vector<size_t> y_shape = common::AnfAlgo::GetOutputInferShape(real_div_node, 0);
y_shape.emplace_back(1);
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(real_div_node, 0)}, {y_shape},
expand_dims_node.get());
if (AnfUtils::IsShapeDynamic(y_shape)) {
auto min_shape = common::AnfAlgo::GetOutputMinShape(real_div_node, 0);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(real_div_node, 0);
min_shape.emplace_back(1);
max_shape.emplace_back(1);
std::vector<int64_t> shape_tmp;
std::transform(y_shape.begin(), y_shape.end(), std::back_inserter(shape_tmp), SizeToLong);
common::AnfAlgo::SetOutputTypeAndDetailShape({common::AnfAlgo::GetOutputInferDataType(real_div_node, 0)},
{std::make_shared<abstract::Shape>(shape_tmp, min_shape, max_shape)},
expand_dims_node.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(real_div_node, 0)}, {y_shape},
expand_dims_node.get());
}
common::AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue(axis), expand_dims_node);
return expand_dims_node;
}
@ -290,8 +337,9 @@ CNodePtr CreateTile(const FuncGraphPtr &graph, const CNodePtr &sparse_softmax_no
auto tile_node = pass.NewCNode(tile_inputs, graph);
MS_EXCEPTION_IF_NULL(tile_node);
tile_node->set_scope(mul_node->scope());
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetPrevNodeOutputInferDataType(mul_node, 1)},
{labels_shape}, tile_node.get());
common::AnfAlgo::SetOutputTypeAndDetailShape({common::AnfAlgo::GetPrevNodeOutputInferDataType(mul_node, 1)},
{common::AnfAlgo::GetPrevNodeOutputDetailShape(sparse_softmax_node, 1)},
tile_node.get());
if (is_convert_const_to_attr) {
common::AnfAlgo::SetNodeAttr(kAttrMultiples, MakeValue(multiples), tile_node);
}

View File

@ -63,8 +63,8 @@ CNodePtr InsertTransposeOp(const FuncGraphPtr &graph, const AnfNodePtr &node, co
auto transpose_op = graph->NewCNode(transpose_input);
// 3.Set the output info of transpose.
auto transpose_type = {common::AnfAlgo::GetPrevNodeOutputInferDataType(used_node, used_node_index)};
auto transpose_shape = {common::AnfAlgo::GetPrevNodeOutputInferShape(used_node, used_node_index)};
common::AnfAlgo::SetOutputInferTypeAndShape(transpose_type, transpose_shape, transpose_op.get());
auto transpose_shape = {common::AnfAlgo::GetPrevNodeOutputDetailShape(used_node, used_node_index)};
common::AnfAlgo::SetOutputTypeAndDetailShape(transpose_type, transpose_shape, transpose_op.get());
common::AnfAlgo::SetNodeAttr(kAttrPerm, MakeValue(transpose_perm), transpose_op);
// 4. Set the new edge of transpose op.
FuncGraphManagerPtr manager = graph->manager();

View File

@ -165,8 +165,8 @@ const AnfNodePtr AdamFusion::Process(const FuncGraphPtr &graph, const AnfNodePtr
auto adam = graph->NewCNode(inputs);
MS_EXCEPTION_IF_NULL(adam);
auto types = {common::AnfAlgo::GetOutputInferDataType(node, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, adam.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, adam.get());
adam->set_scope(node->scope());
auto build_info = GenerateKernelBuildInfo(adam);
AnfAlgo::SetSelectKernelBuildInfo(build_info, adam.get());

View File

@ -170,8 +170,8 @@ const AnfNodePtr AdamWeightDecayFusion::Process(const FuncGraphPtr &graph, const
auto adam_weight_decay = graph->NewCNode(inputs);
MS_EXCEPTION_IF_NULL(adam_weight_decay);
auto types = {common::AnfAlgo::GetOutputInferDataType(node, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, adam_weight_decay.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, adam_weight_decay.get());
adam_weight_decay->set_scope(node->scope());
auto build_info = GenerateKernelBuildInfo(adam_weight_decay);

View File

@ -89,8 +89,8 @@ const AnfNodePtr AddReluGradV2Fusion::Process(const FuncGraphPtr &graph, const A
auto add_relugrad = graph->NewCNode(inputs);
MS_EXCEPTION_IF_NULL(add_relugrad);
auto types = {common::AnfAlgo::GetOutputInferDataType(node, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, add_relugrad.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, add_relugrad.get());
add_relugrad->set_scope(node->scope());
auto build_info = GenerateKernelBuildInfo(add_relugrad);

View File

@ -88,14 +88,14 @@ const AnfNodePtr AddReluV2Fusion::Process(const FuncGraphPtr &graph, const AnfNo
MS_EXCEPTION_IF_NULL(add_relu);
std::vector<TypeId> types;
std::vector<std::vector<size_t>> shapes;
std::vector<BaseShapePtr> shapes;
size_t output_num = common::AnfAlgo::GetOutputTensorNum(node);
for (size_t i = 0; i < output_num; i++) {
types.push_back(common::AnfAlgo::GetOutputInferDataType(node, i));
shapes.push_back(common::AnfAlgo::GetOutputInferShape(node, i));
shapes.push_back(common::AnfAlgo::GetOutputDetailShape(node, i));
}
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, add_relu.get());
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, add_relu.get());
add_relu->set_scope(node->scope());
auto build_info = GenerateKernelBuildInfo(add_relu);

View File

@ -69,8 +69,20 @@ CNodePtr CreateSplitNode(const FuncGraphPtr &graph, const CNodePtr &all_to_all)
// Set Split CNode outputs type and shape, and CNode attributes.
std::vector<TypeId> dtypes(split_count, dtype);
std::vector<std::vector<size_t>> shapes(split_count, shape);
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, split.get());
if (AnfUtils::IsShapeDynamic(shape)) {
ShapeVector shape_tmp;
auto min_shape = common::AnfAlgo::GetOutputMinShape(all_to_all_input, 0);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(all_to_all_input, 0);
min_shape[LongToSize(split_dim)] /= split_count;
max_shape[LongToSize(split_dim)] /= split_count;
std::transform(shape.begin(), shape.end(), std::back_inserter(shape_tmp), SizeToLong);
std::vector<BaseShapePtr> shapes(split_count, std::make_shared<abstract::Shape>(shape_tmp, min_shape, max_shape));
common::AnfAlgo::SetOutputTypeAndDetailShape(dtypes, shapes, split.get());
} else {
std::vector<std::vector<size_t>> shapes(split_count, shape);
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, split.get());
}
common::AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue<int64_t>(split_dim), split);
common::AnfAlgo::SetNodeAttr(kAttrOutputNum, MakeValue<int64_t>(split_count), split);
return split;
@ -95,11 +107,11 @@ CNodePtr CreateAllToAllvNode(const FuncGraphPtr &graph, const CNodePtr &all_to_a
MS_EXCEPTION_IF_NULL(all_to_all_v);
// Prepare dtypes, shapes and ranks vectors.
auto single_shape = common::AnfAlgo::GetOutputInferShape(split_outputs[0], 0);
auto single_shape = common::AnfAlgo::GetOutputDetailShape(split_outputs[0], 0);
auto single_type = common::AnfAlgo::GetOutputInferDataType(split_outputs[0], 0);
std::vector<TypeId> dtypes(split_count, single_type);
std::vector<std::vector<size_t>> shapes(split_count, single_shape);
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, all_to_all_v.get());
std::vector<BaseShapePtr> shapes(split_count, single_shape);
common::AnfAlgo::SetOutputTypeAndDetailShape(dtypes, shapes, all_to_all_v.get());
uint32_t rank_size = device::gpu::CollectiveInitializer::instance().GetGroupSize(group);
std::vector<int64_t> rank_ids(rank_size, 0);
for (uint32_t i = 0; i < rank_size; ++i) {
@ -141,8 +153,20 @@ CNodePtr CreateConcatNode(const FuncGraphPtr &graph, const CNodePtr &all_to_all,
// Set Concat CNode outputs and attributes.
single_shape[LongToSize(concat_dim)] *= split_count;
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(all_to_all_v_outputs[0], 0)},
{single_shape}, concat.get());
if (AnfUtils::IsShapeDynamic(single_shape)) {
ShapeVector shape_tmp;
auto min_shape = common::AnfAlgo::GetOutputMinShape(all_to_all_v_outputs[0], 0);
auto max_shape = common::AnfAlgo::GetOutputMaxShape(all_to_all_v_outputs[0], 0);
min_shape[LongToSize(concat_dim)] *= split_count;
max_shape[LongToSize(concat_dim)] *= split_count;
std::transform(single_shape.begin(), single_shape.end(), std::back_inserter(shape_tmp), SizeToLong);
common::AnfAlgo::SetOutputTypeAndDetailShape({common::AnfAlgo::GetOutputInferDataType(all_to_all_v_outputs[0], 0)},
{std::make_shared<abstract::Shape>(shape_tmp)}, concat.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(all_to_all_v_outputs[0], 0)},
{single_shape}, concat.get());
}
common::AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue<int64_t>(concat_dim), concat);
common::AnfAlgo::SetNodeAttr(kAttrInputNums, MakeValue(split_count), concat);
std::vector<int64_t> dyn_input_size{split_count};

View File

@ -89,8 +89,8 @@ const AnfNodePtr ApplyMomentumScaleFusion::Process(const FuncGraphPtr &graph, co
auto replace_node = graph->NewCNode(inputs);
MS_EXCEPTION_IF_NULL(replace_node);
auto types = {common::AnfAlgo::GetOutputInferDataType(node, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, replace_node.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, replace_node.get());
replace_node->set_scope(node->scope());
return replace_node;
}

View File

@ -61,8 +61,8 @@ const AnfNodePtr ApplyMomentumWeightDecayFusion::Process(const FuncGraphPtr &gra
auto replace_node = graph->NewCNode(inputs);
MS_EXCEPTION_IF_NULL(replace_node);
auto types = {common::AnfAlgo::GetOutputInferDataType(node, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, replace_node.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, replace_node.get());
replace_node->set_scope(node->scope());
return replace_node;
}

View File

@ -126,8 +126,8 @@ const AnfNodePtr ApplyMomentumWeightDecayScaleFusion::Process(const FuncGraphPtr
auto replace_node = graph->NewCNode(inputs);
MS_EXCEPTION_IF_NULL(replace_node);
auto types = {common::AnfAlgo::GetOutputInferDataType(node, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, replace_node.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, replace_node.get());
replace_node->set_scope(node->scope());
return replace_node;
}

View File

@ -99,13 +99,13 @@ const AnfNodePtr BatchNormAddReluFusion::Process(const FuncGraphPtr &graph, cons
MS_EXCEPTION_IF_NULL(fused_batch_norm_with_add_relu);
std::vector<TypeId> outputs_type;
std::vector<std::vector<size_t>> outputs_shape;
std::vector<BaseShapePtr> outputs_shape;
auto output_num = common::AnfAlgo::GetOutputTensorNum(batch_norm);
for (size_t i = 0; i < output_num; i++) {
outputs_type.push_back(common::AnfAlgo::GetOutputInferDataType(batch_norm, i));
outputs_shape.push_back(common::AnfAlgo::GetOutputInferShape(batch_norm, i));
outputs_shape.push_back(common::AnfAlgo::GetOutputDetailShape(batch_norm, i));
}
common::AnfAlgo::SetOutputInferTypeAndShape(outputs_type, outputs_shape, fused_batch_norm_with_add_relu.get());
common::AnfAlgo::SetOutputTypeAndDetailShape(outputs_type, outputs_shape, fused_batch_norm_with_add_relu.get());
common::AnfAlgo::CopyNodeAttrs(batch_norm, fused_batch_norm_with_add_relu);
auto manager = graph->manager();

View File

@ -69,16 +69,16 @@ bool GetBatchNormOutputs(const FuncGraphPtr &func_graph, const AnfNodePtr &bn, s
void SetShapeAndType(const CNodePtr &bn_add_relu_grad, const AnfNodePtr &bn_grad, const AnfNodePtr &relu_grad) {
// set output shape and dtype
std::vector<TypeId> outputs_type;
std::vector<std::vector<size_t>> outputs_shape;
std::vector<BaseShapePtr> outputs_shape;
auto output_num = common::AnfAlgo::GetOutputTensorNum(bn_grad);
for (size_t i = 0; i < output_num; ++i) {
outputs_type.push_back(common::AnfAlgo::GetOutputInferDataType(bn_grad, i));
outputs_shape.push_back(common::AnfAlgo::GetOutputInferShape(bn_grad, i));
outputs_shape.push_back(common::AnfAlgo::GetOutputDetailShape(bn_grad, i));
}
outputs_type.push_back(common::AnfAlgo::GetOutputInferDataType(relu_grad, 0));
outputs_shape.push_back(common::AnfAlgo::GetOutputInferShape(relu_grad, 0));
common::AnfAlgo::SetOutputInferTypeAndShape(outputs_type, outputs_shape, bn_add_relu_grad.get());
outputs_shape.push_back(common::AnfAlgo::GetOutputDetailShape(relu_grad, 0));
common::AnfAlgo::SetOutputTypeAndDetailShape(outputs_type, outputs_shape, bn_add_relu_grad.get());
}
void ReplaceOutput(const FuncGraphPtr &graph, const AnfNodePtr &bn_grad, const AnfNodePtr &relu_grad,

View File

@ -94,13 +94,13 @@ const AnfNodePtr BatchNormReluFusion::Process(const FuncGraphPtr &graph, const A
MS_EXCEPTION_IF_NULL(fused_batch_norm_with_relu);
std::vector<TypeId> outputs_type;
std::vector<std::vector<size_t>> outputs_shape;
std::vector<BaseShapePtr> outputs_shape;
auto output_num = common::AnfAlgo::GetOutputTensorNum(batch_norm);
for (size_t i = 0; i < output_num; i++) {
outputs_type.push_back(common::AnfAlgo::GetOutputInferDataType(batch_norm, i));
outputs_shape.push_back(common::AnfAlgo::GetOutputInferShape(batch_norm, i));
outputs_shape.push_back(common::AnfAlgo::GetOutputDetailShape(batch_norm, i));
}
common::AnfAlgo::SetOutputInferTypeAndShape(outputs_type, outputs_shape, fused_batch_norm_with_relu.get());
common::AnfAlgo::SetOutputTypeAndDetailShape(outputs_type, outputs_shape, fused_batch_norm_with_relu.get());
common::AnfAlgo::CopyNodeAttrs(batch_norm, fused_batch_norm_with_relu);
auto manager = graph->manager();

View File

@ -96,13 +96,13 @@ const AnfNodePtr BatchNormReluGradFusion::Process(const FuncGraphPtr &graph, con
MS_EXCEPTION_IF_NULL(fused_batch_norm_grad_with_relu);
std::vector<TypeId> outputs_type;
std::vector<std::vector<size_t>> outputs_shape;
std::vector<BaseShapePtr> outputs_shape;
auto output_num = common::AnfAlgo::GetOutputTensorNum(node);
for (size_t i = 0; i < output_num; i++) {
outputs_type.push_back(common::AnfAlgo::GetOutputInferDataType(node, i));
outputs_shape.push_back(common::AnfAlgo::GetOutputInferShape(node, i));
outputs_shape.push_back(common::AnfAlgo::GetOutputDetailShape(node, i));
}
common::AnfAlgo::SetOutputInferTypeAndShape(outputs_type, outputs_shape, fused_batch_norm_grad_with_relu.get());
common::AnfAlgo::SetOutputTypeAndDetailShape(outputs_type, outputs_shape, fused_batch_norm_grad_with_relu.get());
common::AnfAlgo::CopyNodeAttrs(node, fused_batch_norm_grad_with_relu);
device::gpu::SetKernelInfo(fused_batch_norm_grad_with_relu);
return fused_batch_norm_grad_with_relu;

View File

@ -39,8 +39,8 @@ AnfNodePtr AddReduceNode(const FuncGraphPtr &func_graph, const AnfNodePtr &node)
MS_EXCEPTION_IF_NULL(new_cnode);
auto predict_input = cnode->inputs()[1];
auto new_node_dtype = {common::AnfAlgo::GetOutputInferDataType(predict_input, 0)};
auto new_node_shape = {common::AnfAlgo::GetOutputInferShape(predict_input, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(new_node_dtype, new_node_shape, new_cnode.get());
auto new_node_shape = {common::AnfAlgo::GetOutputDetailShape(predict_input, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(new_node_dtype, new_node_shape, new_cnode.get());
// Add reduce node
string reduction = common::AnfAlgo::GetNodeAttr<std::string>(node, kAttrReduction);
@ -57,8 +57,8 @@ AnfNodePtr AddReduceNode(const FuncGraphPtr &func_graph, const AnfNodePtr &node)
auto reduce_node = func_graph->NewCNode(reduce_inputs);
MS_EXCEPTION_IF_NULL(reduce_node);
auto type = common::AnfAlgo::GetOutputInferDataType(node, 0);
auto shape = {common::AnfAlgo::GetOutputInferShape(node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape({type}, shape, reduce_node.get());
auto shape = {common::AnfAlgo::GetOutputDetailShape(node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape({type}, shape, reduce_node.get());
common::AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue(std::vector<int64_t>{}), reduce_node);
common::AnfAlgo::SetNodeAttr("keep_dims", MakeValue(false), reduce_node);
reduce_node->set_scope(cnode->scope());

View File

@ -23,12 +23,14 @@
namespace mindspore::opt {
namespace {
using OutputInfo =
std::tuple<std::vector<TypeId>, std::vector<std::vector<size_t>>, std::vector<std::string>, std::vector<TypeId>>;
using OutputInfo = std::tuple<std::vector<TypeId>, std::vector<std::vector<size_t>>, std::vector<ShapeVector>,
std::vector<ShapeVector>, std::vector<std::string>, std::vector<TypeId>>;
OutputInfo GetNodeOutputInfo(const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(node);
std::vector<TypeId> output_infer_dtype;
std::vector<std::vector<size_t>> output_infer_shape;
std::vector<ShapeVector> output_max_shape;
std::vector<ShapeVector> output_min_shape;
std::vector<std::string> output_format;
std::vector<TypeId> output_device_dtype;
auto type_ptr = node->Type();
@ -41,11 +43,14 @@ OutputInfo GetNodeOutputInfo(const AnfNodePtr &node) {
for (size_t i = 0; i < output_num; i++) {
output_infer_dtype.emplace_back(common::AnfAlgo::GetOutputInferDataType(type_ptr, i));
output_infer_shape.emplace_back(common::AnfAlgo::GetOutputInferShape(node, shape_ptr, i));
output_min_shape.emplace_back(common::AnfAlgo::GetOutputMinShape(node, i));
output_max_shape.emplace_back(common::AnfAlgo::GetOutputMaxShape(node, i));
output_format.emplace_back(build_info->GetOutputFormat(i));
output_device_dtype.emplace_back(build_info->GetOutputDeviceType(i));
}
return {output_infer_dtype, output_infer_shape, output_format, output_device_dtype};
return {output_infer_dtype, output_infer_shape, output_min_shape,
output_max_shape, output_format, output_device_dtype};
}
kernel::KernelBuildInfoPtr GenerateKernelBuildInfo(const AnfNodePtr &concat, const OutputInfo &allgather_output_info,
@ -59,8 +64,8 @@ kernel::KernelBuildInfoPtr GenerateKernelBuildInfo(const AnfNodePtr &concat, con
size_t concat_input_num = common::AnfAlgo::GetInputTensorNum(concat);
for (size_t i = 0; i < concat_input_num; ++i) {
size_t input_index = allgather_input_idx + i * allgather_input_num;
inputs_device_format.emplace_back(std::get<2>(allgather_output_info)[input_index]);
inputs_device_type.emplace_back(std::get<3>(allgather_output_info)[input_index]);
inputs_device_format.emplace_back(std::get<kIndex4>(allgather_output_info)[input_index]);
inputs_device_type.emplace_back(std::get<kIndex5>(allgather_output_info)[input_index]);
}
// Current only support default format & float16
auto cmp_format = inputs_device_format.begin();
@ -101,10 +106,21 @@ AnfNodePtr InsertConcatForOutput(const FuncGraphPtr &func_graph, const AnfNodePt
MS_EXCEPTION_IF_NULL(concat);
MS_EXCEPTION_IF_NULL(new_tuple_getitems[i]);
const std::vector<TypeId> &dtypes = {std::get<0>(output_info)[i]};
const auto &shape = std::get<1>(output_info)[i];
std::vector<std::vector<size_t>> shapes = {shape};
shapes[0][0] *= rank_size;
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, concat.get());
auto shape = std::get<1>(output_info)[i];
shape[0] *= LongToSize(rank_size);
if (AnfUtils::IsShapeDynamic(shape)) {
ShapeVector tensor_shape;
auto min_shape = std::get<kIndex2>(output_info)[i];
auto max_shape = std::get<kIndex3>(output_info)[i];
max_shape[0] *= rank_size;
min_shape[0] *= rank_size;
std::transform(shape.begin(), shape.end(), std::back_inserter(tensor_shape), SizeToLong);
BaseShapePtr base_shape = std::make_shared<abstract::Shape>(tensor_shape, min_shape, max_shape);
common::AnfAlgo::SetOutputTypeAndDetailShape(dtypes, {base_shape}, concat.get());
} else {
common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, {shape}, concat.get());
}
common::AnfAlgo::SetNodeAttr(kAttrAxis, MakeValue(static_cast<int64_t>(0)), concat);
common::AnfAlgo::SetNodeAttr(kAttrInputNums, MakeValue(rank_size), concat);
std::vector<int64_t> dyn_input_size{rank_size};
@ -154,8 +170,8 @@ const AnfNodePtr ConcatOutputsForAllGather::Process(const FuncGraphPtr &func_gra
idx->set_abstract(abstract_scalar);
auto tuple_getitem = func_graph->NewCNode({NewValueNode(prim::kPrimTupleGetItem), node, idx});
MS_EXCEPTION_IF_NULL(tuple_getitem);
common::AnfAlgo::SetOutputInferTypeAndShape({std::get<0>(output_info)[i]}, {std::get<1>(output_info)[i]},
tuple_getitem.get());
auto shape = common::AnfAlgo::GetOutputDetailShape(node, i);
common::AnfAlgo::SetOutputTypeAndDetailShape({std::get<0>(output_info)[i]}, {shape}, tuple_getitem.get());
new_outputs.emplace_back(std::move(tuple_getitem));
}
return InsertConcatForOutput(func_graph, node, output_info, new_outputs, rank_size);

View File

@ -151,12 +151,12 @@ void CopyKernelInfo(AnfNodePtr src, AnfNodePtr dst) {
AnfAlgo::SetSelectKernelBuildInfo(build_info, dst.get());
size_t output_num = common::AnfAlgo::GetOutputTensorNum(src);
std::vector<TypeId> types;
std::vector<std::vector<size_t>> shapes;
std::vector<BaseShapePtr> shapes;
for (size_t i = 0; i < output_num; i++) {
types.emplace_back(common::AnfAlgo::GetOutputInferDataType(src, i));
shapes.emplace_back(common::AnfAlgo::GetOutputInferShape(src, i));
shapes.emplace_back(common::AnfAlgo::GetOutputDetailShape(src, i));
}
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, dst.get());
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, dst.get());
}
void CheckInplaceNodeInputs(std::vector<AnfNodeIndex> *inplace_node, size_t cover_index, const FuncGraphPtr &graph) {

View File

@ -37,8 +37,8 @@ void InsertCast(const FuncGraphPtr &graph, const AnfNodePtr &node, size_t i, con
std::vector<AnfNodePtr> inputs = {NewValueNode(prim), common::AnfAlgo::GetInputNode(utils::cast<CNodePtr>(node), i)};
auto cast = graph->NewCNode(inputs);
MS_EXCEPTION_IF_NULL(cast);
auto cast_shape = {common::AnfAlgo::GetPrevNodeOutputInferShape(node, i)};
common::AnfAlgo::SetOutputInferTypeAndShape({cast_type}, cast_shape, cast.get());
auto cast_shape = {common::AnfAlgo::GetPrevNodeOutputDetailShape(node, i)};
common::AnfAlgo::SetOutputTypeAndDetailShape({cast_type}, cast_shape, cast.get());
FuncGraphManagerPtr manager = graph->manager();
MS_EXCEPTION_IF_NULL(manager);
manager->SetEdge(node, i + 1, cast);
@ -107,12 +107,12 @@ bool InsertCastGPU::Run(const FuncGraphPtr &graph) {
if (IsCasted) {
auto output_types = std::vector<TypeId>(output_num, kNumberTypeFloat32);
std::vector<std::vector<size_t>> output_shapes;
std::vector<BaseShapePtr> output_shapes;
for (size_t output_index = 0; output_index < output_num; ++output_index) {
std::vector<size_t> shape = common::AnfAlgo::GetOutputInferShape(node, output_index);
auto shape = common::AnfAlgo::GetOutputDetailShape(node, output_index);
(void)output_shapes.emplace_back(shape);
}
common::AnfAlgo::SetOutputInferTypeAndShape(output_types, output_shapes, node.get());
common::AnfAlgo::SetOutputTypeAndDetailShape(output_types, output_shapes, node.get());
}
}
return true;

View File

@ -102,7 +102,8 @@ CNodePtr InsertTransposeOp(const FuncGraphPtr &graph, const AnfNodePtr &node, co
// 3.Set the output info of transpose.
auto transpose_type = {common::AnfAlgo::GetPrevNodeOutputInferDataType(used_node, used_node_index)};
auto transpose_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(used_node, used_node_index);
common::AnfAlgo::SetOutputInferTypeAndShape(transpose_type, {transpose_shape}, transpose_op.get());
auto base_shape = common::AnfAlgo::GetPrevNodeOutputDetailShape(used_node, used_node_index);
common::AnfAlgo::SetOutputTypeAndDetailShape(transpose_type, {base_shape}, transpose_op.get());
if (is_fake) {
std::vector<int64_t> shape;
std::transform(transpose_shape.begin(), transpose_shape.end(), std::back_inserter(shape), SizeToLong);

View File

@ -99,8 +99,8 @@ const AnfNodePtr MatMulBiasAddFusion::Process(const FuncGraphPtr &graph, const A
// Copy Abstract and KernelBuildInfo.
auto types = {common::AnfAlgo::GetOutputInferDataType(node, 0)};
auto shapes = {common::AnfAlgo::GetOutputInferShape(node, 0)};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, fused_node.get());
auto shapes = {common::AnfAlgo::GetOutputDetailShape(node, 0)};
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, fused_node.get());
common::AnfAlgo::CopyNodeAttrs(matmul, fused_node);
fused_node->set_scope(node->scope());
auto build_info = GenerateKernelBuildInfo(fused_node);

View File

@ -84,13 +84,13 @@ const AnfNodePtr PostBatchNormAddReluFusion::Process(const FuncGraphPtr &graph,
MS_EXCEPTION_IF_NULL(fused_batch_norm_with_add_relu);
std::vector<TypeId> outputs_type;
std::vector<std::vector<size_t>> outputs_shape;
std::vector<BaseShapePtr> outputs_shape;
auto output_num = common::AnfAlgo::GetOutputTensorNum(batch_norm);
for (size_t i = 0; i < output_num; i++) {
outputs_type.push_back(common::AnfAlgo::GetOutputInferDataType(batch_norm, i));
outputs_shape.push_back(common::AnfAlgo::GetOutputInferShape(batch_norm, i));
outputs_shape.push_back(common::AnfAlgo::GetOutputDetailShape(batch_norm, i));
}
common::AnfAlgo::SetOutputInferTypeAndShape(outputs_type, outputs_shape, fused_batch_norm_with_add_relu.get());
common::AnfAlgo::SetOutputTypeAndDetailShape(outputs_type, outputs_shape, fused_batch_norm_with_add_relu.get());
common::AnfAlgo::CopyNodeAttrs(batch_norm, fused_batch_norm_with_add_relu);
auto manager = graph->manager();

View File

@ -178,13 +178,13 @@ bool PrintReduceFusion::Run(const FuncGraphPtr &graph) {
common::AnfAlgo::SetNodeAttr("value_type_pos", MakeValue<std::vector<int64_t>>(value_type_pos), print_fused);
// set output type and shape
std::vector<TypeId> types;
std::vector<std::vector<size_t>> shapes;
std::vector<BaseShapePtr> shapes;
size_t output_num = common::AnfAlgo::GetOutputTensorNum(cnode);
for (size_t i = 0; i < output_num; i++) {
types.push_back(common::AnfAlgo::GetOutputInferDataType(cnode, i));
shapes.push_back(common::AnfAlgo::GetOutputInferShape(cnode, i));
shapes.push_back(common::AnfAlgo::GetOutputDetailShape(cnode, i));
}
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, print_fused.get());
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, print_fused.get());
// add build info
auto build_info = GenerateKernelBuildInfo(print_fused);
AnfAlgo::SetSelectKernelBuildInfo(build_info, print_fused.get());

View File

@ -83,10 +83,11 @@ CNodePtr CreateReluV2(const FuncGraphPtr &graph, const CNodePtr &relu) {
auto element_num =
std::accumulate(output_shape.begin(), output_shape.end(), static_cast<size_t>(1), std::multiplies<size_t>());
std::vector<size_t> mask_shape = {(element_num + kBitPerUInt - 1) / kBitPerUInt};
auto shapes = {common::AnfAlgo::GetOutputInferShape(relu, 0), mask_shape};
std::vector<int64_t> mask_shape = {SizeToLong((element_num + kBitPerUInt - 1) / kBitPerUInt)};
std::vector<BaseShapePtr> shapes = {common::AnfAlgo::GetOutputDetailShape(relu, 0),
std::make_shared<abstract::Shape>(mask_shape)};
auto types = {common::AnfAlgo::GetOutputInferDataType(relu, 0), kNumberTypeUInt32};
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, new_node.get());
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, new_node.get());
auto build_info = GenerateKernelBuildInfo(new_node);
AnfAlgo::SetSelectKernelBuildInfo(build_info, new_node.get());
@ -106,14 +107,14 @@ CNodePtr CreateReluGradV2(const FuncGraphPtr &graph, const CNodePtr &relu_grad,
new_node->set_abstract(relu_grad->abstract());
std::vector<TypeId> types;
std::vector<std::vector<size_t>> shapes;
std::vector<BaseShapePtr> shapes;
size_t output_num = common::AnfAlgo::GetOutputTensorNum(relu_grad);
for (size_t i = 0; i < output_num; i++) {
types.push_back(common::AnfAlgo::GetOutputInferDataType(relu_grad, i));
shapes.push_back(common::AnfAlgo::GetOutputInferShape(relu_grad, i));
shapes.push_back(common::AnfAlgo::GetOutputDetailShape(relu_grad, i));
}
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, new_node.get());
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, new_node.get());
new_node->set_scope(relu_grad->scope());
auto build_info = GenerateKernelBuildInfo(new_node);

View File

@ -43,10 +43,10 @@ const AnfNodePtr ReplaceAddNFusion::Process(const FuncGraphPtr &graph, const Anf
auto add_new = graph->NewCNode(inputs);
MS_EXCEPTION_IF_NULL(add_new);
std::vector<TypeId> outputs_type;
std::vector<std::vector<size_t>> outputs_shape;
std::vector<BaseShapePtr> outputs_shape;
outputs_type.push_back(common::AnfAlgo::GetOutputInferDataType(A, 0));
outputs_shape.push_back(common::AnfAlgo::GetOutputInferShape(A, 0));
common::AnfAlgo::SetOutputInferTypeAndShape(outputs_type, outputs_shape, add_new.get());
outputs_shape.push_back(common::AnfAlgo::GetOutputDetailShape(A, 0));
common::AnfAlgo::SetOutputTypeAndDetailShape(outputs_type, outputs_shape, add_new.get());
auto manager = graph->manager();
MS_EXCEPTION_IF_NULL(manager);
manager->Replace(utils::cast<CNodePtr>(node), utils::cast<CNodePtr>(add_new));

View File

@ -47,15 +47,15 @@ const AnfNodePtr ReplaceMomentumCastFusion::Process(const FuncGraphPtr &graph, c
MS_EXCEPTION_IF_NULL(manager);
manager->Replace(utils::cast<CNodePtr>(grad_cast), utils::cast<CNodePtr>(grad));
std::vector<TypeId> outputs_type;
std::vector<std::vector<size_t>> outputs_shape;
std::vector<BaseShapePtr> outputs_shape;
auto output_num = common::AnfAlgo::GetOutputTensorNum(node);
for (size_t i = 0; i < output_num; i++) {
outputs_type.push_back(common::AnfAlgo::GetOutputInferDataType(node, i));
outputs_shape.push_back(common::AnfAlgo::GetOutputInferShape(node, i));
outputs_shape.push_back(common::AnfAlgo::GetOutputDetailShape(node, i));
}
outputs_type[kGradIndex] = common::AnfAlgo::GetPrevNodeOutputInferDataType(grad_cast, 0);
common::AnfAlgo::SetOutputInferTypeAndShape(outputs_type, outputs_shape, node.get());
common::AnfAlgo::SetOutputTypeAndDetailShape(outputs_type, outputs_shape, node.get());
return node;
}

View File

@ -36,13 +36,13 @@ namespace opt {
namespace {
void CopyGraphOutputTypeAndShape(const std::vector<session::KernelWithIndex> &graph_outputs, CNodePtr trt_node) {
std::vector<TypeId> types;
std::vector<std::vector<size_t>> shapes;
std::vector<BaseShapePtr> shapes;
for (const auto &item : graph_outputs) {
types.push_back(common::AnfAlgo::GetOutputInferDataType(item.first, item.second));
shapes.push_back(common::AnfAlgo::GetOutputInferShape(item.first, item.second));
shapes.push_back(common::AnfAlgo::GetOutputDetailShape(item.first, item.second));
}
common::AnfAlgo::SetOutputInferTypeAndShape(types, shapes, trt_node.get());
common::AnfAlgo::SetOutputTypeAndDetailShape(types, shapes, trt_node.get());
return;
}