[Mlir-commits] [mlir] [mlir][scf] add unroll-full option to test-loop-unrolling pass. (PR #127158)

[Oleksandr Alex Zinenko]

================
@@ -56,3 +57,59 @@ func.func @scf_loop_unroll_factor_1_promote() -> () {
   // UNROLL-BY-1-NEXT: %[[C0:.*]] = arith.constant 0 : index
   // UNROLL-BY-1-NEXT: %{{.*}} = "test.foo"(%[[C0]]) : (index) -> i32
 }
+
+// UNROLL-FULL-LABEL: func @scf_loop_unroll_full_single(
+// UNROLL-FULL-SAME:    %[[VAL_0:.*]]: index) -> index {
+func.func @scf_loop_unroll_full_single(%arg : index) -> index {
+  %0 = arith.constant 0 : index
+  %1 = arith.constant 1 : index
+  %2 = arith.constant 4 : index
+  %4 = scf.for %iv = %0 to %2 step %1 iter_args(%arg1 = %1) -> index {
+    %3 = arith.addi %arg1, %arg : index
+    scf.yield %3 : index
+  }
+  return %4 : index
+  // UNROLL-FULL: %[[VAL_1:.*]] = arith.constant 1 : index
+  // UNROLL-FULL: %[[VAL_2:.*]] = arith.addi %[[VAL_0]], %[[VAL_1]] : index
+  // UNROLL-FULL: %[[VAL_3:.*]] = arith.addi %[[VAL_2]], %[[VAL_0]] : index
+  // UNROLL-FULL: %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_0]] : index
+  // UNROLL-FULL: %[[VAL_5:.*]] = arith.addi %[[VAL_4]], %[[VAL_0]] : index
+  // UNROLL-FULL: return %[[VAL_5]] : index
+}
+
+// UNROLL-FULL-LABEL: func @scf_loop_unroll_full_outter_loops(
+// UNROLL-FULL-SAME:    %[[VAL_0:.*]]: vector<4x4xindex>) -> index {
+func.func @scf_loop_unroll_full_outter_loops(%arg0: vector<4x4xindex>) -> index {
+  %0 = arith.constant 0 : index
+  %1 = arith.constant 1 : index
+  %2 = arith.constant 4 : index
+  %6 = scf.for %arg1 = %0 to %2 step %1 iter_args(%it0 = %0) -> index {
+    %5 = scf.for %arg2 = %0 to %2 step %1 iter_args(%it1 = %it0) -> index {
+      %3 = vector.extract %arg0[%arg1, %arg2] : index from vector<4x4xindex>
+      %4 = arith.addi %3, %it1 : index
+      scf.yield %3 : index
+    }
+    scf.yield %5 : index
+  }
+  return %6 : index
+  // UNROLL-FULL:   %[[VAL_1:.*]] = arith.constant 0 : index
+  // UNROLL-FULL:   %[[VAL_2:.*]] = arith.constant 1 : index
+  // UNROLL-FULL:   %[[VAL_3:.*]] = arith.constant 4 : index
+  // UNROLL-FULL:   %[[VAL_4:.*]] = scf.for %[[VAL_5:.*]] = %[[VAL_1]] to %[[VAL_3]] step %[[VAL_2]] iter_args(%[[VAL_6:.*]] = %[[VAL_1]]) -> (index) {
+  // UNROLL-FULL:   %[[VAL_7:.*]] = vector.extract %[[VAL_0]][0, %[[VAL_5]]] : index from vector<4x4xindex>
+  // UNROLL-FULL:   scf.yield %[[VAL_7]] : index
----------------
ftynse wrote:

```suggestion
  // UNROLL-FULL:     %[[VAL_7:.*]] = vector.extract %[[VAL_0]][0, %[[VAL_5]]] : index from vector<4x4xindex>
  // UNROLL-FULL:     scf.yield %[[VAL_7]] : index
```

indent IR in tests for readability. Note that tests are often read by humans, make them as easy to read as you would make code.

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