[Vulkan] Vulkan runtime reimplementation (stream approach)

[ajtulloch]

In the Vulkan API, it is recommended to launch multiple compute kernels on a shared command buffer in general, instead of having a 1-1 relationship between compute kernels and buffers. See http://on-demand.gputechconf.com/gtc/2016/events/vulkanday/High_Performance_Vulkan.pdf, https://devblogs.nvidia.com/vulkan-dos-donts/, etc. The extant TVM Vulkan runtime uses one command buffer per kernel, which can lead to significant overheads for smaller-kernels (on the order of half a millisecond on some of the devices I looked at).

An alternative approach leverages an approach similar to a CUDA stream abstraction, where we record commands onto the command buffer, and at synchronization points, submit the command buffer to the queue and wait on the fence. This is non-trivially more efficient - similar to the approach taken by ncnn - there are some useful ideas in there that applied here. In particular it's quite convenient to depend on the KHR push descriptors extension, but that could be removed without too much pain similar to how ncnn does it.

This code isn't production ready, and it's not super clear how much interest there is in the Vulkan side of things. I think it's quite promising and was planning on spending some time looking at codegen stuff, but the difficulty in getting reasonable numbers for small B/W bound kernels was the motivator in working on this to begin with.

If there's interest we could probably figure out a way to merge this into the existing Vulkan runtime, perhaps gated by a feature flag?

Performance improves for simple pointwise kernels as expected, using a script like:

import tvm
import numpy as np

tx = tvm.thread_axis("threadIdx.x")
bx = tvm.thread_axis("blockIdx.x")

num_thread = 256
from tvm import rpc
tracker = rpc.connect_tracker('localhost', 9090)
remote = tracker.request("android", priority=1,
                         session_timeout=6000)
ctx = remote.vulkan(0)
def check_vulkan(dtype, n):
    A = tvm.placeholder((n,), name='A', dtype=dtype)
    B = tvm.compute((n,), lambda i: A[i]+tvm.const(1, A.dtype), name='B')
    s = tvm.create_schedule(B.op)
    xo, xi = s[B].split(B.op.axis[0], factor=num_thread * 4)
    s[B].bind(xo, bx)
    xi, vx = s[B].split(xi, factor=4)
    s[B].bind(xi, tx)
    s[B].vectorize(vx)
    f = tvm.build(
        s, [A, B],
        target="vulkan",
        target_host="llvm -target=arm64-linux-android")
    import os
    os.environ['TVM_NDK_CC'] = os.path.expanduser("~/opt/android-toolchain-arm64/bin/aarch64-linux-android-g++")
    fname = f"dev_lib_vulkan_{np.random.random()}.so"
    path_dso_vulkan = fname
    from tvm.contrib import ndk

    f.export_library(path_dso_vulkan, ndk.create_shared)
    ctx = remote.vulkan(0)
    remote.upload(path_dso_vulkan)
    f1 = remote.load_module(fname)
    a_np = np.random.uniform(size=(n,)).astype(dtype)
    a = tvm.nd.array(a_np, ctx)
    c = tvm.nd.empty((n,), B.dtype, ctx)
    f1(a, c)
    tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + 1)
    te = f1.time_evaluator(f.entry_name, ctx=ctx, min_repeat_ms=500, number=5)
    for _ in range(3):
        perf = te(a, c).mean
        print(f"N: {n}, t: {perf * 1.0e6:.2f}us, GFLOP/s: {n / perf / 1.0e9}")

for log_n in range(10, 20):
    check_vulkan("float32", 2 ** log_n)

[ajtulloch]

cc @jwfromm, @tqchen

[tqchen]

Awesome, this is something that i overlooked and definitely should be the way to go. The main reason that I took the current approach was because there is no clear way to do synchronization other than the fence, and the stream model is not supported by vulkan natively.

The proposed approach(to add stream support via lazy execution) is definitely better and we should just move toward this new runtime IMO. As long as we implement all the features, please feel free to remove the old vulkan runtime.

[ajtulloch]

@nihui would you be interested in taking a look at this PR and this approach? Are there more useful things in NCNN’s Vulkan backend that TVM should learn from?

[nihui]

I also found that it took a lot of time to wait for each operator's result, so I put multiple operators in a single command buffer and apply the lazy execution.
But I'm not familiar with TVM, and the code in the PR seems to be very different from ncnn, so I cannot judge if this PR is the best, but it must be much better than the current TVM approach I think.

If you have any questions about the ncnn vulkan backend, feel free to post an issue to https://github.com/Tencent/ncnn/issues

[ajtulloch]

Thanks @nihui, @tqchen, @jwfromm.

I generalized this approach to handle devices that don't support the push descriptor (see the Stream::LaunchDeferred APIs), and devices that don't support dedicated allocation APIs, and removed the previous runtime implementation as @tqchen suggested. I also added some more tests which (under Vulkan API validation) revealed some thread safety bugs.

I think this should be closer to something that is acceptable, so I'll remove the RFC and WIP tags, this should be ready for review.

File/class structure is certainly up for debate. I'm fine with whatever folks propose.

[tqchen]

@jwfromm @yzhliu @masahi @kazum can you guys help to review this PR?

[jwfromm]

This is really excellent work, I'm super impressed and excited to get better Vulkan performance merged in to TVM. I like that you've now replaced all the old Vulkan runtime code, however, now that it's gone I think you should remove all the 2s from your file, class, and function naming. I've also added a few minor comments, however everything else LGTM.

[jwfromm]

Why do we need to hard-code the version like this? Is it possible to use spv::Version both here and in vulkcan2.cc and check that its >= v1.0?

[ajtulloch]

My understanding is that the semantically correct thing to do is:

• at code-generation time, output the lowest possible version number that supports all the features we require. Currently this is 1.0, we don't use features from 1.1 or higher.
• in vkApplicationInfo.apiVersion, use the minimum required API version we require for the devices. This is also what NCNN uses, see https://github.com/Tencent/ncnn/blob/c65c215c2c1084303849f9adae2f7495a1b99409/src/gpu.cpp#L407.

[ajtulloch]

Thanks @jwfromm, will follow up on your suggestions.

[ajtulloch]

@jwfromm, I believe I've addressed your comments (except the one re: the correct version semantics at SPIR-V emission time and Vulkan instance instantiation time).

[tqchen]

Thanks @ajtulloch @jwfromm @nihui !