From d82036dbbd0136348d5360b0c3f8231154dcf627 Mon Sep 17 00:00:00 2001
From: cloudhan <guangyunhan@microsoft.com>
Date: Tue, 25 Oct 2022 17:44:28 +0800
Subject: [PATCH] Add Pre- and Post-tunning API to allow pre- and
 postprocessing of params (#13411)

Some op will use a buffer for input and output at the same time, so it will do inplace update to it.
If we blindly tune over the `params`, there will be accumulated update to that buffer during FindFastest,
which is an undesired side effect. In this case, we use a proxy params struct for the tuning to avoid this side effect.
---
 .../core/providers/rocm/tunable/gemm.cu       | 18 +-------
 .../providers/rocm/tunable/gemm_tunable.cuh   | 41 +++++++++++++++++++
 .../core/providers/rocm/tunable/tunable.h     | 17 +++++++-
 3 files changed, 58 insertions(+), 18 deletions(-)

diff --git a/onnxruntime/core/providers/rocm/tunable/gemm.cu b/onnxruntime/core/providers/rocm/tunable/gemm.cu
index 5f860356c2..d345b3ca7d 100644
--- a/onnxruntime/core/providers/rocm/tunable/gemm.cu
+++ b/onnxruntime/core/providers/rocm/tunable/gemm.cu
@@ -18,22 +18,6 @@ namespace blas {
 
 namespace row_major {
 
-template <typename T>
-bool IsZero(T v) {
-  return v == 0.0f;
-}
-
-template <>
-bool IsZero(BFloat16 v) {
-  return v.val == 0;
-}
-
-template <>
-bool IsZero(half v) {
-  return __half2float(v) == 0.0f;
-}
-
-
 template <typename T, typename ScalarT>
 inline GEMM(T, ScalarT) {
   GemmParams<T> params;
@@ -63,7 +47,7 @@ inline GEMM(T, ScalarT) {
   params.ldc = ldc;
 
   // TODO: current implementation for beta != 0 will cause repeatedly inplace update in C buffer. Skip them for now.
-  if (IsZero(params.beta) && tunable) {
+  if (tunable) {
     if (opa == BlasOp::N && opb == BlasOp::N) {
       static internal::GemmTunableOp<T, internal::Row, internal::Row> gemm{};
       gemm.EnableTuning();
diff --git a/onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh b/onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh
index dc185c793d..b347a7ff0f 100644
--- a/onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh
+++ b/onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh
@@ -16,6 +16,22 @@ namespace tunable {
 namespace blas {
 namespace internal {
 
+template <typename T>
+bool IsZero(T v) {
+  return v == 0.0f;
+}
+
+template <>
+bool IsZero(BFloat16 v) {
+  return v.val == 0;
+}
+
+template <>
+bool IsZero(half v) {
+  return __half2float(v) == 0.0f;
+}
+
+
 template <typename T, typename ALayout, typename BLayout>
 class GemmTunableOp : public tunable::TunableOp<GemmParams<T>> {
  public:
@@ -26,6 +42,31 @@ class GemmTunableOp : public tunable::TunableOp<GemmParams<T>> {
       this->ops_.emplace_back(std::move(op));
     }
   }
+
+  const GemmParams<T>* PreTuning(const GemmParams<T>* params) override {
+    if (!IsZero(params->beta)) {
+      // When beta != 0, C buffer is used as an input as well as an output. We need to create a proxy params for the
+      // tuning process. Otherwise, tuning will cause the C buffer been updated accumulatedly, say, we tune it for n
+      // iterations, then during tuning C^(1) = alpha A B + beta C^(0), ..., C^(n) = alpha A B + beta C^(n-1). And for
+      // the actual run after tuning, the result will be C^(n+1), whereas what we want is C^(1). This only happens if
+      // the tuning's FindFastest is invoked.
+      //
+      // Note, C^(i) is the C at i-th iteration.
+      GemmParams<T>* proxy = new GemmParams<T>();
+      *proxy = *params;
+      HIP_CALL_THROW(hipMalloc(&(proxy->c), proxy->m * proxy->ldc * sizeof(T)));
+      return proxy;
+    }
+
+    return params;
+  }
+
+  void PostTuning(const GemmParams<T>* params) override {
+    if (!IsZero(params->beta)) {
+      HIP_CALL_THROW(hipFree(params->c));
+      delete params;
+    }
+  }
 };
 
 }  // namespace internal
diff --git a/onnxruntime/core/providers/rocm/tunable/tunable.h b/onnxruntime/core/providers/rocm/tunable/tunable.h
index d50a7b0c3c..66016214cf 100644
--- a/onnxruntime/core/providers/rocm/tunable/tunable.h
+++ b/onnxruntime/core/providers/rocm/tunable/tunable.h
@@ -28,6 +28,7 @@ namespace tunable {
 struct OpParams {
   OpParams() : stream{} {}
   explicit OpParams(hipStream_t stream) : stream(stream) {}
+  virtual ~OpParams() = default;
   virtual std::string Signature() const = 0;
   hipStream_t stream;
 };
@@ -80,7 +81,9 @@ class TunableOp {
     int id;
     if (tuning_) {
       if (kernel_map_.find(params->Signature()) == kernel_map_.end()) {
-        id = FindFastest(params);
+        auto maybe_proxy_params = this->PreTuning(params);
+        id = FindFastest(maybe_proxy_params);
+        PostTuning(maybe_proxy_params);
         kernel_map_.insert({params->Signature(), id});
       } else {
         id = kernel_map_[params->Signature()];
@@ -100,6 +103,18 @@ class TunableOp {
     tuning_ = false;
   }
 
+  // We might want to do some tricks to the `params`, e.g., some op will use a buffer for input and output at the same
+  // time, so it will do inplace update to it. If we blindly tune over the `params`, there will be accumulated update
+  // to that buffer during FindFastest, which is an undesired side effect. In this case, we must prepare a new (proxy)
+  // params struct for the tuning to avoid this side effect.
+  virtual const ParamsT* PreTuning(const ParamsT* params) {
+    return params;
+  }
+
+  virtual void PostTuning(const ParamsT* /*params*/) {
+    // Do nothing if we are not playing around with params
+  }
+
   virtual ~TunableOp() = default;
 
  protected: