Combine parallel dense Optimization pass in ONNX Dialect

[Arkar-Hema]

Combine Parallel Dense

CombineParallelDense is an optimization pass designed to merge multiple parallel ONNXGemmOp (Dense/Fully Connected) operations into a single, more efficient Dense layer. This optimization reduces redundant computations, improves memory efficiency, and enhances hardware utilization.

The pass identifies Dense (Gemm) operations that:

• Share the same input tensor.
• Have identical attributes such as alpha, beta, transA and transB (ensuring compatibility).
• May have different output dimensions (number of neurons) but maintain compatible weight shapes for concatenation.

Lets assume a input case:

• Input Shape: (1, 512)
• Dense A: out_features = 256
• Dense B: out_features = 128
• Dense C: out_features = 64
• Attributes: transB = 0, alpha = 1.0, beta = 1.0

Before Optimization (Three Parallel Gemms)

• Each GEMM does one full matrix multiplication (1×512 × 512×N)
• Three separate weight and bias tensors and produces three outputs
  -Memory Reads: 3 times full input (one for each gemm)
• Post-processing: A Concat(axis=1) merges them into one output: Y (1×448)

After Optimization (Combined Dense)

• Total Output Features: 256 + 128 + 64 = 448
• All weights are concatenated along output channel axis → New weight shape: (512, 448)
• Biases are also concatenated
• A single ONNXGemmOp computes Y (1×448) directly

Improvement in performance metrics

Latency Improvement: 7-15%
Throughput Improvement: 8-14%
Memory Usage Improvement: 10-12%

[jenkins-droid]

Can one of the admins verify this patch?

[jenkins-droid]

Can one of the admins verify this patch?

[jenkins-droid]

Can one of the admins verify this patch?

[tungld]

@Arkar-Hema A general question: in what kind of models have you seen this kind of pattern: multiple Gemm ops followed by a Concat op? and also similar patterns you have recently created PRs for? Just curious on how practical it is. Thanks!

[Arkar-Hema]

@Arkar-Hema A general question: in what kind of models have you seen this kind of pattern: multiple Gemm ops followed by a Concat op? and also similar patterns you have recently created PRs for? Just curious on how practical it is. Thanks!

• Models with the CombineParallelDense pattern (Combine parallel dense Optimization pass in ONNX Dialect #3123):
  These contain multiple Gemm ops, though not always followed by a Concat. I added the Concat condition to the pass so it would still handle those cases gracefully if present. Some models with this pattern include:

1. Bertsquad-8
2. Bertsquad-10
3. Bertsquad-12
4. FasterRCNN-10

• Models with the Reorder ReLU and MaxPool pattern (Reorder relu to maxpool optimization pass in ONNX dialect #3109): This pattern shows up frequently in vision models like:

1. ResNet101-DUC-12
2. ResNet101-DUC-7
3. emotion-ferplus models
4. caffenet models
5. Densenet models
6. googlenet models
7. inception models
8. rcnn-ilsvrc13 models
9. resnet models
10. vgg models

• In the yolo models shows up the merge concat pattern (Merge nested concat Ops optimization pass in ONNX dialect #3111)

• Models with the CombineParallelConv pattern (Combine Parallel Convolution optimization pass in ONNX Dialect #3116): This shows up in many convolution-heavy models like:

1. retinanet models
2. version-RFB-320
3. version-RFB-640
4. googlenet models
5. inception models
6. resnet models
7. squeezenet models

[Arkar-Hema]

@tungld could you please verify this patch?

[tungld]

@Arkar-Hema thank you for the information!

I have some general comments:

• I think that when multiple GEMM ops are followed by a concat, the performance in theory would be better. But, could you run with multiple input sizes to see how the performance benefit in practice?
• When multiple GEMM ops are NOT followed by a concat (this is the case for the models you listed), you need a split and I think the split axis is the innermost dimension. I am not sure how slow the split is and if we can get speedup or not. Could you do a performance comparison to see if you can achieve speedup in this case?
• Are you targeting optimization for CPU or it is beneficial for AI accelerators as well given that AI accelerators may use special data layout which may be not convenient for concat or split.

Thanks.

[Arkar-Hema]

I ran performance benchmarks across a range of input sizes for both the GEMM → Concat and the Combined GEMM → Split cases. Results show that:

• In the Concat case, the optimization provides consistent Latency improvement of 2-7%, and throughput improvement of 1-5%
  
  
  
  

• In the cases where it splits, the optimization provides consistent Latency improvement of 1-7%, and throughput improvement of 1-8%
  
  
  
  

• I’ve currently targeted this pass primarily for CPU backends only.

[tungld]

Thanks @Arkar-Hema for the experiments! Did you compile your programs with -O3?

Since this parallel fusion may not work for accelerators, could you create a compile option to enable this if needed, for example -fuse-parallel-onnx-gemm?

I don't think you need to handle the case where there is a concat after multiple gemms. Just emit a split op, then later you can write a simple canonicalization rule for concat to fuse Split -> Concat.

Below are my first-round comments, most of them are for simplifying the code, making it easy to follow. However, the important thing is you need to check the input C carefully because it's broadcastable.

[jenkins-droid]

Can one of the admins verify this patch?

[Arkar-Hema]

Thanks @Arkar-Hema for the experiments! Did you compile your programs with -O3?

Since this parallel fusion may not work for accelerators, could you create a compile option to enable this if needed, for example -fuse-parallel-onnx-gemm?

I don't think you need to handle the case where there is a concat after multiple gemms. Just emit a split op, then later you can write a simple canonicalization rule for concat to fuse Split -> Concat.

Below are my first-round comments, most of them are for simplifying the code, making it easy to follow. However, the important thing is you need to check the input C carefully because it's broadcastable.

I have added it, Thanks

[jenkins-droid]

Can one of the admins verify this patch?

[jenkins-droid]

Can one of the admins verify this patch?

[jenkins-droid]

Can one of the admins verify this patch?

[jenkins-droid]

Can one of the admins verify this patch?

[AlexandreEichenberger]

@jenkins-droid test this please

[jenkins-droid]

Can one of the admins verify this patch?

[jenkins-droid]

Can one of the admins verify this patch?

[jenkins-droid]

Can one of the admins verify this patch?

[jenkins-droid]

Can one of the admins verify this patch?

[jenkins-droid]

Can one of the admins verify this patch?

[jenkins-droid]

Can one of the admins verify this patch?

[AlexandreEichenberger]

@jenkins-droid test this please

[tungld]

Hi @Arkar-Hema When addressing a comment, could you please provide a brief explanation of how you did so? This will make the review process easier. Thanks!

[jenkins-droid]

Can one of the admins verify this patch?

[tungld]

Both Biases as tensor<1xf32>:
If both biases are of shape tensor<1xf32>, I now check their corresponding Gemm output shapes and ensure their output channels (last dimension) match before considering them compatible. If they differ, the function returns false, as merging them without this check would be invalid.

It does not make sense to me how this can solve the problem. You must check there is no broadcasting here, say the last dim in the output must be 1 also, for example:

if (aOutputShape.back() != 1 || bOutputShape.back() != 1)
   return false;

Also, please do add a lit test for this case, to make sure gemm ops are not merged.

[tungld]

Replace newWeightType by unrankedTensorType . It's redundant to define newWeightType.

[tungld]

Replace newBiasType by unrankedTensorType . It's redundant to define newBiasType .

[tungld]

Replace these by:

ArrayRef<int64_t> aOutputShape = getShape(a.getY().getType());
ArrayRef<int64_t> bOutputShape = getShape(b.getY().getType());

[tungld]

Please replace this by

Value newGemmOutput = create.onnx.gemm(unrankedTensorType, input,
        newWeight, newBias, gemmOp1.getAlphaAttr(), gemmOp1.getBetaAttr(),
        gemmOp1.getTransAAttr(), gemmOp1.getTransBAttr());

[tungld]

Please use: ArrayRef<int64_t> gemmShape = getShape(gemm.getY().getType());

[tungld]

Please replace this by

    SmallVector<Type, 4> splitTypes(splitSizes.size(), unrankedTensorType);
    ValueRange splitResults = create.onnx.split(
        splitTypes, newGemmOutput, create.onnx.constantInt64(splitSizes), splitAxis);

[tungld]

Check if gemmShape[splitAxis] is a static dimension in areCompatible() function. Otherwise, it fails to create a constant tensor here.

[tungld]

When aC is None, do check that the last dim of aOutput is static. Otherwise, it fails when you create a constant tensor of zeros in the later code that you use the last dim of aOutput.

Check the same thing for bC.

Please add a list test for the case where aC or bC is None.

[tungld]

Check this earlier just after you collect all parallelGemms. The reason is you have a return failure() here which may interrupt the whole rewriting while you created new weight, new bias, and new gemm. Moving this check earlier before creating any new ops would make the IR clean.