From 93ca5e2f5476e18b8ba77a48a3b4fd309e7c7538 Mon Sep 17 00:00:00 2001 From: Sunita Nadampalli Date: Sun, 3 Dec 2023 13:09:10 -0600 Subject: [PATCH] [aarch64] patch mkl-dnn to accelerate torch.compile() --- aarch64_linux/aarch64_wheel_ci_build.py | 3 +++ aarch64_linux/build_aarch64_wheel.py | 1 + 2 files changed, 4 insertions(+) diff --git a/aarch64_linux/aarch64_wheel_ci_build.py b/aarch64_linux/aarch64_wheel_ci_build.py index 3b772847c5..63ce98735d 100755 --- 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 index 9efd2e6ae5..b8fae3febf 100755 --- 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]