[msd-intel-gen] Cache context buffer mapping

This is in the hot path for command buffer submission.

Test:
pixelbook:go/magma-tps#L1

Change-Id: I418eb62d79316cf2d561f5a4fdeb8be59672c5f0
diff --git a/drivers/gpu/msd-intel-gen/src/engine_command_streamer.cc b/drivers/gpu/msd-intel-gen/src/engine_command_streamer.cc
index bfd0b3c..a02324d 100644
--- a/drivers/gpu/msd-intel-gen/src/engine_command_streamer.cc
+++ b/drivers/gpu/msd-intel-gen/src/engine_command_streamer.cc
@@ -385,22 +385,17 @@
     if (!context->GetRingbufferGpuAddress(id(), &gpu_addr))
         return DRETF(false, "failed to get ringbuffer gpu address");
 
-    void* cpu_addr;
-    if (!context->get_context_buffer(id())->platform_buffer()->MapCpu(&cpu_addr))
-        return DRETF(false, "failed to map context page 1");
+    uint8_t* cpu_addr = reinterpret_cast<uint8_t*>(context->GetCachedContextBufferCpuAddr(id()));
+    if (!cpu_addr)
+        return DRETF(false, "failed to get cached context buffer cpu address");
 
-    RegisterStateHelper helper(
-        id(), mmio_base_,
-        reinterpret_cast<uint32_t*>(reinterpret_cast<uint8_t*>(cpu_addr) + PAGE_SIZE));
+    RegisterStateHelper helper(id(), mmio_base_, reinterpret_cast<uint32_t*>(cpu_addr + PAGE_SIZE));
 
     DLOG("UpdateContext ringbuffer gpu_addr 0x%lx tail 0x%x", gpu_addr, tail);
 
     helper.write_ring_tail_pointer(tail);
     helper.write_ring_buffer_start(gpu_addr);
 
-    if (!context->get_context_buffer(id())->platform_buffer()->UnmapCpu())
-        DLOG("UnmapPageCpu failed");
-
     return true;
 }
 
diff --git a/drivers/gpu/msd-intel-gen/src/msd_intel_context.h b/drivers/gpu/msd-intel-gen/src/msd_intel_context.h
index a2a1322..6f93a37 100644
--- a/drivers/gpu/msd-intel-gen/src/msd_intel_context.h
+++ b/drivers/gpu/msd-intel-gen/src/msd_intel_context.h
@@ -55,6 +55,21 @@
         return iter == state_map_.end() ? nullptr : iter->second.context_buffer.get();
     }
 
+    void* GetCachedContextBufferCpuAddr(EngineCommandStreamerId id)
+    {
+        auto iter = state_map_.find(id);
+        if (iter == state_map_.end())
+            return nullptr;
+        if (!iter->second.context_buffer_cpu_addr) {
+            MsdIntelBuffer* context_buffer = iter->second.context_buffer.get();
+            if (!context_buffer)
+                return nullptr;
+            if (!context_buffer->platform_buffer()->MapCpu(&iter->second.context_buffer_cpu_addr))
+                return DRETP(nullptr, "Failed to map context buffer");
+        }
+        return iter->second.context_buffer_cpu_addr;
+    }
+
     Ringbuffer* get_ringbuffer(EngineCommandStreamerId id)
     {
         auto iter = state_map_.find(id);
@@ -75,6 +90,7 @@
         std::shared_ptr<MsdIntelBuffer> context_buffer;
         std::unique_ptr<GpuMapping> context_mapping;
         std::unique_ptr<Ringbuffer> ringbuffer;
+        void* context_buffer_cpu_addr = nullptr;
     };
 
     std::map<EngineCommandStreamerId, PerEngineState> state_map_;
diff --git a/drivers/gpu/msd-intel-gen/tests/unit_tests/test_context.cc b/drivers/gpu/msd-intel-gen/tests/unit_tests/test_context.cc
index edef586..25ad04a 100644
--- a/drivers/gpu/msd-intel-gen/tests/unit_tests/test_context.cc
+++ b/drivers/gpu/msd-intel-gen/tests/unit_tests/test_context.cc
@@ -51,6 +51,7 @@
 
     void Map(bool global)
     {
+        // Arbitrary
         constexpr uint32_t base = 0x10000;
 
         std::weak_ptr<MsdIntelConnection> connection;
@@ -93,6 +94,35 @@
         EXPECT_FALSE(context->Unmap(RENDER_COMMAND_STREAMER));
     }
 
+    void CachedMapping()
+    {
+        // Arbitrary
+        constexpr uint32_t base = 0x10000;
+
+        std::unique_ptr<MsdIntelBuffer> buffer(MsdIntelBuffer::Create(PAGE_SIZE, "test"));
+
+        auto ringbuffer = std::make_unique<Ringbuffer>(
+            std::unique_ptr<MsdIntelBuffer>(MsdIntelBuffer::Create(PAGE_SIZE, "test")));
+
+        auto address_space_owner = std::make_unique<AddressSpaceOwner>();
+        auto address_space = std::make_shared<MockAddressSpace>(address_space_owner.get(), base,
+                                                                buffer->platform_buffer()->size() +
+                                                                    ringbuffer->size());
+
+        auto context =
+            std::make_unique<ClientContext>(std::weak_ptr<MsdIntelConnection>(), address_space);
+
+        void* cpu_addr = context->GetCachedContextBufferCpuAddr(RENDER_COMMAND_STREAMER);
+        EXPECT_EQ(nullptr, cpu_addr);
+
+        context->SetEngineState(RENDER_COMMAND_STREAMER, std::move(buffer), std::move(ringbuffer));
+
+        cpu_addr = context->GetCachedContextBufferCpuAddr(RENDER_COMMAND_STREAMER);
+        EXPECT_NE(nullptr, cpu_addr);
+        // Returned address shouldn't change.
+        EXPECT_EQ(cpu_addr, context->GetCachedContextBufferCpuAddr(RENDER_COMMAND_STREAMER));
+    }
+
     static void SubmitCommandBuffer(uint32_t command_buffer_count, uint32_t semaphore_count)
     {
         DLOG("SubmitCommandBuffer command_buffer_count %u semaphore_count %u", command_buffer_count,
@@ -226,6 +256,8 @@
     test.Map(false);
 }
 
+TEST(MsdIntelContext, CachedMapping) { TestContext().CachedMapping(); }
+
 TEST(GlobalContext, GlobalMap)
 {
     TestContext test;