From d2d95a8948b1d6b1ac132bd5ae56aa10d5cc1cec Mon Sep 17 00:00:00 2001
From: David Yat Sin <David.YatSin@amd.com>
Date: Fri, 16 Feb 2024 22:27:12 +0000
Subject: [PATCH] rocrtst: add test for contiguous mem allocations

This test may fail when run on non-upstream versions of KFD as this
feature will not be upstreamed.

Change-Id: I7131e1f50984739c0df12e4c9afe790bd7e4cdfa
---
 .../suites/functional/memory_allocation.cc    | 77 +++++++++++++++++++
 rocrtst/suites/functional/memory_allocation.h |  2 +
 rocrtst/suites/test_common/main.cc            |  7 ++
 3 files changed, 86 insertions(+)
 mode change 100755 => 100644 rocrtst/suites/functional/memory_allocation.cc
 mode change 100755 => 100644 rocrtst/suites/functional/memory_allocation.h

diff --git a/rocrtst/suites/functional/memory_allocation.cc b/rocrtst/suites/functional/memory_allocation.cc
old mode 100755
new mode 100644
index 7a47ec6925..a842a78d2c
--- a/rocrtst/suites/functional/memory_allocation.cc
+++ b/rocrtst/suites/functional/memory_allocation.cc
@@ -58,6 +58,7 @@
 #include "common/hsatimer.h"
 #include "gtest/gtest.h"
 #include "hsa/hsa.h"
+#include "hsa/hsa_ext_amd.h"
 
 static const uint32_t kNumBufferElements = 256;
 static const int kValue = 5;
@@ -486,3 +487,79 @@ void MemoryAllocationTest::MemoryBasicAllocationAndFree(void) {
   }
 }
 
+void MemoryAllocationTest::MemoryAllocateContiguousTest(hsa_agent_t agent,
+                                                        hsa_amd_memory_pool_t pool) {
+  rocrtst::pool_info_t pool_i;
+  hsa_device_type_t ag_type;
+  ASSERT_SUCCESS(rocrtst::AcquirePoolInfo(pool, &pool_i));
+
+  if (verbosity() > 0) PrintAgentNameAndType(agent);
+
+  ASSERT_SUCCESS(hsa_agent_get_info(agent, HSA_AGENT_INFO_DEVICE, &ag_type));
+
+  // if allocation is allowed in this pool allocate the memory
+  // and then free it
+  if (ag_type != HSA_DEVICE_TYPE_GPU || !pool_i.alloc_allowed || !pool_i.alloc_granule ||
+      !pool_i.alloc_alignment) {
+    return;
+  }
+
+  if (verbosity() > 0) PrintSegmentNameAndType(pool_i.segment);
+
+  // find all gpu agents
+  std::vector<hsa_agent_t> gpus;
+  ASSERT_SUCCESS(hsa_iterate_agents(rocrtst::IterateGPUAgents, &gpus));
+
+  const size_t alloc_size = pool_i.alloc_granule * 1024;
+
+  char* memoryPtr;
+
+  ASSERT_SUCCESS(hsa_amd_memory_pool_allocate(pool, alloc_size, HSA_AMD_MEMORY_POOL_CONTIGUOUS_FLAG,
+                                              reinterpret_cast<void**>(&memoryPtr)));
+  if (!memoryPtr) return;
+
+  int dmabuf = -1;
+  uint64_t offset;
+  ASSERT_SUCCESS(hsa_amd_portable_export_dmabuf(memoryPtr, alloc_size, &dmabuf, &offset));
+
+  void* importedPtr;
+  size_t importedSz;
+
+  ASSERT_SUCCESS(hsa_amd_interop_map_buffer(gpus.size(), gpus.data(), dmabuf, 0, &importedSz,
+                                                   &importedPtr, 0, NULL));
+
+  ASSERT_NE((uint64_t)importedPtr, 0);
+  ASSERT_EQ(importedSz, alloc_size);
+
+  close(dmabuf);
+
+  ASSERT_SUCCESS(hsa_amd_interop_unmap_buffer(importedPtr));
+
+  ASSERT_SUCCESS(hsa_amd_memory_pool_free(memoryPtr));
+  return;
+}
+
+void MemoryAllocationTest::MemoryAllocateContiguousTest(void) {
+  hsa_status_t err;
+  std::vector<std::shared_ptr<rocrtst::agent_pools_t>> agent_pools;
+  if (verbosity() > 0) {
+    PrintMemorySubtestHeader("MemoryAllocateContiguousTest");
+  }
+
+  ASSERT_SUCCESS(rocrtst::GetAgentPools(&agent_pools));
+
+  auto pool_idx = 0;
+  for (auto a : agent_pools) {
+    for (auto p : a->pools) {
+      if (verbosity() > 0) {
+        std::cout << "  Pool " << pool_idx++ << ":" << std::endl;
+      }
+      MemoryAllocateContiguousTest(a->agent, p);
+    }
+  }
+
+  if (verbosity() > 0) {
+    std::cout << "subtest Passed" << std::endl;
+    std::cout << kSubTestSeparator << std::endl;
+  }
+}
diff --git a/rocrtst/suites/functional/memory_allocation.h b/rocrtst/suites/functional/memory_allocation.h
old mode 100755
new mode 100644
index 3ca04114d7..db43d81787
--- a/rocrtst/suites/functional/memory_allocation.h
+++ b/rocrtst/suites/functional/memory_allocation.h
@@ -76,12 +76,14 @@ class MemoryAllocationTest : public TestBase {
   void GroupMemoryDynamicAllocation(void);
 
   void MemoryBasicAllocationAndFree(void);
+  void MemoryAllocateContiguousTest(void);
 
 
  private:
   void GroupMemoryDynamicAllocation(hsa_agent_t cpuAgent, hsa_agent_t gpuAgent);
   void MemoryBasicAllocationAndFree(hsa_agent_t agent,
                                                hsa_amd_memory_pool_t pool);
+  void MemoryAllocateContiguousTest(hsa_agent_t agent, hsa_amd_memory_pool_t pool);
 
   void WriteAQLPktToQueue(hsa_queue_t* q);
 };
diff --git a/rocrtst/suites/test_common/main.cc b/rocrtst/suites/test_common/main.cc
index ed897acac5..c9406a69f6 100644
--- a/rocrtst/suites/test_common/main.cc
+++ b/rocrtst/suites/test_common/main.cc
@@ -152,6 +152,13 @@ TEST(rocrtstFunc, MemoryAllocateAndFreeTest) {
   RunCustomTestEpilog(&ma);
 }
 
+TEST(rocrtstFunc, MemoryAllocateContiguousTest) {
+  MemoryAllocationTest ma(false, true);
+  RunCustomTestProlog(&ma);
+  ma.MemoryAllocateContiguousTest();
+  RunCustomTestEpilog(&ma);
+}
+
 TEST(rocrtstFunc, Concurrent_Init_Test) {
   ConcurrentInitTest ci;
   RunCustomTestProlog(&ci);