[mlir] Add integer range analysis to SCCP

[krzysz00]

Summary

This integrates the scalar integer analysis recently added to MLIR
into the SCCP pass, allowing for more substantial constant
propagation (as values that can be logically proven constant will be
made constant) in a way that mirrors LLVM's SCCP pass, and eliminating
the need for a test pass.

A motivating example for the new logic is the comparison
operations (and their associated selects) emitted by
expandAffineMap() for the modulus operator, which can be substantially
in the presence of proof that the input to the modulus argument is
non-negative.

Depends on D124022

[Mogball]

delete the pass itself?

[krzysz00]

I don't think we can delete old SCCP - range analysis doesn't handle vectors, for example

[Mogball]

Oh I meant the test pass. I didn't see that the file was removed :P

[Mogball]

mlir/lib/Transforms/SCCP.cpp

288

rangeAnalysis.run(op)?

[Mogball]

mlir/lib/Transforms/SCCP.cpp

288

Ah, nevermind. I didn't see this was changed.

[rriddle]

From a quick first pass we should really avoid building multiple analysis if we can. Each analysis is going to do full passes of the IR. I wonder if we could work it so that the analysis are built simultaneously, so we don't have to do two full passes of the IR.

[krzysz00]

Since @Mogball is handling this in his dataflow analysis rework, abandoning.

mlir/include/mlir/Transforms/Passes.td

def SCCP : Pass<"sccp"> {
   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).
+
+    In addition, this pass uses the IntRangeAnalysis to determine which
+    scalar integer values are constants.
   }];
   let constructor = "mlir::createSCCPPass()";
 }
 
 def StripDebugInfo : Pass<"strip-debuginfo"> {
   let summary = "Strip debug info from all operations";
   let description = [{
     This pass strips the IR of any location information, by replacing all

mlir/lib/Transforms/SCCP.cpp

 // information throughout the IR, and replaces them. This is done with an
 // optimistic dataflow analysis that assumes that all values are constant until
 // proven otherwise.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/Analysis/DataFlowAnalysis.h"
+#include "mlir/Analysis/IntRangeAnalysis.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/Support/Debug.h"

 
 //===----------------------------------------------------------------------===//
 // SCCP Rewrites
 //===----------------------------------------------------------------------===//
 
 /// Replace the given value with a constant if the corresponding lattice
 /// represents a constant. Returns success if the value was replaced, failure
 /// otherwise.
+static LogicalResult replaceWithConstant(IntRangeAnalysis &analysis,
+                                         OpBuilder &b, OperationFolder &folder,
+                                         Value value) {
+  Optional<ConstantIntRanges> maybeInferredRange = analysis.getResult(value);
+  if (!maybeInferredRange)
+    return failure();
+  const ConstantIntRanges &inferredRange = maybeInferredRange.getValue();
+  Optional<APInt> maybeConstValue = inferredRange.getConstantValue();
+  if (!maybeConstValue.hasValue())
+    return failure();
+
+  Operation *maybeDefiningOp = value.getDefiningOp();
+  Dialect *valueDialect =
+      maybeDefiningOp ? maybeDefiningOp->getDialect()
+                      : value.getParentRegion()->getParentOp()->getDialect();
+  Attribute constAttr = b.getIntegerAttr(value.getType(), *maybeConstValue);
+  Value constant = folder.getOrCreateConstant(b, valueDialect, constAttr,
+                                              value.getType(), value.getLoc());
+  if (!constant)
+    return failure();
+
+  value.replaceAllUsesWith(constant);
+  return success();
+}
+
+/// Replace the given value with a constant if the corresponding lattice
+/// represents a constant. Returns success if the value was replaced, failure
+/// otherwise.
 static LogicalResult replaceWithConstant(SCCPAnalysis &analysis,
                                          OpBuilder &builder,
                                          OperationFolder &folder, Value value) {
   LatticeElement<SCCPLatticeValue> *lattice =
       analysis.lookupLatticeElement(value);
   if (!lattice)
     return failure();
   SCCPLatticeValue &latticeValue = lattice->getValue();
static LogicalResult replaceWithConstant(SCCPAnalysis &analysis,
 
   value.replaceAllUsesWith(constant);
   return success();
 }
 
 /// Rewrite the given regions using the computing analysis. This replaces the
 /// uses of all values that have been computed to be constant, and erases as
 /// many newly dead operations.
-static void rewrite(SCCPAnalysis &analysis, MLIRContext *context,
+static void rewrite(IntRangeAnalysis &rangeAnalysis, SCCPAnalysis &sccpAnalysis,
+                    MLIRContext *context,
                     MutableArrayRef<Region> initialRegions) {
   SmallVector<Block *> worklist;
   auto addToWorklist = [&](MutableArrayRef<Region> regions) {
     for (Region &region : regions)
       for (Block &block : llvm::reverse(region))
         worklist.push_back(&block);
   };
 
   // An operation folder used to create and unique constants.
   OperationFolder folder(context);
   OpBuilder builder(context);
 
   addToWorklist(initialRegions);
   while (!worklist.empty()) {
     Block *block = worklist.pop_back_val();
 
     for (Operation &op : llvm::make_early_inc_range(*block)) {
       builder.setInsertionPoint(&op);
 
       // Replace any result with constants.
       bool replacedAll = op.getNumResults() != 0;
       for (Value res : op.getResults())
-        replacedAll &=
-            succeeded(replaceWithConstant(analysis, builder, folder, res));
+        replacedAll &= (succeeded(replaceWithConstant(rangeAnalysis, builder,
+                                                      folder, res)) ||
+                        succeeded(replaceWithConstant(sccpAnalysis, builder,
+                                                      folder, res)));
 
       // If all of the results of the operation were replaced, try to erase
       // the operation completely.
       if (replacedAll && wouldOpBeTriviallyDead(&op)) {
         assert(op.use_empty() && "expected all uses to be replaced");
         op.erase();
         continue;
       }
 
       // Add any the regions of this operation to the worklist.
       addToWorklist(op.getRegions());
     }
 
     // Replace any block arguments with constants.
     builder.setInsertionPointToStart(block);
     for (BlockArgument arg : block->getArguments())
-      (void)replaceWithConstant(analysis, builder, folder, arg);
+      if (failed(replaceWithConstant(rangeAnalysis, builder, folder, arg)))
+        (void)replaceWithConstant(sccpAnalysis, builder, folder, arg);
   }
 }
 
 //===----------------------------------------------------------------------===//
 // SCCP Pass
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct SCCP : public SCCPBase<SCCP> {
   void runOnOperation() override;
 };
 } // namespace
 
 void SCCP::runOnOperation() {
   Operation *op = getOperation();
 
+  IntRangeAnalysis rangeAnalysis(op);
   SCCPAnalysis analysis(op->getContext());
   analysis.run(op);
-  rewrite(analysis, op->getContext(), op->getRegions());
+  rewrite(rangeAnalysis, analysis, op->getContext(), op->getRegions());
 }
 
 std::unique_ptr<Pass> mlir::createSCCPPass() {
   return std::make_unique<SCCP>();
 }

mlir/test/Dialect/Arithmetic/int-range-interface.mlir

-// RUN: mlir-opt -test-int-range-inference -canonicalize %s | FileCheck %s
+// RUN: mlir-opt -sccp -canonicalize %s | FileCheck %s
 
 // CHECK-LABEL: func @add_min_max
 // CHECK: %[[c3:.*]] = arith.constant 3 : index
 // CHECK: return %[[c3]]
 func.func @add_min_max(%a: index, %b: index) -> index {
     %c1 = arith.constant 1 : index
     %c2 = arith.constant 2 : index
     %0 = arith.minsi %a, %c1 : index

mlir/test/Interfaces/InferIntRangeInterface/infer-int-range-test-ops.mlir

-// RUN: mlir-opt -test-int-range-inference %s | FileCheck %s
+// RUN: mlir-opt -sccp %s | FileCheck %s
 
 // CHECK-LABEL: func @constant
 // CHECK: %[[cst:.*]] = "test.constant"() {value = 3 : index}
 // CHECK: return %[[cst]]
 func.func @constant() -> index {
   %0 = test.with_bounds { umin = 3 : index, umax = 3 : index,
                                smin = 3 : index, smax = 3 : index}
   func.return %0 : index

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestConstantFold.cpp
   TestControlFlowSink.cpp
   TestInlining.cpp
-  TestIntRangeInference.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   LINK_LIBS PUBLIC
   MLIRAnalysis

mlir/test/lib/Transforms/TestIntRangeInference.cpp

-//===- TestIntRangeInference.cpp - Create consts from range inference ---===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-// TODO: This pass is needed to test integer range inference until that
-// functionality has been integrated into SCCP.
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Analysis/IntRangeAnalysis.h"
-#include "mlir/Interfaces/SideEffectInterfaces.h"
-#include "mlir/Pass/Pass.h"
-#include "mlir/Pass/PassRegistry.h"
-#include "mlir/Support/TypeID.h"
-#include "mlir/Transforms/FoldUtils.h"
-
-using namespace mlir;
-
-/// Patterned after SCCP
-static LogicalResult replaceWithConstant(IntRangeAnalysis &analysis,
-                                         OpBuilder &b, OperationFolder &folder,
-                                         Value value) {
-  Optional<ConstantIntRanges> maybeInferredRange = analysis.getResult(value);
-  if (!maybeInferredRange)
-    return failure();
-  const ConstantIntRanges &inferredRange = maybeInferredRange.getValue();
-  Optional<APInt> maybeConstValue = inferredRange.getConstantValue();
-  if (!maybeConstValue.hasValue())
-    return failure();
-
-  Operation *maybeDefiningOp = value.getDefiningOp();
-  Dialect *valueDialect =
-      maybeDefiningOp ? maybeDefiningOp->getDialect()
-                      : value.getParentRegion()->getParentOp()->getDialect();
-  Attribute constAttr = b.getIntegerAttr(value.getType(), *maybeConstValue);
-  Value constant = folder.getOrCreateConstant(b, valueDialect, constAttr,
-                                              value.getType(), value.getLoc());
-  if (!constant)
-    return failure();
-
-  value.replaceAllUsesWith(constant);
-  return success();
-}
-
-static void rewrite(IntRangeAnalysis &analysis, MLIRContext *context,
-                    MutableArrayRef<Region> initialRegions) {
-  SmallVector<Block *> worklist;
-  auto addToWorklist = [&](MutableArrayRef<Region> regions) {
-    for (Region &region : regions)
-      for (Block &block : llvm::reverse(region))
-        worklist.push_back(&block);
-  };
-
-  OpBuilder builder(context);
-  OperationFolder folder(context);
-
-  addToWorklist(initialRegions);
-  while (!worklist.empty()) {
-    Block *block = worklist.pop_back_val();
-
-    for (Operation &op : llvm::make_early_inc_range(*block)) {
-      builder.setInsertionPoint(&op);
-
-      // Replace any result with constants.
-      bool replacedAll = op.getNumResults() != 0;
-      for (Value res : op.getResults())
-        replacedAll &=
-            succeeded(replaceWithConstant(analysis, builder, folder, res));
-
-      // If all of the results of the operation were replaced, try to erase
-      // the operation completely.
-      if (replacedAll && wouldOpBeTriviallyDead(&op)) {
-        assert(op.use_empty() && "expected all uses to be replaced");
-        op.erase();
-        continue;
-      }
-
-      // Add any the regions of this operation to the worklist.
-      addToWorklist(op.getRegions());
-    }
-
-    // Replace any block arguments with constants.
-    builder.setInsertionPointToStart(block);
-    for (BlockArgument arg : block->getArguments())
-      (void)replaceWithConstant(analysis, builder, folder, arg);
-  }
-}
-
-namespace {
-struct TestIntRangeInference
-    : PassWrapper<TestIntRangeInference, OperationPass<>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestIntRangeInference)
-
-  StringRef getArgument() const final { return "test-int-range-inference"; }
-  StringRef getDescription() const final {
-    return "Test integer range inference analysis";
-  }
-
-  void runOnOperation() override {
-    Operation *op = getOperation();
-    IntRangeAnalysis analysis(op);
-    rewrite(analysis, op->getContext(), op->getRegions());
-  }
-};
-} // end anonymous namespace
-
-namespace mlir {
-namespace test {
-void registerTestIntRangeInference() {
-  PassRegistration<TestIntRangeInference>();
-}
-} // end namespace test
-} // end namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestDataLayoutQuery();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDiagnosticsPass();
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestExpandMathPass();
 void registerTestComposeSubView();
 void registerTestMultiBuffering();
-void registerTestIntRangeInference();
 void registerTestIRVisitorsPass();
 void registerTestGenericIRVisitorsPass();
 void registerTestGenericIRVisitorsInterruptPass();
 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();
 void registerTestLinalgElementwiseFusion();
 void registerTestLinalgFusionTransforms();
 void registerTestLinalgTensorFusionTransforms();
#endif
   mlir::test::registerTestDecomposeCallGraphTypes();
   mlir::test::registerTestDataFlowPass();
   mlir::test::registerTestDataLayoutQuery();
   mlir::test::registerTestDominancePass();
   mlir::test::registerTestDynamicPipelinePass();
   mlir::test::registerTestExpandMathPass();
   mlir::test::registerTestComposeSubView();
   mlir::test::registerTestMultiBuffering();
-  mlir::test::registerTestIntRangeInference();
   mlir::test::registerTestIRVisitorsPass();
   mlir::test::registerTestGenericIRVisitorsPass();
   mlir::test::registerTestInterfaces();
   mlir::test::registerTestLinalgCodegenStrategy();
   mlir::test::registerTestLinalgElementwiseFusion();
   mlir::test::registerTestLinalgFusionTransforms();
   mlir::test::registerTestLinalgTensorFusionTransforms();
   mlir::test::registerTestLinalgTiledLoopFusionTransforms();