[vulkan] Refactor Command (#80720)
Part of a refactor of the Vulkan codebase.
This diff introduces the new `CommandBuffer` class as well as a revamped `CommandPool` class.
## `CommandPool`
This class is a wrapper around a `VkCommandPool` that is responsible for allocating command buffers. It only has one function of interest, `get_new_cmd()`, which will fetch a command buffer from the pool. At the end of each inference, `flush()` will reset all command buffers in the pool, allowing them to be re-used.
## `CommandBuffer`
This class will represent a command buffer and will provide function calls necessary to encode commands into a command buffer.
This class has an internal state that tracks what stage of recording it's in. Note that this is separate from the State of the `VkCommandBuffer` that it owns, which is defined by the Vulkan API.
```
enum State {
INVALID, // Used to indicate the command buffer is moved from
NEW, // Set during constructor
RECORDING, // Set during call to begin(), dispatch(), and copy_texture_to_texture()
PIPELINE_BOUND, // Set during call to bind_pipeline()
DESCRIPTORS_BOUND, // Set during call to bind_descriptors()
BARRIERS_INSERTED, // Set during call to insert_barrier()
READY, // Set during call to end()
SUBMITTED, // Set during call to get_submit_handle()
};
```
The usage of a `CommandBuffer` object can be described below.
```
// Fetch a new command buffer from the command pool
CommandBuffer cmd = command_pool.get_new_cmd()
// The state is now NEW
// This calls vkBeginCommandBuffer
cmd.begin()
// The state is now RECORDING
// This calls vkCmdBindPipeline to bind a compute pipeline to the command buffer
cmd.bind_pipeline(...)
// The state is now PIPELINE_BOUND
// This calls vkCmdBindDescriptors to bind a descriptor set to the command buffer
cmd.bind_descriptors(...)
// The state is now DESCRIPTORS_BOUND
// This calls vkCmdPipelineBarrier and inserts a pipeline barrier in the command buffer
cmd.insert_barrier(...)
// The state is now BARRERS_INSERTED
// This calls vkCmdDispatch and records a compute dispatch job into the command buffer
cmd.dispatch(...)
// The state is now RECORDING
// From here, we can repeat the process from bind_pipeline() to record a new compute dispatch into the command buffer.
...
// When we are ready to submit the command buffer, this calls vkEndCommandBuffer
cmd.end()
// The state is now READY
// Relinquish the VkCommandBuffer cmd owns and returns it to be used in a vkQueueSubmit
VkCommandBuffer handle = cmd.get_submit_handle()
// The state is now SUBMITTED
```
Differential Revision: [D36824003](https://our.internmc.facebook.com/intern/diff/D36824003/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/80720
Approved by: https://github.com/kimishpatel
