[llvm] 290ddb9 - [VPlan][PseudoProbe] Fix `pseudoprobe` duplication when `VF=1` (#185238)

[via llvm-commits]

Author: Henry Jiang
Date: 2026-04-08T10:04:57-07:00
New Revision: 290ddb9238ec569f8518614ddbc48f8f0445777f

URL: https://github.com/llvm/llvm-project/commit/290ddb9238ec569f8518614ddbc48f8f0445777f
DIFF: https://github.com/llvm/llvm-project/commit/290ddb9238ec569f8518614ddbc48f8f0445777f.diff

LOG: [VPlan][PseudoProbe] Fix `pseudoprobe` duplication when `VF=1` (#185238)

Fix assertion in `loop-vectorize` on loops that contains
`llvm.pseudoprobe` at VF=1, UF=2. Minimal Reproducer:
https://godbolt.org/z/nrcMWWqMx

Originally in https://reviews.llvm.org/D144066, Pseudoprobes were marked
non-uniform in `isUniformAfterVectorization` even for VF=1 that allows
the `REPLICATE call @llvm.pseudoprobe` to survive until the plan is
executed when VF=1, UF=2, causing the crash.

Instead, `isUniformAfterVectorization` as true even for pseudoprobe when
`VF.isScalar()`.

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
    llvm/test/Transforms/LoopVectorize/pseudoprobe.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index e907b5621b817..fe9b676e36c9d 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -980,15 +980,16 @@ class LoopVectorizationCostModel {
     assert(
         TheLoop->isInnermost() &&
         "cost-model should not be used for outer loops (in VPlan-native path)");
-    // Pseudo probe needs to be duplicated for each unrolled iteration and
-    // vector lane so that profiled loop trip count can be accurately
-    // accumulated instead of being under counted.
-    if (isa<PseudoProbeInst>(I))
-      return false;
 
+    // If VF is scalar, then all instructions are trivially uniform.
     if (VF.isScalar())
       return true;
 
+    // Pseudo probes must be duplicated per vector lane so that the
+    // profiled loop trip count is not undercounted.
+    if (isa<PseudoProbeInst>(I))
+      return false;
+
     auto UniformsPerVF = Uniforms.find(VF);
     assert(UniformsPerVF != Uniforms.end() &&
            "VF not yet analyzed for uniformity");

diff  --git a/llvm/test/Transforms/LoopVectorize/pseudoprobe.ll b/llvm/test/Transforms/LoopVectorize/pseudoprobe.ll
index c07d7c84e1bf1..ac59f1c35965a 100644
--- a/llvm/test/Transforms/LoopVectorize/pseudoprobe.ll
+++ b/llvm/test/Transforms/LoopVectorize/pseudoprobe.ll
@@ -1,8 +1,48 @@
-; RUN: opt < %s -passes=loop-vectorize -force-vector-interleave=1 -force-vector-width=4 -S | FileCheck %s
+; RUN: opt < %s -passes=loop-vectorize -force-vector-interleave=1 -force-vector-width=1 -S | FileCheck %s --check-prefix=VF1UF1
+; RUN: opt < %s -passes=loop-vectorize -force-vector-interleave=2 -force-vector-width=1 -S | FileCheck %s --check-prefix=VF1UF2
+; RUN: opt < %s -passes=loop-vectorize -force-vector-interleave=1 -force-vector-width=2 -S | FileCheck %s --check-prefix=VF2UF1
+; RUN: opt < %s -passes=loop-vectorize -force-vector-interleave=1 -force-vector-width=4 -S | FileCheck %s --check-prefix=VF4UF1
+; RUN: opt < %s -passes=loop-vectorize -force-vector-interleave=2 -force-vector-width=4 -S | FileCheck %s --check-prefix=VF4UF2
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 
 define i32 @test1(ptr nocapture %a, ptr nocapture readonly %b) {
+; VF1UF1-LABEL:  @test1
+; VF1UF1:        for.body:
+; VF1UF1-COUNT-1:  call void @llvm.pseudoprobe
+; VF1UF1-NOT:      call void @llvm.pseudoprobe
+; VF1UF1:          br i1
+;
+; VF1UF2-LABEL:  @test1
+; VF1UF2:        vector.body:
+; VF1UF2-COUNT-2:  call void @llvm.pseudoprobe
+; VF1UF2-NOT:      call void @llvm.pseudoprobe
+; VF1UF2:          %index.next = add nuw i64 %index, 2
+;
+; VF2UF1-LABEL:  @test1
+; VF2UF1:        vector.body:
+; VF2UF1:          load <2 x float>, ptr %{{.*}}
+; VF2UF1:          store <2 x i32> %{{.*}}, ptr %{{.*}}
+; VF2UF1-COUNT-2:  call void @llvm.pseudoprobe
+; VF2UF1-NOT:      call void @llvm.pseudoprobe
+; VF2UF1:          %index.next = add nuw i64 %index, 2
+;
+; VF4UF1-LABEL:  @test1
+; VF4UF1:        vector.body:
+; VF4UF1:          load <4 x float>, ptr %{{.*}}
+; VF4UF1:          store <4 x i32> %{{.*}}, ptr %{{.*}}
+; VF4UF1-COUNT-4:  call void @llvm.pseudoprobe
+; VF4UF1-NOT:      call void @llvm.pseudoprobe
+; VF4UF1:          %index.next = add nuw i64 %index, 4
+;
+; VF4UF2-LABEL:  @test1
+; VF4UF2:        vector.body:
+; VF4UF2:          load <4 x float>, ptr %{{.*}}
+; VF4UF2:          store <4 x i32> %{{.*}}, ptr %{{.*}}
+; VF4UF2-COUNT-8:  call void @llvm.pseudoprobe
+; VF4UF2-NOT:      call void @llvm.pseudoprobe
+; VF4UF2:          %index.next = add nuw i64 %index, 8
+;
 entry:
   call void @llvm.pseudoprobe(i64 3666282617048535130, i64 1, i32 0, i64 -1)
   br label %for.body
@@ -24,17 +64,6 @@ for.end:                                          ; preds = %for.body
   ret i32 0
 }
 
-
-; CHECK-LABEL:  @test1
-; CHECK:        vector.body:
-; CHECK:          load <4 x float>, ptr %{{.*}}
-; CHECK:          store <4 x i32> %{{.*}}, ptr %{{.*}}
-; CHECK-COUNT-4:  call void @llvm.pseudoprobe(i64 3666282617048535130, i64 2, i32 0, i64 -1)
-; CHECK:          %index.next = add nuw i64 %index, 4
-
-
-
-
 declare void @llvm.pseudoprobe(i64, i64, i32, i64)
 
 !llvm.pseudo_probe_desc = !{!0}