[Vulkan] Support uniform buffer object for passing many scalar arguments

[masahi]

We are using push constants to pass scalar arguments to SPIR-V kernels. The vulkan spec ensures that the size of push constants storage is at least 128 bytes. Since we pass each scalar via 64 bit union, that means we can only pass 16 parameters via push constants. Unfortunately, for fused, dynamic input kernel, the number of any_dim and other params can easily go beyond that limit. See for example this crazy kernel that needs 20 scalars to be passed in https://gist.github.com/masahi/ce51c0d51c6115109203b3732f185aab

This PR enables passing parameters beyond push constants limit, using uniform buffer object. In particular, with this PR I was able to run PyTorch MaskRCNN end to end on Vulkan!

Running NMS tests in onnx/test_forward.py serves as a test case.

There are some minor TODOs left, but I'd like to get some feedback. Supporting UBO requires making non trivial changes to both codegen and runtime, but hopefully they are straightforward.

@tqchen @tmoreau89 @jwfromm @ajtulloch

TODO

• Query platform specific push constant size limit rather than using a hard coded constant
• Fix segfault on deleting UBO host buffer (Seems I shouldn't allocate/deallocate the host buf explicitly, rather use vkMapMemory/vkUnMapMemory)
• Add doc comments in ir_builder.h

[masahi]

@tqchen This is ready for review.

On the codegen side, UBO is declared in a similar way as push constants, so I made DeclareStorageVariable function to declare both storage type.

On the runtime side, I introduced CreateBuffer function to declare both device storage buffers and uniform buffers. The change should be straightforward.

One potentially controversial point could be the use of runtime function MaxPushConstantsSize() during codegen, to query the HW param and decide if we should use push constants or UBO.

[tqchen]

cc @ajtulloch @antinucleon please help to see if you can take a look

[tmoreau89]

Thank you @masahi for these changes, I left a comment on cosmetics

[tmoreau89]

Thanks you @masahi your PR has been merged

[tqchen]

Sorry that i didn't get a chance to review. Here are some items that needs to be fixed as followup PRs:

• Let us not use GetMaxPushConstantSize, as it is dynamic across devices, and the code won't behave correctly if we are doing cross compilation(where we compile on a machine that does not have vulkan), we should use a minimum TVM default (e.g. 128 bytes).
• I see a few places where delete is used, always prefer std::unique_ptr when possible instead of delete

[masahi]

@tqchen

• I added one delete but I just followed the existing use of delete in

  tvm/src/runtime/vulkan/vulkan.cc

  Lines 239 to 241 in 5bc1cec

  	vkDestroyBuffer(vctx.device, pbuf->buffer, nullptr);
  	vkFreeMemory(vctx.device, pbuf->memory, nullptr);
  	delete pbuf;

  Since we need to manually delete or free associated vk related resources anyway, I don't mind this delete too much

• Is there a way to make GetMaxPushConstantSize work for cross compilation use cases? Because for vulkan in particular, we definitely need to query HW related information like this at compile time. For example, the max thread block sizes, warp size, what extentions are available (subgroup ops extensions correspond to warp shuffle ops in CUDA but not all vk devices support them) etc. The situation with warp size is the worst because we don't have warp_size specified for vulkan target https://github.com/apache/tvm/blob/main/src/target/target_kind.cc#L273-L275 but some topi ops uses thread_warp_size

  tvm/python/tvm/topi/cuda/scatter.py

  Line 196 in 8131364

  warp_size = tvm.target.Target.current(False).thread_warp_size

  so right now we cannot compile these ops on vulkan.

I'm ok if we decide to use a fixed max push constant size, but for other HW params we need to have a solution for runtime querying.

[tqchen]

Thanks @masahi. I agree that the use of delete is minor.

In the case of push constant, we can try to obtain these constant and put them into target. The main problem is the code as it is we are querying on the local machine but we need parameters on the remote. It would be useful to consider de-coupling to avoid such inconsistency(of information obtained by compiler and information in the runtime). e.g. The kernel should declare whether UBO or push constant is being used in its function meta-data, so we do not need to make dependency on the runtime parameters (e.g. the compiler decides whether to use UBO by arg size, and write that to an optional meta data of the shader to enable UBO). e.g. we can add a bitmask flag to VulkanShader::flag to indicate the difference

The runtime simply read that information and run dispatch, this will avoid inconsistency during compile time and runtime. We can then enhance the target registration to include such informations of max push constant per target, while defaults to the minimum value.

After taking a closer look, there are a few more potential things that needs to be fixed:

• Vulkan runtime have deferred mode and immediate mode that needs to be executed in separate ways(see the two places using push constants), we will need to do the same thing for UBO logic, right now seems that the code will work for immediate mode but not deferred mode
• Right now the UBO argument buffer simply re-allocates without deleting the original buffer, this could cause memory leaks. We will need to have a logic that re-allocates if necessary(also think about relation to synchronization in deferred case) and use the buffer if the size is large enough.

[masahi]

I confirmed that it also works for deferred mode (by always returning false from UseImmediate()). The relevent changes are

• https://github.com/masahi/tvm/blob/1a3dbee99c9a2c362373707678d5657e59ea6827/src/runtime/vulkan/vulkan.cc#L1200-L1204
• https://github.com/masahi/tvm/blob/1a3dbee99c9a2c362373707678d5657e59ea6827/src/runtime/vulkan/vulkan.cc#L1215

Need to look at the reallocate problem

[tqchen]

Sorry I was wrong about the re-allocation, seems you are allocating a single UBO per-kernel, so it might be fine. My original thought was along the line of single UBO per runtime.

[masahi]

Yes, one UBO per one pipeline https://github.com/masahi/tvm/blob/1a3dbee99c9a2c362373707678d5657e59ea6827/src/runtime/vulkan/vulkan.cc#L108

[tqchen]

To followup on the deffered execution case:

Because memcpy happens here: https://github.com/masahi/tvm/blob/1a3dbee99c9a2c362373707678d5657e59ea6827/src/runtime/vulkan/vulkan.cc#L1150

Imagine there are two consecutive launch of the same kernel A(uses UBO) that are deferred. Then the memcpy of the second kernel launch would override the value of the first launch(both of which hasn't yet been launched). So we want to move memcpy to the lambda in the immediate and deferred kernels(besides the push constant calls)

[tqchen]

To followup on the per pipeline UBO, after thinking a bit more about this design. Per pipeline UBO is fine for single threaded case, but can also be problematic under a multi-threaded setting when multiple threads are launching the same kernel A.

To make the runtime thread-safe, we normally needs to divide the data structure into constants(e.g. pipeline) and runtime structure(e.g. staging buffer, streams). The runtime structure part belongs to VulkanThreadEntry that comes with a thread-local copy to avoid threading issue. So we would want to do that, and uses a similar logic as StagingBuffer to create staging buffer for UBO.

Finally, right now seems we relies on being able to have an unified memory(host mapped memory) for UBO. Will all GPU support such kind of memory? Shall we allow explicit data copy into the UBO buffer when such memory is not supported?

[tqchen]

Thanks @masahi for great discussions. It seems to me that we would have benefited more reviews and discussion :) Great that we have catched these points for improvements during the discussion after merging. Given it is already merged, we could consider:

• A0: Revert the PR, and bring another one, do another round of thorough reviews there.
• A1: Bring a fix PR as soon as possible.

To summarize the items that needs to be addressed:

• P0: Add the use of UBO as meta-data to shader, so we won't have compiler runtime in-consistentcy problem
• P1: Solve the deferred execution case when a UBO buffer can be overriden by the followup kernel calls(need to do the copy in deferred stream as well).
• P2: Solve the multi-thread safety issue when multiple threads calls into UBO at the same time.

[masahi]

ok sent the revert #7821

[tmoreau89]

@masahi @tqchen Thanks for the comments, it appears as though the merge might have been to hasty. I'm good with the decision to go with A0; thanks @masahi for issuing the revert.