From 5c62322b190065371737105c6c63a565036de8e3 Mon Sep 17 00:00:00 2001
From: Jianhua Zheng <zhengjianhua@oneflow.org>
Date: Tue, 31 Dec 2024 12:08:09 +0800
Subject: [PATCH] add support for xpu/KunLunXin device (#10540)

---
 oneflow/core/common/device_type.proto       | 1 +
 oneflow/ir/lib/OneFlow/UserOpConversion.cpp | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/oneflow/core/common/device_type.proto b/oneflow/core/common/device_type.proto
index 08e702b3232..2d39fdd18ed 100644
--- a/oneflow/core/common/device_type.proto
+++ b/oneflow/core/common/device_type.proto
@@ -9,4 +9,5 @@ enum DeviceType {
   kMeta = 4;
   kMLU = 5;  // Cambricon MLU
   kNPU = 6;  // Ascend NPU
+  kXPU = 7;  // KunLunXin
 }
diff --git a/oneflow/ir/lib/OneFlow/UserOpConversion.cpp b/oneflow/ir/lib/OneFlow/UserOpConversion.cpp
index da289cdc6fa..0b040420268 100644
--- a/oneflow/ir/lib/OneFlow/UserOpConversion.cpp
+++ b/oneflow/ir/lib/OneFlow/UserOpConversion.cpp
@@ -488,6 +488,8 @@ ::oneflow::DeviceType getDeviceTypeFromAttrDictionary(DictionaryAttr attributes)
     return ::oneflow::DeviceType::kMLU;
   } else if (device_tag.str() == "npu") {
     return ::oneflow::DeviceType::kNPU;
+  } else if (device_tag.str() == "xpu") {
+    return ::oneflow::DeviceType::kXPU;
   } else {
     LOG(FATAL) << "unsupported device tag: " << device_tag.str();
     return ::oneflow::DeviceType::kInvalidDevice;