From d2aa2109c05f4008b09eaf0de1907072bc055644 Mon Sep 17 00:00:00 2001 From: cloudhan Date: Thu, 15 Sep 2022 21:11:27 +0800 Subject: [PATCH] Make TunableOp follow stream semantics (#12856) --- .../contrib_ops/rocm/bert/tunable_op.h | 2 +- onnxruntime/contrib_ops/rocm/bert/util.cc | 6 +++--- onnxruntime/contrib_ops/rocm/bert/util.h | 3 ++- .../kernel_explorer_interface.h | 2 +- .../tools/kernel_explorer/kernels/gemm_ck.cc | 1 + .../kernel_explorer/kernels/gemm_rocblas.cc | 1 + .../kernel_explorer/kernels/gemm_rocblas.h | 20 +++++++++++++++++++ .../kernel_explorer/kernels/gemm_tunable.cc | 1 + 8 files changed, 30 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/rocm/bert/tunable_op.h b/onnxruntime/contrib_ops/rocm/bert/tunable_op.h index cebda8848c..39bec11300 100644 --- a/onnxruntime/contrib_ops/rocm/bert/tunable_op.h +++ b/onnxruntime/contrib_ops/rocm/bert/tunable_op.h @@ -116,7 +116,7 @@ class TunableOp { static double Profile(Op& op, const ParamsT* param) { const int num_iter = 100; - Timer timer{}; + Timer timer{param->stream}; timer.Start(); for (int i = 0; i < num_iter; i++) { ORT_THROW_IF_ERROR(op(param)); diff --git a/onnxruntime/contrib_ops/rocm/bert/util.cc b/onnxruntime/contrib_ops/rocm/bert/util.cc index 2ed0a8d162..04874d1204 100644 --- a/onnxruntime/contrib_ops/rocm/bert/util.cc +++ b/onnxruntime/contrib_ops/rocm/bert/util.cc @@ -12,18 +12,18 @@ int CeilingDivision(int n, int m) { return r; } -Timer::Timer() { +Timer::Timer(hipStream_t stream): stream_(stream) { HIP_CHECK(hipEventCreate(&start_)); HIP_CHECK(hipEventCreate(&end_)); } void Timer::Start() { HIP_CHECK(hipDeviceSynchronize()); - HIP_CHECK(hipEventRecord(start_, nullptr)); + HIP_CHECK(hipEventRecord(start_, stream_)); } void Timer::End() { - HIP_CHECK(hipEventRecord(end_, nullptr)); + HIP_CHECK(hipEventRecord(end_, stream_)); HIP_CHECK(hipEventSynchronize(end_)); } diff --git a/onnxruntime/contrib_ops/rocm/bert/util.h b/onnxruntime/contrib_ops/rocm/bert/util.h index 9c5c2900e0..16fe3d538f 100644 --- a/onnxruntime/contrib_ops/rocm/bert/util.h +++ b/onnxruntime/contrib_ops/rocm/bert/util.h @@ -29,13 +29,14 @@ struct alignas(sizeof(T) * VecSize) AlignedVector { class Timer { public: - Timer(); + explicit Timer(hipStream_t stream); void Start(); void End(); float Duration(); ~Timer(); private: + hipStream_t stream_; hipEvent_t start_; hipEvent_t end_; }; diff --git a/onnxruntime/python/tools/kernel_explorer/kernel_explorer_interface.h b/onnxruntime/python/tools/kernel_explorer/kernel_explorer_interface.h index c09d621417..32285ab5b8 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernel_explorer_interface.h +++ b/onnxruntime/python/tools/kernel_explorer/kernel_explorer_interface.h @@ -22,7 +22,7 @@ class IKernelExplorer { for (int i = 0; i < 5; i++) { Run(); } - Timer timer; + Timer timer{Stream()}; timer.Start(); for (int i = 0; i < repeats_; i++) { Run(); diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.cc b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.cc index 39f9bdbc8b..188aa979bd 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.cc +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.cc @@ -30,6 +30,7 @@ class CKGemm : public IKernelExplorer { auto supports_b = opb == BlasOp::N ? std::is_same_v : std::is_same_v; ORT_ENFORCE(supports_a && supports_b); + params_.stream = Stream(); // rocblas handle is not used for ck params_.handle = nullptr; params_.opa = opa; diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.cc b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.cc index 1a23a401ca..f2da8513d7 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.cc +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.cc @@ -27,6 +27,7 @@ class RocBlasGemm : public IKernelExplorer { double beta, DeviceArray& c, int64_t ldc) { ROCBLAS_CALL_THROW(rocblas_create_handle(&rocblas_handle_)); + params_.stream = Stream(); params_.handle = rocblas_handle_; params_.opa = opa; params_.opb = opb; diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.h b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.h index f6c5c733e6..4c4518d4d2 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.h +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.h @@ -13,9 +13,29 @@ namespace py = pybind11; namespace onnxruntime { +// RAII style guard to set stream and restore original stream for rocblas_handle +class RocblasHandleStreamGuard { + public: + RocblasHandleStreamGuard(rocblas_handle handle, hipStream_t stream) : handle_{handle} { + ROCBLAS_CALL_THROW(rocblas_get_stream(handle_, &original_stream_)); + ROCBLAS_CALL_THROW(rocblas_set_stream(handle_, stream)); + } + + ~RocblasHandleStreamGuard() { + ROCBLAS_CALL_THROW(rocblas_set_stream(handle_, original_stream_)); + } + + ORT_DISALLOW_COPY_AND_ASSIGNMENT(RocblasHandleStreamGuard); + + private: + rocblas_handle handle_; + hipStream_t original_stream_; +}; + // to be moved to onnxruntime once we have a monolithicly tunable gemm wrapper and it is enabled for onnxruntime template Status RocBlasGemmOp(const GemmParams* params) { + RocblasHandleStreamGuard guard(params->handle, params->stream); // NOTE: rocblas assumes the storage is column-majored, swapping A and B makes it have the same interface // as those with row-majored convention. That is, if you treat the storage as row-majored but view the matrices as // transposed, then by using the property Transpose(A*B) = Tranpose(B)*Transpose(A), the correctness is obvious. diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_tunable.cc b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_tunable.cc index e33f3e0e29..e93759a856 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_tunable.cc +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_tunable.cc @@ -40,6 +40,7 @@ class GemmTunable : public IKernelExplorer { double beta, DeviceArray& c, int64_t ldc) { ROCBLAS_CALL_THROW(rocblas_create_handle(&rocblas_handle_)); + params_.stream = Stream(); params_.handle = rocblas_handle_; params_.opa = opa; params_.opb = opb;