Indicate that lttng-ust 2.13+ is required

[deliverable/lttng-ust-mpi.git] / lttng-auto-mpi-wrappers

#!/usr/bin/env python3
#
# SPDX-License-Identifier: MIT
#
# SPDX-FileCopyrightText: 2023 EfficiOS, Inc.
#
# Author: Olivier Dion <odion@efficios.com>
#
# Auto-generate lttng-ust tracepoints for OpenMPI.
#
# Require: python-clang (libclang)

import argparse
import os
import re

from string import Template

import clang.cindex

def list_function_declarations(root):
    return [ child
             for child in root.get_children()
             if child.kind == clang.cindex.CursorKind.FUNCTION_DECL ]

def parse_header(header_file):
    return clang.cindex.Index.create().parse(header_file).cursor

def list_functions(root):
    return [
        fn
        for fn in list_function_declarations(root)
        if fn.spelling.startswith("MPI_") and fn.spelling
    ]

def exact_definition(arg):
    m = re.search(r'(\[[0-9]*\])+', arg.type.spelling)
    if m:
        return f"{arg.type.spelling[:m.start(0)]} {arg.spelling}{m.group(0)}"
    else:
        return f"{arg.type.spelling} {arg.spelling}"

forbiden_list = {
    "MPI_Pcontrol"
}

extra_works = {
    "MPI_Init": """
	if (MPI_SUCCESS == ret) {
		int (*mpi_comm_rank)(MPI_Comm, int *rank);
		MPI_Comm mpi_comm_world;
#ifdef CRAY_MPICH_VERSION
		mpi_comm_world = MPI_COMM_WORLD;
#else
		mpi_comm_world = *(void**)resolve_or_die("ompi_mpi_comm_world_addr");
#endif
		mpi_comm_rank  = resolve_or_die("PMPI_Comm_rank");
		mpi_comm_rank(mpi_comm_world, &mpi_rank);
		mpi_provider.priv = lttng_ust_context_provider_register(&mpi_provider);
	}
""",
    "MPI_Finalize": """
	if (mpi_provider.priv) {
		lttng_ust_context_provider_unregister(mpi_provider.priv);
	}
""",
}

def main():

    if os.getenv("LTTNG_UST_MPI_CLANG_LIBRARY_FILE", None) is not None:
        clang.cindex.Config.set_library_file(os.getenv("LTTNG_UST_MPI_CLANG_LIBRARY_FILE"))

    parser = argparse.ArgumentParser(prog="lttng-ust-auto-mpi")

    parser.add_argument("api",
                        help="MPI API header")

    parser.add_argument("wrappers",
                        help="Path to MPI wrappers")

    args = parser.parse_args()

    fn_tpl = Template("""
${ret_type} ${fn_name}(${fn_arguments})
{
	${ret_type} ret;
	{
		static ${ret_type}(*real_fn)(${fn_arguments}) = NULL;
		if (unlikely(NULL == __atomic_load_n(&real_fn, __ATOMIC_RELAXED))) {
			void *result = resolve_or_die("P${fn_name}");
			__atomic_store_n(&real_fn, result, __ATOMIC_RELAXED);
		}
		LTTNG_MAKE_API_OBJECT(${fn_name}${fn_rest_argument_names});
		ret = real_fn(${fn_pass_argument_names});
		LTTNG_MARK_RETURN_API_OBJECT(ret);
	}
$extra_work
	return ret;
}
""")

    with open(args.wrappers, "w") as output:
        output.write("""/* Auto-generated */
#define _GNU_SOURCE

#include <stdlib.h>
#include <string.h>

#include <dlfcn.h>

#include <mpi.h>

#include <lttng/ust-events.h>
#include <lttng/ust-ringbuffer-context.h>

#include "lttng/ust-context-provider.h"

#include "lttng-ust-mpi-states.h"

#define likely(x)   __builtin_expect(!!(x), 1)
#define unlikely(x) __builtin_expect(!!(x), 0)

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

static void *resolve_or_die(const char *symbol)
{
	void *ret = dlsym(RTLD_NEXT, symbol);
	if (unlikely(!ret)) {
		die("could not resolve `%s': %s", symbol, dlerror());
	}
	return ret;
}

static inline int streq(const char *A, const char *B)
{
	return 0 == strcmp(A, B);
}

static inline char *context_type(struct lttng_ust_app_context *uctx)
{
	char *suffix = index(uctx->ctx_name, ':');

	if (likely(suffix)) {
		suffix = &suffix[1]; /* Skip ':' */
	}

	return suffix;
}

static int mpi_rank = -1;

static size_t mpi_provider_get_size(void *uctx,
				    struct lttng_ust_probe_ctx *probe_ctx __attribute__((unused)),
				    size_t offset)
{
	size_t size = 0;
	char *type = context_type(uctx);

	size += lttng_ust_ring_buffer_align(offset, lttng_ust_rb_alignof(char));
	size += sizeof(char);

	if (unlikely(!type)) {
		goto error;
	}

	if (streq(type, "rank")) {
		size += lttng_ust_ring_buffer_align(offset, lttng_ust_rb_alignof(int64_t));
		size += sizeof(int64_t);

	} else {
	error:
		/* Unknown context. */
		(void) size;
	}

	return size;
}

static void mpi_provider_record(void *uctx,
				struct lttng_ust_probe_ctx *probe_ctx __attribute__((unused)),
				struct lttng_ust_ring_buffer_ctx *ctx,
				struct lttng_ust_channel_buffer *lttng_chan_buf)
{
	int sel;
	char sel_char;
	char *type = context_type(uctx);

	if (unlikely(!type)) {
		goto error;
	}

	if (streq(type, "rank")) {
		int64_t v;
		sel      = LTTNG_UST_DYNAMIC_TYPE_S64;
		sel_char = (char) sel;
		v        = (int64_t) mpi_rank;
		lttng_chan_buf->ops->event_write(ctx, &sel_char, sizeof(sel_char),
						 lttng_ust_rb_alignof(char));

		lttng_chan_buf->ops->event_write(ctx, &v, sizeof(v), lttng_ust_rb_alignof(v));
	} else {
	error:
		sel	 = LTTNG_UST_DYNAMIC_TYPE_NONE;
		sel_char = (char) sel;
		lttng_chan_buf->ops->event_write(ctx, &sel_char, sizeof(sel_char),
						 lttng_ust_rb_alignof(char));
	}
}

static void mpi_provider_get_value(void *uctx,
				   struct lttng_ust_probe_ctx *probe_ctx __attribute__((unused)),
				   struct lttng_ust_ctx_value *value)
{
	char *type = context_type(uctx);

	if (unlikely(!type)) {
		goto error;
	}

	if (streq(type, "rank")) {
		value->sel   = LTTNG_UST_DYNAMIC_TYPE_S64;
		value->u.s64 = (int64_t) mpi_rank;
	} else {
	error:
		value->sel = LTTNG_UST_DYNAMIC_TYPE_NONE;
	}
}

static struct lttng_ust_context_provider mpi_provider = {
	.struct_size = sizeof(struct lttng_ust_context_provider),
	.name        = "$app.MPI",
	.get_size    = mpi_provider_get_size,
	.record      = mpi_provider_record,
	.get_value   = mpi_provider_get_value
};
""")
        for fn in list_functions(parse_header(args.api)):

            if fn.spelling in forbiden_list:
                continue

            args = list(fn.get_arguments())
            fn_pass_argument_names = ", ".join([
                f"{arg.spelling}"
                for arg in args
            ])

            if args:
                fn_rest_argument_names = ", " + ", ".join([
                    "(%s)%s" % (re.sub(r'\[[0-9]*\]', '*', arg.type.spelling),
                                arg.spelling)
                    for arg in args
                ])
            else:
                fn_rest_argument_names=""

            if fn.spelling in extra_works:
                extra_work = extra_works[fn.spelling]
            else:
                extra_work = ""

            output.write(fn_tpl.substitute(fn_name=fn.spelling,
                                           fn_arguments=", ".join([
                                               exact_definition(arg)
                                               for arg in fn.get_arguments()
                                           ]),
                                           fn_pass_argument_names=fn_pass_argument_names,
                                           fn_rest_argument_names=fn_rest_argument_names,
                                           ret_type=fn.type.get_result().spelling,
                                           extra_work=extra_work))

if __name__ == "__main__":
    main()