draft of rocprofiler-sdk into XLA

third_party/tsl/third_party/gpus/rocm_configure.bzl

@@ -206,6 +206,9 @@ def _rocm_include_path(repository_ctx, rocm_config, bash_bin):
     inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/17/include")
     inc_dirs.append(rocm_toolkit_path + "/llvm/lib/clang/18/include")
 
+    if int(rocm_config.rocm_version_number) >= 60200:
+        inc_dirs.append(rocm_toolkit_path + "/lib/llvm/lib/clang/18/include")
+
     # Support hcc based off clang 10.0.0 (for ROCm 3.3)
     inc_dirs.append(rocm_toolkit_path + "/hcc/compiler/lib/clang/10.0.0/include/")
     inc_dirs.append(rocm_toolkit_path + "/hcc/lib/clang/10.0.0/include")

third_party/tsl/tsl/platform/default/dso_loader.cc

@@ -170,7 +170,8 @@ absl::StatusOr<void*> GetHipsolverDsoHandle() {
 #endif
 
 absl::StatusOr<void*> GetRoctracerDsoHandle() {
-  return GetDsoHandle("roctracer64", "");
+  // return GetDsoHandle("roctracer64", "");
+  return GetDsoHandle("rocprofiler-sdk", "");
 }
 
 absl::StatusOr<void*> GetHipsparseDsoHandle() {

xla/backends/profiler/gpu/BUILD

@@ -31,7 +31,7 @@ tsl_gpu_library(
         ":cupti_collector",
         ":cupti_tracer",
         ":cupti_wrapper",
-        ":rocm_collector",
+        # ":rocm_collector",
         ":rocm_tracer",
     ],
     deps = [
@@ -213,6 +213,7 @@ tsl_gpu_library(
     ],
 )
 
+
 tsl_gpu_library(
     name = "rocm_collector",
     srcs = if_rocm(["rocm_collector.cc"]),
@@ -256,6 +257,7 @@ tsl_gpu_library(
     visibility = ["//visibility:public"],
     deps = [
         ":rocm_collector",
+        "//xla/backends/profiler/gpu/common:common_headers",
         "//xla/stream_executor/rocm:roctracer_wrapper",
         "@com_google_absl//absl/container:fixed_array",
         "@com_google_absl//absl/container:flat_hash_map",

xla/backends/profiler/gpu/common/BUILD

@@ -0,0 +1,26 @@
+load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm")
+load(
+    "@tsl//tsl/platform:build_config.bzl",
+    "tf_additional_device_tracer_srcs",
+)
+load("@tsl//tsl/profiler/builds:build_config.bzl", "tf_profiler_copts")
+load("//xla/tests:build_defs.bzl", "xla_test")
+load(
+    "//xla/tsl:tsl.bzl",
+    "internal_visibility",
+    "tsl_copts",
+    "tsl_gpu_library",
+)
+
+cc_library(
+    name = "common_headers",
+    hdrs = [
+        "call_stack.hpp",
+        "defines.hpp",
+        "filesystem.hpp",
+        "name_info.hpp",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//xla/stream_executor/rocm:roctracer_wrapper"]
+)

xla/backends/profiler/gpu/common/call_stack.hpp

@@ -0,0 +1,97 @@
+// MIT License
+//
+// Copyright (c) 2023 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 "filesystem.hpp"
+
+#include <cstdint>
+#include <fstream>
+#include <functional>
+#include <iomanip>
+#include <iostream>
+#include <string>
+#include <vector>
+
+#include "tsl/platform/env.h"
+#include "tsl/platform/errors.h"
+#include "tsl/platform/logging.h"
+#include "tsl/platform/macros.h"
+
+
+namespace xla {
+namespace common {
+struct source_location
+{
+    std::string function = {};
+    std::string file     = {};
+    uint32_t    line     = 0;
+    std::string context  = {};
+};
+
+using call_stack_t = std::vector<source_location>;
+
+inline void
+print_call_stack(std::string         ofname,
+                 const call_stack_t& _call_stack,
+                 const char*         env_variable = "ROCPROFILER_SAMPLE_OUTPUT_FILE")
+{
+    if(auto* eofname = getenv(env_variable)) ofname = eofname;
+
+    std::ostream* ofs     = nullptr;
+    auto          cleanup = std::function<void(std::ostream*&)>{};
+
+    if(ofname == "stdout")
+        ofs = &std::cout;
+    else if(ofname == "stderr")
+        ofs = &std::cerr;
+    else
+    {
+        ofs = new std::ofstream{ofname};
+        if(ofs && *ofs)
+            cleanup = [](std::ostream*& _os) { delete _os; };
+        else
+        {
+            std::cerr << "Error outputting to " << ofname << ". Redirecting to stderr...\n";
+            ofname = "stderr";
+            ofs    = &std::cerr;
+        }
+    }
+
+    // LOG(INFO) << "Outputting collected data to " << ofname;
+
+    size_t n = 0;
+    for(const auto& itr : _call_stack)
+    {
+        *ofs << std::left << std::setw(2) << ++n << "/" << std::setw(2) << _call_stack.size()
+             << " [" << common::fs::path{itr.file}.filename() << ":" << itr.line << "] "
+             << std::setw(20) << itr.function;
+        if(!itr.context.empty()) *ofs << " :: " << itr.context;
+        *ofs << "\n";
+    }
+
+    *ofs << std::flush;
+
+    if(cleanup) cleanup(ofs);
+}
+}  // namespace common
+}  // namespace xla

xla/backends/profiler/gpu/common/defines.hpp

@@ -0,0 +1,83 @@
+// MIT License
+//
+// Copyright (c) 2023 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 "xla/stream_executor/rocm/roctracer_wrapper.h"
+namespace se = ::stream_executor;
+
+#define ROCPROFILER_VAR_NAME_COMBINE(X, Y) X##Y
+#define ROCPROFILER_VARIABLE(X, Y)         ROCPROFILER_VAR_NAME_COMBINE(X, Y)
+
+#define ROCPROFILER_WARN(result)                                                                   \
+    {                                                                                              \
+        rocprofiler_status_t ROCPROFILER_VARIABLE(CHECKSTATUS, __LINE__) = result;                 \
+        if(ROCPROFILER_VARIABLE(CHECKSTATUS, __LINE__) != ROCPROFILER_STATUS_SUCCESS)              \
+        {                                                                                          \
+            std::string status_msg =                                                               \
+                se::wrap::rocprofiler_get_status_string(ROCPROFILER_VARIABLE(CHECKSTATUS, __LINE__));        \
+            std::cerr << "[" << __FILE__ << ":" << __LINE__ << "] " << #result                     \
+                      << " returned error code " << ROCPROFILER_VARIABLE(CHECKSTATUS, __LINE__)    \
+                      << ": " << status_msg << ". This is just a warning!" << std::endl;           \
+        }                                                                                          \
+    }
+
+#define ROCPROFILER_CHECK(result)                                                                  \
+    {                                                                                              \
+        rocprofiler_status_t ROCPROFILER_VARIABLE(CHECKSTATUS, __LINE__) = result;                 \
+        if(ROCPROFILER_VARIABLE(CHECKSTATUS, __LINE__) != ROCPROFILER_STATUS_SUCCESS)              \
+        {                                                                                          \
+            std::string status_msg =                                                               \
+                se::wrap::rocprofiler_get_status_string(ROCPROFILER_VARIABLE(CHECKSTATUS, __LINE__));        \
+            std::stringstream errmsg{};                                                            \
+            errmsg << "[" << __FILE__ << ":" << __LINE__ << "] " << #result                        \
+                   << " failed with error code " << ROCPROFILER_VARIABLE(CHECKSTATUS, __LINE__)    \
+                   << " :: " << status_msg;                                                        \
+            throw std::runtime_error(errmsg.str());                                                \
+        }                                                                                          \
+    }
+
+#define ROCPROFILER_CALL(result, msg)                                                              \
+    {                                                                                              \
+        rocprofiler_status_t ROCPROFILER_VARIABLE(CHECKSTATUS, __LINE__) = result;                 \
+        if(ROCPROFILER_VARIABLE(CHECKSTATUS, __LINE__) != ROCPROFILER_STATUS_SUCCESS)              \
+        {                                                                                          \
+            std::string status_msg =                                                               \
+                se::wrap::rocprofiler_get_status_string(ROCPROFILER_VARIABLE(CHECKSTATUS, __LINE__));        \
+            std::cerr << "[" #result "][" << __FILE__ << ":" << __LINE__ << "] " << msg            \
+                      << " failed with error code " << ROCPROFILER_VARIABLE(CHECKSTATUS, __LINE__) \
+                      << ": " << status_msg << std::endl;                                          \
+            std::stringstream errmsg{};                                                            \
+            errmsg << "[" #result "][" << __FILE__ << ":" << __LINE__ << "] " << msg " failure ("  \
+                   << status_msg << ")";                                                           \
+        }                                                                                          \
+    }
+
+#if HIP_VERSION >= 60300000
+#    define HIP_HOST_ALLOC_FUNC hipExtHostAlloc
+#    define HIP_HOST_FREE_FUNC  hipFreeHost
+#else
+#    define HIP_HOST_ALLOC_FUNC hipHostMalloc
+#    define HIP_HOST_FREE_FUNC  hipHostFree
+#endif
+
+// throw std::runtime_error(errmsg.str());

xla/backends/profiler/gpu/common/filesystem.hpp

@@ -0,0 +1,78 @@
+// MIT License
+//
+// Copyright (c) 2023 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
+
+#if !defined(ROCPROFILER_SAMPLES_HAS_GHC_LIB_FILESYSTEM)
+#    if defined __has_include
+#        if __has_include(<ghc/filesystem.hpp>)
+#            define ROCPROFILER_SAMPLES_HAS_GHC_LIB_FILESYSTEM 1
+#        else
+#            define ROCPROFILER_SAMPLES_HAS_GHC_LIB_FILESYSTEM 0
+#        endif
+#    else
+#        define ROCPROFILER_SAMPLES_HAS_GHC_LIB_FILESYSTEM 0
+#    endif
+#endif
+
+#if ROCPROFILER_SAMPLES_HAS_GHC_LIB_FILESYSTEM == 0
+#    if defined __has_include
+#        if __has_include(<version>)
+#            include <version>
+#        endif
+#    endif
+
+#    if defined(__cpp_lib_filesystem)
+#        define ROCPROFILER_SAMPLES_HAS_CPP_LIB_FILESYSTEM 1
+#    else
+#        if defined __has_include
+#            if __has_include(<filesystem>)
+#                define ROCPROFILER_SAMPLES_HAS_CPP_LIB_FILESYSTEM 1
+#            endif
+#        endif
+#    endif
+#endif
+
+// include the correct filesystem header
+#if defined(ROCPROFILER_SAMPLES_HAS_GHC_LIB_FILESYSTEM) &&                                         \
+    ROCPROFILER_SAMPLES_HAS_GHC_LIB_FILESYSTEM > 0
+#    include <ghc/filesystem.hpp>
+#elif defined(ROCPROFILER_SAMPLES_HAS_CPP_LIB_FILESYSTEM) &&                                       \
+    ROCPROFILER_SAMPLES_HAS_CPP_LIB_FILESYSTEM > 0
+#    include <filesystem>
+#else
+#    include <experimental/filesystem>
+#endif
+
+namespace xla {
+namespace common {
+#if defined(ROCPROFILER_SAMPLES_HAS_GHC_LIB_FILESYSTEM) &&                                         \
+    ROCPROFILER_SAMPLES_HAS_GHC_LIB_FILESYSTEM > 0
+namespace fs = ::ghc::filesystem;  // NOLINT(misc-unused-alias-decls)
+#elif defined(ROCPROFILER_SAMPLES_HAS_CPP_LIB_FILESYSTEM) &&                                       \
+    ROCPROFILER_SAMPLES_HAS_CPP_LIB_FILESYSTEM > 0
+namespace fs = ::std::filesystem;  // NOLINT(misc-unused-alias-decls)
+#else
+namespace fs = ::std::experimental::filesystem;  // NOLINT(misc-unused-alias-decls)
+#endif
+}  // namespace common
+}  // namespace xla

xla/backends/profiler/gpu/common/name_info.hpp

@@ -0,0 +1,54 @@
+// MIT License
+//
+// Copyright (c) 2023 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 "xla/stream_executor/rocm/roctracer_wrapper.h"
+
+#include "defines.hpp"
+
+#include <iostream>
+#include <map>
+#include <sstream>
+#include <stdexcept>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+
+namespace xla {
+namespace common {
+using callback_name_info = rocprofiler::sdk::callback_name_info;
+using buffer_name_info   = rocprofiler::sdk::buffer_name_info;
+
+inline auto
+get_buffer_tracing_names()
+{
+    return rocprofiler::sdk::get_buffer_tracing_names();
+}
+
+inline auto
+get_callback_id_names()
+{
+    return rocprofiler::sdk::get_callback_tracing_names();
+}
+}  // namespace common
+}  // namespace xla