Skip to content

[tmva][sofie] memory optimization in convK and imcol for ConvTranspose #18168

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sanjibansg merged 2 commits into from
Apr 4, 2025
Merged

[tmva][sofie] memory optimization in convK and imcol for ConvTranspose #18168

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ce69ea0
feat: add convK and imcol to memory optimization for ConvTranspose
sanjibansg Mar 27, 2025
2de34c2
fix: follow root naming convention for convK and imcol
sanjibansg Apr 1, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 0 additions & 23 deletions tmva/sofie/inc/TMVA/ROperator_Conv.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -300,29 +300,6 @@ public:
return out.str();
}

// Generate code for Session data members (e.g. internal vectors)
virtual std::string GenerateSessionMembersCode(std::string opName) override {

size_t outputChannelSize = fShapeY[2]; // size/channel = D * H * W
size_t kernelSize = fAttrKernelShape[0];
for (size_t i = 1; i < fDim; i++) {
outputChannelSize *= fShapeY[2 + i];
kernelSize *= fAttrKernelShape[i];
}

opName = "op_" + opName;
std::stringstream out;
// matrix with convolution kernels
// out << "std::vector<" << fType << "> fVec_" << opName << "_f = std::vector<" << fType << ">("
// << fShapeW[0] * fShapeW[1] * kernelSize << ");\n";
// // output matrix of im2col
// out << "std::vector<" << fType << "> fVec_" << opName << "_xcol = std::vector<" << fType << ">("
// << fShapeW[1] * kernelSize * outputChannelSize << ");\n";
// out << "\n";

return out.str();
}

std::string Generate(std::string OpName) override {
OpName = "op_" + OpName;

Expand Down
8 changes: 3 additions & 5 deletions tmva/sofie/inc/TMVA/ROperator_ConvTranspose.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,9 @@ private:
std::string fNBroadcastedB;
std::string fNY;

std::string fConvK;
std::string fImcol;

std::vector<size_t> fShapeX;
std::vector<size_t> fShapeW;
std::vector<size_t> fShapeB;
Expand Down Expand Up @@ -113,11 +116,6 @@ public:
*/
std::string GenerateInitCode() override;

/*! \brief Generate code for Session data members (e.g. internal vectors)
* \param opName name of the operator
*/
std::string GenerateSessionMembersCode(std::string /*opName*/) override;

/*! \brief Generate the inference code
* \param opName name of the operator
*/
Expand Down
82 changes: 34 additions & 48 deletions tmva/sofie/inc/TMVA/ROperator_ConvTranspose.icc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,22 @@ void ROperator_ConvTranspose<T>::Initialize(RModel& model){
fNBroadcastedB = fNB;
}
}

size_t kernelSize = 1;
size_t inputSize = 1;
for (size_t i = 0; i < fDim; i++) {
inputSize *= fShapeX[2+ i];
kernelSize *= fAttrKernelShape[i];
}

std::vector<size_t> shape1 = {fShapeW[0], fShapeW[1], kernelSize};
std::vector<size_t> shape2 = {fShapeW[1], kernelSize, inputSize};
model.AddIntermediateTensor(fNX +"_f", ConvertStringToType(fType), shape1 );
model.AddIntermediateTensor(fNX +"_xcol", ConvertStringToType(fType), shape2 );
fConvK = fNX +"_f";
fImcol = fNX +"_xcol";
fOutputTensorNames.emplace_back(fConvK);
fOutputTensorNames.emplace_back(fImcol);
}

template <typename T>
Expand All @@ -198,32 +214,6 @@ std::string ROperator_ConvTranspose<T>::GenerateInitCode()
return out.str();
}

// generate code for Session data members (e.g. internal vectors)
template <typename T>
std::string ROperator_ConvTranspose<T>::GenerateSessionMembersCode(std::string opName)
{

//size_t outputChannelSize = fShapeY[1];
size_t kernelSize = 1;
size_t inputSize = 1;
for (size_t i = 0; i < fDim; i++) {
inputSize *= fShapeX[2+ i];
kernelSize *= fAttrKernelShape[i];
}

opName = "op_" + opName;
std::stringstream out;
// matrix with convolution kernels
out << "std::vector<" << fType << "> fVec_" << opName << "_f = std::vector<" << fType << ">("
<< fShapeW[0] * fShapeW[1] * kernelSize << ");\n";
// output matrix of im2col
out << "std::vector<" << fType << "> fVec_" << opName << "_xcol = std::vector<" << fType << ">("
<< kernelSize * fShapeW[1] * inputSize << ");\n"; // kernelsize * output channel size * input size
out << "\n";

return out.str();
}

template <typename T>
std::string ROperator_ConvTranspose<T>::Generate(std::string OpName)
{
Expand Down Expand Up @@ -251,13 +241,11 @@ std::string ROperator_ConvTranspose<T>::Generate(std::string OpName)
out << "\n//---- operator ConvTranspose " << OpName << "\n";

// create first matrix with convolution kernels
if (fUseSession)
out << SP << fType << " * " << OpName << "_f = fVec_" << OpName << "_f.data();\n";
else {
if (!fUseSession) {
size_t kernelSize = fAttrKernelShape[0];
if (fDim > 1)
kernelSize *= fAttrKernelShape[1];
out << SP << fType << " " << OpName << "_f[" << fShapeW[0] * fShapeW[1] * kernelSize << "] = {0};\n";
out << SP << fType << " tensor_" << fNX << "_f[" << fShapeW[0] * fShapeW[1] * kernelSize << "] = {0};\n";
}

// vectorize the (dilated)convolution kernels into a matrix
Expand Down Expand Up @@ -295,7 +283,7 @@ std::string ROperator_ConvTranspose<T>::Generate(std::string OpName)
out << SP << SP << SP << "for (std::size_t kh = 0; kh < " << kHeight << "; kh++) {\n";
out << SP << SP << SP << SP << "for (std::size_t kw = 0; kw < " << kWidth << "; kw++) {\n";

out << SP << SP << SP << SP << SP << OpName << "_f[ic * " << ocstrideDil << " + oc * " << icstrideDil;
out << SP << SP << SP << SP << SP << "tensor_" << fNX << "_f[ic * " << ocstrideDil << " + oc * " << icstrideDil;
if (fDim > 2)
out << " + kd * " << dstrideDil;
if (fDim > 1)
Expand Down Expand Up @@ -329,10 +317,8 @@ std::string ROperator_ConvTranspose<T>::Generate(std::string OpName)
out << SP << "float " << OpName << "_alpha = 1.0;\n";
out << SP << "float " << OpName << "_beta = 0.0;\n";

if (fUseSession) {
out << SP << fType << " * " << OpName << "_xcol = fVec_" << OpName << "_xcol.data();\n";
} else {
out << SP << fType << " " << OpName << "_xcol[" << fShapeW[0]*icstrideDil * oDepth * oHeight * oWidth << "] = {0};\n";
if (!fUseSession) {
out << SP << fType << " tensor_" << fNX << "_xcol[" << fShapeW[0]*icstrideDil * oDepth * oHeight * oWidth << "] = {0};\n";
}

// Loop on batch size
Expand Down Expand Up @@ -380,14 +366,14 @@ std::string ROperator_ConvTranspose<T>::Generate(std::string OpName)
out << SP << SP << "BLAS::sgemm_(&" << OpName << "_transA, &" << OpName << "_transB, &" << OpName << "_m, &"
<< OpName << "_n, &" << OpName << "_k, &" << OpName << "_alpha, "
<< "tensor_" << fNX << " + x_offset, &" << OpName << "_m,\n"; // use m if op_xcol is not transpose , otherwise k
out << SP << SP << SP << OpName << "_f, &" << OpName << "_n, &" << OpName << "_beta, "
<< OpName << "_xcol, &" << OpName << "_m);\n";
out << SP << SP << SP << "tensor_" << fNX <<"_f, &" << OpName << "_n, &" << OpName << "_beta, tensor_"
<< fNX <<"_xcol, &" << OpName << "_m);\n";

// when using im2col - resulting matrix is transposed, is (input_c * filter_h * filter_w, output_h *
// output_w)
// before using col2im I need to transpose matrix
if (fDim < 3) {
out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::col2im<float>(" << OpName << "_xcol,"
out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::col2im<float>(tensor_" << fNX << "_xcol,"
// channels, height, width, kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w, dilation_h,
// dilation_w,
<< fShapeY[1] << "," << oHeight << "," << oWidth << ",";
Expand All @@ -410,14 +396,14 @@ std::string ROperator_ConvTranspose<T>::Generate(std::string OpName)
<< fShapeX[1] << "," << oDepth << "," << oHeight << "," << oWidth << "," << fAttrKernelShape[0] << ","
<< fAttrKernelShape[1] << "," << fAttrKernelShape[2] << "," << fAttrPads[0] << "," << fAttrPads[1] << ","
<< fAttrPads[2] << "," << fAttrStrides[0] << "," << fAttrStrides[1] << "," << fAttrStrides[2] << ","
<< fAttrDilations[0] << "," << fAttrDilations[1] << "," << fAttrDilations[2] << "," << OpName
<< "_xcol);\n\n ";
<< fAttrDilations[0] << "," << fAttrDilations[1] << "," << fAttrDilations[2] <<
",tensor_" << fNX <<"_xcol);\n\n ";
}
// // BLAS
// out << SP << SP << "BLAS::sgemm_(&" << OpName << "_transA, &" << OpName << "_transB, &" << OpName << "_m, &"
// << OpName << "_n, &" << OpName << "_k, &" << OpName << "_alpha, " << OpName << "_xcol, &" << OpName
// << OpName << "_n, &" << OpName << "_k, &" << OpName << "_alpha, tensor_" << fNX << "_xcol, &" << OpName
// << "_m,\n"; // use m if op_xcol is not transpose , otherwise k
// out << SP << SP << SP << OpName << "_f, &" << OpName << "_k, &" << OpName << "_beta, tensor_" << fNY
// out << SP << SP << SP <<"tensor_" << fNX << "_f, &" << OpName << "_k, &" << OpName << "_beta, tensor_" << fNY
// << " + out_offset, &" << OpName << "_m);\n";
} else {
// case of group transposed convolution
Expand All @@ -434,11 +420,11 @@ std::string ROperator_ConvTranspose<T>::Generate(std::string OpName)
<< OpName << "_n, &" << OpName << "_k, &" << OpName << "_alpha, "
<< "tensor_" << fNX << " + x_offset, &" << OpName
<< "_m,\n"; // use m if op_xcol is not transpose , otherwise k
out << SP << SP << SP << OpName << "_f, &" << OpName << "_n, &" << OpName
<< "_beta, " << OpName << "_xcol , &" << OpName << "_m);\n";
out << SP << SP << SP << "tensor_" << fNX << "_f, &" << OpName << "_n, &" << OpName
<< "_beta, tensor_" << fNX << "_xcol , &" << OpName << "_m);\n";

if (fDim < 3) {
out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::col2im<float>(" << OpName << "_xcol,"
out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::col2im<float>(tensor_" << fNX << "_xcol,"
// channels, height, width, kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w, dilation_h,
// dilation_w,
<< fShapeY[1] << "," << oHeight << "," << oWidth << ",";
Expand All @@ -462,17 +448,17 @@ std::string ROperator_ConvTranspose<T>::Generate(std::string OpName)
<< fShapeX[1] << "," << oDepth << "," << oHeight << "," << oWidth << "," << fAttrKernelShape[0] << ","
<< fAttrKernelShape[1] << "," << fAttrKernelShape[2] << "," << fAttrPads[0] << "," << fAttrPads[1] << ","
<< fAttrPads[2] << "," << fAttrStrides[0] << "," << fAttrStrides[1] << "," << fAttrStrides[2] << ","
<< fAttrDilations[0] << "," << fAttrDilations[1] << "," << fAttrDilations[2] << "," << OpName
<< fAttrDilations[0] << "," << fAttrDilations[1] << "," << fAttrDilations[2] << "," << "tensor_" << fNX
<< "_xcol);\n\n ";
}

// // BLAS
// // offset g must be g * k * n
// out << SP << SP << SP << "size_t offset_f = g * " << fShapeW[0] * fShapeW[1] * icstrideDil / fAttrGroup << ";\n";
// out << SP << SP << "BLAS::sgemm_(&" << OpName << "_transA, &" << OpName << "_transB, &" << OpName << "_m, &"
// << OpName << "_n, &" << OpName << "_k, &" << OpName << "_alpha, " << OpName << "_xcol, &" << OpName
// << OpName << "_n, &" << OpName << "_k, &" << OpName << "_alpha, tensor_" << fNX << "_xcol, &" << OpName
// << "_m,\n"; // use m if op_xcol is not transpose , otherwise k
// out << SP << SP << SP << OpName << "_f + offset_f, &" << OpName << "_k, &" << OpName << "_beta, tensor_" << fNY
// out << SP << SP << SP << "tensor_" << fNX << "_f + offset_f, &" << OpName << "_k, &" << OpName << "_beta, tensor_" << fNY
// << " + out_offset"
// << ", &" << OpName << "_m);\n";

Expand Down
Loading