tensorflow/compiler/mlir/xla/transforms/legalize_tf_communication.cc
// Assigns frontend attributes holding information about data type and
// TensorFlow rendezvous channel name. The TensorFlow rendezvous channel name is
// handled differently as individual names are used per data send and receive.
Assigns frontend attributes holding information about data type and
TensorFlow rendezvous channel name. The TensorFlow rendezvous channel name is
handled differently as individual names are used per data send and receive.
Assigns frontend attributes holding information about data type and TensorFlow rendezvous channel name . The TensorFlow rendezvous channel name is handled differently as individual names are used per data send and receive .
void SetFrontendAttributes(Operation* op, int32_t index, StringRef key,
                           Type type, bool device_to_host,
                           StringRef host_handler_name) {
  MLIRContext* context = op->getContext();
  std::string formatted_key =
      device_to_host ? llvm::formatv("{0}_dtoh_{1}", key, index).str()
                     : llvm::formatv("{0}_htod_{1}", key, index).str();
  auto rendezvous_name = StringAttr::get(context, formatted_key);
  auto rendezvous_name_attr = NamedAttribute(
      StringAttr::get(context, xla::kXlaHostTransferRendezvousNameAttr),
      rendezvous_name);
  auto element_type = getElementTypeOrSelf(type);
  auto xla_element_type = ::xla::TypeToPrimitiveType(element_type);
  const std::string& xla_element_type_str =
      ::xla::primitive_util::LowercasePrimitiveTypeName(xla_element_type);
  auto original_type = StringAttr::get(context, xla_element_type_str);
  auto original_type_attr = NamedAttribute(
      StringAttr::get(context, xla::kXlaHostTransferOriginalTypeAttr),
      original_type);
  auto host_handler_name_value =
      StringAttr::get(context, host_handler_name.str());
  auto host_handler_name_attr = NamedAttribute(
      StringAttr::get(context, xla::kXlaHostTransferHandlerNameAttr),
      host_handler_name_value);
  auto frontend_attributes = DictionaryAttr::get(
      context,
      ArrayRef<NamedAttribute>{rendezvous_name_attr, original_type_attr,
                               host_handler_name_attr});
  op->setAttr(kFrontendAttributesAttr, frontend_attributes);
}
void SetFrontendAttributes ( Operation * op , int32_t index , StringRef key , Type type , bool device_to_host , StringRef host_handler_name ) { MLIRContext * context = op -> getContext ( ) ; std :: string formatted_key = device_to_host ? llvm :: formatv ( " " , key , index ) . str ( ) : llvm :: formatv ( " " , key , index ) . str ( ) ; auto rendezvous_name = StringAttr :: get ( context , formatted_key ) ; auto rendezvous_name_attr = NamedAttribute ( StringAttr :: get ( context , xla :: kXlaHostTransferRendezvousNameAttr ) , rendezvous_name ) ; auto element_type = getElementTypeOrSelf ( type ) ; auto xla_element_type = :: xla :: TypeToPrimitiveType ( element_type ) ; const std :: string & xla_element_type_str = :: xla :: primitive_util :: LowercasePrimitiveTypeName ( xla_element_type ) ; auto original_type = StringAttr :: get ( context , xla_element_type_str ) ; auto original_type_attr = NamedAttribute ( StringAttr :: get ( context , xla :: kXlaHostTransferOriginalTypeAttr ) , original_type ) ; auto host_handler_name_value = StringAttr :: get ( context , host_handler_name . str ( ) ) ; auto host_handler_name_attr = NamedAttribute ( StringAttr :: get ( context , xla :: kXlaHostTransferHandlerNameAttr ) , host_handler_name_value ) ; auto frontend_attributes = DictionaryAttr :: get ( context , ArrayRef < NamedAttribute > { rendezvous_name_attr , original_type_attr , host_handler_name_attr } ) ; op -> setAttr ( kFrontendAttributesAttr , frontend_attributes ) ; }
Assigns frontend attributes holding information about data type and
TensorFlow rendezvous channel name.
Assigns frontend attributes holding information about data type and TensorFlow rendezvous channel name .