fix: set staging buffer writeable for aub and tbx

Signed-off-by: Szymon Morek <szymon.morek@intel.com>

Author: smorek-intel

opencl/source/context/context.cpp

@@ -308,7 +308,8 @@ bool Context::createImpl(const cl_context_properties *properties,
             this->svmAllocsManager = new SVMAllocsManager(this->memoryManager,
                                                           this->areMultiStorageAllocationsPreferred());
             this->svmAllocsManager->initUsmAllocationsCaches(device->getDevice());
-            this->stagingBufferManager = std::make_unique<StagingBufferManager>(svmAllocsManager, rootDeviceIndices, deviceBitfields);
+            auto requiresWritableStaging = device->getDefaultEngine().commandStreamReceiver->getType() != CommandStreamReceiverType::hardware;
+            this->stagingBufferManager = std::make_unique<StagingBufferManager>(svmAllocsManager, rootDeviceIndices, deviceBitfields, requiresWritableStaging);
         }
     }
 

opencl/test/unit_test/mocks/mock_context.cpp

@@ -125,7 +125,7 @@ void MockContext::initializeWithDevices(const ClDeviceVector &devices, bool noSp
         }
         deviceBitfields.insert({rootDeviceIndex, deviceBitfield});
     }
-    stagingBufferManager = std::make_unique<StagingBufferManager>(svmAllocsManager, rootDeviceIndices, deviceBitfields);
+    stagingBufferManager = std::make_unique<StagingBufferManager>(svmAllocsManager, rootDeviceIndices, deviceBitfields, true);
 
     cl_int retVal;
     if (!noSpecialQueue) {

shared/source/utilities/staging_buffer_manager.cpp

@@ -31,7 +31,8 @@ void StagingBufferTracker::freeChunk() const {
     allocator->free(chunkAddress, size);
 }
 
-StagingBufferManager::StagingBufferManager(SVMAllocsManager *svmAllocsManager, const RootDeviceIndicesContainer &rootDeviceIndices, const std::map<uint32_t, DeviceBitfield> &deviceBitfields) : svmAllocsManager(svmAllocsManager), rootDeviceIndices(rootDeviceIndices), deviceBitfields(deviceBitfields) {
+StagingBufferManager::StagingBufferManager(SVMAllocsManager *svmAllocsManager, const RootDeviceIndicesContainer &rootDeviceIndices, const std::map<uint32_t, DeviceBitfield> &deviceBitfields, bool requiresWritable)
+    : svmAllocsManager(svmAllocsManager), rootDeviceIndices(rootDeviceIndices), deviceBitfields(deviceBitfields), requiresWritable(requiresWritable) {
     if (debugManager.flags.StagingBufferSize.get() != -1) {
         chunkSize = debugManager.flags.StagingBufferSize.get() * MemoryConstants::kiloByte;
     }
@@ -240,6 +241,11 @@ std::pair<HeapAllocator *, uint64_t> StagingBufferManager::getExistingBuffer(siz
         allocator = stagingBuffer.getAllocator();
         buffer = allocator->allocate(size);
         if (buffer != 0) {
+            if (requiresWritable) {
+                auto alloc = svmAllocsManager->getSVMAlloc(stagingBuffer.getBaseAddress())->gpuAllocations.getDefaultGraphicsAllocation();
+                alloc->setTbxWritable(true, std::numeric_limits<uint32_t>::max());
+                alloc->setAubWritable(true, std::numeric_limits<uint32_t>::max());
+            }
             break;
         }
     }

shared/source/utilities/staging_buffer_manager.h

@@ -71,7 +71,7 @@ using StagingQueue = std::queue<std::pair<UserDstData, StagingBufferTracker>>;
 
 class StagingBufferManager {
   public:
-    StagingBufferManager(SVMAllocsManager *svmAllocsManager, const RootDeviceIndicesContainer &rootDeviceIndices, const std::map<uint32_t, DeviceBitfield> &deviceBitfields);
+    StagingBufferManager(SVMAllocsManager *svmAllocsManager, const RootDeviceIndicesContainer &rootDeviceIndices, const std::map<uint32_t, DeviceBitfield> &deviceBitfields, bool requiresWritable);
     ~StagingBufferManager();
     StagingBufferManager(StagingBufferManager &&other) noexcept = delete;
     StagingBufferManager(const StagingBufferManager &other) = delete;
@@ -106,6 +106,7 @@ class StagingBufferManager {
     SVMAllocsManager *svmAllocsManager;
     const RootDeviceIndicesContainer rootDeviceIndices;
     const std::map<uint32_t, DeviceBitfield> deviceBitfields;
+    const bool requiresWritable = false;
 };
 
 } // namespace NEO

shared/test/unit_test/utilities/staging_buffer_manager_tests.cpp

@@ -29,7 +29,7 @@ class StagingBufferManagerFixture : public DeviceFixture {
         debugManager.flags.EnableCopyWithStagingBuffers.set(1);
         RootDeviceIndicesContainer rootDeviceIndices = {mockRootDeviceIndex};
         std::map<uint32_t, DeviceBitfield> deviceBitfields{{mockRootDeviceIndex, mockDeviceBitfield}};
-        this->stagingBufferManager = std::make_unique<StagingBufferManager>(svmAllocsManager.get(), rootDeviceIndices, deviceBitfields);
+        this->stagingBufferManager = std::make_unique<StagingBufferManager>(svmAllocsManager.get(), rootDeviceIndices, deviceBitfields, false);
         this->csr = pDevice->commandStreamReceivers[0].get();
     }
 
@@ -211,6 +211,41 @@ TEST_F(StagingBufferManagerTest, givenStagingBufferEnabledWhenValidForImageWrite
     svmAllocsManager->freeSVMAlloc(usmBuffer);
 }
 
+TEST_F(StagingBufferManagerTest, givenStagingBufferWhenPerformCopyOnHwThenDontSetWritable) {
+    constexpr size_t numOfChunkCopies = 8;
+    constexpr size_t remainder = 1024;
+    constexpr size_t totalCopySize = stagingBufferSize * numOfChunkCopies + remainder;
+    copyThroughStagingBuffers(totalCopySize, numOfChunkCopies + 1, 1, csr);
+    auto svmData = svmAllocsManager->svmAllocs.allocations[0].second.get();
+    auto alloc = svmData->gpuAllocations.getDefaultGraphicsAllocation();
+    alloc->setAubWritable(false, std::numeric_limits<uint32_t>::max());
+    alloc->setTbxWritable(false, std::numeric_limits<uint32_t>::max());
+    copyThroughStagingBuffers(totalCopySize, numOfChunkCopies + 1, 0, csr);
+
+    EXPECT_FALSE(alloc->isAubWritable(std::numeric_limits<uint32_t>::max()));
+    EXPECT_FALSE(alloc->isTbxWritable(std::numeric_limits<uint32_t>::max()));
+}
+
+TEST_F(StagingBufferManagerTest, givenStagingBufferWhenPerformCopyOnSimulationThenSetWritable) {
+    constexpr size_t numOfChunkCopies = 8;
+    constexpr size_t remainder = 1024;
+    constexpr size_t totalCopySize = stagingBufferSize * numOfChunkCopies + remainder;
+
+    RootDeviceIndicesContainer rootDeviceIndices = {mockRootDeviceIndex};
+    std::map<uint32_t, DeviceBitfield> deviceBitfields{{mockRootDeviceIndex, mockDeviceBitfield}};
+    stagingBufferManager = std::make_unique<StagingBufferManager>(svmAllocsManager.get(), rootDeviceIndices, deviceBitfields, true);
+
+    copyThroughStagingBuffers(totalCopySize, numOfChunkCopies + 1, 1, csr);
+    auto svmData = svmAllocsManager->svmAllocs.allocations[0].second.get();
+    auto alloc = svmData->gpuAllocations.getDefaultGraphicsAllocation();
+    alloc->setAubWritable(false, std::numeric_limits<uint32_t>::max());
+    alloc->setTbxWritable(false, std::numeric_limits<uint32_t>::max());
+    copyThroughStagingBuffers(totalCopySize, numOfChunkCopies + 1, 0, csr);
+
+    EXPECT_TRUE(alloc->isAubWritable(std::numeric_limits<uint32_t>::max()));
+    EXPECT_TRUE(alloc->isTbxWritable(std::numeric_limits<uint32_t>::max()));
+}
+
 TEST_F(StagingBufferManagerTest, givenStagingBufferWhenPerformCopyThenCopyData) {
     constexpr size_t numOfChunkCopies = 8;
     constexpr size_t remainder = 1024;
@@ -336,7 +371,7 @@ TEST_F(StagingBufferManagerTest, givenStagingBufferWhenChangedBufferSizeThenPerf
 
     RootDeviceIndicesContainer rootDeviceIndices = {mockRootDeviceIndex};
     std::map<uint32_t, DeviceBitfield> deviceBitfields{{mockRootDeviceIndex, mockDeviceBitfield}};
-    stagingBufferManager = std::make_unique<StagingBufferManager>(svmAllocsManager.get(), rootDeviceIndices, deviceBitfields);
+    stagingBufferManager = std::make_unique<StagingBufferManager>(svmAllocsManager.get(), rootDeviceIndices, deviceBitfields, false);
     copyThroughStagingBuffers(totalCopySize, numOfChunkCopies + 1, 1, csr);
 }