[mlir][sparse] fix crash when using pure constant index in indexing mapping (fixes #61530)

[Peiming]

Summary

To address https://github.com/llvm/llvm-project/issues/61530

[aartbik]

can you add a before fail/after pass test?

[wrengr]

If you add "(fixes #61530)" to the title, then when you submit, github will automatically close the ticket and add a comment there linking to the commit

[wrengr]

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

441

This would read more fluently as "If we have an"

[wrengr]

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

446–447

I think it would be better to update the ctor/initialize so that it makes lvlToLoop have the right size in the first place, rather than doing this workaround. Since this looks like it's actually a problem with confusing srcLvl vs dstLvl.

[Peiming]

add test cases.

[Peiming]

revert unintended change.

[Peiming]

address comments.

[Peiming]

revert unintended change.

[Peiming]

revert unintended changes.

[wrengr]

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

270

You should add documentation for what this parameter is for. And should rename the parameter to make clear what sort of "rank" it's the maximum for (srcLvlRank? dstLvlRank? srcDimRank? dstDimRank? other?). Based on the changes to Merger.cpp it looks like it's the maximum srcLvlRank, but based on the code in GenericOpSparsifier it looks instead like it's the maximum dstDimRank.

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

295–296

👍

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

1822–1823

Reiterating my comment from before, this code computes the maximum dstDimRank but Merger.cpp looks like it actually wants the maximum srcLvlRank. So there's some confusion that needs to be cleared up here.

I'm pretty sure you want the level-rank (regardless of the src-vs-dst issue), so you should use SparseTensorType::getLvlRank here instead.

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

223

Just to connect with the previous comments, this code is why I think "maxRank" should actually be "maxLvlRank" (since the lvlToLoop is indexed by (TensorId, Level))

[Peiming]

address comments.

[Peiming]

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

270

There is no src/dstLvlRank difference outside LoopEmitter. The bug is also unrelated to view-based reshaping.

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

1822–1823

There is no src/dstLvlRank difference outside loopEmitter

[wrengr]

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

259

"storage"

[Peiming]

fix typo.

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

public:
   ///
   /// for (const auto c0 : coordinates[0]) {
   ///   if (c0 == d0 + d1) {
   ///      generated_code;
   ///   }
   /// }
   ///
   /// to filter out coordinates that are not equal to the affine expression.
+  ///
+  /// The maxLvlRank specifies the max level rank of all inputs/output tensors.
+  /// It is used to pre-allocate sufficient memory for internal storage.
   //
   // TODO: we want to make the filter loop more efficient in the future,
   // e.g., by avoiding scanning the full list of stored coordinates (keeping
   // the last position in ordered list) or even apply binary search to find
   // the coordinate.
   //
   // TODO: would be cleaner to understand/document if the first argument
   // gave the number of input tensors, instead of the current number of
   // input+output tensors.
   Merger(unsigned numInputOutputTensors, unsigned numNativeLoops,
-         unsigned numFilterLoops);
+         unsigned numFilterLoops, unsigned maxLvlRank);
 
   /// Constructs a new tensor expression, and returns its identifier.
   /// The type of the `e0` argument varies according to the value of the
   /// `k` argument, as described by the `TensorExp` ctor.
   ExprId addExp(TensorExp::Kind k, unsigned e0, ExprId e1 = kInvalidId,
                 Value v = Value(), Operation *op = nullptr);
   ExprId addExp(TensorExp::Kind k, ExprId e, Value v, Operation *op = nullptr) {
     return addExp(k, e, kInvalidId, v, op);
std::optional<Level> getLvl(TensorLoopId b) const {
     return getLvl(tensor(b), loop(b));
   }
 
   /// Sets the level number and level-type of the `t`th tensor on
   /// `i`th loop.
   void setLevelAndType(TensorId t, LoopId i, Level lvl, DimLevelType dlt) {
     assert(t < numTensors && i < numLoops && lvl < lvlToLoop[t].size() &&
            isValidDLT(dlt));
     lvlTypes[t][i] = dlt;
     loopToLvl[t][i] = lvl;
     lvlToLoop[t][lvl] = i;
   }
 
   /// Iterates over a set of `TensorLoopId`s, invoking the callback
   /// for each `TensorLoopId` and passing it the corresponding tensor
   /// identifier, level, and level-type, following with a boolean value
   /// indicating whether it is a dependent index reduction loop condition.
   void foreachTensorLoopId(
       LatPointId p, function_ref<void(TensorLoopId, TensorId,
                                       std::optional<Level>, DimLevelType, bool)>
                         callback) {
     for (const TensorLoopId b : latPoints[p].bits.set_bits()) {
       TensorId t = tensor(b);

mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h

 /// convience methods that keep implementation details transparent
 /// to sparsification while asserting on internal consistency.
 class CodegenEnv {
 public:
   /// Constructs a code generation environment which can be
   /// passed around during sparsification for bookkeeping
   /// together with some consistency asserts.
   CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts,
-             unsigned numTensors, unsigned numLoops, unsigned numFilterLoops);
+             unsigned numTensors, unsigned numLoops, unsigned numFilterLoops,
+             unsigned maxRank);
 
   //
   // General methods.
   //
 
   LogicalResult initTensorExp();
   ExprId getExprId() const { return tensorExp; }
 

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

 }
 
 //===----------------------------------------------------------------------===//
 // Code generation environment constructor and general methods
 //===----------------------------------------------------------------------===//
 
 CodegenEnv::CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts,
                        unsigned numTensors, unsigned numLoops,
-                       unsigned numFilterLoops)
+                       unsigned numFilterLoops, unsigned maxRank)
     : linalgOp(linop), sparseOptions(opts),
-      latticeMerger(numTensors, numLoops, numFilterLoops), loopEmitter(),
-      topSort(), sparseOut(nullptr), outerParNest(-1u), insChain(), expValues(),
-      expFilled(), expAdded(), expCount(), redVal(), redExp(kInvalidId),
-      redCustom(kInvalidId), redValidLexInsert() {}
+      latticeMerger(numTensors, numLoops, numFilterLoops, maxRank),
+      loopEmitter(), topSort(), sparseOut(nullptr), outerParNest(-1u),
+      insChain(), expValues(), expFilled(), expAdded(), expCount(), redVal(),
+      redExp(kInvalidId), redCustom(kInvalidId), redValidLexInsert() {}
 
 LogicalResult CodegenEnv::initTensorExp() {
   // Builds the tensor expression for the Linalg operation in SSA form.
   std::optional<ExprId> optExp = latticeMerger.buildTensorExpFromLinalg(op());
   if (!optExp || !isAdmissibleTensorExp(*optExp))
     return failure();
 
   tensorExp = *optExp;

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

for (TensorId tid = 0; tid < numTensors; tid++) {
     highs[tid].assign(lvlRank, Value());
     segHi[tid].assign(lvlRank, Value());
     posits[tid].assign(lvlRank, Value());
     coords[tid].assign(lvlRank, Value());
     positionsBuffers[tid].assign(lvlRank, Value());
     coordinatesBuffers[tid].assign(lvlRank, Value());
     sliceOffsets[tid].assign(lvlRank, Value());
     sliceStrides[tid].assign(lvlRank, Value());
-
     dependentLvlMap[tid].assign(lvlRank,
                                 std::vector<std::pair<TensorId, Level>>());
-    if (dimGetter)
-      for (Level l = 0; l < lvlRank; l++)
+    if (dimGetter) {
+      auto reassoc = collapseReassoc[tid];
+      Level dstRank = reassoc ? reassoc.size() : lvlRank;
+      for (Level l = 0; l < dstRank; l++) {
         dependentLvlMap[tid][l] = dimGetter(tid, l);
+        // TODO: View-base collapse and dependent index reduction are not
+        // compatible right now.
+        assert(!reassoc || dependentLvlMap[tid][l].empty());
+      }
+    }
   }
 
   // Construct the inverse of the `topSort` from the sparsifier.
   // This is needed to map `AffineDimExpr`s back to the `LoopOrd`
   // used in loop emitter.
   // FIXME: This map should be maintained outside loop emitter.
   for (LoopOrd n = 0; n < numLoops; n++)
     loopIdToOrd[topSort[n]] = n;

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

LogicalResult matchAndRewrite(linalg::GenericOp op,
     // Sets up a code generation environment.
     const unsigned numTensors = op->getNumOperands();
     const unsigned numLoops = op.getNumLoops();
     const unsigned numFilterLoops = getNumNonTrivialIdxExpOnSparseLvls(op);
     // TODO: we should probably always use slice-based codegen whenever
     // possible, we can even intermix slice-based and filter-loop based codegen.
     bool idxReducBased = options.enableIndexReduction && numFilterLoops != 0;
 
+    // If we have indexing map like (d0) -> (0, d0), there might be more
+    // levels then loops because of the constant index, that means we can not
+    // use numLoops as the upper bound for ranks of all tensors.
+    // TODO: Constant indices are currently not support on sparse tensor, but
+    // are allowed in non-annotated dense tensor. Support it, it would be
+    // required for sparse tensor slice rank reducing too.
+    Level maxLvlRank = 0;
+    for (auto operand : op.getOperands()) {
+      if (auto rtp = operand.getType().dyn_cast<RankedTensorType>()) {
+        maxLvlRank = std::max(maxLvlRank, SparseTensorType(rtp).getLvlRank());
+      }
+    }
+
     // If we uses slice based algorithm for affine index, we do not need filter
     // loop.
     CodegenEnv env(op, options, numTensors, numLoops,
-                   /*numFilterLoops=*/idxReducBased ? 0 : numFilterLoops);
+                   /*numFilterLoops=*/idxReducBased ? 0 : numFilterLoops,
+                   maxLvlRank);
 
     // Detects sparse annotations and translates the per-level sparsity
     // information for all tensors to loop indices in the kernel.
     if (!findSparseAnnotations(env, idxReducBased))
       return failure();
 
     // Constructs the tensor expressions tree from `op`, returns failure if the
     // tree can not be built or the tensor expression is inadmissible.

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

LatPoint::LatPoint(unsigned numTensors, unsigned numLoops, TensorId t, LoopId i,
                    ExprId e)
     : bits(numLoops * numTensors, false), exp(e) {
   assert(t < numTensors && i < numLoops);
   const TensorLoopId b = numTensors * i + t;
   bits.set(b);
 }
 
 Merger::Merger(unsigned numInputOutputTensors, unsigned numNativeLoops,
-               unsigned numFilterLoops)
+               unsigned numFilterLoops, unsigned maxLvlRank)
     : outTensor(numInputOutputTensors - 1),
       syntheticTensor(numInputOutputTensors),
       numTensors(numInputOutputTensors + 1), numNativeLoops(numNativeLoops),
       numLoops(numNativeLoops + numFilterLoops), hasSparseOut(false),
       lvlTypes(numTensors,
                std::vector<DimLevelType>(numLoops, DimLevelType::Undef)),
       loopToLvl(numTensors,
                 std::vector<std::optional<Level>>(numLoops, std::nullopt)),
       lvlToLoop(numTensors,
-                std::vector<std::optional<LoopId>>(numLoops, std::nullopt)),
+                std::vector<std::optional<LoopId>>(maxLvlRank, std::nullopt)),
       loopToDependencies(numLoops, std::vector<std::optional<Level>>(
                                        numTensors, std::nullopt)),
       levelToDependentIdx(numTensors, std::vector<std::vector<LoopId>>(
-                                          numLoops, std::vector<LoopId>())),
+                                          maxLvlRank, std::vector<LoopId>())),
       loopBounds(numLoops, std::make_pair(numTensors, numLoops)) {}
 
 //===----------------------------------------------------------------------===//
 // Lattice methods.
 //===----------------------------------------------------------------------===//
 
 ExprId Merger::addExp(TensorExp::Kind k, unsigned x, ExprId y, Value v,
                       Operation *op) {

mlir/test/Dialect/SparseTensor/constant_index_map.mlir

+// Reported by https://github.com/llvm/llvm-project/issues/61530
+
+// RUN: mlir-opt %s -sparsification | FileCheck %s
+
+#map1 = affine_map<(d0) -> (0, d0)>
+#map2 = affine_map<(d0) -> (d0)>
+
+#SpVec = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
+
+// CHECK-LABEL:   func.func @main(
+// CHECK-SAME:      %[[VAL_0:.*0]]: tensor<1x77xi1>,
+// CHECK-SAME:      %[[VAL_1:.*1]]: tensor<1x77xi1>) -> tensor<77xi1, #{{.*}}> {
+// CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 77 : index
+// CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 0 : index
+// CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 1 : index
+// CHECK-DAG:       %[[VAL_5:.*]] = bufferization.alloc_tensor() : tensor<77xi1, #{{.*}}>
+// CHECK-DAG:       %[[VAL_6:.*]] = bufferization.to_memref %[[VAL_0]] : memref<1x77xi1>
+// CHECK-DAG:       %[[VAL_7:.*]] = bufferization.to_memref %[[VAL_1]] : memref<1x77xi1>
+// CHECK:           %[[VAL_8:.*]] = scf.for %[[VAL_9:.*]] = %[[VAL_3]] to %[[VAL_2]] step %[[VAL_4]] iter_args(%[[VAL_10:.*]] = %[[VAL_5]]) -> (tensor<77xi1, #{{.*}}>) {
+// CHECK:             %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_3]], %[[VAL_9]]] : memref<1x77xi1>
+// CHECK:             %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_3]], %[[VAL_9]]] : memref<1x77xi1>
+// CHECK:             %[[VAL_13:.*]] = arith.addi %[[VAL_11]], %[[VAL_12]] : i1
+// CHECK:             %[[VAL_14:.*]] = sparse_tensor.insert %[[VAL_13]] into %[[VAL_10]]{{\[}}%[[VAL_9]]] : tensor<77xi1, #{{.*}}>
+// CHECK:             scf.yield %[[VAL_14]] : tensor<77xi1, #{{.*}}>
+// CHECK:           }
+// CHECK:           %[[VAL_15:.*]] = sparse_tensor.load %[[VAL_16:.*]] hasInserts : tensor<77xi1, #{{.*}}>
+// CHECK:           return %[[VAL_15]] : tensor<77xi1, #{{.*}}>
+// CHECK:         }
+func.func @main(%arg0: tensor<1x77xi1>, %arg1: tensor<1x77xi1>) -> tensor<77xi1, #SpVec> {
+  %0 = bufferization.alloc_tensor() : tensor<77xi1, #SpVec>
+  %1 = linalg.generic {
+    indexing_maps = [#map1, #map1, #map2],
+    iterator_types = ["parallel"]}
+    ins(%arg0, %arg1 : tensor<1x77xi1>, tensor<1x77xi1>)
+    outs(%0 : tensor<77xi1, #SpVec>) {
+  ^bb0(%in: i1, %in_0: i1, %out: i1):
+    %2 = arith.addi %in, %in_0 : i1
+    linalg.yield %2 : i1
+  } -> tensor<77xi1, #SpVec>
+  return %1 : tensor<77xi1, #SpVec>
+}

mlir/test/Dialect/SparseTensor/sparse_affine.mlir

%0 = linalg.generic #trait5
     outs(%argx: tensor<32x16xf64>) {
       ^bb(%a: f64, %b: f64, %x: f64):
         %0 = arith.mulf %a, %b : f64
         %1 = arith.addf %x, %0 : f64
         linalg.yield %1 : f64
   } -> tensor<32x16xf64>
   return %0 : tensor<32x16xf64>
 }
-

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

 FOREVERY_BINOP(IMPL_BINOP_PATTERN)
 
 #undef IMPL_BINOP_PATTERN
 
 class MergerTestBase : public ::testing::Test {
 protected:
   MergerTestBase(unsigned numTensors, unsigned numLoops)
       : numTensors(numTensors), numLoops(numLoops),
-        merger(numTensors, numLoops, /*numFilterLoops=*/0) {}
+        merger(numTensors, numLoops, /*numFilterLoops=*/0,
+               /*maxRank=*/numLoops) {}
 
   ///
   /// Expression construction helpers.
   ///
 
   unsigned tensor(unsigned tensor) {
     return merger.addExp(TensorExp::Kind::kTensor, tensor);
   }