| @@ -152,6 +152,7 @@ class PrimLib: | |||||
| 'ReduceMin': Prim(REDUCE), | 'ReduceMin': Prim(REDUCE), | ||||
| 'make_tuple': Prim(CONTROL), | 'make_tuple': Prim(CONTROL), | ||||
| 'ControlDepend': Prim(CONTROL), | 'ControlDepend': Prim(CONTROL), | ||||
| 'Assign': Prim(ELEMWISE), | |||||
| '@ReduceInit': Prim(ELEMWISE), | '@ReduceInit': Prim(ELEMWISE), | ||||
| } | } | ||||
| @@ -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::kPrimRealDiv, prim::kPrimMul, prim::kPrimMinimum, prim::kPrimMaximum, prim::kPrimLog, | ||||
| prim::kPrimPow, prim::kPrimSub, prim::kPrimRsqrt, prim::kPrimSqrt, prim::kPrimCast, | prim::kPrimPow, prim::kPrimSub, prim::kPrimRsqrt, prim::kPrimSqrt, prim::kPrimCast, | ||||
| prim::kPrimAddN, prim::kPrimEqual, prim::kPrimReciprocal, prim::KPrimTransData, prim::kPrimSelect, | prim::kPrimAddN, prim::kPrimEqual, prim::kPrimReciprocal, prim::KPrimTransData, prim::kPrimSelect, | ||||
| prim::kPrimGreater}; | |||||
| prim::kPrimGreater, prim::kPrimAssign}; | |||||
| #else | #else | ||||
| std::vector<PrimitivePtr> fusable_basic_ops; | std::vector<PrimitivePtr> fusable_basic_ops; | ||||
| #endif | #endif | ||||
| @@ -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::kPrimRealDiv, prim::kPrimMul, prim::kPrimMinimum, prim::kPrimMaximum, prim::kPrimLog, | ||||
| prim::kPrimPow, prim::kPrimSub, prim::kPrimRsqrt, prim::kPrimSqrt, prim::kPrimCast, | prim::kPrimPow, prim::kPrimSub, prim::kPrimRsqrt, prim::kPrimSqrt, prim::kPrimCast, | ||||
| prim::kPrimAddN, prim::kPrimEqual, prim::kPrimReciprocal, prim::KPrimTransData, prim::kPrimSelect, | prim::kPrimAddN, prim::kPrimEqual, prim::kPrimReciprocal, prim::KPrimTransData, prim::kPrimSelect, | ||||
| prim::kPrimGreater}; | |||||
| prim::kPrimGreater, prim::kPrimAssign}; | |||||
| #else | #else | ||||
| std::vector<PrimitivePtr> basic_ops; | std::vector<PrimitivePtr> basic_ops; | ||||
| #endif | #endif | ||||
| @@ -54,5 +54,6 @@ from .sub import _sub_akg | |||||
| from .tile import _tile_akg | from .tile import _tile_akg | ||||
| from .select import _select_akg | from .select import _select_akg | ||||
| from .greater import _greater_akg | from .greater import _greater_akg | ||||
| from .assign import _assign_akg | |||||
| # Please insert op register in lexicographical order of the filename. | # Please insert op register in lexicographical order of the filename. | ||||
| @@ -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 | |||||