Use upstream `BackendDevice` instead of inherit from it

[JackCaoG]

🚀 Feature

Currently we inherit a Device from torch::lazy::BackendDevice but that should be unneceaary.

xla/torch_xla/csrc/device.h

Line 24 in 94d7165

struct Device : public torch::lazy::BackendDevice {

We should use BackendDevice directly in PyTorch/XLA.

Motivation

Going forward we will need to use functions like torch::lazy::GetBackendDevice which returns a c10::optional<BackendDevice>

This is also blocking us from using code-gen.

@alanwaketan @wonjoolee95

[alanwaketan]

Yea, I believe the discussion was about XLA inheriting from BackendDeviceType instead of BackendDevice.

[wonjoo-wj]

One discrepancy I noticed that the torch_xla::ir::Device (https://github.com/pytorch/xla/blob/master/torch_xla/csrc/device.h#L26) has a constructor:

explicit Device(const std::string& device_spec);

This constructor parses the device_spec string, but I noticed that this is marked as deprecated and is umimplemented in the torch::lazy::Device (https://github.com/pytorch/pytorch/blob/master/torch/csrc/lazy/backend/backend_device.h#L48) constructor:

// The string -> Device conversion should be handled by the backend interface.
C10_DEPRECATED explicit BackendDevice(const std::string& device_spec);

I don't think PT/XLA uses this specific constructor directly, so it should be okay but I'll post updates as I get the PR ready.

cc @JackCaoG

[wonjoo-wj]

Upon working on this PR, there seems to multiple places where we need this string parse constructor (ex1, ex2). We'll add a util function ParseDeviceString that takes in std:str device_spec and returns the appropriate torch::lazy::BackendDevice on our side to continue this use case.

[JackCaoG]

@wonjoolee95 Let's also make some update to upstream device

1. We should remove https://github.com/pytorch/pytorch/blob/2e2cfa2ef77ff95741f7ebc21a24846567d8630b/torch/csrc/lazy/backend/backend_device.h#L48 which does not do anything and it caused some very hard to debug issue on our end.
2. I want to init the ordinal_ to -1 instead of 0, this way it will fail early if we ever accidentally call the default constructor.
3. I want to add a is_valid function which pretty much checks if ordinal_ >= 0. This will be helpful in some our pt/xla use case where we want to check if device is valid.

@alanwaketan Do these changes looks ok to you?

[alanwaketan]

@wonjoolee95 Let's also make some update to upstream device

1. We should remove https://github.com/pytorch/pytorch/blob/2e2cfa2ef77ff95741f7ebc21a24846567d8630b/torch/csrc/lazy/backend/backend_device.h#L48 which does not do anything and it caused some very hard to debug issue on our end.
2. I want to init the ordinal_ to -1 instead of 0, this way it will fail early if we ever accidentally call the default constructor.
3. I want to add a is_valid function which pretty much checks if ordinal_ >= 0. This will be helpful in some our pt/xla use case where we want to check if device is valid.

@alanwaketan Do these changes looks ok to you?

I think you basically need a way to set the default ordinal for XLA. I remember that we discussed before that we can provide a backend interface for you to overwrite the ordinal. Do you still love this idea? The reason for that is I want to have BackendDevice() to always construct an object that represents the default setting tailored to each backend.

On the other hand, ‘-1’ might also work. Let me double check if it breaks anything.

[alanwaketan]

Looks like -1 won’t break anything. Then, let’s go ahead with the -1 change.