support builtin download of onednn because homebrew no longer provides onednn 2.x

README.cn.md

@@ -126,7 +126,7 @@ macOS 请使用 [homebrew](https://brew.sh) 安装额外的依赖。
 
 
 ```bash
-brew install onednn libomp
+brew install libomp
 ```
 
 之后您可以通过 pip 安装 jittor，并测试是否可以成功运行。

README.md

@@ -122,7 +122,7 @@ python3.7 -m jittor.test.test_example
 Please first install additional dependencies with [homebrew](https://brew.sh).
 
 ```bash
-brew install onednn libomp
+brew install libomp
 ```
 
 
@@ -134,7 +134,7 @@ python3.7 -m jittor.test.test_example
 ```
 
 
-Currently jittor only supports CPU in macOS.
+Currently jittor only supports CPU on macOS.
 
 
 ### Windows install

README.src.md

@@ -158,7 +158,7 @@ macOS 请使用 [homebrew](https://brew.sh) 安装额外的依赖。
 Please first install additional dependencies with [homebrew](https://brew.sh).
 
 ```bash
-brew install onednn libomp
+brew install libomp
 ```
 
 之后您可以通过 pip 安装 jittor，并测试是否可以成功运行。
@@ -172,7 +172,7 @@ python3.7 -m jittor.test.test_example
 
 目前在 macOS 中，jittor 只支持 CPU 计算。
 
-Currently jittor only supports CPU in macOS.
+Currently jittor only supports CPU on macOS.
 
 ### Windows安装
 

python/jittor/compile_extern.py

@@ -51,6 +51,13 @@ def install_mkl(root_folder):
         # url = "https://github.com/oneapi-src/oneDNN/releases/download/v2.2/dnnl_win_2.2.0_cpu_vcomp.zip"
         filename = "dnnl_win_2.2.0_cpu_vcomp.zip"
         md5 = "fa12c693b2ec07700d174e1e99d60a7e"
+    elif platform.system() == "Darwin":
+        if platform.machine() == "arm64":
+            filename = "dnnl_mac_2.2.0_cpu_omp_arm64.tgz"
+            md5 = "d8fdf56d3cf618685d22d18f08119f88"
+        else:
+            filename = "dnnl_mac_2.2.0_cpu_omp_x86_64"
+            md5 = "6e2f065d6a589c82081536b684768fe6"
     else:
         raise RuntimeError(f"platform.machine()=={platform.machine()} not support yet,"
         " Please contact us on https://github.com/jittor/jittor ")
@@ -61,7 +68,8 @@ def install_mkl(root_folder):
     dirname = os.path.join(root_folder, filename.rsplit(".",1)[0])
 
     if not (os.path.isfile(os.path.join(dirname, "lib", "libmkldnn.so")) or
-        os.path.isfile(os.path.join(dirname, "bin", "dnnl.dll"))):
+        os.path.isfile(os.path.join(dirname, "bin", "dnnl.dll")) or 
+        os.path.isfile(os.path.join(dirname, "lib", "libmkldnn.dylib"))):
         LOG.i("Downloading mkl...")
         download_url_to_local(url, filename, root_folder, md5)
         if fullname.endswith(".zip"):
@@ -81,6 +89,9 @@ def install_mkl(root_folder):
             
             assert 0 == os.system(cmd)
             assert 0 == os.system(f"{dirname}/examples/test")
+        if platform.system() == "Darwin":
+            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 && DYLD_LIBRARY_PATH=../lib/ ./test")
         else:
             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")
@@ -105,48 +116,40 @@ def setup_mkl():
     mkl_include_path = os.environ.get("mkl_include_path")
     mkl_lib_path = os.environ.get("mkl_lib_path")
     
-    if platform.system() == 'Linux' or os.name == 'nt':
+    if mkl_lib_path is None or mkl_include_path is None:
-        if mkl_lib_path is None or mkl_include_path is None:
+        LOG.v("setup mkl...")
-            mkl_install_sh = os.path.join(jittor_path, "script", "install_mkl.sh")
+        # mkl_path = os.path.join(cache_path, "mkl")
-            LOG.v("setup mkl...")
+        # mkl_path decouple with cc_path
-            # mkl_path = os.path.join(cache_path, "mkl")
+        mkl_path = os.path.join(jit_utils.home(), ".cache", "jittor", "mkl")
-            # mkl_path decouple with cc_path
+        
-            mkl_path = os.path.join(jit_utils.home(), ".cache", "jittor", "mkl")
+        make_cache_dir(mkl_path)
-            
+        install_mkl(mkl_path)
-            make_cache_dir(mkl_path)
+        mkl_home = ""
-            install_mkl(mkl_path)
+        for name in os.listdir(mkl_path):
-            mkl_home = ""
+            if name.startswith("dnnl") and os.path.isdir(os.path.join(mkl_path, name)):
-            for name in os.listdir(mkl_path):
+                mkl_home = os.path.join(mkl_path, name)
-                if name.startswith("dnnl") and os.path.isdir(os.path.join(mkl_path, name)):
+                break
-                    mkl_home = os.path.join(mkl_path, name)
+        assert mkl_home!=""
-                    break
+    mkl_include_path = os.path.join(mkl_home, "include")
-            assert mkl_home!=""
+    mkl_lib_path = os.path.join(mkl_home, "lib")
-        mkl_include_path = os.path.join(mkl_home, "include")
-        mkl_lib_path = os.path.join(mkl_home, "lib")
 
-        mkl_lib_name = os.path.join(mkl_lib_path, "libmkldnn.so")
+    mkl_lib_name = os.path.join(mkl_lib_path, "libmkldnn.so")
-        extra_flags = f" -I\"{mkl_include_path}\" -L\"{mkl_lib_path}\" -lmkldnn "
+    extra_flags = f" -I\"{mkl_include_path}\" -L\"{mkl_lib_path}\" -lmkldnn "
-        if os.name == 'nt':
+    if os.name == 'nt':
-            mkl_lib_name = os.path.join(mkl_home, 'bin', 'dnnl.dll')
+        mkl_lib_name = os.path.join(mkl_home, 'bin', 'dnnl.dll')
-            mkl_bin_path = os.path.join(mkl_home, 'bin')
+        mkl_bin_path = os.path.join(mkl_home, 'bin')
-            extra_flags = f" -I\"{mkl_include_path}\"  -L\"{mkl_lib_path}\" -L\"{mkl_bin_path}\" -ldnnl "
+        extra_flags = f" -I\"{mkl_include_path}\"  -L\"{mkl_lib_path}\" -L\"{mkl_bin_path}\" -ldnnl "
-        assert os.path.isdir(mkl_include_path)
+    elif platform.system() == "Darwin":
-        assert os.path.isdir(mkl_lib_path)
+        mkl_lib_name = os.path.join(mkl_lib_path, "libmkldnn.dylib")
-        assert os.path.isfile(mkl_lib_name)
-        LOG.v(f"mkl_include_path: {mkl_include_path}")
-        LOG.v(f"mkl_lib_path: {mkl_lib_path}")
-        LOG.v(f"mkl_lib_name: {mkl_lib_name}")
-        # We do not link manualy, link in custom ops
-        # ctypes.CDLL(mkl_lib_name, dlopen_flags)
 
-    elif platform.system() == 'Darwin':
+    assert os.path.isdir(mkl_include_path)
-        mkl_lib_paths = [
+    assert os.path.isdir(mkl_lib_path)
-            "/usr/local/lib/libdnnl.dylib",       # x86_64
+    assert os.path.isfile(mkl_lib_name)
-            "/opt/homebrew/lib/libdnnl.dylib",    # arm64
+    LOG.v(f"mkl_include_path: {mkl_include_path}")
-        ]
+    LOG.v(f"mkl_lib_path: {mkl_lib_path}")
-        if not any([os.path.exists(lib) for lib in mkl_lib_paths]):
+    LOG.v(f"mkl_lib_name: {mkl_lib_name}")
-            raise RuntimeError("Not found onednn, please install it by the command 'brew install onednn'")
+    # We do not link manualy, link in custom ops
-        extra_flags = f" -ldnnl "
+    # ctypes.CDLL(mkl_lib_name, dlopen_flags)
 
     mkl_op_dir = os.path.join(jittor_path, "extern", "mkl", "ops")
     mkl_op_files = [os.path.join(mkl_op_dir, name) for name in os.listdir(mkl_op_dir)]
@@ -585,10 +588,10 @@ setup_nccl()
 
 setup_cutt()
 
-try:
+# try:
-    setup_mkl()
+setup_mkl()
-except Exception as e:
+# except Exception as e:
-    LOG.w("MKL install failed, msg:", e)
+#     LOG.w("MKL install failed, msg:", e)
 
 setup_cuda_extern()
 

python/jittor/compiler.py

@@ -1034,14 +1034,19 @@ cc_flags += " -fdiagnostics-color=always "
 # 2. Non standard include path
 if platform.system() == 'Darwin':
     # TODO: if not using apple clang, there is no need to add -lomp
-    cc_flags += "-undefined dynamic_lookup -lomp "
+    cc_flags += " -undefined dynamic_lookup -lomp "
-    if os.environ.get('CONDA_PREFIX', None):
+    if os.environ.get("CONDA_PREFIX", None):
         cc_flags += f" -L{os.path.join(os.environ['CONDA_PREFIX'], 'lib')} "
-    if platform.machine() == "arm64":
+    # if platform.machine() == "arm64":
-        cc_flags += " -I/opt/homebrew/include -L/opt/homebrew/lib  "
+    #     cc_flags += " -I/opt/homebrew/include -L/opt/homebrew/lib  "
-        # Homebrew does not symlink the openmp library (libomp >= 15.0.6) into /opt/homebrew/lib
+    # Homebrew does not symlink the openmp library (libomp >= 15.0.6) into /opt/homebrew/lib
-        if os.path.exists('/opt/homebrew/opt/libomp'):
+    homebrew_openmp_paths = [
-            cc_flags += " -I/opt/homebrew/opt/libomp/include -L/opt/homebrew/opt/libomp/lib"  
+        "/opt/homebrew/opt/libomp",
+        "/usr/local/opt/libomp"
+    ]
+    for openmp_path in homebrew_openmp_paths:
+        if os.path.exists(openmp_path):
+            cc_flags += f" -I{openmp_path}/include -L{openmp_path}/lib"
 
 # 3. User specified flags
 if "cc_flags" in os.environ:

python/jittor_utils/__init__.py

@@ -460,16 +460,26 @@ def get_py3_config_path():
         #   python-config path only from sys.executable.
         #   To address this issue, we add predefined paths to search,
         #       - Linux: /usr/bin/python3.x-config
-        #       - macOS (installed via homebrew): /usr/local/bin/python3.x-config
+        #       - macOS:
+        #           - shiped with macOS 13: /Library/Developer/CommandLineTools/Library/Frameworks/
+        #                                   Python3.framework/Versions/3.x/lib/python3.x/config-3.x-darwin/python-config.py
+        #           - installed via homebrew: /usr/local/bin/python3.x-config
         #   There may be issues under other cases, e.g., installed via conda.
         py3_config_paths = [
             os.path.dirname(sys.executable) + f"/python3.{sys.version_info.minor}-config",
             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 platform.system() == "Darwin":
+            if "homebrew" in sys.executable:
+                py3_config_paths.append(f'/opt/homebrew/bin/python3.{sys.version_info.minor}-config')
+            else:
+                py3_config_paths.append(f'/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/'\
+                                        f'Versions/3.{sys.version_info.minor}/lib/python3.{sys.version_info.minor}/'\
+                                        f'config-3.{sys.version_info.minor}-darwin/python-config.py')
+
         if "python_config_path" in os.environ:
             py3_config_paths.insert(0, os.environ["python_config_path"])
 
@@ -498,6 +508,17 @@ def get_py3_include_path():
         ) + '"'
     else:
         _py3_include_path = run_cmd(get_py3_config_path()+" --includes")
+        
+        # macOS (>=13) is shiped with a fake python3-config which outputs wrong include paths
+        # check the include paths and fix them
+        if platform.system() == "Darwin":
+            is_real_path = False
+            for include_path in _py3_include_path.strip().split():
+                if os.path.exists(include_path[2:]):
+                    is_real_path = True
+            if not is_real_path:
+                _py3_include_path = f"-I/Library/Developer/CommandLineTools/Library/Frameworks/"\
+                                    f"Python3.framework/Versions/3.{sys.version_info.minor}/Headers"
     return _py3_include_path
 
 

python/jittor_utils/misc.py

@@ -73,7 +73,7 @@ def download_url_to_local(url, filename, root_folder, md5):
                 os.remove(file_path)
             raise RuntimeError(msg)
     if not check_file_exist(file_path, md5):
-        raise RuntimeError("File downloads failed.")
+        raise RuntimeError(f"MD5 mismatch between the server and the downloaded file {file_path}")