fix inference_dyn_bug of test_llm1.sh

Change-Id: Icab64e11f74ce840c0c7b7d0db38ce592d21ee4e
sophgo · Nov 22, 2024 · 20e20b3 · 20e20b3
1 parent 1b7f942
commit 20e20b3
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git 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
@@ -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) {