Add dynamic shape support for ONNX Conv to Linalg lowering

[kimm240]

Changes

• Added IndexExprBuilderForLinalg
• Integrated ShapeHelper for dynamic shape computation
• Enhanced tensor::EmptyOp creation

Testing

• Added conv_dynamic test case in test/mlir/conversion/onnx_to_linalg/NN/Conv.mlir
  • Tests dynamic input shapes: tensor<?x3x?x?xf32>
  • Verifies dynamic output shapes: tensor<?x2x?x?xf32>
  • Checks that tensor.dim, affine.apply, and tensor.empty with dynamic sizes are generated correctly

Example

Before (static only):

func.func @conv(%arg0: tensor<1x3x5x5xf32>, %arg1: tensor<2x3x3x3xf32>)
    -> tensor<1x2x3x3xf32> {
  %0 = tensor.empty() : tensor<1x2x3x3xf32>
  // ...
}

After (dynamic support):

func.func @conv(%arg0: tensor<?x3x?x?xf32>, %arg1: tensor<2x3x3x3xf32>)
    -> tensor<?x2x?x?xf32> {
  %dim0 = tensor.dim %arg0, %c0 : tensor<?x3x?x?xf32>
  %dim2 = tensor.dim %arg0, %c2 : tensor<?x3x?x?xf32>
  %out_h = affine.apply #map(%dim2)
  %dim3 = tensor.dim %arg0, %c3 : tensor<?x3x?x?xf32>
  %out_w = affine.apply #map1(%dim2, %dim3)
  %0 = tensor.empty(%dim0, %out_h, %out_w) : tensor<?x2x?x?xf32>
  // ...
}

Related

• Related to Optimize ConvOp on CPU #3354 (ONNX Conv to Linalg lowering PR)

[jenkins-droid]

Can one of the admins verify this patch?

[chentong319]

In your example, could you include all the details, such as constant %c* and maps? If a dimension size is not a affine function (not in ConvOp case), what dialect can be used to express the expression? Or when the output tensor shape of an Op can not be expressed with affine, this Op can not be lowered to Linalg. Is it true?

[kimm240]

Hi, @chentong319 ! Thanks for your review.
Here is an answer.

1. Example with All Details (Constants %c* and Maps)

Here is the complete example with all details from the actual generated code:

#map = affine_map<(d0) -> (d0 - 2)>
#map1 = affine_map<(d0, d1) -> (d1 - 2)>
module {
  func.func @conv_dynamic(%arg0: tensor<?x3x?x?xf32>, %arg1: tensor<2x3x3x3xf32>) -> tensor<?x2x?x?xf32> {
    %cst = arith.constant 0.000000e+00 : f32
    %c3 = arith.constant 3 : index
    %c2 = arith.constant 2 : index
    %c0 = arith.constant 0 : index
    %dim = tensor.dim %arg0, %c0 : tensor<?x3x?x?xf32>
    %dim_0 = tensor.dim %arg0, %c2 : tensor<?x3x?x?xf32>
    %0 = affine.apply #map(%dim_0)
    %dim_1 = tensor.dim %arg0, %c3 : tensor<?x3x?x?xf32>
    %1 = affine.apply #map1(%dim_0, %dim_1)
    %2 = tensor.empty(%dim, %0, %1) : tensor<?x2x?x?xf32>
    %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x2x?x?xf32>) -> tensor<?x2x?x?xf32>
    %4 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<?x3x?x?xf32>, tensor<2x3x3x3xf32>) outs(%3 : tensor<?x2x?x?xf32>) -> tensor<?x2x?x?xf32>
    return %4 : tensor<?x2x?x?xf32>
  }
}

You can check it by running ./build/Release/bin/onnx-mlir-opt --convert-onnx-to-linalg='linalg-ops=onnx.Conv' test/mlir/conversion/onnx_to_linalg/NN/Conv.mlir -split-input-file 2>&1.

2. Non-Affine Dimension Size Expressions

When a dimension size is not an affine function (e.g., division by a variable, complex expressions), the IndexExpr library automatically falls back to the arith dialect to express the computation.

How it works:

The IndexExpr library (used by IndexExprBuilderForLinalg) has built-in automatic detection:

1. Automatic Affine Detection (IndexExpr.cpp:575-576):

   bool resIsAffine = resIsLit || (canBeAffine && isAffine() && b.isAffine() &&
                                    (!affineWithLitB || b.isLiteral()));

2. Fallback Mechanism (IndexExpr.cpp:882-884, 926-928, 958-961):

   • Affine case: Uses affine::AffineApplyOp (as shown in the example above)
   • Non-affine case: Automatically creates NonAffineIndexExpr and uses MathBuilder to generate arith dialect operations.

3. Can Operations with Non-Affine Output Shapes be Lowered to Linalg?

Answer: Yes, they can be lowered to Linalg.

The key insight is that Linalg operations themselves use affine indexing maps, but the output tensor shape computation (which happens before creating the tensor) is separate and can use any dialect operations.

1. IndexExprImpl::getValue() (IndexExprDetail.cpp):

   • If the expression is affine: creates affine::AffineApplyOp
   • If the expression is non-affine: the NonAffineIndexExpr already has a Value (created by arith operations), which is returned directly
   • This Value can be used as input to tensor::EmptyOp regardless of whether it came from affine or non-affine computation

2. My implementation (Conv.cpp):

   dynamicSizes.push_back(outputDims[i].getValue());

   The getValue() call automatically handles both cases:

   • Affine → affine::AffineApplyOp → returns Value
   • Non-affine → arith operations → returns Value
   • Both can be passed to tensor::EmptyOp::create()

3. Linalg operations only require that:

   • The tensor types are known (can have dynamic dimensions)
   • The indexing maps are affine (which they are, as they operate on the tensor elements, not the shape)

[chentong319]

I like the idea of extension IndexExpr. Shall we discuss on adding the destination-passing interface on onnx ops in issue #3355? We can schedule a meeting to discuss how to add DestinationStyleOpInterface to onnx op. We can refactor the code that generate the output memref for ops in onnx-to-krnl conversion into reifyResultShapes of onnx ops. That code can be shared by multiple users. You do not need to deal it only for Linalg.

[kimm240]

That sounds like a great plan, @chentong319!
To ensure a productive discussion, could we start the technical dialogue on Slack or within Issue #3355 first? We can share a more detailed design proposal or an RFC-style draft regarding the DestinationStyleOpInterface and #3355.
I’m already in the #onnx-mlir-discussion Slack channel as 'hyun gyu kim'. Discussing it there first will help us align on the core logic asynchronously!

[chentong319]

I lost my slack account for onnx-mlir-discussion. Let's start discussion in issue 3355.