1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #include "tensorflow/lite/delegates/flex/training/training_delegate.h" 17 18 namespace tflite { 19 namespace flex { 20 21 // Corresponding weak declaration found in lite/interpreter_builder.cc. 22 #if TFLITE_HAS_ATTRIBUTE_WEAK 23 // If weak symbol is not supported (Windows), it can use 24 // TF_AcquireFlexDelegate() path instead. AcquireFlexDelegate()25TfLiteDelegateUniquePtr AcquireFlexDelegate() { return {nullptr, nullptr}; } 26 #endif 27 TrainingFlexDelegate()28TrainingFlexDelegate::TrainingFlexDelegate() 29 : delegate_(FlexDelegate::Create()) { 30 cancellation_manager_ = absl::make_unique<tensorflow::CancellationManager>(); 31 auto* flex_delegate = static_cast<FlexDelegate*>(delegate_->data_); 32 auto* delegate_data = flex_delegate->mutable_data(); 33 delegate_data->SetCancellationManager(cancellation_manager_.get()); 34 } 35 Cancel()36void TrainingFlexDelegate::Cancel() { cancellation_manager_->StartCancel(); } 37 ShouldCancel(void * data)38bool TrainingFlexDelegate::ShouldCancel(void* data) { 39 if (data == nullptr) { 40 return false; 41 } 42 43 auto* training_flex_delegate = static_cast<TrainingFlexDelegate*>(data); 44 return training_flex_delegate->cancellation_manager_->IsCancelled(); 45 } 46 47 } // namespace flex 48 } // namespace tflite 49 50 // Exported C interface function which is used by AcquireFlexDelegate() at 51 // interpreter_builder.cc. To export the function name globally, the function 52 // name must be matched with patterns in tf_version_script.lds. In Android, we 53 // don't use this feature so skip building. 54 #if !defined(__ANDROID__) 55 extern "C" { TF_AcquireFlexDelegate()56TFL_CAPI_EXPORT tflite::TfLiteDelegateUniquePtr TF_AcquireFlexDelegate() { 57 return {nullptr, nullptr}; 58 } 59 } // extern "C" 60 #endif // !defined(__ANDROID__) 61