#include "HalideRuntime.h"
#include "device_buffer_utils.h"
extern "C" {
// This function is expected to be provided by LLVM as part of the MSAN implementation.
extern void __msan_unpoison(const volatile void *a, size_t size);
extern void __msan_check_mem_is_initialized(const volatile void *x, size_t size);
extern long __msan_test_shadow(const volatile void *x, size_t size);
WEAK int halide_msan_annotate_memory_is_initialized(void *user_context, const void *ptr, uint64_t len) {
__msan_unpoison(ptr, (size_t)len);
return 0;
}
WEAK int halide_msan_check_memory_is_initialized(void *user_context, const void *ptr, uint64_t len, const char *name) {
long offset = __msan_test_shadow(ptr, (size_t)len);
if (offset >= 0) {
print(user_context) << "MSAN failure detected for " << name << " @ " << ptr << " + " << (int)offset << "\n";
// This is slightly redundant but gives better output.
__msan_check_mem_is_initialized(ptr, (size_t)len);
}
return 0;
}
namespace Halide {
namespace Runtime {
namespace Internal {
WEAK void annotate_helper(void *uc, const device_copy &c, int d, int64_t off) {
while (d >= 0 && c.extent[d] == 1) {
d--;
}
if (d == -1) {
const void *from = (void *)(c.src + off);
(void)halide_msan_annotate_memory_is_initialized(uc, from, c.chunk_size); // ignore errors
} else {
for (uint64_t i = 0; i < c.extent[d]; i++) {
annotate_helper(uc, c, d - 1, off);
off += c.src_stride_bytes[d];
}
}
};
WEAK void check_helper(void *uc, const device_copy &c, int d, int64_t off, const char *buf_name) {
while (d >= 0 && c.extent[d] == 1) {
d--;
}
if (d == -1) {
const void *from = (void *)(c.src + off);
(void)halide_msan_check_memory_is_initialized(uc, from, c.chunk_size, buf_name); // ignore errors
} else {
for (uint64_t i = 0; i < c.extent[d]; i++) {
check_helper(uc, c, d - 1, off, buf_name);
off += c.src_stride_bytes[d];
}
}
};
} // namespace Internal
} // namespace Runtime
} // namespace Halide
// Default implementation marks the data pointed to by the halide_buffer_t as initialized
// (but *not* the halide_buffer_t itself); it takes pains to only mark the active memory ranges
// (skipping padding), and sorting into ranges to always mark the smallest number of
// ranges, in monotonically increasing memory order.
WEAK int halide_msan_annotate_buffer_is_initialized(void *user_context, halide_buffer_t *b) {
if (b == nullptr) {
return 0;
}
Halide::Runtime::Internal::device_copy c = Halide::Runtime::Internal::make_host_to_device_copy(b);
if (c.chunk_size == 0) {
return 0;
}
if (b->device_dirty()) {
// buffer has been computed on a gpu, but not copied back:
// don't annotate as initialized. (We'll assume that subsequent
// calls to halide_copy_to_host will force another call.)
return 0;
}
annotate_helper(user_context, c, MAX_COPY_DIMS - 1, 0);
return 0;
}
WEAK void halide_msan_annotate_buffer_is_initialized_as_destructor(void *user_context, void *b) {
(void)halide_msan_annotate_buffer_is_initialized(user_context, (halide_buffer_t *)b);
}
WEAK int halide_msan_check_buffer_is_initialized(void *user_context, halide_buffer_t *b, const char *buf_name) {
if (b == nullptr) {
return 0;
}
(void)halide_msan_check_memory_is_initialized(user_context, (void *)b, sizeof(*b), buf_name); // ignore errors
(void)halide_msan_check_memory_is_initialized(user_context, (void *)b->dim, b->dimensions * sizeof(b->dim[0]), buf_name); // ignore errors
Halide::Runtime::Internal::device_copy c = Halide::Runtime::Internal::make_host_to_device_copy(b);
if (c.chunk_size == 0) {
return 0;
}
if (b->device_dirty()) {
// buffer has been computed on a gpu, but not copied back:
// don't check. (We'll assume that subsequent calls to
// halide_copy_to_host will force another call.)
return 0;
}
check_helper(user_context, c, MAX_COPY_DIMS - 1, 0, buf_name);
return 0;
}
} // extern "C"