pass references instead of copies

This commit is contained in:
Vignesh Kothapalli 2021-01-07 03:50:16 +05:30
parent 45ed98546c
commit fa75f1142e
No known key found for this signature in database
GPG Key ID: 7C9D6956FBA21DD7
2 changed files with 13 additions and 13 deletions

View File

@ -116,10 +116,10 @@ Status IteratorResource::GetNext(OpKernelContext* ctx,
IteratorContext::Params params(ctx); IteratorContext::Params params(ctx);
params.flr = captured_state->GetFLR(); params.flr = captured_state->GetFLR();
params.function_handle_cache = captured_state->GetFunctionHandleCache(); 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_factory = unbounded_thread_pool_.get_thread_factory();
params.thread_pool = &unbounded_thread_pool_; params.thread_pool = &unbounded_thread_pool_;
params.cancellation_manager = &captured_state->GetCancellationManager(); params.cancellation_manager = captured_state->GetCancellationManager();
std::function<void()> deregister_fn; std::function<void()> deregister_fn;
TF_RETURN_IF_ERROR(RegisterCancellationCallback( TF_RETURN_IF_ERROR(RegisterCancellationCallback(
ctx->cancellation_manager(), ctx->cancellation_manager(),
@ -204,10 +204,10 @@ Status IteratorResource::Restore(OpKernelContext* ctx,
IteratorContext::Params params(ctx); IteratorContext::Params params(ctx);
params.flr = new_state->GetFLR(); params.flr = new_state->GetFLR();
params.function_handle_cache = new_state->GetFunctionHandleCache(); 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_factory = unbounded_thread_pool_.get_thread_factory();
params.thread_pool = &unbounded_thread_pool_; params.thread_pool = &unbounded_thread_pool_;
params.cancellation_manager = &new_state->GetCancellationManager(); params.cancellation_manager = new_state->GetCancellationManager();
std::function<void()> deregister_fn; std::function<void()> deregister_fn;
TF_RETURN_IF_ERROR(RegisterCancellationCallback( TF_RETURN_IF_ERROR(RegisterCancellationCallback(
ctx->cancellation_manager(), ctx->cancellation_manager(),
@ -239,10 +239,10 @@ Status IteratorResource::SetIteratorFromDataset(OpKernelContext* ctx,
IteratorContext::Params params(ctx); IteratorContext::Params params(ctx);
params.flr = new_state->GetFLR(); params.flr = new_state->GetFLR();
params.function_handle_cache = new_state->GetFunctionHandleCache(); 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_factory = unbounded_thread_pool_.get_thread_factory();
params.thread_pool = &unbounded_thread_pool_; params.thread_pool = &unbounded_thread_pool_;
params.cancellation_manager = &new_state->GetCancellationManager(); params.cancellation_manager = new_state->GetCancellationManager();
std::function<void()> deregister_fn; std::function<void()> deregister_fn;
TF_RETURN_IF_ERROR(RegisterCancellationCallback( TF_RETURN_IF_ERROR(RegisterCancellationCallback(
ctx->cancellation_manager(), ctx->cancellation_manager(),

View File

@ -93,23 +93,23 @@ class IteratorResource : public ResourceBase {
iterator.reset(static_cast<DatasetBaseIterator*>(it.release())); iterator.reset(static_cast<DatasetBaseIterator*>(it.release()));
} }
FunctionLibraryDefinition* GetFlibDef() { return flib_def; } std::shared_ptr<FunctionLibraryDefinition> GetFlibDef() { return flib_def; }
FunctionLibraryRuntime* GetFLR() { return flr; } FunctionLibraryRuntime* GetFLR() { return flr; }
ProcessFunctionLibraryRuntime* GetPFLR() { return pflr; } std::shared_ptr<ProcessFunctionLibraryRuntime> GetPFLR() { return pflr; }
FunctionHandleCache GetFunctionHandleCache() { FunctionHandleCache* GetFunctionHandleCache() {
return function_handle_cache.get(); return function_handle_cache.get();
} }
ResourceMgr GetResourceMgr() { return resource_mgr; } ResourceMgr* GetResourceMgr() { return &resource_mgr; }
CancellationManager GetCancellationManager() { CancellationManager* GetCancellationManager() {
return cancellation_manager; return &cancellation_manager;
} }
DatasetBaseIterator* GetIterator() { return iterator; } DatasetBaseIterator* GetIterator() { return iterator.get(); }
private: private:
std::shared_ptr<FunctionLibraryDefinition> flib_def; std::shared_ptr<FunctionLibraryDefinition> flib_def;