Added a bunch of unary ops to the estimator.
PiperOrigin-RevId: 324607213 Change-Id: I24369f36cc29f68caac412a5d3076f5ef43859fe
This commit is contained in:
parent
123af9ba86
commit
8e9f3196fd
@ -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>));
|
||||
|
||||
@ -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) {
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user