diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 5551628a0d..e4d1a85b05 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -222,6 +222,7 @@ option(onnxruntime_ENABLE_ATEN "Enable ATen fallback" OFF) # composable kernel is managed automatically, unless user want to explicitly disable it, it should not be manually set option(onnxruntime_USE_COMPOSABLE_KERNEL "Enable composable kernel for ROCm EP" ON) option(onnxruntime_USE_ROCBLAS_EXTENSION_API "Enable rocblas tuning for ROCm EP" OFF) +option(onnxruntime_USE_TRITON_KERNEL "Enable triton compiled kernel" OFF) option(onnxruntime_BUILD_KERNEL_EXPLORER "Build Kernel Explorer for testing and profiling GPU kernels" OFF) option(onnxruntime_BUILD_CACHE "onnxruntime build with cache" OFF) diff --git a/cmake/onnxruntime_compile_triton_kernel.cmake b/cmake/onnxruntime_compile_triton_kernel.cmake new file mode 100644 index 0000000000..b9482ecaac --- /dev/null +++ b/cmake/onnxruntime_compile_triton_kernel.cmake @@ -0,0 +1,32 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +find_package(Python3 COMPONENTS Interpreter REQUIRED) + +# set all triton kernel ops that need to be compiled +set(triton_kernel_scripts + "onnxruntime/core/providers/rocm/math/softmax.py" +) + +function(compile_triton_kernel out_triton_kernel_obj_file out_triton_kernel_header_dir) + # compile triton kernel, generate .a and .h files + set(triton_kernel_compiler "${REPO_ROOT}/tools/ci_build/compile_triton.py") + set(out_dir "${CMAKE_CURRENT_BINARY_DIR}/triton_kernels") + set(out_obj_file "${out_dir}/triton_kernel_infos.a") + set(header_file "${out_dir}/triton_kernel_infos.h") + + list(TRANSFORM triton_kernel_scripts PREPEND "${REPO_ROOT}/") + + add_custom_command( + OUTPUT ${out_obj_file} ${header_file} + COMMAND Python3::Interpreter ${triton_kernel_compiler} + --header ${header_file} + --script_files ${triton_kernel_scripts} + --obj_file ${out_obj_file} + DEPENDS ${triton_kernel_scripts} ${triton_kernel_compiler} + COMMENT "Triton compile generates: ${out_obj_file}" + ) + add_custom_target(onnxruntime_triton_kernel DEPENDS ${out_obj_file} ${header_file}) + set(${out_triton_kernel_obj_file} ${out_obj_file} PARENT_SCOPE) + set(${out_triton_kernel_header_dir} ${out_dir} PARENT_SCOPE) +endfunction() diff --git a/cmake/onnxruntime_kernel_explorer.cmake b/cmake/onnxruntime_kernel_explorer.cmake index d4ae88a1f6..e05321748e 100644 --- a/cmake/onnxruntime_kernel_explorer.cmake +++ b/cmake/onnxruntime_kernel_explorer.cmake @@ -66,6 +66,9 @@ elseif (onnxruntime_USE_ROCM) target_compile_definitions(kernel_explorer PRIVATE USE_COMPOSABLE_KERNEL) target_link_libraries(kernel_explorer PRIVATE onnxruntime_composable_kernel_includes) endif() + if (onnxruntime_USE_TRITON_KERNEL) + target_compile_definitions(kernel_explorer PRIVATE USE_TRITON_KERNEL) + endif() endif() add_dependencies(kernel_explorer onnxruntime_pybind11_state) diff --git a/cmake/onnxruntime_providers.cmake b/cmake/onnxruntime_providers.cmake index e561088ac9..0df081bbde 100644 --- a/cmake/onnxruntime_providers.cmake +++ b/cmake/onnxruntime_providers.cmake @@ -500,6 +500,18 @@ if (onnxruntime_USE_CUDA) target_link_directories(onnxruntime_providers_cuda PRIVATE ${onnxruntime_CUDNN_HOME}/lib) endif() + if (onnxruntime_USE_TRITON_KERNEL) + # compile triton kernel, generate .a and .h files + include(onnxruntime_compile_triton_kernel.cmake) + compile_triton_kernel(triton_kernel_obj_file triton_kernel_header_dir) + add_dependencies(onnxruntime_providers_cuda onnxruntime_triton_kernel) + target_compile_definitions(onnxruntime_providers_cuda PRIVATE USE_TRITON_KERNEL) + target_include_directories(onnxruntime_providers_cuda PRIVATE ${triton_kernel_header_dir}) + target_link_libraries(onnxruntime_providers_cuda PUBLIC -Wl,--whole-archive ${triton_kernel_obj_file} -Wl,--no-whole-archive) + # lib cuda needed by cuLaunchKernel + target_link_libraries(onnxruntime_providers_cuda PRIVATE cuda) + endif() + if (onnxruntime_USE_FLASH_ATTENTION) include(cutlass) target_include_directories(onnxruntime_providers_cuda PRIVATE ${cutlass_SOURCE_DIR}/include ${cutlass_SOURCE_DIR}/examples) @@ -1619,6 +1631,16 @@ if (onnxruntime_USE_ROCM) target_compile_definitions(onnxruntime_providers_rocm PRIVATE USE_HIPBLASLT) endif() + if (onnxruntime_USE_TRITON_KERNEL) + # compile triton kernel, generate .a and .h files + include(onnxruntime_compile_triton_kernel.cmake) + compile_triton_kernel(triton_kernel_obj_file triton_kernel_header_dir) + add_dependencies(onnxruntime_providers_rocm onnxruntime_triton_kernel) + target_compile_definitions(onnxruntime_providers_rocm PRIVATE USE_TRITON_KERNEL) + target_include_directories(onnxruntime_providers_rocm PRIVATE ${triton_kernel_header_dir}) + target_link_libraries(onnxruntime_providers_rocm PUBLIC -Wl,--whole-archive ${triton_kernel_obj_file} -Wl,--no-whole-archive) + endif() + if (onnxruntime_USE_COMPOSABLE_KERNEL) include(composable_kernel) target_link_libraries(onnxruntime_providers_rocm PRIVATE diff --git a/docs/ORT_use_trtion_kernel.md b/docs/ORT_use_trtion_kernel.md new file mode 100644 index 0000000000..7c9aff4bb5 --- /dev/null +++ b/docs/ORT_use_trtion_kernel.md @@ -0,0 +1,98 @@ +## Description +In some scenarios, the triton written kernels are more performant than CK or other handwritten kernels, so we implement a framework that onnxruntime can use these triton written kernels. + +Here we use `softmax` op as an example to show how to integrate a triton written kernel into onnxruntime cuda/rocm EP. + +#### Write and compile triton kernel +We have implemented a softmax kernel using triton at `onnxruntime/core/providers/rocm/math/softmax.py` + +``` +@triton.jit +def softmax_kernel( + output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols, + BLOCK_SIZE: tl.constexpr +): + # softmax implementings + .... + .... +``` +This is a very simple implementation. The `n_cols` parameter should be smaller than BLOCK_SIZE. And BLOCK_SIZE MUST be determined at compile time. + +In order to support different input shape, we compile multiple kernels with different BLOCK_SIZE. + +Each kernel with different BLOCK_SIZE generates different num_warps and shared memory usage, we call them as `metadata`, and these metadata are needed when launching kernels in onnxruntime. + +We develop a script `tools/ci_build/compile_triton.py` to compile kernel and generate metadata for kernel launching. + +To generate metadta for softmax, it needs to add description info and implement a `get_function_table` function in `softmax.py`: +``` +# kernel dtype and BLOCK_SIZE to generate. +dtypes = ['fp32', 'fp16'] +blocks = [1024, 2048, 4096, 8192, 16384] +name_pattern = 'softmax_{}_{}' +sig_pattern = '*{},*{},i32,i32,i32' +group_pattern = 'softmax_{}' + +""" +SHOULD implement a function that returns a metadata list with format: + +function_table = [ + {'name': xx, + 'group': yy, + 'func': func, + 'sig': sig, + 'kwargs': kwargs + } +] + +The kwargs is a dict of {string: int} which is used for kernel constants. For example, BLOCK_SIZE of softmax. +""" + +def get_funcion_table(): + ...... + ...... +``` + +When compiling onnxruntime with `--use_triton_kernel` flag, this softmax kernel will be compiled and combined into libonnxruntime_providers_rocm.so for rocm or libonnxruntime_providers_cuda.so for cuda. + +### onnxruntime c++ code moidfication +To use the triton kernels in onnxruntime, we need to implement a c++ op that calls these triton kernels. + +And there are many implements for softmax, we implement a tunable op for softmax to choose the best performant one. + +Same as CK, we implement a function that returns all possible triton kernels, and the tunable op will select best one. + +``` +template +auto GetSoftmaxTritonOps() { + std::vector>>> ret; + auto group_name = GetSoftmaxTritonGroupName(); + // here use group_name to get all kernel with same group_name + // for example, 'softmax_fp16' represents a group of kernels with different BLOCK_SIZE for float16 softmax + auto *kernel_list = GetOrtTritonKernelByGroup(group_name); + if (kernel_list == nullptr) { + return ret; + } + + for (auto i : *kernel_list) { + // check params match + ..... + } + return ret; +} +``` + +### Test +Using kernel_explorer, we can test this softmax kernel. +``` +export KERNEL_EXPLORER_BUILD_DIR=/ws/code/onnxruntime/build_rocm/Debug + +python onnxruntime/python/tools/kernel_explorer/kernels/softmax_test.py +``` +and can see the results, here TunableOp select softmax_fp16_2048 which is a triton written kernel, it's better than others. +``` +SoftmaxTunable float16 batch_count=1 softmax_elements=2048 is_log_softmax=0 4.27 us, 1.92 GB/s +softmax_fp16_2048 float16 batch_count=1 softmax_elements=2048 is_log_softmax=0 4.48 us, 1.83 GB/s +.... +.... +``` diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index 2cd0467c77..c0acd365c1 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -19,6 +19,10 @@ #include "orttraining/training_ops/cuda/cuda_training_kernels.h" #endif +#ifdef USE_TRITON_KERNEL +#include "core/providers/cuda/triton_kernel.h" +#endif + #include "core/providers/cuda/cuda_stream_handle.h" using namespace onnxruntime::common; @@ -267,6 +271,10 @@ CUDAExecutionProvider::CUDAExecutionProvider(const CUDAExecutionProviderInfo& in CUDA_CALL_THROW(cudaMemGetInfo(&free, &total)); OverrideTunableOpInfoByEnv(info_); + +#ifdef USE_TRITON_KERNEL + onnxruntime::cuda::LoadOrtTritonKernel(); +#endif } CUDAExecutionProvider::~CUDAExecutionProvider() { diff --git a/onnxruntime/core/providers/cuda/triton_kernel.cu b/onnxruntime/core/providers/cuda/triton_kernel.cu new file mode 100644 index 0000000000..c3210d5dad --- /dev/null +++ b/onnxruntime/core/providers/cuda/triton_kernel.cu @@ -0,0 +1,189 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cuda/triton_kernel.h" +#include "core/framework/tunable.h" +#include +#include + +#ifdef USE_TRITON_KERNEL +#include +#include "triton_kernel_infos.h" +#endif + +#define ORT_TRITON_CHECK(status, msg) \ + if ((status) != CUDA_SUCCESS) { \ + ORT_RETURN_IF(true, msg); \ + } + +#define ORT_TRITON_THROW(status, msg) \ + if ((status) != CUDA_SUCCESS) { \ + ORT_THROW(msg); \ + } + +namespace onnxruntime { +namespace cuda { +namespace { + +// a vector of kernel metadata +static std::vector ort_triton_kernel_metadata; + +// store group_name -> [kernel metadata id vector] +static std::unordered_map> ort_triton_kernel_group_map; + +#ifdef USE_TRITON_KERNEL + +// store func_name -> kernel metadata id +static std::unordered_map ort_triton_kernel_map; + +const int GPU_WARP_SIZE = 32; + +Status GetSymbolFromLibrary(const std::string& symbol_name, void** symbol) { + dlerror(); // clear any old error str + + // USe RTLD_DEFAULT for search current lib.so + // value of RTLD_DEFAULT differs across posix platforms (-2 on macos, 0 on linux). + void* handle = RTLD_DEFAULT; + *symbol = dlsym(handle, symbol_name.c_str()); + + char* error_str = dlerror(); + if (error_str) { + Status status = ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Failed to get symbol " + symbol_name + " with error: " + error_str); + return status; + } + // it's possible to get a NULL symbol in our case when Schemas are not custom. + return Status::OK(); +} +#endif + + +/* + * Try to load HIP kernels that compiled by triton. + * They are in hsaco/cubin format, and should use cuModuleLoad to load these kernels. + */ +void TryToLoadKernel() { + auto status = Status::OK(); + +#ifdef USE_TRITON_KERNEL + // get all kernel symbols from curret lib.so + size_t size = sizeof(kernel_infos) / sizeof(kernel_infos[0]); + + for (int i = 0; i < size; ++i) { + auto k_i = kernel_infos[i]; + + void *buff; + ORT_THROW_IF_ERROR(GetSymbolFromLibrary(k_i.name_start, &buff)); + + // try to load module and get function + CUmodule module; + ORT_TRITON_THROW(cuModuleLoadData(&module, buff), "load module data failed."); + + CUfunction function; + ORT_TRITON_THROW(cuModuleGetFunction(&function, module, k_i.func_name), "get funcion from module failed."); + + // setup kernel metadata + TritonKernelMetaData metadata; + metadata.num_warps = k_i.num_warps; + metadata.shared_mem_size = k_i.shared; + metadata.func = function; + std::string fname = k_i.name; // name is not same as func_name + metadata.name = fname; + std::string group_name = k_i.group_name; + + // pass constants + for (auto &kv : k_i.constants) { + metadata.constants[kv.first] = kv.second; + } + + auto idx = ort_triton_kernel_metadata.size(); + ort_triton_kernel_metadata.push_back(metadata); + ort_triton_kernel_map[fname] = idx; + ort_triton_kernel_group_map[group_name].push_back(idx); + LOGS_DEFAULT(VERBOSE) << "loaded ort triton kernel: " << fname << " idx: " << idx; + } +#endif + + ORT_THROW_IF_ERROR(status); +} + +static std::once_flag load_ort_triton_kernel_flag; + +} // end of namespace + +void LoadOrtTritonKernel() { + // load kernel should be called only once + std::call_once(load_ort_triton_kernel_flag, TryToLoadKernel); +} + +Status LaunchTritonKernel(cudaStream_t stream, std::string fname, int grid0, int grid1, int grid2, void* args, size_t args_size) { +#ifdef USE_TRITON_KERNEL + if (ort_triton_kernel_map.count(fname) == 0) { + // return unsupported status when not found function name in registry + // this error status will be used by tunableOp + std::ostringstream message_stream; + message_stream << "can't find ort triton kernel name: " << fname; + std::string message = message_stream.str(); + TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(true, message); + } + auto idx = ort_triton_kernel_map[fname]; + auto metadata = ort_triton_kernel_metadata[idx]; + + void* config[] = {CU_LAUNCH_PARAM_BUFFER_POINTER, args, CU_LAUNCH_PARAM_BUFFER_SIZE, &args_size, + CU_LAUNCH_PARAM_END}; + + ORT_TRITON_CHECK(cuLaunchKernel(metadata.func, + grid0, grid1, grid2, + GPU_WARP_SIZE * metadata.num_warps, 1, 1, + metadata.shared_mem_size, + stream, + nullptr, + (void**)&config), "launch kernel failed."); +#endif + + return Status::OK(); +} + +Status LaunchTritonKernel(cudaStream_t stream, size_t idx, int grid0, int grid1, int grid2, void* args, size_t args_size) { +#ifdef USE_TRITON_KERNEL + if (idx >= ort_triton_kernel_metadata.size()) { + // return unsupported status when not found function name in registry + // this error status will be used by tunableOp + std::ostringstream message_stream; + message_stream << "can't find ort triton kernel idx: " << idx; + std::string message = message_stream.str(); + TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(true, message); + } + auto metadata = ort_triton_kernel_metadata[idx]; + + void* config[] = {CU_LAUNCH_PARAM_BUFFER_POINTER, args, CU_LAUNCH_PARAM_BUFFER_SIZE, &args_size, + CU_LAUNCH_PARAM_END}; + + ORT_TRITON_CHECK(cuLaunchKernel(metadata.func, + grid0, grid1, grid2, + GPU_WARP_SIZE * metadata.num_warps, 1, 1, + metadata.shared_mem_size, + stream, + nullptr, + (void**)&config), "launch kernel failed."); +#endif + + return Status::OK(); +} + +const TritonKernelMetaData* GetOrtTritonKernelMetadata(size_t idx) { + if (idx >= ort_triton_kernel_metadata.size()) { + return nullptr; + } + return &ort_triton_kernel_metadata[idx]; +} + +const std::vector* GetOrtTritonKernelByGroup(std::string group_name) { + if (ort_triton_kernel_group_map.count(group_name) == 0) { + return nullptr; + } + return &ort_triton_kernel_group_map.at(group_name); +} + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/triton_kernel.h b/onnxruntime/core/providers/cuda/triton_kernel.h new file mode 100644 index 0000000000..a4f2b4f549 --- /dev/null +++ b/onnxruntime/core/providers/cuda/triton_kernel.h @@ -0,0 +1,54 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/cuda_common.h" +#include + +namespace onnxruntime { +namespace cuda { + +struct TritonKernelMetaData { + int num_warps; + int shared_mem_size; + CUfunction func; + std::unordered_map constants; + std::string name; +}; + +namespace { + +template +struct DataTypeToName; + +#define DTYPE_TO_STR(type, name) \ + template <> \ + struct DataTypeToName { \ + constexpr static const char* value = name; \ + }; + +DTYPE_TO_STR(float, "fp32"); +DTYPE_TO_STR(half, "fp16"); +DTYPE_TO_STR(double, "fp64"); +DTYPE_TO_STR(BFloat16, "bf16"); + +} // end of namespace + +template +const std::string GetDataTypeName() { + return DataTypeToName::value; +} + +void LoadOrtTritonKernel(); + +Status LaunchTritonKernel(cudaStream_t stream, std::string fname, int grid0, int grid1, int grid2, void* args, size_t args_size); + +const TritonKernelMetaData* GetOrtTritonKernelMetadata(size_t idx); + +const std::vector* GetOrtTritonKernelByGroup(std::string group_name); + +Status LaunchTritonKernel(cudaStream_t stream, size_t idx, int grid0, int grid1, int grid2, void* args, size_t args_size); + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/version_script.lds b/onnxruntime/core/providers/cuda/version_script.lds index 094abb3329..c02a8e4bcf 100644 --- a/onnxruntime/core/providers/cuda/version_script.lds +++ b/onnxruntime/core/providers/cuda/version_script.lds @@ -1,7 +1,8 @@ #_init and _fini should be local VERS_1.0 { global: - GetProvider; + GetProvider; + _binary_*; # Hide everything else. local: diff --git a/onnxruntime/core/providers/rocm/math/softmax.py b/onnxruntime/core/providers/rocm/math/softmax.py new file mode 100644 index 0000000000..6d9c1634e7 --- /dev/null +++ b/onnxruntime/core/providers/rocm/math/softmax.py @@ -0,0 +1,65 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +import triton +import triton.language as tl + + +@triton.jit +def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols, BLOCK_SIZE: tl.constexpr): + # The rows of the softmax are independent, so we parallelize across those + row_idx = tl.program_id(0) + # The stride represents how much we need to increase the pointer to advance 1 row + row_start_ptr = input_ptr + row_idx * input_row_stride + # The block size is the next power of two greater than n_cols, so we can fit each + # row in a single block + col_offsets = tl.arange(0, BLOCK_SIZE) + input_ptrs = row_start_ptr + col_offsets + # Load the row into SRAM, using a mask since BLOCK_SIZE may be > than n_cols + row = tl.load(input_ptrs, mask=col_offsets < n_cols, other=-float("inf")) + row_f32 = row.to(tl.float32) + # Subtract maximum for numerical stability + row_minus_max = row_f32 - tl.max(row_f32, axis=0) + # Note that exponentials in Triton are fast but approximate (i.e., think __expf in CUDA) + numerator = tl.exp(row_minus_max) + denominator = tl.sum(numerator, axis=0) + softmax_output = numerator / denominator + # Write back output to DRAM + output_row_start_ptr = output_ptr + row_idx * output_row_stride + output_ptrs = output_row_start_ptr + col_offsets + tl.store(output_ptrs, softmax_output.to(row.dtype), mask=col_offsets < n_cols) + + +# function_table = {'name': name, 'func': func, 'sig'=sig, kwargs={}}, + +dtypes = ["fp32", "fp16"] +blocks = [1024, 2048, 4096, 8192, 16384] +name_pattern = "softmax_{}_{}" +sig_pattern = "*{},*{},i32,i32,i32" +group_pattern = "softmax_{}" + + +def get_funcion_table(): + func_table = [] + + def get_num_warps(block_size): + num_warps = 4 + if block_size >= 2048: + num_warps = 8 + if block_size >= 4096: + num_warps = 16 + return num_warps + + for dtype in dtypes: + for b in blocks: + name = name_pattern.format(dtype, b) + group = group_pattern.format(dtype) + sig = sig_pattern.format(dtype, dtype) + num_warps = get_num_warps(b) + kwargs = {"num_warps": num_warps, "constants": {"BLOCK_SIZE": b}} + func_desc = {"name": name, "group": group, "func": softmax_kernel, "sig": sig, "kwargs": kwargs} + func_table.append(func_desc) + + return func_table diff --git a/onnxruntime/core/providers/rocm/math/softmax_triton.cuh b/onnxruntime/core/providers/rocm/math/softmax_triton.cuh new file mode 100644 index 0000000000..11cf2daf20 --- /dev/null +++ b/onnxruntime/core/providers/rocm/math/softmax_triton.cuh @@ -0,0 +1,73 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +#include "core/providers/rocm/math/softmax_common.h" +#include "core/providers/rocm/triton_kernel.h" + +namespace onnxruntime { +namespace rocm { + +#ifdef USE_TRITON_KERNEL + +namespace { + +template +std::string GetSoftmaxTritonGroupName() { + std::string ret = "softmax_"; + ret += GetDataTypeName(); + return ret; +} + +} // end of namespace + +template +auto GetSoftmaxTritonOps() { + std::vector>>> ret; + auto group_name = GetSoftmaxTritonGroupName(); + auto *kernel_list = GetOrtTritonKernelByGroup(group_name); + if (kernel_list == nullptr) { + return ret; + } + + for (auto i : *kernel_list) { + // check params match + auto *metadata = GetOrtTritonKernelMetadata(i); + auto block_size = -1; + const std::string block_name = "BLOCK_SIZE"; + if (metadata->constants.count(block_name) != 0) { + block_size = metadata->constants.at(block_name); + } + auto impl = [i, block_size](const SoftmaxParams *params) -> Status { + if (params->is_log_softmax) { + TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(true, "log_softmax not support"); + } + if (block_size < params->softmax_elements) { + TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(true, "BLOCK_SIZE not support"); + } + // construct args for launch kernel + struct { + void *out; + const void *in; + int in_stride; + int out_stride; + int n_cols; + } args = {(void*)params->output, (const void*)params->input, params->input_stride, params->output_stride, params->softmax_elements}; + + // grid dim is (batch_count, 1, 1) + return LaunchTritonKernel(params->stream, i, params->batch_count, 1, 1, &args, sizeof(args)); + }; + ret.emplace_back(std::make_pair(metadata->name, std::move(impl))); + } + return ret; +} + +#endif // USE_TRITON_KERNEL + +} // namespace rocm +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rocm/math/softmax_tunable_op.cuh b/onnxruntime/core/providers/rocm/math/softmax_tunable_op.cuh index 8c761c5b83..d3c943c315 100644 --- a/onnxruntime/core/providers/rocm/math/softmax_tunable_op.cuh +++ b/onnxruntime/core/providers/rocm/math/softmax_tunable_op.cuh @@ -11,6 +11,7 @@ #include "core/providers/rocm/math/softmax_warpwise_impl.cuh" #include "core/providers/rocm/math/softmax_blockwise_impl.cuh" #include "core/providers/rocm/tunable/rocm_tunable.h" +#include "core/providers/rocm/math/softmax_triton.cuh" namespace onnxruntime { namespace rocm { @@ -129,6 +130,14 @@ class SoftmaxTunableOp : public tunable::TunableOpRegisterOp(std::move(op)); } #endif // USE_COMPOSABLE_KERNEL + +#ifdef USE_TRITON_KERNEL + for (auto && [_, op] : GetSoftmaxTritonOps()) { + ORT_UNUSED_PARAMETER(_); + this->RegisterOp(std::move(op)); + } + // this->RegisterOp(SoftmaxTritonOp); +#endif } }; diff --git a/onnxruntime/core/providers/rocm/rocm_execution_provider.cc b/onnxruntime/core/providers/rocm/rocm_execution_provider.cc index fb42792761..14bad1f3fe 100644 --- a/onnxruntime/core/providers/rocm/rocm_execution_provider.cc +++ b/onnxruntime/core/providers/rocm/rocm_execution_provider.cc @@ -19,6 +19,10 @@ #include "orttraining/training_ops/rocm/rocm_training_kernels.h" #endif +#ifdef USE_TRITON_KERNEL +#include "core/providers/rocm/triton_kernel.h" +#endif + using namespace onnxruntime::common; namespace onnxruntime { @@ -208,6 +212,10 @@ ROCMExecutionProvider::ROCMExecutionProvider(const ROCMExecutionProviderInfo& in HIP_CALL_THROW(hipMemGetInfo(&free, &total)); OverrideTunableOpInfoByEnv(info_); + +#ifdef USE_TRITON_KERNEL + onnxruntime::rocm::LoadOrtTritonKernel(); +#endif } ROCMExecutionProvider::~ROCMExecutionProvider() { diff --git a/onnxruntime/core/providers/rocm/version_script.lds b/onnxruntime/core/providers/rocm/version_script.lds index 2c8e9c4b3e..c02a8e4bcf 100644 --- a/onnxruntime/core/providers/rocm/version_script.lds +++ b/onnxruntime/core/providers/rocm/version_script.lds @@ -2,6 +2,7 @@ VERS_1.0 { global: GetProvider; + _binary_*; # Hide everything else. local: diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/softmax.cu b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/softmax.cu index cf364ef9c3..d1d322e3fd 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/softmax.cu +++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/softmax.cu @@ -14,6 +14,7 @@ #include "core/providers/rocm/shared_inc/accumulation_type.h" #include "python/tools/kernel_explorer/device_array.h" #include "python/tools/kernel_explorer/kernel_explorer_interface.h" +#include "core/providers/rocm/math/softmax_triton.cuh" namespace py = pybind11; @@ -171,6 +172,51 @@ class CKSoftmax : public IKernelExplorer { }; #endif // USE_COMPOSABLE_KERNEL +#ifdef USE_TRITON_KERNEL +template +class SoftmaxTriton : public IKernelExplorer { + public: + SoftmaxTriton(DeviceArray& output, DeviceArray& input, int softmax_elements, + int input_stride, int output_stride, int batch_count, bool is_log_softmax) + : params_(TuningContext(), Stream(), static_cast(output.ptr()), static_cast(input.ptr()), + softmax_elements, input_stride, output_stride, batch_count, is_log_softmax) { + for (auto&& [name, op] : rocm::GetSoftmaxTritonOps()) { + name_strings_.emplace_back(name); + ops_.emplace_back(std::move(op)); + } + } + + void Run() override { + ORT_THROW_IF_ERROR(ops_[selected_op_](¶ms_)); + } + + std::vector ListOps() const { + return name_strings_; + } + + bool SelectOp(const std::string& name) { + for (size_t i = 0; i < ops_.size(); i++) { + if (name_strings_[i] == name) { + selected_op_ = i; + Status status = ops_[i](¶ms_); + return status.IsOK(); + } + } + + ORT_THROW("Cannot find implementation ", name); + } + + private: + using ParamsT = rocm::SoftmaxParams; + using OpT = rocm::tunable::Op; + ParamsT params_{}; + std::vector ops_; + std::vector name_strings_; + size_t selected_op_{}; +}; + +#endif // USE_TRITON_KERNEL + #define REGISTER_OP(name, type, vec_size) \ py::class_>(m, #name "_" #type "_" #vec_size) \ .def(py::init()) \ @@ -217,4 +263,11 @@ KE_REGISTER(m) { } #endif // USE_COMPOSABLE_KERNEL +#ifdef USE_TRITON_KERNEL +KE_REGISTER(m) { + REGISTER_OP_TYPED(SoftmaxTriton, half); + REGISTER_OP_TYPED(SoftmaxTriton, float); +} +#endif // USE_TRITON_KERNEL + } // namespace onnxruntime diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/softmax_test.py b/onnxruntime/python/tools/kernel_explorer/kernels/softmax_test.py index fe8c9711bb..96e009c944 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/softmax_test.py +++ b/onnxruntime/python/tools/kernel_explorer/kernels/softmax_test.py @@ -49,7 +49,7 @@ def _test_softmax(batch_count, softmax_elements, is_log_softmax, dtype, func): softmax_op.Run() y_d.UpdateHostNumpyArray() - np.testing.assert_allclose(y_ref, y, rtol=1e-02) + np.testing.assert_allclose(y_ref, y, rtol=1e-02, err_msg=func) dtypes = ["float16", "float32"] diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 38725ff937..3ad5df6aec 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -686,6 +686,8 @@ def parse_arguments(): parser.add_argument("--use_cache", action="store_true", help="Use compiler cache in CI") + parser.add_argument("--use_triton_kernel", action="store_true", help="Use triton compiled kernels") + if not is_windows(): parser.add_argument( "--allow_running_as_root", @@ -990,6 +992,7 @@ def generate_build_tree( "-Donnxruntime_USE_XNNPACK=" + ("ON" if args.use_xnnpack else "OFF"), "-Donnxruntime_USE_WEBNN=" + ("ON" if args.use_webnn else "OFF"), "-Donnxruntime_USE_CANN=" + ("ON" if args.use_cann else "OFF"), + "-Donnxruntime_USE_TRITON_KERNEL=" + ("ON" if args.use_triton_kernel else "OFF"), ] # By default on Windows we currently support only cross compiling for ARM/ARM64 diff --git a/tools/ci_build/compile_triton.py b/tools/ci_build/compile_triton.py new file mode 100644 index 0000000000..212b23a0e6 --- /dev/null +++ b/tools/ci_build/compile_triton.py @@ -0,0 +1,174 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +import argparse +import importlib.util +import os +import shutil + +import triton + + +def compile(function_table, out_dir): + def compile_one(func, sig, **kwargs): + ret = triton.compile(func, signature=sig, **kwargs) + return ret + + metadata = [] + for func_desc in function_table: + name = func_desc["name"] + group = func_desc["group"] + sig = func_desc["sig"] + func = func_desc["func"] + kwargs = func_desc["kwargs"] + + # print("compile func: ", func_desc) + + ret = compile_one(func, sig, **kwargs) + + compile_res = {} + compile_res["name"] = name + compile_res["group"] = group + compile_res["func_name"] = ret.metadata["name"] + compile_res["num_warps"] = ret.metadata["num_warps"] + compile_res["shared"] = ret.metadata["shared"] + if "constants" in kwargs: + compile_res["constants"] = kwargs["constants"] + + # move tmp kernel file into current dir + if "hsaco_path" in ret.asm and os.path.exists(ret.asm["hsaco_path"]): + # is rocm + lib_name = f"{name}.hsaco" + shutil.copyfile(ret.asm["hsaco_path"], f"{out_dir}/{lib_name}") + elif "cubin" in ret.asm: + # is cuda + lib_name = f"{name}.cubin" + # need to write cubin into file + with open(f"{out_dir}/{lib_name}", "wb") as fp: + fp.write(ret.asm["cubin"]) + else: + raise Exception("not find rocm or cuda compiled kernel") + + compile_res["lib_file"] = lib_name + metadata.append(compile_res) + + return metadata + + +def convert_lib_to_obj(lib_file, out_dir): + obj_file = lib_file.split(".")[0] + ".o" + command = f"cd {out_dir}; objcopy -I binary -O elf64-x86-64 -B i386:x86-64 {lib_file} {obj_file}; cd -" + ret = os.system(command) + + if ret != 0: + raise Exception(f"exec convert command: {command} failed.") + # check file exist + if not os.path.exists(f"{out_dir}/{obj_file}"): + raise Exception(f"the output file not exist, after exec comamnd: {command}") + + return obj_file + + +def archive_obj_files(obj_files, out_dir, out_obj_file): + obj_files = " ".join(obj_files) + command = f"cd {out_dir}; ar rcs {out_obj_file} {obj_files}; cd -" + ret = os.system(command) + + if ret != 0: + raise Exception(f"exec convert command: {command} failed.") + # check file exist + if not os.path.exists(f"{out_dir}/{out_obj_file}"): + raise Exception(f"the output file not exist, after exec comamnd: {command}") + + +def convert_and_save(metadata, header_file, out_dir, out_obj_file): + c_metadata = [] + binary_files = [] + for m in metadata: + meta_ele = [] + obj_file = convert_lib_to_obj(m["lib_file"], out_dir) + binary_files.append(obj_file) + + lib_name = m["lib_file"].replace(".", "_") + meta_ele.append(f'"_binary_{lib_name}_start"') + meta_ele.append(f"\"{m['func_name']}\"") + meta_ele.append(f"\"{m['group']}\"") + meta_ele.append(f"\"{m['name']}\"") + meta_ele.append(str(m["num_warps"])) + meta_ele.append(str(m["shared"])) + + # convert constants + constants = [] + for k, v in m["constants"].items(): + constants.append(f'{{ "{k}", {str(v)}}}') + meta_ele.append(f"{{ { ', '.join(constants) } }}") + + c_metadata.append(f"{{ { ', '.join(meta_ele) } }}") + + archive_obj_files(binary_files, out_dir, out_obj_file) + + code = f""" +#include + +struct _TritonKernelInfo {{ + const char* name_start; + const char* func_name; + const char* group_name; + const char* name; + int num_warps; + int shared; + std::unordered_map constants; +}}; + +const _TritonKernelInfo kernel_infos[] = {{ + { ', '.join(c_metadata) }, +}}; + """ + + with open(header_file, "w") as fp: + fp.write(code) + + +def main(args): + out_obj_file = args.obj_file + out_dir = os.path.dirname(out_obj_file) + out_obj_file = os.path.basename(out_obj_file) + if not os.path.exists(out_dir): + os.mkdir(out_dir) + + metadata = [] + print("[triton kernel] start compile triton kernel.") + for i, f in enumerate(args.script_files): + # import module in f, and call function + spec = importlib.util.spec_from_file_location(f"module_{i}", f) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + func_tb = module.get_funcion_table() + m = compile(func_tb, out_dir) + metadata.extend(m) + + print("[triton kernel] compile triton kernel done.") + + # save metadata into header file + convert_and_save(metadata, args.header, out_dir, out_obj_file) + print("[triton kernel] save into file done.") + + +def get_arges(): + parser = argparse.ArgumentParser(description="PyTorch Template Finetune Example") + parser.add_argument( + "--header", type=str, default="triton_kernel_infos.h", help="the header file that should be generated." + ) + parser.add_argument("--ort_root", type=str, default="onnxruntime", help="the root dir of onnxruntime.") + parser.add_argument("--script_files", type=str, nargs="+", help="the root dir of onnxruntime.") + parser.add_argument("--obj_file", type=str, default="triton_kernel_infos.a", help="output target object files.") + + args = parser.parse_args() + return args + + +if __name__ == "__main__": + args = get_arges() + main(args)