Retain original identifier names for debugging

MLIR
1

This thread is to discuss how we could enable a debug flag to preserve identifier names when generating IR. I do not believe that existing debug passes allow us to keep the names of identifiers.

Right now, the original names of identifiers is lost, and when printing IR, anonymous names are used. This is intended behaviour, and a sensible default, See LangRef.md - “Identifier names for values may be used in an MLIR text file but are not persisted as part of the IR - the printer will give them anonymous names like %42.”

Therefore, something like:

func.func @add_one(%my_input: f64) -> f64 {
    %my_constant = arith.constant 1.00000e+00 : f64
    %my_output = arith.addf %my_input, %my_constant : f64
    return %my_output : f64
}

:arrow_down: becomes

func.func @add_one(%arg0: f64) -> f64 {
  %cst = arith.constant 1.000000e+00 : f64
  %0 = arith.addf %arg0, %cst : f64
  return %0 : f64
}

However there are situations where it could be useful to preserve names:

  • debugging and development, where a developer has given meaningful variable names to their identifiers and wants to understand what has happened to them
  • Code formatting. See this thread, where I am developing a code formatter for MLIR files. One of the requirements here is that the formatter does not change the names

We could add an option such as --debug-retain-identifier-names, to allow this behavior.

What would be an elegant way to design this? What challenges are there? Are there any existing components that we can leverage to achieve this goal?

In my initial investigation, the original name is still available during parsing, but is lost when the Value is initialised.

The generic names are generated at print time, e.g., in AliasInitializer::generateAlias in IR/AsmPrinter.cpp, or a more concerete implementation SSANameState::numberValuesInBlock where we generate the %argX names for block arguments.

I can imagine a relatively generic check in these printers, something like if keepNames == true and nameIsAvailable(value) == true then print the original identifier name. If not, following the existing execution path, otherwise print the custom/original name.

In terms of how to store the original names, there are two main approaches I can imagine. First, is extending the Value class with the field and access methods.
Second, is having some external mapping, say llvm::DenseMap<Value, llvm::StringRef> valueToAlias.

The first approach has the downside of increasing bloat to all Value objects, but in the second approach there is the question of where the mapping should be stored so that it is easily accessible. MLIRContext is the most global object I’m aware of that could fit this requirement.

Another question is how to handle the anonymous naming of newly introduced values. E.g., if we add a pass which transforms the code, we may create some new values. If a SSA would otherwise be %2, but the previous SSAs have been given names, should this be %0 (the first anonymous SSA) or %2 (the third SSA)? We need to ensure that our generated names do not clash with named that the user has defined (since the user could specify names which could also be generated by our alias generators).

Any thoughts or comments?

2

I’d consider that to be what locations are for. Unfortunately they aren’t printed by default so they aren’t even present in your example. But locations already allow for tracking fusions, call stacks and annotating origin with other metadata. It’s not sufficient in some cases where listeners would probably be needed in addition.

Wouldn’t these be built on top of AsmParserState (llvm-project/mlir/include/mlir/AsmParser/AsmParserState.h at b0b491d458962136c696366b8cf535d54511baf3 · llvm/llvm-project · GitHub) rather than be in the IR? (Probably along with some prepopulated SSAState AsmState variant). And so these could be outside the context.

3

Yes, locations can be helpful for understanding the relationship between the original source, and the processed source. However, I believe that locations still require cross-referencing on the part of the developer, and thus I can imagine situations where say the output retaining the original names would be more helpful (since it reduces the need to manually cross-reference). Maybe I could get an implementation of identifier retention by using location information, but my experience thus far with that side of things suggests that might not be the way (e.g., the column numbers don’t always match up with what I’m looking for)

I’m still hammering out the design of the mlir-format tool (I have an initial version of comment retention working at time of writing). Initially, I’m going through the whole IR generation stage, which is more straightforward because it hooks in with the custom IR printers infrasturue already (I’m essentially just working with a modified mlir-opt right now). mlir-format stopping after parsing might make more sense design-wise, but perhaps that discussion is better suited for that thread.

4

You can see in my branch retain-names-dev, I’ve implemented this behaviour in around a dozen lines (without a flag implemented, this is just default behaviour in this branch).

I do my name replacement in AsmPrinter.cpp’s SSANameState::printValueID. I believe this is the most generic place where this replacement can occur.

I am keeping track of the original names with llvm::DenseMap<Operation *, llvm::StringRef> in MLIRContext. Values are inserted at parse time, and I’m using value->getDefiningOp() as the key.

Since the original name information needs to be accessible from the printer, as I said above, I think this should either be a property of the Value class, or in a map in the MLIRContext.

As @jpienaar mentioned, the location info is already in the Value object, so perhaps this kind of debug information is acceptable to include there. I suppose I could parse the original name from the location information at print time, but for now I’m storing it when we parse it the first time. The tradeoff is either adding a new set of data to keep track off, or increasing the code complexity to parse at print time from the location. I’d be interesting in hearing opinions from more experienced folk.

As you can see, to actually support the functionality of retained variable names only requires a small amount of code. The question is where that logic should reside, and if having this as a feature at all in e.g., mlir-opt would be useful.

5

This is a reasonable request, and it has been discussed before. Beyond your usecase, some folk want to tag an operation early in an optimization pipeline and have the name propagate through to ease debugging a compiler pipeline, for example.

There are big questions around what happens during IR transformations: how do the names avoid conflict (eg when an operation is duplicated) and how do you prevent invalidation when an operation is deleted (invalidating the operation*), also how do we make this efficient? There are no perfect answers here, and there is a lot of damage in the LLVM IR implementation of this (which does online autorenaming and a bunch of other complicated stuff).

If you’re interested in pursuing this, I would suggest a different approach. I’d invent a new magic attribute “_ssaName” (or something like that, stored in the existing attribute dictionary on the operation) which can either be a stringattr or an array of strings (for multiple results). Enhance the MLIR asmprinter to check for this attribute after checking to see if an op has an OpAsmOpInterface, but before generating an anonymous one. If present, treat it just as if OpAsmOpInterface returned the name - this means that the uniquing happens only in the asmprinter, not online during the compiler execution.

You can then go ahead and add your MLIR parser flag to produce these - please make it opt-in, not on by default of course!

Given an approach like this, I think you’d get something pretty useful, general, and low impact on the compiler.

-Chris

6

I’d be more interested by a debug mode (build time config?) where we could have the LLVM functionality than anything else.

Do you have a list of the things to “fix” from the LLVM IR implementation?
(which is enabled/disabled at runtime with a flag in the LLVMContext at the moment)

You’ll have to think about Operations being deleted / cloned, and the fact that the compiler is multi-threaded. It’s not clear to me that the MLIRContext is the right place to keep pointers to the IR from a design point of view. At least right now it does not: it is responsible for the immutable objects storage (Types and Attributes), the diagnostic handling, the thread pool, and dialect registrations.

7

The big issue is that it has pervasive cost, and (as you say) does not work with multithreading.

The approach I’m describing above does not - it only puts cost into the asmparser/printer and is full thread safe.

-Chris

8

Thanks for the thoughts both, I’ve got an updated implementation based on your comments, and have added a parser flag.

I’ve shifted the data storage from MLIRContext to the Value base class (with an _ssaValue variable). I hadn’t been considering the multi-threading issue of MLIRContext, and I also think it’s conceptually cleaner for it to just be an attribute of the Value class.

I’ve also pushed the name fetching logic to be earlier, before we gain an anonymous name, and thus we don’t increment the anonymous name counter (is this desirable, or should we increment it anyway for consistency?).

Compared to my previous version, where I fetched the custom SSA name once before printing the value, in this new version I need to fetch the name in two cases: SSANameState::numberValuesInOp and SSANameState::setValueName.

I still need to check if I’m covering all identifiers, e.g., basic blocks.

I’ll try and make a pull request tomorrow (tidying stuff up, checking all the identifiers, and whatnot)

9

I’m not so sure about the multi-threading problem with an implement like the LLVM one: since we’re talking about SSA value and the associated scope of an isolated region it seems to me that it matches perfectly the granularity of MLIR multi-threading.

As of the cost: we should look into this, especially as a debugging feature. There is to consider, I think, these two aspects:

  • overhead of the feature when it is disabled (for example LLVM has runtime-disablement of the feature: what’s the production-path cost of this?)
  • “design” complexity: how much complexity this all add to the infrastructure, and all the client code using it (it this well isolated and invisible to the client code, or does it require pervasive changes to all the users of the IR).
10

hi @Wheest, I’m sorry for the confusion. I mean that the name would be added to the /existing/ attribute dictionary that every operation carries.

I’m not so sure about the multi-threading problem with an implement like the LLVM one: since we’re talking about SSA value and the associated scope of an isolated region it seems to me that it matches perfectly the granularity of MLIR multi-threading.

I’m think you and I are imagining your suggestion as different implementation approaches. Can you explain it in more detail?

The approach out outline above has /zero cost/ when it isn’t used, except in the parser/printer. I don’t see any downside to it at all.

-Chris

11

I was referring to “providing the same feature as LLVM” somehow, isn’t it what you were answering to with “it has pervasive cost, and (as you say) does not work with multithreading”?
(I wasn’t addressing your suggestion at all, I have other concerns with it but not “performance” oriented)

12

Here’s a updated version of this feature, so that the SSA name is stored as an attribute (in the attr-dict sense) of the operation, rather than a member variable of the Value itself. This avoids adding any storage cost unless we enable it.

The printer has two passes, first using SSANameState where we assign our SSA names, and second with OperationPrinter where we actually print things. Previously I was replacing my SSA names in the second stage.

This works for the names of results of operations, but it unfortunately doesn’t cover arguments in function calls, (e.g., func.func @add_one(%my_input: f64)).

These aren’t the result of an Operation, and thus they lose their name, since we don’t have an attribute for it.

The FunctionOpInterface defines argAttrs, which we can use to store attributes of the arguments, which could be the solution here.

However, if I try and dyn_cast<FunctionOpInterface> for my functions within AsmPrinter, the symbols are undefined. This is because we would have a circular dependency — the MLIRFunctionInterfaces library depends on the MLIRIR library (where ASMPrinter is built), so we can’t make MLIRFunctionInterfaces a dependency of MLIRIR. This precludes extracting our argAttrs from our op after casting to a FunctionOpInterface, since we can access those methods.

I think I could avoid this problem by just making the SSA names of arguments an additional attribute of the operation (e.g., _ssaNameArgs), but I feel that it makes more conceptual sense for these names to be in the argAttrs dict defined by FuncOp-style operations. Thoughts?

13

Also there is the question how we store identifiers for other structures, say basic blocks. We lose information for caret identifiers, e.g., ^bb_foo will become ^bb1.

Also with blocks, we also have a similar problem to FuncOp-style operations, namely arguments. However for blocks, and block arguments, we don’t have any attribute dicts we can use as far as I can tell.

I think that every operation would be owned by a block, so a solution could be to again store the block name (and block arg names) as an attribute of (at least one) operation owned by the block.

14

BTW, there is an OpAsmInterface, which allows operation to provide custom names for results, block arguments and blocks. Can it be reused somehow (eg. as external implementation)?

15

Operation that wants to opt-in this behavior can use the OpAsmInterface indeed, there is also already the methods to get the SSA value names from the parser. So the name can be preserved on an opt-in basis, but I believe the goal here would be to preserve them transparently always (which runs a bit at odd with the OpAsmInterface which also can define a name at printing time by the way).

I have some concerns about this behavior: the attributes dictionary is load-bearing, it is part of the IR, and while I understand the convenience I don’t think that something “cosmetic” like SSA value name should be mixed up there. Changing a value name in a text shouldn’t change the in-memory IR!
This would technically impact things like CSE or other transformation looking at attributes for example.

16

I’ll take a look to try and understand the logic of OpAsmInterface and how it handles this case. If OpAsmInterface allows this behaviour, it could be handy, and maybe give a more general solution (i.e., function arguments, blocks, etc).

If there are cases where the printing overrides custom names via OpAsmInterface right now, perhaps to force preservation I could do something like prefixing names during parsing with a magic value (e.g., name = "forcePreserve_" + parsedSsaName;). The printer could check for this, and ensure the printer doesn’t define a new name.

Thinking back to the suggestion about leveraging location information, I’m beginning to see the appeal of re-parsing this at print time. Though it still leaves open the question of how we pass the flag to keep the original identifier names to the printer, as well as the additional complexity of handling operations which have gone through multiple stages of transformation.

Yes, the approach I was taking was to delete the attribute before printing, but if we first pass our IR through several layers of transformation we could end up with unexpected behaviour.

As an illustrative example, when looking at adding the names to the argAttrs of FuncOps, I had to contend with the verifyTrait (Interfaces/FunctionInterfaces.h#L182-L186), which ensured that attributes meet certain conditions. Since dialects can implement arbitrary trait verification, who knows what could happen with the _ssaName attribute.

17

I was approaching this from the position of being “useful” not complete or perfect. As Mehdi points out “re: load bearing”, nothing in this category can be complete or perfect. LLVM IR isn’t either, eg across IR mutations.

I do agree it is useful to go beyond instruction names, e.g. to support function arguments and incoming block arguments for structured statements like “functional if” statements etc. Two thoughts:

  1. You could special case the first block of an operation, and have another “_first_block_arg_names” sort of attribute that you put on the enclosing op. This would cover function arguments and many other common cases.

  2. You could then either a) extend that approach to additional blocks in general, or b) introduce an interface that dialect authors need to opt into for this, or c) find an existing “control flowy” dialect that it close enough that could be built off.

With respect to block arguments, we have to be extra careful to make the code stable to errors. Unlike operation results, block arguments can be mutated/added/deleted and the name list will drift. We wouldn’t want the compiler to crash because a block arg get added or deleted, even though the names would be hopefully out of date.

In any case, my opinion on this topic in general is that we can provide something that is “imperfect but still useful in practice”. I’m not aware of an “llvm equivalent” solution that can be built without impact on MLIR core users who don’t want this functionality.

-Chris

18

Ah, leveraging location information is another good idea. The problem with that is that “preserving names” would break clients of location info :slight_smile:

As an illustrative example, when looking at adding the names to the argAttrs of FuncOps, I had to contend with the verifyTrait (Interfaces/FunctionInterfaces.h#L182-L186), which ensured that attributes meet certain conditions. Since dialects can implement arbitrary trait verification, who knows what could happen with the _ssaName attribute.

While it may be imperfect MLIR officially sanctions installing “other dialect attributes” on ops. Maybe the way to go is to give it something like an “asmprinter” dialect, e.g. naming this “mlir.ssaname” (or somesuch) instead of “_ssaName”. If this doesn’t work with an op, then other dialect attributes won’t either. It could be a good forcing function for people to fix incorrect verifiers.

-Chris

19

Using an attribute approach is going to have a relatively small code footprint and set of dependencies, and thus could be stripped out and changed more easily should a better approach become possible if new core features emerge.

The ways in which using attributes could be a problem have been discussed above, but:

  1. I agree that collecting these attributes under a dialect (e.g., mlir.ssaName) would create a clearer and more consistent way for storing these (and possible future) “magic” attributes that are consumed by some part of the core infrastructure. It also means that verifyTrait functions don’t need to add a special case (see my comment above). It still assumes they at least allow dialect attributes, which I think they always should.
  2. Although care should be taken to not crash if ops are erased, names conflict, etc, actually deployed compilers are not going to use this feature, it is probably just going to be a dev/debug tool. Thus the requirements for strict correctness are not as strong. Refinement will be needed to minimise debug confusion, but naming errors != logic errors.

I think the downsides of the attributes approach have been quite clearly discussed in this thread. I’m going to go for a complete version that leverages this, and if more design issues emerge that are unclear I’ll discuss them here.

Regarding the suggestion around OpAsmInterface, I have the following notes:

  • As I understand it, this is the way in which a given dialect’s ops can override the default naming scheme. See IR/OpAsmInterface.td#L42-L57
  • An example of this override can be seen in for arith.constant, which calls its outputs %cst (see lib/Dialect/Arith/IR/ArithOps.cpp#L155-L175).
  • You can see this function being called in AsmPrinter here (note that the argument we pass is a lambda function which actually registers the name for printing, seen in L1539-L1547).
  • As you can see, this means that the OpAsmInterface doesn’t have any dynamic storage that we could leverage in the parser for custom names. Instead, we sub-classed versions of a function for operations (i.e., the custom naming scheme is defined at LLVM compile time — we want something that can vary per-operation instance at parse time of a given MLIR input).
  • We could extend the OpAsmInterface with a method that allows names to be set at parse time, and then retrieved at print time. However, this has the downside of introducing a constant overhead to all operations all the time, even when this feature is not being used.
20

Okay, I’ve got an implementation on a fresh branch (retain-names-dev-2).

In this design I’m storing everything in the attribute dictionary of operations. I’m storing these attributes under the mlir. namespace for now. Right now I support:

  • Operation result name (mlir.ssaName, still need to support multiple results)
  • The name of the owning block (mlir.blockName, still need to cover blocks without operations, see below)
  • Owning block arguments (mlir.blockArgNames).

I still haven’t enumerated every identifier in MLIR, but will do this in further commits.

One case that I know I’m not handling right now is the names for blocks that don’t have an operation.

If we take the following input (adapted from the example in LangRef.md):

IR showing branches, returns, and block arguments: 
func.func @simple(i64, i1) -> i64 {
^bb_alpha(%a: i64, %cond: i1): // Code dominated by ^bb_alpha may refer to %a
  cf.cond_br %cond, ^bb_beta, ^bb_gamma

^bb_beta:
  cf.br ^bb_delta(%a: i64)    // Branch passes %a as the argument

^bb_gamma:
  %b = arith.addi %a, %a : i64
  cf.br ^bb_delta(%b: i64)    // Branch passes %b as the argument

// ^bb_gamma receives an argument, named %c, from predecessors
// and passes it on to bb4 along with %a. %a is referenced
// directly from its defining operation and is not passed through
// an argument of ^bb_gamma.
^bb_delta(%c: i64):
  cf.br ^bb_eps(%c, %a : i64, i64)

^bb_eps(%d : i64, %e : i64):
  %f = arith.addi %d, %e : i64
  return %f : i64   // Return is also a terminator.
}

We generate the following output :arrow_down: (I highlight names which have been changed with :orange_circle: the first time):

Processed IR showing branches, returns, and block arguments: 
func.func @simple(%arg0🟠: i64, %arg1🟠: i1) -> i64 {
  cf.cond_br %arg1, ^bb1🟠, ^bb_gamma
^bb1:  // pred: ^bb0
  cf.br ^bb3🟠(%arg0 : i64)
^bb_gamma:  // pred: ^bb0
  %b = arith.addi %arg0, %arg0 : i64
  cf.br ^bb3(%b : i64)
^bb3(%0🟠: i64):  // 2 preds: ^bb_gamma, ^bb1
  cf.br ^bb_eps(%0, %arg0 : i64, i64)
^bb_eps(%d: i64, %e: i64):  // pred: ^bb3
  %f = arith.addi %d, %e : i64
  return %f : i64
}

I think I could handle this case if I add attributes to the cf statements, but these aren’t parsed by OperationParser::parseOperation(). I’ll look under the hood at how those are handled.

You can see in the below hidden block what the attributes look like if I don’t cleanup:

IR without name attributes removed 
func.func @simple(%arg0: i64, %arg1: i1) -> i64 {
  cf.cond_br %arg1, ^bb1, ^bb_gamma
^bb1:  // pred: ^bb0
  cf.br ^bb3(%arg0 : i64)
^bb_gamma:  // pred: ^bb0
  %b = arith.addi %arg0, %arg0 {mlir.blockArgNames = [], mlir.blockName = "^bb_gamma", mlir.ssaName = "b"} : i64
  cf.br ^bb3(%b : i64)
^bb3(%0: i64):  // 2 preds: ^bb_gamma, ^bb1
  cf.br ^bb_eps(%0, %arg0 : i64, i64)
^bb_eps(%d: i64, %e: i64):  // pred: ^bb3
  %f = arith.addi %d, %e {mlir.blockArgNames = ["d", "e"], mlir.blockName = "^bb_eps", mlir.ssaName = "f"} : i64
  return %f : i64
}