#!/usr/bin/env python3 # # SPDX-License-Identifier: MIT # # SPDX-FileCopyrightText: 2023 EfficiOS, Inc. # # Author: Olivier Dion # # 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 #include #include #include #include #include #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()