From 74824f659ae25c74c9726711bb09785eb0a01d64 Mon Sep 17 00:00:00 2001
From: "Naklicki, Mateusz" <mateusz.naklicki@intel.com>
Date: Mon, 9 Dec 2024 17:21:39 +0000
Subject: [PATCH] fix: enable scratch pages on xekmd

Signed-off-by: Naklicki, Mateusz <mateusz.naklicki@intel.com>
---
 .../os_interface/linux/xe/ioctl_helper_xe.cpp |  5 ++++-
 .../linux/xe/ioctl_helper_xe_tests.cpp        | 22 ++++++++++++-------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/shared/source/os_interface/linux/xe/ioctl_helper_xe.cpp b/shared/source/os_interface/linux/xe/ioctl_helper_xe.cpp
index 1579e5ddd7..6de401ab34 100644
--- a/shared/source/os_interface/linux/xe/ioctl_helper_xe.cpp
+++ b/shared/source/os_interface/linux/xe/ioctl_helper_xe.cpp
@@ -937,10 +937,13 @@ std::unique_ptr<uint8_t[]> IoctlHelperXe::createVmControlExtRegion(const std::op
 
 uint32_t IoctlHelperXe::getFlagsForVmCreate(bool disableScratch, bool enablePageFault, bool useVmBind) {
     xeLog(" -> IoctlHelperXe::%s %d,%d,%d\n", __FUNCTION__, disableScratch, enablePageFault, useVmBind);
-    uint32_t flags = DRM_XE_VM_CREATE_FLAG_LR_MODE;
+    uint32_t flags = DRM_XE_VM_CREATE_FLAG_LR_MODE | DRM_XE_VM_CREATE_FLAG_SCRATCH_PAGE;
     if (enablePageFault) {
         flags |= DRM_XE_VM_CREATE_FLAG_FAULT_MODE;
     }
+    if (disableScratch) {
+        flags &= ~DRM_XE_VM_CREATE_FLAG_SCRATCH_PAGE;
+    }
     return flags;
 }
 
diff --git a/shared/test/unit_test/os_interface/linux/xe/ioctl_helper_xe_tests.cpp b/shared/test/unit_test/os_interface/linux/xe/ioctl_helper_xe_tests.cpp
index 7df0bbddfd..a9eb48dc5e 100644
--- a/shared/test/unit_test/os_interface/linux/xe/ioctl_helper_xe_tests.cpp
+++ b/shared/test/unit_test/os_interface/linux/xe/ioctl_helper_xe_tests.cpp
@@ -458,14 +458,20 @@ TEST(IoctlHelperXeTest, whenGettingFlagsForVmCreateThenPropertValueIsReturned) {
     DrmMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
     auto xeIoctlHelper = std::make_unique<IoctlHelperXe>(drm);
 
-    EXPECT_EQ(static_cast<uint32_t>(DRM_XE_VM_CREATE_FLAG_LR_MODE | DRM_XE_VM_CREATE_FLAG_FAULT_MODE), xeIoctlHelper->getFlagsForVmCreate(true, true, true));
-    EXPECT_EQ(static_cast<uint32_t>(DRM_XE_VM_CREATE_FLAG_LR_MODE | DRM_XE_VM_CREATE_FLAG_FAULT_MODE), xeIoctlHelper->getFlagsForVmCreate(true, true, false));
-    EXPECT_EQ(static_cast<uint32_t>(DRM_XE_VM_CREATE_FLAG_LR_MODE | DRM_XE_VM_CREATE_FLAG_FAULT_MODE), xeIoctlHelper->getFlagsForVmCreate(false, true, true));
-    EXPECT_EQ(static_cast<uint32_t>(DRM_XE_VM_CREATE_FLAG_LR_MODE | DRM_XE_VM_CREATE_FLAG_FAULT_MODE), xeIoctlHelper->getFlagsForVmCreate(false, true, false));
-    EXPECT_EQ(static_cast<uint32_t>(DRM_XE_VM_CREATE_FLAG_LR_MODE), xeIoctlHelper->getFlagsForVmCreate(true, false, true));
-    EXPECT_EQ(static_cast<uint32_t>(DRM_XE_VM_CREATE_FLAG_LR_MODE), xeIoctlHelper->getFlagsForVmCreate(true, false, false));
-    EXPECT_EQ(static_cast<uint32_t>(DRM_XE_VM_CREATE_FLAG_LR_MODE), xeIoctlHelper->getFlagsForVmCreate(false, false, true));
-    EXPECT_EQ(static_cast<uint32_t>(DRM_XE_VM_CREATE_FLAG_LR_MODE), xeIoctlHelper->getFlagsForVmCreate(false, false, false));
+    for (auto &disableScratch : ::testing::Bool()) {
+        for (auto &enablePageFault : ::testing::Bool()) {
+            for (auto &useVmBind : ::testing::Bool()) {
+                auto flags = xeIoctlHelper->getFlagsForVmCreate(disableScratch, enablePageFault, useVmBind);
+                EXPECT_EQ(static_cast<uint32_t>(DRM_XE_VM_CREATE_FLAG_LR_MODE), (flags & DRM_XE_VM_CREATE_FLAG_LR_MODE));
+                if (enablePageFault) {
+                    EXPECT_EQ(static_cast<uint32_t>(DRM_XE_VM_CREATE_FLAG_FAULT_MODE), (flags & DRM_XE_VM_CREATE_FLAG_FAULT_MODE));
+                }
+                if (!disableScratch) {
+                    EXPECT_EQ(static_cast<uint32_t>(DRM_XE_VM_CREATE_FLAG_SCRATCH_PAGE), (flags & DRM_XE_VM_CREATE_FLAG_SCRATCH_PAGE));
+                }
+            }
+        }
+    }
 }
 
 TEST(IoctlHelperXeTest, whenGettingFlagsForVmBindThenPropertValueIsReturned) {