*** Reason for rollback ***

CL316577397 prevents H2D prefetch overlapping for TF2 MLPerf Resnet model on V100x8
More details in b/159372996

  With the rollback, the prefetching overlapping is back, and the training speed also recovered.
*** Original change description ***

Add DT_BOOL support to GPU variable ops

This is a follow-on to PR  & PR  and resolves remaining ask
in Issue . The original PR tried to add many variable ops on the
GPU including DT_BOOL. However, this caused testCondModifyBoolPred to
fail and thus the DT_BOOL type was removed. The reason for the test
failure is once DT_BOOL variables are supported on the GPU, we need to
ensure the switch ops are also updated to not have ho...

***

PiperOrigin-RevId: 317241338
Change-Id: Id7b7d79622e0537ccb677f081b487014ac4d2395
This commit is contained in:
A. Unique TensorFlower 2020-06-18 21:06:35 -07:00 committed by TensorFlower Gardener
parent 0b0eef4031
commit b5bb616121
4 changed files with 13 additions and 12 deletions

View File

@ -111,17 +111,15 @@ REGISTER_GPU_SWITCH(uint64);
TF_CALL_variant(REGISTER_GPU_SWITCH);
TF_CALL_uint32(REGISTER_GPU_SWITCH);
TF_CALL_uint32(REGISTER_GPU_REF_SWITCH);
TF_CALL_bool(REGISTER_GPU_SWITCH);
TF_CALL_bool(REGISTER_GPU_REF_SWITCH);
#undef REGISTER_CPU_SWITCH
#undef REGISTER_CPU_REF_SWITCH
#undef REGISTER_GPU_SWITCH
#undef REGISTER_GPU_REF_SWITCH
// Special GPU kernels for int32, string & resource handles. Requiring all
// inputs and outputs to be in host memory.
// TODO(b/25387198): Also enable int32 in device memory.
// Special GPU kernels for int32 and string.
// TODO(b/25387198): Also enable int32 in device memory. This kernel
// registration requires all int32 inputs and outputs to be in host memory.
#define REGISTER_GPU_HOST_KERNEL(type) \
REGISTER_KERNEL_BUILDER(Name("Switch") \
.Device(DEVICE_GPU) \
@ -151,6 +149,8 @@ TF_CALL_bool(REGISTER_GPU_REF_SWITCH);
REGISTER_GPU_HOST_KERNEL(int32);
REGISTER_GPU_HOST_REF_KERNEL(int32);
REGISTER_GPU_HOST_KERNEL(bool);
REGISTER_GPU_HOST_REF_KERNEL(bool);
REGISTER_GPU_HOST_KERNEL(tstring);
REGISTER_GPU_HOST_REF_KERNEL(tstring);
REGISTER_GPU_HOST_KERNEL(ResourceHandle);

View File

@ -252,7 +252,8 @@ TF_CALL_GPU_NUMBER_TYPES_NO_HALF(REGISTER_SYCL_KERNEL);
TF_CALL_int64(REGISTER_GPU_KERNELS);
TF_CALL_uint32(REGISTER_GPU_KERNELS);
TF_CALL_GPU_ALL_TYPES(REGISTER_GPU_KERNELS);
TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNELS);
TF_CALL_COMPLEX_TYPES(REGISTER_GPU_KERNELS);
#undef REGISTER_GPU_KERNELS
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM

View File

@ -73,9 +73,9 @@ class ReconstructNonDebugGraphTest(test_util.TensorFlowTestCase):
for attr_key in new_node.attr:
if attr_key == "parallel_iterations":
new_node.attr[attr_key].i = 1
elif new_node.op == "Switch" or new_node.op == "Identity":
# We don't check the inputs to Switch or Identity ops as their inputs
# may be Send/Recv nodes.
elif new_node.op == "Switch":
# We don't check the inputs to Switch ops as their inputs may be
# Send/Recv nodes.
del new_node.input[:]
return output_graph_def

View File

@ -396,10 +396,10 @@ class CondTest(test_util.TensorFlowTestCase):
fn2=lambda: math_ops.add(y, 23))
self.assertEquals(self.evaluate(z), 24)
@test_util.run_v1_only("Exercises Ref variables")
@test_util.run_deprecated_v1
def testCondModifyBoolPred(self):
# We want to use the GPU here because we want to ensure that we can update
# a boolean ref variable on the GPU.
# This test in particular used to fail only when running in GPU, hence
# use_gpu=True.
with test_util.use_gpu():
bool_var = variable_scope.get_variable(
"bool_var", dtype=dtypes.bool, initializer=True)