parameters: (program, default_device="cpu")
return_statement: return origin_porgram, heter_ops, default_ops, program_block_ops
docstring: re-place sum op to fix bug for union forward backward op
docstring_summary: re-place sum op to fix bug for union forward backward op
docstring_tokens: re - place sum op to fix bug for union forward backward op
function: def find_heter_ops(program, default_device="cpu"):
      if default_device not in DEVICE_LIST:
          raise ValueError("Given device {} is not in device list {}".format(
              default_device, DEVICE_LIST))
  
      def _is_heter_op(op, current_heter_device, default_device="cpu"):
          heter_devices = list(DEVICE_LIST)
          heter_devices.remove(default_device)
          op_device = op.attr("op_device")
          op_type = op.type
          if op_device in heter_devices:
              return True
          elif op_type in COMMUNICATE_OPS_TYPE and current_heter_device != default_device:
              # for distributed communciate ops: send & recv & barrier etc.
              # Todo: need update this method
              #op._set_attr('op_device', current_heter_device)
              return True
          elif op_device == None or op_device == default_device:
              op._set_attr('op_device', default_device)
              return False
          return False
  
      def _is_same_device(op, pre_device, default_device="cpu"):
          op_device = op.attr("op_device")
          if op_device == pre_device:
              return True
          if pre_device == default_device:
              return True
          return False
  
      def _append_heter_op(op, current_heter_block_ops, heter_ops):
          op_device = op.attr("op_device")
          if op_device not in heter_ops:
              heter_ops[op_device] = {}
          current_heter_block_ops.append(op)
  
      origin_porgram = program.clone()
      block = program.global_block()
      '''
         re-place sum op to fix bug for union forward backward op
      '''
      var2idx = {}
      op_list = list(block.ops)
      op_size = len(op_list)
  
      for i in range(op_size - 1, -1, -1):
          op_list = list(block.ops)
          op = op_list[i]
          if "_grad" in op.type:
              forward_op_type = op.type.split("_grad")[0]
              if forward_op_type in SPARSE_OP_TYPE_DICT.keys() \
                  and op.attr('remote_prefetch') is True:
                  param_name = op.input(SPARSE_OP_TYPE_DICT[forward_op_type])[0]
                  if param_name in var2idx:
                      ## insert sum op & remove sum op from var2idx and origin place
                      op_list = list(block.ops)
                      sum_op = op_list[var2idx[param_name]]
                      sum_op_inputs = {
                          sum_op.input_names[0]: [
                              block.vars[input]
                              for input in sum_op.input_arg_names
                          ]
                      }
                      sum_op_outputs = {
                          sum_op.output_names[0]: [
                              block.vars[output]
                              for output in sum_op.output_arg_names
                          ]
                      }
                      block._insert_op(
                          index=i + 1,
                          type=sum_op.type,
                          inputs=sum_op_inputs,
                          outputs=sum_op_outputs,
                          attrs=sum_op.all_attrs())
                      block._remove_op(var2idx[param_name] + 1)
                      var2idx.pop(param_name)
                      for var_ in var2idx:
                          var2idx[var_] += 1
              elif forward_op_type == "elementwise_mul":
                  """
                  get output varname of pre op
  
                  """
                  output_vars_no_grad = []
                  for key in op.output_names:
                      for varname in op.output(key):
                          if varname == "@EMPTY@":
                              continue
                          if "lod_tensor_blocking_queue" in varname:
                              continue
                          output_vars_no_grad.append(varname.split("@GRAD")[0])
                  for no_grad_var in output_vars_no_grad:
                      if no_grad_var in var2idx:
                          """
                         insert sum op & remove sum op from var2idx and origin place
    
                         """
                          op_list = list(block.ops)
                          sum_op = op_list[var2idx[no_grad_var]]
                          sum_op_inputs = {
                              sum_op.input_names[0]: [
                                  block.vars[input]
                                  for input in sum_op.input_arg_names
                              ]
                          }
                          sum_op_outputs = {
                              sum_op.output_names[0]: [
                                  block.vars[output]
                                  for output in sum_op.output_arg_names
                              ]
                          }
                          block._insert_op(
                              index=i + 1,
                              type=sum_op.type,
                              inputs=sum_op_inputs,
                              outputs=sum_op_outputs,
                              attrs=sum_op.all_attrs())
                          block._remove_op(var2idx[no_grad_var] + 1)
                          var2idx.pop(no_grad_var)
                          for var_ in var2idx:
                              var2idx[var_] += 1
          else:
              if op.type == "sum":
                  var = op.output("Out")[0]
                  if "@GRAD" in var:
                      origin_var = var.split("@GRAD")[0]
                      pre_op = op_list[i - 1]
                      if "_grad" in pre_op.type:
                          forward_op_type = pre_op.type.split("_grad")[0]
                          if forward_op_type in SPARSE_OP_TYPE_DICT.keys() \
                              and pre_op.attr('remote_prefetch') is True:
                              param_name = pre_op.input(SPARSE_OP_TYPE_DICT[
                                  forward_op_type])[0]
                              if param_name == origin_var and op.attr(
                                      "op_device") == pre_op.attr("op_device"):
                                  continue
                              else:
                                  var2idx[origin_var] = i
                          elif forward_op_type == "elementwise_mul":
                              output_vars = []
                              for key in pre_op.output_names:
                                  for varname in pre_op.output(key):
                                      if varname == "@EMPTY@":
                                          continue
                                      if "lod_tensor_blocking_queue" in varname:
                                          continue
                                      output_vars.append(varname)
                              input_vars = []
                              for key in op.input_names:
                                  for varname in op.input(key):
                                      if varname == "@EMPTY@":
                                          continue
                                      if "lod_tensor_blocking_queue" in varname:
                                          continue
                                      input_vars.append(varname)
                              is_match = False
                              for varname in output_vars:
                                  if varname in input_vars:
                                      is_match = True
                                      break
                              if is_match:
                                  continue
                              else:
                                  var2idx[origin_var] = i
                      else:
                          var2idx[origin_var] = i
  
      origin_porgram = program.clone()
      block = program.global_block()
  
      program_block_ops = []
      default_ops = {default_device: {}}
      heter_ops = {}
      block_index = 0
  
      current_heter_block_ops = []
      current_default_block_ops = []
      current_heter_device = default_device
      is_heter = False
      for op in block.ops:
          if _is_heter_op(op, current_heter_device, default_device):
              # for gpu/xpu-op
              is_heter = True
  
              # for cpu-op block append
              if len(current_default_block_ops) > 1:
                  default_ops[default_device][
                      block_index] = current_default_block_ops
                  program_block_ops.append(current_default_block_ops)
                  current_default_block_ops = []
                  block_index += 1
  
              if _is_same_device(op, current_heter_device, default_device):
                  # for gpu-op, gpu-op -> gpu-op,...
                  current_heter_device = op.attr("op_device")
                  _append_heter_op(op, current_heter_block_ops, heter_ops)
              else:
                  # for gpu-op -> xpu-op, ...
                  op_device = current_heter_block_ops[0].attr("op_device")
                  heter_ops[op_device][block_index] = current_heter_block_ops
                  program_block_ops.append(current_heter_block_ops)
                  block_index += 1
                  current_heter_block_ops = []
                  current_heter_device = op.attr("op_device")
                  _append_heter_op(op, current_heter_block_ops, heter_ops)
  
          elif is_heter:
              # for gpu/xpu-op -> cpu-op
              op_device = current_heter_block_ops[0].attr("op_device")
              heter_ops[op_device][block_index] = current_heter_block_ops
              program_block_ops.append(current_heter_block_ops)
              block_index += 1
              current_heter_block_ops = []
              current_heter_device = default_device
              is_heter = False
              current_default_block_ops.append(op)
          else:
              # for cpu-op
              current_default_block_ops.append(op)
  
      if current_default_block_ops != []:
          default_ops[default_device][block_index] = current_default_block_ops
          program_block_ops.append(current_default_block_ops)
  
      if current_heter_block_ops != []:
          op_device = current_heter_block_ops[0].attr("op_device")
          heter_ops[op_device][block_index] = current_heter_block_ops
          program_block_ops.append(current_heter_block_ops)
  
      if len(heter_ops) == 0:
          warnings.warn(
              "No heterogeneous OP was found in your program , "
              " please using fluid.device_guard() to run OPs on different device.")
  
      total_heter_ops = 0
      heter_blocks = 0
      for device in heter_ops.keys():
          heter_block_dict = heter_ops[device]
          heter_blocks += len(heter_block_dict)
          for _, heter_block in heter_block_dict.items():
              total_heter_ops += len(heter_block)
      print(
          "There are {} OPs in your main_program, and contains {} heter-OPs which is made up of {} heter-blocks.".
          format(len(block.ops), total_heter_ops, heter_blocks))
  
      return origin_porgram, heter_ops, default_ops, program_block_ops
function_tokens: def find_heter_ops ( program , default_device = "cpu" ) : if default_device not in DEVICE_LIST : raise ValueError ( "Given device {} is not in device list {}" . format ( default_device , DEVICE_LIST ) ) def _is_heter_op ( op , current_heter_device , default_device = "cpu" ) : heter_devices = list ( DEVICE_LIST ) heter_devices . remove ( default_device ) op_device = op . attr ( "op_device" ) op_type = op . type if op_device in heter_devices : return True elif op_type in COMMUNICATE_OPS_TYPE and current_heter_device != default_device : # for distributed communciate ops: send & recv & barrier etc. # Todo: need update this method #op._set_attr('op_device', current_heter_device) return True elif op_device == None or op_device == default_device : op . _set_attr ( 'op_device' , default_device ) return False return False def _is_same_device ( op , pre_device , default_device = "cpu" ) : op_device = op . attr ( "op_device" ) if op_device == pre_device : return True if pre_device == default_device : return True return False def _append_heter_op ( op , current_heter_block_ops , heter_ops ) : op_device = op . attr ( "op_device" ) if op_device not in heter_ops : heter_ops [ op_device ] = { } current_heter_block_ops . append ( op ) origin_porgram = program . clone ( ) block = program . global_block ( ) var2idx = { } op_list = list ( block . ops ) op_size = len ( op_list ) for i in range ( op_size - 1 , - 1 , - 1 ) : op_list = list ( block . ops ) op = op_list [ i ] if "_grad" in op . type : forward_op_type = op . type . split ( "_grad" ) [ 0 ] if forward_op_type in SPARSE_OP_TYPE_DICT . keys ( ) and op . attr ( 'remote_prefetch' ) is True : param_name = op . input ( SPARSE_OP_TYPE_DICT [ forward_op_type ] ) [ 0 ] if param_name in var2idx : ## insert sum op & remove sum op from var2idx and origin place op_list = list ( block . ops ) sum_op = op_list [ var2idx [ param_name ] ] sum_op_inputs = { sum_op . input_names [ 0 ] : [ block . vars [ input ] for input in sum_op . input_arg_names ] } sum_op_outputs = { sum_op . output_names [ 0 ] : [ block . vars [ output ] for output in sum_op . output_arg_names ] } block . _insert_op ( index = i + 1 , type = sum_op . type , inputs = sum_op_inputs , outputs = sum_op_outputs , attrs = sum_op . all_attrs ( ) ) block . _remove_op ( var2idx [ param_name ] + 1 ) var2idx . pop ( param_name ) for var_ in var2idx : var2idx [ var_ ] += 1 elif forward_op_type == "elementwise_mul" : """
                  get output varname of pre op
  
                  """ output_vars_no_grad = [ ] for key in op . output_names : for varname in op . output ( key ) : if varname == "@EMPTY@" : continue if "lod_tensor_blocking_queue" in varname : continue output_vars_no_grad . append ( varname . split ( "@GRAD" ) [ 0 ] ) for no_grad_var in output_vars_no_grad : if no_grad_var in var2idx : """
                         insert sum op & remove sum op from var2idx and origin place
    
                         """ op_list = list ( block . ops ) sum_op = op_list [ var2idx [ no_grad_var ] ] sum_op_inputs = { sum_op . input_names [ 0 ] : [ block . vars [ input ] for input in sum_op . input_arg_names ] } sum_op_outputs = { sum_op . output_names [ 0 ] : [ block . vars [ output ] for output in sum_op . output_arg_names ] } block . _insert_op ( index = i + 1 , type = sum_op . type , inputs = sum_op_inputs , outputs = sum_op_outputs , attrs = sum_op . all_attrs ( ) ) block . _remove_op ( var2idx [ no_grad_var ] + 1 ) var2idx . pop ( no_grad_var ) for var_ in var2idx : var2idx [ var_ ] += 1 else : if op . type == "sum" : var = op . output ( "Out" ) [ 0 ] if "@GRAD" in var : origin_var = var . split ( "@GRAD" ) [ 0 ] pre_op = op_list [ i - 1 ] if "_grad" in pre_op . type : forward_op_type = pre_op . type . split ( "_grad" ) [ 0 ] if forward_op_type in SPARSE_OP_TYPE_DICT . keys ( ) and pre_op . attr ( 'remote_prefetch' ) is True : param_name = pre_op . input ( SPARSE_OP_TYPE_DICT [ forward_op_type ] ) [ 0 ] if param_name == origin_var and op . attr ( "op_device" ) == pre_op . attr ( "op_device" ) : continue else : var2idx [ origin_var ] = i elif forward_op_type == "elementwise_mul" : output_vars = [ ] for key in pre_op . output_names : for varname in pre_op . output ( key ) : if varname == "@EMPTY@" : continue if "lod_tensor_blocking_queue" in varname : continue output_vars . append ( varname ) input_vars = [ ] for key in op . input_names : for varname in op . input ( key ) : if varname == "@EMPTY@" : continue if "lod_tensor_blocking_queue" in varname : continue input_vars . append ( varname ) is_match = False for varname in output_vars : if varname in input_vars : is_match = True break if is_match : continue else : var2idx [ origin_var ] = i else : var2idx [ origin_var ] = i origin_porgram = program . clone ( ) block = program . global_block ( ) program_block_ops = [ ] default_ops = { default_device : { } } heter_ops = { } block_index = 0 current_heter_block_ops = [ ] current_default_block_ops = [ ] current_heter_device = default_device is_heter = False for op in block . ops : if _is_heter_op ( op , current_heter_device , default_device ) : # for gpu/xpu-op is_heter = True # for cpu-op block append if len ( current_default_block_ops ) > 1 : default_ops [ default_device ] [ block_index ] = current_default_block_ops program_block_ops . append ( current_default_block_ops ) current_default_block_ops = [ ] block_index += 1 if _is_same_device ( op , current_heter_device , default_device ) : # for gpu-op, gpu-op -> gpu-op,... current_heter_device = op . attr ( "op_device" ) _append_heter_op ( op , current_heter_block_ops , heter_ops ) else : # for gpu-op -> xpu-op, ... op_device = current_heter_block_ops [ 0 ] . attr ( "op_device" ) heter_ops [ op_device ] [ block_index ] = current_heter_block_ops program_block_ops . append ( current_heter_block_ops ) block_index += 1 current_heter_block_ops = [ ] current_heter_device = op . attr ( "op_device" ) _append_heter_op ( op , current_heter_block_ops , heter_ops ) elif is_heter : # for gpu/xpu-op -> cpu-op op_device = current_heter_block_ops [ 0 ] . attr ( "op_device" ) heter_ops [ op_device ] [ block_index ] = current_heter_block_ops program_block_ops . append ( current_heter_block_ops ) block_index += 1 current_heter_block_ops = [ ] current_heter_device = default_device is_heter = False current_default_block_ops . append ( op ) else : # for cpu-op current_default_block_ops . append ( op ) if current_default_block_ops != [ ] : default_ops [ default_device ] [ block_index ] = current_default_block_ops program_block_ops . append ( current_default_block_ops ) if current_heter_block_ops != [ ] : op_device = current_heter_block_ops [ 0 ] . attr ( "op_device" ) heter_ops [ op_device ] [ block_index ] = current_heter_block_ops program_block_ops . append ( current_heter_block_ops ) if len ( heter_ops ) == 0 : warnings . warn ( "No heterogeneous OP was found in your program , " " please using fluid.device_guard() to run OPs on different device." ) total_heter_ops = 0 heter_blocks = 0 for device in heter_ops . keys ( ) : heter_block_dict = heter_ops [ device ] heter_blocks += len ( heter_block_dict ) for _ , heter_block in heter_block_dict . items ( ) : total_heter_ops += len ( heter_block ) print ( "There are {} OPs in your main_program, and contains {} heter-OPs which is made up of {} heter-blocks." . format ( len ( block . ops ) , total_heter_ops , heter_blocks ) ) return origin_porgram , heter_ops , default_ops , program_block_ops