PR #12912: Overload operator<< for XLA_FFI_DataType

Imported from GitHub PR openxla/xla#12912 Copybara import of the project: -- 32908d2d90a89269c655dff6d999bfc974f405a9 by Andrey Portnoy <aportnoy@nvidia.com>: Overload operator<< for XLA_FFI_DataType Merging this change closes #12912 PiperOrigin-RevId: 635960244
tensorflow · May 21, 2024 · 5afa4a8 · 5afa4a8
1 parent e8b4bfb
commit 5afa4a8
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 36 deletions.
diff --git a/third_party/xla/xla/ffi/api/api.h b/third_party/xla/xla/ffi/api/api.h
@@ -1299,6 +1299,46 @@ inline std::ostream& operator<<(std::ostream& os, const XLA_FFI_AttrType type) {
   }
 }
 
+inline std::ostream& operator<<(std::ostream& os,
+                                const XLA_FFI_DataType dtype) {
+  switch (dtype) {
+    case XLA_FFI_DataType_INVALID:
+      return os << "INVALID";
+    case XLA_FFI_DataType_PRED:
+      return os << "PRED";
+    case XLA_FFI_DataType_S8:
+      return os << "S8";
+    case XLA_FFI_DataType_S16:
+      return os << "S16";
+    case XLA_FFI_DataType_S32:
+      return os << "S32";
+    case XLA_FFI_DataType_S64:
+      return os << "S64";
+    case XLA_FFI_DataType_U8:
+      return os << "U8";
+    case XLA_FFI_DataType_U16:
+      return os << "U16";
+    case XLA_FFI_DataType_U32:
+      return os << "U32";
+    case XLA_FFI_DataType_U64:
+      return os << "U64";
+    case XLA_FFI_DataType_F16:
+      return os << "F16";
+    case XLA_FFI_DataType_F32:
+      return os << "F32";
+    case XLA_FFI_DataType_F64:
+      return os << "F64";
+    case XLA_FFI_DataType_BF16:
+      return os << "BF16";
+    case XLA_FFI_DataType_C64:
+      return os << "C64";
+    case XLA_FFI_DataType_C128:
+      return os << "C128";
+    case XLA_FFI_DataType_TOKEN:
+      return os << "TOKEN";
+  }
+}
+
 #define XLA_FFI_REGISTER_SCALAR_ATTR_DECODING(T, TYPE)                \
   template <>                                                         \
   struct AttrDecoding<T> {                                            \

diff --git a/third_party/xla/xla/ffi/api/ffi.h b/third_party/xla/xla/ffi/api/ffi.h
@@ -63,42 +63,7 @@ enum class DataType : uint8_t {
 };
 
 inline std::ostream& operator<<(std::ostream& os, const DataType dtype) {
-  switch (dtype) {
-    case DataType::INVALID:
-      return os << "INVALID";
-    case DataType::PRED:
-      return os << "PRED";
-    case DataType::S8:
-      return os << "S8";
-    case DataType::S16:
-      return os << "S16";
-    case DataType::S32:
-      return os << "S32";
-    case DataType::S64:
-      return os << "S64";
-    case DataType::U8:
-      return os << "U8";
-    case DataType::U16:
-      return os << "U16";
-    case DataType::U32:
-      return os << "U32";
-    case DataType::U64:
-      return os << "U64";
-    case DataType::F16:
-      return os << "F16";
-    case DataType::F32:
-      return os << "F32";
-    case DataType::F64:
-      return os << "F64";
-    case DataType::BF16:
-      return os << "BF16";
-    case DataType::C64:
-      return os << "C64";
-    case DataType::C128:
-      return os << "C128";
-    case DataType::TOKEN:
-      return os << "TOKEN";
-  }
+  return os << static_cast<XLA_FFI_DataType>(dtype);
 }
 
 //===----------------------------------------------------------------------===//