From f9d0076a861bb755567466078de55ec8fffd7c9c Mon Sep 17 00:00:00 2001
From: YdrMaster <ydrml@hotmail.com>
Date: Tue, 14 Feb 2023 16:44:08 +0800
Subject: [PATCH] =?UTF-8?q?opt:=20=E4=BC=98=E5=8C=96=20SliceObj=20?=
 =?UTF-8?q?=E6=9E=84=E9=80=A0=E5=99=A8=E5=AE=9E=E7=8E=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: YdrMaster <ydrml@hotmail.com>
---
 include/operators/slice.h |  6 +++---
 src/operators/slice.cc    | 31 +++++++++++++++----------------
 2 files changed, 18 insertions(+), 19 deletions(-)

diff --git a/include/operators/slice.h b/include/operators/slice.h
index e7262552..7aeb0941 100644
--- a/include/operators/slice.h
+++ b/include/operators/slice.h
@@ -20,14 +20,14 @@ class SliceObj : public OperatorObj {
      * list which has the same length with axis.
      * @param ends The end position to slice at certain axes. `ends` is a list
      * which has the same length with axis.
-     * @param axis The dimensions to slice. If `axis` is empty, it is set to [0,
+     * @param axes The dimensions to slice. If `axis` is empty, it is set to [0,
      * 1, ..., d-1], where d is the number of dimensions of the input tensor.
      * @param steps The step to slice at certain axes. `step` is a list which
      * has the same length with axis.
      */
     SliceObj(GraphObj *graph, Tensor input, Tensor output,
              const vector<int> &starts, const vector<int> &ends,
-             const optional<vector<int>> &axis,
+             const optional<vector<int>> &axes,
              const optional<vector<int>> &steps);
     OP_CLONE(SliceObj);
 
@@ -41,4 +41,4 @@ class SliceObj : public OperatorObj {
     vector<int> getWorkloadVector() const override;
     vector<int> getOpAttrVector() const override;
 };
-} // namespace infini
\ No newline at end of file
+} // namespace infini
diff --git a/src/operators/slice.cc b/src/operators/slice.cc
index 5987531f..e5a5ec8d 100644
--- a/src/operators/slice.cc
+++ b/src/operators/slice.cc
@@ -3,34 +3,33 @@
 namespace infini {
 SliceObj::SliceObj(GraphObj *graph, Tensor input, Tensor output,
                    const vector<int> &starts, const vector<int> &ends,
-                   const optional<vector<int>> &axis,
+                   const optional<vector<int>> &axes,
                    const optional<vector<int>> &steps)
     : OperatorObj(OpType::Slice, {input}, {output}) {
-    if (steps != std::nullopt)
+    if (steps)
         IT_TODO_HALT();
     IT_ASSERT(starts.size() == ends.size());
 
-    if (axis == std::nullopt) {
+    if (!axes) {
         this->starts = starts;
         this->ends = ends;
     } else {
-        int nAxis = (*axis).size();
-        IT_ASSERT((int)starts.size() == nAxis);
+        auto nAxis = (*axes).size();
+        IT_ASSERT(starts.size() == nAxis);
 
-        int nDims = input->getDims().size();
-        vector<int> tmpS(nDims, 0), tmpE;
-        for (int i = 0; i < nDims; ++i) {
-            tmpE.emplace_back(input->getDims()[i] - 1);
-        }
+        auto dims = input->getDims();
+        this->starts = vector<int>(dims.size(), 0);
+        this->ends.resize(dims.size());
+        std::transform(dims.begin(), dims.end(), this->ends.begin(),
+                       [](auto x) { return x - 1; });
 
-        for (int i = 0; i < nAxis; ++i) {
-            if ((*axis)[i] < 0)
+        for (size_t j = 0; j < nAxis; ++j) {
+            auto i = (*axes)[j];
+            if (i < 0)
                 IT_TODO_HALT();
-            tmpS[(*axis)[i]] = starts[i];
-            tmpE[(*axis)[i]] = ends[i];
+            this->starts[i] = starts[j];
+            this->ends[i] = ends[j];
         }
-        this->starts = tmpS;
-        this->ends = tmpE;
     }
     IT_ASSERT(checkValid(graph));
 }