10: Building Your Own Dialect

Welcome to Chapter 10 of the "Implementing a language with MLIR" tutorial. So far, our AST has generated operations from MLIR's existing dialects directly. This has served us well: arith represents arithmetic, func represents functions, scf represents structured control flow, and memref provides storage for mutable variables. In this chapter, you'll create a small custom dialect called Kaleidoscope. We'll add a handful of operations, and they all solve one specific problem.

10.1 The Problem We're Solving

Let's start with a simple function:

def incrementAndSquare(x)
  var y = x in
    (y = y + 1) * y;

Up to now, when our compiler saw x and y, it immediately turned them into generic memory allocations. We lost the x and y. In Chapter 9, we worked around this by saving the parameter names separately and stitching them back together later. That worked, but it was a patch. Ideally, we wouldn't lose the information in the first place. We can fix this with a custom operation for which we need a custom dialect.

Here's what we are aiming to emit in MLIR.

func.func @incrementAndSquare(%arg0: f64) -> f64 {
  %cst = arith.constant 1.000000e+00 : f64
  %0 = kaleidoscope.var "x"(%arg0 : f64) {argumentNumber = 1 : i64} : !kaleidoscope.variable
  %1 = kaleidoscope.read(%0 : !kaleidoscope.variable) : f64
  %2 = kaleidoscope.var "y"(%1 : f64) {argumentNumber = 0 : i64} : !kaleidoscope.variable
  %3 = kaleidoscope.read(%2 : !kaleidoscope.variable) : f64
  %4 = arith.addf %3, %cst : f64
  kaleidoscope.assign(%4 : f64, %2 : !kaleidoscope.variable)
  %5 = kaleidoscope.read(%2 : !kaleidoscope.variable) : f64
  %6 = arith.mulf %4, %5 : f64
  return %6 : f64
}

We'll invent 3 kaleidoscope.* operations.

kaleidoscope.var: Declares a variable and retains its name. You can see how the "x" and "y" are now part of the IR above.
kaleidoscope.read: Reads a variable.
kaleidoscope.assign: Assigns a value to a variable.

Notice that everything else still uses the standard dialects you already know: arith for math, func for functions.

10.2 Round-tripping MLIR

As we'll see in more concrete terms soon, MLIR's textual assembly format is meant to round-trip: print an in-memory IR to text and parse that text back into IR to reconstruct the operations. This way we can save serialize the IR to files and read them back in later if needed. Each custom operation below declares an assemblyFormat in TableGen. MLIR uses it to generate the operation's printer and parser, so the syntax you see in --dump-mlir can be serialized and read back in later. The VariableType definition supplies the text representation !kaleidoscope.variable, and useDefaultTypePrinterParser generates a printer and parser for the dialect.

10.3 Define the Dialect

MLIR dialects are usually defined with TableGen, a compact language for describing types and operations.

I wrote a more relaxed tutorial on TableGen which you can peruse later. For now, just truck on with the examples to build some intuition on why TableGen is important in the first place.

Create a file called KaleidoscopeOps.td and start with the dialect itself:

KaleidoscopeOps.td

// Provides the TableGen definitions for MLIR dialects, types, and operations.
include "mlir/IR/OpBase.td"

def Kaleidoscope_Dialect : Dialect {
  // The prefix for every operation and type in this dialect. An operation
  // with the mnemonic `var` will therefore print as `kaleidoscope.var`.
  let name = "kaleidoscope";

  // Where the generated C++ classes live. Our classes end up in the
  // `mlir::kaleidoscope` namespace.
  let cppNamespace = "::mlir::kaleidoscope";

  let summary = "Operations that preserve Kaleidoscope variable semantics";

  // Ask MLIR to generate the parser and printer for our types based on the
  // assembly format each type declares. We only have one type, and its format
  // is defined below.
  let useDefaultTypePrinterParser = 1;
}

The name gives our operations their kaleidoscope. prefix. The cppNamespace places the generated classes in mlir::kaleidoscope, and useDefaultTypePrinterParser gives our types the usual MLIR parsing and printing support.

10.4 Define the Variable Type

We will build the dialect one piece at a time. For each piece, we will look at the TableGen definition, the C++ that TableGen generates, the C++ that we write to use it, and the MLIR that eventually gets printed.

Start with the type:

KaleidoscopeOps.td

// A type representing a mutable source variable.
//
// TypeDef generates the C++ class `mlir::kaleidoscope::VariableType` from the
// class stem "Variable", and the mnemonic "variable" gives it the textual
// spelling `!kaleidoscope.variable`.
//
// The type deliberately says nothing about how the variable is stored. That
// is a lowering decision, not a property of the source language.
def Kaleidoscope_VariableType
    : TypeDef<Kaleidoscope_Dialect, "Variable"> {
  let mnemonic = "variable";
  let summary = "a mutable Kaleidoscope variable";

  // An empty assembly format means the type has no parameters to print after
  // its mnemonic, so it always appears as the bare `!kaleidoscope.variable`.
  let assemblyFormat = "";
}

From that small definition, mlir-tblgen generates the real VariableType class below. This is an excerpt from KaleidoscopeTypes.h.inc; the generated file also contains parser, printer, and registration machinery.

To see the C++ declaration generated for our type, run:

/path/to/llvm-project/build/bin/mlir-tblgen \
  -I /path/to/llvm-project/mlir/include \
  -gen-typedef-decls \
  KaleidoscopeOps.td

If LLVM was installed with Homebrew instead of built from source, run:

"$(brew --prefix llvm)/bin/mlir-tblgen" \
  -I "$(brew --prefix llvm)/include" \
  -gen-typedef-decls \
  KaleidoscopeOps.td

And hopefully, you can see this printed to the terminal.

#ifdef GET_TYPEDEF_CLASSES
#undef GET_TYPEDEF_CLASSES

namespace mlir {
class AsmParser;
class AsmPrinter;
} // namespace mlir
namespace mlir {
namespace kaleidoscope {
/// a mutable Kaleidoscope variable
class VariableType;
class VariableType : public ::mlir::Type::TypeBase<VariableType, ::mlir::Type, ::mlir::TypeStorage> {
public:
  using Base::Base;
  static constexpr ::llvm::StringLiteral name = "kaleidoscope.variable";
  static constexpr ::llvm::StringLiteral dialectName = "kaleidoscope";
  static constexpr ::llvm::StringLiteral getMnemonic() {
    return {"variable"};
  }

  static ::mlir::Type parse(::mlir::AsmParser &odsParser);
  void print(::mlir::AsmPrinter &odsPrinter) const;
};
} // namespace kaleidoscope
} // namespace mlir
MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::kaleidoscope::VariableType)

#endif  // GET_TYPEDEF_CLASSES
};

TypeBase gives us the normal MLIR type behavior. Types are uniqued inside an MLIRContext, so asking the same context for VariableType twice gives us handles to the same underlying type storage. Our compiler asks for the type like this:

auto variableType =
    kaleidoscope::VariableType::get(TheContext.get());

When MLIR prints that type, it appears as:

!kaleidoscope.variable

The type says only, "this SSA value represents a Kaleidoscope variable." It does not decide how the variable is stored. We make that decision later during lowering.

10.5 Add an Operation Base Class

All three operations in this chapter share the same dialect and basic setup. This small TableGen class saves us from repeating it:

KaleidoscopeOps.td

// A common base for every operation in this dialect. Each concrete operation
// supplies its own mnemonic and, optionally, a list of traits describing its
// behavior.
class Kaleidoscope_Op<string mnemonic, list<Trait> traits = []>
    : Op<Kaleidoscope_Dialect, mnemonic, traits>;

We can now define kaleidoscope.var, kaleidoscope.read, and kaleidoscope.assign one at a time.

10.6 Define `kaleidoscope.var`

The source example begins with a function parameter named x:

def incrementAndSquare(x)

Here is the TableGen definition for the operation that preserves that variable:

KaleidoscopeOps.td

// Declares and initializes a mutable source variable.
//
// This generates the C++ class `kaleidoscope::DeclareOp`, printed in IR as
// `kaleidoscope.var`. It is the operation that carries the source name and
// argument number that would otherwise be lost the moment a variable became
// an anonymous allocation.
def Kaleidoscope_DeclareOp : Kaleidoscope_Op<"var", []> {
  let summary = "declare and initialize a mutable source variable";

  // `(ins ...)` lists everything the operation takes in. The `$` names
  // generate C++ accessors (getInitialValue, getName, getArgumentNumber).
  //
  //   - F64:$initialValue   is an SSA operand, constrained to f64
  //   - StrAttr:$name       is an attribute holding the source name
  //   - I64Attr:$argumentNumber is an attribute; 0 means a local variable,
  //                         and 1+ means a function parameter (DWARF numbers
  //                         parameters starting at 1)
  let arguments = (ins F64:$initialValue, StrAttr:$name,
                       I64Attr:$argumentNumber);

  // The SSA result represents the *variable itself*, not the floating-point
  // value currently stored in it. Subsequent read and assign operations use
  // this value to refer to the variable.
  let results = (outs Kaleidoscope_VariableType:$variable);

  // The textual syntax. Backticks contain literal punctuation, `$name` and
  // `$initialValue` refer to the fields above, and `type(...)` prints the type
  // of the named operand or result. `attr-dict` prints any attributes not
  // already consumed by the format.
  //
  // This produces, for example:
  //   %0 = kaleidoscope.var "x"(%arg0 : f64)
  //       {argumentNumber = 1 : i64} : !kaleidoscope.variable
  let assemblyFormat = "$name `(` $initialValue `:` type($initialValue) `)` attr-dict `:` type($variable)";
}

MLIR conventions in the IR (convention, not rule) choose to put the type after the variable name like so: argumentNumber = 1 : i64. TableGen conventions do the opposite and use this I64Attr:$argumentNumber. Not a deal breaker. Just something you should be aware of and get annoyed with from time to time.

The operation takes an initial f64 value, a source name, and an argument number. It defines one result of our new variable type. The $ names generate accessors such as getInitialValue(), getName(), getArgumentNumber(), and getVariable().

To print the generated operation declarations in the terminal, run:

/path/to/llvm-project/build/bin/mlir-tblgen \
  -I /path/to/llvm-project/mlir/include \
  -gen-op-decls \
  KaleidoscopeOps.td

To write the generated declarations into KaleidoscopeOps.h.inc, run:

/path/to/llvm-project/build/bin/mlir-tblgen \
  -I /path/to/llvm-project/mlir/include \
  -gen-op-decls \
  -o KaleidoscopeOps.h.inc \
  KaleidoscopeOps.td

If LLVM was installed with Homebrew instead of built from source, run:

"$(brew --prefix llvm)/bin/mlir-tblgen" \
  -I "$(brew --prefix llvm)/include" \
  -gen-op-decls \
  KaleidoscopeOps.td

Here is the relevant part of the actual DeclareOp class generated in KaleidoscopeOps.h.inc. The full class contains additional builder overloads and MLIR bookkeeping.

class DeclareOp : public ::mlir::Op<
    DeclareOp,
    ::mlir::OpTrait::OneResult,
    ::mlir::OpTrait::OneTypedResult<VariableType>::Impl,
    ::mlir::OpTrait::OneOperand> {
public:
  static constexpr ::llvm::StringLiteral getOperationName() {
    return "kaleidoscope.var";
  }

  ::mlir::TypedValue<::mlir::FloatType> getInitialValue();
  ::llvm::StringRef getName();
  uint64_t getArgumentNumber();

  ::mlir::TypedValue<VariableType> getVariable();

  static void build(::mlir::OpBuilder &builder,
                    ::mlir::OperationState &state,
                    ::mlir::Type variable,
                    ::mlir::Value initialValue,
                    ::mlir::StringAttr name,
                    ::mlir::IntegerAttr argumentNumber);
};

The parser has already turned our source into an AST by the time we use this class. During AST code generation, this helper asks the builder to create a DeclareOp:

toy.cpp

static Value CreateVariable(StringRef Name, Value InitialValue,
                            int64_t ArgumentNumber = 0) {
  return TheBuilder->create<kaleidoscope::DeclareOp>(
      getLocation(), kaleidoscope::VariableType::get(TheContext.get()),
      InitialValue, TheBuilder->getStringAttr(Name),
      TheBuilder->getI64IntegerAttr(ArgumentNumber));
}

Function parameters call that helper with one-based argument numbers:

toy.cpp

func::FuncOp FunctionAST::codegen() {
  ...
  unsigned Index = 0;
  for (BlockArgument Argument : TheFunction.getArguments()) {
    StringRef Name = P.getArgs()[Index];
    Value Storage = CreateVariable(Name, Argument, Index + 1);
    NamedValues[Name.str()] = Storage;
    ++Index;
  }
  ...
}

Locals and loop variables use the helper's default argument number of 0.

The call to create<DeclareOp> uses the generated build method, inserts the operation into the current MLIR block, and creates its result at the same time. MLIR prints the result as %0 here:

%0 = kaleidoscope.var "x"(%arg0 : f64) {argumentNumber = 1 : i64} : !kaleidoscope.variable

%arg0 : f64 is the initial value received by the function.
"x" is the name from the source.
argumentNumber = 1 : i64 records that x is the first function parameter. Locals use 0.
%0 is the SSA value defined by DeclareOp. Its type is !kaleidoscope.variable.

Other operations use %0 to say exactly which variable they mean. When this dialect is lowered, %0 will be replaced by the address of the storage allocated for x.

10.7 Define `kaleidoscope.read`

The read operation accepts a variable and defines an f64 result containing its current value:

KaleidoscopeOps.td

def Kaleidoscope_ReadOp : Kaleidoscope_Op<"read", []> {
  let summary = "read a mutable source variable";

  // Consumes the variable; produces the f64 value currently stored in it.
  let arguments = (ins Kaleidoscope_VariableType:$variable);
  let results = (outs F64:$value);
  let assemblyFormat = "`(` $variable `:` type($variable) `)` attr-dict `:` type($value)";
}

TableGen generates a ReadOp with typed accessors for both sides of that relationship:

class ReadOp : public ::mlir::Op<
    ReadOp,
    ::mlir::OpTrait::OneResult,
    ::mlir::OpTrait::OneTypedResult<::mlir::FloatType>::Impl,
    ::mlir::OpTrait::OneOperand> {
public:
  static constexpr ::llvm::StringLiteral getOperationName() {
    return "kaleidoscope.read";
  }

  ::mlir::TypedValue<VariableType> getVariable();
  ::mlir::TypedValue<::mlir::FloatType> getValue();

  static void build(::mlir::OpBuilder &builder,
                    ::mlir::OperationState &state,
                    ::mlir::Type value,
                    ::mlir::Value variable);
};

When a VariableExprAST needs the current value of x, its codegen() method creates that operation:

toy.cpp

Value VariableExprAST::codegen() {
  ...
  return TheBuilder->create<kaleidoscope::ReadOp>(
      getLocation(), TheBuilder->getF64Type(), It->second);
}

The resulting MLIR makes both the input and output types explicit:

%1 = kaleidoscope.read(%0 : !kaleidoscope.variable) : f64

This says, "Read the current value from the variable %0 and define the f64 result %1."

ReadOp is not marked Pure. Two reads of the same variable might return different values if an assignment occurs between them. Marking it pure would let common-subexpression elimination incorrectly merge those reads.

10.8 Define `kaleidoscope.assign`

The assign operation accepts a new f64 value and the variable that should receive it:

KaleidoscopeOps.td

def Kaleidoscope_AssignOp : Kaleidoscope_Op<"assign", []> {
  let summary = "assign a new value to a mutable source variable";

  // Consumes the variable and the replacement value. Print both operand types
  // explicitly so it is immediately clear that the value is an f64 and the
  // destination is a !kaleidoscope.variable.
  let arguments = (ins Kaleidoscope_VariableType:$variable, F64:$value);
  let assemblyFormat = "`(` $value `:` type($value) `,` $variable `:` type($variable) `)` attr-dict";
}

Its generated class has two operands and no result:

class AssignOp : public ::mlir::Op<
    AssignOp,
    ::mlir::OpTrait::ZeroResults,
    ::mlir::OpTrait::NOperands<2>::Impl> {
public:
  static constexpr ::llvm::StringLiteral getOperationName() {
    return "kaleidoscope.assign";
  }

  ::mlir::TypedValue<VariableType> getVariable();
  ::mlir::TypedValue<::mlir::FloatType> getValue();

  static void build(::mlir::OpBuilder &builder,
                    ::mlir::OperationState &state,
                    ::mlir::Value variable,
                    ::mlir::Value value);
};

Assignment code generation creates the operation for its side effect and returns AssignedValue separately. That preserves Kaleidoscope's rule that an assignment is also an expression: (y = y + 1) evaluates to the newly assigned value.

toy.cpp

Value BinaryExprAST::codegen() {
  ...
  if (Op == '=') {
    ...
    TheBuilder->create<kaleidoscope::AssignOp>(getLocation(), It->second,
                                               AssignedValue);
    return AssignedValue;
  }
  ...
}

Because AssignOp defines no result, there is no %result = and no result type at the end:

kaleidoscope.assign(%4 : f64, %2 : !kaleidoscope.variable)

Putting the three operations together gives us a complete variable story:

%0 = kaleidoscope.var "x"(%arg0 : f64) {argumentNumber = 1 : i64} : !kaleidoscope.variable
%1 = kaleidoscope.read(%0 : !kaleidoscope.variable) : f64
%2 = arith.addf %1, %cst : f64
kaleidoscope.assign(%2 : f64, %0 : !kaleidoscope.variable)
%3 = kaleidoscope.read(%0 : !kaleidoscope.variable) : f64

DeclareOp defines %0. The two reads and the assignment use %0 to identify the same variable. %1, %2, and %3 are ordinary f64 values.

10.9 Generate and Include the C++ Classes

Add these rules to CMakeLists.txt to generate the dialect, type, and operation declarations and definitions:

CMakeLists.txt

# Every mlir_tablegen invocation below reads this TableGen source file.
set(LLVM_TARGET_DEFINITIONS KaleidoscopeOps.td)

# Generate the KaleidoscopeDialect class declaration and definition from the
# `def Kaleidoscope_Dialect : Dialect` record.
mlir_tablegen(KaleidoscopeDialect.h.inc -gen-dialect-decls
  -dialect=kaleidoscope)
mlir_tablegen(KaleidoscopeDialect.cpp.inc -gen-dialect-defs
  -dialect=kaleidoscope)

# Generate C++ declarations and definitions for our TypeDef records.
mlir_tablegen(KaleidoscopeTypes.h.inc -gen-typedef-decls
  -typedefs-dialect=kaleidoscope)
mlir_tablegen(KaleidoscopeTypes.cpp.inc -gen-typedef-defs
  -typedefs-dialect=kaleidoscope)

# Generate C++ declarations and definitions for our Op records.
mlir_tablegen(KaleidoscopeOps.h.inc -gen-op-decls)
mlir_tablegen(KaleidoscopeOps.cpp.inc -gen-op-defs)

# Give the generated fragments one build target so the executable can depend
# on all of them before compiling files that #include those fragments.
add_custom_target(KaleidoscopeOpsIncGen DEPENDS
  KaleidoscopeDialect.h.inc
  KaleidoscopeDialect.cpp.inc
  KaleidoscopeTypes.h.inc
  KaleidoscopeTypes.cpp.inc
  KaleidoscopeOps.h.inc
  KaleidoscopeOps.cpp.inc
)

The add_custom_target groups the six generated files under one target name. The generated .inc files live in the build directory, so add that directory to the include path and make the executable depend on the generated target:

CMakeLists.txt

target_include_directories(toy PRIVATE
  ${CMAKE_CURRENT_SOURCE_DIR}
  ${CMAKE_CURRENT_BINARY_DIR}
  ${LLVM_INCLUDE_DIRS}
  ${MLIR_INCLUDE_DIRS}
)

CMakeLists.txt

add_dependencies(toy KaleidoscopeOpsIncGen)

Include the generated declarations in the dialect header:

KaleidoscopeDialect.h

// Generated declaration of mlir::kaleidoscope::KaleidoscopeDialect.
#include "KaleidoscopeDialect.h.inc"

// KaleidoscopeTypes.h.inc contains several selectable sections. Defining this
// macro asks it to emit the generated type class declarations at this include.
#define GET_TYPEDEF_CLASSES
#include "KaleidoscopeTypes.h.inc"

// Likewise, select the generated operation class declarations from the
// operation header fragment.
#define GET_OP_CLASSES
#include "KaleidoscopeOps.h.inc"

Include the generated definitions in the implementation file:

KaleidoscopeDialect.cpp

// Generated definitions for the dialect class declared by the matching header
// fragment.
#include "KaleidoscopeDialect.cpp.inc"

// Select the generated C++ definitions for our TypeDef and Op records.
#define GET_TYPEDEF_CLASSES
#include "KaleidoscopeTypes.cpp.inc"

#define GET_OP_CLASSES
#include "KaleidoscopeOps.cpp.inc"

The selector macros tell each generated .inc file which section to provide at that include location.

Finally, register the generated type and operations when the dialect initializes:

KaleidoscopeDialect.cpp

void KaleidoscopeDialect::initialize() {
  // The same generated .inc files also contain lists of every type and
  // operation in the dialect. These macros select those lists so the dialect
  // can register all generated classes with MLIR.
  addTypes<
#define GET_TYPEDEF_LIST
#include "KaleidoscopeTypes.cpp.inc"
      >();
  addOperations<
#define GET_OP_LIST
#include "KaleidoscopeOps.cpp.inc"
      >();
}

Once TheContext->loadDialect<kaleidoscope::KaleidoscopeDialect>() runs, MLIR knows how to create, parse, print, and verify everything in the dialect.

We can also delete the FunctionParameters map from Chapter 9. Each declaration now carries its own name, location, and argument number, so there is nothing to reconstruct later.

10.10 Lower the Dialect to the LLVM Dialect

When lowering time comes, we convert each kaleidoscope.* operation into operations in MLIR's LLVM dialect. Note carefully: the LLVM dialect is still MLIR. It is not LLVM IR. Translation to actual LLVM IR happens later in the pipeline, after all our lowering passes have run. The MLIR LLVM dialect is what we lower to here; the translation step is a separate, final stage.

A declaration becomes an allocation plus an initializing store:

toy.cpp

Value One = LLVM::ConstantOp::create(Rewriter, Loc, Rewriter.getI64Type(),
                                     Rewriter.getI64IntegerAttr(1));
Value Address =
    LLVM::AllocaOp::create(Rewriter, Loc, PointerType, DoubleType, One, 0);
LLVM::StoreOp::create(Rewriter, Loc, Adaptor.getInitialValue(), Address);

We still have the source variable's name, location, and argument number in hand, and we've just produced its final stack address. That means we can create the debug declaration right here, without any reconstruction:

toy.cpp

auto Variable = LLVM::DILocalVariableAttr::get(
    Scope, Op.getName(), Scope.getFile(), Line, Op.getArgumentNumber(),
    /*alignInBits=*/0, VariableType, LLVM::DIFlags::Zero);
LLVM::DbgDeclareOp::create(
    Rewriter, Loc, Address, Variable,
    LLVM::DIExpressionAttr::get(Rewriter.getContext()));

Reads and assignments become simple loads and stores:

toy.cpp

Rewriter.replaceOpWithNewOp<LLVM::LoadOp>(Op, Rewriter.getF64Type(),
                                          Adaptor.getVariable());

toy.cpp

Rewriter.replaceOpWithNewOp<LLVM::StoreOp>(Op, Adaptor.getValue(),
                                           Adaptor.getVariable());

Finally, tell MLIR that leaving a kaleidoscope.* operation unconverted is an error:

toy.cpp

void runOnOperation() override {
  ...
  ConversionTarget Target(Context);
  Target.addIllegalDialect<kaleidoscope::KaleidoscopeDialect>();
  Target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });

  RewritePatternSet Patterns(&Context);
  Patterns.add<DeclareOpLowering, ReadOpLowering, AssignOpLowering>(Converter,
                                                                    &Context);
  if (failed(applyPartialConversion(getOperation(), Target,
                                    std::move(Patterns))))
    signalPassFailure();
}

Without addIllegalDialect, a buggy pattern could silently skip an operation and pass IR with dangling kaleidoscope ops to the LLVM dialect translation step. Declaring the dialect illegal makes that failure loud.

10.11 Wire Up the Pass Pipeline

The order matters here. First, lower the standard high-level dialects down to the LLVM dialect. Then run our Chapter 9 debug pass to create the compile unit and function scopes:

toy.cpp

lowerToLLVM(const llvm::DataLayout &DataLayout) {
  ...
  PassManager DebugPM(TheContext.get());
  DebugPM.addPass(
      createKaleidoscopeDebugInfoPass(InputFilename.getValue(), OptLevel));
  ...
}

With function scopes now available, we can run the variable-lowering pass, which creates both the stack storage and the variable debug declarations together:

toy.cpp

lowerToLLVM(const llvm::DataLayout &DataLayout) {
  ...
  DebugPM.addPass(std::make_unique<LowerKaleidoscopeVariablesPass>());
  ...
}

Then let MLIR's standard pass fill in the remaining scopes on the lowered operations:

toy.cpp

lowerToLLVM(const llvm::DataLayout &DataLayout) {
  ...
  LLVM::DIScopeForLLVMFuncOpPassOptions DebugOptions;
  DebugOptions.emissionKind = LLVM::DIEmissionKind::Full;
  DebugPM.addPass(
      LLVM::createDIScopeForLLVMFuncOpPass(std::move(DebugOptions)));
  ...
}

After all MLIR passes have finished, the module is in the LLVM dialect. From there, the standard translation step turns it into LLVM IR:

%x = alloca double, i64 1, align 8
store double %arg0, ptr %x, align 8
#dbg_declare(ptr %x, !variable, !DIExpression(), !location)

Once we're at LLVM IR, everything downstream, including JIT, object emission, and DWARF generation, proceeds unchanged.

10.12 Try It Out

Configure and build the example first, setting MLIR_DIR to the directory containing MLIRConfig.cmake in your LLVM build:

cmake -S . -B build \
  -DMLIR_DIR=/path/to/llvm-project/build/lib/cmake/mlir
cmake --build build

Now run the interpreter:

$ ./build/toy
ready> def incrementAndSquare(x) var y = x in (y = y + 1) * y;
ready> incrementAndSquare(4);
Evaluated to 25.000000

Now dump the MLIR before lowering to see your named variables:

$ ./build/toy --dump-mlir
ready> def incrementAndSquare(x) var y = x in (y = y + 1) * y;
Read function definition:
func.func @incrementAndSquare(%arg0: f64) -> f64 {
  %cst = arith.constant 1.000000e+00 : f64
  %0 = kaleidoscope.var "x"(%arg0 : f64) {argumentNumber = 1 : i64} : !kaleidoscope.variable
  %1 = kaleidoscope.read(%0 : !kaleidoscope.variable) : f64
  %2 = kaleidoscope.var "y"(%1 : f64) {argumentNumber = 0 : i64} : !kaleidoscope.variable
  %3 = kaleidoscope.read(%2 : !kaleidoscope.variable) : f64
  %4 = arith.addf %3, %cst : f64
  kaleidoscope.assign(%4 : f64, %2 : !kaleidoscope.variable)
  %5 = kaleidoscope.read(%2 : !kaleidoscope.variable) : f64
  %6 = arith.mulf %4, %5 : f64
  return %6 : f64
}

And the LLVM IR after lowering, where variables have become stack allocations with debug declarations attached:

$ ./build/toy --dump-llvm-ir
ready> def incrementAndSquare(x) var y = x in (y = y + 1) * y;
; ModuleID = 'LLVMDialectModule'
source_filename = "LLVMDialectModule"
target datalayout = "e-m:o-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-n32:64-S128-Fn32"

define double @incrementAndSquare(double %0) !dbg !3 {
  %2 = alloca double, i64 1, align 8, !dbg !6
  store double %0, ptr %2, align 8, !dbg !6
    #dbg_declare(ptr %2, !7, !DIExpression(), !6)
  %3 = load double, ptr %2, align 8, !dbg !9
  %4 = alloca double, i64 1, align 8, !dbg !10
  store double %3, ptr %4, align 8, !dbg !10
    #dbg_declare(ptr %4, !11, !DIExpression(), !10)
  %5 = load double, ptr %4, align 8, !dbg !12
  %6 = fadd double %5, 1.000000e+00, !dbg !13
  store double %6, ptr %4, align 8, !dbg !14
  %7 = load double, ptr %4, align 8, !dbg !15
  %8 = fmul double %6, %7, !dbg !16
  ret double %8, !dbg !6
}

; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare void @llvm.dbg.declare(metadata, metadata, metadata) #0

attributes #0 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }

!llvm.dbg.cu = !{!0}
!llvm.module.flags = !{!2}

!0 = distinct !DICompileUnit(language: DW_LANG_C, file: !1, producer: "Kaleidoscope", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug)
!1 = !DIFile(filename: "<stdin>", directory: "")
!2 = !{i32 2, !"Debug Info Version", i32 3}
!3 = distinct !DISubprogram(name: "incrementAndSquare", linkageName: "incrementAndSquare", scope: !1, file: !1, line: 1, type: !4, scopeLine: 1, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
!4 = !DISubroutineType(cc: DW_CC_normal, types: !5)
!5 = !{}
!6 = !DILocation(line: 1, column: 5, scope: !3)
!7 = !DILocalVariable(name: "x", arg: 1, scope: !3, file: !1, line: 1, type: !8)
!8 = !DIBasicType(name: "double", size: 64, encoding: DW_ATE_float)
!9 = !DILocation(line: 1, column: 35, scope: !3)
!10 = !DILocation(line: 1, column: 27, scope: !3)
!11 = !DILocalVariable(name: "y", scope: !3, file: !1, line: 1, type: !8)
!12 = !DILocation(line: 1, column: 45, scope: !3)
!13 = !DILocation(line: 1, column: 47, scope: !3)
!14 = !DILocation(line: 1, column: 43, scope: !3)
!15 = !DILocation(line: 1, column: 54, scope: !3)
!16 = !DILocation(line: 1, column: 52, scope: !3)

10.13 Where the Code Lives

The full implementation is spread across these files:

toy.cpp: the compiler and the variable-lowering pass
KaleidoscopeOps.td: the variable type and operation definitions
KaleidoscopeDialect.h / KaleidoscopeDialect.cpp: connects generated classes to the compiler
KaleidoscopeDebugInfo.h / KaleidoscopeDebugInfo.cpp: creates compile-unit and function scopes
CMakeLists.txt: runs TableGen and builds the executable

10.13.1 Build Configuration

CMakeLists.txt

# CMakeLists.txt

cmake_minimum_required(VERSION 3.20)

project(kaleidoscope-chapter-10 LANGUAGES C CXX)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED YES)
set(CMAKE_CXX_EXTENSIONS NO)

find_package(MLIR REQUIRED CONFIG)

list(APPEND CMAKE_MODULE_PATH "${MLIR_CMAKE_DIR}")
include(TableGen)
include(AddMLIR)
include_directories(${MLIR_INCLUDE_DIRS})

# Every mlir_tablegen invocation below reads this TableGen source file.
set(LLVM_TARGET_DEFINITIONS KaleidoscopeOps.td)

# Generate the KaleidoscopeDialect class declaration and definition from the
# `def Kaleidoscope_Dialect : Dialect` record.
mlir_tablegen(KaleidoscopeDialect.h.inc -gen-dialect-decls
  -dialect=kaleidoscope)
mlir_tablegen(KaleidoscopeDialect.cpp.inc -gen-dialect-defs
  -dialect=kaleidoscope)

# Generate C++ declarations and definitions for our TypeDef records.
mlir_tablegen(KaleidoscopeTypes.h.inc -gen-typedef-decls
  -typedefs-dialect=kaleidoscope)
mlir_tablegen(KaleidoscopeTypes.cpp.inc -gen-typedef-defs
  -typedefs-dialect=kaleidoscope)

# Generate C++ declarations and definitions for our Op records.
mlir_tablegen(KaleidoscopeOps.h.inc -gen-op-decls)
mlir_tablegen(KaleidoscopeOps.cpp.inc -gen-op-defs)

# Give the generated fragments one build target so the executable can depend
# on all of them before compiling files that #include those fragments.
add_custom_target(KaleidoscopeOpsIncGen DEPENDS
  KaleidoscopeDialect.h.inc
  KaleidoscopeDialect.cpp.inc
  KaleidoscopeTypes.h.inc
  KaleidoscopeTypes.cpp.inc
  KaleidoscopeOps.h.inc
  KaleidoscopeOps.cpp.inc
)

add_executable(toy
  toy.cpp
  KaleidoscopeDialect.cpp
  KaleidoscopeDebugInfo.cpp
)
add_dependencies(toy KaleidoscopeOpsIncGen)

# Make symbols in the executable available to the JIT for runtime lookup.
set_target_properties(toy PROPERTIES ENABLE_EXPORTS ON)

target_include_directories(toy PRIVATE
  ${CMAKE_CURRENT_SOURCE_DIR}
  ${CMAKE_CURRENT_BINARY_DIR}
  ${LLVM_INCLUDE_DIRS}
  ${MLIR_INCLUDE_DIRS}
)

target_compile_definitions(toy PRIVATE ${LLVM_DEFINITIONS})

if(NOT LLVM_ENABLE_RTTI)
  target_compile_options(toy PRIVATE -fno-rtti)
endif()

llvm_map_components_to_libnames(LLVM_LIBS
  Core
  CodeGen
  OrcJIT
  Support
  Target
  ${LLVM_TARGETS_TO_BUILD}
)

target_link_libraries(toy PRIVATE
  MLIRArithToLLVM
  MLIRArithDialect
  MLIRBuiltinToLLVMIRTranslation
  MLIRControlFlowDialect
  MLIRControlFlowToLLVM
  MLIRFuncToLLVM
  MLIRFuncDialect
  MLIRLLVMDialect
  MLIRLLVMIRTransforms
  MLIRLLVMToLLVMIRTranslation
  MLIRMemRefDialect
  MLIRMemRefToLLVM
  MLIRReconcileUnrealizedCasts
  MLIRSCFDialect
  MLIRSCFToControlFlow
  MLIRTargetLLVMIRExport
  MLIRTransforms
  ${LLVM_LIBS}
)

10.13.2 Compiler

toy.cpp

// toy.cpp

#include "../include/KaleidoscopeJIT.h"
#include "KaleidoscopeDebugInfo.h"
#include "KaleidoscopeDialect.h"
#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
#include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/Transforms/Passes.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/Verifier.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Export.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/Passes.h"
#include "llvm/ADT/SmallString.h"
#include "llvm/BinaryFormat/Dwarf.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Module.h"
#include "llvm/MC/TargetRegistry.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/Path.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Target/TargetOptions.h"
#include "llvm/TargetParser/Host.h"
#include <cassert>
#include <cctype>
#include <cstdio>
#include <cstdlib>
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <system_error>
#include <utility>
#include <vector>

using namespace mlir;

//===----------------------------------------------------------------------===//
// Lexer
//===----------------------------------------------------------------------===//

// The lexer returns tokens [0-255] if it is an unknown character, otherwise one
// of these for known things.
enum Token {
  tok_eof = -1,

  // commands
  tok_def = -2,
  tok_extern = -3,

  // primary
  tok_identifier = -4,
  tok_number = -5,

  // control
  tok_if = -6,
  tok_then = -7,
  tok_else = -8,
  tok_for = -9,
  tok_in = -10,

  // operators
  tok_binary = -11,
  tok_unary = -12,

  // var definition
  tok_var = -13
};

struct SourceLocation {
  int Line;
  int Col;
};

static SourceLocation CurLoc;
static SourceLocation LexLoc = {1, 0};

static int advance() {
  int LastChar = getchar();
  if (LastChar == '\n' || LastChar == '\r') {
    ++LexLoc.Line;
    LexLoc.Col = 0;
  } else {
    ++LexLoc.Col;
  }
  return LastChar;
}

static std::string IdentifierStr; // Filled in for identifiers and keywords
static double NumVal;             // Filled in if tok_number

/// gettok - Return the next token from standard input.
static int gettok() {
  static int LastChar = ' ';

  // Skip any whitespace.
  while (isspace(LastChar))
    LastChar = advance();

  CurLoc = LexLoc;

  if (isalpha(LastChar)) { // identifier: [a-zA-Z][a-zA-Z0-9]*
    IdentifierStr = LastChar;
    while (isalnum((LastChar = advance())))
      IdentifierStr += LastChar;

    if (IdentifierStr == "def")
      return tok_def;
    if (IdentifierStr == "extern")
      return tok_extern;
    if (IdentifierStr == "if")
      return tok_if;
    if (IdentifierStr == "then")
      return tok_then;
    if (IdentifierStr == "else")
      return tok_else;
    if (IdentifierStr == "for")
      return tok_for;
    if (IdentifierStr == "in")
      return tok_in;
    if (IdentifierStr == "binary")
      return tok_binary;
    if (IdentifierStr == "unary")
      return tok_unary;
    if (IdentifierStr == "var")
      return tok_var;
    return tok_identifier;
  }

  if (isdigit(LastChar) || LastChar == '.') { // Number: [0-9.]+
    std::string NumStr;
    do {
      NumStr += LastChar;
      LastChar = advance();
    } while (isdigit(LastChar) || LastChar == '.');

    NumVal = strtod(NumStr.c_str(), nullptr);
    return tok_number;
  }

  if (LastChar == '#') {
    // Comment until end of line.
    do
      LastChar = advance();
    while (LastChar != EOF && LastChar != '\n' && LastChar != '\r');

    if (LastChar != EOF)
      return gettok();
  }

  // Check for end of file.  Don't eat the EOF.
  if (LastChar == EOF)
    return tok_eof;

  // Otherwise, just return the character as its ascii value.
  int ThisChar = LastChar;
  LastChar = advance();
  return ThisChar;
}

//===----------------------------------------------------------------------===//
// Abstract Syntax Tree (aka Parse Tree)
//===----------------------------------------------------------------------===//

namespace {

/// ExprAST - Base class for all expression nodes.
class ExprAST {
  SourceLocation Loc;

public:
  ExprAST(SourceLocation Loc = CurLoc) : Loc(Loc) {}
  virtual ~ExprAST() = default;

  virtual Value codegen() = 0;
  virtual const std::string *getVariableName() const { return nullptr; }
  SourceLocation getSourceLocation() const { return Loc; }
};

/// NumberExprAST - Expression class for numeric literals like "1.0".
class NumberExprAST : public ExprAST {
  double Val;

public:
  NumberExprAST(double Val) : Val(Val) {}

  Value codegen() override;
};

/// VariableExprAST - Expression class for referencing a variable, like "a".
class VariableExprAST : public ExprAST {
  std::string Name;

public:
  VariableExprAST(SourceLocation Loc, const std::string &Name)
      : ExprAST(Loc), Name(Name) {}

  Value codegen() override;
  const std::string *getVariableName() const override { return &Name; }
};

/// UnaryExprAST - Expression class for a unary operator.
class UnaryExprAST : public ExprAST {
  char Opcode;
  std::unique_ptr<ExprAST> Operand;

public:
  UnaryExprAST(SourceLocation Loc, char Opcode,
               std::unique_ptr<ExprAST> Operand)
      : ExprAST(Loc), Opcode(Opcode), Operand(std::move(Operand)) {}

  Value codegen() override;
};

/// BinaryExprAST - Expression class for a binary operator.
class BinaryExprAST : public ExprAST {
  char Op;
  std::unique_ptr<ExprAST> LHS, RHS;

public:
  BinaryExprAST(SourceLocation Loc, char Op, std::unique_ptr<ExprAST> LHS,
                std::unique_ptr<ExprAST> RHS)
      : ExprAST(Loc), Op(Op), LHS(std::move(LHS)), RHS(std::move(RHS)) {}

  Value codegen() override;
};

/// CallExprAST - Expression class for function calls.
class CallExprAST : public ExprAST {
  std::string Callee;
  std::vector<std::unique_ptr<ExprAST>> Args;

public:
  CallExprAST(SourceLocation Loc, const std::string &Callee,
              std::vector<std::unique_ptr<ExprAST>> Args)
      : ExprAST(Loc), Callee(Callee), Args(std::move(Args)) {}

  Value codegen() override;
};

/// IfExprAST - Expression class for if/then/else.
class IfExprAST : public ExprAST {
  std::unique_ptr<ExprAST> Cond, Then, Else;

public:
  IfExprAST(SourceLocation Loc, std::unique_ptr<ExprAST> Cond,
            std::unique_ptr<ExprAST> Then, std::unique_ptr<ExprAST> Else)
      : ExprAST(Loc), Cond(std::move(Cond)), Then(std::move(Then)),
        Else(std::move(Else)) {}

  Value codegen() override;
};

/// ForExprAST - Expression class for for/in.
class ForExprAST : public ExprAST {
  std::string VarName;
  std::unique_ptr<ExprAST> Start, End, Step, Body;

public:
  ForExprAST(SourceLocation Loc, const std::string &VarName,
             std::unique_ptr<ExprAST> Start, std::unique_ptr<ExprAST> End,
             std::unique_ptr<ExprAST> Step, std::unique_ptr<ExprAST> Body)
      : ExprAST(Loc), VarName(VarName), Start(std::move(Start)),
        End(std::move(End)), Step(std::move(Step)), Body(std::move(Body)) {}

  Value codegen() override;
};

/// VarExprAST - Expression class for var/in.
class VarExprAST : public ExprAST {
  std::vector<std::pair<std::string, std::unique_ptr<ExprAST>>> VarNames;
  std::unique_ptr<ExprAST> Body;

public:
  VarExprAST(
      SourceLocation Loc,
      std::vector<std::pair<std::string, std::unique_ptr<ExprAST>>> VarNames,
      std::unique_ptr<ExprAST> Body)
      : ExprAST(Loc), VarNames(std::move(VarNames)), Body(std::move(Body)) {}

  Value codegen() override;
};

/// PrototypeAST - This class represents the "prototype" for a function,
/// which captures its argument names as well as if it is an operator.
class PrototypeAST {
  std::string Name;
  std::vector<std::string> Args;
  bool IsOperator;
  unsigned Precedence; // Precedence if a binary op.
  SourceLocation Loc;

public:
  PrototypeAST(SourceLocation Loc, const std::string &Name,
               std::vector<std::string> Args, bool IsOperator = false,
               unsigned Prec = 0)
      : Name(Name), Args(std::move(Args)), IsOperator(IsOperator),
        Precedence(Prec), Loc(Loc) {}

  func::FuncOp codegen();
  const std::string &getName() const { return Name; }
  const std::vector<std::string> &getArgs() const { return Args; }

  bool isUnaryOp() const { return IsOperator && Args.size() == 1; }
  bool isBinaryOp() const { return IsOperator && Args.size() == 2; }

  char getOperatorName() const {
    assert(isUnaryOp() || isBinaryOp());
    return Name.back();
  }

  unsigned getBinaryPrecedence() const { return Precedence; }
  SourceLocation getSourceLocation() const { return Loc; }
};

/// FunctionAST - This class represents a function definition itself.
class FunctionAST {
  std::unique_ptr<PrototypeAST> Proto;
  std::unique_ptr<ExprAST> Body;

public:
  FunctionAST(std::unique_ptr<PrototypeAST> Proto,
              std::unique_ptr<ExprAST> Body)
      : Proto(std::move(Proto)), Body(std::move(Body)) {}

  func::FuncOp codegen();
};

} // end anonymous namespace

//===----------------------------------------------------------------------===//
// Parser
//===----------------------------------------------------------------------===//

/// CurTok/getNextToken - Provide a simple token buffer.  CurTok is the current
/// token the parser is looking at.  getNextToken reads another token from the
/// lexer and updates CurTok with its results.
static int CurTok;
static int getNextToken() { return CurTok = gettok(); }

/// BinopPrecedence - This holds the precedence for each binary operator that is
/// defined.
static std::map<char, int> BinopPrecedence;

/// GetTokPrecedence - Get the precedence of the pending binary operator token.
static int GetTokPrecedence() {
  if (!isascii(CurTok))
    return -1;

  // Make sure it's a declared binop.
  int TokPrec = BinopPrecedence[CurTok];
  if (TokPrec <= 0)
    return -1;
  return TokPrec;
}

/// LogError* - These are little helper functions for error handling.
std::unique_ptr<ExprAST> LogError(const char *Str) {
  fprintf(stderr, "Error: %s\n", Str);
  return nullptr;
}
std::unique_ptr<PrototypeAST> LogErrorP(const char *Str) {
  LogError(Str);
  return nullptr;
}

static std::unique_ptr<ExprAST> ParseExpression();

/// numberexpr ::= number
static std::unique_ptr<ExprAST> ParseNumberExpr() {
  auto Result = std::make_unique<NumberExprAST>(NumVal);
  getNextToken(); // consume the number
  return std::move(Result);
}

/// parenexpr ::= '(' expression ')'
static std::unique_ptr<ExprAST> ParseParenExpr() {
  getNextToken(); // eat (.
  auto V = ParseExpression();
  if (!V)
    return nullptr;

  if (CurTok != ')')
    return LogError("expected ')'");
  getNextToken(); // eat ).
  return V;
}

/// identifierexpr
///   ::= identifier
///   ::= identifier '(' expression* ')'
static std::unique_ptr<ExprAST> ParseIdentifierExpr() {
  std::string IdName = IdentifierStr;
  SourceLocation IdLoc = CurLoc;

  getNextToken(); // eat identifier.

  if (CurTok != '(') // Simple variable ref.
    return std::make_unique<VariableExprAST>(IdLoc, IdName);

  // Call.
  getNextToken(); // eat (
  std::vector<std::unique_ptr<ExprAST>> Args;
  if (CurTok != ')') {
    while (true) {
      if (auto Arg = ParseExpression())
        Args.push_back(std::move(Arg));
      else
        return nullptr;

      if (CurTok == ')')
        break;

      if (CurTok != ',')
        return LogError("Expected ')' or ',' in argument list");
      getNextToken();
    }
  }

  // Eat the ')'.
  getNextToken();

  return std::make_unique<CallExprAST>(IdLoc, IdName, std::move(Args));
}

/// ifexpr ::= 'if' expression 'then' expression 'else' expression
static std::unique_ptr<ExprAST> ParseIfExpr() {
  SourceLocation IfLoc = CurLoc;
  getNextToken(); // eat the if.

  auto Cond = ParseExpression();
  if (!Cond)
    return nullptr;

  if (CurTok != tok_then)
    return LogError("expected then");
  getNextToken(); // eat the then.

  auto Then = ParseExpression();
  if (!Then)
    return nullptr;

  if (CurTok != tok_else)
    return LogError("expected else");
  getNextToken(); // eat the else.

  auto Else = ParseExpression();
  if (!Else)
    return nullptr;

  return std::make_unique<IfExprAST>(IfLoc, std::move(Cond), std::move(Then),
                                     std::move(Else));
}

/// forexpr ::= 'for' identifier '=' expr ',' expr (',' expr)? 'in' expression
static std::unique_ptr<ExprAST> ParseForExpr() {
  SourceLocation ForLoc = CurLoc;
  getNextToken(); // eat the for.

  if (CurTok != tok_identifier)
    return LogError("expected identifier after for");

  std::string IdName = IdentifierStr;
  getNextToken(); // eat identifier.

  if (CurTok != '=')
    return LogError("expected '=' after for");
  getNextToken(); // eat '='.

  auto Start = ParseExpression();
  if (!Start)
    return nullptr;
  if (CurTok != ',')
    return LogError("expected ',' after for start value");
  getNextToken();

  auto End = ParseExpression();
  if (!End)
    return nullptr;

  // The step value is optional.
  std::unique_ptr<ExprAST> Step;
  if (CurTok == ',') {
    getNextToken();
    Step = ParseExpression();
    if (!Step)
      return nullptr;
  }

  if (CurTok != tok_in)
    return LogError("expected 'in' after for");
  getNextToken(); // eat the in.

  auto Body = ParseExpression();
  if (!Body)
    return nullptr;

  return std::make_unique<ForExprAST>(ForLoc, IdName, std::move(Start),
                                      std::move(End), std::move(Step),
                                      std::move(Body));
}

/// varexpr ::= 'var' identifier ('=' expression)?
///                    (',' identifier ('=' expression)?)* 'in' expression
static std::unique_ptr<ExprAST> ParseVarExpr() {
  SourceLocation VarLoc = CurLoc;
  getNextToken(); // eat the var.

  std::vector<std::pair<std::string, std::unique_ptr<ExprAST>>> VarNames;
  if (CurTok != tok_identifier)
    return LogError("expected identifier after var");

  while (true) {
    std::string Name = IdentifierStr;
    getNextToken(); // eat identifier.

    std::unique_ptr<ExprAST> Init;
    if (CurTok == '=') {
      getNextToken(); // eat '='.
      Init = ParseExpression();
      if (!Init)
        return nullptr;
    }

    VarNames.emplace_back(Name, std::move(Init));

    if (CurTok != ',')
      break;
    getNextToken(); // eat ','.
    if (CurTok != tok_identifier)
      return LogError("expected identifier list after var");
  }

  if (CurTok != tok_in)
    return LogError("expected 'in' keyword after 'var'");
  getNextToken(); // eat 'in'.

  auto Body = ParseExpression();
  if (!Body)
    return nullptr;

  return std::make_unique<VarExprAST>(VarLoc, std::move(VarNames),
                                      std::move(Body));
}

/// primary
///   ::= identifierexpr
///   ::= numberexpr
///   ::= parenexpr
///   ::= ifexpr
///   ::= forexpr
///   ::= varexpr
static std::unique_ptr<ExprAST> ParsePrimary() {
  switch (CurTok) {
  default:
    return LogError("unknown token when expecting an expression");
  case tok_identifier:
    return ParseIdentifierExpr();
  case tok_number:
    return ParseNumberExpr();
  case '(':
    return ParseParenExpr();
  case tok_if:
    return ParseIfExpr();
  case tok_for:
    return ParseForExpr();
  case tok_var:
    return ParseVarExpr();
  }
}

/// unary
///   ::= primary
///   ::= '!' unary
static std::unique_ptr<ExprAST> ParseUnary() {
  // If the current token is not an operator, it must be a primary expression.
  if (!isascii(CurTok) || CurTok == '(' || CurTok == ',')
    return ParsePrimary();

  // If this is a unary operator, read it.
  int Opc = CurTok;
  SourceLocation UnaryLoc = CurLoc;
  getNextToken();
  if (auto Operand = ParseUnary())
    return std::make_unique<UnaryExprAST>(UnaryLoc, Opc, std::move(Operand));
  return nullptr;
}

/// binoprhs
///   ::= ('+' unary)*
static std::unique_ptr<ExprAST> ParseBinOpRHS(int ExprPrec,
                                              std::unique_ptr<ExprAST> LHS) {
  // If this is a binop, find its precedence.
  while (true) {
    int TokPrec = GetTokPrecedence();

    // If this is a binop that binds at least as tightly as the current binop,
    // consume it, otherwise we are done.
    if (TokPrec < ExprPrec)
      return LHS;

    // Okay, we know this is a binop.
    int BinOp = CurTok;
    SourceLocation BinLoc = CurLoc;
    getNextToken(); // eat binop

    // Parse the unary expression after the binary operator.
    auto RHS = ParseUnary();
    if (!RHS)
      return nullptr;

    // If BinOp binds less tightly with RHS than the operator after RHS, let
    // the pending operator take RHS as its LHS.
    int NextPrec = GetTokPrecedence();
    if (TokPrec < NextPrec) {
      RHS = ParseBinOpRHS(TokPrec + 1, std::move(RHS));
      if (!RHS)
        return nullptr;
    }

    // Merge LHS/RHS.
    LHS = std::make_unique<BinaryExprAST>(BinLoc, BinOp, std::move(LHS),
                                          std::move(RHS));
  }
}

/// expression
///   ::= unary binoprhs
///
static std::unique_ptr<ExprAST> ParseExpression() {
  auto LHS = ParseUnary();
  if (!LHS)
    return nullptr;

  return ParseBinOpRHS(0, std::move(LHS));
}

/// prototype
///   ::= id '(' id* ')'
///   ::= binary LETTER number? (id, id)
///   ::= unary LETTER (id)
static std::unique_ptr<PrototypeAST> ParsePrototype() {
  SourceLocation FnLoc = CurLoc;
  std::string FnName;
  unsigned Kind = 0; // 0 = identifier, 1 = unary, 2 = binary.
  unsigned BinaryPrecedence = 30;

  switch (CurTok) {
  default:
    return LogErrorP("Expected function name in prototype");
  case tok_identifier:
    FnName = IdentifierStr;
    getNextToken();
    break;
  case tok_unary:
    getNextToken();
    if (!isascii(CurTok))
      return LogErrorP("Expected unary operator");
    FnName = "unary";
    FnName += static_cast<char>(CurTok);
    Kind = 1;
    getNextToken();
    break;
  case tok_binary:
    getNextToken();
    if (!isascii(CurTok))
      return LogErrorP("Expected binary operator");
    FnName = "binary";
    FnName += static_cast<char>(CurTok);
    Kind = 2;
    getNextToken();

    // Read the precedence if present.
    if (CurTok == tok_number) {
      if (NumVal < 1 || NumVal > 100)
        return LogErrorP("Invalid precedence: must be 1..100");
      BinaryPrecedence = static_cast<unsigned>(NumVal);
      getNextToken();
    }
    break;
  }

  if (CurTok != '(')
    return LogErrorP("Expected '(' in prototype");

  std::vector<std::string> ArgNames;
  while (getNextToken() == tok_identifier)
    ArgNames.push_back(IdentifierStr);
  if (CurTok != ')')
    return LogErrorP("Expected ')' in prototype");

  // success.
  getNextToken(); // eat ')'.

  // Verify the right number of names for an operator.
  if (Kind && ArgNames.size() != Kind)
    return LogErrorP("Invalid number of operands for operator");

  return std::make_unique<PrototypeAST>(FnLoc, FnName, std::move(ArgNames),
                                        Kind != 0, BinaryPrecedence);
}

/// definition ::= 'def' prototype expression
static std::unique_ptr<FunctionAST> ParseDefinition() {
  getNextToken(); // eat def.
  auto Proto = ParsePrototype();
  if (!Proto)
    return nullptr;

  if (auto E = ParseExpression())
    return std::make_unique<FunctionAST>(std::move(Proto), std::move(E));
  return nullptr;
}

/// toplevelexpr ::= expression
static std::unique_ptr<FunctionAST> ParseTopLevelExpr() {
  SourceLocation FnLoc = CurLoc;
  if (auto E = ParseExpression()) {
    // Make an anonymous proto.
    auto Proto = std::make_unique<PrototypeAST>(FnLoc, "__anon_expr",
                                                std::vector<std::string>());
    return std::make_unique<FunctionAST>(std::move(Proto), std::move(E));
  }
  return nullptr;
}

/// external ::= 'extern' prototype
static std::unique_ptr<PrototypeAST> ParseExtern() {
  getNextToken(); // eat extern.
  return ParsePrototype();
}

//===----------------------------------------------------------------------===//
// Code Generation
//===----------------------------------------------------------------------===//

static std::unique_ptr<MLIRContext> TheContext;
static OwningOpRef<ModuleOp> TheModule;
static std::unique_ptr<OpBuilder> TheBuilder;
static std::unique_ptr<PassManager> ThePM;
static std::map<std::string, Value> NamedValues;
static std::unique_ptr<llvm::orc::KaleidoscopeJIT> TheJIT;
static std::map<std::string, std::unique_ptr<PrototypeAST>> FunctionProtos;
static llvm::ExitOnError ExitOnErr;
static llvm::cl::opt<bool> DumpMLIR("dump-mlir",
                                    llvm::cl::desc("Print generated MLIR"),
                                    llvm::cl::init(false));
static llvm::cl::opt<bool> EmitObject(
    "emit-object",
    llvm::cl::desc(
        "Compile an input file to an object instead of using the JIT"),
    llvm::cl::init(false));
static llvm::cl::opt<std::string> InputFilename(llvm::cl::Positional,
                                                llvm::cl::desc("<input file>"),
                                                llvm::cl::init(""));
static llvm::cl::opt<std::string>
    TargetTripleOption("target",
                       llvm::cl::desc("Target triple for object emission"),
                       llvm::cl::value_desc("triple"), llvm::cl::init(""));
static llvm::cl::opt<std::string>
    OutputFilename("o", llvm::cl::desc("Output filename"),
                   llvm::cl::value_desc("filename"), llvm::cl::init(""));
static llvm::cl::opt<char>
    OptLevel("O", llvm::cl::desc("Optimization level: -O0, -O1, -O2, or -O3"),
             llvm::cl::Prefix, llvm::cl::init('2'));
static llvm::cl::opt<bool>
    DumpLLVMIR("dump-llvm-ir",
               llvm::cl::desc("Print LLVM IR before emitting the object file"),
               llvm::cl::init(false));

static SourceLocation CodegenLoc = {1, 1};

static Location getLocation() {
  llvm::StringRef Filename = "<stdin>";
  if (!InputFilename.empty())
    Filename = InputFilename.getValue();
  return FileLineColLoc::get(TheContext.get(), Filename, CodegenLoc.Line,
                             CodegenLoc.Col);
}

class LocationGuard {
  SourceLocation Previous;

public:
  explicit LocationGuard(SourceLocation Loc) : Previous(CodegenLoc) {
    CodegenLoc = Loc;
  }
  ~LocationGuard() { CodegenLoc = Previous; }
};

Value LogErrorV(const char *Str) {
  LogError(Str);
  return {};
}

func::FuncOp getFunction(const std::string &Name) {
  // First, see if the function has already been added to the current module.
  if (auto Function = TheModule->lookupSymbol<func::FuncOp>(Name))
    return Function;

  // If not, codegen the declaration from an existing prototype.
  auto It = FunctionProtos.find(Name);
  if (It != FunctionProtos.end()) {
    auto Function = It->second->codegen();
    Function.setPrivate();
    return Function;
  }

  return {};
}

static func::FuncOp getCurrentFunction() {
  Operation *Parent = TheBuilder->getInsertionBlock()->getParentOp();
  if (auto Function = dyn_cast<func::FuncOp>(Parent))
    return Function;
  return Parent->getParentOfType<func::FuncOp>();
}

/// CreateVariable - Preserve a source variable until dialect lowering.
static Value CreateVariable(StringRef Name, Value InitialValue,
                            int64_t ArgumentNumber = 0) {
  return TheBuilder->create<kaleidoscope::DeclareOp>(
      getLocation(), kaleidoscope::VariableType::get(TheContext.get()),
      InitialValue, TheBuilder->getStringAttr(Name),
      TheBuilder->getI64IntegerAttr(ArgumentNumber));
}

Value NumberExprAST::codegen() {
  LocationGuard Guard(getSourceLocation());
  return TheBuilder->create<arith::ConstantOp>(
      getLocation(), TheBuilder->getF64FloatAttr(Val));
}

Value VariableExprAST::codegen() {
  LocationGuard Guard(getSourceLocation());
  // Look this variable up in the function.
  auto It = NamedValues.find(Name);
  if (It == NamedValues.end())
    return LogErrorV("Unknown variable name");

  return TheBuilder->create<kaleidoscope::ReadOp>(
      getLocation(), TheBuilder->getF64Type(), It->second);
}

Value UnaryExprAST::codegen() {
  LocationGuard Guard(getSourceLocation());
  Value OperandV = Operand->codegen();
  if (!OperandV)
    return {};

  auto Operator = getFunction(std::string("unary") + Opcode);
  if (!Operator)
    return LogErrorV("Unknown unary operator");

  return TheBuilder->create<func::CallOp>(getLocation(), Operator, OperandV)
      .getResult(0);
}

Value BinaryExprAST::codegen() {
  LocationGuard Guard(getSourceLocation());
  // Assignment stores into the variable's mutable memref slot.
  if (Op == '=') {
    const std::string *Name = LHS->getVariableName();
    if (!Name)
      return LogErrorV("destination of '=' must be a variable");

    Value AssignedValue = RHS->codegen();
    if (!AssignedValue)
      return {};

    auto It = NamedValues.find(*Name);
    if (It == NamedValues.end())
      return LogErrorV("Unknown variable name");

    TheBuilder->create<kaleidoscope::AssignOp>(getLocation(), It->second,
                                               AssignedValue);
    return AssignedValue;
  }

  Value L = LHS->codegen();
  Value R = RHS->codegen();
  if (!L || !R)
    return {};

  switch (Op) {
  case '+':
    return TheBuilder->create<arith::AddFOp>(getLocation(), L, R);
  case '-':
    return TheBuilder->create<arith::SubFOp>(getLocation(), L, R);
  case '*':
    return TheBuilder->create<arith::MulFOp>(getLocation(), L, R);
  case '<': {
    Value Comparison = TheBuilder->create<arith::CmpFOp>(
        getLocation(), arith::CmpFPredicate::ULT, L, R);
    // Convert bool 0/1 to double 0.0 or 1.0.
    return TheBuilder->create<arith::UIToFPOp>(
        getLocation(), TheBuilder->getF64Type(), Comparison);
  }
  default:
    break;
  }

  // If it wasn't a builtin binary operator, it must be a user-defined one.
  auto Operator = getFunction(std::string("binary") + Op);
  if (!Operator)
    return LogErrorV("Unknown binary operator");

  Value Operands[] = {L, R};
  return TheBuilder->create<func::CallOp>(getLocation(), Operator, Operands)
      .getResult(0);
}

Value CallExprAST::codegen() {
  LocationGuard Guard(getSourceLocation());
  // Look up the name in the global module table.
  auto CalleeF = getFunction(Callee);
  if (!CalleeF)
    return LogErrorV("Unknown function referenced");

  // If argument mismatch error.
  if (CalleeF.getNumArguments() != Args.size())
    return LogErrorV("Incorrect # arguments passed");

  std::vector<Value> ArgsV;
  for (auto &Arg : Args) {
    ArgsV.push_back(Arg->codegen());
    if (!ArgsV.back())
      return {};
  }

  return TheBuilder->create<func::CallOp>(getLocation(), CalleeF, ArgsV)
      .getResult(0);
}

Value IfExprAST::codegen() {
  LocationGuard Guard(getSourceLocation());
  Value CondV = Cond->codegen();
  if (!CondV)
    return {};

  // Convert the condition to a boolean by comparing it with 0.0.
  Value Zero = TheBuilder->create<arith::ConstantOp>(
      getLocation(), TheBuilder->getF64FloatAttr(0.0));
  CondV = TheBuilder->create<arith::CmpFOp>(
      getLocation(), arith::CmpFPredicate::ONE, CondV, Zero);

  bool CodegenFailed = false;
  auto IfOp = TheBuilder->create<scf::IfOp>(
      getLocation(), CondV,
      [&](OpBuilder &Builder, Location Loc) {
        Value ThenV = Then->codegen();
        if (!ThenV) {
          CodegenFailed = true;
          ThenV = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(0.0));
        }
        Builder.create<scf::YieldOp>(Loc, ThenV);
      },
      [&](OpBuilder &Builder, Location Loc) {
        Value ElseV = Else->codegen();
        if (!ElseV) {
          CodegenFailed = true;
          ElseV = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(0.0));
        }
        Builder.create<scf::YieldOp>(Loc, ElseV);
      });

  if (CodegenFailed)
    return {};
  return IfOp.getResult(0);
}

Value ForExprAST::codegen() {
  LocationGuard Guard(getSourceLocation());
  // Emit the start value before putting the loop variable in scope.
  Value StartVal = Start->codegen();
  if (!StartVal)
    return {};

  Value Variable = CreateVariable(VarName, StartVal);

  auto OldValue = NamedValues.find(VarName);
  bool HadOldValue = OldValue != NamedValues.end();
  Value SavedValue = HadOldValue ? OldValue->second : Value();
  NamedValues[VarName] = Variable;
  bool CodegenFailed = false;

  // Test the condition before each iteration, then emit the body and step.
  TheBuilder->create<scf::WhileOp>(
      getLocation(), TypeRange{}, ValueRange{},
      [&](OpBuilder &Builder, Location Loc, ValueRange) {
        Value EndCond = End->codegen();
        if (!EndCond) {
          CodegenFailed = true;
          EndCond = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(0.0));
        }

        Value Zero = Builder.create<arith::ConstantOp>(
            Loc, Builder.getF64FloatAttr(0.0));
        EndCond = Builder.create<arith::CmpFOp>(Loc, arith::CmpFPredicate::ONE,
                                                EndCond, Zero);
        Builder.create<scf::ConditionOp>(Loc, EndCond, ValueRange{});
      },
      [&](OpBuilder &Builder, Location Loc, ValueRange) {
        if (!Body->codegen())
          CodegenFailed = true;

        Value StepVal;
        if (Step)
          StepVal = Step->codegen();
        else
          StepVal = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(1.0));
        if (!StepVal) {
          CodegenFailed = true;
          StepVal = Builder.create<arith::ConstantOp>(
              Loc, Builder.getF64FloatAttr(1.0));
        }

        // Reload after the body and step in case either mutated the variable.
        Value Current = Builder.create<kaleidoscope::ReadOp>(
            Loc, Builder.getF64Type(), Variable);
        Value NextVar = Builder.create<arith::AddFOp>(Loc, Current, StepVal);
        Builder.create<kaleidoscope::AssignOp>(Loc, Variable, NextVar);
        Builder.create<scf::YieldOp>(Loc);
      });

  // Restore any variable shadowed by the loop induction variable.
  if (HadOldValue)
    NamedValues[VarName] = SavedValue;
  else
    NamedValues.erase(VarName);

  if (CodegenFailed)
    return {};

  // A for expression always returns 0.0.
  return TheBuilder->create<arith::ConstantOp>(
      getLocation(), TheBuilder->getF64FloatAttr(0.0));
}

Value VarExprAST::codegen() {
  LocationGuard Guard(getSourceLocation());
  std::vector<std::pair<std::string, std::optional<Value>>> OldBindings;

  auto RestoreBindings = [&]() {
    for (auto It = OldBindings.rbegin(); It != OldBindings.rend(); ++It) {
      if (It->second)
        NamedValues[It->first] = *It->second;
      else
        NamedValues.erase(It->first);
    }
  };

  for (auto &Variable : VarNames) {
    const std::string &Name = Variable.first;

    // Generate the initializer before introducing the new binding.
    Value InitialValue;
    if (Variable.second)
      InitialValue = Variable.second->codegen();
    else
      InitialValue = TheBuilder->create<arith::ConstantOp>(
          getLocation(), TheBuilder->getF64FloatAttr(0.0));
    if (!InitialValue) {
      RestoreBindings();
      return {};
    }

    Value Storage = CreateVariable(Name, InitialValue);

    auto Old = NamedValues.find(Name);
    OldBindings.emplace_back(Name, Old == NamedValues.end()
                                       ? std::optional<Value>()
                                       : std::optional<Value>(Old->second));
    NamedValues[Name] = Storage;
  }

  Value BodyValue = Body->codegen();
  RestoreBindings();
  return BodyValue;
}

func::FuncOp PrototypeAST::codegen() {
  LocationGuard Guard(getSourceLocation());
  // Make the function type: double(double, double), etc.
  std::vector<Type> Doubles(Args.size(), TheBuilder->getF64Type());
  auto FunctionType =
      TheBuilder->getFunctionType(Doubles, {TheBuilder->getF64Type()});

  auto Function = func::FuncOp::create(getLocation(), Name, FunctionType);
  TheModule->push_back(Function);
  return Function;
}

func::FuncOp FunctionAST::codegen() {
  // Save the prototype so declarations can be emitted in later modules.
  auto &P = *Proto;
  LocationGuard Guard(P.getSourceLocation());
  FunctionProtos[Proto->getName()] = std::move(Proto);
  auto TheFunction = getFunction(P.getName());

  if (!TheFunction)
    return {};

  if (!TheFunction.isDeclaration()) {
    LogError("Function cannot be redefined.");
    return {};
  }

  // A definition is visible outside the module, even if an earlier extern
  // declaration created the function with private symbol visibility.
  TheFunction.setPublic();

  // If this is a binary operator, install its precedence.
  if (P.isBinaryOp())
    BinopPrecedence[P.getOperatorName()] = P.getBinaryPrecedence();

  // Create a new basic block to start insertion into.
  Block *EntryBlock = TheFunction.addEntryBlock();
  TheBuilder->setInsertionPointToStart(EntryBlock);

  // Give each function argument a mutable storage slot.
  NamedValues.clear();
  unsigned Index = 0;
  for (BlockArgument Argument : TheFunction.getArguments()) {
    StringRef Name = P.getArgs()[Index];
    Value Storage = CreateVariable(Name, Argument, Index + 1);
    NamedValues[Name.str()] = Storage;
    ++Index;
  }

  if (Value RetVal = Body->codegen()) {
    // Finish off the function.
    TheBuilder->create<func::ReturnOp>(getLocation(), RetVal);

    // Validate the generated code, checking for consistency.
    if (succeeded(verify(TheFunction))) {
      // Run the optimizer on the module.
      if (failed(ThePM->run(*TheModule))) {
        LogError("Could not optimize function.");
        TheFunction.erase();
        if (P.isBinaryOp())
          BinopPrecedence.erase(P.getOperatorName());
        return {};
      }
      return TheFunction;
    }
  }

  // Error reading body, remove function.
  TheFunction.erase();
  if (P.isBinaryOp())
    BinopPrecedence.erase(P.getOperatorName());
  return {};
}

//===----------------------------------------------------------------------===//
// Top-Level parsing and JIT Driver
//===----------------------------------------------------------------------===//

static void InitializeModuleAndManagers() {
  // Destroy objects that refer to the old context before replacing it.
  ThePM.reset();
  TheBuilder.reset();
  TheModule = OwningOpRef<ModuleOp>();
  TheContext.reset();

  // Open a new context and module.
  TheContext = std::make_unique<MLIRContext>();
  TheContext->loadDialect<arith::ArithDialect, cf::ControlFlowDialect,
                          func::FuncDialect, kaleidoscope::KaleidoscopeDialect,
                          memref::MemRefDialect, scf::SCFDialect>();
  TheModule = ModuleOp::create(UnknownLoc::get(TheContext.get()));

  // Create a new builder for the module.
  TheBuilder = std::make_unique<OpBuilder>(TheContext.get());

  // Create a pass manager and enable our simple optimizations above -O0.
  ThePM = std::make_unique<PassManager>(TheContext.get());
  if (OptLevel != '0') {
    ThePM->addNestedPass<func::FuncOp>(createCanonicalizerPass());
    ThePM->addNestedPass<func::FuncOp>(createCSEPass());
  }
}

struct LoweredModule {
  std::unique_ptr<llvm::LLVMContext> Context;
  std::unique_ptr<llvm::Module> Module;
};

class DeclareOpLowering : public OpConversionPattern<kaleidoscope::DeclareOp> {
public:
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(kaleidoscope::DeclareOp Op, OpAdaptor Adaptor,
                  ConversionPatternRewriter &Rewriter) const override {
    Location Loc = Op.getLoc();
    Type DoubleType = Rewriter.getF64Type();
    Type PointerType = LLVM::LLVMPointerType::get(Rewriter.getContext());
    Value One = LLVM::ConstantOp::create(Rewriter, Loc, Rewriter.getI64Type(),
                                         Rewriter.getI64IntegerAttr(1));
    Value Address =
        LLVM::AllocaOp::create(Rewriter, Loc, PointerType, DoubleType, One, 0);
    LLVM::StoreOp::create(Rewriter, Loc, Adaptor.getInitialValue(), Address);

    auto Function = Op->getParentOfType<LLVM::LLVMFuncOp>();
    auto ScopeLoc =
        Function.getLoc()
            ->findInstanceOf<FusedLocWith<LLVM::DISubprogramAttr>>();
    if (!ScopeLoc)
      return Rewriter.notifyMatchFailure(Op, "function has no debug scope");

    LLVM::DISubprogramAttr Scope = ScopeLoc.getMetadata();
    int64_t Line = Scope.getLine();
    if (auto FileLoc = Loc->findInstanceOf<FileLineColLoc>())
      Line = FileLoc.getLine();
    auto VariableType = LLVM::DIBasicTypeAttr::get(
        Rewriter.getContext(), llvm::dwarf::DW_TAG_base_type, "double", 64,
        llvm::dwarf::DW_ATE_float);
    auto Variable = LLVM::DILocalVariableAttr::get(
        Scope, Op.getName(), Scope.getFile(), Line, Op.getArgumentNumber(),
        /*alignInBits=*/0, VariableType, LLVM::DIFlags::Zero);
    LLVM::DbgDeclareOp::create(
        Rewriter, Loc, Address, Variable,
        LLVM::DIExpressionAttr::get(Rewriter.getContext()));

    Rewriter.replaceOp(Op, Address);
    return success();
  }
};

class ReadOpLowering : public OpConversionPattern<kaleidoscope::ReadOp> {
public:
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(kaleidoscope::ReadOp Op, OpAdaptor Adaptor,
                  ConversionPatternRewriter &Rewriter) const override {
    Rewriter.replaceOpWithNewOp<LLVM::LoadOp>(Op, Rewriter.getF64Type(),
                                              Adaptor.getVariable());
    return success();
  }
};

class AssignOpLowering : public OpConversionPattern<kaleidoscope::AssignOp> {
public:
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(kaleidoscope::AssignOp Op, OpAdaptor Adaptor,
                  ConversionPatternRewriter &Rewriter) const override {
    Rewriter.replaceOpWithNewOp<LLVM::StoreOp>(Op, Adaptor.getValue(),
                                               Adaptor.getVariable());
    return success();
  }
};

class LowerKaleidoscopeVariablesPass
    : public PassWrapper<LowerKaleidoscopeVariablesPass,
                         OperationPass<ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerKaleidoscopeVariablesPass)

  void runOnOperation() override {
    MLIRContext &Context = getContext();
    LLVMTypeConverter Converter(&Context);
    Converter.addConversion([](kaleidoscope::VariableType InputType) -> Type {
      return LLVM::LLVMPointerType::get(InputType.getContext());
    });

    ConversionTarget Target(Context);
    Target.addIllegalDialect<kaleidoscope::KaleidoscopeDialect>();
    Target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });

    RewritePatternSet Patterns(&Context);
    Patterns.add<DeclareOpLowering, ReadOpLowering, AssignOpLowering>(Converter,
                                                                      &Context);
    if (failed(applyPartialConversion(getOperation(), Target,
                                      std::move(Patterns))))
      signalPassFailure();
  }
};

static llvm::Expected<LoweredModule>
lowerToLLVM(const llvm::DataLayout &DataLayout) {
  // Lower the high-level MLIR operations to the LLVM dialect.
  PassManager LoweringPM(TheContext.get());
  LoweringPM.addPass(createSCFToControlFlowPass());
  LoweringPM.addPass(createConvertFuncToLLVMPass());
  LoweringPM.addPass(createArithToLLVMConversionPass());
  LoweringPM.addPass(createFinalizeMemRefToLLVMConversionPass());
  LoweringPM.addPass(createConvertControlFlowToLLVMPass());

  // Clean up any temporary casts introduced by dialect conversion.
  LoweringPM.addPass(createReconcileUnrealizedCastsPass());
  if (failed(LoweringPM.run(*TheModule)))
    return llvm::make_error<llvm::StringError>(
        "could not lower module to the LLVM dialect",
        llvm::inconvertibleErrorCode());

  // Add compile-unit and function scopes before lowering source variables.
  PassManager DebugPM(TheContext.get());
  DebugPM.addPass(
      createKaleidoscopeDebugInfoPass(InputFilename.getValue(), OptLevel));

  // Lower source variables while their names and locations are still present.
  DebugPM.addPass(std::make_unique<LowerKaleidoscopeVariablesPass>());

  // Fill in the remaining debug scopes on the lowered LLVM operations.
  LLVM::DIScopeForLLVMFuncOpPassOptions DebugOptions;
  DebugOptions.emissionKind = LLVM::DIEmissionKind::Full;
  DebugPM.addPass(
      LLVM::createDIScopeForLLVMFuncOpPass(std::move(DebugOptions)));
  if (failed(DebugPM.run(*TheModule)))
    return llvm::make_error<llvm::StringError>(
        "could not add LLVM debug scopes", llvm::inconvertibleErrorCode());

  // Register the translations from MLIR's LLVM dialect to LLVM IR.
  registerBuiltinDialectTranslation(*TheContext);
  registerLLVMDialectTranslation(*TheContext);

  // Translate the lowered MLIR module into the LLVM IR module consumed by the
  // target's object-file emitter.
  auto LLVMContext = std::make_unique<llvm::LLVMContext>();
  auto LLVMModule = translateModuleToLLVMIR(*TheModule, *LLVMContext);
  if (!LLVMModule)
    return llvm::make_error<llvm::StringError>(
        "could not translate the LLVM dialect to LLVM IR",
        llvm::inconvertibleErrorCode());

  LLVMModule->setDataLayout(DataLayout);

  if (DumpLLVMIR) {
    LLVMModule->print(llvm::errs(), nullptr);
    llvm::errs() << '\n';
  }

  return LoweredModule{std::move(LLVMContext), std::move(LLVMModule)};
}

static void HandleDefinition() {
  if (auto FnAST = ParseDefinition()) {
    if (auto FnIR = FnAST->codegen()) {
      if (DumpMLIR) {
        fprintf(stderr, "Read function definition:\n");
        FnIR.print(llvm::errs(), OpPrintingFlags().assumeVerified());
        fprintf(stderr, "\n");
      }

      if (!EmitObject) {
        auto Lowered = ExitOnErr(lowerToLLVM(TheJIT->getDataLayout()));
        ExitOnErr(TheJIT->addModule(llvm::orc::ThreadSafeModule(
            std::move(Lowered.Module), std::move(Lowered.Context))));
        InitializeModuleAndManagers();
      }
    }
  } else {
    // Skip token for error recovery.
    getNextToken();
  }
}

static void HandleExtern() {
  if (auto ProtoAST = ParseExtern()) {
    if (auto FnIR = ProtoAST->codegen()) {
      FnIR.setPrivate();
      if (DumpMLIR) {
        fprintf(stderr, "Read extern:\n");
        FnIR.print(llvm::errs(), OpPrintingFlags().assumeVerified());
        fprintf(stderr, "\n");
      }
      FunctionProtos[ProtoAST->getName()] = std::move(ProtoAST);
    }
  } else {
    // Skip token for error recovery.
    getNextToken();
  }
}

static void HandleTopLevelExpression() {
  // Evaluate a top-level expression with the JIT, or retain it when compiling.
  if (auto FnAST = ParseTopLevelExpr()) {
    if (auto FnIR = FnAST->codegen()) {
      if (DumpMLIR) {
        fprintf(stderr, "Read top-level expression:\n");
        FnIR.print(llvm::errs(), OpPrintingFlags().assumeVerified());
        fprintf(stderr, "\n");
      }

      if (!EmitObject) {
        auto RT = TheJIT->getMainJITDylib().createResourceTracker();
        auto Lowered = ExitOnErr(lowerToLLVM(TheJIT->getDataLayout()));
        ExitOnErr(TheJIT->addModule(
            llvm::orc::ThreadSafeModule(std::move(Lowered.Module),
                                        std::move(Lowered.Context)),
            RT));
        InitializeModuleAndManagers();

        auto ExprSymbol = ExitOnErr(TheJIT->lookup("__anon_expr"));
        double (*FP)() = ExprSymbol.getAddress().toPtr<double (*)()>();
        fprintf(stderr, "Evaluated to %f\n", FP());

        ExitOnErr(RT->remove());
      }
    }
  } else {
    // Skip token for error recovery.
    getNextToken();
  }
}

//===----------------------------------------------------------------------===//
// "Library" functions that can be "extern'd" from user code.
//===----------------------------------------------------------------------===//

#ifdef _WIN32
#define DLLEXPORT __declspec(dllexport)
#else
#define DLLEXPORT
#endif

/// putchard - putchar that takes a double and returns 0.
extern "C" DLLEXPORT double putchard(double X) {
  fputc((char)X, stderr);
  return 0;
}

/// printd - printf that takes a double, prints it as "%f\n", and returns 0.
extern "C" DLLEXPORT double printd(double X) {
  fprintf(stderr, "%f\n", X);
  return 0;
}

/// top ::= definition | external | expression | ';'
static void MainLoop() {
  while (true) {
    switch (CurTok) {
    case tok_eof:
      return;
    case ';': // ignore top-level semicolons.
      if (!EmitObject)
        fprintf(stderr, "ready> ");
      getNextToken();
      continue;
    case tok_def:
      HandleDefinition();
      break;
    case tok_extern:
      HandleExtern();
      break;
    default:
      HandleTopLevelExpression();
      break;
    }
    if (!EmitObject && CurTok != tok_eof && CurTok != ';')
      fprintf(stderr, "ready> ");
  }
}

//===----------------------------------------------------------------------===//
// Main driver code.
//===----------------------------------------------------------------------===//

int main(int argc, char **argv) {
  llvm::cl::ParseCommandLineOptions(argc, argv,
                                    "Kaleidoscope object file compiler\n");

  if (OptLevel < '0' || OptLevel > '3') {
    llvm::errs() << "Error: optimization level must be -O0, -O1, -O2, or -O3\n";
    return 1;
  }

  if (EmitObject && InputFilename.empty()) {
    llvm::errs() << "Error: --emit-object requires an input file\n";
    return 1;
  }
  if (!EmitObject && !InputFilename.empty()) {
    llvm::errs() << "Error: an input file requires --emit-object\n";
    return 1;
  }
  if (!EmitObject && !TargetTripleOption.empty()) {
    llvm::errs() << "Error: --target requires --emit-object\n";
    return 1;
  }
  if (!EmitObject && !OutputFilename.empty()) {
    llvm::errs() << "Error: -o requires --emit-object\n";
    return 1;
  }
  if (EmitObject && !std::freopen(InputFilename.c_str(), "r", stdin)) {
    std::perror(("Error opening " + InputFilename).c_str());
    return 1;
  }

  llvm::InitializeAllTargetInfos();
  llvm::InitializeAllTargets();
  llvm::InitializeAllTargetMCs();
  llvm::InitializeAllAsmParsers();
  llvm::InitializeAllAsmPrinters();

  // Install standard binary operators.
  // 1 is lowest precedence.
  BinopPrecedence['='] = 2;
  BinopPrecedence['<'] = 10;
  BinopPrecedence['+'] = 20;
  BinopPrecedence['-'] = 20;
  BinopPrecedence['*'] = 40; // highest.

  // Prime the first token.
  if (!EmitObject)
    fprintf(stderr, "ready> ");
  getNextToken();

  if (!EmitObject)
    TheJIT = ExitOnErr(llvm::orc::KaleidoscopeJIT::Create());

  // JIT mode replaces this module as definitions are submitted. Object mode
  // retains it until the entire input has been parsed.
  InitializeModuleAndManagers();

  // Run the main "interpreter loop" now.
  MainLoop();

  if (!EmitObject)
    return 0;

  // Select the host target and configure its object-file emitter.
  std::string TargetTriple = TargetTripleOption.empty()
                                 ? llvm::sys::getDefaultTargetTriple()
                                 : llvm::Triple::normalize(TargetTripleOption);
  std::string Error;
  const llvm::Target *Target =
      llvm::TargetRegistry::lookupTarget(TargetTriple, Error);
  if (!Target) {
    llvm::errs() << Error << '\n';
    return 1;
  }

  llvm::TargetOptions Options;
  llvm::CodeGenOptLevel CodeGenOpt;
  switch (OptLevel) {
  case '0':
    CodeGenOpt = llvm::CodeGenOptLevel::None;
    break;
  case '1':
    CodeGenOpt = llvm::CodeGenOptLevel::Less;
    break;
  case '2':
    CodeGenOpt = llvm::CodeGenOptLevel::Default;
    break;
  case '3':
    CodeGenOpt = llvm::CodeGenOptLevel::Aggressive;
    break;
  }
  std::unique_ptr<llvm::TargetMachine> TargetMachine(
      Target->createTargetMachine(llvm::Triple(TargetTriple), "generic", "",
                                  Options, llvm::Reloc::PIC_, std::nullopt,
                                  CodeGenOpt));
  if (!TargetMachine) {
    llvm::errs() << "Could not create the target machine\n";
    return 1;
  }

  // Lower the complete MLIR module once, then attach the target information
  // required to produce a native object file.
  auto Lowered = ExitOnErr(lowerToLLVM(TargetMachine->createDataLayout()));
  Lowered.Module->setTargetTriple(llvm::Triple(TargetTriple));

  llvm::SmallString<256> Filename;
  if (OutputFilename.empty()) {
    Filename = InputFilename;
    llvm::sys::path::replace_extension(Filename, "o");
  } else {
    Filename = OutputFilename;
  }
  std::error_code EC;
  llvm::raw_fd_ostream Dest(Filename, EC, llvm::sys::fs::OF_None);
  if (EC) {
    llvm::errs() << "Could not open " << Filename << ": " << EC.message()
                 << '\n';
    return 1;
  }

  llvm::legacy::PassManager EmitPM;
  if (TargetMachine->addPassesToEmitFile(EmitPM, Dest, nullptr,
                                         llvm::CodeGenFileType::ObjectFile)) {
    llvm::errs() << "Target machine cannot emit an object file\n";
    return 1;
  }

  EmitPM.run(*Lowered.Module);
  Dest.flush();
  llvm::outs() << "Wrote " << Filename << '\n';

  return 0;
}

10.13.3 Dialect Definitions

KaleidoscopeOps.td

// KaleidoscopeOps.td
//
// Defines the Kaleidoscope dialect: a small set of operations that preserve
// source-level variable semantics (name, location, argument number) in the IR
// until lowering is ready to turn them into stack storage and debug info.

// Provides the TableGen definitions for MLIR dialects, types, and operations.
include "mlir/IR/OpBase.td"

def Kaleidoscope_Dialect : Dialect {
  // The prefix for every operation and type in this dialect. An operation
  // with the mnemonic `var` will therefore print as `kaleidoscope.var`.
  let name = "kaleidoscope";

  // Where the generated C++ classes live. Our classes end up in the
  // `mlir::kaleidoscope` namespace.
  let cppNamespace = "::mlir::kaleidoscope";

  let summary = "Operations that preserve Kaleidoscope variable semantics";

  // Ask MLIR to generate the parser and printer for our types based on the
  // assembly format each type declares. We only have one type, and its format
  // is defined below.
  let useDefaultTypePrinterParser = 1;
}

// A common base for every operation in this dialect. Each concrete operation
// supplies its own mnemonic and, optionally, a list of traits describing its
// behavior.
class Kaleidoscope_Op<string mnemonic, list<Trait> traits = []>
    : Op<Kaleidoscope_Dialect, mnemonic, traits>;

// A type representing a mutable source variable.
//
// TypeDef generates the C++ class `mlir::kaleidoscope::VariableType` from the
// class stem "Variable", and the mnemonic "variable" gives it the textual
// spelling `!kaleidoscope.variable`.
//
// The type deliberately says nothing about how the variable is stored. That
// is a lowering decision, not a property of the source language.
def Kaleidoscope_VariableType
    : TypeDef<Kaleidoscope_Dialect, "Variable"> {
  let mnemonic = "variable";
  let summary = "a mutable Kaleidoscope variable";

  // An empty assembly format means the type has no parameters to print after
  // its mnemonic, so it always appears as the bare `!kaleidoscope.variable`.
  let assemblyFormat = "";
}

// Declares and initializes a mutable source variable.
//
// This generates the C++ class `kaleidoscope::DeclareOp`, printed in IR as
// `kaleidoscope.var`. It is the operation that carries the source name and
// argument number that would otherwise be lost the moment a variable became
// an anonymous allocation.
def Kaleidoscope_DeclareOp : Kaleidoscope_Op<"var", []> {
  let summary = "declare and initialize a mutable source variable";

  // `(ins ...)` lists everything the operation takes in. The `$` names
  // generate C++ accessors (getInitialValue, getName, getArgumentNumber).
  //
  //   - F64:$initialValue   is an SSA operand, constrained to f64
  //   - StrAttr:$name       is an attribute holding the source name
  //   - I64Attr:$argumentNumber is an attribute; 0 means a local variable,
  //                         and 1+ means a function parameter (DWARF numbers
  //                         parameters starting at 1)
  let arguments = (ins F64:$initialValue, StrAttr:$name,
                       I64Attr:$argumentNumber);

  // The SSA result represents the *variable itself*, not the floating-point
  // value currently stored in it. Subsequent read and assign operations use
  // this value to refer to the variable.
  let results = (outs Kaleidoscope_VariableType:$variable);

  // The textual syntax. Backticks contain literal punctuation, `$name` and
  // `$initialValue` refer to the fields above, and `type(...)` prints the type
  // of the named operand or result. `attr-dict` prints any attributes not
  // already consumed by the format.
  //
  // This produces, for example:
  //   %0 = kaleidoscope.var "x"(%arg0 : f64)
  //       {argumentNumber = 1 : i64} : !kaleidoscope.variable
  let assemblyFormat = "$name `(` $initialValue `:` type($initialValue) `)` attr-dict `:` type($variable)";
}

// Reads the current value of a mutable source variable.
//
// Not marked `Pure`: two reads of the same variable are not necessarily equal,
// since an assignment may occur between them. Marking it pure would let CSE
// incorrectly collapse the two reads into one.
def Kaleidoscope_ReadOp : Kaleidoscope_Op<"read", []> {
  let summary = "read a mutable source variable";

  // Consumes the variable; produces the f64 value currently stored in it.
  let arguments = (ins Kaleidoscope_VariableType:$variable);
  let results = (outs F64:$value);
  let assemblyFormat = "`(` $variable `:` type($variable) `)` attr-dict `:` type($value)";
}

// Assigns a new value to a mutable source variable.
//
// The operation itself produces no SSA result. The AST layer returns the
// assigned value separately so that an assignment expression can be used as a
// subexpression (as in `(y = y + 1) * y`).
def Kaleidoscope_AssignOp : Kaleidoscope_Op<"assign", []> {
  let summary = "assign a new value to a mutable source variable";

  // Consumes the variable and the replacement value. Print both operand types
  // explicitly so it is immediately clear that the value is an f64 and the
  // destination is a !kaleidoscope.variable.
  let arguments = (ins Kaleidoscope_VariableType:$variable, F64:$value);
  let assemblyFormat = "`(` $value `:` type($value) `,` $variable `:` type($variable) `)` attr-dict";
}

10.13.4 Dialect Declaration

KaleidoscopeDialect.h

// KaleidoscopeDialect.h

#ifndef KALEIDOSCOPE_DIALECT_H
#define KALEIDOSCOPE_DIALECT_H

#include "mlir/Bytecode/BytecodeOpInterface.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"

// Generated declaration of mlir::kaleidoscope::KaleidoscopeDialect.
#include "KaleidoscopeDialect.h.inc"

// KaleidoscopeTypes.h.inc contains several selectable sections. Defining this
// macro asks it to emit the generated type class declarations at this include.
#define GET_TYPEDEF_CLASSES
#include "KaleidoscopeTypes.h.inc"

// Likewise, select the generated operation class declarations from the
// operation header fragment.
#define GET_OP_CLASSES
#include "KaleidoscopeOps.h.inc"

#endif

10.13.5 Dialect Implementation

KaleidoscopeDialect.cpp

// KaleidoscopeDialect.cpp

#include "KaleidoscopeDialect.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/DialectImplementation.h"
#include "llvm/ADT/TypeSwitch.h"

using namespace mlir;
using namespace mlir::kaleidoscope;

// Generated definitions for the dialect class declared by the matching header
// fragment.
#include "KaleidoscopeDialect.cpp.inc"

// Select the generated C++ definitions for our TypeDef and Op records.
#define GET_TYPEDEF_CLASSES
#include "KaleidoscopeTypes.cpp.inc"

#define GET_OP_CLASSES
#include "KaleidoscopeOps.cpp.inc"

void KaleidoscopeDialect::initialize() {
  // The same generated .inc files also contain lists of every type and
  // operation in the dialect. These macros select those lists so the dialect
  // can register all generated classes with MLIR.
  addTypes<
#define GET_TYPEDEF_LIST
#include "KaleidoscopeTypes.cpp.inc"
      >();
  addOperations<
#define GET_OP_LIST
#include "KaleidoscopeOps.cpp.inc"
      >();
}

10.13.6 Debug-Scope Pass Declaration

KaleidoscopeDebugInfo.h

// KaleidoscopeDebugInfo.h

#ifndef KALEIDOSCOPE_DEBUG_INFO_H
#define KALEIDOSCOPE_DEBUG_INFO_H

#include "mlir/Pass/Pass.h"
#include "llvm/ADT/StringRef.h"
#include <memory>
#include <string>

/// Create the module pass that attaches compile-unit and function debug scopes
/// before the LLVM dialect is translated to LLVM IR.
///
/// Chapter 10 no longer needs a separate parameter-name map here: source
/// variable names are preserved by the Kaleidoscope dialect and consumed by
/// its lowering pass.
std::unique_ptr<mlir::Pass>
createKaleidoscopeDebugInfoPass(llvm::StringRef inputFilename, char optLevel);

#endif

10.13.7 Debug-Scope Pass Implementation

KaleidoscopeDebugInfo.cpp

// KaleidoscopeDebugInfo.cpp
//
// Constructs the compile-unit and function-scope debug metadata that later
// lowering and LLVM IR translation will attach to instructions. Variable-
// level debug info is produced separately by the variable-lowering pass.

#include "KaleidoscopeDebugInfo.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/IR/BuiltinOps.h"
#include "llvm/BinaryFormat/Dwarf.h"
#include "llvm/Support/Path.h"

using namespace mlir;

namespace {

class KaleidoscopeDebugInfoPass
    : public PassWrapper<KaleidoscopeDebugInfoPass, OperationPass<ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(KaleidoscopeDebugInfoPass)

  KaleidoscopeDebugInfoPass(StringRef inputFilename, char optLevel)
      : inputFilename(inputFilename.str()), optLevel(optLevel) {}

  StringRef getArgument() const final { return "kaleidoscope-debug-info"; }
  StringRef getDescription() const final {
    return "Add Kaleidoscope compile-unit and function debug info";
  }

  void runOnOperation() override {
    ModuleOp module = getOperation();
    MLIRContext *context = module.getContext();

    // A DWARF file entry is a (directory, filename) pair, not a single path.
    // Input read interactively from the REPL has no real path, so we invent
    // a synthetic <stdin> entry to describe it.
    StringRef inputPath = inputFilename;
    auto file = inputPath.empty()
                    ? LLVM::DIFileAttr::get(context, "<stdin>", "")
                    : LLVM::DIFileAttr::get(
                          context, llvm::sys::path::filename(inputPath),
                          llvm::sys::path::parent_path(inputPath));

    // Create the top-level debug record for this translation unit. We wrap
    // it in a DistinctAttr so it has identity: two compile units with
    // otherwise identical fields must not be uniqued into one, which would
    // be invalid DWARF.
    auto compileUnit = LLVM::DICompileUnitAttr::get(
        DistinctAttr::create(UnitAttr::get(context)), llvm::dwarf::DW_LANG_C,
        file, StringAttr::get(context, "Kaleidoscope"),
        /*isOptimized=*/optLevel != '0', LLVM::DIEmissionKind::Full);

    // MLIR carries debug metadata on locations rather than in a separate
    // side table. Fusing the compile unit into the module's existing
    // location is how we attach it so that LLVM IR translation can recover
    // both the original source location and the new metadata.
    module->setLoc(FusedLoc::get(context, {module.getLoc()}, compileUnit));

    for (LLVM::LLVMFuncOp function : module.getOps<LLVM::LLVMFuncOp>()) {
      // Prefer the function's own file and line. Fall back to the module's
      // input file and line 1 for compiler-generated functions that carry
      // no concrete source location.
      Location originalLoc = function.getLoc();
      LLVM::DIFileAttr functionFile = file;
      int64_t line = 1;
      if (auto fileLoc = originalLoc->findInstanceOf<FileLineColLoc>()) {
        StringRef functionPath = fileLoc.getFilename().getValue();
        functionFile = LLVM::DIFileAttr::get(
            context, llvm::sys::path::filename(functionPath),
            llvm::sys::path::parent_path(functionPath));
        line = fileLoc.getLine();
      }

      // A definition and an external declaration need different DISubprogram
      // configurations. A definition belongs to this compile unit, carries a
      // distinct identity so it won't be merged with an identical-looking
      // subprogram, and is flagged as a Definition. A declaration has no
      // body emitted by this compile unit, so it gets neither the attachment
      // nor the flag.
      DistinctAttr id;
      LLVM::DICompileUnitAttr functionCompileUnit = compileUnit;
      auto flags = static_cast<LLVM::DISubprogramFlags>(0);
      if (optLevel != '0')
        flags = flags | LLVM::DISubprogramFlags::Optimized;
      if (function.isExternal()) {
        functionCompileUnit = {};
      } else {
        id = DistinctAttr::create(UnitAttr::get(context));
        flags = flags | LLVM::DISubprogramFlags::Definition;
      }

      // Create the function's lexical debug scope. A complete source-level
      // function signature is outside this tutorial's scope, but a
      // DISubprogram still requires a subroutine-type metadata node, so we
      // supply an empty one. Individual parameter variables receive their
      // types later through their debug declarations; the function signature
      // itself remains unspecified.
      auto functionType = LLVM::DISubroutineTypeAttr::get(
          context, llvm::dwarf::DW_CC_normal, {});
      auto name = function.getNameAttr();
      auto scope = LLVM::DISubprogramAttr::get(
          context, id, functionCompileUnit, functionFile, name, name,
          functionFile, line, line, flags, functionType,
          /*retainedNodes=*/{}, /*annotations=*/{});

      // Preserve the original source location while attaching the subprogram
      // scope. Subsequent debug-info lowering and LLVM IR translation use
      // this scope when they need to describe instructions in the function
      // body, so it must be reachable from the function's location.
      function->setLoc(FusedLoc::get(context, {originalLoc}, scope));
    }
  }

private:
  // The pass owns copies of these because execution of the pass manager is
  // not guaranteed to finish within the lifetime of the caller's arguments.
  std::string inputFilename;
  char optLevel;
};

} // namespace

// Hide the concrete pass class behind a factory so the driver only needs the
// standard mlir::Pass interface.
std::unique_ptr<Pass> createKaleidoscopeDebugInfoPass(StringRef inputFilename,
                                                      char optLevel) {
  return std::make_unique<KaleidoscopeDebugInfoPass>(inputFilename, optLevel);
}

10.14 Wrapping Up

Our dialect is small: it has only three operations and one type. It remembers useful details about each variable, like its name, where it appears in the source code, and whether it is a function parameter. We keep those details around until we create the lower-level code, which makes it easier to generate storage and debugging information correctly.

If you want to take this further, the next natural step would be moving the entire Kaleidoscope AST into a dialect. The MLIR Toy tutorial shows what that architecture looks like. But for the problem we set out to solve, three operations were enough.

10.15 Download the Source

Download the complete source code on GitHub.

< 9. Kaleidoscope: Adding Debug Information

11: Conclusion and Other Useful MLIR Tidbits >

10: Building Your Own Dialect

10.1 The Problem We're Solving

10.2 Round-tripping MLIR

10.3 Define the Dialect

10.4 Define the Variable Type

10.5 Add an Operation Base Class

10.6 Define kaleidoscope.var

10.7 Define kaleidoscope.read

10.8 Define kaleidoscope.assign

10.9 Generate and Include the C++ Classes

10.10 Lower the Dialect to the LLVM Dialect

10.11 Wire Up the Pass Pipeline

10.12 Try It Out

10.13 Where the Code Lives

10.13.1 Build Configuration

10.13.2 Compiler

10.13.3 Dialect Definitions

10.13.4 Dialect Declaration

10.13.5 Dialect Implementation

10.13.6 Debug-Scope Pass Declaration

10.13.7 Debug-Scope Pass Implementation

10.14 Wrapping Up

10.15 Download the Source

10.6 Define `kaleidoscope.var`

10.7 Define `kaleidoscope.read`

10.8 Define `kaleidoscope.assign`