tensorflow/compiler/mlir/xla/transforms/legalize_tf_communication.cc
// Creates a `mhlo.recv` op for receiving a value. If `tpu_core` is set, op
// sharding for the respective device will be set.
Creates a `mhlo.recv` op for receiving a value. If `tpu_core` is set, op
sharding for the respective device will be set.
Creates a ` mhlo . recv ` op for receiving a value . If ` tpu_core ` is set op sharding for the respective device will be set .
Value CreateRecvOp(OpBuilder& builder, int64_t& channel_id, Location loc,
                   Value result, StringRef key, size_t index,
                   const Optional<int64_t>& tpu_core, Value token,
                   StringRef host_handler_name) {
  auto channel_handle = ChannelHandle::get(
      builder.getI64IntegerAttr(channel_id++),
      builder.getI64IntegerAttr(3), builder.getContext());
  auto result_type = result.getType();
  SmallVector<Type, 2> recv_result_type = {result_type, token.getType()};
  auto recv =
      builder.create<RecvOp>(loc, recv_result_type, token, channel_handle,
                             builder.getBoolAttr(true));
  SetFrontendAttributes(recv, index, key, result_type,
                        false, host_handler_name);
  if (tpu_core) SetOpSharding(recv, *tpu_core);
  result.replaceAllUsesWith(recv.getResult(0));
  return recv.getResult(1);
}
Value CreateRecvOp ( OpBuilder & builder , int64_t & channel_id , Location loc , Value result , StringRef key , size_t index , const Optional < int64_t > & tpu_core , Value token , StringRef host_handler_name ) { auto channel_handle = ChannelHandle :: get ( builder . getI64IntegerAttr ( channel_id ++ ) , builder . getI64IntegerAttr ( 3 ) , builder . getContext ( ) ) ; auto result_type = result . getType ( ) ; SmallVector < Type , 2 > recv_result_type = { result_type , token . getType ( ) } ; auto recv = builder . create < RecvOp > ( loc , recv_result_type , token , channel_handle , builder . getBoolAttr ( true ) ) ; SetFrontendAttributes ( recv , index , key , result_type , false , host_handler_name ) ; if ( tpu_core ) SetOpSharding ( recv , * tpu_core ) ; result . replaceAllUsesWith ( recv . getResult ( 0 ) ) ; return recv . getResult ( 1 ) ; }
Creates a `mhlo.recv` op for receiving a value.
Creates a ` mhlo . recv ` op for receiving a value .
// type 3 == HOST_TO_DEVICE /*handle=*/ /*type=*/ /*is_host_transfer=*/ /*device_to_host=*/