Add active_forwards to lbt_libinfo_t

This adds a new field to the `lbt_libinfo_t` structure that allows
determining which BLAS/LAPACK library calls are being satisfied by which
loaded libraries.  By inspecting the `active_forwards` field it will be
possible to, for example, walk over all loaded libraries and determine
which one the `dgemm_64_` call is provided by, which may be useful for
applications that, for instance, tune threading thresholds based on
which BLAS provider is being used.

Author: staticfloat

src/config.c

@@ -31,23 +31,56 @@ void clear_loaded_libraries() {
     for (int idx=0; idx<MAX_TRACKED_LIBS; ++idx) {
         if (lbt_config.loaded_libs[idx] != NULL) {
             free(lbt_config.loaded_libs[idx]->libname);
+            free(lbt_config.loaded_libs[idx]->active_forwards);
             close_library(lbt_config.loaded_libs[idx]->handle);
             free(lbt_config.loaded_libs[idx]);
             lbt_config.loaded_libs[idx] = NULL;
         }
     }
 }
 
-void record_library_load(const char * libname, void * handle, const char * suffix, int interface, int f2c) {
+void clear_forwarding_mark(int32_t symbol_idx, int32_t interface) {
+    for (int idx=0; idx<MAX_TRACKED_LIBS; ++idx) {
+        if (lbt_config.loaded_libs[idx] == NULL) {
+            return;
+        }
+        if (lbt_config.loaded_libs[idx]->interface != interface) {
+            continue;
+        }
+
+        BITFIELD_CLEAR(lbt_config.loaded_libs[idx]->active_forwards, symbol_idx);
+    }
+}
+
+void clear_other_forwards(int skip_idx, uint8_t * forwards, int32_t interface) {
+    for (int idx=0; idx<MAX_TRACKED_LIBS; ++idx) {
+        if (lbt_config.loaded_libs[idx] == NULL) {
+            return;
+        }
+        // Skip ourselves and things that don't match our interface
+        if (idx == skip_idx || lbt_config.loaded_libs[idx]->interface != interface) {
+            continue;
+        }
+
+        // Flip off anything in this library that is flipped on in the passed-in forwards
+        for (uint32_t chunk_idx=0; chunk_idx < (NUM_EXPORTED_FUNCS/8)+1; ++chunk_idx) {
+            lbt_config.loaded_libs[idx]->active_forwards[chunk_idx] &= (forwards[chunk_idx] ^ 0xff);
+        }
+    }
+}
+
+void record_library_load(const char * libname, void * handle, const char * suffix, uint8_t * forwards, int interface, int f2c) {
     // Scan for the an empty slot, and also check to see if this library has been loaded before.
     int free_idx = -1;
     for (int idx=0; idx<MAX_TRACKED_LIBS; ++idx) {
         if (lbt_config.loaded_libs[idx] == NULL) {
             free_idx = idx;
             break;
         }
-        // If this library has been loaded before, just early-exit.
+        // If this library has been loaded before, all we do is copy the `forwards` over
         if (lbt_config.loaded_libs[idx]->handle == handle) {
+            memcpy(lbt_config.loaded_libs[idx]->active_forwards, forwards, sizeof(uint8_t)*(NUM_EXPORTED_FUNCS/8 + 1));
+            clear_other_forwards(idx, forwards, interface);
             return;
         }
     }
@@ -63,8 +96,13 @@ void record_library_load(const char * libname, void * handle, const char * suffi
     memcpy(new_libinfo->libname, libname, namelen);
     new_libinfo->handle = handle;
     new_libinfo->suffix = suffix;
+    new_libinfo->active_forwards = (uint8_t *)malloc(sizeof(uint8_t)*(NUM_EXPORTED_FUNCS/8 + 1));
+    memcpy(new_libinfo->active_forwards, forwards, sizeof(uint8_t)*(NUM_EXPORTED_FUNCS/8 + 1));
     new_libinfo->interface = interface;
     new_libinfo->f2c = f2c;
 
     lbt_config.loaded_libs[free_idx] = new_libinfo;
+
+    // Next, we go through and un-set all other loaded libraries of the same interface's `forwards`:
+    clear_other_forwards(free_idx, forwards, interface);
 }

src/libblastrampoline.c

@@ -135,7 +135,13 @@ LBT_DLLEXPORT int32_t lbt_set_forward(const char * symbol_name, const void * add
     if (symbol_idx == -1)
         return -1;
 
-    return set_forward_by_index(symbol_idx, addr, interface, f2c, verbose);
+    int32_t ret = set_forward_by_index(symbol_idx, addr, interface, f2c, verbose);
+    if (ret == 0) {
+        // Un-mark this symbol as being provided by any of our libraries;
+        // if you use the footgun API, you can keep track of who is providing what.
+        clear_forwarding_mark(symbol_idx, interface);
+    }
+    return ret;
 }
 
 /*
@@ -262,6 +268,7 @@ LBT_DLLEXPORT int32_t lbt_forward(const char * libname, int32_t clear, int32_t v
     int32_t nforwards = 0;
     int32_t symbol_idx = 0;
     char symbol_name[MAX_SYMBOL_LEN];
+    uint8_t forwards[(NUM_EXPORTED_FUNCS/8) + 1] = {0};
     for (symbol_idx=0; exported_func_names[symbol_idx] != NULL; ++symbol_idx) {
         // If `clear` is set, zero out all symbols that may have been set so far
         if (clear) {
@@ -274,11 +281,12 @@ LBT_DLLEXPORT int32_t lbt_forward(const char * libname, int32_t clear, int32_t v
         void *addr = lookup_symbol(handle, symbol_name);
         if (addr != NULL) {
             set_forward_by_index(symbol_idx,  addr, interface, f2c, verbose);
+            BITFIELD_SET(forwards, symbol_idx);
             nforwards++;
         }
     }
 
-    record_library_load(libname, handle, lib_suffix, interface, f2c);
+    record_library_load(libname, handle, lib_suffix, &forwards[0], interface, f2c);
     if (verbose) {
         printf("Processed %d symbols; forwarded %d symbols with %d-bit interface and mangling to a suffix of \"%s\"\n", symbol_idx, nforwards, interface, lib_suffix);
     }

src/libblastrampoline.h

@@ -34,6 +34,13 @@ extern "C" {
 # define LBT_HIDDEN    __attribute__ ((visibility("hidden")))
 #endif
 
+#define BF_CHUNK(array, idx)           (array[((uint32_t)(idx/8))])
+#define BF_MASK(idx)                   ((uint8_t)(0x1 << (idx % 8)))
+#define BITFIELD_GET(array, idx)      ((BF_CHUNK(array, idx) &  BF_MASK(idx)) >> (idx % 8))
+#define BITFIELD_CLEAR(array, idx)      BF_CHUNK(array, idx) &= ~(BF_MASK(idx))
+#define BITFIELD_SET(array, idx)        BF_CHUNK(array, idx) |=   BF_MASK(idx)
+
+
 // The metadata stored on each loaded library
 typedef struct {
     // The library name as passed to `lbt_forward()`.
@@ -43,6 +50,11 @@ typedef struct {
     // The suffix used within this library as autodetected by `lbt_forward`.
     // Common values are `""` or `"64_"`.
     const char * suffix;
+    // bitfield (in uint8_t form) representing the active forwards for this library.
+    // Use the `BITFIELD_{SET,GET}` macros to look at particular indices within this field.
+    // Note that if you use the footgun API (e.g. "lbt_set_forward()") these values will be
+    // zeroed out and you must track them manually if you need to.
+    uint8_t * active_forwards;
     // The interface type  as autodetected by `lbt_forward`, see `LBT_INTERFACE_XXX` below
     int32_t interface;
     // The `f2c` status  as autodetected by `lbt_forward`, see `LBT_F2C_XXX` below

src/libblastrampoline_internal.h

@@ -52,7 +52,8 @@ int32_t find_symbol_idx(const char * name);
 // Functions in `config.c`
 void init_config();
 void clear_loaded_libraries();
-void record_library_load(const char * libname, void * handle, const char * suffix, int interface, int f2c);
+void clear_forwarding_mark(int32_t symbol_idx, int32_t interface);
+void record_library_load(const char * libname, void * handle, const char * suffix, uint8_t * forwards, int interface, int f2c);
 
 // Functions in `win_utils.c`
 #ifdef _OS_WINDOWS_

test/direct.jl

@@ -7,14 +7,27 @@ function unpack_loaded_libraries(config::lbt_config_t)
     idx = 1
     lib_ptr = unsafe_load(config.loaded_libs, idx)
     while lib_ptr != C_NULL
-        push!(libs, LBTLibraryInfo(unsafe_load(lib_ptr)))
+        push!(libs, LBTLibraryInfo(unsafe_load(lib_ptr), config.num_exported_symbols))
 
         idx += 1
         lib_ptr = unsafe_load(config.loaded_libs, idx)
     end
     return libs
 end
 
+function find_symbol_offset(config::lbt_config_t, symbol::String)
+    for sym_idx in 1:config.num_exported_symbols
+        if unsafe_string(unsafe_load(config.exported_symbols, sym_idx)) == symbol
+            return UInt32(sym_idx - 1)
+        end
+    end
+    return nothing
+end
+
+function bitfield_get(field::Vector{UInt8}, symbol_idx::UInt32)
+    return field[div(symbol_idx,8)+1] & (UInt8(0x01) << (symbol_idx%8))
+end
+
 lbt_prefix = get_blastrampoline_dir()
 lbt_handle = dlopen("$(lbt_prefix)/$(binlib)/libblastrampoline.$(shlib_ext)", RTLD_GLOBAL | RTLD_DEEPBIND)
 
@@ -38,6 +51,10 @@ lbt_handle = dlopen("$(lbt_prefix)/$(binlib)/libblastrampoline.$(shlib_ext)", RT
         @test (config.build_flags & LBT_BUILDFLAGS_F2C_CAPABLE) != 0
     end
 
+    # Check to make sure that `dgemm_` is part of the exported symbols:
+    dgemm_idx = find_symbol_offset(config, "dgemm_")
+    @test dgemm_idx !== nothing
+
     # Walk the libraries and check we have two
     libs = unpack_loaded_libraries(config)
     @test length(libs) == 2
@@ -52,13 +69,22 @@ lbt_handle = dlopen("$(lbt_prefix)/$(binlib)/libblastrampoline.$(shlib_ext)", RT
         @test libs[1].interface == LBT_INTERFACE_LP64
     end
     @test libs[1].f2c == LBT_F2C_PLAIN
+    @test bitfield_get(libs[1].active_forwards, dgemm_idx) != 0
 
     # Next check OpenBLAS32_jll which is always LP64
     @test libs[2].libname == OpenBLAS32_jll.libopenblas_path
     @test libs[2].suffix == ""
     @test libs[2].interface == LBT_INTERFACE_LP64
     @test libs[2].f2c == LBT_F2C_PLAIN
 
+    # If OpenBLAS32 and OpenBLAS are the same interface (e.g. i686)
+    # then libs[2].active_forwards should be all zero!
+    if libs[1].interface == libs[2].interface
+        @test bitfield_get(libs[2].active_forwards, dgemm_idx) == 0
+    else
+        @test bitfield_get(libs[2].active_forwards, dgemm_idx) != 0
+    end
+
     # Load OpenBLAS32_jll again, but this time clearing it and ensure the config gets cleared too
     lbt_forward(lbt_handle, OpenBLAS32_jll.libopenblas_path; clear=true)
     config = lbt_get_config(lbt_handle)
@@ -68,7 +94,6 @@ lbt_handle = dlopen("$(lbt_prefix)/$(binlib)/libblastrampoline.$(shlib_ext)", RT
     @test libs[1].suffix == ""
     @test libs[1].interface == LBT_INTERFACE_LP64
     @test libs[1].f2c == LBT_F2C_PLAIN
-
 end
 
 @testset "get/set threads" begin
@@ -119,11 +144,25 @@ end
     @test slamch_32 != C_NULL
     @test lbt_get_forward(lbt_handle, "slamch_", LBT_INTERFACE_LP64) == slamch_32
 
+    # Ensure that the libs show that `slamch_` is forwarded by this library:
+    config = lbt_get_config(lbt_handle)
+    libs = unpack_loaded_libraries(config)
+    @test length(libs) == 1
+
+    slamch_idx = find_symbol_offset(config, "slamch_")
+    @test slamch_idx !== nothing
+    @test bitfield_get(libs[1].active_forwards, slamch_idx) != 0
+    orig_forwards = copy(libs[1].active_forwards)
+
     # Now, test that we can muck this up
     my_slamch = @cfunction(record_slamch_args, Float32, (Cstring,))
     @test lbt_set_forward(lbt_handle, "slamch_", my_slamch, LBT_INTERFACE_LP64) == 0
     @test lbt_get_forward(lbt_handle, "slamch_", LBT_INTERFACE_LP64) == my_slamch
 
+    config = lbt_get_config(lbt_handle)
+    libs = unpack_loaded_libraries(config)
+    @test bitfield_get(libs[1].active_forwards, slamch_idx) == 0
+
     # Ensure that we actually overrode the symbol
     @test ccall(dlsym(lbt_handle, "slamch_"), Float32, (Cstring,), "test") == 13.37f0
     @test slamch_args == ["test"]

test/utils.jl

@@ -76,17 +76,19 @@ struct lbt_library_info_t
     libname::Cstring
     handle::Ptr{Cvoid}
     suffix::Cstring
+    active_forwards::Ptr{UInt8}
     interface::Int32
     f2c::Int32
 end
 struct LBTLibraryInfo
     libname::String
     handle::Ptr{Cvoid}
     suffix::String
+    active_forwards::Vector{UInt8}
     interface::Int32
     f2c::Int32
 
-    LBTLibraryInfo(x::lbt_library_info_t) = new(unsafe_string(x.libname), x.handle, unsafe_string(x.suffix), x.interface, x.f2c)
+    LBTLibraryInfo(x::lbt_library_info_t, num_symbols::UInt32) = new(unsafe_string(x.libname), x.handle, unsafe_string(x.suffix), unsafe_wrap(Vector{UInt8}, x.active_forwards, div(num_symbols,8)+1), x.interface, x.f2c)
 end
 const LBT_INTERFACE_LP64 = 32
 const LBT_INTERFACE_ILP64 = 64