[aarch64] patch mkl-dnn to accelerate torch.compile()

pytorch · Dec 3, 2023 · 93ca5e2 · 93ca5e2
1 parent b75ebbd
commit 93ca5e2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/aarch64_linux/aarch64_wheel_ci_build.py b/aarch64_linux/aarch64_wheel_ci_build.py
@@ -108,6 +108,9 @@ def parse_arguments():
     # work around to fix Raspberry pie crash
     print("Applying mkl-dnn patch to fix readdir crash")
     os.system("cd /pytorch/third_party/ideep/mkl-dnn && patch -p1 < /builder/mkldnn_fix/aarch64-fix-readdir-crash.patch")
+    # patch acl inner product to accelerate torch.compile() path
+    print("Applying mkl-dnn patch to acl inner product")
+    os.system("cd /pytorch/third_party/ideep/mkl-dnn && patch -p1 < /builder/mkldnn_fix/onednn-cpu-aarch64-remove-weight-format-checking.patch")
     os.system(f"cd /pytorch; {build_vars} python3 setup.py bdist_wheel")
     pytorch_wheel_name = complete_wheel("pytorch")
     print(f"Build Compelete. Created {pytorch_wheel_name}..")
diff --git a/aarch64_linux/build_aarch64_wheel.py b/aarch64_linux/build_aarch64_wheel.py
@@ -555,6 +555,7 @@ def start_build(host: RemoteHost, *,
         print("build pytorch with mkldnn+acl backend")
         build_vars += " USE_MKLDNN=ON USE_MKLDNN_ACL=ON"
         host.run_cmd(f"cd $HOME && git clone https://github.com/pytorch/builder.git")
+        host.run_cmd(f"cd $HOME/pytorch/third_party/ideep/mkl-dnn && patch -p1 < $HOME/builder/mkldnn_fix/onednn-cpu-aarch64-remove-weight-format-checking.patch")
         host.run_cmd(f"cd $HOME/pytorch && export ACL_ROOT_DIR=$HOME/ComputeLibrary && {build_vars} python3 setup.py bdist_wheel{build_opts}")
         print('Repair the wheel')
         pytorch_wheel_name = host.list_dir("pytorch/dist")[0]