instrumentation: include functions with specific calls (#202)

* instrumentation: include functions with specific calls

Add the option `--caller-include <regex>` or environment variable
`OMNITRACE_REGEX_CALLER_INCLUDE` to instrument functions which contain
call to a set of functions, E.g. `--caller-include foo` instruments any
function which calls `foo`.

* Serialize caller include information

* Add test for caller include

* Tweak to the caller include test

- tweak environment
- tweak pass regexes

* Set rewrite caller example to debug

, to avoid optimizing out the call expressions that it relies on.

Co-authored-by: Jonathan R. Madsen <jonathanrmadsen@gmail.com>

examples/CMakeLists.txt

@@ -45,3 +45,4 @@ add_subdirectory(mpi)
 add_subdirectory(python)
 add_subdirectory(lulesh)
 add_subdirectory(rccl)
+add_subdirectory(rewrite-caller)

examples/rewrite-caller/CMakeLists.txt

@@ -0,0 +1,24 @@
+cmake_minimum_required(VERSION 3.15 FATAL_ERROR)
+
+project(omnitrace-rewrite-caller-example LANGUAGES CXX)
+
+if(OMNITRACE_DISABLE_EXAMPLES)
+    get_filename_component(_DIR ${CMAKE_CURRENT_LIST_DIR} NAME)
+
+    if(${PROJECT_NAME} IN_LIST OMNITRACE_DISABLE_EXAMPLES OR ${_DIR} IN_LIST
+                                                             OMNITRACE_DISABLE_EXAMPLES)
+        return()
+    endif()
+endif()
+
+set(CMAKE_BUILD_TYPE "Debug")
+
+add_executable(rewrite-caller rewrite-caller.cpp)
+target_compile_options(rewrite-caller PRIVATE ${_FLAGS})
+
+if(OMNITRACE_INSTALL_EXAMPLES)
+    install(
+        TARGETS rewrite-caller
+        DESTINATION bin
+        COMPONENT omnitrace-examples)
+endif()

examples/rewrite-caller/rewrite-caller.cpp

@@ -0,0 +1,37 @@
+#include <cstdio>
+#include <cstdlib>
+#include <string>
+
+#define NOINLINE __attribute__((noinline))
+
+int num_calls;
+
+void NOINLINE
+inner()
+{
+    ++num_calls;
+}
+
+void NOINLINE
+outer()
+{
+    inner();
+}
+
+int
+main(int argc, char** argv)
+{
+    int ncalls = 10;
+    if(argc > 1) ncalls = atol(argv[1]);
+
+    std::string _name = argv[0];
+    auto        _pos  = _name.find_last_of('/');
+    if(_pos != std::string::npos) _name = _name.substr(_pos + 1);
+
+    for(int i = 0; i < ncalls; ++i)
+    {
+        outer();
+    }
+
+    printf("[%s] number of calls made = %d\n", _name.c_str(), num_calls);
+}

source/bin/omnitrace/fwd.hpp

@@ -208,6 +208,7 @@ extern regexvec_t       file_include;
 extern regexvec_t       file_exclude;
 extern regexvec_t       file_restrict;
 extern regexvec_t       func_restrict;
+extern regexvec_t       caller_include;
 extern CodeCoverageMode coverage_mode;
 
 // logging

source/bin/omnitrace/module_function.cpp

@@ -333,7 +333,7 @@ module_function::is_user_included() const
         }
     }
 
-    return false;
+    return contains_user_callsite();
 }
 
 bool
@@ -534,6 +534,27 @@ module_function::contains_dynamic_callsites() const
     return false;
 }
 
+bool
+module_function::contains_user_callsite() const
+{
+    if(caller_include.empty()) return false;
+
+    bpvector_t<BPatch_point*> call_points;
+    function->getCallPoints(call_points);
+    for(const auto& call_point : call_points)
+    {
+        if(check_regex_restrictions(
+               std::string(get_name(call_point->getCalledFunction())), caller_include))
+        {
+            messages.emplace_back(2, "Forcing", "function", "caller-include-regex",
+                                  function_name);
+            return true;
+        }
+    }
+
+    return false;
+}
+
 bool
 module_function::is_dynamic_callsite_forced() const
 {

source/bin/omnitrace/module_function.hpp

@@ -112,6 +112,7 @@ private:
     bool is_loop_address_range_constrained() const;  // checks loop addr range constraint
     bool contains_dynamic_callsites() const;
     bool should_instrument(bool _coverage) const;
+    bool contains_user_callsite() const;  // checks user caller regexes
 
 public:
     template <typename ArchiveT>
@@ -198,6 +199,7 @@ module_function::serialize(ArchiveT& ar, const unsigned)
            cereal::make_nvp("is_routine_constrained", is_routine_constrained()),
            cereal::make_nvp("is_user_restricted", is_user_restricted()),
            cereal::make_nvp("is_user_included", is_user_included()),
+           cereal::make_nvp("contains_user_callsite", contains_user_callsite()),
            cereal::make_nvp("is_user_excluded", is_user_excluded()),
            cereal::make_nvp("is_overlapping_constrained", is_overlapping_constrained()),
            cereal::make_nvp("is_entry_trap_constrained", is_entry_trap_constrained()),

source/bin/omnitrace/omnitrace.cpp

@@ -124,6 +124,7 @@ regexvec_t       file_include                  = {};
 regexvec_t       file_exclude                  = {};
 regexvec_t       file_restrict                 = {};
 regexvec_t       func_restrict                 = {};
+regexvec_t       caller_include                = {};
 CodeCoverageMode coverage_mode                 = CODECOV_NONE;
 
 std::unique_ptr<std::ofstream> log_ofs = {};
@@ -633,6 +634,9 @@ main(int argc, char** argv)
     parser.add_argument({ "-R", "--function-restrict" },
                         "Regex(es) for restricting functions only to those "
                         "that match the provided regular-expressions");
+    parser.add_argument({ "--caller-include" },
+                        "Regex(es) for including functions that call the "
+                        "listed functions (despite heuristics)");
     parser.add_argument({ "-MI", "--module-include" },
                         "Regex(es) for selecting modules/files/libraries "
                         "(despite heuristics)");
@@ -1101,6 +1105,8 @@ main(int argc, char** argv)
         add_regex(func_include, tim::get_env<string_t>("OMNITRACE_REGEX_INCLUDE", ""));
         add_regex(func_exclude, tim::get_env<string_t>("OMNITRACE_REGEX_EXCLUDE", ""));
         add_regex(func_restrict, tim::get_env<string_t>("OMNITRACE_REGEX_RESTRICT", ""));
+        add_regex(caller_include,
+                  tim::get_env<string_t>("OMNITRACE_REGEX_CALLER_INCLUDE"));
 
         add_regex(file_include,
                   tim::get_env<string_t>("OMNITRACE_REGEX_MODULE_INCLUDE", ""));
@@ -1123,6 +1129,7 @@ main(int argc, char** argv)
         _parse_regex_option("function-include", func_include);
         _parse_regex_option("function-exclude", func_exclude);
         _parse_regex_option("function-restrict", func_restrict);
+        _parse_regex_option("caller-include", caller_include);
         _parse_regex_option("module-include", file_include);
         _parse_regex_option("module-exclude", file_exclude);
         _parse_regex_option("module-restrict", file_restrict);

tests/CMakeLists.txt

@@ -678,6 +678,27 @@ omnitrace_add_test(
     ENVIRONMENT "${_base_environment};OMNITRACE_CRITICAL_TRACE=OFF"
     REWRITE_FAIL_REGEX "0 instrumented loops in procedure transpose")
 
+omnitrace_add_test(
+    SKIP_PRELOAD SKIP_RUNTIME SKIP_SAMPLING
+    NAME rewrite-caller
+    TARGET rewrite-caller
+    LABELS "caller-include"
+    REWRITE_ARGS
+        -e
+        -i
+        256
+        --caller-include
+        "^inner"
+        -v
+        2
+        --print-instrumented
+        functions
+    RUN_ARGS 17
+    ENVIRONMENT "${_base_environment};OMNITRACE_COUT_OUTPUT=ON"
+    BASELINE_PASS_REGEX "number of calls made = 17"
+    REWRITE_PASS_REGEX "\\[function\\]\\[Forcing\\] caller-include-regex :: 'outer'"
+    REWRITE_RUN_PASS_REGEX ">>> ._outer ([ \\|]+) 17")
+
 set(OMNITRACE_ROCM_EVENTS_TEST
     "GRBM_COUNT,GPUBusy,SQ_WAVES,SQ_INSTS_VALU,VALUInsts,TCC_HIT_sum,TA_TA_BUSY[0]:device=0,TA_TA_BUSY[11]:device=0"
     )