[RFC] Poison semantics for MLIR

Summary

We propose new mechanisms to model deferred undefined behavior in MLIR using the notion of poison. The goal is to allow dialects to better define their op semantics and reason about it in transformations. Initially, we plan to use it in ‘LLVM’, ‘arith’, and ‘vector’ and motivate this RFC with a few examples of arith and vector ops of choice. The exact use of this model in arith, vector, and other dialects is outside of the scope of this RFC.

This proposal generalizes the notion of poison / freeze from LLVM IR [1, 2, 3] so that it can be applied to an open set of types and ops. Poison semantics are opt-in on a per-type/op basis and should not break existing dialects.

We also propose a new dialect ‘ub’ that provides basic poison-related ops for use in documentation, tests, and some transformations. It includes the following ops: ub.poison, ub.freeze, ub.unreachable. Dialects can define their own equivalent ops using the proposed interfaces.

Background

Not all ops have well-defined behavior for all inputs values, for example:

  1. Integer addition: arith.addi can be defined for all possible inputs as mathematical addition of the operands modulo the 2^N where N is the bit-width.
  2. Right shift: arith.shrui can require the shift value to be less than the bit-width, and have no meaningful and consistent interpretation otherwise.
  3. Integer division: arith.divui requires the divisor to be non-zero, and the underlying division instruction may trap.

While A and C can already be reasoned about with existing constructs (e.g., the ConditionallySpeculatable op interface), we need some formalized framework that allows us to precisely define what happens for B. when preconditions of ops are not satisfied and how the unspecified results propagate and interact with ops that use them [4].

This framework will enable us to employ more speculative execution and related transformations (e.g., hoisting, duplication) in a way that is consistent and easy to reason about.

Semantics

We propose to introduce poison as an extra ‘virtual’ state that types define and ops mutate. Poison representation is a property of types which affects op semantics. Poison state allows us to define ops that may not always return values with a well-defined state.

To make the presentation easier and gradually build up the semantics, we consider 4 cases:

  • Builtin scalars, modeled after poison in LLVM
  • Builtin vectors, modeled similarly to LLVM
  • Custom types
  • Compound types consisting of arbitrary subtypes

Scalar Case

We consider the poison state to be an extra ‘virtual’ bit. For example, we would say that i32 has 2^32 possible well-defined values and one extra value called poison which is separate from those plain bitvectors. When the poison bit is on, we say that the entire value is poison. Otherwise the value is well-defined.

A use of a poison value can result in one of the following, depending on the op (inline constants for brevity):

  1. Another poison, e.g., arith.addi %x, poison : i32
  2. A well-defined non-poison value, e.g., arith.select false, poison, 1 : i32
  3. Immediate UB, e.g., arith.divui poison, 0 : i32

The ub.poison : ty op creates a poison value of the given type. Poison propagation is unconditionally stopped by the ub.freeze %x : ty which is a no-op for well-defined values and a nondeterministic choice for poison values. Last, we provide ub.unreachable that when executed always exhibits immediate UB.

The ops above are said to be ‘convenience’ ops because they simply implement the new interfaces (see below). Other dialects may define their own equivalent ops, e.g., llvm.poison.

Vector Case

For some compound types like 1-D vectors, tuples, or other aggregate-like types it would probably be more useful to support per-element poison semantics. In this case, the poison state can be a bitvector with a poison bit per each element. For brevity, we say that a poison value of a vector is a vector with all poison state bits on:

%v = ub.poison : vector<2xi32> // Poison state: <true, true>

while ub.freeze becomes an elementwise operation for vector types.

Vector operations (like in arith.* and vector.*) can then manipulate the bitvector of poison bits. For example, when we perform a vector shift arith.shrui %x, %y : vector<2xi32>, it’s possible that only one of the shift values is greater or equal to the bitwidth. We would hope that the other result element is well-defined despite the other one being poison. This allows us to, for example, scalarize vector operations:

%r  = arith.shrui %x, %y : vector<2xi32>
%e0 = vector.extract %r[0] : vector<2xi32>
%e1 = vector.extract %r[1] : vector<2xi32> 
 ==>
%x0 = vector.extract %x[0] : vector<2xi32>
%x1 = vector.extract %x[1] : vector<2xi32>
%y0 = vector.extract %y[0] : vector<2xi32>
%y1 = vector.extract %y[1] : vector<2xi32>
%e0 = arith.shrui %x0, %y0 : i32
%e1 = arith.shrui %x1, %y1 : i32

Or vice versa, i.e., vectorize scalar operations. Similar concerns can arise with higher-level ops like vector.mask [7] or tcp.elementwise [8].

(Note: the transformation above is correct in LLVM: https://alive2.llvm.org/ce/z/UPVa6x.)

Vectors in MLIR can have more than one dimension (unlike LLVM). There’s also support for scalable vectors. This may necessitate a more elaborate poison state representation than a (flat) bitvector of poison bits.

Custom Type Case

Poison semantics may not be required or make sense for all types. For example, resource handles or synchronization primitives may not have meaningful arbitrary values necessary to freeze poison values. This makes us propose poison semantics to be opt-in for types:

  • By default, types do not have poison semantics. ub.poison and ub.freeze are not defined for these types and fail verification.
  • When opted-in, types can define their poison state representation.
  • Types decide whether to support ub.poison and/or ub.freeze. This is a static property of the given type. We say that a type that supports ub.poison is poisonable and a type that supports ub.freeze is freezable.

Compound Type Case

Let’s consider a hypothetical pair<ty1, ty2> type combining two arbitrary element types. This type could combine the poison semantics of its element types in the following ways:

  1. Per-element poison semantics:
    • One poison state bit per each element type with poison semantics. By definition, only poisonable elements can be in the poison state.
    • Only poisonable when all elements are poisonable.
    • Only freezable when all poisonable element types are also freezable.
  2. Lowest common denominator of element types poison semantics:
    • If either of the element types does not support poison semantics, pair<ty1, ty2> should not either.
    • The poison state is a pair of two poison states of the elements.
    • If either of the types is not poisonable/freezable, pair<ty1, ty2> is not either.
  3. Whole-value poison semantics:
    • The poison state is a single bit.
    • If either of the elements is poison at some point, the whole pair becomes poison.
    • The pair is freezable only when both element types are freezable.

Ultimately, it’s up to the dialect that defines the pair<ty1, ty2> type to define its semantics in a way that is compatible with its ops.

Desired Folds

// 1
%x = ub.poison : i32
%y = arith.add %x, %x : i32
  ==>
%y = ub.poison : i32

// 2
%x = arith.constant 0 : i32
%y = ub.freeze %x : i32
  ==>
%y = arith.constant 0 : i32

// 3
%x = ub.freeze %arg0 : i32
%y = ub.freeze %x : i32
  ==>
%y = ub.freeze %arg0 : i32

// 4
%x = arith.add %arg0, %arg0 : vector<2xi32>
%y = vector.extract %x [0] : vector<2xi32>
%z = vector.extract %x [1] : vector<2xi32> 
  ==>
%a = vector.extract %arg0 [0] : vector<2xi32>
%b = vector.extract %arg0 [1] : vector<2xi32>
%y = arith.add %a, %a : i32
%z = arith.add %b, %b : i32 

// 5
%a = vector.extract %arg0 [0] : vector<2xi32>
%b = vector.extract %arg0 [1] : vector<2xi32>
%y = arith.add %a, %a : i32
%z = arith.add %b, %b : i32 
 ==>
%x = arith.add %arg0, %arg0 : vector<2xi32>
%y = vector.extract %x [0] : vector<2xi32>
%z = vector.extract %x [1] : vector<2xi32>

// 6
%p = ub.poison : i32
%v = arith.constant dense 0 : vector<2xi32>
%y = vector.insert %p into %v [1] : vector<2xi32>
%z = vector.extract %v [0] : vector<2xi32>
 ==>
%z = arith.constant 0 : i32

// 7 -- https://alive2.llvm.org/ce/z/2o8FpN
%v = ub.poison : vector<2xi32>
%c = arith.constant 0 : i32
%y = vector.insert %c into %v [0] : vector<2xi32>
%z = vector.extract %v [0] : vector<2xi32>
 ==>
%z = arith.constant 0 : i32

// 8
%p = ub.poison : i32
%v = vector.broadcast %p : i32 to vector<2xi32>
%x = vector.extract %v [0] : vector<2xi32>
 ==>
%x = ub.poison : i32

Implementation

Goals

  • We want to faithfully capture the poison semantic while making the implementation performant for the open set of possible types and ops.
  • Dialects should be able to opt into poison semantics.
  • Compound types can define the elements with poison semantics.
  • Analyses derive lattices and transfer functions based on dialect definitions.
  • Allow for gradual rollout without breaking existing dialects.

Implementation Strategy

  1. Introduce the concept of poison and poison-related terminology to the MLIR documentation.
  2. Add implementation for lowest-level dialects that will benefit from poison semantics: LLVM, arith, vector.
    1. Introduce poison bottom-up in the dialect lowering chain.
  3. Gradually expand and generalize the implementation. Make it easy for other dialects to opt into poison semantics.
    1. Add traits for most common cases simplify the implementation.
    2. Add a ‘ub’. dialect with ub.poison and ub.freeze ops, so that other dialects do not have to reimplement those.
    3. Add support for a generic poison dataflow analysis.

Later implementation stages (3.) may require separate RFC to agree on the details.

Initial Implementation Idea

The implementation is not 100% clear at this point. We hope that the implementation experience gained when adding support to a couple initial dialects will inform a general and pluggable implementation.

At this point, we propose for dialects to opt into poison semantics through Type and Op interfaces.

Types:

  • By default, types do not have any poison states
  • Type can be poison as a whole or can have its element poison
  • Types with metadata (e.g., memref) can have metadata as poison but elements well-defined
  • The initial type interface has two methods:
    • bool canBePoison() – whether the type is poisonable and ub.poison can be used on this type, is allowed to recursively check the elements and other contents of the type. Intended for compound types that may include element types that do not support poison semantics.
    • bool canBeFrozen() – whether the type is freezable and ub.freeze can be used on this type, is allowed to recursively check the elements and other contents of the type. Intended for types like resource handles that cannot be frozen.

Ops:

  • By default, if any of the input values have poison, the poison state is propagated.
  • Ops can specify a more fine-grained behavior if they wish to.
  • The op interface describes how poison state is propagated from inputs to outputs.

FAQ

Why not undef?

Both LLVM and SPIR-V started with undef values to model uninitialized values or results of instructions whose preconditions were violated without immediately trapping. Every use of undef yields an arbitrary bit pattern, which makes some compiler optimizations difficult to get right. For these reasons, we do not propose a similar op in the ‘ub’ dialect.

Related discussion: [5, 6].

Do all types have a poison state?

No, poison is opt-in for types. We expect it to be most useful to basic types (integers, floats, etc.) and compound types containing the basic ones (vectors, structs, tensors, etc.).

Should all types with poison state support ub.poison?

Some types may combine types with and without poison semantics which makes them not poisonable as a whole, e.g., pair<i32, m_dialect.ResourceHandle>. In this case, we may have custom ops that only poison/freeze the first pair element.

Should all types with poison state support ub.freeze?

It may not be possible to create meaningful arbitrary values of some types, e.g., synchronization primitives or resource handles.

When you add poison state to builtin types like IntegerType, will I have to update my dialect to support poison?

Most likely not if your dialect is higher-level. If none of your ops produce poison, and you do not lower from ops that are known to produce poison, then you don’t have to worry about any poison values. This is why we intend to use a bottom-up strategy and introduce poison to the lowest-level dialects first.

When should I consider giving my types and ops poison semantics?

When your ops can benefit from refusing to define some behavior, e.g., allowing uninitialized values or ops that would otherwise require inserting runtime checks when lowering. In such cases, we can gain extra optimization power at the expense of increased semantic complexity. Using poison semantics has 2 main benefits:

  • It gives us a framework to talk and reason about such cases and define the precise semantics (documentation and dialect definitions)
  • It allows for code hoisting, duplication, and speculation based on these semantics.

We expect this to be most useful in lower-level dialects. In those dialects, we expect the vast majority of ops to simply propagate poison elementwise.

References

[1] https://www.cs.utah.edu/~regehr/papers/undef-pldi17.pdf

[2] [llvm-dev] RFC: Killing undef and spreading poison

[3] LLVM Language Reference Manual — LLVM 16.0.0git documentation

[4] [RFC] Define precise arith semantics

[5] Why there's no "std.undef"?

[6] [RFC] Undef/uninitialized values and control flow

[7] [RFC] Vector Masking Representation in MLIR

[8] [RFC] Elementwise ops in TCP


-Jakub Kuderski (@kuhar), with input from Alex Zinenko (@ftynse), Sanjoy Das (@sanjoyd), and Nuno Lopes (@nlopes)

cc: @regehr, @River707, @clattner, @Mogball

3 Likes

This looks good, I particularly like the desired folds part, this kind of operational explanation is probably the most useful kind of explanation for most readers.

I would phrase “When should I consider giving my types and ops poison semantics” differently. The core question is not whether an operation naturally is defined for all inputs, because it’s generally trivial to define these cases. For example, we can easily define division by zero to return 0 or -42 or whatever. Or see funnel shifts in LLVM, they have no undefined cases, unlike the regular shifts. So the problem is not defining these cases, the problem is deciding whether it is worthwhile to deliberately refuse to define these cases in order to buy some extra optimization power. If I were writing this, I would encourage people to think real hard on this point before committing to returning poison, and only do this when the resulting semantic complexity is offset by increased optimization power.

2 Likes

Thanks for the proposal, would you be up for a presentation at an open meeting? We don’t have anything on the agenda so it can be as soon as next week :slight_smile:

Sure, next week sounds good

Awesome RFC. I enjoyed the read. +1 to making this opt-in. I have some questions and concerns:

  1. How does this interact with constant folding? Are we going to introduce a #ub.poison wildcard which, say, div x, 0 would fold to?
  2. What’s the implementation for the LLVM dialect going to look like, and how will this affect LLVMIR import and export?

Adding a separate ub dialect with interfaces to plug into the behaviour seems like a solid design for me.

  1. How does this interact with constant folding? Are we going to introduce a #ub.poison wildcard which, say, div x, 0 would fold to?

Implementation is not 100% clear to me yet, suggestions welcome :slight_smile: . For C++ folders, I imagined matchers like m_Poison, m_NotPoison, and m_Nondet (for freeze poison) to cover simple cases. Any non-trivial case will probably require a poison dataflow analysis. And as for dix x, 0 specifically, this would be considered immediate UB, so I think we would want to DCE it instead (with ub.unreachable or otherwise).

  1. What’s the implementation for the LLVM dialect going to look like, and how will this affect LLVMIR import and export?

LLVM IR currently has all of: poison, freeze, ureachable, undef, so the LLVM dialect would map 1:1 after we add the missing ops: llvm.poison, llvm.freeze.

Hi,

Maybe you will want to talk about the elements of tensor.empty as well.
Bufferization + lowering to LLVM IR may turn it into an uninitialized malloc (correct me if I’m wrong), and reading uninitialized bytes of malloc may raise UB. Since one cannot turn a well-defined program into undefined one, this implies that the elements of tensor.empty must be poison.

2 Likes

Hey! Has there been further development on this? What is the status? This is very interesting stuff.

1 Like

Hi @Moxinilian,

I went deep into a couple of SPIR-V compilation rabbit holes and haven’t started the implementation of poison semantics yet. I expect to get back to it in a month or so.

The only related development was that we recently spected out the out-of-bounds semantics for vector.gather / scatter / maskedload / maskedstore: https://reviews.llvm.org/D145824, https://reviews.llvm.org/D145975.

Have you run into anything that could use poison semantics or just wanted a status update?

1 Like

Great news!

I am currently doing mem2reg related stuff in MLIR. I am porting the LLVM infrastructure to the LLVM dialect, but I am of course thinking of extending it to be more general-purpose. Undefined states are an important part of alloca-based constructs so having poison semantics to represent illegal-to-access uninitialized memory values would be great.

There is effort in LLVM to make mem2reg do more with poison values instead of undefined values, so having poison semantics is also relevant for the LLVM IR case.

4 Likes

Hello!

It seems my line of thinking on this is very different. I may be more in the MLIR mindest than LLVM by now: constants need to be materialized, and you don’t modify existing language elements, you add new ones. So, I decided the only thing the ub dialect really needed to do for me is to provide an attribute and a materializer.

I kinda needed this right now, so I cobbled something together. Keep in mind this is extremely bare-bones.

ub.poison and ub.frozen

As stated also in your slides, poison is a compile-time manifestation of UB. In other words, it only shows up as a constant value, i.e., an attribute value. “Encountering poison at runtime” is just UB, after all. Consequently, the only place I was ever going to “make use” of poison was folding. Still, to allow partial poisoning of aggregates, poison must have a materializer. The same does not apply to freezing – that must not fold away, and you have to explicitly visit these definitions to make use of them.

My PoisonAttr is either fully-poisoned, or wraps a sourceAttr from a sourceDialect together with a poisonMask. The source dialect is necessary so that an opt-in user may control materialization (via the materializeConstant callback). In my use cases though, that’s enough to get me through. Also, only ShapedType are considered aggregates right now, but this could be remedied with a type interface.

Some definite advantages of your approach are, of course, trait based foldings and interfaces. I generally prefer interfaces everywhere, but I am unsure how, in a non-keyword scenario, types and attributes would actually participate. I am also slightly annoyed by all the monomorphisation I need to do during constant folding, to combine the source attribute values with the mask. I haven’t actually found a good non-template solution, so I did not push any.

What is even more lacking though is one type of folding you described, where container manipulations update the partial poisoning. It can’t be done with my approach, because attributes and values don’t mix. I’d have to use a witness for that, and I don’t want to have to keep track of this.

On a different note: ub.unreachable

I wonder how this would integrate well. We can’t make it a new terminator of existing regions like LLVM (and probably shouldn’t). We also don’t have a PanicSideEffect interface yet that could ensure it stays put. It can’t be a witness-style op either because its control-flow and not value based.

I had a thought that it should be represented differently, and that ub would provide a new region-based op that participates in the existing control flow interfaces. If that depends on values, it would also work in the DAG case. Without a concrete use I couldn’t tell you how though.

I would like to hear what you all think about this, because it seems very different to me.

That’s awesome! What’s you use-case, if you don’t mind sharing?

As stated also in your slides, poison is a compile-time manifestation of UB. In other words, it only shows up as a constant value, i.e., an attribute value. “Encountering poison at runtime” is just UB, after all.

I don’t think I agree on this one. We can have entirely valid program executions where poison appears and is used as an operand and there’s no UB. It’s true that when poison appears in the IR it is a constant, but it can be also a value returned by ops under some conditions. If we had a hypothetical reference mlir interpreter, poison could be an actual value like any other value of <your favorite type>.

My PoisonAttr is either fully-poisoned, or wraps a sourceAttr from a sourceDialect together with a poisonMask . The source dialect is necessary so that an opt-in user may control materialization (via the materializeConstant callback). In my use cases though, that’s enough to get me through. Also, only ShapedType are considered aggregates right now, but this could be remedied with a type interface.

This makes sense to me. I’d expect constant-like ops to be able to materialize PoisonAttr. Perhaps this means that we don’t need the ub.poison op after all and can use the usual dialect-specific constant materialization, e.g., %x = arith.constant #ub.poison : i32?

What is even more lacking though is one type of folding you described, where container manipulations update the partial poisoning. It can’t be done with my approach, because attributes and values don’t mix. I’d have to use a witness for that, and I don’t want to have to keep track of this.

Could we materialize those by a composition of a few ops, e.g., first materializing a fully-poisoned vector and then inserting it into the known-to-be-poison vector element positions? In general, I don’t think we strictly need to be able to express constant of every possible poison configuration – this is nice for writing tests, but an alternative would be to have a poison analysis and use it’s results during compilation only. Or something even simpler, we could agree on the concept and define dialect semantics with poison in mind, but not have any constant/attribute to express it and no analysis in place either.

I wonder how this would integrate well. We can’t make it a new terminator of existing regions like LLVM (and probably shouldn’t).

The llvm instruction is a terminator though LLVM Language Reference Manual — LLVM 18.0.0git documentation; why shouldn’t we? I could see ub.unreachable being a terminator or something similar to noreturn functions, but I haven’t thought too much about it.

It can’t be a witness-style op either because its control-flow and not value based.

What is a ‘witness-style’ op? Something like: 'shape' Dialect - MLIR?

Nothing new, unfortunately. The same as yours I believe. Except that I am part of the “carthago should be destroyed” camp when it comes to arith, and propose a complete replacement.

It was a bit late yesterday, maybe I didn’t get the message across. If you mean the classical example of select %cst_false, poison, %well_defined, then I totally agree, and that should work with this.

I’ve been following that too. While, for performance reasons, it may be desirable to represent intermediate values in the interpreter differently, the most canonical way right now would be attributes, I think. So that’d still be possible.

Yeah, that’s kinda what I started with. I just realized that a completely poisoned value deserves its own op for canonicalization purposes.

I agree. I think the limiting factor for doing this right now is some kind of AggregateTypeInterface that is able to materialize enumeration, insertion and extraction operations. This seems very hard though as some types (think dynamic shapes) are only runtime enumeratable. I imagine that part is also a bottleneck for a generalized analysis.

Yes, that’s what I’m saying. And I say we can’t make it a terminator because that is not a thing in MLIR. Region declaring operations are free to assign any meaning (or none) to their terminator instructions. In fact, most MLIR things I do don’t use SSACFG but DAG instead. So, unless you add a new type of control flow interface that indicates “yes, I do actually mean linearized control flow”, such a terminator is not well-defined IMHO.

Exactly. By splicing a def/use edge using a witness operation, you’re able to flag some property for some value beyond that point. That protects your assertion from being optimized away or reordered. The other way is to define a new kind of side-effect, I suppose. Just that doing this for control-flow and not values requires some creative thinking I believe.

Hello again.

I added an an attribute concept to ub-mlir that shows how I expect users to interface with it. Ideally, they would use it to wrap their own attribute concepts, as I have done for my bit dialect here. That file also declares the three folding helpers, map and zip. This is where the monomorphisation kicks in… You can see their implementation here. This is where it would be cool to have a templated interface so we wouldn’t have to do this for every dialect again and again.

To have something to show you, I added poison semantics to the bit.cast op (you can guess what that is, I’m sure). The following test now passes:

// CHECK-LABEL: func.func @cast_poison(
func.func @cast_poison() -> i32 {
    // CHECK: %[[POISON:.+]] = ub.poison : i32
    %0 = ub.poison : si32
    %1 = bit.cast %0 : si32 to i32
    // CHECK: return %[[POISON]]
    return %1 : i32
}

I’m gonna hope I didn’t mess up (which I probably did somehwere, considering how much code I just pushed) and press on. Partial containers should work, let’s see how that’s going to look like.

PS: Note that the base2-mlir project I linked to is quite out of date design-wise. I decided to focus on the specification first.

I have had some time to think about unreachable and I ended up with these ideas. I’d be really interested if you’re able to relate them to your imagined use cases.

Discardable UnitAttr "ub.unreachable"

The simplest way to add unreachable to MLIR is probably to just add this attribute, and let it be conditionally valid only on Terminator ops. To be more precise, only those defined within an SSACFG block.

The question is how this information is supposed to be used. Without a concrete usage scenario, I wouldn’t know why I’d use this.

cf.unreachable terminator

In LLVM, the unreachable instruction is often emitted for a covered switch (i.e., where the default alternative is known to never be taken). A new block will be emitted for the default label, which only contains the unreachable terminator. The reason why this is necessary in the first place is because the switch instruction requires a default block. One could have just as easily defined an absent default to be unreachable, for instance.

We have cf.switch, which has the same semantics, and also requires a default. I would argue that this use case would probably be better served by a cf.unreachable terminator. It could also be used to canonicalize cf.assert, for example:

  %false = arith.constant false
  cf.assert %false, "unreachable"

into

  cf.br ^unreachable()
  ...
^unreachable():
  cf.unreachable

which ties back to when assert(false && "unreachable") should be the same as __builtin_unreachable() – when assertions are enabled.

!ub.never, #ub.never, ub.never

A stronger version of unreachable is a noreturn function. Some languages model this using a never type, which is defined as the result of a non-returning function. In terms of annotating a function, the difference between a noreturn attribute and a !ub.never result type is that the latter requires callsites to be aware of this.

func.func @loop_forever() -> !ub.never { ... }

// I think discarding all op results is actually valid.
func.call @loop_forever() : () -> !ub.never

I think this can be generalized to dataflow problems as I described in my earlier post. Let’s say we can use ub.never to materialize a !ub.never value. This would be the last stage in reasoning for noreturn.

Before that, we’ll use ub.never to materialize a #ub.never : T attribute, which is similar to poison. It indicates that this result will never be computed, but allows current users to still remain syntactically valid. Alternatively, we could builtin.unrealized_conversion_cast from !ub.never to T, but then users would loose the understanding that #ub.never is actually a compile-time constant thing.

func.func @actually_never() -> i64 { ... }

%x = func.call @actually_never() : () -> i64
%y = arith.muli %x, %x : i64
// Knowing @actually_never has no reachable terminator, this becomes
%y = ub.never : i64

In this case, we could let a propagation pass simply turn every computation in to ub.never that depends on a ub.never operand – instead of allowing custom poison semantics, we require this special value to always proagate. In this way, we could eliminate dead dataflow and use the equivalent of unreachable in non-SSACFG scenarios.

There are two concepts that seem mixed up in your writing (as far as I understand…).
The first one is the concept of poison, which is about “delayed UB”: a value is marked poison but there is no UB happening until it used in very specific conditions.
The second is the unreachable instruction which is really about the control: it tells us that that an entire execution path is impossible. It allows the compiler to reason backward about branch conditions for example. It is useful during optimizations to make a path as impossible without having to update a possible complex control flow, making the transformation “local” and leaving more complex cleanup or other transformations to use the instruction to infer the “impossible path” property.

I don’t think one can replace the other or is a superset of the other, and so we need both?

If this is directed at me, then yes, I agree. I’m not trying to unify them. After previously presenting my no-intrinsics implementation of poison semantics, I just don’t have anything to add to that particular part.

However, I’m trying to challenge the original poster’s design of the ub dialect:

My argument is that an unspecified unreachable terminator is not a well-defined concept in MLIR, where regions, block and terminators can be assigned arbitrary meanings outside of control flow. I argue that currently, only the cf dialect could reasonably define and make use of an unreachable terminator.

I do, however, concede that unreachable code paths that produce dataflow dependencies may be a more general use case that is worth modelling. Hence ub.never, to propagate reachability in SSA dependency graphs.

I was working towards a v0.2 of my ub-mlir dialect that includes unreachability (and adheres to newer MLIR code standards, i.e., free function casting) and LLVM lowering (Big thanks to @Moxinilian for getting the llvm dialect extended!). I will pause this for now, but I can now give a concrete argument for why a ub.unreachable terminator is problematic in the way I thought of it.

In general, any terminator that marks “well-defined” control flow may participate in unreachability analysis:

  • Any BranchOpInterface implementation.
  • Any RegionBranchTerminatorOpInterface implementation.

We can relax this a bit to:

  • Any IsTerminator that is defined inside an SSACFG block.
  • Any isRegionReturnLike(Operation*) (also allows ReturnLike ops).

If we allow attaching a discardable ub.unreachable attribute to all these, we can perform the analysis, including region branch and data flow based unreachability propagation.

There is just one scenario that is not well covered: when we allow sequential control flow reasoning, i.e., an operation is unreachable if a previous op in the block is noreturn or unreachable, we would want to split the block and make the rest unreachable. (We would not allow this in graph regions.) Then, we need to insert a ub.unreachable terminator op (or llvm.unreachable if you will).

^bb0:
  op0
  op1 // <- noreturn
  op2
  term

// Perform split before 'op2' to get:
^bb0:
  op0
  op1
  ub.unreachable
^bb1:
  op2
  term // with 'never' operands, if any

This is not generally legal! The surrounding op may be SSACFG, but still only allow a single block (see scf), in which case splitting is ilegal and we have no way to know. Or, we can split, but our new terminator is actually not a valid ReturnLike for the surrounding op. Either way, we are screwed, because we can’t do this right without new interfaces.

The reason why I’m not moving term, which is guaranteed to be legal, up to replace op2, and erase the other operations, is because I allow unreachble operations to produce values at compile time (i.e., if they are known to have internal control flow that excludes the unreachable never path). If I just deleted them and moved up term, it would be perfectly legal, but would loose this information.

I am tending towards making ub.unreachable just a graph container that is deleted when unreachability is finalized, and move deleted blocks in there. However, this makes llvm.unreachable not emissible by my conversion, because I do not control any terminators anymore, thus defeating the point. Unless of course, all exploitation of unreachability happens in MLIR, and there is no reason to pass it to LLVM anymore.

Again, I think cf.unreachable is the best option to go with for general use.

2 Likes

Thanks for a great write up!

In a non-core dialect with a custom loop and yield ops, do you mean something like:

%a, %b = my.loop {
  op0
  %1 = op1
  %2 = op2
  my.yield %1, %2
}

==>

%a, %b = my.loop {
  op0
  ub.unreachable {
    %3 = op1
    %4 = op2
  }
  %1 = ub.poison
  my.yield %1, %1
}

when you mention ‘graph container’? If my.loop only permits a single terminator type, my.yield, we would have to preserve it and replace all operands by poison (which may be not possible for types that do not have poison values).

I think we still end up with 2 cases:

  • #ub.unreachable used in non-terminator ops for cases where the enclosing region does not support the ub.unreachable op as the terminator
  • ub.unreachable for regions that support ub.unreachable as a terminator type.

Regular transforms and rewrites would only have to add the #ub.unreachable attribute to some op, and later we would have a completely separate dead code elimination pass use that inserts ub.unreachable and deletes dead instructions. For regions that do not support ub.unreachable, those unreachable ops can be eliminated later in the pipeline, after lowering to scf or cf. I’m not sure how much we would gain by introducing a new region for the deleted ops.

In general the reason to keep code around resolves mostly around preserving code invariants informations. For example (in pseudo IR):

%cond = icmp le %x, 0
cond_br %cond, ^bb1, ^bb2
^bb1:
  unreachable
^bb2: 
  op1
  op2
  ...

Can be replaced by:

  op1
  op2
  ...

(that is only bb2 remains)

However we lost the information about the condition here. I believe LLVM would introduce llvm.assume to preserve the information about the condition to not lose it.

From there: are there other kind of information that isn’t captured by the branch condition and that we’d want to capture and where a region is useful?