[Mlir-commits] [mlir] [MLIR][Linalg][Docs] Add forms to Linalg rationale docs (PR #156859)

[Andrzej Warzyński]

================
@@ -506,6 +506,65 @@ potential by introducing lower-level IR ops and *smaller* Linalg ops.
 This gradually reduces the potential, all the way to Loops + VectorOps
 and LLVMIR.
 
+### Interchangeability of Forms<a name="forms"></a>
+
+Linalg's various forms (named, generic) also carry information, and that
+information should be preserved as much as possible during the progressive
+lowering. A `matmul` operation is a special case of a `contract` operation,
+which in turn is a special case of `generic` operation. Transformations on
+the more special forms should not be converted to the more generic ones
+unnecessarily, in the same way that they should not be broken down into
+loops + arithmetic if they can still be represented as a Linalg op.
+
+#### Generic, Category, Named<a name="generic_category_named"></a>
+
+The core Linalg operation tree has three forms:
+* **Generic:** Represented by `linalg.generic` and can encode all perfectly-nested
+loop operations.
+* **Category:** Represented by `linalg.contract` and `linalg.elementwise`,
+which are special (einsum) forms of the `generic` operation.
+* **Named:** All _named_ forms that can lower to either _category_ or
+_generic_ forms. For example, `linalg.matmul`, `linalg.add`, etc.
+
+Unlike lowering to loops, the different Linalg forms that are derived from
+`linalg.generic` are *equivalent*. It should always be possible to convert
+a named operation into a generic and back to named, if the semantics is
+preserved. The various forms in the Linalg dialect are meant to facilitate
+pattern matching (single operations or DAGs) and to be able to consider
+different forms as *canonical* for different transforms.
+
+#### Special Operations<a name="special_ops"></a>
+
+Not all Linalg operations represent perfectly nested loops, and therefore
+cannot be represented as a `linalg.generic`. There are two kinds of Linalg
+operations that fall into this category:
+* **Composite:** Operations that compose multiple Linalg operations, for
+example `linalg.softmax`. These can be converted to a DAG of Linalg operations
+(in any form).
+* **Special Named:** Operations that are usually matched against library calls
+or special lowering, but can only be lowered to a combination of Linalg and
+non-Linalg operations, for example `linalg.*conv*`, `linalg.winograd*`,
+`linalg.pooling*`, etc.
----------------
banach-space wrote:

I think that this needs a bit of refining. 

**Conc <--> Generic**

Note the comment above:

> Not all Linalg operations represent perfectly nested loops, and therefore
cannot be represented as a `linalg.generic`.

Yet:
```bash
$ cat conv.mlir
func.func @depthwise_conv1d_nwc_wc_1x8x3xi8_tensor(%input: tensor<1x8x?xi8>,
                                                   %filter: tensor<1x?xi8>,
                                                   %output: tensor<1x8x?xi8>) -> (tensor<1x8x?xi8>) {
  %res = linalg.depthwise_conv_1d_nwc_wc
    {dilations = dense<1> : vector<1xi64>,
    strides = dense<1> : vector<1xi64>}
    ins(%input, %filter : tensor<1x8x?xi8>, tensor<1x?xi8>)
    outs(%output : tensor<1x8x?xi8>) -> tensor<1x8x?xi8>
  return %res : tensor<1x8x?xi8>
}

$ bin/mlir-opt --linalg-generalize-named-ops conv.mlir
#map = affine_map<(d0, d1, d2, d3) -> (d0, d1 + d3, d2)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
module {
  func.func @depthwise_conv1d_nwc_wc_1x8x3xi8_tensor(%arg0: tensor<1x8x?xi8>, %arg1: tensor<1x?xi8>, %arg2: tensor<1x8x?xi8>) -> tensor<1x8x?xi8> {
    %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<1x8x?xi8>, tensor<1x?xi8>) outs(%arg2 : tensor<1x8x?xi8>) {
    ^bb0(%in: i8, %in_0: i8, %out: i8):
      %1 = arith.muli %in, %in_0 : i8
      %2 = arith.addi %out, %1 : i8
      linalg.yield %2 : i8
    } -> tensor<1x8x?xi8>
    return %0 : tensor<1x8x?xi8>
  }
```

So, convolutions can and are "representable" as `linalg.generic`. IIRC from our conversation on Discord:
* `linalg.conv` --> `linalg.generic` is indeed fine,
* `linalg.generic` --> `linalg.conv` should be possible in practice but will be tricky to implement (we just need a brave volunteer).

**Library calls**
As fot this point:
> Operations that are usually matched against library calls
or special lowering, but can only be lowered to a combination of Linalg and
non-Linalg operations,

_We_ don't use library calls for convs - we lower them and such lowerings don't require going via other `Linalg` Ops - we just rely on the Linalg vectorizer. My point being - while some people may choose to lower these to lib calls, we can't really say that that's the common/preferred practice. I can't really comment on the other Ops as don't use them personally.

**Specific Suggestions**
* From what I can tell, Convs and friends are `TODO` rather than "forms do not apply to these Ops".
* Lowering Convs to lib calls is a viable option, but just an option. To me, that's tangential to "forms".

WDYT?

https://github.com/llvm/llvm-project/pull/156859