Address review.

Author: ysiraichi

test/cpp/test_device.cpp

@@ -1,10 +1,12 @@
 #include <gtest/gtest.h>
-#include <torch/csrc/lazy/backend/backend_device.h>
 
 #include <cstdint>
 #include <string>
 #include <string_view>
 
+#include <torch/csrc/lazy/backend/backend_device.h>
+
+#include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 
 #include "torch_xla/csrc/device.h"
@@ -47,11 +49,11 @@ static void CheckDeviceTypeConstructionWithString(
 }
 
 TEST(DeviceTest, ConstructNativeDeviceTypeWithString) {
-#define XLA_NATIVE_DEVICE_TYPE_CHECK_CONSTRUCTION_WITH_STRING(type, _) \
-  CheckDeviceTypeConstructionWithString(XlaDeviceType::type, #type);
-  XLA_FOR_ALL_NATIVE_DEVICE_TYPES_(
-      XLA_NATIVE_DEVICE_TYPE_CHECK_CONSTRUCTION_WITH_STRING)
-#undef XLA_DEVICE_TYPE_CHECK_CONSTRUCTION_WITH_STRING
+  CheckDeviceTypeConstructionWithString(XlaDeviceType::CPU, "CPU");
+  CheckDeviceTypeConstructionWithString(XlaDeviceType::CUDA, "CUDA");
+  CheckDeviceTypeConstructionWithString(XlaDeviceType::TPU, "TPU");
+  CheckDeviceTypeConstructionWithString(XlaDeviceType::NEURON, "NEURON");
+  CheckDeviceTypeConstructionWithString(XlaDeviceType::SPMD, "SPMD");
 }
 
 TEST(DeviceTest, ConstructPluginDeviceTypeWithString) {
@@ -68,11 +70,11 @@ static void CheckDeviceTypeConstructionWithEnum(
 }
 
 TEST(DeviceTest, ConstructNativeDeviceTypeWithEnum) {
-#define XLA_NATIVE_DEVICE_TYPE_CHECK_CONSTRUCTION_WITH_ENUM(type, _) \
-  CheckDeviceTypeConstructionWithEnum(XlaDeviceType::type, #type);
-  XLA_FOR_ALL_NATIVE_DEVICE_TYPES_(
-      XLA_NATIVE_DEVICE_TYPE_CHECK_CONSTRUCTION_WITH_ENUM)
-#undef XLA_DEVICE_TYPE_CHECK_CONSTRUCTION_WITH_STRING
+  CheckDeviceTypeConstructionWithEnum(XlaDeviceType::CPU, "CPU");
+  CheckDeviceTypeConstructionWithEnum(XlaDeviceType::CUDA, "CUDA");
+  CheckDeviceTypeConstructionWithEnum(XlaDeviceType::TPU, "TPU");
+  CheckDeviceTypeConstructionWithEnum(XlaDeviceType::NEURON, "NEURON");
+  CheckDeviceTypeConstructionWithEnum(XlaDeviceType::SPMD, "SPMD");
 }
 
 TEST(DeviceTest, ConstructPluginDeviceTypeWithEnumError) {

torch_xla/csrc/device.cpp

@@ -5,15 +5,14 @@
 #include <cstddef>
 #include <cstdint>
 #include <exception>
+#include <iterator>
 #include <memory>
 #include <string>
 #include <string_view>
-#include <utility>
 #include <vector>
 
 #include <torch/csrc/lazy/backend/backend_device.h>
 
-#include "absl/log/absl_check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
@@ -30,14 +29,13 @@ namespace {
 static bool spmd_config_is_locked = false;
 static bool use_virtual_device = false;
 
-constexpr std::size_t kNativeXlaDeviceTypeNumber = 5;
-constexpr std::array<std::pair<XlaDeviceType, std::string_view>,
-                     kNativeXlaDeviceTypeNumber>
-    kNativeXlaDeviceTypeWithName = {{
-#define XLA_DEVICE_NAME_PAIR(name, _) {XlaDeviceType::name, #name},
-        XLA_FOR_ALL_NATIVE_DEVICE_TYPES_(XLA_DEVICE_NAME_PAIR)
-#undef XLA_DEVICE_NAME_PAIR
-    }};
+constexpr int8_t kNativeXlaDeviceTypeNumber =
+    static_cast<int8_t>(XlaDeviceType::PLUGIN);
+
+// The elements in this array should match the order in the XlaDeviceType enum
+// declaration. So, if you modify one of them, make sure to keep them in sync.
+constexpr std::array<std::string_view, kNativeXlaDeviceTypeNumber>
+    kNativeXlaDeviceTypeNames = {"CPU", "CUDA", "TPU", "NEURON", "SPMD"};
 
 absl::Status CheckIsNativeXlaDeviceType(int8_t value) {
   if (value < 0 || value >= kNativeXlaDeviceTypeNumber) {
@@ -62,21 +60,20 @@ std::string_view NativeXlaDeviceTypeToString(XlaDeviceType type) {
   //   2. The XlaDeviceType::PLUGIN enum, since it's not considered a "native"
   //      device type
   XLA_CHECK_OK(CheckIsNativeXlaDeviceType(value));
-  return kNativeXlaDeviceTypeWithName[value].second;
+  return kNativeXlaDeviceTypeNames[value];
 }
 
 XlaDeviceType StringToXlaDeviceType(std::string_view type_name) {
-  std::array<std::pair<XlaDeviceType, std::string_view>,
-             kNativeXlaDeviceTypeNumber>::const_iterator it =
-      std::find_if(kNativeXlaDeviceTypeWithName.begin(),
-                   kNativeXlaDeviceTypeWithName.end(),
-                   [=](const std::pair<XlaDeviceType, std::string_view>& pair) {
-                     return pair.second == type_name;
-                   });
-  if (it != kNativeXlaDeviceTypeWithName.end()) {
-    return it->first;
+  std::array<std::string_view, kNativeXlaDeviceTypeNumber>::const_iterator it =
+      std::find(kNativeXlaDeviceTypeNames.begin(),
+                kNativeXlaDeviceTypeNames.end(), type_name);
+
+  if (it == kNativeXlaDeviceTypeNames.end()) {
+    return XlaDeviceType::PLUGIN;
   }
-  return XlaDeviceType::PLUGIN;
+
+  std::size_t index = std::distance(kNativeXlaDeviceTypeNames.begin(), it);
+  return static_cast<XlaDeviceType>(index);
 }
 
 }  // namespace

torch_xla/csrc/device.h

@@ -1,48 +1,20 @@
 #ifndef XLA_TORCH_XLA_CSRC_DEVICE_H_
 #define XLA_TORCH_XLA_CSRC_DEVICE_H_
 
+#include <cstdint>
 #include <string>
 #include <string_view>
 
 #include <torch/csrc/lazy/backend/backend_device.h>
-#include <torch/csrc/lazy/core/hash.h>
-#include <torch/csrc/lazy/core/util.h>
 
 #include "absl/status/statusor.h"
 
 namespace torch_xla {
 
-// Convenient macro for applying another macro to all native device types.
-//
-// Add new device type
-// ===================
-//
-// Add a new line to the macro below:
-//
-//     _(<DEVICE>, <INDEX>)
-//
-// Where <DEVICE> is the enum of the given device, and <INDEX> is the
-// previous number plus 1.
-//
-#define XLA_FOR_ALL_NATIVE_DEVICE_TYPES_(_) \
-  _(CPU, 0)                                 \
-  _(CUDA, 1)                                \
-  _(TPU, 2)                                 \
-  _(NEURON, 3)                              \
-  _(SPMD, 4)
-
 // TODO(yeounoh) `SPMD` is a virtual device that defers data `TransferToDevice`
 // until after the paritioning pass. This avoids transfering  the full input
 // tensor to the device.
-enum class XlaDeviceType : int8_t {
-#define XLA_DECLARE_ENUM(name, value) name = value,
-  XLA_FOR_ALL_NATIVE_DEVICE_TYPES_(XLA_DECLARE_ENUM)
-#undef XLA_DECLARE_ENUM
-
-  // Plugin is not considered a native device type.
-  // It has a special treatment for some functions.
-  PLUGIN,
-};
+enum class XlaDeviceType : int8_t { CPU = 0, CUDA, TPU, NEURON, SPMD, PLUGIN };
 
 struct DeviceType : public torch::lazy::BackendDeviceType {
   DeviceType(XlaDeviceType xla_device_type);