[mlir][sparse] Implement sparse_tensor.select

[jim22k]

Summary

The region within select is used as the runtime criteria
of an scf::IfOp. If the check succeeds, the original value from the
sparse tensor remains in the output at the same location. If it fails,
the location becomes missing in the output.

The value is provided to the region block for the comparison. However,
the indices of the value may also be used during the comparison for
selecting the upper triangle of a matrix, for example.

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

77

would you mind moving this under kUnary (at all relevant places), so it follows the order in which you introduced them in the .td file better?

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

890

can you merge the else/if and save the extra "walk to the right"

894

See comment below.

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

1236

oh, this feels a bit hacky, since I would not expect a side effect from calling this method
the bookkeeping of caching the value should probably live in sparsification?

[jim22k]

• Updates based on feedback

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

1236

I moved the runtime "value saving" to Sparsification. It does feel better there.

[aartbik]

Nice work!

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

885

add an

assert(merger.exp(exp).val);

before the assignment

1054

add { } around the if and ass an

assert(!merger.exp(exp).val)

before the assignment

[jim22k]

• Add assert statements

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

enum Kind {
   kCastUF, // unsigned
   kCastS,  // signed
   kCastU,  // unsigned
   kCastIdx,
   kTruncI,
   kCIm, // complex.im
   kCRe, // complex.re
   kBitCast,
   kBinaryBranch, // semiring unary branch created from a binary op
   kUnary,        // semiring unary op
+  kSelect,       // custom selection criteria
   // Binary operations.
   kMulF,
   kMulC,
   kMulI,
   kDivF,
   kDivC, // complex
   kDivS, // signed
   kDivU, // unsigned
union {
     Children children;
   };
 
   /// Direct link to IR for an invariant or the destination value (to
   /// infer destination type) of a cast operation During code generation,
   /// this field may be used to cache "hoisted" loop invariant tensor loads.
   Value val;
 
-  /// Code blocks used by semirings. For the case of kUnary, kBinary, and
-  /// kReduce, this holds the original operation with all regions. For
+  /// Code blocks used by semirings. For the case of kUnary, kBinary, kReduce,
+  /// and kSelect, this holds the original operation with all regions. For
   /// kBinaryBranch, this holds the YieldOp for the left or right half
   /// to be merged into a nested scf loop.
   Operation *op;
 };
 
 /// Lattice point. Each lattice point consists of a conjunction of tensor
 /// loop indices (encoded in a bitvector) and the index of the corresponding
 /// tensor expression.

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

static void genTensorStore(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   }
   // Store during insertion.
   OpOperand *t = op.getOutputOperand(0);
   if (t == codegen.sparseOut) {
     if (!rhs) {
       // Only unary and binary are allowed to return uninitialized rhs
       // to indicate missing output.
       assert(merger.exp(exp).kind == kUnary || merger.exp(exp).kind == kBinary);
+    } else if (merger.exp(exp).kind == kSelect) {
+      scf::IfOp ifOp = builder.create<scf::IfOp>(loc, rhs);
+      builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+      // Existing value was preserved to be used here.
+      assert(merger.exp(exp).val);
+      Value v0 = merger.exp(exp).val;
+      genInsertionStore(codegen, builder, op, t, v0);
+      merger.exp(exp).val = Value();
+      builder.setInsertionPointAfter(ifOp);
     } else {
       genInsertionStore(codegen, builder, op, t, rhs);
     }
     return;
   }
   // Actual store.
   SmallVector<Value, 4> args;
   Value ptr = genSubscript(codegen, builder, op, t, args);
   if (codegen.curVecLength > 1)
     genVectorStore(codegen, builder, rhs, ptr, args);
   else
     builder.create<memref::StoreOp>(loc, rhs, ptr, args);
 }
static Value genExp(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
   Value v0 =
       genExp(merger, codegen, rewriter, op, merger.exp(exp).children.e0, ldx);
   Value v1 =
       genExp(merger, codegen, rewriter, op, merger.exp(exp).children.e1, ldx);
   Value ee = merger.buildExp(rewriter, loc, exp, v0, v1);
   if (ee && (merger.exp(exp).kind == Kind::kUnary ||
              merger.exp(exp).kind == Kind::kBinary ||
              merger.exp(exp).kind == Kind::kBinaryBranch ||
-             merger.exp(exp).kind == Kind::kReduce))
+             merger.exp(exp).kind == Kind::kReduce ||
+             merger.exp(exp).kind == Kind::kSelect))
     ee = relinkBranch(codegen, rewriter, ee.getParentBlock(), ee, ldx);
 
+  if (merger.exp(exp).kind == kSelect) {
+    assert(!merger.exp(exp).val);
+    merger.exp(exp).val = v0; // Preserve value for later use.
+  }
+
   if (merger.exp(exp).kind == Kind::kReduce) {
     assert(codegen.redCustom != -1u);
     codegen.redCustom = -1u;
   }
 
   return ee;
 }
 

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

TensorExp::TensorExp(Kind k, unsigned x, unsigned y, Value v, Operation *o)
   case kCastIdx:
   case kTruncI:
   case kBitCast:
     assert(x != -1u && y == -1u && v && !o);
     children.e0 = x;
     children.e1 = y;
     break;
   case kBinaryBranch:
+  case kSelect:
     assert(x != -1u && y == -1u && !v && o);
     children.e0 = x;
     children.e1 = y;
     break;
   case kUnary:
     // No assertion on y can be made, as the branching paths involve both
     // a unary (mapSet) and binary (takeDisj) pathway.
     assert(x != -1u && !v && o);
if (opright)
       s1 = mapSet(rtrans, s1, Value(), opright);
     for (unsigned p : latSets[s1])
       latSets[s].push_back(p);
   }
   return s;
 }
 
 unsigned Merger::mapSet(Kind kind, unsigned s0, Value v, Operation *op) {
-  assert(kAbsF <= kind && kind <= kUnary);
+  assert(kAbsF <= kind && kind <= kSelect);
   unsigned s = addSet();
   for (unsigned p : latSets[s0]) {
     unsigned e = addExp(kind, latPoints[p].exp, v, op);
     latPoints.push_back(LatPoint(latPoints[p].bits, e));
     latSets[s].push_back(latPoints.size() - 1);
   }
   return s;
 }
if (p0 != p1 && latGT(p0, p1)) {
       isSingleton = false;
       break;
     }
   }
   // Now apply the two basic rules.
   BitVector simple = latPoints[p0].bits;
   bool reset = isSingleton && hasAnySparse(simple);
   for (unsigned b = 0, be = simple.size(); b < be; b++) {
-    if (simple[b] &&
-        (!isDimLevelType(b, DimLvlType::kCompressed) &&
+    if (simple[b] && (!isDimLevelType(b, DimLvlType::kCompressed) &&
                       !isDimLevelType(b, DimLvlType::kSingleton))) {
       if (reset)
         simple.reset(b);
       reset = true;
     }
   }
   return simple;
 }
 
bool Merger::isSingleCondition(unsigned t, unsigned e) const {
   case kCastIdx:
   case kTruncI:
   case kCIm:
   case kCRe:
   case kBitCast:
     return isSingleCondition(t, tensorExps[e].children.e0);
   case kBinaryBranch:
   case kUnary:
+  case kSelect:
     return false;
   // Binary operations.
   case kDivF: // note: x / c only
   case kDivC:
   case kDivS:
   case kDivU:
     assert(!maybeZero(tensorExps[e].children.e1));
     return isSingleCondition(t, tensorExps[e].children.e0);
static const char *kindToOpSymbol(Kind kind) {
   case kCRe:
     return "complex.re";
   case kBitCast:
     return "cast";
   case kBinaryBranch:
     return "binary_branch";
   case kUnary:
     return "unary";
+  case kSelect:
+    return "select";
   // Binary operations.
   case kMulF:
   case kMulC:
   case kMulI:
     return "*";
   case kDivF:
   case kDivC:
   case kDivS:
void Merger::dumpExp(unsigned e) const {
   case kCastU:
   case kCastIdx:
   case kTruncI:
   case kCIm:
   case kCRe:
   case kBitCast:
   case kBinaryBranch:
   case kUnary:
+  case kSelect:
     llvm::dbgs() << kindToOpSymbol(tensorExps[e].kind) << " ";
     dumpExp(tensorExps[e].children.e0);
     break;
   // Binary operations.
   case kMulF:
   case kMulC:
   case kMulI:
   case kDivF:
case kBitCast:
     // lattice set of the operand through the operator into a new set.
     //
     //  -y|!y | y |
     //  --+---+---+
     //    | 0 |-y |
     return mapSet(kind, buildLattices(tensorExps[e].children.e0, i),
                   tensorExps[e].val);
   case kBinaryBranch:
+  case kSelect:
     // The left or right half of a binary operation which has already
     // been split into separate operations for each region.
     return mapSet(kind, buildLattices(tensorExps[e].children.e0, i), Value(),
                   tensorExps[e].op);
   case kUnary:
     // A custom unary operation.
     //
     //  op y|    !y    |     y      |
if (x.has_value()) {
         return addExp(kCRe, e);
       if (isa<arith::BitcastOp>(def))
         return addExp(kBitCast, e, v);
       if (auto unop = dyn_cast<sparse_tensor::UnaryOp>(def)) {
         if (isAdmissableBranch(unop, unop.getPresentRegion()) &&
             isAdmissableBranch(unop, unop.getAbsentRegion()))
           return addExp(kUnary, e, Value(), def);
       }
+      if (auto selop = dyn_cast<sparse_tensor::SelectOp>(def)) {
+        if (isAdmissableBranch(selop, selop.getRegion()))
+          return addExp(kSelect, e, Value(), def);
+      }
     }
   }
   // Construct binary operations if subexpressions can be built.
   // See buildLattices() for an explanation of rejecting certain
   // division and shift operations.
   if (def->getNumOperands() == 2) {
     auto x = buildTensorExp(op, def->getOperand(0));
     auto y = buildTensorExp(op, def->getOperand(1));
case kXorI:
     return rewriter.create<arith::XOrIOp>(loc, v0, v1);
   case kShrS:
     return rewriter.create<arith::ShRSIOp>(loc, v0, v1);
   case kShrU:
     return rewriter.create<arith::ShRUIOp>(loc, v0, v1);
   case kShlI:
     return rewriter.create<arith::ShLIOp>(loc, v0, v1);
   case kBinaryBranch: // semi-ring ops with custom logic.
     return insertYieldOp(rewriter, loc,
                          *tensorExps[e].op->getBlock()->getParent(), {v0});
   case kUnary:
     return buildUnaryPresent(rewriter, loc, tensorExps[e].op, v0);
+  case kSelect:
+    return insertYieldOp(rewriter, loc,
+                         cast<SelectOp>(tensorExps[e].op).getRegion(), {v0});
   case kBinary:
     return buildBinaryOverlap(rewriter, loc, tensorExps[e].op, v0, v1);
   case kReduce: {
     ReduceOp redOp = cast<ReduceOp>(tensorExps[e].op);
     return insertYieldOp(rewriter, loc, redOp.getRegion(), {v0, v1});
   }
   }
   llvm_unreachable("unexpected expression kind in build");
 }
 
 } // namespace sparse_tensor
 } // namespace mlir

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

+// RUN: mlir-opt %s --sparse-compiler | \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+#CSR = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}>
+#CSC = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  dimOrdering = affine_map<(i,j) -> (j,i)>
+}>
+
+//
+// Traits for tensor operations.
+//
+#trait_vec_select = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>, // A
+    affine_map<(i) -> (i)>  // C (out)
+  ],
+  iterator_types = ["parallel"]
+}
+
+#trait_mat_select = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>,  // A (in)
+    affine_map<(i,j) -> (i,j)>   // X (out)
+  ],
+  iterator_types = ["parallel", "parallel"]
+}
+
+module {
+  func.func @vecSelect(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
+    %c0 = arith.constant 0 : index
+    %cf1 = arith.constant 1.0 : f64
+    %d0 = tensor.dim %arga, %c0 : tensor<?xf64, #SparseVector>
+    %xv = bufferization.alloc_tensor(%d0): tensor<?xf64, #SparseVector>
+    %0 = linalg.generic #trait_vec_select
+      ins(%arga: tensor<?xf64, #SparseVector>)
+      outs(%xv: tensor<?xf64, #SparseVector>) {
+        ^bb(%a: f64, %b: f64):
+          %1 = sparse_tensor.select %a : f64 {
+              ^bb0(%x: f64):
+                %keep = arith.cmpf "oge", %x, %cf1 : f64
+                sparse_tensor.yield %keep : i1
+            }
+          linalg.yield %1 : f64
+    } -> tensor<?xf64, #SparseVector>
+    return %0 : tensor<?xf64, #SparseVector>
+  }
+
+  func.func @matUpperTriangle(%arga: tensor<?x?xf64, #CSR>) -> tensor<?x?xf64, #CSR> {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #CSR>
+    %d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #CSR>
+    %xv = bufferization.alloc_tensor(%d0, %d1): tensor<?x?xf64, #CSR>
+    %0 = linalg.generic #trait_mat_select
+      ins(%arga: tensor<?x?xf64, #CSR>)
+      outs(%xv: tensor<?x?xf64, #CSR>) {
+        ^bb(%a: f64, %b: f64):
+          %row = linalg.index 0 : index
+          %col = linalg.index 1 : index
+          %1 = sparse_tensor.select %a : f64 {
+              ^bb0(%x: f64):
+                %keep = arith.cmpi "ugt", %col, %row : index
+                sparse_tensor.yield %keep : i1
+            }
+          linalg.yield %1 : f64
+    } -> tensor<?x?xf64, #CSR>
+    return %0 : tensor<?x?xf64, #CSR>
+  }
+
+  // Dumps a sparse vector of type f64.
+  func.func @dump_vec(%arg0: tensor<?xf64, #SparseVector>) {
+    // Dump the values array to verify only sparse contents are stored.
+    %c0 = arith.constant 0 : index
+    %d0 = arith.constant -1.0 : f64
+    %0 = sparse_tensor.values %arg0 : tensor<?xf64, #SparseVector> to memref<?xf64>
+    %1 = vector.transfer_read %0[%c0], %d0: memref<?xf64>, vector<8xf64>
+    vector.print %1 : vector<8xf64>
+    // Dump the dense vector to verify structure is correct.
+    %dv = sparse_tensor.convert %arg0 : tensor<?xf64, #SparseVector> to tensor<?xf64>
+    %2 = vector.transfer_read %dv[%c0], %d0: tensor<?xf64>, vector<16xf64>
+    vector.print %2 : vector<16xf64>
+    return
+  }
+
+  // Dump a sparse matrix.
+  func.func @dump_mat(%arg0: tensor<?x?xf64, #CSR>) {
+    // Dump the values array to verify only sparse contents are stored.
+    %c0 = arith.constant 0 : index
+    %d0 = arith.constant -1.0 : f64
+    %0 = sparse_tensor.values %arg0 : tensor<?x?xf64, #CSR> to memref<?xf64>
+    %1 = vector.transfer_read %0[%c0], %d0: memref<?xf64>, vector<16xf64>
+    vector.print %1 : vector<16xf64>
+    %dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #CSR> to tensor<?x?xf64>
+    %2 = vector.transfer_read %dm[%c0, %c0], %d0: tensor<?x?xf64>, vector<5x5xf64>
+    vector.print %2 : vector<5x5xf64>
+    return
+  }
+
+  // Driver method to call and verify vector kernels.
+  func.func @entry() {
+    %c0 = arith.constant 0 : index
+
+    // Setup sparse matrices.
+    %v1 = arith.constant sparse<
+        [ [1], [3], [5], [7], [9] ],
+        [ 1.0, 2.0, -4.0, 0.0, 5.0 ]
+    > : tensor<10xf64>
+    %m1 = arith.constant sparse<
+        [ [0, 3], [1, 4], [2, 1], [2, 3], [3, 3], [3, 4], [4, 2] ],
+        [ 1., 2., 3., 4., 5., 6., 7.]
+    > : tensor<5x5xf64>
+    %sv1 = sparse_tensor.convert %v1 : tensor<10xf64> to tensor<?xf64, #SparseVector>
+    %sm1 = sparse_tensor.convert %m1 : tensor<5x5xf64> to tensor<?x?xf64, #CSR>
+
+    // Call sparse matrix kernels.
+    %1 = call @vecSelect(%sv1) : (tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector>
+    %2 = call @matUpperTriangle(%sm1) : (tensor<?x?xf64, #CSR>) -> tensor<?x?xf64, #CSR>
+
+    //
+    // Verify the results.
+    //
+    // CHECK:      ( 1, 2, -4, 0, 5, -1, -1, -1 )
+    // CHECK-NEXT: ( 0, 1, 0, 2, 0, -4, 0, 0, 0, 5, -1, -1, -1, -1, -1, -1 )
+    // CHECK-NEXT: ( 1, 2, 3, 4, 5, 6, 7, -1, -1, -1, -1, -1, -1, -1, -1, -1 )
+    // CHECK-NEXT: ( ( 0, 0, 0, 1, 0 ), ( 0, 0, 0, 0, 2 ), ( 0, 3, 0, 4, 0 ), ( 0, 0, 0, 5, 6 ), ( 0, 0, 7, 0, 0 ) )
+    // CHECK-NEXT: ( 1, 2, 5, -1, -1, -1, -1, -1 )
+    // CHECK-NEXT: ( 0, 1, 0, 2, 0, 0, 0, 0, 0, 5, -1, -1, -1, -1, -1, -1 )
+    // CHECK-NEXT: ( 1, 2, 4, 6, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1 )
+    // CHECK-NEXT: ( ( 0, 0, 0, 1, 0 ), ( 0, 0, 0, 0, 2 ), ( 0, 0, 0, 4, 0 ), ( 0, 0, 0, 0, 6 ), ( 0, 0, 0, 0, 0 ) )
+    //
+    call @dump_vec(%sv1) : (tensor<?xf64, #SparseVector>) -> ()
+    call @dump_mat(%sm1) : (tensor<?x?xf64, #CSR>) -> ()
+    call @dump_vec(%1) : (tensor<?xf64, #SparseVector>) -> ()
+    call @dump_mat(%2) : (tensor<?x?xf64, #CSR>) -> ()
+
+    // Release the resources.
+    bufferization.dealloc_tensor %sv1 : tensor<?xf64, #SparseVector>
+    bufferization.dealloc_tensor %sm1 : tensor<?x?xf64, #CSR>
+    bufferization.dealloc_tensor %1 : tensor<?xf64, #SparseVector>
+    bufferization.dealloc_tensor %2 : tensor<?x?xf64, #CSR>
+    return
+  }
+}

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

bool compareExpression(unsigned e, const std::shared_ptr<Pattern> &pattern) {
     case kCastUF:
     case kCastS:
     case kCastU:
     case kCastIdx:
     case kTruncI:
     case kCIm:
     case kCRe:
     case kBitCast:
+    case kSelect:
     case kBinaryBranch:
     case kUnary:
       return compareExpression(tensorExp.children.e0, pattern->e0);
     // Binary operations.
     case kMulF:
     case kMulC:
     case kMulI:
     case kDivF: