[Vulkan][Refactor] Move ownership of per-CPU-thread objects to VulkanDeviceAPI

[Lunderberg]

Continuing refactoring from #8188. The goal of this PR is to establish clear ownership of vulkan resources. Previously, resources were owned both by VulkanDeviceAPI and by VulkanThreadEntry, depending on whether it is a per-CPU-thread resource. This PR introduces a ThreadMap<T> class that behaves similarly to thread_local, but can be used for non-static data members, then uses ThreadMap<T> to move all ownership of vulkan resources to be relative to the VulkanDeviceAPI. Individual steps are in separate commits for readability.

• Moved ownership of per-thread active device id to VulkanDeviceAPI

  • Previously, the active device was owned by VulkanThreadEntry, so lookups to multiple global variables were required. Now, everything goes from the VulkanDeviceAPI.

  • Removed VulkanThreadEntry, as all functionality has been moved to either VulkanDevice or VulkanDeviceAPI.

• Moved ownership of per-thread workspace pool to VulkanDeviceAPI

  • Previously, the WorkspacePool was owned by VulkanThreadEntry, and required a lookup from VulkanDeviceAPI::AllocWorkspace. As a result, non-global VulkanDeviceAPI would interact with each other.

• Move ownership of per-thread uniform buffer to VulkanDevice

  • Previously, VulkanUniformBuffer was owned by VulkanThreadEntry, so any use required looking up both the thread entry and the device. Now, thread-specific lookup is handled in the VulkanDevice class.

• Move the VulkanStagingBuffer to be owned by the VulkanDevice

  • Previously, was owned by VulkanThreadEntry, so any use required looking up both the thread entry and the device. Now, thread-specific lookup is handled in the VulkanDevice class.

• Pulled VulkanBuffer allocation/deallocation into constructor/destructor.

  • VulkanBuffer owns the VkBuffer and VkDeviceMemory that it allocates, and deallocates on destruction.

  • VulkanHostVisibleBuffer owns a VulkanBuffer, and additional calls vkUnmapMemory on destruction.

• Moved VulkanStream ownership from VulkanThreadEntry to VulkanDevice

  • Implemented ThreadMap, a container for per-thread objects. Unlike dmlc::ThreadLocalStore, ThreadMap is intended for use as a non-static thread-specific lookup.

  • Added ThreadMap as a member to VulkanDevice, updated all uses.

[Lunderberg]

Potential reviewers: @masahi @tmoreau89

[tqchen]

Thanks @Lunderberg , it would still be useful to think about the overhead. While thread local is a different way of organizing things, the direct TLS lookup have a lower overhead than a lock based cache lookup. Assuming we only have the singleton case it might still be desirable to use TLS when possible(i understand that might restrict us to singleton usecases)

Another minor thing is we usually keep the non-public facing apis in src/ so they are not visible to the downstream deps. This would help reduce the possible future API compact problem

[Lunderberg]

Thank you @tqchen , and that makes sense with the potential overhead issues. I had been assuming that the main performance cost would be in the std::unordered_map lookup. I'd expect that lookup to be roughly equivalent, since while the local map requires a lookup by std::thread::id, the TLS requires a lookup by device_id. Having two threads block should be rare, since the locks are acquired non-exclusively whenever the thread's object already exists, but you are right that I didn't consider the cost of acquiring that shared lock in the first place.

I'll see if I can measure the performance difference, and will also see whether I can refactor to maintain the ownership model of everything deriving from the VulkanDeviceAPI, while keeping the minimal overhead of TLS.

Regarding the non-public apis, that makes sense. It wasn't vulkan-specific, and so I didn't want it all the way in src/runtime/vulkan, but hadn't considered that impact of adding to the public api. I will move thread_map.h (or whatever thread_local variant replaces it) to be in src/runtime instead.

[tqchen]

That sounds good. Thanks @Lunderberg , i agree likely the hit could be negligible considering the launching overhead

[Lunderberg]

Running some performance tests, it looks like the refactor has very little impact on the overall runtime. The plots below show the Q1/median/Q3 runtimes for different low-level tasks that would need to access the thread-specific resources. The only significant difference is for the copying data to the device, which is slightly higher for very small buffer copies.

Benchmarking details: Used pytest-benchmark, mostly with default settings. Number of iterations chosen based on runtime of the first iteration such that each data point is collected in ~1 second. Repeated initialization was allowed to use up to 10 seconds per data point.

[Lunderberg]

On second thought, it looks like that difference was because pytest-benchmark does not include warmup runs by default. That run just happened to be the first that required initialization, which messed up some of the timing. With warmup runs, the performance is matched before/after the refactor.

[masahi]

@Lunderberg @tqchen Is this good to go?

[Lunderberg]

The performance tests are matched pre/post change for cases that access the thread-specific resources, so I think it's good to go.

@tqchen Any follow-up concerns?

[masahi]

Thanks @Lunderberg, surprisingly after this change, the segfault on my APU that was happening on process exit is somehow gone!

[Lunderberg]

Ooh, that's a nice side-effect to get! I guess that confirms that the segfault was caused by out-of-order destruction, though I'm still not quite sure how that happened, given that we checked the order of construction.