<commit_before># IREE CPU/GPU Code Generation Pipeline

This document is intended to provide an overview of the codegen pipeline within
IREE used to generate CPU/GPU code. It intends to give an overview of the main
passes used, the objective of the pass, the current implementation, and what it
is expected to achieve in the long term.

Note that while the code generation pipeline supports dynamic shapes, this work
is very preliminary. The description of this is not covered here.

## Input to the codegen pipeline

The input to the code generation pipeline is the module within the
`hal.executable.target` operation. Functions within this module that do __not__
have `Visibility::Private` are the *entry point* functions of the dispatch
region. These are the functions that are *invoked* by the IREE runtime. In
addition, each dispatch region also contains a `hal.interface` operation that
describes the ABI to use for the dispatch region. Two examples of the input to
the code generation pipeline are shown below. In both of these, a single
dispatch function contains a sequence of MHLO operations that the dispatch
region creation has grouped into a single region. Ideally the grouped operations
are fused into a single kernel.

```mlir
hal.executable.target "vulkan*" {
  module attributes {spv.target_env = ...} {
    func @main_ex_dispatch() {
      %c0 = constant 0 : index
      %0 = hal.interface.load.tensor @legacy_io::@arg0,
             offset = %c0 : tensor<4x5xf32>
      %1 = hal.interface.load.tensor @legacy_io::@arg1,
             offset = %c0 : tensor<5x10xf32>
      %2 = "mhlo.dot"(%0, %1) {precision_config = ["DEFAULT", "DEFAULT"]} :
             (tensor<4x5xf32>, tensor<5x10xf32>) -> tensor<4x10xf32>
      hal.interface.store.tensor %2, @legacy_io::@ret0,
        offset = %c0 : tensor<4x10xf32>
      return
    }
    hal.interface @legacy_io attributes {sym_visibility = "private"} {
      hal.interface.binding @arg0, set=0, binding=0,
        type="StorageBuffer", access="Read"
      hal.interface.binding @arg1, set=0, binding=1,
        type="StorageBuffer", access="Read"
      hal.interface.binding @ret0, set=0, binding=2,
        type="StorageBuffer", access="Write|Discard"
    }
  }
}
```
<a name="snippet1"></a>
Snippet 1 : Dispatch region with matrix-matrix multiply operation.

```mlir
hal.executable.target "vulkan*" {
  module attributes {spv.target_env = ...} {
    func @main_ex_dispatch() {
      %c0 = constant 0 : index
      %0 = hal.interface.load.tensor @legacy_io::@arg0,
             offset = %c0 : tensor<10x5xf32>
      %1 = hal.interface.load.tensor @legacy_io::@arg1,
             offset = %c0 : tensor<10x5xf32>
      %2 = hal.interface.load.tensor @legacy_io::@arg2,
             offset = %c0 : tensor<10x5xf32>
      %3 = "mhlo.add"(%0, %1) :
         (tensor<10x5xf32>, tensor<10x5xf32>) -> tensor<10x5xf32>
      %4 = "mhlo.multiply"(%3, %2) :
         (tensor<10x5xf32>, tensor<10x5xf32>) -> tensor<10x5xf32>
      hal.interface.store.tensor %4, @legacy_io::@ret0,
        offset = %c0 : tensor<10x5xf32>
      return
    }
    hal.interface @legacy_io attributes {sym_visibility = "private"} {
      hal.interface.binding @arg0, set=0, binding=0,
        type="StorageBuffer", access="Read"
      hal.interface.binding @arg1, set=0, binding=1,
        type="StorageBuffer", access="Read"
      hal.interface.binding @arg2, set=0, binding=2,
        type="StorageBuffer", access="Read"
      hal.interface.binding @ret0, set=0, binding=3,
        type="StorageBuffer", access="Write|Discard"
    }
  }
}
```
<a name="snippet2"></a>
Snippet 2 : Dispatch region with element-wise operations.

__Roadmap Note__: The current implementation might not actually fuse the
operations grouped into a dispatch region into a single kernel. It is possible
to end up with multiple kernels per dispatch region. Over time we plan to address
this by using fusion at different levels (see below).

The inputs to the dispatch region are materialized within the entry point
function using the `hal.interface.load.tensor` operation, This operation returns
a `tensor` view of the buffer used to store the inputs. Similarly the result of
the dispatch region are *written* out using the `hal.interface.store.tensor`
operation.

The main constraint that the code generation operates under is that it should
not require additional (temporary) buffers to execute the operations grouped
together within a dispatch region. The rationale behind this constraint is that
buffer allocation/synchronization in IREE happens at the granularity of dispatch
regions, allowing the scheduler to make better decision about where to insert
appropriate synchronizations.

The IR after all the passes used in the lowering from MHLO to SPIR-V for the
above two examples can be found here ([matrix-matrix multiply op][DotAfterAll],
[elementwise ops][PwAfterAll]). Below is a description of the major passes used.

## Conversion from MHLO dialect to Linalg on buffers

The code generation pipeline heavily relies on use of [Structured
Operations][LinalgRationale], specifically the [Linalg Dialect][LinalgDialect].
Both, the Linalg operations on `tensor`s and on `memref`s are central to the
progressive lowering approach followed here. The first part of the code
generation pipeline is to convert the MHLO operations on `tensor`s to Linalg
operation on `memref`s. This part of the pipeline is common to both CPU and GPU
code generation.

The steps involved in this conversion is shown below. Each of the arrows
represents a pass in the pipeline ![MHLO To Linalg on `memref`.
conversion](./hlo_to_linalg.png) The next sections describe each of
these passes in more detail.


### MHLO to Linalg on tensors

The first step is to convert MHLO operations to Linalg on tensors. This is done
using the [HLOToLinalgPass][HLOToLinalgPass] from Tensorflow. An example of the
conversion is shown below, where each of the `mhlo.add` and `mhlo.multiply`
operations are converted to `linalg.generic` operations on tensors.

```mlir
#map0 = affine_map<(d0, d1) -> (d0, d1)>
%3 = linalg.generic
       {args_in = 2 : i64, args_out = 1 : i64,
        indexing_maps = [#map0, #map0, #map0],
        iterator_types = ["parallel", "parallel"]} %0, %1 {
     ^bb0(%arg0: f32, %arg1: f32):  // no predecessors
       %5 = addf %arg0, %arg1 : f32
       linalg.yield %5 : f32
     } : tensor<10x5xf32>, tensor<10x5xf32> -> tensor<10x5xf32>
%4 = linalg.generic
       {args_in = 2 : i64, args_out = 1 : i64,
        indexing_maps = [#map0, #map0, #map0],
        iterator_types = ["parallel", "parallel"]} %3, %2 {
     ^bb0(%arg0: f32, %arg1: f32):  // no predecessors
       %5 = mulf %arg0, %arg1 : f32
       linalg.yield %5 : f32
     }: tensor<10x5xf32>, tensor<10x5xf32> -> tensor<10x5xf32>
```
<a name="snippet3"></a>
Snippet 3 : MHLO to Linalg conversion for [element-wise operations](#snippet2)

At the time of writing the representation of Linalg on `tensor`s does not model
reduction iterator types completely. Specifically, the reduction in Linalg is
modeled using read-modify-write approach, i.e. each iteration of the reduction
loop reads the value stored in the output, adds its contribution, and writes
back to the same location. This means the output has to be *initialized* to the
null element of the reduction operator (i.e. 0 if the reduction is done using
addition). This works for operations on buffers. Since tensors are SSA values
they cannot be updated in-place. As a result, the reduction semantics does not
map as well to `tensor`s. For now it is treated as a convention that when the
Linalg operation is converted to use `memref`s it has to be initialized
appropriately before performing the reduction. Due to this, the conversion from
MHLO op to Linalg op is only done for operations which do not need a *reduction*
iterator type in the converted Linalg op. Consequently, only element-wise
operations, broadcast operations and data movement operations (like copy and
transpose) are converted to Linalg operations at this stage.

__Roadmap note__: One long term solution for the above is to have operations on
tensors that have *reduction* iterator type to take an additional argument that
contains the initial value of the result tensor. When the operation is converted
to use `memref`s, the buffer for the initial value operand can be reused for the
result. The details involved have not been fully worked out yet.

### Fusion of Linalg on tensor operations

The Linalg on `tensor` operations generated at the previous step are fused using
the [LinalgFusionOfTensorOps][LinalgFusionOfTensorOps] from MLIR. Since
`tensor`s are SSA values, fusion at this stage can be done without using alias
analysis or dependence analysis based on reads and writes.  Instead the use-def
chains for the `tensor` values can be used to implement producer-consumer
fusion. This stage fuses most elementwise operations, broadcast operations and
data movement operations. An example of the fused op is shown below.

```mlir
#map0 = affine_map<(d0, d1) -> (d0, d1)>
%3 = linalg.generic
       {args_in = 3 : i64, args_out = 1 : i64,
        indexing_maps = [#map0, #map0, #map0, #map0],
        iterator_types = ["parallel", "parallel"]} %0, %1, %2 {
     ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):  // no predecessors
       %4 = addf %arg0, %arg1 : f32
       %5 = mulf %4, %arg2 : f32
       linalg.yield %5 : f32
     }: tensor<?x5xf32>, tensor<?x5xf32>, tensor<?x5xf32> -> tensor<?x5xf32>
```
<a name="snippet4"></a>
Snippet 4: Fusion of Linalg operation on tensors for element-wise operations
shown in [Snippet 3](#snippet3)

### Conversion of Linalg on tensors to Linalg on buffers

Post fusion all the operation on `tensor`s are converted to analogous operations
on `memref`s. In general, this requires a buffer allocation pass. In IREE,
buffer allocation happens at the granularity of dispatch region, and as
mentioned [earlier](#input-to-the-codegen-pipeline), the dispatch region is not
expected to use any additional temporary buffers. So instead of having another
buffer allocation pass within the code generation pipeline, a simpler approach
is used within IREE:

-   For each `hal.interface.store.tensor` an `iree.placeholder` operation is
    created. The latter uses the same `hal.interface.binding` as the former, but
    returns a `memref` view of the output of the dispatch region instead of a
    `tensor` view. This `iree.placeholder` operation is added to start of the
    entry point function.

-   A map is constructed that for a given `tensor` records the `memref` value to
    use during the conversion. In this map the `tensor` value used in the
    `hal.interface.store.tensor` is mapped to the `memref` value returned by the
    created `iree.placeholder` operation.

-   The Dialect Conversion framework is used to implement a set of patterns that
    convert from operations on `tensor`s to operation on `memref`s,

    -   A `hal.interface.load.tensor`, is replaced with an `iree.placeholder` to
        get the `memref` view of the input to the dispatch region.
    -   All Linalg operation on `tensor`s (expected to be just `linalg.generic`
        or `linalg.indexed_generic` operations) are converted to the
        corresponding operation on `memref`s. Instead of returning a `tensor`
        value the converted operation takes an additional `memref` operand as
        argument. This `memref` is where the result of the operation is
        populated. Current implementation looks for the `memref` to use from the
        map constructed previously. If there is no `memref` associated with the
        result `tensor` the conversion fails.
    -   At this stage, any `mhlo` operation not converted to a Linalg operation
        are directly converted to a Linalg operation on buffers. This is done
        for operations that when converted to Linalg have a *reduction* iterator
        type. Some examples of ops converted this way are

        -   `mhlo.dot`
        -   `mhlo.reduce`
        -   `mhlo.conv`
        -   `mhlo.reduce_window`.

        Since the specification of the Linalg operations require the output
        `memref` to be initialized appropriately, a `linalg.fill` operation is
        used to achieve this.

__Roadmap Note__ : Right now the code-generation pipeline relies on fusion of
operations on tensor level. In the near future, we want to be able to fuse
operations like `linalg.matmul` and `linalg.conv` with consumers/producers that
are element-wise operations using the [fusion of Linalg operation on
`memref`s][LinalgFusionOnBuffers].

At this stage of the compilation all operations must have been converted to
Linalg operations on buffers. Shown below are the IR at the end of this stage
for the two examples in Snippets 1 and 2.

```mlir
func @main_ex_dispatch() {
  %0 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@ret0} : memref<4x10xf32>
  %c0 = constant 0 : index
  %1 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg0} : memref<4x5xf32>
  %2 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg1} : memref<5x10xf32>
  %cst = constant 0.000000e+00 : f32
  linalg.matmul(%1, %2, %0) :
    memref<4x5xf32>, memref<5x10xf32>, memref<4x10xf32>
  return
}
```
<a name="snippet5"></a>
Snippet 5 : Matrix-matrix multiply after conversion to
Linalg operation on `memref`s.

```mlir
#map0 = affine_map<(d0, d1) -> (d0, d1)>
func @main_ex_dispatch() {
  %0 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@ret0} : memref<10x5xf32>
  %c0 = constant 0 : index
  %1 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg0} : memref<10x5xf32>
  %2 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg1} : memref<10x5xf32>
  %3 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg2} : memref<10x5xf32>
  linalg.generic
    {args_in = 3 : i64, args_out = 1 : i64,
     indexing_maps = [#map0, #map0, #map0],
     iterator_types = ["parallel", "parallel"]} %1, %2, %3, %0 {
  ^bb0(%arg0: f32, %arg1: f32, %arg2: f32, %arg3: f32):  // no predecessors
    %4 = addf %arg0, %arg1 : f32
    %5 = mulf %4, %arg2 : f32
    linalg.yield %5 : f32
  }: memref<10x5xf32>, memref<10x5xf32>, memref<10x5xf32>, memref<10x5xf32>
  return
}
```
<a name="snippet6"></a>
Snippet 6 : Elementwise operations after conversion to Linalg operation on
`memref`s

The rest of the code-generation differs on whether the compilation is for CPU
(using LLVM) or for GPU (using SPIR-V).

## Conversion from Linalg on buffers to SPIR-V dialect

The following sections describe the progressive lowering of Linalg operation on
buffers to SPIR-V dialect. Once lowered to the SPIR-V dialect, it can be
serialized into a SPIR-V binary using the [serialization mechanism provided by
the SPIR-V dialect][SpirvSerialization]. The steps involved in the lowering are
described below, with each of the arrows representing a pass. ![Linalg on
`memref` to SPIR-V conversion](./CodegenPasses/linalg_to_spirv.png) These passes
are described below in more detail.

### Tiling and fusion on buffer operations

The GPU hardware typically provides multiple-levels of compute hierarchy, namely
*workgroup* level, *subgroup* level and *workitem* level. These map to blocks,
warps and threads, respectively, in CUDA terminology. Tiling is a way to map the
computations to each level of the compute hierarchy. For example 3-D tiling a
`linalg.matmul` operation decomposes the computation into several tiled
matrix-matrix multiplies. [Tiling transformation in Linalg
dialect][LinalgTiling] generates the outer-loops that iterate over tiled
`linalg.matmul` operations. These outer loops can be mapped to different
workgroups, if they are parallel. The tiled `linalg.matmul` operation can be
further tiled to map to subgroups. Finally, the tiled operation can be lowered
to loops with individual iterations mapped to workitems. The
[LinalgTileAndFusePass][LinalgTileAndFuse] uses the Linalg Tiling patterns
([defined here][LinalgTilingPatterns]) to tile operations like `linalg.matmul`,
`linalg.conv` and `linalg.*_pooling`. The result of tiling the code in Snippet 5
is shown below. As expected there are 2-parallel loops that iterate over tiles
of the original iteration space (i.e. inter-tile loops) and can be distributed
to workgroups.

```mlir
func @main_ex_dispatch_0() 
  attributes {
    spv.entry_point_abi = {local_size = dense<[8, 8, 1]> : vector<3xi32>}} {
  %cst = constant 0.000000e+00 : f32
  %c0 = constant 0 : index
  %c4 = constant 4 : index
  %c10 = constant 10 : index
  %0 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@ret0} : memref<4x10xf32>
  %1 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg0} : memref<4x5xf32>
  %2 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg1} : memref<5x10xf32>
  linalg.fill(%0, %cst) : memref<4x10xf32>, f32
  scf.parallel (%arg0, %arg1) = (%c0, %c0) to (%c4, %c10) step (%c8, %c8) {<commit_msg>[doc] Fix broken picture links (#2636)<commit_after># IREE CPU/GPU Code Generation Pipeline

This document is intended to provide an overview of the codegen pipeline within
IREE used to generate CPU/GPU code. It intends to give an overview of the main
passes used, the objective of the pass, the current implementation, and what it
is expected to achieve in the long term.

Note that while the code generation pipeline supports dynamic shapes, this work
is very preliminary. The description of this is not covered here.

## Input to the codegen pipeline

The input to the code generation pipeline is the module within the
`hal.executable.target` operation. Functions within this module that do __not__
have `Visibility::Private` are the *entry point* functions of the dispatch
region. These are the functions that are *invoked* by the IREE runtime. In
addition, each dispatch region also contains a `hal.interface` operation that
describes the ABI to use for the dispatch region. Two examples of the input to
the code generation pipeline are shown below. In both of these, a single
dispatch function contains a sequence of MHLO operations that the dispatch
region creation has grouped into a single region. Ideally the grouped operations
are fused into a single kernel.

```mlir
hal.executable.target "vulkan*" {
  module attributes {spv.target_env = ...} {
    func @main_ex_dispatch() {
      %c0 = constant 0 : index
      %0 = hal.interface.load.tensor @legacy_io::@arg0,
             offset = %c0 : tensor<4x5xf32>
      %1 = hal.interface.load.tensor @legacy_io::@arg1,
             offset = %c0 : tensor<5x10xf32>
      %2 = "mhlo.dot"(%0, %1) {precision_config = ["DEFAULT", "DEFAULT"]} :
             (tensor<4x5xf32>, tensor<5x10xf32>) -> tensor<4x10xf32>
      hal.interface.store.tensor %2, @legacy_io::@ret0,
        offset = %c0 : tensor<4x10xf32>
      return
    }
    hal.interface @legacy_io attributes {sym_visibility = "private"} {
      hal.interface.binding @arg0, set=0, binding=0,
        type="StorageBuffer", access="Read"
      hal.interface.binding @arg1, set=0, binding=1,
        type="StorageBuffer", access="Read"
      hal.interface.binding @ret0, set=0, binding=2,
        type="StorageBuffer", access="Write|Discard"
    }
  }
}
```
<a name="snippet1"></a>
Snippet 1 : Dispatch region with matrix-matrix multiply operation.

```mlir
hal.executable.target "vulkan*" {
  module attributes {spv.target_env = ...} {
    func @main_ex_dispatch() {
      %c0 = constant 0 : index
      %0 = hal.interface.load.tensor @legacy_io::@arg0,
             offset = %c0 : tensor<10x5xf32>
      %1 = hal.interface.load.tensor @legacy_io::@arg1,
             offset = %c0 : tensor<10x5xf32>
      %2 = hal.interface.load.tensor @legacy_io::@arg2,
             offset = %c0 : tensor<10x5xf32>
      %3 = "mhlo.add"(%0, %1) :
         (tensor<10x5xf32>, tensor<10x5xf32>) -> tensor<10x5xf32>
      %4 = "mhlo.multiply"(%3, %2) :
         (tensor<10x5xf32>, tensor<10x5xf32>) -> tensor<10x5xf32>
      hal.interface.store.tensor %4, @legacy_io::@ret0,
        offset = %c0 : tensor<10x5xf32>
      return
    }
    hal.interface @legacy_io attributes {sym_visibility = "private"} {
      hal.interface.binding @arg0, set=0, binding=0,
        type="StorageBuffer", access="Read"
      hal.interface.binding @arg1, set=0, binding=1,
        type="StorageBuffer", access="Read"
      hal.interface.binding @arg2, set=0, binding=2,
        type="StorageBuffer", access="Read"
      hal.interface.binding @ret0, set=0, binding=3,
        type="StorageBuffer", access="Write|Discard"
    }
  }
}
```
<a name="snippet2"></a>
Snippet 2 : Dispatch region with element-wise operations.

__Roadmap Note__: The current implementation might not actually fuse the
operations grouped into a dispatch region into a single kernel. It is possible
to end up with multiple kernels per dispatch region. Over time we plan to address
this by using fusion at different levels (see below).

The inputs to the dispatch region are materialized within the entry point
function using the `hal.interface.load.tensor` operation, This operation returns
a `tensor` view of the buffer used to store the inputs. Similarly the result of
the dispatch region are *written* out using the `hal.interface.store.tensor`
operation.

The main constraint that the code generation operates under is that it should
not require additional (temporary) buffers to execute the operations grouped
together within a dispatch region. The rationale behind this constraint is that
buffer allocation/synchronization in IREE happens at the granularity of dispatch
regions, allowing the scheduler to make better decision about where to insert
appropriate synchronizations.

The IR after all the passes used in the lowering from MHLO to SPIR-V for the
above two examples can be found here ([matrix-matrix multiply op][DotAfterAll],
[elementwise ops][PwAfterAll]). Below is a description of the major passes used.

## Conversion from MHLO dialect to Linalg on buffers

The code generation pipeline heavily relies on use of [Structured
Operations][LinalgRationale], specifically the [Linalg Dialect][LinalgDialect].
Both, the Linalg operations on `tensor`s and on `memref`s are central to the
progressive lowering approach followed here. The first part of the code
generation pipeline is to convert the MHLO operations on `tensor`s to Linalg
operation on `memref`s. This part of the pipeline is common to both CPU and GPU
code generation.

The steps involved in this conversion is shown below. Each of the arrows
represents a pass in the pipeline:

![MHLO To Linalg on `memref` conversion](./hlo_to_linalg.png)

The next sections describe each of these passes in more detail.


### MHLO to Linalg on tensors

The first step is to convert MHLO operations to Linalg on tensors. This is done
using the [HLOToLinalgPass][HLOToLinalgPass] from Tensorflow. An example of the
conversion is shown below, where each of the `mhlo.add` and `mhlo.multiply`
operations are converted to `linalg.generic` operations on tensors.

```mlir
#map0 = affine_map<(d0, d1) -> (d0, d1)>
%3 = linalg.generic
       {args_in = 2 : i64, args_out = 1 : i64,
        indexing_maps = [#map0, #map0, #map0],
        iterator_types = ["parallel", "parallel"]} %0, %1 {
     ^bb0(%arg0: f32, %arg1: f32):  // no predecessors
       %5 = addf %arg0, %arg1 : f32
       linalg.yield %5 : f32
     } : tensor<10x5xf32>, tensor<10x5xf32> -> tensor<10x5xf32>
%4 = linalg.generic
       {args_in = 2 : i64, args_out = 1 : i64,
        indexing_maps = [#map0, #map0, #map0],
        iterator_types = ["parallel", "parallel"]} %3, %2 {
     ^bb0(%arg0: f32, %arg1: f32):  // no predecessors
       %5 = mulf %arg0, %arg1 : f32
       linalg.yield %5 : f32
     }: tensor<10x5xf32>, tensor<10x5xf32> -> tensor<10x5xf32>
```
<a name="snippet3"></a>
Snippet 3 : MHLO to Linalg conversion for [element-wise operations](#snippet2)

At the time of writing the representation of Linalg on `tensor`s does not model
reduction iterator types completely. Specifically, the reduction in Linalg is
modeled using read-modify-write approach, i.e. each iteration of the reduction
loop reads the value stored in the output, adds its contribution, and writes
back to the same location. This means the output has to be *initialized* to the
null element of the reduction operator (i.e. 0 if the reduction is done using
addition). This works for operations on buffers. Since tensors are SSA values
they cannot be updated in-place. As a result, the reduction semantics does not
map as well to `tensor`s. For now it is treated as a convention that when the
Linalg operation is converted to use `memref`s it has to be initialized
appropriately before performing the reduction. Due to this, the conversion from
MHLO op to Linalg op is only done for operations which do not need a *reduction*
iterator type in the converted Linalg op. Consequently, only element-wise
operations, broadcast operations and data movement operations (like copy and
transpose) are converted to Linalg operations at this stage.

__Roadmap note__: One long term solution for the above is to have operations on
tensors that have *reduction* iterator type to take an additional argument that
contains the initial value of the result tensor. When the operation is converted
to use `memref`s, the buffer for the initial value operand can be reused for the
result. The details involved have not been fully worked out yet.

### Fusion of Linalg on tensor operations

The Linalg on `tensor` operations generated at the previous step are fused using
the [LinalgFusionOfTensorOps][LinalgFusionOfTensorOps] from MLIR. Since
`tensor`s are SSA values, fusion at this stage can be done without using alias
analysis or dependence analysis based on reads and writes.  Instead the use-def
chains for the `tensor` values can be used to implement producer-consumer
fusion. This stage fuses most elementwise operations, broadcast operations and
data movement operations. An example of the fused op is shown below.

```mlir
#map0 = affine_map<(d0, d1) -> (d0, d1)>
%3 = linalg.generic
       {args_in = 3 : i64, args_out = 1 : i64,
        indexing_maps = [#map0, #map0, #map0, #map0],
        iterator_types = ["parallel", "parallel"]} %0, %1, %2 {
     ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):  // no predecessors
       %4 = addf %arg0, %arg1 : f32
       %5 = mulf %4, %arg2 : f32
       linalg.yield %5 : f32
     }: tensor<?x5xf32>, tensor<?x5xf32>, tensor<?x5xf32> -> tensor<?x5xf32>
```
<a name="snippet4"></a>
Snippet 4: Fusion of Linalg operation on tensors for element-wise operations
shown in [Snippet 3](#snippet3)

### Conversion of Linalg on tensors to Linalg on buffers

Post fusion all the operation on `tensor`s are converted to analogous operations
on `memref`s. In general, this requires a buffer allocation pass. In IREE,
buffer allocation happens at the granularity of dispatch region, and as
mentioned [earlier](#input-to-the-codegen-pipeline), the dispatch region is not
expected to use any additional temporary buffers. So instead of having another
buffer allocation pass within the code generation pipeline, a simpler approach
is used within IREE:

-   For each `hal.interface.store.tensor` an `iree.placeholder` operation is
    created. The latter uses the same `hal.interface.binding` as the former, but
    returns a `memref` view of the output of the dispatch region instead of a
    `tensor` view. This `iree.placeholder` operation is added to start of the
    entry point function.

-   A map is constructed that for a given `tensor` records the `memref` value to
    use during the conversion. In this map the `tensor` value used in the
    `hal.interface.store.tensor` is mapped to the `memref` value returned by the
    created `iree.placeholder` operation.

-   The Dialect Conversion framework is used to implement a set of patterns that
    convert from operations on `tensor`s to operation on `memref`s,

    -   A `hal.interface.load.tensor`, is replaced with an `iree.placeholder` to
        get the `memref` view of the input to the dispatch region.
    -   All Linalg operation on `tensor`s (expected to be just `linalg.generic`
        or `linalg.indexed_generic` operations) are converted to the
        corresponding operation on `memref`s. Instead of returning a `tensor`
        value the converted operation takes an additional `memref` operand as
        argument. This `memref` is where the result of the operation is
        populated. Current implementation looks for the `memref` to use from the
        map constructed previously. If there is no `memref` associated with the
        result `tensor` the conversion fails.
    -   At this stage, any `mhlo` operation not converted to a Linalg operation
        are directly converted to a Linalg operation on buffers. This is done
        for operations that when converted to Linalg have a *reduction* iterator
        type. Some examples of ops converted this way are

        -   `mhlo.dot`
        -   `mhlo.reduce`
        -   `mhlo.conv`
        -   `mhlo.reduce_window`.

        Since the specification of the Linalg operations require the output
        `memref` to be initialized appropriately, a `linalg.fill` operation is
        used to achieve this.

__Roadmap Note__ : Right now the code-generation pipeline relies on fusion of
operations on tensor level. In the near future, we want to be able to fuse
operations like `linalg.matmul` and `linalg.conv` with consumers/producers that
are element-wise operations using the [fusion of Linalg operation on
`memref`s][LinalgFusionOnBuffers].

At this stage of the compilation all operations must have been converted to
Linalg operations on buffers. Shown below are the IR at the end of this stage
for the two examples in Snippets 1 and 2.

```mlir
func @main_ex_dispatch() {
  %0 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@ret0} : memref<4x10xf32>
  %c0 = constant 0 : index
  %1 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg0} : memref<4x5xf32>
  %2 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg1} : memref<5x10xf32>
  %cst = constant 0.000000e+00 : f32
  linalg.matmul(%1, %2, %0) :
    memref<4x5xf32>, memref<5x10xf32>, memref<4x10xf32>
  return
}
```
<a name="snippet5"></a>
Snippet 5 : Matrix-matrix multiply after conversion to
Linalg operation on `memref`s.

```mlir
#map0 = affine_map<(d0, d1) -> (d0, d1)>
func @main_ex_dispatch() {
  %0 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@ret0} : memref<10x5xf32>
  %c0 = constant 0 : index
  %1 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg0} : memref<10x5xf32>
  %2 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg1} : memref<10x5xf32>
  %3 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg2} : memref<10x5xf32>
  linalg.generic
    {args_in = 3 : i64, args_out = 1 : i64,
     indexing_maps = [#map0, #map0, #map0],
     iterator_types = ["parallel", "parallel"]} %1, %2, %3, %0 {
  ^bb0(%arg0: f32, %arg1: f32, %arg2: f32, %arg3: f32):  // no predecessors
    %4 = addf %arg0, %arg1 : f32
    %5 = mulf %4, %arg2 : f32
    linalg.yield %5 : f32
  }: memref<10x5xf32>, memref<10x5xf32>, memref<10x5xf32>, memref<10x5xf32>
  return
}
```
<a name="snippet6"></a>
Snippet 6 : Elementwise operations after conversion to Linalg operation on
`memref`s

The rest of the code-generation differs on whether the compilation is for CPU
(using LLVM) or for GPU (using SPIR-V).

## Conversion from Linalg on buffers to SPIR-V dialect

The following sections describe the progressive lowering of Linalg operation on
buffers to SPIR-V dialect. Once lowered to the SPIR-V dialect, it can be
serialized into a SPIR-V binary using the [serialization mechanism provided by
the SPIR-V dialect][SpirvSerialization]. The steps involved in the lowering are
described below, with each of the arrows representing a pass.

![Linalg on `memref` to SPIR-V conversion](./linalg_to_spirv.png)

These passes are described below in more detail.

### Tiling and fusion on buffer operations

The GPU hardware typically provides multiple-levels of compute hierarchy, namely
*workgroup* level, *subgroup* level and *workitem* level. These map to blocks,
warps and threads, respectively, in CUDA terminology. Tiling is a way to map the
computations to each level of the compute hierarchy. For example 3-D tiling a
`linalg.matmul` operation decomposes the computation into several tiled
matrix-matrix multiplies. [Tiling transformation in Linalg
dialect][LinalgTiling] generates the outer-loops that iterate over tiled
`linalg.matmul` operations. These outer loops can be mapped to different
workgroups, if they are parallel. The tiled `linalg.matmul` operation can be
further tiled to map to subgroups. Finally, the tiled operation can be lowered
to loops with individual iterations mapped to workitems. The
[LinalgTileAndFusePass][LinalgTileAndFuse] uses the Linalg Tiling patterns
([defined here][LinalgTilingPatterns]) to tile operations like `linalg.matmul`,
`linalg.conv` and `linalg.*_pooling`. The result of tiling the code in Snippet 5
is shown below. As expected there are 2-parallel loops that iterate over tiles
of the original iteration space (i.e. inter-tile loops) and can be distributed
to workgroups.

```mlir
func @main_ex_dispatch_0()
  attributes {
    spv.entry_point_abi = {local_size = dense<[8, 8, 1]> : vector<3xi32>}} {
  %cst = constant 0.000000e+00 : f32
  %c0 = constant 0 : index
  %c4 = constant 4 : index
  %c10 = constant 10 : index
  %0 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@ret0} : memref<4x10xf32>
  %1 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg0} : memref<4x5xf32>
  %2 = iree.placeholder for "interface buffer"
         {binding = @legacy_io::@arg1} : memref<5x10xf32>
  linalg.fill(%0, %cst) : memref<4x10xf32>, f32
  scf.parallel (%arg0, %arg1) = (%c0, %c0) to (%c4, %c10) step (%c8, %c8) {
{'fraction_non_alphanumeric': 0.08970545348498105, 'fraction_numerical': 0.029512977544473608, 'mean_word_length': 3.8715726665719563, 'pattern_counts': {'":': 0, '<': 125, '<?xml version=': 0, '>': 143, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5376069', 'n_tokens_mistral': 10895, 'n_tokens_neox': 10292, 'n_words': 4227}