add support for xpu/KunLunXin device

Oneflow-Inc · Jul 3, 2024 · cce5d8c · cce5d8c
1 parent 850b4ad
commit cce5d8c
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git 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
@@ -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;