Fix MSan uninitialized reads in FBGEMM depthwise convolution and INT64 GEMM packing

Author: ddelgadovargas-cyber

src/FbgemmI64.cc

@@ -461,13 +461,16 @@ void cblas_gemm_i64_i64acc(
     ldb = N;
   }
 
-  alignas(64) array<int64_t, MCB * KCB> packA;
-  alignas(64) array<int64_t, KCB * NCB> packB;
-  alignas(64) array<int64_t, MCB * NCB> packC;
+  alignas(64) array<int64_t, MCB * KCB> packA = {};
+  alignas(64) array<int64_t, KCB * NCB> packB = {};
+  alignas(64) array<int64_t, MCB * NCB> packC = {};
 
   for (int ic = 0; ic < M; ic += MCB) {
     for (int kc = 0; kc < K; kc += KCB) {
       // pack A
+      if (std::min(MCB, M - ic) < MCB || std::min(K - kc, KCB) < KCB) {
+        packA.fill(0);
+      }
       for (int i = 0; i < std::min(MCB, M - ic); ++i) {
         memcpy(
             &packA[i * KCB],
@@ -477,6 +480,9 @@ void cblas_gemm_i64_i64acc(
 
       for (int jc = 0; jc < N; jc += NCB) {
         // pack B
+        if (std::min(K - kc, KCB) < KCB || std::min(NCB, N - jc) < NCB) {
+          packB.fill(0);
+        }
         for (int i = 0; i < std::min(KCB, K - kc); ++i) {
           memcpy(
               &packB[i * NCB],
@@ -512,6 +518,9 @@ void cblas_gemm_i64_i64acc(
                 std::min(KCB, K - kc),
                 NCB);
           } else {
+            if (std::min(MCB, M - ic) < MCB || std::min(NCB, N - jc) < NCB) {
+              packC.fill(0);
+            }
             for (int i = 0; i < std::min(MCB, M - ic); ++i) {
               memcpy(
                   &packC[i * NCB],

src/FbgemmI8Depthwise2D-inl.h

@@ -11,6 +11,7 @@
 #include "./FbgemmI8DepthwiseUtils.h" // @manual
 #include "./GenerateI8Depthwise.h" // @manual
 #include "./MaskAvx2.h" // @manual
+#include <cstring>
 #include "fbgemm/Utils.h"
 #include "fbgemm/UtilsAvx2.h"
 
@@ -155,6 +156,9 @@ static ALWAYS_INLINE void depthwise_2d_(
   auto row_offsets_owner =
       makeAlignedUniquePtr<int32_t>(64, (IC + 31) / 32 * 32);
   int32_t* row_offsets = row_offsets_owner.get();
+  if (row_offsets) {
+    std::memset(row_offsets, 0, (IC + 31) / 32 * 32 * sizeof(int32_t));
+  }
 
   int64_t n_begin = 0, n_end = 0, h_begin = 0, h_end = 0, w_begin = 0,
           w_end = 0;

src/FbgemmI8Depthwise3D.cc

@@ -9,6 +9,7 @@
 #define FBGEMM_EXPORTS
 #include "fbgemm/FbgemmI8Depthwise.h"
 
+#include <cstring>
 #include <stdexcept> // for logic_error
 #include <string>
 
@@ -170,6 +171,9 @@ static ALWAYS_INLINE void depthwise_3d_same_pad_(
   auto row_offsets_owner =
       makeAlignedUniquePtr<int32_t>(64, (IC + 31) / 32 * 32);
   int32_t* row_offsets = row_offsets_owner.get();
+  if (row_offsets) {
+    std::memset(row_offsets, 0, (IC + 31) / 32 * 32 * sizeof(int32_t));
+  }
 
   int64_t n_begin = 0, n_end = 0, t_begin = 0, t_end = 0, h_begin = 0,
           h_end = 0;