[MLIR][Linalg] Fixed obsolete examples in the MLIR Linalg Dialect doc

[kumasento]

Summary

This diff fixes some obsolete examples in the Linalg dialect documentation: https://mlir.llvm.org/docs/Dialects/Linalg/

These examples are used to explain the basic properties of the Linalg dialect, which are not automatically generated from TableGen and are using out-of-date MLIR/Linalg syntax.

This diff extends each example by adding essential attributes and changing its syntax to make it processible by mlir-opt. There is also a command attached to each example that says how the example can be processed.

[ftynse]

Thanks! Please address the comments and this is good to go.

mlir/docs/Dialects/Linalg.md

74

Nit: we usually indent blocks the same as the parent operation, i.e.

linalg.generic... {
^bb0(...):
  "nested_op"() ...
}

96–111

Nit: while mlir-opt does produce an explicit module op, it is not mandatory in the input (an implicit top-level module is assumed) and superfluous to understanding this example. Let's not introduce it.

97

Let's not drop the layout. It's important to mention that the memref _may_ have one, without necessarily caring which one. You can use #layout1 and mention it's an attribute alias defined somewhere above.

[kumasento]

Finished fixing other examples in the properties section.

[kumasento]

mlir/docs/Dialects/Linalg.md

74

Thank you :) I've changed the styling of those mlir snippets.

96–111

Thank you @ftynse ! I've removed all module mentioned.

[kumasento]

mlir/docs/Dialects/Linalg.md

97

Hi @ftynse

Thank you for your comment. But I'm not sure whether it is still necessary to add the layout after I've already specified the indexing_maps in the #attrs. I kind of assume that layout is recommended to be generated from indexing_maps instead of writing by hand, based on some of the unittests in https://github.com/llvm/llvm-project/blob/master/mlir/test/Dialect/Linalg/standard.mlir, but I'm not very certain about this. If you would prefer explicitly specifying the layout, I will update these examples accordingly :)

[ftynse]

mlir/docs/Dialects/Linalg.md

97

Thank you for your comment. But I'm not sure whether it is still necessary to add the layout after I've already specified the indexing_maps in the #attrs. I kind of assume that layout is recommended to be generated from indexing_maps instead of writing by hand, based on some of the unittests in https://github.com/llvm/llvm-project/blob/master/mlir/test/Dialect/Linalg/standard.mlir, but I'm not very certain about this. If you would prefer explicitly specifying the layout, I will update these examples accordingly :)

The memref layout is independent of indexings. The former describes the property of the data (most use cases fit into the strided layout scheme) while the latter is a property of the operation. Matrix multiplication works equally well on contiguous and non-contiguous memrefs with transpositions. The two properties contribute to the address calculation and have different expressivity trade-offs. For instance, indexings for matmul must be a permutation of inputs so they cannot express strided access, but the memref itself may have a strided layout. This is of utmost importance for transformations because strided layout do appear as the result of tiling and we must make clear Linalg supports memrefs with strided layouts.

[kumasento]

Added memory layouts to the first two examples.

[kumasento]

mlir/docs/Dialects/Linalg.md

97

Thank you @ftynse , I've added layout specification to the memref operands in the first two examples, please check them out when you're available.

I'm so sorry that I was very silly not being able to distinguish between the concept of layout and indexing. Now I fully understand that and I definitely agree with you that these examples should elaborate the possibility to specify memref layout for linalg.generic operation.

Please let me know if there is anything else I should improve, thanks!

[kumasento]

Hi Alex @ftynse , if this diff looks ok would you mind helping me committing it? Thanks!

mlir/docs/Dialects/Linalg.md

 #### Property 1: Input and Output Operands Define The Iteration Space<a name="prop1"></a>
 A `linalg.generic` op fully *derives* the specification of its iteration space
 from its operands.
 The property enforces that a localized IR element (the op) *has* all the information
 needed to synthesize the control-flow required to iterate over its operands,
 according to their type. This notion of IR localization bears some resemblance
 to [URUK](http://icps.u-strasbg.fr/~bastoul/research/papers/GVBCPST06-IJPP.pdf).
 
-Consider the following, partially specified, `linalg.generic` example:
-```
-#attrs = {args_in: 1, args_out: 1}
-func @example(%A: memref<?xf32, layout1>,
-              %B: memref<?xvector<4xf32, layout2>>) {
-  linalg.generic #attrs (%2, %3): memref<?xf32, layout1>,
-                                  memref<?xvector<4xf32, layout2>>
+Consider the following fully specified `linalg.generic` example.
+Here, the first operand is a `memref` of `f32` scalar elements that
+has an ordinary identity layout, and the second one is a `memref` of
+4-element vectors with a 2-strided, 1-offset layout.
+
+```mlir
+// File name: example1.mlir
+#accesses = [
+  affine_map<(m) -> (m)>,
+  affine_map<(m) -> (m)>
+]
+#attrs = {
+  args_in = 1,
+  args_out = 1,
+  indexing_maps = #accesses,
+  iterator_types = ["parallel"]
+}
+// memory layouts
+#identity = affine_map<(d0) -> (d0)>
+
+func @example(%A: memref<?xf32, #identity>,
+              %B: memref<?xvector<4xf32>, offset: 1, strides: [2]>) {
+  linalg.generic #attrs %A, %B {
+  ^bb0(%a: f32, %b: vector<4xf32>):
+    %c = "some_compute"(%a, %b): (f32, vector<4xf32>) -> (vector<4xf32>)
+    linalg.yield %c: vector<4xf32>
+  } : memref<?xf32, #identity>, memref<?xvector<4xf32>, offset: 1, strides: [2]>
   return
 }
 ```
 
 The property "*Input and Output Operands Define The Iteration Space*" is
 materialized by a lowering into a form that will resemble:
-```
-func @example(%A: memref<?xf32, layout1>,
-              %B: memref<?xvector<4xf32, layout2>>) {
-  %M = "dim" %A, 0: index
-  %N = "dim" %B, 0: index
-  %eq = eq %M, %N: i1   // iteration space is consistent with data
-  assert(%eq): (i1) -> ()
-  for %i = 0 to %M {
-    %a = load %A[%i]: memref<?xf32, layout1>
-    %b = load %B[%i]: memref<?xvector<4xf32>, layout2>
-    // compute arg types match elemental tensor types
-    %c = "some_compute"(%a, %b): (f32, vector<4xf32>) -> (vector<4xf32>)
-    store %c, %B[%i]: memref<?xvector<4xf32>, layout2>
+
+```mlir
+// Run: mlir-opt example1.mlir -allow-unregistered-dialect -convert-linalg-to-loops
+// This converted representation is in the `scf` dialect.
+// It's syntax can be found here: https://mlir.llvm.org/docs/Dialects/SCFDialect/
+#map0 = affine_map<(d0) -> (d0 * 2 + 1)>
+
+func @example(%arg0: memref<?xf32>, %arg1: memref<?xvector<4xf32>, #map0>) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %0 = dim %arg0, %c0 : memref<?xf32>
+  scf.for %arg2 = %c0 to %0 step %c1 {
+    %1 = load %arg0[%arg2] : memref<?xf32>
+    %2 = load %arg1[%arg2] : memref<?xvector<4xf32>, #map0>
+    %3 = "some_compute"(%1, %2) : (f32, vector<4xf32>) -> vector<4xf32>
+    store %3, %arg1[%arg2] : memref<?xvector<4xf32>, #map0>
   }
   return
 }
 ```
 
 The property participates in simplifying analyses and transformations. For
 instance, it guarantees no out-of bounds access can occur by construction
 (assuming dynamic operand dimensions agree with each other, which is the

 and [position-dependent
 arrays](https://www.lift-project.org/publications/2019/pizzuti19positiondependentarrays.pdf),
 as well as [TACO](http://tensor-compiler.org/), has shown.
 
 #### Property 2: Reversible Mappings Between Control and Data Structures<a name="prop2"></a>
 A `linalg.generic` *defines* the mapping between the iteration space (i.e. the
 loops) and the data.
 
-Consider the following, partially specified, `linalg.generic` example:
-```
-#indexing_maps = {
-  (i, j) -> (j, i),
-  (i, j) -> (j)
-}
-#attrs = {args_in: 1, args_out: 1, indexings: indexing_maps}
-func @example(%A: memref<8x?xf32, layout1>,
-              %B: memref<?xvector<4xf32, layout2>>) {
-  linalg.generic #attrs (%A, %B): memref<8x?xf32, layout1>,
-                                  memref<?xvector<4xf32, layout2>>
+Consider the following fully specified `linalg.generic` example.
+Here, the first `memref` is a 2-strided one on both of its dimensions,
+and the second `memref` uses an identity layout.
+
+```
+// File name: example2.mlir
+#indexing_maps = [
+  affine_map<(i, j) -> (j, i)>,
+  affine_map<(i, j) -> (j)>
+]
+#attrs = {
+  args_in = 1,
+  args_out = 1,
+  indexing_maps = #indexing_maps,
+  iterator_types = ["parallel", "parallel"]
+}
+
+func @example(%A: memref<8x?xf32, offset: 0, strides: [2, 2]>,
+              %B: memref<?xvector<4xf32>>) {
+  linalg.generic #attrs %A, %B {
+  ^bb0(%a: f32, %b: vector<4xf32>):
+    %c = "some_compute"(%a, %b): (f32, vector<4xf32>) -> (vector<4xf32>)
+    linalg.yield %c: vector<4xf32>
+  }: memref<8x?xf32 , offset: 0, strides: [2, 2]>, memref<?xvector<4xf32>>
   return
 }
 ```
 
 The property "*Reversible Mappings Between Control and Data Structures*" is
 materialized by a lowering into a form that will resemble:
 ```
-#attrs = {args_in: 1, args_out: 1, indexings: indexing_maps}
-func @example(%A: memref<8x?xf32, layout1>,
-              %B: memref<?xvector<4xf32, layout2>>) {
-  // loop bounds determined from data sizes by “inverting the map”
-  %J = "dim" %A, 0: index
-  %I = "dim" %A, 1: index
-  %J2 = "dim" %B, 0: index
-  // iteration space is consistent with data + mapping inference
-  %eq = "eq" %J, %J2: i1
-  "assert" %eq: (i1) -> ()
-  for %i = 0 to %I {           // loop order is fully defined by indexing maps
-    for %j = 0 to %J {         // arbitrary permutations are possible
-      %a = "load" %A, %j, %i: memref<8x?xf32>
-      %b = "load" %B, %j: memref<?xvector<4xf32>>
-      %c = "some_compute"(%a, %b): (f32, vector<4xf32>) -> (vector<4xf32>)
-      "store" %c, %B, %j: memref<?xvector<4xf32>>
+// Run: mlir-opt example2.mlir -allow-unregistered-dialect -convert-linalg-to-loops
+#map0 = affine_map<(d0, d1) -> (d0 * 2 + d1 * 2)>
+
+func @example(%arg0: memref<8x?xf32, #map0>, %arg1: memref<?xvector<4xf32>>) {
+  %c8 = constant 8 : index
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %0 = dim %arg0, %c1 : memref<8x?xf32, #map0>
+  scf.for %arg2 = %c0 to %0 step %c1 {
+    scf.for %arg3 = %c0 to %c8 step %c1 {
+      %1 = load %arg0[%arg3, %arg2] : memref<8x?xf32, #map0>
+      %2 = load %arg1[%arg3] : memref<?xvector<4xf32>>
+      %3 = "some_compute"(%1, %2) : (f32, vector<4xf32>) -> vector<4xf32>
+      store %3, %arg1[%arg3] : memref<?xvector<4xf32>>
     }
   }
   return
 }
 ```
 
 This mapping needs to be reversible because we want to be
 able to go back and forth between the two and answer questions such as:
 - Given a subset of the iteration space, what subset of data does it read and
 write?
 - Given a subset of data read or written, what subset of the iteration space
 is responsible for this read or write?
 
 Answering these `2` questions is one of the main analyses that Linalg uses to
 implement transformations such as tiling, tiled producer-consumer fusion, and
 promotion to temporary buffers in fast memory.
 
-In the current implementation, `linalg.generic` uses a list of [AffineMaps]().
+In the current implementation, `linalg.generic` uses a list of [AffineMaps](https://mlir.llvm.org/docs/LangRef/#affinemap-attribute) (see the `#indexing_maps` attribute in the previous examples).
 This is a pragmatic short-term solution, but in the longer term note that
 this property could be even evaluated dynamically, similarly to
 inspector-executor algorithms.
 
 #### Property 3: The Type Of Iterators is Defined Explicitly<a name="prop3"></a>
 A `linalg.generic` op fully *declares* the type of its iterators. This
 information is used in transformations.
 

 At this time there are no additional restrictions to the region
 semantics. This is meant to allow the exploration of various design tradeoffs
 at the intersection of regions and iterator types.
 In particular, the frontend is responsible for the semantics of iterator types
 to correspond to the operations inside the region: the region can capture
 buffers arbitrarily and write into them. If this conflicts with some parallel
 iterator requirement, this is undefined behavior.
 
-Concretely, consider the following, partially specified, `linalg.generic`
-example:
+Previous examples already ellaborate compute payloads with an unregistered function `"some_compute"`. The following code snippet shows what the result will be when using a concrete operation `addf`:
 ```
-#indexing_maps = {
-  (i, j) -> (i, j),
-  (i, j) -> (i, j)
+// File name: example3.mlir
+#indexing_maps = [
+  affine_map<(i, j) -> (i, j)>,
+  affine_map<(i, j) -> (i, j)>,
+  affine_map<(i, j) -> (i, j)>
+]
+#attrs = {
+  args_in = 2,
+  args_out = 1,
+  indexing_maps = #indexing_maps,
+  iterator_types = ["parallel", "parallel"]
 }
-#attrs = {args_in: 2, args_out: 1, indexings: #indexing_maps}
 func @example(%A: memref<?x?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {
-  linalg.generic #attrs (%A, %B, %C) {
-    ^bb0(%a: f32, %b: f32):
-      %c = addf %a, %b : f32
-      return %c : f32
+  linalg.generic #attrs %A, %B, %C {
+  ^bb0(%a: f32, %b: f32, %c: f32):
+    %d = addf %a, %b : f32
+    linalg.yield %d : f32
   }: memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>
   return
 }
 ```
 
+This function basically element-wise adds up two matrices (`%A` and `%B`) and stores the result into another one (`%C`).
+
 The property "*The Compute Payload is Specified With a Region*" is
 materialized by a lowering into a form that will resemble:
 ```
-func @example(%A: memref<?x?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {
-  %M = dim %A, 0: index
-  %N = dim %B, 1: index
-  for %i = 0 to %M {
-    for %j = 0 to %N {
-      %a = load %A[%i, %j]: memref<?x?xf32>
-      %b = load %B[%i, %j]: memref<?x?xf32>>
-      %c = addf %a, %b : f32
-      store %c, %C[%i, %j]: memref<?x?xf32>
-    }
-  }
+// Run: mlir-opt example3.mlir -convert-linalg-to-loops
+#indexing_maps = [
+  affine_map<(i, j) -> (i, j)>,
+  affine_map<(i, j) -> (i, j)>,
+  affine_map<(i, j) -> (i, j)>
+]
+#attrs = {
+  args_in = 2,
+  args_out = 1,
+  indexing_maps = #indexing_maps,
+  iterator_types = ["parallel", "parallel"]
+func @example(%A: memref<?x?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {
+  linalg.generic #attrs %A, %B, %C {
+  ^bb0(%a: f32, %b: f32, %c: f32):
+    %d = addf %a, %b : f32
+    linalg.yield %d : f32
+  }: memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>
   return
 }
 ```
 
 In the process of lowering to loops and lower-level constructs, similar
 requirements are encountered, as are discussed in the [inlined call op
 proposal](https://llvm.discourse.group/t/introduce-std-inlined-call-op-proposal/282/2).
 We expect to be able to reuse the common lower-level infrastructure provided
 it evolves to support both region arguments and captures.
 
 #### Property 5: May Map To an External Library Call<a name="prop5"></a>
 A `linalg.generic` op may map to an external library call by specifying a
 `SymbolAttr`. At this level of abstraction, the important glue is the ability
 to perform transformations that preserve the structure necessary to ***call
 the external library after different transformations have been applied***.
 
 This involves considerations related to preservation of op semantics
 and integration at the ABI level. Regardless of whether one wants to use
 external library calls or a custom ISA, the problem for codegen is similar:
 preservation of a fixed granularity.
 
-Consider the following, partially specified, `linalg.generic`
-example:
+Consider the following example that adds an additional attribute `library_call="pointwise_add"`
+that specifies the name of an external library call we intend to use:
 ```
-#fun_attr = "pointwise_add"
-#indexing_maps = {
-  (i, j) -> (i, j),
-  (i, j) -> (i, j)
+// File name: example4.mlir
+#indexing_maps = [
+  affine_map<(i, j) -> (i, j)>,
+  affine_map<(i, j) -> (i, j)>,
+  affine_map<(i, j) -> (i, j)>
+]
+#attrs = {
+  args_in = 2,
+  args_out = 1,
+  indexing_maps = #indexing_maps,
+  iterator_types = ["parallel", "parallel"],
+  library_call = "pointwise_add"
 }
-#attrs = {args_in: 2, args_out: 1, indexings: #indexing_maps, fun: #fun_attr}
 func @example(%A: memref<?x?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {
-  linalg.generic #attrs (%A, %B, %C) {
-    ^bb0(%a: f32, %b: f32):
-      %c = addf %a, %b : f32
-      return %c : f32
+  linalg.generic #attrs %A, %B, %C {
+  ^bb0(%a: f32, %b: f32, %c: f32):
+    %d = addf %a, %b : f32
+    linalg.yield %d : f32
   }: memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>
   return
 }
 ```
 
 The property "*Map To an External Library Call*" is
 materialized by a lowering into a form that will resemble:
 
 ```
-func @pointwise_add_sxsxf32_sxsxf32(memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>) -> ()
-
-func @example(%A: memref<?x?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {
-  call @pointwise_add_sxsxf32_sxsxf32 (%A, %B, %C):
-    (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>) -> ()
+// Run: mlir-opt example4.mlir -convert-linalg-to-std
+// Note that we lower the Linalg dialect directly to the Standard dialect.
+// See this doc: https://mlir.llvm.org/docs/Dialects/Standard/
+
+#map0 = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+
+func @example(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
+  %0 = memref_cast %arg0 : memref<?x?xf32> to memref<?x?xf32, #map0>
+  %1 = memref_cast %arg1 : memref<?x?xf32> to memref<?x?xf32, #map0>
+  %2 = memref_cast %arg2 : memref<?x?xf32> to memref<?x?xf32, #map0>
+  call @pointwise_add(%0, %1, %2) : (memref<?x?xf32, #map0>, memref<?x?xf32, #map0>, memref<?x?xf32, #map0>) -> ()
   return
 }
+func @pointwise_add(memref<?x?xf32, #map0>, memref<?x?xf32, #map0>, memref<?x?xf32, #map0>) attributes {llvm.emit_c_interface}
 ```
 
 Which, after lowering to LLVM resembles:
 ```
-func @pointwise_add_sxsxf32_sxsxf32(!llvm<"{ float*, i64, [2 x i64], [3 x i64] }*">,
-                                    !llvm<"{ float*, i64, [2 x i64], [3 x i64] }*">,
-                                    !llvm<"{ float*, i64, [2 x i64], [3 x i64] }*">) -> ()
-
-func @example(%A: !llvm<"{ float*, i64, [2 x i64], [3 x i64] }*">,
-              %B: !llvm<"{ float*, i64, [2 x i64], [3 x i64] }*">,
-              %C: !llvm<"{ float*, i64, [2 x i64], [3 x i64] }*">) {
-  llvm.call @pointwise_add_sxsxf32_sxsxf32 (%A, %B, %C):
-    (!llvm<"{ float*, i64, [2 x i64], [3 x i64] }*">...) -> ()
+// Run: mlir-opt example4.mlir -convert-linalg-to-std | mlir-opt -convert-std-to-llvm
+// Some generated code are omitted here.
+func @example(%arg0: !llvm<"float*">, ...) {
+  ...
+  llvm.call @pointwise_add(...) : (!llvm<"float*">, ...) -> ()
   return
 }
+
+llvm.func @pointwise_add(%arg0: !llvm<"float*">, ...) attributes {llvm.emit_c_interface} {
+  ...
+  llvm.call @_mlir_ciface_pointwise_add(%9, %19, %29) : (!llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }*">, !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }*">, !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }
+*">) -> ()
+  llvm.return
+}
+llvm.func @_mlir_ciface_pointwise_add(!llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }*">, !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }*">, !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }*">) attributes {llvm.emit_c_interface}
 ```
 
 ##### Convention For External Library Interoperability
 The `linalg` dialect adopts a convention that is similar to `BLAS` when
 offloading operations to fast library implementations: pass a non-owning
 pointer to input and output data with additional metadata. This convention
 is also found in libraries such as `MKL`, `OpenBLAS`, `BLIS`, `cuBLAS`,
 `cuDNN`, etc.. and more generally at interface points across language