[mlir][MemRef] Add UB as a dependent dialect and use ub.poison for Mem2Reg (#168066)

This patch adds `ub` as a dependent dialect to `memref`, and uses
`ub.poison` as the default value in `AllocaOp::getDefaultValue` for the
mem2reg pass.

This aligns the behavior of `mem2reg` with LLVM, where loading a value
before having a value should be poison.

---------

Signed-off-by: Fabian Mora <[email protected]>

Author: fabianmcg

mlir/include/mlir/Dialect/MemRef/IR/MemRefBase.td

@@ -19,7 +19,14 @@ def MemRef_Dialect : Dialect {
     manipulation ops, which are not strongly associated with any particular
     other dialect or domain abstraction.
   }];
-  let dependentDialects = ["arith::ArithDialect"];
+  let dependentDialects = [
+    // `arith` is a dependency because it is used to materialize constants,
+    // and in some canonicalization patterns.
+    "arith::ArithDialect",
+    // `ub` is a dependency because `AllocaOp::getDefaultValue` can produce a
+    // `ub.poison` value.
+    "ub::UBDialect"
+  ];
   let hasConstantMaterializer = 1;
 }
 

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

@@ -25,6 +25,7 @@ add_mlir_dialect_library(MLIRMemRefDialect
   MLIRMemorySlotInterfaces
   MLIRShapedOpInterfaces
   MLIRSideEffectInterfaces
+  MLIRUBDialect
   MLIRValueBoundsOpInterface
   MLIRViewLikeInterface
 )

mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp

@@ -10,6 +10,7 @@
 #include "mlir/Conversion/ConvertToLLVM/ToLLVMInterface.h"
 #include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
 #include "mlir/Interfaces/RuntimeVerifiableOpInterface.h"

mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp

@@ -13,6 +13,7 @@
 
 #include "mlir/Dialect/MemRef/IR/MemRefMemorySlot.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
@@ -61,15 +62,8 @@ static void walkIndicesAsAttr(MLIRContext *ctx, ArrayRef<int64_t> shape,
 //  Interfaces for AllocaOp
 //===----------------------------------------------------------------------===//
 
-static bool isSupportedElementType(Type type) {
-  return llvm::isa<MemRefType>(type) ||
-         OpBuilder(type.getContext()).getZeroAttr(type);
-}
-
 SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
   MemRefType type = getType();
-  if (!isSupportedElementType(type.getElementType()))
-    return {};
   if (!type.hasStaticShape())
     return {};
   // Make sure the memref contains only a single element.
@@ -81,16 +75,7 @@ SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
 
 Value memref::AllocaOp::getDefaultValue(const MemorySlot &slot,
                                         OpBuilder &builder) {
-  assert(isSupportedElementType(slot.elemType));
-  // TODO: support more types.
-  return TypeSwitch<Type, Value>(slot.elemType)
-      .Case([&](MemRefType t) {
-        return memref::AllocaOp::create(builder, getLoc(), t);
-      })
-      .Default([&](Type t) {
-        return arith::ConstantOp::create(builder, getLoc(), t,
-                                         builder.getZeroAttr(t));
-      });
+  return ub::PoisonOp::create(builder, getLoc(), slot.elemType);
 }
 
 std::optional<PromotableAllocationOpInterface>

mlir/test/Dialect/MemRef/mem2reg.mlir

@@ -18,7 +18,7 @@ func.func @basic() -> i32 {
 // CHECK-LABEL: func.func @basic_default
 func.func @basic_default() -> i32 {
   // CHECK-NOT: = memref.alloca
-  // CHECK: %[[RES:.*]] = arith.constant 0 : i32
+  // CHECK: %[[RES:.*]] = ub.poison : i32
   // CHECK-NOT: = memref.alloca
   %0 = arith.constant 5 : i32
   %1 = memref.alloca() : memref<i32>