Skip to content

Commit

Permalink
avoid extra copy in PackedGemmMatrixB constructor
Browse files Browse the repository at this point in the history
Summary:
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.

Differential Revision: D69564913
  • Loading branch information
helloguo authored and facebook-github-bot committed Feb 13, 2025
1 parent 1b7789a commit 2398bbc
Showing 1 changed file with 7 additions and 5 deletions.
12 changes: 7 additions & 5 deletions include/fbgemm/FbgemmPackMatrixB.h
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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() {
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -293,6 +294,7 @@ class PackedGemmMatrixB {
int kernel_ncol_blocks_;
T* pmat_;
bool packed_{false};
bool pmat_passed_in{false};
};

} // namespace fbgemm

0 comments on commit 2398bbc

Please sign in to comment.