fix: aub residency in HW+AUB mode

Related-To: NEO-9230
Signed-off-by: Fabian Zwolinski <[email protected]>
Source: 06ec4c3

Author: fzwolinski

opencl/test/unit_test/mt_tests/os_interface/linux/drm_memory_manager_mt_tests.cpp

@@ -40,7 +40,7 @@ TEST(DrmMemoryManagerTest, givenDrmMemoryManagerWhenSharedAllocationIsCreatedFro
     auto mock = new MockDrm(0, *executionEnvironment.rootDeviceEnvironments[0]);
     mock->setupIoctlHelper(defaultHwInfo->platform.eProductFamily);
     executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
 
     auto memoryManager = std::make_unique<TestedDrmMemoryManager>(executionEnvironment);
 
@@ -112,7 +112,7 @@ TEST(DrmMemoryManagerTest, givenMultipleThreadsWhenSharedAllocationIsCreatedThen
     auto mock = new MockDrm(0, *executionEnvironment.rootDeviceEnvironments[0]);
     mock->setupIoctlHelper(defaultHwInfo->platform.eProductFamily);
     executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
 
     auto memoryManager = std::make_unique<TestedDrmMemoryManager>(executionEnvironment);
 

opencl/test/unit_test/os_interface/linux/linux_create_command_queue_with_properties_tests.cpp

@@ -29,7 +29,7 @@ struct ClCreateCommandQueueWithPropertiesLinux : public UltCommandStreamReceiver
         auto osInterface = new OSInterface();
         osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface.reset(osInterface);
-        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, rootDeviceIndex);
+        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, rootDeviceIndex, false);
         executionEnvironment->memoryManager.reset(new TestedDrmMemoryManager(*executionEnvironment));
         mdevice = std::make_unique<MockClDevice>(MockDevice::create<MockDevice>(executionEnvironment, rootDeviceIndex));
 

shared/source/command_stream/command_stream_receiver.h

@@ -16,6 +16,8 @@
 #include "shared/source/helpers/options.h"
 #include "shared/source/utilities/spinlock.h"
 
+#include "aubstream/allocation_params.h"
+
 #include <atomic>
 #include <cstddef>
 #include <cstdint>
@@ -245,6 +247,7 @@ class CommandStreamReceiver {
     virtual bool expectMemory(const void *gfxAddress, const void *srcAddress, size_t length, uint32_t compareOperation);
     MOCKABLE_VIRTUAL bool writeMemory(GraphicsAllocation &gfxAllocation) { return writeMemory(gfxAllocation, false, 0, 0); }
     virtual bool writeMemory(GraphicsAllocation &gfxAllocation, bool isChunkCopy, uint64_t gpuVaChunkOffset, size_t chunkSize) { return false; }
+    virtual void writeMemoryAub(aub_stream::AllocationParams &allocationParams){};
 
     virtual bool isMultiOsContextCapable() const = 0;
 

shared/source/command_stream/command_stream_receiver_simulated_common_hw.h

@@ -7,6 +7,7 @@
 
 #pragma once
 #include "shared/source/command_stream/command_stream_receiver_hw.h"
+#include "shared/source/helpers/hardware_context_controller.h"
 #include "shared/source/memory_manager/memory_banks.h"
 
 #include "aub_mapper_common.h"
@@ -20,7 +21,6 @@ struct AubStream;
 namespace NEO {
 class AddressMapper;
 class GraphicsAllocation;
-class HardwareContextController;
 template <typename GfxFamily>
 class CommandStreamReceiverSimulatedCommonHw : public CommandStreamReceiverHw<GfxFamily> {
   protected:
@@ -59,6 +59,9 @@ class CommandStreamReceiverSimulatedCommonHw : public CommandStreamReceiverHw<Gf
     virtual void writeMemory(uint64_t gpuAddress, void *cpuAddress, size_t size, uint32_t memoryBank, uint64_t entryBits) = 0;
     virtual void writeMemoryWithAubManager(GraphicsAllocation &graphicsAllocation, bool isChunkCopy, uint64_t gpuVaChunkOffset, size_t chunkSize) = 0;
     virtual void writeMMIO(uint32_t offset, uint32_t value) = 0;
+    void writeMemoryAub(aub_stream::AllocationParams &allocationParams) override {
+        hardwareContextController->writeMemory(allocationParams);
+    }
 
     virtual void setAubWritable(bool writable, GraphicsAllocation &graphicsAllocation) = 0;
     virtual bool isAubWritable(GraphicsAllocation &graphicsAllocation) const = 0;

shared/source/os_interface/aub_memory_operations_handler.cpp

@@ -7,9 +7,13 @@
 
 #include "shared/source/os_interface/aub_memory_operations_handler.h"
 
+#include "shared/source/aub/aub_helper.h"
 #include "shared/source/aub_mem_dump/aub_mem_dump.h"
+#include "shared/source/command_stream/command_stream_receiver.h"
+#include "shared/source/device/device.h"
 #include "shared/source/gmm_helper/cache_settings_helper.h"
 #include "shared/source/gmm_helper/gmm.h"
+#include "shared/source/gmm_helper/gmm_helper.h"
 #include "shared/source/memory_manager/graphics_allocation.h"
 
 #include "aubstream/allocation_params.h"
@@ -29,7 +33,12 @@ MemoryOperationsStatus AubMemoryOperationsHandler::makeResident(Device *device,
     auto lock = acquireLock(resourcesLock);
     int hint = AubMemDump::DataTypeHintValues::TraceNotype;
     for (const auto &allocation : gfxAllocations) {
-        aub_stream::AllocationParams params(allocation->getGpuAddress(),
+        if (!isAubWritable(*allocation, device)) {
+            continue;
+        }
+
+        uint64_t gpuAddress = device ? device->getGmmHelper()->decanonize(allocation->getGpuAddress()) : allocation->getGpuAddress();
+        aub_stream::AllocationParams params(gpuAddress,
                                             allocation->getUnderlyingBuffer(),
                                             allocation->getUnderlyingBufferSize(),
                                             allocation->storageInfo.getMemoryBanks(),
@@ -43,10 +52,19 @@ MemoryOperationsStatus AubMemoryOperationsHandler::makeResident(Device *device,
             params.additionalParams.uncached = CacheSettingsHelper::isUncachedType(gmm->resourceParams.Usage);
         }
 
-        aubManager->writeMemory2(params);
+        if (allocation->storageInfo.cloningOfPageTables || !allocation->isAllocatedInLocalMemoryPool()) {
+            aubManager->writeMemory2(params);
+        } else {
+            device->getDefaultEngine().commandStreamReceiver->writeMemoryAub(params);
+        }
+
         if (!allocation->getAubInfo().writeMemoryOnly) {
             residentAllocations.push_back(allocation);
         }
+
+        if (AubHelper::isOneTimeAubWritableAllocationType(allocation->getAllocationType())) {
+            setAubWritable(false, *allocation, device);
+        }
     }
     return MemoryOperationsStatus::SUCCESS;
 }
@@ -79,8 +97,44 @@ MemoryOperationsStatus AubMemoryOperationsHandler::isResident(Device *device, Gr
         return MemoryOperationsStatus::SUCCESS;
     }
 }
+
 void AubMemoryOperationsHandler::setAubManager(aub_stream::AubManager *aubManager) {
     this->aubManager = aubManager;
 }
 
+bool AubMemoryOperationsHandler::isAubWritable(GraphicsAllocation &graphicsAllocation, Device *device) const {
+    if (!device) {
+        return false;
+    }
+    auto bank = static_cast<uint32_t>(getMemoryBanksBitfield(&graphicsAllocation, device).to_ulong());
+    if (bank == 0u || graphicsAllocation.storageInfo.cloningOfPageTables) {
+        bank = GraphicsAllocation::defaultBank;
+    }
+    return graphicsAllocation.isAubWritable(bank);
+}
+
+void AubMemoryOperationsHandler::setAubWritable(bool writable, GraphicsAllocation &graphicsAllocation, Device *device) {
+    if (!device) {
+        return;
+    }
+    auto bank = static_cast<uint32_t>(getMemoryBanksBitfield(&graphicsAllocation, device).to_ulong());
+    if (bank == 0u || graphicsAllocation.storageInfo.cloningOfPageTables) {
+        bank = GraphicsAllocation::defaultBank;
+    }
+    graphicsAllocation.setAubWritable(writable, bank);
+}
+
+DeviceBitfield AubMemoryOperationsHandler::getMemoryBanksBitfield(GraphicsAllocation *allocation, Device *device) const {
+    if (allocation->getMemoryPool() == MemoryPool::LocalMemory) {
+        if (allocation->storageInfo.memoryBanks.any()) {
+            if (allocation->storageInfo.cloningOfPageTables ||
+                device->getDefaultEngine().commandStreamReceiver->isMultiOsContextCapable()) {
+                return allocation->storageInfo.memoryBanks;
+            }
+        }
+        return device->getDeviceBitfield();
+    }
+    return {};
+}
+
 } // namespace NEO

shared/source/os_interface/aub_memory_operations_handler.h

@@ -1,11 +1,12 @@
 /*
- * Copyright (C) 2019-2022 Intel Corporation
+ * Copyright (C) 2019-2023 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
  */
 
 #pragma once
+#include "shared/source/helpers/device_bitfield.h"
 #include "shared/source/memory_manager/memory_operations_handler.h"
 #include "shared/source/utilities/spinlock.h"
 
@@ -30,7 +31,11 @@ class AubMemoryOperationsHandler : public MemoryOperationsHandler {
 
     void setAubManager(aub_stream::AubManager *aubManager);
 
+    bool isAubWritable(GraphicsAllocation &graphicsAllocation, Device *device) const;
+    void setAubWritable(bool writable, GraphicsAllocation &graphicsAllocation, Device *device);
+
   protected:
+    DeviceBitfield getMemoryBanksBitfield(GraphicsAllocation *allocation, Device *device) const;
     [[nodiscard]] MOCKABLE_VIRTUAL std::unique_lock<SpinLock> acquireLock(SpinLock &lock) {
         return std::unique_lock<SpinLock>{lock};
     }

shared/source/os_interface/linux/CMakeLists.txt

@@ -38,6 +38,7 @@ set(NEO_CORE_OS_INTERFACE_LINUX
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_operations_handler_create.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_operations_handler_default.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_operations_handler_default.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_operations_handler_with_aub_dump.h
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_manager_create_multi_host_allocation.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_version.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_wrappers_checks.cpp

shared/source/os_interface/linux/drm_memory_operations_handler.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2022 Intel Corporation
+ * Copyright (C) 2019-2023 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -25,7 +25,7 @@ class DrmMemoryOperationsHandler : public MemoryOperationsHandler {
 
     virtual MemoryOperationsStatus evictUnusedAllocations(bool waitForCompletion, bool isLockNeeded) = 0;
 
-    static std::unique_ptr<DrmMemoryOperationsHandler> create(Drm &drm, uint32_t rootDeviceIndex);
+    static std::unique_ptr<DrmMemoryOperationsHandler> create(Drm &drm, uint32_t rootDeviceIndex, bool withAubDump);
 
     uint32_t getRootDeviceIndex() const {
         return this->rootDeviceIndex;

shared/source/os_interface/linux/drm_memory_operations_handler_create.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2022 Intel Corporation
+ * Copyright (C) 2020-2023 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -8,18 +8,28 @@
 #include "shared/source/execution_environment/root_device_environment.h"
 #include "shared/source/os_interface/linux/drm_memory_operations_handler_bind.h"
 #include "shared/source/os_interface/linux/drm_memory_operations_handler_default.h"
+#include "shared/source/os_interface/linux/drm_memory_operations_handler_with_aub_dump.h"
 #include "shared/source/os_interface/linux/drm_neo.h"
 
 namespace NEO {
 
-std::unique_ptr<DrmMemoryOperationsHandler> DrmMemoryOperationsHandler::create(Drm &drm, uint32_t rootDeviceIndex) {
+std::unique_ptr<DrmMemoryOperationsHandler> DrmMemoryOperationsHandler::create(Drm &drm, uint32_t rootDeviceIndex, bool withAubDump) {
     bool useVmBind = drm.isVmBindAvailable();
+    auto rootDeviceEnv = drm.getRootDeviceEnvironment().executionEnvironment.rootDeviceEnvironments[rootDeviceIndex].get();
 
     if (useVmBind) {
-        return std::make_unique<DrmMemoryOperationsHandlerBind>(drm.getRootDeviceEnvironment(), rootDeviceIndex);
+        if (withAubDump) {
+            return std::make_unique<DrmMemoryOperationsHandlerWithAubDump<DrmMemoryOperationsHandlerBind>>(*rootDeviceEnv, rootDeviceIndex);
+        } else {
+            return std::make_unique<DrmMemoryOperationsHandlerBind>(drm.getRootDeviceEnvironment(), rootDeviceIndex);
+        }
     }
 
-    return std::make_unique<DrmMemoryOperationsHandlerDefault>(rootDeviceIndex);
+    if (withAubDump) {
+        return std::make_unique<DrmMemoryOperationsHandlerWithAubDump<DrmMemoryOperationsHandlerDefault>>(*rootDeviceEnv, rootDeviceIndex);
+    } else {
+        return std::make_unique<DrmMemoryOperationsHandlerDefault>(rootDeviceIndex);
+    }
 }
 
 } // namespace NEO

shared/source/os_interface/linux/drm_memory_operations_handler_default.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2022 Intel Corporation
+ * Copyright (C) 2020-2023 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -12,9 +12,11 @@
 
 namespace NEO {
 class OsContextLinux;
+struct RootDeviceEnvironment;
 class DrmMemoryOperationsHandlerDefault : public DrmMemoryOperationsHandler {
   public:
     DrmMemoryOperationsHandlerDefault(uint32_t rootDeviceIndex);
+    DrmMemoryOperationsHandlerDefault(const RootDeviceEnvironment &rootDeviceEnvironment, uint32_t rootDeviceIndex) : DrmMemoryOperationsHandlerDefault(rootDeviceIndex) {}
     ~DrmMemoryOperationsHandlerDefault() override;
 
     MemoryOperationsStatus makeResidentWithinOsContext(OsContext *osContext, ArrayRef<GraphicsAllocation *> gfxAllocations, bool evictable) override;