[et device 1/n] introduce cuda device type

runtime/core/portable_type/device.h

@@ -8,28 +8,33 @@
 
 #pragma once
 
-#include <executorch/runtime/platform/assert.h>
+#include <cstddef>
+#include <cstdint>
 
 namespace executorch {
 namespace runtime {
 namespace etensor {
 
-/// Denotes the specific genre of compute device.
-/// Subset of https://github.com/pytorch/pytorch/blob/main/c10/core/Device.h
+/// Represents the type of compute device.
+/// Note: ExecuTorch Device is distinct from PyTorch Device.
 enum class DeviceType : int8_t {
   CPU = 0,
+  CUDA = 1,
 };
 
-/// An index representing a specific device; For cpu it should always be -1 or 0
+/// Total number of device types, used for fixed-size registry arrays.
+constexpr size_t kNumDeviceTypes = 2;
+
+/// An index representing a specific device; e.g. GPU 0 vs GPU 1.
+/// -1 means the default/unspecified device for that type.
 using DeviceIndex = int8_t;
 
 /**
  * An abstraction for the compute device on which a tensor is located.
- * ExecuTorch doesn't allow dynamic dispatching based on device, so this type is
- * just a skeleton to allow certain kernels that expect device as an
- * argument to still be run.
  *
- * In ExecuTorch this is always expected to be CPU.
+ * Tensors carry a Device to express where their underlying data resides
+ * (e.g. CPU host memory vs CUDA device memory). The runtime uses this to
+ * dispatch memory allocation to the appropriate device allocator.
  */
 struct Device final {
   using Type = DeviceType;
@@ -39,7 +44,7 @@ struct Device final {
   /* implicit */ Device(DeviceType type, DeviceIndex index = -1)
       : type_(type), index_(index) {}
 
-  /// Returns the type of device this is. Only CPU is supported.
+  /// Returns the type of device the tensor data resides on.
   DeviceType type() const noexcept {
     return type_;
   }
@@ -49,12 +54,19 @@ struct Device final {
     return type_ == DeviceType::CPU;
   }
 
-  /// Returns the device index. Always 0 if specified or -1 if not provided.
+  /// Returns the device index, or -1 if default/unspecified.
   DeviceIndex index() const noexcept {
-    ET_CHECK(index_ == 0 || index_ == -1);
     return index_;
   }
 
+  bool operator==(const Device& other) const noexcept {
+    return type_ == other.type_ && index_ == other.index_;
+  }
+
+  bool operator!=(const Device& other) const noexcept {
+    return !(*this == other);
+  }
+
  private:
   DeviceType type_;
   DeviceIndex index_ = -1;

runtime/core/portable_type/test/device_test.cpp

@@ -0,0 +1,95 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/runtime/core/portable_type/device.h>
+
+#include <gtest/gtest.h>
+
+using executorch::runtime::etensor::Device;
+using executorch::runtime::etensor::DeviceIndex;
+using executorch::runtime::etensor::DeviceType;
+using executorch::runtime::etensor::kNumDeviceTypes;
+
+// --- DeviceType enum ---
+
+TEST(DeviceTypeTest, EnumValues) {
+  EXPECT_EQ(static_cast<int8_t>(DeviceType::CPU), 0);
+  EXPECT_EQ(static_cast<int8_t>(DeviceType::CUDA), 1);
+}
+
+TEST(DeviceTypeTest, NumDeviceTypesCoversAllEnums) {
+  // kNumDeviceTypes must be large enough to index all defined device types.
+  EXPECT_GT(kNumDeviceTypes, static_cast<size_t>(DeviceType::CPU));
+  EXPECT_GT(kNumDeviceTypes, static_cast<size_t>(DeviceType::CUDA));
+}
+
+// --- Device: CPU ---
+
+TEST(DeviceTest, CpuDefaultIndex) {
+  Device d(DeviceType::CPU);
+  EXPECT_TRUE(d.is_cpu());
+  EXPECT_EQ(d.type(), DeviceType::CPU);
+  EXPECT_EQ(d.index(), -1);
+}
+
+TEST(DeviceTest, CpuExplicitIndex) {
+  Device d(DeviceType::CPU, 0);
+  EXPECT_TRUE(d.is_cpu());
+  EXPECT_EQ(d.index(), 0);
+}
+
+// --- Device: CUDA ---
+
+TEST(DeviceTest, CudaDefaultIndex) {
+  Device d(DeviceType::CUDA);
+  EXPECT_FALSE(d.is_cpu());
+  EXPECT_EQ(d.type(), DeviceType::CUDA);
+  EXPECT_EQ(d.index(), -1);
+}
+
+TEST(DeviceTest, CudaExplicitIndex) {
+  Device d(DeviceType::CUDA, 0);
+  EXPECT_EQ(d.index(), 0);
+}
+
+// --- Device: equality ---
+
+TEST(DeviceTest, EqualitySameTypeAndIndex) {
+  EXPECT_EQ(Device(DeviceType::CPU, 0), Device(DeviceType::CPU, 0));
+  EXPECT_EQ(Device(DeviceType::CUDA, 1), Device(DeviceType::CUDA, 1));
+}
+
+TEST(DeviceTest, InequalityDifferentType) {
+  EXPECT_NE(Device(DeviceType::CPU, 0), Device(DeviceType::CUDA, 0));
+}
+
+TEST(DeviceTest, InequalityDifferentIndex) {
+  EXPECT_NE(Device(DeviceType::CUDA, 0), Device(DeviceType::CUDA, 1));
+}
+
+TEST(DeviceTest, EqualityDefaultIndices) {
+  EXPECT_EQ(Device(DeviceType::CPU), Device(DeviceType::CPU));
+  EXPECT_EQ(Device(DeviceType::CUDA), Device(DeviceType::CUDA));
+  EXPECT_NE(Device(DeviceType::CPU), Device(DeviceType::CUDA));
+}
+
+// --- Device: implicit construction ---
+
+TEST(DeviceTest, ImplicitConstructionFromDeviceType) {
+  // Device constructor is implicit, allowing DeviceType → Device conversion.
+  Device d = DeviceType::CUDA;
+  EXPECT_EQ(d.index(), -1);
+}
+
+// --- Deprecated namespace aliases ---
+
+TEST(DeviceTest, DeprecatedNamespaceAliases) {
+  // Verify the torch::executor aliases still work.
+  torch::executor::Device d(torch::executor::DeviceType::CUDA, 0);
+  EXPECT_EQ(d.index(), 0);
+}

runtime/core/portable_type/test/targets.bzl

@@ -47,6 +47,14 @@ def define_common_targets():
         ],
     )
 
+    runtime.cxx_test(
+        name = "device_test",
+        srcs = ["device_test.cpp"],
+        deps = [
+            "//executorch/runtime/core/portable_type:portable_type",
+        ],
+    )
+
     runtime.cxx_test(
         name = "tensor_impl_test",
         srcs = ["tensor_impl_test.cpp"],