experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Adds config proto fields for enabling internal autotuning features.

Browse Source 
Adds an integer config attribute to TPUPartitionedCallOp.

PiperOrigin-RevId: 273788960
This commit is contained in:
Tayo Oguntebi 2019-10-09 11:47:55 -07:00 committed by TensorFlower Gardener
parent 97e7214f00
commit a0a83d7447
7 changed files with 28 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/core/ops/tpu_functional_ops.cc
View File

@ -26,6 +26,7 @@ REGISTER_OP("TPUPartitionedCall")
.Attr("Tin: list(type) >= 0")
.Attr("Tout: list(type) >= 0")
.Attr("f: func")
.Attr("autotuner_thresh: int = 0")
.SetShapeFn(shape_inference::UnknownShape);
} // namespace tensorflow

7
tensorflow/core/protobuf/config.proto
View File

@ -560,6 +560,13 @@ message ConfigProto {
// If this option is set to true when a session is created, the
// `RunOptions.output_partition_graphs` options must not be set.
bool disable_output_partition_graphs = 14;
// Minimum number of batches run through the XLA graph before XLA fusion
// autotuner is enabled. Default value of zero disables the autotuner.
//
// The XLA fusion autotuner can improve performance by executing a heuristic
// search on the compiler parameters.
int64 xla_fusion_autotuner_thresh = 15;
};
Experimental experimental = 16;

6
tensorflow/core/protobuf/tpu/compile_metadata.proto
View File

@ -86,4 +86,10 @@ message TPUCompileMetadataProto {
// The location of step markers that XLA compile will instrument.
xla.DebugOptions.StepMarkerLocation step_marker_location = 12;
// Minimum number of batches run through the XLA graph before XLA fusion
// autotuner is enabled. Default value of zero disables the autotuner.
// The XLA fusion autotuner can improve performance by executing a heuristic
// search on the compiler parameters.
int64 xla_fusion_autotuner_thresh = 13;
}

6
tensorflow/tools/api/golden/v1/tensorflow.-config-proto.-experimental.pbtxt
View File

@ -81,6 +81,12 @@ tf_proto {
label: LABEL_OPTIONAL
type: TYPE_BOOL
}
field {
name: "xla_fusion_autotuner_thresh"
number: 15
label: LABEL_OPTIONAL
type: TYPE_INT64
}
reserved_range {
start: 2
end: 3

6
tensorflow/tools/api/golden/v1/tensorflow.-config-proto.pbtxt
View File

@ -204,6 +204,12 @@ tf_proto {
label: LABEL_OPTIONAL
type: TYPE_BOOL
}
field {
name: "xla_fusion_autotuner_thresh"
number: 15
label: LABEL_OPTIONAL
type: TYPE_INT64
}
reserved_range {
start: 2
end: 3

2
tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
View File

@ -4282,7 +4282,7 @@ tf_module {
}
member_method {
name: "TPUPartitionedCall"
argspec: "args=[\'args\', \'device_ordinal\', \'Tout\', \'f\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
argspec: "args=[\'args\', \'device_ordinal\', \'Tout\', \'f\', \'autotuner_thresh\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'None\'], "
}
member_method {
name: "TPUReplicateMetadata"

2
tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
View File

@ -4282,7 +4282,7 @@ tf_module {
}
member_method {
name: "TPUPartitionedCall"
argspec: "args=[\'args\', \'device_ordinal\', \'Tout\', \'f\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
argspec: "args=[\'args\', \'device_ordinal\', \'Tout\', \'f\', \'autotuner_thresh\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'None\'], "
}
member_method {
name: "TPUReplicateMetadata"