workaround wenet

Change-Id: I4870b90869ad3973b8846450bfd37693ba691b13
sophgo · Nov 4, 2024 · d19605d · d19605d
1 parent 5a602cf
commit d19605d
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 6 deletions.
diff --git a/lib/Dialect/Tpu/Transforms/Codegen/TensorLocation.cpp b/lib/Dialect/Tpu/Transforms/Codegen/TensorLocation.cpp
@@ -206,8 +206,15 @@ json::Object record_tensor(const T val_or_opd, const slice_index &slice_i,
     layout += "_group3d"; // {d * n, c, h, w}
   }
 
+  auto shape = ArrayRef(module::getShape(val));
+  json::Array shapeArray;
+  for (auto dim : shape) {
+    shapeArray.push_back(dim);
+  }
+
   return json::Object{
       {"name", name},     {"address", address}, {"memory_type", memory_type},
+      {"shape", std::move(shapeArray)},
       {"layout", layout}, {"type", type},       {"reshape", reshape},
       {"slice", slice}};
 }
@@ -240,6 +247,7 @@ json::Object record_tensor(Value v, const group_type_t group_type) {
 
   std::string sliceStr = "[...]";
   std::string name = module::getName(v).str();
+  std::vector<int64_t> orgShape;
 
   { // coreGroup slice
     Value value;
@@ -253,9 +261,9 @@ json::Object record_tensor(Value v, const group_type_t group_type) {
         value = joinOp->getResult(0);
     }
 
+    orgShape = std::vector<int64_t>(module::getShape(value ? value : v));
     if (value) {
       auto baseAddr = module::getAddress(value);
-      auto orgShape = module::getShape(value);
       auto shape = ArrayRef(v_spc.shape, v_spc.dims);
       auto fmt_bytes = BM168x::getFmtBytes((DATA_TYPE_T)v_spc.dtype);
       auto offset = ind2sub((address - baseAddr) / fmt_bytes, orgShape);
@@ -270,8 +278,14 @@ json::Object record_tensor(Value v, const group_type_t group_type) {
     }
   }
 
+  json::Array shapeArray;
+  for (auto dim : orgShape) {
+    shapeArray.push_back(dim);
+  }
+
   return json::Object{
       {"name", name},     {"address", address}, {"memory_type", memory_type},
+      {"shape", std::move(shapeArray)},
       {"layout", layout}, {"type", type},       {"reshape", ""},
       {"slice", sliceStr}};
 }

diff --git a/lib/Dialect/Tpu/Transforms/LayerGroup/CycleCalculator.cpp b/lib/Dialect/Tpu/Transforms/LayerGroup/CycleCalculator.cpp
@@ -323,16 +323,16 @@ int64_t CycleCalculator::getGroupCycle(BasicTimeStepPtr &time_step,
   });
   }
 
-  DEBUG_WITH_TYPE("mc_lg_refactor", {
-    if(num_core == 8){
+  if(num_core == 8){
+    if(!getenv("NO_MC")){
       loop_num = loop_num / num_core + (loop_num % num_core > 0);
-          DEBUG_WITH_TYPE("cycle_calc", {
+        DEBUG_WITH_TYPE("cycle_calc", {
         llvm::dbgs() << "; action = cycle_calc"
           << "; step = multi_core_refactor"
           << "; loop_num = " << loop_num << "\n";
       });
     }
-  });
+  }
 
   int64_t filling_cycle = 0, kernel_cycle = 0, draining_cycle = 0;
   int64_t total_layer_cycle = 0, total_gdma_cycle = 0;

diff --git a/lib/Dialect/Tpu/Transforms/LayerGroup/LmemAllocator.cpp b/lib/Dialect/Tpu/Transforms/LayerGroup/LmemAllocator.cpp
@@ -1369,7 +1369,16 @@ void LmemAllocator::sc_method_multi_core_v3(
                      << "; cost = " << _group_cost
                      << "\n";
       });
-
+      if (_group_cost == 96230) {
+        for (auto op : lg_info.group_ops) {
+          if (auto matmul = dyn_cast<tpu::MatMulOp>(op)) {
+            if (matmul.getHdimIsBatch()) {
+              return;
+            }
+          }
+        }
+        return;
+      }
       group_costs.push_back(_group_cost);
       shape_secs_space.push_back(shape_secs);
       return;