[mlir][sparse] Macros to clean up StridedMemRefType in the SparseTensorRuntime

[wrengr]

Summary

In particular, this silences warnings from [-Wsign-compare].

This is a revised version of D137735, which got reverted due to a sign-comparison warning on LLVM's Windows buildbot (which was not on MLIR's Windows buildbot). Differences vs the previous differential:

• vectorToMemref now uses detail::checkOverflowCast to silence the warning that caused the the previous differential to get reverted.
• MEMREF_GET_USIZE now uses detail::checkOverflowCast rather than static_cast
• ASSERT_USIZE_EQ added to abbreviate another common idiom, and to ensure that we use detail::safelyEQ everywhere (to silence a few other warnings)
• A couple for-loops now use index_type for the induction variable, since their upper bound uses that typedef too. (Namely _mlir_ciface_getSparseTensorReaderDimSizes and _mlir_ciface_outSparseTensorWriterNext)

Depends on D138149

[wrengr]

Adding dependency on D138149, and rebasing.

[wrengr]

@stella.stamenova As a heads-up, I'm going to land this now. If you encounter any other issues from LLVM's buildbot, just let me know and I'm more than happy to address them promptly.

[stella.stamenova]

In D137998#3932416, @wrengr wrote:

@stella.stamenova As a heads-up, I'm going to land this now. If you encounter any other issues from LLVM's buildbot, just let me know and I'm more than happy to address them promptly.

One of your changes did break the bot: https://lab.llvm.org/buildbot/#/builders/13/builds/28417. I'm not sure which one.

[wrengr]

Taking a look now

[wrengr]

@stella.stamenova That's due to D138154.

I'm not quite sure what the preprocessor is unhappy about though. The __has_builtin(__builtin_mul_overflow) check is the same as the one used by include/mlir/Analysis/Presburger/MPInt.h and I can't find anywhere that defines __has_builtin if it's missing or anything. I'll post a quick change that does that, just in case that's what's going awry.

[wrengr]

@stella.stamenova I just posted D138167. Can you check to see if that resolves things on your side?

[stella.stamenova]

In D137998#3932524, @wrengr wrote:

@stella.stamenova I just posted D138167. Can you check to see if that resolves things on your side?

I can run a local build to see, but it's already past business hours and I'd hate to leave the buildbot broken overnight pending on a response from me.

[wrengr]

(I'm guessing you're on the east coast then? I'm on the west coast is why I ask)

I can land the patch right away, if you trust that'll fix things. Otherwise, if you need to revert stuff before heading home can you revert only D138154 and not any of the others?

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

 // only visible as an opaque pointer.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/ExecutionEngine/SparseTensorRuntime.h"
 
 #ifdef MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS
 
+#include "mlir/ExecutionEngine/SparseTensor/ArithmeticUtils.h"
 #include "mlir/ExecutionEngine/SparseTensor/COO.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"
 
 #include <cstring>
 #include <numeric>
fromMLIRSparseTensor(const SparseTensorStorage<uint64_t, uint64_t, V> *tensor,
   delete coo;
   *pRank = dimRank;
   *pNse = nse;
   *pShape = shape;
   *pValues = values;
   *pIndices = indices;
 }
 
+//===----------------------------------------------------------------------===//
+//
+// Utilities for manipulating `StridedMemRefType`.
+//
+//===----------------------------------------------------------------------===//
+
+// We shouldn't need to use `detail::safelyEQ` here since the `1` is a literal.
+#define ASSERT_NO_STRIDE(MEMREF)                                               \
+  do {                                                                         \
+    assert((MEMREF) && "Memref is nullptr");                                   \
+    assert(((MEMREF)->strides[0] == 1) && "Memref has non-trivial stride");    \
+  } while (false)
+
+// All our functions use `uint64_t` for ranks, but `StridedMemRefType::sizes`
+// uses `int64_t` on some platforms.  So we explicitly cast this lookup to
+// ensure we get a consistent type, and we use `checkOverflowCast` rather
+// than `static_cast` just to be extremely sure that the casting can't
+// go awry.  (The cast should aways be safe since (1) sizes should never
+// be negative, and (2) the maximum `int64_t` is smaller than the maximum
+// `uint64_t`.  But it's better to be safe than sorry.)
+#define MEMREF_GET_USIZE(MEMREF)                                               \
+  detail::checkOverflowCast<uint64_t>((MEMREF)->sizes[0])
+
+#define ASSERT_USIZE_EQ(MEMREF, SZ)                                            \
+  assert(detail::safelyEQ(MEMREF_GET_USIZE(MEMREF), (SZ)) &&                   \
+         "Memref size mismatch")
+
+#define MEMREF_GET_PAYLOAD(MEMREF) ((MEMREF)->data + (MEMREF)->offset)
+
+// We make this a function rather than a macro mainly for type safety
+// reasons.  This function does not modify the vector, but it cannot
+// be marked `const` because it is stored into the non-`const` memref.
+template <typename T>
+static void vectorToMemref(std::vector<T> &v, StridedMemRefType<T, 1> &ref) {
+  ref.basePtr = ref.data = v.data();
+  ref.offset = 0;
+  using SizeT = typename std::remove_reference_t<decltype(ref.sizes[0])>;
+  ref.sizes[0] = detail::checkOverflowCast<SizeT>(v.size());
+  ref.strides[0] = 1;
+}
+
 } // anonymous namespace
 
 extern "C" {
 
 //===----------------------------------------------------------------------===//
 //
 // Public functions which operate on MLIR buffers (memrefs) to interact
 // with sparse tensors (which are only visible as opaque pointers externally).

 // (3) whether all the arguments are actually used/required.
 void *_mlir_ciface_newSparseTensor( // NOLINT
     StridedMemRefType<index_type, 1> *dimSizesRef,
     StridedMemRefType<index_type, 1> *lvlSizesRef,
     StridedMemRefType<DimLevelType, 1> *lvlTypesRef,
     StridedMemRefType<index_type, 1> *lvl2dimRef,
     StridedMemRefType<index_type, 1> *dim2lvlRef, OverheadType ptrTp,
     OverheadType indTp, PrimaryType valTp, Action action, void *ptr) {
-  assert(dimSizesRef && dimSizesRef->strides[0] == 1);
-  assert(lvlSizesRef && lvlSizesRef->strides[0] == 1);
-  assert(lvlTypesRef && lvlTypesRef->strides[0] == 1);
-  assert(lvl2dimRef && lvl2dimRef->strides[0] == 1);
-  assert(dim2lvlRef && dim2lvlRef->strides[0] == 1);
-  const uint64_t dimRank = dimSizesRef->sizes[0];
-  const uint64_t lvlRank = lvlSizesRef->sizes[0];
-  assert(dim2lvlRef->sizes[0] == (int64_t)dimRank);
-  assert(lvlTypesRef->sizes[0] == (int64_t)lvlRank &&
-         lvl2dimRef->sizes[0] == (int64_t)lvlRank);
-  const index_type *dimSizes = dimSizesRef->data + dimSizesRef->offset;
-  const index_type *lvlSizes = lvlSizesRef->data + lvlSizesRef->offset;
-  const DimLevelType *lvlTypes = lvlTypesRef->data + lvlTypesRef->offset;
-  const index_type *lvl2dim = lvl2dimRef->data + lvl2dimRef->offset;
-  const index_type *dim2lvl = dim2lvlRef->data + dim2lvlRef->offset;
+  ASSERT_NO_STRIDE(dimSizesRef);
+  ASSERT_NO_STRIDE(lvlSizesRef);
+  ASSERT_NO_STRIDE(lvlTypesRef);
+  ASSERT_NO_STRIDE(lvl2dimRef);
+  ASSERT_NO_STRIDE(dim2lvlRef);
+  const uint64_t dimRank = MEMREF_GET_USIZE(dimSizesRef);
+  const uint64_t lvlRank = MEMREF_GET_USIZE(lvlSizesRef);
+  ASSERT_USIZE_EQ(dim2lvlRef, dimRank);
+  ASSERT_USIZE_EQ(lvlTypesRef, lvlRank);
+  ASSERT_USIZE_EQ(lvl2dimRef, lvlRank);
+  const index_type *dimSizes = MEMREF_GET_PAYLOAD(dimSizesRef);
+  const index_type *lvlSizes = MEMREF_GET_PAYLOAD(lvlSizesRef);
+  const DimLevelType *lvlTypes = MEMREF_GET_PAYLOAD(lvlTypesRef);
+  const index_type *lvl2dim = MEMREF_GET_PAYLOAD(lvl2dimRef);
+  const index_type *dim2lvl = MEMREF_GET_PAYLOAD(dim2lvlRef);
 
   // Rewrite kIndex to kU64, to avoid introducing a bunch of new cases.
   // This is safe because of the static_assert above.
   if (ptrTp == OverheadType::kIndex)
     ptrTp = OverheadType::kU64;
   if (indTp == OverheadType::kIndex)
     indTp = OverheadType::kU64;
 

 #undef CASE_SECSAME
 
 #define IMPL_SPARSEVALUES(VNAME, V)                                            \
   void _mlir_ciface_sparseValues##VNAME(StridedMemRefType<V, 1> *ref,          \
                                         void *tensor) {                        \
     assert(ref &&tensor);                                                      \
     std::vector<V> *v;                                                         \
     static_cast<SparseTensorStorageBase *>(tensor)->getValues(&v);             \
-    ref->basePtr = ref->data = v->data();                                      \
-    ref->offset = 0;                                                           \
-    ref->sizes[0] = v->size();                                                 \
-    ref->strides[0] = 1;                                                       \
+    assert(v);                                                                 \
+    vectorToMemref(*v, *ref);                                                  \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_SPARSEVALUES)
 #undef IMPL_SPARSEVALUES
 
 #define IMPL_GETOVERHEAD(NAME, TYPE, LIB)                                      \
   void _mlir_ciface_##NAME(StridedMemRefType<TYPE, 1> *ref, void *tensor,      \
                            index_type d) {                                     \
     assert(ref &&tensor);                                                      \
     std::vector<TYPE> *v;                                                      \
     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v, d);                \
-    ref->basePtr = ref->data = v->data();                                      \
-    ref->offset = 0;                                                           \
-    ref->sizes[0] = v->size();                                                 \
-    ref->strides[0] = 1;                                                       \
+    assert(v);                                                                 \
+    vectorToMemref(*v, *ref);                                                  \
   }
 #define IMPL_SPARSEPOINTERS(PNAME, P)                                          \
   IMPL_GETOVERHEAD(sparsePointers##PNAME, P, getPointers)
 MLIR_SPARSETENSOR_FOREVERY_O(IMPL_SPARSEPOINTERS)
 #undef IMPL_SPARSEPOINTERS
 
 #define IMPL_SPARSEINDICES(INAME, I)                                           \
   IMPL_GETOVERHEAD(sparseIndices##INAME, I, getIndices)
 MLIR_SPARSETENSOR_FOREVERY_O(IMPL_SPARSEINDICES)
 #undef IMPL_SPARSEINDICES
 #undef IMPL_GETOVERHEAD
 
 // TODO: while this API design will work for arbitrary dim2lvl mappings,
 // we should probably move the `dimInd`-to-`lvlInd` computation into codegen
 // (since that could enable optimizations to remove the intermediate memref).
 #define IMPL_ADDELT(VNAME, V)                                                  \
   void *_mlir_ciface_addElt##VNAME(                                            \
       void *lvlCOO, StridedMemRefType<V, 0> *vref,                             \
       StridedMemRefType<index_type, 1> *dimIndRef,                             \
       StridedMemRefType<index_type, 1> *dim2lvlRef) {                          \
-    assert(lvlCOO &&vref &&dimIndRef &&dim2lvlRef);                            \
-    assert(dimIndRef->strides[0] == 1 && dim2lvlRef->strides[0] == 1);         \
-    const uint64_t rank = dimIndRef->sizes[0];                                 \
-    assert(dim2lvlRef->sizes[0] == (int64_t)rank);                             \
-    const index_type *dimInd = dimIndRef->data + dimIndRef->offset;            \
-    const index_type *dim2lvl = dim2lvlRef->data + dim2lvlRef->offset;         \
+    assert(lvlCOO &&vref);                                                     \
+    ASSERT_NO_STRIDE(dimIndRef);                                               \
+    ASSERT_NO_STRIDE(dim2lvlRef);                                              \
+    const uint64_t rank = MEMREF_GET_USIZE(dimIndRef);                         \
+    ASSERT_USIZE_EQ(dim2lvlRef, rank);                                         \
+    const index_type *dimInd = MEMREF_GET_PAYLOAD(dimIndRef);                  \
+    const index_type *dim2lvl = MEMREF_GET_PAYLOAD(dim2lvlRef);                \
     std::vector<index_type> lvlInd(rank);                                      \
     for (uint64_t d = 0; d < rank; ++d)                                        \
       lvlInd[dim2lvl[d]] = dimInd[d];                                          \
-    V *value = vref->data + vref->offset;                                      \
+    V *value = MEMREF_GET_PAYLOAD(vref);                                       \
     static_cast<SparseTensorCOO<V> *>(lvlCOO)->add(lvlInd, *value);            \
     return lvlCOO;                                                             \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_ADDELT)
 #undef IMPL_ADDELT
 
 #define IMPL_GETNEXT(VNAME, V)                                                 \
   bool _mlir_ciface_getNext##VNAME(void *iter,                                 \
                                    StridedMemRefType<index_type, 1> *iref,     \
                                    StridedMemRefType<V, 0> *vref) {            \
-    assert(iter &&iref &&vref);                                                \
-    assert(iref->strides[0] == 1);                                             \
-    index_type *indx = iref->data + iref->offset;                              \
-    V *value = vref->data + vref->offset;                                      \
-    const uint64_t isize = iref->sizes[0];                                     \
+    assert(iter &&vref);                                                       \
+    ASSERT_NO_STRIDE(iref);                                                    \
+    index_type *indx = MEMREF_GET_PAYLOAD(iref);                               \
+    V *value = MEMREF_GET_PAYLOAD(vref);                                       \
+    const uint64_t isize = MEMREF_GET_USIZE(iref);                             \
     const Element<V> *elem =                                                   \
         static_cast<SparseTensorIterator<V> *>(iter)->getNext();               \
     if (elem == nullptr)                                                       \
       return false;                                                            \
     for (uint64_t r = 0; r < isize; r++)                                       \
       indx[r] = elem->indices[r];                                              \
     *value = elem->value;                                                      \
     return true;                                                               \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_GETNEXT)
 #undef IMPL_GETNEXT
 
 #define IMPL_LEXINSERT(VNAME, V)                                               \
   void _mlir_ciface_lexInsert##VNAME(void *tensor,                             \
                                      StridedMemRefType<index_type, 1> *cref,   \
                                      StridedMemRefType<V, 0> *vref) {          \
-    assert(tensor &&cref &&vref);                                              \
-    assert(cref->strides[0] == 1);                                             \
-    index_type *cursor = cref->data + cref->offset;                            \
+    assert(tensor &&vref);                                                     \
+    ASSERT_NO_STRIDE(cref);                                                    \
+    index_type *cursor = MEMREF_GET_PAYLOAD(cref);                             \
     assert(cursor);                                                            \
-    V *value = vref->data + vref->offset;                                      \
+    V *value = MEMREF_GET_PAYLOAD(vref);                                       \
     static_cast<SparseTensorStorageBase *>(tensor)->lexInsert(cursor, *value); \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_LEXINSERT)
 #undef IMPL_LEXINSERT
 
 #define IMPL_EXPINSERT(VNAME, V)                                               \
   void _mlir_ciface_expInsert##VNAME(                                          \
       void *tensor, StridedMemRefType<index_type, 1> *cref,                    \
       StridedMemRefType<V, 1> *vref, StridedMemRefType<bool, 1> *fref,         \
       StridedMemRefType<index_type, 1> *aref, index_type count) {              \
-    assert(tensor &&cref &&vref &&fref &&aref);                                \
-    assert(cref->strides[0] == 1);                                             \
-    assert(vref->strides[0] == 1);                                             \
-    assert(fref->strides[0] == 1);                                             \
-    assert(aref->strides[0] == 1);                                             \
-    assert(vref->sizes[0] == fref->sizes[0]);                                  \
-    index_type *cursor = cref->data + cref->offset;                            \
-    V *values = vref->data + vref->offset;                                     \
-    bool *filled = fref->data + fref->offset;                                  \
-    index_type *added = aref->data + aref->offset;                             \
+    assert(tensor);                                                            \
+    ASSERT_NO_STRIDE(cref);                                                    \
+    ASSERT_NO_STRIDE(vref);                                                    \
+    ASSERT_NO_STRIDE(fref);                                                    \
+    ASSERT_NO_STRIDE(aref);                                                    \
+    ASSERT_USIZE_EQ(vref, MEMREF_GET_USIZE(fref));                             \
+    index_type *cursor = MEMREF_GET_PAYLOAD(cref);                             \
+    V *values = MEMREF_GET_PAYLOAD(vref);                                      \
+    bool *filled = MEMREF_GET_PAYLOAD(fref);                                   \
+    index_type *added = MEMREF_GET_PAYLOAD(aref);                              \
     static_cast<SparseTensorStorageBase *>(tensor)->expInsert(                 \
         cursor, values, filled, added, count);                                 \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_EXPINSERT)
 #undef IMPL_EXPINSERT
 
 void _mlir_ciface_getSparseTensorReaderDimSizes(
     void *p, StridedMemRefType<index_type, 1> *dref) {
-  assert(p && dref);
-  assert(dref->strides[0] == 1);
-  index_type *dimSizes = dref->data + dref->offset;
+  assert(p);
+  ASSERT_NO_STRIDE(dref);
+  index_type *dimSizes = MEMREF_GET_PAYLOAD(dref);
   SparseTensorReader &file = *static_cast<SparseTensorReader *>(p);
   const index_type *sizes = file.getDimSizes();
   index_type rank = file.getRank();
-  for (uint64_t r = 0; r < rank; ++r)
+  for (index_type r = 0; r < rank; ++r)
     dimSizes[r] = sizes[r];
 }
 
 #define IMPL_GETNEXT(VNAME, V)                                                 \
   void _mlir_ciface_getSparseTensorReaderNext##VNAME(                          \
       void *p, StridedMemRefType<index_type, 1> *iref,                         \
       StridedMemRefType<V, 0> *vref) {                                         \
-    assert(p &&iref &&vref);                                                   \
-    assert(iref->strides[0] == 1);                                             \
-    index_type *indices = iref->data + iref->offset;                           \
+    assert(p &&vref);                                                          \
+    ASSERT_NO_STRIDE(iref);                                                    \
+    index_type *indices = MEMREF_GET_PAYLOAD(iref);                            \
     SparseTensorReader *stfile = static_cast<SparseTensorReader *>(p);         \
     index_type rank = stfile->getRank();                                       \
-    V *value = vref->data + vref->offset;                                      \
+    V *value = MEMREF_GET_PAYLOAD(vref);                                       \
     *value = stfile->readCOOElement<V>(rank, indices);                         \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_GETNEXT)
 #undef IMPL_GETNEXT
 
 void _mlir_ciface_outSparseTensorWriterMetaData(
     void *p, index_type rank, index_type nnz,
     StridedMemRefType<index_type, 1> *dref) {
-  assert(p && dref);
-  assert(dref->strides[0] == 1);
+  assert(p);
+  ASSERT_NO_STRIDE(dref);
   assert(rank != 0);
-  index_type *dimSizes = dref->data + dref->offset;
+  index_type *dimSizes = MEMREF_GET_PAYLOAD(dref);
   SparseTensorWriter &file = *static_cast<SparseTensorWriter *>(p);
   file << rank << " " << nnz << std::endl;
   for (index_type r = 0; r < rank - 1; ++r)
     file << dimSizes[r] << " ";
   file << dimSizes[rank - 1] << std::endl;
 }
 
 #define IMPL_OUTNEXT(VNAME, V)                                                 \
   void _mlir_ciface_outSparseTensorWriterNext##VNAME(                          \
       void *p, index_type rank, StridedMemRefType<index_type, 1> *iref,        \
       StridedMemRefType<V, 0> *vref) {                                         \
-    assert(p &&iref &&vref);                                                   \
-    assert(iref->strides[0] == 1);                                             \
-    index_type *indices = iref->data + iref->offset;                           \
+    assert(p &&vref);                                                          \
+    ASSERT_NO_STRIDE(iref);                                                    \
+    index_type *indices = MEMREF_GET_PAYLOAD(iref);                            \
     SparseTensorWriter &file = *static_cast<SparseTensorWriter *>(p);          \
-    for (uint64_t r = 0; r < rank; ++r)                                        \
+    for (index_type r = 0; r < rank; ++r)                                      \
       file << (indices[r] + 1) << " ";                                         \
-    V *value = vref->data + vref->offset;                                      \
+    V *value = MEMREF_GET_PAYLOAD(vref);                                       \
     file << *value << std::endl;                                               \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_OUTNEXT)
 #undef IMPL_OUTNEXT
 
 //===----------------------------------------------------------------------===//
 //
 // Public functions which accept only C-style data structures to interact
void delSparseTensorWriter(void *p) {
   file->flush();
   assert(file->good());
   if (file != &std::cout)
     delete file;
 }
 
 } // extern "C"
 
+#undef MEMREF_GET_PAYLOAD
+#undef ASSERT_USIZE_EQ
+#undef MEMREF_GET_USIZE
+#undef ASSERT_NO_STRIDE
+
 #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS