Method ComputeAsync

tensorflow/core/kernels/stack.cc:221–281

}

void StackPushOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
  // Get the stack from the handle.
  Stack* stack = nullptr;
  OP_REQUIRES_OK_ASYNC(ctx, GetStack(ctx, &stack), done);
  core::ScopedUnref unref(stack);

  if (ctx->input_dtype(1) != stack->ElemType()) {
    ctx->CtxFailure(errors::InvalidArgument("Must have type ",
                                            stack->ElemType(), " but got ",
                                            ctx->input_dtype(1)));
    done();
    return;
  }

  // Push the tensor onto the stack. Swap the tensor to CPU if instructed.
  const Tensor& tensor = ctx->input(1);
  AllocatorAttributes alloc_attrs = ctx->input_alloc_attr(1);
  // For now, we use a simple heuristic for swapping: A GPU tensor is moved
  // to CPU if the tensor has more than kCopyThreshold bytes and the GPU
  // allocator says more than kOccupancy of the memory is in use.
  static constexpr int kCopyThreshold = 2048;
  static constexpr double kOccupancy = 0.7;
  if (swap_memory_ && !alloc_attrs.on_host() &&
      tensor.TotalBytes() > kCopyThreshold && stack->IsUsefulToSwap(tensor)) {
    DeviceContext* device_ctxt = ctx->op_device_context();
    auto device = static_cast<tensorflow::Device*>(ctx->device());
    Allocator* allocator = device->GetAllocator(alloc_attrs);
    absl::optional<AllocatorStats> stats = allocator->GetStats();
    if (stats && *stats->bytes_limit &&
        stats->bytes_in_use > (*stats->bytes_limit * kOccupancy)) {
      // Asynchronously copy the tensor from GPU to CPU memory.
      // TODO(yuanbyu): Swap the oldest tensor first.
      AllocatorAttributes host_alloc_attrs;
      host_alloc_attrs.set_gpu_compatible(true);
      host_alloc_attrs.set_on_host(true);
      Allocator* cpu_allocator = device->GetAllocator(host_alloc_attrs);
      Tensor* cpu_tensor =
          new Tensor(cpu_allocator, tensor.dtype(), tensor.shape());
      device_ctxt->CopyDeviceTensorToCPU(
          &tensor, "StackPush", device, cpu_tensor,
          [cpu_tensor, stack, ctx, done](const Status& s) {
            ctx->SetStatus(s);
            if (s.ok()) {
              AllocatorAttributes alloc_attrs = ctx->input_alloc_attr(1);
              ctx->SetStatus(stack->Push({*cpu_tensor, alloc_attrs, true}));
            }
            if (ctx->status().ok()) {
              ctx->set_output(0, *cpu_tensor);
            }
            done();
            delete cpu_tensor;
          });
      return;
    }
  }

  // Execute synchronously if not swapped.
  OP_REQUIRES_OK_ASYNC(ctx, stack->Push({tensor, alloc_attrs, false}), done);