From 79989e8f79288eac95b42e49c28acc7c0a2ce1e0 Mon Sep 17 00:00:00 2001
From: Braden Stefanuk <121893577+bstefanuk@users.noreply.github.com>
Date: Wed, 18 Sep 2024 18:04:48 -0500
Subject: [PATCH] feat: add profiling to tensile

---
 tensilelite/Tensile/TensileCreateLibrary.py |  4 +-
 tensilelite/Tensile/Utilities/Profile.py    | 77 +++++++++++++++++++++
 2 files changed, 79 insertions(+), 2 deletions(-)
 create mode 100644 tensilelite/Tensile/Utilities/Profile.py

diff --git a/tensilelite/Tensile/TensileCreateLibrary.py b/tensilelite/Tensile/TensileCreateLibrary.py
index b8cea8455..1bc24bd1d 100644
--- a/tensilelite/Tensile/TensileCreateLibrary.py
+++ b/tensilelite/Tensile/TensileCreateLibrary.py
@@ -41,6 +41,7 @@
 from .SolutionLibrary import MasterSolutionLibrary
 from .SolutionStructs import Solution
 from .CustomYamlLoader import load_logic_gfx_arch
+from .Utilities.Profile import profile
 
 import argparse
 import collections
@@ -1233,7 +1234,7 @@ def validateLibrary(masterLibraries: MasterSolutionLibrary,
 ################################################################################
 # Tensile Create Library
 ################################################################################
-@timing
+@profile
 def TensileCreateLibrary():
   print1("")
   print1(HR)
@@ -1558,7 +1559,6 @@ def param(key, value):
 
   print1("# Check if generated files exists.")
 
-  @timing
   def checkFileExistence(files):
     for filePath in files:
       if not os.path.exists(filePath):
diff --git a/tensilelite/Tensile/Utilities/Profile.py b/tensilelite/Tensile/Utilities/Profile.py
new file mode 100644
index 000000000..cc3c7eb44
--- /dev/null
+++ b/tensilelite/Tensile/Utilities/Profile.py
@@ -0,0 +1,77 @@
+################################################################################
+#
+# Copyright (C) 2016-2024 Advanced Micro Devices, Inc. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+################################################################################
+
+import cProfile
+import pstats
+import os
+
+from pathlib import Path
+from datetime import datetime, timezone
+from typing import Callable, Tuple
+
+PROFILE_ENV_VAR = "TENSILE_PROFILE"
+
+def profile(func: Callable) -> Callable:
+    """Profiling decorator. 
+
+    Add ``@profile`` to mark a function for profiling; set the environment variable 
+    TENSILE_PROFILE=ON to enable profiling decorated functions.
+    """
+    if not envVariableIsSet(PROFILE_ENV_VAR):
+        return func
+    def wrapper(*args, **kwargs):
+        path, filename = initProfileArtifacts(func.__name__)
+
+        prof = cProfile.Profile()
+        output = prof.runcall(func, *args, **kwargs)
+        result = pstats.Stats(prof)
+        result.sort_stats(pstats.SortKey.TIME)
+        result.dump_stats(path/filename)
+
+        return output
+    return wrapper
+
+def envVariableIsSet(varName: str) -> bool:
+    """Checks if the provided environment variable is set to "ON", "TRUE", or "1"
+    Args:
+        varName: Environment variable name.
+    Returns:
+        True if the environment variable is set, otherwise False.
+    """
+    value = os.environ.get(varName, "").upper()
+    return True if value in ["ON", "TRUE", "1"] else False
+
+def initProfileArtifacts(funcName: str) -> Tuple[Path, str]:
+    """Initializes filenames and paths for profiling artifacts based on the current datetime
+    Args:
+        funcName: The name of the function being profiled, nominally passed via func.__name__ 
+    Returns:
+        A tuple (path, filename) where the path is the artifact directory and filename is
+        a .prof file with the profiling results.
+    """
+    dt = datetime.now(timezone.utc)
+    filename = f"{funcName}-{dt.strftime('%Y-%m-%dT%H-%M-%SZ')}.prof"
+    path = Path().cwd()/f"profiling-results-{dt.strftime('%Y-%m-%d')}"
+    path.mkdir(exist_ok=True)
+    return path, filename