@annarev @penpornk @yisitu This is the second PR which is for subdevice type support in TensorFlow, please help to review it. Thanks

@jzhoulon Can you please check @annarev's comments and keep us posted ? Thanks!

@gbaned , sorry to late response, we are in a long public holiday last week, will address all comments soon. Thanks!

@annarev @gbaned all the comments are addressed, please help to review. Thanks.

@annarev It seems that device_factory moved back to core/common_runtime and cause the conflict, should we add the header dependency back? thanks

@annarev It seems that device_factory moved back to core/common_runtime and cause the conflict, should we add the header dependency back? thanks

Yep, sorry for confusion. I had to roll it back due to a build failure. However, it has been fixed today and re-committed: c60ced5

@jzhoulon Can you please check @annarev's comments and keep us posted ? Thanks!

Thanks @sanjoy to bring the idea and discussion, seems alias device type is still the better way we can see currently, removing alias device type will bring some additional complex design, like eager support (need model change, not transparent to user, also front end "GPU" device redirect to pluggable device need to be added), graph optimization execution is also needed to change. Any idea?

@annarev @penpornk Do you have any comments?

The other concern is that since the grappler pass for Modular TF Graph happens at the last step

Separately, have you considered putting the modular TF Graph pass at the first step of the pipeline? That will probably be more future proof -- as we convert more things in the pipeline to use MLIR, it would be harder to have a GraphDef representation of a TF graph that has gone through some stages of optimization. OTOH we will have to support SavedModel for a long time to come anyway, so having a GraphDef representation for in the first stage of the optimization pipeline is probably OK.

CC @jpienaar

Yeah, so if this were to run as part of import passes then yes you'd still have GraphDef (well for SavedModel you would at loading, during execution you have the similar but different tensorflow::Graph). As GraphDef gets used less for internal graph representation and optimizations, we can avoid extra overhead of converting to GraphDef - IMHO we already convert too often between the two formats, and I'd rather iterate to where we only have GraphDef as resting input format.

The other concern is that since the grappler pass for Modular TF Graph happens at the last step

Separately, have you considered putting the modular TF Graph pass at the first step of the pipeline? That will probably be more future proof -- as we convert more things in the pipeline to use MLIR, it would be harder to have a GraphDef representation of a TF graph that has gone through some stages of optimization. OTOH we will have to support SavedModel for a long time to come anyway, so having a GraphDef representation for in the first stage of the optimization pipeline is probably OK.
CC @jpienaar

Yeah, so if this were to run as part of import passes then yes you'd still have GraphDef (well for SavedModel you would at loading, during execution you have the similar but different tensorflow::Graph). As GraphDef gets used less for internal graph representation and optimizations, we can avoid extra overhead of converting to GraphDef - IMHO we already convert too often between the two formats, and I'd rather iterate to where we only have GraphDef as resting input format.

Putting the modular TF Graph pass at the first step of the pipeline won't work for pluggable devices. The pluggable devices need to get a chance to conduct target-depedent optimization, which may invalidate target-independent grappler passes.

What's the underlying reason to alias over the existing name? I don't think we encode any particular meaning in the fact that something is a "GPU" vs "TPU" vs "XPU".

What's the underlying reason to alias over the existing name? I don't think we encode any particular meaning in the fact that something is a "GPU" vs "TPU" vs "XPU".

One reason that comes to mind is conditional checks. We might be checking if something is GPU to run general GPU logic, or we might be checking that something is GPU to run CUDA specific logic. We could probably differentiate whether something is a pluggable device or not instead (if pluggable, don't run CUDA-specific logic).

Another reason is supporting running code only on specific type of GPU. For e.g having with tf.device('/INTEL_GPU:0') (say, fail if INTEL_GPU is not available in this case). I don't know though if there is a known usecase for this though.

What's the underlying reason to alias over the existing name? I don't think we encode any particular meaning in the fact that something is a "GPU" vs "TPU" vs "XPU".

One reason that comes to mind is conditional checks. We might be checking if something is GPU to run general GPU logic, or we might be checking that something is GPU to run CUDA specific logic. We could probably differentiate whether something is a pluggable device or not instead (if pluggable, don't run CUDA-specific logic).

Another reason is supporting running code only on specific type of GPU. For e.g having with tf.device('/INTEL_GPU:0') (say, fail if INTEL_GPU is not available in this case). I don't know though if there is a known usecase for this though.

And also, alias over existing name will give pluggable device the opportunity to run legacy model (with tf.device('/gpu:0')) without changing any user code.

I just want to clarify: is aliasing necessary for custom devices to work at all? Or is this a convenience to support existing models? (I don't think it's required, because things like TPUs work, but please correct me if I'm wrong). Existing code that checks for "device:GPU" is likely specialized for CUDA GPUs, so I don't think it's correct in general to assume we can alias over it.

If it's not a hard requirement, I'd recommend we move ahead with the other PRs related to pluggable devices until we understand this a bit better.

I just want to clarify: is aliasing necessary for custom devices to work at all? Or is this a convenience to support existing models? (I don't think it's required, because things like TPUs work, but please correct me if I'm wrong). Existing code that checks for "device:GPU" is likely specialized for CUDA GPUs, so I don't think it's correct in general to assume we can alias over it.

If it's not a hard requirement, I'd recommend we move ahead with the other PRs related to pluggable devices until we understand this a bit better.

I would say yes, we do see a need for us to run the existing TF application with "device:GPU" as is. GPU has a broad market so we expect more entry-level users and we want to enable them transparently. If TF's "device:GPU" implies CUDA, it is against the spirit of modular TF, which intends to provide a vendor-neutral interface. We are working on CUDA specific code to make it more general.

I just want to clarify: is aliasing necessary for custom devices to work at all? Or is this a convenience to support existing models? (I don't think it's required, because things like TPUs work, but please correct me if I'm wrong). Existing code that checks for "device:GPU" is likely specialized for CUDA GPUs, so I don't think it's correct in general to assume we can alias over it.
If it's not a hard requirement, I'd recommend we move ahead with the other PRs related to pluggable devices until we understand this a bit better.

I would say yes, we do see a need for us to run the existing TF application with "device:GPU" as is. GPU has a broad market so we expect more entry-level users and we want to enable them transparently. If TF's "device:GPU" implies CUDA, it is against the spirit of modular TF, which intends to provide a vendor-neutral interface. We are working on CUDA specific code to make it more general.

I joined the discussion late, so please correct me if I missed some context here. As @Jianhui-Li mentioned it would be great if can enable existing TF applications on GPU seamlessly. As we are adding Metal backend support we noticed this dependency as well and as part of pluggable design it gives an opportunity to straighten it out.

@rjpower Is the concern here that the change has wider impact and will break existing code-paths in DEVICE_GPU?

@rjpower Any update on this PR? Please. Thanks!

@kulinseth Sorry for the late reply! I believe the concern here is that whatever we choose here (having a device_alias or not), it might restrict how TensorFlow does device placement in the future. And there have been several existing issues with device placement to begin with. So some of us wanted to explore the impact on that side more before we proceed here. This PR is not blocking other PluggableDevice PRs and will likely be merged last. (We are still aiming for TF 2.5.)

@gbaned This PR will be on hold for a while until we can make a decision on how we want to move forward.

@penpornk Thanks for the reply. Makes sense and since this is not blocking other PRs, we can take a look later and not hold them up.

(Removing the API review tag for now. Will add it back once the PR is ready for API review again.)

@penpornk Any update on this PR? Please. Thanks!

@gbaned It's still on hold. Sorry!

@penpornk Any update on this PR? Please. Thanks!

@gbaned We are exploring an alternative. I'll temporarily close this PR. We can reopen it if later on if we want to fall back to this one. Sorry for the inconveniences!