[MLIR][Presburger] introduce MPInt to support fast arbitrary precision in Presburger

[arjunp]

Summary

This uses an int64_t-based fastpath for the common case and falls back to
SlowMPInt to handle the rare cases where larger numbers occur.
It uses __builtin_* for performance through the support in LLVM MathExtras.

Using this in the Presburger library results in a minor performance
*improvement* over any commit hash before sequence of patches
starting at d5e31cf38adfc2c240fb9717989792537cc9e819.

[arjunp]

Fix windows build by removing unnecessary include.

[arjunp]

• remove unused include
• add missing always_inlines
• Fix performance regression due to switching to llvm MathExtras (see top of file)

[arjunp]

Add some more missing always_inlines

[Groverkss]

mlir/include/mlir/Analysis/Presburger/MPInt.h

74–75

significantly impact in what ways? better or worse?

385

Did you plan to fix this before landing this patch?

mlir/unittests/Analysis/Presburger/MPIntTest.cpp

2

nit: please make this the right length.

[arjunp]

• use the terms small/large uniformly
• address kunwar's comments
• remove getAsAP and just use the cast
• simplify mod, floorDiv, ceilDiv implementations
• add assert messages
• deduplicate division overflow checking
• add some more comments

[arjunp]

mlir/include/mlir/Analysis/Presburger/MPInt.h

385

Yes, fixed, thanks

[Groverkss]

mlir/include/mlir/Analysis/Presburger/MPInt.h

111–127

I think these asserts should be more descriptive. For example, "Cannot call getSmall on a large value".

[arjunp]

• add more tests
• deduplicate the tests using googletest features

[arjunp]

Use LLVM_ATTRIBUTE_ALWAYS_INLINE instead of attributes directly.

[arjunp]

fix typo

[arjunp]

• improve assert error

[arjunp]

• format

[Groverkss]

LGTM.

[arjunp]

fix comment

[arjunp]

I accidentally pushed an outdated diff. Revert to diff 441903.

[arjunp]

• really fix comment

[ftynse]

mlir/include/mlir/Analysis/Presburger/MPInt.h

28–33

Have you considered proposing LLVM to add the attribute instead?

[arjunp]

mlir/include/mlir/Analysis/Presburger/MPInt.h

28–33

Yes, will send a patch afterwards.

[clementval]

This is breaking a buildbot. https://lab.llvm.org/buildbot/#/builders/172/builds/14767 Can you have a look or revert?

[arjunp]

Taking a look.

[Meinersbur]

This error is:

C:\PROGRA~1\MICROS~4\2022\COMMUN~1\VC\Tools\MSVC\1431~1.311\bin\Hostx64\x64\cl.exe  /nologo /TP -DGTEST_HAS_RTTI=0 -DMLIR_CUDA_CONVERSIONS_ENABLED=0 -DMLIR_ROCM_CONVERSIONS_ENABLED=0 -DUNICODE -D_CRT_NONSTDC_NO_DEPRECATE -D_CRT_NONSTDC_NO_WARNINGS -D_CRT_SECURE_NO_DEPRECATE -D_CRT_SECURE_NO_WARNINGS -D_HAS_EXCEPTIONS=0 -D_SCL_SECURE_NO_DEPRECATE -D_SCL_SECURE_NO_WARNINGS -D_UNICODE -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -IC:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\build\tools\mlir\lib\Analysis\Presburger -IC:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\lib\Analysis\Presburger -IC:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\build\include -IC:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\llvm\include -IC:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\include -IC:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\build\tools\mlir\include /DWIN32 /D_WINDOWS   /Zc:inline /Zc:__cplusplus /Oi /bigobj /permissive- /W4 -wd4141 -wd4146 -wd4244 -wd4267 -wd4291 -wd4351 -wd4456 -wd4457 -wd4458 -wd4459 -wd4503 -wd4624 -wd4722 -wd4100 -wd4127 -wd4512 -wd4505 -wd4610 -wd4510 -wd4702 -wd4245 -wd4706 -wd4310 -wd4701 -wd4703 -wd4389 -wd4611 -wd4805 -wd4204 -wd4577 -wd4091 -wd4592 -wd4319 -wd4709 -wd4324 -w14062 -we4238 /Gw /MD /O2 /Ob2  /EHs-c- /GR- -UNDEBUG -std:c++17 /showIncludes /Fotools\mlir\lib\Analysis\Presburger\CMakeFiles\obj.MLIRPresburger.dir\MPInt.cpp.obj /Fdtools\mlir\lib\Analysis\Presburger\CMakeFiles\obj.MLIRPresburger.dir\ /FS -c C:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\lib\Analysis\Presburger\MPInt.cpp
C:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\lib\Analysis\Presburger\MPInt.cpp(15): error C2039: 'hash_value': is not a member of 'mlir::presburger'
C:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\include\mlir/Analysis/Presburger/MPInt.h(24): note: see declaration of 'mlir::presburger'
C:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\lib\Analysis\Presburger\MPInt.cpp(16): error C2248: 'mlir::presburger::MPInt::isSmall': cannot access private member declared in class 'mlir::presburger::MPInt'
C:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\include\mlir/Analysis/Presburger/MPInt.h(108): note: see declaration of 'mlir::presburger::MPInt::isSmall'
C:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\include\mlir/Analysis/Presburger/MPInt.h(75): note: see declaration of 'mlir::presburger::MPInt'
C:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\lib\Analysis\Presburger\MPInt.cpp(17): error C2248: 'mlir::presburger::MPInt::valSmall': cannot access private member declared in class 'mlir::presburger::MPInt'
C:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\include\mlir/Analysis/Presburger/MPInt.h(78): note: see declaration of 'mlir::presburger::MPInt::valSmall'
C:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\include\mlir/Analysis/Presburger/MPInt.h(75): note: see declaration of 'mlir::presburger::MPInt'
C:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\lib\Analysis\Presburger\MPInt.cpp(18): error C2248: 'mlir::presburger::MPInt::valLarge': cannot access private member declared in class 'mlir::presburger::MPInt'
C:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\include\mlir/Analysis/Presburger/MPInt.h(79): note: see declaration of 'mlir::presburger::MPInt::valLarge'
C:\Users\buildbot-worker\minipc-ryzen-win\flang-x86_64-windows\llvm-project\mlir\include\mlir/Analysis/Presburger/MPInt.h(75): note: see declaration of 'mlir::presburger::MPInt'

From here:

llvm::hash_code mlir::presburger::hash_value(const MPInt &x) {
  if (x.isSmall())
[...]
}

MPInt::isSmall is indeed private, but hash_value also declared as friend of MPInt. However, friend llvm::hash_code hash_value(const MPInt &x) is declared without namespace, so I am not sure what the compiler is supposed to make it a forward declaration of.

[arjunp]

It should be fixed already: https://lab.llvm.org/buildbot/#/builders/172/builds/14775

The friend declaration becomes part of the innermost enclosing namespace:

Names introduced by friend declarations within a non-local class X become members of the innermost enclosing namespace of X, but they do not become visible to ordinary name lookup (neither unqualified nor qualified) unless a matching declaration is provided at namespace scope, either before or after the class definition.

(from https://en.cppreference.com/w/cpp/language/namespace)

Strangely it works under clang++ with no warning whereas g++ errors out: https://godbolt.org/z/59T35h6cW

[mehdi_amini]

mlir/include/mlir/Analysis/Presburger/MPInt.h

50

Why these "always inline"?

Can these be made static inline instead?

[arjunp]

mlir/include/mlir/Analysis/Presburger/MPInt.h

50

It is important for performance reasons that these actually get inlined. See the comment block at lines 27-32.

mlir/include/mlir/Analysis/Presburger/MPInt.h

+//===- MPInt.h - MLIR MPInt Class -------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This is a simple class to represent arbitrary precision signed integers.
+// Unlike APInt, one does not have to specify a fixed maximum size, and the
+// integer can take on any arbitrary values. This is optimized for small-values
+// by providing fast-paths for the cases when the value stored fits in 64-bits.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_PRESBURGER_MPINT_H
+#define MLIR_ANALYSIS_PRESBURGER_MPINT_H
+
+#include "mlir/Analysis/Presburger/SlowMPInt.h"
+#include "mlir/Support/MathExtras.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace mlir {
+namespace presburger {
+
+namespace detail {
+// If builtin intrinsics for overflow-checked arithmetic are available,
+// use them. Otherwise, call through to LLVM's overflow-checked arithmetic
+// functionality. Those functions also have such macro-gated uses of intrinsics,
+// however they are not always_inlined, which is important for us to achieve
+// high-performance; calling the functions directly would result in a slowdown
+// of 1.15x.
+__attribute__((always_inline)) inline bool addOverflow(int64_t x, int64_t y,
+                                                       int64_t &result) {
+#if __has_builtin(__builtin_add_overflow)
+  return __builtin_add_overflow(x, y, &result);
+#else
+  return llvm::AddOverflow(x, y, result);
+#endif
+}
+__attribute__((always_inline)) inline bool subOverflow(int64_t x, int64_t y,
+                                                       int64_t &result) {
+#if __has_builtin(__builtin_sub_overflow)
+  return __builtin_sub_overflow(x, y, &result);
+#else
+  return llvm::subOverflow(x, y, result);
+#endif
+}
+__attribute__((always_inline)) inline bool mulOverflow(int64_t x, int64_t y,
+                                                       int64_t &result) {
+#if __has_builtin(__builtin_mul_overflow)
+  return __builtin_mul_overflow(x, y, &result);
+#else
+  return llvm::MulOverflow(x, y, result);
+#endif
+}
+} // namespace detail
+
+/// This class provides support for multi-precision arithmetic.
+///
+/// Unlike APInt, this extends the precision as necessary to prevent overflows
+/// and supports operations between objects with differing internal precisions.
+///
+/// This is optimized for small-values by providing fast-paths for the cases
+/// when the value stored fits in 64-bits. We annotate all fastpaths by using
+/// the LLVM_LIKELY/LLVM_UNLIKELY annotations. Removing these would result in
+/// a 1.2x performance slowdown.
+///
+/// We always_inline all operations; removing these results in a 1.5x
+/// performance slowdown.
+///
+/// When holdsLarge is true, a SlowMPInt is held in the union. If it is false,
+/// the int64_t is held. Using std::variant instead would lead to significantly
+/// worse performance.
+class MPInt {
+private:
+  union {
+    int64_t valSmall;
+    detail::SlowMPInt valLarge;
+  };
+  unsigned holdsLarge;
+
+  __attribute__((always_inline)) void initSmall(int64_t o) {
+    if (LLVM_UNLIKELY(isLarge()))
+      valLarge.detail::SlowMPInt::~SlowMPInt();
+    valSmall = o;
+    holdsLarge = false;
+  }
+  __attribute__((always_inline)) void initLarge(const detail::SlowMPInt &o) {
+    if (LLVM_LIKELY(isSmall())) {
+      // The data in memory could be in an arbitrary state, not necessarily
+      // corresponding to any valid state of valLarge; we cannot call any member
+      // functions, e.g. the assignment operator on it, as they may access the
+      // invalid internal state. We instead construct a new object using
+      // placement new.
+      new (&valLarge) detail::SlowMPInt(o);
+    } else {
+      // In this case, we need to use the assignment operator, because if we use
+      // placement-new as above we would lose track of allocated memory
+      // and leak it.
+      valLarge = o;
+    }
+    holdsLarge = true;
+  }
+
+  __attribute__((always_inline)) explicit MPInt(const detail::SlowMPInt &val)
+      : valLarge(val), holdsLarge(true) {}
+  __attribute__((always_inline)) bool isSmall() const { return !holdsLarge; }
+  __attribute__((always_inline)) bool isLarge() const { return holdsLarge; }
+  /// Get the stored value. For getSmall/Large,
+  /// the stored value should be small/large.
+  __attribute__((always_inline)) int64_t getSmall() const {
+    assert(isSmall() && "Wrong accessor called!");
+    return valSmall;
+  }
+  __attribute__((always_inline)) int64_t &getSmall() {
+    assert(isSmall() && "Wrong accessor called!");
+    return valSmall;
+  }
+  __attribute__((always_inline)) const detail::SlowMPInt &getLarge() const {
+    assert(isLarge() && "Wrong accessor called!");
+    return valLarge;
+  }
+  __attribute__((always_inline)) detail::SlowMPInt &getLarge() {
+    assert(isLarge() && "Wrong accessor called!");
+    return valLarge;
+  }
+  explicit operator detail::SlowMPInt() const {
+    if (isSmall())
+      return detail::SlowMPInt(getSmall());
+    return getLarge();
+  }
+
+public:
+  __attribute__((always_inline)) explicit MPInt(int64_t val)
+      : valSmall(val), holdsLarge(false) {}
+  __attribute__((always_inline)) MPInt() : MPInt(0) {}
+  __attribute__((always_inline)) ~MPInt() {
+    if (LLVM_UNLIKELY(isLarge()))
+      valLarge.detail::SlowMPInt::~SlowMPInt();
+  }
+  __attribute__((always_inline)) MPInt(const MPInt &o)
+      : valSmall(o.valSmall), holdsLarge(false) {
+    if (LLVM_UNLIKELY(o.isLarge()))
+      initLarge(o.valLarge);
+  }
+  __attribute__((always_inline)) MPInt &operator=(const MPInt &o) {
+    if (LLVM_LIKELY(o.isSmall())) {
+      initSmall(o.valSmall);
+      return *this;
+    }
+    initLarge(o.valLarge);
+    return *this;
+  }
+  __attribute__((always_inline)) MPInt &operator=(int x) {
+    initSmall(x);
+    return *this;
+  }
+  __attribute__((always_inline)) explicit operator int64_t() const {
+    if (isSmall())
+      return getSmall();
+    return static_cast<int64_t>(getLarge());
+  }
+
+  bool operator==(const MPInt &o) const;
+  bool operator!=(const MPInt &o) const;
+  bool operator>(const MPInt &o) const;
+  bool operator<(const MPInt &o) const;
+  bool operator<=(const MPInt &o) const;
+  bool operator>=(const MPInt &o) const;
+  MPInt operator+(const MPInt &o) const;
+  MPInt operator-(const MPInt &o) const;
+  MPInt operator*(const MPInt &o) const;
+  MPInt operator/(const MPInt &o) const;
+  MPInt operator%(const MPInt &o) const;
+  MPInt &operator+=(const MPInt &o);
+  MPInt &operator-=(const MPInt &o);
+  MPInt &operator*=(const MPInt &o);
+  MPInt &operator/=(const MPInt &o);
+  MPInt &operator%=(const MPInt &o);
+  MPInt operator-() const;
+  MPInt &operator++();
+  MPInt &operator--();
+
+  // Divide by a number that is known to be positive.
+  // This is slightly more efficient because it saves an overflow check.
+  MPInt divByPositive(const MPInt &o) const;
+  MPInt &divByPositiveInPlace(const MPInt &o);
+
+  friend MPInt abs(const MPInt &x);
+  friend MPInt gcdRange(ArrayRef<MPInt> range);
+  friend MPInt ceilDiv(const MPInt &lhs, const MPInt &rhs);
+  friend MPInt floorDiv(const MPInt &lhs, const MPInt &rhs);
+  friend MPInt gcd(const MPInt &a, const MPInt &b);
+  friend MPInt lcm(const MPInt &a, const MPInt &b);
+  friend MPInt mod(const MPInt &lhs, const MPInt &rhs);
+
+  llvm::raw_ostream &print(llvm::raw_ostream &os) const;
+  void dump() const;
+
+  /// ---------------------------------------------------------------------------
+  /// Convenience operator overloads for int64_t.
+  /// ---------------------------------------------------------------------------
+  friend MPInt &operator+=(MPInt &a, int64_t b);
+  friend MPInt &operator-=(MPInt &a, int64_t b);
+  friend MPInt &operator*=(MPInt &a, int64_t b);
+  friend MPInt &operator/=(MPInt &a, int64_t b);
+  friend MPInt &operator%=(MPInt &a, int64_t b);
+
+  friend bool operator==(const MPInt &a, int64_t b);
+  friend bool operator!=(const MPInt &a, int64_t b);
+  friend bool operator>(const MPInt &a, int64_t b);
+  friend bool operator<(const MPInt &a, int64_t b);
+  friend bool operator<=(const MPInt &a, int64_t b);
+  friend bool operator>=(const MPInt &a, int64_t b);
+  friend MPInt operator+(const MPInt &a, int64_t b);
+  friend MPInt operator-(const MPInt &a, int64_t b);
+  friend MPInt operator*(const MPInt &a, int64_t b);
+  friend MPInt operator/(const MPInt &a, int64_t b);
+  friend MPInt operator%(const MPInt &a, int64_t b);
+
+  friend bool operator==(int64_t a, const MPInt &b);
+  friend bool operator!=(int64_t a, const MPInt &b);
+  friend bool operator>(int64_t a, const MPInt &b);
+  friend bool operator<(int64_t a, const MPInt &b);
+  friend bool operator<=(int64_t a, const MPInt &b);
+  friend bool operator>=(int64_t a, const MPInt &b);
+  friend MPInt operator+(int64_t a, const MPInt &b);
+  friend MPInt operator-(int64_t a, const MPInt &b);
+  friend MPInt operator*(int64_t a, const MPInt &b);
+  friend MPInt operator/(int64_t a, const MPInt &b);
+  friend MPInt operator%(int64_t a, const MPInt &b);
+
+  friend llvm::hash_code hash_value(const MPInt &x); // NOLINT
+};
+
+/// This just calls through to the operator int64_t, but it's useful when a
+/// function pointer is required. (Although this is marked inline, it is still
+/// possible to obtain and use a function pointer to this.)
+__attribute__((always_inline)) inline int64_t int64FromMPInt(const MPInt &x) {
+  return int64_t(x);
+}
+
+llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const MPInt &x);
+
+// The RHS is always expected to be positive, and the result
+/// is always non-negative.
+MPInt mod(const MPInt &lhs, const MPInt &rhs);
+
+namespace detail {
+// Division overflows only when trying to negate the minimal signed value.
+__attribute((always_inline)) inline bool divWouldOverflow(int64_t x,
+                                                          int64_t y) {
+  return x == std::numeric_limits<int64_t>::min() && y == -1;
+}
+} // namespace detail
+
+/// We define the operations here in the header to facilitate inlining.
+
+/// ---------------------------------------------------------------------------
+/// Comparison operators.
+/// ---------------------------------------------------------------------------
+__attribute__((always_inline)) inline bool
+MPInt::operator==(const MPInt &o) const {
+  if (LLVM_LIKELY(isSmall() && o.isSmall()))
+    return getSmall() == o.getSmall();
+  return detail::SlowMPInt(*this) == detail::SlowMPInt(o);
+}
+__attribute__((always_inline)) inline bool
+MPInt::operator!=(const MPInt &o) const {
+  if (LLVM_LIKELY(isSmall() && o.isSmall()))
+    return getSmall() != o.getSmall();
+  return detail::SlowMPInt(*this) != detail::SlowMPInt(o);
+}
+__attribute__((always_inline)) inline bool
+MPInt::operator>(const MPInt &o) const {
+  if (LLVM_LIKELY(isSmall() && o.isSmall()))
+    return getSmall() > o.getSmall();
+  return detail::SlowMPInt(*this) > detail::SlowMPInt(o);
+}
+__attribute__((always_inline)) inline bool
+MPInt::operator<(const MPInt &o) const {
+  if (LLVM_LIKELY(isSmall() && o.isSmall()))
+    return getSmall() < o.getSmall();
+  return detail::SlowMPInt(*this) < detail::SlowMPInt(o);
+}
+__attribute__((always_inline)) inline bool
+MPInt::operator<=(const MPInt &o) const {
+  if (LLVM_LIKELY(isSmall() && o.isSmall()))
+    return getSmall() <= o.getSmall();
+  return detail::SlowMPInt(*this) <= detail::SlowMPInt(o);
+}
+__attribute__((always_inline)) inline bool
+MPInt::operator>=(const MPInt &o) const {
+  if (LLVM_LIKELY(isSmall() && o.isSmall()))
+    return getSmall() >= o.getSmall();
+  return detail::SlowMPInt(*this) >= detail::SlowMPInt(o);
+}
+
+/// ---------------------------------------------------------------------------
+/// Arithmetic operators.
+/// ---------------------------------------------------------------------------
+__attribute__((always_inline)) inline MPInt
+MPInt::operator+(const MPInt &o) const {
+  if (LLVM_LIKELY(isSmall() && o.isSmall())) {
+    MPInt result;
+    bool overflow =
+        detail::addOverflow(getSmall(), o.getSmall(), result.getSmall());
+    if (LLVM_LIKELY(!overflow))
+      return result;
+    return MPInt(detail::SlowMPInt(*this) + detail::SlowMPInt(o));
+  }
+  return MPInt(detail::SlowMPInt(*this) + detail::SlowMPInt(o));
+}
+__attribute__((always_inline)) inline MPInt
+MPInt::operator-(const MPInt &o) const {
+  if (LLVM_LIKELY(isSmall() && o.isSmall())) {
+    MPInt result;
+    bool overflow =
+        detail::subOverflow(getSmall(), o.getSmall(), result.getSmall());
+    if (LLVM_LIKELY(!overflow))
+      return result;
+    return MPInt(detail::SlowMPInt(*this) - detail::SlowMPInt(o));
+  }
+  return MPInt(detail::SlowMPInt(*this) - detail::SlowMPInt(o));
+}
+__attribute__((always_inline)) inline MPInt
+MPInt::operator*(const MPInt &o) const {
+  if (LLVM_LIKELY(isSmall() && o.isSmall())) {
+    MPInt result;
+    bool overflow =
+        detail::mulOverflow(getSmall(), o.getSmall(), result.getSmall());
+    if (LLVM_LIKELY(!overflow))
+      return result;
+    return MPInt(detail::SlowMPInt(*this) * detail::SlowMPInt(o));
+  }
+  return MPInt(detail::SlowMPInt(*this) * detail::SlowMPInt(o));
+}
+
+// Division overflows only occur when negating the minimal possible value.
+__attribute__((always_inline)) inline MPInt
+MPInt::divByPositive(const MPInt &o) const {
+  assert(o > 0);
+  if (LLVM_LIKELY(isSmall() && o.isSmall()))
+    return MPInt(getSmall() / o.getSmall());
+  return MPInt(detail::SlowMPInt(*this) / detail::SlowMPInt(o));
+}
+
+__attribute__((always_inline)) inline MPInt
+MPInt::operator/(const MPInt &o) const {
+  if (LLVM_LIKELY(isSmall() && o.isSmall())) {
+    // Division overflows only occur when negating the minimal possible value.
+    if (LLVM_UNLIKELY(detail::divWouldOverflow(getSmall(), o.getSmall())))
+      return -*this;
+    return MPInt(getSmall() / o.getSmall());
+  }
+  return MPInt(detail::SlowMPInt(*this) / detail::SlowMPInt(o));
+}
+
+inline MPInt abs(const MPInt &x) { return MPInt(x >= 0 ? x : -x); }
+// Division overflows only occur when negating the minimal possible value.
+inline MPInt ceilDiv(const MPInt &lhs, const MPInt &rhs) {
+  if (LLVM_LIKELY(lhs.isSmall() && rhs.isSmall())) {
+    if (LLVM_UNLIKELY(detail::divWouldOverflow(lhs.getSmall(), rhs.getSmall())))
+      return -lhs;
+  }
+  return MPInt(ceilDiv(detail::SlowMPInt(lhs), detail::SlowMPInt(rhs)));
+}
+inline MPInt floorDiv(const MPInt &lhs, const MPInt &rhs) {
+  if (LLVM_LIKELY(lhs.isSmall() && rhs.isSmall())) {
+    if (LLVM_UNLIKELY(detail::divWouldOverflow(lhs.getSmall(), rhs.getSmall())))
+      return -lhs;
+    return MPInt(mlir::floorDiv(lhs.getSmall(), rhs.getSmall()));
+  }
+  return MPInt(floorDiv(detail::SlowMPInt(lhs), detail::SlowMPInt(rhs)));
+}
+// The RHS is always expected to be positive, and the result
+/// is always non-negative.
+inline MPInt mod(const MPInt &lhs, const MPInt &rhs) {
+  if (LLVM_LIKELY(lhs.isSmall() && rhs.isSmall()))
+    return MPInt(mlir::mod(lhs.getSmall(), rhs.getSmall()));
+  return MPInt(mod(detail::SlowMPInt(lhs), detail::SlowMPInt(rhs)));
+}
+
+__attribute__((always_inline)) inline MPInt gcd(const MPInt &a,
+                                                const MPInt &b) {
+  if (LLVM_LIKELY(a.isSmall() && b.isSmall()))
+    return MPInt(
+        int64_t(llvm::greatestCommonDivisor(a.getSmall(), b.getSmall())));
+  return MPInt(gcd(detail::SlowMPInt(a), detail::SlowMPInt(b)));
+}
+
+/// Returns the least common multiple of 'a' and 'b'.
+inline MPInt lcm(const MPInt &a, const MPInt &b) {
+  MPInt x = abs(a);
+  MPInt y = abs(b);
+  return (x * y) / gcd(x, y);
+}
+
+/// This operation cannot overflow.
+inline MPInt MPInt::operator%(const MPInt &o) const {
+  if (LLVM_LIKELY(isSmall() && o.isSmall()))
+    return MPInt(getSmall() % o.getSmall());
+  return MPInt(detail::SlowMPInt(*this) % detail::SlowMPInt(o));
+}
+
+inline MPInt MPInt::operator-() const {
+  if (LLVM_LIKELY(isSmall())) {
+    if (LLVM_LIKELY(getSmall() != std::numeric_limits<int64_t>::min()))
+      return MPInt(-getSmall());
+    return MPInt(-detail::SlowMPInt(*this));
+  }
+  return MPInt(-detail::SlowMPInt(*this));
+}
+
+/// ---------------------------------------------------------------------------
+/// Assignment operators, preincrement, predecrement.
+/// ---------------------------------------------------------------------------
+__attribute__((always_inline)) inline MPInt &MPInt::operator+=(const MPInt &o) {
+  if (LLVM_LIKELY(isSmall() && o.isSmall())) {
+    int64_t result = getSmall();
+    bool overflow = detail::addOverflow(getSmall(), o.getSmall(), result);
+    if (LLVM_LIKELY(!overflow)) {
+      getSmall() = result;
+      return *this;
+    }
+    // Note: this return is not strictly required but
+    // removing it leads to a performance regression.
+    return *this = MPInt(detail::SlowMPInt(*this) + detail::SlowMPInt(o));
+  }
+  return *this = MPInt(detail::SlowMPInt(*this) + detail::SlowMPInt(o));
+}
+__attribute__((always_inline)) inline MPInt &MPInt::operator-=(const MPInt &o) {
+  if (LLVM_LIKELY(isSmall() && o.isSmall())) {
+    int64_t result = getSmall();
+    bool overflow = detail::subOverflow(getSmall(), o.getSmall(), result);
+    if (LLVM_LIKELY(!overflow)) {
+      getSmall() = result;
+      return *this;
+    }
+    // Note: this return is not strictly required but
+    // removing it leads to a performance regression.
+    return *this = MPInt(detail::SlowMPInt(*this) - detail::SlowMPInt(o));
+  }
+  return *this = MPInt(detail::SlowMPInt(*this) - detail::SlowMPInt(o));
+}
+__attribute__((always_inline)) inline MPInt &MPInt::operator*=(const MPInt &o) {
+  if (LLVM_LIKELY(isSmall() && o.isSmall())) {
+    int64_t result = getSmall();
+    bool overflow = detail::mulOverflow(getSmall(), o.getSmall(), result);
+    if (LLVM_LIKELY(!overflow)) {
+      getSmall() = result;
+      return *this;
+    }
+    // Note: this return is not strictly required but
+    // removing it leads to a performance regression.
+    return *this = MPInt(detail::SlowMPInt(*this) * detail::SlowMPInt(o));
+  }
+  return *this = MPInt(detail::SlowMPInt(*this) * detail::SlowMPInt(o));
+}
+__attribute__((always_inline)) inline MPInt &MPInt::operator/=(const MPInt &o) {
+  if (LLVM_LIKELY(isSmall() && o.isSmall())) {
+    // Division overflows only occur when negating the minimal possible value.
+    if (LLVM_UNLIKELY(detail::divWouldOverflow(getSmall(), o.getSmall())))
+      return *this = -*this;
+    getSmall() /= o.getSmall();
+    return *this;
+  }
+  return *this = MPInt(detail::SlowMPInt(*this) / detail::SlowMPInt(o));
+}
+
+// Division overflows only occur when the divisor is -1.
+__attribute__((always_inline)) inline MPInt &
+MPInt::divByPositiveInPlace(const MPInt &o) {
+  assert(o > 0);
+  if (LLVM_LIKELY(isSmall() && o.isSmall())) {
+    getSmall() /= o.getSmall();
+    return *this;
+  }
+  return *this = MPInt(detail::SlowMPInt(*this) / detail::SlowMPInt(o));
+}
+
+__attribute__((always_inline)) inline MPInt &MPInt::operator%=(const MPInt &o) {
+  return *this = *this % o;
+}
+__attribute__((always_inline)) inline MPInt &MPInt::operator++() {
+  return *this += 1;
+}
+__attribute__((always_inline)) inline MPInt &MPInt::operator--() {
+  return *this -= 1;
+}
+
+/// ----------------------------------------------------------------------------
+/// Convenience operator overloads for int64_t.
+/// ----------------------------------------------------------------------------
+__attribute__((always_inline)) inline MPInt &operator+=(MPInt &a, int64_t b) {
+  return a = a + b;
+}
+__attribute__((always_inline)) inline MPInt &operator-=(MPInt &a, int64_t b) {
+  return a = a - b;
+}
+__attribute__((always_inline)) inline MPInt &operator*=(MPInt &a, int64_t b) {
+  return a = a * b;
+}
+__attribute__((always_inline)) inline MPInt &operator/=(MPInt &a, int64_t b) {
+  return a = a / b;
+}
+__attribute__((always_inline)) inline MPInt &operator%=(MPInt &a, int64_t b) {
+  return a = a % b;
+}
+inline MPInt operator+(const MPInt &a, int64_t b) { return a + MPInt(b); }
+inline MPInt operator-(const MPInt &a, int64_t b) { return a - MPInt(b); }
+inline MPInt operator*(const MPInt &a, int64_t b) { return a * MPInt(b); }
+inline MPInt operator/(const MPInt &a, int64_t b) { return a / MPInt(b); }
+inline MPInt operator%(const MPInt &a, int64_t b) { return a % MPInt(b); }
+inline MPInt operator+(int64_t a, const MPInt &b) { return MPInt(a) + b; }
+inline MPInt operator-(int64_t a, const MPInt &b) { return MPInt(a) - b; }
+inline MPInt operator*(int64_t a, const MPInt &b) { return MPInt(a) * b; }
+inline MPInt operator/(int64_t a, const MPInt &b) { return MPInt(a) / b; }
+inline MPInt operator%(int64_t a, const MPInt &b) { return MPInt(a) % b; }
+
+/// We provide special implementations of the comparison operators rather than
+/// calling through as above, as this would result in a 1.2x slowdown.
+inline bool operator==(const MPInt &a, int64_t b) {
+  if (LLVM_LIKELY(a.isSmall()))
+    return a.getSmall() == b;
+  return a.getLarge() == b;
+}
+inline bool operator!=(const MPInt &a, int64_t b) {
+  if (LLVM_LIKELY(a.isSmall()))
+    return a.getSmall() != b;
+  return a.getLarge() != b;
+}
+inline bool operator>(const MPInt &a, int64_t b) {
+  if (LLVM_LIKELY(a.isSmall()))
+    return a.getSmall() > b;
+  return a.getLarge() > b;
+}
+inline bool operator<(const MPInt &a, int64_t b) {
+  if (LLVM_LIKELY(a.isSmall()))
+    return a.getSmall() < b;
+  return a.getLarge() < b;
+}
+inline bool operator<=(const MPInt &a, int64_t b) {
+  if (LLVM_LIKELY(a.isSmall()))
+    return a.getSmall() <= b;
+  return a.getLarge() <= b;
+}
+inline bool operator>=(const MPInt &a, int64_t b) {
+  if (LLVM_LIKELY(a.isSmall()))
+    return a.getSmall() >= b;
+  return a.getLarge() >= b;
+}
+inline bool operator==(int64_t a, const MPInt &b) {
+  if (LLVM_LIKELY(b.isSmall()))
+    return a == b.getSmall();
+  return a == b.getLarge();
+}
+inline bool operator!=(int64_t a, const MPInt &b) {
+  if (LLVM_LIKELY(b.isSmall()))
+    return a != b.getSmall();
+  return a != b.getLarge();
+}
+inline bool operator>(int64_t a, const MPInt &b) {
+  if (LLVM_LIKELY(b.isSmall()))
+    return a > b.getSmall();
+  return a > b.getLarge();
+}
+inline bool operator<(int64_t a, const MPInt &b) {
+  if (LLVM_LIKELY(b.isSmall()))
+    return a < b.getSmall();
+  return a < b.getLarge();
+}
+inline bool operator<=(int64_t a, const MPInt &b) {
+  if (LLVM_LIKELY(b.isSmall()))
+    return a <= b.getSmall();
+  return a <= b.getLarge();
+}
+inline bool operator>=(int64_t a, const MPInt &b) {
+  if (LLVM_LIKELY(b.isSmall()))
+    return a >= b.getSmall();
+  return a >= b.getLarge();
+}
+
+} // namespace presburger
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_PRESBURGER_MPINT_H

mlir/lib/Analysis/Presburger/CMakeLists.txt

 add_mlir_library(MLIRPresburger
   IntegerRelation.cpp
   LinearTransform.cpp
   Matrix.cpp
+  MPInt.cpp
   PresburgerRelation.cpp
   PresburgerSpace.cpp
   PWMAFunction.cpp
   Simplex.cpp
   SlowMPInt.cpp
   Utils.cpp
 
   LINK_LIBS PUBLIC
   MLIRSupport
   )

mlir/lib/Analysis/Presburger/MPInt.cpp

+//===- MPInt.cpp - MLIR MPInt Class ---------------------------------------===//
+//
+// 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/Analysis/Presburger/MPInt.h"
+#include "llvm/Support/MathExtras.h"
+
+using namespace mlir;
+using namespace presburger;
+
+llvm::hash_code mlir::presburger::hash_value(const MPInt &x) {
+  if (x.isSmall())
+    return llvm::hash_value(x.valSmall);
+  return detail::hash_value(x.valLarge);
+}
+
+/// ---------------------------------------------------------------------------
+/// Printing.
+/// ---------------------------------------------------------------------------
+llvm::raw_ostream &MPInt::print(llvm::raw_ostream &os) const {
+  if (isSmall())
+    return os << valSmall;
+  return os << valLarge;
+}
+
+void MPInt::dump() const { print(llvm::errs()); }
+
+llvm::raw_ostream &mlir::presburger::operator<<(llvm::raw_ostream &os,
+                                                const MPInt &x) {
+  x.print(os);
+  return os;
+}

mlir/unittests/Analysis/Presburger/CMakeLists.txt

 add_mlir_unittest(MLIRPresburgerTests
   IntegerPolyhedronTest.cpp
   IntegerRelationTest.cpp
   LinearTransformTest.cpp
   MatrixTest.cpp
+  MPIntTest.cpp
   PresburgerSetTest.cpp
   PresburgerSpaceTest.cpp
   PWMAFunctionTest.cpp
   SimplexTest.cpp
   SlowMPIntTest.cpp
   ../../Dialect/Affine/Analysis/AffineStructuresParser.cpp
 )
 
 target_link_libraries(MLIRPresburgerTests
   PRIVATE MLIRPresburger
   MLIRAffineAnalysis
   MLIRParser
   )

mlir/unittests/Analysis/Presburger/MPIntTest.cpp

+//===- MPIntTest.cpp - Tests for MPInt ------------------------------------===//
+//
+// 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/Analysis/Presburger/MPInt.h"
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+using namespace mlir;
+using namespace presburger;
+
+TEST(MPIntTest, ops) {
+  MPInt two(2), five(5), seven(7), ten(10);
+  EXPECT_EQ(five + five, ten);
+  EXPECT_EQ(five * five, 2 * ten + five);
+  EXPECT_EQ(five * five, 3 * ten - five);
+  EXPECT_EQ(five * two, ten);
+  EXPECT_EQ(five / two, two);
+  EXPECT_EQ(five % two, two / two);
+
+  EXPECT_EQ(-ten % seven, -10 % 7);
+  EXPECT_EQ(ten % -seven, 10 % -7);
+  EXPECT_EQ(-ten % -seven, -10 % -7);
+  EXPECT_EQ(ten % seven, 10 % 7);
+
+  EXPECT_EQ(-ten / seven, -10 / 7);
+  EXPECT_EQ(ten / -seven, 10 / -7);
+  EXPECT_EQ(-ten / -seven, -10 / -7);
+  EXPECT_EQ(ten / seven, 10 / 7);
+
+  MPInt x = ten;
+  x += five;
+  EXPECT_EQ(x, 15);
+  x *= two;
+  EXPECT_EQ(x, 30);
+  x /= seven;
+  EXPECT_EQ(x, 4);
+  x -= two * 10;
+  EXPECT_EQ(x, -16);
+  x *= 2 * two;
+  EXPECT_EQ(x, -64);
+  x /= two / -2;
+  EXPECT_EQ(x, 64);
+
+  EXPECT_LE(ten, ten);
+  EXPECT_GE(ten, ten);
+  EXPECT_EQ(ten, ten);
+  EXPECT_FALSE(ten != ten);
+  EXPECT_FALSE(ten < ten);
+  EXPECT_FALSE(ten > ten);
+  EXPECT_LT(five, ten);
+  EXPECT_GT(ten, five);
+}
+
+TEST(MPIntTest, ops64Overloads) {
+  MPInt two(2), five(5), seven(7), ten(10);
+  EXPECT_EQ(five + 5, ten);
+  EXPECT_EQ(five + 5, 5 + five);
+  EXPECT_EQ(five * 5, 2 * ten + 5);
+  EXPECT_EQ(five * 5, 3 * ten - 5);
+  EXPECT_EQ(five * two, ten);
+  EXPECT_EQ(5 / two, 2);
+  EXPECT_EQ(five / 2, 2);
+  EXPECT_EQ(2 % two, 0);
+  EXPECT_EQ(2 - two, 0);
+  EXPECT_EQ(2 % two, two % 2);
+
+  MPInt x = ten;
+  x += 5;
+  EXPECT_EQ(x, 15);
+  x *= 2;
+  EXPECT_EQ(x, 30);
+  x /= 7;
+  EXPECT_EQ(x, 4);
+  x -= 20;
+  EXPECT_EQ(x, -16);
+  x *= 4;
+  EXPECT_EQ(x, -64);
+  x /= -1;
+  EXPECT_EQ(x, 64);
+
+  EXPECT_LE(ten, 10);
+  EXPECT_GE(ten, 10);
+  EXPECT_EQ(ten, 10);
+  EXPECT_FALSE(ten != 10);
+  EXPECT_FALSE(ten < 10);
+  EXPECT_FALSE(ten > 10);
+  EXPECT_LT(five, 10);
+  EXPECT_GT(ten, 5);
+
+  EXPECT_LE(10, ten);
+  EXPECT_GE(10, ten);
+  EXPECT_EQ(10, ten);
+  EXPECT_FALSE(10 != ten);
+  EXPECT_FALSE(10 < ten);
+  EXPECT_FALSE(10 > ten);
+  EXPECT_LT(5, ten);
+  EXPECT_GT(10, five);
+}
+
+TEST(MPIntTest, overflows) {
+  MPInt x(1ll << 60);
+  EXPECT_EQ((x * x - x * x * x * x) / (x * x * x), 1 - (1ll << 60));
+  MPInt y(1ll << 62);
+  EXPECT_EQ((y + y + y + y + y + y) / y, 6);
+  EXPECT_EQ(-(2 * (-y)), 2 * y); // -(-2^63) overflow.
+}