[Linalg] Add matchers to infer which Convolution Op a given linalg.generic is

[Abhishek-Varma]

-- This commit adds matchers to infer which linalg.*conv/pooling* ops a given linalg.generic op is.
-- The goal is directed towards addressing the aim of [RFC] Op explosion in Linalg iteratively for *conv*/*pooling* ops.
-- It makes use of the matchers as part of linalg-specialize-generic-ops to specialize a linalg.generic to the corresponding linalg.*conv/pooling* op and thus demonstrates the matchers' working.

Signed-off-by: Abhishek Varma abhvarma@amd.com

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[Copilot]

Pull Request Overview

This PR adds comprehensive pattern matching capabilities to identify and specialize generic linalg operations to specific convolution and pooling named operations. The implementation includes matchers for numerous convolution variants (1D, 2D, 3D), depthwise convolutions, and pooling operations across different data layouts, enabling automatic specialization of linalg.generic operations to their corresponding named operation counterparts.

• Adds 45+ matcher functions for various convolution and pooling operation patterns
• Implements specialization logic to convert linalg.generic operations to named convolution/pooling operations
• Provides comprehensive test coverage with roundtrip tests for all supported operation variants

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
mlir/test/Dialect/Linalg/roundtrip-linalg-convolution-named-ops.mlir	New test file with comprehensive roundtrip tests for all convolution/pooling operation variants
mlir/lib/Dialect/Linalg/Utils/Utils.cpp	Core implementation of pattern matching logic for convolution and pooling operations
mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp	Integration of convolution matchers into the specialization pipeline
mlir/include/mlir/Dialect/Linalg/Utils/Utils.h	Public API declarations for all convolution/pooling matcher functions

[Copilot]

The CHECK comment indicates linalg.pooling_nhwc_min but the function being tested is pooling_nhwc_min_unsigned. This appears to be a copy-paste error that should be corrected to match the actual expected operation name.

Suggested change

	// CHECK: linalg.pooling_nhwc_min
	// CHECK: linalg.pooling_nhwc_min_unsigned

[Abhishek-Varma]

So I wanted to flag this for max/min_unsigned ops - glad Copilot reminded me!

Step 1: Use the following input IR :-

func.func @pooling_nhwc_max_unsigned_float(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?xf32>, %init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
  %0 = linalg.pooling_nhwc_max_unsigned {dilations = dense<1> : tensor<2xi64>,
                                strides = dense<1> : tensor<2xi64>}
     ins (%input, %filter: tensor<?x?x?x?xf32>, tensor<?x?xf32>)
    outs (%init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
  return %0 : tensor<?x?x?x?xf32>
}
func.func @pooling_nhwc_max_signed_float(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?xf32>, %init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
  %0 = linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>,
                                strides = dense<1> : tensor<2xi64>}
     ins (%input, %filter: tensor<?x?x?x?xf32>, tensor<?x?xf32>)
    outs (%init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
  return %0 : tensor<?x?x?x?xf32>
}

Step 2: Run mlir-opt --linalg-generalize-named-ops input.mlir

We'll observe the following :-

  func.func @pooling_nhwc_max_unsigned_float(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
    %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%arg0, %arg1 : tensor<?x?x?x?xf32>, tensor<?x?xf32>) outs(%arg2 : tensor<?x?x?x?xf32>) {
    ^bb0(%in: f32, %in_0: f32, %out: f32):
      %1 = arith.maximumf %out, %in : f32
      linalg.yield %1 : f32
    } -> tensor<?x?x?x?xf32>
    return %0 : tensor<?x?x?x?xf32>
  }
  func.func @pooling_nhwc_max_signed_float(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
    %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%arg0, %arg1 : tensor<?x?x?x?xf32>, tensor<?x?xf32>) outs(%arg2 : tensor<?x?x?x?xf32>) {
    ^bb0(%in: f32, %in_0: f32, %out: f32):
      %1 = arith.maximumf %out, %in : f32
      linalg.yield %1 : f32
    } -> tensor<?x?x?x?xf32>
    return %0 : tensor<?x?x?x?xf32>
  }

Since both the linalg.generic ops are IDENTICAL - there's no way to distinguish between these two and therefore we get linalg.pooling_nhwc_max (a less constrained version) for both the cases when using --linalg-specialize-generic-ops. Basically the matcher, currently, will match both - just the matter of which one is preferred more (as is demonstrated via Specialize.cpp).

NOTE: This only happens when the input is a float. In case of integer, the body of the linalg.generic has either arith.maxui or arith.maxsi - therefore let's us distinguish between linalg.pooling_nhwc_max_unsigned and linalg.pooling_nhwc_max respectively.

CC: @hanhanW @rengolin - let me know how do we want to proceed in this ops' case.

[banach-space]

Surely linalg.pooling_nhwc_max_unsigned should only allow integers, no?

[rengolin]

Could we avoid the explosion in exported function names? They're almost identical, modulo the one thing different.

[hanhanW]

Thanks @Abhishek-Varma for the patch. I'll help review it. Let's add some motivation for the change; I'll add my thoughts below.

Could we avoid the explosion in exported function names? They're almost identical, modulo the one thing different.

The PR is related to [RFC] Op explosion in Linalg. We are making progress on matmul side in upstream. E.g., the transpose variants are removed from named ops; we have matchers for the deprecated ops. [https://github.com//pull/147961]. The revision aims to do something similar, especially for matcher parts. However, we don't have generic conv/depthwise_conv/pooling ops today. Thus, the implementation looks different. I did a quick survey in the beginning, and I felt that the "explosion" is likely needed in the transition period. They are stronger matchers comparing to isa<linalg::*Conv*>. I'm happy to learn more if you have an idea about avoiding the explosion. My naive thinking is that we'll need this anyway in the transition period, like other matmul matchers.

The motivation of the change: our goal is to enhance existing patterns to work with generic op form; it requires matchers. Hopefully, the matching logic can be reused when we have the generic conv/depthwise_conv/pooling ops.

Mirroring some context from IREE issue, and the plan of upstream pattern enhancement:

• DecomposeConvolutionToLowerDimOpsPass: this is the pass that downscale 2D convolution to 1D convolution when either H or W dimensions are ones. The work include updating the upstream decompose convolution patterns to handle generic form.
• Vectorization: There are named ops in the file like linalg::DepthwiseConv1DNwcWcOp. We need to switch the isa to a matcher. I did not check all the code, but the core vectorizer is using linalg::LinalgOp which is good.

So I think there are two major tasks in the work stream:

• Implement the matchers and use it in IREE and Vectorization.
• Adapt the convolution patterns, that are mentioned in the issue, to accept generalized convolution ops.

[Abhishek-Varma]

I tried exposing a single API, so now with just one function name exposed we should be able to achieve the goal of this PR :-

template <typename ConvOpTy>
bool isaConvolutionOpOfType(LinalgOp op,
                            SmallVector<int64_t> *dilations = nullptr,
                            SmallVector<int64_t> *strides = nullptr);

Let me know if this looks okay. Have updated the description of the PR slightly.

CC: @rengolin @hanhanW

[rengolin]

Thanks @hanhanW for the context around IREE. This is one of the key design decisions behind the Linalg Forms.

Adding my side to it: While the background motivation for this particular PR is to improve IREE, the matchers themselves are useful on their own for all other upstream and downstream users to be able to convert between the forms and simplify their matching infrastructure.

I'll let @javedabsar1 review the code and approve, but I like your public API using just one template argument and two function arguments. While the (static) implementation is still very verbose, the outward facing design is very clean. We can improve the private implementation with time, independent of being able to use this already upstream. Thanks!

[banach-space]

Thanks for working of this, clearly a lot of effort has gone into this!

Before we proceed, I suggest splitting this into multiple PRs. This is nearly 3k LOC and just impossible to review as a single PR. From what I can tell, splitting this into multiple PRs should be trivial:

• PR1 - high-level structure (the required minimum) + 1 representative example (e.g. 1D+2D+3D depthwise conv, NWC format)
• PR2 - Depthwise convs - NCW format
• PR3 - Regular convs
• ...

Or something similar.

Two other high-level point re testing:

• Lets follow a trend initiated by @javedabsar for element-wise Ops and move Convs-related tests into a dedicated directly (see https://github.com/llvm/llvm-project/tree/main/mlir/test/Dialect/Linalg/elementwise for Elementwise).
• Lets make sure to use consistent naming for such roundtrip tests, see e.g. https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/Linalg/roundtrip-broadcast.mlir and https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/Linalg/roundtrip-transpose.mlir

How does it sound?

[banach-space]

IMO, a switch statement will be easier to follow. Could you try this:

llvm-project/llvm/include/llvm/ADT/TypeSwitch.h

Line 88 in b228a18

class TypeSwitch : public detail::TypeSwitchBase<TypeSwitch<T, ResultT>, T> {

?

[banach-space]

Surely linalg.pooling_nhwc_max_unsigned should only allow integers, no?

[Abhishek-Varma]

Thank you!

@banach-space @javedabsar1 @rengolin @hanhanW - I've raised the first PR that aims to break the current PR into reviewable chunk.

Here is the PR's link : #163724

I'm closing this PR - we can shift any discussion henceforth to that PR.