[mlir][sparse] introduce LevelFormat which do not encoding level properties of a sparse tensor level.

[Peiming]

update comments.

[Peiming]

update comments.

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

161

storage formats (plural)

249

make a note that adding properties will l need to be addressed with new parameters

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

164

can we do without?

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

161–163

please update this documentation per my discussion on the previous patch

243

"LevelFormat"

248

documentation needs updating

253–255

Can you use return isValidDLT(dlt) ? dlt : std::nullopt; instead

258

per my comment on the previous patch, can you add a comment about why you want to assert this? I assume it's so that the next block of asserions can use the operator== without casting. But it's good to make it clear whether the assertion is just for that reason or whether it's something we actually intend the API to promise.

277–300

I think it'd be better to use some macros to help simplify this, but I can do that in a separate patch

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

249

Actually, I'm thinking it'd be nicer to just define a LevelProperties type with its own operator| to set the various property bits. But I can make a separate patch for that

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

161–163

Ah, looks like that got resolved between when I loaded the page and when my comment posted :)

[Peiming]

address comment

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

249

SG

258

removed

[Peiming]

update comments.

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

162

"properties"

162

"the"

252

I'd change this to "TODO: factor out a new LevelProperties type so we can add new properties without changing this function's signature"

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

244

"Returns"

252

"are"

252

"for"

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

252

"Returns"

[Peiming]

address comments.

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

259

Do you really need to pass this template parameter? The compiler should be able to infer it from the dlt variable.

[Peiming]

address comments.

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

 #ifndef MLIR_DIALECT_SPARSETENSOR_IR_ENUMS_H
 #define MLIR_DIALECT_SPARSETENSOR_IR_ENUMS_H
 
 // NOTE: Client code will need to include "mlir/ExecutionEngine/Float16bits.h"
 // if they want to use the `MLIR_SPARSETENSOR_FOREVERY_V` macro.
 
 #include <cinttypes>
 #include <complex>
+#include <optional>
 
 namespace mlir {
 namespace sparse_tensor {
 
 /// This type is used in the public API at all places where MLIR expects
 /// values with the built-in type "index".  For now, we simply assume that
 /// type is 64-bit, but targets with different "index" bit widths should
 /// link with an alternatively built runtime support library.
enum class DimLevelType : uint8_t {
   CompressedNo = 10,   // 0b010_10
   CompressedNuNo = 11, // 0b010_11
   Singleton = 16,      // 0b100_00
   SingletonNu = 17,    // 0b100_01
   SingletonNo = 18,    // 0b100_10
   SingletonNuNo = 19,  // 0b100_11
 };
 
+/// This enum defines all the storage formats supported by the sparse compiler,
+/// without the level properties.
+enum class LevelFormat : uint8_t {
+  Dense = 4,      // 0b001_00
+  Compressed = 8, // 0b010_00
+  Singleton = 16, // 0b100_00
+};
+
 /// Returns string representation of the given dimension level type.
 inline std::string toMLIRString(DimLevelType dlt) {
   switch (dlt) {
   // TODO: should probably raise an error instead of printing it...
   case DimLevelType::Undef:
     return "\"undef\"";
   case DimLevelType::Dense:
     return "\"dense\"";
constexpr bool isOrderedDLT(DimLevelType dlt) {
   return !(static_cast<uint8_t>(dlt) & 2);
 }
 
 /// Check if the `DimLevelType` is unique (regardless of storage format).
 constexpr bool isUniqueDLT(DimLevelType dlt) {
   return !(static_cast<uint8_t>(dlt) & 1);
 }
 
+/// Convert a DimLevelType to its corresponding LevelFormat.
+/// Returns std::nullopt when input dlt is Undef.
+constexpr std::optional<LevelFormat> getLevelFormat(DimLevelType dlt) {
+  if (dlt == DimLevelType::Undef)
+    return std::nullopt;
+  return static_cast<LevelFormat>(static_cast<uint8_t>(dlt) & ~3);
+}
+
+/// Convert a LevelFormat to its corresponding DimLevelType with the given
+/// properties. Returns std::nullopt when the properties are not applicable for
+/// the input level format.
+/// TODO: factor out a new LevelProperties type so we can add new properties
+/// without changing this function's signature
+constexpr std::optional<DimLevelType>
+getDimLevelType(LevelFormat lf, bool ordered, bool unique) {
+  auto dlt = static_cast<DimLevelType>(static_cast<uint8_t>(lf) |
+                                       (ordered ? 0 : 2) | (unique ? 0 : 1));
+  return isValidDLT(dlt) ? std::optional(dlt) : std::nullopt;
+}
+
+/// Ensure the above conversion works as intended.
+static_assert(
+    (getLevelFormat(DimLevelType::Undef) == std::nullopt &&
+     *getLevelFormat(DimLevelType::Dense) == LevelFormat::Dense &&
+     *getLevelFormat(DimLevelType::Compressed) == LevelFormat::Compressed &&
+     *getLevelFormat(DimLevelType::CompressedNu) == LevelFormat::Compressed &&
+     *getLevelFormat(DimLevelType::CompressedNo) == LevelFormat::Compressed &&
+     *getLevelFormat(DimLevelType::CompressedNuNo) == LevelFormat::Compressed &&
+     *getLevelFormat(DimLevelType::Singleton) == LevelFormat::Singleton &&
+     *getLevelFormat(DimLevelType::SingletonNu) == LevelFormat::Singleton &&
+     *getLevelFormat(DimLevelType::SingletonNo) == LevelFormat::Singleton &&
+     *getLevelFormat(DimLevelType::SingletonNuNo) == LevelFormat::Singleton),
+    "getLevelFormat conversion is broken");
+
+static_assert(
+    (getDimLevelType(LevelFormat::Dense, false, true) == std::nullopt &&
+     getDimLevelType(LevelFormat::Dense, true, false) == std::nullopt &&
+     getDimLevelType(LevelFormat::Dense, false, false) == std::nullopt &&
+     *getDimLevelType(LevelFormat::Dense, true, true) == DimLevelType::Dense &&
+     *getDimLevelType(LevelFormat::Compressed, true, true) ==
+         DimLevelType::Compressed &&
+     *getDimLevelType(LevelFormat::Compressed, true, false) ==
+         DimLevelType::CompressedNu &&
+     *getDimLevelType(LevelFormat::Compressed, false, true) ==
+         DimLevelType::CompressedNo &&
+     *getDimLevelType(LevelFormat::Compressed, false, false) ==
+         DimLevelType::CompressedNuNo &&
+     *getDimLevelType(LevelFormat::Singleton, true, true) ==
+         DimLevelType::Singleton &&
+     *getDimLevelType(LevelFormat::Singleton, true, false) ==
+         DimLevelType::SingletonNu &&
+     *getDimLevelType(LevelFormat::Singleton, false, true) ==
+         DimLevelType::SingletonNo &&
+     *getDimLevelType(LevelFormat::Singleton, false, false) ==
+         DimLevelType::SingletonNuNo),
+    "getDimLevelType conversion is broken");
+
 // Ensure the above predicates work as intended.
 static_assert((isValidDLT(DimLevelType::Undef) &&
                isValidDLT(DimLevelType::Dense) &&
                isValidDLT(DimLevelType::Compressed) &&
                isValidDLT(DimLevelType::CompressedNu) &&
                isValidDLT(DimLevelType::CompressedNo) &&
                isValidDLT(DimLevelType::CompressedNuNo) &&
                isValidDLT(DimLevelType::Singleton) &&
                isValidDLT(DimLevelType::SingletonNu) &&