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

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>