[mlir][sparse] Surface syntax change in parsing

[yinying-lisa-li]

Summary

Example: compressed(nonunique, nonordered) or compressed(nonordered, nonunique) instead of compressed_nu_no.

[yinying-lisa-li]

Remove isSet variable.

[aartbik]

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

204

I am not super thrilled with the name, but have no good alternatives
since indeed we have default prop = 0 and non-default prop = 1 for the bit position

but okay for now

mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp

394

since we are not round tripping, this is not required yet right?
but okay to keep, just observing

mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.h

22–23

I would still keep the trivial constructor (now without the map)

LvlTypeParser() {}

22–23

I think you can remove these TODOs, we know pretty well what we still need for migraton

[yinying-lisa-li]

Addressing comments.

[yinying-lisa-li]

mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.h

22–23

Sounds good. Cider is recommending LvlTypeParser() = default;, so I'll use that.

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

 ///
 /// The `Undef` "format" is a special value used internally for cases
 /// where we need to store an undefined or indeterminate `DimLevelType`.
 /// It should not be used externally, since it does not indicate an
 /// actual/representable format.
 ///
 // TODO: We should generalize TwoOutOfFour to N out of M and use property to
 // encode the value of N and M.
+// TODO: Update DimLevelType to use lower 8 bits for storage formats and the
+// higher 4 bits to store level properties. Consider CompressedWithHi and
+// TwoOutOfFour as properties instead of formats.
 enum class DimLevelType : uint8_t {
   Undef = 0,                 // 0b00000_00
   Dense = 4,                 // 0b00001_00
   Compressed = 8,            // 0b00010_00
   CompressedNu = 9,          // 0b00010_01
   CompressedNo = 10,         // 0b00010_10
   CompressedNuNo = 11,       // 0b00010_11
   Singleton = 16,            // 0b00100_00

 enum class LevelFormat : uint8_t {
   Dense = 4,             // 0b00001_00
   Compressed = 8,        // 0b00010_00
   Singleton = 16,        // 0b00100_00
   CompressedWithHi = 32, // 0b01000_00
   TwoOutOfFour = 64,     // 0b10000_00
 };
 
+/// This enum defines all the nondefault properties for storage formats.
+enum class LevelNondefaultProperty : uint8_t {
+  Nonunique = 1,  // 0b00000_01
+  Nonordered = 2, // 0b00000_10
+  High = 32,      // 0b01000_00
+  Block2_4 = 64   // 0b10000_00
+};
+
 /// Returns string representation of the given dimension level type.
 constexpr const char *toMLIRString(DimLevelType dlt) {
   switch (dlt) {
   case DimLevelType::Undef:
     return "undef";
   case DimLevelType::Dense:
     return "dense";
   case DimLevelType::Compressed:

mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp

if (symRank != 0) {
     os << ']';
   }
 
   // LvlVar forward-declarations.
   if (mustPrintLvlVars) {
     os << '{';
     llvm::interleaveComma(
         lvlSpecs, os, [&](LvlSpec const &spec) { os << spec.getBoundVar(); });
-    os << '}';
+    os << "} ";
   }
 
   // Dimension specifiers.
   os << '(';
   llvm::interleaveComma(
       dimSpecs, os, [&](DimSpec const &spec) { spec.print(os, wantElision); });
   os << ") -> (";
   // Level specifiers.
   wantElision = wantElision && !mustPrintLvlVars;
   llvm::interleaveComma(
       lvlSpecs, os, [&](LvlSpec const &spec) { spec.print(os, wantElision); });
   os << ')';
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMapParser.cpp

ParseResult DimLvlMapParser::parseLvlSpec(bool requireLvlVarBinding) {
   // field every time `parseLvlSpec` is called.
   FAILURE_IF_FAILED(parser.parseAffineExpr(dimsAndSymbols, affine))
   LvlExpr expr{affine};
 
   FAILURE_IF_FAILED(parser.parseColon())
   const auto type = lvlTypeParser.parseLvlType(parser);
   FAILURE_IF_FAILED(type)
 
-  lvlSpecs.emplace_back(var, expr, *type);
+  lvlSpecs.emplace_back(var, expr, static_cast<DimLevelType>(*type));
   return success();
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.h

 //===- LvlTypeParser.h - `DimLevelType` parser ------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPARSETENSOR_IR_DETAIL_LVLTYPEPARSER_H
 #define MLIR_DIALECT_SPARSETENSOR_IR_DETAIL_LVLTYPEPARSER_H
 
-#include "mlir/Dialect/SparseTensor/IR/Enums.h"
 #include "mlir/IR/OpImplementation.h"
-#include "llvm/ADT/StringMap.h"
 
 namespace mlir {
 namespace sparse_tensor {
 namespace ir_detail {
 
-//===----------------------------------------------------------------------===//
-// These macros are for generating a C++ expression of type
-// `std::initializer_list<std::pair<StringRef,DimLevelType>>` since there's
-// no way to construct an object of that type directly via C++ code.
-#define FOREVERY_LEVELTYPE(DO)                                                 \
-  DO(DimLevelType::Dense)                                                      \
-  DO(DimLevelType::Compressed)                                                 \
-  DO(DimLevelType::CompressedNu)                                               \
-  DO(DimLevelType::CompressedNo)                                               \
-  DO(DimLevelType::CompressedNuNo)                                             \
-  DO(DimLevelType::Singleton)                                                  \
-  DO(DimLevelType::SingletonNu)                                                \
-  DO(DimLevelType::SingletonNo)                                                \
-  DO(DimLevelType::SingletonNuNo)                                              \
-  DO(DimLevelType::CompressedWithHi)                                           \
-  DO(DimLevelType::CompressedWithHiNu)                                         \
-  DO(DimLevelType::CompressedWithHiNo)                                         \
-  DO(DimLevelType::CompressedWithHiNuNo)                                       \
-  DO(DimLevelType::TwoOutOfFour)
-#define LEVELTYPE_INITLIST_ELEMENT(lvlType)                                    \
-  std::make_pair(StringRef(toMLIRString(lvlType)), lvlType),
-#define LEVELTYPE_INITLIST                                                     \
-  { FOREVERY_LEVELTYPE(LEVELTYPE_INITLIST_ELEMENT) }
-
-// TODO(wrengr): Since this parser is non-trivial to construct, is there
-// any way to hook into the parsing process so that we construct it only once
-// at the begining of parsing and then destroy it once parsing has finished?
 class LvlTypeParser {
-  const llvm::StringMap<DimLevelType> map;
-
 public:
-  explicit LvlTypeParser() : map(LEVELTYPE_INITLIST) {}
-#undef LEVELTYPE_INITLIST
-#undef LEVELTYPE_INITLIST_ELEMENT
-#undef FOREVERY_LEVELTYPE
-
-  std::optional<DimLevelType> lookup(StringRef str) const;
-  std::optional<DimLevelType> lookup(StringAttr str) const;
-  ParseResult parseLvlType(AsmParser &parser, DimLevelType &out) const;
-  FailureOr<DimLevelType> parseLvlType(AsmParser &parser) const;
-  // TODO(wrengr): `parseOptionalLvlType`?
-  // TODO(wrengr): `parseLvlTypeList`?
+  LvlTypeParser() = default;
+  FailureOr<uint8_t> parseLvlType(AsmParser &parser) const;
+
+private:
+  ParseResult parseProperty(AsmParser &parser, uint8_t *properties) const;
 };
 
 } // namespace ir_detail
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_IR_DETAIL_LVLTYPEPARSER_H

mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.cpp

 //===- LvlTypeParser.h - `DimLevelType` parser ----------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "LvlTypeParser.h"
+#include "mlir/Dialect/SparseTensor/IR/Enums.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 using namespace mlir::sparse_tensor::ir_detail;
 
 //===----------------------------------------------------------------------===//
 // TODO(wrengr): rephrase these to do the trick for gobbling up any trailing
 // semicolon
#define ERROR_IF(COND, MSG)                                                    \
   if (COND) {                                                                  \
     return parser.emitError(loc, MSG);                                         \
   }
 
 //===----------------------------------------------------------------------===//
 // `LvlTypeParser` implementation.
 //===----------------------------------------------------------------------===//
 
-std::optional<DimLevelType> LvlTypeParser::lookup(StringRef str) const {
-  // NOTE: `StringMap::lookup` will return a default-constructed value if
-  // the key isn't found; which for enums means zero, and therefore makes
-  // it impossible to distinguish between actual zero-DimLevelType vs
-  // not-found.  Whereas `StringMap::at` asserts that the key is found,
-  // which we don't want either.
-  const auto it = map.find(str);
-  return it == map.end() ? std::nullopt : std::make_optional(it->second);
-}
-
-std::optional<DimLevelType> LvlTypeParser::lookup(StringAttr str) const {
-  return str ? lookup(str.getValue()) : std::nullopt;
-}
+FailureOr<uint8_t> LvlTypeParser::parseLvlType(AsmParser &parser) const {
+  StringRef base;
+  FAILURE_IF_FAILED(parser.parseOptionalKeyword(&base));
+  uint8_t properties = 0;
+  const auto loc = parser.getCurrentLocation();
 
-FailureOr<DimLevelType> LvlTypeParser::parseLvlType(AsmParser &parser) const {
-  DimLevelType out;
-  FAILURE_IF_FAILED(parseLvlType(parser, out))
-  return out;
+  ParseResult res = parser.parseCommaSeparatedList(
+      mlir::OpAsmParser::Delimiter::OptionalParen,
+      [&]() -> ParseResult { return parseProperty(parser, &properties); },
+      " in level property list");
+  FAILURE_IF_FAILED(res)
+
+  // Set the base bit for properties.
+  if (base.compare("dense") == 0) {
+    properties |= static_cast<uint8_t>(LevelFormat::Dense);
+  } else if (base.compare("compressed") == 0) {
+    // TODO: Remove this condition once dimLvlType enum is refactored. Current
+    // enum treats High and TwoOutOfFour as formats instead of properties.
+    if (!(properties & static_cast<uint8_t>(LevelNondefaultProperty::High) ||
+          properties &
+              static_cast<uint8_t>(LevelNondefaultProperty::Block2_4))) {
+      properties |= static_cast<uint8_t>(LevelFormat::Compressed);
+    }
+  } else if (base.compare("singleton") == 0) {
+    properties |= static_cast<uint8_t>(LevelFormat::Singleton);
+  } else {
+    parser.emitError(loc, "unknown level format");
+    return failure();
+  }
+
+  ERROR_IF(!isValidDLT(static_cast<DimLevelType>(properties)),
+           "invalid level type");
+  return properties;
 }
 
-ParseResult LvlTypeParser::parseLvlType(AsmParser &parser,
-                                        DimLevelType &out) const {
-  const auto loc = parser.getCurrentLocation();
+ParseResult LvlTypeParser::parseProperty(AsmParser &parser,
+                                         uint8_t *properties) const {
   StringRef strVal;
   FAILURE_IF_FAILED(parser.parseOptionalKeyword(&strVal));
-  const auto lvlType = lookup(strVal);
-  ERROR_IF(!lvlType, "unknown level-type '" + strVal + "'")
-  out = *lvlType;
+  if (strVal.compare("nonunique") == 0) {
+    *properties |= static_cast<uint8_t>(LevelNondefaultProperty::Nonunique);
+  } else if (strVal.compare("nonordered") == 0) {
+    *properties |= static_cast<uint8_t>(LevelNondefaultProperty::Nonordered);
+  } else if (strVal.compare("high") == 0) {
+    *properties |= static_cast<uint8_t>(LevelNondefaultProperty::High);
+  } else if (strVal.compare("block2_4") == 0) {
+    *properties |= static_cast<uint8_t>(LevelNondefaultProperty::Block2_4);
+  } else {
+    parser.emitError(parser.getCurrentLocation(), "unknown level property");
+    return failure();
+  }
   return success();
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir

 
 // CHECK-LABEL: func private @sparse_dcsc(
 // CHECK-SAME: tensor<?x?xf32, #sparse_tensor.encoding<{ lvlTypes = [ "compressed", "compressed" ], dimToLvl = affine_map<(d0, d1) -> (d1, d0)>, crdWidth = 64 }>>)
 func.func private @sparse_dcsc(tensor<?x?xf32, #DCSC>)
 
 // -----
 
 #COO = #sparse_tensor.encoding<{
-  lvlTypes = [ "compressed_nu_no", "singleton_no" ]
+  map = (d0, d1) -> (d0 : compressed(nonunique, nonordered), d1 : singleton(nonordered))
 }>
 
 // CHECK-LABEL: func private @sparse_coo(
 // CHECK-SAME: tensor<?x?xf32, #sparse_tensor.encoding<{ lvlTypes = [ "compressed_nu_no", "singleton_no" ] }>>)
 func.func private @sparse_coo(tensor<?x?xf32, #COO>)
 
 // -----
 
 #BCOO = #sparse_tensor.encoding<{
-  lvlTypes = [ "dense", "compressed_hi_nu", "singleton" ]
+  map = (d0, d1, d2) -> (d0 : dense, d1 : compressed(nonunique, high), d2 : singleton)
 }>
 
 // CHECK-LABEL: func private @sparse_bcoo(
 // CHECK-SAME: tensor<?x?x?xf32, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed_hi_nu", "singleton" ] }>>)
 func.func private @sparse_bcoo(tensor<?x?x?xf32, #BCOO>)
 
 // -----
 
 #SortedCOO = #sparse_tensor.encoding<{
-  lvlTypes = [ "compressed_nu", "singleton" ]
+  map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton)
 }>
 
 // CHECK-LABEL: func private @sparse_sorted_coo(
 // CHECK-SAME: tensor<10x10xf64, #sparse_tensor.encoding<{ lvlTypes = [ "compressed_nu", "singleton" ] }>>)
 func.func private @sparse_sorted_coo(tensor<10x10xf64, #SortedCOO>)
 
 // -----
 

 func.func private @sparse_slice(tensor<?x?xf64, #CSR_SLICE>)
 
 // -----
 
 // TODO: It is probably better to use [dense, dense, 2:4] (see NV_24 defined using new syntax
 // below) to encode a 2D matrix, but it would require dim2lvl mapping which is not ready yet.
 // So we take the simple path for now.
 #NV_24= #sparse_tensor.encoding<{
-  lvlTypes = [ "dense", "compressed24" ],
+  map = (d0, d1) -> (d0 : dense, d1 : compressed(block2_4))
 }>
 
 // CHECK-LABEL: func private @sparse_2_out_of_4(
 // CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed24" ] }>>
 func.func private @sparse_2_out_of_4(tensor<?x?xf64, #NV_24>)
 
 // -----
 

 }
 
 // -----
 
 #NV_24 = #sparse_tensor.encoding<{
   map = ( i, j ) ->
   ( i            : dense,
     j floordiv 4 : dense,
-    j mod 4      : compressed24
+    j mod 4      : compressed(block2_4)
   )
 }>
 
 // CHECK-LABEL: func private @NV_24(
 // CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "dense", "compressed24" ], dimToLvl = affine_map<(d0, d1) -> (d0, d1 floordiv 4, d1 mod 4)> }>>
 func.func private @NV_24(%arg0: tensor<?x?xf64, #NV_24>) {
   return
 }