1# Chapter 6: Lowering to LLVM and CodeGeneration
2
3[TOC]
4
5In the [previous chapter](Ch-5.md), we introduced the
6[dialect conversion](../../DialectConversion.md) framework and partially lowered
7many of the `Toy` operations to affine loop nests for optimization. In this
8chapter, we will finally lower to LLVM for code generation.
9
10## Lowering to LLVM
11
12For this lowering, we will again use the dialect conversion framework to perform
13the heavy lifting. However, this time, we will be performing a full conversion
14to the [LLVM dialect](../../Dialects/LLVM.md). Thankfully, we have already
15lowered all but one of the `toy` operations, with the last being `toy.print`.
16Before going over the conversion to LLVM, let's lower the `toy.print` operation.
17We will lower this operation to a non-affine loop nest that invokes `printf` for
18each element. Note that, because the dialect conversion framework supports
19[transitive lowering](../../../getting_started/Glossary.md/#transitive-lowering),
20we don't need to directly emit operations in the LLVM dialect. By transitive
21lowering, we mean that the conversion framework may apply multiple patterns to
22fully legalize an operation. In this example, we are generating a structured
23loop nest instead of the branch-form in the LLVM dialect. As long as we then
24have a lowering from the loop operations to LLVM, the lowering will still
25succeed.
26
27During lowering we can get, or build, the declaration for printf as so:
28
29```c++
30/// Return a symbol reference to the printf function, inserting it into the
31/// module if necessary.
32static FlatSymbolRefAttr getOrInsertPrintf(PatternRewriter &rewriter,
33                                           ModuleOp module,
34                                           LLVM::LLVMDialect *llvmDialect) {
35  auto *context = module.getContext();
36  if (module.lookupSymbol<LLVM::LLVMFuncOp>("printf"))
37    return SymbolRefAttr::get("printf", context);
38
39  // Create a function declaration for printf, the signature is:
40  //   * `i32 (i8*, ...)`
41  auto llvmI32Ty = IntegerType::get(context, 32);
42  auto llvmI8PtrTy =
43      LLVM::LLVMPointerType::get(IntegerType::get(context, 8));
44  auto llvmFnType = LLVM::LLVMFunctionType::get(llvmI32Ty, llvmI8PtrTy,
45                                                /*isVarArg=*/true);
46
47  // Insert the printf function into the body of the parent module.
48  PatternRewriter::InsertionGuard insertGuard(rewriter);
49  rewriter.setInsertionPointToStart(module.getBody());
50  rewriter.create<LLVM::LLVMFuncOp>(module.getLoc(), "printf", llvmFnType);
51  return SymbolRefAttr::get("printf", context);
52}
53```
54
55Now that the lowering for the printf operation has been defined, we can specify
56the components necessary for the lowering. These are largely the same as the
57components defined in the [previous chapter](Ch-5.md).
58
59### Conversion Target
60
61For this conversion, aside from the top-level module, we will be lowering
62everything to the LLVM dialect.
63
64```c++
65  mlir::ConversionTarget target(getContext());
66  target.addLegalDialect<mlir::LLVMDialect>();
67  target.addLegalOp<mlir::ModuleOp>();
68```
69
70### Type Converter
71
72This lowering will also transform the MemRef types which are currently being
73operated on into a representation in LLVM. To perform this conversion, we use a
74TypeConverter as part of the lowering. This converter specifies how one type
75maps to another. This is necessary now that we are performing more complicated
76lowerings involving block arguments. Given that we don't have any
77Toy-dialect-specific types that need to be lowered, the default converter is
78enough for our use case.
79
80```c++
81  LLVMTypeConverter typeConverter(&getContext());
82```
83
84### Conversion Patterns
85
86Now that the conversion target has been defined, we need to provide the patterns
87used for lowering. At this point in the compilation process, we have a
88combination of `toy`, `affine`, `arith`, and `std` operations. Luckily, the
89`affine`, `arith`, and `std` dialects already provide the set of patterns needed
90to transform them into LLVM dialect. These patterns allow for lowering the IR in
91multiple stages by relying on
92[transitive lowering](../../../getting_started/Glossary.md/#transitive-lowering).
93
94```c++
95  mlir::RewritePatternSet patterns(&getContext());
96  mlir::populateAffineToStdConversionPatterns(patterns, &getContext());
97  mlir::populateLoopToStdConversionPatterns(patterns, &getContext());
98  mlir::arith::populateArithmeticToLLVMConversionPatterns(typeConverter,
99                                                          patterns);
100  mlir::populateStdToLLVMConversionPatterns(typeConverter, patterns);
101
102  // The only remaining operation, to lower from the `toy` dialect, is the
103  // PrintOp.
104  patterns.add<PrintOpLowering>(&getContext());
105```
106
107### Full Lowering
108
109We want to completely lower to LLVM, so we use a `FullConversion`. This ensures
110that only legal operations will remain after the conversion.
111
112```c++
113  mlir::ModuleOp module = getOperation();
114  if (mlir::failed(mlir::applyFullConversion(module, target, patterns)))
115    signalPassFailure();
116```
117
118Looking back at our current working example:
119
120```mlir
121func @main() {
122  %0 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>
123  %2 = toy.transpose(%0 : tensor<2x3xf64>) to tensor<3x2xf64>
124  %3 = toy.mul %2, %2 : tensor<3x2xf64>
125  toy.print %3 : tensor<3x2xf64>
126  toy.return
127}
128```
129
130We can now lower down to the LLVM dialect, which produces the following code:
131
132```mlir
133llvm.func @free(!llvm<"i8*">)
134llvm.func @printf(!llvm<"i8*">, ...) -> i32
135llvm.func @malloc(i64) -> !llvm<"i8*">
136llvm.func @main() {
137  %0 = llvm.mlir.constant(1.000000e+00 : f64) : f64
138  %1 = llvm.mlir.constant(2.000000e+00 : f64) : f64
139
140  ...
141
142^bb16:
143  %221 = llvm.extractvalue %25[0 : index] : !llvm<"{ double*, i64, [2 x i64], [2 x i64] }">
144  %222 = llvm.mlir.constant(0 : index) : i64
145  %223 = llvm.mlir.constant(2 : index) : i64
146  %224 = llvm.mul %214, %223 : i64
147  %225 = llvm.add %222, %224 : i64
148  %226 = llvm.mlir.constant(1 : index) : i64
149  %227 = llvm.mul %219, %226 : i64
150  %228 = llvm.add %225, %227 : i64
151  %229 = llvm.getelementptr %221[%228] : (!llvm."double*">, i64) -> !llvm<"f64*">
152  %230 = llvm.load %229 : !llvm<"double*">
153  %231 = llvm.call @printf(%207, %230) : (!llvm<"i8*">, f64) -> i32
154  %232 = llvm.add %219, %218 : i64
155  llvm.br ^bb15(%232 : i64)
156
157  ...
158
159^bb18:
160  %235 = llvm.extractvalue %65[0 : index] : !llvm<"{ double*, i64, [2 x i64], [2 x i64] }">
161  %236 = llvm.bitcast %235 : !llvm<"double*"> to !llvm<"i8*">
162  llvm.call @free(%236) : (!llvm<"i8*">) -> ()
163  %237 = llvm.extractvalue %45[0 : index] : !llvm<"{ double*, i64, [2 x i64], [2 x i64] }">
164  %238 = llvm.bitcast %237 : !llvm<"double*"> to !llvm<"i8*">
165  llvm.call @free(%238) : (!llvm<"i8*">) -> ()
166  %239 = llvm.extractvalue %25[0 : index] : !llvm<"{ double*, i64, [2 x i64], [2 x i64] }">
167  %240 = llvm.bitcast %239 : !llvm<"double*"> to !llvm<"i8*">
168  llvm.call @free(%240) : (!llvm<"i8*">) -> ()
169  llvm.return
170}
171```
172
173See [Conversion to the LLVM IR Dialect](../../ConversionToLLVMDialect.md) for
174more in-depth details on lowering to the LLVM dialect.
175
176## CodeGen: Getting Out of MLIR
177
178At this point we are right at the cusp of code generation. We can generate code
179in the LLVM dialect, so now we just need to export to LLVM IR and setup a JIT to
180run it.
181
182### Emitting LLVM IR
183
184Now that our module is comprised only of operations in the LLVM dialect, we can
185export to LLVM IR. To do this programmatically, we can invoke the following
186utility:
187
188```c++
189  std::unique_ptr<llvm::Module> llvmModule = mlir::translateModuleToLLVMIR(module);
190  if (!llvmModule)
191    /* ... an error was encountered ... */
192```
193
194Exporting our module to LLVM IR generates:
195
196```llvm
197define void @main() {
198  ...
199
200102:
201  %103 = extractvalue { double*, i64, [2 x i64], [2 x i64] } %8, 0
202  %104 = mul i64 %96, 2
203  %105 = add i64 0, %104
204  %106 = mul i64 %100, 1
205  %107 = add i64 %105, %106
206  %108 = getelementptr double, double* %103, i64 %107
207  %109 = memref.load double, double* %108
208  %110 = call i32 (i8*, ...) @printf(i8* getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double %109)
209  %111 = add i64 %100, 1
210  br label %99
211
212  ...
213
214115:
215  %116 = extractvalue { double*, i64, [2 x i64], [2 x i64] } %24, 0
216  %117 = bitcast double* %116 to i8*
217  call void @free(i8* %117)
218  %118 = extractvalue { double*, i64, [2 x i64], [2 x i64] } %16, 0
219  %119 = bitcast double* %118 to i8*
220  call void @free(i8* %119)
221  %120 = extractvalue { double*, i64, [2 x i64], [2 x i64] } %8, 0
222  %121 = bitcast double* %120 to i8*
223  call void @free(i8* %121)
224  ret void
225}
226```
227
228If we enable optimization on the generated LLVM IR, we can trim this down quite
229a bit:
230
231```llvm
232define void @main()
233  %0 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 1.000000e+00)
234  %1 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 1.600000e+01)
235  %putchar = tail call i32 @putchar(i32 10)
236  %2 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 4.000000e+00)
237  %3 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 2.500000e+01)
238  %putchar.1 = tail call i32 @putchar(i32 10)
239  %4 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 9.000000e+00)
240  %5 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 3.600000e+01)
241  %putchar.2 = tail call i32 @putchar(i32 10)
242  ret void
243}
244```
245
246The full code listing for dumping LLVM IR can be found in
247`examples/toy/Ch6/toy.cpp` in the `dumpLLVMIR()` function:
248
249```c++
250
251int dumpLLVMIR(mlir::ModuleOp module) {
252  // Translate the module, that contains the LLVM dialect, to LLVM IR. Use a
253  // fresh LLVM IR context. (Note that LLVM is not thread-safe and any
254  // concurrent use of a context requires external locking.)
255  llvm::LLVMContext llvmContext;
256  auto llvmModule = mlir::translateModuleToLLVMIR(module, llvmContext);
257  if (!llvmModule) {
258    llvm::errs() << "Failed to emit LLVM IR\n";
259    return -1;
260  }
261
262  // Initialize LLVM targets.
263  llvm::InitializeNativeTarget();
264  llvm::InitializeNativeTargetAsmPrinter();
265  mlir::ExecutionEngine::setupTargetTriple(llvmModule.get());
266
267  /// Optionally run an optimization pipeline over the llvm module.
268  auto optPipeline = mlir::makeOptimizingTransformer(
269      /*optLevel=*/EnableOpt ? 3 : 0, /*sizeLevel=*/0,
270      /*targetMachine=*/nullptr);
271  if (auto err = optPipeline(llvmModule.get())) {
272    llvm::errs() << "Failed to optimize LLVM IR " << err << "\n";
273    return -1;
274  }
275  llvm::errs() << *llvmModule << "\n";
276  return 0;
277}
278```
279
280### Setting up a JIT
281
282Setting up a JIT to run the module containing the LLVM dialect can be done using
283the `mlir::ExecutionEngine` infrastructure. This is a utility wrapper around
284LLVM's JIT that accepts `.mlir` as input. The full code listing for setting up
285the JIT can be found in `Ch6/toyc.cpp` in the `runJit()` function:
286
287```c++
288int runJit(mlir::ModuleOp module) {
289  // Initialize LLVM targets.
290  llvm::InitializeNativeTarget();
291  llvm::InitializeNativeTargetAsmPrinter();
292
293  // An optimization pipeline to use within the execution engine.
294  auto optPipeline = mlir::makeOptimizingTransformer(
295      /*optLevel=*/EnableOpt ? 3 : 0, /*sizeLevel=*/0,
296      /*targetMachine=*/nullptr);
297
298  // Create an MLIR execution engine. The execution engine eagerly JIT-compiles
299  // the module.
300  auto maybeEngine = mlir::ExecutionEngine::create(module,
301      /*llvmModuleBuilder=*/nullptr, optPipeline);
302  assert(maybeEngine && "failed to construct an execution engine");
303  auto &engine = maybeEngine.get();
304
305  // Invoke the JIT-compiled function.
306  auto invocationResult = engine->invoke("main");
307  if (invocationResult) {
308    llvm::errs() << "JIT invocation failed\n";
309    return -1;
310  }
311
312  return 0;
313}
314```
315
316You can play around with it from the build directory:
317
318```shell
319$ echo 'def main() { print([[1, 2], [3, 4]]); }' | ./bin/toyc-ch6 -emit=jit
3201.000000 2.000000
3213.000000 4.000000
322```
323
324You can also play with `-emit=mlir`, `-emit=mlir-affine`, `-emit=mlir-llvm`, and
325`-emit=llvm` to compare the various levels of IR involved. Also try options like
326[`--print-ir-after-all`](../../PassManagement.md/#ir-printing) to track the
327evolution of the IR throughout the pipeline.
328
329The example code used throughout this section can be found in
330test/Examples/Toy/Ch6/llvm-lowering.mlir.
331
332So far, we have worked with primitive data types. In the
333[next chapter](Ch-7.md), we will add a composite `struct` type.
334