Linalg convolutions - who's converting ConvInputNHWCFilterHWCFOp to ConvOp?

I’m looking into the code that transforms Linalg (and in particular convolutions) to affine loops. It’s pretty nice, but there’s a catch:

The C++ code (inside llvm-project/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp) only works on generic ConvOp operations, whereas my code contains only specialized ConvInputNHWCFilterHWCFOp operations.

These two classes are not related by inheritance, so someone must perform these conversions somewhere else, and I can’t find where. It would be interesting to know on which other operations such generalizations are performed, how to enable them, and how to disable them.

I think all the TC named ops are handled in a different way, ie, they are not hard-coded. Since we have all the indexing maps and region builder, it is fair to have a unified way to convert TC named ops to loops. I could not find the implementation now, but I tested with one of TC ops and it’s working. E.g.,

func @depthwise_conv_2d_input_nhwc_filter_hwcf_memref(%input: memref<2x4x5x2xf32>, %filter: memref<2x2x2x3xf32>, %output: memref<2x3x4x2x3xf32>) {
  linalg.depthwise_conv_2d_input_nhwc_filter_hwcf
     { strides = dense<1> : tensor<2xi64> }
     ins(%input, %filter : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>)
    outs(%output : memref<2x3x4x2x3xf32>)
  return
}

It will be lowered to loops if you run mlir-opt -convert-linalg-to-loops conv.mlir.

#map = affine_map<(d0, d1) -> (d0 + d1)>
module  {
  func @depthwise_conv_2d_input_nhwc_filter_hwcf_memref(%arg0: memref<2x4x5x2xf32>, %arg1: memref<2x2x2x3xf32>, %arg2: memref<2x3x4x2x3xf32>) {
    %c2 = constant 2 : index
    %c3 = constant 3 : index
    %c4 = constant 4 : index
    %c0 = constant 0 : index
    %c1 = constant 1 : index
    scf.for %arg3 = %c0 to %c2 step %c1 {
      scf.for %arg4 = %c0 to %c3 step %c1 {
        scf.for %arg5 = %c0 to %c4 step %c1 {
          scf.for %arg6 = %c0 to %c2 step %c1 {
            scf.for %arg7 = %c0 to %c3 step %c1 {
              scf.for %arg8 = %c0 to %c2 step %c1 {
                scf.for %arg9 = %c0 to %c2 step %c1 {
                  %0 = affine.apply #map(%arg4, %arg8)
                  %1 = affine.apply #map(%arg5, %arg9)
                  %2 = memref.load %arg0[%arg3, %0, %1, %arg6] : memref<2x4x5x2xf32>
                  %3 = memref.load %arg1[%arg8, %arg9, %arg6, %arg7] : memref<2x2x2x3xf32>
                  %4 = memref.load %arg2[%arg3, %arg4, %arg5, %arg6, %arg7] : memref<2x3x4x2x3xf32>
                  %5 = mulf %2, %3 : f32
                  %6 = addf %4, %5 : f32
                  memref.store %6, %arg2[%arg3, %arg4, %arg5, %arg6, %arg7] : memref<2x3x4x2x3xf32>
                }
              }
            }
          }
        }
      }
    }
    return
  }
}

The definition of TC ops can be found at llvm-project/LinalgNamedStructuredOpsSpec.tc at main · llvm/llvm-project · GitHub

note: they will be migrated to a new mechanism: llvm-project/LinalgNamedStructuredOps.yaml at main · llvm/llvm-project · GitHub

I know the operations, and I know mlir-opt can convert them. But when I’m looking into the pass definition, I cannot find the transformation itself. Only the one for ConvOp can be easily found. This is why I’m asking. My assumption was that some hidden, implicit translate layer converts all ConvXXX operations into ConvOp before lowering.

I found the implementation. A LinalgOp is a structured Linalg op. All the TC ops can be casted to a LinalgOp.

Thus the actual implementation for LinalgOp is

side note: I don’t think you should use ConvOp and Pooling*Op. They will be deprecated at some point, and they only work for memref types. I think we will have a ConvOpInterface or something similar that you could use to anchor transformations. In IREE, we already migrated all the conv ops and pooling ops to TC ops. This also enables tile and fuse and distribute in Linalg on tensors.

Can you explain how I can cast a ConvInputNHWCFilterHWCFOp into a
ConvOp? This is what I was looking for in my original post. The problem is that the two classes are not related by inheritance, so there must be a conversion op (e.g. a constructor) somewhere. And I was not able to find it.

No, you can not cast a ConvInputNHWCFilterHWCFOp into a ConvOp. AFAIK, there is no pattern for it.

So then how does mlir-opt --convert-linalg-to-affine-loops manage to convert the linalg ConvInputNHWCFilterHWCFOp operation? Where is the conversion pattern implemented? I cound not find it in the code.

I think I explained it above.ConvInputNHWCFilterHWCFOp is casted into a LinalgOp and there is a method to lower it to loops.