[mlir][sparse] Factoring out SparseTensorFile::canReadAs predicate

[wrengr]

Summary

This is a followup to the refactoring of D133462, D133830, D133831, and D133833.

Depends On D133833

[wrengr]

rebase

[wrengr]

rebase

[stella.stamenova]

Looks like one of your changes, most likely this one, broke the windows mlir bot: https://lab.llvm.org/buildbot/#/builders/13/builds/26434

[wrengr]

I just posted https://reviews.llvm.org/D134925 for review, which should fix the LLVM Windows buildbot

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

public:
   /// Reads and parses the file's header.
   void readHeader();
 
   ValueKind getValueKind() const { return valueKind_; }
 
   /// Checks if a header has been successfully read.
   bool isValid() const { return valueKind_ != ValueKind::kInvalid; }
 
+  /// Checks if the file's ValueKind can be converted into the given
+  /// tensor PrimaryType.  Is only valid after parsing the header.
+  bool canReadAs(PrimaryType valTy) const;
+
   /// Gets the MME "pattern" property setting.  Is only valid after
   /// parsing the header.
   bool isPattern() const {
     assert(isValid() && "Attempt to isPattern() before readHeader()");
     return valueKind_ == ValueKind::kPattern;
   }
 
   /// Gets the MME "symmetric" property setting.  Is only valid after

 template <typename V>
 inline SparseTensorCOO<V> *
 openSparseTensorCOO(const char *filename, uint64_t rank, const uint64_t *shape,
                     const uint64_t *perm, PrimaryType valTp) {
   SparseTensorFile stfile(filename);
   stfile.openFile();
   stfile.readHeader();
   // Check tensor element type against the value type in the input file.
-  SparseTensorFile::ValueKind valueKind = stfile.getValueKind();
-  bool tensorIsInteger =
-      (valTp >= PrimaryType::kI64 && valTp <= PrimaryType::kI8);
-  bool tensorIsReal = (valTp >= PrimaryType::kF64 && valTp <= PrimaryType::kI8);
-  if ((valueKind == SparseTensorFile::ValueKind::kReal && tensorIsInteger) ||
-      (valueKind == SparseTensorFile::ValueKind::kComplex && tensorIsReal)) {
+  if (!stfile.canReadAs(valTp))
     MLIR_SPARSETENSOR_FATAL(
         "Tensor element type %d not compatible with values in file %s\n",
         static_cast<int>(valTp), filename);
-  }
   stfile.assertMatchesShape(rank, shape);
   // Prepare sparse tensor object with per-dimension sizes
   // and the number of nonzeros as initial capacity.
   uint64_t nnz = stfile.getNNZ();
   auto *coo = SparseTensorCOO<V>::newSparseTensorCOO(rank, stfile.getDimSizes(),
                                                      perm, nnz);
   // Read all nonzero elements.
   std::vector<uint64_t> indices(rank);

mlir/lib/ExecutionEngine/SparseTensor/File.cpp

 void SparseTensorFile::assertMatchesShape(uint64_t rank,
                                           const uint64_t *shape) const {
   assert(rank == getRank() && "Rank mismatch");
   for (uint64_t r = 0; r < rank; ++r)
     assert((shape[r] == 0 || shape[r] == idata[2 + r]) &&
            "Dimension size mismatch");
 }
 
+bool SparseTensorFile::canReadAs(PrimaryType valTy) const {
+  switch (valueKind_) {
+  case ValueKind::kInvalid:
+    assert(false && "Must readHeader() before calling canReadAs()");
+    return false; // In case assertions are disabled.
+  case ValueKind::kPattern:
+    return true;
+  case ValueKind::kInteger:
+    // When the file is specified to store integer values, we still
+    // allow implicitly converting those to floating primary-types.
+    return isRealPrimaryType(valTy);
+  case ValueKind::kReal:
+    // When the file is specified to store real/floating values, then
+    // we disallow implicit conversion to integer primary-types.
+    return isFloatingPrimaryType(valTy);
+  case ValueKind::kComplex:
+    // When the file is specified to store complex values, then we
+    // require a complex primary-type.
+    return isComplexPrimaryType(valTy);
+  case ValueKind::kUndefined:
+    // The "extended" FROSTT format doesn't specify a ValueKind.
+    // So we allow implicitly converting the stored values to both
+    // integer and floating primary-types.
+    return isRealPrimaryType(valTy);
+  }
+}
+
 /// Helper to convert C-style strings (i.e., '\0' terminated) to lower case.
 static inline char *toLower(char *token) {
   for (char *c = token; *c; ++c)
     *c = tolower(*c);
   return token;
 }
 
 /// Read the MME header of a general sparse matrix of type real.