[mlir][index] Add folders for index ops

This patch adds folders for `index` dialect ops. Ths folders are
careful to ensure that fold results are valid on both 32-bit and 64-bit
targets.

Depends on D135689

Reviewed By: rriddle

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

Author: Jeff Niu

mlir/include/mlir/Dialect/Index/IR/IndexDialect.td

@@ -81,6 +81,7 @@ def IndexDialect : Dialect {
     void registerOperations();
   }];
 
+  let hasConstantMaterializer = 1;
   let useDefaultAttributePrinterParser = 1;
 }
 

mlir/include/mlir/Dialect/Index/IR/IndexOps.td

@@ -34,6 +34,7 @@ class IndexBinaryOp<string mnemonic, list<Trait> traits = []>
   let arguments = (ins Index:$lhs, Index:$rhs);
   let results = (outs Index:$result);
   let assemblyFormat = "$lhs `,` $rhs attr-dict";
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
@@ -378,6 +379,7 @@ def Index_CmpOp : IndexOp<"cmp"> {
   let arguments = (ins IndexCmpPredicateAttr:$pred, Index:$lhs, Index:$rhs);
   let results = (outs I1:$result);
   let assemblyFormat = "`` $pred `(` $lhs `,` $rhs `)` attr-dict";
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
@@ -422,6 +424,7 @@ def Index_ConstantOp : IndexOp<"constant", [ConstantLike]> {
   let arguments = (ins IndexAttr:$value);
   let results = (outs Index:$result);
   let assemblyFormat = "attr-dict $value";
+  let hasFolder = 1;
 
   let builders = [OpBuilder<(ins "int64_t":$value)>];
 }
@@ -449,6 +452,7 @@ def Index_BoolConstantOp : IndexOp<"bool.constant", [ConstantLike]> {
   let arguments = (ins BoolAttr:$value);
   let results = (outs I1:$result);
   let assemblyFormat = "attr-dict $value";
+  let hasFolder = 1;
 }
 
 #endif // INDEX_OPS

mlir/include/mlir/Support/LLVM.h

@@ -99,6 +99,7 @@ namespace mlir {
 using llvm::cast;
 using llvm::cast_or_null;
 using llvm::dyn_cast;
+using llvm::dyn_cast_if_present;
 using llvm::dyn_cast_or_null;
 using llvm::isa;
 using llvm::isa_and_nonnull;

mlir/lib/Dialect/Index/IR/IndexOps.cpp

@@ -26,6 +26,264 @@ void IndexDialect::registerOperations() {
       >();
 }
 
+Operation *IndexDialect::materializeConstant(OpBuilder &b, Attribute value,
+                                             Type type, Location loc) {
+  // Materialize bool constants as `i1`.
+  if (auto boolValue = dyn_cast<BoolAttr>(value)) {
+    if (!type.isSignlessInteger(1))
+      return nullptr;
+    return b.create<BoolConstantOp>(loc, type, boolValue);
+  }
+
+  // Materialize integer attributes as `index`.
+  if (auto indexValue = dyn_cast<IntegerAttr>(value)) {
+    if (!indexValue.getType().isa<IndexType>() || !type.isa<IndexType>())
+      return nullptr;
+    assert(indexValue.getValue().getBitWidth() ==
+           IndexType::kInternalStorageBitWidth);
+    return b.create<ConstantOp>(loc, indexValue);
+  }
+
+  return nullptr;
+}
+
+//===----------------------------------------------------------------------===//
+// Fold Utilities
+//===----------------------------------------------------------------------===//
+
+/// Fold an index operation irrespective of the target bitwidth. The
+/// operation must satisfy the property:
+///
+/// ```
+/// trunc(f(a, b)) = f(trunc(a), trunc(b))
+/// ```
+///
+/// For all values of `a` and `b`. The function accepts a lambda that computes
+/// the integer result, which in turn must satisfy the above property.
+static OpFoldResult foldBinaryOpUnchecked(
+    ArrayRef<Attribute> operands,
+    function_ref<APInt(const APInt &, const APInt &)> calculate) {
+  assert(operands.size() == 2 && "binary operation expected 2 operands");
+  auto lhs = dyn_cast_if_present<IntegerAttr>(operands[0]);
+  auto rhs = dyn_cast_if_present<IntegerAttr>(operands[1]);
+  if (!lhs || !rhs)
+    return {};
+
+  APInt result = calculate(lhs.getValue(), rhs.getValue());
+  assert(result.trunc(32) ==
+         calculate(lhs.getValue().trunc(32), rhs.getValue().trunc(32)));
+  return IntegerAttr::get(IndexType::get(lhs.getContext()), std::move(result));
+}
+
+/// Fold an index operation only if the truncated 64-bit result matches the
+/// 32-bit result for operations that don't satisfy the above property. These
+/// are operations where the upper bits of the operands can affect the lower
+/// bits of the results.
+///
+/// The function accepts a lambda that computes the integer result in both
+/// 64-bit and 32-bit. If either call returns `None`, the operation is not
+/// folded.
+static OpFoldResult foldBinaryOpChecked(
+    ArrayRef<Attribute> operands,
+    function_ref<Optional<APInt>(const APInt &, const APInt &lhs)> calculate) {
+  assert(operands.size() == 2 && "binary operation expected 2 operands");
+  auto lhs = dyn_cast_if_present<IntegerAttr>(operands[0]);
+  auto rhs = dyn_cast_if_present<IntegerAttr>(operands[1]);
+  // Only fold index operands.
+  if (!lhs || !rhs)
+    return {};
+
+  // Compute the 64-bit result and the 32-bit result.
+  Optional<APInt> result64 = calculate(lhs.getValue(), rhs.getValue());
+  if (!result64)
+    return {};
+  Optional<APInt> result32 =
+      calculate(lhs.getValue().trunc(32), rhs.getValue().trunc(32));
+  if (!result32)
+    return {};
+  // Compare the truncated 64-bit result to the 32-bit result.
+  if (result64->trunc(32) != *result32)
+    return {};
+  // The operation can be folded for these particular operands.
+  return IntegerAttr::get(IndexType::get(lhs.getContext()),
+                          std::move(*result64));
+}
+
+//===----------------------------------------------------------------------===//
+// AddOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult AddOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpUnchecked(
+      operands, [](const APInt &lhs, const APInt &rhs) { return lhs + rhs; });
+}
+
+//===----------------------------------------------------------------------===//
+// SubOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult SubOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpUnchecked(
+      operands, [](const APInt &lhs, const APInt &rhs) { return lhs - rhs; });
+}
+
+//===----------------------------------------------------------------------===//
+// MulOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult MulOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpUnchecked(
+      operands, [](const APInt &lhs, const APInt &rhs) { return lhs * rhs; });
+}
+
+//===----------------------------------------------------------------------===//
+// DivSOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult DivSOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpChecked(
+      operands, [](const APInt &lhs, const APInt &rhs) -> Optional<APInt> {
+        // Don't fold division by zero.
+        if (rhs.isZero())
+          return None;
+        return lhs.sdiv(rhs);
+      });
+}
+
+//===----------------------------------------------------------------------===//
+// DivUOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult DivUOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpChecked(
+      operands, [](const APInt &lhs, const APInt &rhs) -> Optional<APInt> {
+        // Don't fold division by zero.
+        if (rhs.isZero())
+          return None;
+        return lhs.udiv(rhs);
+      });
+}
+
+//===----------------------------------------------------------------------===//
+// CeilDivSOp
+//===----------------------------------------------------------------------===//
+
+/// Compute `ceildivs(n, m)` as `x = m > 0 ? -1 : 1` and then
+/// `n*m > 0 ? (n+x)/m + 1 : -(-n/m)`.
+static Optional<APInt> calculateCeilDivS(const APInt &n, const APInt &m) {
+  // Don't fold division by zero.
+  if (m.isZero())
+    return None;
+  // Short-circuit the zero case.
+  if (n.isZero())
+    return n;
+
+  bool mGtZ = m.sgt(0);
+  if (n.sgt(0) != mGtZ) {
+    // If the operands have different signs, compute the negative result. Signed
+    // division overflow is not possible, since if `m == -1`, `n` can be at most
+    // `INT_MAX`, and `-INT_MAX != INT_MIN` in two's complement.
+    return -(-n).sdiv(m);
+  }
+  // Otherwise, compute the positive result. Signed division overflow is not
+  // possible since if `m == -1`, `x` will be `1`.
+  int64_t x = mGtZ ? -1 : 1;
+  return (n + x).sdiv(m) + 1;
+}
+
+OpFoldResult CeilDivSOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpChecked(operands, calculateCeilDivS);
+}
+
+//===----------------------------------------------------------------------===//
+// CeilDivUOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult CeilDivUOp::fold(ArrayRef<Attribute> operands) {
+  // Compute `ceildivu(n, m)` as `n == 0 ? 0 : (n-1)/m + 1`.
+  return foldBinaryOpChecked(
+      operands, [](const APInt &n, const APInt &m) -> Optional<APInt> {
+        // Don't fold division by zero.
+        if (m.isZero())
+          return None;
+        // Short-circuit the zero case.
+        if (n.isZero())
+          return n;
+
+        return (n - 1).udiv(m) + 1;
+      });
+}
+
+//===----------------------------------------------------------------------===//
+// FloorDivSOp
+//===----------------------------------------------------------------------===//
+
+/// Compute `floordivs(n, m)` as `x = m < 0 ? 1 : -1` and then
+/// `n*m < 0 ? -1 - (x-n)/m : n/m`.
+static Optional<APInt> calculateFloorDivS(const APInt &n, const APInt &m) {
+  // Don't fold division by zero.
+  if (m.isZero())
+    return None;
+  // Short-circuit the zero case.
+  if (n.isZero())
+    return n;
+
+  bool mLtZ = m.slt(0);
+  if (n.slt(0) == mLtZ) {
+    // If the operands have the same sign, compute the positive result.
+    return n.sdiv(m);
+  }
+  // If the operands have different signs, compute the negative result. Signed
+  // division overflow is not possible since if `m == -1`, `x` will be 1 and
+  // `n` can be at most `INT_MAX`.
+  int64_t x = mLtZ ? 1 : -1;
+  return -1 - (x - n).sdiv(m);
+}
+
+OpFoldResult FloorDivSOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpChecked(operands, calculateFloorDivS);
+}
+
+//===----------------------------------------------------------------------===//
+// RemSOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult RemSOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpChecked(operands, [](const APInt &lhs, const APInt &rhs) {
+    return lhs.srem(rhs);
+  });
+}
+
+//===----------------------------------------------------------------------===//
+// RemUOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult RemUOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpChecked(operands, [](const APInt &lhs, const APInt &rhs) {
+    return lhs.urem(rhs);
+  });
+}
+
+//===----------------------------------------------------------------------===//
+// MaxSOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult MaxSOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpChecked(operands, [](const APInt &lhs, const APInt &rhs) {
+    return lhs.sgt(rhs) ? lhs : rhs;
+  });
+}
+
+//===----------------------------------------------------------------------===//
+// MaxUOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult MaxUOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpChecked(operands, [](const APInt &lhs, const APInt &rhs) {
+    return lhs.ugt(rhs) ? lhs : rhs;
+  });
+}
+
 //===----------------------------------------------------------------------===//
 // CastSOp
 //===----------------------------------------------------------------------===//
@@ -42,6 +300,74 @@ bool CastUOp::areCastCompatible(TypeRange lhsTypes, TypeRange rhsTypes) {
   return lhsTypes.front().isa<IndexType>() != rhsTypes.front().isa<IndexType>();
 }
 
+//===----------------------------------------------------------------------===//
+// CmpOp
+//===----------------------------------------------------------------------===//
+
+/// Compare two integers according to the comparison predicate.
+bool compareIndices(const APInt &lhs, const APInt &rhs,
+                    IndexCmpPredicate pred) {
+  switch (pred) {
+  case IndexCmpPredicate::EQ:
+    return lhs.eq(rhs);
+  case IndexCmpPredicate::NE:
+    return lhs.ne(rhs);
+  case IndexCmpPredicate::SGE:
+    return lhs.sge(rhs);
+  case IndexCmpPredicate::SGT:
+    return lhs.sgt(rhs);
+  case IndexCmpPredicate::SLE:
+    return lhs.sle(rhs);
+  case IndexCmpPredicate::SLT:
+    return lhs.slt(rhs);
+  case IndexCmpPredicate::UGE:
+    return lhs.uge(rhs);
+  case IndexCmpPredicate::UGT:
+    return lhs.ugt(rhs);
+  case IndexCmpPredicate::ULE:
+    return lhs.ule(rhs);
+  case IndexCmpPredicate::ULT:
+    return lhs.ult(rhs);
+  }
+  llvm_unreachable("unhandled IndexCmpPredicate predicate");
+}
+
+OpFoldResult CmpOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.size() == 2 && "compare expected 2 operands");
+  auto lhs = dyn_cast_if_present<IntegerAttr>(operands[0]);
+  auto rhs = dyn_cast_if_present<IntegerAttr>(operands[1]);
+  if (!lhs || !rhs)
+    return {};
+
+  // Perform the comparison in 64-bit and 32-bit.
+  bool result64 = compareIndices(lhs.getValue(), rhs.getValue(), getPred());
+  bool result32 = compareIndices(lhs.getValue().trunc(32),
+                                 rhs.getValue().trunc(32), getPred());
+  if (result64 != result32)
+    return {};
+  return BoolAttr::get(getContext(), result64);
+}
+
+//===----------------------------------------------------------------------===//
+// ConstantOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult ConstantOp::fold(ArrayRef<Attribute> operands) {
+  return getValueAttr();
+}
+
+void ConstantOp::build(OpBuilder &b, OperationState &state, int64_t value) {
+  build(b, state, b.getIndexType(), b.getIndexAttr(value));
+}
+
+//===----------------------------------------------------------------------===//
+// BoolConstantOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult BoolConstantOp::fold(ArrayRef<Attribute> operands) {
+  return getValueAttr();
+}
+
 //===----------------------------------------------------------------------===//
 // ODS-Generated Definitions
 //===----------------------------------------------------------------------===//