onnxruntime

1
# Copyright (c) Microsoft Corporation. All rights reserved.
2
# Licensed under the MIT License.
3

4
include(CheckLanguage)
5

6
if(NOT onnxruntime_ENABLE_PYTHON)
7
  message(FATAL_ERROR "python is required but is not enabled")
8
endif()
9

10
set(KERNEL_EXPLORER_ROOT ${ONNXRUNTIME_ROOT}/python/tools/kernel_explorer)
11

12
if (onnxruntime_USE_CUDA)
13
  check_language(CUDA)
14
  set(LANGUAGE CUDA)
15
  set(BERT_DIR ${ONNXRUNTIME_ROOT}/contrib_ops/cuda/bert)
16
elseif(onnxruntime_USE_ROCM)
17
  check_language(HIP)
18
  set(LANGUAGE HIP)
19
  if (onnxruntime_USE_COMPOSABLE_KERNEL)
20
    include(composable_kernel)
21
  endif()
22
  if (onnxruntime_USE_HIPBLASLT)
23
    find_package(hipblaslt REQUIRED)
24
  endif()
25
  set(BERT_DIR ${ONNXRUNTIME_ROOT}/contrib_ops/rocm/bert)
26
endif()
27

28
file(GLOB kernel_explorer_srcs CONFIGURE_DEPENDS
29
  "${KERNEL_EXPLORER_ROOT}/*.cc"
30
  "${KERNEL_EXPLORER_ROOT}/*.h"
31
)
32

33
file(GLOB kernel_explorer_kernel_srcs CONFIGURE_DEPENDS
34
  "${KERNEL_EXPLORER_ROOT}/kernels/*.cc"
35
  "${KERNEL_EXPLORER_ROOT}/kernels/*.h"
36
  "${KERNEL_EXPLORER_ROOT}/kernels/*.cu"
37
  "${KERNEL_EXPLORER_ROOT}/kernels/*.cuh"
38
)
39

40
onnxruntime_add_shared_library_module(kernel_explorer ${kernel_explorer_srcs} ${kernel_explorer_kernel_srcs})
41
set_target_properties(kernel_explorer PROPERTIES PREFIX "_")
42
target_include_directories(kernel_explorer PUBLIC
43
  $<TARGET_PROPERTY:onnxruntime_pybind11_state,INCLUDE_DIRECTORIES>
44
  ${KERNEL_EXPLORER_ROOT})
45
target_link_libraries(kernel_explorer PRIVATE $<TARGET_PROPERTY:onnxruntime_pybind11_state,LINK_LIBRARIES>)
46
target_compile_definitions(kernel_explorer PRIVATE $<TARGET_PROPERTY:onnxruntime_pybind11_state,COMPILE_DEFINITIONS>)
47
target_compile_options(kernel_explorer PRIVATE -Wno-sign-compare)
48

49
if (onnxruntime_USE_CUDA)
50
  file(GLOB kernel_explorer_cuda_kernel_srcs CONFIGURE_DEPENDS
51
    "${KERNEL_EXPLORER_ROOT}/kernels/cuda/*.cc"
52
    "${KERNEL_EXPLORER_ROOT}/kernels/cuda/*.h"
53
    "${KERNEL_EXPLORER_ROOT}/kernels/cuda/*.cu"
54
    "${KERNEL_EXPLORER_ROOT}/kernels/cuda/*.cuh"
55
  )
56
  target_sources(kernel_explorer PRIVATE ${kernel_explorer_cuda_kernel_srcs})
57
  target_include_directories(kernel_explorer PUBLIC ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
58
elseif (onnxruntime_USE_ROCM)
59
  file(GLOB kernel_explorer_rocm_kernel_srcs CONFIGURE_DEPENDS
60
    "${KERNEL_EXPLORER_ROOT}/kernels/rocm/*.cc"
61
    "${KERNEL_EXPLORER_ROOT}/kernels/rocm/*.h"
62
    "${KERNEL_EXPLORER_ROOT}/kernels/rocm/*.cu"
63
    "${KERNEL_EXPLORER_ROOT}/kernels/rocm/*.cuh"
64
  )
65
  auto_set_source_files_hip_language(${kernel_explorer_kernel_srcs} ${kernel_explorer_rocm_kernel_srcs})
66
  target_sources(kernel_explorer PRIVATE ${kernel_explorer_rocm_kernel_srcs})
67
  target_compile_definitions(kernel_explorer PRIVATE __HIP_PLATFORM_AMD__=1 __HIP_PLATFORM_HCC__=1)
68
  if (onnxruntime_USE_COMPOSABLE_KERNEL)
69
    target_compile_definitions(kernel_explorer PRIVATE USE_COMPOSABLE_KERNEL)
70
    if (onnxruntime_USE_COMPOSABLE_KERNEL_CK_TILE)
71
      target_compile_definitions(kernel_explorer PRIVATE USE_COMPOSABLE_KERNEL_CK_TILE)
72
    endif()
73
    target_link_libraries(kernel_explorer PRIVATE onnxruntime_composable_kernel_includes)
74
  endif()
75
  if (onnxruntime_USE_TRITON_KERNEL)
76
    target_compile_definitions(kernel_explorer PRIVATE USE_TRITON_KERNEL)
77
  endif()
78
  if (onnxruntime_USE_HIPBLASLT)
79
    target_compile_definitions(kernel_explorer PRIVATE USE_HIPBLASLT)
80
  endif()
81
  if (onnxruntime_USE_ROCBLAS_EXTENSION_API)
82
    target_compile_definitions(kernel_explorer PRIVATE USE_ROCBLAS_EXTENSION_API)
83
    target_compile_definitions(kernel_explorer PRIVATE ROCBLAS_NO_DEPRECATED_WARNINGS)
84
    target_compile_definitions(kernel_explorer PRIVATE ROCBLAS_BETA_FEATURES_API)
85
  endif()
86
endif()
87

88
add_dependencies(kernel_explorer onnxruntime_pybind11_state)
89

90
enable_testing()
91
find_package(Python COMPONENTS Interpreter REQUIRED)
92
# add_test(NAME test_kernels COMMAND ${Python_EXECUTABLE} -m pytest ..)
93