Add explicit MIT license and copyright

[deliverable/exatracer.git] / src / lttng-ust-exatracer.cpp

// SPDX-FileCopyrightText: 2024 Advanced Micro Devices, Inc. All rights reserved.
// SPDX-FileCopyrightText: 2024 EfficiOS, Inc.
//
// SPDX-License-Identifier: MIT

#include <cinttypes>

#include <hip/hip_runtime.h>
#include <hip/hip_gl_interop.h>

/*
 * Do not re-order!  hip_api_trace.h header referenced some R0000 types which
 * are only defined in the hip_deprecated.h header.
 */
#include <hip/hip_deprecated.h>
#include <hip/amd_detail/hip_api_trace.hpp>

/*
 * Without AMD_INTERNAL_BUILD, hsa_api_trace.h tries to include
 * "inc/*", which is not installed.
 */
#define AMD_INTERNAL_BUILD
#include <hsa/hsa.h>
#include <hsa/hsa_api_trace.h>

/*
 * Rocprofiler integration.
 */
#include <rocprofiler-sdk/registration.h>
#include <rocprofiler-sdk/rocprofiler.h>
#include <rocprofiler-sdk-roctx/api_trace.h>

/*
 * This is the table with LTTng tracepoints for HIP.
 */
static struct HipDispatchTable lttng_hip_table;

/*
 * This is an internal copy of the orignal table used by the LTTng wrappers to
 * call the next function for HIP.
 */
static struct HipDispatchTable next_hip_table;

/* Auto-generated by gen-hip-wrappers. */
#include "lttng-ust-hip-wrappers.cpp"


/*
 * This is the table with LTTng tracepoints for HSA.
 */
static struct CoreApiTable lttng_hsa_core_table;

/*
 * This is an internal copy of the orignal table used by the LTTng wrappers to
 * call the next function for HSA.
 */
static struct CoreApiTable next_hsa_core_table;

/* Auto-generated by gen-hsa-wrappers. */
#include "lttng-ust-hsa-wrappers.cpp"


/*
 * This is the table with LTTng tracepoints for roctx.
 */
static roctxCoreApiTable_t lttng_roctx_core_table;

/*
 * This is an internal copy of the orignal table used by the LTTng wrappers to
 * call the next function for roctx.
 */
static roctxCoreApiTable_t next_roctx_core_table;

/* LTTng events for roctx. */
#include "lttng-roctx.c"

#define die(FMT, ...)                                           \
    do {                                                        \
        fprintf(stderr, FMT "\n" __VA_OPT__(,) __VA_ARGS__);    \
        exit(EXIT_FAILURE);                                     \
    } while(0)

#define error(FMT, ...)                                 \
    fprintf(stderr, FMT "\n" __VA_OPT__(,) __VA_ARGS__)

static void
register_hip_table(uint64_t lib_instance, uint64_t num_tables, void **tables)
{
    /* Add instrumentation only to first library instance. */
    if (0 != lib_instance) {
        return;
    }

    /* Ensure that there is at least one table. */
    if (num_tables < 1) {
        return;
    }

    /*
     * This could be done at compile time in some way if C++ could support
     * designated initializers.
     */
    lttng_hip_install_wrappers();

    auto original_hip_table = static_cast<HipDispatchTable*>(tables[0]);

    /* Swap tables. */
    if (original_hip_table) {
        auto size = std::min(sizeof(next_hip_table), original_hip_table->size);
        memcpy(&next_hip_table, original_hip_table, size);
        memcpy(original_hip_table, &lttng_hip_table, size);
    }
}

static void
register_hsa_core_table(uint64_t lib_instance, uint64_t num_tables, void **tables)
{
    /* Add instrumentation only to first library instance. */
    if (0 != lib_instance) {
        return;
    }

    /* Ensure that there is at least one table. */
    if (num_tables < 1) {
        return;
    }

    /*
     * This could be done at compile time in some way if C++ could support
     * designated initializers.
     */
    lttng_hsa_install_wrappers();

    auto original_hsa_table = static_cast<HsaApiTable*>(tables[0]);

    /* Swap tables. */
    if (original_hsa_table) {
        auto core = original_hsa_table->core_;
        auto size = sizeof(*core);
        memcpy(&next_hsa_core_table, core, size);
        memcpy(core, &lttng_hsa_core_table, size);
    }
}

static void
register_roctx_core_table(uint64_t lib_instance, uint64_t num_tables, void **tables)
{
    /* Add instrumentation only to first library instance. */
    if (0 != lib_instance) {
        return;
    }

    /* Ensure that there is at least one table. */
    if (num_tables < 1) {
        return;
    }

    /*
     * This could be done at compile time in some way if C++ could support
     * designated initializers.
     */
    lttng_roctx_install_wrappers();

    auto original_roctx_core_table = static_cast<roctxCoreApiTable_t*>(tables[0]);

    /* Swap tables. */
    if (original_roctx_core_table) {
        auto size = std::min(original_roctx_core_table->size, sizeof(roctxCoreApiTable_t));
        memcpy(&next_roctx_core_table, original_roctx_core_table, size);
        memcpy(original_roctx_core_table, &lttng_roctx_core_table, size);
    }
}

static void
api_registration_callback(rocprofiler_intercept_table_t type,
                          uint64_t lib_version,
                          uint64_t lib_instance,
                          void**   tables,
                          uint64_t num_tables,
                          void*    user_data)
{
    (void) user_data;
    (void) lib_version;

    /*
     * We only want HIP, HSA or ROC-TX tables.  If we get something else, there
     * is a bug somewhere.
     */
    switch (type) {
    case ROCPROFILER_HIP_RUNTIME_TABLE:
        register_hip_table(lib_instance, num_tables, tables);
        break;
    case ROCPROFILER_HSA_TABLE:
        register_hsa_core_table(lib_instance, num_tables, tables);
        break;
    case ROCPROFILER_MARKER_CORE_TABLE:
        register_roctx_core_table(lib_instance, num_tables, tables);
        break;
    defualt:
        die("unexpected library type: %d", type);
    }
}

extern "C"
__attribute__((visibility("default")))
rocprofiler_tool_configure_result_t *rocprofiler_configure(uint32_t version,
							   const char *runtime_version,
							   uint32_t priority,
							   rocprofiler_client_id_t *id)
{
    (void) priority;
    (void) runtime_version;
    (void) version;

    /*
     * Maybe Exatracer instead?
     */
    id->name = "LTTng-UST";

    if (ROCPROFILER_STATUS_ERROR_NOT_IMPLEMENTED ==
        rocprofiler_at_intercept_table_registration(api_registration_callback,
                                                    ROCPROFILER_HIP_RUNTIME_TABLE |
						    ROCPROFILER_HSA_TABLE |
						    ROCPROFILER_MARKER_CORE_TABLE,
                                                    nullptr)) {
        die("Trying to register API interception : " "NOT IMPLEMENTED");
    }

    static auto cfg = rocprofiler_tool_configure_result_t {
        sizeof(rocprofiler_tool_configure_result_t),
        nullptr,
        nullptr,
        nullptr
    };

    return &cfg;
}