Photo by USGS on Unsplash

A Primer on MLIR’s Symbol Table

One key feature in MLIR is the symbol table, which enables efficient management of named operations like functions, variables, and constants.

In this blog, we’ll explore what MLIR’s symbol table is, how it works, and how you can use it in your custom dialect. By the end, you’ll have a solid grasp of MLIR’s symbol system and be able to use it effectively in your own projects.

Understanding MLIR’s Symbol Infrastructure

1. The Symbol Trait

In MLIR, operations that should be part of the symbol table must be marked with the Symbol trait. For example, in a custom dialect:

def MyDialect_Const : MyDialectOp<"const", [Symbol]> {
  let summary = "Defines a named constant";
  let arguments = (ins SymbolNameAttr:$sym_name);
  let results = (outs AnyType);
}

The [Symbol] trait makes this operation a named entity in MLIR, which means:

• It must have a SymbolNameAttr (sym_name) attribute.
• It can be inserted into a SymbolTable for lookup.

2. The SymbolTable Trait

Operations that hold other symbol operations (like modules or functions) must have the SymbolTable trait:

def MyDialect_FuncOp : MyDialectOp<"func", [IsolatedFromAbove, SymbolTable, Symbol]> {
  let summary = "Defines a function in MyDialect";
  let arguments = (ins SymbolNameAttr:$sym_name, TypeAttrOf<FunctionType>:$funcType);
  let regions = (region AnyRegion:$body);
}

Using the Symbol Table while codegen

1. Creating an MLIR Context and Module

Every MLIR operation lives within a context. We start by creating a module and registering our dialect:

// Create an empty module
OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(&context);

2. Defining a Function

We define a function in our custom dialect and insert it into the module:

auto funcType = builder.getFunctionType({}, {});
auto funcOp = builder.create<mydialect::FuncOp>(builder.getUnknownLoc(), "main", funcType);
module->push_back(funcOp);

Since MyDialect_FuncOp has the SymbolTable trait, it can contain symbols like constants.

3. Adding a Constant to the Function

Let’s insert a named constant (my_var) inside the function:

auto &entryBlock = funcOp.getBody().emplaceBlock();
builder.setInsertionPointToStart(&entryBlock);

// Create a symbol table for the function
mlir::SymbolTable symbolTable(funcOp);

// Define a constant and insert it into the symbol table
auto constOp = builder.create<mydialect::Const>(builder.getUnknownLoc(), builder.getI32Type(), builder.getStringAttr("my_var"));
symbolTable.insert(constOp);

Here:

• We create a Const operation with the symbol name "my_var".
• We insert it into funcOp's symbol table.

4. Looking Up a Symbol

Now, let’s retrieve my_var from the symbol table:

if (auto found = symbolTable.lookup("my_var")) {
    llvm::outs() << "Found: " << found->getName() << "\n";
}

The complete code is available in our GitHub repository: Blog CodeBase Repo.