[mlir][sparse] fix crashes when using custom reduce with unary operation.

[Peiming]

Summary

The tests case is directly copied from https://reviews.llvm.org/D152179 authored by @aartbik

[aartbik]

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

522

You can do this even before the yield = ... statement

and perhaps let's even test pedantically

if (block.getOperations().size() <= 1)
   return false

so that we never crash getting the yield either

[Peiming]

address comments.

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

 #undef BINOP
 
 /// This method is called twice to analyze and rewrite the given for-loop.
 /// The first call (!codegen) does the analysis. Then, on success, the second
 /// call (codegen) rewriters the IR into vector form. This mechanism ensures
 /// that analysis and rewriting code stay in sync.
 static bool vectorizeStmt(PatternRewriter &rewriter, scf::ForOp forOp, VL vl,
                           bool codegen) {
-  Location loc = forOp.getLoc();
   Block &block = forOp.getRegion().front();
+  // For loops with single yield statement (as below) could be generated
+  // when custom reduce is used with unary operation.
+  // for (...)
+  //   yield c_0
+  if (block.getOperations().size() <= 1)
+    return false;
+
+  Location loc = forOp.getLoc();
   scf::YieldOp yield = cast<scf::YieldOp>(block.getTerminator());
   auto &last = *++block.rbegin();
   scf::ForOp forOpNew;
 
   // Perform initial set up during codegen (we know that the first analysis
   // pass was successful). For reductions, we need to construct a completely
   // new for-loop, since the incoming and outgoing reduction type
   // changes into SIMD form. For stores, we can simply adjust the stride
   // and insert in the existing for-loop. In both cases, we set up a vector
   // mask for all operations which takes care of confining vectors to
   // the original iteration space (later cleanup loops or other
   // optimizations can take care of those).
   Value vmask;
   if (codegen) {

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_prod.mlir

#trait_reduction = {
   ],
   iterator_types = ["reduction"],
   doc = "x += PROD_CUSTOM_i a(i)"
 }
 
 // An example of vector reductions.
 module {
 
+  // Custom prod reduction: stored i32 elements only.
   func.func @prod_dreduction_i32(%arga: tensor<32xi32, #DV>,
                                  %argx: tensor<i32>) -> tensor<i32> {
     %c = tensor.extract %argx[] : tensor<i32>
     %0 = linalg.generic #trait_reduction
       ins(%arga: tensor<32xi32, #DV>)
       outs(%argx: tensor<i32>) {
         ^bb(%a: i32, %b: i32):
           %1 = sparse_tensor.reduce %a, %b, %c : i32 {
             ^bb0(%x: i32, %y: i32):
               %2 = arith.muli %x, %y : i32
               sparse_tensor.yield %2 : i32
           }
           linalg.yield %1 : i32
     } -> tensor<i32>
     return %0 : tensor<i32>
   }
 
+  // Custom prod reduction: stored f32 elements only.
   func.func @prod_dreduction_f32(%arga: tensor<32xf32, #DV>,
                                  %argx: tensor<f32>) -> tensor<f32> {
     %c = tensor.extract %argx[] : tensor<f32>
     %0 = linalg.generic #trait_reduction
       ins(%arga: tensor<32xf32, #DV>)
       outs(%argx: tensor<f32>) {
         ^bb(%a: f32, %b: f32):
           %1 = sparse_tensor.reduce %a, %b, %c : f32 {
             ^bb0(%x: f32, %y: f32):
               %2 = arith.mulf %x, %y : f32
               sparse_tensor.yield %2 : f32
           }
           linalg.yield %1 : f32
     } -> tensor<f32>
     return %0 : tensor<f32>
   }
 
+  // Custom prod reduction: stored i32 elements only.
   func.func @prod_sreduction_i32(%arga: tensor<32xi32, #SV>,
                                  %argx: tensor<i32>) -> tensor<i32> {
     %c = tensor.extract %argx[] : tensor<i32>
     %0 = linalg.generic #trait_reduction
       ins(%arga: tensor<32xi32, #SV>)
       outs(%argx: tensor<i32>) {
         ^bb(%a: i32, %b: i32):
           %1 = sparse_tensor.reduce %a, %b, %c : i32 {
             ^bb0(%x: i32, %y: i32):
               %2 = arith.muli %x, %y : i32
               sparse_tensor.yield %2 : i32
           }
           linalg.yield %1 : i32
     } -> tensor<i32>
     return %0 : tensor<i32>
   }
 
+  // Custom prod reduction: stored f32 elements only.
   func.func @prod_sreduction_f32(%arga: tensor<32xf32, #SV>,
                                  %argx: tensor<f32>) -> tensor<f32> {
     %c = tensor.extract %argx[] : tensor<f32>
     %0 = linalg.generic #trait_reduction
       ins(%arga: tensor<32xf32, #SV>)
       outs(%argx: tensor<f32>) {
         ^bb(%a: f32, %b: f32):
           %1 = sparse_tensor.reduce %a, %b, %c : f32 {
             ^bb0(%x: f32, %y: f32):
               %2 = arith.mulf %x, %y : f32
               sparse_tensor.yield %2 : f32
           }
           linalg.yield %1 : f32
     } -> tensor<f32>
     return %0 : tensor<f32>
   }
 
+  // Custom prod reduction: stored i32 elements and implicit zeros.
+  //
+  // NOTE: this is a somewhat strange operation, since for most sparse
+  //       situations the outcome would always be zero; it is added
+  //       to test full functionality and illustrate the subtle differences
+  //       between the various custom operations; it would make a bit more
+  //       sense for e.g. a min/max reductions, although it still would
+  //       "densify" the iteration space.
+  //
+  func.func @prod_xreduction_i32(%arga: tensor<32xi32, #SV>,
+                                 %argx: tensor<i32>) -> tensor<i32> {
+    %c = tensor.extract %argx[] : tensor<i32>
+    %0 = linalg.generic #trait_reduction
+      ins(%arga: tensor<32xi32, #SV>)
+      outs(%argx: tensor<i32>) {
+        ^bb(%a: i32, %b: i32):
+           %u = sparse_tensor.unary %a : i32 to i32
+           present={
+             ^bb0(%x: i32):
+             sparse_tensor.yield %x : i32
+           } absent={
+             ^bb0:
+             %c0 = arith.constant 0 : i32
+             sparse_tensor.yield %c0 : i32
+          }
+          %1 = sparse_tensor.reduce %u, %b, %c : i32 {
+            ^bb0(%x: i32, %y: i32):
+              %2 = arith.muli %x, %y : i32
+              sparse_tensor.yield %2 : i32
+          }
+          linalg.yield %1 : i32
+    } -> tensor<i32>
+    return %0 : tensor<i32>
+  }
+
+
   func.func @dump_i32(%arg0 : tensor<i32>) {
     %v = tensor.extract %arg0[] : tensor<i32>
     vector.print %v : i32
     return
   }
 
   func.func @dump_f32(%arg0 : tensor<f32>) {
     %v = tensor.extract %arg0[] : tensor<f32>
func.func @entry() {
     %1 = call @prod_dreduction_f32(%d0_f32, %rf) : (tensor<32xf32, #DV>, tensor<f32>) -> tensor<f32>
     %2 = call @prod_sreduction_i32(%s0_i32, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
     %3 = call @prod_sreduction_f32(%s0_f32, %rf) : (tensor<32xf32, #SV>, tensor<f32>) -> tensor<f32>
     %4 = call @prod_dreduction_i32(%d1_i32, %ri) : (tensor<32xi32, #DV>, tensor<i32>) -> tensor<i32>
     %5 = call @prod_dreduction_f32(%d1_f32, %rf) : (tensor<32xf32, #DV>, tensor<f32>) -> tensor<f32>
     %6 = call @prod_sreduction_i32(%s1_i32, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
     %7 = call @prod_sreduction_f32(%s1_f32, %rf) : (tensor<32xf32, #SV>, tensor<f32>) -> tensor<f32>
     %8 = call @prod_sreduction_i32(%s0,     %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
+    %9 = call @prod_xreduction_i32(%s0_i32, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
+    %10 = call @prod_xreduction_i32(%s1_i32, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
 
     // Verify results. Note that the custom reduction gave permission
     // to treat an explicit vs implicit zero differently to compute the
     // full product reduction over stored elements. A "standard" product
     // reduction would have to return 0 for any implicit zero occurrence
     // too. An explicit zero nullifies the product, though, as requested.
     //
     // CHECK: 0
     // CHECK: 0
     // CHECK: 3087
     // CHECK: 14
     // CHECK: 3087
     // CHECK: 168
     // CHECK: 3087
     // CHECK: 168
     // CHECK: 0
+    // CHECK: 0
+    // CHECK: 3087
     //
     call @dump_i32(%0) : (tensor<i32>) -> ()
     call @dump_f32(%1) : (tensor<f32>) -> ()
     call @dump_i32(%2) : (tensor<i32>) -> ()
     call @dump_f32(%3) : (tensor<f32>) -> ()
     call @dump_i32(%4) : (tensor<i32>) -> ()
     call @dump_f32(%5) : (tensor<f32>) -> ()
     call @dump_i32(%6) : (tensor<i32>) -> ()
     call @dump_f32(%7) : (tensor<f32>) -> ()
     call @dump_i32(%8) : (tensor<i32>) -> ()
+    call @dump_i32(%9) : (tensor<i32>) -> ()
+    call @dump_i32(%10) : (tensor<i32>) -> ()
 
     // Release the resources.
     bufferization.dealloc_tensor %d0_i32 : tensor<32xi32, #DV>
     bufferization.dealloc_tensor %d0_f32 : tensor<32xf32, #DV>
     bufferization.dealloc_tensor %s0_i32 : tensor<32xi32, #SV>
     bufferization.dealloc_tensor %s0_f32 : tensor<32xf32, #SV>
     bufferization.dealloc_tensor %d1_i32 : tensor<32xi32, #DV>
     bufferization.dealloc_tensor %d1_f32 : tensor<32xf32, #DV>
     bufferization.dealloc_tensor %s1_i32 : tensor<32xi32, #SV>
     bufferization.dealloc_tensor %s1_f32 : tensor<32xf32, #SV>
     bufferization.dealloc_tensor %s0     : tensor<32xi32, #SV>
 
     return
   }
 }