[Ascend] fuj / fix-index-error-for-ascend (#1353)

* fix index error for ascend

* fix range-based for loop in index

* add unittest for index

* Update diopi_configs.py

* fix dim check for AscendTensor squeeze

---------

Co-authored-by: zhangzefeng92 <zhang_zefeng@foxmail.com>

Author: jingguo-st

diopi_test/python/configs/diopi_configs.py

@@ -4953,6 +4953,33 @@
         ),
     ),
 
+    'index_mask': dict(
+        name=["index"],
+        interface=["CustomizedTest"],
+        # input[:, mask]
+        tensor_para=dict(
+            args=[
+                {
+                    "ins": ['input'],
+                    "shape": ((3,4,5), (4,5,6,7), (5,6,7,8,9)),
+                    "gen_fn": 'Genfunc.randn',
+                    "dtype": [np.int16, np.int32, np.int64, np.uint8, np.int8, np.bool_, np.float16, np.float32, np.float64],
+                },
+                {
+                    "ins": ['idx1'],
+                    "shape": (None, None, None),
+                    "dtype": [np.int64]
+                },
+                {
+                    "ins": ['idx2'],
+                    "shape": ((4,5), (5,6,7), (6,7,8,9)),
+                    "gen_fn": 'Genfunc.mask',
+                    "dtype": [np.bool_]
+                }
+            ]
+        )
+    ),
+
     'sgd': dict(
         name=["sgd"],
         interface=["CustomizedTest"],

impl/ascend/ascend_tensor.cpp

@@ -199,6 +199,20 @@ AscendTensor& AscendTensor::unsqueeze(int dim) {
     return *this;
 }
 
+AscendTensor& AscendTensor::squeeze(int dim) {
+    auto shape = this->shape();
+    if (shape[dim] != 1) {
+        return *this;
+    }
+    auto strides = this->stride();
+
+    shape.erase(shape.begin() + dim);
+    strides.erase(strides.begin() + dim);
+
+    this->asStrided(shape, strides);
+    return *this;
+}
+
 AscendTensor& AscendTensor::view(const std::vector<int64_t>& shape) {
     // must be contiguous
     ASCEND_CHECK_ABORT(this->isContiguous(), "now only contiguous tensor support view by shape.");

impl/ascend/ascend_tensor.hpp

@@ -244,6 +244,7 @@ class AscendTensor final {
     // Those methods may change the class attribute.
     AscendTensor& asStrided(const std::vector<int64_t>& shape, const std::vector<int64_t>& stride);
     AscendTensor& unsqueeze(int dim);
+    AscendTensor& squeeze(int dim);
     AscendTensor& view(const std::vector<int64_t>& shape);
     AscendTensor& resize(const std::vector<int64_t>& shape);
     AscendTensor& select(int64_t dim, int64_t index);

impl/ascend/functions/index.cpp

@@ -62,7 +62,7 @@ static AscendTensor nonZeroTensor(diopiContextHandle_t ctx, const AscendTensor&
 
     auto aclNZTensor = ::aclCreateTensor(
         nShape.data(), nShape.size(), aclDataType::ACL_INT64, nStride.data(), 0, aclFormat::ACL_FORMAT_ND, &numELem, 1, const_cast<void*>(nzTensor.data()));
-    DIOPI_ASCEND_CALL_ACLNN(aclnnNonzero, ctx, self, aclNZTensor);
+    DIOPI_ASCEND_CALL_ACLNN_SYNC(aclnnNonzero, ctx, self, aclNZTensor);
 
     int64_t* vDims = nullptr;
     uint64_t vDimsNum = 0;
@@ -106,8 +106,23 @@ static std::vector<AscendTensor> expandIndicesTensors(diopiContextHandle_t ctx,
                                        srcIdx);
                 }
                 AscendTensor non = nonZeroTensor(ctx, t);
-                for (int64_t j = 0; j < t.dim(); j++) {
-                    result.push_back(non.select(0, j));
+
+                auto shape = non.shape();
+                shape[0] = 1;
+                diopiSize_t size = vectorToDiopiSize(shape);
+                std::vector<diopiTensorHandle_t> nons;
+
+                for (int i = 0; i < non.shape(0); i++) {
+                    diopiTensorHandle_t tmp = nullptr;
+                    diopiRequireTensor(ctx, &tmp, &size, nullptr, diopi_dtype_int64, diopi_device);
+                    nons.push_back(tmp);
+                }
+                std::vector<int64_t> splitSize(non.shape(0), 1);
+                diopiSize_t splitSizeDiopi = vectorToDiopiSize(splitSize);
+                DIOPI_ASCEND_CALL_ACLNN(aclnnSplitWithSize, ctx, non, splitSizeDiopi, 0, nons);
+                for (const auto nj : nons) {
+                    AscendTensor njTensor(nj);
+                    result.push_back(njTensor.squeeze(0));
                 }
             } else {
                 result.push_back(t);