[roctx] Python bindings for rocprofiler-sdk-roctx (#402)

* [roctx] Python bindings for rocprofiler-sdk-roctx * Update CHANGELOG --------- Co-authored-by: Jonathan R. Madsen <jonathanrmadsen@gmail.com>
2025-05-19 20:02:51 -05:00
@@ -0,0 +1,9 @@
+#
+#
+#
+
+foreach(_PYTHON_VERSION ${ROCPROFILER_PYTHON_VERSIONS})
+    rocprofiler_roctx_python_bindings(${_PYTHON_VERSION})
+endforeach()
+
+rocprofiler_reset_python3_cache()
@@ -0,0 +1,89 @@
+###############################################################################
+# MIT License
+#
+# Copyright (c) 2023 Advanced Micro Devices, Inc.
+#
+# 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.
+###############################################################################
+
+
+from . import libpyroctx
+from . import context_decorators
+
+__all__ = [
+    "mark",
+    "profilerPause",
+    "profilerResume",
+    "getThreadId",
+    "rangePush",
+    "rangePop",
+    "rangeStart",
+    "rangeStop",
+    "nameOsThread",
+    "nameHipDevice",
+    "context_decorators",
+]
+
+
+def mark(msg):
+    return libpyroctx.roctxMark(msg) if msg is not None else None
+
+
+def profilerPause(tid=0):
+    return libpyroctx.roctxProfilerPause(tid)
+
+
+def profilerResume(tid=0):
+    return libpyroctx.roctxProfilerResume(tid)
+
+
+def getThreadId():
+    return libpyroctx.roctxGetThreadId()
+
+
+def rangePush(msg):
+    return libpyroctx.roctxRangePush(msg)
+
+
+def rangePop():
+    return libpyroctx.roctxRangePop()
+
+
+def rangeStart(msg):
+    return libpyroctx.roctxRangeStart(msg) if msg is not None else None
+
+
+def rangeStop(id=0):
+    return libpyroctx.roctxRangeStop(id) if id is not None else None
+
+
+def nameOsThread(name):
+    return libpyroctx.roctxNameOsThread(name)
+
+
+# def nameHsaAgent(name, agent):
+#     return libpyroctx.roctxNameHsaAgent(name, agent)
+
+
+def nameHipDevice(name, device_id=0):
+    return libpyroctx.roctxNameHipDevice(name, device_id)
+
+
+# def nameHipStream(name, stream):
+#     return libpyroctx.roctxNameHipStream(name, stream)
@@ -0,0 +1,101 @@
+###############################################################################
+# MIT License
+#
+# Copyright (c) 2023 Advanced Micro Devices, Inc.
+#
+# 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.
+###############################################################################
+
+
+from . import libpyroctx
+from functools import wraps
+
+
+class RoctxRange:
+    """Provides decorators and context-manager for roctx range"""
+
+    def __init__(self, msg=None):
+        """Initialize with a message"""
+        self.msg = msg
+
+    def __call__(self, func):
+        """Decorator"""
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            libpyroctx.roctxRangePush(self.msg)
+            try:
+                return func(*args, **kwargs)
+            finally:
+                libpyroctx.roctxRangePop()
+
+        return wrapper
+
+    def __enter__(self):
+        """Context manager start function"""
+        if self.msg is not None:
+            self.a = libpyroctx.roctxRangePush(self.msg)
+            return self.a
+        return self
+
+    def __exit__(self, exc_type, exc_value, tb):
+        """Context manager stop function"""
+        if self.msg is not None:
+            libpyroctx.roctxRangePop()
+
+        if exc_type is not None and exc_value is not None and tb is not None:
+            import traceback
+
+            traceback.print_exception(exc_type, exc_value, tb, limit=5)
+
+
+class RoctxProfiler:
+    """Provides decorators and context-manager for roctx profiler"""
+
+    def __init__(self, tid=0):
+        """Initialize with a tid"""
+        self.tid = tid
+
+    def __call__(self, func):
+        """Decorator"""
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            libpyroctx.roctxProfilerResume(self.tid)
+            try:
+                return func(*args, **kwargs)
+            finally:
+                libpyroctx.roctxProfilerPause(self.tid)
+
+        return wrapper
+
+    def __enter__(self):
+        """Context manager start function"""
+        self.a = libpyroctx.roctxProfilerResume(self.tid)
+        return self.a
+
+    def __exit__(self, exc_type, exc_value, tb):
+        """Context manager stop function"""
+
+        libpyroctx.roctxProfilerPause(self.tid)
+
+        if exc_type is not None and exc_value is not None and tb is not None:
+            import traceback
+
+            traceback.print_exception(exc_type, exc_value, tb, limit=5)
@@ -0,0 +1,97 @@
+// MIT License
+//
+// Copyright (c) 2022 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.
+
+#include "libpyroctx.hpp"
+
+#include <rocprofiler-sdk-roctx/roctx.h>
+#include <rocprofiler-sdk-roctx/types.h>
+
+namespace py = ::pybind11;
+
+PYBIND11_MODULE(libpyroctx, pyroctx)
+{
+    py::doc("Rocprofiler-SDK ROCTx Python bindings");
+
+    pyroctx.def(
+        "roctxMark",
+        [](const std::string& _msg) { roctxMarkA(_msg.c_str()); },
+        "Mark an event in any attached profiler");
+
+    pyroctx.def(
+        "roctxProfilerPause",
+        [](roctx_thread_id_t tid) { return roctxProfilerPause(tid); },
+        "Pause data collection in any attached profiler");
+
+    pyroctx.def(
+        "roctxProfilerResume",
+        [](roctx_thread_id_t tid) { return roctxProfilerResume(tid); },
+        "Resume data collection in any attached profiler");
+
+    pyroctx.def(
+        "roctxGetThreadId",
+        []() {
+            auto _tid = roctx_thread_id_t{0};
+            roctxGetThreadId(&_tid);
+            return _tid;
+        },
+        "Get the current thread ID");
+
+    pyroctx.def(
+        "roctxRangePush",
+        [](const std::string& _msg) { return roctxRangePushA(_msg.c_str()); },
+        "Start a new nested range");
+
+    pyroctx.def(
+        "roctxRangePop", []() { return roctxRangePop(); }, "Stop the current nested range");
+
+    pyroctx.def(
+        "roctxRangeStart",
+        [](const std::string& _msg) { return roctxRangeStartA(_msg.c_str()); },
+        "Start a process range");
+
+    pyroctx.def(
+        "roctxRangeStop", [](roctx_range_id_t id) { roctxRangeStop(id); }, "Stop a process range");
+
+    pyroctx.def(
+        "roctxNameOsThread",
+        [](const std::string& name) { return roctxNameOsThread(name.c_str()); },
+        "Label the current CPU OS thread with the provided name");
+
+    // pyroctx.def(
+    //     "roctxNameHsaAgent",
+    //     [](const std::string& name, const struct hsa_agent_s* agent) { return
+    //     roctxNameHsaAgent(name.c_str(), agent); }, "Label the given HSA agent with the provided
+    //     name");
+
+    pyroctx.def(
+        "roctxNameHipDevice",
+        [](const std::string& name, int device_id) {
+            return roctxNameHipDevice(name.c_str(), device_id);
+        },
+        "Label the given HIP device id with the provided name");
+
+    // pyroctx.def(
+    //     "roctxNameHipStream",
+    //     [](const std::string& name, const struct ihipStream_t* stream) { return
+    //     roctxNameHipStream(name.c_str(), stream); }, "Label the given HIP stream with the
+    //     provided name");
+}
@@ -0,0 +1,36 @@
+// MIT License
+//
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include <pybind11/cast.h>
+#include <pybind11/detail/common.h>
+#include <pybind11/embed.h>
+#include <pybind11/eval.h>
+#include <pybind11/functional.h>
+#include <pybind11/iostream.h>
+#include <pybind11/numpy.h>
+#include <pybind11/operators.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
+#include <pybind11/stl.h>
+#include <pyerrors.h>