diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td --- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td @@ -12,7 +12,45 @@ include "mlir/Pass/PassBase.td" def Sparsification : Pass<"sparsification", "ModuleOp"> { - let summary = "Automatically generate sparse tensor code from annotations"; + let summary = "Automatically generate sparse tensor code from sparse tensor types"; + let description = [{ + A pass that implements the core functionality of a **sparse compiler**. + Each Linalg operation (MLIR's tensor index notation) that operates on + sparse tensor types is converted into code in which the sparsity is + explicit both in terms of co-iterating looping logic as well as + selected sparse storage schemes. + + See the `SparseTensor` dialect documentation for more background. + + Example input: + + ```mlir + #matvec = { + indexing_maps = [ + affine_map<(i,j) -> (i,j)>, // A + affine_map<(i,j) -> (j)>, // b + affine_map<(i,j) -> (i)> // x (out) + ], + iterator_types = ["parallel", "reduction"], + doc = "X(i) += A(i,j) * B(j)" + } + + // Multiply a sparse matrix A with a dense vector b into a dense vector x. + func @kernel_matvec(%arga: tensor, + %argb: tensor, + %argx: tensor) -> tensor { + %0 = linalg.generic #matvec + ins(%arga, %argb: tensor, tensor) + outs(%argx: tensor) { + ^bb(%a: f64, %b: f64, %x: f64): + %0 = mulf %a, %b : f64 + %1 = addf %x, %0 : f64 + linalg.yield %1 : f64 + } -> tensor + return %0 : tensor + } + ``` + }]; let constructor = "mlir::createSparsificationPass()"; let dependentDialects = [ "LLVM::LLVMDialect", @@ -34,7 +72,34 @@ } def SparseTensorConversion : Pass<"sparse-tensor-conversion", "ModuleOp"> { - let summary = "Apply conversion rules to sparse tensors"; + let summary = "Apply conversion rules to sparse tensor primitives and types"; + let description = [{ + A pass that converts sparse tensor primitives to calls into a runtime + support library. All sparse tensor types are converted into opaque + pointers to the underlying sparse storage schemes. + + Note that this is a current implementation choice to keep the conversion + relatively simple. In principle, these primitives could also be + converted to actual elaborate IR code that implements the primitives + on the selected sparse tensor storage schemes. + + Example of the conversion: + + ```mlir + Before: + %c1 = constant 1 : index + %0 = sparse_tensor.pointers %arg0, %c1 + : tensor<8x8xf32, #sparse_tensor.encoding<{ + dimLevelType = [ "dense", "compressed" ], + pointerBitWidth = 0, + indexBitWidth = 0 + }>> to memref + + After: + %c1 = constant 1 : index + %0 = call @sparsePointers(%arg0, %c1) : (!llvm.ptr, index) -> memref + ``` + }]; let constructor = "mlir::createSparseTensorConversionPass()"; let dependentDialects = [ "LLVM::LLVMDialect", @@ -46,4 +111,3 @@ } #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES -