Fix the ATen mode build of ETDump (pytorch#22135)

### Summary

ETDump maps the device a tensor lives on to the runtime device type so
it can pick the right allocator. In ATen mode it did that by calling
`torch_to_executorch_device` from `extension/aten_util/aten_bridge.h`.

That header cannot be included from a file compiled with `USE_ATEN_LIB`.
It declares functions over `torch::executor::ScalarType` and
`torch::executor::Tensor`, and those two names only exist when
`runtime/core/exec_aten/exec_aten.h` takes its non ATen branch. So the
include fired in exactly the one mode where it does not compile, and the
ATen flavor of the ETDump library failed to build with 7 errors like:

```
no type named 'ScalarType' in namespace 'torch::executor'
```

Nothing in the CMake build compiles these files in ATen mode, so this is
invisible in the CMake based CI jobs and only shows up for the Buck
targets that carry the `_aten` suffix.

### What this changes

* New header `extension/aten_util/aten_device.h` holds the two device
conversions. They are small and self contained, so they are now `inline`
functions that need only `c10/core/Device.h` and the portable `Device`
type, and they compile in both modes.
* `aten_bridge.h` includes the new header and keeps exporting the same
names, so its callers do not change. It also gains a short note saying
it is portable mode only, so the same include does not come back later.
* `etdump_flatcc.cpp` includes the new header instead of the whole
portable mode bridge, and the ATen only Buck dep points at the new
target.
* The portable `Device` type gets its own Buck target, so a caller that
wants only `Device` does not have to depend on every portable type. This
mirrors the existing `:scalar_type` split in the same file.
* The new header is wrapped in
`C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wswitch-enum")`.
`c10::DeviceType` has over twenty values, and `-Wswitch-enum` asks for
every one to be listed even when the switch already has a default.
`extension/tensor/tensor_ptr.h` already does the same thing for its own
switch over a c10 enum. Plain `-Wswitch` is left alone, so an unhandled
new ExecuTorch device type is still a compile error.

No behavior change. The mapping itself is byte for byte the same code,
just moved.

### Test plan

Compiled with `-Wall -Wextra -Wswitch-enum -Werror` against PyTorch
headers on Linux x86_64:

| what | ATen mode | portable mode |
| --- | --- | --- |
| `aten_device.h` | clean | clean |
| the ETDump device helper | clean | clean |
| `aten_bridge.cpp` | not applicable | clean |
| the call shapes the existing `aten_bridge` test uses | not applicable
| clean |

Each check was confirmed to fail when its fix is removed, so none of
them is a test that cannot fail:

* Without the move, the ATen mode compile of the ETDump helper fails
with the 7 errors above.
* Without the `-Wswitch-enum` guard, `aten_device.h` fails with 20
errors.
* Adding a third value to the portable `DeviceType` enum still fails the
build through plain `-Wswitch`, so the guard does not hide a missing
case.

What is not covered: no job in this repository builds the `_aten` Buck
targets, so CI on this pull request exercises the portable build only.
The ATen mode results above come from compiling those translation units
by hand.

Author: shoumikhin

devtools/etdump/etdump_flatcc.cpp

@@ -22,7 +22,7 @@
 #include <executorch/runtime/platform/assert.h>
 
 #ifdef USE_ATEN_LIB
-#include <executorch/extension/aten_util/aten_bridge.h>
+#include <executorch/extension/aten_util/aten_device.h>
 #endif
 
 #include <flatcc/flatcc_types.h>

devtools/etdump/targets.bzl

@@ -135,7 +135,7 @@ def define_common_targets():
                 "//executorch/runtime/core:device_allocator",
                 "//executorch/runtime/platform:platform",
             ] + ([
-                "//executorch/extension/aten_util:aten_bridge",
+                "//executorch/extension/aten_util:aten_device",
             ] if aten_mode else []),
             exported_deps = [
                 ":etdump_schema_flatcc",

extension/aten_util/aten_bridge.cpp

@@ -134,36 +134,6 @@ c10::ScalarType executorch_to_torch_scalar_type(
   return static_cast<c10::ScalarType>(intermediate);
 }
 
-c10::Device executorch_to_torch_device(
-    executorch::runtime::etensor::Device device) {
-  switch (device.type()) {
-    case executorch::runtime::etensor::DeviceType::CPU:
-      return c10::Device(c10::DeviceType::CPU);
-    case executorch::runtime::etensor::DeviceType::CUDA:
-      return c10::Device(c10::DeviceType::CUDA, device.index());
-  }
-  // Aborting on an unknown device rather than falling back to CPU: a wrong
-  // label made the host read accelerator memory and segfault.
-  ET_CHECK_MSG(
-      false,
-      "Tensor reports device type %d, which this build cannot map to a PyTorch device",
-      static_cast<int>(device.type()));
-}
-
-std::optional<executorch::runtime::etensor::Device> torch_to_executorch_device(
-    c10::Device device) {
-  switch (device.type()) {
-    case c10::DeviceType::CPU:
-      return executorch::runtime::etensor::Device(
-          executorch::runtime::etensor::DeviceType::CPU);
-    case c10::DeviceType::CUDA:
-      return executorch::runtime::etensor::Device(
-          executorch::runtime::etensor::DeviceType::CUDA, device.index());
-    default:
-      return std::nullopt;
-  }
-}
-
 /*
  * Following makes two assumptions:
  * 1. aten_tensor's lifetime is longer than the liftime within which mutable_et

extension/aten_util/aten_bridge.h

@@ -8,17 +8,20 @@
 
 #pragma once
 
+// Portable mode only: the declarations below name torch::executor::ScalarType
+// and torch::executor::Tensor, which exec_aten.h defines only when it is not
+// building against ATen. For device conversions either way, use aten_device.h.
+
+#include <executorch/extension/aten_util/aten_device.h>
 #include <executorch/extension/tensor/tensor.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
-#include <executorch/runtime/core/portable_type/device.h>
 
 #include <ATen/Functions.h> // @manual=//caffe2/aten:ATen-cpu
 #include <ATen/Tensor.h> // @manual=//caffe2/aten:ATen-core
 #include <ATen/core/functional.h> // @manual=//caffe2/aten:ATen-core
 #include <c10/core/ScalarTypeToTypeMeta.h> // @manual=//caffe2/c10:c10
 
 #include <memory>
-#include <optional>
 #include <vector>
 
 namespace executorch {
@@ -30,19 +33,6 @@ torch::executor::ScalarType torch_to_executorch_scalar_type(
 c10::ScalarType executorch_to_torch_scalar_type(
     torch::executor::ScalarType type);
 
-c10::Device executorch_to_torch_device(
-    executorch::runtime::etensor::Device device);
-
-/**
- * Maps a PyTorch device onto the ExecuTorch device naming the same location.
- *
- * Returns nothing for a device this runtime has no type for. That is a valid
- * thing for a caller to ask, so it is reported rather than fatal, and the
- * caller adds the context it has before failing.
- */
-std::optional<executorch::runtime::etensor::Device> torch_to_executorch_device(
-    c10::Device device);
-
 /*
  * @param[in] aten_tensor Input at::Tensor
  * @param[in,out] mutable_et ETensor whose underlying memory now will alias to

extension/aten_util/aten_device.h

@@ -0,0 +1,72 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <executorch/runtime/core/portable_type/device.h>
+#include <executorch/runtime/platform/assert.h>
+
+#include <c10/core/Device.h> // @manual=//caffe2/c10:c10
+#include <c10/macros/Macros.h> // @manual=//caffe2/c10:c10
+
+#include <optional>
+
+// Kept out of aten_bridge.h, which compiles in portable mode only, so that a
+// translation unit built with USE_ATEN_LIB can still convert devices.
+
+// c10::DeviceType has over twenty values and -Wswitch-enum wants every one
+// listed even with a default. -Wswitch still covers the switch below.
+C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wswitch-enum")
+
+namespace executorch {
+namespace extension {
+
+/**
+ * Maps an ExecuTorch device onto the PyTorch device naming the same location.
+ */
+inline c10::Device executorch_to_torch_device(
+    executorch::runtime::etensor::Device device) {
+  switch (device.type()) {
+    case executorch::runtime::etensor::DeviceType::CPU:
+      return c10::Device(c10::DeviceType::CPU);
+    case executorch::runtime::etensor::DeviceType::CUDA:
+      return c10::Device(c10::DeviceType::CUDA, device.index());
+  }
+  // Aborting on an unknown device rather than falling back to CPU: a wrong
+  // label made the host read accelerator memory and segfault.
+  ET_CHECK_MSG(
+      false,
+      "Tensor reports device type %d, which this build cannot map to a PyTorch device",
+      static_cast<int>(device.type()));
+}
+
+/**
+ * Maps a PyTorch device onto the ExecuTorch device naming the same location.
+ *
+ * Returns nothing for a device this runtime has no type for. That is a valid
+ * thing for a caller to ask, so it is reported rather than fatal, and the
+ * caller adds the context it has before failing.
+ */
+inline std::optional<executorch::runtime::etensor::Device>
+torch_to_executorch_device(c10::Device device) {
+  switch (device.type()) {
+    case c10::DeviceType::CPU:
+      return executorch::runtime::etensor::Device(
+          executorch::runtime::etensor::DeviceType::CPU);
+    case c10::DeviceType::CUDA:
+      return executorch::runtime::etensor::Device(
+          executorch::runtime::etensor::DeviceType::CUDA, device.index());
+    default:
+      return std::nullopt;
+  }
+}
+
+} // namespace extension
+} // namespace executorch
+
+C10_DIAGNOSTIC_POP()

extension/aten_util/targets.bzl

@@ -7,6 +7,21 @@ def define_common_targets():
     TARGETS and BUCK files that call this function.
     """
 
+    # Split out of :aten_bridge, which compiles in portable mode only, so that
+    # ATen mode code can still convert devices.
+    runtime.cxx_library(
+        name = "aten_device",
+        exported_headers = ["aten_device.h"],
+        visibility = ["PUBLIC"],
+        exported_deps = [
+            "//executorch/runtime/core/portable_type:device",
+            "//executorch/runtime/platform:platform",
+        ],
+        exported_external_deps = [
+            "libtorch",
+        ],
+    )
+
     runtime.cxx_library(
         name = "aten_bridge",
         srcs = ["aten_bridge.cpp"],
@@ -23,6 +38,7 @@ def define_common_targets():
         ],
         visibility = ["PUBLIC"],
         exported_deps = [
+            ":aten_device",
             "//executorch/extension/kernel_util:kernel_util",
             "//executorch/extension/tensor:tensor",
             "//executorch/runtime/core:core",

runtime/core/portable_type/targets.bzl

@@ -19,7 +19,6 @@ def define_common_targets():
             "tensor.h",
             "tensor_impl.h",
             "string_view.h",
-            "device.h",
         ],
         # Only should be depended on by kernel_types:kernel_types, but various suffixes like Android and Static
         # mean I cant just expose visibility to a single rule.
@@ -36,6 +35,7 @@ def define_common_targets():
             "//executorch/runtime/core/portable_type/c10/c10:c10",
         ],
         exported_deps = [
+            ":device",
             ":scalar_type",
             "//executorch/runtime/core:core",
             "//executorch/runtime/core:tensor_shape_dynamism",
@@ -46,6 +46,20 @@ def define_common_targets():
         ],
     )
 
+    # device.h is standalone and is the one portable type header that ATen mode
+    # code also needs, so it is split out like :scalar_type below rather than
+    # making a caller that wants only Device depend on every portable type.
+    runtime.cxx_library(
+        name = "device",
+        exported_headers = [
+            "device.h",
+        ],
+        visibility = [
+            "//executorch/extension/aten_util/...",
+            "//executorch/runtime/core/...",
+        ],
+    )
+
     # Set up a specific exported library for scalar_type to avoid circle dependency in ScalarTypeUtil.h
     runtime.cxx_library(
         name = "scalar_type",