From 20e20b3ca6f46d8eb8b8ec2c894ee10ff1924834 Mon Sep 17 00:00:00 2001
From: "chang.zhao" <chang.zhao@sophgo.com>
Date: Fri, 22 Nov 2024 11:46:58 +0800
Subject: [PATCH] fix inference_dyn_bug of test_llm1.sh

Change-Id: Icab64e11f74ce840c0c7b7d0db38ce592d21ee4e
---
 lib/Dialect/Top/Interfaces/Slice.cpp        | 7 +++++--
 lib/Dialect/Tpu/Interfaces/Common/Slice.cpp | 7 +++++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/lib/Dialect/Top/Interfaces/Slice.cpp b/lib/Dialect/Top/Interfaces/Slice.cpp
index 160a240a6..82e57e5cf 100644
--- a/lib/Dialect/Top/Interfaces/Slice.cpp
+++ b/lib/Dialect/Top/Interfaces/Slice.cpp
@@ -58,6 +58,7 @@ LogicalResult top::SliceOp::inference(InferenceParameter &p) {
   auto offset_v = module::getI64Array(getOffset());
   auto steps_v = module::getI64Array(getSteps());
   std::vector<int64_t> in_shape = module::getShape(getInput());
+  auto in_dims = in_shape.size();
   if (out_num_elem == 0) {
     return success();
   }
@@ -66,11 +67,14 @@ LogicalResult top::SliceOp::inference(InferenceParameter &p) {
   auto axes = module::getI64Array(getHasparamConvertAxesAttr());
   auto slice_n = axes->size();
   auto ends_v = ends_v_old;
-  if(slice_n) {
+  if (slice_n) {
     ends_v = std::make_shared<std::vector<int64_t>>(in_shape);
   }
   for (int i = 0; i < slice_n; ++i) {
     int axis = axes->at(i);
+    if (axis < 0) {
+      axis += in_dims;
+    }
     int step = steps_v->at(axis);
     int64_t end = ends_v_old->at(axis);
     int64_t offset = offset_v->at(axis);
@@ -83,7 +87,6 @@ LogicalResult top::SliceOp::inference(InferenceParameter &p) {
       offset_v->at(i) += in_shape[i];
     }
   }
-  auto in_dims = in_shape.size();
   std::vector<int64_t> out_shape(in_dims);
   for (size_t i = 0; i < in_dims; ++i) {
     if (i < slice_dims) {
diff --git a/lib/Dialect/Tpu/Interfaces/Common/Slice.cpp b/lib/Dialect/Tpu/Interfaces/Common/Slice.cpp
index 0a8cde92f..34fb91106 100644
--- a/lib/Dialect/Tpu/Interfaces/Common/Slice.cpp
+++ b/lib/Dialect/Tpu/Interfaces/Common/Slice.cpp
@@ -140,14 +140,18 @@ LogicalResult tpu::SliceOp::inference(InferenceParameter &p) {
   auto ends_v_old = module::getI64Array(getEnds());
   const size_t slice_dims = offset_v->size();
   std::vector<int64_t> in_shape = module::getShape(getInput());
+  auto in_dims = in_shape.size();
   auto axes = module::getI64Array(getHasparamConvertAxesAttr());
   auto slice_n = axes->size();
   auto ends_v = ends_v_old;
-  if(slice_n) {
+  if (slice_n) {
     ends_v = std::make_shared<std::vector<int64_t>>(in_shape);
   }
   for (int i = 0; i < slice_n; ++i) {
     int axis = axes->at(i);
+    if (axis < 0) {
+      axis += in_dims;
+    }
     int step = steps_v->at(axis);
     int64_t end = ends_v_old->at(axis);
     int64_t offset = offset_v->at(axis);
@@ -160,7 +164,6 @@ LogicalResult tpu::SliceOp::inference(InferenceParameter &p) {
       offset_v->at(i) += in_shape[i];
     }
   }
-  auto in_dims = in_shape.size();
   std::vector<int64_t> output_shape(in_dims);
   for (size_t i = 0; i < in_dims; ++i) {
     if (i < slice_dims) {