tensorflow/compiler/xla/service/gpu/all_reduce_blueconnect.cc
// Attempts to decompose all-reduces as described by the BlueConnect paper.
//
// If possible, the all-reduce will be transformed into:
// 1. reduce-scatter
// 2. all-reduce
// 3. all-gather
//
// If the all-reduce replica groups have more than one device within the same
// host, the reduce-scatter will be performed over all devices with each host.
// Otherwise, the reduce-scatter will be performed between pairs of devices on
// different hosts.
//
// When applied repeatedly, this transformation will reproduce the same pattern
// as described in the BlueConnect paper.
Attempts to decompose all-reduces as described by the BlueConnect paper.
If possible, the all-reduce will be transformed into:
1. reduce-scatter
2. all-reduce
3. all-gather

If the all-reduce replica groups have more than one device within the same
host, the reduce-scatter will be performed over all devices with each host.
Otherwise, the reduce-scatter will be performed between pairs of devices on
different hosts.

When applied repeatedly, this transformation will reproduce the same pattern
as described in the BlueConnect paper.
Attempts to decompose all - reduces as described by the BlueConnect paper . If possible the all - reduce will be transformed into : 1 . reduce - scatter 2 . all - reduce 3 . all - gather If the all - reduce replica groups have more than one device within the same host the reduce - scatter will be performed over all devices with each host . Otherwise the reduce - scatter will be performed between pairs of devices on different hosts . When applied repeatedly this transformation will reproduce the same pattern as described in the BlueConnect paper .
StatusOr<bool> TryDecomposeAllReduce(HloAllReduceInstruction* all_reduce,
                                     size_t num_devices_per_host) {
  TF_RET_CHECK(all_reduce);
  TF_RET_CHECK(!all_reduce->has_sharding());
  HloComputation& computation = *all_reduce->parent();  
  PrimitiveType element_type = all_reduce->operand(0)->shape().element_type();
  TF_ASSIGN_OR_RETURN(
      absl::optional<DecomposedReplicaGroups> decomposed_groups,
      TryDecomposeReplicaGroups(*all_reduce, num_devices_per_host));
  if (!decomposed_groups) return false;
  std::vector<HloInstruction*> flat_operands;
  flat_operands.reserve(all_reduce->operand_count());
  std::vector<Shape> flat_shapes;
  flat_shapes.reserve(all_reduce->operand_count());
  std::vector<Shape> scattered_shapes;
  scattered_shapes.reserve(all_reduce->operand_count());
  int scatter_group_size =
      decomposed_groups->scatter_gather_groups[0].replica_ids_size();
  for (HloInstruction* operand : all_reduce->operands()) {
    TF_RET_CHECK(operand->shape().IsArray());
    int64_t num_elements = ShapeUtil::ElementsIn(operand->shape());
    Shape flat_shape = ShapeUtil::MakeShape(element_type, {num_elements});
    flat_operands.push_back(computation.AddInstruction(
        HloInstruction::CreateBitcast(flat_shape, operand)));
    flat_shapes.push_back(std::move(flat_shape));
    scattered_shapes.push_back(ShapeUtil::MakeShape(
        element_type, {num_elements / scatter_group_size}));
  }
  Shape reduce_scatter_shape = ShapeUtil::MakeMaybeTupleShape(scattered_shapes);
  HloInstruction* reduce_scatter =
      computation.AddInstruction(HloInstruction::CreateReduceScatter(
          reduce_scatter_shape, flat_operands, all_reduce->to_apply(),
          decomposed_groups->scatter_gather_groups, false,
          all_reduce->channel_id(), all_reduce->use_global_device_ids(),
          0));
  HloInstruction* new_all_reduce =
      computation.AddInstruction(HloInstruction::CreateAllReduce(
          reduce_scatter_shape, GetOutputs(*reduce_scatter),
          all_reduce->to_apply(), decomposed_groups->new_all_reduce_groups,
          false, all_reduce->channel_id(),
          all_reduce->use_global_device_ids()));
  HloInstruction* all_gather =
      computation.AddInstruction(HloInstruction::CreateAllGather(
          ShapeUtil::MakeMaybeTupleShape(flat_shapes),
          GetOutputs(*new_all_reduce),
          0, decomposed_groups->scatter_gather_groups,
          false, all_reduce->channel_id(),
          all_reduce->use_global_device_ids()));
  std::vector<HloInstruction*> outputs = GetOutputs(*all_gather);
  for (int64_t i = 0; i < outputs.size(); ++i) {
    outputs[i] = computation.AddInstruction(HloInstruction::CreateBitcast(
        all_reduce->operand(i)->shape(), outputs[i]));
  }
  TF_RETURN_IF_ERROR(computation.ReplaceInstruction(
      all_reduce,
      (outputs.size() == 1)
          ? outputs[0]
          : computation.AddInstruction(HloInstruction::CreateTuple(outputs))));
  TF_RETURN_IF_ERROR(
      TryDecomposeAllReduce(Cast<HloAllReduceInstruction>(new_all_reduce),
                            num_devices_per_host)
          .status());
  return true;
}
StatusOr < bool > TryDecomposeAllReduce ( HloAllReduceInstruction * all_reduce , size_t num_devices_per_host ) { TF_RET_CHECK ( all_reduce ) ; TF_RET_CHECK ( ! all_reduce -> has_sharding ( ) ) ; HloComputation & computation = * all_reduce -> parent ( ) ; PrimitiveType element_type = all_reduce -> operand ( 0 ) -> shape ( ) . element_type ( ) ; TF_ASSIGN_OR_RETURN ( absl :: optional < DecomposedReplicaGroups > decomposed_groups , TryDecomposeReplicaGroups ( * all_reduce , num_devices_per_host ) ) ; if ( ! decomposed_groups ) return false ; std :: vector < HloInstruction * > flat_operands ; flat_operands . reserve ( all_reduce -> operand_count ( ) ) ; std :: vector < Shape > flat_shapes ; flat_shapes . reserve ( all_reduce -> operand_count ( ) ) ; std :: vector < Shape > scattered_shapes ; scattered_shapes . reserve ( all_reduce -> operand_count ( ) ) ; int scatter_group_size = decomposed_groups -> scatter_gather_groups [ 0 ] . replica_ids_size ( ) ; for ( HloInstruction * operand : all_reduce -> operands ( ) ) { TF_RET_CHECK ( operand -> shape ( ) . IsArray ( ) ) ; int64_t num_elements = ShapeUtil :: ElementsIn ( operand -> shape ( ) ) ; Shape flat_shape = ShapeUtil :: MakeShape ( element_type , { num_elements } ) ; flat_operands . push_back ( computation . AddInstruction ( HloInstruction :: CreateBitcast ( flat_shape , operand ) ) ) ; flat_shapes . push_back ( std :: move ( flat_shape ) ) ; scattered_shapes . push_back ( ShapeUtil :: MakeShape ( element_type , { num_elements / scatter_group_size } ) ) ; } Shape reduce_scatter_shape = ShapeUtil :: MakeMaybeTupleShape ( scattered_shapes ) ; HloInstruction * reduce_scatter = computation . AddInstruction ( HloInstruction :: CreateReduceScatter ( reduce_scatter_shape , flat_operands , all_reduce -> to_apply ( ) , decomposed_groups -> scatter_gather_groups , false , all_reduce -> channel_id ( ) , all_reduce -> use_global_device_ids ( ) , 0 ) ) ; HloInstruction * new_all_reduce = computation . AddInstruction ( HloInstruction :: CreateAllReduce ( reduce_scatter_shape , GetOutputs ( * reduce_scatter ) , all_reduce -> to_apply ( ) , decomposed_groups -> new_all_reduce_groups , false , all_reduce -> channel_id ( ) , all_reduce -> use_global_device_ids ( ) ) ) ; HloInstruction * all_gather = computation . AddInstruction ( HloInstruction :: CreateAllGather ( ShapeUtil :: MakeMaybeTupleShape ( flat_shapes ) , GetOutputs ( * new_all_reduce ) , 0 , decomposed_groups -> scatter_gather_groups , false , all_reduce -> channel_id ( ) , all_reduce -> use_global_device_ids ( ) ) ) ; std :: vector < HloInstruction * > outputs = GetOutputs ( * all_gather ) ; for ( int64_t i = 0 ; i < outputs . size ( ) ; ++ i ) { outputs [ i ] = computation . AddInstruction ( HloInstruction :: CreateBitcast ( all_reduce -> operand ( i ) -> shape ( ) , outputs [ i ] ) ) ; } TF_RETURN_IF_ERROR ( computation . ReplaceInstruction ( all_reduce , ( outputs . size ( ) == 1 ) ? outputs [ 0 ] : computation . AddInstruction ( HloInstruction :: CreateTuple ( outputs ) ) ) ) ; TF_RETURN_IF_ERROR ( TryDecomposeAllReduce ( Cast < HloAllReduceInstruction > ( new_all_reduce ) , num_devices_per_host ) . status ( ) ) ; return true ; }
Attempts to decompose all-reduces as described by the BlueConnect paper.
Attempts to decompose all - reduces as described by the BlueConnect paper .
// never null // Bitcast operands to 1D to guarantee that first dimension is divisible by // scatter group size (we checked num elements was divisible above). /*constrain_layout=*/ /*scatter_dimension=*/ /*constrain_layout=*/ /*all_gather_dimension=*/ /*constrain_layout=*/ // Bitcast back to the original shapes and replace all-reduce with decomposed // implementation. // Try to apply decomposition recursively.