libhsakmt: Add support functions for ASAN

Add support functions to remap the first page of device memory (GPU/GTT) to share host ASAN logic. Signed-off-by: David Yat Sin <David.YatSin@amd.com> Change-Id: I4c27d5417ba80a172dccb0a079a597c5dc1c8f85 [ROCm/ROCR-Runtime commit: 1e6d728730]
2023-05-01 15:20:21 -07:00
@@ -913,6 +913,30 @@ hsaKmtOpenSMI(
    int *fd             // OUT: anonymous file handle
 );

+/**
+   If this is GPU Mapped memory, remap the first page at this address to be normal system memory
+
+   This is used in ASAN mode to remap the first page of device memory to share host ASAN logic.
+   This function is only supported when libhsakmt is compiled in ASAN mode.
+*/
+HSAKMT_STATUS
+HSAKMTAPI
+hsaKmtReplaceAsanHeaderPage(
+    void *addr     // IN: Start of othe virtual address page
+);
+
+/**
+   If this is GPU Mapped memory, remap the first page back to the original GPU memory
+
+   This is used in ASAN mode to remap the first page back to its original mapping.
+   This function is only supported when libhsakmt is compiled in ASAN mode.
+*/
+HSAKMT_STATUS
+HSAKMTAPI
+hsaKmtReturnAsanHeaderPage(
+    void *addr     // IN: Start of othe virtual address page
+);
+
 #ifdef __cplusplus
 }   //extern "C"
 #endif
@@ -23,9 +23,9 @@
 * DEALINGS IN THE SOFTWARE.
 */

+#include "libhsakmt.h"
 #include "fmm.h"
 #include "linux/kfd_ioctl.h"
-#include "libhsakmt.h"
 #include <stdlib.h>
 #include <stdio.h>
 #include <string.h>
@@ -106,6 +106,11 @@ struct vm_object {
 	void *user_data;
 	/* Flag to indicate imported KFD buffer */
 	bool is_imported_kfd_bo;
+#ifdef SANITIZER_AMDGPU
+	int mmap_flags;
+	int mmap_fd;
+	off_t mmap_offset;
+#endif
 };
 typedef struct vm_object vm_object_t;

@@ -340,6 +345,9 @@ static vm_object_t *vm_create_and_init_object(void *start, uint64_t size,
 		object->is_imported_kfd_bo = false;
 		object->node.key = rbtree_key((unsigned long)start, size);
 		object->user_node.key = rbtree_key(0, 0);
+#ifdef SANITIZER_AMDGPU
+		object->mmap_fd = 0;
+#endif
 	}

 	return object;
@@ -1495,6 +1503,13 @@ void *fmm_allocate_device(uint32_t gpu_id, uint32_t node_id, void *address,
 			__fmm_release(vm_obj, aperture);
 			return NULL;
 		}
+#ifdef SANITIZER_AMDGPU
+		if (vm_obj) {
+			vm_obj->mmap_flags = mflags.ui32.HostAccess ? PROT_READ | PROT_WRITE : PROT_NONE;
+			vm_obj->mmap_fd = gpu_mem[gpu_mem_id].drm_render_fd;
+			vm_obj->mmap_offset = mmap_offset;
+		}
+#endif
 	}

 	return mem;
@@ -1738,6 +1753,14 @@ static void *fmm_allocate_host_gpu(uint32_t node_id, void *address,
 				return NULL;
 			}
 		}
+
+#ifdef SANITIZER_AMDGPU
+		if (mem && vm_obj) {
+			vm_obj->mmap_flags = mflags.ui32.HostAccess ? PROT_READ | PROT_WRITE : PROT_NONE;
+			vm_obj->mmap_fd = gpu_drm_fd;
+			vm_obj->mmap_offset = mmap_offset;
+		}
+#endif
 	}

 	if (mem && vm_obj) {
@@ -3846,6 +3869,64 @@ HSAKMT_STATUS fmm_get_mem_info(const void *address, HsaPointerInfo *info)
 	return ret;
 }

+#ifdef SANITIZER_AMDGPU
+HSAKMT_STATUS fmm_replace_asan_header_page(void* address)
+{
+	HSAKMT_STATUS ret = HSAKMT_STATUS_SUCCESS;
+	manageable_aperture_t* aperture;
+	vm_object_t* vm_obj;
+
+	vm_obj = vm_find_object(address, UINT64_MAX, &aperture);
+	if (!vm_obj)
+		return HSAKMT_STATUS_ERROR;
+	/* Successful vm_find_object returns with the aperture locked */
+
+	/* If this is a GPU-mapped memory, remap the first page to be normal system memory*/
+	if (vm_obj->mmap_fd) {
+		void* p = mmap(address,
+				PAGE_SIZE,
+				PROT_WRITE | PROT_READ, MAP_ANONYMOUS | MAP_PRIVATE | MAP_FIXED,
+				-1,
+				0);
+
+		if (p == MAP_FAILED)
+			ret = HSAKMT_STATUS_ERROR;
+	}
+
+	pthread_mutex_unlock(&aperture->fmm_mutex);
+	return ret;
+}
+
+HSAKMT_STATUS fmm_return_asan_header_page(void* address)
+{
+	HSAKMT_STATUS ret = HSAKMT_STATUS_SUCCESS;
+	manageable_aperture_t* aperture;
+	vm_object_t* vm_obj;
+
+	vm_obj = vm_find_object(address, UINT64_MAX, &aperture);
+	if (!vm_obj)
+		return HSAKMT_STATUS_ERROR;
+	/* Successful vm_find_object returns with the aperture locked */
+
+	/* If this is a GPU-mapped memory, remap the first page back to the original GPU memory*/
+	if (vm_obj->mmap_fd) {
+		off_t mmap_offset = vm_obj->mmap_offset + ((char*)address - (char*)vm_obj->start);
+		void* p = mmap(address,
+				PAGE_SIZE,
+				vm_obj->mmap_flags,
+				MAP_SHARED | MAP_FIXED,
+				vm_obj->mmap_fd,
+				mmap_offset);
+
+		if (p == MAP_FAILED)
+			ret = HSAKMT_STATUS_ERROR;
+	}
+
+	pthread_mutex_unlock(&aperture->fmm_mutex);
+	return ret;
+}
+#endif
+
 HSAKMT_STATUS fmm_set_mem_user_data(const void *mem, void *usr_data)
 {
 	manageable_aperture_t *aperture;
@@ -62,6 +62,10 @@ int fmm_unmap_from_gpu(void *address);
 bool fmm_get_handle(void *address, uint64_t *handle);
 HSAKMT_STATUS fmm_get_mem_info(const void *address, HsaPointerInfo *info);
 HSAKMT_STATUS fmm_set_mem_user_data(const void *mem, void *usr_data);
+#ifdef SANITIZER_AMDGPU
+HSAKMT_STATUS fmm_replace_asan_header_page(void* address);
+HSAKMT_STATUS fmm_return_asan_header_page(void* address);
+#endif

 /* Topology interface*/
 HSAKMT_STATUS fmm_node_added(HSAuint32 gpu_id);
@@ -43,6 +43,12 @@ extern HsaVersionInfo kfd_version_info;
 #undef HSAKMTAPI
 #define HSAKMTAPI __attribute__((visibility ("default")))

+#if defined(__clang__)
+#if __has_feature(address_sanitizer)
+#define SANITIZER_AMDGPU 1
+#endif
+#endif
+
 /*Avoid pointer-to-int-cast warning*/
 #define PORT_VPTR_TO_UINT64(vptr) ((uint64_t)(unsigned long)(vptr))

@@ -588,3 +588,27 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtSetMemoryUserData(const void *Pointer,

 	return fmm_set_mem_user_data(Pointer, UserData);
 }
+
+HSAKMT_STATUS HSAKMTAPI hsaKmtReplaceAsanHeaderPage(void *addr)
+{
+#ifdef SANITIZER_AMDGPU
+	pr_debug("[%s] address %p\n", __func__, addr);
+	CHECK_KFD_OPEN();
+
+	return fmm_replace_asan_header_page(addr);
+#else
+	return HSAKMT_STATUS_NOT_SUPPORTED;
+#endif
+}
+
+HSAKMT_STATUS HSAKMTAPI hsaKmtReturnAsanHeaderPage(void *addr)
+{
+#ifdef SANITIZER_AMDGPU
+	pr_debug("[%s] address %p\n", __func__, addr);
+	CHECK_KFD_OPEN();
+
+	return fmm_return_asan_header_page(addr);
+#else
+	return HSAKMT_STATUS_NOT_SUPPORTED;
+#endif
+}