[mlir] Add `ub` dialect and `poison` op.

[Hardcode84]

Summary

Add new dialect boilerplate and poison op definition.

Discussion: https://discourse.llvm.org/t/rfc-poison-semantics-for-mlir/66245/24

[tschuett]

mlir/include/mlir/Dialect/Ub/IR/UbOps.td

19 ↗

(On Diff #536349)

Func dialect?

[Hardcode84]

fix comment

[Hardcode84]

mlir/include/mlir/Dialect/Ub/IR/UbOps.td

19 ↗

(On Diff #536349)

Copypaste )

[kuhar]

What happened with the idea of having a poison attribute? I think this would be required for poison constants.

mlir/include/mlir/Dialect/CMakeLists.txt

2–4

Could you land this sorting change as a separate NFC commit? This seems uncontroversial and provides clear value IMO, but would be worth separating from this PR.

36

Can we follow the existing capitalization rules and call this UB?

mlir/include/mlir/Dialect/Ub/IR/UbOps.td

27 ↗

(On Diff #536353)

Do we want this op to be constant-like? I guess that's not possible without a new attribute.

35 ↗

(On Diff #536353)

What are the semantics of 'composite' poison types? E.g., vector<4xi32>, struct<...>, tensor<?x...>.

Are they TBD, defined by their dialect, or something else?

[Hardcode84]

mlir/include/mlir/Dialect/Ub/IR/UbOps.td

27 ↗

(On Diff #536353)

With constant-like, rewriters will merge multiple poison into one, am I right? We should probably do this.

35 ↗

(On Diff #536353)

I think, it should be up to the specific dialect/lowering how to interpret poison for its types.

[Hardcode84]

mlir/include/mlir/Dialect/CMakeLists.txt

2–4

https://reviews.llvm.org/D154267

[kuhar]

mlir/include/mlir/Dialect/Ub/IR/UbOps.td

27 ↗

(On Diff #536353)

I think declaring it as 'Pure' should be sufficient for that and handled by CSE.

But with this definition, it's impossible to create composite poison constants. For example, for vector<2xi32> whose only element 1 is poison, we'd have to first create a constant and then insert a poison. If we do this another way round and first poison the whole vector and, as-is, we do not know what happens if we insert a non-poison element. I think this story needs to be fleshed out before landing anything.

I can see temporary solution without poison attributes, but we should at least agree on the semantics.

35 ↗

(On Diff #536353)

I think that for builtin types, the interpretation is universal. For scalars poison is an additional possible value. For composite types this is less obvious and should be carefully explained and agreed on.

[Hardcode84]

mlir/include/mlir/Dialect/Ub/IR/UbOps.td

35 ↗

(On Diff #536353)

I think, current approach in llvm is to create whole poison vector and unpoison individual elements via insertElement, but I can't find any proper documentation on this.

I suggest to not introduce any special handling for aggregate/composite types in ub dialect (And I also think 'composite' term itself in not very well defined in MLIR, are memrefs composite or not).

If some dialect need some special handling for poison, it should introduce its own ops/attributes and refuse to lower ub.poison for its types.

[kuhar]

mlir/include/mlir/Dialect/Ub/IR/UbOps.td

35 ↗

(On Diff #536353)

I think, current approach in llvm is to create whole poison vector and unpoison individual elements via insertElement, but I can't find any proper documentation on this.

LLVM allows you to create vector constants with only some elements being poison, some example:
https://github.com/llvm/llvm-project/blob/15ef9b26adeb8c9dd98228fc26757966d8355986/llvm/test/Transforms/InstCombine/insert-extract-shuffle-inseltpoison.ll#L47

<2 x i32> <i32 0, i32 poison>

If some dialect need some special handling for poison, it should introduce its own ops/attributes and refuse to lower ub.poison for its types.

I'm not sure we agree on this. I think the end goal for poison would be to be an equally valid value as other ones.

Why wouldn't this be a property of types? Ops can decide which values to accept and produce, but I think of types being orthogonal to that. For example you could have an my.even_constant n : i32 op that only returns even constants, and fails to verify for odd numbers, but this does mean that 5 should now be considered a 'special' integer value.

(And I also think 'composite' term itself in not very well defined in MLIR, are memrefs composite or not).

I understand wanting to implement this by incrementally increasing the scope, but don't understand what the proposed strategy is. Could you briefly describe why we shouldn't do something like Karl's prototype with an attribute for poison values? It seems odd to me to have an op that is conceptually a constant without being an actual constant.

[Hardcode84]

mlir/include/mlir/Dialect/Ub/IR/UbOps.td

35 ↗

(On Diff #536353)

Could you briefly describe why we shouldn't do something like Karl's prototype with an attribute for poison values? It seems odd to me to have an op that is conceptually a constant without being an actual constant.

For my usecase (control flow transformations) I need to create poisons of arbitrary type. If we decide to use arith.constant with poison attribute, I will have to hardcode arith dialect (and probably all dialects with custom types) dependency into my pass.

We can approach it with type interfaces (createPoison on type), but it will introduce dependency from builtin types to arith dialect.

[Hardcode84]

mlir/include/mlir/Dialect/Ub/IR/UbOps.td

35 ↗

(On Diff #536353)

After some more thought:
I think I can add poison op and poison attr to ub dialect, but without any composite semantics (always treat as fully poisoned), dialects can potentially introduce derived attributes to handle composites (and will have to do dyn_cast to check if it's base (fully-poisoned) or derived (partially))

[Hardcode84]

rebase, review comments, add poison attr

[kuhar]

Thanks for pushing this forward, I think it's the right first step

mlir/include/mlir/Dialect/UB/IR/UBOps.td

57

Is this the full syntax? Where is the type?

mlir/test/Dialect/UB/ops.mlir

10

Since the poison op can return any type, should we add a couple of more types than just i32? I'm thinking vector/tensor/complex

[Hardcode84]

@kuhar so what's your opinion on typed vs typeless PoisonAttr (from discussion thread)

[Hardcode84]

rebase, fix asm format, more tests

[Hardcode84]

mlir/include/mlir/Dialect/UB/IR/UBOps.td

57

FIxed.
Also, slightly changed asm format, because I was getting strange parsing errors with comma as separator

[kuhar]

In D154248#4494992, @Hardcode84 wrote:

@kuhar so what's your opinion on typed vs typeless PoisonAttr (from discussion thread)

I think we will need more poison attributes in the future, but for all or nothing poison a single untyped attribute should work fine. And if we do need a typed full poison, this doesn't seem too hard to add to me. I'd probably start with a typed one from the beginning, following Karl's prototype, but I'm not the one doing the work so I don't get to complain about details like this.

[kuhar]

Please try to get at least one more approval before submitting. I'd like to make sure that we have enough consensus in the community to move forward.

[Hardcode84]

add PoisonAttrInterface

[Hardcode84]

fix namespace

[Hardcode84]

move interface to separate file

[Hardcode84]

fix license header

[Hardcode84]

fix more license headers

mlir/include/mlir/Dialect/CMakeLists.txt

 add_subdirectory(Affine)
 add_subdirectory(AMDGPU)
 add_subdirectory(AMX)
 add_subdirectory(Arith)
 add_subdirectory(ArmNeon)
 add_subdirectory(ArmSME)
 add_subdirectory(ArmSVE)
 add_subdirectory(Async)
 add_subdirectory(Bufferization)
 add_subdirectory(Complex)
 add_subdirectory(ControlFlow)
 add_subdirectory(DLTI)

 add_subdirectory(Quant)
 add_subdirectory(SCF)
 add_subdirectory(Shape)
 add_subdirectory(SparseTensor)
 add_subdirectory(SPIRV)
 add_subdirectory(Tensor)
 add_subdirectory(Tosa)
 add_subdirectory(Transform)
+add_subdirectory(UB)
 add_subdirectory(Utils)
 add_subdirectory(Vector)
 add_subdirectory(X86Vector)

mlir/include/mlir/Dialect/UB/CMakeLists.txt

+add_subdirectory(IR)

mlir/include/mlir/Dialect/UB/IR/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS UBOps.td)
+mlir_tablegen(UBOps.h.inc -gen-op-decls)
+mlir_tablegen(UBOps.cpp.inc -gen-op-defs)
+mlir_tablegen(UBOpsDialect.h.inc -gen-dialect-decls)
+mlir_tablegen(UBOpsDialect.cpp.inc -gen-dialect-defs)
+mlir_tablegen(UBOpsAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=ub)
+mlir_tablegen(UBOpsAttributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=ub)
+add_public_tablegen_target(MLIRUBOpsIncGen)
+
+add_mlir_doc(UBOps UBOps Dialects/ -gen-op-doc)

mlir/include/mlir/Dialect/UB/IR/UBOps.h

+//===- FuncOps.h - UB Dialect Operations ----------------------*--- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_UB_IR_OPS_H
+#define MLIR_DIALECT_UB_IR_OPS_H
+
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/UB/IR/UBOpsAttributes.h.inc"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/UB/IR/UBOps.h.inc"
+
+#include "mlir/Dialect/UB/IR/UBOpsDialect.h.inc"
+
+#endif // MLIR_DIALECT_UB_IR_OPS_H

mlir/include/mlir/Dialect/UB/IR/UBOps.td

+//===- FuncOps.td - UB operation definitions ---------------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_UB_IR_UBOPS_TD
+#define MLIR_DIALECT_UB_IR_UBOPS_TD
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/AttrTypeBase.td"
+
+def UB_Dialect : Dialect {
+  let name = "ub";
+  let cppNamespace = "::mlir::ub";
+
+  let hasConstantMaterializer = 1;
+  let useDefaultAttributePrinterParser = 1;
+}
+
+// Base class for UB dialect attributes.
+class UB_Attr<string name, string attrMnemonic, list<Trait> traits = []> :
+    AttrDef<UB_Dialect, name, traits> {
+  let mnemonic = attrMnemonic;
+}
+
+// Base class for UB dialect ops.
+class UB_Op<string mnemonic, list<Trait> traits = []> :
+    Op<UB_Dialect, mnemonic, traits>;
+
+
+def PoisonAttr : UB_Attr<"Poison", "poison"> {
+
+  // To make DefaultValuedAttr happy.
+  let constBuilderCall = cppNamespace # "::" # cppClassName #
+                         "::get($_builder.getContext())";
+}
+
+//===----------------------------------------------------------------------===//
+// PoisonOp
+//===----------------------------------------------------------------------===//
+
+def PoisonOp : UB_Op<"poison", [ConstantLike, Pure]> {
+  let summary = "Poisoned constant operation.";
+  let description = [{
+    The `poison` operation materializes a compile-time poisoned constant value
+    to indicate deferred undefined behavior.
+    `value` attirbute is needed to indicate an optional additional poison
+    semantics (e.g. partially poisoned vectors), default value indicates results
+    is fully poisoned.
+
+    Syntax:
+
+    ```
+    poison-op ::= `poison` $value
+    ```
+
+    Examples:
+
+    ```
+    // Short form
+    %0 = ub.poison : i32
+    // Long form
+    %1 = ub.poison #<custom_poison_elements_attr> : vector<rxi64>
+    ```
+  }];
+
+  let arguments = (ins DefaultValuedAttr<PoisonAttr, "{}">:$value);
+  let results = (outs AnyType:$result);
+
+  let assemblyFormat = "attr-dict (`,` $value^)? `:` type($result)";
+
+  let hasFolder = 1;
+}
+
+#endif // MLIR_DIALECT_UB_IR_UBOPS_TD

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
 #include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/PDLExtension/PDLExtension.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
 #include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/X86Vector/X86VectorDialect.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 
 namespace mlir {
registry.insert<acc::OpenACCDialect,
                   ROCDL::ROCDLDialect,
                   scf::SCFDialect,
                   shape::ShapeDialect,
                   sparse_tensor::SparseTensorDialect,
                   spirv::SPIRVDialect,
                   tensor::TensorDialect,
                   tosa::TosaDialect,
                   transform::TransformDialect,
+                  ub::UBDialect,
                   vector::VectorDialect,
                   x86vector::X86VectorDialect>();
   // clang-format on
 
   // Register all dialect extensions.
   affine::registerTransformDialectExtension(registry);
   bufferization::registerTransformDialectExtension(registry);
   gpu::registerTransformDialectExtension(registry);

mlir/lib/Dialect/CMakeLists.txt

 add_subdirectory(Quant)
 add_subdirectory(SCF)
 add_subdirectory(Shape)
 add_subdirectory(SparseTensor)
 add_subdirectory(SPIRV)
 add_subdirectory(Tensor)
 add_subdirectory(Tosa)
 add_subdirectory(Transform)
+add_subdirectory(UB)
 add_subdirectory(Utils)
 add_subdirectory(Vector)
 add_subdirectory(X86Vector)
 
 set(LLVM_OPTIONAL_SOURCES
   Traits.cpp
 )
 

mlir/lib/Dialect/UB/CMakeLists.txt

+add_subdirectory(IR)

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

+add_mlir_dialect_library(MLIRUBDialect
+  UBOps.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/UB
+
+  DEPENDS
+  MLIRUBOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  )

mlir/lib/Dialect/UB/IR/UBOps.cpp

+//===- FuncOps.cpp - UB Dialect 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/UB/IR/UBOps.h"
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+#include "mlir/Dialect/UB/IR/UBOpsDialect.cpp.inc"
+
+using namespace mlir;
+using namespace mlir::ub;
+
+//===----------------------------------------------------------------------===//
+// UBDialect
+//===----------------------------------------------------------------------===//
+
+void UBDialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "mlir/Dialect/UB/IR/UBOps.cpp.inc"
+      >();
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "mlir/Dialect/UB/IR/UBOpsAttributes.cpp.inc"
+      >();
+}
+
+Operation *UBDialect::materializeConstant(OpBuilder &builder, Attribute value,
+                                          Type type, Location loc) {
+  if (auto attr = dyn_cast<PoisonAttr>(value))
+    return builder.create<PoisonOp>(loc, type, attr);
+
+  return nullptr;
+}
+
+OpFoldResult PoisonOp::fold(FoldAdaptor /*adaptor*/) { return getValue(); }
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/UB/IR/UBOpsAttributes.cpp.inc"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/UB/IR/UBOps.cpp.inc"

mlir/test/Dialect/UB/canonicalize.mlir

+// RUN: mlir-opt %s -canonicalize="test-convergence" --split-input-file | FileCheck %s
+
+
+// CHECK-LABEL: func @merge_poison()
+//       CHECK:   %[[RES:.*]] = ub.poison : i32
+//       CHECK:   return %[[RES]], %[[RES]]
+func.func @merge_poison() -> (i32, i32) {
+  %0 = ub.poison : i32
+  %1 = ub.poison : i32
+  return %0, %1 : i32, i32
+}

mlir/test/Dialect/UB/ops.mlir

+// RUN: mlir-opt %s | FileCheck %s
+// Verify the printed output can be parsed.
+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+// Verify the generic form can be parsed.
+// RUN: mlir-opt -mlir-print-op-generic %s | mlir-opt | FileCheck %s
+
+// CHECK-LABEL: func @poison()
+//       CHECK:   %{{.*}} = ub.poison : i32
+func.func @poison() -> i32 {
+  %0 = ub.poison : i32
+  return %0 : i32
+}