[mlir][emitc] Fix emitc.for verification crash

[CoTinker]

This PR adds block arguments check to prevent a crash in emitc.for verifier. Fixes #159950.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-emitc

Author: Longsheng Mou (CoTinker)

Changes

This PR adds block arguments check to prevent a crash in emitc.for verifier. Fixes #159950.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/EmitC/IR/EmitC.cpp (+4)
• (modified) mlir/test/Dialect/EmitC/invalid_ops.mlir (+25)

diff --git a/mlir/lib/Dialect/EmitC/IR/EmitC.cpp b/mlir/lib/Dialect/EmitC/IR/EmitC.cpp
index 4754f0bfe895e..77a56fb463243 100644
--- a/mlir/lib/Dialect/EmitC/IR/EmitC.cpp
+++ b/mlir/lib/Dialect/EmitC/IR/EmitC.cpp
@@ -584,6 +584,10 @@ void ForOp::print(OpAsmPrinter &p) {
 LogicalResult ForOp::verifyRegions() {
   // Check that the body defines as single block argument for the induction
   // variable.
+  if (getBody()->getNumArguments() == 0)
+    return emitOpError("expected body to have a single block argument for the "
+                       "induction variable");
+
   if (getInductionVar().getType() != getLowerBound().getType())
     return emitOpError(
         "expected induction variable to be same type as bounds and step");
diff --git a/mlir/test/Dialect/EmitC/invalid_ops.mlir b/mlir/test/Dialect/EmitC/invalid_ops.mlir
index 5f594fb08c43f..f30a614a93de0 100644
--- a/mlir/test/Dialect/EmitC/invalid_ops.mlir
+++ b/mlir/test/Dialect/EmitC/invalid_ops.mlir
@@ -876,3 +876,28 @@ func.func @test_do(%arg0 : !emitc.ptr<i32>) {
 
   return
 }
+
+// -----
+
+func.func @test_for_none_block_argument(%arg0: index) {
+  // expected-error@+1 {{expected body to have a single block argument for the induction variable}}
+  "emitc.for"(%arg0, %arg0, %arg0) (
+    {
+      emitc.yield
+    }
+  ) : (index, index, index) -> ()
+  return
+}
+
+// -----
+
+func.func @test_for_unmatch_type(%arg0: index) {
+  // expected-error@+1 {{expected induction variable to be same type as bounds}}
+  "emitc.for"(%arg0, %arg0, %arg0) (
+    {
+    ^bb0(%i0 : f32):
+      emitc.yield
+    }
+  ) : (index, index, index) -> ()
+  return
+}

[aniragil]

Thanks for fixing this!

[aniragil]

Suggested change

	if (getBody()->getNumArguments() == 0)
	if (getBody()->getNumArguments() != 1)

[aniragil]

Worth adding a third case with more than one argument.