support M1 chip

README.md

@@ -95,7 +95,7 @@ We provide some jupyter notebooks to help you quick start with Jittor.
 
 Jittor environment requirements:
 
-* System: **Linux**(e.g. Ubuntu/CentOS/Arch), **macOS** (x86_64), or **Windows Subsystem of Linux (WSL)**
+* System: **Linux**(e.g. Ubuntu/CentOS/Arch), **macOS**, or **Windows Subsystem of Linux (WSL)**
 * Python version >= 3.7
 * CPU compiler (require at least one of the following)
     * g++ (>=5.4.0)

README.src.md

@@ -132,7 +132,7 @@ Jittor 提供了三种安装方法：docker，pip和手动安装：
 
 Jittor environment requirements:
 
-* System: **Linux**(e.g. Ubuntu/CentOS/Arch), **macOS** (x86_64), or **Windows Subsystem of Linux (WSL)**
+* System: **Linux**(e.g. Ubuntu/CentOS/Arch), **macOS**, or **Windows Subsystem of Linux (WSL)**
 * Python version >= 3.7
 * CPU compiler (require at least one of the following)
     * g++ (>=5.4.0)

python/jittor/compile_extern.py

@@ -102,8 +102,12 @@ def setup_mkl():
         extra_flags = f" -I'{mkl_include_path}' -L'{mkl_lib_path}' -lmkldnn -Wl,-rpath='{mkl_lib_path}' "
 
     elif platform.system() == 'Darwin':
-        mkl_lib_name = "/usr/local/lib/libmkldnn.dylib"
-        assert os.path.exists(mkl_lib_name), "Not found onednn, please install it by the command 'brew install onednn@2.2.3'"
+        mkl_lib_paths = [
+            "/usr/local/lib/libmkldnn.dylib",       # x86_64
+            "/opt/homebrew/lib/libmkldnn.dylib",    # arm64
+        ]
+        if not any([os.path.exists(lib) for lib in mkl_lib_paths]):
+            raise RuntimeError("Not found onednn, please install it by the command 'brew install onednn@2.2.3'")
         extra_flags = f" -lmkldnn "
 
     mkl_op_dir = os.path.join(jittor_path, "extern", "mkl", "ops")

python/jittor/compiler.py

@@ -903,14 +903,24 @@ gdb_path = try_find_exe('gdb')
 addr2line_path = try_find_exe('addr2line')
 has_pybt = check_pybt(gdb_path, python_path)
 
-cc_flags += " -Wall -Werror -Wno-unknown-pragmas -std=c++14 -fPIC -march=native "
+cc_flags += " -Wall -Werror -Wno-unknown-pragmas -std=c++14 -fPIC "
+# 1. Arch/CPU specific optimization
+if platform.machine() == "x86_64":
+    cc_flags += " -march=native " 
+elif platform.machine() == 'arm64' and platform.system() == "Darwin":
+    cc_flags += " -mcpu=apple-a14 "
 cc_flags += " -fdiagnostics-color=always "
+# 2. Non standard include path
+if platform.system() == 'Darwin' and platform.machine() == 'arm64':
+    cc_flags += " -I/opt/homebrew/include "
+# 3. User specified flags
 if "cc_flags" in os.environ:
     cc_flags += os.environ["cc_flags"] + ' '
+
 link_flags = " -lstdc++ -ldl -shared "
 if platform.system() == 'Darwin':
     # TODO: if not using apple clang, there is no need to add -lomp
-    link_flags += "-undefined dynamic_lookup -lomp "
+    link_flags += "-undefined dynamic_lookup -lomp -L/opt/homebrew/lib"
 
 core_link_flags = ""
 opt_flags = ""
@@ -1042,8 +1052,10 @@ if platform.system() == 'Linux':
             a = line.split('=')
             if len(a) != 2: continue
             os_release[a[0]] = a[1].replace("\"", "")
+        os_arch = ''
 elif platform.system() == 'Darwin':
     os_release = {'ID' : 'macos'}
+    os_arch = platform.machine()
 
 os_type = {
     "ubuntu": "ubuntu",
@@ -1061,7 +1073,8 @@ if os.path.isfile(version_file) and not os.path.isdir(os.path.join(jittor_path,
     # key = f"{version}-{cc_type}-{'cuda' if has_cuda else 'cpu'}.o"
     key = f"{version}-g++-cpu"
     os_id = os_release["ID"]
-    os_key = os_type.get(os_id, "ubuntu")
+    os_key = os_type.get(os_id, "ubuntu") 
+    os_key += '-' + os_arch if os_arch else ''
     if "os_key" in os.environ:
         os_key = os.environ['os_key']
     LOG.i("OS type:", os_id, " OS key:", os_key)

python/jittor/src/jit_key.cc

@@ -6,13 +6,12 @@
 // ***************************************************************
 #include <sys/mman.h>
 #include <sstream>
+#include <unistd.h>
 #include "jit_key.h"
 #include "utils/str_utils.h"
 
 namespace jittor {
 
-const int page_size = 4*1024;
-
 extern thread_local size_t protected_page;
 
 static size_t get_buffer_end_page(size_t buffer_end) {
@@ -21,23 +20,23 @@ static size_t get_buffer_end_page(size_t buffer_end) {
     //  |       |       |       |       |
     //  buffer:    xxxxxxxxxxxxxxxxxxxxxxxx
     //                          ^  buffer_end_page
-    size_t buffer_end_page = buffer_end - buffer_end % page_size;
-    if (buffer_end_page + page_size-1 > buffer_end)
-        buffer_end_page -= page_size;
+    size_t buffer_end_page = buffer_end - buffer_end % getpagesize();
+    if (buffer_end_page + getpagesize()-1 > buffer_end)
+        buffer_end_page -= getpagesize();
     return buffer_end_page;
 }
 
 JitKey::JitKey() {
     auto buffer_end_page = get_buffer_end_page((size_t)&buffer[buffer_size-1]);
     LOGvv << "protect page" << (void*)buffer_end_page;
-    ASSERT(0==mprotect((void*)buffer_end_page, page_size, PROT_NONE));
+    ASSERT(0==mprotect((void*)buffer_end_page, getpagesize(), PROT_NONE));
     protected_page = buffer_end_page;
 }
 
 JitKey::~JitKey() {
     auto buffer_end_page = get_buffer_end_page((size_t)&buffer[buffer_size-1]);
     LOGvv << "un-protect page" << (void*)buffer_end_page;
-    mprotect((void*)buffer_end_page, page_size, PROT_READ|PROT_WRITE|PROT_EXEC);
+    mprotect((void*)buffer_end_page, getpagesize(), PROT_READ|PROT_WRITE|PROT_EXEC);
     protected_page = 0;
 }
 

python/jittor/utils/polish.py

@@ -63,6 +63,7 @@ os_name_system_dict = {
 for os_name, os_type in os_name_system_dict.items():
     if platform.system() != os_type:
         continue
+    os_arch = platform.machine() if os_type == 'macos' else ''
 
     for cc_type in ["g++"]:
         for device in ["cpu"]:
@@ -72,13 +73,15 @@ for os_name, os_type in os_name_system_dict.items():
             env += cname
             # use core2 arch, avoid using avx instructions
             # TODO: support more archs, such as arm, or use ir(GIMPLE or LLVM)
-            env += " cc_flags='-march=core2' "
+            if platform.machine() == "x86_64":
+                env += " cc_flags='-march=core2' "
             if device == "cpu":
                 env += "nvcc_path='' "
             elif jt.flags.nvcc_path == "":
                 env = "unset nvcc_path && " + env
             cmd = f"{env} {sys.executable} -c 'import jittor'"
             if key != 'ubuntu': key += '-' + os_name
+            if os_arch : key += '-' + os_arch
             if os_name == 'centos':
                 run_in_centos(env)
                 obj_path = home + f"/.cache/centos/build/{cc_type}/{device}/{cname}/obj_files"

python/jittor_utils/__init__.py

@@ -316,6 +316,7 @@ py3_config_paths = [
     sys.executable + "-config",
     f"/usr/bin/python3.{sys.version_info.minor}-config",
     f"/usr/local/bin/python3.{sys.version_info.minor}-config",
+    f'/opt/homebrew/bin/python3.{sys.version_info.minor}-config',
     os.path.dirname(sys.executable) + "/python3-config",
 ]
 if "python_config_path" in os.environ: