pytorch

caa14951 - [vulkan] Refactor Command (#80720)

[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