[Mlir-commits] [mlir] [mlir][sparse] Fix a crash if block not have terminator (PR #180741)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-sparse

Author: Longsheng Mou (CoTinker)

<details>
<summary>Changes</summary>

This PR fixes a crash in `verifyNumBlockArgs` if region not end with a terminator. Fixes #<!-- -->180720.

---
Full diff: https://github.com/llvm/llvm-project/pull/180741.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp (+6-1) 
- (modified) mlir/test/Dialect/SparseTensor/invalid.mlir (+14) 


``````````diff
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index f6a707f53eb1b..3a34ad90941b0 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -1746,7 +1746,12 @@ static LogicalResult verifyNumBlockArgs(T *op, Region &region,
       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 4c37fc6882bab..30c74bbcdbf90 100644
--- a/mlir/test/Dialect/SparseTensor/invalid.mlir
+++ b/mlir/test/Dialect/SparseTensor/invalid.mlir
@@ -631,6 +631,20 @@ func.func @invalid_reduce_wrong_yield(%arg0: f64, %arg1: f64) -> f64 {
 
 // -----
 
+func.func @invalid_reduce_wrong_terminator(%arg0: f64, %arg1: f64) -> f64 {
+  %cf1 = arith.constant 1.0 : f64
+  // expected-error at +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 at +1 {{select region must have exactly 1 arguments}}
   %r = sparse_tensor.select %arg0 : f64 {

``````````

</details>


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