From 0e2597085301813d24407461e52ecb6a054dc499 Mon Sep 17 00:00:00 2001
From: Jack Myers <jack.myers@intel.com>
Date: Tue, 18 Mar 2025 20:46:23 +0000
Subject: [PATCH] fix: re-add switch case for once writable query

A change related to the tbx fault manager
incorrectly removed a switch case from
`AubHelper::isOneTimeAubWritableAllocationType`.

This fixes that and refactors some APIs to prevent
similar mistakes from happening again by cleaning
up logic.

Addresses show stopper for pre-si pytorch workflows.

Resolves: NEO-14399
Signed-off-by: Jack Myers <jack.myers@intel.com>
---
 shared/source/aub/aub_helper.cpp                           | 4 +++-
 .../command_stream/tbx_command_stream_receiver_hw.inl      | 4 +++-
 shared/source/debug_settings/debug_settings_manager.h      | 7 ++++---
 shared/test/unit_test/aub/aub_helper_tests.cpp             | 3 ++-
 4 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/shared/source/aub/aub_helper.cpp b/shared/source/aub/aub_helper.cpp
index 713358063d..4a151caf16 100644
--- a/shared/source/aub/aub_helper.cpp
+++ b/shared/source/aub/aub_helper.cpp
@@ -41,8 +41,10 @@ bool AubHelper::isOneTimeAubWritableAllocationType(const AllocationType &type) {
     case AllocationType::assertBuffer:
     case AllocationType::tagBuffer:
     case AllocationType::syncDispatchToken:
+        return true;
     case AllocationType::bufferHostMemory:
-        return (NEO::debugManager.flags.SetBufferHostMemoryAlwaysAubWritable.get() ? false : true) || (NEO::debugManager.flags.EnableTbxPageFaultManager.get() == 1);
+        return NEO::debugManager.isTbxPageFaultManagerEnabled() ||
+               (NEO::debugManager.flags.SetBufferHostMemoryAlwaysAubWritable.get() == false);
     default:
         return false;
     }
diff --git a/shared/source/command_stream/tbx_command_stream_receiver_hw.inl b/shared/source/command_stream/tbx_command_stream_receiver_hw.inl
index fc0524bc7d..bda0b5a76d 100644
--- a/shared/source/command_stream/tbx_command_stream_receiver_hw.inl
+++ b/shared/source/command_stream/tbx_command_stream_receiver_hw.inl
@@ -85,7 +85,9 @@ TbxCommandStreamReceiverHw<GfxFamily>::~TbxCommandStreamReceiverHw() {
 template <typename GfxFamily>
 bool TbxCommandStreamReceiverHw<GfxFamily>::isAllocTbxFaultable(GraphicsAllocation *gfxAlloc) {
     // indicates host memory not managed by the driver
-    if (gfxAlloc->getDriverAllocatedCpuPtr() == nullptr || !debugManager.isTbxPageFaultManagerEnabled() || this->getTbxPageFaultManager() == nullptr) {
+    if ((gfxAlloc->getDriverAllocatedCpuPtr() == nullptr) ||
+        (debugManager.isTbxPageFaultManagerEnabled() == false) ||
+        (this->getTbxPageFaultManager() == nullptr)) {
         return false;
     }
     auto allocType = gfxAlloc->getAllocationType();
diff --git a/shared/source/debug_settings/debug_settings_manager.h b/shared/source/debug_settings/debug_settings_manager.h
index 7bf9d42ffb..b9d8151787 100644
--- a/shared/source/debug_settings/debug_settings_manager.h
+++ b/shared/source/debug_settings/debug_settings_manager.h
@@ -162,9 +162,10 @@ class DebugSettingsManager : NEO::NonCopyableAndNonMovableClass {
 
     inline bool isTbxPageFaultManagerEnabled() {
         auto setCsr = flags.SetCommandStreamReceiver.get();
-        auto tbxMngrFlag = flags.EnableTbxPageFaultManager.get();
-        auto isTbxMode = (setCsr == static_cast<int32_t>(CommandStreamReceiverType::tbx)) || (setCsr == static_cast<int32_t>(CommandStreamReceiverType::tbxWithAub));
-        return tbxMngrFlag && isTbxMode;
+        auto isTbxMode = (setCsr == static_cast<int32_t>(CommandStreamReceiverType::tbx)) ||
+                         (setCsr == static_cast<int32_t>(CommandStreamReceiverType::tbxWithAub));
+        auto isFaultManagerEnabledInEnvVars = flags.EnableTbxPageFaultManager.get();
+        return isFaultManagerEnabledInEnvVars && isTbxMode;
     }
 
   protected:
diff --git a/shared/test/unit_test/aub/aub_helper_tests.cpp b/shared/test/unit_test/aub/aub_helper_tests.cpp
index 9ce95f9cf0..d7e3473f1b 100644
--- a/shared/test/unit_test/aub/aub_helper_tests.cpp
+++ b/shared/test/unit_test/aub/aub_helper_tests.cpp
@@ -129,9 +129,10 @@ TEST(AubHelper, givenAllocationTypeWhenAskingIfOneTimeWritableThenReturnCorrectR
     }
 }
 
-TEST(AubHelper, givenSetBufferHostMemoryAlwaysAubWritableAndDisabledTbxFaultMngrWhenAskingIfBufferHostMemoryAllocationIsOneTimeAubWritableThenReturnCorrectResult) {
+TEST(AubHelper, givenAlwaysAubWritableAndEnableTbxFaultManagerSetExternallyThenAllocationIsOneTimeAubWritableShouldReturnCorrectResult) {
     DebugManagerStateRestore stateRestore;
     NEO::debugManager.flags.EnableTbxPageFaultManager.set(0);
+    NEO::debugManager.flags.SetCommandStreamReceiver.set(2);
 
     for (auto isAlwaysAubWritable : {false, true}) {
         for (auto isTbxFaultManagerEnabled : {false, true}) {