[mlir][sparse] Add checking parent op of SortOp

[sotoshigoto]

Summary

Fix crash with segmentation fault caused by setting a parent operator
that is not func::FuncOp with sparse_tensor SortOp.

fixes https://github.com/llvm/llvm-project/issues/59988

[aartbik]

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

1023 ↗

(On Diff #496834)

Thanks for filing the bug and looking into this. The code crashes since the rewriting indeed assumes to find a func op during codegen, in particular

auto insertPoint = op->template getParentOfType<func::FuncOp>();

during SparseBufferRewriting. However, it seems we need to fix this at the rewriting level, not in verification of the op (which is really a user facing validation).

[sotoshigoto]

Move a fix point from verification of op to rewriting level

[sotoshigoto]

remove unnecessary an include

[wrengr]

This comment was removed by wrengr.

[wrengr]

Since matchAndRewriteSortOp was already getting the insertPoint, it seems perfectly reasonable to have the function return failure if it can't do so. I also like that it catches the error exactly where it is introduced: if we changed the function to compute insertPoint differently, then the check is still valid (and if we change the function to not compute insertPoint at all, then it's also clear that we should remove the check).

Thanks :)

[sotoshigoto]

@aartbik @wrengr Could you land this patch? I don't have commit access.

[aartbik]

In D143874#4133682, @sotoshigoto wrote:

@aartbik @wrengr Could you land this patch? I don't have commit access.

Thanks Alex for landing! I missed this request.

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

if (!mtp.isDynamicDim(0)) {
       auto newMtp =
           MemRefType::get({ShapedType::kDynamic}, mtp.getElementType());
       v = rewriter.create<memref::CastOp>(loc, newMtp, v);
     }
     operands.push_back(v);
   }
 
   auto insertPoint = op->template getParentOfType<func::FuncOp>();
+  if (!insertPoint)
+    return failure();
+
   SmallString<32> funcName;
   FuncGeneratorType funcGenerator;
   uint32_t nTrailingP = 0;
   switch (op.getAlgorithm()) {
   case SparseTensorSortKind::HybridQuickSort: {
     funcName = kHybridQuickSortFuncNamePrefix;
     funcGenerator = createQuickSortFunc;
     nTrailingP = 1;