python demo code for setting multi-device

python/paddle/v2/fluid/framework.py

@@ -8,7 +8,8 @@
 __all__ = [
     'Block', 'Variable', 'Program', 'Operator', 'default_startup_program',
     'default_main_program', 'program_guard', 'switch_startup_program',
-    'switch_main_program'
+    'switch_main_program', 'Device', 'main_op_device_vec',
+    'start_op_device_vec', 'append_op_device_vec'
 ]
 
 
@@ -741,3 +742,40 @@ def program_guard(main_program, startup_program=None):
     switch_main_program(main_program)
     if startup_program is not None:
         switch_startup_program(startup_program)
+
+
+_device_stack_ = []
+
+
+class Device(object):
+    def __init__(self, device):
+        self.device = device
+
+    def __enter__(self):
+        _device_stack_.append(self.device)
+
+    def __exit__(self, type, value, traceback):
+        _device_stack_.pop()
+
+
+_start_op_device_vec_ = []
+_main_op_device_vec_ = []
+
+
+def start_op_device_vec():
+    return _start_op_device_vec_
+
+
+def main_op_device_vec():
+    return _main_op_device_vec_
+
+
+def append_op_device_vec(helper):
+    start_ops = helper.start_ops()
+    main_ops = helper.main_ops()
+
+    for op in start_ops:
+        _start_op_device_vec_.append(_device_stack_[-1])
+
+    for op in main_ops:
+        _main_op_device_vec_.append(_device_stack_[-1])

python/paddle/v2/fluid/layer_helper.py

@@ -15,6 +15,15 @@ def __init__(self, layer_type, **kwargs):
         if name is None:
             self.kwargs['name'] = unique_name(self.layer_type)
 
+        self.start_prog_ops = []
+        self.main_prog_ops = []
+
+    def start_ops(self):
+        return self.start_prog_ops
+
+    def main_ops(self):
+        return self.main_prog_ops
+
     @property
     def name(self):
         return self.kwargs['name']
@@ -36,7 +45,9 @@ def startup_program(self):
             return prog
 
     def append_op(self, *args, **kwargs):
-        return self.main_program.current_block().append_op(*args, **kwargs)
+        op = self.main_program.current_block().append_op(*args, **kwargs)
+        self.main_prog_ops.append(op)
+        return op
 
     def multiple_input(self, input_param_name='input'):
         inputs = self.kwargs.get(input_param_name, [])
@@ -117,8 +128,9 @@ def create_parameter(self,
         if attr.name is None:
             attr.name = unique_name(".".join([self.name, suffix]))
 
-        self.startup_program.global_block().create_parameter(
+        param = self.startup_program.global_block().create_parameter(
             dtype=dtype, shape=shape, **attr.to_kwargs(with_initializer=True))
+        self.start_prog_ops.append(param.op)
         return self.main_program.global_block().create_parameter(
             dtype=dtype, shape=shape, **attr.to_kwargs())
 
@@ -166,8 +178,9 @@ def append_bias_op(self, input_var, dim_start=1, dim_end=None):
 
         b = self.create_parameter(
             attr=bias_attr, shape=size, dtype=input_var.dtype, is_bias=True)
+        self.start_prog_ops.append(b.op)
         tmp = self.create_tmp_variable(dtype=input_var.dtype)
-        self.append_op(
+        op = self.append_op(
             type='elementwise_add',
             inputs={'X': [input_var],
                     'Y': [b]},
@@ -183,7 +196,7 @@ def append_activation(self, input_var):
             act = {'type': act}
         tmp = self.create_tmp_variable(dtype=input_var.dtype)
         act_type = act.pop('type')
-        self.append_op(
+        op = self.append_op(
             type=act_type,
             inputs={"X": [input_var]},
             outputs={"Y": [tmp]},