[mlir][arith] Optimize arith.cmpi based on integer range analysis.

[Hardcode84]

Summary

Add a pass which do arith dialect ops optimization based on integer range analysis (only cmpi for now).

[Hardcode84]

fromUnsigned -> fromSigned

[Mogball]

mlir/include/mlir/Dialect/Arith/Transforms/Passes.td

52

Can you merge this pass with ArithUnsignedWhenEquivalentPass? Both of them use the results of integer range inference and it seems that pass can be considered an "integer range optimization" anyways.

mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp

17–22

28

please document this function: what it does and its purpose

117

do you need a pattern for this? Why not just walk the top-level operation of the pass and find all arith.cmpi ops?

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

1160 ↗

(On Diff #485114)

Can you add an op to TestOps.td? I think this creates a better separation of concerns with this patch. You can implement the interface on Dialect operations as necessary in subsequent patches.

[Hardcode84]

mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp

117

Downstream we are building custom pass which combines integer range analysis with some shape analysis and we need this populateIntRangeOptimizationsPatterns function as a building block, standalone pass needed just to test things in isolation. I can remove dedicated pass and just add these patterns to UnsignedWhenEquivalentPass (although I don't think UnsignedWhenEquivalent is strictly an optimization pass, as backends can have different preferences between signed and unsigned ops) but I need to keep either populateIntRangeOptimizationsPatterns or some other way to express these optimizations separated from pass.

[Mogball]

mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp

117

Okay, this rationale makes sense. Thanks. You can keep the separate pass

[Hardcode84]

review comments

[bondhugula]

mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp

182

if (failed(...) return failure()`

This should converge.

182

-> signalPassFailure()

[Hardcode84]

check applyPatternsAndFoldGreedily result

[Mogball]

mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp

149

please just use a switch statement for this

[Hardcode84]

use switch

[Hardcode84]

mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp

149

I did it, but it just makes code 2 time bigger and actually less readable

[Mogball]

mlir/include/mlir/Dialect/Arith/Transforms/Passes.h

20

You can do this for the new pass gen style

mlir/include/mlir/Dialect/Arith/Transforms/Passes.td

59

Can you omit the constructor and use the new style of pass generation?

mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp

27–31

?

34

please return FailureOr<bool> for these functions instead. Also, just FYI but you can drop the llvm:: on Optional here

149

If you want to compress the code size, you can outline this switch and do

case Pred::slt:
  return handleSlt(lhsValue, rhsValue);
case Pred::sle: ...

Also, the compiler will turn the switch into a jump table, so I don't think you should be manually creating that.

[Hardcode84]

mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp

149

I've tried nested lambda with return but it is even uglier. And I was using table not because I wanted to jump table in asm but because it is a common idiom. switch not worth it, I'm switching back to table.

[Hardcode84]

review comments

[krzysz00]

Very delayedly, I'd like to point out that a more general version of this code lives out in mlir/test/lib/Transforms/TestIntRangeInference.cpp and thet perhaps we might just want to make that a non-test pass?

[krzysz00]

... Ok, that and UnsignedWhenEquivalent

[Hardcode84]

In D140629#4322492, @krzysz00 wrote:

Very delayedly, I'd like to point out that a more general version of this code lives out in mlir/test/lib/Transforms/TestIntRangeInference.cpp and thet perhaps we might just want to make that a non-test pass?

Make sense, only requirements from our side is that it must be composable, as we have our custom IntegerRangeAnalysisEx derived from IntegerRangeAnalysis and we adding some additional custom analyses to DataFlowSolver.

[krzysz00]

Entirely fair - I'm curious about what extensions you've made and if they
can be pulled upstream.

I'd be happy to discuss whether there's something we can unify on here
since it looks like we both have downstream code that uses this analysis.

[Hardcode84]

Here is our pass https://github.com/numba/numba-mlir/blob/main/mlir/lib/Transforms/ShapeIntegerRangePropagation.cpp

We have extended int range analysis for shaped ops (tensor, memref, and our custom ntensor dialects) and propagating int range for each dim independently coupling it with existing analysis on shaped.dim ops.
We were planning to upstream this eventually but don't have much throughput at the moment.

mlir/include/mlir/Dialect/Arith/Transforms/Passes.h

 //===- Passes.h - Pass Entrypoints ------------------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_ARITH_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_ARITH_TRANSFORMS_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
+class DataFlowSolver;
+
 namespace arith {
 
 #define GEN_PASS_DECL
 #include "mlir/Dialect/Arith/Transforms/Passes.h.inc"
 
 class WideIntEmulationConverter;
 
 /// Create a pass to bufferize Arith ops.
 std::unique_ptr<Pass> createArithBufferizePass();
 
 /// Create a pass to bufferize arith.constant ops.
 std::unique_ptr<Pass> createConstantBufferizePass(uint64_t alignment = 0);

 
 /// Create a pass to legalize Arith ops.
 std::unique_ptr<Pass> createArithExpandOpsPass();
 
 /// Create a pass to replace signed ops with unsigned ones where they are proven
 /// equivalent.
 std::unique_ptr<Pass> createArithUnsignedWhenEquivalentPass();
 
+/// Add patterns for int range based optimizations.
+void populateIntRangeOptimizationsPatterns(RewritePatternSet &patterns,
+                                           DataFlowSolver &solver);
+
+/// Create a pass which do optimizations based on integer range analysis.
+std::unique_ptr<Pass> createIntRangeOptimizationsPass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Arith/Transforms/Passes.h.inc"
 
 } // namespace arith
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARITH_TRANSFORMS_PASSES_H_

mlir/include/mlir/Dialect/Arith/Transforms/Passes.td

let description = [{
     range [0, signed_max(type)].
 
     The affect ops include division, remainder, shifts, min, max, and integer
     comparisons.
   }];
   let constructor = "mlir::arith::createArithUnsignedWhenEquivalentPass()";
 }
 
+def ArithIntRangeOpts : Pass<"int-range-optimizations"> {
+  let summary = "Do optimizations based on integer range analysis";
+  let description = [{
+    This pass runs integer range analysis and apllies optimizations based on its
+    results. e.g. replace arith.cmpi with const if it can be inferred from
+    args ranges.
+  }];
+  let constructor = "mlir::arith::createIntRangeOptimizationsPass()";
+}
+
 def ArithEmulateWideInt : Pass<"arith-emulate-wide-int"> {
   let summary = "Emulate 2*N-bit integer operations using N-bit operations";
   let description = [{
     Emulate arith integer operations that use too wide integer types with
     equivalent operations on supported narrow integer types. This is done by
     splitting original integer values into two halves.
 
     This pass is intended preserve semantics but not necessarily provide the

mlir/lib/Dialect/Arith/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRArithTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   EmulateWideInt.cpp
   ExpandOps.cpp
+  IntRangeOptimizations.cpp
   UnsignedWhenEquivalent.cpp
 
   ADDITIONAL_HEADER_DIRS
   {$MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arith/Transforms
 
   DEPENDS
   MLIRArithTransformsIncGen
 

mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp

+//===- IntRangeOptimizations.cpp - Optimizations based on integer ranges --===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/Transforms/Passes.h"
+
+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/Analysis/DataFlow/IntegerRangeAnalysis.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::arith {
+#define GEN_PASS_DEF_ARITHINTRANGEOPTS
+#include "mlir/Dialect/Arith/Transforms/Passes.h.inc"
+} // namespace mlir::arith
+
+using namespace mlir;
+using namespace mlir::arith;
+using namespace mlir::dataflow;
+
+/// Returns true if 2 integer ranges have intersection.
+static bool intersects(ConstantIntRanges lhs, ConstantIntRanges rhs) {
+  if ((lhs.smax().slt(rhs.smin()) || lhs.smin().sgt(rhs.smax())) &&
+      (lhs.umax().ult(rhs.umin()) || lhs.umin().ugt(rhs.umax())))
+    return false;
+
+  return true;
+}
+
+static llvm::Optional<bool> handleEq(ConstantIntRanges lhs,
+                                     ConstantIntRanges rhs) {
+  if (!intersects(lhs, rhs))
+    return false;
+
+  return std::nullopt;
+}
+
+static llvm::Optional<bool> handleNe(ConstantIntRanges lhs,
+                                     ConstantIntRanges rhs) {
+  if (!intersects(lhs, rhs))
+    return true;
+
+  return std::nullopt;
+}
+
+static llvm::Optional<bool> handleSlt(ConstantIntRanges lhs,
+                                      ConstantIntRanges rhs) {
+  if (lhs.smax().slt(rhs.smin()))
+    return true;
+
+  if (lhs.smin().sge(rhs.smax()))
+    return false;
+
+  return std::nullopt;
+}
+
+static llvm::Optional<bool> handleSle(ConstantIntRanges lhs,
+                                      ConstantIntRanges rhs) {
+  if (lhs.smax().sle(rhs.smin()))
+    return true;
+
+  if (lhs.smin().sgt(rhs.smax()))
+    return false;
+
+  return std::nullopt;
+}
+
+static llvm::Optional<bool> handleSgt(ConstantIntRanges lhs,
+                                      ConstantIntRanges rhs) {
+  return handleSlt(rhs, lhs);
+}
+
+static llvm::Optional<bool> handleSge(ConstantIntRanges lhs,
+                                      ConstantIntRanges rhs) {
+  return handleSle(rhs, lhs);
+}
+
+static llvm::Optional<bool> handleUlt(ConstantIntRanges lhs,
+                                      ConstantIntRanges rhs) {
+  if (lhs.umax().ult(rhs.umin()))
+    return true;
+
+  if (lhs.umin().uge(rhs.umax()))
+    return false;
+
+  return std::nullopt;
+}
+
+static llvm::Optional<bool> handleUle(ConstantIntRanges lhs,
+                                      ConstantIntRanges rhs) {
+  if (lhs.umax().ule(rhs.umin()))
+    return true;
+
+  if (lhs.umin().ugt(rhs.umax()))
+    return false;
+
+  return std::nullopt;
+}
+
+static llvm::Optional<bool> handleUgt(ConstantIntRanges lhs,
+                                      ConstantIntRanges rhs) {
+  return handleUlt(rhs, lhs);
+}
+
+static llvm::Optional<bool> handleUge(ConstantIntRanges lhs,
+                                      ConstantIntRanges rhs) {
+  return handleUle(rhs, lhs);
+}
+
+namespace {
+struct ConvertCmpOp : public OpRewritePattern<arith::CmpIOp> {
+
+  ConvertCmpOp(MLIRContext *context, DataFlowSolver &s)
+      : OpRewritePattern<arith::CmpIOp>(context), solver(s) {}
+
+  LogicalResult matchAndRewrite(arith::CmpIOp op,
+                                PatternRewriter &rewriter) const override {
+    auto *lhsResult =
+        solver.lookupState<dataflow::IntegerValueRangeLattice>(op.getLhs());
+    if (!lhsResult || lhsResult->getValue().isUninitialized())
+      return failure();
+
+    auto *rhsResult =
+        solver.lookupState<dataflow::IntegerValueRangeLattice>(op.getRhs());
+    if (!rhsResult || rhsResult->getValue().isUninitialized())
+      return failure();
+
+    ConstantIntRanges lhsValue = lhsResult->getValue().getValue();
+    ConstantIntRanges rhsValue = rhsResult->getValue().getValue();
+    llvm::Optional<bool> result;
+
+    using Pred = arith::CmpIPredicate;
+    switch (op.getPredicate()) {
+    case Pred::eq:
+      result = handleEq(lhsValue, rhsValue);
+      break;
+    case Pred::ne:
+      result = handleNe(lhsValue, rhsValue);
+      break;
+    case Pred::slt:
+      result = handleSlt(lhsValue, rhsValue);
+      break;
+    case Pred::sle:
+      result = handleSle(lhsValue, rhsValue);
+      break;
+    case Pred::sgt:
+      result = handleSgt(lhsValue, rhsValue);
+      break;
+    case Pred::sge:
+      result = handleSge(lhsValue, rhsValue);
+      break;
+    case Pred::ult:
+      result = handleUlt(lhsValue, rhsValue);
+      break;
+    case Pred::ule:
+      result = handleUle(lhsValue, rhsValue);
+      break;
+    case Pred::ugt:
+      result = handleUgt(lhsValue, rhsValue);
+      break;
+    case Pred::uge:
+      result = handleUge(lhsValue, rhsValue);
+      break;
+    }
+
+    if (!result)
+      return failure();
+
+    rewriter.replaceOpWithNewOp<arith::ConstantIntOp>(
+        op, static_cast<int64_t>(*result), /*width*/ 1);
+    return success();
+  }
+
+private:
+  DataFlowSolver &solver;
+};
+
+struct IntRangeOptimizationsPass
+    : public arith::impl::ArithIntRangeOptsBase<IntRangeOptimizationsPass> {
+
+  void runOnOperation() override {
+    Operation *op = getOperation();
+    MLIRContext *ctx = op->getContext();
+    DataFlowSolver solver;
+    solver.load<DeadCodeAnalysis>();
+    solver.load<IntegerRangeAnalysis>();
+    if (failed(solver.initializeAndRun(op)))
+      return signalPassFailure();
+
+    RewritePatternSet patterns(ctx);
+    populateIntRangeOptimizationsPatterns(patterns, solver);
+
+    if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+void mlir::arith::populateIntRangeOptimizationsPatterns(
+    RewritePatternSet &patterns, DataFlowSolver &solver) {
+  patterns.add<ConvertCmpOp>(patterns.getContext(), solver);
+}
+
+std::unique_ptr<Pass> mlir::arith::createIntRangeOptimizationsPass() {
+  return std::make_unique<IntRangeOptimizationsPass>();
+}

mlir/test/Dialect/Arith/int-range-opts.mlir

+// RUN: mlir-opt -int-range-optimizations --split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: func @test
+//       CHECK:   %[[C:.*]] = arith.constant false
+//       CHECK:   return %[[C]]
+func.func @test() -> i1 {
+  %cst1 = arith.constant -1 : index
+  %0 = test.with_bounds { umin = 0 : index, umax = 0x7fffffffffffffff : index, smin = 0 : index, smax = 0x7fffffffffffffff : index }
+  %1 = arith.cmpi eq, %0, %cst1 : index
+  return %1: i1
+}
+
+// -----
+
+// CHECK-LABEL: func @test
+//       CHECK:   %[[C:.*]] = arith.constant true
+//       CHECK:   return %[[C]]
+func.func @test() -> i1 {
+  %cst1 = arith.constant -1 : index
+  %0 = test.with_bounds { umin = 0 : index, umax = 0x7fffffffffffffff : index, smin = 0 : index, smax = 0x7fffffffffffffff : index }
+  %1 = arith.cmpi ne, %0, %cst1 : index
+  return %1: i1
+}
+
+// -----
+
+
+// CHECK-LABEL: func @test
+//       CHECK:   %[[C:.*]] = arith.constant true
+//       CHECK:   return %[[C]]
+func.func @test() -> i1 {
+  %cst = arith.constant 0 : index
+  %0 = test.with_bounds { umin = 0 : index, umax = 0x7fffffffffffffff : index, smin = 0 : index, smax = 0x7fffffffffffffff : index }
+  %1 = arith.cmpi sge, %0, %cst : index
+  return %1: i1
+}
+
+// -----
+
+// CHECK-LABEL: func @test
+//       CHECK:   %[[C:.*]] = arith.constant false
+//       CHECK:   return %[[C]]
+func.func @test() -> i1 {
+  %cst = arith.constant 0 : index
+  %0 = test.with_bounds { umin = 0 : index, umax = 0x7fffffffffffffff : index, smin = 0 : index, smax = 0x7fffffffffffffff : index }
+  %1 = arith.cmpi slt, %0, %cst : index
+  return %1: i1
+}
+
+// -----
+
+
+// CHECK-LABEL: func @test
+//       CHECK:   %[[C:.*]] = arith.constant true
+//       CHECK:   return %[[C]]
+func.func @test() -> i1 {
+  %cst1 = arith.constant -1 : index
+  %0 = test.with_bounds { umin = 0 : index, umax = 0x7fffffffffffffff : index, smin = 0 : index, smax = 0x7fffffffffffffff : index }
+  %1 = arith.cmpi sgt, %0, %cst1 : index
+  return %1: i1
+}
+
+// -----
+
+// CHECK-LABEL: func @test
+//       CHECK:   %[[C:.*]] = arith.constant false
+//       CHECK:   return %[[C]]
+func.func @test() -> i1 {
+  %cst1 = arith.constant -1 : index
+  %0 = test.with_bounds { umin = 0 : index, umax = 0x7fffffffffffffff : index, smin = 0 : index, smax = 0x7fffffffffffffff : index }
+  %1 = arith.cmpi sle, %0, %cst1 : index
+  return %1: i1
+}