LLVM Discussion Forums

[RFC] Better support for dialects that want to make SSA names "load bearing"

Hi all,

I’m working on a dialect where arguments and many operations have a “load bearing” (required for correctness) name associated with them. This is nicely implementable by using an optional “name” attribute that carries a string along with it, e.g.:

xyz.func @my_function(%arg0 : i8 { xyz.name = "first_arg" }) {
  ...
  %42 = xyz.add %40, %arg0 {xyz.name = "foo"} : i32
  %43 = xyz.mul %42, %50 {xyz.name = "bar"} : i32

Note that I’m using “xyz.name” instead of “name” as the attribute name,
because arguments to regions in FunctionLike ops cannot currently have
un-namespaced attributes.

Desired end result

Given the prevalence of these names, it becomes quickly annoying to have to read this. Instead, I would like to get:

xyz.func @my_function(%first_arg : i8) {
  ...
  %foo = xyz.add %40, %first_arg : i32
  %bar = xyz.mul %foo, %50 : i32

It is worth noting that I’m not using any builtin dialects - this behavior will be completely specific to the XYZ dialect, which is one of the really nice things that MLIR enables :-). Making this happens is almost completely possible with existing MLIR support, but there are a few rough edges, and it required writing too much code in some places. As such, I have a few patches that I need to upstream in order to support this properly.

To help provide context, I thought that I would send this RFC as a head’s up and to allow discussion of this, along with some of my proposed solutions. I’d love feedback or thoughts on alternate solutions to these challenges. Assuming there are no major concerns, I’ll send in some patches to improve these areas.

Here are a few issues I’ve encountered:

Problem: not all strings are valid SSA names

We currently have nice interfaces for dialects/ops to overload the asm names of an operation (e.g. OpAsmDialectInterface). However, when an invalid MLIR symbol name is returned (e.g. “foo bar”, the asmprinter doesn’t notice it and generates invalid IR. The fix is to escape invalid names, which isn’t a big issue.

Problem: These names are not necessarily unique, but need to be preserved

These names are not stored in a symbol table, and they can be duplicates. For example, this
is valid:

  %42 = xyz.add %40, %arg0 {xyz.name = "foo"} : i32
  %43 = xyz.mul %42, %50 {xyz.name = "foo"} : i32

and it would be invalid to produce two %foo results. As it turns out, the AsmPrinter already handles this correctly, and autorenames the second one to %foo_1. However, relying on this breaks round tripping. The solution is fairly simple: if the name is autorenamed, print out to attribute explicitly. In this case, we’d get something like this:

  %foo = xyz.add %40, %arg0 : i32
  %foo_1 = xyz.mul %foo, %50 {xyz.name = "foo"} : i32

Which is perfectly fine - the attribute can be explicitly listed in these (very rare) cases, and the operation can omit the attribute from printing iff the names match.

We need one small extension though, we need a method on OpAsmPrinter that returns the name for a Value, so implementations of the hooks can see whether something is getting autorenamed. This just requires adding one method to OpAsmPrinter.

Problem: Custom op parser needs to be able to query names of op results

When parsing an operation with a meaningful name, the op parser hook has to know what name the current operation had. For example, in order to parse this:

 %foo = xyz.add %40, %arg0 : i32

into an Operation that has an xyz.name attribute, the operation parsing hook has to be able to retrieve the string %foo. Solution: add a method to OpAsmParser, allowing the parser hook to query the names of any results for the current operation.

Suboptimality: Printer/parser logic for FunctionLike ops isn’t very flexible

The FunctionLike trait is really awesome, and provides a lot of functionality out of the box. Unfortunately, it isn’t very parameterizable. For example, a FunctionLike op with an implied terminator cannot use the mlir::impl::printFunctionLikeOp interface, because the printBlockTerminators argument to printRegion is forced to true. FunctionLike ops aren’t required to use this, but duplicating all this code seems sad.

Since other things are currently configurable (e.g. whether variadic functions are allowed by the parser), the simple fix is to widen the interfaces (e.g. pass in a struct of options).

There are other things that need similar extension, e.g. when the argument names match, I don’t want to print out the argument names in the parameter attribute dictionary either.

-Chris

It seems to me that other than bugs fixes or minor changes, the most important aspect to discuss here in what you are proposing is that it offers a way to give a semantic meaning to SSA value name.

You mentioned the change to OpAsmPrinter to access the printed result names in the custom op printer, but unless I’m missing something you will also need a way to access the parsed LHS of an operation in the custom parser: I seem to remember that the OpAsmParser does not expose this (and that it was intentional to have the custom parser be limited to parsing the RHS without context on the LHS).

Can you comment more on the utility of this in the first place? I guess I’m missing why this would be useful.

Hi Stephen,

It is the same reason as any work to improve the ‘pretty’ form of the IR dumper: humans end up looking at IR dumps (e.g. when debugging the compiler) a lot, and easing the ability to write test cases makes it more likely that good tests will be written. Eliminating obvious redundancy from the dump makes both simpler and better.

-Chris

Not quite. This doesn’t change the core IR semantics at all - the names are just stored as a string attribute. I’m discussing providing a sugared form to an already existing in-memory representation.

I don’t think there is a need to access the operands of an instruction like this. I think you just need to know if the “autorenamed” name given to an instruction matches what it thinks it has.

Perhaps I’m missing something big though, can you explain what you’re thinking?

-Chris

Yeah, I get that part. I’m trying to understand what the extra ‘load bearing’ semantics are trying to provide.

I think what Mehdi is referring to is how you populate these names when parsing the IR back in. The format of the operation is intended, at least at the moment, from the ssa names. There is no way ATM to get the names given to the results of an operation during parsing. What you are proposing here would require semantically tying the two together.

On a minor note, I’m trying to understand why the roundtripping breaks when the autorenaming is done correctly. The uses don’t get renamed correctly?

Why even have xyz.name when you have locations? Just have a NamedLoc. If you then need more info, you could always just create a FusedLoc with a NamedLoc and other loc (and metadata to make it clear it is the load bearing name if needed). You’ll notice we have used locations pretty extensively in TensorFlow and TFLite dialect where name attribute was pretty pervasive. It avoids needing to teach CSE, for example, about name attribute.

Oh right, I forgot that part. I added a section above name “Problem: Custom op parser needs to be able to query names of op results”. This can be handled by adding one method to OpAsmParser.

All I meant by “load bearing” (which I admit is not a great term :slight_smile: ) in that the parser and printing hooks for operations have access to the parsed/printed result names of the current operation, allowing them to implement custom behavior. That’s all this RFC boils down to.

Uses get renamed correctly. Recall that these names are semantically meaningful. Therefore if this operation:

%42 = xyz.add %40, %arg0 {xyz.name = "foo"} : i32

Got printed out as:

%foo_42 = xyz.add %40, %arg0 : i32

Then when it is parsed back in, the xyz.name attribute would be set to “foo_42” instead of “foo”. This is what I meant by “breaking round tripping”. The solution for this is to explicitly place the attribute in its attribute dictionary in this narrow case. This eliminates the syntactic niceness, but maintains correctness.

I hadn’t thought about using locations for this, but I don’t think this helps. These are semantically part of the IR, so they should block CSE and other things. Also, locations aren’t printed by default by mlir-opt, which is another bad thing. Furthermore, when locations are enabled, we would end up with the same verbosity problem.

In my view, the pretty syntax features that MLIR provides are a way for op/dialect implementations to eliminate redundancy. This is fundamentally why std.add wants to print a type of “i32” instead of “(i32, i32) -> i32”. Eliminating redundancy increases clarity, increases understandability, and makes it easier to write testcases. My goal here is to do the same thing with a string attribute, encoding it into syntax that already exists.

Makes sense now - thanks!

This is where the names now have semantic meaning. This makes the result assignment part of the custom assembly format, which it currently is not. For example, how would such an operation handle the many different ways of naming a result? Consider the example of a simple variadic operation, there could be zero names, one name, or even multiple names assigned to the result of that operation. What does the parser of the operation do in any of these cases? handle all of them? fail?

A couple potential interactions of the proposal. Would like to know what folks think.

For a while I’ve been wanting the IR to better round-trip names, which makes testing a lot easier.

E.g. currently this function:

func @f(%arg0: f32) {
br ^loop(%arg0: f32)
^loop(%loop: f32):
br ^loop(%loop: f32)
}

Round-trips as:
func @f(%arg0: f32) {
br ^bb1(%arg0 : f32)
^bb1(%0: f32): // 2 preds: ^bb0, ^bb1
br ^bb1(%0 : f32)
}

That makes testing kind of annoying. It would be nice if the names of blocks and block arguments were preserved.

Currently I see great productivity gains by just being able to hardcode arg0 in tests (reduces the amount of filecheck capturing needed), and preserving a custom name there would make tests even more readable. Block names and SSA value names in the body being preserved (as long as we know that the pass under test isn’t touching the associated ops) could really help a lot with making tests easier to write and more readable.

I’m curious if we want to push further in this direction. If we do, then it feels like it would interact with what Chris wants here. I really think this could significantly reduce test maintenance costs in the MLIR ecosystem.

Also, Chris, are the names logically on the ops, or on the Value’s? (or is there an implicit single-result assumption making the question moot?)

What are the semantics of two ops/values with the same name? I’m failing to come up with a load-bearing semantic that would make sense in that case.

On the specific IR required here, I’m also worried that it won’t be obvious at a glance at the IR which names are load bearing vs not. Could we use a different sigil or something to indicate it?

Btw, for some prior art in this domain, IREE’s vm.import op does have load-bearing names on the arguments. (done with a vm.name arg attribute)

In my opinion, the parser should hand it to the operation’s parser implementation (when queried through the new method) and the operation parser can decide what to do with it.

I think it would be a bad idea to do this. LLVM does this, and many people find it to be appealing, but this is a bad tradeoff for production. Those strings are expensive to propagate and need to be uniqued, which makes IR manipulation (the core part of the compiler) slow.

The proposal I’m floating here are just a few tiny extensions that make existing structures more flexible. They don’t have any effect at all on the core semantics of the IR.

As River points out, an asmparser hook that wants to do something with names should be prepared to handle the various spellings. If it is %foo:4 = or %foo,%bar =, then the parser hook should know that.

This should remain a parse error. The operation parser hook should not be able to override core aspects of the IR like dominance rules etc.

How is this different than any other dialect specific asmparser/printer rules? The things we do for region based control flow statements are way more “interpretive” than this :slight_smile:

-Chris