avoid extra copy in PackedGemmMatrixB constructor (#3691)

Summary:
Pull Request resolved: #3691

X-link: facebookresearch/FBGEMM#767

This diff modifies one of the constructors, to allow pass in pmat instead of always copying the memory. It's uses' responsibility to make sure the passed in parameter is valid.

Reviewed By: hl475

Differential Revision: D69564913

fbshipit-source-id: 64dcc64cce642d3b37212eb5cc5baea365085329

include/fbgemm/FbgemmPackMatrixB.h

@@ -108,7 +108,7 @@ class PackedGemmMatrixB {
       const int nbcol,
       const uint64_t size,
       const int kernel_ncol_blocks,
-      const void* pmat)
+      void* pmat)
       : nrow_(nrow),
         ncol_(ncol),
         brow_(brow),
@@ -118,10 +118,9 @@ class PackedGemmMatrixB {
         nbcol_(nbcol),
         size_(size),
         kernel_ncol_blocks_(kernel_ncol_blocks) {
-    pmat_ =
-        static_cast<T*>(fbgemmAlignedAlloc(PMAT_ALIGNMENT, size * sizeof(T)));
-    memcpy(pmat_, pmat, size * sizeof(T));
+    pmat_ = static_cast<T*>(pmat);
     packed_ = true;
+    pmat_passed_in = true;
   }
 
   void initializeParam() {
@@ -166,7 +165,9 @@ class PackedGemmMatrixB {
   }
 
   ~PackedGemmMatrixB() {
-    fbgemmAlignedFree(pmat_);
+    if (pmat_passed_in == false) {
+      fbgemmAlignedFree(pmat_);
+    }
   }
 
   void unpackFromSrc(const matrix_op_t trans, T* src_mat) {
@@ -293,6 +294,7 @@ class PackedGemmMatrixB {
   int kernel_ncol_blocks_;
   T* pmat_;
   bool packed_{false};
+  bool pmat_passed_in{false};
 };
 
 } // namespace fbgemm