feat: support MemorySpace on Iluvatar, Moore Threads, and Cambricon

src/cambricon/checks.h

@@ -0,0 +1,32 @@
+#ifndef INFINI_CCL_CAMBRICON_CHECKS_H_
+#define INFINI_CCL_CAMBRICON_CHECKS_H_
+
+#include <cnrt.h>
+
+#include <iostream>
+
+#include "return_status_impl.h"
+
+#define INFINI_CHECK_CNRT(result) \
+  ::infini::ccl::detail::CheckCnrtImpl((result), __FILE__, __LINE__)
+
+namespace infini::ccl {
+
+namespace detail {
+
+inline ReturnStatus CheckCnrtImpl(cnrtRet_t cnrt_result, const char *file,
+                                  int line) {
+  if (cnrt_result != cnrtSuccess) {
+    cnrtGetLastError();
+    std::cerr << "CNRT error code: " << cnrt_result << " at line " << line
+              << " in " << file << std::endl;
+    std::abort();
+  }
+  return ReturnStatus::kSuccess;
+}
+
+}  // namespace detail
+
+}  // namespace infini::ccl
+
+#endif  // INFINI_CCL_CAMBRICON_CHECKS_H_

src/cambricon/device_.h

@@ -1,13 +1,32 @@
 #ifndef INFINI_CCL_CAMBRICON_DEVICE__H_
 #define INFINI_CCL_CAMBRICON_DEVICE__H_
 
+#include <cnrt.h>
+
+#include "checks.h"
 #include "device.h"
 
 namespace infini::ccl {
 
 template <>
 struct DeviceEnabled<Device::Type::kCambricon> : std::true_type {};
 
+template <>
+MemorySpace GetMemorySpace<Device::Type::kCambricon>(const void* ptr) {
+  if (!ptr) {
+    return MemorySpace::kHost;
+  }
+
+  cnrtPointerAttributes_t attr;
+  INFINI_CHECK_CNRT(cnrtPointerGetAttributes(&attr, ptr));
+
+  if (attr.type == cnrtMemTypeDevice) {
+    return MemorySpace::kDevice;
+  }
+
+  return MemorySpace::kHost;
+}
+
 }  // namespace infini::ccl
 
 #endif  // INFINI_CCL_CAMBRICON_DEVICE__H_

src/nvidia/checks.h → src/cuda/checks.h

@@ -1,5 +1,5 @@
-#ifndef INFINI_CCL_NVIDIA_CHECKS_H_
-#define INFINI_CCL_NVIDIA_CHECKS_H_
+#ifndef INFINI_CCL_CUDA_CHECKS_H_
+#define INFINI_CCL_CUDA_CHECKS_H_
 
 #include <cuda_runtime.h>
 
@@ -29,4 +29,4 @@ inline ReturnStatus CheckCudaImpl(cudaError_t cuda_result, const char *file,
 
 }  // namespace infini::ccl
 
-#endif  // INFINI_CCL_NVIDIA_CHECKS_H_
+#endif  // INFINI_CCL_CUDA_CHECKS_H_

src/iluvatar/device_.h

@@ -1,13 +1,32 @@
 #ifndef INFINI_CCL_ILUVATAR_DEVICE__H_
 #define INFINI_CCL_ILUVATAR_DEVICE__H_
 
+#include <cuda_runtime.h>
+
+#include "cuda/checks.h"
 #include "device.h"
 
 namespace infini::ccl {
 
 template <>
 struct DeviceEnabled<Device::Type::kIluvatar> : std::true_type {};
 
+template <>
+MemorySpace GetMemorySpace<Device::Type::kIluvatar>(const void* ptr) {
+  if (!ptr) {
+    return MemorySpace::kHost;
+  }
+
+  cudaPointerAttributes attr;
+  INFINI_CHECK_CUDA(cudaPointerGetAttributes(&attr, ptr));
+
+  if (attr.type == cudaMemoryTypeDevice || attr.type == cudaMemoryTypeManaged) {
+    return MemorySpace::kDevice;
+  }
+
+  return MemorySpace::kHost;
+}
+
 }  // namespace infini::ccl
 
 #endif  // INFINI_CCL_ILUVATAR_DEVICE__H_

src/moore/checks.h

@@ -0,0 +1,32 @@
+#ifndef INFINI_CCL_MOORE_CHECKS_H_
+#define INFINI_CCL_MOORE_CHECKS_H_
+
+#include <musa_runtime.h>
+
+#include <iostream>
+
+#include "return_status_impl.h"
+
+#define INFINI_CHECK_MUSA(result) \
+  ::infini::ccl::detail::CheckMusaImpl((result), __FILE__, __LINE__)
+
+namespace infini::ccl {
+
+namespace detail {
+
+inline ReturnStatus CheckMusaImpl(musaError_t musa_result, const char *file,
+                                  int line) {
+  if (musa_result != musaSuccess) {
+    musaGetLastError();
+    std::cerr << "MUSA error code: " << musa_result << " at line " << line
+              << " in " << file << std::endl;
+    std::abort();
+  }
+  return ReturnStatus::kSuccess;
+}
+
+}  // namespace detail
+
+}  // namespace infini::ccl
+
+#endif  // INFINI_CCL_MOORE_CHECKS_H_

src/moore/device_.h

@@ -1,13 +1,32 @@
 #ifndef INFINI_CCL_MOORE_DEVICE__H_
 #define INFINI_CCL_MOORE_DEVICE__H_
 
+#include <musa_runtime.h>
+
+#include "checks.h"
 #include "device.h"
 
 namespace infini::ccl {
 
 template <>
 struct DeviceEnabled<Device::Type::kMoore> : std::true_type {};
 
+template <>
+MemorySpace GetMemorySpace<Device::Type::kMoore>(const void* ptr) {
+  if (!ptr) {
+    return MemorySpace::kHost;
+  }
+
+  musaPointerAttributes attr;
+  INFINI_CHECK_MUSA(musaPointerGetAttributes(&attr, ptr));
+
+  if (attr.type == musaMemoryTypeDevice || attr.type == musaMemoryTypeManaged) {
+    return MemorySpace::kDevice;
+  }
+
+  return MemorySpace::kHost;
+}
+
 }  // namespace infini::ccl
 
 #endif  // INFINI_CCL_MOORE_DEVICE__H_

src/nvidia/device_.h

@@ -3,7 +3,7 @@
 
 #include <cuda_runtime.h>
 
-#include "checks.h"
+#include "cuda/checks.h"
 #include "device.h"
 
 namespace infini::ccl {