From f6677628c0dae754f67b741efa43bd4c744c42e0 Mon Sep 17 00:00:00 2001 From: yanzhenxiang2020 Date: Sat, 12 Dec 2020 15:29:16 +0800 Subject: [PATCH] add StackPush and StackPop for aicpu --- .../kernel_compiler/aicpu/aicpu_util.h | 8 +- mindspore/ops/_op_impl/aicpu/__init__.py | 4 + .../ops/_op_impl/aicpu/stack_push_pop.py | 87 ++++++++++++ mindspore/ops/operations/_inner_ops.py | 127 +++++++++++++++++ .../ops/ascend/test_aicpu_ops/test_stack.py | 128 ++++++++++++++++++ 5 files changed, 352 insertions(+), 2 deletions(-) create mode 100644 mindspore/ops/_op_impl/aicpu/stack_push_pop.py create mode 100644 tests/st/ops/ascend/test_aicpu_ops/test_stack.py diff --git a/mindspore/ccsrc/backend/kernel_compiler/aicpu/aicpu_util.h b/mindspore/ccsrc/backend/kernel_compiler/aicpu/aicpu_util.h index cd6a8ce9fc..ebb6b0ff43 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/aicpu/aicpu_util.h +++ b/mindspore/ccsrc/backend/kernel_compiler/aicpu/aicpu_util.h @@ -44,6 +44,10 @@ constexpr auto kSeed2 = "seed2"; constexpr auto kTopK = "TopK"; constexpr auto kTopKV2 = "TopKV2"; constexpr auto kStack = "Stack"; +constexpr auto kStackInit = "StackInit"; +constexpr auto kStackPush = "StackPush"; +constexpr auto kStackPop = "StackPop"; +constexpr auto kStackDestroy = "StackDestroy"; constexpr auto kEditDistance = "EditDistance"; constexpr auto kGatherD = "GatherD"; constexpr auto kIdentity = "Identity"; @@ -55,8 +59,8 @@ constexpr auto kCustRunApi = "RunCpuKernel"; constexpr auto kDropout2D = "Dropout2D"; constexpr auto kDropout3D = "Dropout3D"; const std::set kCustAiCpuKernelOps{kIdentity}; -const std::set kCacheKernelOps{kUpdateCache, kCacheSwapTable, kSubAndFilter, - kPadAndShift, kDropout3D, kDropout2D}; +const std::set kCacheKernelOps{kUpdateCache, kCacheSwapTable, kSubAndFilter, kPadAndShift, kDropout3D, + kDropout2D, kStackInit, kStackPush, kStackPop, kStackDestroy}; struct AicpuParamHead { uint32_t length; // Total length: include cunstom message diff --git a/mindspore/ops/_op_impl/aicpu/__init__.py b/mindspore/ops/_op_impl/aicpu/__init__.py index d5e18adf3c..e50cbe9584 100644 --- a/mindspore/ops/_op_impl/aicpu/__init__.py +++ b/mindspore/ops/_op_impl/aicpu/__init__.py @@ -70,4 +70,8 @@ from .fused_sparse_ftrl import _fused_sparse_ftrl_aicpu from .fused_sparse_proximal_adagrad import _fused_sparse_proximal_adagrad_aicpu from .meshgrid import _meshgrid_aicpu from .trans_data import _trans_data_aicpu +from .stack_push_pop import _stack_init_aicpu +from .stack_push_pop import _stack_push_aicpu +from .stack_push_pop import _stack_pop_aicpu +from .stack_push_pop import _stack_destroy_aicpu from .ctc_greedy_decoder import _ctc_greedy_decoder_aicpu diff --git a/mindspore/ops/_op_impl/aicpu/stack_push_pop.py b/mindspore/ops/_op_impl/aicpu/stack_push_pop.py new file mode 100644 index 0000000000..7916bfe801 --- /dev/null +++ b/mindspore/ops/_op_impl/aicpu/stack_push_pop.py @@ -0,0 +1,87 @@ +# 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. +# ============================================================================ + +"""StackPush and StackPop op""" +from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType + +stack_init_op_info = AiCPURegOp("StackInit") \ + .fusion_type("OPAQUE") \ + .attr("index", "int") \ + .get_op_info() + +stack_push_op_info = AiCPURegOp("StackPush") \ + .fusion_type("OPAQUE") \ + .input(0, "src", "required") \ + .attr("index", "int") \ + .dtype_format(DataType.U8_Default) \ + .dtype_format(DataType.U16_Default) \ + .dtype_format(DataType.U32_Default) \ + .dtype_format(DataType.U64_Default) \ + .dtype_format(DataType.I8_Default) \ + .dtype_format(DataType.I16_Default) \ + .dtype_format(DataType.I32_Default) \ + .dtype_format(DataType.I64_Default) \ + .dtype_format(DataType.F16_Default) \ + .dtype_format(DataType.F32_Default) \ + .dtype_format(DataType.F64_Default) \ + .dtype_format(DataType.BOOL_Default) \ + .get_op_info() + +stack_pop_op_info = AiCPURegOp("StackPop") \ + .fusion_type("OPAQUE") \ + .output(0, "dst", "required") \ + .attr("index", "int") \ + .dtype_format(DataType.U8_Default) \ + .dtype_format(DataType.U16_Default) \ + .dtype_format(DataType.U32_Default) \ + .dtype_format(DataType.U64_Default) \ + .dtype_format(DataType.I8_Default) \ + .dtype_format(DataType.I16_Default) \ + .dtype_format(DataType.I32_Default) \ + .dtype_format(DataType.I64_Default) \ + .dtype_format(DataType.F16_Default) \ + .dtype_format(DataType.F32_Default) \ + .dtype_format(DataType.F64_Default) \ + .dtype_format(DataType.BOOL_Default) \ + .get_op_info() + +stack_destroy_op_info = AiCPURegOp("StackDestroy") \ + .fusion_type("OPAQUE") \ + .attr("index", "int") \ + .get_op_info() + + +@op_info_register(stack_init_op_info) +def _stack_init_aicpu(): + """StackInit aicpu register""" + return + + +@op_info_register(stack_push_op_info) +def _stack_push_aicpu(): + """StackPush aicpu register""" + return + + +@op_info_register(stack_pop_op_info) +def _stack_pop_aicpu(): + """StackPop aicpu register""" + return + + +@op_info_register(stack_destroy_op_info) +def _stack_destroy_aicpu(): + """StackDestroy aicpu register""" + return diff --git a/mindspore/ops/operations/_inner_ops.py b/mindspore/ops/operations/_inner_ops.py index 4e663e4d06..68e18d7eee 100644 --- a/mindspore/ops/operations/_inner_ops.py +++ b/mindspore/ops/operations/_inner_ops.py @@ -15,6 +15,7 @@ """Inner operators.""" +import numpy as np from ..._checkparam import Rel from ..._checkparam import Validator as validator from ... import context @@ -882,3 +883,129 @@ class Centralization(PrimitiveWithInfer): 'dtype': x_dtype, 'value': None} return out + + +class StackInit(PrimitiveWithInfer): + """ + Create a stack that produces tensors in first-in last-out order. + + After `StackInit`, a tensor can be pushed onto the stack using `StackPush`, and popped + at the top of the stack using `StackPop`. Finally, the stack should be destroyed with `StackDestroy`. + + Args: + index (int): The index of the stack. + + Supported Platforms: + ``Ascend`` + + Examples: + >>> x = Tensor(np.array([[1, 3], [2, 0]])) + >>> index = 0 + >>> stack = ops.StackInit(index) + >>> push = ops.StackPush(index) + >>> pop = ops.StackPop(index, x.shape, x.dtype) + >>> destroy = ops.StackDestroy(index) + >>> stack() + >>> push(x) + >>> y = pop() + >>> destroy() + >>> print(y) + [[1 3] + [2 0]] + """ + @prim_attr_register + def __init__(self, index=1): + """StackInit""" + validator.check_value_type("index", index, [int], self.name) + + +class StackPush(PrimitiveWithInfer): + """ + Push a tensor onto the stack. + + Before `StackPush`, the stack should be created using `StackInit`. + Please refer to the usage in source code of `StackInit`. + + Args: + index (int): The index of the stack. + + Inputs: + - **input** (Tensor) - A tensor to be pushed onto the stack. + + Supported Platforms: + ``Ascend`` + + Examples: + Please refer to the usage of `StackInit`. + """ + @prim_attr_register + def __init__(self, index=1): + """StackPush""" + validator.check_value_type("index", index, [int], self.name) + self.init_prim_io_names(inputs=['input'], outputs=[]) + + +class StackPop(PrimitiveWithInfer): + """ + Pop the tensor at the top of the stack. + + Before `StackPop`, the stack should be created using `StackInit`. + Please refer to the usage in source code of `StackInit`. + + Args: + index (int): The index of the stack. + shape (tuple): The shape of the tensor at the top of the stack. + dtype (mindspore.dtype): The type of the tensor at the top of the stack. + + Outputs: + - **output** (Tensor) - The tensor at the top of the stack. + + Supported Platforms: + ``Ascend`` + + Examples: + Please refer to the usage of `StackInit`. + """ + @prim_attr_register + def __init__(self, index=1, shape=(1,), dtype=mstype.float32): + """StackPop""" + validator.check_value_type("index", index, [int], self.name) + + validator.check_value_type('shape type', shape, [list, tuple], self.name) + validator.check_int(len(np.array(shape).shape), 1, Rel.EQ, "dim of shape", self.name) + for elem in shape: + validator.check_int(elem, 1, Rel.GE, 'shape element', self.name) + validator.check_value_type('type of shape element', elem, [int], self.name) + + validator.check_type_name("dtype", dtype, (mstype.bool_,) + mstype.number_type, self.name) + self.shape = shape + self.dtype = dtype + + self.init_prim_io_names(inputs=[], outputs=['output']) + + def __infer__(self): + return {'shape': (list(self.shape)), + 'dtype': (self.dtype), + 'value': None} + + +class StackDestroy(PrimitiveWithInfer): + """ + Destroy the stack. + + Before `StackDestroy`, the stack should be created using `StackInit`. + Please refer to the usage in source code of `StackInit`. + + Args: + index (int): The index of the stack. + + Supported Platforms: + ``Ascend`` + + Examples: + Please refer to the usage of `StackInit`. + """ + @prim_attr_register + def __init__(self, index=1): + """StackDestroy""" + validator.check_value_type("index", index, [int], self.name) diff --git a/tests/st/ops/ascend/test_aicpu_ops/test_stack.py b/tests/st/ops/ascend/test_aicpu_ops/test_stack.py new file mode 100644 index 0000000000..035bf9348c --- /dev/null +++ b/tests/st/ops/ascend/test_aicpu_ops/test_stack.py @@ -0,0 +1,128 @@ +# 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. +# ============================================================================ +import numpy as np +import mindspore.nn as nn +import mindspore.context as context +from mindspore import Tensor +from mindspore.ops.operations import _inner_ops as P + +context.set_context(mode=context.PYNATIVE_MODE, device_target="Ascend") + + +class Net(nn.Cell): + def __init__(self, index=0, shapes_and_types=None): + super(Net, self).__init__() + shapes_and_types.reverse() + self.init = P.StackInit(index) + self.push = P.StackPush(index) + self.pop = [P.StackPop(index, shape, dtype) for (shape, dtype) in shapes_and_types] + self.destroy = P.StackDestroy(index) + + def construct(self, x1, x2, x3): + self.init() + self.push(x1) + self.push(x2) + self.push(x3) + y1 = self.pop[0]() + y2 = self.pop[1]() + y3 = self.pop[2]() + self.destroy() + return y1, y2, y3 + + +class NetTwoStack(nn.Cell): + def __init__(self, index=0, shapes_and_types=None): + super(NetTwoStack, self).__init__() + self.init_0 = P.StackInit(index) + self.push_0 = P.StackPush(index) + self.pop_0 = [P.StackPop(index, shape, dtype) for (shape, dtype) in shapes_and_types] + self.destroy_0 = P.StackDestroy(index) + + index += 1 + self.init_1 = P.StackInit(index) + self.push_1 = P.StackPush(index) + self.pop_1 = [P.StackPop(index, shape, dtype) for (shape, dtype) in shapes_and_types] + self.destroy_1 = P.StackDestroy(index) + + def construct(self, x1, x2, x3): + self.init_0() + self.init_1() + + self.push_0(x1) + self.push_1(x3) + y1 = self.pop_0[0]() + z1 = self.pop_1[2]() + self.push_0(x2) + self.push_0(x3) + self.push_1(x1) + self.push_1(x2) + y2 = self.pop_0[2]() + z2 = self.pop_1[1]() + y3 = self.pop_0[1]() + z3 = self.pop_1[0]() + + self.destroy_0() + self.destroy_1() + return y1, y2, y3, z1, z2, z3 + + +def test_net(): + x1 = Tensor(np.random.randn(4,).astype(np.float64)) + x2 = Tensor(np.random.randn(4, 6).astype(np.float32)) + x3 = Tensor(np.random.randint(100, size=(3, 4, 5)).astype(np.int32)) + + shapes_and_types = [] + shapes_and_types.append((x1.shape, x1.dtype)) + shapes_and_types.append((x2.shape, x2.dtype)) + shapes_and_types.append((x3.shape, x3.dtype)) + + net = Net(2018, shapes_and_types) + y1, y2, y3 = net(x1, x2, x3) + print(x1) + print(x2) + print(x3) + print(y1) + print(y2) + print(y3) + assert np.array_equal(y1.asnumpy(), x3.asnumpy()) + assert np.array_equal(y2.asnumpy(), x2.asnumpy()) + assert np.array_equal(y3.asnumpy(), x1.asnumpy()) + + +def test_net_tow_stack(): + x1 = Tensor(np.random.randn(4,).astype(np.float64)) + x2 = Tensor(np.random.randn(4, 6).astype(np.float32)) + x3 = Tensor(np.random.randint(100, size=(3, 4, 5)).astype(np.int32)) + + shapes_and_types = [] + shapes_and_types.append((x1.shape, x1.dtype)) + shapes_and_types.append((x2.shape, x2.dtype)) + shapes_and_types.append((x3.shape, x3.dtype)) + + net = NetTwoStack(1998, shapes_and_types) + y1, y2, y3, z1, z2, z3 = net(x1, x2, x3) + print(x1) + print(x2) + print(x3) + print(y1) + print(y2) + print(y3) + assert np.array_equal(y1.asnumpy(), x1.asnumpy()) + assert np.array_equal(y2.asnumpy(), x3.asnumpy()) + assert np.array_equal(y3.asnumpy(), x2.asnumpy()) + + assert np.array_equal(z1.asnumpy(), x3.asnumpy()) + assert np.array_equal(z2.asnumpy(), x2.asnumpy()) + assert np.array_equal(z3.asnumpy(), x1.asnumpy())