Added a bunch of unary ops to the estimator.

PiperOrigin-RevId: 324607213
Change-Id: I24369f36cc29f68caac412a5d3076f5ef43859fe
This commit is contained in:
A. Unique TensorFlower 2020-08-03 08:22:33 -07:00 committed by TensorFlower Gardener
parent 123af9ba86
commit 8e9f3196fd
2 changed files with 30 additions and 17 deletions

View File

@ -522,6 +522,8 @@ OpLevelCostEstimator::OpLevelCostEstimator() {
// Unary ops alphabetically sorted
elementwise_ops_.emplace("Acos", EIGEN_COST(scalar_acos_op<float>));
elementwise_ops_.emplace("All", EIGEN_COST(scalar_boolean_and_op));
elementwise_ops_.emplace("ArgMax", EIGEN_COST(scalar_max_op<float>));
elementwise_ops_.emplace("Asin", EIGEN_COST(scalar_asin_op<float>));
elementwise_ops_.emplace("Atan", EIGEN_COST(scalar_atan_op<float>));
elementwise_ops_.emplace("Atan2", EIGEN_COST(scalar_quotient_op<float>) +
@ -546,7 +548,10 @@ OpLevelCostEstimator::OpLevelCostEstimator() {
elementwise_ops_.emplace("Lgamma", 1);
elementwise_ops_.emplace("Log", EIGEN_COST(scalar_log_op<float>));
elementwise_ops_.emplace("Log1p", EIGEN_COST(scalar_log1p_op<float>));
elementwise_ops_.emplace("Max", EIGEN_COST(scalar_max_op<float>));
elementwise_ops_.emplace("Min", EIGEN_COST(scalar_min_op<float>));
elementwise_ops_.emplace("Neg", EIGEN_COST(scalar_opposite_op<float>));
elementwise_ops_.emplace("Prod", EIGEN_COST(scalar_product_op<float>));
elementwise_ops_.emplace("QuantizeAndDequantizeV2",
quantize_and_dequantize_v2_cost);
elementwise_ops_.emplace("QuantizedSigmoid",
@ -554,6 +559,7 @@ OpLevelCostEstimator::OpLevelCostEstimator() {
elementwise_ops_.emplace("QuantizeV2", quantize_v2_cost);
elementwise_ops_.emplace("Reciprocal", EIGEN_COST(scalar_inverse_op<float>));
elementwise_ops_.emplace("Relu", EIGEN_COST(scalar_max_op<float>));
elementwise_ops_.emplace("Relu6", EIGEN_COST(scalar_max_op<float>));
elementwise_ops_.emplace("Rint", 1);
elementwise_ops_.emplace("Round", EIGEN_COST(scalar_round_op<float>));
elementwise_ops_.emplace("Rsqrt", EIGEN_COST(scalar_rsqrt_op<float>));
@ -562,8 +568,10 @@ OpLevelCostEstimator::OpLevelCostEstimator() {
elementwise_ops_.emplace("Sin", EIGEN_COST(scalar_sin_op<float>));
elementwise_ops_.emplace("Sqrt", EIGEN_COST(scalar_sqrt_op<float>));
elementwise_ops_.emplace("Square", EIGEN_COST(scalar_square_op<float>));
elementwise_ops_.emplace("Sum", EIGEN_COST(scalar_sum_op<float>));
elementwise_ops_.emplace("Tan", EIGEN_COST(scalar_tan_op<float>));
elementwise_ops_.emplace("Tanh", EIGEN_COST(scalar_tanh_op<float>));
elementwise_ops_.emplace("TopKV2", EIGEN_COST(scalar_max_op<float>));
// Binary ops alphabetically sorted
elementwise_ops_.emplace("Add", EIGEN_COST(scalar_sum_op<float>));
elementwise_ops_.emplace("AddV2", EIGEN_COST(scalar_sum_op<float>));

View File

@ -939,24 +939,29 @@ TEST_F(OpLevelCostEstimatorTest, SquaredDifferenceExecutionTime) {
EXPECT_EQ(cost.num_ops_with_unknown_shapes, 0);
}
TEST_F(OpLevelCostEstimatorTest, ReluExecutionTime) {
auto cost = PredictCosts(DescribeUnaryOp("Relu", 1000));
EXPECT_EQ(Costs::Duration(800), cost.memory_time);
EXPECT_EQ(Costs::Duration(100), cost.compute_time);
EXPECT_EQ(Costs::Duration(900), cost.execution_time);
EXPECT_EQ(1, cost.num_ops_total);
EXPECT_FALSE(cost.inaccurate);
EXPECT_EQ(0, cost.num_ops_with_unknown_shapes);
}
TEST_F(OpLevelCostEstimatorTest, UnaryOpExecutionTime) {
std::vector<std::pair<std::string, int>> unary_ops = {
{"All", 1}, {"ArgMax", 1}, {"Cast", 1}, {"Max", 1}, {"Min", 1},
{"Prod", 1}, {"Relu", 1}, {"Relu6", 1}, {"Sum", 1}, {"TopKV2", 1}};
TEST_F(OpLevelCostEstimatorTest, CastExecutionTime) {
auto cost = PredictCosts(DescribeUnaryOp("Cast", 1000));
EXPECT_EQ(Costs::Duration(800), cost.memory_time);
EXPECT_EQ(Costs::Duration(100), cost.compute_time);
EXPECT_EQ(Costs::Duration(900), cost.execution_time);
EXPECT_EQ(1, cost.num_ops_total);
EXPECT_FALSE(cost.inaccurate);
EXPECT_EQ(0, cost.num_ops_with_unknown_shapes);
const int kTensorSize = 1000;
for (auto unary_op : unary_ops) {
OpContext op_context = DescribeUnaryOp(unary_op.first, kTensorSize);
const int kExpectedMemoryTime = 800;
int expected_compute_time = std::ceil(
unary_op.second * kTensorSize /
estimator_.GetDeviceInfo(op_context.op_info.device()).gigaops);
auto cost = PredictCosts(op_context);
EXPECT_EQ(cost.memory_time, Costs::Duration(kExpectedMemoryTime));
EXPECT_EQ(cost.compute_time, Costs::Duration(expected_compute_time));
EXPECT_EQ(cost.execution_time,
Costs::Duration(expected_compute_time + kExpectedMemoryTime));
EXPECT_EQ(cost.num_ops_total, 1);
EXPECT_EQ(cost.num_ops_with_unknown_shapes, 0);
EXPECT_FALSE(cost.inaccurate);
}
}
TEST_F(OpLevelCostEstimatorTest, BroadcastAddExecutionTime) {