[mlir][sparse] Factored out a "FATAL" macro for unrecoverable assertion failure

[wrengr]

Summary

Depends On D126019

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

350–373

can we use one of your nifty looping macros for this too (in follow up CL)?

[wrengr]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

350–373

Yep, in my workspace I add some x-macros for overhead types to clean this up; I'm just trying to break that mega-CL into more manageable chunks :)

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

 
 /// A version of `operator*` on `uint64_t` which checks for overflows.
 static 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;
 }
 
+// This macro helps minimize repetition of this idiom, as well as ensuring
+// we have some additional output indicating where the error is coming from.
+// (Since `fprintf` doesn't provide a stacktrace, this helps make it easier
+// to track down whether an error is coming from our code vs somewhere else
+// in MLIR.)
+#define FATAL(...)                                                             \
+  {                                                                            \
+    fprintf(stderr, "SparseTensorUtils: " __VA_ARGS__);                        \
+    exit(1);                                                                   \
+  }
+
 // TODO: adjust this so it can be used by `openSparseTensorCOO` too.
 // That version doesn't have the permutation, and the `dimSizes` are
 // a pointer/C-array rather than `std::vector`.
 //
 /// Asserts that the `dimSizes` (in target-order) under the `perm` (mapping
 /// semantic-order to target-order) are a refinement of the desired `shape`
 /// (in semantic-order).
 ///
#define FOREVERY_V(DO)                                                         \
   FOREVERY_SIMPLEX_V(DO)                                                       \
   DO(C64, complex64)                                                           \
   DO(C32, complex32)
 
 // Forward.
 template <typename V>
 class SparseTensorEnumeratorBase;
 
+// Helper macro for generating error messages when some
+// `SparseTensorStorage<P,I,V>` is cast to `SparseTensorStorageBase`
+// and then the wrong "partial method specialization" is called.
+#define FATAL_PIV(NAME) FATAL("<P,I,V> type mismatch for: " #NAME);
+
 /// Abstract base class for `SparseTensorStorage<P,I,V>`.  This class
 /// takes responsibility for all the `<P,I,V>`-independent aspects
 /// of the tensor (e.g., shape, sparsity, permutation).  In addition,
 /// we use function overloading to implement "partial" method
 /// specialization, which the C-API relies on to catch type errors
 /// arising from our use of opaque pointers.
 class SparseTensorStorageBase {
 public:
bool isCompressedDim(uint64_t d) const {
     assert(d < getRank());
     return (dimTypes[d] == DimLevelType::kCompressed);
   }
 
   /// Allocate a new enumerator.
 #define DECL_NEWENUMERATOR(VNAME, V)                                           \
   virtual void newEnumerator(SparseTensorEnumeratorBase<V> **, uint64_t,       \
                              const uint64_t *) const {                         \
-    fatal("newEnumerator" #VNAME);                                             \
+    FATAL_PIV("newEnumerator" #VNAME);                                         \
   }
   FOREVERY_V(DECL_NEWENUMERATOR)
 #undef DECL_NEWENUMERATOR
 
   /// Overhead storage.
-  virtual void getPointers(std::vector<uint64_t> **, uint64_t) { fatal("p64"); }
-  virtual void getPointers(std::vector<uint32_t> **, uint64_t) { fatal("p32"); }
-  virtual void getPointers(std::vector<uint16_t> **, uint64_t) { fatal("p16"); }
-  virtual void getPointers(std::vector<uint8_t> **, uint64_t) { fatal("p8"); }
-  virtual void getIndices(std::vector<uint64_t> **, uint64_t) { fatal("i64"); }
-  virtual void getIndices(std::vector<uint32_t> **, uint64_t) { fatal("i32"); }
-  virtual void getIndices(std::vector<uint16_t> **, uint64_t) { fatal("i16"); }
-  virtual void getIndices(std::vector<uint8_t> **, uint64_t) { fatal("i8"); }
+  virtual void getPointers(std::vector<uint64_t> **, uint64_t) {
+    FATAL_PIV("p64");
+  }
+  virtual void getPointers(std::vector<uint32_t> **, uint64_t) {
+    FATAL_PIV("p32");
+  }
+  virtual void getPointers(std::vector<uint16_t> **, uint64_t) {
+    FATAL_PIV("p16");
+  }
+  virtual void getPointers(std::vector<uint8_t> **, uint64_t) {
+    FATAL_PIV("p8");
+  }
+  virtual void getIndices(std::vector<uint64_t> **, uint64_t) {
+    FATAL_PIV("i64");
+  }
+  virtual void getIndices(std::vector<uint32_t> **, uint64_t) {
+    FATAL_PIV("i32");
+  }
+  virtual void getIndices(std::vector<uint16_t> **, uint64_t) {
+    FATAL_PIV("i16");
+  }
+  virtual void getIndices(std::vector<uint8_t> **, uint64_t) {
+    FATAL_PIV("i8");
+  }
 
   /// Primary storage.
 #define DECL_GETVALUES(VNAME, V)                                               \
-  virtual void getValues(std::vector<V> **) { fatal("getValues" #VNAME); }
+  virtual void getValues(std::vector<V> **) { FATAL_PIV("getValues" #VNAME); }
   FOREVERY_V(DECL_GETVALUES)
 #undef DECL_GETVALUES
 
   /// Element-wise insertion in lexicographic index order.
 #define DECL_LEXINSERT(VNAME, V)                                               \
-  virtual void lexInsert(const uint64_t *, V) { fatal("lexInsert" #VNAME); }
+  virtual void lexInsert(const uint64_t *, V) { FATAL_PIV("lexInsert" #VNAME); }
   FOREVERY_V(DECL_LEXINSERT)
 #undef DECL_LEXINSERT
 
   /// Expanded insertion.
 #define DECL_EXPINSERT(VNAME, V)                                               \
   virtual void expInsert(uint64_t *, V *, bool *, uint64_t *, uint64_t) {      \
-    fatal("expInsert" #VNAME);                                                 \
+    FATAL_PIV("expInsert" #VNAME);                                             \
   }
   FOREVERY_V(DECL_EXPINSERT)
 #undef DECL_EXPINSERT
 
   /// Finishes insertion.
   virtual void endInsert() = 0;
 
 protected:
   // Since this class is virtual, we must disallow public copying in
   // order to avoid "slicing".  Since this class has data members,
   // that means making copying protected.
   // <https://github.com/isocpp/CppCoreGuidelines/blob/master/CppCoreGuidelines.md#Rc-copy-virtual>
   SparseTensorStorageBase(const SparseTensorStorageBase &) = default;
   // Copy-assignment would be implicitly deleted (because `dimSizes`
   // is const), so we explicitly delete it for clarity.
   SparseTensorStorageBase &operator=(const SparseTensorStorageBase &) = delete;
 
 private:
-  static void fatal(const char *tp) {
-    fprintf(stderr, "unsupported %s\n", tp);
-    exit(1);
-  }
-
   const std::vector<uint64_t> dimSizes;
   std::vector<uint64_t> rev;
   const std::vector<DimLevelType> dimTypes;
 };
 
+#undef FATAL_PIV
+
 // Forward.
 template <typename P, typename I, typename V>
 class SparseTensorEnumerator;
 
 /// A memory-resident sparse tensor using a storage scheme based on
 /// per-dimension sparse/dense annotations. This data structure provides a
 /// bufferized form of a sparse tensor type. In contrast to generating setup
 /// methods for each differently annotated sparse tensor, this method provides
static void readMMEHeader(FILE *file, char *filename, char *line,
                           uint64_t *idata, bool *isPattern, bool *isSymmetric) {
   char header[64];
   char object[64];
   char format[64];
   char field[64];
   char symmetry[64];
   // Read header line.
   if (fscanf(file, "%63s %63s %63s %63s %63s\n", header, object, format, field,
-             symmetry) != 5) {
-    fprintf(stderr, "Corrupt header in %s\n", filename);
-    exit(1);
-  }
+             symmetry) != 5)
+    FATAL("Corrupt header in %s\n", filename);
   // Set properties
   *isPattern = (strcmp(toLower(field), "pattern") == 0);
   *isSymmetric = (strcmp(toLower(symmetry), "symmetric") == 0);
   // Make sure this is a general sparse matrix.
   if (strcmp(toLower(header), "%%matrixmarket") ||
       strcmp(toLower(object), "matrix") ||
       strcmp(toLower(format), "coordinate") ||
       (strcmp(toLower(field), "real") && !(*isPattern)) ||
-      (strcmp(toLower(symmetry), "general") && !(*isSymmetric))) {
-    fprintf(stderr, "Cannot find a general sparse matrix in %s\n", filename);
-    exit(1);
-  }
+      (strcmp(toLower(symmetry), "general") && !(*isSymmetric)))
+    FATAL("Cannot find a general sparse matrix in %s\n", filename);
   // Skip comments.
   while (true) {
-    if (!fgets(line, kColWidth, file)) {
-      fprintf(stderr, "Cannot find data in %s\n", filename);
-      exit(1);
-    }
+    if (!fgets(line, kColWidth, file))
+      FATAL("Cannot find data in %s\n", filename);
     if (line[0] != '%')
       break;
   }
   // Next line contains M N NNZ.
   idata[0] = 2; // rank
   if (sscanf(line, "%" PRIu64 "%" PRIu64 "%" PRIu64 "\n", idata + 2, idata + 3,
-             idata + 1) != 3) {
-    fprintf(stderr, "Cannot find size in %s\n", filename);
-    exit(1);
-  }
+             idata + 1) != 3)
+    FATAL("Cannot find size in %s\n", filename);
 }
 
 /// Read the "extended" FROSTT header. Although not part of the documented
 /// format, we assume that the file starts with optional comments followed
 /// by two lines that define the rank, the number of nonzeros, and the
 /// dimensions sizes (one per rank) of the sparse tensor.
 static void readExtFROSTTHeader(FILE *file, char *filename, char *line,
                                 uint64_t *idata) {
   // Skip comments.
   while (true) {
-    if (!fgets(line, kColWidth, file)) {
-      fprintf(stderr, "Cannot find data in %s\n", filename);
-      exit(1);
-    }
+    if (!fgets(line, kColWidth, file))
+      FATAL("Cannot find data in %s\n", filename);
     if (line[0] != '#')
       break;
   }
   // Next line contains RANK and NNZ.
-  if (sscanf(line, "%" PRIu64 "%" PRIu64 "\n", idata, idata + 1) != 2) {
-    fprintf(stderr, "Cannot find metadata in %s\n", filename);
-    exit(1);
-  }
+  if (sscanf(line, "%" PRIu64 "%" PRIu64 "\n", idata, idata + 1) != 2)
+    FATAL("Cannot find metadata in %s\n", filename);
   // Followed by a line with the dimension sizes (one per rank).
-  for (uint64_t r = 0; r < idata[0]; r++) {
-    if (fscanf(file, "%" PRIu64, idata + 2 + r) != 1) {
-      fprintf(stderr, "Cannot find dimension size %s\n", filename);
-      exit(1);
-    }
-  }
+  for (uint64_t r = 0; r < idata[0]; r++)
+    if (fscanf(file, "%" PRIu64, idata + 2 + r) != 1)
+      FATAL("Cannot find dimension size %s\n", filename);
   fgets(line, kColWidth, file); // end of line
 }
 
 /// Reads a sparse tensor with the given filename into a memory-resident
 /// sparse tensor in coordinate scheme.
 template <typename V>
 static SparseTensorCOO<V> *openSparseTensorCOO(char *filename, uint64_t rank,
                                                const uint64_t *shape,
                                                const uint64_t *perm) {
   // Open the file.
-  FILE *file = fopen(filename, "r");
-  if (!file) {
-    assert(filename && "Received nullptr for filename");
-    fprintf(stderr, "Cannot find file %s\n", filename);
-    exit(1);
-  }
+  FILE *file = fopen(filename, "r");
+  if (!file)
+    FATAL("Cannot find file %s\n", filename);
   // Perform some file format dependent set up.
   char line[kColWidth];
   uint64_t idata[512];
   bool isPattern = false;
   bool isSymmetric = false;
   if (strstr(filename, ".mtx")) {
     readMMEHeader(file, filename, line, idata, &isPattern, &isSymmetric);
   } else if (strstr(filename, ".tns")) {
     readExtFROSTTHeader(file, filename, line, idata);
   } else {
-    fprintf(stderr, "Unknown format %s\n", filename);
-    exit(1);
+    FATAL("Unknown format %s\n", filename);
   }
   // Prepare sparse tensor object with per-dimension sizes
   // and the number of nonzeros as initial capacity.
   assert(rank == idata[0] && "rank mismatch");
   uint64_t nnz = idata[1];
   for (uint64_t r = 0; r < rank; r++)
     assert((shape[r] == 0 || shape[r] == idata[2 + r]) &&
            "dimension size mismatch");
   SparseTensorCOO<V> *tensor =
       SparseTensorCOO<V>::newSparseTensorCOO(rank, idata + 2, perm, nnz);
   // Read all nonzero elements.
   std::vector<uint64_t> indices(rank);
   for (uint64_t k = 0; k < nnz; k++) {
-    if (!fgets(line, kColWidth, file)) {
-      fprintf(stderr, "Cannot find next line of data in %s\n", filename);
-      exit(1);
-    }
+    if (!fgets(line, kColWidth, file))
+      FATAL("Cannot find next line of data in %s\n", filename);
     char *linePtr = line;
     for (uint64_t r = 0; r < rank; r++) {
       uint64_t idx = strtoul(linePtr, &linePtr, 10);
       // Add 0-based index.
       indices[perm[r]] = idx - 1;
     }
     // The external formats always store the numerical values with the type
     // double, but we cast these values to the sparse tensor object type.

 static SparseTensorStorage<uint64_t, uint64_t, V> *
 toMLIRSparseTensor(uint64_t rank, uint64_t nse, uint64_t *shape, V *values,
                    uint64_t *indices, uint64_t *perm, uint8_t *sparse) {
   const DimLevelType *sparsity = (DimLevelType *)(sparse);
 #ifndef NDEBUG
   // Verify that perm is a permutation of 0..(rank-1).
   std::vector<uint64_t> order(perm, perm + rank);
   std::sort(order.begin(), order.end());
-  for (uint64_t i = 0; i < rank; ++i) {
-    if (i != order[i]) {
-      fprintf(stderr, "Not a permutation of 0..%" PRIu64 "\n", rank);
-      exit(1);
-    }
-  }
+  for (uint64_t i = 0; i < rank; ++i)
+    if (i != order[i])
+      FATAL("Not a permutation of 0..%" PRIu64 "\n", rank);
 
   // Verify that the sparsity values are supported.
-  for (uint64_t i = 0; i < rank; ++i) {
+  for (uint64_t i = 0; i < rank; ++i)
     if (sparsity[i] != DimLevelType::kDense &&
-        sparsity[i] != DimLevelType::kCompressed) {
-      fprintf(stderr, "Unsupported sparsity value %d\n",
-              static_cast<int>(sparsity[i]));
-      exit(1);
-    }
-  }
+        sparsity[i] != DimLevelType::kCompressed)
+      FATAL("Unsupported sparsity value %d\n", static_cast<int>(sparsity[i]));
 #endif
 
   // Convert external format to internal COO.
   auto *coo = SparseTensorCOO<V>::newSparseTensorCOO(rank, shape, perm, nse);
   std::vector<uint64_t> idx(rank);
   for (uint64_t i = 0, base = 0; i < nse; i++) {
     for (uint64_t r = 0; r < rank; r++)
       idx[perm[r]] = indices[base + r];
_mlir_ciface_newSparseTensor(StridedMemRefType<DimLevelType, 1> *aref, // NOLINT
   CASE_SECSAME(OverheadType::kU8, PrimaryType::kI16, uint8_t, int16_t);
   CASE_SECSAME(OverheadType::kU8, PrimaryType::kI8, uint8_t, int8_t);
 
   // Complex matrices with wide overhead.
   CASE_SECSAME(OverheadType::kU64, PrimaryType::kC64, uint64_t, complex64);
   CASE_SECSAME(OverheadType::kU64, PrimaryType::kC32, uint64_t, complex32);
 
   // Unsupported case (add above if needed).
-  fputs("unsupported combination of types\n", stderr);
-  exit(1);
+  // TODO: better pretty-printing of enum values!
+  FATAL("unsupported combination of types: <P=%d, I=%d, V=%d>\n",
+        static_cast<int>(ptrTp), static_cast<int>(indTp),
+        static_cast<int>(valTp));
 }
 #undef CASE
 #undef CASE_SECSAME
 
 /// Methods that provide direct access to values.
 #define IMPL_SPARSEVALUES(VNAME, V)                                            \
   void _mlir_ciface_sparseValues##VNAME(StridedMemRefType<V, 1> *ref,          \
                                         void *tensor) {                        \

 //===----------------------------------------------------------------------===//
 
 /// Helper method to read a sparse tensor filename from the environment,
 /// defined with the naming convention ${TENSOR0}, ${TENSOR1}, etc.
 char *getTensorFilename(index_type id) {
   char var[80];
   sprintf(var, "TENSOR%" PRIu64, id);
   char *env = getenv(var);
-  if (!env) {
-    fprintf(stderr, "Environment variable %s is not set\n", var);
-    exit(1);
-  }
+  if (!env)
+    FATAL("Environment variable %s is not set\n", var);
   return env;
 }
 
 /// Returns size of sparse tensor in given dimension.
 index_type sparseDimSize(void *tensor, index_type d) {
   return static_cast<SparseTensorStorageBase *>(tensor)->getDimSize(d);
 }