[mlir] Add msan memory unpoisoning macros to mlir ExecutionEngine

Adding annotations on as-needed bases, currently only for memrefCopy, but in general all C API functions that take pointers to memory allocated/initialized inside the jit-compiled code must be annotated, to be able to run with msan.

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D123557

mlir/include/mlir/ExecutionEngine/Msan.h

@@ -0,0 +1,35 @@
+//===- Msan.h - Utils related to the memory sanitizer ---------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares and defines macros related to msan.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_EXECUTIONENGINE_MSAN_H
+#define MLIR_EXECUTIONENGINE_MSAN_H
+
+// Memory sanitizer currently can't be enabled for the jit-compiled code, and
+// to suppress msan warnings we need to unpoison pointers and pointed-to
+// datastructures before they can be accessed.
+
+#ifndef __has_feature
+#define __has_feature(x) 0
+#endif
+
+#if __has_feature(memory_sanitizer) && !defined(MLIR_MEMORY_SANITIZER)
+#define MLIR_MEMORY_SANITIZER
+#endif
+
+#if defined(MLIR_MEMORY_SANITIZER)
+#include <sanitizer/msan_interface.h>
+#define MLIR_MSAN_MEMORY_IS_INITIALIZED(p, s) __msan_unpoison((p), (s))
+#else // Memory sanitizer: OFF
+#define MLIR_MSAN_MEMORY_IS_INITIALIZED(p, s)
+#endif // MLIR_MEMORY_SANITIZER
+
+#endif // MLIR_EXECUTIONENGINE_MSAN_H

mlir/lib/ExecutionEngine/CRunnerUtils.cpp

@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
+#include "mlir/ExecutionEngine/Msan.h"
 
 #ifndef _WIN32
 #if defined(__FreeBSD__) || defined(__NetBSD__) || defined(__OpenBSD__)
@@ -51,6 +52,8 @@ memrefCopy(int64_t elemSize, UnrankedMemRefType<char> *srcArg,
   DynamicMemRefType<char> dst(*dstArg);
 
   int64_t rank = src.rank;
+  MLIR_MSAN_MEMORY_IS_INITIALIZED(src.sizes, rank * sizeof(int64_t));
+
   // Handle empty shapes -> nothing to copy.
   for (int rankp = 0; rankp < rank; ++rankp)
     if (src.sizes[rankp] == 0)

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

@@ -6078,6 +6078,7 @@ cc_library(
     ],
     hdrs = [
         "include/mlir/ExecutionEngine/CRunnerUtils.h",
+        "include/mlir/ExecutionEngine/Msan.h",
         "include/mlir/ExecutionEngine/SparseTensorUtils.h",
     ],
     includes = ["include"],