support aarch64 uos

python/jittor/__init__.py

@@ -9,7 +9,7 @@
 # file 'LICENSE.txt', which is part of this source code package.
 # ***************************************************************
 
-__version__ = '1.2.3.32'
+__version__ = '1.2.3.33'
 from jittor_utils import lock
 with lock.lock_scope():
     ori_int = int

python/jittor/compile_extern.py

@@ -23,29 +23,35 @@ def search_file(dirs, name, prefer_version=()):
 def install_mkl(root_folder):
     # origin url is
     # url = "https://github.com/intel/mkl-dnn/releases/download/v1.0.2/mkldnn_lnx_1.0.2_cpu_gomp.tgz"
-    url = "https://cloud.tsinghua.edu.cn/f/da02bf62b55b4aa3b8ee/?dl=1"
+    import platform
-    filename = "mkldnn_lnx_1.0.2_cpu_gomp.tgz"
+    if platform.system()=="Linux":
-    # newest version for oneDNN
+        if platform.machine()=='x86_64':
-    # url = "https://github.com/oneapi-src/oneDNN/releases/download/v2.2/dnnl_lnx_2.2.0_cpu_gomp.tgz"
+            filename = "dnnl_lnx_2.2.0_cpu_gomp.tgz"
-    # filename = "dnnl_lnx_2.2.0_cpu_gomp.tgz"
+            md5 = "35bbbdf550a9d8ad54db798e372000f6"
+        elif platform.machine()=='aarch64':
+            filename = "dnnl_lnx_2.2.0_cpu_gomp_aarch64.tgz"
+            md5 = "72cf9b0b8fd6c3c786d35a9daaee22b8"
+        else:
+            raise RuntimeError(f"platform.machine()=={platform.machine()} not support yet,"
+            " Please contact us on https://github.com/jittor/jittor ")
+    else:
+        raise RuntimeError(f"platform.machine()=={platform.machine()} not support yet,"
+        " Please contact us on https://github.com/jittor/jittor ")
+
+    url = "https://cg.cs.tsinghua.edu.cn/jittor/assets/" + filename
     fullname = os.path.join(root_folder, filename)
     dirname = os.path.join(root_folder, filename.replace(".tgz",""))
 
-    if not os.path.isfile(os.path.join(dirname, "examples", "test")):
+    if not os.path.isfile(os.path.join(dirname, "lib", "libmkldnn.so")):
         LOG.i("Downloading mkl...")
-        download_url_to_local(url, filename, root_folder, "47187284ede27ad3bd64b5f0e7d5e730")
+        download_url_to_local(url, filename, root_folder, md5)
-        # newest version for oneDNN
-        # download_url_to_local(url, filename, root_folder, "35bbbdf550a9d8ad54db798e372000f6")
         import tarfile
 
         with tarfile.open(fullname, "r") as tar:
             tar.extractall(root_folder)
 
         assert 0 == os.system(f"cd {dirname}/examples && "
-            f"{cc_path} -std=c++14 cpu_cnn_inference_f32.cpp -Ofast -lmkldnn -I ../include -L ../lib -o test && LD_LIBRARY_PATH=../lib/ ./test")
+            f"{cc_path} -std=c++14 cnn_inference_f32.cpp -Ofast -lmkldnn -I ../include -L ../lib -o test && LD_LIBRARY_PATH=../lib/ ./test")
-        # newest version for oneDNN
-        # assert 0 == os.system(f"cd {dirname}/examples && "
-        #     f"{cc_path} -std=c++14 cnn_inference_f32.cpp -Ofast -lmkldnn -I ../include -L ../lib -o test && LD_LIBRARY_PATH=../lib/ ./test")
 
 def setup_mkl():
     global mkl_ops, use_mkl
@@ -80,7 +86,7 @@ def setup_mkl():
         install_mkl(mkl_path)
         mkl_home = ""
         for name in os.listdir(mkl_path):
-            if name.startswith("mkldnn_lnx") and os.path.isdir(os.path.join(mkl_path, name)):
+            if name.startswith("dnnl") and os.path.isdir(os.path.join(mkl_path, name)):
                 mkl_home = os.path.join(mkl_path, name)
                 break
         assert mkl_home!=""
@@ -197,8 +203,14 @@ def setup_cuda_lib(lib_name, link=True, extra_flags=""):
 
         if lib_name == "cublas" and nvcc_version[0] >= 10:
             # manual link libcublasLt.so
-            cublas_lt_lib_path = search_file([cuda_lib, extra_lib_path, "/usr/lib/x86_64-linux-gnu", "/usr/lib"], f"libcublasLt.so", nvcc_version)
+            try:
-            ctypes.CDLL(cublas_lt_lib_path, dlopen_flags)
+                cublas_lt_lib_path = search_file([cuda_lib, extra_lib_path, "/usr/lib/x86_64-linux-gnu", "/usr/lib"], f"libcublasLt.so", nvcc_version)
+                ctypes.CDLL(cublas_lt_lib_path, dlopen_flags)
+            except:
+                # some aarch64 os, such as uos with FT2000 cpu,
+                # it's cuda 10 doesn't have libcublasLt.so
+                pass
+
 
 
         if lib_name == "cudnn":

python/jittor/compiler.py

@@ -12,6 +12,7 @@ import inspect
 import datetime
 import threading
 import ctypes
+import platform
 from ctypes import cdll
 from ctypes.util import find_library
 
@@ -1038,6 +1039,8 @@ if os.path.isfile(version_file) and not os.path.isdir(os.path.join(jittor_path,
     os_key = os_type.get(os_id, "ubuntu")
     if "os_key" in os.environ:
         os_key = os.environ['os_key']
+    if platform.machine()=='aarch64':
+        os_key += '-aarch64'
     LOG.i("OS type:", os_id, " OS key:", os_key)
     key += '-' + os_key + '.o'
     # TODO: open the website

python/jittor/script/build_aarch64_mkl.sh

@@ -0,0 +1,31 @@
+# wget https://github.com/oneapi-src/oneDNN/archive/refs/tags/v2.2.zip
+# extract zip
+# cd to root folder
+
+mkdir -p build
+cd build
+make clean
+export CC=aarch64-linux-gnu-gcc-8
+export CXX=aarch64-linux-gnu-g++-8
+cmake .. \
+    -DCMAKE_SYSTEM_NAME=Linux \
+    -DCMAKE_SYSTEM_PROCESSOR=AARCH64 \
+    -DCMAKE_LIBRARY_PATH=/usr/aarch64-linux-gnu/lib \
+    -DCMAKE_BUILD_TYPE=Release
+    # -DCMAKE_SHARED_LINKER_FLAGS=' -lm ' \
+make -j8
+
+name=dnnl_lnx_2.2.0_cpu_gomp_aarch64
+mkdir -p $name
+cp -r ../include ./$name/
+mkdir -p ./$name/lib
+cp ./src/libmkldnn.so ./$name/lib/libmkldnn.so
+cp -r ../examples ./$name/
+cp ./include/oneapi/dnnl/* ./$name/include/oneapi/dnnl/
+
+tar -acvf $name.tgz ./$name/
+
+rsync -avPu $name.tgz jittor-web:Documents/jittor-blog/assets/
+ssh jittor-web Documents/jittor-blog.git/hooks/post-update
+echo "https://cg.cs.tsinghua.edu.cn/jittor/assets/$name.tgz"
+md5sum $name.tgz

python/jittor/src/utils/log.cc

@@ -322,6 +322,34 @@ but you can hot fix it by this command:
         )";
 }
 
+static inline void check_cuda_gcc_version(const string& output) {
+    /*  if such error occur: 
+    error: identifier "__is_assignable" is undefined
+    this means your gcc version is not match with nvcc,
+    for example, nvcc 10 support gcc<=7, nvcc 11 support gcc<=9,
+
+    https://gist.github.com/ax3l/9489132
+    */
+    string pat = "__is_assignable";
+    auto id = output.find(pat);
+    if (id == string::npos) return;
+    LOGf << output << R"(
+!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+Dear user, your nvcc and gcc version are still not match
+after dirty hack, your should install the correct version of g++
+or nvcc, for example, nvcc 10 support g++<=7, nvcc 11 support g++<=9,
+here is the NVCC Compatibility Matrix:
+    https://gist.github.com/ax3l/9489132
+Please install correct version of gcc, for example:
+    >>> sudo apt install g++-7
+After your g++ is installed, using enviroment variable `cc_path` to
+tell jittor use the correct version of g++, for example:
+    >>> cc_path='g++-7' python3.7 -m jittor.test.test_core
+If you still have problems, please contact us:
+    https://github.com/Jittor/jittor/issues
+    )";
+}
+
 int system_popen(const char* cmd) {
     char buf[BUFSIZ];
     string cmd2;
@@ -342,6 +370,7 @@ int system_popen(const char* cmd) {
     }
     if (ret) {
         check_cuda_unsupport_version(output);
+        check_cuda_gcc_version(output);
     }
     return ret;
 }