Add the reset of the memory stuff from the default overrides.

Author: aarongreig

scripts/templates/mockddi.cpp.mako

@@ -70,6 +70,31 @@ namespace driver
                 %else:
                     *phNative${func_class} = reinterpret_cast<ur_native_handle_t>(h${func_class});
                 %endif
+            ## These special cases handle memory stuff. Use verbose regex matching
+            ## to limit the possibility of unintentional stuff getting generated for
+            ## future entry points with similar names.
+            %elif re.search(r"MemBufferCreate$", fname):
+                if (pProperties && (pProperties)->pHost &&
+                    flags & UR_MEM_FLAG_USE_HOST_POINTER) {
+                    *phBuffer = mock::createDummyHandleWithData<ur_mem_handle_t>(
+                        reinterpret_cast<unsigned char *>((pProperties)->pHost),
+                        size);
+                } else {
+                    *phBuffer =
+                        mock::createDummyHandle<ur_mem_handle_t>(size);
+                }
+            %elif re.search(r"EnqueueMemBufferMap$", fname):
+                if(phEvent) {
+                    *phEvent = mock::createDummyHandle<ur_event_handle_t>();
+                }
+
+                auto parentDummyHandle =
+                    reinterpret_cast<mock::dummy_handle_t>(hBuffer);
+                *ppRetMap = (void *)(parentDummyHandle->MData);
+            %elif re.search(r"USM(Host|Device|Shared)Alloc$", fname):
+                *ppMem = mock::createDummyHandle<void *>(size);
+            %elif re.search(r"USMPitchedAllocExp$", fname):
+                *ppMem = mock::createDummyHandle<void *>(widthInBytes * height);
             %else:
                 %if fname == 'urAdapterGet' or fname == 'urDeviceGet' or fname == 'urPlatformGet':
                     <%

source/adapters/mock/ur_mockddi.cpp

@@ -1446,7 +1446,13 @@ __urdlllocal ur_result_t UR_APICALL urMemBufferCreate(
         result = replaceCallback(&params);
     } else {
 
-        *phBuffer = mock::createDummyHandle<ur_mem_handle_t>();
+        if (pProperties && (pProperties)->pHost &&
+            flags & UR_MEM_FLAG_USE_HOST_POINTER) {
+            *phBuffer = mock::createDummyHandleWithData<ur_mem_handle_t>(
+                reinterpret_cast<unsigned char *>((pProperties)->pHost), size);
+        } else {
+            *phBuffer = mock::createDummyHandle<ur_mem_handle_t>(size);
+        }
         result = UR_RESULT_SUCCESS;
     }
 
@@ -2186,6 +2192,7 @@ __urdlllocal ur_result_t UR_APICALL urUSMHostAlloc(
         result = replaceCallback(&params);
     } else {
 
+        *ppMem = mock::createDummyHandle<void *>(size);
         result = UR_RESULT_SUCCESS;
     }
 
@@ -2237,6 +2244,7 @@ __urdlllocal ur_result_t UR_APICALL urUSMDeviceAlloc(
         result = replaceCallback(&params);
     } else {
 
+        *ppMem = mock::createDummyHandle<void *>(size);
         result = UR_RESULT_SUCCESS;
     }
 
@@ -2288,6 +2296,7 @@ __urdlllocal ur_result_t UR_APICALL urUSMSharedAlloc(
         result = replaceCallback(&params);
     } else {
 
+        *ppMem = mock::createDummyHandle<void *>(size);
         result = UR_RESULT_SUCCESS;
     }
 
@@ -6281,10 +6290,13 @@ __urdlllocal ur_result_t UR_APICALL urEnqueueMemBufferMap(
         result = replaceCallback(&params);
     } else {
 
-        // optional output handle
         if (phEvent) {
             *phEvent = mock::createDummyHandle<ur_event_handle_t>();
         }
+
+        auto parentDummyHandle =
+            reinterpret_cast<mock::dummy_handle_t>(hBuffer);
+        *ppRetMap = (void *)(parentDummyHandle->MData);
         result = UR_RESULT_SUCCESS;
     }
 
@@ -7069,6 +7081,7 @@ __urdlllocal ur_result_t UR_APICALL urUSMPitchedAllocExp(
         result = replaceCallback(&params);
     } else {
 
+        *ppMem = mock::createDummyHandle<void *>(widthInBytes * height);
         result = UR_RESULT_SUCCESS;
     }