Skip to content

[tmva][sofie] Softmax contiguous memory fast path #19743

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

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

[tmva][sofie] Softmax contiguous memory fast path #19743

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5960569
Softmax contiguous memory
Jul 23, 2025
1848c97
comments solved
Aug 27, 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
212 changes: 63 additions & 149 deletions tmva/sofie/inc/TMVA/ROperator_Softmax.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,199 +58,113 @@ public:
throw std::runtime_error("TMVA SOFIE Operator Softmax called to Generate without being initialized first");
}
std::stringstream out;
int size = fShape.size();
auto length = ConvertDimShapeToLength(fShape);
auto stride = UTILITY::ComputeStrideFromShape(fShape);
int axis = fAttrAxis < 0 ? size + fAttrAxis : fAttrAxis;
out << "\n" << SP << "//------ SOFTMAX - " << size << " " << length << " " << axis << "\n";
// use safe numerically implementation by subtracting max of tensor
if (size == 1) {
out << SP << fType << " vmax = tensor_" << fNX << "[0];\n";
out << SP << "for (size_t i = 1; i < " << length << " ; i++){\n";
out << SP << SP << "if (tensor_" << fNX << "[i] > vmax) vmax = tensor_" << fNX << "[i];\n";
out << SP << "}\n";
out << SP << fType << " sum = 0.0;\n";
out << SP << "for (size_t i = 0; i < " << length << " ; i++){\n";
out << SP << SP << "tensor_" << fNY << "[i] = std::exp(tensor_" << fNX << "[i] - vmax);\n";
out << SP << SP << "sum += tensor_" << fNY << "[i];\n";
out << SP << "}\n";
out << SP << "for (size_t i = 0; i < " << length << " ; i++){\n";
out << SP << SP << "tensor_" << fNY << "[i] /= sum;\n";
size_t size = fShape.size();
auto length_str = ConvertDimShapeToLength(fShape);
size_t axis = fAttrAxis < 0 ? size + fAttrAxis : fAttrAxis;

// Check if this is the special case where memory is contiguous.
if (axis == size - 1) {
std::string axis_size = fShape[axis].GetVal();
std::string num_rows;
if (IsInteger(length_str) && IsInteger(axis_size)) {
num_rows = std::to_string(std::stoul(length_str) / std::stoul(axis_size));
} else {
num_rows = "(" + length_str + ") / (" + axis_size + ")";
}

out << "\n" << SP << "//------ SOFTMAX - " << size << " " << length_str << " " << axis << "\n";
out << SP << "for (int i = 0; i < " << num_rows << "; ++i) {\n";
out << SP << SP << "size_t offset = i * " << axis_size << ";\n";
out << SP << SP << fType << " const * x_ptr = &tensor_" << fNX << "[offset];\n";
out << SP << SP << fType << " * y_ptr = &tensor_" << fNY << "[offset];\n";

out << SP << SP << fType << " vmax = x_ptr[0];\n";
out << SP << SP << "for (int j = 1; j < " << axis_size << "; ++j) {\n";
out << SP << SP << SP << "if (x_ptr[j] > vmax) vmax = x_ptr[j];\n";
out << SP << SP << "}\n";

out << SP << SP << fType << " sum = 0.0;\n";
out << SP << SP << "for (int j = 0; j < " << axis_size << "; ++j) {\n";
out << SP << SP << SP << "y_ptr[j] = std::exp(x_ptr[j] - vmax);\n";
out << SP << SP << SP << "sum += y_ptr[j];\n";
out << SP << SP << "}\n";

out << SP << SP << fType << " inv_sum = 1.0f / sum;\n";
out << SP << SP << "for (int j = 0; j < " << axis_size << "; ++j) {\n";
out << SP << SP << SP << "y_ptr[j] *= inv_sum;\n";
out << SP << SP << "}\n";
out << SP << "}\n";

} else {
int k = 0;
auto stride = UTILITY::ComputeStrideFromShape(fShape);
size_t k = 0;
std::vector<std::string> l(size);
for (int i = 0; i < size; i++) {
for (size_t i = 0; i < size; i++) {
if (i != axis) {
for (int j = 0; j < k; j++) out << SP;
for (size_t j = 0; j < k; j++) out << SP;
l[i] = std::string("i") + std::to_string(i);
out << "for (int " << l[i] << " = 0; " << l[i] << " < " << fShape[i] << "; " << l[i] << "++) {\n";
k++;
}
}
for (int j = 0; j < size-1; j++) out << SP;
for (size_t j = 0; j < size-1; j++) out << SP;
out << fType << " sum = 0.;\n";
for (int j = 0; j < size-1; j++) out << SP;
for (size_t j = 0; j < size-1; j++) out << SP;
out << "size_t index = ";
for (int i = 0; i < size; i++) {
bool first = true;
for (size_t i = 0; i < size; i++) {
if (i == axis) continue;
if ((i > 0 && axis != 0) || i > 1 ) out << "+";
if (!first) out << " + ";
if (stride[i].GetVal() != "1")
out << stride[i] << "*";
out << l[i];
first = false;
}
out << ";\n";
// find maximum looping along reduced axix
for (int j = 0; j < size-1; j++) out << SP;
// find maximum looping along reduced axis
for (size_t j = 0; j < size-1; j++) out << SP;
out << fType << " vmax = tensor_" << fNX << "[index];\n";
for (int j = 0; j < size-1; j++) out << SP;
out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n";
for (int j = 0; j < size; j++) out << SP;
for (size_t j = 0; j < size-1; j++) out << SP;
out << "for (int i = 1; i < " << fShape[axis] << "; i++) {\n";
for (size_t j = 0; j < size; j++) out << SP;
out << fType << " x = tensor_" << fNX << "[index + i";
if (stride[axis].GetVal() != "1") out << "*(" << stride[axis] << ")";
out << "];\n";
for (int j = 0; j < size; j++) out << SP;
for (size_t j = 0; j < size; j++) out << SP;
out << "if (x > vmax) vmax = x;\n";
for (int j = 0; j < size-1; j++) out << SP;
for (size_t j = 0; j < size-1; j++) out << SP;
out << "}\n";
// compute softmax
for (int j = 0; j < size-1; j++) out << SP;
for (size_t j = 0; j < size-1; j++) out << SP;
out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n";
for (int j = 0; j < size; j++) out << SP;
for (size_t j = 0; j < size; j++) out << SP;
out << "size_t id = index + i";
if (stride[axis].GetVal() != "1") out << "*(" << stride[axis] << ")";
out << ";\n";
for (int j = 0; j < size; j++) out << SP;
for (size_t j = 0; j < size; j++) out << SP;
out << "tensor_" << fNY << "[id] = std::exp(tensor_" << fNX << "[id] - vmax);\n";
for (int j = 0; j < size; j++) out << SP;
for (size_t j = 0; j < size; j++) out << SP;
out << "sum += tensor_" << fNY << "[id];\n";
for (int j = 0; j < size-1; j++) out << SP;
for (size_t j = 0; j < size-1; j++) out << SP;
out << "}\n";
// normalize
for (int j = 0; j < size-1; j++) out << SP;
for (size_t j = 0; j < size-1; j++) out << SP;
out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n";
for (int j = 0; j < size; j++) out << SP;
for (size_t j = 0; j < size; j++) out << SP;
out << "tensor_" << fNY << "[index + i";
if (stride[axis].GetVal() != "1") out << "*(" << stride[axis] << ")";
out << "] /= sum;\n";
for (int j = 0; j < size-1; j++) out << SP;
for (size_t j = 0; j < size-1; j++) out << SP;
out << "}\n";
//end loops
for (int i = size-2; i >=0; i--) {
for (int i = static_cast<int>(k) - 1; i >= 0; i--) {
for (int j = 0; j < i; j++) out << SP;
out << "}\n";
}

#if 0
size_t batch = fShape[0];
size_t channel = fShape[1];
size_t width = (size > 2) ? fShape[size - 1] : 1;
size_t height = (size > 3) ? fShape[size - 2] : 1;
size_t depth = (size > 4) ? fShape[size - 3] : 1;
size_t hStride = width;
size_t dStride = height * width;
size_t cStride = depth * dStride;
size_t bStride = channel * cStride;

size_t N = 0; // Size of the axis
size_t iStride = 0;
if (axis == 0) {
N = batch;
iStride = bStride;
} else if (axis == 1) {
N = channel;
iStride = cStride;
} else if (axis == size - 1) {
N = width;
iStride = 1;
} else if (size > 3 && axis == size - 2) {
N = height;
iStride = hStride;
} else if (size == 5 && axis == size - 3) {
N = depth;
iStride = dStride;
} else {
throw
std::runtime_error("TMVA::SOFIE - Softmax operator along the axis "
+ std::to_string(fAttrAxis) + " with " + std::to_string(size)
+ "d input tensor not supported.");
}

bool notBatch = axis != 0;
bool notChannel = axis != 1;
bool notDepth = (size == 5 && axis != 2);
bool notHeight = (size == 5 && axis != 3) || (size == 4 && axis != 2);
bool notWidth = (size == 5 && axis != 4) || (size == 4 && axis != 3) || (size == 3 && axis != 2);

if (notBatch) {
out << SP << "for (size_t n = 0; n < " << batch << " ; n++){\n";
}
if (notChannel) {
out << SP << SP << "for (size_t c = 0; c < " << channel << " ; c++){\n";
}
if (notDepth) {
out << SP << SP << "for (size_t d = 0; d < " << depth << " ; d++){\n";
}
if (notHeight) {
out << SP << SP << "for (size_t h = 0; h < " << height << " ; h++){\n";
}
if (notWidth) {
out << SP << SP << "for (size_t w = 0; w < " << width << " ; w++){\n";
}
out << SP << SP << SP << fType << " sum = 0.;\n";
out << SP << SP << SP << "size_t index = 0";
if (notBatch) {
out << " + n * " << bStride;
}
if (notChannel) {
out << "+ c * " << cStride;
}
if (notDepth) {
out << " + d * " << dStride;
}
if (notHeight) {
out << " + h * " << hStride;
}
if (notWidth) {
out << " + w";
}
out << ";\n";
// apply softmax along the axis - find first maximum value for numerical stability
if (N == 0)
throw std::runtime_error("TMVA::SOFIE - Softmax operator is along axis with zero elements");
out << SP << SP << SP << fType << " vmax = tensor_" << fNX << "[index];\n";
out << SP << SP << SP << "for (size_t i = 1; i < " << N << "; i++) {\n";
out << SP << SP << SP << SP << "if (tensor_" << fNX << "[index + i*" << iStride << "] > vmax)\n";
out << SP << SP << SP << SP << SP << "vmax = tensor_" << fNX << "[index + i*" << iStride << "];\n";
out << SP << SP << SP << "}\n";
out << SP << SP << SP << "for (size_t i = 0; i < " << N << "; i++) {\n";
out << SP << SP << SP << SP << "tensor_" << fNY << "[index + i*" << iStride << "] = std::exp(tensor_" << fNX
<< "[index + i*" << iStride << "] - vmax);\n";
out << SP << SP << SP << SP << "sum += tensor_" << fNY << "[index + i*" << iStride << "];\n";
out << SP << SP << SP << "}\n";
out << SP << SP << SP << "for (size_t i = 0; i < " << N << "; i++) {\n";
out << SP << SP << SP << SP << "tensor_" << fNY << "[index + i*" << iStride << "] /= sum;\n";
out << SP << SP << SP << "}\n";
if (notWidth) {
out << SP << SP << "}\n"; // end w
}
if (notHeight) {
out << SP << SP << "}\n"; // end h
}
if (notDepth) {
out << SP << SP << "}\n"; // end d
}
if (notChannel) {
out << SP << SP << "}\n"; // end c
}
if (notBatch) {
out << SP << "}\n"; // end n
}

#endif
}
return out.str();
}
std::vector<std::string> GetStdLibs() override { return { std::string("cmath") }; }
};

} // namespace SOFIE
Expand Down
Loading