Implement the logic to handle failures with cufile.

liuliu · Dec 4, 2024 · 2a59318 · 2a59318
1 parent 98eb262
commit 2a59318
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/lib/nnc/gpu/ccv_nnc_compat.cu b/lib/nnc/gpu/ccv_nnc_compat.cu
@@ -1,6 +1,7 @@
 #include "ccv_nnc_compat.h"
 #include <cufile.h> // For GPUDirect Storage
 extern "C" {
+#include <sys/mman.h>
 #include <nnc/ccv_nnc_easy.h>
 #include <nnc/_ccv_nnc_stream.h>
 #include "3rdparty/khash/khash.h"
@@ -316,7 +317,16 @@ void cufileread(const int fd, const off_t file_offset, void* const buf, const si
 		.fs_ops = 0,
 	};
 	CUfileHandle_t file_handle;
-	CUFILE_ENFORCE(cuFileHandleRegister(&file_handle, &file_descr));
+	const CUfileError_t status = cuFileHandleRegister(&file_handle, &file_descr);
+	if (status.err != CU_FILE_SUCCESS)
+	{
+		PRINT(CCV_CLI_ERROR, "[%s:%d]:CUFILE - Error: %s\n", __FILE__, __LINE__, CUFILE_ERRSTR(status.err));
+		void* bufptr = mmap(0, size, PROT_READ, MAP_PRIVATE, fd, file_offset);
+		madvise(bufptr, size, MADV_SEQUENTIAL | MADV_WILLNEED);
+		cumemcpy(buf, CCV_TENSOR_GPU_MEMORY, bufptr, CCV_TENSOR_CPU_MEMORY, size);
+		munmap(bufptr, size);
+		return;
+	}
 	cuFileRead(file_handle, buf, size, file_offset, 0);
 	cuFileHandleDeregister(file_handle);
 }

diff --git a/lib/nnc/gpu/ccv_nnc_compat.h b/lib/nnc/gpu/ccv_nnc_compat.h
@@ -127,7 +127,7 @@ CCV_WARN_UNUSED(size_t) ccv_nnc_cublas_workspace_size_in_bytes(const ccv_nnc_ten
 #define CUFILE_ENFORCE(status) do {                               \
 	const CUfileError_t __status = status;                        \
 	if (__status.err != CU_FILE_SUCCESS) {                        \
-		printf("[%s:%d]:CUDA - Error: %s\n",                      \
+		printf("[%s:%d]:CUFILE - Error: %s\n",                    \
 				__FILE__, __LINE__, CUFILE_ERRSTR(__status.err)); \
 		assert(0);                                                \
 		exit(EXIT_FAILURE);                                       \