[Kernel C API] Implementation of variable ops RFC.

[kulinseth]

The implementation for the Variable Ops RFC.
https://github.com/tensorflow/community/blob/master/rfcs/20210504-kernel-extension-variable-ops.md

@penpornk , @reedwm , @saxenasaurabh , @jzhoulon

[kulinseth]

I am moving these to c_api_experimental file (ran into some build issues). I wanted to get started with PR, to get feedback.

[jzhoulon]

I think at the end of the function, TF_DeleteTensor() is needed, TF_TensorFromTensor will new a TF_Tensor struct and will cause memory leak if TF_DeleteTensor() is not invoked.

/ Non-static for testing.
TF_Tensor* TF_TensorFromTensor(const tensorflow::Tensor& src, Status* status) {
  *status = tensorflow::Status::OK();
  if (!src.IsInitialized()) {
    *status = FailedPrecondition(
        "attempt to use a tensor with an uninitialized value");
    return nullptr;
  }
  if (src.NumElements() == 0) {
    return EmptyTensor(static_cast<TF_DataType>(src.dtype()), src.shape());
  }
  if (src.dtype() == tensorflow::DT_RESOURCE) {
    if (src.shape().dims() != 0) {
      *status = InvalidArgument(
          "Unexpected non-scalar DT_RESOURCE tensor seen (shape: ",
          src.shape().DebugString(),
          "). Please file a bug at "
          "https://github.com/tensorflow/tensorflow/issues/new, "
          "ideally with a "
          "short code snippet that reproduces this error.");
      return nullptr;
    }
    const string str =
        src.scalar<tensorflow::ResourceHandle>()().SerializeAsString();
    TF_Tensor* t = TF_AllocateTensor(TF_RESOURCE, {}, 0, str.size());
    std::memcpy(TF_TensorData(t), str.c_str(), str.size());
    return t;
  }

  Tensor tensor;
  if (!tensor.CopyFrom(src, src.shape())) {
    return nullptr;
  }
  return new TF_Tensor{new tensorflow::TensorInterface(std::move(tensor))};
}

[kulinseth]

Done.

[yanzhang-dev]

Do we need to check the two status ?

[kulinseth]

Maybe we can return the status back to caller of the PrepareUpdateVariable().

[saxenasaurabh]

nit: Rename values to dims and max_vals to num_dims to be consistent with TF_GraphGetTensorShape.

[kulinseth]

Done.

[saxenasaurabh]

Is it possible to share code with the existing impl for EnsureSparseVariableAccess in training_op_helpers.h? Same comment for other helpers below.

[kulinseth]

Thanks @saxenasaurabh for the review. I have used the helper functions where possible like LookupResource. With EnsureSparseVariableAccess , there is Device dependency which we are passing in as Copy functors. I do agree there is duplication of code which can be avoided. One possibility is to refactor the core helper functions to remove this dependency and we can adopt it here. Maybe we can do it as a followup cleanup as it will require invasive changes in the core which can break things. What do you think ?

[saxenasaurabh]

Sounds good. Please try to clean this is up as a follow-up. That would give test coverage for free as well.

[kulinseth]

Makes sense, will do.

[kulinseth]

We are expecting plugin to call TF_DeleteTensor in copyFunc. This would be more in line with rest of the TF, where we release the tensors in the Compute.

[saxenasaurabh]

Do we need to expose isVariantType in the API or can that be inferred from the tensor? I believe this is equivalent to TF_TensorType(tensor) == TF_VARIANT.

[kulinseth]

@saxenasaurabh thanks! (sorry for the delay), you are right , we can probably skip the isVariantType in the API. Currently in our metal-plugin release we are using this API. If its not too much of a concern, can we keep this way?

[saxenasaurabh]

Could this be in c_api_experimental as well?