[mlir][sparse] clean up a few TODOs and layout

[aartbik]

Summary

Removed some TODOs that are not likely to be done
or are already recorded through git tracker.

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

 add_mlir_dialect(SparseTensorOps sparse_tensor)
 add_mlir_doc(SparseTensorOps SparseTensorOps Dialects/ -gen-dialect-doc)
 
 set(LLVM_TARGET_DEFINITIONS SparseTensorAttrDefs.td)
 mlir_tablegen(SparseTensorAttrDefs.h.inc -gen-attrdef-decls)
 mlir_tablegen(SparseTensorAttrDefs.cpp.inc -gen-attrdef-defs)
 mlir_tablegen(SparseTensorAttrEnums.h.inc -gen-enum-decls)
 mlir_tablegen(SparseTensorAttrEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRSparseTensorAttrDefsIncGen)
 
 set(LLVM_TARGET_DEFINITIONS SparseTensorTypes.td)
 mlir_tablegen(SparseTensorTypes.h.inc -gen-typedef-decls)
 mlir_tablegen(SparseTensorTypes.cpp.inc -gen-typedef-defs)
 add_public_tablegen_target(MLIRSparseTensorTypesIncGen)
-

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

 
 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" bitwidths should
 /// link with an alternatively built runtime support library.
-// TODO: support such targets?
 using index_type = uint64_t;
 
 /// Encoding of overhead types (both position overhead and coordinate
 /// overhead), for "overloading" @newSparseTensor.
 enum class OverheadType : uint32_t {
   kIndex = 0,
   kU64 = 1,
   kU32 = 2,
enum class PrimaryType : uint32_t {
   kI32 = 6,
   kI16 = 7,
   kI8 = 8,
   kC64 = 9,
   kC32 = 10
 };
 
 // This x-macro includes all `V` types.
-// TODO: We currently split out the non-variadic version from the variadic
-// version. Using ##__VA_ARGS__ to avoid the split gives
-//   warning: token pasting of ',' and __VA_ARGS__ is a GNU extension
-//   [-Wgnu-zero-variadic-macro-arguments]
-// and __VA_OPT__(, ) __VA_ARGS__ requires c++20.
 #define MLIR_SPARSETENSOR_FOREVERY_V(DO)                                       \
   DO(F64, double)                                                              \
   DO(F32, float)                                                               \
   DO(F16, f16)                                                                 \
   DO(BF16, bf16)                                                               \
   DO(I64, int64_t)                                                             \
   DO(I32, int32_t)                                                             \
   DO(I16, int16_t)                                                             \
enum class LevelFormat : uint8_t {
   Singleton = 16,        // 0b00100_00
   CompressedWithHi = 32, // 0b01000_00
   TwoOutOfFour = 64,     // 0b10000_00
 };
 
 /// Returns string representation of the given dimension level type.
 constexpr const char *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";
   case DimLevelType::Compressed:
     return "compressed";
   case DimLevelType::CompressedNu:
     return "compressed-nu";

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

- An optional affine map from dimension-coordinates to level-coordinates;
       defaulting to the identity map.  For example, given a 2-d tensor:
       `(i, j) -> (i, j)` specifies row-wise storage, `(i, j) -> (j, i)`
       specifies column-wise storage, and
       `(i, j) -> (i floordiv 2, j floordiv 3, i mod 2, j mod 3)`
       specifies 2x3 block-sparsity.  For block-sparsity, blocks are typically
       stored with compression while dense storage is used within each block
       (although hybrid schemes are possible as well).
 
-      TODO: the following example is out-of-date and will be implemented
-      in a different manner than described here.
-      (This will be corrected in an upcoming change that completely
+      (The following will be corrected in an upcoming change that completely
       overhauls the syntax of this attribute.)
 
       The dimToLvl mapping also provides a notion of "counting a
       dimension", where every stored element with the same coordinate
       is mapped to a new slice.  For instance, ELL storage of a 2-d
       tensor can be defined with the mapping `(i, j) -> (#i, i, j)`
       using the notation of [Chou20].  Lacking the `#` symbol in MLIR's
       affine mapping, we use a free symbol `c` to define such counting,
class RankedSparseTensorOf<list<Type> allowedTypes>
   : RankedTensorOf<allowedTypes, [IsSparseTensorPred], "ranked sparse tensor">;
 
 def AnyRankedSparseTensor : RankedSparseTensorOf<[AnyType]>;
 
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Sorting Algorithm Attribute.
 //===----------------------------------------------------------------------===//
 
-// TODO: Currently, we only provide four implementations, and expose the
-// implementations via attribute algorithm. In the future, if we will need
-// to support both stable and non-stable quick sort, we may add
-// quick_sort_nonstable enum to the attribute. Alternative, we may use two
-// attributes, (stable|nonstable, algorithm), to specify a sorting
+// Currently, we only provide four implementations, and expose the
+// implementations via attribute algorithm. In the future, if we will
+// need to support both stable and non-stable quick sort, we may add
+// quick_sort_nonstable enum to the attribute. Alternative, we may use
+// two attributes, (stable|nonstable, algorithm), to specify a sorting
 // implementation.
 //
 // --------------------------------------------------------------------------
 // |           | hybrid_qsort| insertion_sort | qsort       | heap_sort.    |
 // |non-stable | Impl        | X              |  Impl       | Impl          |
 // |stable     | X           | Impl           |  Not Impl   | X             |
 // --------------------------------------------------------------------------
 

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

def SparseTensor_OutOp : SparseTensor_Op<"out", []>,
   let assemblyFormat = "$tensor `,` $dest attr-dict `:` type($tensor) `,` type($dest)";
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Sorting Operations.
 //===----------------------------------------------------------------------===//
 
 def SparseTensor_SortOp : SparseTensor_Op<"sort", [AttrSizedOperandSegments]>,
-    // TODO: May want to extend tablegen with
-    // class NonemptyVariadic<Type type> : Variadic<type> { let minSize = 1; }
-    // and then use NonemptyVariadic<...>:$xs here.
-    //
-    // TODO: Currently tablegen doesn't support the assembly syntax when
-    // `algorithm` is an optional enum attribute. We may want to use an optional
-    // enum attribute when this is fixed in tablegen.
-    //
     Arguments<(ins Index:$n,
                Variadic<StridedMemRefRankOf<[AnyInteger, Index], [1]>>:$xs,
                Variadic<StridedMemRefRankOf<[AnyType], [1]>>:$ys,
                SparseTensorSortKindAttr:$algorithm)>  {
   string summary = "Sorts the arrays in xs and ys lexicographically on the "
                    "integral values found in the xs list";
   string description = [{
     Lexicographically sort the first `n` values in `xs` along with the values in

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

 /// self-documenting; unfortunately it is not type-checked, so it only
 /// provides documentation rather than doing anything to prevent mixups.
 using FieldIndex = unsigned;
 
 /// Provides methods to access fields of a sparse tensor with the given
 /// encoding.
 class StorageLayout {
 public:
-  // TODO: Functions/methods marked with [NUMFIELDS] might should use
+  // TODO: Functions/methods marked with [NUMFIELDS] should use
   // `FieldIndex` for their return type, via the same reasoning for why
   // `Dimension`/`Level` are used both for identifiers and ranks.
   explicit StorageLayout(const SparseTensorType &stt)
       : StorageLayout(stt.getEncoding()) {}
   explicit StorageLayout(SparseTensorEncodingAttr enc) : enc(enc) {
     assert(enc);
   }
 

 private:
   const SparseTensorEncodingAttr enc;
 };
 
 //
 // Wrapper functions to invoke StorageLayout-related method.
 //
 
-// TODO: See note [NUMFIELDS].
+// See note [NUMFIELDS].
 inline unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc) {
   return StorageLayout(enc).getNumFields();
 }
 
-// TODO: See note [NUMFIELDS].
+// See note [NUMFIELDS].
 inline unsigned getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc) {
   return StorageLayout(enc).getNumDataFields();
 }
 
 inline void foreachFieldInSparseTensor(
     SparseTensorEncodingAttr enc,
     llvm::function_ref<bool(FieldIndex, SparseTensorFieldKind, Level,
                             DimLevelType)>

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

 /// outermost loop only), and (2) the generated code is an actual for-loop
 /// (and not a co-iterating while-loop).
 enum class SparseParallelizationStrategy {
   kNone,
   kDenseOuterLoop,
   kAnyStorageOuterLoop,
   kDenseAnyLoop,
   kAnyStorageAnyLoop
-  // TODO: support reduction parallelization too?
 };
 
 // TODO : Zero copy is disabled due to correctness bugs.Tracker #64316
 enum class GPUDataTransferStrategy { kRegularDMA, kZeroCopy, kPinnedDMA };
 
 #define GEN_PASS_DECL
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"