[Mlir-commits] [mlir] [mlir][SCF] Handle dropped operands in structural type conversion (PR #220473)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-scf

Author: Akun (AkunCube)

<details>
<summary>Changes</summary>

SCF structural type conversion assumes that `scf.for` bounds and `scf.if` conditions are converted to exactly one value. A 1:N conversion can drop or split these operands, causing `getSingleElement` to trigger an assertion.

Reject conversions with non-single-valued bounds or conditions and add regression coverage.

Fixes: https://github.com/llvm/llvm-project/issues/220459

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


3 Files Affected:

- (modified) mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp (+10) 
- (modified) mlir/test/Transforms/test-legalize-type-conversion.mlir (+27) 
- (modified) mlir/test/lib/Dialect/Test/TestPatterns.cpp (+9) 


``````````diff
diff --git a/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp b/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
index 9468927021495..3816007771ebe 100644
--- a/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
@@ -89,6 +89,14 @@ class ConvertForOpTypes
   std::optional<ForOp> convertSourceOp(ForOp op, OneToNOpAdaptor adaptor,
                                        ConversionPatternRewriter &rewriter,
                                        TypeRange dstTypes) const {
+    // Loop bounds are single operands in the SCF IR. A 1:N conversion may
+    // produce zero or multiple values, in which case this operation cannot be
+    // reconstructed.
+    if (!llvm::hasSingleElement(adaptor.getLowerBound()) ||
+        !llvm::hasSingleElement(adaptor.getUpperBound()) ||
+        !llvm::hasSingleElement(adaptor.getStep()))
+      return std::nullopt;
+
     // Create a empty new op and inline the regions from the old op.
     //
     // This is a little bit tricky. We have two concerns here:
@@ -141,6 +149,8 @@ class ConvertIfOpTypes
   std::optional<IfOp> convertSourceOp(IfOp op, OneToNOpAdaptor adaptor,
                                       ConversionPatternRewriter &rewriter,
                                       TypeRange dstTypes) const {
+    if (!llvm::hasSingleElement(adaptor.getCondition()))
+      return std::nullopt;
 
     IfOp newOp =
         IfOp::create(rewriter, op.getLoc(), dstTypes,
diff --git a/mlir/test/Transforms/test-legalize-type-conversion.mlir b/mlir/test/Transforms/test-legalize-type-conversion.mlir
index 7a8bf06a3a060..352483520ff74 100644
--- a/mlir/test/Transforms/test-legalize-type-conversion.mlir
+++ b/mlir/test/Transforms/test-legalize-type-conversion.mlir
@@ -113,6 +113,33 @@ func.func @test_signature_conversion_no_converter() {
 
 // -----
 
+// Do not assert when type conversion drops an SCF loop bound.
+func.func @test_scf_for_dropped_bound(%arg0: i32) {
+  // expected-error at +1 {{failed to legalize operation 'scf.for' that was explicitly marked illegal}}
+  %0 = scf.for %iv = %arg0 to %arg0 step %arg0 iter_args(%arg = %arg0) -> (i32) : i32 {
+    scf.yield %arg : i32
+  }
+  return
+}
+
+// -----
+
+// Do not assert when type conversion drops an SCF if condition.
+func.func @test_scf_if_dropped_condition() {
+  %cond = "test.context_op"() {drop_i1} : () -> i1
+  // expected-error at +1 {{failed to legalize operation 'scf.if' that was explicitly marked illegal}}
+  %result = scf.if %cond -> (i32) {
+    %value = "test.context_op"() : () -> i32
+    scf.yield %value : i32
+  } else {
+    %value = "test.context_op"() : () -> i32
+    scf.yield %value : i32
+  }
+  return
+}
+
+// -----
+
 // CHECK-LABEL: @recursive_type_conversion
 func.func @recursive_type_conversion() {
   // CHECK:  !test.test_rec<outer_converted_type, smpla>
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index fbc5ab76d9ba1..06517910e270e 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -2213,6 +2213,15 @@ struct TestTypeConversionDriver
       return IntegerType::get(v.getContext(),
                               intType.getWidth() + incrementAttr.getInt());
     });
+    converter.addConversion(
+        [](Value v, SmallVectorImpl<Type> &) -> std::optional<LogicalResult> {
+          // Test dropping an i1 value used as an SCF condition.
+          Operation *definingOp = v.getDefiningOp();
+          if (!definingOp || !v.getType().isInteger(1) ||
+              !definingOp->hasAttr("drop_i1"))
+            return std::nullopt;
+          return success();
+        });
 
     /// Add the legal set of type materializations.
     converter.addSourceMaterialization([](OpBuilder &builder, Type resultType,

``````````

</details>


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