MemorySanitizerJIT

Introduction

MemorySanitizer is implemented as an LLVM pass, and is used by the Clang compiler front-end to instrument statically compiled C/C++ programs by specifying the -fsanitize=memory option. Support for other languages is not readily available, but integration is feasible for other LLVM-based compilers.

This page details how programming language developers can implement MemorySanitizer support for Just-In-Time compiled code generated by the LLVM JIT. We'll use the Building a JIT in LLVM tutorial as our starting point.

Instrumenting the host program and libraries

JIT-compiled languages are executed by programs and/or libraries which contain their compiler (LLVM) and runtime. While for a standalone language this would typically be called the "Virtual Machine" or VM, it can be part of a larger program, like e.g. a browser. Data gets exchanged between this statically compiled code, and the dynamically compiled code, so to be able to detect uses of uninitialized values within the JIT-compiled code we must also have MemorySanitizer instrument the static 'host' program which is written in C++.

MemorySanitizerLibcxxHowTo details how to compile an instrumented C++ standard library to be used when building the host program. However, one should build it with -DLLVM_USE_SANITIZER=Memory instead of the suggested MemoryWithOrigins, and similarly ensure the host program is built without -fsanitize-memory-track-origins (or assign it the value 0). That's because origin tracking requires static debug information for stack unwinding, which isn't available for the dynamically generated code.

The MemorySanitizer pass

The header for the MemorySanitizer pass in the LLVM code base is llvm/include/llvm/Transforms/Instrumentation/MemorySanitizer.h, while its implementation resides at llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp.

Basing ourselves on http://llvm.org/docs/tutorial/BuildingAJIT2.html, the MemorySanitizer pass can be added before the other optimization passes by using:

if (__has_feature(memory_sanitizer)) {
  FPM->add(createMemorySanitizerLegacyPassPass());
}

Adding the pass before optimization passes helps catch more bugs early because the optimizations can transform things like branches into arithmetic code, which would not trigger an MSan error. Note also that we're guarding this code by the C++ 'host' program itself being compiled with MemorySanitizer enabled.

Enabling instrumentation

Just adding the pass is not enough. MemorySanitizer requires opting in to emitting code which will perform the checks for uses of uninitialized values. Right after creating each function, we must add this attribute:

F->addFnAttr(llvm::Attribute::SanitizeMemory);

Omitting this would correspond to adding __attribute__((no_sanitize("memory"))) to C++ functions. The shadow memory still gets updated, but doesn't trigger an error when used within this function.

Setting the target triple

MemorySanitizer requires knowing which OS we're running on, to determine where to map the 'shadow' memory. This can be set using:

TheModule->setTargetTriple(LLVM_HOST_TRIPLE);

Implementing Emulated TLS

MemorySanitizer relies on several Thread-local Storage (TLS) variables for storing the shadow value of function parameters and return values. The LLVM JIT defaults to using an 'Emulated TLS' implementation. This makes it emit calls to __emutls_get_address() to obtain the address of MemorySanitizer's TLS variables, passing it the address of an __emutls_v.* control structure unique to each variable. While this 'Emulated TLS' is implemented by compiler-rt, that shouldn't actually be used. The issue is that the TLS variables MemorySanitizer uses are instantiated as part of the host program, which does not use Emulated TLS, and hence the __emutls_v.* structures do not exist. Also note we want to make sure we use the same TLS variables for our JIT-compiled code as for the statically compiled code, to get correct shadow values for data passed between the two.

While the mismatch between TLS implementations could potentially be solved by compiling the host program with -femulated-tls, this can have a significant effect on performance (remember the host program includes at least all of LLVM and the runtime). At the time of writing the JIT does not support other TLS implementations, and adding it would not be straightforward.

Fortunately, we can actually use Emulated TLS to our advantage by providing our own implementation which uses the native TLS support of the host program. Both the address of the __emutls_get_address() function, and the addresses of the __emutls_v.* structures are resolved through the DefinitionGenerator interface. The KaleidoscopeJIT tutorial just uses the DynamicLibrarySearchGenerator implementation of this symbol resolver interface, which looks for the address in the global symbol table of the program by using dlsym(), but we can provide alternative/additional ones as well.

This enables us to return the address of a different function when the JIT queries for "__emutls_get_address", like the GetTLSAddress function below.

// Forward declare the real TLS variables used by MemorySanitizer. These are
// defined in llvm-project/compiler-rt/lib/msan/msan.cpp.
extern __thread unsigned long long __msan_param_tls[];
extern __thread unsigned long long __msan_retval_tls[];
extern __thread unsigned long long __msan_va_arg_tls[];
extern __thread unsigned long long __msan_va_arg_overflow_size_tls;

// These enums each represent one of the TLS variables above, and will be
// passed to our GetTLSAddress(void *) function below, which replaces
// __emutls_get_address(__emutls_control *).
enum class MSanTLS {
  param = 1,            // __msan_param_tls
  retval,               // __msan_retval_tls
  va_arg,               // __msan_va_arg_tls
  va_arg_overflow_size  // __msan_va_arg_overflow_size_tls
};

// This function takes a pointer, but we actually get LLVM to pass it the value
// of one of the enums representing MSan's TLS variables. Then for each enum we
// simply let this statically compiled C++ code determine the address of the
// corresponding TLS variable (note this will differ depending on which thread
// makes the call).
static void *GetTLSAddress(void *control) {
  auto tlsIndex = static_cast<MSanTLS>(reinterpret_cast<uintptr_t>(control));
  switch (tlsIndex) {
    case MSanTLS::param:
      return reinterpret_cast<void *>(&__msan_param_tls);
    case MSanTLS::retval:
      return reinterpret_cast<void *>(&__msan_retval_tls);
    case MSanTLS::va_arg:
      return reinterpret_cast<void *>(&__msan_va_arg_tls);
    case MSanTLS::va_arg_overflow_size:
      return reinterpret_cast<void *>(&__msan_va_arg_overflow_size_tls);
    default:
      assert(false && "MemorySanitizer used an unrecognized TLS variable");
      return nullptr;
  }
}

All that's left to do now is to have our custom DefinitionGenerator perform the following mappings from a symbol string to an address:

if (Symbol == "__emutls_get_address")
  return reinterpret_cast<void *>(GetTLSAddress));
if (Symbol == "__emutls_v.__msan_retval_tls")
  return reinterpret_cast<void *>(static_cast<uintptr_t>(MSanTLS::retval)));
if (Symbol == "__emutls_v.__msan_param_tls")
  return reinterpret_cast<void *>(static_cast<uintptr_t>(MSanTLS::param)));
if (Symbol == "__emutls_v.__msan_va_arg_tls")
  return reinterpret_cast<void *>(static_cast<uintptr_t>(MSanTLS::va_arg)));
if (Symbol == "__emutls_v.__msan_va_arg_overflow_size_tls")
  return reinterpret_cast<void *>(static_cast<uintptr_t>(MSanTLS::va_arg_overflow_size)));

Resolving other symbols

When the MSan instrumentation detects a use of an uninitialized value, it calls the __msan_warning_with_origin_noreturn function. Thus we must ensure the JIT can also resolve its address. The DynamicLibrarySearchGenerator used by the KaleidoscopeJIT tutorial takes care of that, and any other symbols that might be used (which can vary between versions of Clang).

Note that most actual JIT implementations don't use DynamicLibrarySearchGenerator, because it requires the symbols to be visible to dlsym(), i.e. the program would have to not have its symbols stripped. While that's usually not an option for release builds shipped to end users, MSan builds by design make their symbols available, to ensure a single instance of the TLS variables is used across all modules.

Hence the use of DynamicLibrarySearchGenerator or calling dlsym() ourselves in a custom symbol resolver, for MSan builds only, is appropriate.

Disabling optimization to improve performance

The MemorySanitizer instrumentation can make LLVM's 'CodeGen' optimizations substantially slower. This can be worked around by disabling all optimization:

JTMB.setCodeGenOptLevel(llvm::CodeGenOpt::None);

Some IR-level optimizations are useful though; see Clang's addGeneralOptsForMemorySanitizer.

Example

The above MemorySanitizer integration for JIT-compilation has been implemented in Reactor, the embedded language which powers the dynamic code generation of the SwiftShader graphics driver. A more detailed account of this particular integration can be found in MemorySanitizer for Reactor.