[mlir] Add ub dialect and poison op.

Add new dialect boilerplate and `poison` op definition.

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

Differential Revision: https://reviews.llvm.org/D154248

mlir/include/mlir/Dialect/CMakeLists.txt

@@ -33,6 +33,7 @@ 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

@@ -0,0 +1 @@
+add_subdirectory(IR)

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

@@ -0,0 +1,15 @@
+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-dialect-doc)
+
+set(LLVM_TARGET_DEFINITIONS UBOpsInterfaces.td)
+mlir_tablegen(UBOpsInterfaces.h.inc -gen-attr-interface-decls)
+mlir_tablegen(UBOpsInterfaces.cpp.inc -gen-attr-interface-defs)
+add_public_tablegen_target(MLIRUBOpsInterfacesIncGen)

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

@@ -0,0 +1,26 @@
+//===- UBOps.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"
+
+#include "mlir/Dialect/UB/IR/UBOpsInterfaces.h.inc"
+
+#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

@@ -0,0 +1,75 @@
+//===- UBOps.td - UB operations 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"
+
+include "UBOpsInterfaces.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", [PoisonAttrInterface]> {
+}
+
+//===----------------------------------------------------------------------===//
+// 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 `>`)? `:` type
+    ```
+
+    Examples:
+
+    ```
+    // Short form
+    %0 = ub.poison : i32
+    // Long form
+    %1 = ub.poison <#custom_poison_elements_attr> : vector<4xi64>
+    ```
+  }];
+
+  let arguments = (ins DefaultValuedAttr<PoisonAttrInterface, "{}">:$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/Dialect/UB/IR/UBOpsInterfaces.td

@@ -0,0 +1,24 @@
+//===- UBOpsInterfaces.td - UB interfaces 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_UBOPSINTERFACES_TD
+#define MLIR_DIALECT_UB_IR_UBOPSINTERFACES_TD
+
+
+include "mlir/IR/OpBase.td"
+
+def PoisonAttrInterface : AttrInterface<"PoisonAttrInterface"> {
+  let cppNamespace = "::mlir::ub";
+  // No methods for now.
+
+  // To make DefaultValuedAttr happy.
+  let constBuilderCall = cppNamespace # "::" # "PoisonAttr" #
+                         "::get($_builder.getContext())";
+}
+
+#endif // MLIR_DIALECT_UB_IR_UBOPSINTERFACES_TD

mlir/include/mlir/InitAllDialects.h

@@ -79,6 +79,7 @@
 #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"
@@ -128,6 +129,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
                   tensor::TensorDialect,
                   tosa::TosaDialect,
                   transform::TransformDialect,
+                  ub::UBDialect,
                   vector::VectorDialect,
                   x86vector::X86VectorDialect>();
   // clang-format on

mlir/lib/Dialect/CMakeLists.txt

@@ -33,6 +33,7 @@ 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/lib/Dialect/UB/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(IR)

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

@@ -0,0 +1,13 @@
+add_mlir_dialect_library(MLIRUBDialect
+  UBOps.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/UB
+
+  DEPENDS
+  MLIRUBOpsIncGen
+  MLIRUBOpsInterfacesIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  )

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

@@ -0,0 +1,51 @@
+//===- UBOps.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(); }
+
+#include "mlir/Dialect/UB/IR/UBOpsInterfaces.cpp.inc"
+
+#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

@@ -0,0 +1,11 @@
+// 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

@@ -0,0 +1,40 @@
+// 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
+}
+
+// CHECK-LABEL: func @poison_full_form()
+//       CHECK:   %{{.*}} = ub.poison : i32
+func.func @poison_full_form() -> i32 {
+  %0 = ub.poison <#ub.poison> : i32
+  return %0 : i32
+}
+
+// CHECK-LABEL: func @poison_complex()
+//       CHECK:   %{{.*}} = ub.poison : complex<f32>
+func.func @poison_complex() -> complex<f32> {
+  %0 = ub.poison : complex<f32>
+  return %0 : complex<f32>
+}
+
+// CHECK-LABEL: func @poison_vec()
+//       CHECK:   %{{.*}} = ub.poison : vector<4xi64>
+func.func @poison_vec() -> vector<4xi64> {
+  %0 = ub.poison : vector<4xi64>
+  return %0 : vector<4xi64>
+}
+
+// CHECK-LABEL: func @poison_tensor()
+//       CHECK:   %{{.*}} = ub.poison : tensor<8x?xf64>
+func.func @poison_tensor() -> tensor<8x?xf64> {
+  %0 = ub.poison : tensor<8x?xf64>
+  return %0 : tensor<8x?xf64>
+}