fix lg sort

Change-Id: I404233372c5659251a51012876e620085fda1808
sophgo · Jun 19, 2023 · b46c21b · b46c21b
1 parent 5721ed0
commit b46c21b
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/lib/Dialect/Tpu/Transforms/LayerGroup/GroupOps.cpp b/lib/Dialect/Tpu/Transforms/LayerGroup/GroupOps.cpp
@@ -165,14 +165,6 @@ void GroupOps::buildGroupOp(const LgInfo &lg_info,
   int64_t id = 0;
   for (int64_t stg = 0; stg < 3; ++stg) {
     for (size_t ts = 0; ts < time_step->get_timestep_num(); ++ts) {
-      if (stg == 1) {
-        auto cur_ts_layers = time_step->getLayers(ts);
-        for (auto op : cur_ts_layers) {
-          UpdateOpLgParam(op, tensor_infos, id++, lg_info.type);
-          op->moveAfter(current_op_);
-          current_op_ = op;
-        }
-      }
 
       auto cur_ts_tensors = time_step->getTensors(ts);
       for (auto tensor : cur_ts_tensors) {
@@ -193,6 +185,15 @@ void GroupOps::buildGroupOp(const LgInfo &lg_info,
           }
         }
       }
+
+      if (stg == 1) {
+        auto cur_ts_layers = time_step->getLayers(ts);
+        for (auto op : cur_ts_layers) {
+          UpdateOpLgParam(op, tensor_infos, id++, lg_info.type);
+          op->moveAfter(current_op_);
+          current_op_ = op;
+        }
+      }
     }
   }
   builder.setInsertionPointAfter(current_op_);