[AutoParallel] Add op_role_guard for PIR (#68796)

JZ-LIANG · winter-wang · web-flow · commit 9d5cd38f0278 · 2024-10-22T11:00:33.000+08:00
* c++ end

* python api

* unitest

* style

* optimize graph compile time

* code style

* style

* revise name

* revise name

* update style

* unitest cmake

---------

Co-authored-by: winter-wang &lt;1030748926@qq.com&gt;
diff --git a/paddle/fluid/pir/dialect/operator/ir/api_builder.h b/paddle/fluid/pir/dialect/operator/ir/api_builder.h
@@ -75,6 +75,9 @@ class ApiBuilder {
   // pop the insertion point and set it to the current insertion point.
   void LoadInsertionPoint();
 
+  void SetOpRole(int op_role) { builder_->set_op_role(op_role); }
+  int GetOpRole() const { return builder_->op_role(); }
+
  private:
   ApiBuilder();
 
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -2132,6 +2132,9 @@ void BindUtils(pybind11::module *m) {
          []() { ApiBuilder::Instance().ResetInsertionPointToStart(); });
   m->def("reset_insertion_point_to_end",
          []() { ApiBuilder::Instance().ResetInsertionPointToEnd(); });
+  m->def("set_op_role",
+         [](int op_role) { ApiBuilder::Instance().SetOpRole(op_role); });
+  m->def("get_op_role", []() { return ApiBuilder::Instance().GetOpRole(); });
   m->def("register_paddle_dialect", []() {
     pir::IrContext::Instance()
         ->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
diff --git a/paddle/pir/include/core/builder.h b/paddle/pir/include/core/builder.h
@@ -115,7 +115,9 @@ class Builder {
         common::errors::PreconditionNotMet("argument of block is nullptr"));
     set_insertion_point(block, block->end());
   }
-
+  /// Set/Get the op_role
+  void set_op_role(int op_role) { op_role_ = op_role; }
+  int op_role() const { return op_role_; }
   IrContext *ir_context() const { return context_; }
 
   Block *block() const { return insertion_point_.first; }
@@ -172,6 +174,10 @@ class Builder {
   InsertionPoint insertion_point_;
 
   bool forbid_insert_without_position_;
+
+  // by now the op_role is used by autoparallel for distinguish the op in fw,
+  // bw, opt region.
+  int op_role_ = -1;
 };
 
 template <typename OpTy, typename... Args>
diff --git a/paddle/pir/src/core/builder.cc b/paddle/pir/src/core/builder.cc
@@ -23,7 +23,13 @@
 namespace pir {
 /// Create an operation given the fields represented as an OperationState.
 Operation *Builder::Build(OperationArgument &&argument) {
-  return Insert(Operation::Create(std::move(argument)));
+  Operation *op = Insert(Operation::Create(std::move(argument)));
+  // TODO(ljz): Generalize here to be a hook function in the future.
+  // we add op_role attribute only when it is not equal to -1.
+  if (op_role_ != -1) {
+    op->set_attribute("op_role", Int32Attribute::get(context_, op_role_));
+  }
+  return op;
 }
 
 /// Creates an operation with the given fields.
diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py
@@ -8461,3 +8461,19 @@ def set_op_roles(block, op_role, always_forward_ops):
         if paddle.framework.in_pir_mode() and is_dist_block(block):
             always_forward_ops = ["pd_op.data", "builtin.parameter"]
             set_op_roles(block, op_role, always_forward_ops)
+
+
+# set op when op_role when it is add by apibuilder
+# pir_op_role_guard could not distinguish "always_forward_ops", therefore if
+# there would be always_forward_ops in your region, you should use "auto_complete_op_role"
+@signature_safe_contextmanager
+def pir_op_role_guard(op_role: int - 1) -> Generator[None, None, None]:
+
+    if paddle.framework.in_pir_mode():
+        original_op_rope = pir.get_op_role()
+        pir.set_op_role(op_role)
+    try:
+        yield
+    finally:
+        if paddle.framework.in_pir_mode():
+            pir.set_op_role(original_op_rope)
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -20,7 +20,7 @@
 import paddle
 from paddle import pir
 from paddle.autograd.backward_utils import ValueDict
-from paddle.base.framework import auto_complete_op_role
+from paddle.base.framework import pir_op_role_guard
 from paddle.base.log_helper import get_logger
 from paddle.distributed.fleet.meta_optimizers.common import OpRole
 from paddle.distributed.passes.pass_base import PassContext, new_pass
@@ -65,7 +65,7 @@ def reshard_single_value(program, op, operand, attr):
     if prev_var.is_dist() and prev_var.dist_attr() != attr:
         operand_attr = attr.as_tensor_dist_attr()
         paddle.pir.set_insertion_point(op)
-        with auto_complete_op_role(program, op.op_role):
+        with pir_op_role_guard(op.op_role):
             # fold reshard
             if prev_var.get_defining_op().name() == 'dist_op.reshard':
                 prev_reshard = prev_var.get_defining_op()
@@ -101,7 +101,7 @@ def reshard_combine_value(program, op, operand, attr):
         )
 
     paddle.pir.set_insertion_point(op)
-    with auto_complete_op_role(program, op.op_role):
+    with pir_op_role_guard(op.op_role):
         combine_value = paddle._C_ops.builtin_combine(reshard_vars)
     return combine_value
 
@@ -136,7 +136,7 @@ def apply_partition_pass(program):
 
             # reshard input
             paddle.pir.set_insertion_point(op)
-            with auto_complete_op_role(program, ref_op_role):
+            with pir_op_role_guard(ref_op_role):
                 reshard_var = paddle._C_ops.reshard_v2(prev_var, operand_attr)
                 operand.set_source(reshard_var)
 
@@ -151,7 +151,7 @@ def apply_partition_pass(program):
             old_dist_attr = result.dist_attr()
             result.update_dist_attr(result_attr)
 
-            with auto_complete_op_role(program, ref_op_role):
+            with pir_op_role_guard(ref_op_role):
                 prev_op = prev_var.get_defining_op()
 
                 # reshard output to assign out input
@@ -171,7 +171,7 @@ def apply_partition_pass(program):
 
             reshard_var_2 = reshard_var_1
             if old_dist_attr != reshard_var_1.dist_attr():
-                with auto_complete_op_role(program, ref_op_role):
+                with pir_op_role_guard(ref_op_role):
                     reshard_var_2 = paddle._C_ops.reshard_v2(
                         result, old_dist_attr
                     )
@@ -201,7 +201,7 @@ def apply_partition_pass(program):
                 var.update_dist_attr(attr.as_tensor_dist_attr())
 
                 # insert reshard
-                with auto_complete_op_role(program, op.op_role):
+                with pir_op_role_guard(op.op_role):
                     reshard_var = paddle._C_ops.reshard_v2(var, old_dist_attr)
                     var.replace_all_uses_with(reshard_var)
                     reshard_var.get_defining_op().operand(0).set_source(var)
@@ -266,7 +266,7 @@ def reshard_op_pass(dist_program, params_grads=[]):
                 paddle.pir.set_insertion_point(op)
                 ref_op_role = op.op_role
 
-                with auto_complete_op_role(dist_program, ref_op_role):
+                with pir_op_role_guard(ref_op_role):
                     out_value = reshard_func.reshard(
                         src_dist_attr,
                         dst_dist_attr,
diff --git a/python/paddle/pir/__init__.py b/python/paddle/pir/__init__.py
@@ -26,6 +26,7 @@
     create_shaped_type,
     fake_value,
     get_current_insertion_point,
+    get_op_role,
     is_fake_value,
     parse_program,
     register_dist_dialect,
@@ -35,6 +36,7 @@
     set_insertion_point,
     set_insertion_point_after,
     set_insertion_point_to_block_end,
+    set_op_role,
     translate_to_pir,
     translate_to_pir_with_param_map,
 )
diff --git a/test/auto_parallel/pir/CMakeLists.txt b/test/auto_parallel/pir/CMakeLists.txt
@@ -5,6 +5,9 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
   py_test_modules(test_ir_dist_attr MODULES test_ir_dist_attr ENVS
                   FLAGS_enable_pir_api=1)
   py_test_modules(test_static_pir_program MODULES test_static_pir_program)
+  py_test_modules(test_op_role MODULES test_op_role)
+  set_tests_properties(test_op_role PROPERTIES ENVIRONMENT
+                                               "FLAGS_enable_pir_api=1")
   py_test_modules(test_pir_elementwise_spmd MODULES test_elementwise_spmd_rule
                   ENVS FLAGS_enable_pir_api=1)
   py_test_modules(test_pir_relu_spmd MODULES test_relu_spmd_rule ENVS
diff --git a/test/auto_parallel/pir/test_op_role.py b/test/auto_parallel/pir/test_op_role.py
@@ -0,0 +1,149 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+import paddle.distributed as dist
+from paddle.base.framework import auto_complete_op_role, pir_op_role_guard
+from paddle.distributed import Replicate, Shard
+from paddle.distributed.auto_parallel.static.mix_to_dist_pass import (
+    apply_mix2dist_pass,
+)
+from paddle.distributed.auto_parallel.static.pir_pass import (
+    ReshardPasses,
+    apply_partition_pass,
+)
+
+
+class TestOpRole(unittest.TestCase):
+    def test_signle(self):
+        paddle.enable_static()
+        with paddle.pir_utils.IrGuard():
+            main_program = paddle.base.Program()
+            with paddle.base.program_guard(main_program):
+
+                # op_role = -1
+                x0 = paddle.static.data(name='x0', shape=[1, 128, 512])
+                x1 = paddle.nn.functional.relu(x0)
+                x2 = paddle.nn.functional.relu(x1)
+
+                with pir_op_role_guard(1):
+                    y0 = paddle.static.data(name='y0', shape=[1, 128, 512])
+                    y1 = paddle.nn.functional.relu(y0)
+                    z0 = paddle.add(y1, x2)
+                    z0 = z0 * 3.0
+                with pir_op_role_guard(3):
+                    z1 = paddle.nn.functional.relu(z0)
+                    z2 = paddle.add(y0, z1)
+                    z4 = paddle.split(z0, num_or_sections=[8, 100, 20], axis=1)
+
+                with pir_op_role_guard(0):
+                    z3 = paddle.add(y1, z2)
+
+                # op_role = -1
+                z4 = paddle.add(y0, z3)
+
+        # check global shape
+        std_ops = [
+            "pd_op.data:-1",
+            "pd_op.data:1",
+            "pd_op.relu:-1",
+            "pd_op.relu:-1",
+            "pd_op.relu:1",
+            "pd_op.add:1",
+            "pd_op.full:1",
+            "pd_op.scale:1",
+            "pd_op.relu:3",
+            "pd_op.add:3",
+            "pd_op.full_int_array:3",
+            "pd_op.full:3",
+            "pd_op.split:3",
+            "builtin.split:3",
+            "pd_op.add:0",
+            "pd_op.add:-1",
+        ]
+
+        cur_ops = [
+            f"{op.name()}:{op.op_role}"
+            for op in main_program.global_block().ops
+        ]
+        self.assertEqual(cur_ops, std_ops)
+
+    def test_dist(self):
+        paddle.enable_static()
+
+        mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
+
+        with paddle.pir_utils.IrGuard():
+            main_program = paddle.base.Program()
+            with paddle.base.program_guard(main_program):
+
+                with auto_complete_op_role(main_program, 0):
+                    x0 = paddle.static.data(name='x0', shape=[1, 128, 512])
+                    x0 = dist.shard_tensor(
+                        x0, mesh, [Shard(1), Replicate()], stop_gradient=False
+                    )
+                    x1 = x0 / 2.0
+
+                    with pir_op_role_guard(3):
+                        x2 = dist.reshard(x1, mesh, [Shard(2), Replicate()])
+                        with pir_op_role_guard(1):
+                            x3 = dist.reshard(
+                                x2, mesh, [Replicate(), Replicate()]
+                            )
+                        x4 = dist.reshard(x3, mesh, [Shard(1), Replicate()])
+
+                    x5 = dist.reshard(x4, mesh, [Replicate(), Replicate()])
+
+        apply_mix2dist_pass(main_program)
+        apply_partition_pass(main_program)
+        ReshardPasses.apply_reshard_pass(main_program, [])
+
+        std_ops = [
+            'pd_op.data:0',
+            'pd_op.full:0',
+            'pd_op.scale:0',
+            'pd_op.all_gather:3',
+            'pd_op.full:3',
+            'pd_op.split_with_num:3',
+            'pd_op.full:3',
+            'pd_op.concat:3',
+            'pd_op.full_int_array:3',
+            'pd_op.full_int_array:3',
+            'pd_op.slice:3',
+            'pd_op.all_gather:1',
+            'pd_op.full:1',
+            'pd_op.split_with_num:1',
+            'pd_op.full:1',
+            'pd_op.concat:1',
+            'pd_op.full_int_array:3',
+            'pd_op.full_int_array:3',
+            'pd_op.slice:3',
+            'pd_op.all_gather:0',
+            'pd_op.full:0',
+            'pd_op.split_with_num:0',
+            'pd_op.full:0',
+            'pd_op.concat:0',
+        ]
+
+        cur_ops = [
+            f"{op.name()}:{op.op_role}"
+            for op in main_program.global_block().ops
+        ]
+        self.assertEqual(cur_ops, std_ops)
+
+
+if __name__ == "__main__":
+    unittest.main()
