Update torch::lazy::BackendDevice to have a new default ordinal (pytorch#76264)

Summary:
Fixes pytorch/xla#3490. Updates `torch::lazy::BackendDevice` with changes below:

1. Remove the no-op string constructor.
2. Update default ordinal to `-1`.
3. Add a `is_valid` function to check if `ordinal` is valid/non-default (`ordinal >= 0`).

Pull Request resolved: pytorch#76264

Reviewed By: mrshenli

Differential Revision: D35860266

Pulled By: alanwaketan

fbshipit-source-id: 554ebe16a0683d37b00270c4f35163bf690bfe28

Author: wonjoolee95

test/cpp/lazy/test_backend_device.cpp

@@ -20,8 +20,9 @@ TEST(BackendDeviceTest, Basic1) {
   auto device = BackendDevice();
 
   EXPECT_EQ(device.type(), 0);
-  EXPECT_EQ(device.ordinal(), 0);
-  EXPECT_STREQ(device.toString().c_str(), "Unknown0");
+  EXPECT_EQ(device.ordinal(), -1);
+  EXPECT_FALSE(device.has_index());
+  EXPECT_STREQ(device.toString().c_str(), "Unknown");
 }
 
 TEST(BackendDeviceTest, Basic2) {
@@ -31,6 +32,7 @@ TEST(BackendDeviceTest, Basic2) {
 
   EXPECT_EQ(device.type(), 1);
   EXPECT_EQ(device.ordinal(), 1);
+  EXPECT_TRUE(device.has_index());
   EXPECT_STREQ(device.toString().c_str(), "Unknown1");
 }
 
@@ -43,6 +45,7 @@ TEST(BackendDeviceTest, Basic3) {
 
   EXPECT_EQ(device.type(), 0);
   EXPECT_EQ(device.ordinal(), 1);
+  EXPECT_TRUE(device.has_index());
   EXPECT_STREQ(device.toString().c_str(), "Test1");
 }
 
@@ -86,8 +89,8 @@ TEST(BackendDeviceTest, FromAten) {
 TEST(BackendDeviceTest, ToAten) {
   auto device = backendDeviceToAtenDevice(BackendDevice());
   EXPECT_EQ(device.type(), c10::kLazy);
-  EXPECT_TRUE(device.has_index());
-  EXPECT_EQ(device.index(), 0);
+  EXPECT_FALSE(device.has_index());
+  EXPECT_EQ(device.index(), -1);
 }
 
 // TODO(alanwaketan): Update the following test once we have TorchScript backend upstreamed.

torch/csrc/lazy/backend/backend_device.cpp

@@ -10,24 +10,24 @@ namespace torch {
 namespace lazy {
 
 // TODO(alanwaketan): Use the backend API to get the default device type.
-// In the future, we should also get the default device ordinal.
 BackendDevice::BackendDevice()
-  : type_(std::make_shared<BackendDeviceType>()) {}
+  : type_(std::make_shared<BackendDeviceType>()), ordinal_(-1) {}
 
 BackendDevice::BackendDevice(std::shared_ptr<BackendDeviceType>&& type, int64_t ordinal)
   : type_(std::move(type)), ordinal_(ordinal) {}
 
-BackendDevice::BackendDevice(const std::string& device_spec)
-  : BackendDevice::BackendDevice() {}
-
 int8_t BackendDevice::type() const {
   TORCH_INTERNAL_ASSERT(type_);
   return type_->type;
 }
 
 std::string BackendDevice::toString() const {
   TORCH_INTERNAL_ASSERT(type_);
-  return c10::str(type_->toString(), ordinal_);
+  std::string str = type_->toString();
+  if (has_index()) {
+    str.append(std::to_string(ordinal_));
+  }
+  return str;
 }
 
 int BackendDevice::compare(const BackendDevice& rhs) const {
@@ -42,10 +42,9 @@ std::ostream& operator<<(std::ostream& os, const BackendDevice& device) {
   return os;
 }
 
-// TODO(whc) refactor this: we need to support non-zero default ordinal for torch/XLA.
 BackendDevice atenDeviceToBackendDevice(const c10::Device& device) {
   TORCH_CHECK(device.type() == at::kLazy, device);
-  int64_t ordinal = device.has_index() ? device.index() : 0;
+  int64_t ordinal = device.has_index() ? device.index() : -1;
   return BackendDevice(getBackend()->GetDefaultDeviceType(), ordinal);
 }
 

torch/csrc/lazy/backend/backend_device.h

@@ -36,16 +36,15 @@ class TORCH_API BackendDevice {
   BackendDevice(std::shared_ptr<BackendDeviceType>&& type, int64_t ordinal);
 
   int8_t type() const;
-  int64_t ordinal() const { return ordinal_;  }
+  int64_t ordinal() const { return ordinal_; }
 
   bool operator==(const BackendDevice& other) const { return compare(other) == 0; }
   bool operator!=(const BackendDevice& other) const { return compare(other) != 0; }
   bool operator<(const BackendDevice& rhs) const { return compare(rhs) < 0; }
 
-  std::string toString() const;
+  bool has_index() const { return ordinal_ >= 0; }
 
-  // The string -> Device conversion should be handled by the backend interface.
-  C10_DEPRECATED explicit BackendDevice(const std::string& device_spec);
+  std::string toString() const;
 
  private:
   int compare(const BackendDevice& rhs) const;