message: Only initialize a single CPU as needed for eager context in folding
diff: @@ -128,21 +128,66 @@ LogicalResult ConstantFoldFallbackHook(
  // should instead be per module/set from the Graph being executed in TF (if
  // any) so that the value of variables in the context could be read.
  // Note: Sharing the context is fine as ops are side-effect free.
  - auto initialize = []() {
  - TF_Status* status = TF_NewStatus();
  - // The TFE_Context is created without an accompanying delete due to current
  - // lifetime. This does not result in memory leaks reported (see totw/110).
  - TFE_ContextOptions* opts = TFE_NewContextOptions();
  + auto initialize = []() -> TFE_Context* {
  + std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
  + TF_NewStatus(), TF_DeleteStatus);
  + std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)>
  + opts(TFE_NewContextOptions(), TFE_DeleteContextOptions);
  + // Only initialize single CPU.
  + tensorflow::ConfigProto config_proto;
  + // This is conceptually equal to what we do in python/eager/context.py but
  + // with all GPU devices ignored and CPU only set to 1.
  + (*config_proto.mutable_device_count())["CPU"] = 1;
  + (*config_proto.mutable_device_count())["GPU"] = 0;
  + std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
  + TF_NewBuffer(), TF_DeleteBuffer);
  + DCHECK(config->data == nullptr);
  +
  + // Copy config_proto into config.
  + {
  + const size_t proto_size = config_proto.ByteSizeLong();
  + void* buf = tensorflow::port::Malloc(proto_size);
  + if (buf == nullptr) {
  + LOG(ERROR) << "Failed to allocate memory to serialize ConfigProto "
  + "while creating context options for constant folding";
  + return nullptr;
  + }
  + if (!config_proto.SerializeWithCachedSizesToArray(
  + static_cast<uint8_t*>(buf))) {
  + tensorflow::port::Free(buf);
  + LOG(ERROR) << "Unable to serialize ConfigProto while creating context "
  + "options for constant folding";
  + return nullptr;
  + }
  + config->data = buf;
  + config->length = proto_size;
  + config->data_deallocator = [](void* data, size_t length) {
  + tensorflow::port::Free(data);
  + };
  + }
  +
  + TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
  + status.get());
  + if (TF_GetCode(status.get()) != TF_OK) {
  + LOG(ERROR) << "Failed to set context options for constant folding: "
  + << status.get();
  + return nullptr;
  + }
  +
  // Input tensors are placed on the host CPU so use the explicit device
  // policy to fail if no CPU kernels are available for the op.
  - TFE_ContextOptionsSetDevicePlacementPolicy(opts,
  + TFE_ContextOptionsSetDevicePlacementPolicy(opts.get(),
  TFE_DEVICE_PLACEMENT_EXPLICIT);
  - auto ctx = TFE_NewContext(opts, status);
  - TFE_DeleteContextOptions(opts);
  - TF_DeleteStatus(status);
  + auto ctx = TFE_NewContext(opts.get(), status.get());
  + if (TF_GetCode(status.get()) != TF_OK) {
  + LOG(ERROR) << "Failed to create context for constant folding: "
  + << status.get();
  + return nullptr;
  + }
  return ctx;
  };
  static TFE_Context* ctx = initialize();
  + if (!ctx) return failure();
  // Returns directly if any of the operands is not an elements attributes.
  if (std::any_of(operands.begin(), operands.end(), [](Attribute attr) {
  