Merge pull request #970 from mangupta/swdev-172995

hipExtMallocWithFlags implementation

[ROCm/hip commit: 30b5c02ec4]

projects/hip/src/hip_memory.cpp

@@ -264,6 +264,42 @@ 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();
+
+#if (__hcc_workweek__ >= 19115)
+    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;
+        }
+    }
+#else
+    hipError_t hip_status = hipErrorMemoryAllocation;
+#endif
+
+    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);