# Named ops refresh

The primary aim of this document is to collaborate on getting to a better state with respect to Linalg named operations. As of writing of this document we have 28 different Convolution operations, 17 different Pooling operations, 14 different Matmul-like operations for a total of 93 different Linalg operations. The list keeps growing. The reason for this increase is that the OpDSL based Linalg named ops is positioned to be a common target for "front-end" dialects like Torch-MLIR, TOSA, StableHLO, etc.

The primary reason for explosion in the named ops (refered to as OpDSL-specified named ops) is that these ops are operations that implement the [LinalgStructuredInterface](https://github.com/llvm/llvm-project/blob/de2fad32513f7420988df1cf99aff90e0a067469/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td#L226) and have
1. Fixed iterator types
1. Fixed Indexing maps
1. Fixed region.

Consequently, for each of the variation of the input operand "layout" a new operation is added. A further limitation is that the named operations cannot easily fold in broadcasting behaviour. For example consider

%lhs = linalg.broadcast ins(%input : tensor<?x?xf32>)
  outs(%empty : tensor<?x?x?xf32>) dimensions = [0]
%bmm = linalg.batch_matmul ins(%lhs, %rhs : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
  outs(%init : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>

while this could be represented more succinctly as

%lhs = linalg.generic {
  indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d1, d3)>,
                   affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>,
                   affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>],
  iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
  ins(%input, %rhs : tensor<?x?xf32>, tensor<?x?x?xf32>)
  outs(%init : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
        ...
}

As a result front-end dialects when targeting `batch_matmul` end up introducing broadcasts, which pessimizes the performance unless the rest of the compilation stack folds it away. The only way currently this is feasible is to generalize named ops and fuse the broadcast operation with the generalized `linalg.batch_matmul` to end up with the `linalg.generic` shown above. This isnt ideal either cause now the semantic information about the region computation carried by the name of the operation is also lost.

A further limitation of the current implementation of named ops is that the named ops carry a hidden region. This region represents the computation within the inner most loop when the op is lowered to loops. The presence of this hidden region makes it impossible to use PDL/PDLL to match named ops (unless those can be extended to handle operations with regions). This region is not serving any purpose apart from making the implementation of the generalization to `linalg.generic` operation easier.

## Solution ##

One way out of these issues is to have a named op defined in TableGen where unlike the OpDSL specified named ops these operations have only a fixed region, i.e. the operation contain indexing maps for each operand, and the iterator types attribute list explicitly. The "name" of the operation therefore just represents the region. The region is not explicitly defined, but the operation carries enough information required to generate the region if needed (say during generalization).

So the above `broadcast` + `batch_matmul` could be represented by a new TableGen defined named operation

%contraction = linalg.contraction {
  indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d1, d3)>,
                   affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>,
                   affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>],
  iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
  ins(%input, %rhs : tensor<?x?xf32>, tensor<?x?x?xf32>)
  outs(%init : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>

One piece of information that these operations will need to carry that isnt there in current OpDSL based named ops or `LinalgStructuredOps` is a list of attributes that can cast from input operand element type to output operand element type. A list of enums can be used to codify the cast semantics. For example, there are two forms of casts supported today that could be represented by two enums.

enum LinalgCastEnum : int32_t {
  CAST_SIGNED = 0,
  CAST_UNSIGNED = 1,
};

So a contraction with mixed types would be

%contraction = linalg.contraction {
  indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d1, d3)>,
                   affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>,
                   affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>],
  iterator_types = ["parallel", "parallel", "parallel", "reduction"],
  casting_fn = ["CAST_UNSIGNED", "CAST_UNSIGNED"]}
  ins(%input, %rhs : tensor<?x?xi8>, tensor<?x?x?xi8>)
  outs(%init : tensor<?x?x?xi32>) -> tensor<?x?x?xi32>

All the 14 Matrix-multiply-like operations can be lowered to a single `linalg.contraction`

## Convolutions ##

The explosion of operations is particularly egriguous for convolutions ops. These are the current convolution operations defined using OpDSL

Conv1DNcwFcwOp
Conv1DNwcWcfOp
Conv1DOp
Conv2DNchwFchwOp
Conv2DNchwFchwQOp
Conv2DNgchwFgchwOp
Conv2DNgchwGfchwOp
Conv2DNgchwGfchwQOp
Conv2DNhwcFhwcOp
Conv2DNhwcFhwcQOp
Conv2DNhwcHwcfOp
Conv2DNhwcHwcfQOp
Conv2DOp
Conv3DNcdhwFcdhwOp
Conv3DNdhwcDhwcfOp
Conv3DNdhwcDhwcfQOp
Conv3DOp
DepthwiseConv1DNcwCwOp
DepthwiseConv1DNwcWcOp
DepthwiseConv1DNwcWcmOp
DepthwiseConv2DNchwChwOp
DepthwiseConv2DNhwcHwcOp
DepthwiseConv2DNhwcHwcQOp
DepthwiseConv2DNhwcHwcmOp
DepthwiseConv2DNhwcHwcmQOp
DepthwiseConv3DNcdhwCdhwOp
DepthwiseConv3DNdhwcDhwcOp
DepthwiseConv3DNdhwcDhwcmOp

All of these operations can be represented using a single convolution operation

%convolution = linalg.convolution {
  indexing_maps = [...],
  iterator_types = [...],
  casting_fn = [...]}
  ins(%image, %filter : ...) outs(%init : ...) -> ...

An analysis of the indexing maps is enough to distinguish between the different dimensions.

1. N: Batch dimension. Infered as the dimension(s) that is
- used as an `AffineDimExpr` in the indexing map for `%image`
- used as an `AffineDimExpr` in the indexing map for `%init`
- not used in the indexing map for `%filter`
1. D, H, W : Output image dimensions. Infered as the dimension(s)
- used in the expression of form `a * b + c * d` in the indexing map for `%image`
- not used in the indexing map for the `%filter`
- used as an `AffineDimExpr` in the indexing map for `%init`
1. C: Input channel dimension. Infered as the dimension(s)
- used as an `AffineDimExpr` in the indexing map for `%image`
- used as an `AffineDimExpr` in the indexing map for `%filter`
- not used in the indexing map for the `%init`
1. F: Output channel dimension. Infered as the dimension(s)
- not used in the indexing map for the `%image`
- used as an `AffineDimExpr` in the indexing map for `%image`
- used as an `AffineDimExpr` in the indexing map for `%init`
1. G: Group dimension. Infered as the dimension(s)
- used as an `AffineDimExpr` in the indexing map for `%image`
- used as an `AffineDimExpr` in the indexing map for `%filter`
- used as an `AffineDimExpr` in the indexing map for `%init`

A depthwise convolution is a grouped convolution with missing `C` based on the taxonomy above (and the `G` dimension is refered to as the channel dimension of input, filter and output).

The approach described here folds all the convolutions ops into a single operation. A similar approach can be used to fold all the pooling ops defined using OpDSL into a single operation. So the 93 OpDSL based named ops can be represented using a handful (less than 10) TableGen-based named ops.

## OpDSL Named ops vs (new) TableGen defined named ops.

The OpDSL named ops and TableGen-based named ops defined here are intended for different purposes. The former are more "upward" facing, i.e. interfacing with front-end dialects. The TableGen-based named ops proposed here are a more restricted set of operations that are meant for transformations within the compiler. They allow for optimizations like folding transposes/broadcasts into the operation without having to generalize the operations. As such they serve slightly complementary purpose. It would be worth considering keeping the current OpDSL based named ops as "a layer above" the new TableGen-based named ops. Lowering from the OpDSL based named ops to these TableGen-based named ops is a strict lowering. In essense the TableGen-based named ops are creating new operations that satisfy the [ContractionOpInterface](https://github.com/llvm/llvm-project/blob/82d5dd28b4de7245088f7ed40da37f8cf80461e4/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td#L21) or [ConvolutionOpInterface](https://github.com/llvm/llvm-project/blob/82d5dd28b4de7245088f7ed40da37f8cf80461e4/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td#L136).

### Using PDL/PDLL for matching to libraries ###

A common use of Linalg named ops is to match a DAG of operations to a library call. Today, the community at large is relying on OpDSL named ops to match a particular sequence of operations. But the same taxonomy can be obtained by specifying the indexing maps that represent the different convolution operations (essentially the inverse of the above taxonomy), and such a match could be done using PDL/PDLL which is a more efficient way of doing such matches + offloading.