[mlir][Transforms] Add pass to perform sparse conditional constant propagation

This revision adds the initial pass for performing SCCP generically in MLIR. SCCP is an algorithm for propagating constants across control flow, and optimistically assumes all values to be constant unless proven otherwise. It currently supports branching control, with support for regions and inter-procedural propagation being added in followups.

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

Author: River707

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

@@ -596,6 +596,10 @@ def BranchOp : Std_Op<"br",
 
     /// Erase the operand at 'index' from the operand list.
     void eraseOperand(unsigned index);
+
+    /// Returns the successor that would be chosen with the given constant
+    /// operands. Returns nullptr if a single successor could not be chosen.
+    Block *getSuccessorForOperands(ArrayRef<Attribute>);
   }];
 
   let hasCanonicalizer = 1;
@@ -1092,6 +1096,10 @@ def CondBranchOp : Std_Op<"cond_br",
       eraseSuccessorOperand(falseIndex, index);
     }
 
+    /// Returns the successor that would be chosen with the given constant
+    /// operands. Returns nullptr if a single successor could not be chosen.
+    Block *getSuccessorForOperands(ArrayRef<Attribute> operands);
+
   private:
     /// Get the index of the first true destination operand.
     unsigned getTrueDestOperandIndex() { return 1; }

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

@@ -68,6 +68,14 @@ def BranchOpInterface : OpInterface<"BranchOpInterface"> {
         }
         return llvm::None;
       }]
+    >,
+    InterfaceMethod<[{
+        Returns the successor that would be chosen with the given constant
+        operands. Returns nullptr if a single successor could not be chosen.
+      }],
+      "Block *", "getSuccessorForOperands",
+      (ins "ArrayRef<Attribute>":$operands), [{}],
+      /*defaultImplementation=*/[{ return nullptr; }]
     >
   ];
 

mlir/include/mlir/Transforms/FoldUtils.h

@@ -119,6 +119,11 @@ class OperationFolder {
   /// Clear out any constants cached inside of the folder.
   void clear();
 
+  /// Get or create a constant using the given builder. On success this returns
+  /// the constant operation, nullptr otherwise.
+  Value getOrCreateConstant(OpBuilder &builder, Dialect *dialect,
+                            Attribute value, Type type, Location loc);
+
 private:
   /// This map keeps track of uniqued constants by dialect, attribute, and type.
   /// A constant operation materializes an attribute with a type. Dialects may

mlir/include/mlir/Transforms/Passes.h

@@ -76,6 +76,10 @@ std::unique_ptr<OperationPass<ModuleOp>> createPrintOpStatsPass();
 /// the CallGraph.
 std::unique_ptr<Pass> createInlinerPass();
 
+/// Creates a pass which performs sparse conditional constant propagation over
+/// nested operations.
+std::unique_ptr<Pass> createSCCPPass();
+
 /// Creates a pass which delete symbol operations that are unreachable. This
 /// pass may *only* be scheduled on an operation that defines a SymbolTable.
 std::unique_ptr<Pass> createSymbolDCEPass();

mlir/include/mlir/Transforms/Passes.td

@@ -273,6 +273,20 @@ def PrintOp : Pass<"print-op-graph", "ModuleOp"> {
   let constructor = "mlir::createPrintOpGraphPass()";
 }
 
+def SCCP : Pass<"sccp"> {
+  let summary = "Sparse Conditional Constant Propagation";
+  let description = [{
+    This pass implements a general algorithm for sparse conditional constant
+    propagation. This algorithm detects values that are known to be constant and
+    optimistically propagates this throughout the IR. Any values proven to be
+    constant are replaced, and removed if possible.
+
+    This implementation is based on the algorithm described by Wegman and Zadeck
+    in [“Constant Propagation with Conditional Branches”](https://dl.acm.org/doi/10.1145/103135.103136) (1991).
+  }];
+  let constructor = "mlir::createSCCPPass()";
+}
+
 def StripDebugInfo : Pass<"strip-debuginfo"> {
   let summary = "Strip debug info from all operations";
   let description = [{

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

@@ -597,6 +597,8 @@ Optional<OperandRange> BranchOp::getSuccessorOperands(unsigned index) {
 
 bool BranchOp::canEraseSuccessorOperand() { return true; }
 
+Block *BranchOp::getSuccessorForOperands(ArrayRef<Attribute>) { return dest(); }
+
 //===----------------------------------------------------------------------===//
 // CallOp
 //===----------------------------------------------------------------------===//
@@ -863,6 +865,14 @@ Optional<OperandRange> CondBranchOp::getSuccessorOperands(unsigned index) {
 
 bool CondBranchOp::canEraseSuccessorOperand() { return true; }
 
+Block *CondBranchOp::getSuccessorForOperands(ArrayRef<Attribute> operands) {
+  if (BoolAttr condAttr = operands.front().dyn_cast_or_null<BoolAttr>())
+    return condAttr.getValue() ? trueDest() : falseDest();
+  if (IntegerAttr condAttr = operands.front().dyn_cast_or_null<IntegerAttr>())
+    return condAttr.getValue().isOneValue() ? trueDest() : falseDest();
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // Constant*Op
 //===----------------------------------------------------------------------===//

mlir/lib/Transforms/CMakeLists.txt

@@ -13,6 +13,7 @@ add_mlir_library(MLIRTransforms
   OpStats.cpp
   ParallelLoopCollapsing.cpp
   PipelineDataTransfer.cpp
+  SCCP.cpp
   StripDebugInfo.cpp
   SymbolDCE.cpp
   ViewOpGraph.cpp