From dd48f10c3d1897d3bd0e4e201a88d781fd2bd122 Mon Sep 17 00:00:00 2001 From: lingyunli63 Date: Mon, 28 Sep 2020 10:17:01 +0800 Subject: [PATCH] add assign ops in composite_topi --- .../_extends/graph_kernel/model/model.py | 1 + .../graph_kernel/basic_ops_fusion.cc | 2 +- .../graph_kernel/composite_ops_fusion.cc | 2 +- mindspore/ops/_op_impl/akg/gpu/__init__.py | 1 + mindspore/ops/_op_impl/akg/gpu/assign.py | 33 +++++++++++++++++++ 5 files changed, 37 insertions(+), 2 deletions(-) create mode 100644 mindspore/ops/_op_impl/akg/gpu/assign.py diff --git a/mindspore/_extends/graph_kernel/model/model.py b/mindspore/_extends/graph_kernel/model/model.py index bbf7285a48..c1eed10702 100644 --- a/mindspore/_extends/graph_kernel/model/model.py +++ b/mindspore/_extends/graph_kernel/model/model.py @@ -152,6 +152,7 @@ class PrimLib: 'ReduceMin': Prim(REDUCE), 'make_tuple': Prim(CONTROL), 'ControlDepend': Prim(CONTROL), + 'Assign': Prim(ELEMWISE), '@ReduceInit': Prim(ELEMWISE), } diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/basic_ops_fusion.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/basic_ops_fusion.cc index e5e8de50af..da7d0bc985 100644 --- a/mindspore/ccsrc/backend/optimizer/graph_kernel/basic_ops_fusion.cc +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/basic_ops_fusion.cc @@ -48,7 +48,7 @@ bool IsBasicOp(const AnfNodePtr &node, bool is_before_kernel_select) { prim::kPrimRealDiv, prim::kPrimMul, prim::kPrimMinimum, prim::kPrimMaximum, prim::kPrimLog, prim::kPrimPow, prim::kPrimSub, prim::kPrimRsqrt, prim::kPrimSqrt, prim::kPrimCast, prim::kPrimAddN, prim::kPrimEqual, prim::kPrimReciprocal, prim::KPrimTransData, prim::kPrimSelect, - prim::kPrimGreater}; + prim::kPrimGreater, prim::kPrimAssign}; #else std::vector fusable_basic_ops; #endif diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/composite_ops_fusion.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/composite_ops_fusion.cc index f27f4de1fc..97741f455c 100644 --- a/mindspore/ccsrc/backend/optimizer/graph_kernel/composite_ops_fusion.cc +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/composite_ops_fusion.cc @@ -54,7 +54,7 @@ bool IsBasicFuseOp(const AnfNodePtr &node, bool is_before_kernel_select) { prim::kPrimRealDiv, prim::kPrimMul, prim::kPrimMinimum, prim::kPrimMaximum, prim::kPrimLog, prim::kPrimPow, prim::kPrimSub, prim::kPrimRsqrt, prim::kPrimSqrt, prim::kPrimCast, prim::kPrimAddN, prim::kPrimEqual, prim::kPrimReciprocal, prim::KPrimTransData, prim::kPrimSelect, - prim::kPrimGreater}; + prim::kPrimGreater, prim::kPrimAssign}; #else std::vector basic_ops; #endif diff --git a/mindspore/ops/_op_impl/akg/gpu/__init__.py b/mindspore/ops/_op_impl/akg/gpu/__init__.py index 91327fefa7..ef7227119d 100644 --- a/mindspore/ops/_op_impl/akg/gpu/__init__.py +++ b/mindspore/ops/_op_impl/akg/gpu/__init__.py @@ -54,5 +54,6 @@ from .sub import _sub_akg from .tile import _tile_akg from .select import _select_akg from .greater import _greater_akg +from .assign import _assign_akg # Please insert op register in lexicographical order of the filename. diff --git a/mindspore/ops/_op_impl/akg/gpu/assign.py b/mindspore/ops/_op_impl/akg/gpu/assign.py new file mode 100644 index 0000000000..c2c0cf0fc3 --- /dev/null +++ b/mindspore/ops/_op_impl/akg/gpu/assign.py @@ -0,0 +1,33 @@ +# Copyright 2020 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +"""Assign op""" +from mindspore.ops.op_info_register import op_info_register, AkgGpuRegOp, DataType as DT + +op_info = AkgGpuRegOp("Assign") \ + .fusion_type("ELEMWISE") \ + .input(0, "x") \ + .input(1, "y") \ + .output(0, "output") \ + .dtype_format(DT.F16_Default, DT.F16_Default, DT.F16_Default) \ + .dtype_format(DT.F32_Default, DT.F32_Default, DT.F32_Default) \ + .dtype_format(DT.I32_Default, DT.I32_Default, DT.I32_Default) \ + .get_op_info() + + +@op_info_register(op_info) +def _assign_akg(): + """Assign Akg register""" + return