[mlir][sparse] Adding safe comparison functions to MLIRSparseTensorRuntime.

[wrengr]

Summary

Different platforms use different signedness for StridedMemRefType::sizes and std::vector::size_type, and this has been causing a lot of portability issues re [-Wsign-compare] warnings. These new functions ensure that we need never worry about those signedness warnings ever again.

Also merging CheckedMul.h into ArithmeticUtils.h

[aartbik]

Looks like I got my wish on CheckedMul after all ;-)

mlir/include/mlir/ExecutionEngine/SparseTensor/ArithmeticUtils.h

+//===- ArithmeticUtils.h - Arithmetic helper functions ----------*- 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 header is not part of the public API.  It is placed in the
+// includes directory only because that's required by the implementations
+// of template-classes.
+//
+// This file is part of the lightweight runtime support library for sparse
+// tensor manipulations.  The functionality of the support library is meant
+// to simplify benchmarking, testing, and debugging MLIR code operating on
+// sparse tensors.  However, the provided functionality is **not** part of
+// core MLIR itself.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_EXECUTIONENGINE_SPARSETENSOR_ARITHMETICUTILS_H
+#define MLIR_EXECUTIONENGINE_SPARSETENSOR_ARITHMETICUTILS_H
+
+#include <cassert>
+#include <cinttypes>
+#include <limits>
+
+namespace mlir {
+namespace sparse_tensor {
+namespace detail {
+
+//===----------------------------------------------------------------------===//
+//
+// Safe comparison functions.
+//
+// Variants of the `==`, `!=`, `<`, `<=`, `>`, and `>=` operators which
+// are careful to ensure that negatives are always considered strictly
+// less than non-negatives regardless of the signedness of the types of
+// the two arguments.  They are "safe" in that they guarantee to *always*
+// give an output and that that output is correct; in particular this means
+// they never use assertions or other mechanisms for "returning an error".
+//
+// These functions are C++17-compatible backports of the safe comparison
+// functions added in C++20, and the implementations are based on the
+// sample implementations provided by the standard:
+// <https://en.cppreference.com/w/cpp/utility/intcmp>.
+//
+//===----------------------------------------------------------------------===//
+
+template <typename T, typename U>
+constexpr bool safelyEQ(T t, U u) noexcept {
+  using UT = std::make_unsigned_t<T>;
+  using UU = std::make_unsigned_t<U>;
+  if constexpr (std::is_signed_v<T> == std::is_signed_v<U>)
+    return t == u;
+  else if constexpr (std::is_signed_v<T>)
+    return t < 0 ? false : static_cast<UT>(t) == u;
+  else
+    return u < 0 ? false : t == static_cast<UU>(u);
+}
+
+template <typename T, typename U>
+constexpr bool safelyNE(T t, U u) noexcept {
+  return !safelyEQ(t, u);
+}
+
+template <typename T, typename U>
+constexpr bool safelyLT(T t, U u) noexcept {
+  using UT = std::make_unsigned_t<T>;
+  using UU = std::make_unsigned_t<U>;
+  if constexpr (std::is_signed_v<T> == std::is_signed_v<U>)
+    return t < u;
+  else if constexpr (std::is_signed_v<T>)
+    return t < 0 ? true : static_cast<UT>(t) < u;
+  else
+    return u < 0 ? false : t < static_cast<UU>(u);
+}
+
+template <typename T, typename U>
+constexpr bool safelyGT(T t, U u) noexcept {
+  return safelyLT(u, t);
+}
+
+template <typename T, typename U>
+constexpr bool safelyLE(T t, U u) noexcept {
+  return !safelyGT(t, u);
+}
+
+template <typename T, typename U>
+constexpr bool safelyGE(T t, U u) noexcept {
+  return !safelyLT(t, u);
+}
+
+//===----------------------------------------------------------------------===//
+//
+// Overflow checking functions.
+//
+// These functions use assertions to ensure correctness with respect to
+// overflow/underflow.  Unlike the "safe" functions above, these "checked"
+// functions only guarantee that *if* they return an answer then that answer
+// is correct.  When assertions are enabled, they do their best to remain
+// as fast as possible (since MLIR keeps assertions enabled by default,
+// even for optimized builds).  When assertions are disabled, they use the
+// standard unchecked implementations.
+//
+//===----------------------------------------------------------------------===//
+
+// TODO: we would like to be able to pass in custom error messages, to
+// improve the user experience.  We should be able to use something like
+// `assert(((void)(msg ? msg : defaultMsg), cond))`; but I'm not entirely
+// sure that'll work as intended when done within a function-definition
+// rather than within a macro-definition.
+
+/// A version of `static_cast<To>` which checks for overflow/underflow.
+/// The implementation avoids performing runtime assertions whenever
+/// the types alone are sufficient to statically prove that overflow
+/// cannot happen.
+template <typename To, typename From>
+[[nodiscard]] inline To checkOverflowCast(From x) {
+  // Check the lower bound. (For when casting from signed types.)
+  constexpr To minTo = std::numeric_limits<To>::min();
+  constexpr From minFrom = std::numeric_limits<From>::min();
+  if constexpr (!safelyGE(minFrom, minTo))
+    assert(safelyGE(x, minTo) && "cast would underflow");
+  // Check the upper bound.
+  constexpr To maxTo = std::numeric_limits<To>::max();
+  constexpr From maxFrom = std::numeric_limits<From>::max();
+  if constexpr (!safelyLE(maxFrom, maxTo))
+    assert(safelyLE(x, maxTo) && "cast would overflow");
+  // Now do the cast itself.
+  return static_cast<To>(x);
+}
+
+// TODO: would be better to use various architectures' intrinsics to
+// detect the overflow directly, instead of doing the assertion beforehand
+// (which requires an expensive division).
+//
+/// A version of `operator*` on `uint64_t` which guards against overflows
+/// (when assertions are enabled).
+inline uint64_t checkedMul(uint64_t lhs, uint64_t rhs) {
+  assert((lhs == 0 || rhs <= std::numeric_limits<uint64_t>::max() / lhs) &&
+         "Integer overflow");
+  return lhs * rhs;
+}
+
+} // namespace detail
+} // namespace sparse_tensor
+} // namespace mlir
+
+#endif // MLIR_EXECUTIONENGINE_SPARSETENSOR_ARITHMETICUTILS_H

mlir/include/mlir/ExecutionEngine/SparseTensor/CheckedMul.h

-//===- CheckedMul.h - multiplication that checks for overflow ---*- 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 header is not part of the public API.  It is placed in the
-// includes directory only because that's required by the implementations
-// of template-classes.
-//
-// This file is part of the lightweight runtime support library for sparse
-// tensor manipulations.  The functionality of the support library is meant
-// to simplify benchmarking, testing, and debugging MLIR code operating on
-// sparse tensors.  However, the provided functionality is **not** part of
-// core MLIR itself.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_EXECUTIONENGINE_SPARSETENSOR_CHECKEDMUL_H
-#define MLIR_EXECUTIONENGINE_SPARSETENSOR_CHECKEDMUL_H
-
-#include <cassert>
-#include <cinttypes>
-#include <limits>
-
-namespace mlir {
-namespace sparse_tensor {
-namespace detail {
-
-// TODO: would be better to use various architectures' intrinsics to
-// detect the overflow directly, instead of doing the assertion beforehand
-// (which requires an expensive division).
-//
-/// A version of `operator*` on `uint64_t` which guards against overflows
-/// (when assertions are enabled).
-inline uint64_t checkedMul(uint64_t lhs, uint64_t rhs) {
-  assert((lhs == 0 || rhs <= std::numeric_limits<uint64_t>::max() / lhs) &&
-         "Integer overflow");
-  return lhs * rhs;
-}
-
-} // namespace detail
-} // namespace sparse_tensor
-} // namespace mlir
-
-#endif // MLIR_EXECUTIONENGINE_SPARSETENSOR_CHECKEDMUL_H

mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_EXECUTIONENGINE_SPARSETENSOR_STORAGE_H
 #define MLIR_EXECUTIONENGINE_SPARSETENSOR_STORAGE_H
 
 #include "mlir/Dialect/SparseTensor/IR/Enums.h"
 #include "mlir/ExecutionEngine/Float16bits.h"
+#include "mlir/ExecutionEngine/SparseTensor/ArithmeticUtils.h"
 #include "mlir/ExecutionEngine/SparseTensor/Attributes.h"
 #include "mlir/ExecutionEngine/SparseTensor/COO.h"
-#include "mlir/ExecutionEngine/SparseTensor/CheckedMul.h"
 #include "mlir/ExecutionEngine/SparseTensor/ErrorHandling.h"
 
 namespace mlir {
 namespace sparse_tensor {
 
 //===----------------------------------------------------------------------===//
 // This forward decl is sufficient to split `SparseTensorStorageBase` into
 // its own header, but isn't sufficient for `SparseTensorStorage` to join it.

 
 private:
   /// Appends an arbitrary new position to `pointers[l]`.  This method
   /// checks that `pos` is representable in the `P` type; however, it
   /// does not check that `pos` is semantically valid (i.e., larger than
   /// the previous position and smaller than `indices[l].capacity()`).
   void appendPointer(uint64_t l, uint64_t pos, uint64_t count = 1) {
     ASSERT_COMPRESSED_LVL(l);
-    assert(pos <= std::numeric_limits<P>::max() &&
-           "Pointer value is too large for the P-type");
-    pointers[l].insert(pointers[l].end(), count, static_cast<P>(pos));
+    // TODO: we'd like to recover the nicer error message:
+    // "Pointer value is too large for the P-type"
+    pointers[l].insert(pointers[l].end(), count,
+                       detail::checkOverflowCast<P>(pos));
   }
 
   /// Appends index `i` to level `l`, in the semantically general sense.
   /// For non-dense levels, that means appending to the `indices[l]` array,
   /// checking that `i` is representable in the `I` type; however, we do
   /// not verify other semantic requirements (e.g., that `i` is in bounds
   /// for `lvlSizes[l]`, and not previously occurring in the same segment).
   /// For dense levels, this method instead appends the appropriate number
   /// of zeros to the `values` array, where `full` is the number of "entries"
   /// already written to `values` for this segment (aka one after the highest
   /// index previously appended).
   void appendIndex(uint64_t l, uint64_t full, uint64_t i) {
     const auto dlt = getLvlType(l); // Avoid redundant bounds checking.
     if (isCompressedDLT(dlt) || isSingletonDLT(dlt)) {
-      assert(i <= std::numeric_limits<I>::max() &&
-             "Index value is too large for the I-type");
-      indices[l].push_back(static_cast<I>(i));
+      // TODO: we'd like to recover the nicer error message:
+      // "Index value is too large for the I-type"
+      indices[l].push_back(detail::checkOverflowCast<I>(i));
     } else { // Dense dimension.
       ASSERT_DENSE_DLT(dlt);
       assert(i >= full && "Index was already filled");
       if (i == full)
         return; // Short-circuit, since it'll be a nop.
       if (l + 1 == getLvlRank())
         values.insert(values.end(), i - full, 0);
       else
         finalizeSegment(l + 1, 0, i - full);
     }
   }
 
   /// Writes the given coordinate to `indices[l][pos]`.  This method
   /// checks that `i` is representable in the `I` type; however, it
   /// does not check that `i` is semantically valid (i.e., in bounds
   /// for `dimSizes[l]` and not elsewhere occurring in the same segment).
   void writeIndex(uint64_t l, uint64_t pos, uint64_t i) {
     ASSERT_COMPRESSED_OR_SINGLETON_LVL(l);
     // Subscript assignment to `std::vector` requires that the `pos`-th
     // entry has been initialized; thus we must be sure to check `size()`
     // here, instead of `capacity()` as would be ideal.
     assert(pos < indices[l].size() && "Index position is out of bounds");
-    assert(i <= std::numeric_limits<I>::max() &&
-           "Index value is too large for the I-type");
-    indices[l][pos] = static_cast<I>(i);
+    // TODO: we'd like to recover the nicer error message:
+    // "Index value is too large for the I-type"
+    indices[l][pos] = detail::checkOverflowCast<I>(i);
   }
 
   /// Computes the assembled-size associated with the `l`-th level,
   /// given the assembled-size associated with the `(l-1)`-th level.
   /// "Assembled-sizes" correspond to the (nominal) sizes of overhead
   /// storage, as opposed to "level-sizes" which are the cardinality
   /// of possible coordinates for that level.
   ///

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     name = "SparseTensorRuntime",
     srcs = [
         "lib/ExecutionEngine/SparseTensor/File.cpp",
         "lib/ExecutionEngine/SparseTensor/NNZ.cpp",
         "lib/ExecutionEngine/SparseTensor/PermutationRef.cpp",
         "lib/ExecutionEngine/SparseTensor/Storage.cpp",
     ],
     hdrs = [
+        "include/mlir/ExecutionEngine/SparseTensor/ArithmeticUtils.h",
         "include/mlir/ExecutionEngine/SparseTensor/Attributes.h",
         "include/mlir/ExecutionEngine/SparseTensor/COO.h",
-        "include/mlir/ExecutionEngine/SparseTensor/CheckedMul.h",
         "include/mlir/ExecutionEngine/SparseTensor/ErrorHandling.h",
         "include/mlir/ExecutionEngine/SparseTensor/File.h",
         "include/mlir/ExecutionEngine/SparseTensor/PermutationRef.h",
         "include/mlir/ExecutionEngine/SparseTensor/Storage.h",
     ],
     includes = ["include"],
     deps = [
         ":SparseTensorEnums",