onnxruntime

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

4
find_package(Python3 COMPONENTS Interpreter REQUIRED)
5

6
# set all triton kernel ops that need to be compiled
7
if(onnxruntime_USE_ROCM)
8
  set(triton_kernel_scripts
9
      "onnxruntime/core/providers/rocm/math/softmax_triton.py"
10
      "onnxruntime/contrib_ops/rocm/diffusion/group_norm_triton.py"
11
  )
12
endif()
13

14
function(compile_triton_kernel out_triton_kernel_obj_file out_triton_kernel_header_dir)
15
  # compile triton kernel, generate .a and .h files
16
  set(triton_kernel_compiler "${REPO_ROOT}/tools/ci_build/compile_triton.py")
17
  set(out_dir "${CMAKE_CURRENT_BINARY_DIR}/triton_kernels")
18
  set(out_obj_file "${out_dir}/triton_kernel_infos.a")
19
  set(header_file "${out_dir}/triton_kernel_infos.h")
20

21
  list(TRANSFORM triton_kernel_scripts PREPEND "${REPO_ROOT}/")
22

23
  add_custom_command(
24
    OUTPUT ${out_obj_file} ${header_file}
25
    COMMAND Python3::Interpreter ${triton_kernel_compiler}
26
            --header ${header_file}
27
            --script_files ${triton_kernel_scripts}
28
            --obj_file ${out_obj_file}
29
    DEPENDS ${triton_kernel_scripts} ${triton_kernel_compiler}
30
    COMMENT "Triton compile generates: ${out_obj_file}"
31
  )
32
  add_custom_target(onnxruntime_triton_kernel DEPENDS ${out_obj_file} ${header_file})
33
  set(${out_triton_kernel_obj_file} ${out_obj_file} PARENT_SCOPE)
34
  set(${out_triton_kernel_header_dir} ${out_dir} PARENT_SCOPE)
35
endfunction()
36