From b75463345818e57d4d7c5aba59fd47b863ab667d Mon Sep 17 00:00:00 2001
From: "Nichols A. Romero" <165712832+naromero77amd@users.noreply.github.com>
Date: Tue, 11 Feb 2025 12:48:05 -0600
Subject: [PATCH] [rocm6.4_internal_testing] [ROCm][TunableOp] Close offline
 tuning results file when offline tuning is disabled. (#1889)

This is cherry-pick of an upstream PR that has been approved but is
unmerged (due to CI delays). The PR did previously pass ROCm tests.
https://github.com/pytorch/pytorch/pull/146574

This PR is to fix UT breakage that has been reported internally and is
considered high priority. When tunable.record_untuned_enable(False) is
invoked, we flush the results of the untuned gemm file.

Offline tuning I/O currently doesn't have a set untuned results filename
member function or untuned results write to file member function. When
performing back-to-back unit tests, the same ofstream ends up getting
reused between UTs. Due to the way the UT are executed, this can lead to
unexpected failures.
---
 aten/src/ATen/cuda/tunable/Tunable.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/aten/src/ATen/cuda/tunable/Tunable.cpp b/aten/src/ATen/cuda/tunable/Tunable.cpp
index 1ef425b617c3a..ff45d43b3d8b3 100644
--- a/aten/src/ATen/cuda/tunable/Tunable.cpp
+++ b/aten/src/ATen/cuda/tunable/Tunable.cpp
@@ -458,6 +458,8 @@ void TuningContext::EnableRecordUntuned(bool value) {
     TUNABLE_LOG1("Enable Record Untuned for TunableOp");
   } else {
     TUNABLE_LOG1("Disable Record Untuned for TunableOp");
+    TUNABLE_LOG1("Closing Untuned GEMM Results File");
+    untuned_file_.close();
   }
 }