issue/988 - adapt to ali ppu

Author: wooway777

include/infinicore.h

@@ -47,6 +47,7 @@ typedef enum {
     INFINI_DEVICE_KUNLUN = 7,
     INFINI_DEVICE_HYGON = 8,
     INFINI_DEVICE_QY = 9,
+    INFINI_DEVICE_ALI = 10,
     INFINI_DEVICE_TYPE_COUNT
 } infiniDevice_t;
 

include/infinicore/device.hpp

@@ -22,6 +22,7 @@ class Device {
         KUNLUN = INFINI_DEVICE_KUNLUN,
         HYGON = INFINI_DEVICE_HYGON,
         QY = INFINI_DEVICE_QY,
+        ALI = INFINI_DEVICE_ALI,
         COUNT = INFINI_DEVICE_TYPE_COUNT,
     };
 

python/infinicore/device.py

@@ -82,6 +82,7 @@ def _from_infinicore_device(infinicore_device: _infinicore.Device):
     _infinicore.Device.Type.KUNLUN: "cuda",
     _infinicore.Device.Type.HYGON: "cuda",
     _infinicore.Device.Type.QY: "cuda",
+    _infinicore.Device.Type.ALI: "cuda",
 }
 
 

src/infiniccl-test/main.cpp

@@ -12,7 +12,7 @@ void printUsage() {
     std::cout << "infiniccl-test --<device>" << std::endl
               << std::endl;
     std::cout << "  --<device>" << std::endl;
-    std::cout << "    Specify the device type --(nvidia|cambricon|ascend|metax|moore|iluvatar|qy|kunlun|hygon)." << std::endl
+    std::cout << "    Specify the device type --(nvidia|cambricon|ascend|metax|moore|iluvatar|qy|kunlun|hygon|ali)." << std::endl
               << std::endl;
     std::cout << "The program will run tests on all visible devices of the specified device type."
               << " Use Environmental Variables such as CUDA_VSIBLE_DEVICES to limit visible device IDs.";
@@ -46,6 +46,7 @@ ParsedArgs parseArgs(int argc, char *argv[]) {
         else PARSE_DEVICE("--qy", INFINI_DEVICE_QY)
         else PARSE_DEVICE("--kunlun", INFINI_DEVICE_KUNLUN)
         else PARSE_DEVICE("--hygon", INFINI_DEVICE_HYGON)
+        else PARSE_DEVICE("--ali", INFINI_DEVICE_ALI)
         else {
             printUsage();
         }

src/infiniccl/cuda/infiniccl_cuda.h

@@ -4,7 +4,7 @@
 #include "../infiniccl_impl.h"
 
 // Windows does not support CUDA
-#if (defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API)) && defined(ENABLE_CCL) && !defined(_WIN32)
+#if (defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API)) || defined(ENABLE_ALI_API) && defined(ENABLE_CCL) && !defined(_WIN32)
 INFINICCL_DEVICE_API_IMPL(cuda)
 #else
 INFINICCL_DEVICE_API_NOOP(cuda)

src/infiniccl/infiniccl.cc

@@ -27,6 +27,7 @@ __C infiniStatus_t infinicclCommInitAll(
         COMM_INIT_ALL(INFINI_DEVICE_METAX, metax);
         COMM_INIT_ALL(INFINI_DEVICE_MOORE, moore);
         COMM_INIT_ALL(INFINI_DEVICE_KUNLUN, kunlun);
+        COMM_INIT_ALL(INFINI_DEVICE_ALI, cuda);
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
     }
@@ -53,6 +54,7 @@ __C infiniStatus_t infinicclCommDestroy(infinicclComm_t comm) {
         COMM_DESTROY(INFINI_DEVICE_METAX, metax);
         COMM_DESTROY(INFINI_DEVICE_MOORE, moore);
         COMM_DESTROY(INFINI_DEVICE_KUNLUN, kunlun);
+        COMM_DESTROY(INFINI_DEVICE_ALI, cuda);
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
     }
@@ -86,6 +88,7 @@ __C infiniStatus_t infinicclAllReduce(
         ALL_REDUCE(INFINI_DEVICE_METAX, metax);
         ALL_REDUCE(INFINI_DEVICE_MOORE, moore);
         ALL_REDUCE(INFINI_DEVICE_KUNLUN, kunlun);
+        ALL_REDUCE(INFINI_DEVICE_ALI, cuda);
 
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;

src/infinicore-test/README.md

@@ -66,6 +66,7 @@ xmake build infinicore-test
 ./infinicore-test --qy
 ./infinicore-test --kunlun
 ./infinicore-test --hygon
+./infinicore-test --ali
 ```
 
 ### Customize Test Parameters

src/infinicore-test/main.cc

@@ -42,6 +42,7 @@ void printUsage() {
               << "  qy" << std::endl
               << "  kunlun" << std::endl
               << "  hygon" << std::endl
+              << "  ali" << std::endl
               << std::endl
               << "Available tests:" << std::endl
               << "  basic       - Basic memory allocation and deallocation tests" << std::endl
@@ -84,6 +85,8 @@ ParsedArgs parseArgs(int argc, char *argv[]) {
             args.device_type = INFINI_DEVICE_KUNLUN;
         } else if (arg == "--hygon") {
             args.device_type = INFINI_DEVICE_HYGON;
+        } else if (arg == "--ali") {
+            args.device_type = INFINI_DEVICE_ALI;
         } else if (arg == "--test") {
             if (i + 1 >= argc) {
                 std::cerr << "Error: --test requires a test name" << std::endl;

src/infinicore/device.cc

@@ -41,6 +41,8 @@ std::string Device::toString(const Type &type) {
         return "KUNLUN";
     case Type::HYGON:
         return "HYGON";
+    case Type::ALI:
+        return "ALI";
     case Type::COUNT:
         return "COUNT";
     default:

src/infinicore/nn/embedding.cc

@@ -45,7 +45,7 @@ Embedding::Embedding(size_t num_embeddings,
 Tensor Embedding::forward(const Tensor &indices) const {
     // TODO: Implement on-device embedding for all devices, then remove the condition and the classic approach
     auto device_type = device_.getType();
-    if (device_type == Device::Type::NVIDIA || device_type == Device::Type::ILUVATAR || device_type == Device::Type::METAX || device_type == Device::Type::MOORE) {
+    if (device_type == Device::Type::NVIDIA || device_type == Device::Type::ILUVATAR || device_type == Device::Type::METAX || device_type == Device::Type::MOORE || device_type == Device::Type::KUNLUN || device_type == Device::Type::HYGON || device_type == Device::Type::QY || device_type == Device::Type::ALI) {
         // Use op::embedding which supports device-side input and batch dimension
         return op::embedding(indices->contiguous()->to(device_), weight_);
     }