Method Compute

tensorflow/compiler/xrt/kernels/xrt_state_ops.h:363–401

  XRTSubTupleOp& operator=(const XRTSubTupleOp&) = delete;

  void Compute(OpKernelContext* ctx) override {
    VLOG(1) << "XRTSubTupleOp::Compute";

    const Tensor& handle_tensor = ctx->input(0);
    OP_REQUIRES(
        ctx, TensorShapeUtils::IsScalar(handle_tensor.shape()),
        errors::Internal("computation input should be an int64 scalar"));
    int64 allocation_handle = handle_tensor.scalar<int64>()();

    const Tensor& subtuple_info = ctx->input(1);
    OP_REQUIRES(
        ctx, TensorShapeUtils::IsVector(subtuple_info.shape()),
        errors::Internal("tuple index input should be an int32 vector"));
    xla::ShapeIndex shape_index;
    for (int i = 0; i < subtuple_info.dim_size(0); ++i) {
      shape_index.push_back(subtuple_info.vec<int32>()(i));
    }

    ResourceMgr* rm;
    OP_REQUIRES_OK(ctx, DeviceAccessor::GetResourceManager(ctx, &rm));

    RefPtr<XRTMemoryManager> memory_manager = XRTMemoryManager::Get(rm);
    RefPtr<XRTTupleAllocation> allocation;
    OP_REQUIRES_OK(ctx, memory_manager->Lookup(allocation_handle, &allocation));

    if (discard_) {
      VLOG(2) << "Releasing handle " << allocation_handle;
      OP_REQUIRES_OK(ctx, memory_manager->Release(allocation_handle));
    }

    XRTTupleAllocation* suballocation;
    OP_REQUIRES_OK(
        ctx, XRTTupleAllocation::MakeSubBuffer(allocation.get(), shape_index,
                                               &suballocation, !discard_));

    Tensor output(DT_INT64, TensorShape({}));
    output.scalar<int64>()() = memory_manager->Register(suballocation);
    ctx->set_output(0, output);
  }
};

// Op that allocates memory for a literal and transfers it to the device.