[mlir][sparse] Fix a crash if block not have terminator (#180741)
This PR fixes a crash in `verifyNumBlockArgs` if region not end with a
terminator. Fixes #180720.
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index f6a707f..3a34ad9 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -1746,7 +1746,12 @@
return op->emitError() << regionName << " region argument " << (i + 1)
<< " type mismatch";
}
- Operation *term = region.front().getTerminator();
+ Block &block = region.front();
+ if (!block.mightHaveTerminator())
+ return op->emitError() << regionName
+ << " region must end with a terminator";
+
+ Operation *term = block.getTerminator();
YieldOp yield = dyn_cast<YieldOp>(term);
if (!yield)
return op->emitError() << regionName
diff --git a/mlir/test/Dialect/SparseTensor/invalid.mlir b/mlir/test/Dialect/SparseTensor/invalid.mlir
index 4c37fc6..30c74bb 100644
--- a/mlir/test/Dialect/SparseTensor/invalid.mlir
+++ b/mlir/test/Dialect/SparseTensor/invalid.mlir
@@ -631,6 +631,20 @@
// -----
+func.func @invalid_reduce_wrong_terminator(%arg0: f64, %arg1: f64) -> f64 {
+ %cf1 = arith.constant 1.0 : f64
+ // expected-error@+1 {{reduce region must end with a terminator}}
+ %r = sparse_tensor.reduce %arg0, %arg1, %cf1 : f64 {
+ ^bb0(%arg2: f64, %arg3: f64):
+ %0 = arith.addf %arg2, %arg3 : f64
+ sparse_tensor.yield %0 : f64
+ %1 = arith.fptosi %0 : f64 to i32
+ }
+ return %r : f64
+}
+
+// -----
+
func.func @invalid_select_num_args_mismatch(%arg0: f64) -> f64 {
// expected-error@+1 {{select region must have exactly 1 arguments}}
%r = sparse_tensor.select %arg0 : f64 {