From fa75f1142e92d96ff79d9d0f8fbee2d9af6b0539 Mon Sep 17 00:00:00 2001
From: Vignesh Kothapalli <vikoth18@in.ibm.com>
Date: Thu, 7 Jan 2021 03:50:16 +0530
Subject: [PATCH] pass references instead of copies

---
 tensorflow/core/kernels/data/iterator_ops.cc | 12 ++++++------
 tensorflow/core/kernels/data/iterator_ops.h  | 14 +++++++-------
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/tensorflow/core/kernels/data/iterator_ops.cc b/tensorflow/core/kernels/data/iterator_ops.cc
index 7bce3850a09..2754e69b847 100644
--- a/tensorflow/core/kernels/data/iterator_ops.cc
+++ b/tensorflow/core/kernels/data/iterator_ops.cc
@@ -116,10 +116,10 @@ Status IteratorResource::GetNext(OpKernelContext* ctx,
   IteratorContext::Params params(ctx);
   params.flr = captured_state->GetFLR();
   params.function_handle_cache = captured_state->GetFunctionHandleCache();
-  params.resource_mgr = &captured_state->GetResourceMgr();
+  params.resource_mgr = captured_state->GetResourceMgr();
   params.thread_factory = unbounded_thread_pool_.get_thread_factory();
   params.thread_pool = &unbounded_thread_pool_;
-  params.cancellation_manager = &captured_state->GetCancellationManager();
+  params.cancellation_manager = captured_state->GetCancellationManager();
   std::function<void()> deregister_fn;
   TF_RETURN_IF_ERROR(RegisterCancellationCallback(
       ctx->cancellation_manager(),
@@ -204,10 +204,10 @@ Status IteratorResource::Restore(OpKernelContext* ctx,
   IteratorContext::Params params(ctx);
   params.flr = new_state->GetFLR();
   params.function_handle_cache = new_state->GetFunctionHandleCache();
-  params.resource_mgr = &new_state->GetResourceMgr();
+  params.resource_mgr = new_state->GetResourceMgr();
   params.thread_factory = unbounded_thread_pool_.get_thread_factory();
   params.thread_pool = &unbounded_thread_pool_;
-  params.cancellation_manager = &new_state->GetCancellationManager();
+  params.cancellation_manager = new_state->GetCancellationManager();
   std::function<void()> deregister_fn;
   TF_RETURN_IF_ERROR(RegisterCancellationCallback(
       ctx->cancellation_manager(),
@@ -239,10 +239,10 @@ Status IteratorResource::SetIteratorFromDataset(OpKernelContext* ctx,
   IteratorContext::Params params(ctx);
   params.flr = new_state->GetFLR();
   params.function_handle_cache = new_state->GetFunctionHandleCache();
-  params.resource_mgr = &new_state->GetResourceMgr();
+  params.resource_mgr = new_state->GetResourceMgr();
   params.thread_factory = unbounded_thread_pool_.get_thread_factory();
   params.thread_pool = &unbounded_thread_pool_;
-  params.cancellation_manager = &new_state->GetCancellationManager();
+  params.cancellation_manager = new_state->GetCancellationManager();
   std::function<void()> deregister_fn;
   TF_RETURN_IF_ERROR(RegisterCancellationCallback(
       ctx->cancellation_manager(),
diff --git a/tensorflow/core/kernels/data/iterator_ops.h b/tensorflow/core/kernels/data/iterator_ops.h
index c2c08eb3f22..6b20d7ea756 100644
--- a/tensorflow/core/kernels/data/iterator_ops.h
+++ b/tensorflow/core/kernels/data/iterator_ops.h
@@ -93,23 +93,23 @@ class IteratorResource : public ResourceBase {
       iterator.reset(static_cast<DatasetBaseIterator*>(it.release()));
     }
 
-    FunctionLibraryDefinition* GetFlibDef() { return flib_def; }
+    std::shared_ptr<FunctionLibraryDefinition> GetFlibDef() { return flib_def; }
 
     FunctionLibraryRuntime* GetFLR() { return flr; }
 
-    ProcessFunctionLibraryRuntime* GetPFLR() { return pflr; }
+    std::shared_ptr<ProcessFunctionLibraryRuntime> GetPFLR() { return pflr; }
 
-    FunctionHandleCache GetFunctionHandleCache() {
+    FunctionHandleCache* GetFunctionHandleCache() {
       return function_handle_cache.get();
     }
 
-    ResourceMgr GetResourceMgr() { return resource_mgr; }
+    ResourceMgr* GetResourceMgr() { return &resource_mgr; }
 
-    CancellationManager GetCancellationManager() {
-      return cancellation_manager;
+    CancellationManager* GetCancellationManager() {
+      return &cancellation_manager;
     }
 
-    DatasetBaseIterator* GetIterator() { return iterator; }
+    DatasetBaseIterator* GetIterator() { return iterator.get(); }
 
    private:
     std::shared_ptr<FunctionLibraryDefinition> flib_def;