[llvm-branch-commits] [llvm] [LoopInterchange] Reject if outer loop reduction has extra user (PR #203241)

[via llvm-branch-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Ryotaro Kasuga (kasuga-fj)

<details>
<summary>Changes</summary>

As demonstrated by the case reported in #<!-- -->203104, there are cases where interchanging the loops is unsafe because the outer loop reduction is used within the inner loop by something other than the PHI node that receives the value. Such cases should be rejected during the legality check phase.
This patch adds a check to ensure that the outer loop reduction is only used by the inner loop reduction PHI node. This check might be too conservative and could cause missed optimizations, but at the moment I can't find such a case.

Fixes #<!-- -->203104.

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/LoopInterchange.cpp (+20) 
- (modified) llvm/test/Transforms/LoopInterchange/reduction-extra-use-in-inner-loop.ll (+13-23) 


``````````diff
diff --git a/llvm/lib/Transforms/Scalar/LoopInterchange.cpp b/llvm/lib/Transforms/Scalar/LoopInterchange.cpp
index 2a35575eb2d4c..42d01c3024c3a 100644
--- a/llvm/lib/Transforms/Scalar/LoopInterchange.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopInterchange.cpp
@@ -1280,6 +1280,26 @@ bool LoopInterchangeLegality::checkInductionsAndReductions(Loop *OuterLoop) {
           });
           return false;
         }
+
+        // Reject if PHI has users other than InnerRedPhi. The typical case is
+        // as follows:
+        //
+        //   o.header:
+        //     %red.o = phi [ 0, ... ], [ %red.next, %o.latch ]
+        //     br label %i.header
+        //
+        //   i.header:
+        //     %red.i = phi [ %red.o, %o.header ], [ %red.next, %i.latch ]
+        //     br label %i.body
+        //
+        //   i.body:
+        //     store %red.o to %mem
+        //     ...
+        //
+        if (!all_of(PHI.users(),
+                    [InnerRedPhi](User *U) { return U == InnerRedPhi; }))
+          return false;
+
         OuterInnerReductions.insert(&PHI);
         OuterInnerReductions.insert(InnerRedPhi);
       } else {
diff --git a/llvm/test/Transforms/LoopInterchange/reduction-extra-use-in-inner-loop.ll b/llvm/test/Transforms/LoopInterchange/reduction-extra-use-in-inner-loop.ll
index 12493e954444f..82ba8e234c5da 100644
--- a/llvm/test/Transforms/LoopInterchange/reduction-extra-use-in-inner-loop.ll
+++ b/llvm/test/Transforms/LoopInterchange/reduction-extra-use-in-inner-loop.ll
@@ -26,40 +26,30 @@
 define i8 @extra_reduction_use_in_inner0(ptr noalias %A, ptr noalias %B) {
 ; CHECK-LABEL: define i8 @extra_reduction_use_in_inner0(
 ; CHECK-SAME: ptr noalias [[A:%.*]], ptr noalias [[B:%.*]]) {
-; CHECK-NEXT:  [[INNER_PREHEADER:.*:]]
+; CHECK-NEXT:  [[INNER_PREHEADER:.*]]:
 ; CHECK-NEXT:    br label %[[INNER:.*]]
-; CHECK:       [[OUTER_HEADER_PREHEADER1:.*]]:
+; CHECK:       [[INNER]]:
+; CHECK-NEXT:    [[I:%.*]] = phi i64 [ 0, %[[INNER_PREHEADER]] ], [ [[I_NEXT:%.*]], %[[OUTER_LATCH:.*]] ]
+; CHECK-NEXT:    [[SUM_OUTER:%.*]] = phi i8 [ 0, %[[INNER_PREHEADER]] ], [ [[SUM_LCSSA:%.*]], %[[OUTER_LATCH]] ]
 ; CHECK-NEXT:    br label %[[OUTER_HEADER_PREHEADER:.*]]
 ; CHECK:       [[OUTER_HEADER_PREHEADER]]:
-; CHECK-NEXT:    [[I:%.*]] = phi i64 [ [[I_NEXT:%.*]], %[[OUTER_LATCH:.*]] ], [ 0, %[[OUTER_HEADER_PREHEADER1]] ]
-; CHECK-NEXT:    [[SUM_INNER:%.*]] = phi i8 [ [[SUM_NEXT:%.*]], %[[OUTER_LATCH]] ], [ [[SUM:%.*]], %[[OUTER_HEADER_PREHEADER1]] ]
-; CHECK-NEXT:    br label %[[INNER_SPLIT1:.*]]
-; CHECK:       [[INNER]]:
-; CHECK-NEXT:    br label %[[INNER1:.*]]
-; CHECK:       [[INNER1]]:
-; CHECK-NEXT:    [[J:%.*]] = phi i64 [ [[TMP2:%.*]], %[[INNER_SPLIT:.*]] ], [ 0, %[[INNER]] ]
-; CHECK-NEXT:    [[SUM]] = phi i8 [ [[SUM_LCSSA:%.*]], %[[INNER_SPLIT]] ], [ 0, %[[INNER]] ]
-; CHECK-NEXT:    br label %[[OUTER_HEADER_PREHEADER1]]
-; CHECK:       [[INNER_SPLIT1]]:
+; CHECK-NEXT:    [[J:%.*]] = phi i64 [ 0, %[[INNER]] ], [ [[TMP0:%.*]], %[[OUTER_HEADER_PREHEADER]] ]
+; CHECK-NEXT:    [[SUM_INNER:%.*]] = phi i8 [ [[SUM_OUTER]], %[[INNER]] ], [ [[SUM_INNER_NEXT:%.*]], %[[OUTER_HEADER_PREHEADER]] ]
 ; CHECK-NEXT:    [[GEP_A:%.*]] = getelementptr [2 x i8], ptr [[A]], i64 [[J]], i64 [[I]]
 ; CHECK-NEXT:    [[A:%.*]] = load i8, ptr [[GEP_A]], align 1
-; CHECK-NEXT:    [[SUM_NEXT]] = add i8 [[SUM_INNER]], [[A]]
+; CHECK-NEXT:    [[SUM_INNER_NEXT]] = add i8 [[SUM_INNER]], [[A]]
 ; CHECK-NEXT:    [[GEP_B:%.*]] = getelementptr [2 x i8], ptr [[B]], i64 [[J]], i64 [[I]]
-; CHECK-NEXT:    store i8 [[SUM]], ptr [[GEP_B]], align 1
-; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[J]], 1
+; CHECK-NEXT:    store i8 [[SUM_OUTER]], ptr [[GEP_B]], align 1
+; CHECK-NEXT:    [[TMP0]] = add i64 [[J]], 1
 ; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i64 [[TMP0]], 2
-; CHECK-NEXT:    br label %[[OUTER_LATCH]]
-; CHECK:       [[INNER_SPLIT]]:
-; CHECK-NEXT:    [[SUM_LCSSA]] = phi i8 [ [[SUM_NEXT]], %[[OUTER_LATCH]] ]
-; CHECK-NEXT:    [[TMP2]] = add i64 [[J]], 1
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i64 [[TMP2]], 2
-; CHECK-NEXT:    br i1 [[TMP3]], label %[[EXIT:.*]], label %[[INNER1]]
+; CHECK-NEXT:    br i1 [[TMP1]], label %[[OUTER_LATCH]], label %[[OUTER_HEADER_PREHEADER]]
 ; CHECK:       [[OUTER_LATCH]]:
+; CHECK-NEXT:    [[SUM_LCSSA]] = phi i8 [ [[SUM_INNER_NEXT]], %[[OUTER_HEADER_PREHEADER]] ]
 ; CHECK-NEXT:    [[I_NEXT]] = add i64 [[I]], 1
 ; CHECK-NEXT:    [[EC_I:%.*]] = icmp eq i64 [[I_NEXT]], 2
-; CHECK-NEXT:    br i1 [[EC_I]], label %[[INNER_SPLIT]], label %[[OUTER_HEADER_PREHEADER]]
+; CHECK-NEXT:    br i1 [[EC_I]], label %[[EXIT:.*]], label %[[INNER]]
 ; CHECK:       [[EXIT]]:
-; CHECK-NEXT:    [[RES:%.*]] = phi i8 [ [[SUM_LCSSA]], %[[INNER_SPLIT]] ]
+; CHECK-NEXT:    [[RES:%.*]] = phi i8 [ [[SUM_LCSSA]], %[[OUTER_LATCH]] ]
 ; CHECK-NEXT:    ret i8 [[RES]]
 ;
 entry:

``````````

</details>


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