hmem/fabtests: Use neuron nrt_get_dmabuf_fd_v2 if available

nrt_get_dmabuf_fd() does not return an offset value through its api. This
means that if a non-aligned VA is given to it, the call will fail because
the dmabuf fd has to point to an aligned address region and there is no way
to return to the caller the offset that their non-aligned VA will have in
the region.

nrt_get_dmabuf_fd_v2() fixes this by adding an offset.

The new code will look to see if it can link to nrt_get_dmabuf_fd_v2() and
if it can't, it will fall back to nrt_get_dmabuf_fd(), and set the offset
to 0, as that is the only case in which it will work.

Signed-off-by: Charles Stoll <stollcha@amazon.com>

Author: charlesstoll

fabtests/common/hmem_neuron.c

@@ -53,13 +53,41 @@ struct neuron_ops {
 	void *(*nrt_tensor_get_va)(const nrt_tensor_t *tensor);
 	NRT_STATUS (*nrt_tensor_read)(const nrt_tensor_t *tensor, void *buf, size_t offset, size_t size);
 	NRT_STATUS (*nrt_tensor_write)(nrt_tensor_t *tensor, const void *buf, size_t offset, size_t size);
-	NRT_STATUS (*nrt_get_dmabuf_fd)(uint64_t va, uint64_t size, int* fd);
+	NRT_STATUS (*nrt_get_dmabuf_fd)(uint64_t va, uint64_t size, int* fd, uint64_t *offset);
 	NRT_STATUS (*nrt_init)(nrt_framework_type_t framework, const char *fw_version, const char *fal_version);
 };
 
-static void *neuron_handle = NULL;
+struct neuron_ops_internal {
+	NRT_STATUS (*nrt_get_dmabuf_fd)(uint64_t va, uint64_t size, int* fd);
+} neuron_ops_internal;
+
+static void *neuron_handle;
 static struct neuron_ops neuron_ops;
 
+NRT_STATUS nrt_get_dmabuf_fd_proxy(uint64_t va, uint64_t size, int* fd, uint64_t *offset) {
+	*offset = 0;
+	return neuron_ops_internal.nrt_get_dmabuf_fd(va, size, fd);
+}
+
+static void ft_setup_nrt_get_dmabuf_fd(void)
+{
+	neuron_ops.nrt_get_dmabuf_fd = dlsym(neuron_handle, "nrt_get_dmabuf_fd_v2");
+	if (neuron_ops.nrt_get_dmabuf_fd)
+		return;
+
+	FI_INFO(&core_prov, FI_LOG_CORE,
+		"Failed to find nrt_get_dmabuf_fd_v2, falling back to nrt_get_dmabuf_fd");
+
+	neuron_ops_internal.nrt_get_dmabuf_fd = dlsym(neuron_handle, "nrt_get_dmabuf_fd");
+	if (!neuron_ops_internal.nrt_get_dmabuf_fd) {
+		FI_INFO(&core_prov, FI_LOG_CORE,
+			"Failed to find nrt_get_dmabuf_fd, "
+			"dmabuf feature will not be used for Neuron devices\n");
+		return;
+	}
+	neuron_ops.nrt_get_dmabuf_fd = &nrt_get_dmabuf_fd_proxy;
+}
+
 /*
  * List to lookup the handle based on the pointer. Not optimal, but probably
  * fine for fabtests and this is better than changing the alloc/free functions
@@ -123,11 +151,7 @@ int ft_neuron_init(void)
 		goto err;
 	}
 
-	neuron_ops.nrt_get_dmabuf_fd = dlsym(neuron_handle, "nrt_get_dmabuf_fd");
-	if (!neuron_ops.nrt_get_dmabuf_fd) {
-		FT_INFO("Failed to find nrt_get_dmabuf_fd, "
-			"dmabuf feature will not be used for Neuron devices\n");
-	}
+	ft_setup_nrt_get_dmabuf_fd();
 
 	dlist_init(&neuron_alloc_list);
 
@@ -345,18 +369,9 @@ int ft_neuron_get_dmabuf_fd(void *addr, size_t size, int *fd,
 			 uint64_t *offset)
 {
 	int ret = NRT_SUCCESS;
-	struct neuron_allocation *region;
-
-	/*
-	 * The assumption is that nrt_get_dmabuf_fd() would fail for
-	 * any addr that is not the starting address of the dma-buf
-	 * object. Otherwise we need a low level op to get the base
-	 * address of the dma-buf object.
-	 */
-	*offset = ft_neuron_find_region(addr, &region);
 
 	if (neuron_ops.nrt_get_dmabuf_fd) {
-		ret = neuron_ops.nrt_get_dmabuf_fd((uintptr_t)region->ptr, size, fd);
+		ret = neuron_ops.nrt_get_dmabuf_fd((uint64_t)addr, size, fd, offset);
 		if (ret != NRT_SUCCESS) {
 			FT_WARN("failed to get dmabuf fd\n");
 			return -FI_EIO;

src/hmem_neuron.c

@@ -51,13 +51,41 @@ struct neuron_ops {
 	void (*nrt_tensor_free)(nrt_tensor_t **tensor);
 	void *(*nrt_tensor_get_va)(const nrt_tensor_t *tensor);
 	NRT_STATUS (*nrt_memcpy_to_device)(void *dest, const void *src, size_t size);
-	NRT_STATUS (*nrt_get_dmabuf_fd)(uint64_t va, uint64_t size, int* fd);
+	NRT_STATUS (*nrt_get_dmabuf_fd)(uint64_t va, uint64_t size, int* fd, uint64_t *offset);
 	NRT_STATUS (*nrt_get_total_nc_count)(uint32_t *nc_count);
 };
 
+struct neuron_ops_internal {
+	NRT_STATUS (*nrt_get_dmabuf_fd)(uint64_t va, uint64_t size, int* fd);
+} neuron_ops_internal;
+
 static void *neuron_handle;
 static struct neuron_ops neuron_ops;
 
+NRT_STATUS nrt_get_dmabuf_fd_proxy(uint64_t va, uint64_t size, int* fd, uint64_t *offset) {
+	*offset = 0;
+	return neuron_ops_internal.nrt_get_dmabuf_fd(va, size, fd);
+}
+
+static void setup_nrt_get_dmabuf_fd(void)
+{
+	neuron_ops.nrt_get_dmabuf_fd = dlsym(neuron_handle, "nrt_get_dmabuf_fd_v2");
+	if (neuron_ops.nrt_get_dmabuf_fd)
+		return;
+
+	FI_INFO(&core_prov, FI_LOG_CORE,
+		"Failed to find nrt_get_dmabuf_fd_v2, falling back to nrt_get_dmabuf_fd");
+
+	neuron_ops_internal.nrt_get_dmabuf_fd = dlsym(neuron_handle, "nrt_get_dmabuf_fd");
+	if (!neuron_ops_internal.nrt_get_dmabuf_fd) {
+		FI_INFO(&core_prov, FI_LOG_CORE,
+			"Failed to find nrt_get_dmabuf_fd, "
+			"dmabuf feature will not be used for Neuron devices\n");
+		return;
+	}
+	neuron_ops.nrt_get_dmabuf_fd = &nrt_get_dmabuf_fd_proxy;
+}
+
 static int neuron_dl_init(void)
 {
 	neuron_handle = dlopen("libnrt.so.1", RTLD_NOW);
@@ -91,12 +119,7 @@ static int neuron_dl_init(void)
 		goto err;
 	}
 
-	neuron_ops.nrt_get_dmabuf_fd = dlsym(neuron_handle, "nrt_get_dmabuf_fd");
-	if (!neuron_ops.nrt_get_dmabuf_fd) {
-		FI_INFO(&core_prov, FI_LOG_CORE,
-			"Failed to find nrt_get_dmabuf_fd, "
-			"dmabuf feature will not be used for Neuron devices\n");
-	}
+	setup_nrt_get_dmabuf_fd();
 
 	neuron_ops.nrt_get_total_nc_count = dlsym(neuron_handle, "nrt_get_total_nc_count");
 	if (!neuron_ops.nrt_get_total_nc_count) {
@@ -228,16 +251,9 @@ int neuron_get_dmabuf_fd(const void *addr, uint64_t size, int *fd,
 		return -FI_EOPNOTSUPP;
 	}
 
-	ret = neuron_ops.nrt_get_dmabuf_fd((uintptr_t)addr, size, fd);
+	ret = neuron_ops.nrt_get_dmabuf_fd((uintptr_t)addr, size, fd, offset);
 
 	if (ret == NRT_SUCCESS) {
-		/*
-		 * The assumption is that nrt_get_dmabuf_fd() would fail for
-		 * any addr that is not the starting address of the dma-buf
-		 * object. Otherwise we need a low level op to get the base
-		 * address of the dma-buf object.
-		 */
-		*offset = 0;
 		return FI_SUCCESS;
 	} else if (ret == NRT_RESOURCE) {
 		/* real error from Neuron */