Adding description of IREE Codegen pass pipeline. This change adds description for the IREE Codegen pass pipeline used to generate SPIR-V code. The LLVM side is TODO. Also accompanied minor change to make it consistent with the description.
diff --git a/docs/design_docs/codegen_passes.md b/docs/design_docs/codegen_passes.md new file mode 100644 index 0000000..5177d37 --- /dev/null +++ b/docs/design_docs/codegen_passes.md
@@ -0,0 +1,636 @@ +# 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  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.  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) { + scf.for %arg2 = %c0 to %c5 step %c4 { + ... + %5 = subview %1[%arg0, %arg2]... + ... + %8 = subview %2[%arg2, %arg1]... + ... + %11 = subview %0[%arg0, %arg1].. + linalg.matmul {__internal_linalg_transform__ = "workgroup"} %5, %8, %11... + } + scf.yield + } + return +} +``` +<a name="snippet7"></a> +Snippet 7 : `linalg.matmul` after tiling. + +#### Tile Size and Workgroup Size + +When operations that are to be tiled exist within the dispatch function (like +`linalg.matmul` or `linalg.conv`), this pass also decides the +1. Tile size to be used for the tiling. +1. The workgroup size to be used. + +The tile size and workgroup size are closely linked since the code within the +tiled loops are to be collectively executed by the entire workgroup. In other +words, all workitems in the workgroup collaborate to execute the tiled +`linalg.matmul`. + +__Roadmap Note__ : Currently the tile sizes used in this pass are hard-wired. +Not much effort has been put into finding ideal tile size for each operation on +different hardware. The value used is meant to be a baseline to test +functionality, with performance considerations addressed over time. + +#### Markers + +Downstream passes have to handle tiled Linalg operations and untiled Linalg +operation that might exist in the same function in different ways. For example, +while the former are to be executed collectively by workitems within a +workgroup, the latter have to be executed by all workitems across +workgroups. One way to distinguish these two operations is to use the marker +mechanism in Linalg ([LinalgMarker][LinalgTilingPatterns]). This is a `StrAttr` +whose value can be used to encode the scope of the operation. For example, in +Snippet 7 above, the tiled `linalg.matmul` operation has a marker `workgroup` to +indicate that this operation needs to be executed by a workgroup in a collective +manner. At this time, the code-generation pipeline uses only the `workgroup` +marker. + +__Roadmap Note__ : Markers are meant to be short-lived, ideally set and consumed +within the same pass. In the current pipeline the lifetime spans passes to allow +lowering to different hierarchies. The separate passes that implement the +lowering from Linalg to SPIR-V can be combined into a single pass, relying A -> +B -> C translation mechanism of the Dialect Conversion framework to implement +the progressive lowering. In interest of separation of concerns and for better +debuggability these passes are kept separate at the cost of having lifetimes of +markers span passes. + +#### Promoting subviews to use workgroup local memory and use of synchronizations + +`Workgroup` memory (or `shared memory` in CUDA terminology) can be used to +prefetch the inputs to the tiled operation. For example in the matrix-matrix +multiply case, the same data row (column) of the LHS (RHS) matrix is read by +multiple workitems. Prefetching the data into `Workgroup` memory can reduce the +number of loads to `StorageClass` memory by an order of magnitude. This +transformation can be achieved by using the [`Linalg +Promotion`][LinalgPromotionPatterns] which modifies the `subview`s that are the +operands to the tiled Linalg operation to use a new `memref` object. The size of +this `memref` is computed from the size of the `subview`. This `memref` object +is later lowered to use `Workgroup` memory Storage Class. The snippet below +shows this transformation when applied to `linalg.matmul` (along with +tiling). The newly created `memref` objects are annotated with the memory space +`3` to indicate that they are to be lowered to use `Workgroup` memory. The copy +of data from the original `memref` into the new `memref`, as well as the +necessary synchronization constructs are generated as well. Note the memory +space annotation used here is consistent with what [address space annotations +used in NVVM][NVVMAddressSpace]. + +```mlir +func @matmul_tile() + attributes { + spv.entry_point_abi = {local_size = dense<[8, 8, 1]> : vector<3xi32>}} { + %c96 = constant 96 : index + %c4 = constant 4 : index + %c8 = constant 8 : index + %c0 = constant 0 : index + %c1 = constant 1 : index + %0 = iree.placeholder for "interface buffer" + {binding = @legacy_io::@arg0} : memref<96x96xf32> + %1 = iree.placeholder for "interface buffer" + {binding = @legacy_io::@arg1} : memref<96x96xf32> + %2 = iree.placeholder for "interface buffer" + {binding = @legacy_io::@ret0} : memref<96x96xf32> + scf.parallel (%arg0, %arg1) = (%c0, %c0) to (%c96, %c96) step (%c8, %c8) { + scf.for %arg2 = %c0 to %c96 step %c4 { + ... + %5 = subview %0[%arg0, %arg2]... + ... + %8 = subview %1[%arg2, %arg1]... + ... + %11 = subview %2[%arg0, %arg1]... + %12 = alloc(%c8, %c4) : memref<?x?xf32, 3> + %13 = subview %12[%c0, %c0]... + %14 = alloc(%c4, %c8) : memref<?x?xf32, 3> + %15 = subview %14[%c0, %c0]... + linalg.copy(%5, %13) {__internal_linalg_transform__ = "workgroup"} + : memref<?x?xf32, #map2>, memref<?x?xf32, #map2, 3> + spv.ControlBarrier "Workgroup", "Workgroup", "AcquireRelease" + linalg.copy(%8, %15) {__internal_linalg_transform__ = "workgroup"} + : memref<?x?xf32, #map2>, memref<?x?xf32, #map2, 3> + spv.ControlBarrier "Workgroup", "Workgroup", "AcquireRelease" + linalg.matmul {__internal_linalg_transform__ = "workgroup"} %13, %15, %11... + spv.ControlBarrier "Workgroup", "Workgroup", "AcquireRelease" + dealloc %12 : memref<?x?xf32, 3> + dealloc %14 : memref<?x?xf32, 3> + } + scf.yield + } + return +} +``` + +<a name="snippet8"></a> +Snippet 8: `linalg.matmul` after tiling and promotion of operand subviews to use +`Workgroup` memory. + +### Distributing to workgroups and workitems + +After tiling the operations within the dispatch functions are either +`scf.parallel` operations or Linalg operations. + +- The outer `scf.parallel` operations represent parallel loops that are to be + distributed across workgroups. The distribution here assumes that the number + of workgroups along each dimension is equal to the number of iterations of the + `scf.parallel` operation. + +- Linalg operations that are not tiled, and are therefore __not within__ `scf` + operations, are lowered to loops. The resulting outer `scf.parallel` operations are + collapsed to have a single induction variable. This loop is then distributed + across workitems using their `GlobalInvocationId`, (which is same as + `blockIdx * blockDim + threadIdx` in CUDA terminology). + +- Linalg operations that are tiled, and are therefore __within__ `scf` + operations, are lowered to loops and the iterations of the `scf.parallel` + operations are mapped to workitems using their `LocalInvocationId` (which is + same as `threadIdx` in CUDA terminology). Note that these operations are + tagged with the `workgroup` marker which makes it easy to disambiguate from + the case where Linalg operations are outside of `scf` operations. Here too, + the distribution assumes that the workgroup size is greater than or equal to + the number of iterations of the partitioned loop. + +These transformations are applied by the +[`ConvertToGPUPass`][ConvertToGPU]. Below is the result of applying this pass to +Snippet 7. The outer `scf.parallel` loop is distributed across workgroups. The +tiled `linalg.matmul` operation is lowered to loops, and the outer +`scf.parallel` operation generated during this lowering are distributed across +workitems within the workgroup. + +```mlir +func @main_ex_dispatch_0_dispatch_1() + attributes { + spv.entry_point_abi = {local_size = dense<[8, 8, 1]> : vector<3xi32>}} { + %c5 = constant 5 : index + %c8 = constant 8 : index + %c4 = constant 4 : index + %c0 = constant 0 : index + %c1 = constant 1 : 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> + %3 = "gpu.block_id"() {dimension = "x"} : () -> index + %4 = muli %3, %c8 : index + scf.for %arg0 = %c0 to %c5 step %c4 { + ... + %9 = subview %1[0, %arg0] + ... + %14 = subview %2[%arg0, %4] + %15 = subview %0[0, %4] + %16 = "gpu.thread_id"() {dimension = "x"} : () -> index + %17 = "gpu.thread_id"() {dimension = "y"} : () -> index + %18 = cmpi "slt", %17, %c4 : index + %19 = cmpi "slt", %16, %13 : index + %20 = and %18, %19 : i1 + scf.if %20 { + scf.for %arg1 = %c0 to %8 step %c1 { + %21 = load %9[%17, %arg1] : memref<4x?xf32, #map0> + %22 = load %14[%arg1, %16] : memref<?x?xf32, #map1> + %23 = load %15[%17, %16] : memref<4x?xf32, #map1> + %24 = mulf %21, %22 : f32 + %25 = addf %23, %24 : f32 + store %25, %15[%17, %16] : memref<4x?xf32, #map1> + } + } + } + return +} +``` +<a name="snippet9"></a> +Snippet 9: `linalg.matmul` after distributing parallel inter-tile loops to +workgroups and intra-tile loops to workitems. + +[Snippet 6](#snippet6) shows the fused element-wise operations represented using +a `linalg.generic` operation. This operation is not tiled in the +`LinalgTileAndFusePass`. So the `ConvertToGPUPass` lowers this operation to +`scf.parallel` loops, which are collapsed into a `scf.parallel` operation with a +single induction variable. This loop is then distributed across workitems using +the `GlobalInvocationId`. The resulting IR is shown below. + +```mlir +func @main_ex_dispatch_0() + attributes { + spv.entry_point_abi = {local_size = dense<[32, 1, 1]> : vector<3xi32>}} { + %c50 = constant 50 : index + %c5 = constant 5 : index + %0 = iree.placeholder for "interface buffer" + {binding = @legacy_io::@ret0} : memref<10x5xf32> + %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> + %4 = "gpu.block_id"() {dimension = "x"} : () -> index + %5 = "gpu.block_dim"() {dimension = "x"} : () -> index + %6 = "gpu.thread_id"() {dimension = "x"} : () -> index + %7 = muli %4, %5 : index + %8 = addi %7, %6 : index + %9 = cmpi "slt", %8, %c50 : index + scf.if %9 { + %10 = divi_signed %8, %c5 : index + %11 = remi_signed %8, %c5 : index + %12 = load %1[%10, %11] : memref<10x5xf32> + %13 = load %2[%10, %11] : memref<10x5xf32> + %14 = load %3[%10, %11] : memref<10x5xf32> + %15 = addf %12, %13 : f32 + %16 = mulf %15, %14 : f32 + store %16, %0[%10, %11] : memref<10x5xf32> + } + return +} +``` +<a name="snippet10"></a> +Snippet 10: Distributing the iterations for pointwise operations for GPU execution. + +### Lowering to SPIR-V dialect + +The last step is to take the result of the previous pass and lowering it to +SPIR-V dialect. Since SPIR-V dialect is *closed*, i.e. it has a separate type +system, its best to lower all the operations to SPIR-V in one step. This is done +by applying all the patterns that lower all the different IR constructs into +SPIR-V within the [`ConvertToSPIRVPass`][ConvertToSPIRV]. These are + +- [GPU dialect to SPIR-V conversion][GPUToSPIRV]. +- [SCF dialect to SPIR-V conversion][SCFToSPIRV]. +- [Standard dialect to SPIR-V conversion][StandardToSPIRV]. +- Patterns that lower the `iree.placeholder` instruction into a SPIR-V. + +Once applied the resulting IR is in SPIR-V dialect that can be serialized to a +SPIR-V binary. + +[ConvertToGPU]: https://github.com/google/iree/blob/main/iree/compiler/Conversion/LinalgToSPIRV/ConvertToGPUPass.cpp +[ConvertToSPIRV]: https://github.com/google/iree/blob/main/iree/compiler/Conversion/LinalgToSPIRV/ConvertToSPIRVPass.cpp +[DotAfterAll]: https://gist.github.com/MaheshRavishankar/9e2d406296f469515c4a79bf1e7eef44 +[GPUToSPIRV]: https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.h +[HLOToLinalgPass]: https://github.com/tensorflow/tensorflow/blob/75c40f6bff2faa3d90a375dfa4025b2e6e2d7a3d/tensorflow/compiler/mlir/xla/transforms/passes.h#L67 +[LinalgDialect]: https://mlir.llvm.org/docs/Dialects/Linalg/ +[LinalgFusionOnBuffers]: https://github.com/llvm/llvm-project/blob/ef868a848e6def288d2df7a1b3ebe09463afc8d0/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h#L86 +[LinalgFusionOfTensorOps]: https://github.com/llvm/llvm-project/blob/80cb25cbd555f9634836b766c86aead435b60eaa/mlir/include/mlir/Dialect/Linalg/Passes.td#L30 +[LinalgPromotionPatterns]: https://github.com/llvm/llvm-project/blob/303a7f7a26e2aae1cb85f49dccbc0b5d14e0b2e0/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h#L358 +[LinalgRationale]: https://mlir.llvm.org/docs/Rationale/RationaleLinalgDialect/ +[LinalgTileAndFuse]: https://github.com/google/iree/blob/main/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp +[LinalgTiling]: https://mlir.llvm.org/docs/Dialects/Linalg/#set-of-key-transformationsa-namekey_transformationsa +[LinalgTilingPatterns]: https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h +[NVVMAddressSpace]: https://docs.nvidia.com/cuda/nvvm-ir-spec/index.html#address-space +[PwAfterAll]: https://gist.github.com/MaheshRavishankar/02cdd22f7c99e568f933244b5a679510 +[SCFToSPIRV]: https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/Conversion/SCFToSPIRV/SCFToSPIRV.h +[SpirvSerialization]: https://mlir.llvm.org/docs/Dialects/SPIR-V/#serialization-and-deserialization +[StandardToSPIRV]: https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.h
diff --git a/docs/design_docs/hlo_to_linalg.png b/docs/design_docs/hlo_to_linalg.png new file mode 100755 index 0000000..469ed26 --- /dev/null +++ b/docs/design_docs/hlo_to_linalg.png Binary files differ
diff --git a/docs/design_docs/linalg_to_spirv.png b/docs/design_docs/linalg_to_spirv.png new file mode 100755 index 0000000..fd6aee7 --- /dev/null +++ b/docs/design_docs/linalg_to_spirv.png Binary files differ
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/ConvertToGPUPass.cpp b/iree/compiler/Conversion/LinalgToSPIRV/ConvertToGPUPass.cpp index 98e91fa..81514ec 100644 --- a/iree/compiler/Conversion/LinalgToSPIRV/ConvertToGPUPass.cpp +++ b/iree/compiler/Conversion/LinalgToSPIRV/ConvertToGPUPass.cpp
@@ -563,7 +563,7 @@ ConversionPatternRewriter &rewriter) const override { // Check for marker that specifies that the linalg op is to be partitioned // across threads within a workgroup. - if (!hasWorkItemMarker(linalgOp)) return failure(); + if (!hasWorkGroupMarker(linalgOp)) return failure(); Optional<linalg::LinalgLoops> loops = linalg::linalgLowerOpToLoops<scf::ParallelOp>(rewriter, linalgOp); if (!loops) return failure(); @@ -587,7 +587,7 @@ LogicalResult matchAndRewrite( LinalgOpTy linalgOp, ArrayRef<Value> operands, ConversionPatternRewriter &rewriter) const override { - if (!hasWorkItemMarker(linalgOp)) return failure(); + if (!hasWorkGroupMarker(linalgOp)) return failure(); Optional<linalg::LinalgLoops> loops = linalg::linalgLowerOpToLoops<scf::ParallelOp>(rewriter, linalgOp); if (!loops) return failure();
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp b/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp index e9dddd6..934e5ae 100644 --- a/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp +++ b/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp
@@ -314,7 +314,7 @@ LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) const override { - if (!hasWorkItemMarker(op)) return failure(); + if (!hasWorkGroupMarker(op)) return failure(); return linalg::LinalgPromotionPattern<linalg::MatmulOp>::matchAndRewrite( op, rewriter); } @@ -365,7 +365,7 @@ .setLoopType(linalg::LinalgTilingLoopType::ParallelLoops), tileSizeCalculator.getWorkGroupSize(), linalg::LinalgMarker(ArrayRef<Identifier>(), - Identifier::get(getWorkItemMarker(), context))); + Identifier::get(getWorkGroupMarker(), context))); applyPatternsAndFoldGreedily(getOperation(), tilingPatterns); if (useWorkgroupMemory) { @@ -385,7 +385,7 @@ [&](OpBuilder &b, Value src, Value dst) -> LogicalResult { return copyToFromWorkgroupMemory(b, src, dst); }), - linalg::LinalgMarker(Identifier::get(getWorkItemMarker(), context), + linalg::LinalgMarker(Identifier::get(getWorkGroupMarker(), context), Identifier::get(PromotionMarker, context))); applyPatternsAndFoldGreedily(getOperation(), promotionPatterns); } @@ -394,7 +394,7 @@ OpBuilder builder(context); funcOp.walk([&builder](linalg::LinalgOp linalgOp) { if (hasMarker(linalgOp, PromotionMarker)) { - setWorkItemMarker(linalgOp); + setWorkGroupMarker(linalgOp); insertBarrierAfter(builder, linalgOp.getLoc(), linalgOp); } });
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/MarkerUtils.cpp b/iree/compiler/Conversion/LinalgToSPIRV/MarkerUtils.cpp index c874234..47747de 100644 --- a/iree/compiler/Conversion/LinalgToSPIRV/MarkerUtils.cpp +++ b/iree/compiler/Conversion/LinalgToSPIRV/MarkerUtils.cpp
@@ -36,8 +36,6 @@ StringRef getWorkGroupMarker() { return "workgroup"; } -StringRef getWorkItemMarker() { return "workitem"; } - bool hasMarker(Operation *op, StringRef marker) { return checkMarkerValue(op, marker); } @@ -46,10 +44,6 @@ return checkMarkerValue(op, getWorkGroupMarker()); } -bool hasWorkItemMarker(Operation *op) { - return checkMarkerValue(op, getWorkItemMarker()); -} - void setMarker(Operation *op, StringRef marker) { op->setAttr(linalg::LinalgTransforms::kLinalgTransformMarker, StringAttr::get(marker, op->getContext())); @@ -57,6 +51,5 @@ void setWorkGroupMarker(Operation *op) { setMarker(op, getWorkGroupMarker()); } -void setWorkItemMarker(Operation *op) { setMarker(op, getWorkItemMarker()); } } // namespace iree_compiler } // namespace mlir
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/MarkerUtils.h b/iree/compiler/Conversion/LinalgToSPIRV/MarkerUtils.h index 36dccca..e512ead 100644 --- a/iree/compiler/Conversion/LinalgToSPIRV/MarkerUtils.h +++ b/iree/compiler/Conversion/LinalgToSPIRV/MarkerUtils.h
@@ -31,7 +31,7 @@ namespace iree_compiler { /// Marker to denote that a linalg operation is to be partitioned to workitems. -StringRef getWorkItemMarker(); +StringRef getWorkGroupMarker(); /// Returns true if an operation has the specified `marker`. When `marker` is /// empty, returns true if the operation has any marker. @@ -39,14 +39,14 @@ /// Returns true if an operation has marker to denote that it is to be /// partitioned to workitems. -bool hasWorkItemMarker(Operation *); +bool hasWorkGroupMarker(Operation *); /// Sets a given marker on an operation. void setMarker(Operation *, StringRef); /// Sets marker to denote that a linalg operation is to be partitioned to /// workitems. -void setWorkItemMarker(Operation *); +void setWorkGroupMarker(Operation *); } // namespace iree_compiler } // namespace mlir
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_gpu.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_gpu.mlir index 679f523..64621f3 100644 --- a/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_gpu.mlir +++ b/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_gpu.mlir
@@ -162,7 +162,7 @@ %12 = dim %arg2, %c1 : memref<?x?xf32> %13 = affine.min #map0(%arg4)[%12] %14 = subview %arg2[%arg3, %arg4] [%11, %13] [1, 1] : memref<?x?xf32> to memref<?x?xf32, #map2> - linalg.matmul %5, %9, %14 {__internal_linalg_transform__ = "workitem"} : (memref<?x?xf32, #map2>, memref<?x?xf32, #map2>, memref<?x?xf32, #map2>) + linalg.matmul %5, %9, %14 {__internal_linalg_transform__ = "workgroup"} : (memref<?x?xf32, #map2>, memref<?x?xf32, #map2>, memref<?x?xf32, #map2>) } scf.yield } @@ -235,7 +235,7 @@ %13 = affine.min #map5(%arg5)[%4] %14 = dim %arg2, %c3 : memref<?x?x?x?xf32> %15 = subview %arg2[%arg3, %arg4, %arg5, 0] [%11, %12, %13, %14] [1, 1, 1, 1] : memref<?x?x?x?xf32> to memref<?x?x?x?xf32, #map3> - linalg.conv(%arg0, %9, %15) {__internal_linalg_transform__ = "workitem", dilations = [1, 1], strides = [1, 1]} : memref<?x?x?x?xf32>, memref<?x?x?x?xf32, #map3>, memref<?x?x?x?xf32, #map3> + linalg.conv(%arg0, %9, %15) {__internal_linalg_transform__ = "workgroup", dilations = [1, 1], strides = [1, 1]} : memref<?x?x?x?xf32>, memref<?x?x?x?xf32, #map3>, memref<?x?x?x?xf32, #map3> scf.yield } return @@ -364,7 +364,7 @@ %9 = affine.min #map3(%arg3)[%2] %10 = affine.min #map4(%arg4)[%3] %11 = subview %arg2[%arg3, %arg4] [%9, %10] [1, 1] : memref<?x?xf32> to memref<?x?xf32, #map2> - linalg.pooling_max(%8, %arg1, %11) {__internal_linalg_transform__ = "workitem", dilations = [1, 1], strides = [1, 1]} : memref<?x?xf32, #map2>, memref<?x?xf32>, memref<?x?xf32, #map2> + linalg.pooling_max(%8, %arg1, %11) {__internal_linalg_transform__ = "workgroup", dilations = [1, 1], strides = [1, 1]} : memref<?x?xf32, #map2>, memref<?x?xf32>, memref<?x?xf32, #map2> scf.yield } return
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_gpu_option.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_gpu_option.mlir index 1701535..63f8aa5 100644 --- a/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_gpu_option.mlir +++ b/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_gpu_option.mlir
@@ -32,7 +32,7 @@ %13 = affine.min #map5(%arg5)[%4] %14 = dim %arg2, %c3 : memref<?x?x?x?xf32> %15 = subview %arg2[%arg3, %arg4, %arg5, 0] [%11, %12, %13, %14] [1, 1, 1, 1] : memref<?x?x?x?xf32> to memref<?x?x?x?xf32, #map3> - linalg.conv(%arg0, %9, %15) {__internal_linalg_transform__ = "workitem", dilations = [1, 1], strides = [1, 1]} : memref<?x?x?x?xf32>, memref<?x?x?x?xf32, #map3>, memref<?x?x?x?xf32, #map3> + linalg.conv(%arg0, %9, %15) {__internal_linalg_transform__ = "workgroup", dilations = [1, 1], strides = [1, 1]} : memref<?x?x?x?xf32>, memref<?x?x?x?xf32, #map3>, memref<?x?x?x?xf32, #map3> scf.yield } return
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/cyclic_to_workgroup.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/cyclic_to_workgroup.mlir index 110ac24..cac18ab 100644 --- a/iree/compiler/Conversion/LinalgToSPIRV/test/cyclic_to_workgroup.mlir +++ b/iree/compiler/Conversion/LinalgToSPIRV/test/cyclic_to_workgroup.mlir
@@ -27,7 +27,7 @@ %12 = dim %arg2, %c1 : memref<?x?xf32> %13 = affine.min #map0(%arg4)[%12] %14 = subview %arg2[%arg3, %arg4] [%11, %13] [1, 1] : memref<?x?xf32> to memref<?x?xf32, #map2> - linalg.matmul %5, %9, %14 {__internal_linalg_transform__ = "workitem"} : (memref<?x?xf32, #map2>, memref<?x?xf32, #map2>, memref<?x?xf32, #map2>) + linalg.matmul %5, %9, %14 {__internal_linalg_transform__ = "workgroup"} : (memref<?x?xf32, #map2>, memref<?x?xf32, #map2>, memref<?x?xf32, #map2>) } scf.yield }
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/linalg_tile_and_fuse.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/linalg_tile_and_fuse.mlir index 0e2fe6d..1728d35 100644 --- a/iree/compiler/Conversion/LinalgToSPIRV/test/linalg_tile_and_fuse.mlir +++ b/iree/compiler/Conversion/LinalgToSPIRV/test/linalg_tile_and_fuse.mlir
@@ -51,7 +51,7 @@ // CHECK: %[[VIEW2:.+]] = subview %[[ARG2]] // CHECK: linalg.conv // CHECK-SAME: %[[ARG0]], %[[VIEW1]], %[[VIEW2]] -// CHECK-SAME: "workitem" +// CHECK-SAME: "workgroup" // ----- @@ -81,7 +81,7 @@ // CHECK: %[[VIEW1:.+]] = subview %[[ARG1]] // CHECK: %[[VIEW2:.+]] = subview %[[ARG2]] // CHECK: linalg.matmul -// CHECK-SAME: "workitem" +// CHECK-SAME: "workgroup" // CHECK-SAME: %[[VIEW0]], %[[VIEW1]], %[[VIEW2]] // ----- @@ -111,4 +111,4 @@ // CHECK: %[[VIEW2:.+]] = subview %[[ARG2]] // CHECK: linalg.pooling_max // CHECK-SAME: %[[VIEW0]], %[[ARG1]], %[[VIEW2]] -// CHECK-SAME: "workitem" +// CHECK-SAME: "workgroup"
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/workgroup_memory_promotion.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/workgroup_memory_promotion.mlir index 76cfcb8..a24c77b 100644 --- a/iree/compiler/Conversion/LinalgToSPIRV/test/workgroup_memory_promotion.mlir +++ b/iree/compiler/Conversion/LinalgToSPIRV/test/workgroup_memory_promotion.mlir
@@ -36,12 +36,12 @@ // CHECK: %[[ALLOC2:.+]] = alloc(%[[C4]], %[[C8]]) : memref<?x?xf32, 3> // CHECK: %[[SUBVIEW2:.+]] = subview %[[ALLOC2]] // CHECK: linalg.copy(%[[ARG0SV]], %[[SUBVIEW1]]) -// CHECK-SAME: "workitem" +// CHECK-SAME: "workgroup" // CHECK: spv.ControlBarrier "Workgroup", "Workgroup", "AcquireRelease" // CHECK: linalg.copy(%[[ARG1SV]], %[[SUBVIEW2]]) -// CHECK-SAME: "workitem" +// CHECK-SAME: "workgroup" // CHECK: spv.ControlBarrier "Workgroup", "Workgroup", "AcquireRelease" -// CHECK: linalg.matmul {{.*}}"workitem"{{.*}} %[[SUBVIEW1]], %[[SUBVIEW2]], %[[RET0SV]] +// CHECK: linalg.matmul {{.*}}"workgroup"{{.*}} %[[SUBVIEW1]], %[[SUBVIEW2]], %[[RET0SV]] // CHECK: spv.ControlBarrier "Workgroup", "Workgroup", "AcquireRelease" // CHECK-DAG: dealloc %[[ALLOC1]] : memref<?x?xf32, 3> // CHECK-DAG: dealloc %[[ALLOC2]] : memref<?x?xf32, 3>
