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

Resubmitting the change to make iterator deleter not on host memory

Browse Source 
PiperOrigin-RevId: 310590461
Change-Id: Iea22ee212637252f66ccc9bc1e96dffe8b3cc3a7
This commit is contained in:
Alexandre Passos 2020-05-08 11:07:04 -07:00 committed by TensorFlower Gardener
parent a37a3569f4
commit 54fa7e4419
3 changed files with 10 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/compiler/jit/xla_device_ops.h
View File

@ -180,11 +180,9 @@ class XlaAssignVariableOp : public OpKernel {
data::MakeIteratorOp); \ data::MakeIteratorOp); \
REGISTER_KERNEL_BUILDER(Name("AnonymousIterator").Device(DEVICE), \ REGISTER_KERNEL_BUILDER(Name("AnonymousIterator").Device(DEVICE), \
data::AnonymousIteratorHandleOp); \ data::AnonymousIteratorHandleOp); \
REGISTER_KERNEL_BUILDER( \ REGISTER_KERNEL_BUILDER(Name("AnonymousIteratorV2").Device(DEVICE), \
Name("AnonymousIteratorV2").Device(DEVICE).HostMemory("deleter"), \
data::AnonymousIteratorHandleOp); \ data::AnonymousIteratorHandleOp); \
REGISTER_KERNEL_BUILDER( \ REGISTER_KERNEL_BUILDER(Name("DeleteIterator").Device(DEVICE), \
Name("DeleteIterator").Device(DEVICE).HostMemory("deleter"), \
data::DeleteIteratorOp); \ data::DeleteIteratorOp); \
REGISTER_KERNEL_BUILDER(Name("IteratorGetNext").Device(DEVICE), \ REGISTER_KERNEL_BUILDER(Name("IteratorGetNext").Device(DEVICE), \
data::IteratorGetNextOp); \ data::IteratorGetNextOp); \

5
tensorflow/core/kernels/data/dataset_utils.h
View File

@ -63,7 +63,10 @@ class AnonymousResourceOp : public OpKernel {
if (create_deleter_) { if (create_deleter_) {
Tensor* deleter_t; Tensor* deleter_t;
OP_REQUIRES_OK(ctx, ctx->allocate_output(1, TensorShape({}), &deleter_t)); AllocatorAttributes attr;
attr.set_on_host(true);
OP_REQUIRES_OK(
ctx, ctx->allocate_output(1, TensorShape({}), &deleter_t, attr));
deleter_t->scalar<Variant>()() = deleter_t->scalar<Variant>()() =
ResourceDeleter(handle, ctx->resource_manager()); ResourceDeleter(handle, ctx->resource_manager());
} }

4
tensorflow/core/kernels/data/iterator_ops.cc
View File

@ -1102,8 +1102,7 @@ REGISTER_KERNEL_BUILDER(
MakeIteratorOp); MakeIteratorOp);
REGISTER_KERNEL_BUILDER(Name("DeleteIterator").Device(DEVICE_CPU).Priority(2), REGISTER_KERNEL_BUILDER(Name("DeleteIterator").Device(DEVICE_CPU).Priority(2),
DeleteIteratorOp); DeleteIteratorOp);
REGISTER_KERNEL_BUILDER( REGISTER_KERNEL_BUILDER(Name("DeleteIterator").Device(DEVICE_GPU).Priority(1),
Name("DeleteIterator").Device(DEVICE_GPU).HostMemory("deleter").Priority(1),
DeleteIteratorOp); DeleteIteratorOp);
REGISTER_KERNEL_BUILDER( REGISTER_KERNEL_BUILDER(
Name("AnonymousIterator").Device(DEVICE_CPU).Priority(2), Name("AnonymousIterator").Device(DEVICE_CPU).Priority(2),
@ -1116,7 +1115,6 @@ REGISTER_KERNEL_BUILDER(
AnonymousIteratorHandleOp); AnonymousIteratorHandleOp);
REGISTER_KERNEL_BUILDER(Name("AnonymousIteratorV2") REGISTER_KERNEL_BUILDER(Name("AnonymousIteratorV2")
.Device(DEVICE_GPU) .Device(DEVICE_GPU)
.HostMemory("deleter")
.Priority(1), .Priority(1),
AnonymousIteratorHandleOp); AnonymousIteratorHandleOp);
REGISTER_KERNEL_BUILDER(Name("DatasetToSingleElement").Device(DEVICE_CPU), REGISTER_KERNEL_BUILDER(Name("DatasetToSingleElement").Device(DEVICE_CPU),