[Mlir-commits] [mlir] [mlir][sparse] introduce sparse_tensor.lvl operation. (PR #69978)

[Yinying Li]

================
@@ -1208,6 +1208,84 @@ LogicalResult CrdTranslateOp::fold(FoldAdaptor adaptor,
   return success();
 }
 
+LogicalResult LvlOp::verify() {
+  if (std::optional<uint64_t> lvl = getConstantLvlIndex()) {
+    auto stt = getSparseTensorType(getSource());
+    if (static_cast<uint64_t>(lvl.value()) >= stt.getLvlRank())
+      emitError("Level index exceeds the rank of the input sparse tensor");
+  }
+  return success();
+}
+
+std::optional<uint64_t> LvlOp::getConstantLvlIndex() {
+  return getConstantIntValue(getIndex());
+}
+
+Speculation::Speculatability LvlOp::getSpeculatability() {
+  auto constantIndex = getConstantLvlIndex();
+  if (!constantIndex)
+    return Speculation::NotSpeculatable;
+
+  assert(constantIndex <
+         cast<RankedTensorType>(getSource().getType()).getRank());
+  return Speculation::Speculatable;
+}
+
+OpFoldResult LvlOp::fold(FoldAdaptor adaptor) {
+  auto lvlIndex = llvm::dyn_cast_if_present<IntegerAttr>(adaptor.getIndex());
+  if (!lvlIndex)
+    return {};
+
+  Level lvl = lvlIndex.getAPSInt().getZExtValue();
+  auto stt = getSparseTensorType(getSource());
+  if (lvl >= stt.getLvlRank()) {
+    // Follows the same convention used by tensor.dim operation. Out of bound
+    // indices produce undefined behavior but are still valid IR. Don't choke on
+    // them.
+    return {};
+  }
+
+  // Helper lambda to build an IndexAttr;
----------------
yinying-lisa-li wrote:

nit: `.`?

https://github.com/llvm/llvm-project/pull/69978