[llvm] 9de327c - [LV] Generalize predication checks from 2c8836c899 for operands.

[Florian Hahn via llvm-commits]

Author: Florian Hahn
Date: 2024-10-02T20:16:41+01:00
New Revision: 9de327c94d0c995803b6485fb28ac4cad1e53bbe

URL: https://github.com/llvm/llvm-project/commit/9de327c94d0c995803b6485fb28ac4cad1e53bbe
DIFF: https://github.com/llvm/llvm-project/commit/9de327c94d0c995803b6485fb28ac4cad1e53bbe.diff

LOG: [LV] Generalize predication checks from 2c8836c899 for operands.

This fixes another case where the VPlan-based and legacy cost models
disagree. If any of the operands is predicated, it can't be trivially
hoisted and we should consider the cost for evaluating it each loop
iteration.

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

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
    llvm/test/Transforms/LoopVectorize/X86/predicated-instruction-cost.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index de3b981a4fe390..6e082b1c134dee 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -6551,14 +6551,17 @@ LoopVectorizationCostModel::getInstructionCost(Instruction *I,
       Op2 = cast<SCEVConstant>(PSE.getSCEV(Op2))->getValue();
     }
     auto Op2Info = TTI.getOperandInfo(Op2);
-    auto IsInvariant = [this](Value *Op) {
+    std::function<bool(Value *)> IsInvariant =
+        [this, &IsInvariant](Value *Op) -> bool {
       if (!Legal->isInvariant(Op))
         return false;
-      // Consider Op2 invariant, if it is not a predicated instruction in the
-      // loop. In that case, it is not trivially hoistable.
+      // Consider Op2invariant, if it or its operands aren't predicated
+      // instruction in the loop. In that case, it is not trivially hoistable.
       return !isa<Instruction>(Op) ||
              !TheLoop->contains(cast<Instruction>(Op)) ||
-             !isPredicatedInst(cast<Instruction>(Op));
+             (!isPredicatedInst(cast<Instruction>(Op)) &&
+              all_of(cast<Instruction>(Op)->operands(),
+                     [&IsInvariant](Value *Op) { return IsInvariant(Op); }));
     };
     if (Op2Info.Kind == TargetTransformInfo::OK_AnyValue && IsInvariant(Op2))
       Op2Info.Kind = TargetTransformInfo::OK_UniformValue;

diff  --git a/llvm/test/Transforms/LoopVectorize/X86/predicated-instruction-cost.ll b/llvm/test/Transforms/LoopVectorize/X86/predicated-instruction-cost.ll
index 0072dd95bd0983..7ae20a07290c35 100644
--- a/llvm/test/Transforms/LoopVectorize/X86/predicated-instruction-cost.ll
+++ b/llvm/test/Transforms/LoopVectorize/X86/predicated-instruction-cost.ll
@@ -52,3 +52,52 @@ loop.latch:
 exit:
   ret void
 }
+
+; Test case for https://github.com/llvm/llvm-project/issues/108697.
+define void @test_wide_shift_uses_predicated_invariant_instruction(i32 %d, i1 %c, ptr %dst) {
+; CHECK-LABEL: define void @test_wide_shift_uses_predicated_invariant_instruction(
+; CHECK-SAME: i32 [[D:%.*]], i1 [[C:%.*]], ptr [[DST:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br label %[[LOOP_HEADER:.*]]
+; CHECK:       [[LOOP_HEADER]]:
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, %[[ENTRY]] ], [ [[IV_NEXT:%.*]], %[[LOOP_LATCH:.*]] ]
+; CHECK-NEXT:    br i1 [[C]], label %[[LOOP_LATCH]], label %[[ELSE:.*]]
+; CHECK:       [[ELSE]]:
+; CHECK-NEXT:    [[REM:%.*]] = urem i32 100, [[D]]
+; CHECK-NEXT:    [[SEXT:%.*]] = shl i32 [[REM]], 12
+; CHECK-NEXT:    [[SHL_I:%.*]] = shl i32 999, [[SEXT]]
+; CHECK-NEXT:    br label %[[LOOP_LATCH]]
+; CHECK:       [[LOOP_LATCH]]:
+; CHECK-NEXT:    [[P:%.*]] = phi i32 [ [[SHL_I]], %[[ELSE]] ], [ 0, %[[LOOP_HEADER]] ]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, ptr [[DST]], i32 [[IV]]
+; CHECK-NEXT:    store i32 [[P]], ptr [[GEP]], align 4
+; CHECK-NEXT:    [[IV_NEXT]] = add i32 [[IV]], 1
+; CHECK-NEXT:    [[EC:%.*]] = icmp eq i32 [[IV_NEXT]], 100
+; CHECK-NEXT:    br i1 [[EC]], label %[[EXIT:.*]], label %[[LOOP_HEADER]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop.header
+
+loop.header:
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop.latch ]
+  br i1 %c, label %loop.latch, label %else
+
+else:
+  %rem = urem i32 100, %d
+  %sext = shl i32 %rem, 12
+  %shl.i = shl i32 999, %sext
+  br label %loop.latch
+
+loop.latch:
+  %p = phi i32 [ %shl.i, %else ], [ 0, %loop.header ]
+  %gep = getelementptr inbounds i32, ptr %dst, i32 %iv
+  store i32 %p, ptr %gep
+  %iv.next = add i32 %iv, 1
+  %ec = icmp eq i32 %iv.next, 100
+  br i1 %ec, label %exit, label %loop.header
+
+exit:
+  ret void
+}