[AOTI] support freezing for MKLDNN

[chunyuan-w]

Stack from ghstack (oldest at bottom):

• -> [AOTI] support freezing for MKLDNN #124350

Description

Fixes #114450. This PR builds upon the work from @imzhuhl done in #114451.

This PR requires #126137 to land firstly.

We leverage the serialization and deserialization API from oneDNN v3.4.1 to save the opaque MKLDNN tensor during the compilation and restore the opaque tensor when loading the compiled .so.
ideep version is updated so that we won't break any pipeline even if third_party/ideep is not updated at the same time.

Test plan:

python -u test/inductor/test_aot_inductor.py -k AOTInductorTestNonABICompatibleCpu.test_freezing_non_abi_compatible_cpu
python -u test/inductor/test_aot_inductor.py -k AOTInductorTestNonABICompatibleCpu.test_conv_freezing_non_abi_compatible_cpu
python -u test/inductor/test_aot_inductor.py -k AOTInductorTestNonABICompatibleCpu.test_deconv_freezing_non_abi_compatible_cpu
python -u test/inductor/test_aot_inductor.py -k AOTInductorTestNonABICompatibleCpu.test_linear_freezing_non_abi_compatible_cpu

TODOs in follow-up PRs

1. We found that using AOTI_TORCH_CHECK will cause performance drop on several models (DistillGPT2, MBartForConditionalGeneration, T5ForConditionalGeneration, T5Small) compared with JIT Inductor which uses TORCH_CHECK. This may need further discussion how to address (AOTI_TORCH_CHECK is introduced in
   [aot_inductor] replace TORCH_CHECK with AOTI_CHECK in the generate cpp code #119220).
2. Freezing in non-ABI compatible mode will work with the support in this PR. While for ABI compatible mode, we need to firstly address this issue: AssertionError: None, i.e. optional output is not supported.
   

   pytorch/torch/_inductor/codegen/cpp_wrapper_cpu.py

   Lines 2023 to 2024 in 6c4f43f

   	def extract_output_name(out):
   	assert out is not None, "None, i.e. optional output is not supported"

cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @voznesenskym @penguinwu @EikanWang @Guobing-Chen @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @amjames @desertfire @chauhang

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/124350

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 39ed894 with merge base 5f15110 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[clee2000]

@pytorchbot revert -m "Seems to have broken inductor/test_aot_inductor.py::AOTInductorTestNonABICompatibleCpu::test_freezing_non_abi_compatible_cpu https://hud.pytorch.org/pytorch/pytorch/commit/654afb6f3ae3ddbd926a753f9af95a6f6e22131c https://github.com/pytorch/pytorch/actions/runs/9224838183/job/25382780192" -c landrace

[pytorchmergebot]

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

[pytorchmergebot]

@chunyuan-w your PR has been successfully reverted.

[chunyuan-w]

@pytorchbot revert -m "Seems to have broken inductor/test_aot_inductor.py::AOTInductorTestNonABICompatibleCpu::test_freezing_non_abi_compatible_cpu https://hud.pytorch.org/pytorch/pytorch/commit/654afb6f3ae3ddbd926a753f9af95a6f6e22131c https://github.com/pytorch/pytorch/actions/runs/9224838183/job/25382780192" -c landrace

This should be caused by conflicts with #126068 which also landed yesterday at similar time. When CI was executed on this PR yesterday, the change from #126068 is not yet included in this PR so this issue is not exposed at that moment.

I fixed this issue in this commit 39ed894.

[chunyuan-w]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here