added docs

llvm · Oct 17, 2023 · 808d878 · 808d878
1 parent 8d28731
commit 808d878
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h b/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);
 
 //

diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/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;