[VectorCombine] allow peeking through GEPs when creating a vector load

[spatel]

Summary

This is an enhancement motivated by https://llvm.org/PR16739 (see D92858 for another).

We can look through a GEP to find a base pointer that may be safe to use for a vector load. If so, then we shuffle (shift) the necessary vector element over to index 0.

Alive2 proof based on 1 of the regression tests:
https://alive2.llvm.org/ce/z/yPJLkh

The vector translation is independent of endian (verify by changing to leading 'E' in the datalayout string).

[spatel]

llvm/test/Transforms/VectorCombine/X86/load.ll

283

Note that the SSE2 cost model is conservatively giving:

{TTI::SK_PermuteSingleSrc, MVT::v8i16, 5}, // 2*pshuflw + 2*pshufhw
                                            // + pshufd/unpck

...so that's why we do not vectorize the tests here.

[lebedev.ri]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

148–150

I strongly suspect that you need to recalculate the Alignment here,
because i don't think the Offset-less pointer is guaranteed to still be Alignment-aligned.

[spatel]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

148–150

That's an interesting question. I think what we're doing here (even without this patch) is comparing the alignment of the original scalar load with the alignment of the pointer of a new vector load. But that's not very meaningful is it? For example, if we are loading an i16 with align 2 then does it matter whether the original pointer is at least align 2 for a load of v8i16?

I haven't looked at alignment requirements much. If there's another related transform that we can use as a template, let me know.

[lebedev.ri]

I'm having trouble coming up with an example because there appears to be a preexisting soundness problems, example: (CC @nlopes @aqjune)

define <8 x i16> @t(i8* align 128 dereferenceable(128) %base) {
  %ptr = getelementptr inbounds i8, i8* %base, i64 1
  %p = bitcast i8* %ptr to <8 x i16>*

  %gep = getelementptr inbounds <8 x i16>, <8 x i16>* %p, i64 0, i64 1
  %s = load i16, i16* %gep, align 1
  %r = insertelement <8 x i16> undef, i16 %s, i64 0
  ret <8 x i16> %r
}

/builddirs/llvm-project/build-Clang11-unknown$ /builddirs/llvm-project/build-Clang11-unknown/bin/opt -load /repositories/alive2/build-Clang-release/tv/tv.so -tv -vector-combine -mtriple=x86_64-- -mattr=avx2 -tv -o /dev/null --tv-smt-to=60000 /tmp/D93229.ll 

----------------------------------------
define <8 x i16> @t(* dereferenceable(128) align(128) %base) {
%0:
  %ptr = gep inbounds * dereferenceable(128) align(128) %base, 1 x i64 1
  %p = bitcast * %ptr to *
  %gep = gep inbounds * %p, 16 x i64 0, 2 x i64 1
  %s = load i16, * %gep, align 1
  %r = insertelement <8 x i16> undef, i16 %s, i64 0
  ret <8 x i16> %r
}
=>
define <8 x i16> @t(* dereferenceable(128) align(128) %base) {
%0:
  %ptr = gep inbounds * dereferenceable(128) align(128) %base, 1 x i64 1
  %p = bitcast * %ptr to *
  %gep = gep inbounds * %p, 16 x i64 0, 2 x i64 1
  %1 = bitcast * %gep to *
  %r = load <8 x i16>, * %1, align 1
  ret <8 x i16> %r
}
Transformation doesn't verify!
ERROR: Target is more poisonous than source

Example:
* dereferenceable(128) align(128) %base = pointer(non-local, block_id=1, offset=1664)

Source:
* %ptr = pointer(non-local, block_id=1, offset=1665)
* %p = pointer(non-local, block_id=1, offset=1665)
* %gep = pointer(non-local, block_id=1, offset=1667)
i16 %s = poison
<8 x i16> %r = < poison, any, any, any, any, any, any, any >

SOURCE MEMORY STATE
===================
NON-LOCAL BLOCKS:
Block 0 >       size: 0 align: 1        alloc type: 0
Block 1 >       size: 2048      align: 128      alloc type: 0

Target:
* %ptr = pointer(non-local, block_id=1, offset=1665)
* %p = pointer(non-local, block_id=1, offset=1665)
* %gep = pointer(non-local, block_id=1, offset=1667)
* %1 = pointer(non-local, block_id=1, offset=1667)
<8 x i16> %r = < poison, poison, poison, poison, poison, poison, poison, poison >
Source value: < poison, any, any, any, any, any, any, any >
Target value: < poison, poison, poison, poison, poison, poison, poison, poison >

Alive2: Transform doesn't verify!

[spatel]

In D93229#2452695, @lebedev.ri wrote:

I'm having trouble coming up with an example because there appears to be a preexisting soundness problems, example: (CC @nlopes @aqjune)

define <8 x i16> @t(i8* align 128 dereferenceable(128) %base) {
  %ptr = getelementptr inbounds i8, i8* %base, i64 1
  %p = bitcast i8* %ptr to <8 x i16>*

  %gep = getelementptr inbounds <8 x i16>, <8 x i16>* %p, i64 0, i64 1
  %s = load i16, i16* %gep, align 1
  %r = insertelement <8 x i16> undef, i16 %s, i64 0
  ret <8 x i16> %r
}

/builddirs/llvm-project/build-Clang11-unknown$ /builddirs/llvm-project/build-Clang11-unknown/bin/opt -load /repositories/alive2/build-Clang-release/tv/tv.so -tv -vector-combine -mtriple=x86_64-- -mattr=avx2 -tv -o /dev/null --tv-smt-to=60000 /tmp/D93229.ll 

----------------------------------------
define <8 x i16> @t(* dereferenceable(128) align(128) %base) {
%0:
  %ptr = gep inbounds * dereferenceable(128) align(128) %base, 1 x i64 1
  %p = bitcast * %ptr to *
  %gep = gep inbounds * %p, 16 x i64 0, 2 x i64 1
  %s = load i16, * %gep, align 1
  %r = insertelement <8 x i16> undef, i16 %s, i64 0
  ret <8 x i16> %r
}
=>
define <8 x i16> @t(* dereferenceable(128) align(128) %base) {
%0:
  %ptr = gep inbounds * dereferenceable(128) align(128) %base, 1 x i64 1
  %p = bitcast * %ptr to *
  %gep = gep inbounds * %p, 16 x i64 0, 2 x i64 1
  %1 = bitcast * %gep to *
  %r = load <8 x i16>, * %1, align 1
  ret <8 x i16> %r
}
Transformation doesn't verify!
ERROR: Target is more poisonous than source

Example:
* dereferenceable(128) align(128) %base = pointer(non-local, block_id=1, offset=1664)

Source:
* %ptr = pointer(non-local, block_id=1, offset=1665)
* %p = pointer(non-local, block_id=1, offset=1665)
* %gep = pointer(non-local, block_id=1, offset=1667)
i16 %s = poison
<8 x i16> %r = < poison, any, any, any, any, any, any, any >

SOURCE MEMORY STATE
===================
NON-LOCAL BLOCKS:
Block 0 >       size: 0 align: 1        alloc type: 0
Block 1 >       size: 2048      align: 128      alloc type: 0

Target:
* %ptr = pointer(non-local, block_id=1, offset=1665)
* %p = pointer(non-local, block_id=1, offset=1665)
* %gep = pointer(non-local, block_id=1, offset=1667)
* %1 = pointer(non-local, block_id=1, offset=1667)
<8 x i16> %r = < poison, poison, poison, poison, poison, poison, poison, poison >
Source value: < poison, any, any, any, any, any, any, any >
Target value: < poison, poison, poison, poison, poison, poison, poison, poison >

Alive2: Transform doesn't verify!

IIUC, this is a question of allowing poison (from the unused loaded memory elements) to propagate?
So we have to freeze or explicitly make those elements undef again?
https://alive2.llvm.org/ce/z/LKqBVW

[lebedev.ri]

In D93229#2452751, @spatel wrote:

In D93229#2452695, @lebedev.ri wrote:

I'm having trouble coming up with an example because there appears to be a preexisting soundness problems, example: (CC @nlopes @aqjune)

define <8 x i16> @t(i8* align 128 dereferenceable(128) %base) {
  %ptr = getelementptr inbounds i8, i8* %base, i64 1
  %p = bitcast i8* %ptr to <8 x i16>*

  %gep = getelementptr inbounds <8 x i16>, <8 x i16>* %p, i64 0, i64 1
  %s = load i16, i16* %gep, align 1
  %r = insertelement <8 x i16> undef, i16 %s, i64 0
  ret <8 x i16> %r
}

/builddirs/llvm-project/build-Clang11-unknown$ /builddirs/llvm-project/build-Clang11-unknown/bin/opt -load /repositories/alive2/build-Clang-release/tv/tv.so -tv -vector-combine -mtriple=x86_64-- -mattr=avx2 -tv -o /dev/null --tv-smt-to=60000 /tmp/D93229.ll 

----------------------------------------
define <8 x i16> @t(* dereferenceable(128) align(128) %base) {
%0:
  %ptr = gep inbounds * dereferenceable(128) align(128) %base, 1 x i64 1
  %p = bitcast * %ptr to *
  %gep = gep inbounds * %p, 16 x i64 0, 2 x i64 1
  %s = load i16, * %gep, align 1
  %r = insertelement <8 x i16> undef, i16 %s, i64 0
  ret <8 x i16> %r
}
=>
define <8 x i16> @t(* dereferenceable(128) align(128) %base) {
%0:
  %ptr = gep inbounds * dereferenceable(128) align(128) %base, 1 x i64 1
  %p = bitcast * %ptr to *
  %gep = gep inbounds * %p, 16 x i64 0, 2 x i64 1
  %1 = bitcast * %gep to *
  %r = load <8 x i16>, * %1, align 1
  ret <8 x i16> %r
}
Transformation doesn't verify!
ERROR: Target is more poisonous than source

Example:
* dereferenceable(128) align(128) %base = pointer(non-local, block_id=1, offset=1664)

Source:
* %ptr = pointer(non-local, block_id=1, offset=1665)
* %p = pointer(non-local, block_id=1, offset=1665)
* %gep = pointer(non-local, block_id=1, offset=1667)
i16 %s = poison
<8 x i16> %r = < poison, any, any, any, any, any, any, any >

SOURCE MEMORY STATE
===================
NON-LOCAL BLOCKS:
Block 0 >       size: 0 align: 1        alloc type: 0
Block 1 >       size: 2048      align: 128      alloc type: 0

Target:
* %ptr = pointer(non-local, block_id=1, offset=1665)
* %p = pointer(non-local, block_id=1, offset=1665)
* %gep = pointer(non-local, block_id=1, offset=1667)
* %1 = pointer(non-local, block_id=1, offset=1667)
<8 x i16> %r = < poison, poison, poison, poison, poison, poison, poison, poison >
Source value: < poison, any, any, any, any, any, any, any >
Target value: < poison, poison, poison, poison, poison, poison, poison, poison >

Alive2: Transform doesn't verify!

IIUC, this is a question of allowing poison (from the unused loaded memory elements) to propagate?
So we have to freeze or explicitly make those elements undef again?
https://alive2.llvm.org/ce/z/LKqBVW

That is how i read it, yes. That will be gone in codegen, so no need to cost that extra legality shuffle.

[nlopes]

In D93229#2452695, @lebedev.ri wrote:

/builddirs/llvm-project/build-Clang11-unknown$ /builddirs/llvm-project/build-Clang11-unknown/bin/opt -load /repositories/alive2/build-Clang-release/tv/tv.so -tv -vector-combine -mtriple=x86_64-- -mattr=avx2 -tv -o /dev/null --tv-smt-to=60000 /tmp/D93229.ll 

----------------------------------------
define <8 x i16> @t(* dereferenceable(128) align(128) %base) {
%0:
  %ptr = gep inbounds * dereferenceable(128) align(128) %base, 1 x i64 1
  %p = bitcast * %ptr to *
  %gep = gep inbounds * %p, 16 x i64 0, 2 x i64 1
  %s = load i16, * %gep, align 1
  %r = insertelement <8 x i16> undef, i16 %s, i64 0
  ret <8 x i16> %r
}
=>
define <8 x i16> @t(* dereferenceable(128) align(128) %base) {
%0:
  %ptr = gep inbounds * dereferenceable(128) align(128) %base, 1 x i64 1
  %p = bitcast * %ptr to *
  %gep = gep inbounds * %p, 16 x i64 0, 2 x i64 1
  %1 = bitcast * %gep to *
  %r = load <8 x i16>, * %1, align 1
  ret <8 x i16> %r
}
Transformation doesn't verify!

Makes sense; it replaces an undef vector with a potential poison vector that might be in memory.
We need to switch to using poison as vector placeholders rather than undef so these problems go away. I guess clang needs a bit of patching (or IRBuilder, or both?).

[aqjune]

Yes, we might want to know who's generating insertvalue undef and replace it with insertvalue poison.
The shufflevector pattern works, but using poison as a placeholder will remove latent bugs.
Maybe it is time to use poison?

[spatel]

Patch updated:
Rebased after D93238 (fix poison bug) and added code to update the alignment + extra test to confirm.
There was an existing test with conflicting alignment specifiers that also changes now. Stepping through that, I noticed that getPointerAlignment() does not peek through gep, so there may still be room for improvement.

[spatel]

In D93229#2453538, @aqjune wrote:

Yes, we might want to know who's generating insertvalue undef and replace it with insertvalue poison.
The shufflevector pattern works, but using poison as a placeholder will remove latent bugs.
Maybe it is time to use poison?

If we switch to initialize with poison elements, I think we would have to swap undef with poison in regression tests and pattern matching in instsimplify/instcombine at the same time. Otherwise, we will be testing/folding the wrong patterns?

[aqjune]

In D93229#2454698, @spatel wrote:

In D93229#2453538, @aqjune wrote:

Yes, we might want to know who's generating insertvalue undef and replace it with insertvalue poison.
The shufflevector pattern works, but using poison as a placeholder will remove latent bugs.
Maybe it is time to use poison?

If we switch to initialize with poison elements, I think we would have to swap undef with poison in regression tests and pattern matching in instsimplify/instcombine at the same time. Otherwise, we will be testing/folding the wrong patterns?

Yes, I think so. What do you think about having three patches
(1) Poison placeholder patch (2) InstSimpify/InstCombine patch (3) The regression tests update patch.
and landing those consecutively when all of those are accepted? Patch (1) and (2) should still pass ninja check. It will show which tests are related to what.
I think I can prepare (2) and (3). The changes in regression tests (3) can be syntactically done with a script.
I tried this and it worked: sed -i.backup 's/insertelement <\(.*\)> undef,/insertelement <\1> poison,/g' (filename)
(sorry, I meant insertelement, not insertvalue)

[lebedev.ri]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

148–150

I think we don't need to ever ask isSafeToLoadUnconditionally() about alignment
(i.e. just pass 0/1, whatever is more permissive),
because we already know the alignment of the load.

We simply need to recalculate the alignment after chopping off the offset
(see commonAlignment(Align A, uint64_t Offset))

[spatel]

In D93229#2456676, @aqjune wrote:

In D93229#2454698, @spatel wrote:

In D93229#2453538, @aqjune wrote:

Yes, we might want to know who's generating insertvalue undef and replace it with insertvalue poison.
The shufflevector pattern works, but using poison as a placeholder will remove latent bugs.
Maybe it is time to use poison?

If we switch to initialize with poison elements, I think we would have to swap undef with poison in regression tests and pattern matching in instsimplify/instcombine at the same time. Otherwise, we will be testing/folding the wrong patterns?

Yes, I think so. What do you think about having three patches
(1) Poison placeholder patch (2) InstSimpify/InstCombine patch (3) The regression tests update patch.
and landing those consecutively when all of those are accepted? Patch (1) and (2) should still pass ninja check. It will show which tests are related to what.
I think I can prepare (2) and (3). The changes in regression tests (3) can be syntactically done with a script.
I tried this and it worked: sed -i.backup 's/insertelement <\(.*\)> undef,/insertelement <\1> poison,/g' (filename)

We can use an alternate plan to make incremental progress, but it will cause redundancy while we transition:

1. Replicate all of the regression tests with the poison constant instead of undef. If we add some unique TODO text marker on all of those new tests, we can then grep to make sure everything that we expect to get updated is actually updated in later steps.
2. Update instcombine/simplify/vectorizer folds to match poison patterns (create a m_UndefOrPoison() pattern matcher?)
3. Run Alive2 through all of the updated regression tests to verify.
4. Update codegen to deal with poison constant?
5. Update folds from step 2 to create poison rather than undef (if that's what they were doing)?
6. Change IRBuilder or other instruction creators to create poison from the start.

[spatel]

Patch updated:
D93397 / D93406 improved the basic alignment logic, so that is adjusted here to include the effect of gep offset.
In the last test diff, note that the final alignment is not the same as either of the starting alignments.

[lebedev.ri]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

159

So we had SrcPtr, and split it into a base pointer SrcPtr', and an offset Offset.
But i think it's SrcPtr = SrcPtr' + Offset,
so shouldn't the offset be negative here,
because the alignment is known for SrcPtr, not SrcPtr'?

[aqjune]

In D93229#2457696, @spatel wrote:

We can use an alternate plan to make incremental progress, but it will cause redundancy while we transition:

1. Replicate all of the regression tests with the poison constant instead of undef. If we add some unique TODO text marker on all of those new tests, we can then grep to make sure everything that we expect to get updated is actually updated in later steps.
2. Update instcombine/simplify/vectorizer folds to match poison patterns (create a m_UndefOrPoison() pattern matcher?)

m_Undef already matches poison because PoisonValue is a subclass of UndefValue :)

3. Run Alive2 through all of the updated regression tests to verify.
4. Update codegen to deal with poison constant?
5. Update folds from step 2 to create poison rather than undef (if that's what they were doing)?
6. Change IRBuilder or other instruction creators to create poison from the start.

I think we can split the goal and first work on replacing the placeholder value that is used by insertelement only.
In this case, actually InstSimplify/InstCombine change isn't necessary as well. It can be done only when suboptimal assembly is being generated.
What do you think?

[spatel]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

159

Good point!
Intuitively, I couldn't see how it would make a difference if we negated the offset, so I looked at the implementation of MinAlign():
https://github.com/llvm/llvm-project/blob/4c8276cdc120c24410dcd62a9986f04e7327fc2f/llvm/include/llvm/Support/MathExtras.h#L673

The optimized IR for that is:

define i64 @MinAlign(i64 %A, i64 %B) local_unnamed_addr #0 {
entry:
  %or = or i64 %B, %A
  %add = sub i64 0, %or
  %and = and i64 %or, %add
  ret i64 %and
}

Double-check to make sure I didn't mess this up:
https://alive2.llvm.org/ce/z/wRSjPD

So it's logically equivalent either way? At the least, I'll add a code comment.

[lebedev.ri]

Memory is confusing, pointing out some more potential issues.

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

154–156

Is there a test coverage for edge case[s]?
Doesn't this check only ensure that we can load a single byte, not the entire element?

158

Don't we need to ensure that the byte offset is a multiple of element size?

159

STGM

195

Assert that the division is exact?

[lebedev.ri]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

195

And also assert that Mask[0] < MinVecNumElts.

[spatel]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

158

Yes, good catch. We have potentially looked through casts at this point, so anything is possible. I'll add test(s).

[spatel]

Patch updated:

1. Added constraints on address offset with respect to vector element size.
2. Added negative tests for those.
3. Added asserts to make sure address assumptions are valid.
4. Added code comment about offset math (negation) logic.

[spatel]

In D93229#2459265, @aqjune wrote:

In D93229#2457696, @spatel wrote:

We can use an alternate plan to make incremental progress, but it will cause redundancy while we transition:

1. Replicate all of the regression tests with the poison constant instead of undef. If we add some unique TODO text marker on all of those new tests, we can then grep to make sure everything that we expect to get updated is actually updated in later steps.
2. Update instcombine/simplify/vectorizer folds to match poison patterns (create a m_UndefOrPoison() pattern matcher?)

m_Undef already matches poison because PoisonValue is a subclass of UndefValue :)

3. Run Alive2 through all of the updated regression tests to verify.
4. Update codegen to deal with poison constant?
5. Update folds from step 2 to create poison rather than undef (if that's what they were doing)?
6. Change IRBuilder or other instruction creators to create poison from the start.

I think we can split the goal and first work on replacing the placeholder value that is used by insertelement only.
In this case, actually InstSimplify/InstCombine change isn't necessary as well. It can be done only when suboptimal assembly is being generated.
What do you think?

Ok - can you create that patch? (We should move this conversation to llvm-dev or another review to reduce confusion.)

[lebedev.ri]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

151–152

I think this should be done later, after cheaper checks.

164–166

We want to load a single element of X bytes,
by instead loading Y such elements at once.
We've dissected the pointer into a base pointer, and a byte offset.
We know that the byte offset is a multiple of element size.
We need to ensure that if we load Y elements from the base pointer,
we still load the element we are after.
I think approaching that check from byte count is highly confusing.
(at least i already spent too much time trying to check this)
Let's do a much more obvious, element count based check.

[spatel]

Patch updated:
Adjusted order/spelling of offset checks as suggested. Also added a test (gep012) just below the offset cut-off, so we have positive and negative tests at the boundary condition.

[lebedev.ri]

Ok, this looks about right to me.

I'm not sure if we'll want to lift the restrictions that

(1) the offset must be a multiple of element size

and (2) that we try to load from the base pointer, which only works if the offset is small enough.

I've checked, and alive2 claims that this is endianness-agnostic.

Thanks.

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

141–161

And now that we already calculate NumEltsInOffset,
shall we just use/record NumEltsInOffset instead of recalculating it from OffsetInBits later?
(It them might use a better name, something like ElementIndex maybe)

195

(If we no longer recompute them here, the asserts can obviously go away to)

[aqjune]

In D93229#2461069, @spatel wrote:

In D93229#2459265, @aqjune wrote:

In D93229#2457696, @spatel wrote:

We can use an alternate plan to make incremental progress, but it will cause redundancy while we transition:

1. Replicate all of the regression tests with the poison constant instead of undef. If we add some unique TODO text marker on all of those new tests, we can then grep to make sure everything that we expect to get updated is actually updated in later steps.
2. Update instcombine/simplify/vectorizer folds to match poison patterns (create a m_UndefOrPoison() pattern matcher?)

m_Undef already matches poison because PoisonValue is a subclass of UndefValue :)

3. Run Alive2 through all of the updated regression tests to verify.
4. Update codegen to deal with poison constant?
5. Update folds from step 2 to create poison rather than undef (if that's what they were doing)?
6. Change IRBuilder or other instruction creators to create poison from the start.

I think we can split the goal and first work on replacing the placeholder value that is used by insertelement only.
In this case, actually InstSimplify/InstCombine change isn't necessary as well. It can be done only when suboptimal assembly is being generated.
What do you think?

Ok - can you create that patch? (We should move this conversation to llvm-dev or another review to reduce confusion.)

Sure. I'll make a patch this weekend and send a mail to llvm-dev.

[spatel]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

141–161

Yes - no need to carry the offset as both an index and a bit value. Keeping both makes it harder to read. Thanks for the thorough review!

[aqjune]

I made a patch here: D93586.
It touches InstCombinerImpl::SimplifyDemandedVectorElts only to keep the size of diff manageable.
Sent a mail to llvm-dev too

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

 
 static void replaceValue(Value &Old, Value &New) {
   Old.replaceAllUsesWith(&New);
   New.takeName(&Old);
 }
 
 bool VectorCombine::vectorizeLoadInsert(Instruction &I) {
   // Match insert into fixed vector of scalar value.
+  // TODO: Handle non-zero insert index.
   auto *Ty = dyn_cast<FixedVectorType>(I.getType());
   Value *Scalar;
   if (!Ty || !match(&I, m_InsertElt(m_Undef(), m_Value(Scalar), m_ZeroInt())) ||
       !Scalar->hasOneUse())
     return false;
 
   // Optionally match an extract from another vector.
   Value *X;
   bool HasExtract = match(Scalar, m_ExtractElt(m_Value(X), m_ZeroInt()));
   if (!HasExtract)
     X = Scalar;
 
   // Match source value as load of scalar or vector.
   // Do not vectorize scalar load (widening) if atomic/volatile or under
   // asan/hwasan/memtag/tsan. The widened load may load data from dirty regions
   // or create data races non-existent in the source.
   auto *Load = dyn_cast<LoadInst>(X);
   if (!Load || !Load->isSimple() || !Load->hasOneUse() ||
       Load->getFunction()->hasFnAttribute(Attribute::SanitizeMemTag) ||
       mustSuppressSpeculation(*Load))
     return false;
 
-  // TODO: Extend this to match GEP with constant offsets.
   const DataLayout &DL = I.getModule()->getDataLayout();
   Value *SrcPtr = Load->getPointerOperand()->stripPointerCasts();
   assert(isa<PointerType>(SrcPtr->getType()) && "Expected a pointer type");
 
   // If original AS != Load's AS, we can't bitcast the original pointer and have
   // to use Load's operand instead. Ideally we would want to strip pointer casts
   // without changing AS, but there's no API to do that ATM.
   unsigned AS = Load->getPointerAddressSpace();
   if (AS != SrcPtr->getType()->getPointerAddressSpace())
     SrcPtr = Load->getPointerOperand();
 
   Type *ScalarTy = Scalar->getType();
   uint64_t ScalarSize = ScalarTy->getPrimitiveSizeInBits();
   unsigned MinVectorSize = TTI.getMinVectorRegisterBitWidth();
   if (!ScalarSize || !MinVectorSize || MinVectorSize % ScalarSize != 0)
     return false;
 
   // Check safety of replacing the scalar load with a larger vector load.
   unsigned MinVecNumElts = MinVectorSize / ScalarSize;
   auto *MinVecTy = VectorType::get(ScalarTy, MinVecNumElts, false);
+  // TODO: also check getPointerAlignment() of SrcPtr?
   Align Alignment = Load->getAlign();
-  if (!isSafeToLoadUnconditionally(SrcPtr, MinVecTy, Alignment, DL, Load, &DT))
+  unsigned OffsetInBits = 0;
+  if (!isSafeToLoadUnconditionally(SrcPtr, MinVecTy, Alignment, DL, Load,
+                                   &DT)) {
+    // It is not safe to load directly from the pointer, but we can still peek
+    // through gep offsets and check if it safe to load from a base address with
+    // updated alignment. If it is, we can shuffle the element(s) into place
+    // after loading.
+    unsigned OffsetBitWidth = DL.getIndexTypeSizeInBits(SrcPtr->getType());
+    APInt Offset(OffsetBitWidth, 0);
+    SrcPtr = SrcPtr->stripAndAccumulateInBoundsConstantOffsets(DL, Offset);
+    Alignment = SrcPtr->getPointerAlignment(DL);
+    if (!isSafeToLoadUnconditionally(SrcPtr, MinVecTy, Alignment, DL, Load,
+                                     &DT))
       return false;
 
+    // The offset must be within a vector-length to allow shuffling into place.
+    if (Offset.uge(MinVectorSize / 8))
+      return false;
+
+    OffsetInBits = Offset.getZExtValue() * 8;
+  }
+
   // Original pattern: insertelt undef, load [free casts of] PtrOp, 0
   Type *LoadTy = Load->getType();
   int OldCost = TTI.getMemoryOpCost(Instruction::Load, LoadTy, Alignment, AS);
   APInt DemandedElts = APInt::getOneBitSet(MinVecNumElts, 0);
   OldCost += TTI.getScalarizationOverhead(MinVecTy, DemandedElts,
                                           /* Insert */ true, HasExtract);
 
   // New pattern: load VecPtr
   int NewCost = TTI.getMemoryOpCost(Instruction::Load, MinVecTy, Alignment, AS);
+  // Optionally, we are shuffling the loaded vector element(s) into place.
+  if (OffsetInBits)
+    NewCost += TTI.getShuffleCost(TTI::SK_PermuteSingleSrc, MinVecTy);
 
   // We can aggressively convert to the vector form because the backend can
   // invert this transform if it does not result in a performance win.
   if (OldCost < NewCost)
     return false;
 
   // It is safe and potentially profitable to load a vector directly:
   // inselt undef, load Scalar, 0 --> load VecPtr
   IRBuilder<> Builder(Load);
   Value *CastedPtr = Builder.CreateBitCast(SrcPtr, MinVecTy->getPointerTo(AS));
   Value *VecLd = Builder.CreateAlignedLoad(MinVecTy, CastedPtr, Alignment);
 
   // Set everything but element 0 to undef to prevent poison from propagating
   // from the extra loaded memory. This will also optionally shrink/grow the
   // vector from the loaded size to the output size.
-  // We assume this operation has no cost in codegen.
+  // We assume this operation has no cost in codegen if there was no offset.
   // Note that we could use freeze to avoid poison problems, but then we might
   // still need a shuffle to change the vector size.
   unsigned OutputNumElts = Ty->getNumElements();
   SmallVector<int, 16> Mask(OutputNumElts, UndefMaskElem);
-  Mask[0] = 0;
+  Mask[0] = OffsetInBits / ScalarSize;
   VecLd = Builder.CreateShuffleVector(VecLd, Mask);
 
   replaceValue(I, *VecLd);
   ++NumVecLoad;
   return true;
 }
 
 /// Determine which, if any, of the inputs should be replaced by a shuffle

llvm/test/Transforms/VectorCombine/X86/load.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt < %s -vector-combine -S -mtriple=x86_64-- -mattr=sse2 | FileCheck %s
-; RUN: opt < %s -vector-combine -S -mtriple=x86_64-- -mattr=avx2 | FileCheck %s
+; RUN: opt < %s -vector-combine -S -mtriple=x86_64-- -mattr=sse2 | FileCheck %s --check-prefixes=CHECK,SSE2
+; RUN: opt < %s -vector-combine -S -mtriple=x86_64-- -mattr=avx2 | FileCheck %s --check-prefixes=CHECK,AVX2
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 
 define float @matching_fp_scalar(float* align 16 dereferenceable(16) %p) {
 ; CHECK-LABEL: @matching_fp_scalar(
 ; CHECK-NEXT:    [[R:%.*]] = load float, float* [[P:%.*]], align 16
 ; CHECK-NEXT:    ret float [[R]]
 ;

 ; CHECK-NEXT:    ret <8 x i16> [[R]]
 ;
   %gep = getelementptr inbounds <8 x i16>, <8 x i16>* %p, i64 0, i64 1
   %s = load i16, i16* %gep, align 2
   %r = insertelement <8 x i16> undef, i16 %s, i64 0
   ret <8 x i16> %r
 }
 
-; Negative test - can't safely load the offset vector, but could load+shuffle.
+; Can't safely load the offset vector, but can load+shuffle if it is profitable.
 
 define <8 x i16> @gep01_load_i16_insert_v8i16_deref(<8 x i16>* align 16 dereferenceable(17) %p) {
-; CHECK-LABEL: @gep01_load_i16_insert_v8i16_deref(
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds <8 x i16>, <8 x i16>* [[P:%.*]], i64 0, i64 1
-; CHECK-NEXT:    [[S:%.*]] = load i16, i16* [[GEP]], align 2
-; CHECK-NEXT:    [[R:%.*]] = insertelement <8 x i16> undef, i16 [[S]], i64 0
-; CHECK-NEXT:    ret <8 x i16> [[R]]
+; SSE2-LABEL: @gep01_load_i16_insert_v8i16_deref(
+; SSE2-NEXT:    [[GEP:%.*]] = getelementptr inbounds <8 x i16>, <8 x i16>* [[P:%.*]], i64 0, i64 1
+; SSE2-NEXT:    [[S:%.*]] = load i16, i16* [[GEP]], align 2
+; SSE2-NEXT:    [[R:%.*]] = insertelement <8 x i16> undef, i16 [[S]], i64 0
+; SSE2-NEXT:    ret <8 x i16> [[R]]
+;
+; AVX2-LABEL: @gep01_load_i16_insert_v8i16_deref(
+; AVX2-NEXT:    [[TMP1:%.*]] = load <8 x i16>, <8 x i16>* [[P:%.*]], align 16
+; AVX2-NEXT:    [[R:%.*]] = shufflevector <8 x i16> [[TMP1]], <8 x i16> undef, <8 x i32> <i32 1, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>
+; AVX2-NEXT:    ret <8 x i16> [[R]]
 ;
   %gep = getelementptr inbounds <8 x i16>, <8 x i16>* %p, i64 0, i64 1
   %s = load i16, i16* %gep, align 2
   %r = insertelement <8 x i16> undef, i16 %s, i64 0
   ret <8 x i16> %r
 }
 
-; TODO: Verify that alignment of the new load is not over-specified.
+; Verify that alignment of the new load is not over-specified.
 
 define <8 x i16> @gep01_load_i16_insert_v8i16_deref_minalign(<8 x i16>* align 2 dereferenceable(16) %p) {
-; CHECK-LABEL: @gep01_load_i16_insert_v8i16_deref_minalign(
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds <8 x i16>, <8 x i16>* [[P:%.*]], i64 0, i64 1
-; CHECK-NEXT:    [[S:%.*]] = load i16, i16* [[GEP]], align 8
-; CHECK-NEXT:    [[R:%.*]] = insertelement <8 x i16> undef, i16 [[S]], i64 0
-; CHECK-NEXT:    ret <8 x i16> [[R]]
+; SSE2-LABEL: @gep01_load_i16_insert_v8i16_deref_minalign(
+; SSE2-NEXT:    [[GEP:%.*]] = getelementptr inbounds <8 x i16>, <8 x i16>* [[P:%.*]], i64 0, i64 1
+; SSE2-NEXT:    [[S:%.*]] = load i16, i16* [[GEP]], align 8
+; SSE2-NEXT:    [[R:%.*]] = insertelement <8 x i16> undef, i16 [[S]], i64 0
+; SSE2-NEXT:    ret <8 x i16> [[R]]
+;
+; AVX2-LABEL: @gep01_load_i16_insert_v8i16_deref_minalign(
+; AVX2-NEXT:    [[TMP1:%.*]] = load <8 x i16>, <8 x i16>* [[P:%.*]], align 2
+; AVX2-NEXT:    [[R:%.*]] = shufflevector <8 x i16> [[TMP1]], <8 x i16> undef, <8 x i32> <i32 1, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>
+; AVX2-NEXT:    ret <8 x i16> [[R]]
 ;
   %gep = getelementptr inbounds <8 x i16>, <8 x i16>* %p, i64 0, i64 1
   %s = load i16, i16* %gep, align 8
   %r = insertelement <8 x i16> undef, i16 %s, i64 0
   ret <8 x i16> %r
 }
 
 ; If there are enough dereferenceable bytes, we can offset the vector load.

 ; CHECK-NEXT:    [[R:%.*]] = insertelement <4 x float> undef, float [[S]], i32 0
 ; CHECK-NEXT:    ret <4 x float> [[R]]
 ;
   %s = load volatile float, float* %p, align 4
   %r = insertelement <4 x float> undef, float %s, i32 0
   ret <4 x float> %r
 }
 
-; Negative test? - pointer is not as aligned as load.
+; Pointer is not as aligned as load, but that's ok.
+; TODO: Should we choose the maximum alignment instead?
 
 define <4 x float> @load_f32_insert_v4f32_align(float* align 1 dereferenceable(16) %p) {
 ; CHECK-LABEL: @load_f32_insert_v4f32_align(
-; CHECK-NEXT:    [[S:%.*]] = load float, float* [[P:%.*]], align 4
-; CHECK-NEXT:    [[R:%.*]] = insertelement <4 x float> undef, float [[S]], i32 0
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast float* [[P:%.*]] to <4 x float>*
+; CHECK-NEXT:    [[TMP2:%.*]] = load <4 x float>, <4 x float>* [[TMP1]], align 1
+; CHECK-NEXT:    [[R:%.*]] = shufflevector <4 x float> [[TMP2]], <4 x float> undef, <4 x i32> <i32 0, i32 undef, i32 undef, i32 undef>
 ; CHECK-NEXT:    ret <4 x float> [[R]]
 ;
   %s = load float, float* %p, align 4
   %r = insertelement <4 x float> undef, float %s, i32 0
   ret <4 x float> %r
 }
 
 ; Negative test - not enough bytes.

 ;
   %l = load <1 x i32>, <1 x i32>* %p, align 4
   store <1 x i32> %l, <1 x i32>* %store_ptr
   %s = extractelement <1 x i32> %l, i32 0
   %r = insertelement <8 x i32> undef, i32 %s, i32 0
   ret <8 x i32> %r
 }
 
-; TODO: Can't safely load the offset vector, but can load+shuffle if it is profitable.
+; Can't safely load the offset vector, but can load+shuffle if it is profitable.
 
 define <8 x i16> @gep1_load_v2i16_extract_insert_v8i16(<2 x i16>* align 16 dereferenceable(16) %p) {
-; CHECK-LABEL: @gep1_load_v2i16_extract_insert_v8i16(
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds <2 x i16>, <2 x i16>* [[P:%.*]], i64 1
-; CHECK-NEXT:    [[L:%.*]] = load <2 x i16>, <2 x i16>* [[GEP]], align 2
-; CHECK-NEXT:    [[S:%.*]] = extractelement <2 x i16> [[L]], i32 0
-; CHECK-NEXT:    [[R:%.*]] = insertelement <8 x i16> undef, i16 [[S]], i64 0
-; CHECK-NEXT:    ret <8 x i16> [[R]]
+; SSE2-LABEL: @gep1_load_v2i16_extract_insert_v8i16(
+; SSE2-NEXT:    [[GEP:%.*]] = getelementptr inbounds <2 x i16>, <2 x i16>* [[P:%.*]], i64 1
+; SSE2-NEXT:    [[L:%.*]] = load <2 x i16>, <2 x i16>* [[GEP]], align 2
+; SSE2-NEXT:    [[S:%.*]] = extractelement <2 x i16> [[L]], i32 0
+; SSE2-NEXT:    [[R:%.*]] = insertelement <8 x i16> undef, i16 [[S]], i64 0
+; SSE2-NEXT:    ret <8 x i16> [[R]]
+;
+; AVX2-LABEL: @gep1_load_v2i16_extract_insert_v8i16(
+; AVX2-NEXT:    [[TMP1:%.*]] = bitcast <2 x i16>* [[P:%.*]] to <8 x i16>*
+; AVX2-NEXT:    [[TMP2:%.*]] = load <8 x i16>, <8 x i16>* [[TMP1]], align 16
+; AVX2-NEXT:    [[R:%.*]] = shufflevector <8 x i16> [[TMP2]], <8 x i16> undef, <8 x i32> <i32 2, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>
+; AVX2-NEXT:    ret <8 x i16> [[R]]
 ;
   %gep = getelementptr inbounds <2 x i16>, <2 x i16>* %p, i64 1
   %l = load <2 x i16>, <2 x i16>* %gep, align 2
   %s = extractelement <2 x i16> %l, i32 0
   %r = insertelement <8 x i16> undef, i16 %s, i64 0
   ret <8 x i16> %r
 }