[AutoParallel-PIR] Remove program clone in Pass (#64137)

JZ-LIANG · web-flow · commit fd7c1ca933e9 · 2024-05-10T15:48:44.000+08:00
* remove clone

* revise typoes

* update unitest
diff --git a/paddle/fluid/pir/dialect/distributed/transforms/dist_to_dense_pass.cc b/paddle/fluid/pir/dialect/distributed/transforms/dist_to_dense_pass.cc
@@ -109,27 +109,27 @@ void VerifyDenseBlock(pir::Block* block) {
   }
 }
 
-std::shared_ptr<pir::Program> DistToDensePass(pir::Program* prog) {
+void DistToDensePass(pir::Program* prog) {
   if (FLAGS_print_ir) {
     VLOG(0) << "IR before DistToDense Pass = " << *prog;
   }
 
   pir::IrMapping mapper;
-  auto new_prog = prog->Clone(mapper);
+  // auto new_prog = prog->Clone(mapper);
 
   pir::IrContext* ctx = pir::IrContext::Instance();
   ctx->GetOrRegisterDialect<OperatorDialect>();
   ctx->GetOrRegisterDialect<DistDialect>();
 
-  ProcessDistBlock(new_prog->block());
-  VLOG(6) << "IR before VerifyDenseBlock Pass = " << *new_prog;
-  VerifyDenseBlock(new_prog->block());
+  ProcessDistBlock(prog->block());
+  VLOG(6) << "IR before VerifyDenseBlock Pass = " << *prog;
+  VerifyDenseBlock(prog->block());
 
   if (FLAGS_print_ir) {
-    VLOG(0) << "IR after DistToDense Pass = " << *new_prog;
+    VLOG(0) << "IR after DistToDense Pass = " << *prog;
   }
 
-  return new_prog;
+  // return prog;
 }
 
 }  // namespace dialect
diff --git a/paddle/fluid/pir/dialect/distributed/transforms/dist_to_dense_pass.h b/paddle/fluid/pir/dialect/distributed/transforms/dist_to_dense_pass.h
@@ -18,7 +18,7 @@
 namespace paddle {
 namespace dialect {
 
-TEST_API std::shared_ptr<pir::Program> DistToDensePass(pir::Program* prog);
+TEST_API void DistToDensePass(pir::Program* prog);
 
 void ProcessDistBlock(pir::Block* block);
 
diff --git a/paddle/fluid/pir/dialect/distributed/transforms/mix_to_dist_pass.cc b/paddle/fluid/pir/dialect/distributed/transforms/mix_to_dist_pass.cc
@@ -132,26 +132,26 @@ void VerifyDistBlock(pir::Block* block) {
   }
 }
 
-std::shared_ptr<pir::Program> MixToDistPass(pir::Program* prog) {
+void MixToDistPass(pir::Program* prog) {
   if (FLAGS_print_ir) {
-    std::cout << "IR before MixToDist Pass = " << *prog << std::endl;
+    VLOG(0) << "IR before MixToDist Pass = " << *prog << std::endl;
   }
 
   pir::IrMapping mapper;
-  auto new_prog = prog->Clone(mapper);
+  // auto new_prog = prog->Clone(mapper);
 
   pir::IrContext* ctx = pir::IrContext::Instance();
   ctx->GetOrRegisterDialect<OperatorDialect>();
   ctx->GetOrRegisterDialect<DistDialect>();
 
-  ProcessMixBlock(new_prog->block());
-  VerifyDistBlock(new_prog->block());
+  ProcessMixBlock(prog->block());
+  VerifyDistBlock(prog->block());
 
   if (FLAGS_print_ir) {
-    std::cout << "IR after MixToDist Pass = " << *new_prog << std::endl;
+    VLOG(0) << "IR after MixToDist Pass = " << *prog << std::endl;
   }
 
-  return new_prog;
+  // return prog;
 }
 
 }  // namespace dialect
diff --git a/paddle/fluid/pir/dialect/distributed/transforms/mix_to_dist_pass.h b/paddle/fluid/pir/dialect/distributed/transforms/mix_to_dist_pass.h
@@ -20,7 +20,7 @@ namespace dialect {
 
 // pir::Type ConvertOpTypeToKernelType(pir::Type op_type);
 
-TEST_API std::shared_ptr<pir::Program> MixToDistPass(pir::Program* prog);
+TEST_API void MixToDistPass(pir::Program* prog);
 
 void ProcessMixBlock(pir::Block* block);
 
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -632,9 +632,10 @@ def _parallel_pir(self, mode):
         # Part 1: Complete program
         # Step 1.1: Mix2Dense Pass
         # TODO(JZ-LIANG) regulization pass with pass management.
-        dist_program = paddle.base.libpaddle.pir.apply_mix2dist_pass(
-            mix_fw_program
-        )
+        dist_program = mix_fw_program
+        self._fwd_main_progs[mode] = mix_fw_program.clone()
+        paddle.base.libpaddle.pir.apply_mix2dist_pass(dist_program)
+
         # Step 1.2: pir backward
         if mode == "train" and self._loss and self._optimizer:
             loss = dist_program.get_output_value_by_name(self._loss_names[0])
@@ -699,9 +700,9 @@ def _parallel_pir(self, mode):
 
         # TODO(JZ-LIANG) Step 4.4 Dist2Dense Pass
         # NOTE All optimization pass that need dist_attr info should be called before Dist2Dense Pass.
-        dense_program = paddle.base.libpaddle.pir.apply_dist2dense_pass(
-            dist_program
-        )
+        dense_program = dist_program
+        dist_program = dist_program.clone()
+        paddle.base.libpaddle.pir.apply_dist2dense_pass(dense_program)
 
         self._pir_dense_main_progs[mode] = dense_program
         self._pir_dist_main_progs[mode] = dist_program
diff --git a/test/auto_parallel/pir/test_learning_rate.py b/test/auto_parallel/pir/test_learning_rate.py
@@ -101,9 +101,7 @@ def test_copy_between_mesh(self):
         engine = dist_model._engine
         engine._build("train")
         dist_program = engine._fwd_main_progs["train"]
-        dist_program = paddle.base.libpaddle.pir.apply_mix2dist_pass(
-            dist_program
-        )
+        paddle.base.libpaddle.pir.apply_mix2dist_pass(dist_program)
         loss = dist_program.get_output_value_by_name(engine._loss_names[0])
         with paddle.static.program_guard(dist_program):
             params_grads = paddle.autograd.ir_backward.append_backward(loss)
diff --git a/test/auto_parallel/pir/test_reshard.py b/test/auto_parallel/pir/test_reshard.py
@@ -53,9 +53,7 @@ def test_to_static_program(self):
         engine = dist_model._engine
         engine._build("train")
         dist_program = engine._fwd_main_progs["train"]
-        dist_program = paddle.base.libpaddle.pir.apply_mix2dist_pass(
-            dist_program
-        )
+        paddle.base.libpaddle.pir.apply_mix2dist_pass(dist_program)
         loss = dist_program.get_output_value_by_name(engine._loss_names[0])
         with paddle.static.program_guard(dist_program):
             params_grads = paddle.autograd.ir_backward.append_backward(loss)
diff --git a/test/auto_parallel/test_pir_mix2dist_pass.py b/test/auto_parallel/test_pir_mix2dist_pass.py
@@ -41,10 +41,7 @@ def test_build_api(self):
                     initializer=paddle.nn.initializer.Uniform(),
                 )
 
-        dist_program = paddle.base.libpaddle.pir.apply_mix2dist_pass(
-            main_program
-        )
-        print(dist_program)
+        paddle.base.libpaddle.pir.apply_mix2dist_pass(main_program)
 
 
 if __name__ == "__main__":
diff --git a/test/cpp/pir/distributed/dist_dialect_test.cc b/test/cpp/pir/distributed/dist_dialect_test.cc
@@ -591,9 +591,8 @@ TEST(mix_to_dist_pass_test, base) {
             (uint32_t)1);
 
   // Apply Pass
-  std::shared_ptr<pir::Program> new_program =
-      paddle::dialect::MixToDistPass(&program);
-  pir::Block* new_block = new_program->block();
+  paddle::dialect::MixToDistPass(&program);
+  pir::Block* new_block = program.block();
   EXPECT_EQ(2, static_cast<int>(new_block->num_ops()));
   std::vector<pir::Operation*> ops;
   for (auto& op : *new_block) {

Original file line number	Diff line number	Diff line change
`@@ -109,27 +109,27 @@ void VerifyDenseBlock(pir::Block* block) {`
`109`	`109`	`}`
`110`	`110`	`}`
`111`	`111`
`112`		`-std::shared_ptr<pir::Program> DistToDensePass(pir::Program* prog) {`
	`112`	`+void DistToDensePass(pir::Program* prog) {`
`113`	`113`	`if (FLAGS_print_ir) {`
`114`	`114`	`VLOG(0) << "IR before DistToDense Pass = " << *prog;`
`115`	`115`	`}`
`116`	`116`
`117`	`117`	`pir::IrMapping mapper;`
`118`		`- auto new_prog = prog->Clone(mapper);`
	`118`	`+ // auto new_prog = prog->Clone(mapper);`
`119`	`119`
`120`	`120`	`pir::IrContext* ctx = pir::IrContext::Instance();`
`121`	`121`	`ctx->GetOrRegisterDialect<OperatorDialect>();`
`122`	`122`	`ctx->GetOrRegisterDialect<DistDialect>();`
`123`	`123`
`124`		`- ProcessDistBlock(new_prog->block());`
`125`		`- VLOG(6) << "IR before VerifyDenseBlock Pass = " << *new_prog;`
`126`		`- VerifyDenseBlock(new_prog->block());`
	`124`	`+ ProcessDistBlock(prog->block());`
	`125`	`+ VLOG(6) << "IR before VerifyDenseBlock Pass = " << *prog;`
	`126`	`+ VerifyDenseBlock(prog->block());`
`127`	`127`
`128`	`128`	`if (FLAGS_print_ir) {`
`129`		`- VLOG(0) << "IR after DistToDense Pass = " << *new_prog;`
	`129`	`+ VLOG(0) << "IR after DistToDense Pass = " << *prog;`
`130`	`130`	`}`
`131`	`131`
`132`		`- return new_prog;`
	`132`	`+ // return prog;`
`133`	`133`	`}`
`134`	`134`
`135`	`135`	`} // namespace dialect`
Original file line number	Diff line number	Diff line change
`@@ -132,26 +132,26 @@ void VerifyDistBlock(pir::Block* block) {`
`132`	`132`	`}`
`133`	`133`	`}`
`134`	`134`
`135`		`-std::shared_ptr<pir::Program> MixToDistPass(pir::Program* prog) {`
	`135`	`+void MixToDistPass(pir::Program* prog) {`
`136`	`136`	`if (FLAGS_print_ir) {`
`137`		`- std::cout << "IR before MixToDist Pass = " << *prog << std::endl;`
	`137`	`+ VLOG(0) << "IR before MixToDist Pass = " << *prog << std::endl;`
`138`	`138`	`}`
`139`	`139`
`140`	`140`	`pir::IrMapping mapper;`
`141`		`- auto new_prog = prog->Clone(mapper);`
	`141`	`+ // auto new_prog = prog->Clone(mapper);`
`142`	`142`
`143`	`143`	`pir::IrContext* ctx = pir::IrContext::Instance();`
`144`	`144`	`ctx->GetOrRegisterDialect<OperatorDialect>();`
`145`	`145`	`ctx->GetOrRegisterDialect<DistDialect>();`
`146`	`146`
`147`		`- ProcessMixBlock(new_prog->block());`
`148`		`- VerifyDistBlock(new_prog->block());`
	`147`	`+ ProcessMixBlock(prog->block());`
	`148`	`+ VerifyDistBlock(prog->block());`
`149`	`149`
`150`	`150`	`if (FLAGS_print_ir) {`
`151`		`- std::cout << "IR after MixToDist Pass = " << *new_prog << std::endl;`
	`151`	`+ VLOG(0) << "IR after MixToDist Pass = " << *prog << std::endl;`
`152`	`152`	`}`
`153`	`153`
`154`		`- return new_prog;`
	`154`	`+ // return prog;`
`155`	`155`	`}`
`156`	`156`
`157`	`157`	`} // namespace dialect`
Original file line number	Diff line number	Diff line change
`@@ -41,10 +41,7 @@ def test_build_api(self):`
`41`	`41`	`initializer=paddle.nn.initializer.Uniform(),`
`42`	`42`	`)`
`43`	`43`
`44`		`- dist_program = paddle.base.libpaddle.pir.apply_mix2dist_pass(`
`45`		`- main_program`
`46`		`- )`
`47`		`- print(dist_program)`
	`44`	`+ paddle.base.libpaddle.pir.apply_mix2dist_pass(main_program)`
`48`	`45`
`49`	`46`
`50`	`47`	`if __name__ == "__main__":`