Skip to content

Commit

Permalink
Use probability sampling over periodic sampling (#213)
Browse files Browse the repository at this point in the history 
This PR introduces probabilistic sampling into kokkos_sampler, enabling the avoidance of bias in the sampled operations.
Two new options are introduced: KOKKOS_TOOLS_SAMPLER_PROB and KOKKOS_TOOLS_RANDOM_SEED which set the rate and a random seed respectively. 
The PR also introduces tests.
  • Loading branch information
@vlkale
vlkale authored May 23, 2024
1 parent baf9ae3 commit 81d3f3a
Show file tree
Hide file tree
Showing 15 changed files with 570 additions and 150 deletions.
11 changes: 9 additions & 2 deletions common/kokkos-sampler/README.md
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,11 @@
This is a sampler utility that is intended to complement other tools in the Kokkos Tools set. This utility allows for sampling (rather than collecting) of profiling or debugging data gathered from a particular tool of the Kokkos Tools set. The Kokkos Tools user provides a sampling rate via the environment variable KOKKOS_TOOLS_SAMPLER_SKIP.
This is a sampler utility that is intended to complement other tools in the Kokkos Tools set. This utility allows for sampling (rather than collecting) of profiling or debugging data gathered from a particular tool of the Kokkos Tools set.

To use this utility, a Kokkos Tools user provides a sampling probability by setting the environment variable `KOKKOS_TOOLS_SAMPLER_PROB` to a positive real number between 0.0 and 100.0. The user can alternatively set a sampling skip rate, i.e., the number of Kokkos kernel invocations to skip before the next sample is taken. The user does so by setting the environment variable `KOKKOS_TOOLS_SAMPLER_SKIP` to a non-negative integer.

If both sampling probability and sampling skip rate are set by the user, this sampling utility only uses the sampling probability for sampling; the utility sets the sampling skip rate to 1, incorporating no pre-defined periodicity in sampling. If neither sampling probability nor the sampling skip rate are set by the user, then randomized sampling is done, with the sampler's probability being 10.0 percent. The sampler is periodic only if the sampling probability is not set by the user and the sampling skip rate is set by the user.

For randomized sampling, the user can ensure reproducibility of this tool's output across multiple runs of a Kokkos application by setting `KOKKOS_TOOLS_RANDOM_SEED` to an integer value before all of the runs. If this environment variable is not set, the seed is based on the C time function.

In order for the state of the sampled profiling and logging data in memory to be captured at the time of the utility's callback invocation, it might be important to enforce fences. However, this also means that there are more synchronization points compared with running the program without the tool.
This fencing behavior can be controlled by setting the environment variable `KOKKOS_TOOLS_GLOBALFENCES`. A non-zero value implies global fences on invocation of the tool. The default is not to introduce extra fences.
This fencing behavior can be controlled by setting the environment variable `KOKKOS_TOOLS_GLOBALFENCES`. A non-zero value implies global fences on invocation of the tool. The default is not to introduce extra fences.

193 changes: 135 additions & 58 deletions common/kokkos-sampler/kp_sampler_skip.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,14 +6,18 @@
#include <dlfcn.h>
#include "../../profiling/all/kp_core.hpp"
#include "kp_config.hpp"
#include <ctime>
#include <limits>
#include <iostream>

namespace KokkosTools {
namespace Sampler {
static uint64_t uniqID = 0;
static uint64_t kernelSampleSkip = 101;
static uint64_t kernelSampleSkip = std::numeric_limits<uint64_t>::max();
static double tool_prob_num = -1.0;
static int tool_verbosity = 0;
static int tool_globFence = 0;
static int tool_seed = -1;

// a hash table mapping kID to nestedkID
static std::unordered_map<uint64_t, uint64_t> infokIDSample;
Expand Down Expand Up @@ -77,6 +81,8 @@ void kokkosp_init_library(const int loadSeq, const uint64_t interfaceVer,
const uint32_t devInfoCount, void* deviceInfo) {
const char* tool_verbose_str = getenv("KOKKOS_TOOLS_SAMPLER_VERBOSE");
const char* tool_globFence_str = getenv("KOKKOS_TOOLS_GLOBALFENCES");
const char* tool_seed_str = getenv("KOKKOS_TOOLS_RANDOM_SEED");

if (NULL != tool_verbose_str) {
tool_verbosity = atoi(tool_verbose_str);
} else {
Expand All @@ -88,11 +94,15 @@ void kokkosp_init_library(const int loadSeq, const uint64_t interfaceVer,
tool_globFence = 0;
}

if (NULL != tool_seed_str) {
tool_seed = atoi(tool_seed_str);
}

char* profileLibrary = getenv("KOKKOS_TOOLS_LIBS");
if (NULL == profileLibrary) {
printf(
"Checking KOKKOS_PROFILE_LIBRARY. WARNING: This is a depreciated "
"variable. Please use KOKKOS_TOOLS_LIBS\n");
std::cout << "Checking KOKKOS_PROFILE_LIBRARY. WARNING: This is a "
"deprecated variable. Please use KOKKOS_TOOLS_LIBS\n";

profileLibrary = getenv("KOKKOS_PROFILE_LIBRARY");
if (NULL == profileLibrary) {
std::cout << "KokkosP: FATAL: No library to call in " << profileLibrary
Expand All @@ -113,8 +123,7 @@ void kokkosp_init_library(const int loadSeq, const uint64_t interfaceVer,
nextLibrary = strtok(NULL, ";");

if (NULL == nextLibrary) {
std::cout << "KokkosP: FATAL: No child library of sampler utility library "
"to call in "
std::cout << "KokkosP: FATAL: No child library to call in "
<< profileLibrary << "!\n";
exit(-1);
} else {
Expand Down Expand Up @@ -158,15 +167,15 @@ void kokkosp_init_library(const int loadSeq, const uint64_t interfaceVer,
std::cout << "KokkosP: Function Status:\n";
std::cout << "KokkosP: begin-parallel-for: "
<< ((beginForCallee == NULL) ? "no" : "yes") << "\n";
std::cout << "KokkosP: begin-parallel-scan: "
std::cout << "KokkosP: begin-parallel-scan: "
<< ((beginScanCallee == NULL) ? "no" : "yes") << "\n";
std::cout << "KokkosP: begin-parallel-reduce: "
std::cout << "KokkosP: begin-parallel-reduce: "
<< ((beginReduceCallee == NULL) ? "no" : "yes") << "\n";
std::cout << "KokkosP: end-parallel-for: "
std::cout << "KokkosP: end-parallel-for: "
<< ((endForCallee == NULL) ? "no" : "yes") << "\n";
std::cout << "KokkosP: end-parallel-scan: "
std::cout << "KokkosP: end-parallel-scan: "
<< ((endScanCallee == NULL) ? "no" : "yes") << "\n";
std::cout << "KokkosP: end-parallel-reduce: "
std::cout << "KokkosP: end-parallel-reduce: "
<< ((endReduceCallee == NULL) ? "no" : "yes") << "\n";
}
}
Expand All @@ -176,15 +185,78 @@ void kokkosp_init_library(const int loadSeq, const uint64_t interfaceVer,

uniqID = 1;

if (0 > tool_seed) {
srand(time(NULL));
if (tool_verbosity > 0) {
std::cout << "KokkosP: Seeding random number generator using clock for "
"random sampling.\n";
}
} else {
srand(tool_seed);
if (tool_verbosity > 0) {
std::cout << "KokkosP: Seeding random number generator using seed "
<< tool_seed << " for random sampling.\n";
}
}

const char* tool_probability = getenv("KOKKOS_TOOLS_SAMPLER_PROB");

if (NULL != tool_probability) {
// Read sampling probability as a float between 0 and 100, representing
// a percentage that data should be gathered.
// Utility reasons about probability as a double between 0.0 and 1.0.
tool_prob_num = atof(tool_probability);
if (tool_prob_num > 100.0) {
std::cout << "KokkosP: The sampling probability value is set to be "
"greater than 100.0. The probability for the sampler will "
"be set to 100 percent; all of the invocations of a Kokkos "
"kernel will be profiled.\n";
tool_prob_num = 100.0;
} else if (tool_prob_num < 0.0) {
std::cout
<< "KokkosP: The sampling probability value is set to be a negative "
"number. The sampler's probability will be set to 0 percent; none "
"of the invocations of a Kokkos kernel will be profiled.\n";
tool_prob_num = 0.0;
}
if (tool_verbosity > 0) {
std::cout << "KokkosP: Probability for the sampler set to: "
<< tool_prob_num << "\n";
}
kernelSampleSkip = 1;
}

const char* tool_sample = getenv("KOKKOS_TOOLS_SAMPLER_SKIP");
if (NULL != tool_sample) {
if ((NULL != tool_sample) && (tool_prob_num == -1.0)) {
// If the user touched the sample skip rate variable
// and the tool probability is set to -1 (no probability sampling
// desired), then use only sampler skip rate.
tool_prob_num = 100.0;
kernelSampleSkip = atoi(tool_sample) + 1;
if (tool_verbosity > 0) {
std::cout << "KokkosP: Sampling rate set to: " << tool_sample << "\n";
}
}

if (tool_verbosity > 0) {
std::cout << "KokkosP: Sampling rate set to: " << tool_sample << "\n";
if (tool_prob_num == -1.0) {
// If the tool probability is set to -1 (no probability sampling
// desired) and the user also didn't set
// skip rate, then use a default with a probability sampling of 10%.

if (tool_verbosity > 0) {
std::cout << "KokkosP: Neither the probability nor the skip rate for "
"sampling were set...\n";
}
tool_prob_num = 10.0;
kernelSampleSkip = 1;
if (tool_verbosity > 0) {
std::cout << "KokkosP: The probability for the sampler is set to the "
"default of "
<< tool_prob_num
<< " percent. The skip rate for sampler will not be used.\n";
}
}
}
} // end kokkosp_init_library

void kokkosp_finalize_library() {
if (NULL != finalizeProfileLibrary) (*finalizeProfileLibrary)();
Expand All @@ -196,22 +268,23 @@ void kokkosp_begin_parallel_for(const char* name, const uint32_t devID,
static uint64_t invocationNum = 0;
++invocationNum;
if ((invocationNum % kernelSampleSkip) == 0) {
if (tool_verbosity > 0) {
std::cout << "KokkosP: sample " << *kID
<< " calling child-begin function...\n";
}

if (NULL != beginForCallee) {
if (tool_globFence) {
invoke_ktools_fence(0);
}
uint64_t nestedkID = 0;
(*beginForCallee)(name, devID, &nestedkID);
if (tool_verbosity > 0) {
std::cout << "KokkosP: sample " << *kID
<< " finished with child-begin function.\n";
if ((rand() / (1.0 * RAND_MAX)) < (tool_prob_num / 100.0)) {
if (NULL != beginForCallee) {
if (tool_verbosity > 0) {
std::cout << "KokkosP: sample " << *kID
<< " calling child-begin function...\n";
}
if (tool_globFence) {
invoke_ktools_fence(0);
}
uint64_t nestedkID = 0;
(*beginForCallee)(name, devID, &nestedkID);
if (tool_verbosity > 0) {
std::cout << "KokkosP: sample " << *kID
<< " finished with child-begin function.\n";
}
infokIDSample.insert({*kID, nestedkID});
}
infokIDSample.insert({*kID, nestedkID});
}
}
}
Expand Down Expand Up @@ -244,21 +317,23 @@ void kokkosp_begin_parallel_scan(const char* name, const uint32_t devID,
static uint64_t invocationNum = 0;
++invocationNum;
if ((invocationNum % kernelSampleSkip) == 0) {
if (tool_verbosity > 0) {
std::cout << "KokkosP: sample " << *kID
<< " calling child-begin function...\n";
}
if (NULL != beginScanCallee) {
uint64_t nestedkID = 0;
if (tool_globFence) {
invoke_ktools_fence(0);
if ((rand() / (1.0 * RAND_MAX)) < (tool_prob_num / 100.0)) {
if (NULL != beginScanCallee) {
if (tool_verbosity > 0) {
std::cout << "KokkosP: sample " << *kID
<< " calling child-begin function...\n";
}
uint64_t nestedkID = 0;
if (tool_globFence) {
invoke_ktools_fence(0);
}
(*beginScanCallee)(name, devID, &nestedkID);
if (tool_verbosity > 0) {
std::cout << "KokkosP: sample " << *kID
<< " finished with child-begin function.\n";
}
infokIDSample.insert({*kID, nestedkID});
}
(*beginScanCallee)(name, devID, &nestedkID);
if (tool_verbosity > 0) {
std::cout << "KokkosP: sample " << *kID
<< " finished with child-begin function.\n";
}
infokIDSample.insert({*kID, nestedkID});
}
}
}
Expand Down Expand Up @@ -290,21 +365,23 @@ void kokkosp_begin_parallel_reduce(const char* name, const uint32_t devID,
static uint64_t invocationNum = 0;
++invocationNum;
if ((invocationNum % kernelSampleSkip) == 0) {
if (tool_verbosity > 0) {
std::cout << "KokkosP: sample " << *kID
<< " calling child-begin function...\n";
}
if (NULL != beginReduceCallee) {
uint64_t nestedkID = 0;
if (tool_globFence) {
invoke_ktools_fence(0);
}
(*beginReduceCallee)(name, devID, &nestedkID);
if (tool_verbosity > 0) {
std::cout << "KokkosP: sample " << *kID
<< " finished with child-begin function.\n";
if ((rand() / (1.0 * RAND_MAX)) < (tool_prob_num / 100.0)) {
if (NULL != beginReduceCallee) {
if (tool_verbosity > 0) {
std::cout << "KokkosP: sample " << *kID
<< " calling child-begin function...\n";
}
uint64_t nestedkID = 0;
if (tool_globFence) {
invoke_ktools_fence(0);
}
(*beginReduceCallee)(name, devID, &nestedkID);
if (tool_verbosity > 0) {
std::cout << "KokkosP: sample " << *kID
<< " finished with child-begin function.\n";
}
infokIDSample.insert({*kID, nestedkID});
}
infokIDSample.insert({*kID, nestedkID});
}
}
}
Expand Down
23 changes: 21 additions & 2 deletions tests/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,13 @@
# to the target file of this argument (optional)
# KOKKOS_TOOLS_SAMPLER_VERBOSE : the test environment will receive the variable 'KOKKOS_TOOLS_SAMPLER_VERBOSE' that is set as the value of 1 for printing the sample has been taken
# KOKKOS_TOOLS_GLOBALFENCES : test environment receives the variable 'KOKKOS_TOOLS_GLOBALFENCES' that is set as the value of 1 to turn the tool's auto-fencing on.
# KOKKOS_TOOLS_RANDOM_SEED : test environment receives the variable 'KOKKOS_TOOLS_RANDOM_SEED' that is set as the value for a seed of the random number generator (used for testing repeatability).
# KOKKOS_TOOLS_SAMPLER_SKIP : test environment receives the variable 'KOKKOS_TOOLS_SAMPLER_SKIP' that is set as the value of the number of Kokkos kernel invocations to skip before a tooling activity is invoked.
# KOKKOS_TOOLS_SAMPLER_PROB : test environment receives the variable 'KOKKOS_TOOLS_SAMPLER_PROB' that is set as the probability that a Kokkos kernel invocation has a tooling activity invoked for it.

function(kp_add_executable_and_test)
cmake_parse_arguments(kaeat_args "" "TARGET_NAME;SOURCE_FILE;KOKKOS_TOOLS_SAMPLER_VERBOSE;KOKKOS_TOOLS_GLOBALFENCES;KOKKOS_TOOLS_SAMPLER_SKIP" "KOKKOS_TOOLS_LIBS" ${ARGN})

function(kp_add_executable_and_test)
cmake_parse_arguments(kaeat_args "" "TARGET_NAME;SOURCE_FILE;KOKKOS_TOOLS_SAMPLER_VERBOSE;KOKKOS_TOOLS_GLOBALFENCES;KOKKOS_TOOLS_SAMPLER_SKIP;KOKKOS_TOOLS_SAMPLER_PROB;KOKKOS_TOOLS_RANDOM_SEED" "KOKKOS_TOOLS_LIBS" ${ARGN})
if(NOT DEFINED kaeat_args_TARGET_NAME)
message(FATAL_ERROR "'TARGET_NAME' is a required argument.")
endif()
Expand Down Expand Up @@ -78,7 +80,24 @@ function(kp_add_executable_and_test)
PROPERTY
ENVIRONMENT "KOKKOS_TOOLS_SAMPLER_SKIP=${kaeat_args_KOKKOS_TOOLS_SAMPLER_SKIP}"
)
endif()

if (DEFINED kaeat_args_KOKKOS_TOOLS_SAMPLER_PROB)
set_property(
TEST ${kaeat_args_TARGET_NAME}
APPEND
PROPERTY
ENVIRONMENT "KOKKOS_TOOLS_SAMPLER_PROB=${kaeat_args_KOKKOS_TOOLS_SAMPLER_PROB}"
)
endif()

if (DEFINED kaeat_args_KOKKOS_TOOLS_RANDOM_SEED)
set_property(
TEST ${kaeat_args_TARGET_NAME}
APPEND
PROPERTY
ENVIRONMENT "KOKKOS_TOOLS_RANDOM_SEED=${kaeat_args_KOKKOS_TOOLS_RANDOM_SEED}"
)
endif()

endfunction(kp_add_executable_and_test)
Expand Down
Loading

0 comments on commit 81d3f3a

Please sign in to comment.