Merge pull request openucx#6020 from petro-rudenko/jucx/memtype-const…

…ants JUCX: UCS Memory type constants.
kunpengcompute · Dec 14, 2020 · 475ab4e · 475ab4e
2 parents 0580ef0 + 2779ba6
commit 475ab4e
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/bindings/java/src/main/java/org/openucx/jucx/ucs/UcsConstants.java b/bindings/java/src/main/java/org/openucx/jucx/ucs/UcsConstants.java
@@ -72,6 +72,20 @@ public static class STATUS {
         public static int UCS_ERR_LAST;
     }
 
+    public static class MEMORY_TYPE {
+        static {
+            load();
+        }
+
+        public static int UCS_MEMORY_TYPE_HOST;          // Default system memory
+        public static int UCS_MEMORY_TYPE_CUDA;          // NVIDIA CUDA memory
+        public static int UCS_MEMORY_TYPE_CUDA_MANAGED;  // NVIDIA CUDA managed (or unified) memory
+        public static int UCS_MEMORY_TYPE_ROCM;          // AMD ROCM memory
+        public static int UCS_MEMORY_TYPE_ROCM_MANAGED;  // AMD ROCM managed system memory
+        public static int UCS_MEMORY_TYPE_LAST;
+        public static int UCS_MEMORY_TYPE_UNKNOWN;
+    }
+
     private static void load() {
         NativeLibs.load();
         loadConstants();

diff --git a/bindings/java/src/main/native/ucs_constants.cc b/bindings/java/src/main/native/ucs_constants.cc
@@ -54,4 +54,14 @@ Java_org_openucx_jucx_ucs_UcsConstants_loadConstants(JNIEnv *env, jclass ucs_cla
     JUCX_DEFINE_INT_CONSTANT(UCS_ERR_ENDPOINT_TIMEOUT);
     JUCX_DEFINE_INT_CONSTANT(UCS_ERR_LAST_ENDPOINT_FAILURE);
     JUCX_DEFINE_INT_CONSTANT(UCS_ERR_LAST);
+
+    // Memory type
+    cls = env->FindClass("org/openucx/jucx/ucs/UcsConstants$MEMORY_TYPE");
+    JUCX_DEFINE_INT_CONSTANT(UCS_MEMORY_TYPE_HOST);
+    JUCX_DEFINE_INT_CONSTANT(UCS_MEMORY_TYPE_CUDA);
+    JUCX_DEFINE_INT_CONSTANT(UCS_MEMORY_TYPE_CUDA_MANAGED);
+    JUCX_DEFINE_INT_CONSTANT(UCS_MEMORY_TYPE_ROCM);
+    JUCX_DEFINE_INT_CONSTANT(UCS_MEMORY_TYPE_ROCM_MANAGED);
+    JUCX_DEFINE_INT_CONSTANT(UCS_MEMORY_TYPE_LAST);
+    JUCX_DEFINE_INT_CONSTANT(UCS_MEMORY_TYPE_UNKNOWN);
 }