[SE] Cache CUDA modules

[jhen]

Summary

Instead of reloading a module if the same kernel is requested multiple times,
cache the loaded module and return the cached value.

The CUDAPlatformDevice now also keeps handles to all its modules so they can be
unloaded if the device is cleared.

[jlebar]

streamexecutor/lib/platforms/cuda/CUDAPlatformDevice.cpp

130

Hm. This makes a copy of "Code" in the map. And also, every time we do a lookup, we're going to have to compare the whole PTX strings. Which are potentially very long.

Is there no other identifier we could use as the map key?

[jhen]

streamexecutor/lib/platforms/cuda/CUDAPlatformDevice.cpp

130

Unfortunately, I don't think there is currently any other identifier that won't ever lead to false matches. Using the whole string as the key is what the original developers did because they couldn't find a better solution, so I'm mostly just following their lead here.

There are some things we could do with randomly generated UUIDs that would work for all practical purposes, but I don't want to worry about correctly generating UUIDs.

I also have the following idea that seems a little complex. Does it seem too complex to you?:

Use a static integer with atomic increments (or mutex or whatever) to give a unique ID to each MultiKernelLoaderSpec instance in the process and then have each instance assign a unique ID to each piece of code that is registered with it. The pair of MultiKernelLoaderSpec ID and code ID will uniquely identify a piece of code and can be used as a key in the module cache.

[jhen]

We've decided to come at this problem from a different angle, so I'm abandoning this revision.

streamexecutor/include/streamexecutor/platforms/cuda/CUDAPlatformDevice.h

 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef STREAMEXECUTOR_PLATFORMS_CUDA_CUDAPLATFORMDEVICE_H
 #define STREAMEXECUTOR_PLATFORMS_CUDA_CUDAPLATFORMDEVICE_H
 
 #include "streamexecutor/PlatformDevice.h"
 
+#include "llvm/Support/Mutex.h"
+
+#include <map>
+
+struct CUfunc_st;
+struct CUmod_st;
+
 namespace streamexecutor {
 namespace cuda {
 
 Error CUresultToError(int CUResult, const llvm::Twine &Message);
 
 class CUDAPlatformDevice : public PlatformDevice {
 public:
   static Expected<CUDAPlatformDevice> create(size_t DeviceIndex);
public:
   Error synchronousCopyD2D(const void *DeviceDstHandle, size_t DstByteOffset,
                            const void *DeviceSrcHandle, size_t SrcByteOffset,
                            size_t ByteCount) override;
 
 private:
   CUDAPlatformDevice(size_t DeviceIndex) : DeviceIndex(DeviceIndex) {}
 
   int DeviceIndex;
+  llvm::sys::Mutex Mutex;
+  std::map<std::string, std::pair<CUmod_st *, CUfunc_st *>> LoadedModules;
 };
 
 } // namespace cuda
 } // namespace streamexecutor
 
 #endif // STREAMEXECUTOR_PLATFORMS_CUDA_CUDAPLATFORMDEVICE_H

streamexecutor/lib/platforms/cuda/CUDAPlatformDevice.cpp

 CUDAPlatformDevice::~CUDAPlatformDevice() {
   CUresult Result = cuDevicePrimaryCtxRelease(DeviceIndex);
   (void)Result;
   // TODO(jhen): Log error.
 }
 
 Expected<const void *>
 CUDAPlatformDevice::createKernel(const MultiKernelLoaderSpec &Spec) {
-  // TODO(jhen): Maybe first check loaded modules?
   if (!Spec.hasCUDAPTXInMemory())
     return make_error("no CUDA code available to create kernel");
 
   CUdevice Device = static_cast<int>(DeviceIndex);
   int ComputeCapabilityMajor = 0;
   int ComputeCapabilityMinor = 0;
   if (CUresult Result = cuDeviceGetAttribute(
           &ComputeCapabilityMajor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
CUDAPlatformDevice::createKernel(const MultiKernelLoaderSpec &Spec) {
   const char *Code = Spec.getCUDAPTXInMemory().getCode(ComputeCapabilityMajor,
                                                        ComputeCapabilityMinor);
 
   if (!Code)
     return make_error("no suitable CUDA source found for compute capability " +
                       llvm::Twine(ComputeCapabilityMajor) + "." +
                       llvm::Twine(ComputeCapabilityMinor));
 
-  CUmodule Module;
+  CUfunction Function = nullptr;
+  {
+    llvm::sys::ScopedLock Lock(Mutex);
+    auto Iterator = LoadedModules.find(Code);
+    if (Iterator == LoadedModules.end()) {
+      CUmodule Module = nullptr;
       if (CUresult Result = cuModuleLoadData(&Module, Code))
         return CUresultToError(Result, "cuModuleLoadData");
-
-  CUfunction Function;
-  if (CUresult Result =
-          cuModuleGetFunction(&Function, Module, Spec.getKernelName().c_str()))
+      if (CUresult Result = cuModuleGetFunction(&Function, Module,
+                                                Spec.getKernelName().c_str()))
         return CUresultToError(Result, "cuModuleGetFunction");
-
-  // TODO(jhen): Should I save this function pointer in case someone asks for
-  // it again?
-
-  // TODO(jhen): Should I save the module pointer so I can unload it when I
-  // destroy this device?
+      LoadedModules.emplace(Code, std::make_pair(Module, Function));
+    } else
+      Function = Iterator->second.second;
+  }
 
   return static_cast<const void *>(Function);
 }
 
 Error CUDAPlatformDevice::destroyKernel(const void *Handle) {
-  // TODO(jhen): Maybe keep track of kernels for each module and unload the
-  // module after they are all destroyed.
   return Error::success();
 }
 
 Expected<const void *> CUDAPlatformDevice::createStream() {
   CUstream Stream;
   if (CUresult Result = cuStreamCreate(&Stream, CU_STREAM_DEFAULT))
     return CUresultToError(Result, "cuStreamCreate");
   return Stream;