[Mlir-commits] [mlir] [mlir][sparse] remove LevelType enum, construct LevelType from LevelF… (PR #81799)

[Aart Bik]

================
@@ -153,374 +153,246 @@ enum class Action : uint32_t {
   kSortCOOInPlace = 8,
 };
 
-/// This enum defines all the sparse representations supportable by
-/// the SparseTensor dialect. We use a lightweight encoding to encode
-/// the "format" per se (dense, compressed, singleton, loose_compressed,
-/// n-out-of-m), the "properties" (ordered, unique) as well as n and m when
-/// the format is NOutOfM.
-/// The encoding is chosen for performance of the runtime library, and thus may
-/// change in future versions; consequently, client code should use the
-/// predicate functions defined below, rather than relying on knowledge
-/// about the particular binary encoding.
-///
-/// The `Undef` "format" is a special value used internally for cases
-/// where we need to store an undefined or indeterminate `LevelType`.
-/// It should not be used externally, since it does not indicate an
-/// actual/representable format.
-///
-/// Bit manipulations for LevelType:
-///
-/// | 8-bit n | 8-bit m | 16-bit LevelFormat | 16-bit LevelProperty |
-///
-enum class LevelType : uint64_t {
-  Undef = 0x000000000000,
-  Dense = 0x000000010000,
-  Compressed = 0x000000020000,
-  CompressedNu = 0x000000020001,
-  CompressedNo = 0x000000020002,
-  CompressedNuNo = 0x000000020003,
-  Singleton = 0x000000040000,
-  SingletonNu = 0x000000040001,
-  SingletonNo = 0x000000040002,
-  SingletonNuNo = 0x000000040003,
-  LooseCompressed = 0x000000080000,
-  LooseCompressedNu = 0x000000080001,
-  LooseCompressedNo = 0x000000080002,
-  LooseCompressedNuNo = 0x000000080003,
-  NOutOfM = 0x000000100000,
-};
-
 /// This enum defines all supported storage format without the level properties.
 enum class LevelFormat : uint64_t {
+  Undef = 0x00000000,
   Dense = 0x00010000,
   Compressed = 0x00020000,
   Singleton = 0x00040000,
   LooseCompressed = 0x00080000,
   NOutOfM = 0x00100000,
 };
 
+/// Returns string representation of the given level format.
+constexpr const char *toFormatString(LevelFormat lvlFmt) {
+  switch (lvlFmt) {
+  case LevelFormat::Undef:
+    return "undef";
+  case LevelFormat::Dense:
+    return "dense";
+  case LevelFormat::Compressed:
+    return "compressed";
+  case LevelFormat::Singleton:
+    return "singleton";
+  case LevelFormat::LooseCompressed:
+    return "loose_compressed";
+  case LevelFormat::NOutOfM:
+    return "structured";
+  }
+  return "";
+}
+
 /// This enum defines all the nondefault properties for storage formats.
-enum class LevelPropertyNondefault : uint64_t {
+enum class LevelPropNonDefault : uint64_t {
   Nonunique = 0x0001,
   Nonordered = 0x0002,
 };
 
-/// Get N of NOutOfM level type.
-constexpr uint64_t getN(LevelType lt) {
-  return (static_cast<uint64_t>(lt) >> 32) & 0xff;
+/// Returns string representation of the given level properties.
+constexpr const char *toPropString(LevelPropNonDefault lvlProp) {
+  switch (lvlProp) {
+  case LevelPropNonDefault::Nonunique:
+    return "nonunique";
+  case LevelPropNonDefault::Nonordered:
+    return "nonordered";
+  }
+  return "";
 }
 
-/// Get M of NOutOfM level type.
-constexpr uint64_t getM(LevelType lt) {
-  return (static_cast<uint64_t>(lt) >> 40) & 0xff;
-}
+/// This enum defines all the sparse representations supportable by
+/// the SparseTensor dialect. We use a lightweight encoding to encode
+/// the "format" per se (dense, compressed, singleton, loose_compressed,
+/// n-out-of-m), the "properties" (ordered, unique) as well as n and m when
+/// the format is NOutOfM.
+/// The encoding is chosen for performance of the runtime library, and thus may
+/// change in future versions; consequently, client code should use the
+/// predicate functions defined below, rather than relying on knowledge
+/// about the particular binary encoding.
+///
+/// The `Undef` "format" is a special value used internally for cases
+/// where we need to store an undefined or indeterminate `LevelType`.
+/// It should not be used externally, since it does not indicate an
+/// actual/representable format.
 
-/// Convert N of NOutOfM level type to the stored bits.
-constexpr uint64_t nToBits(uint64_t n) { return n << 32; }
+struct LevelType {
+public:
+  /// Check that the `LevelType` contains a valid (possibly undefined) value.
+  static constexpr bool isValidLvlBits(uint64_t lvlBits) {
+    const uint64_t formatBits = lvlBits & 0xffff0000;
+    const uint64_t propertyBits = lvlBits & 0xffff;
+    // If undefined/dense/NOutOfM, then must be unique and ordered.
+    // Otherwise, the format must be one of the known ones.
+    return (formatBits <= 0x10000 || formatBits == 0x100000)
+               ? (propertyBits == 0)
+               : (formatBits == 0x20000 || formatBits == 0x40000 ||
+                  formatBits == 0x80000);
+  }
 
-/// Convert M of NOutOfM level type to the stored bits.
-constexpr uint64_t mToBits(uint64_t m) { return m << 40; }
+  /// Convert a LevelFormat to its corresponding LevelType with the given
+  /// properties. Returns std::nullopt when the properties are not applicable
+  /// for the input level format.
+  static std::optional<LevelType>
+  buildLvlType(LevelFormat lf,
+               const std::vector<LevelPropNonDefault> &properties,
+               uint64_t n = 0, uint64_t m = 0) {
+    uint64_t newN = n << 32;
----------------
aartbik wrote:

to make sure we don't corrupt other bits, this should be (n & 0xff) << 32
(and perhaps an assert that n &0xff == n), same for m

https://github.com/llvm/llvm-project/pull/81799