| 335 | } |
| 336 | |
| 337 | void FunctionMetadata::ValidateMultiDevice() { |
| 338 | const FunctionDef* fdef = lib_def_->Find(func_.name()); |
| 339 | if (is_multi_device_function_) { |
| 340 | auto attr = fdef->attr().find(FunctionLibraryDefinition::kIntsOnDeviceAttr); |
| 341 | if (attr != fdef->attr().end() && attr->second.b()) { |
| 342 | LOG(WARNING) |
| 343 | << "Disabling multi-device execution for a function that uses the " |
| 344 | << FunctionLibraryDefinition::kIntsOnDeviceAttr << " attribute."; |
| 345 | is_multi_device_function_ = false; |
| 346 | return; |
| 347 | } |
| 348 | auto validate_arg = [this](const OpDef::ArgDef& arg) { |
| 349 | if (!arg.number_attr().empty() || !arg.type_list_attr().empty()) { |
| 350 | LOG(WARNING) << "Disabling multi-device execution for a function with " |
| 351 | "a vector argument " |
| 352 | << arg.name() << "."; |
| 353 | is_multi_device_function_ = false; |
| 354 | } |
| 355 | }; |
| 356 | for (const auto& arg : fdef->signature().input_arg()) { |
| 357 | validate_arg(arg); |
| 358 | } |
| 359 | for (const auto& arg : fdef->signature().output_arg()) { |
| 360 | validate_arg(arg); |
| 361 | } |
| 362 | } |
| 363 | } |
| 364 | |
| 365 | /* static */ |
| 366 | Status CapturedFunction::Create( |