Method ComputeAsync

tensorflow/core/kernels/data/map_defun_op.cc:179–238

}

void MapDefunOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
  ComputeOptions* compute_opts = nullptr;

  OP_REQUIRES_OK_ASYNC(ctx, SetupArgs(ctx, &compute_opts), done);

  Status s = SetupOutputs(ctx, compute_opts);
  if (!s.ok()) delete compute_opts;
  OP_REQUIRES_OK_ASYNC(ctx, s, done);

  FunctionLibraryRuntime::Options opts;
  SetRunOptions(ctx, &opts, compute_opts, /*always_collect_stats=*/false);

  // Run loop
  StatusCallback callback = std::bind(
      [](OpKernelContext* ctx, ComputeOptions* compute_opts, DoneCallback& done,
         const Status& status) {
        delete compute_opts;
        ctx->SetStatus(status);
        done();
      },
      ctx, compute_opts, std::move(done), std::placeholders::_1);

  auto* refcounted = new ReffedStatusCallback(std::move(callback));

  CancellationManager* parent_mgr = ctx->cancellation_manager();

  for (size_t i = 0; i < static_cast<size_t>(compute_opts->batch_size); ++i) {
    // We use a different cancellation manager each time the function is run
    // to avoid the race condition between a function run error and other
    // functions being cancelled as a result.
    CancellationManager* c_mgr = new CancellationManager();
    CancellationToken token = parent_mgr->get_cancellation_token();
    const bool success = parent_mgr->RegisterCallback(
        token, [c_mgr]() { c_mgr->StartCancel(); });

    opts.cancellation_manager = c_mgr;
    if (!success) {
      delete c_mgr;
      refcounted->UpdateStatus(errors::Cancelled(
          "MapDefunOp functions cancelled because parent graph cancelled"));
      break;
    }

    auto* call_frame = new MapFunctionCallFrame(compute_opts, this, i);

    refcounted->Ref();
    ctx->function_library()->Run(opts, func_handle_, call_frame,
                                 [call_frame, refcounted, c_mgr, parent_mgr,
                                  token](const Status& func_status) {
                                   parent_mgr->DeregisterCallback(token);
                                   delete c_mgr;
                                   delete call_frame;
                                   refcounted->UpdateStatus(func_status);
                                   refcounted->Unref();
                                 });
  }

  // Unref 1 because refcounted is initialized with refcount = 1