Skip to content

Commit

Permalink
[NPU] Support npu kernel for shape op (#31427)
Browse files Browse the repository at this point in the history
* add shape npu

* fix

* fix
  • Loading branch information
frankwhzhang authored Mar 4, 2021
1 parent ac3d821 commit b3c88e9
Show file tree
Hide file tree
Showing 2 changed files with 118 additions and 0 deletions.
61 changes: 61 additions & 0 deletions paddle/fluid/operators/shape_op_npu.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#include <memory>
#include <string>

#include "paddle/fluid/operators/npu_op_runner.h"
#include "paddle/fluid/operators/shape_op.h"

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
using SelectedRows = framework::SelectedRows;

template <typename DeviceContext, typename T>
class ShapeNPUKernel : public framework::OpKernel<T> {
public:
void Compute(const framework::ExecutionContext& ctx) const override {
auto* in_var = ctx.InputVar("Input");
framework::DDim in_dims;
if (in_var->IsType<SelectedRows>()) {
in_dims = in_var->Get<SelectedRows>().value().dims();
} else {
in_dims = in_var->Get<LoDTensor>().dims();
}
auto* out_t = ctx.Output<Tensor>("Out");
out_t->Resize({in_dims.size()});
// to do: cpuplace?
auto out_data = out_t->mutable_data<int32_t>(platform::CPUPlace());
for (int i = 0; i < in_dims.size(); ++i) {
out_data[i] = in_dims[i];
}
}
};

} // namespace operators
} // namespace paddle

namespace ops = paddle::operators;

REGISTER_OP_NPU_KERNEL(
shape, ops::ShapeNPUKernel<paddle::platform::NPUDeviceContext, bool>,
ops::ShapeNPUKernel<paddle::platform::NPUDeviceContext, int>,
ops::ShapeNPUKernel<paddle::platform::NPUDeviceContext, int8_t>,
ops::ShapeNPUKernel<paddle::platform::NPUDeviceContext, uint8_t>,
ops::ShapeNPUKernel<paddle::platform::NPUDeviceContext, int64_t>,
ops::ShapeNPUKernel<paddle::platform::NPUDeviceContext, float>,
ops::ShapeNPUKernel<paddle::platform::NPUDeviceContext, double>);
#endif
57 changes: 57 additions & 0 deletions python/paddle/fluid/tests/unittests/npu/test_shape_op_npu.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# 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.

from __future__ import print_function

import numpy as np
import unittest
import sys
sys.path.append("..")
from op_test import OpTest
import paddle
import paddle.fluid as fluid

paddle.enable_static()
SEED = 2021


@unittest.skipIf(not paddle.is_compiled_with_npu(),
"core is not compiled with NPU")
class TestShape(OpTest):
def setUp(self):
self.set_npu()
self.op_type = "shape"
self.place = paddle.NPUPlace(0)

self.init_dtype()
np.random.seed(SEED)
x = np.random.uniform(1, 2, [5, 10]).astype(self.dtype)
out = np.array([5, 10])

self.inputs = {'Input': OpTest.np_dtype_to_fluid_dtype(x)}
self.attrs = {}
self.outputs = {'Out': out}

def set_npu(self):
self.__class__.use_npu = True

def init_dtype(self):
self.dtype = np.float32

def test_check_output(self):
self.check_output_with_place(self.place, check_dygraph=False)


if __name__ == '__main__':
unittest.main()

0 comments on commit b3c88e9

Please sign in to comment.