[Moore] Fix mem2reg implementation (#7498)

Fix a few issues in the mem2reg interface implementations of VariableOp,
ReadOp, and BlockingAssignOp. Add tests reduced from the Snitch core
that used to fail before this fix.

lib/Dialect/Moore/MooreOps.cpp

@@ -259,37 +259,19 @@ void VariableOp::getAsmResultNames(OpAsmSetValueNameFn setNameFn) {
     setNameFn(getResult(), *getName());
 }
 
-llvm::SmallVector<MemorySlot> VariableOp::getPromotableSlots() {
-  if (isa<SVModuleOp>(getOperation()->getParentOp()))
-    return {};
-  if (getInitial())
-    return {};
-  return {MemorySlot{getResult(), getType()}};
-}
-
-Value VariableOp::getDefaultValue(const MemorySlot &slot, OpBuilder &builder) {
-  if (auto value = getInitial())
-    return value;
-  return builder.create<ConstantOp>(
-      getLoc(),
-      cast<moore::IntType>(cast<RefType>(slot.elemType).getNestedType()), 0);
-}
-
-void VariableOp::handleBlockArgument(const MemorySlot &slot,
-                                     BlockArgument argument,
-                                     OpBuilder &builder) {}
-
-::std::optional<::mlir::PromotableAllocationOpInterface>
-VariableOp::handlePromotionComplete(const MemorySlot &slot, Value defaultValue,
-                                    OpBuilder &builder) {
-  if (defaultValue.use_empty())
-    defaultValue.getDefiningOp()->erase();
-  this->erase();
-  return std::nullopt;
-}
-
 LogicalResult VariableOp::canonicalize(VariableOp op,
                                        PatternRewriter &rewriter) {
+  // If the variable is embedded in an SSACFG region, move the initial value
+  // into an assignment immediately after the variable op. This allows the
+  // mem2reg pass which cannot handle variables with initial values.
+  auto initial = op.getInitial();
+  if (initial && mlir::mayHaveSSADominance(*op->getParentRegion())) {
+    rewriter.modifyOpInPlace(op, [&] { op.getInitialMutable().clear(); });
+    rewriter.setInsertionPointAfter(op);
+    rewriter.create<BlockingAssignOp>(initial.getLoc(), op, initial);
+    return success();
+  }
+
   // Check if the variable has one unique continuous assignment to it, all other
   // uses are reads, and that all uses are in the same block as the variable
   // itself.
@@ -334,6 +316,46 @@ LogicalResult VariableOp::canonicalize(VariableOp op,
   return success();
 }
 
+SmallVector<MemorySlot> VariableOp::getPromotableSlots() {
+  // We cannot promote variables with an initial value, since that value may not
+  // dominate the location where the default value needs to be constructed.
+  if (mlir::mayBeGraphRegion(*getOperation()->getParentRegion()) ||
+      getInitial())
+    return {};
+  return {MemorySlot{getResult(), getType().getNestedType()}};
+}
+
+Value VariableOp::getDefaultValue(const MemorySlot &slot, OpBuilder &builder) {
+  auto packedType = dyn_cast<PackedType>(slot.elemType);
+  if (!packedType)
+    return {};
+  auto bitWidth = packedType.getBitSize();
+  if (!bitWidth)
+    return {};
+  auto fvint = packedType.getDomain() == Domain::FourValued
+                   ? FVInt::getAllX(*bitWidth)
+                   : FVInt::getZero(*bitWidth);
+  Value value = builder.create<ConstantOp>(
+      getLoc(), IntType::get(getContext(), *bitWidth, packedType.getDomain()),
+      fvint);
+  if (value.getType() != packedType)
+    builder.create<ConversionOp>(getLoc(), packedType, value);
+  return value;
+}
+
+void VariableOp::handleBlockArgument(const MemorySlot &slot,
+                                     BlockArgument argument,
+                                     OpBuilder &builder) {}
+
+std::optional<mlir::PromotableAllocationOpInterface>
+VariableOp::handlePromotionComplete(const MemorySlot &slot, Value defaultValue,
+                                    OpBuilder &builder) {
+  if (defaultValue && defaultValue.use_empty())
+    defaultValue.getDefiningOp()->erase();
+  this->erase();
+  return {};
+}
+
 SmallVector<DestructurableMemorySlot> VariableOp::getDestructurableSlots() {
   if (isa<SVModuleOp>(getOperation()->getParentOp()))
     return {};
@@ -1054,8 +1076,7 @@ bool BlockingAssignOp::canUsesBeRemoved(
     return false;
   Value blockingUse = (*blockingUses.begin())->get();
   return blockingUse == slot.ptr && getDst() == slot.ptr &&
-         getSrc() != slot.ptr &&
-         getSrc().getType() == cast<RefType>(slot.elemType).getNestedType();
+         getSrc() != slot.ptr && getSrc().getType() == slot.elemType;
 }
 
 DeletionKind BlockingAssignOp::removeBlockingUses(
@@ -1070,7 +1091,7 @@ DeletionKind BlockingAssignOp::removeBlockingUses(
 //===----------------------------------------------------------------------===//
 
 bool ReadOp::loadsFrom(const MemorySlot &slot) {
-  return getOperand() == slot.ptr;
+  return getInput() == slot.ptr;
 }
 
 bool ReadOp::storesTo(const MemorySlot &slot) { return false; }
@@ -1089,7 +1110,7 @@ bool ReadOp::canUsesBeRemoved(const MemorySlot &slot,
     return false;
   Value blockingUse = (*blockingUses.begin())->get();
   return blockingUse == slot.ptr && getOperand() == slot.ptr &&
-         getResult().getType() == cast<RefType>(slot.elemType).getNestedType();
+         getResult().getType() == slot.elemType;
 }
 
 DeletionKind

lib/Dialect/Moore/Transforms/SimplifyProcedures.cpp

@@ -43,8 +43,8 @@ void SimplifyProceduresPass::runOnOperation() {
 
     // Use to collect blocking assignments that have been replaced by a "shadow"
     // variable.
-    DenseSet<Operation *> assignOps;
     procedureOp.walk([&](Operation *op) {
+      SmallVector<std::tuple<BlockingAssignOp, Value, Value>> assignOps;
       auto &nestedOp = *op;
       // Only create a "shadow" varaible for the global variable used by other
       // operations in the procedure body.
@@ -67,15 +67,16 @@ void SimplifyProceduresPass::runOnOperation() {
               nestedOp.getLoc(), cast<RefType>(resultType).getNestedType(),
               varOp.getResult());
           auto newVarOp = builder.create<VariableOp>(
-              nestedOp.getLoc(), resultType, StringAttr{}, readOp);
+              nestedOp.getLoc(), resultType, StringAttr{}, Value{});
+          builder.create<BlockingAssignOp>(nestedOp.getLoc(), newVarOp, readOp);
           builder.clearInsertionPoint();
 
           // Replace the users of the global variable with a corresponding
           // "shadow" variable.
           for (auto *user : users) {
             user->replaceUsesOfWith(user->getOperand(0), newVarOp);
-            if (isa<BlockingAssignOp>(user))
-              assignOps.insert(user);
+            if (auto assignOp = dyn_cast<BlockingAssignOp>(user))
+              assignOps.push_back({assignOp, newVarOp, varOp});
           }
         }
       }
@@ -83,20 +84,12 @@ void SimplifyProceduresPass::runOnOperation() {
       // Ensure the global variable has the correct value. So needing to create
       // a blocking assign for the global variable when the "shadow" variable
       // has a new value.
-      for (auto *assignOp : assignOps)
-        if (auto localVarOp = llvm::dyn_cast_or_null<VariableOp>(
-                assignOp->getOperand(0).getDefiningOp())) {
-          auto resultType = localVarOp.getResult().getType();
-          builder.setInsertionPointAfter(assignOp);
-          auto readOp = builder.create<ReadOp>(
-              localVarOp.getLoc(), cast<RefType>(resultType).getNestedType(),
-              localVarOp.getResult());
-          builder.create<BlockingAssignOp>(
-              nestedOp.getLoc(),
-              localVarOp.getInitial().getDefiningOp()->getOperand(0), readOp);
-          builder.clearInsertionPoint();
-          assignOps.erase(assignOp);
-        }
+      for (auto [assignOp, localVar, var] : assignOps) {
+        builder.setInsertionPointAfter(assignOp);
+        auto readOp = builder.create<ReadOp>(assignOp.getLoc(), localVar);
+        builder.create<BlockingAssignOp>(assignOp.getLoc(), var, readOp);
+        builder.clearInsertionPoint();
+      }
     });
   });
 }

test/Dialect/Moore/canonicalizers.mlir

@@ -280,3 +280,14 @@ func.func @Pow(%arg0 : !moore.l32) -> (!moore.l32, !moore.l32, !moore.l32, !moor
   // CHECK-NEXT: return [[V0]], [[V0]], [[V0]], [[V0]], [[V4]], [[V5]] :
   return %3, %4, %5, %6, %7, %8 : !moore.l32, !moore.l32, !moore.l32, !moore.l32, !moore.l32, !moore.l32
 }
+
+// CHECK-LABEL: func.func @MoveInitialOutOfSSAVariable
+func.func @MoveInitialOutOfSSAVariable() {
+  // CHECK: [[TMP:%.+]] = moore.constant 9001
+  %0 = moore.constant 9001 : i42
+  // CHECK: [[VAR:%.+]] = moore.variable : <i42>
+  // CHECK-NEXT: moore.blocking_assign [[VAR]], [[TMP]]
+  %1 = moore.variable %0 : <i42>
+  func.call @useRef(%1) : (!moore.ref<i42>) -> ()
+  return
+}

test/Dialect/Moore/mem2reg.mlir

@@ -1,34 +1,48 @@
-// RUN: circt-opt --mem2reg %s | FileCheck %s
+// RUN: circt-opt --mem2reg --verify-diagnostics %s | FileCheck %s
 
-// CHECK-LABEL: moore.module @LocalVar() {
-moore.module @LocalVar() {
-  // CHECK: %x = moore.variable : <i32>
-  // CHECK: %y = moore.variable : <i32>
-  // CHECK: %z = moore.variable : <i32>
-  %x = moore.variable : <i32>
-  %y = moore.variable : <i32>
-  %z = moore.variable : <i32>
-  moore.procedure always_comb {
-    // CHECK: %0 = moore.read %x
-    // CHECK: %1 = moore.constant 1 : i32
-    // CHECK: %2 = moore.add %0, %1 : i32
-    // CHECK: moore.blocking_assign %z, %2 : i32
-    // CHECK: %3 = moore.constant 1 : i32
-    // CHECK: %4 = moore.add %2, %3 : i32
-    // CHECK: moore.blocking_assign %y, %4 : i32
-    %a = moore.variable : <i32>
-    %0 = moore.read %x : <i32>
-    %1 = moore.constant 1 : i32
-    %2 = moore.add %0, %1 : i32
-    moore.blocking_assign %a, %2 : i32
-    %3 = moore.read %a : <i32>
-    moore.blocking_assign %z, %3 : i32
-    %4 = moore.read %a : <i32>
-    %5 = moore.constant 1 : i32
-    %6 = moore.add %4, %5 : i32
-    moore.blocking_assign %a, %6 : i32
-    %7 = moore.read %a : <i32>
-    moore.blocking_assign %y, %7 : i32
-    moore.return
-  }
+// CHECK-LABEL: func.func @Basic(
+func.func @Basic() -> !moore.i42 {
+  // CHECK: [[TMP:%.*]] = moore.constant 9001 : i42
+  // CHECK-NOT: = moore.variable
+  // CHECK-NOT: = moore.blocking_assign
+  // CHECK-NOT: = moore.read
+  %0 = moore.constant 9001 : i42
+  %1 = moore.variable : <i42>
+  moore.blocking_assign %1, %0 : i42
+  %2 = moore.read %1 : <i42>
+  // CHECK: return [[TMP]] : !moore.i42
+  return %2 : !moore.i42
+}
+
+// CHECK-LABEL: func.func @ControlFlow(
+func.func @ControlFlow(%arg0: i1, %arg1: !moore.l8) -> !moore.l8 {
+  // CHECK-NOT: moore.variable
+  // CHECK: [[DEFAULT:%.+]] = moore.constant hXX : l8
+  %0 = moore.variable : <l8>
+  // CHECK: cf.cond_br %arg0, ^[[BB1:.+]], ^[[BB2:.+]]([[DEFAULT]] : !moore.l8)
+  cf.cond_br %arg0, ^bb1, ^bb2
+^bb1:
+  // CHECK-NOT: moore.blocking_assign
+  // CHECK: cf.br ^[[BB2]](%arg1 : !moore.l8)
+  moore.blocking_assign %0, %arg1 : l8
+  cf.br ^bb2
+^bb2:
+  // CHECK: ^[[BB2]]([[TMP:%.+]]: !moore.l8):
+  // CHECK-NOT: moore.read
+  // CHECK: return [[TMP]]
+  %1 = moore.read %0 : <l8>
+  return %1 : !moore.l8
+}
+
+// CHECK-LABEL: func.func @InitialValueDoesNotDominateDefault(
+func.func @InitialValueDoesNotDominateDefault() {
+  cf.br ^bb1
+^bb1:
+  %0 = moore.constant 0 : i32
+  %1 = moore.variable %0 : <i32>
+  cf.br ^bb2
+^bb2:
+  %2 = moore.read %1 : <i32>
+  moore.blocking_assign %1, %2 : i32
+  cf.br ^bb1
 }

test/Dialect/Moore/simplify-procedures.mlir

@@ -9,7 +9,8 @@ moore.module @Foo() {
   // CHECK: moore.procedure always_comb
   moore.procedure always_comb {
     // CHECK: [[TMP:%.+]] = moore.read %a
-    // CHECK: [[LOCAL_A:%.+]] = moore.variable [[TMP]]
+    // CHECK: [[LOCAL_A:%.+]] = moore.variable
+    // CHECK: moore.blocking_assign [[LOCAL_A]], [[TMP]]
 
     // CHECK: [[C1:%.+]] = moore.constant 1
     // CHECK: moore.blocking_assign [[LOCAL_A]], [[C1]]
@@ -40,7 +41,8 @@ moore.module @Foo() {
   // CHECK: moore.procedure always_comb
   moore.procedure always_comb {
     // CHECK: [[TMP:%.+]] = moore.read %a
-    // CHECK: [[LOCAL_A:%.+]] = moore.variable %0
+    // CHECK: [[LOCAL_A:%.+]] = moore.variable
+    // CHECK: moore.blocking_assign [[LOCAL_A]], [[TMP]]
 
     // CHECK: [[TMP:%.+]] = moore.read [[LOCAL_A]]
     // CHECK: moore.blocking_assign %y, [[TMP]]

tools/circt-reduce/circt-reduce.cpp

@@ -19,6 +19,7 @@
 #include "circt/Reduce/GenericReductions.h"
 #include "circt/Reduce/Tester.h"
 #include "circt/Support/Version.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
@@ -424,7 +425,8 @@ int main(int argc, char **argv) {
   // Register all the dialects and create a context to work wtih.
   mlir::DialectRegistry registry;
   registerAllDialects(registry);
-  registry.insert<func::FuncDialect, scf::SCFDialect, LLVM::LLVMDialect>();
+  registry.insert<func::FuncDialect, scf::SCFDialect, cf::ControlFlowDialect,
+                  LLVM::LLVMDialect>();
   arc::registerReducePatternDialectInterface(registry);
   firrtl::registerReducePatternDialectInterface(registry);
   hw::registerReducePatternDialectInterface(registry);