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 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_COMPILE_TIME_CAP_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_COMPILE_TIME_CAP_H_ 18 #include <algorithm> 19 20 #include "absl/strings/string_view.h" 21 #include "tensorflow/compiler/xla/service/hlo_module.h" 22 23 namespace xla { 24 // Provide a common way to bound compiler analyses that potentially have 25 // overhead that is non-linear to the number of instructions in a module. 26 class BoundNonLinearCompilerAnalysis { 27 public: 28 // Sampling_rate specifies the proportion of all instructions expected to be 29 // analyzed. e.g., if sampling_rate_=2, then every other instructions are 30 // expected to be analyzed. If sample_rate <= 0, the analysis will be always 31 // allowed to complete. Each analysis is allowed at least a constant number of 32 // abstract cost units, before it is considered for early termination. BoundNonLinearCompilerAnalysis(HloModule * m,absl::string_view pass_name,std::optional<int64_t> sampling_rate)33 explicit BoundNonLinearCompilerAnalysis(HloModule* m, 34 absl::string_view pass_name, 35 std::optional<int64_t> sampling_rate) 36 : analysis_allowance_( 37 (!sampling_rate.has_value() || sampling_rate.value() <= 0 || 38 m->config().GetAnalysisAllowance(pass_name) < 0) 39 ? -1 40 : std::max(m->config().GetAnalysisAllowance(pass_name), 41 m->instruction_count() / sampling_rate.value())) {} 42 // Return whether the cost is deducted successfully. If not, the analysis 43 // should be terminated as its overhead is too high. DeductCost(int64_t cost_now)44 bool DeductCost(int64_t cost_now) { 45 if (analysis_allowance_ > 0 && cost_now > 0) { 46 analysis_allowance_ -= cost_now; 47 if (analysis_allowance_ < 0) { 48 analysis_allowance_ = 0; 49 } 50 } 51 return analysis_allowance_ != 0; 52 } 53 ContinueAnalysis()54 bool ContinueAnalysis() const { return analysis_allowance_ != 0; } analysis_allowance()55 int64_t analysis_allowance() const { return analysis_allowance_; } 56 57 private: 58 int64_t analysis_allowance_; 59 }; 60 61 }; // namespace xla 62 63 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_COMPILE_TIME_CAP_H_ 64