Add an integer divisibility analysis. (#18727)

* Also extends the numeric optimization test to elide arith.remui that
exactly divides (a reasonable optimization but primarily used for
testing in this patch).
* Only implements the analysis for `arith.constant` and
`util.int.assume` in this patch.
* Renames assumption `divisor` to `udiv` to match terminology elsewhere
wrt signed/unsigned analysis.
* The lattice tracks unsigned and signed interpretations separately as
this is needed for propagation through signed ops (but this is not
implemented here).

---------

Signed-off-by: Stella Laurenzo <stellaraccident@gmail.com>

compiler/plugins/input/Torch/InputConversion/test/bind_symbolic_shapes.mlir

@@ -19,7 +19,7 @@ module @basic_example {
     // CHECK: %[[ARG1_DIM0_RANGE:.*]] = util.assume.int %[[DIM0]]<umin = 1, umax = 1024>
     // CHECK: %[[MULTIPLIER0:.*]] = arith.constant 2 : index
     // CHECK: %[[ARG1_DIM1:.*]] = arith.muli %[[DIM1]], %[[MULTIPLIER0]]
-    // CHECK: %[[ARG1_DIM1_RANGE:.*]] = util.assume.int %[[ARG1_DIM1]]<umin = 2, umax = 2048, divisor = 2> : index
+    // CHECK: %[[ARG1_DIM1_RANGE:.*]] = util.assume.int %[[ARG1_DIM1]]<umin = 2, umax = 2048, udiv = 2> : index
     // CHECK: %[[ARG1_TIE:.*]] = flow.tensor.tie_shape %[[ARG1_ANCHOR]] : tensor<?x?xf32>{%[[ARG1_DIM0_RANGE]], %[[ARG1_DIM1_RANGE]]}
     // CHECK: %[[ARG1_EXPORT:.*]] = torch_c.from_builtin_tensor %[[ARG1_TIE]]
     %0 = torch.symbolic_int "s0" {min_val = 0, max_val = 1024} : !torch.int
@@ -49,10 +49,10 @@ module @basic_example {
 module @unbacked_symbol {
   func.func @main(%arg0: !torch.vtensor<[?,?],f32>, %arg1: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
     // CHECK: util.assume.int{{.*}}<umin = 1, umax = 1024>
-    // CHECK: util.assume.int{{.*}}<umin = 2, umax = 2048, divisor = 2>
+    // CHECK: util.assume.int{{.*}}<umin = 2, umax = 2048, udiv = 2>
     // CHECK: tie_shape
     // CHECK: util.assume.int{{.*}}<umin = 1, umax = 1024>
-    // CHECK: util.assume.int{{.*}}<umin = 4, umax = 4096, divisor = 4>
+    // CHECK: util.assume.int{{.*}}<umin = 4, umax = 4096, udiv = 4>
     // CHECK: tie_shape
     %0 = torch.symbolic_int "s0" {min_val = 0, max_val = 1024} : !torch.int
     %1 = torch.symbolic_int "2*s4" {min_val = 0, max_val = 2048} : !torch.int
@@ -100,7 +100,7 @@ module @all_bindings_dropped {
 module @add_expr {
   func.func @main(%arg0: !torch.vtensor<[?,?],f32>, %arg1: !torch.vtensor<[?,?],f32>) {
     // CHECK: addi
-    // CHECK-NOT: divisor
+    // CHECK-NOT: udiv
     %0 = torch.symbolic_int "s0" {min_val = 0, max_val = 1024} : !torch.int
     %1 = torch.symbolic_int "s1" {min_val = 0, max_val = 1024} : !torch.int
     torch.bind_symbolic_shape %arg0, [%0, %1], affine_map<()[s0, s1] -> (s0, s1)> : !torch.vtensor<[?,?],f32>
@@ -114,7 +114,7 @@ module @add_expr {
 module @mod_expr {
   func.func @main(%arg0: !torch.vtensor<[?,?],f32>, %arg1: !torch.vtensor<[?,?],f32>) {
     // CHECK: remui
-    // CHECK-NOT: divisor
+    // CHECK-NOT: udiv
     %0 = torch.symbolic_int "s0" {min_val = 0, max_val = 1024} : !torch.int
     %1 = torch.symbolic_int "s1" {min_val = 0, max_val = 1024} : !torch.int
     torch.bind_symbolic_shape %arg0, [%0, %1], affine_map<()[s0, s1] -> (s0, s1)> : !torch.vtensor<[?,?],f32>
@@ -128,7 +128,7 @@ module @mod_expr {
 module @floordiv_expr {
   func.func @main(%arg0: !torch.vtensor<[?,?],f32>, %arg1: !torch.vtensor<[?,?],f32>) {
     // CHECK: divui
-    // CHECK-NOT: divisor
+    // CHECK-NOT: udiv
     %0 = torch.symbolic_int "s0" {min_val = 0, max_val = 1024} : !torch.int
     %1 = torch.symbolic_int "s1" {min_val = 0, max_val = 1024} : !torch.int
     torch.bind_symbolic_shape %arg0, [%0, %1], affine_map<()[s0, s1] -> (s0, s1)> : !torch.vtensor<[?,?],f32>

compiler/src/iree/compiler/Dialect/Util/Analysis/BUILD.bazel

@@ -17,11 +17,13 @@ iree_compiler_cc_library(
     srcs = [
         "Explorer.cpp",
         "GlobalTable.cpp",
+        "IntegerDivisibilityAnalysis.cpp",
         "Position.cpp",
     ],
     hdrs = [
         "Explorer.h",
         "GlobalTable.h",
+        "IntegerDivisibilityAnalysis.h",
         "Position.h",
     ],
     deps = [

compiler/src/iree/compiler/Dialect/Util/Analysis/CMakeLists.txt

@@ -16,10 +16,12 @@ iree_cc_library(
   HDRS
     "Explorer.h"
     "GlobalTable.h"
+    "IntegerDivisibilityAnalysis.h"
     "Position.h"
   SRCS
     "Explorer.cpp"
     "GlobalTable.cpp"
+    "IntegerDivisibilityAnalysis.cpp"
     "Position.cpp"
   DEPS
     LLVMSupport

compiler/src/iree/compiler/Dialect/Util/Analysis/IntegerDivisibilityAnalysis.cpp

@@ -0,0 +1,69 @@
+// Copyright 2024 The IREE Authors
+//
+// Licensed 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 "iree/compiler/Dialect/Util/Analysis/IntegerDivisibilityAnalysis.h"
+
+#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "iree-util-int-divisibility-analysis"
+
+using llvm::dbgs;
+
+namespace mlir::iree_compiler::IREE::Util {
+
+void IntegerDivisibilityAnalysis::setToEntryState(
+    IntegerDivisibilityLattice *lattice) {
+  propagateIfChanged(lattice,
+                     lattice->join(IntegerDivisibility::getMinDivisibility()));
+}
+
+LogicalResult IntegerDivisibilityAnalysis::visitOperation(
+    Operation *op, ArrayRef<const IntegerDivisibilityLattice *> operands,
+    ArrayRef<IntegerDivisibilityLattice *> results) {
+  auto inferrable = dyn_cast<InferIntDivisibilityOpInterface>(op);
+  if (!inferrable) {
+    setAllToEntryStates(results);
+    return success();
+  }
+
+  LLVM_DEBUG(dbgs() << "Inferring divisibility for " << *op << "\n");
+  auto argDivs = llvm::map_to_vector(
+      operands, [](const IntegerDivisibilityLattice *lattice) {
+        return lattice->getValue();
+      });
+  auto joinCallback = [&](Value v, const IntegerDivisibility &newDiv) {
+    auto result = dyn_cast<OpResult>(v);
+    if (!result)
+      return;
+    assert(llvm::is_contained(op->getResults(), result));
+
+    LLVM_DEBUG(dbgs() << "Inferred divisibility " << newDiv << "\n");
+    IntegerDivisibilityLattice *lattice = results[result.getResultNumber()];
+    IntegerDivisibility oldDiv = lattice->getValue();
+
+    ChangeResult changed = lattice->join(newDiv);
+
+    // Catch loop results with loop variant bounds and conservatively make
+    // them [-inf, inf] so we don't circle around infinitely often (because
+    // the dataflow analysis in MLIR doesn't attempt to work out trip counts
+    // and often can't).
+    bool isYieldedResult = llvm::any_of(v.getUsers(), [](Operation *op) {
+      return op->hasTrait<OpTrait::IsTerminator>();
+    });
+    if (isYieldedResult && !oldDiv.isUninitialized() &&
+        !(lattice->getValue() == oldDiv)) {
+      LLVM_DEBUG(llvm::dbgs() << "Loop variant loop result detected\n");
+      changed |= lattice->join(IntegerDivisibility::getMinDivisibility());
+    }
+    propagateIfChanged(lattice, changed);
+  };
+
+  inferrable.inferResultDivisibility(argDivs, joinCallback);
+  return success();
+}
+
+} // namespace mlir::iree_compiler::IREE::Util

compiler/src/iree/compiler/Dialect/Util/Analysis/IntegerDivisibilityAnalysis.h

@@ -0,0 +1,42 @@
+// Copyright 2024 The IREE Authors
+//
+// Licensed 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 IREE_COMPILER_DIALECT_UTIL_INTEGER_DIVISIBILITY_ANALYSIS_H_
+#define IREE_COMPILER_DIALECT_UTIL_INTEGER_DIVISIBILITY_ANALYSIS_H_
+
+#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
+#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
+
+#include <optional>
+
+namespace mlir::iree_compiler::IREE::Util {
+
+class IntegerDivisibilityLattice
+    : public dataflow::Lattice<IntegerDivisibility> {
+public:
+  using Lattice::Lattice;
+};
+
+class IntegerDivisibilityAnalysis
+    : public dataflow::SparseForwardDataFlowAnalysis<
+          IntegerDivisibilityLattice> {
+public:
+  using SparseForwardDataFlowAnalysis::SparseForwardDataFlowAnalysis;
+
+  // At an entry point, set the lattice to the most pessimistic state,
+  // indicating that no further reasoning can be done.
+  void setToEntryState(IntegerDivisibilityLattice *lattice) override;
+
+  // Visit an operation, invoking the transfer function.
+  LogicalResult
+  visitOperation(Operation *op,
+                 ArrayRef<const IntegerDivisibilityLattice *> operands,
+                 ArrayRef<IntegerDivisibilityLattice *> results) override;
+};
+
+} // namespace mlir::iree_compiler::IREE::Util
+
+#endif // IREE_COMPILER_DIALECT_UTIL_INTEGER_DIVISIBILITY_ANALYSIS_H_

compiler/src/iree/compiler/Dialect/Util/IR/UtilAttrs.td

@@ -30,10 +30,10 @@ def Util_IntAssumptionAttr : AttrDef<Util_Dialect, "IntAssumption", []
   let parameters = (ins
     DefaultValuedParameter<"std::optional<uint64_t>", "std::nullopt">:$umin,
     DefaultValuedParameter<"std::optional<uint64_t>", "std::nullopt">:$umax,
-    DefaultValuedParameter<"std::optional<uint64_t>", "std::nullopt">:$divisor
+    DefaultValuedParameter<"std::optional<uint64_t>", "std::nullopt">:$udiv
   );
   let assemblyFormat = [{
-    `<` struct($umin, $umax, $divisor) `>`
+    `<` struct($umin, $umax, $udiv) `>`
   }];
 }
 

compiler/src/iree/compiler/Dialect/Util/IR/UtilInterfaces.td

@@ -94,6 +94,32 @@ def Util_ClosureOpInterface : OpInterface<"ClosureOpInterface"> {
   ];
 }
 
+//===----------------------------------------------------------------------===//
+// IREE::Util::InferIntDivisibilityOpInterface
+//===----------------------------------------------------------------------===//
+
+def InferIntDivisibilityOpInterface :
+    OpInterface<"InferIntDivisibilityOpInterface"> {
+  let cppNamespace = "::mlir::iree_compiler::IREE::Util";
+
+  let description = [{
+    Allows operations to participate in integer divisibility analysis.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+    /*desc=*/[{
+
+    }],
+    /*retTy=*/"void",
+    /*methodName=*/"inferResultDivisibility",
+    /*args=*/(ins
+        "::llvm::ArrayRef<::mlir::iree_compiler::IREE::Util::IntegerDivisibility>":$argDivs,
+        "::mlir::iree_compiler::IREE::Util::SetIntDivisibilityFn":$setResultDivs)
+    >
+  ];
+}
+
 //===----------------------------------------------------------------------===//
 // IREE::Util::InitializerOpInterface
 //===----------------------------------------------------------------------===//

compiler/src/iree/compiler/Dialect/Util/IR/UtilOps.cpp

@@ -1141,11 +1141,12 @@ AssumeIntOp::getUnionedUnsignedRange(unsigned operandIndex) {
 // Gets the unioned divisor for an operand. If there are multiple divisor
 // assumptions, the gcd of all of them is returned. If there are no
 // divisor assumptions, std::nullopt is returned.
-std::optional<uint64_t> AssumeIntOp::getUnionedDivisor(unsigned operandIndex) {
+std::optional<uint64_t>
+AssumeIntOp::getUnionedUnsignedDivisor(unsigned operandIndex) {
   auto assumptions = getOperandAssumptions(operandIndex);
   std::optional<uint64_t> divisorUnion;
   for (auto assumption : assumptions) {
-    auto divisor = assumption.getDivisor();
+    auto divisor = assumption.getUdiv();
     if (divisor) {
       if (divisorUnion)
         divisorUnion = std::gcd(*divisor, *divisorUnion);
@@ -1176,6 +1177,20 @@ void AssumeIntOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
   }
 }
 
+void AssumeIntOp::inferResultDivisibility(ArrayRef<IntegerDivisibility> argDivs,
+                                          SetIntDivisibilityFn setResultDivs) {
+  for (auto [index, result] : llvm::enumerate(getResults())) {
+    Type type = result.getType();
+    if (!isa<IndexType>(type) && !isa<IntegerType>(type))
+      continue;
+    auto udiv = getUnionedUnsignedDivisor(index);
+    if (udiv) {
+      setResultDivs(result,
+                    ConstantIntDivisibility(/*udiv=*/*udiv, /*sdiv=*/*udiv));
+    }
+  }
+}
+
 void AssumeIntOp::build(OpBuilder &builder, OperationState &state,
                         Value singleOperand,
                         IntAssumptionAttr singleAssumption) {

compiler/src/iree/compiler/Dialect/Util/IR/UtilOps.td

@@ -460,7 +460,8 @@ def OpGroupCompilerHintOps : OpDocGroup {
 let opDocGroup = OpGroupCompilerHintOps in {
 
 def Util_AssumeIntOp : Util_PureOp<"assume.int", [
-    DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>
+    DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>,
+    DeclareOpInterfaceMethods<InferIntDivisibilityOpInterface, ["inferResultRanges"]>
 ]> {
   let summary = "memorializes assumptions about index/integer values.";
   let description = [{
@@ -507,7 +508,7 @@ def Util_AssumeIntOp : Util_PureOp<"assume.int", [
     // Gets the unioned divisor for an operand. If there are multiple divisor
     // assumptions, the gcd of all of them is returned. If there are no
     // divisor assumptions, std::nullopt is returned.
-    std::optional<uint64_t> getUnionedDivisor(unsigned operandIndex);
+    std::optional<uint64_t> getUnionedUnsignedDivisor(unsigned operandIndex);
   }];
 
   let hasCustomAssemblyFormat = 1;

compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.h

@@ -22,6 +22,8 @@
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 
+#include <numeric>
+
 // clang-format off: must be included after all LLVM/MLIR headers.
 #include "iree/compiler/Dialect/Util/IR/UtilEnums.h.inc" // IWYU pragma: keep
 // clang-format on
@@ -155,6 +157,94 @@ void excludeTiedOperandAndResultIndices(
     ArrayRef<unsigned> excludedResultIndices,
     SmallVector<int64_t> &tiedOperandIndices);
 
+//===----------------------------------------------------------------------===//
+// Forward defines for InferIntDivisibilityOpInterface
+// See implementations in IntegerDivisibility.h.
+//===----------------------------------------------------------------------===//
+
+class ConstantIntDivisibility {
+public:
+  ConstantIntDivisibility() = default;
+  ConstantIntDivisibility(uint64_t udiv, uint64_t sdiv)
+      : udivVal(udiv), sdivVal(sdiv) {}
+
+  bool operator==(const ConstantIntDivisibility &other) const {
+    return udivVal == other.udivVal && sdivVal == other.sdivVal;
+  }
+
+  uint64_t udiv() const { return this->udivVal; }
+  uint64_t sdiv() const { return this->sdivVal; }
+
+  // Returns the union (computed separately for signed and unsigned bounds)
+  // for this range and `other`.
+  ConstantIntDivisibility getUnion(const ConstantIntDivisibility &other) const {
+    return ConstantIntDivisibility(
+        /*udiv=*/std::gcd(udiv(), other.udiv()),
+        /*sdiv=*/std::gcd(sdiv(), other.sdiv()));
+  }
+
+private:
+  uint64_t udivVal;
+  uint64_t sdivVal;
+
+  friend raw_ostream &operator<<(raw_ostream &os,
+                                 const ConstantIntDivisibility &div);
+};
+
+inline raw_ostream &operator<<(raw_ostream &os,
+                               const ConstantIntDivisibility &div) {
+  os << "ConstantIntDivisibility(udiv = " << div.udivVal
+     << ", sdiv = " << div.sdivVal << ")";
+  return os;
+}
+
+class IntegerDivisibility {
+public:
+  IntegerDivisibility(ConstantIntDivisibility value)
+      : value(std::move(value)) {}
+  IntegerDivisibility(
+      std::optional<ConstantIntDivisibility> value = std::nullopt)
+      : value(std::move(value)) {}
+  // Gets the minimum divisibility of 1 that is used to indicate that the value
+  // cannot be analyzed further.
+  static IntegerDivisibility getMinDivisibility() {
+    return IntegerDivisibility(ConstantIntDivisibility(1, 1));
+  }
+
+  bool isUninitialized() const { return !value.has_value(); }
+  const ConstantIntDivisibility &getValue() const {
+    assert(!isUninitialized());
+    return *value;
+  }
+
+  bool operator==(const IntegerDivisibility &rhs) const {
+    return value == rhs.value;
+  }
+
+  static IntegerDivisibility join(const IntegerDivisibility &lhs,
+                                  const IntegerDivisibility &rhs) {
+    if (lhs.isUninitialized())
+      return rhs;
+    if (rhs.isUninitialized())
+      return lhs;
+    return IntegerDivisibility(lhs.getValue().getUnion(rhs.getValue()));
+  }
+
+  void print(raw_ostream &os) const { os << value; }
+
+private:
+  std::optional<ConstantIntDivisibility> value;
+};
+
+inline raw_ostream &operator<<(raw_ostream &os,
+                               const IntegerDivisibility &div) {
+  div.print(os);
+  return os;
+}
+
+using SetIntDivisibilityFn =
+    llvm::function_ref<void(Value, const ConstantIntDivisibility &)>;
+
 //===----------------------------------------------------------------------===//
 // Shape-aware interface utilities
 //===----------------------------------------------------------------------===//