[mlir][sparse] skip sparsification for unannotated (or unhandled) cases

[aartbik]

Summary

Skip the sparsification pass for Linalg ops without annotated tensors
(or cases that are not properly handled yet).

[bixia]

Shall we add tests to check that we skip certain cases? Rather than this, it looks good to me.

[aartbik]

In D102787#2769444, @bixia wrote:

Shall we add tests to check that we skip certain cases? Rather than this, it looks good to me.

Thanks!

I am actually already working on accepting the cases, but just wanted to make sure external projects are not confused until that time ;-)

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

if (enc) {
     if (tp == SparseTensorEncodingAttr::DimLevelType::Singleton)
       return Dim::kSingle;
   }
   return Dim::kDense;
 }
 
 /// Helper method to inspect sparse encodings in the tensor types.
 /// Fills the per-dimension sparsity information for all tensors.
-static void findSparseAnnotations(Merger &merger, linalg::GenericOp op) {
+static bool findSparseAnnotations(Merger &merger, linalg::GenericOp op) {
+  bool annotated = false;
   unsigned numTensors = op.getNumShapedOperands();
+  unsigned lhs = numTensors - 1;
   for (unsigned t = 0; t < numTensors; t++) {
     auto map = op.getIndexingMap(t);
     unsigned rank = op.getShapedType(t).getRank();
     auto enc = getSparseTensorEncoding(op.getShapedType(t));
+    if (enc) {
+      annotated = true;
+      if (enc.getDimOrdering() && !enc.getDimOrdering().isIdentity())
+        return false; // TODO: handle permutations
+      if (t == lhs)
+        return false; // TODO: handle sparse outputs
+    }
     for (unsigned d = 0; d < rank; d++) {
       unsigned idx = map.getDimPosition(d);
       merger.setDim(t, idx, toDim(enc, d));
     }
   }
+  return annotated;
 }
 
 /// A DFS helper to compute a topological sort. Note that recursion is
 /// bounded by the number of implicit loops, which is always small.
 /// Returns false when a cycle is detected.
 static bool topSortDFS(unsigned i, std::vector<unsigned> &visit,
                        std::vector<unsigned> &topSort,
                        std::vector<std::vector<bool>> &adjM) {
public:
   LogicalResult matchAndRewrite(linalg::GenericOp op,
                                 PatternRewriter &rewriter) const override {
     // Detects sparse annotations and translate the per-dimension sparsity
     // information for all tensors to loop indices in the kernel.
     assert(op.getNumOutputs() == 1);
     unsigned numTensors = op.getNumShapedOperands();
     unsigned numLoops = op.iterator_types().getValue().size();
     Merger merger(numTensors, numLoops);
-    findSparseAnnotations(merger, op);
+    if (!findSparseAnnotations(merger, op))
+      return failure();
 
     // Computes a topologically sorted iteration graph to ensure
     // tensors are visited in natural index order. Fails on cycles.
     // This assumes that higher-level passes have already put the
     // tensors in each tensor expression in a feasible order.
     std::vector<unsigned> topSort;
     if (!computeIterationGraph(merger, op, topSort, /*sparseOnly=*/false) &&
         !computeIterationGraph(merger, op, topSort, /*sparseOnly=*/true))

mlir/test/Dialect/SparseTensor/sparse_3d.mlir

 // NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
 // RUN: mlir-opt %s -sparsification | FileCheck %s
 
+#Td = #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>
+
 #Tddd = #sparse_tensor.encoding<{ dimLevelType = [ "dense",      "dense",      "dense"      ] }>
 #Tdds = #sparse_tensor.encoding<{ dimLevelType = [ "dense",      "dense",      "compressed" ] }>
 #Tdsd = #sparse_tensor.encoding<{ dimLevelType = [ "dense",      "compressed", "dense"      ] }>
 #Tdss = #sparse_tensor.encoding<{ dimLevelType = [ "dense",      "compressed", "compressed" ] }>
 #Tsdd = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense",      "dense"      ] }>
 #Tsds = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense",      "compressed" ] }>
 #Tssd = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense"      ] }>
 #Tsss = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>
indexing_maps = [
     affine_map<(i,j,k) -> ()>        // x (scalar out)
   ],
   iterator_types = ["reduction", "reduction", "reduction"],
   doc = "x += SUM_i A(i,j,k) * b(i)"
 }
 
 // CHECK-LABEL:   func @sum_reduction_inv(
 // CHECK-SAME:                            %[[VAL_0:.*]]: tensor<?x?x?xf32>,
-// CHECK-SAME:                            %[[VAL_1:.*]]: tensor<?xf32>,
+// CHECK-SAME:                            %[[VAL_1:.*]]: tensor<?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ], pointerBitWidth = 0, indexBitWidth = 0 }>>,
 // CHECK-SAME:                            %[[VAL_2:.*]]: tensor<f32>) -> tensor<f32> {
 // CHECK:           %[[VAL_3:.*]] = constant 2 : index
 // CHECK:           %[[VAL_4:.*]] = constant 0 : index
 // CHECK:           %[[VAL_5:.*]] = constant 1 : index
 // CHECK:           %[[VAL_6:.*]] = memref.dim %[[VAL_0]], %[[VAL_5]] : tensor<?x?x?xf32>
 // CHECK:           %[[VAL_7:.*]] = memref.dim %[[VAL_0]], %[[VAL_3]] : tensor<?x?x?xf32>
 // CHECK:           %[[VAL_8:.*]] = memref.buffer_cast %[[VAL_0]] : memref<?x?x?xf32>
-// CHECK:           %[[VAL_9:.*]] = memref.dim %[[VAL_1]], %[[VAL_4]] : tensor<?xf32>
-// CHECK:           %[[VAL_10:.*]] = memref.buffer_cast %[[VAL_1]] : memref<?xf32>
+// CHECK:           %[[VAL_9:.*]] = memref.dim %[[VAL_1]], %[[VAL_4]] : tensor<?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ], pointerBitWidth = 0, indexBitWidth = 0 }>>
+// CHECK:           %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ], pointerBitWidth = 0, indexBitWidth = 0 }>> to memref<?xf32>
 // CHECK:           %[[VAL_11:.*]] = memref.buffer_cast %[[VAL_2]] : memref<f32>
 // CHECK:           %[[VAL_12:.*]] = memref.alloc() : memref<f32>
 // CHECK:           linalg.copy(%[[VAL_11]], %[[VAL_12]]) : memref<f32>, memref<f32>
 // CHECK:           scf.for %[[VAL_13:.*]] = %[[VAL_4]] to %[[VAL_9]] step %[[VAL_5]] {
 // CHECK:             %[[VAL_14:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_13]]] : memref<?xf32>
 // CHECK:             scf.for %[[VAL_15:.*]] = %[[VAL_4]] to %[[VAL_6]] step %[[VAL_5]] {
 // CHECK:               %[[VAL_16:.*]] = memref.load %[[VAL_12]][] : memref<f32>
 // CHECK:               %[[VAL_17:.*]] = scf.for %[[VAL_18:.*]] = %[[VAL_4]] to %[[VAL_7]] step %[[VAL_5]] iter_args(%[[VAL_19:.*]] = %[[VAL_16]]) -> (f32) {
 // CHECK:                 %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_13]], %[[VAL_15]], %[[VAL_18]]] : memref<?x?x?xf32>
 // CHECK:                 %[[VAL_21:.*]] = mulf %[[VAL_20]], %[[VAL_14]] : f32
 // CHECK:                 %[[VAL_22:.*]] = addf %[[VAL_19]], %[[VAL_21]] : f32
 // CHECK:                 scf.yield %[[VAL_22]] : f32
 // CHECK:               }
 // CHECK:               memref.store %[[VAL_23:.*]], %[[VAL_12]][] : memref<f32>
 // CHECK:             }
 // CHECK:           }
 // CHECK:           %[[VAL_24:.*]] = memref.tensor_load %[[VAL_12]] : memref<f32>
 // CHECK:           return %[[VAL_24]] : tensor<f32>
 // CHECK:         }
 func @sum_reduction_inv(%arga: tensor<?x?x?xf32>,
-                        %argb: tensor<?xf32>,
+                        %argb: tensor<?xf32, #Td>,
             %argx: tensor<f32>) -> tensor<f32> {
   %0 = linalg.generic #trait_sum_reduction_inv
-    ins(%arga, %argb: tensor<?x?x?xf32>, tensor<?xf32>)
+    ins(%arga, %argb: tensor<?x?x?xf32>, tensor<?xf32, #Td>)
     outs(%argx: tensor<f32>) {
       ^bb(%a: f32, %b: f32, %x: f32):
         %0 = mulf %a, %b : f32
         %1 = addf %x, %0 : f32
         linalg.yield %1 : f32
   } -> tensor<f32>
   return %0 : tensor<f32>
 }
 
 #trait_invariants = {
   indexing_maps = [
     affine_map<(i,j,k) -> (i)>,      // a
     affine_map<(i,j,k) -> (j)>,      // b
     affine_map<(i,j,k) -> (k)>,      // c
     affine_map<(i,j,k) -> (i,j,k)>   // X (out)
   ],
   iterator_types = ["parallel", "parallel", "parallel"],
   doc = "X(i,j,k) = a(i) * b(j) * c(k)"
 }
 
 // CHECK-LABEL:   func @invariants(
-// CHECK-SAME:                     %[[VAL_0:.*]]: tensor<10xf32>,
+// CHECK-SAME:                     %[[VAL_0:.*]]: tensor<10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ], pointerBitWidth = 0, indexBitWidth = 0 }>>,
 // CHECK-SAME:                     %[[VAL_1:.*]]: tensor<20xf32>,
 // CHECK-SAME:                     %[[VAL_2:.*]]: tensor<30xf32>,
 // CHECK-SAME:                     %[[VAL_3:.*]]: tensor<10x20x30xf32>) -> tensor<10x20x30xf32> {
 // CHECK:           %[[VAL_4:.*]] = constant 10 : index
 // CHECK:           %[[VAL_5:.*]] = constant 20 : index
 // CHECK:           %[[VAL_6:.*]] = constant 30 : index
 // CHECK:           %[[VAL_7:.*]] = constant 0 : index
 // CHECK:           %[[VAL_8:.*]] = constant 1 : index
-// CHECK:           %[[VAL_9:.*]] = memref.buffer_cast %[[VAL_0]] : memref<10xf32>
+// CHECK:           %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ], pointerBitWidth = 0, indexBitWidth = 0 }>> to memref<?xf32>
 // CHECK:           %[[VAL_10:.*]] = memref.buffer_cast %[[VAL_1]] : memref<20xf32>
 // CHECK:           %[[VAL_11:.*]] = memref.buffer_cast %[[VAL_2]] : memref<30xf32>
 // CHECK:           %[[VAL_12:.*]] = memref.buffer_cast %[[VAL_3]] : memref<10x20x30xf32>
 // CHECK:           %[[VAL_13:.*]] = memref.alloc() : memref<10x20x30xf32>
 // CHECK:           linalg.copy(%[[VAL_12]], %[[VAL_13]]) : memref<10x20x30xf32>, memref<10x20x30xf32>
 // CHECK:           scf.for %[[VAL_14:.*]] = %[[VAL_7]] to %[[VAL_4]] step %[[VAL_8]] {
-// CHECK:             %[[VAL_15:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_14]]] : memref<10xf32>
+// CHECK:             %[[VAL_15:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_14]]] : memref<?xf32>
 // CHECK:             scf.for %[[VAL_16:.*]] = %[[VAL_7]] to %[[VAL_5]] step %[[VAL_8]] {
 // CHECK:               %[[VAL_17:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_16]]] : memref<20xf32>
 // CHECK:               scf.for %[[VAL_18:.*]] = %[[VAL_7]] to %[[VAL_6]] step %[[VAL_8]] {
 // CHECK:                 %[[VAL_19:.*]] = mulf %[[VAL_15]], %[[VAL_17]] : f32
 // CHECK:                 %[[VAL_20:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_18]]] : memref<30xf32>
 // CHECK:                 %[[VAL_21:.*]] = mulf %[[VAL_19]], %[[VAL_20]] : f32
 // CHECK:                 memref.store %[[VAL_21]], %[[VAL_13]]{{\[}}%[[VAL_14]], %[[VAL_16]], %[[VAL_18]]] : memref<10x20x30xf32>
 // CHECK:               }
 // CHECK:             }
 // CHECK:           }
 // CHECK:           %[[VAL_22:.*]] = memref.tensor_load %[[VAL_13]] : memref<10x20x30xf32>
 // CHECK:           return %[[VAL_22]] : tensor<10x20x30xf32>
 // CHECK:         }
-func @invariants(%arga: tensor<10xf32>,
+func @invariants(%arga: tensor<10xf32, #Td>,
                  %argb: tensor<20xf32>,
                  %argc: tensor<30xf32>,
                  %argx: tensor<10x20x30xf32>) -> tensor<10x20x30xf32> {
   %0 = linalg.generic #trait_invariants
-     ins(%arga, %argb, %argc : tensor<10xf32>, tensor<20xf32>, tensor<30xf32>)
+     ins(%arga, %argb, %argc : tensor<10xf32, #Td>, tensor<20xf32>, tensor<30xf32>)
     outs(%argx: tensor<10x20x30xf32>) {
       ^bb(%a: f32, %b: f32, %c: f32, %x: f32):
         %0 = mulf %a, %b : f32
         %1 = mulf %0, %c : f32
         linalg.yield %1 : f32
   } -> tensor<10x20x30xf32>
   return %0 : tensor<10x20x30xf32>
 }

mlir/test/Dialect/SparseTensor/sparse_parallel.mlir

 // RUN: mlir-opt %s -sparsification="parallelization-strategy=0" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR0
 // RUN: mlir-opt %s -sparsification="parallelization-strategy=1" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR1
 // RUN: mlir-opt %s -sparsification="parallelization-strategy=2" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR2
 // RUN: mlir-opt %s -sparsification="parallelization-strategy=3" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR3
 // RUN: mlir-opt %s -sparsification="parallelization-strategy=4" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR4
 
+#DenseMatrix = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "dense" ]
+}>
+
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ]
 }>
 
 #CSR = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ]
 }>
 

 // CHECK-PAR3:           scf.parallel
 // CHECK-PAR3:         return
 //
 // CHECK-PAR4-LABEL: func @scale_dd
 // CHECK-PAR4:         scf.parallel
 // CHECK-PAR4:           scf.parallel
 // CHECK-PAR4:         return
 //
-func @scale_dd(%scale: f32, %arga: tensor<?x?xf32>, %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {
+func @scale_dd(%scale: f32,
+               %arga: tensor<?x?xf32, #DenseMatrix>,
+         %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {
   %0 = linalg.generic #trait_dd
-     ins(%arga: tensor<?x?xf32>)
+     ins(%arga: tensor<?x?xf32, #DenseMatrix>)
     outs(%argx: tensor<?x?xf32>) {
       ^bb(%a: f32, %x: f32):
         %0 = mulf %a, %scale : f32
         linalg.yield %0 : f32
   } -> tensor<?x?xf32>
   return %0 : tensor<?x?xf32>
 }
 

 // CHECK-PAR3:           scf.for
 // CHECK-PAR3:         return
 //
 // CHECK-PAR4-LABEL: func @scale_ss
 // CHECK-PAR4:         scf.parallel
 // CHECK-PAR4:           scf.parallel
 // CHECK-PAR4:         return
 //
-func @scale_ss(%scale: f32, %arga: tensor<?x?xf32, #SparseMatrix>, %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {
+func @scale_ss(%scale: f32,
+               %arga: tensor<?x?xf32, #SparseMatrix>,
+         %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {
   %0 = linalg.generic #trait_ss
      ins(%arga: tensor<?x?xf32, #SparseMatrix>)
     outs(%argx: tensor<?x?xf32>) {
       ^bb(%a: f32, %x: f32):
         %0 = mulf %a, %scale : f32
         linalg.yield %0 : f32
   } -> tensor<?x?xf32>
   return %0 : tensor<?x?xf32>

 // CHECK-PAR3:           scf.for
 // CHECK-PAR3:         return
 //
 // CHECK-PAR4-LABEL: func @matvec
 // CHECK-PAR4:         scf.parallel
 // CHECK-PAR4:           scf.for
 // CHECK-PAR4:         return
 //
-func @matvec(%argA: tensor<16x32xf32, #CSR>, %argb: tensor<32xf32>, %argx: tensor<16xf32>) -> tensor<16xf32> {
+func @matvec(%arga: tensor<16x32xf32, #CSR>,
+             %argb: tensor<32xf32>,
+       %argx: tensor<16xf32>) -> tensor<16xf32> {
   %0 = linalg.generic #trait_matvec
-      ins(%argA, %argb : tensor<16x32xf32, #CSR>, tensor<32xf32>)
+      ins(%arga, %argb : tensor<16x32xf32, #CSR>, tensor<32xf32>)
      outs(%argx: tensor<16xf32>) {
     ^bb(%A: f32, %b: f32, %x: f32):
       %0 = mulf %A, %b : f32
       %1 = addf %0, %x : f32
       linalg.yield %1 : f32
   } -> tensor<16xf32>
   return %0 : tensor<16xf32>
 }