diff --git a/mlir/include/mlir/ExecutionEngine/Msan.h b/mlir/include/mlir/ExecutionEngine/Msan.h new file mode 100644 --- /dev/null +++ b/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 +#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 diff --git a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp b/mlir/lib/ExecutionEngine/CRunnerUtils.cpp --- a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp +++ b/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 @@ DynamicMemRefType 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) diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel --- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel @@ -6077,6 +6077,7 @@ ], hdrs = [ "include/mlir/ExecutionEngine/CRunnerUtils.h", + "include/mlir/ExecutionEngine/Msan.h", "include/mlir/ExecutionEngine/SparseTensorUtils.h", ], includes = ["include"],