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 

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

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

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

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

@ -63,7 +63,10 @@ class AnonymousResourceOp : public OpKernel {
if (create_deleter_) {
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>()() =
ResourceDeleter(handle, ctx->resource_manager());
}

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

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