[mlir][SparseTensor] Fix insertion point in createQuickSort

[matthias-springer]

createQuickSort used to generate invalid IR:

"func.func"() <{function_type = (index, index, memref<?xindex>, memref<?xf32>, memref<?xi32>) -> (), sym_name = "_sparse_qsort_0_1_index_coo_1_f32_i32", sym_visibility = "private"}> ({
^bb0(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>):
  %0:2 = "scf.while"(%arg0, %arg1) ({
  ^bb0(%arg5: index, %arg6: index):
    // ...
    "scf.condition"(%3, %arg5, %arg6) : (i1, index, index) -> ()
  }, {
  ^bb0(%arg5: index, %arg6: index):
    // ...
    %7:2 = "scf.if"(%6) ({
      %8 = "arith.cmpi"(%2, %3) <{predicate = 7 : i64}> : (index, index) -> i1
      // ...
      "scf.yield"(%9#0, %9#1) : (index, index) -> ()
      %10 = "arith.constant"() <{value = 0 : index}> : () -> index
    }, {
      "scf.yield"(%arg5, %arg5) : (index, index) -> ()
    }) : (i1) -> (index, index)
    "scf.yield"(%7#0, %7#1) : (index, index) -> ()
  }) : (index, index) -> (index, index)
  "func.return"() : () -> ()
}) : () -> ()
within split at mlir/test/Dialect/SparseTensor/buffer_rewriting.mlir:76 offset :11:1: error: 'scf.yield' op must be the last operation in the parent block

This commit fixes tests such as mlir/test/Dialect/SparseTensor/buffer_rewriting.mlir when verifying the IR after each pattern application (#74270).

[llvmbot]

@llvm/pr-subscribers-mlir-sparse

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

createQuickSort used to generate invalid IR:

"func.func"() <{function_type = (index, index, memref<?xindex>, memref<?xf32>, memref<?xi32>) -> (), sym_name = "_sparse_qsort_0_1_index_coo_1_f32_i32", sym_visibility = "private"}> ({
^bb0(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>):
  %0:2 = "scf.while"(%arg0, %arg1) ({
  ^bb0(%arg5: index, %arg6: index):
    // ...
    "scf.condition"(%3, %arg5, %arg6) : (i1, index, index) -> ()
  }, {
  ^bb0(%arg5: index, %arg6: index):
    // ...
    %7:2 = "scf.if"(%6) ({
      %8 = "arith.cmpi"(%2, %3) <{predicate = 7 : i64}> : (index, index) -> i1
      // ...
      "scf.yield"(%9#<!-- -->0, %9#<!-- -->1) : (index, index) -&gt; ()
      %10 = "arith.constant"() &lt;{value = 0 : index}&gt; : () -&gt; index
    }, {
      "scf.yield"(%arg5, %arg5) : (index, index) -&gt; ()
    }) : (i1) -&gt; (index, index)
    "scf.yield"(%7#<!-- -->0, %7#<!-- -->1) : (index, index) -&gt; ()
  }) : (index, index) -&gt; (index, index)
  "func.return"() : () -&gt; ()
}) : () -&gt; ()
within split at mlir/test/Dialect/SparseTensor/buffer_rewriting.mlir:76 offset :11:1: error: 'scf.yield' op must be the last operation in the parent block

This commit fixes tests such as mlir/test/Dialect/SparseTensor/buffer_rewriting.mlir when verifying the IR after each pattern application (#74270).

---

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

1 Files Affected:

• (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp (+1-1)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp
index 463a49f52283a..cdbf4f048a00f 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp
@@ -952,9 +952,9 @@ createQuickSort(OpBuilder &builder, ModuleOp module, func::FuncOp func,
   Value cond = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ule,
                                              lenLow, lenHigh);
 
+  Value c0 = constantIndex(builder, loc, 0);
   scf::IfOp ifOp = builder.create<scf::IfOp>(loc, types, cond, /*else=*/true);
 
-  Value c0 = constantIndex(builder, loc, 0);
   auto mayRecursion = [&](Value low, Value high, Value len) {
     Value cond =
         builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ne, len, c0);