From 73ec5d54b541a8b3ea07a182a17d7d117187efcb Mon Sep 17 00:00:00 2001
From: Maneesh Gupta <maneesh.gupta@amd.com>
Date: Tue, 19 Mar 2019 11:59:22 +0530
Subject: [PATCH] hipExtMallocWithFlags implementation

Change-Id: Iee9e119796472200b2933d5e23be60813f33bc75
---
 include/hip/hcc_detail/hip_runtime_api.h | 18 ++++++++++++++
 src/hip_memory.cpp                       | 31 ++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/include/hip/hcc_detail/hip_runtime_api.h b/include/hip/hcc_detail/hip_runtime_api.h
index ea93d28bfe..cd6648c623 100644
--- a/include/hip/hcc_detail/hip_runtime_api.h
+++ b/include/hip/hcc_detail/hip_runtime_api.h
@@ -173,6 +173,8 @@ enum hipLimit_t {
     0x80000000  ///< Allocate non-coherent memory. Overrides HIP_COHERENT_HOST_ALLOC for specific
                 ///< allocation.
 
+#define hipDeviceMallocDefault 0x0
+#define hipDeviceMallocFinegrained 0x1  ///< Memory is allocated in fine grained region of device.
 
 //! Flags that can be used with hipHostRegister
 #define hipHostRegisterDefault 0x0   ///< Memory is Mapped and Portable
@@ -1051,6 +1053,22 @@ hipError_t hipPointerGetAttributes(hipPointerAttribute_t* attributes, const void
  */
 hipError_t hipMalloc(void** ptr, size_t size);
 
+/**
+ *  @brief Allocate memory on the default accelerator
+ *
+ *  @param[out] ptr Pointer to the allocated memory
+ *  @param[in]  size Requested memory size
+ *  @param[in]  flags Type of memory allocation
+ *
+ *  If size is 0, no memory is allocated, *ptr returns nullptr, and hipSuccess is returned.
+ *
+ *  @return #hipSuccess, #hipErrorMemoryAllocation, #hipErrorInvalidValue (bad context, null *ptr)
+ *
+ *  @see hipMallocPitch, hipFree, hipMallocArray, hipFreeArray, hipMalloc3D, hipMalloc3DArray,
+ * hipHostFree, hipHostMalloc
+ */
+hipError_t hipExtMallocWithFlags(void** ptr, size_t sizeBytes, unsigned int flags);
+
 /**
  *  @brief Allocate pinned host memory [Deprecated]
  *
diff --git a/src/hip_memory.cpp b/src/hip_memory.cpp
index 9eebcfb28c..1875b2219c 100644
--- a/src/hip_memory.cpp
+++ b/src/hip_memory.cpp
@@ -264,6 +264,37 @@ hipError_t hipMalloc(void** ptr, size_t sizeBytes) {
     return ihipLogStatus(hip_status);
 }
 
+hipError_t hipExtMallocWithFlags(void** ptr, size_t sizeBytes, unsigned int flags) {
+    HIP_INIT_SPECIAL_API(hipExtMallocWithFlags, (TRACE_MEM), ptr, sizeBytes, flags);
+    HIP_SET_DEVICE();
+    hipError_t hip_status = hipSuccess;
+
+    auto ctx = ihipGetTlsDefaultCtx();
+    // return NULL pointer when malloc size is 0
+    if (sizeBytes == 0) {
+        *ptr = NULL;
+        hip_status = hipSuccess;
+    } else if ((ctx == nullptr) || (ptr == nullptr)) {
+        hip_status = hipErrorInvalidValue;
+    } else {
+        unsigned amFlags = 0;
+        if (flags & hipDeviceMallocFinegrained) {
+            amFlags = amDeviceFinegrained;
+        } else if (flags != hipDeviceMallocDefault) {
+            hip_status = hipErrorInvalidValue;
+            return ihipLogStatus(hip_status);
+        }
+        auto device = ctx->getWriteableDevice();
+        *ptr = hip_internal::allocAndSharePtr("device_mem", sizeBytes, ctx, false /*shareWithAll*/,
+                                              amFlags /*amFlags*/, 0 /*hipFlags*/, 0);
+
+        if (sizeBytes && (*ptr == NULL)) {
+            hip_status = hipErrorMemoryAllocation;
+        }
+    }
+
+    return ihipLogStatus(hip_status);
+}
 
 hipError_t hipHostMalloc(void** ptr, size_t sizeBytes, unsigned int flags) {
     HIP_INIT_SPECIAL_API(hipHostMalloc, (TRACE_MEM), ptr, sizeBytes, flags);