[mlir][sparse] Use the runtime DimLevelType instead of a separate tablegen enum

[wrengr]

Summary

This differential replaces all uses of SparseTensorEncodingAttr::DimLevelType with DimLevelType. The next differential will break out a separate library for the DimLevelType enum, so that the Dialect code doesn't need to depend on the rest of the runtime

Depends On D135995

[wrengr]

git-clang-format

[aartbik]

This is a pretty sweet cleanup of unnecessary codedup, Wren. Thanks!

[wrengr]

rebase

[wrengr]

correcting rebase error

[qcolombet]

Hi @wrengr ,

Posting it here in case you didn't see the notification on the commit itself:

I have a bunch of compiler errors because of that change.
Could you fix or revert, please?

E.g.,

In file included from mlir/include/mlir/InitAllDialects.h:58,                  
                 from mlir/unittests/ExecutionEngine/Invoke.cpp:23:            
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h: In function ‘constexpr bool mlir::sparse_tensor::isDenseDim(mlir::RankedTensorType, uint64_t)’:                                                                      
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h:53:36: error: call to non-‘constexpr’ function ‘mlir::sparse_tensor::DimLevelType mlir::sparse_tensor::getDimLevelType(mlir::RankedTensorType, uint64_t)’             
   53 |   return isDenseDLT(getDimLevelType(type, d));                                                                                          
      |                     ~~~~~~~~~~~~~~~^~~~~~~~~                                                                                            
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h:44:21: note: ‘mlir::sparse_tensor::DimLevelType mlir::sparse_tensor::getDimLevelType(mlir::RankedTensorType, uint64_t)’ declared here                                 
   44 | inline DimLevelType getDimLevelType(RankedTensorType type, uint64_t d) {                                                                
      |                     ^~~~~~~~~~~~~~~

[wrengr]

Hi @qcolombet,

sorry for the breakage. That should be fixed by D136123 (which also landed today). For some reason our buildbot didn't catch those warnings on this differential but did so on D136123. Can you roll things forward to where D136123 landed, or do you still need a fix?

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

 //===- SparseTensor.h - Sparse tensor dialect -------------------*- 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_SPARSETENSOR_H_
 #define MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSOR_H_
 
+#include "mlir/ExecutionEngine/SparseTensor/Enums.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TensorEncoding.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 

 /// Convenience method to get a sparse encoding attribute from a type.
 /// Returns null-attribute for any type without an encoding.
 SparseTensorEncodingAttr getSparseTensorEncoding(Type type);
 
 //
 // Dimension level types.
 //
 
-bool isDenseDim(SparseTensorEncodingAttr::DimLevelType dltp);
-bool isCompressedDim(SparseTensorEncodingAttr::DimLevelType dltp);
-bool isSingletonDim(SparseTensorEncodingAttr::DimLevelType dltp);
-
-/// Convenience method to test for dense dimension (0 <= d < rank).
-bool isDenseDim(RankedTensorType type, uint64_t d);
-
-/// Convenience method to test for compressed dimension (0 <= d < rank).
-bool isCompressedDim(RankedTensorType type, uint64_t d);
-
-/// Convenience method to test for singleton dimension (0 <= d < rank).
-bool isSingletonDim(RankedTensorType type, uint64_t d);
+// Cannot be constexpr, because `getRank` isn't constexpr.  However,
+// for some strange reason, the wrapper functions below don't trigger
+// the same [-Winvalid-constexpr] warning (despite this function not
+// being constexpr).
+inline DimLevelType getDimLevelType(RankedTensorType type, uint64_t d) {
+  assert(d < static_cast<uint64_t>(type.getRank()));
+  if (auto enc = getSparseTensorEncoding(type))
+    return enc.getDimLevelType()[d];
+  return DimLevelType::Dense; // unannotated tensor is dense
+}
+
+/// Convenience function to test for dense dimension (0 <= d < rank).
+constexpr bool isDenseDim(RankedTensorType type, uint64_t d) {
+  return isDenseDLT(getDimLevelType(type, d));
+}
+
+/// Convenience function to test for compressed dimension (0 <= d < rank).
+constexpr bool isCompressedDim(RankedTensorType type, uint64_t d) {
+  return isCompressedDLT(getDimLevelType(type, d));
+}
+
+/// Convenience function to test for singleton dimension (0 <= d < rank).
+constexpr bool isSingletonDim(RankedTensorType type, uint64_t d) {
+  return isSingletonDLT(getDimLevelType(type, d));
+}
 
 //
 // Dimension level properties.
 //
 
-bool isOrderedDim(SparseTensorEncodingAttr::DimLevelType dltp);
-bool isUniqueDim(SparseTensorEncodingAttr::DimLevelType dltp);
-
-/// Convenience method to test for ordered property in the
+/// Convenience function to test for ordered property in the
 /// given dimension (0 <= d < rank).
-bool isOrderedDim(RankedTensorType type, uint64_t d);
+constexpr bool isOrderedDim(RankedTensorType type, uint64_t d) {
+  return isOrderedDLT(getDimLevelType(type, d));
+}
 
-/// Convenience method to test for unique property in the
+/// Convenience function to test for unique property in the
 /// given dimension (0 <= d < rank).
-bool isUniqueDim(RankedTensorType type, uint64_t d);
+constexpr bool isUniqueDim(RankedTensorType type, uint64_t d) {
+  return isUniqueDLT(getDimLevelType(type, d));
+}
 
 //
 // Reordering.
 //
 
 uint64_t toOrigDim(const SparseTensorEncodingAttr &enc, uint64_t d);
 uint64_t toStoredDim(const SparseTensorEncodingAttr &enc, uint64_t d);
 

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

let description = [{
     ```
   }];
 
   // Data in sparse tensor encoding.
   let parameters = (
     ins
     // A dimension level type for each dimension of the tensor type.
     ArrayRefParameter<
-      "SparseTensorEncodingAttr::DimLevelType",
+      "::mlir::sparse_tensor::DimLevelType",
       "per dimension level type"
       >: $dimLevelType,
     // A dimension order on the indices of this tensor type.
     "AffineMap":$dimOrdering,
     // A mapping between the original and higher-ordering index space.
     "AffineMap":$higherOrdering,
     // The required bit width for pointer storage.
     "unsigned":$pointerBitWidth,
     // The required bit width for index storage.
     "unsigned":$indexBitWidth
   );
 
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
-
-  let extraClassDeclaration = [{
-    // Dimension level types. By default, each type has the unique and
-    // ordered properties. Alternatives properties are indicated by
-    // Nu (not-unique) and No (not-ordered).
-    //
-    // TODO: separate type and property in encoding
-    //
-    enum class DimLevelType : uint8_t {
-      Dense            = 4,  // 0b001_00
-      Compressed       = 8,  // 0b010_00
-      CompressedNu     = 9,  // 0b010_01
-      CompressedNo     = 10, // 0b010_10
-      CompressedNuNo   = 11, // 0b010_11
-      Singleton        = 16, // 0b100_00
-      SingletonNu      = 17, // 0b100_01
-      SingletonNo      = 18, // 0b100_10
-      SingletonNuNo    = 19, // 0b100_11
-    };
-  }];
 }
 
 def IsSparseTensorPred
   : CPred<"!!::mlir::sparse_tensor::getSparseTensorEncoding($_self)">;
 
 // The following four follow the same idiom as `TensorOf`, `AnyTensor`,
 // `RankedTensorOf`, `AnyRankedTensor`.
 

mlir/lib/CAPI/Dialect/SparseTensor.cpp

 using namespace mlir::sparse_tensor;
 
 MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(SparseTensor, sparse_tensor,
                                       mlir::sparse_tensor::SparseTensorDialect)
 
 // Ensure the C-API enums are int-castable to C++ equivalents.
 static_assert(
     static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE) ==
-            static_cast<int>(SparseTensorEncodingAttr::DimLevelType::Dense) &&
+            static_cast<int>(DimLevelType::Dense) &&
         static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED) ==
-            static_cast<int>(
-                SparseTensorEncodingAttr::DimLevelType::Compressed) &&
+            static_cast<int>(DimLevelType::Compressed) &&
         static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU) ==
-            static_cast<int>(
-                SparseTensorEncodingAttr::DimLevelType::CompressedNu) &&
+            static_cast<int>(DimLevelType::CompressedNu) &&
         static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NO) ==
-            static_cast<int>(
-                SparseTensorEncodingAttr::DimLevelType::CompressedNo) &&
+            static_cast<int>(DimLevelType::CompressedNo) &&
         static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU_NO) ==
-            static_cast<int>(
-                SparseTensorEncodingAttr::DimLevelType::CompressedNuNo) &&
+            static_cast<int>(DimLevelType::CompressedNuNo) &&
         static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON) ==
-            static_cast<int>(
-                SparseTensorEncodingAttr::DimLevelType::Singleton) &&
+            static_cast<int>(DimLevelType::Singleton) &&
         static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU) ==
-            static_cast<int>(
-                SparseTensorEncodingAttr::DimLevelType::SingletonNu) &&
+            static_cast<int>(DimLevelType::SingletonNu) &&
         static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NO) ==
-            static_cast<int>(
-                SparseTensorEncodingAttr::DimLevelType::SingletonNo) &&
+            static_cast<int>(DimLevelType::SingletonNo) &&
         static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU_NO) ==
-            static_cast<int>(
-                SparseTensorEncodingAttr::DimLevelType::SingletonNuNo),
+            static_cast<int>(DimLevelType::SingletonNuNo),
     "MlirSparseTensorDimLevelType (C-API) and DimLevelType (C++) mismatch");
 
 bool mlirAttributeIsASparseTensorEncodingAttr(MlirAttribute attr) {
   return unwrap(attr).isa<SparseTensorEncodingAttr>();
 }
 
 MlirAttribute mlirSparseTensorEncodingAttrGet(
     MlirContext ctx, intptr_t numDimLevelTypes,
     MlirSparseTensorDimLevelType const *dimLevelTypes,
     MlirAffineMap dimOrdering, MlirAffineMap higherOrdering,
     int pointerBitWidth, int indexBitWidth) {
-  SmallVector<SparseTensorEncodingAttr::DimLevelType> cppDimLevelTypes;
+  SmallVector<DimLevelType> cppDimLevelTypes;
   cppDimLevelTypes.resize(numDimLevelTypes);
   for (intptr_t i = 0; i < numDimLevelTypes; ++i)
-    cppDimLevelTypes[i] =
-        static_cast<SparseTensorEncodingAttr::DimLevelType>(dimLevelTypes[i]);
+    cppDimLevelTypes[i] = static_cast<DimLevelType>(dimLevelTypes[i]);
   return wrap(SparseTensorEncodingAttr::get(
       unwrap(ctx), cppDimLevelTypes, unwrap(dimOrdering),
       unwrap(higherOrdering), pointerBitWidth, indexBitWidth));
 }
 
 MlirAffineMap mlirSparseTensorEncodingAttrGetDimOrdering(MlirAttribute attr) {
   return wrap(unwrap(attr).cast<SparseTensorEncodingAttr>().getDimOrdering());
 }

mlir/lib/Dialect/SparseTensor/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRSparseTensorDialect
   SparseTensorDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
   DEPENDS
   MLIRSparseTensorAttrDefsIncGen
   MLIRSparseTensorOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRDialect
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRSupport
+  MLIRSparseTensorRuntime
   )

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

if (failed(parser.parseLess()))
     return {};
   // Parse the data as a dictionary.
   DictionaryAttr dict;
   if (failed(parser.parseAttribute(dict)))
     return {};
   if (failed(parser.parseGreater()))
     return {};
   // Process the data from the parsed dictionary value into struct-like data.
-  SmallVector<SparseTensorEncodingAttr::DimLevelType, 4> dlt;
+  SmallVector<DimLevelType, 4> dlt;
   AffineMap dimOrd = {};
   AffineMap higherOrd = {};
   unsigned ptr = 0;
   unsigned ind = 0;
   for (const NamedAttribute &attr : dict) {
     if (attr.getName() == "dimLevelType") {
       auto arrayAttr = attr.getValue().dyn_cast<ArrayAttr>();
       if (!arrayAttr) {
         parser.emitError(parser.getNameLoc(),
                          "expected an array for dimension level types");
         return {};
       }
       for (auto i : arrayAttr) {
         auto strAttr = i.dyn_cast<StringAttr>();
         if (!strAttr) {
           parser.emitError(parser.getNameLoc(),
                            "expected a string value in dimension level types");
           return {};
         }
         auto strVal = strAttr.getValue();
         if (strVal == "dense") {
-          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::Dense);
+          dlt.push_back(DimLevelType::Dense);
         } else if (strVal == "compressed") {
-          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::Compressed);
+          dlt.push_back(DimLevelType::Compressed);
         } else if (strVal == "compressed-nu") {
-          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::CompressedNu);
+          dlt.push_back(DimLevelType::CompressedNu);
         } else if (strVal == "compressed-no") {
-          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::CompressedNo);
+          dlt.push_back(DimLevelType::CompressedNo);
         } else if (strVal == "compressed-nu-no") {
-          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::CompressedNuNo);
+          dlt.push_back(DimLevelType::CompressedNuNo);
         } else if (strVal == "singleton") {
-          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::Singleton);
+          dlt.push_back(DimLevelType::Singleton);
         } else if (strVal == "singleton-nu") {
-          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::SingletonNu);
+          dlt.push_back(DimLevelType::SingletonNu);
         } else if (strVal == "singleton-no") {
-          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::SingletonNo);
+          dlt.push_back(DimLevelType::SingletonNo);
         } else if (strVal == "singleton-nu-no") {
-          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::SingletonNuNo);
+          dlt.push_back(DimLevelType::SingletonNuNo);
         } else {
           parser.emitError(parser.getNameLoc(),
                            "unexpected dimension level type: ")
               << strVal;
           return {};
         }
       }
     } else if (attr.getName() == "dimOrdering") {

 
 SparseTensorEncodingAttr
 mlir::sparse_tensor::getSparseTensorEncoding(Type type) {
   if (auto ttp = type.dyn_cast<RankedTensorType>())
     return ttp.getEncoding().dyn_cast_or_null<SparseTensorEncodingAttr>();
   return nullptr;
 }
 
-bool mlir::sparse_tensor::isDenseDim(
-    SparseTensorEncodingAttr::DimLevelType dltp) {
-  return dltp == SparseTensorEncodingAttr::DimLevelType::Dense;
-}
-
-bool mlir::sparse_tensor::isCompressedDim(
-    SparseTensorEncodingAttr::DimLevelType dltp) {
-  switch (dltp) {
-  case SparseTensorEncodingAttr::DimLevelType::Compressed:
-  case SparseTensorEncodingAttr::DimLevelType::CompressedNu:
-  case SparseTensorEncodingAttr::DimLevelType::CompressedNo:
-  case SparseTensorEncodingAttr::DimLevelType::CompressedNuNo:
-    return true;
-  default:
-    return false;
-  }
-}
-
-bool mlir::sparse_tensor::isSingletonDim(
-    SparseTensorEncodingAttr::DimLevelType dltp) {
-  switch (dltp) {
-  case SparseTensorEncodingAttr::DimLevelType::Singleton:
-  case SparseTensorEncodingAttr::DimLevelType::SingletonNu:
-  case SparseTensorEncodingAttr::DimLevelType::SingletonNo:
-  case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
-    return true;
-  default:
-    return false;
-  }
-}
-
-bool mlir::sparse_tensor::isDenseDim(RankedTensorType type, uint64_t d) {
-  assert(d < static_cast<uint64_t>(type.getRank()));
-  if (auto enc = getSparseTensorEncoding(type))
-    return isDenseDim(enc.getDimLevelType()[d]);
-  return true; // unannotated tensor is dense
-}
-
-bool mlir::sparse_tensor::isCompressedDim(RankedTensorType type, uint64_t d) {
-  assert(d < static_cast<uint64_t>(type.getRank()));
-  if (auto enc = getSparseTensorEncoding(type))
-    return isCompressedDim(enc.getDimLevelType()[d]);
-  return false; // unannotated tensor is dense
-}
-
-bool mlir::sparse_tensor::isSingletonDim(RankedTensorType type, uint64_t d) {
-  assert(d < static_cast<uint64_t>(type.getRank()));
-  if (auto enc = getSparseTensorEncoding(type))
-    return isSingletonDim(enc.getDimLevelType()[d]);
-  return false; // unannotated tensor is dense
-}
-
-bool mlir::sparse_tensor::isOrderedDim(
-    SparseTensorEncodingAttr::DimLevelType dltp) {
-  switch (dltp) {
-  case SparseTensorEncodingAttr::DimLevelType::CompressedNo:
-  case SparseTensorEncodingAttr::DimLevelType::CompressedNuNo:
-  case SparseTensorEncodingAttr::DimLevelType::SingletonNo:
-  case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
-    return false;
-  default:
-    return true;
-  }
-}
-
-bool mlir::sparse_tensor::isUniqueDim(
-    SparseTensorEncodingAttr::DimLevelType dltp) {
-  switch (dltp) {
-  case SparseTensorEncodingAttr::DimLevelType::CompressedNu:
-  case SparseTensorEncodingAttr::DimLevelType::CompressedNuNo:
-  case SparseTensorEncodingAttr::DimLevelType::SingletonNu:
-  case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
-    return false;
-  default:
-    return true;
-  }
-}
-
-bool mlir::sparse_tensor::isOrderedDim(RankedTensorType type, uint64_t d) {
-  assert(d < static_cast<uint64_t>(type.getRank()));
-  if (auto enc = getSparseTensorEncoding(type))
-    return isOrderedDim(enc.getDimLevelType()[d]);
-  return true; // unannotated tensor is dense (and thus ordered)
-}
-
-bool mlir::sparse_tensor::isUniqueDim(RankedTensorType type, uint64_t d) {
-  assert(d < static_cast<uint64_t>(type.getRank()));
-  if (auto enc = getSparseTensorEncoding(type))
-    return isUniqueDim(enc.getDimLevelType()[d]);
-  return true; // unannotated tensor is dense (and thus unique)
-}
-
 uint64_t mlir::sparse_tensor::toOrigDim(const SparseTensorEncodingAttr &enc,
                                         uint64_t d) {
   if (enc) {
     auto order = enc.getDimOrdering();
     if (order) {
       assert(order.isPermutation());
       return order.getDimPosition(d);
     }

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

private:
   /// Return false if tid[dim] is a dense dimension that does not need to be
   /// prepared (to be used by sparsification for needUniv).
   bool prepareLoopOverTensorAtDim(OpBuilder &builder, Location loc, size_t tid,
                                   size_t dim);
 
   /// Input (TODO: and output) tensors.
   std::vector<Value> tensors;
   /// The dim type array for each tensor.
-  std::vector<std::vector<SparseTensorEncodingAttr::DimLevelType>> dims;
+  std::vector<std::vector<DimLevelType>> dims;
   /// Sparse iteration information (by tensor and dim). These arrays
   /// are updated to remain current within the current loop.
   std::vector<std::vector<Value>> pidxs;
   std::vector<std::vector<Value>> coord;
   std::vector<std::vector<Value>> highs;
   /// Universal dense indices and upper bounds (by index). The sizes array is
   /// set once with the inferred dimension sizes.
   std::vector<std::vector<Value>> sizes;

 PrimaryType primaryTypeEncoding(Type elemTp);
 
 /// Convert PrimaryType to its function-name suffix.
 StringRef primaryTypeFunctionSuffix(PrimaryType pt);
 
 /// Converts a primary storage type to its function-name suffix.
 StringRef primaryTypeFunctionSuffix(Type elemTp);
 
-/// Converts the IR's dimension level type to its internal type-encoding.
-DimLevelType dimLevelTypeEncoding(SparseTensorEncodingAttr::DimLevelType dlt);
-
 //===----------------------------------------------------------------------===//
 // Misc code generators and utilities.
 //===----------------------------------------------------------------------===//
 
 /// Generates a 1-valued attribute of the given type.  This supports
 /// all the same types as `getZeroAttr`; however, unlike `getZeroAttr`,
 /// for unsupported types we raise `llvm_unreachable` rather than
 /// returning a null attribute.

 /// Generates a constant of the internal type-encoding for primary storage.
 inline Value constantPrimaryTypeEncoding(OpBuilder &builder, Location loc,
                                          Type elemTp) {
   return constantI32(builder, loc,
                      static_cast<uint32_t>(primaryTypeEncoding(elemTp)));
 }
 
 /// Generates a constant of the internal dimension level type encoding.
-inline Value
-constantDimLevelTypeEncoding(OpBuilder &builder, Location loc,
-                             SparseTensorEncodingAttr::DimLevelType dlt) {
-  return constantI8(builder, loc,
-                    static_cast<uint8_t>(dimLevelTypeEncoding(dlt)));
+inline Value constantDimLevelTypeEncoding(OpBuilder &builder, Location loc,
+                                          DimLevelType dlt) {
+  return constantI8(builder, loc, static_cast<uint8_t>(dlt));
 }
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

if (!rtp) // a scalar (0-dimension tensors)
       continue;
 
     auto rank = static_cast<size_t>(rtp.getRank());
     auto enc = getSparseTensorEncoding(rtp);
     if (enc)
       for (auto dimTp : enc.getDimLevelType())
         dims[i].push_back(dimTp);
     else
-      dims[i].assign(rank, SparseTensorEncodingAttr::DimLevelType::Dense);
+      dims[i].assign(rank, DimLevelType::Dense);
 
     // Initialize using empty value.
     pidxs[i].assign(rank, Value());
     coord[i].assign(rank, Value());
     highs[i].assign(rank, Value());
     sizes[i].assign(rank, Value());
     ptrBuffer[i].assign(rank, Value());
     idxBuffer[i].assign(rank, Value());
for (size_t t = 0, e = tensors.size(); t < e; t++) {
     auto enc = getSparseTensorEncoding(rtp);
     auto dynShape = {ShapedType::kDynamicSize};
     // Scan all dimensions of current tensor.
     for (int64_t d = 0; d < rank; d++) {
       // This should be called only once at beginning.
       assert(!ptrBuffer[t][d] && !idxBuffer[t][d] && !sizes[t][d] &&
              !highs[t][d]);
       // Handle sparse storage schemes.
-      if (isCompressedDim(dims[t][d])) {
+      if (isCompressedDLT(dims[t][d])) {
         auto ptrTp =
             MemRefType::get(dynShape, getPointerOverheadType(builder, enc));
         auto indTp =
             MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
         auto dim = builder.getIndexAttr(d);
         // Generate sparse primitives to obtains pointer and indices.
         ptrBuffer[t][d] = builder.create<ToPointersOp>(loc, ptrTp, tensor, dim);
         idxBuffer[t][d] = builder.create<ToIndicesOp>(loc, indTp, tensor, dim);
-      } else if (isSingletonDim(dims[t][d])) {
+      } else if (isSingletonDLT(dims[t][d])) {
         // Singleton dimension, fetch indices.
         auto indTp =
             MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
         auto dim = builder.getIndexAttr(d);
         idxBuffer[t][d] = builder.create<ToIndicesOp>(loc, indTp, tensor, dim);
       } else {
         // Dense dimension, nothing to fetch.
-        assert(isDenseDim(dims[t][d]));
+        assert(isDenseDLT(dims[t][d]));
       }
 
       // Find upper bound in current dimension.
       unsigned p = toOrigDim(enc, d);
       Value up = mlir::linalg::createOrFoldDimOp(builder, loc, tensor, p);
       sizes[t][d] = highs[t][d] = up;
     }
     // Perform the required bufferization. Dense inputs materialize

 Operation *SparseTensorLoopEmitter::enterLoopOverTensorAtDim(
     OpBuilder &builder, Location loc, size_t tid, size_t dim,
     ArrayRef<Value> reduc) {
   assert(dims[tid].size() > dim);
   // We can not re-enter the same level.
   assert(!coord[tid][dim]);
   Value step = constantIndex(builder, loc, 1);
   auto dimType = dims[tid][dim];
-  bool isSparse = isCompressedDim(dimType) || isSingletonDim(dimType);
-  assert(isDenseDim(dimType) || isCompressedDim(dimType) ||
-         isSingletonDim(dimType));
+  bool isSparse = isCompressedDLT(dimType) || isSingletonDLT(dimType);
+  assert(isDenseDLT(dimType) || isCompressedDLT(dimType) ||
+         isSingletonDLT(dimType));
 
   Value lo = isSparse ? pidxs[tid][dim] : constantIndex(builder, loc, 0);
   Value hi = highs[tid][dim];
 
   // TODO: support reduction.
   if (!reduc.empty())
     llvm_unreachable("TODO: not implemented yet");
 
bool SparseTensorLoopEmitter::prepareLoopOverTensorAtDim(OpBuilder &builder,
                                                          Location loc,
                                                          size_t tid,
                                                          size_t dim) {
   // TODO: generate loop iteration on output tensor based on the shape
   // instead of pointer/indices arrays.
   assert(dims[tid].size() > dim);
   auto dimType = dims[tid][dim];
 
-  if (isDenseDim(dimType))
+  if (isDenseDLT(dimType))
     return false;
 
   // Either the first dimension, or the previous dimension has been set.
   assert(dim == 0 || pidxs[tid][dim - 1]);
   Value c0 = constantIndex(builder, loc, 0);
   Value c1 = constantIndex(builder, loc, 1);
-  if (isCompressedDim(dimType)) {
+  if (isCompressedDLT(dimType)) {
     Value ptr = ptrBuffer[tid][dim];
 
     Value pLo = dim == 0 ? c0 : pidxs[tid][dim - 1];
     Value pHi = builder.create<arith::AddIOp>(loc, pLo, c1);
 
     pidxs[tid][dim] = genIndexLoad(builder, loc, ptr, pLo);
     highs[tid][dim] = genIndexLoad(builder, loc, ptr, pHi);
     return true;
   }
-  if (isSingletonDim(dimType)) {
+  if (isSingletonDLT(dimType)) {
     Value pLo = dim == 0 ? c0 : pidxs[tid][dim - 1];
     Value pHi = builder.create<arith::AddIOp>(loc, pLo, c1);
 
     pidxs[tid][dim] = pLo;
     highs[tid][dim] = pHi;
     return true;
   }
 
void SparseTensorLoopEmitter::exitCurrentLoop() {
   assert(loopInfo.tensors.size() == loopInfo.dims.size());
   for (auto info : llvm::zip(loopInfo.tensors, loopInfo.dims)) {
     auto tid = std::get<0>(info);
     auto dim = std::get<1>(info);
     assert(pidxs[tid][dim] && coord[tid][dim] && highs[tid][dim]);
     // Reset to null.
     pidxs[tid][dim] = Value();
     coord[tid][dim] = Value();
-    if (!isDenseDim(dims[tid][dim]))
+    if (!isDenseDLT(dims[tid][dim]))
       // Dense dimension, high is fixed.
       highs[tid][dim] = Value();
   }
   loopStack.pop_back();
 }
 
 //===----------------------------------------------------------------------===//
 // ExecutionEngine/SparseTensorUtils helper functions.
#undef CASE
   }
   llvm_unreachable("Unknown PrimaryType");
 }
 
 StringRef mlir::sparse_tensor::primaryTypeFunctionSuffix(Type elemTp) {
   return primaryTypeFunctionSuffix(primaryTypeEncoding(elemTp));
 }
 
-DimLevelType mlir::sparse_tensor::dimLevelTypeEncoding(
-    SparseTensorEncodingAttr::DimLevelType dlt) {
-  switch (dlt) {
-  case SparseTensorEncodingAttr::DimLevelType::Dense:
-    return DimLevelType::Dense;
-  case SparseTensorEncodingAttr::DimLevelType::Compressed:
-    return DimLevelType::Compressed;
-  case SparseTensorEncodingAttr::DimLevelType::CompressedNu:
-    return DimLevelType::CompressedNu;
-  case SparseTensorEncodingAttr::DimLevelType::CompressedNo:
-    return DimLevelType::CompressedNo;
-  case SparseTensorEncodingAttr::DimLevelType::CompressedNuNo:
-    return DimLevelType::CompressedNuNo;
-  case SparseTensorEncodingAttr::DimLevelType::Singleton:
-    return DimLevelType::Singleton;
-  case SparseTensorEncodingAttr::DimLevelType::SingletonNu:
-    return DimLevelType::SingletonNu;
-  case SparseTensorEncodingAttr::DimLevelType::SingletonNo:
-    return DimLevelType::SingletonNo;
-  case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
-    return DimLevelType::SingletonNuNo;
-  }
-  llvm_unreachable("Unknown SparseTensorEncodingAttr::DimLevelType");
-}
-
 //===----------------------------------------------------------------------===//
 // Misc code generators.
 //===----------------------------------------------------------------------===//
 
 mlir::Attribute mlir::sparse_tensor::getOneAttr(Builder &builder, Type tp) {
   if (tp.isa<FloatType>())
     return builder.getFloatAttr(tp, 1.0);
   if (tp.isa<IndexType>())

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

 
 /// Populates parameters required to call the "swiss army knife" method of the
 /// sparse runtime support library for materializing sparse tensors into the
 /// computation.
 static void newParams(OpBuilder &builder, SmallVector<Value, 8> &params,
                       Location loc, ShapedType stp,
                       SparseTensorEncodingAttr &enc, Action action,
                       ValueRange szs, Value ptr = Value()) {
-  ArrayRef<SparseTensorEncodingAttr::DimLevelType> dlt = enc.getDimLevelType();
+  ArrayRef<DimLevelType> dlt = enc.getDimLevelType();
   unsigned sz = dlt.size();
   // Sparsity annotations.
   SmallVector<Value, 4> attrs;
   for (unsigned i = 0; i < sz; i++)
     attrs.push_back(constantDimLevelTypeEncoding(builder, loc, dlt[i]));
   params.push_back(genBuffer(builder, loc, attrs));
   // Dimension sizes array of the enveloping tensor. Useful for either
   // verification of external data, or for construction of internal data.
static void insertScalarIntoDenseTensor(OpBuilder &builder, Location loc,
                                         Value elemPtr, Value tensor,
                                         ValueRange ivs) {
   Value elemV = builder.create<memref::LoadOp>(loc, elemPtr);
   builder.create<memref::StoreOp>(loc, elemV, tensor, ivs);
 }
 
 /// Determine if the runtime library supports direct conversion to the
 /// given target `dimTypes`.
-static bool canUseDirectConversion(
-    ArrayRef<SparseTensorEncodingAttr::DimLevelType> dimTypes) {
+static bool canUseDirectConversion(ArrayRef<DimLevelType> dimTypes) {
   bool alreadyCompressed = false;
   for (uint64_t rank = dimTypes.size(), r = 0; r < rank; r++) {
-    const DimLevelType dlt = dimLevelTypeEncoding(dimTypes[r]);
+    const DimLevelType dlt = dimTypes[r];
     if (isCompressedDLT(dlt)) {
       if (alreadyCompressed)
         return false; // Multiple compressed dimensions not yet supported.
       alreadyCompressed = true;
     } else if (isDenseDLT(dlt)) {
       if (alreadyCompressed)
         return false; // Dense after Compressed not yet supported.
     } else if (isSingletonDLT(dlt)) {
if (!encDst && encSrc) {
       RankedTensorType srcTensorTp = srcType.cast<RankedTensorType>();
       unsigned rank = dstTensorTp.getRank();
       Type elemTp = dstTensorTp.getElementType();
       // Fabricate a no-permutation encoding for newParams().
       // The pointer/index types must be those of `src`.
       // The dimLevelTypes aren't actually used by Action::kToIterator.
       encDst = SparseTensorEncodingAttr::get(
           op->getContext(),
-          SmallVector<SparseTensorEncodingAttr::DimLevelType>(
-              rank, SparseTensorEncodingAttr::DimLevelType::Dense),
-          AffineMap(), AffineMap(), encSrc.getPointerBitWidth(),
-          encSrc.getIndexBitWidth());
+          SmallVector<DimLevelType>(rank, DimLevelType::Dense), AffineMap(),
+          AffineMap(), encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
       SmallVector<Value, 4> sizes;
       SmallVector<Value, 8> params;
       sizesFromPtr(rewriter, sizes, loc, encSrc, srcTensorTp, src);
       newParams(rewriter, params, loc, dstTensorTp, encDst, Action::kToIterator,
                 sizes, src);
       Value iter = genNewCall(rewriter, loc, params);
       Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
       Value elemPtr = genAllocaScalar(rewriter, loc, elemTp);

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

 // Helper method to match any typed zero.
 static bool isZeroValue(Value val) {
   return matchPattern(val, m_Zero()) || matchPattern(val, m_AnyZeroFloat());
 }
 
 // Helper to detect a sparse tensor type operand.
 static bool isSparseTensor(OpOperand *op) {
   if (auto enc = getSparseTensorEncoding(op->get().getType())) {
-    if (llvm::is_contained(enc.getDimLevelType(),
-                           SparseTensorEncodingAttr::DimLevelType::Compressed))
+    if (llvm::is_contained(enc.getDimLevelType(), DimLevelType::Compressed))
       return true;
   }
   return false;
 }
 
 // Helper method to find zero/uninitialized allocation.
 static bool isAlloc(OpOperand *op, bool isZero) {
   Value val = op->get();

 }
 
 // TODO: The dim level property of the COO type relies on input tensors, the
 // shape relies on the output tensor
 // Helpers to setup a COO type.
 static RankedTensorType getUnorderedCOOFromType(RankedTensorType src) {
   auto *ctx = src.getContext();
   auto rank = src.getRank();
-  SmallVector<SparseTensorEncodingAttr::DimLevelType, 4> dims;
+  SmallVector<DimLevelType, 4> dims;
 
   // An unordered and non-unique compressed dim at beginning unless the tensor
   // is a 1D tensor.
   if (rank > 1)
-    dims.push_back(SparseTensorEncodingAttr::DimLevelType::CompressedNuNo);
+    dims.push_back(DimLevelType::CompressedNuNo);
 
   // TODO: it is actually ordered at the level for ordered input.
   // Followed by unordered non-unique n-2 singleton levels.
-  std::fill_n(std::back_inserter(dims), rank - 2,
-              SparseTensorEncodingAttr::DimLevelType::SingletonNuNo);
+  std::fill_n(std::back_inserter(dims), rank - 2, DimLevelType::SingletonNuNo);
   // TODO: only if all the inputs (for concatentate) are unique at the last
   // level should the COO has a unique level at the end. Ends by a unordered
   // unique singleton level.
-  dims.push_back(SparseTensorEncodingAttr::DimLevelType::SingletonNo);
+  dims.push_back(DimLevelType::SingletonNo);
   SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(src);
   // TODO: Maybe pick the bitwidth based on input/output tensors (probably the
   // largest one among them) in the original operation instead of using the
   // default value.
   auto enc = SparseTensorEncodingAttr::get(
       ctx, dims, AffineMap::getMultiDimIdentityMap(rank, ctx), AffineMap(),
       encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
   return RankedTensorType::get(src.getShape(), src.getElementType(), enc);

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

 /// Helper method to obtain the dimension level format from the encoding.
 //
 //  TODO: note that we store, but currently completely *ignore* the properties
 //
 static DimLevelFormat toDimLevelFormat(const SparseTensorEncodingAttr &enc,
                                        unsigned d) {
   if (enc) {
     switch (enc.getDimLevelType()[d]) {
-    case SparseTensorEncodingAttr::DimLevelType::Dense:
+    case DimLevelType::Dense:
       return DimLevelFormat(DimLvlType::kDense);
-    case SparseTensorEncodingAttr::DimLevelType::Compressed:
+    case DimLevelType::Compressed:
       return DimLevelFormat(DimLvlType::kCompressed);
-    case SparseTensorEncodingAttr::DimLevelType::CompressedNu:
+    case DimLevelType::CompressedNu:
       return DimLevelFormat(DimLvlType::kCompressed, true, false);
-    case SparseTensorEncodingAttr::DimLevelType::CompressedNo:
+    case DimLevelType::CompressedNo:
       return DimLevelFormat(DimLvlType::kCompressed, false, true);
-    case SparseTensorEncodingAttr::DimLevelType::CompressedNuNo:
+    case DimLevelType::CompressedNuNo:
       return DimLevelFormat(DimLvlType::kCompressed, false, false);
-    case SparseTensorEncodingAttr::DimLevelType::Singleton:
+    case DimLevelType::Singleton:
       return DimLevelFormat(DimLvlType::kSingleton);
-    case SparseTensorEncodingAttr::DimLevelType::SingletonNu:
+    case DimLevelType::SingletonNu:
       return DimLevelFormat(DimLvlType::kSingleton, true, false);
-    case SparseTensorEncodingAttr::DimLevelType::SingletonNo:
+    case DimLevelType::SingletonNo:
       return DimLevelFormat(DimLvlType::kSingleton, false, true);
-    case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
+    case DimLevelType::SingletonNuNo:
       return DimLevelFormat(DimLvlType::kSingleton, false, false);
     }
   }
   return DimLevelFormat(DimLvlType::kDense);
 }
 
 /// Helper method to inspect affine expressions. Rejects cases where the
 /// same index is used more than once. Also rejects compound affine

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     hdrs = ["include/mlir/Dialect/SparseTensor/IR/SparseTensor.h"],
     includes = ["include"],
     deps = [
         ":ArithDialect",
         ":IR",
         ":InferTypeOpInterface",
         ":SparseTensorAttrDefsIncGen",
         ":SparseTensorOpsIncGen",
+        ":SparseTensorRuntime",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "SparseTensorUtils",
     srcs = glob(["lib/Dialect/SparseTensor/Utils/*.cpp"]),
     hdrs = glob(["include/mlir/Dialect/SparseTensor/Utils/*.h"]),