mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Suppress MSAN warnings from SVD that are showing up in CI.

Browse Source 
In our MSAN CI, the copy of LAPACK we use is not MSAN-instrumented, leading to false positives. Suppress those false-positives via annotations.

PiperOrigin-RevId: 712607044
This commit is contained in:
Peter Hawkins 2025-01-06 11:48:33 -08:00 committed by jax authors
parent 52cc5c7f05
commit 61dd041225
1 changed files with 29 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
jaxlib/cpu/lapack_kernels.cc
View File

@ -758,6 +758,35 @@ static ffi::Error SvdKernel(
work_data.get(), &workspace_dim_v,
iwork_data.get(), info_data);
}
// Suppress MSAN warnings when using a copy of LAPACK uninstrumented by
// MSAN.
using T [[maybe_unused]] = typename svd::SVDType<dtype>::ValueType;
ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(info_data, sizeof(*info_data));
ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(x_out_data,
x_cols_v * x_leading_dim_v * sizeof(T));
ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(
singular_values_data, std::min(x_rows_v, x_cols_v) * sizeof(RealType));
if (mode_v == 'A') {
ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(
u_data, u_leading_dim_v * x_rows_v * sizeof(T));
ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(
vt_data, vt_leading_dim_v * x_cols_v * sizeof(T));
} else if (mode_v == 'O') {
if (x_rows_v < x_cols_v) {
ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(
u_data, u_leading_dim_v * x_rows_v * sizeof(T));
} else {
ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(
vt_data, vt_leading_dim_v * x_cols_v * sizeof(T));
}
} else if (mode_v == 'S') {
ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(
u_data, u_leading_dim_v * std::min(x_rows_v, x_cols_v) * sizeof(T));
ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(
vt_data, vt_leading_dim_v * x_cols_v * sizeof(T));
}
x_out_data += x_out_step;
singular_values_data += singular_values_step;
u_data += u_step;