From 19f2cc6fb6d57324057006771120d005cbce24e1 Mon Sep 17 00:00:00 2001
From: George Nash <george.nash@intel.com>
Date: Wed, 10 May 2023 15:50:59 -0700
Subject: [PATCH] Check the AMX tile configuration is unchanged (#15387)

Don't assume the AMX tile configuration will always remain unchanged It
is possible that other code will change the AMX tile configuration.

This change will read the current tile configuration
 - if the tile is un-configured it will be configured
- if the tile is configured but does not match the expected
configuration it will be configured for the expected configuration

This resolves issues seen in unit tests when building OneDNN ep.

### Description
<!-- Describe your changes. -->



### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

---------

Signed-off-by: George Nash <george.nash@intel.com>
---
 onnxruntime/core/mlas/lib/qgemm_kernel_amx.cpp | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/onnxruntime/core/mlas/lib/qgemm_kernel_amx.cpp b/onnxruntime/core/mlas/lib/qgemm_kernel_amx.cpp
index fd9d8c7b58..7c8743026b 100644
--- a/onnxruntime/core/mlas/lib/qgemm_kernel_amx.cpp
+++ b/onnxruntime/core/mlas/lib/qgemm_kernel_amx.cpp
@@ -168,9 +168,12 @@ MlasGemmQuantCopyPackB<MLAS_GEMM_U8S8_KERNEL_AMX>(
 // Tile configure structure
 struct tileconfig_t {
     uint8_t palette_id = 0;
-    uint8_t reserved[15] = {0};
-    uint16_t colb[16] = {0};
-    uint8_t rows[16] = {0};
+    uint8_t start_row = 0;
+    uint8_t reserved1[14] = {0};
+    uint16_t colb[8] = {0};
+    uint8_t reserved2[16] = {0};
+    uint8_t rows[8] = {0};
+    uint8_t reserved3[8] = {0};
 };
 
 template <>
@@ -197,9 +200,12 @@ MlasGemmQuantThreadInit<MLAS_GEMM_U8S8_KERNEL_AMX>()
 
     MlasThreadedBufAlloc(bufsize);
 
-    static thread_local bool tile_configured = false;
     static thread_local struct tileconfig_t tc = {0};
-    if (!tile_configured) {
+    struct tileconfig_t current_tc = {0};
+    _tile_storeconfig(&current_tc);
+
+    if (tc.palette_id == 0 || (std::memcmp(&current_tc.colb, &tc.colb, sizeof(uint16_t) * 8) != 0 &&
+                               std::memcmp(&current_tc.rows, &tc.rows, sizeof(uint8_t) * 8) != 0)) {
         // Filling tile configure structure.
         tc.palette_id = 1;
         for (int t = 0; t < 8; t++) {
@@ -207,7 +213,6 @@ MlasGemmQuantThreadInit<MLAS_GEMM_U8S8_KERNEL_AMX>()
             tc.colb[t] = 64;
         }
         _tile_loadconfig(&tc);
-        tile_configured = true;
     }
 }