(tvm) branch unity updated: [Unity] Update FlashInfer (#16208)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/unity by this push:
     new d0504027bb [Unity] Update FlashInfer (#16208)
d0504027bb is described below

commit d0504027bb72c877f3b5e482e5dd45f5caa5f5ca
Author: Yaxing Cai <caiyaxing...@gmail.com>
AuthorDate: Thu Dec 7 07:06:59 2023 -0800

    [Unity] Update FlashInfer (#16208)
    
    This PR updates the FlashInfer submodule for latest changes.
---
 3rdparty/flashinfer                 | 2 +-
 src/runtime/cuda/cuda_device_api.cc | 8 ++++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/3rdparty/flashinfer b/3rdparty/flashinfer
index 4aa8130e05..11364ca4c3 160000
--- a/3rdparty/flashinfer
+++ b/3rdparty/flashinfer
@@ -1 +1 @@
-Subproject commit 4aa8130e05a2cd9e1f4077979535de8c1c8c71cd
+Subproject commit 11364ca4c3ce651dd544efff3225906fe15c5b8a
diff --git a/src/runtime/cuda/cuda_device_api.cc 
b/src/runtime/cuda/cuda_device_api.cc
index 5cf6cd7342..f01fd60476 100644
--- a/src/runtime/cuda/cuda_device_api.cc
+++ b/src/runtime/cuda/cuda_device_api.cc
@@ -310,5 +310,13 @@ 
TVM_REGISTER_GLOBAL("runtime.get_cuda_stream").set_body_typed([]() {
   return static_cast<void*>(CUDAThreadEntry::ThreadLocal()->stream);
 });
 
+TVM_DLL int GetCudaDeviceCount() {
+  int count;
+  CUDA_CALL(cudaGetDeviceCount(&count));
+  return count;
+}
+
+TVM_REGISTER_GLOBAL("runtime.GetCudaDeviceCount").set_body_typed(GetCudaDeviceCount);
+
 }  // namespace runtime
 }  // namespace tvm