saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[ROCm] Some fixes in tunable (#18575)

Browse source 
### Description

* Fix workspace size for hipBLASLt algos at 32M
* Update according to API changes
This commit is contained in:
mindest 2023-11-24 17:22:00 +08:00 committed by GitHub
parent 62f00ad8e7
commit b9c935f605
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 9 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/contrib_ops/rocm/diffusion/group_norm_triton.cuh
View file

@ -77,7 +77,7 @@ auto GetTritonGroupNormNHWCTypeStringAndOps() {
params->epsilon};
// Grid dim is (batch_count, groups, 1)
return LaunchTritonKernel(params->stream, i, params->n, params->groups, 1, &args, sizeof(args));
return LaunchTritonKernel(params->StreamHandle(), i, params->n, params->groups, 1, &args, sizeof(args));
};
ret.emplace_back(std::make_pair(metadata->name, std::move(impl)));
}

2
onnxruntime/core/providers/rocm/math/softmax_triton.cuh
View file

@ -60,7 +60,7 @@ auto GetSoftmaxTritonOps() {
} 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));
return LaunchTritonKernel(params->StreamHandle(), i, params->batch_count, 1, 1, &args, sizeof(args));
};
ret.emplace_back(std::make_pair(metadata->name, std::move(impl)));
}

7
onnxruntime/core/providers/rocm/tunable/gemm_hipblaslt.h
View file

@ -26,6 +26,10 @@ using onnxruntime::contrib::rocm::blas::GemmFastGeluParams;
#ifdef USE_HIPBLASLT
// For large K and small M/N, K dim will be split to multiple workgroups and buffers,
// which will require additional workspace. Here we set the max workspace size to 32MB.
constexpr const size_t kHipBlasLtMaxWorkSpaceSizeInBytes = 32 * 1024 * 1024;
enum ActivationType {
NONE = 0,
RELU = 1,
@ -225,6 +229,9 @@ auto GetHipBlasLtTypeStringAndOps(ActivationType activation_type = ActivationTyp
IAllocatorUniquePtr<void> workspace_buffer;
if (workspace_size > 0) {
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(workspace_size > kHipBlasLtMaxWorkSpaceSizeInBytes,
"Workspace size exceeds limit (32M): ", workspace_size);
workspace_size = kHipBlasLtMaxWorkSpaceSizeInBytes;
workspace_buffer = params->tuning_ctx->GetScratchBuffer(workspace_size, params->stream);
}