RFC: Kernel C API extension for Variable ops.

[kulinseth]

The RFC will be open for comment until Tuesday, May 18, 2021

Status	Proposed
RFC #	387
Author(s)	Kulin Seth (Apple), Charles Brissart (Apple)
Sponsor	Saurabh Saksena (srbs@google.com)
Updated	2021-05-04

The proposal extends the current Kernel C API to enable plugin writers using Modular API to add support for :

• Optimizer operations like SGD, Adam etc.
• Variable updates like Assign, AssignUpdate

[penpornk]

If you are interested in attending the design review meeting for this RFC, here are the call details:

Time: Tuesday, May 18th, 2021 from 10:00-11:00am PT Meeting link: Google Meet

Meeting participants are expected to read the RFC ahead of time as the meeting will focus on the remaining issues/questions.

[ematejska]

This was accepted at the design review meeting today.

[penpornk]

Design review meeting notes

@kulinseth: We started by exposing each helper function for optimizer / Resource variables through C API. Based on @saxenasaurabh's comments, GetInputTensorFromVariable can be used for Resource variables too. We'll refine them in the actual implementation PR.

@saxenasaurabh: How do you scope the set of ops? Do you have a fallback mechanism? @kulinseth: We looked through 14-15 networks and dumped out all the ops. For TensorArray, we fall back to CPU.

@kulinseth: @jzhoulon, what models do you need TensorArray and UnbatchResource for? @jzhoulon / @Jianhui-Li: No specific models. Just for ops coverage / completeness. We looked at TensorArray, Variables, ResourceMgr, BatchResource-related ops. The API you exposed is cleaner.

@kulinseth: Are ResourceManager and TensorArray needed?

• @saxenasaurabh: For Resource variable alone, the low-level API makes sense.
  • TensorArray ops are used mostly in TF 1. We use TensorList* ops in v2.
  • It depends on what models you care about.
• @kulinseth: We have been looking at v2 networks.
  • Rely on DEVICE_DEFAULT for TensorList*.

@kulinseth: In the PR, we will break up input tensor from Variable.

• @saxenasaurabh: Longer term exposing low level mutex / refcounting may be more beneficial.
  • Keep high-level API under c_api_experimental for now.
  • Expand set of ops -- revisit if more low-level stuff is easier to maintain.
  • Doesn't need to support TensorArray and ResourceMgr in this RFC. Can follow-up.

@kulinseth: We recently added a helper function TF_GetInputByName to the RFC.

@Jianhui-Li: TensorArray is replaced by TensorList* in v2. What about UnbatchResource?

• @kulinseth: Also, what about async stuff in general?
• @saxenasaurabh: Not familiar with this (UnbatchResource). Will have to check.
• @reedwm: I don't think v2 usually uses batched ops.
• @kulinseth: Any other ops with async op kernels?
• @reedwm: Some GPU ops perhaps incorrectly use async op kernel to do some work / transfer.
• @saxenasaurabh: How do you handle control flow? Functional control flow ops are async in TF for CPU.
• @kulinseth: We rely on default TF implementation and run the lowered Switch, Merge etc ops.

@kulinseth: We can send out a PR today so everybody can take a closer look.

[kulinseth]

Thanks @penpornk for detailed notes. I will update here with PR.

[penpornk]

Here is @kulinseth's implementation PR: https://github.com/tensorflow/tensorflow/pull/49717