From b3a44f2c0282af6b27cf1f5baaa02b7278522a15 Mon Sep 17 00:00:00 2001
From: Patrick Gartung <gartung@fnal.gov>
Date: Tue, 4 Mar 2025 20:08:44 +0100
Subject: [PATCH 1/2] Backport commit 2ae24507835242a8d217b40fde629ef671b904de
 from oneDNN/main to allow stack unwinding

---
 tensorflow/workspace2.bzl                     |  1 +
 third_party/mkl_dnn/onednn_unwind_stack.patch | 33 +++++++++++++++++++
 2 files changed, 34 insertions(+)
 create mode 100644 third_party/mkl_dnn/onednn_unwind_stack.patch

diff --git a/tensorflow/workspace2.bzl b/tensorflow/workspace2.bzl
index e906f65b7efb1d..917d9b6c0d9327 100644
--- a/tensorflow/workspace2.bzl
+++ b/tensorflow/workspace2.bzl
@@ -195,6 +195,7 @@ def _tf_repositories():
             "//third_party/mkl_dnn:onednn_acl_fixed_format_kernels.patch",
             "//third_party/mkl_dnn:onednn_acl_depthwise_convolution.patch",
             "//third_party/mkl_dnn:onednn_acl_threadpool_scheduler.patch",
+            "//third_party/mkl_dnn:onednn_unwind_stack.patch",
         ],
         sha256 = "a50993aa6265b799b040fe745e0010502f9f7103cc53a9525d59646aef006633",
         strip_prefix = "oneDNN-2.7.3",
diff --git a/third_party/mkl_dnn/onednn_unwind_stack.patch b/third_party/mkl_dnn/onednn_unwind_stack.patch
new file mode 100644
index 00000000000000..2920d7e5685703
--- /dev/null
+++ b/third_party/mkl_dnn/onednn_unwind_stack.patch
@@ -0,0 +1,33 @@
+commit 2ae24507835242a8d217b40fde629ef671b904de
+Author: Wierschem, Keola <keola.wierschem@intel.com>
+Date:   Thu Sep 15 15:31:30 2022 -0700
+
+    cpu: x64: update RBP on entry to jit code
+
+diff --git a/src/cpu/x64/jit_generator.hpp b/src/cpu/x64/jit_generator.hpp
+index 60ec65211c..6399bc64d1 100644
+--- a/src/cpu/x64/jit_generator.hpp
++++ b/src/cpu/x64/jit_generator.hpp
+@@ -103,8 +103,8 @@ static inline void tc_configure_tile(
+ 
+ #ifdef XBYAK64
+ constexpr Xbyak::Operand::Code abi_save_gpr_regs[] = {
+-        Xbyak::Operand::RBX,
+         Xbyak::Operand::RBP,
++        Xbyak::Operand::RBX,
+         Xbyak::Operand::R12,
+         Xbyak::Operand::R13,
+         Xbyak::Operand::R14,
+@@ -194,8 +194,11 @@ public:
+                 uni_vmovdqu(ptr[rsp + i * xmm_len],
+                         Xbyak::Xmm(xmm_to_preserve_start + i));
+         }
+-        for (size_t i = 0; i < num_abi_save_gpr_regs; ++i)
++        for (size_t i = 0; i < num_abi_save_gpr_regs; ++i) {
+             push(Xbyak::Reg64(abi_save_gpr_regs[i]));
++            // Stack magic: save rsp into rbp state to be able to unwind stack.
++            if (i == 0) mov(rbp, rsp);
++        }
+ #ifndef DNNL_ENABLE_MEM_DEBUG
+         // do not use RBP in mem debug mode to enable backtracing from jit code
+         if (is_valid_isa(avx512_core)) {

From 3b6f86cfe7ed46eec396e8b2a2d1366258db33c4 Mon Sep 17 00:00:00 2001
From: Patrick Gartung <gartung@fnal.gov>
Date: Tue, 4 Mar 2025 20:10:45 +0100
Subject: [PATCH 2/2] Revert "    Enable memory debug"

This reverts commit 8ea181ff8b5bb56ae5de02c7ac1837cc8afee049.
---
 third_party/mkl_dnn/mkldnn_v1.BUILD | 1 -
 1 file changed, 1 deletion(-)

diff --git a/third_party/mkl_dnn/mkldnn_v1.BUILD b/third_party/mkl_dnn/mkldnn_v1.BUILD
index 1cd1f8be452dfd..26c61c7bffe09d 100644
--- a/third_party/mkl_dnn/mkldnn_v1.BUILD
+++ b/third_party/mkl_dnn/mkldnn_v1.BUILD
@@ -101,7 +101,6 @@ _COPTS_LIST = select({
     "-UUSE_MKL",
     "-UUSE_CBLAS",
     "-DDNNL_ENABLE_MAX_CPU_ISA",
-    "-DDNNL_ENABLE_MEM_DEBUG",
 ] + tf_openmp_copts()
 
 _INCLUDES_LIST = [