From 340e80bb1afda5399103724c90b621675dc878f3 Mon Sep 17 00:00:00 2001
From: "Klein, Yaniv" <yaniv.klein@intel.com>
Date: Fri, 1 Nov 2019 11:00:26 -0700
Subject: [PATCH] added support for 2 stages write memory of aub_manager

1. write ppgtt pages only by calling write memory with pagesOnly set to true
2. write the actual memory pages by calling writePhysicalMemoryPages

Change-Id: Id3766704e4b8f84f77336f78da6b9270bfd44b70
---
 third_party/aub_stream/headers/aub_manager.h |  8 +++++++-
 third_party/aub_stream/headers/page_info.h   | 16 ++++++++++++++++
 unit_tests/mocks/mock_aub_manager.h          | 12 ++++++++++++
 3 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 third_party/aub_stream/headers/page_info.h

diff --git a/third_party/aub_stream/headers/aub_manager.h b/third_party/aub_stream/headers/aub_manager.h
index 68f5562f58..4c7c7d04a1 100644
--- a/third_party/aub_stream/headers/aub_manager.h
+++ b/third_party/aub_stream/headers/aub_manager.h
@@ -8,6 +8,8 @@
 #pragma once
 #include <cstdint>
 #include <string>
+#include <vector>
+#include "page_info.h"
 
 namespace aub_stream {
 
@@ -26,7 +28,11 @@ class AubManager {
     virtual void pause(bool onoff) = 0;
 
     virtual void addComment(const char *message) = 0;
-    virtual void writeMemory(uint64_t gfxAddress, const void *memory, size_t size, uint32_t memoryBanks, int hint, size_t pageSize) = 0;
+    virtual void writeMemory(uint64_t gfxAddress, const void *memory, size_t size, uint32_t memoryBanks,int hint, size_t pageSize) = 0;
+    virtual void writePageTableEntries(uint64_t gfxAddress, const void *memory, size_t size, uint32_t memoryBanks, int hint,
+                                       std::vector<PageInfo> &lastLevelPages, size_t pageSize) = 0;
+
+    virtual void writePhysicalMemoryPages(const void *memory, std::vector<PageInfo> &pages, size_t size, int hint) = 0;
     virtual void freeMemory(uint64_t gfxAddress, size_t size) = 0;
 
     static AubManager *create(uint32_t productFamily, uint32_t devicesCount, uint64_t memoryBankSize, bool localMemorySupported, uint32_t streamMode);
diff --git a/third_party/aub_stream/headers/page_info.h b/third_party/aub_stream/headers/page_info.h
new file mode 100644
index 0000000000..f307d64fc8
--- /dev/null
+++ b/third_party/aub_stream/headers/page_info.h
@@ -0,0 +1,16 @@
+/*
+ * Copyright (C) 2018-2019 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+#pragma once
+namespace aub_stream {
+
+struct PageInfo {
+    uint64_t physicalAddress;
+    size_t size;
+    bool isLocalMemory;
+    uint32_t memoryBank;
+};
+} // namespace aub_stream
\ No newline at end of file
diff --git a/unit_tests/mocks/mock_aub_manager.h b/unit_tests/mocks/mock_aub_manager.h
index 210f4aa76a..98756ee6aa 100644
--- a/unit_tests/mocks/mock_aub_manager.h
+++ b/unit_tests/mocks/mock_aub_manager.h
@@ -51,6 +51,7 @@ struct MockHardwareContext : public aub_stream::HardwareContext {
 
 class MockAubManager : public aub_stream::AubManager {
     using HardwareContext = aub_stream::HardwareContext;
+    using PageInfo = aub_stream::PageInfo;
 
   public:
     MockAubManager(){};
@@ -100,6 +101,15 @@ class MockAubManager : public aub_stream::AubManager {
         writeMemoryPageSizePassed = pageSize;
     }
 
+    void writePageTableEntries(uint64_t gfxAddress, const void *memory, size_t size, uint32_t memoryBanks, int hint,
+                               std::vector<PageInfo> &lastLevelPages, size_t pageSize) override {
+        writePageTableEntriesCalled = true;
+    }
+
+    void writePhysicalMemoryPages(const void *memory, std::vector<PageInfo> &pages, size_t size, int hint) override {
+        writePhysicalMemoryPagesCalled = true;
+    }
+
     void freeMemory(uint64_t gfxAddress, size_t size) override {
         freeMemoryCalled = true;
     }
@@ -113,6 +123,8 @@ class MockAubManager : public aub_stream::AubManager {
     bool addCommentCalled = false;
     std::string receivedComment = "";
     bool writeMemoryCalled = false;
+    bool writePageTableEntriesCalled = false;
+    bool writePhysicalMemoryPagesCalled = false;
     bool freeMemoryCalled = false;
     uint32_t contextFlags = 0;
     int hintToWriteMemory = 0;