[mlir][sparse] Modify the pivot selection method for quick sort.

[bixia]

Summary

Previously, we choose the median of three values. We now choose the median of
five values when the number of values being sorted exceed a threshold
(currently 100). This is similar to std::sort.

[aartbik]

We should perhaps start thinking about moving all sorted related codegen to a more general place and in its own file.
(not know, but this is becoming a powerful utility that others may want to use too)

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp

438

The new insertion point is right ....

[bixia]

Fixed comment. Rebased. Fixed FileCheck test.

[bixia]

Fix a comment.

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp

createScanLoop(OpBuilder &builder, ModuleOp module, func::FuncOp func,
   compareOperands[0] = i;
   compareOperands[1] = p;
   Value compareEq =
       createInlinedEqCompare(builder, loc, compareOperands, nx, ny, isCoo);
 
   return std::make_pair(whileOp.getResult(0), compareEq);
 }
 
-/// Creates a code block to swap the values so that data[mi] is the median among
-/// data[lo], data[hi], and data[mi].
-//  The generated code corresponds to this C-like algorithm:
-//  median = mi
-//  if (data[mi] < data[lo]).                               (if1)
-//    if (data[hi] < data[lo])                              (if2)
-//       median = data[hi] < data[mi] ? mi : hi
-//    else
-//       median = lo
-//  else
-//    if data[hi] < data[mi]                                (if3)
-//      median = data[hi] < data[lo] ? lo : hi
-//  if median != mi swap data[median] with data[mi]
+/// Creates and returns an IfOp to compare two elements and swap the elements
+/// if compareFunc(data[b], data[a]) returns true. The new insertion point is
+/// right after the swap instructions.
+static scf::IfOp createCompareThenSwap(OpBuilder &builder, Location loc,
+                                       uint64_t nx, uint64_t ny, bool isCoo,
+                                       SmallVectorImpl<Value> &swapOperands,
+                                       SmallVectorImpl<Value> &compareOperands,
+                                       Value a, Value b) {
+  // Compare(data[b], data[a]).
+  compareOperands[0] = b;
+  compareOperands[1] = a;
+  Value cond =
+      createInlinedLessThan(builder, loc, compareOperands, nx, ny, isCoo);
+  scf::IfOp ifOp = builder.create<scf::IfOp>(loc, cond, /*else=*/false);
+  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+  swapOperands[0] = b;
+  swapOperands[1] = a;
+  createSwap(builder, loc, swapOperands, nx, ny, isCoo);
+  return ifOp;
+}
+
+/// Creates code to insert the 3rd element to a list of two sorted elements.
+static void createInsert3rd(OpBuilder &builder, Location loc, uint64_t nx,
+                            uint64_t ny, bool isCoo,
+                            SmallVectorImpl<Value> &swapOperands,
+                            SmallVectorImpl<Value> &compareOperands, Value v0,
+                            Value v1, Value v2) {
+  scf::IfOp ifOp = createCompareThenSwap(builder, loc, nx, ny, isCoo,
+                                         swapOperands, compareOperands, v1, v2);
+  createCompareThenSwap(builder, loc, nx, ny, isCoo, swapOperands,
+                        compareOperands, v0, v1);
+  builder.setInsertionPointAfter(ifOp);
+}
+
+/// Creates code to sort 3 elements.
+static void createSort3(OpBuilder &builder, Location loc, uint64_t nx,
+                        uint64_t ny, bool isCoo,
+                        SmallVectorImpl<Value> &swapOperands,
+                        SmallVectorImpl<Value> &compareOperands, Value v0,
+                        Value v1, Value v2) {
+  // Sort the first 2 elements.
+  scf::IfOp ifOp1 = createCompareThenSwap(
+      builder, loc, nx, ny, isCoo, swapOperands, compareOperands, v0, v1);
+  builder.setInsertionPointAfter(ifOp1);
+
+  // Insert the 3th element.
+  createInsert3rd(builder, loc, nx, ny, isCoo, swapOperands, compareOperands,
+                  v0, v1, v2);
+}
+
+/// Creates code to sort 5 elements.
+static void createSort5(OpBuilder &builder, Location loc, uint64_t nx,
+                        uint64_t ny, bool isCoo,
+                        SmallVectorImpl<Value> &swapOperands,
+                        SmallVectorImpl<Value> &compareOperands, Value v0,
+                        Value v1, Value v2, Value v3, Value v4) {
+  // Sort the first 3 elements.
+  createSort3(builder, loc, nx, ny, isCoo, swapOperands, compareOperands, v0,
+              v1, v2);
+
+  auto insert4th = [&]() {
+    scf::IfOp ifOp = createCompareThenSwap(
+        builder, loc, nx, ny, isCoo, swapOperands, compareOperands, v2, v3);
+    createInsert3rd(builder, loc, nx, ny, isCoo, swapOperands, compareOperands,
+                    v0, v1, v2);
+    builder.setInsertionPointAfter(ifOp);
+  };
+
+  // Insert the 4th element.
+  insert4th();
+
+  // Insert the 5th element.
+  scf::IfOp ifOp = createCompareThenSwap(builder, loc, nx, ny, isCoo,
+                                         swapOperands, compareOperands, v3, v4);
+  insert4th();
+  builder.setInsertionPointAfter(ifOp);
+}
+
+/// Creates a code block to swap the values in indices lo, mi, and hi so that
+/// data[lo], data[mi] and data[hi] are sorted in non-decreasing values. When
+/// the number of values in range [lo, hi) is more than a threshold, we also
+/// include the middle of [lo, mi) and [mi, hi) and sort a total of five values.
 static void createChoosePivot(OpBuilder &builder, ModuleOp module,
                               func::FuncOp func, uint64_t nx, uint64_t ny,
                               bool isCoo, Value lo, Value hi, Value mi,
                               ValueRange args) {
   SmallVector<Value> compareOperands{mi, lo};
   uint64_t numXBuffers = isCoo ? 1 : nx;
   compareOperands.append(args.begin() + xStartIdx,
                          args.begin() + xStartIdx + numXBuffers);
-  Type i1Type = IntegerType::get(module.getContext(), 1, IntegerType::Signless);
-  SmallVector<Type, 1> cmpTypes{i1Type};
+  SmallVector<Value> swapOperands{mi, lo};
+  swapOperands.append(args.begin() + xStartIdx, args.end());
   Location loc = func.getLoc();
-  // Compare data[mi] < data[lo].
-  Value cond1 =
-      createInlinedLessThan(builder, loc, compareOperands, nx, ny, isCoo);
-  SmallVector<Type, 1> ifTypes{lo.getType()};
-  scf::IfOp ifOp1 =
-      builder.create<scf::IfOp>(loc, ifTypes, cond1, /*else=*/true);
-
-  // Generate an if-stmt to find the median value, assuming we already know that
-  // data[b] < data[a] and we haven't compare data[c] yet.
-  auto createFindMedian = [&](Value a, Value b, Value c) -> scf::IfOp {
-    compareOperands[0] = c;
-    compareOperands[1] = a;
-    // Compare data[c] < data[b].
-    Value cond2 =
-        createInlinedLessThan(builder, loc, compareOperands, nx, ny, isCoo);
-    scf::IfOp ifOp2 =
-        builder.create<scf::IfOp>(loc, ifTypes, cond2, /*else=*/true);
-    builder.setInsertionPointToStart(&ifOp2.getThenRegion().front());
-    compareOperands[0] = c;
-    compareOperands[1] = b;
-    // Compare data[c] < data[b].
-    Value cond3 =
-        createInlinedLessThan(builder, loc, compareOperands, nx, ny, isCoo);
-    builder.create<scf::YieldOp>(
-        loc, ValueRange{builder.create<arith::SelectOp>(loc, cond3, b, c)});
-    builder.setInsertionPointToStart(&ifOp2.getElseRegion().front());
-    builder.create<scf::YieldOp>(loc, ValueRange{a});
-    return ifOp2;
-  };
-
-  builder.setInsertionPointToStart(&ifOp1.getThenRegion().front());
-  scf::IfOp ifOp2 = createFindMedian(lo, mi, hi);
-  builder.setInsertionPointAfter(ifOp2);
-  builder.create<scf::YieldOp>(loc, ValueRange{ifOp2.getResult(0)});
-
-  builder.setInsertionPointToStart(&ifOp1.getElseRegion().front());
-  scf::IfOp ifOp3 = createFindMedian(mi, lo, hi);
+  Value c1 = constantIndex(builder, loc, 1);
+  Value hiP1 = builder.create<arith::AddIOp>(loc, hi, c1);
+  Value len = builder.create<arith::SubIOp>(loc, hiP1, lo);
+  Value lenThreshold = constantIndex(builder, loc, 1000);
+  Value lenCond = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
+                                                len, lenThreshold);
+  scf::IfOp lenIf = builder.create<scf::IfOp>(loc, lenCond, /*else=*/true);
 
-  builder.setInsertionPointAfter(ifOp3);
-  builder.create<scf::YieldOp>(loc, ValueRange{ifOp3.getResult(0)});
+  // When len < 1000, choose pivot from median of 3 values.
+  builder.setInsertionPointToStart(&lenIf.getThenRegion().front());
+  createSort3(builder, loc, nx, ny, isCoo, swapOperands, compareOperands, lo,
+              mi, hi);
 
-  builder.setInsertionPointAfter(ifOp1);
-  Value median = ifOp1.getResult(0);
-  Value cond =
-      builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ne, mi, median);
-  scf::IfOp ifOp =
-      builder.create<scf::IfOp>(loc, TypeRange(), cond, /*else=*/false);
+  // When len >= 1000, choose pivot from median of 5 values.
+  builder.setInsertionPointToStart(&lenIf.getElseRegion().front());
+  Value miP1 = builder.create<arith::AddIOp>(loc, hi, c1);
+  Value a = builder.create<arith::AddIOp>(loc, lo, miP1);
+  // Value a is the middle between [loc, mi].
+  a = builder.create<arith::ShRUIOp>(loc, a, c1);
+  Value b = builder.create<arith::AddIOp>(loc, mi, hiP1);
+  // Value b is the middle between [mi, hi].
+  b = builder.create<arith::ShRUIOp>(loc, b, c1);
+  createSort5(builder, loc, nx, ny, isCoo, swapOperands, compareOperands, lo, a,
+              mi, b, hi);
 
-  SmallVector<Value> swapOperands{median, mi};
-  swapOperands.append(args.begin() + xStartIdx, args.end());
-  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
-  createSwap(builder, loc, swapOperands, nx, ny, isCoo);
-  builder.setInsertionPointAfter(ifOp);
+  builder.setInsertionPointAfter(lenIf);
 }
 
 /// Creates a function to perform quick sort partition on the values in the
 /// range of index [lo, hi), assuming lo < hi.
 //
 // The generated IR corresponds to this C like algorithm:
 // int partition(lo, hi, xs) {
 //   p = (lo+hi)/2  // pivot index

mlir/test/Dialect/SparseTensor/buffer_rewriting.mlir

 func.func @sparse_push_back_inbound(%arg0: index, %arg1: memref<?xf64>, %arg2: f64) -> (memref<?xf64>, index) {
   %0:2 = sparse_tensor.push_back inbounds %arg0, %arg1, %arg2 : index, memref<?xf64>, f64
   return %0#0, %0#1 : memref<?xf64>, index
 }
 
 // -----
 
 // CHECK-LABEL:   func.func private @_sparse_partition_1_i8_f32_index(
-// CHECK-SAME:    %[[VAL_0:.*0]]: index,
-// CHECK-SAME:    %[[VAL_1:.*1]]: index,
-// CHECK-SAME:    %[[VAL_2:.*2]]: memref<?xi8>,
-// CHECK-SAME:    %[[VAL_3:.*3]]: memref<?xf32>,
-// CHECK-SAME:    %[[VAL_4:.*4]]: memref<?xindex>) -> index {
+// CHECK-SAME:                                                        %[[VAL_0:.*0]]: index,
+// CHECK-SAME:                                                        %[[VAL_1:.*1]]: index,
+// CHECK-SAME:                                                        %[[VAL_2:.*2]]: memref<?xi8>,
+// CHECK-SAME:                                                        %[[VAL_3:.*3]]: memref<?xf32>,
+// CHECK-SAME:                                                        %[[VAL_4:.*4]]: memref<?xindex>) -> index {
 // CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 1
-// CHECK-DAG:       %[[VAL_6:.*]] = arith.constant -1
-// CHECK:           %[[VAL_7:.*]] = arith.addi %[[VAL_0]], %[[VAL_1]]
-// CHECK:           %[[VAL_8:.*]] = arith.shrui %[[VAL_7]], %[[VAL_5]]
-// CHECK:           %[[VAL_9:.*]] = arith.subi %[[VAL_1]], %[[VAL_5]]
-// CHECK:           %[[VAL_10:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_8]]]
-// CHECK:           %[[VAL_11:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
-// CHECK:           %[[VAL_12:.*]] = arith.cmpi ult, %[[VAL_10]], %[[VAL_11]]
-// CHECK:           %[[VAL_13:.*]] = scf.if %[[VAL_12]] -> (index) {
-// CHECK:             %[[VAL_14:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
-// CHECK:             %[[VAL_15:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
-// CHECK:             %[[VAL_16:.*]] = arith.cmpi ult, %[[VAL_14]], %[[VAL_15]]
-// CHECK:             %[[VAL_17:.*]] = scf.if %[[VAL_16]] -> (index) {
-// CHECK:               %[[VAL_18:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
-// CHECK:               %[[VAL_19:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_8]]]
-// CHECK:               %[[VAL_20:.*]] = arith.cmpi ult, %[[VAL_18]], %[[VAL_19]]
-// CHECK:               %[[VAL_21:.*]] = arith.select %[[VAL_20]], %[[VAL_8]], %[[VAL_9]]
-// CHECK:               scf.yield %[[VAL_21]]
-// CHECK:             } else {
-// CHECK:               scf.yield %[[VAL_0]]
+// CHECK-DAG:       %[[VAL_6:.*]] = arith.constant 1000
+// CHECK-DAG:       %[[VAL_7:.*]] = arith.constant -1
+// CHECK:           %[[VAL_8:.*]] = arith.addi %[[VAL_0]], %[[VAL_1]]
+// CHECK:           %[[VAL_9:.*]] = arith.shrui %[[VAL_8]], %[[VAL_5]]
+// CHECK:           %[[VAL_10:.*]] = arith.subi %[[VAL_1]], %[[VAL_5]]
+// CHECK:           %[[VAL_11:.*]] = arith.subi %[[VAL_1]], %[[VAL_0]]
+// CHECK:           %[[VAL_12:.*]] = arith.cmpi ult, %[[VAL_11]], %[[VAL_6]]
+// CHECK:           scf.if %[[VAL_12]] {
+// CHECK:             %[[VAL_13:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:             %[[VAL_14:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:             %[[VAL_15:.*]] = arith.cmpi ult, %[[VAL_13]], %[[VAL_14]]
+// CHECK:             scf.if %[[VAL_15]] {
+// CHECK:               %[[VAL_16:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_17:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:               memref.store %[[VAL_17]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_16]], %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:               %[[VAL_18:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_19:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_0]]]
+// CHECK:               memref.store %[[VAL_19]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_18]], %[[VAL_3]]{{\[}}%[[VAL_0]]]
+// CHECK:               %[[VAL_20:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_21:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_0]]]
+// CHECK:               memref.store %[[VAL_21]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_20]], %[[VAL_4]]{{\[}}%[[VAL_0]]]
 // CHECK:             }
-// CHECK:             scf.yield %[[VAL_22:.*]]
-// CHECK:           } else {
+// CHECK:             %[[VAL_22:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_10]]]
 // CHECK:             %[[VAL_23:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
-// CHECK:             %[[VAL_24:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_8]]]
-// CHECK:             %[[VAL_25:.*]] = arith.cmpi ult, %[[VAL_23]], %[[VAL_24]]
-// CHECK:             %[[VAL_26:.*]] = scf.if %[[VAL_25]] -> (index) {
-// CHECK:               %[[VAL_27:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
-// CHECK:               %[[VAL_28:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
-// CHECK:               %[[VAL_29:.*]] = arith.cmpi ult, %[[VAL_27]], %[[VAL_28]]
-// CHECK:               %[[VAL_30:.*]] = arith.select %[[VAL_29]], %[[VAL_0]], %[[VAL_9]]
-// CHECK:               scf.yield %[[VAL_30]]
-// CHECK:             } else {
-// CHECK:               scf.yield %[[VAL_8]]
+// CHECK:             %[[VAL_24:.*]] = arith.cmpi ult, %[[VAL_22]], %[[VAL_23]]
+// CHECK:             scf.if %[[VAL_24]] {
+// CHECK:               %[[VAL_25:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_10]]]
+// CHECK:               %[[VAL_26:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_26]], %[[VAL_2]]{{\[}}%[[VAL_10]]]
+// CHECK:               memref.store %[[VAL_25]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_27:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_10]]]
+// CHECK:               %[[VAL_28:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_28]], %[[VAL_3]]{{\[}}%[[VAL_10]]]
+// CHECK:               memref.store %[[VAL_27]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_29:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_10]]]
+// CHECK:               %[[VAL_30:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_30]], %[[VAL_4]]{{\[}}%[[VAL_10]]]
+// CHECK:               memref.store %[[VAL_29]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_31:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_32:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:               %[[VAL_33:.*]] = arith.cmpi ult, %[[VAL_31]], %[[VAL_32]]
+// CHECK:               scf.if %[[VAL_33]] {
+// CHECK:                 %[[VAL_34:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_35:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:                 memref.store %[[VAL_35]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_34]], %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:                 %[[VAL_36:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_37:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_0]]]
+// CHECK:                 memref.store %[[VAL_37]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_36]], %[[VAL_3]]{{\[}}%[[VAL_0]]]
+// CHECK:                 %[[VAL_38:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_39:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_0]]]
+// CHECK:                 memref.store %[[VAL_39]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_38]], %[[VAL_4]]{{\[}}%[[VAL_0]]]
+// CHECK:               }
+// CHECK:             }
+// CHECK:           } else {
+// CHECK:             %[[VAL_40:.*]] = arith.addi %[[VAL_9]], %[[VAL_1]]
+// CHECK:             %[[VAL_41:.*]] = arith.shrui %[[VAL_40]], %[[VAL_5]]
+// CHECK:             %[[VAL_42:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:             %[[VAL_43:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:             %[[VAL_44:.*]] = arith.cmpi ult, %[[VAL_42]], %[[VAL_43]]
+// CHECK:             scf.if %[[VAL_44]] {
+// CHECK:               %[[VAL_45:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_46:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:               memref.store %[[VAL_46]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_45]], %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:               %[[VAL_47:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_48:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_0]]]
+// CHECK:               memref.store %[[VAL_48]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_47]], %[[VAL_3]]{{\[}}%[[VAL_0]]]
+// CHECK:               %[[VAL_49:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_50:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_0]]]
+// CHECK:               memref.store %[[VAL_50]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_49]], %[[VAL_4]]{{\[}}%[[VAL_0]]]
+// CHECK:             }
+// CHECK:             %[[VAL_51:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:             %[[VAL_52:.*]] = arith.cmpi ult, %[[VAL_51]], %[[VAL_51]]
+// CHECK:             scf.if %[[VAL_52]] {
+// CHECK:               %[[VAL_53:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_53]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_53]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_54:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_54]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_54]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_55:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_55]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_55]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_56:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_57:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:               %[[VAL_58:.*]] = arith.cmpi ult, %[[VAL_56]], %[[VAL_57]]
+// CHECK:               scf.if %[[VAL_58]] {
+// CHECK:                 %[[VAL_59:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_60:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:                 memref.store %[[VAL_60]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_59]], %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:                 %[[VAL_61:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_62:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_0]]]
+// CHECK:                 memref.store %[[VAL_62]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_61]], %[[VAL_3]]{{\[}}%[[VAL_0]]]
+// CHECK:                 %[[VAL_63:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_64:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_0]]]
+// CHECK:                 memref.store %[[VAL_64]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_63]], %[[VAL_4]]{{\[}}%[[VAL_0]]]
+// CHECK:               }
+// CHECK:             }
+// CHECK:             %[[VAL_65:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_41]]]
+// CHECK:             %[[VAL_66:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:             %[[VAL_67:.*]] = arith.cmpi ult, %[[VAL_65]], %[[VAL_66]]
+// CHECK:             scf.if %[[VAL_67]] {
+// CHECK:               %[[VAL_68:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_41]]]
+// CHECK:               %[[VAL_69:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_69]], %[[VAL_2]]{{\[}}%[[VAL_41]]]
+// CHECK:               memref.store %[[VAL_68]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_70:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_41]]]
+// CHECK:               %[[VAL_71:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_71]], %[[VAL_3]]{{\[}}%[[VAL_41]]]
+// CHECK:               memref.store %[[VAL_70]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_72:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_41]]]
+// CHECK:               %[[VAL_73:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:               memref.store %[[VAL_73]], %[[VAL_4]]{{\[}}%[[VAL_41]]]
+// CHECK:               memref.store %[[VAL_72]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_74:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_75:.*]] = arith.cmpi ult, %[[VAL_74]], %[[VAL_74]]
+// CHECK:               scf.if %[[VAL_75]] {
+// CHECK:                 %[[VAL_76:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_76]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_76]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_77:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_77]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_77]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_78:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_78]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_78]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_79:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_80:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:                 %[[VAL_81:.*]] = arith.cmpi ult, %[[VAL_79]], %[[VAL_80]]
+// CHECK:                 scf.if %[[VAL_81]] {
+// CHECK:                   %[[VAL_82:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                   %[[VAL_83:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:                   memref.store %[[VAL_83]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                   memref.store %[[VAL_82]], %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:                   %[[VAL_84:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                   %[[VAL_85:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_0]]]
+// CHECK:                   memref.store %[[VAL_85]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                   memref.store %[[VAL_84]], %[[VAL_3]]{{\[}}%[[VAL_0]]]
+// CHECK:                   %[[VAL_86:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                   %[[VAL_87:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_0]]]
+// CHECK:                   memref.store %[[VAL_87]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                   memref.store %[[VAL_86]], %[[VAL_4]]{{\[}}%[[VAL_0]]]
+// CHECK:                 }
+// CHECK:               }
+// CHECK:             }
+// CHECK:             %[[VAL_88:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_10]]]
+// CHECK:             %[[VAL_89:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_41]]]
+// CHECK:             %[[VAL_90:.*]] = arith.cmpi ult, %[[VAL_88]], %[[VAL_89]]
+// CHECK:             scf.if %[[VAL_90]] {
+// CHECK:               %[[VAL_91:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_10]]]
+// CHECK:               %[[VAL_92:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_41]]]
+// CHECK:               memref.store %[[VAL_92]], %[[VAL_2]]{{\[}}%[[VAL_10]]]
+// CHECK:               memref.store %[[VAL_91]], %[[VAL_2]]{{\[}}%[[VAL_41]]]
+// CHECK:               %[[VAL_93:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_10]]]
+// CHECK:               %[[VAL_94:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_41]]]
+// CHECK:               memref.store %[[VAL_94]], %[[VAL_3]]{{\[}}%[[VAL_10]]]
+// CHECK:               memref.store %[[VAL_93]], %[[VAL_3]]{{\[}}%[[VAL_41]]]
+// CHECK:               %[[VAL_95:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_10]]]
+// CHECK:               %[[VAL_96:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_41]]]
+// CHECK:               memref.store %[[VAL_96]], %[[VAL_4]]{{\[}}%[[VAL_10]]]
+// CHECK:               memref.store %[[VAL_95]], %[[VAL_4]]{{\[}}%[[VAL_41]]]
+// CHECK:               %[[VAL_97:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_41]]]
+// CHECK:               %[[VAL_98:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:               %[[VAL_99:.*]] = arith.cmpi ult, %[[VAL_97]], %[[VAL_98]]
+// CHECK:               scf.if %[[VAL_99]] {
+// CHECK:                 %[[VAL_100:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_41]]]
+// CHECK:                 %[[VAL_101:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_101]], %[[VAL_2]]{{\[}}%[[VAL_41]]]
+// CHECK:                 memref.store %[[VAL_100]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_102:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_41]]]
+// CHECK:                 %[[VAL_103:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_103]], %[[VAL_3]]{{\[}}%[[VAL_41]]]
+// CHECK:                 memref.store %[[VAL_102]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_104:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_41]]]
+// CHECK:                 %[[VAL_105:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                 memref.store %[[VAL_105]], %[[VAL_4]]{{\[}}%[[VAL_41]]]
+// CHECK:                 memref.store %[[VAL_104]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_106:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                 %[[VAL_107:.*]] = arith.cmpi ult, %[[VAL_106]], %[[VAL_106]]
+// CHECK:                 scf.if %[[VAL_107]] {
+// CHECK:                   %[[VAL_108:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                   memref.store %[[VAL_108]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                   memref.store %[[VAL_108]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                   %[[VAL_109:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                   memref.store %[[VAL_109]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                   memref.store %[[VAL_109]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                   %[[VAL_110:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                   memref.store %[[VAL_110]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                   memref.store %[[VAL_110]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                   %[[VAL_111:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                   %[[VAL_112:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:                   %[[VAL_113:.*]] = arith.cmpi ult, %[[VAL_111]], %[[VAL_112]]
+// CHECK:                   scf.if %[[VAL_113]] {
+// CHECK:                     %[[VAL_114:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                     %[[VAL_115:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:                     memref.store %[[VAL_115]], %[[VAL_2]]{{\[}}%[[VAL_9]]]
+// CHECK:                     memref.store %[[VAL_114]], %[[VAL_2]]{{\[}}%[[VAL_0]]]
+// CHECK:                     %[[VAL_116:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                     %[[VAL_117:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_0]]]
+// CHECK:                     memref.store %[[VAL_117]], %[[VAL_3]]{{\[}}%[[VAL_9]]]
+// CHECK:                     memref.store %[[VAL_116]], %[[VAL_3]]{{\[}}%[[VAL_0]]]
+// CHECK:                     %[[VAL_118:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                     %[[VAL_119:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_0]]]
+// CHECK:                     memref.store %[[VAL_119]], %[[VAL_4]]{{\[}}%[[VAL_9]]]
+// CHECK:                     memref.store %[[VAL_118]], %[[VAL_4]]{{\[}}%[[VAL_0]]]
+// CHECK:                   }
+// CHECK:                 }
+// CHECK:               }
 // CHECK:             }
-// CHECK:             scf.yield %[[VAL_31:.*]]
 // CHECK:           }
-// CHECK:           %[[VAL_32:.*]] = arith.cmpi ne, %[[VAL_8]], %[[VAL_13:.*]]
-// CHECK:           scf.if %[[VAL_32]] {
-// CHECK:             %[[VAL_34:.*]] = memref.load %[[VAL_2]]{{\[}}
-// CHECK:             %[[VAL_35:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_8]]]
-// CHECK:             memref.store %[[VAL_35]], %[[VAL_2]]
-// CHECK:             memref.store %[[VAL_34]], %[[VAL_2]]{{\[}}%[[VAL_8]]]
-// CHECK:             %[[VAL_36:.*]] = memref.load %[[VAL_3]]
-// CHECK:             %[[VAL_37:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_8]]]
-// CHECK:             memref.store %[[VAL_37]], %[[VAL_3]]
-// CHECK:             memref.store %[[VAL_36]], %[[VAL_3]]{{\[}}%[[VAL_8]]]
-// CHECK:             %[[VAL_38:.*]] = memref.load %[[VAL_4]]
-// CHECK:             %[[VAL_39:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_8]]]
-// CHECK:             memref.store %[[VAL_39]], %[[VAL_4]]
-// CHECK:             memref.store %[[VAL_38]], %[[VAL_4]]{{\[}}%[[VAL_8]]]
-// CHECK:           }
-// CHECK:           %[[VAL_40:.*]]:3 = scf.while (%[[VAL_41:.*]] = %[[VAL_0]], %[[VAL_42:.*]] = %[[VAL_9]], %[[VAL_43:.*]] = %[[VAL_8]])
-// CHECK:             %[[VAL_44:.*]] = arith.cmpi ult, %[[VAL_41]], %[[VAL_42]]
-// CHECK:             scf.condition(%[[VAL_44]]) %[[VAL_41]], %[[VAL_42]], %[[VAL_43]]
+// CHECK:           %[[VAL_120:.*]]:3 = scf.while (%[[VAL_121:.*]] = %[[VAL_0]], %[[VAL_122:.*]] = %[[VAL_10]], %[[VAL_123:.*]] = %[[VAL_9]])
+// CHECK:             %[[VAL_124:.*]] = arith.cmpi ult, %[[VAL_121]], %[[VAL_122]]
+// CHECK:             scf.condition(%[[VAL_124]]) %[[VAL_121]], %[[VAL_122]], %[[VAL_123]]
 // CHECK:           } do {
-// CHECK:           ^bb0(%[[VAL_45:.*]]: index, %[[VAL_46:.*]]: index, %[[VAL_47:.*]]: index):
-// CHECK:             %[[VAL_48:.*]] = scf.while (%[[VAL_49:.*]] = %[[VAL_45]]) : (index) -> index {
-// CHECK:               %[[VAL_50:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_49]]]
-// CHECK:               %[[VAL_51:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_47]]]
-// CHECK:               %[[VAL_52:.*]] = arith.cmpi ult, %[[VAL_50]], %[[VAL_51]]
-// CHECK:               scf.condition(%[[VAL_52]]) %[[VAL_49]]
+// CHECK:           ^bb0(%[[VAL_125:.*]]: index, %[[VAL_126:.*]]: index, %[[VAL_127:.*]]: index)
+// CHECK:             %[[VAL_128:.*]] = scf.while (%[[VAL_129:.*]] = %[[VAL_125]])
+// CHECK:               %[[VAL_130:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_129]]]
+// CHECK:               %[[VAL_131:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_127]]]
+// CHECK:               %[[VAL_132:.*]] = arith.cmpi ult, %[[VAL_130]], %[[VAL_131]]
+// CHECK:               scf.condition(%[[VAL_132]]) %[[VAL_129]]
 // CHECK:             } do {
-// CHECK:             ^bb0(%[[VAL_53:.*]]: index):
-// CHECK:               %[[VAL_54:.*]] = arith.addi %[[VAL_53]], %[[VAL_5]]
-// CHECK:               scf.yield %[[VAL_54]]
-// CHECK:             }
-// CHECK:             %[[VAL_55:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_56:.*]]]
-// CHECK:             %[[VAL_57:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_47]]]
-// CHECK:             %[[VAL_58:.*]] = arith.cmpi eq, %[[VAL_55]], %[[VAL_57]]
-// CHECK:             %[[VAL_59:.*]] = scf.while (%[[VAL_60:.*]] = %[[VAL_46]]) : (index) -> index {
-// CHECK:               %[[VAL_61:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_47]]]
-// CHECK:               %[[VAL_62:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_60]]]
-// CHECK:               %[[VAL_63:.*]] = arith.cmpi ult, %[[VAL_61]], %[[VAL_62]]
-// CHECK:               scf.condition(%[[VAL_63]]) %[[VAL_60]]
+// CHECK:             ^bb0(%[[VAL_133:.*]]: index):
+// CHECK:               %[[VAL_134:.*]] = arith.addi %[[VAL_133]], %[[VAL_5]]
+// CHECK:               scf.yield %[[VAL_134]]
+// CHECK:             }
+// CHECK:             %[[VAL_135:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_136:.*]]]
+// CHECK:             %[[VAL_137:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_127]]]
+// CHECK:             %[[VAL_138:.*]] = arith.cmpi eq, %[[VAL_135]], %[[VAL_137]]
+// CHECK:             %[[VAL_139:.*]] = scf.while (%[[VAL_140:.*]] = %[[VAL_126]])
+// CHECK:               %[[VAL_141:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_127]]]
+// CHECK:               %[[VAL_142:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_140]]]
+// CHECK:               %[[VAL_143:.*]] = arith.cmpi ult, %[[VAL_141]], %[[VAL_142]]
+// CHECK:               scf.condition(%[[VAL_143]]) %[[VAL_140]]
 // CHECK:             } do {
-// CHECK:             ^bb0(%[[VAL_64:.*]]: index):
-// CHECK:               %[[VAL_65:.*]] = arith.addi %[[VAL_64]], %[[VAL_6]]
-// CHECK:               scf.yield %[[VAL_65]]
-// CHECK:             }
-// CHECK:             %[[VAL_66:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_67:.*]]]
-// CHECK:             %[[VAL_68:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_47]]]
-// CHECK:             %[[VAL_69:.*]] = arith.cmpi eq, %[[VAL_66]], %[[VAL_68]]
-// CHECK:             %[[VAL_70:.*]] = arith.cmpi ult, %[[VAL_56]], %[[VAL_67]]
-// CHECK:             %[[VAL_71:.*]]:3 = scf.if %[[VAL_70]] -> (index, index, index) {
-// CHECK:               %[[VAL_72:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_56]]]
-// CHECK:               %[[VAL_73:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_67]]]
-// CHECK:               memref.store %[[VAL_73]], %[[VAL_2]]{{\[}}%[[VAL_56]]]
-// CHECK:               memref.store %[[VAL_72]], %[[VAL_2]]{{\[}}%[[VAL_67]]]
-// CHECK:               %[[VAL_74:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_56]]]
-// CHECK:               %[[VAL_75:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_67]]]
-// CHECK:               memref.store %[[VAL_75]], %[[VAL_3]]{{\[}}%[[VAL_56]]]
-// CHECK:               memref.store %[[VAL_74]], %[[VAL_3]]{{\[}}%[[VAL_67]]]
-// CHECK:               %[[VAL_76:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_56]]]
-// CHECK:               %[[VAL_77:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_67]]]
-// CHECK:               memref.store %[[VAL_77]], %[[VAL_4]]{{\[}}%[[VAL_56]]]
-// CHECK:               memref.store %[[VAL_76]], %[[VAL_4]]{{\[}}%[[VAL_67]]]
-// CHECK:               %[[VAL_78:.*]] = arith.cmpi eq, %[[VAL_56]], %[[VAL_47]]
-// CHECK:               %[[VAL_79:.*]] = scf.if %[[VAL_78]] -> (index) {
-// CHECK:                 scf.yield %[[VAL_67]]
+// CHECK:             ^bb0(%[[VAL_144:.*]]: index):
+// CHECK:               %[[VAL_145:.*]] = arith.addi %[[VAL_144]], %[[VAL_7]]
+// CHECK:               scf.yield %[[VAL_145]]
+// CHECK:             }
+// CHECK:             %[[VAL_146:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_147:.*]]]
+// CHECK:             %[[VAL_148:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_127]]]
+// CHECK:             %[[VAL_149:.*]] = arith.cmpi eq, %[[VAL_146]], %[[VAL_148]]
+// CHECK:             %[[VAL_150:.*]] = arith.cmpi ult, %[[VAL_136]], %[[VAL_147]]
+// CHECK:             %[[VAL_151:.*]]:3 = scf.if %[[VAL_150]]
+// CHECK:               %[[VAL_152:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_136]]]
+// CHECK:               %[[VAL_153:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_147]]]
+// CHECK:               memref.store %[[VAL_153]], %[[VAL_2]]{{\[}}%[[VAL_136]]]
+// CHECK:               memref.store %[[VAL_152]], %[[VAL_2]]{{\[}}%[[VAL_147]]]
+// CHECK:               %[[VAL_154:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_136]]]
+// CHECK:               %[[VAL_155:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_147]]]
+// CHECK:               memref.store %[[VAL_155]], %[[VAL_3]]{{\[}}%[[VAL_136]]]
+// CHECK:               memref.store %[[VAL_154]], %[[VAL_3]]{{\[}}%[[VAL_147]]]
+// CHECK:               %[[VAL_156:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_136]]]
+// CHECK:               %[[VAL_157:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_147]]]
+// CHECK:               memref.store %[[VAL_157]], %[[VAL_4]]{{\[}}%[[VAL_136]]]
+// CHECK:               memref.store %[[VAL_156]], %[[VAL_4]]{{\[}}%[[VAL_147]]]
+// CHECK:               %[[VAL_158:.*]] = arith.cmpi eq, %[[VAL_136]], %[[VAL_127]]
+// CHECK:               %[[VAL_159:.*]] = scf.if %[[VAL_158]]
+// CHECK:                 scf.yield %[[VAL_147]]
 // CHECK:               } else {
-// CHECK:                 %[[VAL_80:.*]] = arith.cmpi eq, %[[VAL_67]], %[[VAL_47]]
-// CHECK:                 %[[VAL_81:.*]] = arith.select %[[VAL_80]], %[[VAL_56]], %[[VAL_47]]
-// CHECK:                 scf.yield %[[VAL_81]]
-// CHECK:               }
-// CHECK:               %[[VAL_82:.*]] = arith.andi %[[VAL_58]], %[[VAL_69]]
-// CHECK:               %[[VAL_83:.*]]:2 = scf.if %[[VAL_82]] -> (index, index) {
-// CHECK:                 %[[VAL_84:.*]] = arith.addi %[[VAL_56]], %[[VAL_5]]
-// CHECK:                 %[[VAL_85:.*]] = arith.subi %[[VAL_67]], %[[VAL_5]]
-// CHECK:                 scf.yield %[[VAL_84]], %[[VAL_85]]
+// CHECK:                 %[[VAL_160:.*]] = arith.cmpi eq, %[[VAL_147]], %[[VAL_127]]
+// CHECK:                 %[[VAL_161:.*]] = arith.select %[[VAL_160]], %[[VAL_136]], %[[VAL_127]]
+// CHECK:                 scf.yield %[[VAL_161]]
+// CHECK:               }
+// CHECK:               %[[VAL_162:.*]] = arith.andi %[[VAL_138]], %[[VAL_149]] : i1
+// CHECK:               %[[VAL_163:.*]]:2 = scf.if %[[VAL_162]]
+// CHECK:                 %[[VAL_164:.*]] = arith.addi %[[VAL_136]], %[[VAL_5]]
+// CHECK:                 %[[VAL_165:.*]] = arith.subi %[[VAL_147]], %[[VAL_5]]
+// CHECK:                 scf.yield %[[VAL_164]], %[[VAL_165]]
 // CHECK:               } else {
-// CHECK:                 scf.yield %[[VAL_56]], %[[VAL_67]]
+// CHECK:                 scf.yield %[[VAL_136]], %[[VAL_147]]
 // CHECK:               }
-// CHECK:               scf.yield %[[VAL_86:.*]]#0, %[[VAL_86]]#1, %[[VAL_87:.*]]
+// CHECK:               scf.yield %[[VAL_166:.*]]#0, %[[VAL_166]]#1, %[[VAL_167:.*]]
 // CHECK:             } else {
-// CHECK:               scf.yield %[[VAL_56]], %[[VAL_67]], %[[VAL_47]]
+// CHECK:               scf.yield %[[VAL_136]], %[[VAL_147]], %[[VAL_127]]
 // CHECK:             }
-// CHECK:             scf.yield %[[VAL_88:.*]]#0, %[[VAL_88]]#1, %[[VAL_88]]#2
+// CHECK:             scf.yield %[[VAL_168:.*]]#0, %[[VAL_168]]#1, %[[VAL_168]]#2
 // CHECK:           }
-// CHECK:           return %[[VAL_89:.*]]#2
+// CHECK:           return %[[VAL_169:.*]]#2
 // CHECK:         }
 
 // CHECK-LABEL:   func.func private @_sparse_qsort_1_i8_f32_index(
 // CHECK-SAME:                                                   %[[L:arg0]]: index,
 // CHECK-SAME:                                                   %[[H:.*]]: index,
 // CHECK-SAME:                                                   %[[X0:.*]]: memref<?xi8>,
 // CHECK-SAME:                                                   %[[Y0:.*]]: memref<?xf32>,
 // CHECK-SAME:                                                   %[[Y1:.*]]: memref<?xindex>) {

 //
 // CHECK-DAG:     func.func private @_sparse_shift_down_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>, %arg5: index) {
 // CHECK-DAG:     func.func private @_sparse_heap_sort_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) {
 // CHECK-LABEL:   func.func @sparse_sort_coo_heap
 func.func @sparse_sort_coo_heap(%arg0: index, %arg1: memref<100xindex>, %arg2: memref<?xf32>, %arg3: memref<10xi32>) -> (memref<100xindex>, memref<?xf32>, memref<10xi32>) {
   sparse_tensor.sort_coo heap_sort %arg0, %arg1 jointly %arg2, %arg3 {nx = 2 : index, ny = 1: index} : memref<100xindex> jointly memref<?xf32>, memref<10xi32>
   return %arg1, %arg2, %arg3 : memref<100xindex>, memref<?xf32>, memref<10xi32>
 }
 No newline at end of file