aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
diff options
context:
space:
mode:
Diffstat (limited to 'tensorflow/core/grappler/costs/op_level_cost_estimator.cc')
-rw-r--r--tensorflow/core/grappler/costs/op_level_cost_estimator.cc135
1 files changed, 71 insertions, 64 deletions
diff --git a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
index d34eecd009..5b303f6ccb 100644
--- a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
+++ b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
@@ -65,6 +65,7 @@ constexpr char kAvgPool[] = "AvgPool";
constexpr char kAvgPoolGrad[] = "AvgPoolGrad";
constexpr char kFusedBatchNorm[] = "FusedBatchNorm";
constexpr char kFusedBatchNormGrad[] = "FusedBatchNormGrad";
+constexpr char kQuantizedMatMulV2[] = "QuantizedMatMulV2";
static const Costs::Duration kMinComputeTime(1);
@@ -226,6 +227,7 @@ OpLevelCostEstimator::OpLevelCostEstimator() {
{kMatMul, wrap(&OpLevelCostEstimator::PredictMatMul)},
{kSparseMatMul, wrap(&OpLevelCostEstimator::PredictMatMul)},
{kBatchMatMul, wrap(&OpLevelCostEstimator::PredictBatchMatMul)},
+ {kQuantizedMatMulV2, wrap(&OpLevelCostEstimator::PredictMatMul)},
{kNoOp, wrap(&OpLevelCostEstimator::PredictNoOp)},
{kGuaranteeConst, wrap(&OpLevelCostEstimator::PredictNoOp)},
@@ -268,67 +270,70 @@ OpLevelCostEstimator::OpLevelCostEstimator() {
EIGEN_COST(scalar_product_op<float>) + EIGEN_COST(scalar_max_op<float>) +
EIGEN_COST(scalar_min_op<float>) + EIGEN_COST(scalar_round_op<float>);
- elementwise_ops_ = {// Unary ops alphabetically sorted
- {"Acos", EIGEN_COST(scalar_acos_op<float>)},
- {"Asin", EIGEN_COST(scalar_asin_op<float>)},
- {"Atan", EIGEN_COST(scalar_atan_op<float>)},
- {"Atan2", EIGEN_COST(scalar_quotient_op<float>) +
- EIGEN_COST(scalar_atan_op<float>)},
- {"Ceil", EIGEN_COST(scalar_ceil_op<float>)},
- {"Cos", EIGEN_COST(scalar_cos_op<float>)},
- {"Dequantize", EIGEN_COST(scalar_product_op<float>)},
- {"Erf", 1},
- {"Erfc", 1},
- {"Exp", EIGEN_COST(scalar_exp_op<float>)},
- {"Expm1", EIGEN_COST(scalar_expm1_op<float>)},
- {"Floor", EIGEN_COST(scalar_floor_op<float>)},
- {"Inv", EIGEN_COST(scalar_inverse_op<float>)},
- {"InvGrad", 1},
- {"Lgamma", 1},
- {"Log", EIGEN_COST(scalar_log_op<float>)},
- {"Log1p", EIGEN_COST(scalar_log1p_op<float>)},
- {"Neg", EIGEN_COST(scalar_opposite_op<float>)},
- {"QuantizeV2", quantize_v2_cost},
- {"Reciprocal", EIGEN_COST(scalar_inverse_op<float>)},
- {"Rint", 1},
- {"Round", EIGEN_COST(scalar_round_op<float>)},
- {"Rsqrt", EIGEN_COST(scalar_rsqrt_op<float>)},
- {"Sqrt", EIGEN_COST(scalar_sqrt_op<float>)},
- {"Square", EIGEN_COST(scalar_square_op<float>)},
- {"Tanh", EIGEN_COST(scalar_tanh_op<float>)},
- {"Relu", EIGEN_COST(scalar_max_op<float>)},
- {"Sigmoid", EIGEN_COST(scalar_sigmoid_op<float>)},
- {"Sign", EIGEN_COST(scalar_sign_op<float>)},
- {"Sin", EIGEN_COST(scalar_sin_op<float>)},
- {"Tan", EIGEN_COST(scalar_tan_op<float>)},
- // Binary ops alphabetically sorted
- {"Add", EIGEN_COST(scalar_sum_op<float>)},
- {"ApproximateEqual", 1},
- {"BiasAdd", EIGEN_COST(scalar_sum_op<float>)},
- {"Div", EIGEN_COST(scalar_quotient_op<float>)},
- {"Equal", 1},
- {"FloorDiv", EIGEN_COST(scalar_quotient_op<float>)},
- {"FloorMod", EIGEN_COST(scalar_mod_op<float>)},
- {"Greater", 1},
- {"GreaterEqual", 1},
- {"Less", 1},
- {"LessEqual", 1},
- {"LogicalAnd", EIGEN_COST(scalar_boolean_and_op)},
- {"LogicalNot", 1},
- {"LogicalOr", EIGEN_COST(scalar_boolean_or_op)},
- {"Maximum", EIGEN_COST(scalar_max_op<float>)},
- {"Minimum", EIGEN_COST(scalar_min_op<float>)},
- {"Mod", EIGEN_COST(scalar_mod_op<float>)},
- {"Mul", EIGEN_COST(scalar_product_op<float>)},
- {"NotEqual", 1},
- {"QuantizedAdd", EIGEN_COST(scalar_sum_op<float>)},
- {"QuantizedMul", EIGEN_COST(scalar_product_op<float>)},
- {"RealDiv", EIGEN_COST(scalar_quotient_op<float>)},
- {"ReluGrad", EIGEN_COST(scalar_max_op<float>)},
- {"SquareDifference", 1},
- {"Sub", EIGEN_COST(scalar_difference_op<float>)},
- {"TruncateDiv", EIGEN_COST(scalar_quotient_op<float>)},
- {"TruncateMod", EIGEN_COST(scalar_mod_op<float>)}};
+ elementwise_ops_ = {
+ // Unary ops alphabetically sorted
+ {"Acos", EIGEN_COST(scalar_acos_op<float>)},
+ {"Asin", EIGEN_COST(scalar_asin_op<float>)},
+ {"Atan", EIGEN_COST(scalar_atan_op<float>)},
+ {"Atan2", EIGEN_COST(scalar_quotient_op<float>) +
+ EIGEN_COST(scalar_atan_op<float>)},
+ {"Ceil", EIGEN_COST(scalar_ceil_op<float>)},
+ {"Cos", EIGEN_COST(scalar_cos_op<float>)},
+ {"Dequantize", EIGEN_COST(scalar_product_op<float>)},
+ {"Erf", 1},
+ {"Erfc", 1},
+ {"Exp", EIGEN_COST(scalar_exp_op<float>)},
+ {"Expm1", EIGEN_COST(scalar_expm1_op<float>)},
+ {"Floor", EIGEN_COST(scalar_floor_op<float>)},
+ {"Inv", EIGEN_COST(scalar_inverse_op<float>)},
+ {"InvGrad", 1},
+ {"Lgamma", 1},
+ {"Log", EIGEN_COST(scalar_log_op<float>)},
+ {"Log1p", EIGEN_COST(scalar_log1p_op<float>)},
+ {"Neg", EIGEN_COST(scalar_opposite_op<float>)},
+ {"QuantizeV2", quantize_v2_cost},
+ {"Reciprocal", EIGEN_COST(scalar_inverse_op<float>)},
+ {"Rint", 1},
+ {"Round", EIGEN_COST(scalar_round_op<float>)},
+ {"Rsqrt", EIGEN_COST(scalar_rsqrt_op<float>)},
+ {"Sqrt", EIGEN_COST(scalar_sqrt_op<float>)},
+ {"Square", EIGEN_COST(scalar_square_op<float>)},
+ {"Tanh", EIGEN_COST(scalar_tanh_op<float>)},
+ {"Relu", EIGEN_COST(scalar_max_op<float>)},
+ {"Sigmoid", EIGEN_COST(scalar_sigmoid_op<float>)},
+ {"QuantizedSigmoid", EIGEN_COST(scalar_sigmoid_op<float>)},
+ {"Sign", EIGEN_COST(scalar_sign_op<float>)},
+ {"Sin", EIGEN_COST(scalar_sin_op<float>)},
+ {"Tan", EIGEN_COST(scalar_tan_op<float>)},
+ // Binary ops alphabetically sorted
+ {"Add", EIGEN_COST(scalar_sum_op<float>)},
+ {"ApproximateEqual", 1},
+ {"BiasAdd", EIGEN_COST(scalar_sum_op<float>)},
+ {"QuantizedBiasAdd", EIGEN_COST(scalar_sum_op<float>)},
+ {"Div", EIGEN_COST(scalar_quotient_op<float>)},
+ {"Equal", 1},
+ {"FloorDiv", EIGEN_COST(scalar_quotient_op<float>)},
+ {"FloorMod", EIGEN_COST(scalar_mod_op<float>)},
+ {"Greater", 1},
+ {"GreaterEqual", 1},
+ {"Less", 1},
+ {"LessEqual", 1},
+ {"LogicalAnd", EIGEN_COST(scalar_boolean_and_op)},
+ {"LogicalNot", 1},
+ {"LogicalOr", EIGEN_COST(scalar_boolean_or_op)},
+ {"Maximum", EIGEN_COST(scalar_max_op<float>)},
+ {"Minimum", EIGEN_COST(scalar_min_op<float>)},
+ {"Mod", EIGEN_COST(scalar_mod_op<float>)},
+ {"Mul", EIGEN_COST(scalar_product_op<float>)},
+ {"NotEqual", 1},
+ {"QuantizedAdd", EIGEN_COST(scalar_sum_op<float>)},
+ {"QuantizedMul", EIGEN_COST(scalar_product_op<float>)},
+ {"RealDiv", EIGEN_COST(scalar_quotient_op<float>)},
+ {"ReluGrad", EIGEN_COST(scalar_max_op<float>)},
+ {"SquareDifference", 1},
+ {"Sub", EIGEN_COST(scalar_difference_op<float>)},
+ {"TruncateDiv", EIGEN_COST(scalar_quotient_op<float>)},
+ {"TruncateMod", EIGEN_COST(scalar_mod_op<float>)}};
#undef EIGEN_COST
@@ -675,7 +680,7 @@ int64 OpLevelCostEstimator::CountMatMulOperations(
}
ops = m_dim * n_dim * k_dim * 2;
- VLOG(1) << "Operations for Matmul" << ops;
+ VLOG(1) << "Operations for Matmul: " << ops;
if (mat_mul != nullptr) {
mat_mul->m = m_dim;
@@ -972,8 +977,10 @@ int64 OpLevelCostEstimator::CalculateTensorElementCount(
int64 OpLevelCostEstimator::CalculateTensorSize(
const OpInfo::TensorProperties& tensor, bool* found_unknown_shapes) const {
- return CalculateTensorElementCount(tensor, found_unknown_shapes) *
- DataTypeSize(BaseType(tensor.dtype()));
+ int64 count = CalculateTensorElementCount(tensor, found_unknown_shapes);
+ int size = DataTypeSize(BaseType(tensor.dtype()));
+ VLOG(2) << "Count: " << count << " DataTypeSize: " << size;
+ return count * size;
}
int64 OpLevelCostEstimator::CalculateInputSize(
generated by cgit v1.2.3 (git 2.39.1) at 2024-06-04 15:17:11 +0000