[VectorCombine] Avoid inserting freeze when scalarizing extend-extract if all extracts would lead to UB on poison. #164683

juliannagele · 2025-10-22T18:56:08Z

This change aims to avoid inserting a freeze instruction between the load and bitcast when scalarizing extend-extract. This is particularly useful in combination with #164682, which can then potentially further scalarize, provided there is no freeze.

alive2 proof: https://alive2.llvm.org/ce/z/W-GD88

…freeze.

… if all extracts would lead to UB on poison.

llvmbot · 2025-10-22T18:56:56Z

@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-llvm-transforms

Author: Julian Nagele (juliannagele)

Changes

This change aims to avoid inserting a freeze instruction between the load and bitcast when scalarizing extend-extract. This is particularly useful in combination with #164682, which can then potentially further scalarize, provided there is no freeze.

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

2 Files Affected:

(modified) llvm/lib/Transforms/Vectorize/VectorCombine.cpp (+18-2)
(modified) llvm/test/Transforms/VectorCombine/AArch64/ext-extract.ll (+34)

diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index d6eb00da11dc8..4a0c98e03203d 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -2017,8 +2017,24 @@ bool VectorCombine::scalarizeExtExtract(Instruction &I) {
 
   Value *ScalarV = Ext->getOperand(0);
   if (!isGuaranteedNotToBePoison(ScalarV, &AC, dyn_cast<Instruction>(ScalarV),
-                                 &DT))
-    ScalarV = Builder.CreateFreeze(ScalarV);
+                                 &DT)) {
+    // Check if all lanes are extracted and all extracts trigger UB on poison.
+    // If so, we do not need to insert a freeze.
+    SmallDenseSet<uint64_t, 8> ExtractedLanes;
+    bool AllExtractsHaveUB = true;
+    for (User *U : Ext->users()) {
+      auto *Extract = cast<ExtractElementInst>(U);
+      uint64_t Idx =
+          cast<ConstantInt>(Extract->getIndexOperand())->getZExtValue();
+      ExtractedLanes.insert(Idx);
+      if (!programUndefinedIfPoison(Extract)) {
+        AllExtractsHaveUB = false;
+        break;
+      }
+    }
+    if (!AllExtractsHaveUB || ExtractedLanes.size() != SrcTy->getNumElements())
+      ScalarV = Builder.CreateFreeze(ScalarV);
+  }
   ScalarV = Builder.CreateBitCast(
       ScalarV,
       IntegerType::get(SrcTy->getContext(), DL->getTypeSizeInBits(SrcTy)));
diff --git a/llvm/test/Transforms/VectorCombine/AArch64/ext-extract.ll b/llvm/test/Transforms/VectorCombine/AArch64/ext-extract.ll
index 60700412686ea..29d4ddd3d0ac8 100644
--- a/llvm/test/Transforms/VectorCombine/AArch64/ext-extract.ll
+++ b/llvm/test/Transforms/VectorCombine/AArch64/ext-extract.ll
@@ -346,3 +346,37 @@ entry:
   call void @use.i32(i32 %ext.3)
   ret void
 }
+
+define noundef i32 @zext_v4i8_all_lanes_used_no_freeze(<4 x i8> %src) {
+; CHECK-LABEL: define noundef i32 @zext_v4i8_all_lanes_used_no_freeze(
+; CHECK-SAME: <4 x i8> [[SRC:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast <4 x i8> [[SRC]] to i32
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[TMP0]], 24
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr i32 [[TMP0]], 16
+; CHECK-NEXT:    [[TMP3:%.*]] = and i32 [[TMP2]], 255
+; CHECK-NEXT:    [[TMP4:%.*]] = lshr i32 [[TMP0]], 8
+; CHECK-NEXT:    [[TMP5:%.*]] = and i32 [[TMP4]], 255
+; CHECK-NEXT:    [[TMP6:%.*]] = and i32 [[TMP0]], 255
+; CHECK-NEXT:    [[EXT:%.*]] = zext nneg <4 x i8> [[SRC]] to <4 x i32>
+; CHECK-NEXT:    [[EXT_0:%.*]] = extractelement <4 x i32> [[EXT]], i64 0
+; CHECK-NEXT:    [[EXT_1:%.*]] = extractelement <4 x i32> [[EXT]], i64 1
+; CHECK-NEXT:    [[EXT_2:%.*]] = extractelement <4 x i32> [[EXT]], i64 2
+; CHECK-NEXT:    [[EXT_3:%.*]] = extractelement <4 x i32> [[EXT]], i64 3
+; CHECK-NEXT:    [[ADD1:%.*]] = add i32 [[TMP6]], [[TMP5]]
+; CHECK-NEXT:    [[ADD2:%.*]] = add i32 [[ADD1]], [[TMP3]]
+; CHECK-NEXT:    [[ADD3:%.*]] = add i32 [[ADD2]], [[TMP1]]
+; CHECK-NEXT:    ret i32 [[ADD3]]
+;
+entry:
+  %ext = zext nneg <4 x i8> %src to <4 x i32>
+  %ext.0 = extractelement <4 x i32> %ext, i64 0
+  %ext.1 = extractelement <4 x i32> %ext, i64 1
+  %ext.2 = extractelement <4 x i32> %ext, i64 2
+  %ext.3 = extractelement <4 x i32> %ext, i64 3
+
+  %add1 = add i32 %ext.0, %ext.1
+  %add2 = add i32 %add1, %ext.2
+  %add3 = add i32 %add2, %ext.3
+  ret i32 %add3
+}

fhahn

Could you add an alive2 proof to the description?

fhahn · 2025-10-23T08:43:31Z

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

+    // Check if all lanes are extracted and all extracts trigger UB on poison.
+    // If so, we do not need to insert a freeze.
+    SmallDenseSet<uint64_t, 8> ExtractedLanes;
+    bool AllExtractsHaveUB = true;
+    for (User *U : Ext->users()) {
+      auto *Extract = cast<ExtractElementInst>(U);
+      uint64_t Idx =
+          cast<ConstantInt>(Extract->getIndexOperand())->getZExtValue();
+      ExtractedLanes.insert(Idx);
+      if (!programUndefinedIfPoison(Extract)) {
+        AllExtractsHaveUB = false;
+        break;
+      }
+    }


It might be worth to see if this logic can be moved to ValueTracking?

Would probably make sense to first add an initial implementation here, and then check if there are other cases where this would help via ValueTracking as follow-up

Sounds good, I've addressed the comments in VectorCombine for now. I think I have a proposal on how to move the logic to ValueTracking, will post it as a follow-up.

nikic · 2025-10-23T08:50:19Z

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

+      uint64_t Idx =
+          cast<ConstantInt>(Extract->getIndexOperand())->getZExtValue();
+      ExtractedLanes.insert(Idx);
+      if (!programUndefinedIfPoison(Extract)) {


This is under the assumption that the extract will be executed. The extract might be conditionally executed in a different block. I think the rest of the code does not exclude that possibility?

At least initially, it should probably be enough to restrict this to cases where the extracts are in the same block as the extend, and they are guaranteed to execute, if the extend executes

Yes, thanks! I added checks that the extracts are in the same block as the extend and that they are executed, and tests showing the freeze will still be inserted if those checks fail.

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

…extract if all extracts would lead to UB on poison.

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

…extend-extract if all extracts would lead to UB on poison.

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

…rizing extend-extract if all extracts would lead to UB on poison.

dtcxzyw

LGTM. Thanks!

fhahn

LGTM, thanks

fhahn · 2025-11-02T16:02:02Z

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

+    // are guaranteed to execute if Ext executes.
+    // If so, we do not need to insert a freeze.


Suggested change

// are guaranteed to execute if Ext executes.

// If so, we do not need to insert a freeze.

// are guaranteed to execute if Ext executes. If so, we do not need to insert a freeze.

…n scalarizing extend-extract if all extracts would lead to UB on poison.

llvm-ci · 2025-11-04T15:05:59Z

LLVM Buildbot has detected a new failure on builder ppc64le-mlir-rhel-clang running on ppc64le-mlir-rhel-test while building llvm at step 6 "test-build-check-mlir-build-only-check-mlir".

Full details are available at: https://lab.llvm.org/buildbot/#/builders/129/builds/32398

Here is the relevant piece of the build log for the reference

Step 6 (test-build-check-mlir-build-only-check-mlir) failure: 1200 seconds without output running [b'ninja', b'check-mlir'], attempting to kill
...
PASS: MLIR :: Pass/pipeline-parsing.mlir (3622 of 3633)
PASS: MLIR-Unit :: Interfaces/./MLIRInterfacesTests/12/22 (3623 of 3633)
PASS: MLIR :: mlir-tblgen/op-interface.td (3624 of 3633)
PASS: MLIR :: mlir-translate/import-diagnostics.ll (3625 of 3633)
PASS: MLIR :: mlir-tblgen/llvm-intrinsics.td (3626 of 3633)
PASS: MLIR :: mlir-reduce/dce-test.mlir (3627 of 3633)
PASS: MLIR :: mlir-runner/simple.mlir (3628 of 3633)
PASS: MLIR :: Pass/pipeline-options-parsing.mlir (3629 of 3633)
PASS: MLIR :: mlir-tblgen/rewriter-errors.td (3630 of 3633)
PASS: MLIR :: mlir-tblgen/op-error.td (3631 of 3633)
command timed out: 1200 seconds without output running [b'ninja', b'check-mlir'], attempting to kill
process killed by signal 9
program finished with exit code -1
elapsedTime=1987.543915

juliannagele added 2 commits October 22, 2025 15:25

[VectorCombine] Add test showing exisiting behaviour adding unneeded …

553cfa8

…freeze.

[VecorCombine] Avoid inserting freeze when scalarizing extend-extract…

b8b84d2

… if all extracts would lead to UB on poison.

juliannagele changed the title ~~[VecorCombine] Avoid inserting freeze when scalarizing extend-extract if all extracts would lead to UB on poison.~~ [VectorCombine] Avoid inserting freeze when scalarizing extend-extract if all extracts would lead to UB on poison. Oct 22, 2025

llvmbot added vectorizers llvm:transforms llvm:vectorcombine labels Oct 22, 2025

fhahn requested review from RKSimon, davemgreen, dtcxzyw and nikic October 23, 2025 08:41

fhahn reviewed Oct 23, 2025

View reviewed changes

nikic reviewed Oct 23, 2025

View reviewed changes

dtcxzyw mentioned this pull request Oct 24, 2025

Task submission dtcxzyw/llvm-opt-benchmark#1312

Open

This was referenced Oct 24, 2025

pre-commit: PR164683 dtcxzyw/llvm-opt-benchmark#2975

Closed

pre-commit: PR164683 dtcxzyw/llvm-opt-benchmark#2976

Closed

dtcxzyw reviewed Oct 24, 2025

View reviewed changes

llvm/lib/Transforms/Vectorize/VectorCombine.cpp Outdated Show resolved Hide resolved

fixup! [VecorCombine] Avoid inserting freeze when scalarizing extend-…

78259eb

…extract if all extracts would lead to UB on poison.

fhahn reviewed Oct 30, 2025

View reviewed changes

llvm/lib/Transforms/Vectorize/VectorCombine.cpp Outdated Show resolved Hide resolved

fixup! fixup! [VecorCombine] Avoid inserting freeze when scalarizing …

4b931b5

…extend-extract if all extracts would lead to UB on poison.

dtcxzyw reviewed Nov 1, 2025

View reviewed changes

llvm/lib/Transforms/Vectorize/VectorCombine.cpp Outdated Show resolved Hide resolved

llvm/lib/Transforms/Vectorize/VectorCombine.cpp Outdated Show resolved Hide resolved

fixup! fixup! fixup! [VecorCombine] Avoid inserting freeze when scala…

610cc37

…rizing extend-extract if all extracts would lead to UB on poison.

dtcxzyw approved these changes Nov 1, 2025

View reviewed changes

fhahn approved these changes Nov 2, 2025

View reviewed changes

fixup! fixup! fixup! fixup! [VecorCombine] Avoid inserting freeze whe…

c197ff6

…n scalarizing extend-extract if all extracts would lead to UB on poison.

juliannagele merged commit 28a20b4 into llvm:main Nov 4, 2025
10 checks passed

juliannagele mentioned this pull request Nov 5, 2025

[VectorCombine] Try to scalarize vector loads feeding bitcast instructions. #164682

Open

		// are guaranteed to execute if Ext executes.
		// If so, we do not need to insert a freeze.

[VectorCombine] Avoid inserting freeze when scalarizing extend-extract if all extracts would lead to UB on poison. #164683

[VectorCombine] Avoid inserting freeze when scalarizing extend-extract if all extracts would lead to UB on poison. #164683

Uh oh!

Conversation

juliannagele commented Oct 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Oct 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

fhahn left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dtcxzyw left a comment

Choose a reason for hiding this comment

Uh oh!

fhahn left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

llvm-ci commented Nov 4, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

juliannagele commented Oct 22, 2025 •

edited

Loading

llvmbot commented Oct 22, 2025 •

edited

Loading