tensorflow/compiler/xla/service/hlo_replication_analysis.cc
// Determines whether an HLO instruction is replicated at index based on current
// knowledge in hlo_replication.
Determines whether an HLO instruction is replicated at index based on current
knowledge in hlo_replication.
Determines whether an HLO instruction is replicated at index based on current knowledge in hlo_replication .
bool DetermineHloInstructionIsReplicated(
    const HloInstruction* hlo, const ShapeIndex& index,
    bool cross_partition_spmd,
    const absl::flat_hash_map<const HloInstruction*, ShapeTree<bool>>&
        hlo_replication) {
  const auto all_operands_replicated =
      [&hlo_replication](const HloInstruction* inst) {
        for (auto operand : inst->operands()) {
          auto operand_it = hlo_replication.find(operand);
          if (operand_it == hlo_replication.end() ||
              !operand_it->second.element({})) {
            return false;
          }
        }
        return true;
      };
  if (hlo->opcode() == HloOpcode::kAllReduce ||
      hlo->opcode() == HloOpcode::kAllGather) {
    if (all_operands_replicated(hlo)) {
      return true;
    }
    if (!hlo->channel_id().has_value()) {
      if (cross_partition_spmd) {
        return false;
      }
      return hlo->replica_groups().empty() || hlo->replica_groups().size() == 1;
    } else {
      bool global_id;
      if (hlo->opcode() == HloOpcode::kAllReduce) {
        global_id = Cast<HloAllReduceInstruction>(hlo)->use_global_device_ids();
      } else {
        global_id = Cast<HloAllGatherInstruction>(hlo)->use_global_device_ids();
      }
      if (global_id) {
        bool replicated_across_partitions = true;
        bool replicated_across_replicas = true;
        const int64_t num_partitions =
            hlo->GetModule()->config().num_partitions();
        for (const auto& group : hlo->replica_groups()) {
          absl::flat_hash_set<int64_t> visited_partitions;
          absl::flat_hash_set<int64_t> visited_replicas;
          for (int64_t id : group.replica_ids()) {
            int64_t rid = id / num_partitions;
            int64_t pid = id % num_partitions;
            visited_partitions.insert(pid);
            visited_replicas.insert(rid);
          }
          replicated_across_partitions &=
              visited_partitions.size() == num_partitions;
          replicated_across_replicas &=
              visited_replicas.size() ==
              hlo->GetModule()->config().replica_count();
        }
        return cross_partition_spmd ? replicated_across_partitions
                                    : replicated_across_replicas;
      }
      return cross_partition_spmd ? true
                                  : hlo->replica_groups().empty() ||
                                        hlo->replica_groups().size() == 1;
    }
  }
  if (hlo->HasSideEffectNoRecurse()) {
    return false;
  }
  if (hlo->opcode() == HloOpcode::kReplicaId) {
    return cross_partition_spmd;
  }
  if (hlo->opcode() == HloOpcode::kPartitionId) {
    return !cross_partition_spmd;
  }
  auto it = hlo_replication.find(hlo);
  if (hlo->opcode() == HloOpcode::kParameter) {
    return it != hlo_replication.end() && it->second.element(index);
  }
  if (it != hlo_replication.end() && !it->second.element(index)) {
    return false;
  }
  if (hlo->opcode() == HloOpcode::kConstant) {
    return true;
  }
  if (hlo->opcode() == HloOpcode::kCustomCall &&
      (hlo->custom_call_target() == "X64SplitLow" ||
       hlo->custom_call_target() == "X64SplitHigh" ||
       hlo->custom_call_target() == "X64Combine")) {
    return all_operands_replicated(hlo);
  }
  if (hlo->IsElementwise() ||                             
      hlo->opcode() == HloOpcode::kConcatenate ||         
      hlo->opcode() == HloOpcode::kConvolution ||         
      hlo->opcode() == HloOpcode::kDot ||                 
      hlo->opcode() == HloOpcode::kReduce ||              
      hlo->opcode() == HloOpcode::kBroadcast ||           
      hlo->opcode() == HloOpcode::kTranspose ||           
      hlo->opcode() == HloOpcode::kReshape ||             
      hlo->opcode() == HloOpcode::kBitcast ||             
      hlo->opcode() == HloOpcode::kReverse ||             
      hlo->opcode() == HloOpcode::kGather ||              
      hlo->opcode() == HloOpcode::kScatter ||             
      hlo->opcode() == HloOpcode::kIota ||                
      hlo->opcode() == HloOpcode::kPad ||                 
      hlo->opcode() == HloOpcode::kSlice ||               
      hlo->opcode() == HloOpcode::kDynamicSlice ||        
      hlo->opcode() == HloOpcode::kDynamicUpdateSlice ||  
      hlo->opcode() == HloOpcode::kReduceWindow ||        
      hlo->opcode() == HloOpcode::kCopy) {
    return all_operands_replicated(hlo);
  }
  return false;
}
bool DetermineHloInstructionIsReplicated ( const HloInstruction * hlo , const ShapeIndex & index , bool cross_partition_spmd , const absl :: flat_hash_map < const HloInstruction * , ShapeTree < bool > > & hlo_replication ) { const auto all_operands_replicated = [ & hlo_replication ] ( const HloInstruction * inst ) { for ( auto operand : inst -> operands ( ) ) { auto operand_it = hlo_replication . find ( operand ) ; if ( operand_it == hlo_replication . end ( ) || ! operand_it -> second . element ( { } ) ) { return false ; } } return true ; } ; if ( hlo -> opcode ( ) == HloOpcode :: kAllReduce || hlo -> opcode ( ) == HloOpcode :: kAllGather ) { if ( all_operands_replicated ( hlo ) ) { return true ; } if ( ! hlo -> channel_id ( ) . has_value ( ) ) { if ( cross_partition_spmd ) { return false ; } return hlo -> replica_groups ( ) . empty ( ) || hlo -> replica_groups ( ) . size ( ) == 1 ; } else { bool global_id ; if ( hlo -> opcode ( ) == HloOpcode :: kAllReduce ) { global_id = Cast < HloAllReduceInstruction > ( hlo ) -> use_global_device_ids ( ) ; } else { global_id = Cast < HloAllGatherInstruction > ( hlo ) -> use_global_device_ids ( ) ; } if ( global_id ) { bool replicated_across_partitions = true ; bool replicated_across_replicas = true ; const int64_t num_partitions = hlo -> GetModule ( ) -> config ( ) . num_partitions ( ) ; for ( const auto & group : hlo -> replica_groups ( ) ) { absl :: flat_hash_set < int64_t > visited_partitions ; absl :: flat_hash_set < int64_t > visited_replicas ; for ( int64_t id : group . replica_ids ( ) ) { int64_t rid = id / num_partitions ; int64_t pid = id % num_partitions ; visited_partitions . insert ( pid ) ; visited_replicas . insert ( rid ) ; } replicated_across_partitions &= visited_partitions . size ( ) == num_partitions ; replicated_across_replicas &= visited_replicas . size ( ) == hlo -> GetModule ( ) -> config ( ) . replica_count ( ) ; } return cross_partition_spmd ? replicated_across_partitions : replicated_across_replicas ; } return cross_partition_spmd ? true : hlo -> replica_groups ( ) . empty ( ) || hlo -> replica_groups ( ) . size ( ) == 1 ; } } if ( hlo -> HasSideEffectNoRecurse ( ) ) { return false ; } if ( hlo -> opcode ( ) == HloOpcode :: kReplicaId ) { return cross_partition_spmd ; } if ( hlo -> opcode ( ) == HloOpcode :: kPartitionId ) { return ! cross_partition_spmd ; } auto it = hlo_replication . find ( hlo ) ; if ( hlo -> opcode ( ) == HloOpcode :: kParameter ) { return it != hlo_replication . end ( ) && it -> second . element ( index ) ; } if ( it != hlo_replication . end ( ) && ! it -> second . element ( index ) ) { return false ; } if ( hlo -> opcode ( ) == HloOpcode :: kConstant ) { return true ; } if ( hlo -> opcode ( ) == HloOpcode :: kCustomCall && ( hlo -> custom_call_target ( ) == " " || hlo -> custom_call_target ( ) == " " || hlo -> custom_call_target ( ) == " " ) ) { return all_operands_replicated ( hlo ) ; } if ( hlo -> IsElementwise ( ) || hlo -> opcode ( ) == HloOpcode :: kConcatenate || hlo -> opcode ( ) == HloOpcode :: kConvolution || hlo -> opcode ( ) == HloOpcode :: kDot || hlo -> opcode ( ) == HloOpcode :: kReduce || hlo -> opcode ( ) == HloOpcode :: kBroadcast || hlo -> opcode ( ) == HloOpcode :: kTranspose || hlo -> opcode ( ) == HloOpcode :: kReshape || hlo -> opcode ( ) == HloOpcode :: kBitcast || hlo -> opcode ( ) == HloOpcode :: kReverse || hlo -> opcode ( ) == HloOpcode :: kGather || hlo -> opcode ( ) == HloOpcode :: kScatter || hlo -> opcode ( ) == HloOpcode :: kIota || hlo -> opcode ( ) == HloOpcode :: kPad || hlo -> opcode ( ) == HloOpcode :: kSlice || hlo -> opcode ( ) == HloOpcode :: kDynamicSlice || hlo -> opcode ( ) == HloOpcode :: kDynamicUpdateSlice || hlo -> opcode ( ) == HloOpcode :: kReduceWindow || hlo -> opcode ( ) == HloOpcode :: kCopy ) { return all_operands_replicated ( hlo ) ; } return false ; }
Determines whether an HLO instruction is replicated at index based on current
knowledge in hlo_replication.
Determines whether an HLO instruction is replicated at index based on current knowledge in hlo_replication .
// Returns true if all operands are known to be replicated. // All-reduce/all-gather returns same values across partitions/replicas as // long as its operands are replicated. // This is cross-replica-only. // Only all-reduce/all-gather across all cores are replicated, which means // there is only one subgroup. // ReplicaId returns the same value for all partitions in each replica. // PartitionId returns the same value for all replicas in each partition. // Parameters should have been processed. // The HLO is already marked as non-replicated. // // // // // // // // // // // // // // // // // //