[mlir][sparse] support iteration over compressed-hi dimension level in loop emitter

[aartbik]

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

1132

Ah, I see why you prefer the LO=POS(2i) and HI=POS(2i+1) so much ;-)

For the records, we may still consider storing this as LO/HI or LO/LEN parallel arrays at some point, perhaps with even more level types or properties, just to give the choice.

[Peiming]

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

1132

Not only that, but also because I do not need to change the sparse tensor layout/specifier

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

for (TensorId t = 0, numTensors = getNumTensors(); t < numTensors; t++) {
     const Level cooStart = enc ? getCOOStart(enc) : lvlRank;
     // Scan all levels of current tensor.
     for (Level l = 0; l < lvlRank; l++) {
       // This should be called only once at beginning.
       assert(!positionsBuffers[t][l] && !coordinatesBuffers[t][l] &&
              !highs[t][l]);
       const auto lvlTp = lvlTypes[t][l];
       // Handle sparse storage schemes.
-      if (isCompressedDLT(lvlTp)) {
+      if (isCompressedDLT(lvlTp) || isCompressedWithHiDLT(lvlTp)) {
         // Generate sparse primitives to obtain positions and coordinates.
         positionsBuffers[t][l] = genToPositions(builder, loc, tensor, l);
         coordinatesBuffers[t][l] =
             genToCoordinates(builder, loc, tensor, l, cooStart);
       } else if (isSingletonDLT(lvlTp)) {
         // Singleton level, fetch coordinates.
         coordinatesBuffers[t][l] =
             genToCoordinates(builder, loc, tensor, l, cooStart);
default:
     llvm_unreachable("unexpected affine subscript");
   }
 }
 
 Operation *LoopEmitter::emitForLoopOverTensorAtLvl(
     OpBuilder &builder, Location loc, TensorId tid, Level dstLvl, Value lo,
     Value hi, MutableArrayRef<Value> reduc, bool isParallel) {
   bool isSparseCond = isCompressedDLT(lvlTypes[tid][dstLvl]) ||
+                      isCompressedWithHiDLT(lvlTypes[tid][dstLvl]) ||
                       isSingletonDLT(lvlTypes[tid][dstLvl]);
 
   const auto reassoc = getCollapseReassociation(tid, dstLvl);
   // TODO: support dynamic slices.
   // Uses the first dimension here to build the loop bound (which is also the
   // biggest range).
   const Level srcLvl = reassoc.front();
   Value step = C_IDX(1);
Operation *LoopEmitter::enterLoopOverTensorAtLvl(
   // the tensor levels, there is at most one sparse tensor level.
   for (auto [t, l] : llvm::zip(tids, lvls)) {
     assert(lvlTypes[t].size() > l);         // Must be a valid tid, dim pair
     assert(!coords[t][l] ||                 // We cannot re-enter the same level
            !dependentLvlMap[t][l].empty()); // unless it is a slice-driver loop
     auto lvlType = lvlTypes[t][l];
     // Must be a recognizable DLT.
     assert(isDenseDLT(lvlType) || isCompressedDLT(lvlType) ||
-           isSingletonDLT(lvlType));
+           isCompressedWithHiDLT(lvlType) || isSingletonDLT(lvlType));
 
     // This is a slice-driven loop on sparse level.
     if (!dependentLvlMap[t][l].empty() && !isDenseDLT(lvlType)) {
       assert(!isSparseSliceCond && !isSparseCond);
       isSparseSliceCond = true;
       tid = t;
       lvl = l;
       continue;
Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
   SmallVector<Type> types;
   SmallVector<Value> operands;
   // Construct the while-loop with a parameter for each coordinate.
   const Type indexType = builder.getIndexType();
   for (auto [tid, lvl] : llvm::zip(tids, lvls)) {
     // TODO: support coiteration with slice driven tensors.
     const auto lvlTp = lvlTypes[tid][lvl];
     assert(dependentLvlMap[tid][lvl].empty() && "TODO: not yet implemented");
-    if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp)) {
+    if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
+        isCompressedWithHiDLT(lvlTp)) {
       const auto reassoc = getCollapseReassociation(tid, lvl);
       for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
         if (!isUniqueDLT(lvlTypes[tid][reassoc[i]])) {
           // This is the segment high for each non-unique levels.
           types.push_back(indexType);
           operands.push_back(C_IDX(0));
         }
       }
Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
   // Build the "before" region, which effectively consists
   // of a conjunction of "i < upper" tests on all induction.
   builder.setInsertionPointToStart(&whileOp.getBefore().front());
   Value cond;
   unsigned o = 0;
   for (auto [t, lvl] : llvm::zip(tids, lvls)) {
     const TensorId tid = t; // Why `t` can not be captured by lambda?
     const auto lvlTp = lvlTypes[tid][lvl];
-    if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp)) {
+    if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
+        isCompressedWithHiDLT(lvlTp)) {
       const auto reassoc = getCollapseReassociation(tid, lvl);
       assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
       for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
         if (!isUniqueDLT(lvlTypes[tid][reassoc[i]])) {
           // Links the SSA chain for segHi.
           segHi[tid][reassoc[i]] = after->getArgument(o++);
         }
       }
Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
   // Generates while body.
   builder.setInsertionPointToStart(&whileOp.getAfter().front());
 
   SmallVector<std::pair<Value, unsigned>> slicesPreds;
   unsigned i = 0;
   for (auto [tid, lvl] : llvm::zip(tids, lvls)) {
     // Prepares for next level.
     const auto lvlTp = lvlTypes[tid][lvl];
-    if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp)) {
+    if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
+        isCompressedWithHiDLT(lvlTp)) {
       coords[tid][lvl] = genSparseCrd(builder, loc, tid, lvl);
       if (isSparseSlices[tid]) {
         auto [trans, pred] =
             genSliceLegitPredicate(builder, loc, coords[tid][lvl], tid, lvl);
         slicesPreds.emplace_back(pred, i);
         // Updates to the relative coordinate to the slice.
         coords[tid][lvl] = trans;
       }
if (!slicesPreds.empty()) {
     builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
   }
 
   Value min;
   // Finds the minimum coordinate
   if (!needsUniv) {
     for (auto [tid, lvl] : llvm::zip(tids, lvls)) {
       const auto lvlTp = lvlTypes[tid][lvl];
-      if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp)) {
+      if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
+          isCompressedWithHiDLT(lvlTp)) {
         const auto crd = coords[tid][lvl];
         if (min) {
           Value cmp = CMPI(ult, coords[tid][lvl], min);
           min = SELECT(cmp, coords[tid][lvl], min);
         } else {
           min = crd;
         }
       }
if (isDenseDLT(lvlTp))
     return;
 
   const Value c0 = C_IDX(0);
   const Value c1 = C_IDX(1);
   for (const Level srcLvl : getCollapseReassociation(tid, dstLvl)) {
     // Either the first level, or the previous level has been set.
     /// FIXME: See the [CLARIFY_POSITS_LVL] note in the header.
     assert(srcLvl == 0 || posits[tid][srcLvl - 1]);
-    if (!isCompressedDLT(lvlTp) && !isSingletonDLT(lvlTp))
+    if (isDenseDLT(lvlTp))
       continue;
-    if (isCompressedDLT(lvlTp)) {
+    if (isCompressedDLT(lvlTp) || isCompressedWithHiDLT(lvlTp)) {
       const Value mem = positionsBuffers[tid][srcLvl];
 
-      const Value pLo = srcLvl == 0 ? c0 : posits[tid][srcLvl - 1];
+      Value pLo = srcLvl == 0 ? c0 : posits[tid][srcLvl - 1];
+      if (isCompressedWithHiDLT(lvlTp))
+        pLo = builder.create<arith::MulIOp>(loc, pLo, C_IDX(2));
       posits[tid][srcLvl] = genIndexLoad(builder, loc, mem, pLo);
 
       const Value pHi = ADDI(pLo, c1);
       highs[tid][srcLvl] = genIndexLoad(builder, loc, mem, pHi);
       return;
     }
     if (isSingletonDLT(lvlTp)) {
       const Value pLo = srcLvl == 0 ? c0 : posits[tid][srcLvl - 1];
for (auto [tid, lvl, resolved] : llvm::zip(
     // It is because we are actually generating yield in the IfOp inside the
     // whileOp to only iterates over inbound coordinates within the slices.
     delta += 2;
   };
 
   Value one = C_IDX(1);
   for (auto [tid, dstLvl] : llvm::zip(loopInfo.tids, loopInfo.lvls)) {
     const auto lvlTp = lvlTypes[tid][dstLvl];
-    if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp)) {
+    if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp) ||
+        isCompressedWithHiDLT(lvlTp)) {
       const auto reassoc = getCollapseReassociation(tid, dstLvl);
       assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
       for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
         const Level srcLvl = reassoc[i];
         if (!isUniqueDLT(lvlTypes[tid][srcLvl])) {
           operands.push_back(segHi[tid][srcLvl]);
           o++;
         }

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

 
 /// Generations insertion finalization code.
 static void genEndInsert(OpBuilder &builder, Location loc,
                          SparseTensorDescriptor desc) {
   const SparseTensorType stt(desc.getRankedTensorType());
   const Level lvlRank = stt.getLvlRank();
   for (Level l = 0; l < lvlRank; l++) {
     const auto dlt = stt.getLvlType(l);
+    if (isCompressedWithHiDLT(dlt))
+      llvm_unreachable("TODO: Not yet implemented");
     if (isCompressedDLT(dlt)) {
       // Compressed dimensions need a position cleanup for all entries
       // that were not visited during the insertion pass.
       //
       // TODO: avoid cleanup and keep compressed scheme consistent at all
       // times?
       //
       if (l > 0) {

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

#define RETURN_ON_FALSE(fidx, kind, dim, dlt)                                  \
   const Level end = cooStart == lvlRank ? cooStart : cooStart + 1;
   FieldIndex fieldIdx = kDataFieldStartingIdx;
   // Per-dimension storage.
   for (Level l = 0; l < end; l++) {
     // Dimension level types apply in order to the reordered dimension.
     // As a result, the compound type can be constructed directly in the given
     // order.
     const auto dlt = lvlTypes[l];
-    if (isCompressedDLT(dlt)) {
+    if (isCompressedDLT(dlt) || isCompressedWithHiDLT(dlt)) {
       RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::PosMemRef, l, dlt);
       RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::CrdMemRef, l, dlt);
     } else if (isSingletonDLT(dlt)) {
       RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::CrdMemRef, l, dlt);
     } else {
       assert(isDenseDLT(dlt)); // no fields
     }
   }

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

for (OpOperand &t : env.op()->getOpOperands()) {
 
     // Push unrelated loops into sparse iteration space, so these
     // will be skipped more often.
     // TODO: Do we really need this?
     if (includesUndef(mask)) {
       const TensorId tid = env.makeTensorId(t.getOperandNumber());
       for (LoopId i = 0; i < numLoops; i++) {
         const auto dltI = env.dlt(tid, i);
-        if (isCompressedDLT(dltI) || isSingletonDLT(dltI)) {
+        if (isCompressedDLT(dltI) || isCompressedWithHiDLT(dltI) ||
+            isSingletonDLT(dltI)) {
           for (LoopId j = 0; j < numLoops; j++)
             if (isUndefDLT(env.dlt(tid, j))) {
               adjM[i][j] = true;
               inDegree[j]++;
             }
         } else {
           assert(isDenseDLT(dltI) || isUndefDLT(dltI));
         }
static scf::IfOp genIf(CodegenEnv &env, OpBuilder &builder, LoopId ldx,
   SmallVector<Type> types;
   Value cond;
   env.merger().foreachTensorLoopId(
       p, /*simple=*/true,
       [&](TensorLoopId b, TensorId tid, std::optional<Level> lvl,
           DimLevelType dlt, bool /*unused*/) {
         assert(ldx == env.merger().loop(b));
         Value clause;
-        if (isCompressedDLT(dlt) || isSingletonDLT(dlt)) {
+        if (isCompressedDLT(dlt) || isSingletonDLT(dlt) || isCompressedWithHiDLT(dlt)) {
           assert(lvl.has_value());
           const Value crd = env.emitter().getCoords()[tid][*lvl];
           const Value lvar = env.getLoopVar(ldx);
           clause = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
                                                  crd, lvar);
         } else {
           assert(isDenseDLT(dlt) || isUndefDLT(dlt));
           clause = constantI1(builder, loc, true);

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

BitVector Merger::simplifyCond(LatSetId s0, LatPointId p0) {
 
   // Now apply the two basic rules. We also iterate the bits reversely to always
   // keep the rightmost bit (which could possibly be a synthetic tensor).
   for (unsigned b = be - 1 - offset, i = 0; i < be;
        b = b == 0 ? be - 1 : b - 1, i++) {
     // Slice on dense level has `locate` property as well, and can be optimized.
     if (simple[b] && !isSparseLvlWithNonTrivialIdxExp(b)) {
       const auto dlt = getDimLevelType(b);
-      if (!isCompressedDLT(dlt) && !isSingletonDLT(dlt)) {
+      if (!isCompressedDLT(dlt) && !isSingletonDLT(dlt) && !isCompressedWithHiDLT(dlt)) {
         if (reset)
           simple.reset(b);
         reset = true;
       }
     }
   }
   return simple;
 }
case TensorExp::Kind::kReduce:
     return false;
   }
   llvm_unreachable("unexpected kind");
 }
 
 bool Merger::hasAnySparse(const BitVector &bits) const {
   for (TensorLoopId b : bits.set_bits()) {
     const auto dlt = getDimLevelType(b);
-    if (isCompressedDLT(dlt) || isSingletonDLT(dlt))
+    if (isCompressedDLT(dlt) || isSingletonDLT(dlt) || isCompressedWithHiDLT(dlt))
       return true;
   }
   return hasSparseIdxReduction(bits);
 }
 
 bool Merger::hasSparseIdxReduction(const BitVector &bits) const {
   for (TensorLoopId b : bits.set_bits())
     if (isSparseLvlWithNonTrivialIdxExp(b))

mlir/test/Dialect/SparseTensor/sparse_foreach.mlir

 // CHECK:           return
 //
 func.func @foreach_print_slice(%A: tensor<4x4xf64, #CSR_SLICE>) {
   sparse_tensor.foreach in %A : tensor<4x4xf64, #CSR_SLICE> do {
   ^bb0(%1: index, %2: index, %v: f64) :
     "test.use" (%v) : (f64) -> ()
   }
   return
-}
+}
+
+#BCOO = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed-hi-nu", "singleton" ],
+}>
+
+// CHECK-LABEL:   func.func @foreach_bcoo(
+// CHECK-SAME:      %[[VAL_0:.*]]: tensor<4x4x4xf64, #sparse_tensor.encoding<{{.*}}>>) {
+// CHECK-DAG:       %[[VAL_1:.*]] = arith.constant 4 : index
+// CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 0 : index
+// CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 1 : index
+// CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 2 : index
+// CHECK-DAG:       %[[VAL_5:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 1 : index} : tensor<4x4x4xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
+// CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<4x4x4xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xf64>
+// CHECK:           scf.for %[[VAL_7:.*]] = %[[VAL_2]] to %[[VAL_1]] step %[[VAL_3]] {
+// CHECK:             %[[VAL_8:.*]] = arith.muli %[[VAL_7]], %[[VAL_4]] : index
+// CHECK:             %[[VAL_9:.*]] = arith.addi %[[VAL_8]], %[[VAL_3]] : index
+// CHECK:             %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_9]]] : memref<?xindex>
+// CHECK:             scf.for %[[VAL_11:.*]] = %[[VAL_2]] to %[[VAL_10]] step %[[VAL_3]] {
+// CHECK:               %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf64>
+// CHECK:               "test.use"(%[[VAL_12]]) : (f64) -> ()
+// CHECK:             }
+// CHECK:           }
+// CHECK:           return
+func.func @foreach_bcoo(%A: tensor<4x4x4xf64, #BCOO>) {
+  sparse_tensor.foreach in %A : tensor<4x4x4xf64, #BCOO> do {
+  ^bb0(%1: index, %2: index, %3: index,  %v: f64) :
+    "test.use" (%v) : (f64) -> ()
+  }
+  return
+}