Repository: microsoft/onnxruntime
File: onnxruntime/core/mlas/lib/convolve.cpp (C++)

Code before:
    ptrdiff_t Index
)
{
    MLAS_CONV_WORK_BLOCK* WorkBlock = (MLAS_CONV_WORK_BLOCK*)Context;

    const MLAS_CONV_PARAMETERS* Parameters = WorkBlock->Parameters;

    const size_t GroupCount = Parameters->GroupCount;
    const size_t BatchGroupCount = Parameters->BatchCount * GroupCount;

    const size_t TargetThreadCount = WorkBlock->TargetThreadCount;

    const size_t BatchGroupCountPerThread = BatchGroupCount / TargetThreadCount;
    const size_t BatchGroupCountExtra = BatchGroupCount % TargetThreadCount;

    size_t BatchGroupStart;
    size_t BatchGroupEnd;

    if (static_cast<uint32_t>(Index) < BatchGroupCountExtra) {
        BatchGroupStart = (BatchGroupCountPerThread + 1) * Index;
        BatchGroupEnd = BatchGroupStart + BatchGroupCountPerThread + 1;
    } else {
        BatchGroupStart = BatchGroupCountPerThread * Index + BatchGroupCountExtra;
        BatchGroupEnd = BatchGroupStart + BatchGroupCountPerThread;
    }

    const size_t FilterCount = Parameters->FilterCount;
    const size_t OutputSize = Parameters->OutputSize;
    const size_t K = Parameters->K;

    const size_t InputGroupSize = Parameters->InputChannels * Parameters->InputSize;
    const size_t OutputGroupSize = FilterCount * OutputSize;
    const size_t FilterGroupSize = FilterCount * K;

    for (size_t bg = BatchGroupStart; bg < BatchGroupEnd; bg++) {
        size_t group = bg % GroupCount;

        const float* input = WorkBlock->Input + bg * InputGroupSize;
        const float* filter = WorkBlock->Filter + group * FilterGroupSize;
        float* output = WorkBlock->Output + bg * OutputGroupSize;
        const float* bias = WorkBlock->Bias;
        if (bias != nullptr) {
            bias += group * FilterCount;
        }
        float* ColumnBuffer = WorkBlock->WorkingBuffer + Index * OutputSize * K;

        MlasConvOperation(Parameters, input, filter, bias, ColumnBuffer, output, 0, OutputSize);
    }
}

inline

Review comment (suggestion):
could we add some documentation for this function? in particular, it might be useful to describe how it differs from `MlasConvTryMultithread()`

Code after:
    ptrdiff_t Index
)
/*++

Routine Description:

    This routine is invoked from a worker thread to execute a segment of a
    convolution operation.

    If using this, the entire convolution operation is parallelized on the
    (batch size * group count) parameter and this routine has logic to
    perform a specific thread's shard of the entire Convolution operation.

Arguments:

    Context - Supplies the pointer to the context for the threaded operation.

    Index - Supplies the current index of the threaded operation.

Return Value:

    None.

--*/

{
    MLAS_CONV_WORK_BLOCK* WorkBlock = (MLAS_CONV_WORK_BLOCK*)Context;

    const MLAS_CONV_PARAMETERS* Parameters = WorkBlock->Parameters;

    const size_t GroupCount = Parameters->GroupCount;
    const size_t BatchGroupCount = Parameters->BatchCount * GroupCount;

    const size_t TargetThreadCount = WorkBlock->TargetThreadCount;

    const size_t BatchGroupCountPerThread = BatchGroupCount / TargetThreadCount;
    const size_t BatchGroupCountExtra = BatchGroupCount % TargetThreadCount;

    size_t BatchGroupStart;
    size_t BatchGroupEnd;

    if (static_cast<size_t>(Index) < BatchGroupCountExtra) {
        BatchGroupStart = (BatchGroupCountPerThread + 1) * Index;
        BatchGroupEnd = BatchGroupStart + BatchGroupCountPerThread + 1;
    } else {
        BatchGroupStart = BatchGroupCountPerThread * Index + BatchGroupCountExtra;
        BatchGroupEnd = BatchGroupStart + BatchGroupCountPerThread;
    }

    const size_t FilterCount = Parameters->FilterCount;
    const size_t OutputSize = Parameters->OutputSize;