fix tsm concat merge prob

Change-Id: Id58b5eabc21a8f96201797eb2acbcb41eb4e2f4a
sophgo · Aug 12, 2024 · c06fb66 · c06fb66
1 parent dc3f8ef
commit c06fb66
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lib/Dialect/Tpu/Transforms/AddressAssign.cpp b/lib/Dialect/Tpu/Transforms/AddressAssign.cpp
@@ -108,13 +108,14 @@ class ConcatMergePattern : public OpRewritePattern<tpu::ConcatOp> {
       return failure();
     }
     auto relu = op.getDoRelu();
+    auto axis = op.getAxis();
     std::vector<Operation *> cat_ops;
     std::vector<Value> operands;
     bool fix = false;
     for (auto in : op.getInputs()) {
       auto concat_in = dyn_cast_or_null<tpu::ConcatOp>(in.getDefiningOp());
       if (concat_in && concat_in->hasOneUse() &&
-          concat_in.getDoRelu() == relu) {
+          concat_in.getDoRelu() == relu && concat_in.getAxis() == axis) {
         for (auto pre_in : concat_in.getInputs()) {
           operands.push_back(pre_in);
         }