From 45e3266d0417bb82de621d28b634765eff72d270 Mon Sep 17 00:00:00 2001
From: Ruihang Lai <ruihangl@cs.cmu.edu>
Date: Sun, 7 May 2023 02:59:02 -0400
Subject: [PATCH] [CMAKE] Add Vulkan header for Android

This PR enhances the Vulkan search for Android NDK, so that we can
find the Vulkan header for Android builds.

Co-authored-by: tqchen <tqchenml@gmail.com>
---
 cmake/utils/FindVulkan.cmake | 57 ++++++++++++++++++++----------------
 1 file changed, 32 insertions(+), 25 deletions(-)

diff --git a/cmake/utils/FindVulkan.cmake b/cmake/utils/FindVulkan.cmake
index a2a6fb60352d..032ff1dffa21 100644
--- a/cmake/utils/FindVulkan.cmake
+++ b/cmake/utils/FindVulkan.cmake
@@ -37,36 +37,43 @@ macro(find_vulkan use_vulkan use_khronos_spirv)
   if(IS_DIRECTORY ${__use_vulkan})
     set(__vulkan_sdk ${__use_vulkan})
     message(STATUS "Custom Vulkan SDK PATH=" ${__use_vulkan})
-   elseif(IS_DIRECTORY $ENV{VULKAN_SDK})
-     set(__vulkan_sdk $ENV{VULKAN_SDK})
-   else()
-     set(__vulkan_sdk "")
-   endif()
+  elseif(IS_DIRECTORY $ENV{VULKAN_SDK})
+    set(__vulkan_sdk $ENV{VULKAN_SDK})
+  else()
+    set(__vulkan_sdk "")
+  endif()
 
 
-   if(IS_DIRECTORY ${use_khronos_spirv})
-     set(__use_khronos_spirv ${use_khronos_spirv})
-     message(STATUS "Custom khronos spirv PATH=" ${__use_khronos_spirv})
-   else()
-     set(__use_khronos_spirv "")
-   endif()
+  if(IS_DIRECTORY ${use_khronos_spirv})
+    set(__use_khronos_spirv ${use_khronos_spirv})
+    message(STATUS "Custom khronos spirv PATH=" ${__use_khronos_spirv})
+  else()
+    set(__use_khronos_spirv "")
+  endif()
 
+  if(CMAKE_SYSTEM_NAME STREQUAL "Android")
+    set(VULKAN_NDK_SRC ${CMAKE_ANDROID_NDK}/sources/third_party/vulkan/src)
+    set(Vulkan_INCLUDE_DIRS ${VULKAN_NDK_SRC}/include)
+    set(Vulkan_FOUND TRUE)
+    message(STATUS "Android Vulkan_INCLUDE_DIRS=" ${Vulkan_INCLUDE_DIRS})
+    message(STATUS "Skip finding SPIRV in Android, make sure you only build tvm runtime.")
+    return()
+  endif()
 
-   if(__vulkan_sdk)
-     set(Vulkan_INCLUDE_DIRS ${__vulkan_sdk}/include)
-     find_library(Vulkan_LIBRARY NAMES vulkan vulkan-1 PATHS ${__vulkan_sdk}/lib)
-     if(Vulkan_LIBRARY)
-       set(Vulkan_FOUND TRUE)
-     endif()
-   endif(__vulkan_sdk)
+  if(__vulkan_sdk)
+    set(Vulkan_INCLUDE_DIRS ${__vulkan_sdk}/include)
+    find_library(Vulkan_LIBRARY NAMES vulkan vulkan-1 PATHS ${__vulkan_sdk}/lib)
+    if(Vulkan_LIBRARY)
+      set(Vulkan_FOUND TRUE)
+    endif()
+  endif(__vulkan_sdk)
 
-   # resort to find vulkan of option is on
-   if(NOT Vulkan_FOUND)
-     if(${__use_vulkan} MATCHES ${IS_TRUE_PATTERN})
-       find_package(Vulkan QUIET)
-     endif()
-   endif()
-   # additional libraries
+  # resort to find vulkan of option is on
+  if(NOT Vulkan_FOUND)
+    if(${__use_vulkan} MATCHES ${IS_TRUE_PATTERN})
+      find_package(Vulkan QUIET)
+    endif()
+  endif()
 
   if(Vulkan_FOUND)
     get_filename_component(VULKAN_LIBRARY_PATH ${Vulkan_LIBRARY} DIRECTORY)