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_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_COST_MODEL_H_ 17 #define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_COST_MODEL_H_ 18 19 #include <string> 20 21 #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project 22 #include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h" 23 24 namespace mlir { 25 namespace TFL { 26 namespace tac { 27 28 // TODO(renjieliu): We need to come up with a better strategy to do cost 29 // estimatation. Maybe build a big lookup table for all the ops. 30 31 // TODO(renjieliu): We need to consider what's the default value if we cannot 32 // analyze the cost. 33 34 // ================== Interface ======================== 35 36 // Get the estimated cost for the op under the given hardware spec senario. 37 float GetCostForOp(Operation* op, const std::string& hardware); 38 39 // Get the estimated cost for the whole function under the given hardware. 40 float GetCostForFunc(func::FuncOp* func, const std::string& hardware); 41 42 // Get the transfer cost given from & to hardware info. 43 // We will only calculate for the "necessary" tensor transferred. 44 // from_graph & to_graph are used to compute the "necessary" tensors. 45 // from_graph 46 // / \ \ 47 // out1 out2 out3 48 // \ / 49 // to_graph 50 // So only out2 & out3 are counted. 51 float GetTransferCost(const std::string& from_hardware_str, 52 const std::string& to_hardware_str, 53 func::CallOp from_graph, func::CallOp to_graph); 54 55 // Get the cross quantization/dequantization boundary cost. 56 float GetQuantDequantCost(InferenceType from_inference_type, 57 InferenceType to_inference_type, 58 func::CallOp from_graph, func::CallOp to_graph); 59 60 } // namespace tac 61 } // namespace TFL 62 } // namespace mlir 63 64 #endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_TAC_TRANSFORMS_COST_MODEL_H_ 65