add one-hop buffer reuse propogation

alibaba · Jan 22, 2024 · 60b282d · 60b282d
1 parent ddcc491
commit 60b282d
Showing 1 changed file with 44 additions and 2 deletions.
diff --git a/tao_compiler/mlir/disc/transforms/disc_input_output_alias.cc b/tao_compiler/mlir/disc/transforms/disc_input_output_alias.cc
@@ -141,17 +141,59 @@ struct DiscInputOutputAliasPass
       if (outputs[outputs_index[i]] == params[params_index[i]]) {
         continue;
       }
-      // DISC now only support one-hop buffer sharing.
+      // Inplace buffer reuse.
+      bool inplace_reuse = false;
       auto defineOp = outputs[outputs_index[i]].getDefiningOp();
       for (const auto& value : defineOp->getOperands()) {
         if (params[params_index[i]] == value) {
           builder.setInsertionPointAfterValue(outputs[outputs_index[i]]);
-          builder.create<mhlo_disc::ArgsMutationOp>(main_func.getLoc(),
+          builder.create<mhlo_disc::ArgsMutationOp>(outputs[outputs_index[i]].getLoc(),
                                                     outputs[outputs_index[i]],
                                                     params[params_index[i]]);
+          inplace_reuse = true;
           break;
         }
       }
+
+      // Try one-hop buffer sharing propogation
+      if(!inplace_reuse) {
+        OneHopBufferReusePropogation(params[params_index[i]], outputs[outputs_index[i]], builder);
+      }
+
+
+    }
+  }
+ private:
+ /*
+                  A = op(src)
+ A = op(src) =>   args_mutation(A, src)
+ B = op(A)   =>   B = op(A)
+                  args_mutation(B, A)
+ */
+  void OneHopBufferReusePropogation(Value src, Value dst, OpBuilder& builder) {
+    auto dst_op = dst.getDefiningOp();
+    auto user_begin = src.user_begin();
+    auto user_end = src.user_end();
+    auto users_cnt = std::distance(user_begin, user_end)
+
+    if(users_cnt > 1 || user_begin->getNumResults() > 1) {
+      return;
+    }
+
+    auto user_result = user_begin->getResult(0);
+    for(const auto& operand : dst_op->getOperands()) {
+      if(operand == user_result) {
+        builder.setInsertionPointAfterValue(user_result);
+        builder.create<mhlo_disc::ArgsMutationOp>(user_result.getLoc(),
+                                                  user_result,
+                                                  src);
+
+        builder.setInsertionPointAfterValue(dst);
+        builder.create<mhlo_disc::ArgsMutationOp>(dst.getLoc(),
+                                                  dst,
+                                                  user_result);
+        break;
+      }
     }
   }
 };