aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/grappler/costs/op_level_cost_estimator.h
diff options
context:
space:
mode:
authorGravatar Yao Zhang <yaozhang@google.com>2017-06-09 17:45:58 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2017-06-09 17:50:06 -0700
commit19b4ccd95fa667adf6240d24368022594f320e73 (patch)
tree077038a26ec4fbdf70d76cec519cbf3d81d5ed85 /tensorflow/core/grappler/costs/op_level_cost_estimator.h
parentef8b5fd1876ba2be3a10a363e121e6d2005bb480 (diff)
Estimate cost for element wise ops and a minimum compute cost for dummy execution time.
PiperOrigin-RevId: 158585464
Diffstat (limited to 'tensorflow/core/grappler/costs/op_level_cost_estimator.h')
-rw-r--r--tensorflow/core/grappler/costs/op_level_cost_estimator.h21
1 files changed, 16 insertions, 5 deletions
diff --git a/tensorflow/core/grappler/costs/op_level_cost_estimator.h b/tensorflow/core/grappler/costs/op_level_cost_estimator.h
index ec7f21622f..28d49a7703 100644
--- a/tensorflow/core/grappler/costs/op_level_cost_estimator.h
+++ b/tensorflow/core/grappler/costs/op_level_cost_estimator.h
@@ -89,17 +89,26 @@ class OpLevelCostEstimator {
ConvolutionDimensions* conv_info,
bool* found_unknown_shapes) const;
- // Calculate the total size in bytes of a single input to a TensorFlow op.
- int64 CalculateSingleInputSize(const OpInfo::TensorProperties& input,
- bool* found_unknown_shapes) const;
+ // Calculate the element count of an input/output tensor.
+ int64 CalculateTensorElementCount(const OpInfo::TensorProperties& tensor,
+ bool* found_unknown_shapes) const;
+
+ // Calculate the total size in bytes of an input/output tensor.
+ int64 CalculateTensorSize(const OpInfo::TensorProperties& tensor,
+ bool* found_unknown_shapes) const;
+
+ // Calculate the element count of the largest
+ // input of specified TensorFlow op.
+ int64 CalculateLargestInputCount(const OpInfo& op_features,
+ bool* found_unknown_shapes) const;
// Calculate the total size in bytes of the all
- // the inputs of specified TensorFlow Op
+ // the inputs of specified TensorFlow op.
int64 CalculateInputSize(const OpInfo& op_features,
bool* found_unknown_shapes) const;
// Calculate the total size in bytes of the all
- // the outputs of specified TensorFlow Op
+ // the outputs of specified TensorFlow op.
int64 CalculateOutputSize(const OpInfo& op_features,
bool* found_unknown_shapes) const;
@@ -114,6 +123,7 @@ class OpLevelCostEstimator {
// execution_time is optional, depending on the
// device.
Costs PredictConv2D(const OpInfo& op_features) const;
+ Costs PredictCwiseOp(const OpInfo& op_features) const;
Costs PredictConv2DBackPropInput(const OpInfo& op_features) const;
Costs PredictConv2DBackPropFilter(const OpInfo& op_features) const;
Costs PredictMatMul(const OpInfo& op_features) const;
@@ -136,6 +146,7 @@ class OpLevelCostEstimator {
bool* found_unknown_shapes);
protected:
+ std::map<string, int> elementwise_ops_;
typedef std::function<Costs(const OpInfo& op_feature)> CostImpl;
std::map<string, CostImpl> device_cost_impl_;