tensorflow/compiler/mlir/tensorflow/transforms/tpu_sharding_identification_pass.cc
// Extracts input/output sharding configuration of `cluster_func` by parsing
// XlaSharding ops inside the `cluster_func`.
Extracts input/output sharding configuration of `cluster_func` by parsing
XlaSharding ops inside the `cluster_func`.
Extracts input / output sharding configuration of ` cluster_func ` by parsing XlaSharding ops inside the ` cluster_func ` .
void IdentifyXlaShardingForTPUComputation(
    Builder* builder, tf_device::ClusterFuncOp cluster_func) {
  FuncOp func = cluster_func->getParentOfType<ModuleOp>().lookupSymbol<FuncOp>(
      cluster_func.func());
  const std::string logical_core_0_sharding =
      xla::sharding_builder::AssignDevice(0).SerializeAsString();
  bool use_spmd = false;
  if (auto use_spmd_attr = cluster_func->getAttrOfType<BoolAttr>(kUseSpmdAttr))
    use_spmd = use_spmd_attr.getValue();
  llvm::SmallVector<llvm::StringRef, 8> sharding_for_args;
  IdentifyXlaShardingForComputationInputs(logical_core_0_sharding, use_spmd,
                                          true,
                                          cluster_func, func, builder,
                                          sharding_for_args);
  llvm::SmallVector<llvm::StringRef, 8> sharding_for_rets;
  IdentifyXlaShardingForComputationOutputs(
      logical_core_0_sharding, use_spmd, true,
      cluster_func, func, builder, sharding_for_args, sharding_for_rets);
  auto has_maximal_sharding = [](llvm::StringRef sharding_string) -> bool {
    xla::OpSharding sharding;
    sharding.ParseFromString(sharding_string.str());
    return sharding.type() == xla::OpSharding::MAXIMAL;
  };
  if ((use_spmd && (absl::c_any_of(sharding_for_args, has_maximal_sharding) ||
                    absl::c_any_of(sharding_for_rets, has_maximal_sharding))) ||
      failed(VerifyShardings(func, sharding_for_args, sharding_for_rets))) {
    LOG(WARNING) << "XLA SPMD only supports cases where all inputs/outputs "
                    "exist on every partition (sharded or replicated). If any "
                    "of the inputs/outputs have maximal sharding, then "
                    "fallback to MPMD.";
    sharding_for_args.clear();
    sharding_for_rets.clear();
    cluster_func->setAttr(kUseSpmdAttr, builder->getBoolAttr(false));
    IdentifyXlaShardingForComputationInputs(
        logical_core_0_sharding,
        false, false, cluster_func,
        func, builder, sharding_for_args);
    IdentifyXlaShardingForComputationOutputs(
        logical_core_0_sharding,
        false, false, cluster_func,
        func, builder, sharding_for_args, sharding_for_rets);
  }
  Block& function_block = func.front();
  for (auto sharding_and_arg :
       llvm::zip(sharding_for_args, function_block.getArguments())) {
    StringRef sharding = std::get<0>(sharding_and_arg);
    BlockArgument arg = std::get<1>(sharding_and_arg);
    func.setArgAttr(arg.getArgNumber(), kShardingAttr,
                    builder->getStringAttr(sharding));
  }
  Operation* terminator = function_block.getTerminator();
  for (auto sharding_and_retval :
       llvm::zip(sharding_for_rets, terminator->getOpOperands())) {
    StringRef sharding = std::get<0>(sharding_and_retval);
    OpOperand& retval = std::get<1>(sharding_and_retval);
    func.setResultAttr(retval.getOperandNumber(), kShardingAttr,
                       builder->getStringAttr(sharding));
  }
  cluster_func->setAttr(tensorflow::kInputShardingAttr,
                        builder->getStrArrayAttr(sharding_for_args));
  cluster_func->setAttr(tensorflow::kOutputShardingAttr,
                        builder->getStrArrayAttr(sharding_for_rets));
}
void IdentifyXlaShardingForTPUComputation ( Builder * builder , tf_device :: ClusterFuncOp cluster_func ) { FuncOp func = cluster_func -> getParentOfType < ModuleOp > ( ) . lookupSymbol < FuncOp > ( cluster_func . func ( ) ) ; const std :: string logical_core_0_sharding = xla :: sharding_builder :: AssignDevice ( 0 ) . SerializeAsString ( ) ; bool use_spmd = false ; if ( auto use_spmd_attr = cluster_func -> getAttrOfType < BoolAttr > ( kUseSpmdAttr ) ) use_spmd = use_spmd_attr . getValue ( ) ; llvm :: SmallVector < llvm :: StringRef , 8 > sharding_for_args ; IdentifyXlaShardingForComputationInputs ( logical_core_0_sharding , use_spmd , true , cluster_func , func , builder , sharding_for_args ) ; llvm :: SmallVector < llvm :: StringRef , 8 > sharding_for_rets ; IdentifyXlaShardingForComputationOutputs ( logical_core_0_sharding , use_spmd , true , cluster_func , func , builder , sharding_for_args , sharding_for_rets ) ; auto has_maximal_sharding = [ ] ( llvm :: StringRef sharding_string ) -> bool { xla :: OpSharding sharding ; sharding . ParseFromString ( sharding_string . str ( ) ) ; return sharding . type ( ) == xla :: OpSharding :: MAXIMAL ; } ; if ( ( use_spmd && ( absl :: c_any_of ( sharding_for_args , has_maximal_sharding ) || absl :: c_any_of ( sharding_for_rets , has_maximal_sharding ) ) ) || failed ( VerifyShardings ( func , sharding_for_args , sharding_for_rets ) ) ) { LOG ( WARNING ) << " " " " " " " " ; sharding_for_args . clear ( ) ; sharding_for_rets . clear ( ) ; cluster_func -> setAttr ( kUseSpmdAttr , builder -> getBoolAttr ( false ) ) ; IdentifyXlaShardingForComputationInputs ( logical_core_0_sharding , false , false , cluster_func , func , builder , sharding_for_args ) ; IdentifyXlaShardingForComputationOutputs ( logical_core_0_sharding , false , false , cluster_func , func , builder , sharding_for_args , sharding_for_rets ) ; } Block & function_block = func . front ( ) ; for ( auto sharding_and_arg : llvm :: zip ( sharding_for_args , function_block . getArguments ( ) ) ) { StringRef sharding = std :: get < 0 > ( sharding_and_arg ) ; BlockArgument arg = std :: get < 1 > ( sharding_and_arg ) ; func . setArgAttr ( arg . getArgNumber ( ) , kShardingAttr , builder -> getStringAttr ( sharding ) ) ; } Operation * terminator = function_block . getTerminator ( ) ; for ( auto sharding_and_retval : llvm :: zip ( sharding_for_rets , terminator -> getOpOperands ( ) ) ) { StringRef sharding = std :: get < 0 > ( sharding_and_retval ) ; OpOperand & retval = std :: get < 1 > ( sharding_and_retval ) ; func . setResultAttr ( retval . getOperandNumber ( ) , kShardingAttr , builder -> getStringAttr ( sharding ) ) ; } cluster_func -> setAttr ( tensorflow :: kInputShardingAttr , builder -> getStrArrayAttr ( sharding_for_args ) ) ; cluster_func -> setAttr ( tensorflow :: kOutputShardingAttr , builder -> getStrArrayAttr ( sharding_for_rets ) ) ; }
Extracts input/output sharding configuration of `cluster_func` by parsing
XlaSharding ops inside the `cluster_func`.
Extracts input / output sharding configuration of ` cluster_func ` by parsing XlaSharding ops inside the ` cluster_func ` .
// Look up function definition from module. // By default inputs/outputs have maximal sharding and are assigned to logical // core 0 if no sharding is defined. /*infer_from_computation=*/ /*infer_from_computation=*/ // XLA SPMD only supports cases where all inputs/outputs exist on every // partition (sharded or replicated). If any of the inputs/outputs have // maximal sharding, then fallback to MPMD. Also fall back if any of the // shardings aren't compatible with the rank of their tensor. /*use_spmd=*/ /*infer_from_computation=*/ /*use_spmd=*/ /*infer_from_computation=*/ // Update sharding on function arguments and returns. // Update input/output sharding attributes on tf_device.cluster_func op.