added docs

Author: yinying-lisa-li

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

@@ -164,7 +164,13 @@ inline bool hasAnySparseOperandOrResult(Operation *op) {
 // Inference.
 //
 
+/// Given the dimToLvl map, infers the lvlToDim map, or returns
+/// empty Affine map when inference fails.
 AffineMap inferLvlToDim(AffineMap dimToLvl, MLIRContext *context);
+
+/// Returns the lvlToDim map for the given dimToLvl map specific
+/// to the block sparse cases.
+/// Asserts on failure (so only use when known to succeed).
 AffineMap inverseBlockSparsity(AffineMap dimToLvl, MLIRContext *context);
 
 //

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

@@ -759,6 +759,7 @@ AffineMap mlir::sparse_tensor::inferLvlToDim(AffineMap dimToLvl,
   } else if (map.isPermutation()) {
     lvlToDim = inversePermutation(map);
   } else {
+    // TODO: check if it's block sparsity
     lvlToDim = inverseBlockSparsity(map, context);
   }
   return lvlToDim;