Refactor: move cufile header checks to the shim layer

cpp/include/kvikio/buffer.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -50,7 +50,6 @@ inline void buffer_register(const void* devPtr_base,
                             const std::vector<int>& errors_to_ignore = std::vector<int>())
 {
   if (defaults::compat_mode()) { return; }
-#ifdef KVIKIO_CUFILE_FOUND
   CUfileError_t status = cuFileAPI::instance().BufRegister(devPtr_base, size, flags);
   if (status.err != CU_FILE_SUCCESS) {
     // Check if `status.err` is in `errors_to_ignore`
@@ -59,7 +58,6 @@ inline void buffer_register(const void* devPtr_base,
       CUFILE_TRY(status);
     }
   }
-#endif
 }
 
 /**
@@ -70,9 +68,7 @@ inline void buffer_register(const void* devPtr_base,
 inline void buffer_deregister(const void* devPtr_base)
 {
   if (defaults::compat_mode()) { return; }
-#ifdef KVIKIO_CUFILE_FOUND
   CUFILE_TRY(cuFileAPI::instance().BufDeregister(devPtr_base));
-#endif
 }
 
 /**

cpp/include/kvikio/error.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -56,7 +56,6 @@ struct CUfileException : public std::runtime_error {
 #define CUDA_DRIVER_TRY_1(_call) CUDA_DRIVER_TRY_2(_call, kvikio::CUfileException)
 #endif
 
-#ifdef KVIKIO_CUFILE_FOUND
 #ifndef CUFILE_TRY
 #define CUFILE_TRY(...)                                         \
   GET_CUFILE_TRY_MACRO(__VA_ARGS__, CUFILE_TRY_2, CUFILE_TRY_1) \
@@ -77,15 +76,13 @@ struct CUfileException : public std::runtime_error {
   } while (0)
 #define CUFILE_TRY_1(_call) CUFILE_TRY_2(_call, kvikio::CUfileException)
 #endif
-#endif
 
 #ifndef CUFILE_CHECK_STREAM_IO
 #define CUFILE_CHECK_STREAM_IO(...)                                  \
   GET_CUFILE_CHECK_STREAM_IO_MACRO(                                  \
     __VA_ARGS__, CUFILE_CHECK_STREAM_IO_2, CUFILE_CHECK_STREAM_IO_1) \
   (__VA_ARGS__)
 #define GET_CUFILE_CHECK_STREAM_IO_MACRO(_1, _2, NAME, ...) NAME
-#ifdef KVIKIO_CUFILE_FOUND
 #define CUFILE_CHECK_STREAM_IO_2(_nbytes_done, _exception_type)                            \
   do {                                                                                     \
     auto const _nbytes = *(_nbytes_done);                                                  \
@@ -94,12 +91,6 @@ struct CUfileException : public std::runtime_error {
                              KVIKIO_STRINGIFY(__LINE__) + ": " + std::to_string(_nbytes)}; \
     }                                                                                      \
   } while (0)
-#else
-// if cufile isn't available, we don't do anything in the body
-#define CUFILE_CHECK_STREAM_IO_2(_nbytes_done, _exception_type) \
-  do {                                                          \
-  } while (0)
-#endif
 #define CUFILE_CHECK_STREAM_IO_1(_call) CUFILE_CHECK_STREAM_IO_2(_call, kvikio::CUfileException)
 #endif
 

cpp/include/kvikio/file_handle.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -169,14 +169,13 @@ class FileHandle {
       _compat_mode = true;  // Fall back to compat mode if we cannot open the file with O_DIRECT
     }
 
-    if (_compat_mode) { return; }
-#ifdef KVIKIO_CUFILE_FOUND
-    CUfileDescr_t desc{};  // It is important to set to zero!
-    desc.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD;
-    // NOLINTNEXTLINE(cppcoreguidelines-pro-type-union-access)
-    desc.handle.fd = _fd_direct_on;
-    CUFILE_TRY(cuFileAPI::instance().HandleRegister(&_handle, &desc));
-#endif
+    if (!_compat_mode) {
+      CUfileDescr_t desc{};  // It is important to set to zero!
+      desc.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD;
+      // NOLINTNEXTLINE(cppcoreguidelines-pro-type-union-access)
+      desc.handle.fd = _fd_direct_on;
+      CUFILE_TRY(cuFileAPI::instance().HandleRegister(&_handle, &desc));
+    }
   }
 
   /**
@@ -214,11 +213,7 @@ class FileHandle {
   {
     if (closed()) { return; }
 
-    if (!_compat_mode) {
-#ifdef KVIKIO_CUFILE_FOUND
-      cuFileAPI::instance().HandleDeregister(_handle);
-#endif
-    }
+    if (!_compat_mode) { cuFileAPI::instance().HandleDeregister(_handle); }
     ::close(_fd_direct_off);
     if (_fd_direct_on != -1) { ::close(_fd_direct_on); }
     _fd_direct_on  = -1;
@@ -310,7 +305,6 @@ class FileHandle {
     if (_compat_mode) {
       return posix_device_read(_fd_direct_off, devPtr_base, size, file_offset, devPtr_offset);
     }
-#ifdef KVIKIO_CUFILE_FOUND
     ssize_t ret = cuFileAPI::instance().Read(
       _handle, devPtr_base, size, convert_size2off(file_offset), convert_size2off(devPtr_offset));
     if (ret == -1) {
@@ -321,9 +315,6 @@ class FileHandle {
                             KVIKIO_STRINGIFY(__LINE__) + ": " + CUFILE_ERRSTR(ret));
     }
     return ret;
-#else
-    throw CUfileException("KvikIO not compiled with cuFile.h");
-#endif
   }
 
   /**
@@ -360,7 +351,6 @@ class FileHandle {
     if (_compat_mode) {
       return posix_device_write(_fd_direct_off, devPtr_base, size, file_offset, devPtr_offset);
     }
-#ifdef KVIKIO_CUFILE_FOUND
     ssize_t ret = cuFileAPI::instance().Write(
       _handle, devPtr_base, size, convert_size2off(file_offset), convert_size2off(devPtr_offset));
     if (ret == -1) {
@@ -371,9 +361,6 @@ class FileHandle {
                             KVIKIO_STRINGIFY(__LINE__) + ": " + CUFILE_ERRSTR(ret));
     }
     return ret;
-#else
-    throw CUfileException("KvikIO not compiled with cuFile.h");
-#endif
   }
 
   /**
@@ -539,14 +526,11 @@ class FileHandle {
                   ssize_t* bytes_read_p,
                   CUstream stream)
   {
-#ifdef KVIKIO_CUFILE_STREAM_API_FOUND
     if (kvikio::is_batch_and_stream_available() && !_compat_mode) {
       CUFILE_TRY(cuFileAPI::instance().ReadAsync(
         _handle, devPtr_base, size_p, file_offset_p, devPtr_offset_p, bytes_read_p, stream));
       return;
     }
-#endif
-
     CUDA_DRIVER_TRY(cudaAPI::instance().StreamSynchronize(stream));
     *bytes_read_p =
       static_cast<ssize_t>(read(devPtr_base, *size_p, *file_offset_p, *devPtr_offset_p));
@@ -632,14 +616,11 @@ class FileHandle {
                    ssize_t* bytes_written_p,
                    CUstream stream)
   {
-#ifdef KVIKIO_CUFILE_STREAM_API_FOUND
     if (kvikio::is_batch_and_stream_available() && !_compat_mode) {
       CUFILE_TRY(cuFileAPI::instance().WriteAsync(
         _handle, devPtr_base, size_p, file_offset_p, devPtr_offset_p, bytes_written_p, stream));
       return;
     }
-#endif
-
     CUDA_DRIVER_TRY(cudaAPI::instance().StreamSynchronize(stream));
     *bytes_written_p =
       static_cast<ssize_t>(write(devPtr_base, *size_p, *file_offset_p, *devPtr_offset_p));

cpp/include/kvikio/shim/cufile.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,8 +23,6 @@
 
 namespace kvikio {
 
-#ifdef KVIKIO_CUFILE_FOUND
-
 /**
  * @brief Shim layer of the cuFile C-API
  *
@@ -46,24 +44,20 @@ class cuFileAPI {
   decltype(cuFileDriverSetPollMode)* DriverSetPollMode{nullptr};
   decltype(cuFileDriverSetMaxCacheSize)* DriverSetMaxCacheSize{nullptr};
   decltype(cuFileDriverSetMaxPinnedMemSize)* DriverSetMaxPinnedMemSize{nullptr};
-
-#ifdef KVIKIO_CUFILE_BATCH_API_FOUND
   decltype(cuFileBatchIOSetUp)* BatchIOSetUp{nullptr};
   decltype(cuFileBatchIOSubmit)* BatchIOSubmit{nullptr};
   decltype(cuFileBatchIOGetStatus)* BatchIOGetStatus{nullptr};
   decltype(cuFileBatchIOCancel)* BatchIOCancel{nullptr};
   decltype(cuFileBatchIODestroy)* BatchIODestroy{nullptr};
-#endif
-
-#ifdef KVIKIO_CUFILE_STREAM_API_FOUND
   decltype(cuFileReadAsync)* ReadAsync{nullptr};
   decltype(cuFileWriteAsync)* WriteAsync{nullptr};
   decltype(cuFileStreamRegister)* StreamRegister{nullptr};
   decltype(cuFileStreamDeregister)* StreamDeregister{nullptr};
-#endif
+
   bool stream_available = false;
 
  private:
+#ifdef KVIKIO_CUFILE_FOUND
   cuFileAPI()
   {
     // CUDA versions before CUDA 11.7.1 did not ship libcufile.so.0, so this is
@@ -128,10 +122,15 @@ class cuFileAPI {
                 << std::endl;
     }
   }
+#else
+  cuFileAPI() { throw std::runtime_error(CUFILE_ERRSTR(0)); }
+#endif
 
  public:
-  cuFileAPI(cuFileAPI const&)      = delete;
-  void operator=(cuFileAPI const&) = delete;
+  cuFileAPI(cuFileAPI const&)       = delete;
+  void operator=(cuFileAPI const&)  = delete;
+  cuFileAPI(cuFileAPI const&&)      = delete;
+  void operator=(cuFileAPI const&&) = delete;
 
   static cuFileAPI& instance()
   {
@@ -140,8 +139,6 @@ class cuFileAPI {
   }
 };
 
-#endif
-
 /**
  * @brief Check if the cuFile library is available
  *

cpp/include/kvikio/shim/cufile_h_wrapper.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,6 +15,9 @@
  */
 #pragma once
 
+#include <cuda.h>
+#include <sys/types.h>
+
 /**
  * In order to support compilation when `cufile.h` isn't available, we
  * wrap all use of cufile in a `#ifdef KVIKIO_CUFILE_FOUND` guard.
@@ -25,15 +28,55 @@
 #ifdef KVIKIO_CUFILE_FOUND
 #include <cufile.h>
 #else
+
+// If cuFile isn't defined, we define some of the data types here.
+// Notice, this doesn't need to be ABI compatible with the cufile definitions.
+
+using CUfileHandle_t = void*;
+using CUfileOpError  = int;
+#define CUFILE_ERRSTR(x)          ("KvikIO not compiled with cuFile.h")
+#define CU_FILE_SUCCESS           0
+#define CU_FILE_CUDA_DRIVER_ERROR 1
+
+struct CUfileError_t {
+  CUfileOpError err;  // cufile error
+  CUresult cu_err;    // cuda driver error
+};
+
 using CUfileDriverControlFlags_t = enum CUfileDriverControlFlags {
-  CU_FILE_USE_POLL_MODE     = 0, /*!< use POLL mode. properties.use_poll_mode*/
-  CU_FILE_ALLOW_COMPAT_MODE = 1  /*!< allow COMPATIBILITY mode. properties.allow_compat_mode*/
+  CU_FILE_USE_POLL_MODE     = 0,
+  CU_FILE_ALLOW_COMPAT_MODE = 1
 };
-using CUfileHandle_t = void*;
+
+enum CUfileFileHandleType { CU_FILE_HANDLE_TYPE_OPAQUE_FD = 1 };
+
+struct CUfileDescr_t {
+  enum CUfileFileHandleType type;
+  struct handle_t {
+    int fd;
+  } handle;
+};
+
+static inline const char* cufileop_status_error(CUfileOpError err) { return CUFILE_ERRSTR(err); };
+CUfileError_t cuFileHandleRegister(...);
+CUfileError_t cuFileHandleDeregister(...);
+ssize_t cuFileRead(...);
+ssize_t cuFileWrite(...);
+CUfileError_t cuFileBufRegister(...);
+CUfileError_t cuFileBufDeregister(...);
+CUfileError_t cuFileDriverOpen(...);
+CUfileError_t cuFileDriverClose(...);
+CUfileError_t cuFileDriverGetProperties(...);
+CUfileError_t cuFileDriverSetPollMode(...);
+CUfileError_t cuFileDriverSetMaxCacheSize(...);
+CUfileError_t cuFileDriverSetMaxPinnedMemSize(...);
+
 #endif
 
 // If the Batch API isn't defined, we define some of the data types here.
-// Notice, this doesn't need to be ABI compatible with the cufile definitions.
+// Notice, this doesn't need to be ABI compatible with the cufile definitions and
+// the lack of definitions is not a problem because the linker will never look for
+// these symbols because the "real" function calls are made through the shim instance.
 #ifndef KVIKIO_CUFILE_BATCH_API_FOUND
 typedef enum CUfileOpcode { CUFILE_READ = 0, CUFILE_WRITE } CUfileOpcode_t;
 
@@ -52,4 +95,18 @@ typedef struct CUfileIOEvents {
   CUfileStatus_t status; /* status of the operation */
   size_t ret;            /* -ve error or amount of I/O done. */
 } CUfileIOEvents_t;
+
+CUfileError_t cuFileBatchIOSetUp(...);
+CUfileError_t cuFileBatchIOSubmit(...);
+CUfileError_t cuFileBatchIOGetStatus(...);
+CUfileError_t cuFileBatchIOCancel(...);
+CUfileError_t cuFileBatchIODestroy(...);
+#endif
+
+// If the Stream API isn't defined, we define some of the data types here.
+#ifndef KVIKIO_CUFILE_STREAM_API_FOUND
+CUfileError_t cuFileReadAsync(...);
+CUfileError_t cuFileWriteAsync(...);
+CUfileError_t cuFileStreamRegister(...);
+CUfileError_t cuFileStreamDeregister(...);
 #endif