[llvm] [InstCombine] Sink widening casts used only by extractelements (PR #201529)

Barbara Mitic via llvm-commits llvm-commits at lists.llvm.org
Fri Jun 12 04:17:02 PDT 2026


https://github.com/barbara-amd updated https://github.com/llvm/llvm-project/pull/201529

>From 5518064f10f177e96cafa6eed54adf3886dfecb9 Mon Sep 17 00:00:00 2001
From: Barbara Mitic <Barbara.Mitic at amd.com>
Date: Wed, 3 Jun 2026 15:25:22 +0200
Subject: [PATCH] [InstCombine] Sink widening casts used only by
 extractelements

InstCombine already canonicalizes extractelement(cast X) into
cast(extractelement X), but only when the cast has a single use. Extend
this to casts with multiple uses, provided the cast is a widening cast
(sext/zext/fpext) and every user is an extractelement that can itself
be sunk. Requiring every user to be sinkable guarantees the wide cast
is fully removed: the narrower operand keeps each scalar cast cheap,
and once all extracts are rewritten the wide cast has no remaining
users. If any user cannot be sunk (e.g. a variable-index extract in a
loop), the wide cast would survive next to the new scalar casts, so the
transform is skipped. Narrowing and fp-to-int casts keep the existing
single-use restriction, where multi-use scalarization could duplicate
bitcasts or turn a vector op into scalar ops in a loop. The existing
gate (bitcast and loop-profitability checks) is unchanged and moved
into a canSinkCastThroughExtract() helper for clarity.
---
 .../InstCombine/InstCombineVectorOps.cpp      | 52 ++++++++++---
 .../Transforms/InstCombine/ExtractCast.ll     | 77 +++++++++++++++++++
 .../Transforms/LoopVectorize/histograms.ll    |  7 +-
 .../AArch64/scalarize-load-ext-extract.ll     | 10 +--
 4 files changed, 125 insertions(+), 21 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index 6a3ee719db899..03de57c6bae71 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -410,6 +410,44 @@ static ConstantInt *getPreferredVectorIndex(ConstantInt *IndexC) {
                           IndexC->getValue().zextOrTrunc(64));
 }
 
+/// Decide whether extractelement(cast X), Index can be rewritten as
+/// cast(extractelement X, Index), where \p Cast is the source of the
+/// extractelement \p EI.
+static bool canSinkCastThroughExtract(CastInst *Cast, ExtractElementInst &EI,
+                                      Value *Index) {
+  // Bitcasts can change the number of vector elements; they are handled by
+  // foldBitcastExtElt instead.
+  if (Cast->getOpcode() == Instruction::BitCast)
+    return false;
+
+  // Sinking would move the cast next to the extractelement. Avoid sinking it
+  // into a loop: only do so when the index is constant or the cast and the
+  // extractelement are in the same basic block.
+  if (EI.getParent() != Cast->getParent() && !isa<ConstantInt>(Index))
+    return false;
+
+  // A single-use cast is always replaced (its only use is this extract).
+  if (Cast->hasOneUse())
+    return true;
+
+  // With multiple uses, only sink widening casts (sext/zext/fpext) whose every
+  // user is a sinkable extractelement (same block as the cast or a constant
+  // index). This ensures the wide cast is fully removed rather than left
+  // alongside the new scalar casts. Other casts are left to the single-use case.
+  switch (Cast->getOpcode()) {
+  case Instruction::SExt:
+  case Instruction::ZExt:
+  case Instruction::FPExt:
+    return all_of(Cast->users(), [&](User *U) {
+      auto *EE = dyn_cast<ExtractElementInst>(U);
+      return EE && (EE->getParent() == Cast->getParent() ||
+                    isa<ConstantInt>(EE->getIndexOperand()));
+    });
+  default:
+    return false;
+  }
+}
+
 Instruction *InstCombinerImpl::visitExtractElementInst(ExtractElementInst &EI) {
   Value *SrcVec = EI.getVectorOperand();
   Value *Index = EI.getIndexOperand();
@@ -600,18 +638,8 @@ Instruction *InstCombinerImpl::visitExtractElementInst(ExtractElementInst &EI) {
         }
       }
     } else if (auto *CI = dyn_cast<CastInst>(I)) {
-      // Canonicalize extractelement(cast) -> cast(extractelement).
-      // Bitcasts can change the number of vector elements, and they cost
-      // nothing.
-      // If the CI has only one use, but that use is inside a loop, this
-      // canonicalization is not profitable because it would turn a vector
-      // operation into scalar operations inside the loop. Apply the transform
-      // when:
-      //  - the index is constant and CI has one use, or
-      //  - the CI and EI are in the same basic block, so the cast won't be sunk
-      //    into a loop.
-      if (CI->hasOneUse() && (CI->getOpcode() != Instruction::BitCast) &&
-          (EI.getParent() == CI->getParent() || isa<ConstantInt>(Index))) {
+      // Canonicalize extractelement(cast X) -> cast(extractelement X).
+      if (canSinkCastThroughExtract(CI, EI, Index)) {
         Value *EE = Builder.CreateExtractElement(CI->getOperand(0), Index);
         return CastInst::Create(CI->getOpcode(), EE, EI.getType());
       }
diff --git a/llvm/test/Transforms/InstCombine/ExtractCast.ll b/llvm/test/Transforms/InstCombine/ExtractCast.ll
index d5ad40f0caf8b..e6b880fc3cd42 100644
--- a/llvm/test/Transforms/InstCombine/ExtractCast.ll
+++ b/llvm/test/Transforms/InstCombine/ExtractCast.ll
@@ -23,3 +23,80 @@ define i32 @b(<4 x float> %I) {
   ret i32 %K
 }
 
+define { i32, i32 } @sext_all_extracts(<2 x i16> %v) {
+; CHECK-LABEL: @sext_all_extracts(
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x i16> [[V:%.*]], i64 0
+; CHECK-NEXT:    [[X:%.*]] = sext i16 [[TMP1]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = extractelement <2 x i16> [[V]], i64 1
+; CHECK-NEXT:    [[Y:%.*]] = sext i16 [[TMP2]] to i32
+; CHECK-NEXT:    [[R0:%.*]] = insertvalue { i32, i32 } poison, i32 [[X]], 0
+; CHECK-NEXT:    [[R1:%.*]] = insertvalue { i32, i32 } [[R0]], i32 [[Y]], 1
+; CHECK-NEXT:    ret { i32, i32 } [[R1]]
+;
+  %e = sext <2 x i16> %v to <2 x i32>
+  %x = extractelement <2 x i32> %e, i64 0
+  %y = extractelement <2 x i32> %e, i64 1
+  %r0 = insertvalue { i32, i32 } poison, i32 %x, 0
+  %r1 = insertvalue { i32, i32 } %r0, i32 %y, 1
+  ret { i32, i32 } %r1
+}
+
+define { i32, i32 } @zext_all_extracts(<2 x i16> %v) {
+; CHECK-LABEL: @zext_all_extracts(
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x i16> [[V:%.*]], i64 0
+; CHECK-NEXT:    [[X:%.*]] = zext i16 [[TMP1]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = extractelement <2 x i16> [[V]], i64 1
+; CHECK-NEXT:    [[Y:%.*]] = zext i16 [[TMP2]] to i32
+; CHECK-NEXT:    [[R0:%.*]] = insertvalue { i32, i32 } poison, i32 [[X]], 0
+; CHECK-NEXT:    [[R1:%.*]] = insertvalue { i32, i32 } [[R0]], i32 [[Y]], 1
+; CHECK-NEXT:    ret { i32, i32 } [[R1]]
+;
+  %e = zext <2 x i16> %v to <2 x i32>
+  %x = extractelement <2 x i32> %e, i64 0
+  %y = extractelement <2 x i32> %e, i64 1
+  %r0 = insertvalue { i32, i32 } poison, i32 %x, 0
+  %r1 = insertvalue { i32, i32 } %r0, i32 %y, 1
+  ret { i32, i32 } %r1
+}
+
+define i32 @sext_mixed_use(<2 x i16> %v, ptr %p) {
+; CHECK-LABEL: @sext_mixed_use(
+; CHECK-NEXT:    [[E:%.*]] = sext <2 x i16> [[V:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[X:%.*]] = extractelement <2 x i32> [[E]], i64 0
+; CHECK-NEXT:    store <2 x i32> [[E]], ptr [[P:%.*]], align 8
+; CHECK-NEXT:    ret i32 [[X]]
+;
+  %e = sext <2 x i16> %v to <2 x i32>
+  %x = extractelement <2 x i32> %e, i64 0
+  store <2 x i32> %e, ptr %p
+  ret i32 %x
+}
+
+define void @zext_unsinkable_user(<2 x i32> %v, i64 %idx, i1 %cond, ptr %p, ptr %q) {
+; CHECK-LABEL: @zext_unsinkable_user(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[E:%.*]] = zext <2 x i32> [[V:%.*]] to <2 x i64>
+; CHECK-NEXT:    [[C:%.*]] = extractelement <2 x i64> [[E]], i64 0
+; CHECK-NEXT:    store i64 [[C]], ptr [[P:%.*]], align 4
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[OTHER:%.*]], label [[EXIT:%.*]]
+; CHECK:       other:
+; CHECK-NEXT:    [[A:%.*]] = extractelement <2 x i64> [[E]], i64 [[IDX:%.*]]
+; CHECK-NEXT:    store i64 [[A]], ptr [[Q:%.*]], align 4
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %e = zext <2 x i32> %v to <2 x i64>
+  %c = extractelement <2 x i64> %e, i64 0
+  store i64 %c, ptr %p
+  br i1 %cond, label %other, label %exit
+
+other:
+  %a = extractelement <2 x i64> %e, i64 %idx
+  store i64 %a, ptr %q
+  br label %exit
+
+exit:
+  ret void
+}
diff --git a/llvm/test/Transforms/LoopVectorize/histograms.ll b/llvm/test/Transforms/LoopVectorize/histograms.ll
index 5850ac3195c39..92bcbb49f2ccc 100644
--- a/llvm/test/Transforms/LoopVectorize/histograms.ll
+++ b/llvm/test/Transforms/LoopVectorize/histograms.ll
@@ -14,9 +14,10 @@ define void @simple_histogram(ptr noalias %buckets, ptr readonly %indices, i64 %
 ; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, [[ENTRY]] ], [ [[IV_NEXT:%.*]], [[FOR_BODY]] ]
 ; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds [4 x i8], ptr [[INDICES]], i64 [[IV]]
 ; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <2 x i32>, ptr [[TMP0]], align 4
-; CHECK-NEXT:    [[TMP1:%.*]] = zext <2 x i32> [[WIDE_LOAD]] to <2 x i64>
-; CHECK-NEXT:    [[TMP2:%.*]] = extractelement <2 x i64> [[TMP1]], i64 0
-; CHECK-NEXT:    [[TMP4:%.*]] = extractelement <2 x i64> [[TMP1]], i64 1
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x i32> [[WIDE_LOAD]], i64 0
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i32 [[TMP1]] to i64
+; CHECK-NEXT:    [[TMP9:%.*]] = extractelement <2 x i32> [[WIDE_LOAD]], i64 1
+; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[TMP9]] to i64
 ; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr inbounds nuw [4 x i8], ptr [[BUCKETS]], i64 [[TMP2]]
 ; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr inbounds nuw [4 x i8], ptr [[BUCKETS]], i64 [[TMP4]]
 ; CHECK-NEXT:    [[TMP6:%.*]] = insertelement <2 x ptr> poison, ptr [[TMP3]], i64 0
diff --git a/llvm/test/Transforms/PhaseOrdering/AArch64/scalarize-load-ext-extract.ll b/llvm/test/Transforms/PhaseOrdering/AArch64/scalarize-load-ext-extract.ll
index 8af7ae1b0ac77..e3c5bfbd43c0d 100644
--- a/llvm/test/Transforms/PhaseOrdering/AArch64/scalarize-load-ext-extract.ll
+++ b/llvm/test/Transforms/PhaseOrdering/AArch64/scalarize-load-ext-extract.ll
@@ -5,12 +5,10 @@ define noundef i32 @load_ext_extract(ptr %src) {
 ; CHECK-LABEL: define noundef range(i32 0, 1021) i32 @load_ext_extract(
 ; CHECK-SAME: ptr nofree readonly captures(none) [[SRC:%.*]]) local_unnamed_addr #[[ATTR0:[0-9]+]] {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[TMP14:%.*]] = load i32, ptr [[SRC]], align 4
-; CHECK-NEXT:    [[TMP0:%.*]] = insertelement <4 x i32> poison, i32 [[TMP14]], i64 0
-; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i32> [[TMP0]], <4 x i32> poison, <4 x i32> zeroinitializer
-; CHECK-NEXT:    [[TMP2:%.*]] = lshr <4 x i32> [[TMP1]], <i32 0, i32 8, i32 16, i32 24>
-; CHECK-NEXT:    [[TMP5:%.*]] = and <4 x i32> [[TMP2]], <i32 255, i32 255, i32 255, i32 -1>
-; CHECK-NEXT:    [[ADD3:%.*]] = tail call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> [[TMP5]])
+; CHECK-NEXT:    [[TMP0:%.*]] = load <4 x i8>, ptr [[SRC]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = zext <4 x i8> [[TMP0]] to <4 x i16>
+; CHECK-NEXT:    [[TMP2:%.*]] = tail call i16 @llvm.vector.reduce.add.v4i16(<4 x i16> [[TMP1]])
+; CHECK-NEXT:    [[ADD3:%.*]] = zext nneg i16 [[TMP2]] to i32
 ; CHECK-NEXT:    ret i32 [[ADD3]]
 ;
 entry:



More information about the llvm-commits mailing list