[mlir][ub] Add poison support to CommonFolders.h

[Hardcode84]

Summary

Return poison from foldBinary/unary if argument(s) is poison. Add ub dialect as dependency to affected dialects (arith, math, spirv, shape).
Add poison materialization to dialects. Add tests for some ops from each dialect.
Not all affected ops are covered as it will involve a huge copypaste.

[zero9178]

I think the layering here is a bit problematic. The generic Dialect/CommonFolders.h suddenly depends on the UB dialect.
We are now also suddenly forcing the semantics of poisons on all users of these folders.
I am not sure what the best solution would be, but maybe versions/wrappers of those folders within the UB dialect? That way all op folders/dialects can also opt-into the poison behaviour. This also ensures correctness.

As a sidenote, since your new materializeConstant implementations may now create a ub.poison op, you need to add the UB dialect as a dialect dependency in the TableGen of the respective dialects as well.

[Hardcode84]

In D159013#4622371, @zero9178 wrote:

I think the layering here is a bit problematic. The generic Dialect/CommonFolders.h suddenly depends on the UB dialect.
We are now also suddenly forcing the semantics of poisons on all users of these folders.
I am not sure what the best solution would be, but maybe versions/wrappers of those folders within the UB dialect? That way all op folders/dialects can also opt-into the poison behaviour. This also ensures correctness.

As a sidenote, since your new materializeConstant implementations may now create a ub.poison op, you need to add the UB dialect as a dialect dependency in the TableGen of the respective dialects as well.

CommonFolders.h is used in limited number of places (4) and all of them will benefit from poison IMO, so I'm not sure if it's worth keeping both poison and non-poison version. But I agree that current place may not be the best after introducing dependency on ub dialect, any suggestions for better place?

Regarding runtime dependency on ub dialect, technically, things will still work, as we can only generate poison if it was already present in IR, meaning dialect was already loaded, but that's fragile, yes.

[mehdi_amini]

mlir/include/mlir/Dialect/CommonFolders.h

46

To break the dependency on the UB dialect, I think we can handle this with a template parameter:

template <class AttrElementT,
          class ElementValueT = typename AttrElementT::ValueType,
          class CalculationT = function_ref<
              std::optional<ElementValueT>(ElementValueT, ElementValueT)>
          class PoisonAttr = void>
Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
                                       Type resultType,
                                       const CalculationT &calculate) {
  assert(operands.size() == 2 && "binary op takes two operands");
  if (isa_and_nonnull<PoisonAttr>(operands[0]))
    return operands[0];

  if (isa_and_nonnull<PoisonAttr>(operands[1]))
    return operands[1];

I don't know if the default "void" will work out-of-the-box, but there is likely a trick to have it work here.

Now that makes it not use Poison by default, so another way could be a forward decl:

namespace ub {
class PoisonAttr;
}
template <class AttrElementT,
          class ElementValueT = typename AttrElementT::ValueType,
          class CalculationT = function_ref<
              std::optional<ElementValueT>(ElementValueT, ElementValueT)>
          class PoisonAttr = ub::PoisonAttr>

And that would make it opt-out for cases where we wouldn't want to depend on the UB dialect...

[mehdi_amini]

LG otherwise, please try to build with SHARED_LIBS=ON to ensure correct build dependencies

[zero9178]

In D159013#4622547, @Hardcode84 wrote:

In D159013#4622371, @zero9178 wrote:

I think the layering here is a bit problematic. The generic Dialect/CommonFolders.h suddenly depends on the UB dialect.
We are now also suddenly forcing the semantics of poisons on all users of these folders.
I am not sure what the best solution would be, but maybe versions/wrappers of those folders within the UB dialect? That way all op folders/dialects can also opt-into the poison behaviour. This also ensures correctness.

As a sidenote, since your new materializeConstant implementations may now create a ub.poison op, you need to add the UB dialect as a dialect dependency in the TableGen of the respective dialects as well.

CommonFolders.h is used in limited number of places (4) and all of them will benefit from poison IMO, so I'm not sure if it's worth keeping both poison and non-poison version. But I agree that current place may not be the best after introducing dependency on ub dialect, any suggestions for better place?

There are likely downstream users of these methods (I believe these are quite old), so I do think this would have an impact. At the very least it'd create a link dependency for them.

mlir/include/mlir/Dialect/CommonFolders.h

46

I think the first option is a great idea!
You can make it work with the default void by just wrapping these two ifs into a if constexpr (!std::is_void_v<PoisonAttr>).

[Hardcode84]

make poison handling optional, add ub dialect as runtime dependency

[Hardcode84]

Made PoisonAttr a template parameter as suggested but still left it on by default as I think in most cases this is a desired behavior. Replaced ub dialect include with forward decl, so any potential users will now need to add explicitly dependency on ub dialect or opt out to successfully compile.

[jpienaar]

mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td

42 ↗

(On Diff #554266)

Shape dialect should not depend on UB dialect.

[zero9178]

Implementation looks good to me, but I rather strongly believe this behaviour should be opt-in rather than opt-out. Otherwise you're by default redefining the semantics of all ops using these. Poison attributes are not the obvious way to handle undefined behaviour but rather a choice in the IR design.

mlir/include/mlir/Dialect/CommonFolders.h

33

Can we update the documentation of at least one of these so users can tell whether they need to at least care about PoisonAttr?

mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp

953

nit: use dyn_cast here and all other similar cases that have isa followed by cast

[Hardcode84]

mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td

42 ↗

(On Diff #554266)

It will depend on it transitively through arith dialect anyway, added ihere explicitly because it is used in materializeConstant. Its there any specific reason why you don't want to handle poison in shape dialect?

[zero9178]

In D159013#4624504, @zero9178 wrote:

Implementation looks good to me, but I rather strongly believe this behaviour should be opt-in rather than opt-out. Otherwise you're by default redefining the semantics of all ops using these. Poison attributes are not the obvious way to handle undefined behaviour but rather a choice in the IR design.

Just for a list of users potentially impacted by opt-out, see: https://github.com/search?q=%22mlir%2FDialect%2FCommonFolders.h%22+language%3AC%2B%2B&type=code
This includes projects like Flang, ClangIR, Circt and IREE.

[Hardcode84]

make poison opt-in

[Hardcode84]

In D159013#4624519, @zero9178 wrote:

In D159013#4624504, @zero9178 wrote:

Implementation looks good to me, but I rather strongly believe this behaviour should be opt-in rather than opt-out. Otherwise you're by default redefining the semantics of all ops using these. Poison attributes are not the obvious way to handle undefined behaviour but rather a choice in the IR design.

Just for a list of users potentially impacted by opt-out, see: https://github.com/search?q=%22mlir%2FDialect%2FCommonFolders.h%22+language%3AC%2B%2B&type=code
This includes projects like Flang, ClangIR, Circt and IREE.

OK, that's fair. Although, with opt-in most of these users will never know about this poison support and will never update.

[kuhar]

Awesome!

[mehdi_amini]

There are likely downstream users of these methods (I believe these are quite old), so I do think this would have an impact. At the very least it'd create a link dependency for them.

It's up to them to opt-out, we don't optimize to avoid breaking downstream, we strive for the best default behavior moving forward.

mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td

42 ↗

(On Diff #554266)

materializeConstant does not require a dependent dialect because it only propagate poison and won't create any out-of-thin air. So please remove this everywhere except where there are canonicalization that could break this.
(and if something does not show up with a grep in the dialect, it should be documented as well)

[zero9178]

In D159013#4627031, @mehdi_amini wrote:

There are likely downstream users of these methods (I believe these are quite old), so I do think this would have an impact. At the very least it'd create a link dependency for them.

It's up to them to opt-out, we don't optimize to avoid breaking downstream, we strive for the best default behavior moving forward.

That's very fair, so maybe I worded this poorly. My concern is not any "churn" downstream in the sense of them requiring actions to keep compiling their code or other refactoring. I like breaking them reguarly as well :-)
Rather, I don't think this is the best default behaviour for such generic functions. My first concern is that always returning poison if any operand is a poison can lead to miscompiles. Obvious cases include things like select operations or other "predicated/masked"-like operations.
Always returning poison therefore can lead to a miscompile. Not propagating poison is only ever a missed optimization, but not an issue of correctness.
It is also not obvious to me that poison semantics, aka deferred UB, is the obvious way that all dialects should be modelling UB. Making it the default would IMO be overfitting to specific kinds of dialects.

I therefore am in favour of the current version of the patch.

[Hardcode84]

remove runtime dependecy on ub dialect

[zero9178]

This version LGTM but please wait for @jpienaar for the Shape dialect changes/decision and potentially a day for others to comment on it as well.

[mehdi_amini]

In D159013#4627450, @zero9178 wrote:

In D159013#4627031, @mehdi_amini wrote:

There are likely downstream users of these methods (I believe these are quite old), so I do think this would have an impact. At the very least it'd create a link dependency for them.

It's up to them to opt-out, we don't optimize to avoid breaking downstream, we strive for the best default behavior moving forward.

That's very fair, so maybe I worded this poorly. My concern is not any "churn" downstream in the sense of them requiring actions to keep compiling their code or other refactoring. I like breaking them reguarly as well :-)

OK :)

Rather, I don't think this is the best default behaviour for such generic functions. My first concern is that always returning poison if any operand is a poison can lead to miscompiles. Obvious cases include things like select operations or other "predicated/masked"-like operations.

But it's not like we're doing it for any op generically in the folder here: this is "opt-in"!
People don't have to use these helpers...

Always returning poison therefore can lead to a miscompile. Not propagating poison is only ever a missed optimization, but not an issue of correctness.
It is also not obvious to me that poison semantics, aka deferred UB, is the obvious way that all dialects should be modelling UB. Making it the default would IMO be overfitting to specific kinds of dialects.

It's perfectly fine to "overfit" to "specific dialects" when these "specific dialects" includes basically "all" of the dialect we have in-tree :)
(especially when we're only ever talking about the default settings)

That is: I'm not sure I understand really the motivation behind your rebuttal here, I don't quite get how to optimize anything without deferred UB.

mlir/include/mlir/Dialect/CommonFolders.h

117

Typo

117

(everywhere else as well)

[mehdi_amini]

@Hardcode84 : don't take this as blocking you from landing, but I'd still like to get to end of the argument with @zero9178 here, this is pretty important for any future direction about how we see UB in MLIR upstream (this "default" deferred UB question in particular may come up all the time!)

[zero9178]

In D159013#4630081, @mehdi_amini wrote:

In D159013#4627450, @zero9178 wrote:

In D159013#4627031, @mehdi_amini wrote:

There are likely downstream users of these methods (I believe these are quite old), so I do think this would have an impact. At the very least it'd create a link dependency for them.

It's up to them to opt-out, we don't optimize to avoid breaking downstream, we strive for the best default behavior moving forward.

That's very fair, so maybe I worded this poorly. My concern is not any "churn" downstream in the sense of them requiring actions to keep compiling their code or other refactoring. I like breaking them reguarly as well :-)

OK :)

Rather, I don't think this is the best default behaviour for such generic functions. My first concern is that always returning poison if any operand is a poison can lead to miscompiles. Obvious cases include things like select operations or other "predicated/masked"-like operations.

But it's not like we're doing it for any op generically in the folder here: this is "opt-in"!
People don't have to use these helpers...

That is true, but people are already using these helpers. From that perspective this is not "opt-in", but rather silently changing the semantics of existing functions. The scenario I am scarred about is a fold implementation, using these functions and incorrectly returning a poison attribute.
This is the kind of bug that in my experience is very hard to track down as it carries through all layers of the compilation and successfully generates runnable code as well. Only when running the code and potentially hitting that code path does one notice. Trying to further narrow down to the specific optimization going wrong is doable, but unnecessarily time consuming.

Always returning poison therefore can lead to a miscompile. Not propagating poison is only ever a missed optimization, but not an issue of correctness.
It is also not obvious to me that poison semantics, aka deferred UB, is the obvious way that all dialects should be modelling UB. Making it the default would IMO be overfitting to specific kinds of dialects.

It's perfectly fine to "overfit" to "specific dialects" when these "specific dialects" includes basically "all" of the dialect we have in-tree :)
(especially when we're only ever talking about the default settings)

If its a default setting, that's fair.

That is: I'm not sure I understand really the motivation behind your rebuttal here, I don't quite get how to optimize anything without deferred UB.

Let me try to summarize my position: I am warry of silently redefining a function already used in optimizations to potentially create hard to track down miscompiles. If these were not already used, but newly defined or there was an effective way of communicating to users about the change of semantics it'd be easier. Build breakage with new APIs tend to get people to look throught the code and check what the right thing to do is e.g. A PSA would fit this definition as well but that is more likely to be ignored and more effort than just not silently redefining the function.

Regarding handling of UB in the future: I guess my fear was mostly about the silently redefining behaviour in a detremental way. I think the ub dialect is the way forward and a good default and that users will end up having to care about it, worst-case when they use in-tree dialects like arith.

[Hardcode84]

rebase, fix typo

[Hardcode84]

In D159013#4630084, @mehdi_amini wrote:

@Hardcode84 : don't take this as blocking you from landing, but I'd still like to get to end of the argument with @zero9178 here, this is pretty important for any future direction about how we see UB in MLIR upstream (this "default" deferred UB question in particular may come up all the time!)

It's not critical so I can wait. Regarding default value, I'm fine with either variant.

[kuhar]

In D159013#4630341, @zero9178 wrote:

That is true, but people are already using these helpers. From that perspective this is not "opt-in", but rather silently changing the semantics of existing functions. The scenario I am scarred about is a fold implementation, using these functions and incorrectly returning a poison attribute.
This is the kind of bug that in my experience is very hard to track down as it carries through all layers of the compilation and successfully generates runnable code as well. Only when running the code and potentially hitting that code path does one notice. Trying to further narrow down to the specific optimization going wrong is doable, but unnecessarily time consuming.

+1, this really resonates with me and is the reason why I approved the patch in the current form. While we do not promise to keep any downstream code working forever, this could result in surprising and difficult to track down change of behavior, while the alternative of having a few more functions is low cost to pay for the upstream IMO.

[mehdi_amini]

In D159013#4630341, @zero9178 wrote:

In D159013#4630081, @mehdi_amini wrote:

In D159013#4627450, @zero9178 wrote:

In D159013#4627031, @mehdi_amini wrote:

There are likely downstream users of these methods (I believe these are quite old), so I do think this would have an impact. At the very least it'd create a link dependency for them.

It's up to them to opt-out, we don't optimize to avoid breaking downstream, we strive for the best default behavior moving forward.

That's very fair, so maybe I worded this poorly. My concern is not any "churn" downstream in the sense of them requiring actions to keep compiling their code or other refactoring. I like breaking them reguarly as well :-)

OK :)

Rather, I don't think this is the best default behaviour for such generic functions. My first concern is that always returning poison if any operand is a poison can lead to miscompiles. Obvious cases include things like select operations or other "predicated/masked"-like operations.

But it's not like we're doing it for any op generically in the folder here: this is "opt-in"!
People don't have to use these helpers...

That is true, but people are already using these helpers. From that perspective this is not "opt-in", but rather silently changing the semantics of existing functions. The scenario I am scarred about is a fold implementation, using these functions and incorrectly returning a poison attribute.

It can't be silent: they would have a build breakage and decide to depend on the UB dialect!
(also back to the previous point about breaking users)

That is: I'm not sure I understand really the motivation behind your rebuttal here, I don't quite get how to optimize anything without deferred UB.

Let me try to summarize my position: I am warry of silently redefining a function already used in optimizations to potentially create hard to track down miscompiles.

I agree with this concern, I assumed here that this would be a build breakage (we forward declare the UB attribute, folks must include the UB dialect and link to it)

[kuhar]

In D159013#4632140, @mehdi_amini wrote:

In D159013#4630341, @zero9178 wrote:

In D159013#4630081, @mehdi_amini wrote:

In D159013#4627450, @zero9178 wrote:

In D159013#4627031, @mehdi_amini wrote:

There are likely downstream users of these methods (I believe these are quite old), so I do think this would have an impact. At the very least it'd create a link dependency for them.

It's up to them to opt-out, we don't optimize to avoid breaking downstream, we strive for the best default behavior moving forward.

That's very fair, so maybe I worded this poorly. My concern is not any "churn" downstream in the sense of them requiring actions to keep compiling their code or other refactoring. I like breaking them reguarly as well :-)

OK :)

Rather, I don't think this is the best default behaviour for such generic functions. My first concern is that always returning poison if any operand is a poison can lead to miscompiles. Obvious cases include things like select operations or other "predicated/masked"-like operations.

But it's not like we're doing it for any op generically in the folder here: this is "opt-in"!
People don't have to use these helpers...

That is true, but people are already using these helpers. From that perspective this is not "opt-in", but rather silently changing the semantics of existing functions. The scenario I am scarred about is a fold implementation, using these functions and incorrectly returning a poison attribute.

It can't be silent: they would have a build breakage and decide to depend on the UB dialect!
(also back to the previous point about breaking users)

That is: I'm not sure I understand really the motivation behind your rebuttal here, I don't quite get how to optimize anything without deferred UB.

Let me try to summarize my position: I am warry of silently redefining a function already used in optimizations to potentially create hard to track down miscompiles.

I agree with this concern, I assumed here that this would be a build breakage (we forward declare the UB attribute, folks must include the UB dialect and link to it)

What I was concerned about was downstream dialects that use these fold helpers -- with no code change downstream, they would end up propagating poison values created by upstream rewrites. Right now we do not produce poison in any folds/canon patterns AFAIK, but if this changes at some point, it seems risky to me without any explicit opt-in mechanism for everything that ends up consuming results of these upstream dialects.

[mehdi_amini]

In D159013#4632419, @kuhar wrote:

Let me try to summarize my position: I am warry of silently redefining a function already used in optimizations to potentially create hard to track down miscompiles.

I agree with this concern, I assumed here that this would be a build breakage (we forward declare the UB attribute, folks must include the UB dialect and link to it)

What I was concerned about was downstream dialects that use these fold helpers -- with no code change downstream, they would end up propagating poison values created by upstream rewrites. Right now we do not produce poison in any folds/canon patterns AFAIK, but if this changes at some point, it seems risky to me without any explicit opt-in mechanism for everything that ends up consuming results of these upstream dialects.

Sure, we all agree on this, my point was that we should make the default be to use poison and make it a build break so that what you're describing does not exist.

[zero9178]

In D159013#4632563, @mehdi_amini wrote:

In D159013#4632419, @kuhar wrote:

Let me try to summarize my position: I am warry of silently redefining a function already used in optimizations to potentially create hard to track down miscompiles.

I agree with this concern, I assumed here that this would be a build breakage (we forward declare the UB attribute, folks must include the UB dialect and link to it)

What I was concerned about was downstream dialects that use these fold helpers -- with no code change downstream, they would end up propagating poison values created by upstream rewrites. Right now we do not produce poison in any folds/canon patterns AFAIK, but if this changes at some point, it seems risky to me without any explicit opt-in mechanism for everything that ends up consuming results of these upstream dialects.

Sure, we all agree on this, my point was that we should make the default be to use poison and make it a build break so that what you're describing does not exist.

True, while this is likely to lead to a build break, I don't think the prior version was 1) reliable and 2) descriptive enough to warn users. Linkage wise, they're unlikely to get an error if already transitively linking against the UB dialect (e.g. the arith dialect). Compilation wise, the error is also not very helpful: https://godbolt.org/z/5YW439Pva One gets a template error 5 layers deeper saying that ub::PoisonAttr is an incomplete type. I'd argue the reflex of most C++ programms in such a scenario is to include the header defining that type. I have rather little trust that people would see this error, go through the backtrace to find that its coming from their usage of these fold methods, and then start thinking whether returning poison in their given uses is the right thing or not.
They'll rather just try and get it to build, run their tests and sign-off only for it to potentially become a problem in the future.

Besides the redfining issue, there's also an argument one can make that the default general function should be safe by default. Correct me if I am wrong, but the argument for making the poison behaviour the default is that the empirical usage upstream of these folders has so far been safe, rather than it being a conscious design decision. I'd rather have a safe-by-default design.

[mehdi_amini]

In D159013#4632419, @kuhar wrote:

True, while this is likely to lead to a build break, I don't think the prior version was 1) reliable and 2) descriptive enough to warn users.

OK but now I feel we're goal post shifting a bit from the previous objections (silent miscompile).

Linkage wise, they're unlikely to get an error if already transitively linking against the UB dialect (e.g. the arith dialect).

Sure, the goal is to have a compile-time failure, I didn't mention a "linker failure" anywhere. Before getting to the linker, you need to include the headers.

Compilation wise, the error is also not very helpful: https://godbolt.org/z/5YW439Pva One gets a template error 5 layers deeper saying that ub::PoisonAttr is an incomplete type. I'd argue the reflex of most C++ programms in such a scenario is to include the header defining that type. I have rather little trust that people would see this error, go through the backtrace to find that its coming from their usage of these fold methods, and then start thinking whether returning poison in their given uses is the right thing or not.
They'll rather just try and get it to build, run their tests and sign-off only for it to potentially become a problem in the future.

Sure, we can argue that's not ideal, but we're very far from what you were discussing before!

Besides the redfining issue, there's also an argument one can make that the default general function should be safe by default. Correct me if I am wrong, but the argument for making the poison behaviour the default is that the empirical usage upstream of these folders has so far been safe, rather than it being a conscious design decision. I'd rather have a safe-by-default design.

"safe-by-default" is a good principle, it cannot be the only driving principle, in particular if "safe" goes against "useful". There is a balance and good APIs should be "easy to use, hard to misuse".

Just going for safe here cannot be a satisfactory answer.
(especially since the "safe" aspect of it is quite... arguable really)

[Hardcode84]

So, what the final decision?

[mehdi_amini]

In D159013#4638176, @Hardcode84 wrote:

So, what the final decision?

How does the build break looks like for the existing users with forward declaration and default to poison?

[Hardcode84]

[1/27] Building CXX object tools\mlir\lib\Dialect\Math\IR\CMakeFiles\obj.MLIRMathDialect.dir\MathOps.cpp.obj
FAILED: tools/mlir/lib/Dialect/Math/IR/CMakeFiles/obj.MLIRMathDialect.dir/MathOps.cpp.obj
C:\PROGRA~2\MICROS~2\2022\BUILDT~1\VC\Tools\MSVC\1435~1.322\bin\Hostx64\x64\cl.exe  /nologo /TP -DBUILD_EXAMPLES -DGTEST_HAS_RTTI=0 -DMLIR_CUDA_CONVERSIONS_ENABLED=0 -DMLIR_ROCM_CONVERSIONS_ENABLED=0 -DUNICODE -D_CRT_NONSTDC_NO_DEPRECATE -D_CRT_NONSTDC_NO_WARNINGS -D_CRT_SECURE_NO_DEPRECATE -D_CRT_SECURE_NO_WARNINGS -D_GLIBCXX_ASSERTIONS -D_HAS_EXCEPTIONS=0 -D_LIBCPP_ENABLE_HARDENED_MODE -D_SCL_SECURE_NO_DEPRECATE -D_SCL_SECURE_NO_WARNINGS -D_UNICODE -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -ID:\projs\llvm\llvm-build\tools\mlir\lib\Dialect\Math\IR -ID:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR -ID:\projs\llvm\llvm-build\include -ID:\projs\llvm\llvm-project\llvm\include -ID:\projs\llvm\llvm-project\mlir\include -ID:\projs\llvm\llvm-build\tools\mlir\include /DWIN32 /D_WINDOWS   /Zc:inline /Zc:preprocessor /Zc:__cplusplus /Oi /bigobj /permissive- /W4 -wd4141 -wd4146 -wd4244 -wd4267 -wd4291 -wd4351 -wd4456 -wd4457 -wd4458 -wd4459 -wd4503 -wd4624 -wd4722 -wd4100 -wd4127 -wd4512 -wd4505 -wd4610 -wd4510 -wd4702 -wd4245 -wd4706 -wd4310 -wd4701 -wd4703 -wd4389 -wd4611 -wd4805 -wd4204 -wd4577 -wd4091 -wd4592 -wd4319 -wd4709 -wd5105 -wd4324 -w14062 -we4238 /Gw /O2 /Ob2 -std:c++17 -MD  /EHs-c- /GR- -UNDEBUG /showIncludes /Fotools\mlir\lib\Dialect\Math\IR\CMakeFiles\obj.MLIRMathDialect.dir\MathOps.cpp.obj /Fdtools\mlir\lib\Dialect\Math\IR\CMakeFiles\obj.MLIRMathDialect.dir\ /FS -c D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp
D:\projs\llvm\llvm-project\mlir\include\mlir/IR/BuiltinAttributes.h(359): warning C4996: 'std::complex<llvm::APFloat>::complex': warning STL4037: The effect of instantiating the template std::complex for any type other than float, double, or long double is unspecified. You can define _SILENCE_NONFLOATING_COMPLEX_DEPRECATION_WARNING to suppress this warning.
D:\projs\llvm\llvm-project\llvm\include\llvm/Support/Casting.h(657): error C2027: use of undefined type 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\include\mlir/Dialect/CommonFolders.h(26): note: see declaration of 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp(526): note: see reference to function template instantiation 'decltype(auto) llvm::dyn_cast<mlir::ub::PoisonAttr,mlir::Attribute>(From &)' being compiled
        with
        [
            From=mlir::Attribute
        ]
D:\projs\llvm\llvm-project\llvm\include\llvm/Support/Casting.h(655): error C3487: 'void': all return expressions must deduce to the same type: previously it was 'To'
        with
        [
            To=mlir::ub::PoisonAttr
        ]
D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp(526): error C2027: use of undefined type 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\include\mlir/Dialect/CommonFolders.h(26): note: see declaration of 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp(526): error C2027: use of undefined type 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\include\mlir/Dialect/CommonFolders.h(26): note: see declaration of 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp(526): error C2451: a conditional expression of type 'To' is not valid
        with
        [
            To=mlir::ub::PoisonAttr
        ]
D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp(526): note: use of undefined type 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\include\mlir/Dialect/CommonFolders.h(26): note: see declaration of 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp(527): error C2039: 'PoisonOp': is not a member of 'mlir::ub'
D:\projs\llvm\llvm-project\mlir\include\mlir/Dialect/CommonFolders.h(25): note: see declaration of 'mlir::ub'
D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp(527): error C2065: 'PoisonOp': undeclared identifier
D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp(527): error C2672: 'mlir::OpBuilder::create': no matching overloaded function found
D:\projs\llvm\llvm-project\mlir\include\mlir/IR/Builders.h(491): note: could be 'OpTy mlir::OpBuilder::create(mlir::Location,Args &&...)'
D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp(527): note: 'mlir::OpBuilder::create': invalid template argument for 'OpTy', type expected
ninja: build stopped: subcommand failed.

VS 2022

[Hardcode84]

[1/27] Building CXX object tools\mlir\lib\Dialect\Math\IR\CMakeFiles\obj.MLIRMathDialect.dir\MathOps.cpp.obj
FAILED: tools/mlir/lib/Dialect/Math/IR/CMakeFiles/obj.MLIRMathDialect.dir/MathOps.cpp.obj
C:\PROGRA~2\MICROS~2\2022\BUILDT~1\VC\Tools\MSVC\1435~1.322\bin\Hostx64\x64\cl.exe  /nologo /TP -DBUILD_EXAMPLES -DGTEST_HAS_RTTI=0 -DMLIR_CUDA_CONVERSIONS_ENABLED=0 -DMLIR_ROCM_CONVERSIONS_ENABLED=0 -DUNICODE -D_CRT_NONSTDC_NO_DEPRECATE -D_CRT_NONSTDC_NO_WARNINGS -D_CRT_SECURE_NO_DEPRECATE -D_CRT_SECURE_NO_WARNINGS -D_GLIBCXX_ASSERTIONS -D_HAS_EXCEPTIONS=0 -D_LIBCPP_ENABLE_HARDENED_MODE -D_SCL_SECURE_NO_DEPRECATE -D_SCL_SECURE_NO_WARNINGS -D_UNICODE -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -ID:\projs\llvm\llvm-build\tools\mlir\lib\Dialect\Math\IR -ID:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR -ID:\projs\llvm\llvm-build\include -ID:\projs\llvm\llvm-project\llvm\include -ID:\projs\llvm\llvm-project\mlir\include -ID:\projs\llvm\llvm-build\tools\mlir\include /DWIN32 /D_WINDOWS   /Zc:inline /Zc:preprocessor /Zc:__cplusplus /Oi /bigobj /permissive- /W4 -wd4141 -wd4146 -wd4244 -wd4267 -wd4291 -wd4351 -wd4456 -wd4457 -wd4458 -wd4459 -wd4503 -wd4624 -wd4722 -wd4100 -wd4127 -wd4512 -wd4505 -wd4610 -wd4510 -wd4702 -wd4245 -wd4706 -wd4310 -wd4701 -wd4703 -wd4389 -wd4611 -wd4805 -wd4204 -wd4577 -wd4091 -wd4592 -wd4319 -wd4709 -wd5105 -wd4324 -w14062 -we4238 /Gw /O2 /Ob2 -std:c++17 -MD  /EHs-c- /GR- -UNDEBUG /showIncludes /Fotools\mlir\lib\Dialect\Math\IR\CMakeFiles\obj.MLIRMathDialect.dir\MathOps.cpp.obj /Fdtools\mlir\lib\Dialect\Math\IR\CMakeFiles\obj.MLIRMathDialect.dir\ /FS -c D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp
D:\projs\llvm\llvm-project\mlir\include\mlir/IR/BuiltinAttributes.h(359): warning C4996: 'std::complex<llvm::APFloat>::complex': warning STL4037: The effect of instantiating the template std::complex for any type other than float, double, or long double is unspecified. You can define _SILENCE_NONFLOATING_COMPLEX_DEPRECATION_WARNING to suppress this warning.
D:\projs\llvm\llvm-project\mlir\include\mlir/IR/Attributes.h(406): error C2027: use of undefined type 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\include\mlir/Dialect/CommonFolders.h(26): note: see declaration of 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\include\mlir/IR/Attributes.h(399): note: while compiling class template member function 'bool llvm::CastInfo<mlir::ub::PoisonAttr,const T,void>::isPossible(mlir::Attribute)'
        with
        [
            T=mlir::Attribute
        ]
D:\projs\llvm\llvm-project\llvm\include\llvm/Support/Casting.h(549): note: see reference to function template instantiation 'bool llvm::CastInfo<mlir::ub::PoisonAttr,const T,void>::isPossible(mlir::Attribute)' being compiled
        with
        [
            T=mlir::Attribute
        ]
D:\projs\llvm\llvm-project\llvm\include\llvm/Support/Casting.h(548): note: see reference to class template instantiation 'llvm::CastInfo<mlir::ub::PoisonAttr,const T,void>' being compiled
        with
        [
            T=mlir::Attribute
        ]
D:\projs\llvm\llvm-project\mlir\include\mlir/Dialect/CommonFolders.h(186): note: see reference to function template instantiation 'bool llvm::isa<mlir::ub::PoisonAttr,T>(const From &)' being compiled
        with
        [
            T=mlir::Attribute,
            From=mlir::Attribute
        ]
D:\projs\llvm\llvm-project\mlir\include\mlir/Dialect/CommonFolders.h(338): note: see reference to function template instantiation 'mlir::Attribute mlir::constFoldUnaryOpConditional<mlir::FloatAttr,mlir::FloatAttr::ValueType,mlir::ub::PoisonAttr,mlir::math::AtanOp::fold::<lambda_1>>(llvm::ArrayRef<T>,CalculationT &&)' being compiled
        with
        [
            T=mlir::Attribute,
            CalculationT=mlir::math::AtanOp::fold::<lambda_1>
        ]
D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp(59): note: see reference to function template instantiation 'mlir::Attribute mlir::constFoldUnaryOpConditionalPoison<mlir::FloatAttr,mlir::FloatAttr::ValueType,mlir::math::AtanOp::fold::<lambda_1>>(llvm::ArrayRef<T>,CalculationT &&)' being compiled
        with
        [
            T=mlir::Attribute,
            CalculationT=mlir::math::AtanOp::fold::<lambda_1>
        ]
D:\projs\llvm\llvm-project\mlir\include\mlir/IR/Attributes.h(406): error C3861: 'classof': identifier not found
ninja: build stopped: subcommand failed.

forgot to comment out materializer

[mehdi_amini]

This comment was removed by mehdi_amini.

[mehdi_amini]

Thanks!
To me seeing:

D:\projs\llvm\llvm-project\mlir\include\mlir/IR/Attributes.h(406): error C2027: use of undefined type 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\include\mlir/Dialect/CommonFolders.h(26): note: see declaration of 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\include\mlir/IR/Attributes.h(399): note: while compiling class template member function 'bool llvm::CastInfo<mlir::ub::PoisonAttr,const T,void>::isPossible(mlir::Attribute)'
        with

Is on-par with what I expect people to see and trigger them to figure out whether this is what they want.
Worth mentioning also that even if they don't think further and just add the right includes, it'll only kick in if they have poison in their IR already.

@zero9178 : what kind of error would you want to see here? Do you have a suggestion that would make this harder to misuse for the users while preserving the "right" default in general?

[zero9178]

In D159013#4638950, @mehdi_amini wrote:

Thanks!
To me seeing:

D:\projs\llvm\llvm-project\mlir\include\mlir/IR/Attributes.h(406): error C2027: use of undefined type 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\include\mlir/Dialect/CommonFolders.h(26): note: see declaration of 'mlir::ub::PoisonAttr'
D:\projs\llvm\llvm-project\mlir\include\mlir/IR/Attributes.h(399): note: while compiling class template member function 'bool llvm::CastInfo<mlir::ub::PoisonAttr,const T,void>::isPossible(mlir::Attribute)'
        with

Is on-par with what I expect people to see and trigger them to figure out whether this is what they want.
Worth mentioning also that even if they don't think further and just add the right includes, it'll only kick in if they have poison in their IR already.

@zero9178 : what kind of error would you want to see here? Do you have a suggestion that would make this harder to misuse for the users while preserving the "right" default in general?

I think a verbose message would be a good compromise. One could detect it with something liek https://devblogs.microsoft.com/oldnewthing/20190710-00/?p=102678 (usual sfinae pattern used on sizeof to check whether the type is complete) if that is not too much effort.
The error message could read something like "PoisonAttr must be a complete type. If any of the operands are an attribute of this kind it will be propagated to the result.".
This also has the advantage of bringing the error location right to the offending function and it being a non-standard C++ error that users shouldn't have a reflex fix for.

I think this would be an okay compromise so that we can land this

[Hardcode84]

I'm not a big fan of those SFINAE tricks.

[kuhar]

In D159013#4639471, @Hardcode84 wrote:

I'm not a big fan of those SFINAE tricks.

Would you be more comfortable with it if instead of doing an one-of thing here, we first added a type completeness check to STLExtras?

[Hardcode84]

In D159013#4639651, @kuhar wrote:

In D159013#4639471, @Hardcode84 wrote:

I'm not a big fan of those SFINAE tricks.

Would you be more comfortable with it if instead of doing an one-of thing here, we first added a type completeness check to STLExtras?

I'm fine with it as long as ensuring it works consistently across different compilers will not be my responsibility :)

Also it will probably add (small) compilation overhead for everyone, I'm fine with this too, but others may argue.

[kuhar]

In D159013#4639664, @Hardcode84 wrote:

In D159013#4639651, @kuhar wrote:

In D159013#4639471, @Hardcode84 wrote:

I'm not a big fan of those SFINAE tricks.

Would you be more comfortable with it if instead of doing an one-of thing here, we first added a type completeness check to STLExtras?

I'm fine with it as long as ensuring it works consistently across different compilers will not be my responsibility :)

Also it will probably add (small) compilation overhead for everyone, I'm fine with this too, but others may argue.

Here you go: https://github.com/llvm/llvm-project/pull/65495

[kuhar]

@Hardcode84 I've just landed llvm::is_incomplete_v

[Hardcode84]

better error message

[Hardcode84]

Now build failure looks like this:

[1/27] Building CXX object tools\mlir\lib\Dialect\Math\IR\CMakeFiles\obj.MLIRMathDialect.dir\MathOps.cpp.obj
FAILED: tools/mlir/lib/Dialect/Math/IR/CMakeFiles/obj.MLIRMathDialect.dir/MathOps.cpp.obj
C:\PROGRA~2\MICROS~2\2022\BUILDT~1\VC\Tools\MSVC\1435~1.322\bin\Hostx64\x64\cl.exe  /nologo /TP -DBUILD_EXAMPLES -DGTEST_HAS_RTTI=0 -DMLIR_CUDA_CONVERSIONS_ENABLED=0 -DMLIR_ROCM_CONVERSIONS_ENABLED=0 -DUNICODE -D_CRT_NONSTDC_NO_DEPRECATE -D_CRT_NONSTDC_NO_WARNINGS -D_CRT_SECURE_NO_DEPRECATE -D_CRT_SECURE_NO_WARNINGS -D_GLIBCXX_ASSERTIONS -D_HAS_EXCEPTIONS=0 -D_LIBCPP_ENABLE_HARDENED_MODE -D_SCL_SECURE_NO_DEPRECATE -D_SCL_SECURE_NO_WARNINGS -D_UNICODE -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -ID:\projs\llvm\llvm-build\tools\mlir\lib\Dialect\Math\IR -ID:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR -ID:\projs\llvm\llvm-build\include -ID:\projs\llvm\llvm-project\llvm\include -ID:\projs\llvm\llvm-project\mlir\include -ID:\projs\llvm\llvm-build\tools\mlir\include /DWIN32 /D_WINDOWS   /Zc:inline /Zc:preprocessor /Zc:__cplusplus /Oi /bigobj /permissive- /W4 -wd4141 -wd4146 -wd4244 -wd4267 -wd4291 -wd4351 -wd4456 -wd4457 -wd4458 -wd4459 -wd4503 -wd4624 -wd4722 -wd4100 -wd4127 -wd4512 -wd4505 -wd4610 -wd4510 -wd4702 -wd4245 -wd4706 -wd4310 -wd4701 -wd4703 -wd4389 -wd4611 -wd4805 -wd4204 -wd4577 -wd4091 -wd4592 -wd4319 -wd4709 -wd5105 -wd4324 -w14062 -we4238 /Gw /O2 /Ob2 -std:c++17 -MD  /EHs-c- /GR- -UNDEBUG /showIncludes /Fotools\mlir\lib\Dialect\Math\IR\CMakeFiles\obj.MLIRMathDialect.dir\MathOps.cpp.obj /Fdtools\mlir\lib\Dialect\Math\IR\CMakeFiles\obj.MLIRMathDialect.dir\ /FS -c D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp
D:\projs\llvm\llvm-project\mlir\include\mlir/IR/BuiltinAttributes.h(359): warning C4996: 'std::complex<llvm::APFloat>::complex': warning STL4037: The effect of instantiating the template std::complex for any type other than float, double, or long double is unspecified. You can define _SILENCE_NONFLOATING_COMPLEX_DEPRECATION_WARNING to suppress this warning.
D:\projs\llvm\llvm-project\mlir\include\mlir/Dialect/CommonFolders.h(200): error C2338: static_assert failed: 'PoisonAttr is undefined, either add a dependency on UB dialect or pass void as template argument to opt-out from poison semantics.'
D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp(59): note: see reference to function template instantiation 'mlir::Attribute mlir::constFoldUnaryOpConditional<mlir::FloatAttr,mlir::FloatAttr::ValueType,mlir::ub::PoisonAttr,mlir::math::AtanOp::fold::<lambda_1>>(llvm::ArrayRef<T>,CalculationT &&)' being compiled
        with
        [
            T=mlir::Attribute,
            CalculationT=mlir::math::AtanOp::fold::<lambda_1>
        ]
D:\projs\llvm\llvm-project\mlir\include\mlir/Dialect/CommonFolders.h(127): error C2338: static_assert failed: 'PoisonAttr is undefined, either add a dependency on UB dialect or pass void as template argument to opt-out from poison semantics.'
D:\projs\llvm\llvm-project\mlir\lib\Dialect\Math\IR\MathOps.cpp(82): note: see reference to function template instantiation 'mlir::Attribute mlir::constFoldBinaryOpConditional<mlir::FloatAttr,mlir::FloatAttr::ValueType,mlir::ub::PoisonAttr,mlir::math::Atan2Op::fold::<lambda_1>>(llvm::ArrayRef<T>,CalculationT &&)' being compiled
        with
        [
            T=mlir::Attribute,
            CalculationT=mlir::math::Atan2Op::fold::<lambda_1>
        ]
ninja: build stopped: subcommand failed.

[mehdi_amini]

Sweet!

mlir/include/mlir/Dialect/CommonFolders.h

 
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include <optional>
 
 namespace mlir {
+namespace ub {
+class PoisonAttr;
+}
 /// Performs constant folding `calculate` with element-wise behavior on the two
 /// attributes in `operands` and returns the result if possible.
 /// Uses `resultType` for the type of the returned attribute.
+/// Optional PoisonAttr template argument allows to specify 'poison' attribute
+/// which will be directly propagated to result.
 template <class AttrElementT,
           class ElementValueT = typename AttrElementT::ValueType,
+          class PoisonAttr = ub::PoisonAttr,
           class CalculationT = function_ref<
               std::optional<ElementValueT>(ElementValueT, ElementValueT)>>
 Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
                                        Type resultType,
-                                       const CalculationT &calculate) {
+                                       CalculationT &&calculate) {
   assert(operands.size() == 2 && "binary op takes two operands");
+  static_assert(
+      std::is_void_v<PoisonAttr> || !llvm::is_incomplete_v<PoisonAttr>,
+      "PoisonAttr is undefined, either add a dependency on UB dialect or pass "
+      "void as template argument to opt-out from poison semantics.");
+  if constexpr (!std::is_void_v<PoisonAttr>) {
+    if (isa_and_nonnull<PoisonAttr>(operands[0]))
+      return operands[0];
+
+    if (isa_and_nonnull<PoisonAttr>(operands[1]))
+      return operands[1];
+  }
+
   if (!resultType || !operands[0] || !operands[1])
     return {};
 
   if (isa<AttrElementT>(operands[0]) && isa<AttrElementT>(operands[1])) {
     auto lhs = cast<AttrElementT>(operands[0]);
     auto rhs = cast<AttrElementT>(operands[1]);
     if (lhs.getType() != rhs.getType())
       return {};
Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
   }
   return {};
 }
 
 /// Performs constant folding `calculate` with element-wise behavior on the two
 /// attributes in `operands` and returns the result if possible.
 /// Uses the operand element type for the element type of the returned
 /// attribute.
+/// Optional PoisonAttr template argument allows to specify 'poison' attribute
+/// which will be directly propagated to result.
 template <class AttrElementT,
           class ElementValueT = typename AttrElementT::ValueType,
+          class PoisonAttr = ub::PoisonAttr,
           class CalculationT = function_ref<
               std::optional<ElementValueT>(ElementValueT, ElementValueT)>>
 Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
-                                       const CalculationT &calculate) {
+                                       CalculationT &&calculate) {
   assert(operands.size() == 2 && "binary op takes two operands");
+  static_assert(
+      std::is_void_v<PoisonAttr> || !llvm::is_incomplete_v<PoisonAttr>,
+      "PoisonAttr is undefined, either add a dependency on UB dialect or pass "
+      "void as template argument to opt-out from poison semantics.");
+  if constexpr (!std::is_void_v<PoisonAttr>) {
+    if (isa_and_nonnull<PoisonAttr>(operands[0]))
+      return operands[0];
+
+    if (isa_and_nonnull<PoisonAttr>(operands[1]))
+      return operands[1];
+  }
+
   auto getResultType = [](Attribute attr) -> Type {
     if (auto typed = dyn_cast_or_null<TypedAttr>(attr))
       return typed.getType();
     return {};
   };
 
   Type lhsType = getResultType(operands[0]);
   Type rhsType = getResultType(operands[1]);
   if (!lhsType || !rhsType)
     return {};
   if (lhsType != rhsType)
     return {};
 
-  return constFoldBinaryOpConditional<AttrElementT, ElementValueT,
-                                      CalculationT>(operands, lhsType,
-                                                    calculate);
+  return constFoldBinaryOpConditional<AttrElementT, ElementValueT, PoisonAttr,
+                                      CalculationT>(
+      operands, lhsType, std::forward<CalculationT>(calculate));
 }
 
 template <class AttrElementT,
           class ElementValueT = typename AttrElementT::ValueType,
+          class PoisonAttr = void,
           class CalculationT =
               function_ref<ElementValueT(ElementValueT, ElementValueT)>>
 Attribute constFoldBinaryOp(ArrayRef<Attribute> operands, Type resultType,
-                            const CalculationT &calculate) {
-  return constFoldBinaryOpConditional<AttrElementT>(
+                            CalculationT &&calculate) {
+  return constFoldBinaryOpConditional<AttrElementT, ElementValueT, PoisonAttr>(
       operands, resultType,
       [&](ElementValueT a, ElementValueT b) -> std::optional<ElementValueT> {
         return calculate(a, b);
       });
 }
 
 template <class AttrElementT,
           class ElementValueT = typename AttrElementT::ValueType,
+          class PoisonAttr = ub::PoisonAttr,
           class CalculationT =
               function_ref<ElementValueT(ElementValueT, ElementValueT)>>
 Attribute constFoldBinaryOp(ArrayRef<Attribute> operands,
-                            const CalculationT &calculate) {
-  return constFoldBinaryOpConditional<AttrElementT>(
+                            CalculationT &&calculate) {
+  return constFoldBinaryOpConditional<AttrElementT, ElementValueT, PoisonAttr>(
       operands,
       [&](ElementValueT a, ElementValueT b) -> std::optional<ElementValueT> {
         return calculate(a, b);
       });
 }
 
 /// Performs constant folding `calculate` with element-wise behavior on the one
 /// attributes in `operands` and returns the result if possible.
+/// Optional PoisonAttr template argument allows to specify 'poison' attribute
+/// which will be directly propagated to result.
 template <class AttrElementT,
           class ElementValueT = typename AttrElementT::ValueType,
+          class PoisonAttr = ub::PoisonAttr,
           class CalculationT =
               function_ref<std::optional<ElementValueT>(ElementValueT)>>
 Attribute constFoldUnaryOpConditional(ArrayRef<Attribute> operands,
-                                      const CalculationT &&calculate) {
+                                      CalculationT &&calculate) {
   assert(operands.size() == 1 && "unary op takes one operands");
   if (!operands[0])
     return {};
 
+  static_assert(
+      std::is_void_v<PoisonAttr> || !llvm::is_incomplete_v<PoisonAttr>,
+      "PoisonAttr is undefined, either add a dependency on UB dialect or pass "
+      "void as template argument to opt-out from poison semantics.");
+  if constexpr (!std::is_void_v<PoisonAttr>) {
+    if (isa<PoisonAttr>(operands[0]))
+      return operands[0];
+  }
+
   if (isa<AttrElementT>(operands[0])) {
     auto op = cast<AttrElementT>(operands[0]);
 
     auto res = calculate(op.getValue());
     if (!res)
       return {};
     return AttrElementT::get(op.getType(), *res);
   }
if (isa<SplatElementsAttr>(operands[0])) {
     }
     return DenseElementsAttr::get(op.getShapedType(), elementResults);
   }
   return {};
 }
 
 template <class AttrElementT,
           class ElementValueT = typename AttrElementT::ValueType,
+          class PoisonAttr = ub::PoisonAttr,
           class CalculationT = function_ref<ElementValueT(ElementValueT)>>
 Attribute constFoldUnaryOp(ArrayRef<Attribute> operands,
-                           const CalculationT &&calculate) {
-  return constFoldUnaryOpConditional<AttrElementT>(
+                           CalculationT &&calculate) {
+  return constFoldUnaryOpConditional<AttrElementT, ElementValueT, PoisonAttr>(
       operands, [&](ElementValueT a) -> std::optional<ElementValueT> {
         return calculate(a);
       });
 }
 
 template <
     class AttrElementT, class TargetAttrElementT,
     class ElementValueT = typename AttrElementT::ValueType,
     class TargetElementValueT = typename TargetAttrElementT::ValueType,
+    class PoisonAttr = ub::PoisonAttr,
     class CalculationT = function_ref<TargetElementValueT(ElementValueT, bool)>>
 Attribute constFoldCastOp(ArrayRef<Attribute> operands, Type resType,
-                          const CalculationT &calculate) {
+                          CalculationT &&calculate) {
   assert(operands.size() == 1 && "Cast op takes one operand");
   if (!operands[0])
     return {};
 
+  static_assert(
+      std::is_void_v<PoisonAttr> || !llvm::is_incomplete_v<PoisonAttr>,
+      "PoisonAttr is undefined, either add a dependency on UB dialect or pass "
+      "void as template argument to opt-out from poison semantics.");
+  if constexpr (!std::is_void_v<PoisonAttr>) {
+    if (isa<PoisonAttr>(operands[0]))
+      return operands[0];
+  }
+
   if (isa<AttrElementT>(operands[0])) {
     auto op = cast<AttrElementT>(operands[0]);
     bool castStatus = true;
     auto res = calculate(op.getValue(), castStatus);
     if (!castStatus)
       return {};
     return TargetAttrElementT::get(resType, res);
   }
for (size_t i = 0, e = op.getNumElements(); i < e; ++i, ++opIt) {
         return {};
       elementResults.push_back(elt);
     }
 
     return DenseElementsAttr::get(cast<ShapedType>(resType), elementResults);
   }
   return {};
 }
-
 } // namespace mlir
 
 #endif // MLIR_DIALECT_COMMONFOLDERS_H

mlir/lib/Dialect/Arith/IR/ArithDialect.cpp

 //===- ArithDialect.cpp - MLIR Arith dialect implementation -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::arith;
 
#include "mlir/Dialect/Arith/IR/ArithOpsAttributes.cpp.inc"
       >();
   addInterfaces<ArithInlinerInterface>();
 }
 
 /// Materialize an integer or floating point constant.
 Operation *arith::ArithDialect::materializeConstant(OpBuilder &builder,
                                                     Attribute value, Type type,
                                                     Location loc) {
+  if (auto poison = dyn_cast<ub::PoisonAttr>(value))
+    return builder.create<ub::PoisonOp>(loc, type, poison);
+
   return ConstantOp::materialize(builder, value, type, loc);
 }

mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp

 //===- ControlFlowOps.cpp - ControlFlow Operations ------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/CommonFolders.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"

mlir/lib/Dialect/Math/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRMathDialect
   MathOps.cpp
   MathDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Math
 
   DEPENDS
   MLIRMathOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRDialect
   MLIRIR
+  MLIRUBDialect
   )

mlir/lib/Dialect/Math/IR/MathDialect.cpp

 //===- MathDialect.cpp - MLIR dialect for Math implementation -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/Transforms/InliningUtils.h"
 
 using namespace mlir;
 using namespace mlir::math;
 
 #include "mlir/Dialect/Math/IR/MathOpsDialect.cpp.inc"
 
 namespace {

mlir/lib/Dialect/Math/IR/MathOps.cpp

 //===- MathOps.cpp - MLIR operations for math implementation --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/CommonFolders.h"
 #include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/IR/Builders.h"
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::math;
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
return constFoldUnaryOpConditional<FloatAttr>(
         }
       });
 }
 
 /// Materialize an integer or floating point constant.
 Operation *math::MathDialect::materializeConstant(OpBuilder &builder,
                                                   Attribute value, Type type,
                                                   Location loc) {
+  if (auto poison = dyn_cast<ub::PoisonAttr>(value))
+    return builder.create<ub::PoisonOp>(loc, type, poison);
+
   return arith::ConstantOp::materialize(builder, value, type, loc);
 }

mlir/lib/Dialect/SPIRV/IR/CMakeLists.txt

add_mlir_dialect_library(MLIRSPIRVDialect
   LINK_LIBS PUBLIC
   MLIRControlFlowInterfaces
   MLIRFunctionInterfaces
   MLIRIR
   MLIRParser
   MLIRSideEffectInterfaces
   MLIRSupport
   MLIRTransforms
+  MLIRUBDialect
 )

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

 #include <optional>
 #include <utility>
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 
 #include "mlir/Dialect/CommonFolders.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVectorExtras.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp

 
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 
 #include "SPIRVParsingUtils.h"
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Parser/Parser.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Sequence.h"

 
 //===----------------------------------------------------------------------===//
 // Constant
 //===----------------------------------------------------------------------===//
 
 Operation *SPIRVDialect::materializeConstant(OpBuilder &builder,
                                              Attribute value, Type type,
                                              Location loc) {
+  if (auto poison = dyn_cast<ub::PoisonAttr>(value))
+    return builder.create<ub::PoisonOp>(loc, type, poison);
+
   if (!spirv::ConstantOp::isBuildableWith(type))
     return nullptr;
 
   return builder.create<spirv::ConstantOp>(loc, type, value);
 }
 
 //===----------------------------------------------------------------------===//
 // Shader Interface ABI

mlir/lib/Dialect/Shape/IR/CMakeLists.txt

add_mlir_dialect_library(MLIRShapeDialect
   MLIRControlFlowInterfaces
   MLIRDialect
   MLIRFuncDialect
   MLIRFunctionInterfaces
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRSideEffectInterfaces
   MLIRTensorDialect
+  MLIRUBDialect
   )

mlir/lib/Dialect/Shape/IR/Shape.cpp

 //===- Shape.cpp - MLIR Shape Operations ----------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include <utility>
 
 #include "mlir/Dialect/Shape/IR/Shape.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/CommonFolders.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Traits.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/FunctionImplementation.h"
 #include "mlir/Transforms/InliningUtils.h"
#include "mlir/Dialect/Shape/IR/ShapeOpsTypes.cpp.inc"
   // still evolving it makes it simple to start with an unregistered ops and
   // try different variants before actually defining the op.
   allowUnknownOperations();
 }
 
 Operation *ShapeDialect::materializeConstant(OpBuilder &builder,
                                              Attribute value, Type type,
                                              Location loc) {
+  if (auto poison = dyn_cast<ub::PoisonAttr>(value))
+    return builder.create<ub::PoisonOp>(loc, type, poison);
+
   if (llvm::isa<ShapeType>(type) || isExtentTensorType(type))
     return builder.create<ConstShapeOp>(
         loc, type, llvm::cast<DenseIntElementsAttr>(value));
   if (llvm::isa<SizeType>(type))
     return builder.create<ConstSizeOp>(loc, type,
                                        llvm::cast<IntegerAttr>(value));
   if (llvm::isa<WitnessType>(type))
     return builder.create<ConstWitnessOp>(loc, type,
                                           llvm::cast<BoolAttr>(value));
+
   return arith::ConstantOp::materialize(builder, value, type, loc);
 }
 
 LogicalResult ShapeDialect::verifyOperationAttribute(Operation *op,
                                                      NamedAttribute attribute) {
   // Verify shape.lib attribute.
   if (attribute.getName() == "shape.lib") {
     if (!op->hasTrait<OpTrait::SymbolTable>())

mlir/test/Dialect/Arith/canonicalize.mlir

func.func @selectOfPoison(%cond : i1, %arg: i32) -> (i32, i32, i32, i32) {
 
   // Check that constant folding is applied prior to poison handling.
   %true = arith.constant true
   %false = arith.constant false
   %select3 = arith.select %true, %poison, %arg : i32
   %select4 = arith.select %false, %poison, %arg : i32
   return %select1, %select2, %select3, %select4 : i32, i32, i32, i32
 }
+
+// CHECK-LABEL: @addi_poison1
+//       CHECK:   %[[P:.*]] = ub.poison : i32
+//       CHECK:   return %[[P]]
+func.func @addi_poison1(%arg: i32) -> i32 {
+  %0 = ub.poison : i32
+  %1 = arith.addi %0, %arg : i32
+  return %1 : i32
+}
+
+// CHECK-LABEL: @addi_poison2
+//       CHECK:   %[[P:.*]] = ub.poison : i32
+//       CHECK:   return %[[P]]
+func.func @addi_poison2(%arg: i32) -> i32 {
+  %0 = ub.poison : i32
+  %1 = arith.addi %arg, %0 : i32
+  return %1 : i32
+}
+
+// CHECK-LABEL: @addf_poison1
+//       CHECK:   %[[P:.*]] = ub.poison : f32
+//       CHECK:   return %[[P]]
+func.func @addf_poison1(%arg: f32) -> f32 {
+  %0 = ub.poison : f32
+  %1 = arith.addf %0, %arg : f32
+  return %1 : f32
+}
+
+// CHECK-LABEL: @addf_poison2
+//       CHECK:   %[[P:.*]] = ub.poison : f32
+//       CHECK:   return %[[P]]
+func.func @addf_poison2(%arg: f32) -> f32 {
+  %0 = ub.poison : f32
+  %1 = arith.addf %arg, %0 : f32
+  return %1 : f32
+}
+
+
+// CHECK-LABEL: @negf_poison
+//       CHECK:   %[[P:.*]] = ub.poison : f32
+//       CHECK:   return %[[P]]
+func.func @negf_poison() -> f32 {
+  %0 = ub.poison : f32
+  %1 = arith.negf %0 : f32
+  return %1 : f32
+}
+
+// CHECK-LABEL: @extsi_poison
+//       CHECK:   %[[P:.*]] = ub.poison : i64
+//       CHECK:   return %[[P]]
+func.func @extsi_poison() -> i64 {
+  %0 = ub.poison : i32
+  %1 = arith.extsi %0 : i32 to i64
+  return %1 : i64
+}

mlir/test/Dialect/Math/canonicalize.mlir

 // CHECK-LABEL: @erf_fold_vec
 // CHECK-NEXT: %[[cst:.+]] = arith.constant dense<[0.000000e+00, {{0.84[0-9]+|8.4[0-9]+e-01}}, 0.000000e+00, {{0.84[0-9]+|8.4[0-9]+e-01}}]> : vector<4xf32>
 // CHECK-NEXT:   return %[[cst]]
 func.func @erf_fold_vec() -> (vector<4xf32>) {
   %v1 = arith.constant dense<[0.0, 1.0, 0.0, 1.0]> : vector<4xf32>
   %0 = math.erf %v1 : vector<4xf32>
   return %0 : vector<4xf32>
 }
+
+// CHECK-LABEL: @abs_poison
+//       CHECK:   %[[P:.*]] = ub.poison : f32
+//       CHECK:   return %[[P]]
+func.func @abs_poison() -> f32 {
+  %0 = ub.poison : f32
+  %1 = math.absf %0 : f32
+  return %1 : f32
+}

mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir

func.func @const_fold_vector_iadd() -> vector<3xi32> {
   %vc1 = spirv.Constant dense<[42, -55, 127]> : vector<3xi32>
   %vc2 = spirv.Constant dense<[-3, -15, 28]> : vector<3xi32>
 
   // CHECK: spirv.Constant dense<[39, -70, 155]>
   %0 = spirv.IAdd %vc1, %vc2 : vector<3xi32>
   return %0: vector<3xi32>
 }
 
+// CHECK-LABEL: @iadd_poison
+//       CHECK:   %[[P:.*]] = ub.poison : i32
+//       CHECK:   return %[[P]]
+func.func @iadd_poison(%arg0: i32) -> i32 {
+  %0 = ub.poison : i32
+  %1 = spirv.IAdd %arg0, %0 : i32
+  return %1: i32
+}
+
 // -----
 
 //===----------------------------------------------------------------------===//
 // spirv.IMul
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @imul_zero_one
 // CHECK-SAME: (%[[ARG:.*]]: i32)

mlir/test/Dialect/Shape/canonicalize.mlir

  %c1 = arith.constant 1 : index
 // CHECK:        %[[C1:.*]] = arith.constant 1
  %shape = shape.shape_of %arg0 : tensor<?x?xf64> -> tensor<2xindex>
 // CHECK:        %[[DIM:.*]] = tensor.dim %[[ARG0]], %[[C1]]
  %result = tensor.extract %shape[%c1] : tensor<2xindex>
 // CHECK:        return %[[DIM]]
  return %result : index
 }
+
+
+// -----
+
+// CHECK-LABEL: @add_poison
+//       CHECK:   %[[P:.*]] = ub.poison : !shape.siz
+//       CHECK:   return %[[P]]
+func.func @add_poison() -> !shape.size {
+  %1 = shape.const_size 2
+  %2 = ub.poison : !shape.size
+  %result = shape.add %1, %2 : !shape.size, !shape.size -> !shape.size
+  return %result : !shape.size
+}