1616#pragma once
1717
1818#include < stdexcept>
19+ #include < string>
1920
20- #include < iostream>
2121#include < kvikio/shim/cufile_h_wrapper.hpp>
2222#include < kvikio/shim/utils.hpp>
2323
@@ -38,8 +38,6 @@ class cuFileAPI {
3838 decltype (cuFileWrite)* Write{nullptr };
3939 decltype (cuFileBufRegister)* BufRegister{nullptr };
4040 decltype (cuFileBufDeregister)* BufDeregister{nullptr };
41- decltype (cuFileDriverOpen)* DriverOpen{nullptr };
42- decltype (cuFileDriverClose)* DriverClose{nullptr };
4341 decltype (cuFileDriverGetProperties)* DriverGetProperties{nullptr };
4442 decltype (cuFileDriverSetPollMode)* DriverSetPollMode{nullptr };
4543 decltype (cuFileDriverSetMaxCacheSize)* DriverSetMaxCacheSize{nullptr };
@@ -54,6 +52,12 @@ class cuFileAPI {
5452 decltype (cuFileStreamRegister)* StreamRegister{nullptr };
5553 decltype (cuFileStreamDeregister)* StreamDeregister{nullptr };
5654
55+ private:
56+ // Don't call driver open and close directly, use `.driver_open()` and `.driver_close()`.
57+ decltype (cuFileDriverOpen)* DriverOpen{nullptr };
58+ decltype (cuFileDriverClose)* DriverClose{nullptr };
59+
60+ public:
5761 bool stream_available = false ;
5862
5963 private:
@@ -105,25 +109,25 @@ class cuFileAPI {
105109 }
106110#endif
107111
108- // cuFile is supposed to open and close the driver automatically but because of a bug in
109- // CUDA 11.8, it sometimes segfault. See <https://github.com/rapidsai/kvikio/issues/159>.
110- CUfileError_t const error = DriverOpen ();
111- if (error.err != CU_FILE_SUCCESS ) {
112- throw std::runtime_error (std::string{" cuFile error at: " " :"
113- KVIKIO_STRINGIFY (__LINE__) + " : "
114- cufileop_status_error (error.err ));
115- }
112+ // cuFile is supposed to open and close the driver automatically but
113+ // because of a bug in cuFile v1.4 (CUDA v11.8) it sometimes segfaults:
114+ // <https://github.com/rapidsai/kvikio/issues/159>.
115+ // We use the stream API as a version indicator of cuFile since it was introduced
116+ // in cuFile v1.7 (CUDA v12.2).
117+ if (!stream_available) { driver_open (); }
116118 }
119+
120+ // Notice, we have to close the driver at program exit (if we opened it) even though we are
121+ // not allowed to call CUDA after main[1]. This is because, cuFile will segfault if the
122+ // driver isn't closed on program exit i.e. we are doomed if we do, doomed if we don't, but
123+ // this seems to be the lesser of two evils.
124+ // [1] <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#initialization>
117125 ~cuFileAPI ()
118126 {
119- CUfileError_t const error = DriverClose ();
120- if (error.err != CU_FILE_SUCCESS ) {
121- std::cerr << " Unable to close GDS file driver: " cufileop_status_error (error.err )
122- << std::endl;
123- }
127+ if (!stream_available) { driver_close (); }
124128 }
125129#else
126- cuFileAPI () { throw std::runtime_error (CUFILE_ERRSTR ( 0 ) ); }
130+ cuFileAPI () { throw std::runtime_error (" KvikIO not compiled with cuFile.h "
127131#endif
128132
129133 public:
@@ -137,6 +141,33 @@ class cuFileAPI {
137141 static cuFileAPI _instance;
138142 return _instance;
139143 }
144+
145+ /* *
146+ * @brief Open the cuFile driver
147+ *
148+ * cuFile allows multiple calls to `cufileDriverOpen()`, only the first call opens
149+ * the driver, but every call should have a matching call to `cufileDriverClose()`.
150+ */
151+ void driver_open ()
152+ {
153+ CUfileError_t const error = DriverOpen ();
154+ if (error.err != CU_FILE_SUCCESS ) {
155+ throw std::runtime_error (std::string{" Unable to open GDS file driver: "
156+ cufileop_status_error (error.err ));
157+ }
158+ }
159+
160+ /* *
161+ * @brief Close the cuFile driver
162+ */
163+ void driver_close ()
164+ {
165+ CUfileError_t const error = DriverClose ();
166+ if (error.err != CU_FILE_SUCCESS ) {
167+ throw std::runtime_error (std::string{" Unable to close GDS file driver: "
168+ cufileop_status_error (error.err ));
169+ }
170+ }
140171};
141172
142173/* *
0 commit comments