saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-24 19:43:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

integrate triton into ort (#15862)

Browse source 
### 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.

This PR is to integrate triton into ort, so that ort can use kernels
that written and compiled by triton.

The main change focus on two part:
1. a build part to compile triton written kernel and combine these
kernels into libonnxruntime_providers_rocm.so
2. a loader and launcher in c++, for loading and launch triton written
kernels.

#### Build

To compile triton written kernel, add a script
`tools/ci_build/compile_triton.py`. This script will dynamic load all
kernel files, compile them, and generate `triton_kernel_infos.a` and
`triton_kernel_infos.h`.

`triton_kernel_infos.a` contains all compiled kernel instructions, this
file will be combined into libonnxruntime_providers_rocm.so, using
--whole-archive flag.

`triton_kernel_infos.h` defines a const array that contains all the
metadata for each compiled kernel. These metadata will be used for load
and launch. So this header file is included by 'triton_kernel.cu' which
defines load and launch functions.

Add a build flag in build.py and CMakeList.txt, when building rocm
provider, it will call triton_kernel build command, and generate all
necessary files.

#### C++ Load and Launch

On c++ part, we implement load and launch functions in triton_kernel.cu
and triton_kernel.h.

These two files located in `providers/cuda`, and when compiling rocm,
they will be hipified. so this part supports both cuda and rocm. But
currently we only call triton kernel in rocm.

We also implement a softmax triton op for example. Because there will
generate many kernels for different input shape of softmax, we use
TunableOp to select the best one.


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
kailums 2023-05-17 09:35:28 +08:00 committed by GitHub
parent a7ad859e3a
commit f62f722c70
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
18 changed files with 796 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
cmake/CMakeLists.txt
View file

@ -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)

32
cmake/onnxruntime_compile_triton_kernel.cmake Normal file
View file

@ -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()

3
cmake/onnxruntime_kernel_explorer.cmake
View file

@ -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)

22
cmake/onnxruntime_providers.cmake
View file

@ -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

98
docs/ORT_use_trtion_kernel.md Normal file
View file

@ -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 <typename T, typename OutputT>
auto GetSoftmaxTritonOps() {
std::vector<std::pair<std::string, tunable::Op<SoftmaxParams<T, OutputT>>>> ret;
auto group_name = GetSoftmaxTritonGroupName<T>();
// 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
....
....
```

8
onnxruntime/core/providers/cuda/cuda_execution_provider.cc
View file

@ -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() {

189
onnxruntime/core/providers/cuda/triton_kernel.cu Normal file
View file

@ -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 <fstream>
#include <thread>
#ifdef USE_TRITON_KERNEL
#include <dlfcn.h>
#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<TritonKernelMetaData> ort_triton_kernel_metadata;
// store group_name -> [kernel metadata id vector]
static std::unordered_map<std::string, std::vector<int>> ort_triton_kernel_group_map;
#ifdef USE_TRITON_KERNEL
// store func_name -> kernel metadata id
static std::unordered_map<std::string, int> 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<int>* 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

54
onnxruntime/core/providers/cuda/triton_kernel.h Normal file
View file

@ -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 <cuda.h>
namespace onnxruntime {
namespace cuda {
struct TritonKernelMetaData {
int num_warps;
int shared_mem_size;
CUfunction func;
std::unordered_map<std::string, int> constants;
std::string name;
};
namespace {
template <typename T>
struct DataTypeToName;
#define DTYPE_TO_STR(type, name) \
template <> \
struct DataTypeToName<type> { \
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 <typename T>
const std::string GetDataTypeName() {
return DataTypeToName<T>::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<int>* 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

3
onnxruntime/core/providers/cuda/version_script.lds
View file

@ -1,7 +1,8 @@
#_init and _fini should be local
VERS_1.0 {
global:
GetProvider;
GetProvider;
_binary_*;
# Hide everything else.
local:

65
onnxruntime/core/providers/rocm/math/softmax.py Normal file
View file

@ -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

73
onnxruntime/core/providers/rocm/math/softmax_triton.cuh Normal file
View file

@ -0,0 +1,73 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <string>
#include <utility>
#include <vector>
#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 <typename T>
std::string GetSoftmaxTritonGroupName() {
std::string ret = "softmax_";
ret += GetDataTypeName<T>();
return ret;
}
} // end of namespace
template <typename T, typename OutputT>
auto GetSoftmaxTritonOps() {
std::vector<std::pair<std::string, tunable::Op<SoftmaxParams<T, OutputT>>>> ret;
auto group_name = GetSoftmaxTritonGroupName<T>();
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<T, OutputT> *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

9
onnxruntime/core/providers/rocm/math/softmax_tunable_op.cuh
View file

@ -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::TunableOp<SoftmaxParams<InputT, OutputT
this->RegisterOp(std::move(op));
}
#endif // USE_COMPOSABLE_KERNEL
#ifdef USE_TRITON_KERNEL
for (auto && [_, op] : GetSoftmaxTritonOps<InputT, OutputT>()) {
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
}
// this->RegisterOp(SoftmaxTritonOp<InputT, OutputT>);
#endif
}
};

8
onnxruntime/core/providers/rocm/rocm_execution_provider.cc
View file

@ -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() {

1
onnxruntime/core/providers/rocm/version_script.lds
View file

@ -2,6 +2,7 @@
VERS_1.0 {
global:
GetProvider;
_binary_*;
# Hide everything else.
local:

53
onnxruntime/python/tools/kernel_explorer/kernels/rocm/softmax.cu
View file

@ -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 <typename T>
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<T*>(output.ptr()), static_cast<T*>(input.ptr()),
softmax_elements, input_stride, output_stride, batch_count, is_log_softmax) {
for (auto&& [name, op] : rocm::GetSoftmaxTritonOps<T, T>()) {
name_strings_.emplace_back(name);
ops_.emplace_back(std::move(op));
}
}
void Run() override {
ORT_THROW_IF_ERROR(ops_[selected_op_](&params_));
}
std::vector<std::string> 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](&params_);
return status.IsOK();
}
}
ORT_THROW("Cannot find implementation ", name);
}
private:
using ParamsT = rocm::SoftmaxParams<T, T>;
using OpT = rocm::tunable::Op<ParamsT>;
ParamsT params_{};
std::vector<OpT> ops_;
std::vector<std::string> name_strings_;
size_t selected_op_{};
};
#endif // USE_TRITON_KERNEL
#define REGISTER_OP(name, type, vec_size) \
py::class_<name<type, vec_size>>(m, #name "_" #type "_" #vec_size) \
.def(py::init<DeviceArray&, DeviceArray&, int, int, int, int, bool>()) \
@ -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

2
onnxruntime/python/tools/kernel_explorer/kernels/softmax_test.py
View file

@ -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"]

3
tools/ci_build/build.py
View file

@ -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

174
tools/ci_build/compile_triton.py Normal file
View file

@ -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 <unordered_map>
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<std::string, int> 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)