kernel_registry.cc revision bc225bfaa534acc25047fe844f19edc333b7a76a
1783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower/* Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower 3783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlowerLicensed under the Apache License, Version 2.0 (the "License"); 4783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFloweryou may not use this file except in compliance with the License. 5783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlowerYou may obtain a copy of the License at 6783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower 7783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower http://www.apache.org/licenses/LICENSE-2.0 8783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower 9783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlowerUnless required by applicable law or agreed to in writing, software 10783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlowerdistributed under the License is distributed on an "AS IS" BASIS, 11783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlowerWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlowerSee the License for the specific language governing permissions and 13783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlowerlimitations under the License. 14783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower==============================================================================*/ 15783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower#include "tensorflow/python/util/kernel_registry.h" 16783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower 17783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower#include "tensorflow/core/framework/node_def_util.h" 18783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower#include "tensorflow/core/framework/op.h" 19783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower#include "tensorflow/core/framework/op_kernel.h" 20783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower#include "tensorflow/core/framework/types.h" 21783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower#include "tensorflow/core/lib/core/status.h" 22783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower#include "tensorflow/core/util/device_name_utils.h" 23783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower 24783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlowernamespace tensorflow { 25783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlowernamespace swig { 26783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower 27783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlowerstring TryFindKernelClass(const string& serialized_node_def) { 28783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower tensorflow::NodeDef node_def; 29783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower if (!node_def.ParseFromString(serialized_node_def)) { 30783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower LOG(WARNING) << "Error parsing node_def"; 31783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower return ""; 32783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower } 33783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower 34783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower const tensorflow::OpRegistrationData* op_reg_data; 35783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower auto status = 36783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower tensorflow::OpRegistry::Global()->LookUp(node_def.op(), &op_reg_data); 37783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower if (!status.ok()) { 38783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower LOG(WARNING) << "Op " << node_def.op() << " not found: " << status; 39783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower return ""; 40783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower } 41783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower AddDefaultsToNodeDef(op_reg_data->op_def, &node_def); 42783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower 43783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower tensorflow::DeviceNameUtils::ParsedName parsed_name; 44783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower if (!tensorflow::DeviceNameUtils::ParseFullName(node_def.device(), 45783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower &parsed_name)) { 46783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower LOG(WARNING) << "Failed to parse device from node_def: " 47783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower << node_def.ShortDebugString(); 48783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower return ""; 49783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower } 50783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower string class_name = ""; 51783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower tensorflow::FindKernelDef(tensorflow::DeviceType(parsed_name.type.c_str()), 52bc225bfaa534acc25047fe844f19edc333b7a76aPeter Hawkins node_def, nullptr /* kernel_def */, &class_name) 53bc225bfaa534acc25047fe844f19edc333b7a76aPeter Hawkins .IgnoreError(); 54783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower return class_name; 55783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower} 56783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower 57783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower} // namespace swig 58783e8388db67f3fefc6c714d479615821a9dc5e7A. Unique TensorFlower} // namespace tensorflow 59