From 1f28968b57bf1e10a865d8cd565ef7af50145f95 Mon Sep 17 00:00:00 2001
From: Leo Chen <chenqiuliang@baidu.com>
Date: Fri, 9 Jul 2021 12:39:46 +0800
Subject: [PATCH] [NPU] Fix vector overflow in slice grad npu op (#34032)

* fix vector overflow

* refine code

* refine ut
---
 paddle/fluid/operators/slice_op_npu.cc           | 12 ++++++------
 .../tests/unittests/npu/test_slice_op_npu.py     | 16 ++++++++--------
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/paddle/fluid/operators/slice_op_npu.cc b/paddle/fluid/operators/slice_op_npu.cc
index 8e0d4b4a01..9e6e6f04ed 100644
--- a/paddle/fluid/operators/slice_op_npu.cc
+++ b/paddle/fluid/operators/slice_op_npu.cc
@@ -25,15 +25,16 @@ namespace operators {
 
 using Tensor = framework::Tensor;
 
-void UpdateAttr(const framework::DDim in_dims, const std::vector<int> axes,
+void UpdateAttr(const framework::DDim& in_dims, const std::vector<int> axes,
                 const std::vector<int> starts, const std::vector<int> ends,
                 std::vector<int>* offsets, std::vector<int>* size) {
   int cnt = 0;
   for (int i = 0; i < in_dims.size(); ++i) {
     int start = 0;
     int end = in_dims[i];
-    int axis = axes[cnt];
-
+    // NOTE(zhiqiu): Becareful that cnt may > axes.size() and result in
+    // overflow.
+    int axis = cnt < static_cast<int>(axes.size()) ? axes[cnt] : -1;
     if (axis == i) {
       start = starts[cnt];
       if (start < 0) {
@@ -63,10 +64,10 @@ class SliceNPUKernel : public framework::OpKernel<T> {
     auto axes = ctx.Attr<std::vector<int>>("axes");
     auto starts = ctx.Attr<std::vector<int>>("starts");
     auto ends = ctx.Attr<std::vector<int>>("ends");
+    const auto& in_dims = input->dims();
 
     out->mutable_data<T>(ctx.GetPlace());
 
-    auto in_dims = input->dims();
     std::vector<int> offsets(in_dims.size());
     std::vector<int> size(in_dims.size());
 
@@ -93,8 +94,7 @@ class SliceGradNPUKernel : public framework::OpKernel<T> {
     auto axes = ctx.Attr<std::vector<int>>("axes");
     auto starts = ctx.Attr<std::vector<int>>("starts");
     auto ends = ctx.Attr<std::vector<int>>("ends");
-
-    auto in_dims = input->dims();
+    const auto& in_dims = input->dims();
     int rank = in_dims.size();
 
     std::vector<int> offsets(rank);
diff --git a/python/paddle/fluid/tests/unittests/npu/test_slice_op_npu.py b/python/paddle/fluid/tests/unittests/npu/test_slice_op_npu.py
index 500618f509..c57758dca8 100644
--- a/python/paddle/fluid/tests/unittests/npu/test_slice_op_npu.py
+++ b/python/paddle/fluid/tests/unittests/npu/test_slice_op_npu.py
@@ -71,12 +71,12 @@ class TestSliceOp(OpTest):
 
 class TestSliceOp2(TestSliceOp):
     def config(self):
-        self.input = np.random.random([3, 4, 5, 6]).astype(self.dtype)
-        self.starts = [1, 0, -3]
-        self.ends = [3, 3, -1]
-        self.axes = [0, 1, 2]
-        self.infer_flags = [1, 1, 1]
-        self.out = self.input[1:3, 0:3, -3:-1, :]
+        self.input = np.random.random([10, 5, 6]).astype(self.dtype)
+        self.starts = [0]
+        self.ends = [1]
+        self.axes = [1]
+        self.infer_flags = [1]
+        self.out = self.input[:, 0:1, :]
 
 
 @unittest.skipIf(not paddle.is_compiled_with_npu(),
@@ -118,8 +118,8 @@ class TestSliceNet(unittest.TestCase):
 
             prediction = paddle.static.nn.fc(z, size=2, activation='softmax')
 
-            cost = paddle.nn.functional.cross_entropy(
-                input=prediction, label=label)
+            cost = paddle.fluid.layers.softmax_with_cross_entropy(
+                logits=prediction, label=label)
             loss = paddle.mean(cost)
             sgd = paddle.optimizer.SGD(learning_rate=0.01)
             sgd.minimize(loss)
-- 
GitLab