[mlir][sparse] Improve handling of NEW_SYNTAX

[wrengr]

Summary

Improves the conversion from DimLvlMap to STEA, in order to correct rank-mismatch issues in the roundtrip tests.

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td

def SparseTensorDimSliceAttr : SparseTensor_Attr<"SparseTensorDimSlice", []> {
 
   let parameters = (
     ins
     "int64_t" : $offset,
     "int64_t" : $size,
     "int64_t" : $stride
   );
 
+  let builders = [
+    // The nop slice (i.e., that includes everything).
+    AttrBuilder<(ins), [{ return $_get($_ctxt, 0, kDynamic, 1); }]>
+  ];
+
   let extraClassDeclaration = [{
     void print(llvm::raw_ostream &os) const;
 
     /// Special value for dynamic offset/size/stride.
     static constexpr int64_t kDynamic = -1;
     static constexpr bool isDynamic(int64_t v) { return v == kDynamic; }
     static std::optional<uint64_t> getStatic(int64_t v);
     static std::string getStaticString(int64_t v);

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

return isDynamic(getOffset()) && isDynamic(getStride()) &&
          isDynamic(getSize());
 }
 
 std::string SparseTensorDimSliceAttr::getStaticString(int64_t v) {
   return isDynamic(v) ? "?" : std::to_string(v);
 }
 
 void SparseTensorDimSliceAttr::print(llvm::raw_ostream &os) const {
+  assert(getImpl() && "Uninitialized SparseTensorDimSliceAttr");
   os << '(';
   os << getStaticString(getOffset());
   os << ", ";
   os << getStaticString(getSize());
   os << ", ";
   os << getStaticString(getStride());
   os << ')';
 }
if (attrName == "lvlTypes") {
       // syntax. While this is ongoing we use the temporary "NEW_SYNTAX = ...."
       // to switch to the new parser. This allows us to gradually migrate
       // examples over to the new surface syntax before making the complete
       // switch once work is completed.
       // TODO: replace everything here with new STEA surface syntax parser
       ir_detail::DimLvlMapParser cParser(parser);
       auto res = cParser.parseDimLvlMap();
       RETURN_ON_FAIL(res);
-      // Proof of concept result.
-      // TODO: use DimLvlMap directly as storage representation
-      for (Level lvl = 0, lvlRank = res->getLvlRank(); lvl < lvlRank; lvl++)
-        lvlTypes.push_back(res->getLvlType(lvl));
+      // TODO: use DimLvlMap directly as storage representation, rather
+      // than converting things over.
+      const auto &dlm = *res;
+
+      ERROR_IF(!lvlTypes.empty(), "Cannot mix `lvlTypes` with `NEW_SYNTAX`")
+      const Level lvlRank = dlm.getLvlRank();
+      for (Level lvl = 0; lvl < lvlRank; lvl++)
+        lvlTypes.push_back(dlm.getLvlType(lvl));
+
+      ERROR_IF(!dimSlices.empty(), "Cannot mix `dimSlices` with `NEW_SYNTAX`")
+      const Dimension dimRank = dlm.getDimRank();
+      for (Dimension dim = 0; dim < dimRank; dim++)
+        dimSlices.push_back(dlm.getDimSlice(dim));
+      // NOTE: the old syntax requires an all-or-nothing approach to
+      // `dimSlices`; therefore, if any slice actually exists then we need
+      // to convert null-DSA into default/nop DSA.
+      const auto isDefined = [](SparseTensorDimSliceAttr slice) {
+        return static_cast<bool>(slice.getImpl());
+      };
+      if (llvm::any_of(dimSlices, isDefined)) {
+        const auto defaultSlice =
+            SparseTensorDimSliceAttr::get(parser.getContext());
+        for (Dimension dim = 0; dim < dimRank; dim++)
+          if (!isDefined(dimSlices[dim]))
+            dimSlices[dim] = defaultSlice;
+      } else {
+        dimSlices.clear();
+      }
+
+      ERROR_IF(dimToLvl, "Cannot mix `dimToLvl` with `NEW_SYNTAX`")
+      dimToLvl = dlm.getDimToLvlMap(parser.getContext());
     }
 
     // Only the last item can omit the comma
     if (parser.parseOptionalComma().failed())
       break;
   }
 
   RETURN_ON_FAIL(parser.parseRBrace())

mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir

 
 // -----
 
 #CSR_implicit = #sparse_tensor.encoding<{
   NEW_SYNTAX =
   (d0, d1) -> (d0 : dense, d1 : compressed)
 }>
 
-// CHECK-LABEL: func private @foo(
+// CHECK-LABEL: func private @CSR_implicit(
 // CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ] }>>
-func.func private @foo(%arg0: tensor<?x?xf64, #CSR_implicit>) {
+func.func private @CSR_implicit(%arg0: tensor<?x?xf64, #CSR_implicit>) {
   return
 }
 
 // -----
 
 #CSR_explicit = #sparse_tensor.encoding<{
   NEW_SYNTAX =
   {l0, l1} (d0 = l0, d1 = l1) -> (l0 = d0 : dense, l1 = d1 : compressed)
 }>
 
-// CHECK-LABEL: func private @foo(
+// CHECK-LABEL: func private @CSR_explicit(
 // CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ] }>>
-func.func private @foo(%arg0: tensor<?x?xf64, #CSR_explicit>) {
+func.func private @CSR_explicit(%arg0: tensor<?x?xf64, #CSR_explicit>) {
   return
 }
 
 // -----
 
 #BCSR_implicit = #sparse_tensor.encoding<{
   NEW_SYNTAX =
   ( i, j ) ->
   ( i floordiv 2 : compressed,
     j floordiv 3 : compressed,
     i mod 2      : dense,
     j mod 3      : dense
   )
 }>
 
-// FIXME: should not have to use 4 dims ;-)
-//
-// CHECK-LABEL: func private @foo(
-// CHECK-SAME: tensor<?x?x?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "compressed", "compressed", "dense", "dense" ] }>>
-func.func private @foo(%arg0: tensor<?x?x?x?xf64, #BCSR_implicit>) {
+// CHECK-LABEL: func private @BCSR_implicit(
+// CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "compressed", "compressed", "dense", "dense" ], dimToLvl = affine_map<(d0, d1) -> (d0 floordiv 2, d1 floordiv 3, d0 mod 2, d1 mod 3)> }>>
+func.func private @BCSR_implicit(%arg0: tensor<?x?xf64, #BCSR_implicit>) {
   return
 }
 
 // -----
 
 #BCSR_explicit = #sparse_tensor.encoding<{
   NEW_SYNTAX =
   {il, jl, ii, jj}
   ( i = il * 2 + ii,
     j = jl * 3 + jj
   ) ->
   ( il = i floordiv 2 : compressed,
     jl = j floordiv 3 : compressed,
     ii = i mod 2      : dense,
     jj = j mod 3      : dense
   )
 }>
 
-// FIXME: should not have to use 4 dims ;-)
-//
-// CHECK-LABEL: func private @foo(
-// CHECK-SAME: tensor<?x?x?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "compressed", "compressed", "dense", "dense" ] }>>
-func.func private @foo(%arg0: tensor<?x?x?x?xf64, #BCSR_explicit>) {
+// CHECK-LABEL: func private @BCSR_explicit(
+// CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "compressed", "compressed", "dense", "dense" ], dimToLvl = affine_map<(d0, d1) -> (d0 floordiv 2, d1 floordiv 3, d0 mod 2, d1 mod 3)> }>>
+func.func private @BCSR_explicit(%arg0: tensor<?x?xf64, #BCSR_explicit>) {
   return
 }
 
 // -----
 
 #NV_24 = #sparse_tensor.encoding<{
   NEW_SYNTAX =
   ( i, j ) ->
   ( i            : dense,
     j floordiv 4 : dense,
     j mod 4      : compressed24
   )
 }>
 
-//
-// CHECK-LABEL: func private @foo_2_out_of_4(
-// CHECK-SAME: tensor<?x?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "dense", "compressed24" ] }>>
-func.func private @foo_2_out_of_4(%arg0: tensor<?x?x?xf64, #NV_24>) {
+// 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
 }