Make TunableOp follow stream semantics (#12856)

onnxruntime/contrib_ops/rocm/bert/tunable_op.h

@@ -116,7 +116,7 @@ class TunableOp {
 
   static double Profile(Op<ParamsT>& 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));

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_));
 }
 

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_;
 };

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

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<BLayout, Row> : std::is_same_v<BLayout, Col>;
     ORT_ENFORCE(supports_a && supports_b);
 
+    params_.stream = Stream();
     // rocblas handle is not used for ck
     params_.handle = nullptr;
     params_.opa = opa;

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;

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 <typename T>
 Status RocBlasGemmOp(const GemmParams<T>* 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.

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;