Update hlo cost function visibility

PiperOrigin-RevId: 348050486
Change-Id: I37343bcdfc9d7b86ad110a1fd4cfadb4f42729fb
This commit is contained in:
A. Unique TensorFlower 2020-12-17 10:43:59 -08:00 committed by TensorFlower Gardener
parent e38610399a
commit 1ea36b6b53

View File

@ -183,6 +183,12 @@ class HloCostAnalysis : public ConstDfsHloVisitor {
return GetProperty(key, per_second_rates_);
}
// Return the key that is used to index into Properties for the specified
// input/output at the shape index.
static std::string GetOperandBytesAccessedKey(int64 operand_num,
ShapeIndex index = {});
static std::string GetOutputBytesAccessedKey(ShapeIndex index = {});
protected:
typedef std::unordered_map<const HloInstruction*, Properties> HloToProperties;
@ -229,12 +235,6 @@ class HloCostAnalysis : public ConstDfsHloVisitor {
void SetOutputBytesAccessed(float value);
void SetOutputBytesAccessed(ShapeIndex index, float value);
// Return the key that is used to index into Properties for the specified
// input/output at the shape index.
static std::string GetOperandBytesAccessedKey(int64 operand_num,
ShapeIndex index = {});
static std::string GetOutputBytesAccessedKey(ShapeIndex index = {});
// Function which computes the size of the top-level of a given shape (not
// including nested elements, if any). If null then bytes_accessed methods
// return an error.