Setup KernelFallbackRequestState in c_api_tfrt in order to use KernelFallback in TF-TFRT integration.

PiperOrigin-RevId: 343595096
Change-Id: If828f781337824bd4118963ea3206a846cb8b948
This commit is contained in:
Xiao Yu 2020-11-20 17:30:04 -08:00 committed by TensorFlower Gardener
parent 03ab8d341c
commit 5d8eae5093

View File

@ -94,14 +94,16 @@ class ScopedStepContainer {
// prefix: optional string prefix to disambiguate step containers.
ScopedStepContainer(const int64 step_id,
std::function<void(const string&)> cleanup)
: container_(strings::StrCat("__per_step_", step_id)),
: step_id_(step_id),
container_(strings::StrCat("__per_step_", step_id)),
cleanup_(cleanup),
dirty_(false) {}
ScopedStepContainer(const int64 step_id,
std::function<void(const string&)> cleanup,
const std::string& prefix)
: container_(strings::StrCat("__", prefix, "_per_step_", step_id)),
: step_id_(step_id),
container_(strings::StrCat("__", prefix, "_per_step_", step_id)),
cleanup_(cleanup),
dirty_(false) {}
@ -141,8 +143,10 @@ class ScopedStepContainer {
template <typename T>
Status LookupOrCreate(ResourceMgr* rm, const std::string& name, T** resource,
std::function<Status(T**)> creator) TF_MUST_USE_RESULT;
int64 StepId() const { return step_id_; }
private:
const int64 step_id_;
const std::string container_;
const std::function<void(const string&)> cleanup_;
mutex mu_;